Model card: evaluation results, gallery, figures, eval script

.gitattributes

@@ -84,3 +84,15 @@ samples/base/047.png filter=lfs diff=lfs merge=lfs -text
 samples/base/048.png filter=lfs diff=lfs merge=lfs -text
 samples/base/049.png filter=lfs diff=lfs merge=lfs -text
 samples/base/grid.jpg filter=lfs diff=lfs merge=lfs -text
+images/closest-pairs.jpg filter=lfs diff=lfs merge=lfs -text
+images/fail-010.jpg filter=lfs diff=lfs merge=lfs -text
+images/fail-013.jpg filter=lfs diff=lfs merge=lfs -text
+images/fail-021.jpg filter=lfs diff=lfs merge=lfs -text
+images/lora-strength.jpg filter=lfs diff=lfs merge=lfs -text
+images/pair-002.jpg filter=lfs diff=lfs merge=lfs -text
+images/pair-006.jpg filter=lfs diff=lfs merge=lfs -text
+images/pair-015.jpg filter=lfs diff=lfs merge=lfs -text
+images/pair-033.jpg filter=lfs diff=lfs merge=lfs -text
+images/pair-036.jpg filter=lfs diff=lfs merge=lfs -text
+images/pair-049.jpg filter=lfs diff=lfs merge=lfs -text
+images/progression.jpg filter=lfs diff=lfs merge=lfs -text

README.md

@@ -1,134 +1,467 @@
 ---
-base_model: runwayml/stable-diffusion-v1-5
+base_model: stable-diffusion-v1-5/stable-diffusion-v1-5
+base_model_relation: adapter
 library_name: diffusers
-license: creativeml-openrail-m
 pipeline_tag: text-to-image
+license: creativeml-openrail-m
+datasets:
+- whosouravsharma/text-to-image-diffusiondb-2M
 tags:
 - text-to-image
 - stable-diffusion
+- stable-diffusion-diffusers
 - lora
 - diffusers
-datasets:
-- whosouravsharma/text-to-image-diffusiondb-2M
+thumbnail: https://huggingface.co/whosouravsharma/diffusiondb-sd15-lora/resolve/main/images/pair-002.jpg
+widget:
+- text: a anthropomorphic lion wizard, diffuse lighting, fantasy, intricate, elegant, highly detailed, lifelike, photorealistic, digital painting, artstation, illustration, concept art, smooth, sharp focus, naturalism, trending on byron's - muse, by greg rutkowski and greg staples
+  output:
+    url: images/pair-002.jpg
+- text: a painting of a happy frog under the rain wearing a rainy coat by kazuo oga
+  output:
+    url: images/pair-006.jpg
+- text: mushrooms growing through a human skull in the woods
+  output:
+    url: images/pair-036.jpg
+- text: robots creating new robots in a factory, oil painting by justin gerard, deviantart, hd, 8 k
+  output:
+    url: images/pair-015.jpg
+- text: a painting of a silhouette in the water in front of a eclipse!!! with a red sky, by jeffrey smith, noah bradley, peter mohrbacher, behance contest winner, symbolism, darksynth, poster art, apocalypse art, hellish background
+  output:
+    url: images/pair-033.jpg
+- text: a giant skull with intricate rune carvings and glowing eyes with symmetrically braided lovecraftian tentacles haunting the cosmos by dan mumford, twirling smoke trail, a twisting vortex of dying galaxies, digital art, vivid colors, highly detailed
+  output:
+    url: images/pair-049.jpg
+model-index:
+- name: diffusiondb-sd15-lora (checkpoint-4240)
+  results:
+  - task:
+      type: text-to-image
+    dataset:
+      name: DiffusionDB held-out prompts (report split, 506 images)
+      type: whosouravsharma/text-to-image-diffusiondb-2M
+      split: validation
+      revision: v2-clean
+    metrics:
+    - name: KID (×10³, vs. real held-out images)
+      type: kid
+      value: 2.88
+    - name: FID (vs. real held-out images)
+      type: fid
+      value: 98.0
+    - name: CLIP score (ViT-L/14)
+      type: clip_score
+      value: 28.98
 ---
 
 # diffusiondb-sd15-lora
 
