Upload folder using huggingface_hub

.gitattributes

@@ -1,35 +1,2 @@
-*.7z filter=lfs diff=lfs merge=lfs -text
-*.arrow filter=lfs diff=lfs merge=lfs -text
-*.bin filter=lfs diff=lfs merge=lfs -text
-*.bz2 filter=lfs diff=lfs merge=lfs -text
-*.ckpt filter=lfs diff=lfs merge=lfs -text
-*.ftz filter=lfs diff=lfs merge=lfs -text
-*.gz filter=lfs diff=lfs merge=lfs -text
-*.h5 filter=lfs diff=lfs merge=lfs -text
-*.joblib filter=lfs diff=lfs merge=lfs -text
-*.lfs.* filter=lfs diff=lfs merge=lfs -text
-*.mlmodel filter=lfs diff=lfs merge=lfs -text
-*.model filter=lfs diff=lfs merge=lfs -text
-*.msgpack filter=lfs diff=lfs merge=lfs -text
-*.npy filter=lfs diff=lfs merge=lfs -text
-*.npz filter=lfs diff=lfs merge=lfs -text
-*.onnx filter=lfs diff=lfs merge=lfs -text
-*.ot filter=lfs diff=lfs merge=lfs -text
-*.parquet filter=lfs diff=lfs merge=lfs -text
-*.pb filter=lfs diff=lfs merge=lfs -text
-*.pickle filter=lfs diff=lfs merge=lfs -text
-*.pkl filter=lfs diff=lfs merge=lfs -text
-*.pt filter=lfs diff=lfs merge=lfs -text
-*.pth filter=lfs diff=lfs merge=lfs -text
-*.rar filter=lfs diff=lfs merge=lfs -text
 *.safetensors filter=lfs diff=lfs merge=lfs -text
-saved_model/**/* filter=lfs diff=lfs merge=lfs -text
-*.tar.* filter=lfs diff=lfs merge=lfs -text
-*.tar filter=lfs diff=lfs merge=lfs -text
-*.tflite filter=lfs diff=lfs merge=lfs -text
-*.tgz filter=lfs diff=lfs merge=lfs -text
-*.wasm filter=lfs diff=lfs merge=lfs -text
-*.xz filter=lfs diff=lfs merge=lfs -text
-*.zip filter=lfs diff=lfs merge=lfs -text
-*.zst filter=lfs diff=lfs merge=lfs -text
-*tfevents* filter=lfs diff=lfs merge=lfs -text
+*.png filter=lfs diff=lfs merge=lfs -text

README.md

