scheduling_flow_map_euler_discrete

class FlowMapEulerDiscreteScheduler(BaseScheduler):
    """Minimal flow-map scheduler.
Parameters
----------
num_train_timesteps:
    Discretization granularity for training. ``t`` is expressed in
    absolute units in ``[0, num_train_timesteps]``.
shift:
    Flow-matching shift parameter (Wan video default: ``5.0``). Set
    to ``1.0`` for an identity shift.
"""

order: int = 1

def __init__(
    self,
    *,
    num_train_timesteps: int = 1000,
    shift: float = 1.0,
) -> None:
    self.num_train_timesteps = int(num_train_timesteps)
    self.shift = float(shift)
    self.timesteps: torch.Tensor = torch.empty(0)
    self.sigmas: torch.Tensor = torch.empty(0)
    super().__init__()

# ------------------------------------------------------------------
# BaseScheduler abstract surface

def set_shift(self, shift: float) -> None:
    self.shift = float(shift)

def scale_model_input(
    self,
    sample: torch.Tensor,
    timestep: int | None = None,
) -> torch.Tensor:
    # Flow-matching has no per-step input scaling; pass through.
    del timestep
    return sample

# ------------------------------------------------------------------
# Public helpers used by the AnyFlow pretrain method.

def apply_shift(
    self,
    t: torch.Tensor,
    *,
    shift: float | None = None,
) -> torch.Tensor:
    """Apply the flow-matching shift: ``t' = s * t / (1 + (s - 1) * t)``.

    Operates in the normalized ``[0, 1]`` domain — callers should pass
    ``t / num_train_timesteps`` (or sample ``t`` directly from
    ``[0, 1]``).
    """
    s = self.shift if shift is None else float(shift)
    if s == 1.0:
        return t
    return s * t / (1.0 + (s - 1.0) * t)

def get_train_weight(
    self,
    t: torch.Tensor,
    *,
    weight_type: WeightType = "beta08",
) -> torch.Tensor:
    """Per-timestep training weight, renormalized so the total weight
    mass equals ``num_train_timesteps`` (matching AnyFlow reference's
    ``scheduling_flowmap_euler_discrete.py``).

    ``beta08``: ``w(t) = t * sqrt(1 - t)`` (in normalized t-space).
    """
    # Auto-detect domain: if t was given in absolute units, normalize.
    t_f = t.float()
    max_val = t_f.max() if t_f.numel() > 0 else torch.tensor(0.0)
    if max_val > 1.0 + 1e-6:
        t_norm = t_f / self.num_train_timesteps
    else:
        t_norm = t_f
    t_norm = t_norm.clamp(min=0.0, max=1.0)

    if weight_type == "uniform":
        w = torch.ones_like(t_norm)
    elif weight_type == "gaussian":
        w = torch.exp(-0.5 * ((t_norm - 0.5) / 0.2)**2)
    elif weight_type == "beta08":
        w = t_norm.pow(1.0) * (1.0 - t_norm).clamp_min(0.0).pow(0.5)
    else:
        raise ValueError(f"Unknown weight_type: {weight_type!r}")

    denom = w.sum().clamp_min(1e-8)
    return w * (float(self.num_train_timesteps) / denom)

# ------------------------------------------------------------------

def set_timesteps(
    self,
    *,
    num_inference_steps: int,
    device: torch.device | str = "cpu",
    custom_timesteps: list[float] | torch.Tensor | None = None,
) -> None:
    """Build a descending timestep schedule ending at 0.

    With ``num_inference_steps=N`` the schedule has ``N + 1`` entries
    ``[T_max, ..., 0]`` so a rollout consumes ``N`` Euler steps.

    ``custom_timesteps`` overrides the linspace+shift schedule with
    a pinned list (in absolute train-timestep units), useful for the
    AnyFlow paper's hand-tuned ``[999, 937, 833, 624, 0]`` schedule.
    """
    if num_inference_steps <= 0:
        raise ValueError("num_inference_steps must be positive, "
                         f"got {num_inference_steps}")
    device = torch.device(device)

    if custom_timesteps is not None:
        ts = torch.as_tensor(custom_timesteps, dtype=torch.float32, device=device)
        if ts.ndim != 1:
            raise ValueError("custom_timesteps must be 1-D, got shape "
                             f"{tuple(ts.shape)}")
        if not torch.all(ts[:-1] >= ts[1:]):
            raise ValueError("custom_timesteps must be descending (largest first)")
    else:
        ts_norm = torch.linspace(1.0, 0.0, num_inference_steps + 1, device=device)
        ts_norm = self.apply_shift(ts_norm)
        ts = ts_norm * self.num_train_timesteps

    self.timesteps = ts
    self.sigmas = ts / self.num_train_timesteps

