diff --git a/litellm/llms/azure/azure.py b/litellm/llms/azure/azure.py
index 95afa53b62..76fa713ca8 100644
--- a/litellm/llms/azure/azure.py
+++ b/litellm/llms/azure/azure.py
@@ -1164,6 +1164,7 @@ class AzureChatCompletion(BaseAzureLLM, BaseLLM):
             ## BASE MODEL CHECK
             if (
                 model_response is not None
+                and litellm_params is not None
                 and litellm_params.get("base_model", None) is not None
             ):
                 model_response._hidden_params["model"] = litellm_params.get("base_model", None)
@@ -1172,7 +1173,8 @@ class AzureChatCompletion(BaseAzureLLM, BaseLLM):
             extra_body = optional_params.pop("extra_body", {})
             flattened_params = {**optional_params, **extra_body}
             
-            data = {"model": litellm_params.get("base_model", None) or model, "prompt": prompt, **flattened_params}
+            base_model = litellm_params.get("base_model", None) if litellm_params else None
+            data = {"model": base_model or model, "prompt": prompt, **flattened_params}
             max_retries = data.pop("max_retries", 2)
             if not isinstance(max_retries, int):
                 raise AzureOpenAIError(
diff --git a/tests/test_litellm/llms/azure/image_generation/test_azure_image_generation_init.py b/tests/test_litellm/llms/azure/image_generation/test_azure_image_generation_init.py
index bf8165c690..44bcc9f954 100644
--- a/tests/test_litellm/llms/azure/image_generation/test_azure_image_generation_init.py
+++ b/tests/test_litellm/llms/azure/image_generation/test_azure_image_generation_init.py
@@ -315,22 +315,18 @@ def test_azure_image_generation_base_model_vs_deployment_name():
         logging_obj.post_call = MagicMock()
         
         # Call the image_generation method
-        try:
-            response = azure_chat_completion.image_generation(
-                prompt=prompt,
-                timeout=60.0,
-                optional_params=optional_params,
-                logging_obj=logging_obj,
-                headers={},
-                model=model,
-                api_key=api_key,
-                api_base=api_base,
-                api_version=api_version,
-                litellm_params=litellm_params,
-            )
-        except Exception as e:
-            # If there's an error, we still want to check the mock calls
-            pass
+        response = azure_chat_completion.image_generation(
+            prompt=prompt,
+            timeout=60.0,
+            optional_params=optional_params,
+            logging_obj=logging_obj,
+            headers={},
+            model=model,
+            api_key=api_key,
+            api_base=api_base,
+            api_version=api_version,
+            litellm_params=litellm_params,
+        )
         
         # Verify the mock was called
         assert mock_request.called, "HTTP request should have been made"
@@ -417,21 +413,17 @@ async def test_azure_aimage_generation_base_model_vs_deployment_name():
         logging_obj.post_call = MagicMock()
         
         # Call the aimage_generation method
-        try:
-            response = await azure_chat_completion.aimage_generation(
-                data=data,
-                model_response=None,
-                azure_client_params=azure_client_params,
-                api_key=api_key,
-                input=[],
-                logging_obj=logging_obj,
-                headers={},
-                model=model,  # Pass the deployment name
-                timeout=60.0,
-            )
-        except Exception as e:
-            # If there's an error, we still want to check the mock calls
-            pass
+        response = await azure_chat_completion.aimage_generation(
+            data=data,
+            model_response=None,
+            azure_client_params=azure_client_params,
+            api_key=api_key,
+            input=[],
+            logging_obj=logging_obj,
+            headers={},
+            model=model,  # Pass the deployment name
+            timeout=60.0,
+        )
         
         # Verify the mock was called
         assert mock_request.called, "HTTP request should have been made"