[Bugfix] Fix Wan2.2 omitted guidance scale resolution (#5615)

Signed-off-by: lishunyang12 <lishunyang03@gmail.com>
Signed-off-by: SYLAR <125541396+lishunyang12@users.noreply.github.com>
Co-authored-by: lishunyang12 <lishunyang03@gmail.com>

tests/diffusion/models/wan2_2/test_wan22_pipeline_diffuse.py

@@ -9,7 +9,9 @@ import torch
 from torch import nn
 
 from vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2 import Wan22Pipeline
+from vllm_omni.diffusion.request import OmniDiffusionRequest
 from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch
+from vllm_omni.inputs.data import OmniDiffusionSamplingParams
 
 pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion]
 
@@ -91,7 +93,20 @@ def _make_pipeline() -> Wan22Pipeline:
     return pipeline
 
 
-def test_forward_delegates_denoising_to_diffuse() -> None:
+@pytest.mark.parametrize(
+    ("sampling_params_kwargs", "expected_low", "expected_high"),
+    [
+        ({}, 4.0, 4.0),
+        ({"guidance_scale": 0.0}, 0.0, 0.0),
+        ({"guidance_scale": 1.0}, 1.0, 1.0),
+        ({"guidance_scale": 3.0, "guidance_scale_2": 5.0}, 3.0, 5.0),
+    ],
+)
+def test_forward_delegates_denoising_to_diffuse(
+    sampling_params_kwargs: dict[str, float],
+    expected_low: float,
+    expected_high: float,
+) -> None:
     pipeline = _make_pipeline()
     captured: dict[str, object] = {}
 
@@ -101,25 +116,15 @@ def test_forward_delegates_denoising_to_diffuse() -> None:
 
     pipeline.diffuse = _fake_diffuse  # type: ignore[method-assign]
 
-    mock_req = SimpleNamespace(
+    mock_req = OmniDiffusionRequest(
         prompt="prompt",
         request_id="test-req",
-        sampling_params=SimpleNamespace(
-            height=None,
-            width=None,
+        sampling_params=OmniDiffusionSamplingParams(
             num_frames=1,
             num_inference_steps=2,
-            guidance_scale_provided=True,
-            guidance_scale=1.0,
-            guidance_scale_2=None,
-            boundary_ratio=None,
-            generator=None,
-            seed=None,
-            num_outputs_per_prompt=1,
             max_sequence_length=32,
-            latents=None,
             output_type="latent",
-            extra_args={},
+            **sampling_params_kwargs,
         ),
     )
     batch = DiffusionRequestBatch(requests=[mock_req])
@@ -129,8 +134,8 @@ def test_forward_delegates_denoising_to_diffuse() -> None:
     assert torch.equal(output.output, torch.ones((1, 4, 1, 8, 8)))
     assert torch.equal(captured["prompt_embeds"], torch.zeros(1, 32, 8))
     assert torch.equal(captured["timesteps"], pipeline.scheduler.timesteps)
-    assert captured["guidance_low"] == 1.0
-    assert captured["guidance_high"] == 1.0
+    assert captured["guidance_low"] == expected_low
+    assert captured["guidance_high"] == expected_high
     assert captured["boundary_timestep"] == pytest.approx(875.0)
     assert captured["latent_condition"] is None
     assert captured["first_frame_mask"] is None

vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2.py

@@ -37,7 +37,7 @@ from vllm_omni.diffusion.postprocess import interpolate_video_tensor
 from vllm_omni.diffusion.profiler.diffusion_pipeline_profiler import DiffusionPipelineProfilerMixin
 from vllm_omni.diffusion.request import OmniDiffusionRequest
 from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch
-from vllm_omni.inputs.data import OmniTextPrompt
+from vllm_omni.inputs.data import OmniDiffusionSamplingParams, OmniTextPrompt
 from vllm_omni.platforms import current_omni_platform
 
 logger = logging.getLogger(__name__)
@@ -84,6 +84,24 @@ def resolve_wan_flow_shift(req: OmniDiffusionRequest, od_config: OmniDiffusionCo
         raise ValueError(f"Invalid flow_shift={raw_flow_shift!r}. flow_shift must be a float.") from exc
 
 
+def resolve_wan_guidance_scales(
+    sampling_params: OmniDiffusionSamplingParams,
+    default_guidance_scale: float,
+) -> tuple[float, float]:
+    guidance_scale = (
+        sampling_params.guidance_scale if sampling_params.guidance_scale_provided else default_guidance_scale
+    )
+    guidance_low = guidance_scale if isinstance(guidance_scale, (int, float)) else guidance_scale[0]
+    guidance_high = (
+        sampling_params.guidance_scale_2
+        if sampling_params.guidance_scale_2_provided and sampling_params.guidance_scale_2 is not None
+        else (
+            guidance_scale[1] if isinstance(guidance_scale, (list, tuple)) and len(guidance_scale) > 1 else guidance_low
+        )
+    )
+    return guidance_low, guidance_high
+
+
 def retrieve_latents(
     encoder_output: torch.Tensor,
     generator: torch.Generator | None = None,
@@ -571,25 +589,10 @@ class Wan22Pipeline(
         width = (width // mod_value) * mod_value
         num_steps = 40 if req.sampling_params.num_inference_steps is None else req.sampling_params.num_inference_steps
 
-        # Respect per-request guidance_scale when explicitly provided.
-        if req.sampling_params.guidance_scale_provided:
-            guidance_scale = req.sampling_params.guidance_scale
-        else:
-            guidance_scale = 4.0
-
         output_type = req.sampling_params.output_type or "np"
         attention_kwargs: dict | None = None
 
-        guidance_low = guidance_scale if isinstance(guidance_scale, (int, float)) else guidance_scale[0]
-        guidance_high = (
-            req.sampling_params.guidance_scale_2
-            if req.sampling_params.guidance_scale_2 is not None
-            else (
-                guidance_scale[1]
-                if isinstance(guidance_scale, (list, tuple)) and len(guidance_scale) > 1
-                else guidance_low
-            )
-        )
+        guidance_low, guidance_high = resolve_wan_guidance_scales(req.sampling_params, default_guidance_scale=4.0)
 
         # record guidance for properties
         self._guidance_scale = guidance_low

vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2_i2v.py

@@ -36,6 +36,7 @@ from vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2 import (
     create_transformer_from_config,
     load_transformer_config,
     resolve_wan_flow_shift,
+    resolve_wan_guidance_scales,
     resolve_wan_sample_solver,
     retrieve_latents,
 )
@@ -475,26 +476,10 @@ class Wan22I2VPipeline(
         num_frames = req.sampling_params.num_frames or 81
         num_steps = 40 if req.sampling_params.num_inference_steps is None else req.sampling_params.num_inference_steps
 
-        # Respect per-request guidance_scale when explicitly provided.
-        if req.sampling_params.guidance_scale_provided:
-            guidance_scale = req.sampling_params.guidance_scale
-        else:
-            guidance_scale = 5.0
-
         output_type = req.sampling_params.output_type or "np"
         attention_kwargs: dict | None = None
 
-        # Handle guidance scales
-        guidance_low = guidance_scale if isinstance(guidance_scale, (int, float)) else guidance_scale[0]
-        guidance_high = (
-            req.sampling_params.guidance_scale_2
-            if req.sampling_params.guidance_scale_2 is not None
-            else (
-                guidance_scale[1]
-                if isinstance(guidance_scale, (list, tuple)) and len(guidance_scale) > 1
-                else guidance_low
-            )
-        )
+        guidance_low, guidance_high = resolve_wan_guidance_scales(req.sampling_params, default_guidance_scale=5.0)
 
         self._guidance_scale = guidance_low
         self._guidance_scale_2 = guidance_high

vllm_omni/diffusion/request.py

@@ -41,11 +41,10 @@ class OmniDiffusionRequest:
         if self.sampling_params.generator is None and self.sampling_params.seed is None:
             self.sampling_params.seed = random.randint(0, 2**31 - 1)
 
-        # Detect whether user explicitly provided guidance_scale.
-        # The unset sentinel is None (see OmniDiffusionSamplingParams); any value
-        # the caller set -- including an explicit 0.0 (disable CFG) -- must be
-        # honored. Resolve before auto-filling guidance_scale_2, else the sentinel
-        # leaks into guidance_scale_2.
+        # Detect whether the caller explicitly provided guidance_scale before
+        # resolving the omitted value.  ``0.0`` is API-valid and must not be
+        # treated as omission because that would unexpectedly enable CFG in
+        # pipelines with a model-specific default above zero.
         if self.sampling_params.guidance_scale is not None:
             self.sampling_params.guidance_scale_provided = True
         else: