diff --git a/requirements.txt b/requirements.txt index 5101b28..fd9f702 100644 --- a/requirements.txt +++ b/requirements.txt @@ -4,4 +4,5 @@ numpy torch google-genai opencv-python -torchaudio \ No newline at end of file +torchaudio +av \ No newline at end of file diff --git a/veo_vertex.py b/veo_vertex.py index 58fb7a6..35716ae 100644 --- a/veo_vertex.py +++ b/veo_vertex.py @@ -11,6 +11,7 @@ from typing import Optional from google import genai from google.genai import types import cv2 +import av class GoogleVeoVertexVideoGenerator: @@ -47,6 +48,8 @@ class GoogleVeoVertexVideoGenerator: "aspect_ratio": (["16:9", "9:16"], {"default": "16:9"}), "duration_seconds": ("INT", {"default": 4, "min": 4, "max": 8, "step": 1}), "seed": ("INT", {"default": 69, "min": 1, "max": 2147483646, "step": 1}), + "generate_audio": ("BOOLEAN", {"default": False}), + "fps": (["24"], {"default": "24"}), }, "optional": { "negative_prompt": ("STRING", {"multiline": True, "default": ""}), @@ -55,11 +58,11 @@ class GoogleVeoVertexVideoGenerator: } } - RETURN_TYPES = ("IMAGE",) - RETURN_NAMES = ("frames",) + RETURN_TYPES = ("IMAGE", "AUDIO") + RETURN_NAMES = ("frames", "audio") FUNCTION = "generate_video" CATEGORY = "video/generation" - OUTPUT_IS_LIST = (True,) + OUTPUT_IS_LIST = (True, False) def setup_client(self, service_account_json, project_id, location): """Setup Vertex AI client with service account JSON content""" @@ -126,6 +129,49 @@ class GoogleVeoVertexVideoGenerator: 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)] + container.close() + + 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 + 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 @@ -142,7 +188,7 @@ class GoogleVeoVertexVideoGenerator: 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, + 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): @@ -155,6 +201,8 @@ class GoogleVeoVertexVideoGenerator: "resolution": resolution, "aspect_ratio": aspect_ratio, "duration_seconds": duration_seconds, + "generate_audio": generate_audio, + "fps": int(fps), } if seed != -1: @@ -191,7 +239,6 @@ class GoogleVeoVertexVideoGenerator: setattr(video_config, 'last_frame', last_frame_img) print("Last frame image 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}") @@ -203,24 +250,23 @@ class GoogleVeoVertexVideoGenerator: print(".", end="", flush=True) print("") - # Check for errors if operation.error: raise Exception(f"Operation failed: {operation.error}") - # Retrieve video 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") - print("Video generated successfully. Extracting frames...") frames_tensor = self.video_to_frames(video_result.video_bytes) - print(f"Extracted {frames_tensor.shape[0]} frames.") - return ([frames_tensor],) + audio_output = None + if generate_audio: + audio_output = self.extract_audio_from_video(video_result.video_bytes) + + return ([frames_tensor], audio_output) NODE_CLASS_MAPPINGS = {