Spaces:
Running
Running
File size: 3,491 Bytes
0534c31 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 |
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()
|