Looped-DiT B/16
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/16 at step 580k:
- patch size 16, 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 |
|---|---|---|---|---|---|---|
| 87.4 | 87.0 | 67.0 | 53.5 | 54.6 | 79.7 | 71.5 |
Usage
Download the checkpoint, then sample with the Looped-DiT code (link coming soon):
hf download sensenova/Looped-DiT-B16 looped-dit-b16.pt --local-dir checkpoints
python -m looped_dit.sample --checkpoint checkpoints/looped-dit-b16.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