OpenSenseNova / Looped-DiT

Looped-DiT: Looped Diffusion Transformer 
Official PyTorch implementation of **Looped-DiT**.
Looped-DiT scales the computation of a text-to-image diffusion transformer by running a
shared group of transformer blocks several times within each denoising step: the model
gets deeper without getting larger. Naive looping does not reliably help, so Looped-DiT
adds two components. **Deep Supervision** decodes the state after every loop through the
post-loop blocks and trains each of these predictions on the same clean-image target.
**Self-Modulating Attention** (exclusive self-attention, XSA, or a head-wise attention
gate) 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 includes:
- **Training** code for Looped-DiT B/32, B/16 and L/16 (pretraining and fine-tuning).
- **Inference** at any loop depth.
- **Evaluation** on GenEval, DPG-Bench, PRISM-Bench, T2I-CoReBench, SpatialGenEval and TIIF-Short.
- **Data preparation** scripts for every training set.
## Table of Contents
- Model Zoo
- Repository Layout
- Installation
- Inference
 - Recommended Inference Settings
- Evaluation
 - Evaluation Metrics
 - Evaluate Checkpoints
- Full Training
 - Configurations
 - Dataset Preparation
 - Pretraining
 - Fine-Tuning
 - Ablation Toggles
- Acknowledgments
## Model Zoo
| Model | Patch | GenEval | DPG | PRISM | CoRe | Spatial | TIIF | Avg | Checkpoint |
| --- | ---: | ---: | ---: | ---: | ---: | ---: | ---: | ---: | --- |
| Looped-DiT B/32 | 32 | 85.1 | 85.3 | 54.4 | 44.5 | 52.3 | 76.1 | 66.3 | sensenova/Looped-DiT-B32 |
| Looped-DiT B/16 | 16 | 87.4 | 87.0 | 67.0 | 53.5 | 54.6 | 79.7 | 71.5 | sensenova/Looped-DiT-B16 |
Scores use the EMA weights, 100 Euler steps, guidance 6.0 and loop depth 4.
Spatial denotes SpatialGenEval. TIIF denotes the short-prompt version of TIIF-Bench.
## Repository Layout
 [code] 
## Installation
Create a Python environment with a CUDA build of PyTorch (2.1 or newer), then install the
dependencies:
 [code] 
Run all commands from the repository root. Configs refer to `${DATA_ROOT}` (training data),
`${ASSET_ROOT}` (benchmark files) and `${OUTPUT_ROOT}` (checkpoints and results in the
evaluation config). They default to `data/`, `eval_assets/` and `outputs/` in the
repository and can be set as environment variables.
Sanity-check the environment with the unit tests (CPU, under a minute):
 [code] 
## Inference
Download a checkpoint from the Model Zoo, then generate. `--loops` sets the
loop depth, and several depths give one row each:
 [code] 
From Python:
 [code] 
### Recommended Inference Settings
These are the settings of the paper's main results and the defaults of the scripts.
| Setting | Value |
| --- | --- |
| Sampler | Euler, 100 steps |
| Classifier-free guidance scale | 6.0 |
| Loop depth | 4, the trained depth (other depths work without retraining) |
| Weights / precision | EMA, bfloat16 |
| Resolution | 512 x 512 |
## Evaluation
##