|
1 | 1 | # Infinity $\infty$: Scaling Bitwise AutoRegressive Modeling for High-Resolution Image Synthesis |
2 | 2 |
|
3 | 3 | <div align="center"> |
4 | | - |
| 4 | + |
| 5 | +[](https://opensource.bytedance.com/gmpt/t2i/invite) |
5 | 6 | [](https://foundationvision.github.io/infinity.project/) |
6 | 7 | [](https://arxiv.org/abs/2412.04431) |
7 | | - |
| 8 | +[](https://huggingface.co/FoundationVision/infinity) |
8 | 9 |
|
9 | 10 | </div> |
10 | 11 | <p align="center" style="font-size: larger;"> |
|
17 | 18 | <p> |
18 | 19 |
|
19 | 20 | ## 🔥 Updates!! |
| 21 | +* Dec 24, 2024: 🔥 Training and Testing Codes && Checkpoints && Demo released! |
20 | 22 | * Dec 12, 2024: 💻 Add Project Page |
21 | 23 | * Dec 5, 2024: 🤗 Paper release |
22 | 24 |
|
| 25 | +## 🕹️ Try and Play with Infinity! |
| 26 | + |
| 27 | +We provide a [demo website](https://opensource.bytedance.com/gmpt/t2i/invite) for you to play with Infinity and generate images interactively. Enjoy the fun of bitwise autoregressive modeling! |
| 28 | + |
| 29 | +We also provide [interactive_infer.ipynb](tools/interactive_infer.ipynb) for you to see more technical details about Infinity. |
| 30 | + |
23 | 31 | ## 📑 Open-Source Plan |
| 32 | + - [ ] Infinity-20B Checkpoints |
| 33 | + - [x] Training Code |
| 34 | + - [x] Web Demo |
| 35 | + - [x] Inference Code |
| 36 | + - [x] Infinity-2B Checkpoints |
| 37 | + - [x] VAE Checkpoints |
24 | 38 |
|
25 | | -- Infinity-2B (Text-to-Image Model) |
26 | | - - [ ] Web Demo |
27 | | - - [ ] Inference |
28 | | - - [ ] Checkpoints |
29 | 39 |
|
30 | 40 |
|
31 | 41 | ## 📖 Introduction |
32 | | -We present Infinity, a Bitwise Visual AutoRegressive Modeling capable of generating high-resolution, photorealistic images following language instruction. Infinity refactors visual autoregressive model under a bitwise token prediction framework with an infinite-vocabulary classifier and bitwise self-correction mechanism. By theoretically expanding the tokenizer vocabulary size to infinity in Transformer, our method significantly unleashes powerful scaling capabilities to infinity compared to vanilla VAR. Extensive experiments indicate Infinity outperforms AutoRegressive Text-to-Image models by large margins, matches or exceeds leading diffusion models. Without extra optimization, Infinity generates a 1024 $\times$ 1024 image in 0.8s, 2.6 $\times$ faster than SD3-Medium, making it the fastest Text-to-Image model. Models and codes are released to promote further exploration of Infinity for visual generation. |
| 42 | +We present Infinity, a Bitwise Visual AutoRegressive Modeling capable of generating high-resolution and photorealistic images. Infinity redefines visual autoregressive model under a bitwise token prediction framework with an infinite-vocabulary tokenizer & classifier and bitwise self-correction. Theoretically scaling the tokenizer vocabulary size to infinity and concurrently scaling the transformer size, our method significantly unleashes powerful scaling capabilities. Infinity sets a new record for autoregressive text-to-image models, outperforming top-tier diffusion models like SD3-Medium and SDXL. Notably, Infinity surpasses SD3-Medium by improving the GenEval benchmark score from 0.62 to 0.73 and the ImageReward benchmark score from 0.87 to 0.96, achieving a win rate of 66%. Without extra optimization, Infinity generates a high-quality 1024×1024 image in 0.8 seconds, making it 2.6× faster than SD3-Medium and establishing it as the fastest text-to-image model. |
| 43 | + |
| 44 | +### 🔥 Redefines VAR under a bitwise token prediction framework 🚀: |
| 45 | + |
| 46 | +<p align="center"> |
| 47 | +<img src="assets/framework_row.png" width=95%> |
| 48 | +<p> |
| 49 | + |
| 50 | +Infinite-Vocabulary Tokenizer✨: We proposes a new bitwise multi-scale residual quantizer, which significantly reduces memory usage, enabling the training of extremely large vocabulary, e.g. $V_d = 2^{32}$ or $V_d = 2^{64}$. |
| 51 | + |
| 52 | +Infinite-Vocabulary Classifier✨: Conventional classifier predicts $2^d$ indices. IVC predicts $d$ bits instead. Slight perturbations to near-zero values in continuous features cause a complete change of indices labels. Bit labels change subtly and still provide steady supervision. Besides, if d = 32 and h = 2048, a conventional classifier requires 8.8T parameters. IVC only requires 0.13M. |
| 53 | + |
| 54 | +Bitwise Self-Correction✨: Teacher-forcing training in AR brings severe train-test discrepancy. It lets the transformer only refine features without recognizing and correcting mistakes. Mistakes will be propagated and amplified, finally messing up generated images. We propose Bitwise Self-Correction (BSC) to mitigate the train-test discrepancy. |
| 55 | + |
| 56 | +### 🔥 Scaling Vocabulary benefits Reconstruction and Generation 📈: |
| 57 | + |
| 58 | +<p align="center"> |
| 59 | +<img src="assets/scaling_vocabulary.png" width=95%> |
| 60 | +<p> |
| 61 | + |
| 62 | +### 🔥 Discovering Scaling Laws in Infinity transformers 📈: |
| 63 | + |
| 64 | +<p align="center"> |
| 65 | +<img src="assets/scaling_models.png" width=95%> |
| 66 | +<p> |
| 67 | + |
| 68 | +## Infinity Model ZOO |
| 69 | +We provide Infinity models for you to play with, which are on <a href='https://huggingface.co/FoundationVision/infinity'><img src='https://img.shields.io/badge/%F0%9F%A4%97%20weights-FoundationVision/Infinity-yellow'></a> or can be downloaded from the following links: |
| 70 | + |
| 71 | +### Visual Tokenizer |
| 72 | + |
| 73 | +| vocabulary | stride | IN-256 rFID $\downarrow$ | IN-256 PSNR $\uparrow$ | IN-512 rFID $\downarrow$ | IN-512 PSNR $\uparrow$ | HF weights🤗 | |
| 74 | +|:----------:|:-----:|:--------:|:---------:|:-------:|:-------:|:------------------------------------------------------------------------------------| |
| 75 | +| $V_d=2^{16}$ | 16 | 1.22 | 20.9 | 0.31 | 22.6 | [infinity_vae_d16.pth](https://huggingface.co/FoundationVision/infinity/blob/main/infinity_vae_d16.pth) | |
| 76 | +| $V_d=2^{24}$ | 16 | 0.75 | 22.0 | 0.30 | 23.5 | [infinity_vae_d24.pth](https://huggingface.co/FoundationVision/infinity/blob/main/infinity_vae_d24.pth) | |
| 77 | +| $V_d=2^{32}$ | 16 | 0.61 | 22.7 | 0.23 | 24.4 | [infinity_vae_d32.pth](https://huggingface.co/FoundationVision/infinity/blob/main/infinity_vae_d32.pth) | |
| 78 | +| $V_d=2^{64}$ | 16 | 0.33 | 24.9 | 0.15 | 26.4 | [infinity_vae_d64.pth](https://huggingface.co/FoundationVision/infinity/blob/main/infinity_vae_d64.pth) | |
| 79 | +| $V_d=2^{32}$ | 16 | 0.75 | 21.9 | 0.32 | 23.6 | [infinity_vae_d32_reg.pth](https://huggingface.co/FoundationVision/infinity/blob/main/infinity_vae_d32_reg.pth) | |
| 80 | + |
| 81 | +### Infinity |
| 82 | +| model | Resolution | GenEval | DPG | HPSv2.1 | HF weights🤗 | |
| 83 | +|:----------:|:-----:|:--------:|:---------:|:-------:|:------------------------------------------------------------------------------------| |
| 84 | +| Infinity-2B | 1024 | 0.69 / 0.73 $^{\dagger}$ | 83.5 | 32.2 | [infinity_2B.pth](https://huggingface.co/FoundationVision/var/resolve/main/infinity_2b_reg.pth) | |
| 85 | +| Infinity-20B | 1024 | - | - | - | [Coming Soon](TBD) | |
| 86 | + |
| 87 | +${\dagger}$ result is tested with a [prompt rewriter](tools/prompt_rewriter.py). |
| 88 | + |
| 89 | +You can load these models to generate images via the codes in [interactive_infer.ipynb](tools/interactive_infer.ipynb). Note: you need to download [infinity_vae_d32reg.pth](https://huggingface.co/FoundationVision/var/resolve/main/var_d30.pth) and [flan-t5-xl](https://huggingface.co/google/flan-t5-xl) first. |
| 90 | + |
| 91 | + |
| 92 | +## Installation |
| 93 | +1. We use FlexAttention to speedup training, which requires `torch>=2.5.1`. |
| 94 | +2. Install other pip packages via `pip3 install -r requirements.txt`. |
| 95 | + |
| 96 | +## Data Preparation |
| 97 | +The structure of the training dataset is listed as bellow. The training dataset contains a list of json files with name "[h_div_w_template1]_[num_examples].jsonl". Here [h_div_w_template] is a float number, which is the template ratio of height to width of the image. [num_examples] is the number of examples where $h/w$ is around h_div_w_template. [dataset_t2i_iterable.py](infinity/dataset/dataset_t2i_iterable.py) supports traing with >100M examples. But we have to specify the number of examples for each h/w template ratio in the filename. |
| 98 | + |
| 99 | + ``` |
| 100 | + /path/to/dataset/: |
| 101 | + [h_div_w_template1]_[num_examples].jsonl |
| 102 | + [h_div_w_template2]_[num_examples].jsonl |
| 103 | + [h_div_w_template3]_[num_examples].jsonl |
| 104 | + ``` |
| 105 | + |
| 106 | +Each "[h_div_w_template1]_[num_examples].jsonl" file contains lines of dumped json item. Each json item contains the following information: |
| 107 | + ``` |
| 108 | + { |
| 109 | + "image_path": "path/to/image, required", |
| 110 | + "h_div_w": "float value of h_div_w for the image, required", |
| 111 | + "long_caption": long_caption of the image, required", |
| 112 | + "long_caption_type": "InternVL 2.0, required", |
| 113 | + "short_caption": "short of the image, optional", |
| 114 | + "short_caption_type": "user prompt, , optional" |
| 115 | + } |
| 116 | + ``` |
| 117 | + |
| 118 | + Still have questions about the data preparation? Easy, we have provided a toy dataset with 10 images. You can prepare your dataset by referring [this](data/infinity_toy_data). |
| 119 | + |
| 120 | + |
| 121 | +## Training Scripts |
| 122 | +We provide [train.sh](scripts/train.sh) for train Infinity-2B with one command |
| 123 | +```shell |
| 124 | +bash scripts/train.sh |
| 125 | +``` |
| 126 | + |
| 127 | +To train Infinity with different model sizes {125M, 1B, 2B} and different {256/512/1024} resolutions, you can run the following command: |
| 128 | +```shell |
| 129 | +# 125M, layer12, pixel number = 256 x 256 = 0.06M Pixels |
| 130 | +torchrun --nproc_per_node=8 --nnodes=... --node_rank=... --master_addr=... --master_port=... train.py \ |
| 131 | + --model=layer12c4 --pn 0.06M --exp_name=infinity_125M_pn_0.06M \ |
| 132 | +# 1B, layer24, pixel number = 256 x 256 = 0.06M Pixels |
| 133 | +torchrun --nproc_per_node=8 --nnodes=... --node_rank=... --master_addr=... --master_port=... train.py \ |
| 134 | + --model=layer24c4 --pn 0.06M --exp_name=infinity_1B_pn_0.06M \ |
| 135 | +# 2B, layer32, pixel number = 256 x 256 = 0.06M Pixels |
| 136 | +torchrun --nproc_per_node=8 --nnodes=... --node_rank=... --master_addr=... --master_port=... train.py \ |
| 137 | + --model=2bc8 --pn 0.06M --exp_name=infinity_2B_pn_0.06M \ |
| 138 | +# 2B, layer32, pixel number = 512 x 512 = 0.25M Pixels |
| 139 | +torchrun --nproc_per_node=8 --nnodes=... --node_rank=... --master_addr=... --master_port=... train.py \ |
| 140 | + --model=2bc8 --pn 0.25M --exp_name=infinity_2B_pn_0.25M \ |
| 141 | +# 2B, layer32, pixel number = 1024 x 1024 = 1M Pixels |
| 142 | +torchrun --nproc_per_node=8 --nnodes=... --node_rank=... --master_addr=... --master_port=... train.py \ |
| 143 | + --model=2bc8 --pn 1M --exp_name=infinity_2B_pn_1M \ |
| 144 | +``` |
| 145 | +A folder named `local_output` will be created to save the checkpoints and logs. |
| 146 | +You can monitor the training process by checking the logs in `local_output/log.txt` and `local_output/stdout.txt`. We highly recommend you use [wandb](https://wandb.ai/site/) for detailed logging. |
| 147 | + |
| 148 | +If your experiment is interrupted, just rerun the command, and the training will **automatically resume** from the last checkpoint in `local_output/ckpt*.pth`. |
| 149 | + |
| 150 | +## Evaluation |
| 151 | +We provide [eval.sh](scripts/eval.sh) for evaluation on various benchmarks with only one command. In particular, [eval.sh](scripts/eval.sh) supports evaluation on commonly used metrics such as [GenEval](https://github.com/djghosh13/geneval), [ImageReward](https://github.com/THUDM/ImageReward), [HPSv2.1](https://github.com/tgxs002/HPSv2), FID and Validation Loss. Please refer to [evaluation/README.md](evaluation/README.md) for more details. |
| 152 | +```shell |
| 153 | +bash scripts/eval.sh |
| 154 | +``` |
| 155 | + |
| 156 | +## One More Thing: Infinity-20B is coming soon 📆 |
| 157 | +Infinity shows strong scaling capabilities as illustrated before. Thus we are encouraged to continue to scale up the model size to 20B. Here we present the side-by-side comparison results between Infinity-2B and Infinity-20B. |
| 158 | + |
| 159 | +| Prompt | Infinity (# params=2B) | Infinity (# params=20B) | |
| 160 | +| ------------ | -------- | -------- | |
| 161 | +| Create an image with the text "Always Priority" on a wooden sign |  |  | |
| 162 | +| Show the text 'Driver Unknown Hard Clearly' in a surreal, imaginative style with a dreamlike landscape backdrop. |  |  | |
| 163 | +| A photograph of a quaint two-story house with a distinctive red-tiled gable roof. The house is painted in a light, sandy color, which contrasts with the vibrant red roof. |  |  | |
| 164 | +| A group of students in a class |  |  | |
| 165 | + |
33 | 166 |
|
34 | 167 |
|
| 168 | +Currently, Infinity-20B is still on the training phrase. We will release Infinity-20B once the training is completed. |
35 | 169 |
|
36 | 170 | ## License |
37 | 171 | This project is licensed under the MIT License - see the [LICENSE](LICENSE) file for details. |
0 commit comments