@@ -0,0 +1,230 @@
+---
+library_name: pytorch
+pipeline_tag: text-to-image
+datasets:
+  - ylecun/mnist
+tags:
+  - diffusion
+  - ddim
+  - mnist
+  - handwritten-digits
+  - tiny-model
+---
+
+# Tiny Digit Diffusion 3M
+
+Tiny Digit Diffusion 3M is a **2,767,529-parameter** conditional diffusion
+model that renders a prompt containing one to eight ASCII digits as a
+handwritten grayscale image. It was trained from scratch on dynamically
+composed MNIST digit strings and contains no weights from another generative
+model.
+
+This is a deliberately narrow tiny-model experiment, not a general-purpose
+text-to-image model. Its complete prompt language is a numeric string such as
+`7`, `2026`, or `31415926`.
+
+![Generated samples](sample_grid.png)
+
+## Repository contents
+
+- `model/model.safetensors`: FP32 EMA inference weights
+- `model/config.json`: architecture and diffusion configuration
+- `tiny_digit_diffusion.py`: model, schedule, DDIM sampler, and weight loader
+- `generate_tiny_digits.py`: command-line generation example
+- `train_tiny_digit_diffusion.py`: training script
+- `artifact_metadata.json`: training settings and evaluation summary
+- `sample_grid.png`: generations from the final checkpoint
+
+Optimizer checkpoints, MNIST files, per-epoch samples, and the training log are
+not included in the distribution package.
+
+## Quick start
+
+Create an environment and install the inference dependencies:
+
+```bash
+python -m venv .venv
+source .venv/bin/activate
+pip install -r requirements.txt
+```
+
+Generate an eight-digit image:
+
+```bash
+python generate_tiny_digits.py 31415926 --output 31415926.png
+```
+
+The script accepts the following useful options:
+
+```bash
+python generate_tiny_digits.py 2026 \
+  --steps 50 \
+  --guidance-scale 1.0 \
+  --seed 0 \
+  --device auto \
+  --output 2026.png
+```
+
+The output is a fixed `32 x 256` grayscale PNG. Short prompts are centered on
+the eight available 32-pixel slots.
+
+## Python example
+
+```python
+import torch
+from PIL import Image
+
+from tiny_digit_diffusion import ddim_sample, load_model
+
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+model = load_model("model", device)
+
+image = ddim_sample(
+    model,
+    prompts=["31415926"],
+    device=device,
+    sampling_steps=50,
+    guidance_scale=1.0,
+    seed=0,
+)[0, 0]
+
+pixels = ((image.cpu() + 1) * 127.5).round().clamp(0, 255).byte().numpy()
+Image.fromarray(pixels, mode="L").save("31415926.png")
+```
+
+Prompts must contain only one to eight ASCII digits. Spaces, signs, decimal
+points, letters, and strings longer than eight digits are rejected.
+
+## Recommended inference settings
+
+```yaml
+sampling_method: deterministic DDIM
+sampling_steps: 50
+guidance_scale: 1.0
+output_size: 32 x 256
+```
+
+Use `guidance_scale=1.0`. Larger classifier-free-guidance values are not
+helpful for this checkpoint. In particular, the original diagnostic setting
+of `3.0` over-amplifies the conditioning difference and creates digit-like
+artifacts in padding slots. At `1.0`, the sampler uses the conditional model
+directly and skips the unnecessary unconditional pass.
+
+## Architecture
+
+```yaml
+model: custom conditional diffusion U-Net
+parameter_count: 2,767,529
+weight_dtype: float32
+weight_file_size: approximately 11 MB
+
+image_channels: 1
+image_height: 32
+image_width: 256
+maximum_digits: 8
+slot_width: 32
+
+base_channels: 48
+channel_multipliers: [1, 1.5, 2, 2.6667]
+embedding_dim: 192
+token_embedding_dim: 32
+spatial_condition_channels: 8
+attention_heads: 4
+diffusion_steps: 400
+```
+
+The denoiser has four resolution levels with residual blocks, skip
+connections, a bottleneck self-attention layer, sinusoidal timestep
+conditioning, global prompt conditioning, and slot-aligned spatial digit
+conditioning. Prompt tokens consist of digits `0` through `9`, a padding token,
+and a classifier-free null token.
+
+## Training
+
+Each training example was assembled dynamically from MNIST:
+
+1. Sample a prompt length uniformly from one to eight digits.
+2. Sample an MNIST image independently for every digit.
+3. Center the string on a `32 x 256` canvas.
+4. Apply small horizontal and vertical offsets and an intensity variation.
+5. Train the model to predict noise at a random cosine-schedule timestep.
+
+Training configuration:
+
+```yaml
+dtype: float32
+epochs: 30
+steps: 28,110
+batch_size: 64
+samples_processed: 1,799,040
+optimizer: AdamW
+learning_rate: 2.0e-4
+warmup_steps: 500
+minimum_learning_rate: 2.0e-5
+weight_decay: 0.0
+gradient_clip: 1.0
+condition_dropout: 0.1
+ema_decay: 0.999
+final_recent_noise_mse: 0.015345
+training_time: approximately 70 minutes
+```
+
+The distributed checkpoint is the EMA model, not the raw final optimizer
+weights.
+
+To train from scratch, install the additional dependencies and choose an output
+directory:
+
+```bash
+pip install -r requirements-train.txt
+python train_tiny_digit_diffusion.py \
+  --output-dir runs/tiny_digit_diffusion_3m \
+  --data-dir data/mnist \
+  --epochs 30 \
+  --batch-size 64 \
+  --device cuda
+```
+
+## Evaluation
+
+A separate small MNIST classifier was trained only for automated readability
+measurement. It reached 98.68% accuracy on the MNIST test split. The final
+generator was then evaluated on 240 random prompts: 30 prompts for every length
+from one through eight digits, using 50 DDIM steps and guidance 1.0.
+
+| Prompt length | Exact string | Per-digit accuracy |
+|---:|---:|---:|
+| 1 | 29/30 (96.7%) | 96.7% |
+| 2 | 26/30 (86.7%) | 91.7% |
+| 3 | 26/30 (86.7%) | 95.6% |
+| 4 | 22/30 (73.3%) | 93.3% |
+| 5 | 22/30 (73.3%) | 93.3% |
+| 6 | 14/30 (46.7%) | 88.9% |
+| 7 | 22/30 (73.3%) | 95.7% |
+| 8 | 13/30 (43.3%) | 90.4% |
+| **Overall** | **174/240 (72.5%)** | **92.6%** |
+
+These are OCR-proxy measurements, not human ratings. A generated digit can be
+legible to a person while being classified differently, and the classifier can
+also be confidently wrong. The samples should therefore be considered
+alongside the numeric results.
+
+## Limitations
+
+- Longer strings compound individual digit errors; eight-digit exact match is
+  substantially lower than single-digit accuracy.
+- Some handwritten `3`, `5`, `7`, `8`, and `9` shapes can be ambiguous.
+- The output canvas and maximum string length are fixed by the architecture.
+- The model supports only ASCII digits and does not render signs, punctuation,
+  decimal values, mathematical expressions, or arbitrary text.
+- This checkpoint reproduces MNIST-like handwriting only. It is not suitable
+  for OCR security testing, document generation, or realistic typography.
+- Results vary with the random seed even though each DDIM trajectory is
+  deterministic for a fixed seed.
+
+## Scope
+
+This checkpoint is intended for education, architecture experiments, tests,
+and demonstrations of a complete conditional image generator at a very small
+parameter count. It should not be interpreted as an official MNIST model or as
+an image-generation counterpart of any production diffusion system.

