V 2.0.0 - Universal api #66 - VEO

components/API/references/Gemini_VEO.py

@@ -0,0 +1,20 @@
+from __future__ import annotations
+
+import io
+from typing import Any
+import numpy as np
+from PIL import Image
+from google.genai import types
+
+
+def handle_reference_images(img_binary_api: Any = None, temp_file_ref: str = "", loaded_client_for_upload: Any = None, **_: Any):
+    if img_binary_api is None:
+        return {}
+    if type(img_binary_api).__name__ == "Tensor":
+        img_array = (img_binary_api[0].cpu().numpy() * 255).astype(np.uint8)
+        pil_image = Image.fromarray(img_array)
+        buffer = io.BytesIO()
+        pil_image.save(buffer, format="PNG")
+        image_bytes = buffer.getvalue()
+        return {"first_image": types.Image(image_bytes=image_bytes, mime_type="image/png")}
+    return {}

components/API/responses/Gemini_VEO.py

@@ -0,0 +1,30 @@
+from __future__ import annotations
+
+import time
+import os
+import random
+import folder_paths
+
+
+def handle_response(api_result, schema=None, loaded_client=None, response_url=None, client=None, sdk_context=None):
+    if api_result is None:
+        return [None, None]
+
+    operation = api_result
+    while not operation.done:
+        time.sleep(10)
+        operation = client.operations.get(operation)
+
+    generated_videos = getattr(getattr(operation, 'response', None), 'generated_videos', None)
+    if not generated_videos:
+        return [None, None]
+
+    generated_video = generated_videos[0]
+    tmp_path = os.path.join(folder_paths.temp_directory, f"veo_{random.randint(10000, 99999)}.mp4")
+    client.files.download(file=generated_video.video)
+    generated_video.video.save(tmp_path)
+
+    with open(tmp_path, 'rb') as f:
+        video_bytes = f.read()
+
+    return ["video_result", video_bytes]