def step(
    self,
    model_output: torch.Tensor,
    *,
    sample: torch.Tensor,
    timestep: torch.Tensor,
    r_timestep: torch.Tensor,
) -> torch.Tensor:
    """One Euler step from ``t`` to ``r``.

    ``model_output`` is the average-velocity prediction
    ``u_θ(x_t, t, r)``. Both ``timestep`` and ``r_timestep`` are in
    absolute train-timestep units (``[0, num_train_timesteps]``).
    """
    t = timestep.to(sample.device, dtype=sample.dtype)
    r = r_timestep.to(sample.device, dtype=sample.dtype)
    dt_norm = (t - r) / float(self.num_train_timesteps)
    # Broadcast dt over channel/spatial dims.
    view: list[int] = [-1] + [1] * (sample.ndim - 1)
    return sample - dt_norm.view(*view) * model_output

def add_noise(
    self,
    original_samples: torch.Tensor,
    noise: torch.Tensor,
    timestep: torch.Tensor,
) -> torch.Tensor:
    """Linear flow-matching interpolation: ``x_t = (1 - σ) * x_0 + σ * ε``,
    where ``σ = t / num_train_timesteps``.
    """
    sigma = (timestep.to(original_samples.device, dtype=original_samples.dtype) / float(self.num_train_timesteps))
    view: list[int] = [-1] + [1] * (original_samples.ndim - 1)
    sigma = sigma.view(*view)
    return (1.0 - sigma) * original_samples + sigma * noise</code></pre></details>
fastvideo.models.schedulers.scheduling_flow_map_euler_discrete.FlowMapEulerDiscreteScheduler.set_shift
set_shift(shift: float) -> None
Source
def set_shift(self, shift: float) -> None:
    self.shift = float(shift)
fastvideo.models.schedulers.scheduling_flow_map_euler_discrete.FlowMapEulerDiscreteScheduler.scale_model_input
scale_model_input(sample: torch.Tensor, timestep: int | None = None) -> torch.Tensor
Source
def scale_model_input(
    self,
    sample: torch.Tensor,
    timestep: int | None = None,
) -> torch.Tensor:
    # Flow-matching has no per-step input scaling; pass through.
    del timestep
    return sample
fastvideo.models.schedulers.scheduling_flow_map_euler_discrete.FlowMapEulerDiscreteScheduler.apply_shift
apply_shift(t: torch.Tensor, *, shift: float | None = None) -> torch.Tensor
Apply the flow-matching shift: t' = s * t / (1 + (s - 1) * t).
Operates in the normalized [0, 1] domain — callers should pass
t / num_train_timesteps (or sample t directly from
[0, 1]).
Source
def apply_shift(
    self,
    t: torch.Tensor,
    *,
    shift: float | None = None,
) -> torch.Tensor:
    """Apply the flow-matching shift: ``t' = s * t / (1 + (s - 1) * t)``.
Operates in the normalized ``&#91;0, 1&#93;`` domain — callers should pass
``t / num_train_timesteps`` (or sample ``t`` directly from
``&#91;0, 1&#93;``).
&quot;&quot;&quot;
s = self.shift if shift is None else float(shift)
if s == 1.0:
    return t
return s * t / (1.0 + (s - 1.0) * t)</code></pre></details>
fastvideo.models.schedulers.scheduling_flow_map_euler_discrete.FlowMapEulerDiscreteScheduler.get_train_weight
get_train_weight(t: torch.Tensor, *, weight_type: WeightType = 'beta08') -> torch.Tensor
Per-timestep training weight, renormalized so the total weight
mass equals num_train_timesteps (matching AnyFlow reference's
scheduling_flowmap_euler_discrete.py).
beta08: w(t) = t * sqrt(1 - t) (in normalized t-space).
Source
def get_train_weight(
    self,
    t: torch.Tensor,
    *,
    weight_type: WeightType = "beta08",
) -> torch.Tensor:
    """Per-timestep training weight, renormalized so the total weight
    mass equals ``num_train_timesteps`` (matching AnyFlow reference's
    ``scheduling_flowmap_euler_discrete.py``).
``beta08``: ``w(t) = t * sqrt(1 - t)`` (in normalized t-space).
&quot;&quot;&quot;
# Auto-detect domain: if t was given in absolute units, normalize.
t_f = t.float()
max_val = t_f.max() if t_f.numel() &gt; 0 else torch.tensor(0.0)
if max_val &gt; 1.0 + 1e-6:
    t_norm = t_f / self.num_train_timesteps
else:
    t_norm = t_f
t_norm = t_norm.clamp(min=0.0, max=1.0)

if weight_type == &quot;uniform&quot;:
    w = torch.ones_like(t_norm)
elif weight_type == &quot;gaussian&quot;:
    w = torch.exp(-0.5 * ((t_norm - 0.5) / 0.2)**2)
elif weight_type == &quot;beta08&quot;:
    w = t_norm.pow(1.0) * (1.0 - t_norm).clamp_min(0.0).pow(0.5)
