[Bugfix][Frontend] Preserve generated audio in ComfyUI video output

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,

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:

apps/ComfyUI-vLLM-Omni/comfyui_vllm_omni/utils/format.py

@@ -152,48 +152,52 @@ 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,
                 frame_rate=frame_rate,
                 audio=audio,
                 metadata=container.metadata if container.metadata else None,
             )
+
     except Exception as e:
         raise RuntimeError(f"Failed to decode video: {e}")