tiny_digit_diffusion.py

#!/usr/bin/env python3
"""Tiny conditional diffusion model for 1-8 digit MNIST strings."""

from __future__ import annotations

import json
import math
from dataclasses import asdict, dataclass
from pathlib import Path

import torch
import torch.nn.functional as F
from PIL import Image, ImageDraw
from safetensors.torch import load_file, save_file
from torch import nn


PAD_TOKEN = 10
NULL_TOKEN = 11


@dataclass
class ModelConfig:
    image_height: int = 32
    image_width: int = 256
    max_digits: int = 8
    slot_width: int = 32
    base_channels: int = 48
    channel_mults: tuple[int, ...] = (1, 1.5, 2, 8 / 3)
    embedding_dim: int = 192
    token_embedding_dim: int = 32
    condition_channels: int = 8
    attention_heads: int = 4
    diffusion_steps: int = 400

    @classmethod
    def from_json(cls, path: str | Path) -> "ModelConfig":
        values = json.loads(Path(path).read_text(encoding="utf-8"))
        if "channel_mults" in values:
            values["channel_mults"] = tuple(values["channel_mults"])
        return cls(**values)

    def save(self, path: str | Path) -> None:
        Path(path).write_text(json.dumps(asdict(self), indent=2), encoding="utf-8")


def prompt_to_tokens(prompt: str, config: ModelConfig) -> tuple[torch.Tensor, int]:
    prompt = prompt.strip()
    if not prompt or len(prompt) > config.max_digits or not prompt.isascii() or not prompt.isdigit():
        raise ValueError(f"Prompt must contain 1-{config.max_digits} ASCII digits, got {prompt!r}.")
    tokens = torch.full((config.max_digits,), PAD_TOKEN, dtype=torch.long)
    start = (config.max_digits - len(prompt)) // 2
    tokens[start : start + len(prompt)] = torch.tensor([int(ch) for ch in prompt])
    return tokens, len(prompt)


class SinusoidalTimeEmbedding(nn.Module):
    def __init__(self, dim: int):
        super().__init__()
        self.dim = dim

    def forward(self, timesteps: torch.Tensor) -> torch.Tensor:
        half = self.dim // 2
        scale = math.log(10_000) / max(half - 1, 1)
        frequencies = torch.exp(-scale * torch.arange(half, device=timesteps.device))
        angles = timesteps.float()[:, None] * frequencies[None]
        embedding = torch.cat([angles.sin(), angles.cos()], dim=-1)
        if self.dim % 2:
            embedding = F.pad(embedding, (0, 1))
        return embedding


class ResBlock(nn.Module):
    def __init__(self, in_channels: int, out_channels: int, embedding_dim: int):
        super().__init__()
        self.norm1 = nn.GroupNorm(8, in_channels)
        self.conv1 = nn.Conv2d(in_channels, out_channels, 3, padding=1)
        self.embedding = nn.Linear(embedding_dim, out_channels)
        self.norm2 = nn.GroupNorm(8, out_channels)
        self.conv2 = nn.Conv2d(out_channels, out_channels, 3, padding=1)
        self.skip = nn.Conv2d(in_channels, out_channels, 1) if in_channels != out_channels else nn.Identity()

    def forward(self, x: torch.Tensor, embedding: torch.Tensor) -> torch.Tensor:
        h = self.conv1(F.silu(self.norm1(x)))
        h = h + self.embedding(F.silu(embedding))[:, :, None, None]
        h = self.conv2(F.silu(self.norm2(h)))
        return h + self.skip(x)


class SelfAttention2d(nn.Module):
    def __init__(self, channels: int, heads: int):
        super().__init__()
        if channels % heads:
            raise ValueError("Attention channels must be divisible by the number of heads.")
        self.heads = heads
        self.norm = nn.GroupNorm(8, channels)
        self.qkv = nn.Conv2d(channels, channels * 3, 1)
        self.proj = nn.Conv2d(channels, channels, 1)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        batch, channels, height, width = x.shape
        head_dim = channels // self.heads
        qkv = self.qkv(self.norm(x)).view(batch, 3, self.heads, head_dim, height * width)
        q, k, v = qkv.unbind(dim=1)
        q = q.transpose(-2, -1)
        k = k.transpose(-2, -1)
        v = v.transpose(-2, -1)
        h = F.scaled_dot_product_attention(q, k, v)
        h = h.transpose(-2, -1).reshape(batch, channels, height, width)
        return x + self.proj(h)


