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

Signed-off-by: ayushag <ayushag@nvidia.com>

tests/diffusion/models/dmd2/test_dmd2_request_sanitization.py

@@ -0,0 +1,180 @@
+# 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.ltx2.pipeline_ltx2 import LTX2Pipeline, LTX2T2VDMD2Pipeline
+from vllm_omni.diffusion.models.ltx2.pipeline_ltx2_image2video import LTX2I2VDMD2Pipeline, LTX2ImageToVideoPipeline
+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 subclass → immediate base pipeline whose __init__ loads model weights (mocked in tests).
+_DMD2_BASE = {
+    WanT2VDMD2Pipeline: Wan22Pipeline,
+    WanI2VDMD2Pipeline: Wan22I2VPipeline,
+    LTX2T2VDMD2Pipeline: LTX2Pipeline,
+    LTX2I2VDMD2Pipeline: LTX2ImageToVideoPipeline,
+}
+
+
+def _make_pipeline(cls):
+    """Run the DMD2 __init__ with the base pipeline mocked out (no model weights loaded)."""
+
+    base = _DMD2_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=list(_DMD2_BASE.keys()),
+    ids=["wan_t2v", "wan_i2v", "ltx2_t2v", "ltx2_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/dmd2/test_dmd2_scheduler.py

@@ -0,0 +1,93 @@
+# 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.ltx2.pipeline_ltx2 import LTX2Pipeline, LTX2T2VDMD2Pipeline
+from vllm_omni.diffusion.models.ltx2.pipeline_ltx2_image2video import LTX2I2VDMD2Pipeline, LTX2ImageToVideoPipeline
+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]
+
+# DMD2 subclass → immediate base pipeline whose __init__ loads model weights (mocked in tests).
+_DMD2_BASE = {
+    WanT2VDMD2Pipeline: Wan22Pipeline,
+    WanI2VDMD2Pipeline: Wan22I2VPipeline,
+    LTX2T2VDMD2Pipeline: LTX2Pipeline,
+    LTX2I2VDMD2Pipeline: LTX2ImageToVideoPipeline,
+}
+
+
+def _make_pipeline(cls):
+    """Run the DMD2 __init__ (including __init_dmd2__) with the base pipeline mocked."""
+
+    base = _DMD2_BASE[cls]
+    od_config = MagicMock()
+    od_config.model = "/nonexistent"
+
+    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=list(_DMD2_BASE.keys()),
+    ids=["wan_t2v", "wan_i2v", "ltx2_t2v", "ltx2_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 = _DMD2_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_idempotent_across_calls(pipeline):
+    """Successive forward() calls must not cause scheduler state to drift."""
+    parent = _DMD2_BASE[type(pipeline)]
+
+    with patch.object(parent, "forward", _fake_parent_forward):
+        pipeline.forward(_make_request())
+        pipeline.forward(_make_request())
+
+    assert pipeline.scheduler.timesteps.long().tolist() == _DMD2_TIMESTEPS

vllm_omni/diffusion/models/dmd2/__init__.py

@@ -0,0 +1,8 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from vllm_omni.diffusion.models.dmd2.mixin import DMD2PipelineMixin
+
+__all__ = [
+    "DMD2PipelineMixin",
+]

vllm_omni/diffusion/models/dmd2/mixin.py

@@ -0,0 +1,88 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from __future__ import annotations
+
+import logging
+import os
+
+from vllm_omni.diffusion.data import DiffusionOutput
+from vllm_omni.diffusion.models.schedulers import DMD2EulerScheduler
+from vllm_omni.diffusion.models.utils import _load_json
+from vllm_omni.diffusion.request import OmniDiffusionRequest
+
+logger = logging.getLogger(__name__)
+
+
+class DMD2PipelineMixin:
+    """Mixin for FastGen DMD2-distilled models. Must appear before the base pipeline in MRO."""
+
+    def __init_dmd2__(self) -> None:
+        """Call after super().__init__() to apply DMD2 scheduler and read model_index."""
+        local_files_only = os.path.exists(self.od_config.model)
+        try:
+            model_index = _load_json(self.od_config.model, "model_index.json", local_files_only)
+        except Exception:
+            model_index = {}
+
+        dmd2_timesteps = model_index.get("dmd2_denoising_timesteps", [999, 937, 833, 624])
+        self.num_inference_steps = model_index.get("dmd2_num_inference_steps", 4)
+        shift = model_index.get("dmd2_scheduler_shift", 1.0)
+        self.dmd2_guidance_scale = model_index.get("dmd2_guidance_scale", 1.0)
+
+        self.scheduler = DMD2EulerScheduler(
+            num_train_timesteps=1000,
+            shift=shift,
+            dmd2_timesteps=dmd2_timesteps,
+        )
+
+    def _sanitize_dmd2_request(self, req: OmniDiffusionRequest) -> None:
+        """Sanitize CFG-related fields in-place. Mutates req.sampling_params and req.prompts."""
+        sp = req.sampling_params
+
+        if sp.num_inference_steps and sp.num_inference_steps != self.num_inference_steps:
+            logger.warning(
+                "DMD2: ignoring num_inference_steps=%d, forcing %d.",
+                sp.num_inference_steps,
+                self.num_inference_steps,
+            )
+        sp.num_inference_steps = self.num_inference_steps
+
+        if sp.guidance_scale_provided and sp.guidance_scale != self.dmd2_guidance_scale:
+            logger.warning(
+                "DMD2: ignoring guidance_scale=%.2f, forcing %.2f.",
+                sp.guidance_scale,
+                self.dmd2_guidance_scale,
+            )
+        sp.guidance_scale = self.dmd2_guidance_scale
+        sp.guidance_scale_provided = False
+
+        if sp.guidance_scale_2 is not None:
+            logger.warning("DMD2: ignoring guidance_scale_2.")
+            sp.guidance_scale_2 = None
+
+        if sp.true_cfg_scale is not None:
+            logger.warning("DMD2: ignoring true_cfg_scale.")
+            sp.true_cfg_scale = None
+
+        sp.do_classifier_free_guidance = False
+        sp.is_cfg_negative = False
+
+        fixed = []
+        for p in req.prompts:
+            if isinstance(p, dict) and "negative_prompt" in p:
+                logger.warning("DMD2: ignoring negative_prompt.")
+                p = {k: v for k, v in p.items() if k != "negative_prompt"}
+            fixed.append(p)
+        req.prompts = fixed
+
+    def forward(self, req: OmniDiffusionRequest, **kwargs) -> DiffusionOutput:
+        self._sanitize_dmd2_request(req)
+        kwargs.pop("guidance_scale", None)
+        kwargs.pop("num_inference_steps", None)
+        return super().forward(
+            req,
+            guidance_scale=self.dmd2_guidance_scale,
+            num_inference_steps=self.num_inference_steps,
+            **kwargs,
+        )

vllm_omni/diffusion/models/ltx2/__init__.py

@@ -4,12 +4,14 @@
 from vllm_omni.diffusion.models.ltx2.ltx2_transformer import LTX2VideoTransformer3DModel
 from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import (
     LTX2Pipeline,
+    LTX2T2VDMD2Pipeline,
     LTX2TwoStagesPipeline,
     create_transformer_from_config,
     get_ltx2_post_process_func,
     load_transformer_config,
 )
 from vllm_omni.diffusion.models.ltx2.pipeline_ltx2_image2video import (
+    LTX2I2VDMD2Pipeline,
     LTX2ImageToVideoPipeline,
     LTX2ImageToVideoTwoStagesPipeline,
 )
@@ -17,7 +19,9 @@ from vllm_omni.diffusion.models.ltx2.pipeline_ltx2_latent_upsample import LTX2La
 
 __all__ = [
     "LTX2Pipeline",
+    "LTX2T2VDMD2Pipeline",
     "LTX2ImageToVideoPipeline",
+    "LTX2I2VDMD2Pipeline",
     "LTX2LatentUpsamplePipeline",
     "LTX2TwoStagesPipeline",
     "LTX2ImageToVideoTwoStagesPipeline",

vllm_omni/diffusion/models/ltx2/pipeline_ltx2.py

@@ -33,6 +33,7 @@ from vllm_omni.diffusion.distributed.parallel_state import (
 from vllm_omni.diffusion.distributed.utils import get_local_device
 from vllm_omni.diffusion.lora.manager import DiffusionLoRAManager
 from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader
+from vllm_omni.diffusion.models.dmd2 import DMD2PipelineMixin
 from vllm_omni.diffusion.models.progress_bar import ProgressBarMixin
 from vllm_omni.diffusion.request import OmniDiffusionRequest
 from vllm_omni.lora.request import LoRARequest
@@ -1304,3 +1305,11 @@ class LTX2TwoStagesPipeline(nn.Module):
     def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
         loader = AutoWeightsLoader(self)
         return loader.load_weights(weights)
+
+
+class LTX2T2VDMD2Pipeline(DMD2PipelineMixin, LTX2Pipeline):
+    """LTX-2 T2V pipeline for FastGen DMD2-distilled models."""
+
+    def __init__(self, *, od_config: OmniDiffusionConfig, prefix: str = ""):
+        super().__init__(od_config=od_config, prefix=prefix)
+        self.__init_dmd2__()

vllm_omni/diffusion/models/ltx2/pipeline_ltx2_image2video.py

@@ -25,6 +25,7 @@ from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig
 from vllm_omni.diffusion.distributed.utils import get_local_device
 from vllm_omni.diffusion.lora.manager import DiffusionLoRAManager
 from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader
+from vllm_omni.diffusion.models.dmd2 import DMD2PipelineMixin
 from vllm_omni.diffusion.request import OmniDiffusionRequest
 from vllm_omni.lora.request import LoRARequest
 
@@ -889,3 +890,11 @@ class LTX2ImageToVideoTwoStagesPipeline(nn.Module):
     def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
         loader = AutoWeightsLoader(self)
         return loader.load_weights(weights)
+
+
+class LTX2I2VDMD2Pipeline(DMD2PipelineMixin, LTX2ImageToVideoPipeline):
+    """LTX-2 I2V pipeline for FastGen DMD2-distilled models."""
+
+    def __init__(self, *, od_config: OmniDiffusionConfig, prefix: str = ""):
+        super().__init__(od_config=od_config, prefix=prefix)
+        self.__init_dmd2__()

vllm_omni/diffusion/models/magi_human/pipeline_magi_human.py

@@ -48,6 +48,7 @@ from vllm_omni.diffusion.model_loader.diffusers_loader import (
 )
 from vllm_omni.diffusion.models.progress_bar import ProgressBarMixin
 from vllm_omni.diffusion.models.t5_encoder.t5_gemma_encoder import T5GemmaEncoderModelTP
+from vllm_omni.diffusion.models.utils import _load_json
 from vllm_omni.diffusion.profiler.diffusion_pipeline_profiler import (
     DiffusionPipelineProfilerMixin,
 )
@@ -1640,20 +1641,6 @@ def get_magi_human_post_process_func(*args, **kwargs):
 # ===========================================================================
 
 
-def _load_json(model_path: str, filename: str, local_files_only: bool = True) -> dict:
-    """Load a JSON config file from a local path or HuggingFace Hub repo."""
-    if local_files_only:
-        path = os.path.join(model_path, *filename.split("/"))
-        with open(path) as f:
-            return json.load(f)
-    else:
-        from huggingface_hub import hf_hub_download
-
-        cached = hf_hub_download(repo_id=model_path, filename=filename)
-        with open(cached) as f:
-            return json.load(f)
-
-
 def _resolve_subdir(
     model_path: str,
     subfolder: str,

vllm_omni/diffusion/models/schedulers/__init__.py

@@ -1,10 +1,12 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
+from vllm_omni.diffusion.models.schedulers.scheduling_dmd2_euler import DMD2EulerScheduler
 from vllm_omni.diffusion.models.schedulers.scheduling_flow_unipc_multistep import (
     FlowUniPCMultistepScheduler,
 )
 
 __all__ = [
+    "DMD2EulerScheduler",
     "FlowUniPCMultistepScheduler",
 ]

vllm_omni/diffusion/models/schedulers/scheduling_dmd2_euler.py

@@ -0,0 +1,23 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from __future__ import annotations
+
+import torch
+from diffusers import FlowMatchEulerDiscreteScheduler
+
+
+class DMD2EulerScheduler(FlowMatchEulerDiscreteScheduler):
+    """Euler scheduler that always uses the fixed DMD2 training timestep schedule."""
+
+    def __init__(self, *args, dmd2_timesteps: list[int], **kwargs):
+        super().__init__(*args, **kwargs)
+        self._dmd2_timesteps = dmd2_timesteps
+
+    def set_timesteps(
+        self,
+        num_inference_steps: int | None = None,
+        device: str | torch.device | None = None,
+        **kwargs,
+    ) -> None:
+        super().set_timesteps(timesteps=self._dmd2_timesteps, device=device)

vllm_omni/diffusion/models/utils.py

@@ -0,0 +1,21 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from __future__ import annotations
+
+import json
+import os
+
+
+def _load_json(model_path: str, filename: str, local_files_only: bool = True) -> dict:
+    """Load a JSON config file from a local path or HuggingFace Hub repo."""
+    if local_files_only:
+        path = os.path.join(model_path, *filename.split("/"))
+        with open(path) as f:
+            return json.load(f)
+    else:
+        from huggingface_hub import hf_hub_download
+
+        cached = hf_hub_download(repo_id=model_path, filename=filename)
+        with open(cached) as f:
+            return json.load(f)

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 import (
 )
 from .pipeline_wan2_2_i2v import (
     Wan22I2VPipeline,
+    WanI2VDMD2Pipeline,
     get_wan22_i2v_post_process_func,
     get_wan22_i2v_pre_process_func,
 )
@@ -28,6 +30,7 @@ from .wan2_2_transformer import WanTransformer3DModel
 from .wan2_2_vace_transformer import VaceWanTransformerBlock, WanVACETransformer3DModel
 
 __all__ = [
+    "WanT2VDMD2Pipeline",
     "Wan22Pipeline",
     "get_wan22_post_process_func",
     "get_wan22_pre_process_func",
@@ -35,6 +38,7 @@ __all__ = [
     "load_transformer_config",
     "create_transformer_from_config",
     "Wan22I2VPipeline",
+    "WanI2VDMD2Pipeline",
     "get_wan22_i2v_post_process_func",
     "get_wan22_i2v_pre_process_func",
     "Wan22TI2VPipeline",

vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2.py

@@ -22,6 +22,7 @@ from vllm_omni.diffusion.distributed.autoencoders.autoencoder_kl_wan import Dist
 from vllm_omni.diffusion.distributed.cfg_parallel import CFGParallelMixin
 from vllm_omni.diffusion.distributed.utils import get_local_device
 from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader
+from vllm_omni.diffusion.models.dmd2 import DMD2PipelineMixin
 from vllm_omni.diffusion.models.progress_bar import ProgressBarMixin, _is_rank_zero
 from vllm_omni.diffusion.models.schedulers import FlowUniPCMultistepScheduler
 from vllm_omni.diffusion.models.wan2_2.scheduling_wan_euler import WanEulerScheduler
@@ -955,3 +956,16 @@ class Wan22Pipeline(nn.Module, CFGParallelMixin, ProgressBarMixin, DiffusionPipe
 
         if boundary_ratio is None and guidance_scale_2 is not None:
             raise ValueError("`guidance_scale_2` is only supported when `boundary_ratio` is set.")
+
+
+# ---------------------------------------------------------------------------
+# DMD2-distilled variant
+# ---------------------------------------------------------------------------
+
+
+class WanT2VDMD2Pipeline(DMD2PipelineMixin, Wan22Pipeline):
+    """Wan 2.x T2V pipeline for FastGen DMD2-distilled models."""
+
+    def __init__(self, *, od_config: OmniDiffusionConfig, prefix: str = ""):
+        super().__init__(od_config=od_config, prefix=prefix)
+        self.__init_dmd2__()

vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2_i2v.py

@@ -23,8 +23,10 @@ from vllm_omni.diffusion.distributed.autoencoders.autoencoder_kl_wan import Dist
 from vllm_omni.diffusion.distributed.cfg_parallel import CFGParallelMixin
 from vllm_omni.diffusion.distributed.utils import get_local_device
 from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader
+from vllm_omni.diffusion.models.dmd2 import DMD2PipelineMixin
 from vllm_omni.diffusion.models.interface import SupportImageInput
 from vllm_omni.diffusion.models.progress_bar import ProgressBarMixin, _is_rank_zero
+from vllm_omni.diffusion.models.utils import _load_json
 from vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2 import (
     build_wan_scheduler,
     create_transformer_from_config,
@@ -43,29 +45,6 @@ logger = logging.getLogger(__name__)
 DEBUG_PERF = False
 
 
-def _load_model_index(model: str, local_files_only: bool) -> dict:
-    """Load model_index.json from local path or HF Hub."""
-    if local_files_only:
-        model_index_path = os.path.join(model, "model_index.json")
-        if os.path.exists(model_index_path):
-            import json
-
-            with open(model_index_path) as f:
-                return json.load(f)
-    else:
-        try:
-            import json
-
-            from huggingface_hub import hf_hub_download
-
-            model_index_path = hf_hub_download(model, "model_index.json")
-            with open(model_index_path) as f:
-                return json.load(f)
-        except Exception:
-            pass
-    return {}
-
-
 def get_wan22_i2v_post_process_func(
     od_config: OmniDiffusionConfig,
 ):
@@ -196,7 +175,10 @@ class Wan22I2VPipeline(
         ]
 
         # Load model_index.json to detect available components
-        model_index = _load_model_index(model, local_files_only)
+        try:
+            model_index = _load_json(model, "model_index.json", local_files_only)
+        except Exception:
+            model_index = {}
 
         # Check if this is a two-stage model (MoE with transformer_2)
         self.has_transformer_2 = "transformer_2" in model_index
@@ -903,3 +885,16 @@ class Wan22I2VPipeline(
         """Load weights using AutoWeightsLoader for vLLM integration."""
         loader = AutoWeightsLoader(self)
         return loader.load_weights(weights)
+
+
+# ---------------------------------------------------------------------------
+# DMD2-distilled variant
+# ---------------------------------------------------------------------------
+
+
+class WanI2VDMD2Pipeline(DMD2PipelineMixin, Wan22I2VPipeline):
+    """Wan 2.x I2V pipeline for FastGen DMD2-distilled models."""
+
+    def __init__(self, *, od_config: OmniDiffusionConfig, prefix: str = ""):
+        super().__init__(od_config=od_config, prefix=prefix)
+        self.__init_dmd2__()

vllm_omni/diffusion/registry.py

@@ -83,6 +83,16 @@ _DIFFUSION_MODELS = {
         "pipeline_ltx2_image2video",
         "LTX2ImageToVideoTwoStagesPipeline",
     ),
+    "LTX2T2VDMD2Pipeline": (
+        "ltx2",
+        "pipeline_ltx2",
+        "LTX2T2VDMD2Pipeline",
+    ),
+    "LTX2I2VDMD2Pipeline": (
+        "ltx2",
+        "pipeline_ltx2_image2video",
+        "LTX2I2VDMD2Pipeline",
+    ),
     "StableAudioPipeline": (
         "stable_audio",
         "pipeline_stable_audio",
@@ -93,6 +103,16 @@ _DIFFUSION_MODELS = {
         "pipeline_wan2_2_i2v",
         "Wan22I2VPipeline",
     ),
+    "WanT2VDMD2Pipeline": (
+        "wan2_2",
+        "pipeline_wan2_2",
+        "WanT2VDMD2Pipeline",
+    ),
+    "WanI2VDMD2Pipeline": (
+        "wan2_2",
+        "pipeline_wan2_2_i2v",
+        "WanI2VDMD2Pipeline",
+    ),
     "LongCatImagePipeline": (
         "longcat_image",
         "pipeline_longcat_image",
@@ -357,8 +377,12 @@ _DIFFUSION_POST_PROCESS_FUNCS = {
     "LTX2TwoStagesPipeline": "get_ltx2_post_process_func",
     "LTX2ImageToVideoPipeline": "get_ltx2_post_process_func",
     "LTX2ImageToVideoTwoStagesPipeline": "get_ltx2_post_process_func",
+    "LTX2T2VDMD2Pipeline": "get_ltx2_post_process_func",
+    "LTX2I2VDMD2Pipeline": "get_ltx2_post_process_func",
     "StableAudioPipeline": "get_stable_audio_post_process_func",
     "WanImageToVideoPipeline": "get_wan22_i2v_post_process_func",
+    "WanT2VDMD2Pipeline": "get_wan22_post_process_func",
+    "WanI2VDMD2Pipeline": "get_wan22_i2v_post_process_func",
     "LongCatImagePipeline": "get_longcat_image_post_process_func",
     "BagelPipeline": "get_bagel_post_process_func",
     "LongCatImageEditPipeline": "get_longcat_image_post_process_func",
@@ -390,6 +414,8 @@ _DIFFUSION_PRE_PROCESS_FUNCS = {
     "WanPipeline": "get_wan22_pre_process_func",
     "WanVACEPipeline": "get_wan22_vace_pre_process_func",
     "WanImageToVideoPipeline": "get_wan22_i2v_pre_process_func",
+    "WanT2VDMD2Pipeline": "get_wan22_pre_process_func",
+    "WanI2VDMD2Pipeline": "get_wan22_i2v_pre_process_func",
     "OmniGen2Pipeline": "get_omnigen2_pre_process_func",
     "HeliosPipeline": "get_helios_pre_process_func",
     "HeliosPyramidPipeline": "get_helios_pre_process_func",