Blame for lecture_3/ddpm.py

"""DDPM, the reverse-time equation written on a finite grid.

A U-Net reads a corrupted image and the step it sits at, and regresses the Gaussian draw that
corrupted it. The reverse transition rebuilds the previous state from that one prediction, with a
mean that comes from completing the square and the posterior variance as its spread. The sampler
starts at a standard Gaussian and applies that transition K times. The three pieces match the
derivation one to one, and we keep the chapter's names. `forward_diffusion.py` covers the schedule
and the forward corruption, and we rebuild both here so that this file stands alone.

In Chapter 3 step k = 0 holds a clean image and corruption grows with k, up to pure noise at
k = K.

Chapter 3, Sections 3.19, 3.20 and 3.24 through 3.28 of the notes. Trained on 28-by-28 images, the
offline toy set by default and MNIST with --dataset mnist. Prints L_simple over the run, the
reverse posterior variance and the ELBO weight formed from it, the chapter's two-step reverse
transition, and the range of the generated images.

    python lecture_3/ddpm.py                    # 300 steps on the offline images
    python lecture_3/ddpm.py --steps 20         # a few seconds, for a first look
    python lecture_3/ddpm.py --dataset mnist    # the real digits, downloaded on first use
"""

# %% 01. Imports and configuration
import math

import torch
import torch.nn as nn
import torch.nn.functional as F

from cis6270.nets import count_parameters
from cis6270.runner import common_parser, pick_device, print_report, save_run, set_seed, train

NUM_STEPS = 1000  # Diffusion steps K.
BETA_START = 1e-4  # Noise variance added at step 1.
BETA_END = 0.02  # Noise variance added at step K.


# %% 02. The schedule tensors and the forward corruption
#
# alpha_k = 1 - beta_k is the one-step retention, abar_k its running product, and beta-tilde_k the
# variance of the reverse posterior q(x_{k-1} | x_k, x_0). We prefix a zero to the betas, so
# Python index k holds beta_k and the cumulative product starts at abar_0 = 1, the clean state.
# One Gaussian draw then reaches any level, x_k = a_k x_0 + sigma_k eps with a_k = sqrt(abar_k) and
# sigma_k = sqrt(1 - abar_k).

betas = torch.cat([
    torch.zeros(1),  # Placeholder so that index k holds beta_k.
    torch.linspace(BETA_START, BETA_END, NUM_STEPS),  # beta_1 ... beta_K.
])
alphas = 1.0 - betas  # One-step signal retention alpha_k.
alphas_cumprod = torch.cumprod(alphas, dim=0)  # Cumulative retention.
alphas_cumprod_prev = torch.cat([
    torch.ones(1),  # abar_0 = 1.
    alphas_cumprod[:-1],  # Index k now holds abar_{k-1}.
])
posterior_variance = (
    betas * (1.0 - alphas_cumprod_prev) / (1.0 - alphas_cumprod)
)  # Reverse variance beta-tilde_k.


def extract(schedule, k, shape):
    # One schedule value per item, reshaped to broadcast over C, H, W.
    out = schedule.to(k.device).gather(0, k)
    return out.reshape(k.shape[0], *((1,) * (len(shape) - 1)))


def q_sample(x_0, k, noise=None):
    if noise is None:
        noise = torch.randn_like(x_0)  # Mean 0, covariance I, shape of x_0.
    abar = extract(alphas_cumprod, k, x_0.shape)
    a_k = abar.sqrt()  # Signal coefficient a_k.
    sigma_k = (1.0 - abar).sqrt()  # Noise standard deviation sigma_k.
    return a_k * x_0 + sigma_k * noise, noise  # x_k and the exact noise.


# %% 03. Sinusoidal step embedding
#
# A raw index is a poor input to a linear layer, so we spread the step over many frequencies
# first.


def timestep_embedding(k, dim):
    half = dim // 2  # Half the coordinates sine, half cosine.
    freqs = torch.exp(
        -math.log(10000.0)
        * torch.arange(half, device=k.device).to(torch.get_default_dtype()) / half
    )
    args = k.to(torch.get_default_dtype())[:, None] * freqs[None, :]  # [B, half].
    return torch.cat([torch.sin(args), torch.cos(args)], dim=1)  # [B, dim].


