fix(xai): honor deployment image pixel pricing (#50490)

Co-authored-by: devin-ai-integration[bot] <158243242+devin-ai-integration[bot]@users.noreply.github.com>

litellm/llms/xai/image_generation/cost_calculator.py

@@ -28,6 +28,9 @@ def cost_calculator(
     call_type: str | None = None,
     model_info: ModelInfo | None = None,
 ) -> float | None:
+    if model_info is not None and model_info.get("input_cost_per_pixel") is not None:
+        return None
+
     caller_set_prices: Final = model_info is not None and any(
         model_info.get(key) is not None
         for key in (

tests/unit/llms/xai/image_generation/test_cost_calculator.py

@@ -1,3 +1,4 @@
+from datetime import datetime
 from typing import Final
 
 import httpx
@@ -6,6 +7,7 @@ from openai import AsyncOpenAI, OpenAI
 
 import litellm
 from litellm.litellm_core_utils.llm_cost_calc.utils import CostCalculatorUtils
+from litellm.litellm_core_utils.litellm_logging import Logging
 from litellm.litellm_core_utils.llm_response_utils.convert_dict_to_response import convert_to_model_response_object
 from litellm.types.utils import ImageResponse
 
@@ -141,6 +143,92 @@ def test_deployment_image_price_overrides_billed_ticks(price_key: str, resolutio
     ) == pytest.approx(0.83 * 3)
 
 
+@pytest.mark.parametrize("ticks", (None, -1, True, "600000000", 0, 600000000))
+@pytest.mark.parametrize("pixel_rate", (0.0, 0.000001))
+@pytest.mark.parametrize(
+    "size,pixels,n", ((None, 1024 * 1024, 3), ("1024x1024", 1024 * 1024, 1), ("512x256", 512 * 256, 2))
+)
+def test_deployment_pixel_price_overrides_ticks_and_resolution_prices(
+    ticks: object, pixel_rate: float, size: str | None, pixels: int, n: int
+) -> None:
+    assert CostCalculatorUtils.route_image_generation_cost_calculator(
+        model="grok-imagine-image-2.0",
+        custom_llm_provider="xai",
+        completion_response=_response(ticks),
+        size=size,
+        quality="low",
+        n=n,
+        model_info={"input_cost_per_pixel": pixel_rate, "output_cost_per_image_2K": 0.29},
+        optional_params={"resolution": "2k"},
+        call_type="image_generation",
+    ) == pytest.approx(pixel_rate * pixels * n)
+
+
+@pytest.mark.parametrize("ticks", (None, 600000000))
+@pytest.mark.parametrize("price_key", ("input_cost_per_image", "output_cost_per_image"))
+def test_deployment_flat_price_still_precedes_pixel_price(ticks: object, price_key: str) -> None:
+    assert CostCalculatorUtils.route_image_generation_cost_calculator(
+        model="unmapped-image-fixture",
+        custom_llm_provider="xai",
+        completion_response=_response(ticks),
+        n=2,
+        quality="low",
+        model_info={price_key: 0.83, "input_cost_per_pixel": 0.000001, "output_cost_per_image_2K": 0.29},
+        optional_params={"resolution": "2k"},
+    ) == pytest.approx(0.83 * 2)
+
+
+@pytest.mark.asyncio
+async def test_router_preserves_deployment_pixel_price_through_image_normalization() -> None:
+    ticks: Final = 600000000
+    transport: Final = httpx.MockTransport(
+        lambda request: httpx.Response(
+            200,
+            json={
+                "created": 123,
+                "data": [{"url": "https://image.test/output.png"}],
+                "usage": {"cost_in_usd_ticks": ticks},
+            },
+        )
+    )
+    router: Final = litellm.Router(
+        model_list=[
+            {
+                "model_name": "pixel-fixture",
+                "litellm_params": {
+                    "model": "xai/grok-imagine-image-2.0",
+                    "api_key": "test",
+                    "api_base": "https://image.test/v1",
+                },
+                "model_info": {"input_cost_per_pixel": 0.000001},
+            }
+        ]
+    )
+    logging_obj: Final = Logging(
+        model="grok-imagine-image-2.0",
+        messages=[],
+        stream=False,
+        call_type="aimage_generation",
+        start_time=datetime(2026, 1, 1),
+        litellm_call_id="pixel-override",
+        function_id="pixel-override",
+    )
+    async with AsyncOpenAI(
+        api_key="test", base_url="https://image.test/v1", http_client=httpx.AsyncClient(transport=transport)
+    ) as client:
+        response: Final = await router.aimage_generation(
+            model="pixel-fixture",
+            prompt="circle",
+            quality="low",
+            n=1,
+            extra_body={"resolution": "2k"},
+            client=client,
+            litellm_logging_obj=logging_obj,
+        )
+        assert getattr(response.usage, "cost_in_usd_ticks") == ticks
+        assert logging_obj.response_cost_calculator(result=response) == pytest.approx(1024 * 1024 * 0.000001)
+
+
 def test_shared_resolution_prices_are_used_when_ticks_are_absent() -> None:
     pricing: Final = litellm.get_model_info("xai/grok-imagine-image-2.0")
     assert CostCalculatorUtils.route_image_generation_cost_calculator(