license: mit
tags:
- diffusion
- image-generation
- text-to-image
- representation-alignment
RECON checkpoints
Checkpoints trained with the RECON code (saeidrazavi/RECON).
RECON trains diffusion transformers with REPA-style representation alignment and semantic tokens:
at block encoder_depth the model predicts the features of a frozen vision encoder (the teacher),
pools them into K tokens and appends them to the sequence for the remaining blocks. Sampling pools
the tokens from the model's own prediction, so no teacher is needed at inference.
| File | Task | Model | Data | Training |
|---|---|---|---|---|
c2i/ep-00080.pt |
Class-to-image | SiT-XL/2 | ImageNet 256×256 | 80 epochs |
c2i/ep-00020.pt |
Class-to-image (weak model for autoguidance) | SiT-XL/2 | ImageNet 256×256 | 20 epochs |
c2i/ep-00480.pt |
Class-to-image | SiT-XL/2 | ImageNet 256×256 | 480 epochs |
c2i/ep-00025.pt |
Class-to-image (weak model for autoguidance) | SiT-XL/2 | ImageNet 256×256 | 25 epochs |
t2i-mscoco/0150000.pt |
Text-to-image | MMDiT | MS-COCO 256×256 | 150,000 steps |
Each file is a full training checkpoint with the keys model, ema, opt, args and steps.
The evaluation scripts use the EMA weights and read the model configuration from args, so no
configuration flags are needed. Training can resume from these files.
The 480-epoch model was trained without label dropout, so it has no null class: use autoguidance, not CFG.
Results
Class-to-image (c2i/ep-00080.pt): 50k samples, SDE sampler with 250 steps, no guidance,
--seed 0, ADM reference batch.
| FID | IS | Precision | Recall |
|---|---|---|---|
| 3.2349 | 203.02 | 0.8168 | 0.5632 |
Class-to-image with autoguidance: 50k samples, SDE sampler, --seed 0, ADM reference batch.
"Improved" is projected autoguidance with Semantic Token Extrapolation (--ag-projected --ste-lambda).
| Model | Weak model | Steps | Guidance | w | Interval | λ | FID |
|---|---|---|---|---|---|---|---|
ep-00080.pt |
ep-00020.pt |
50 | Autoguidance | 2 | (0, 0.7) | - | 2.26 |
ep-00080.pt |
ep-00020.pt |
50 | Improved | 2.5 | (0, 0.7) | 0.05 | 2.10 |
ep-00480.pt |
ep-00025.pt |
100 | Autoguidance | 1.4 | (0, 0.85) | - | 1.80 |
ep-00480.pt |
ep-00025.pt |
100 | Improved | 2 | (0, 0.85) | 0.05 | 1.50 |
Text-to-image (t2i-mscoco/0150000.pt): 40,192 MS-COCO val captions, ODE sampler (Heun) with
50 steps, U-ViT FID reference.
| CFG scale | FID | CLIP score |
|---|---|---|
| 1.0 | 8.2671 | 0.2311 |
| 2.0 | 4.5077 | 0.2465 |
Training settings
| Setting | Class-to-image (80 epochs) | Text-to-image |
|---|---|---|
| Model | SiT-XL/2 | MMDiT |
| VAE (latents) | sdvae-ft-mse-f8d4 |
sdvae-ft-ema-f8d4 |
| Text condition | — | CLIP ViT-L/14 features (U-ViT preprocessing) |
| Teacher | dinov2-vit-b |
dinov2-vit-b |
| Batch size / learning rate | 256 / 1e-4 | 256 / 1e-4 |
| Mixed precision | bf16 | fp16 |
| EMA decay | 0.9999 | 0.9999 |
--proj-coeff |
1.5 | 1.5 |
--encoder-depth |
8 | 8 |
| Semantic tokens | K=1, zeropad head, visible for t ≤ 0.5 |
K=1, zeropad head, visible for t ≤ 0.5 |
Caption dropout (--cfg-prob) |
— | 0.1 |
Usage
Set up the code and data as in the code repository's README. Then download the checkpoints
(to download only some of them, give their file names after SaeedRazavi/RECON):
uv run hf download SaeedRazavi/RECON --local-dir checkpoints/recon
Evaluate them:
# Class-to-image
REFERENCE_NPZ=data/imagenet/VIRTUAL_imagenet256_labeled.npz \
uv run accelerate launch --num_processes 4 --mixed_precision bf16 c2i_eval.py \
--ckpt checkpoints/recon/c2i/ep-00080.pt
# Class-to-image, improved autoguidance (80 epochs)
REFERENCE_NPZ=data/imagenet/VIRTUAL_imagenet256_labeled.npz \
uv run accelerate launch --num_processes 4 --mixed_precision bf16 c2i_eval.py \
--ckpt checkpoints/recon/c2i/ep-00080.pt \
--weak-ckpt checkpoints/recon/c2i/ep-00020.pt \
--num-steps 50 --ag-scale 2.5 --ag-projected --ste-lambda 0.05 --guidance-high 0.7
# Class-to-image, improved autoguidance (480 epochs)
REFERENCE_NPZ=data/imagenet/VIRTUAL_imagenet256_labeled.npz \
uv run accelerate launch --num_processes 4 --mixed_precision bf16 c2i_eval.py \
--ckpt checkpoints/recon/c2i/ep-00480.pt \
--weak-ckpt checkpoints/recon/c2i/ep-00025.pt \
--num-steps 100 --ag-scale 2 --ag-projected --ste-lambda 0.05 --guidance-high 0.85
# Text-to-image at CFG 2.0 (reads data/mscoco256 as in training; without --cfg-scale, no guidance)
uv run accelerate launch --num_processes 4 t2i_eval.py --ckpt checkpoints/recon/t2i-mscoco/0150000.pt --cfg-scale 2.0