init opensource code

.gitignore

@@ -0,0 +1,21 @@
+__pycache__/
+*.py[cod]
+.env
+.venv/
+.DS_Store
+*.local.*
+/outputs
+/cache
+/data
+/release
+wandb/
+*.ckpt
+*.pt
+*.pth
+*.npz
+*.pb
+*.log
+
+# Model assets distributed on Hugging Face.
+!/checkpoints/*.pt
+!/refs/*.npz

LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2025 Bowen Zheng, The Chinese University of Hong Kong, Shenzhen
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

README.md

@@ -1,3 +1,130 @@
----
-license: mit
----
+# wAR-Tok: Learning Discrete Autoregressive Priors with Wasserstein Gradient Flow
+
+CIFAR-10 implementation of [our paper](https://arxiv.org/abs/2605.06148).
+
+## Features
+
+- Discrete image tokenization with Wasserstein gradient flow regularization.
+- Transformer encoder and decoder with cross-attention, RMSNorm and GEGLU.
+- Six variants: baseline, tail dropout, flow matching, FM + tail dropout, FM + tied tail dropout, and wAR-Tok.
+- Two-stage training: tokenizer training followed by an autoregressive prior.
+- Mixed precision, multi-GPU training and optional W&B logging.
+
+## Quick Start
+
+### Installation
+
+Use Python 3.12 and CUDA. The training environment uses PyTorch 2.8.0 / CUDA 12.8.
+
+```bash
+python -m venv .venv
+source .venv/bin/activate
+pip install -r requirements.txt
+```
+
+### Data Preparation
+
+```bash
+python scripts/prepare_cifar10.py --download
+```
+
+This prepares the CIFAR-10 training split under `data/cifar10`. Download the [FID reference](https://huggingface.co/Zyriix/wARTok/resolve/main/refs/cifar10-32x32_local.npz?download=true) to `refs/cifar10-32x32_local.npz`.
+
+### Training
+
+The default configuration uses 8 GPUs, global batch size 512 and BF16. Each method has a `*_tokenizer.yaml` and a `*_ar.yaml` in [configs/cifar10](configs/cifar10): train the tokenizer for 400K steps, then the AR prior for 200K steps.
+
+**Stage 1: Train the tokenizer**
+
+```bash
+python train.py --config=configs/cifar10/war_tokenizer.yaml
+```
+
+**Stage 2: Train the AR prior**
+
+```bash
+python train.py --config=configs/cifar10/war_ar.yaml \
+  ae_ckpt_path=./outputs/war_tokenizer/000/ckpts/last.ckpt
+```
+
+The AR configuration defaults to the matching tokenizer's `000/ckpts/last.ckpt`. Set `WANDB_MODE=online` to enable experiment logging.
+
+### Evaluation
+
+Use the matching tokenizer configuration for reconstruction and the AR configuration for generation. FM reconstruction uses 25 Euler steps; generation uses a single decoder step. Evaluation uses FP32; generation uses AR EMA, temperature 1, CFG 0 and no top-k/top-p filtering.
+
+**Reconstruction**
+
+```bash
+CUDA_VISIBLE_DEVICES=0 python evaluate.py \
+  --config configs/cifar10/war_tokenizer.yaml --checkpoint checkpoints/war.pt \
+  --mode reconstruct --data data/cifar10 --num-samples 50000 \
+  --output outputs/war-reconstruction --fid-reference refs/cifar10-32x32_local.npz
+```
+
+**Generation**
+
+```bash
+CUDA_VISIBLE_DEVICES=0,1,2,3,4,5,6,7 torchrun --standalone --nproc-per-node=8 evaluate.py \
+  --config configs/cifar10/war_ar.yaml --checkpoint checkpoints/war.pt \
+  --num-samples 50000 --batch-size 64 --seed 42 --no-legacy \
+  --output outputs/war-generation --fid-reference refs/cifar10-32x32_local.npz
+```
+
+Non-legacy mode (`--no-legacy`) uses a different seed on each GPU and generally yields [lower gFID scores](#checkpoints), without changing the relative ranking of the methods or the conclusions supported by the results.
+
+To evaluate all six checkpoints with `--no-legacy`:
+
+```bash
+python scripts/evaluate_all.py --gpus 0,1,2,3,4,5,6,7 \
+  --output outputs/cifar10-nonlegacy
+```
+
+Use `--methods baseline war` for a subset or `--resume` to reuse completed evaluations.
+
+### Exporting Checkpoints
+
+```bash
+python scripts/export_checkpoint.py --source path/to/raw-ar.ckpt \
+  --config configs/cifar10/war_ar.yaml --output checkpoints/war.pt
+```
+
+The export includes tokenizer and raw/EMA AR weights without optimizer states.
+
+## Checkpoints
+
+Name checkpoints after their configuration prefix (e.g. `baseline_ar.yaml` → `checkpoints/baseline.pt`, `war_ar.yaml` → `checkpoints/war.pt`); each includes the tokenizer and raw/EMA AR weights. Hugging Face repository: [wARTok](https://huggingface.co/Zyriix/wARTok).
+
+| Method | rFID | Historical gFID | Non-legacy gFID | Download |
+|---|---:|---:|---:|---|
+| [Baseline](configs/cifar10/baseline_ar.yaml) | 2.3015 | 52.5343 | 49.4191 | [Download](https://huggingface.co/Zyriix/wARTok/resolve/main/checkpoints/baseline.pt?download=true) |
+| [Tail Dropout](configs/cifar10/tail_dropout_ar.yaml) | 3.3567 | 23.3713 | 19.8997 | [Download](https://huggingface.co/Zyriix/wARTok/resolve/main/checkpoints/tail_dropout.pt?download=true) |
+| [FM](configs/cifar10/flow_matching_ar.yaml) | 3.4194 | 38.6309 | 35.8593 | [Download](https://huggingface.co/Zyriix/wARTok/resolve/main/checkpoints/flow_matching.pt?download=true) |
+| [FM + Tail](configs/cifar10/flow_matching_tail_dropout_ar.yaml) | 4.4247 | 22.7613 | 19.0569 | [Download](https://huggingface.co/Zyriix/wARTok/resolve/main/checkpoints/flow_matching_tail_dropout.pt?download=true) |
+| [FM + Tail + Tied](configs/cifar10/flow_matching_tied_tail_dropout_ar.yaml) | 6.4911 | 44.7971 | 41.1334 | [Download](https://huggingface.co/Zyriix/wARTok/resolve/main/checkpoints/flow_matching_tied_tail_dropout.pt?download=true) |
+| [wAR-Tok](configs/cifar10/war_ar.yaml) | 3.0648 | 11.3059 | 7.3749 | [Download](https://huggingface.co/Zyriix/wARTok/resolve/main/checkpoints/war.pt?download=true) |
+
+Download the [FID reference](https://huggingface.co/Zyriix/wARTok/resolve/main/refs/cifar10-32x32_local.npz?download=true) to `refs/cifar10-32x32_local.npz`.
+
+## Project Structure
+
+```text
+wAR/
+├── configs/cifar10/   # Training configurations
+├── checkpoints/      # Downloaded weights
+├── scripts/          # Data preparation and evaluation tools
+├── models.py         # Model architectures
+├── train.py          # Tokenizer and AR training
+├── evaluate.py       # Reconstruction and generation
+├── dataset.py        # Data loading
+├── lpips.py          # Perceptual loss
+└── calc_fid.py        # ADM FID evaluation
+```
+
+## References
+
+This project builds on [EDM](https://github.com/NVlabs/edm), [ADM](https://github.com/openai/guided-diffusion), and [LPIPS](https://github.com/richzhang/PerceptualSimilarity).
+
+## License
+
+Project-authored code is covered by [MIT](LICENSE). Third-party components retain their original licenses.

calc_fid.py

@@ -0,0 +1,34 @@
+"""ADM FID, preserving the input reference archive."""
+import argparse
+import math
+import numpy as np
+import tensorflow.compat.v1 as tf
+from evaluator import Evaluator
+
+
+def main():
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument('--ref_batch', required=True)
+    parser.add_argument('--sample_batch', required=True)
+    parser.add_argument('--batch_size', type=int, default=64)
+    parser.add_argument('--write-reference-stats', help='Optional NEW cache file; never overwrites the reference')
+    args = parser.parse_args()
+    if args.batch_size < 1: parser.error('batch_size must be positive')
+    config = tf.ConfigProto(allow_soft_placement=True, intra_op_parallelism_threads=8, inter_op_parallelism_threads=8)
+    config.gpu_options.allow_growth = True
+    with tf.Session(config=config) as session:
+        evaluator = Evaluator(session, batch_size=args.batch_size)
+        evaluator.warmup()
+        ref_acts = evaluator.read_activations(args.ref_batch)
+        ref, spatial = evaluator.read_statistics(args.ref_batch, ref_acts)
+        sample_acts = evaluator.read_activations(args.sample_batch)
+        sample, _ = evaluator.read_statistics(args.sample_batch, sample_acts)
+        fid = sample.frechet_distance(ref)
+        if not math.isfinite(fid): raise ValueError(f'Non-finite FID: {fid}')
+        if args.write_reference_stats:
+            with open(args.write_reference_stats, 'xb') as f:
+                np.savez_compressed(f, mu=ref.mu, sigma=ref.sigma, mu_s=spatial.mu, sigma_s=spatial.sigma)
+        print(f'FID_RESULT:{fid}')
+
+
+if __name__ == '__main__': main()

checkpoints/baseline.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:ec76675438aa159dfcddf8b4f153ed5f418fc8d2cb01d2f38f082ad32c3e0d85
+size 887697419

checkpoints/flow_matching.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:6d9a7c13f0b199b07fec4f500296a7d5f1f815aec2d81f8635e6ee37a2d30e61
+size 887699762

checkpoints/flow_matching_tail_dropout.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:043dc2f5e32100d2b531007bbdf33c72d57a721d4826672742aa1a85a871a7aa
+size 887699762

checkpoints/flow_matching_tied_tail_dropout.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:490c18b0542a92a7950a370b26b3a9660597b1e1a0977c8fc593ee0bfceab728
+size 887699762

checkpoints/tail_dropout.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:8c790a89b69de1d3bb55acc14f72bc46f723bbf13e5d5753a2fef1b8d201e0b3
+size 887697419

checkpoints/war.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:9c58f2a9520de64e1a29f0e431be15f069b16a88d5a5938da1f36b7b049fbf5c
+size 887697419

configs/cifar10/baseline_ar.yaml

@@ -0,0 +1,198 @@
+wandb_project: wAR
+wandb_name: baseline_ar
+data_dir: ./data/cifar10
+xflip: true
+crop_type: center
+image_size: 32
+num_channels: 3
+num_classes: 11
+patch_size: 4
+use_label: false
+patch_dim: 48
+x_len: 64
+diffusion_decoder: false
+train_ae: false
+train_ar: true
+phases: 200000:DO_PRIOR_AR:1
+batch_size: 512
+eval_batch_size: 512
+lr_enc: 0.00015
+lr_dec: 0.00015
+lr_ar: 0.00015
+lr_gan_loss: 0.00015
+weight_decay_enc: 0.0
+weight_decay_dec: 0.0
+weight_decay_ar: 0.0001
+weight_decay_gan: 0.0001
+lr_enc_min: 3.0e-05
+lr_dec_min: 3.0e-05
+lr_ar_min: 3.0e-05
+lr_gan_loss_min: 3.0e-05
+use_tf32: false
+lr_scheduler: none
+drop: 0.0
+ar_drop: 0.1
+grad_clip: 0.0
+seed: 42
+beta1: 0.5
+beta2: 0.95
+ar_beta1: 0.9
+ar_beta2: 0.99
+label_drop_prob: 0.1
+precision: bf16-mixed
+replace_cur_with_ema: false
+num_workers: 4
+devices: 8
+reverse_t: true
+use_noise_query: false
+use_tail_dropout: false
+l2_v_target: false
+tail_dropout_t2k: 1.0
+tied_timestep: false
+time_stages: 200,400,600,800,1000
+k_per_stage: 192,184,72,48,16
+load_ar_model: false
+load_ar_model_ema: false
+prior_G_loss: false
+use_l1_loss: false
+use_l2_loss: false
+use_lpips_loss: false
+use_gan_loss: false
+use_prior_loss: false
+teacher_forcing_ae: false
+disc_adaptive_weight: false
+lpips_weight: 0.1
+gan_G_weight: 0.1
+prior_enc_weight: 1.0
+l1_weight: 0.1
+l2_weight: 1.0
+gan_start: 5000
+prior_start: 0
+tail_dropout_start: 0
+tail_dropout_p: 0.0
+torch_compile: true
+compile_ae: true
+compile_ar: true
+compile_gan: false
+ar_ema_halflife: 2000
+ae_ema_halflife: 0
+ema_sampling: true
+ema_reconstruction: false
+cache_kv: true
+topK: 0.0
+topP: 0.0
+cfg: 0.0
+cfg_schedule: cosine
+cfg_power: 2.75
+temperature: 1.0
+eval_data_dir: ./data/cifar10
+eval_crop_type: center
+eval_fid_ref_path: ./refs/cifar10-32x32_local.npz
+save_dir: ./outputs/baseline_ar
+log_freq: 25
+visualize_freq: 5000
+visualize_img_num: 8
+ckpt_freq: 25000
+eval_freq: 25000
+save_ckpt: true
+ar_ckpt_path: ''
+ae_ckpt_path: ./outputs/baseline_tokenizer/000/ckpts/last.ckpt
+resume_optimizer: false
+cond_emb_dim: 32
+greedy_coding_prob: 0.0
+z_len: 64
+z_dim: 512
+vq_type: VQ
+codebook_size: 1024
+norm_layer: RMSNorm
+ffn_type: geglu
+context_type: cross_attn
+hidden_dim: 512
+Encoder:
+  input_dim: 48
+  output_dim: 512
+  cond_dim: 32
+  cond_input_dim: 11
+  max_seq_len: 64
+  ctx_max_seq_len: 64
+  dim: 512
+  layer_num: 10
+  mlp_ratio: 4.0
+  heads: 8
+  emb_dropout: 0.0
+  drop: 0.0
+  l2_norm: false
+  context_type: cross_attn
+  rope: false
+  abs_pos: true
+  causal: false
+  norm_layer: RMSNorm
+  ffn_type: geglu
+Quantizer:
+  codebook_size: 1024
+  z_dim: 12
+  dim: 512
+  length: 64
+  cond_dim: 32
+  cond_affine_norm: false
+  vq_beta: 0.25
+  use_norm: true
+  norm_layer: l2
+  ffn_type: geglu
+  temperature: 1.0
+Decoder:
+  input_dim: 48
+  ctx_input_dim: 512
+  output_dim: 48
+  cond_dim: 32
+  cond_input_dim: 11
+  max_seq_len: 64
+  ctx_max_seq_len: 64
+  dim: 512
+  layer_num: 10
+  heads: 8
+  emb_dropout: 0.0
+  drop: 0.0
+  l2_norm: false
+  context_type: cross_attn
+  rope: false
+  abs_pos: true
+  causal: false
+  norm_layer: RMSNorm
+  ffn_type: geglu
+  mlp_ratio: 4.0
+  zero_out: true
+  query_type: learnable
+ARModel:
+  input_dim: 512
+  output_dim: 1024
+  max_seq_len: 64
+  cond_dim: 32
+  cond_input_dim: 11
+  dim: 512
+  layer_num: 13
+  heads: 8
+  emb_dropout: 0.0
+  drop: 0.1
+  l2_norm: false
+  context_type: none
+  rope: false
+  abs_pos: true
+  causal: true
+  out_act: false
+  norm_layer: RMSNorm
+  ffn_type: geglu
+  mlp_ratio: 4.0
+  zero_out: false
+  temperature: 1.0
+GANLoss:
+  disc_loss: vanilla
+  disc_dim: 64
+  disc_type: patchgan
+  image_size: 32
+  disc_num_layers: 3
+  disc_in_channels: 3
+  gen_adv_loss: vanilla
+  lecam_loss_weight: 0.001
+  norm_type: bn
+  aug_prob: 1.0

configs/cifar10/baseline_tokenizer.yaml

@@ -0,0 +1,198 @@
+wandb_project: wAR
+wandb_name: baseline_tokenizer
+data_dir: ./data/cifar10
+xflip: true
+crop_type: center
+image_size: 32
+num_channels: 3
+num_classes: 11
+patch_size: 4
+use_label: false
+patch_dim: 48
+x_len: 64
+diffusion_decoder: false
+train_ae: true
+train_ar: false
+phases: 400000:DO_L1-DO_L2-DO_LPIPS:1
+batch_size: 512
+eval_batch_size: 512
+lr_enc: 0.00015
+lr_dec: 0.00015
+lr_ar: 0.00015
+lr_gan_loss: 0.00015
+weight_decay_enc: 0.0
+weight_decay_dec: 0.0
+weight_decay_ar: 0.0
+weight_decay_gan: 0.0001
+lr_enc_min: 3.0e-05
+lr_dec_min: 3.0e-05
+lr_ar_min: 3.0e-05
+lr_gan_loss_min: 3.0e-05
+use_tf32: false
+lr_scheduler: none
+drop: 0.0
+ar_drop: 0.0
+grad_clip: 0.0
+seed: 42
+beta1: 0.5
+beta2: 0.95
+ar_beta1: 0.5
+ar_beta2: 0.95
+label_drop_prob: 0.1
+precision: bf16-mixed
+replace_cur_with_ema: false
+num_workers: 4
+devices: 8
+reverse_t: false
+use_noise_query: false
+use_tail_dropout: false
+l2_v_target: false
+tail_dropout_t2k: 1.0
+tied_timestep: false
+time_stages: 200,400,600,800,1000
+k_per_stage: 192,184,72,48,16
+load_ar_model: false
+load_ar_model_ema: false
+prior_G_loss: false
+use_l1_loss: true
+use_l2_loss: true
+use_lpips_loss: true
+use_gan_loss: false
+use_prior_loss: false
+teacher_forcing_ae: false
+disc_adaptive_weight: false
+lpips_weight: 0.1
+gan_G_weight: 0.1
+prior_enc_weight: 1.0
+l1_weight: 0.1
+l2_weight: 1.0
+gan_start: 5000
+prior_start: 0
+tail_dropout_start: 0
+tail_dropout_p: 0
+torch_compile: true
+compile_ae: true
+compile_ar: true
+compile_gan: false
+ar_ema_halflife: 2500
+ae_ema_halflife: 1000
+ema_sampling: false
+ema_reconstruction: false
+cache_kv: true
+topK: 0.0
+topP: 0.0
+cfg: 0.0
+cfg_schedule: cosine
+cfg_power: 2.75
+temperature: 1.0
+eval_data_dir: ./data/cifar10
+eval_crop_type: center
+eval_fid_ref_path: ./refs/cifar10-32x32_local.npz
+save_dir: ./outputs/baseline_tokenizer
+log_freq: 25
+visualize_freq: 5000
+visualize_img_num: 8
+ckpt_freq: 25000
+eval_freq: 25000
+save_ckpt: true
+ar_ckpt_path: ''
+ae_ckpt_path: ''
+resume_optimizer: false
+cond_emb_dim: 32
+greedy_coding_prob: 0.0
+z_len: 64
+z_dim: 512
+vq_type: VQ
+codebook_size: 1024
+norm_layer: RMSNorm
+ffn_type: geglu
+context_type: cross_attn
+hidden_dim: 512
+Encoder:
+  input_dim: 48
+  output_dim: 512
+  cond_dim: 32
+  cond_input_dim: 11
+  max_seq_len: 64
+  ctx_max_seq_len: 64
+  dim: 512
+  layer_num: 10
+  mlp_ratio: 4.0
+  heads: 8
+  emb_dropout: 0.0
+  drop: 0.0
+  l2_norm: false
+  context_type: cross_attn
+  rope: false
+  abs_pos: true
+  causal: false
+  norm_layer: RMSNorm
+  ffn_type: geglu
+Quantizer:
+  codebook_size: 1024
+  z_dim: 12
+  dim: 512
+  length: 64
+  cond_dim: 32
+  cond_affine_norm: false
+  vq_beta: 0.25
+  use_norm: true
+  norm_layer: l2
+  ffn_type: geglu
+  temperature: 1.0
+Decoder:
+  input_dim: 48
+  ctx_input_dim: 512
+  output_dim: 48
+  cond_dim: 32
+  cond_input_dim: 11
+  max_seq_len: 64
+  ctx_max_seq_len: 64
+  dim: 512
+  layer_num: 10
+  heads: 8
+  emb_dropout: 0.0
+  drop: 0.0
+  l2_norm: false
+  context_type: cross_attn
+  rope: false
+  abs_pos: true
+  causal: false
+  norm_layer: RMSNorm
+  ffn_type: geglu
+  mlp_ratio: 4.0
+  zero_out: true
+  query_type: learnable
+ARModel:
+  input_dim: 512
+  output_dim: 1024
+  max_seq_len: 64
+  cond_dim: 32
+  cond_input_dim: 11
+  dim: 512
+  layer_num: 6
+  heads: 8
+  emb_dropout: 0.0
+  drop: 0.0
+  l2_norm: false
+  context_type: none
+  rope: false
+  abs_pos: true
+  causal: true
+  out_act: false
+  norm_layer: RMSNorm
+  ffn_type: geglu
+  mlp_ratio: 4
+  zero_out: false
+  temperature: 1.0
+GANLoss:
+  disc_loss: vanilla
+  disc_dim: 64
+  disc_type: patchgan
+  image_size: 32
+  disc_num_layers: 3
+  disc_in_channels: 3
+  gen_adv_loss: vanilla
+  lecam_loss_weight: 0.001
+  norm_type: bn
+  aug_prob: 1.0

configs/cifar10/flow_matching_ar.yaml

@@ -0,0 +1,198 @@
+wandb_project: wAR
+wandb_name: flow_matching_ar
+data_dir: ./data/cifar10
+xflip: true
+crop_type: center
+image_size: 32
+num_channels: 3
+num_classes: 11
+patch_size: 4
+use_label: false
+patch_dim: 48
+x_len: 64
+diffusion_decoder: false
+train_ae: false
+train_ar: true
+phases: 200000:DO_PRIOR_AR:1
+batch_size: 512
+eval_batch_size: 512
+lr_enc: 0.00015
+lr_dec: 0.00015
+lr_ar: 0.00015
+lr_gan_loss: 0.00015
+weight_decay_enc: 0.0
+weight_decay_dec: 0.0
+weight_decay_ar: 0.0001
+weight_decay_gan: 0.0001
+lr_enc_min: 3.0e-05
+lr_dec_min: 3.0e-05
+lr_ar_min: 3.0e-05
+lr_gan_loss_min: 3.0e-05
+use_tf32: false
+lr_scheduler: none
+drop: 0.0
+ar_drop: 0.1
+grad_clip: 0.0
+seed: 42
+beta1: 0.5
+beta2: 0.95
+ar_beta1: 0.9
+ar_beta2: 0.99
+label_drop_prob: 0.1
+precision: bf16-mixed
+replace_cur_with_ema: false
+num_workers: 4
+devices: 8
+reverse_t: true
+use_noise_query: true
+use_tail_dropout: false
+l2_v_target: true
+tail_dropout_t2k: 1.0
+tied_timestep: false
+time_stages: 200,400,600,800,1000
+k_per_stage: 192,184,72,48,16
+load_ar_model: false
+load_ar_model_ema: false
+prior_G_loss: false
+use_l1_loss: false
+use_l2_loss: false
+use_lpips_loss: false
+use_gan_loss: false
+use_prior_loss: false
+teacher_forcing_ae: false
+disc_adaptive_weight: false
+lpips_weight: 0.1
+gan_G_weight: 0.1
+prior_enc_weight: 1.0
+l1_weight: 0.1
+l2_weight: 1.0
+gan_start: 5000
+prior_start: 0
+tail_dropout_start: 0
+tail_dropout_p: 0.0
+torch_compile: true
+compile_ae: true
+compile_ar: true
+compile_gan: false
+ar_ema_halflife: 2000
+ae_ema_halflife: 0
+ema_sampling: true
+ema_reconstruction: false
+cache_kv: true
+topK: 0.0
+topP: 0.0
+cfg: 0.0
+cfg_schedule: cosine
+cfg_power: 2.75
+temperature: 1.0
+eval_data_dir: ./data/cifar10
+eval_crop_type: center
+eval_fid_ref_path: ./refs/cifar10-32x32_local.npz
+save_dir: ./outputs/flow_matching_ar
+log_freq: 25
+visualize_freq: 5000
+visualize_img_num: 8
+ckpt_freq: 25000
+eval_freq: 25000
+save_ckpt: true
+ar_ckpt_path: ''
+ae_ckpt_path: ./outputs/flow_matching_tokenizer/000/ckpts/last.ckpt
+resume_optimizer: false
+cond_emb_dim: 32
+greedy_coding_prob: 0.0
+z_len: 64
+z_dim: 512
+vq_type: VQ
+codebook_size: 1024
+norm_layer: RMSNorm
+ffn_type: geglu
+context_type: cross_attn
+hidden_dim: 512
+Encoder:
+  input_dim: 48
+  output_dim: 512
+  cond_dim: 32
+  cond_input_dim: 11
+  max_seq_len: 64
+  ctx_max_seq_len: 64
+  dim: 512
+  layer_num: 10
+  mlp_ratio: 4.0
+  heads: 8
+  emb_dropout: 0.0
+  drop: 0.0
+  l2_norm: false
+  context_type: cross_attn
+  rope: false
+  abs_pos: true
+  causal: false
+  norm_layer: RMSNorm
+  ffn_type: geglu
+Quantizer:
+  codebook_size: 1024
+  z_dim: 12
+  dim: 512
+  length: 64
+  cond_dim: 32
+  cond_affine_norm: false
+  vq_beta: 0.25
+  use_norm: true
+  norm_layer: l2
+  ffn_type: geglu
+  temperature: 1.0
+Decoder:
+  input_dim: 48
+  ctx_input_dim: 512
+  output_dim: 48
+  cond_dim: 32
+  cond_input_dim: 11
+  max_seq_len: 64
+  ctx_max_seq_len: 64
+  dim: 512
+  layer_num: 10
+  heads: 8
+  emb_dropout: 0.0
+  drop: 0.0
+  l2_norm: false
+  context_type: cross_attn
+  rope: false
+  abs_pos: true
+  causal: false
+  norm_layer: RMSNorm
+  ffn_type: geglu
+  mlp_ratio: 4.0
+  zero_out: true
+  query_type: noise
+ARModel:
+  input_dim: 512
+  output_dim: 1024
+  max_seq_len: 64
+  cond_dim: 32
+  cond_input_dim: 11
+  dim: 512
+  layer_num: 13
+  heads: 8
+  emb_dropout: 0.0
+  drop: 0.1
+  l2_norm: false
+  context_type: none
+  rope: false
+  abs_pos: true
+  causal: true
+  out_act: false
+  norm_layer: RMSNorm
+  ffn_type: geglu
+  mlp_ratio: 4.0
+  zero_out: false
+  temperature: 1.0
+GANLoss:
+  disc_loss: vanilla
+  disc_dim: 64
+  disc_type: patchgan
+  image_size: 32
+  disc_num_layers: 3
+  disc_in_channels: 3
+  gen_adv_loss: vanilla
+  lecam_loss_weight: 0.001
+  norm_type: bn
+  aug_prob: 1.0

configs/cifar10/flow_matching_tail_dropout_ar.yaml

@@ -0,0 +1,198 @@
+wandb_project: wAR
+wandb_name: flow_matching_tail_dropout_ar
+data_dir: ./data/cifar10
+xflip: true
+crop_type: center
+image_size: 32
+num_channels: 3
+num_classes: 11
+patch_size: 4
+use_label: false
+patch_dim: 48
+x_len: 64
+diffusion_decoder: false
+train_ae: false
+train_ar: true
+phases: 200000:DO_PRIOR_AR:1
+batch_size: 512
+eval_batch_size: 512
+lr_enc: 0.00015
+lr_dec: 0.00015
+lr_ar: 0.00015
+lr_gan_loss: 0.00015
+weight_decay_enc: 0.0
+weight_decay_dec: 0.0
+weight_decay_ar: 0.0001
+weight_decay_gan: 0.0001
+lr_enc_min: 3.0e-05
+lr_dec_min: 3.0e-05
+lr_ar_min: 3.0e-05
+lr_gan_loss_min: 3.0e-05
+use_tf32: false
+lr_scheduler: none
+drop: 0.0
+ar_drop: 0.1
+grad_clip: 0.0
+seed: 42
+beta1: 0.5
+beta2: 0.95
+ar_beta1: 0.9
+ar_beta2: 0.99
+label_drop_prob: 0.1
+precision: bf16-mixed
+replace_cur_with_ema: false
+num_workers: 4
+devices: 8
+reverse_t: true
+use_noise_query: true
+use_tail_dropout: true
+l2_v_target: true
+tail_dropout_t2k: 1.0
+tied_timestep: false
+time_stages: 200,400,600,800,1000
+k_per_stage: 192,184,72,48,16
+load_ar_model: false
+load_ar_model_ema: false
+prior_G_loss: false
+use_l1_loss: false
+use_l2_loss: false
+use_lpips_loss: false
+use_gan_loss: false
+use_prior_loss: false
+teacher_forcing_ae: false
+disc_adaptive_weight: false
+lpips_weight: 0.1
+gan_G_weight: 0.1
+prior_enc_weight: 1.0
+l1_weight: 0.1
+l2_weight: 1.0
+gan_start: 5000
+prior_start: 0
+tail_dropout_start: 0
+tail_dropout_p: 0.5
+torch_compile: true
+compile_ae: true
+compile_ar: true
+compile_gan: false
+ar_ema_halflife: 2000
+ae_ema_halflife: 0
+ema_sampling: true
+ema_reconstruction: false
+cache_kv: true
+topK: 0.0
+topP: 0.0
+cfg: 0.0
+cfg_schedule: cosine
+cfg_power: 2.75
+temperature: 1.0
+eval_data_dir: ./data/cifar10
+eval_crop_type: center
+eval_fid_ref_path: ./refs/cifar10-32x32_local.npz
+save_dir: ./outputs/flow_matching_tail_dropout_ar
+log_freq: 25
+visualize_freq: 5000
+visualize_img_num: 8
+ckpt_freq: 25000
+eval_freq: 25000
+save_ckpt: true
+ar_ckpt_path: ''
+ae_ckpt_path: ./outputs/flow_matching_tail_dropout_tokenizer/000/ckpts/last.ckpt
+resume_optimizer: false
+cond_emb_dim: 32
+greedy_coding_prob: 0.0
+z_len: 64
+z_dim: 512
+vq_type: VQ
+codebook_size: 1024
+norm_layer: RMSNorm
+ffn_type: geglu
+context_type: cross_attn
+hidden_dim: 512
+Encoder:
+  input_dim: 48
+  output_dim: 512
+  cond_dim: 32
+  cond_input_dim: 11
+  max_seq_len: 64
+  ctx_max_seq_len: 64
+  dim: 512
+  layer_num: 10
+  mlp_ratio: 4.0
+  heads: 8
+  emb_dropout: 0.0
+  drop: 0.0
+  l2_norm: false
+  context_type: cross_attn
+  rope: false
+  abs_pos: true
+  causal: false
+  norm_layer: RMSNorm
+  ffn_type: geglu
+Quantizer:
+  codebook_size: 1024
+  z_dim: 12
+  dim: 512
+  length: 64
+  cond_dim: 32
+  cond_affine_norm: false
+  vq_beta: 0.25
+  use_norm: true
+  norm_layer: l2
+  ffn_type: geglu
+  temperature: 1.0
+Decoder:
+  input_dim: 48
+  ctx_input_dim: 512
+  output_dim: 48
+  cond_dim: 32
+  cond_input_dim: 11
+  max_seq_len: 64
+  ctx_max_seq_len: 64
+  dim: 512
+  layer_num: 10
+  heads: 8
+  emb_dropout: 0.0
+  drop: 0.0
+  l2_norm: false
+  context_type: cross_attn
+  rope: false
+  abs_pos: true
+  causal: false
+  norm_layer: RMSNorm
+  ffn_type: geglu
+  mlp_ratio: 4.0
+  zero_out: true
+  query_type: noise
+ARModel:
+  input_dim: 512
+  output_dim: 1024
+  max_seq_len: 64
+  cond_dim: 32
+  cond_input_dim: 11
+  dim: 512
+  layer_num: 13
+  heads: 8
+  emb_dropout: 0.0
+  drop: 0.1
+  l2_norm: false
+  context_type: none
+  rope: false
+  abs_pos: true
+  causal: true
+  out_act: false
+  norm_layer: RMSNorm
+  ffn_type: geglu
+  mlp_ratio: 4.0
+  zero_out: false
+  temperature: 1.0
+GANLoss:
+  disc_loss: vanilla
+  disc_dim: 64
+  disc_type: patchgan
+  image_size: 32
+  disc_num_layers: 3
+  disc_in_channels: 3
+  gen_adv_loss: vanilla
+  lecam_loss_weight: 0.001
+  norm_type: bn
+  aug_prob: 1.0

configs/cifar10/flow_matching_tail_dropout_tokenizer.yaml

@@ -0,0 +1,198 @@
+wandb_project: wAR
+wandb_name: flow_matching_tail_dropout_tokenizer
+data_dir: ./data/cifar10
+xflip: true
+crop_type: center
+image_size: 32
+num_channels: 3
+num_classes: 11
+patch_size: 4
+use_label: false
+patch_dim: 48
+x_len: 64
+diffusion_decoder: true
+train_ae: true
+train_ar: false
+phases: 400000:DO_L2-DO_LPIPS:1
+batch_size: 512
+eval_batch_size: 512
+lr_enc: 0.00015
+lr_dec: 0.00015
+lr_ar: 0.00015
+lr_gan_loss: 0.00015
+weight_decay_enc: 0.0
+weight_decay_dec: 0.0
+weight_decay_ar: 0.0
+weight_decay_gan: 0.0001
+lr_enc_min: 3.0e-05
+lr_dec_min: 3.0e-05
+lr_ar_min: 3.0e-05
+lr_gan_loss_min: 3.0e-05
+use_tf32: false
+lr_scheduler: none
+drop: 0.0
+ar_drop: 0.0
+grad_clip: 0.0
+seed: 42
+beta1: 0.5
+beta2: 0.95
+ar_beta1: 0.5
+ar_beta2: 0.95
+label_drop_prob: 0.1
+precision: bf16-mixed
+replace_cur_with_ema: false
+num_workers: 4
+devices: 8
+reverse_t: false
+use_noise_query: true
+use_tail_dropout: true
+l2_v_target: true
+tail_dropout_t2k: 1.0
+tied_timestep: false
+time_stages: 200,400,600,800,1000
+k_per_stage: 192,184,72,48,16
+load_ar_model: false
+load_ar_model_ema: false
+prior_G_loss: false
+use_l1_loss: false
+use_l2_loss: true
+use_lpips_loss: true
+use_gan_loss: false
+use_prior_loss: false
+teacher_forcing_ae: false
+disc_adaptive_weight: false
+lpips_weight: 0.1
+gan_G_weight: 0.1
+prior_enc_weight: 1.0
+l1_weight: 0.1
+l2_weight: 1.0
+gan_start: 5000
+prior_start: 0
+tail_dropout_start: 0
+tail_dropout_p: 0.5
+torch_compile: true
+compile_ae: true
+compile_ar: true
+compile_gan: false
+ar_ema_halflife: 2500
+ae_ema_halflife: 1000
+ema_sampling: false
+ema_reconstruction: false
+cache_kv: true
+topK: 0.0
+topP: 0.0
+cfg: 0.0
+cfg_schedule: cosine
+cfg_power: 2.75
+temperature: 1.0
+eval_data_dir: ./data/cifar10
+eval_crop_type: center
+eval_fid_ref_path: ./refs/cifar10-32x32_local.npz
+save_dir: ./outputs/flow_matching_tail_dropout_tokenizer
+log_freq: 25
+visualize_freq: 5000
+visualize_img_num: 8
+ckpt_freq: 25000
+eval_freq: 25000
+save_ckpt: true
+ar_ckpt_path: ''
+ae_ckpt_path: ''
+resume_optimizer: false
+cond_emb_dim: 32
+greedy_coding_prob: 0.0
+z_len: 64
+z_dim: 512
+vq_type: VQ
+codebook_size: 1024
+norm_layer: RMSNorm
+ffn_type: geglu
+context_type: cross_attn
+hidden_dim: 512
+Encoder:
+  input_dim: 48
+  output_dim: 512
+  cond_dim: 32
+  cond_input_dim: 11
+  max_seq_len: 64
+  ctx_max_seq_len: 64
+  dim: 512
+  layer_num: 10
+  mlp_ratio: 4.0
+  heads: 8
+  emb_dropout: 0.0
+  drop: 0.0
+  l2_norm: false
+  context_type: cross_attn
+  rope: false
+  abs_pos: true
+  causal: false
+  norm_layer: RMSNorm
+  ffn_type: geglu
+Quantizer:
+  codebook_size: 1024
+  z_dim: 12
+  dim: 512
+  length: 64
+  cond_dim: 32
+  cond_affine_norm: false
+  vq_beta: 0.25
+  use_norm: true
+  norm_layer: l2
+  ffn_type: geglu
+  temperature: 1.0
+Decoder:
+  input_dim: 48
+  ctx_input_dim: 512
+  output_dim: 48
+  cond_dim: 32
+  cond_input_dim: 11
+  max_seq_len: 64
+  ctx_max_seq_len: 64
+  dim: 512
+  layer_num: 10
+  heads: 8
+  emb_dropout: 0.0
+  drop: 0.0
+  l2_norm: false
+  context_type: cross_attn
+  rope: false
+  abs_pos: true
+  causal: false
+  norm_layer: RMSNorm
+  ffn_type: geglu
+  mlp_ratio: 4.0
+  zero_out: true
+  query_type: noise
+ARModel:
+  input_dim: 512
+  output_dim: 1024
+  max_seq_len: 64
+  cond_dim: 32
+  cond_input_dim: 11
+  dim: 512
+  layer_num: 6
+  heads: 8
+  emb_dropout: 0.0
+  drop: 0.0
+  l2_norm: false
+  context_type: none
+  rope: false
+  abs_pos: true
+  causal: true
+  out_act: false
+  norm_layer: RMSNorm
+  ffn_type: geglu
+  mlp_ratio: 4
+  zero_out: false
+  temperature: 1.0
+GANLoss:
+  disc_loss: vanilla
+  disc_dim: 64
+  disc_type: patchgan
+  image_size: 32
+  disc_num_layers: 3
+  disc_in_channels: 3
+  gen_adv_loss: vanilla
+  lecam_loss_weight: 0.001
+  norm_type: bn
+  aug_prob: 1.0

configs/cifar10/flow_matching_tied_tail_dropout_ar.yaml

@@ -0,0 +1,198 @@
+wandb_project: wAR
+wandb_name: flow_matching_tied_tail_dropout_ar
+data_dir: ./data/cifar10
+xflip: true
+crop_type: center
+image_size: 32
+num_channels: 3
+num_classes: 11
+patch_size: 4
+use_label: false
+patch_dim: 48
+x_len: 64
+diffusion_decoder: false
+train_ae: false
+train_ar: true
+phases: 200000:DO_PRIOR_AR:1
+batch_size: 512
+eval_batch_size: 512
+lr_enc: 0.00015
+lr_dec: 0.00015
+lr_ar: 0.00015
+lr_gan_loss: 0.00015
+weight_decay_enc: 0.0
+weight_decay_dec: 0.0
+weight_decay_ar: 0.0001
+weight_decay_gan: 0.0001
+lr_enc_min: 3.0e-05
+lr_dec_min: 3.0e-05
+lr_ar_min: 3.0e-05
+lr_gan_loss_min: 3.0e-05
+use_tf32: false
+lr_scheduler: none
+drop: 0.0
+ar_drop: 0.1
+grad_clip: 0.0
+seed: 42
+beta1: 0.5
+beta2: 0.95
+ar_beta1: 0.9
+ar_beta2: 0.99
+label_drop_prob: 0.1
+precision: bf16-mixed
+replace_cur_with_ema: false
+num_workers: 4
+devices: 8
+reverse_t: true
+use_noise_query: true
+use_tail_dropout: true
+l2_v_target: true
+tail_dropout_t2k: 1.0
+tied_timestep: true
+time_stages: 200,400,600,800,1000
+k_per_stage: 192,184,72,48,16
+load_ar_model: false
+load_ar_model_ema: false
+prior_G_loss: false
+use_l1_loss: false
+use_l2_loss: false
+use_lpips_loss: false
+use_gan_loss: false
+use_prior_loss: false
+teacher_forcing_ae: false
+disc_adaptive_weight: false
+lpips_weight: 0.1
+gan_G_weight: 0.1
+prior_enc_weight: 1.0
+l1_weight: 0.1
+l2_weight: 1.0
+gan_start: 5000
+prior_start: 0
+tail_dropout_start: 0
+tail_dropout_p: 1.0
+torch_compile: true
+compile_ae: true
+compile_ar: true
+compile_gan: false
+ar_ema_halflife: 2000
+ae_ema_halflife: 0
+ema_sampling: true
+ema_reconstruction: false
+cache_kv: true
+topK: 0.0
+topP: 0.0
+cfg: 0.0
+cfg_schedule: cosine
+cfg_power: 2.75
+temperature: 1.0
+eval_data_dir: ./data/cifar10
+eval_crop_type: center
+eval_fid_ref_path: ./refs/cifar10-32x32_local.npz
+save_dir: ./outputs/flow_matching_tied_tail_dropout_ar
+log_freq: 25
+visualize_freq: 5000
+visualize_img_num: 8
+ckpt_freq: 25000
+eval_freq: 25000
+save_ckpt: true
+ar_ckpt_path: ''
+ae_ckpt_path: ./outputs/flow_matching_tied_tail_dropout_tokenizer/000/ckpts/last.ckpt
+resume_optimizer: false
+cond_emb_dim: 32
+greedy_coding_prob: 0.0
+z_len: 64
+z_dim: 512
+vq_type: VQ
+codebook_size: 1024
+norm_layer: RMSNorm
+ffn_type: geglu
+context_type: cross_attn
+hidden_dim: 512
+Encoder:
+  input_dim: 48
+  output_dim: 512
+  cond_dim: 32
+  cond_input_dim: 11
+  max_seq_len: 64
+  ctx_max_seq_len: 64
+  dim: 512
+  layer_num: 10
+  mlp_ratio: 4.0
+  heads: 8
+  emb_dropout: 0.0
+  drop: 0.0
+  l2_norm: false
+  context_type: cross_attn
+  rope: false
+  abs_pos: true
+  causal: false
+  norm_layer: RMSNorm
+  ffn_type: geglu
+Quantizer:
+  codebook_size: 1024
+  z_dim: 12
+  dim: 512
+  length: 64
+  cond_dim: 32
+  cond_affine_norm: false
+  vq_beta: 0.25
+  use_norm: true
+  norm_layer: l2
+  ffn_type: geglu
+  temperature: 1.0
+Decoder:
+  input_dim: 48
+  ctx_input_dim: 512
+  output_dim: 48
+  cond_dim: 32
+  cond_input_dim: 11
+  max_seq_len: 64
+  ctx_max_seq_len: 64
+  dim: 512
+  layer_num: 10
+  heads: 8
+  emb_dropout: 0.0
+  drop: 0.0
+  l2_norm: false
+  context_type: cross_attn
+  rope: false
+  abs_pos: true
+  causal: false
+  norm_layer: RMSNorm
+  ffn_type: geglu
+  mlp_ratio: 4.0
+  zero_out: true
+  query_type: noise
+ARModel:
+  input_dim: 512
+  output_dim: 1024
+  max_seq_len: 64
+  cond_dim: 32
+  cond_input_dim: 11
+  dim: 512
+  layer_num: 13
+  heads: 8
+  emb_dropout: 0.0
+  drop: 0.1
+  l2_norm: false
+  context_type: none
+  rope: false
+  abs_pos: true
+  causal: true
+  out_act: false
+  norm_layer: RMSNorm
+  ffn_type: geglu
+  mlp_ratio: 4.0
+  zero_out: false
+  temperature: 1.0
+GANLoss:
+  disc_loss: vanilla
+  disc_dim: 64
+  disc_type: patchgan
+  image_size: 32
+  disc_num_layers: 3
+  disc_in_channels: 3
+  gen_adv_loss: vanilla
+  lecam_loss_weight: 0.001
+  norm_type: bn
+  aug_prob: 1.0

configs/cifar10/flow_matching_tied_tail_dropout_tokenizer.yaml

@@ -0,0 +1,198 @@
+wandb_project: wAR
+wandb_name: flow_matching_tied_tail_dropout_tokenizer
+data_dir: ./data/cifar10
+xflip: true
+crop_type: center
+image_size: 32
+num_channels: 3
+num_classes: 11
+patch_size: 4
+use_label: false
+patch_dim: 48
+x_len: 64
+diffusion_decoder: true
+train_ae: true
+train_ar: false
+phases: 400000:DO_L2-DO_LPIPS:1
+batch_size: 512
+eval_batch_size: 512
+lr_enc: 0.00015
+lr_dec: 0.00015
+lr_ar: 0.00015
+lr_gan_loss: 0.00015
+weight_decay_enc: 0.0
+weight_decay_dec: 0.0
+weight_decay_ar: 0.0
+weight_decay_gan: 0.0001
+lr_enc_min: 3.0e-05
+lr_dec_min: 3.0e-05
+lr_ar_min: 3.0e-05
+lr_gan_loss_min: 3.0e-05
+use_tf32: false
+lr_scheduler: none
+drop: 0.0
+ar_drop: 0.0
+grad_clip: 0.0
+seed: 42
+beta1: 0.5
+beta2: 0.95
+ar_beta1: 0.5
+ar_beta2: 0.95
+label_drop_prob: 0.1
+precision: bf16-mixed
+replace_cur_with_ema: false
+num_workers: 4
+devices: 8
+reverse_t: false
+use_noise_query: true
+use_tail_dropout: true
+l2_v_target: true
+tail_dropout_t2k: 1.0
+tied_timestep: true
+time_stages: 200,400,600,800,1000
+k_per_stage: 192,184,72,48,16
+load_ar_model: false
+load_ar_model_ema: false
+prior_G_loss: false
+use_l1_loss: false
+use_l2_loss: true
+use_lpips_loss: true
+use_gan_loss: false
+use_prior_loss: false
+teacher_forcing_ae: false
+disc_adaptive_weight: false
+lpips_weight: 0.1
+gan_G_weight: 0.1
+prior_enc_weight: 1.0
+l1_weight: 0.1
+l2_weight: 1.0
+gan_start: 5000
+prior_start: 0
+tail_dropout_start: 0
+tail_dropout_p: 1.0
+torch_compile: true
+compile_ae: true
+compile_ar: true
+compile_gan: false
+ar_ema_halflife: 2500
+ae_ema_halflife: 1000
+ema_sampling: false
+ema_reconstruction: false
+cache_kv: true
+topK: 0.0
+topP: 0.0
+cfg: 0.0
+cfg_schedule: cosine
+cfg_power: 2.75
+temperature: 1.0
+eval_data_dir: ./data/cifar10
+eval_crop_type: center
+eval_fid_ref_path: ./refs/cifar10-32x32_local.npz
+save_dir: ./outputs/flow_matching_tied_tail_dropout_tokenizer
+log_freq: 25
+visualize_freq: 5000
+visualize_img_num: 8
+ckpt_freq: 25000
+eval_freq: 25000
+save_ckpt: true
+ar_ckpt_path: ''
+ae_ckpt_path: ''
+resume_optimizer: false
+cond_emb_dim: 32
+greedy_coding_prob: 0.0
+z_len: 64
+z_dim: 512
+vq_type: VQ
+codebook_size: 1024
+norm_layer: RMSNorm
+ffn_type: geglu
+context_type: cross_attn
+hidden_dim: 512
+Encoder:
+  input_dim: 48
+  output_dim: 512
+  cond_dim: 32
+  cond_input_dim: 11
+  max_seq_len: 64
+  ctx_max_seq_len: 64
+  dim: 512
+  layer_num: 10
+  mlp_ratio: 4.0
+  heads: 8
+  emb_dropout: 0.0
+  drop: 0.0
+  l2_norm: false
+  context_type: cross_attn
+  rope: false
+  abs_pos: true
+  causal: false
+  norm_layer: RMSNorm
+  ffn_type: geglu
+Quantizer:
+  codebook_size: 1024
+  z_dim: 12
+  dim: 512
+  length: 64
+  cond_dim: 32
+  cond_affine_norm: false
+  vq_beta: 0.25
+  use_norm: true
+  norm_layer: l2
+  ffn_type: geglu
+  temperature: 1.0
+Decoder:
+  input_dim: 48
+  ctx_input_dim: 512
+  output_dim: 48
+  cond_dim: 32
+  cond_input_dim: 11
+  max_seq_len: 64
+  ctx_max_seq_len: 64
+  dim: 512
+  layer_num: 10
+  heads: 8
+  emb_dropout: 0.0
+  drop: 0.0
+  l2_norm: false
+  context_type: cross_attn
+  rope: false
+  abs_pos: true
+  causal: false
+  norm_layer: RMSNorm
+  ffn_type: geglu
+  mlp_ratio: 4.0
+  zero_out: true
+  query_type: noise
+ARModel:
+  input_dim: 512
+  output_dim: 1024
+  max_seq_len: 64
+  cond_dim: 32
+  cond_input_dim: 11
+  dim: 512
+  layer_num: 6
+  heads: 8
+  emb_dropout: 0.0
+  drop: 0.0
+  l2_norm: false
+  context_type: none
+  rope: false
+  abs_pos: true
+  causal: true
+  out_act: false
+  norm_layer: RMSNorm
+  ffn_type: geglu
+  mlp_ratio: 4
+  zero_out: false
+  temperature: 1.0
+GANLoss:
+  disc_loss: vanilla
+  disc_dim: 64
+  disc_type: patchgan
+  image_size: 32
+  disc_num_layers: 3
+  disc_in_channels: 3
+  gen_adv_loss: vanilla
+  lecam_loss_weight: 0.001
+  norm_type: bn
+  aug_prob: 1.0

configs/cifar10/flow_matching_tokenizer.yaml

@@ -0,0 +1,198 @@
+wandb_project: wAR
+wandb_name: flow_matching_tokenizer
+data_dir: ./data/cifar10
+xflip: true
+crop_type: center
+image_size: 32
+num_channels: 3
+num_classes: 11
+patch_size: 4
+use_label: false
+patch_dim: 48
+x_len: 64
+diffusion_decoder: true
+train_ae: true
+train_ar: false
+phases: 400000:DO_L2-DO_LPIPS:1
+batch_size: 512
+eval_batch_size: 512
+lr_enc: 0.00015
+lr_dec: 0.00015
+lr_ar: 0.00015
+lr_gan_loss: 0.00015
+weight_decay_enc: 0.0
+weight_decay_dec: 0.0
+weight_decay_ar: 0.0
+weight_decay_gan: 0.0001
+lr_enc_min: 3.0e-05
+lr_dec_min: 3.0e-05
+lr_ar_min: 3.0e-05
+lr_gan_loss_min: 3.0e-05
+use_tf32: false
+lr_scheduler: none
+drop: 0.0
+ar_drop: 0.0
+grad_clip: 0.0
+seed: 42
+beta1: 0.5
+beta2: 0.95
+ar_beta1: 0.5
+ar_beta2: 0.95
+label_drop_prob: 0.1
+precision: bf16-mixed
+replace_cur_with_ema: false
+num_workers: 4
+devices: 8
+reverse_t: false
+use_noise_query: true
+use_tail_dropout: true
+l2_v_target: true
+tail_dropout_t2k: 1.0
+tied_timestep: false
+time_stages: 200,400,600,800,1000
+k_per_stage: 192,184,72,48,16
+load_ar_model: false
+load_ar_model_ema: false
+prior_G_loss: false
+use_l1_loss: false
+use_l2_loss: true
+use_lpips_loss: true
+use_gan_loss: false
+use_prior_loss: false
+teacher_forcing_ae: false
+disc_adaptive_weight: false
+lpips_weight: 0.1
+gan_G_weight: 0.1
+prior_enc_weight: 1.0
+l1_weight: 0.1
+l2_weight: 1.0
+gan_start: 5000
+prior_start: 0
+tail_dropout_start: 0
+tail_dropout_p: 0
+torch_compile: true
+compile_ae: true
+compile_ar: true
+compile_gan: false
+ar_ema_halflife: 2500
+ae_ema_halflife: 1000
+ema_sampling: false
+ema_reconstruction: false
+cache_kv: true
+topK: 0.0
+topP: 0.0
+cfg: 0.0
+cfg_schedule: cosine
+cfg_power: 2.75
+temperature: 1.0
+eval_data_dir: ./data/cifar10
+eval_crop_type: center
+eval_fid_ref_path: ./refs/cifar10-32x32_local.npz
+save_dir: ./outputs/flow_matching_tokenizer
+log_freq: 25
+visualize_freq: 5000
+visualize_img_num: 8
+ckpt_freq: 25000
+eval_freq: 25000
+save_ckpt: true
+ar_ckpt_path: ''
+ae_ckpt_path: ''
+resume_optimizer: false
+cond_emb_dim: 32
+greedy_coding_prob: 0.0
+z_len: 64
+z_dim: 512
+vq_type: VQ
+codebook_size: 1024
+norm_layer: RMSNorm
+ffn_type: geglu
+context_type: cross_attn
+hidden_dim: 512
+Encoder:
+  input_dim: 48
+  output_dim: 512
+  cond_dim: 32
+  cond_input_dim: 11
+  max_seq_len: 64
+  ctx_max_seq_len: 64
+  dim: 512
+  layer_num: 10
+  mlp_ratio: 4.0
+  heads: 8
+  emb_dropout: 0.0
+  drop: 0.0
+  l2_norm: false
+  context_type: cross_attn
+  rope: false
+  abs_pos: true
+  causal: false
+  norm_layer: RMSNorm
+  ffn_type: geglu
+Quantizer:
+  codebook_size: 1024
+  z_dim: 12
+  dim: 512
+  length: 64
+  cond_dim: 32
+  cond_affine_norm: false
+  vq_beta: 0.25
+  use_norm: true
+  norm_layer: l2
+  ffn_type: geglu
+  temperature: 1.0
+Decoder:
+  input_dim: 48
+  ctx_input_dim: 512
+  output_dim: 48
+  cond_dim: 32
+  cond_input_dim: 11
+  max_seq_len: 64
+  ctx_max_seq_len: 64
+  dim: 512
+  layer_num: 10
+  heads: 8
+  emb_dropout: 0.0
+  drop: 0.0
+  l2_norm: false
+  context_type: cross_attn
+  rope: false
+  abs_pos: true
+  causal: false
+  norm_layer: RMSNorm
+  ffn_type: geglu
+  mlp_ratio: 4.0
+  zero_out: true
+  query_type: noise
+ARModel:
+  input_dim: 512
+  output_dim: 1024
+  max_seq_len: 64
+  cond_dim: 32
+  cond_input_dim: 11
+  dim: 512
+  layer_num: 6
+  heads: 8
+  emb_dropout: 0.0
+  drop: 0.0
+  l2_norm: false
+  context_type: none
+  rope: false
+  abs_pos: true
+  causal: true
+  out_act: false
+  norm_layer: RMSNorm
+  ffn_type: geglu
+  mlp_ratio: 4
+  zero_out: false
+  temperature: 1.0
+GANLoss:
+  disc_loss: vanilla
+  disc_dim: 64
+  disc_type: patchgan
+  image_size: 32
+  disc_num_layers: 3
+  disc_in_channels: 3
+  gen_adv_loss: vanilla
+  lecam_loss_weight: 0.001
+  norm_type: bn
+  aug_prob: 1.0

configs/cifar10/tail_dropout_ar.yaml

@@ -0,0 +1,198 @@
+wandb_project: wAR
+wandb_name: tail_dropout_ar
+data_dir: ./data/cifar10
+xflip: true
+crop_type: center
+image_size: 32
+num_channels: 3
+num_classes: 11
+patch_size: 4
+use_label: false
+patch_dim: 48
+x_len: 64
+diffusion_decoder: false
+train_ae: false
+train_ar: true
+phases: 200000:DO_PRIOR_AR:1
+batch_size: 512
+eval_batch_size: 512
+lr_enc: 0.00015
+lr_dec: 0.00015
+lr_ar: 0.00015
+lr_gan_loss: 0.00015
+weight_decay_enc: 0.0
+weight_decay_dec: 0.0
+weight_decay_ar: 0.0001
+weight_decay_gan: 0.0001
+lr_enc_min: 3.0e-05
+lr_dec_min: 3.0e-05
+lr_ar_min: 3.0e-05
+lr_gan_loss_min: 3.0e-05
+use_tf32: false
+lr_scheduler: none
+drop: 0.0
+ar_drop: 0.1
+grad_clip: 0.0
+seed: 42
+beta1: 0.5
+beta2: 0.95
+ar_beta1: 0.9
+ar_beta2: 0.99
+label_drop_prob: 0.1
+precision: bf16-mixed
+replace_cur_with_ema: false
+num_workers: 4
+devices: 8
+reverse_t: true
+use_noise_query: false
+use_tail_dropout: true
+l2_v_target: false
+tail_dropout_t2k: 1.0
+tied_timestep: false
+time_stages: 200,400,600,800,1000
+k_per_stage: 192,184,72,48,16
+load_ar_model: false
+load_ar_model_ema: false
+prior_G_loss: false
+use_l1_loss: false
+use_l2_loss: false
+use_lpips_loss: false
+use_gan_loss: false
+use_prior_loss: false
+teacher_forcing_ae: false
+disc_adaptive_weight: false
+lpips_weight: 0.1
+gan_G_weight: 0.1
+prior_enc_weight: 1.0
+l1_weight: 0.1
+l2_weight: 1.0
+gan_start: 5000
+prior_start: 0
+tail_dropout_start: 0
+tail_dropout_p: 0.5
+torch_compile: true
+compile_ae: true
+compile_ar: true
+compile_gan: false
+ar_ema_halflife: 2000
+ae_ema_halflife: 0
+ema_sampling: true
+ema_reconstruction: false
+cache_kv: true
+topK: 0.0
+topP: 0.0
+cfg: 0.0
+cfg_schedule: cosine
+cfg_power: 2.75
+temperature: 1.0
+eval_data_dir: ./data/cifar10
+eval_crop_type: center
+eval_fid_ref_path: ./refs/cifar10-32x32_local.npz
+save_dir: ./outputs/tail_dropout_ar
+log_freq: 25
+visualize_freq: 5000
+visualize_img_num: 8
+ckpt_freq: 25000
+eval_freq: 25000
+save_ckpt: true
+ar_ckpt_path: ''
+ae_ckpt_path: ./outputs/tail_dropout_tokenizer/000/ckpts/last.ckpt
+resume_optimizer: false
+cond_emb_dim: 32
+greedy_coding_prob: 0.0
+z_len: 64
+z_dim: 512
+vq_type: VQ
+codebook_size: 1024
+norm_layer: RMSNorm
+ffn_type: geglu
+context_type: cross_attn
+hidden_dim: 512
+Encoder:
+  input_dim: 48
+  output_dim: 512
+  cond_dim: 32
+  cond_input_dim: 11
+  max_seq_len: 64
+  ctx_max_seq_len: 64
+  dim: 512
+  layer_num: 10
+  mlp_ratio: 4.0
+  heads: 8
+  emb_dropout: 0.0
+  drop: 0.0
+  l2_norm: false
+  context_type: cross_attn
+  rope: false
+  abs_pos: true
+  causal: false
+  norm_layer: RMSNorm
+  ffn_type: geglu
+Quantizer:
+  codebook_size: 1024
+  z_dim: 12
+  dim: 512
+  length: 64
+  cond_dim: 32
+  cond_affine_norm: false
+  vq_beta: 0.25
+  use_norm: true
+  norm_layer: l2
+  ffn_type: geglu
+  temperature: 1.0
+Decoder:
+  input_dim: 48
+  ctx_input_dim: 512
+  output_dim: 48
+  cond_dim: 32
+  cond_input_dim: 11
+  max_seq_len: 64
+  ctx_max_seq_len: 64
+  dim: 512
+  layer_num: 10
+  heads: 8
+  emb_dropout: 0.0
+  drop: 0.0
+  l2_norm: false
+  context_type: cross_attn
+  rope: false
+  abs_pos: true
+  causal: false
+  norm_layer: RMSNorm
+  ffn_type: geglu
+  mlp_ratio: 4.0
+  zero_out: true
+  query_type: learnable
+ARModel:
+  input_dim: 512
+  output_dim: 1024
+  max_seq_len: 64
+  cond_dim: 32
+  cond_input_dim: 11
+  dim: 512
+  layer_num: 13
+  heads: 8
+  emb_dropout: 0.0
+  drop: 0.1
+  l2_norm: false
+  context_type: none
+  rope: false
+  abs_pos: true
+  causal: true
+  out_act: false
+  norm_layer: RMSNorm
+  ffn_type: geglu
+  mlp_ratio: 4.0
+  zero_out: false
+  temperature: 1.0
+GANLoss:
+  disc_loss: vanilla
+  disc_dim: 64
+  disc_type: patchgan
+  image_size: 32
+  disc_num_layers: 3
+  disc_in_channels: 3
+  gen_adv_loss: vanilla
+  lecam_loss_weight: 0.001
+  norm_type: bn
+  aug_prob: 1.0

configs/cifar10/tail_dropout_tokenizer.yaml

@@ -0,0 +1,198 @@
+wandb_project: wAR
+wandb_name: tail_dropout_tokenizer
+data_dir: ./data/cifar10
+xflip: true
+crop_type: center
+image_size: 32
+num_channels: 3
+num_classes: 11
+patch_size: 4
+use_label: false
+patch_dim: 48
+x_len: 64
+diffusion_decoder: false
+train_ae: true
+train_ar: false
+phases: 400000:DO_L1-DO_L2-DO_LPIPS:1
+batch_size: 512
+eval_batch_size: 512
+lr_enc: 0.00015
+lr_dec: 0.00015
+lr_ar: 0.00015
+lr_gan_loss: 0.00015
+weight_decay_enc: 0.0
+weight_decay_dec: 0.0
+weight_decay_ar: 0.0
+weight_decay_gan: 0.0001
+lr_enc_min: 3.0e-05
+lr_dec_min: 3.0e-05
+lr_ar_min: 3.0e-05
+lr_gan_loss_min: 3.0e-05
+use_tf32: false
+lr_scheduler: none
+drop: 0.0
+ar_drop: 0.0
+grad_clip: 0.0
+seed: 42
+beta1: 0.5
+beta2: 0.95
+ar_beta1: 0.5
+ar_beta2: 0.95
+label_drop_prob: 0.1
+precision: bf16-mixed
+replace_cur_with_ema: false
+num_workers: 4
+devices: 8
+reverse_t: false
+use_noise_query: false
+use_tail_dropout: true
+l2_v_target: false
+tail_dropout_t2k: 1.0
+tied_timestep: false
+time_stages: 200,400,600,800,1000
+k_per_stage: 192,184,72,48,16
+load_ar_model: false
+load_ar_model_ema: false
+prior_G_loss: false
+use_l1_loss: true
+use_l2_loss: true
+use_lpips_loss: true
+use_gan_loss: false
+use_prior_loss: false
+teacher_forcing_ae: false
+disc_adaptive_weight: false
+lpips_weight: 0.1
+gan_G_weight: 0.1
+prior_enc_weight: 1.0
+l1_weight: 0.1
+l2_weight: 1.0
+gan_start: 5000
+prior_start: 0
+tail_dropout_start: 0
+tail_dropout_p: 0.5
+torch_compile: true
+compile_ae: true
+compile_ar: true
+compile_gan: false
+ar_ema_halflife: 2500
+ae_ema_halflife: 1000
+ema_sampling: false
+ema_reconstruction: false
+cache_kv: true
+topK: 0.0
+topP: 0.0
+cfg: 0.0
+cfg_schedule: cosine
+cfg_power: 2.75
+temperature: 1.0
+eval_data_dir: ./data/cifar10
+eval_crop_type: center
+eval_fid_ref_path: ./refs/cifar10-32x32_local.npz
+save_dir: ./outputs/tail_dropout_tokenizer
+log_freq: 25
+visualize_freq: 5000
+visualize_img_num: 8
+ckpt_freq: 25000
+eval_freq: 25000
+save_ckpt: true
+ar_ckpt_path: ''
+ae_ckpt_path: ''
+resume_optimizer: false
+cond_emb_dim: 32
+greedy_coding_prob: 0.0
+z_len: 64
+z_dim: 512
+vq_type: VQ
+codebook_size: 1024
+norm_layer: RMSNorm
+ffn_type: geglu
+context_type: cross_attn
+hidden_dim: 512
+Encoder:
+  input_dim: 48
+  output_dim: 512
+  cond_dim: 32
+  cond_input_dim: 11
+  max_seq_len: 64
+  ctx_max_seq_len: 64
+  dim: 512
+  layer_num: 10
+  mlp_ratio: 4.0
+  heads: 8
+  emb_dropout: 0.0
+  drop: 0.0
+  l2_norm: false
+  context_type: cross_attn
+  rope: false
+  abs_pos: true
+  causal: false
+  norm_layer: RMSNorm
+  ffn_type: geglu
+Quantizer:
+  codebook_size: 1024
+  z_dim: 12
+  dim: 512
+  length: 64
+  cond_dim: 32
+  cond_affine_norm: false
+  vq_beta: 0.25
+  use_norm: true
+  norm_layer: l2
+  ffn_type: geglu
+  temperature: 1.0
+Decoder:
+  input_dim: 48
+  ctx_input_dim: 512
+  output_dim: 48
+  cond_dim: 32
+  cond_input_dim: 11
+  max_seq_len: 64
+  ctx_max_seq_len: 64
+  dim: 512
+  layer_num: 10
+  heads: 8
+  emb_dropout: 0.0
+  drop: 0.0
+  l2_norm: false
+  context_type: cross_attn
+  rope: false
+  abs_pos: true
+  causal: false
+  norm_layer: RMSNorm
+  ffn_type: geglu
+  mlp_ratio: 4.0
+  zero_out: true
+  query_type: learnable
+ARModel:
+  input_dim: 512
+  output_dim: 1024
+  max_seq_len: 64
+  cond_dim: 32
+  cond_input_dim: 11
+  dim: 512
+  layer_num: 6
+  heads: 8
+  emb_dropout: 0.0
+  drop: 0.0
+  l2_norm: false
+  context_type: none
+  rope: false
+  abs_pos: true
+  causal: true
+  out_act: false
+  norm_layer: RMSNorm
+  ffn_type: geglu
+  mlp_ratio: 4
+  zero_out: false
+  temperature: 1.0
+GANLoss:
+  disc_loss: vanilla
+  disc_dim: 64
+  disc_type: patchgan
+  image_size: 32
+  disc_num_layers: 3
+  disc_in_channels: 3
+  gen_adv_loss: vanilla
+  lecam_loss_weight: 0.001
+  norm_type: bn
+  aug_prob: 1.0

configs/cifar10/war_ar.yaml

@@ -0,0 +1,201 @@
+wandb_project: wAR
+wandb_name: war_ar
+data_dir: ./data/cifar10
+xflip: true
+crop_type: center
+image_size: 32
+num_channels: 3
+num_classes: 11
+patch_size: 4
+use_label: false
+patch_dim: 48
+x_len: 64
+diffusion_decoder: false
+train_ae: false
+train_ar: true
+phases: 200000:DO_PRIOR_AR:1
+batch_size: 512
+eval_batch_size: 512
+lr_enc: 0.00015
+lr_dec: 0.00015
+lr_ar: 0.00015
+lr_gan_loss: 0.00015
+weight_decay_enc: 0.0
+weight_decay_dec: 0.0
+weight_decay_ar: 0.0001
+weight_decay_gan: 0.0001
+lr_enc_min: 3.0e-05
+lr_dec_min: 3.0e-05
+lr_ar_min: 3.0e-05
+lr_gan_loss_min: 3.0e-05
+use_tf32: false
+lr_scheduler: none
+drop: 0.0
+ar_drop: 0.1
+grad_clip: 0.0
+seed: 42
+beta1: 0.5
+beta2: 0.95
+ar_beta1: 0.9
+ar_beta2: 0.99
+label_drop_prob: 0.1
+precision: bf16-mixed
+replace_cur_with_ema: false
+num_workers: 4
+devices: 8
+reverse_t: true
+use_noise_query: false
+use_tail_dropout: false
+l2_v_target: false
+tail_dropout_t2k: 1.0
+tied_timestep: false
+time_stages: 200,400,600,800,1000
+k_per_stage: 192,184,72,48,16
+load_ar_model: false
+load_ar_model_ema: false
+prior_G_loss: false
+use_l1_loss: false
+use_l2_loss: false
+use_lpips_loss: false
+use_gan_loss: false
+use_prior_loss: false
+teacher_forcing_ae: false
+disc_adaptive_weight: false
+lpips_weight: 0.1
+gan_G_weight: 0.1
+prior_enc_weight: 1.0
+l1_weight: 0.1
+l2_weight: 1.0
+gan_start: 5000
+prior_start: 0
+tail_dropout_start: 0
+tail_dropout_p: 0.0
+torch_compile: true
+compile_ae: true
+compile_ar: true
+compile_gan: false
+ar_ema_halflife: 2000
+ae_ema_halflife: 0
+ema_sampling: true
+ema_reconstruction: false
+cache_kv: true
+topK: 0.0
+topP: 0.0
+cfg: 0.0
+cfg_schedule: cosine
+cfg_power: 2.75
+temperature: 1.0
+eval_data_dir: ./data/cifar10
+eval_crop_type: center
+eval_fid_ref_path: ./refs/cifar10-32x32_local.npz
+save_dir: ./outputs/war_ar
+log_freq: 25
+visualize_freq: 5000
+visualize_img_num: 8
+ckpt_freq: 25000
+eval_freq: 25000
+save_ckpt: true
+ar_ckpt_path: ''
+ae_ckpt_path: ./outputs/war_tokenizer/000/ckpts/last.ckpt
+resume_optimizer: false
+cond_emb_dim: 32
+greedy_coding_prob: 0.0
+z_len: 64
+z_dim: 512
+vq_type: VQ
+codebook_size: 1024
+norm_layer: RMSNorm
+ffn_type: geglu
+context_type: cross_attn
+hidden_dim: 512
+Encoder:
+  input_dim: 48
+  output_dim: 512
+  cond_dim: 32
+  cond_input_dim: 11
+  max_seq_len: 64
+  ctx_max_seq_len: 64
+  dim: 512
+  layer_num: 10
+  mlp_ratio: 4.0
+  heads: 8
+  emb_dropout: 0.0
+  drop: 0.0
+  l2_norm: false
+  context_type: cross_attn
+  rope: false
+  abs_pos: true
+  causal: false
+  norm_layer: RMSNorm
+  ffn_type: geglu
+Quantizer:
+  codebook_size: 1024
+  z_dim: 12
+  dim: 512
+  length: 64
+  cond_dim: 32
+  cond_affine_norm: false
+  vq_beta: 0.25
+  use_norm: true
+  norm_layer: l2
+  ffn_type: geglu
+  temperature: 0.05
+Decoder:
+  input_dim: 48
+  ctx_input_dim: 512
+  output_dim: 48
+  cond_dim: 32
+  cond_input_dim: 11
+  max_seq_len: 64
+  ctx_max_seq_len: 64
+  dim: 512
+  layer_num: 10
+  heads: 8
+  emb_dropout: 0.0
+  drop: 0.0
+  l2_norm: false
+  context_type: cross_attn
+  rope: false
+  abs_pos: true
+  causal: false
+  norm_layer: RMSNorm
+  ffn_type: geglu
+  mlp_ratio: 4.0
+  zero_out: true
+  query_type: learnable
+ARModel:
+  input_dim: 512
+  output_dim: 1024
+  max_seq_len: 64
+  cond_dim: 32
+  cond_input_dim: 11
+  dim: 512
+  layer_num: 13
+  heads: 8
+  emb_dropout: 0.0
+  drop: 0.1
+  l2_norm: false
+  context_type: none
+  rope: false
+  abs_pos: true
+  causal: true
+  out_act: false
+  norm_layer: RMSNorm
+  ffn_type: geglu
+  mlp_ratio: 4.0
+  zero_out: false
+  temperature: 1.0
+GANLoss:
+  disc_loss: vanilla
+  disc_dim: 64
+  disc_type: patchgan
+  image_size: 32
+  disc_num_layers: 3
+  disc_in_channels: 3
+  gen_adv_loss: vanilla
+  lecam_loss_weight: 0.001
+  norm_type: bn
+  aug_prob: 1.0
+prior_ar_temp_low: 0.1
+prior_ar_temp_high: 1.0
+prior_ar_temp_schedule: constant

configs/cifar10/war_tokenizer.yaml

@@ -0,0 +1,201 @@
+wandb_project: wAR
+wandb_name: war_tokenizer
+data_dir: ./data/cifar10
+xflip: true
+crop_type: center
+image_size: 32
+num_channels: 3
+num_classes: 11
+patch_size: 4
+use_label: false
+patch_dim: 48
+x_len: 64
+diffusion_decoder: false
+train_ae: true
+train_ar: true
+phases: 400000:DO_L1-DO_L2-DO_LPIPS-DO_PRIOR_ENC-DO_PRIOR_AR:1
+batch_size: 512
+eval_batch_size: 512
+lr_enc: 0.00015
+lr_dec: 0.00015
+lr_ar: 0.00015
+lr_gan_loss: 0.00015
+weight_decay_enc: 0.0
+weight_decay_dec: 0.0
+weight_decay_ar: 0.0
+weight_decay_gan: 0.0001
+lr_enc_min: 3.0e-05
+lr_dec_min: 3.0e-05
+lr_ar_min: 3.0e-05
+lr_gan_loss_min: 3.0e-05
+use_tf32: false
+lr_scheduler: none
+drop: 0.0
+ar_drop: 0.0
+grad_clip: 0.0
+seed: 42
+beta1: 0.5
+beta2: 0.95
+ar_beta1: 0.5
+ar_beta2: 0.95
+label_drop_prob: 0.1
+precision: bf16-mixed
+replace_cur_with_ema: false
+num_workers: 4
+devices: 8
+reverse_t: true
+use_noise_query: false
+use_tail_dropout: false
+l2_v_target: false
+tail_dropout_t2k: 1.0
+tied_timestep: false
+time_stages: 200,400,600,800,1000
+k_per_stage: 192,184,72,48,16
+load_ar_model: false
+load_ar_model_ema: false
+prior_G_loss: false
+use_l1_loss: true
+use_l2_loss: true
+use_lpips_loss: true
+use_gan_loss: false
+use_prior_loss: true
+teacher_forcing_ae: false
+disc_adaptive_weight: false
+lpips_weight: 0.1
+gan_G_weight: 0.1
+prior_enc_weight: 1.0
+l1_weight: 0.1
+l2_weight: 1.0
+gan_start: 5000
+prior_start: 0
+tail_dropout_start: 0
+tail_dropout_p: 0.0
+torch_compile: true
+compile_ae: true
+compile_ar: true
+compile_gan: false
+ar_ema_halflife: 0
+ae_ema_halflife: 0
+ema_sampling: false
+ema_reconstruction: false
+cache_kv: true
+topK: 0.0
+topP: 0.0
+cfg: 0.0
+cfg_schedule: cosine
+cfg_power: 2.75
+temperature: 1.0
+eval_data_dir: ./data/cifar10
+eval_crop_type: center
+eval_fid_ref_path: ./refs/cifar10-32x32_local.npz
+save_dir: ./outputs/war_tokenizer
+log_freq: 25
+visualize_freq: 5000
+visualize_img_num: 8
+ckpt_freq: 25000
+eval_freq: 25000
+save_ckpt: true
+ar_ckpt_path: ''
+ae_ckpt_path: ''
+resume_optimizer: false
+cond_emb_dim: 32
+greedy_coding_prob: 0.0
+z_len: 64
+z_dim: 512
+vq_type: VQ
+codebook_size: 1024
+norm_layer: RMSNorm
+ffn_type: geglu
+context_type: cross_attn
+hidden_dim: 512
+Encoder:
+  input_dim: 48
+  output_dim: 512
+  cond_dim: 32
+  cond_input_dim: 11
+  max_seq_len: 64
+  ctx_max_seq_len: 64
+  dim: 512
+  layer_num: 10
+  mlp_ratio: 4.0
+  heads: 8
+  emb_dropout: 0.0
+  drop: 0.0
+  l2_norm: false
+  context_type: cross_attn
+  rope: false
+  abs_pos: true
+  causal: false
+  norm_layer: RMSNorm
+  ffn_type: geglu
+Quantizer:
+  codebook_size: 1024
+  z_dim: 12
+  dim: 512
+  length: 64
+  cond_dim: 32
+  cond_affine_norm: false
+  vq_beta: 0.25
+  use_norm: true
+  norm_layer: l2
+  ffn_type: geglu
+  temperature: 0.05
+Decoder:
+  input_dim: 48
+  ctx_input_dim: 512
+  output_dim: 48
+  cond_dim: 32
+  cond_input_dim: 11
+  max_seq_len: 64
+  ctx_max_seq_len: 64
+  dim: 512
+  layer_num: 10
+  heads: 8
+  emb_dropout: 0.0
+  drop: 0.0
+  l2_norm: false
+  context_type: cross_attn
+  rope: false
+  abs_pos: true
+  causal: false
+  norm_layer: RMSNorm
+  ffn_type: geglu
+  mlp_ratio: 4.0
+  zero_out: true
+  query_type: learnable
+ARModel:
+  input_dim: 512
+  output_dim: 1024
+  max_seq_len: 64
+  cond_dim: 32
+  cond_input_dim: 11
+  dim: 512
+  layer_num: 4
+  heads: 8
+  emb_dropout: 0.0
+  drop: 0.0
+  l2_norm: false
+  context_type: none
+  rope: false
+  abs_pos: true
+  causal: true
+  out_act: false
+  norm_layer: RMSNorm
+  ffn_type: geglu
+  mlp_ratio: 4
+  zero_out: false
+  temperature: 1.0
+GANLoss:
+  disc_loss: vanilla
+  disc_dim: 64
+  disc_type: patchgan
+  image_size: 32
+  disc_num_layers: 3
+  disc_in_channels: 3
+  gen_adv_loss: vanilla
+  lecam_loss_weight: 0.001
+  norm_type: bn
+  aug_prob: 1.0
+prior_ar_temp_low: 0.1
+prior_ar_temp_high: 1.0
+prior_ar_temp_schedule: constant

dataset.py

@@ -0,0 +1,415 @@
+# Copyright (c) 2022, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+#
+# Modified work Copyright 2025 Bowen Zheng
+# The Chinese University of Hongkong, Shenzhen
+#
+# This work is licensed under a Creative Commons
+# Attribution-NonCommercial-ShareAlike 4.0 International License.
+# You should have received a copy of the license along with this
+# work. If not, see http://creativecommons.org/licenses/by-nc-sa/4.0/
+
+import os
+import sys
+import math
+import subprocess
+import numpy as np
+import zipfile
+import PIL.Image
+import json
+import torch
+import dnnlib
+from pathlib import Path
+
+try:
+    import pyspng
+except ImportError:
+    pyspng = None
+
+#----------------------------------------------------------------------------
+# Simple augmentations (center crop / random crop) adapted from ADM
+
+def center_crop_arr(pil_image, image_size):
+    """
+    Center cropping implementation from ADM.
+    https://github.com/openai/guided-diffusion/blob/8fb3ad9197f16bbc40620447b2742e13458d2831/guided_diffusion/image_datasets.py#L126
+    """
+    while min(*pil_image.size) >= 2 * image_size:
+        pil_image = pil_image.resize(
+            tuple(x // 2 for x in pil_image.size), resample=PIL.Image.BOX
+        )
+
+    scale = image_size / min(*pil_image.size)
+    pil_image = pil_image.resize(
+        tuple(round(x * scale) for x in pil_image.size), resample=PIL.Image.BICUBIC
+    )
+
+    arr = np.array(pil_image)
+    crop_y = (arr.shape[0] - image_size) // 2
+    crop_x = (arr.shape[1] - image_size) // 2
+    return PIL.Image.fromarray(arr[crop_y: crop_y + image_size, crop_x: crop_x + image_size])
+
+
+def random_crop_arr(pil_image, image_size, min_crop_frac=0.8, max_crop_frac=1.0, rng=None):
+    min_smaller_dim_size = math.ceil(image_size / max_crop_frac)
+    max_smaller_dim_size = math.ceil(image_size / min_crop_frac)
+    # Use an explicit RNG when provided (e.g., deterministic per-sample cropping).
+    # Otherwise fall back to the global NumPy RNG (which can be controlled via seeding).
+    rng = np.random if rng is None else rng
+    smaller_dim_size = rng.randint(min_smaller_dim_size, max_smaller_dim_size + 1)
+
+    # Repeated BOX downsampling preserves the original crop preprocessing.
+    while min(*pil_image.size) >= 2 * smaller_dim_size:
+        pil_image = pil_image.resize(
+            tuple(x // 2 for x in pil_image.size), resample=PIL.Image.BOX
+        )
+
+    scale = smaller_dim_size / min(*pil_image.size)
+    pil_image = pil_image.resize(
+        tuple(round(x * scale) for x in pil_image.size), resample=PIL.Image.BICUBIC
+    )
+
+    arr = np.array(pil_image)
+    crop_y = rng.randint(0, arr.shape[0] - image_size + 1)
+    crop_x = rng.randint(0, arr.shape[1] - image_size + 1)
+    return PIL.Image.fromarray(arr[crop_y: crop_y + image_size, crop_x: crop_x + image_size])
+
+class Dataset(torch.utils.data.Dataset):
+    def __init__(self,
+        name,                   # Name of the dataset.
+        raw_shape,              # Shape of the raw image data (NCHW).
+        max_size    = None,     # Artificially limit the size of the dataset. None = no limit. Applied before xflip.
+        use_label  = False,    # Enable conditioning labels? False = label dimension is zero.
+        xflip       = False,    # Artificially double the size of the dataset via x-flips. Applied after max_size.
+        random_seed = 0,        # Random seed to use when applying max_size.
+            ):
+        self._name = name
+        self._raw_shape = list(raw_shape)
+        self._use_label = use_label
+        self._raw_labels = None
+        self._label_shape = None
+
+        # Apply max_size.
+        self._raw_idx = np.arange(self._raw_shape[0], dtype=np.int64)
+        if (max_size is not None) and (self._raw_idx.size > max_size):
+            np.random.RandomState(random_seed % (1 << 31)).shuffle(self._raw_idx)
+            self._raw_idx = np.sort(self._raw_idx[:max_size])
+
+        # Apply xflip.
+        self._xflip = np.zeros(self._raw_idx.size, dtype=np.uint8)
+        if xflip:
+            self._raw_idx = np.tile(self._raw_idx, 2)
+            self._xflip = np.concatenate([self._xflip, np.ones_like(self._xflip)])
+
+    def _get_raw_labels(self):
+        if self._raw_labels is None:
+            self._raw_labels = self._load_raw_labels() if self._use_label else None
+            if self._raw_labels is None:
+                self._raw_labels = np.zeros([self._raw_shape[0], 0], dtype=np.float32)
+            assert isinstance(self._raw_labels, np.ndarray)
+            assert self._raw_labels.shape[0] == self._raw_shape[0]
+            assert self._raw_labels.dtype in [np.float32, np.int64]
+            if self._raw_labels.dtype == np.int64:
+                assert self._raw_labels.ndim == 1
+                assert np.all(self._raw_labels >= 0)
+        return self._raw_labels
+
+    def close(self): # to be overridden by subclass
+        pass
+
+    def _load_raw_image(self, raw_idx): # to be overridden by subclass
+        raise NotImplementedError
+
+    def _load_raw_labels(self): # to be overridden by subclass
+        raise NotImplementedError
+
+    def __getstate__(self):
+        return dict(self.__dict__, _raw_labels=None)
+
+    def __del__(self):
+        try:
+            self.close()
+        except:
+            pass
+
+    def __len__(self):
+        return self._raw_idx.size
+
+    def __getitem__(self, idx):
+        raw_idx = self._raw_idx[idx]
+        image = self._load_raw_image(raw_idx)
+        assert isinstance(image, np.ndarray)
+        assert list(image.shape) == self.image_shape
+        assert image.dtype == np.uint8
+        if self._xflip[idx]:
+            assert image.ndim == 3 # CHW
+            image = image[:, :, ::-1]
+        return image.copy(), self.get_label(idx)
+
+    def get_label(self, idx):
+        label = self._get_raw_labels()[self._raw_idx[idx]]
+        if label.dtype == np.int64:
+            onehot = np.zeros(self.label_shape, dtype=np.float32)
+            onehot[label] = 1
+            label = onehot
+        return label.copy()
+
+    def get_details(self, idx):
+        d = dnnlib.EasyDict()
+        d.raw_idx = int(self._raw_idx[idx])
+        d.xflip = (int(self._xflip[idx]) != 0)
+        d.raw_label = self._get_raw_labels()[d.raw_idx].copy()
+        return d
+
+    @property
+    def name(self):
+        return self._name
+
+    @property
+    def image_shape(self):
+        return list(self._raw_shape[1:])
+
+    @property
+    def num_channels(self):
+        assert len(self.image_shape) == 3 # CHW
+        return self.image_shape[0]
+
+    @property
+    def resolution(self):
+        assert len(self.image_shape) == 3 # CHW
+        assert self.image_shape[1] == self.image_shape[2]
+        return self.image_shape[1]
+
+    @property
+    def label_shape(self):
+        if self._label_shape is None:
+            raw_labels = self._get_raw_labels()
+            if raw_labels.dtype == np.int64:
+                self._label_shape = [int(np.max(raw_labels)) + 1]
+            else:
+                self._label_shape = raw_labels.shape[1:]
+        return list(self._label_shape)
+
+    @property
+    def label_dim(self):
+        assert len(self.label_shape) == 1
+        return self.label_shape[0]
+
+    @property
+    def has_labels(self):
+        return any(x != 0 for x in self.label_shape)
+
+    @property
+    def has_onehot_labels(self):
+        return self._get_raw_labels().dtype == np.int64
+
+    # ------------------------------------------------------------------------
+    # Simple augmentations (center crop / random crop) as Dataset helpers,
+    def _center_crop_arr(self, pil_image, image_size):
+        """
+        Center cropping implementation from ADM.
+        https://github.com/openai/guided-diffusion/blob/8fb3ad9197f16bbc40620447b2742e13458d2831/guided_diffusion/image_datasets.py#L126
+        """
+        while min(*pil_image.size) >= 2 * image_size:
+            pil_image = pil_image.resize(
+                tuple(x // 2 for x in pil_image.size), resample=PIL.Image.BOX
+            )
+
+        scale = image_size / min(*pil_image.size)
+        pil_image = pil_image.resize(
+            tuple(round(x * scale) for x in pil_image.size), resample=PIL.Image.BICUBIC
+        )
+
+        arr = np.array(pil_image)
+        crop_y = (arr.shape[0] - image_size) // 2
+        crop_x = (arr.shape[1] - image_size) // 2
+        return PIL.Image.fromarray(arr[crop_y: crop_y + image_size, crop_x: crop_x + image_size])
+
+    def _random_crop_arr(self, pil_image, image_size, min_crop_frac=0.8, max_crop_frac=1.0, rng=None):
+        min_smaller_dim_size = math.ceil(image_size / max_crop_frac)
+        max_smaller_dim_size = math.ceil(image_size / min_crop_frac)
+        rng = np.random if rng is None else rng
+        smaller_dim_size = rng.randint(min_smaller_dim_size, max_smaller_dim_size + 1)
+
+        # Repeated BOX downsampling preserves the original crop preprocessing.
+        while min(*pil_image.size) >= 2 * smaller_dim_size:
+            pil_image = pil_image.resize(
+                tuple(x // 2 for x in pil_image.size), resample=PIL.Image.BOX
+            )
+
+        scale = smaller_dim_size / min(*pil_image.size)
+        pil_image = pil_image.resize(
+            tuple(round(x * scale) for x in pil_image.size), resample=PIL.Image.BICUBIC
+        )
+
+        arr = np.array(pil_image)
+        crop_y = rng.randint(0, arr.shape[0] - image_size + 1)
+        crop_x = rng.randint(0, arr.shape[1] - image_size + 1)
+        return PIL.Image.fromarray(arr[crop_y: crop_y + image_size, crop_x: crop_x + image_size])
+
+class ImageFolderDataset(Dataset):
+    def __init__(self,
+        path,                   # Path to directory or zip.
+        resolution      = None, # Ensure specific resolution, None = highest available.
+        use_pyspng      = True, # Use pyspng if available?
+        crop_type       = None, # None / 'center' / 'random'
+        deterministic_crop = False, # If True and crop_type=='random', crop depends only on (crop_seed, raw_idx).
+        crop_seed       = 0,
+        **super_kwargs,         # Additional arguments for the Dataset base class.
+    ):
+        self._path = path
+        self._use_pyspng = use_pyspng
+        self._zipfile = None
+        self._crop_type = crop_type
+        self._resolution = resolution
+        self._deterministic_crop = bool(deterministic_crop)
+        self._crop_seed = int(crop_seed)
+
+        # If labels are requested and no dataset.json is present, try to
+        # create it in-place using dataset_tools.py (labels-only mode).
+        use_label = bool(super_kwargs.get('use_label', False))
+        if use_label:
+            self._ensure_dataset_json()
+
+        if os.path.isdir(self._path):
+            self._type = 'dir'
+            self._all_fnames = {os.path.relpath(os.path.join(root, fname), start=self._path) for root, _dirs, files in os.walk(self._path) for fname in files}
+        elif self._file_ext(self._path) == '.zip':
+            self._type = 'zip'
+            self._all_fnames = set(self._get_zipfile().namelist())
+        else:
+            raise IOError('Path must point to a directory or zip')
+
+        PIL.Image.init()
+        self._image_fnames = sorted(fname for fname in self._all_fnames if self._file_ext(fname) in PIL.Image.EXTENSION)
+        if len(self._image_fnames) == 0:
+            raise IOError('No image files found in the specified path')
+
+        name = os.path.splitext(os.path.basename(self._path))[0]
+        raw_shape = [len(self._image_fnames)] + list(self._load_raw_image(0).shape)
+        if resolution is not None and (raw_shape[2] != resolution or raw_shape[3] != resolution):
+            raise IOError('Image files do not match the specified resolution')
+        super().__init__(name=name, raw_shape=raw_shape, **super_kwargs)
+
+    @staticmethod
+    def _file_ext(fname):
+        return os.path.splitext(fname)[1].lower()
+
+    def _get_zipfile(self):
+        assert self._type == 'zip'
+        if self._zipfile is None:
+            self._zipfile = zipfile.ZipFile(self._path)
+        return self._zipfile
+
+    def _open_file(self, fname):
+        if self._type == 'dir':
+            return open(os.path.join(self._path, fname), 'rb')
+        if self._type == 'zip':
+            return self._get_zipfile().open(fname, 'r')
+        return None
+
+    def close(self):
+        try:
+            if self._zipfile is not None:
+                self._zipfile.close()
+        finally:
+            self._zipfile = None
+
+    def __getstate__(self):
+        return dict(super().__getstate__(), _zipfile=None)
+
+    def _ensure_dataset_json(self):
+        """Ensure dataset.json with labels exists for this dataset path.
+
+        If missing, call dataset_tools.py in labels-only mode to generate it.
+        """
+        # Directory: check for dataset.json at root.
+        if os.path.isdir(self._path):
+            json_path = os.path.join(self._path, 'dataset.json')
+            if os.path.isfile(json_path):
+                return
+        # Zip: check if dataset.json is inside archive.
+        elif self._file_ext(self._path) == '.zip':
+            try:
+                with zipfile.ZipFile(self._path, mode='r') as z:
+                    if 'dataset.json' in z.namelist():
+                        return
+            except Exception:
+                return
+        else:
+            return
+
+        # Use the dataset preparation tool bundled with this package.
+        project_root = Path(__file__).resolve().parent
+        tools_path = project_root / 'scripts' / 'dataset_tools.py'
+        if not tools_path.is_file():
+            print(f"Warning: dataset_tools.py not found at {tools_path}, cannot auto-generate dataset.json labels.")
+            return
+
+        cmd = [
+            sys.executable,
+            str(tools_path),
+            '--source', str(self._path),
+            '--dest', str(self._path),
+            '--labels-only',
+        ]
+        try:
+            subprocess.run(cmd, check=True)
+        except Exception as e:
+            print(f"Warning: failed to generate dataset.json via dataset_tools.py: {e}")
+
+    def _load_raw_image(self, raw_idx):
+        fname = self._image_fnames[raw_idx]
+
+        # If no crop augmentation is requested, keep original fast path
+        # (optionally using pyspng for PNGs).
+        if self._crop_type is None:
+            with self._open_file(fname) as f:
+                if self._use_pyspng and pyspng is not None and self._file_ext(fname) == '.png':
+                    image = pyspng.load(f.read())
+                else:
+                    image = np.array(PIL.Image.open(f))
+        else:
+            # Use PIL path so we can apply center/random crop.
+            if self._resolution is None:
+                raise ValueError("resolution must be set when using crop_type for ImageFolderDataset.")
+            with self._open_file(fname) as f:
+                pil_img = PIL.Image.open(f).convert("RGB")
+                if self._crop_type == 'center':
+                    pil_img = self._center_crop_arr(pil_img, self._resolution)
+                elif self._crop_type == 'random':
+                    rng = None
+                    if self._deterministic_crop:
+                        # Make cropping independent of call order / other RNG consumers.
+                        # Note: this fixes the crop for a given raw_idx across epochs.
+                        seed = (self._crop_seed + int(raw_idx)) % (1 << 31)
+                        rng = np.random.RandomState(seed)
+                    pil_img = self._random_crop_arr(pil_img, self._resolution, rng=rng)
+                else:
+                    raise ValueError(f"Unknown crop_type: {self._crop_type}")
+                image = np.array(pil_img)
+
+        # Normalize channel count to 3 (RGB) where possible.
+        if image.ndim == 2:
+            image = np.stack([image] * 3, axis=-1)  # HW => HWC (3-channel grayscale)
+        elif image.ndim == 3:
+            if image.shape[2] == 1:
+                image = np.repeat(image, 3, axis=2)
+            elif image.shape[2] > 3:
+                image = image[:, :, :3]
+
+        image = image.transpose(2, 0, 1) # HWC => CHW
+        return image
+
+    def _load_raw_labels(self):
+        fname = 'dataset.json'
+        if fname not in self._all_fnames:
+            return None
+        with self._open_file(fname) as f:
+            labels = json.load(f)['labels']
+        if labels is None:
+            return None
+        labels = dict(labels)
+        labels = [labels[fname.replace('\\', '/')] for fname in self._image_fnames]
+        labels = np.array(labels)
+        labels = labels.astype({1: np.int64, 2: np.float32}[labels.ndim])
+        return labels

dnnlib/__init__.py

@@ -0,0 +1,8 @@
+# Copyright (c) 2022, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+#
+# This work is licensed under a Creative Commons
+# Attribution-NonCommercial-ShareAlike 4.0 International License.
+# You should have received a copy of the license along with this
+# work. If not, see http://creativecommons.org/licenses/by-nc-sa/4.0/
+
+from .util import EasyDict, make_cache_dir_path

dnnlib/util.py

@@ -0,0 +1,491 @@
+# Copyright (c) 2022, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+#
+# This work is licensed under a Creative Commons
+# Attribution-NonCommercial-ShareAlike 4.0 International License.
+# You should have received a copy of the license along with this
+# work. If not, see http://creativecommons.org/licenses/by-nc-sa/4.0/
+
+"""Miscellaneous utility classes and functions."""
+
+import ctypes
+import fnmatch
+import importlib
+import inspect
+import numpy as np
+import os
+import shutil
+import sys
+import types
+import io
+import pickle
+import re
+import requests
+import html
+import hashlib
+import glob
+import tempfile
+import urllib
+import urllib.request
+import uuid
+
+from distutils.util import strtobool
+from typing import Any, List, Tuple, Union, Optional
+
+
+# Util classes
+# ------------------------------------------------------------------------------------------
+
+
+class EasyDict(dict):
+    """Convenience class that behaves like a dict but allows access with the attribute syntax."""
+
+    def __getattr__(self, name: str) -> Any:
+        try:
+            return self[name]
+        except KeyError:
+            raise AttributeError(name)
+
+    def __setattr__(self, name: str, value: Any) -> None:
+        self[name] = value
+
+    def __delattr__(self, name: str) -> None:
+        del self[name]
+
+
+class Logger(object):
+    """Redirect stderr to stdout, optionally print stdout to a file, and optionally force flushing on both stdout and the file."""
+
+    def __init__(self, file_name: Optional[str] = None, file_mode: str = "w", should_flush: bool = True):
+        self.file = None
+
+        if file_name is not None:
+            self.file = open(file_name, file_mode)
+
+        self.should_flush = should_flush
+        self.stdout = sys.stdout
+        self.stderr = sys.stderr
+
+        sys.stdout = self
+        sys.stderr = self
+
+    def __enter__(self) -> "Logger":
+        return self
+
+    def __exit__(self, exc_type: Any, exc_value: Any, traceback: Any) -> None:
+        self.close()
+
+    def write(self, text: Union[str, bytes]) -> None:
+        """Write text to stdout (and a file) and optionally flush."""
+        if isinstance(text, bytes):
+            text = text.decode()
+        if len(text) == 0: # workaround for a bug in VSCode debugger: sys.stdout.write(''); sys.stdout.flush() => crash
+            return
+
+        if self.file is not None:
+            self.file.write(text)
+
+        self.stdout.write(text)
+
+        if self.should_flush:
+            self.flush()
+
+    def flush(self) -> None:
+        """Flush written text to both stdout and a file, if open."""
+        if self.file is not None:
+            self.file.flush()
+
+        self.stdout.flush()
+
+    def close(self) -> None:
+        """Flush, close possible files, and remove stdout/stderr mirroring."""
+        self.flush()
+
+        # if using multiple loggers, prevent closing in wrong order
+        if sys.stdout is self:
+            sys.stdout = self.stdout
+        if sys.stderr is self:
+            sys.stderr = self.stderr
+
+        if self.file is not None:
+            self.file.close()
+            self.file = None
+
+
+# Cache directories
+# ------------------------------------------------------------------------------------------
+
+_dnnlib_cache_dir = None
+
+def set_cache_dir(path: str) -> None:
+    global _dnnlib_cache_dir
+    _dnnlib_cache_dir = path
+
+def make_cache_dir_path(*paths: str) -> str:
+    if _dnnlib_cache_dir is not None:
+        return os.path.join(_dnnlib_cache_dir, *paths)
+    if 'DNNLIB_CACHE_DIR' in os.environ:
+        return os.path.join(os.environ['DNNLIB_CACHE_DIR'], *paths)
+    if 'HOME' in os.environ:
+        return os.path.join(os.environ['HOME'], '.cache', 'dnnlib', *paths)
+    if 'USERPROFILE' in os.environ:
+        return os.path.join(os.environ['USERPROFILE'], '.cache', 'dnnlib', *paths)
+    return os.path.join(tempfile.gettempdir(), '.cache', 'dnnlib', *paths)
+
+# Small util functions
+# ------------------------------------------------------------------------------------------
+
+
+def format_time(seconds: Union[int, float]) -> str:
+    """Convert the seconds to human readable string with days, hours, minutes and seconds."""
+    s = int(np.rint(seconds))
+
+    if s < 60:
+        return "{0}s".format(s)
+    elif s < 60 * 60:
+        return "{0}m {1:02}s".format(s // 60, s % 60)
+    elif s < 24 * 60 * 60:
+        return "{0}h {1:02}m {2:02}s".format(s // (60 * 60), (s // 60) % 60, s % 60)
+    else:
+        return "{0}d {1:02}h {2:02}m".format(s // (24 * 60 * 60), (s // (60 * 60)) % 24, (s // 60) % 60)
+
+
+def format_time_brief(seconds: Union[int, float]) -> str:
+    """Convert the seconds to human readable string with days, hours, minutes and seconds."""
+    s = int(np.rint(seconds))
+
+    if s < 60:
+        return "{0}s".format(s)
+    elif s < 60 * 60:
+        return "{0}m {1:02}s".format(s // 60, s % 60)
+    elif s < 24 * 60 * 60:
+        return "{0}h {1:02}m".format(s // (60 * 60), (s // 60) % 60)
+    else:
+        return "{0}d {1:02}h".format(s // (24 * 60 * 60), (s // (60 * 60)) % 24)
+
+
+def ask_yes_no(question: str) -> bool:
+    """Ask the user the question until the user inputs a valid answer."""
+    while True:
+        try:
+            print("{0} [y/n]".format(question))
+            return strtobool(input().lower())
+        except ValueError:
+            pass
+
+
+def tuple_product(t: Tuple) -> Any:
+    """Calculate the product of the tuple elements."""
+    result = 1
+
+    for v in t:
+        result *= v
+
+    return result
+
+
+_str_to_ctype = {
+    "uint8": ctypes.c_ubyte,
+    "uint16": ctypes.c_uint16,
+    "uint32": ctypes.c_uint32,
+    "uint64": ctypes.c_uint64,
+    "int8": ctypes.c_byte,
+    "int16": ctypes.c_int16,
+    "int32": ctypes.c_int32,
+    "int64": ctypes.c_int64,
+    "float32": ctypes.c_float,
+    "float64": ctypes.c_double
+}
+
+
+def get_dtype_and_ctype(type_obj: Any) -> Tuple[np.dtype, Any]:
+    """Given a type name string (or an object having a __name__ attribute), return matching Numpy and ctypes types that have the same size in bytes."""
+    type_str = None
+
+    if isinstance(type_obj, str):
+        type_str = type_obj
+    elif hasattr(type_obj, "__name__"):
+        type_str = type_obj.__name__
+    elif hasattr(type_obj, "name"):
+        type_str = type_obj.name
+    else:
+        raise RuntimeError("Cannot infer type name from input")
+
+    assert type_str in _str_to_ctype.keys()
+
+    my_dtype = np.dtype(type_str)
+    my_ctype = _str_to_ctype[type_str]
+
+    assert my_dtype.itemsize == ctypes.sizeof(my_ctype)
+
+    return my_dtype, my_ctype
+
+
+def is_pickleable(obj: Any) -> bool:
+    try:
+        with io.BytesIO() as stream:
+            pickle.dump(obj, stream)
+        return True
+    except:
+        return False
+
+
+# Functionality to import modules/objects by name, and call functions by name
+# ------------------------------------------------------------------------------------------
+
+def get_module_from_obj_name(obj_name: str) -> Tuple[types.ModuleType, str]:
+    """Searches for the underlying module behind the name to some python object.
+    Returns the module and the object name (original name with module part removed)."""
+
+    # allow convenience shorthands, substitute them by full names
+    obj_name = re.sub("^np.", "numpy.", obj_name)
+    obj_name = re.sub("^tf.", "tensorflow.", obj_name)
+
+    # list alternatives for (module_name, local_obj_name)
+    parts = obj_name.split(".")
+    name_pairs = [(".".join(parts[:i]), ".".join(parts[i:])) for i in range(len(parts), 0, -1)]
+
+    # try each alternative in turn
+    for module_name, local_obj_name in name_pairs:
+        try:
+            module = importlib.import_module(module_name) # may raise ImportError
+            get_obj_from_module(module, local_obj_name) # may raise AttributeError
+            return module, local_obj_name
+        except:
+            pass
+
+    # maybe some of the modules themselves contain errors?
+    for module_name, _local_obj_name in name_pairs:
+        try:
+            importlib.import_module(module_name) # may raise ImportError
+        except ImportError:
+            if not str(sys.exc_info()[1]).startswith("No module named '" + module_name + "'"):
+                raise
+
+    # maybe the requested attribute is missing?
+    for module_name, local_obj_name in name_pairs:
+        try:
+            module = importlib.import_module(module_name) # may raise ImportError
+            get_obj_from_module(module, local_obj_name) # may raise AttributeError
+        except ImportError:
+            pass
+
+    # we are out of luck, but we have no idea why
+    raise ImportError(obj_name)
+
+
+def get_obj_from_module(module: types.ModuleType, obj_name: str) -> Any:
+    """Traverses the object name and returns the last (rightmost) python object."""
+    if obj_name == '':
+        return module
+    obj = module
+    for part in obj_name.split("."):
+        obj = getattr(obj, part)
+    return obj
+
+
+def get_obj_by_name(name: str) -> Any:
+    """Finds the python object with the given name."""
+    module, obj_name = get_module_from_obj_name(name)
+    return get_obj_from_module(module, obj_name)
+
+
+def call_func_by_name(*args, func_name: str = None, **kwargs) -> Any:
+    """Finds the python object with the given name and calls it as a function."""
+    assert func_name is not None
+    func_obj = get_obj_by_name(func_name)
+    assert callable(func_obj)
+    return func_obj(*args, **kwargs)
+
+
+def construct_class_by_name(*args, class_name: str = None, **kwargs) -> Any:
+    """Finds the python class with the given name and constructs it with the given arguments."""
+    return call_func_by_name(*args, func_name=class_name, **kwargs)
+
+
+def get_module_dir_by_obj_name(obj_name: str) -> str:
+    """Get the directory path of the module containing the given object name."""
+    module, _ = get_module_from_obj_name(obj_name)
+    return os.path.dirname(inspect.getfile(module))
+
+
+def is_top_level_function(obj: Any) -> bool:
+    """Determine whether the given object is a top-level function, i.e., defined at module scope using 'def'."""
+    return callable(obj) and obj.__name__ in sys.modules[obj.__module__].__dict__
+
+
+def get_top_level_function_name(obj: Any) -> str:
+    """Return the fully-qualified name of a top-level function."""
+    assert is_top_level_function(obj)
+    module = obj.__module__
+    if module == '__main__':
+        module = os.path.splitext(os.path.basename(sys.modules[module].__file__))[0]
+    return module + "." + obj.__name__
+
+
+# File system helpers
+# ------------------------------------------------------------------------------------------
+
+def list_dir_recursively_with_ignore(dir_path: str, ignores: List[str] = None, add_base_to_relative: bool = False) -> List[Tuple[str, str]]:
+    """List all files recursively in a given directory while ignoring given file and directory names.
+    Returns list of tuples containing both absolute and relative paths."""
+    assert os.path.isdir(dir_path)
+    base_name = os.path.basename(os.path.normpath(dir_path))
+
+    if ignores is None:
+        ignores = []
+
+    result = []
+
+    for root, dirs, files in os.walk(dir_path, topdown=True):
+        for ignore_ in ignores:
+            dirs_to_remove = [d for d in dirs if fnmatch.fnmatch(d, ignore_)]
+
+            # dirs need to be edited in-place
+            for d in dirs_to_remove:
+                dirs.remove(d)
+
+            files = [f for f in files if not fnmatch.fnmatch(f, ignore_)]
+
+        absolute_paths = [os.path.join(root, f) for f in files]
+        relative_paths = [os.path.relpath(p, dir_path) for p in absolute_paths]
+
+        if add_base_to_relative:
+            relative_paths = [os.path.join(base_name, p) for p in relative_paths]
+
+        assert len(absolute_paths) == len(relative_paths)
+        result += zip(absolute_paths, relative_paths)
+
+    return result
+
+
+def copy_files_and_create_dirs(files: List[Tuple[str, str]]) -> None:
+    """Takes in a list of tuples of (src, dst) paths and copies files.
+    Will create all necessary directories."""
+    for file in files:
+        target_dir_name = os.path.dirname(file[1])
+
+        # will create all intermediate-level directories
+        if not os.path.exists(target_dir_name):
+            os.makedirs(target_dir_name)
+
+        shutil.copyfile(file[0], file[1])
+
+
+# URL helpers
+# ------------------------------------------------------------------------------------------
+
+def is_url(obj: Any, allow_file_urls: bool = False) -> bool:
+    """Determine whether the given object is a valid URL string."""
+    if not isinstance(obj, str) or not "://" in obj:
+        return False
+    if allow_file_urls and obj.startswith('file://'):
+        return True
+    try:
+        res = requests.compat.urlparse(obj)
+        if not res.scheme or not res.netloc or not "." in res.netloc:
+            return False
+        res = requests.compat.urlparse(requests.compat.urljoin(obj, "/"))
+        if not res.scheme or not res.netloc or not "." in res.netloc:
+            return False
+    except:
+        return False
+    return True
+
+
+def open_url(url: str, cache_dir: str = None, num_attempts: int = 10, verbose: bool = True, return_filename: bool = False, cache: bool = True) -> Any:
+    """Download the given URL and return a binary-mode file object to access the data."""
+    assert num_attempts >= 1
+    assert not (return_filename and (not cache))
+
+    # Doesn't look like an URL scheme so interpret it as a local filename.
+    if not re.match('^[a-z]+://', url):
+        return url if return_filename else open(url, "rb")
+
+    # Handle file URLs.  This code handles unusual file:// patterns that
+    # arise on Windows:
+    #
+    # file:///c:/foo.txt
+    #
+    # which would translate to a local '/c:/foo.txt' filename that's
+    # invalid.  Drop the forward slash for such pathnames.
+    #
+    # If you touch this code path, you should test it on both Linux and
+    # Windows.
+    #
+    # Some internet resources suggest using urllib.request.url2pathname() but
+    # but that converts forward slashes to backslashes and this causes
+    # its own set of problems.
+    if url.startswith('file://'):
+        filename = urllib.parse.urlparse(url).path
+        if re.match(r'^/[a-zA-Z]:', filename):
+            filename = filename[1:]
+        return filename if return_filename else open(filename, "rb")
+
+    assert is_url(url)
+
+    # Lookup from cache.
+    if cache_dir is None:
+        cache_dir = make_cache_dir_path('downloads')
+
+    url_md5 = hashlib.md5(url.encode("utf-8")).hexdigest()
+    if cache:
+        cache_files = glob.glob(os.path.join(cache_dir, url_md5 + "_*"))
+        if len(cache_files) == 1:
+            filename = cache_files[0]
+            return filename if return_filename else open(filename, "rb")
+
+    # Download.
+    url_name = None
+    url_data = None
+    with requests.Session() as session:
+        if verbose:
+            print("Downloading %s ..." % url, end="", flush=True)
+        for attempts_left in reversed(range(num_attempts)):
+            try:
+                with session.get(url) as res:
+                    res.raise_for_status()
+                    if len(res.content) == 0:
+                        raise IOError("No data received")
+
+                    if len(res.content) < 8192:
+                        content_str = res.content.decode("utf-8")
+                        if "download_warning" in res.headers.get("Set-Cookie", ""):
+                            links = [html.unescape(link) for link in content_str.split('"') if "export=download" in link]
+                            if len(links) == 1:
+                                url = requests.compat.urljoin(url, links[0])
+                                raise IOError("Google Drive virus checker nag")
+                        if "Google Drive - Quota exceeded" in content_str:
+                            raise IOError("Google Drive download quota exceeded -- please try again later")
+
+                    match = re.search(r'filename="([^"]*)"', res.headers.get("Content-Disposition", ""))
+                    url_name = match[1] if match else url
+                    url_data = res.content
+                    if verbose:
+                        print(" done")
+                    break
+            except KeyboardInterrupt:
+                raise
+            except:
+                if not attempts_left:
+                    if verbose:
+                        print(" failed")
+                    raise
+                if verbose:
+                    print(".", end="", flush=True)
+
+    # Save to cache.
+    if cache:
+        safe_name = re.sub(r"[^0-9a-zA-Z-._]", "_", url_name)
+        safe_name = safe_name[:min(len(safe_name), 128)]
+        cache_file = os.path.join(cache_dir, url_md5 + "_" + safe_name)
+        temp_file = os.path.join(cache_dir, "tmp_" + uuid.uuid4().hex + "_" + url_md5 + "_" + safe_name)
+        os.makedirs(cache_dir, exist_ok=True)
+        with open(temp_file, "wb") as f:
+            f.write(url_data)
+        os.replace(temp_file, cache_file) # atomic
+        if return_filename:
+            return cache_file
+
+    # Return data as file object.
+    assert not return_filename
+    return io.BytesIO(url_data)

evaluate.py

@@ -0,0 +1,181 @@
+"""Pure evaluation: no optimizers, training steps, or EMA updates."""
+import argparse
+import json
+import hashlib
+import os
+from pathlib import Path
+import subprocess
+import sys
+import numpy as np
+import torch
+from torch.utils.data import DataLoader, Subset
+from torchvision.utils import save_image
+from tqdm import tqdm
+from dataset import ImageFolderDataset
+from omegaconf import OmegaConf
+import torch.nn.functional as F
+from models import Encoder, Decoder, ARModel
+from train import (seed_everything, patchify, img_unint8_to_norm, DiTiScheduler,
+                   reconstruction, sampling, euler_recon_loop, euler_sample_loop,
+                   img_norm_to_uint8, unpatchify)
+
+
+def state_dict(state):
+    return {k.removeprefix('_orig_mod.'): v for k, v in state.items()}
+
+
+def load_models(config, checkpoint, device='cuda', ema=True, need_ar=True):
+    state = torch.load(checkpoint, map_location='cpu', weights_only=True, mmap=True)
+    enc, dec = Encoder(config), Decoder(config)
+    enc.load_state_dict(state_dict(state['enc']), strict=True)
+    dec.load_state_dict(state_dict(state['dec']), strict=True)
+    ar = None
+    if need_ar:
+        ar = ARModel(config)
+        key = 'ar_model_ema' if ema else 'ar_model'
+        if state.get(key) is None:
+            raise ValueError(f'Checkpoint has no {key}; select a matching AR checkpoint')
+        ar.load_state_dict(state_dict(state[key]), strict=True)
+    for model in (enc, dec, ar):
+        if model is not None:
+            model.to(device).eval().requires_grad_(False)
+    return enc, dec, ar
+
+
+def scheduler(config):
+    if config.get('time_stages') and config.get('k_per_stage'):
+        return DiTiScheduler(1000, config.z_len, config.time_stages,
+                             config.k_per_stage, config.reverse_t)
+    return None
+
+
+def labels_for(config, batch_size, device):
+    if config.use_label:
+        raise ValueError('The released CIFAR-10 protocols are unconditional')
+    return F.one_hot(torch.full((batch_size,), config.num_classes-1,
+                               device=device, dtype=torch.long), config.num_classes).float()
+
+
+@torch.no_grad()
+def decode_batch(config, enc, dec, ar, x, labels, mode, diti=None):
+    """x is a patchified batch; its shape also defines unconditional generation."""
+    t = torch.ones(x.shape[0], device=x.device) if config.use_noise_query else None
+    # wARTok draws reconstruction noise before the generation branch, even for AR-only eval.
+    noise_recon = torch.randn_like(x) if config.use_noise_query else None
+    if mode == 'reconstruct':
+        if config.use_noise_query and config.diffusion_decoder:
+            result, _ = euler_recon_loop(enc, dec, x, labels, config=config, diti_scheduler=diti)
+        else:
+            result, _ = reconstruction(enc, dec, x, labels, query_tokens=noise_recon, t=t)
+            if config.l2_v_target:
+                result = noise_recon - result
+    elif mode == 'generate':
+        kwargs = dict(bz=x.shape[0], class_label=labels, temperature=config.temperature,
+                      topK=config.topK, topP=config.topP, cfg=config.cfg,
+                      cfg_schedule=config.cfg_schedule, cfg_power=config.cfg_power,
+                      cache_kv=config.cache_kv, config=config)
+        if config.use_noise_query and config.diffusion_decoder:
+            result = euler_sample_loop(enc, dec, ar, **kwargs, diti_scheduler=diti)
+        else:
+            noise = torch.randn_like(x) if config.use_noise_query else None
+            result = sampling(enc, dec, ar, **kwargs, t=t, query_tokens=noise)
+            if config.l2_v_target:
+                result = noise - result
+    else:
+        raise ValueError(mode)
+    return img_norm_to_uint8(unpatchify(result, config.image_size, config.patch_size))
+
+
+def build_parser():
+    parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.ArgumentDefaultsHelpFormatter)
+    parser.add_argument('--config', required=True)
+    parser.add_argument('--checkpoint', required=True)
+    parser.add_argument('--mode', choices=['generate', 'reconstruct'], default='generate')
+    parser.add_argument('--data', help='Required for reconstruction; use the CIFAR-10 training split')
+    parser.add_argument('--output', required=True)
+    parser.add_argument('--num-samples', type=int, default=50000)
+    parser.add_argument('--batch-size', type=int, default=64, help='Per process batch size (original evaluation: 64)')
+    parser.add_argument('--seed', type=int, default=42)
+    parser.add_argument('--legacy', action=argparse.BooleanOptionalAction, default=True,
+                        help='Use the same seed on every rank (wARTok behavior); --no-legacy uses seed + rank')
+    parser.add_argument('--device', default='cuda')
+    parser.add_argument('--no-ema', action='store_true')
+    parser.add_argument('--fid-reference')
+    return parser
+
+
+def main():
+    parser = build_parser()
+    args = parser.parse_args()
+    if args.num_samples < 1 or args.batch_size < 1:
+        parser.error('sample count and batch size must be positive')
+    world = int(os.environ.get('WORLD_SIZE', 1)); rank = int(os.environ.get('RANK', 0))
+    if world > 1:
+        if args.device != 'cuda': parser.error('distributed evaluation requires CUDA')
+        torch.cuda.set_device(int(os.environ['LOCAL_RANK']))
+        # Ranks exchange files, not CUDA tensors; CPU barriers avoid requiring GPU P2P.
+        torch.distributed.init_process_group('gloo')
+    c = OmegaConf.load(args.config)
+    torch.backends.cuda.matmul.allow_tf32 = bool(c.use_tf32)
+    torch.backends.cudnn.allow_tf32 = bool(c.use_tf32)
+    seed_everything(args.seed)
+    enc, dec, ar = load_models(c, args.checkpoint, args.device, not args.no_ema,
+                               args.mode == 'generate')
+    diti = scheduler(c)
+    # Strided partition, without DistributedSampler padding: exactly N images globally.
+    indices = list(range(rank, args.num_samples, world))
+    loader = None
+    if args.mode == 'reconstruct':
+        if not args.data: parser.error('--data is required for reconstruction')
+        dataset = ImageFolderDataset(path=args.data, resolution=c.image_size,
+                                     use_label=False, xflip=False, crop_type=c.eval_crop_type)
+        if args.num_samples > len(dataset): parser.error('num-samples exceeds dataset size')
+        loader = iter(DataLoader(Subset(dataset, indices), batch_size=args.batch_size, shuffle=False))
+    # A fixed evaluation seed is explicit; historical training-loop RNG was not checkpointed.
+    seed_everything(args.seed + (0 if args.legacy else rank))
+    samples = []
+    with torch.inference_mode():
+        for start in tqdm(range(0, len(indices), args.batch_size), desc=args.mode, disable=rank != 0):
+            n = min(args.batch_size, len(indices)-start)
+            if loader is None:
+                x = torch.empty(n, c.x_len, c.patch_dim, device=args.device)
+            else:
+                images, _ = next(loader)
+                x = patchify(img_unint8_to_norm(images.to(args.device)), c.patch_size)
+            labels = labels_for(c, n, args.device)
+            images = decode_batch(c, enc, dec, ar, x, labels, args.mode, diti)
+            samples.append(images.permute(0,2,3,1).cpu().numpy())
+    arr = np.concatenate(samples) if samples else np.empty((0,c.image_size,c.image_size,3),dtype=np.uint8)
+    out = Path(args.output); out.mkdir(parents=True, exist_ok=True)
+    if world > 1:
+        np.savez(out/f'rank-{rank}.npz', images=arr, indices=np.asarray(indices))
+        torch.distributed.barrier()
+    if rank == 0:
+        if world > 1:
+            arr = np.empty((args.num_samples,c.image_size,c.image_size,3),dtype=np.uint8)
+            for r in range(world):
+                with np.load(out/f'rank-{r}.npz') as shard: arr[shard['indices']]=shard['images']
+        assert len(arr) == args.num_samples
+        path = out/'samples.npz'; np.savez_compressed(path, arr)
+        save_image(torch.from_numpy(arr[:64]).permute(0,3,1,2).float()/255, out/'preview.png',nrow=8)
+        report = dict(mode=args.mode, num_samples=len(arr),
+                      unique_images=len({hashlib.sha256(im.tobytes()).digest() for im in arr}), seed=args.seed,
+                      world_size=world, batch_size=args.batch_size, legacy=args.legacy, rank_seed_mode='shared' if args.legacy else 'independent', config=args.config,
+                      checkpoint=args.checkpoint, ar_ema=(not args.no_ema) if args.mode == 'generate' else None,
+                      decoder='euler25' if c.use_noise_query and c.diffusion_decoder else 'single_step',
+                      precision='32-true', fid=None, torch_version=torch.__version__,
+                      config_sha256=hashlib.sha256(Path(args.config).read_bytes()).hexdigest())
+        if args.fid_reference:
+            report['reference_sha256']=hashlib.sha256(Path(args.fid_reference).read_bytes()).hexdigest()
+            result = subprocess.run([sys.executable, str(Path(__file__).with_name('calc_fid.py')),
+                '--ref_batch', args.fid_reference, '--sample_batch', str(path), '--batch_size','64'],
+                text=True, check=True, stdout=subprocess.PIPE)
+            print(result.stdout)
+            report['fid']=float(next(line.split(':',1)[1] for line in result.stdout.splitlines() if line.startswith('FID_RESULT:')))
+        (out/'metrics.json').write_text(json.dumps(report,indent=2)+'\n')
+        print(json.dumps(report, indent=2))
+    if world > 1:
+        torch.distributed.barrier(); torch.distributed.destroy_process_group()
+
+
+if __name__ == '__main__': main()

evaluator.py

@@ -0,0 +1,675 @@
+import argparse
+import io
+import os
+import random
+import warnings
+import zipfile
+from abc import ABC, abstractmethod
+from contextlib import contextmanager
+from functools import partial
+from multiprocessing import cpu_count
+from multiprocessing.pool import ThreadPool
+from typing import Iterable, Optional, Tuple, Union
+
+import numpy as np
+import requests
+import tensorflow.compat.v1 as tf
+from scipy import linalg
+from tqdm.auto import tqdm
+
+INCEPTION_V3_URL = "https://openaipublic.blob.core.windows.net/diffusion/jul-2021/ref_batches/classify_image_graph_def.pb"
+INCEPTION_V3_PATH = "classify_image_graph_def.pb"
+
+FID_POOL_NAME = "pool_3:0"
+FID_SPATIAL_NAME = "mixed_6/conv:0"
+
+
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("ref_batch", help="path to reference batch npz file")
+    parser.add_argument("sample_batch", help="path to sample batch npz file")
+    args = parser.parse_args()
+
+    config = tf.ConfigProto(
+        allow_soft_placement=True  # allows DecodeJpeg to run on CPU in Inception graph
+    )
+    config.gpu_options.allow_growth = True
+    evaluator = Evaluator(tf.Session(config=config))
+
+    print("warming up TensorFlow...")
+    # This will cause TF to print a bunch of verbose stuff now rather
+    # than after the next print(), to help prevent confusion.
+    evaluator.warmup()
+
+    print("computing reference batch activations...")
+    ref_acts = evaluator.read_activations(args.ref_batch)
+    print("computing/reading reference batch statistics...")
+    ref_stats, ref_stats_spatial = evaluator.read_statistics(args.ref_batch, ref_acts)
+
+    print("computing sample batch activations...")
+    sample_acts = evaluator.read_activations(args.sample_batch)
+    print("computing/reading sample batch statistics...")
+    sample_stats, sample_stats_spatial = evaluator.read_statistics(args.sample_batch, sample_acts)
+
+    print("Computing evaluations...")
+    print("Inception Score:", evaluator.compute_inception_score(sample_acts[0]))
+    print("FID:", sample_stats.frechet_distance(ref_stats))
+    print("sFID:", sample_stats_spatial.frechet_distance(ref_stats_spatial))
+    prec, recall = evaluator.compute_prec_recall(ref_acts[0], sample_acts[0])
+    print("Precision:", prec)
+    print("Recall:", recall)
+
+
+class InvalidFIDException(Exception):
+    pass
+
+
+class FIDStatistics:
+    def __init__(self, mu: np.ndarray, sigma: np.ndarray):
+        self.mu = mu
+        self.sigma = sigma
+
+    def frechet_distance(self, other, eps=1e-6):
+        """
+        Compute the Frechet distance between two sets of statistics.
+        """
+        # https://github.com/bioinf-jku/TTUR/blob/73ab375cdf952a12686d9aa7978567771084da42/fid.py#L132
+        mu1, sigma1 = self.mu, self.sigma
+        mu2, sigma2 = other.mu, other.sigma
+
+        mu1 = np.atleast_1d(mu1)
+        mu2 = np.atleast_1d(mu2)
+
+        sigma1 = np.atleast_2d(sigma1)
+        sigma2 = np.atleast_2d(sigma2)
+
+        assert (
+            mu1.shape == mu2.shape
+        ), f"Training and test mean vectors have different lengths: {mu1.shape}, {mu2.shape}"
+        assert (
+            sigma1.shape == sigma2.shape
+        ), f"Training and test covariances have different dimensions: {sigma1.shape}, {sigma2.shape}"
+
+        diff = mu1 - mu2
+
+        # product might be almost singular
+        covmean, _ = linalg.sqrtm(sigma1.dot(sigma2), disp=False)
+        if not np.isfinite(covmean).all():
+            msg = (
+                "fid calculation produces singular product; adding %s to diagonal of cov estimates"
+                % eps
+            )
+            warnings.warn(msg)
+            offset = np.eye(sigma1.shape[0]) * eps
+            covmean = linalg.sqrtm((sigma1 + offset).dot(sigma2 + offset))
+
+        # 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)
+
+        return diff.dot(diff) + np.trace(sigma1) + np.trace(sigma2) - 2 * tr_covmean
+
+
+class Evaluator:
+    def __init__(
+        self,
+        session,
+        batch_size=64,
+        softmax_batch_size=512,
+    ):
+        self.sess = session
+        self.batch_size = batch_size
+        self.softmax_batch_size = softmax_batch_size
+        self.manifold_estimator = ManifoldEstimator(session)
+        with self.sess.graph.as_default():
+            self.image_input = tf.placeholder(tf.float32, shape=[None, None, None, 3])
+            self.softmax_input = tf.placeholder(tf.float32, shape=[None, 2048])
+            self.pool_features, self.spatial_features = _create_feature_graph(self.image_input)
+            self.softmax = _create_softmax_graph(self.softmax_input)
+
+    def warmup(self):
+        self.compute_activations(np.zeros([1, 8, 64, 64, 3]))
+
+    def read_activations(self, npz_path: Union[str, np.ndarray]) -> Tuple[np.ndarray, np.ndarray]:
+        if isinstance(npz_path, str):
+            # Archives contain either images (arr_0) or precomputed FID statistics.
+            with np.load(npz_path) as obj:
+                keys = set(obj.files)
+
+            # Statistics-only archives do not require feature extraction.
+            if "mu" in keys and "sigma" in keys and "arr_0" not in keys:
+                dummy = np.zeros((0, 1), dtype=np.float32)
+                return dummy, dummy
+
+            # Image archives store the image batch in arr_0.
+            with open_npz_array(npz_path, "arr_0") as reader:
+                return self.compute_activations(reader.read_batches(self.batch_size))
+        elif isinstance(npz_path, np.ndarray):
+            # If npz_path is a numpy array, split it into batches manually
+            batches = np.array_split(npz_path, range(self.batch_size, npz_path.shape[0], self.batch_size))
+            return self.compute_activations(batches)
+        else:
+            raise ValueError("npz_path must be either a file path (str) or a numpy array (np.ndarray)")
+
+
+    def compute_activations(self, batches: Iterable[np.ndarray]) -> Tuple[np.ndarray, np.ndarray]:
+        """
+        Compute image features for downstream evals.
+
+        :param batches: a iterator over NHWC numpy arrays in [0, 255].
+        :return: a tuple of numpy arrays of shape [N x X], where X is a feature
+                 dimension. The tuple is (pool_3, spatial).
+        """
+        preds = []
+        spatial_preds = []
+        for batch in tqdm(batches):
+            batch = batch.astype(np.float32)
+            pred, spatial_pred = self.sess.run(
+                [self.pool_features, self.spatial_features], {self.image_input: batch}
+            )
+            preds.append(pred.reshape([pred.shape[0], -1]))
+            spatial_preds.append(spatial_pred.reshape([spatial_pred.shape[0], -1]))
+        return (
+            np.concatenate(preds, axis=0),
+            np.concatenate(spatial_preds, axis=0),
+        )
+
+    def read_statistics(
+        self, npz_path: Union[str, np.ndarray], activations: Tuple[np.ndarray, np.ndarray]
+    ) -> Tuple[FIDStatistics, FIDStatistics]:
+        if isinstance(npz_path, str):
+            obj = np.load(npz_path)
+            if "mu" in list(obj.keys()):
+                return FIDStatistics(obj["mu"], obj["sigma"]), FIDStatistics(
+                    obj["mu_s"], obj["sigma_s"]
+                )
+        elif isinstance(npz_path, np.ndarray):
+            obj = npz_path
+        else:
+            raise ValueError("npz_path must be either a file path (str) or a numpy array (np.ndarray)")
+        return tuple(self.compute_statistics(x) for x in activations)
+
+    def compute_statistics(self, activations: np.ndarray) -> FIDStatistics:
+        mu = np.mean(activations, axis=0)
+        sigma = np.cov(activations, rowvar=False)
+        return FIDStatistics(mu, sigma)
+
+    def compute_inception_score(self, activations: np.ndarray, split_size: int = 5000) -> float:
+        softmax_out = []
+        for i in range(0, len(activations), self.softmax_batch_size):
+            acts = activations[i : i + self.softmax_batch_size]
+            softmax_out.append(self.sess.run(self.softmax, feed_dict={self.softmax_input: acts}))
+        preds = np.concatenate(softmax_out, axis=0)
+        # https://github.com/openai/improved-gan/blob/4f5d1ec5c16a7eceb206f42bfc652693601e1d5c/inception_score/model.py#L46
+        scores = []
+        for i in range(0, len(preds), split_size):
+            part = preds[i : i + split_size]
+            kl = part * (np.log(part) - np.log(np.expand_dims(np.mean(part, 0), 0)))
+            kl = np.mean(np.sum(kl, 1))
+            scores.append(np.exp(kl))
+        return float(np.mean(scores))
+
+    def compute_prec_recall(
+        self, activations_ref: np.ndarray, activations_sample: np.ndarray
+    ) -> Tuple[float, float]:
+        radii_1 = self.manifold_estimator.manifold_radii(activations_ref)
+        radii_2 = self.manifold_estimator.manifold_radii(activations_sample)
+        pr = self.manifold_estimator.evaluate_pr(
+            activations_ref, radii_1, activations_sample, radii_2
+        )
+        return (float(pr[0][0]), float(pr[1][0]))
+
+
+class ManifoldEstimator:
+    """
+    A helper for comparing manifolds of feature vectors.
+
+    Adapted from https://github.com/kynkaat/improved-precision-and-recall-metric/blob/f60f25e5ad933a79135c783fcda53de30f42c9b9/precision_recall.py#L57
+    """
+
+    def __init__(
+        self,
+        session,
+        row_batch_size=10000,
+        col_batch_size=10000,
+        nhood_sizes=(3,),
+        clamp_to_percentile=None,
+        eps=1e-5,
+    ):
+        """
+        Estimate the manifold of given feature vectors.
+
+        :param session: the TensorFlow session.
+        :param row_batch_size: row batch size to compute pairwise distances
+                               (parameter to trade-off between memory usage and performance).
+        :param col_batch_size: column batch size to compute pairwise distances.
+        :param nhood_sizes: number of neighbors used to estimate the manifold.
+        :param clamp_to_percentile: prune hyperspheres that have radius larger than
+                                    the given percentile.
+        :param eps: small number for numerical stability.
+        """
+        self.distance_block = DistanceBlock(session)
+        self.row_batch_size = row_batch_size
+        self.col_batch_size = col_batch_size
+        self.nhood_sizes = nhood_sizes
+        self.num_nhoods = len(nhood_sizes)
+        self.clamp_to_percentile = clamp_to_percentile
+        self.eps = eps
+
+    def warmup(self):
+        feats, radii = (
+            np.zeros([1, 2048], dtype=np.float32),
+            np.zeros([1, 1], dtype=np.float32),
+        )
+        self.evaluate_pr(feats, radii, feats, radii)
+
+    def manifold_radii(self, features: np.ndarray) -> np.ndarray:
+        num_images = len(features)
+
+        # Estimate manifold of features by calculating distances to k-NN of each sample.
+        radii = np.zeros([num_images, self.num_nhoods], dtype=np.float32)
+        distance_batch = np.zeros([self.row_batch_size, num_images], dtype=np.float32)
+        seq = np.arange(max(self.nhood_sizes) + 1, dtype=np.int32)
+
+        for begin1 in range(0, num_images, self.row_batch_size):
+            end1 = min(begin1 + self.row_batch_size, num_images)
+            row_batch = features[begin1:end1]
+
+            for begin2 in range(0, num_images, self.col_batch_size):
+                end2 = min(begin2 + self.col_batch_size, num_images)
+                col_batch = features[begin2:end2]
+
+                # Compute distances between batches.
+                distance_batch[
+                    0 : end1 - begin1, begin2:end2
+                ] = self.distance_block.pairwise_distances(row_batch, col_batch)
+
+            # Find the k-nearest neighbor from the current batch.
+            radii[begin1:end1, :] = np.concatenate(
+                [
+                    x[:, self.nhood_sizes]
+                    for x in _numpy_partition(distance_batch[0 : end1 - begin1, :], seq, axis=1)
+                ],
+                axis=0,
+            )
+
+        if self.clamp_to_percentile is not None:
+            max_distances = np.percentile(radii, self.clamp_to_percentile, axis=0)
+            radii[radii > max_distances] = 0
+        return radii
+
+    def evaluate(self, features: np.ndarray, radii: np.ndarray, eval_features: np.ndarray):
+        """
+        Evaluate if new feature vectors are at the manifold.
+        """
+        num_eval_images = eval_features.shape[0]
+        num_ref_images = radii.shape[0]
+        distance_batch = np.zeros([self.row_batch_size, num_ref_images], dtype=np.float32)
+        batch_predictions = np.zeros([num_eval_images, self.num_nhoods], dtype=np.int32)
+        max_realism_score = np.zeros([num_eval_images], dtype=np.float32)
+        nearest_indices = np.zeros([num_eval_images], dtype=np.int32)
+
+        for begin1 in range(0, num_eval_images, self.row_batch_size):
+            end1 = min(begin1 + self.row_batch_size, num_eval_images)
+            feature_batch = eval_features[begin1:end1]
+
+            for begin2 in range(0, num_ref_images, self.col_batch_size):
+                end2 = min(begin2 + self.col_batch_size, num_ref_images)
+                ref_batch = features[begin2:end2]
+
+                distance_batch[
+                    0 : end1 - begin1, begin2:end2
+                ] = self.distance_block.pairwise_distances(feature_batch, ref_batch)
+
+            # From the minibatch of new feature vectors, determine if they are in the estimated manifold.
+            # If a feature vector is inside a hypersphere of some reference sample, then
+            # the new sample lies at the estimated manifold.
+            # The radii of the hyperspheres are determined from distances of neighborhood size k.
+            samples_in_manifold = distance_batch[0 : end1 - begin1, :, None] <= radii
+            batch_predictions[begin1:end1] = np.any(samples_in_manifold, axis=1).astype(np.int32)
+
+            max_realism_score[begin1:end1] = np.max(
+                radii[:, 0] / (distance_batch[0 : end1 - begin1, :] + self.eps), axis=1
+            )
+            nearest_indices[begin1:end1] = np.argmin(distance_batch[0 : end1 - begin1, :], axis=1)
+
+        return {
+            "fraction": float(np.mean(batch_predictions)),
+            "batch_predictions": batch_predictions,
+            "max_realisim_score": max_realism_score,
+            "nearest_indices": nearest_indices,
+        }
+
+    def evaluate_pr(
+        self,
+        features_1: np.ndarray,
+        radii_1: np.ndarray,
+        features_2: np.ndarray,
+        radii_2: np.ndarray,
+    ) -> Tuple[np.ndarray, np.ndarray]:
+        """
+        Evaluate precision and recall efficiently.
+
+        :param features_1: [N1 x D] feature vectors for reference batch.
+        :param radii_1: [N1 x K1] radii for reference vectors.
+        :param features_2: [N2 x D] feature vectors for the other batch.
+        :param radii_2: [N x K2] radii for other vectors.
+        :return: a tuple of arrays for (precision, recall):
+                 - precision: an np.ndarray of length K1
+                 - recall: an np.ndarray of length K2
+        """
+        features_1_status = np.zeros([len(features_1), radii_2.shape[1]], dtype=np.bool)
+        features_2_status = np.zeros([len(features_2), radii_1.shape[1]], dtype=np.bool)
+        for begin_1 in range(0, len(features_1), self.row_batch_size):
+            end_1 = begin_1 + self.row_batch_size
+            batch_1 = features_1[begin_1:end_1]
+            for begin_2 in range(0, len(features_2), self.col_batch_size):
+                end_2 = begin_2 + self.col_batch_size
+                batch_2 = features_2[begin_2:end_2]
+                batch_1_in, batch_2_in = self.distance_block.less_thans(
+                    batch_1, radii_1[begin_1:end_1], batch_2, radii_2[begin_2:end_2]
+                )
+                features_1_status[begin_1:end_1] |= batch_1_in
+                features_2_status[begin_2:end_2] |= batch_2_in
+        return (
+            np.mean(features_2_status.astype(np.float64), axis=0),
+            np.mean(features_1_status.astype(np.float64), axis=0),
+        )
+
+
+class DistanceBlock:
+    """
+    Calculate pairwise distances between vectors.
+
+    Adapted from https://github.com/kynkaat/improved-precision-and-recall-metric/blob/f60f25e5ad933a79135c783fcda53de30f42c9b9/precision_recall.py#L34
+    """
+
+    def __init__(self, session):
+        self.session = session
+
+        # Initialize TF graph to calculate pairwise distances.
+        with session.graph.as_default():
+            self._features_batch1 = tf.placeholder(tf.float32, shape=[None, None])
+            self._features_batch2 = tf.placeholder(tf.float32, shape=[None, None])
+            distance_block_16 = _batch_pairwise_distances(
+                tf.cast(self._features_batch1, tf.float16),
+                tf.cast(self._features_batch2, tf.float16),
+            )
+            self.distance_block = tf.cond(
+                tf.reduce_all(tf.math.is_finite(distance_block_16)),
+                lambda: tf.cast(distance_block_16, tf.float32),
+                lambda: _batch_pairwise_distances(self._features_batch1, self._features_batch2),
+            )
+
+            # Extra logic for less thans.
+            self._radii1 = tf.placeholder(tf.float32, shape=[None, None])
+            self._radii2 = tf.placeholder(tf.float32, shape=[None, None])
+            dist32 = tf.cast(self.distance_block, tf.float32)[..., None]
+            self._batch_1_in = tf.math.reduce_any(dist32 <= self._radii2, axis=1)
+            self._batch_2_in = tf.math.reduce_any(dist32 <= self._radii1[:, None], axis=0)
+
+    def pairwise_distances(self, U, V):
+        """
+        Evaluate pairwise distances between two batches of feature vectors.
+        """
+        return self.session.run(
+            self.distance_block,
+            feed_dict={self._features_batch1: U, self._features_batch2: V},
+        )
+
+    def less_thans(self, batch_1, radii_1, batch_2, radii_2):
+        return self.session.run(
+            [self._batch_1_in, self._batch_2_in],
+            feed_dict={
+                self._features_batch1: batch_1,
+                self._features_batch2: batch_2,
+                self._radii1: radii_1,
+                self._radii2: radii_2,
+            },
+        )
+
+
+def _batch_pairwise_distances(U, V):
+    """
+    Compute pairwise distances between two batches of feature vectors.
+    """
+    with tf.variable_scope("pairwise_dist_block"):
+        # Squared norms of each row in U and V.
+        norm_u = tf.reduce_sum(tf.square(U), 1)
+        norm_v = tf.reduce_sum(tf.square(V), 1)
+
+        # norm_u as a column and norm_v as a row vectors.
+        norm_u = tf.reshape(norm_u, [-1, 1])
+        norm_v = tf.reshape(norm_v, [1, -1])
+
+        # Pairwise squared Euclidean distances.
+        D = tf.maximum(norm_u - 2 * tf.matmul(U, V, False, True) + norm_v, 0.0)
+
+    return D
+
+
+class NpzArrayReader(ABC):
+    @abstractmethod
+    def read_batch(self, batch_size: int) -> Optional[np.ndarray]:
+        pass
+
+    @abstractmethod
+    def remaining(self) -> int:
+        pass
+
+    def read_batches(self, batch_size: int) -> Iterable[np.ndarray]:
+        def gen_fn():
+            while True:
+                batch = self.read_batch(batch_size)
+                if batch is None:
+                    break
+                yield batch
+
+        rem = self.remaining()
+        num_batches = rem // batch_size + int(rem % batch_size != 0)
+        return BatchIterator(gen_fn, num_batches)
+
+
+class BatchIterator:
+    def __init__(self, gen_fn, length):
+        self.gen_fn = gen_fn
+        self.length = length
+
+    def __len__(self):
+        return self.length
+
+    def __iter__(self):
+        return self.gen_fn()
+
+
+class StreamingNpzArrayReader(NpzArrayReader):
+    def __init__(self, arr_f, shape, dtype):
+        self.arr_f = arr_f
+        self.shape = shape
+        self.dtype = dtype
+        self.idx = 0
+
+    def read_batch(self, batch_size: int) -> Optional[np.ndarray]:
+        if self.idx >= self.shape[0]:
+            return None
+
+        bs = min(batch_size, self.shape[0] - self.idx)
+        self.idx += bs
+
+        if self.dtype.itemsize == 0:
+            return np.ndarray([bs, *self.shape[1:]], dtype=self.dtype)
+
+        read_count = bs * np.prod(self.shape[1:])
+        read_size = int(read_count * self.dtype.itemsize)
+        data = _read_bytes(self.arr_f, read_size, "array data")
+        return np.frombuffer(data, dtype=self.dtype).reshape([bs, *self.shape[1:]])
+
+    def remaining(self) -> int:
+        return max(0, self.shape[0] - self.idx)
+
+
+class MemoryNpzArrayReader(NpzArrayReader):
+    def __init__(self, arr):
+        self.arr = arr
+        self.idx = 0
+
+    @classmethod
+    def load(cls, path: str, arr_name: str):
+        with open(path, "rb") as f:
+            arr = np.load(f)[arr_name]
+        return cls(arr)
+
+    def read_batch(self, batch_size: int) -> Optional[np.ndarray]:
+        if self.idx >= self.arr.shape[0]:
+            return None
+
+        res = self.arr[self.idx : self.idx + batch_size]
+        self.idx += batch_size
+        return res
+
+    def remaining(self) -> int:
+        return max(0, self.arr.shape[0] - self.idx)
+
+
+@contextmanager
+def open_npz_array(path: str, arr_name: str) -> NpzArrayReader:
+    with _open_npy_file(path, arr_name) as arr_f:
+        version = np.lib.format.read_magic(arr_f)
+        if version == (1, 0):
+            header = np.lib.format.read_array_header_1_0(arr_f)
+        elif version == (2, 0):
+            header = np.lib.format.read_array_header_2_0(arr_f)
+        else:
+            yield MemoryNpzArrayReader.load(path, arr_name)
+            return
+        shape, fortran, dtype = header
+        if fortran or dtype.hasobject:
+            yield MemoryNpzArrayReader.load(path, arr_name)
+        else:
+            yield StreamingNpzArrayReader(arr_f, shape, dtype)
+
+
+def _read_bytes(fp, size, error_template="ran out of data"):
+    """
+    Copied from: https://github.com/numpy/numpy/blob/fb215c76967739268de71aa4bda55dd1b062bc2e/numpy/lib/format.py#L788-L886
+
+    Read from file-like object until size bytes are read.
+    Raises ValueError if not EOF is encountered before size bytes are read.
+    Non-blocking objects only supported if they derive from io objects.
+    Required as e.g. ZipExtFile in python 2.6 can return less data than
+    requested.
+    """
+    data = bytes()
+    while True:
+        # io files (default in python3) return None or raise on
+        # would-block, python2 file will truncate, probably nothing can be
+        # done about that.  note that regular files can't be non-blocking
+        try:
+            r = fp.read(size - len(data))
+            data += r
+            if len(r) == 0 or len(data) == size:
+                break
+        except io.BlockingIOError:
+            pass
+    if len(data) != size:
+        msg = "EOF: reading %s, expected %d bytes got %d"
+        raise ValueError(msg % (error_template, size, len(data)))
+    else:
+        return data
+
+
+@contextmanager
+def _open_npy_file(path: str, arr_name: str):
+    with open(path, "rb") as f:
+        with zipfile.ZipFile(f, "r") as zip_f:
+            if f"{arr_name}.npy" not in zip_f.namelist():
+                raise ValueError(f"missing {arr_name} in npz file")
+            with zip_f.open(f"{arr_name}.npy", "r") as arr_f:
+                yield arr_f
+
+
+def _download_inception_model():
+    if os.path.exists(INCEPTION_V3_PATH):
+        return
+    print("downloading InceptionV3 model...")
+    with requests.get(INCEPTION_V3_URL, stream=True) as r:
+        r.raise_for_status()
+        tmp_path = INCEPTION_V3_PATH + ".tmp"
+        with open(tmp_path, "wb") as f:
+            for chunk in tqdm(r.iter_content(chunk_size=8192)):
+                f.write(chunk)
+        os.rename(tmp_path, INCEPTION_V3_PATH)
+
+
+def _create_feature_graph(input_batch):
+    _download_inception_model()
+    prefix = f"{random.randrange(2**32)}_{random.randrange(2**32)}"
+    with open(INCEPTION_V3_PATH, "rb") as f:
+        graph_def = tf.GraphDef()
+        graph_def.ParseFromString(f.read())
+    pool3, spatial = tf.import_graph_def(
+        graph_def,
+        input_map={f"ExpandDims:0": input_batch},
+        return_elements=[FID_POOL_NAME, FID_SPATIAL_NAME],
+        name=prefix,
+    )
+    _update_shapes(pool3)
+    spatial = spatial[..., :7]
+    return pool3, spatial
+
+
+def _create_softmax_graph(input_batch):
+    _download_inception_model()
+    prefix = f"{random.randrange(2**32)}_{random.randrange(2**32)}"
+    with open(INCEPTION_V3_PATH, "rb") as f:
+        graph_def = tf.GraphDef()
+        graph_def.ParseFromString(f.read())
+    (matmul,) = tf.import_graph_def(
+        graph_def, return_elements=[f"softmax/logits/MatMul"], name=prefix
+    )
+    w = matmul.inputs[1]
+    logits = tf.matmul(input_batch, w)
+    return tf.nn.softmax(logits)
+
+
+def _update_shapes(pool3):
+    # https://github.com/bioinf-jku/TTUR/blob/73ab375cdf952a12686d9aa7978567771084da42/fid.py#L50-L63
+    ops = pool3.graph.get_operations()
+    for op in ops:
+        for o in op.outputs:
+            shape = o.get_shape()
+            if shape._dims is not None:  # pylint: disable=protected-access
+                shape = [s for s in shape]  # TF 2.x
+                new_shape = []
+                for j, s in enumerate(shape):
+                    if s == 1 and j == 0:
+                        new_shape.append(None)
+                    else:
+                        new_shape.append(s)
+                o.__dict__["_shape_val"] = tf.TensorShape(new_shape)
+    return pool3
+
+
+def _numpy_partition(arr, kth, **kwargs):
+    num_workers = min(cpu_count(), len(arr))
+    chunk_size = len(arr) // num_workers
+    extra = len(arr) % num_workers
+
+    start_idx = 0
+    batches = []
+    for i in range(num_workers):
+        size = chunk_size + (1 if i < extra else 0)
+        batches.append(arr[start_idx : start_idx + size])
+        start_idx += size
+
+    with ThreadPool(num_workers) as pool:
+        return list(pool.map(partial(np.partition, kth=kth, **kwargs), batches))
+
+
+if __name__ == "__main__":
+    main()

gan_loss.py

@@ -0,0 +1,1229 @@
+"""
+Single-file GAN loss module with discriminators and perceptual loss.
+"""
+
+import functools
+import hashlib
+import math
+import os
+import random
+from collections import namedtuple
+from typing import List, Tuple
+
+import numpy as np
+import requests
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch.nn.utils.spectral_norm import SpectralNorm
+from torchvision import models
+from torchvision.transforms import RandomCrop
+from tqdm import tqdm
+
+###############################################################################
+# PatchGAN discriminator and ActNorm
+###############################################################################
+
+
+class NLayerDiscriminator(nn.Module):
+    """Defines a PatchGAN discriminator as in Pix2Pix.
+
+    Reference:
+        https://github.com/junyanz/pytorch-CycleGAN-and-pix2pix/blob/master/models/networks.py
+    """
+
+    def __init__(self, input_nc=3, ndf=64, n_layers=3, use_actnorm=False):
+        """Construct a PatchGAN discriminator.
+
+        Args:
+            input_nc (int): number of channels in input images.
+            ndf (int): number of filters in the last conv layer.
+            n_layers (int): number of conv layers in the discriminator.
+            use_actnorm (bool): whether to use ActNorm instead of BatchNorm.
+        """
+        super().__init__()
+        if not use_actnorm:
+            norm_layer = nn.BatchNorm2d
+        else:
+            norm_layer = ActNorm
+
+        if isinstance(norm_layer, functools.partial):
+            # no need to use bias as BatchNorm2d has affine parameters
+            use_bias = norm_layer.func != nn.BatchNorm2d
+        else:
+            use_bias = norm_layer != nn.BatchNorm2d
+
+        kw = 4
+        padw = 1
+        sequence = [
+            nn.Conv2d(input_nc, ndf, kernel_size=kw, stride=2, padding=padw),
+            nn.LeakyReLU(0.2, True),
+        ]
+        nf_mult = 1
+        nf_mult_prev = 1
+        # gradually increase the number of filters
+        for n in range(1, n_layers):
+            nf_mult_prev = nf_mult
+            nf_mult = min(2 ** n, 8)
+            sequence += [
+                nn.Conv2d(
+                    ndf * nf_mult_prev,
+                    ndf * nf_mult,
+                    kernel_size=kw,
+                    stride=2,
+                    padding=padw,
+                    bias=use_bias,
+                ),
+                norm_layer(ndf * nf_mult),
+                nn.LeakyReLU(0.2, True),
+            ]
+
+        nf_mult_prev = nf_mult
+        nf_mult = min(2 ** n_layers, 8)
+        sequence += [
+            nn.Conv2d(
+                ndf * nf_mult_prev,
+                ndf * nf_mult,
+                kernel_size=kw,
+                stride=1,
+                padding=padw,
+                bias=use_bias,
+            ),
+            norm_layer(ndf * nf_mult),
+            nn.LeakyReLU(0.2, True),
+        ]
+
+        # output 1 channel prediction map
+        sequence += [
+            nn.Conv2d(
+                ndf * nf_mult,
+                1,
+                kernel_size=kw,
+                stride=1,
+                padding=padw,
+            )
+        ]
+        self.main = nn.Sequential(*sequence)
+
+        self.apply(self._init_weights)
+
+    def _init_weights(self, module):
+        if isinstance(module, nn.Conv2d):
+            nn.init.normal_(module.weight.data, 0.0, 0.02)
+        elif isinstance(module, nn.BatchNorm2d):
+            nn.init.normal_(module.weight.data, 1.0, 0.02)
+            nn.init.constant_(module.bias.data, 0)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        """Standard forward."""
+        return self.main(x)
+
+
+class ActNorm(nn.Module):
+    def __init__(self, num_features, logdet=False, affine=True, allow_reverse_init=False):
+        assert affine
+        super().__init__()
+        self.logdet = logdet
+        self.loc = nn.Parameter(torch.zeros(1, num_features, 1, 1))
+        self.scale = nn.Parameter(torch.ones(1, num_features, 1, 1))
+        self.allow_reverse_init = allow_reverse_init
+
+        self.register_buffer("initialized", torch.tensor(0, dtype=torch.uint8))
+
+    def initialize(self, x: torch.Tensor):
+        with torch.no_grad():
+            flatten = x.permute(1, 0, 2, 3).contiguous().view(x.shape[1], -1)
+            mean = (
+                flatten.mean(1)
+                .unsqueeze(1)
+                .unsqueeze(2)
+                .unsqueeze(3)
+                .permute(1, 0, 2, 3)
+            )
+            std = (
+                flatten.std(1)
+                .unsqueeze(1)
+                .unsqueeze(2)
+                .unsqueeze(3)
+                .permute(1, 0, 2, 3)
+            )
+
+            self.loc.data.copy_(-mean)
+            self.scale.data.copy_(1 / (std + 1e-6))
+
+    def forward(self, x: torch.Tensor, reverse: bool = False):
+        if reverse:
+            return self.reverse(x)
+        if len(x.shape) == 2:
+            x = x[:, :, None, None]
+            squeeze = True
+        else:
+            squeeze = False
+
+        _, _, height, width = x.shape
+
+        if self.training and self.initialized.item() == 0:
+            self.initialize(x)
+            self.initialized.fill_(1)
+
+        h = self.scale * (x + self.loc)
+
+        if squeeze:
+            h = h.squeeze(-1).squeeze(-1)
+
+        if self.logdet:
+            log_abs = torch.log(torch.abs(self.scale))
+            logdet = height * width * torch.sum(log_abs)
+            logdet = logdet * torch.ones(x.shape[0]).to(x)
+            return h, logdet
+
+        return h
+
+    def reverse(self, y: torch.Tensor):
+        if self.training and self.initialized.item() == 0:
+            if not self.allow_reverse_init:
+                raise RuntimeError(
+                    "Initializing ActNorm in reverse direction is disabled by default. "
+                    "Use allow_reverse_init=True to enable."
+                )
+            else:
+                self.initialize(y)
+                self.initialized.fill_(1)
+
+        if len(y.shape) == 2:
+            y = y[:, :, None, None]
+            squeeze = True
+        else:
+            squeeze = False
+
+        h = y / self.scale - self.loc
+
+        if squeeze:
+            h = h.squeeze(-1).squeeze(-1)
+        return h
+
+
+###############################################################################
+# StyleGAN-like image discriminator
+###############################################################################
+
+try:
+    from kornia.filters import filter2d
+except Exception:
+    filter2d = None
+
+
+def leaky_relu(p: float = 0.2):
+    return nn.LeakyReLU(p, inplace=True)
+
+
+def exists(val):
+    return val is not None
+
+
+class Blur(nn.Module):
+    def __init__(self):
+        super().__init__()
+        f = torch.Tensor([1, 2, 1])
+        self.register_buffer("f", f)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        if filter2d is None:
+            return x
+        f = self.f
+        f = f[None, None, :] * f[None, :, None]
+        return filter2d(x, f, normalized=True)
+
+
+class DiscriminatorBlock(nn.Module):
+    def __init__(self, input_channels, filters, downsample=True):
+        super().__init__()
+        self.conv_res = nn.Conv2d(input_channels, filters, 1, stride=(2 if downsample else 1))
+
+        self.net = nn.Sequential(
+            nn.Conv2d(input_channels, filters, 3, padding=1),
+            leaky_relu(),
+            nn.Conv2d(filters, filters, 3, padding=1),
+            leaky_relu(),
+        )
+
+        self.downsample = (
+            nn.Sequential(
+                Blur(),
+                nn.Conv2d(filters, filters, 3, padding=1, stride=2),
+            )
+            if downsample
+            else None
+        )
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        res = self.conv_res(x)
+        x = self.net(x)
+        if exists(self.downsample):
+            x = self.downsample(x)
+        x = (x + res) * (1 / math.sqrt(2))
+        return x
+
+
+class StyleGANDiscriminator(nn.Module):
+    def __init__(self, input_nc=3, ndf=64, n_layers=3, channel_multiplier=1, image_size=256):
+        super().__init__()
+        channels = {
+            4: 512,
+            8: 512,
+            16: 512,
+            32: 512,
+            64: 256 * channel_multiplier,
+            128: 128 * channel_multiplier,
+            256: 64 * channel_multiplier,
+            512: 32 * channel_multiplier,
+            1024: 16 * channel_multiplier,
+        }
+
+        log_size = int(math.log(image_size, 2))
+        in_channel = channels[image_size]
+
+        blocks = [nn.Conv2d(input_nc, in_channel, 3, padding=1), leaky_relu()]
+        for i in range(log_size, 2, -1):
+            out_channel = channels[2 ** (i - 1)]
+            blocks.append(DiscriminatorBlock(in_channel, out_channel))
+            in_channel = out_channel
+        self.blocks = nn.ModuleList(blocks)
+
+        self.final_conv = nn.Sequential(
+            nn.Conv2d(in_channel, channels[4], 3, padding=1),
+            leaky_relu(),
+        )
+        self.final_linear = nn.Sequential(
+            nn.Linear(channels[4] * 4 * 4, channels[4]),
+            leaky_relu(),
+            nn.Linear(channels[4], 1),
+        )
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        for block in self.blocks:
+            x = block(x)
+        x = self.final_conv(x)
+        x = x.view(x.shape[0], -1)
+        x = self.final_linear(x)
+        return x
+
+
+###############################################################################
+# DiffAug (data augmentation for discriminator)
+###############################################################################
+
+
+class DiffAug(object):
+    def __init__(self, prob: float = 1.0, cutout: float = 0.2):
+        self.grids = {}
+        self.prob = abs(prob)
+        self.using_cutout = prob > 0
+        self.cutout = cutout
+        self.img_channels = -1
+        self.last_blur_radius = -1
+        self.last_blur_kernel_h = None
+        self.last_blur_kernel_w = None
+
+    def get_grids(self, B, x, y, dev):
+        if (B, x, y) in self.grids:
+            return self.grids[(B, x, y)]
+
+        self.grids[(B, x, y)] = ret = torch.meshgrid(
+            torch.arange(B, dtype=torch.long, device=dev),
+            torch.arange(x, dtype=torch.long, device=dev),
+            torch.arange(y, dtype=torch.long, device=dev),
+            indexing="ij",
+        )
+        return ret
+
+    def aug(self, BCHW: torch.Tensor, warmup_blur_schedule: float = 0) -> torch.Tensor:
+        # warmup blurring
+        if BCHW.dtype != torch.float32:
+            BCHW = BCHW.float()
+        if warmup_blur_schedule > 0:
+            self.img_channels = BCHW.shape[1]
+            sigma0 = (BCHW.shape[-2] * 0.5) ** 0.5
+            sigma = sigma0 * warmup_blur_schedule
+            blur_radius = math.floor(sigma * 3)  # 3-sigma is enough for Gaussian
+            if blur_radius >= 1:
+                if self.last_blur_radius != blur_radius:
+                    self.last_blur_radius = blur_radius
+                    gaussian = torch.arange(
+                        -blur_radius, blur_radius + 1, dtype=torch.float32, device=BCHW.device
+                    )
+                    gaussian = gaussian.mul_(1 / sigma).square_().neg_().exp2_()
+                    gaussian.div_(gaussian.sum())  # normalize
+                    self.last_blur_kernel_h = (
+                        gaussian.view(1, 1, 2 * blur_radius + 1, 1)
+                        .repeat(self.img_channels, 1, 1, 1)
+                        .contiguous()
+                    )
+                    self.last_blur_kernel_w = (
+                        gaussian.view(1, 1, 1, 2 * blur_radius + 1)
+                        .repeat(self.img_channels, 1, 1, 1)
+                        .contiguous()
+                    )
+
+                BCHW = F.pad(BCHW, [blur_radius, blur_radius, blur_radius, blur_radius], mode="reflect")
+                BCHW = F.conv2d(
+                    input=BCHW, weight=self.last_blur_kernel_h, bias=None, groups=self.img_channels
+                )
+                BCHW = F.conv2d(
+                    input=BCHW, weight=self.last_blur_kernel_w, bias=None, groups=self.img_channels
+                )
+
+        if self.prob < 1e-6:
+            return BCHW
+        trans, color, cut = torch.rand(3) <= self.prob
+        trans, color, cut = trans.item(), color.item(), cut.item()
+        B, dev = BCHW.shape[0], BCHW.device
+        rand01 = torch.rand(7, B, 1, 1, device=dev) if (trans or color or cut) else None
+
+        raw_h, raw_w = BCHW.shape[-2:]
+        if trans:
+            ratio = 0.125
+            delta_h = round(raw_h * ratio)
+            delta_w = round(raw_w * ratio)
+            translation_h = rand01[0].mul(delta_h + delta_h + 1).floor().long() - delta_h
+            translation_w = rand01[1].mul(delta_w + delta_w + 1).floor().long() - delta_w
+
+            grid_B, grid_h, grid_w = self.get_grids(B, raw_h, raw_w, dev)
+            grid_h = (grid_h + translation_h).add_(1).clamp_(0, raw_h + 1)
+            grid_w = (grid_w + translation_w).add_(1).clamp_(0, raw_w + 1)
+            bchw_pad = F.pad(BCHW, [1, 1, 1, 1, 0, 0, 0, 0])
+            BCHW = (
+                bchw_pad.permute(0, 2, 3, 1)
+                .contiguous()[grid_B, grid_h, grid_w]
+                .permute(0, 3, 1, 2)
+                .contiguous()
+            )
+
+        if color:
+            BCHW = BCHW.add(rand01[2].unsqueeze(-1).sub(0.5))
+            bchw_mean = BCHW.mean(dim=1, keepdim=True)
+            BCHW = BCHW.sub(bchw_mean).mul(rand01[3].unsqueeze(-1).mul(2)).add_(bchw_mean)
+            bchw_mean = BCHW.mean(dim=(1, 2, 3), keepdim=True)
+            BCHW = BCHW.sub(bchw_mean).mul(rand01[4].unsqueeze(-1).add(0.5)).add_(bchw_mean)
+
+        if self.using_cutout and cut:
+            ratio = self.cutout
+            cutout_h = round(raw_h * ratio)
+            cutout_w = round(raw_w * ratio)
+            offset_h = rand01[5].mul(raw_h + (1 - cutout_h % 2)).floor().long()
+            offset_w = rand01[6].mul(raw_w + (1 - cutout_w % 2)).floor().long()
+
+            grid_B, grid_h, grid_w = self.get_grids(B, cutout_h, cutout_w, dev)
+            grid_h = (grid_h + offset_h).sub_(cutout_h // 2).clamp(min=0, max=raw_h - 1)
+            grid_w = (grid_w + offset_w).sub_(cutout_w // 2).clamp(min=0, max=raw_w - 1)
+            mask = torch.ones(B, raw_h, raw_w, dtype=BCHW.dtype, device=dev)
+            mask[grid_B, grid_h, grid_w] = 0
+            BCHW = BCHW.mul(mask.unsqueeze(1))
+
+        return BCHW
+
+
+###############################################################################
+# LPIPS perceptual similarity loss
+###############################################################################
+
+URL_MAP = {
+    "vgg_lpips": "https://heibox.uni-heidelberg.de/f/607503859c864bc1b30b/?dl=1",
+}
+
+CKPT_MAP = {
+    "vgg_lpips": "vgg.pth",
+}
+
+MD5_MAP = {
+    "vgg_lpips": "d507d7349b931f0638a25a48a722f98a",
+}
+
+
+def download(url, local_path, chunk_size=1024):
+    os.makedirs(os.path.split(local_path)[0], exist_ok=True)
+    with requests.get(url, stream=True) as r:
+        total_size = int(r.headers.get("content-length", 0))
+        with tqdm(total=total_size, unit="B", unit_scale=True) as pbar:
+            with open(local_path, "wb") as f:
+                for data in r.iter_content(chunk_size=chunk_size):
+                    if data:
+                        f.write(data)
+                        pbar.update(chunk_size)
+
+
+def md5_hash(path):
+    with open(path, "rb") as f:
+        content = f.read()
+    return hashlib.md5(content).hexdigest()
+
+
+def get_ckpt_path(name, root, check=False):
+    assert name in URL_MAP
+    path = os.path.join(root, CKPT_MAP[name])
+    if not os.path.exists(path) or (check and not md5_hash(path) == MD5_MAP[name]):
+        print(f"Downloading {name} model from {URL_MAP[name]} to {path}")
+        download(URL_MAP[name], path)
+        md5 = md5_hash(path)
+        assert md5 == MD5_MAP[name], md5
+    return path
+
+
+class ScalingLayer(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.register_buffer("shift", torch.Tensor([-.030, -.088, -.188])[None, :, None, None])
+        self.register_buffer("scale", torch.Tensor([.458, .448, .450])[None, :, None, None])
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return (x - self.shift) / self.scale
+
+
+class NetLinLayer(nn.Module):
+    """A single linear layer which does a 1x1 conv."""
+
+    def __init__(self, chn_in, chn_out=1, use_dropout=False):
+        super().__init__()
+        layers = [nn.Dropout()] if use_dropout else []
+        layers += [nn.Conv2d(chn_in, chn_out, 1, stride=1, padding=0, bias=False)]
+        self.model = nn.Sequential(*layers)
+
+
+class VGG16Features(nn.Module):
+    def __init__(self, requires_grad=False, pretrained=True):
+        super().__init__()
+        vgg_pretrained_features = models.vgg16(pretrained=pretrained).features
+        self.slice1 = nn.Sequential()
+        self.slice2 = nn.Sequential()
+        self.slice3 = nn.Sequential()
+        self.slice4 = nn.Sequential()
+        self.slice5 = nn.Sequential()
+        for x in range(4):
+            self.slice1.add_module(str(x), vgg_pretrained_features[x])
+        for x in range(4, 9):
+            self.slice2.add_module(str(x), vgg_pretrained_features[x])
+        for x in range(9, 16):
+            self.slice3.add_module(str(x), vgg_pretrained_features[x])
+        for x in range(16, 23):
+            self.slice4.add_module(str(x), vgg_pretrained_features[x])
+        for x in range(23, 30):
+            self.slice5.add_module(str(x), vgg_pretrained_features[x])
+        if not requires_grad:
+            for param in self.parameters():
+                param.requires_grad = False
+
+    def forward(self, x: torch.Tensor):
+        h = self.slice1(x)
+        h_relu1_2 = h
+        h = self.slice2(h)
+        h_relu2_2 = h
+        h = self.slice3(h)
+        h_relu3_3 = h
+        h = self.slice4(h)
+        h_relu4_3 = h
+        h = self.slice5(h)
+        h_relu5_3 = h
+        vgg_outputs = namedtuple(
+            "VggOutputs", ["relu1_2", "relu2_2", "relu3_3", "relu4_3", "relu5_3"]
+        )
+        out = vgg_outputs(h_relu1_2, h_relu2_2, h_relu3_3, h_relu4_3, h_relu5_3)
+        return out
+
+
+def normalize_tensor(x, eps: float = 1e-10):
+    norm_factor = torch.sqrt(torch.sum(x ** 2, dim=1, keepdim=True))
+    return x / (norm_factor + eps)
+
+
+def spatial_average(x, keepdim: bool = True):
+    return x.mean([2, 3], keepdim=keepdim)
+
+
+class LPIPS(nn.Module):
+    """Learned perceptual image patch similarity."""
+
+    def __init__(self, use_dropout: bool = True):
+        super().__init__()
+        self.scaling_layer = ScalingLayer()
+        self.chns = [64, 128, 256, 512, 512]  # VGG16 feature sizes
+        self.net = VGG16Features(pretrained=True, requires_grad=False)
+        self.lin0 = NetLinLayer(self.chns[0], use_dropout=use_dropout)
+        self.lin1 = NetLinLayer(self.chns[1], use_dropout=use_dropout)
+        self.lin2 = NetLinLayer(self.chns[2], use_dropout=use_dropout)
+        self.lin3 = NetLinLayer(self.chns[3], use_dropout=use_dropout)
+        self.lin4 = NetLinLayer(self.chns[4], use_dropout=use_dropout)
+        self.load_from_pretrained()
+        for param in self.parameters():
+            param.requires_grad = False
+
+    def load_from_pretrained(self, name: str = "vgg_lpips"):
+        ckpt = get_ckpt_path(
+            name, os.path.join(os.path.dirname(os.path.abspath(__file__)), "cache")
+        )
+        self.load_state_dict(torch.load(ckpt, map_location=torch.device("cpu")), strict=False)
+        print(f"loaded pretrained LPIPS loss from {ckpt}")
+
+    def forward(self, x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
+        x_in, y_in = self.scaling_layer(x), self.scaling_layer(y)
+        outs0, outs1 = self.net(x_in), self.net(y_in)
+        feats0, feats1, diffs = {}, {}, {}
+        lins = [self.lin0, self.lin1, self.lin2, self.lin3, self.lin4]
+        for kk in range(len(self.chns)):
+            feats0[kk], feats1[kk] = normalize_tensor(outs0[kk]), normalize_tensor(outs1[kk])
+            diffs[kk] = (feats0[kk] - feats1[kk]) ** 2
+
+        res = [spatial_average(lins[kk].model(diffs[kk]), keepdim=True) for kk in range(len(self.chns))]
+        val = res[0]
+        for l in range(1, len(self.chns)):
+            val += res[l]
+        return val
+
+
+###############################################################################
+# DINO-based discriminator (DINODiscriminator)
+###############################################################################
+
+
+class MLPNoDrop(nn.Module):
+    def __init__(self, in_features, hidden_features=None, out_features=None, fused_if_available=True):
+        super().__init__()
+        try:
+            from flash_attn.ops.fused_dense import fused_mlp_func as _fused_mlp_func
+        except Exception:
+            _fused_mlp_func = None
+
+        self.fused_mlp_func = (
+            _fused_mlp_func if (torch.cuda.is_available() and fused_if_available) else None
+        )
+        out_features = out_features or in_features
+        hidden_features = hidden_features or in_features
+        self.fc1 = nn.Linear(in_features, hidden_features)
+        self.act = nn.GELU(approximate="tanh")
+        self.fc2 = nn.Linear(hidden_features, out_features)
+
+    def forward(self, x):
+        if self.fused_mlp_func is not None:
+            return self.fused_mlp_func(
+                x=x,
+                weight1=self.fc1.weight,
+                weight2=self.fc2.weight,
+                bias1=self.fc1.bias,
+                bias2=self.fc2.bias,
+                activation="gelu_approx",
+                save_pre_act=self.training,
+                return_residual=False,
+                checkpoint_lvl=0,
+                heuristic=0,
+                process_group=None,
+            )
+        else:
+            return self.fc2(self.act(self.fc1(x)))
+
+
+class SelfAttentionNoDrop(nn.Module):
+    def __init__(self, block_idx, embed_dim=768, num_heads=12, flash_if_available=True):
+        super().__init__()
+        assert embed_dim % num_heads == 0
+        self.block_idx = block_idx
+        self.num_heads = num_heads
+        self.head_dim = embed_dim // num_heads
+        self.scale = 1 / math.sqrt(self.head_dim)
+        self.qkv = nn.Linear(embed_dim, embed_dim * 3, bias=True)
+        self.proj = nn.Linear(embed_dim, embed_dim, bias=True)
+
+        try:
+            from flash_attn import flash_attn_qkvpacked_func as _flash_attn_qkvpacked_func
+        except Exception:
+            _flash_attn_qkvpacked_func = None
+
+        self.using_flash_attn = (
+            torch.cuda.is_available() and flash_if_available and _flash_attn_qkvpacked_func is not None
+        )
+        self._flash_attn_qkvpacked_func = _flash_attn_qkvpacked_func
+
+    def forward(self, x):
+        B, L, C = x.shape
+        qkv = self.qkv(x).view(B, L, 3, self.num_heads, self.head_dim)
+        if self.using_flash_attn and qkv.dtype != torch.float32:
+            oup = self._flash_attn_qkvpacked_func(qkv, softmax_scale=self.scale).view(B, L, C)
+        else:
+            q, k, v = qkv.permute(2, 0, 3, 1, 4).unbind(dim=0)  # BHLc
+            scale = self.scale
+            attn = q.mul(scale) @ k.transpose(-2, -1)
+            attn = attn.softmax(dim=-1)
+            oup = (attn @ v).transpose(1, 2).reshape(B, L, C)
+        return self.proj(oup)
+
+
+class SABlockNoDrop(nn.Module):
+    def __init__(self, block_idx, embed_dim, num_heads, mlp_ratio, norm_eps):
+        super().__init__()
+        self.norm1 = nn.LayerNorm(embed_dim, eps=norm_eps)
+        self.attn = SelfAttentionNoDrop(
+            block_idx=block_idx,
+            embed_dim=embed_dim,
+            num_heads=num_heads,
+            flash_if_available=True,
+        )
+        self.norm2 = nn.LayerNorm(embed_dim, eps=norm_eps)
+        self.mlp = MLPNoDrop(
+            in_features=embed_dim,
+            hidden_features=round(embed_dim * mlp_ratio),
+            fused_if_available=True,
+        )
+
+    def forward(self, x):
+        x = x + self.attn(self.norm1(x))
+        x = x + self.mlp(self.norm2(x))
+        return x
+
+
+class ResidualBlock(nn.Module):
+    def __init__(self, fn):
+        super().__init__()
+        self.fn = fn
+        self.ratio = 1 / np.sqrt(2)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return (self.fn(x).add(x)).mul_(self.ratio)
+
+
+class SpectralConv1d(nn.Conv1d):
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        SpectralNorm.apply(self, name="weight", n_power_iterations=1, dim=0, eps=1e-12)
+
+
+class BatchNormLocal(nn.Module):
+    def __init__(self, num_features: int, affine: bool = True, virtual_bs: int = 8, eps: float = 1e-6):
+        super().__init__()
+        self.virtual_bs = virtual_bs
+        self.eps = eps
+        self.affine = affine
+
+        if self.affine:
+            self.weight = nn.Parameter(torch.ones(num_features))
+            self.bias = nn.Parameter(torch.zeros(num_features))
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        shape = x.size()
+        x = x.float()
+
+        # Reshape batch into groups.
+        G = int(np.ceil(x.size(0) / self.virtual_bs))
+        x = x.view(G, -1, x.size(-2), x.size(-1))
+
+        # Calculate stats.
+        mean = x.mean([1, 3], keepdim=True)
+        var = x.var([1, 3], keepdim=True, unbiased=False)
+        x = (x - mean) / (torch.sqrt(var + self.eps))
+
+        if self.affine:
+            x = x * self.weight[None, :, None] + self.bias[None, :, None]
+
+        return x.view(shape)
+
+
+def make_block(channels: int, kernel_size: int, norm_type: str, norm_eps: float, using_spec_norm: bool):
+    if norm_type == "bn":
+        norm = BatchNormLocal(channels, eps=norm_eps)
+    elif norm_type == "sbn":
+        norm = nn.SyncBatchNorm(channels, eps=norm_eps, process_group=None)
+    elif norm_type in {"lbn", "hbn"}:
+        # Fallback to SyncBatchNorm without a custom local machine group.
+        norm = nn.SyncBatchNorm(channels, eps=norm_eps, process_group=None)
+    elif norm_type == "gn":
+        norm = nn.GroupNorm(num_groups=32, num_channels=channels, eps=norm_eps, affine=True)
+    else:
+        raise NotImplementedError
+
+    conv_cls = SpectralConv1d if using_spec_norm else nn.Conv1d
+    return nn.Sequential(
+        conv_cls(channels, channels, kernel_size=kernel_size, padding=kernel_size // 2, padding_mode="circular"),
+        norm,
+        nn.LeakyReLU(negative_slope=0.2, inplace=True),
+    )
+
+
+class PatchEmbed(nn.Module):
+    def __init__(self, img_size=224, patch_size=16, in_chans=3, embed_dim=768, norm_layer=None, flatten=True):
+        super().__init__()
+        self.img_size = img_size
+        self.patch_size = patch_size
+        self.num_patches = (img_size // patch_size) ** 2
+        self.flatten = flatten
+
+        self.proj = nn.Conv2d(in_chans, embed_dim, kernel_size=patch_size, stride=patch_size)
+        self.norm = norm_layer(embed_dim) if norm_layer else nn.Identity()
+
+    def forward(self, x):
+        x = self.proj(x).flatten(2).transpose(1, 2)  # BCHW => BCL => BLC
+        return self.norm(x)
+
+
+class FrozenDINOSmallNoDrop(nn.Module):
+    """
+    Frozen DINO ViT without any dropout or droppath layers (eval-only),
+    based on timm.create_model('vit_small_patch16_224', pretrained=False, num_classes=0).
+    """
+
+    def __init__(
+        self,
+        depth=12,
+        key_depths=(2, 5, 8, 11),
+        norm_eps=1e-6,
+        patch_size=16,
+        in_chans=3,
+        num_classes=0,
+        embed_dim=384,
+        num_heads=6,
+        mlp_ratio=4.0,
+    ):
+        super().__init__()
+        self.num_classes = num_classes
+        self.num_features = self.embed_dim = embed_dim
+
+        self.img_size = 224
+        self.patch_embed = PatchEmbed(
+            img_size=self.img_size,
+            patch_size=patch_size,
+            in_chans=in_chans,
+            embed_dim=embed_dim,
+        )
+        self.patch_size = patch_size
+        self.patch_nums = self.img_size // patch_size
+
+        # x \in [-1, 1]
+        m = torch.tensor((0.485, 0.456, 0.406))
+        s = torch.tensor((0.229, 0.224, 0.225))
+        self.register_buffer("x_scale", (0.5 / s).reshape(1, 3, 1, 1))
+        self.register_buffer("x_shift", ((0.5 - m) / s).reshape(1, 3, 1, 1))
+        self.crop = RandomCrop(self.img_size)
+
+        self.cls_token = nn.Parameter(torch.zeros(1, 1, embed_dim))
+        self.dist_token = None
+        self.pos_embed = nn.Parameter(
+            torch.zeros(1, self.patch_nums * self.patch_nums + 1, embed_dim)
+        )  # +1: for cls
+
+        self.key_depths = set(d for d in key_depths if d < depth)
+        self.blocks = nn.Sequential(
+            *[
+                SABlockNoDrop(
+                    block_idx=i,
+                    embed_dim=embed_dim,
+                    num_heads=num_heads,
+                    mlp_ratio=mlp_ratio,
+                    norm_eps=norm_eps,
+                )
+                for i in range(max(depth, 1 + max(self.key_depths)))
+            ]
+        )
+        self.norm = nn.LayerNorm(embed_dim, eps=norm_eps)
+
+        # eval mode only
+        self.eval()
+        for p in self.parameters():
+            p.requires_grad_(False)
+
+    def inter_pos_embed(self, patch_nums=(14, 14)):
+        if patch_nums[0] == self.patch_nums and patch_nums[1] == self.patch_nums:
+            return self.pos_embed
+        pe_cls, pe_grid = self.pos_embed[:, :1], self.pos_embed[0, 1:]
+        pe_grid = pe_grid.reshape(1, self.patch_nums, self.patch_nums, -1).permute(0, 3, 1, 2)
+        pe_grid = F.interpolate(
+            pe_grid, size=(patch_nums[0], patch_nums[1]), mode="bilinear", align_corners=False
+        )
+        pe_grid = pe_grid.permute(0, 2, 3, 1).reshape(1, patch_nums[0] * patch_nums[1], -1)
+        return torch.cat([pe_cls, pe_grid], dim=1)
+
+    def forward(self, x, grad_ckpt: bool = False):
+        with torch.cuda.amp.autocast(enabled=False):
+            x = (self.x_scale * x.float()).add_(self.x_shift)
+            H, W = x.shape[-2], x.shape[-1]
+            if H > self.img_size and W > self.img_size and random.random() <= 0.5:
+                x = self.crop(x)
+            else:
+                x = F.interpolate(
+                    x,
+                    size=(self.img_size, self.img_size),
+                    mode="area" if H > self.img_size else "bicubic",
+                )
+
+        x = self.patch_embed(x)
+
+        with torch.cuda.amp.autocast(enabled=False):
+            x = torch.cat((self.cls_token.expand(x.shape[0], -1, -1), x.float()), dim=1)
+            x = x + self.pos_embed
+            activations = [(x[:, 1:] + x[:, :1]).transpose_(1, 2)]
+        for i, b in enumerate(self.blocks):
+            if not grad_ckpt:
+                x = b(x)
+            else:
+                x = torch.utils.checkpoint.checkpoint(b, x, use_reentrant=False)
+            if i in self.key_depths:
+                activations.append((x[:, 1:].float() + x[:, :1].float()).transpose_(1, 2))
+        return activations
+
+
+class DinoDisc(nn.Module):
+    def __init__(
+        self,
+        dino_ckpt_path="https://dl.fbaipublicfiles.com/dino/dino_deitsmall16_pretrain/dino_deitsmall16_pretrain.pth",
+        device="cuda",
+        ks=9,
+        depth=12,
+        key_depths=(2, 5, 8, 11),
+        norm_type="bn",
+        using_spec_norm=True,
+        norm_eps=1e-6,
+    ):
+        super().__init__()
+        # load state
+        state = torch.hub.load_state_dict_from_url(dino_ckpt_path, map_location="cpu")
+        for k in sorted(state.keys()):
+            if ".attn.qkv.bias" in k:
+                bias = state[k]
+                C = bias.numel() // 3
+                bias[C : 2 * C].zero_()  # zero out k_bias
+
+        key_depths = tuple(d for d in key_depths if d < depth)
+        d = FrozenDINOSmallNoDrop(depth=depth, key_depths=key_depths, norm_eps=norm_eps)
+        missing, unexpected = d.load_state_dict(state, strict=False)
+        missing = [
+            m
+            for m in missing
+            if all(x not in m for x in {"x_scale", "x_shift"})
+        ]
+        if torch.cuda.is_available():
+            assert len(missing) == 0, f"missing keys: {missing}"
+            assert len(unexpected) == 0, f"unexpected keys: {unexpected}"
+
+        self.dino_proxy: Tuple[FrozenDINOSmallNoDrop] = (d.to(device=device),)
+        dino_C = self.dino_proxy[0].embed_dim
+
+        self.heads = nn.ModuleList(
+            [
+                nn.Sequential(
+                    make_block(
+                        dino_C,
+                        kernel_size=1,
+                        norm_type=norm_type,
+                        norm_eps=norm_eps,
+                        using_spec_norm=using_spec_norm,
+                    ),
+                    ResidualBlock(
+                        make_block(
+                            dino_C,
+                            kernel_size=ks,
+                            norm_type=norm_type,
+                            norm_eps=norm_eps,
+                            using_spec_norm=using_spec_norm,
+                        )
+                    ),
+                    (
+                        SpectralConv1d
+                        if using_spec_norm
+                        else nn.Conv1d
+                    )(dino_C, 1, kernel_size=1, padding=0),
+                )
+                for _ in range(len(key_depths) + 1)
+            ]
+        )
+
+    def reinit(
+        self,
+        dino_ckpt_path="https://dl.fbaipublicfiles.com/dino/dino_deitsmall16_pretrain/dino_deitsmall16_pretrain.pth",
+        device="cuda",
+        ks=9,
+        depth=12,
+        key_depths=(2, 5, 8, 11),
+        norm_type="bn",
+        using_spec_norm=True,
+        norm_eps=1e-6,
+    ):
+        dino_C = self.dino_proxy[0].embed_dim
+        heads = nn.ModuleList(
+            [
+                nn.Sequential(
+                    make_block(
+                        dino_C,
+                        kernel_size=1,
+                        norm_type=norm_type,
+                        norm_eps=norm_eps,
+                        using_spec_norm=using_spec_norm,
+                    ),
+                    ResidualBlock(
+                        make_block(
+                            dino_C,
+                            kernel_size=ks,
+                            norm_type=norm_type,
+                            norm_eps=norm_eps,
+                            using_spec_norm=using_spec_norm,
+                        )
+                    ),
+                    (
+                        SpectralConv1d
+                        if using_spec_norm
+                        else nn.Conv1d
+                    )(dino_C, 1, kernel_size=1, padding=0),
+                )
+                for _ in range(len(key_depths) + 1)
+            ]
+        )
+
+        self.heads.load_state_dict(heads.state_dict())
+
+    def forward(self, x_in_pm1: torch.Tensor, grad_ckpt: bool = False) -> torch.Tensor:
+        # x_in_pm1: image tensor normalized to [-1, 1]
+        dino_grad_ckpt = grad_ckpt and x_in_pm1.requires_grad
+        FrozenDINOSmallNoDrop.forward
+        activations: List[torch.Tensor] = self.dino_proxy[0](
+            x_in_pm1.float(), grad_ckpt=dino_grad_ckpt
+        )
+        B = x_in_pm1.shape[0]
+        return torch.cat(
+            [
+                (
+                    h(act)
+                    if not grad_ckpt
+                    else torch.utils.checkpoint.checkpoint(h, act, use_reentrant=False)
+                ).view(B, -1)
+                for h, act in zip(self.heads, activations)
+            ],
+            dim=1,
+        )
+
+
+###############################################################################
+# GAN loss helpers
+###############################################################################
+
+
+def hinge_d_loss(logits_real: torch.Tensor, logits_fake: torch.Tensor) -> torch.Tensor:
+    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 vanilla_d_loss(logits_real: torch.Tensor, logits_fake: torch.Tensor) -> torch.Tensor:
+    loss_real = torch.mean(F.softplus(-logits_real))
+    loss_fake = torch.mean(F.softplus(logits_fake))
+    d_loss = loss_real + loss_fake
+    return d_loss
+
+
+def non_saturating_d_loss(logits_real: torch.Tensor, logits_fake: torch.Tensor) -> torch.Tensor:
+    loss_real = torch.mean(
+        F.binary_cross_entropy_with_logits(torch.ones_like(logits_real), logits_real)
+    )
+    loss_fake = torch.mean(
+        F.binary_cross_entropy_with_logits(torch.zeros_like(logits_fake), logits_fake)
+    )
+    d_loss = 0.5 * (loss_real + loss_fake)
+    return d_loss
+
+
+def hinge_gen_loss(logit_fake: torch.Tensor) -> torch.Tensor:
+    return -torch.mean(logit_fake)
+
+def vanilla_gen_loss(logit_fake: torch.Tensor) -> torch.Tensor:
+    return torch.mean(F.softplus(-logit_fake))
+
+def non_saturating_gen_loss(logit_fake: torch.Tensor) -> torch.Tensor:
+    return torch.mean(
+        F.binary_cross_entropy_with_logits(torch.ones_like(logit_fake), logit_fake)
+    )
+
+
+def adopt_weight(weight: float, global_step: int, threshold: int = 0, value: float = 0.0):
+    if global_step < threshold:
+        weight = value
+    return weight
+
+
+class LeCAM_EMA(object):
+    def __init__(self, init: float = 0.0, decay: float = 0.999):
+        self.logits_real_ema = init
+        self.logits_fake_ema = init
+        self.decay = decay
+
+    def update(self, logits_real: torch.Tensor, logits_fake: torch.Tensor):
+        self.logits_real_ema = (
+            self.logits_real_ema * self.decay
+            + torch.mean(logits_real).item() * (1 - self.decay)
+        )
+        self.logits_fake_ema = (
+            self.logits_fake_ema * self.decay
+            + torch.mean(logits_fake).item() * (1 - self.decay)
+        )
+
+
+def lecam_reg(real_pred: torch.Tensor, fake_pred: torch.Tensor, lecam_ema: LeCAM_EMA):
+    reg = torch.mean(F.relu(real_pred - lecam_ema.logits_fake_ema).pow(2)) + torch.mean(
+        F.relu(lecam_ema.logits_real_ema - fake_pred).pow(2)
+    )
+    return reg
+
+
+###############################################################################
+# Combined GAN loss (generator + discriminator) with perceptual & VQ terms
+###############################################################################
+
+
+class GANLoss(nn.Module):
+    """
+    Combined loss used for training VQ-style autoencoders with GAN objectives.
+
+    This is functionally similar to the original `VQLoss` but packaged so it can
+    live in a single file together with its discriminators and LPIPS.
+    """
+
+    def __init__(
+        self,
+        disc_loss: str = "hinge",
+        disc_dim: int = 64,
+        disc_type: str = "patchgan",
+        image_size: int = 256,
+        disc_num_layers: int = 3,
+        disc_in_channels: int = 3,
+        disc_adaptive_weight: bool = False,
+        gen_adv_loss: str = "hinge",
+        reconstruction_loss: str = "l2",
+        reconstruction_weight: float = 1.0,
+        codebook_weight: float = 1.0,
+        perceptual_weight: float = 1.0,
+        lecam_loss_weight: float = None,
+        norm_type: str = "bn",
+        aug_prob: float = 1.0,
+    ):
+        super().__init__()
+        # discriminator type & architecture
+        assert disc_type in ["patchgan", "stylegan", "dinodisc"]
+        assert disc_loss in ["hinge", "vanilla", "non-saturating"]
+        self.disc_type = disc_type
+        if disc_type == "patchgan":
+            self.discriminator = NLayerDiscriminator(
+                input_nc=disc_in_channels,
+                n_layers=disc_num_layers,
+                ndf=disc_dim,
+            )
+        elif disc_type == "stylegan":
+            self.discriminator = StyleGANDiscriminator(
+                input_nc=disc_in_channels,
+                image_size=image_size,
+            )
+        elif disc_type == "dinodisc":
+            self.discriminator = DinoDisc(norm_type=norm_type)  # default 224 otherwise crop
+            self.daug = DiffAug(prob=aug_prob, cutout=0.2)
+        else:
+            raise ValueError(f"Unknown GAN discriminator type '{disc_type}'.")
+
+        if disc_loss == "hinge":
+            self.disc_loss = hinge_d_loss
+        elif disc_loss == "vanilla":
+            self.disc_loss = vanilla_d_loss
+        elif disc_loss == "non-saturating":
+            self.disc_loss = non_saturating_d_loss
+        else:
+            raise ValueError(f"Unknown GAN discriminator loss '{disc_loss}'.")
+
+
+        # generator adversarial loss
+        if gen_adv_loss == "hinge":
+            self.gen_adv_loss = hinge_gen_loss
+        elif gen_adv_loss == "non-saturating":
+            self.gen_adv_loss = non_saturating_gen_loss
+        elif gen_adv_loss == 'vanilla':
+            self.gen_adv_loss = vanilla_gen_loss
+        else:
+            raise ValueError(f"Unknown GAN generator loss '{gen_adv_loss}'.")
+
+        # perceptual loss
+        self.perceptual_loss = LPIPS().eval()
+        self.perceptual_weight = perceptual_weight
+
+        # reconstruction loss
+        if reconstruction_loss == "l1":
+            self.rec_loss = F.l1_loss
+        elif reconstruction_loss == "l2":
+            self.rec_loss = F.mse_loss
+        else:
+            raise ValueError(f"Unknown rec loss '{reconstruction_loss}'.")
+        self.rec_weight = reconstruction_weight
+
+        # codebook loss weight (used externally when aggregating)
+        self.codebook_weight = codebook_weight
+
+        # LeCAM regularization
+        self.lecam_loss_weight = lecam_loss_weight
+        if self.lecam_loss_weight is not None:
+            self.lecam_ema = LeCAM_EMA()
+
+    def calculate_adaptive_weight(self, nll_loss, g_loss, last_layer):
+        with fabric.no_backward_sync(last_layer):
+            nll_grads = torch.autograd.grad(nll_loss, last_layer, retain_graph=True)[0]
+            g_grads = torch.autograd.grad(g_loss, last_layer, retain_graph=True)[0]
+
+        d_weight = torch.norm(nll_grads) / (torch.norm(g_grads) + 1e-4)
+        d_weight = torch.clamp(d_weight, 0.0, 1e4).detach()
+        return d_weight.detach()
+
+    def forward(
+        self,
+        x,
+        x_hat,
+        global_step,
+        loss='G'
+    ):
+        """
+        Args are kept compatible with the original `VQLoss` for drop-in usage.
+
+        optimizer_idx = 0: generator (autoencoder) update
+        optimizer_idx = 1: discriminator update
+        """
+        # generator update
+        if loss == 'G':          
+            # discriminator loss (for generator)
+            if self.disc_type == "dinodisc":
+                if fade_blur_schedule < 1e-6:
+                    fade_blur_schedule = 0.0
+                logits_fake = self.discriminator(
+                    self.daug.aug(x_hat.contiguous(), fade_blur_schedule)
+                )
+            else:
+                logits_fake = self.discriminator(x_hat.contiguous())
+            generator_adv_loss = self.gen_adv_loss(logits_fake)
+
+            loss = generator_adv_loss
+            return (loss,{'gen_G_loss_vanilla':generator_adv_loss.item()})
+
+        # discriminator update
+        if loss == 'D':
+            if self.disc_type == "dinodisc":
+                logits_fake = self.discriminator(
+                    self.daug.aug(x_hat.contiguous().detach(), 0.)
+                )
+                logits_real = self.discriminator(
+                    self.daug.aug(x.contiguous().detach(), 0.)
+                )
+            else:
+                logits_fake = self.discriminator(x_hat.contiguous().detach())
+                logits_real = self.discriminator(x.contiguous().detach())
+
+         
+
+            if self.lecam_loss_weight is not None:
+                self.lecam_ema.update(logits_real, logits_fake)
+                lecam_loss = self.lecam_loss_weight * lecam_reg(logits_real, logits_fake, self.lecam_ema)
+                non_saturate_d = self.disc_loss(logits_real, logits_fake)
+                d_adversarial_loss = lecam_loss + non_saturate_d
+                
+            else:
+                d_adversarial_loss = self.disc_weight * self.disc_loss(logits_real, logits_fake)
+
+            logits_real_log = logits_real.detach()
+            logits_fake_log = logits_fake.detach()
+
+            return (d_adversarial_loss, { "disc_loss":non_saturate_d.item(), "lecam_loss": lecam_loss.item(), "prob_real":logits_real_log.sigmoid().mean().item(), "prob_fake":logits_fake_log.sigmoid().mean().item()})

lpips.py

@@ -0,0 +1,164 @@
+"""Stripped version of https://github.com/richzhang/PerceptualSimilarity/tree/master/models"""
+
+import os, hashlib
+import requests
+from tqdm import tqdm
+
+import torch
+import torch.nn as nn
+from torchvision import models
+from collections import namedtuple
+
+URL_MAP = {
+    "vgg_lpips": "https://heibox.uni-heidelberg.de/f/607503859c864bc1b30b/?dl=1"
+}
+
+CKPT_MAP = {
+    "vgg_lpips": "vgg.pth"
+}
+
+MD5_MAP = {
+    "vgg_lpips": "d507d7349b931f0638a25a48a722f98a"
+}
+
+def download(url, local_path, chunk_size=1024):
+    os.makedirs(os.path.split(local_path)[0], exist_ok=True)
+    with requests.get(url, stream=True) as r:
+        total_size = int(r.headers.get("content-length", 0))
+        with tqdm(total=total_size, unit="B", unit_scale=True) as pbar:
+            with open(local_path, "wb") as f:
+                for data in r.iter_content(chunk_size=chunk_size):
+                    if data:
+                        f.write(data)
+                        pbar.update(chunk_size)
+
+
+def md5_hash(path):
+    with open(path, "rb") as f:
+        content = f.read()
+    return hashlib.md5(content).hexdigest()
+
+
+def get_ckpt_path(name, root, check=False):
+    assert name in URL_MAP
+    path = os.path.join(root, CKPT_MAP[name])
+    if not os.path.exists(path) or (check and not md5_hash(path) == MD5_MAP[name]):
+        print("Downloading {} model from {} to {}".format(name, URL_MAP[name], path))
+        download(URL_MAP[name], path)
+        md5 = md5_hash(path)
+        assert md5 == MD5_MAP[name], md5
+    return path
+
+
+class LPIPS(nn.Module):
+    # Learned perceptual metric
+    def __init__(self, use_dropout=True):
+        super().__init__()
+        self.scaling_layer = ScalingLayer()
+        self.chns = [64, 128, 256, 512, 512]  # vg16 features
+        self.net = vgg16(pretrained=True, requires_grad=False)
+        self.lin0 = NetLinLayer(self.chns[0], use_dropout=use_dropout)
+        self.lin1 = NetLinLayer(self.chns[1], use_dropout=use_dropout)
+        self.lin2 = NetLinLayer(self.chns[2], use_dropout=use_dropout)
+        self.lin3 = NetLinLayer(self.chns[3], use_dropout=use_dropout)
+        self.lin4 = NetLinLayer(self.chns[4], use_dropout=use_dropout)
+        self.load_from_pretrained()
+        for param in self.parameters():
+            param.requires_grad = False
+
+    def load_from_pretrained(self, name="vgg_lpips"):
+        ckpt = get_ckpt_path(name, os.path.join(os.path.dirname(os.path.abspath(__file__)), "cache"))
+        self.load_state_dict(torch.load(ckpt, map_location=torch.device("cpu")), strict=False)
+        print("loaded pretrained LPIPS loss from {}".format(ckpt))
+
+    @classmethod
+    def from_pretrained(cls, name="vgg_lpips"):
+        if name != "vgg_lpips":
+            raise NotImplementedError
+        model = cls()
+        ckpt = get_ckpt_path(name, os.path.join(os.path.dirname(os.path.abspath(__file__)), "cache"))
+        model.load_state_dict(torch.load(ckpt, map_location=torch.device("cpu")), strict=False)
+        return model
+
+    def forward(self, input, target):
+        in0_input, in1_input = (self.scaling_layer(input), self.scaling_layer(target))
+        outs0, outs1 = self.net(in0_input), self.net(in1_input)
+        feats0, feats1, diffs = {}, {}, {}
+        lins = [self.lin0, self.lin1, self.lin2, self.lin3, self.lin4]
+        for kk in range(len(self.chns)):
+            feats0[kk], feats1[kk] = normalize_tensor(outs0[kk]), normalize_tensor(outs1[kk])
+            diffs[kk] = (feats0[kk] - feats1[kk]) ** 2
+
+        res = [spatial_average(lins[kk].model(diffs[kk]), keepdim=True) for kk in range(len(self.chns))]
+        val = res[0]
+        for l in range(1, len(self.chns)):
+            val += res[l]
+        return val
+
+
+class ScalingLayer(nn.Module):
+    def __init__(self):
+        super(ScalingLayer, self).__init__()
+        self.register_buffer('shift', torch.Tensor([-.030, -.088, -.188])[None, :, None, None])
+        self.register_buffer('scale', torch.Tensor([.458, .448, .450])[None, :, None, None])
+
+    def forward(self, inp):
+        return (inp - self.shift) / self.scale
+
+
+class NetLinLayer(nn.Module):
+    """ A single linear layer which does a 1x1 conv """
+    def __init__(self, chn_in, chn_out=1, use_dropout=False):
+        super(NetLinLayer, self).__init__()
+        layers = [nn.Dropout(), ] if (use_dropout) else []
+        layers += [nn.Conv2d(chn_in, chn_out, 1, stride=1, padding=0, bias=False), ]
+        self.model = nn.Sequential(*layers)
+
+
+class vgg16(torch.nn.Module):
+    def __init__(self, requires_grad=False, pretrained=True):
+        super(vgg16, self).__init__()
+        vgg_pretrained_features = models.vgg16(pretrained=pretrained).features
+        self.slice1 = torch.nn.Sequential()
+        self.slice2 = torch.nn.Sequential()
+        self.slice3 = torch.nn.Sequential()
+        self.slice4 = torch.nn.Sequential()
+        self.slice5 = torch.nn.Sequential()
+        self.N_slices = 5
+        for x in range(4):
+            self.slice1.add_module(str(x), vgg_pretrained_features[x])
+        for x in range(4, 9):
+            self.slice2.add_module(str(x), vgg_pretrained_features[x])
+        for x in range(9, 16):
+            self.slice3.add_module(str(x), vgg_pretrained_features[x])
+        for x in range(16, 23):
+            self.slice4.add_module(str(x), vgg_pretrained_features[x])
+        for x in range(23, 30):
+            self.slice5.add_module(str(x), vgg_pretrained_features[x])
+        if not requires_grad:
+            for param in self.parameters():
+                param.requires_grad = False
+
+    def forward(self, X):
+        h = self.slice1(X)
+        h_relu1_2 = h
+        h = self.slice2(h)
+        h_relu2_2 = h
+        h = self.slice3(h)
+        h_relu3_3 = h
+        h = self.slice4(h)
+        h_relu4_3 = h
+        h = self.slice5(h)
+        h_relu5_3 = h
+        vgg_outputs = namedtuple("VggOutputs", ['relu1_2', 'relu2_2', 'relu3_3', 'relu4_3', 'relu5_3'])
+        out = vgg_outputs(h_relu1_2, h_relu2_2, h_relu3_3, h_relu4_3, h_relu5_3)
+        return out
+
+
+def normalize_tensor(x,eps=1e-10):
+    norm_factor = torch.sqrt(torch.sum(x**2,dim=1,keepdim=True))
+    return x/(norm_factor+eps)
+
+
+def spatial_average(x, keepdim=True):
+    return x.mean([2,3],keepdim=keepdim)

models.py

@@ -0,0 +1,1013 @@
+import math
+from math import pi
+
+import torch
+import torch.nn as nn
+from torch import einsum, broadcast_tensors, Tensor
+import torch.nn.functional as F
+from torch.nn import Module
+from torch.amp import autocast
+from torch import nn
+from einops import rearrange, repeat
+import einops
+import numpy as np
+import os
+import copy
+
+# Well Inited Linear
+def weight_init(shape, mode, fan_in, fan_out):
+    if mode == 'xavier_uniform': return np.sqrt(6 / (fan_in + fan_out)) * (torch.rand(*shape) * 2 - 1)
+    if mode == 'xavier_normal':  return np.sqrt(2 / (fan_in + fan_out)) * torch.randn(*shape)
+    if mode == 'kaiming_uniform': return np.sqrt(3 / fan_in) * (torch.rand(*shape) * 2 - 1)
+    if mode == 'kaiming_normal':  return np.sqrt(1 / fan_in) * torch.randn(*shape)
+    if mode == 'uniform': return torch.rand() * 2 - 1
+    raise ValueError(f'Invalid init mode "{mode}"')
+
+class Linear(torch.nn.Module):
+    def __init__(self, in_features, out_features, bias=True, init_mode='kaiming_uniform', init_weight=1, init_bias=0):
+        super().__init__()
+        self.in_features = in_features
+        self.out_features = out_features
+        init_kwargs = dict(mode=init_mode, fan_in=in_features, fan_out=out_features)
+        self.weight = torch.nn.Parameter(weight_init([out_features, in_features], **init_kwargs) * init_weight)
+        self.bias = torch.nn.Parameter(weight_init([out_features], **init_kwargs) * init_bias) if bias else None
+
+    def forward(self, x):
+        x = x @ self.weight.to(x.dtype).t()
+        if self.bias is not None:
+            x = x.add_(self.bias.to(x.dtype))
+        return x
+
+
+class AdaRMSNorm(nn.Module):
+    def __init__(self,dim, cond_dim=None, elementwise_affine=True, cond_based_affine=True, centering=False, eps=None):
+        super(AdaRMSNorm, self).__init__()        
+        self.dim = dim
+        self.eps = eps
+        self.cond_based_affine = cond_based_affine      
+        self.elementwise_affine = elementwise_affine
+        self.centering = centering  
+        assert not(cond_dim is None and cond_based_affine), 'cond_dim must be provided if cond_based_affine is True'
+        if elementwise_affine:
+            if cond_based_affine and cond_dim is not None:
+                self.affine = Linear(cond_dim, dim, init_weight=1e-5)
+            else:
+                self.weight = nn.Parameter(torch.ones(self.dim))
+        else:
+            self.register_parameter("weight", None)
+
+    def forward(self, x, cond_emb=None):   
+        eps = torch.finfo(x.dtype).eps if self.eps is None else self.eps
+        rms = torch.rsqrt(torch.square(x).mean(dim=-1, keepdim=True) + eps)
+        out = ((x - x.mean(dim=-1, keepdim=True)) * rms) if self.centering else (x * rms)
+        if self.elementwise_affine:
+            if self.cond_based_affine:
+                out = out.mul(1. + self.affine(cond_emb).unsqueeze(1))
+            else:
+                out = out.mul(self.weight)
+        return out
+
+class AdaLN(nn.Module):
+    def __init__(self,dim, cond_dim=None, elementwise_affine=True, cond_based_affine=True, bias=False):
+        super(AdaLN, self).__init__()        
+        self.norm = nn.LayerNorm(dim, elementwise_affine=elementwise_affine and not cond_based_affine)
+        self.cond_based_affine = cond_based_affine
+        self.bias = bias
+        assert not(cond_dim is None and cond_based_affine), 'cond_dim must be provided if cond_based_affine is True'
+        self.affine = Linear(cond_dim,  2 * dim if bias else dim, init_weight=1e-5) if cond_based_affine and cond_dim is not None else None      
+    def forward(self, x, cond_emb=None):  
+        x = self.norm(x)
+        if self.cond_based_affine:
+            if self.bias:
+                shift, scale = self.affine(cond_emb).unsqueeze(1).chunk(2, dim=-1)
+                x = x.mul(1. + scale).add_(shift)
+            else:
+                scale = self.affine(cond_emb).unsqueeze(1)
+                x = x.mul(1. + scale)
+        return x
+
+class TimestepEmbedder(nn.Module):
+    """
+    Embeds scalar timesteps into vector representations.
+    """
+    def __init__(self, hidden_size, frequency_embedding_size=256):
+        super().__init__()
+        self.mlp = nn.Sequential(
+            Linear(frequency_embedding_size, hidden_size, init_weight=1e-5),
+            nn.SiLU(),
+            Linear(hidden_size, hidden_size, init_weight=1e-5),
+        )
+        self.frequency_embedding_size = frequency_embedding_size
+
+    @staticmethod
+    def timestep_embedding(t, dim, max_period=10000):
+        """
+        Create sinusoidal timestep embeddings.
+        :param t: a 1-D Tensor of N indices, one per batch element.
+                          These may be fractional.
+        :param dim: the dimension of the output.
+        :param max_period: controls the minimum frequency of the embeddings.
+        :return: an (N, D) Tensor of positional embeddings.
+        """
+        # https://github.com/openai/glide-text2im/blob/main/glide_text2im/nn.py
+        half = dim // 2
+        freqs = torch.exp(
+            -math.log(max_period) * torch.arange(start=0, end=half, dtype=torch.float32) / half
+        ).to(device=t.device)
+        args = t[:, None].float() * freqs[None]
+        embedding = torch.cat([torch.cos(args), torch.sin(args)], dim=-1)
+        if dim % 2:
+            embedding = torch.cat([embedding, torch.zeros_like(embedding[:, :1])], dim=-1)
+        return embedding
+
+    def forward(self, t):
+        t_freq = self.timestep_embedding(t, self.frequency_embedding_size)
+        t_emb = self.mlp(t_freq)
+        return t_emb
+
+# ROPE
+def exists(val):
+    return val is not None
+
+def default(val, d):
+    return val if exists(val) else d
+
+def broadcat(tensors, dim = -1):
+    broadcasted_tensors = broadcast_tensors(*tensors)
+    return torch.cat(broadcasted_tensors, dim = dim)
+
+def slice_at_dim(t, dim_slice: slice, *, dim):
+    dim += (t.ndim if dim < 0 else 0)
+    colons = [slice(None)] * t.ndim
+    colons[dim] = dim_slice
+    return t[tuple(colons)]
+
+def rotate_half(x):
+    x = rearrange(x, '... (d r) -> ... d r', r = 2)
+    x1, x2 = x.unbind(dim = -1)
+    x = torch.stack((-x2, x1), dim = -1)
+    return rearrange(x, '... d r -> ... (d r)')
+
+def apply_rotary_emb(
+    freqs,
+    t,
+    start_index = 0,
+    scale = 1.,
+    seq_dim = -2,
+    freqs_seq_dim = None
+):
+    dtype = t.dtype
+    with torch.amp.autocast('cuda',enabled=False):
+        freqs = freqs.float()
+        t = t.float()
+        if not exists(freqs_seq_dim):
+            if freqs.ndim == 2 or t.ndim == 3:
+                freqs_seq_dim = 0
+
+        if t.ndim == 3 or exists(freqs_seq_dim):
+            seq_len = t.shape[seq_dim]
+            freqs = slice_at_dim(freqs, slice(-seq_len, None), dim = freqs_seq_dim)
+
+        rot_dim = freqs.shape[-1]
+        end_index = start_index + rot_dim
+
+        assert rot_dim <= t.shape[-1], f'feature dimension {t.shape[-1]} is not of sufficient size to rotate in all the positions {rot_dim}'
+
+        # Split t into three parts: left, middle (to be transformed), and right
+        t_left = t[..., :start_index]
+        t_middle = t[..., start_index:end_index]
+        t_right = t[..., end_index:]
+
+        # Apply rotary embeddings without modifying t in place    
+        t_transformed = (t_middle * freqs.cos() * scale) + (rotate_half(t_middle) * freqs.sin() * scale)
+            
+        out = torch.cat((t_left, t_transformed, t_right), dim=-1)
+
+    return out.type(dtype)
+
+def apply_learned_rotations(rotations, t, start_index = 0, freq_ranges = None):
+    if exists(freq_ranges):
+        rotations = einsum('..., f -> ... f', rotations, freq_ranges)
+        rotations = rearrange(rotations, '... r f -> ... (r f)')
+
+    rotations = repeat(rotations, '... n -> ... (n r)', r = 2)
+    return apply_rotary_emb(rotations, t, start_index = start_index)
+
+class RotaryEmbedding(Module):
+    def __init__(
+        self,
+        dim,
+        custom_freqs= None,
+        freqs_for= 'lang',
+        theta = 10000,
+        max_freq = 10,
+        num_freqs = 1,
+        learned_freq = False,
+        use_xpos = True,
+        xpos_scale_base = 512,
+        interpolate_factor = 1.,
+        theta_rescale_factor = 1.,
+        seq_before_head_dim = False,
+        cache_if_possible = True,
+        cache_max_seq_len = 8192
+    ):
+        super().__init__()
+        theta *= theta_rescale_factor ** (dim / (dim - 2))
+
+        self.freqs_for = freqs_for
+
+        if exists(custom_freqs):
+            freqs = custom_freqs
+        elif freqs_for == 'lang':
+            freqs = 1. / (theta ** (torch.arange(0, dim, 2)[:(dim // 2)].float() / dim))
+        elif freqs_for == 'pixel':
+            freqs = torch.linspace(1., max_freq / 2, dim // 2) * pi
+        elif freqs_for == 'constant':
+            freqs = torch.ones(num_freqs).float()
+
+        self.cache_if_possible = cache_if_possible
+        self.cache_max_seq_len = cache_max_seq_len
+
+        self.register_buffer('cached_freqs', torch.zeros(cache_max_seq_len, dim), persistent = False)
+        self.cached_freqs_seq_len = 0
+
+        self.freqs = nn.Parameter(freqs, requires_grad = learned_freq)
+
+        self.learned_freq = learned_freq
+        
+   
+        # dummy for device
+        self.register_buffer('dummy', torch.tensor(0), persistent = False)
+
+        # default sequence dimension
+        self.seq_before_head_dim = seq_before_head_dim
+        self.default_seq_dim = -3 if seq_before_head_dim else -2
+
+        # interpolation factors
+        assert interpolate_factor >= 1.
+        self.interpolate_factor = interpolate_factor
+
+        # xpos
+        self.use_xpos = use_xpos
+
+        if not use_xpos:
+            return
+
+        scale = (torch.arange(0, dim, 2) + 0.4 * dim) / (1.4 * dim)
+        self.scale_base = xpos_scale_base
+
+        self.register_buffer('scale', scale, persistent = False)
+        self.register_buffer('cached_scales', torch.zeros(cache_max_seq_len, dim), persistent = False)
+        self.cached_scales_seq_len = 0
+
+
+
+    @property
+    def device(self):
+        return self.dummy.device
+
+    def get_seq_pos(self, seq_len, device, dtype, offset = 0):
+        return (torch.arange(seq_len, device = device, dtype = dtype) + offset) / self.interpolate_factor
+
+    def rotate_queries_or_keys(self, t, seq_dim = None, offset = 0, scale = None):
+        seq_dim = default(seq_dim, self.default_seq_dim)
+
+        assert not self.use_xpos or exists(scale), 'you must use `.rotate_queries_and_keys` method instead and pass in both queries and keys, for length extrapolatable rotary embeddings'
+
+        device, dtype, seq_len = t.device, t.dtype, t.shape[seq_dim]
+
+        seq = self.get_seq_pos(seq_len, device = device, dtype = dtype, offset = offset)
+
+        freqs = self.forward(seq, seq_len = seq_len, offset = offset)
+
+        if seq_dim == -3:
+            freqs = rearrange(freqs, 'n d -> n 1 d')
+
+        return self.apply_rotary_emb(freqs, t, scale = default(scale, 1.), seq_dim = seq_dim)
+
+    def rotate_queries_with_cached_keys(self, q, k, seq_dim = None, offset = 0):
+        dtype, device, seq_dim = q.dtype, q.device, default(seq_dim, self.default_seq_dim)
+
+        q_len, k_len = q.shape[seq_dim], k.shape[seq_dim]
+        assert q_len <= k_len
+
+        q_scale = k_scale = 1.
+
+        if self.use_xpos:
+            seq = self.get_seq_pos(k_len, dtype = dtype, device = device)
+
+            q_scale = self.get_scale(seq[-q_len:]).type(dtype)
+            k_scale = self.get_scale(seq).type(dtype)
+
+        rotated_q = self.rotate_queries_or_keys(q, seq_dim = seq_dim, scale = q_scale, offset = k_len - q_len + offset)
+        rotated_k = self.rotate_queries_or_keys(k, seq_dim = seq_dim, scale = k_scale ** -1)
+
+        rotated_q = rotated_q.type(q.dtype)
+        rotated_k = rotated_k.type(k.dtype)
+
+        return rotated_q, rotated_k
+
+    def rotate_queries_and_keys(self, q, k, seq_dim = None, offset = 0):
+        seq_dim = default(seq_dim, self.default_seq_dim)
+
+        device, dtype, seq_len = q.device, q.dtype, q.shape[seq_dim]
+
+        seq = self.get_seq_pos(seq_len, dtype = dtype, device = device, offset = offset)
+
+        freqs = self.forward(seq, seq_len = seq_len, offset = offset)
+        scale = self.get_scale(seq, seq_len = seq_len, offset = offset).to(dtype)
+
+        if seq_dim == -3:
+            freqs = rearrange(freqs, 'n d -> n 1 d')
+            scale = rearrange(scale, 'n d -> n 1 d')
+
+        rotated_q = apply_rotary_emb(freqs, q, scale = scale, seq_dim = seq_dim)
+        rotated_k = apply_rotary_emb(freqs, k, scale = scale ** -1, seq_dim = seq_dim)
+
+        rotated_q = rotated_q.type(q.dtype)
+        rotated_k = rotated_k.type(k.dtype)
+
+        return rotated_q, rotated_k
+
+    def get_scale(
+        self,
+        t: Tensor,
+        seq_len= None,
+        offset = 0
+    ):
+
+        should_cache = (
+            self.cache_if_possible and
+            exists(seq_len) and
+            (offset + seq_len) <= self.cache_max_seq_len
+        )
+
+        if (
+            should_cache and \
+            exists(self.cached_scales) and \
+            (seq_len + offset) <= self.cached_scales_seq_len
+        ):
+            return self.cached_scales[offset:(offset + seq_len)]
+
+        scale = 1.
+        if self.use_xpos:
+            power = (t - len(t) // 2) / self.scale_base
+            scale = self.scale ** rearrange(power, 'n -> n 1')
+            scale = repeat(scale, 'n d -> n (d r)', r = 2)
+
+        if should_cache and offset == 0:
+            self.cached_scales[:seq_len] = scale.detach()
+            self.cached_scales_seq_len = seq_len
+
+        return scale
+
+    def get_axial_freqs(self, *dims):
+        Colon = slice(None)
+        all_freqs = []
+
+        for ind, dim in enumerate(dims):
+            if self.freqs_for == 'pixel':
+                pos = torch.linspace(-1, 1, steps = dim, device = self.device)
+            else:
+                pos = torch.arange(dim, device = self.device)
+
+            freqs = self.forward(pos, seq_len = dim)
+
+            all_axis = [None] * len(dims)
+            all_axis[ind] = Colon
+
+            new_axis_slice = (Ellipsis, *all_axis, Colon)
+            all_freqs.append(freqs[new_axis_slice])
+
+        all_freqs = broadcast_tensors(*all_freqs)
+        return torch.cat(all_freqs, dim = -1)
+
+    def forward(
+        self,
+        t,
+        seq_len = None,
+        offset = 0
+    ):
+        dtype = t.dtype
+        with torch.amp.autocast('cuda',enabled=False):
+            t = t.float()
+            freqs = self.freqs.float()
+
+            should_cache = (
+                self.cache_if_possible and
+                not self.learned_freq and
+                exists(seq_len) and
+                self.freqs_for != 'pixel' and
+                (offset + seq_len) <= self.cache_max_seq_len
+            )
+
+            if (
+                should_cache and \
+                exists(self.cached_freqs) and \
+                (offset + seq_len) <= self.cached_freqs_seq_len
+            ):
+                return self.cached_freqs[offset:(offset + seq_len)].detach()
+
+
+            freqs = einsum('..., f -> ... f', t, freqs)
+            freqs = repeat(freqs, '... n -> ... (n r)', r = 2)
+
+            if should_cache and offset == 0:
+                self.cached_freqs[:seq_len] = freqs.detach()
+                self.cached_freqs_seq_len = seq_len
+
+        return freqs.type(dtype)
+
+# DropPath & FFN
+def drop_path(x, drop_prob: float = 0., training: bool = False, scale_by_keep: bool = True):    # taken from timm
+    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 = x.new_empty(shape).bernoulli_(keep_prob)
+    if keep_prob > 0.0 and scale_by_keep:
+        random_tensor.div_(keep_prob)
+    return x * random_tensor
+
+class DropPath(nn.Module):  # taken from timm
+    def __init__(self, drop_prob: float = 0., scale_by_keep: bool = True):
+        super(DropPath, self).__init__()
+        self.drop_prob = drop_prob
+        self.scale_by_keep = scale_by_keep
+    
+    def forward(self, x):
+        return drop_path(x, self.drop_prob, self.training, self.scale_by_keep)
+    
+    def extra_repr(self):
+        return f'(drop_prob=...)'
+
+class FFN(nn.Module):
+    def __init__(self, in_features, hidden_features=None, out_features=None, drop=0., ffn_type='geglu'):
+        super().__init__()
+        out_features = out_features or in_features
+        hidden_features = hidden_features or in_features
+        self.ffn_type = ffn_type
+        if self.ffn_type == 'geglu':
+            self.fc1 = Linear(in_features, 2*hidden_features)
+            self.act = nn.GELU(approximate='tanh')
+            self.fc2 = Linear(hidden_features, out_features)
+        elif self.ffn_type == 'ffn':
+            self.fc1 = Linear(in_features, hidden_features)
+            self.act = nn.GELU(approximate='tanh')
+            self.fc2 = Linear(hidden_features, out_features)
+        self.drop = nn.Dropout(drop, inplace=True) if drop > 0 else nn.Identity()
+    
+    def forward(self, x):
+        if self.ffn_type == 'geglu':
+            gate, value = self.fc1(x).chunk(2, dim=-1)
+            return self.drop(self.fc2(self.act(value)*gate))
+        elif self.ffn_type == 'ffn':
+            return self.drop(self.fc2(self.act(self.fc1(x))))
+    
+
+class Attention(nn.Module):
+    def __init__(
+        self, block_idx, embed_dim=768, num_heads=12,
+        attn_drop=0., proj_drop=0., attn_l2_norm=False, 
+        rope=False
+    ):
+        super().__init__()
+        assert embed_dim % num_heads == 0
+        self.block_idx = block_idx
+        self.num_heads = num_heads
+        self.head_dim = embed_dim // num_heads
+        self.attn_l2_norm = attn_l2_norm
+        self.rope = RotaryEmbedding(embed_dim//num_heads) if rope else None
+        
+        if self.attn_l2_norm:
+            self.scale = 1
+            self.scale_mul_1H11 = nn.Parameter(torch.full(size=(1, self.num_heads, 1, 1), fill_value=4.0).log(), requires_grad=True)
+            self.max_scale_mul = torch.log(torch.tensor(100)).item()
+        else:
+            self.scale = 0.25 / math.sqrt(self.head_dim)
+        
+        self.to_q = Linear(embed_dim, embed_dim, bias=False)
+        self.to_kv = Linear(embed_dim, embed_dim * 2, bias=False)
+
+        self.register_buffer('zero_k_bias', torch.zeros(embed_dim))
+        
+        self.proj = Linear(embed_dim, embed_dim)
+        self.proj_drop = nn.Dropout(proj_drop, inplace=True) if proj_drop > 0 else nn.Identity()
+        self.attn_drop: float = attn_drop    
+    
+    def forward(self, x, context_emb=None, causal=False, attn_bias=None, cache_kv=False, past_kvs=None):
+        B, L, C = x.shape
+
+        q = self.to_q(x) 
+        q = einops.rearrange(q, 'b l (h d) -> b l h d', h=self.num_heads)
+        
+        kv = self.to_kv(x) if context_emb is None else self.to_kv(context_emb)
+        kv = einops.rearrange(kv, 'b l (h d) -> b l h d', h=self.num_heads)
+        k, v = kv.chunk(2, dim=-1)
+
+        if self.rope is not None:
+            past_len = 0
+            if past_kvs is not None:
+                past_len = past_kvs[0].shape[2]
+            q, k = self.rope.rotate_queries_and_keys(q, k, offset=past_len)
+        
+        q = q.permute(0, 2, 1, 3)  # [B, H, L, D]
+        k = k.permute(0, 2, 1, 3)  # [B, H, L, D]
+        v = v.permute(0, 2, 1, 3)  # [B, H, L, D]
+
+        if past_kvs is not None:
+            past_keys, past_values = past_kvs
+            k = torch.cat((past_keys, k), dim=2)
+            v = torch.cat((past_values, v), dim=2)
+    
+        dropout_p = self.attn_drop if self.training else 0.0
+        oup = F.scaled_dot_product_attention(
+            query=q, 
+            key=k, 
+            value=v, 
+            scale=self.scale, 
+            is_causal=causal and past_kvs is None,
+            dropout_p=dropout_p,
+            attn_mask=attn_bias
+        ).transpose(1, 2).reshape(B, L, -1)
+
+        out = self.proj_drop(self.proj(oup))
+        if not cache_kv:
+            return out
+        else:
+            return out, (k,v)
+
+class TransformerLayer(nn.Module):
+    def __init__(
+        self, block_idx, embed_dim, cond_dim,
+        num_heads, mlp_ratio=8/3, drop=0., attn_drop=0., drop_path=0., attn_l2_norm=False, cross_attn=False, rope=False,
+        ffn_type='geglu', norm_layer=AdaRMSNorm
+    ):
+        super(TransformerLayer, self).__init__()
+        self.block_idx = block_idx
+        self.drop_path = DropPath(drop_path) if drop_path > 0. else nn.Identity()
+
+        self.adaln_1 = norm_layer(embed_dim, cond_dim)
+        self.attn1 = Attention(block_idx=block_idx, embed_dim=embed_dim, num_heads=num_heads, attn_drop=attn_drop, proj_drop=drop, attn_l2_norm=attn_l2_norm, rope=rope)
+        self.gate1 = Linear(cond_dim, embed_dim, init_weight=1e-5)
+
+        if cross_attn:
+            self.has_cross_attn = True
+            self.adaln_2 = norm_layer(embed_dim, cond_dim)
+            self.adaln_ctx = norm_layer(embed_dim, cond_dim)
+
+            self.attn2 = Attention(block_idx=block_idx, embed_dim=embed_dim, num_heads=num_heads, attn_drop=attn_drop, proj_drop=drop, attn_l2_norm=attn_l2_norm, rope=rope)
+            self.gate2 = Linear(cond_dim, embed_dim, init_weight=1e-5)
+        else:
+            self.has_cross_attn = False
+
+        self.adaln_mlp = norm_layer(embed_dim, cond_dim)
+        mlp_ratio = eval(mlp_ratio) if isinstance(mlp_ratio, str) else mlp_ratio
+        self.ffn = FFN(in_features=embed_dim, hidden_features=round(embed_dim * mlp_ratio), drop=drop, ffn_type=ffn_type)
+        self.gate_mlp = Linear(cond_dim, embed_dim, init_weight=1e-5)
+        
+    def forward(self, x,context_emb=None,  cond_emb=None, causal=False, attn_mask=None, cache_kv=False, past_kvs=None):
+        # If we are in "concat" mode, `context_emb` is None and masking should apply to self-attention.
+        # If we are in "cross_attn" mode, masking should apply to cross-attention (see below).
+        self_attn_mask = attn_mask if (attn_mask is not None and context_emb is None) else None
+
+        attn_out1 = self.attn1(
+            self.adaln_1(x, cond_emb),
+            context_emb=None,
+            causal=causal,
+            attn_bias=self_attn_mask,
+            cache_kv=cache_kv,
+            past_kvs=past_kvs,
+        )
+        if cache_kv:
+            attn_out1, new_kvs = attn_out1[0], attn_out1[1]
+        
+        x = x + self.gate1(cond_emb).unsqueeze(1) * self.drop_path(attn_out1)
+        if context_emb is not None and self.has_cross_attn:
+            x = x + self.gate2(cond_emb).unsqueeze(1) * self.drop_path(self.attn2(self.adaln_2(x, cond_emb), context_emb=self.adaln_ctx(context_emb,cond_emb), causal=False, attn_bias=attn_mask))
+        x = x + self.gate_mlp(cond_emb).unsqueeze(1) * self.drop_path(self.ffn(self.adaln_mlp(x, cond_emb)))
+    
+        if cache_kv:
+            return x, new_kvs
+        else:
+            return x
+
+class Transformer(nn.Module):
+    def __init__(
+        self, 
+        layer_num, 
+        input_dim, dim, output_dim, max_seq_len, heads,
+        cond_input_dim, cond_dim, 
+        context_type='none', ctx_input_dim=None, ctx_max_seq_len=None, 
+        emb_dropout=0., drop=0., 
+        rope=False, abs_pos=False, 
+        l2_norm=False,  causal=False, 
+        zero_out=False, 
+        query_type='learnable', # 'learnable', 'noise', 'none'
+        out_act=True,
+        ffn_type='geglu', # ffn or geglu
+        norm_layer='RMSNorm', # LayerNorm or RMSNorm
+        mlp_ratio=8/3, **params
+    ):
+        super(Transformer, self).__init__()
+        self.input_dim = input_dim
+        self.output_dim = output_dim
+        
+
+       
+        self.query_type = query_type
+            
+        self.cond_input_dim = cond_input_dim
+        self.dim = dim
+        self.causal = causal
+        self.cross_attn = context_type=='cross_attn'
+        norm_layer = AdaRMSNorm if norm_layer=="RMSNorm" else AdaLN
+        self.layers = nn.ModuleList([
+            TransformerLayer(block_idx=i, embed_dim=dim, 
+            cond_dim=cond_dim, num_heads=heads, mlp_ratio=mlp_ratio, 
+            drop=drop, attn_drop=drop, attn_l2_norm=l2_norm, rope=rope, cross_attn=self.cross_attn, 
+            ffn_type=ffn_type, norm_layer=norm_layer)
+            for i in range(layer_num)
+        ])
+
+        out_weight = 1e-5 if zero_out else 1
+        scale = np.sqrt(dim)
+        self.adaln_out = norm_layer(dim, cond_dim)
+
+        self.out_activate = nn.GELU(approximate='tanh') if out_act else None
+        self.out = Linear(dim, output_dim,init_weight=out_weight)
+
+        if self.query_type!='noise':
+            self.cond_proj = nn.Sequential(Linear(cond_input_dim, cond_dim),norm_layer(cond_dim, cond_based_affine=False), nn.GELU(approximate='tanh'))
+        else:
+            self.cond_proj = Linear(cond_input_dim, cond_dim)
+            self.cond_norm_act = nn.Sequential(norm_layer(cond_dim, cond_based_affine=False), nn.GELU(approximate='tanh'))
+        self.max_seq_len = max_seq_len
+        self.context_type = context_type # "concat, cross_attn, none"
+        if abs_pos:
+            if context_type == 'none':
+                self.abs_pos = nn.Embedding(max_seq_len, dim)
+            elif context_type == 'concat':
+                ctx_max_seq_len = default(ctx_max_seq_len, max_seq_len)
+                self.abs_pos = nn.Embedding(max_seq_len + ctx_max_seq_len, dim)
+            elif context_type == 'cross_attn':
+                ctx_max_seq_len = default(ctx_max_seq_len, max_seq_len)
+                self.abs_pos = nn.Embedding(max_seq_len, dim)
+                self.ctx_pos = nn.Embedding(ctx_max_seq_len, dim)
+            else:
+                raise ValueError(f"Invalid context_type: {context_type}")
+        else:
+            self.abs_pos = None
+
+        if context_type == 'concat':
+            self.input_layer = Linear(input_dim, dim)
+            self.learnable_query_param = nn.Parameter(torch.zeros(1, max_seq_len, input_dim).normal_(0, 1)) if self.query_type == 'learnable' else None
+            ctx_input_dim = default(ctx_input_dim, input_dim)
+            self.ctx_input_dim = ctx_input_dim
+            self.ctx_input_layer =  Linear(ctx_input_dim, dim)
+        elif context_type == 'cross_attn':
+            self.input_layer = Linear(input_dim, dim)
+            ctx_input_dim = default(ctx_input_dim, input_dim)
+            self.ctx_input_dim = ctx_input_dim
+            self.learnable_query_param = nn.Parameter(torch.zeros(1, max_seq_len, input_dim).normal_(0, 1)) if self.query_type == 'learnable' else None
+            self.ctx_input_layer =  Linear(ctx_input_dim, dim)
+        else:
+            self.input_layer = Linear(input_dim, dim)
+            
+        if self.query_type == 'noise':
+            self.t_embedder = TimestepEmbedder(cond_input_dim)
+            self.t_proj = Linear(cond_input_dim,cond_dim)
+
+    def forward(self, x, condition, attn_mask=None, cache_kv=False, past_kvs=None, query_tokens=None, t=None):
+        
+        # Inject time embedding if using noise query
+        if self.query_type == 'noise':
+            cond_emb = self.cond_proj(condition)
+            if t is not None:
+                t_emb = self.t_embedder(t)
+                t_emb = self.t_proj(t_emb)
+                cond_emb = cond_emb + t_emb
+            cond_emb = self.cond_norm_act(cond_emb)
+        else:
+            cond_emb = self.cond_proj(condition)
+        if self.context_type == 'none':
+            x, context_emb = self.input_layer(x.contiguous()), None
+        elif self.context_type == 'concat':
+            B = x.size(0)
+            if query_tokens is None:
+                if self.query_type == 'learnable':
+                    query_tokens = self.learnable_query_param.expand(B, -1, -1)
+                else:
+                    # Should be provided for noise query, or error if not available
+                    raise ValueError(f"query_tokens is required when query_type is '{self.query_type}'.")
+
+            query_tokens = self.input_layer(query_tokens.contiguous())
+            x = self.ctx_input_layer(x.contiguous())
+            x =  torch.cat((query_tokens,x), dim=1)
+            x, context_emb = x, None
+        else:
+            # cross_attn: x is context (keys/values), query_tokens are queries
+            B = x.size(0)
+            if query_tokens is None:
+                if self.query_type == 'learnable':
+                    query_tokens = self.learnable_query_param.expand(B, -1, -1)
+                else:
+                    raise ValueError(f"query_tokens is required when query_type is '{self.query_type}'.")
+
+            x, context_emb = self.input_layer(query_tokens.contiguous()), self.ctx_input_layer(x.contiguous())
+            
+        if self.abs_pos is not None:
+            if self.context_type == 'none':
+                pos_ids = torch.arange(x.size(1), device=x.device)
+                if past_kvs is not None:
+                    pos_ids = pos_ids + past_kvs[0][0].shape[2]
+                x = x + self.abs_pos(pos_ids)
+            elif self.context_type == 'concat':
+                pos_ids = torch.arange(x.size(1), device=x.device)
+                x = x + self.abs_pos(pos_ids)
+            elif self.context_type == 'cross_attn':
+                pos_ids = torch.arange(x.size(1), device=x.device)
+                x = x + self.abs_pos(pos_ids)
+                if context_emb is not None and self.ctx_pos is not None:
+                    context_emb = context_emb + self.ctx_pos(torch.arange(context_emb.size(1), device=context_emb.device))
+
+        kv_caches = [] 
+        for i,layer in enumerate(self.layers):
+            layer_output = layer(x, context_emb, cond_emb, self.causal, attn_mask, cache_kv, past_kvs[i] if past_kvs is not None else None)
+            if cache_kv:
+                x, new_kvs = layer_output[0], layer_output[1]
+                kv_caches.append(new_kvs)
+            else:
+                x = layer_output
+        if self.context_type == 'concat':
+            x = x[:, :self.max_seq_len]
+        out = self.adaln_out(x, cond_emb)
+        if self.out_activate is not None:   
+            out = self.out_activate(out)
+        out = self.out(out)
+        
+        if cache_kv:
+            return out, kv_caches
+        else:
+            return out
+
+
+class Encoder(nn.Module):
+    def __init__(self, config) -> None:
+        super().__init__()
+        self.enc = Transformer(**config["Encoder"])
+        if config["vq_type"] == "HGQ":
+            self.quantizer = HGQ(**config["Quantizer"])
+        elif config["vq_type"] == "VQ":
+            self.quantizer = VQ(**config["Quantizer"])
+        else:
+            raise ValueError(f"Invalid quantizer type: {config['vq_type']}")
+    def forward(self, x, labels, training=True):
+        h = self.enc(x, labels)
+        quant, idx, one_hot, vqloss = self.quantizer(h, labels, training=training)
+        return quant, idx, one_hot, vqloss
+    
+    
+class Decoder(nn.Module):
+    def __init__(self, config) -> None:
+        super().__init__()
+        self.dec = Transformer(**config["Decoder"])
+    
+    
+    def forward(self, quant, labels, attn_mask=None, query_tokens=None, t=None):
+        x_hat = self.dec(quant, labels, attn_mask=attn_mask, query_tokens=query_tokens, t=t)
+        return x_hat
+
+
+class ARModel(nn.Module):
+    def __init__(self, config) -> None:
+        super().__init__()
+        self.bos = nn.Parameter(torch.zeros(1, 1, config["ARModel"]["input_dim"]).normal_(0,1))
+        self.semantic_emb = nn.Embedding(config["Quantizer"]["codebook_size"], config["ARModel"]["input_dim"])
+        self.ar_model = Transformer(**config["ARModel"])
+
+        self.temperature = config["ARModel"]["temperature"]
+        self.max_length = config["ARModel"]["max_seq_len"]
+
+       
+    def forward(self,idx,labels, temperature=None):
+        bos = self.bos.expand(idx.shape[0], -1, -1)
+        shift_input = torch.cat([bos, self.semantic_emb(idx)[:,:-1]],dim=1)
+        logits = self.ar_model(shift_input, labels)
+        
+        if temperature is not None:
+            logits = logits / temperature
+        elif self.temperature is not None:
+            logits = logits / self.temperature
+        return logits
+
+    @torch.no_grad()
+    @torch._dynamo.disable
+    def sampling(self, bz, class_label=None, temperature=1.0,  topK=None, topP=None, cfg=16.0, cfg_schedule='cosine', cfg_power=2.75, cache_kv=False):
+        cfg = 0. if class_label is None else cfg
+        if cfg>0. :
+            num_classes = class_label.shape[1]
+            uncond_class_label = F.one_hot(torch.full((bz,), num_classes - 1, device=class_label.device, dtype=torch.long), num_classes=num_classes).float()
+        quant_input = self.bos.expand(bz, -1, -1)
+        quant_output = []
+        past_kvs, uncond_past_kvs = None, None
+        for step in range(self.max_length):
+            ar_out = self.ar_model(quant_input, class_label, cache_kv=cache_kv, past_kvs=past_kvs)
+            if cache_kv:
+                logits, past_kvs = ar_out
+            else:
+                logits = ar_out
+            logits = logits[:,-1]
+            if cfg>0.:
+                if cfg_schedule == 'constant':
+                    cfg_scale = cfg
+                elif cfg_schedule == 'linear':
+                    cfg_scale = 1.0 * (1-step/self.max_length) + cfg * (step/self.max_length)
+                elif cfg_schedule == 'cosine':
+                    cfg_scale = (1 - math.cos(
+                        ((step / self.max_length) ** cfg_power) * math.pi)) * 1/2
+                    cfg_scale = (cfg - 1) * cfg_scale + 1
+                else:
+                    raise ValueError(f"Invalid cfg_schedule: {cfg_schedule}")
+                uncond_ar_out = self.ar_model(quant_input, uncond_class_label, cache_kv=cache_kv, past_kvs=uncond_past_kvs)
+                if cache_kv:
+                    uncond_logits, uncond_past_kvs = uncond_ar_out
+                else:
+                    uncond_logits = uncond_ar_out
+                uncond_logits = uncond_logits[:,-1]
+                logits = cfg_scale * logits + (1 - cfg_scale) * uncond_logits
+            logits = logits / temperature
+
+            if topK is not None and topK>0.:
+                top_logits, top_indices = logits.topk(topK, dim=-1)
+                logits = torch.full_like(logits, float('-inf'))
+                logits.scatter_(dim=-1, index=top_indices, src=top_logits)
+            if topP is not None and 0.<topP<1.:
+                sorted_logits, sorted_indices = torch.sort(logits, dim=-1, descending=True)
+                probs_sum = sorted_logits.softmax(dim=-1).cumsum(dim=-1)
+                mask = probs_sum > topP
+                mask[..., 1:] = mask[...,:-1].clone()
+                mask[...,0] = False
+                sorted_logits[mask] = float('-inf')
+                logits = torch.full_like(logits, float('-inf'))
+                logits.scatter_(dim=-1, index=sorted_indices, src=sorted_logits)
+            next_idx = torch.multinomial(F.softmax(logits, dim=-1), 1)
+            quant_output.append(next_idx)
+            if not cache_kv:
+                quant_input = torch.cat((quant_input, self.semantic_emb(next_idx)), dim=1)
+            else:
+                quant_input = self.semantic_emb(next_idx)
+        quant_output = torch.cat(quant_output, dim=1)
+        return quant_output
+
+
+
+# deterministic Hard-Gumbel Softmax Quantizer   
+class HGQ(nn.Module):
+    def __init__(self, codebook_size, dim, z_dim, cond_dim, cond_based_affine=False, norm_layer='RMSNorm',temperature=1.0, **params):
+        super().__init__()
+        self.codebook_size = codebook_size  # codebook size (K)
+        self.dim = dim  # embedding dimension (D)
+        self.z_dim = z_dim
+        # Initialize embedding table
+        self.proj_in = Linear(dim, z_dim)
+        self.codebook = nn.Embedding(codebook_size, z_dim)
+        self.codebook.weight.data.uniform_(-1e-4, 1e-4)
+        self.norm_layer = norm_layer
+        if norm_layer == 'RMSNorm':
+            self.z_norm =  AdaRMSNorm(z_dim, cond_dim, elementwise_affine=True, cond_based_affine=cond_based_affine)
+            self.code_norm =  AdaRMSNorm(z_dim, cond_dim, elementwise_affine=True, cond_based_affine=cond_based_affine)
+        elif norm_layer == 'LN':
+            self.z_norm =  AdaLN(z_dim, cond_dim, elementwise_affine=True, cond_based_affine=cond_based_affine)
+            self.code_norm =  AdaLN(z_dim, cond_dim, elementwise_affine=True, cond_based_affine=cond_based_affine)
+        
+        self.out_proj = Linear(z_dim, dim, bias=False)
+        self.total_cnt = 0
+        self.book_cnt = torch.zeros(codebook_size)
+        self.temperature=temperature
+        self.float()
+
+    def forward(self, x, labels=None, training=True, log_usage=False, indices=None):
+        B, L, D = x.shape
+        x = x.float()
+        z = self.proj_in(x)
+
+        if self.norm_layer != 'none':
+            if self.norm_layer=='l2':
+                z= F.normalize(z, dim=-1)
+                codes= F.normalize(self.codebook.weight, dim=-1)
+            elif self.norm_layer=='RMSNorm':
+                z= self.z_norm(z, labels)
+                codes= self.code_norm(self.codebook.weight, labels)
+            elif self.norm_layer=='LN':
+                z= self.z_norm(z, labels)
+                codes= self.code_norm(self.codebook.weight, labels)
+            else:
+                raise ValueError(f"Invalid norm layer: {self.norm_layer}")
+
+        logits = torch.einsum('bld,nd->bln', z, codes)
+        prob = F.softmax(logits/self.temperature, dim=-1)
+
+        indices = torch.argmax(prob, dim=-1) if indices is None else indices
+        one_hot_ng = F.one_hot(indices, prob.shape[-1]).view(B, L, -1).to(z.device).to(z.dtype)
+        one_hot = prob + (one_hot_ng - prob).detach()
+ 
+        quant = torch.einsum('bln,nd->bld', one_hot, codes)
+        loss = 0.
+        quant_out = self.out_proj(quant)
+       
+        return quant_out, indices, one_hot, loss
+        
+    def get_codes_w_onehot(self, one_hot, labels=None, **params):
+        codes = self.codebook.weight
+        quant = torch.einsum('bln,nd->bld', one_hot, codes)
+
+        if self.norm_layer != 'none':
+            if self.norm_layer=='l2':
+                quant= F.normalize(quant, dim=-1)
+            elif self.norm_layer=='RMSNorm':
+                quant= self.code_norm(quant, labels)
+            elif self.norm_layer=='LN':
+                quant= self.code_norm(quant, labels)
+            else:
+                raise ValueError(f"Invalid norm layer: {self.norm_layer}")
+        return self.out_proj(quant)
+    def get_codes_w_indices(self, indices, labels=None, **params):
+        codes = self.codebook(indices)
+        if self.norm_layer != 'none':
+            if self.norm_layer=='l2':
+                codes= F.normalize(codes, dim=-1)
+            elif self.norm_layer=='RMSNorm':
+                codes= self.code_norm(codes, labels)
+            elif self.norm_layer=='LN':
+                codes= self.code_norm(codes, labels)
+            else:
+                raise ValueError(f"Invalid norm layer: {self.norm_layer}")
+        return self.out_proj(codes)
+
+
+class VQ(torch.nn.Module):
+    def __init__(self,
+                 codebook_size: int = 1024,
+                 dim: int = 256,
+                 z_dim: int = 256,
+                 beta: float = 0.25,
+                 use_norm: bool = False,
+                 temperature=1.0, **params
+                 ):
+        super().__init__()
+        self.codebook_size = codebook_size
+        self.dim = dim
+        self.z_dim = z_dim
+        self.beta = beta
+        self.total_cnt = 0
+        self.book_cnt = torch.zeros(codebook_size)
+        self.proj_in = Linear(dim, z_dim)
+        self.embedding = torch.nn.Embedding(codebook_size, z_dim)
+        self.embedding.weight.data.uniform_(-1.0 / codebook_size, 1.0 / codebook_size)
+        self.use_norm = use_norm
+        self.proj_out = Linear(z_dim, dim)
+        self.register_buffer('ema_vocab_hit_SV', torch.full((1, self.codebook_size), fill_value=0.0))
+        self.record_hit = 0
+        self.temperature=temperature
+        self.float()
+
+
+    def forward(self, x: torch.Tensor, labels=None, training=True):
+        x=x.float()
+        z = self.proj_in(x)
+
+        if self.use_norm:
+            z = F.normalize(z, dim=-1)
+            embedding = F.normalize(self.embedding.weight, dim=-1)
+            dist = -1 * torch.einsum('bld,nd->bln', z, embedding)
+        else:
+            embedding = self.embedding.weight
+
+            dist = torch.sum(z**2, dim=-1, keepdim=True) + \
+                torch.sum(embedding**2, dim=1) - 2 * \
+                torch.einsum('bld,nd->bln', z, embedding)
+
+        indices = torch.argmin(dist, dim=-1) # num_ele
+        ste_prob = (-dist/self.temperature).softmax(dim=-1)
+        one_hot_ng = F.one_hot(indices, dist.shape[-1]).view(x.shape[0], x.shape[1], -1).to(z.device).to(z.dtype) 
+        one_hot = ste_prob + (one_hot_ng - ste_prob).detach()
+
+        quant = embedding[indices]
+
+        loss = self.beta * torch.mean((quant.detach() - z) **2) +  torch.mean((quant - z.detach()) **2) if training else 0.
+
+
+        # preserve gradients
+        quant = z + (quant - z).detach()
+        quant = self.proj_out(quant)
+
+        return quant, indices, one_hot, loss
+    def get_codes_w_onehot(self, one_hot, labels=None, **params):
+        codes = self.embedding.weight
+        quant = torch.einsum('bln,nd->bld', one_hot, codes)
+        quant = F.normalize(quant, dim=-1) if self.use_norm else quant
+        return self.proj_out(quant)
+    def get_codes_w_indices(self, indices):
+        quant = self.embedding(indices)
+        quant = F.normalize(quant, dim=-1) if self.use_norm else quant
+        return self.proj_out(quant)

refs/cifar10-32x32_local.npz

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:5493968c4f29b2878f64d6a640351f377e32e86553db56eca8147a0ae50d93f2
+size 63557107

requirements.txt

@@ -0,0 +1,15 @@
+torch==2.8.0
+torchvision==0.23.0
+lightning==2.5.5
+omegaconf==2.3.0
+einops==0.8.1
+numpy==2.1.3
+scipy==1.16.3
+Pillow==11.1.0
+PyYAML==6.0.3
+wandb==0.23.0
+tqdm==4.67.1
+requests==2.32.5
+click==8.3.0
+setuptools==80.9.0
+tensorflow==2.19.1

scripts/dataset_tools.py

@@ -0,0 +1,618 @@
+# Copyright (c) 2022, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+#
+# This work is licensed under a Creative Commons
+# Attribution-NonCommercial-ShareAlike 4.0 International License.
+# You should have received a copy of the license along with this
+# work. If not, see http://creativecommons.org/licenses/by-nc-sa/4.0/
+
+"""Tool for creating ZIP/PNG based datasets."""
+
+import functools
+import gzip
+import io
+import json
+import os
+import pickle
+import re
+import sys
+import tarfile
+import zipfile
+from pathlib import Path
+from typing import Callable, Optional, Tuple, Union
+import click
+import numpy as np
+import PIL.Image
+from tqdm import tqdm
+
+#----------------------------------------------------------------------------
+# Parse a 'M,N' or 'MxN' integer tuple.
+# Example: '4x2' returns (4,2)
+
+def parse_tuple(s: str) -> Tuple[int, int]:
+    m = re.match(r'^(\d+)[x,](\d+)$', s)
+    if m:
+        return int(m.group(1)), int(m.group(2))
+    raise click.ClickException(f'cannot parse tuple {s}')
+
+#----------------------------------------------------------------------------
+
+def maybe_min(a: int, b: Optional[int]) -> int:
+    if b is not None:
+        return min(a, b)
+    return a
+
+#----------------------------------------------------------------------------
+
+def file_ext(name: Union[str, Path]) -> str:
+    return str(name).split('.')[-1]
+
+#----------------------------------------------------------------------------
+
+def is_image_ext(fname: Union[str, Path]) -> bool:
+    ext = file_ext(fname).lower()
+    return f'.{ext}' in PIL.Image.EXTENSION
+
+#----------------------------------------------------------------------------
+
+def open_image_folder(source_dir, *, max_images: Optional[int]):
+    input_images = [str(f) for f in sorted(Path(source_dir).rglob('*')) if is_image_ext(f) and os.path.isfile(f)]
+    arch_fnames = {fname: os.path.relpath(fname, source_dir).replace('\\', '/') for fname in input_images}
+    max_idx = maybe_min(len(input_images), max_images)
+
+    # Load labels.
+    labels = dict()
+    meta_fname = os.path.join(source_dir, 'dataset.json')
+    if os.path.isfile(meta_fname):
+        with open(meta_fname, 'r') as file:
+            data = json.load(file)['labels']
+            if data is not None:
+                labels = {x[0]: x[1] for x in data}
+
+    # No labels available => determine from top-level directory names.
+    if len(labels) == 0:
+        toplevel_names = {arch_fname: arch_fname.split('/')[0] if '/' in arch_fname else '' for arch_fname in arch_fnames.values()}
+        toplevel_indices = {toplevel_name: idx for idx, toplevel_name in enumerate(sorted(set(toplevel_names.values())))}
+        if len(toplevel_indices) > 1:
+            labels = {arch_fname: toplevel_indices[toplevel_name] for arch_fname, toplevel_name in toplevel_names.items()}
+
+    def iterate_images():
+        for idx, fname in enumerate(input_images):
+            img = np.array(PIL.Image.open(fname))
+            yield dict(img=img, label=labels.get(arch_fnames.get(fname)))
+            if idx >= max_idx - 1:
+                break
+    return max_idx, iterate_images()
+
+#----------------------------------------------------------------------------
+
+def open_image_zip(source, *, max_images: Optional[int]):
+    with zipfile.ZipFile(source, mode='r') as z:
+        input_images = [str(f) for f in sorted(z.namelist()) if is_image_ext(f)]
+        max_idx = maybe_min(len(input_images), max_images)
+
+        # Load labels.
+        labels = dict()
+        if 'dataset.json' in z.namelist():
+            with z.open('dataset.json', 'r') as file:
+                data = json.load(file)['labels']
+                if data is not None:
+                    labels = {x[0]: x[1] for x in data}
+
+    def iterate_images():
+        with zipfile.ZipFile(source, mode='r') as z:
+            for idx, fname in enumerate(input_images):
+                with z.open(fname, 'r') as file:
+                    img = np.array(PIL.Image.open(file))
+                yield dict(img=img, label=labels.get(fname))
+                if idx >= max_idx - 1:
+                    break
+    return max_idx, iterate_images()
+
+#----------------------------------------------------------------------------
+
+def open_lmdb(lmdb_dir: str, *, max_images: Optional[int]):
+    import cv2  # pyright: ignore [reportMissingImports] # pip install opencv-python
+    import lmdb  # pyright: ignore [reportMissingImports] # pip install lmdb
+
+    with lmdb.open(lmdb_dir, readonly=True, lock=False).begin(write=False) as txn:
+        max_idx = maybe_min(txn.stat()['entries'], max_images)
+
+    def iterate_images():
+        with lmdb.open(lmdb_dir, readonly=True, lock=False).begin(write=False) as txn:
+            for idx, (_key, value) in enumerate(txn.cursor()):
+                try:
+                    try:
+                        img = cv2.imdecode(np.frombuffer(value, dtype=np.uint8), 1)
+                        if img is None:
+                            raise IOError('cv2.imdecode failed')
+                        img = img[:, :, ::-1] # BGR => RGB
+                    except IOError:
+                        img = np.array(PIL.Image.open(io.BytesIO(value)))
+                    yield dict(img=img, label=None)
+                    if idx >= max_idx - 1:
+                        break
+                except:
+                    print(sys.exc_info()[1])
+
+    return max_idx, iterate_images()
+
+#----------------------------------------------------------------------------
+
+def open_cifar10(tarball: str, *, max_images: Optional[int]):
+    images = []
+    labels = []
+
+    with tarfile.open(tarball, 'r:gz') as tar:
+        for batch in range(1, 6):
+            member = tar.getmember(f'cifar-10-batches-py/data_batch_{batch}')
+            with tar.extractfile(member) as file:
+                data = pickle.load(file, encoding='latin1')
+            images.append(data['data'].reshape(-1, 3, 32, 32))
+            labels.append(data['labels'])
+
+    images = np.concatenate(images)
+    labels = np.concatenate(labels)
+    images = images.transpose([0, 2, 3, 1]) # NCHW -> NHWC
+    assert images.shape == (50000, 32, 32, 3) and images.dtype == np.uint8
+    assert labels.shape == (50000,) and labels.dtype in [np.int32, np.int64]
+    assert np.min(images) == 0 and np.max(images) == 255
+    assert np.min(labels) == 0 and np.max(labels) == 9
+
+    max_idx = maybe_min(len(images), max_images)
+
+    def iterate_images():
+        for idx, img in enumerate(images):
+            yield dict(img=img, label=int(labels[idx]))
+            if idx >= max_idx - 1:
+                break
+
+    return max_idx, iterate_images()
+
+#----------------------------------------------------------------------------
+
+def open_mnist(images_gz: str, *, max_images: Optional[int]):
+    labels_gz = images_gz.replace('-images-idx3-ubyte.gz', '-labels-idx1-ubyte.gz')
+    assert labels_gz != images_gz
+    images = []
+    labels = []
+
+    with gzip.open(images_gz, 'rb') as f:
+        images = np.frombuffer(f.read(), np.uint8, offset=16)
+    with gzip.open(labels_gz, 'rb') as f:
+        labels = np.frombuffer(f.read(), np.uint8, offset=8)
+
+    images = images.reshape(-1, 28, 28)
+    images = np.pad(images, [(0,0), (2,2), (2,2)], 'constant', constant_values=0)
+    assert images.shape == (60000, 32, 32) and images.dtype == np.uint8
+    assert labels.shape == (60000,) and labels.dtype == np.uint8
+    assert np.min(images) == 0 and np.max(images) == 255
+    assert np.min(labels) == 0 and np.max(labels) == 9
+
+    max_idx = maybe_min(len(images), max_images)
+
+    def iterate_images():
+        for idx, img in enumerate(images):
+            yield dict(img=img, label=int(labels[idx]))
+            if idx >= max_idx - 1:
+                break
+
+    return max_idx, iterate_images()
+
+#----------------------------------------------------------------------------
+
+def make_transform(
+    transform: Optional[str],
+    output_width: Optional[int],
+    output_height: Optional[int]
+) -> Callable[[np.ndarray], Optional[np.ndarray]]:
+    def scale(width, height, img):
+        w = img.shape[1]
+        h = img.shape[0]
+        if width == w and height == h:
+            return img
+        img = PIL.Image.fromarray(img)
+        ww = width if width is not None else w
+        hh = height if height is not None else h
+        img = img.resize((ww, hh), PIL.Image.Resampling.LANCZOS)
+        return np.array(img)
+
+    def center_crop(width, height, img):
+        crop = np.min(img.shape[:2])
+        img = img[(img.shape[0] - crop) // 2 : (img.shape[0] + crop) // 2, (img.shape[1] - crop) // 2 : (img.shape[1] + crop) // 2]
+        if img.ndim == 2:
+            img = img[:, :, np.newaxis].repeat(3, axis=2)
+        img = PIL.Image.fromarray(img, 'RGB')
+        img = img.resize((width, height), PIL.Image.Resampling.LANCZOS)
+        return np.array(img)
+
+    def center_crop_wide(width, height, img):
+        ch = int(np.round(width * img.shape[0] / img.shape[1]))
+        if img.shape[1] < width or ch < height:
+            return None
+
+        img = img[(img.shape[0] - ch) // 2 : (img.shape[0] + ch) // 2]
+        if img.ndim == 2:
+            img = img[:, :, np.newaxis].repeat(3, axis=2)
+        img = PIL.Image.fromarray(img, 'RGB')
+        img = img.resize((width, height), PIL.Image.Resampling.LANCZOS)
+        img = np.array(img)
+
+        canvas = np.zeros([width, width, 3], dtype=np.uint8)
+        canvas[(width - height) // 2 : (width + height) // 2, :] = img
+        return canvas
+
+    if transform is None:
+        return functools.partial(scale, output_width, output_height)
+    if transform == 'center-crop':
+        if output_width is None or output_height is None:
+            raise click.ClickException('must specify --resolution=WxH when using ' + transform + 'transform')
+        return functools.partial(center_crop, output_width, output_height)
+    if transform == 'center-crop-wide':
+        if output_width is None or output_height is None:
+            raise click.ClickException('must specify --resolution=WxH when using ' + transform + ' transform')
+        return functools.partial(center_crop_wide, output_width, output_height)
+    assert False, 'unknown transform'
+
+#----------------------------------------------------------------------------
+
+def open_dataset(source, *, max_images: Optional[int]):
+    if os.path.isdir(source):
+        if source.rstrip('/').endswith('_lmdb'):
+            return open_lmdb(source, max_images=max_images)
+        else:
+            return open_image_folder(source, max_images=max_images)
+    elif os.path.isfile(source):
+        if os.path.basename(source) == 'cifar-10-python.tar.gz':
+            return open_cifar10(source, max_images=max_images)
+        elif os.path.basename(source) == 'train-images-idx3-ubyte.gz':
+            return open_mnist(source, max_images=max_images)
+        elif file_ext(source) == 'zip':
+            return open_image_zip(source, max_images=max_images)
+        else:
+            assert False, 'unknown archive type'
+    else:
+        raise click.ClickException(f'Missing input file or directory: {source}')
+
+#----------------------------------------------------------------------------
+
+def create_labels_only(source: str) -> None:
+    """Create a StyleGAN-style dataset.json with image->class labels in-place.
+
+    - For directories: writes dataset.json under the source root.
+    - For ZIP files: writes dataset.json inside the ZIP archive.
+    - Labels are inferred from an existing dataset.json if present, otherwise
+      from top-level directory names (ImageNet-style). If labels cannot be
+      determined for all images, labels is set to None.
+    """
+    def _maybe_load_imagenet_val_labels(rel_fnames):
+        """Try to infer ImageNet val labels from LOC_val_solution.csv + LOC_synset_mapping.txt.
+
+        rel_fnames are relative paths within `source` (for dirs) or within the zip.
+        This function returns a mapping {rel_fname: class_idx} or an empty dict if
+        the required metadata files are not found or parsing fails.
+        """
+        root = Path(__file__).resolve().parent
+        val_csv = root / 'LOC_val_solution.csv'
+        synset_map = root / 'LOC_synset_mapping.txt'
+        if not (val_csv.is_file() and synset_map.is_file()):
+            return {}
+
+        # Build synset -> class index mapping, consistent with train (sorted by synset).
+        synsets = []
+        with open(synset_map, 'r') as f:
+            for line in f:
+                line = line.strip()
+                if not line:
+                    continue
+                syn = line.split()[0]
+                synsets.append(syn)
+        synsets = sorted(set(synsets))
+        class_to_idx = {syn: idx for idx, syn in enumerate(synsets)}
+
+        # Build image_id -> synset mapping from LOC_val_solution.csv
+        img_to_syn = {}
+        with open(val_csv, 'r') as f:
+            # Skip header
+            header = f.readline()
+            for line in f:
+                line = line.strip()
+                if not line:
+                    continue
+                parts = line.split(',', 1)
+                if len(parts) != 2:
+                    continue
+                image_id, pred_str = parts
+                pred_str = pred_str.strip()
+                if not pred_str:
+                    continue
+                syn = pred_str.split()[0]
+                img_to_syn[image_id] = syn
+
+        labels = {}
+        for rel in rel_fnames:
+            base = os.path.basename(rel)
+            stem, _ext = os.path.splitext(base)
+            syn = img_to_syn.get(stem, None)
+            if syn is None:
+                continue
+            cls_idx = class_to_idx.get(syn, None)
+            if cls_idx is None:
+                continue
+            labels[rel] = cls_idx
+        return labels
+
+    # Directory source: infer labels based on folder structure or metadata.
+    if os.path.isdir(source):
+        input_images = [
+            str(f) for f in sorted(Path(source).rglob('*'))
+            if is_image_ext(f) and os.path.isfile(f)
+        ]
+        arch_fnames = {
+            fname: os.path.relpath(fname, source).replace('\\', '/')
+            for fname in input_images
+        }
+
+        labels = dict()
+        meta_fname = os.path.join(source, 'dataset.json')
+        if os.path.isfile(meta_fname):
+            with open(meta_fname, 'r') as file:
+                data = json.load(file).get('labels', None)
+                if data is not None:
+                    labels = {x[0]: x[1] for x in data}
+
+        # No labels available => determine from top-level directory names (train-style),
+        # or fall back to ImageNet val metadata if available.
+        if len(labels) == 0:
+            toplevel_names = {
+                arch_fname: arch_fname.split('/')[0] if '/' in arch_fname else ''
+                for arch_fname in arch_fnames.values()
+            }
+            unique_toplevel = sorted(set(toplevel_names.values()))
+            if len(unique_toplevel) > 1:
+                toplevel_indices = {
+                    toplevel_name: idx
+                    for idx, toplevel_name in enumerate(unique_toplevel)
+                }
+                labels = {
+                    arch_fname: toplevel_indices[toplevel_names[arch_fname]]
+                    for arch_fname in arch_fnames.values()
+                }
+            else:
+                # Likely a flat ImageNet val-style directory; try metadata-based labels.
+                labels = _maybe_load_imagenet_val_labels(list(arch_fnames.values()))
+
+        labels_list = []
+        all_labeled = True
+        for arch_fname in arch_fnames.values():
+            lbl = labels.get(arch_fname, None)
+            if lbl is None:
+                all_labeled = False
+            labels_list.append([arch_fname, lbl])
+
+        metadata = {'labels': labels_list if all_labeled else None}
+        with open(meta_fname, 'w') as f:
+            json.dump(metadata, f)
+        return
+
+    # ZIP source: write dataset.json inside the archive.
+    if os.path.isfile(source) and file_ext(source) == 'zip':
+        with zipfile.ZipFile(source, mode='r') as z:
+            namelist = z.namelist()
+            if 'dataset.json' in namelist:
+                return
+            input_images = [f for f in sorted(namelist) if is_image_ext(f)]
+
+            labels = dict()
+            # Determine from top-level directory names inside zip, or fall back to
+            # ImageNet val metadata if available.
+            toplevel_names = {
+                fname: fname.split('/')[0] if '/' in fname else ''
+                for fname in input_images
+            }
+            unique_toplevel = sorted(set(toplevel_names.values()))
+            if len(unique_toplevel) > 1:
+                toplevel_indices = {
+                    toplevel_name: idx
+                    for idx, toplevel_name in enumerate(unique_toplevel)
+                }
+                labels = {
+                    fname: toplevel_indices[toplevel_names[fname]]
+                    for fname in input_images
+                }
+            else:
+                # Flat layout: try ImageNet val metadata (filenames like ILSVRC2012_val_*.JPEG).
+                labels = _maybe_load_imagenet_val_labels(input_images)
+
+            labels_list = []
+            all_labeled = True
+            for fname in input_images:
+                lbl = labels.get(fname, None)
+                if lbl is None:
+                    all_labeled = False
+                labels_list.append([fname, lbl])
+
+            metadata = {'labels': labels_list if all_labeled else None}
+
+        # Append dataset.json to the zip.
+        with zipfile.ZipFile(source, mode='a') as z:
+            z.writestr('dataset.json', json.dumps(metadata))
+        return
+
+    raise click.ClickException(
+        f'labels-only currently supports directory or zip sources, got: {source}'
+    )
+
+#----------------------------------------------------------------------------
+
+def open_dest(dest: str) -> Tuple[str, Callable[[str, Union[bytes, str]], None], Callable[[], None]]:
+    dest_ext = file_ext(dest)
+
+    if dest_ext == 'zip':
+        if os.path.dirname(dest) != '':
+            os.makedirs(os.path.dirname(dest), exist_ok=True)
+        zf = zipfile.ZipFile(file=dest, mode='w', compression=zipfile.ZIP_STORED)
+        def zip_write_bytes(fname: str, data: Union[bytes, str]):
+            zf.writestr(fname, data)
+        return '', zip_write_bytes, zf.close
+    else:
+        # If the output folder already exists, check that is is
+        # empty.
+        #
+        # Note: creating the output directory is not strictly
+        # necessary as folder_write_bytes() also mkdirs, but it's better
+        # to give an error message earlier in case the dest folder
+        # somehow cannot be created.
+        if os.path.isdir(dest) and len(os.listdir(dest)) != 0:
+            raise click.ClickException('--dest folder must be empty')
+        os.makedirs(dest, exist_ok=True)
+
+        def folder_write_bytes(fname: str, data: Union[bytes, str]):
+            os.makedirs(os.path.dirname(fname), exist_ok=True)
+            with open(fname, 'wb') as fout:
+                if isinstance(data, str):
+                    data = data.encode('utf8')
+                fout.write(data)
+        return dest, folder_write_bytes, lambda: None
+
+#----------------------------------------------------------------------------
+
+@click.command()
+@click.option('--source',     help='Input directory or archive name', metavar='PATH',   type=str, required=True)
+@click.option('--dest',       help='Output directory or archive name', metavar='PATH',  type=str, required=True)
+@click.option('--max-images', help='Maximum number of images to output', metavar='INT', type=int)
+@click.option('--transform',  help='Input crop/resize mode', metavar='MODE',            type=click.Choice(['center-crop', 'center-crop-wide']))
+@click.option('--resolution', help='Output resolution (e.g., 512x512)', metavar='WxH',  type=parse_tuple)
+@click.option('--labels-only', is_flag=True, help='Only generate dataset.json labels at source root; do not copy or re-encode images.')
+
+def main(
+    source: str,
+    dest: str,
+    max_images: Optional[int],
+    transform: Optional[str],
+    resolution: Optional[Tuple[int, int]],
+    labels_only: bool,
+):
+    """Convert an image dataset into a dataset archive usable with StyleGAN2 ADA PyTorch,
+    or only generate dataset.json labels in-place.
+
+    The input dataset format is guessed from the --source argument:
+
+    \b
+    --source *_lmdb/                    Load LSUN dataset
+    --source cifar-10-python.tar.gz     Load CIFAR-10 dataset
+    --source train-images-idx3-ubyte.gz Load MNIST dataset
+    --source path/                      Recursively load all images from path/
+    --source dataset.zip                Recursively load all images from dataset.zip
+
+    Specifying the output format and path:
+
+    \b
+    --dest /path/to/dir                 Save output files under /path/to/dir
+    --dest /path/to/dataset.zip         Save output files into /path/to/dataset.zip
+
+    The output dataset format can be either an image folder or an uncompressed zip archive.
+    Zip archives makes it easier to move datasets around file servers and clusters, and may
+    offer better training performance on network file systems.
+
+    Images within the dataset archive will be stored as uncompressed PNG.
+    Uncompresed PNGs can be efficiently decoded in the training loop.
+
+    Class labels are stored in a file called 'dataset.json' that is stored at the
+    dataset root folder.  This file has the following structure:
+
+    \b
+    {
+        "labels": [
+            ["00000/img00000000.png",6],
+            ["00000/img00000001.png",9],
+            ... repeated for every image in the datase
+            ["00049/img00049999.png",1]
+        ]
+    }
+
+    If the 'dataset.json' file cannot be found, class labels are determined from
+    top-level directory names.
+
+    Image scale/crop and resolution requirements:
+
+    Output images must be square-shaped and they must all have the same power-of-two
+    dimensions.
+
+    To scale arbitrary input image size to a specific width and height, use the
+    --resolution option.  Output resolution will be either the original
+    input resolution (if resolution was not specified) or the one specified with
+    --resolution option.
+
+    Use the --transform=center-crop or --transform=center-crop-wide options to apply a
+    center crop transform on the input image.  These options should be used with the
+    --resolution option.  For example:
+
+    \b
+    python scripts/dataset_tools.py --source data/cat_lmdb --dest data/cat_png \\
+        --transform=center-crop-wide --resolution=512x384
+    """
+
+    # Labels-only mode: just write dataset.json next to the original data.
+    if labels_only:
+        PIL.Image.init()
+        create_labels_only(source)
+        return
+
+    PIL.Image.init()
+
+    if dest == '':
+        raise click.ClickException('--dest output filename or directory must not be an empty string')
+
+    num_files, input_iter = open_dataset(source, max_images=max_images)
+    archive_root_dir, save_bytes, close_dest = open_dest(dest)
+
+    if resolution is None: resolution = (None, None)
+    transform_image = make_transform(transform, *resolution)
+
+    dataset_attrs = None
+
+    labels = []
+    for idx, image in tqdm(enumerate(input_iter), total=num_files):
+        idx_str = f'{idx:08d}'
+        archive_fname = f'{idx_str[:5]}/img{idx_str}.png'
+
+        # Apply crop and resize.
+        img = transform_image(image['img'])
+        if img is None:
+            continue
+
+        # Error check to require uniform image attributes across
+        # the whole dataset.
+        channels = img.shape[2] if img.ndim == 3 else 1
+        cur_image_attrs = {'width': img.shape[1], 'height': img.shape[0], 'channels': channels}
+        if dataset_attrs is None:
+            dataset_attrs = cur_image_attrs
+            width = dataset_attrs['width']
+            height = dataset_attrs['height']
+            if width != height:
+                raise click.ClickException(f'Image dimensions after scale and crop are required to be square.  Got {width}x{height}')
+            if dataset_attrs['channels'] not in [1, 3]:
+                raise click.ClickException('Input images must be stored as RGB or grayscale')
+            if width != 2 ** int(np.floor(np.log2(width))):
+                raise click.ClickException('Image width/height after scale and crop are required to be power-of-two')
+        elif dataset_attrs != cur_image_attrs:
+            err = [f'  dataset {k}/cur image {k}: {dataset_attrs[k]}/{cur_image_attrs[k]}' for k in dataset_attrs.keys()]
+            raise click.ClickException(f'Image {archive_fname} attributes must be equal across all images of the dataset.  Got:\n' + '\n'.join(err))
+
+        # Save the image as an uncompressed PNG.
+        img = PIL.Image.fromarray(img, {1: 'L', 3: 'RGB'}[channels])
+        image_bits = io.BytesIO()
+        img.save(image_bits, format='png', compress_level=0, optimize=False)
+        save_bytes(os.path.join(archive_root_dir, archive_fname), image_bits.getbuffer())
+        labels.append([archive_fname, image['label']] if image['label'] is not None else None)
+
+    metadata = {'labels': labels if all(x is not None for x in labels) else None}
+    save_bytes(os.path.join(archive_root_dir, 'dataset.json'), json.dumps(metadata))
+    close_dest()
+
+#----------------------------------------------------------------------------
+
+if __name__ == "__main__":
+    main()
+
+#----------------------------------------------------------------------------

scripts/evaluate_all.py

@@ -0,0 +1,130 @@
+"""Evaluate CIFAR-10 checkpoints with independent per-GPU seeds."""
+import argparse
+import csv
+import hashlib
+import io
+import json
+import os
+from pathlib import Path
+import subprocess
+import sys
+import time
+
+ROOT = Path(__file__).resolve().parents[1]
+
+# Method names match configuration prefixes and checkpoint filenames.
+CONFIGS = {
+    "baseline": "baseline_ar.yaml",
+    "tail_dropout": "tail_dropout_ar.yaml",
+    "flow_matching": "flow_matching_ar.yaml",
+    "flow_matching_tail_dropout": "flow_matching_tail_dropout_ar.yaml",
+    "flow_matching_tied_tail_dropout": "flow_matching_tied_tail_dropout_ar.yaml",
+    "war": "war_ar.yaml"
+}
+
+
+def sha256(path):
+    digest = hashlib.sha256()
+    with open(path, 'rb') as f:
+        for block in iter(lambda: f.read(8 * 1024 * 1024), b''):
+            digest.update(block)
+    return digest.hexdigest()
+
+
+def check_free(gpus):
+    output = subprocess.check_output(
+        ['nvidia-smi', '--query-gpu=index,memory.used,utilization.gpu',
+         '--format=csv,noheader,nounits'], text=True)
+    devices = {int(row[0]): (int(row[1]), int(row[2]))
+               for row in csv.reader(io.StringIO(output))}
+    busy = {gpu: devices.get(gpu) for gpu in gpus
+            if gpu not in devices or devices[gpu][0] > 100 or devices[gpu][1] > 0}
+    if busy:
+        raise RuntimeError(f'Requested GPUs are not free: {busy}')
+
+
+def save_summary(output, rows):
+    (output / 'results.json').write_text(json.dumps(rows, indent=2) + '\n')
+    if rows:
+        with (output / 'results.csv').open('w', newline='') as f:
+            writer = csv.DictWriter(f, fieldnames=list(rows[0]))
+            writer.writeheader()
+            writer.writerows(rows)
+
+
+def main():
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument('--gpus', default='0,1,2,3,4,5,6,7')
+    parser.add_argument('--output', default='outputs/cifar10-nonlegacy')
+    parser.add_argument('--methods', nargs='+', choices=CONFIGS,
+                        help='Optional subset, e.g. baseline war')
+    parser.add_argument('--resume', action='store_true', help='Reuse matching completed evaluations')
+    args = parser.parse_args()
+    gpus = [int(x) for x in args.gpus.split(',')]
+    if not gpus or len(gpus) != len(set(gpus)):
+        parser.error('GPU IDs must be nonempty and distinct')
+    configs = {key: ROOT / 'configs/cifar10' / name for key, name in CONFIGS.items()}
+    for config in configs.values():
+        if not config.is_file():
+            raise FileNotFoundError(config)
+    if args.methods:
+        if set(args.methods) - configs.keys():
+            parser.error('Unknown method')
+        configs = {key: path for key, path in configs.items() if key in args.methods}
+    checkpoints = {key: ROOT / 'checkpoints' / f'{key}.pt' for key in configs}
+    for checkpoint in checkpoints.values():
+        if not checkpoint.is_file():
+            raise FileNotFoundError(f'Place the downloaded checkpoint at {checkpoint}')
+    reference = ROOT / 'refs/cifar10-32x32_local.npz'
+    reference_hash = sha256(reference)
+    if reference_hash != '5493968c4f29b2878f64d6a640351f377e32e86553db56eca8147a0ae50d93f2':
+        raise ValueError('The CIFAR-10 FID reference checksum does not match')
+    output = Path(args.output).resolve()
+    output.mkdir(parents=True, exist_ok=True)
+    env = os.environ.copy()
+    env.update(CUDA_VISIBLE_DEVICES=','.join(map(str, gpus)), PYTHONDONTWRITEBYTECODE='1',
+               WANDB_MODE='disabled', OMP_NUM_THREADS='8', MKL_NUM_THREADS='8', TF_CPP_MIN_LOG_LEVEL='2')
+    rows = []
+    for key, config in configs.items():
+        checkpoint = checkpoints[key]
+        checkpoint_hash, config_hash = sha256(checkpoint), sha256(config)
+        case = output / key
+        result_path = case / 'metrics.json'
+        if result_path.exists() and not args.resume:
+            raise FileExistsError(result_path)
+        if not result_path.exists():
+            check_free(gpus)
+            case.mkdir(exist_ok=True)
+            cmd = [sys.executable, '-m', 'torch.distributed.run', '--standalone',
+                   f'--nproc-per-node={len(gpus)}', 'evaluate.py', '--config', str(config),
+                   '--checkpoint', str(checkpoint), '--num-samples', '50000', '--batch-size', '64',
+                   '--seed', '42', '--no-legacy', '--output', str(case), '--fid-reference', str(reference)]
+            request = dict(command=cmd, checkpoint_sha256=checkpoint_hash,
+                           config_sha256=config_hash, legacy=False)
+            (case / 'request.json').write_text(json.dumps(request, indent=2) + '\n')
+            print(f'START {key}', flush=True)
+            started = time.monotonic()
+            with (case / 'run.log').open('w') as log:
+                subprocess.run(cmd, cwd=ROOT, env=env, stdout=log, stderr=subprocess.STDOUT, check=True)
+            print(f'DONE {key} in {time.monotonic() - started:.1f}s', flush=True)
+        result = json.loads(result_path.read_text())
+        request = json.loads((case / 'request.json').read_text())
+        if request['checkpoint_sha256'] != checkpoint_hash:
+            raise ValueError(f'{key}: checkpoint changed since the recorded evaluation')
+        expected = dict(legacy=False, num_samples=50000, seed=42, world_size=len(gpus),
+                        batch_size=64, ar_ema=True, precision='32-true', decoder='single_step',
+                        config_sha256=config_hash, reference_sha256=reference_hash)
+        if any(result.get(name) != value for name, value in expected.items()):
+            raise ValueError(f'{key}: recorded evaluation settings do not match')
+        row = dict(method=key, nonlegacy_gfid=result['fid'], unique_images=result['unique_images'],
+                   num_samples=result['num_samples'], legacy=False,
+                   checkpoint_sha256=checkpoint_hash, config_sha256=config_hash,
+                   reference_sha256=reference_hash)
+        rows.append(row)
+        save_summary(output, rows)
+        print(f'RESULT {key}: {result["fid"]:.8f}', flush=True)
+    print(f'COMPLETE: {len(rows)} non-legacy evaluations', flush=True)
+
+
+if __name__ == '__main__':
+    main()

scripts/export_checkpoint.py

@@ -0,0 +1,47 @@
+"""Export inference states without optimizer, paths, or experiment tracker metadata."""
+import argparse
+import hashlib
+import json
+from pathlib import Path
+import sys
+sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
+import torch
+from omegaconf import OmegaConf
+from models import Encoder, Decoder, ARModel
+
+
+def sha256(path):
+    digest=hashlib.sha256()
+    with open(path,'rb') as f:
+        for block in iter(lambda:f.read(8*1024*1024),b''): digest.update(block)
+    return digest.hexdigest()
+
+
+def export(source, config, output):
+    source_state=torch.load(source, map_location='cpu', weights_only=True, mmap=True)
+    c=OmegaConf.load(config)
+    states={key:{name.removeprefix('_orig_mod.'): tensor for name,tensor in source_state[key].items()}
+            for key in ('enc','dec','ar_model','ar_model_ema')}
+    for cls,key in ((Encoder,'enc'),(Decoder,'dec'),(ARModel,'ar_model'),(ARModel,'ar_model_ema')):
+        # Assignment checks all keys/shapes without allocating a second model's storage.
+        with torch.device('meta'): model=cls(c)
+        model.load_state_dict(states[key],strict=True,assign=True)
+        del model
+    out=Path(output); out.parent.mkdir(parents=True,exist_ok=True)
+    if out.exists(): raise FileExistsError(out)
+    torch.save(dict(format_version=1,**states),out)
+    saved=torch.load(out,map_location='cpu',weights_only=True,mmap=True)
+    for key,state in states.items():
+        for name,tensor in state.items():
+            if not torch.equal(tensor,saved[key][name]): raise AssertionError((key,name))
+    return dict(filename=out.name,sha256=sha256(out),bytes=out.stat().st_size)
+
+
+def main():
+    p=argparse.ArgumentParser(description=__doc__)
+    p.add_argument('--source',required=True); p.add_argument('--config',required=True); p.add_argument('--output',required=True)
+    a=p.parse_args(); result=export(a.source,a.config,a.output)
+    Path(a.output+'.json').write_text(json.dumps(result,indent=2)+'\n'); print(json.dumps(result))
+
+
+if __name__=='__main__': main()

scripts/prepare_cifar10.py

@@ -0,0 +1,31 @@
+"""Convert the official CIFAR-10 training split into the wARTok PNG layout."""
+import argparse
+import json
+from pathlib import Path
+import numpy as np
+from PIL import Image
+from torchvision.datasets import CIFAR10
+
+
+def main():
+    p=argparse.ArgumentParser(description=__doc__)
+    p.add_argument('--download-dir',default='./data/downloads')
+    p.add_argument('--output',default='./data/cifar10')
+    p.add_argument('--download',action='store_true')
+    p.add_argument('--reference-images',help='Optional NPZ of the same 50K training images for ADM statistics')
+    a=p.parse_args(); out=Path(a.output)
+    if out.exists() and any(out.iterdir()): p.error('output must be empty to prevent mixing splits')
+    ds=CIFAR10(a.download_dir,train=True,download=a.download)
+    assert len(ds)==50000
+    labels=[]
+    for i,(image,label) in enumerate(ds):
+        name=f'{i//1000:05d}/img{i:08d}.png'
+        path=out/name; path.parent.mkdir(parents=True,exist_ok=True)
+        image.save(path); labels.append([name,label])
+    (out/'dataset.json').write_text(json.dumps({'labels':labels}))
+    if a.reference_images:
+        with open(a.reference_images,'xb') as f: np.savez_compressed(f,ds.data)
+    print(f'Wrote {len(ds)} training images to {out}')
+
+
+if __name__=='__main__': main()

train.py

@@ -0,0 +1,1209 @@
+import itertools
+import math
+import os
+import random
+import subprocess
+import sys
+from pathlib import Path
+from typing import Iterable, Iterator, List, NamedTuple
+os.environ.setdefault("WANDB_MODE", "disabled")
+import wandb
+import numpy as np
+import torch
+from torch import nn
+import torch.nn.functional as F
+import torchvision.utils
+from einops import rearrange
+from lightning import Fabric
+from wandb.integration.lightning.fabric import WandbLogger
+from dataset import ImageFolderDataset
+from gan_loss import GANLoss
+from lpips import LPIPS
+from models import ARModel, Encoder, Decoder, Linear
+from tqdm import tqdm
+from omegaconf import OmegaConf
+from torch.utils.data import DataLoader
+import copy
+import torch.distributed as dist
+from glob import glob
+try:
+    from thop import profile
+    from thop.vision.basic_hooks import count_linear
+except ImportError:
+    profile = None
+    count_linear = None
+
+
+torch.backends.cudnn.benchmark = True
+torch.backends.cudnn.deterministic = False
+
+# Phase utils
+class Target(NamedTuple):
+    DO_AE: bool = False
+    DO_L2: bool = False
+    DO_L1: bool = False
+    DO_LPIPS: bool = False
+    DO_GAN_G: bool = False
+    DO_GAN_D: bool = False
+    DO_PRIOR_ENC: bool = False
+    DO_PRIOR_AR: bool = False
+
+class Phase(NamedTuple):
+    num_steps: int
+    targets: List[Target]       
+    internal_steps: List[int]   
+
+# Parse phases such as "1000:DO_L2-DO_PRIOR_AR:1 7000:DO_PRIOR_AR:1".
+def parse_phases(phases_str):
+    phases = [] # list of Phase objects
+    for phase_str in phases_str.split(' '):
+        num_steps, targets_str, internal_steps_str = phase_str.split(':')
+        num_steps = int(num_steps)
+        targets = [Target(**{k: True for obj in target_str.split(',') for k in obj.split('-') }) for target_str in targets_str.split(',')]
+        internal_steps = [int(step) for step in internal_steps_str.split(',')]
+        phases.append(Phase(num_steps, targets, internal_steps))
+    return phases
+
+def get_phase(global_step, phases, phase_step_accum, config):
+    target=None
+    for phase_idx, phase_step in enumerate(phase_step_accum):
+        if global_step <= phase_step:
+            internel_step = (global_step - phase_step_accum[phase_idx-1]) if phase_idx > 0 else global_step
+            internel_accumulate = list(itertools.accumulate(phases[phase_idx].internal_steps))
+            internel_step = internel_step % internel_accumulate[-1]
+
+            for inner_idx in range(len(internel_accumulate)):
+                if internel_step < internel_accumulate[inner_idx]:
+                    target = phases[phase_idx].targets[inner_idx]
+                    break
+            if target is not None:
+                break
+
+    DO_AE = any([target.DO_L2, target.DO_L1, target.DO_LPIPS, target.DO_GAN_G])
+
+    # assert inconsistency configuration
+    assert not (DO_AE and  not config.train_ae), f"train_ae={config.train_ae} is inconsistent with DO_AE={DO_AE} in phases"
+    assert not (target.DO_PRIOR_AR and  not config.train_ar), f"train_ar={config.train_ar} is inconsistent with DO_PRIOR_AR={target.DO_PRIOR_AR} in phases"
+    assert not ((target.DO_GAN_G or target.DO_GAN_D) and  not config.use_gan_loss), f"use_gan_loss={config.use_gan_loss} is inconsistent with DO_GAN_G or DO_GAN_D={(target.DO_GAN_G or target.DO_GAN_D)} in phases"
+    assert not (target.DO_PRIOR_ENC and  not config.use_prior_loss), f"use_prior_loss={config.use_prior_loss} is inconsistent with DO_PRIOR_ENC={target.DO_PRIOR_ENC} in phases"
+    assert not (target.DO_LPIPS and  not config.use_lpips_loss), f"use_lpips_loss={config.use_lpips_loss} is inconsistent with DO_LPIPS={target.DO_LPIPS} in phases"
+
+    # update target with global_step
+    target = Target(DO_L1=target.DO_L1, 
+                    DO_L2=target.DO_L2, 
+                    DO_LPIPS=target.DO_LPIPS, 
+                    DO_GAN_G=target.DO_GAN_G and (global_step >= config.gan_start),
+                    DO_GAN_D=target.DO_GAN_D and (global_step >= config.gan_start),
+                    DO_PRIOR_ENC=target.DO_PRIOR_ENC and (global_step >= config.prior_start),
+                    DO_PRIOR_AR=target.DO_PRIOR_AR,
+                    DO_AE=DO_AE)
+    return phase_idx, inner_idx, target, internel_step
+
+class InfiniteIterator(Iterator):
+    def __init__(self, iterable: Iterable):
+        self.iterable = iterable
+        self._it = iter(iterable)
+
+    def __iter__(self):
+        return self
+
+    def __next__(self):
+        try:
+            return next(self._it)
+        except StopIteration:
+            self._it = iter(self.iterable)
+            return next(self._it)
+
+# Image Processing Utils
+def patchify(x, patch_size):
+    x = rearrange(x, 'b c (h p1) (w p2) -> b (h w) (p1 p2 c)', p1=patch_size, p2=patch_size)
+    return x
+
+def unpatchify(x, image_size, patch_size):
+    x = rearrange(x, 'b (h w) (p1 p2 c) -> b c (h p1) (w p2)', p1=patch_size, p2=patch_size, h=image_size//patch_size, w=image_size//patch_size)
+    return x
+
+def img_unint8_to_norm(x):
+    return x.float() / 127.5 - 1.0
+
+def img_denormalize(x):
+    return x.clamp(-1,1) * 0.5 + 0.5
+
+def img_norm_to_uint8(x):
+    return torch.clamp(127.5 * x + 128.0, 0, 255).byte()
+
+# -----------------------------
+# Tail-dropout + noise-query utils
+# -----------------------------
+
+def _get_decoder_context_type(dec: nn.Module) -> str:
+    # `dec` is Decoder, `dec.dec` is Transformer
+    ctx = getattr(getattr(dec, "dec", None), "context_type", None)
+    return ctx if ctx is not None else "none"
+
+def _get_decoder_query_len(dec: nn.Module) -> int:
+    return int(getattr(getattr(dec, "dec", None), "max_seq_len", 0))
+
+def _t_to_k(t: torch.Tensor, K: int, t2k: float = 1.0) -> torch.Tensor:
+    """
+    Map t in [0,1] to discrete k in [0, K-1].
+    """
+    if K <= 1:
+        return torch.zeros_like(t, dtype=torch.long)
+    t_tmp = (t2k * t).clamp(0.0, 1.0)
+    k = torch.floor(t_tmp * K).to(torch.long).clamp(0, K - 1)
+    return k
+
+class DiTiScheduler:
+    """
+    Piecewise linear time scheduler from Selftok, adapted for wARTok.
+    """
+    def __init__(
+        self,
+        n_timesteps: int,
+        K: int,
+        stages: str | List[int],
+        k_per_stage: str | List[int],
+        reverse_t: bool = False,
+    ):
+        self.n_timesteps = n_timesteps
+        self.K = K
+        self.reverse_t = reverse_t
+
+        if isinstance(stages, str):
+            stages = [int(s.strip()) for s in stages.split(',')]
+        if isinstance(k_per_stage, str):
+            k_per_stage = [int(k.strip()) for k in k_per_stage.split(',')]
+        self.stages = stages
+        
+        # Scaling logic: if sum(k_per_stage) != K, scale it
+        total_k = sum(k_per_stage)
+        if total_k != K:
+            print(f"[DiTiScheduler] Warning: sum(k_per_stage)={total_k} does not match K={K}. Scaling k_per_stage...")
+            scale_factor = K / total_k
+            scaled_k = [int(k * scale_factor) for k in k_per_stage]
+            # Fix rounding errors to ensure sum is exactly K
+            diff = K - sum(scaled_k)
+            if diff != 0:
+                scaled_k[-1] += diff
+            k_per_stage = scaled_k
+            print(f"[DiTiScheduler] Scaled k_per_stage: {k_per_stage}")
+            
+        self.k_per_stage = k_per_stage
+        
+        self.t_to_idx = torch.zeros(n_timesteps).long()
+        
+        # Precompute lookup table
+        current_stage = 0
+        sum_indices = 0
+        # stages should be cumulative, e.g., 200, 400, 600...
+        
+        for t in range(n_timesteps):
+            # Check if we moved to next stage
+            if current_stage < len(self.stages) and t >= self.stages[current_stage]:
+                sum_indices += self.k_per_stage[current_stage]
+                current_stage += 1
+            
+            if current_stage >= len(self.stages):
+                 idx = self.K - 1
+            else:
+                prev_stage_t = self.stages[current_stage - 1] if current_stage > 0 else 0
+                current_stage_t = self.stages[current_stage]
+                
+                current_steps = float(current_stage_t - prev_stage_t)
+                current_k = float(self.k_per_stage[current_stage])
+                
+                t_adj = t - prev_stage_t
+                idx = int(float(t_adj) / current_steps * current_k + sum_indices)
+            
+            self.t_to_idx[t] = idx
+
+    def to_indices(self, t: torch.Tensor) -> torch.Tensor:
+        """
+        Convert t in [0, 1] to indices [0, K-1] using the lookup table.
+        """
+        if self.reverse_t:
+            t = 1.0-t
+        device = t.device
+        # Map [0, 1] to [0, n_timesteps]
+        t_idx = (t * self.n_timesteps).long().clamp(0, self.n_timesteps - 1)
+        # Move the lookup table to the input device.
+        return self.t_to_idx.to(device)[t_idx].clamp(0, self.K - 1)
+
+
+def build_taildrop_attn_mask(
+    *,
+    t: torch.Tensor,
+    context_len: int,
+    query_len: int,
+    context_type: str,
+    t2k: float = 1.0,
+    device=None,
+    diti_scheduler: DiTiScheduler = None,
+) -> torch.Tensor | None:
+    """
+    Build a boolean attention mask (True for keep, False for masked).
+
+    - cross_attn: returns (B, 1, query_len, context_len)
+    - concat:     returns (B, 1, total_len, total_len) where total_len = query_len + context_len
+    """
+    device = device or t.device
+    B = t.shape[0]
+    if context_len <= 0:
+        return None
+
+    if diti_scheduler is not None:
+        k = diti_scheduler.to_indices(t) # (B,)
+    else:
+        k = _t_to_k(t, context_len, t2k=t2k)  # (B,)
+        
+    keep = (torch.arange(context_len, device=device)[None, :] <= k[:, None])  # (B, context_len)
+
+
+    if context_type == "cross_attn":
+        # (B, 1, query_len, context_len)
+        return keep[:, None, :].unsqueeze(1).expand(B, 1, query_len, context_len)
+
+    if context_type == "concat":
+        total_len = query_len + context_len
+        keep_total = torch.cat(
+            [torch.ones((B, query_len), device=device, dtype=torch.bool), keep],
+            dim=1,
+        )  # (B, total_len)
+        # (B, 1, total_len, total_len)
+        return keep_total[:, None, :].unsqueeze(1).expand(B, 1, total_len, total_len)
+
+    return None
+
+def make_noise_image_patch_query(x: torch.Tensor, t: torch.Tensor, query_len: int) -> torch.Tensor:
+    """
+    Replace dummy query with a noised version of the (clean) image patches x.
+    x: (B, L, D)
+    t: (B,)
+    Returns noised patches, noise, and velocity, each of shape (B, L, D).
+    """
+    B, L, D = x.shape
+    noise = torch.randn_like(x)
+    t = t.view(B, 1, 1).to(x.dtype)
+    x_noised = (1.0 - t) * x + t * noise
+    return x_noised, noise, noise-x
+
+# Seed utils
+def seed_everything(seed):
+    # manually call these to ensure reproducibility
+    random.seed(seed)
+    np.random.seed(seed)
+    torch.manual_seed(seed)
+    torch.cuda.manual_seed_all(seed)
+    os.environ['PYTHONHASHSEED'] = str(seed)
+    # fabric seed everything
+    Fabric.seed_everything(seed, workers=True)
+
+def worker_init_fn(worker_id):
+    worker_seed = torch.initial_seed() % 2**32
+    np.random.seed(worker_seed)
+    random.seed(worker_seed)
+
+
+def make_worker_init_fn(base_seed: int):
+    base_seed = int(base_seed) % (2**32)
+
+    def _init(worker_id: int):
+        rank = dist.get_rank() if dist.is_available() and dist.is_initialized() else 0
+        worker_seed = (base_seed + worker_id + 1000 * rank) % (2**32)
+        np.random.seed(worker_seed)
+        random.seed(worker_seed)
+        torch.manual_seed(worker_seed)
+
+    return _init
+
+
+# FID Utils                    
+def adm_fid_evaluator(sample_cached_path, gt_cache_path, config, fabric):
+    if not os.path.exists(gt_cache_path):
+        raise FileNotFoundError(f"Ground-truth cache not found: {gt_cache_path}")
+    if not os.path.exists(sample_cached_path):
+        raise FileNotFoundError(f"Sample cache not found: {sample_cached_path}")
+
+    fid_script = str(Path(__file__).with_name('calc_fid.py'))
+    env = os.environ.copy()
+    cmd = [sys.executable,fid_script,"--ref_batch",gt_cache_path,"--sample_batch",sample_cached_path,"--batch_size",str(config.eval_batch_size),]
+    fabric.print(f"Running FID evaluation via {fid_script}...")
+
+    FID = None
+    process = subprocess.Popen(cmd,stdout=subprocess.PIPE,stderr=subprocess.STDOUT,text=True,bufsize=1,env=env,)
+    assert process.stdout is not None
+    for line in process.stdout:
+        line = line.rstrip("\n")
+        if line:
+            fabric.print(line, flush=True)
+        if line.startswith("FID_RESULT:"):
+            try:
+                FID = float(line.split("FID_RESULT:")[1].strip())
+            except ValueError:
+                pass
+    retcode = process.wait()
+    if retcode != 0 or FID is None or not math.isfinite(FID):
+        raise RuntimeError(f"FID evaluation failed: exit={retcode}, result={FID}")
+
+    return FID
+
+
+# Training Utils
+@torch.no_grad()
+@torch._dynamo.disable
+def ema_update(model, ema_model, ema_rate):
+    if model is None or ema_model is None:
+        return
+    
+    for p,ema_p in zip(model.parameters(),ema_model.parameters()):
+            ema_p.copy_(p.detach().lerp(ema_p, ema_rate))
+    
+
+
+@torch.no_grad()
+@torch._dynamo.disable
+def sampling(enc, dec, ar_model,  bz, class_label=None, temperature=1.0, topK=0, topP=0., cfg=0.0, cfg_schedule=None, cfg_power=None, cache_kv=False, config=None,t=None,query_tokens=None, **params):
+    token_ids = ar_model.sampling(bz, class_label, temperature, topK, topP, cfg, cfg_schedule, cfg_power, cache_kv)
+    quant = enc.quantizer.get_codes_w_indices(token_ids)
+    decoded = dec(quant, class_label, attn_mask=None, query_tokens=query_tokens,t=t)
+    return decoded
+
+@torch.no_grad()
+@torch._dynamo.disable
+def reconstruction(enc, dec, x, labels, return_idx: bool = False, query_tokens=None,t=None):
+    quant, idx, _, _ = enc(x,labels, training=False)
+    decoded = dec(quant, labels, attn_mask=None, query_tokens=query_tokens, t=t)
+    return decoded, idx
+
+@torch.no_grad()
+@torch._dynamo.disable
+def euler_sample_loop(enc, dec, ar_model,  bz, class_label=None, temperature=1.0, topK=None, topP=None, cfg=1.0, cfg_schedule=None, cfg_power=None, cache_kv=False, config=None, diti_scheduler=None):
+    """
+    Euler solver for Flow Matching with dynamic token masking (DiTi).
+    """
+    if config is None:
+        raise ValueError("config is required for euler_sample_loop")
+    # Derive device / shape from config to avoid relying on outer-scope globals.
+    device = next(dec.parameters()).device
+    patch_dim = config.patch_dim
+    token_ids = ar_model.sampling(bz, class_label, temperature, topK, topP, cfg, cfg_schedule, cfg_power, cache_kv)
+    quant = enc.quantizer.get_codes_w_indices(token_ids)
+    query_len = config.x_len
+
+    num_steps = 25
+    x = torch.randn((bz, query_len, patch_dim), device=device)
+    context_len = quant.shape[1]
+    timesteps = torch.linspace(1.0, 0.0, num_steps + 1, device=device)
+    dt = 1.0 / num_steps
+    
+    for i in range(num_steps):
+        t_curr = timesteps[i]
+        t_tensor = torch.full((bz,), t_curr, device=device)
+        
+        attn_bias = None
+        if config.use_tail_dropout and config.tied_timestep:
+            if diti_scheduler is not None:
+                attn_bias = build_taildrop_attn_mask(
+                    t=t_tensor, context_len=context_len, query_len=query_len,
+                    context_type=config.context_type, device=device, diti_scheduler=diti_scheduler,
+                )
+            else:
+                attn_bias = build_taildrop_attn_mask(
+                    t=t_tensor, context_len=context_len, query_len=query_len,
+                    context_type=config.context_type, t2k=config.tail_dropout_t2k, device=device,
+                )
+        decoded = dec(quant, class_label, attn_mask=attn_bias, query_tokens=x, t=t_tensor)
+        if config.l2_v_target:
+            v_pred = decoded
+            x = x - v_pred * dt
+        else:
+            x0_pred = decoded
+            t_frac = dt/t_tensor
+            t_frac = t_frac[:,None,None]
+            x = x * (1-t_frac) +  x0_pred * t_frac
+    return x
+
+@torch.no_grad()
+@torch._dynamo.disable
+def euler_recon_loop(enc, dec, x, class_label=None, config=None, diti_scheduler=None):
+    """
+    Euler solver for Flow Matching with dynamic token masking (DiTi).
+    """
+    if config is None:
+        raise ValueError("config is required for euler_recon_loop")
+    # Derive device / shape from config to avoid relying on outer-scope globals.
+    device = next(dec.parameters()).device
+    B = x.shape[0]
+
+    quant, idx, _, _ = enc(x,class_label, training=False)
+    context_len = quant.shape[1]
+    query_len = x.shape[1]
+
+    num_steps = 25
+    timesteps = torch.linspace(1.0, 0.0, num_steps + 1, device=device)
+    dt = 1.0 / num_steps
+    xt = torch.randn_like(x)
+    for i in range(num_steps):
+        t_tensor = torch.full((B,), timesteps[i], device=device)
+        attn_bias = None
+        if config.use_tail_dropout and config.tied_timestep:
+            if diti_scheduler is not None:
+                attn_bias = build_taildrop_attn_mask(
+                    t=t_tensor, context_len=context_len, query_len=query_len,
+                    context_type=config.context_type, device=device, diti_scheduler=diti_scheduler,
+                )
+            else:
+                attn_bias = build_taildrop_attn_mask(
+                    t=t_tensor, context_len=context_len, query_len=query_len,
+                    context_type=config.context_type, t2k=config.tail_dropout_t2k, device=device,
+                )
+        decoded = dec(quant, class_label, attn_mask=attn_bias, query_tokens=xt, t=t_tensor)
+        if config.l2_v_target:
+            v_pred = decoded
+            xt = xt - v_pred * dt
+        else:
+            x0_pred = decoded
+            t_frac = dt/t_tensor
+            t_frac = t_frac[:,None,None]
+            xt = xt * (1-t_frac) +  x0_pred * t_frac
+    return xt, idx
+
+def toogle_train_eval(model, train=True):
+    if model is None: return
+    if train:
+        model.train()
+    else:
+        model.eval()
+
+def toogle_require_grad(model, grads=True):
+    if model is None:
+        return
+    model.requires_grad_(grads)
+
+def zero_nan_gradients(model):
+    for name, param in model.named_parameters():
+        if param.grad is not None:
+            param.grad.nan_to_num_(nan=0.0, posinf=1e5, neginf=-1e5)
+
+# Model Prams and GFLOPS Utils
+def print_model_info(fabric,model, inputs, name="Model"):
+    if model is None: return
+    if profile is None:
+        fabric.print(f"[{name}] thop not installed, skipping GFLOPS calculation.")
+        params = sum(p.numel() for p in model.parameters())
+        fabric.print(f"[{name}] Params: {params/1e6:.2f}M")
+        return
+
+    # Move inputs to device
+    device = next(model.parameters()).device
+    inputs = tuple(i.to(device) if isinstance(i, torch.Tensor) else i for i in inputs)
+
+    try:
+        model_copy = copy.deepcopy(model)
+        model_copy.eval()
+        custom_ops = {}
+        if count_linear is not None:
+            custom_ops[Linear] = count_linear
+
+        macs, params = profile(model_copy, inputs=inputs, custom_ops=custom_ops, verbose=False)
+        gflops = macs * 2 / 1e9
+        fabric.print(f"[{name}] Params: {params/1e6:.2f}M, GFLOPS: {gflops:.2f}")
+        del model_copy
+    except Exception as e:
+        fabric.print(f"[{name}] Failed to calculate GFLOPS: {e}")
+        params = sum(p.numel() for p in model.parameters())
+        fabric.print(f"[{name}] Params: {params/1e6:.2f}M")
+        
+def print_model_stats(fabric,config, enc, dec, ar_model):
+    x_len = int(config.x_len)
+    patch_dim = int(config.patch_dim)
+    num_classes = int(config.num_classes)
+    z_len = int(config.z_len)
+    z_dim = int(config.z_dim)
+    codebook_size = int(config.codebook_size)
+    
+    x_dummy = torch.randn(1, x_len, patch_dim)
+    label_dummy = torch.zeros(1, num_classes)
+    
+    print_model_info(fabric,enc, (x_dummy, label_dummy), "Encoder")
+    
+    z_dummy = torch.randn(1, z_len, z_dim)
+    # If Decoder uses noise query, it requires query_tokens (and optionally t) at forward time.
+    query_type = getattr(getattr(dec, "dec", None), "query_type", "learnable")
+    if query_type == "noise":
+        query_len = int(getattr(config, "x_len", z_len))
+        query_dummy = torch.randn(1, query_len, patch_dim)
+        t_dummy = torch.zeros(1)
+        print_model_info(fabric, dec, (z_dummy, label_dummy, None, query_dummy, t_dummy), "Decoder")
+    else:
+        print_model_info(fabric,dec, (z_dummy, label_dummy), "Decoder")
+    
+    if ar_model is not None:
+        idx_dummy = torch.randint(0, codebook_size, (1, z_len))
+        print_model_info(fabric,ar_model, (idx_dummy, label_dummy), "ARModel")
+
+
+def train(config):
+    if config.use_tf32:
+        torch.backends.cuda.matmul.allow_tf32 = True
+        torch.backends.cudnn.allow_tf32 = True
+    # Parse devices
+    devices_value = config.devices
+    if isinstance(devices_value, str) and "," in devices_value:
+        devices_for_fabric = [int(d.strip()) for d in devices_value.split(',')]
+        num_devices = len(devices_for_fabric)
+    else:
+        try:
+            num_devices = int(devices_value)
+            devices_for_fabric = num_devices
+        except (TypeError, ValueError):
+            raise ValueError(f"Invalid value for devices: {devices_value}. Expected an integer or comma-separated string like '0,1,2'.")
+
+    if config.batch_size % num_devices or config.batch_size < num_devices:
+        raise ValueError("Global batch_size must be divisible by the number of devices")
+    if config.eval_data_dir is not None and not Path(config.eval_fid_ref_path).is_file():
+        raise FileNotFoundError("Prepare eval_fid_ref_path before training, or set eval_data_dir=null to disable periodic evaluation")
+
+    # prepare seed and save dir
+    if config.seed is not None:
+        seed = int(config.seed) if config.seed is not None else 0
+    seed_everything(seed)
+    dl_generator = torch.Generator()
+    dl_generator.manual_seed(seed)
+    worker_init = make_worker_init_fn(seed)
+
+    if "EXPERIMENT_SAVE_DIR" in os.environ:
+        save_dir = os.environ["EXPERIMENT_SAVE_DIR"]
+    else:
+        experiment_index = len(glob(f"{config.save_dir}/*"))
+        save_dir = config.save_dir+f"/{experiment_index:03d}"
+        os.environ["EXPERIMENT_SAVE_DIR"] = save_dir
+    logger = WandbLogger(project=config.wandb_project, save_dir=save_dir, name=config.wandb_name)
+
+    # initialize fabric
+    fabric = Fabric(accelerator="gpu", devices=devices_for_fabric, precision=config.precision, loggers=logger)
+    fabric.launch()
+    fabric.print(config)
+    fabric.print(f"Global seed set to {config.seed}")
+
+    # Initialize DiTi Scheduler if config provided
+    diti_scheduler = None
+    if getattr(config, "time_stages", None) and getattr(config, "k_per_stage", None):
+        fabric.print("Initializing DiTi Scheduler for piecewise linear time-to-k mapping...")
+        diti_scheduler = DiTiScheduler(
+            n_timesteps=1000, 
+            K=int(config.z_len),
+            stages=config.time_stages,
+            k_per_stage=config.k_per_stage,
+            reverse_t=config.reverse_t
+        )
+
+    # Models and optimizers
+    enc = Encoder(config).to(fabric.device) # always initialize encoder
+    dec = Decoder(config).to(fabric.device) 
+    if not config.train_ae:
+        enc = enc.requires_grad_(False)
+        dec = dec.requires_grad_(False)
+    ar_model = ARModel(config).to(fabric.device) if config.train_ar else None 
+    gan_loss = GANLoss(**config.GANLoss).to(fabric.device) if config.use_gan_loss else None
+    perceptual_loss = LPIPS().to(fabric.device).eval().requires_grad_(False) if config.use_lpips_loss else None
+
+
+    # Load Ckpt
+    if config.ae_ckpt_path != "":
+        fabric.print(f'resume ae from ckpt {config.ae_ckpt_path}')
+
+        ckpt_states = fabric.load(config.ae_ckpt_path)
+        if enc is not None and ckpt_states.get('enc', None) is not None:
+            enc.load_state_dict({k.split('_orig_mod.')[-1]: v for k, v in ckpt_states['enc'].items() })
+            fabric.print("loaded enc")
+        if dec is not None and ckpt_states.get('dec', None) is not None:
+            dec.load_state_dict({k.split('_orig_mod.')[-1]: v for k, v in ckpt_states['dec'].items() })
+            fabric.print("loaded dec")
+
+        if gan_loss is not None and ckpt_states.get('gan_loss', None) is not None:
+            gan_loss.load_state_dict({k.split('_orig_mod.')[-1]: v for k, v in ckpt_states['gan_loss'].items() })
+
+    if config.ar_ckpt_path!="":
+        fabric.print(f'resume ar from ckpt {config.ar_ckpt_path}')
+        ar_ckpt_states = fabric.load(config.ar_ckpt_path)
+        if config.load_ar_model and ar_model is not None and ar_ckpt_states.get('ar_model', None) is not None:
+            ar_model.load_state_dict({k.split('_orig_mod.')[-1]: v for k, v in ar_ckpt_states['ar_model'].items() } )
+            fabric.print("loaded ar")
+
+    # for ema
+    ae_ema_rate = math.pow(0.5,1/config.ae_ema_halflife) if config.ae_ema_halflife>0 else 0.
+    ar_ema_rate = math.pow(0.5,1/config.ar_ema_halflife) if config.ar_ema_halflife>0 else 0.
+    enc_ema = copy.deepcopy(enc).requires_grad_(False).eval() if ae_ema_rate>0 else None
+    dec_ema = copy.deepcopy(dec).requires_grad_(False).eval() if ae_ema_rate>0 else None
+    freeze_ar_model = copy.deepcopy(ar_model).requires_grad_(False).eval() if config.train_ar else None
+    ar_model_ema = copy.deepcopy(ar_model).requires_grad_(False).eval() if config.train_ar and ar_ema_rate>0 else None
+    
+    if config.ae_ckpt_path != "":
+        fabric.print(f'resume ae ema from ckpt {config.ae_ckpt_path}')
+        ckpt_states = fabric.load(config.ae_ckpt_path)
+        if enc_ema is not None and ckpt_states.get('enc_ema', None) is not None:
+            enc_ema.load_state_dict({k.split('_orig_mod.')[-1]: v for k, v in ckpt_states['enc_ema'].items() })
+        if dec_ema is not None and ckpt_states.get('dec_ema', None) is not None:
+            dec_ema.load_state_dict({k.split('_orig_mod.')[-1]: v for k, v in ckpt_states['dec_ema'].items() })
+        print("loaded ae ema")
+
+    if config.ar_ckpt_path!="" and config.load_ar_model_ema and ar_model_ema is not None and ar_ckpt_states.get('ar_model_ema', None) is not None:
+        ar_model_ema.load_state_dict({k.split('_orig_mod.')[-1]: v for k, v in ar_ckpt_states['ar_model_ema'].items() })
+        print("loaded ar ema")
+    if config.ar_ckpt_path!="" and config.load_ar_model_ema and freeze_ar_model is not None and ar_ckpt_states.get('freeze_ar_model', None) is not None:
+        freeze_ar_model.load_state_dict({k.split('_orig_mod.')[-1]: v for k, v in ar_ckpt_states['freeze_ar_model'].items() })
+        print("loaded frozen ar")
+
+    # for inference
+    enc_raw = enc
+    dec_raw = dec
+    ar_model_raw = ar_model
+    enc_ema_raw = enc_ema
+    dec_ema_raw = dec_ema
+    ar_model_ema_raw = ar_model_ema
+
+    if fabric.is_global_zero:
+        fabric.print("Calculating model stats...")
+        print_model_stats(fabric,config, enc_raw, dec_raw, ar_model_raw)
+        root_dir = Path(save_dir)
+        root_dir.mkdir(exist_ok=True, parents=True)
+        OmegaConf.save(config, root_dir / "config.yaml", resolve=True)
+        img_dir  = Path(save_dir+ "/images")
+        img_dir.mkdir(exist_ok=True, parents=True)
+        ckpt_dir  = Path(save_dir+ "/ckpts")
+        ckpt_dir.mkdir(exist_ok=True, parents=True)
+    
+    fabric.barrier()
+
+    if config.torch_compile:
+        if enc is not None and config.compile_ae:
+            enc = torch.compile(enc)
+        if enc_ema is not None and config.compile_ae:
+            enc_ema = torch.compile(enc_ema)
+        if dec is not None and config.compile_ae:
+            dec = torch.compile(dec)
+        if dec_ema is not None and config.compile_ae:
+            dec_ema = torch.compile(dec_ema)
+        if ar_model is not None and config.compile_ar:
+            ar_model = torch.compile(ar_model)
+        if ar_model_ema is not None and config.compile_ar:
+            ar_model_ema = torch.compile(ar_model_ema)
+        if freeze_ar_model is not None and config.compile_ar:
+            freeze_ar_model= torch.compile(freeze_ar_model)
+        if gan_loss is not None and config.compile_gan:
+            gan_loss = torch.compile(gan_loss)
+
+    opt_enc = torch.optim.AdamW(enc.parameters() , lr=config.lr_enc, betas=(config.beta1, config.beta2), weight_decay=config.weight_decay_enc) if config.train_ae else None
+    opt_ar = torch.optim.AdamW(ar_model.parameters(), lr=config.lr_ar, betas=(config.ar_beta1, config.ar_beta2), weight_decay=config.weight_decay_ar) if config.train_ar else None
+    opt_dec = torch.optim.AdamW(dec.parameters(), lr=config.lr_dec, betas=(config.beta1, config.beta2), weight_decay=config.weight_decay_dec) if config.train_ae else None
+    opt_gan_loss = torch.optim.AdamW(gan_loss.parameters(), lr=config.lr_gan_loss, betas=(config.beta1, config.beta2), weight_decay=config.weight_decay_gan) if config.use_gan_loss else None
+    if  config.resume_optimizer:
+        if config.train_ae:       
+            fabric.print("resumed AE optimizer")
+            opt_enc.load_state_dict(ckpt_states['opt_enc'])
+            opt_dec.load_state_dict(ckpt_states['opt_dec'])
+        if config.train_ar:
+            fabric.print("resumed AR optimizer")
+            ar_optimizer_states = ar_ckpt_states if config.ar_ckpt_path else ckpt_states
+            opt_ar.load_state_dict(ar_optimizer_states['opt_ar'])
+        if config.use_gan_loss:
+            fabric.print("resumed GAN optimizer")
+            opt_gan_loss.load_state_dict(ckpt_states['opt_gan_loss'])
+
+    # Datasets
+    dataset = ImageFolderDataset(path=config['data_dir'],
+        resolution=config.image_size,
+        use_label=config.use_label,
+        max_size=None,
+        xflip=config.xflip,
+        crop_type=config.crop_type,
+        deterministic_crop=bool(getattr(config, "deterministic_crop", False)),
+        crop_seed=int(config.seed) if config.seed is not None else 0,
+    )
+    train_loader = DataLoader(
+        dataset,
+        batch_size=config.batch_size//num_devices,
+        shuffle=True,
+        num_workers=config.num_workers,
+        pin_memory=True,
+        drop_last=True,
+        persistent_workers=config.num_workers > 0,
+        worker_init_fn=worker_init,
+        generator=dl_generator,
+    )
+    eval_loader=None
+    if config.eval_data_dir is not None:
+        eval_dataset = ImageFolderDataset(path=config['eval_data_dir'],
+            resolution=config.image_size,
+            use_label=config.use_label,
+            max_size=None,
+            xflip=False,
+            crop_type=config.eval_crop_type,
+            deterministic_crop=bool(getattr(config, "deterministic_crop", False)),
+            crop_seed=int(config.seed) if config.seed is not None else 0,
+        )
+        eval_loader = DataLoader(
+            eval_dataset,
+            batch_size=config.batch_size//num_devices,
+            shuffle=False,
+            pin_memory=True,
+            drop_last=False,
+            worker_init_fn=worker_init,
+            generator=dl_generator,
+        )
+    
+    # Setup Model, Optimizer, DataLoader with Fabric
+    if config.train_ae:
+        enc, opt_enc = fabric.setup(enc, opt_enc)
+        dec, opt_dec = fabric.setup(dec, opt_dec)
+        if enc_ema is not None:
+            enc_ema = enc_ema.to(fabric.device)
+        if dec_ema is not None:
+            dec_ema = dec_ema.to(fabric.device)
+    if config.train_ar:
+        ar_model, opt_ar = fabric.setup(ar_model, opt_ar)
+        if ar_model_ema is not None:
+            ar_model_ema = ar_model_ema.to(fabric.device)
+        if freeze_ar_model is not None:
+            freeze_ar_model = freeze_ar_model.to(fabric.device)
+    if config.use_gan_loss:
+        gan_loss, opt_gan_loss = fabric.setup(gan_loss, opt_gan_loss)
+
+    train_loader = fabric.setup_dataloaders(train_loader)
+    train_loader = InfiniteIterator(train_loader)
+    if eval_loader is not None:
+        eval_loader = fabric.setup_dataloaders(eval_loader)
+
+
+    if config.train_ar:
+        ar_model.mark_forward_method('sampling')
+
+    global_step = 0
+
+    phases = parse_phases(config.phases)
+    fabric.print("Training Phases: ", phases)
+    total_phase_steps = sum(phase.num_steps for phase in phases)
+    total_phase = len(phases)
+
+    prev_phase_idx=-1
+    phase_step_accum = list(itertools.accumulate([phase.num_steps for phase in phases]))
+    pbar = tqdm(total=total_phase_steps, disable=not fabric.is_global_zero, dynamic_ncols=True, file=sys.stdout)
+    pbar.set_description(f"Total Steps {total_phase_steps}")
+    data_buffer = []
+    while global_step < total_phase_steps:
+        # Assign gradient and optimizer based on current phase and target
+        phase_idx, inner_idx, target, internel_step = get_phase(global_step, phases, phase_step_accum, config)
+        if not any(target):
+            phase_idx, inner_idx, target, internel_step = get_phase(global_step+1, phases, phase_step_accum, config)
+        if phase_idx != prev_phase_idx:
+            if config.lr_scheduler != 'none':
+                scheduler_enc = torch.optim.lr_scheduler.CosineAnnealingLR(opt_enc, T_max=phases[phase_idx].num_steps/len(phases[phase_idx].internal_steps),eta_min=config.lr_enc_min)  if config.train_ae else None
+                scheduler_dec = torch.optim.lr_scheduler.CosineAnnealingLR(opt_dec, T_max=phases[phase_idx].num_steps/len(phases[phase_idx].internal_steps),eta_min=config.lr_dec_min)  if config.train_ae else None
+                scheduler_ar = torch.optim.lr_scheduler.CosineAnnealingLR(opt_ar, T_max=phases[phase_idx].num_steps/len(phases[phase_idx].internal_steps),eta_min=config.lr_ar_min)  if config.train_ar else None
+                scheduler_gan_loss = torch.optim.lr_scheduler.CosineAnnealingLR(opt_gan_loss, T_max=phases[phase_idx].num_steps/len(phases[phase_idx].internal_steps),eta_min=config.lr_gan_loss_min)  if config.use_gan_loss else None
+            else:
+                scheduler_enc = None
+                scheduler_dec = None
+                scheduler_ar = None
+                scheduler_gan_loss = None
+
+
+        enc_grad = target.DO_AE or target.DO_PRIOR_ENC
+        dec_grad = target.DO_AE
+        prior_grad = target.DO_PRIOR_AR
+        disc_grad = target.DO_GAN_D 
+        toogle_require_grad(enc, enc_grad)
+        toogle_require_grad(dec, dec_grad)
+        toogle_require_grad(ar_model, prior_grad)
+        toogle_require_grad(gan_loss, disc_grad)
+        opts = []
+        if enc_grad:
+            opts.append((enc, opt_enc, scheduler_enc))
+        if dec_grad:
+            opts.append((dec, opt_dec, scheduler_dec))
+        if prior_grad:
+            opts.append((ar_model, opt_ar, scheduler_ar))
+        if disc_grad:
+            opts.append((gan_loss, opt_gan_loss, scheduler_gan_loss))
+
+        # input
+        # Clear buffer at start of cycle
+        if internel_step == 0:
+            data_buffer = []
+
+        if inner_idx == 0:
+            batch = next(train_loader)
+            data_buffer.append(batch)
+        else:
+            if len(data_buffer) > 0:
+                batch = data_buffer.pop(0)
+            else:
+                batch = next(train_loader)
+
+        imgs, labels = batch if isinstance(batch, (list, tuple)) and len(batch) == 2 else (batch, None)
+        imgs = img_unint8_to_norm(imgs)
+
+        # no labels or drop labels
+        uncond_labels = F.one_hot(torch.full((imgs.shape[0],), config.num_classes - 1, device=fabric.device, dtype=torch.long), num_classes=config.num_classes).float()
+
+        if labels is None or len(labels)==0 or not config.use_label:
+            labels = uncond_labels
+        elif config.label_drop_prob > 0:
+            labels = torch.where(torch.rand(imgs.shape[0], device=fabric.device) < config.label_drop_prob, uncond_labels, labels)
+
+        x = patchify(imgs, config.patch_size)
+
+        
+        # forward, build need computation graph for gradient calculation
+        with torch.set_grad_enabled(enc_grad):
+            quant, idx, one_hot, vqloss = enc(x, labels, training=enc_grad)
+        
+        if target.DO_AE:
+            
+            # For Diffusion/FM decoder
+            t_fm,xt,noise,v = None,None,None,None
+            if config.use_noise_query:
+                t_fm = torch.rand((x.shape[0],), device=fabric.device) 
+                xt, noise, v = make_noise_image_patch_query(x, t_fm, config.z_len) 
+
+            # For Tail Dropout
+            attn_bias = None
+            if config.use_tail_dropout and global_step>config.tail_dropout_start:
+                if config.tied_timestep:
+                    t_tail_drop = t_fm
+                else:
+                    t_tail_drop = torch.rand((x.shape[0],), device=fabric.device)
+                    no_drop_t = torch.ones_like(t_tail_drop)
+                    t_tail_drop = torch.where(torch.rand_like(t_tail_drop)<=config.tail_dropout_p,t_tail_drop,no_drop_t)
+                attn_bias = build_taildrop_attn_mask(t=t_tail_drop,context_len=int(quant.shape[1]),query_len=config.x_len,
+                                    context_type=config.context_type,t2k=config.tail_dropout_t2k,
+                                    device=fabric.device,
+                                    diti_scheduler=diti_scheduler,
+                                                    ) if config.use_tail_dropout else None
+
+            decoded = dec(quant, labels, attn_mask=attn_bias, query_tokens=xt, t=t_fm)
+
+            v_hat = decoded if config.l2_v_target else None
+            x_hat = xt - decoded * t_fm[:,None,None] if config.l2_v_target else decoded
+
+            imgs_hat = unpatchify(x_hat, config.image_size, config.patch_size)
+            
+            l2_target = v if config.l2_v_target else x
+            l2_input = v_hat if config.l2_v_target else x_hat
+
+        if target.DO_PRIOR_ENC or target.DO_PRIOR_AR:
+            with torch.set_grad_enabled(prior_grad):
+                ar_logits = ar_model(idx,labels)
+
+        if target.DO_PRIOR_ENC:
+            with torch.no_grad():
+                ema_ar_logits = freeze_ar_model(idx,labels)
+                with torch.amp.autocast('cuda',enabled=False):
+                    ar_logits_detach = ar_logits.detach().float()
+                    temp_low = config.prior_ar_temp_low
+                    temp_high = config.prior_ar_temp_high
+                    steps = config.z_len
+                    if config.prior_ar_temp_schedule == 'cosine':
+                        t = torch.linspace(0, 1, steps, device=fabric.device)
+                        temp = temp_low + 0.5 * (temp_high - temp_low) * (1 - torch.cos(math.pi * t))
+                    elif config.prior_ar_temp_schedule == 'linear':
+                        temp = torch.linspace(temp_low, temp_high, steps, device=fabric.device)
+                    elif config.prior_ar_temp_schedule == 'constant':
+                        temp = torch.full((steps,), temp_low, device=fabric.device)
+                    else:
+                        raise NotImplementedError(f"Unknown schedule: {config.prior_ar_temp_schedule}")
+
+                    b, n = ema_ar_logits.shape[0], ema_ar_logits.shape[1]
+                    temp = temp.view(1, n, 1).expand(b, n, 1)
+
+                    ema_ar_logits = ema_ar_logits.detach().float() / temp
+                    prior_ema_ar_loss = F.cross_entropy(ema_ar_logits.reshape(-1, ema_ar_logits.shape[-1]).float(), idx.reshape(-1)) 
+
+                    ema_ar_log_probs= torch.log_softmax(ema_ar_logits, dim=-1)
+                    ar_log_probs = torch.log_softmax(ar_logits_detach, dim=-1)
+                    entropy_ar_log_probs = -(ar_log_probs.exp() * ar_log_probs).sum(dim=-1)
+                    entropy_ema_ar_log_probs = -(ema_ar_log_probs.exp() * ema_ar_log_probs).sum(dim=-1)
+                    log_diff = ar_log_probs - ema_ar_log_probs
+
+
+        # caculate the loss
+        l2_loss = config.l2_weight * F.mse_loss(l2_input, l2_target) if target.DO_L2 else 0.
+        l1_loss = config.l1_weight * F.l1_loss(x_hat, x) if target.DO_L1 else 0.
+        lpips_loss = config.lpips_weight * perceptual_loss(imgs_hat, imgs).mean() if target.DO_LPIPS else 0.
+        ae_loss = l2_loss + l1_loss + lpips_loss
+
+        gan_G_loss, gan_G_loss_dict = gan_loss(imgs_hat, imgs, global_step=global_step, loss='G') if target.DO_GAN_G else (0., {})
+        gan_D_loss, gan_D_loss_dict = gan_loss(imgs_hat, imgs, global_step=global_step, loss='D') if target.DO_GAN_D else (0., {})
+        gan_G_loss_weight = (config.gan_G_weight * min(1., gan_loss.calculate_adaptive_weight(ae_loss,gan_G_loss,dec.dec.out.weight))) if target.DO_GAN_G and config.disc_adaptive_weight else config.gan_G_weight
+        gan_G_loss = gan_G_loss * gan_G_loss_weight
+        gan_D_loss = gan_D_loss * config.gan_G_weight
+        correct_tokens = (ar_logits.argmax(dim=-1) == idx).detach().float().sum().item() if target.DO_PRIOR_AR else 0.
+        total_tokens = torch.tensor(idx.numel()).item() if target.DO_PRIOR_AR else 0.
+        correct_token_rate = correct_tokens / total_tokens if target.DO_PRIOR_AR else 0.
+        prior_enc_loss = config.prior_enc_weight * (one_hot * log_diff.detach()).mean() if target.DO_PRIOR_ENC else 0.
+        prior_ar_loss = F.cross_entropy(ar_logits.reshape(-1, ar_logits.shape[-1]).float(), idx.reshape(-1)) if target.DO_PRIOR_AR else 0.
+
+        loss = ae_loss + prior_enc_loss + prior_ar_loss + gan_G_loss + gan_D_loss + vqloss 
+
+        # backward and optimization
+        fabric.backward(loss)
+        
+        for model, opt, scheduler in opts:
+            if model is not None and opt is not None:
+                zero_nan_gradients(model)
+                if config.grad_clip>0:
+                    fabric.clip_gradients(model, opt, max_norm=config.grad_clip)
+                opt.step()
+                opt.zero_grad(set_to_none=True)
+            if scheduler is not None:
+                scheduler.step() 
+        if config.train_ae:
+            if enc_grad:
+                ema_update(enc, enc_ema, ae_ema_rate)
+            if dec_grad:
+                ema_update(dec, dec_ema, ae_ema_rate)
+            if config.use_gan_loss and disc_grad:
+                ema_update(gan_loss, gan_loss_ema, ae_ema_rate)
+        if config.train_ar and prior_grad:
+            ema_update(ar_model, ar_model_ema, ar_ema_rate)
+
+        #visualization
+        if (global_step + 1) % config.visualize_freq == 0:
+            grid = x[:config.visualize_img_num]
+            if not config.ema_reconstruction:
+                if config.train_ae:
+                    toogle_train_eval(enc_raw,train=False) 
+                    toogle_train_eval(dec_raw,train=False) 
+                eval_enc = enc_raw
+                eval_dec = dec_raw
+            else:
+                eval_enc = enc_ema_raw
+                eval_dec = dec_ema_raw
+            if not config.ema_sampling:
+                if config.train_ar:
+                    toogle_train_eval(ar_model_raw,train=False) 
+                eval_ar = ar_model_raw
+            else:
+                eval_ar = ar_model_ema_raw
+            with torch.no_grad():
+                noise_t = torch.ones((x.shape[0],), device=fabric.device) if config.use_noise_query else None
+                if config.train_ae:
+                    noise_recon = torch.randn_like(x) if config.use_noise_query else None
+                    if config.use_noise_query and config.diffusion_decoder:
+                        all_recon_images,_ = euler_recon_loop(eval_enc, eval_dec, x, labels, config=config, diti_scheduler=diti_scheduler)
+                    else:
+                        decoded_recon,_ = reconstruction(eval_enc, eval_dec, x, labels, query_tokens=noise_recon, t=noise_t)
+                        all_recon_images = decoded_recon if not config.l2_v_target else noise_recon-decoded_recon
+                    grid = torch.cat([grid, all_recon_images[:config.visualize_img_num]], dim=0)
+
+                if config.train_ar:
+                    noise_sample = torch.randn_like(x) if config.use_noise_query else None
+                    if config.use_noise_query and config.diffusion_decoder:
+                        all_sample_images = euler_sample_loop(eval_enc, eval_dec, eval_ar, bz=x.shape[0], 
+                                            class_label=labels,temperature=config.temperature, 
+                                            topK=config.topK, topP=config.topP, 
+                                            cfg=config.cfg, cfg_schedule=config.cfg_schedule, cfg_power=config.cfg_power, 
+                                            cache_kv=config.cache_kv, config=config, diti_scheduler=diti_scheduler)
+                    else:
+                        decoded_sample = sampling(eval_enc, eval_dec, eval_ar, bz=x.shape[0], 
+                                                class_label=labels,temperature=config.temperature, 
+                                                topK=config.topK, topP=config.topP, 
+                                                cfg=config.cfg, cfg_schedule=config.cfg_schedule, cfg_power=config.cfg_power, 
+                                                cache_kv=config.cache_kv, config=config, t=noise_t ,query_tokens=noise_sample)
+                        all_sample_images = decoded_sample if not config.l2_v_target else noise_sample-decoded_sample
+                    grid = torch.cat([grid, all_sample_images[:config.visualize_img_num]], dim=0)
+                grid = img_denormalize(unpatchify(grid, config.image_size, config.patch_size))
+                grid = torchvision.utils.make_grid(grid, nrow=config.visualize_img_num, normalize=False)
+                fabric.log_dict({"visualization":wandb.Image(grid),'global_step':global_step+1})
+            if not config.ema_reconstruction and config.train_ae:
+                toogle_train_eval(enc_raw,train=True) 
+                toogle_train_eval(dec_raw,train=True) 
+            if not config.ema_sampling and config.train_ar:
+                toogle_train_eval(ar_model_raw,train=True) 
+           
+        # eval and save ckpt
+        rFID = 0.0
+        gFID = 0.0
+        if eval_loader is not None and (global_step + 1) % config.eval_freq == 0:
+             
+            sample_cached_path = os.path.join(save_dir, "sample_images.npz")
+            recon_cached_path = os.path.join(save_dir, "recon_images.npz")
+            gt_cache_path = config.eval_fid_ref_path
+
+            samples_buf = []
+            recons_buf = []
+            if not config.ema_reconstruction:
+                if config.train_ae:
+                    toogle_train_eval(enc_raw,train=False) 
+                    toogle_train_eval(dec_raw,train=False) 
+                eval_enc = enc_raw
+                eval_dec = dec_raw
+            else:
+                eval_enc = enc_ema_raw
+                eval_dec = dec_ema_raw
+            
+            if not config.ema_sampling:
+                if config.train_ar:
+                    toogle_train_eval(ar_model_raw, train=False)
+                eval_ar = ar_model_raw
+            else:
+                eval_ar = ar_model_ema_raw
+
+
+            with torch.no_grad():
+
+                fabric.barrier()
+                for i, batch in enumerate(tqdm(eval_loader,disable=not fabric.is_global_zero,dynamic_ncols=True,file=sys.stdout,desc="Evaluating")):
+                    imgs, labels = batch if isinstance(batch, (list, tuple)) and len(batch) == 2 else (batch, None)
+                    x = img_unint8_to_norm(imgs)
+                    x = patchify(x, config.patch_size)
+                    uncond_labels = F.one_hot(torch.full((imgs.shape[0],), config.num_classes - 1, device=fabric.device, dtype=torch.long), num_classes=config.num_classes).float()
+                    noise_t = torch.ones((labels.shape[0],), device=fabric.device) if config.use_noise_query else None
+
+                    if labels is None or len(labels)==0 or not config.use_label:
+                        labels = uncond_labels
+                    idx = None
+                    noise_recon = torch.randn_like(x) if config.use_noise_query else None
+
+                    if config.train_ae:
+                        if config.use_noise_query and config.diffusion_decoder:
+                            recon_images, idx = euler_recon_loop(eval_enc, eval_dec, x, labels, config=config, diti_scheduler=diti_scheduler)
+                        else:
+                            decoded_recon, idx = reconstruction(eval_enc, eval_dec, x, labels, query_tokens=noise_recon,t = noise_t)
+                            recon_images = decoded_recon if not config.l2_v_target else noise_recon-decoded_recon
+
+                        recon_images = img_norm_to_uint8(unpatchify(recon_images, config.image_size, config.patch_size))
+                        if fabric.world_size>1:
+                            recon_images = fabric.all_gather(recon_images) 
+                            recon_images = recon_images.flatten(0,1)
+                        recons_buf.append(recon_images.permute(0,2,3,1).cpu().numpy())
+                    if config.train_ar:
+                        if config.use_noise_query and config.diffusion_decoder:
+                            sample_images = euler_sample_loop(eval_enc, eval_dec, eval_ar, bz=x.shape[0], 
+                                                class_label=labels,temperature=config.temperature, 
+                                                topK=config.topK, topP=config.topP, 
+                                                cfg=config.cfg, cfg_schedule=config.cfg_schedule, cfg_power=config.cfg_power, 
+                                                cache_kv=config.cache_kv, config=config, diti_scheduler=diti_scheduler)
+                        else:
+                            noise_sample = torch.randn_like(x) if config.use_noise_query else None
+                            decoded_sample = sampling(eval_enc, eval_dec, eval_ar, 
+                                            bz=x.shape[0], class_label=labels,temperature=config.temperature, 
+                                            topK=config.topK, topP=config.topP, cfg=config.cfg, cfg_schedule=config.cfg_schedule, 
+                                            cfg_power=config.cfg_power, cache_kv=config.cache_kv, config=config,t=noise_t,query_tokens=noise_sample)
+                            sample_images = decoded_sample if not config.l2_v_target else noise_sample-decoded_sample
+
+                        sample_images = img_norm_to_uint8(unpatchify(sample_images, config.image_size, config.patch_size))
+                        if fabric.world_size>1:
+                            sample_images = fabric.all_gather(sample_images) 
+                            sample_images = sample_images.flatten(0,1)
+                        samples_buf.append(sample_images.permute(0,2,3,1).cpu().numpy())
+
+        
+                fabric.barrier()
+                if fabric.is_global_zero:
+                    if config.train_ar:
+                        np.savez_compressed(sample_cached_path, np.concatenate(samples_buf,axis=0))
+                    if config.train_ae:
+                        np.savez_compressed(recon_cached_path, np.concatenate(recons_buf,axis=0))
+                    
+                    if config.train_ar:
+                        gFID = adm_fid_evaluator(sample_cached_path, gt_cache_path, config, fabric)
+                    if config.train_ae:
+                        rFID = adm_fid_evaluator(recon_cached_path, gt_cache_path, config, fabric)
+                fabric.barrier()
+
+
+            if not config.ema_reconstruction and config.train_ae:
+                toogle_train_eval(enc_raw,train=True) 
+                toogle_train_eval(dec_raw,train=True) 
+            if not config.ema_sampling and config.train_ar:
+                toogle_train_eval(ar_model_raw,train=True) 
+
+        # metrics
+        if (global_step +1) % config.log_freq == 0:
+            # current learning rates
+            lr_enc = opt_enc.param_groups[0]['lr'] if opt_enc is not None else 0.
+            lr_dec = opt_dec.param_groups[0]['lr'] if opt_dec is not None else 0.
+            lr_ar = opt_ar.param_groups[0]['lr'] if opt_ar is not None else 0.
+            lr_gan_loss = opt_gan_loss.param_groups[0]['lr'] if opt_gan_loss is not None else 0.
+            metrics = {
+                        'Phase':phase_idx,
+                        'ae_loss/l1_loss':l1_loss,
+                        'ae_loss/l2_loss':l2_loss,
+                        'ae_loss/lpips_loss':lpips_loss,
+                        'ae_loss/loss':ae_loss, 
+                        'ae_loss/vqloss': vqloss,
+                        'ae_loss/gan_G_loss':gan_G_loss,
+                        'ae_loss/gan_D_loss':gan_D_loss,
+                        'prior_loss/logq-logp':log_diff.mean() if target.DO_PRIOR_ENC else 0., 
+                        'prior_loss/logq':ar_log_probs.mean() if target.DO_PRIOR_ENC else 0., 
+                        'prior_loss/logp':ema_ar_log_probs.mean() if target.DO_PRIOR_ENC else 0., 
+                        'prior_loss/entropy_q':entropy_ar_log_probs.mean() if target.DO_PRIOR_ENC else 0.,
+                        'prior_loss/entropy_p':entropy_ema_ar_log_probs.mean() if target.DO_PRIOR_ENC else 0.,
+                        'prior_loss/ar_loss':prior_ar_loss, 
+                        'prior_loss/ema_ar_loss':prior_ema_ar_loss  if target.DO_PRIOR_ENC else 0.,
+                        'prior_loss/correct_token_rate': correct_token_rate,
+                        'FID/rFID':rFID,
+                        'FID/gFID':gFID,
+                        'lr/lr_enc': lr_enc,
+                        'lr/lr_dec': lr_dec,
+                        'lr/lr_ar': lr_ar,
+                        'lr/lr_gan_loss': lr_gan_loss,
+                      }
+            gan_loss_dict = {**gan_G_loss_dict, **gan_D_loss_dict, 'gan_G_loss_weight':gan_G_loss_weight}
+            metrics.update(gan_loss_dict)
+            metrics = {k: fabric.all_reduce(v).item() if isinstance(v, torch.Tensor) else v for k, v in metrics.items() }
+            metrics_logger = {k: v for k, v in metrics.items() if v!=0. }
+            metrics_4f = {k: f"{v:.4f}" if k!='Phase' else int(v) for k, v in metrics.items() }
+            fabric.log_dict(metrics_logger,step=global_step+1)
+            pbar.set_postfix(metrics_4f)
+            pbar.update(config.log_freq)
+
+        if config.save_ckpt and (global_step==0 or (global_step+1) % config.ckpt_freq == 0):
+            ckpt_name = f'Phase={phase_idx}-Step={global_step+1}-rFID={rFID:.4f}-gFID={gFID:.4f}' if rFID is not None and gFID is not None else "_".join([f'{k}={v:.4f}' for k, v in metrics.items()]) 
+            ckpt_path = f'{save_dir}/ckpts/{ckpt_name}.ckpt'
+            fabric.save(ckpt_path, {'enc':enc,
+                        'ar_model':ar_model,
+                        'freeze_ar_model': freeze_ar_model,
+                        'dec':dec,
+                        'gan_loss':gan_loss,
+                        'enc_ema':enc_ema,
+                        'dec_ema':dec_ema,
+                        'ar_model_ema':ar_model_ema,
+                        'opt_enc':opt_enc,
+                        'opt_ar':opt_ar,
+                        'opt_dec':opt_dec,
+                        'opt_gan_loss':opt_gan_loss})
+        global_step += 1
+        prev_phase_idx = phase_idx
+    if config.save_ckpt:
+        ckpt_path = f'{save_dir}/ckpts/last.ckpt'
+        fabric.save(ckpt_path, {'enc':enc,
+                    'ar_model':ar_model,
+                    'dec':dec,
+                    'gan_loss':gan_loss,
+                    'freeze_ar_model':freeze_ar_model,
+                    'enc_ema':enc_ema,
+                    'dec_ema':dec_ema,
+                    'ar_model_ema':ar_model_ema,
+                    'opt_enc':opt_enc,
+                    'opt_ar':opt_ar,
+                    'opt_dec':opt_dec,
+                    'opt_gan_loss':opt_gan_loss})
+
+    logger.finalize("Training completed")
+
+if __name__ == "__main__":
+    OmegaConf.register_new_resolver("eval", eval)
+    cli_conf = OmegaConf.from_cli()
+
+    config_path = cli_conf.get("--config")
+    conf = OmegaConf.load(config_path)
+    if cli_conf is not None:
+        for k, v in cli_conf.items():
+            OmegaConf.update(conf, k, v)
+
+    train(conf)
+