def num_groups(channels):
    for g in (8, 4, 2, 1):  # Largest valid GroupNorm group count.
        if channels % g == 0:
            return g


# %% 04. The time-conditioned residual block
#
# The step embedding enters as a per-channel offset between the two convolutions, so every layer
# has the current noise level.


class ResidualBlock(nn.Module):
    def __init__(self, in_channels, out_channels, time_embed_dim):
        super().__init__()
        self.norm1 = nn.GroupNorm(num_groups(in_channels), in_channels)
        self.conv1 = nn.Conv2d(in_channels, out_channels,
                               kernel_size=3, padding=1)
        self.time_proj = nn.Linear(time_embed_dim, out_channels)
        self.norm2 = nn.GroupNorm(num_groups(out_channels), out_channels)
        self.conv2 = nn.Conv2d(out_channels, out_channels,
                               kernel_size=3, padding=1)
        self.skip = (
            nn.Identity() if in_channels == out_channels
            else nn.Conv2d(in_channels, out_channels, kernel_size=1)
        )
        self.activation = nn.SiLU()

    def forward(self, x, t_emb):
        h = self.conv1(self.activation(self.norm1(x)))  # GN, SiLU, Conv_1.
        h = h + self.time_proj(t_emb)[:, :, None, None]  # Offset per channel.
        h = self.conv2(self.activation(self.norm2(h)))
        return h + self.skip(x)  # Residual connection.


# %% 05. The U-Net
#
# Three scales, 28 by 28 down to 7 by 7 and back, with a skip at each resolution. The output has
# the shape of the input because it predicts one noise coordinate per pixel.


class UNet(nn.Module):
    def __init__(self, in_channels=1, base_channels=64, time_embed_dim=128):
        super().__init__()
        self.time_embed_dim = time_embed_dim
        self.time_mlp = nn.Sequential(  # Learned representation of step k.
            nn.Linear(time_embed_dim, time_embed_dim),
            nn.SiLU(),
            nn.Linear(time_embed_dim, time_embed_dim),
        )
        c1 = base_channels  # 64 channels at 28 x 28.
        c2 = base_channels * 2  # 128 channels at 14 x 14.
        c3 = base_channels * 4  # 256 channels at 7 x 7.

        self.in_conv = nn.Conv2d(in_channels, c1, kernel_size=3, padding=1)
        self.down1 = ResidualBlock(c1, c1, time_embed_dim)
        self.pool1 = nn.Conv2d(c1, c1, kernel_size=3, stride=2, padding=1)
        self.down2 = ResidualBlock(c1, c2, time_embed_dim)
        self.pool2 = nn.Conv2d(c2, c2, kernel_size=3, stride=2, padding=1)
        self.mid = ResidualBlock(c2, c3, time_embed_dim)
        self.up1 = nn.ConvTranspose2d(c3, c2, kernel_size=4,
                                      stride=2, padding=1)
        self.dec1 = ResidualBlock(c2 + c2, c2, time_embed_dim)
        self.up2 = nn.ConvTranspose2d(c2, c1, kernel_size=4,
                                      stride=2, padding=1)
        self.dec2 = ResidualBlock(c1 + c1, c1, time_embed_dim)
        self.out_conv = nn.Conv2d(c1, in_channels, kernel_size=1)

    def forward(self, x, k):
        t_emb = self.time_mlp(
            timestep_embedding(k, self.time_embed_dim)
        )  # Learned representation of diffusion step k.
        h = self.in_conv(x)
        h1 = self.down1(h, t_emb)  # 28 x 28 encoder features.
        h = self.pool1(h1)  # 28 -> 14.
        h2 = self.down2(h, t_emb)  # 14 x 14 encoder features.
        h = self.pool2(h2)  # 14 -> 7.
        h = self.mid(h, t_emb)  # Compressed representation.
        h = self.up1(h)  # 7 -> 14.
        h = self.dec1(torch.cat([h, h2], dim=1), t_emb)  # Join 14 x 14 skip.
        h = self.up2(h)  # 14 -> 28.
        h = self.dec2(torch.cat([h, h1], dim=1), t_emb)  # Join 28 x 28 skip.
        return self.out_conv(h)  # [B, 1, 28, 28] predicted noise.


