feat: Support video extension for Veo on Gemini Developer API

PiperOrigin-RevId: 819532215

Author: matthew29tang

google/genai/_operations_converters.py

@@ -201,11 +201,11 @@ def _GeneratedVideo_from_mldev(
     parent_object: Optional[dict[str, Any]] = None,
 ) -> dict[str, Any]:
   to_object: dict[str, Any] = {}
-  if getv(from_object, ['_self']) is not None:
+  if getv(from_object, ['video']) is not None:
     setv(
         to_object,
         ['video'],
-        _Video_from_mldev(getv(from_object, ['_self']), to_object),
+        _Video_from_mldev(getv(from_object, ['video']), to_object),
     )
 
   return to_object
@@ -261,14 +261,14 @@ def _Video_from_mldev(
     parent_object: Optional[dict[str, Any]] = None,
 ) -> dict[str, Any]:
   to_object: dict[str, Any] = {}
-  if getv(from_object, ['video', 'uri']) is not None:
-    setv(to_object, ['uri'], getv(from_object, ['video', 'uri']))
+  if getv(from_object, ['uri']) is not None:
+    setv(to_object, ['uri'], getv(from_object, ['uri']))
 
-  if getv(from_object, ['video', 'encodedVideo']) is not None:
+  if getv(from_object, ['encodedVideo']) is not None:
     setv(
         to_object,
         ['video_bytes'],
-        base_t.t_bytes(getv(from_object, ['video', 'encodedVideo'])),
+        base_t.t_bytes(getv(from_object, ['encodedVideo'])),
     )
 
   if getv(from_object, ['encoding']) is not None:

google/genai/models.py

@@ -2262,11 +2262,11 @@ def _GeneratedVideo_from_mldev(
     parent_object: Optional[dict[str, Any]] = None,
 ) -> dict[str, Any]:
   to_object: dict[str, Any] = {}
-  if getv(from_object, ['_self']) is not None:
+  if getv(from_object, ['video']) is not None:
     setv(
         to_object,
         ['video'],
-        _Video_from_mldev(getv(from_object, ['_self']), to_object),
+        _Video_from_mldev(getv(from_object, ['video']), to_object),
     )
 
   return to_object
@@ -3680,14 +3680,14 @@ def _Video_from_mldev(
     parent_object: Optional[dict[str, Any]] = None,
 ) -> dict[str, Any]:
   to_object: dict[str, Any] = {}
-  if getv(from_object, ['video', 'uri']) is not None:
-    setv(to_object, ['uri'], getv(from_object, ['video', 'uri']))
+  if getv(from_object, ['uri']) is not None:
+    setv(to_object, ['uri'], getv(from_object, ['uri']))
 
-  if getv(from_object, ['video', 'encodedVideo']) is not None:
+  if getv(from_object, ['encodedVideo']) is not None:
     setv(
         to_object,
         ['video_bytes'],
-        base_t.t_bytes(getv(from_object, ['video', 'encodedVideo'])),
+        base_t.t_bytes(getv(from_object, ['encodedVideo'])),
     )
 
   if getv(from_object, ['encoding']) is not None:
@@ -3723,12 +3723,12 @@ def _Video_to_mldev(
 ) -> dict[str, Any]:
   to_object: dict[str, Any] = {}
   if getv(from_object, ['uri']) is not None:
-    setv(to_object, ['video', 'uri'], getv(from_object, ['uri']))
+    setv(to_object, ['uri'], getv(from_object, ['uri']))
 
   if getv(from_object, ['video_bytes']) is not None:
     setv(
         to_object,
-        ['video', 'encodedVideo'],
+        ['encodedVideo'],
         base_t.t_bytes(getv(from_object, ['video_bytes'])),
     )
 
@@ -5448,6 +5448,36 @@ def generate_videos(
           'Source and prompt/image/video are mutually exclusive.'
           + ' Please only use source.'
       )
