(feat) add cost tracking stable diffusion 3 on Bedrock (#6676)

* add cost tracking for sd3

* test_image_generation_bedrock

* fix get model info for image cost

* add cost_calculator for stability 1 models

* add unit testing for bedrock image cost calc

* test_cost_calculator_with_no_optional_params

* add test_cost_calculator_basic

* correctly allow size Optional

* fix cost_calculator

* sd3 unit tests cost calc

litellm/cost_calculator.py

@@ -28,6 +28,9 @@ from litellm.llms.azure_ai.cost_calculator import (
 from litellm.llms.AzureOpenAI.cost_calculation import (
     cost_per_token as azure_openai_cost_per_token,
 )
+from litellm.llms.bedrock.image.cost_calculator import (
+    cost_calculator as bedrock_image_cost_calculator,
+)
 from litellm.llms.cohere.cost_calculator import (
     cost_per_query as cohere_rerank_cost_per_query,
 )
@@ -521,12 +524,13 @@ def completion_cost(  # noqa: PLR0915
     custom_llm_provider=None,
     region_name=None,  # used for bedrock pricing
     ### IMAGE GEN ###
-    size=None,
+    size: Optional[str] = None,
     quality=None,
     n=None,  # number of images
     ### CUSTOM PRICING ###
     custom_cost_per_token: Optional[CostPerToken] = None,
     custom_cost_per_second: Optional[float] = None,
+    optional_params: Optional[dict] = None,
 ) -> float:
     """
     Calculate the cost of a given completion call fot GPT-3.5-turbo, llama2, any litellm supported llm.
@@ -667,7 +671,17 @@ def completion_cost(  # noqa: PLR0915
                 # https://cloud.google.com/vertex-ai/generative-ai/pricing
                 # Vertex Charges Flat $0.20 per image
                 return 0.020
-
+            elif custom_llm_provider == "bedrock":
+                if isinstance(completion_response, ImageResponse):
+                    return bedrock_image_cost_calculator(
+                        model=model,
+                        size=size,
+                        image_response=completion_response,
+                        optional_params=optional_params,
+                    )
+                raise TypeError(
+                    "completion_response must be of type ImageResponse for bedrock image cost calculation"
+                )
             if size is None:
                 size = "1024-x-1024"  # openai default
             # fix size to match naming convention
@@ -677,9 +691,9 @@ def completion_cost(  # noqa: PLR0915
             image_gen_model_name_with_quality = image_gen_model_name
             if quality is not None:
                 image_gen_model_name_with_quality = f"{quality}/{image_gen_model_name}"
-            size = size.split("-x-")
-            height = int(size[0])  # if it's 1024-x-1024 vs. 1024x1024
-            width = int(size[1])
+            size_parts = size.split("-x-")
+            height = int(size_parts[0])  # if it's 1024-x-1024 vs. 1024x1024
+            width = int(size_parts[1])
             verbose_logger.debug(f"image_gen_model_name: {image_gen_model_name}")
             verbose_logger.debug(
                 f"image_gen_model_name_with_quality: {image_gen_model_name_with_quality}"
@@ -844,6 +858,7 @@ def response_cost_calculator(
                     model=model,
                     call_type=call_type,
                     custom_llm_provider=custom_llm_provider,
+                    optional_params=optional_params,
                 )
             else:
                 if custom_pricing is True:  # override defaults if custom pricing is set

litellm/llms/bedrock/image/cost_calculator.py

@@ -0,0 +1,41 @@
+from typing import Optional
+
+import litellm
+from litellm.types.utils import ImageResponse
+
+
+def cost_calculator(
+    model: str,
+    image_response: ImageResponse,
+    size: Optional[str] = None,
+    optional_params: Optional[dict] = None,
+) -> float:
+    """
+    Bedrock image generation cost calculator
+
+    Handles both Stability 1 and Stability 3 models
+    """
+    if litellm.AmazonStability3Config()._is_stability_3_model(model=model):
+        pass
+    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(
+        model=model,
+        custom_llm_provider="bedrock",
+    )
+
+    output_cost_per_image: float = _model_info.get("output_cost_per_image") or 0.0
+    num_images: int = len(image_response.data)
+    return output_cost_per_image * num_images

litellm/utils.py

@@ -4636,6 +4636,7 @@ def get_model_info(  # noqa: PLR0915
                     "output_cost_per_character_above_128k_tokens", None
                 ),
                 output_cost_per_second=_model_info.get("output_cost_per_second", None),
+                output_cost_per_image=_model_info.get("output_cost_per_image", None),
                 output_vector_size=_model_info.get("output_vector_size", None),
                 litellm_provider=_model_info.get(
                     "litellm_provider", custom_llm_provider

tests/image_gen_tests/test_bedrock_image_gen_unit_tests.py

@@ -9,12 +9,14 @@ from openai.types.image import Image
 logging.basicConfig(level=logging.DEBUG)
 load_dotenv()
 import asyncio
-import os
 
 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.types.utils import ImageResponse, ImageObject
+import os
 
 import litellm
 from litellm.llms.bedrock.image.amazon_stability3_transformation import (
@@ -27,7 +29,6 @@ from litellm.types.llms.bedrock import (
     AmazonStability3TextToImageRequest,
     AmazonStability3TextToImageResponse,
 )
-from litellm.types.utils import ImageResponse
 from unittest.mock import MagicMock, patch
 from litellm.llms.bedrock.image.image_handler import (
     BedrockImageGeneration,
@@ -149,7 +150,7 @@ def test_get_request_body_stability():
     handler = BedrockImageGeneration()
     prompt = "A beautiful sunset"
     optional_params = {"cfg_scale": 7}
-    model = "stability.stable-diffusion-xl"
+    model = "stability.stable-diffusion-xl-v1"
 
     result = handler._get_request_body(
         model=model, prompt=prompt, optional_params=optional_params
@@ -185,3 +186,80 @@ def test_transform_response_dict_to_openai_response_stability3():
     assert len(result.data) == 2
     assert all(hasattr(img, "b64_json") for img in result.data)
     assert [img.b64_json for img in result.data] == ["base64_image_1", "base64_image_2"]
+
+
+def test_cost_calculator_stability3():
+    # Mock image response
+    image_response = ImageResponse(
+        data=[
+            ImageObject(b64_json="base64_image_1"),
+            ImageObject(b64_json="base64_image_2"),
+        ]
+    )
+
+    cost = cost_calculator(
+        model="stability.sd3-large-v1:0",
+        size="1024-x-1024",
+        image_response=image_response,
+    )
+
+    print("cost", cost)
+
+    # Assert cost is calculated correctly for 2 images
+    assert isinstance(cost, float)
+    assert cost > 0
+
+
+def test_cost_calculator_stability1():
+    # Mock image response
+    image_response = ImageResponse(data=[ImageObject(b64_json="base64_image_1")])
+
+    # Test with different step configurations
+    cost_default_steps = cost_calculator(
+        model="stability.stable-diffusion-xl-v1",
+        size="1024-x-1024",
+        image_response=image_response,
+        optional_params={"steps": 50},
+    )
+
+    cost_max_steps = cost_calculator(
+        model="stability.stable-diffusion-xl-v1",
+        size="1024-x-1024",
+        image_response=image_response,
+        optional_params={"steps": 51},
+    )
+
+    # Assert costs are calculated correctly
+    assert isinstance(cost_default_steps, float)
+    assert isinstance(cost_max_steps, float)
+    assert cost_default_steps > 0
+    assert cost_max_steps > 0
+    # Max steps should be more expensive
+    assert cost_max_steps > cost_default_steps
+
+
+def test_cost_calculator_with_no_optional_params():
+    image_response = ImageResponse(data=[ImageObject(b64_json="base64_image_1")])
+
+    cost = cost_calculator(
+        model="stability.stable-diffusion-xl-v0",
+        size="512-x-512",
+        image_response=image_response,
+        optional_params=None,
+    )
+
+    assert isinstance(cost, float)
+    assert cost > 0
+
+
+def test_cost_calculator_basic():
+    image_response = ImageResponse(data=[ImageObject(b64_json="base64_image_1")])
+
+    cost = cost_calculator(
+        model="stability.stable-diffusion-xl-v1",
+        image_response=image_response,
+        optional_params=None,
+    )
+
+    assert isinstance(cost, float)
+    assert cost > 0

tests/image_gen_tests/test_image_generation.py

@@ -253,6 +253,9 @@ def test_image_generation_bedrock(model):
         )
 
         print(f"response: {response}")
+        print("response hidden params", response._hidden_params)
+
+        assert response._hidden_params["response_cost"] is not None
         from openai.types.images_response import ImagesResponse
 
         ImagesResponse.model_validate(response.model_dump())