code: def image_generation(
      self,
      prompt: str,
      timeout: float,
      model: Optional[str] = None,
      api_key: Optional[str] = None,
      api_base: Optional[str] = None,
      api_version: Optional[str] = None,
      model_response: Optional[litellm.utils.ImageResponse] = None,
      azure_ad_token: Optional[str] = None,
      logging_obj=None,
      optional_params=None,
      client=None,
      aimg_generation=None,
  ):
      exception_mapping_worked = False
      try:
          if model and len(model) > 0:
              model = model
          else:
              model = None
          data = {"model": model, "prompt": prompt, **optional_params}
          max_retries = data.pop("max_retries", 2)
          if not isinstance(max_retries, int):
              raise AzureOpenAIError(
                  status_code=422, message="max retries must be an int"
              )
  
          # init AzureOpenAI Client
          azure_client_params = {
              "api_version": api_version,
              "azure_endpoint": api_base,
              "azure_deployment": model,
              "max_retries": max_retries,
              "timeout": timeout,
          }
          if api_key is not None:
              azure_client_params["api_key"] = api_key
          elif azure_ad_token is not None:
              azure_client_params["azure_ad_token"] = azure_ad_token
  
          if aimg_generation == True:
              response = self.aimage_generation(data=data, input=input, logging_obj=logging_obj, model_response=model_response, api_key=api_key, client=client, azure_client_params=azure_client_params, timeout=timeout)  # type: ignore
              return response
  
          if client is None:
              client_session = litellm.client_session or httpx.Client(
                  transport=CustomHTTPTransport(),
              )
              azure_client = AzureOpenAI(http_client=client_session, **azure_client_params)  # type: ignore
          else:
              azure_client = client
  
          ## LOGGING
          logging_obj.pre_call(
              input=prompt,
              api_key=azure_client.api_key,
              additional_args={
                  "headers": {
                      "Authorization": f"Bearer {azure_client.api_key}"
                  },
                  "api_base": azure_client._base_url._uri_reference,
                  "acompletion": False,
                  "complete_input_dict": data,
              },
          )
  
          ## COMPLETION CALL
          response = azure_client.images.generate(**data, timeout=timeout)  # type: ignore
          ## LOGGING
          logging_obj.post_call(
              input=input,
              api_key=api_key,
              additional_args={"complete_input_dict": data},
              original_response=response,
          )
          # return response
          return convert_to_model_response_object(response_object=response.model_dump(), model_response_object=model_response, response_type="image_generation")  # type: ignore
      except AzureOpenAIError as e:
          exception_mapping_worked = True
          raise e
      except Exception as e:
          if hasattr(e, "status_code"):
              raise AzureOpenAIError(
                  status_code=e.status_code, message=str(e)
              )
          else:
              raise AzureOpenAIError(status_code=500, message=str(e))
  
imports: from typing import Optional, Union, Any
  import types, requests
  from .base import BaseLLM
  from litellm.utils import ModelResponse, Choices, Message, CustomStreamWrapper, convert_to_model_response_object
  from typing import Callable, Optional
  from litellm import OpenAIConfig
  import litellm, json
  import httpx
  from .custom_httpx.azure_dall_e_2 import CustomHTTPTransport, AsyncCustomHTTPTransport
  from openai import AzureOpenAI, AsyncAzureOpenAI
label_desc: Use image node_id 6 for calling the AzureChatCompletion obj's underlying member method code with example usage: obj.image_generation(prompt, timeout, model, api_key, api_base, api_version, model_response, azure_ad_token, logging_obj, optional_params, client, aimg_generation) and returns: convert_to_model_response_object, response