chore: initial public snapshot for github upload
This commit is contained in:
Your Name
@@ -0,0 +1,13 @@
|
||||
from litellm.llms.base_llm.image_generation.transformation import (
|
||||
BaseImageGenerationConfig,
|
||||
)
|
||||
|
||||
from .transformation import AimlImageGenerationConfig
|
||||
|
||||
__all__ = [
|
||||
"AimlImageGenerationConfig",
|
||||
]
|
||||
|
||||
|
||||
def get_aiml_image_generation_config(model: str) -> BaseImageGenerationConfig:
|
||||
return AimlImageGenerationConfig()
|
||||
@@ -0,0 +1,27 @@
|
||||
from typing import Any
|
||||
|
||||
import litellm
|
||||
from litellm.types.utils import ImageResponse
|
||||
|
||||
|
||||
def cost_calculator(
|
||||
model: str,
|
||||
image_response: Any,
|
||||
) -> float:
|
||||
"""
|
||||
AI/ML flux image generation cost calculator
|
||||
"""
|
||||
_model_info = litellm.get_model_info(
|
||||
model=model,
|
||||
custom_llm_provider=litellm.LlmProviders.AIML.value,
|
||||
)
|
||||
output_cost_per_image: float = _model_info.get("output_cost_per_image") or 0.0
|
||||
num_images: int = 0
|
||||
if isinstance(image_response, ImageResponse):
|
||||
if image_response.data:
|
||||
num_images = len(image_response.data)
|
||||
return output_cost_per_image * num_images
|
||||
else:
|
||||
raise ValueError(
|
||||
f"image_response must be of type ImageResponse got type={type(image_response)}"
|
||||
)
|
||||
@@ -0,0 +1,234 @@
|
||||
from typing import TYPE_CHECKING, Any, List, Optional
|
||||
|
||||
import httpx
|
||||
|
||||
from litellm.llms.base_llm.image_generation.transformation import (
|
||||
BaseImageGenerationConfig,
|
||||
)
|
||||
from litellm.secret_managers.main import get_secret_str
|
||||
from litellm.types.llms.aiml import AimlImageGenerationRequestParams
|
||||
from litellm.types.llms.openai import (
|
||||
AllMessageValues,
|
||||
OpenAIImageGenerationOptionalParams,
|
||||
)
|
||||
from litellm.types.utils import ImageObject, ImageResponse
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from litellm.litellm_core_utils.litellm_logging import Logging as _LiteLLMLoggingObj
|
||||
|
||||
LiteLLMLoggingObj = _LiteLLMLoggingObj
|
||||
else:
|
||||
LiteLLMLoggingObj = Any
|
||||
|
||||
|
||||
class AimlImageGenerationConfig(BaseImageGenerationConfig):
|
||||
DEFAULT_BASE_URL: str = "https://api.aimlapi.com"
|
||||
IMAGE_GENERATION_ENDPOINT: str = "v1/images/generations"
|
||||
|
||||
def get_supported_openai_params(
|
||||
self, model: str
|
||||
) -> List[OpenAIImageGenerationOptionalParams]:
|
||||
"""
|
||||
https://api.aimlapi.com/v1/images/generations
|
||||
"""
|
||||
return ["n", "response_format", "size"]
|
||||
|
||||
def map_openai_params(
|
||||
self,
|
||||
non_default_params: dict,
|
||||
optional_params: dict,
|
||||
model: str,
|
||||
drop_params: bool,
|
||||
) -> dict:
|
||||
supported_params = self.get_supported_openai_params(model)
|
||||
|
||||
for k in non_default_params.keys():
|
||||
if k not in optional_params.keys():
|
||||
if k in supported_params:
|
||||
# Map OpenAI params to AI/ML params
|
||||
if k == "n":
|
||||
optional_params["num_images"] = non_default_params[k]
|
||||
elif k == "response_format":
|
||||
optional_params["output_format"] = non_default_params[k]
|
||||
elif k == "size":
|
||||
# Map OpenAI size format to AI/ML image_size
|
||||
size_value = non_default_params[k]
|
||||
if isinstance(size_value, str):
|
||||
# Handle standard OpenAI sizes like "1024x1024"
|
||||
if "x" in size_value:
|
||||
width, height = map(int, size_value.split("x"))
|
||||
optional_params["image_size"] = {
|
||||
"width": width,
|
||||
"height": height,
|
||||
}
|
||||
else:
|
||||
# Pass through predefined sizes
|
||||
optional_params["image_size"] = size_value
|
||||
else:
|
||||
optional_params["image_size"] = size_value
|
||||
else:
|
||||
optional_params[k] = non_default_params[k]
|
||||
elif drop_params:
|
||||
pass
|
||||
else:
|
||||
raise ValueError(
|
||||
f"Parameter {k} is not supported for model {model}. Supported parameters are {supported_params}. Set drop_params=True to drop unsupported parameters."
|
||||
)
|
||||
|
||||
return optional_params
|
||||
|
||||
def get_complete_url(
|
||||
self,
|
||||
api_base: Optional[str],
|
||||
api_key: Optional[str],
|
||||
model: str,
|
||||
optional_params: dict,
|
||||
litellm_params: dict,
|
||||
stream: Optional[bool] = None,
|
||||
) -> str:
|
||||
"""
|
||||
Get the complete url for the request
|
||||
"""
|
||||
complete_url: str = (
|
||||
api_base or get_secret_str("AIML_API_BASE") or self.DEFAULT_BASE_URL
|
||||
)
|
||||
|
||||
complete_url = complete_url.rstrip("/")
|
||||
# Strip /v1 suffix if present since IMAGE_GENERATION_ENDPOINT already includes v1
|
||||
if complete_url.endswith("/v1"):
|
||||
complete_url = complete_url[:-3]
|
||||
complete_url = f"{complete_url}/{self.IMAGE_GENERATION_ENDPOINT}"
|
||||
return complete_url
|
||||
|
||||
def validate_environment(
|
||||
self,
|
||||
headers: dict,
|
||||
model: str,
|
||||
messages: List[AllMessageValues],
|
||||
optional_params: dict,
|
||||
litellm_params: dict,
|
||||
api_key: Optional[str] = None,
|
||||
api_base: Optional[str] = None,
|
||||
) -> dict:
|
||||
final_api_key: Optional[str] = (
|
||||
api_key
|
||||
or get_secret_str("AIML_API_KEY")
|
||||
or get_secret_str("AIMLAPI_KEY") # Alternative name
|
||||
)
|
||||
if not final_api_key:
|
||||
raise ValueError("AIML_API_KEY or AIMLAPI_KEY is not set")
|
||||
|
||||
headers["Authorization"] = f"Bearer {final_api_key}"
|
||||
headers["Content-Type"] = "application/json"
|
||||
return headers
|
||||
|
||||
def transform_image_generation_request(
|
||||
self,
|
||||
model: str,
|
||||
prompt: str,
|
||||
optional_params: dict,
|
||||
litellm_params: dict,
|
||||
headers: dict,
|
||||
) -> dict:
|
||||
"""
|
||||
Transform the image generation request to the AI/ML flux image generation request body
|
||||
|
||||
https://api.aimlapi.com/v1/images/generations
|
||||
"""
|
||||
aiml_image_generation_request_body: AimlImageGenerationRequestParams = (
|
||||
AimlImageGenerationRequestParams(
|
||||
prompt=prompt,
|
||||
model=model,
|
||||
**optional_params,
|
||||
)
|
||||
)
|
||||
return dict(aiml_image_generation_request_body)
|
||||
|
||||
def transform_image_generation_response(
|
||||
self,
|
||||
model: str,
|
||||
raw_response: httpx.Response,
|
||||
model_response: ImageResponse,
|
||||
logging_obj: LiteLLMLoggingObj,
|
||||
request_data: dict,
|
||||
optional_params: dict,
|
||||
litellm_params: dict,
|
||||
encoding: Any,
|
||||
api_key: Optional[str] = None,
|
||||
json_mode: Optional[bool] = None,
|
||||
) -> ImageResponse:
|
||||
"""
|
||||
Transform the image generation response to the litellm image response
|
||||
|
||||
https://api.aimlapi.com/v1/images/generations
|
||||
"""
|
||||
try:
|
||||
response_data = raw_response.json()
|
||||
except Exception as e:
|
||||
raise self.get_error_class(
|
||||
error_message=f"Error transforming image generation response: {e}",
|
||||
status_code=raw_response.status_code,
|
||||
headers=raw_response.headers,
|
||||
)
|
||||
|
||||
if not model_response.data:
|
||||
model_response.data = []
|
||||
|
||||
# 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 "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,
|
||||
)
|
||||
)
|
||||
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
|
||||
Reference in New Issue
Block a user