-A LoRA adapter for [Stable Diffusion 1.5](https://huggingface.co/runwayml/stable-diffusion-v1-5),
-fine-tuned on a cleaned, safety-filtered slice of
-[DiffusionDB](https://huggingface.co/datasets/poloclub/diffusiondb).
-
-> **Status: not yet trained.** This repository was created ahead of the first
-> run. Checkpoints and sample renders appear here as training proceeds.
+A rank-32 **style LoRA** for Stable Diffusion 1.5, trained on 13,598 prompt–image
+pairs from [DiffusionDB](https://huggingface.co/datasets/poloclub/diffusiondb).
+It gives SD 1.5 a **more saturated, higher-contrast, illustrative look** while
+following prompts just as well as the base model.
 
-## What this is for
+<Gallery />
 
-DiffusionDB is itself Stable Diffusion 1.x output, collected from the official
-Stable Diffusion Discord. Fine-tuning SD 1.5 on it therefore shifts the model
-toward the **DiffusionDB aesthetic** — the keyword-heavy `artstation /
-intricate / octane render` idiom that its users prompted with. It does not
-push image quality past what SD 1.5 already does, and it is not intended to.
-Judge it on style adherence, not on "is it better than the base model".
+*Each image: plain SD 1.5 on the left, the same prompt and seed with the LoRA on
+the right.*
 
-## Training data
-
-[`whosouravsharma/text-to-image-diffusiondb-2M`](https://huggingface.co/datasets/whosouravsharma/text-to-image-diffusiondb-2M)
-at revision `v2-clean` — 14,598 images from `part_id` 1–20.
-
-| split | examples |
+| | |
 |---|---|
-| train | 13,598 |
-| validation | 1,000 |
+| **What it does** | Adds bolder colour, stronger contrast and more dramatic lighting. Painterly prompts come out more like illustration |
+| **Prompt-following** | Unchanged: CLIP score 28.98 vs. 28.94 for base |
+| **Base model** | [`stable-diffusion-v1-5/stable-diffusion-v1-5`](https://huggingface.co/stable-diffusion-v1-5/stable-diffusion-v1-5) |
+| **Recommended checkpoint** | `checkpoints/checkpoint-4240` (25.5 MB) |
+| **Trigger word** | None. It applies to every prompt |
+| **Try it** | [DiffusionDB SD 1.5 LoRA Space](https://huggingface.co/spaces/whosouravsharma/diffusiondb-sd15-lora). The strength slider gives a live before/after |
+| **License** | CreativeML OpenRAIL-M |
+
+> **Finding worth knowing.** The adapter was trained on DiffusionDB in order to
+> reproduce DiffusionDB's look. It doesn't. DiffusionDB is itself Stable
+> Diffusion 1.x output, and plain SD 1.5 already matches it statistically
+> (KID ≈ 0). The LoRA instead adds a distinct style of its own, and measurably
+> moves outputs *away* from DiffusionDB. See [Evaluation](#evaluation).
 
-The validation split is separated by **normalized prompt group**, not by row.
-DiffusionDB users re-roll the same prompt at many seeds, so a random row split
-leaks: in an earlier revision, 21% of validation prompts had already been seen
-in training. Every image sharing a prompt now lands on the same side.
+## Table of contents
 
-Filters applied upstream: prompts under 4 words dropped, short side ≥ 384 px
-and area ≥ 262,144 px, `image_nsfw` and `prompt_nsfw` below 0.2, at most 2
-images per normalized prompt, exact SHA-256 duplicates removed.
+- [Model details](#model-details)
+- [Uses](#uses)
+- [How to get started](#how-to-get-started)
+- [Evaluation](#evaluation)
+- [Bias, risks, and limitations](#bias-risks-and-limitations)
+- [Training details](#training-details)
+- [Environmental impact](#environmental-impact)
+- [Technical specifications](#technical-specifications)
+- [Reproduce](#reproduce)
+- [Citation](#citation)
 
-## Configuration
+## Model details
 
 | | |
 |---|---|
-| base | `runwayml/stable-diffusion-v1-5` |
-| adapter | LoRA rank 32, alpha 32, on UNet `to_q`/`to_k`/`to_v`/`to_out.0` |
-| text encoder | frozen |
-| resolution | 512×512, centre crop |
-| VAE | `stabilityai/sd-vae-ft-mse` (latents cached ahead of training) |
-| effective batch | 32 (8 × 4 grad accumulation) |
-| optimizer | AdamW, lr 1e-4, cosine schedule, 500 warmup |
-| caption dropout | 10%, for classifier-free guidance |
-| precision | fp16 |
+| **Developed by** | [whosouravsharma](https://huggingface.co/whosouravsharma) |
+| **Model type** | LoRA adapter for a latent diffusion text-to-image model |
+| **Base model** | `stable-diffusion-v1-5/stable-diffusion-v1-5` (trained under its former name, `runwayml/stable-diffusion-v1-5`) |
+| **Adapter** | LoRA, rank 32, alpha 32, on the UNet attention projections (`to_q`, `to_k`, `to_v`, `to_out.0`) |
+| **Trainable parameters** | About 6.4M, against a UNet of about 860M |
+| **Language** | English prompts |
+| **License** | CreativeML OpenRAIL-M, the same as the base model |
 
-Each checkpoint carries a `state.json` recording the exact values it was
-trained with, so the table above can be checked rather than trusted.
+| Resource | Link |
+|---|---|
+| Demo | [Space: diffusiondb-sd15-lora](https://huggingface.co/spaces/whosouravsharma/diffusiondb-sd15-lora) |
+| Inference backend | [Space: diffusiondb-sd15-lora-inference](https://huggingface.co/spaces/whosouravsharma/diffusiondb-sd15-lora-inference) |
+| Training data | [whosouravsharma/text-to-image-diffusiondb-2M](https://huggingface.co/datasets/whosouravsharma/text-to-image-diffusiondb-2M) |
+| Training and evaluation code | [`training/`](./training) and [`eval/`](./eval) in this repository |
 
-## Repository layout
+## Uses
 
-```
-checkpoints/
-  checkpoint-<step>/
-    pytorch_lora_weights.safetensors   the adapter
-    optimizer.pt                       optimizer + scaler state, for resuming
-    state.json                         step, epoch, hyperparameters
-training/                              the scripts that produced all of this
-samples/
-  base/                              vanilla SD 1.5, the comparison baseline
-  checkpoint-<step>/
-    grid.jpg                         contact sheet, all eval prompts
-    000.png … 049.png                individual renders
-    prompts.json                     prompt list, seed, steps, guidance
-```
+### Direct use
 
-## Evaluation
+- Giving SD 1.5 images a punchier, more saturated, illustrative finish, at an
+  adjustable strength.
+- Studying how a fine-tune differs from its base. At strength 0 the output is
+  pixel-identical to plain SD 1.5 at the same seed, so every difference comes
+  from the adapter.
+- As a documented example of a LoRA fine-tune evaluated end to end: a
+  leak-free split, a fixed-seed evaluation set, distribution metrics, a safety
+  check and a memorization check.
 
-50 prompts held out from training entirely (drawn from validation groups) are
-rendered at every checkpoint with a fixed per-prompt seed, so differences
-between contact sheets come from the weights rather than from noise.
-`samples/base/` is vanilla SD 1.5 on the same prompts.
+### Out-of-scope use
 
-Validation loss is logged per epoch, but on a diffusion fine-tune it tracks
-output quality only loosely — it is there to catch divergence, not to rank
-checkpoints.
+- Any use that the CreativeML OpenRAIL-M license prohibits.
+- Reproducing the DiffusionDB distribution. Plain SD 1.5 already does that
+  better (see [Evaluation](#evaluation)).
+- Muted, pastel or low-contrast palettes. The adapter tends to override them
+  (see [Failure cases](#failure-cases)).
+- Resolutions other than 512×512.
+- Photorealistic images of real, identifiable people.
+- Any setting where unfiltered output reaches users with no moderation step.
 
-## Usage
+## How to get started
 
 ```python
 import torch
 from diffusers import StableDiffusionPipeline
 
 pipe = StableDiffusionPipeline.from_pretrained(
-    "runwayml/stable-diffusion-v1-5", torch_dtype=torch.float16
+    "stable-diffusion-v1-5/stable-diffusion-v1-5",
+    torch_dtype=torch.float16,
+    variant="fp16",
 ).to("cuda")
+
 pipe.load_lora_weights(
-    "whosouravsharma/diffusiondb-sd15-lora", subfolder="checkpoints/checkpoint-4240"
+    "whosouravsharma/diffusiondb-sd15-lora",
+    subfolder="checkpoints/checkpoint-4240",
+    weight_name="pytorch_lora_weights.safetensors",
+    adapter_name="diffusiondb",
 )
+pipe.set_adapters(["diffusiondb"], adapter_weights=[1.0])  # 0.0 = plain SD 1.5
 
 image = pipe(
     "a steampunk owl inside a glass jar, intricate detail",
-    num_inference_steps=30, guidance_scale=7.5,
+    num_inference_steps=25,
+    guidance_scale=7.5,
+    generator=torch.Generator("cuda").manual_seed(42),
 ).images[0]
+image.save("owl.png")
+```
+
+This snippet is run verbatim as part of the evaluation (diffusers 0.31.0,
+torch 2.4.0).
+
+**Strength.** 1.0 gives the full effect. Around 0.5 keeps the style while toning
+down the colour when it gets too strong. 1.5 exaggerates the style and starts to change the
+composition.
+
+![LoRA strength sweep](images/lora-strength.jpg)
+
+*The same seed at strength 0, 0.5, 1.0 and 1.5. Strength 0 is plain SD 1.5.*
+
+## Evaluation
+
+Everything below comes from one evaluation run on a single A10G (1 h 53 min).
+The full outputs, including every render and metric, can be reproduced with
+[`eval/eval_job.py`](./eval/eval_job.py).
+
+### Protocol
+
+- **Held-out data.** The 1,000 validation images were split by prompt group,
+  so the model never trained on these prompts. Those 1,000 were then split
+  again by prompt hash: a **selection half** (494 images) to choose a
+  checkpoint, and a **report half** (506 images) for the numbers below. The
+  reported numbers are therefore free of selection bias.
+- **Rendering.** For every model: 30 steps, guidance 7.5, the pipeline's
+  default PNDM scheduler, and seed 42 + the validation row index. Every model
+  sees the same prompts and seeds.
+- **Metrics.**
+  - **KID and FID** against the real held-out images (centre-cropped to 512,
+    the same view as training). Lower means closer to DiffusionDB. KID is the
+    primary metric because FID is biased at this sample size.
+  - **CLIP score** (ViT-L/14) of each image against its own prompt, which
+    measures prompt-following.
+
+### Results
+
+| Model (report half, n = 506) | KID ×10³ ↓ | FID ↓ | CLIP score |
+|---|---|---|---|
+| SD 1.5 base | **0.14 ± 0.32** | **95.4** | 28.94 |
+| + LoRA `checkpoint-2000` | 2.87 ± 0.64 | 98.7 | 28.79 |
+| + LoRA `checkpoint-4240` | 2.88 ± 0.61 | 98.0 | 28.98 |
+| Real held-out images (reference) | — | — | 28.95 |
+
+*KID is the mean ± standard deviation over 100 subsets of 253 images each.
+CLIP scores have a standard deviation of about 3.9 per image, which makes the
+standard error of each mean about 0.17.*
+
+**What the numbers say**
+
+1. **Base SD 1.5 already matches DiffusionDB.** Its KID is indistinguishable
+   from zero. That's expected, because DiffusionDB was generated with SD 1.x.
+2. **The LoRA moves outputs away from DiffusionDB.** KID rises by about
+   2.7×10⁻³, roughly four times the combined spread of the two KID estimates. This is the style shift visible in
+   the gallery, and it is not an approach to the training distribution.
+3. **Prompt-following is unchanged.** All CLIP scores are within noise of each
+   other and of the real images.
+4. **Checkpoints from step 2,000 onward are statistically identical.** On the
+   selection half, step 2,000 won narrowly (KID 2.50 vs. 2.62 for step 4,240
+   and 2.90 for step 3,000, all ± about 0.6). On the report half the two are
+   tied. `checkpoint-4240` stays the recommended checkpoint.
+
+**A likely cause, not yet tested.** The added saturation looks like
+classifier-free guidance overshooting. The adapter may have sharpened the gap
+between the conditional and unconditional predictions, so guidance 7.5
+behaves like a higher value. Sweeping the guidance scale with the adapter would
+confirm or rule this out.
+
+### Validation loss
+
+![Validation loss by checkpoint](images/val-loss.png)
+
+This is noise-prediction MSE on the 1,000 held-out images, with identical noise
+and timesteps for every model. The loss falls 0.9% below the base model
+(0.1496 → 0.1482), mostly in the first 2,000 steps, and is flat after about
+3,000. There's no sign of divergence or overfitting. As is common for diffusion
+fine-tunes, a lower denoising loss did not translate into samples closer to the
+data (see KID above).
+
+### How the style develops during training
+
+![Samples at each checkpoint](images/progression.jpg)
+
+*The same prompts and seeds at every checkpoint. Most of the style is in place
+by step 1,500. Later checkpoints refine it rather than change it.*
+
+### Failure cases
+
+![Moon prompt](images/fail-013.jpg)
+*"shattering of the moon's surface, digital art, illustration". The detailed
+scene collapses into a flat, sparse composition.*
+
+![Vodka bottle prompt](images/fail-010.jpg)
+*"glass vodka bottle by shusei nagaoka, kaws, david rudnick, airbrush on
+canvas, **pastel colors**, cell-shaded, 8 k". The adapter's saturation
+overrides the requested pastel palette.*
+
+![Polar bear prompt](images/fail-021.jpg)
+*"long distance shot of a tiny cute polar bear on a tiny iceberg in the middle
+of the ocean, sunset, atmospheric, hazy". The animal's anatomy degrades, and
+the hazy mood is lost.*
+
+### Correctness checks
+
+| Check | Result |
+|---|---|
+| Strength 0 equals plain SD 1.5 (same seed) | ✅ Max pixel difference 0 |
+| Unloading the adapter restores plain SD 1.5 | ✅ Max pixel difference 0 |
+| Same seed gives the same image | ✅ Max pixel difference 0 |
+| The adapter changes the output | ✅ Mean pixel difference 41.7 / 255 |
+| The usage snippet above runs as written | ✅ |
+
+## Bias, risks, and limitations
+
+### Safety
+
+The SD safety checker was run afterwards over the report-half renders; it is
+disabled during generation in the demo.
+
+| Images (n = 506) | Flagged |
+|---|---|
+| Real held-out images | 7 (1.4%) |
+| SD 1.5 base | 13 (2.6%) |
+| + LoRA `checkpoint-2000` | 10 (2.0%) |
+| + LoRA `checkpoint-4240` | 13 (2.6%) |
+
+The adapter doesn't raise the flag rate over the base model. The counts are
+small, so treat these as rough rates. The training data was filtered with
+DiffusionDB's own NSFW scores (below 0.2), but that classifier is noisy, so the
+filtering is not a guarantee.
+
+### Memorization
+
+For each report-half render, we found the most similar of the 13,598 training
+images by CLIP ViT-L/14 image-embedding cosine. Real held-out images, which
+have different prompts but the same style, give the baseline for how similar
+two unrelated DiffusionDB images can be.
+
+| Images (n = 506) | Median nearest-neighbour cosine | ≥ 0.95 | Max |
+|---|---|---|---|
+| Real held-out images | 0.845 | 12 | 0.995 |
+| SD 1.5 base | 0.833 | 3 | 0.971 |
+| + LoRA `checkpoint-2000` | 0.829 | 2 | 0.959 |
+| + LoRA `checkpoint-4240` | 0.833 | 2 | 0.961 |
+
+The LoRA's renders are no closer to the training images than the base model's.
+
+![Closest render / training-image pairs](images/closest-pairs.jpg)
+
+*The 8 closest pairs for `checkpoint-2000`. They share a subject and style, not
+a composition, so none is a copy. CLIP similarity is a proxy for copying, not
+proof either way.*
+
+### Limitations
+
+- **The style is fixed.** Saturation and contrast go up on every prompt,
+  including ones that ask for muted or pastel colours.
+- **Occasional composition loss.** Some detailed scenes are simplified (see
+  [Failure cases](#failure-cases)).
+- **Small training set.** 13,598 images is enough to fine-tune, not to train
+  from scratch.
+- **It inherits SD 1.x flaws.** The training images were themselves SD 1.x
+  outputs, including their artifacts.
+- **Centre crop only.** There was no aspect-ratio bucketing, and roughly half
+  the source images aren't square, so their edges were lost in training.
+- **The text encoder was frozen,** so the model understands prompts exactly as
+  SD 1.5 does.
+
+### Risks and recommendations
+
+- **Turn the safety checker back on.** The demo Space disables it. Put a
+  safety checker or other moderation step in front of any public-facing use.
+- **It inherits SD 1.5's biases.** SD 1.5 and its LAION training data carry
+  social and cultural biases. This adapter doesn't reduce them.
+- **Artist names in prompts.** Many DiffusionDB prompts name living artists,
+  and the adapter learned from those prompt–image pairs.
+
+## Training details
+
+### Training data
+
+[`whosouravsharma/text-to-image-diffusiondb-2M`](https://huggingface.co/datasets/whosouravsharma/text-to-image-diffusiondb-2M)
+at revision `v2-clean`: 13,598 train and 1,000 validation images, built from
+parts 1–20 of [`poloclub/diffusiondb`](https://huggingface.co/datasets/poloclub/diffusiondb).
+The validation split is made **by prompt group**: DiffusionDB users often ran
+the same prompt at many seeds, so every image sharing a normalized prompt goes
+to the same side.
+
+<details>
+<summary>Filtering, stage by stage (from the dataset's <code>manifest.json</code>)</summary>
+
+| Stage | Rows kept | Removed |
+|---|---|---|
+| Source metadata | 20,000 | — |
+| Prompt has at least 4 words | 19,513 | 487 |
+| Short side ≥ 384 px and area ≥ 262,144 px | 19,426 | 87 |
+| Image and prompt NSFW scores below 0.2 | 15,451 | 3,975 |
+| At most 2 images per normalized prompt | 14,606 | 845 |
+| Exact SHA-256 duplicates removed | **14,598** | 8 |
+
+</details>
+
+### Training procedure
+
+Every image was encoded once with
+[`stabilityai/sd-vae-ft-mse`](https://huggingface.co/stabilityai/sd-vae-ft-mse)
+after resizing and a centre crop to 512×512. The results are stored as fp16
+latents: both the posterior mean and its log-variance, so each training step
+samples a fresh latent. The text embeddings were not cached, which keeps 10%
+caption dropout possible for classifier-free guidance.
+
+<details>
+<summary>Hyperparameters</summary>
+
+| | |
+|---|---|
+| Objective | ε-prediction (MSE on the noise), DDPM noise schedule |
+| Epochs | 10 (424 optimizer steps per epoch, 4,240 in total) |
+| Batch size | 8 per step × 4 gradient-accumulation steps = 32 |
+| Optimizer | AdamW, lr 1e-4, weight decay 1e-2, grad-norm clip 1.0 |
+| Schedule | 500 linear warmup steps, then cosine decay |
+| Caption dropout | 10% |
+| LoRA init | Gaussian |
+| Precision | Frozen weights in fp16; LoRA weights in fp32; forward pass under fp16 autocast |
+| Seed | 42 |
+| Checkpoints | Every 500 steps, plus the final step (4,240) |
+
+Each checkpoint has a `state.json` recording the exact values it was trained
+with.
+
+</details>
+
+## Environmental impact
+
+| | |
+|---|---|
+| Hardware | 1× NVIDIA A10G (Hugging Face Jobs, `a10g-small`) |
+| Training | 2.68 h (plus 0.25 h caching latents and 0.04 h baseline samples) |
+| Evaluation | 1.88 h |
+| Total | About 4.9 GPU-hours |
+| Carbon emitted | Not measured |
+
+## Technical specifications
+
+SD 1.5 latent diffusion: CLIP ViT-L/14 text encoder (frozen), a UNet of about
+860M parameters (frozen, with LoRA on the attention projections), and a VAE
+with 8× downsampling. Every stage runs as a standalone PEP 723 script on
+Hugging Face Jobs. Each job pulls its inputs from the Hub, and training can
+resume from any checkpoint with `RESUME_FROM`.
+
+<details>
+<summary>Repository layout</summary>
+
 ```
+checkpoints/
+  checkpoint-<step>/                  step = 500, 1000, … 4000, 4240
+    pytorch_lora_weights.safetensors  the adapter
+    optimizer.pt                      optimizer and grad-scaler state, for resuming
+    state.json                        step, epoch, hyperparameters
+samples/base/                         plain SD 1.5 on the 50 eval prompts
+images/                               figures used in this card
+training/                             data caching, training and sampling scripts
+eval/eval_job.py                      the evaluation that produced the numbers above
+```
+
+</details>
+
+## Reproduce
+
+```bash
+# training (run from training/)
+python3 main.py baseline     # plain SD 1.5 samples, for comparison
+python3 main.py latents      # VAE-encode the dataset once
+python3 main.py train        # LoRA fine-tune
+python3 main.py sample checkpoint-4240
+
+# evaluation: needs no token, writes only to the mounted bucket
+hf jobs uv run --flavor a10g-small --timeout 4h \
+  -v hf://buckets/<you>/<bucket>:/out eval/eval_job.py
+```
+
+## Citation
+
+Training data comes from DiffusionDB (CC0 1.0):
+
+```bibtex
+@article{wangDiffusionDBLargescalePrompt2022,
+  title   = {DiffusionDB: A Large-Scale Prompt Gallery Dataset for Text-to-Image Generative Models},
+  author  = {Wang, Zijie J. and Montoya, Evan and Munechika, David and Yang, Haoyang and Hoover, Benjamin and Chau, Duen Horng},
+  journal = {arXiv:2210.14896 [cs]},
+  year    = {2022},
+  url     = {https://arxiv.org/abs/2210.14896}
+}
+```
+
+## Model card contact
 
-## Limitations
-
-- Trained on ~13.6k images — fine-tuning scale, not from-scratch scale.
-- Source images are SD 1.x generations, so artifacts of that model are
-  reproduced along with its style.
-- 512×512 centre crop discards roughly 16% of the source image area; ~10% of
-  images lose more than 40% of their frame.
-- NSFW filtering relies on DiffusionDB's own classifier scores at a 0.2
-  threshold. That classifier is noisy, so the training set is filtered, not
-  guaranteed clean.
-- Prompts carry heavy style boilerplate (~22% mention `artstation`), which the
-  adapter will have learned as part of the aesthetic.
-
-## License & attribution
-
-Adapter released under CreativeML OpenRAIL-M, matching the SD 1.5 base model.
-Training data derives from
-[poloclub/diffusiondb](https://huggingface.co/datasets/poloclub/diffusiondb),
-CC0-1.0. Wang et al., 2022, [arXiv:2210.14896](https://arxiv.org/abs/2210.14896).
+Open a discussion in the
+[Community tab](https://huggingface.co/whosouravsharma/diffusiondb-sd15-lora/discussions).

eval/eval_job.py

@@ -0,0 +1,826 @@
+# /// script
+# requires-python = ">=3.10"
+# dependencies = [
+#   "torch==2.4.0",
+#   "diffusers==0.31.0",
+#   "transformers==4.44.2",
+#   "peft==0.13.2",
+#   "accelerate==0.34.2",
+#   "safetensors",
+#   "huggingface_hub>=0.24,<1.0",
+#   "datasets>=2.19,<4",
+#   "torchmetrics[image]>=1.4",
+#   "torch-fidelity",
+#   "scipy",
+#   "numpy<2",
+#   "pillow>=10.1",
+#   "matplotlib",
+# ]
+# ///
+"""Evaluate the DiffusionDB SD 1.5 LoRA without writing to any repo.
+
+    hf jobs uv run --flavor a10g-small --timeout 4h \\
+        -v hf://buckets/whosouravsharma/jobs-artifacts:/out \\
+        eval/eval_job.py
+
+Deliberately launched WITHOUT --secrets HF_TOKEN. Every input is public, so
+the job has no credentials that could write to the model repo, the dataset
+repo or either Space. The only writable location is the mounted bucket, and
+everything lands under /out/eval/<RUN_NAME>/.
+
+Reads (all public, read-only):
+    whosouravsharma/diffusiondb-sd15-lora            checkpoints
+    whosouravsharma/text-to-image-diffusiondb-2M     v2-clean images + prompts,
+                                                     latents-512 validation latents
+    stable-diffusion-v1-5/stable-diffusion-v1-5      base model, safety checker
+    openai/clip-vit-large-patch14                    CLIP score + similarity
+
+Stages (STAGES env var, comma-separated; default: all, in this order):
+    snippet       run the model card's usage snippet verbatim
+    checks        LoRA-scale-0 == base, unload restores base, same seed == same image
+    loss          validation loss for base + every checkpoint (fixed noise)
+    metrics       select a checkpoint on half the validation set (KID), then report
+                  KID / FID / CLIP score on the other half, base vs selected (+ final)
+    grids         the 50 eval prompts, base vs selected (+ final), same seeds
+    sweep         LoRA strength 0 / 0.5 / 1.0 / 1.5 on a few prompts
+    progression   base + every checkpoint on a few prompts
+    safety        SD safety-checker flag rate on the report-half renders and real images
+    memorization  nearest training image (CLIP cosine) for each report-half render
+
+SMOKE=1 shrinks every stage to a few examples, to catch bugs cheaply.
+Each stage is independent: a failure is recorded in status.json and the next
+stage still runs.
+"""
+
+from __future__ import annotations
+
+import hashlib
+import io
+import json
+import math
+import os
+import re
+import time
+import traceback
+from pathlib import Path
+
+import numpy as np
+import torch
+import torch.nn.functional as F
+from PIL import Image, ImageDraw, ImageFont
+
+# ---------------------------------------------------------------------------
+# configuration
+# ---------------------------------------------------------------------------
+
+MODEL_REPO = "whosouravsharma/diffusiondb-sd15-lora"
+DATASET_REPO = "whosouravsharma/text-to-image-diffusiondb-2M"
+DATASET_REVISION = "v2-clean"
+LATENTS_REVISION = "latents-512"
+BASE_MODEL = "stable-diffusion-v1-5/stable-diffusion-v1-5"
+CLIP_MODEL = "openai/clip-vit-large-patch14"
+CHECKPOINT_DIR = "checkpoints"
+ADAPTER = "diffusiondb"
+FINAL = "checkpoint-4240"
+
+# Same render settings as training/sample_job.py, so grids line up with it.
+SEED = 42
+STEPS = 30
+GUIDANCE = 7.5
+RENDER_BATCH = int(os.environ.get("RENDER_BATCH", 4))
+LOSS_BATCH = 8
+
+SMOKE = os.environ.get("SMOKE") == "1"
+RUN_NAME = os.environ.get("RUN_NAME") or time.strftime("%Y%m%dT%H%M%S") + ("-smoke" if SMOKE else "")
+OUT = Path(os.environ.get("OUT_ROOT", "/out/eval")) / RUN_NAME
+
+ALL_STAGES = ["snippet", "checks", "loss", "metrics", "grids", "sweep",
+              "progression", "safety", "memorization"]
+STAGES = [s.strip() for s in os.environ.get("STAGES", ",".join(ALL_STAGES)).split(",") if s.strip()]
+
+CANDIDATES = os.environ.get(
+    "CANDIDATES", "checkpoint-2000,checkpoint-3000,checkpoint-4240").split(",")
+SWEEP_SCALES = [0.0, 0.5, 1.0, 1.5]
+# Indices into eval_prompts.json.
+SWEEP_PROMPTS = [int(i) for i in os.environ.get("SWEEP_PROMPTS", "2,6,8,0").split(",")]
+PROGRESSION_PROMPTS = [int(i) for i in os.environ.get("PROGRESSION_PROMPTS", "2,6,8").split(",")]
+
+# Smoke-test limits.
+LIMIT_LOSS = 32 if SMOKE else None
+LIMIT_HALF = 16 if SMOKE else None
+LIMIT_GRID = 4 if SMOKE else None
+LIMIT_TRAIN = 256 if SMOKE else None
+LIMIT_CKPTS = 2 if SMOKE else None
+
+DEVICE = "cuda"
+DTYPE = torch.float16
+
+STATUS: dict = {"run": RUN_NAME, "smoke": SMOKE, "stages": {}, "started": time.time()}
+STATE: dict = {}  # shared between stages: selected checkpoint, report renders, ...
+
+
+def log(message: str) -> None:
+    print(f"[{time.strftime('%H:%M:%S')}] {message}", flush=True)
+
+
+def write_json(path: Path, payload) -> None:
+    path.parent.mkdir(parents=True, exist_ok=True)
+    path.write_text(json.dumps(payload, indent=2, default=str))
+
+
+def save_status() -> None:
+    write_json(OUT / "status.json", STATUS)
+
+
+# ---------------------------------------------------------------------------
+# inputs
+# ---------------------------------------------------------------------------
+
+def list_checkpoints() -> list[str]:
+    from huggingface_hub import HfApi
+
+    names = {
+        f.split("/")[1] for f in HfApi().list_repo_files(MODEL_REPO)
+        if f.startswith(f"{CHECKPOINT_DIR}/checkpoint-")
+    }
+    ordered = sorted(names, key=lambda n: int(n.split("-")[1]))
+    return ordered[:LIMIT_CKPTS] if LIMIT_CKPTS else ordered
+
+
+def eval_prompts() -> list[str]:
+    from huggingface_hub import hf_hub_download
+
+    path = hf_hub_download(DATASET_REPO, "eval_prompts.json",
+                           repo_type="dataset", revision=DATASET_REVISION)
+    return json.loads(Path(path).read_text())
+
+
+def center_crop(image: Image.Image, size: int = 512) -> Image.Image:
+    """Same view as training: resize the short side, crop the centre square."""
+    image = image.convert("RGB")
+    width, height = image.size
+    scale = size / min(width, height)
+    image = image.resize((max(size, round(width * scale)), max(size, round(height * scale))),
+                         Image.BICUBIC)
+    width, height = image.size
+    left, top = (width - size) // 2, (height - size) // 2
+    return image.crop((left, top, left + size, top + size))
+
+
+def normalize_prompt(prompt: str) -> str:
+    return re.sub(r"[^a-z0-9]+", " ", prompt.lower()).strip()
+
+
+def validation_halves():
+    """Split the 1,000 validation rows into a selection half and a report half.
+
+    Split by prompt hash so both images of a prompt group land on the same
+    side. The halves are used for different purposes: choosing a checkpoint
+    on one and reporting numbers on the other keeps the reported numbers
+    free of selection bias.
+    """
+    from datasets import load_dataset
+
+    data = load_dataset(DATASET_REPO, split="validation", revision=DATASET_REVISION)
+    halves = {"select": [], "report": []}
+    for index, row in enumerate(data):
+        key = normalize_prompt(row["prompt"])
+        half = "select" if int(hashlib.sha1(key.encode()).hexdigest(), 16) % 2 == 0 else "report"
+        halves[half].append({
+            "index": index,
+            "prompt": row["prompt"],
+            "image": center_crop(row["image"]),
+            "seed": SEED + index,  # same seed for every model on this row
+        })
+    if LIMIT_HALF:
+        halves = {k: v[:LIMIT_HALF] for k, v in halves.items()}
+    log(f"validation halves: select {len(halves['select'])}, report {len(halves['report'])}")
+    return halves
+
+
+# ---------------------------------------------------------------------------
+# pipeline
+# ---------------------------------------------------------------------------
+
+class Runner:
+    """One SD 1.5 pipeline; LoRA checkpoints are swapped in and out of it."""
+
+    def __init__(self) -> None:
+        from diffusers import StableDiffusionPipeline
+
+        self.pipe = StableDiffusionPipeline.from_pretrained(
+            BASE_MODEL, torch_dtype=DTYPE, variant="fp16", use_safetensors=True,
+            safety_checker=None, requires_safety_checker=False,
+        ).to(DEVICE)
+        self.pipe.set_progress_bar_config(disable=True)
+        self.current: str | None = None
+
+    def use(self, checkpoint: str | None, scale: float = 1.0) -> None:
+        if checkpoint in (None, "base"):
+            if self.current:
+                self.pipe.unload_lora_weights()
+                self.current = None
+            return
+        if self.current != checkpoint:
+            if self.current:
+                self.pipe.unload_lora_weights()
+            self.pipe.load_lora_weights(
+                MODEL_REPO, subfolder=f"{CHECKPOINT_DIR}/{checkpoint}",
+                weight_name="pytorch_lora_weights.safetensors", adapter_name=ADAPTER,
+            )
+            self.current = checkpoint
+        self.pipe.set_adapters([ADAPTER], adapter_weights=[float(scale)])
+
+    @torch.no_grad()
+    def render(self, prompts: list[str], seeds: list[int]) -> list[Image.Image]:
+        images: list[Image.Image] = []
+        for start in range(0, len(prompts), RENDER_BATCH):
+            batch = prompts[start:start + RENDER_BATCH]
+            generators = [torch.Generator(DEVICE).manual_seed(int(s))
+                          for s in seeds[start:start + RENDER_BATCH]]
+            images += self.pipe(batch, num_inference_steps=STEPS, guidance_scale=GUIDANCE,
+                                generator=generators).images
+        return images
+
+
+_RUNNER: Runner | None = None
+
+
+def runner() -> Runner:
+    global _RUNNER
+    if _RUNNER is None:
+        log(f"loading {BASE_MODEL}")
+        _RUNNER = Runner()
+    return _RUNNER
+
+
+def render_model(model: str, prompts: list[str], seeds: list[int], scale: float = 1.0):
+    run = runner()
+    run.use(model, scale)
+    started = time.time()
+    images = run.render(prompts, seeds)
+    log(f"  rendered {len(images)} with {model} (scale {scale}) "
+        f"in {time.time() - started:.0f}s")
+    return images
+
+
+def selected() -> str:
+    return STATE.get("selected", FINAL)
+
+
+def report_models() -> list[str]:
+    models = ["base", selected()]
+    if selected() != FINAL:
+        models.append(FINAL)
+    return models
+
+
+# ---------------------------------------------------------------------------
+# image helpers
+# ---------------------------------------------------------------------------
+
+def font(size: int = 16):
+    try:
+        return ImageFont.load_default(size=size)
+    except TypeError:
+        return ImageFont.load_default()
+
+
+def labelled_grid(rows: list[list[Image.Image]], column_labels: list[str],
+                  row_labels: list[str] | None = None, thumb: int = 256) -> Image.Image:
+    top, left, pad = 32, (220 if row_labels else 0), 6
+    width = left + len(column_labels) * (thumb + pad) + pad
+    height = top + len(rows) * (thumb + pad) + pad
+    sheet = Image.new("RGB", (width, height), (252, 252, 251))
+    draw = ImageDraw.Draw(sheet)
+    for c, label in enumerate(column_labels):
+        draw.text((left + pad + c * (thumb + pad) + 4, 8), label, fill=(11, 11, 11), font=font(16))
+    for r, row in enumerate(rows):
+        y = top + pad + r * (thumb + pad)
+        if row_labels:
+            text = row_labels[r]
+            lines = [text[i:i + 26] for i in range(0, min(len(text), 26 * 6), 26)]
+            draw.multiline_text((8, y + 4), "\n".join(lines), fill=(82, 81, 78), font=font(13))
+        for c, image in enumerate(row):
+            sheet.paste(image.resize((thumb, thumb), Image.LANCZOS), (left + pad + c * (thumb + pad), y))
+    return sheet
+
+
+def contact_sheet(images: list[Image.Image], columns: int = 5, thumb: int = 320) -> Image.Image:
+    rows = (len(images) + columns - 1) // columns
+    sheet = Image.new("RGB", (columns * thumb, rows * thumb), (18, 18, 22))
+    for i, image in enumerate(images):
+        sheet.paste(image.resize((thumb, thumb), Image.LANCZOS), ((i % columns) * thumb, (i // columns) * thumb))
+    return sheet
+
+
+def to_tensor(images: list[Image.Image]) -> torch.Tensor:
+    """PIL list -> float tensor in [0, 1], shape (N, 3, H, W)."""
+    array = np.stack([np.asarray(im.convert("RGB"), dtype=np.uint8) for im in images])
+    return torch.from_numpy(array).permute(0, 3, 1, 2).float() / 255.0
+
+
+def max_pixel_diff(a: Image.Image, b: Image.Image) -> tuple[int, float]:
+    x = np.asarray(a, dtype=np.int16)
+    y = np.asarray(b, dtype=np.int16)
+    d = np.abs(x - y)
+    return int(d.max()), float(d.mean())
+
+
+# ---------------------------------------------------------------------------
+# CLIP
+# ---------------------------------------------------------------------------
+
+_CLIP = None
+
+
+def clip():
+    global _CLIP
+    if _CLIP is None:
+        from transformers import CLIPModel, CLIPProcessor
+
+        _CLIP = (CLIPModel.from_pretrained(CLIP_MODEL, torch_dtype=DTYPE).to(DEVICE).eval(),
+                 CLIPProcessor.from_pretrained(CLIP_MODEL))
+    return _CLIP
+
+
+@torch.no_grad()
+def clip_image_embeddings(images: list[Image.Image], batch: int = 64) -> torch.Tensor:
+    model, processor = clip()
+    out = []
+    for start in range(0, len(images), batch):
+        pixels = processor(images=images[start:start + batch], return_tensors="pt")["pixel_values"]
+        emb = model.get_image_features(pixel_values=pixels.to(DEVICE, DTYPE))
+        out.append(F.normalize(emb.float(), dim=-1))
+    return torch.cat(out)
+
+
+@torch.no_grad()
+def clip_scores(images: list[Image.Image], prompts: list[str], batch: int = 64) -> np.ndarray:
+    """CLIP score as in torchmetrics: 100 * max(cos(image, text), 0).
+
+    Prompts are truncated to CLIP's 77-token limit, as the SD text encoder does.
+    """
+    model, processor = clip()
+    image_emb = clip_image_embeddings(images, batch)
+    scores = []
+    for start in range(0, len(prompts), batch):
+        tokens = processor(text=prompts[start:start + batch], return_tensors="pt",
+                           padding=True, truncation=True, max_length=77)
+        text = model.get_text_features(input_ids=tokens["input_ids"].to(DEVICE),
+                                       attention_mask=tokens["attention_mask"].to(DEVICE))
+        text = F.normalize(text.float(), dim=-1)
+        cos = (image_emb[start:start + batch] * text).sum(-1)
+        scores.append((100 * cos.clamp(min=0)).cpu().numpy())
+    return np.concatenate(scores)
+
+
+# ---------------------------------------------------------------------------
+# stages
+# ---------------------------------------------------------------------------
+
+def stage_snippet() -> dict:
+    """The model card's "How to get started" code, verbatim (plus a save path)."""
+    import diffusers
+    from diffusers import StableDiffusionPipeline
+
+    pipe = StableDiffusionPipeline.from_pretrained(
+        "stable-diffusion-v1-5/stable-diffusion-v1-5",
+        torch_dtype=torch.float16,
+        variant="fp16",
+    ).to("cuda")
+
+    pipe.load_lora_weights(
+        "whosouravsharma/diffusiondb-sd15-lora",
+        subfolder="checkpoints/checkpoint-4240",
+        weight_name="pytorch_lora_weights.safetensors",
+        adapter_name="diffusiondb",
+    )
+    pipe.set_adapters(["diffusiondb"], adapter_weights=[1.0])  # 0.0 = plain SD 1.5
+
+    image = pipe(
+        "a steampunk owl inside a glass jar, intricate detail",
+        num_inference_steps=25,
+        guidance_scale=7.5,
+        generator=torch.Generator("cuda").manual_seed(42),
+    ).images[0]
+
+    path = OUT / "snippet" / "owl.png"
+    path.parent.mkdir(parents=True, exist_ok=True)
+    image.save(path)
+    del pipe
+    torch.cuda.empty_cache()
+    return {"passed": True, "diffusers": diffusers.__version__, "torch": torch.__version__,
+            "image": str(path.relative_to(OUT))}
+
+
+def stage_checks() -> dict:
+    prompt = "a steampunk owl inside a glass jar, intricate detail"
+    run = runner()
+
+    run.use("base")
+    base = run.render([prompt], [SEED])[0]
+    run.use(FINAL, 0.0)
+    scale_zero = run.render([prompt], [SEED])[0]
+    run.use(FINAL, 1.0)
+    lora_a = run.render([prompt], [SEED])[0]
+    lora_b = run.render([prompt], [SEED])[0]
+    run.use("base")
+    unloaded = run.render([prompt], [SEED])[0]
+
+    folder = OUT / "checks"
+    folder.mkdir(parents=True, exist_ok=True)
+    for name, image in [("base", base), ("scale0", scale_zero), ("lora_a", lora_a),
+                        ("lora_b", lora_b), ("after_unload", unloaded)]:
+        image.save(folder / f"{name}.png")
+
+    results = {}
+    for name, (x, y) in {
+        "scale0_equals_base": (base, scale_zero),
+        "unload_restores_base": (base, unloaded),
+        "same_seed_is_deterministic": (lora_a, lora_b),
+        "lora_changes_output": (base, lora_a),
+    }.items():
+        max_diff, mean_diff = max_pixel_diff(x, y)
+        expected_equal = name != "lora_changes_output"
+        passed = (max_diff <= 1) if expected_equal else (mean_diff > 1.0)
+        results[name] = {"max_pixel_diff": max_diff, "mean_pixel_diff": round(mean_diff, 4),
+                         "passed": passed}
+    results["passed"] = all(r["passed"] for r in results.values() if isinstance(r, dict))
+    return results
+
+
+@torch.no_grad()
+def stage_loss() -> dict:
+    from datasets import load_dataset
+    from diffusers import DDPMScheduler
+    from huggingface_hub import hf_hub_download
+
+    manifest = json.loads(Path(hf_hub_download(
+        DATASET_REPO, "data/manifest.json", repo_type="dataset", revision=LATENTS_REVISION,
+    )).read_text())
+    scaling = manifest["scaling_factor"]
+    shape = tuple(manifest["latent_shape"])
+
+    data = load_dataset(DATASET_REPO, split="validation", revision=LATENTS_REVISION)
+    if LIMIT_LOSS:
+        data = data.select(range(LIMIT_LOSS))
+
+    def decode(blobs):
+        return torch.from_numpy(np.stack(
+            [np.frombuffer(b, dtype=np.float16).reshape(shape) for b in blobs]).copy())
+
+    scheduler = DDPMScheduler.from_pretrained(BASE_MODEL, subfolder="scheduler")
+    run = runner()
+    pipe = run.pipe
+    losses = {}
+
+    for model in ["base"] + list_checkpoints():
+        run.use(model, 1.0)
+        generator = torch.Generator(DEVICE).manual_seed(SEED)  # identical noise per model
+        total, batches = 0.0, 0
+        for start in range(0, len(data), LOSS_BATCH):
+            rows = data[start:start + LOSS_BATCH]
+            mean = decode(rows["latent_mean"]).to(DEVICE, torch.float32)
+            logvar = decode(rows["latent_logvar"]).to(DEVICE, torch.float32).clamp(-30.0, 20.0)
+            eps = torch.randn(mean.shape, device=DEVICE, generator=generator)
+            latents = (mean + torch.exp(0.5 * logvar) * eps) * scaling
+            noise = torch.randn(latents.shape, device=DEVICE, generator=generator)
+            steps = torch.randint(0, scheduler.config.num_train_timesteps, (latents.shape[0],),
+                                  device=DEVICE, generator=generator)
+            noisy = scheduler.add_noise(latents, noise, steps)
+            tokens = pipe.tokenizer(rows["prompt"], padding="max_length", truncation=True,
+                                    max_length=pipe.tokenizer.model_max_length,
+                                    return_tensors="pt").input_ids.to(DEVICE)
+            encoded = pipe.text_encoder(tokens)[0]
+            predicted = pipe.unet(noisy.to(DTYPE), steps, encoded).sample
+            total += F.mse_loss(predicted.float(), noise.float()).item()
+            batches += 1
+        losses[model] = total / max(batches, 1)
+        log(f"  val loss {model}: {losses[model]:.5f}")
+
+    write_json(OUT / "loss" / "val_loss.json", {"examples": len(data), "seed": SEED, "loss": losses})
+    plot_loss(losses, OUT / "loss" / "val_loss.png")
+    return {"examples": len(data), "loss": losses}
+
+
+def plot_loss(losses: dict, path: Path) -> None:
+    import matplotlib
+    matplotlib.use("Agg")
+    import matplotlib.pyplot as plt
+
+    surface, ink, muted, grid, series = "#fcfcfb", "#0b0b0b", "#52514e", "#e6e5e1", "#2a78d6"
+    points = [(int(k.split("-")[1]), v) for k, v in losses.items() if k != "base"]
+    points.sort()
+    steps = [p[0] for p in points]
+    values = [p[1] for p in points]
+
+    fig, ax = plt.subplots(figsize=(8, 4.2), dpi=150)
+    fig.patch.set_facecolor(surface)
+    ax.set_facecolor(surface)
+    ax.plot(steps, values, color=series, linewidth=2, marker="o", markersize=6,
+            markeredgecolor=surface, markeredgewidth=2, zorder=3)
+    if "base" in losses:
+        ax.axhline(losses["base"], color=muted, linewidth=1.5, linestyle=(0, (4, 3)), zorder=2)
+        ax.annotate("SD 1.5 base (no LoRA)", xy=(steps[0] if steps else 0, losses["base"]),
+                    xytext=(0, 6), textcoords="offset points", color=muted, fontsize=9)
+    if points:
+        ax.annotate(f"{values[-1]:.4f}", xy=(steps[-1], values[-1]), xytext=(6, -12),
+                    textcoords="offset points", color=ink, fontsize=9)
+    ax.set_title("Validation loss by checkpoint (1,000 held-out images, fixed noise)",
+                 color=ink, fontsize=11, loc="left")
+    ax.set_xlabel("training step", color=muted, fontsize=9)
+    ax.set_ylabel("MSE (noise prediction)", color=muted, fontsize=9)
+    ax.tick_params(colors=muted, labelsize=8)
+    ax.grid(axis="y", color=grid, linewidth=0.8)
+    for side in ("top", "right", "left"):
+        ax.spines[side].set_visible(False)
+    ax.spines["bottom"].set_color(grid)
+    fig.tight_layout()
+    path.parent.mkdir(parents=True, exist_ok=True)
+    fig.savefig(path, facecolor=surface)
+    plt.close(fig)
+
+
+def inception_metrics(real: torch.Tensor):
+    from torchmetrics.image.fid import FrechetInceptionDistance
+    from torchmetrics.image.kid import KernelInceptionDistance
+
+    n = real.shape[0]
+    subset = max(2, min(1000, n // 2))
+    fid = FrechetInceptionDistance(feature=2048, normalize=True, reset_real_features=False).to(DEVICE)
+    kid = KernelInceptionDistance(subset_size=subset, subsets=100, normalize=True,
+                                  reset_real_features=False).to(DEVICE)
+    for start in range(0, n, 50):
+        chunk = real[start:start + 50].to(DEVICE)
+        fid.update(chunk, real=True)
+        kid.update(chunk, real=True)
+    return fid, kid
+
+
+def score_fake(fid, kid, fake: torch.Tensor) -> dict:
+    fid.reset()
+    kid.reset()
+    for start in range(0, fake.shape[0], 50):
+        chunk = fake[start:start + 50].to(DEVICE)
+        fid.update(chunk, real=False)
+        kid.update(chunk, real=False)
+    kid_mean, kid_std = kid.compute()
+    return {"fid": float(fid.compute()), "kid": float(kid_mean), "kid_std": float(kid_std)}
+
+
+def stage_metrics() -> dict:
+    halves = validation_halves()
+    result: dict = {"n_select": len(halves["select"]), "n_report": len(halves["report"])}
+
+    # --- selection half: pick the checkpoint closest to DiffusionDB (lowest KID)
+    rows = halves["select"]
+    prompts, seeds = [r["prompt"] for r in rows], [r["seed"] for r in rows]
+    fid, kid = inception_metrics(to_tensor([r["image"] for r in rows]))
+    candidates = CANDIDATES[:LIMIT_CKPTS] if LIMIT_CKPTS else CANDIDATES
+    selection = {}
+    for model in candidates:
+        images = render_model(model, prompts, seeds)
+        selection[model] = score_fake(fid, kid, to_tensor(images))
+        selection[model]["clip"] = float(clip_scores(images, prompts).mean())
+        log(f"  select {model}: {selection[model]}")
+    winner = min(selection, key=lambda m: selection[m]["kid"])
+    STATE["selected"] = winner
+    result["selection"] = {"criterion": "lowest KID on the selection half",
+                           "candidates": selection, "selected": winner}
+    del fid, kid
+    torch.cuda.empty_cache()
+
+    # --- report half: the numbers that go in the model card
+    rows = halves["report"]
+    prompts, seeds = [r["prompt"] for r in rows], [r["seed"] for r in rows]
+    real_images = [r["image"] for r in rows]
+    fid, kid = inception_metrics(to_tensor(real_images))
+    report = {"real_images": {"clip": float(clip_scores(real_images, prompts).mean()),
+                              "clip_std": float(clip_scores(real_images, prompts).std())}}
+    renders = {"real": real_images}
+    for model in report_models():
+        images = render_model(model, prompts, seeds)
+        renders[model] = images
+        scores = clip_scores(images, prompts)
+        report[model] = score_fake(fid, kid, to_tensor(images))
+        report[model].update({"clip": float(scores.mean()), "clip_std": float(scores.std())})
+        folder = OUT / "renders" / model
+        folder.mkdir(parents=True, exist_ok=True)
+        for row, image in zip(rows, images):
+            image.save(folder / f"{row['index']:04}.jpg", quality=92)
+        log(f"  report {model}: {report[model]}")
+    STATE["report_rows"] = rows
+    STATE["renders"] = renders
+    result["report"] = report
+    result["notes"] = {
+        "reference": "real validation images from the report half, centre-cropped to 512",
+        "kid_subset_size": max(2, min(1000, len(rows) // 2)),
+        "render": {"steps": STEPS, "guidance": GUIDANCE, "scheduler": "PNDM (pipeline default)",
+                   "seed": "42 + validation row index"},
+        "clip_model": CLIP_MODEL,
+    }
+    write_json(OUT / "metrics" / "metrics.json", result)
+    return result
+
+
+def stage_grids() -> dict:
+    prompts = eval_prompts()
+    if LIMIT_GRID:
+        prompts = prompts[:LIMIT_GRID]
+    seeds = [SEED + i for i in range(len(prompts))]
+    models = report_models()
+    outputs = {}
+    for model in models:
+        images = render_model(model, prompts, seeds)
+        outputs[model] = images
+        folder = OUT / "samples" / model
+        folder.mkdir(parents=True, exist_ok=True)
+        for i, image in enumerate(images):
+            image.save(folder / f"{i:03}.png")
+        contact_sheet(images).save(folder / "grid.jpg", quality=92)
+        write_json(folder / "prompts.json", {"checkpoint": model, "steps": STEPS, "guidance": GUIDANCE,
+                                             "seed": SEED, "prompts": prompts})
+    pairs = OUT / "samples" / "pairs"
+    pairs.mkdir(parents=True, exist_ok=True)
+    for i in range(len(prompts)):
+        labelled_grid([[outputs["base"][i], outputs[selected()][i]]],
+                      ["SD 1.5 base", f"+ LoRA ({selected()})"], thumb=384).save(
+            pairs / f"{i:03}.jpg", quality=92)
+    return {"prompts": len(prompts), "models": models}
+
+
+def stage_sweep() -> dict:
+    prompts = eval_prompts()
+    chosen = [prompts[i] for i in SWEEP_PROMPTS][:1 if SMOKE else None]
+    rows = []
+    for p_index, prompt in zip(SWEEP_PROMPTS, chosen):
+        row = []
+        for scale in SWEEP_SCALES:
+            row += render_model(selected(), [prompt], [SEED + p_index], scale)
+        rows.append(row)
+    path = OUT / "sweep" / "lora_strength.jpg"
+    path.parent.mkdir(parents=True, exist_ok=True)
+    labelled_grid(rows, [f"strength {s}" for s in SWEEP_SCALES], chosen).save(path, quality=92)
+    return {"checkpoint": selected(), "scales": SWEEP_SCALES, "prompts": chosen}
+
+
+def stage_progression() -> dict:
+    prompts = eval_prompts()
+    chosen = [prompts[i] for i in PROGRESSION_PROMPTS][:1 if SMOKE else None]
+    models = ["base"] + list_checkpoints()
+    columns = {m: render_model(m, chosen, [SEED + i for i in PROGRESSION_PROMPTS[:len(chosen)]])
+               for m in models}
+    rows = [[columns[m][r] for m in models] for r in range(len(chosen))]
+    labels = ["base"] + [m.split("-")[1] for m in models[1:]]
+    path = OUT / "progression" / "checkpoints.jpg"
+    path.parent.mkdir(parents=True, exist_ok=True)
+    labelled_grid(rows, labels, chosen, thumb=192).save(path, quality=92)
+    return {"prompts": chosen, "columns": labels}
+
+
+def report_renders() -> dict:
+    """Report-half renders from the metrics stage, or reloaded from disk."""
+    if "renders" in STATE:
+        return STATE["renders"]
+    raise RuntimeError("safety and memorization need the metrics stage in the same run")
+
+
+@torch.no_grad()
+def stage_safety() -> dict:
+    from diffusers.pipelines.stable_diffusion.safety_checker import StableDiffusionSafetyChecker
+    from transformers import CLIPImageProcessor
+
+    checker = StableDiffusionSafetyChecker.from_pretrained(BASE_MODEL, subfolder="safety_checker").to(DEVICE).eval()
+    processor = CLIPImageProcessor.from_pretrained(BASE_MODEL, subfolder="feature_extractor")
+    result = {}
+    for name, images in report_renders().items():
+        flagged = 0
+        for start in range(0, len(images), 32):
+            chunk = images[start:start + 32]
+            pixels = processor(chunk, return_tensors="pt").pixel_values.to(DEVICE)
+            _, has_nsfw = checker(images=np.zeros((len(chunk), 1, 1, 3)), clip_input=pixels)
+            flagged += int(sum(bool(x) for x in has_nsfw))
+        result[name] = {"flagged": flagged, "total": len(images),
+                        "rate": round(flagged / max(len(images), 1), 4)}
+        log(f"  safety {name}: {result[name]}")
+    del checker
+    torch.cuda.empty_cache()
+    write_json(OUT / "safety" / "safety.json", result)
+    return result
+
+
+@torch.no_grad()
+def stage_memorization() -> dict:
+    """Nearest training image for every report-half render, by CLIP cosine.
+
+    CLIP similarity is a proxy for copying, not proof either way. Real
+    validation images (distinct prompts, same style) give the baseline for
+    what "similar" means in this dataset.
+    """
+    from datasets import load_dataset
+
+    renders = report_renders()
+    names = list(renders)
+    queries = {n: clip_image_embeddings(renders[n]) for n in names}
+    best = {n: torch.full((len(renders[n]),), -1.0, device=DEVICE) for n in names}
+    best_thumb = {n: [None] * len(renders[n]) for n in names}
+    best_prompt = {n: [None] * len(renders[n]) for n in names}
+
+    stream = load_dataset(DATASET_REPO, split="train", revision=DATASET_REVISION, streaming=True)
+    batch_images, batch_prompts, seen = [], [], 0
+
+    def flush():
+        nonlocal batch_images, batch_prompts
+        if not batch_images:
+            return
+        emb = clip_image_embeddings(batch_images)
+        for n in names:
+            sims = queries[n] @ emb.T
+            top, arg = sims.max(dim=1)
+            improved = (top > best[n]).nonzero().flatten().tolist()
+            best[n] = torch.maximum(best[n], top)
+            for q in improved:
+                thumb = batch_images[arg[q].item()].resize((256, 256), Image.LANCZOS)
+                buffer = io.BytesIO()
+                thumb.save(buffer, format="JPEG", quality=85)
+                best_thumb[n][q] = buffer.getvalue()
+                best_prompt[n][q] = batch_prompts[arg[q].item()]
+        batch_images, batch_prompts = [], []
+
+    for row in stream:
+        batch_images.append(center_crop(row["image"]))
+        batch_prompts.append(row["prompt"])
+        seen += 1
+        if len(batch_images) == 64:
+            flush()
+        if LIMIT_TRAIN and seen >= LIMIT_TRAIN:
+            break
+        if seen % 2000 == 0:
+            log(f"  memorization: {seen} training images embedded")
+    flush()
+
+    result = {"training_images": seen, "similarity": "CLIP ViT-L/14 image cosine"}
+    for n in names:
+        values = best[n].cpu().numpy()
+        result[n] = {
+            "mean": round(float(values.mean()), 4), "median": round(float(np.median(values)), 4),
+            "p95": round(float(np.percentile(values, 95)), 4), "max": round(float(values.max()), 4),
+            "count_ge_0.90": int((values >= 0.90).sum()), "count_ge_0.95": int((values >= 0.95).sum()),
+        }
+        log(f"  memorization {n}: {result[n]}")
+
+    # The closest render/training pairs for the selected checkpoint, to look at.
+    model = selected()
+    order = np.argsort(-best[model].cpu().numpy())[:8]
+    rows = [[renders[model][q], Image.open(io.BytesIO(best_thumb[model][q]))] for q in order]
+    labels = [f"cos {best[model][q].item():.3f}" for q in order]
+    path = OUT / "memorization" / "closest_pairs.jpg"
+    path.parent.mkdir(parents=True, exist_ok=True)
+    labelled_grid(rows, ["render", "nearest training image"], labels, thumb=256).save(path, quality=90)
+    write_json(OUT / "memorization" / "memorization.json", result)
+    return result
+
+
+STAGE_FUNCTIONS = {
+    "snippet": stage_snippet, "checks": stage_checks, "loss": stage_loss,
+    "metrics": stage_metrics, "grids": stage_grids, "sweep": stage_sweep,
+    "progression": stage_progression, "safety": stage_safety,
+    "memorization": stage_memorization,
+}
+
+
+def main() -> None:
+    if not torch.cuda.is_available():
+        raise SystemExit("No GPU. Run this on a GPU flavor, e.g. a10g-small.")
+    if os.environ.get("HF_TOKEN"):
+        log("NOTE: HF_TOKEN is set. This job needs no token; launch it without --secrets HF_TOKEN.")
+    OUT.mkdir(parents=True, exist_ok=True)
+    log(f"run {RUN_NAME} -> {OUT} | stages: {', '.join(STAGES)} | smoke={SMOKE}")
+    STATUS["config"] = {"base": BASE_MODEL, "model": MODEL_REPO, "dataset": f"{DATASET_REPO}@{DATASET_REVISION}",
+                        "steps": STEPS, "guidance": GUIDANCE, "seed": SEED, "render_batch": RENDER_BATCH,
+                        "candidates": CANDIDATES, "gpu": torch.cuda.get_device_name(0)}
+    save_status()
+
+    for name in STAGES:
+        started = time.time()
+        log(f"=== {name}")
+        try:
+            result = STAGE_FUNCTIONS[name]()
+            STATUS["stages"][name] = {"ok": True, "seconds": round(time.time() - started), "result": result}
+        except Exception as error:
+            traceback.print_exc()
+            STATUS["stages"][name] = {"ok": False, "seconds": round(time.time() - started),
+                                      "error": f"{type(error).__name__}: {error}"}
+        save_status()
+        log(f"=== {name} {'ok' if STATUS['stages'][name]['ok'] else 'FAILED'} "
+            f"({STATUS['stages'][name]['seconds']}s)")
+
+    STATUS["seconds"] = round(time.time() - STATUS["started"])
+    save_status()
+    failed = [n for n, s in STATUS["stages"].items() if not s["ok"]]
+    log(f"done in {STATUS['seconds']}s; failed stages: {failed or 'none'}; results in {OUT}")
+
+
+if __name__ == "__main__":
+    main()