diff --git a/components/src/dynamo/common/tests/test_video_utils.py b/components/src/dynamo/common/tests/test_video_utils.py index fab867fb611c..0ce3a025e61e 100644 --- a/components/src/dynamo/common/tests/test_video_utils.py +++ b/components/src/dynamo/common/tests/test_video_utils.py @@ -3,6 +3,7 @@ """Unit tests for dynamo.common.utils.video_utils module.""" +import io from unittest.mock import MagicMock, patch import numpy as np @@ -15,142 +16,127 @@ ] -def make_frames(n=3, h=8, w=8) -> np.ndarray: - """Return a small uint8 frame array (n, h, w, 3).""" +def make_frames(n=3, h=16, w=16) -> np.ndarray: + """Return a small uint8 frame array (n, h, w, 3). Even dims for libx264.""" return np.zeros((n, h, w, 3), dtype=np.uint8) # --------------------------------------------------------------------------- -# encode_to_video_bytes +# encode_to_video_bytes — PyAV/libx264 software encode (no NVENC, no imageio) # --------------------------------------------------------------------------- class TestEncodeToVideoBytes: - """Tests for encode_to_video_bytes().""" - - def _mock_iio_v3(self): - """Return a mock that looks like imageio.v3 (has imwrite).""" - iio = MagicMock() - iio.imwrite = MagicMock() - return iio - - def _mock_iio_v2(self): - """Return a mock that looks like imageio v2 (no imwrite, has get_writer).""" - iio = MagicMock(spec=[]) # no attributes by default - writer = MagicMock() - iio.get_writer = MagicMock(return_value=writer) - return iio, writer - - def test_mp4_selects_libx264_codec(self): - from dynamo.common.utils.video_utils import encode_to_video_bytes - - iio = self._mock_iio_v3() - with patch("dynamo.common.utils.video_utils.io") as mock_io, patch( - "imageio.v3", iio, create=True - ), patch.dict("sys.modules", {"imageio.v3": iio}): - buf = MagicMock() - buf.getvalue.return_value = b"fake-mp4" - mock_io.BytesIO.return_value = buf + """Tests for encode_to_video_bytes(): software-encodes via PyAV so it works on + NVENC-less datacenter GPUs. These exercise the real encoder (skipped if PyAV is + unavailable) rather than mocking, since the whole point is that the bytes decode.""" - encode_to_video_bytes(make_frames(), fps=8, output_format="mp4") + def test_mp4_returns_h264_bytes(self): + av = pytest.importorskip("av") + from dynamo.common.utils.video_utils import encode_to_video_bytes - iio.imwrite.assert_called_once() - _, kwargs = iio.imwrite.call_args - assert kwargs.get("codec") == "libx264" - assert kwargs.get("fps") == 8 + out = encode_to_video_bytes(make_frames(n=5), fps=8, output_format="mp4") + assert isinstance(out, bytes) and len(out) > 0 + with av.open(io.BytesIO(out)) as container: + assert container.streams.video[0].codec_context.name == "h264" - def test_webm_selects_libvpx_vp9_codec(self): + def test_webm_returns_vp9_bytes(self): + av = pytest.importorskip("av") from dynamo.common.utils.video_utils import encode_to_video_bytes - iio = self._mock_iio_v3() - with patch("dynamo.common.utils.video_utils.io") as mock_io, patch( - "imageio.v3", iio, create=True - ), patch.dict("sys.modules", {"imageio.v3": iio}): - buf = MagicMock() - buf.getvalue.return_value = b"fake-webm" - mock_io.BytesIO.return_value = buf + out = encode_to_video_bytes(make_frames(n=5), fps=8, output_format="webm") + assert isinstance(out, bytes) and len(out) > 0 + with av.open(io.BytesIO(out)) as container: + assert container.streams.video[0].codec_context.name in ( + "vp9", + "libvpx-vp9", + ) - encode_to_video_bytes(make_frames(), fps=16, output_format="webm") + def test_mp4_tags_bt709(self): + av = pytest.importorskip("av") + from dynamo.common.utils.video_utils import encode_to_video_bytes - iio.imwrite.assert_called_once() - _, kwargs = iio.imwrite.call_args - assert kwargs.get("codec") == "libvpx-vp9" + out = encode_to_video_bytes(make_frames(n=5), fps=16, output_format="mp4") + with av.open(io.BytesIO(out)) as container: + cc = container.streams.video[0].codec_context + # BT.709 primaries/transfer/colorspace so players don't render washed-out. + assert int(cc.color_primaries) == 1 + assert int(cc.color_trc) == 1 + assert int(cc.colorspace) == 1 - def test_mp4_passes_extension_to_imwrite(self): + def test_squeezes_5d_batch_dim(self): + av = pytest.importorskip("av") from dynamo.common.utils.video_utils import encode_to_video_bytes - iio = self._mock_iio_v3() - with patch("dynamo.common.utils.video_utils.io") as mock_io, patch( - "imageio.v3", iio, create=True - ), patch.dict("sys.modules", {"imageio.v3": iio}): - buf = MagicMock() - buf.getvalue.return_value = b"bytes" - mock_io.BytesIO.return_value = buf + # MediaOutput.video is (B, T, H, W, C) since TRT-LLM rc9; encode must squeeze B=1. + frames_5d = make_frames(n=5)[None] # (1, 5, 16, 16, 3) + out = encode_to_video_bytes(frames_5d, fps=8, output_format="mp4") + with av.open(io.BytesIO(out)) as container: + assert container.streams.video[0].codec_context.name == "h264" - encode_to_video_bytes(make_frames(), output_format="mp4") + def test_unsupported_format_raises_value_error(self): + pytest.importorskip("av") + from dynamo.common.utils.video_utils import encode_to_video_bytes - _, kwargs = iio.imwrite.call_args - assert kwargs.get("extension") == ".mp4" + with pytest.raises(ValueError): + encode_to_video_bytes(make_frames(), output_format="avi") - def test_webm_passes_extension_to_imwrite(self): - from dynamo.common.utils.video_utils import encode_to_video_bytes - iio = self._mock_iio_v3() - with patch("dynamo.common.utils.video_utils.io") as mock_io, patch( - "imageio.v3", iio, create=True - ), patch.dict("sys.modules", {"imageio.v3": iio}): - buf = MagicMock() - buf.getvalue.return_value = b"bytes" - mock_io.BytesIO.return_value = buf +# --------------------------------------------------------------------------- +# limit_video_worker_threads / video_encode_threads — shared video-worker caps +# --------------------------------------------------------------------------- - encode_to_video_bytes(make_frames(), output_format="webm") - _, kwargs = iio.imwrite.call_args - assert kwargs.get("extension") == ".webm" +class TestVideoWorkerThreadCaps: + """The three env-overridable CPU-thread knobs + the OpenMP wait policy.""" - def test_unsupported_format_raises_value_error(self): - from dynamo.common.utils.video_utils import encode_to_video_bytes + def test_encode_threads_default(self, monkeypatch): + monkeypatch.delenv("DI_VIDEO_ENCODE_THREADS", raising=False) + from dynamo.common.utils.video_utils import ( + DEFAULT_VIDEO_ENCODE_THREADS, + video_encode_threads, + ) - iio = self._mock_iio_v3() - with patch("dynamo.common.utils.video_utils.io") as mock_io, patch( - "imageio.v3", iio, create=True - ), patch.dict("sys.modules", {"imageio.v3": iio}): - mock_io.BytesIO.return_value = MagicMock() + assert video_encode_threads() == DEFAULT_VIDEO_ENCODE_THREADS - # ValueError is wrapped into RuntimeError by the except block - with pytest.raises(RuntimeError, match="Video encoding to bytes failed"): - encode_to_video_bytes(make_frames(), output_format="avi") + def test_encode_threads_env_override(self, monkeypatch): + monkeypatch.setenv("DI_VIDEO_ENCODE_THREADS", "7") + from dynamo.common.utils.video_utils import video_encode_threads - def test_returns_bytes_from_buffer(self): - from dynamo.common.utils.video_utils import encode_to_video_bytes + assert video_encode_threads() == 7 - expected = b"\x00\x01\x02" - iio = self._mock_iio_v3() - with patch("dynamo.common.utils.video_utils.io") as mock_io, patch( - "imageio.v3", iio, create=True - ), patch.dict("sys.modules", {"imageio.v3": iio}): - buf = MagicMock() - buf.getvalue.return_value = expected - mock_io.BytesIO.return_value = buf + def test_limit_sets_wait_policy_compile_and_torch(self, monkeypatch): + import os - result = encode_to_video_bytes(make_frames(), output_format="mp4") + for k in ( + "OMP_WAIT_POLICY", + "KMP_BLOCKTIME", + "TORCHINDUCTOR_COMPILE_THREADS", + "DI_VIDEO_COMPILE_THREADS", + "DI_VIDEO_TORCH_THREADS", + ): + monkeypatch.delenv(k, raising=False) + from dynamo.common.utils import video_utils - assert result == expected + fake_torch = MagicMock() + with patch.dict("sys.modules", {"torch": fake_torch}): + video_utils.limit_video_worker_threads() - def test_v2_api_fallback_writes_all_frames(self): - """When imageio.v3.imwrite is absent, falls back to get_writer loop.""" - from dynamo.common.utils.video_utils import encode_to_video_bytes + assert os.environ["OMP_WAIT_POLICY"] == "PASSIVE" + assert os.environ["KMP_BLOCKTIME"] == "0" + assert os.environ["TORCHINDUCTOR_COMPILE_THREADS"] == str( + video_utils.DEFAULT_VIDEO_COMPILE_THREADS + ) + fake_torch.set_num_threads.assert_called_once_with( + video_utils.DEFAULT_VIDEO_TORCH_THREADS + ) - iio_v2, writer = self._mock_iio_v2() - with patch("dynamo.common.utils.video_utils.io") as mock_io, patch( - "imageio.v3", iio_v2, create=True - ), patch.dict("sys.modules", {"imageio.v3": iio_v2}): - buf = MagicMock() - buf.getvalue.return_value = b"v2-bytes" - mock_io.BytesIO.return_value = buf + def test_limit_torch_threads_env_override(self, monkeypatch): + monkeypatch.setenv("DI_VIDEO_TORCH_THREADS", "9") + from dynamo.common.utils import video_utils - frames = make_frames(n=4) - encode_to_video_bytes(frames, output_format="mp4") + fake_torch = MagicMock() + with patch.dict("sys.modules", {"torch": fake_torch}): + video_utils.limit_video_worker_threads() - assert writer.append_data.call_count == 4 - writer.close.assert_called_once() + fake_torch.set_num_threads.assert_called_once_with(9) diff --git a/components/src/dynamo/common/utils/video_utils.py b/components/src/dynamo/common/utils/video_utils.py index 37326d3280bc..61024f295e13 100644 --- a/components/src/dynamo/common/utils/video_utils.py +++ b/components/src/dynamo/common/utils/video_utils.py @@ -7,7 +7,6 @@ video frames to MP4 format. """ -import io import logging import os from typing import Tuple @@ -175,6 +174,73 @@ def encode_to_mp4( raise RuntimeError(f"Video encoding failed: {e}") from e +# Video workers have three distinct CPU consumers that peak in different phases and want +# different thread counts, so each gets its own env-overridable knob (defaults below): +# - Inductor compile pool (boot / first request: parallel Triton kernel compilation) +# - torch intra-op pool (generation: CPU-side ops around a GPU-bound diffusion) +# - libx264 encode (encode: ~12 is the measured PyAV sweet spot for 720p) +# On shared GPU nodes an uncapped worker is a noisy neighbor: torch sizes its pool to the +# physical core count (~112 on a 224-thread box) and libx264 grabs every core, which both +# thrashes (a 720p/81f encode took 244s vs ~1s) and starves co-located pods. The three +# phases are sequential, so with OMP_WAIT_POLICY=PASSIVE / KMP_BLOCKTIME=0 an idle pool +# sleeps and hands its cores to the active phase instead of spinning — they don't sum. +DEFAULT_VIDEO_TORCH_THREADS = 12 +DEFAULT_VIDEO_ENCODE_THREADS = 12 +DEFAULT_VIDEO_COMPILE_THREADS = 32 + + +def video_encode_threads() -> int: + """CPU threads for the libx264 software encode (env: DI_VIDEO_ENCODE_THREADS).""" + return int(os.getenv("DI_VIDEO_ENCODE_THREADS", str(DEFAULT_VIDEO_ENCODE_THREADS))) + + +def limit_video_worker_threads() -> None: + """Cap a video worker's CPU threads so it stays a good neighbor on shared nodes. + + Call once at the very top of a video worker's startup, before torch loads. Caps + torch's intra-op pool (via torch.set_num_threads, which always applies), bounds the + Inductor compile pool, and sets the OpenMP wait policy so idle pools sleep between + the sequential compile -> generate -> encode phases rather than spinning. The encoder + is capped separately via video_encode_threads(). Every video worker entry calls this + once — that is the general pattern. No-op if torch is unavailable. + + Caveat: OMP_WAIT_POLICY / KMP_BLOCKTIME only bind if set before the OpenMP runtime + initializes (≈ first torch import), so for guaranteed effect set them in the launch + env. torch.set_num_threads and TORCHINDUCTOR_COMPILE_THREADS apply regardless. + """ + import sys + + if "torch" in sys.modules: + logger.warning( + "limit_video_worker_threads() ran after torch was already imported — " + "OMP_WAIT_POLICY/KMP_BLOCKTIME may not take effect; set them in the launch " + "env for reliability (torch.set_num_threads still applies)." + ) + # Import-time knobs: set before torch imports so OpenMP / Inductor pick them up. Idle + # pools then sleep instead of spin (covers both libgomp and Intel OpenMP / MKL). + os.environ.setdefault("OMP_WAIT_POLICY", "PASSIVE") + os.environ.setdefault("KMP_BLOCKTIME", "0") + os.environ.setdefault( + "TORCHINDUCTOR_COMPILE_THREADS", + os.getenv("DI_VIDEO_COMPILE_THREADS", str(DEFAULT_VIDEO_COMPILE_THREADS)), + ) + torch_threads = int( + os.getenv("DI_VIDEO_TORCH_THREADS", str(DEFAULT_VIDEO_TORCH_THREADS)) + ) + try: + import torch + + torch.set_num_threads(torch_threads) + except Exception: # noqa: BLE001 + logger.warning("could not cap torch CPU threads to %s", torch_threads) + logger.info( + "video worker CPU budget: torch=%s encode=%s compile=%s (OMP_WAIT_POLICY=PASSIVE)", + torch_threads, + video_encode_threads(), + os.environ.get("TORCHINDUCTOR_COMPILE_THREADS"), + ) + + def encode_to_video_bytes( frames: np.ndarray, fps: int = 16, @@ -192,51 +258,73 @@ def encode_to_video_bytes( Encoded video as bytes. Raises: - ImportError: If imageio is not available. + ImportError: If PyAV (``av``) is not available. + ValueError: If output_format has no known codec. RuntimeError: If encoding fails. """ + import tempfile + + import av + + # Defensive squeeze: MediaOutput.video is (B, T, H, W, C) since TRT-LLM rc9. + if frames.ndim == 5 and frames.shape[0] == 1: + frames = frames[0] + frames = np.ascontiguousarray(frames, dtype=np.uint8) + num_frames, height, width, _ = frames.shape + + if output_format == "mp4": + codec = "libx264" + elif output_format == "webm": + codec = "libvpx-vp9" + else: + raise ValueError(f"No codec specified for response format: {output_format}") + + logger.info( + f"Encoding {num_frames} frames to {output_format} ({codec}, software) at {fps} fps" + ) + + # Software encode via PyAV: the worker runs on NVENC-less datacenter GPUs + # (B200/H100/A100) where h264_nvenc has no capable device, and the in-tree + # imageio ffmpeg has no software h264. PyAV bundles libx264. Encode to a temp + # file (not an in-memory pipe) for reliability, then read the bytes back. + tmp = tempfile.NamedTemporaryFile(suffix=f".{output_format}", delete=False) + tmp.close() try: - import imageio.v3 as iio - except ImportError: + container = av.open(tmp.name, mode="w") try: - import imageio as iio # type: ignore[no-redef] - except ImportError: - raise ImportError( - "imageio is required for video encoding. " - "Install with: pip install imageio[ffmpeg]" - ) - - logger.info(f"Encoding {len(frames)} frames to {output_format} bytes at {fps} fps") - - try: - buffer = io.BytesIO() - - kwargs: dict = {"fps": fps} - if output_format == "webm": - kwargs["codec"] = "libvpx-vp9" - elif output_format == "mp4": - kwargs["codec"] = "libx264" - else: - raise ValueError(f"No codec specified for response format: {output_format}") - - if hasattr(iio, "imwrite"): - # v3 API - iio.imwrite(buffer, frames, extension=f".{output_format}", **kwargs) - else: - # v2 API - writer = iio.get_writer( # type: ignore[attr-defined] - buffer, format="FFMPEG", mode="I", **kwargs - ) + stream = container.add_stream(codec, rate=fps) + stream.width = width + stream.height = height + stream.pix_fmt = "yuv420p" + stream.codec_context.thread_count = video_encode_threads() + if codec == "libx264": + stream.options = {"crf": "18", "preset": "veryfast"} + # BT.709 color tags so players don't render the clip washed-out. try: - for frame in frames: - writer.append_data(frame) - finally: - writer.close() - - video_bytes = buffer.getvalue() + cc = stream.codec_context + cc.color_primaries = 1 # AVCOL_PRI_BT709 + cc.color_trc = 1 # AVCOL_TRC_BT709 + cc.colorspace = 1 # AVCOL_SPC_BT709 + cc.color_range = 1 # AVCOL_RANGE_MPEG (limited / "tv") + except Exception: # noqa: BLE001 + logger.warning("could not set BT.709 color tags on the stream") + for i in range(num_frames): + frame = av.VideoFrame.from_ndarray(frames[i], format="rgb24") + for packet in stream.encode(frame): + container.mux(packet) + for packet in stream.encode(): # flush encoder + container.mux(packet) + finally: + container.close() + with open(tmp.name, "rb") as fh: + video_bytes = fh.read() logger.info(f"Encoded video to {len(video_bytes)} bytes") return video_bytes - except Exception as e: logger.error(f"Failed to encode video to bytes: {e}") raise RuntimeError(f"Video encoding to bytes failed: {e}") from e + finally: + try: + os.unlink(tmp.name) + except OSError: + pass diff --git a/components/src/dynamo/trtllm/backend_args.py b/components/src/dynamo/trtllm/backend_args.py index 8c345e78dba7..8558adb24c5e 100644 --- a/components/src/dynamo/trtllm/backend_args.py +++ b/components/src/dynamo/trtllm/backend_args.py @@ -277,6 +277,15 @@ def _add_diffusion_arguments(self, parser: argparse.ArgumentParser) -> None: arg_type=float, help="TeaCache threshold.", ) + add_argument( + diffusion_group, + flag_name="--cache-backend", + env_var="DYN_TRTLLM_CACHE_BACKEND", + default="none", + arg_type=str, + choices=["none", "teacache", "cache_dit"], + help="Step-caching backend: none, teacache, or cache_dit (Wan 2.2 supports cache_dit, not teacache).", + ) add_argument( diffusion_group, flag_name="--torch-dtype", @@ -449,6 +458,24 @@ def _add_diffusion_request_arguments(self, parser: argparse.ArgumentParser) -> N arg_type=float, help="Default CFG guidance scale.", ) + add_argument( + diffusion_request_group, + flag_name="--default-guidance-scale-2", + env_var="DYN_TRTLLM_DEFAULT_GUIDANCE_SCALE_2", + default=None, + arg_type=float, + help="Default second-stage CFG guidance scale (Wan 2.2 MoE low-noise " + "expert). None = use the pipeline default (single guidance).", + ) + add_argument( + diffusion_request_group, + flag_name="--default-boundary-ratio", + env_var="DYN_TRTLLM_DEFAULT_BOUNDARY_RATIO", + default=None, + arg_type=float, + help="Default timestep boundary ratio for switching guidance scales " + "(Wan 2.2). None = use the pipeline/model default.", + ) # Video specific args add_argument( diffusion_request_group, @@ -506,8 +533,11 @@ class DynamoTrtllmConfig(ConfigBase): default_num_images_per_prompt: int default_num_inference_steps: int default_guidance_scale: float + default_guidance_scale_2: Optional[float] = None + default_boundary_ratio: Optional[float] = None torch_dtype: str revision: Optional[str] = None + cache_backend: str enable_teacache: bool teacache_use_ret_steps: bool teacache_thresh: float diff --git a/components/src/dynamo/trtllm/configs/diffusion_config.py b/components/src/dynamo/trtllm/configs/diffusion_config.py index ffabe3f109ae..7a1340ebd4d4 100644 --- a/components/src/dynamo/trtllm/configs/diffusion_config.py +++ b/components/src/dynamo/trtllm/configs/diffusion_config.py @@ -70,6 +70,11 @@ class DiffusionConfig: default_seconds: int = 4 # Default video duration when only fps is specified default_num_inference_steps: int = 50 default_guidance_scale: float = 5.0 + # Second-stage guidance (Wan 2.2 MoE low-noise expert) + the timestep boundary at + # which guidance switches from stage-1 to stage-2. None → pipeline/model default + # (i.e. single guidance). Set both to run Wan 2.2 dual guidance (e.g. 4.0 / 3.0 @ 0.875). + default_guidance_scale_2: Optional[float] = None + default_boundary_ratio: Optional[float] = None # ── Pipeline optimization config (maps to PipelineConfig) ── disable_torch_compile: bool = False @@ -97,6 +102,13 @@ class DiffusionConfig: # Enable dynamic weight quantization (quantize BF16 weights on-the-fly during loading) quant_dynamic: bool = True + # ── Cache-acceleration backend ── + # Step-caching backend: "none" (default), "teacache", or "cache_dit". cache_dit + # (DBCache/TaylorSeer/SCM) is the ~2x step-skip Wan 2.2 supports -- it REJECTS + # teacache -- and VisualGen's CacheDiTConfig defaults are already tuned for + # few-step runs, so no extra knobs are needed here. + cache_backend: str = "none" + # ── TeaCache optimization config (maps to TeaCacheConfig) ── enable_teacache: bool = False teacache_use_ret_steps: bool = True @@ -149,6 +161,10 @@ def __str__(self) -> str: f"default_num_frames={self.default_num_frames}, " f"default_num_images_per_prompt={self.default_num_images_per_prompt}, " f"default_num_inference_steps={self.default_num_inference_steps}, " + f"default_guidance_scale={self.default_guidance_scale}, " + f"default_guidance_scale_2={self.default_guidance_scale_2}, " + f"default_boundary_ratio={self.default_boundary_ratio}, " + f"cache_backend={self.cache_backend}, " f"enable_teacache={self.enable_teacache}, " f"attn_backend={self.attn_backend}, " f"quant_algo={self.quant_algo}, " diff --git a/components/src/dynamo/trtllm/engines/diffusion_engine.py b/components/src/dynamo/trtllm/engines/diffusion_engine.py index fdacb0ebe09c..b2bb49d20f59 100644 --- a/components/src/dynamo/trtllm/engines/diffusion_engine.py +++ b/components/src/dynamo/trtllm/engines/diffusion_engine.py @@ -159,7 +159,10 @@ def _build_diffusion_args(self) -> "VisualGenArgs": TorchCompileConfig, VisualGenArgs, ) - from tensorrt_llm._torch.visual_gen.config import AttentionConfig + from tensorrt_llm._torch.visual_gen.config import ( + AttentionConfig, + CacheDiTConfig, + ) # Build quant_config dict if quantization is requested # VisualGenArgs accepts a dict in ModelOpt format and parses it via model_validator @@ -202,7 +205,12 @@ def _build_diffusion_args(self) -> "VisualGenArgs": ) # Add optional fields - if self.config.enable_teacache: + if self.config.cache_backend == "cache_dit": + # DBCache/TaylorSeer/SCM step-skip. VisualGen's CacheDiTConfig defaults + # are tuned for few-step runs (warmup 4, L1 thresh 0.24, <=3 continuous + # cached) -> ~2x skip. Wan 2.2 supports cache_dit and rejects teacache. + args_kwargs["cache"] = CacheDiTConfig() + elif self.config.enable_teacache or self.config.cache_backend == "teacache": args_kwargs["cache"] = TeaCacheConfig( use_ret_steps=self.config.teacache_use_ret_steps, teacache_thresh=self.config.teacache_thresh, @@ -224,6 +232,8 @@ def generate( num_images_per_prompt: int = 1, num_inference_steps: int = 50, guidance_scale: float = 5.0, + guidance_scale_2: Optional[float] = None, + boundary_ratio: Optional[float] = None, seed: Optional[int] = None, ) -> "MediaOutput": """Generate video/image frames from text prompt. @@ -286,6 +296,18 @@ def generate( params=params, ) + # Wan 2.2 dual-guidance overrides. guidance_scale_2 / boundary_ratio are + # extra_param_specs that default to None (single guidance); set them explicitly + # here — before the spec-default merge below (setdefault won't clobber a value + # already present) — so goff (e.g. 4.0/3.0 @ boundary 0.875) actually takes effect. + if guidance_scale_2 is not None or boundary_ratio is not None: + if req.params.extra_params is None: + req.params.extra_params = {} + if guidance_scale_2 is not None: + req.params.extra_params["guidance_scale_2"] = guidance_scale_2 + if boundary_ratio is not None: + req.params.extra_params["boundary_ratio"] = boundary_ratio + # Replicate the TRTLLM's visual_gen executor's _merge_defaults: fill None fields in # req.params with pipeline-specific defaults (universal + extra_param # specs), since we call pipeline.infer() directly instead of going diff --git a/components/src/dynamo/trtllm/request_handlers/diffusion/video_handler.py b/components/src/dynamo/trtllm/request_handlers/diffusion/video_handler.py index d1f73846c964..41de45a0d52c 100644 --- a/components/src/dynamo/trtllm/request_handlers/diffusion/video_handler.py +++ b/components/src/dynamo/trtllm/request_handlers/diffusion/video_handler.py @@ -205,10 +205,20 @@ async def generate( if nvext.guidance_scale is not None else self.config.default_guidance_scale ) - + # Dual-guidance / boundary are config-only (not carried in the /v1/videos + # nvext); None → engine/pipeline default (single guidance). + guidance_scale_2 = self.config.default_guidance_scale_2 + boundary_ratio = self.config.default_boundary_ratio + + guidance_desc = ( + f"{guidance_scale}/{guidance_scale_2}@{boundary_ratio}" + if guidance_scale_2 is not None + else f"{guidance_scale} (single)" + ) logger.info( f"Request {request_id}: prompt='{req.prompt[:50]}...', " - f"size={width}x{height}, frames={num_frames}, steps={num_inference_steps}" + f"size={width}x{height}, frames={num_frames}, steps={num_inference_steps}, " + f"guidance={guidance_desc}" ) # Run generation in thread pool (blocking operation). @@ -227,6 +237,8 @@ async def generate( num_frames=num_frames, num_inference_steps=num_inference_steps, guidance_scale=guidance_scale, + guidance_scale_2=guidance_scale_2, + boundary_ratio=boundary_ratio, seed=nvext.seed, ) @@ -234,7 +246,11 @@ async def generate( raise RuntimeError("Pipeline returned no output (MediaOutput is None)") # Determine output format - response_format = req.response_format or "url" + # Default to inline b64_json: deepinfra's deepapi consumer reads the video from + # data[0].b64_json and never sends response_format, so a "url" default would return + # a file reference it can't read (-> ERR_MODEL "No video data"). Callers that want a + # hosted URL can still pass response_format="url" explicitly. + response_format = req.response_format or "b64_json" if response_format not in ("url", "b64_json"): raise ValueError( f"Unsupported response_format: {response_format!r}; expected 'url' or 'b64_json'" @@ -248,13 +264,18 @@ async def generate( # Encode media based on what the pipeline returned if output.video is not None: - # MediaOutput.video is (B, T, H, W, C) uint8 since TRT-LLM rc9; - # squeeze the batch dim to get (T, H, W, C) for MP4 encoding. + # MediaOutput.video is (B, T, H, W, C) since TRT-LLM rc9; accept the + # batch-1 5D form or a bare 4D (T, H, W, C). Explicit error (not a bare + # assert, which -O strips) on anything unexpected. video = output.video - assert ( - video.ndim == 5 and video.shape[0] == 1 - ), f"Expected video shape (1, T, H, W, C), got {video.shape}" - frames_np = video[0].cpu().numpy() + if video.ndim == 5 and video.shape[0] == 1: + video = video[0] + elif video.ndim != 4: + raise RuntimeError( + f"Unexpected video tensor shape {tuple(video.shape)}; " + "expected (T, H, W, C) or (1, T, H, W, C)" + ) + frames_np = video.cpu().numpy() logger.info( f"Request {request_id}: encoding video output " f"(shape={frames_np.shape}) to MP4 at {fps} fps" diff --git a/components/src/dynamo/trtllm/tests/test_trtllm_video_diffusion.py b/components/src/dynamo/trtllm/tests/test_trtllm_video_diffusion.py index 21b8e50fd25f..6a68cc82988c 100644 --- a/components/src/dynamo/trtllm/tests/test_trtllm_video_diffusion.py +++ b/components/src/dynamo/trtllm/tests/test_trtllm_video_diffusion.py @@ -846,8 +846,10 @@ async def test_b64_response_format(self): assert decoded == b"fake_mp4_bytes" @pytest.mark.asyncio - async def test_default_response_format_is_url(self): - """Test that generate() defaults to url response format.""" + async def test_default_response_format_is_b64_json(self): + """generate() defaults to b64_json: deepinfra's deepapi reads the clip from + data[0].b64_json and never sends response_format, so a "url" default returned a + reference it couldn't read (-> ERR_MODEL "No video data").""" handler = self._make_handler() request = { @@ -868,9 +870,49 @@ async def test_default_response_format_is_url(self): results.append(result) assert len(results) == 1 - # Default should be "url" format, so upload_to_fs should be called - mock_upload.assert_called_once() - assert results[0]["data"][0]["url"] is not None + # Default is b64_json now: no upload, clip returned inline. + mock_upload.assert_not_called() + assert results[0]["data"][0]["b64_json"] is not None + assert results[0]["data"][0].get("url") is None + + @pytest.mark.asyncio + async def test_guidance_defaults_forwarded_to_engine(self): + """The handler forwards config default_guidance_scale_2 / default_boundary_ratio + to engine.generate() -- dual guidance isn't carried in the /v1/videos nvext.""" + from dynamo.trtllm.request_handlers.diffusion.video_handler import ( + VideoGenerationHandler, + ) + + mock_output = SimpleNamespace( + video=torch.zeros((1, 4, 64, 64, 3), dtype=torch.uint8), + image=None, + audio=None, + ) + mock_engine = MagicMock() + mock_engine.generate = MagicMock(return_value=mock_output) + config = DiffusionConfig( + media_output_fs_url="file:///tmp/test_media", + default_guidance_scale_2=3.0, + default_boundary_ratio=0.875, + ) + with patch( + "dynamo.trtllm.request_handlers.diffusion.video_handler.get_fs", + return_value=MagicMock(), + ): + handler = VideoGenerationHandler(engine=mock_engine, config=config) + + with patch( + "dynamo.trtllm.request_handlers.diffusion.video_handler.encode_to_video_bytes", + return_value=b"fake_mp4", + ): + async for _ in handler.generate( + {"prompt": "p", "model": "m", "response_format": "b64_json"}, + MagicMock(), + ): + pass + + assert mock_engine.generate.call_args.kwargs["guidance_scale_2"] == 3.0 + assert mock_engine.generate.call_args.kwargs["boundary_ratio"] == 0.875 @pytest.mark.asyncio async def test_error_response_on_failure(self): diff --git a/components/src/dynamo/trtllm/workers/video_diffusion_worker.py b/components/src/dynamo/trtllm/workers/video_diffusion_worker.py index 417adc4207ed..48af300d670e 100644 --- a/components/src/dynamo/trtllm/workers/video_diffusion_worker.py +++ b/components/src/dynamo/trtllm/workers/video_diffusion_worker.py @@ -33,6 +33,12 @@ async def init_video_diffusion_worker( shutdown_event: Event to signal shutdown. shutdown_endpoints: Optional list to populate with endpoints for graceful shutdown. """ + from dynamo.common.utils.video_utils import limit_video_worker_threads + + # Good-neighbor CPU budget on shared GPU nodes (bounds torch's intra-op pool; the + # encoder is bounded via the same budget). Set before the engine/torch load. + limit_video_worker_threads() + # Check tensorrt_llm visual_gen availability early with a clear error message. # visual_gen is part of TensorRT-LLM (tensorrt_llm._torch.visual_gen). # Without this check, users would get a cryptic ImportError deep inside @@ -67,6 +73,19 @@ async def init_video_diffusion_worker( # Build DiffusionConfig from the main Config diffusion_config = DiffusionConfig.from_config(config, skip_components) + # Make the guidance mode explicit in the logs: dual-guidance (goff) is driven by + # config-only flags (nvext can't carry them), so a missing --default-guidance-scale-2 + # silently falls back to single guidance. Log it prominently instead of hiding it in + # the full config dump. + logging.info( + "Video guidance: %s (g=%s, g2=%s, boundary=%s); cache_backend=%s", + "DUAL" if diffusion_config.default_guidance_scale_2 is not None else "SINGLE", + diffusion_config.default_guidance_scale, + diffusion_config.default_guidance_scale_2, + diffusion_config.default_boundary_ratio, + diffusion_config.cache_backend, + ) + # Get the endpoint from the runtime endpoint = runtime.endpoint( f"{config.namespace}.{config.component}.{config.endpoint}"