140 lines
5.0 KiB
Python
140 lines
5.0 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
"""On-device (MLX) DMD sampling for the FastWan runtime.
|
|
|
|
The hybrid proof-of-concept ran the FastWan DiT in MLX but bounced every
|
|
denoising step back through torch/NumPy to run the DMD scheduler math
|
|
(``MLX -> np.array -> torch (CPU) -> np.array -> MLX``). That host round-trip
|
|
forces a full device sync per step and defeats MLX's lazy graph execution.
|
|
|
|
This module mirrors the exact DMD arithmetic from
|
|
``fastvideo/models/utils.py::pred_noise_to_pred_video`` and
|
|
``FlowMatchEulerDiscreteScheduler.add_noise`` while keeping every large tensor
|
|
on the MLX device. The schedule lookup (``argmin`` over the ~1000-entry
|
|
training schedule) is done once on the host in NumPy: it is tiny, it is the
|
|
same value torch would compute, and it sidesteps the reduction-index quirk that
|
|
affects ``argmin`` on the Metal/MPS backends (see the CPU fallbacks in
|
|
``fastvideo/models/utils.py`` and ``scheduling_flow_match_euler_discrete.py``).
|
|
|
|
Because the DMD loop applies a single scalar timestep per step, ``sigma`` is a
|
|
scalar and the update is a plain elementwise affine combination — no
|
|
permute/flatten reshaping is required.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
from dataclasses import dataclass
|
|
from typing import TYPE_CHECKING, Any
|
|
|
|
import numpy as np
|
|
|
|
if TYPE_CHECKING: # pragma: no cover - typing only
|
|
import mlx.core as mx
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class MLXDMDSchedule:
|
|
"""Host-side copy of a flow-match scheduler's ``(sigmas, timesteps)``.
|
|
|
|
Holds the full training schedule so a DMD timestep (e.g. one of
|
|
``1000, 757, 522``) can be mapped to its flow-match ``sigma`` with the same
|
|
nearest-timestep lookup the torch path uses.
|
|
"""
|
|
|
|
sigmas: np.ndarray
|
|
timesteps: np.ndarray
|
|
|
|
@classmethod
|
|
def from_torch_scheduler(cls, scheduler: Any) -> MLXDMDSchedule:
|
|
"""Snapshot ``scheduler.sigmas`` / ``scheduler.timesteps`` to NumPy.
|
|
|
|
Matches ``pred_noise_to_pred_video`` / ``add_noise``, which index the
|
|
scheduler's *full* training schedule (not the per-inference subset).
|
|
"""
|
|
sigmas = scheduler.sigmas.detach().to("cpu").double().numpy()
|
|
timesteps = scheduler.timesteps.detach().to("cpu").double().numpy()
|
|
return cls(sigmas=np.asarray(sigmas), timesteps=np.asarray(timesteps))
|
|
|
|
def sigma_for(self, timestep: float) -> float:
|
|
"""
|
|
Find the sigma associated with the scheduled timestep nearest to the given timestep.
|
|
|
|
Parameters:
|
|
timestep (float): Timestep for which to find the nearest scheduled sigma.
|
|
|
|
Returns:
|
|
float: Sigma associated with the nearest scheduled timestep.
|
|
"""
|
|
idx = int(np.argmin(np.abs(self.timesteps - float(timestep))))
|
|
return float(self.sigmas[idx])
|
|
|
|
|
|
def pred_noise_to_pred_video(
|
|
pred_noise: mx.array,
|
|
noise_input_latent: mx.array,
|
|
sigma: float,
|
|
) -> mx.array:
|
|
"""
|
|
Compute the clean latent prediction from a flow-matching noise prediction.
|
|
|
|
Parameters:
|
|
pred_noise (mx.array): Predicted noise.
|
|
noise_input_latent (mx.array): Noised latent input.
|
|
sigma (float): Noise level used for the prediction.
|
|
|
|
Returns:
|
|
mx.array: Predicted clean latent.
|
|
"""
|
|
return noise_input_latent - sigma * pred_noise
|
|
|
|
|
|
def add_noise(
|
|
clean_latent: mx.array,
|
|
noise: mx.array,
|
|
sigma: float,
|
|
) -> mx.array:
|
|
"""Flow-match forward noising, mirroring the scheduler's ``add_noise``.
|
|
|
|
``sample = (1 - sigma) * clean_latent + sigma * noise``.
|
|
"""
|
|
return (1.0 - sigma) * clean_latent + sigma * noise
|
|
|
|
|
|
def dmd_step(
|
|
*,
|
|
latents: mx.array,
|
|
noise_input_latent: mx.array,
|
|
pred_noise: mx.array,
|
|
schedule: MLXDMDSchedule,
|
|
timestep: float,
|
|
next_timestep: float | None,
|
|
noise: mx.array | None = None,
|
|
) -> mx.array:
|
|
"""
|
|
Compute one DMD sampling update, optionally re-noising the clean latent prediction.
|
|
|
|
Args:
|
|
latents: Retained for call-site compatibility and not used in the update.
|
|
noise_input_latent: Noisy latent used to compute the clean prediction.
|
|
pred_noise: Predicted noise or velocity.
|
|
schedule: Flow-matching schedule used to map timesteps to sigmas.
|
|
timestep: Current sampling timestep.
|
|
next_timestep: Timestep for the next update, or `None` for the final step.
|
|
noise: Fresh noise used for re-noising intermediate steps.
|
|
|
|
Returns:
|
|
The re-noised latent for the next step or the clean latent prediction on
|
|
the final step.
|
|
|
|
Raises:
|
|
ValueError: If `next_timestep` is provided without `noise`.
|
|
"""
|
|
del latents # symmetry with the torch loop; not needed for the math.
|
|
sigma = schedule.sigma_for(timestep)
|
|
pred_video = pred_noise_to_pred_video(pred_noise, noise_input_latent, sigma)
|
|
if next_timestep is None:
|
|
return pred_video
|
|
if noise is None:
|
|
raise ValueError("dmd_step requires `noise` when `next_timestep` is set (re-noise step).")
|
|
sigma_next = schedule.sigma_for(next_timestep)
|
|
return add_noise(pred_video, noise, sigma_next)
|