SHA256SUMS

@@ -0,0 +1,11 @@
+32c3d5d9a15f591030c3a116ba7d068f88aa3e59f3f50edfa173e9174fa8d627  .gitattributes
+1f1c500038095b00eef5c89dac5fa260f7f3b132052b4788281874e08329dbf3  README.md
+334a86682c1450ba57900e4ac690557b5ccab31feb67e780b0d8c56eeaabe1ab  artifact_metadata.json
+3cf8877a074fc77f0e50197a800b3614a532dbba1386ed0b30ae5252004747cc  generate_tiny_digits.py
+421b80e7ba8f6b89125b7d6f65bb2356a1df2c7f39b7c125001a093a09d20795  model/config.json
+035ccfb2876e5d78a637e282288e6bfaff227d7a18b3e9fe4562ff2a441458c2  model/model.safetensors
+e18e17c2ae10c8fa2a3947abc31ea4a0eac86fc5c2ecad39578b9fe027999aca  requirements-train.txt
+7db726c6b928d639520fbeb88c6146fb7b33cb6960dc8a01e2cc8380822c51dc  requirements.txt
+d382d9dae68f91c93c9da353ee49a746428ab0c64a4b3e43ac88024765bd4a79  sample_grid.png
+38e6cb6b1daa5b27a6663a6c29455cfdd0255d43ae1b4f487d48417334d211bb  tiny_digit_diffusion.py
+8c70ce74440d2a4930ddef0199d68c191de598fca7220241ac75d46e9b72fb89  train_tiny_digit_diffusion.py

artifact_metadata.json

