anodyine/vllm-omni
1
0 Fork 0
Code Issues 4 Pull Requests 2 Actions Packages Projects Releases Wiki Activity

[Bugfix][Wan2.2] Apply UniPC flow shift once (#7576)

Browse Source 
Signed-off-by: HAAZZZEEEE <wang-taicheng@qq.com>
This commit is contained in:
HAAZZZEEEE authored and GitHub committed 2026-09-16 10:06:41 +08:00
1 parent 2185c4e26f
commit 3eb5c3b4e6
11 files changed
+298 -34
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

tests/diffusion/models/wan2_2/conftest.py
+4 -1
View File
@@ -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))
tests/diffusion/models/wan2_2/fixtures/native_unipc_shift12.json
+97
View File
@@ -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"
}
tests/diffusion/models/wan2_2/test_wan22_i2v_pipeline.py
+26 -8
View File
@@ -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]
tests/diffusion/models/wan2_2/test_wan22_pipeline_diffuse.py
+28 -5
View File
@@ -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:
tests/diffusion/models/wan2_2/test_wan22_s2v_pipeline.py
+40 -4
View File
@@ -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)
tests/diffusion/models/wan2_2/test_wan22_scheduler.py
+54
View File
@@ -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)
tests/diffusion/models/wan2_2/test_wan22_vace_pipeline.py
+31 -9
View File
@@ -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]
vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2.py
+7 -2
View File
@@ -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
vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2_i2v.py
+4 -1
View File
@@ -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)
vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2_s2v.py
+3 -3
View File
@@ -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)
vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2_vace.py
+4 -1
View File
@@ -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)