diff --git a/examples/inference/basic/README.md b/examples/inference/basic/README.md
index 9257dbd57..633e7b5b8 100644
--- a/examples/inference/basic/README.md
+++ b/examples/inference/basic/README.md
@@ -18,6 +18,16 @@ git clone https://github.com/hao-ai-lab/FastVideo.git && cd FastVideo
 python examples/inference/basic/basic.py
 ```
 
+For an example on Apple silicon: 
+```
+python examples/inference/basic/basic_mps.py
+```
+
+For an example running DMD+VSA inference:
+```
+python examples/inference/basic/basic_dmd.py
+```
+
 ## Basic Walkthrough
 
 All you need to generate videos using multi-gpus from state-of-the-art diffusion pipelines is the following few lines!
diff --git a/examples/inference/basic/basic_dmd.py b/examples/inference/basic/basic_dmd.py
new file mode 100644
index 000000000..7724a1cf6
--- /dev/null
+++ b/examples/inference/basic/basic_dmd.py
@@ -0,0 +1,46 @@
+import os
+import time
+from fastvideo import VideoGenerator
+
+from fastvideo.v1.configs.sample import SamplingParam
+
+OUTPUT_PATH = "video_samples_dmd2"
+def main():
+    os.environ["FASTVIDEO_ATTENTION_BACKEND"] = "VIDEO_SPARSE_ATTN"
+
+    load_start_time = time.perf_counter()
+    model_name = "FastVideo/FastWan2.1-T2V-1.3B-Diffusers"
+    generator = VideoGenerator.from_pretrained(
+        model_name,
+        # FastVideo will automatically handle distributed setup
+        num_gpus=1,
+        use_fsdp_inference=True,
+        # Adjust these offload parameters if you have < 32GB of VRAM
+        text_encoder_offload=False,
+        use_cpu_offload=False,
+        VSA_sparsity=0.8,
+    )
+    load_end_time = time.perf_counter()
+    load_time = load_end_time - load_start_time
+
+
+    sampling_param = SamplingParam.from_pretrained(model_name)
+
+    prompt = (
+        "A curious raccoon peers through a vibrant field of yellow sunflowers, its eyes "
+        "wide with interest. The playful yet serene atmosphere is complemented by soft "
+        "natural light filtering through the petals. Mid-shot, warm and cheerful tones."
+    )
+    start_time = time.perf_counter()
+    video = generator.generate_video(prompt, output_path=OUTPUT_PATH, save_video=True, sampling_param=sampling_param)
+    end_time = time.perf_counter()
+    gen_time = end_time - start_time
+    e2e_gen_time = end_time - start_time
+
+    print(f"Time taken to load model: {load_time} seconds")
+    print(f"Time taken to generate video: {gen_time} seconds")
+    print(f"Time taken for e2e generation: {e2e_gen_time} seconds")
+
+
+if __name__ == "__main__":
+    main()
diff --git a/fastvideo/v1/configs/pipelines/base.py b/fastvideo/v1/configs/pipelines/base.py
index 5875370e8..92aa7aefc 100644
--- a/fastvideo/v1/configs/pipelines/base.py
+++ b/fastvideo/v1/configs/pipelines/base.py
@@ -83,6 +83,9 @@ class PipelineConfig:
     STA_mode: STA_Mode = STA_Mode.STA_INFERENCE
     skip_time_steps: int = 15
 
+    # DMD parameters
+    dmd_denoising_steps: list[int] | None = field(default=None)
+
     # Compilation
     # enable_torch_compile: bool = False
 
@@ -200,6 +203,15 @@ class PipelineConfig:
             "Bool for applying scheduler scale in set_timesteps, used in stepvideo",
         )
 
+        # DMD parameters
+        parser.add_argument(
+            f"--{prefix_with_dot}dmd-denoising-steps",
+            type=parse_int_list,
+            default=PipelineConfig.dmd_denoising_steps,
+            help=
+            "Comma-separated list of denoising steps (e.g., '1000,757,522')",
+        )
+
         # Add VAE configuration arguments
         from fastvideo.v1.configs.models.vaes.base import VAEConfig
         VAEConfig.add_cli_args(parser, prefix=f"{prefix_with_dot}vae-config")
@@ -381,3 +393,10 @@ class SlidingTileAttnConfig(PipelineConfig):
     # Additional configuration specific to sliding tile attention
     pad_to_square: bool = False
     use_overlap_optimization: bool = True
+
+
+def parse_int_list(value: str) -> list[int]:
+    """Parse a comma-separated string of integers into a list."""
+    if not value:
+        return []
+    return [int(x.strip()) for x in value.split(",")]
diff --git a/fastvideo/v1/configs/pipelines/registry.py b/fastvideo/v1/configs/pipelines/registry.py
index dd7f1fddc..801d34c61 100644
--- a/fastvideo/v1/configs/pipelines/registry.py
+++ b/fastvideo/v1/configs/pipelines/registry.py
@@ -8,7 +8,8 @@ from fastvideo.v1.configs.pipelines.base import PipelineConfig
 from fastvideo.v1.configs.pipelines.hunyuan import (FastHunyuanConfig,
                                                     HunyuanConfig)
 from fastvideo.v1.configs.pipelines.stepvideo import StepVideoT2VConfig
-from fastvideo.v1.configs.pipelines.wan import (WanI2V480PConfig,
+from fastvideo.v1.configs.pipelines.wan import (FastWanT2V480PConfig,
+                                                WanI2V480PConfig,
                                                 WanI2V720PConfig,
                                                 WanT2V480PConfig,
                                                 WanT2V720PConfig)
@@ -27,6 +28,7 @@ PIPE_NAME_TO_CONFIG: dict[str, type[PipelineConfig]] = {
     "Wan-AI/Wan2.1-I2V-14B-480P-Diffusers": WanI2V480PConfig,
     "Wan-AI/Wan2.1-I2V-14B-720P-Diffusers": WanI2V720PConfig,
     "Wan-AI/Wan2.1-T2V-14B-Diffusers": WanT2V720PConfig,
+    "FastVideo/FastWan2.1-T2V-1.3B-Diffusers": FastWanT2V480PConfig,
     "FastVideo/stepvideo-t2v-diffusers": StepVideoT2VConfig,
     # Add other specific weight variants
 }
@@ -36,6 +38,7 @@ PIPELINE_DETECTOR: dict[str, Callable[[str], bool]] = {
     "hunyuan": lambda id: "hunyuan" in id.lower(),
     "wanpipeline": lambda id: "wanpipeline" in id.lower(),
     "wanimagetovideo": lambda id: "wanimagetovideo" in id.lower(),
+    "wandmdpipeline": lambda id: "wandmdpipeline" in id.lower(),
     "stepvideo": lambda id: "stepvideo" in id.lower(),
     # Add other pipeline architecture detectors
 }
@@ -47,6 +50,7 @@ PIPELINE_FALLBACK_CONFIG: dict[str, type[PipelineConfig]] = {
     "wanpipeline":
     WanT2V480PConfig,  # Base Wan config as fallback for any Wan variant
     "wanimagetovideo": WanI2V480PConfig,
+    "wandmdpipeline": FastWanT2V480PConfig,
     "stepvideo": StepVideoT2VConfig
     # Other fallbacks by architecture
 }
diff --git a/fastvideo/v1/configs/pipelines/wan.py b/fastvideo/v1/configs/pipelines/wan.py
index 6309afc1f..c9d23ac0f 100644
--- a/fastvideo/v1/configs/pipelines/wan.py
+++ b/fastvideo/v1/configs/pipelines/wan.py
@@ -95,3 +95,19 @@ class WanI2V720PConfig(WanI2V480PConfig):
 
     # Denoising stage
     flow_shift: int = 5
+
+
+@dataclass
+class FastWanT2V480PConfig(WanT2V480PConfig):
+    """Base configuration for FastWan T2V 1.3B 480P pipeline architecture with DMD"""
+
+    # WanConfig-specific parameters with defaults
+
+    # Denoising stage
+    flow_shift: int = 8
+    dmd_denoising_steps: list[int] | None = field(
+        default_factory=lambda: [1000, 757, 522])
+
+    def __post_init__(self) -> None:
+        self.vae_config.load_encoder = True
+        self.vae_config.load_decoder = True
diff --git a/fastvideo/v1/configs/sample/registry.py b/fastvideo/v1/configs/sample/registry.py
index 15d8445c9..b1df34675 100644
--- a/fastvideo/v1/configs/sample/registry.py
+++ b/fastvideo/v1/configs/sample/registry.py
@@ -6,7 +6,8 @@ from typing import Any
 from fastvideo.v1.configs.sample.hunyuan import (FastHunyuanSamplingParam,
                                                  HunyuanSamplingParam)
 from fastvideo.v1.configs.sample.stepvideo import StepVideoT2VSamplingParam
-from fastvideo.v1.configs.sample.wan import (WanI2V_14B_480P_SamplingParam,
+from fastvideo.v1.configs.sample.wan import (FastWanT2V480PConfig,
+                                             WanI2V_14B_480P_SamplingParam,
                                              WanI2V_14B_720P_SamplingParam,
                                              WanT2V_1_3B_SamplingParam,
                                              WanT2V_14B_SamplingParam)
@@ -24,6 +25,7 @@ SAMPLING_PARAM_REGISTRY: dict[str, Any] = {
     "Wan-AI/Wan2.1-I2V-14B-480P-Diffusers": WanI2V_14B_480P_SamplingParam,
     "Wan-AI/Wan2.1-I2V-14B-720P-Diffusers": WanI2V_14B_720P_SamplingParam,
     "FastVideo/stepvideo-t2v-diffusers": StepVideoT2VSamplingParam,
+    "FastVideo/FastWan2.1-T2V-1.3B-Diffusers": FastWanT2V480PConfig,
     # Add other specific weight variants
 }
 
diff --git a/fastvideo/v1/configs/sample/wan.py b/fastvideo/v1/configs/sample/wan.py
index 2ca49e4d4..c3e6b6ed3 100644
--- a/fastvideo/v1/configs/sample/wan.py
+++ b/fastvideo/v1/configs/sample/wan.py
@@ -94,3 +94,14 @@ class WanI2V_14B_720P_SamplingParam(WanT2V_14B_SamplingParam):
                 -5784.54975374, 5449.50911966, -1811.16591783, 256.27178429,
                 -13.02252404
             ]))
+
+
+@dataclass
+class FastWanT2V480PConfig(WanT2V_1_3B_SamplingParam):
+    # DMD parameters
+    # dmd_denoising_steps: list[int] | None = field(default_factory=lambda: [1000, 757, 522])
+    num_inference_steps: int = 3
+    num_frames: int = 61
+    height: int = 448
+    width: int = 832
+    fps: int = 16
diff --git a/fastvideo/v1/fastvideo_args.py b/fastvideo/v1/fastvideo_args.py
index 63a0274c5..365f64006 100644
--- a/fastvideo/v1/fastvideo_args.py
+++ b/fastvideo/v1/fastvideo_args.py
@@ -95,6 +95,9 @@ class FastVideoArgs:
         "vae": True,
     })
 
+    # # DMD parameters
+    # dmd_denoising_steps: List[int] | None = field(default=None)
+
     @property
     def training_mode(self) -> bool:
         return not self.inference_mode
diff --git a/fastvideo/v1/models/registry.py b/fastvideo/v1/models/registry.py
index 6643b90e0..eeb42c1ef 100644
--- a/fastvideo/v1/models/registry.py
+++ b/fastvideo/v1/models/registry.py
@@ -54,7 +54,7 @@ _VAE_MODELS = {
 _SCHEDULERS = {
     "FlowMatchEulerDiscreteScheduler":
     ("schedulers", "scheduling_flow_match_euler_discrete",
-     "FlowMatchDiscreteScheduler"),
+     "FlowMatchEulerDiscreteScheduler"),
     "UniPCMultistepScheduler":
     ("schedulers", "scheduling_unipc_multistep", "UniPCMultistepScheduler"),
 }
diff --git a/fastvideo/v1/models/schedulers/base.py b/fastvideo/v1/models/schedulers/base.py
index a6cabb458..8a516ec0b 100644
--- a/fastvideo/v1/models/schedulers/base.py
+++ b/fastvideo/v1/models/schedulers/base.py
@@ -3,7 +3,6 @@
 from abc import ABC, abstractmethod
 
 import torch
-from diffusers.utils import BaseOutput
 
 
 class BaseScheduler(ABC):
@@ -33,13 +32,3 @@ class BaseScheduler(ABC):
                           sample: torch.Tensor,
                           timestep: int | None = None) -> torch.Tensor:
         pass
-
-    @abstractmethod
-    def step(
-        self,
-        model_output: torch.Tensor,
-        timestep: int | torch.Tensor,
-        sample: torch.Tensor,
-        return_dict: bool = True,
-    ) -> BaseOutput | tuple:
-        pass
diff --git a/fastvideo/v1/models/schedulers/scheduling_flow_match_euler_discrete.py b/fastvideo/v1/models/schedulers/scheduling_flow_match_euler_discrete.py
index e3d3fff08..b7792c9ad 100644
--- a/fastvideo/v1/models/schedulers/scheduling_flow_match_euler_discrete.py
+++ b/fastvideo/v1/models/schedulers/scheduling_flow_match_euler_discrete.py
@@ -18,10 +18,12 @@
 # Modified from diffusers==0.29.2
 #
 # ==============================================================================
-
+import math
 from dataclasses import dataclass
 from typing import Any
 
+import numpy as np
+import scipy.stats
 import torch
 from diffusers.configuration_utils import ConfigMixin, register_to_config
 from diffusers.schedulers.scheduling_utils import SchedulerMixin
@@ -34,7 +36,7 @@ logger = init_logger(__name__)
 
 
 @dataclass
-class FlowMatchDiscreteSchedulerOutput(BaseOutput):
+class FlowMatchEulerDiscreteSchedulerOutput(BaseOutput):
     """
     Output class for the scheduler's `step` function output.
 
@@ -47,7 +49,8 @@ class FlowMatchDiscreteSchedulerOutput(BaseOutput):
     prev_sample: torch.FloatTensor
 
 
-class FlowMatchDiscreteScheduler(SchedulerMixin, ConfigMixin, BaseScheduler):
+class FlowMatchEulerDiscreteScheduler(SchedulerMixin, ConfigMixin,
+                                      BaseScheduler):
     """
     Euler scheduler.
 
@@ -57,13 +60,34 @@ class FlowMatchDiscreteScheduler(SchedulerMixin, ConfigMixin, BaseScheduler):
     Args:
         num_train_timesteps (`int`, defaults to 1000):
             The number of diffusion steps to train the model.
-        timestep_spacing (`str`, defaults to `"linspace"`):
-            The way the timesteps should be scaled. Refer to Table 2 of the [Common Diffusion Noise Schedules and
-            Sample Steps are Flawed](https://huggingface.co/papers/2305.08891) for more information.
         shift (`float`, defaults to 1.0):
             The shift value for the timestep schedule.
-        reverse (`bool`, defaults to `True`):
-            Whether to reverse the timestep schedule.
+        use_dynamic_shifting (`bool`, defaults to False):
+            Whether to apply timestep shifting on-the-fly based on the image resolution.
+        base_shift (`float`, defaults to 0.5):
+            Value to stabilize image generation. Increasing `base_shift` reduces variation and image is more consistent
+            with desired output.
+        max_shift (`float`, defaults to 1.15):
+            Value change allowed to latent vectors. Increasing `max_shift` encourages more variation and image may be
+            more exaggerated or stylized.
+        base_image_seq_len (`int`, defaults to 256):
+            The base image sequence length.
+        max_image_seq_len (`int`, defaults to 4096):
+            The maximum image sequence length.
+        invert_sigmas (`bool`, defaults to False):
+            Whether to invert the sigmas.
+        shift_terminal (`float`, defaults to None):
+            The end value of the shifted timestep schedule.
+        use_karras_sigmas (`bool`, defaults to False):
+            Whether to use Karras sigmas for step sizes in the noise schedule during sampling.
+        use_exponential_sigmas (`bool`, defaults to False):
+            Whether to use exponential sigmas for step sizes in the noise schedule during sampling.
+        use_beta_sigmas (`bool`, defaults to False):
+            Whether to use beta sigmas for step sizes in the noise schedule during sampling.
+        time_shift_type (`str`, defaults to "exponential"):
+            The type of dynamic resolution-dependent timestep shifting to apply. Either "exponential" or "linear".
+        stochastic_sampling (`bool`, defaults to False):
+            Whether to use stochastic sampling.
     """
 
     _compatibles: list[Any] = []
@@ -74,48 +98,76 @@ class FlowMatchDiscreteScheduler(SchedulerMixin, ConfigMixin, BaseScheduler):
         self,
         num_train_timesteps: int = 1000,
         shift: float = 1.0,
-        reverse: bool = True,
-        solver: str = "euler",
-        n_tokens: int | None = None,
-        **kwargs,
+        use_dynamic_shifting: bool = False,
+        base_shift: float | None = 0.5,
+        max_shift: float | None = 1.15,
+        base_image_seq_len: int | None = 256,
+        max_image_seq_len: int | None = 4096,
+        invert_sigmas: bool = False,
+        shift_terminal: float | None = None,
+        use_karras_sigmas: bool | None = False,
+        use_exponential_sigmas: bool | None = False,
+        use_beta_sigmas: bool | None = False,
+        time_shift_type: str = "exponential",
+        stochastic_sampling: bool = False,
     ):
-        sigmas = torch.linspace(1, 0, num_train_timesteps + 1)
+        if sum([
+                self.config.use_beta_sigmas, self.config.use_exponential_sigmas,
+                self.config.use_karras_sigmas
+        ]) > 1:
+            raise ValueError(
+                "Only one of `config.use_beta_sigmas`, `config.use_exponential_sigmas`, `config.use_karras_sigmas` can be used."
+            )
+        if time_shift_type not in {"exponential", "linear"}:
+            raise ValueError(
+                "`time_shift_type` must either be 'exponential' or 'linear'.")
 
-        if not reverse:
-            sigmas = sigmas.flip(0)
+        timesteps = np.linspace(1,
+                                num_train_timesteps,
+                                num_train_timesteps,
+                                dtype=np.float32)[::-1].copy()
+        timesteps = torch.from_numpy(timesteps).to(dtype=torch.float32)
 
-        self.sigmas = sigmas
-        # the value fed to model
-        self.timesteps = (sigmas[:-1] *
-                          num_train_timesteps).to(dtype=torch.float32)
+        sigmas = timesteps / num_train_timesteps
+        if not use_dynamic_shifting:
+            # when use_dynamic_shifting is True, we apply the timestep shifting on the fly based on the image resolution
+            sigmas = shift * sigmas / (1 + (shift - 1) * sigmas)
+
+        self.timesteps = sigmas * num_train_timesteps
 
         self._step_index: int | None = None
-        self._begin_index = 0
+        self._begin_index: int | None = None
 
-        self.supported_solver = ["euler"]
-        if solver not in self.supported_solver:
-            raise ValueError(
-                f"Solver {solver} not supported. Supported solvers: {self.supported_solver}"
-            )
+        self._shift = shift
 
-        BaseScheduler.__init__(self)
+        self.sigmas = sigmas.to(
+            "cpu")  # to avoid too much CPU/GPU communication
+        self.sigma_min = self.sigmas[-1].item()
+        self.sigma_max = self.sigmas[0].item()
 
     @property
-    def step_index(self):
+    def shift(self) -> float:
+        """
+        The value used for shifting.
+        """
+        return self._shift
+
+    @property
+    def step_index(self) -> int | None:
         """
         The index counter for current timestep. It will increase 1 after each scheduler step.
         """
         return self._step_index
 
     @property
-    def begin_index(self):
+    def begin_index(self) -> int | None:
         """
         The index for the first timestep. It should be set from pipeline with `set_begin_index` method.
         """
         return self._begin_index
 
     # Copied from diffusers.schedulers.scheduling_dpmsolver_multistep.DPMSolverMultistepScheduler.set_begin_index
-    def set_begin_index(self, begin_index: int = 0):
+    def set_begin_index(self, begin_index: int = 0) -> None:
         """
         Sets the begin index for the scheduler. This function should be run from pipeline before the inference.
 
@@ -125,44 +177,242 @@ class FlowMatchDiscreteScheduler(SchedulerMixin, ConfigMixin, BaseScheduler):
         """
         self._begin_index = begin_index
 
-    def _sigma_to_t(self, sigma):
+    def set_shift(self, shift: float) -> None:
+        self._shift = shift
+
+    def scale_noise(
+        self,
+        sample: torch.FloatTensor,
+        timestep: float | torch.FloatTensor,
+        noise: torch.FloatTensor | None = None,
+    ) -> torch.FloatTensor:
+        """
+        Forward process in flow-matching
+
+        Args:
+            sample (`torch.FloatTensor`):
+                The input sample.
+            timestep (`int`, *optional*):
+                The current timestep in the diffusion chain.
+
+        Returns:
+            `torch.FloatTensor`:
+                A scaled input sample.
+        """
+        # Make sure sigmas and timesteps have the same device and dtype as original_samples
+        sigmas = self.sigmas.to(device=sample.device, dtype=sample.dtype)
+
+        if sample.device.type == "mps" and torch.is_floating_point(timestep):
+            # mps does not support float64
+            schedule_timesteps = self.timesteps.to(sample.device,
+                                                   dtype=torch.float32)
+            assert isinstance(timestep, torch.Tensor)
+            timestep = timestep.to(sample.device, dtype=torch.float32)
+        else:
+            schedule_timesteps = self.timesteps.to(sample.device)
+            assert isinstance(timestep, torch.Tensor)
+            timestep = timestep.to(sample.device)
+
+        # self.begin_index is None when scheduler is used for training, or pipeline does not implement set_begin_index
+        if self.begin_index is None:
+            step_indices = [
+                self.index_for_timestep(t, schedule_timesteps) for t in timestep
+            ]
+        elif self.step_index is not None:
+            # add_noise is called after first denoising step (for inpainting)
+            step_indices = [self.step_index] * timestep.shape[0]
+        else:
+            # add noise is called before first denoising step to create initial latent(img2img)
+            step_indices = [self.begin_index] * timestep.shape[0]
+
+        sigma = sigmas[step_indices].flatten()
+        while len(sigma.shape) < len(sample.shape):
+            sigma = sigma.unsqueeze(-1)
+
+        sample = sigma * noise + (1.0 - sigma) * sample
+
+        return sample
+
+    def _sigma_to_t(self, sigma: float) -> float:
         return sigma * self.config.num_train_timesteps
 
+    def time_shift(self, mu: float, sigma: float,
+                   t: torch.Tensor | np.ndarray) -> torch.Tensor | np.ndarray:
+        if self.config.time_shift_type == "exponential":
+            return self._time_shift_exponential(mu, sigma, t)
+        elif self.config.time_shift_type == "linear":
+            return self._time_shift_linear(mu, sigma, t)
+        else:
+            raise ValueError(
+                f"Unknown time_shift_type: {self.config.time_shift_type}")
+
+    def stretch_shift_to_terminal(self, t: torch.Tensor) -> torch.Tensor:
+        r"""
+        Stretches and shifts the timestep schedule to ensure it terminates at the configured `shift_terminal` config
+        value.
+
+        Reference:
+        https://github.com/Lightricks/LTX-Video/blob/a01a171f8fe3d99dce2728d60a73fecf4d4238ae/ltx_video/schedulers/rf.py#L51
+
+        Args:
+            t (`torch.Tensor`):
+                A tensor of timesteps to be stretched and shifted.
+
+        Returns:
+            `torch.Tensor`:
+                A tensor of adjusted timesteps such that the final value equals `self.config.shift_terminal`.
+        """
+        one_minus_z = 1 - t
+        scale_factor = one_minus_z[-1] / (1 - self.config.shift_terminal)
+        stretched_t = 1 - (one_minus_z / scale_factor)
+        return stretched_t
+
     def set_timesteps(
         self,
-        num_inference_steps: int,
+        num_inference_steps: int | None = None,
         device: str | torch.device = None,
-        n_tokens: int = 0,
-    ):
+        sigmas: list[float] | None = None,
+        mu: float | None = None,
+        timesteps: list[float] | None = None,
+    ) -> None:
         """
         Sets the discrete timesteps used for the diffusion chain (to be run before inference).
 
         Args:
-            num_inference_steps (`int`):
+            num_inference_steps (`int`, *optional*):
                 The number of diffusion steps used when generating samples with a pre-trained model.
             device (`str` or `torch.device`, *optional*):
                 The device to which the timesteps should be moved to. If `None`, the timesteps are not moved.
-            n_tokens (`int`, *optional*):
-                Number of tokens in the input sequence.
+            sigmas (`List[float]`, *optional*):
+                Custom values for sigmas to be used for each diffusion step. If `None`, the sigmas are computed
+                automatically.
+            mu (`float`, *optional*):
+                Determines the amount of shifting applied to sigmas when performing resolution-dependent timestep
+                shifting.
+            timesteps (`List[float]`, *optional*):
+                Custom values for timesteps to be used for each diffusion step. If `None`, the timesteps are computed
+                automatically.
         """
+        if self.config.use_dynamic_shifting and mu is None:
+            raise ValueError(
+                "`mu` must be passed when `use_dynamic_shifting` is set to be `True`"
+            )
+
+        if sigmas is not None and timesteps is not None and len(sigmas) != len(
+                timesteps):
+            raise ValueError(
+                "`sigmas` and `timesteps` should have the same length")
+
+        if num_inference_steps is not None:
+            if (sigmas is not None and len(sigmas) != num_inference_steps) or (
+                    timesteps is not None
+                    and len(timesteps) != num_inference_steps):
+                raise ValueError(
+                    "`sigmas` and `timesteps` should have the same length as num_inference_steps, if `num_inference_steps` is provided"
+                )
+        else:
+            if sigmas is not None:
+                num_inference_steps = len(sigmas)
+            elif timesteps is not None:
+                num_inference_steps = len(timesteps)
+            else:
+                raise ValueError(
+                    "Either num_inference_steps, sigmas, or timesteps must be provided"
+                )
+
         self.num_inference_steps = num_inference_steps
 
-        sigmas = torch.linspace(1, 0, num_inference_steps + 1)
-        sigmas = self.sd3_time_shift(sigmas)
+        # 1. Prepare default sigmas
+        is_timesteps_provided = timesteps is not None
 
-        if not self.config.reverse:
-            sigmas = 1 - sigmas
+        timesteps_array: np.ndarray | None = None
+        if is_timesteps_provided:
+            assert timesteps is not None
+            timesteps_array = np.array(timesteps).astype(np.float32)
 
-        self.sigmas = sigmas
-        if not getattr(self.config, "timesteps_scale", True):
-            self.timesteps = sigmas[:-1]  # for stepvideo
+        sigmas_array: np.ndarray
+        if sigmas is None:
+            if timesteps_array is None:
+                timesteps_array = np.linspace(self._sigma_to_t(self.sigma_max),
+                                              self._sigma_to_t(self.sigma_min),
+                                              num_inference_steps)
+            sigmas_array = timesteps_array / self.config.num_train_timesteps
         else:
-            self.timesteps = (sigmas[:-1] * self.config.num_train_timesteps).to(
-                dtype=torch.float32, device=device)
-        # Reset step index
-        self._step_index = None
+            sigmas_array = np.array(sigmas).astype(np.float32)
+            num_inference_steps = len(sigmas_array)
 
-    def index_for_timestep(self, timestep, schedule_timesteps=None) -> int:
+        # 2. Perform timestep shifting. Either no shifting is applied, or resolution-dependent shifting of
+        #    "exponential" or "linear" type is applied
+        if self.config.use_dynamic_shifting:
+            assert mu is not None, "mu cannot be None when use_dynamic_shifting is True"
+            sigmas_array = self.time_shift(mu, 1.0, sigmas_array)
+        else:
+            sigmas_array = self.shift * sigmas_array / (
+                1 + (self.shift - 1) * sigmas_array)
+
+        # 3. If required, stretch the sigmas schedule to terminate at the configured `shift_terminal` value
+        if self.config.shift_terminal:
+            sigmas_tensor = torch.from_numpy(sigmas_array).to(
+                dtype=torch.float32)
+            sigmas_tensor = self.stretch_shift_to_terminal(sigmas_tensor)
+            sigmas_array = sigmas_tensor.numpy()
+
+        # 4. If required, convert sigmas to one of karras, exponential, or beta sigma schedules
+        if self.config.use_karras_sigmas:
+            sigmas_tensor = torch.from_numpy(sigmas_array).to(
+                dtype=torch.float32)
+            sigmas_tensor = self._convert_to_karras(
+                in_sigmas=sigmas_tensor,
+                num_inference_steps=num_inference_steps)
+            sigmas_array = sigmas_tensor.numpy()
+        elif self.config.use_exponential_sigmas:
+            sigmas_tensor = torch.from_numpy(sigmas_array).to(
+                dtype=torch.float32)
+            sigmas_tensor = self._convert_to_exponential(
+                in_sigmas=sigmas_tensor,
+                num_inference_steps=num_inference_steps)
+            sigmas_array = sigmas_tensor.numpy()
+        elif self.config.use_beta_sigmas:
+            sigmas_tensor = torch.from_numpy(sigmas_array).to(
+                dtype=torch.float32)
+            sigmas_tensor = self._convert_to_beta(
+                in_sigmas=sigmas_tensor,
+                num_inference_steps=num_inference_steps)
+            sigmas_array = sigmas_tensor.numpy()
+
+        # 5. Convert sigmas and timesteps to tensors and move to specified device
+        sigmas_tensor = torch.from_numpy(sigmas_array).to(dtype=torch.float32,
+                                                          device=device)
+        if not is_timesteps_provided:
+            timesteps_tensor = sigmas_tensor * self.config.num_train_timesteps
+        else:
+            assert timesteps_array is not None
+            timesteps_tensor = torch.from_numpy(timesteps_array).to(
+                dtype=torch.float32, device=device)
+
+        # 6. Append the terminal sigma value.
+        #    If a model requires inverted sigma schedule for denoising but timesteps without inversion, the
+        #    `invert_sigmas` flag can be set to `True`. This case is only required in Mochi
+        if self.config.invert_sigmas:
+            sigmas_tensor = 1.0 - sigmas_tensor
+            timesteps_tensor = sigmas_tensor * self.config.num_train_timesteps
+            sigmas_tensor = torch.cat(
+                [sigmas_tensor,
+                 torch.ones(1, device=sigmas_tensor.device)])
+        else:
+            sigmas_tensor = torch.cat(
+                [sigmas_tensor,
+                 torch.zeros(1, device=sigmas_tensor.device)])
+
+        self.timesteps = timesteps_tensor
+        self.sigmas = sigmas_tensor
+        self._step_index = None
+        self._begin_index = None
+
+    def index_for_timestep(
+            self,
+            timestep: float | torch.FloatTensor,
+            schedule_timesteps: torch.Tensor | None = None) -> int:
         if schedule_timesteps is None:
             schedule_timesteps = self.timesteps
 
@@ -174,17 +424,9 @@ class FlowMatchDiscreteScheduler(SchedulerMixin, ConfigMixin, BaseScheduler):
         # case we start in the middle of the denoising schedule (e.g. for image-to-image)
         pos = 1 if len(indices) > 1 else 0
 
-        idx: int = indices[pos].item()
+        return indices[pos].item()
 
-        return idx
-
-    def set_shift(self, shift: float) -> None:
-        self.config.shift = shift
-
-    def set_timesteps_scale(self, timesteps_scale: bool) -> None:
-        self.config.timesteps_scale = timesteps_scale
-
-    def _init_step_index(self, timestep) -> None:
+    def _init_step_index(self, timestep: float | torch.FloatTensor) -> None:
         if self.begin_index is None:
             if isinstance(timestep, torch.Tensor):
                 timestep = timestep.to(self.timesteps.device)
@@ -192,22 +434,19 @@ class FlowMatchDiscreteScheduler(SchedulerMixin, ConfigMixin, BaseScheduler):
         else:
             self._step_index = self._begin_index
 
-    def scale_model_input(self,
-                          sample: torch.Tensor,
-                          timestep: int | None = None) -> torch.Tensor:
-        return sample
-
-    def sd3_time_shift(self, t: torch.Tensor):
-        return (self.config.shift * t) / (1 + (self.config.shift - 1) * t)
-
     def step(
         self,
         model_output: torch.FloatTensor,
-        timestep: float | torch.FloatTensor,
+        timestep: int | torch.Tensor,
         sample: torch.FloatTensor,
+        s_churn: float = 0.0,
+        s_tmin: float = 0.0,
+        s_tmax: float = float("inf"),
+        s_noise: float = 1.0,
+        generator: torch.Generator | None = None,
+        per_token_timesteps: torch.Tensor | None = None,
         return_dict: bool = True,
-        **kwargs,
-    ) -> FlowMatchDiscreteSchedulerOutput | tuple:
+    ) -> FlowMatchEulerDiscreteSchedulerOutput | tuple[torch.FloatTensor, ...]:
         """
         Predict the sample from the previous timestep by reversing the SDE. This function propagates the diffusion
         process from the learned model outputs (most often the predicted noise).
@@ -215,28 +454,34 @@ class FlowMatchDiscreteScheduler(SchedulerMixin, ConfigMixin, BaseScheduler):
         Args:
             model_output (`torch.FloatTensor`):
                 The direct output from learned diffusion model.
-            timestep (`float`):
+            timestep (`int` or `torch.Tensor`):
                 The current discrete timestep in the diffusion chain.
             sample (`torch.FloatTensor`):
                 A current instance of a sample created by the diffusion process.
+            s_churn (`float`):
+            s_tmin  (`float`):
+            s_tmax  (`float`):
+            s_noise (`float`, defaults to 1.0):
+                Scaling factor for noise added to the sample.
             generator (`torch.Generator`, *optional*):
                 A random number generator.
-            n_tokens (`int`, *optional*):
-                Number of tokens in the input sequence.
+            per_token_timesteps (`torch.Tensor`, *optional*):
+                The timesteps for each token in the sample.
             return_dict (`bool`):
-                Whether or not to return a [`~schedulers.scheduling_euler_discrete.EulerDiscreteSchedulerOutput`] or
-                tuple.
+                Whether or not to return a
+                [`~schedulers.scheduling_flow_match_euler_discrete.FlowMatchEulerDiscreteSchedulerOutput`] or tuple.
 
         Returns:
-            [`~schedulers.scheduling_euler_discrete.EulerDiscreteSchedulerOutput`] or `tuple`:
-                If return_dict is `True`, [`~schedulers.scheduling_euler_discrete.EulerDiscreteSchedulerOutput`] is
-                returned, otherwise a tuple is returned where the first element is the sample tensor.
+            [`~schedulers.scheduling_flow_match_euler_discrete.FlowMatchEulerDiscreteSchedulerOutput`] or `tuple`:
+                If return_dict is `True`,
+                [`~schedulers.scheduling_flow_match_euler_discrete.FlowMatchEulerDiscreteSchedulerOutput`] is returned,
+                otherwise a tuple is returned where the first element is the sample tensor.
         """
 
-        if isinstance(timestep, (int | torch.IntTensor | torch.LongTensor)):
+        if (isinstance(timestep, int | torch.IntTensor | torch.LongTensor)):
             raise ValueError((
                 "Passing integer indices (e.g. from `enumerate(timesteps)`) as timesteps to"
-                " `EulerDiscreteScheduler.step()` is not supported. Make sure to pass"
+                " `FlowMatchEulerDiscreteScheduler.step()` is not supported. Make sure to pass"
                 " one of the `scheduler.timesteps` as a timestep."), )
 
         if self.step_index is None:
@@ -245,24 +490,161 @@ class FlowMatchDiscreteScheduler(SchedulerMixin, ConfigMixin, BaseScheduler):
         # Upcast to avoid precision issues when computing prev_sample
         sample = sample.to(torch.float32)
 
-        assert self.step_index is not None
-        dt = self.sigmas[self.step_index + 1] - self.sigmas[self.step_index]
+        if per_token_timesteps is not None:
+            per_token_sigmas = per_token_timesteps / self.config.num_train_timesteps
 
-        if self.config.solver == "euler":
-            prev_sample = sample + model_output.to(torch.float32) * dt
+            sigmas = self.sigmas[:, None, None]
+            lower_mask = sigmas < per_token_sigmas[None] - 1e-6
+            lower_sigmas = lower_mask * sigmas
+            lower_sigmas, _ = lower_sigmas.max(dim=0)
+
+            current_sigma = per_token_sigmas[..., None]
+            next_sigma = lower_sigmas[..., None]
+            dt = current_sigma - next_sigma
         else:
-            raise ValueError(
-                f"Solver {self.config.solver} not supported. Supported solvers: {self.supported_solver}"
-            )
+            assert self.step_index is not None, "step_index should not be None"
+            sigma_idx = self.step_index
+            sigma = self.sigmas[sigma_idx]
+            sigma_next = self.sigmas[sigma_idx + 1]
+
+            current_sigma = sigma
+            next_sigma = sigma_next
+            dt = sigma_next - sigma
+
+        if self.config.stochastic_sampling:
+            x0 = sample - current_sigma * model_output
+            noise = torch.randn_like(sample)
+            prev_sample = (1.0 - next_sigma) * x0 + next_sigma * noise
+        else:
+            prev_sample = sample + dt * model_output
 
         # upon completion increase step index by one
-        assert self._step_index is not None
+        assert self._step_index is not None, "_step_index should not be None"
         self._step_index += 1
+        if per_token_timesteps is None:
+            # Cast sample back to model compatible dtype
+            prev_sample = prev_sample.to(model_output.dtype)
 
         if isinstance(prev_sample, torch.Tensor | float) and not return_dict:
             return (prev_sample, )
 
-        return FlowMatchDiscreteSchedulerOutput(prev_sample=prev_sample)
+        return FlowMatchEulerDiscreteSchedulerOutput(prev_sample=prev_sample)
 
-    def __len__(self):
+    # Copied from diffusers.schedulers.scheduling_euler_discrete.EulerDiscreteScheduler._convert_to_karras
+    def _convert_to_karras(self, in_sigmas: torch.Tensor,
+                           num_inference_steps: int) -> torch.Tensor:
+        """Constructs the noise schedule of Karras et al. (2022)."""
+
+        # Hack to make sure that other schedulers which copy this function don't break
+        # TODO: Add this logic to the other schedulers
+        if hasattr(self.config, "sigma_min"):
+            sigma_min = self.config.sigma_min
+        else:
+            sigma_min = None
+
+        if hasattr(self.config, "sigma_max"):
+            sigma_max = self.config.sigma_max
+        else:
+            sigma_max = None
+
+        sigma_min = sigma_min if sigma_min is not None else in_sigmas[-1].item()
+        sigma_max = sigma_max if sigma_max is not None else in_sigmas[0].item()
+
+        rho = 7.0  # 7.0 is the value used in the paper
+        ramp = np.linspace(0, 1, num_inference_steps)
+        min_inv_rho = sigma_min**(1 / rho)
+        max_inv_rho = sigma_max**(1 / rho)
+        sigmas = (max_inv_rho + ramp * (min_inv_rho - max_inv_rho))**rho
+        return sigmas
+
+    # Copied from diffusers.schedulers.scheduling_euler_discrete.EulerDiscreteScheduler._convert_to_exponential
+    def _convert_to_exponential(self, in_sigmas: torch.Tensor,
+                                num_inference_steps: int) -> torch.Tensor:
+        """Constructs an exponential noise schedule."""
+
+        # Hack to make sure that other schedulers which copy this function don't break
+        # TODO: Add this logic to the other schedulers
+        if hasattr(self.config, "sigma_min"):
+            sigma_min = self.config.sigma_min
+        else:
+            sigma_min = None
+
+        if hasattr(self.config, "sigma_max"):
+            sigma_max = self.config.sigma_max
+        else:
+            sigma_max = None
+
+        sigma_min = sigma_min if sigma_min is not None else in_sigmas[-1].item()
+        sigma_max = sigma_max if sigma_max is not None else in_sigmas[0].item()
+
+        sigmas = np.exp(
+            np.linspace(math.log(sigma_max), math.log(sigma_min),
+                        num_inference_steps))
+        return sigmas
+
+    # Copied from diffusers.schedulers.scheduling_euler_discrete.EulerDiscreteScheduler._convert_to_beta
+    def _convert_to_beta(self,
+                         in_sigmas: torch.Tensor,
+                         num_inference_steps: int,
+                         alpha: float = 0.6,
+                         beta: float = 0.6) -> torch.Tensor:
+        """From "Beta Sampling is All You Need" [arXiv:2407.12173] (Lee et. al, 2024)"""
+
+        # Hack to make sure that other schedulers which copy this function don't break
+        # TODO: Add this logic to the other schedulers
+        if hasattr(self.config, "sigma_min"):
+            sigma_min = self.config.sigma_min
+        else:
+            sigma_min = None
+
+        if hasattr(self.config, "sigma_max"):
+            sigma_max = self.config.sigma_max
+        else:
+            sigma_max = None
+
+        sigma_min = sigma_min if sigma_min is not None else in_sigmas[-1].item()
+        sigma_max = sigma_max if sigma_max is not None else in_sigmas[0].item()
+
+        sigmas = np.array([
+            sigma_min + (ppf * (sigma_max - sigma_min)) for ppf in [
+                scipy.stats.beta.ppf(timestep, alpha, beta)
+                for timestep in 1 - np.linspace(0, 1, num_inference_steps)
+            ]
+        ])
+        return sigmas
+
+    def _time_shift_exponential(
+            self, mu: float, sigma: float,
+            t: torch.Tensor | np.ndarray) -> torch.Tensor | np.ndarray:
+        if isinstance(t, np.ndarray):
+            return np.exp(mu) / (np.exp(mu) + (1 / t - 1)**sigma)
+        else:
+            return math.exp(mu) / (math.exp(mu) + (1 / t - 1)**sigma)
+
+    def _time_shift_linear(
+            self, mu: float, sigma: float,
+            t: torch.Tensor | np.ndarray) -> torch.Tensor | np.ndarray:
+        return mu / (mu + (1 / t - 1)**sigma)
+
+    def add_noise(
+        self,
+        clean_latent: torch.Tensor,
+        noise: torch.Tensor,
+        timestep: torch.IntTensor,
+    ) -> torch.Tensor:
+        self.sigmas = self.sigmas.to(noise.device)
+        timestep = timestep.expand(clean_latent.shape[0])
+        self.timesteps = self.timesteps.to(noise.device)
+        timestep_id = torch.argmin(
+            (self.timesteps.unsqueeze(0) - timestep.unsqueeze(1)).abs(), dim=1)
+        sigma = self.sigmas[timestep_id].reshape(-1, 1, 1, 1)
+        sample = (1 - sigma) * clean_latent + sigma * noise
+        return sample.type_as(noise)
+
+    def scale_model_input(self,
+                          sample: torch.Tensor,
+                          timestep: int | None = None) -> torch.Tensor:
+        return sample
+
+    def __len__(self) -> int:
         return self.config.num_train_timesteps
diff --git a/fastvideo/v1/models/schedulers/scheduling_flow_unipc_multistep.py b/fastvideo/v1/models/schedulers/scheduling_flow_unipc_multistep.py
index 85a4f0f73..40f86b634 100644
--- a/fastvideo/v1/models/schedulers/scheduling_flow_unipc_multistep.py
+++ b/fastvideo/v1/models/schedulers/scheduling_flow_unipc_multistep.py
@@ -12,13 +12,10 @@ from diffusers.configuration_utils import ConfigMixin, register_to_config
 from diffusers.schedulers.scheduling_utils import (KarrasDiffusionSchedulers,
                                                    SchedulerMixin,
                                                    SchedulerOutput)
-from diffusers.utils import deprecate, is_scipy_available
+from diffusers.utils import deprecate
 
 from fastvideo.v1.models.schedulers.base import BaseScheduler
 
-if is_scipy_available():
-    pass
-
 
 class FlowUniPCMultistepScheduler(SchedulerMixin, ConfigMixin, BaseScheduler):
     """
diff --git a/fastvideo/v1/models/schedulers/scheduling_unipc_multistep.py b/fastvideo/v1/models/schedulers/scheduling_unipc_multistep.py
index 1e3966354..2c06ef66c 100644
--- a/fastvideo/v1/models/schedulers/scheduling_unipc_multistep.py
+++ b/fastvideo/v1/models/schedulers/scheduling_unipc_multistep.py
@@ -25,18 +25,16 @@
 import math
 
 import numpy as np
+import scipy.stats
 import torch
 from diffusers.configuration_utils import ConfigMixin, register_to_config
 from diffusers.schedulers.scheduling_utils import (KarrasDiffusionSchedulers,
                                                    SchedulerMixin,
                                                    SchedulerOutput)
-from diffusers.utils import deprecate, is_scipy_available
+from diffusers.utils import deprecate
 
 from fastvideo.v1.models.schedulers.base import BaseScheduler
 
-if is_scipy_available():
-    import scipy.stats
-
 
 # Copied from diffusers.schedulers.scheduling_ddpm.betas_for_alpha_bar
 def betas_for_alpha_bar(
@@ -223,9 +221,6 @@ class UniPCMultistepScheduler(SchedulerMixin, ConfigMixin, BaseScheduler):
         final_sigmas_type: str | None = "zero",  # "zero", "sigma_min"
         rescale_betas_zero_snr: bool = False,
     ):
-        if self.config.use_beta_sigmas and not is_scipy_available():
-            raise ImportError(
-                "Make sure to install scipy if you want to use beta sigmas.")
         if sum([
                 self.config.use_beta_sigmas, self.config.use_exponential_sigmas,
                 self.config.use_karras_sigmas
diff --git a/fastvideo/v1/pipelines/stages/__init__.py b/fastvideo/v1/pipelines/stages/__init__.py
index 8471da71d..af94f505f 100644
--- a/fastvideo/v1/pipelines/stages/__init__.py
+++ b/fastvideo/v1/pipelines/stages/__init__.py
@@ -9,7 +9,8 @@ complete diffusion pipelines.
 from fastvideo.v1.pipelines.stages.base import PipelineStage
 from fastvideo.v1.pipelines.stages.conditioning import ConditioningStage
 from fastvideo.v1.pipelines.stages.decoding import DecodingStage
-from fastvideo.v1.pipelines.stages.denoising import DenoisingStage
+from fastvideo.v1.pipelines.stages.denoising import (DenoisingStage,
+                                                     DmdDenoisingStage)
 from fastvideo.v1.pipelines.stages.encoding import EncodingStage
 from fastvideo.v1.pipelines.stages.image_encoding import ImageEncodingStage
 from fastvideo.v1.pipelines.stages.input_validation import InputValidationStage
@@ -28,6 +29,7 @@ __all__ = [
     "LatentPreparationStage",
     "ConditioningStage",
     "DenoisingStage",
+    "DmdDenoisingStage",
     "EncodingStage",
     "DecodingStage",
     "ImageEncodingStage",
diff --git a/fastvideo/v1/pipelines/stages/base.py b/fastvideo/v1/pipelines/stages/base.py
index 7a6650fb5..6f54ef4bb 100644
--- a/fastvideo/v1/pipelines/stages/base.py
+++ b/fastvideo/v1/pipelines/stages/base.py
@@ -147,21 +147,20 @@ class PipelineStage(ABC):
         # Execute the actual stage logic
         # envs.ENABLE_STAGE_LOGGING
         if False:
-            self._logger.info("[%s] Starting execution", stage_name)
+            logger.info("[%s] Starting execution", stage_name)
             start_time = time.perf_counter()
 
             try:
                 result = self.forward(batch, fastvideo_args)
                 execution_time = time.perf_counter() - start_time
-                self._logger.info("[%s] Execution completed in %s ms",
-                                  stage_name, execution_time * 1000)
+                logger.info("[%s] Execution completed in %s ms", stage_name,
+                            execution_time * 1000)
             except Exception as e:
                 execution_time = time.perf_counter() - start_time
-                self._logger.error(
-                    "[%s] Error during execution after %s ms: %s", stage_name,
-                    execution_time * 1000, e)
-                self._logger.error("[%s] Traceback: %s", stage_name,
-                                   traceback.format_exc())
+                logger.error("[%s] Error during execution after %s ms: %s",
+                             stage_name, execution_time * 1000, e)
+                logger.error("[%s] Traceback: %s", stage_name,
+                             traceback.format_exc())
                 raise
         else:
             # Direct execution (current behavior)
diff --git a/fastvideo/v1/pipelines/stages/denoising.py b/fastvideo/v1/pipelines/stages/denoising.py
index 57955f552..07a4aad15 100644
--- a/fastvideo/v1/pipelines/stages/denoising.py
+++ b/fastvideo/v1/pipelines/stages/denoising.py
@@ -24,6 +24,8 @@ from fastvideo.v1.fastvideo_args import FastVideoArgs
 from fastvideo.v1.forward_context import set_forward_context
 from fastvideo.v1.logger import init_logger
 from fastvideo.v1.models.loader.component_loader import TransformerLoader
+from fastvideo.v1.models.schedulers.scheduling_flow_match_euler_discrete import (
+    FlowMatchEulerDiscreteScheduler)
 from fastvideo.v1.pipelines.pipeline_batch_info import ForwardBatch
 from fastvideo.v1.pipelines.stages.base import PipelineStage
 from fastvideo.v1.pipelines.stages.validators import StageValidators as V
@@ -567,3 +569,226 @@ class DenoisingStage(PipelineStage):
         result.add_check("latents", batch.latents,
                          [V.is_tensor, V.with_dims(5)])
         return result
+
+
+class DmdDenoisingStage(DenoisingStage):
+    """
+    Denoising stage for DMD.
+    """
+
+    def __init__(self, transformer, scheduler) -> None:
+        super().__init__(transformer, scheduler)
+        self.scheduler = FlowMatchEulerDiscreteScheduler(shift=8.0)
+
+    def forward(
+        self,
+        batch: ForwardBatch,
+        fastvideo_args: FastVideoArgs,
+    ) -> ForwardBatch:
+        """
+        Run the denoising loop.
+        
+        Args:
+            batch: The current batch information.
+            fastvideo_args: The inference arguments.
+            
+        Returns:
+            The batch with denoised latents.
+        """
+        # Setup precision and autocast settings
+        # TODO(will): make the precision configurable for inference
+        # target_dtype = PRECISION_TO_TYPE[fastvideo_args.precision]
+        target_dtype = torch.bfloat16
+        autocast_enabled = (target_dtype != torch.float32
+                            ) and not fastvideo_args.disable_autocast
+
+        # Get timesteps and calculate warmup steps
+        timesteps = batch.timesteps
+
+        # TODO(will): remove this once we add input/output validation for stages
+        if timesteps is None:
+            raise ValueError("Timesteps must be provided")
+        num_inference_steps = batch.num_inference_steps
+        num_warmup_steps = len(
+            timesteps) - num_inference_steps * self.scheduler.order
+
+        # Prepare image latents and embeddings for I2V generation
+        image_embeds = batch.image_embeds
+        if len(image_embeds) > 0:
+            assert torch.isnan(image_embeds[0]).sum() == 0
+            image_embeds = [
+                image_embed.to(target_dtype) for image_embed in image_embeds
+            ]
+
+        image_kwargs = self.prepare_extra_func_kwargs(
+            self.transformer.forward,
+            {
+                "encoder_hidden_states_image": image_embeds,
+                "mask_strategy": dict_to_3d_list(
+                    None, t_max=50, l_max=60, h_max=24)
+            },
+        )
+
+        pos_cond_kwargs = self.prepare_extra_func_kwargs(
+            self.transformer.forward,
+            {
+                "encoder_hidden_states_2": batch.clip_embedding_pos,
+                "encoder_attention_mask": batch.prompt_attention_mask,
+            },
+        )
+
+        # Prepare STA parameters
+        if st_attn_available and self.attn_backend == SlidingTileAttentionBackend:
+            self.prepare_sta_param(batch, fastvideo_args)
+
+        # Get latents and embeddings
+        assert batch.latents is not None, "latents must be provided"
+        latents = batch.latents
+        # TODO(yongqi) hard code prepare latents
+        latents = torch.randn(
+            latents.permute(0, 2, 1, 3, 4).shape,
+            dtype=torch.bfloat16,
+            device="cuda",
+            generator=torch.Generator(device="cuda").manual_seed(42))
+        video_raw_latent_shape = latents.shape
+        prompt_embeds = batch.prompt_embeds
+        assert torch.isnan(prompt_embeds[0]).sum() == 0
+        timesteps = torch.tensor(
+            fastvideo_args.pipeline_config.dmd_denoising_steps,
+            dtype=torch.long,
+            device=get_local_torch_device())
+
+        # Handle sequence parallelism if enabled
+        sp_world_size, rank_in_sp_group = get_sp_world_size(
+        ), get_sp_parallel_rank()
+        sp_group = sp_world_size > 1
+        if sp_group:
+            latents = rearrange(latents,
+                                "b (n t) c h w -> b n t c h w",
+                                n=sp_world_size).contiguous()
+            latents = latents[:, rank_in_sp_group, :, :, :, :]
+            if batch.image_latent is not None:
+                image_latent = rearrange(batch.image_latent,
+                                         "b c (n t) h w -> b c n t h w",
+                                         n=sp_world_size).contiguous()
+
+                image_latent = image_latent[:, :, rank_in_sp_group, :, :, :]
+                batch.image_latent = image_latent
+
+        # Run denoising loop
+        with self.progress_bar(total=len(timesteps)) as progress_bar:
+            for i, t in enumerate(timesteps):
+                # Skip if interrupted
+                if hasattr(self, 'interrupt') and self.interrupt:
+                    continue
+                # Expand latents for I2V
+                noise_latents = latents.clone()
+                latent_model_input = latents.to(target_dtype)
+
+                if batch.image_latent is not None:
+                    latent_model_input = torch.cat([
+                        latent_model_input,
+                        batch.image_latent.permute(0, 2, 1, 3, 4)
+                    ],
+                                                   dim=2).to(target_dtype)
+                assert torch.isnan(latent_model_input).sum() == 0
+
+                # Prepare inputs for transformer
+                t_expand = t.repeat(latent_model_input.shape[0])
+                guidance_expand = (
+                    torch.tensor(
+                        [fastvideo_args.pipeline_config.embedded_cfg_scale] *
+                        latent_model_input.shape[0],
+                        dtype=torch.float32,
+                        device=get_local_torch_device(),
+                    ).to(target_dtype) *
+                    1000.0 if fastvideo_args.pipeline_config.embedded_cfg_scale
+                    is not None else None)
+
+                # Predict noise residual
+                with torch.autocast(device_type="cuda",
+                                    dtype=target_dtype,
+                                    enabled=autocast_enabled):
+                    if (st_attn_available
+                            and self.attn_backend == SlidingTileAttentionBackend
+                        ) or (vsa_available and self.attn_backend
+                              == VideoSparseAttentionBackend):
+                        self.attn_metadata_builder_cls = self.attn_backend.get_builder_cls(
+                        )
+
+                        if self.attn_metadata_builder_cls is not None:
+                            self.attn_metadata_builder = self.attn_metadata_builder_cls(
+                            )
+                            # TODO(will): clean this up
+                            attn_metadata = self.attn_metadata_builder.build(
+                                current_timestep=i,
+                                forward_batch=batch,
+                                fastvideo_args=fastvideo_args,
+                            )
+                            assert attn_metadata is not None, "attn_metadata cannot be None"
+                        else:
+                            attn_metadata = None
+                    else:
+                        attn_metadata = None
+
+                    batch.is_cfg_negative = False
+                    with set_forward_context(
+                            current_timestep=i,
+                            attn_metadata=attn_metadata,
+                            forward_batch=batch,
+                            # fastvideo_args=fastvideo_args
+                    ):
+                        # Run transformer
+                        pred_noise = self.transformer(
+                            latent_model_input.permute(0, 2, 1, 3, 4),
+                            prompt_embeds,
+                            t_expand,
+                            guidance=guidance_expand,
+                            **image_kwargs,
+                            **pos_cond_kwargs,
+                        ).permute(0, 2, 1, 3, 4)
+
+                    from fastvideo.v1.training.training_utils import (
+                        pred_noise_to_pred_video)
+                    pred_video = pred_noise_to_pred_video(
+                        pred_noise=pred_noise.flatten(0, 1),
+                        noise_input_latent=noise_latents.flatten(0, 1),
+                        timestep=t_expand,
+                        scheduler=self.scheduler).unflatten(
+                            0, pred_noise.shape[:2])
+
+                    if i < len(timesteps) - 1:
+                        next_timestep = timesteps[i + 1] * torch.ones(
+                            pred_video.shape[:2],
+                            dtype=torch.long,
+                            device=pred_video.device)
+                        noise = torch.randn(video_raw_latent_shape,
+                                            device=self.device,
+                                            dtype=pred_video.dtype)
+                        if sp_group:
+                            noise = rearrange(noise,
+                                              "b (n t) c h w -> b n t c h w",
+                                              n=sp_world_size).contiguous()
+                            noise = noise[:, rank_in_sp_group, :, :, :, :]
+                        latents = self.scheduler.add_noise(
+                            pred_video.flatten(0, 1), noise.flatten(0, 1),
+                            next_timestep.flatten(0, 1)).unflatten(
+                                0, pred_video.shape[:2])
+                    else:
+                        latents = pred_video
+
+                    # Update progress bar
+                    if i == len(timesteps) - 1 or (
+                        (i + 1) > num_warmup_steps and
+                        (i + 1) % self.scheduler.order == 0
+                            and progress_bar is not None):
+                        progress_bar.update()
+
+        # Gather results if using sequence parallelism
+        if sp_group:
+            latents = sequence_model_parallel_all_gather(latents, dim=1)
+        latents = latents.permute(0, 2, 1, 3, 4)
+        # Update batch with final latents
+        batch.latents = latents
+
+        return batch
diff --git a/fastvideo/v1/pipelines/wan/wan_dmd_pipeline.py b/fastvideo/v1/pipelines/wan/wan_dmd_pipeline.py
new file mode 100644
index 000000000..79cc4bda3
--- /dev/null
+++ b/fastvideo/v1/pipelines/wan/wan_dmd_pipeline.py
@@ -0,0 +1,71 @@
+# SPDX-License-Identifier: Apache-2.0
+"""
+Wan video diffusion pipeline implementation.
+
+This module contains an implementation of the Wan video diffusion pipeline
+using the modular pipeline architecture.
+"""
+
+from fastvideo.v1.fastvideo_args import FastVideoArgs
+from fastvideo.v1.logger import init_logger
+from fastvideo.v1.models.schedulers.scheduling_flow_match_euler_discrete import (
+    FlowMatchEulerDiscreteScheduler)
+from fastvideo.v1.pipelines import ComposedPipelineBase, LoRAPipeline
+
+# isort: off
+from fastvideo.v1.pipelines.stages import (
+    ConditioningStage, DecodingStage, DmdDenoisingStage, InputValidationStage,
+    LatentPreparationStage, TextEncodingStage, TimestepPreparationStage)
+# isort: on
+
+logger = init_logger(__name__)
+
+
+class WanDMDPipeline(LoRAPipeline, ComposedPipelineBase):
+    """
+    Wan video diffusion pipeline with LoRA support.
+    """
+
+    _required_config_modules = [
+        "text_encoder", "tokenizer", "vae", "transformer", "scheduler"
+    ]
+
+    def initialize_pipeline(self, fastvideo_args: FastVideoArgs):
+        # We use UniPCMScheduler from Wan2.1 official repo, not the one in diffusers.
+        self.modules["scheduler"] = FlowMatchEulerDiscreteScheduler(
+            shift=fastvideo_args.pipeline_config.flow_shift)
+
+    def create_pipeline_stages(self, fastvideo_args: FastVideoArgs) -> None:
+        """Set up pipeline stages with proper dependency injection."""
+
+        self.add_stage(stage_name="input_validation_stage",
+                       stage=InputValidationStage())
+
+        self.add_stage(stage_name="prompt_encoding_stage",
+                       stage=TextEncodingStage(
+                           text_encoders=[self.get_module("text_encoder")],
+                           tokenizers=[self.get_module("tokenizer")],
+                       ))
+
+        self.add_stage(stage_name="conditioning_stage",
+                       stage=ConditioningStage())
+
+        self.add_stage(stage_name="timestep_preparation_stage",
+                       stage=TimestepPreparationStage(
+                           scheduler=self.get_module("scheduler")))
+
+        self.add_stage(stage_name="latent_preparation_stage",
+                       stage=LatentPreparationStage(
+                           scheduler=self.get_module("scheduler"),
+                           transformer=self.get_module("transformer", None)))
+
+        self.add_stage(stage_name="denoising_stage",
+                       stage=DmdDenoisingStage(
+                           transformer=self.get_module("transformer"),
+                           scheduler=self.get_module("scheduler")))
+
+        self.add_stage(stage_name="decoding_stage",
+                       stage=DecodingStage(vae=self.get_module("vae")))
+
+
+EntryClass = WanDMDPipeline
diff --git a/fastvideo/v1/tests/ssim/L40S_reference_videos/FastHunyuan-diffusers/FLASH_ATTN/Will Smith casually eats noodles, his relaxed demeanor contrasting with the energetic background of .mp4 b/fastvideo/v1/tests/ssim/L40S_reference_videos/FastHunyuan-diffusers/FLASH_ATTN/Will Smith casually eats noodles, his relaxed demeanor contrasting with the energetic background of .mp4
index 95f1162d2..5cfa926b0 100644
Binary files a/fastvideo/v1/tests/ssim/L40S_reference_videos/FastHunyuan-diffusers/FLASH_ATTN/Will Smith casually eats noodles, his relaxed demeanor contrasting with the energetic background of .mp4 and b/fastvideo/v1/tests/ssim/L40S_reference_videos/FastHunyuan-diffusers/FLASH_ATTN/Will Smith casually eats noodles, his relaxed demeanor contrasting with the energetic background of .mp4 differ
diff --git a/fastvideo/v1/tests/ssim/L40S_reference_videos/FastHunyuan-diffusers/TORCH_SDPA/Will Smith casually eats noodles, his relaxed demeanor contrasting with the energetic background of .mp4 b/fastvideo/v1/tests/ssim/L40S_reference_videos/FastHunyuan-diffusers/TORCH_SDPA/Will Smith casually eats noodles, his relaxed demeanor contrasting with the energetic background of .mp4
index d27b708fa..797ae1d61 100644
Binary files a/fastvideo/v1/tests/ssim/L40S_reference_videos/FastHunyuan-diffusers/TORCH_SDPA/Will Smith casually eats noodles, his relaxed demeanor contrasting with the energetic background of .mp4 and b/fastvideo/v1/tests/ssim/L40S_reference_videos/FastHunyuan-diffusers/TORCH_SDPA/Will Smith casually eats noodles, his relaxed demeanor contrasting with the energetic background of .mp4 differ
diff --git a/fastvideo/v1/tests/ssim/L40S_reference_videos/Wan2.1-I2V-14B-480P-Diffusers/FLASH_ATTN/An astronaut hatching from an egg, on the surface of the moon, the darkness and depth of space reali.mp4 b/fastvideo/v1/tests/ssim/L40S_reference_videos/Wan2.1-I2V-14B-480P-Diffusers/FLASH_ATTN/An astronaut hatching from an egg, on the surface of the moon, the darkness and depth of space reali.mp4
index c8632bf89..49dac8f7f 100644
Binary files a/fastvideo/v1/tests/ssim/L40S_reference_videos/Wan2.1-I2V-14B-480P-Diffusers/FLASH_ATTN/An astronaut hatching from an egg, on the surface of the moon, the darkness and depth of space reali.mp4 and b/fastvideo/v1/tests/ssim/L40S_reference_videos/Wan2.1-I2V-14B-480P-Diffusers/FLASH_ATTN/An astronaut hatching from an egg, on the surface of the moon, the darkness and depth of space reali.mp4 differ
diff --git a/fastvideo/v1/tests/ssim/L40S_reference_videos/Wan2.1-I2V-14B-480P-Diffusers/TORCH_SDPA/An astronaut hatching from an egg, on the surface of the moon, the darkness and depth of space reali.mp4 b/fastvideo/v1/tests/ssim/L40S_reference_videos/Wan2.1-I2V-14B-480P-Diffusers/TORCH_SDPA/An astronaut hatching from an egg, on the surface of the moon, the darkness and depth of space reali.mp4
index 849841b47..6cf33b1db 100644
Binary files a/fastvideo/v1/tests/ssim/L40S_reference_videos/Wan2.1-I2V-14B-480P-Diffusers/TORCH_SDPA/An astronaut hatching from an egg, on the surface of the moon, the darkness and depth of space reali.mp4 and b/fastvideo/v1/tests/ssim/L40S_reference_videos/Wan2.1-I2V-14B-480P-Diffusers/TORCH_SDPA/An astronaut hatching from an egg, on the surface of the moon, the darkness and depth of space reali.mp4 differ
diff --git a/fastvideo/v1/tests/ssim/L40S_reference_videos/Wan2.1-T2V-1.3B-Diffusers/FLASH_ATTN/Will Smith casually eats noodles, his relaxed demeanor contrasting with the energetic background of .mp4 b/fastvideo/v1/tests/ssim/L40S_reference_videos/Wan2.1-T2V-1.3B-Diffusers/FLASH_ATTN/Will Smith casually eats noodles, his relaxed demeanor contrasting with the energetic background of .mp4
index 3a9fdac79..a454aee9e 100644
Binary files a/fastvideo/v1/tests/ssim/L40S_reference_videos/Wan2.1-T2V-1.3B-Diffusers/FLASH_ATTN/Will Smith casually eats noodles, his relaxed demeanor contrasting with the energetic background of .mp4 and b/fastvideo/v1/tests/ssim/L40S_reference_videos/Wan2.1-T2V-1.3B-Diffusers/FLASH_ATTN/Will Smith casually eats noodles, his relaxed demeanor contrasting with the energetic background of .mp4 differ
diff --git a/fastvideo/v1/tests/ssim/L40S_reference_videos/Wan2.1-T2V-1.3B-Diffusers/TORCH_SDPA/Will Smith casually eats noodles, his relaxed demeanor contrasting with the energetic background of .mp4 b/fastvideo/v1/tests/ssim/L40S_reference_videos/Wan2.1-T2V-1.3B-Diffusers/TORCH_SDPA/Will Smith casually eats noodles, his relaxed demeanor contrasting with the energetic background of .mp4
index 7f158dc72..b0da7496d 100644
Binary files a/fastvideo/v1/tests/ssim/L40S_reference_videos/Wan2.1-T2V-1.3B-Diffusers/TORCH_SDPA/Will Smith casually eats noodles, his relaxed demeanor contrasting with the energetic background of .mp4 and b/fastvideo/v1/tests/ssim/L40S_reference_videos/Wan2.1-T2V-1.3B-Diffusers/TORCH_SDPA/Will Smith casually eats noodles, his relaxed demeanor contrasting with the energetic background of .mp4 differ
diff --git a/fastvideo/v1/tests/ssim/update_reference_videos.sh b/fastvideo/v1/tests/ssim/update_reference_videos.sh
index 85d132f15..e137c1b44 100644
--- a/fastvideo/v1/tests/ssim/update_reference_videos.sh
+++ b/fastvideo/v1/tests/ssim/update_reference_videos.sh
@@ -7,7 +7,7 @@ set -e  # Exit on any error
 
 # Define directory paths
 GENERATED_DIR="generated_videos"
-REFERENCE_DIR="set_me_to_correct_path"
+REFERENCE_DIR="L40S_reference_videos"
 
 # Colors for output
 RED='\033[0;31m'
diff --git a/fastvideo/v1/training/training_utils.py b/fastvideo/v1/training/training_utils.py
index c2e8576a4..afc94ee69 100644
--- a/fastvideo/v1/training/training_utils.py
+++ b/fastvideo/v1/training/training_utils.py
@@ -555,3 +555,24 @@ def custom_to_hf_state_dict(
             new_state_dict[training_key] = v
 
     return new_state_dict
+
+
+def pred_noise_to_pred_video(pred_noise: torch.Tensor,
+                             noise_input_latent: torch.Tensor,
+                             timestep: torch.Tensor,
+                             scheduler: Any) -> torch.Tensor:
+    """
+    Convert predicted noise to clean latent.
+    """
+    timestep = timestep.expand(noise_input_latent.shape[0])
+    dtype = pred_noise.dtype
+    device = pred_noise.device
+    pred_noise = pred_noise.float().to(device)
+    noise_input_latent = noise_input_latent.float().to(device)
+    sigmas = scheduler.sigmas.float().to(device)
+    timesteps = scheduler.timesteps.float().to(device)
+    timestep_id = torch.argmin(
+        (timesteps.unsqueeze(0) - timestep.unsqueeze(1)).abs(), dim=1)
+    sigma_t = sigmas[timestep_id].reshape(-1, 1, 1, 1)
+    pred_video = noise_input_latent - sigma_t * pred_noise
+    return pred_video.to(dtype)