[Model] feat: FastGen DMD2-distilled Wan 2.1 pipelines (T2V, I2V)

tests/diffusion/models/wan2_2/test_wan_dmd2_request_sanitization.py

@@ -0,0 +1,173 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+from unittest.mock import MagicMock, patch
+
+import pytest
+import torch
+
+from vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2 import Wan22Pipeline, WanT2VDMD2Pipeline
+from vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2_i2v import Wan22I2VPipeline, WanI2VDMD2Pipeline
+from vllm_omni.diffusion.request import OmniDiffusionRequest, OmniDiffusionSamplingParams
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+# Wan base pipeline whose __init__ loads model weights — mocked in tests.
+_WAN_BASE = {
+    WanT2VDMD2Pipeline: Wan22Pipeline,
+    WanI2VDMD2Pipeline: Wan22I2VPipeline,
+}
+
+
+def _make_pipeline(cls):
+    """Run the DMD2 __init__ with the Wan base mocked out (no model weights loaded)."""
+
+    base = _WAN_BASE[cls]
+    od_config = MagicMock()
+    od_config.model = "/nonexistent"
+
+    def _mock_base_init(self, *a, **kw):
+        self.od_config = od_config
+
+    with patch.object(base, "__init__", _mock_base_init):
+        pipeline = object.__new__(cls)
+        torch.nn.Module.__init__(pipeline)
+        cls.__init__(pipeline, od_config=od_config)
+    return pipeline
+
+
+def _make_request(prompts=None, **sp_kwargs) -> OmniDiffusionRequest:
+    sp = OmniDiffusionSamplingParams(**sp_kwargs)
+    return OmniDiffusionRequest(
+        prompts=prompts or [{"prompt": "a cat dancing"}],
+        sampling_params=sp,
+    )
+
+
+@pytest.fixture(params=[WanT2VDMD2Pipeline, WanI2VDMD2Pipeline], ids=["t2v", "i2v"])
+def pipeline(request):
+    return _make_pipeline(request.param)
+
+
+# ---------------------------------------------------------------------------
+# num_inference_steps
+# ---------------------------------------------------------------------------
+
+
+def test_num_inference_steps_forced_to_dmd2_value(pipeline):
+    req = _make_request(num_inference_steps=40)
+    pipeline._sanitize_dmd2_request(req)
+    assert req.sampling_params.num_inference_steps == pipeline.num_inference_steps
+
+
+def test_num_inference_steps_already_correct(pipeline):
+    req = _make_request(num_inference_steps=pipeline.num_inference_steps)
+    pipeline._sanitize_dmd2_request(req)
+    assert req.sampling_params.num_inference_steps == pipeline.num_inference_steps
+
+
+# ---------------------------------------------------------------------------
+# guidance_scale
+# ---------------------------------------------------------------------------
+
+
+def test_guidance_scale_forced_to_one(pipeline):
+    req = _make_request(guidance_scale=5.0, guidance_scale_provided=True)
+    pipeline._sanitize_dmd2_request(req)
+    assert req.sampling_params.guidance_scale == pipeline.dmd2_guidance_scale
+    assert req.sampling_params.guidance_scale_provided is False
+
+
+def test_guidance_scale_already_correct(pipeline):
+    req = _make_request(guidance_scale=pipeline.dmd2_guidance_scale, guidance_scale_provided=False)
+    pipeline._sanitize_dmd2_request(req)
+    assert req.sampling_params.guidance_scale == pipeline.dmd2_guidance_scale
+
+
+def test_guidance_scale_provided_flag_cleared(pipeline):
+    """guidance_scale_provided=True must be cleared even if scale is already dmd2_guidance_scale."""
+    req = _make_request(guidance_scale=pipeline.dmd2_guidance_scale, guidance_scale_provided=True)
+    pipeline._sanitize_dmd2_request(req)
+    assert req.sampling_params.guidance_scale_provided is False
+
+
+def test_guidance_scale_2_cleared(pipeline):
+    req = _make_request(guidance_scale_2=3.0)
+    pipeline._sanitize_dmd2_request(req)
+    assert req.sampling_params.guidance_scale_2 is None
+
+
+def test_guidance_scale_2_unset_unchanged(pipeline):
+    req = _make_request()
+    pipeline._sanitize_dmd2_request(req)
+    assert req.sampling_params.guidance_scale_2 is None
+
+
+def test_true_cfg_scale_cleared(pipeline):
+    req = _make_request(true_cfg_scale=2.0)
+    pipeline._sanitize_dmd2_request(req)
+    assert req.sampling_params.true_cfg_scale is None
+
+
+def test_do_classifier_free_guidance_forced_false(pipeline):
+    req = _make_request(do_classifier_free_guidance=True)
+    pipeline._sanitize_dmd2_request(req)
+    assert req.sampling_params.do_classifier_free_guidance is False
+
+
+def test_is_cfg_negative_forced_false(pipeline):
+    req = _make_request(is_cfg_negative=True)
+    pipeline._sanitize_dmd2_request(req)
+    assert req.sampling_params.is_cfg_negative is False
+
+
+def test_negative_prompt_stripped_from_prompt_dict(pipeline):
+    req = _make_request(prompts=[{"prompt": "a cat", "negative_prompt": "blurry"}])
+    pipeline._sanitize_dmd2_request(req)
+    assert "negative_prompt" not in req.prompts[0]
+    assert req.prompts[0]["prompt"] == "a cat"
+
+
+def test_no_negative_prompt_unchanged(pipeline):
+    req = _make_request(prompts=[{"prompt": "a cat"}])
+    pipeline._sanitize_dmd2_request(req)
+    assert req.prompts[0] == {"prompt": "a cat"}
+
+
+def test_string_prompt_not_mutated(pipeline):
+    """String prompts (not dicts) must pass through unchanged."""
+    req = _make_request(prompts=["a cat dancing"])
+    pipeline._sanitize_dmd2_request(req)
+    assert req.prompts == ["a cat dancing"]
+
+
+def test_multiple_prompts_all_sanitized(pipeline):
+    req = _make_request(
+        prompts=[
+            {"prompt": "a cat", "negative_prompt": "blurry"},
+            {"prompt": "a dog", "negative_prompt": "ugly"},
+        ]
+    )
+    pipeline._sanitize_dmd2_request(req)
+    for p in req.prompts:
+        assert "negative_prompt" not in p
+
+
+# ---------------------------------------------------------------------------
+# Clean request — nothing changes
+# ---------------------------------------------------------------------------
+
+
+def test_clean_request_no_changes(pipeline):
+    req = _make_request(
+        guidance_scale=pipeline.dmd2_guidance_scale,
+        guidance_scale_provided=False,
+        do_classifier_free_guidance=False,
+        is_cfg_negative=False,
+    )
+    pipeline._sanitize_dmd2_request(req)
+    assert req.sampling_params.guidance_scale == pipeline.dmd2_guidance_scale
+    assert req.sampling_params.guidance_scale_provided is False
+    assert req.sampling_params.guidance_scale_2 is None
+    assert req.sampling_params.true_cfg_scale is None
+    assert req.sampling_params.do_classifier_free_guidance is False
+    assert req.sampling_params.is_cfg_negative is False