class TinyDigitDiffusion(nn.Module):
    def __init__(self, config: ModelConfig):
        super().__init__()
        self.config = config
        channels = [int(config.base_channels * mult) for mult in config.channel_mults]
        if any(c % 8 for c in channels):
            raise ValueError(f"All channels must be divisible by 8, got {channels}.")

        time_dim = 64
        self.time_embedding = nn.Sequential(
            SinusoidalTimeEmbedding(time_dim),
            nn.Linear(time_dim, config.embedding_dim),
            nn.SiLU(),
            nn.Linear(config.embedding_dim, config.embedding_dim),
        )
        self.token_embedding = nn.Embedding(12, config.token_embedding_dim)
        self.position_embedding = nn.Parameter(
            torch.randn(config.max_digits, config.token_embedding_dim) * 0.02
        )
        self.length_embedding = nn.Embedding(config.max_digits + 1, config.embedding_dim)
        self.condition_mlp = nn.Sequential(
            nn.Linear(config.max_digits * config.token_embedding_dim, config.embedding_dim),
            nn.SiLU(),
            nn.Linear(config.embedding_dim, config.embedding_dim),
        )
        self.spatial_condition = nn.Embedding(12, config.condition_channels)

        self.input_conv = nn.Conv2d(1 + config.condition_channels, channels[0], 3, padding=1)
        self.down_blocks = nn.ModuleList()
        self.downsamples = nn.ModuleList()
        for i, channel in enumerate(channels):
            self.down_blocks.append(ResBlock(channel, channel, config.embedding_dim))
            if i < len(channels) - 1:
                self.downsamples.append(nn.Conv2d(channel, channels[i + 1], 4, stride=2, padding=1))

        self.mid1 = ResBlock(channels[-1], channels[-1], config.embedding_dim)
        self.mid_attention = SelfAttention2d(channels[-1], config.attention_heads)
        self.mid2 = ResBlock(channels[-1], channels[-1], config.embedding_dim)

        self.upsamples = nn.ModuleList()
        self.up_blocks = nn.ModuleList()
        for i in range(len(channels) - 1, 0, -1):
            self.upsamples.append(nn.ConvTranspose2d(channels[i], channels[i - 1], 4, stride=2, padding=1))
            self.up_blocks.append(ResBlock(channels[i - 1] * 2, channels[i - 1], config.embedding_dim))

        self.output_norm = nn.GroupNorm(8, channels[0])
        self.output_conv = nn.Conv2d(channels[0], 1, 3, padding=1)

    def _condition(
        self, tokens: torch.Tensor, lengths: torch.Tensor, height: int, width: int
    ) -> tuple[torch.Tensor, torch.Tensor]:
        token_features = self.token_embedding(tokens) + self.position_embedding[None]
        global_condition = self.condition_mlp(token_features.flatten(1))
        global_condition = global_condition + self.length_embedding(lengths)

        spatial = self.spatial_condition(tokens).transpose(1, 2)
        spatial = spatial.repeat_interleave(self.config.slot_width, dim=-1)
        if spatial.shape[-1] != width:
            spatial = F.interpolate(spatial, size=width, mode="nearest")
        spatial = spatial[:, :, None, :].expand(-1, -1, height, -1)
        return global_condition, spatial

    def forward(
        self,
        noisy_images: torch.Tensor,
        timesteps: torch.Tensor,
        tokens: torch.Tensor,
        lengths: torch.Tensor,
    ) -> torch.Tensor:
        condition, spatial = self._condition(tokens, lengths, noisy_images.shape[-2], noisy_images.shape[-1])
        embedding = self.time_embedding(timesteps) + condition
        h = self.input_conv(torch.cat([noisy_images, spatial], dim=1))

        skips = []
        for i, block in enumerate(self.down_blocks):
            h = block(h, embedding)
            skips.append(h)
            if i < len(self.downsamples):
                h = self.downsamples[i](h)

        h = self.mid2(self.mid_attention(self.mid1(h, embedding)), embedding)
        for upsample, block, skip in zip(self.upsamples, self.up_blocks, reversed(skips[:-1])):
            h = upsample(h)
            h = block(torch.cat([h, skip], dim=1), embedding)
        return self.output_conv(F.silu(self.output_norm(h)))


def cosine_beta_schedule(steps: int, s: float = 0.008) -> torch.Tensor:
    x = torch.linspace(0, steps, steps + 1, dtype=torch.float64)
    cumulative = torch.cos(((x / steps + s) / (1 + s)) * math.pi * 0.5).square()
    cumulative = cumulative / cumulative[0]
    betas = 1 - cumulative[1:] / cumulative[:-1]
    return betas.clamp(1e-5, 0.999).float()


