fix imagegeneration output_cost_per_image on model cost map

This commit is contained in:
Ishaan Jaff 2024-11-14 17:46:42 -08:00
parent 0585fd56d6
commit 0ee884d4da
4 changed files with 39 additions and 9 deletions

View file

@ -46,6 +46,9 @@ from litellm.llms.OpenAI.cost_calculation import (
from litellm.llms.OpenAI.cost_calculation import cost_per_token as openai_cost_per_token
from litellm.llms.OpenAI.cost_calculation import cost_router as openai_cost_router
from litellm.llms.together_ai.cost_calculator import get_model_params_and_category
from litellm.llms.vertex_ai_and_google_ai_studio.image_generation.cost_calculator import (
cost_calculator as vertex_ai_image_cost_calculator,
)
from litellm.types.llms.openai import HttpxBinaryResponseContent
from litellm.types.rerank import RerankResponse
from litellm.types.router import SPECIAL_MODEL_INFO_PARAMS
@ -667,9 +670,11 @@ def completion_cost( # noqa: PLR0915
):
### IMAGE GENERATION COST CALCULATION ###
if custom_llm_provider == "vertex_ai":
# https://cloud.google.com/vertex-ai/generative-ai/pricing
# Vertex Charges Flat $0.20 per image
return 0.020
if isinstance(completion_response, ImageResponse):
return vertex_ai_image_cost_calculator(
model=model,
image_response=completion_response,
)
elif custom_llm_provider == "bedrock":
if isinstance(completion_response, ImageResponse):
return bedrock_image_cost_calculator(

View file

@ -0,0 +1,25 @@
"""
Vertex AI Image Generation Cost Calculator
"""
from typing import Optional
import litellm
from litellm.types.utils import ImageResponse
def cost_calculator(
model: str,
image_response: ImageResponse,
) -> float:
"""
Vertex AI Image Generation Cost Calculator
"""
_model_info = litellm.get_model_info(
model=model,
custom_llm_provider="vertex_ai",
)
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

View file

@ -2986,19 +2986,19 @@
"supports_function_calling": true
},
"vertex_ai/imagegeneration@006": {
"cost_per_image": 0.020,
"output_cost_per_image": 0.020,
"litellm_provider": "vertex_ai-image-models",
"mode": "image_generation",
"source": "https://cloud.google.com/vertex-ai/generative-ai/pricing"
},
"vertex_ai/imagen-3.0-generate-001": {
"cost_per_image": 0.04,
"output_cost_per_image": 0.04,
"litellm_provider": "vertex_ai-image-models",
"mode": "image_generation",
"source": "https://cloud.google.com/vertex-ai/generative-ai/pricing"
},
"vertex_ai/imagen-3.0-fast-generate-001": {
"cost_per_image": 0.02,
"output_cost_per_image": 0.02,
"litellm_provider": "vertex_ai-image-models",
"mode": "image_generation",
"source": "https://cloud.google.com/vertex-ai/generative-ai/pricing"

View file

@ -2986,19 +2986,19 @@
"supports_function_calling": true
},
"vertex_ai/imagegeneration@006": {
"cost_per_image": 0.020,
"output_cost_per_image": 0.020,
"litellm_provider": "vertex_ai-image-models",
"mode": "image_generation",
"source": "https://cloud.google.com/vertex-ai/generative-ai/pricing"
},
"vertex_ai/imagen-3.0-generate-001": {
"cost_per_image": 0.04,
"output_cost_per_image": 0.04,
"litellm_provider": "vertex_ai-image-models",
"mode": "image_generation",
"source": "https://cloud.google.com/vertex-ai/generative-ai/pricing"
},
"vertex_ai/imagen-3.0-fast-generate-001": {
"cost_per_image": 0.02,
"output_cost_per_image": 0.02,
"litellm_provider": "vertex_ai-image-models",
"mode": "image_generation",
"source": "https://cloud.google.com/vertex-ai/generative-ai/pricing"