Update README.md

README.md

@@ -13,11 +13,32 @@ datasets:
 pipeline_tag: unconditional-image-generation
 ---
 
-# Dancing Stick Figures — baseline checkpoints (v0.1)
+# Dancing Stick Figures — reference baselines and checkpoints
 
-Reference models trained on [sprited/dancing-stick-figures](https://huggingface.co/datasets/sprited/dancing-stick-figures)
-with the code in <https://github.com/sprited-ai/dancing-stick-figures>. All unconditional, RGBA, EMA weights included,
-optimizer state stripped. Load with the trainers' `--init` (warm-start) or `--resume`, or with `eval/score_images.py`.
+**Companion baselines for [Dancing Stick Figures](https://huggingface.co/datasets/sprited/dancing-stick-figures), a small teaching dataset for learning video generation from scratch.** The dataset is the main project; this repository contains reference checkpoints, sample outputs, and evaluation results.
+
+<p align="center">
+  <img src="unet_ar64_rollout.gif" alt="Eight generated stick-figure dances rolled out for 5.6 seconds" width="640">
+</p>
+
+<p align="center"><em>Eight unconditional samples from the 64×64 autoregressive reference baseline, rolled out for 5.6 seconds.</em></p>
+
+**Start here:** [Dataset](https://huggingface.co/datasets/sprited/dancing-stick-figures) · [Colab](https://colab.research.google.com/github/sprited-ai/dancing-stick-figures/blob/main/notebooks/dancing_stick_figures_colab.ipynb) · [Code](https://github.com/sprited-ai/dancing-stick-figures)
+
+## What this repository is for
+
+These checkpoints are reference points, not the dataset itself. They show that the data can support a complete learning path: train a small image generator, warm-start an 8-frame video model, roll it out autoregressively, and compare it with real clips.
+
+- **Dataset:** 1,430 six-second clips / 514,800 frames
+- **Small configuration:** 64×64, 0.85 GB
+- **Labels:** exact 2D + 3D pose for 27 joints, depth, normals, part segmentation, camera parameters, and raw motion
+- **Licenses:** CC0 data and MIT code
+- **Reference architectures:** factorised UNet and DiT flow matching
+- **Reproducible path:** download → cache → train → roll out → evaluate
+
+All checkpoints are RGBA and include EMA weights with optimizer state removed. Use a checkpoint with the trainers' `--init` or `--resume`, or evaluate it with `eval/score_images.py`.
+
+## Image checkpoints
 
 | file | arch | res | steps | recipe | oracle tvr / lie / clean (floor) |
 |---|---|---|---|---|---|
@@ -28,7 +49,9 @@ optimizer state stripped. Load with the trainers' `--init` (warm-start) or `--re
 | `unet_img128.pt` | UNet, min-SNR-5 | 128² | 20k | batch 128 | .226 / .073 / .22 (.203 / .047 / .23) |
 | `dit_img128_p4.pt` | DiT-FM patch 4 | 128² | 40k | batch 128 | .251 / .065 / .23 (.209 / .048 / .21) |
 
-**Video (64² × 8 frames @ 20 fps, UNet 46 M, batch 16, unconditional) — finished runs.** Oracle on 64 samples × 2 seeds vs 64 real
+## Short-video checkpoints
+
+**64² × 8 frames @ 20 fps, UNet 46 M, batch 16, unconditional.** Oracle on 64 samples × 2 seeds vs 64 real
 8-frame windows (floor); FVD with a frame-repeated I3D (real-vs-real ≈ 110–120 in this setup):
 
 | file | init | steps | tvr↓ | lie↓ | head-jitter↓ | angle-jerk↓ | FVD↓ |
@@ -39,18 +62,26 @@ optimizer state stripped. Load with the trainers' `--init` (warm-start) or `--re
 
 Read: warm-starting from the image model gets there **~2.5× sooner** (loss at 4k ≈ scratch at 10k) but ends in the same place — a
 convergence-speed win, not a quality win. Per-frame anatomy is within ~0.03 of real frames; temporal jitter ~1.2× real; FVD ~80–100
-above the real-vs-real floor. Sample GIFs: `unet_vid64_scratch_samples.gif`, `unet_vid64_from_img_samples.gif`.
+above the real-vs-real floor. Sample outputs:
+
+<table><tr><td align="center"><img src="unet_vid64_scratch_samples.gif" width="360"><br><em>trained from scratch</em></td><td align="center"><img src="unet_vid64_from_img_samples.gif" width="360"><br><em>warm-started from the image checkpoint</em></td></tr></table>
 
-**Autoregressive video (the route's step 4) — `unet_ar64.pt`**: same UNet, `--ar_ctx 8 --frames 8 --stride 2` (8 context + 8 new
+## Long autoregressive rollout
+
+**`unet_ar64.pt`**: same UNet, `--ar_ctx 8 --frames 8 --stride 2` (8 context + 8 new
 frames per chunk, 10 fps), initialised from `unet_img64.pt`, 60k steps on one H100 (batch 16, ~5.5 h). Rolls out to any length:
 `python scripts/rollout.py --ckpt unet_ar64.pt --seconds 5`. Oracle on 32 five-second rollouts vs 32 real 5 s clips —
 tvr .149 / .128, lie .121 / .103, head-jitter .69 / .57, angle-jerk .19 / .14 (model / real). `unet_ar64_rollout.gif` = 8 samples × 5.6 s.
 `unet_ar64_interim.pt` = the same run at 30k.
 
+## Additional checkpoints
+
 DiT track (interim, stopped at 44k/42k of 61k): `dit_vid64_from_img_interim.pt` (init from `dit_img64_p2.pt`, patch 2, shift 2, 10 % image
 batches, 20 % I2V) and `dit_vid64_scratch_interim.pt`; class-conditional image models: `unet_img64_cond.pt`, `dit_img64_p2_cond.pt` (5 prompt groups, CFG).
 
-**Compare your model against these in one command:**
+## Compare your run
+
+Compare your model against these references in one command:
 ```bash
 python scripts/compare.py --ckpt runs/img64/ckpt.pt --cache data/cache --ref unet_img64     # image models
 python -m eval.run_ckpt --run runs/vid64 --cache data/cache --n 64                          # video: temporal oracle + FVD