diff --git a/python/sglang/multimodal_gen/runtime/entrypoints/utils.py b/python/sglang/multimodal_gen/runtime/entrypoints/utils.py index d084bd9dc..2c5b5839c 100644 --- a/python/sglang/multimodal_gen/runtime/entrypoints/utils.py +++ b/python/sglang/multimodal_gen/runtime/entrypoints/utils.py @@ -8,11 +8,15 @@ This module provides a consolidated interface for generating videos using diffusion models. """ +import atexit import json +import mmap import os import shutil import subprocess import tempfile +import threading +from contextlib import contextmanager from copy import copy from dataclasses import dataclass, field from typing import Any, Callable, List, Optional, Sequence, Union @@ -43,6 +47,114 @@ from sglang.srt.observability.trace import TraceReqContext logger = init_logger(__name__) +_MAX_CACHED_CUDA_VIDEO_BUFFER_BYTES = 1024 * 1024 * 1024 +_cuda_video_buffer_cache_lock = threading.Lock() +_cached_cuda_video_buffer: "_CudaMemfdVideoBuffer | None" = None + + +class _CudaMemfdVideoBuffer: + """CUDA-registered memfd used as direct ffmpeg raw-video input.""" + + def __init__(self, shape: tuple[int, ...]): + self.shape = shape + self.nbytes = int(np.prod(shape, dtype=np.int64)) + self.fd = -1 + self.mapping: mmap.mmap | None = None + self.array: np.ndarray | None = None + self.tensor: torch.Tensor | None = None + self._registered = False + + try: + self.fd = os.memfd_create( + "sglang-video-frames", + flags=getattr(os, "MFD_CLOEXEC", 0), + ) + os.ftruncate(self.fd, self.nbytes) + self.mapping = mmap.mmap( + self.fd, + self.nbytes, + flags=mmap.MAP_SHARED, + prot=mmap.PROT_READ | mmap.PROT_WRITE, + ) + self.array = np.ndarray(shape, dtype=np.uint8, buffer=self.mapping) + error = torch.cuda.cudart().cudaHostRegister( + self.array.ctypes.data, + self.nbytes, + 0, + ) + if error != 0: + raise RuntimeError(f"cudaHostRegister failed: {error}") + self._registered = True + self.tensor = torch.from_numpy(self.array) + if not self.tensor.is_pinned(): + raise RuntimeError("CUDA-registered memfd is not pinned") + except Exception: + self.close() + raise + + def close(self) -> None: + if self._registered and self.array is not None: + try: + torch.cuda.cudart().cudaHostUnregister(self.array.ctypes.data) + except Exception: + pass + self._registered = False + self.tensor = None + self.array = None + if self.mapping is not None: + self.mapping.close() + self.mapping = None + if self.fd >= 0: + os.close(self.fd) + self.fd = -1 + + +@contextmanager +def _acquire_cuda_video_buffer(shape: tuple[int, ...]): + global _cached_cuda_video_buffer + + buffer = None + stale_buffer = None + with _cuda_video_buffer_cache_lock: + if _cached_cuda_video_buffer is not None: + if _cached_cuda_video_buffer.shape == shape: + buffer = _cached_cuda_video_buffer + else: + stale_buffer = _cached_cuda_video_buffer + _cached_cuda_video_buffer = None + + if stale_buffer is not None: + stale_buffer.close() + + if buffer is None: + buffer = _CudaMemfdVideoBuffer(shape) + + try: + yield buffer + finally: + with _cuda_video_buffer_cache_lock: + if ( + buffer.nbytes <= _MAX_CACHED_CUDA_VIDEO_BUFFER_BYTES + and _cached_cuda_video_buffer is None + ): + _cached_cuda_video_buffer = buffer + buffer = None + if buffer is not None: + buffer.close() + + +def _close_cached_cuda_video_buffer() -> None: + global _cached_cuda_video_buffer + + with _cuda_video_buffer_cache_lock: + buffer = _cached_cuda_video_buffer + _cached_cuda_video_buffer = None + if buffer is not None: + buffer.close() + + +atexit.register(_close_cached_cuda_video_buffer) + @dataclass class SetLoraReq: @@ -335,6 +447,157 @@ def _resolve_ffmpeg_exe() -> str: return ffmpeg_exe +def _x264_auto_thread_count(height: int) -> int: + """Match x264's auto frame-thread count for progressive video.""" + try: + cpu_count = len(os.sched_getaffinity(0)) + except (AttributeError, OSError): + cpu_count = os.cpu_count() or 1 + cpu_limit = max(1, cpu_count * 3 // 2) + macroblock_rows = max(1, (height + 15) // 16) + row_limit = max(1, macroblock_rows // 2) + return min(cpu_limit, row_limit, 128) + + +def _try_save_cuda_video_direct( + *, + save_file_path: str, + sample: Any, + fps: int, + audio_sample_rate: Optional[int], + output_compression: Optional[int], +) -> bool: + """Save a CUDA RGB video through a registered memfd instead of a Python pipe.""" + if not hasattr(os, "memfd_create") or not os.path.isdir("/proc/self/fd"): + return False + + sample_without_audio, audio = _split_sample_audio(sample) + if not ( + isinstance(sample_without_audio, torch.Tensor) + and sample_without_audio.device.type == "cuda" + and sample_without_audio.dim() in (3, 4) + ): + return False + if os.path.splitext(save_file_path)[1].lower() != ".mp4": + return False + + video = sample_without_audio + if video.dim() == 3: + video = video.unsqueeze(1) + if video.shape[0] != 3: + return False + + frames = (video * 255).clamp(0, 255).to(torch.uint8) + frames = frames.permute(1, 2, 3, 0).contiguous() + num_frames, height, width, _ = frames.shape + + quality = output_compression / 10 if output_compression is not None else 5 + if not 1 <= quality <= 10: + return False + crf = int((1 - quality / 10.0) * 51) + + audio_np = _normalize_audio_to_numpy(audio) + tmp_wav_path = None + try: + if audio_np is not None: + if scipy_wavfile is None: + return False + selected_sr = _pick_audio_sample_rate( + audio_np=audio_np, + audio_sample_rate=audio_sample_rate, + fps=fps, + num_frames=num_frames, + ) + with tempfile.NamedTemporaryFile(suffix=".wav", delete=False) as f: + tmp_wav_path = f.name + scipy_wavfile.write(tmp_wav_path, selected_sr, audio_np) + + ffmpeg_exe = _resolve_ffmpeg_exe() + shape = tuple(frames.shape) + with _acquire_cuda_video_buffer(shape) as buffer: + assert buffer.tensor is not None + buffer.tensor.copy_(frames, non_blocking=True) + torch.cuda.current_stream(frames.device).synchronize() + del frames + os.lseek(buffer.fd, 0, os.SEEK_SET) + + command = [ + ffmpeg_exe, + "-y", + "-f", + "rawvideo", + "-vcodec", + "rawvideo", + "-s", + f"{width}x{height}", + "-pix_fmt", + "rgb24", + "-r", + f"{fps:.02f}", + "-i", + f"/proc/self/fd/{buffer.fd}", + ] + if tmp_wav_path is None: + command += ["-an"] + else: + command += ["-i", tmp_wav_path] + command += [ + "-vcodec", + "libx264", + "-pix_fmt", + "yuv420p", + "-crf", + str(crf), + ] + + macro_block_size = 16 + if width % macro_block_size or height % macro_block_size: + output_width = ( + width + if width % macro_block_size == 0 + else width + macro_block_size - width % macro_block_size + ) + output_height = ( + height + if height % macro_block_size == 0 + else height + macro_block_size - height % macro_block_size + ) + command += ["-vf", f"scale={output_width}:{output_height}"] + + command += ["-threads", str(_x264_auto_thread_count(height))] + if tmp_wav_path is not None: + command += [ + "-acodec", + "aac", + "-map", + "0:v:0", + "-map", + "1:a:0", + ] + command += ["-v", "warning", save_file_path] + + subprocess.run( + command, + check=True, + pass_fds=(buffer.fd,), + stdout=subprocess.DEVNULL, + stderr=subprocess.PIPE, + ) + return True + except Exception as e: + logger.warning( + "Direct CUDA video save failed; falling back to imageio: %s", + str(e), + ) + return False + finally: + if tmp_wav_path: + try: + os.remove(tmp_wav_path) + except OSError: + pass + + def _mux_audio_np_into_mp4( *, save_file_path: str, @@ -423,6 +686,62 @@ def _maybe_mux_audio_into_mp4( ) +def _try_save_video_with_audio( + *, + save_file_path: str, + frames: list, + fps: int, + audio: Any, + audio_sample_rate: Optional[int], + output_format: str, + quality: float, +) -> bool: + """Encode video and audio in one ffmpeg pass when audio is available.""" + audio_np = _normalize_audio_to_numpy(audio) + if audio_np is None: + return False + + selected_sr = _pick_audio_sample_rate( + audio_np=audio_np, + audio_sample_rate=audio_sample_rate, + fps=fps, + num_frames=len(frames), + ) + tmp_wav_path = None + try: + if scipy_wavfile is None: + raise RuntimeError( + "scipy is required to mux audio into mp4 (pip install scipy)" + ) + with tempfile.NamedTemporaryFile(suffix=".wav", delete=False) as f: + tmp_wav_path = f.name + scipy_wavfile.write(tmp_wav_path, selected_sr, audio_np) + imageio.mimsave( + save_file_path, + frames, + fps=fps, + format=output_format, + codec="libx264", + quality=quality, + audio_path=tmp_wav_path, + audio_codec="aac", + ) + return True + except Exception as e: + logger.warning( + "Failed to encode video and audio in one pass; " + "falling back to the compatible two-pass path: %s", + str(e), + ) + return False + finally: + if tmp_wav_path: + try: + os.remove(tmp_wav_path) + except OSError: + pass + + def prepare_request( server_args: ServerArgs, sampling_params: SamplingParams, @@ -502,7 +821,23 @@ def _sample_to_uint8_frames(sample: Any) -> list[Any]: if sample.dim() == 3: sample = sample.unsqueeze(1) sample = (sample * 255).clamp(0, 255).to(torch.uint8) - videos = sample.permute(1, 2, 3, 0).contiguous().cpu().numpy() + videos = sample.permute(1, 2, 3, 0).contiguous() + if videos.device.type == "cuda": + try: + host_videos = torch.empty( + videos.shape, + dtype=videos.dtype, + device="cpu", + pin_memory=True, + ) + except RuntimeError: + videos = videos.cpu().numpy() + else: + host_videos.copy_(videos, non_blocking=True) + torch.cuda.current_stream(videos.device).synchronize() + videos = host_videos.numpy() + else: + videos = videos.cpu().numpy() return list(videos) if not isinstance(sample, np.ndarray): @@ -592,22 +927,33 @@ def save_materialized_output( os.makedirs(os.path.dirname(save_file_path), exist_ok=True) if data_type == DataType.VIDEO: quality = output_compression / 10 if output_compression is not None else 5 - imageio.mimsave( - save_file_path, - materialized.frames, - fps=materialized.fps, - format=data_type.get_default_extension(), - codec="libx264", - quality=quality, - ) - - _maybe_mux_audio_into_mp4( + output_format = data_type.get_default_extension() + saved_with_audio = _try_save_video_with_audio( save_file_path=save_file_path, - audio=materialized.audio, frames=materialized.frames, fps=materialized.fps, + audio=materialized.audio, audio_sample_rate=audio_sample_rate, + output_format=output_format, + quality=quality, ) + if not saved_with_audio: + imageio.mimsave( + save_file_path, + materialized.frames, + fps=materialized.fps, + format=output_format, + codec="libx264", + quality=quality, + ) + + _maybe_mux_audio_into_mp4( + save_file_path=save_file_path, + audio=materialized.audio, + frames=materialized.frames, + fps=materialized.fps, + audio_sample_rate=audio_sample_rate, + ) else: quality = output_compression if output_compression is not None else 75 if len(materialized.frames) > 1: @@ -681,6 +1027,28 @@ def save_outputs( if data_type == DataType.VIDEO: sample = attach_audio_to_video_sample(sample, audio, idx) + if ( + save_output + and save_file_path + and frames_out is None + and not enable_frame_interpolation + and not enable_upscaling + ): + os.makedirs(os.path.dirname(save_file_path) or ".", exist_ok=True) + if _try_save_cuda_video_direct( + save_file_path=save_file_path, + sample=sample, + fps=fps, + audio_sample_rate=audio_sample_rate, + output_compression=output_compression, + ): + if samples_out is not None: + samples_out.append(sample) + if audios_out is not None: + audios_out.append(select_output_audio(audio, idx)) + output_paths.append(save_file_path) + logger.info(f"Output saved to {CYAN}{save_file_path}{RESET}") + continue frames = post_process_sample( sample, diff --git a/python/sglang/multimodal_gen/test/unit/test_output_saving.py b/python/sglang/multimodal_gen/test/unit/test_output_saving.py index c3113331b..5157816a6 100644 --- a/python/sglang/multimodal_gen/test/unit/test_output_saving.py +++ b/python/sglang/multimodal_gen/test/unit/test_output_saving.py @@ -1,10 +1,15 @@ import numpy as np import pytest +import torch from PIL import Image import sglang.multimodal_gen.runtime.entrypoints.utils as output_utils from sglang.multimodal_gen.configs.sample.sampling_params import DataType -from sglang.multimodal_gen.runtime.entrypoints.utils import post_process_sample +from sglang.multimodal_gen.runtime.entrypoints.utils import ( + MaterializedOutput, + post_process_sample, + save_materialized_output, +) def _rgb_frame() -> np.ndarray: @@ -66,3 +71,131 @@ def test_png_output_saving_uses_fast_pillow_path( ) assert save_calls == [("PNG", expected_compress_level)] + + +def test_video_with_audio_uses_single_pass_encoder(tmp_path, monkeypatch): + output_path = tmp_path / "sample.mp4" + calls = [] + + class FakeWavFile: + @staticmethod + def write(*_args, **_kwargs): + pass + + def mimsave_spy(path, frames, **kwargs): + calls.append((path, frames, kwargs)) + assert kwargs["audio_path"].endswith(".wav") + assert kwargs["audio_codec"] == "aac" + + def fail_legacy_mux(**_kwargs): + raise AssertionError("the two-pass mux path should not run") + + monkeypatch.setattr(output_utils.imageio, "mimsave", mimsave_spy) + monkeypatch.setattr(output_utils, "scipy_wavfile", FakeWavFile) + monkeypatch.setattr(output_utils, "_maybe_mux_audio_into_mp4", fail_legacy_mux) + + materialized = MaterializedOutput( + sample=None, + frames=[_rgb_frame()], + audio=np.zeros((320, 2), dtype=np.float32), + fps=24, + ) + save_materialized_output( + materialized, + DataType.VIDEO, + str(output_path), + audio_sample_rate=32000, + ) + + assert len(calls) == 1 + + +def test_video_audio_single_pass_failure_falls_back(tmp_path, monkeypatch): + output_path = tmp_path / "sample.mp4" + calls = [] + mux_calls = [] + + class FakeWavFile: + @staticmethod + def write(*_args, **_kwargs): + pass + + def mimsave_spy(path, frames, **kwargs): + calls.append((path, frames, kwargs)) + if "audio_path" in kwargs: + raise RuntimeError("unsupported audio input") + + monkeypatch.setattr(output_utils.imageio, "mimsave", mimsave_spy) + monkeypatch.setattr(output_utils, "scipy_wavfile", FakeWavFile) + monkeypatch.setattr( + output_utils, + "_maybe_mux_audio_into_mp4", + lambda **kwargs: mux_calls.append(kwargs), + ) + + materialized = MaterializedOutput( + sample=None, + frames=[_rgb_frame()], + audio=np.zeros((320, 2), dtype=np.float32), + fps=24, + ) + save_materialized_output( + materialized, + DataType.VIDEO, + str(output_path), + audio_sample_rate=32000, + ) + + assert len(calls) == 2 + assert "audio_path" in calls[0][2] + assert "audio_path" not in calls[1][2] + assert len(mux_calls) == 1 + + +@pytest.mark.parametrize( + ("height", "available_cpus", "expected_threads"), + [ + (768, 256, 24), + (720, 256, 22), + (2160, 16, 24), + (4320, 256, 128), + (16, 1, 1), + ], +) +def test_x264_auto_thread_count(monkeypatch, height, available_cpus, expected_threads): + monkeypatch.setattr( + output_utils.os, + "sched_getaffinity", + lambda _pid: set(range(available_cpus)), + ) + + assert output_utils._x264_auto_thread_count(height) == expected_threads + + +def test_video_direct_save_short_circuits_materialization(tmp_path, monkeypatch): + output_path = tmp_path / "sample.mp4" + direct_calls = [] + + monkeypatch.setattr( + output_utils, + "_try_save_cuda_video_direct", + lambda **kwargs: direct_calls.append(kwargs) or True, + ) + monkeypatch.setattr( + output_utils, + "post_process_sample", + lambda *_args, **_kwargs: pytest.fail( + "successful direct save should skip frame materialization" + ), + ) + + paths = output_utils.save_outputs( + [torch.zeros((3, 1, 2, 3))], + DataType.VIDEO, + fps=24, + save_output=True, + build_output_path=lambda _idx: str(output_path), + ) + + assert paths == [str(output_path)] + assert len(direct_calls) == 1