reyhanehesi/minidog-checkpoints

MiniDog checkpoints

The two checkpoints of MiniDog's best recipe: a 22M-parameter LightningDiT trained from scratch for dog-breed text-to-image generation on four RTX 3090s. The pretraining checkpoint also stores the 2.7M-parameter iREPA projector, which is only used during pretraining. Code: Reyhaneesmailizadeh/minidog. Data: reyhanehesi/dog-t2i-diffusion-data.

File	Stage	Config	Result
irepa-eupe-mse-ep200.pt	pretraining, 200 epochs (20.2k steps), 55 min	configs/pretrain_e2e-invae_128tok_mse_irepa_eupe.yaml	FID 8.32
sft-eupe-norepa-ep100.pt	fine-tuning, 100 epochs (700 steps), 15 min	configs/sft_e2e-invae_128tok_mse_norepa.yaml	HPSv2 0.2427, PickScore 19.61

FID is measured against the 26k pretraining images. The preference scores are averages over 500 held-out prompts; the pretrained checkpoint scores 0.1625 / 18.42 on the same prompts, so fine-tuning moves HPSv2 by +0.0802 and PickScore by +1.19.

The recipe

Both checkpoints share the architecture: 12 layers, 384 wide, 6 heads, patch size 1, on 16x16x32 latents from a frozen E2E-INVAE tokenizer, conditioned on Qwen3-0.6B caption embeddings (128 tokens), velocity-prediction rectified flow.

Pretraining aligns DiT block 4 to EUPE ViT-B/16 patch features with iREPA: a 3x3 convolutional projector and spatial normalisation (gamma 0.7), squared-error distance, lambda 0.5. This beat plain REPA with DINOv2 (FID 8.80) and every other target representation tested.

Fine-tuning drops the alignment term entirely and trains on the flow-matching loss alone, for 100 epochs on 2,000 FLUX.1-schnell images of the same 20 breeds.

Contents

Each file is a torch.save dict with model, ema, optimizer, scheduler, epoch, step. Evaluation uses the ema weights. The optimizer and scheduler state are included so a run can be resumed exactly; warm-starting a new fine-tune with --init-weights-only ignores them.

Reproducing

git clone -b irepa-eupe https://github.com/Reyhaneesmailizadeh/minidog && cd minidog && uv sync
export DATA=data/dog-t2i-diffusion-data
uv run hf download reyhanehesi/dog-t2i-diffusion-data --local-dir $DATA --repo-type dataset
for NAME in dogs_recaptioned_wds dogs_synthetic_2k_wds; do
  mkdir -p $DATA/$NAME && tar -xzf $DATA/$NAME.tar.gz -C $DATA/$NAME
done
uv run hf download reyhanehesi/minidog-checkpoints --local-dir ckpts/released

Generate from both checkpoints on the same prompts and the same initial noise, then score:

uv run python -m minidog.generate --config configs/pretrain_e2e-invae_128tok_mse_irepa_eupe.yaml \
    --checkpoint ckpts/released/irepa-eupe-mse-ep200.pt \
    --captions-json $DATA/captions_500.json \
    --output-dir results/samples/pretrain \
    --noise-file results/samples/shared_noise_500.pt --group-by-breed

uv run python -m minidog.generate --config configs/sft_e2e-invae_128tok_mse_norepa.yaml \
    --checkpoint ckpts/released/sft-eupe-norepa-ep100.pt \
    --captions-json $DATA/captions_500.json \
    --output-dir results/samples/sft \
    --noise-file results/samples/shared_noise_500.pt --group-by-breed

uv run python -m minidog.score_absolute --dirs results/samples/pretrain results/samples/sft

--noise-file draws the 500 noise tensors once and reuses them, so each pair of images differs only by the checkpoint. Expect 0.1625 -> 0.2427 (HPSv2) and 18.42 -> 19.61 (PickScore).

To train them yourself rather than download, see minidog/README.md: precompute the EUPE latents, then run the two configs above. Pretraining takes 55 minutes on 4x3090, fine-tuning 15 more, not counting the FID evaluation during training (set eval.eval_interval: 0 to turn it off).

Caveats

• FID is measured against the 26k pretraining images, so it reports fit to the training distribution rather than generalisation.
• HPSv2 and PickScore are models trained to predict human preference, not human judgments. A higher score is evidence, not proof, that people prefer the images.
• The model is small, trained on 20 breeds at 256x256. The conclusions may not hold at larger scale.