From b7e68eccdd6413e8b34725dbe117f38c9776fa19 Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Sun, 12 Jan 2025 16:41:18 -0800 Subject: [PATCH] fixes for img gen cost cal --- litellm/cost_calculator.py | 15 +++++++++++++-- 1 file changed, 13 insertions(+), 2 deletions(-) diff --git a/litellm/cost_calculator.py b/litellm/cost_calculator.py index 14b1f036c13..ff735bb8bf0 100644 --- a/litellm/cost_calculator.py +++ b/litellm/cost_calculator.py @@ -922,8 +922,19 @@ def default_image_cost_calculator( elif base_model_name in litellm.model_cost: cost_info = litellm.model_cost[base_model_name] else: - raise Exception( - f"Model not found in cost map. Tried {model_name_with_quality} and {base_model_name}" + # Try without provider prefix + model_without_provider = f"{size_str}/{model.split('/')[-1]}" + model_with_quality_without_provider = ( + f"{quality}/{model_without_provider}" if quality else model_without_provider ) + if model_with_quality_without_provider in litellm.model_cost: + cost_info = litellm.model_cost[model_with_quality_without_provider] + elif model_without_provider in litellm.model_cost: + cost_info = litellm.model_cost[model_without_provider] + else: + raise Exception( + f"Model not found in cost map. Tried {model_name_with_quality}, {base_model_name}, {model_with_quality_without_provider}, and {model_without_provider}" + ) + return cost_info["input_cost_per_pixel"] * height * width * n