upload

.gitattributes

@@ -33,3 +33,9 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+assets/random_vis_l32.png filter=lfs diff=lfs merge=lfs -text
+assets/rar_overview.png filter=lfs diff=lfs merge=lfs -text
+assets/recon_w_model_size_num_token.png filter=lfs diff=lfs merge=lfs -text
+assets/vis1.png filter=lfs diff=lfs merge=lfs -text
+assets/vis2.png filter=lfs diff=lfs merge=lfs -text
+assets/vis3.png filter=lfs diff=lfs merge=lfs -text

LICENSE

@@ -0,0 +1,201 @@
+                                 Apache License
+                           Version 2.0, January 2004
+                        http://www.apache.org/licenses/
+
+   TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
+
+   1. Definitions.
+
+      "License" shall mean the terms and conditions for use, reproduction,
+      and distribution as defined by Sections 1 through 9 of this document.
+
+      "Licensor" shall mean the copyright owner or entity authorized by
+      the copyright owner that is granting the License.
+
+      "Legal Entity" shall mean the union of the acting entity and all
+      other entities that control, are controlled by, or are under common
+      control with that entity. For the purposes of this definition,
+      "control" means (i) the power, direct or indirect, to cause the
+      direction or management of such entity, whether by contract or
+      otherwise, or (ii) ownership of fifty percent (50%) or more of the
+      outstanding shares, or (iii) beneficial ownership of such entity.
+
+      "You" (or "Your") shall mean an individual or Legal Entity
+      exercising permissions granted by this License.
+
+      "Source" form shall mean the preferred form for making modifications,
+      including but not limited to software source code, documentation
+      source, and configuration files.
+
+      "Object" form shall mean any form resulting from mechanical
+      transformation or translation of a Source form, including but
+      not limited to compiled object code, generated documentation,
+      and conversions to other media types.
+
+      "Work" shall mean the work of authorship, whether in Source or
+      Object form, made available under the License, as indicated by a
+      copyright notice that is included in or attached to the work
+      (an example is provided in the Appendix below).
+
+      "Derivative Works" shall mean any work, whether in Source or Object
+      form, that is based on (or derived from) the Work and for which the
+      editorial revisions, annotations, elaborations, or other modifications
+      represent, as a whole, an original work of authorship. For the purposes
+      of this License, Derivative Works shall not include works that remain
+      separable from, or merely link (or bind by name) to the interfaces of,
+      the Work and Derivative Works thereof.
+
+      "Contribution" shall mean any work of authorship, including
+      the original version of the Work and any modifications or additions
+      to that Work or Derivative Works thereof, that is intentionally
+      submitted to Licensor for inclusion in the Work by the copyright owner
+      or by an individual or Legal Entity authorized to submit on behalf of
+      the copyright owner. For the purposes of this definition, "submitted"
+      means any form of electronic, verbal, or written communication sent
+      to the Licensor or its representatives, including but not limited to
+      communication on electronic mailing lists, source code control systems,
+      and issue tracking systems that are managed by, or on behalf of, the
+      Licensor for the purpose of discussing and improving the Work, but
+      excluding communication that is conspicuously marked or otherwise
+      designated in writing by the copyright owner as "Not a Contribution."
+
+      "Contributor" shall mean Licensor and any individual or Legal Entity
+      on behalf of whom a Contribution has been received by Licensor and
+      subsequently incorporated within the Work.
+
+   2. Grant of Copyright License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      copyright license to reproduce, prepare Derivative Works of,
+      publicly display, publicly perform, sublicense, and distribute the
+      Work and such Derivative Works in Source or Object form.
+
+   3. Grant of Patent License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      (except as stated in this section) patent license to make, have made,
+      use, offer to sell, sell, import, and otherwise transfer the Work,
+      where such license applies only to those patent claims licensable
+      by such Contributor that are necessarily infringed by their
+      Contribution(s) alone or by combination of their Contribution(s)
+      with the Work to which such Contribution(s) was submitted. If You
+      institute patent litigation against any entity (including a
+      cross-claim or counterclaim in a lawsuit) alleging that the Work
+      or a Contribution incorporated within the Work constitutes direct
+      or contributory patent infringement, then any patent licenses
+      granted to You under this License for that Work shall terminate
+      as of the date such litigation is filed.
+
+   4. Redistribution. You may reproduce and distribute copies of the
+      Work or Derivative Works thereof in any medium, with or without
+      modifications, and in Source or Object form, provided that You
+      meet the following conditions:
+
+      (a) You must give any other recipients of the Work or
+          Derivative Works a copy of this License; and
+
+      (b) You must cause any modified files to carry prominent notices
+          stating that You changed the files; and
+
+      (c) You must retain, in the Source form of any Derivative Works
+          that You distribute, all copyright, patent, trademark, and
+          attribution notices from the Source form of the Work,
+          excluding those notices that do not pertain to any part of
+          the Derivative Works; and
+
+      (d) If the Work includes a "NOTICE" text file as part of its
+          distribution, then any Derivative Works that You distribute must
+          include a readable copy of the attribution notices contained
+          within such NOTICE file, excluding those notices that do not
+          pertain to any part of the Derivative Works, in at least one
+          of the following places: within a NOTICE text file distributed
+          as part of the Derivative Works; within the Source form or
+          documentation, if provided along with the Derivative Works; or,
+          within a display generated by the Derivative Works, if and
+          wherever such third-party notices normally appear. The contents
+          of the NOTICE file are for informational purposes only and
+          do not modify the License. You may add Your own attribution
+          notices within Derivative Works that You distribute, alongside
+          or as an addendum to the NOTICE text from the Work, provided
+          that such additional attribution notices cannot be construed
+          as modifying the License.
+
+      You may add Your own copyright statement to Your modifications and
+      may provide additional or different license terms and conditions
+      for use, reproduction, or distribution of Your modifications, or
+      for any such Derivative Works as a whole, provided Your use,
+      reproduction, and distribution of the Work otherwise complies with
+      the conditions stated in this License.
+
+   5. Submission of Contributions. Unless You explicitly state otherwise,
+      any Contribution intentionally submitted for inclusion in the Work
+      by You to the Licensor shall be under the terms and conditions of
+      this License, without any additional terms or conditions.
+      Notwithstanding the above, nothing herein shall supersede or modify
+      the terms of any separate license agreement you may have executed
+      with Licensor regarding such Contributions.
+
+   6. Trademarks. This License does not grant permission to use the trade
+      names, trademarks, service marks, or product names of the Licensor,
+      except as required for reasonable and customary use in describing the
+      origin of the Work and reproducing the content of the NOTICE file.
+
+   7. Disclaimer of Warranty. Unless required by applicable law or
+      agreed to in writing, Licensor provides the Work (and each
+      Contributor provides its Contributions) on an "AS IS" BASIS,
+      WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
+      implied, including, without limitation, any warranties or conditions
+      of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A
+      PARTICULAR PURPOSE. You are solely responsible for determining the
+      appropriateness of using or redistributing the Work and assume any
+      risks associated with Your exercise of permissions under this License.
+
+   8. Limitation of Liability. In no event and under no legal theory,
+      whether in tort (including negligence), contract, or otherwise,
+      unless required by applicable law (such as deliberate and grossly
+      negligent acts) or agreed to in writing, shall any Contributor be
+      liable to You for damages, including any direct, indirect, special,
+      incidental, or consequential damages of any character arising as a
+      result of this License or out of the use or inability to use the
+      Work (including but not limited to damages for loss of goodwill,
+      work stoppage, computer failure or malfunction, or any and all
+      other commercial damages or losses), even if such Contributor
+      has been advised of the possibility of such damages.
+
+   9. Accepting Warranty or Additional Liability. While redistributing
+      the Work or Derivative Works thereof, You may choose to offer,
+      and charge a fee for, acceptance of support, warranty, indemnity,
+      or other liability obligations and/or rights consistent with this
+      License. However, in accepting such obligations, You may act only
+      on Your own behalf and on Your sole responsibility, not on behalf
+      of any other Contributor, and only if You agree to indemnify,
+      defend, and hold each Contributor harmless for any liability
+      incurred by, or claims asserted against, such Contributor by reason
+      of your accepting any such warranty or additional liability.
+
+   END OF TERMS AND CONDITIONS
+
+   APPENDIX: How to apply the Apache License to your work.
+
+      To apply the Apache License to your work, attach the following
+      boilerplate notice, with the fields enclosed by brackets "[]"
+      replaced with your own identifying information. (Don't include
+      the brackets!)  The text should be enclosed in the appropriate
+      comment syntax for the file format. We also recommend that a
+      file or class name and description of purpose be included on the
+      same "printed page" as the copyright notice for easier
+      identification within third-party archives.
+
+   Copyright [yyyy] [name of copyright owner]
+
+   Licensed under the Apache License, Version 2.0 (the "License");
+   you may not use this file except in compliance with the License.
+   You may obtain a copy of the License at
+
+       http://www.apache.org/licenses/LICENSE-2.0
+
+   Unless required by applicable law or agreed to in writing, software
+   distributed under the License is distributed on an "AS IS" BASIS,
+   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+   See the License for the specific language governing permissions and
+   limitations under the License.

README.md

