Looped-DiT B/32
Looped-DiT is a text-to-image diffusion transformer that runs a shared group of transformer blocks several times within each denoising step, so the model gets deeper without getting larger. Deep supervision trains the prediction after every loop, and exclusive self-attention (XSA) regulates the attention updates inside the loop. The backbone is the pixel-space MMDiT of MiniT2I, conditioned on a frozen FLAN-T5-Large.
This repository holds the EMA weights of Looped-DiT B/32 at step 290k:
- patch size 32, 512 x 512 images
- [6,5,6] blocks (pre-loop, looped, post-loop), with the 5 looped blocks run 4 times
- trained with deep supervision, with XSA in the looped blocks
Results
100 Euler steps, classifier-free guidance 6.0, loop depth 4. TIIF-Short is TIIF-Bench scored on its short prompts.
| GenEval | DPG-Bench | PRISM | CoReBench | SpatialGenEval | TIIF-Short | Avg |
|---|---|---|---|---|---|---|
| 85.1 | 85.3 | 54.4 | 44.5 | 52.3 | 76.1 | 66.3 |
Usage
Download the checkpoint, then sample with the Looped-DiT code (link coming soon):
hf download sensenova/Looped-DiT-B32 looped-dit-b32.pt --local-dir checkpoints
python -m looped_dit.sample --checkpoint checkpoints/looped-dit-b32.pt \
--prompt "a red cube on top of a blue sphere" --out sample.png
Add --loops 1 2 3 4 to render the prompt at several loop depths. Other depths work
without retraining.
License
MIT