linting

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

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"
         )

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,15 +69,26 @@ 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))
             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}))
+        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))