diff --git a/tests/diffusion/models/dmd2/__init__.py b/tests/diffusion/models/dmd2/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/tests/diffusion/models/dmd2/test_dmd2_request_sanitization.py b/tests/diffusion/models/dmd2/test_dmd2_request_sanitization.py new file mode 100644 index 000000000..e270390bd --- /dev/null +++ b/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 diff --git a/tests/diffusion/models/dmd2/test_dmd2_scheduler.py b/tests/diffusion/models/dmd2/test_dmd2_scheduler.py new file mode 100644 index 000000000..32d00dbf1 --- /dev/null +++ b/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 diff --git a/tests/diffusion/models/wan2_2/__init__.py b/tests/diffusion/models/wan2_2/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/vllm_omni/diffusion/models/dmd2/__init__.py b/vllm_omni/diffusion/models/dmd2/__init__.py new file mode 100644 index 000000000..d0c8219d4 --- /dev/null +++ b/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", +] diff --git a/vllm_omni/diffusion/models/dmd2/mixin.py b/vllm_omni/diffusion/models/dmd2/mixin.py new file mode 100644 index 000000000..60c4b95ba --- /dev/null +++ b/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, + ) diff --git a/vllm_omni/diffusion/models/ltx2/__init__.py b/vllm_omni/diffusion/models/ltx2/__init__.py index 9f9d70f01..2a78b61ba 100644 --- a/vllm_omni/diffusion/models/ltx2/__init__.py +++ b/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", diff --git a/vllm_omni/diffusion/models/ltx2/pipeline_ltx2.py b/vllm_omni/diffusion/models/ltx2/pipeline_ltx2.py index c60b192f0..f06ffab16 100644 --- a/vllm_omni/diffusion/models/ltx2/pipeline_ltx2.py +++ b/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__() diff --git a/vllm_omni/diffusion/models/ltx2/pipeline_ltx2_image2video.py b/vllm_omni/diffusion/models/ltx2/pipeline_ltx2_image2video.py index 65e7454b7..50a71a54b 100644 --- a/vllm_omni/diffusion/models/ltx2/pipeline_ltx2_image2video.py +++ b/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__() diff --git a/vllm_omni/diffusion/models/magi_human/pipeline_magi_human.py b/vllm_omni/diffusion/models/magi_human/pipeline_magi_human.py index 881c72edc..c1abdf91f 100644 --- a/vllm_omni/diffusion/models/magi_human/pipeline_magi_human.py +++ b/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, diff --git a/vllm_omni/diffusion/models/schedulers/__init__.py b/vllm_omni/diffusion/models/schedulers/__init__.py index 6f8df78eb..e683ed272 100644 --- a/vllm_omni/diffusion/models/schedulers/__init__.py +++ b/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", ] diff --git a/vllm_omni/diffusion/models/schedulers/scheduling_dmd2_euler.py b/vllm_omni/diffusion/models/schedulers/scheduling_dmd2_euler.py new file mode 100644 index 000000000..01447a41d --- /dev/null +++ b/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) diff --git a/vllm_omni/diffusion/models/utils.py b/vllm_omni/diffusion/models/utils.py new file mode 100644 index 000000000..ba0d8dda2 --- /dev/null +++ b/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) diff --git a/vllm_omni/diffusion/models/wan2_2/__init__.py b/vllm_omni/diffusion/models/wan2_2/__init__.py index d418001d9..97808df29 100644 --- a/vllm_omni/diffusion/models/wan2_2/__init__.py +++ b/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", 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 5b0676144..c74b72441 100644 --- a/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2.py +++ b/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__() 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 4eedd9d20..726d8853a 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 @@ -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__() diff --git a/vllm_omni/diffusion/registry.py b/vllm_omni/diffusion/registry.py index 0bf8c0451..4001109cc 100644 --- a/vllm_omni/diffusion/registry.py +++ b/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",