| --- |
| license: mit |
| tags: |
| - diffusion |
| - image-generation |
| - text-to-image |
| - representation-alignment |
| --- |
| |
| # RECON checkpoints |
|
|
| Checkpoints trained with the RECON code ([saeidrazavi/RECON](https://github.com/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`): |
|
|
| ```bash |
| uv run hf download SaeedRazavi/RECON --local-dir checkpoints/recon |
| ``` |
|
|
| Evaluate them: |
|
|
| ```bash |
| # 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 |
| ``` |
|
|