@@ -0,0 +1,66 @@
+{
+  "parameter_count": 2767529,
+  "training_steps": 28110,
+  "epochs_completed": 30,
+  "final_recent_loss": 0.01534514181315899,
+  "training_seconds": 4215.443364993001,
+  "args": {
+    "output_dir": "runs/tiny_digit_diffusion_3m",
+    "data_dir": "/home/codex/data/mnist",
+    "epochs": 30,
+    "samples_per_epoch": 60000,
+    "batch_size": 64,
+    "learning_rate": 0.0002,
+    "weight_decay": 0.0,
+    "warmup_steps": 500,
+    "grad_clip": 1.0,
+    "condition_dropout": 0.1,
+    "ema_decay": 0.999,
+    "num_workers": 4,
+    "max_steps": 0,
+    "log_steps": 50,
+    "sample_every": 1,
+    "sample_steps": 40,
+    "guidance_scale": 3.0,
+    "checkpoint_every": 5,
+    "seed": 1234,
+    "device": "cuda"
+  },
+  "config": {
+    "image_height": 32,
+    "image_width": 256,
+    "max_digits": 8,
+    "slot_width": 32,
+    "base_channels": 48,
+    "channel_mults": [
+      1,
+      1.5,
+      2,
+      2.6666666666666665
+    ],
+    "embedding_dim": 192,
+    "token_embedding_dim": 32,
+    "condition_channels": 8,
+    "attention_heads": 4,
+    "diffusion_steps": 400
+  },
+  "sample_prompts": [
+    "0",
+    "7",
+    "42",
+    "2026",
+    "12345678",
+    "99999999"
+  ],
+  "recommended_inference": {
+    "sampling_steps": 50,
+    "guidance_scale": 1.0
+  },
+  "evaluation": {
+    "classifier_test_accuracy": 0.9868,
+    "generated_digit_accuracy": 0.926,
+    "generated_sequence_exact_match": 0.725,
+    "evaluation_prompts": 240,
+    "prompts_per_length_1_to_8": 30
+  }
+}

generate_tiny_digits.py

@@ -0,0 +1,53 @@
+#!/usr/bin/env python3
+"""Generate a handwritten 1-8 digit image from a numeric prompt."""
+
+from __future__ import annotations
+
+import argparse
+from pathlib import Path
+
+import torch
+from PIL import Image
+
+from tiny_digit_diffusion import ddim_sample, load_model
+
+
+def parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser()
+    parser.add_argument("prompt", help="A numeric prompt containing 1-8 digits, for example 2026.")
+    parser.add_argument(
+        "--model-dir",
+        default=str(Path(__file__).resolve().parent / "model"),
+    )
+    parser.add_argument("--output", default="generated_digits.png")
+    parser.add_argument("--steps", type=int, default=50)
+    parser.add_argument("--guidance-scale", type=float, default=1.0)
+    parser.add_argument("--seed", type=int, default=0)
+    parser.add_argument("--device", choices=["auto", "cpu", "cuda"], default="auto")
+    return parser.parse_args()
+
+
+def main() -> None:
+    args = parse_args()
+    device = torch.device(
+        "cuda" if args.device == "auto" and torch.cuda.is_available() else
+        "cpu" if args.device == "auto" else args.device
+    )
+    model = load_model(args.model_dir, device)
+    image = ddim_sample(
+        model,
+        [args.prompt],
+        device,
+        sampling_steps=args.steps,
+        guidance_scale=args.guidance_scale,
+        seed=args.seed,
+    )[0, 0]
+    pixels = ((image.cpu() + 1) * 127.5).round().clamp(0, 255).byte().numpy()
+    output = Path(args.output)
+    output.parent.mkdir(parents=True, exist_ok=True)
+    Image.fromarray(pixels, mode="L").save(output)
+    print(output.resolve())
+
+
+if __name__ == "__main__":
+    main()

model/config.json

@@ -0,0 +1,18 @@
+{
+  "image_height": 32,
+  "image_width": 256,
+  "max_digits": 8,
+  "slot_width": 32,
+  "base_channels": 48,
+  "channel_mults": [
+    1,
+    1.5,
+    2,
+    2.6666666666666665
+  ],
+  "embedding_dim": 192,
+  "token_embedding_dim": 32,
+  "condition_channels": 8,
+  "attention_heads": 4,
+  "diffusion_steps": 400
+}

model/model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:035ccfb2876e5d78a637e282288e6bfaff227d7a18b3e9fe4562ff2a441458c2
+size 11081924

