google-genai 1.2.0__py3-none-any.whl → 1.3.0__py3-none-any.whl

google/genai/models.py

@@ -28,6 +28,8 @@ from ._common import get_value_by_path as getv
 from ._common import set_value_by_path as setv
 from .pagers import AsyncPager, Pager
 
+logger = logging.getLogger('google_genai.models')
+
 
 def _Part_to_mldev(
     api_client: ApiClient,
@@ -2789,6 +2791,202 @@ def _ComputeTokensParameters_to_vertex(
   return to_object
 
 
+def _GenerateVideosConfig_to_mldev(
+    api_client: ApiClient,
+    from_object: Union[dict, object],
+    parent_object: dict = None,
+) -> dict:
+  to_object = {}
+
+  if getv(from_object, ['number_of_videos']) is not None:
+    setv(
+        parent_object,
+        ['parameters', 'sampleCount'],
+        getv(from_object, ['number_of_videos']),
+    )
+
+  if getv(from_object, ['output_gcs_uri']) is not None:
+    raise ValueError('output_gcs_uri parameter is not supported in Gemini API.')
+
+  if getv(from_object, ['fps']) is not None:
+    raise ValueError('fps parameter is not supported in Gemini API.')
+
+  if getv(from_object, ['duration_seconds']) is not None:
+    raise ValueError(
+        'duration_seconds parameter is not supported in Gemini API.'
+    )
+
+  if getv(from_object, ['seed']) is not None:
+    raise ValueError('seed parameter is not supported in Gemini API.')
+
+  if getv(from_object, ['aspect_ratio']) is not None:
+    setv(
+        parent_object,
+        ['parameters', 'aspectRatio'],
+        getv(from_object, ['aspect_ratio']),
+    )
+
+  if getv(from_object, ['resolution']) is not None:
+    raise ValueError('resolution parameter is not supported in Gemini API.')
+
+  if getv(from_object, ['person_generation']) is not None:
+    setv(
+        parent_object,
+        ['parameters', 'personGeneration'],
+        getv(from_object, ['person_generation']),
+    )
+
+  if getv(from_object, ['pubsub_topic']) is not None:
+    raise ValueError('pubsub_topic parameter is not supported in Gemini API.')
+
+  if getv(from_object, ['negative_prompt']) is not None:
+    setv(
+        parent_object,
+        ['parameters', 'negativePrompt'],
+        getv(from_object, ['negative_prompt']),
+    )
+
+  if getv(from_object, ['enhance_prompt']) is not None:
+    raise ValueError('enhance_prompt parameter is not supported in Gemini API.')
+
+  return to_object
+
+
+def _GenerateVideosConfig_to_vertex(
+    api_client: ApiClient,
+    from_object: Union[dict, object],
+    parent_object: dict = None,
+) -> dict:
+  to_object = {}
+
+  if getv(from_object, ['number_of_videos']) is not None:
+    setv(
+        parent_object,
+        ['parameters', 'sampleCount'],
+        getv(from_object, ['number_of_videos']),
+    )
+
+  if getv(from_object, ['output_gcs_uri']) is not None:
+    setv(
+        parent_object,
+        ['parameters', 'storageUri'],
+        getv(from_object, ['output_gcs_uri']),
+    )
+
+  if getv(from_object, ['fps']) is not None:
+    setv(parent_object, ['parameters', 'fps'], getv(from_object, ['fps']))
+
+  if getv(from_object, ['duration_seconds']) is not None:
+    setv(
+        parent_object,
+        ['parameters', 'durationSeconds'],
+        getv(from_object, ['duration_seconds']),
+    )
+
+  if getv(from_object, ['seed']) is not None:
+    setv(parent_object, ['parameters', 'seed'], getv(from_object, ['seed']))
+
+  if getv(from_object, ['aspect_ratio']) is not None:
+    setv(
+        parent_object,
+        ['parameters', 'aspectRatio'],
+        getv(from_object, ['aspect_ratio']),
+    )
+
+  if getv(from_object, ['resolution']) is not None:
+    setv(
+        parent_object,
+        ['parameters', 'resolution'],
+        getv(from_object, ['resolution']),
+    )
+
+  if getv(from_object, ['person_generation']) is not None:
+    setv(
+        parent_object,
+        ['parameters', 'personGeneration'],
+        getv(from_object, ['person_generation']),
+    )
+
+  if getv(from_object, ['pubsub_topic']) is not None:
+    setv(
+        parent_object,
+        ['parameters', 'pubsubTopic'],
+        getv(from_object, ['pubsub_topic']),
+    )
+
+  if getv(from_object, ['negative_prompt']) is not None:
+    setv(
+        parent_object,
+        ['parameters', 'negativePrompt'],
+        getv(from_object, ['negative_prompt']),
+    )
+
+  if getv(from_object, ['enhance_prompt']) is not None:
+    setv(
+        parent_object,
+        ['parameters', 'enhancePrompt'],
+        getv(from_object, ['enhance_prompt']),
+    )
+
+  return to_object
+
+
+def _GenerateVideosParameters_to_mldev(
+    api_client: ApiClient,
+    from_object: Union[dict, object],
+    parent_object: dict = None,
+) -> dict:
+  to_object = {}
+  if getv(from_object, ['model']) is not None:
+    setv(
+        to_object,
+        ['_url', 'model'],
+        t.t_model(api_client, getv(from_object, ['model'])),
+    )
+
+  if getv(from_object, ['prompt']) is not None:
+    setv(to_object, ['instances[0]', 'prompt'], getv(from_object, ['prompt']))
+
+  if getv(from_object, ['config']) is not None:
+    setv(
+        to_object,
+        ['config'],
+        _GenerateVideosConfig_to_mldev(
+            api_client, getv(from_object, ['config']), to_object
+        ),
+    )
+
+  return to_object
+
+
+def _GenerateVideosParameters_to_vertex(
+    api_client: ApiClient,
+    from_object: Union[dict, object],
+    parent_object: dict = None,
+) -> dict:
+  to_object = {}
+  if getv(from_object, ['model']) is not None:
+    setv(
+        to_object,
+        ['_url', 'model'],
+        t.t_model(api_client, getv(from_object, ['model'])),
+    )
+
+  if getv(from_object, ['prompt']) is not None:
+    setv(to_object, ['instances[0]', 'prompt'], getv(from_object, ['prompt']))
+
+  if getv(from_object, ['config']) is not None:
+    setv(
+        to_object,
+        ['config'],
+        _GenerateVideosConfig_to_vertex(
+            api_client, getv(from_object, ['config']), to_object
+        ),
+    )
+
+  return to_object
+
+
 def _SafetyFilterLevel_to_mldev_enum_validate(enum_value: Any):
   if enum_value in set(['BLOCK_NONE']):
     raise ValueError(f'{enum_value} enum value is not supported in Gemini API.')
@@ -3788,6 +3986,216 @@ def _ComputeTokensResponse_from_vertex(
   return to_object
 
 
+def _Video_from_mldev(
+    api_client: ApiClient,
+    from_object: Union[dict, object],
+    parent_object: dict = None,
+) -> dict:
+  to_object = {}
+  if getv(from_object, ['uri']) is not None:
+    setv(to_object, ['uri'], getv(from_object, ['uri']))
+
+  if getv(from_object, ['encodedVideo']) is not None:
+    setv(
+        to_object,
+        ['video_bytes'],
+        t.t_bytes(api_client, getv(from_object, ['encodedVideo'])),
+    )
+
+  if getv(from_object, ['encoding']) is not None:
+    setv(to_object, ['mime_type'], getv(from_object, ['encoding']))
+
+  return to_object
+
+
+def _Video_from_vertex(
+    api_client: ApiClient,
+    from_object: Union[dict, object],
+    parent_object: dict = None,
+) -> dict:
+  to_object = {}
+  if getv(from_object, ['gcsUri']) is not None:
+    setv(to_object, ['uri'], getv(from_object, ['gcsUri']))
+
+  if getv(from_object, ['bytesBase64Encoded']) is not None:
+    setv(
+        to_object,
+        ['video_bytes'],
+        t.t_bytes(api_client, getv(from_object, ['bytesBase64Encoded'])),
+    )
+
+  if getv(from_object, ['mimeType']) is not None:
+    setv(to_object, ['mime_type'], getv(from_object, ['mimeType']))
+
+  return to_object
+
+
+def _GeneratedVideo_from_mldev(
+    api_client: ApiClient,
+    from_object: Union[dict, object],
+    parent_object: dict = None,
+) -> dict:
+  to_object = {}
+  if getv(from_object, ['_self']) is not None:
+    setv(
+        to_object,
+        ['video'],
+        _Video_from_mldev(api_client, getv(from_object, ['_self']), to_object),
+    )
+
+  return to_object
+
+
+def _GeneratedVideo_from_vertex(
+    api_client: ApiClient,
+    from_object: Union[dict, object],
+    parent_object: dict = None,
+) -> dict:
+  to_object = {}
+  if getv(from_object, ['_self']) is not None:
+    setv(
+        to_object,
+        ['video'],
+        _Video_from_vertex(api_client, getv(from_object, ['_self']), to_object),
+    )
+
+  return to_object
+
+
+def _GenerateVideosResponse_from_mldev(
+    api_client: ApiClient,
+    from_object: Union[dict, object],
+    parent_object: dict = None,
+) -> dict:
+  to_object = {}
+  if getv(from_object, ['videos']) is not None:
+    setv(
+        to_object,
+        ['generated_videos'],
+        [
+            _GeneratedVideo_from_mldev(api_client, item, to_object)
+            for item in getv(from_object, ['videos'])
+        ],
+    )
+
+  if getv(from_object, ['raiMediaFilteredCount']) is not None:
+    setv(
+        to_object,
+        ['rai_media_filtered_count'],
+        getv(from_object, ['raiMediaFilteredCount']),
+    )
+
+  if getv(from_object, ['raiMediaFilteredReasons']) is not None:
+    setv(
+        to_object,
+        ['rai_media_filtered_reasons'],
+        getv(from_object, ['raiMediaFilteredReasons']),
+    )
+
+  return to_object
+
+
+def _GenerateVideosResponse_from_vertex(
+    api_client: ApiClient,
+    from_object: Union[dict, object],
+    parent_object: dict = None,
+) -> dict:
+  to_object = {}
+  if getv(from_object, ['videos']) is not None:
+    setv(
+        to_object,
+        ['generated_videos'],
+        [
+            _GeneratedVideo_from_vertex(api_client, item, to_object)
+            for item in getv(from_object, ['videos'])
+        ],
+    )
+
+  if getv(from_object, ['raiMediaFilteredCount']) is not None:
+    setv(
+        to_object,
+        ['rai_media_filtered_count'],
+        getv(from_object, ['raiMediaFilteredCount']),
+    )
+
+  if getv(from_object, ['raiMediaFilteredReasons']) is not None:
+    setv(
+        to_object,
+        ['rai_media_filtered_reasons'],
+        getv(from_object, ['raiMediaFilteredReasons']),
+    )
+
+  return to_object
+
+
+def _GenerateVideosOperation_from_mldev(
+    api_client: ApiClient,
+    from_object: Union[dict, object],
+    parent_object: dict = None,
+) -> dict:
+  to_object = {}
+  if getv(from_object, ['name']) is not None:
+    setv(to_object, ['name'], getv(from_object, ['name']))
+
+  if getv(from_object, ['metadata']) is not None:
+    setv(to_object, ['metadata'], getv(from_object, ['metadata']))
+
+  if getv(from_object, ['done']) is not None:
+    setv(to_object, ['done'], getv(from_object, ['done']))
+
+  if getv(from_object, ['error']) is not None:
+    setv(to_object, ['error'], getv(from_object, ['error']))
+
+  if getv(from_object, ['response']) is not None:
+    setv(to_object, ['response'], getv(from_object, ['response']))
+
+  if getv(from_object, ['response', 'generateVideoResponse']) is not None:
+    setv(
+        to_object,
+        ['result'],
+        _GenerateVideosResponse_from_mldev(
+            api_client,
+            getv(from_object, ['response', 'generateVideoResponse']),
+            to_object,
+        ),
+    )
+
+  return to_object
+
+
+def _GenerateVideosOperation_from_vertex(
+    api_client: ApiClient,
+    from_object: Union[dict, object],
+    parent_object: dict = None,
+) -> dict:
+  to_object = {}
+  if getv(from_object, ['name']) is not None:
+    setv(to_object, ['name'], getv(from_object, ['name']))
+
+  if getv(from_object, ['metadata']) is not None:
+    setv(to_object, ['metadata'], getv(from_object, ['metadata']))
+
+  if getv(from_object, ['done']) is not None:
+    setv(to_object, ['done'], getv(from_object, ['done']))
+
+  if getv(from_object, ['error']) is not None:
+    setv(to_object, ['error'], getv(from_object, ['error']))
+
+  if getv(from_object, ['response']) is not None:
+    setv(to_object, ['response'], getv(from_object, ['response']))
+
+  if getv(from_object, ['response']) is not None:
+    setv(
+        to_object,
+        ['result'],
+        _GenerateVideosResponse_from_vertex(
+            api_client, getv(from_object, ['response']), to_object
+        ),
+    )
+
+  return to_object
+
+
 class Models(_api_module.BaseModule):
 
   def _generate_content(
@@ -3915,7 +4323,7 @@ class Models(_api_module.BaseModule):
       contents: Union[types.ContentListUnion, types.ContentListUnionDict],
       config: Optional[types.EmbedContentConfigOrDict] = None,
   ) -> types.EmbedContentResponse:
-    """Calculates embeddings for the given contents(only text is supported).
+    """Calculates embeddings for the given contents. Only text is supported.
 
     Args:
       model (str): The model to use.
@@ -4007,7 +4415,7 @@ class Models(_api_module.BaseModule):
     .. code-block:: python
 
       response = client.models.generate_images(
-        model='imagen-3.0-generate-001',
+        model='imagen-3.0-generate-002',
         prompt='Man with a dog',
         config=types.GenerateImagesConfig(
             number_of_images= 1,
@@ -4447,10 +4855,11 @@ class Models(_api_module.BaseModule):
   ) -> types.CountTokensResponse:
     """Counts the number of tokens in the given content.
 
+    Multimodal input is supported for Gemini models.
+
     Args:
       model (str): The model to use for counting tokens.
       contents (list[types.Content]): The content to count tokens for.
-        Multimodal input is supported for Gemini models.
       config (CountTokensConfig): The configuration for counting tokens.
 
     Usage:
@@ -4522,14 +4931,15 @@ class Models(_api_module.BaseModule):
       contents: Union[types.ContentListUnion, types.ContentListUnionDict],
       config: Optional[types.ComputeTokensConfigOrDict] = None,
   ) -> types.ComputeTokensResponse:
-    """Return a list of tokens based on the input text.
+    """Return a list of tokens based on the input contents.
+
+    Only text is supported.
 
     This method is not supported by the Gemini Developer API.
 
     Args:
       model (str): The model to use.
-      contents (list[shared.Content]): The content to compute tokens for. Only
-        text is supported.
+      contents (list[shared.Content]): The content to compute tokens for.
 
     Usage:
 
@@ -4592,6 +5002,88 @@ class Models(_api_module.BaseModule):
     self._api_client._verify_response(return_value)
     return return_value
 
+  @_common.experimental_warning(
+      'This method is experimental and may change in future versions.'
+  )
+  def generate_videos(
+      self,
+      *,
+      model: str,
+      prompt: Optional[str] = None,
+      config: Optional[types.GenerateVideosConfigOrDict] = None,
+  ) -> types.GenerateVideosOperation:
+    """Generates videos based on a text description and configuration.
+
+    Args:
+      model: The model to use.
+      instances: A list of prompts, images and videos to generate videos from.
+      config: Configuration for generation.
+
+    Usage:
+
+      ```
+      operation = client.models.generate_videos(
+          model="veo-2.0-generate-001",
+          prompt="A neon hologram of a cat driving at top speed",
+      )
+      while not operation.done:
+          time.sleep(10)
+          operation = client.operations.get(operation)
+
+      operation.result.generated_videos[0].video.uri
+      ```
+    """
+
+    parameter_model = types._GenerateVideosParameters(
+        model=model,
+        prompt=prompt,
+        config=config,
+    )
+
+    if self._api_client.vertexai:
+      request_dict = _GenerateVideosParameters_to_vertex(
+          self._api_client, parameter_model
+      )
+      path = '{model}:predictLongRunning'.format_map(request_dict.get('_url'))
+    else:
+      request_dict = _GenerateVideosParameters_to_mldev(
+          self._api_client, parameter_model
+      )
+      path = '{model}:predictLongRunning'.format_map(request_dict.get('_url'))
+    query_params = request_dict.get('_query')
+    if query_params:
+      path = f'{path}?{urlencode(query_params)}'
+    # TODO: remove the hack that pops config.
+    request_dict.pop('config', None)
+
+    http_options = None
+    if isinstance(config, dict):
+      http_options = config.get('http_options', None)
+    elif hasattr(config, 'http_options'):
+      http_options = config.http_options
+
+    request_dict = _common.convert_to_dict(request_dict)
+    request_dict = _common.encode_unserializable_types(request_dict)
+
+    response_dict = self._api_client.request(
+        'post', path, request_dict, http_options
+    )
+
+    if self._api_client.vertexai:
+      response_dict = _GenerateVideosOperation_from_vertex(
+          self._api_client, response_dict
+      )
+    else:
+      response_dict = _GenerateVideosOperation_from_mldev(
+          self._api_client, response_dict
+      )
+
+    return_value = types.GenerateVideosOperation._from_response(
+        response=response_dict, kwargs=parameter_model
+    )
+    self._api_client._verify_response(return_value)
+    return return_value
+
   def generate_content(
       self,
       *,
@@ -4601,19 +5093,19 @@ class Models(_api_module.BaseModule):
   ) -> types.GenerateContentResponse:
     """Makes an API request to generate content using a model.
 
-    For the `model` parameter, supported format for Vertex AI API includes:
-    - the Gemini model ID, for example: 'gemini-1.5-flash-002'
-    - the full resource name starts with 'projects/', for example:
+    For the `model` parameter, supported formats for Vertex AI API include:
+    - The Gemini model ID, for example: 'gemini-1.5-flash-002'
+    - The full resource name starts with 'projects/', for example:
       'projects/my-project-id/locations/us-central1/publishers/google/models/gemini-1.5-flash-002'
-    - the partial resource name with 'publishers/', for example:
+    - The partial resource name with 'publishers/', for example:
       'publishers/google/models/gemini-1.5-flash-002' or
       'publishers/meta/models/llama-3.1-405b-instruct-maas'
     - `/` separated publisher and model name, for example:
       'google/gemini-1.5-flash-002' or 'meta/llama-3.1-405b-instruct-maas'
 
-    For the `model` parameter, supported format for Gemini API includes:
-    - the Gemini model ID, for example: 'gemini-1.5-flash-002'
-    - the model name starts with 'models/', for example:
+    For the `model` parameter, supported formats for Gemini API include:
+    - The Gemini model ID, for example: 'gemini-1.5-flash-002'
+    - The model name starts with 'models/', for example:
       'models/gemini-1.5-flash-002'
     - if you would like to use a tuned model, the model name starts with
       'tunedModels/', for example:
@@ -4661,7 +5153,7 @@ class Models(_api_module.BaseModule):
           model=model, contents=contents, config=config
       )
     remaining_remote_calls_afc = _extra_utils.get_max_remote_calls_afc(config)
-    logging.info(
+    logger.info(
         f'AFC is enabled with max remote calls: {remaining_remote_calls_afc}.'
     )
     automatic_function_calling_history = []
@@ -4672,10 +5164,10 @@ class Models(_api_module.BaseModule):
       response = self._generate_content(
           model=model, contents=contents, config=config
       )
-      logging.info(f'AFC remote call {i} is done.')
+      logger.info(f'AFC remote call {i} is done.')
       remaining_remote_calls_afc -= 1
       if remaining_remote_calls_afc == 0:
-        logging.info('Reached max remote calls for automatic function calling.')
+        logger.info('Reached max remote calls for automatic function calling.')
 
       function_map = _extra_utils.get_function_map(config)
       if not function_map:
@@ -4720,21 +5212,21 @@ class Models(_api_module.BaseModule):
   ) -> Iterator[types.GenerateContentResponse]:
     """Makes an API request to generate content using a model and yields the model's response in chunks.
 
-    For the `model` parameter, supported format for Vertex AI API includes:
-    - the Gemini model ID, for example: 'gemini-1.5-flash-002'
-    - the full resource name starts with 'projects/', for example:
+    For the `model` parameter, supported formats for Vertex AI API include:
+    - The Gemini model ID, for example: 'gemini-1.5-flash-002'
+    - The full resource name starts with 'projects/', for example:
       'projects/my-project-id/locations/us-central1/publishers/google/models/gemini-1.5-flash-002'
-    - the partial resource name with 'publishers/', for example:
+    - The partial resource name with 'publishers/', for example:
       'publishers/google/models/gemini-1.5-flash-002' or
       'publishers/meta/models/llama-3.1-405b-instruct-maas'
     - `/` separated publisher and model name, for example:
       'google/gemini-1.5-flash-002' or 'meta/llama-3.1-405b-instruct-maas'
 
-    For the `model` parameter, supported format for Gemini API includes:
-    - the Gemini model ID, for example: 'gemini-1.5-flash-002'
-    - the model name starts with 'models/', for example:
+    For the `model` parameter, supported formats for Gemini API include:
+    - The Gemini model ID, for example: 'gemini-1.5-flash-002'
+    - The model name starts with 'models/', for example:
       'models/gemini-1.5-flash-002'
-    - if you would like to use a tuned model, the model name starts with
+    - If you would like to use a tuned model, the model name starts with
       'tunedModels/', for example:
       'tunedModels/1234567890123456789'
 
@@ -4782,7 +5274,7 @@ class Models(_api_module.BaseModule):
       return
 
     remaining_remote_calls_afc = _extra_utils.get_max_remote_calls_afc(config)
-    logging.info(
+    logger.info(
         f'AFC is enabled with max remote calls: {remaining_remote_calls_afc}.'
     )
     automatic_function_calling_history = []
@@ -4794,10 +5286,10 @@ class Models(_api_module.BaseModule):
       response = self._generate_content_stream(
           model=model, contents=contents, config=config
       )
-      logging.info(f'AFC remote call {i} is done.')
+      logger.info(f'AFC remote call {i} is done.')
       remaining_remote_calls_afc -= 1
       if remaining_remote_calls_afc == 0:
-        logging.info('Reached max remote calls for automatic function calling.')
+        logger.info('Reached max remote calls for automatic function calling.')
 
       function_map = _extra_utils.get_function_map(config)
 
@@ -4809,6 +5301,12 @@ class Models(_api_module.BaseModule):
           if not function_map:
             yield chunk
           else:
+            if (
+                not chunk.candidates
+                or not chunk.candidates[0].content
+                or not chunk.candidates[0].content.parts
+            ):
+              break
             func_response_parts = _extra_utils.get_function_response_parts(
                 chunk, function_map
             )
@@ -4823,20 +5321,16 @@ class Models(_api_module.BaseModule):
                 automatic_function_calling_history
             )
           yield chunk
+        if (
+            not chunk.candidates
+            or not chunk.candidates[0].content
+            or not chunk.candidates[0].content.parts
+        ):
+          break
         func_response_parts = _extra_utils.get_function_response_parts(
             chunk, function_map
         )
 
-      if not chunk:
-        break
-      if (
-          not chunk
-          or not chunk.candidates
-          or not chunk.candidates[0].content
-          or not chunk.candidates[0].content.parts
-      ):
-        break
-
       if not function_map:
         break
       if not func_response_parts:
@@ -4898,7 +5392,7 @@ class Models(_api_module.BaseModule):
 
     # Convert to API config.
     config = config or {}
-    config_dct = config if isinstance(config, dict) else config.dict()
+    config_dct = config if isinstance(config, dict) else config.model_dump()
     api_config = types._UpscaleImageAPIConfigDict(**config_dct)  # pylint: disable=protected-access
 
     # Provide default values through API config.
@@ -5095,7 +5589,7 @@ class AsyncModels(_api_module.BaseModule):
       contents: Union[types.ContentListUnion, types.ContentListUnionDict],
       config: Optional[types.EmbedContentConfigOrDict] = None,
   ) -> types.EmbedContentResponse:
-    """Calculates embeddings for the given contents(only text is supported).
+    """Calculates embeddings for the given contents. Only text is supported.
 
     Args:
       model (str): The model to use.
@@ -5106,7 +5600,7 @@ class AsyncModels(_api_module.BaseModule):
 
     .. code-block:: python
 
-      embeddings = client.models.embed_content(
+      embeddings = await client.aio.models.embed_content(
           model= 'text-embedding-004',
           contents=[
               'What is your name?',
@@ -5186,8 +5680,8 @@ class AsyncModels(_api_module.BaseModule):
 
     .. code-block:: python
 
-      response = client.models.generate_images(
-        model='imagen-3.0-generate-001',
+      response = await client.aio.models.generate_images(
+        model='imagen-3.0-generate-002',
         prompt='Man with a dog',
         config=types.GenerateImagesConfig(
             number_of_images= 1,
@@ -5284,7 +5778,7 @@ class AsyncModels(_api_module.BaseModule):
             mask_dilation=0.06,
         ),
       )
-      response = client.models.edit_image(
+      response = await client.aio.models.edit_image(
         model='imagen-3.0-capability-preview-0930',
         prompt='man with dog',
         reference_images=[raw_ref_image, mask_ref_image],
@@ -5627,17 +6121,18 @@ class AsyncModels(_api_module.BaseModule):
   ) -> types.CountTokensResponse:
     """Counts the number of tokens in the given content.
 
+    Multimodal input is supported for Gemini models.
+
     Args:
       model (str): The model to use for counting tokens.
       contents (list[types.Content]): The content to count tokens for.
-        Multimodal input is supported for Gemini models.
       config (CountTokensConfig): The configuration for counting tokens.
 
     Usage:
 
     .. code-block:: python
 
-      response = client.models.count_tokens(
+      response = await client.aio.models.count_tokens(
           model='gemini-1.5-flash',
           contents='What is your name?',
       )
@@ -5702,20 +6197,21 @@ class AsyncModels(_api_module.BaseModule):
       contents: Union[types.ContentListUnion, types.ContentListUnionDict],
       config: Optional[types.ComputeTokensConfigOrDict] = None,
   ) -> types.ComputeTokensResponse:
-    """Return a list of tokens based on the input text.
+    """Return a list of tokens based on the input contents.
+
+    Only text is supported.
 
     This method is not supported by the Gemini Developer API.
 
     Args:
       model (str): The model to use.
-      contents (list[shared.Content]): The content to compute tokens for. Only
-        text is supported.
+      contents (list[shared.Content]): The content to compute tokens for.
 
     Usage:
 
     .. code-block:: python
 
-      response = client.models.compute_tokens(
+      response = await client.aio.models.compute_tokens(
           model='gemini-1.5-flash',
           contents='What is your name?',
       )
@@ -5772,6 +6268,88 @@ class AsyncModels(_api_module.BaseModule):
     self._api_client._verify_response(return_value)
     return return_value
 
+  @_common.experimental_warning(
+      'This method is experimental and may change in future versions.'
+  )
+  async def generate_videos(
+      self,
+      *,
+      model: str,
+      prompt: Optional[str] = None,
+      config: Optional[types.GenerateVideosConfigOrDict] = None,
+  ) -> types.GenerateVideosOperation:
+    """Generates videos based on a text description and configuration.
+
+    Args:
+      model: The model to use.
+      instances: A list of prompts, images and videos to generate videos from.
+      config: Configuration for generation.
+
+    Usage:
+
+      ```
+      operation = client.models.generate_videos(
+          model="veo-2.0-generate-001",
+          prompt="A neon hologram of a cat driving at top speed",
+      )
+      while not operation.done:
+          time.sleep(10)
+          operation = client.operations.get(operation)
+
+      operation.result.generated_videos[0].video.uri
+      ```
+    """
+
+    parameter_model = types._GenerateVideosParameters(
+        model=model,
+        prompt=prompt,
+        config=config,
+    )
+
+    if self._api_client.vertexai:
+      request_dict = _GenerateVideosParameters_to_vertex(
+          self._api_client, parameter_model
+      )
+      path = '{model}:predictLongRunning'.format_map(request_dict.get('_url'))
+    else:
+      request_dict = _GenerateVideosParameters_to_mldev(
+          self._api_client, parameter_model
+      )
+      path = '{model}:predictLongRunning'.format_map(request_dict.get('_url'))
+    query_params = request_dict.get('_query')
+    if query_params:
+      path = f'{path}?{urlencode(query_params)}'
+    # TODO: remove the hack that pops config.
+    request_dict.pop('config', None)
+
+    http_options = None
+    if isinstance(config, dict):
+      http_options = config.get('http_options', None)
+    elif hasattr(config, 'http_options'):
+      http_options = config.http_options
+
+    request_dict = _common.convert_to_dict(request_dict)
+    request_dict = _common.encode_unserializable_types(request_dict)
+
+    response_dict = await self._api_client.async_request(
+        'post', path, request_dict, http_options
+    )
+
+    if self._api_client.vertexai:
+      response_dict = _GenerateVideosOperation_from_vertex(
+          self._api_client, response_dict
+      )
+    else:
+      response_dict = _GenerateVideosOperation_from_mldev(
+          self._api_client, response_dict
+      )
+
+    return_value = types.GenerateVideosOperation._from_response(
+        response=response_dict, kwargs=parameter_model
+    )
+    self._api_client._verify_response(return_value)
+    return return_value
+
   async def generate_content(
       self,
       *,
@@ -5813,7 +6391,7 @@ class AsyncModels(_api_module.BaseModule):
           model=model, contents=contents, config=config
       )
     remaining_remote_calls_afc = _extra_utils.get_max_remote_calls_afc(config)
-    logging.info(
+    logger.info(
         f'AFC is enabled with max remote calls: {remaining_remote_calls_afc}.'
     )
     automatic_function_calling_history = []
@@ -5824,7 +6402,7 @@ class AsyncModels(_api_module.BaseModule):
       )
       remaining_remote_calls_afc -= 1
       if remaining_remote_calls_afc == 0:
-        logging.info('Reached max remote calls for automatic function calling.')
+        logger.info('Reached max remote calls for automatic function calling.')
 
       function_map = _extra_utils.get_function_map(config)
       if not function_map:
@@ -5870,21 +6448,21 @@ class AsyncModels(_api_module.BaseModule):
   ) -> Awaitable[AsyncIterator[types.GenerateContentResponse]]:
     """Makes an API request to generate content using a model and yields the model's response in chunks.
 
-    For the `model` parameter, supported format for Vertex AI API includes:
-    - the Gemini model ID, for example: 'gemini-1.5-flash-002'
-    - the full resource name starts with 'projects/', for example:
+    For the `model` parameter, supported formats for Vertex AI API include:
+    - The Gemini model ID, for example: 'gemini-1.5-flash-002'
+    - The full resource name starts with 'projects/', for example:
       'projects/my-project-id/locations/us-central1/publishers/google/models/gemini-1.5-flash-002'
-    - the partial resource name with 'publishers/', for example:
+    - The partial resource name with 'publishers/', for example:
       'publishers/google/models/gemini-1.5-flash-002' or
       'publishers/meta/models/llama-3.1-405b-instruct-maas'
     - `/` separated publisher and model name, for example:
       'google/gemini-1.5-flash-002' or 'meta/llama-3.1-405b-instruct-maas'
 
-    For the `model` parameter, supported format for Gemini API includes:
-    - the Gemini model ID, for example: 'gemini-1.5-flash-002'
-    - the model name starts with 'models/', for example:
+    For the `model` parameter, supported formats for Gemini API include:
+    - The Gemini model ID, for example: 'gemini-1.5-flash-002'
+    - The model name starts with 'models/', for example:
       'models/gemini-1.5-flash-002'
-    - if you would like to use a tuned model, the model name starts with
+    - If you would like to use a tuned model, the model name starts with
       'tunedModels/', for example:
       'tunedModels/1234567890123456789'
 
@@ -5938,10 +6516,11 @@ class AsyncModels(_api_module.BaseModule):
 
     async def async_generator(model, contents, config):
       remaining_remote_calls_afc = _extra_utils.get_max_remote_calls_afc(config)
-      logging.info(
+      logger.info(
           f'AFC is enabled with max remote calls: {remaining_remote_calls_afc}.'
       )
       automatic_function_calling_history = []
+      func_response_parts = None
       chunk = None
       i = 0
       while remaining_remote_calls_afc > 0:
@@ -5949,10 +6528,10 @@ class AsyncModels(_api_module.BaseModule):
         response = await self._generate_content_stream(
             model=model, contents=contents, config=config
         )
-        logging.info(f'AFC remote call {i} is done.')
+        logger.info(f'AFC remote call {i} is done.')
         remaining_remote_calls_afc -= 1
         if remaining_remote_calls_afc == 0:
-          logging.info(
+          logger.info(
               'Reached max remote calls for automatic function calling.'
           )
 
@@ -5966,6 +6545,12 @@ class AsyncModels(_api_module.BaseModule):
             if not function_map:
               yield chunk
             else:
+              if (
+                  not chunk.candidates
+                  or not chunk.candidates[0].content
+                  or not chunk.candidates[0].content.parts
+              ):
+                break
               func_response_parts = _extra_utils.get_function_response_parts(
                   chunk, function_map
               )
@@ -5981,18 +6566,15 @@ class AsyncModels(_api_module.BaseModule):
                   automatic_function_calling_history
               )
             yield chunk
+          if (
+              not chunk.candidates
+              or not chunk.candidates[0].content
+              or not chunk.candidates[0].content.parts
+          ):
+            break
           func_response_parts = _extra_utils.get_function_response_parts(
               chunk, function_map
           )
-        if not chunk:
-          break
-        if (
-            not chunk
-            or not chunk.candidates
-            or not chunk.candidates[0].content
-            or not chunk.candidates[0].content.parts
-        ):
-          break
         if not function_map:
           break