Add workspace abstraction (#2982)

* Add workspace abstraction

* Remove old workspace implementation

* Extract path resolution to a helper function

* Add api key requirements to new tests
This commit is contained in:
James Collins
2023-04-23 12:36:04 -07:00
committed by GitHub
parent da48f9c972
commit dcd6aa912b
19 changed files with 379 additions and 196 deletions

View File

@@ -9,7 +9,6 @@ from PIL import Image
from autogpt.commands.command import command
from autogpt.config import Config
from autogpt.workspace import path_in_workspace
CFG = Config()
@@ -25,7 +24,7 @@ def generate_image(prompt: str, size: int = 256) -> str:
Returns:
str: The filename of the image
"""
filename = f"{str(uuid.uuid4())}.jpg"
filename = f"{CFG.workspace_path}/{str(uuid.uuid4())}.jpg"
# DALL-E
if CFG.image_provider == "dalle":
@@ -72,7 +71,7 @@ def generate_image_with_hf(prompt: str, filename: str) -> str:
image = Image.open(io.BytesIO(response.content))
print(f"Image Generated for prompt:{prompt}")
image.save(path_in_workspace(filename))
image.save(filename)
return f"Saved to disk:{filename}"
@@ -109,7 +108,7 @@ def generate_image_with_dalle(prompt: str, filename: str, size: int) -> str:
image_data = b64decode(response["data"][0]["b64_json"])
with open(path_in_workspace(filename), mode="wb") as png:
with open(filename, mode="wb") as png:
png.write(image_data)
return f"Saved to disk:{filename}"
@@ -160,6 +159,6 @@ def generate_image_with_sd_webui(
response = response.json()
b64 = b64decode(response["images"][0].split(",", 1)[0])
image = Image.open(io.BytesIO(b64))
image.save(path_in_workspace(filename))
image.save(filename)
return f"Saved to disk:{filename}"