Denoising Diffusion (DDPM)

"""DDPM, transcribed the way the paper reads.

Two processes:

  FORWARD (fixed, no learning):  gradually add noise until nothing is left
      q(x_t | x_0) = N(sqrt(alpha_bar_t) x_0, (1 - alpha_bar_t) I)

  REVERSE (learned):             remove a little noise, one step at a time
      loss = || eps - eps_theta(x_t, t) ||^2

The closed form for q(x_t | x_0) is the structural key: because the noise
accumulates as a product of Gaussians, ANY timestep can be sampled directly
from the clean data. Training therefore never simulates a trajectory - all the
sequential cost is deferred to sampling, which is the whole cost profile of
this model in one sentence.

Plain loops, no libraries.
"""

import math
import random

SEED = 29


def linear_beta_schedule(steps, start=1e-4, end=0.02):
    """The original schedule. Included to be compared against, not used.

    A linear schedule destroys information too quickly at the end, so most of
    the trajectory is spent at noise levels where the input is already
    indistinguishable from noise - capacity spent on steps that carry nothing.
    """
    return [start + (end - start) * t / (steps - 1) for t in range(steps)]


def cosine_beta_schedule(steps, offset=0.008):
    """The schedule that actually works, and it is not a minor difference.

    Defined through the CUMULATIVE signal retention rather than the per-step
    noise, so information is destroyed at a roughly constant rate in
    perceptual terms. That redistributes the trajectory toward the noise
    levels where structure lives.
    """
    def f(t):
        angle = ((t / steps) + offset) / (1.0 + offset) * math.pi / 2.0
        return math.cos(angle) ** 2

    alpha_bars = [f(t) / f(0) for t in range(steps + 1)]
    betas = []
    for t in range(steps):
        beta = 1.0 - alpha_bars[t + 1] / alpha_bars[t]
        # Clipped: an unclipped cosine schedule produces betas near one at the
        # end, which makes the reverse step numerically unstable.
        betas.append(min(beta, 0.999))
    return betas


def build_schedule(betas):
    """Precompute every coefficient the forward and reverse steps need.

    All of these are FIXED - nothing here is learned. That is worth stating
    plainly, because it is the difference from a VAE: the encoder is not a
    network, it is a schedule chosen in advance.
    """
    alphas = [1.0 - beta for beta in betas]

    alpha_bars = []
    running = 1.0
    for alpha in alphas:
        running *= alpha
        alpha_bars.append(running)

    return {
        'betas': betas,
        'alphas': alphas,
        'alpha_bars': alpha_bars,
        'sqrt_alpha_bars': [math.sqrt(a) for a in alpha_bars],
        'sqrt_one_minus_alpha_bars': [math.sqrt(1.0 - a) for a in alpha_bars],
    }


def forward_diffuse(x0, t, schedule, rng):
    """Jump straight to timestep t. No intervening steps are simulated.

    x_t = sqrt(alpha_bar_t) * x_0 + sqrt(1 - alpha_bar_t) * eps

    The noise eps is returned because it IS the training target - the loop
    added it, so the label is known exactly. That is what makes diffusion
    training so much more stable than anything adversarial: the supervision is
    a known quantity rather than the output of another network that is also
    learning.
    """
    signal = schedule['sqrt_alpha_bars'][t]
    noise_scale = schedule['sqrt_one_minus_alpha_bars'][t]

    epsilon = [rng.gauss(0.0, 1.0) for _ in x0]
    x_t = [signal * value + noise_scale * e for value, e in zip(x0, epsilon)]
    return x_t, epsilon


def denoiser(x_t, t, params, num_steps):
    """The network: given a noisy input and a timestep, predict the noise.

    The timestep enters as a sinusoidal embedding rather than a raw number,
    because a single scalar spanning a thousand values is a poor input to a
    network - the embedding gives it a smooth multi-scale representation.
    """
    embedding = timestep_embedding(t, len(params['w_time']), num_steps)

    hidden = []
    for i, row in enumerate(params['w1']):
        total = params['b1'][i] + params['w_time'][i] * embedding[i % len(embedding)]
        for weight, value in zip(row, x_t):
            total += weight * value
        hidden.append(math.tanh(total))

    return [
        bias + sum(w * h for w, h in zip(row, hidden))
        for row, bias in zip(params['w2'], params['b2'])
    ]


