diff --git a/autogpt/commands/image_gen.py b/autogpt/commands/image_gen.py index 0832a067c995..0809fcdd3e38 100644 --- a/autogpt/commands/image_gen.py +++ b/autogpt/commands/image_gen.py @@ -57,7 +57,7 @@ def generate_image_with_hf(prompt: str, filename: str) -> str: ) headers = { "Authorization": f"Bearer {CFG.huggingface_api_token}", - "X-Use-Cache": "false" + "X-Use-Cache": "false", } response = requests.post( @@ -91,7 +91,9 @@ def generate_image_with_dalle(prompt: str, filename: str) -> str: # Check for supported image sizes if size not in [256, 512, 1024]: closest = min([256, 512, 1024], key=lambda x: abs(x - size)) - print(f"DALL-E only supports image sizes of 256x256, 512x512, or 1024x1024. Setting to {closest}, was {size}.") + print( + f"DALL-E only supports image sizes of 256x256, 512x512, or 1024x1024. Setting to {closest}, was {size}." + ) size = closest response = openai.Image.create( @@ -111,7 +113,13 @@ def generate_image_with_dalle(prompt: str, filename: str) -> str: return f"Saved to disk:{filename}" -def generate_image_with_sd_webui(prompt: str, filename: str, size: int = 512, negative_prompt: str = "", extra: dict = {}) -> str: +def generate_image_with_sd_webui( + prompt: str, + filename: str, + size: int = 512, + negative_prompt: str = "", + extra: dict = {}, +) -> str: """Generate an image with Stable Diffusion webui. Args: prompt (str): The prompt to use diff --git a/autogpt/config/config.py b/autogpt/config/config.py index e19c85d5239f..4b53df10e8d2 100644 --- a/autogpt/config/config.py +++ b/autogpt/config/config.py @@ -87,7 +87,9 @@ def __init__(self) -> None: self.image_provider = os.getenv("IMAGE_PROVIDER") self.image_size = int(os.getenv("IMAGE_SIZE", 256)) self.huggingface_api_token = os.getenv("HUGGINGFACE_API_TOKEN") - self.huggingface_image_model = os.getenv("HUGGINGFACE_IMAGE_MODEL", "CompVis/stable-diffusion-v1-4") + self.huggingface_image_model = os.getenv( + "HUGGINGFACE_IMAGE_MODEL", "CompVis/stable-diffusion-v1-4" + ) self.huggingface_audio_to_text_model = os.getenv( "HUGGINGFACE_AUDIO_TO_TEXT_MODEL" ) diff --git a/tests/test_image_gen.py b/tests/test_image_gen.py index c1cb3f925dcd..19c57e427d5c 100644 --- a/tests/test_image_gen.py +++ b/tests/test_image_gen.py @@ -1,9 +1,11 @@ -import unittest import hashlib -from PIL import Image import os -from autogpt.config import Config +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 @@ -67,7 +69,14 @@ def test_sd_webui(self): 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})) + 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)) @@ -75,7 +84,11 @@ def test_sd_webui(self): 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})) + 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))