diff --git a/litellm/__init__.py b/litellm/__init__.py index a1a12fc9313..9b69beccd79 100644 --- a/litellm/__init__.py +++ b/litellm/__init__.py @@ -1203,9 +1203,9 @@ from .llms.bedrock.chat.invoke_transformations.amazon_openai_transformation impo AmazonBedrockOpenAIConfig, ) -from .llms.bedrock.image.amazon_stability1_transformation import AmazonStabilityConfig -from .llms.bedrock.image.amazon_stability3_transformation import AmazonStability3Config -from .llms.bedrock.image.amazon_nova_canvas_transformation import AmazonNovaCanvasConfig +from .llms.bedrock.image_generation.amazon_stability1_transformation import AmazonStabilityConfig +from .llms.bedrock.image_generation.amazon_stability3_transformation import AmazonStability3Config +from .llms.bedrock.image_generation.amazon_nova_canvas_transformation import AmazonNovaCanvasConfig from .llms.bedrock.embed.amazon_titan_g1_transformation import AmazonTitanG1Config from .llms.bedrock.embed.amazon_titan_multimodal_transformation import ( AmazonTitanMultimodalEmbeddingG1Config, diff --git a/litellm/images/utils.py b/litellm/images/utils.py index fdf240ba2af..fa271b61b6a 100644 --- a/litellm/images/utils.py +++ b/litellm/images/utils.py @@ -82,7 +82,6 @@ class ImageEditRequestUtils: filtered_params = { k: v for k, v in params.items() if k in valid_keys and v is not None } - return cast(ImageEditOptionalRequestParams, filtered_params) @staticmethod diff --git a/litellm/llms/bedrock/image/cost_calculator.py b/litellm/llms/bedrock/image/cost_calculator.py index bc1a57b8aec..b04acc3e809 100644 --- a/litellm/llms/bedrock/image/cost_calculator.py +++ b/litellm/llms/bedrock/image/cost_calculator.py @@ -1,6 +1,6 @@ from typing import Optional -from litellm.llms.bedrock.image.image_handler import BedrockImageGeneration +from litellm.llms.bedrock.image_generation.image_handler import BedrockImageGeneration from litellm.types.utils import ImageResponse diff --git a/litellm/llms/bedrock/image/image_handler.py b/litellm/llms/bedrock/image/image_handler.py index 2e76596eefe..0a4cde90b27 100644 --- a/litellm/llms/bedrock/image/image_handler.py +++ b/litellm/llms/bedrock/image/image_handler.py @@ -9,13 +9,13 @@ from pydantic import BaseModel import litellm from litellm._logging import verbose_logger from litellm.litellm_core_utils.litellm_logging import Logging as LitellmLogging -from litellm.llms.bedrock.image.amazon_nova_canvas_transformation import ( +from litellm.llms.bedrock.image_generation.amazon_nova_canvas_transformation import ( AmazonNovaCanvasConfig, ) -from litellm.llms.bedrock.image.amazon_stability3_transformation import ( +from litellm.llms.bedrock.image_generation.amazon_stability3_transformation import ( AmazonStability3Config, ) -from litellm.llms.bedrock.image.amazon_titan_transformation import ( +from litellm.llms.bedrock.image_generation.amazon_titan_transformation import ( AmazonTitanImageGenerationConfig, ) from litellm.llms.custom_httpx.http_handler import ( diff --git a/litellm/llms/custom_httpx/llm_http_handler.py b/litellm/llms/custom_httpx/llm_http_handler.py index a777ea3c42e..34ea598a655 100644 --- a/litellm/llms/custom_httpx/llm_http_handler.py +++ b/litellm/llms/custom_httpx/llm_http_handler.py @@ -3761,7 +3761,7 @@ class BaseLLMHTTPHandler: input=prompt, api_key="", additional_args={ - "complete_input_dict": data, + "complete_input_dict": files, "api_base": api_base, "headers": headers, }, diff --git a/litellm/llms/stability/image_edit/__init__.py b/litellm/llms/stability/image_edit/__init__.py new file mode 100644 index 00000000000..7710f080168 --- /dev/null +++ b/litellm/llms/stability/image_edit/__init__.py @@ -0,0 +1,37 @@ +""" +Stability AI Image Edit Module + +Factory function for getting the appropriate config class. +""" + +from litellm.llms.base_llm.image_edit.transformation import ( + BaseImageEditConfig, +) + +from .transformation import StabilityImageEditConfig + +__all__ = [ + "StabilityImageEditConfig", + "get_stability_image_edit_config", +] + + +def get_stability_image_edit_config(model: str) -> BaseImageEditConfig: + """ + Get the appropriate Stability AI config for the given model. + + Currently all models use the same config class, but this factory + allows for model-specific configs in the future. + + Args: + model: The model name (e.g., "stability/inpaint", "stability/outpaint") + + Returns: + BaseImageEditConfig instance for Stability AI + """ + # For now, all models use the same config + # In the future, we could have model-specific configs: + # - StabilityInpaintConfig for Inpaint models + # - StabilityOutpaintConfig for Outpaint models + # - etc. + return StabilityImageEditConfig() diff --git a/litellm/llms/stability/image_edit/transformations.py b/litellm/llms/stability/image_edit/transformations.py new file mode 100644 index 00000000000..fabd1e72054 --- /dev/null +++ b/litellm/llms/stability/image_edit/transformations.py @@ -0,0 +1,291 @@ +""" +Stability AI Image Edit Config + +Handles transformation between OpenAI-compatible format and Stability AI API format. + +API Reference: https://platform.stability.ai/docs/api-reference +""" + +from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple + +import httpx +from httpx._types import RequestFiles + +from litellm.llms.base_llm.image_edit.transformation import BaseImageEditConfig +from litellm.secret_managers.main import get_secret_str +from litellm.types.llms.openai import ( + OpenAIImageEditOptionalParams, +) +from litellm.types.llms.stability import ( + OPENAI_SIZE_TO_STABILITY_ASPECT_RATIO, + STABILITY_EDIT_ENDPOINTS, + StabilityImageEditRequest, +) +from litellm.types.utils import FileTypes, ImageObject, ImageResponse +from litellm.utils import get_model_info + +if TYPE_CHECKING: + from litellm.litellm_core_utils.litellm_logging import Logging as _LiteLLMLoggingObj + + LiteLLMLoggingObj = _LiteLLMLoggingObj +else: + LiteLLMLoggingObj = Any + + +class StabilityImageEditConfig(BaseImageEditConfig): + """ + Configuration for Stability AI image edit. + + Supports: + - Stable Diffusion 3 (SD3, SD3.5) Image Edit + """ + + DEFAULT_BASE_URL: str = "https://api.stability.ai" + + def get_supported_openai_params( + self, model: str + ) -> List[OpenAIImageEditOptionalParams]: + """ + Return list of OpenAI params supported by Stability AI. + + https://platform.stability.ai/docs/api-reference + """ + return [ + "n", # Number of images (Stability always returns 1, we can loop) + "size", # Maps to aspect_ratio + "response_format", # b64_json or url (Stability only returns b64) + "mask" + ] + + def map_openai_params( + self, + image_edit_optional_params: dict, + model: str, + drop_params: bool, + ) -> dict: + """ + Map OpenAI parameters to Stability AI parameters. + + OpenAI -> Stability mappings: + - size -> aspect_ratio + - n -> (handled separately, Stability returns 1 image per request) + """ + supported_params = self.get_supported_openai_params(model) + # Define mapping from OpenAI params to Stability params + param_mapping = { + "size": "aspect_ratio", + # "n" and "response_format" are handled separately + } + + # Create a copy to not mutate original + mapped_params = image_edit_optional_params.copy() + + for k, v in image_edit_optional_params.items(): + if k in param_mapping: + # Map param if mapping exists and value is valid + if k == "size" and v in OPENAI_SIZE_TO_STABILITY_ASPECT_RATIO: + mapped_params[param_mapping[k]] = OPENAI_SIZE_TO_STABILITY_ASPECT_RATIO[v] + # Don't copy "size" itself to final dict + elif k == "n": + # Store for logic but do not add to outgoing params + mapped_params["_n"] = v + elif k == "response_format": + # Only b64 supported at Stability; store for postprocessing + mapped_params["_response_format"] = v + elif k not in supported_params: + if not drop_params: + raise ValueError( + f"Parameter {k} is not supported for model {model}. " + f"Supported parameters are {supported_params}. " + f"Set drop_params=True to drop unsupported parameters." + ) + # Otherwise, param will simply be dropped + else: + # param is supported and not mapped, keep as-is + continue + + # Remove OpenAI params that have been mapped unless they're in stability + for mapped in ["size", "n", "response_format"]: + if mapped in mapped_params: + del mapped_params[mapped] + + return mapped_params + + def _get_model_endpoint(self, model: str) -> str: + """ + Get the API endpoint for a given model. + """ + # Remove "stability/" prefix if present + model_name = model.lower() + if model_name.startswith("stability/"): + model_name = model_name[10:] # Remove "stability/" prefix + + # Check if model is in our mapping + for key, endpoint in STABILITY_EDIT_ENDPOINTS.items(): + if key in model_name: + return endpoint + + # Default to SD3 endpoint + return "/v2beta/stable-image/edit/inpaint" + + def get_complete_url( + self, + api_base: Optional[str], + model: str, + litellm_params: dict, + ) -> str: + """ + Get the complete URL for the Stability AI API request. + """ + base_url: str = ( + api_base + or get_secret_str("STABILITY_API_BASE") + or litellm_params.get("api_base", None) + or self.DEFAULT_BASE_URL + ) + base_url = base_url.rstrip("/") + + endpoint = self._get_model_endpoint(model) + return f"{base_url}{endpoint}" + + def validate_environment( + self, + headers: dict, + model: str, + api_key: Optional[str] = None, + ) -> dict: + """ + Validate environment and set up headers for Stability AI. + """ + final_api_key: Optional[str] = api_key or get_secret_str("STABILITY_API_KEY") + + if not final_api_key: + raise ValueError( + "STABILITY_API_KEY is not set. " + "Please set it via environment variable or pass api_key parameter." + ) + + headers["Authorization"] = f"Bearer {final_api_key}" + headers["Accept"] = "application/json" + return headers + + def transform_image_edit_request( + self, + model: str, + prompt: str, + image: FileTypes, + image_edit_optional_request_params: OpenAIImageEditOptionalParams, + litellm_params: dict, + headers: dict, + ) -> Tuple[Dict, RequestFiles]: + """ + Transform OpenAI-style request to Stability AI request format. + + Note: Stability AI uses multipart/form-data, but the HTTP handler + will handle the conversion from dict to form data. + """ + # Build Stability request + # Populate multipart form-data as separate text fields (data) and files. + # Stability expects prompt/output_format/etc. as normal form fields, not file parts. + data: StabilityImageEditRequest = { + "prompt": prompt, + "output_format": "png", # Default to PNG + } + files: Dict[str, Any] = {"image": image[0]} + + # Add optional params (already mapped in map_openai_params) + for key, value in image_edit_optional_request_params.items(): + # Skip internal params (prefixed with _) + if key.startswith("_") or value is None: + continue + + # File-like optional param + if key == "mask": + files["mask"] = value # type: ignore + continue + + # Supported text fields + if key in [ + "negative_prompt", + "aspect_ratio", + "seed", + "output_format", + "model", + "mode", + "strength", + "style_preset", + ]: + data[key] = value # type: ignore + + return data, files + + def transform_image_edit_response( + self, + model: str, + raw_response: httpx.Response, + logging_obj: LiteLLMLoggingObj, + api_key: Optional[str] = None, + json_mode: Optional[bool] = None, + ) -> ImageResponse: + """ + Transform Stability AI response to OpenAI-compatible ImageResponse. + + Stability returns: {"image": "base64...", "finish_reason": "SUCCESS", "seed": 123} + OpenAI expects: {"data": [{"b64_json": "base64..."}], "created": timestamp} + """ + try: + response_data = raw_response.json() + except Exception as e: + raise self.get_error_class( + error_message=f"Error parsing Stability AI response: {e}", + status_code=raw_response.status_code, + headers=raw_response.headers, + ) + + # Check for errors in response + if "errors" in response_data: + raise self.get_error_class( + error_message=f"Stability AI error: {response_data['errors']}", + status_code=raw_response.status_code, + headers=raw_response.headers, + ) + + # Check finish_reason + finish_reason = response_data.get("finish_reason", "") + if finish_reason == "CONTENT_FILTERED": + raise self.get_error_class( + error_message="Content was filtered by Stability AI safety systems", + status_code=400, + headers=raw_response.headers, + ) + + model_response = ImageResponse() + if not model_response.data: + model_response.data = [] + + # Extract image from response + image_b64 = response_data.get("image") + if image_b64: + model_response.data.append( + ImageObject( + b64_json=image_b64, + url=None, + revised_prompt=None, + ) + ) + + if not hasattr(model_response, "_hidden_params"): + model_response._hidden_params = {} + if "additional_headers" not in model_response._hidden_params: + model_response._hidden_params["additional_headers"] = {} + # Override: fetch model-cost from model_cost map based on the provided model name + model_info = get_model_info(model, custom_llm_provider="stability") + cost_per_image = model_info.get("output_cost_per_image", 0) + model_response._hidden_params["additional_headers"]["llm_provider-x-litellm-response-cost"] = float(cost_per_image) + return model_response + + def use_multipart_form_data(self) -> bool: + """ + Stability AI requires multipart/form-data for image generation. + """ + return True diff --git a/litellm/main.py b/litellm/main.py index b46208cc432..570709d8a60 100644 --- a/litellm/main.py +++ b/litellm/main.py @@ -165,7 +165,7 @@ from .llms.azure_ai.anthropic.handler import AzureAnthropicChatCompletion from .llms.azure_ai.embed import AzureAIEmbedding from .llms.bedrock.chat import BedrockConverseLLM, BedrockLLM from .llms.bedrock.embed.embedding import BedrockEmbedding -from .llms.bedrock.image.image_handler import BedrockImageGeneration +from .llms.bedrock.image_generation.image_handler import BedrockImageGeneration from .llms.bytez.chat.transformation import BytezChatConfig from .llms.clarifai.chat.transformation import ClarifaiConfig from .llms.codestral.completion.handler import CodestralTextCompletion diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 0f5f61e708d..19b18747c31 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -24483,6 +24483,90 @@ "output_cost_per_image": 0.08, "supported_endpoints": ["/v1/images/generations"] }, + "stability/inpaint": { + "litellm_provider": "stability", + "mode": "image_edit", + "output_cost_per_image": 0.005, + "supported_endpoints": ["/v1/images/edits"] + }, + "stability/outpaint": { + "litellm_provider": "stability", + "mode": "image_edit", + "output_cost_per_image": 0.004, + "supported_endpoints": ["/v1/images/edits"] + }, + "stability/erase": { + "litellm_provider": "stability", + "mode": "image_edit", + "output_cost_per_image": 0.005, + "supported_endpoints": ["/v1/images/edits"] + }, + "stability/search-and-replace": { + "litellm_provider": "stability", + "mode": "image_edit", + "output_cost_per_image": 0.005, + "supported_endpoints": ["/v1/images/edits"] + }, + "stability/search-and-recolor": { + "litellm_provider": "stability", + "mode": "image_edit", + "output_cost_per_image": 0.005, + "supported_endpoints": ["/v1/images/edits"] + }, + "stability/remove-background": { + "litellm_provider": "stability", + "mode": "image_edit", + "output_cost_per_image": 0.005, + "supported_endpoints": ["/v1/images/edits"] + }, + "stability/replace-background-and-relight": { + "litellm_provider": "stability", + "mode": "image_edit", + "output_cost_per_image": 0.008, + "supported_endpoints": ["/v1/images/edits"] + }, + "stability/sketch": { + "litellm_provider": "stability", + "mode": "image_edit", + "output_cost_per_image": 0.005, + "supported_endpoints": ["/v1/images/edits"] + }, + "stability/structure": { + "litellm_provider": "stability", + "mode": "image_edit", + "output_cost_per_image": 0.005, + "supported_endpoints": ["/v1/images/edits"] + }, + "stability/style": { + "litellm_provider": "stability", + "mode": "image_edit", + "output_cost_per_image": 0.005, + "supported_endpoints": ["/v1/images/edits"] + }, + "stability/style-transfer": { + "litellm_provider": "stability", + "mode": "image_edit", + "output_cost_per_image": 0.008, + "supported_endpoints": ["/v1/images/edits"] + }, + "stability/fast": { + "litellm_provider": "stability", + "mode": "image_edit", + "output_cost_per_image": 0.002, + "supported_endpoints": ["/v1/images/edits"] + }, + "stability/conservative": { + "litellm_provider": "stability", + "mode": "image_edit", + "output_cost_per_image": 0.04, + "supported_endpoints": ["/v1/images/edits"] + }, + "stability/creative": { + "litellm_provider": "stability", + "mode": "image_edit", + "output_cost_per_image": 0.06, + "supported_endpoints": ["/v1/images/edits"] + }, "stability/stable-image-core": { "litellm_provider": "stability", "mode": "image_generation", diff --git a/litellm/types/llms/openai.py b/litellm/types/llms/openai.py index 502bc7e4fe0..ceeae958a80 100644 --- a/litellm/types/llms/openai.py +++ b/litellm/types/llms/openai.py @@ -1029,6 +1029,19 @@ OpenAIImageGenerationOptionalParams = Literal[ "user", ] +OpenAIImageEditOptionalParams = Literal[ + "background", + "n", + "mask" + "output_compression", + "output_format", + "quality", + "partial_images", + "response_format", + "size", + "style", + "user", +] class ComputerToolParam(TypedDict, total=False): display_height: Required[float] diff --git a/litellm/types/llms/stability.py b/litellm/types/llms/stability.py index 33199ff769d..7dd92e380c7 100644 --- a/litellm/types/llms/stability.py +++ b/litellm/types/llms/stability.py @@ -29,6 +29,13 @@ class StabilityImageGenerationRequest(TypedDict, total=False): strength: Optional[float] # How much to transform the image (0-1) style_preset: Optional[str] # Style preset name +class StabilityImageEditRequest(StabilityImageGenerationRequest): + """ + Request parameters for Stability AI image edit endpoint. + + Endpoint: /v2beta/stable-image/edit/inpaint + """ + mask: Optional[str] # Base64-encoded mask (white = edit, black = keep) class StabilityImageGenerationResponse(TypedDict, total=False): """ @@ -197,16 +204,12 @@ STABILITY_EDIT_ENDPOINTS = { "search-and-replace": "/v2beta/stable-image/edit/search-and-replace", "search-and-recolor": "/v2beta/stable-image/edit/search-and-recolor", "remove-background": "/v2beta/stable-image/edit/remove-background", -} - -STABILITY_UPSCALE_ENDPOINTS = { + "replace-background-and-relight": "/v2beta/stable-image/edit/replace-background-and-relight", "fast": "/v2beta/stable-image/upscale/fast", "conservative": "/v2beta/stable-image/upscale/conservative", "creative": "/v2beta/stable-image/upscale/creative", -} - -STABILITY_CONTROL_ENDPOINTS = { "sketch": "/v2beta/stable-image/control/sketch", "structure": "/v2beta/stable-image/control/structure", "style": "/v2beta/stable-image/control/style", + "style-transfer": "/v2beta/stable-image/control/style-transfer", } diff --git a/litellm/utils.py b/litellm/utils.py index 5fb78e141c1..238a6a207b3 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -7966,6 +7966,12 @@ class ProviderConfigManager: ) return get_vertex_ai_image_edit_config(model) + elif LlmProviders.STABILITY == provider: + from litellm.llms.stability.image_edit import ( + get_stability_image_edit_config, + ) + + return get_stability_image_edit_config(model) return None @staticmethod diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 0f5f61e708d..19b18747c31 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -24483,6 +24483,90 @@ "output_cost_per_image": 0.08, "supported_endpoints": ["/v1/images/generations"] }, + "stability/inpaint": { + "litellm_provider": "stability", + "mode": "image_edit", + "output_cost_per_image": 0.005, + "supported_endpoints": ["/v1/images/edits"] + }, + "stability/outpaint": { + "litellm_provider": "stability", + "mode": "image_edit", + "output_cost_per_image": 0.004, + "supported_endpoints": ["/v1/images/edits"] + }, + "stability/erase": { + "litellm_provider": "stability", + "mode": "image_edit", + "output_cost_per_image": 0.005, + "supported_endpoints": ["/v1/images/edits"] + }, + "stability/search-and-replace": { + "litellm_provider": "stability", + "mode": "image_edit", + "output_cost_per_image": 0.005, + "supported_endpoints": ["/v1/images/edits"] + }, + "stability/search-and-recolor": { + "litellm_provider": "stability", + "mode": "image_edit", + "output_cost_per_image": 0.005, + "supported_endpoints": ["/v1/images/edits"] + }, + "stability/remove-background": { + "litellm_provider": "stability", + "mode": "image_edit", + "output_cost_per_image": 0.005, + "supported_endpoints": ["/v1/images/edits"] + }, + "stability/replace-background-and-relight": { + "litellm_provider": "stability", + "mode": "image_edit", + "output_cost_per_image": 0.008, + "supported_endpoints": ["/v1/images/edits"] + }, + "stability/sketch": { + "litellm_provider": "stability", + "mode": "image_edit", + "output_cost_per_image": 0.005, + "supported_endpoints": ["/v1/images/edits"] + }, + "stability/structure": { + "litellm_provider": "stability", + "mode": "image_edit", + "output_cost_per_image": 0.005, + "supported_endpoints": ["/v1/images/edits"] + }, + "stability/style": { + "litellm_provider": "stability", + "mode": "image_edit", + "output_cost_per_image": 0.005, + "supported_endpoints": ["/v1/images/edits"] + }, + "stability/style-transfer": { + "litellm_provider": "stability", + "mode": "image_edit", + "output_cost_per_image": 0.008, + "supported_endpoints": ["/v1/images/edits"] + }, + "stability/fast": { + "litellm_provider": "stability", + "mode": "image_edit", + "output_cost_per_image": 0.002, + "supported_endpoints": ["/v1/images/edits"] + }, + "stability/conservative": { + "litellm_provider": "stability", + "mode": "image_edit", + "output_cost_per_image": 0.04, + "supported_endpoints": ["/v1/images/edits"] + }, + "stability/creative": { + "litellm_provider": "stability", + "mode": "image_edit", + "output_cost_per_image": 0.06, + "supported_endpoints": ["/v1/images/edits"] + }, "stability/stable-image-core": { "litellm_provider": "stability", "mode": "image_generation",