else:
    raise ValueError(f&quot;Unknown weight_type: &#123;weight_type!r&#125;&quot;)

denom = w.sum().clamp_min(1e-8)
return w * (float(self.num_train_timesteps) / denom)</code></pre></details>
fastvideo.models.schedulers.scheduling_flow_map_euler_discrete.FlowMapEulerDiscreteScheduler.set_timesteps
set_timesteps(*, num_inference_steps: int, device: torch.device | str = 'cpu', custom_timesteps: list[float] | torch.Tensor | None = None) -> None
Build a descending timestep schedule ending at 0.
With num_inference_steps=N the schedule has N + 1 entries
[T_max, ..., 0] so a rollout consumes N Euler steps.
custom_timesteps overrides the linspace+shift schedule with
a pinned list (in absolute train-timestep units), useful for the
AnyFlow paper's hand-tuned [999, 937, 833, 624, 0] schedule.
Source
def set_timesteps(
    self,
    *,
    num_inference_steps: int,
    device: torch.device | str = "cpu",
    custom_timesteps: list[float] | torch.Tensor | None = None,
) -> None:
    """Build a descending timestep schedule ending at 0.
With ``num_inference_steps=N`` the schedule has ``N + 1`` entries
``&#91;T_max, ..., 0&#93;`` so a rollout consumes ``N`` Euler steps.

``custom_timesteps`` overrides the linspace+shift schedule with
a pinned list (in absolute train-timestep units), useful for the
AnyFlow paper&#x27;s hand-tuned ``&#91;999, 937, 833, 624, 0&#93;`` schedule.
&quot;&quot;&quot;
if num_inference_steps &lt;= 0:
    raise ValueError(&quot;num_inference_steps must be positive, &quot;
                     f&quot;got &#123;num_inference_steps&#125;&quot;)
device = torch.device(device)

if custom_timesteps is not None:
    ts = torch.as_tensor(custom_timesteps, dtype=torch.float32, device=device)
    if ts.ndim != 1:
        raise ValueError(&quot;custom_timesteps must be 1-D, got shape &quot;
                         f&quot;&#123;tuple(ts.shape)&#125;&quot;)
    if not torch.all(ts&#91;:-1&#93; &gt;= ts&#91;1:&#93;):
        raise ValueError(&quot;custom_timesteps must be descending (largest first)&quot;)
else:
    ts_norm = torch.linspace(1.0, 0.0, num_inference_steps + 1, device=device)
    ts_norm = self.apply_shift(ts_norm)
    ts = ts_norm * self.num_train_timesteps

self.timesteps = ts
self.sigmas = ts / self.num_train_timesteps</code></pre></details>
fastvideo.models.schedulers.scheduling_flow_map_euler_discrete.FlowMapEulerDiscreteScheduler.step
step(model_output: torch.Tensor, *, sample: torch.Tensor, timestep: torch.Tensor, r_timestep: torch.Tensor) -> torch.Tensor
One Euler step from t to r.
model_output is the average-velocity prediction
u_θ(x_t, t, r). Both timestep and r_timestep are in
absolute train-timestep units ([0, num_train_timesteps]).
Source
def step(
    self,
    model_output: torch.Tensor,
    *,
    sample: torch.Tensor,
    timestep: torch.Tensor,
    r_timestep: torch.Tensor,
) -> torch.Tensor:
    """One Euler step from ``t`` to ``r``.
``model_output`` is the average-velocity prediction
``u_θ(x_t, t, r)``. Both ``timestep`` and ``r_timestep`` are in
absolute train-timestep units (``&#91;0, num_train_timesteps&#93;``).
&quot;&quot;&quot;
t = timestep.to(sample.device, dtype=sample.dtype)
r = r_timestep.to(sample.device, dtype=sample.dtype)
dt_norm = (t - r) / float(self.num_train_timesteps)
# Broadcast dt over channel/spatial dims.
view: list&#91;int&#93; = &#91;-1&#93; + &#91;1&#93; * (sample.ndim - 1)
return sample - dt_norm.view(*view) * model_output</code></pre></details>
fastvideo.models.schedulers.scheduling_flow_map_euler_discrete.FlowMapEulerDiscreteScheduler.add_noise
add_noise(original_samples: torch.Tensor, noise: torch.Tensor, timestep: torch.Tensor) -> torch.Tensor
Linear flow-matching interpolation: x_t = (1 - σ) * x_0 + σ * ε,
where σ = t / num_train_timesteps.
Source
def add_noise(
    self,
    original_samples: torch.Tensor,
    noise: torch.Tensor,
    timestep: torch.Tensor,
) -> torch.Tensor:
    """Linear flow-matching interpolation: ``x_t = (1 - σ) * x_0 + σ * ε``,
    where ``σ = t / num_train_timesteps``.
    """
    sigma = (timestep.to(original_samples.device, dtype=original_samples.dtype) / float(self.num_train_timesteps))
    view: list[int] = [-1] + [1] * (original_samples.ndim - 1)
    sigma = sigma.view(*view)
    return (1.0 - sigma) * original_samples + sigma * noise