From 1c21b2e2851fe79e45d19eec9668cd6706e2cb0a Mon Sep 17 00:00:00 2001 From: Ishaan Jaffer Date: Thu, 16 Oct 2025 17:52:37 -0700 Subject: [PATCH] fixes: AI/ML API --- .../aiml/image_generation/transformation.py | 26 +++++++++++++++---- 1 file changed, 21 insertions(+), 5 deletions(-) diff --git a/litellm/llms/aiml/image_generation/transformation.py b/litellm/llms/aiml/image_generation/transformation.py index 3b586689ea..006a2c16d7 100644 --- a/litellm/llms/aiml/image_generation/transformation.py +++ b/litellm/llms/aiml/image_generation/transformation.py @@ -172,16 +172,32 @@ class AimlImageGenerationConfig(BaseImageGenerationConfig): if not model_response.data: model_response.data = [] - # 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) + # AI/ML API can return images in multiple formats: + # 1. Top-level data array with url (OpenAI-like format) + # 2. output.choices array with image_base64 + # 3. images array with url (and optional width, height, content_type) - if "output" in response_data and "choices" in response_data["output"]: + if "data" in response_data and isinstance(response_data["data"], list): + # Handle OpenAI-like format: {"data": [{"url": "...", "width": 1024, "height": 768, "content_type": "image/jpeg"}]} + for image in response_data["data"]: + if "url" in image: + model_response.data.append(ImageObject( + b64_json=None, + url=image["url"], + revised_prompt=image.get("revised_prompt"), + )) + elif "b64_json" in image or "image_base64" in image: + model_response.data.append(ImageObject( + b64_json=image.get("b64_json") or image.get("image_base64"), + url=None, + revised_prompt=image.get("revised_prompt"), + )) + elif "output" in response_data and "choices" in response_data["output"]: for choice in response_data["output"]["choices"]: if "image_base64" in choice: model_response.data.append(ImageObject( b64_json=choice["image_base64"], - url=None, # AI/ML API returns base64, not URLs + url=None, )) elif "url" in choice: model_response.data.append(ImageObject(