sensenova/Looped-DiT-B32

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