diff --git a/apps/ComfyUI-vLLM-Omni/comfyui_vllm_omni/nodes.py b/apps/ComfyUI-vLLM-Omni/comfyui_vllm_omni/nodes.py index 8fd00cd9c6c..a83043d2197 100644 --- a/apps/ComfyUI-vLLM-Omni/comfyui_vllm_omni/nodes.py +++ b/apps/ComfyUI-vLLM-Omni/comfyui_vllm_omni/nodes.py @@ -163,6 +163,7 @@ def INPUT_TYPES(cls): "height": ("INT", {"default": 480, "min": 1}), "fps": ("INT", {"default": 16, "min": 1}), "num_frames": ("INT", {"default": 41, "min": 1}), + "generate_sound": ("BOOLEAN", {"default": False}), }, "optional": { "frame": ("IMAGE",), @@ -195,6 +196,7 @@ async def generate( height: int, fps: int, num_frames: int, + generate_sound: bool, negative_prompt: str | None = None, frame: torch.Tensor | None = None, references: dict | None = None, @@ -234,6 +236,7 @@ async def generate( height=height, num_frames=num_frames, fps=fps, + generate_sound=generate_sound, negative_prompt=negative_prompt, sampling_params=sampling_params, lora=lora, diff --git a/apps/ComfyUI-vLLM-Omni/comfyui_vllm_omni/utils/api_client.py b/apps/ComfyUI-vLLM-Omni/comfyui_vllm_omni/utils/api_client.py index cd509244a32..44f86fd1bfc 100644 --- a/apps/ComfyUI-vLLM-Omni/comfyui_vllm_omni/utils/api_client.py +++ b/apps/ComfyUI-vLLM-Omni/comfyui_vllm_omni/utils/api_client.py @@ -265,6 +265,7 @@ async def generate_video( height: int, num_frames: int, fps: int, + generate_sound: bool = False, negative_prompt: str | None = None, frame: torch.Tensor | None = None, references: dict | None = None, @@ -284,6 +285,10 @@ async def generate_video( form.add_field("height", str(height)) form.add_field("num_frames", str(num_frames)) form.add_field("fps", str(fps)) + form.add_field( + "generate_sound", + "true" if generate_sound else "false", + ) if negative_prompt: form.add_field("negative_prompt", negative_prompt) if sampling_params is not None: diff --git a/apps/ComfyUI-vLLM-Omni/comfyui_vllm_omni/utils/format.py b/apps/ComfyUI-vLLM-Omni/comfyui_vllm_omni/utils/format.py index 42d396f0694..6d2c53e153e 100644 --- a/apps/ComfyUI-vLLM-Omni/comfyui_vllm_omni/utils/format.py +++ b/apps/ComfyUI-vLLM-Omni/comfyui_vllm_omni/utils/format.py @@ -152,41 +152,44 @@ def bytes_to_video(video_bytes: bytes) -> VideoInput: if video_stream is None: raise ValueError("No video stream found in decoded payload.") + audio_stream = container.streams.audio[-1] if len(container.streams.audio) else None + frames = [] - for frame in container.decode(video_stream): - if not isinstance(frame, VideoFrame): - continue - image = frame.to_ndarray(format="rgb24") - frames.append(torch.from_numpy(image).float() / 255.0) + audio_frames = [] + resampler = AudioResampler(format="fltp") if audio_stream is not None else None - if len(frames) == 0: - raise ValueError("No video frames found in decoded payload.") + streams = [video_stream] + if audio_stream is not None: + streams.append(audio_stream) - images = torch.stack(frames, dim=0) - frame_rate = Fraction(video_stream.average_rate) if video_stream.average_rate else Fraction(1) + for frame in container.decode(*streams): + if isinstance(frame, VideoFrame): + image = frame.to_ndarray(format="rgb24") + frames.append(torch.from_numpy(image).float() / 255.0) - audio: AudioInput | None = None - if len(container.streams.audio): - audio_stream = container.streams.audio[-1] - audio_frames = [] - resampler = AudioResampler(format="fltp") - for frame in container.decode(audio_stream): - if not isinstance(frame, AudioFrame): - continue + elif isinstance(frame, AudioFrame) and resampler is not None: resampled = resampler.resample(frame) if not isinstance(resampled, list): resampled = [resampled] + for audio_frame in resampled: if audio_frame is not None: audio_frames.append(audio_frame.to_ndarray()) - if len(audio_frames) > 0: - audio_data = np.concatenate(audio_frames, axis=1) - sample_rate = int(audio_stream.sample_rate) if audio_stream.sample_rate else 1 - audio = { - "waveform": torch.from_numpy(audio_data).unsqueeze(0), - "sample_rate": sample_rate, - } + if len(frames) == 0: + raise ValueError("No video frames found in decoded payload.") + + images = torch.stack(frames, dim=0) + frame_rate = Fraction(video_stream.average_rate) if video_stream.average_rate else Fraction(1) + + audio: AudioInput | None = None + if audio_stream is not None and len(audio_frames) > 0: + audio_data = np.concatenate(audio_frames, axis=1) + sample_rate = int(audio_stream.sample_rate) if audio_stream.sample_rate else 1 + audio = { + "waveform": torch.from_numpy(audio_data).unsqueeze(0), + "sample_rate": sample_rate, + } components = Types.VideoComponents( images=images, @@ -194,6 +197,7 @@ def bytes_to_video(video_bytes: bytes) -> VideoInput: audio=audio, metadata=container.metadata if container.metadata else None, ) + except Exception as e: raise RuntimeError(f"Failed to decode video: {e}")