# %% 06. One reverse step and the sampler


@torch.no_grad()
def p_sample(model, x_k, k):
    eps_theta = model(x_k, k)  # Predicted noise at step k.
    beta_k = extract(betas, k, x_k.shape)
    alpha_k = extract(alphas, k, x_k.shape)
    abar_k = extract(alphas_cumprod, k, x_k.shape)
    mean = (
        x_k - beta_k / (1.0 - abar_k).sqrt() * eps_theta
    ) / alpha_k.sqrt()  # Reverse mean mu_theta.
    var = extract(posterior_variance, k, x_k.shape)  # beta-tilde_k.
    z = torch.randn_like(x_k)  # Fresh Gaussian draw.
    keep = (k > 1).to(torch.get_default_dtype()).reshape(-1, *((1,) * (x_k.dim() - 1)))
    return mean + keep * var.sqrt() * z  # No noise at the final step.


@torch.no_grad()
def sample(model, num_images, device):
    x = torch.randn(num_images, 1, 28, 28, device=device)  # X_K ~ N(0, I).
    for step in range(NUM_STEPS, 0, -1):  # k = K down to k = 1.
        k = torch.full((num_images,), step, device=device, dtype=torch.long)
        x = p_sample(model, x, k)
    return x  # Generated images in the training range [-1, 1].


# %% 07. The dropped weight and one complete reverse transition
#
# One interior term of the variational bound is beta_k^2 / (2 beta-tilde_k alpha_k (1 - abar_k))
# multiplied by the squared noise error, and L_simple sets every such factor to one. On the
# two-step schedule beta_1 = beta_2 = 0.2 the whole transition is small enough to check by hand.
# The mean there has two equal forms, one built from the noise prediction and one from the clean
# estimate it implies.


def elbo_weight(beta_k, posterior, alpha_k, abar_k) -> float:
    """The schedule constant in front of the squared noise error at one step."""
    return beta_k ** 2 / (2.0 * posterior * alpha_k * (1.0 - abar_k))


def two_step_transition(beta: float, x_k: float, eps_theta: float, z: float) -> dict:
    """The reverse step at k = 2 of the two-step schedule, through both forms of the mean."""
    alpha = 1.0 - beta  # alpha_2.
    abar_1, abar_2 = alpha, alpha ** 2  # 0.8 and 0.64 at beta = 0.2.
    sigma = math.sqrt(1.0 - abar_2)  # 0.6, the corruption standard deviation at k = 2.
    mean = (x_k - beta / sigma * eps_theta) / math.sqrt(alpha)  # The noise form of mu_theta.
    x_0_hat = (x_k - sigma * eps_theta) / math.sqrt(abar_2)  # The clean estimate it implies.
    weights = (math.sqrt(abar_1) * beta / (1.0 - abar_2),  # Weight on the clean estimate.
               math.sqrt(alpha) * (1.0 - abar_1) / (1.0 - abar_2))  # Weight on x_k.
    variance = beta * (1.0 - abar_1) / (1.0 - abar_2)  # beta-tilde_2 = 1/9.
    return {
        "mean_from_noise": mean, "mean_from_clean": weights[0] * x_0_hat + weights[1] * x_k,
        "x_0_hat": x_0_hat, "weights": weights, "variance": variance,
        "x_{k-1}": mean + math.sqrt(variance) * z,
        "elbo_weight": elbo_weight(beta, variance, alpha, abar_2),
    }


