diff --git a/tests/diffusion/models/wan2_2/conftest.py b/tests/diffusion/models/wan2_2/conftest.py index f836fa545..cc95914c0 100644 --- a/tests/diffusion/models/wan2_2/conftest.py +++ b/tests/diffusion/models/wan2_2/conftest.py @@ -1,3 +1,6 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + from __future__ import annotations from contextlib import contextmanager @@ -33,7 +36,7 @@ class StubScheduler: self.config = SimpleNamespace(num_train_timesteps=1000) self.set_timesteps_calls: list[tuple[int, torch.device]] = [] - def set_timesteps(self, num_steps: int, device: torch.device) -> None: + def set_timesteps(self, num_steps: int, device: torch.device, shift: float | None = None) -> None: self.set_timesteps_calls.append((num_steps, device)) diff --git a/tests/diffusion/models/wan2_2/fixtures/native_unipc_shift12.json b/tests/diffusion/models/wan2_2/fixtures/native_unipc_shift12.json new file mode 100644 index 000000000..9b6fd2433 --- /dev/null +++ b/tests/diffusion/models/wan2_2/fixtures/native_unipc_shift12.json @@ -0,0 +1,97 @@ +{ + "source_url": "https://github.com/Wan-Video/Wan2.2/blob/42bf4cfaa384bc21833865abc2f9e6c0e67233dc/wan/utils/fm_solvers_unipc.py", + "source_sha256": "0dec8c7ed17f6f2049275c6848113314da6ccec1c8db5bdc89df43c05c6038d9", + "git_head": "42bf4cfaa384bc21833865abc2f9e6c0e67233dc", + "num_train_timesteps": 1000, + "use_dynamic_shifting": false, + "reproduction": "Load FlowUniPCMultistepScheduler from source_url; construct with num_train_timesteps=1000, shift=1.0, use_dynamic_shifting=False; call set_timesteps(40, device='cpu', shift=12.0); record timesteps.tolist() and sigmas.tolist().", + "constructor_shift": 1.0, + "runtime_shift": 12.0, + "steps": 40, + "timesteps": [ + 999, + 997, + 995, + 993, + 990, + 988, + 985, + 982, + 979, + 976, + 972, + 969, + 965, + 961, + 956, + 952, + 947, + 941, + 936, + 929, + 922, + 915, + 907, + 898, + 888, + 877, + 865, + 852, + 837, + 819, + 799, + 776, + 749, + 717, + 678, + 631, + 571, + 492, + 386, + 235 + ], + "sigmas": [ + 0.9999166131019592, + 0.9977826476097107, + 0.9955461621284485, + 0.9931996464729309, + 0.9907346963882446, + 0.988142192363739, + 0.9854119420051575, + 0.9825326800346375, + 0.9794917702674866, + 0.9762753844261169, + 0.9728676676750183, + 0.9692512154579163, + 0.9654061794281006, + 0.961310088634491, + 0.9569376707077026, + 0.95225989818573, + 0.9472436308860779, + 0.9418508410453796, + 0.9360373020172119, + 0.9297519326210022, + 0.9229347705841064, + 0.9155154228210449, + 0.9074103832244873, + 0.8985199332237244, + 0.888724148273468, + 0.8778772950172424, + 0.865800678730011, + 0.8522725701332092, + 0.8370144367218018, + 0.8196719288825989, + 0.7997865080833435, + 0.7767546772956848, + 0.7497654557228088, + 0.717703104019165, + 0.6789888143539429, + 0.6313128471374512, + 0.5711563229560852, + 0.49288034439086914, + 0.38684695959091187, + 0.23510953783988953, + 0.0 + ], + "implementation": "Wan-Video/Wan2.2 wan/utils/fm_solvers_unipc.py" +} diff --git a/tests/diffusion/models/wan2_2/test_wan22_i2v_pipeline.py b/tests/diffusion/models/wan2_2/test_wan22_i2v_pipeline.py index 26e78ddab..70bb146ab 100644 --- a/tests/diffusion/models/wan2_2/test_wan22_i2v_pipeline.py +++ b/tests/diffusion/models/wan2_2/test_wan22_i2v_pipeline.py @@ -1,15 +1,16 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project from types import SimpleNamespace +import numpy as np import pytest import torch from PIL import Image from torch import nn from tests.diffusion.models.wan2_2.conftest import StubScheduler, StubTransformer, StubVAE, noop_progress_bar -from vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2 import _WAN_TEXT_ENCODER_OFFLOAD_PLAN +from vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2 import _WAN_TEXT_ENCODER_OFFLOAD_PLAN, build_wan_scheduler from vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2_i2v import ( Wan22I2VPipeline, get_wan22_i2v_post_process_func, @@ -50,7 +51,7 @@ def _make_i2v_pipeline(*, expand_timesteps: bool) -> Wan22I2VPipeline: def _make_i2v_sampling(**overrides): - values = { + values: dict[str, object] = { "height": 16, "width": 16, "num_frames": 5, @@ -275,16 +276,18 @@ def test_i2v_prepare_latents_preserves_batched_image_conditions() -> None: assert first_frame_mask.shape == (2, 1, 2, 2, 2) -def test_i2v_forward_batches_conditions_random_inputs_and_outputs(monkeypatch) -> None: +@pytest.mark.parametrize("solver", ["unipc", "euler"]) +@pytest.mark.parametrize("shift", [3.0, 5.0, 12.0]) +def test_i2v_forward_batches_conditions_random_inputs_and_outputs(monkeypatch, solver: str, shift: float) -> None: monkeypatch.setattr( "vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2_i2v.current_omni_platform", SimpleNamespace(is_available=lambda: False), ) pipeline = _make_i2v_pipeline(expand_timesteps=True) - pipeline.scheduler = StubScheduler([9]) - pipeline.od_config = SimpleNamespace(flow_shift=5.0) - pipeline._sample_solver = "unipc" - pipeline._flow_shift = 5.0 + pipeline.scheduler = build_wan_scheduler(solver, shift) + pipeline.od_config = SimpleNamespace(flow_shift=shift) + pipeline._sample_solver = solver + pipeline._flow_shift = shift pipeline.boundary_ratio = 0.875 pipeline.has_image_encoder = False pipeline._guidance_scale = None @@ -327,6 +330,8 @@ def test_i2v_forward_batches_conditions_random_inputs_and_outputs(monkeypatch) - generator=gen_a, latents=latents_a, num_outputs_per_prompt=2, + num_inference_steps=50, + extra_args={"sample_solver": solver, "flow_shift": shift}, ), ), SimpleNamespace( @@ -336,12 +341,25 @@ def test_i2v_forward_batches_conditions_random_inputs_and_outputs(monkeypatch) - generator=gen_b, latents=latents_b, num_outputs_per_prompt=2, + num_inference_steps=50, + extra_args={"sample_solver": solver, "flow_shift": shift}, ), ), ] ) outputs = pipeline.forward(batch) + if solver == "unipc": + sigmas = np.linspace(float(np.float32(0.999)), 0.0, 51)[:-1] + sigmas = shift * sigmas / (1.0 + (shift - 1.0) * sigmas) + torch.testing.assert_close( + pipeline.scheduler.sigmas, torch.tensor(np.append(sigmas, 0.0), dtype=torch.float32), rtol=0, atol=0 + ) + assert pipeline.scheduler.timesteps.tolist() == (sigmas * 1000).astype(np.int64).tolist() + else: + reference = build_wan_scheduler("euler", shift) + reference.set_timesteps(50, device="cpu") + torch.testing.assert_close(pipeline.scheduler.sigmas, reference.sigmas, rtol=0, atol=0) assert prepare_call["batch_size"] == 4 assert prepare_call["generator"] == [gen_a, gen_a, gen_b, gen_b] 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 6f7833ff8..526eeeb15 100644 --- a/tests/diffusion/models/wan2_2/test_wan22_pipeline_diffuse.py +++ b/tests/diffusion/models/wan2_2/test_wan22_pipeline_diffuse.py @@ -10,7 +10,7 @@ import torch from torch import nn from vllm_omni.diffusion.media import VideoTensorEncoding, VideoTensorLayout, VideoValueRange -from vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2 import Wan22Pipeline +from vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2 import Wan22Pipeline, build_wan_scheduler from vllm_omni.diffusion.models.wan2_2.wan2_2_transformer import WanSelfAttention from vllm_omni.diffusion.request import OmniDiffusionRequest from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch @@ -19,6 +19,12 @@ from vllm_omni.inputs.data import OmniDiffusionSamplingParams pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion] +@pytest.fixture(autouse=True) +def _cpu_platform(monkeypatch): + module = importlib.import_module("vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2") + monkeypatch.setattr(module.current_omni_platform, "is_available", lambda: False) + + class _StubTransformer(nn.Module): @property def dtype(self) -> torch.dtype: @@ -51,10 +57,10 @@ class _StubScheduler: def __init__(self, timesteps: list[int]) -> None: self.timesteps = torch.tensor(timesteps, dtype=torch.int64) self.config = SimpleNamespace(num_train_timesteps=1000) - self.set_timesteps_calls: list[tuple[int, torch.device]] = [] + self.set_timesteps_calls: list[tuple[int, torch.device, float | None]] = [] - def set_timesteps(self, num_steps: int, device: torch.device) -> None: - self.set_timesteps_calls.append((num_steps, device)) + def set_timesteps(self, num_steps: int, device: torch.device, shift: float | None = None) -> None: + self.set_timesteps_calls.append((num_steps, device, shift)) @contextmanager @@ -184,7 +190,24 @@ def test_forward_delegates_denoising_to_diffuse( assert captured["boundary_timestep"] == pytest.approx(875.0) assert captured["latent_condition"] is None assert captured["first_frame_mask"] is None - assert pipeline.scheduler.set_timesteps_calls == [(2, torch.device("cpu"))] + assert pipeline.scheduler.set_timesteps_calls == [(2, torch.device("cpu"), 5.0)] + + +@pytest.mark.parametrize("solver", ["unipc", "euler"]) +def test_forward_passes_request_shift_without_mutating_scheduler_config(solver: str) -> None: + pipeline = _make_pipeline() + pipeline.diffuse = lambda **kwargs: kwargs["latents"] + for shift in (3.0, 12.0, 5.0): + sampling = _make_sampling(num_inference_steps=5, extra_args={"sample_solver": solver, "flow_shift": shift}) + request = OmniDiffusionRequest(prompt="prompt", request_id="schedule", sampling_params=sampling) + pipeline.forward(DiffusionRequestBatch(requests=[request])) + reference = build_wan_scheduler(solver, shift) + if solver == "unipc": + reference.set_timesteps(5, device="cpu", shift=shift) + assert pipeline.scheduler.config.shift == pipeline.scheduler.config["shift"] == 1.0 + else: + reference.set_timesteps(5, device="cpu") + torch.testing.assert_close(pipeline.scheduler.sigmas, reference.sigmas, rtol=0, atol=0) def test_forward_batches_text_generators_latents_and_splits_outputs() -> None: diff --git a/tests/diffusion/models/wan2_2/test_wan22_s2v_pipeline.py b/tests/diffusion/models/wan2_2/test_wan22_s2v_pipeline.py index 3b49d7edb..94f8b86fd 100644 --- a/tests/diffusion/models/wan2_2/test_wan22_s2v_pipeline.py +++ b/tests/diffusion/models/wan2_2/test_wan22_s2v_pipeline.py @@ -1,3 +1,6 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + from types import SimpleNamespace from unittest.mock import MagicMock, patch @@ -7,6 +10,7 @@ import pytest import torch import torch.nn as nn +from vllm_omni.diffusion.models.schedulers import FlowUniPCMultistepScheduler from vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2_s2v import ( Wan22S2VPipeline, _make_clip_generators, @@ -14,11 +18,32 @@ from vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2_s2v import ( from vllm_omni.diffusion.models.wan2_2.wan2_2_s2v_transformer import WanS2VTransformer3DModel from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch -pytestmark = [pytest.mark.core_model, pytest.mark.cpu] +pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion] + + +@pytest.mark.parametrize("configured_shift", [None, 12.0]) +def test_s2v_constructor_preserves_unshifted_endpoints(configured_shift: float | None) -> None: + module = "vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2_s2v" + config = SimpleNamespace(model="unused", flow_shift=configured_shift, enable_diffusion_pipeline_profiler=False) + + def init_components(pipeline, *args): + pipeline.vae = SimpleNamespace(config=SimpleNamespace(scale_factor_temporal=4, scale_factor_spatial=8)) + + with ( + patch(f"{module}.get_local_device", return_value=torch.device("cpu")), + patch(f"{module}._resolve_model_path", return_value="unused"), + patch(f"{module}._is_diffusers_format", return_value=True), + patch.object(Wan22S2VPipeline, "_init_diffusers_format", init_components), + patch.object(Wan22S2VPipeline, "setup_diffusion_pipeline_profiler"), + ): + pipeline = Wan22S2VPipeline(od_config=config) + assert pipeline._flow_shift == (3.0 if configured_shift is None else configured_shift) + assert pipeline.scheduler.config.shift == pipeline.scheduler.config["shift"] == 1.0 + assert pipeline.scheduler.sigma_max == float(np.float32(0.999)) def _make_s2v_sampling(**overrides): - values = { + values: dict[str, object] = { "height": 16, "width": 16, "num_frames": 8, @@ -211,7 +236,8 @@ def test_s2v_forward_rejects_different_raw_audio_shapes() -> None: pipeline.forward(batch) -def test_s2v_forward_batches_request_local_inputs_and_splits_outputs() -> None: +@pytest.mark.parametrize("shift", [3.0, 12.0]) +def test_s2v_forward_batches_request_local_inputs_and_splits_outputs(shift: float) -> None: pipeline = object.__new__(Wan22S2VPipeline) nn.Module.__init__(pipeline) pipeline.device = torch.device("cpu") @@ -227,7 +253,8 @@ def test_s2v_forward_batches_request_local_inputs_and_splits_outputs() -> None: enable_cpu_offload=False, parallel_config=SimpleNamespace(use_hsdp=False), ) - pipeline.scheduler = MagicMock(timesteps=torch.tensor([1.0])) + pipeline._flow_shift = shift + pipeline.scheduler = FlowUniPCMultistepScheduler(shift=1.0) pipeline.vae_scale_factor_spatial = 8 pipeline.resolution_divisor = 16 pipeline.motion_frames = 7 @@ -262,6 +289,7 @@ def test_s2v_forward_batches_request_local_inputs_and_splits_outputs() -> None: "multi_modal_data": {"image": image, "audio": audio_a}, }, sampling_params=_make_s2v_sampling( + num_inference_steps=5, num_outputs_per_prompt=2, generator=generators_a, latents=latents_a, @@ -275,6 +303,7 @@ def test_s2v_forward_batches_request_local_inputs_and_splits_outputs() -> None: "multi_modal_data": {"image": image, "audio": audio_b}, }, sampling_params=_make_s2v_sampling( + num_inference_steps=5, num_outputs_per_prompt=2, generator=generators_b, latents=latents_b, @@ -287,6 +316,12 @@ def test_s2v_forward_batches_request_local_inputs_and_splits_outputs() -> None: platform.is_available.return_value = False outputs = pipeline.forward(batch) + expected = np.linspace(float(np.float32(0.999)), 0.0, 6)[:-1] + expected = shift * expected / (1.0 + (shift - 1.0) * expected) + torch.testing.assert_close( + pipeline.scheduler.sigmas, torch.tensor(np.append(expected, 0.0), dtype=torch.float32), rtol=0, atol=0 + ) + assert pipeline.scheduler.config.shift == pipeline.scheduler.config["shift"] == 1.0 assert len(outputs) == 2 assert outputs[0].output[0].shape[0] == 2 assert outputs[1].output[0].shape[0] == 2 @@ -519,6 +554,7 @@ def test_s2v_pipeline_hsdp_forward_complete_process(): mock_scheduler.timesteps = torch.linspace(999, 0, 5) mock_scheduler.step = MagicMock(return_value=(torch.zeros(1, 16, 20, 88, 128),)) pipeline.scheduler = mock_scheduler + pipeline._flow_shift = 3.0 # -- Bind methods from the real class -- pipeline.encode_prompt = Wan22S2VPipeline.encode_prompt.__get__(pipeline) diff --git a/tests/diffusion/models/wan2_2/test_wan22_scheduler.py b/tests/diffusion/models/wan2_2/test_wan22_scheduler.py new file mode 100644 index 000000000..87cc24670 --- /dev/null +++ b/tests/diffusion/models/wan2_2/test_wan22_scheduler.py @@ -0,0 +1,54 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Compare actual Wan schedules to native Wan outputs and independent sigma arithmetic.""" + +import json +from pathlib import Path + +import numpy as np +import numpy.typing as npt +import pytest +import torch + +from vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2 import ( + FASTWAN_DMD_SCHEDULER_SHIFT, + build_wan_scheduler, +) + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion] + + +def test_wan_unipc_matches_captured_native_full_schedule() -> None: + # The fixture pins the official Wan source and records how to regenerate it. + golden = json.loads((Path(__file__).parent / "fixtures/native_unipc_shift12.json").read_text()) + scheduler = build_wan_scheduler("unipc", golden["runtime_shift"]) + for steps in (golden["steps"], 20, golden["steps"]): + scheduler.set_timesteps(steps, device="cpu", shift=golden["runtime_shift"]) + if steps == golden["steps"]: + assert scheduler.timesteps.tolist() == golden["timesteps"] + torch.testing.assert_close(scheduler.sigmas, torch.tensor(golden["sigmas"]), rtol=0, atol=0) + + +@pytest.mark.parametrize("shift", [1.0, 3.0, 5.0, 12.0]) +@pytest.mark.parametrize("steps", [20, 40]) +def test_wan_unipc_shifts_unmodified_training_endpoints(shift: float, steps: int) -> None: + # Native Wan: float32 training sigma_max=(1-1/1000), sigma_min=0; + # interpolate in float64, apply the runtime shift, then cast sigma to FP32. + unshifted = np.linspace(float(np.float32(0.999)), 0.0, steps + 1)[:-1] + shifted = shift * unshifted / (1.0 + (shift - 1.0) * unshifted) + expected_timesteps = (shifted * 1000).astype(np.int64) + expected_sigmas: npt.NDArray[np.float32] = np.append(shifted, 0.0).astype(np.float32) + scheduler = build_wan_scheduler("unipc", shift) + scheduler.set_timesteps(steps, device="cpu", shift=shift) + assert scheduler.config.shift == scheduler.config["shift"] == 1.0 + torch.testing.assert_close(scheduler.timesteps, torch.from_numpy(expected_timesteps), rtol=0, atol=0) + torch.testing.assert_close(scheduler.sigmas, torch.from_numpy(expected_sigmas), rtol=0, atol=0) + + +def test_wan_euler_keeps_dmd_shift_semantics() -> None: + scheduler = build_wan_scheduler("euler", FASTWAN_DMD_SCHEDULER_SHIFT) + scheduler.set_timesteps(3, device="cpu") + # The unchanged Euler branch starts at sigma=1, with shift=8 for DMD. + expected = torch.tensor([1.0, 16.0 / 17.0, 0.8, 0.0]) + torch.testing.assert_close(scheduler.sigmas, expected, rtol=0, atol=1e-7) + torch.testing.assert_close(scheduler.timesteps, expected[:-1] * 1000, rtol=0, atol=1e-4) diff --git a/tests/diffusion/models/wan2_2/test_wan22_vace_pipeline.py b/tests/diffusion/models/wan2_2/test_wan22_vace_pipeline.py index 5f712bc8a..75cf868ce 100644 --- a/tests/diffusion/models/wan2_2/test_wan22_vace_pipeline.py +++ b/tests/diffusion/models/wan2_2/test_wan22_vace_pipeline.py @@ -1,15 +1,17 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project from dataclasses import dataclass, field from types import SimpleNamespace +import numpy as np import pytest import torch from PIL import Image from torch import nn from tests.diffusion.models.wan2_2.conftest import StubScheduler, StubTransformer, StubVAE, noop_progress_bar +from vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2 import build_wan_scheduler from vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2_vace import ( Wan22VACEPipeline, create_vace_transformer_from_config, @@ -33,6 +35,8 @@ def test_vace_default_flow_shift_preserves_compact_offload_runtime_state(monkeyp ) enable_layerwise_offload: bool = True post_init_calls: int = 0 + _resolved_diffusion_offload: object = None + _diffusion_offload_flags_materialized: bool = False def __post_init__(self) -> None: self.post_init_calls += 1 @@ -41,7 +45,7 @@ def test_vace_default_flow_shift_preserves_compact_offload_runtime_state(monkeyp resolved_offload = object() config._resolved_diffusion_offload = resolved_offload config._diffusion_offload_flags_materialized = True - captured = {} + captured: dict[str, object] = {} def capture_base_init(self, *, od_config, prefix="") -> None: captured.update(od_config=od_config, prefix=prefix) @@ -51,6 +55,7 @@ def test_vace_default_flow_shift_preserves_compact_offload_runtime_state(monkeyp Wan22VACEPipeline(od_config=config, prefix="stage") forwarded = captured["od_config"] + assert isinstance(forwarded, RuntimeConfig) assert forwarded is not config assert config.flow_shift is None assert forwarded.flow_shift == 3.0 @@ -75,7 +80,7 @@ def _make_vace_pipeline() -> Wan22VACEPipeline: def _make_vace_sampling(**overrides): - values = { + values: dict[str, object] = { "height": 16, "width": 16, "num_frames": 5, @@ -123,7 +128,7 @@ def test_vace_preprocess_collects_reference_video_and_mask_inputs() -> None: def test_create_vace_transformer_from_config_maps_vace_specific_keys(monkeypatch) -> None: - captured = {} + captured: dict[str, object] = {} class FakeVACETransformer: def __init__(self, **kwargs) -> None: @@ -203,7 +208,9 @@ def test_vace_diffuse_passes_context_and_scale_to_cfg_branches() -> None: torch.testing.assert_close(result, torch.ones_like(latents)) -def test_vace_forward_batches_random_inputs_and_splits_outputs(monkeypatch) -> None: +@pytest.mark.parametrize("solver", ["unipc", "euler"]) +@pytest.mark.parametrize("shift", [3.0, 5.0, 12.0]) +def test_vace_forward_batches_random_inputs_and_splits_outputs(monkeypatch, solver: str, shift: float) -> None: monkeypatch.setattr( "vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2_vace.current_omni_platform", SimpleNamespace(is_available=lambda: False), @@ -211,10 +218,10 @@ def test_vace_forward_batches_random_inputs_and_splits_outputs(monkeypatch) -> N pipeline = _make_vace_pipeline() pipeline.transformer.vace_patch_embedding = None pipeline.transformer_2 = None - pipeline.scheduler = StubScheduler([9]) - pipeline.od_config = SimpleNamespace(flow_shift=3.0) - pipeline._sample_solver = "unipc" - pipeline._flow_shift = 3.0 + pipeline.scheduler = build_wan_scheduler(solver, shift) + pipeline.od_config = SimpleNamespace(flow_shift=shift) + pipeline._sample_solver = solver + pipeline._flow_shift = shift pipeline.boundary_ratio = None pipeline._guidance_scale = None pipeline._num_timesteps = None @@ -247,6 +254,8 @@ def test_vace_forward_batches_random_inputs_and_splits_outputs(monkeypatch) -> N generator=gen_a, latents=latents_a, num_outputs_per_prompt=2, + num_inference_steps=50, + extra_args={"sample_solver": solver, "flow_shift": shift}, ), ), SimpleNamespace( @@ -256,12 +265,25 @@ def test_vace_forward_batches_random_inputs_and_splits_outputs(monkeypatch) -> N generator=gen_b, latents=latents_b, num_outputs_per_prompt=2, + num_inference_steps=50, + extra_args={"sample_solver": solver, "flow_shift": shift}, ), ), ] ) outputs = pipeline.forward(batch) + if solver == "unipc": + sigmas = np.linspace(float(np.float32(0.999)), 0.0, 51)[:-1] + sigmas = shift * sigmas / (1.0 + (shift - 1.0) * sigmas) + torch.testing.assert_close( + pipeline.scheduler.sigmas, torch.tensor(np.append(sigmas, 0.0), dtype=torch.float32), rtol=0, atol=0 + ) + assert pipeline.scheduler.timesteps.tolist() == (sigmas * 1000).astype(np.int64).tolist() + else: + reference = build_wan_scheduler("euler", shift) + reference.set_timesteps(50, device="cpu") + torch.testing.assert_close(pipeline.scheduler.sigmas, reference.sigmas, rtol=0, atol=0) assert prepare_call["batch_size"] == 4 assert prepare_call["generator"] == [gen_a, gen_a, gen_b, gen_b] 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 a285e0009..15b036864 100644 --- a/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2.py +++ b/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2.py @@ -59,9 +59,11 @@ FASTWAN_DMD_SCHEDULER_SHIFT = 8.0 def build_wan_scheduler(sample_solver: str, flow_shift: float) -> Any: if sample_solver == "unipc": + # Keep native Wan's unshifted training endpoints; set_timesteps applies + # the requested shift to the interpolated inference sigmas. return FlowUniPCMultistepScheduler( num_train_timesteps=1000, - shift=flow_shift, + shift=1.0, prediction_type="flow_prediction", ) if sample_solver == "euler": @@ -774,7 +776,10 @@ class Wan22Pipeline( self._sample_solver = sample_solver self._flow_shift = flow_shift - self.scheduler.set_timesteps(num_steps, device=device) + if sample_solver == "unipc": + self.scheduler.set_timesteps(num_steps, device=device, shift=flow_shift) + else: + self.scheduler.set_timesteps(num_steps, device=device) timesteps = self.scheduler.timesteps self._num_timesteps = len(timesteps) boundary_timestep = None 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 908a35020..01bd53bb8 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 @@ -629,7 +629,10 @@ class Wan22I2VPipeline( self._flow_shift = flow_shift # Timesteps - self.scheduler.set_timesteps(num_steps, device=device) + if sample_solver == "unipc": + self.scheduler.set_timesteps(num_steps, device=device, shift=flow_shift) + else: + self.scheduler.set_timesteps(num_steps, device=device) timesteps = self.scheduler.timesteps self._num_timesteps = len(timesteps) diff --git a/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2_s2v.py b/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2_s2v.py index ea99e065f..1d28b425d 100644 --- a/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2_s2v.py +++ b/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2_s2v.py @@ -545,10 +545,10 @@ class Wan22S2VPipeline( self._init_original_format(model_path, dtype) # -- Scheduler -- - flow_shift = od_config.flow_shift if od_config.flow_shift is not None else 3.0 # S2V default + self._flow_shift = od_config.flow_shift if od_config.flow_shift is not None else 3.0 # S2V default self.scheduler = FlowUniPCMultistepScheduler( num_train_timesteps=1000, - shift=flow_shift, + shift=1.0, prediction_type="flow_prediction", ) @@ -1342,7 +1342,7 @@ class Wan22S2VPipeline( ) # -- Scheduler -- - self.scheduler.set_timesteps(num_steps, device=device) + self.scheduler.set_timesteps(num_steps, device=device, shift=self._flow_shift) timesteps = self.scheduler.timesteps self._num_timesteps = len(timesteps) diff --git a/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2_vace.py b/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2_vace.py index 0d48dffb5..aeecff608 100644 --- a/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2_vace.py +++ b/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2_vace.py @@ -690,7 +690,10 @@ class Wan22VACEPipeline(Wan22Pipeline, SupportImageInput): ) # Set up scheduler - self.scheduler.set_timesteps(num_inference_steps, device=device) + if self._sample_solver == "unipc": + self.scheduler.set_timesteps(num_inference_steps, device=device, shift=self._flow_shift) + else: + self.scheduler.set_timesteps(num_inference_steps, device=device) timesteps = self.scheduler.timesteps self._num_timesteps = len(timesteps)