Make Bedrock image generation more consistent (#17021)

komarovd95 · web-flow · commit 046b7efbbeb1 · 2025-11-24T20:58:01.000-08:00
diff --git 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 openai.types.image import Image
 
+from litellm import get_model_info
 from litellm.types.llms.bedrock import (
     AmazonNovaCanvasColorGuidedGenerationParams,
     AmazonNovaCanvasColorGuidedRequest,
@@ -197,3 +198,22 @@ def transform_response_dict_to_openai_response(
 
         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_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 @@ def map_openai_params(
 
         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 @@ def transform_response_dict_to_openai_response(
         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
diff --git a/litellm/llms/bedrock/image/amazon_stability3_transformation.py b/litellm/llms/bedrock/image/amazon_stability3_transformation.py
@@ -3,6 +3,8 @@
 
 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 @@ def _is_stability_3_model(cls, model: Optional[str] = None) -> bool:
 
     @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 @@ def transform_response_dict_to_openai_response(
         """
 
         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
@@ -103,16 +103,16 @@ def map_openai_params(
         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
@@ -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
@@ -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 @@ def _prepare_request(
             model=model,
             prompt=prompt,
             optional_params=optional_params,
-            bedrock_provider=bedrock_provider,
         )
 
         # Make POST Request
@@ -241,7 +256,6 @@ def _prepare_request(
     def _get_request_body(
         self,
         model: str,
-        bedrock_provider: Optional[BEDROCK_INVOKE_PROVIDERS_LITERAL],
         prompt: str,
         optional_params: dict,
     ) -> dict:
@@ -253,49 +267,9 @@ def _get_request_body(
         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 @@ def _transform_response_dict_to_openai_response(
         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
@@ -2631,16 +2631,7 @@ def _check_valid_arg(supported_params):
     ):
         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