+    # Gemini Developer API does not support video bytes.
+    video_dct: dict[str, Any] = {}
+    if not self._api_client.vertexai and video:
+      if isinstance(video, types.Video):
+        video_dct = video.model_dump()
+      else:
+        video_dct = dict(video)
+
+      if video_dct.get('uri') and video_dct.get('video_bytes'):
+        video = types.Video(
+            uri=video_dct.get('uri'), mime_type=video_dct.get('mime_type')
+        )
+    elif not self._api_client.vertexai and source:
+      if isinstance(source, types.GenerateVideosSource):
+        source_dct = source.model_dump()
+        video_dct = source_dct.get('video', {})
+      else:
+        source_dct = dict(source)
+        if isinstance(source_dct.get('video'), types.Video):
+          video_obj: types.Video = source_dct.get('video', types.Video())
+          video_dct = video_obj.model_dump()
+      if video_dct and video_dct.get('uri') and video_dct.get('video_bytes'):
+        source = types.GenerateVideosSource(
+            prompt=source_dct.get('prompt'),
+            image=source_dct.get('image'),
+            video=types.Video(
+                uri=video_dct.get('uri'),
+                mime_type=video_dct.get('mime_type'),
+            ),
+        )
     return self._generate_videos(
         model=model,
         prompt=prompt,

google/genai/tests/models/test_generate_videos.py

@@ -310,7 +310,7 @@ def test_text_and_image_to_video_poll(client):
 
 
 def test_video_to_video_poll(client):
-  # Video extension is only supported in Vertex AI.
+  # GCS URI video input is only supported in Vertex AI.
   if not client.vertexai:
     return
 
@@ -333,7 +333,7 @@ def test_video_to_video_poll(client):
 
 
 def test_text_and_video_to_video_poll(client):
-  # Video extension is only supported in Vertex AI.
+  # GCS URI video input is only supported in Vertex AI.
   if not client.vertexai:
     return
 
@@ -356,6 +356,139 @@ def test_text_and_video_to_video_poll(client):
   assert operation.result.generated_videos[0].video.uri
 
 
+def test_generated_video_extension_poll(client):
+  # Gemini API only supports video extension on generated videos.
+  if client.vertexai:
+    return
+
+  operation1 = client.models.generate_videos(
+      model="veo-3-exp",
+      prompt="Rain",
+      config=types.GenerateVideosConfig(
+          number_of_videos=1,
+      ),
+  )
+  while not operation1.done:
+    # Skip the sleep when in replay mode.
+    if client._api_client._mode not in ("replay", "auto"):
+      time.sleep(20)
+    operation1 = client.operations.get(operation=operation1)
+
+  video1 = operation1.result.generated_videos[0].video
+  assert video1.uri
+  client.files.download(file=video1)
+  assert video1.video_bytes
+
+  operation2 = client.models.generate_videos(
+      model="veo-3-exp",
+      prompt="Sun",
+      video=video1,
+      config=types.GenerateVideosConfig(
+          number_of_videos=1,
+      ),
+  )
+  while not operation2.done:
+    # Skip the sleep when in replay mode.
+    if client._api_client._mode not in ("replay", "auto"):
+      time.sleep(20)
+    operation2 = client.operations.get(operation=operation2)
+
+  video2 = operation2.result.generated_videos[0].video
+  assert video2.uri
+  client.files.download(file=video2)
+  assert video2.video_bytes
+
+
+def test_generated_video_extension_from_source_poll(client):
+  # Gemini API only supports video extension on generated videos.
+  if client.vertexai:
+    return
+
+  operation1 = client.models.generate_videos(
+      model="veo-3-exp",
+      prompt="Rain",
+      config=types.GenerateVideosConfig(
+          number_of_videos=1,
+      ),
+  )
+  while not operation1.done:
+    # Skip the sleep when in replay mode.
+    if client._api_client._mode not in ("replay", "auto"):
+      time.sleep(20)
+    operation1 = client.operations.get(operation=operation1)
+
+  video1 = operation1.result.generated_videos[0].video
+  assert video1.uri
+  client.files.download(file=video1)
+  assert video1.video_bytes
+
+  operation2 = client.models.generate_videos(
+      model="veo-3-exp",
+      source=types.GenerateVideosSource(
+          prompt="Sun",
+          video=video1
+      ),
+      config=types.GenerateVideosConfig(
+          number_of_videos=1,
+      ),
+  )
+  while not operation2.done:
+    # Skip the sleep when in replay mode.
+    if client._api_client._mode not in ("replay", "auto"):
+      time.sleep(20)
+    operation2 = client.operations.get(operation=operation2)
+
+  video2 = operation2.result.generated_videos[0].video
+  assert video2.uri
+  client.files.download(file=video2)
+  assert video2.video_bytes
+
+
+def test_generated_video_extension_from_source_dict_poll(client):
+  # Gemini API only supports video extension on generated videos.
+  if client.vertexai:
+    return
+
+  operation1 = client.models.generate_videos(
+      model="veo-3-exp",
+      prompt="Rain",
+      config=types.GenerateVideosConfig(
+          number_of_videos=1,
+      ),
+  )
+  while not operation1.done:
+    # Skip the sleep when in replay mode.
+    if client._api_client._mode not in ("replay", "auto"):
+      time.sleep(20)
+    operation1 = client.operations.get(operation=operation1)
+
+  video1 = operation1.result.generated_videos[0].video
+  assert video1.uri
+  client.files.download(file=video1)
+  assert video1.video_bytes
+
+  operation2 = client.models.generate_videos(
+      model="veo-3-exp",
+      source={
+          "prompt": "Sun",
+          "video": video1,
+      },
+      config=types.GenerateVideosConfig(
+          number_of_videos=1,
+      ),
+  )
+  while not operation2.done:
+    # Skip the sleep when in replay mode.
+    if client._api_client._mode not in ("replay", "auto"):
+      time.sleep(20)
+    operation2 = client.operations.get(operation=operation2)
+
+  video2 = operation2.result.generated_videos[0].video
+  assert video2.uri
+  client.files.download(file=video2)
+  assert video2.video_bytes
+
+
 def test_image_to_video_frame_interpolation_poll(client):
   operation = client.models.generate_videos(
       model="veo-2.0-generate-exp" if client.vertexai else "veo-3-exp",