trtllm diffusion: dual guidance (guidance_scale_2 / boundary_ratio) for Wan 2.2

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

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

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

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,14 +231,20 @@ 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,
                 )
 
             if output is None:
                 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'"

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):