src/flow/paths.py

"""
Flow matching and score matching algorithms for generative modeling

This module implements three training algorithms for ablation:
1. ScoreMatchingDiffusion - Baseline (predict score, VP-SDE diffusion path)
2. FlowMatchingDiffusion - Flow matching with diffusion path (predict velocity)
3. FlowMatchingOT - Flow matching with optimal transport path (predict velocity)

Each algorithm encapsulates both the probability path and loss computation.
"""

import math

import torch
import torch.nn.functional as F
from abc import ABC, abstractmethod

# SD3-style resolution-dependent time shift.
# SD3 (Esser et al., 2024, Eq. 23) derives the shift from matching the
# uncertainty about the clean image across resolutions:
#     alpha = sqrt(m / n),
# where n and m are the pixel counts of the reference and target resolutions.
# With a fixed VAE and patch size, token count is proportional to pixel count,
# so the same ratio applies to token counts. Reference: 256 tokens (256×256
# image → 32×32 latent, patch_size=2 → 16×16 patches) gets shift 1.0.
# SD3 empirically found little quality difference among shift values above
# 1.5 and used 3.0 at 1024×1024, where the formula gives 4.0 — so the formula
# values used here sit inside the range they found acceptable.
_SHIFT_BASE_TOKENS = 256      # 256×256 → 32×32 latent → 16×16 patches


def resolution_time_shift(z: torch.Tensor, patch_size: int = 2) -> float:
    """Compute SD3-style time shift from a noise/latent tensor [B, C, H, W]."""
    _, _, h, w = z.shape
    n_tokens = (h * w) / (patch_size ** 2)
    if n_tokens <= _SHIFT_BASE_TOKENS:
        return 1.0
    return math.sqrt(n_tokens / _SHIFT_BASE_TOKENS)


def apply_time_shift(t: torch.Tensor, shift: float) -> torch.Tensor:
    """Apply the SD3 shift transform in this repo's timestep convention.

    SD3 (Esser et al., 2024, Eq. 23) uses t=0 for data and t=1 for noise and
    pushes timesteps toward the noisy end at higher resolutions:
        u' = (s * u) / (1 + (s - 1) * u).
    This repo uses the opposite convention (t=0 noise, t=1 data), so the same
    transform applied to u = 1 - t yields:
        t' = t / (s - (s - 1) * t),
    which decreases t (moves toward noise) when shift > 1.
    """
    return t / (shift - (shift - 1) * t)


def shift_timesteps(
    t: torch.Tensor,
    z_like: torch.Tensor,
    *,
    patch_size: int = 2,
    time_shift: float | None = None,
) -> torch.Tensor:
    """Shift timesteps using the repo's resolution-dependent convention.

    This is the shared entry point for both training and inference to ensure
    the model always sees the same "t_used" that is implied by the resolution.

    Args:
        t: Base timesteps in [0, 1], shape [B] or scalar tensor.
        z_like: A latent/noise tensor with shape [B, C, H, W] used to infer
            the resolution-dependent shift when time_shift is None.
        patch_size: Patch size used to convert latent HxW into token count.
        time_shift: Optional explicit shift scalar to override auto-compute.
    """
    if time_shift is None:
        time_shift = resolution_time_shift(z_like, patch_size=patch_size)
    return apply_time_shift(t, time_shift)


def sample_weighted_mse(
    error: torch.Tensor, sample_weights: torch.Tensor
) -> torch.Tensor:
    """Weighted mean of the per-sample mean squared errors of ``error``.

    The samples are combined as ``sum(w_i * L_i) / sum(w_i)``: a weighted mean,
    not a weighted sum.  A sum would scale a step's gradient with the weights
    themselves, which is a learning-rate change wearing a reweighting's
    clothes, and two runs would no longer be comparable at equal step counts.

    The reduction runs in float32 whatever the compute dtype is, because the
    weight sum is a normalization: a denominator rounded to the compute dtype
    would bias every loss the weights touch.
    """
    if sample_weights.shape != (error.shape[0],):
        raise ValueError(
            "sample weights must hold one multiplier per sample, got "
            f"{tuple(sample_weights.shape)} for a batch of {error.shape[0]}"
        )
    weights = sample_weights.float()
    per_sample = (error.float() ** 2).flatten(1).mean(dim=1)
    return (per_sample * weights).sum() / weights.sum()


class BaseAlgorithm(ABC):
    @abstractmethod
    def sample_zt(
        self, z0: torch.Tensor, z1: torch.Tensor, t: torch.Tensor
    ) -> torch.Tensor:
        """
        Sample z_t from the probability path p_t(z_t | z0, z1).
        Args:
            z0: Source samples (noise), [B, C, H, W]
            z1: Target samples (data), [B, C, H, W]
            t: Timesteps, [B] or [B, 1, 1, 1]
        Returns:
            Interpolated samples z_t, [B, C, H, W]
        """
        pass

    @abstractmethod
    def compute_loss(
        self,
        model_output: torch.Tensor,
        z0: torch.Tensor,
        z1: torch.Tensor,
        t: torch.Tensor,
        sample_weights: torch.Tensor | None = None,
    ) -> torch.Tensor:
        """
        Compute the loss for the algorithm.
        Args:
            model_output: Model predictions, [B, C, H, W]
            z0: Source samples (noise), [B, C, H, W]
            z1: Target samples (data), [B, C, H, W]
            t: Timesteps, [B] or [B, 1, 1, 1]
            sample_weights: Optional per-sample multipliers, one per row of the
                batch. With them the returned scalar is the weighted mean of
                the per-sample losses instead of their plain mean, which is how
                a run raises the share of its gradient that long captions
                receive; without them the computation is unchanged.
        Returns:
            Scalar loss value
        """
        pass


class FlowMatchingOT(BaseAlgorithm):
    """
    Flow matching with Optimal Transport path.
    """

    def sample_zt(
        self, z0: torch.Tensor, z1: torch.Tensor, t: torch.Tensor
    ) -> torch.Tensor:
        if t.dim() == 1:
            t = t.view(-1, 1, 1, 1)

        z_t = (1.0 - t) * z0 + t * z1
        return z_t

    def compute_loss(
        self,
        model_output: torch.Tensor,
        z0: torch.Tensor,
        z1: torch.Tensor,
        t: torch.Tensor,
        sample_weights: torch.Tensor | None = None,
    ) -> torch.Tensor:
        velocity_target = z1 - z0
        if sample_weights is not None:
            return sample_weighted_mse(model_output - velocity_target, sample_weights)
        return F.mse_loss(model_output, velocity_target)