From 3404de63a8b0edd589e1235fc3e35bb6e950958c Mon Sep 17 00:00:00 2001 From: thot experiment <94414189+thot-experiment@users.noreply.github.com> Date: Wed, 30 Apr 2025 13:14:02 -0700 Subject: [PATCH] change default poll interval (#76), rework veo2 --- comfy_api_nodes/apis/client.py | 2 +- comfy_api_nodes/nodes_veo2.py | 178 +++++++++++++++------------------ 2 files changed, 84 insertions(+), 96 deletions(-) diff --git a/comfy_api_nodes/apis/client.py b/comfy_api_nodes/apis/client.py index bbb70519a..22e011f47 100644 --- a/comfy_api_nodes/apis/client.py +++ b/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 diff --git a/comfy_api_nodes/nodes_veo2.py b/comfy_api_nodes/nodes_veo2.py index ae3329b8d..9233944b5 100644 --- a/comfy_api_nodes/nodes_veo2.py +++ b/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): - samples = image.movedim(-1,1) - # Downscaling input images to roughly the same size as the outputs - 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) +from comfy_api_nodes.apinode_utils import ( + downscale_image_tensor, + tensor_to_base64_string +) - s = common_upscale(samples, width, height, "lanczos", "disabled") - s = s.movedim(1,-1) - return s +def convert_image_to_base64(image: torch.Tensor): + if image is None: + 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: - poll_operation = SynchronousOperation( - endpoint=ApiEndpoint( - path="/proxy/veo/poll", - method=HttpMethod.POST, - request_model=Veo2GenVidPollRequest, - response_model=Veo2GenVidPollResponse - ), - request=Veo2GenVidPollRequest( - operationName=operation_name - ), - auth_token=auth_token - ) + 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] - poll_response = poll_operation.execute() - - # 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) - - if poll_response.done: - # 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.)" - - 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) + # 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) + 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 + else: + raise Exception("Video returned but no data or URL was provided") + else: + raise Exception("Video generation completed but no video was returned") if not video_data: raise Exception("No video data was returned")