def timestep_embedding(t, width, num_steps):
    """Sinusoidal position encoding over the timestep.

    Same construction as a transformer's positional encoding, and for the same
    reason: a smooth multi-scale representation of an integer position.
    """
    half = width // 2
    return [
        math.sin(t / (10000.0 ** (2 * i / width))) if i < half
        else math.cos(t / (10000.0 ** (2 * (i - half) / width)))
        for i in range(width)
    ]


def training_loss(x0, params, schedule, num_steps, rng):
    """One training example. Note there is no rollout anywhere.

    A random timestep, one noise draw, one network evaluation. Training is
    embarrassingly parallel across examples AND across timesteps, which is
    entirely due to the closed form above.
    """
    t = rng.randrange(num_steps)
    x_t, epsilon = forward_diffuse(x0, t, schedule, rng)
    predicted = denoiser(x_t, t, params, num_steps)

    # A plain mean squared error on noise. The DERIVATION is a variational
    # bound of exactly the VAE's form, with a weighting per timestep - and
    # that weighting is DROPPED here. The simplified objective samples much
    # better and is no longer a bound on the likelihood.
    return sum((e - p) ** 2 for e, p in zip(epsilon, predicted)) / len(epsilon)


def ddpm_sample_step(x_t, t, params, schedule, num_steps, rng, stochastic=True):
    """One reverse step. This is where all the cost lives.

    x_{t-1} = (x_t - (1-alpha_t)/sqrt(1-alpha_bar_t) * eps_theta) / sqrt(alpha_t)
              + sigma_t * z
    """
    predicted = denoiser(x_t, t, params, num_steps)

    alpha = schedule['alphas'][t]
    alpha_bar = schedule['alpha_bars'][t]
    coefficient = (1.0 - alpha) / math.sqrt(1.0 - alpha_bar)

    mean = [
        (value - coefficient * p) / math.sqrt(alpha)
        for value, p in zip(x_t, predicted)
    ]

    if t == 0 or not stochastic:
        # Setting the injected noise to zero makes the sampler DETERMINISTIC,
        # which is DDIM - and a deterministic trajectory can be traversed in
        # far fewer steps, because its intermediate points can be skipped.
        return mean

    sigma = math.sqrt(schedule['betas'][t])
    return [m + sigma * rng.gauss(0.0, 1.0) for m in mean]


def sample(dimension, params, schedule, num_steps, rng, stochastic=True):
    """Start from pure noise and walk back. S sequential network evaluations.

    This is the whole cost profile: training is one evaluation per example,
    sampling is S. At a thousand steps that is three orders of magnitude more
    expensive per sample than a GAN, and every practical advance in this area
    has been about reducing S.
    """
    x = [rng.gauss(0.0, 1.0) for _ in range(dimension)]
    for t in range(num_steps - 1, -1, -1):
        x = ddpm_sample_step(x, t, params, schedule, num_steps, rng, stochastic)
    return x


def strided_sample(dimension, params, schedule, num_steps, stride, rng):
    """Skip steps. Only valid for the DETERMINISTIC sampler.

    Once the reverse process has no injected noise, its intermediate states
    are points on a trajectory - so skipping them is exactly what a solver
    with a larger step size does, and the whole numerical-integration
    literature becomes applicable.
    """
    x = [rng.gauss(0.0, 1.0) for _ in range(dimension)]
    for t in range(num_steps - 1, -1, -stride):
        x = ddpm_sample_step(x, t, params, schedule, num_steps, rng, stochastic=False)
    return x