@@ -1,14 +1,82 @@
----
-title: RAR
-emoji: 📈
-colorFrom: purple
-colorTo: purple
-sdk: gradio
-sdk_version: 5.6.0
-app_file: app.py
-pinned: false
-license: apache-2.0
-short_description: Randomized Autoregressive Visual Generation
----
-
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+# 1D Visual Tokenization and Generation
+
+This repo hosts the code and models for the following projects:
+
+- RAR: [Randomized Autoregressive Visual Generation](https://yucornetto.github.io/projects/rar.html)
+
+- TiTok: [An Image is Worth 32 Tokens for Reconstruction and Generation](https://yucornetto.github.io/projects/titok.html)
+
+
+## Short Intro on [Randomized Autoregressive Visual Generation](https://arxiv.org/abs/2406.07550) ([README](README_RAR.md))
+
+RAR is a an autoregressive (AR) image generator with full compatibility to language modeling. It introduces a randomness annealing strategy with permuted objective at no additional cost, which enhances the model's ability to learn bidirectional contexts while leaving the autoregressive framework intact. RAR sets a FID score 1.48, demonstrating state-of-the-art performance on ImageNet-256 benchmark and significantly outperforming prior AR image generators.
+
+<p>
+<img src="assets/rar_overview.png" alt="teaser" width=90% height=90%>
+</p>
+<p>
+<img src="assets/perf_comp.png" alt="teaser" width=90% height=90%>
+</p>
+
+See more details at [README_RAR](README_RAR.md).
+
+## Short Intro on [An Image is Worth 32 Tokens for Reconstruction and Generation](https://arxiv.org/abs/2406.07550) ([README](README_TiTok.md))
+
+We present a compact 1D tokenizer which can represent an image with as few as 32 discrete tokens. As a result, it leads to a substantial speed-up on the sampling process (e.g., **410 × faster** than DiT-XL/2) while obtaining a competitive generation quality.
+
+<p>
+<img src="assets/titok_teaser.png" alt="teaser" width=90% height=90%>
+</p>
+<p>
+<img src="assets/speed_vs_perf.png" alt="teaser" width=90% height=90%>
+</p>
+
+See more details at [README_TiTok](README_TiTok.md).
+
+## Updates
+- 11/04/2024: We release the [tech report](https://arxiv.org/abs/2411.00776) and code for RAR models.
+- 10/16/2024: We update a set of TiTok tokenizer weights trained with an updated single-stage recipe, leading to easier training and better performance. We release the weight of different model size for both VQ and VAE variants TiTok, which we hope could facilitate the research in this area. More details will be available in a tech report later. 
+- 09/25/2024: TiTok is accepted by NeurIPS 2024.
+- 09/11/2024: Release the training codes of generator based on TiTok. 
+- 08/28/2024: Release the training codes of TiTok.
+- 08/09/2024: Better support on loading pretrained weights from huggingface models, thanks for the help from [@NielsRogge](https://github.com/NielsRogge)！
+- 07/03/2024: Evaluation scripts for reproducing the results reported in the paper, checkpoints of TiTok-B64 and TiTok-S128 are available.
+- 06/21/2024: Demo code and TiTok-L-32 checkpoints release. 
+- 06/11/2024: The [tech report](https://arxiv.org/abs/2406.07550) of TiTok is available.
+
+
+## Installation
+```shell
+pip3 install -r requirements.txt
+```
+
+## Citing
+If you use our work in your research, please use the following BibTeX entry.
+
+```BibTeX
+@article{yu2024randomized,
+  author    = {Qihang Yu and Ju He and Xueqing Deng and Xiaohui Shen and Liang-Chieh Chen},
+  title     = {Randomized Autoregressive Visual Generation},
+  journal   = {arXiv preprint arXiv:2411.00776},
+  year      = {2024}
+}
+```
+
+```BibTeX
+@article{yu2024an,
+  author    = {Qihang Yu and Mark Weber and Xueqing Deng and Xiaohui Shen and Daniel Cremers and Liang-Chieh Chen},
+  title     = {An Image is Worth 32 Tokens for Reconstruction and Generation},
+  journal   = {NeurIPS},
+  year      = {2024}
+}
+```
+
+## Acknowledgement
+
+[MaskGIT](https://github.com/google-research/maskgit)
+
+[Taming-Transformers](https://github.com/CompVis/taming-transformers)
+
+[Open-MUSE](https://github.com/huggingface/open-muse)
+
+[MUSE-Pytorch](https://github.com/baaivision/MUSE-Pytorch)

README_RAR.md

@@ -0,0 +1,221 @@
+# Randomized Autoregressive Visual Generation
+
+
+RAR is a an autoregressive (AR) image generator with full compatibility to language modeling. It introduces a randomness annealing strategy with permuted objective at no additional cost, which enhances the model's ability to learn bidirectional contexts while leaving the autoregressive framework intact. RAR sets a FID score 1.48, demonstrating state-of-the-art performance on ImageNet-256 benchmark and significantly outperforming prior AR image generators.
+
+
+<p>
+<img src="assets/rar_overview.png" alt="teaser" width=90% height=90%>
+</p>
+<p>
+<img src="assets/perf_comp.png" alt="teaser" width=90% height=90%>
+</p>
+
+
+## 🚀 Contributions
+
+#### We introduce RAR, an improved training strategy enabling standard autoregressive image generator to achieve state-of-the-art performance.
+
+#### The proposed RAR is extremly simple yet effective: During training, we randomly permute the input token sequence with probability r, where r will starts at 1.0 and linearly decays to 0.0 over the course of training. This simple strategy enbales better bidirectional representation learning which is missing in standard raster-order-based AR image generator training.
+
+#### RAR keeps the AR framework intact, and thus it is totally compatible to the LLM optimization techniques, such as KV-cache, leading to a significantly faster sampling speed compared to MAR-H or MaskBit while maintaining a better performance.
+
+## Model Zoo
+| Model | Link | FID |
+| ------------- | ------------- | ------------- |
+| RAR-B | [checkpoint](https://huggingface.co/yucornetto/RAR/blob/main/rar_b.bin)| 1.95 (generation) |
+| RAR-L | [checkpoint](https://huggingface.co/yucornetto/RAR/blob/main/rar_l.bin)| 1.70 (generation) |
+| RAR-XL | [checkpoint](https://huggingface.co/yucornetto/RAR/blob/main/rar_xl.bin)| 1.50 (generation) |
+| RAR-XXL | [checkpoint](https://huggingface.co/yucornetto/RAR/blob/main/rar_xxl.bin)| 1.48 (generation) |
+
+Please note that these models are trained only on limited academic dataset ImageNet, and they are only for research purposes.
+
+## Installation
+```shell
+pip3 install -r requirements.txt
+```
+
+## Get Started
+```python
+import torch
+from PIL import Image
+import numpy as np
+import demo_util
+from huggingface_hub import hf_hub_download
+from utils.train_utils import create_pretrained_tokenizer
+
+
+# Choose one from ["rar_b_imagenet", "rar_l_imagenet", "rar_xl_imagenet", "rar_xxl_imagenet"]
+rar_model_name = ["rar_b", "rar_l", "rar_xl", "rar_xxl"][3]
+
+# download the maskgit-vq tokenizer
+hf_hub_download(repo_id="fun-research/TiTok", filename=f"maskgit-vqgan-imagenet-f16-256.bin", local_dir="./")
+# download the rar generator weight
+hf_hub_download(repo_id="yucornetto/RAR", filename=f"{rar_model_name}.bin", local_dir="./")
+
+# load config
+# config = demo_util.get_config("configs/infer/titok_l32.yaml")
+# titok_tokenizer = demo_util.get_titok_tokenizer(config)
+# titok_generator = demo_util.get_titok_generator(config)
+
+device = "cuda"
+# maskgit-vq as tokenizer
+tokenizer = create_pretrained_tokenizer(config)
+generator = demo_util.get_rar_generator(config)
+tokenizer.to(device)
+generator.to(device)
+
+# generate an image
+sample_labels = [torch.randint(0, 999, size=(1,)).item()] # random IN-1k class
+generated_image = demo_util.sample_fn(
+    generator=generator,
+    tokenizer=tokenizer,
+    labels=sample_labels,
+    randomize_temperature=1.0,
+    guidance_scale=4.0,
+    guidance_scale_pow=0.0, # constant cfg
+    device=device
+)
+Image.fromarray(generated_image[0]).save(f"assets/rar_generated_{sample_labels[0]}.png")
+```
+
+## Testing on ImageNet-1K Benchmark
+
+We provide a [sampling script](./sample_imagenet_rar.py) for reproducing the generation results on ImageNet-1K benchmark.
+```bash
+# Prepare ADM evaluation script
+git clone https://github.com/openai/guided-diffusion.git
+
+wget https://openaipublic.blob.core.windows.net/diffusion/jul-2021/ref_batches/imagenet/256/VIRTUAL_imagenet256_labeled.npz
+```
+```python
+# Reproducing RAR-B
+torchrun --nnodes=1 --nproc_per_node=8 --rdzv-endpoint=localhost:9999 sample_imagenet_rar.py config=configs/training/generator/rar.yaml \
+    experiment.output_dir="rar_b" \
+    experiment.generator_checkpoint="rar_b.bin" \
+    model.generator.hidden_size=768 \
+    model.generator.num_hidden_layers=24 \
+    model.generator.num_attention_heads=16 \
+    model.generator.intermediate_size=3072 \
+    model.generator.randomize_temperature=1.0 \
+    model.generator.guidance_scale=16.0 \
+    model.generator.guidance_scale_pow=2.75
+# Run eval script. The result FID should be ~1.95
+python3 guided-diffusion/evaluations/evaluator.py VIRTUAL_imagenet256_labeled.npz rar_b.npz
+
+# Reproducing RAR-L
+torchrun --nnodes=1 --nproc_per_node=8 --rdzv-endpoint=localhost:9999 sample_imagenet_rar.py config=configs/training/generator/rar.yaml \
+    experiment.output_dir="rar_l" \
+    experiment.generator_checkpoint="rar_l.bin" \
+    model.generator.hidden_size=1024 \
+    model.generator.num_hidden_layers=24 \
+    model.generator.num_attention_heads=16 \
+    model.generator.intermediate_size=4096 \
+    model.generator.randomize_temperature=1.02 \
+    model.generator.guidance_scale=15.5 \
+    model.generator.guidance_scale_pow=2.5
+# Run eval script. The result FID should be ~1.70
+python3 guided-diffusion/evaluations/evaluator.py VIRTUAL_imagenet256_labeled.npz rar_l.npz
+
+# Reproducing RAR-XL
+torchrun --nnodes=1 --nproc_per_node=8 --rdzv-endpoint=localhost:9999 sample_imagenet_rar.py config=configs/training/generator/rar.yaml \
+    experiment.output_dir="rar_xl" \
+    experiment.generator_checkpoint="rar_xl.bin" \
+    model.generator.hidden_size=1280 \
+    model.generator.num_hidden_layers=32 \
+    model.generator.num_attention_heads=16 \
+    model.generator.intermediate_size=5120 \
+    model.generator.randomize_temperature=1.02 \
+    model.generator.guidance_scale=6.9 \
+    model.generator.guidance_scale_pow=1.5
+# Run eval script. The result FID should be ~1.50
+python3 guided-diffusion/evaluations/evaluator.py VIRTUAL_imagenet256_labeled.npz rar_xl.npz
+
+# Reproducing RAR-XXL
+torchrun --nnodes=1 --nproc_per_node=8 --rdzv-endpoint=localhost:9999 sample_imagenet_rar.py config=configs/training/generator/rar.yaml \
+    experiment.output_dir="rar_xxl" \
+    experiment.generator_checkpoint="rar_xxl.bin" \
+    model.generator.hidden_size=1408 \
+    model.generator.num_hidden_layers=40 \
+    model.generator.num_attention_heads=16 \
+    model.generator.intermediate_size=6144 \
+    model.generator.randomize_temperature=1.02 \
+    model.generator.guidance_scale=8.0 \
+    model.generator.guidance_scale_pow=1.2
+# Run eval script. The result FID should be ~1.48
+python3 guided-diffusion/evaluations/evaluator.py VIRTUAL_imagenet256_labeled.npz rar_xxl.npz
+```
+## Training Preparation
+We pretokenize the whole dataset for speed-up the training process. We have uploaded [it](https://huggingface.co/yucornetto/RAR/blob/main/maskgitvq.jsonl) so you can train RAR directly. The training script will download the prerequisite checkpoints and dataset automatically.
+
+## Training
+We provide example commands to train RAR as follows:
+```bash
+# Training for RAR-B
+WANDB_MODE=offline accelerate launch --num_machines=4 --num_processes=32 --machine_rank=${MACHINE_RANK} --main_process_ip=${ROOT_IP} --main_process_port=${ROOT_PORT} --same_network scripts/train_rar.py config=configs/training/generator/rar.yaml \
+    experiment.project="rar" \
+    experiment.name="rar_b" \
+    experiment.output_dir="rar_b" \
+    model.generator.hidden_size=768 \
+    model.generator.num_hidden_layers=24 \
+    model.generator.num_attention_heads=16 \
+    model.generator.intermediate_size=3072
+
+# Training for RAR-L
+WANDB_MODE=offline accelerate launch --num_machines=4 --num_processes=32 --machine_rank=${MACHINE_RANK} --main_process_ip=${ROOT_IP} --main_process_port=${ROOT_PORT} --same_network scripts/train_rar.py config=configs/training/generator/rar.yaml \
+    experiment.project="rar" \
+    experiment.name="rar_l" \
+    experiment.output_dir="rar_l" \
+    model.generator.hidden_size=1024 \
+    model.generator.num_hidden_layers=24 \
+    model.generator.num_attention_heads=16 \
+    model.generator.intermediate_size=4096
+
+# Training for RAR-XL
+WANDB_MODE=offline accelerate launch --num_machines=4 --num_processes=32 --machine_rank=${MACHINE_RANK} --main_process_ip=${ROOT_IP} --main_process_port=${ROOT_PORT} --same_network scripts/train_rar.py config=configs/training/generator/rar.yaml \
+    experiment.project="rar" \
+    experiment.name="rar_xl" \
+    experiment.output_dir="rar_xl" \
+    model.generator.hidden_size=1280 \
+    model.generator.num_hidden_layers=32 \
+    model.generator.num_attention_heads=16 \
+    model.generator.intermediate_size=5120
+
+# Training for RAR-XXL
+WANDB_MODE=offline accelerate launch --num_machines=4 --num_processes=32 --machine_rank=${MACHINE_RANK} --main_process_ip=${ROOT_IP} --main_process_port=${ROOT_PORT} --same_network scripts/train_rar.py config=configs/training/generator/rar.yaml \
+    experiment.project="rar" \
+    experiment.name="rar_xxl" \
+    experiment.output_dir="rar_xxl" \
+    model.generator.hidden_size=1408 \
+    model.generator.num_hidden_layers=40 \
+    model.generator.num_attention_heads=16 \
+    model.generator.intermediate_size=6144
+```
+You may remove the flag "WANDB_MODE=offline" to support online wandb logging, if you have configured it.
+
+Notably, you can enable grad checkpointing by adding the flag "model.generator.use_checkpoint=True" and adjust the machine number & GPU number based on your own need. All RAR checkpoints were trained with a global batchsize = 2048.
+
+
+## Visualizations
+<p>
+<img src="assets/vis1.png" alt="teaser" width=90% height=90%>
+</p>
+<p>
+<img src="assets/vis2.png" alt="teaser" width=90% height=90%>
+</p>
+<p>
+<img src="assets/vis3.png" alt="teaser" width=90% height=90%>
+</p>
+
+
+## Citing
+If you use our work in your research, please use the following BibTeX entry.
+
+```BibTeX
+@inproceedings{yu2024randomized,
+  author    = {Qihang Yu and Ju He and Xueqing Deng and Xiaohui Shen and Liang-Chieh Chen},
+  title     = {Randomized Autoregressive Visual Generation},
+  journal   = {arXiv preprint arXiv:2411.00776},
+  year      = {2024}
+}
+```

README_TiTok.md

@@ -0,0 +1,213 @@
+# (NeurIPS 2024) Compact and Mighty - Image Tokenization with Only 32 Tokens for both Reconstruction and Generation!
+
+<div align="center">
+
+[![demo](https://img.shields.io/badge/%F0%9F%A4%97%20Hugging%20Face-Online_Demo-blue)](https://huggingface.co/spaces/fun-research/TiTok)&nbsp;&nbsp;
+[![Website](https://img.shields.io/badge/Project-Website-87CEEB)](https://yucornetto.github.io/projects/titok.html)&nbsp;&nbsp;
+[![paper](https://img.shields.io/badge/arXiv-Paper-<COLOR>.svg)](https://arxiv.org/abs/2406.07550)&nbsp;&nbsp;
+
+</div>
+
+We present a compact 1D tokenizer which can represent an image with as few as 32 discrete tokens. As a result, it leads to a substantial speed-up on the sampling process (e.g., **410 × faster** than DiT-XL/2) while obtaining a competitive generation quality.
+
+
+<p>
+<img src="assets/titok_teaser.png" alt="teaser" width=90% height=90%>
+</p>
+<p>
+<img src="assets/speed_vs_perf.png" alt="teaser" width=90% height=90%>
+</p>
+
+
+## 🚀 Contributions
+
+#### We introduce a novel 1D image tokenization framework that breaks grid constraints existing in 2D tokenization methods, leading to a much more flexible and compact image latent representation.
+
+#### The proposed 1D tokenizer can tokenize a 256 × 256 image into as few as 32 discrete tokens, leading to a significant speed-up (hundreds times faster than diffusion models) in generation process, while maintaining state-of-the-art generation quality.
+
+#### We conduct a series of experiments to probe the properties of rarely studied 1D image tokenization, paving the path towards compact latent space for efficient and effective image representation.
+
+## Model Zoo
+| Model | Link | FID |
+| ------------- | ------------- | ------------- |
+| TiTok-L-32 Tokenizer | [checkpoint](https://huggingface.co/yucornetto/tokenizer_titok_l32_imagenet)| 2.21 (reconstruction) |
+| TiTok-B-64 Tokenizer | [checkpoint](https://huggingface.co/yucornetto/tokenizer_titok_b64_imagenet) | 1.70 (reconstruction) |
+| TiTok-S-128 Tokenizer | [checkpoint](https://huggingface.co/yucornetto/tokenizer_titok_s128_imagenet) | 1.71 (reconstruction) |
+| TiTok-L-32 Generator | [checkpoint](https://huggingface.co/yucornetto/generator_titok_l32_imagenet) | 2.77 (generation) |
+| TiTok-B-64 Generator | [checkpoint](https://huggingface.co/yucornetto/generator_titok_b64_imagenet) | 2.48 (generation) |
+| TiTok-S-128 Generator | [checkpoint](https://huggingface.co/yucornetto/generator_titok_s128_imagenet) | 1.97 (generation) |
+| TiTok-BL-64 VQ Tokenizer | [checkpoint](https://huggingface.co/yucornetto/tokenizer_titok_bl64_vq8k_imagenet)| 2.06 (reconstruction) |
+| TiTok-BL-128 VQ Tokenizer | [checkpoint](https://huggingface.co/yucornetto/tokenizer_titok_bl128_vq8k_imagenet)| 1.49 (reconstruction) |
+| TiTok-SL-256 VQ Tokenizer | [checkpoint](https://huggingface.co/yucornetto/tokenizer_titok_sl256_vq8k_imagenet)| 1.03 (reconstruction) |
+| TiTok-LL-32 VAE Tokenizer | [checkpoint](https://huggingface.co/yucornetto/tokenizer_titok_ll32_vae_c16_imagenet)| 1.61 (reconstruction) |
+| TiTok-BL-64 VAE Tokenizer | [checkpoint](https://huggingface.co/yucornetto/tokenizer_titok_bl64_vae_c16_imagenet)| 1.25 (reconstruction) |
+| TiTok-BL-128 VAE Tokenizer | [checkpoint](https://huggingface.co/yucornetto/tokenizer_titok_bl128_vae_c16_imagenet)| 0.84 (reconstruction) |
+
+Please note that these models are trained only on limited academic dataset ImageNet, and they are only for research purposes.
+
+## Installation
+```shell
+pip3 install -r requirements.txt
+```
+
+## Get Started
+```python
+import torch
+from PIL import Image
+import numpy as np
+import demo_util
+from huggingface_hub import hf_hub_download
+from modeling.maskgit import ImageBert
+from modeling.titok import TiTok
+
+# Choose one from ["tokenizer_titok_l32_imagenet", "tokenizer_titok_b64_imagenet",
+#  "tokenizer_titok_s128_imagenet", "tokenizer_titok_bl128_vae_c16_imagenet", tokenizer_titok_bl64_vae_c16_imagenet",
+# "tokenizer_titok_ll32_vae_c16_imagenet", "tokenizer_titok_sl256_vq8k_imagenet", "tokenizer_titok_bl128_vq8k_imagenet",
+# "tokenizer_titok_bl64_vq8k_imagenet",]
+titok_tokenizer = TiTok.from_pretrained("yucornetto/tokenizer_titok_l32_imagenet")
+titok_tokenizer.eval()
+titok_tokenizer.requires_grad_(False)
+titok_generator = ImageBert.from_pretrained("yucornetto/generator_titok_l32_imagenet")
+titok_generator.eval()
+titok_generator.requires_grad_(False)
+
+# or alternatively, downloads from hf
+# hf_hub_download(repo_id="fun-research/TiTok", filename="tokenizer_titok_l32.bin", local_dir="./")
+# hf_hub_download(repo_id="fun-research/TiTok", filename="generator_titok_l32.bin", local_dir="./")
+
+# load config
+# config = demo_util.get_config("configs/infer/titok_l32.yaml")
+# titok_tokenizer = demo_util.get_titok_tokenizer(config)
+# titok_generator = demo_util.get_titok_generator(config)
+
+device = "cuda"
+titok_tokenizer = titok_tokenizer.to(device)
+titok_generator = titok_generator.to(device)
+
+# reconstruct an image. I.e., image -> 32 tokens -> image
+img_path = "assets/ILSVRC2012_val_00010240.png"
+image = torch.from_numpy(np.array(Image.open(img_path)).astype(np.float32)).permute(2, 0, 1).unsqueeze(0) / 255.0
+# tokenization
+if titok_tokenizer.quantize_mode == "vq":
+    encoded_tokens = titok_tokenizer.encode(image.to(device))[1]["min_encoding_indices"]
+elif titok_tokenizer.quantize_mode == "vae":
+    posteriors = titok_tokenizer.encode(image.to(device))[1]
+    encoded_tokens = posteriors.sample()
+else:
+    raise NotImplementedError
+# image assets/ILSVRC2012_val_00010240.png is encoded into tokens tensor([[[ 887, 3979,  349,  720, 2809, 2743, 2101,  603, 2205, 1508, 1891, 4015, 1317, 2956, 3774, 2296,  484, 2612, 3472, 2330, 3140, 3113, 1056, 3779,  654, 2360, 1901, 2908, 2169,  953, 1326, 2598]]], device='cuda:0'), with shape torch.Size([1, 1, 32])
+print(f"image {img_path} is encoded into tokens {encoded_tokens}, with shape {encoded_tokens.shape}")
+# de-tokenization
+reconstructed_image = titok_tokenizer.decode_tokens(encoded_tokens)
+reconstructed_image = torch.clamp(reconstructed_image, 0.0, 1.0)
+reconstructed_image = (reconstructed_image * 255.0).permute(0, 2, 3, 1).to("cpu", dtype=torch.uint8).numpy()[0]
+reconstructed_image = Image.fromarray(reconstructed_image).save("assets/ILSVRC2012_val_00010240_recon.png")
+
+# generate an image
+sample_labels = [torch.randint(0, 999, size=(1,)).item()] # random IN-1k class
+generated_image = demo_util.sample_fn(
+    generator=titok_generator,
+    tokenizer=titok_tokenizer,
+    labels=sample_labels,
+    guidance_scale=4.5,
+    randomize_temperature=1.0,
+    num_sample_steps=8,
+    device=device
+)
+Image.fromarray(generated_image[0]).save(f"assets/generated_{sample_labels[0]}.png")
+```
+
+We also provide a [jupyter notebook](demo.ipynb) for a quick tutorial on reconstructing and generating images with TiTok-L-32.
+
+We also support TiTok with [HuggingFace 🤗 Demo](https://huggingface.co/spaces/fun-research/TiTok)!
+
+## Testing on ImageNet-1K Benchmark
+
+We provide a [sampling script](./sample_imagenet_titok.py) for reproducing the generation results on ImageNet-1K benchmark.
+```bash
+# Prepare ADM evaluation script
+git clone https://github.com/openai/guided-diffusion.git
+
+wget https://openaipublic.blob.core.windows.net/diffusion/jul-2021/ref_batches/imagenet/256/VIRTUAL_imagenet256_labeled.npz
+```
+```python
+# Reproducing TiTok-L-32
+torchrun --nnodes=1 --nproc_per_node=8 --rdzv-endpoint=localhost:9999 sample_imagenet_titok.py config=configs/infer/titok_l32.yaml experiment.output_dir="titok_l_32"
+# Run eval script. The result FID should be ~2.77
+python3 guided-diffusion/evaluations/evaluator.py VIRTUAL_imagenet256_labeled.npz titok_l_32.npz
+
+# Reproducing TiTok-B-64
+torchrun --nnodes=1 --nproc_per_node=8 --rdzv-endpoint=localhost:9999 sample_imagenet_titok.py config=configs/infer/titok_b64.yaml experiment.output_dir="titok_b_64"
+# Run eval script. The result FID should be ~2.48
+python3 guided-diffusion/evaluations/evaluator.py VIRTUAL_imagenet256_labeled.npz titok_b_64.npz
+
+# Reproducing TiTok-S-128
+torchrun --nnodes=1 --nproc_per_node=8 --rdzv-endpoint=localhost:9999 sample_imagenet_titok.py config=configs/infer/titok_s128.yaml experiment.output_dir="titok_s_128"
+# Run eval script. The result FID should be ~1.97
+python3 guided-diffusion/evaluations/evaluator.py VIRTUAL_imagenet256_labeled.npz titok_s_128.npz
+```
+## Training Preparation
+We use [webdataset](https://github.com/webdataset/webdataset) format for data loading. To begin with, it is needed to convert the dataset into webdataset format. An example script to convert ImageNet to wds format is provided [here](./data/convert_imagenet_to_wds.py).
+
+Furthermore, the stage1 training relies on a pre-trained MaskGIT-VQGAN to generate proxy codes as learning targets. You can convert the [official Jax weight](https://github.com/google-research/maskgit) to PyTorch version using [this script](https://github.com/huggingface/open-muse/blob/main/scripts/convert_maskgit_vqgan.py). Alternatively, we provided a converted version at [HuggingFace](https://huggingface.co/fun-research/TiTok/blob/main/maskgit-vqgan-imagenet-f16-256.bin) and [Google Drive](https://drive.google.com/file/d/1DjZqzJrUt2hwpmUPkjGSBTFEJcOkLY-Q/view?usp=sharing). The MaskGIT-VQGAN's weight will be automatically downloaded when you run the training script.
+
+## Training
+We provide example commands to train TiTok as follows:
+```bash
+# Training for TiTok-B64
+# Stage 1
+WANDB_MODE=offline accelerate launch --num_machines=1 --num_processes=8 --machine_rank=0 --main_process_ip=127.0.0.1 --main_process_port=9999 --same_network scripts/train_titok.py config=configs/training/stage1/titok_b64.yaml \
+    experiment.project="titok_b64_stage1" \
+    experiment.name="titok_b64_stage1_run1" \
+    experiment.output_dir="titok_b64_stage1_run1" \
+    training.per_gpu_batch_size=32
+
+# Stage 2
+WANDB_MODE=offline accelerate launch --num_machines=1 --num_processes=8 --machine_rank=0 --main_process_ip=127.0.0.1 --main_process_port=9999 --same_network scripts/train_titok.py config=configs/training/stage2/titok_b64.yaml \
+    experiment.project="titok_b64_stage2" \
+    experiment.name="titok_b64_stage2_run1" \
+    experiment.output_dir="titok_b64_stage2_run1" \
+    training.per_gpu_batch_size=32 \
+    experiment.init_weight=${PATH_TO_STAGE1_WEIGHT}
+
+# Train Generator (TiTok-B64 as example)
+WANDB_MODE=offline accelerate launch --num_machines=4 --num_processes=32 --machine_rank=${MACHINE_RANK} --main_process_ip=${ROOT_IP}--main_process_port=${ROOT_PORT} --same_network scripts/train_maskgit.py config=configs/training/generator/maskgit.yaml \
+    experiment.project="titok_generation" \
+    experiment.name="titok_b64_maskgit" \
+    experiment.output_dir="titok_b64_maskgit" \
+    experiment.tokenizer_checkpoint=${PATH_TO_STAGE1_or_STAGE2_WEIGHT}
+```
+You may remove the flag "WANDB_MODE=offline" to support online wandb logging, if you have configured it.
+
+The config _titok_b64.yaml_ can be replaced with _titok_s128.yaml_ or _titok_l32.yaml_ for other TiTok variants.
+
+## Visualizations
+<p>
+<img src="assets/recon_w_model_size_num_token.png" alt="teaser" width=90% height=90%>
+</p>
+<p>
+<img src="assets/random_vis_l32.png" alt="teaser" width=90% height=90%>
+</p>
+
+
+## Citing
+If you use our work in your research, please use the following BibTeX entry.
+
+```BibTeX
+@inproceedings{yu2024an,
+  author    = {Qihang Yu and Mark Weber and Xueqing Deng and Xiaohui Shen and Daniel Cremers and Liang-Chieh Chen},
+  title     = {An Image is Worth 32 Tokens for Reconstruction and Generation},
+  journal   = {NeurIPS},
+  year      = {2024}
+}
+```
+
+## Acknowledgement
+
+[MaskGIT](https://github.com/google-research/maskgit)
+
+[Taming-Transformers](https://github.com/CompVis/taming-transformers)
+
+[Open-MUSE](https://github.com/huggingface/open-muse)
+
+[MUSE-Pytorch](https://github.com/baaivision/MUSE-Pytorch)

configs/infer/titok_b64.yaml

@@ -0,0 +1,39 @@
+experiment:
+    tokenizer_checkpoint: "tokenizer_titok_b64.bin"
+    generator_checkpoint: "generator_titok_b64.bin"
+    output_dir: "titok_b_64"
+model:
+    vq_model:
+        codebook_size: 4096
+        token_size: 12
+        use_l2_norm: True
+        commitment_cost: 0.25
+        # vit arch
+        vit_enc_model_size: "base"
+        vit_dec_model_size: "base"
+        vit_enc_patch_size: 16
+        vit_dec_patch_size: 16
+        num_latent_tokens: 64
+        finetune_decoder: True
+    
+    generator:
+        model_type: "ViT"
+        hidden_size: 768
+        num_hidden_layers: 24
+        num_attention_heads: 16
+        intermediate_size: 3072
+        dropout: 0.1
+        attn_drop: 0.1
+        num_steps: 8
+        class_label_dropout: 0.1
+        image_seq_len: ${model.vq_model.num_latent_tokens}
+        condition_num_classes: 1000
+
+        # sampling hyper-params
+        randomize_temperature: 11.0
+        guidance_scale: 3.0
+        guidance_decay: "linear"
+
+dataset:
+    preprocessing:
+        crop_size: 256

configs/infer/titok_bl128_vae_c16.yaml

@@ -0,0 +1,19 @@
+experiment:
+    tokenizer_checkpoint: "titok_bl128_vae_c16.bin"
+    output_dir: "titok_bl128_vae_c16"
+model:
+    vq_model:
+        quantize_mode: "vae"
+        token_size: 16 
+        # vit arch
+        vit_enc_model_size: "base"
+        vit_dec_model_size: "large"
+        vit_enc_patch_size: 16
+        vit_dec_patch_size: 16
+        num_latent_tokens: 128
+        finetune_decoder: False
+        is_legacy: False
+
+dataset:
+    preprocessing:
+        crop_size: 256

configs/infer/titok_bl128_vq8k.yaml

@@ -0,0 +1,21 @@
+experiment:
+    tokenizer_checkpoint: "tokenizer_titok_bl128_vq8k.bin"
+    output_dir: "titok_bl128_vq8k"
+model:
+    vq_model:
+        codebook_size: 8192
+        token_size: 64
+        use_l2_norm: False
+        commitment_cost: 0.25
+        # vit arch
+        vit_enc_model_size: "base"
+        vit_dec_model_size: "large"
+        vit_enc_patch_size: 16
+        vit_dec_patch_size: 16
+        num_latent_tokens: 128
+        finetune_decoder: False
+        is_legacy: False
+
+dataset:
+    preprocessing:
+        crop_size: 256

configs/infer/titok_bl64_vae_c16.yaml

@@ -0,0 +1,19 @@
+experiment:
+    tokenizer_checkpoint: "titok_bl64_vae_c16.bin"
+    output_dir: "titok_bl64_vae_c16"
+model:
+    vq_model:
+        quantize_mode: "vae"
+        token_size: 16 
+        # vit arch
+        vit_enc_model_size: "base"
+        vit_dec_model_size: "large"
+        vit_enc_patch_size: 16
+        vit_dec_patch_size: 16
+        num_latent_tokens: 64
+        finetune_decoder: False
+        is_legacy: False
+
+dataset:
+    preprocessing:
+        crop_size: 256

configs/infer/titok_bl64_vq8k.yaml

@@ -0,0 +1,21 @@
+experiment:
+    tokenizer_checkpoint: "tokenizer_titok_bl64_vq8k.bin"
+    output_dir: "titok_bl64_vq8k"
+model:
+    vq_model:
+        codebook_size: 8192
+        token_size: 64
+        use_l2_norm: False
+        commitment_cost: 0.25
+        # vit arch
+        vit_enc_model_size: "base"
+        vit_dec_model_size: "large"
+        vit_enc_patch_size: 16
+        vit_dec_patch_size: 16
+        num_latent_tokens: 64
+        finetune_decoder: False
+        is_legacy: False
+
+dataset:
+    preprocessing:
+        crop_size: 256

configs/infer/titok_l32.yaml

@@ -0,0 +1,40 @@
+experiment:
+    tokenizer_checkpoint: "tokenizer_titok_l32.bin"
+    generator_checkpoint: "generator_titok_l32.bin"
+    output_dir: "titok_l_32"
+model:
+    vq_model:
+        codebook_size: 4096
+        token_size: 12
+        use_l2_norm: True
+        commitment_cost: 0.25
+        # vit arch
+        vit_enc_model_size: "large"
+        vit_dec_model_size: "large"
+        vit_enc_patch_size: 16
+        vit_dec_patch_size: 16
+        num_latent_tokens: 32
+        finetune_decoder: True
+    
+    generator:
+        model_type: "ViT"
+        hidden_size: 768
+        num_hidden_layers: 24
+        num_attention_heads: 16
+        intermediate_size: 3072
+        dropout: 0.1
+        attn_drop: 0.1
+        num_steps: 8
+        class_label_dropout: 0.1
+        image_seq_len: ${model.vq_model.num_latent_tokens}
+        condition_num_classes: 1000
+
+        # sampling hyper-params
+        randomize_temperature: 9.5
+        guidance_scale: 4.5
+        guidance_decay: "linear"
+        
+
+dataset:
+    preprocessing:
+        crop_size: 256

configs/infer/titok_ll32_vae_c16.yaml

@@ -0,0 +1,19 @@
+experiment:
+    tokenizer_checkpoint: "titok_ll32_vae_c16.bin"
+    output_dir: "titok_ll32_vae_c16"
+model:
+    vq_model:
+        quantize_mode: "vae"
+        token_size: 16 
+        # vit arch
+        vit_enc_model_size: "large"
+        vit_dec_model_size: "large"
+        vit_enc_patch_size: 16
+        vit_dec_patch_size: 16
+        num_latent_tokens: 32
+        finetune_decoder: False
+        is_legacy: False
+
+dataset:
+    preprocessing:
+        crop_size: 256

configs/infer/titok_s128.yaml

@@ -0,0 +1,39 @@
+experiment:
+    tokenizer_checkpoint: "tokenizer_titok_s128.bin"
+    generator_checkpoint: "generator_titok_s128.bin"
+    output_dir: "titok_s_128"
+model:
+    vq_model:
+        codebook_size: 4096
+        token_size: 12
+        use_l2_norm: True
+        commitment_cost: 0.25
+        # vit arch
+        vit_enc_model_size: "small"
+        vit_dec_model_size: "small"
+        vit_enc_patch_size: 16
+        vit_dec_patch_size: 16
+        num_latent_tokens: 128
+        finetune_decoder: True
+    
+    generator:
+        model_type: "UViT"
+        hidden_size: 1024
+        num_hidden_layers: 20
+        num_attention_heads: 16
+        intermediate_size: 4096
+        dropout: 0.1
+        attn_drop: 0.1
+        num_steps: 64
+        class_label_dropout: 0.1
+        image_seq_len: ${model.vq_model.num_latent_tokens}
+        condition_num_classes: 1000
+
+        # sampling hyper-params
+        randomize_temperature: 2.8
+        guidance_scale: 6.9
+        guidance_decay: "power-cosine"
+
+dataset:
+    preprocessing:
+        crop_size: 256

configs/infer/titok_sl256_vq8k.yaml

@@ -0,0 +1,21 @@
+experiment:
+    tokenizer_checkpoint: "tokenizer_titok_sl256_vq8k.bin"
+    output_dir: "titok_sl256_vq8k"
+model:
+    vq_model:
+        codebook_size: 8192
+        token_size: 64
+        use_l2_norm: False
+        commitment_cost: 0.25
+        # vit arch
+        vit_enc_model_size: "small"
+        vit_dec_model_size: "large"
+        vit_enc_patch_size: 16
+        vit_dec_patch_size: 16
+        num_latent_tokens: 256
+        finetune_decoder: False
+        is_legacy: False
+
+dataset:
+    preprocessing:
+        crop_size: 256

configs/training/generator/maskgit.yaml

@@ -0,0 +1,86 @@
+experiment:
+    project: "titok_generation"
+    name: "titok_b64_maskgit"
+    max_train_examples: 1_281_167
+    save_every: 50_000
+    eval_every: 50_000
+    generate_every: 5_000
+    log_every: 50
+    log_grad_norm_every: 1_000
+    resume: True
+    tokenizer_checkpoint: "tokenizer_titok_b64.bin"
+
+model:
+    vq_model:
+        codebook_size: 4096
+        token_size: 12
+        use_l2_norm: True
+        commitment_cost: 0.25
+        # vit arch
+        vit_enc_model_size: "base"
+        vit_dec_model_size: "base"
+        vit_enc_patch_size: 16
+        vit_dec_patch_size: 16
+        num_latent_tokens: 64
+        finetune_decoder: True
+    
+    generator:
+        model_type: "ViT"
+        hidden_size: 768
+        num_hidden_layers: 24
+        num_attention_heads: 16
+        intermediate_size: 3072
+        dropout: 0.1
+        attn_drop: 0.1
+        num_steps: 8
+        class_label_dropout: 0.1
+        image_seq_len: ${model.vq_model.num_latent_tokens}
+        condition_num_classes: 1000
+
+        # sampling hyper-params on the flight
+        randomize_temperature: 1.0
+        guidance_scale: 4.5
+        guidance_decay: "constant"
+
+losses:
+    label_smoothing: 0.1
+    loss_weight_unmasked_token: 0.1
+
+dataset:
+    params:
+        train_shards_path_or_url: "imagenet_sharded/train/imagenet-train-{0000..0252}.tar"
+        eval_shards_path_or_url: "imagenet_sharded/val/imagenet-val-{0000..0009}.tar"
+        num_workers_per_gpu: 12
+    preprocessing:
+        resize_shorter_edge: 256
+        crop_size: 256
+        random_crop: False
+        random_flip: True
+
+optimizer:
+    name: adamw 
+    params:
+        learning_rate: 2e-4
+        beta1: 0.9
+        beta2: 0.96
+        weight_decay: 0.03
+
+
+lr_scheduler:
+    scheduler: "cosine"
+    params:
+        learning_rate: ${optimizer.params.learning_rate}
+        warmup_steps: 10_000
+        end_lr: 1e-5
+
+
+training:
+    gradient_accumulation_steps: 1
+    per_gpu_batch_size: 64 # 32 GPU, total batch size 2048
+    mixed_precision: "bf16"
+    enable_tf32: True
+    enable_wandb: True
+    use_ema: True
+    seed: 42
+    max_train_steps: 500_000
+    max_grad_norm: 1.0

configs/training/generator/rar.yaml

@@ -0,0 +1,78 @@
+experiment:
+    project: "rar_generation"
+    name: "rar-b"
+    max_train_examples: 1_281_167
+    save_every: 50_000
+    eval_every: 50_000
+    generate_every: 5_000
+    log_every: 50
+    log_grad_norm_every: 1_000
+    resume: True
+
+model:
+    vq_model:
+        codebook_size: 1024
+        token_size: 256
+        num_latent_tokens: 256
+        finetune_decoder: False
+        pretrained_tokenizer_weight: "maskgit-vqgan-imagenet-f16-256.bin"
+    
+    generator:
+        hidden_size: 768
+        num_hidden_layers: 24
+        num_attention_heads: 16
+        intermediate_size: 3072
+        dropout: 0.1
+        attn_drop: 0.1
+        class_label_dropout: 0.1
+        image_seq_len: 256
+        condition_num_classes: 1000
+
+        # sampling hyper-params for RAR-B
+        randomize_temperature: 1.0
+        guidance_scale: 16.0
+        guidance_scale_pow: 2.75
+        use_checkpoint: False # True to save memory
+
+        randomness_anneal_start: 125000 # 200 epoch
+        randomness_anneal_end: 187500 # 300 epoch
+
+dataset:
+    params:
+        # use pretokenized dataset for speed-up
+        pretokenization: "maskgitvq.jsonl"
+        train_shards_path_or_url: "imagenet_sharded/train/imagenet-train-{0000..0252}.tar"
+        eval_shards_path_or_url: "imagenet_sharded/val/imagenet-val-{0000..0009}.tar"
+        num_workers_per_gpu: 12
+    preprocessing:
+        resize_shorter_edge: 256
+        crop_size: 256
+        random_crop: False
+        random_flip: True
+
+optimizer:
+    name: adamw 
+    params:
+        learning_rate: 4e-4
+        beta1: 0.9
+        beta2: 0.96
+        weight_decay: 0.03
+
+
+lr_scheduler:
+    scheduler: "cosine"
+    params:
+        learning_rate: ${optimizer.params.learning_rate}
+        warmup_steps: 62_500 # 100 epochs with bsz 2048
+        end_lr: 1e-5
+
+training:
+    gradient_accumulation_steps: 1
+    per_gpu_batch_size: 64 # 32 GPU, total batch size 2048
+    mixed_precision: "bf16"
+    enable_tf32: True
+    enable_wandb: True
+    use_ema: False
+    seed: 42
+    max_train_steps: 250_000
+    max_grad_norm: 1.0

configs/training/stage1/titok_b64.yaml

@@ -0,0 +1,70 @@
+experiment:
+    project: "titok_b_64_stage1"
+    name: "titok_b_64_stage1_run1"
+    output_dir: "titok_b_64_stage1_run1"
+    max_train_examples: 1_281_167
+    save_every: 50_000
+    eval_every: 50_000
+    generate_every: 5_000
+    log_every: 50
+    log_grad_norm_every: 1_000
+    resume: True
+    init_weight: ""
+
+
+model:
+    vq_model:
+        codebook_size: 4096
+        token_size: 12
+        use_l2_norm: True
+        commitment_cost: 0.25
+        # vit arch
+        vit_enc_model_size: "base"
+        vit_dec_model_size: "base"
+        vit_enc_patch_size: 16
+        vit_dec_patch_size: 16
+        num_latent_tokens: 64
+        finetune_decoder: False
+        pretrained_tokenizer_weight: "maskgit-vqgan-imagenet-f16-256.bin"
+
+losses:
+    quantizer_weight: 1.0
+
+dataset:
+    params:
+        train_shards_path_or_url: "imagenet_sharded/train/imagenet-train-{0000..0252}.tar"
+        eval_shards_path_or_url: "imagenet_sharded/val/imagenet-val-{0000..0009}.tar"
+        num_workers_per_gpu: 12
+    preprocessing:
+        resize_shorter_edge: 256
+        crop_size: 256
+        random_crop: True
+        random_flip: True
+
+
+optimizer:
+    name: adamw 
+    params:
+        learning_rate: 1e-4
+        beta1: 0.9
+        beta2: 0.99
+        weight_decay: 1e-4
+
+lr_scheduler:
+    scheduler: "cosine"
+    params:
+        learning_rate: ${optimizer.params.learning_rate}
+        warmup_steps: 10_000
+        end_lr: 1e-5
+
+training:
+    gradient_accumulation_steps: 1
+    per_gpu_batch_size: 32
+    mixed_precision: "fp16"
+    enable_tf32: True
+    enable_wandb: True
+    use_ema: True
+    seed: 42
+    max_train_steps: 1_000_000
+    num_generated_images: 2
+    max_grad_norm: 1.0

configs/training/stage1/titok_l32.yaml

@@ -0,0 +1,70 @@
+experiment:
+    project: "titok_l_32_stage1"
+    name: "titok_l_32_stage1_run1"
+    output_dir: "titok_l_32_stage1_run1"
+    max_train_examples: 1_281_167
+    save_every: 50_000
+    eval_every: 50_000
+    generate_every: 5_000
+    log_every: 50
+    log_grad_norm_every: 1_000
+    resume: True
+    init_weight: ""
+
+
+model:
+    vq_model:
+        codebook_size: 4096
+        token_size: 12
+        use_l2_norm: True
+        commitment_cost: 0.25
+        # vit arch
+        vit_enc_model_size: "large"
+        vit_dec_model_size: "large"
+        vit_enc_patch_size: 16
+        vit_dec_patch_size: 16
+        num_latent_tokens: 32
+        finetune_decoder: False
+        pretrained_tokenizer_weight: "maskgit-vqgan-imagenet-f16-256.bin"
+
+losses:
+    quantizer_weight: 1.0
+
+dataset:
+    params:
+        train_shards_path_or_url: "imagenet_sharded/train/imagenet-train-{0000..0252}.tar"
+        eval_shards_path_or_url: "imagenet_sharded/val/imagenet-val-{0000..0009}.tar"
+        num_workers_per_gpu: 12
+    preprocessing:
+        resize_shorter_edge: 256
+        crop_size: 256
+        random_crop: True
+        random_flip: True
+
+optimizer:
+    name: adamw 
+    params:
+        learning_rate: 1e-4
+        beta1: 0.9
+        beta2: 0.99
+        weight_decay: 1e-4
+        epsilon: 1e-8
+
+lr_scheduler:
+    scheduler: "cosine"
+    params:
+        learning_rate: ${optimizer.params.learning_rate}
+        warmup_steps: 10_000
+        end_lr: 1e-5
+
+training:
+    gradient_accumulation_steps: 1
+    per_gpu_batch_size: 32
+    mixed_precision: "fp16"
+    enable_tf32: True
+    enable_wandb: True
+    use_ema: True
+    seed: 42
+    max_train_steps: 1_000_000
+    num_generated_images: 2
+    max_grad_norm: 1.0

configs/training/stage1/titok_s128.yaml

@@ -0,0 +1,70 @@
+experiment:
+    project: "titok_s_128_stage1"
+    name: "titok_s_128_stage1_run1"
+    output_dir: "titok_s_128_stage1_run1"
+    max_train_examples: 1_281_167
+    save_every: 50_000
+    eval_every: 50_000
+    generate_every: 5_000
+    log_every: 50
+    log_grad_norm_every: 1_000
+    resume: True
+    init_weight: ""
+
+
+model:
+    vq_model:
+        codebook_size: 4096
+        token_size: 12
+        use_l2_norm: True
+        commitment_cost: 0.25
+        # vit arch
+        vit_enc_model_size: "small"
+        vit_dec_model_size: "small"
+        vit_enc_patch_size: 16
+        vit_dec_patch_size: 16
+        num_latent_tokens: 128
+        finetune_decoder: False
+        pretrained_tokenizer_weight: "maskgit-vqgan-imagenet-f16-256.bin"
+    
+losses:
+    quantizer_weight: 1.0
+
+dataset:
+    params:
+        train_shards_path_or_url: "imagenet_sharded/train/imagenet-train-{0000..0252}.tar"
+        eval_shards_path_or_url: "imagenet_sharded/val/imagenet-val-{0000..0009}.tar"
+        num_workers_per_gpu: 12
+    preprocessing:
+        resize_shorter_edge: 256
+        crop_size: 256
+        random_crop: True
+        random_flip: True
+
+optimizer:
+    name: adamw 
+    params:
+        learning_rate: 1e-4
+        beta1: 0.9
+        beta2: 0.99
+        weight_decay: 1e-4
+        epsilon: 1e-8
+
+lr_scheduler:
+    scheduler: "cosine"
+    params:
+        learning_rate: ${optimizer.params.learning_rate}
+        warmup_steps: 10_000
+        end_lr: 1e-5
+
+training:
+    gradient_accumulation_steps: 1
+    per_gpu_batch_size: 32
+    mixed_precision: "fp16"
+    enable_tf32: True
+    enable_wandb: True
+    use_ema: True
+    seed: 42
+    max_train_steps: 1_000_000
+    num_generated_images: 2
+    max_grad_norm: 1.0

configs/training/stage2/titok_b64.yaml

@@ -0,0 +1,80 @@
+experiment:
+    project: "titok_b_64_stage2"
+    name: "titok_b_64_stage2_run1"
+    output_dir: "titok_b_64_stage2_run1"
+    max_train_examples: 1_281_167
+    save_every: 50_000
+    eval_every: 50_000
+    generate_every: 5_000
+    log_every: 50
+    log_grad_norm_every: 1_000
+    resume: True
+    init_weight: "titok_b_64_stage1.bin"
+
+
+model:
+    vq_model:
+        codebook_size: 4096
+        token_size: 12
+        use_l2_norm: True
+        commitment_cost: 0.0
+        # vit arch
+        vit_enc_model_size: "base"
+        vit_dec_model_size: "base"
+        vit_enc_patch_size: 16
+        vit_dec_patch_size: 16
+        num_latent_tokens: 64
+        finetune_decoder: True
+        pretrained_tokenizer_weight: "maskgit-vqgan-imagenet-f16-256.bin"
+
+losses:
+    discriminator_start: 20_000
+    quantizer_weight: 0.0
+    discriminator_factor: 1.0
+    discriminator_weight: 0.01
+    perceptual_loss: "convnext_s"
+    perceptual_weight: 0.1
+    reconstruction_loss: "l2"
+    reconstruction_weight: 1.0
+    lecam_regularization_weight: 0.001
+
+
+dataset:
+    params:
+        train_shards_path_or_url: "imagenet_sharded/train/imagenet-train-{0000..0252}.tar"
+        eval_shards_path_or_url: "imagenet_sharded/val/imagenet-val-{0000..0009}.tar"
+        num_workers_per_gpu: 12
+    preprocessing:
+        resize_shorter_edge: 256
+        crop_size: 256
+        random_crop: True
+        random_flip: True
+
+
+optimizer:
+    name: adamw 
+    params:
+        learning_rate: 1e-4
+        discriminator_learning_rate: 1e-4
+        beta1: 0.9
+        beta2: 0.999
+        weight_decay: 1e-4
+
+lr_scheduler:
+    scheduler: "cosine"
+    params:
+        learning_rate: ${optimizer.params.learning_rate}
+        warmup_steps: 5_000
+        end_lr: 1e-5
+
+training:
+    gradient_accumulation_steps: 1
+    per_gpu_batch_size: 32
+    mixed_precision: "fp16"
+    enable_tf32: True
+    enable_wandb: True
+    use_ema: True
+    seed: 42
+    max_train_steps: 500_000
+    num_generated_images: 2
+    max_grad_norm: 1.0

configs/training/stage2/titok_l32.yaml

@@ -0,0 +1,79 @@
+experiment:
+    project: "titok_l_32_stage2"
+    name: "titok_l_32_stage2_run1"
+    output_dir: "titok_l_32_stage2_run1"
+    max_train_examples: 1_281_167
+    save_every: 50_000
+    eval_every: 50_000
+    generate_every: 5_000
+    log_every: 50
+    log_grad_norm_every: 1_000
+    resume: True
+    init_weight: "titok_l_32_stage1.bin"
+
+
+model:
+    vq_model:
+        codebook_size: 4096
+        token_size: 12
+        use_l2_norm: True
+        commitment_cost: 0.0
+        # vit arch
+        vit_enc_model_size: "large"
+        vit_dec_model_size: "large"
+        vit_enc_patch_size: 16
+        vit_dec_patch_size: 16
+        num_latent_tokens: 32
+        finetune_decoder: True
+        pretrained_tokenizer_weight: "maskgit-vqgan-imagenet-f16-256.bin"
+
+losses:
+    discriminator_start: 20_000
+    quantizer_weight: 0.0
+    discriminator_factor: 1.0
+    discriminator_weight: 0.01
+    perceptual_loss: "convnext_s"
+    perceptual_weight: 0.1
+    reconstruction_loss: "l2"
+    reconstruction_weight: 1.0
+    lecam_regularization_weight: 0.001
+
+dataset:
+    params:
+        train_shards_path_or_url: "imagenet_sharded/train/imagenet-train-{0000..0252}.tar"
+        eval_shards_path_or_url: "imagenet_sharded/val/imagenet-val-{0000..0009}.tar"
+        num_workers_per_gpu: 12
+    preprocessing:
+        resize_shorter_edge: 256
+        crop_size: 256
+        random_crop: True
+        random_flip: True
+
+
+optimizer:
+    name: adamw 
+    params:
+        learning_rate: 1e-4
+        discriminator_learning_rate: 1e-4
+        beta1: 0.9
+        beta2: 0.999
+        weight_decay: 1e-4
+
+lr_scheduler:
+    scheduler: "cosine"
+    params:
+        learning_rate: ${optimizer.params.learning_rate}
+        warmup_steps: 5_000
+        end_lr: 1e-5
+
+training:
+    gradient_accumulation_steps: 1
+    per_gpu_batch_size: 32
+    mixed_precision: "fp16"
+    enable_tf32: True
+    enable_wandb: True
+    use_ema: True
+    seed: 42
+    max_train_steps: 500_000
+    num_generated_images: 2
+    max_grad_norm: 1.0

configs/training/stage2/titok_s128.yaml

@@ -0,0 +1,79 @@
+experiment:
+    project: "titok_s_128_stage2"
+    name: "titok_s_128_stage2_run1"
+    output_dir: "titok_s_128_stage2_run1"
+    max_train_examples: 1_281_167
+    save_every: 50_000
+    eval_every: 50_000
+    generate_every: 5_000
+    log_every: 50
+    log_grad_norm_every: 1_000
+    resume: True
+    init_weight: "titok_s_128_stage1.bin"
+
+
+model:
+    vq_model:
+        codebook_size: 4096
+        token_size: 12
+        use_l2_norm: True
+        commitment_cost: 0.0
+        # vit arch
+        vit_enc_model_size: "small"
+        vit_dec_model_size: "small"
+        vit_enc_patch_size: 16
+        vit_dec_patch_size: 16
+        num_latent_tokens: 128
+        finetune_decoder: True
+        pretrained_tokenizer_weight: "maskgit-vqgan-imagenet-f16-256.bin"
+    
+losses:
+    discriminator_start: 20_000
+    quantizer_weight: 0.0
+    discriminator_factor: 1.0
+    discriminator_weight: 0.01
+    perceptual_loss: "convnext_s"
+    perceptual_weight: 0.1
+    reconstruction_loss: "l2"
+    reconstruction_weight: 1.0
+    lecam_regularization_weight: 0.001
+
+dataset:
+    params:
+        train_shards_path_or_url: "imagenet_sharded/train/imagenet-train-{0000..0252}.tar"
+        eval_shards_path_or_url: "imagenet_sharded/val/imagenet-val-{0000..0009}.tar"
+        num_workers_per_gpu: 12
+    preprocessing:
+        resize_shorter_edge: 256
+        crop_size: 256
+        random_crop: True
+        random_flip: True
+
+
+optimizer:
+    name: adamw 
+    params:
+        learning_rate: 1e-4
+        discriminator_learning_rate: 1e-4
+        beta1: 0.9
+        beta2: 0.999
+        weight_decay: 1e-4
+
+lr_scheduler:
+    scheduler: "cosine"
+    params:
+        learning_rate: ${optimizer.params.learning_rate}
+        warmup_steps: 5_000
+        end_lr: 1e-5
+
+training:
+    gradient_accumulation_steps: 1
+    per_gpu_batch_size: 32
+    mixed_precision: "fp16"
+    enable_tf32: True
+    enable_wandb: True
+    use_ema: True
+    seed: 42
+    max_train_steps: 500_000
+    num_generated_images: 2
+    max_grad_norm: 1.0

data/__init__.py

@@ -0,0 +1 @@
+from .webdataset_reader import SimpleImageDataset, PretoeknizedDataSetJSONL

data/convert_imagenet_to_wds.py

@@ -0,0 +1,68 @@
+# coding=utf-8
+# Copyright 2023 The HuggingFace Inc. team.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+# Adapted from https://github.com/webdataset/webdataset-imagenet/blob/main/convert-imagenet.py
+
+import argparse
+import os
+import sys
+import time
+
+import webdataset as wds
+from datasets import load_dataset
+
+
+def convert_imagenet_to_wds(output_dir, max_train_samples_per_shard, max_val_samples_per_shard):
+    assert not os.path.exists(os.path.join(output_dir, "imagenet-train-000000.tar"))
+    assert not os.path.exists(os.path.join(output_dir, "imagenet-val-000000.tar"))
+
+    opat = os.path.join(output_dir, "imagenet-train-%06d.tar")
+    output = wds.ShardWriter(opat, maxcount=max_train_samples_per_shard)
+    dataset = load_dataset("imagenet-1k", streaming=True, split="train", use_auth_token=True)
+    now = time.time()
+    for i, example in enumerate(dataset):
+        if i % max_train_samples_per_shard == 0:
+            print(i, file=sys.stderr)
+        img, label = example["image"], example["label"]
+        output.write({"__key__": "%08d" % i, "jpg": img.convert("RGB"), "cls": label})
+    output.close()
+    time_taken = time.time() - now
+    print(f"Wrote {i+1} train examples in {time_taken // 3600} hours.")
+
+    opat = os.path.join(output_dir, "imagenet-val-%06d.tar")
+    output = wds.ShardWriter(opat, maxcount=max_val_samples_per_shard)
+    dataset = load_dataset("imagenet-1k", streaming=True, split="validation", use_auth_token=True)
+    now = time.time()
+    for i, example in enumerate(dataset):
+        if i % max_val_samples_per_shard == 0:
+            print(i, file=sys.stderr)
+        img, label = example["image"], example["label"]
+        output.write({"__key__": "%08d" % i, "jpg": img.convert("RGB"), "cls": label})
+    output.close()
+    time_taken = time.time() - now
+    print(f"Wrote {i+1} val examples in {time_taken // 60} min.")
+
+
+if __name__ == "__main__":
+    # create parase object
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--output_dir", type=str, required=True, help="Path to the output directory.")
+    parser.add_argument("--max_train_samples_per_shard", type=int, default=4000, help="Path to the output directory.")
+    parser.add_argument("--max_val_samples_per_shard", type=int, default=1000, help="Path to the output directory.")
+    args = parser.parse_args()
+
+    # create output directory
+    os.makedirs(args.output_dir, exist_ok=True)
+    convert_imagenet_to_wds(args.output_dir, args.max_train_samples_per_shard, args.max_val_samples_per_shard)

data/webdataset_reader.py

@@ -0,0 +1,227 @@
+"""This file contains the definition of data loader using webdataset.
+
+This file may have been modified by Bytedance Ltd. and/or its affiliates (“Bytedance's Modifications”).
+All Bytedance's Modifications are Copyright (year) Bytedance Ltd. and/or its affiliates. 
+
+Reference:
+    https://github.com/mlfoundations/open_clip/blob/main/src/training/data.py
+    https://github.com/huggingface/open-muse/blob/main/training/data.py
+"""
+
+import math
+from typing import List, Union, Text
+import webdataset as wds
+import torch
+from torch.utils.data import default_collate
+from torchvision import transforms
+from torch.utils.data import Dataset
+import linecache
+import json
+
+
+def filter_keys(key_set):
+    def _f(dictionary):
+        return {k: v for k, v in dictionary.items() if k in key_set}
+
+    return _f
+
+
+class ImageTransform:
+    def __init__(self,
+                 resize_shorter_edge: int = 256,
+                 crop_size: int = 256,
+                 random_crop: bool = True,
+                 random_flip: bool = True,
+                 normalize_mean: List[float] = [0., 0., 0.],
+                 normalize_std: List[float] = [1., 1., 1.]):
+        """Initializes the WebDatasetReader with specified augmentation parameters.
+
+        Args:
+            resize_shorter_edge: An integer, the shorter edge size to resize the input image to.
+            crop_size: An integer, the size to crop the input image to.
+            random_crop: A boolean, whether to use random crop augmentation during training.
+            random_flip: A boolean, whether to use random flipping augmentation during training.
+            normalize_mean: A list of float, the normalization mean used to normalize the image tensor.
+            normalize_std: A list of float, the normalization std used to normalize the image tensor.
+        
+        Raises:
+            NotImplementedError: If the interpolation mode is not one of ["bicubic", "bilinear"].
+        """
+        train_transform = []
+        interpolation = transforms.InterpolationMode.BICUBIC
+
+        train_transform.append(
+            transforms.Resize(resize_shorter_edge, interpolation=interpolation, antialias=True))
+        if random_crop:
+            train_transform.append(transforms.RandomCrop(crop_size))
+        else:
+            train_transform.append(transforms.CenterCrop(crop_size))
+        if random_flip:
+            train_transform.append(transforms.RandomHorizontalFlip())
+        train_transform.append(transforms.ToTensor())
+        # normalize_mean = [0, 0, 0] and normalize_std = [1, 1, 1] will normalize images into [0, 1],
+        # normalize_mean = [0.5, 0.5, 0.5] and normalize_std = [0.5, 0.5, 0.5] will normalize images into [-1, 1].
+        train_transform.append(transforms.Normalize(normalize_mean, normalize_std))
+
+        self.train_transform = transforms.Compose(train_transform)
+        self.eval_transform = transforms.Compose(
+            [
+                # Note that we always resize to crop_size during eval to ensure the results
+                # can be compared against reference numbers on ImageNet etc.
+                transforms.Resize(crop_size, interpolation=interpolation, antialias=True),
+                transforms.CenterCrop(crop_size),
+                transforms.ToTensor(),
+                transforms.Normalize(normalize_mean, normalize_std)
+            ]
+        )
+        print(f"self.train_transform: {self.train_transform}")
+        print(f"self.eval_transform: {self.eval_transform}")
+
+
+class SimpleImageDataset:
+    def __init__(
+        self,
+        train_shards_path: Union[Text, List[Text]],
+        eval_shards_path: Union[Text, List[Text]],
+        num_train_examples: int,
+        per_gpu_batch_size: int,
+        global_batch_size: int,
+        num_workers_per_gpu: int = 12,
+        resize_shorter_edge: int = 256,
+        crop_size: int = 256,
+        random_crop = True,
+        random_flip = True,
+        normalize_mean: List[float] = [0., 0., 0.],
+        normalize_std: List[float] = [1., 1., 1.],
+    ):
+        """Initializes the WebDatasetReader class.
+
+        Args:
+            train_shards_path: A string or list of string, path to the training data shards in webdataset format.
+            eval_shards_path: A string or list of string, path to the evaluation data shards in webdataset format.
+            num_train_examples: An integer, total number of training examples.
+            per_gpu_batch_size: An integer, number of examples per GPU batch.
+            global_batch_size: An integer, total number of examples in a batch across all GPUs.
+            num_workers_per_gpu: An integer, number of workers per GPU.
+            resize_shorter_edge: An integer, the shorter edge size to resize the input image to.
+            crop_size: An integer, the size to crop the input image to.
+            random_crop: A boolean, whether to use random crop augmentation during training.
+            random_flip: A boolean, whether to use random flipping augmentation during training.
+            normalize_mean: A list of float, the normalization mean used to normalize the image tensor.
+            normalize_std: A list of float, the normalization std used to normalize the image tensor.
+        """
+        transform = ImageTransform(
+            resize_shorter_edge, crop_size, random_crop, random_flip,
+            normalize_mean, normalize_std)
+
+        train_processing_pipeline = [
+            wds.decode(wds.autodecode.ImageHandler("pil", extensions=["webp", "png", "jpg", "jpeg"])),
+            wds.rename(
+                image="jpg;png;jpeg;webp",
+                class_id="cls",
+                handler=wds.warn_and_continue,
+                ),
+            wds.map(filter_keys(set(["image", "class_id", "filename"]))),
+            wds.map_dict(
+                image=transform.train_transform,
+                class_id=lambda x: int(x),
+                handler=wds.warn_and_continue,
+            ),
+        ]
+
+        test_processing_pipeline = [
+            wds.decode(wds.autodecode.ImageHandler("pil", extensions=["webp", "png", "jpg", "jpeg"])),
+            wds.rename(
+                image="jpg;png;jpeg;webp",
+                class_id="cls",
+                handler=wds.warn_and_continue,
+                ),
+            wds.map(filter_keys(set(["image", "class_id", "filename"]))),
+            wds.map_dict(
+                image=transform.eval_transform,
+                class_id=lambda x: int(x),
+                handler=wds.warn_and_continue,
+            ),
+        ]
+
+        # Create train dataset and loader.
+        pipeline = [
+            wds.ResampledShards(train_shards_path),
+            wds.tarfile_to_samples(handler=wds.warn_and_continue),
+            wds.shuffle(bufsize=5000,
+                        initial=1000),
+            *train_processing_pipeline,
+            wds.batched(per_gpu_batch_size, partial=False, collation_fn=default_collate),
+        ]
+
+        num_batches = math.ceil(num_train_examples / global_batch_size)
+        num_worker_batches = math.ceil(num_train_examples / 
+            (global_batch_size * num_workers_per_gpu))
+        num_batches = num_worker_batches * num_workers_per_gpu
+        num_samples = num_batches * global_batch_size
+
+        # Each worker is iterating over the complete dataset.
+        self._train_dataset = wds.DataPipeline(*pipeline).with_epoch(num_worker_batches)
+        self._train_dataloader = wds.WebLoader(
+            self._train_dataset,
+            batch_size=None,
+            shuffle=False,
+            num_workers=num_workers_per_gpu,
+            pin_memory=True,
+            persistent_workers=True,
+        )
+        # Add meta-data to dataloader instance for convenience.
+        self._train_dataloader.num_batches = num_batches
+        self._train_dataloader.num_samples = num_samples
+
+        # Create eval dataset and loader.
+        pipeline = [
+            wds.SimpleShardList(eval_shards_path),
+            wds.split_by_worker,
+            wds.tarfile_to_samples(handler=wds.ignore_and_continue),
+            *test_processing_pipeline,
+            wds.batched(per_gpu_batch_size, partial=True, collation_fn=default_collate),
+        ]
+        self._eval_dataset = wds.DataPipeline(*pipeline)
+        self._eval_dataloader = wds.WebLoader(
+            self._eval_dataset,
+            batch_size=None,
+            shuffle=False,
+            num_workers=num_workers_per_gpu,
+            pin_memory=True,
+            persistent_workers=True,
+        )
+
+    @property
+    def train_dataset(self):
+        return self._train_dataset
+
+    @property
+    def train_dataloader(self):
+        return self._train_dataloader
+
+    @property
+    def eval_dataset(self):
+        return self._eval_dataset
+
+    @property
+    def eval_dataloader(self):
+        return self._eval_dataloader
+    
+
+class PretoeknizedDataSetJSONL(Dataset):
+    def __init__(self, data_path):
+        super().__init__()
+        self.jsonl_file = data_path
+        self.num_lines = sum(1 for _ in open(self.jsonl_file))
+        # Ensure the file is cached
+        linecache.checkcache(self.jsonl_file)
+        print("Number of data:", self.num_lines)
+
+    def __len__(self):
+        return self.num_lines
+
+    def __getitem__(self, idx):
+        line = linecache.getline(self.jsonl_file, idx + 1).strip()
+        data = json.loads(line)
+        return torch.tensor(data["class_id"]), torch.tensor(data["tokens"])

demo_util.py

@@ -0,0 +1,108 @@
+"""Demo file for sampling images from TiTok.
+
+Copyright (2024) Bytedance Ltd. and/or its affiliates
+
+Licensed under the Apache License, Version 2.0 (the "License"); 
+you may not use this file except in compliance with the License. 
+You may obtain a copy of the License at 
+
+    http://www.apache.org/licenses/LICENSE-2.0 
+
+Unless required by applicable law or agreed to in writing, software 
+distributed under the License is distributed on an "AS IS" BASIS, 
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 
+See the License for the specific language governing permissions and 
+limitations under the License. 
+"""
+
+
+import torch
+
+from omegaconf import OmegaConf
+from modeling.titok import TiTok
+from modeling.maskgit import ImageBert, UViTBert
+from modeling.rar import RAR
+
+
+def get_config_cli():
+    cli_conf = OmegaConf.from_cli()
+
+    yaml_conf = OmegaConf.load(cli_conf.config)
+    conf = OmegaConf.merge(yaml_conf, cli_conf)
+
+    return conf
+
+def get_config(config_path):
+    conf = OmegaConf.load(config_path)
+    return conf
+
+def get_titok_tokenizer(config):
+    tokenizer = TiTok(config)
+    tokenizer.load_state_dict(torch.load(config.experiment.tokenizer_checkpoint, map_location="cpu"))
+    tokenizer.eval()
+    tokenizer.requires_grad_(False)
+    return tokenizer
+
+def get_titok_generator(config):
+    if config.model.generator.model_type == "ViT":
+        model_cls = ImageBert
+    elif config.model.generator.model_type == "UViT":
+        model_cls = UViTBert
+    else:
+        raise ValueError(f"Unsupported model type {config.model.generator.model_type}")
+    generator = model_cls(config)
+    generator.load_state_dict(torch.load(config.experiment.generator_checkpoint, map_location="cpu"))
+    generator.eval()
+    generator.requires_grad_(False)
+    return generator
+
+def get_rar_generator(config):
+    model_cls = RAR
+    generator = model_cls(config)
+    generator.load_state_dict(torch.load(config.experiment.generator_checkpoint, map_location="cpu"))
+    generator.eval()
+    generator.requires_grad_(False)
+    generator.set_random_ratio(0)
+    return generator
+
+
+@torch.no_grad()
+def sample_fn(generator,
+              tokenizer,
+              labels=None,
+              guidance_scale=3.0,
+              guidance_decay="constant",
+              guidance_scale_pow=3.0,
+              randomize_temperature=2.0,
+              softmax_temperature_annealing=False,
+              num_sample_steps=8,
+              device="cuda",
+              return_tensor=False):
+    generator.eval()
+    tokenizer.eval()
+    if labels is None:
+        # goldfish, chicken, tiger, cat, hourglass, ship, dog, race car, airliner, teddy bear, random
+        labels = [1, 7, 282, 604, 724, 179, 751, 404, 850, torch.randint(0, 999, size=(1,))]
+
+    if not isinstance(labels, torch.Tensor):
+        labels = torch.LongTensor(labels).to(device)
+
+    generated_tokens = generator.generate(
+        condition=labels,
+        guidance_scale=guidance_scale,
+        guidance_decay=guidance_decay,
+        guidance_scale_pow=guidance_scale_pow,
+        randomize_temperature=randomize_temperature,
+        softmax_temperature_annealing=softmax_temperature_annealing,
+        num_sample_steps=num_sample_steps)
+    
+    generated_image = tokenizer.decode_tokens(
+        generated_tokens.view(generated_tokens.shape[0], -1)
+    )
+    if return_tensor:
+        return generated_image
+
+    generated_image = torch.clamp(generated_image, 0.0, 1.0)
+    generated_image = (generated_image * 255.0).permute(0, 2, 3, 1).to("cpu", dtype=torch.uint8).numpy()
+
+    return generated_image

evaluator/__init__.py

@@ -0,0 +1 @@
+from .evaluator import VQGANEvaluator

evaluator/evaluator.py

@@ -0,0 +1,245 @@
+"""This file contains a class to evalute the reconstruction results.
+
+Copyright (2024) Bytedance Ltd. and/or its affiliates
+
+Licensed under the Apache License, Version 2.0 (the "License"); 
+you may not use this file except in compliance with the License. 
+You may obtain a copy of the License at 
+
+    http://www.apache.org/licenses/LICENSE-2.0 
+
+Unless required by applicable law or agreed to in writing, software 
+distributed under the License is distributed on an "AS IS" BASIS, 
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 
+See the License for the specific language governing permissions and 
+limitations under the License. 
+"""
+
+import warnings
+
+from typing import Sequence, Optional, Mapping, Text
+import numpy as np
+from scipy import linalg
+import torch
+import torch.nn.functional as F
+
+from .inception import get_inception_model
+
+
+def get_covariance(sigma: torch.Tensor, total: torch.Tensor, num_examples: int) -> torch.Tensor:
+    """Computes covariance of the input tensor.
+
+    Args:
+        sigma: A torch.Tensor, sum of outer products of input features.
+        total: A torch.Tensor, sum of all input features.
+        num_examples: An integer, number of examples in the input tensor.
+    Returns:
+        A torch.Tensor, covariance of the input tensor.
+    """
+    if num_examples == 0:
+        return torch.zeros_like(sigma)
+
+    sub_matrix = torch.outer(total, total)
+    sub_matrix = sub_matrix / num_examples
+
+    return (sigma - sub_matrix) / (num_examples - 1)
+
+
+class VQGANEvaluator:
+    def __init__(
+        self,
+        device,
+        enable_rfid: bool = True,
+        enable_inception_score: bool = True,
+        enable_codebook_usage_measure: bool = False,
+        enable_codebook_entropy_measure: bool = False,
+        num_codebook_entries: int = 1024
+    ):
+        """Initializes VQGAN Evaluator.
+
+        Args:
+            device: The device to use for evaluation.
+            enable_rfid: A boolean, whether enabling rFID score.
+            enable_inception_score: A boolean, whether enabling Inception Score.
+            enable_codebook_usage_measure: A boolean, whether enabling codebook usage measure.
+            enable_codebook_entropy_measure: A boolean, whether enabling codebook entropy measure.
+            num_codebook_entries: An integer, the number of codebook entries.
+        """
+        self._device = device
+
+        self._enable_rfid = enable_rfid
+        self._enable_inception_score = enable_inception_score
+        self._enable_codebook_usage_measure = enable_codebook_usage_measure
+        self._enable_codebook_entropy_measure = enable_codebook_entropy_measure
+        self._num_codebook_entries = num_codebook_entries
+
+        # Variables related to Inception score and rFID.
+        self._inception_model = None
+        self._is_num_features = 0
+        self._rfid_num_features = 0
+        if self._enable_inception_score or self._enable_rfid:
+            self._rfid_num_features = 2048
+            self._is_num_features = 1008
+            self._inception_model = get_inception_model().to(self._device)
+            self._inception_model.eval()
+        self._is_eps = 1e-16
+        self._rfid_eps = 1e-6
+
+        self.reset_metrics()
+
+    def reset_metrics(self):
+        """Resets all metrics."""
+        self._num_examples = 0
+        self._num_updates = 0
+
+        self._is_prob_total = torch.zeros(
+            self._is_num_features, dtype=torch.float64, device=self._device
+        )
+        self._is_total_kl_d = torch.zeros(
+            self._is_num_features, dtype=torch.float64, device=self._device
+        )
+        self._rfid_real_sigma = torch.zeros(
+            (self._rfid_num_features, self._rfid_num_features),
+            dtype=torch.float64, device=self._device
+        )
+        self._rfid_real_total = torch.zeros(
+            self._rfid_num_features, dtype=torch.float64, device=self._device
+        )
+        self._rfid_fake_sigma = torch.zeros(
+            (self._rfid_num_features, self._rfid_num_features),
+            dtype=torch.float64, device=self._device
+        )
+        self._rfid_fake_total = torch.zeros(
+            self._rfid_num_features, dtype=torch.float64, device=self._device
+        )
+
+        self._set_of_codebook_indices = set()
+        self._codebook_frequencies = torch.zeros((self._num_codebook_entries), dtype=torch.float64, device=self._device)
+
+    def update(
+        self,
+        real_images: torch.Tensor,
+        fake_images: torch.Tensor,
+        codebook_indices: Optional[torch.Tensor] = None
+    ):
+        """Updates the metrics with the given images.
+
+        Args:
+            real_images: A torch.Tensor, the real images.
+            fake_images: A torch.Tensor, the fake images.
+            codebook_indices: A torch.Tensor, the indices of the codebooks for each image.
+
+        Raises:
+            ValueError: If the fake images is not in RGB (3 channel).
+            ValueError: If the fake and real images have different shape.
+        """
+
+        batch_size = real_images.shape[0]
+        dim = tuple(range(1, real_images.ndim))
+        self._num_examples += batch_size
+        self._num_updates += 1
+
+        if self._enable_inception_score or self._enable_rfid:
+            # Quantize to uint8 as a real image.
+            fake_inception_images = (fake_images * 255).to(torch.uint8)
+            features_fake = self._inception_model(fake_inception_images)
+            inception_logits_fake = features_fake["logits_unbiased"]
+            inception_probabilities_fake = F.softmax(inception_logits_fake, dim=-1)
+        
+        if self._enable_inception_score:
+            probabiliies_sum = torch.sum(inception_probabilities_fake, 0, dtype=torch.float64)
+
+            log_prob = torch.log(inception_probabilities_fake + self._is_eps)
+            if log_prob.dtype != inception_probabilities_fake.dtype:
+                log_prob = log_prob.to(inception_probabilities_fake)
+            kl_sum = torch.sum(inception_probabilities_fake * log_prob, 0, dtype=torch.float64)
+
+            self._is_prob_total += probabiliies_sum
+            self._is_total_kl_d += kl_sum
+
+        if self._enable_rfid:
+            real_inception_images = (real_images * 255).to(torch.uint8)
+            features_real = self._inception_model(real_inception_images)
+            if (features_real['2048'].shape[0] != features_fake['2048'].shape[0] or
+                features_real['2048'].shape[1] != features_fake['2048'].shape[1]):
+                raise ValueError(f"Number of features should be equal for real and fake.")
+
+            for f_real, f_fake in zip(features_real['2048'], features_fake['2048']):
+                self._rfid_real_total += f_real
+                self._rfid_fake_total += f_fake
+
+                self._rfid_real_sigma += torch.outer(f_real, f_real)
+                self._rfid_fake_sigma += torch.outer(f_fake, f_fake)
+
+        if self._enable_codebook_usage_measure:
+            self._set_of_codebook_indices |= set(torch.unique(codebook_indices, sorted=False).tolist())
+
+        if self._enable_codebook_entropy_measure:
+            entries, counts = torch.unique(codebook_indices, sorted=False, return_counts=True)
+            self._codebook_frequencies.index_add_(0, entries.int(), counts.double())
+
+
+    def result(self) -> Mapping[Text, torch.Tensor]:
+        """Returns the evaluation result."""
+        eval_score = {}
+
+        if self._num_examples < 1:
+            raise ValueError("No examples to evaluate.")
+        
+        if self._enable_inception_score:
+            mean_probs = self._is_prob_total / self._num_examples
+            log_mean_probs = torch.log(mean_probs + self._is_eps)
+            if log_mean_probs.dtype != self._is_prob_total.dtype:
+                log_mean_probs = log_mean_probs.to(self._is_prob_total)
+            excess_entropy = self._is_prob_total * log_mean_probs
+            avg_kl_d = torch.sum(self._is_total_kl_d - excess_entropy) / self._num_examples
+
+            inception_score = torch.exp(avg_kl_d).item()
+            eval_score["InceptionScore"] = inception_score
+
+        if self._enable_rfid:
+            mu_real = self._rfid_real_total / self._num_examples
+            mu_fake = self._rfid_fake_total / self._num_examples
+            sigma_real = get_covariance(self._rfid_real_sigma, self._rfid_real_total, self._num_examples)
+            sigma_fake = get_covariance(self._rfid_fake_sigma, self._rfid_fake_total, self._num_examples)
+
+            mu_real, mu_fake = mu_real.cpu(), mu_fake.cpu()
+            sigma_real, sigma_fake = sigma_real.cpu(), sigma_fake.cpu()
+
+            diff = mu_real - mu_fake
+
+            # Product might be almost singular.
+            covmean, _ = linalg.sqrtm(sigma_real.mm(sigma_fake).numpy(), disp=False)
+            # Numerical error might give slight imaginary component.
+            if np.iscomplexobj(covmean):
+                if not np.allclose(np.diagonal(covmean).imag, 0, atol=1e-3):
+                    m = np.max(np.abs(covmean.imag))
+                    raise ValueError("Imaginary component {}".format(m))
+                covmean = covmean.real
+
+            tr_covmean = np.trace(covmean)
+
+            if not np.isfinite(covmean).all():
+                tr_covmean = np.sum(np.sqrt((
+                    (np.diag(sigma_real) * self._rfid_eps) * (np.diag(sigma_fake) * self._rfid_eps))
+                    / (self._rfid_eps * self._rfid_eps)
+                ))
+
+            rfid = float(diff.dot(diff).item() + torch.trace(sigma_real) + torch.trace(sigma_fake) 
+                - 2 * tr_covmean
+            )
+            if torch.isnan(torch.tensor(rfid)) or torch.isinf(torch.tensor(rfid)):
+                warnings.warn("The product of covariance of train and test features is out of bounds.")
+
+            eval_score["rFID"] = rfid
+
+        if self._enable_codebook_usage_measure:
+            usage = float(len(self._set_of_codebook_indices)) / self._num_codebook_entries
+            eval_score["CodebookUsage"] = usage
+
+        if self._enable_codebook_entropy_measure:
+            probs = self._codebook_frequencies / self._codebook_frequencies.sum()
+            entropy = (-torch.log2(probs + 1e-8) * probs).sum()
+            eval_score["CodebookEntropy"] = entropy
+
+        return eval_score

evaluator/inception.py

@@ -0,0 +1,231 @@
+"""This file is for Inception model borrowed from torch metrics / fidelity.
+
+This file may have been modified by Bytedance Ltd. and/or its affiliates (“Bytedance's Modifications”).
+All Bytedance's Modifications are Copyright (year) Bytedance Ltd. and/or its affiliates. 
+
+Reference:
+    https://github.com/mseitzer/pytorch-fid/blob/master/src/pytorch_fid/inception.py
+"""
+# Copyright The Lightning team.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+import torch
+import torch.nn.functional as F
+
+from torch_fidelity.feature_extractor_base import FeatureExtractorBase
+from torch_fidelity.helpers import vassert
+from torch_fidelity.feature_extractor_inceptionv3 import BasicConv2d, InceptionA, InceptionB, InceptionC, InceptionD, InceptionE_1, InceptionE_2
+from torch_fidelity.interpolate_compat_tensorflow import interpolate_bilinear_2d_like_tensorflow1x
+
+try:
+    from torchvision.models.utils import load_state_dict_from_url
+except ImportError:
+    from torch.utils.model_zoo import load_url as load_state_dict_from_url
+
+
+# Note: Compared shasum and models should be the same.
+FID_WEIGHTS_URL = 'https://github.com/mseitzer/pytorch-fid/releases/download/fid_weights/pt_inception-2015-12-05-6726825d.pth'
+
+class FeatureExtractorInceptionV3(FeatureExtractorBase):
+    INPUT_IMAGE_SIZE = 299
+
+    def __init__(
+            self,
+            name,
+            features_list,
+            **kwargs,
+    ):
+        """
+        InceptionV3 feature extractor for 2D RGB 24bit images.
+
+        Args:
+
+            name (str): Unique name of the feature extractor, must be the same as used in
+                :func:`register_feature_extractor`.
+
+            features_list (list): A list of the requested feature names, which will be produced for each input. This
+                feature extractor provides the following features:
+
+                - '64'
+                - '192'
+                - '768'
+                - '2048'
+                - 'logits_unbiased'
+                - 'logits'
+
+        """
+        super(FeatureExtractorInceptionV3, self).__init__(name, features_list)
+        self.feature_extractor_internal_dtype = torch.float64
+
+        self.Conv2d_1a_3x3 = BasicConv2d(3, 32, kernel_size=3, stride=2)
+        self.Conv2d_2a_3x3 = BasicConv2d(32, 32, kernel_size=3)
+        self.Conv2d_2b_3x3 = BasicConv2d(32, 64, kernel_size=3, padding=1)
+        self.MaxPool_1 = torch.nn.MaxPool2d(kernel_size=3, stride=2)
+
+        self.Conv2d_3b_1x1 = BasicConv2d(64, 80, kernel_size=1)
+        self.Conv2d_4a_3x3 = BasicConv2d(80, 192, kernel_size=3)
+        self.MaxPool_2 = torch.nn.MaxPool2d(kernel_size=3, stride=2)
+
+        self.Mixed_5b = InceptionA(192, pool_features=32)
+        self.Mixed_5c = InceptionA(256, pool_features=64)
+        self.Mixed_5d = InceptionA(288, pool_features=64)
+        self.Mixed_6a = InceptionB(288)
+        self.Mixed_6b = InceptionC(768, channels_7x7=128)
+        self.Mixed_6c = InceptionC(768, channels_7x7=160)
+        self.Mixed_6d = InceptionC(768, channels_7x7=160)
+        self.Mixed_6e = InceptionC(768, channels_7x7=192)
+
+        self.Mixed_7a = InceptionD(768)
+        self.Mixed_7b = InceptionE_1(1280)
+        self.Mixed_7c = InceptionE_2(2048)
+        self.AvgPool = torch.nn.AdaptiveAvgPool2d(output_size=(1, 1))
+
+        self.fc = torch.nn.Linear(2048, 1008)
+
+        state_dict = load_state_dict_from_url(FID_WEIGHTS_URL, progress=True)
+        #state_dict = torch.load(FID_WEIGHTS_URL, map_location='cpu')
+        self.load_state_dict(state_dict)
+
+        self.to(self.feature_extractor_internal_dtype)
+        self.requires_grad_(False)
+        self.eval()
+
+    def forward(self, x):
+        vassert(torch.is_tensor(x) and x.dtype == torch.uint8, 'Expecting image as torch.Tensor with dtype=torch.uint8')
+        vassert(x.dim() == 4 and x.shape[1] == 3, f'Input is not Bx3xHxW: {x.shape}')
+        features = {}
+        remaining_features = self.features_list.copy()
+
+        x = x.to(self.feature_extractor_internal_dtype)
+        # N x 3 x ? x ?
+
+        x = interpolate_bilinear_2d_like_tensorflow1x(
+            x,
+            size=(self.INPUT_IMAGE_SIZE, self.INPUT_IMAGE_SIZE),
+            align_corners=False,
+        )
+        # N x 3 x 299 x 299
+
+        # x = (x - 128) * torch.tensor(0.0078125, dtype=torch.float32, device=x.device)  # really happening in graph
+        x = (x - 128) / 128  # but this gives bit-exact output _of this step_ too
+        # N x 3 x 299 x 299
+
+        x = self.Conv2d_1a_3x3(x)
+        # N x 32 x 149 x 149
+        x = self.Conv2d_2a_3x3(x)
+        # N x 32 x 147 x 147
+        x = self.Conv2d_2b_3x3(x)
+        # N x 64 x 147 x 147
+        x = self.MaxPool_1(x)
+        # N x 64 x 73 x 73
+
+        if '64' in remaining_features:
+            features['64'] = F.adaptive_avg_pool2d(x, output_size=(1, 1)).squeeze(-1).squeeze(-1)
+            remaining_features.remove('64')
+            if len(remaining_features) == 0:
+                return features
+
+        x = self.Conv2d_3b_1x1(x)
+        # N x 80 x 73 x 73
+        x = self.Conv2d_4a_3x3(x)
+        # N x 192 x 71 x 71
+        x = self.MaxPool_2(x)
+        # N x 192 x 35 x 35
+
+        if '192' in remaining_features:
+            features['192'] = F.adaptive_avg_pool2d(x, output_size=(1, 1)).squeeze(-1).squeeze(-1)
+            remaining_features.remove('192')
+            if len(remaining_features) == 0:
+                return features
+
+        x = self.Mixed_5b(x)
+        # N x 256 x 35 x 35
+        x = self.Mixed_5c(x)
+        # N x 288 x 35 x 35
+        x = self.Mixed_5d(x)
+        # N x 288 x 35 x 35
+        x = self.Mixed_6a(x)
+        # N x 768 x 17 x 17
+        x = self.Mixed_6b(x)
+        # N x 768 x 17 x 17
+        x = self.Mixed_6c(x)
+        # N x 768 x 17 x 17
+        x = self.Mixed_6d(x)
+        # N x 768 x 17 x 17
+        x = self.Mixed_6e(x)
+        # N x 768 x 17 x 17
+
+        if '768' in remaining_features:
+            features['768'] = F.adaptive_avg_pool2d(x, output_size=(1, 1)).squeeze(-1).squeeze(-1).to(torch.float32)
+            remaining_features.remove('768')
+            if len(remaining_features) == 0:
+                return features
+
+        x = self.Mixed_7a(x)
+        # N x 1280 x 8 x 8
+        x = self.Mixed_7b(x)
+        # N x 2048 x 8 x 8
+        x = self.Mixed_7c(x)
+        # N x 2048 x 8 x 8
+        x = self.AvgPool(x)
+        # N x 2048 x 1 x 1
+
+        x = torch.flatten(x, 1)
+        # N x 2048
+
+        if '2048' in remaining_features:
+            features['2048'] = x
+            remaining_features.remove('2048')
+            if len(remaining_features) == 0:
+                return features
+
+        if 'logits_unbiased' in remaining_features:
+            x = x.mm(self.fc.weight.T)
+            # N x 1008 (num_classes)
+            features['logits_unbiased'] = x
+            remaining_features.remove('logits_unbiased')
+            if len(remaining_features) == 0:
+                return features
+
+            x = x + self.fc.bias.unsqueeze(0)
+        else:
+            x = self.fc(x)
+            # N x 1008 (num_classes)
+
+        features['logits'] = x
+        return features
+
+    @staticmethod
+    def get_provided_features_list():
+        return '64', '192', '768', '2048', 'logits_unbiased', 'logits'
+
+    @staticmethod
+    def get_default_feature_layer_for_metric(metric):
+        return {
+            'isc': 'logits_unbiased',
+            'fid': '2048',
+            'kid': '2048',
+            'prc': '2048',
+        }[metric]
+
+    @staticmethod
+    def can_be_compiled():
+        return True
+
+    @staticmethod
+    def get_dummy_input_for_compile():
+        return (torch.rand([1, 3, 4, 4]) * 255).to(torch.uint8)
+
+def get_inception_model():
+    model = FeatureExtractorInceptionV3("inception_model", ["2048", "logits_unbiased"])
+    return model

imagenet_classes.py

@@ -0,0 +1,1001 @@
+imagenet_idx2classname = {
+ 0: 'tench, Tinca tinca',
+ 1: 'goldfish, Carassius auratus',
+ 2: 'great white shark, white shark, man-eater, man-eating shark, Carcharodon carcharias',
+ 3: 'tiger shark, Galeocerdo cuvieri',
+ 4: 'hammerhead, hammerhead shark',
+ 5: 'electric ray, crampfish, numbfish, torpedo',
+ 6: 'stingray',
+ 7: 'cock',
+ 8: 'hen',
+ 9: 'ostrich, Struthio camelus',
+ 10: 'brambling, Fringilla montifringilla',
+ 11: 'goldfinch, Carduelis carduelis',
+ 12: 'house finch, linnet, Carpodacus mexicanus',
+ 13: 'junco, snowbird',
+ 14: 'indigo bunting, indigo finch, indigo bird, Passerina cyanea',
+ 15: 'robin, American robin, Turdus migratorius',
+ 16: 'bulbul',
+ 17: 'jay',
+ 18: 'magpie',
+ 19: 'chickadee',
+ 20: 'water ouzel, dipper',
+ 21: 'kite',
+ 22: 'bald eagle, American eagle, Haliaeetus leucocephalus',
+ 23: 'vulture',
+ 24: 'great grey owl, great gray owl, Strix nebulosa',
+ 25: 'European fire salamander, Salamandra salamandra',
+ 26: 'common newt, Triturus vulgaris',
+ 27: 'eft',
+ 28: 'spotted salamander, Ambystoma maculatum',
+ 29: 'axolotl, mud puppy, Ambystoma mexicanum',
+ 30: 'bullfrog, Rana catesbeiana',
+ 31: 'tree frog, tree-frog',
+ 32: 'tailed frog, bell toad, ribbed toad, tailed toad, Ascaphus trui',
+ 33: 'loggerhead, loggerhead turtle, Caretta caretta',
+ 34: 'leatherback turtle, leatherback, leathery turtle, Dermochelys coriacea',
+ 35: 'mud turtle',
+ 36: 'terrapin',
+ 37: 'box turtle, box tortoise',
+ 38: 'banded gecko',
+ 39: 'common iguana, iguana, Iguana iguana',
+ 40: 'American chameleon, anole, Anolis carolinensis',
+ 41: 'whiptail, whiptail lizard',
+ 42: 'agama',
+ 43: 'frilled lizard, Chlamydosaurus kingi',
+ 44: 'alligator lizard',
+ 45: 'Gila monster, Heloderma suspectum',
+ 46: 'green lizard, Lacerta viridis',
+ 47: 'African chameleon, Chamaeleo chamaeleon',
+ 48: 'Komodo dragon, Komodo lizard, dragon lizard, giant lizard, Varanus komodoensis',
+ 49: 'African crocodile, Nile crocodile, Crocodylus niloticus',
+ 50: 'American alligator, Alligator mississipiensis',
+ 51: 'triceratops',
+ 52: 'thunder snake, worm snake, Carphophis amoenus',
+ 53: 'ringneck snake, ring-necked snake, ring snake',
+ 54: 'hognose snake, puff adder, sand viper',
+ 55: 'green snake, grass snake',
+ 56: 'king snake, kingsnake',
+ 57: 'garter snake, grass snake',
+ 58: 'water snake',
+ 59: 'vine snake',
+ 60: 'night snake, Hypsiglena torquata',
+ 61: 'boa constrictor, Constrictor constrictor',
+ 62: 'rock python, rock snake, Python sebae',
+ 63: 'Indian cobra, Naja naja',
+ 64: 'green mamba',
+ 65: 'sea snake',
+ 66: 'horned viper, cerastes, sand viper, horned asp, Cerastes cornutus',
+ 67: 'diamondback, diamondback rattlesnake, Crotalus adamanteus',
+ 68: 'sidewinder, horned rattlesnake, Crotalus cerastes',
+ 69: 'trilobite',
+ 70: 'harvestman, daddy longlegs, Phalangium opilio',
+ 71: 'scorpion',
+ 72: 'black and gold garden spider, Argiope aurantia',
+ 73: 'barn spider, Araneus cavaticus',
+ 74: 'garden spider, Aranea diademata',
+ 75: 'black widow, Latrodectus mactans',
+ 76: 'tarantula',
+ 77: 'wolf spider, hunting spider',
+ 78: 'tick',
+ 79: 'centipede',
+ 80: 'black grouse',
+ 81: 'ptarmigan',
+ 82: 'ruffed grouse, partridge, Bonasa umbellus',
+ 83: 'prairie chicken, prairie grouse, prairie fowl',
+ 84: 'peacock',
+ 85: 'quail',
+ 86: 'partridge',
+ 87: 'African grey, African gray, Psittacus erithacus',
+ 88: 'macaw',
+ 89: 'sulphur-crested cockatoo, Kakatoe galerita, Cacatua galerita',
+ 90: 'lorikeet',
+ 91: 'coucal',
+ 92: 'bee eater',
+ 93: 'hornbill',
+ 94: 'hummingbird',
+ 95: 'jacamar',
+ 96: 'toucan',
+ 97: 'drake',
+ 98: 'red-breasted merganser, Mergus serrator',
+ 99: 'goose',
+ 100: 'black swan, Cygnus atratus',
+ 101: 'tusker',
+ 102: 'echidna, spiny anteater, anteater',
+ 103: 'platypus, duckbill, duckbilled platypus, duck-billed platypus, Ornithorhynchus anatinus',
+ 104: 'wallaby, brush kangaroo',
+ 105: 'koala, koala bear, kangaroo bear, native bear, Phascolarctos cinereus',
+ 106: 'wombat',
+ 107: 'jellyfish',
+ 108: 'sea anemone, anemone',
+ 109: 'brain coral',
+ 110: 'flatworm, platyhelminth',
+ 111: 'nematode, nematode worm, roundworm',
+ 112: 'conch',
+ 113: 'snail',
+ 114: 'slug',
+ 115: 'sea slug, nudibranch',
+ 116: 'chiton, coat-of-mail shell, sea cradle, polyplacophore',
+ 117: 'chambered nautilus, pearly nautilus, nautilus',
+ 118: 'Dungeness crab, Cancer magister',
+ 119: 'rock crab, Cancer irroratus',
+ 120: 'fiddler crab',
+ 121: 'king crab, Alaska crab, Alaskan king crab, Alaska king crab, Paralithodes camtschatica',
+ 122: 'American lobster, Northern lobster, Maine lobster, Homarus americanus',
+ 123: 'spiny lobster, langouste, rock lobster, crawfish, crayfish, sea crawfish',
+ 124: 'crayfish, crawfish, crawdad, crawdaddy',
+ 125: 'hermit crab',
+ 126: 'isopod',
+ 127: 'white stork, Ciconia ciconia',
+ 128: 'black stork, Ciconia nigra',
+ 129: 'spoonbill',
+ 130: 'flamingo',
+ 131: 'little blue heron, Egretta caerulea',
+ 132: 'American egret, great white heron, Egretta albus',
+ 133: 'bittern',
+ 134: 'crane',
+ 135: 'limpkin, Aramus pictus',
+ 136: 'European gallinule, Porphyrio porphyrio',
+ 137: 'American coot, marsh hen, mud hen, water hen, Fulica americana',
+ 138: 'bustard',
+ 139: 'ruddy turnstone, Arenaria interpres',
+ 140: 'red-backed sandpiper, dunlin, Erolia alpina',
+ 141: 'redshank, Tringa totanus',
+ 142: 'dowitcher',
+ 143: 'oystercatcher, oyster catcher',
+ 144: 'pelican',
+ 145: 'king penguin, Aptenodytes patagonica',
+ 146: 'albatross, mollymawk',
+ 147: 'grey whale, gray whale, devilfish, Eschrichtius gibbosus, Eschrichtius robustus',
+ 148: 'killer whale, killer, orca, grampus, sea wolf, Orcinus orca',
+ 149: 'dugong, Dugong dugon',
+ 150: 'sea lion',
+ 151: 'Chihuahua',
+ 152: 'Japanese spaniel',
+ 153: 'Maltese dog, Maltese terrier, Maltese',
+ 154: 'Pekinese, Pekingese, Peke',
+ 155: 'Shih-Tzu',
+ 156: 'Blenheim spaniel',
+ 157: 'papillon',
+ 158: 'toy terrier',
+ 159: 'Rhodesian ridgeback',
+ 160: 'Afghan hound, Afghan',
+ 161: 'basset, basset hound',
+ 162: 'beagle',
+ 163: 'bloodhound, sleuthhound',
+ 164: 'bluetick',
+ 165: 'black-and-tan coonhound',
+ 166: 'Walker hound, Walker foxhound',
+ 167: 'English foxhound',
+ 168: 'redbone',
+ 169: 'borzoi, Russian wolfhound',
+ 170: 'Irish wolfhound',
+ 171: 'Italian greyhound',
+ 172: 'whippet',
+ 173: 'Ibizan hound, Ibizan Podenco',
+ 174: 'Norwegian elkhound, elkhound',
+ 175: 'otterhound, otter hound',
+ 176: 'Saluki, gazelle hound',
+ 177: 'Scottish deerhound, deerhound',
+ 178: 'Weimaraner',
+ 179: 'Staffordshire bullterrier, Staffordshire bull terrier',
+ 180: 'American Staffordshire terrier, Staffordshire terrier, American pit bull terrier, pit bull terrier',
+ 181: 'Bedlington terrier',
+ 182: 'Border terrier',
+ 183: 'Kerry blue terrier',
+ 184: 'Irish terrier',
+ 185: 'Norfolk terrier',
+ 186: 'Norwich terrier',
+ 187: 'Yorkshire terrier',
+ 188: 'wire-haired fox terrier',
+ 189: 'Lakeland terrier',
+ 190: 'Sealyham terrier, Sealyham',
+ 191: 'Airedale, Airedale terrier',
+ 192: 'cairn, cairn terrier',
+ 193: 'Australian terrier',
+ 194: 'Dandie Dinmont, Dandie Dinmont terrier',
+ 195: 'Boston bull, Boston terrier',
+ 196: 'miniature schnauzer',
+ 197: 'giant schnauzer',
+ 198: 'standard schnauzer',
+ 199: 'Scotch terrier, Scottish terrier, Scottie',
+ 200: 'Tibetan terrier, chrysanthemum dog',
+ 201: 'silky terrier, Sydney silky',
+ 202: 'soft-coated wheaten terrier',
+ 203: 'West Highland white terrier',
+ 204: 'Lhasa, Lhasa apso',
+ 205: 'flat-coated retriever',
+ 206: 'curly-coated retriever',
+ 207: 'golden retriever',
+ 208: 'Labrador retriever',
+ 209: 'Chesapeake Bay retriever',
+ 210: 'German short-haired pointer',
+ 211: 'vizsla, Hungarian pointer',
+ 212: 'English setter',
+ 213: 'Irish setter, red setter',
+ 214: 'Gordon setter',
+ 215: 'Brittany spaniel',
+ 216: 'clumber, clumber spaniel',
+ 217: 'English springer, English springer spaniel',
+ 218: 'Welsh springer spaniel',
+ 219: 'cocker spaniel, English cocker spaniel, cocker',
+ 220: 'Sussex spaniel',
+ 221: 'Irish water spaniel',
+ 222: 'kuvasz',
+ 223: 'schipperke',
+ 224: 'groenendael',
+ 225: 'malinois',
+ 226: 'briard',
+ 227: 'kelpie',
+ 228: 'komondor',
+ 229: 'Old English sheepdog, bobtail',
+ 230: 'Shetland sheepdog, Shetland sheep dog, Shetland',
+ 231: 'collie',
+ 232: 'Border collie',
+ 233: 'Bouvier des Flandres, Bouviers des Flandres',
+ 234: 'Rottweiler',
+ 235: 'German shepherd, German shepherd dog, German police dog, alsatian',
+ 236: 'Doberman, Doberman pinscher',
+ 237: 'miniature pinscher',
+ 238: 'Greater Swiss Mountain dog',
+ 239: 'Bernese mountain dog',
+ 240: 'Appenzeller',
+ 241: 'EntleBucher',
+ 242: 'boxer',
+ 243: 'bull mastiff',
+ 244: 'Tibetan mastiff',
+ 245: 'French bulldog',
+ 246: 'Great Dane',
+ 247: 'Saint Bernard, St Bernard',
+ 248: 'Eskimo dog, husky',
+ 249: 'malamute, malemute, Alaskan malamute',
+ 250: 'Siberian husky',
+ 251: 'dalmatian, coach dog, carriage dog',
+ 252: 'affenpinscher, monkey pinscher, monkey dog',
+ 253: 'basenji',
+ 254: 'pug, pug-dog',
+ 255: 'Leonberg',
+ 256: 'Newfoundland, Newfoundland dog',
+ 257: 'Great Pyrenees',
+ 258: 'Samoyed, Samoyede',
+ 259: 'Pomeranian',
+ 260: 'chow, chow chow',
+ 261: 'keeshond',
+ 262: 'Brabancon griffon',
+ 263: 'Pembroke, Pembroke Welsh corgi',
+ 264: 'Cardigan, Cardigan Welsh corgi',
+ 265: 'toy poodle',
+ 266: 'miniature poodle',
+ 267: 'standard poodle',
+ 268: 'Mexican hairless',
+ 269: 'timber wolf, grey wolf, gray wolf, Canis lupus',
+ 270: 'white wolf, Arctic wolf, Canis lupus tundrarum',
+ 271: 'red wolf, maned wolf, Canis rufus, Canis niger',
+ 272: 'coyote, prairie wolf, brush wolf, Canis latrans',
+ 273: 'dingo, warrigal, warragal, Canis dingo',
+ 274: 'dhole, Cuon alpinus',
+ 275: 'African hunting dog, hyena dog, Cape hunting dog, Lycaon pictus',
+ 276: 'hyena, hyaena',
+ 277: 'red fox, Vulpes vulpes',
+ 278: 'kit fox, Vulpes macrotis',
+ 279: 'Arctic fox, white fox, Alopex lagopus',
+ 280: 'grey fox, gray fox, Urocyon cinereoargenteus',
+ 281: 'tabby, tabby cat',
+ 282: 'tiger cat',
+ 283: 'Persian cat',
+ 284: 'Siamese cat, Siamese',
+ 285: 'Egyptian cat',
+ 286: 'cougar, puma, catamount, mountain lion, painter, panther, Felis concolor',
+ 287: 'lynx, catamount',
+ 288: 'leopard, Panthera pardus',
+ 289: 'snow leopard, ounce, Panthera uncia',
+ 290: 'jaguar, panther, Panthera onca, Felis onca',
+ 291: 'lion, king of beasts, Panthera leo',
+ 292: 'tiger, Panthera tigris',
+ 293: 'cheetah, chetah, Acinonyx jubatus',
+ 294: 'brown bear, bruin, Ursus arctos',
+ 295: 'American black bear, black bear, Ursus americanus, Euarctos americanus',
+ 296: 'ice bear, polar bear, Ursus Maritimus, Thalarctos maritimus',
+ 297: 'sloth bear, Melursus ursinus, Ursus ursinus',
+ 298: 'mongoose',
+ 299: 'meerkat, mierkat',
+ 300: 'tiger beetle',
+ 301: 'ladybug, ladybeetle, lady beetle, ladybird, ladybird beetle',
+ 302: 'ground beetle, carabid beetle',
+ 303: 'long-horned beetle, longicorn, longicorn beetle',
+ 304: 'leaf beetle, chrysomelid',
+ 305: 'dung beetle',
+ 306: 'rhinoceros beetle',
+ 307: 'weevil',
+ 308: 'fly',
+ 309: 'bee',
+ 310: 'ant, emmet, pismire',
+ 311: 'grasshopper, hopper',
+ 312: 'cricket',
+ 313: 'walking stick, walkingstick, stick insect',
+ 314: 'cockroach, roach',
+ 315: 'mantis, mantid',
+ 316: 'cicada, cicala',
+ 317: 'leafhopper',
+ 318: 'lacewing, lacewing fly',
+ 319: "dragonfly, darning needle, devil's darning needle, sewing needle, snake feeder, snake doctor, mosquito hawk, skeeter hawk",
+ 320: 'damselfly',
+ 321: 'admiral',
+ 322: 'ringlet, ringlet butterfly',
+ 323: 'monarch, monarch butterfly, milkweed butterfly, Danaus plexippus',
+ 324: 'cabbage butterfly',
+ 325: 'sulphur butterfly, sulfur butterfly',
+ 326: 'lycaenid, lycaenid butterfly',
+ 327: 'starfish, sea star',
+ 328: 'sea urchin',
+ 329: 'sea cucumber, holothurian',
+ 330: 'wood rabbit, cottontail, cottontail rabbit',
+ 331: 'hare',
+ 332: 'Angora, Angora rabbit',
+ 333: 'hamster',
+ 334: 'porcupine, hedgehog',
+ 335: 'fox squirrel, eastern fox squirrel, Sciurus niger',
+ 336: 'marmot',
+ 337: 'beaver',
+ 338: 'guinea pig, Cavia cobaya',
+ 339: 'sorrel',
+ 340: 'zebra',
+ 341: 'hog, pig, grunter, squealer, Sus scrofa',
+ 342: 'wild boar, boar, Sus scrofa',
+ 343: 'warthog',
+ 344: 'hippopotamus, hippo, river horse, Hippopotamus amphibius',
+ 345: 'ox',
+ 346: 'water buffalo, water ox, Asiatic buffalo, Bubalus bubalis',
+ 347: 'bison',
+ 348: 'ram, tup',
+ 349: 'bighorn, bighorn sheep, cimarron, Rocky Mountain bighorn, Rocky Mountain sheep, Ovis canadensis',
+ 350: 'ibex, Capra ibex',
+ 351: 'hartebeest',
+ 352: 'impala, Aepyceros melampus',
+ 353: 'gazelle',
+ 354: 'Arabian camel, dromedary, Camelus dromedarius',
+ 355: 'llama',
+ 356: 'weasel',
+ 357: 'mink',
+ 358: 'polecat, fitch, foulmart, foumart, Mustela putorius',
+ 359: 'black-footed ferret, ferret, Mustela nigripes',
+ 360: 'otter',
+ 361: 'skunk, polecat, wood pussy',
+ 362: 'badger',
+ 363: 'armadillo',
+ 364: 'three-toed sloth, ai, Bradypus tridactylus',
+ 365: 'orangutan, orang, orangutang, Pongo pygmaeus',
+ 366: 'gorilla, Gorilla gorilla',
+ 367: 'chimpanzee, chimp, Pan troglodytes',
+ 368: 'gibbon, Hylobates lar',
+ 369: 'siamang, Hylobates syndactylus, Symphalangus syndactylus',
+ 370: 'guenon, guenon monkey',
+ 371: 'patas, hussar monkey, Erythrocebus patas',
+ 372: 'baboon',
+ 373: 'macaque',
+ 374: 'langur',
+ 375: 'colobus, colobus monkey',
+ 376: 'proboscis monkey, Nasalis larvatus',
+ 377: 'marmoset',
+ 378: 'capuchin, ringtail, Cebus capucinus',
+ 379: 'howler monkey, howler',
+ 380: 'titi, titi monkey',
+ 381: 'spider monkey, Ateles geoffroyi',
+ 382: 'squirrel monkey, Saimiri sciureus',
+ 383: 'Madagascar cat, ring-tailed lemur, Lemur catta',
+ 384: 'indri, indris, Indri indri, Indri brevicaudatus',
+ 385: 'Indian elephant, Elephas maximus',
+ 386: 'African elephant, Loxodonta africana',
+ 387: 'lesser panda, red panda, panda, bear cat, cat bear, Ailurus fulgens',
+ 388: 'giant panda, panda, panda bear, coon bear, Ailuropoda melanoleuca',
+ 389: 'barracouta, snoek',
+ 390: 'eel',
+ 391: 'coho, cohoe, coho salmon, blue jack, silver salmon, Oncorhynchus kisutch',
+ 392: 'rock beauty, Holocanthus tricolor',
+ 393: 'anemone fish',
+ 394: 'sturgeon',
+ 395: 'gar, garfish, garpike, billfish, Lepisosteus osseus',
+ 396: 'lionfish',
+ 397: 'puffer, pufferfish, blowfish, globefish',
+ 398: 'abacus',
+ 399: 'abaya',
+ 400: "academic gown, academic robe, judge's robe",
+ 401: 'accordion, piano accordion, squeeze box',
+ 402: 'acoustic guitar',
+ 403: 'aircraft carrier, carrier, flattop, attack aircraft carrier',
+ 404: 'airliner',
+ 405: 'airship, dirigible',
+ 406: 'altar',
+ 407: 'ambulance',
+ 408: 'amphibian, amphibious vehicle',
+ 409: 'analog clock',
+ 410: 'apiary, bee house',
+ 411: 'apron',
+ 412: 'ashcan, trash can, garbage can, wastebin, ash bin, ash-bin, ashbin, dustbin, trash barrel, trash bin',
+ 413: 'assault rifle, assault gun',
+ 414: 'backpack, back pack, knapsack, packsack, rucksack, haversack',
+ 415: 'bakery, bakeshop, bakehouse',
+ 416: 'balance beam, beam',
+ 417: 'balloon',
+ 418: 'ballpoint, ballpoint pen, ballpen, Biro',
+ 419: 'Band Aid',
+ 420: 'banjo',
+ 421: 'bannister, banister, balustrade, balusters, handrail',
+ 422: 'barbell',
+ 423: 'barber chair',
+ 424: 'barbershop',
+ 425: 'barn',
+ 426: 'barometer',
+ 427: 'barrel, cask',
+ 428: 'barrow, garden cart, lawn cart, wheelbarrow',
+ 429: 'baseball',
+ 430: 'basketball',
+ 431: 'bassinet',
+ 432: 'bassoon',
+ 433: 'bathing cap, swimming cap',
+ 434: 'bath towel',
+ 435: 'bathtub, bathing tub, bath, tub',
+ 436: 'beach wagon, station wagon, wagon, estate car, beach waggon, station waggon, waggon',
+ 437: 'beacon, lighthouse, beacon light, pharos',
+ 438: 'beaker',
+ 439: 'bearskin, busby, shako',
+ 440: 'beer bottle',
+ 441: 'beer glass',
+ 442: 'bell cote, bell cot',
+ 443: 'bib',
+ 444: 'bicycle-built-for-two, tandem bicycle, tandem',
+ 445: 'bikini, two-piece',
+ 446: 'binder, ring-binder',
+ 447: 'binoculars, field glasses, opera glasses',
+ 448: 'birdhouse',
+ 449: 'boathouse',
+ 450: 'bobsled, bobsleigh, bob',
+ 451: 'bolo tie, bolo, bola tie, bola',
+ 452: 'bonnet, poke bonnet',
+ 453: 'bookcase',
+ 454: 'bookshop, bookstore, bookstall',
+ 455: 'bottlecap',
+ 456: 'bow',
+ 457: 'bow tie, bow-tie, bowtie',
+ 458: 'brass, memorial tablet, plaque',
+ 459: 'brassiere, bra, bandeau',
+ 460: 'breakwater, groin, groyne, mole, bulwark, seawall, jetty',
+ 461: 'breastplate, aegis, egis',
+ 462: 'broom',
+ 463: 'bucket, pail',
+ 464: 'buckle',
+ 465: 'bulletproof vest',
+ 466: 'bullet train, bullet',
+ 467: 'butcher shop, meat market',
+ 468: 'cab, hack, taxi, taxicab',
+ 469: 'caldron, cauldron',
+ 470: 'candle, taper, wax light',
+ 471: 'cannon',
+ 472: 'canoe',
+ 473: 'can opener, tin opener',
+ 474: 'cardigan',
+ 475: 'car mirror',
+ 476: 'carousel, carrousel, merry-go-round, roundabout, whirligig',
+ 477: "carpenter's kit, tool kit",
+ 478: 'carton',
+ 479: 'car wheel',
+ 480: 'cash machine, cash dispenser, automated teller machine, automatic teller machine, automated teller, automatic teller, ATM',
+ 481: 'cassette',
+ 482: 'cassette player',
+ 483: 'castle',
+ 484: 'catamaran',
+ 485: 'CD player',
+ 486: 'cello, violoncello',
+ 487: 'cellular telephone, cellular phone, cellphone, cell, mobile phone',
+ 488: 'chain',
+ 489: 'chainlink fence',
+ 490: 'chain mail, ring mail, mail, chain armor, chain armour, ring armor, ring armour',
+ 491: 'chain saw, chainsaw',
+ 492: 'chest',
+ 493: 'chiffonier, commode',
+ 494: 'chime, bell, gong',
+ 495: 'china cabinet, china closet',
+ 496: 'Christmas stocking',
+ 497: 'church, church building',
+ 498: 'cinema, movie theater, movie theatre, movie house, picture palace',
+ 499: 'cleaver, meat cleaver, chopper',
+ 500: 'cliff dwelling',
+ 501: 'cloak',
+ 502: 'clog, geta, patten, sabot',
+ 503: 'cocktail shaker',
+ 504: 'coffee mug',
+ 505: 'coffeepot',
+ 506: 'coil, spiral, volute, whorl, helix',
+ 507: 'combination lock',
+ 508: 'computer keyboard, keypad',
+ 509: 'confectionery, confectionary, candy store',
+ 510: 'container ship, containership, container vessel',
+ 511: 'convertible',
+ 512: 'corkscrew, bottle screw',
+ 513: 'cornet, horn, trumpet, trump',
+ 514: 'cowboy boot',
+ 515: 'cowboy hat, ten-gallon hat',
+ 516: 'cradle',
+ 517: 'crane',
+ 518: 'crash helmet',
+ 519: 'crate',
+ 520: 'crib, cot',
+ 521: 'Crock Pot',
+ 522: 'croquet ball',
+ 523: 'crutch',
+ 524: 'cuirass',
+ 525: 'dam, dike, dyke',
+ 526: 'desk',
+ 527: 'desktop computer',
+ 528: 'dial telephone, dial phone',
+ 529: 'diaper, nappy, napkin',
+ 530: 'digital clock',
+ 531: 'digital watch',
+ 532: 'dining table, board',
+ 533: 'dishrag, dishcloth',
+ 534: 'dishwasher, dish washer, dishwashing machine',
+ 535: 'disk brake, disc brake',
+ 536: 'dock, dockage, docking facility',
+ 537: 'dogsled, dog sled, dog sleigh',
+ 538: 'dome',
+ 539: 'doormat, welcome mat',
+ 540: 'drilling platform, offshore rig',
+ 541: 'drum, membranophone, tympan',
+ 542: 'drumstick',
+ 543: 'dumbbell',
+ 544: 'Dutch oven',
+ 545: 'electric fan, blower',
+ 546: 'electric guitar',
+ 547: 'electric locomotive',
+ 548: 'entertainment center',
+ 549: 'envelope',
+ 550: 'espresso maker',
+ 551: 'face powder',
+ 552: 'feather boa, boa',
+ 553: 'file, file cabinet, filing cabinet',
+ 554: 'fireboat',
+ 555: 'fire engine, fire truck',
+ 556: 'fire screen, fireguard',
+ 557: 'flagpole, flagstaff',
+ 558: 'flute, transverse flute',
+ 559: 'folding chair',
+ 560: 'football helmet',
+ 561: 'forklift',
+ 562: 'fountain',
+ 563: 'fountain pen',
+ 564: 'four-poster',
+ 565: 'freight car',
+ 566: 'French horn, horn',
+ 567: 'frying pan, frypan, skillet',
+ 568: 'fur coat',
+ 569: 'garbage truck, dustcart',
+ 570: 'gasmask, respirator, gas helmet',
+ 571: 'gas pump, gasoline pump, petrol pump, island dispenser',
+ 572: 'goblet',
+ 573: 'go-kart',
+ 574: 'golf ball',
+ 575: 'golfcart, golf cart',
+ 576: 'gondola',
+ 577: 'gong, tam-tam',
+ 578: 'gown',
+ 579: 'grand piano, grand',
+ 580: 'greenhouse, nursery, glasshouse',
+ 581: 'grille, radiator grille',
+ 582: 'grocery store, grocery, food market, market',
+ 583: 'guillotine',
+ 584: 'hair slide',
+ 585: 'hair spray',
+ 586: 'half track',
+ 587: 'hammer',
+ 588: 'hamper',
+ 589: 'hand blower, blow dryer, blow drier, hair dryer, hair drier',
+ 590: 'hand-held computer, hand-held microcomputer',
+ 591: 'handkerchief, hankie, hanky, hankey',
+ 592: 'hard disc, hard disk, fixed disk',
+ 593: 'harmonica, mouth organ, harp, mouth harp',
+ 594: 'harp',
+ 595: 'harvester, reaper',
+ 596: 'hatchet',
+ 597: 'holster',
+ 598: 'home theater, home theatre',
+ 599: 'honeycomb',
+ 600: 'hook, claw',
+ 601: 'hoopskirt, crinoline',
+ 602: 'horizontal bar, high bar',
+ 603: 'horse cart, horse-cart',
+ 604: 'hourglass',
+ 605: 'iPod',
+ 606: 'iron, smoothing iron',
+ 607: "jack-o'-lantern",
+ 608: 'jean, blue jean, denim',
+ 609: 'jeep, landrover',
+ 610: 'jersey, T-shirt, tee shirt',
+ 611: 'jigsaw puzzle',
+ 612: 'jinrikisha, ricksha, rickshaw',
+ 613: 'joystick',
+ 614: 'kimono',
+ 615: 'knee pad',
+ 616: 'knot',
+ 617: 'lab coat, laboratory coat',
+ 618: 'ladle',
+ 619: 'lampshade, lamp shade',
+ 620: 'laptop, laptop computer',
+ 621: 'lawn mower, mower',
+ 622: 'lens cap, lens cover',
+ 623: 'letter opener, paper knife, paperknife',
+ 624: 'library',
+ 625: 'lifeboat',
+ 626: 'lighter, light, igniter, ignitor',
+ 627: 'limousine, limo',
+ 628: 'liner, ocean liner',
+ 629: 'lipstick, lip rouge',
+ 630: 'Loafer',
+ 631: 'lotion',
+ 632: 'loudspeaker, speaker, speaker unit, loudspeaker system, speaker system',
+ 633: "loupe, jeweler's loupe",
+ 634: 'lumbermill, sawmill',
+ 635: 'magnetic compass',
+ 636: 'mailbag, postbag',
+ 637: 'mailbox, letter box',
+ 638: 'maillot',
+ 639: 'maillot, tank suit',
+ 640: 'manhole cover',
+ 641: 'maraca',
+ 642: 'marimba, xylophone',
+ 643: 'mask',
+ 644: 'matchstick',
+ 645: 'maypole',
+ 646: 'maze, labyrinth',
+ 647: 'measuring cup',
+ 648: 'medicine chest, medicine cabinet',
+ 649: 'megalith, megalithic structure',
+ 650: 'microphone, mike',
+ 651: 'microwave, microwave oven',
+ 652: 'military uniform',
+ 653: 'milk can',
+ 654: 'minibus',
+ 655: 'miniskirt, mini',
+ 656: 'minivan',
+ 657: 'missile',
+ 658: 'mitten',
+ 659: 'mixing bowl',
+ 660: 'mobile home, manufactured home',
+ 661: 'Model T',
+ 662: 'modem',
+ 663: 'monastery',
+ 664: 'monitor',
+ 665: 'moped',
+ 666: 'mortar',
+ 667: 'mortarboard',
+ 668: 'mosque',
+ 669: 'mosquito net',
+ 670: 'motor scooter, scooter',
+ 671: 'mountain bike, all-terrain bike, off-roader',
+ 672: 'mountain tent',
+ 673: 'mouse, computer mouse',
+ 674: 'mousetrap',
+ 675: 'moving van',
+ 676: 'muzzle',
+ 677: 'nail',
+ 678: 'neck brace',
+ 679: 'necklace',
+ 680: 'nipple',
+ 681: 'notebook, notebook computer',
+ 682: 'obelisk',
+ 683: 'oboe, hautboy, hautbois',
+ 684: 'ocarina, sweet potato',
+ 685: 'odometer, hodometer, mileometer, milometer',
+ 686: 'oil filter',
+ 687: 'organ, pipe organ',
+ 688: 'oscilloscope, scope, cathode-ray oscilloscope, CRO',
+ 689: 'overskirt',
+ 690: 'oxcart',
+ 691: 'oxygen mask',
+ 692: 'packet',
+ 693: 'paddle, boat paddle',
+ 694: 'paddlewheel, paddle wheel',
+ 695: 'padlock',
+ 696: 'paintbrush',
+ 697: "pajama, pyjama, pj's, jammies",
+ 698: 'palace',
+ 699: 'panpipe, pandean pipe, syrinx',
+ 700: 'paper towel',
+ 701: 'parachute, chute',
+ 702: 'parallel bars, bars',
+ 703: 'park bench',
+ 704: 'parking meter',
+ 705: 'passenger car, coach, carriage',
+ 706: 'patio, terrace',
+ 707: 'pay-phone, pay-station',
+ 708: 'pedestal, plinth, footstall',
+ 709: 'pencil box, pencil case',
+ 710: 'pencil sharpener',
+ 711: 'perfume, essence',
+ 712: 'Petri dish',
+ 713: 'photocopier',
+ 714: 'pick, plectrum, plectron',
+ 715: 'pickelhaube',
+ 716: 'picket fence, paling',
+ 717: 'pickup, pickup truck',
+ 718: 'pier',
+ 719: 'piggy bank, penny bank',
+ 720: 'pill bottle',
+ 721: 'pillow',
+ 722: 'ping-pong ball',
+ 723: 'pinwheel',
+ 724: 'pirate, pirate ship',
+ 725: 'pitcher, ewer',
+ 726: "plane, carpenter's plane, woodworking plane",
+ 727: 'planetarium',
+ 728: 'plastic bag',
+ 729: 'plate rack',
+ 730: 'plow, plough',
+ 731: "plunger, plumber's helper",
+ 732: 'Polaroid camera, Polaroid Land camera',
+ 733: 'pole',
+ 734: 'police van, police wagon, paddy wagon, patrol wagon, wagon, black Maria',
+ 735: 'poncho',
+ 736: 'pool table, billiard table, snooker table',
+ 737: 'pop bottle, soda bottle',
+ 738: 'pot, flowerpot',
+ 739: "potter's wheel",
+ 740: 'power drill',
+ 741: 'prayer rug, prayer mat',
+ 742: 'printer',
+ 743: 'prison, prison house',
+ 744: 'projectile, missile',
+ 745: 'projector',
+ 746: 'puck, hockey puck',
+ 747: 'punching bag, punch bag, punching ball, punchball',
+ 748: 'purse',
+ 749: 'quill, quill pen',
+ 750: 'quilt, comforter, comfort, puff',
+ 751: 'racer, race car, racing car',
+ 752: 'racket, racquet',
+ 753: 'radiator',
+ 754: 'radio, wireless',
+ 755: 'radio telescope, radio reflector',
+ 756: 'rain barrel',
+ 757: 'recreational vehicle, RV, R.V.',
+ 758: 'reel',
+ 759: 'reflex camera',
+ 760: 'refrigerator, icebox',
+ 761: 'remote control, remote',
+ 762: 'restaurant, eating house, eating place, eatery',
+ 763: 'revolver, six-gun, six-shooter',
+ 764: 'rifle',
+ 765: 'rocking chair, rocker',
+ 766: 'rotisserie',
+ 767: 'rubber eraser, rubber, pencil eraser',
+ 768: 'rugby ball',
+ 769: 'rule, ruler',
+ 770: 'running shoe',
+ 771: 'safe',
+ 772: 'safety pin',
+ 773: 'saltshaker, salt shaker',
+ 774: 'sandal',
+ 775: 'sarong',
+ 776: 'sax, saxophone',
+ 777: 'scabbard',
+ 778: 'scale, weighing machine',
+ 779: 'school bus',
+ 780: 'schooner',
+ 781: 'scoreboard',
+ 782: 'screen, CRT screen',
+ 783: 'screw',
+ 784: 'screwdriver',
+ 785: 'seat belt, seatbelt',
+ 786: 'sewing machine',
+ 787: 'shield, buckler',
+ 788: 'shoe shop, shoe-shop, shoe store',
+ 789: 'shoji',
+ 790: 'shopping basket',
+ 791: 'shopping cart',
+ 792: 'shovel',
+ 793: 'shower cap',
+ 794: 'shower curtain',
+ 795: 'ski',
+ 796: 'ski mask',
+ 797: 'sleeping bag',
+ 798: 'slide rule, slipstick',
+ 799: 'sliding door',
+ 800: 'slot, one-armed bandit',
+ 801: 'snorkel',
+ 802: 'snowmobile',
+ 803: 'snowplow, snowplough',
+ 804: 'soap dispenser',
+ 805: 'soccer ball',
+ 806: 'sock',
+ 807: 'solar dish, solar collector, solar furnace',
+ 808: 'sombrero',
+ 809: 'soup bowl',
+ 810: 'space bar',
+ 811: 'space heater',
+ 812: 'space shuttle',
+ 813: 'spatula',
+ 814: 'speedboat',
+ 815: "spider web, spider's web",
+ 816: 'spindle',
+ 817: 'sports car, sport car',
+ 818: 'spotlight, spot',
+ 819: 'stage',
+ 820: 'steam locomotive',
+ 821: 'steel arch bridge',
+ 822: 'steel drum',
+ 823: 'stethoscope',
+ 824: 'stole',
+ 825: 'stone wall',
+ 826: 'stopwatch, stop watch',
+ 827: 'stove',
+ 828: 'strainer',
+ 829: 'streetcar, tram, tramcar, trolley, trolley car',
+ 830: 'stretcher',
+ 831: 'studio couch, day bed',
+ 832: 'stupa, tope',
+ 833: 'submarine, pigboat, sub, U-boat',
+ 834: 'suit, suit of clothes',
+ 835: 'sundial',
+ 836: 'sunglass',
+ 837: 'sunglasses, dark glasses, shades',
+ 838: 'sunscreen, sunblock, sun blocker',
+ 839: 'suspension bridge',
+ 840: 'swab, swob, mop',
+ 841: 'sweatshirt',
+ 842: 'swimming trunks, bathing trunks',
+ 843: 'swing',
+ 844: 'switch, electric switch, electrical switch',
+ 845: 'syringe',
+ 846: 'table lamp',
+ 847: 'tank, army tank, armored combat vehicle, armoured combat vehicle',
+ 848: 'tape player',
+ 849: 'teapot',
+ 850: 'teddy, teddy bear',
+ 851: 'television, television system',
+ 852: 'tennis ball',
+ 853: 'thatch, thatched roof',
+ 854: 'theater curtain, theatre curtain',
+ 855: 'thimble',
+ 856: 'thresher, thrasher, threshing machine',
+ 857: 'throne',
+ 858: 'tile roof',
+ 859: 'toaster',
+ 860: 'tobacco shop, tobacconist shop, tobacconist',
+ 861: 'toilet seat',
+ 862: 'torch',
+ 863: 'totem pole',
+ 864: 'tow truck, tow car, wrecker',
+ 865: 'toyshop',
+ 866: 'tractor',
+ 867: 'trailer truck, tractor trailer, trucking rig, rig, articulated lorry, semi',
+ 868: 'tray',
+ 869: 'trench coat',
+ 870: 'tricycle, trike, velocipede',
+ 871: 'trimaran',
+ 872: 'tripod',
+ 873: 'triumphal arch',
+ 874: 'trolleybus, trolley coach, trackless trolley',
+ 875: 'trombone',
+ 876: 'tub, vat',
+ 877: 'turnstile',
+ 878: 'typewriter keyboard',
+ 879: 'umbrella',
+ 880: 'unicycle, monocycle',
+ 881: 'upright, upright piano',
+ 882: 'vacuum, vacuum cleaner',
+ 883: 'vase',
+ 884: 'vault',
+ 885: 'velvet',
+ 886: 'vending machine',
+ 887: 'vestment',
+ 888: 'viaduct',
+ 889: 'violin, fiddle',
+ 890: 'volleyball',
+ 891: 'waffle iron',
+ 892: 'wall clock',
+ 893: 'wallet, billfold, notecase, pocketbook',
+ 894: 'wardrobe, closet, press',
+ 895: 'warplane, military plane',
+ 896: 'washbasin, handbasin, washbowl, lavabo, wash-hand basin',
+ 897: 'washer, automatic washer, washing machine',
+ 898: 'water bottle',
+ 899: 'water jug',
+ 900: 'water tower',
+ 901: 'whiskey jug',
+ 902: 'whistle',
+ 903: 'wig',
+ 904: 'window screen',
+ 905: 'window shade',
+ 906: 'Windsor tie',
+ 907: 'wine bottle',
+ 908: 'wing',
+ 909: 'wok',
+ 910: 'wooden spoon',
+ 911: 'wool, woolen, woollen',
+ 912: 'worm fence, snake fence, snake-rail fence, Virginia fence',
+ 913: 'wreck',
+ 914: 'yawl',
+ 915: 'yurt',
+ 916: 'web site, website, internet site, site',
+ 917: 'comic book',
+ 918: 'crossword puzzle, crossword',
+ 919: 'street sign',
+ 920: 'traffic light, traffic signal, stoplight',
+ 921: 'book jacket, dust cover, dust jacket, dust wrapper',
+ 922: 'menu',
+ 923: 'plate',
+ 924: 'guacamole',
+ 925: 'consomme',
+ 926: 'hot pot, hotpot',
+ 927: 'trifle',
+ 928: 'ice cream, icecream',
+ 929: 'ice lolly, lolly, lollipop, popsicle',
+ 930: 'French loaf',
+ 931: 'bagel, beigel',
+ 932: 'pretzel',
+ 933: 'cheeseburger',
+ 934: 'hotdog, hot dog, red hot',
+ 935: 'mashed potato',
+ 936: 'head cabbage',
+ 937: 'broccoli',
+ 938: 'cauliflower',
+ 939: 'zucchini, courgette',
+ 940: 'spaghetti squash',
+ 941: 'acorn squash',
+ 942: 'butternut squash',
+ 943: 'cucumber, cuke',
+ 944: 'artichoke, globe artichoke',
+ 945: 'bell pepper',
+ 946: 'cardoon',
+ 947: 'mushroom',
+ 948: 'Granny Smith',
+ 949: 'strawberry',
+ 950: 'orange',
+ 951: 'lemon',
+ 952: 'fig',
+ 953: 'pineapple, ananas',
+ 954: 'banana',
+ 955: 'jackfruit, jak, jack',
+ 956: 'custard apple',
+ 957: 'pomegranate',
+ 958: 'hay',
+ 959: 'carbonara',
+ 960: 'chocolate sauce, chocolate syrup',
+ 961: 'dough',
+ 962: 'meat loaf, meatloaf',
+ 963: 'pizza, pizza pie',
+ 964: 'potpie',
+ 965: 'burrito',
+ 966: 'red wine',
+ 967: 'espresso',
+ 968: 'cup',
+ 969: 'eggnog',
+ 970: 'alp',
+ 971: 'bubble',
+ 972: 'cliff, drop, drop-off',
+ 973: 'coral reef',
+ 974: 'geyser',
+ 975: 'lakeside, lakeshore',
+ 976: 'promontory, headland, head, foreland',
+ 977: 'sandbar, sand bar',
+ 978: 'seashore, coast, seacoast, sea-coast',
+ 979: 'valley, vale',
+ 980: 'volcano',
+ 981: 'ballplayer, baseball player',
+ 982: 'groom, bridegroom',
+ 983: 'scuba diver',
+ 984: 'rapeseed',
+ 985: 'daisy',
+ 986: "yellow lady's slipper, yellow lady-slipper, Cypripedium calceolus, Cypripedium parviflorum",
+ 987: 'corn',
+ 988: 'acorn',
+ 989: 'hip, rose hip, rosehip',
+ 990: 'buckeye, horse chestnut, conker',
+ 991: 'coral fungus',
+ 992: 'agaric',
+ 993: 'gyromitra',
+ 994: 'stinkhorn, carrion fungus',
+ 995: 'earthstar',
+ 996: 'hen-of-the-woods, hen of the woods, Polyporus frondosus, Grifola frondosa',
+ 997: 'bolete',
+ 998: 'ear, spike, capitulum',
+ 999: 'toilet tissue, toilet paper, bathroom tissue'}

modeling/__init__.py

@@ -0,0 +1,15 @@
+"""
+Copyright (2024) Bytedance Ltd. and/or its affiliates
+
+Licensed under the Apache License, Version 2.0 (the "License"); 
+you may not use this file except in compliance with the License. 
+You may obtain a copy of the License at 
+
+    http://www.apache.org/licenses/LICENSE-2.0 
+
+Unless required by applicable law or agreed to in writing, software 
+distributed under the License is distributed on an "AS IS" BASIS, 
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 
+See the License for the specific language governing permissions and 
+limitations under the License.
+"""

modeling/maskgit.py

@@ -0,0 +1,282 @@
+"""This file contains implementation for MaskGIT model.
+
+Copyright (2024) Bytedance Ltd. and/or its affiliates
+
+Licensed under the Apache License, Version 2.0 (the "License"); 
+you may not use this file except in compliance with the License. 
+You may obtain a copy of the License at 
+
+    http://www.apache.org/licenses/LICENSE-2.0 
+
+Unless required by applicable law or agreed to in writing, software 
+distributed under the License is distributed on an "AS IS" BASIS, 
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 
+See the License for the specific language governing permissions and 
+limitations under the License. 
+
+Reference: 
+    https://github.com/huggingface/open-muse
+    https://github.com/baaivision/MUSE-Pytorch
+    https://github.com/sail-sg/MDT/blob/main/masked_diffusion/models.py
+"""
+
+import torch
+from torch import nn
+import numpy as np
+import math
+import torch.utils.checkpoint
+from transformers import BertConfig, BertModel
+from einops import rearrange
+
+import json
+from huggingface_hub import PyTorchModelHubMixin
+from omegaconf import OmegaConf
+from pathlib import Path
+
+from modeling.modules.base_model import BaseModel
+from modeling.modules.blocks import UViTBlock
+
+
+class ImageBert(BaseModel, PyTorchModelHubMixin, tags=["arxiv:2406.07550", "image-generation"], repo_url="https://github.com/bytedance/1d-tokenizer", license="apache-2.0"):
+    def __init__(self, config):
+
+        if isinstance(config, dict):
+            config = OmegaConf.create(config)
+
+        super().__init__()
+        self.config = config
+        self.target_codebook_size = config.model.vq_model.codebook_size
+        self.condition_num_classes = config.model.generator.condition_num_classes
+        self.image_seq_len = config.model.generator.image_seq_len
+        self.mask_token_id = self.target_codebook_size
+        self.hidden_size = config.model.generator.hidden_size
+        self.num_hidden_layers = config.model.generator.num_hidden_layers
+        self.num_attention_heads = config.model.generator.num_attention_heads
+        self.intermediate_size = config.model.generator.intermediate_size
+
+        self.model = BertModel(BertConfig(
+            vocab_size=self.target_codebook_size + self.condition_num_classes + 2,
+            hidden_size=self.hidden_size,
+            num_hidden_layers=self.num_hidden_layers,
+            num_attention_heads=self.num_attention_heads,
+            intermediate_size=self.intermediate_size,
+            hidden_act='gelu',
+            hidden_dropout_prob=config.model.generator.dropout,
+            attention_probs_dropout_prob=config.model.generator.attn_drop,
+            max_position_embeddings=config.model.generator.image_seq_len + 1,
+            initializer_range=0.02,
+            layer_norm_eps=1e-12,
+            pad_token_id=None,
+            position_embedding_type="absolute",
+            use_cache=True
+        ), add_pooling_layer=False)
+        self.model.lm_head = nn.Linear(self.hidden_size, self.target_codebook_size, bias=True)
+        
+        self.model.post_init()
+
+    def _save_pretrained(self, save_directory: Path) -> None:
+        """Save weights and config to a local directory."""
+        # Assume 'self.config' is your DictConfig object
+        # Convert to a regular dictionary
+        dict_config = OmegaConf.to_container(self.config)
+        # Save as JSON
+        file_path = Path(save_directory) / "config.json"
+        with open(file_path, 'w') as json_file:
+            json.dump(dict_config, json_file, indent=4)
+        super()._save_pretrained(save_directory)
+    
+    def forward(self, input_ids=None, condition=None, cond_drop_prob=0.1):
+        # Token space:
+        #  [0, codebook_size - 1]                       : those are the learned quantized image tokens
+        #  codebook_size                                : the mask token used to mask image tokens
+        #  [codebook_size + 1, codebook_size + nclass]  : the imagenet class tokens
+        #  codebook_size + 1 + nclass                   : the class drop label
+        drop_label_mask = torch.rand_like(condition, dtype=torch.float) < cond_drop_prob
+        # Shift the classes
+        condition = condition + self.target_codebook_size + 1  # [0, 999] -> [codebook_size + 1, codebook_size + 999]
+        condition[drop_label_mask] = self.condition_num_classes + self.target_codebook_size + 1
+        # prepend condition token
+        if input_ids is not None:
+            input_ids = torch.cat([condition.view(condition.shape[0], -1),
+                                   input_ids.view(input_ids.shape[0], -1),], dim=1)
+        else:
+            # at least there should be masked token
+            raise NotImplementedError
+        model_output = self.model(input_ids=input_ids)
+        model_output = model_output[0]
+        return self.model.lm_head(model_output[:, 1:]) # remove cond
+    
+    # ref: https://github.com/baaivision/MUSE-Pytorch/blob/master/libs/muse.py#L40
+    @torch.no_grad()
+    def generate(self,
+                 condition,
+                 guidance_scale=3.0,
+                 guidance_decay="constant",
+                 guidance_scale_pow=3.0,
+                 randomize_temperature=4.5,
+                 softmax_temperature_annealing=False,
+                 num_sample_steps=8):
+        if guidance_decay not in ["constant", "linear", "power-cosine"]:
+            # contstant: constant guidance scale
+            # linear: linear increasing the guidance scale as in MUSE
+            # power-cosine: the guidance schedule from MDT
+            raise ValueError(f"Unsupported guidance decay {guidance_decay}")
+        device = condition.device
+        ids = torch.full((condition.shape[0], self.image_seq_len),
+                          self.mask_token_id, device=device)
+
+        cfg_scale = guidance_scale if guidance_decay == "constant" else 0.
+
+        for step in range(num_sample_steps):
+            ratio = 1. * (step + 1) / num_sample_steps
+            annealed_temp = randomize_temperature * (1.0 - ratio)
+            is_mask = (ids == self.mask_token_id)
+
+            if guidance_decay == "power-cosine":
+                # ref: https://github.com/sail-sg/MDT/blob/main/masked_diffusion/models.py#L501
+                guidance_scale_pow = torch.ones((1), device=device) * guidance_scale_pow
+                scale_step = (1 - torch.cos(((step / num_sample_steps) ** guidance_scale_pow) * torch.pi)) * 1/2
+                cfg_scale = (guidance_scale - 1) * scale_step + 1
+
+            if cfg_scale != 0:
+                cond_logits = self.forward(
+                    ids, condition, cond_drop_prob=0.0
+                )
+                uncond_logits = self.forward(
+                    ids, condition, cond_drop_prob=1.0
+                )
+                if guidance_decay == "power-cosine":
+                    logits = uncond_logits + (cond_logits - uncond_logits) * cfg_scale
+                else:
+                    logits = cond_logits + (cond_logits - uncond_logits) * cfg_scale
+            else:
+                logits = self.forward(
+                    ids, condition, cond_drop_prob=0.0
+                )
+
+            if softmax_temperature_annealing:
+                softmax_temperature = 0.5 + 0.8 * (1 - ratio)
+                logits = logits / softmax_temperature
+
+            # Add gumbel noise
+            def log(t, eps=1e-20):
+                return torch.log(t.clamp(min=eps))
+            def gumbel_noise(t):
+                noise = torch.zeros_like(t).uniform_(0, 1)
+                return -log(-log(noise))
+            def add_gumbel_noise(t, temperature):
+                return t + temperature * gumbel_noise(t)
+
+            sampled_ids = add_gumbel_noise(logits, annealed_temp).argmax(dim=-1)
+            sampled_logits = torch.squeeze(
+                torch.gather(logits, dim=-1, index=torch.unsqueeze(sampled_ids, -1)), -1)
+            sampled_ids = torch.where(is_mask, sampled_ids, ids)
+            sampled_logits = torch.where(is_mask, sampled_logits, +np.inf).float()
+            # masking
+            mask_ratio = np.arccos(ratio) / (math.pi * 0.5)
+
+            mask_len = torch.Tensor([np.floor(self.image_seq_len * mask_ratio)]).to(device)
+            mask_len = torch.maximum(torch.Tensor([1]).to(device),
+                                     torch.minimum(torch.sum(is_mask, dim=-1, keepdims=True) - 1,
+                                                   mask_len))[0].squeeze()
+            confidence = add_gumbel_noise(sampled_logits, annealed_temp)
+            sorted_confidence, _ = torch.sort(confidence, axis=-1)
+            cut_off = sorted_confidence[:, mask_len.long() - 1:mask_len.long()]
+            masking = (confidence <= cut_off)
+            if step == num_sample_steps - 1:
+                ids = sampled_ids
+            else:
+                ids = torch.where(masking, self.mask_token_id, sampled_ids)
+
+            if guidance_decay == "linear":
+                cfg_scale = ratio * guidance_scale
+        return ids
+
+    def masking_input_tokens(self, input_tokens):
+        batch_size, seq_len = input_tokens.shape
+        device = input_tokens.device
+
+        timesteps = torch.zeros((batch_size,), device=device).float().uniform_(0, 1.0)
+        mask_ratio = torch.acos(timesteps) / (math.pi * 0.5) # arccos schedule
+        mask_ratio = torch.clamp(mask_ratio, min=1e-6, max=1.)
+        num_token_masked = (seq_len * mask_ratio).round().clamp(min=1)
+        batch_randperm = torch.rand(batch_size, seq_len, device=device).argsort(dim=-1)
+        masks = batch_randperm < rearrange(num_token_masked, 'b -> b 1')
+        masked_tokens = torch.where(masks, self.mask_token_id, input_tokens)
+        return masked_tokens, masks
+
+
+class UViTBert(ImageBert):
+    def __init__(self, config):
+        super().__init__(config=config)
+
+        del self.model
+
+        self.embeddings = nn.Embedding(
+            self.target_codebook_size + self.condition_num_classes + 2,
+            self.hidden_size)
+        
+        self.pos_embed = nn.init.trunc_normal_(
+            nn.Parameter(torch.zeros(1, self.config.model.generator.image_seq_len + 1, self.hidden_size)), 0., 0.02)
+ 
+        self.in_blocks = nn.ModuleList([
+            UViTBlock(
+                dim=self.hidden_size, num_heads=self.num_attention_heads, mlp_ratio=(self.intermediate_size / self.hidden_size),
+                qkv_bias=False, qk_scale=None, norm_layer=nn.LayerNorm, use_checkpoint=False)
+            for _ in range(self.num_hidden_layers // 2)])
+
+        self.mid_block = UViTBlock(
+                dim=self.hidden_size, num_heads=self.num_attention_heads, mlp_ratio=(self.intermediate_size / self.hidden_size),
+                qkv_bias=False, qk_scale=None, norm_layer=nn.LayerNorm, use_checkpoint=False)
+
+        self.out_blocks = nn.ModuleList([
+            UViTBlock(
+                dim=self.hidden_size, num_heads=self.num_attention_heads, mlp_ratio=(self.intermediate_size / self.hidden_size),
+                qkv_bias=False, qk_scale=None, norm_layer=nn.LayerNorm, skip=True, use_checkpoint=False)
+            for _ in range(self.num_hidden_layers // 2)])
+
+        self.norm = nn.LayerNorm(self.hidden_size)
+        self.lm_head = nn.Linear(self.hidden_size,
+                                 self.target_codebook_size, bias=True)
+        self.apply(self._init_weights)
+
+    def _init_weights(self, m):
+        if isinstance(m, nn.Linear):
+            nn.init.trunc_normal_(m.weight, std=.02)
+            if isinstance(m, nn.Linear) and m.bias is not None:
+                nn.init.constant_(m.bias, 0)
+        elif isinstance(m, nn.Embedding):
+            m.weight.data = nn.init.trunc_normal_(m.weight.data, mean=0.0, std=0.02)
+        elif isinstance(m, nn.LayerNorm):
+            nn.init.constant_(m.bias, 0)
+            nn.init.constant_(m.weight, 1.0)
+
+    def forward(self, input_ids=None, condition=None, cond_drop_prob=0.1):
+        # Token space:
+        #  [0, codebook_size - 1]                       : those are the learned quantized image tokens
+        #  codebook_size                                : the mask token used to mask image tokens
+        #  [codebook_size + 1, codebook_size + nclass]  : the imagenet class tokens
+        #  codebook_size + 1 + nclass                   : the class drop label
+        drop_label_mask = torch.rand_like(condition, dtype=torch.float) < cond_drop_prob
+        # Shift the classes
+        condition = condition + self.target_codebook_size + 1  # [0, 999] -> [codebook_size + 1, codebook_size + 999]
+        condition[drop_label_mask] = self.condition_num_classes + self.target_codebook_size + 1
+        # prepend condition token
+        if input_ids is not None:
+            input_ids = torch.cat([condition.view(condition.shape[0], -1),
+                                   input_ids.view(input_ids.shape[0], -1),], dim=1)
+        else:
+            # at least there should be masked token
+            raise NotImplementedError
+        # UViT forward
+        embeddings = self.embeddings(input_ids)
+        x = embeddings + self.pos_embed[:, :embeddings.shape[1]]
+        skips = []
+        for blk in self.in_blocks:
+            x = blk(x)
+            skips.append(x)
+        x = self.mid_block(x)
+        for blk in self.out_blocks:
+            x = blk(x, skips.pop())
+        x = self.norm(x)
+        return self.lm_head(x[:, 1:]) # remove cond

modeling/modules/__init__.py

@@ -0,0 +1,6 @@
+from .base_model import BaseModel
+from .ema_model import EMAModel
+from .losses import ReconstructionLoss_Stage1, ReconstructionLoss_Stage2, MLMLoss, ARLoss
+from .blocks import TiTokEncoder, TiTokDecoder, UViTBlock
+from .maskgit_vqgan import Decoder as Pixel_Decoder
+from .maskgit_vqgan import VectorQuantizer as Pixel_Quantizer

modeling/modules/base_model.py

@@ -0,0 +1,127 @@
+"""This file contains some base class implementation for models.
+
+This file may have been modified by Bytedance Ltd. and/or its affiliates (“Bytedance's Modifications”).
+All Bytedance's Modifications are Copyright (year) Bytedance Ltd. and/or its affiliates. 
+
+Reference:
+    https://github.com/huggingface/open-muse/blob/main/muse/modeling_utils.py
+"""
+import os
+from typing import Union, Callable, Dict, Optional
+
+import torch
+
+
+class BaseModel(torch.nn.Module):
+
+    def __init__(self):
+        super().__init__()
+
+    def save_pretrained_weight(
+        self,
+        save_directory: Union[str, os.PathLike],
+        save_function: Callable = None,
+        state_dict: Optional[Dict[str, torch.Tensor]] = None,
+    ):
+        """Saves a model and its configuration file to a directory.
+
+        Args:
+            save_directory: A string or os.PathLike, directory to which to save. 
+                Will be created if it doesn't exist.
+            save_function: A Callable function, the function to use to save the state dictionary.
+                Useful on distributed training like TPUs when one need to replace `torch.save` by
+                another method. Can be configured with the environment variable `DIFFUSERS_SAVE_MODE`.
+            state_dict: A dictionary from str to torch.Tensor, the state dictionary to save.
+                If `None`, the model's state dictionary will be saved.
+        """
+        if os.path.isfile(save_directory):
+            print(f"Provided path ({save_directory}) should be a directory, not a file")
+            return
+
+        if save_function is None:
+            save_function = torch.save
+
+        os.makedirs(save_directory, exist_ok=True)
+
+        model_to_save = self
+
+        if state_dict is None:
+            state_dict = model_to_save.state_dict()
+        weights_name = "pytorch_model.bin"
+
+        save_function(state_dict, os.path.join(save_directory, weights_name))
+
+        print(f"Model weights saved in {os.path.join(save_directory, weights_name)}")
+
+    def load_pretrained_weight(
+        self,
+        pretrained_model_path: Union[str, os.PathLike],
+        strict_loading: bool = True,
+        torch_dtype: Optional[torch.dtype] = None
+    ):
+        r"""Instantiates a pretrained pytorch model from a pre-trained model configuration.
+
+        The model is set in evaluation mode by default using `model.eval()` (Dropout modules are deactivated). To train
+        the model, you should first set it back in training mode with `model.train()`.
+
+        Args:
+            pretrained_model_path: A string or os.PathLike, a path to a *directory* or *file* containing model weights.
+
+        Raises:
+            ValueError: If pretrained_model_path does not exist.
+        """
+        # If pretrained_model_path is a file, set model_file to this file.
+        if os.path.isfile(pretrained_model_path):
+            model_file = pretrained_model_path
+        # If pretrained_model_path is a directory, set model_file to the path of the 
+        # file "pytorch_model.bin" in this directory.
+        elif os.path.isdir(pretrained_model_path):
+            pretrained_model_path = os.path.join(pretrained_model_path, "pytorch_model.bin")
+            if os.path.isfile(pretrained_model_path):
+                model_file = pretrained_model_path
+            else:
+                raise ValueError(f"{pretrained_model_path} does not exist")
+        else:
+            raise ValueError(f"{pretrained_model_path} does not exist")
+
+        # Load model state from checkpoint.
+        checkpoint = torch.load(model_file, map_location="cpu")
+        # Load state dictionary into self.
+        msg = self.load_state_dict(checkpoint, strict=strict_loading)
+        # Print information about loading weights.
+        print(f"loading weight from {model_file}, msg: {msg}")
+        # If torch_dtype is specified and is a valid torch.dtype, convert self to this dtype.
+        if torch_dtype is not None and not isinstance(torch_dtype, torch.dtype):
+            raise ValueError(
+                f"{torch_dtype} needs to be of type `torch.dtype`, e.g. `torch.float16`, but is {type(torch_dtype)}."
+            )
+        elif torch_dtype is not None:
+            self.to(torch_dtype)
+
+        # Set model in evaluation mode to deactivate DropOut modules by default.
+        self.eval()
+
+    def num_parameters(self, only_trainable: bool = False, exclude_embeddings: bool = False) -> int:
+        """Gets the number of parameters in the module.
+
+        Args:
+            only_trainable: A boolean, whether to only include trainable parameters.
+            exclude_embeddings: A boolean, whether to exclude parameters associated with embeddings.
+
+        Returns:
+            An integer, the number of parameters.
+        """
+
+        if exclude_embeddings:
+            embedding_param_names = [
+                f"{name}.weight"
+                for name, module_type in self.named_modules()
+                if isinstance(module_type, torch.nn.Embedding)
+            ]
+            non_embedding_parameters = [
+                parameter for name, parameter in self.named_parameters() if name not in embedding_param_names
+            ]
+            return sum(p.numel() for p in non_embedding_parameters if p.requires_grad or not only_trainable)
+        else:
+            return sum(p.numel() for p in self.parameters() if p.requires_grad or not only_trainable)
+

modeling/modules/blocks.py

@@ -0,0 +1,376 @@
+"""Building blocks for TiTok.
+
+Copyright (2024) Bytedance Ltd. and/or its affiliates
+
+Licensed under the Apache License, Version 2.0 (the "License"); 
+you may not use this file except in compliance with the License. 
+You may obtain a copy of the License at 
+
+    http://www.apache.org/licenses/LICENSE-2.0 
+
+Unless required by applicable law or agreed to in writing, software 
+distributed under the License is distributed on an "AS IS" BASIS, 
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 
+See the License for the specific language governing permissions and 
+limitations under the License. 
+
+Reference: 
+    https://github.com/mlfoundations/open_clip/blob/main/src/open_clip/transformer.py
+    https://github.com/baofff/U-ViT/blob/main/libs/timm.py
+"""
+
+import torch
+import torch.nn as nn
+from collections import OrderedDict
+import einops
+from einops.layers.torch import Rearrange
+
+
+class ResidualAttentionBlock(nn.Module):
+    def __init__(
+            self,
+            d_model,
+            n_head,
+            mlp_ratio = 4.0,
+            act_layer = nn.GELU,
+            norm_layer = nn.LayerNorm
+        ):
+        super().__init__()
+
+        self.ln_1 = norm_layer(d_model)
+        self.attn = nn.MultiheadAttention(d_model, n_head)
+        self.mlp_ratio = mlp_ratio
+        # optionally we can disable the FFN
+        if mlp_ratio > 0:
+            self.ln_2 = norm_layer(d_model)
+            mlp_width = int(d_model * mlp_ratio)
+            self.mlp = nn.Sequential(OrderedDict([
+                ("c_fc", nn.Linear(d_model, mlp_width)),
+                ("gelu", act_layer()),
+                ("c_proj", nn.Linear(mlp_width, d_model))
+            ]))
+
+    def attention(
+            self,
+            x: torch.Tensor
+    ):
+        return self.attn(x, x, x, need_weights=False)[0]
+
+    def forward(
+            self,
+            x: torch.Tensor,
+    ):
+        attn_output = self.attention(x=self.ln_1(x))
+        x = x + attn_output
+        if self.mlp_ratio > 0:
+            x = x + self.mlp(self.ln_2(x))
+        return x
+
+if hasattr(torch.nn.functional, 'scaled_dot_product_attention'):
+    ATTENTION_MODE = 'flash'
+else:
+    try:
+        import xformers
+        import xformers.ops
+        ATTENTION_MODE = 'xformers'
+    except:
+        ATTENTION_MODE = 'math'
+print(f'attention mode is {ATTENTION_MODE}')
+
+
+class Attention(nn.Module):
+    def __init__(self, dim, num_heads=8, qkv_bias=False, qk_scale=None, attn_drop=0., proj_drop=0.):
+        super().__init__()
+        self.num_heads = num_heads
+        head_dim = dim // num_heads
+        self.scale = qk_scale or head_dim ** -0.5
+        self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias)
+        self.attn_drop = nn.Dropout(attn_drop)
+        self.proj = nn.Linear(dim, dim)
+        self.proj_drop = nn.Dropout(proj_drop)
+
+    def forward(self, x):
+        B, L, C = x.shape
+
+        qkv = self.qkv(x)
+        if ATTENTION_MODE == 'flash':
+            qkv = einops.rearrange(qkv, 'B L (K H D) -> K B H L D', K=3, H=self.num_heads).float()
+            q, k, v = qkv[0], qkv[1], qkv[2]  # B H L D
+            x = torch.nn.functional.scaled_dot_product_attention(q, k, v)
+            x = einops.rearrange(x, 'B H L D -> B L (H D)')
+        elif ATTENTION_MODE == 'xformers':
+            qkv = einops.rearrange(qkv, 'B L (K H D) -> K B L H D', K=3, H=self.num_heads)
+            q, k, v = qkv[0], qkv[1], qkv[2]  # B L H D
+            x = xformers.ops.memory_efficient_attention(q, k, v)
+            x = einops.rearrange(x, 'B L H D -> B L (H D)', H=self.num_heads)
+        elif ATTENTION_MODE == 'math':
+            qkv = einops.rearrange(qkv, 'B L (K H D) -> K B H L D', K=3, H=self.num_heads)
+            q, k, v = qkv[0], qkv[1], qkv[2]  # B H L D
+            attn = (q @ k.transpose(-2, -1)) * self.scale
+            attn = attn.softmax(dim=-1)
+            attn = self.attn_drop(attn)
+            x = (attn @ v).transpose(1, 2).reshape(B, L, C)
+        else:
+            raise NotImplemented
+
+        x = self.proj(x)
+        x = self.proj_drop(x)
+        return x
+
+
+def drop_path(x, drop_prob: float = 0., training: bool = False):
+    """Drop paths (Stochastic Depth) per sample (when applied in main path of residual blocks).
+
+    This is the same as the DropConnect impl I created for EfficientNet, etc networks, however,
+    the original name is misleading as 'Drop Connect' is a different form of dropout in a separate paper...
+    See discussion: https://github.com/tensorflow/tpu/issues/494#issuecomment-532968956 ... I've opted for
+    changing the layer and argument names to 'drop path' rather than mix DropConnect as a layer name and use
+    'survival rate' as the argument.
+
+    """
+    if drop_prob == 0. or not training:
+        return x
+    keep_prob = 1 - drop_prob
+    shape = (x.shape[0],) + (1,) * (x.ndim - 1)  # work with diff dim tensors, not just 2D ConvNets
+    random_tensor = keep_prob + torch.rand(shape, dtype=x.dtype, device=x.device)
+    random_tensor.floor_()  # binarize
+    output = x.div(keep_prob) * random_tensor
+    return output
+
+
+class DropPath(nn.Module):
+    """Drop paths (Stochastic Depth) per sample  (when applied in main path of residual blocks).
+    """
+    def __init__(self, drop_prob=None):
+        super(DropPath, self).__init__()
+        self.drop_prob = drop_prob
+
+    def forward(self, x):
+        return drop_path(x, self.drop_prob, self.training)
+
+
+class Mlp(nn.Module):
+    def __init__(self, in_features, hidden_features=None, out_features=None, act_layer=nn.GELU, drop=0.):
+        super().__init__()
+        out_features = out_features or in_features
+        hidden_features = hidden_features or in_features
+        self.fc1 = nn.Linear(in_features, hidden_features)
+        self.act = act_layer()
+        self.fc2 = nn.Linear(hidden_features, out_features)
+        self.drop = nn.Dropout(drop)
+
+    def forward(self, x):
+        x = self.fc1(x)
+        x = self.act(x)
+        x = self.drop(x)
+        x = self.fc2(x)
+        x = self.drop(x)
+        return x
+
+
+class UViTBlock(nn.Module):
+
+    def __init__(self, dim, num_heads, mlp_ratio=4., qkv_bias=False, qk_scale=None, drop=0., attn_drop=0.,
+                 drop_path=0., act_layer=nn.GELU, norm_layer=nn.LayerNorm, skip=False, use_checkpoint=False):
+        super().__init__()
+        self.norm1 = norm_layer(dim)
+        self.attn = Attention(
+            dim, num_heads=num_heads, qkv_bias=qkv_bias, qk_scale=qk_scale, attn_drop=attn_drop, proj_drop=drop)
+        # NOTE: drop path for stochastic depth, we shall see if this is better than dropout here
+        self.drop_path = DropPath(drop_path) if drop_path > 0. else nn.Identity()
+        self.norm2 = norm_layer(dim)
+        mlp_hidden_dim = int(dim * mlp_ratio)
+        self.mlp = Mlp(in_features=dim, hidden_features=mlp_hidden_dim, act_layer=act_layer, drop=drop)
+        self.skip_linear = nn.Linear(2 * dim, dim) if skip else None
+        self.use_checkpoint = use_checkpoint
+
+    def forward(self, x, skip=None):
+        if self.use_checkpoint:
+            return torch.utils.checkpoint.checkpoint(self._forward, x, skip)
+        else:
+            return self._forward(x, skip)
+
+    def _forward(self, x, skip=None):
+        if self.skip_linear is not None:
+            x = self.skip_linear(torch.cat([x, skip], dim=-1))
+        x = x + self.drop_path(self.attn(self.norm1(x)))
+        x = x + self.drop_path(self.mlp(self.norm2(x)))
+        return x
+    
+
+def _expand_token(token, batch_size: int):
+    return token.unsqueeze(0).expand(batch_size, -1, -1)
+
+
+class TiTokEncoder(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        self.config = config
+        self.image_size = config.dataset.preprocessing.crop_size 
+        self.patch_size = config.model.vq_model.vit_enc_patch_size
+        self.grid_size = self.image_size // self.patch_size
+        self.model_size = config.model.vq_model.vit_enc_model_size
+        self.num_latent_tokens = config.model.vq_model.num_latent_tokens
+        self.token_size = config.model.vq_model.token_size
+
+        if config.model.vq_model.get("quantize_mode", "vq") == "vae":
+            self.token_size = self.token_size * 2 # needs to split into mean and std
+
+        self.is_legacy = config.model.vq_model.get("is_legacy", True)
+
+        self.width = {
+                "small": 512,
+                "base": 768,
+                "large": 1024,
+            }[self.model_size]
+        self.num_layers = {
+                "small": 8,
+                "base": 12,
+                "large": 24,
+            }[self.model_size]
+        self.num_heads = {
+                "small": 8,
+                "base": 12,
+                "large": 16,
+            }[self.model_size]
+        
+        self.patch_embed = nn.Conv2d(
+            in_channels=3, out_channels=self.width,
+              kernel_size=self.patch_size, stride=self.patch_size, bias=True)
+        
+        scale = self.width ** -0.5
+        self.class_embedding = nn.Parameter(scale * torch.randn(1, self.width))
+        self.positional_embedding = nn.Parameter(
+                scale * torch.randn(self.grid_size ** 2 + 1, self.width))
+        self.latent_token_positional_embedding = nn.Parameter(
+            scale * torch.randn(self.num_latent_tokens, self.width))
+        self.ln_pre = nn.LayerNorm(self.width)
+        self.transformer = nn.ModuleList()
+        for i in range(self.num_layers):
+            self.transformer.append(ResidualAttentionBlock(
+                self.width, self.num_heads, mlp_ratio=4.0
+            ))
+        self.ln_post = nn.LayerNorm(self.width)
+        self.conv_out = nn.Conv2d(self.width, self.token_size, kernel_size=1, bias=True)
+
+    def forward(self, pixel_values, latent_tokens):
+        batch_size = pixel_values.shape[0]
+        x = pixel_values
+        x = self.patch_embed(x)
+        x = x.reshape(x.shape[0], x.shape[1], -1)
+        x = x.permute(0, 2, 1) # shape = [*, grid ** 2, width]
+        # class embeddings and positional embeddings
+        x = torch.cat([_expand_token(self.class_embedding, x.shape[0]).to(x.dtype), x], dim=1)
+        x = x + self.positional_embedding.to(x.dtype) # shape = [*, grid ** 2 + 1, width]
+        
+
+        latent_tokens = _expand_token(latent_tokens, x.shape[0]).to(x.dtype)
+        latent_tokens = latent_tokens + self.latent_token_positional_embedding.to(x.dtype)
+        x = torch.cat([x, latent_tokens], dim=1)
+
+        x = self.ln_pre(x)
+        x = x.permute(1, 0, 2)  # NLD -> LND
+        for i in range(self.num_layers):
+            x = self.transformer[i](x)
+        x = x.permute(1, 0, 2)  # LND -> NLD
+        
+        latent_tokens = x[:, 1+self.grid_size**2:]
+        latent_tokens = self.ln_post(latent_tokens)
+        # fake 2D shape
+        if self.is_legacy:
+            latent_tokens = latent_tokens.reshape(batch_size, self.width, self.num_latent_tokens, 1)
+        else:
+            # Fix legacy problem.
+            latent_tokens = latent_tokens.reshape(batch_size, self.num_latent_tokens, self.width, 1).permute(0, 2, 1, 3)
+        latent_tokens = self.conv_out(latent_tokens)
+        latent_tokens = latent_tokens.reshape(batch_size, self.token_size, 1, self.num_latent_tokens)
+        return latent_tokens
+    
+
+class TiTokDecoder(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        self.config = config
+        self.image_size = config.dataset.preprocessing.crop_size
+        self.patch_size = config.model.vq_model.vit_dec_patch_size
+        self.grid_size = self.image_size // self.patch_size
+        self.model_size = config.model.vq_model.vit_dec_model_size
+        self.num_latent_tokens = config.model.vq_model.num_latent_tokens
+        self.token_size = config.model.vq_model.token_size
+        self.is_legacy = config.model.vq_model.get("is_legacy", True)
+        self.width = {
+                "small": 512,
+                "base": 768,
+                "large": 1024,
+            }[self.model_size]
+        self.num_layers = {
+                "small": 8,
+                "base": 12,
+                "large": 24,
+            }[self.model_size]
+        self.num_heads = {
+                "small": 8,
+                "base": 12,
+                "large": 16,
+            }[self.model_size]
+
+        self.decoder_embed = nn.Linear(
+            self.token_size, self.width, bias=True)
+        scale = self.width ** -0.5
+        self.class_embedding = nn.Parameter(scale * torch.randn(1, self.width))
+        self.positional_embedding = nn.Parameter(
+                scale * torch.randn(self.grid_size ** 2 + 1, self.width))
+        # add mask token and query pos embed
+        self.mask_token = nn.Parameter(scale * torch.randn(1, 1, self.width))
+        self.latent_token_positional_embedding = nn.Parameter(
+            scale * torch.randn(self.num_latent_tokens, self.width))
+        self.ln_pre = nn.LayerNorm(self.width)
+        self.transformer = nn.ModuleList()
+        for i in range(self.num_layers):
+            self.transformer.append(ResidualAttentionBlock(
+                self.width, self.num_heads, mlp_ratio=4.0
+            ))
+        self.ln_post = nn.LayerNorm(self.width)
+
+        if self.is_legacy:
+            self.ffn = nn.Sequential(
+                nn.Conv2d(self.width, 2 * self.width, 1, padding=0, bias=True),
+                nn.Tanh(),
+                nn.Conv2d(2 * self.width, 1024, 1, padding=0, bias=True),
+            )
+            self.conv_out = nn.Identity()
+        else:
+            # Directly predicting RGB pixels
+            self.ffn = nn.Sequential(
+                nn.Conv2d(self.width, self.patch_size * self.patch_size * 3, 1, padding=0, bias=True),
+                Rearrange('b (p1 p2 c) h w -> b c (h p1) (w p2)',
+                    p1 = self.patch_size, p2 = self.patch_size),)
+            self.conv_out = nn.Conv2d(3, 3, 3, padding=1, bias=True)
+    
+    def forward(self, z_quantized):
+        N, C, H, W = z_quantized.shape
+        assert H == 1 and W == self.num_latent_tokens, f"{H}, {W}, {self.num_latent_tokens}"
+        x = z_quantized.reshape(N, C*H, W).permute(0, 2, 1) # NLD
+        x = self.decoder_embed(x)
+
+        batchsize, seq_len, _ = x.shape
+
+        mask_tokens = self.mask_token.repeat(batchsize, self.grid_size**2, 1).to(x.dtype)
+        mask_tokens = torch.cat([_expand_token(self.class_embedding, mask_tokens.shape[0]).to(mask_tokens.dtype),
+                                    mask_tokens], dim=1)
+        mask_tokens = mask_tokens + self.positional_embedding.to(mask_tokens.dtype)
+        x = x + self.latent_token_positional_embedding[:seq_len]
+        x = torch.cat([mask_tokens, x], dim=1)
+        
+        x = self.ln_pre(x)
+        x = x.permute(1, 0, 2)  # NLD -> LND
+        for i in range(self.num_layers):
+            x = self.transformer[i](x)
+        x = x.permute(1, 0, 2)  # LND -> NLD
+        x = x[:, 1:1+self.grid_size**2] # remove cls embed
+        x = self.ln_post(x)
+        # N L D -> N D H W
+        x = x.permute(0, 2, 1).reshape(batchsize, self.width, self.grid_size, self.grid_size)
+        x = self.ffn(x.contiguous())
+        x = self.conv_out(x)
+        return x

modeling/modules/discriminator.py

@@ -0,0 +1,141 @@
+"""This file contains some base implementation for discrminators.
+
+Copyright (2024) Bytedance Ltd. and/or its affiliates
+
+Licensed under the Apache License, Version 2.0 (the "License"); 
+you may not use this file except in compliance with the License. 
+You may obtain a copy of the License at 
+
+    http://www.apache.org/licenses/LICENSE-2.0 
+
+Unless required by applicable law or agreed to in writing, software 
+distributed under the License is distributed on an "AS IS" BASIS, 
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 
+See the License for the specific language governing permissions and 
+limitations under the License. 
+
+TODO: Add reference to Mark Weber's tech report on the improved discriminator architecture.
+"""
+import functools
+import math
+from typing import Tuple
+
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from .maskgit_vqgan import Conv2dSame
+
+
+class BlurBlock(torch.nn.Module):
+    def __init__(self,
+                 kernel: Tuple[int] = (1, 3, 3, 1)
+                 ):
+        super().__init__()
+
+        kernel = torch.tensor(kernel, dtype=torch.float32, requires_grad=False)
+        kernel = kernel[None, :] * kernel[:, None]
+        kernel /= kernel.sum()
+        kernel = kernel.unsqueeze(0).unsqueeze(0)
+        self.register_buffer("kernel", kernel)
+
+    def calc_same_pad(self, i: int, k: int, s: int) -> int:
+        return max((math.ceil(i / s) - 1) * s + (k - 1) + 1 - i, 0)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        ic, ih, iw = x.size()[-3:]
+        pad_h = self.calc_same_pad(i=ih, k=4, s=2)
+        pad_w = self.calc_same_pad(i=iw, k=4, s=2)
+        if pad_h > 0 or pad_w > 0:
+            x = F.pad(x, [pad_w // 2, pad_w - pad_w // 2, pad_h // 2, pad_h - pad_h // 2])
+
+        weight = self.kernel.expand(ic, -1, -1, -1)
+
+        out = F.conv2d(input=x, weight=weight, stride=2, groups=x.shape[1])
+        return out
+
+
+class NLayerDiscriminator(torch.nn.Module):
+    def __init__(
+        self,
+        num_channels: int = 3,
+        hidden_channels: int = 128,
+        num_stages: int = 3,
+        blur_resample: bool = True,
+        blur_kernel_size: int = 4
+    ):
+        """ Initializes the NLayerDiscriminator.
+
+        Args:
+            num_channels -> int: The number of input channels.
+            hidden_channels -> int: The number of hidden channels.
+            num_stages -> int: The number of stages.
+            blur_resample -> bool: Whether to use blur resampling.
+            blur_kernel_size -> int: The blur kernel size.
+        """
+        super().__init__()
+        assert num_stages > 0, "Discriminator cannot have 0 stages"
+        assert (not blur_resample) or (blur_kernel_size >= 3 and blur_kernel_size <= 5), "Blur kernel size must be in [3,5] when sampling]"
+
+        in_channel_mult = (1,) + tuple(map(lambda t: 2**t, range(num_stages)))
+        init_kernel_size = 5
+        activation = functools.partial(torch.nn.LeakyReLU, negative_slope=0.1)
+
+        self.block_in = torch.nn.Sequential(
+            Conv2dSame(
+                num_channels,
+                hidden_channels,
+                kernel_size=init_kernel_size
+            ),
+            activation(),
+        )
+
+        BLUR_KERNEL_MAP = {
+            3: (1,2,1),
+            4: (1,3,3,1),
+            5: (1,4,6,4,1),
+        }
+
+        discriminator_blocks = []
+        for i_level in range(num_stages):
+            in_channels = hidden_channels * in_channel_mult[i_level]
+            out_channels = hidden_channels * in_channel_mult[i_level + 1]
+            block = torch.nn.Sequential(
+                Conv2dSame(
+                    in_channels,
+                    out_channels,
+                    kernel_size=3,
+                ),
+                torch.nn.AvgPool2d(kernel_size=2, stride=2) if not blur_resample else BlurBlock(BLUR_KERNEL_MAP[blur_kernel_size]),
+                torch.nn.GroupNorm(32, out_channels),
+                activation(),
+            )
+            discriminator_blocks.append(block)
+
+        self.blocks = torch.nn.ModuleList(discriminator_blocks)
+
+        self.pool = torch.nn.AdaptiveMaxPool2d((16, 16))
+
+        self.to_logits = torch.nn.Sequential(
+            Conv2dSame(out_channels, out_channels, 1),
+            activation(),
+            Conv2dSame(out_channels, 1, kernel_size=5)
+        )
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        """ Forward pass.
+
+        Args:
+            x -> torch.Tensor: The input tensor.
+
+        Returns:
+            output -> torch.Tensor: The output tensor.
+        """
+        hidden_states = self.block_in(x)
+        for block in self.blocks:
+            hidden_states = block(hidden_states)
+
+        hidden_states = self.pool(hidden_states)
+
+        return self.to_logits(hidden_states)

modeling/modules/ema_model.py

@@ -0,0 +1,244 @@
+"""This file contains some base class implementation for EMA.
+
+This file may have been modified by Bytedance Ltd. and/or its affiliates (“Bytedance's Modifications”).
+All Bytedance's Modifications are Copyright (year) Bytedance Ltd. and/or its affiliates. 
+
+Reference:
+    https://github.com/huggingface/open-muse/blob/64e1afe033717d795866ab8204484705cd4dc3f7/muse/modeling_ema.py#L8
+"""
+
+
+import copy
+from typing import Any, Iterable, Optional, Union
+
+import torch
+
+
+class EMAModel:
+    """Exponential Moving Average of models weights."""
+    def __init__(
+        self,
+        parameters: Iterable[torch.nn.Parameter],
+        decay: float = 0.9999,
+        min_decay: float = 0.0,
+        update_after_step: int = 0,
+        update_every: int = 1,
+        current_step: int = 0,
+        use_ema_warmup: bool = False,
+        inv_gamma: Union[float, int] = 1.0,
+        power: Union[float, int] = 2 / 3,
+        model_cls: Optional[Any] = None,
+        **model_config_kwargs
+    ):
+        """
+        Args:
+            parameters (Iterable[torch.nn.Parameter]): The parameters to track.
+            decay (float): The decay factor for the exponential moving average.
+            min_decay (float): The minimum decay factor for the exponential moving average.
+            update_after_step (int): The number of steps to wait before starting to update the EMA weights.
+            update_every (int): The number of steps between each EMA update.
+            current_step (int): The current training step.
+            use_ema_warmup (bool): Whether to use EMA warmup.
+            inv_gamma (float):
+                Inverse multiplicative factor of EMA warmup. Default: 1. Only used if `use_ema_warmup` is True.
+            power (float): Exponential factor of EMA warmup. Default: 2/3. Only used if `use_ema_warmup` is True.
+
+        notes on EMA Warmup:
+            If gamma=1 and power=1, implements a simple average. gamma=1, power=2/3 are good values for models you plan
+            to train for a million or more steps (reaches decay factor 0.999 at 31.6K steps, 0.9999 at 1M steps),
+            gamma=1, power=3/4 for models you plan to train for less (reaches decay factor 0.999 at 10K steps, 0.9999
+            at 215.4k steps).
+        """
+
+        parameters = list(parameters)
+        self.shadow_params = [p.clone().detach() for p in parameters]
+        self.temp_stored_params = None
+
+        self.decay = decay
+        self.min_decay = min_decay
+        self.update_after_step = update_after_step
+        self.update_every = update_every
+        self.use_ema_warmup = use_ema_warmup
+        self.inv_gamma = inv_gamma
+        self.power = power
+        self.optimization_step = current_step
+        self.cur_decay_value = None  # set in `step()`
+
+        self.model_cls = model_cls
+        self.model_config_kwargs = model_config_kwargs
+
+    @classmethod
+    def from_pretrained(cls, checkpoint, model_cls, **model_config_kwargs) -> "EMAModel":
+        model = model_cls(**model_config_kwargs)
+        model.load_pretrained_weight(checkpoint)
+
+        ema_model = cls(model.parameters(), model_cls=model_cls, **model_config_kwargs)
+        return ema_model
+
+    def save_pretrained(self, path):
+        if self.model_cls is None:
+            raise ValueError("`save_pretrained` can only be used if `model_cls` was defined at __init__.")
+
+        if self.model_config_kwargs is None:
+            raise ValueError("`save_pretrained` can only be used if `model_config_kwargs` was defined at __init__.")
+
+        model = self.model_cls(**self.model_config_kwargs)
+        self.copy_to(model.parameters())
+        model.save_pretrained_weight(path)
+    
+    def set_step(self, optimization_step: int):
+        self.optimization_step = optimization_step
+
+    def get_decay(self, optimization_step: int) -> float:
+        """Computes the decay factor for the exponential moving average."""
+        step = max(0, optimization_step - self.update_after_step - 1)
+
+        if step <= 0:
+            return 0.0
+
+        if self.use_ema_warmup:
+            cur_decay_value = 1 - (1 + step / self.inv_gamma) ** -self.power
+        else:
+            cur_decay_value = (1 + step) / (10 + step)
+
+        cur_decay_value = min(cur_decay_value, self.decay)
+        # Make sure decay is not smaller than min_decay.
+        cur_decay_value = max(cur_decay_value, self.min_decay)
+        return cur_decay_value
+
+    @torch.no_grad()
+    def step(self, parameters: Iterable[torch.nn.Parameter]):
+        parameters = list(parameters)
+
+        self.optimization_step += 1
+
+        if (self.optimization_step - 1) % self.update_every != 0:
+            return
+
+        # Compute the decay factor for the exponential moving average.
+        decay = self.get_decay(self.optimization_step)
+        self.cur_decay_value = decay
+        one_minus_decay = 1 - decay
+
+        for s_param, param in zip(self.shadow_params, parameters):
+            if param.requires_grad:
+                s_param.sub_(one_minus_decay * (s_param - param))
+            else:
+                s_param.copy_(param)
+
+    def copy_to(self, parameters: Iterable[torch.nn.Parameter]) -> None:
+        """Copies current averaged parameters into given collection of parameters.
+
+        Args:
+            parameters: Iterable of `torch.nn.Parameter`; the parameters to be
+                updated with the stored moving averages. If `None`, the parameters with which this
+                `ExponentialMovingAverage` was initialized will be used.
+        """
+        parameters = list(parameters)
+        for s_param, param in zip(self.shadow_params, parameters):
+            param.data.copy_(s_param.to(param.device).data)
+
+    def to(self, device=None, dtype=None) -> None:
+        r"""Moves internal buffers of the ExponentialMovingAverage to `device`.
+
+        Args:
+            device: like `device` argument to `torch.Tensor.to`
+        """
+        # .to() on the tensors handles None correctly
+        self.shadow_params = [
+            p.to(device=device, dtype=dtype) if p.is_floating_point() else p.to(device=device)
+            for p in self.shadow_params
+        ]
+
+    def state_dict(self) -> dict:
+        r"""Returns the state of the ExponentialMovingAverage as a dict. This method is used by accelerate during
+        checkpointing to save the ema state dict.
+        """
+        # Following PyTorch conventions, references to tensors are returned:
+        # "returns a reference to the state and not its copy!" -
+        # https://pytorch.org/tutorials/beginner/saving_loading_models.html#what-is-a-state-dict
+        return {
+            "decay": self.decay,
+            "min_decay": self.min_decay,
+            "optimization_step": self.optimization_step,
+            "update_after_step": self.update_after_step,
+            "use_ema_warmup": self.use_ema_warmup,
+            "inv_gamma": self.inv_gamma,
+            "power": self.power,
+            "shadow_params": self.shadow_params,
+        }
+
+    def store(self, parameters: Iterable[torch.nn.Parameter]) -> None:
+        r"""
+        Args:
+        Save the current parameters for restoring later.
+            parameters: Iterable of `torch.nn.Parameter`; the parameters to be
+                temporarily stored.
+        """
+        self.temp_stored_params = [param.detach().cpu().clone() for param in parameters]
+
+    def restore(self, parameters: Iterable[torch.nn.Parameter]) -> None:
+        r"""Restores the parameters stored with the `store` method. Useful to validate
+          the model with EMA parameters without affecting the original optimization process.
+            Store the parameters before the `copy_to()` method. After validation (or
+              model saving), use this to restore the former parameters.
+        
+        Args:
+            parameters: Iterable of `torch.nn.Parameter`; the parameters to be
+                updated with the stored parameters. If `None`, the parameters with which this
+                `ExponentialMovingAverage` was initialized will be used.
+        """
+        if self.temp_stored_params is None:
+            raise RuntimeError("This ExponentialMovingAverage has no `store()`ed weights to `restore()`")
+        for c_param, param in zip(self.temp_stored_params, parameters):
+            param.data.copy_(c_param.data)
+
+        # Better memory-wise.
+        self.temp_stored_params = None
+
+    def load_state_dict(self, state_dict: dict) -> None:
+        r"""Loads the ExponentialMovingAverage state. This method is used by accelerate during checkpointing to save the
+        ema state dict.
+        
+        Args: 
+            state_dict (dict): EMA state. Should be an object returned
+                from a call to :meth:`state_dict`.
+        """
+        # Deepcopy, to be consistent with module API
+        state_dict = copy.deepcopy(state_dict)
+
+        self.decay = state_dict.get("decay", self.decay)
+        if self.decay < 0.0 or self.decay > 1.0:
+            raise ValueError("Decay must be between 0 and 1")
+
+        self.min_decay = state_dict.get("min_decay", self.min_decay)
+        if not isinstance(self.min_decay, float):
+            raise ValueError("Invalid min_decay")
+
+        self.optimization_step = state_dict.get("optimization_step", self.optimization_step)
+        if not isinstance(self.optimization_step, int):
+            raise ValueError("Invalid optimization_step")
+
+        self.update_after_step = state_dict.get("update_after_step", self.update_after_step)
+        if not isinstance(self.update_after_step, int):
+            raise ValueError("Invalid update_after_step")
+
+        self.use_ema_warmup = state_dict.get("use_ema_warmup", self.use_ema_warmup)
+        if not isinstance(self.use_ema_warmup, bool):
+            raise ValueError("Invalid use_ema_warmup")
+
+        self.inv_gamma = state_dict.get("inv_gamma", self.inv_gamma)
+        if not isinstance(self.inv_gamma, (float, int)):
+            raise ValueError("Invalid inv_gamma")
+
+        self.power = state_dict.get("power", self.power)
+        if not isinstance(self.power, (float, int)):
+            raise ValueError("Invalid power")
+
+        shadow_params = state_dict.get("shadow_params", None)
+        if shadow_params is not None:
+            self.shadow_params = shadow_params
+            if not isinstance(self.shadow_params, list):
+                raise ValueError("shadow_params must be a list")
+            if not all(isinstance(p, torch.Tensor) for p in self.shadow_params):
+                raise ValueError("shadow_params must all be Tensors")

modeling/modules/losses.py

@@ -0,0 +1,293 @@
+"""This files contains training loss implementation.
+
+Copyright (2024) Bytedance Ltd. and/or its affiliates
+
+Licensed under the Apache License, Version 2.0 (the "License"); 
+you may not use this file except in compliance with the License. 
+You may obtain a copy of the License at 
+
+    http://www.apache.org/licenses/LICENSE-2.0 
+
+Unless required by applicable law or agreed to in writing, software 
+distributed under the License is distributed on an "AS IS" BASIS, 
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 
+See the License for the specific language governing permissions and 
+limitations under the License. 
+
+Ref:
+    https://github.com/CompVis/taming-transformers/blob/master/taming/modules/losses/vqperceptual.py
+"""
+from typing import Mapping, Text, Tuple
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from einops import rearrange
+from torch.cuda.amp import autocast
+from .perceptual_loss import PerceptualLoss
+from .discriminator import NLayerDiscriminator
+
+
+def hinge_d_loss(logits_real: torch.Tensor, logits_fake: torch.Tensor) -> torch.Tensor:
+    """Hinge loss for discrminator.
+
+    This function is borrowed from
+    https://github.com/CompVis/taming-transformers/blob/master/taming/modules/losses/vqperceptual.py#L20
+    """
+    loss_real = torch.mean(F.relu(1.0 - logits_real))
+    loss_fake = torch.mean(F.relu(1.0 + logits_fake))
+    d_loss = 0.5 * (loss_real + loss_fake)
+    return d_loss
+
+
+def compute_lecam_loss(
+    logits_real_mean: torch.Tensor,
+    logits_fake_mean: torch.Tensor,
+    ema_logits_real_mean: torch.Tensor,
+    ema_logits_fake_mean: torch.Tensor
+) -> torch.Tensor:
+    """Computes the LeCam loss for the given average real and fake logits.
+
+    Args:
+        logits_real_mean -> torch.Tensor: The average real logits.
+        logits_fake_mean -> torch.Tensor: The average fake logits.
+        ema_logits_real_mean -> torch.Tensor: The EMA of the average real logits.
+        ema_logits_fake_mean -> torch.Tensor: The EMA of the average fake logits.
+
+    Returns:
+        lecam_loss -> torch.Tensor: The LeCam loss.
+    """
+    lecam_loss = torch.mean(torch.pow(F.relu(logits_real_mean - ema_logits_fake_mean), 2))
+    lecam_loss += torch.mean(torch.pow(F.relu(ema_logits_real_mean - logits_fake_mean), 2))
+    return lecam_loss
+
+
+class ReconstructionLoss_Stage1(torch.nn.Module):
+    def __init__(
+        self,
+        config
+    ):
+        super().__init__()
+        loss_config = config.losses
+        self.quantizer_weight = loss_config.quantizer_weight
+        self.target_codebook_size = 1024
+
+    def forward(self,
+                target_codes: torch.Tensor,
+                reconstructions: torch.Tensor,
+                quantizer_loss: torch.Tensor,
+                ) -> Tuple[torch.Tensor, Mapping[Text, torch.Tensor]]:
+        return self._forward_generator(target_codes, reconstructions, quantizer_loss)
+
+    def _forward_generator(self,
+                           target_codes: torch.Tensor,
+                           reconstructions: torch.Tensor,
+                           quantizer_loss: Mapping[Text, torch.Tensor],
+                           ) -> Tuple[torch.Tensor, Mapping[Text, torch.Tensor]]:
+        reconstructions = reconstructions.contiguous()
+        loss_fct = nn.CrossEntropyLoss(reduction="mean")
+        batch_size = reconstructions.shape[0]
+        reconstruction_loss = loss_fct(reconstructions.view(batch_size, self.target_codebook_size, -1),
+                                        target_codes.view(batch_size, -1))
+        total_loss = reconstruction_loss + \
+            self.quantizer_weight * quantizer_loss["quantizer_loss"]
+
+        loss_dict = dict(
+            total_loss=total_loss.clone().detach(),
+            reconstruction_loss=reconstruction_loss.detach(),
+            quantizer_loss=(self.quantizer_weight * quantizer_loss["quantizer_loss"]).detach(),
+            commitment_loss=quantizer_loss["commitment_loss"].detach(),
+            codebook_loss=quantizer_loss["codebook_loss"].detach(),
+        )
+
+        return total_loss, loss_dict
+
+
+class ReconstructionLoss_Stage2(torch.nn.Module):
+    def __init__(
+        self,
+        config
+    ):
+        """Initializes the losses module.
+
+        Args:
+            config: A dictionary, the configuration for the model and everything else.
+        """
+        super().__init__()
+        loss_config = config.losses
+        self.discriminator = NLayerDiscriminator()
+
+        self.reconstruction_loss = loss_config.reconstruction_loss
+        self.reconstruction_weight = loss_config.reconstruction_weight
+        self.quantizer_weight = loss_config.quantizer_weight
+        self.perceptual_loss = PerceptualLoss(
+            loss_config.perceptual_loss).eval()
+        self.perceptual_weight = loss_config.perceptual_weight
+        self.discriminator_iter_start = loss_config.discriminator_start
+
+        self.discriminator_factor = loss_config.discriminator_factor
+        self.discriminator_weight = loss_config.discriminator_weight
+        self.lecam_regularization_weight = loss_config.lecam_regularization_weight
+        self.lecam_ema_decay = loss_config.get("lecam_ema_decay", 0.999)
+        if self.lecam_regularization_weight > 0.0:
+            self.register_buffer("ema_real_logits_mean", torch.zeros((1)))
+            self.register_buffer("ema_fake_logits_mean", torch.zeros((1)))
+
+        self.config = config
+
+    @autocast(enabled=False)
+    def forward(self,
+                inputs: torch.Tensor,
+                reconstructions: torch.Tensor,
+                extra_result_dict: Mapping[Text, torch.Tensor],
+                global_step: int,
+                mode: str = "generator",
+                ) -> Tuple[torch.Tensor, Mapping[Text, torch.Tensor]]:
+        # Both inputs and reconstructions are in range [0, 1].
+        inputs = inputs.float()
+        reconstructions = reconstructions.float()
+
+        if mode == "generator":
+            return self._forward_generator(inputs, reconstructions, extra_result_dict, global_step)
+        elif mode == "discriminator":
+            return self._forward_discriminator(inputs, reconstructions, global_step)
+        else:
+            raise ValueError(f"Unsupported mode {mode}")
+   
+    def should_discriminator_be_trained(self, global_step : int):
+        return global_step >= self.discriminator_iter_start
+
+    def _forward_generator(self,
+                           inputs: torch.Tensor,
+                           reconstructions: torch.Tensor,
+                           extra_result_dict: Mapping[Text, torch.Tensor],
+                           global_step: int
+                           ) -> Tuple[torch.Tensor, Mapping[Text, torch.Tensor]]:
+        """Generator training step."""
+        inputs = inputs.contiguous()
+        reconstructions = reconstructions.contiguous()
+        if self.reconstruction_loss == "l1":
+            reconstruction_loss = F.l1_loss(inputs, reconstructions, reduction="mean")
+        elif self.reconstruction_loss == "l2":
+            reconstruction_loss = F.mse_loss(inputs, reconstructions, reduction="mean")
+        else:
+            raise ValueError(f"Unsuppored reconstruction_loss {self.reconstruction_loss}")
+        reconstruction_loss *= self.reconstruction_weight
+
+        # Compute perceptual loss.
+        perceptual_loss = self.perceptual_loss(inputs, reconstructions).mean()
+
+        # Compute discriminator loss.
+        generator_loss = torch.zeros((), device=inputs.device)
+        discriminator_factor = self.discriminator_factor if self.should_discriminator_be_trained(global_step) else 0
+        d_weight = 1.0
+        if discriminator_factor > 0.0 and self.discriminator_weight > 0.0:
+            # Disable discriminator gradients.
+            for param in self.discriminator.parameters():
+                param.requires_grad = False
+            logits_fake = self.discriminator(reconstructions)
+            generator_loss = -torch.mean(logits_fake)
+
+        d_weight *= self.discriminator_weight
+
+        # Compute quantizer loss.
+        quantizer_loss = extra_result_dict["quantizer_loss"]
+        total_loss = (
+            reconstruction_loss
+            + self.perceptual_weight * perceptual_loss
+            + self.quantizer_weight * quantizer_loss
+            + d_weight * discriminator_factor * generator_loss
+        )
+        loss_dict = dict(
+            total_loss=total_loss.clone().detach(),
+            reconstruction_loss=reconstruction_loss.detach(),
+            perceptual_loss=(self.perceptual_weight * perceptual_loss).detach(),
+            quantizer_loss=(self.quantizer_weight * quantizer_loss).detach(),
+            weighted_gan_loss=(d_weight * discriminator_factor * generator_loss).detach(),
+            discriminator_factor=torch.tensor(discriminator_factor),
+            commitment_loss=extra_result_dict["commitment_loss"].detach(),
+            codebook_loss=extra_result_dict["codebook_loss"].detach(),
+            d_weight=d_weight,
+            gan_loss=generator_loss.detach(),
+        )
+
+        return total_loss, loss_dict
+
+    def _forward_discriminator(self,
+                               inputs: torch.Tensor,
+                               reconstructions: torch.Tensor,
+                               global_step: int,
+                               ) -> Tuple[torch.Tensor, Mapping[Text, torch.Tensor]]:
+        """Discrminator training step."""
+        discriminator_factor = self.discriminator_factor if self.should_discriminator_be_trained(global_step) else 0
+        loss_dict = {}
+        # Turn the gradients on.
+        for param in self.discriminator.parameters():
+            param.requires_grad = True
+
+        real_images = inputs.detach().requires_grad_(True)
+        logits_real = self.discriminator(real_images)
+        logits_fake = self.discriminator(reconstructions.detach())
+
+        discriminator_loss = discriminator_factor * hinge_d_loss(logits_real=logits_real, logits_fake=logits_fake)
+
+        # optional lecam regularization
+        lecam_loss = torch.zeros((), device=inputs.device)
+        if self.lecam_regularization_weight > 0.0:
+            lecam_loss = compute_lecam_loss(
+                torch.mean(logits_real),
+                torch.mean(logits_fake),
+                self.ema_real_logits_mean,
+                self.ema_fake_logits_mean
+            ) * self.lecam_regularization_weight
+
+            self.ema_real_logits_mean = self.ema_real_logits_mean * self.lecam_ema_decay + torch.mean(logits_real).detach()  * (1 - self.lecam_ema_decay)
+            self.ema_fake_logits_mean = self.ema_fake_logits_mean * self.lecam_ema_decay + torch.mean(logits_fake).detach()  * (1 - self.lecam_ema_decay)
+        
+        discriminator_loss += lecam_loss
+
+        loss_dict = dict(
+            discriminator_loss=discriminator_loss.detach(),
+            logits_real=logits_real.detach().mean(),
+            logits_fake=logits_fake.detach().mean(),
+            lecam_loss=lecam_loss.detach(),
+        )
+        return discriminator_loss, loss_dict
+
+
+class MLMLoss(torch.nn.Module):
+    def __init__(self,
+                 config):
+        super().__init__()
+        self.label_smoothing = config.losses.label_smoothing
+        self.loss_weight_unmasked_token = config.losses.loss_weight_unmasked_token
+        self.criterion = torch.nn.CrossEntropyLoss(label_smoothing=self.label_smoothing,
+                                                   reduction="none")
+    
+    def forward(self, inputs: torch.Tensor, targets: torch.Tensor,
+                weights=None) -> Tuple[torch.Tensor, Mapping[Text, torch.Tensor]]:
+        inputs = rearrange(inputs, "b n c -> b c n")
+        loss = self.criterion(inputs, targets)
+        weights = weights.to(loss)
+        loss_weights = (1.0 - weights) * self.loss_weight_unmasked_token + weights # set 0 to self.loss_weight_unasked_token
+        loss = (loss * loss_weights).sum() / (loss_weights.sum() + 1e-8)
+        # we only compute correct tokens on masked tokens
+        correct_tokens = ((torch.argmax(inputs, dim=1) == targets) * weights).sum(dim=1) / (weights.sum(1) + 1e-8)
+        return loss, {"loss": loss, "correct_tokens": correct_tokens.mean()}
+    
+
+class ARLoss(torch.nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        self.target_vocab_size = config.model.vq_model.codebook_size
+        self.criterion = torch.nn.CrossEntropyLoss(reduction="mean")
+    
+    def forward(self, logits: torch.Tensor, labels: torch.Tensor) -> Tuple[torch.Tensor, Mapping[Text, torch.Tensor]]:
+        shift_logits = logits[..., :-1, :].permute(0, 2, 1).contiguous() # NLC->NCL
+        shift_labels = labels.contiguous()
+        shift_logits = shift_logits.view(shift_logits.shape[0], self.target_vocab_size, -1)
+        shift_labels = shift_labels.view(shift_labels.shape[0], -1)
+        shift_labels = shift_labels.to(shift_logits.device)
+        loss = self.criterion(shift_logits, shift_labels)
+        correct_tokens = (torch.argmax(shift_logits, dim=1) == shift_labels).sum(dim=1) / shift_labels.size(1)
+        return loss, {"loss": loss, "correct_tokens": correct_tokens.mean()}