From aeb1178a47906a11c0017813e44c81fea3d0c5ae Mon Sep 17 00:00:00 2001
From: BillSchumacher <34168009+BillSchumacher@users.noreply.github.com>
Date: Tue, 18 Apr 2023 19:26:18 -0500
Subject: [PATCH] linting

---
 autogpt/commands/image_gen.py | 14 +++++++++++---
 autogpt/config/config.py      |  4 +++-
 tests/test_image_gen.py       | 23 ++++++++++++++++++-----
 3 files changed, 32 insertions(+), 9 deletions(-)

diff --git a/autogpt/commands/image_gen.py b/autogpt/commands/image_gen.py
index 0832a067..0809fcdd 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 e19c85d5..4b53df10 100644
--- a/autogpt/config/config.py
+++ b/autogpt/config/config.py
@@ -87,7 +87,9 @@ class Config(metaclass=Singleton):
         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 c1cb3f92..19c57e42 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 @@ class TestImageGen(unittest.TestCase):
         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 @@ class TestImageGen(unittest.TestCase):
         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))