"""The same model, with the schedule as a type and the sampler as a choice.

Two things change. The noise schedule becomes an object that validates its own
invariants and precomputes every derived coefficient once, because it is a
FIXED part of the model — baked into the checkpoint, and a sampler using a
different one produces degraded output with no error anywhere. And the sampler
becomes an explicit type, because step count and determinism are serving
parameters that change the output without changing the model, which makes them
exactly the sort of thing that should not be a loose argument.
"""

from __future__ import annotations

import math
import random
from dataclasses import dataclass
from enum import Enum
from typing import NamedTuple, Sequence


class ScheduleKind(Enum):
    """Which noise schedule. Not interchangeable at serving time.

    LINEAR destroys information too quickly at the end, so most of the
    trajectory sits at noise levels where the input carries nothing. COSINE
    redistributes it toward the levels where structure lives, and the
    difference is large rather than marginal.
    """

    LINEAR = 'linear'
    COSINE = 'cosine'


class SamplerKind(Enum):
    """Stochastic or deterministic, which is a bigger choice than it looks.

    DDPM injects noise at every step; DDIM does not. A deterministic
    trajectory can be traversed in far fewer steps because its intermediate
    points are skippable, and it makes the same input give the same output.
    """

    DDPM = 'ddpm'
    DDIM = 'ddim'


class ScheduleMismatch(ValueError):
    """Raised when a sampler's schedule differs from the trained one.

    Its own type because the schedule is part of the MODEL, not the sampler:
    the network learned to denoise at specific noise levels, and sampling with
    different ones produces degraded output with nothing to indicate it.
    """


class DegenerateSchedule(ValueError):
    """Raised when a schedule is not monotone or leaves signal at the end.

    Cumulative signal retention must decrease to near zero, or the reverse
    process starts from something that is not the prior it assumes.
    """


class ShapeMismatch(ValueError):
    """Raised on a dimension violation instead of computing nonsense."""


class NoisedSample(NamedTuple):
    """The noisy input AND the noise that was added.

    Both come back because the noise IS the training target — the loop added
    it, so the label is known exactly. That is what makes diffusion training
    stable where adversarial training is not: the supervision is a known
    quantity rather than the output of a network that is also learning.
    """

    x_t: list[float]
    epsilon: list[float]
    timestep: int


