coding-alt/AutoGPT

1import hashlib2import os3import unittest4 5from PIL import Image6 7from autogpt.commands.image_gen import generate_image, generate_image_with_sd_webui8from autogpt.config import Config9from autogpt.workspace import path_in_workspace10 11 12def lst(txt):13    return txt.split(":")[1].strip()14 15 16@unittest.skipIf(os.getenv("CI"), "Skipping image generation tests")17class TestImageGen(unittest.TestCase):18    def setUp(self):19        self.config = Config()20 21    def test_dalle(self):22        self.config.image_provider = "dalle"23 24        # Test using size 25625        result = lst(generate_image("astronaut riding a horse", 256))26        image_path = path_in_workspace(result)27        self.assertTrue(image_path.exists())28        with Image.open(image_path) as img:29            self.assertEqual(img.size, (256, 256))30        image_path.unlink()31 32        # Test using size 51233        result = lst(generate_image("astronaut riding a horse", 512))34        image_path = path_in_workspace(result)35        with Image.open(image_path) as img:36            self.assertEqual(img.size, (512, 512))37        image_path.unlink()38 39    def test_huggingface(self):40        self.config.image_provider = "huggingface"41 42        # Test usin SD 1.4 model and size 51243        self.config.huggingface_image_model = "CompVis/stable-diffusion-v1-4"44        result = lst(generate_image("astronaut riding a horse", 512))45        image_path = path_in_workspace(result)46        self.assertTrue(image_path.exists())47        with Image.open(image_path) as img:48            self.assertEqual(img.size, (512, 512))49        image_path.unlink()50 51        # Test using SD 2.1 768 model and size 76852        self.config.huggingface_image_model = "stabilityai/stable-diffusion-2-1"53        result = lst(generate_image("astronaut riding a horse", 768))54        image_path = path_in_workspace(result)55        with Image.open(image_path) as img:56            self.assertEqual(img.size, (768, 768))57        image_path.unlink()58 59    def test_sd_webui(self):60        self.config.image_provider = "sd_webui"61        return62 63        # Test using size 12864        result = lst(generate_image_with_sd_webui("astronaut riding a horse", 128))65        image_path = path_in_workspace(result)66        self.assertTrue(image_path.exists())67        with Image.open(image_path) as img:68            self.assertEqual(img.size, (128, 128))69        image_path.unlink()70 71        # Test using size 64 and negative prompt72        result = lst(73            generate_image_with_sd_webui(74                "astronaut riding a horse",75                negative_prompt="horse",76                size=64,77                extra={"seed": 123},78            )79        )80        image_path = path_in_workspace(result)81        with Image.open(image_path) as img:82            self.assertEqual(img.size, (64, 64))83            neg_image_hash = hashlib.md5(img.tobytes()).hexdigest()84        image_path.unlink()85 86        # Same test as above but without the negative prompt87        result = lst(88            generate_image_with_sd_webui(89                "astronaut riding a horse", image_size=64, size=1, extra={"seed": 123}90            )91        )92        image_path = path_in_workspace(result)93        with Image.open(image_path) as img:94            self.assertEqual(img.size, (64, 64))95            image_hash = hashlib.md5(img.tobytes()).hexdigest()96        image_path.unlink()97 98        self.assertNotEqual(image_hash, neg_image_hash)99 100 101if __name__ == "__main__":102    unittest.main()103