| """Flow-matching convention — the ONLY place it is defined (docs/04-architecture.md). |
| |
| * ``x0 ~ N(0, I)`` is noise, ``x1`` is the data latent, ``t ∈ [0, 1]``; |
| **t = 0 is noise, t = 1 is data**. |
| * ``x_t = (1 − t) · x0 + t · x1``; the target velocity is ``v = x1 − x0``. |
| * Loss: masked MSE ``‖v̂(x_t, t, c) − v‖²`` over the frames to generate (``loss_mask``). |
| * Training ``t``: logit-normal by default (``sigmoid(N(mean, std))``), or uniform. |
| * Sampling integrates ``dx/dt = v̂`` from ``t = 0`` to ``t = 1`` (Euler or midpoint), on a |
| time grid optionally warped by "sway" (``t + s · (cos(π t / 2) − 1 + t)``, F5-TTS). |
| |
| Shapes: latents are ``[batch, frames, channels]``; ``t`` is ``[batch]``; masks are |
| ``[batch, frames]`` booleans (True = counted). |
| """ |
|
|
| from __future__ import annotations |
|
|
| import math |
| from collections.abc import Callable |
| from typing import Literal |
|
|
| import torch |
| from torch import Tensor |
|
|
| __all__ = [ |
| "VelocityFn", |
| "interpolate", |
| "masked_mse", |
| "sample_ode", |
| "sample_t", |
| "target_velocity", |
| "time_grid", |
| ] |
|
|
| VelocityFn = Callable[[Tensor, Tensor], Tensor] |
|
|
|
|
| def _check_t(t: Tensor, batch: int) -> None: |
| if t.ndim != 1 or t.shape[0] != batch: |
| raise ValueError(f"t must have shape [{batch}], got {tuple(t.shape)}") |
|
|
|
|
| def interpolate(x0: Tensor, x1: Tensor, t: Tensor) -> Tensor: |
| """``x_t = (1 − t) · x0 + t · x1`` with ``t`` broadcast over frames and channels.""" |
| if x0.shape != x1.shape or x0.ndim != 3: |
| raise ValueError(f"x0 and x1 must be equal [B,T,C], got {x0.shape} and {x1.shape}") |
| _check_t(t, x0.shape[0]) |
| tt = t.to(x0.dtype)[:, None, None] |
| return (1.0 - tt) * x0 + tt * x1 |
|
|
|
|
| def target_velocity(x0: Tensor, x1: Tensor) -> Tensor: |
| """``v = x1 − x0`` (constant along the straight path).""" |
| if x0.shape != x1.shape: |
| raise ValueError(f"shape mismatch {x0.shape} vs {x1.shape}") |
| return x1 - x0 |
|
|
|
|
| def masked_mse(pred: Tensor, target: Tensor, loss_mask: Tensor) -> Tensor: |
| """Mean squared error over the masked frames (and all channels). |
| |
| Raises: |
| ValueError: shape mismatch, or a mask that selects no frame at all (the loss would |
| be 0/0 — a silent NaN or a silent zero are both wrong). |
| """ |
| if pred.shape != target.shape or pred.ndim != 3: |
| raise ValueError(f"pred/target must be equal [B,T,C], got {pred.shape}, {target.shape}") |
| if loss_mask.shape != pred.shape[:2] or loss_mask.dtype != torch.bool: |
| raise ValueError(f"loss_mask must be bool [B,T], got {loss_mask.dtype} {loss_mask.shape}") |
| count = loss_mask.sum() |
| if int(count) == 0: |
| raise ValueError("loss_mask selects no frame") |
| |
| |
| diff = torch.where( |
| loss_mask[..., None], pred.float() - target.float(), torch.zeros((), device=pred.device) |
| ) |
| return diff.pow(2).mean(dim=-1).sum() / count |
|
|
|
|
| def sample_t( |
| batch: int, |
| *, |
| kind: Literal["logit_normal", "uniform"], |
| mean: float = 0.0, |
| std: float = 1.0, |
| generator: torch.Generator | None = None, |
| device: torch.device | None = None, |
| ) -> Tensor: |
| """Training time steps in ``(0, 1)``, shape ``[batch]``.""" |
| if batch < 1: |
| raise ValueError("batch must be >= 1") |
| if kind == "logit_normal": |
| if not (math.isfinite(std) and std > 0 and math.isfinite(mean)): |
| raise ValueError(f"invalid logit-normal parameters mean={mean} std={std}") |
| normal = torch.randn(batch, generator=generator, device=device) * std + mean |
| return torch.sigmoid(normal) |
| if kind == "uniform": |
| return torch.rand(batch, generator=generator, device=device) |
| raise ValueError(f"unknown t sampling {kind!r}") |
|
|
|
|
| def time_grid(steps: int, *, sway: float = 0.0, device: torch.device | None = None) -> Tensor: |
| """``steps + 1`` increasing times from 0 to 1, optionally sway-warped. |
| |
| Sway (F5-TTS): ``t' = t + s · (cos(π t / 2) − 1 + t)``; ``s < 0`` puts more steps near |
| ``t = 0`` (noise end). Endpoints stay exactly 0 and 1; the grid must stay increasing. |
| |
| Raises: |
| ValueError: ``steps < 1`` or a sway that makes the grid non-monotonic. |
| """ |
| if steps < 1: |
| raise ValueError("steps must be >= 1") |
| t = torch.linspace(0.0, 1.0, steps + 1, dtype=torch.float64, device=device) |
| if sway != 0.0: |
| t = t + sway * (torch.cos(math.pi / 2 * t) - 1 + t) |
| t[0], t[-1] = 0.0, 1.0 |
| if not bool(torch.all(t[1:] > t[:-1])): |
| raise ValueError(f"sway {sway} gives a non-increasing time grid") |
| return t |
|
|
|
|
| @torch.no_grad() |
| def sample_ode( |
| velocity: VelocityFn, |
| x0: Tensor, |
| *, |
| steps: int, |
| method: Literal["euler", "midpoint"] = "euler", |
| sway: float = 0.0, |
| ) -> Tensor: |
| """Integrate from noise ``x0`` (t = 0) to data (t = 1). |
| |
| With the oracle ``v̂ = x1 − x0`` both methods return ``x1`` exactly (up to float error): |
| the path is a straight line (unit test of the convention). |
| """ |
| if x0.ndim != 3: |
| raise ValueError(f"x0 must be [B,T,C], got {tuple(x0.shape)}") |
| if method not in ("euler", "midpoint"): |
| raise ValueError(f"unknown ODE method {method!r}") |
| grid = time_grid(steps, sway=sway, device=x0.device) |
| x = x0 |
| batch = x0.shape[0] |
| for i in range(steps): |
| t0, t1 = float(grid[i]), float(grid[i + 1]) |
| dt = t1 - t0 |
| t_now = torch.full((batch,), t0, dtype=x0.dtype, device=x0.device) |
| v = velocity(x, t_now) |
| if method == "euler": |
| x = x + dt * v |
| elif method == "midpoint": |
| t_mid = torch.full((batch,), t0 + dt / 2, dtype=x0.dtype, device=x0.device) |
| x = x + dt * velocity(x + (dt / 2) * v, t_mid) |
| else: |
| raise ValueError(f"unknown ODE method {method!r}") |
| return x |
|
|
