From 046b7efbbeb1928c3f94d64eb85c0b7cd7cc5b13 Mon Sep 17 00:00:00 2001 From: Dmitrii Komarov Date: Tue, 25 Nov 2025 05:58:01 +0100 Subject: [PATCH] Make Bedrock image generation more consistent (#17021) --- .../amazon_nova_canvas_transformation.py | 20 +++++ .../image/amazon_stability1_transformation.py | 59 +++++++++++++ .../image/amazon_stability3_transformation.py | 31 ++++++- .../image/amazon_titan_transformation.py | 6 +- litellm/llms/bedrock/image/cost_calculator.py | 41 ++------- litellm/llms/bedrock/image/image_handler.py | 85 +++++-------------- litellm/utils.py | 11 +-- .../test_bedrock_image_gen_unit_tests.py | 35 +++++--- 8 files changed, 164 insertions(+), 124 deletions(-) diff --git a/litellm/llms/bedrock/image/amazon_nova_canvas_transformation.py b/litellm/llms/bedrock/image/amazon_nova_canvas_transformation.py index cd33e62af1..f2b94b617c 100644 --- a/litellm/llms/bedrock/image/amazon_nova_canvas_transformation.py +++ b/litellm/llms/bedrock/image/amazon_nova_canvas_transformation.py @@ -3,6 +3,7 @@ from typing import Any, Dict, List, Optional from openai.types.image import Image +from litellm import get_model_info from litellm.types.llms.bedrock import ( AmazonNovaCanvasColorGuidedGenerationParams, AmazonNovaCanvasColorGuidedRequest, @@ -197,3 +198,22 @@ class AmazonNovaCanvasConfig: model_response.data = openai_images return model_response + + @classmethod + def cost_calculator( + cls, + model: str, + image_response: ImageResponse, + size: Optional[str] = None, + optional_params: Optional[dict] = None, + ) -> float: + model_info = get_model_info( + model=model, + custom_llm_provider="bedrock", + ) + + output_cost_per_image: float = model_info.get("output_cost_per_image") or 0.0 + num_images: int = 0 + if image_response.data: + num_images = len(image_response.data) + return output_cost_per_image * num_images \ No newline at end of file diff --git a/litellm/llms/bedrock/image/amazon_stability1_transformation.py b/litellm/llms/bedrock/image/amazon_stability1_transformation.py index 698ecca94b..63af32f3f5 100644 --- a/litellm/llms/bedrock/image/amazon_stability1_transformation.py +++ b/litellm/llms/bedrock/image/amazon_stability1_transformation.py @@ -1,8 +1,11 @@ +import copy +import os import types from typing import List, Optional from openai.types.image import Image +from litellm import get_model_info from litellm.types.utils import ImageResponse @@ -90,6 +93,31 @@ class AmazonStabilityConfig: return optional_params + @classmethod + def transform_request_body( + cls, + text: str, + optional_params: dict, + ) -> dict: + inference_params = copy.deepcopy(optional_params) + inference_params.pop( + "user", None + ) # make sure user is not passed in for bedrock call + + prompt = text.replace(os.linesep, " ") + ## LOAD CONFIG + config = cls.get_config() + for k, v in config.items(): + if ( + k not in inference_params + ): # completion(top_k=3) > anthropic_config(top_k=3) <- allows for dynamic variables to be passed in + inference_params[k] = v + + return { + "text_prompts": [{"text": prompt, "weight": 1}], + **inference_params, + } + @classmethod def transform_response_dict_to_openai_response( cls, model_response: ImageResponse, response_dict: dict @@ -102,3 +130,34 @@ class AmazonStabilityConfig: model_response.data = image_list return model_response + + @classmethod + def cost_calculator( + cls, + model: str, + image_response: ImageResponse, + size: Optional[str] = None, + optional_params: Optional[dict] = None, + ) -> float: + optional_params = optional_params or {} + + # see model_prices_and_context_window.json for details on how steps is used + # Reference pricing by steps for stability 1: https://aws.amazon.com/bedrock/pricing/ + _steps = optional_params.get("steps", 50) + steps = "max-steps" if _steps > 50 else "50-steps" + + # size is stored in model_prices_and_context_window.json as 1024-x-1024 + # current size has 1024x1024 + size = size or "1024-x-1024" + model = f"{size}/{steps}/{model}" + + model_info = get_model_info( + model=model, + custom_llm_provider="bedrock", + ) + + output_cost_per_image: float = model_info.get("output_cost_per_image") or 0.0 + num_images: int = 0 + if image_response.data: + num_images = len(image_response.data) + return output_cost_per_image * num_images \ No newline at end of file diff --git a/litellm/llms/bedrock/image/amazon_stability3_transformation.py b/litellm/llms/bedrock/image/amazon_stability3_transformation.py index 06e0620979..445a2fe110 100644 --- a/litellm/llms/bedrock/image/amazon_stability3_transformation.py +++ b/litellm/llms/bedrock/image/amazon_stability3_transformation.py @@ -3,6 +3,8 @@ from typing import List, Optional from openai.types.image import Image +from litellm import get_model_info +from litellm.llms.bedrock.common_utils import BedrockError from litellm.types.llms.bedrock import ( AmazonStability3TextToImageRequest, AmazonStability3TextToImageResponse, @@ -66,12 +68,12 @@ class AmazonStability3Config: @classmethod def transform_request_body( - cls, prompt: str, optional_params: dict + cls, text: str, optional_params: dict ) -> AmazonStability3TextToImageRequest: """ Transform the request body for the Stability 3 models """ - data = AmazonStability3TextToImageRequest(prompt=prompt, **optional_params) + data = AmazonStability3TextToImageRequest(prompt=text, **optional_params) return data @classmethod @@ -92,9 +94,34 @@ class AmazonStability3Config: """ stability_3_response = AmazonStability3TextToImageResponse(**response_dict) + + finish_reasons = stability_3_response.get("finish_reasons", []) + finish_reasons = [reason for reason in finish_reasons if reason] + if len(finish_reasons) > 0: + raise BedrockError(status_code=400, message="; ".join(finish_reasons)) + openai_images: List[Image] = [] for _img in stability_3_response.get("images", []): openai_images.append(Image(b64_json=_img)) model_response.data = openai_images return model_response + + @classmethod + def cost_calculator( + cls, + model: str, + image_response: ImageResponse, + size: Optional[str] = None, + optional_params: Optional[dict] = None, + ) -> float: + model_info = get_model_info( + model=model, + custom_llm_provider="bedrock", + ) + + output_cost_per_image: float = model_info.get("output_cost_per_image") or 0.0 + num_images: int = 0 + if image_response.data: + num_images = len(image_response.data) + return output_cost_per_image * num_images diff --git a/litellm/llms/bedrock/image/amazon_titan_transformation.py b/litellm/llms/bedrock/image/amazon_titan_transformation.py index 2709f406df..bed9ad0c30 100644 --- a/litellm/llms/bedrock/image/amazon_titan_transformation.py +++ b/litellm/llms/bedrock/image/amazon_titan_transformation.py @@ -103,16 +103,16 @@ class AmazonTitanImageGenerationConfig: return optional_params @classmethod - def _transform_request( + def transform_request_body( cls, - input: str, + text: str, optional_params: dict, ) -> AmazonTitanImageGenerationRequestBody: from typing import Any, Dict image_generation_config = optional_params.pop("imageGenerationConfig", {}) negative_text = optional_params.pop("negativeText", None) - text_to_image_params: Dict[str, Any] = {"text": input} + text_to_image_params: Dict[str, Any] = {"text": text} if negative_text: text_to_image_params["negativeText"] = negative_text task_type = optional_params.pop("taskType", "TEXT_IMAGE") diff --git a/litellm/llms/bedrock/image/cost_calculator.py b/litellm/llms/bedrock/image/cost_calculator.py index 9b2ae8782c..bc1a57b8ae 100644 --- a/litellm/llms/bedrock/image/cost_calculator.py +++ b/litellm/llms/bedrock/image/cost_calculator.py @@ -1,9 +1,6 @@ from typing import Optional -import litellm -from litellm.llms.bedrock.image.amazon_titan_transformation import ( - AmazonTitanImageGenerationConfig, -) +from litellm.llms.bedrock.image.image_handler import BedrockImageGeneration from litellm.types.utils import ImageResponse @@ -18,36 +15,10 @@ def cost_calculator( Handles both Stability 1 and Stability 3 models """ - if litellm.AmazonStability3Config()._is_stability_3_model(model=model): - pass - elif AmazonTitanImageGenerationConfig._is_titan_model(model=model): - return AmazonTitanImageGenerationConfig.cost_calculator( - model=model, - image_response=image_response, - size=size, - optional_params=optional_params, - ) - else: - # Stability 1 models - optional_params = optional_params or {} - - # see model_prices_and_context_window.json for details on how steps is used - # Reference pricing by steps for stability 1: https://aws.amazon.com/bedrock/pricing/ - _steps = optional_params.get("steps", 50) - steps = "max-steps" if _steps > 50 else "50-steps" - - # size is stored in model_prices_and_context_window.json as 1024-x-1024 - # current size has 1024x1024 - size = size or "1024-x-1024" - model = f"{size}/{steps}/{model}" - - _model_info = litellm.get_model_info( + config_class = BedrockImageGeneration.get_config_class(model=model) + return config_class.cost_calculator( model=model, - custom_llm_provider="bedrock", + image_response=image_response, + size=size, + optional_params=optional_params, ) - - output_cost_per_image: float = _model_info.get("output_cost_per_image") or 0.0 - num_images: int = 0 - if image_response.data: - num_images = len(image_response.data) - return output_cost_per_image * num_images diff --git a/litellm/llms/bedrock/image/image_handler.py b/litellm/llms/bedrock/image/image_handler.py index 313a1dc17b..0825aecc85 100644 --- a/litellm/llms/bedrock/image/image_handler.py +++ b/litellm/llms/bedrock/image/image_handler.py @@ -1,13 +1,10 @@ -import copy import json -import os from typing import TYPE_CHECKING, Any, Optional, Union import httpx from pydantic import BaseModel import litellm -from litellm import BEDROCK_INVOKE_PROVIDERS_LITERAL 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 ( @@ -47,11 +44,30 @@ class BedrockImagePreparedRequest(BaseModel): data: dict +BedrockImageConfigClass = Union[ + type[AmazonTitanImageGenerationConfig], + type[AmazonNovaCanvasConfig], + type[AmazonStability3Config], + type[litellm.AmazonStabilityConfig], +] + + class BedrockImageGeneration(BaseAWSLLM): """ Bedrock Image Generation handler """ + @classmethod + def get_config_class(cls, model: str | None) -> BedrockImageConfigClass: + if AmazonTitanImageGenerationConfig._is_titan_model(model): + return AmazonTitanImageGenerationConfig + elif AmazonNovaCanvasConfig._is_nova_model(model): + return AmazonNovaCanvasConfig + elif AmazonStability3Config._is_stability_3_model(model): + return AmazonStability3Config + else: + return litellm.AmazonStabilityConfig + def image_generation( self, model: str, @@ -202,7 +218,6 @@ class BedrockImageGeneration(BaseAWSLLM): model=model, prompt=prompt, optional_params=optional_params, - bedrock_provider=bedrock_provider, ) # Make POST Request @@ -241,7 +256,6 @@ class BedrockImageGeneration(BaseAWSLLM): def _get_request_body( self, model: str, - bedrock_provider: Optional[BEDROCK_INVOKE_PROVIDERS_LITERAL], prompt: str, optional_params: dict, ) -> dict: @@ -253,49 +267,9 @@ class BedrockImageGeneration(BaseAWSLLM): Returns: dict: The request body to use for the Bedrock Image Generation API """ - if bedrock_provider == "amazon" or bedrock_provider == "nova": - # Handle Amazon Nova Canvas models - provider = "amazon" - elif bedrock_provider == "stability": - provider = "stability" - else: - # Fallback to original logic for backward compatibility - provider = model.split(".")[0] - inference_params = copy.deepcopy(optional_params) - inference_params.pop( - "user", None - ) # make sure user is not passed in for bedrock call - data = {} - if provider == "stability": - if litellm.AmazonStability3Config._is_stability_3_model(model): - request_body = litellm.AmazonStability3Config.transform_request_body( - prompt=prompt, optional_params=optional_params - ) - return dict(request_body) - else: - prompt = prompt.replace(os.linesep, " ") - ## LOAD CONFIG - config = litellm.AmazonStabilityConfig.get_config() - for k, v in config.items(): - if ( - k not in inference_params - ): # completion(top_k=3) > anthropic_config(top_k=3) <- allows for dynamic variables to be passed in - inference_params[k] = v - data = { - "text_prompts": [{"text": prompt, "weight": 1}], - **inference_params, - } - elif provider == "amazon": - return dict( - litellm.AmazonNovaCanvasConfig.transform_request_body( - text=prompt, optional_params=optional_params - ) - ) - else: - raise BedrockError( - status_code=422, message=f"Unsupported model={model}, passed in" - ) - return data + config_class = self.get_config_class(model=model) + request_body = config_class.transform_request_body(text=prompt, optional_params=optional_params) + return dict(request_body) def _transform_response_dict_to_openai_response( self, @@ -323,20 +297,7 @@ class BedrockImageGeneration(BaseAWSLLM): if response_dict is None: raise ValueError("Error in response object format, got None") - config_class: Union[ - type[AmazonTitanImageGenerationConfig], - type[AmazonNovaCanvasConfig], - type[AmazonStability3Config], - type[litellm.AmazonStabilityConfig], - ] - if AmazonTitanImageGenerationConfig._is_titan_model(model=model): - config_class = AmazonTitanImageGenerationConfig - elif AmazonNovaCanvasConfig._is_nova_model(model=model): - config_class = AmazonNovaCanvasConfig - elif AmazonStability3Config._is_stability_3_model(model=model): - config_class = AmazonStability3Config - else: - config_class = litellm.AmazonStabilityConfig + config_class = self.get_config_class(model=model) config_class.transform_response_dict_to_openai_response( model_response=model_response, diff --git a/litellm/utils.py b/litellm/utils.py index 78ed4170f4..302e2ec630 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -2631,16 +2631,7 @@ def get_optional_params_image_gen( ): optional_params = non_default_params elif custom_llm_provider == "bedrock": - # use stability3 config class if model is a stability3 model - config_class = ( - litellm.AmazonStability3Config - if litellm.AmazonStability3Config._is_stability_3_model(model=model) - else ( - litellm.AmazonNovaCanvasConfig - if litellm.AmazonNovaCanvasConfig._is_nova_model(model=model) - else litellm.AmazonStabilityConfig - ) - ) + config_class = litellm.BedrockImageGeneration.get_config_class(model=model) supported_params = config_class.get_supported_openai_params(model=model) _check_valid_arg(supported_params=supported_params) optional_params = config_class.map_openai_params( 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 d3a5ade1ce..5526f22cd5 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 @@ -119,6 +119,20 @@ def test_transform_response_dict_to_openai_response(): assert [img.b64_json for img in result.data] == response_dict["images"] +def test_transform_response_dict_to_openai_response_from_stability_3_models_with_no_null_finish_reason(): + # Create a mock response + response_dict = {"finish_reasons": ["Filter reason: prompt"]} + model_response = ImageResponse() + + with pytest.raises(BedrockError) as exc_info: + AmazonStability3Config.transform_response_dict_to_openai_response( + model_response, response_dict + ) + + assert exc_info.value.status_code == 400 + assert exc_info.value.message == "Filter reason: prompt" + + def test_amazon_stability_get_supported_openai_params(): result = AmazonStabilityConfig.get_supported_openai_params() assert result == ["size"] @@ -168,7 +182,7 @@ def test_get_request_body_stability3(): model = "stability.sd3-large" result = handler._get_request_body( - model=model, bedrock_provider=None, prompt=prompt, optional_params=optional_params + model=model, prompt=prompt, optional_params=optional_params ) assert result["prompt"] == prompt @@ -181,7 +195,7 @@ def test_get_request_body_stability(): model = "stability.stable-diffusion-xl-v1" result = handler._get_request_body( - model=model, bedrock_provider=None, prompt=prompt, optional_params=optional_params + model=model, prompt=prompt, optional_params=optional_params ) assert result["text_prompts"][0]["text"] == prompt @@ -239,7 +253,7 @@ def test_get_request_body_nova_canvas_default(): model = "amazon.nova-canvas-v1" result = handler._get_request_body( - model=model, bedrock_provider=None, prompt=prompt, optional_params=optional_params + model=model, prompt=prompt, optional_params=optional_params ) assert result["taskType"] == "TEXT_IMAGE" @@ -254,7 +268,7 @@ def test_get_request_body_nova_canvas_text_image(): model = "amazon.nova-canvas-v1" result = handler._get_request_body( - model=model, bedrock_provider=None, prompt=prompt, optional_params=optional_params + model=model, prompt=prompt, optional_params=optional_params ) assert result["taskType"] == "TEXT_IMAGE" @@ -273,7 +287,7 @@ def test_get_request_body_nova_canvas_color_guided_generation(): model = "amazon.nova-canvas-v1" result = handler._get_request_body( - model=model, bedrock_provider=None, prompt=prompt, optional_params=optional_params + model=model, prompt=prompt, optional_params=optional_params ) assert result["taskType"] == "COLOR_GUIDED_GENERATION" @@ -437,7 +451,7 @@ def test_get_request_body_nova_canvas_inference_profile_arn(): bedrock_provider = handler.get_bedrock_invoke_provider(model=nova_model) result = handler._get_request_body( - model=nova_model, bedrock_provider=bedrock_provider, prompt=prompt, optional_params=optional_params + model=nova_model, prompt=prompt, optional_params=optional_params ) assert result["taskType"] == "TEXT_IMAGE" @@ -453,7 +467,7 @@ def test_get_request_body_nova_canvas_with_model_id_param(): model = "amazon.nova-canvas-v1" result = handler._get_request_body( - model=model, bedrock_provider=None, prompt=prompt, optional_params=optional_params + model=model, prompt=prompt, optional_params=optional_params ) # After fix, model_id should not appear in the result @@ -488,12 +502,9 @@ def test_get_request_body_cross_region_inference_profile(): # Cross-region inference profile format model = "us.amazon.nova-canvas-v1:0" - # Get the provider using the method from the handler - bedrock_provider = handler.get_bedrock_invoke_provider(model=model) - # This should work after the fix - cross-region format should be detected as 'nova' result = handler._get_request_body( - model=model, bedrock_provider=bedrock_provider, prompt=prompt, optional_params=optional_params + model=model, prompt=prompt, optional_params=optional_params ) assert result["taskType"] == "TEXT_IMAGE" @@ -508,7 +519,7 @@ def test_backward_compatibility_regular_nova_model(): model = "amazon.nova-canvas-v1" result = handler._get_request_body( - model=model, bedrock_provider=None, prompt=prompt, optional_params=optional_params + model=model, prompt=prompt, optional_params=optional_params ) assert result["taskType"] == "TEXT_IMAGE"