Dog-Breed Text-to-image Flow Matching from Pretraining to Fine-Tuning on 4 Consumer GPUs

What it claims

A 22.06M-parameter, 12-block LightningDiT trained with rectified flow on 20 dog breeds at 256×256 reaches gFID 8.32 after 200 epochs (~20k steps, 55 min on 4× RTX 3090). A further 100-epoch SFT (~700 steps, 15 min) raises HPSv2 from 0.1625 to 0.2427 and PickScore from 18.42 to 19.61. The authors state they claim no technical novelty. The value is a cheap, complete, ablated recipe, and the paper is framed as course material for advanced undergraduate and graduate students.

How it works

Everything except the DiT is frozen. Image latents and text tokens are precomputed once and cached.

flowchart LR
  C[Caption] --> Q[Qwen3-0.6B ❄]
  I[Image] --> V[E2E-INVAE ❄]
  I --> E[EUPE ViT-B/16 ❄]
  Q --> S[Text + time + noisy latent<br/>one sequence]
  V -->|+ noise| S
  S --> D[LightningDiT 12 blocks]
  D --> FM[Flow-matching loss]
  D -->|Block 4 features| A[iREPA loss, λ=0.5]
  E --> A

• Conditioning: Qwen3-0.6B text tokens, the timestep token and the noisy latents are concatenated into one sequence (in-context conditioning). There is no cross-attention.
• Alignment: iREPA aligns Block 4 hidden states to EUPE features. It differs from REPA in two ways:
• It uses a 3×3 conv projector over the 16×16 token grid instead of a linear layer.
• It normalises the target features per image across patches (mean scaled by γ=0.7, ε=1e-6).
• SFT stage: the REPA loss is turned off and only flow matching remains.
• Data:
• Pretraining uses 26k recaptioned ImageNet images, each with a short (≤64 tokens) and a long (≤128 tokens) caption.
• SFT uses 2k FLUX.1-schnell images generated from prompts written by GPT-5.5.
• Evaluation uses 500 held-out prompts.

Evidence

All gFID numbers below use E2E-INVAE and 128-token captions unless the row varies that setting.

The best encoder depends on the alignment method. DINOv2 is best for REPA, while EUPE is best for iREPA.

How to use it

• Code, models and data: all are stated to be open-sourced. The code is on GitHub (built on DiffusionBench), and the three datasets are on HuggingFace. The page gives no commands, repo URL or licence.
• Single-GPU cost: on one RTX 3090, pretraining takes about 3 h and SFT about 21 min.
• The default config (micro-batch 64, 4 gradient-accumulation steps, bf16) needs at least 10 GiB VRAM, or 15 GiB with online FID.
• Feature caching: a one-off step of about 15 min on one 3090.
• Online evaluation: this doubles pretraining time and adds about 1 h to SFT.
• Swappable components:
• Tokenizer: E2E-INVAE or E2E-VAVAE
• Alignment method: REPA or iREPA
• Alignment target encoder: DINOv2, DINOv3, PE-Spatial or EUPE

Limits

• gFID is not held-out. It compares 26k generations from the training prompts against the 26k training images. It measures fit to the training distribution, not generalisation.
• SFT makes FID worse. The authors attribute this to the synthetic FLUX data pulling the model off the pretraining distribution. They do not report the size of the increase.
• Preference gains are judged by models only. The SFT improvements are measured solely by reward models (HPSv2, PickScore), and the sample figure is explicitly cherry-picked.
• The setting is narrow. Results cover a single domain (20 dog breeds) at 256×256 only. Nothing is shown about scaling beyond that.
• Some ablation margins are small. For example, 8.32 vs 8.37 vs 8.46 come from single runs, and the page reports no variance.
• Unverified details: no licence is stated, and the fetched text is truncated partway through the appendices. The full hyperparameters (Appendix C) were not visible.