From 931a72523eb4ec0f990bc681439b55e2cd1d78a9 Mon Sep 17 00:00:00 2001
From: Ishaan Jaff <ishaanjaffer0324@gmail.com>
Date: Wed, 19 Jun 2024 15:27:46 -0700
Subject: [PATCH] fix - use router level client on proxy img gen

---
 litellm/main.py | 3 +++
 1 file changed, 3 insertions(+)

diff --git a/litellm/main.py b/litellm/main.py
index 7c775310aef..efc7df0e411 100644
--- a/litellm/main.py
+++ b/litellm/main.py
@@ -3944,6 +3944,7 @@ def image_generation(
         proxy_server_request = kwargs.get("proxy_server_request", None)
         model_info = kwargs.get("model_info", None)
         metadata = kwargs.get("metadata", {})
+        client = kwargs.get("client", None)
 
         model_response = litellm.utils.ImageResponse()
         if model is not None or custom_llm_provider is not None:
@@ -4082,6 +4083,7 @@ def image_generation(
                 model_response=model_response,
                 api_version=api_version,
                 aimg_generation=aimg_generation,
+                client=client,
             )
         elif custom_llm_provider == "openai":
             model_response = openai_chat_completions.image_generation(
@@ -4094,6 +4096,7 @@ def image_generation(
                 optional_params=optional_params,
                 model_response=model_response,
                 aimg_generation=aimg_generation,
+                client=client,
             )
         elif custom_llm_provider == "bedrock":
             if model is None: