mirror of
https://github.com/comfyanonymous/ComfyUI.git
synced 2025-09-10 11:35:40 +00:00
eb39019daa
* convert Google Veo API node to the V3 schema * use own full io.Schema for Veo3VideoGenerationNode * fixed typo * use auth_kwargs instead of auth_token/comfy_api_key
422 lines
15 KiB
Python
422 lines
15 KiB
Python
import logging
|
|
import base64
|
|
import aiohttp
|
|
import torch
|
|
from io import BytesIO
|
|
from typing import Optional
|
|
from typing_extensions import override
|
|
|
|
from comfy_api.latest import ComfyExtension, io as comfy_io
|
|
from comfy_api.input_impl.video_types import VideoFromFile
|
|
from comfy_api_nodes.apis import (
|
|
VeoGenVidRequest,
|
|
VeoGenVidResponse,
|
|
VeoGenVidPollRequest,
|
|
VeoGenVidPollResponse,
|
|
)
|
|
from comfy_api_nodes.apis.client import (
|
|
ApiEndpoint,
|
|
HttpMethod,
|
|
SynchronousOperation,
|
|
PollingOperation,
|
|
)
|
|
|
|
from comfy_api_nodes.apinode_utils import (
|
|
downscale_image_tensor,
|
|
tensor_to_base64_string,
|
|
)
|
|
|
|
AVERAGE_DURATION_VIDEO_GEN = 32
|
|
|
|
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)
|
|
|
|
|
|
def get_video_url_from_response(poll_response: VeoGenVidPollResponse) -> Optional[str]:
|
|
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]
|
|
else:
|
|
return None
|
|
if hasattr(video, "gcsUri") and video.gcsUri:
|
|
return str(video.gcsUri)
|
|
return None
|
|
|
|
|
|
class VeoVideoGenerationNode(comfy_io.ComfyNode):
|
|
"""
|
|
Generates videos from text prompts using Google's Veo API.
|
|
|
|
This node can create videos from text descriptions and optional image inputs,
|
|
with control over parameters like aspect ratio, duration, and more.
|
|
"""
|
|
|
|
@classmethod
|
|
def define_schema(cls):
|
|
return comfy_io.Schema(
|
|
node_id="VeoVideoGenerationNode",
|
|
display_name="Google Veo 2 Video Generation",
|
|
category="api node/video/Veo",
|
|
description="Generates videos from text prompts using Google's Veo 2 API",
|
|
inputs=[
|
|
comfy_io.String.Input(
|
|
"prompt",
|
|
multiline=True,
|
|
default="",
|
|
tooltip="Text description of the video",
|
|
),
|
|
comfy_io.Combo.Input(
|
|
"aspect_ratio",
|
|
options=["16:9", "9:16"],
|
|
default="16:9",
|
|
tooltip="Aspect ratio of the output video",
|
|
),
|
|
comfy_io.String.Input(
|
|
"negative_prompt",
|
|
multiline=True,
|
|
default="",
|
|
tooltip="Negative text prompt to guide what to avoid in the video",
|
|
optional=True,
|
|
),
|
|
comfy_io.Int.Input(
|
|
"duration_seconds",
|
|
default=5,
|
|
min=5,
|
|
max=8,
|
|
step=1,
|
|
display_mode=comfy_io.NumberDisplay.number,
|
|
tooltip="Duration of the output video in seconds",
|
|
optional=True,
|
|
),
|
|
comfy_io.Boolean.Input(
|
|
"enhance_prompt",
|
|
default=True,
|
|
tooltip="Whether to enhance the prompt with AI assistance",
|
|
optional=True,
|
|
),
|
|
comfy_io.Combo.Input(
|
|
"person_generation",
|
|
options=["ALLOW", "BLOCK"],
|
|
default="ALLOW",
|
|
tooltip="Whether to allow generating people in the video",
|
|
optional=True,
|
|
),
|
|
comfy_io.Int.Input(
|
|
"seed",
|
|
default=0,
|
|
min=0,
|
|
max=0xFFFFFFFF,
|
|
step=1,
|
|
display_mode=comfy_io.NumberDisplay.number,
|
|
control_after_generate=True,
|
|
tooltip="Seed for video generation (0 for random)",
|
|
optional=True,
|
|
),
|
|
comfy_io.Image.Input(
|
|
"image",
|
|
tooltip="Optional reference image to guide video generation",
|
|
optional=True,
|
|
),
|
|
comfy_io.Combo.Input(
|
|
"model",
|
|
options=["veo-2.0-generate-001"],
|
|
default="veo-2.0-generate-001",
|
|
tooltip="Veo 2 model to use for video generation",
|
|
optional=True,
|
|
),
|
|
],
|
|
outputs=[
|
|
comfy_io.Video.Output(),
|
|
],
|
|
hidden=[
|
|
comfy_io.Hidden.auth_token_comfy_org,
|
|
comfy_io.Hidden.api_key_comfy_org,
|
|
comfy_io.Hidden.unique_id,
|
|
],
|
|
is_api_node=True,
|
|
)
|
|
|
|
@classmethod
|
|
async def execute(
|
|
cls,
|
|
prompt,
|
|
aspect_ratio="16:9",
|
|
negative_prompt="",
|
|
duration_seconds=5,
|
|
enhance_prompt=True,
|
|
person_generation="ALLOW",
|
|
seed=0,
|
|
image=None,
|
|
model="veo-2.0-generate-001",
|
|
generate_audio=False,
|
|
):
|
|
# Prepare the instances for the request
|
|
instances = []
|
|
|
|
instance = {
|
|
"prompt": prompt
|
|
}
|
|
|
|
# Add image if provided
|
|
if image is not None:
|
|
image_base64 = convert_image_to_base64(image)
|
|
if image_base64:
|
|
instance["image"] = {
|
|
"bytesBase64Encoded": image_base64,
|
|
"mimeType": "image/png"
|
|
}
|
|
|
|
instances.append(instance)
|
|
|
|
# Create parameters dictionary
|
|
parameters = {
|
|
"aspectRatio": aspect_ratio,
|
|
"personGeneration": person_generation,
|
|
"durationSeconds": duration_seconds,
|
|
"enhancePrompt": enhance_prompt,
|
|
}
|
|
|
|
# Add optional parameters if provided
|
|
if negative_prompt:
|
|
parameters["negativePrompt"] = negative_prompt
|
|
if seed > 0:
|
|
parameters["seed"] = seed
|
|
# Only add generateAudio for Veo 3 models
|
|
if "veo-3.0" in model:
|
|
parameters["generateAudio"] = generate_audio
|
|
|
|
auth = {
|
|
"auth_token": cls.hidden.auth_token_comfy_org,
|
|
"comfy_api_key": cls.hidden.api_key_comfy_org,
|
|
}
|
|
# Initial request to start video generation
|
|
initial_operation = SynchronousOperation(
|
|
endpoint=ApiEndpoint(
|
|
path=f"/proxy/veo/{model}/generate",
|
|
method=HttpMethod.POST,
|
|
request_model=VeoGenVidRequest,
|
|
response_model=VeoGenVidResponse
|
|
),
|
|
request=VeoGenVidRequest(
|
|
instances=instances,
|
|
parameters=parameters
|
|
),
|
|
auth_kwargs=auth,
|
|
)
|
|
|
|
initial_response = await initial_operation.execute()
|
|
operation_name = initial_response.name
|
|
|
|
logging.info(f"Veo generation started with operation name: {operation_name}")
|
|
|
|
# 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=f"/proxy/veo/{model}/poll",
|
|
method=HttpMethod.POST,
|
|
request_model=VeoGenVidPollRequest,
|
|
response_model=VeoGenVidPollResponse
|
|
),
|
|
completed_statuses=["completed"],
|
|
failed_statuses=[], # No failed statuses, we'll handle errors after polling
|
|
status_extractor=status_extractor,
|
|
progress_extractor=progress_extractor,
|
|
request=VeoGenVidPollRequest(
|
|
operationName=operation_name
|
|
),
|
|
auth_kwargs=auth,
|
|
poll_interval=5.0,
|
|
result_url_extractor=get_video_url_from_response,
|
|
node_id=cls.hidden.unique_id,
|
|
estimated_duration=AVERAGE_DURATION_VIDEO_GEN,
|
|
)
|
|
|
|
# Execute the polling operation
|
|
poll_response = await 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
|
|
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)
|
|
elif hasattr(video, 'gcsUri') and video.gcsUri:
|
|
# Download from URL
|
|
async with aiohttp.ClientSession() as session:
|
|
async with session.get(video.gcsUri) as video_response:
|
|
video_data = await video_response.content.read()
|
|
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")
|
|
|
|
logging.info("Video generation completed successfully")
|
|
|
|
# Convert video data to BytesIO object
|
|
video_io = BytesIO(video_data)
|
|
|
|
# Return VideoFromFile object
|
|
return comfy_io.NodeOutput(VideoFromFile(video_io))
|
|
|
|
|
|
class Veo3VideoGenerationNode(VeoVideoGenerationNode):
|
|
"""
|
|
Generates videos from text prompts using Google's Veo 3 API.
|
|
|
|
Supported models:
|
|
- veo-3.0-generate-001
|
|
- veo-3.0-fast-generate-001
|
|
|
|
This node extends the base Veo node with Veo 3 specific features including
|
|
audio generation and fixed 8-second duration.
|
|
"""
|
|
|
|
@classmethod
|
|
def define_schema(cls):
|
|
return comfy_io.Schema(
|
|
node_id="Veo3VideoGenerationNode",
|
|
display_name="Google Veo 3 Video Generation",
|
|
category="api node/video/Veo",
|
|
description="Generates videos from text prompts using Google's Veo 3 API",
|
|
inputs=[
|
|
comfy_io.String.Input(
|
|
"prompt",
|
|
multiline=True,
|
|
default="",
|
|
tooltip="Text description of the video",
|
|
),
|
|
comfy_io.Combo.Input(
|
|
"aspect_ratio",
|
|
options=["16:9", "9:16"],
|
|
default="16:9",
|
|
tooltip="Aspect ratio of the output video",
|
|
),
|
|
comfy_io.String.Input(
|
|
"negative_prompt",
|
|
multiline=True,
|
|
default="",
|
|
tooltip="Negative text prompt to guide what to avoid in the video",
|
|
optional=True,
|
|
),
|
|
comfy_io.Int.Input(
|
|
"duration_seconds",
|
|
default=8,
|
|
min=8,
|
|
max=8,
|
|
step=1,
|
|
display_mode=comfy_io.NumberDisplay.number,
|
|
tooltip="Duration of the output video in seconds (Veo 3 only supports 8 seconds)",
|
|
optional=True,
|
|
),
|
|
comfy_io.Boolean.Input(
|
|
"enhance_prompt",
|
|
default=True,
|
|
tooltip="Whether to enhance the prompt with AI assistance",
|
|
optional=True,
|
|
),
|
|
comfy_io.Combo.Input(
|
|
"person_generation",
|
|
options=["ALLOW", "BLOCK"],
|
|
default="ALLOW",
|
|
tooltip="Whether to allow generating people in the video",
|
|
optional=True,
|
|
),
|
|
comfy_io.Int.Input(
|
|
"seed",
|
|
default=0,
|
|
min=0,
|
|
max=0xFFFFFFFF,
|
|
step=1,
|
|
display_mode=comfy_io.NumberDisplay.number,
|
|
control_after_generate=True,
|
|
tooltip="Seed for video generation (0 for random)",
|
|
optional=True,
|
|
),
|
|
comfy_io.Image.Input(
|
|
"image",
|
|
tooltip="Optional reference image to guide video generation",
|
|
optional=True,
|
|
),
|
|
comfy_io.Combo.Input(
|
|
"model",
|
|
options=["veo-3.0-generate-001", "veo-3.0-fast-generate-001"],
|
|
default="veo-3.0-generate-001",
|
|
tooltip="Veo 3 model to use for video generation",
|
|
optional=True,
|
|
),
|
|
comfy_io.Boolean.Input(
|
|
"generate_audio",
|
|
default=False,
|
|
tooltip="Generate audio for the video. Supported by all Veo 3 models.",
|
|
optional=True,
|
|
),
|
|
],
|
|
outputs=[
|
|
comfy_io.Video.Output(),
|
|
],
|
|
hidden=[
|
|
comfy_io.Hidden.auth_token_comfy_org,
|
|
comfy_io.Hidden.api_key_comfy_org,
|
|
comfy_io.Hidden.unique_id,
|
|
],
|
|
is_api_node=True,
|
|
)
|
|
|
|
|
|
class VeoExtension(ComfyExtension):
|
|
@override
|
|
async def get_node_list(self) -> list[type[comfy_io.ComfyNode]]:
|
|
return [
|
|
VeoVideoGenerationNode,
|
|
Veo3VideoGenerationNode,
|
|
]
|
|
|
|
async def comfy_entrypoint() -> VeoExtension:
|
|
return VeoExtension()
|