diff --git a/tests/diffusion/test_diffusion_scheduler.py b/tests/diffusion/test_diffusion_scheduler.py index 0a317f1fb..d6d17c18f 100644 --- a/tests/diffusion/test_diffusion_scheduler.py +++ b/tests/diffusion/test_diffusion_scheduler.py @@ -276,12 +276,16 @@ class TestGetRequestBatchSamplingParamsKey: generator: torch.Generator | None = None, extra_args: dict | None = None, condition_key: tuple | None = None, + guidance_scale: float | None = None, + guidance_scale_2: float | None = None, ) -> OmniDiffusionRequest: sp = OmniDiffusionSamplingParams( num_inference_steps=num_inference_steps, seed=seed, generator=generator, extra_args=extra_args or {}, + guidance_scale=guidance_scale, + guidance_scale_2=guidance_scale_2, ) return OmniDiffusionRequest( prompt="prompt", @@ -356,6 +360,21 @@ class TestGetRequestBatchSamplingParamsKey: self._make(condition_key=("wan22_s2v_condition", True)) ) != scheduler._build_sampling_params_key(self._make(condition_key=("wan22_s2v_condition", False))) + def test_distinguishes_explicit_guidance_scale_2(self) -> None: + # An omitted guidance_scale_2 is auto-filled from guidance_scale, so a + # request that omits it and one that passes the same value explicitly end + # up with an identical numeric guidance_scale_2 but different + # guidance_scale_2_provided. Pipelines read guidance_scale_2_provided from + # the batch's first request to gate image guidance, so the two must not + # share a request batch. + scheduler = RequestScheduler() + omitted = self._make(guidance_scale=2.0) + explicit = self._make(guidance_scale=2.0, guidance_scale_2=2.0) + + assert omitted.sampling_params.guidance_scale_2 == explicit.sampling_params.guidance_scale_2 + assert omitted.sampling_params.guidance_scale_2_provided != explicit.sampling_params.guidance_scale_2_provided + assert scheduler._build_sampling_params_key(omitted) != scheduler._build_sampling_params_key(explicit) + class TestRequestScheduler: def setup_method(self) -> None: diff --git a/vllm_omni/diffusion/sched/interface.py b/vllm_omni/diffusion/sched/interface.py index d97b5cbd0..122a97ec2 100644 --- a/vllm_omni/diffusion/sched/interface.py +++ b/vllm_omni/diffusion/sched/interface.py @@ -119,6 +119,7 @@ class RequestBatchSamplingParamsKey: guidance_scale: float = 0.0 guidance_scale_provided: bool = False guidance_scale_2: float | None = None + guidance_scale_2_provided: bool = False guidance_rescale: float = 0.0 true_cfg_scale: float | None = None cfg_normalize: bool = False