diff --git a/tests/diffusion/models/wan2_2/test_wan22_pipeline_diffuse.py b/tests/diffusion/models/wan2_2/test_wan22_pipeline_diffuse.py index 198c31908..8c06df034 100644 --- a/tests/diffusion/models/wan2_2/test_wan22_pipeline_diffuse.py +++ b/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 diff --git a/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2.py b/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2.py index e3617f84d..0fc9aca65 100644 --- a/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2.py +++ b/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 diff --git a/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2_i2v.py b/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2_i2v.py index d8c1d4a8c..180f8bb7a 100644 --- a/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2_i2v.py +++ b/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 diff --git a/vllm_omni/diffusion/request.py b/vllm_omni/diffusion/request.py index 9ea879212..e1147c3d6 100644 --- a/vllm_omni/diffusion/request.py +++ b/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: