---
license: mit
pipeline_tag: text-to-image
tags:
- text-to-image
- diffusion
- pytorch
- looped-dit
---

# 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](https://github.com/PeppaKing8/minit2i-jax), conditioned on a frozen
[FLAN-T5-Large](https://huggingface.co/google/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):

```bash
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