requirements-train.txt

@@ -0,0 +1,3 @@
+-r requirements.txt
+torchvision>=0.18
+tqdm>=4.66

requirements.txt

@@ -0,0 +1,3 @@
+torch>=2.3
+safetensors>=0.4
+Pillow>=10.0

tiny_digit_diffusion.py

@@ -0,0 +1,318 @@
+#!/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())

train_tiny_digit_diffusion.py

@@ -0,0 +1,273 @@
+#!/usr/bin/env python3
+"""Train TinyDigitDiffusion on dynamically composed multi-digit MNIST images."""
+
+from __future__ import annotations
+
+import argparse
+import copy
+import json
+import math
+import random
+import time
+from pathlib import Path
+
+import torch
+import torch.nn.functional as F
+from torch.utils.data import DataLoader, Dataset
+from torchvision.datasets import MNIST
+from tqdm import tqdm
+
+from tiny_digit_diffusion import (
+    NULL_TOKEN,
+    PAD_TOKEN,
+    DiffusionSchedule,
+    ModelConfig,
+    TinyDigitDiffusion,
+    count_parameters,
+    ddim_sample,
+    save_prompt_sheet,
+    save_weights,
+)
+
+
+class MultiDigitMNIST(Dataset):
+    def __init__(self, root: str | Path, max_digits: int, samples_per_epoch: int, train: bool = True):
+        self.mnist = MNIST(root=str(root), train=train, download=True)
+        self.images = self.mnist.data
+        self.labels = self.mnist.targets
+        self.max_digits = max_digits
+        self.samples_per_epoch = samples_per_epoch
+
+    def __len__(self) -> int:
+        return self.samples_per_epoch
+
+    def __getitem__(self, index: int):
+        del index
+        length = int(torch.randint(1, self.max_digits + 1, ()).item())
+        tokens = torch.full((self.max_digits,), PAD_TOKEN, dtype=torch.long)
+        start_slot = (self.max_digits - length) // 2
+        canvas = torch.zeros(1, 32, self.max_digits * 32, dtype=torch.float32)
+
+        chosen = torch.randint(0, len(self.images), (length,))
+        for offset, image_index in enumerate(chosen):
+            digit = self.images[image_index].float().div(255)
+            label = int(self.labels[image_index])
+            slot = start_slot + offset
+            tokens[slot] = label
+            x = slot * 32 + 2 + int(torch.randint(-2, 3, ()).item())
+            y = 2 + int(torch.randint(-2, 3, ()).item())
+            x = min(max(x, slot * 32), slot * 32 + 4)
+            y = min(max(y, 0), 4)
+            intensity = float(torch.empty(()).uniform_(0.85, 1.0))
+            canvas[0, y : y + 28, x : x + 28] = torch.maximum(
+                canvas[0, y : y + 28, x : x + 28], digit * intensity
+            )
+        return canvas.mul(2).sub(1), tokens, torch.tensor(length, dtype=torch.long)
+
+
+def parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--output-dir", required=True)
+    parser.add_argument("--data-dir", default="data/mnist")
+    parser.add_argument("--epochs", type=int, default=30)
+    parser.add_argument("--samples-per-epoch", type=int, default=60_000)
+    parser.add_argument("--batch-size", type=int, default=64)
+    parser.add_argument("--learning-rate", type=float, default=2e-4)
+    parser.add_argument("--weight-decay", type=float, default=0.0)
+    parser.add_argument("--warmup-steps", type=int, default=500)
+    parser.add_argument("--grad-clip", type=float, default=1.0)
+    parser.add_argument("--condition-dropout", type=float, default=0.1)
+    parser.add_argument("--ema-decay", type=float, default=0.999)
+    parser.add_argument("--num-workers", type=int, default=4)
+    parser.add_argument("--max-steps", type=int, default=0)
+    parser.add_argument("--log-steps", type=int, default=50)
+    parser.add_argument("--sample-every", type=int, default=1)
+    parser.add_argument("--sample-steps", type=int, default=40)
+    parser.add_argument("--guidance-scale", type=float, default=1.0)
+    parser.add_argument("--checkpoint-every", type=int, default=5)
+    parser.add_argument("--seed", type=int, default=1234)
+    parser.add_argument("--device", choices=["auto", "cpu", "cuda"], default="auto")
+    return parser.parse_args()
+
+
+def update_ema(ema_model: torch.nn.Module, model: torch.nn.Module, decay: float) -> None:
+    with torch.no_grad():
+        for ema_parameter, parameter in zip(ema_model.parameters(), model.parameters()):
+            ema_parameter.lerp_(parameter, 1 - decay)
+        for ema_buffer, buffer in zip(ema_model.buffers(), model.buffers()):
+            ema_buffer.copy_(buffer)
+
+
+def main() -> None:
+    args = parse_args()
+    random.seed(args.seed)
+    torch.manual_seed(args.seed)
+    if torch.cuda.is_available():
+        torch.cuda.manual_seed_all(args.seed)
+    device = torch.device(
+        "cuda" if args.device == "auto" and torch.cuda.is_available() else
+        "cpu" if args.device == "auto" else args.device
+    )
+
+    output_dir = Path(args.output_dir).resolve()
+    output_dir.mkdir(parents=True, exist_ok=True)
+    model_dir = output_dir / "model"
+    sample_dir = output_dir / "samples"
+    checkpoint_dir = output_dir / "checkpoints"
+    model_dir.mkdir(exist_ok=True)
+    sample_dir.mkdir(exist_ok=True)
+    checkpoint_dir.mkdir(exist_ok=True)
+
+    config = ModelConfig()
+    config.save(model_dir / "config.json")
+    model = TinyDigitDiffusion(config).to(device=device, dtype=torch.float32)
+    ema_model = copy.deepcopy(model).requires_grad_(False).eval()
+    parameters = count_parameters(model)
+    print("Device:", device)
+    print("Dtype: float32")
+    print("Parameters:", f"{parameters:,}")
+    print("Image size:", f"{config.image_height}x{config.image_width}")
+    print("Maximum digits:", config.max_digits)
+
+    dataset = MultiDigitMNIST(args.data_dir, config.max_digits, args.samples_per_epoch)
+    loader = DataLoader(
+        dataset,
+        batch_size=args.batch_size,
+        shuffle=False,
+        num_workers=args.num_workers,
+        pin_memory=device.type == "cuda",
+        persistent_workers=args.num_workers > 0,
+        drop_last=True,
+    )
+    steps_per_epoch = len(loader)
+    total_steps = args.max_steps if args.max_steps > 0 else args.epochs * steps_per_epoch
+    if total_steps <= args.warmup_steps:
+        args.warmup_steps = max(0, total_steps // 10)
+    print("Steps per epoch:", steps_per_epoch)
+    print("Training steps:", total_steps)
+
+    optimizer = torch.optim.AdamW(model.parameters(), lr=args.learning_rate, weight_decay=args.weight_decay)
+
+    def lr_factor(step: int) -> float:
+        if args.warmup_steps and step < args.warmup_steps:
+            return max((step + 1) / args.warmup_steps, 1 / args.warmup_steps)
+        progress = (step - args.warmup_steps) / max(total_steps - args.warmup_steps, 1)
+        return 0.1 + 0.9 * 0.5 * (1 + math.cos(progress * math.pi))
+
+    scheduler = torch.optim.lr_scheduler.LambdaLR(optimizer, lr_factor)
+    diffusion = DiffusionSchedule(config.diffusion_steps, device)
+    history: list[dict] = []
+    recent_losses: list[float] = []
+    global_step = 0
+    started = time.monotonic()
+    sample_prompts = ["0", "7", "42", "2026", "12345678", "99999999"]
+
+    model.train()
+    stop = False
+    for epoch in range(1, args.epochs + 1):
+        progress = tqdm(loader, desc=f"epoch {epoch}/{args.epochs}")
+        for clean, tokens, lengths in progress:
+            clean = clean.to(device, non_blocking=True)
+            tokens = tokens.to(device, non_blocking=True)
+            lengths = lengths.to(device, non_blocking=True)
+            drop = torch.rand(len(clean), device=device) < args.condition_dropout
+            tokens = tokens.clone()
+            lengths = lengths.clone()
+            tokens[drop] = NULL_TOKEN
+            lengths[drop] = 0
+
+            timesteps = torch.randint(0, config.diffusion_steps, (len(clean),), device=device)
+            noise = torch.randn_like(clean)
+            noisy = diffusion.add_noise(clean, noise, timesteps)
+            predicted = model(noisy, timesteps, tokens, lengths)
+            loss = F.mse_loss(predicted, noise)
+            if not torch.isfinite(loss):
+                raise RuntimeError(f"Non-finite loss at step {global_step + 1}: {loss}")
+
+            optimizer.zero_grad(set_to_none=True)
+            loss.backward()
+            if args.grad_clip > 0:
+                torch.nn.utils.clip_grad_norm_(model.parameters(), args.grad_clip)
+            optimizer.step()
+            scheduler.step()
+            global_step += 1
+            effective_ema_decay = min(
+                args.ema_decay, (1 + global_step) / (10 + global_step)
+            )
+            update_ema(ema_model, model, effective_ema_decay)
+            recent_losses.append(float(loss.detach().cpu()))
+
+            if global_step % args.log_steps == 0:
+                average = sum(recent_losses[-args.log_steps:]) / min(len(recent_losses), args.log_steps)
+                record = {
+                    "step": global_step,
+                    "epoch": epoch,
+                    "loss": average,
+                    "learning_rate": scheduler.get_last_lr()[0],
+                }
+                history.append(record)
+                progress.set_postfix(loss=f"{average:.4f}", lr=f"{record['learning_rate']:.2e}")
+            if global_step >= total_steps:
+                stop = True
+                break
+
+        if args.sample_every > 0 and (epoch % args.sample_every == 0 or stop):
+            samples = ddim_sample(
+                ema_model, sample_prompts, device,
+                sampling_steps=args.sample_steps,
+                guidance_scale=args.guidance_scale,
+                seed=args.seed + epoch,
+            )
+            save_prompt_sheet(samples, sample_prompts, sample_dir / f"epoch_{epoch:03d}.png")
+        if args.checkpoint_every > 0 and epoch % args.checkpoint_every == 0 and not stop:
+            torch.save(
+                {
+                    "epoch": epoch,
+                    "step": global_step,
+                    "model": model.state_dict(),
+                    "ema_model": ema_model.state_dict(),
+                    "optimizer": optimizer.state_dict(),
+                    "scheduler": scheduler.state_dict(),
+                    "args": vars(args),
+                },
+                checkpoint_dir / f"epoch_{epoch:03d}.pt",
+            )
+        (output_dir / "training_history.json").write_text(
+            json.dumps(history, indent=2), encoding="utf-8"
+        )
+        if stop:
+            break
+
+    save_weights(ema_model, model_dir / "model.safetensors")
+    elapsed = time.monotonic() - started
+    metadata = {
+        "parameter_count": parameters,
+        "training_steps": global_step,
+        "epochs_completed": epoch,
+        "final_recent_loss": sum(recent_losses[-100:]) / min(len(recent_losses), 100),
+        "training_seconds": elapsed,
+        "args": vars(args),
+        "config": vars(config),
+        "sample_prompts": sample_prompts,
+        "recommended_inference": {
+            "sampling_steps": 50,
+            "guidance_scale": 1.0,
+        },
+    }
+    (output_dir / "artifact_metadata.json").write_text(
+        json.dumps(metadata, indent=2), encoding="utf-8"
+    )
+    final_samples = ddim_sample(
+        ema_model, sample_prompts, device,
+        sampling_steps=max(args.sample_steps, 50),
+        guidance_scale=args.guidance_scale,
+        seed=0,
+    )
+    save_prompt_sheet(final_samples, sample_prompts, output_dir / "final_samples.png")
+    print("Done:", output_dir)
+    print("Final recent loss:", metadata["final_recent_loss"])
+    print("Training seconds:", elapsed)
+
+
+if __name__ == "__main__":
+    main()