From a7cc274b375c898810109bb32cf4851328da4afc Mon Sep 17 00:00:00 2001
From: Yan Zhu <hackzhuyan@gmail.com>
Date: Thu, 13 Aug 2026 18:35:22 +0800
Subject: [PATCH] fix(images): honor custom_llm_provider in aimage_generation

Same routing bug as atranscription: unprefixed proxy models with
custom_llm_provider were dropped before get_llm_provider ran.

Co-authored-by: Cursor <cursoragent@cursor.com>
---
 litellm/images/main.py                        |  8 +++--
 .../test_hosted_vllm_audio_transcription.py   | 33 ++++++++++++++++++-
 2 files changed, 38 insertions(+), 3 deletions(-)

diff --git a/litellm/images/main.py b/litellm/images/main.py
index ae4818b1967..84bbd2b904f 100644
--- a/litellm/images/main.py
+++ b/litellm/images/main.py
@@ -92,7 +92,7 @@ async def aimage_generation(*args, **kwargs) -> ImageResponse:
     model: Final = args[0] if len(args) > 0 else kwargs["model"]
     ### PASS ARGS TO Image Generation ###
     kwargs["aimg_generation"] = True
-    custom_llm_provider = None
+    custom_llm_provider = kwargs.get("custom_llm_provider", None)
     try:
         # Use a partial function to pass your keyword arguments
         func: Final = partial(image_generation, *args, **kwargs)
@@ -101,7 +101,11 @@ async def aimage_generation(*args, **kwargs) -> ImageResponse:
         ctx: Final = contextvars.copy_context()
         func_with_context: Final = partial(ctx.run, func)
 
-        _, custom_llm_provider, _, _ = get_llm_provider(model=model, api_base=kwargs.get("api_base", None))
+        _, custom_llm_provider, _, _ = get_llm_provider(
+            model=model,
+            custom_llm_provider=custom_llm_provider,
+            api_base=kwargs.get("api_base", None),
+        )
 
         # Await normally
         init_response: Final = await loop.run_in_executor(None, func_with_context)
diff --git a/tests/test_litellm/llms/hosted_vllm/transcriptions/test_hosted_vllm_audio_transcription.py b/tests/test_litellm/llms/hosted_vllm/transcriptions/test_hosted_vllm_audio_transcription.py
index 5caea4e30a0..644cd4c17ea 100644
--- a/tests/test_litellm/llms/hosted_vllm/transcriptions/test_hosted_vllm_audio_transcription.py
+++ b/tests/test_litellm/llms/hosted_vllm/transcriptions/test_hosted_vllm_audio_transcription.py
@@ -10,7 +10,7 @@ import litellm
 from litellm.llms.hosted_vllm.transcriptions.transformation import (
     HostedVLLMAudioTranscriptionConfig,
 )
-from litellm.types.utils import TranscriptionResponse
+from litellm.types.utils import ImageResponse, TranscriptionResponse
 
 
 def _complete_url(api_base: str | None) -> str:
@@ -105,3 +105,34 @@ class TestAtranscriptionCustomLlmProvider:
             )
 
         mock_speech.assert_called()
+
+
+class TestAimageGenerationCustomLlmProvider:
+    @pytest.mark.asyncio
+    async def test_unprefixed_model_uses_custom_llm_provider(self) -> None:
+        mock_response = ImageResponse(
+            created=1234567890,
+            data=[{"url": "https://example.com/image.png"}],
+        )
+        with patch(
+            "litellm.images.main.image_generation",
+            return_value=mock_response,
+        ) as mock_image_generation:
+            response = await litellm.aimage_generation(
+                model="unprefixed-image-model",
+                prompt="a cat",
+                custom_llm_provider="openai",
+                api_base="http://vllm.example.com/v1",
+            )
+
+        assert response.data is not None
+        mock_image_generation.assert_called()
+
+    @pytest.mark.asyncio
+    async def test_unprefixed_model_without_provider_still_fails(self) -> None:
+        with pytest.raises(Exception, match="LLM Provider NOT provided"):
+            await litellm.aimage_generation(
+                model="unprefixed-image-model",
+                prompt="a cat",
+                api_base="http://vllm.example.com/v1",
+            )