From 96ec2a5a191e307340d8c9c26317599b055de1df Mon Sep 17 00:00:00 2001
From: Ishaan Jaff <ishaanjaffer0324@gmail.com>
Date: Tue, 12 Nov 2024 17:46:27 -0800
Subject: [PATCH] fix response_cost_calculator

---
 litellm/cost_calculator.py | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/litellm/cost_calculator.py b/litellm/cost_calculator.py
index 03a86fb13..6ebec8529 100644
--- a/litellm/cost_calculator.py
+++ b/litellm/cost_calculator.py
@@ -854,6 +854,8 @@ def response_cost_calculator(
             if isinstance(response_object, BaseModel):
                 response_object._hidden_params["optional_params"] = optional_params
             if isinstance(response_object, ImageResponse):
+                if base_model is not None:
+                    model = base_model
                 response_cost = completion_cost(
                     completion_response=response_object,
                     model=model,