From 835879bc8582154e2a37766ac40d14f5291f63d5 Mon Sep 17 00:00:00 2001
From: Ishaan Jaff <ishaanjaffer0324@gmail.com>
Date: Sat, 23 Aug 2025 14:03:11 -0700
Subject: [PATCH] fix: img gen

---
 litellm/images/main.py                        | 44 +++++++++----------
 .../aiml/image_generation/transformation.py   | 24 +++++++++-
 2 files changed, 44 insertions(+), 24 deletions(-)

diff --git a/litellm/images/main.py b/litellm/images/main.py
index b75a536beb4..953a8012e49 100644
--- a/litellm/images/main.py
+++ b/litellm/images/main.py
@@ -335,6 +335,28 @@ def image_generation(  # noqa: PLR0915
                 headers=headers,
                 litellm_params=litellm_params_dict,
             )
+        #########################################################
+        # Providers using llm_http_handler
+        #########################################################
+        elif custom_llm_provider in (
+            litellm.LlmProviders.RECRAFT,
+            litellm.LlmProviders.AIML,
+            litellm.LlmProviders.GEMINI,
+        ):
+            if image_generation_config is None:
+                raise ValueError(f"image generation config is not supported for {custom_llm_provider}")
+            
+            return llm_http_handler.image_generation_handler(
+                model=model,
+                prompt=prompt,
+                image_generation_provider_config=image_generation_config,
+                image_generation_optional_request_params=optional_params,
+                custom_llm_provider=custom_llm_provider,
+                litellm_params=litellm_params_dict,
+                logging_obj=litellm_logging_obj,
+                timeout=timeout,
+                client=client,
+            )
         elif custom_llm_provider == "azure_ai":
             from litellm.llms.azure_ai.common_utils import AzureFoundryModelInfo
             api_base = AzureFoundryModelInfo.get_api_base(api_base)
@@ -439,28 +461,6 @@ def image_generation(  # noqa: PLR0915
                 api_base=api_base,
                 client=client,
             )
-        #########################################################
-        # Providers using llm_http_handler
-        #########################################################
-        elif custom_llm_provider in (
-            litellm.LlmProviders.RECRAFT,
-            litellm.LlmProviders.AIML,
-            litellm.LlmProviders.GEMINI,
-        ):
-            if image_generation_config is None:
-                raise ValueError(f"image generation config is not supported for {custom_llm_provider}")
-            
-            return llm_http_handler.image_generation_handler(
-                model=model,
-                prompt=prompt,
-                image_generation_provider_config=image_generation_config,
-                image_generation_optional_request_params=optional_params,
-                custom_llm_provider=custom_llm_provider,
-                litellm_params=litellm_params_dict,
-                logging_obj=litellm_logging_obj,
-                timeout=timeout,
-                client=client,
-            )
         elif (
             custom_llm_provider in litellm._custom_providers
         ):  # Assume custom LLM provider
diff --git a/litellm/llms/aiml/image_generation/transformation.py b/litellm/llms/aiml/image_generation/transformation.py
index 2686f7a267c..3b586689ea7 100644
--- a/litellm/llms/aiml/image_generation/transformation.py
+++ b/litellm/llms/aiml/image_generation/transformation.py
@@ -172,7 +172,10 @@ class AimlImageGenerationConfig(BaseImageGenerationConfig):
         if not model_response.data:
             model_response.data = []
         
-        # AI/ML API returns images in output.choices array with image_base64
+        # AI/ML API can return images in two different formats:
+        # 1. output.choices array with image_base64
+        # 2. images array with url (and optional width, height, content_type)
+        
         if "output" in response_data and "choices" in response_data["output"]:
             for choice in response_data["output"]["choices"]:
                 if "image_base64" in choice:
@@ -180,5 +183,22 @@ class AimlImageGenerationConfig(BaseImageGenerationConfig):
                         b64_json=choice["image_base64"],
                         url=None,  # AI/ML API returns base64, not URLs
                     ))
-        
+                elif "url" in choice:
+                    model_response.data.append(ImageObject(
+                        b64_json=None,
+                        url=choice["url"],
+                    ))
+        elif "images" in response_data:
+            # Handle alternative format: {"images": [{"url": "...", "width": 1024, "height": 768, "content_type": "image/jpeg"}]}
+            for image in response_data["images"]:
+                if "url" in image:
+                    model_response.data.append(ImageObject(
+                        b64_json=None,
+                        url=image["url"],
+                    ))
+                elif "image_base64" in image:
+                    model_response.data.append(ImageObject(
+                        b64_json=image["image_base64"],
+                        url=None,
+                    ))
         return model_response