nodes/Aryan185-ComfyUI-ExternalAPI-Helpers
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Streamlined Veo nodes

Browse Source
This commit is contained in:
Aryan185
2026-01-02 11:50:02 +00:00
parent c7c7ba3467
commit 3e0ecfc40d
2 changed files with 130 additions and 339 deletions
Download Patch File Download Diff File Expand all files Collapse all files
veo.py
+81 -217
View File
@@ -1,18 +1,13 @@
import time
import os
import io
import json
import tempfile
import torch
import numpy as np
import av
from PIL import Image
import uuid
from typing import Optional
from google import genai
from google.genai import types
import cv2
import av
class VeoVertexVideoGenerator:
@classmethod
@@ -36,13 +31,9 @@ class VeoVertexVideoGenerator:
], {"default": "us-central1"}),
"service_account": ("STRING", {"multiline": True, "default": ""}),
"model": ([
"veo-2.0-generate-001",
"veo-2.0-generate-exp",
"veo-2.0-generate-preview",
"veo-3.0-generate-001",
"veo-3.0-fast-generate-001",
"veo-3.1-generate-001",
"veo-3.1-fast-generate-001"
"veo-2.0-generate-001", "veo-2.0-generate-exp", "veo-2.0-generate-preview",
"veo-3.0-generate-001", "veo-3.0-fast-generate-001",
"veo-3.1-generate-001", "veo-3.1-fast-generate-001"
], {"default": "veo-3.0-generate-001"}),
"resolution": (["720p", "1080p"], {"default": "720p"}),
"aspect_ratio": (["16:9", "9:16"], {"default": "16:9"}),
@@ -64,215 +55,88 @@ class VeoVertexVideoGenerator:
CATEGORY = "video/generation"
OUTPUT_IS_LIST = (True, False)
def setup_client(self, service_account_json, project_id, location):
"""Setup Vertex AI client with service account JSON content"""
if not service_account_json.strip():
raise ValueError("Service account JSON content is required.")
def generate_video(self, prompt, project_id, location, service_account, model, resolution, aspect_ratio,
duration_seconds, seed, generate_audio, fps, negative_prompt=None,
first_frame=None, last_frame=None):
if not project_id.strip():
raise ValueError("Project ID is required.")
# Validate and write JSON content to temporary file
creds_file = tempfile.NamedTemporaryFile(mode='w', suffix='.json', delete=False)
creds_file.write(service_account.strip())
creds_file.close()
os.environ['GOOGLE_APPLICATION_CREDENTIALS'] = creds_file.name
try:
json.loads(service_account_json) # Validate JSON format
except json.JSONDecodeError as e:
raise ValueError(f"Invalid JSON content: {str(e)}")
# Create temporary file with JSON content
temp_file = tempfile.NamedTemporaryFile(mode='w', suffix='.json', delete=False)
temp_file.write(service_account_json.strip())
temp_file.close()
# Set credentials path
os.environ['GOOGLE_APPLICATION_CREDENTIALS'] = temp_file.name
return genai.Client(
vertexai=True,
project=project_id.strip(),
location=location.strip()
)
def pil_to_tensor(self, pil_image):
if pil_image.mode != 'RGB':
pil_image = pil_image.convert('RGB')
numpy_image = np.array(pil_image).astype(np.float32) / 255.0
return torch.from_numpy(numpy_image).unsqueeze(0)
def video_to_frames(self, video_bytes):
temp_video_path = os.path.join(tempfile.gettempdir(), f"temp_video_{uuid.uuid4().hex}.mp4")
try:
with open(temp_video_path, 'wb') as f:
f.write(video_bytes)
client = genai.Client(vertexai=True, project=project_id, location=location)
cap = cv2.VideoCapture(temp_video_path)
config = types.GenerateVideosConfig(
resolution=resolution,
aspect_ratio=aspect_ratio,
duration_seconds=duration_seconds,
generate_audio=generate_audio,
fps=int(fps),
seed=seed if seed != -1 else None,
negative_prompt=negative_prompt.strip() if negative_prompt else None
)
# Helper: Tensor -> Bytes
def tensor_to_bytes(t):
arr = (t.cpu().numpy()[0] * 255).astype(np.uint8) if len(t.shape) == 4 else (t.cpu().numpy() * 255).astype(np.uint8)
b = io.BytesIO()
Image.fromarray(arr).save(b, format="PNG")
return b.getvalue()
gen_kwargs = {"model": model, "prompt": prompt, "config": config}
if first_frame is not None:
gen_kwargs["image"] = types.Image(image_bytes=tensor_to_bytes(first_frame), mime_type="image/png")
if last_frame is not None:
setattr(config, 'last_frame', types.Image(image_bytes=tensor_to_bytes(last_frame), mime_type="image/png"))
op = client.models.generate_videos(**gen_kwargs)
print(f"Veo Operation: {op.name}")
while not op.done:
time.sleep(5)
op = client.operations.get(op)
if op.error: raise Exception(f"Veo Error: {op.error}")
if not op.result.generated_videos: raise Exception("No videos generated")
video_bytes = io.BytesIO(op.result.generated_videos[0].video.video_bytes)
container = av.open(video_bytes)
frames = []
while True:
ret, frame = cap.read()
if not ret:
break
frame_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
tensor_frame = self.pil_to_tensor(Image.fromarray(frame_rgb))
frames.append(tensor_frame)
cap.release()
if not frames:
raise ValueError("No frames extracted from video")
return torch.cat(frames, dim=0)
finally:
if os.path.exists(temp_video_path):
os.remove(temp_video_path)
def extract_audio_from_video(self, video_bytes):
temp_video_path = os.path.join(tempfile.gettempdir(), f"temp_video_audio_{uuid.uuid4().hex}.mp4")
try:
with open(temp_video_path, 'wb') as f:
f.write(video_bytes)
container = av.open(temp_video_path)
if not container.streams.audio:
print("Warning: No audio track found in video.")
container.close()
return None
audio_stream = container.streams.audio[0]
audio_frames = [frame.to_ndarray() for frame in container.decode(audio=0)]
for frame in container.decode(video=0):
img = frame.to_rgb().to_ndarray().astype(np.float32) / 255.0
frames.append(torch.from_numpy(img).unsqueeze(0))
container.close()
# Decode Audio
audio = None
if generate_audio:
video_bytes.seek(0)
container = av.open(video_bytes)
if container.streams.audio:
audio_data = [f.to_ndarray() for f in container.decode(audio=0)]
if audio_data:
waveform = torch.from_numpy(np.concatenate(audio_data, axis=1)).float()
# Normalize 16/32-bit audio
if audio_data[0].dtype == np.int16: waveform /= 32768.0
elif audio_data[0].dtype == np.int32: waveform /= 2147483648.0
audio = {
"waveform": waveform.unsqueeze(0),
"sample_rate": container.streams.audio[0].rate
}
container.close()
if not frames: raise Exception("Failed to decode video frames")
if not audio_frames:
print("Warning: No audio frames extracted.")
return None
audio_array = np.concatenate(audio_frames, axis=1)
# Normalize based on dtype
if audio_array.dtype == np.int16:
waveform = torch.from_numpy(audio_array).float() / 32768.0
elif audio_array.dtype == np.int32:
waveform = torch.from_numpy(audio_array).float() / 2147483648.0
else:
waveform = torch.from_numpy(audio_array).float()
return {"waveform": waveform.unsqueeze(0), "sample_rate": audio_stream.rate}
except Exception as e:
print(f"Warning: Failed to extract audio: {str(e)}")
return None
return ([torch.cat(frames, dim=0)], audio)
finally:
if os.path.exists(temp_video_path):
try:
os.remove(temp_video_path)
except:
pass
def tensor_to_image_bytes(self, image_tensor):
"""Convert ComfyUI image tensor to bytes"""
img_array = image_tensor.cpu().numpy() if isinstance(image_tensor, torch.Tensor) else image_tensor
if len(img_array.shape) == 4:
img_array = img_array[0]
if img_array.dtype in [np.float32, np.float64]:
img_array = (img_array * 255).astype(np.uint8)
buffered = io.BytesIO()
Image.fromarray(img_array).save(buffered, format="PNG")
return buffered.getvalue()
def generate_video(self, prompt: str, project_id: str, location: str,
service_account: str, model: str, resolution: str, aspect_ratio: str,
duration_seconds: int, seed: int, generate_audio: bool, fps: str,
negative_prompt: Optional[str] = None,
first_frame: Optional[torch.Tensor] = None,
last_frame: Optional[torch.Tensor] = None):
# Initialize Vertex AI client
client = self.setup_client(service_account, project_id, location)
# Configure video generation
config_params = {
"resolution": resolution,
"aspect_ratio": aspect_ratio,
"duration_seconds": duration_seconds,
"generate_audio": generate_audio,
"fps": int(fps),
}
if seed != -1:
config_params["seed"] = seed
if negative_prompt and negative_prompt.strip():
config_params["negative_prompt"] = negative_prompt.strip()
video_config = types.GenerateVideosConfig(**config_params)
# Prepare generation parameters
generation_params = {
"model": model,
"prompt": prompt,
"config": video_config,
}
# Handle first frame image
if first_frame is not None:
image_bytes = self.tensor_to_image_bytes(first_frame)
generation_params["image"] = types.Image(
image_bytes=image_bytes,
mime_type="image/png"
)
print("First frame image provided for video generation")
# Handle last frame image (dynamic attribute for preview SDK)
if last_frame is not None:
last_image_bytes = self.tensor_to_image_bytes(last_frame)
last_frame_img = types.Image(
image_bytes=last_image_bytes,
mime_type="image/png"
)
setattr(video_config, 'last_frame', last_frame_img)
print("Last frame image provided for video generation")
operation = client.models.generate_videos(**generation_params)
print(f"Operation started: {operation.name}")
# Poll for completion
print("Waiting for video generation to complete...")
while not operation.done:
time.sleep(10)
operation = client.operations.get(operation)
print(".", end="", flush=True)
print("")
if operation.error:
raise Exception(f"Operation failed: {operation.error}")
if not operation.result or not operation.result.generated_videos:
raise Exception("No videos were generated.")
video_result = operation.result.generated_videos[0].video
if not video_result.video_bytes:
raise Exception("No video bytes returned from API")
frames_tensor = self.video_to_frames(video_result.video_bytes)
audio_output = None
if generate_audio:
audio_output = self.extract_audio_from_video(video_result.video_bytes)
return ([frames_tensor], audio_output)
# Cleanup Auth File
if os.path.exists(creds_file.name):
os.remove(creds_file.name)
NODE_CLASS_MAPPINGS = {
"VeoVertexVideoGenerator": VeoVertexVideoGenerator
}
NODE_DISPLAY_NAME_MAPPINGS = {
"VeoVertexVideoGenerator": "Veo (Vertex AI)"
}
NODE_CLASS_MAPPINGS = {"VeoVertexVideoGenerator": VeoVertexVideoGenerator}
NODE_DISPLAY_NAME_MAPPINGS = {"VeoVertexVideoGenerator": "Veo (Vertex AI)"}
veo_api.py
+49 -122
View File
@@ -1,15 +1,12 @@
import time
import os
import io
import torch
import numpy as np
import av
from PIL import Image
import tempfile
import uuid
from typing import Optional
from google import genai
from google.genai import types
import cv2
class VeoGeminiVideoGenerator:
@classmethod
@@ -35,124 +32,54 @@ class VeoGeminiVideoGenerator:
CATEGORY = "video/generation"
OUTPUT_IS_LIST = (True,)
def pil_to_tensor(self, pil_image):
if pil_image.mode != 'RGB':
pil_image = pil_image.convert('RGB')
def generate_video(self, seed, prompt, model, aspect_ratio, duration_seconds, api_key,
negative_prompt=None, image=None):
numpy_image = np.array(pil_image).astype(np.float32) / 255.0
return torch.from_numpy(numpy_image).unsqueeze(0)
def video_to_frames(self, video_file):
temp_video_path = os.path.join(tempfile.gettempdir(), f"temp_video_{uuid.uuid4().hex}.mp4")
try:
video_file.save(temp_video_path)
cap = cv2.VideoCapture(temp_video_path)
frames = []
while True:
ret, frame = cap.read()
if not ret:
break
frame_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
tensor_frame = self.pil_to_tensor(Image.fromarray(frame_rgb))
frames.append(tensor_frame)
cap.release()
if not frames:
raise ValueError("No frames extracted from video")
return torch.cat(frames, dim=0)
finally:
if os.path.exists(temp_video_path):
os.remove(temp_video_path)
def tensor_to_image_bytes(self, image_tensor):
"""Convert ComfyUI image tensor to bytes"""
img_array = image_tensor.cpu().numpy() if isinstance(image_tensor, torch.Tensor) else image_tensor
if len(img_array.shape) == 4:
img_array = img_array[0]
if img_array.dtype in [np.float32, np.float64]:
img_array = (img_array * 255).astype(np.uint8)
buffered = io.BytesIO()
Image.fromarray(img_array).save(buffered, format="PNG")
return buffered.getvalue()
def generate_video(self, seed: int, prompt: str, model: str, aspect_ratio: str, duration_seconds: int, api_key: str,
negative_prompt: Optional[str] = None, image: Optional[torch.Tensor] = None):
try:
key = api_key.strip() or os.environ.get("GEMINI_API_KEY")
if not key:
raise ValueError("Error: No API key provided. Set GEMINI_API_KEY or provide api_key parameter.")
client = genai.Client(http_options={"api_version": "v1beta"}, api_key=key)
comfy_seed = seed
config_params = {
"aspect_ratio": aspect_ratio,
"number_of_videos": 1,
"duration_seconds": duration_seconds,
}
if negative_prompt and negative_prompt.strip():
config_params["negative_prompt"] = negative_prompt.strip()
generation_params = {
"model": model,
"prompt": prompt,
"config": types.GenerateVideosConfig(**config_params),
}
if image is not None:
generation_params["image"] = types.Image(
image_bytes=self.tensor_to_image_bytes(image),
mime_type="image/png"
)
print("Image input provided for video generation")
print(f"Starting video generation with model {model}...")
operation = client.models.generate_videos(**generation_params)
print(f"Operation started: {operation.name}")
# Wait for operation to complete by polling
print("Waiting for video generation to complete...")
import time
while not operation.done:
time.sleep(10)
operation = client.operations.get(operation)
if not operation.response or not operation.response.generated_videos:
raise Exception("No videos were generated.")
generated_videos = operation.response.generated_videos
print(f"Generated {len(generated_videos)} video(s).")
all_frames = []
for n, generated_video in enumerate(generated_videos):
print(f"Processing video {n+1}/{len(generated_videos)}")
client.files.download(file=generated_video.video)
frames_tensor = self.video_to_frames(generated_video.video)
all_frames.append(frames_tensor)
print(f"Video {n+1} processed. Extracted {frames_tensor.shape[0]} frames.")
return (all_frames,)
except Exception as e:
print(f"Error in video generation: {str(e)}")
return ([torch.zeros((1, 64, 64, 3), dtype=torch.float32)],)
# 1. Setup Client
key = api_key.strip() or os.environ.get("GEMINI_API_KEY")
if not key: raise ValueError("API Key required")
client = genai.Client(http_options={"api_version": "v1beta"}, api_key=key)
# 2. Prepare Config (Seed used for ComfyUI caching only)
gen_kwargs = {
"model": model,
"prompt": prompt,
"config": types.GenerateVideosConfig(
aspect_ratio=aspect_ratio,
duration_seconds=duration_seconds,
negative_prompt=negative_prompt.strip() if negative_prompt else None,
)
}
NODE_CLASS_MAPPINGS = {
"VeoGeminiVideoGenerator": VeoGeminiVideoGenerator
}
# 3. Handle Image (In-Memory, ComfyUI Tensors are always [B,H,W,C])
if image is not None:
buf = io.BytesIO()
Image.fromarray((image[0].cpu().numpy() * 255).astype(np.uint8)).save(buf, format="PNG")
gen_kwargs["image"] = types.Image(image_bytes=buf.getvalue(), mime_type="image/png")
NODE_DISPLAY_NAME_MAPPINGS = {
"VeoGeminiVideoGenerator": "Veo Video Generator (Gemini API)"
}
# 4. Generate & Poll
op = client.models.generate_videos(**gen_kwargs)
print(f"Gemini Veo Operation: {op.name}")
while not op.done:
time.sleep(4)
op = client.operations.get(op)
if not op.response or not op.response.generated_videos:
raise Exception("No videos generated")
# 5. Download & Decode (Direct Memory Stream)
video_bytes = io.BytesIO(client.files.download(file=op.response.generated_videos[0].video))
container = av.open(video_bytes)
frames = []
for frame in container.decode(video=0):
frames.append(torch.from_numpy(frame.to_rgb().to_ndarray().astype(np.float32) / 255.0).unsqueeze(0))
container.close()
if not frames: raise Exception("Failed to decode video frames")
return ([torch.cat(frames, dim=0)],)
NODE_CLASS_MAPPINGS = {"VeoGeminiVideoGenerator": VeoGeminiVideoGenerator}
NODE_DISPLAY_NAME_MAPPINGS = {"VeoGeminiVideoGenerator": "Veo (Gemini API)"}