diff --git a/litellm/images/main.py b/litellm/images/main.py index b711aa31c05..49d2026100e 100644 --- a/litellm/images/main.py +++ b/litellm/images/main.py @@ -33,6 +33,7 @@ from litellm.main import ( base_llm_aiohttp_handler, base_llm_http_handler, bedrock_image_generation, + bedrock_image_edit, openai_chat_completions, openai_image_variations, ) @@ -812,6 +813,25 @@ def image_edit( custom_llm_provider=custom_llm_provider, ) + # Route bedrock to its specific handler (AWS signing required) + if custom_llm_provider == "bedrock": + if model is None: + raise Exception("Model needs to be set for bedrock") + return bedrock_image_edit.image_edit( # type: ignore + model=model, + image=images, + prompt=prompt, + timeout=timeout, + logging_obj=litellm_logging_obj, + optional_params=image_edit_request_params, + model_response=ImageResponse(), + aimage_edit=_is_async, + client=kwargs.get("client"), + api_base=kwargs.get("api_base"), + extra_headers=extra_headers, + api_key=kwargs.get("api_key"), + ) + # Call the handler with _is_async flag instead of directly calling the async handler return base_llm_http_handler.image_edit_handler( model=model, diff --git a/litellm/litellm_core_utils/llm_cost_calc/utils.py b/litellm/litellm_core_utils/llm_cost_calc/utils.py index ef2183a4556..232d9bfc5d1 100644 --- a/litellm/litellm_core_utils/llm_cost_calc/utils.py +++ b/litellm/litellm_core_utils/llm_cost_calc/utils.py @@ -674,7 +674,7 @@ class CostCalculatorUtils: from litellm.llms.azure_ai.image_generation.cost_calculator import ( cost_calculator as azure_ai_image_cost_calculator, ) - from litellm.llms.bedrock.image.cost_calculator import ( + from litellm.llms.bedrock.image_generation.cost_calculator import ( cost_calculator as bedrock_image_cost_calculator, ) from litellm.llms.gemini.image_generation.cost_calculator import ( diff --git a/litellm/llms/bedrock/image_edit/__init__.py b/litellm/llms/bedrock/image_edit/__init__.py new file mode 100644 index 00000000000..f3a0e61067d --- /dev/null +++ b/litellm/llms/bedrock/image_edit/__init__.py @@ -0,0 +1,10 @@ +""" +Bedrock Image Edit Module + +Handles image edit operations for Bedrock stability models. +""" + +from .handler import BedrockImageEdit + +__all__ = ["BedrockImageEdit"] + diff --git a/litellm/llms/bedrock/image_edit/handler.py b/litellm/llms/bedrock/image_edit/handler.py new file mode 100644 index 00000000000..dfc3b8fb449 --- /dev/null +++ b/litellm/llms/bedrock/image_edit/handler.py @@ -0,0 +1,310 @@ +""" +Bedrock Image Edit Handler + +Handles image edit requests for Bedrock stability models. +""" + +from __future__ import annotations + +import json +from typing import TYPE_CHECKING, Any, Optional, Union + +import httpx +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_edit.stability_transformation import ( + BedrockStabilityImageEditConfig, +) +from litellm.llms.custom_httpx.http_handler import ( + AsyncHTTPHandler, + HTTPHandler, + _get_httpx_client, + get_async_httpx_client, +) +from litellm.types.utils import ImageResponse + +from ..base_aws_llm import BaseAWSLLM +from ..common_utils import BedrockError + +if TYPE_CHECKING: + from botocore.awsrequest import AWSPreparedRequest +else: + AWSPreparedRequest = Any + + +class BedrockImageEditPreparedRequest(BaseModel): + """ + Internal/Helper class for preparing the request for bedrock image edit + """ + + endpoint_url: str + prepped: AWSPreparedRequest + body: bytes + data: dict + + +class BedrockImageEdit(BaseAWSLLM): + """ + Bedrock Image Edit handler + """ + + @classmethod + def get_config_class(cls, model: str | None): + if BedrockStabilityImageEditConfig._is_stability_edit_model(model): + return BedrockStabilityImageEditConfig + else: + raise ValueError(f"Unsupported model for bedrock image edit: {model}") + + def image_edit( + self, + model: str, + image: list, + prompt: str, + model_response: ImageResponse, + optional_params: dict, + logging_obj: LitellmLogging, + timeout: Optional[Union[float, httpx.Timeout]], + aimage_edit: bool = False, + api_base: Optional[str] = None, + extra_headers: Optional[dict] = None, + client: Optional[Union[HTTPHandler, AsyncHTTPHandler]] = None, + api_key: Optional[str] = None, + ): + prepared_request = self._prepare_request( + model=model, + image=image, + prompt=prompt, + optional_params=optional_params, + api_base=api_base, + extra_headers=extra_headers, + logging_obj=logging_obj, + api_key=api_key, + ) + + if aimage_edit is True: + return self.async_image_edit( + prepared_request=prepared_request, + timeout=timeout, + model=model, + logging_obj=logging_obj, + prompt=prompt, + model_response=model_response, + client=( + client + if client is not None and isinstance(client, AsyncHTTPHandler) + else None + ), + ) + + if client is None or not isinstance(client, HTTPHandler): + client = _get_httpx_client() + try: + response = client.post(url=prepared_request.endpoint_url, headers=prepared_request.prepped.headers, data=prepared_request.body) # type: ignore + response.raise_for_status() + except httpx.HTTPStatusError as err: + error_code = err.response.status_code + raise BedrockError(status_code=error_code, message=err.response.text) + except httpx.TimeoutException: + raise BedrockError(status_code=408, message="Timeout error occurred.") + + ### FORMAT RESPONSE TO OPENAI FORMAT ### + model_response = self._transform_response_dict_to_openai_response( + model_response=model_response, + model=model, + logging_obj=logging_obj, + prompt=prompt, + response=response, + data=prepared_request.data, + ) + return model_response + + async def async_image_edit( + self, + prepared_request: BedrockImageEditPreparedRequest, + timeout: Optional[Union[float, httpx.Timeout]], + model: str, + logging_obj: LitellmLogging, + prompt: str, + model_response: ImageResponse, + client: Optional[AsyncHTTPHandler] = None, + ) -> ImageResponse: + """ + Asynchronous handler for bedrock image edit + """ + async_client = client or get_async_httpx_client( + llm_provider=litellm.LlmProviders.BEDROCK, + params={"timeout": timeout}, + ) + + try: + response = await async_client.post(url=prepared_request.endpoint_url, headers=prepared_request.prepped.headers, data=prepared_request.body) # type: ignore + response.raise_for_status() + except httpx.HTTPStatusError as err: + error_code = err.response.status_code + raise BedrockError(status_code=error_code, message=err.response.text) + except httpx.TimeoutException: + raise BedrockError(status_code=408, message="Timeout error occurred.") + + ### FORMAT RESPONSE TO OPENAI FORMAT ### + model_response = self._transform_response_dict_to_openai_response( + model=model, + logging_obj=logging_obj, + prompt=prompt, + response=response, + data=prepared_request.data, + model_response=model_response, + ) + return model_response + + def _prepare_request( + self, + model: str, + image: list, + prompt: str, + optional_params: dict, + api_base: Optional[str], + extra_headers: Optional[dict], + logging_obj: LitellmLogging, + api_key: Optional[str], + ) -> BedrockImageEditPreparedRequest: + """ + Prepare the request body, headers, and endpoint URL for the Bedrock Image Edit API + + Args: + model (str): The model to use for the image edit + image (list): The images to edit + prompt (str): The prompt for the edit + optional_params (dict): The optional parameters for the image edit + api_base (Optional[str]): The base URL for the Bedrock API + extra_headers (Optional[dict]): The extra headers to include in the request + logging_obj (LitellmLogging): The logging object to use for logging + api_key (Optional[str]): The API key to use + + Returns: + BedrockImageEditPreparedRequest: The prepared request object + """ + boto3_credentials_info = self._get_boto_credentials_from_optional_params( + optional_params, model + ) + + # Use the existing ARN-aware provider detection method + bedrock_provider = self.get_bedrock_invoke_provider(model) + ### SET RUNTIME ENDPOINT ### + modelId = self.get_bedrock_model_id( + model=model, + provider=bedrock_provider, + optional_params=optional_params, + ) + _, proxy_endpoint_url = self.get_runtime_endpoint( + api_base=api_base, + aws_bedrock_runtime_endpoint=boto3_credentials_info.aws_bedrock_runtime_endpoint, + aws_region_name=boto3_credentials_info.aws_region_name, + ) + proxy_endpoint_url = f"{proxy_endpoint_url}/model/{modelId}/invoke" + data = self._get_request_body( + model=model, + image=image, + prompt=prompt, + optional_params=optional_params, + ) + + # Make POST Request + body = json.dumps(data).encode("utf-8") + headers = {"Content-Type": "application/json"} + if extra_headers is not None: + headers = {"Content-Type": "application/json", **extra_headers} + + prepped = self.get_request_headers( + credentials=boto3_credentials_info.credentials, + aws_region_name=boto3_credentials_info.aws_region_name, + extra_headers=extra_headers, + endpoint_url=proxy_endpoint_url, + data=body, + headers=headers, + api_key=api_key, + ) + + ## LOGGING + logging_obj.pre_call( + input=prompt, + api_key="", + additional_args={ + "complete_input_dict": data, + "api_base": proxy_endpoint_url, + "headers": prepped.headers, + }, + ) + return BedrockImageEditPreparedRequest( + endpoint_url=proxy_endpoint_url, + prepped=prepped, + body=body, + data=data, + ) + + def _get_request_body( + self, + model: str, + image: list, + prompt: str, + optional_params: dict, + ) -> dict: + """ + Get the request body for the Bedrock Image Edit API + + Checks the model/provider and transforms the request body accordingly + + Returns: + dict: The request body to use for the Bedrock Image Edit API + """ + config_class = self.get_config_class(model=model) + config_instance = config_class() + request_body = config_instance.transform_image_edit_request( + model=model, + prompt=prompt, + image=image[0] if image else None, + image_edit_optional_request_params=optional_params, + litellm_params={}, + headers={}, + ) + return dict(request_body) + + def _transform_response_dict_to_openai_response( + self, + model_response: ImageResponse, + model: str, + logging_obj: LitellmLogging, + prompt: str, + response: httpx.Response, + data: dict, + ) -> ImageResponse: + """ + Transforms the Image Edit response from Bedrock to OpenAI format + """ + + ## LOGGING + if logging_obj is not None: + logging_obj.post_call( + input=prompt, + api_key="", + original_response=response.text, + additional_args={"complete_input_dict": data}, + ) + verbose_logger.debug("raw model_response: %s", response.text) + response_dict = response.json() + if response_dict is None: + raise ValueError("Error in response object format, got None") + + config_class = self.get_config_class(model=model) + config_instance = config_class() + + config_instance.transform_image_edit_response( + model_response=model_response, + raw_response=response, + logging_obj=logging_obj, + ) + + return model_response + diff --git a/litellm/llms/bedrock/image_edit/stability_transformation.py b/litellm/llms/bedrock/image_edit/stability_transformation.py new file mode 100644 index 00000000000..0484bbab613 --- /dev/null +++ b/litellm/llms/bedrock/image_edit/stability_transformation.py @@ -0,0 +1,404 @@ +""" +Bedrock Stability AI Image Edit Transformation + +Handles transformation between OpenAI-compatible format and Bedrock Stability AI Image Edit API format. + +Supported models: +- stability.stable-conservative-upscale-v1:0 +- stability.stable-creative-upscale-v1:0 +- stability.stable-fast-upscale-v1:0 +- stability.stable-outpaint-v1:0 +- stability.stable-image-control-sketch-v1:0 +- stability.stable-image-control-structure-v1:0 +- stability.stable-image-erase-object-v1:0 +- stability.stable-image-inpaint-v1:0 +- stability.stable-image-remove-background-v1:0 +- stability.stable-image-search-recolor-v1:0 +- stability.stable-image-search-replace-v1:0 +- stability.stable-image-style-guide-v1:0 +- stability.stable-style-transfer-v1:0 + +API Reference: https://docs.aws.amazon.com/bedrock/latest/userguide/model-parameters.html +""" + +import json +import base64 +from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple + +import httpx + +from litellm.llms.base_llm.image_edit.transformation import BaseImageEditConfig +from litellm.types.llms.openai import ( + OpenAIImageEditOptionalParams, +) +from litellm.types.llms.stability import ( + OPENAI_SIZE_TO_STABILITY_ASPECT_RATIO, + STABILITY_EDIT_ENDPOINTS, +) +from litellm.types.utils import FileTypes, ImageObject, ImageResponse +from litellm.utils import get_model_info +from litellm.llms.bedrock.common_utils import BedrockError + +if TYPE_CHECKING: + from litellm.litellm_core_utils.litellm_logging import Logging as _LiteLLMLoggingObj + + LiteLLMLoggingObj = _LiteLLMLoggingObj +else: + LiteLLMLoggingObj = Any + + +class BedrockStabilityImageEditConfig(BaseImageEditConfig): + """ + Configuration for Bedrock Stability AI image edit. + + Supports all Stability image edit operations through Bedrock. + """ + + @classmethod + def _is_stability_edit_model(cls, model: Optional[str] = None) -> bool: + """ + Returns True if the model is a Bedrock Stability edit model. + + Bedrock Stability edit models follow this pattern: + stability.stable-conservative-upscale-v1:0 + stability.stable-creative-upscale-v1:0 + stability.stable-fast-upscale-v1:0 + stability.stable-outpaint-v1:0 + stability.stable-image-inpaint-v1:0 + stability.stable-image-erase-object-v1:0 + etc. + """ + if model: + model_lower = model.lower() + if "stability." in model_lower and any([ + "upscale" in model_lower, + "outpaint" in model_lower, + "inpaint" in model_lower, + "erase" in model_lower, + "remove-background" in model_lower, + "search-recolor" in model_lower, + "search-replace" in model_lower, + "control-sketch" in model_lower, + "control-structure" in model_lower, + "style-guide" in model_lower, + "style-transfer" in model_lower, + ]): + return True + return False + + def get_supported_openai_params( + self, model: str + ) -> List[OpenAIImageEditOptionalParams]: + """ + Return list of OpenAI params supported by Bedrock Stability. + """ + 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 Bedrock Stability 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_operation(self, model: str) -> str: + """ + Get the operation type for a given model. + """ + model_lower = model.lower() + + if "upscale" in model_lower: + if "fast" in model_lower: + return "fast" + elif "conservative" in model_lower: + return "conservative" + elif "creative" in model_lower: + return "creative" + elif "outpaint" in model_lower: + return "outpaint" + elif "inpaint" in model_lower: + return "inpaint" + elif "erase" in model_lower: + return "erase" + elif "remove-background" in model_lower: + return "remove-background" + elif "search-recolor" in model_lower: + return "search-and-recolor" + elif "search-replace" in model_lower: + return "search-and-replace" + elif "control-sketch" in model_lower: + return "sketch" + elif "control-structure" in model_lower: + return "structure" + elif "style-guide" in model_lower: + return "style" + elif "style-transfer" in model_lower: + return "style-transfer" + + # Default to inpaint + return "inpaint" + + def transform_image_edit_request( + self, + model: str, + prompt: str, + image: FileTypes, + image_edit_optional_request_params: OpenAIImageEditOptionalParams, + litellm_params: dict, + headers: dict, + ) -> Dict[str, Any]: + """ + Transform OpenAI-style request to Bedrock Stability request format. + + Returns the request body dict that will be JSON-encoded by the handler. + """ + # Build Bedrock Stability request + data: Dict[str, Any] = { + "prompt": prompt, + "output_format": "png", # Default to PNG + } + + # Convert image to base64 + image_b64: str + if hasattr(image, 'read'): + # File-like object (e.g., BufferedReader from open()) + image_bytes = image.read() + image_b64 = base64.b64encode(image_bytes).decode('utf-8') + elif isinstance(image, bytes): + # Raw bytes + image_b64 = base64.b64encode(image).decode('utf-8') + elif isinstance(image, str): + # Already a base64 string + image_b64 = image + else: + # Try to handle as bytes + image_b64 = base64.b64encode(image).decode('utf-8') + + data["image"] = image_b64 + + # 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": + if hasattr(value, 'read'): + mask_bytes = value.read() + else: + mask_bytes = value + if isinstance(mask_bytes, bytes): + mask_b64 = base64.b64encode(mask_bytes).decode('utf-8') + else: + mask_b64 = mask_bytes + data["mask"] = mask_b64 + continue + + # Supported text fields + if key in [ + "negative_prompt", + "aspect_ratio", + "seed", + "output_format", + "model", + "mode", + "strength", + "style_preset", + "creativity", + "control_strength", + "grow_mask", + "left", + "right", + "up", + "down", + "select_prompt", + "search_prompt", + "fidelity", + "composition_fidelity", + "style_strength", + "change_strength", + "init_image", + "style_image", + ]: + data[key] = value # type: ignore + + return data + + 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 Bedrock Stability response to OpenAI-compatible ImageResponse. + + Bedrock returns: {"images": ["base64..."], "finish_reasons": [null], "seeds": [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 Bedrock Stability 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"Bedrock Stability error: {response_data['errors']}", + status_code=raw_response.status_code, + headers=raw_response.headers, + ) + + # Check finish_reasons + finish_reasons = response_data.get("finish_reasons", []) + if finish_reasons and finish_reasons[0]: + raise self.get_error_class( + error_message=f"Bedrock Stability error: {finish_reasons[0]}", + status_code=400, + headers=raw_response.headers, + ) + + model_response = ImageResponse() + if not model_response.data: + model_response.data = [] + + # Extract images from response + images = response_data.get("images", []) + if images: + for image_b64 in images: + 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"] = {} + + # Set cost based on model + model_info = get_model_info(model, custom_llm_provider="bedrock") + 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: + """ + Bedrock Stability uses JSON format, not multipart/form-data. + """ + return False + + def get_complete_url( + self, + api_base: Optional[str], + model: str, + litellm_params: dict, + ) -> str: + """ + Get the complete URL for the Bedrock Image Edit API. + + For Bedrock, this is handled by the handler which constructs the endpoint URL + based on the model ID and AWS region. This method is required by the base class + but the actual URL construction happens in BedrockImageEdit.image_edit(). + + Returns a placeholder - the real endpoint is constructed in the handler. + """ + # Bedrock URLs are constructed in the handler using boto3 + # This is a placeholder for the abstract method requirement + return "bedrock://image-edit" + + def validate_environment( + self, + headers: dict, + model: str, + api_key: Optional[str] = None, + ) -> dict: + """ + Validate environment for Bedrock Stability image edit. + + For Bedrock, AWS credentials are managed by the BaseAWSLLM class. + This method validates that headers are properly set up. + + Args: + headers: The request headers to validate/update + model: The model name being used + api_key: Optional API key (not used for Bedrock, which uses AWS credentials) + + Returns: + Updated headers dict + """ + if headers is None: + headers = {} + + # Bedrock uses AWS credentials, not API keys + # Headers are set up by the handler's get_request_headers() method + # This just ensures basic headers are present + if "Content-Type" not in headers: + headers["Content-Type"] = "application/json" + + return headers + diff --git a/litellm/llms/bedrock/image/amazon_nova_canvas_transformation.py b/litellm/llms/bedrock/image_generation/amazon_nova_canvas_transformation.py similarity index 100% rename from litellm/llms/bedrock/image/amazon_nova_canvas_transformation.py rename to litellm/llms/bedrock/image_generation/amazon_nova_canvas_transformation.py diff --git a/litellm/llms/bedrock/image/amazon_stability1_transformation.py b/litellm/llms/bedrock/image_generation/amazon_stability1_transformation.py similarity index 100% rename from litellm/llms/bedrock/image/amazon_stability1_transformation.py rename to litellm/llms/bedrock/image_generation/amazon_stability1_transformation.py diff --git a/litellm/llms/bedrock/image/amazon_stability3_transformation.py b/litellm/llms/bedrock/image_generation/amazon_stability3_transformation.py similarity index 100% rename from litellm/llms/bedrock/image/amazon_stability3_transformation.py rename to litellm/llms/bedrock/image_generation/amazon_stability3_transformation.py diff --git a/litellm/llms/bedrock/image/amazon_titan_transformation.py b/litellm/llms/bedrock/image_generation/amazon_titan_transformation.py similarity index 100% rename from litellm/llms/bedrock/image/amazon_titan_transformation.py rename to litellm/llms/bedrock/image_generation/amazon_titan_transformation.py diff --git a/litellm/llms/bedrock/image/cost_calculator.py b/litellm/llms/bedrock/image_generation/cost_calculator.py similarity index 100% rename from litellm/llms/bedrock/image/cost_calculator.py rename to litellm/llms/bedrock/image_generation/cost_calculator.py diff --git a/litellm/llms/bedrock/image/image_handler.py b/litellm/llms/bedrock/image_generation/image_handler.py similarity index 100% rename from litellm/llms/bedrock/image/image_handler.py rename to litellm/llms/bedrock/image_generation/image_handler.py diff --git a/litellm/main.py b/litellm/main.py index 570709d8a60..0715dd8e61b 100644 --- a/litellm/main.py +++ b/litellm/main.py @@ -166,6 +166,7 @@ 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_generation.image_handler import BedrockImageGeneration +from .llms.bedrock.image_edit.handler import BedrockImageEdit from .llms.bytez.chat.transformation import BytezChatConfig from .llms.clarifai.chat.transformation import ClarifaiConfig from .llms.codestral.completion.handler import CodestralTextCompletion @@ -271,6 +272,7 @@ codestral_text_completions = CodestralTextCompletion() bedrock_converse_chat_completion = BedrockConverseLLM() bedrock_embedding = BedrockEmbedding() bedrock_image_generation = BedrockImageGeneration() +bedrock_image_edit = BedrockImageEdit() vertex_chat_completion = VertexLLM() vertex_embedding = VertexEmbedding() vertex_multimodal_embedding = VertexMultimodalEmbedding() diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 19b18747c31..2312876f068 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -24615,6 +24615,84 @@ "mode": "image_generation", "output_cost_per_image": 0.14 }, + "stability.stable-conservative-upscale-v1:0": { + "litellm_provider": "bedrock", + "max_input_tokens": 77, + "mode": "image_edit", + "output_cost_per_image": 0.40 + }, + "stability.stable-creative-upscale-v1:0": { + "litellm_provider": "bedrock", + "max_input_tokens": 77, + "mode": "image_edit", + "output_cost_per_image": 0.60 + }, + "stability.stable-fast-upscale-v1:0": { + "litellm_provider": "bedrock", + "max_input_tokens": 77, + "mode": "image_edit", + "output_cost_per_image": 0.03 + }, + "stability.stable-outpaint-v1:0": { + "litellm_provider": "bedrock", + "max_input_tokens": 77, + "mode": "image_edit", + "output_cost_per_image": 0.06 + }, + "stability.stable-image-control-sketch-v1:0": { + "litellm_provider": "bedrock", + "max_input_tokens": 77, + "mode": "image_edit", + "output_cost_per_image": 0.07 + }, + "stability.stable-image-control-structure-v1:0": { + "litellm_provider": "bedrock", + "max_input_tokens": 77, + "mode": "image_edit", + "output_cost_per_image": 0.07 + }, + "stability.stable-image-erase-object-v1:0": { + "litellm_provider": "bedrock", + "max_input_tokens": 77, + "mode": "image_edit", + "output_cost_per_image": 0.07 + }, + "stability.stable-image-inpaint-v1:0": { + "litellm_provider": "bedrock", + "max_input_tokens": 77, + "mode": "image_edit", + "output_cost_per_image": 0.07 + }, + "stability.stable-image-remove-background-v1:0": { + "litellm_provider": "bedrock", + "max_input_tokens": 77, + "mode": "image_edit", + "output_cost_per_image": 0.07 + }, + "stability.stable-image-search-recolor-v1:0": { + "litellm_provider": "bedrock", + "max_input_tokens": 77, + "mode": "image_edit", + "output_cost_per_image": 0.07 + }, + "stability.stable-image-search-replace-v1:0": { + "litellm_provider": "bedrock", + "max_input_tokens": 77, + "mode": "image_edit", + "output_cost_per_image": 0.07 + }, + "stability.stable-image-style-guide-v1:0": { + "litellm_provider": "bedrock", + "max_input_tokens": 77, + "mode": "image_edit", + "output_cost_per_image": 0.07 + }, + "stability.stable-style-transfer-v1:0": { + "litellm_provider": "bedrock", + "max_input_tokens": 77, + "mode": "image_edit", + "output_cost_per_image": 0.08 + }, "standard/1024-x-1024/dall-e-3": { "input_cost_per_pixel": 3.81469e-08, "litellm_provider": "openai", diff --git a/litellm/utils.py b/litellm/utils.py index 238a6a207b3..581f91d6201 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -7972,6 +7972,12 @@ class ProviderConfigManager: ) return get_stability_image_edit_config(model) + elif LlmProviders.BEDROCK == provider: + from litellm.llms.bedrock.image_edit.stability_transformation import ( + BedrockStabilityImageEditConfig, + ) + + return BedrockStabilityImageEditConfig() return None @staticmethod diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 19b18747c31..fd9c5d0c618 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -24594,6 +24594,84 @@ "mode": "image_generation", "output_cost_per_image": 0.04 }, + "stability.stable-conservative-upscale-v1:0": { + "litellm_provider": "bedrock", + "max_input_tokens": 77, + "mode": "image_edit", + "output_cost_per_image": 0.40 + }, + "stability.stable-creative-upscale-v1:0": { + "litellm_provider": "bedrock", + "max_input_tokens": 77, + "mode": "image_edit", + "output_cost_per_image": 0.60 + }, + "stability.stable-fast-upscale-v1:0": { + "litellm_provider": "bedrock", + "max_input_tokens": 77, + "mode": "image_edit", + "output_cost_per_image": 0.03 + }, + "stability.stable-outpaint-v1:0": { + "litellm_provider": "bedrock", + "max_input_tokens": 77, + "mode": "image_edit", + "output_cost_per_image": 0.06 + }, + "stability.stable-image-control-sketch-v1:0": { + "litellm_provider": "bedrock", + "max_input_tokens": 77, + "mode": "image_edit", + "output_cost_per_image": 0.07 + }, + "stability.stable-image-control-structure-v1:0": { + "litellm_provider": "bedrock", + "max_input_tokens": 77, + "mode": "image_edit", + "output_cost_per_image": 0.07 + }, + "stability.stable-image-erase-object-v1:0": { + "litellm_provider": "bedrock", + "max_input_tokens": 77, + "mode": "image_edit", + "output_cost_per_image": 0.07 + }, + "stability.stable-image-inpaint-v1:0": { + "litellm_provider": "bedrock", + "max_input_tokens": 77, + "mode": "image_edit", + "output_cost_per_image": 0.07 + }, + "stability.stable-image-remove-background-v1:0": { + "litellm_provider": "bedrock", + "max_input_tokens": 77, + "mode": "image_edit", + "output_cost_per_image": 0.07 + }, + "stability.stable-image-search-recolor-v1:0": { + "litellm_provider": "bedrock", + "max_input_tokens": 77, + "mode": "image_edit", + "output_cost_per_image": 0.07 + }, + "stability.stable-image-search-replace-v1:0": { + "litellm_provider": "bedrock", + "max_input_tokens": 77, + "mode": "image_edit", + "output_cost_per_image": 0.07 + }, + "stability.stable-image-style-guide-v1:0": { + "litellm_provider": "bedrock", + "max_input_tokens": 77, + "mode": "image_edit", + "output_cost_per_image": 0.07 + }, + "stability.stable-style-transfer-v1:0": { + "litellm_provider": "bedrock", + "max_input_tokens": 77, + "mode": "image_edit", + "output_cost_per_image": 0.08 + }, "stability.stable-image-core-v1:1": { "litellm_provider": "bedrock", "max_input_tokens": 77, diff --git a/tests/image_gen_tests/test_bedrock_image_gen_unit_tests.py b/tests/image_gen_tests/test_bedrock_image_gen_unit_tests.py index 73331547772..459f3b01aa4 100644 --- a/tests/image_gen_tests/test_bedrock_image_gen_unit_tests.py +++ b/tests/image_gen_tests/test_bedrock_image_gen_unit_tests.py @@ -10,7 +10,7 @@ sys.path.insert( 0, os.path.abspath("../..") ) # Adds the parent directory to the system path -from litellm.llms.bedrock.image.amazon_nova_canvas_transformation import ( +from litellm.llms.bedrock.image_generation.amazon_nova_canvas_transformation import ( AmazonNovaCanvasConfig, ) @@ -22,15 +22,15 @@ sys.path.insert( 0, os.path.abspath("../..") ) # Adds the parent directory to the system path import pytest -from litellm.llms.bedrock.image.cost_calculator import cost_calculator +from litellm.llms.bedrock.image_generation.cost_calculator import cost_calculator from litellm.types.utils import ImageResponse, ImageObject import os import litellm -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_stability1_transformation import ( +from litellm.llms.bedrock.image_generation.amazon_stability1_transformation import ( AmazonStabilityConfig, ) from litellm.types.llms.bedrock import ( @@ -38,7 +38,7 @@ from litellm.types.llms.bedrock import ( AmazonStability3TextToImageResponse, ) from unittest.mock import MagicMock, patch -from litellm.llms.bedrock.image.image_handler import ( +from litellm.llms.bedrock.image_generation.image_handler import ( BedrockImageGeneration, BedrockImagePreparedRequest, ) diff --git a/tests/test_litellm/llms/bedrock/image/test_amazon_nova_canvas_transformation.py b/tests/test_litellm/llms/bedrock/image/test_amazon_nova_canvas_transformation.py index 0dd0b80f36f..122d3e44364 100644 --- a/tests/test_litellm/llms/bedrock/image/test_amazon_nova_canvas_transformation.py +++ b/tests/test_litellm/llms/bedrock/image/test_amazon_nova_canvas_transformation.py @@ -1,5 +1,5 @@ import pytest -from litellm.llms.bedrock.image.amazon_nova_canvas_transformation import AmazonNovaCanvasConfig +from litellm.llms.bedrock.image_generation.amazon_nova_canvas_transformation import AmazonNovaCanvasConfig from litellm.types.utils import ImageResponse def test_transform_request_body_text_to_image(): diff --git a/tests/test_litellm/llms/bedrock/image/test_amazon_stability3_transformation.py b/tests/test_litellm/llms/bedrock/image/test_amazon_stability3_transformation.py index 1cf1747b8c7..a758202d74f 100644 --- a/tests/test_litellm/llms/bedrock/image/test_amazon_stability3_transformation.py +++ b/tests/test_litellm/llms/bedrock/image/test_amazon_stability3_transformation.py @@ -10,7 +10,7 @@ sys.path.insert( ) # Adds the parent directory to the system path from unittest.mock import MagicMock, patch -from litellm.llms.bedrock.image.amazon_stability3_transformation import ( +from litellm.llms.bedrock.image_generation.amazon_stability3_transformation import ( AmazonStability3Config, ) diff --git a/tests/test_litellm/llms/bedrock/image/test_bedrock_image_prepare_request.py b/tests/test_litellm/llms/bedrock/image/test_bedrock_image_prepare_request.py index 22dc0cc8a48..6c56ccc1ef7 100644 --- a/tests/test_litellm/llms/bedrock/image/test_bedrock_image_prepare_request.py +++ b/tests/test_litellm/llms/bedrock/image/test_bedrock_image_prepare_request.py @@ -1,6 +1,6 @@ from unittest.mock import patch, MagicMock -from litellm.llms.bedrock.image.image_handler import BedrockImageGeneration +from litellm.llms.bedrock.image_generation.image_handler import BedrockImageGeneration def test_bedrock_image_prepare_request_with_arn() -> None: dummy_arn = "arn:aws:bedrock:us-east-1:123456789012:application-inference-profile/abcdefghi123"