change default poll interval (#76), rework veo2

comfy_api_nodes/apis/client.py

@@ -496,7 +496,7 @@ class PollingOperation(Generic[T, R]):
         request: Optional[T] = None,
         api_base: str | None = None,
         auth_token: Optional[str] = None,
-        poll_interval: float = 1.0,
+        poll_interval: float = 5.0,
     ):
         self.poll_endpoint = poll_endpoint
         self.request = request

comfy_api_nodes/nodes_veo2.py

@@ -2,13 +2,9 @@ import io
 import logging
 import base64
 import requests
-import math
 import torch
-import numpy as np
-from PIL import Image
 
 from comfy.comfy_types.node_typing import IO, ComfyNodeABC
-from comfy.utils import common_upscale
 from comfy_api.input_impl.video_types import VideoFromFile
 from comfy_api_nodes.apis import (
     Veo2GenVidRequest,
@@ -20,21 +16,20 @@ from comfy_api_nodes.apis.client import (
     ApiEndpoint,
     HttpMethod,
     SynchronousOperation,
+    PollingOperation,
 )
 
-def downscale_input(image, total_pixels=1536*1024):
+from comfy_api_nodes.apinode_utils import (
-    samples = image.movedim(-1,1)
+    downscale_image_tensor,
-    # Downscaling input images to roughly the same size as the outputs
+    tensor_to_base64_string
-    total = int(total_pixels)
+)
-    scale_by = math.sqrt(total / (samples.shape[3] * samples.shape[2]))
-    if scale_by >= 1:
-        return image
-    width = round(samples.shape[3] * scale_by)
-    height = round(samples.shape[2] * scale_by)
 
-    s = common_upscale(samples, width, height, "lanczos", "disabled")
+def convert_image_to_base64(image: torch.Tensor):
-    s = s.movedim(1,-1)
+    if image is None:
-    return s
+        return None
+
+    scaled_image = downscale_image_tensor(image, total_pixels=2048*2048)
+    return tensor_to_base64_string(scaled_image)
 
 class VeoVideoGenerationNode(ComfyNodeABC):
     """
@@ -128,25 +123,6 @@ class VeoVideoGenerationNode(ComfyNodeABC):
     DESCRIPTION = "Generates videos from text prompts using Google's Veo API"
     API_NODE = True
 
-    def _convert_image_to_base64(self, image: torch.Tensor):
-        if image is None:
-            return None
-
-        scaled_image = downscale_input(image, total_pixels=2048*2048)
-
-        # Remove batch dimension if present
-        if len(scaled_image.shape) > 3:
-            scaled_image = scaled_image[0]
-
-        # Convert to numpy array and then to PIL Image
-        image_np = (scaled_image.numpy() * 255).astype(np.uint8)
-        img = Image.fromarray(image_np)
-
-        # Convert to base64
-        buffer = io.BytesIO()
-        img.save(buffer, format="PNG")
-        return base64.b64encode(buffer.getvalue()).decode('utf-8')
-
     def generate_video(
         self,
         prompt,
@@ -168,7 +144,7 @@ class VeoVideoGenerationNode(ComfyNodeABC):
 
         # Add image if provided
         if image is not None:
-            image_base64 = self._convert_image_to_base64(image)
+            image_base64 = convert_image_to_base64(image)
             if image_base64:
                 instance["image"] = {
                     "bytesBase64Encoded": image_base64,
@@ -211,67 +187,79 @@ class VeoVideoGenerationNode(ComfyNodeABC):
 
         logging.info(f"Veo generation started with operation name: {operation_name}")
 
-        # Poll until operation is complete
+        # Define status extractor function
+        def status_extractor(response):
+            # Only return "completed" if the operation is done, regardless of success or failure
+            # We'll check for errors after polling completes
+            return "completed" if response.done else "pending"
+
+        # Define progress extractor function
+        def progress_extractor(response):
+            # Could be enhanced if the API provides progress information
+            return None
+
+        # Define the polling operation
+        poll_operation = PollingOperation(
+            poll_endpoint=ApiEndpoint(
+                path="/proxy/veo/poll",
+                method=HttpMethod.POST,
+                request_model=Veo2GenVidPollRequest,
+                response_model=Veo2GenVidPollResponse
+            ),
+            completed_statuses=["completed"],
+            failed_statuses=[],  # No failed statuses, we'll handle errors after polling
+            status_extractor=status_extractor,
+            progress_extractor=progress_extractor,
+            request=Veo2GenVidPollRequest(
+                operationName=operation_name
+            ),
+            auth_token=auth_token,
+            poll_interval=5.0
+        )
+
+        # Execute the polling operation
+        poll_response = poll_operation.execute()
+
+        # Now check for errors in the final response
+        # Check for error in poll response
+        if hasattr(poll_response, 'error') and poll_response.error:
+            error_message = f"Veo API error: {poll_response.error.message} (code: {poll_response.error.code})"
+            logging.error(error_message)
+            raise Exception(error_message)
+
+        # Check for RAI filtered content
+        if (hasattr(poll_response.response, 'raiMediaFilteredCount') and
+            poll_response.response.raiMediaFilteredCount > 0):
+
+            # Extract reason message if available
+            if (hasattr(poll_response.response, 'raiMediaFilteredReasons') and
+                poll_response.response.raiMediaFilteredReasons):
+                reason = poll_response.response.raiMediaFilteredReasons[0]
+                error_message = f"Content filtered by Google's Responsible AI practices: {reason} ({poll_response.response.raiMediaFilteredCount} videos filtered.)"
+            else:
+                error_message = f"Content filtered by Google's Responsible AI practices ({poll_response.response.raiMediaFilteredCount} videos filtered.)"
+
+            logging.error(error_message)
+            raise Exception(error_message)
+
+        # Extract video data
         video_data = None
-        while True:
+        if poll_response.response and hasattr(poll_response.response, 'videos') and poll_response.response.videos and len(poll_response.response.videos) > 0:
-            poll_operation = SynchronousOperation(
+            video = poll_response.response.videos[0]
-                endpoint=ApiEndpoint(
-                    path="/proxy/veo/poll",
-                    method=HttpMethod.POST,
-                    request_model=Veo2GenVidPollRequest,
-                    response_model=Veo2GenVidPollResponse
-                ),
-                request=Veo2GenVidPollRequest(
-                    operationName=operation_name
-                ),
-                auth_token=auth_token
-            )
 
-            poll_response = poll_operation.execute()
+            # Check if video is provided as base64 or URL
-
+            if hasattr(video, 'bytesBase64Encoded') and video.bytesBase64Encoded:
-            # Check for error in poll response
+                # Decode base64 string to bytes
-            if hasattr(poll_response, 'error') and poll_response.error:
+                video_data = base64.b64decode(video.bytesBase64Encoded)
-                error_message = f"Veo API error: {poll_response.error.message} (code: {poll_response.error.code})"
+            elif hasattr(video, 'gcsUri') and video.gcsUri:
-                logging.error(error_message)
+                # Download from URL
-                raise Exception(error_message)
+                video_url = video.gcsUri
-
+                video_response = requests.get(video_url)
-            if poll_response.done:
+                video_data = video_response.content
-                # Check for RAI filtered content
+            else:
-                if (hasattr(poll_response.response, 'raiMediaFilteredCount') and
+                raise Exception("Video returned but no data or URL was provided")
-                    poll_response.response.raiMediaFilteredCount > 0):
+        else:
-
+            raise Exception("Video generation completed but no video was returned")
-                    # Extract reason message if available
-                    if (hasattr(poll_response.response, 'raiMediaFilteredReasons') and
-                        poll_response.response.raiMediaFilteredReasons):
-                        reason = poll_response.response.raiMediaFilteredReasons[0]
-                        error_message = f"Content filtered by Google's Responsible AI practices: {reason} ({poll_response.response.raiMediaFilteredCount} videos filtered.)"
-
-                    logging.error(error_message)
-                    raise Exception(error_message)
-
-                # Process successful response
-                if poll_response.response and hasattr(poll_response.response, 'videos') and poll_response.response.videos and len(poll_response.response.videos) > 0:
-                    video = poll_response.response.videos[0]
-
-                    # Check if video is provided as base64 or URL
-                    if hasattr(video, 'bytesBase64Encoded') and video.bytesBase64Encoded:
-                        # Decode base64 string to bytes
-                        video_data = base64.b64decode(video.bytesBase64Encoded)
-                        break
-                    elif hasattr(video, 'gcsUri') and video.gcsUri:
-                        # Download from URL
-                        video_url = video.gcsUri
-                        video_response = requests.get(video_url)
-                        video_data = video_response.content
-                        break
-                    else:
-                        raise Exception("Video returned but no data or URL was provided")
-                else:
-                    raise Exception("Video generation completed but no video was returned")
-
-            # Wait before polling again
-            import time
-            time.sleep(5)
 
         if not video_data:
             raise Exception("No video data was returned")