# %% 08. The report
def main() -> None:
    parser = common_parser(__doc__.split("\n")[0])
    parser.add_argument("--dataset", choices=["toy", "mnist"], default="toy")
    parser.add_argument("--base-channels", type=int, default=8,
                        help="U-Net base width; the notes uses 64")
    parser.set_defaults(steps=300, batch_size=64, lr=2e-3, samples=4)
    args = parser.parse_args()
    set_seed(args.seed)
    device = pick_device(args.device)

    if args.dataset == "toy":
        from cis6270.data import toy_images

        images = toy_images(512, 28, args.seed).to(device)  # [512, 1, 28, 28] offline images.
        draw = lambda n: images[torch.randint(0, len(images), (n,), device=device)]  # noqa: E731
    else:
        from cis6270.data import mnist_loader

        loader = mnist_loader(batch_size=args.batch_size)
        stream = iter(loader)

        def draw(n):  # One MNIST batch per call, restarting the loader at the end of an epoch.
            nonlocal stream
            batch = next(stream, None)
            if batch is None:
                stream = iter(loader)
                batch = next(stream)
            return batch[0].to(device)

    model = UNet(in_channels=1, base_channels=args.base_channels, time_embed_dim=128).to(device)
    if not args.quiet:
        print(f"device={device} dataset={args.dataset} parameters={count_parameters(model):,}")

    def loss_fn(step: int) -> torch.Tensor:
        x_0 = draw(args.batch_size)  # [B, 1, 28, 28] clean images in [-1, 1].
        k = torch.randint(1, NUM_STEPS + 1, (x_0.shape[0],), device=device)  # One level per item.
        x_k, noise = q_sample(x_0, k)  # Corrupt directly to step k.
        return F.mse_loss(model(x_k, k), noise)  # L_simple on this batch.

    # The notes trains at 2e-4 for 20 epochs; 2e-3 for a few hundred steps is the
    # classroom version of the same run and reaches a usable sampler in under a minute.
    losses = train(model, loss_fn, args.steps, args.lr, quiet=args.quiet)
    images_out = sample(model, args.samples, device)  # [S, 1, 28, 28] after K reverse steps.

    window = min(100, args.steps)  # A short run averages over everything it has.
    reverse = two_step_transition(0.2, 7.4, -1.0, 0.3)
    weight_500 = elbo_weight(betas[500], posterior_variance[500],
                             alphas[500], alphas_cumprod[500])  # The dropped factor at k = 500.
    report = {
        "parameters": count_parameters(model),
        f"loss, first {window} steps": sum(losses[:window]) / window,
        f"loss, last {window} steps": sum(losses[-window:]) / window,
        "beta-tilde at k=500": f"{float(posterior_variance[500]):.5f}",  # 0.01003.
        # The weight L_simple drops, on the thousand-step grid and on the two-step hand schedule.
        "ELBO weight, k=500 and two-step": [round(float(weight_500), 6),
                                            round(reverse["elbo_weight"], 4)],  # 0.005508, 0.625.
        "two-step reverse mean, two forms": [round(reverse["mean_from_noise"], 4),
                                             round(reverse["mean_from_clean"], 4)],  # 8.646.
        "two-step x_0 estimate and weights":
            [round(reverse["x_0_hat"], 4)] + [round(w, 4) for w in reverse["weights"]],
        "two-step reverse variance and x_1": [round(reverse["variance"], 4),
                                              round(reverse["x_{k-1}"], 4)],  # 1/9 and 8.746.
        "generated range": [round(float(images_out.min()), 3), round(float(images_out.max()), 3)],
        "generated standard deviation": float(images_out.std()),
    }
    print_report(f"DDPM on {args.dataset} images", report)
    saved = save_run(args.out, model=model, config=vars(args), report=report, losses=losses)
    if saved is not None:
        try:  # Without torchvision the run still trains, samples and saves its weights.
            from torchvision.utils import save_image

            display = ((images_out.cpu() + 1.0) / 2.0).clamp(0.0, 1.0)  # [-1, 1] -> [0, 1].
            save_image(display, saved / "samples.png", nrow=4)  # A grid of the generated images.
        except ImportError:
            pass


# %% 09. Run the script
if __name__ == "__main__":
    main()