scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-07 02:59:05 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 103

fix: img gen

Browse source
This commit is contained in:
Ishaan Jaff 2025-08-23 14:03:11 -07:00
parent f1c710a608
commit 835879bc85
2 changed files with 44 additions and 24 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

44
litellm/images/main.py
View file

@ -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

24
litellm/llms/aiml/image_generation/transformation.py
View file

@ -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