From 567e3a69c48448f1b60c2d325a9d2d9672b3ae10 Mon Sep 17 00:00:00 2001 From: Johan de Ruiter Date: Thu, 23 Jul 2026 16:29:07 +0000 Subject: [PATCH 1/3] trtllm diffusion: add guidance_scale_2 / boundary_ratio to DiffusionConfig Wan 2.2 A14B is a two-expert MoE that needs dual guidance (e.g. 4.0/3.0 with a 0.875 timestep boundary). The diffusion worker only exposed a single --default-guidance-scale; guidance_scale_2 / boundary_ratio (Wan extra_param_specs) defaulted to None => single guidance. Add: - DiffusionConfig.default_guidance_scale_2 / default_boundary_ratio + the CLI flags --default-guidance-scale-2 / --default-boundary-ratio (backend_args). - DiffusionEngine.generate() injects them into req.params.extra_params BEFORE the spec-default merge so they take effect. - video_handler forwards the config defaults (config-only; the /v1/videos nvext doesn't carry them). Both default to None, so existing models (Flux, Wan single-guidance) are unaffected. Co-Authored-By: Claude Opus 4.8 (1M context) --- components/src/dynamo/trtllm/backend_args.py | 20 +++++++++++++++++++ .../dynamo/trtllm/configs/diffusion_config.py | 5 +++++ .../dynamo/trtllm/engines/diffusion_engine.py | 14 +++++++++++++ .../diffusion/video_handler.py | 6 ++++++ 4 files changed, 45 insertions(+) diff --git a/components/src/dynamo/trtllm/backend_args.py b/components/src/dynamo/trtllm/backend_args.py index 8c345e78dba7..d4cdbc31003c 100644 --- a/components/src/dynamo/trtllm/backend_args.py +++ b/components/src/dynamo/trtllm/backend_args.py @@ -449,6 +449,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,6 +524,8 @@ 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 enable_teacache: bool diff --git a/components/src/dynamo/trtllm/configs/diffusion_config.py b/components/src/dynamo/trtllm/configs/diffusion_config.py index ffabe3f109ae..5357a3e34cd0 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 diff --git a/components/src/dynamo/trtllm/engines/diffusion_engine.py b/components/src/dynamo/trtllm/engines/diffusion_engine.py index fdacb0ebe09c..d453185d9cdc 100644 --- a/components/src/dynamo/trtllm/engines/diffusion_engine.py +++ b/components/src/dynamo/trtllm/engines/diffusion_engine.py @@ -224,6 +224,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 +288,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..541cd3bd0c52 100644 --- a/components/src/dynamo/trtllm/request_handlers/diffusion/video_handler.py +++ b/components/src/dynamo/trtllm/request_handlers/diffusion/video_handler.py @@ -205,6 +205,10 @@ 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 logger.info( f"Request {request_id}: prompt='{req.prompt[:50]}...', " @@ -227,6 +231,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, ) From 6a660a66a78a773f8286c43668d092638bd4db44 Mon Sep 17 00:00:00 2001 From: Johan de Ruiter Date: Fri, 24 Jul 2026 21:46:25 +0000 Subject: [PATCH 2/3] trtllm diffusion: default video response to inline b64_json The video handler defaulted response_format to "url", which uploads the clip to media_output_fs_url and returns a file reference. deepinfra's deepapi reads the video from data[0].b64_json and never sends response_format, so it received no video (ERR_MODEL "No video data"). Default to b64_json so the clip is returned inline; callers can still pass response_format="url" explicitly. Co-Authored-By: Claude Opus 4.8 --- .../trtllm/request_handlers/diffusion/video_handler.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) 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 541cd3bd0c52..8de0d6457dd3 100644 --- a/components/src/dynamo/trtllm/request_handlers/diffusion/video_handler.py +++ b/components/src/dynamo/trtllm/request_handlers/diffusion/video_handler.py @@ -240,7 +240,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'" From b4b4a3568aba960a849332acf6f7ef1f94fca964 Mon Sep 17 00:00:00 2001 From: Johan de Ruiter Date: Sat, 25 Jul 2026 03:57:51 +0000 Subject: [PATCH 3/3] trtllm diffusion tests: fix default-response-format test for b64 + cover guidance forwarding The b64_json default change flipped VideoGenerationHandler's default from "url" to "b64_json", which broke test_default_response_format_is_url (it asserted upload_to_fs was called and data[0].url was set). Rewrote it to test_default_response_format_is_b64_json (no upload; data[0].b64_json set) so the new default is pinned. Also added test_guidance_defaults_forwarded_to_engine: the handler must forward config.default_guidance_scale_2 / default_boundary_ratio to engine.generate() (dual guidance isn't carried in the /v1/videos nvext) -- the core of this PR had no test. Co-Authored-By: Claude Opus 4.8 --- .../tests/test_trtllm_video_diffusion.py | 49 +++++++++++++++++-- 1 file changed, 44 insertions(+), 5 deletions(-) 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..18a3f5b54c9c 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,46 @@ 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):