aljaz/Auto-GPT
1
0
Fork 0
mirror of https://github.com/aljazceru/Auto-GPT.git synced 2025-12-24 01:14:22 +01:00
Code Issues Packages Projects Releases Wiki Activity

Add workspace abstraction (#2982)

Browse Source 
* 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
Download Patch File Download Diff File Expand all files Collapse all files

30
tests/test_image_gen.py
View File

@@ -1,39 +1,46 @@
import hashlib
import os
import shutil
import unittest
from pathlib import Path
from PIL import Image
from autogpt.commands.image_gen import generate_image, generate_image_with_sd_webui
from autogpt.config import Config
from autogpt.workspace import path_in_workspace
from autogpt.workspace import Workspace
from tests.utils import requires_api_key
def lst(txt):
return txt.split(":")[1].strip()
return Path(txt.split(":")[1].strip())
@unittest.skipIf(os.getenv("CI"), "Skipping image generation tests")
@unittest.skip("Skipping image generation tests")
class TestImageGen(unittest.TestCase):
def setUp(self):
self.config = Config()
workspace_path = os.path.join(os.path.dirname(__file__), "workspace")
self.workspace_path = Workspace.make_workspace(workspace_path)
self.config.workspace_path = workspace_path
self.workspace = Workspace(workspace_path, restrict_to_workspace=True)
def tearDown(self) -> None:
shutil.rmtree(self.workspace_path)
@requires_api_key("OPENAI_API_KEY")
def test_dalle(self):
self.config.image_provider = "dalle"
# Test using size 256
result = lst(generate_image("astronaut riding a horse", 256))
image_path = path_in_workspace(result)
image_path = lst(generate_image("astronaut riding a horse", 256))
self.assertTrue(image_path.exists())
with Image.open(image_path) as img:
self.assertEqual(img.size, (256, 256))
image_path.unlink()
# Test using size 512
result = lst(generate_image("astronaut riding a horse", 512))
image_path = path_in_workspace(result)
image_path = lst(generate_image("astronaut riding a horse", 512))
with Image.open(image_path) as img:
self.assertEqual(img.size, (512, 512))
image_path.unlink()
@@ -44,8 +51,7 @@ class TestImageGen(unittest.TestCase):
# Test usin SD 1.4 model and size 512
self.config.huggingface_image_model = "CompVis/stable-diffusion-v1-4"
result = lst(generate_image("astronaut riding a horse", 512))
image_path = path_in_workspace(result)
image_path = lst(generate_image("astronaut riding a horse", 512))
self.assertTrue(image_path.exists())
with Image.open(image_path) as img:
self.assertEqual(img.size, (512, 512))
@@ -53,8 +59,7 @@ class TestImageGen(unittest.TestCase):
# Test using SD 2.1 768 model and size 768
self.config.huggingface_image_model = "stabilityai/stable-diffusion-2-1"
result = lst(generate_image("astronaut riding a horse", 768))
image_path = path_in_workspace(result)
image_path = lst(generate_image("astronaut riding a horse", 768))
with Image.open(image_path) as img:
self.assertEqual(img.size, (768, 768))
image_path.unlink()
@@ -64,8 +69,7 @@ class TestImageGen(unittest.TestCase):
return
# Test using size 128
result = lst(generate_image_with_sd_webui("astronaut riding a horse", 128))
image_path = path_in_workspace(result)
image_path = lst(generate_image_with_sd_webui("astronaut riding a horse", 128))
self.assertTrue(image_path.exists())
with Image.open(image_path) as img:
self.assertEqual(img.size, (128, 128))