fail ai test fix

litellm/litellm_core_utils/llm_cost_calc/utils.py

@@ -11,8 +11,8 @@ from litellm.types.utils import (
     ImageResponse,
     ModelInfo,
     PassthroughCallTypes,
-    Usage,
     ServiceTier,
+    Usage,
 )
 from litellm.utils import get_model_info
 
@@ -716,6 +716,15 @@ class CostCalculatorUtils:
                 model=model,
                 image_response=completion_response,
             )
+        elif custom_llm_provider == litellm.LlmProviders.FAL_AI.value:
+            from litellm.llms.fal_ai.cost_calculator import (
+                cost_calculator as fal_ai_image_cost_calculator,
+            )
+
+            return fal_ai_image_cost_calculator(
+                model=model,
+                image_response=completion_response,
+            )
         else:
             return default_image_cost_calculator(
                 model=model,