diff --git a/jigsawstack/image_generation.py b/jigsawstack/image_generation.py
index 08cf81c..d3bc860 100644
--- a/jigsawstack/image_generation.py
+++ b/jigsawstack/image_generation.py
@@ -3,7 +3,7 @@
 from typing_extensions import Literal, NotRequired, Required, TypedDict
 
 from ._config import ClientConfig
-from .async_request import AsyncRequest
+from .async_request import AsyncRequest, AsyncRequestConfig
 from .request import Request, RequestConfig
 
 
@@ -103,6 +103,7 @@ def image_generation(
         self, params: ImageGenerationParams
     ) -> Union[ImageGenerationResponse, bytes]:
         path = "/ai/image_generation"
+
         resp = Request(
             config=self.config,
             path=path,
@@ -113,7 +114,7 @@ def image_generation(
 
 
 class AsyncImageGeneration(ClientConfig):
-    config: RequestConfig
+    config: AsyncRequestConfig
 
     def __init__(
         self,
@@ -122,7 +123,7 @@ def __init__(
         headers: Union[Dict[str, str], None] = None,
     ):
         super().__init__(api_key, base_url, headers)
-        self.config = RequestConfig(
+        self.config = AsyncRequestConfig(
             base_url=base_url,
             api_key=api_key,
             headers=headers,
diff --git a/tests/test_image_generation.py b/tests/test_image_generation.py
index e11bf76..72c35f2 100644
--- a/tests/test_image_generation.py
+++ b/tests/test_image_generation.py
@@ -28,10 +28,10 @@
 )
 
 IMAGE_URL = "https://images.unsplash.com/photo-1494588024300-e9df7ff98d78?q=80&w=1284&auto=format&fit=crop&ixlib=rb-4.1.0&ixid=M3wxMjA3fDB8MHxwaG90by1wYWdlfHx8fGVufDB8fHx8fA%3D%3D"
-FILE_STORE_KEY = jigsaw.store.upload(
-    requests.get(IMAGE_URL).content,
-    {"filename": "test_image.jpg", "content_type": "image/jpeg", "overwrite": True},
-)
+# FILE_STORE_KEY = jigsaw.store.upload(
+#     requests.get(IMAGE_URL).content,
+#     {"filename": "test_image.jpg", "content_type": "image/jpeg", "overwrite": True},
+# )
 
 TEST_CASES = [
     {
@@ -110,13 +110,13 @@
             "return_type": "base64",
         },
     },
-    {
-        "name": "with_file_store_key",
-        "params": {
-            "prompt": "Apply a cyberpunk style to this image",
-            "file_store_key": FILE_STORE_KEY,
-        },
-    },
+    # {
+    #     "name": "with_file_store_key",
+    #     "params": {
+    #         "prompt": "Apply a cyberpunk style to this image",
+    #         "file_store_key": FILE_STORE_KEY,
+    #     },
+    # },
 ]