[diffusion] optimization: accelerate CUDA video output finalization (#32784)
This commit is contained in:
Mick
@@ -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,
|
||||
|
||||
@@ -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
|
||||
|
||||
Reference in New Issue
Block a user