class DiffusionSchedule:
    def __init__(self, steps: int, device: torch.device):
        self.steps = steps
        self.betas = cosine_beta_schedule(steps).to(device)
        self.alphas = 1 - self.betas
        self.alpha_bars = self.alphas.cumprod(dim=0)

    def add_noise(
        self, clean: torch.Tensor, noise: torch.Tensor, timesteps: torch.Tensor
    ) -> torch.Tensor:
        alpha_bar = self.alpha_bars[timesteps][:, None, None, None]
        return alpha_bar.sqrt() * clean + (1 - alpha_bar).sqrt() * noise


@torch.inference_mode()
def ddim_sample(
    model: TinyDigitDiffusion,
    prompts: list[str],
    device: torch.device,
    sampling_steps: int = 50,
    guidance_scale: float = 3.0,
    seed: int = 0,
) -> torch.Tensor:
    config = model.config
    if not 2 <= sampling_steps <= config.diffusion_steps:
        raise ValueError(f"sampling_steps must be in [2, {config.diffusion_steps}].")
    encoded = [prompt_to_tokens(prompt, config) for prompt in prompts]
    tokens = torch.stack([item[0] for item in encoded]).to(device)
    lengths = torch.tensor([item[1] for item in encoded], device=device)
    null_tokens = torch.full_like(tokens, NULL_TOKEN)
    null_lengths = torch.zeros_like(lengths)

    generator = torch.Generator(device=device).manual_seed(seed)
    images = torch.randn(
        len(prompts), 1, config.image_height, config.image_width,
        generator=generator, device=device,
    )
    schedule = DiffusionSchedule(config.diffusion_steps, device)
    timesteps = torch.linspace(
        config.diffusion_steps - 1, 0, sampling_steps, device=device
    ).round().long()

    was_training = model.training
    model.eval()
    for i, timestep in enumerate(timesteps):
        t = torch.full((len(prompts),), int(timestep), device=device, dtype=torch.long)
        if guidance_scale == 1.0:
            # At exactly 1.0, CFG reduces algebraically to the conditional
            # prediction.  Avoiding the unconditional half nearly halves
            # inference memory and compute and also makes the recommended path
            # straightforward.
            predicted_noise = model(images, t, tokens, lengths)
        else:
            model_input = torch.cat([images, images], dim=0)
            time_input = torch.cat([t, t], dim=0)
            token_input = torch.cat([null_tokens, tokens], dim=0)
            length_input = torch.cat([null_lengths, lengths], dim=0)
            eps_unconditional, eps_conditional = model(
                model_input, time_input, token_input, length_input
            ).chunk(2)
            predicted_noise = eps_unconditional + guidance_scale * (
                eps_conditional - eps_unconditional
            )

        alpha_bar = schedule.alpha_bars[timestep]
        if i + 1 < len(timesteps):
            previous_alpha_bar = schedule.alpha_bars[timesteps[i + 1]]
        else:
            previous_alpha_bar = torch.ones((), device=device)
        predicted_clean = (
            images - (1 - alpha_bar).sqrt() * predicted_noise
        ) / alpha_bar.sqrt()
        predicted_clean = predicted_clean.clamp(-1, 1)
        images = previous_alpha_bar.sqrt() * predicted_clean + (
            1 - previous_alpha_bar
        ).sqrt() * predicted_noise
    if was_training:
        model.train()
    return images.clamp(-1, 1)


def save_prompt_sheet(images: torch.Tensor, prompts: list[str], path: str | Path) -> None:
    images = ((images.detach().cpu() + 1) * 127.5).round().clamp(0, 255).byte()
    row_height = images.shape[-2] + 10
    label_width = 90
    sheet = Image.new("L", (label_width + images.shape[-1], row_height * len(prompts)), 255)
    draw = ImageDraw.Draw(sheet)
    for i, (image, prompt) in enumerate(zip(images, prompts)):
        y = i * row_height
        draw.text((4, y + 10), prompt, fill=0)
        digit_image = Image.fromarray(image[0].numpy(), mode="L")
        sheet.paste(digit_image, (label_width, y + 5))
    Path(path).parent.mkdir(parents=True, exist_ok=True)
    sheet.save(path)


def save_weights(model: TinyDigitDiffusion, path: str | Path) -> None:
    tensors = {name: value.detach().cpu().contiguous() for name, value in model.state_dict().items()}
    save_file(tensors, str(path))


def load_model(model_dir: str | Path, device: torch.device) -> TinyDigitDiffusion:
    model_dir = Path(model_dir)
    config = ModelConfig.from_json(model_dir / "config.json")
    model = TinyDigitDiffusion(config)
    model.load_state_dict(load_file(model_dir / "model.safetensors", device=str(device)))
    return model.to(device).eval()


def count_parameters(model: nn.Module) -> int:
    return sum(parameter.numel() for parameter in model.parameters())