class NoiseSchedule:
    """Fixed, not learned, and part of the model artifact.

    Worth stating plainly because it is the structural difference from a VAE:
    the encoder here is not a network, it is a schedule chosen in advance. All
    derived coefficients are precomputed once, since they are constants of the
    schedule and recomputing them per step is pure waste.
    """

    def __init__(self, kind: ScheduleKind, steps: int, cosine_offset: float = 0.008) -> None:
        if steps < 2:
            raise ShapeMismatch('a diffusion process needs at least two steps')

        self._kind = kind
        self._steps = steps
        self._betas = (
            self._linear(steps) if kind is ScheduleKind.LINEAR
            else self._cosine(steps, cosine_offset)
        )

        alphas = [1.0 - beta for beta in self._betas]
        self._alpha_bars: list[float] = []
        running = 1.0
        for alpha in alphas:
            running *= alpha
            self._alpha_bars.append(running)

        # Guard clause: the reverse process starts from a standard Gaussian,
        # so the forward process must actually reach one.
        if self._alpha_bars[-1] > 0.01:
            raise DegenerateSchedule(
                f'cumulative signal retention ends at {self._alpha_bars[-1]:.3f}; the '
                'forward process does not reach the prior the sampler assumes'
            )
        if any(a >= b for a, b in zip(self._alpha_bars, self._alpha_bars[1:])):
            raise DegenerateSchedule('signal retention must decrease monotonically')

        self._alphas = alphas
        self._sqrt_alpha_bars = [math.sqrt(a) for a in self._alpha_bars]
        self._sqrt_one_minus = [math.sqrt(1.0 - a) for a in self._alpha_bars]

    @staticmethod
    def _linear(steps: int, start: float = 1e-4, end: float = 0.02) -> list[float]:
        return [start + (end - start) * t / (steps - 1) for t in range(steps)]

    @staticmethod
    def _cosine(steps: int, offset: float) -> list[float]:
        """Defined through CUMULATIVE retention rather than per-step noise.

        That is the whole point: information is destroyed at a roughly
        constant perceptual rate rather than collapsing at the end.
        """
        def f(t: int) -> float:
            angle = ((t / steps) + offset) / (1.0 + offset) * math.pi / 2.0
            return math.cos(angle) ** 2

        bars = [f(t) / f(0) for t in range(steps + 1)]
        # Clipped: an unclipped cosine schedule produces betas near one at the
        # end, which makes the reverse step numerically unstable.
        return [min(1.0 - bars[t + 1] / bars[t], 0.999) for t in range(steps)]

    @property
    def steps(self) -> int:
        return self._steps

    def fingerprint(self) -> tuple[str, int]:
        """Comparable across training and serving, so a mismatch is catchable."""
        return (self._kind.value, self._steps)

    def assert_matches(self, trained: NoiseSchedule) -> None:
        """Guard clause for a mismatch with no other symptom."""
        if self.fingerprint() != trained.fingerprint():
            raise ScheduleMismatch(
                f'sampling schedule {self.fingerprint()} differs from the trained '
                f'{trained.fingerprint()}; the network learned different noise levels'
            )

    def diffuse(self, x0: Sequence[float], t: int, rng: random.Random) -> NoisedSample:
        """Jump straight to timestep t. No intervening steps are simulated.

        The closed form is the structural key: because the noise accumulates
        as a product of Gaussians, any timestep is directly sampleable from
        the clean data. Training therefore never simulates a trajectory, and
        all the sequential cost is deferred to sampling.
        """
        if not 0 <= t < self._steps:
            raise ShapeMismatch(f'timestep {t} outside [0, {self._steps})')

        epsilon = [rng.gauss(0.0, 1.0) for _ in x0]
        x_t = [
            self._sqrt_alpha_bars[t] * value + self._sqrt_one_minus[t] * e
            for value, e in zip(x0, epsilon)
        ]
        return NoisedSample(x_t=x_t, epsilon=epsilon, timestep=t)

    def reverse_mean(
        self, x_t: Sequence[float], predicted_noise: Sequence[float], t: int
    ) -> list[float]:
        coefficient = (1.0 - self._alphas[t]) / self._sqrt_one_minus[t]
        root_alpha = math.sqrt(self._alphas[t])
        return [
            (value - coefficient * p) / root_alpha
            for value, p in zip(x_t, predicted_noise)
        ]

    def sampling_sigma(self, t: int) -> float:
        return math.sqrt(self._betas[t])


