Merge branch 'master' of https://github.com/BillSchumacher/Auto-GPT into plugin-support

tests/test_image_gen.py

@@ -0,0 +1,102 @@
+import hashlib
+import os
+import unittest
+
+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
+
+
+def lst(txt):
+    return txt.split(":")[1].strip()
+
+
+@unittest.skipIf(os.getenv("CI"), "Skipping image generation tests")
+class TestImageGen(unittest.TestCase):
+    def setUp(self):
+        self.config = Config()
+
+    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)
+        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)
+        with Image.open(image_path) as img:
+            self.assertEqual(img.size, (512, 512))
+        image_path.unlink()
+
+    def test_huggingface(self):
+        self.config.image_provider = "huggingface"
+
+        # 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)
+        self.assertTrue(image_path.exists())
+        with Image.open(image_path) as img:
+            self.assertEqual(img.size, (512, 512))
+        image_path.unlink()
+
+        # 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)
+        with Image.open(image_path) as img:
+            self.assertEqual(img.size, (768, 768))
+        image_path.unlink()
+
+    def test_sd_webui(self):
+        self.config.image_provider = "sd_webui"
+        return
+
+        # Test using size 128
+        result = lst(generate_image_with_sd_webui("astronaut riding a horse", 128))
+        image_path = path_in_workspace(result)
+        self.assertTrue(image_path.exists())
+        with Image.open(image_path) as img:
+            self.assertEqual(img.size, (128, 128))
+        image_path.unlink()
+
+        # Test using size 64 and negative prompt
+        result = lst(
+            generate_image_with_sd_webui(
+                "astronaut riding a horse",
+                negative_prompt="horse",
+                size=64,
+                extra={"seed": 123},
+            )
+        )
+        image_path = path_in_workspace(result)
+        with Image.open(image_path) as img:
+            self.assertEqual(img.size, (64, 64))
+            neg_image_hash = hashlib.md5(img.tobytes()).hexdigest()
+        image_path.unlink()
+
+        # Same test as above but without the negative prompt
+        result = lst(
+            generate_image_with_sd_webui(
+                "astronaut riding a horse", image_size=64, size=1, extra={"seed": 123}
+            )
+        )
+        image_path = path_in_workspace(result)
+        with Image.open(image_path) as img:
+            self.assertEqual(img.size, (64, 64))
+            image_hash = hashlib.md5(img.tobytes()).hexdigest()
+        image_path.unlink()
+
+        self.assertNotEqual(image_hash, neg_image_hash)
+
+
+if __name__ == "__main__":
+    unittest.main()

tests/test_json_parser.py

@@ -1,7 +1,7 @@
 import unittest
 
 import tests.context
-from autogpt.json_fixes.parsing import fix_and_parse_json
+from autogpt.json_utils.json_fix_llm import fix_and_parse_json
 
 
 class TestParseJson(unittest.TestCase):

tests/unit/json_tests.py

@@ -1,6 +1,6 @@
 import unittest
 
-from autogpt.json_parser import fix_and_parse_json
+from autogpt.json_utils.json_fix_llm import fix_and_parse_json
 
 
 class TestParseJson(unittest.TestCase):