diff --git a/veo.py b/veo.py index 7a197ef..c6b7a73 100644 --- a/veo.py +++ b/veo.py @@ -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)" -} \ No newline at end of file +NODE_CLASS_MAPPINGS = {"VeoVertexVideoGenerator": VeoVertexVideoGenerator} +NODE_DISPLAY_NAME_MAPPINGS = {"VeoVertexVideoGenerator": "Veo (Vertex AI)"} \ No newline at end of file diff --git a/veo_api.py b/veo_api.py index d63ae7a..450f178 100644 --- a/veo_api.py +++ b/veo_api.py @@ -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)" -} \ No newline at end of file + # 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)"} \ No newline at end of file