@dataclass(frozen=True)
class SamplerConfig:
    """Serving parameters that change the output without changing the model.

    Frozen and explicit because that is exactly what makes them dangerous: two
    deployments of one checkpoint at different step counts produce materially
    different quality, and comparing them compares samplers rather than models.
    """

    kind: SamplerKind = SamplerKind.DDIM
    steps: int = 50
    guidance_scale: float = 1.0

    def __post_init__(self) -> None:
        if self.steps < 1:
            raise ShapeMismatch('at least one sampling step is required')
        # Guard clause for the combination that is simply invalid: skipping
        # steps assumes a deterministic trajectory whose intermediate points
        # can be omitted, which stochastic sampling does not have.
        if self.kind is SamplerKind.DDPM and self.steps < 100:
            raise ShapeMismatch(
                'stochastic sampling cannot skip steps; use DDIM for low step counts'
            )
        if self.guidance_scale < 1.0:
            raise ShapeMismatch('guidance below 1.0 pushes away from the condition')

    def timesteps(self, schedule: NoiseSchedule) -> list[int]:
        """Which timesteps to visit, evenly strided and descending."""
        stride = max(schedule.steps // self.steps, 1)
        return list(range(schedule.steps - 1, -1, -stride))


class SampleResult(NamedTuple):
    """The sample AND the settings that produced it.

    Bundled because quality and cost are on a dial here: a sample reported
    without its step count and guidance scale is not comparable to any other,
    and that is the most common way diffusion results become meaningless.
    """

    values: list[float]
    steps_taken: int
    sampler: SamplerKind
    guidance_scale: float


def guided_prediction(
    conditional: Sequence[float],
    unconditional: Sequence[float],
    scale: float,
) -> list[float]:
    """Classifier-free guidance: extrapolate away from the unconditional.

    A serving-time knob trading DIVERSITY for ADHERENCE. Raising it makes each
    sample match its condition better and makes the collection less varied,
    which is easy to over-tune because the failure looks like improvement.
    """
    return [
        u + scale * (c - u)
        for c, u in zip(conditional, unconditional)
    ]

"""Batched training, and a sampler that skips.

Two changes of completely different character, which is the point:

  TRAINING batches trivially. The closed-form forward process means every
  example can draw its own random timestep and be diffused directly, so a
  batch is one gather from a precomputed coefficient table and one network
  pass. There is no rollout to parallelize because there is no rollout.

  SAMPLING does not batch away. It is inherently sequential, and the only
  real optimization is to TAKE FEWER STEPS - which is available because the
  deterministic sampler follows a trajectory whose intermediate points can be
  skipped. That reframes sampling as solving an ODE, where step count is a
  solver choice rather than an approximation to the model.

The asymmetry between those two paragraphs is the whole cost profile.
"""

from __future__ import annotations

import numpy as np
from numpy.typing import NDArray

FLOAT = np.float32


class BatchedSchedule:
    """Every coefficient precomputed as an array, indexed by timestep.

    All of these are constants of the schedule - nothing is learned. Storing
    them as arrays means a batch with a different timestep per example is one
    fancy-index gather rather than a loop.
    """

    def __init__(self, betas: NDArray[np.float32]) -> None:
        self._betas = np.ascontiguousarray(betas, dtype=FLOAT)
        alphas = 1.0 - self._betas
        self._alphas = alphas
        # cumprod, not a Python loop: the cumulative product is the one place
        # the schedule has a sequential dependency, and NumPy already has it.
        self._alpha_bars = np.cumprod(alphas, dtype=FLOAT)
        self._sqrt_alpha_bars = np.sqrt(self._alpha_bars)
        self._sqrt_one_minus = np.sqrt(1.0 - self._alpha_bars)
        self._reverse_coefficient = (1.0 - alphas) / self._sqrt_one_minus
        self._root_alphas = np.sqrt(alphas)

    @property
    def steps(self) -> int:
        return len(self._betas)

    def diffuse(
        self,
        x0: NDArray[np.float32],
        timesteps: NDArray[np.int32],
        rng: np.random.Generator,
    ) -> tuple[NDArray[np.float32], NDArray[np.float32]]:
        """Every example at its OWN random timestep, in one operation.

        The gather is the whole trick: a different t per example would be a
        loop in the literal version, and here it is one indexed read from the
        coefficient table broadcast over the feature axis.
        """
        signal = self._sqrt_alpha_bars[timesteps][:, None]
        noise_scale = self._sqrt_one_minus[timesteps][:, None]

        epsilon = rng.standard_normal(x0.shape, dtype=FLOAT)
        # Fused: the noisy input is built in place from the scaled clean data
        # rather than materializing both scaled terms separately.
        x_t = x0 * signal
        x_t += noise_scale * epsilon
        return x_t, epsilon

    def reverse_mean(
        self,
        x_t: NDArray[np.float32],
        predicted: NDArray[np.float32],
        t: int,
    ) -> NDArray[np.float32]:
        """One reverse step, in place over the predicted-noise buffer."""
        out = predicted
        out *= -self._reverse_coefficient[t]
        out += x_t
        out /= self._root_alphas[t]
        return out

    def sigma(self, t: int) -> float:
        return float(np.sqrt(self._betas[t]))


def cosine_betas(steps: int, offset: float = 0.008) -> NDArray[np.float32]:
    """Defined through CUMULATIVE retention rather than per-step noise.

    Information is destroyed at a roughly constant perceptual rate rather than
    collapsing at the end, which is where a linear schedule wastes most of its
    trajectory.
    """
    t = np.arange(steps + 1, dtype=np.float64)
    angle = ((t / steps) + offset) / (1.0 + offset) * np.pi / 2.0
    bars = np.cos(angle) ** 2
    bars /= bars[0]
    # Clipped: an unclipped cosine schedule produces betas near one at the
    # end, which makes the reverse step numerically unstable.
    return np.clip(1.0 - bars[1:] / bars[:-1], 0.0, 0.999).astype(FLOAT)


def timestep_embedding(
    timesteps: NDArray[np.int32], width: int
) -> NDArray[np.float32]:
    """Sinusoidal encoding for a whole batch of timesteps at once.

    Same construction as a transformer's positional encoding, and for the same
    reason: a smooth multi-scale representation of an integer position. The
    frequencies are computed once and broadcast, not rebuilt per example.
    """
    half = width // 2
    frequencies = np.exp(
        -np.log(10000.0) * np.arange(half, dtype=FLOAT) / half
    )
    angles = timesteps[:, None].astype(FLOAT) * frequencies[None, :]
    return np.concatenate((np.sin(angles), np.cos(angles)), axis=1)


def training_step(
    x0: NDArray[np.float32],
    schedule: BatchedSchedule,
    denoiser,
    rng: np.random.Generator,
) -> np.float32:
    """One batched training step. Note there is no rollout anywhere.

    A random timestep per example, one noise draw, one network pass. Training
    is embarrassingly parallel across examples AND timesteps, entirely because
    the forward process has a closed form.
    """
    timesteps = rng.integers(0, schedule.steps, size=len(x0)).astype(np.int32)
    x_t, epsilon = schedule.diffuse(x0, timesteps, rng)

    predicted = denoiser(x_t, timestep_embedding(timesteps, 128))

    # A plain mean squared error on noise. The DERIVATION is a variational
    # bound with a per-timestep weighting, and that weighting is DROPPED here:
    # the simplified objective samples much better and is no longer a bound.
    residual = predicted
    residual -= epsilon
    return np.einsum('ij,ij->', residual, residual, optimize=True) / residual.size


def ddim_sample(
    shape: tuple[int, int],
    schedule: BatchedSchedule,
    denoiser,
    steps: int,
    rng: np.random.Generator,
) -> NDArray[np.float32]:
    """Deterministic sampling with STRIDED timesteps.

    This is the only real optimization available for sampling, and it works
    because the deterministic reverse process is a trajectory: its
    intermediate states are points on a path, so skipping them is exactly what
    a solver with a larger step size does. Twenty steps out of a thousand is a
    fifty-fold reduction, and the quality degrades gradually rather than
    falling off a cliff.

    Note what is NOT here: no batching over steps, because the dependency is
    real. Sampling is sequential and the only lever is step count.
    """
    stride = max(schedule.steps // steps, 1)
    visited = np.arange(schedule.steps - 1, -1, -stride, dtype=np.int32)

    x = rng.standard_normal(shape, dtype=FLOAT)
    # The batch dimension IS parallel even though the steps are not: many
    # samples advance together through the same sequential trajectory.
    embeddings = timestep_embedding(visited, 128)

    for index, t in enumerate(visited):
        predicted = denoiser(x, np.broadcast_to(embeddings[index], (shape[0], 128)))
        x = schedule.reverse_mean(x, predicted, int(t))

    return x


def guided_prediction(
    conditional: NDArray[np.float32],
    unconditional: NDArray[np.float32],
    scale: float,
) -> NDArray[np.float32]:
    """Classifier-free guidance, computed in place on the conditional buffer.

    A serving-time knob trading DIVERSITY for ADHERENCE. Raising it makes each
    sample match its condition better and makes the collection less varied,
    which is easy to over-tune because the failure looks like improvement.

    Note the cost: this needs BOTH a conditional and an unconditional
    prediction, so guided sampling is twice the network evaluations of
    unguided sampling at the same step count.
    """
    out = conditional
    out -= unconditional
    out *= FLOAT(scale)
    out += unconditional
    return out