tests/diffusion/models/wan2_2/test_wan_dmd2_scheduler.py

@@ -0,0 +1,88 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+from unittest.mock import MagicMock, patch
+
+import pytest
+import torch
+
+from vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2 import Wan22Pipeline, WanT2VDMD2Pipeline
+from vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2_i2v import Wan22I2VPipeline, WanI2VDMD2Pipeline
+from vllm_omni.diffusion.request import OmniDiffusionRequest, OmniDiffusionSamplingParams
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+_DMD2_TIMESTEPS = [999, 937, 833, 624]
+
+# Wan base pipeline whose __init__ loads model weights — mocked in tests.
+_WAN_BASE = {
+    WanT2VDMD2Pipeline: Wan22Pipeline,
+    WanI2VDMD2Pipeline: Wan22I2VPipeline,
+}
+
+
+def _make_pipeline(cls):
+    """Run the DMD2 __init__ (including __init_dmd2__) with the Wan base mocked."""
+
+    base = _WAN_BASE[cls]
+    od_config = MagicMock()
+    od_config.model = "/nonexistent"  # _load_model_index returns {} → uses inline defaults
+
+    def _mock_base_init(self, *a, **kw):
+        self.od_config = od_config  # __init_dmd2__ needs this
+
+    with patch.object(base, "__init__", _mock_base_init):
+        pipeline = object.__new__(cls)
+        torch.nn.Module.__init__(pipeline)
+        cls.__init__(pipeline, od_config=od_config)
+    return pipeline
+
+
+def _make_request(**sp_kwargs) -> OmniDiffusionRequest:
+    sp = OmniDiffusionSamplingParams(**sp_kwargs)
+    return OmniDiffusionRequest(prompts=[{"prompt": "a cat"}], sampling_params=sp)
+
+
+@pytest.fixture(params=[WanT2VDMD2Pipeline, WanI2VDMD2Pipeline], ids=["t2v", "i2v"])
+def pipeline(request):
+    return _make_pipeline(request.param)
+
+
+# ---------------------------------------------------------------------------
+# forward() timestep injection
+# ---------------------------------------------------------------------------
+
+
+def _fake_parent_forward(self, req, *args, num_inference_steps=40, **kwargs):
+    """Stub that calls set_timesteps as the real parent does."""
+    self.scheduler.set_timesteps(num_inference_steps, device="cpu")
+    return MagicMock()
+
+
+def test_forward_timesteps_match_dmd2_schedule(pipeline):
+    """After forward() runs, scheduler.timesteps must equal the DMD2 training schedule."""
+    parent = _WAN_BASE[type(pipeline)]
+
+    # Baseline: calling set_timesteps(40) without the DMD2 override gives a different schedule
+    pipeline.scheduler.set_timesteps(40, device="cpu")
+    default_timesteps = pipeline.scheduler.timesteps.long().tolist()
+    assert default_timesteps == _DMD2_TIMESTEPS, (
+        "DMD2EulerScheduler should always return DMD2 timesteps regardless of num_steps"
+    )
+
+    with patch.object(parent, "forward", _fake_parent_forward):
+        pipeline.forward(_make_request())
+
+    assert pipeline.scheduler.timesteps.long().tolist() == _DMD2_TIMESTEPS
+
+
+def test_forward_timesteps_fixed_across_num_steps(pipeline):
+    """scheduler.timesteps is always the DMD2 schedule regardless of num_steps passed."""
+    parent = _WAN_BASE[type(pipeline)]
+
+    for num_steps in [1, 4, 10, 40, 100]:
+        with patch.object(parent, "forward", _fake_parent_forward):
+            pipeline.forward(_make_request())
+
+        assert pipeline.scheduler.timesteps.long().tolist() == _DMD2_TIMESTEPS, (
+            f"num_steps={num_steps}: got {pipeline.scheduler.timesteps.tolist()}"
+        )

vllm_omni/diffusion/models/wan2_2/__init__.py

@@ -3,6 +3,7 @@
 
 from .pipeline_wan2_2 import (
     Wan22Pipeline,
+    WanT2VDMD2Pipeline,
     create_transformer_from_config,
     get_wan22_post_process_func,
     get_wan22_pre_process_func,
@@ -11,6 +12,7 @@
 )
 from .pipeline_wan2_2_i2v import (
     Wan22I2VPipeline,
+    WanI2VDMD2Pipeline,
     get_wan22_i2v_post_process_func,
     get_wan22_i2v_pre_process_func,
 )
@@ -28,13 +30,15 @@
 from .wan2_2_vace_transformer import VaceWanTransformerBlock, WanVACETransformer3DModel
 
 __all__ = [
+    "WanT2VDMD2Pipeline",
     "Wan22Pipeline",
     "get_wan22_post_process_func",
     "get_wan22_pre_process_func",
     "retrieve_latents",
     "load_transformer_config",
     "create_transformer_from_config",
     "Wan22I2VPipeline",
+    "WanI2VDMD2Pipeline",
     "get_wan22_i2v_post_process_func",
     "get_wan22_i2v_pre_process_func",
     "Wan22TI2VPipeline",