diff --git a/litellm/llms/vertex_ai_and_google_ai_studio/image_generation/image_generation_handler.py b/litellm/llms/vertex_ai_and_google_ai_studio/image_generation/image_generation_handler.py
index dac4f08b693..440d0841a84 100644
--- a/litellm/llms/vertex_ai_and_google_ai_studio/image_generation/image_generation_handler.py
+++ b/litellm/llms/vertex_ai_and_google_ai_studio/image_generation/image_generation_handler.py
@@ -16,7 +16,7 @@ class VertexImageGeneration(VertexLLM):
         self,
         json_response: Dict[str, Any],
         model_response: litellm.ImageResponse,
-        model: str,
+        model: Optional[str] = None,
     ) -> litellm.ImageResponse:
         if "predictions" not in json_response:
             raise litellm.InternalServerError(
diff --git a/litellm/proxy/pass_through_endpoints/success_handler.py b/litellm/proxy/pass_through_endpoints/success_handler.py
index f29129df1e8..5d315ae3dc4 100644
--- a/litellm/proxy/pass_through_endpoints/success_handler.py
+++ b/litellm/proxy/pass_through_endpoints/success_handler.py
@@ -1,5 +1,6 @@
 import re
 from datetime import datetime
+from typing import Union
 
 import httpx
 
@@ -117,7 +118,9 @@ class PassThroughEndpointLogging:
             model = self.extract_model_from_url(url_route)
             _json_response = httpx_response.json()
 
-            litellm_model_response = litellm.ModelResponse()
+            litellm_model_response: Union[
+                litellm.ModelResponse, litellm.EmbeddingResponse, litellm.ImageResponse
+            ] = litellm.ModelResponse()
             if vertex_image_generation_class.is_image_generation_response(
                 _json_response
             ):