From 25bae4cc237a02c8d46d189898ef631c0935a5a6 Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Mon, 11 Nov 2024 20:21:44 -0800 Subject: [PATCH] (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 | 25 ++++-- litellm/llms/bedrock/image/cost_calculator.py | 41 +++++++++ litellm/utils.py | 1 + .../test_bedrock_image_gen_unit_tests.py | 84 ++++++++++++++++++- .../image_gen_tests/test_image_generation.py | 3 + 5 files changed, 146 insertions(+), 8 deletions(-) create mode 100644 litellm/llms/bedrock/image/cost_calculator.py diff --git a/litellm/cost_calculator.py b/litellm/cost_calculator.py index 0be7f1d38..2aff3b04c 100644 --- a/litellm/cost_calculator.py +++ b/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 diff --git a/litellm/llms/bedrock/image/cost_calculator.py b/litellm/llms/bedrock/image/cost_calculator.py new file mode 100644 index 000000000..0a20b44cb --- /dev/null +++ b/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 diff --git a/litellm/utils.py b/litellm/utils.py index b10c94859..1e8025be4 100644 --- a/litellm/utils.py +++ b/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 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 e04eb2a1a..10845a895 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 @@ -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 diff --git a/tests/image_gen_tests/test_image_generation.py b/tests/image_gen_tests/test_image_generation.py index cf46f90bb..e94d62c1f 100644 --- a/tests/image_gen_tests/test_image_generation.py +++ b/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())