Create model card and add pipeline tag
#1
by nielsr HF Staff - opened
README.md
ADDED
|
@@ -0,0 +1,72 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
---
|
| 2 |
+
pipeline_tag: unconditional-image-generation
|
| 3 |
+
---
|
| 4 |
+
|
| 5 |
+
# Registers Matter for Pixel-Space Diffusion Transformers
|
| 6 |
+
|
| 7 |
+
This repository contains the checkpoints for the paper [Registers Matter for Pixel-Space Diffusion Transformers](https://huggingface.co/papers/2605.16147).
|
| 8 |
+
|
| 9 |
+
* **Project Page:** [https://quickjkee.github.io/registers-project-page/](https://quickjkee.github.io/registers-project-page/)
|
| 10 |
+
* **GitHub Repository:** [https://github.com/quickjkee/register-guidance](https://github.com/quickjkee/register-guidance)
|
| 11 |
+
* **Paper:** [https://arxiv.org/abs/2605.16147](https://arxiv.org/abs/2605.16147)
|
| 12 |
+
|
| 13 |
+
<p align="center">
|
| 14 |
+
<img src="https://raw.githubusercontent.com/quickjkee/register-guidance/main/demo/grid.png" width="100%">
|
| 15 |
+
</p>
|
| 16 |
+
|
| 17 |
+
## Introduction
|
| 18 |
+
|
| 19 |
+
Diffusion Transformers (DiTs) trained in pixel-space benefit significantly from register tokens. In this work, the authors show that register tokens produce cleaner feature maps at high noise levels, which improves visual coherence. To amplify this effect, they propose **Register Guidance**, a technique that amplifies the contribution of register tokens responsible for improving visual structure and coherence in pixel-space generation.
|
| 20 |
+
|
| 21 |
+
## Available Models
|
| 22 |
+
|
| 23 |
+
The paper provides two types of pretrained models:
|
| 24 |
+
1. Trained from scratch on ImageNet.
|
| 25 |
+
2. Fine-tuned from the pretrained [JiT checkpoints](https://github.com/LTH14/JiT) with register tokens.
|
| 26 |
+
|
| 27 |
+
| Setting | Model | FID, RG | FID, CFG+RG |
|
| 28 |
+
|:---|:---|---:|---:|
|
| 29 |
+
| Training from scratch | JiT-B/16 | 3.46 | 3.32 |
|
| 30 |
+
| Training from scratch | JiT-L/16 | 2.18 | 2.16 |
|
| 31 |
+
| Training from scratch | JiT-B/32 | 4.01 | 3.69 |
|
| 32 |
+
| Training from scratch | JiT-L/32 | 2.47 | 2.41 |
|
| 33 |
+
| Fine-tuning | JiT-B/16 | 3.25 | 2.96 |
|
| 34 |
+
| Fine-tuning | JiT-L/16 | 2.00 | 2.13 |
|
| 35 |
+
| Fine-tuning | JiT-H/16 | 1.85 | 1.80 |
|
| 36 |
+
| Fine-tuning | JiT-B/32 | 3.62 | 3.34 |
|
| 37 |
+
| Fine-tuning | JiT-L/32 | 2.24 | 2.23 |
|
| 38 |
+
| Fine-tuning | JiT-H/32, RG | 2.01 | -- |
|
| 39 |
+
| Fine-tuning | JiT-H/32, CFG | -- | 1.87 |
|
| 40 |
+
|
| 41 |
+
## Evaluation
|
| 42 |
+
|
| 43 |
+
To evaluate a pretrained JiT-B/16 (or JiT-B/32) with Register Guidance, clone the official [GitHub repository](https://github.com/quickjkee/register-guidance) and run:
|
| 44 |
+
|
| 45 |
+
```bash
|
| 46 |
+
torchrun --nproc_per_node=8 --nnodes=1 --node_rank=0 \
|
| 47 |
+
main_jit.py \
|
| 48 |
+
--model JiT-B/16 \
|
| 49 |
+
--img_size 256 --noise_scale 1.0 \
|
| 50 |
+
--gen_bsz 256 --num_images 50000 \
|
| 51 |
+
--guidance_type rg \
|
| 52 |
+
--rg 3.5 \
|
| 53 |
+
--interval_min_rg 0.03 --interval_max_rg 0.9 \
|
| 54 |
+
--output_dir output \
|
| 55 |
+
--resume ${CKPT_PATH} \
|
| 56 |
+
--in_context_start 4 \
|
| 57 |
+
--in_context_len 32 \
|
| 58 |
+
--evaluate_gen
|
| 59 |
+
```
|
| 60 |
+
|
| 61 |
+
For more details on evaluation configs and training scripts, please refer to the [GitHub repository](https://github.com/quickjkee/register-guidance).
|
| 62 |
+
|
| 63 |
+
## Citation
|
| 64 |
+
|
| 65 |
+
```bibtex
|
| 66 |
+
@article{starodubcev2026registers,
|
| 67 |
+
title={Registers Matter for Pixel-Space Diffusion Transformers},
|
| 68 |
+
author={Starodubcev, Nikita and Sudakov, Ilia and Drobyshevskiy, Ilya and Babenko, Artem and Baranchuk, Dmitry},
|
| 69 |
+
journal={arXiv preprint arXiv:2605.16147},
|
| 70 |
+
year={2026}
|
| 71 |
+
}
|
| 72 |
+
```
|