update

.gitattributes

@@ -33,3 +33,10 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+assets/app_examples/0.png filter=lfs diff=lfs merge=lfs -text
+assets/app_examples/1.png filter=lfs diff=lfs merge=lfs -text
+assets/overview.png filter=lfs diff=lfs merge=lfs -text
+assets/app_examples/2.png filter=lfs diff=lfs merge=lfs -text
+assets/app_examples/4.png filter=lfs diff=lfs merge=lfs -text
+assets/comparison_of_generation.png filter=lfs diff=lfs merge=lfs -text
+*.png filter=lfs diff=lfs merge=lfs -text

.gitignore

@@ -0,0 +1,6 @@
+__pycache__/
+*.ckpt
+checkpoints/
+results/
+VTBench_models/
+README.md

app.py

@@ -0,0 +1,166 @@
+import os
+import spaces
+import subprocess
+import sys
+
+# REQUIREMENTS_FILE = "requirements.txt"
+# if os.path.exists(REQUIREMENTS_FILE):
+#     try:
+#         print("Installing dependencies from requirements.txt...")
+#         subprocess.check_call([sys.executable, "-m", "pip", "install", "-r", REQUIREMENTS_FILE])
+#         print("Dependencies installed successfully.")
+#     except subprocess.CalledProcessError as e:
+#         print(f"Failed to install dependencies: {e}")
+# else:
+#     print("requirements.txt not found.")
+
+import gradio as gr
+from src.data_processing import pil_to_tensor, tensor_to_pil
+from PIL import Image
+from src.model_processing import get_model
+from huggingface_hub import snapshot_download
+import torch
+
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+print(f"Running on: {device}")
+
+MODEL_DIR = "./VTBench_models"
+if not os.path.exists(MODEL_DIR):
+    print("Downloading VTBench_models from Hugging Face...")
+    snapshot_download(
+        repo_id="huaweilin/VTBench_models",
+        local_dir=MODEL_DIR,
+        local_dir_use_symlinks=False
+    )
+    print("Download complete.")
+
+example_image_paths = [f"assets/app_examples/{i}.png" for i in range(0, 5)]
+
+model_name_mapping = {
+    "SD3.5L": "SD3.5L",
+    "chameleon": "Chameleon",
+    # "flowmo_lo": "FlowMo Lo",
+    # "flowmo_hi": "FlowMo Hi",
+    # "gpt4o": "GPT-4o",
+    "janus_pro_1b": "Janus Pro 1B/7B",
+    # "llamagen-ds8": "LlamaGen ds8",
+    # "llamagen-ds16": "LlamaGen ds16",
+    # "llamagen-ds16-t2i": "LlamaGen ds16 T2I",
+    # "maskbit_16bit": "MaskBiT 16bit",
+    # "maskbit_18bit": "MaskBiT 18bit",
+    # "open_magvit2": "OpenMagViT",
+    # "titok_b64": "Titok-b64",
+    # "titok_bl64": "Titok-bl64",
+    # "titok_s128": "Titok-s128",
+    # "titok_bl128": "Titok-bl128",
+    # "titok_l32": "Titok-l32",
+    # "titok_sl256": "Titok-sl256",
+    # "var_256": "VAR-256",
+    # "var_512": "VAR-512",
+    # "FLUX.1-dev": "FLUX.1-dev",
+    # "infinity_d32": "Infinity-d32",
+    # "infinity_d64": "Infinity-d64",
+    # "bsqvit": "BSQ-VIT",
+}
+
+def load_model(model_name):
+    model, data_params = get_model(MODEL_DIR, model_name)
+    model = model.to(device)
+    model.eval()
+    return model, data_params
+
+model_dict = {
+    model_name: load_model(model_name)
+    for model_name in model_name_mapping
+}
+
+placeholder_image = Image.new("RGBA", (512, 512), (0, 0, 0, 0))
+
+@spaces.GPU
+def process_selected_models(uploaded_image, selected_models):
+    results = []
+    for model_name in model_name_mapping:
+        if uploaded_image is None:
+            results.append(gr.update(value=placeholder_image, label=f"{model_name_mapping[model_name]} (No input)"))
+        elif model_name in selected_models:
+            try:
+                model, data_params = model_dict[model_name]
+                pixel_values = pil_to_tensor(uploaded_image, **data_params).unsqueeze(0).to(device)
+                output = model(pixel_values)[0]
+                reconstructed_image = tensor_to_pil(output[0].cpu(), **data_params)
+                results.append(gr.update(value=reconstructed_image, label=model_name_mapping[model_name]))
+            except Exception as e:
+                print(f"Error in model {model_name}: {e}")
+                results.append(gr.update(value=placeholder_image, label=f"{model_name_mapping[model_name]} (Error)"))
+        else:
+            results.append(gr.update(value=placeholder_image, label=f"{model_name_mapping[model_name]} (Not selected)"))
+    return results
+
+with gr.Blocks() as demo:
+    gr.Markdown("## VTBench")
+
+    gr.Markdown("---")
+
+    image_input = gr.Image(
+        type="pil",
+        label="Upload an image",
+        width=512,
+        height=512,
+    )
+
+    gr.Markdown("### Click on an example image to use it as input:")
+    example_rows = [example_image_paths[i:i+5] for i in range(0, len(example_image_paths), 5)]
+    for row in example_rows:
+        with gr.Row():
+            for path in row:
+                ex_img = gr.Image(
+                    value=path,
+                    show_label=False,
+                    interactive=True,
+                    width=256,
+                    height=256,
+                )
+
+                def make_loader(p=path):
+                    def load_img():
+                        return Image.open(p)
+                    return load_img
+        
+                ex_img.select(fn=make_loader(), outputs=image_input)
+
+    gr.Markdown("---")
+
+    gr.Markdown("⚠️ **The more models you select, the longer the processing time will be.**")
+    model_selector = gr.CheckboxGroup(
+        choices=list(model_name_mapping.keys()),
+        label="Select models to run",
+        value=["SD3.5L", "chameleon", "janus_pro_1b"],
+        interactive=True,
+    )
+    run_button = gr.Button("Start Processing")
+
+    image_outputs = []
+    model_items = list(model_name_mapping.items())
+
+    n_columns = 5
+    output_rows = [model_items[i:i+n_columns] for i in range(0, len(model_items), n_columns)]
+
+    with gr.Column():
+        for row in output_rows:
+            with gr.Row():
+                for model_name, display_name in row:
+                    out_img = gr.Image(
+                        label=f"{display_name} (Not run)",
+                        value=placeholder_image,
+                        width=512,
+                        height=512,
+                    )
+                    image_outputs.append(out_img)
+
+    run_button.click(
+        fn=process_selected_models,
+        inputs=[image_input, model_selector],
+        outputs=image_outputs
+    )
+
+demo.launch()

evaluations/character_error_rate.py

@@ -0,0 +1,27 @@
+import torch
+from torchmetrics import Metric
+from ocr import OCR
+import Levenshtein
+
+
+class CharacterErrorRate(Metric):
+    def __init__(self, ocr, dist_sync_on_step=False):
+        # super().__init__(dist_sync_on_step=dist_sync_on_step)
+        super().__init__()
+        self.add_state("total_errors", default=torch.tensor(0.0), dist_reduce_fx="sum")
+        self.add_state("total_chars", default=torch.tensor(0.0), dist_reduce_fx="sum")
+        self.ocr = ocr
+
+    def update(self, pred_images, target_images):
+        for pred_img, target_img in zip(pred_images, target_images):
+            pred_text = self.ocr.predict(pred_img)
+            target_text = self.ocr.predict(target_img)
+
+            dist = Levenshtein.distance(pred_text, target_text)
+            self.total_errors += dist
+            self.total_chars += len(target_text)
+
+    def compute(self):
+        if self.total_chars == 0:
+            return torch.tensor(0.0)
+        return self.total_errors / self.total_chars

evaluations/evaluate_images.py

@@ -0,0 +1,130 @@
+import os
+import argparse
+from PIL import Image
+from tqdm import tqdm
+from torchvision import transforms
+from torch.utils.data import Dataset, DataLoader
+import torch
+import torch.nn.functional as F
+from ocr import OCR
+from character_error_rate import CharacterErrorRate
+from word_error_rate import WordErrorRate
+from torchmetrics.image import (
+    PeakSignalNoiseRatio,
+    StructuralSimilarityIndexMeasure,
+    LearnedPerceptualImagePatchSimilarity,
+    FrechetInceptionDistance,
+)
+
+
+class ImageFolderPairDataset(Dataset):
+    def __init__(self, dir1, dir2, transform=None):
+        self.dir1 = dir1
+        self.dir2 = dir2
+        self.filenames = sorted(os.listdir(dir1))
+        self.transform = transform
+
+    def __len__(self):
+        return len(self.filenames)
+
+    def __getitem__(self, idx):
+        name = self.filenames[idx]
+        img1 = Image.open(os.path.join(self.dir1, name)).convert("RGB")
+        img2 = Image.open(os.path.join(self.dir2, name)).convert("RGB")
+        if self.transform:
+            img1 = self.transform(img1)
+            img2 = self.transform(img2)
+        return img1, img2
+
+
+def evaluate(args):
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    print(f"Using device: {device}")
+
+    transform = transforms.Compose(
+        [transforms.Resize((args.image_size, args.image_size)), transforms.ToTensor()]
+    )
+
+    dataset = ImageFolderPairDataset(
+        args.original_dir, args.reconstructed_dir, transform
+    )
+    loader = DataLoader(
+        dataset, batch_size=args.batch_size, shuffle=False, num_workers=args.num_workers
+    )
+
+    if "cer" in args.metrics or "wer" in args.metrics:
+        ocr = OCR(device)
+
+    # Metrics init
+    metrics = {}
+
+    if "psnr" in args.metrics:
+        metrics["psnr"] = PeakSignalNoiseRatio().to(device)
+    if "ssim" in args.metrics:
+        metrics["ssim"] = StructuralSimilarityIndexMeasure().to(device)
+    if "lpips" in args.metrics:
+        metrics["lpips"] = LearnedPerceptualImagePatchSimilarity().to(device)
+    if "fid" in args.metrics:
+        metrics["fid"] = FrechetInceptionDistance().to(device)
+    if "cer" in args.metrics:
+        metrics["cer"] = CharacterErrorRate(ocr)
+    if "wer" in args.metrics:
+        metrics["wer"] = WordErrorRate(ocr)
+
+    for batch in tqdm(loader, desc="Evaluating"):
+        # img1, img1_path, img2, img2_path = [b.to(device) for b in batch]
+        img1, img2 = [b.to(device) for b in batch]
+
+        if "psnr" in metrics:
+            metrics["psnr"].update(img2, img1)
+        if "ssim" in metrics:
+            metrics["ssim"].update(img2, img1)
+        if "lpips" in metrics:
+            metrics["lpips"].update(img2, img1)
+        if "cer" in metrics:
+            metrics["cer"].update(img2, img1)
+        if "wer" in metrics:
+            metrics["wer"].update(img2, img1)
+        if "fid" in metrics:
+            img1_uint8 = (img1 * 255).clamp(0, 255).to(torch.uint8)
+            img2_uint8 = (img2 * 255).clamp(0, 255).to(torch.uint8)
+            metrics["fid"].update(img1_uint8, real=True)
+            metrics["fid"].update(img2_uint8, real=False)
+
+    print("\nResults:")
+    for name, metric in metrics.items():
+        print(f"{name.upper()}", end="\t")
+    print()
+    for name, metric in metrics.items():
+        result = metric.compute().item()
+        print(f"{result:.4f}", end="\t")
+    print()
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "--original_dir", type=str, required=True, help="Path to original images"
+    )
+    parser.add_argument(
+        "--reconstructed_dir",
+        type=str,
+        required=True,
+        help="Path to reconstructed images",
+    )
+    parser.add_argument(
+        "--metrics",
+        nargs="+",
+        default=["psnr", "ssim", "lpips", "fid"],
+        help="Metrics to compute: psnr, ssim, lpips, fid",
+    )
+    parser.add_argument(
+        "--batch_size", type=int, default=8, help="Batch size for processing"
+    )
+    parser.add_argument("--image_size", type=int, default=256, help="Image resize size")
+    parser.add_argument(
+        "--num_workers", type=int, default=4, help="Number of workers for DataLoader"
+    )
+    args = parser.parse_args()
+
+    evaluate(args)

evaluations/ocr.py

@@ -0,0 +1,44 @@
+from PIL import Image
+from transformers import AutoProcessor, AutoModelForImageTextToText
+import torch
+
+
+class OCR:
+    def __init__(self, device="cpu"):
+        self.device = torch.device(device)
+        self.model = AutoModelForImageTextToText.from_pretrained(
+            "google/gemma-3-12b-it",
+            torch_dtype=torch.bfloat16,
+        ).to(self.device)
+        self.processor = AutoProcessor.from_pretrained("google/gemma-3-12b-it")
+
+        self.messages = [
+            {
+                "role": "user",
+                "content": [
+                    {"type": "image"},
+                    {
+                        "type": "text",
+                        "text": "Extract and output only the text from the image in its original language. If there is no text, return nothing.",
+                    },
+                ],
+            },
+        ]
+
+    def predict(self, image):
+        image = (
+            (image * 255).clamp(0, 255).to(torch.uint8).permute((1, 2, 0)).cpu().numpy()
+        )
+        image = Image.fromarray(image).convert("RGB").resize((1024, 1024))
+        prompt = self.processor.apply_chat_template(
+            self.messages, add_generation_prompt=True
+        )
+        inputs = self.processor(text=prompt, images=[image], return_tensors="pt").to(
+            self.device
+        )
+        with torch.no_grad():
+            generated_ids = self.model.generate(**inputs, max_new_tokens=1024)
+        generated_text = self.processor.batch_decode(
+            generated_ids[:, inputs.input_ids.shape[-1] :], skip_special_tokens=True
+        )[0]
+        return generated_text

evaluations/word_error_rate.py

@@ -0,0 +1,30 @@
+import torch
+from torchmetrics import Metric
+import Levenshtein
+
+
+class WordErrorRate(Metric):
+    def __init__(self, ocr, dist_sync_on_step=False):
+        # super().__init__(dist_sync_on_step=dist_sync_on_step)
+        super().__init__()
+        self.ocr = ocr
+        self.add_state("total_errors", default=torch.tensor(0.0), dist_reduce_fx="sum")
+        self.add_state("total_words", default=torch.tensor(0.0), dist_reduce_fx="sum")
+
+    def update(self, pred_images, target_images):
+        for pred_img, target_img in zip(pred_images, target_images):
+            pred_text = self.ocr.predict(pred_img)
+            target_text = self.ocr.predict(target_img)
+
+            pred_words = pred_text.strip().split()
+            target_words = target_text.strip().split()
+
+            dist = Levenshtein.distance(" ".join(pred_words), " ".join(target_words))
+
+            self.total_errors += dist
+            self.total_words += len(target_words)
+
+    def compute(self):
+        if self.total_words == 0:
+            return torch.tensor(0.0)
+        return self.total_errors / self.total_words

examples/get_result.py

@@ -0,0 +1,94 @@
+import os
+import pandas as pd
+
+root_dir = "./"
+
+model_name_mapping = {
+    "flowmo_lo": "FlowMo Lo",
+    "flowmo_hi": "FlowMo Hi",
+    "gpt4o": "GPT-4o",
+    "janus_pro_1b": "Janus Pro 1B/7B",
+    "llamagen-ds8": "LlamaGen ds8",
+    "llamagen-ds16": "LlamaGen ds16",
+    "llamagen-ds16-t2i": "LlamaGen ds16 T2I",
+    "maskbit_16bit": "MaskBiT 16bit",
+    "maskbit_18bit": "MaskBiT 18bit",
+    "open_magvit2": "OpenMagViT",
+    "titok_b64": "Titok-b64",
+    "titok_bl64": "Titok-bl64",
+    "titok_s128": "Titok-s128",
+    "titok_bl128": "Titok-bl128",
+    "titok_l32": "Titok-l32",
+    "titok_sl256": "Titok-sl256",
+    "var_256": "VAR-256",
+    "var_512": "VAR-512",
+    "SD3.5L": "SD3.5L",
+    "FLUX.1-dev": "FLUX.1-dev",
+    "infinity_d32": "Infinity-d32",
+    "infinity_d64": "Infinity-d64",
+    "chameleon": "Chameleon",
+    "bsqvit": "BSQ-VIT",
+}
+
+output_order = [
+    "FlowMo Lo",
+    "FlowMo Hi",
+    "MaskBiT 16bit",
+    "MaskBiT 18bit",
+    "Titok-l32",
+    "Titok-b64",
+    "Titok-s128",
+    "Titok-bl64",
+    "Titok-bl128",
+    "Titok-sl256",
+    "OpenMagViT",
+    "LlamaGen ds8",
+    "BSQ-VIT",
+    "VAR-256",
+    "Janus Pro 1B/7B",
+    "Chameleon",
+    "LlamaGen ds16",
+    "LlamaGen ds16 T2I",
+    "VAR-512",
+    "Infinity-d32",
+    "Infinity-d64",
+    "SD3.5L",
+    "FLUX.1-dev",
+    "GPT-4o",
+]
+
+for dataset_name in os.listdir(root_dir):
+    dataset_path = os.path.join(root_dir, dataset_name)
+    if not os.path.isdir(dataset_path):
+        continue
+
+    results = {}
+
+    for model_dir in os.listdir(dataset_path):
+        model_path = os.path.join(dataset_path, model_dir)
+        result_file = os.path.join(model_path, "result.txt")
+
+        if os.path.isfile(result_file):
+            with open(result_file, "r", encoding="utf-8") as f:
+                lines = f.readlines()
+
+            if len(lines) >= 2:
+                metrics_line = lines[-2].strip()
+                values_line = lines[-1].strip()
+
+                metrics = metrics_line.split()
+                values = values_line.split()
+
+                mapped_name = model_name_mapping.get(model_dir, model_dir)
+                results[mapped_name] = values
+
+    if results:
+        header = "\t".join(metrics)
+        print(f"{dataset_name}\t{header}")
+        for model_name in output_order:
+            if model_name in results:
+                values = results[model_name]
+                print(f"{model_name}\t" + "\t".join(values))
+            else:
+                print(f"{model_name}\t" + "no result")
+        print()

examples/run.sh

@@ -0,0 +1,85 @@
+#!/bin/bash
+
+dataset_name_list=("task1-imagenet" "task1-high-resolution" "task1-varying-resolution" "task2-detail-preservation" "task3-movie-posters" "task3-arxiv-abstracts" "task3-multilingual_Chinese" "task3-multilingual_Hindi" "task3-multilingual_Japanese" "task3-multilingual_Korean")
+model_name_list=("chameleon" "llamagen-ds16" "llamagen-ds8" "flowmo_lo" "flowmo_hi" "open_magvit2" "titok_l32" "titok_b64" "titok_s128" "titok_bl64" "titok_bl128" "titok_sl256" "janus_pro_1b" "maskbit_18bit" "maskbit_16bit" "var_256" "var_512" "SD3.5L" "gpt4o" "llamagen-ds16-t2i" "infinity_d32" "infinity_d64" "bsqvit" "FLUX.1-dev")
+
+batch_size=1
+
+if command -v sbatch >/dev/null 2>&1; then
+  has_slurm=true
+else
+  has_slurm=false
+fi
+
+shell_dir=$(cd "$(dirname "$0")";pwd)
+echo "shell_dir: ${shell_dir}"
+base_path="${shell_dir}/../"
+
+for dataset_name in "${dataset_name_list[@]}"
+do
+  cd ${shell_dir}
+  folder_dir="${dataset_name}"
+  mkdir ${folder_dir}
+
+  metrics="fid ssim psnr lpips"
+  split_name="test"
+  n_take=-1
+
+  if [[ $dataset_name == task3-multilingual_* ]]; then
+    split_name="${dataset_name##*_}"
+    dataset_name="${dataset_name%_*}"
+  fi
+  if [ "$dataset_name" = "task1-imagenet" ]; then
+    split_name="val"
+  fi
+
+  if [ "$dataset_name" = "task1-varying-resolution" ]; then
+    batch_size=1
+  fi
+  if [ "$dataset_name" = "task3-movie-posters" ]; then
+    metrics="fid ssim psnr lpips cer wer"
+  fi
+  if [ "$dataset_name" = "task3-arxiv-abstracts" ]; then
+    metrics="fid ssim psnr lpips cer wer"
+  fi
+  if [ "$dataset_name" = "task3-multilingual" ]; then
+    metrics="fid ssim psnr lpips cer"
+  fi
+
+  for model_name in "${model_name_list[@]}"
+  do
+    if [ "$dataset_name" = "task1-imagenet" ] && [ "$model_name" = "gpt4o" ]; then
+      n_take=100
+    fi
+    cd ${shell_dir}
+
+    work_dir="${folder_dir}/${model_name}"
+    echo "model_name: ${model_name}, work_dir: ${work_dir}"
+    mkdir ${work_dir}
+
+    cp submit.sh ${work_dir}
+  
+    cd ${work_dir}
+    sed -i "s|{model_name}|${model_name}|g" submit.sh
+    sed -i "s|{split_name}|${split_name}|g" submit.sh
+    sed -i "s|{dataset_name}|${dataset_name}|g" submit.sh
+    sed -i "s|{batch_size}|${batch_size}|g" submit.sh
+    sed -i "s|{base_path}|${base_path}|g" submit.sh
+    sed -i "s|{metrics}|${metrics}|g" submit.sh
+    sed -i "s|{n_take}|${n_take}|g" submit.sh
+  
+#     if [ "$has_slurm" = true ]; then
+#       res=$(sbatch ./submit.sh)
+#       res=($res)
+#       task_id=${res[-1]}
+#       echo "task_id: ${task_id}"
+#       touch "task_id_${task_id}"
+#     else
+#       echo "Slurm not detected, running with bash..."
+#       bash ./submit.sh
+#     fi
+
+    bash ./submit.sh
+
+  done
+done

examples/submit.sh

@@ -0,0 +1,10 @@
+#!/bin/bash
+# Put your slurm commands here
+
+accelerate launch --num_processes=1 {base_path}/main.py --batch_size {batch_size} --model_name {model_name} --split_name {split_name} --dataset_name {dataset_name} --output_dir {model_name}_results --n_take {n_take}
+python {base_path}/evaluations/evaluate_images.py \
+    --original_dir {model_name}_results/original_images \
+    --reconstructed_dir {model_name}_results/reconstructed_images/ \
+    --metrics {metrics} \
+    --batch_size 16 \
+    --num_workers 8 | tee result.txt

main.py

@@ -0,0 +1,99 @@
+import numpy as np
+import os
+import PIL
+import pickle
+import torch
+import argparse
+import json
+from PIL import Image
+import torch.nn as nn
+import torch
+from transformers import AutoProcessor, AutoModelForImageTextToText
+from src.data_loader import DataCollatorForSupervisedDataset, get_dataset
+from src.data_processing import tensor_to_pil
+from src.model_processing import get_model
+from PIL import Image
+from accelerate import Accelerator
+from torch.utils.data import DataLoader
+from tqdm import tqdm
+from concurrent.futures import ThreadPoolExecutor
+
+parser = argparse.ArgumentParser()
+parser.add_argument("--model_name", type=str, default="chameleon")
+parser.add_argument("--model_path", type=str, default=None)
+parser.add_argument("--dataset_name", type=str, default="task3-movie-posters")
+parser.add_argument("--split_name", type=str, default="test")
+parser.add_argument("--batch_size", default=8, type=int)
+parser.add_argument("--output_dir", type=str, default=None)
+parser.add_argument("--begin_id", default=0, type=int)
+parser.add_argument("--n_take", default=-1, type=int)
+args = parser.parse_args()
+
+batch_size = args.batch_size
+output_dir = args.output_dir
+
+accelerator = Accelerator()
+
+if accelerator.is_main_process and output_dir is not None:
+    os.makedirs(output_dir, exist_ok=True)
+    os.makedirs(f"{output_dir}/original_images", exist_ok=True)
+    os.makedirs(f"{output_dir}/reconstructed_images", exist_ok=True)
+    os.makedirs(f"{output_dir}/results", exist_ok=True)
+
+model, data_params = get_model(args.model_path, args.model_name)
+dataset = get_dataset(args.dataset_name, args.split_name, None if args.n_take <= 0 else args.n_take)
+data_collator = DataCollatorForSupervisedDataset(args.dataset_name, **data_params)
+dataloader = DataLoader(
+    dataset, batch_size=batch_size, num_workers=0, collate_fn=data_collator
+)
+
+model, dataloader = accelerator.prepare(model, dataloader)
+print("Model prepared...")
+
+
+def save_results(
+    pixel_values, reconstructed_image, idx, output_dir, data_params
+):
+    if reconstructed_image is None:
+        return
+
+    ori_img = tensor_to_pil(pixel_values, **data_params)
+    rec_img = tensor_to_pil(reconstructed_image, **data_params)
+
+    ori_img.save(f"{output_dir}/original_images/{idx:08d}.png")
+    rec_img.save(f"{output_dir}/reconstructed_images/{idx:08d}.png")
+
+    result = {
+        "ori_img": ori_img,
+        "rec_img": rec_img,
+    }
+
+    with open(f"{output_dir}/results/{idx:08d}.pickle", "wb") as fw:
+        pickle.dump(result, fw)
+
+
+executor = ThreadPoolExecutor(max_workers=16)
+with torch.no_grad():
+    print("Begin data loading...")
+    for batch in tqdm(dataloader):
+        pixel_values = batch["image"]
+        reconstructed_images = model(pixel_values)
+        if isinstance(reconstructed_images, tuple):
+            reconstructed_images = reconstructed_images[0]
+
+        if output_dir is not None:
+            idx_list = batch["idx"]
+            original_images = pixel_values.detach().cpu()
+            if not isinstance(reconstructed_images, list):
+                reconstructed_images = reconstructed_images.detach().cpu()
+            for i in range(pixel_values.shape[0]):
+                executor.submit(
+                    save_results,
+                    original_images[i],
+                    reconstructed_images[i],
+                    idx_list[i],
+                    output_dir,
+                    data_params,
+                )
+
+executor.shutdown(wait=True)

requirements.txt

@@ -0,0 +1,21 @@
+numpy
+mup==1.0.0
+einops
+omegaconf
+lightning==2.3.3
+piq
+python-Levenshtein
+verovio
+pytorch_fid
+transformers
+torch-fidelity
+accelerate
+datasets
+git+https://github.com/deepseek-ai/Janus.git
+diffusers
+openai
+imageio
+huggingface_hub
+gradio
+torch
+torchvision

src/data_loader.py

@@ -0,0 +1,61 @@
+import PIL
+from PIL import Image
+from dataclasses import dataclass, field
+from datasets import load_dataset
+import torch
+from .data_processing import pil_to_tensor
+
+
+@dataclass
+class DataCollatorForSupervisedDataset(object):
+    """Collate examples for supervised fine-tuning."""
+
+    def __init__(self, dataset_name, **kwargs):
+        override_params = {}
+        if dataset_name == "DIV2K":
+            override_params = {
+                "target_image_size": -1,
+                "lock_ratio": True,
+                "center_crop": False,
+                "padding": False,
+            }
+        if dataset_name == "imagenet":
+            override_params = {"center_crop": True, "padding": False}
+        if dataset_name == "movie_posters":
+            override_params = {"center_crop": True, "padding": False}
+        if dataset_name == "high_quality_1024":
+            override_params = {"target_image_size": (1024, 1024)}
+
+        self.data_params = {**kwargs, **override_params}
+
+    def __call__(self, instances):
+        images = torch.stack(
+            [
+                pil_to_tensor(instance["image"], **self.data_params)
+                for instance in instances
+            ],
+            dim=0,
+        )
+        idx = [instance["idx"] for instance in instances]
+        return dict(image=images, idx=idx)
+
+
+class ImagenetDataset(torch.utils.data.Dataset):
+    def __init__(self, dataset_name, split_name="test", n_take=None):
+        print(dataset_name, split_name)
+        ds = load_dataset("huaweilin/VTBench", name=dataset_name, split=split_name if n_take is None else f"{split_name}[:{n_take}]")
+        self.image_list = ds["image"]
+
+    def __len__(self):
+        return len(self.image_list)
+
+    def __getitem__(self, idx):
+        return dict(
+            image=self.image_list[idx],
+            idx=idx,
+        )
+
+
+def get_dataset(dataset_name, split_name, n_take):
+    dataset = ImagenetDataset(dataset_name, split_name, n_take)
+    return dataset

src/data_processing.py

@@ -0,0 +1,89 @@
+import numpy as np
+import PIL
+from PIL import Image
+import torch
+
+
+def pil_to_tensor(
+    img: Image.Image,
+    target_image_size=512,
+    lock_ratio=True,
+    center_crop=True,
+    padding=False,
+    standardize=True,
+    **kwarg
+) -> torch.Tensor:
+    if img.mode != "RGB":
+        img = img.convert("RGB")
+
+    if isinstance(target_image_size, int):
+        target_size = (target_image_size, target_image_size)
+        if target_image_size < 0:
+            target_size = img.size
+    else:
+        target_size = target_image_size  # (width, height)
+
+    if lock_ratio:
+        original_width, original_height = img.size
+        target_width, target_height = target_size
+
+        scale_w = target_width / original_width
+        scale_h = target_height / original_height
+
+        if center_crop:
+            scale = max(scale_w, scale_h)
+        elif padding:
+            scale = min(scale_w, scale_h)
+        else:
+            scale = 1.0  # fallback
+
+        new_size = (round(original_width * scale), round(original_height * scale))
+        img = img.resize(new_size, Image.LANCZOS)
+
+        if center_crop:
+            left = (img.width - target_width) // 2
+            top = (img.height - target_height) // 2
+            img = img.crop((left, top, left + target_width, top + target_height))
+        elif padding:
+            new_img = Image.new("RGB", target_size, (0, 0, 0))
+            left = (target_width - img.width) // 2
+            top = (target_height - img.height) // 2
+            new_img.paste(img, (left, top))
+            img = new_img
+    else:
+        img = img.resize(target_size, Image.LANCZOS)
+
+    np_img = np.array(img) / 255.0  # Normalize to [0, 1]
+    if standardize:
+        np_img = np_img * 2 - 1  # Scale to [-1, 1]
+    tensor_img = torch.from_numpy(np_img).permute(2, 0, 1).float()  # (C, H, W)
+
+    return tensor_img
+
+
+def tensor_to_pil(chw_tensor: torch.Tensor, standardize=True, **kwarg) -> PIL.Image:
+    # Ensure detachment and move tensor to CPU.
+    detached_chw_tensor = chw_tensor.detach().cpu()
+
+    # Normalize tensor to [0, 1] range from [-1, 1] range.
+    if standardize:
+        normalized_chw_tensor = (
+            torch.clamp(detached_chw_tensor, -1.0, 1.0) + 1.0
+        ) / 2.0
+    else:
+        normalized_chw_tensor = torch.clamp(detached_chw_tensor, 0.0, 1.0)
+
+    # Permute CHW tensor to HWC format and convert to NumPy array.
+    hwc_array = normalized_chw_tensor.permute(1, 2, 0).numpy()
+
+    # Convert to an 8-bit unsigned integer format.
+    image_array_uint8 = (hwc_array * 255).astype(np.uint8)
+
+    # Convert NumPy array to PIL Image.
+    pil_image = Image.fromarray(image_array_uint8)
+
+    # Convert image to RGB if it is not already.
+    if pil_image.mode != "RGB":
+        pil_image = pil_image.convert("RGB")
+
+    return pil_image

src/model_processing.py

@@ -0,0 +1,409 @@
+import requests
+import os
+import yaml
+from .utils import get_ckpt, get_yaml_config
+
+
+def download_ckpt_yaml(model_path, model_name, ckpt_path, yaml_url=None):
+    def download_file(url, save_path):
+        response = requests.get(url)
+        response.raise_for_status()
+        with open(save_path, 'wb') as f:
+            f.write(response.content)
+
+    # os.makedirs(model_path, exist_ok=True)
+    local_dir = os.path.join(model_path, model_name)
+    os.makedirs(local_dir, exist_ok=True)
+
+    ckpt_name = ckpt_path.split("/")[-1]
+    local_ckpt_path = os.path.join(local_dir, ckpt_name)
+    if not os.path.exists(local_ckpt_path):
+        print(f"Downloading CKPT to {local_ckpt_path}")
+        download_file(ckpt_path, local_ckpt_path)
+
+    if yaml_url:
+        yaml_name = yaml_url.split("/")[-1]
+        local_yaml_path = os.path.join(local_dir, yaml_name)
+        if not os.path.exists(local_yaml_path):
+            print(f"Downloading YAML to {local_yaml_path}")
+            download_file(yaml_url, local_yaml_path)
+        return local_ckpt_path, local_yaml_path
+
+    return local_ckpt_path, None
+
+
+def get_model(model_path, model_name):
+    model = None
+    data_params = {
+        "target_image_size": (512, 512),
+        "lock_ratio": True,
+        "center_crop": True,
+        "padding": False,
+    }
+
+    if model_name.lower() == "anole":
+        from src.vqvaes.anole.anole import VQModel
+        yaml_url = "https://huggingface.co/GAIR/Anole-7b-v0.1/resolve/main/tokenizer/vqgan.yaml"
+        ckpt_path = "https://huggingface.co/GAIR/Anole-7b-v0.1/resolve/main/tokenizer/vqgan.ckpt"
+
+        if model_path is not None:
+            ckpt_path, yaml_url = download_ckpt_yaml(model_path, "anole", ckpt_path, yaml_url)
+        config = get_yaml_config(yaml_url)
+
+        params = config["model"]["params"]
+        if "lossconfig" in params:
+            del params["lossconfig"]
+        params["ckpt_path"] = ckpt_path
+        model = VQModel(**params)
+        data_params = {
+            "target_image_size": (512, 512),
+            "lock_ratio": True,
+            "center_crop": True,
+            "padding": False,
+        }
+
+    elif model_name.lower() == "chameleon":
+        from src.vqvaes.anole.anole import VQModel
+
+        yaml_url = "https://huggingface.co/huaweilin/chameleon_vqvae/resolve/main/vqgan.yaml"
+        ckpt_path = "https://huggingface.co/huaweilin/chameleon_vqvae/resolve/main/vqgan.ckpt"
+        if model_path is not None:
+            ckpt_path, yaml_url = download_ckpt_yaml(model_path, "chameleon", ckpt_path, yaml_url)
+        config = get_yaml_config(yaml_url)
+
+        params = config["model"]["params"]
+        if "lossconfig" in params:
+            del params["lossconfig"]
+        params["ckpt_path"] = ckpt_path
+        model = VQModel(**params)
+        data_params = {
+            "target_image_size": (512, 512),
+            "lock_ratio": True,
+            "center_crop": True,
+            "padding": False,
+        }
+
+    elif model_name.lower() == "llamagen-ds16":
+        from src.vqvaes.llamagen.llamagen import VQ_models
+        ckpt_path = "https://huggingface.co/FoundationVision/LlamaGen/resolve/main/vq_ds16_c2i.pt"
+        if model_path is not None:
+            ckpt_path, _ = download_ckpt_yaml(model_path, "llamagen-ds16", ckpt_path, None)
+
+        model = VQ_models["VQ-16"](codebook_size=16384, codebook_embed_dim=8)
+        model.load_state_dict(get_ckpt(ckpt_path, key="model"))
+        data_params = {
+            "target_image_size": (512, 512),
+            "lock_ratio": True,
+            "center_crop": True,
+            "padding": False,
+        }
+
+    elif model_name.lower() == "llamagen-ds16-t2i":
+        from src.vqvaes.llamagen.llamagen import VQ_models
+        ckpt_path = "https://huggingface.co/peizesun/llamagen_t2i/resolve/main/vq_ds16_t2i.pt"
+        if model_path is not None:
+            ckpt_path, _ = download_ckpt_yaml(model_path, "llamagen-ds16-t2i", ckpt_path, None)
+
+        model = VQ_models["VQ-16"](codebook_size=16384, codebook_embed_dim=8)
+        model.load_state_dict(get_ckpt(ckpt_path, key="model"))
+        data_params = {
+            "target_image_size": (512, 512),
+            "lock_ratio": True,
+            "center_crop": True,
+            "padding": False,
+        }
+
+    elif model_name.lower() == "llamagen-ds8":
+        from src.vqvaes.llamagen.llamagen import VQ_models
+        ckpt_path = "https://huggingface.co/FoundationVision/LlamaGen/resolve/main/vq_ds8_c2i.pt"
+        if model_path is not None:
+            ckpt_path, _ = download_ckpt_yaml(model_path, "llamagen-ds8", ckpt_path, None)
+
+        model = VQ_models["VQ-8"](codebook_size=16384, codebook_embed_dim=8)
+        model.load_state_dict(get_ckpt(ckpt_path, key="model"))
+        data_params = {
+            "target_image_size": (256, 256),
+            "lock_ratio": True,
+            "center_crop": True,
+            "padding": False,
+        }
+
+    elif model_name.lower() == "flowmo_lo":
+        from src.vqvaes.flowmo.flowmo import build_model
+        yaml_url = "https://raw.githubusercontent.com/kylesargent/FlowMo/refs/heads/main/flowmo/configs/base.yaml"
+        ckpt_path = "https://huggingface.co/ksarge/FlowMo/resolve/main/flowmo_lo.pth"
+        if model_path is not None:
+            ckpt_path, yaml_url = download_ckpt_yaml(model_path, "flowmo_lo", ckpt_path, yaml_url)
+        config = get_yaml_config(yaml_url)
+
+        config.model.context_dim = 18
+        model = build_model(config)
+        model.load_state_dict(
+            get_ckpt(ckpt_path, key="model_ema_state_dict")
+        )
+        data_params = {
+            "target_image_size": (256, 256),
+            "lock_ratio": True,
+            "center_crop": True,
+            "padding": False,
+        }
+
+    elif model_name.lower() == "flowmo_hi":
+        from src.vqvaes.flowmo.flowmo import build_model
+
+        yaml_url = "https://raw.githubusercontent.com/kylesargent/FlowMo/refs/heads/main/flowmo/configs/base.yaml"
+        ckpt_path = "https://huggingface.co/ksarge/FlowMo/resolve/main/flowmo_hi.pth"
+        if model_path is not None:
+            ckpt_path, yaml_url = download_ckpt_yaml(model_path, "flowmo_hi", ckpt_path, yaml_url)
+        config = get_yaml_config(yaml_url)
+
+        config.model.context_dim = 56
+        config.model.codebook_size_for_entropy = 14
+        model = build_model(config)
+        model.load_state_dict(
+            get_ckpt(ckpt_path, key="model_ema_state_dict")
+        )
+        data_params = {
+            "target_image_size": (256, 256),
+            "lock_ratio": True,
+            "center_crop": True,
+            "padding": False,
+        }
+
+    elif model_name.lower() == "open_magvit2":
+        from src.vqvaes.open_magvit2.open_magvit2 import VQModel
+
+        yaml_url = "https://raw.githubusercontent.com/TencentARC/SEED-Voken/refs/heads/main/configs/Open-MAGVIT2/gpu/imagenet_lfqgan_256_L.yaml"
+        ckpt_path = "https://huggingface.co/TencentARC/Open-MAGVIT2-Tokenizer-256-resolution/resolve/main/imagenet_256_L.ckpt"
+        if model_path is not None:
+            ckpt_path, yaml_url = download_ckpt_yaml(model_path, "open_magvit2", ckpt_path, yaml_url)
+        config = get_yaml_config(yaml_url)
+
+        model = VQModel(**config.model.init_args)
+        model.load_state_dict(get_ckpt(ckpt_path, key="state_dict"))
+        data_params = {
+            "target_image_size": (256, 256),
+            "lock_ratio": True,
+            "center_crop": True,
+            "padding": False,
+        }
+
+    elif "maskbit" in model_name.lower():
+        from src.vqvaes.maskbit.maskbit import ConvVQModel
+
+        if "16bit" in model_name.lower():
+            yaml_url = "https://raw.githubusercontent.com/markweberdev/maskbit/refs/heads/main/configs/tokenizer/maskbit_tokenizer_16bit.yaml"
+            ckpt_path = "https://huggingface.co/markweber/maskbit_tokenizer_16bit/resolve/main/maskbit_tokenizer_16bit.bin"
+            if model_path is not None:
+                ckpt_path, yaml_url = download_ckpt_yaml(model_path, "maskbit-16bit", ckpt_path, yaml_url)
+        elif "18bit" in model_name.lower():
+            yaml_url = "https://raw.githubusercontent.com/markweberdev/maskbit/refs/heads/main/configs/tokenizer/maskbit_tokenizer_18bit.yaml"
+            ckpt_path = "https://huggingface.co/markweber/maskbit_tokenizer_18bit/resolve/main/maskbit_tokenizer_18bit.bin"
+            if model_path is not None:
+                ckpt_path, yaml_url = download_ckpt_yaml(model_path, "maskbit-18bit", ckpt_path, yaml_url)
+        else:
+            raise Exception(f"Unsupported model: {model_name}")
+
+        config = get_yaml_config(yaml_url)
+        model = ConvVQModel(config.model.vq_model, legacy=False)
+        model.load_pretrained(get_ckpt(ckpt_path, key=None))
+        data_params = {
+            "target_image_size": (256, 256),
+            "lock_ratio": True,
+            "center_crop": True,
+            "padding": False,
+            "standardize": False,
+        }
+
+    elif "bsqvit" in model_name.lower():
+        from src.vqvaes.bsqvit.bsqvit import VITBSQModel
+
+        yaml_url = "https://huggingface.co/huaweilin/bsqvit_256x256/resolve/main/config.yaml"
+        ckpt_path = "https://huggingface.co/huaweilin/bsqvit_256x256/resolve/main/checkpoint.pt"
+        if model_path is not None:
+            ckpt_path, yaml_url = download_ckpt_yaml(model_path, "bsqvit", ckpt_path, yaml_url)
+
+        config = get_yaml_config(yaml_url)
+        model = VITBSQModel(**config["model"]["params"])
+        model.init_from_ckpt(get_ckpt(ckpt_path, key="state_dict"))
+        data_params = {
+            "target_image_size": (256, 256),
+            "lock_ratio": True,
+            "center_crop": True,
+            "padding": False,
+            "standardize": False,
+        }
+
+    elif "titok" in model_name.lower():
+        from src.vqvaes.titok.titok import TiTok
+
+        ckpt_path = None
+        if "bl64" in model_name.lower():
+            ckpt_path = "yucornetto/tokenizer_titok_bl64_vq8k_imagenet"
+        elif "bl128" in model_name.lower():
+            ckpt_path = "yucornetto/tokenizer_titok_bl128_vq8k_imagenet"
+        elif "sl256" in model_name.lower():
+            ckpt_path = "yucornetto/tokenizer_titok_sl256_vq8k_imagenet"
+        elif "l32" in model_name.lower():
+            ckpt_path = "yucornetto/tokenizer_titok_l32_imagenet"
+        elif "b64" in model_name.lower():
+            ckpt_path = "yucornetto/tokenizer_titok_b64_imagenet"
+        elif "s128" in model_name.lower():
+            ckpt_path = "yucornetto/tokenizer_titok_s128_imagenet"
+        else:
+            raise Exception(f"Unsupported model: {model_name}")
+
+        model = TiTok.from_pretrained(ckpt_path)
+        data_params = {
+            "target_image_size": (256, 256),
+            "lock_ratio": True,
+            "center_crop": True,
+            "padding": False,
+            "standardize": False,
+        }
+
+    elif "janus_pro" in model_name.lower():
+        from janus.models import MultiModalityCausalLM
+        from src.vqvaes.janus_pro.janus_pro import forward
+        import types
+
+        model = MultiModalityCausalLM.from_pretrained(
+            "deepseek-ai/Janus-Pro-7B", trust_remote_code=True
+        ).gen_vision_model
+        model.forward = types.MethodType(forward, model)
+        data_params = {
+            "target_image_size": (384, 384),
+            "lock_ratio": True,
+            "center_crop": False,
+            "padding": True,
+        }
+
+    elif "var" in model_name.lower():
+        from src.vqvaes.var.var_vq import VQVAE
+
+        ckpt_path = "https://huggingface.co/FoundationVision/var/resolve/main/vae_ch160v4096z32.pth"
+        if model_path is not None:
+            ckpt_path, _ = download_ckpt_yaml(model_path, "var", ckpt_path, None)
+
+        v_patch_nums = (1, 2, 3, 4, 5, 6, 8, 10, 13, 16)
+        if "512" in model_name.lower():
+            v_patch_nums = (1, 2, 3, 4, 6, 9, 13, 18, 24, 32)
+        model = VQVAE(
+            vocab_size=4096,
+            z_channels=32,
+            ch=160,
+            test_mode=True,
+            share_quant_resi=4,
+            v_patch_nums=v_patch_nums,
+        )
+        model.load_state_dict(get_ckpt(ckpt_path, key=None))
+        data_params = {
+            "target_image_size": (
+                (512, 512) if "512" in model_name.lower() else (256, 256)
+            ),
+            "lock_ratio": True,
+            "center_crop": False,
+            "padding": True,
+            "standardize": False,
+        }
+
+    elif (
+        "infinity" in model_name.lower()
+    ):  # "infinity_d32", "infinity_d64", "infinity_d56_f8_14_patchify"
+        from src.vqvaes.infinity.vae import vae_model
+
+        if "d32" in model_name:
+            ckpt_path = "https://huggingface.co/FoundationVision/Infinity/resolve/main/infinity_vae_d32.pth"
+            codebook_dim = 32
+            if model_path is not None:
+                ckpt_path, _ = download_ckpt_yaml(model_path, "infinity-d32", ckpt_path, None)
+        elif "d64" in model_name:
+            ckpt_path = "https://huggingface.co/FoundationVision/Infinity/resolve/main/infinity_vae_d64.pth"
+            codebook_dim = 64
+            if model_path is not None:
+                ckpt_path, _ = download_ckpt_yaml(model_path, "infinity-d64", ckpt_path, None)
+
+        schedule_mode = "dynamic"
+        codebook_size = 2**codebook_dim
+        patch_size = 16
+        encoder_ch_mult = [1, 2, 4, 4, 4]
+        decoder_ch_mult = [1, 2, 4, 4, 4]
+
+        ckpt = get_ckpt(ckpt_path, key=None)
+        model = vae_model(
+            ckpt,
+            schedule_mode,
+            codebook_dim,
+            codebook_size,
+            patch_size=patch_size,
+            encoder_ch_mult=encoder_ch_mult,
+            decoder_ch_mult=decoder_ch_mult,
+            test_mode=True,
+        )
+
+        data_params = {
+            "target_image_size": (1024, 1024),
+            "lock_ratio": True,
+            "center_crop": False,
+            "padding": True,
+            "standardize": False,
+        }
+
+    elif "sd3.5l" in model_name.lower():  # SD3.5L
+        from src.vaes.stable_diffusion.vae import forward
+        from diffusers import AutoencoderKL
+        import types
+
+        model = AutoencoderKL.from_pretrained(
+            "huaweilin/stable-diffusion-3.5-large-vae", subfolder="vae"
+        )
+        model.forward = types.MethodType(forward, model)
+        data_params = {
+            "target_image_size": (1024, 1024),
+            "lock_ratio": True,
+            "center_crop": False,
+            "padding": True,
+            "standardize": True,
+        }
+
+    elif "FLUX.1-dev".lower() in model_name.lower():  # SD3.5L
+        from src.vaes.stable_diffusion.vae import forward
+        from diffusers import AutoencoderKL
+        import types
+
+        model = AutoencoderKL.from_pretrained(
+            "black-forest-labs/FLUX.1-dev", subfolder="vae"
+        )
+        model.forward = types.MethodType(forward, model)
+        data_params = {
+            "target_image_size": (1024, 1024),
+            "lock_ratio": True,
+            "center_crop": False,
+            "padding": True,
+            "standardize": True,
+        }
+
+    elif "gpt4o" in model_name.lower():
+        from src.vaes.gpt_image.gpt_image import GPTImage
+
+        data_params = {
+            "target_image_size": (1024, 1024),
+            "lock_ratio": True,
+            "center_crop": False,
+            "padding": True,
+            "standardize": False,
+        }
+        model = GPTImage(data_params)
+
+    else:
+        raise Exception(f"Unsupported model: \"{model_name}\"")
+
+    try:
+        trainable_params = sum(p.numel() for p in model.parameters())
+        print("trainable_params:", trainable_params)
+    except Exception as e:
+        print(e)
+        pass
+
+    model.eval()
+    return model, data_params

src/utils.py

@@ -0,0 +1,47 @@
+import os
+from omegaconf import OmegaConf
+import torch
+import tempfile
+from safetensors.torch import load_file
+import requests
+import yaml
+
+def get_ckpt(path, key="state_dict"):
+    is_url = path.startswith("http://") or path.startswith("https://")
+    suffix = os.path.splitext(path)[-1]
+
+    if is_url:
+        print(f"Loading checkpoint from URL: {path}")
+        with tempfile.NamedTemporaryFile(suffix=suffix) as tmp_file:
+            response = requests.get(path)
+            response.raise_for_status()
+            tmp_file.write(response.content)
+            tmp_file.flush()
+            ckpt_path = tmp_file.name
+
+            if suffix == ".safetensors":
+                checkpoint = load_file(ckpt_path)
+            else:
+                checkpoint = torch.load(ckpt_path, map_location="cpu", weights_only=False)
+    else:
+        print(f"Loading checkpoint from local path: {path}")
+        if suffix == ".safetensors":
+            checkpoint = load_file(path)
+        else:
+            checkpoint = torch.load(path, map_location="cpu", weights_only=False)
+
+    if key is not None and key in checkpoint:
+        checkpoint = checkpoint[key]
+
+    return checkpoint
+
+
+def get_yaml_config(path):
+    if path.startswith("http://") or path.startswith("https://"):
+        response = requests.get(path)
+        response.raise_for_status()
+        config = OmegaConf.create(response.text)
+    else:
+        with open(path, 'r') as f:
+            config = OmegaConf.load(f)
+    return config

src/vaes/gpt_image/gpt_image.py

@@ -0,0 +1,48 @@
+import base64
+from torchvision.transforms.functional import to_pil_image
+from openai import OpenAI
+import io
+import torch
+import numpy as np
+from PIL import Image
+from ...data_processing import tensor_to_pil, pil_to_tensor
+
+
+class GPTImage:
+    def __init__(self, data_params):
+        self.client = OpenAI(organization="org-xZTnLOf1k9s04LEoKKjl4jOB")
+        self.prompt = "Please recreate the exact same image without any alterations. Please preserve the original resolution (1024*1024)."
+        self.data_params = data_params
+
+    def eval(self):
+        pass
+
+    def __call__(self, *args, **kwargs):
+        return self.forward(*args, **kwargs)
+
+    def forward(self, input):
+        results = []
+        for image in input:
+            image = tensor_to_pil(image, **self.data_params)
+            buffer = io.BytesIO()
+            image.save(buffer, format="PNG")
+            buffer.seek(0)
+            image_file = ("image.png", buffer, "image/png")
+
+            try:
+                result = self.client.images.edit(
+                    model="gpt-image-1",
+                    image=image_file,
+                    prompt=self.prompt,
+                    n=1,
+                    size="1024x1024",
+                )
+                image_base64 = result.data[0].b64_json
+                image_bytes = base64.b64decode(image_base64)
+                image = Image.open(io.BytesIO(image_bytes))
+                results.append(pil_to_tensor(image, **self.data_params))
+            except Exception as e:
+                print("💥 Unexpected error occurred:", e)
+                results.append(None)
+
+        return results, None, None

src/vaes/stable_diffusion/vae.py

@@ -0,0 +1,23 @@
+def forward(
+    self,
+    sample,
+    sample_posterior=False,
+    return_dict=True,
+    generator=None,
+):
+    r"""
+    Args:
+        sample (`torch.Tensor`): Input sample.
+        sample_posterior (`bool`, *optional*, defaults to `False`):
+            Whether to sample from the posterior.
+        return_dict (`bool`, *optional*, defaults to `True`):
+            Whether or not to return a [`DecoderOutput`] instead of a plain tuple.
+    """
+    x = sample
+    posterior = self.encode(x).latent_dist
+    if sample_posterior:
+        z = posterior.sample(generator=generator)
+    else:
+        z = posterior.mode()
+    dec = self.decode(z).sample
+    return dec, None, None

src/vqvaes/anole/anole.py

@@ -0,0 +1,706 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+
+# This source code is licensed under the Chameleon License found in the
+# LICENSE file in the root directory of this source tree.
+
+"""
+Contents of this file are taken from https://github.com/CompVis/taming-transformers/blob/3ba01b241669f5ade541ce990f7650a3b8f65318/taming/models/vqgan.py
+[with minimal dependencies]
+
+This implementation is inference-only -- training steps and optimizer components
+introduce significant additional dependencies
+"""
+
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from ...utils import get_ckpt
+
+
+class VectorQuantizer2(nn.Module):
+    """
+    Improved version over VectorQuantizer, can be used as a drop-in replacement. Mostly
+    avoids costly matrix multiplications and allows for post-hoc remapping of indices.
+    """
+
+    # NOTE: due to a bug the beta term was applied to the wrong term. for
+    # backwards compatibility we use the buggy version by default, but you can
+    # specify legacy=False to fix it.
+    def __init__(
+        self,
+        n_e,
+        e_dim,
+        beta,
+        remap=None,
+        unknown_index="random",
+        sane_index_shape=False,
+        legacy=True,
+    ):
+        super().__init__()
+        self.n_e = n_e
+        self.e_dim = e_dim
+        self.beta = beta
+        self.legacy = legacy
+
+        self.embedding = nn.Embedding(self.n_e, self.e_dim)
+        self.embedding.weight.data.uniform_(-1.0 / self.n_e, 1.0 / self.n_e)
+
+        self.remap = remap
+        if self.remap is not None:
+            self.register_buffer("used", torch.tensor(np.load(self.remap)))
+            self.re_embed = self.used.shape[0]
+            self.unknown_index = unknown_index  # "random" or "extra" or integer
+            if self.unknown_index == "extra":
+                self.unknown_index = self.re_embed
+                self.re_embed = self.re_embed + 1
+            print(
+                f"Remapping {self.n_e} indices to {self.re_embed} indices. "
+                f"Using {self.unknown_index} for unknown indices."
+            )
+        else:
+            self.re_embed = n_e
+
+        self.sane_index_shape = sane_index_shape
+
+    def remap_to_used(self, inds):
+        ishape = inds.shape
+        assert len(ishape) > 1
+        inds = inds.reshape(ishape[0], -1)
+        used = self.used.to(inds)
+        match = (inds[:, :, None] == used[None, None, ...]).long()
+        new = match.argmax(-1)
+        unknown = match.sum(2) < 1
+        if self.unknown_index == "random":
+            new[unknown] = torch.randint(0, self.re_embed, size=new[unknown].shape).to(
+                device=new.device
+            )
+        else:
+            new[unknown] = self.unknown_index
+        return new.reshape(ishape)
+
+    def unmap_to_all(self, inds):
+        ishape = inds.shape
+        assert len(ishape) > 1
+        inds = inds.reshape(ishape[0], -1)
+        used = self.used.to(inds)
+        if self.re_embed > self.used.shape[0]:  # extra token
+            inds[inds >= self.used.shape[0]] = 0  # simply set to zero
+        back = torch.gather(used[None, :][inds.shape[0] * [0], :], 1, inds)
+        return back.reshape(ishape)
+
+    def forward(self, z, temp=None, rescale_logits=False, return_logits=False):
+        assert temp is None or temp == 1.0, "Only for interface compatible with Gumbel"
+        assert rescale_logits is False, "Only for interface compatible with Gumbel"
+        assert return_logits is False, "Only for interface compatible with Gumbel"
+        # reshape z -> (batch, height, width, channel) and flatten
+        z = z.permute(0, 2, 3, 1).contiguous()
+        z_flattened = z.view(-1, self.e_dim)
+        # distances from z to embeddings e_j (z - e)^2 = z^2 + e^2 - 2 e * z
+
+        d = (
+            torch.sum(z_flattened**2, dim=1, keepdim=True)
+            + torch.sum(self.embedding.weight**2, dim=1)
+            - 2
+            * torch.einsum(
+                "bd,dn->bn", z_flattened, self.embedding.weight.transpose(0, 1)
+            )
+        )
+
+        min_encoding_indices = torch.argmin(d, dim=1)
+        z_q = self.embedding(min_encoding_indices).view(z.shape)
+        perplexity = None
+        min_encodings = None
+
+        # compute loss for embedding
+        if not self.legacy:
+            loss = self.beta * torch.mean((z_q.detach() - z) ** 2) + torch.mean(
+                (z_q - z.detach()) ** 2
+            )
+        else:
+            loss = torch.mean((z_q.detach() - z) ** 2) + self.beta * torch.mean(
+                (z_q - z.detach()) ** 2
+            )
+
+        # preserve gradients
+        z_q = z + (z_q - z).detach()
+
+        # reshape back to match original input shape
+        z_q = z_q.permute(0, 3, 1, 2).contiguous()
+
+        if self.remap is not None:
+            min_encoding_indices = min_encoding_indices.reshape(
+                z.shape[0], -1
+            )  # add batch axis
+            min_encoding_indices = self.remap_to_used(min_encoding_indices)
+            min_encoding_indices = min_encoding_indices.reshape(-1, 1)  # flatten
+
+        if self.sane_index_shape:
+            min_encoding_indices = min_encoding_indices.reshape(
+                z_q.shape[0], z_q.shape[2], z_q.shape[3]
+            )
+
+        return z_q, loss, (perplexity, min_encodings, min_encoding_indices)
+
+    def get_codebook_entry(self, indices, shape):
+        # shape specifying (batch, height, width, channel)
+        if self.remap is not None:
+            indices = indices.reshape(shape[0], -1)  # add batch axis
+            indices = self.unmap_to_all(indices)
+            indices = indices.reshape(-1)  # flatten again
+
+        # get quantized latent vectors
+        z_q = self.embedding(indices)
+
+        if shape is not None:
+            z_q = z_q.view(shape)
+            # reshape back to match original input shape
+            z_q = z_q.permute(0, 3, 1, 2).contiguous()
+
+        return z_q
+
+
+# Alias
+VectorQuantizer = VectorQuantizer2
+
+
+def nonlinearity(x):
+    # swish
+    return x * torch.sigmoid(x)
+
+
+def Normalize(in_channels, num_groups=32):
+    return torch.nn.GroupNorm(
+        num_groups=num_groups, num_channels=in_channels, eps=1e-6, affine=True
+    )
+
+
+class Upsample(nn.Module):
+    def __init__(self, in_channels, with_conv):
+        super().__init__()
+        self.with_conv = with_conv
+        if self.with_conv:
+            self.conv = torch.nn.Conv2d(
+                in_channels, in_channels, kernel_size=3, stride=1, padding=1
+            )
+
+    def forward(self, x):
+        x = F.interpolate(x, scale_factor=2.0, mode="nearest")
+        if self.with_conv:
+            x = self.conv(x)
+        return x
+
+
+class Downsample(nn.Module):
+    def __init__(self, in_channels, with_conv):
+        super().__init__()
+        self.with_conv = with_conv
+        if self.with_conv:
+            # no asymmetric padding in torch conv, must do it ourselves
+            self.conv = torch.nn.Conv2d(
+                in_channels, in_channels, kernel_size=3, stride=2, padding=0
+            )
+
+    def forward(self, x):
+        if self.with_conv:
+            pad = (0, 1, 0, 1)
+            x = F.pad(x, pad, mode="constant", value=0)
+            x = self.conv(x)
+        else:
+            x = F.avg_pool2d(x, kernel_size=2, stride=2)
+        return x
+
+
+class ResnetBlock(nn.Module):
+    def __init__(
+        self,
+        *,
+        in_channels,
+        out_channels=None,
+        conv_shortcut=False,
+        dropout,
+        temb_channels=512,
+    ):
+        super().__init__()
+        self.in_channels = in_channels
+        out_channels = in_channels if out_channels is None else out_channels
+        self.out_channels = out_channels
+        self.use_conv_shortcut = conv_shortcut
+
+        self.norm1 = Normalize(in_channels)
+        self.conv1 = torch.nn.Conv2d(
+            in_channels, out_channels, kernel_size=3, stride=1, padding=1
+        )
+        if temb_channels > 0:
+            self.temb_proj = torch.nn.Linear(temb_channels, out_channels)
+        self.norm2 = Normalize(out_channels)
+        self.dropout = torch.nn.Dropout(dropout)
+        self.conv2 = torch.nn.Conv2d(
+            out_channels, out_channels, kernel_size=3, stride=1, padding=1
+        )
+        if self.in_channels != self.out_channels:
+            if self.use_conv_shortcut:
+                self.conv_shortcut = torch.nn.Conv2d(
+                    in_channels, out_channels, kernel_size=3, stride=1, padding=1
+                )
+            else:
+                self.nin_shortcut = torch.nn.Conv2d(
+                    in_channels, out_channels, kernel_size=1, stride=1, padding=0
+                )
+
+    def forward(self, x, temb):
+        h = x
+        h = self.norm1(h)
+        h = nonlinearity(h)
+        h = self.conv1(h)
+
+        if temb is not None:
+            h = h + self.temb_proj(nonlinearity(temb))[:, :, None, None]
+
+        h = self.norm2(h)
+        h = nonlinearity(h)
+        h = self.dropout(h)
+        h = self.conv2(h)
+
+        if self.in_channels != self.out_channels:
+            if self.use_conv_shortcut:
+                x = self.conv_shortcut(x)
+            else:
+                x = self.nin_shortcut(x)
+
+        return x + h
+
+
+class AttnBlock(nn.Module):
+    def __init__(self, in_channels):
+        super().__init__()
+        self.in_channels = in_channels
+
+        self.norm = Normalize(in_channels)
+        self.q = torch.nn.Conv2d(
+            in_channels, in_channels, kernel_size=1, stride=1, padding=0
+        )
+        self.k = torch.nn.Conv2d(
+            in_channels, in_channels, kernel_size=1, stride=1, padding=0
+        )
+        self.v = torch.nn.Conv2d(
+            in_channels, in_channels, kernel_size=1, stride=1, padding=0
+        )
+        self.proj_out = torch.nn.Conv2d(
+            in_channels, in_channels, kernel_size=1, stride=1, padding=0
+        )
+
+    def forward(self, x):
+        h_ = x
+        h_ = self.norm(h_)
+        q = self.q(h_)
+        k = self.k(h_)
+        v = self.v(h_)
+
+        # compute attention
+        b, c, h, w = q.shape
+        q = q.reshape(b, c, h * w)
+        q = q.permute(0, 2, 1)  # b,hw,c
+        k = k.reshape(b, c, h * w)  # b,c,hw
+        w_ = torch.bmm(q, k)  # b,hw,hw    w[b,i,j]=sum_c q[b,i,c]k[b,c,j]
+        w_ = w_ * (int(c) ** (-0.5))
+        w_ = F.softmax(w_, dim=2)
+
+        # attend to values
+        v = v.reshape(b, c, h * w)
+        w_ = w_.permute(0, 2, 1)  # b,hw,hw (first hw of k, second of q)
+        h_ = torch.bmm(v, w_)  # b, c,hw (hw of q) h_[b,c,j] = sum_i v[b,c,i] w_[b,i,j]
+        h_ = h_.reshape(b, c, h, w)
+
+        h_ = self.proj_out(h_)
+
+        return x + h_
+
+
+def make_attn(in_channels, attn_type="vanilla"):
+    assert attn_type in ["vanilla", "linear", "none"], f"attn_type {attn_type} unknown"
+    # print(f"making attention of type '{attn_type}' with {in_channels} in_channels")
+    if attn_type == "vanilla":
+        return AttnBlock(in_channels)
+    elif attn_type == "none":
+        return nn.Identity(in_channels)
+    else:
+        raise ValueError("Unexpected attention type")
+
+
+class Encoder(nn.Module):
+    def __init__(
+        self,
+        *,
+        ch,
+        out_ch,
+        ch_mult=(1, 2, 4, 8),
+        num_res_blocks,
+        attn_resolutions,
+        dropout=0.0,
+        resamp_with_conv=True,
+        in_channels,
+        resolution,
+        z_channels,
+        double_z=True,
+        use_linear_attn=False,
+        attn_type="vanilla",
+        **ignore_kwargs,
+    ):
+        super().__init__()
+        if use_linear_attn:
+            attn_type = "linear"
+        self.ch = ch
+        self.temb_ch = 0
+        self.num_resolutions = len(ch_mult)
+        self.num_res_blocks = num_res_blocks
+        self.resolution = resolution
+        self.in_channels = in_channels
+
+        # downsampling
+        self.conv_in = torch.nn.Conv2d(
+            in_channels, self.ch, kernel_size=3, stride=1, padding=1
+        )
+
+        curr_res = resolution
+        in_ch_mult = (1,) + tuple(ch_mult)
+        self.in_ch_mult = in_ch_mult
+        self.down = nn.ModuleList()
+        for i_level in range(self.num_resolutions):
+            block = nn.ModuleList()
+            attn = nn.ModuleList()
+            block_in = ch * in_ch_mult[i_level]
+            block_out = ch * ch_mult[i_level]
+            for i_block in range(self.num_res_blocks):
+                block.append(
+                    ResnetBlock(
+                        in_channels=block_in,
+                        out_channels=block_out,
+                        temb_channels=self.temb_ch,
+                        dropout=dropout,
+                    )
+                )
+                block_in = block_out
+                if curr_res in attn_resolutions:
+                    attn.append(make_attn(block_in, attn_type=attn_type))
+            down = nn.Module()
+            down.block = block
+            down.attn = attn
+            if i_level != self.num_resolutions - 1:
+                down.downsample = Downsample(block_in, resamp_with_conv)
+                curr_res = curr_res // 2
+            self.down.append(down)
+
+        # middle
+        self.mid = nn.Module()
+        self.mid.block_1 = ResnetBlock(
+            in_channels=block_in,
+            out_channels=block_in,
+            temb_channels=self.temb_ch,
+            dropout=dropout,
+        )
+        self.mid.attn_1 = make_attn(block_in, attn_type=attn_type)
+        self.mid.block_2 = ResnetBlock(
+            in_channels=block_in,
+            out_channels=block_in,
+            temb_channels=self.temb_ch,
+            dropout=dropout,
+        )
+
+        # end
+        self.norm_out = Normalize(block_in)
+        self.conv_out = torch.nn.Conv2d(
+            block_in,
+            2 * z_channels if double_z else z_channels,
+            kernel_size=3,
+            stride=1,
+            padding=1,
+        )
+
+    def forward(self, x):
+        # timestep embedding
+        temb = None
+
+        # downsampling
+        hs = [self.conv_in(x)]
+        for i_level in range(self.num_resolutions):
+            for i_block in range(self.num_res_blocks):
+                h = self.down[i_level].block[i_block](hs[-1], temb)
+                if len(self.down[i_level].attn) > 0:
+                    h = self.down[i_level].attn[i_block](h)
+                hs.append(h)
+            if i_level != self.num_resolutions - 1:
+                hs.append(self.down[i_level].downsample(hs[-1]))
+
+        # middle
+        h = hs[-1]
+        h = self.mid.block_1(h, temb)
+        h = self.mid.attn_1(h)
+        h = self.mid.block_2(h, temb)
+
+        # end
+        h = self.norm_out(h)
+        h = nonlinearity(h)
+        h = self.conv_out(h)
+        return h
+
+
+class Decoder(nn.Module):
+    def __init__(
+        self,
+        *,
+        ch,
+        out_ch,
+        ch_mult=(1, 2, 4, 8),
+        num_res_blocks,
+        attn_resolutions,
+        dropout=0.0,
+        resamp_with_conv=True,
+        in_channels,
+        resolution,
+        z_channels,
+        give_pre_end=False,
+        tanh_out=False,
+        use_linear_attn=False,
+        attn_type="vanilla",
+        **ignorekwargs,
+    ):
+        super().__init__()
+        if use_linear_attn:
+            attn_type = "linear"
+        self.ch = ch
+        self.temb_ch = 0
+        self.num_resolutions = len(ch_mult)
+        self.num_res_blocks = num_res_blocks
+        self.resolution = resolution
+        self.in_channels = in_channels
+        self.give_pre_end = give_pre_end
+        self.tanh_out = tanh_out
+
+        # compute in_ch_mult, block_in and curr_res at lowest res
+        block_in = ch * ch_mult[self.num_resolutions - 1]
+        curr_res = resolution // 2 ** (self.num_resolutions - 1)
+        self.z_shape = (1, z_channels, curr_res, curr_res)
+
+        # z to block_in
+        self.conv_in = torch.nn.Conv2d(
+            z_channels, block_in, kernel_size=3, stride=1, padding=1
+        )
+
+        # middle
+        self.mid = nn.Module()
+        self.mid.block_1 = ResnetBlock(
+            in_channels=block_in,
+            out_channels=block_in,
+            temb_channels=self.temb_ch,
+            dropout=dropout,
+        )
+        self.mid.attn_1 = make_attn(block_in, attn_type=attn_type)
+        self.mid.block_2 = ResnetBlock(
+            in_channels=block_in,
+            out_channels=block_in,
+            temb_channels=self.temb_ch,
+            dropout=dropout,
+        )
+
+        # upsampling
+        self.up = nn.ModuleList()
+        for i_level in reversed(range(self.num_resolutions)):
+            block = nn.ModuleList()
+            attn = nn.ModuleList()
+            block_out = ch * ch_mult[i_level]
+            for i_block in range(self.num_res_blocks + 1):
+                block.append(
+                    ResnetBlock(
+                        in_channels=block_in,
+                        out_channels=block_out,
+                        temb_channels=self.temb_ch,
+                        dropout=dropout,
+                    )
+                )
+                block_in = block_out
+                if curr_res in attn_resolutions:
+                    attn.append(make_attn(block_in, attn_type=attn_type))
+            up = nn.Module()
+            up.block = block
+            up.attn = attn
+            if i_level != 0:
+                up.upsample = Upsample(block_in, resamp_with_conv)
+                curr_res = curr_res * 2
+            self.up.insert(0, up)  # prepend to get consistent order
+
+        # end
+        self.norm_out = Normalize(block_in)
+        self.conv_out = torch.nn.Conv2d(
+            block_in, out_ch, kernel_size=3, stride=1, padding=1
+        )
+
+    def forward(self, z):
+        # assert z.shape[1:] == self.z_shape[1:]
+        self.last_z_shape = z.shape
+
+        # timestep embedding
+        temb = None
+
+        # z to block_in
+        h = self.conv_in(z)
+
+        # middle
+        h = self.mid.block_1(h, temb)
+        h = self.mid.attn_1(h)
+        h = self.mid.block_2(h, temb)
+
+        # upsampling
+        for i_level in reversed(range(self.num_resolutions)):
+            for i_block in range(self.num_res_blocks + 1):
+                h = self.up[i_level].block[i_block](h, temb)
+                if len(self.up[i_level].attn) > 0:
+                    h = self.up[i_level].attn[i_block](h)
+            if i_level != 0:
+                h = self.up[i_level].upsample(h)
+
+        # end
+        if self.give_pre_end:
+            return h
+
+        h = self.norm_out(h)
+        h = nonlinearity(h)
+        h = self.conv_out(h)
+        if self.tanh_out:
+            h = torch.tanh(h)
+        return h
+
+
+class VQModel(nn.Module):
+    def __init__(
+        self,
+        ddconfig,
+        n_embed,
+        embed_dim,
+        ckpt_path=None,
+        ignore_keys=[],
+        image_key="image",
+        colorize_nlabels=None,
+        monitor=None,
+        scheduler_config=None,
+        lr_g_factor=1.0,
+        remap=None,
+        sane_index_shape=False,  # tell vector quantizer to return indices as bhw
+    ):
+        super().__init__()
+        self.image_key = image_key
+        self.encoder = Encoder(**ddconfig)
+        self.decoder = Decoder(**ddconfig)
+        self.quantize = VectorQuantizer(
+            n_embed,
+            embed_dim,
+            beta=0.25,
+            remap=remap,
+            sane_index_shape=sane_index_shape,
+        )
+        self.quant_conv = torch.nn.Conv2d(ddconfig["z_channels"], embed_dim, 1)
+        self.post_quant_conv = torch.nn.Conv2d(embed_dim, ddconfig["z_channels"], 1)
+        if ckpt_path is not None:
+            self.init_from_ckpt(ckpt_path, ignore_keys=ignore_keys)
+        self.image_key = image_key
+        if colorize_nlabels is not None:
+            assert isinstance(colorize_nlabels, int)
+            self.register_buffer("colorize", torch.randn(3, colorize_nlabels, 1, 1))
+        if monitor is not None:
+            self.monitor = monitor
+        self.scheduler_config = scheduler_config
+        self.lr_g_factor = lr_g_factor
+
+    def init_from_ckpt(self, path, ignore_keys=list()):
+        if path.startswith("http://") or path.startswith("https://"):
+            sd = get_ckpt(path)
+        else:
+            print(f"Loading checkpoint from local path: {path}")
+            sd = torch.load(path, map_location="cpu")["state_dict"]
+
+        keys = list(sd.keys())
+        for k in keys:
+            for ik in ignore_keys:
+                if k.startswith(ik):
+                    print(f"Deleting key {k} from state_dict.")
+                    del sd[k]
+
+        self.load_state_dict(sd, strict=False)
+        print(f"VQModel loaded from {path}")
+
+    def encode(self, x):
+        h = self.encoder(x)
+        h = self.quant_conv(h)
+        quant, emb_loss, info = self.quantize(h)
+        return quant, emb_loss, info
+
+    def decode(self, quant):
+        quant = self.post_quant_conv(quant)
+        dec = self.decoder(quant)
+        return dec
+
+    def decode_code(self, code_b):
+        quant_b = self.quantize.embed_code(code_b)
+        dec = self.decode(quant_b)
+        return dec
+
+    #     def forward(self, input):
+    #         quant, diff, _ = self.encode(input)
+    #         dec = self.decode(quant)
+    #         return dec, diff
+
+    def forward(self, input):
+        quant, diff, [_, _, img_toks] = self.encode(input)
+
+        batch_size, n_channel, height, width = (
+            input.shape[0],
+            quant.shape[-1],
+            quant.shape[-2],
+            quant.shape[-3],
+        )
+        codebook_entry = self.quantize.get_codebook_entry(
+            img_toks, (batch_size, n_channel, height, width)
+        )
+        pixels = self.decode(codebook_entry)
+
+        return pixels, img_toks, quant
+
+    def get_input(self, batch, k):
+        x = batch[k]
+        if len(x.shape) == 3:
+            x = x[..., None]
+        x = x.permute(0, 3, 1, 2).to(memory_format=torch.contiguous_format)
+        return x.float()
+
+    def get_last_layer(self):
+        return self.decoder.conv_out.weight
+
+    def log_images(self, batch, **kwargs):
+        log = dict()
+        x = self.get_input(batch, self.image_key)
+        x = x.to(self.device)
+        xrec, _ = self(x)
+        if x.shape[1] > 3:
+            # colorize with random projection
+            assert xrec.shape[1] > 3
+            x = self.to_rgb(x)
+            xrec = self.to_rgb(xrec)
+        log["inputs"] = x
+        log["reconstructions"] = xrec
+        return log
+
+    @property
+    def device(self):
+        return next(self.parameters()).device
+
+    @property
+    def dtype(self):
+        return next(self.parameters()).dtype
+
+    def to_rgb(self, x):
+        assert self.image_key == "segmentation"
+        if not hasattr(self, "colorize"):
+            self.register_buffer("colorize", torch.randn(3, x.shape[1], 1, 1).to(x))
+        x = F.conv2d(x, weight=self.colorize)
+        x = 2.0 * (x - x.min()) / (x.max() - x.min()) - 1.0
+        return x

src/vqvaes/bsqvit/attention_mask.py

@@ -0,0 +1,42 @@
+import torch
+
+
+def get_attention_mask(sequence_length, device, mask_type="block-causal", **kwargs):
+    if mask_type.lower() == 'none' or mask_type is None:
+        return None
+    elif mask_type.lower() == 'block-causal':
+        return _block_caulsal_mask_impl(sequence_length, device, **kwargs)
+    elif mask_type.lower() == 'causal':
+        return _caulsal_mask_impl(sequence_length, device, **kwargs)
+    else:
+        raise NotImplementedError(f"Mask type {mask_type} not implemented")
+
+
+def _block_caulsal_mask_impl(sequence_length, device, block_size=16, **kwargs):
+    """
+    Create a block-causal mask
+    """
+    assert sequence_length % block_size == 0, "for block causal masks sequence length must be divisible by block size"
+    blocks = torch.ones(sequence_length // block_size, block_size, block_size, device=device)
+    block_diag_enable_mask = torch.block_diag(*blocks)
+    causal_enable_mask = torch.ones(sequence_length, sequence_length, device=device).tril_(0)
+    disable_mask = ((block_diag_enable_mask + causal_enable_mask) < 0.5)
+    return disable_mask
+
+
+def _caulsal_mask_impl(sequence_length, device, **kwargs):
+    """
+    Create a causal mask
+    """
+    causal_disable_mask = torch.triu(
+        torch.full((sequence_length, sequence_length), float('-inf'), dtype=torch.float32, device=device),
+        diagonal=1,
+    )
+    return causal_disable_mask
+
+
+if __name__ == '__main__':
+    mask = get_attention_mask(9, "cuda", mask_type="block-causal", block_size=3)
+    print(mask)
+    mask = get_attention_mask(9, "cuda", mask_type="causal")
+    print(mask)

src/vqvaes/bsqvit/bsqvit.py

@@ -0,0 +1,150 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from .quantizer.bsq import BinarySphericalQuantizer
+from .quantizer.vq import VectorQuantizer
+from .transformer import TransformerDecoder, TransformerEncoder
+
+
+class VITVQModel(nn.Module):
+    def __init__(self, vitconfig, n_embed, embed_dim,
+                 l2_norm=False, logit_laplace=False, ckpt_path=None, ignore_keys=[],
+                 grad_checkpointing=False, selective_checkpointing=False,
+                 clamp_range=(0, 1),
+                 dvitconfig=None,
+                 ):
+        super().__init__()
+        self.encoder = TransformerEncoder(**vitconfig)
+        dvitconfig = vitconfig if dvitconfig is None else dvitconfig
+        self.decoder = TransformerDecoder(**dvitconfig, logit_laplace=logit_laplace)
+        if self.training and grad_checkpointing:
+            self.encoder.set_grad_checkpointing(True, selective=selective_checkpointing)
+            self.decoder.set_grad_checkpointing(True, selective=selective_checkpointing)
+        
+        self.n_embed = n_embed
+        self.embed_dim = embed_dim
+        self.l2_norm = l2_norm
+        self.setup_quantizer()
+        
+        self.quant_embed = nn.Linear(in_features=vitconfig['width'], out_features=embed_dim)
+        self.post_quant_embed = nn.Linear(in_features=embed_dim, out_features=dvitconfig['width'])
+        self.l2_norm = l2_norm
+        self.logit_laplace = logit_laplace
+        self.clamp_range = clamp_range
+
+        if ckpt_path is not None:
+            self.init_from_ckpt(ckpt_path, ignore_keys=ignore_keys)
+    
+    def setup_quantizer(self):
+        self.quantize = VectorQuantizer(self.n_embed, self.embed_dim, l2_norm=self.l2_norm, beta=0.25, input_format='blc')
+
+    # def init_from_ckpt(self, ckpt_path, ignore_keys=[]):
+    def init_from_ckpt(self, state_dict, ignore_keys=[]):
+        state_dict = {k[7:]: v for k, v in state_dict.items() if k.startswith('module.')}
+        filtered_state_dict = {k: v for k, v in state_dict.items() if all([not k.startswith(ig) for ig in ignore_keys])}
+        missing_keys, unexpected_keys = self.load_state_dict(filtered_state_dict, strict=False)
+        print(f"missing_keys: {missing_keys}")
+        print(f"unexpected_keys: {unexpected_keys}")
+
+    def encode(self, x, skip_quantize=False):
+        h = self.encoder(x)
+        h = self.quant_embed(h)
+        if skip_quantize:
+            assert not self.training, 'skip_quantize should be used in eval mode only.'
+            if self.l2_norm:
+                h = F.normalize(h, dim=-1)
+            return h, {}, {}
+        quant, loss, info = self.quantize(h)
+        return quant, loss, info
+
+    def decode(self, quant):
+        h = self.post_quant_embed(quant)
+        x = self.decoder(h)
+        return x
+
+    def clamp(self, x):
+        if self.logit_laplace:
+            dec, _ = x.chunk(2, dim=1)
+            x = self.logit_laplace_loss.unmap(F.sigmoid(dec))
+        else:
+            x = x.clamp_(self.clamp_range[0], self.clamp_range[1])
+        return x
+
+    def forward(self, input, skip_quantize=False):
+        if self.logit_laplace:
+            input = self.logit_laplace_loss.inmap(input)
+        quant, loss, info = self.encode(input, skip_quantize=skip_quantize)
+        dec = self.decode(quant)
+        if self.logit_laplace:
+            dec, lnb = dec.chunk(2, dim=1)
+            logit_laplace_loss = self.logit_laplace_loss(dec, lnb, input)
+            info.update({'logit_laplace_loss': logit_laplace_loss})
+            dec = self.logit_laplace_loss.unmap(F.sigmoid(dec))
+        else:
+            dec = dec.clamp_(self.clamp_range[0], self.clamp_range[1])
+        return dec, loss, info
+
+    def get_last_layer(self):
+        return self.decoder.conv_out.weight
+
+
+class VITBSQModel(VITVQModel):
+    def __init__(self, vitconfig, embed_dim, embed_group_size=9,
+                 l2_norm=False, logit_laplace=False, ckpt_path=None, ignore_keys=[],
+                 grad_checkpointing=False, selective_checkpointing=False,
+                 clamp_range=(0, 1),
+                 dvitconfig=None, beta=0., gamma0=1.0, gamma=1.0, zeta=1.0,
+                 persample_entropy_compute='group',
+                 cb_entropy_compute='group',
+                 post_q_l2_norm=False,
+                 inv_temperature=1.,
+                 ):
+        # set quantizer params
+        self.beta = beta      # commit loss
+        self.gamma0 = gamma0  # entropy
+        self.gamma = gamma    # entropy penalty
+        self.zeta = zeta      # lpips
+        self.embed_group_size = embed_group_size
+        self.persample_entropy_compute = persample_entropy_compute
+        self.cb_entropy_compute = cb_entropy_compute
+        self.post_q_l2_norm = post_q_l2_norm
+        self.inv_temperature = inv_temperature
+        
+        # call init
+        super().__init__(
+            vitconfig,
+            2 ** embed_dim,
+            embed_dim,
+            l2_norm=l2_norm,
+            logit_laplace=logit_laplace,
+            ckpt_path=ckpt_path,
+            ignore_keys=ignore_keys,
+            grad_checkpointing=grad_checkpointing,
+            selective_checkpointing=selective_checkpointing,
+            clamp_range=clamp_range,
+            dvitconfig=dvitconfig,
+        )
+        
+
+    def setup_quantizer(self):
+        self.quantize = BinarySphericalQuantizer(
+            self.embed_dim, self.beta, self.gamma0, self.gamma, self.zeta,
+            group_size=self.embed_group_size,
+            persample_entropy_compute=self.persample_entropy_compute,
+            cb_entropy_compute=self.cb_entropy_compute,
+            input_format='blc',
+            l2_norm=self.post_q_l2_norm,
+            inv_temperature=self.inv_temperature,
+        )
+
+    def encode(self, x, skip_quantize=False):
+        h = self.encoder(x)
+        h = self.quant_embed(h)
+        if self.l2_norm:
+            h = F.normalize(h, dim=-1)
+        if skip_quantize:
+            assert not self.training, 'skip_quantize should be used in eval mode only.'
+            return h, {}, {}
+        quant, loss, info = self.quantize(h)
+        return quant, loss, info

src/vqvaes/bsqvit/quantizer/bsq.py

@@ -0,0 +1,223 @@
+from einops import rearrange, reduce
+import torch
+import torch.nn as nn
+from torch.autograd import Function
+
+
+class DifferentiableEntropyFunction(Function):
+    @staticmethod
+    def forward(ctx, zq, basis, K, eps):
+        zb = (zq + 1) / 2
+        zi = ((zb * basis).sum(-1)).to(torch.int64)
+        cnt = torch.scatter_reduce(torch.zeros(2**K, device=zq.device, dtype=zq.dtype), 
+                                   0, 
+                                   zi.flatten(), 
+                                   torch.ones_like(zi.flatten()).to(zq.dtype),
+                                    'sum')
+        prob = (cnt + eps) / (cnt + eps).sum()
+        H = -(prob * torch.log(prob)).sum()
+        ctx.save_for_backward(zq, zi, prob)
+        ctx.K = K
+        return H
+    
+    @staticmethod
+    def backward(ctx, grad_output):
+        zq, zi, prob= ctx.saved_tensors
+        grad_array = -grad_output * (torch.log(prob) + 1) / zi.numel() / ctx.K
+        reord_grad = grad_array[zi.flatten()].reshape(zi.shape)
+        grad_input = reord_grad.unsqueeze(-1) * zq
+        return grad_input, None, None, None, None
+
+
+def codebook_entropy(zq, basis, K, eps=1e-4):
+    return DifferentiableEntropyFunction.apply(zq, basis, K, eps)
+
+
+class BinarySphericalQuantizer(nn.Module):
+    def __init__(self, embed_dim, beta, gamma0, gamma, zeta,
+                 input_format='bchw', 
+                 soft_entropy=True, group_size=9,
+                 persample_entropy_compute='group',
+                 cb_entropy_compute='group',
+                 l2_norm=False,
+                 inv_temperature=1):
+        super().__init__()
+        self.embed_dim = embed_dim
+        self.beta = beta    # loss weight for commit loss
+        self.gamma0 = gamma0  # loss weight for entropy penalty
+        self.gamma = gamma  # loss weight for entropy penalty
+        self.zeta = zeta    # loss weight for entire entropy penalty
+        self.input_format = input_format
+        assert self.embed_dim % group_size == 0, "embed_dim must be divisible by group_size"
+        self.num_groups = self.embed_dim // group_size
+        self.group_size = group_size
+        assert persample_entropy_compute in ['group', 'analytical'], "persample_entropy_compute must be either 'group' or 'analytical'"
+        assert cb_entropy_compute in ['group', 'nce'], "cb_entropy_compute must be either 'group' or 'nce'"
+        self.persample_entropy_compute = persample_entropy_compute
+        self.cb_entropy_compute = cb_entropy_compute
+        self.l2_norm = l2_norm
+        self.inv_temperature = inv_temperature
+
+        self.register_buffer('basis', 2 ** torch.arange(embed_dim - 1, -1, -1))
+        self.register_buffer('group_basis', 2 ** torch.arange(group_size - 1, -1, -1))
+
+        self.num_dimensions = 2 ** embed_dim
+        self.bits_per_index = embed_dim
+
+        # we only need to keep the codebook portion up to the group size 
+        # because we approximate the H loss with this subcode
+        group_codes = torch.arange(2 ** self.group_size)
+        group_codebook = self.indexes_to_codes(group_codes).float()[:, -group_size:]
+        self.register_buffer('group_codebook', group_codebook, persistent=False)
+
+        self.soft_entropy = soft_entropy  # soft_entropy: Sec 3.2 of https://arxiv.org/pdf/1911.05894.pdf
+
+    def quantize(self, z):
+        assert z.shape[-1] == self.embed_dim, f"Expected {self.embed_dim} dimensions, got {z.shape[-1]}"
+
+        zhat = torch.where(z > 0, 
+                           torch.tensor(1, dtype=z.dtype, device=z.device), 
+                           torch.tensor(-1, dtype=z.dtype, device=z.device))
+        return z + (zhat - z).detach()
+
+    def forward(self, z):
+        if self.input_format == 'bchw':
+            z = rearrange(z, 'b c h w -> b h w c')
+        zq = self.quantize(z)
+
+        indices = self.codes_to_indexes(zq.detach())
+        group_indices = self.codes_to_group_indexes(zq.detach())
+        if not self.training:
+            used_codes = torch.unique(indices, return_counts=False)
+        else:
+            used_codes = None
+
+        q_scale = 1. / (self.embed_dim ** 0.5) if self.l2_norm else 1.
+
+        if self.soft_entropy:
+            persample_entropy, cb_entropy, avg_prob = self.soft_entropy_loss(z)
+            entropy_penalty = self.gamma0 * persample_entropy - self.gamma * cb_entropy
+        else:    
+            zb_by_sample= ((zq + 1)/2).reshape(z.shape[0], -1, z.shape[-1]).to(torch.float32)    
+            persample_entropy = self.get_hard_per_sample_entropy(zb_by_sample)  
+            cb_entropy = codebook_entropy(zq, self.basis, self.embed_dim)
+            entropy_penalty = self.gamma0 * persample_entropy - self.gamma * cb_entropy
+
+        zq = zq * q_scale
+
+        # commit loss
+        commit_loss = self.beta * torch.mean(((zq.detach() - z) ** 2).sum(dim=-1))
+
+        if self.input_format == 'bchw':
+            zq = rearrange(zq, 'b h w c -> b c h w')
+
+        return (
+            zq,
+            commit_loss + self.zeta * entropy_penalty / self.inv_temperature,
+            {"H": cb_entropy, "used_codes": used_codes, "indices": indices, "group_indices": group_indices,
+             "avg_prob": avg_prob}
+        )
+    
+    def soft_entropy_loss(self, z):
+        # if we divide the code in subgroups of size group_size, the codebook will be of size 2 ** group_size
+        # the sub-code is the last group_size bits of the full code
+        group_code_book = self.group_codebook / (self.embed_dim ** 0.5 if self.l2_norm else 1)
+        divided_z = rearrange(z, '... (g c) -> ... g c', c=self.group_size)
+
+        # we calculate the distance between the divided_z and the codebook for each subgroup
+        distance =  - 2 * torch.einsum('... g c, d c ->... g d', divided_z, group_code_book)  
+        prob = (-distance * self.inv_temperature).softmax(dim = -1)
+        if self.persample_entropy_compute == 'analytical':
+            if self.l2_norm:
+                p = torch.sigmoid(-4 * z / (self.embed_dim ** 0.5) * self.inv_temperature)
+            else:
+                p = torch.sigmoid(-4 * z * self.inv_temperature)
+            prob = torch.stack([p, 1-p], dim=-1)
+            per_sample_entropy = self.get_entropy(prob, dim=-1, normalize=False).sum(dim=-1).mean()
+        else:
+            per_sample_entropy = self.get_entropy(prob, dim=-1, normalize=False).sum(dim=-1).mean()
+
+        # macro average of the probability of each subgroup
+        avg_prob = reduce(prob, '... g d ->g d', 'mean')
+        codebook_entropy = self.get_entropy(avg_prob, dim=-1, normalize=False)
+
+        # the approximation of the entropy is the sum of the entropy of each subgroup
+        return per_sample_entropy, codebook_entropy.sum(), avg_prob
+        
+    def get_hard_per_sample_entropy(self, zb_by_sample):
+        probs_per_dim = zb_by_sample.sum(1) / zb_by_sample.shape[1]
+        persample_entropy = - probs_per_dim * torch.log(probs_per_dim + 1e-8) - (1 - probs_per_dim) * torch.log(1 - probs_per_dim + 1e-8)
+        persample_entropy = persample_entropy.sum(-1)
+        return persample_entropy.mean()
+
+    def codes_to_indexes(self, zhat):
+        """Converts a `code` to an index in the codebook.
+        Args:
+            zhat: A tensor of shape (B, ..., C) containing the codes. must be in {-1, 1}
+        """
+        assert zhat.shape[-1] == self.embed_dim, f"Expected {self.embed_dim} dimensions, got {zhat.shape[-1]}"
+        return ((zhat + 1) / 2 * self.basis).sum(axis=-1).to(torch.int64)
+
+    def codes_to_group_indexes(self, zhat):
+        """Converts a `code` to a list of indexes (in groups) in the codebook.
+        Args:
+            zhat: A tensor of shape (B, ..., C) containing the codes. must be in {-1, 1}
+        """
+        zhat_in_group = rearrange(zhat, 'b ... (g c) -> b ... g c', c=self.group_size)
+        return ((zhat_in_group + 1) / 2 * self.group_basis).sum(axis=-1).to(torch.int64)
+
+    def indexes_to_codes(self, indices):
+        """Inverse of `indexes_to_codes`."""
+        indices = indices.unsqueeze(-1)
+        codes_non_centered = torch.remainder(
+            torch.floor_divide(indices, self.basis), 2
+        )
+        return codes_non_centered * 2 - 1
+
+    def group_indexes_to_codes(self, group_indices):
+        """Inverse of `group_indexes_to_codes`."""
+        group_indices = group_indices.unsqueeze(-1)
+        codes_non_centered = torch.remainder(
+            torch.floor_divide(group_indices, self.group_basis), 2
+        )
+        codes_non_centered = rearrange(codes_non_centered, 'b ... g c -> b ... (g c)')
+        return codes_non_centered * 2 - 1
+
+    def get_entropy(self, count, dim=-1, eps=1e-4, normalize=True):
+        if normalize:
+            probs = (count + eps) / (count + eps).sum(dim=dim, keepdim =True)
+        else:
+            probs = count
+        H = -(probs * torch.log(probs + 1e-8)).sum(dim=dim)
+        return H
+
+    def get_group_codebook_entry(self, group_indices):
+        z_q = self.group_indexes_to_codes(group_indices)
+        q_scale = 1. / (self.embed_dim ** 0.5) if self.l2_norm else 1.
+        z_q = z_q * q_scale
+        if self.input_format == 'bchw':
+            h, w = int(z_q.shape[1] ** 0.5)
+            assert h * w == z_q.shape[1], 'Invalid sequence length'
+            z_q = rearrange(z_q, 'b (h w) c -> b c h w', h=h)
+        return z_q
+
+    def get_codebook_entry(self, indices):
+        z_q = self.indexes_to_codes(indices)
+        q_scale = 1. / (self.embed_dim ** 0.5) if self.l2_norm else 1.
+        z_q = z_q * q_scale
+        if self.input_format == 'bchw':
+            h, w = int(z_q.shape[1] ** 0.5)
+            assert h * w == z_q.shape[1], 'Invalid sequence length'
+            z_q = rearrange(z_q, 'b (h w) c -> b c h w', h=h)
+        return z_q
+
+
+if __name__ == "__main__":
+    K = 8
+    # zq = torch.randint(0, 2, (4, 32, K), dtype=torch.bfloat16, device='cuda') * 2 - 1
+    zq = torch.zeros((4, 32, K), dtype=torch.bfloat16, device='cuda') * 2 - 1
+    basis = (2 ** torch.arange(K - 1, -1, -1)).to(torch.bfloat16).cuda()
+    zq.requires_grad = True
+    h = codebook_entropy(zq, basis, K)
+    h.backward()
+    print(zq.grad, zq)

src/vqvaes/bsqvit/quantizer/vq.py

@@ -0,0 +1,152 @@
+from einops import rearrange
+import numpy as np
+import torch
+import torch.distributed as dist
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+class VectorQuantizer(nn.Module):
+    def __init__(self, n_embed, embed_dim, l2_norm, beta, input_format='bchw'):
+        super().__init__()
+
+        self.n_embed = n_embed
+        self.embed_dim = embed_dim
+        self.l2_norm = l2_norm
+        self.beta = beta
+        assert input_format in ['bchw', 'blc']
+        self.input_format = input_format
+
+        self.embedding = nn.Embedding(n_embed, embed_dim)
+        self.embedding.weight.data.uniform_(-1 / n_embed, 1 / n_embed)
+        self.bits_per_index = int(np.ceil(np.log2(n_embed)))
+
+    def forward(self, z):
+        batch = z.shape[0]
+        if self.input_format == 'bchw':
+            z = rearrange(z, 'b c h w -> b h w c')
+
+        if self.l2_norm:
+            z = F.normalize(z, dim=-1)
+            z_flatten = z.reshape(-1, self.embed_dim)
+            embedding_weight = F.normalize(self.embedding.weight, dim=-1)
+            d = -z_flatten @ embedding_weight.t()
+        else:
+            z_flatten = z.reshape(-1, self.embed_dim)
+            d = torch.sum(z_flatten ** 2, dim=1, keepdim=True) + torch.sum(self.embedding.weight ** 2, dim=1) - 2 * z_flatten @ self.embedding.weight.t()
+
+        min_encoding_indices = torch.argmin(d.detach(), dim=1)
+        if not self.training:
+            used_codes = torch.unique(min_encoding_indices, return_counts=False)
+        else:
+            used_codes = None
+        cb_usage = F.one_hot(min_encoding_indices, self.n_embed).sum(0)
+        cb_entropy = self.get_entropy(cb_usage)
+
+        z_q = self.embedding(min_encoding_indices).view(z.shape)
+        if self.l2_norm:
+            z_q = F.normalize(z_q, dim=-1)
+
+        # fix the issue with loss scaling
+        # loss weight should not associate with the dimensionality of words
+        # loss = self.beta * torch.mean((z_q.detach() - z) ** 2) + torch.mean((z_q - z.detach()) ** 2)
+        loss = self.beta * torch.mean(((z_q.detach() - z) ** 2).sum(dim=-1)) + torch.mean(((z_q - z.detach()) ** 2).sum(dim=-1))
+
+        z_q = z + (z_q - z).detach()
+        if self.input_format == 'bchw':
+            z_q = rearrange(z_q, 'b h w c -> b c h w')
+        return z_q, loss, {"H":cb_entropy, "used_codes": used_codes, 'indices': min_encoding_indices.view(batch, -1)}
+
+    def get_entropy(self, count, eps=1e-4):
+        probs = (count + eps) / (count + eps).sum()
+        H = -(probs * torch.log(probs)).sum()
+        return H
+
+
+    def get_codebook_entry(self, indices):
+        z_q = self.embedding(indices)
+        if self.l2_norm:
+            z_q = F.normalize(z_q, dim=-1)
+
+        if self.input_format == 'bchw':
+            h = w = int(z_q.shape[1] ** 0.5)
+            assert h * w == z_q.shape[1], 'Invalid sequence length'
+            z_q = rearrange(z_q, 'b (h w) c -> b c h w', h=h)
+        return z_q
+
+
+class EMAVectorQuantizer(nn.Module):
+    def __init__(self, n_embed, embed_dim, l2_norm, beta, decay=0.99, eps=1e-5, random_restart=True, restart_threshold=1.0, input_format='bchw'):
+        super().__init__()
+
+        self.n_embed = n_embed
+        self.embed_dim = embed_dim
+        self.l2_norm = l2_norm
+        self.beta = beta
+        self.decay = decay
+        self.eps = eps
+        self.random_restart = random_restart
+        self.restart_threshold = restart_threshold
+        self.input_format = input_format
+
+        self.embedding = nn.Embedding(n_embed, embed_dim)
+        self.embedding.weight.data.uniform_(-1 / n_embed, 1 / n_embed) # TODO (yzhao): test other initialization methods 
+        self.register_buffer("ema_cluster_size", torch.zeros(self.n_embed))
+        self.embedding_avg = nn.Parameter(torch.Tensor(self.n_embed, self.embed_dim))
+        self.embedding_avg.data.copy_(self.embedding.weight.data)
+
+    def _tile(self, z):
+        n_z, embedding_dim = z.shape
+        if n_z < self.n_embed:
+            n_repeats = (self.n_embed + n_z - 1) // n_z
+            std = 0.01 / np.sqrt(embedding_dim)
+            z = z.repeat(n_repeats, 1)
+            z = z + torch.randn_like(z) * std
+        return z
+
+    def forward(self, z):
+        if self.input_format == 'bchw':
+            z = rearrange(z, 'b c h w -> b h w c')
+        z_flatten = z.reshape(-1, self.embed_dim)
+
+        d = torch.sum(z_flatten ** 2, dim=1, keepdim=True) + torch.sum(self.embedding.weight ** 2, dim=1) - 2 * z_flatten @ self.embedding.weight.t()
+
+        encoding_indices = torch.argmin(d, dim=1).unsqueeze(1)
+        encodings = torch.zeros(encoding_indices.size(0), self.n_embed, device=z.device)
+        encodings.scatter_(1, encoding_indices, 1)
+
+        z_q = self.embedding(encoding_indices).view(z.shape)
+        if self.l2_norm:
+            z = F.normalize(z, dim=-1)
+            z_q = F.normalize(z_q, dim=-1)
+
+        if self.training:
+            # EMA update cluster size
+            encodings_sum = encodings.sum(0)
+            if dist.is_initialized(): dist.all_reduce(encodings_sum)
+            self.ema_cluster_size.data.mul_(self.decay).add_(encodings_sum, alpha=1-self.decay)
+
+            # EMA update of the embedding vectors
+            dw = encodings.t() @ z_flatten
+            if dist.is_initialized(): dist.all_reduce(dw)
+            self.embedding_avg.data.mul_(self.decay).add_(dw, alpha=1-self.decay)
+ 
+            # Laplace smoothing of the cluster size
+            n = torch.sum(self.ema_cluster_size)
+            weights = (self.ema_cluster_size + self.eps) / (n + self.n_embed * self.eps) * n
+            self.embedding.weight.data = self.embedding_avg.data / weights.unsqueeze(1)
+
+            if self.random_restart:
+                zz = self._tile(z_flatten)
+                _k_rand = zz[torch.randperm(zz.size(0))][:self.n_embed]
+                if dist.is_initialized(): dist.broadcast(_k_rand, 0)
+                usage = (self.ema_cluster_size.view(-1, 1) > self.restart_threshold).float()
+                self.embedding.weight.data.mul_(usage).add_(_k_rand * (1 - usage))
+
+        loss = self.beta * torch.mean((z_q.detach() - z) ** 2)
+
+        z_q = z + (z_q - z).detach()
+        if self.input_format == 'bchw':
+            z_q = rearrange(z_q, 'b h w c -> b c h w')
+        # TODO (yzhao): monitor utility of the dictionary
+        return z_q, loss, {}

src/vqvaes/bsqvit/stylegan_utils/custom_ops.py

@@ -0,0 +1,126 @@
+# Copyright (c) 2021, NVIDIA CORPORATION.  All rights reserved.
+#
+# NVIDIA CORPORATION and its licensors retain all intellectual property
+# and proprietary rights in and to this software, related documentation
+# and any modifications thereto.  Any use, reproduction, disclosure or
+# distribution of this software and related documentation without an express
+# license agreement from NVIDIA CORPORATION is strictly prohibited.
+
+import os
+import glob
+import torch
+import torch.utils.cpp_extension
+import importlib
+import hashlib
+import shutil
+from pathlib import Path
+
+from torch.utils.file_baton import FileBaton
+
+#----------------------------------------------------------------------------
+# Global options.
+
+verbosity = 'brief' # Verbosity level: 'none', 'brief', 'full'
+
+#----------------------------------------------------------------------------
+# Internal helper funcs.
+
+def _find_compiler_bindir():
+    patterns = [
+        'C:/Program Files (x86)/Microsoft Visual Studio/*/Professional/VC/Tools/MSVC/*/bin/Hostx64/x64',
+        'C:/Program Files (x86)/Microsoft Visual Studio/*/BuildTools/VC/Tools/MSVC/*/bin/Hostx64/x64',
+        'C:/Program Files (x86)/Microsoft Visual Studio/*/Community/VC/Tools/MSVC/*/bin/Hostx64/x64',
+        'C:/Program Files (x86)/Microsoft Visual Studio */vc/bin',
+    ]
+    for pattern in patterns:
+        matches = sorted(glob.glob(pattern))
+        if len(matches):
+            return matches[-1]
+    return None
+
+#----------------------------------------------------------------------------
+# Main entry point for compiling and loading C++/CUDA plugins.
+
+_cached_plugins = dict()
+
+def get_plugin(module_name, sources, **build_kwargs):
+    assert verbosity in ['none', 'brief', 'full']
+
+    # Already cached?
+    if module_name in _cached_plugins:
+        return _cached_plugins[module_name]
+
+    # Print status.
+    if verbosity == 'full':
+        print(f'Setting up PyTorch plugin "{module_name}"...')
+    elif verbosity == 'brief':
+        print(f'Setting up PyTorch plugin "{module_name}"... ', end='', flush=True)
+
+    try: # pylint: disable=too-many-nested-blocks
+        # Make sure we can find the necessary compiler binaries.
+        if os.name == 'nt' and os.system("where cl.exe >nul 2>nul") != 0:
+            compiler_bindir = _find_compiler_bindir()
+            if compiler_bindir is None:
+                raise RuntimeError(f'Could not find MSVC/GCC/CLANG installation on this computer. Check _find_compiler_bindir() in "{__file__}".')
+            os.environ['PATH'] += ';' + compiler_bindir
+
+        # Compile and load.
+        verbose_build = (verbosity == 'full')
+
+        # Incremental build md5sum trickery.  Copies all the input source files
+        # into a cached build directory under a combined md5 digest of the input
+        # source files.  Copying is done only if the combined digest has changed.
+        # This keeps input file timestamps and filenames the same as in previous
+        # extension builds, allowing for fast incremental rebuilds.
+        #
+        # This optimization is done only in case all the source files reside in
+        # a single directory (just for simplicity) and if the TORCH_EXTENSIONS_DIR
+        # environment variable is set (we take this as a signal that the user
+        # actually cares about this.)
+        source_dirs_set = set(os.path.dirname(source) for source in sources)
+        if len(source_dirs_set) == 1 and ('TORCH_EXTENSIONS_DIR' in os.environ):
+            all_source_files = sorted(list(x for x in Path(list(source_dirs_set)[0]).iterdir() if x.is_file()))
+
+            # Compute a combined hash digest for all source files in the same
+            # custom op directory (usually .cu, .cpp, .py and .h files).
+            hash_md5 = hashlib.md5()
+            for src in all_source_files:
+                with open(src, 'rb') as f:
+                    hash_md5.update(f.read())
+            build_dir = torch.utils.cpp_extension._get_build_directory(module_name, verbose=verbose_build) # pylint: disable=protected-access
+            digest_build_dir = os.path.join(build_dir, hash_md5.hexdigest())
+
+            if not os.path.isdir(digest_build_dir):
+                os.makedirs(digest_build_dir, exist_ok=True)
+                baton = FileBaton(os.path.join(digest_build_dir, 'lock'))
+                if baton.try_acquire():
+                    try:
+                        for src in all_source_files:
+                            shutil.copyfile(src, os.path.join(digest_build_dir, os.path.basename(src)))
+                    finally:
+                        baton.release()
+                else:
+                    # Someone else is copying source files under the digest dir,
+                    # wait until done and continue.
+                    baton.wait()
+            digest_sources = [os.path.join(digest_build_dir, os.path.basename(x)) for x in sources]
+            torch.utils.cpp_extension.load(name=module_name, build_directory=build_dir,
+                verbose=verbose_build, sources=digest_sources, **build_kwargs)
+        else:
+            torch.utils.cpp_extension.load(name=module_name, verbose=verbose_build, sources=sources, **build_kwargs)
+        module = importlib.import_module(module_name)
+
+    except:
+        if verbosity == 'brief':
+            print('Failed!')
+        raise
+
+    # Print status and add to cache.
+    if verbosity == 'full':
+        print(f'Done setting up PyTorch plugin "{module_name}".')
+    elif verbosity == 'brief':
+        print('Done.')
+    _cached_plugins[module_name] = module
+    return module
+
+#----------------------------------------------------------------------------

src/vqvaes/bsqvit/stylegan_utils/misc.py

@@ -0,0 +1,40 @@
+import torch
+import warnings
+
+
+#----------------------------------------------------------------------------
+# Symbolic assert.
+
+try:
+    symbolic_assert = torch._assert # 1.8.0a0 # pylint: disable=protected-access
+except AttributeError:
+    symbolic_assert = torch.Assert # 1.7.0
+
+#----------------------------------------------------------------------------
+# Context manager to suppress known warnings in torch.jit.trace().
+
+class suppress_tracer_warnings(warnings.catch_warnings):
+    def __enter__(self):
+        super().__enter__()
+        warnings.simplefilter('ignore', category=torch.jit.TracerWarning)
+        return self
+
+#----------------------------------------------------------------------------
+# Assert that the shape of a tensor matches the given list of integers.
+# None indicates that the size of a dimension is allowed to vary.
+# Performs symbolic assertion when used in torch.jit.trace().
+
+def assert_shape(tensor, ref_shape):
+    if tensor.ndim != len(ref_shape):
+        raise AssertionError(f'Wrong number of dimensions: got {tensor.ndim}, expected {len(ref_shape)}')
+    for idx, (size, ref_size) in enumerate(zip(tensor.shape, ref_shape)):
+        if ref_size is None:
+            pass
+        elif isinstance(ref_size, torch.Tensor):
+            with suppress_tracer_warnings(): # as_tensor results are registered as constants
+                symbolic_assert(torch.equal(torch.as_tensor(size), ref_size), f'Wrong size for dimension {idx}')
+        elif isinstance(size, torch.Tensor):
+            with suppress_tracer_warnings(): # as_tensor results are registered as constants
+                symbolic_assert(torch.equal(size, torch.as_tensor(ref_size)), f'Wrong size for dimension {idx}: expected {ref_size}')
+        elif size != ref_size:
+            raise AssertionError(f'Wrong size for dimension {idx}: got {size}, expected {ref_size}')

src/vqvaes/bsqvit/stylegan_utils/ops/bias_act.cpp

@@ -0,0 +1,99 @@
+// Copyright (c) 2021, NVIDIA CORPORATION.  All rights reserved.
+//
+// NVIDIA CORPORATION and its licensors retain all intellectual property
+// and proprietary rights in and to this software, related documentation
+// and any modifications thereto.  Any use, reproduction, disclosure or
+// distribution of this software and related documentation without an express
+// license agreement from NVIDIA CORPORATION is strictly prohibited.
+
+#include <torch/extension.h>
+#include <ATen/cuda/CUDAContext.h>
+#include <c10/cuda/CUDAGuard.h>
+#include "bias_act.h"
+
+//------------------------------------------------------------------------
+
+static bool has_same_layout(torch::Tensor x, torch::Tensor y)
+{
+    if (x.dim() != y.dim())
+        return false;
+    for (int64_t i = 0; i < x.dim(); i++)
+    {
+        if (x.size(i) != y.size(i))
+            return false;
+        if (x.size(i) >= 2 && x.stride(i) != y.stride(i))
+            return false;
+    }
+    return true;
+}
+
+//------------------------------------------------------------------------
+
+static torch::Tensor bias_act(torch::Tensor x, torch::Tensor b, torch::Tensor xref, torch::Tensor yref, torch::Tensor dy, int grad, int dim, int act, float alpha, float gain, float clamp)
+{
+    // Validate arguments.
+    TORCH_CHECK(x.is_cuda(), "x must reside on CUDA device");
+    TORCH_CHECK(b.numel() == 0 || (b.dtype() == x.dtype() && b.device() == x.device()), "b must have the same dtype and device as x");
+    TORCH_CHECK(xref.numel() == 0 || (xref.sizes() == x.sizes() && xref.dtype() == x.dtype() && xref.device() == x.device()), "xref must have the same shape, dtype, and device as x");
+    TORCH_CHECK(yref.numel() == 0 || (yref.sizes() == x.sizes() && yref.dtype() == x.dtype() && yref.device() == x.device()), "yref must have the same shape, dtype, and device as x");
+    TORCH_CHECK(dy.numel() == 0 || (dy.sizes() == x.sizes() && dy.dtype() == x.dtype() && dy.device() == x.device()), "dy must have the same dtype and device as x");
+    TORCH_CHECK(x.numel() <= INT_MAX, "x is too large");
+    TORCH_CHECK(b.dim() == 1, "b must have rank 1");
+    TORCH_CHECK(b.numel() == 0 || (dim >= 0 && dim < x.dim()), "dim is out of bounds");
+    TORCH_CHECK(b.numel() == 0 || b.numel() == x.size(dim), "b has wrong number of elements");
+    TORCH_CHECK(grad >= 0, "grad must be non-negative");
+
+    // Validate layout.
+    TORCH_CHECK(x.is_non_overlapping_and_dense(), "x must be non-overlapping and dense");
+    TORCH_CHECK(b.is_contiguous(), "b must be contiguous");
+    TORCH_CHECK(xref.numel() == 0 || has_same_layout(xref, x), "xref must have the same layout as x");
+    TORCH_CHECK(yref.numel() == 0 || has_same_layout(yref, x), "yref must have the same layout as x");
+    TORCH_CHECK(dy.numel() == 0 || has_same_layout(dy, x), "dy must have the same layout as x");
+
+    // Create output tensor.
+    const at::cuda::OptionalCUDAGuard device_guard(device_of(x));
+    torch::Tensor y = torch::empty_like(x);
+    TORCH_CHECK(has_same_layout(y, x), "y must have the same layout as x");
+
+    // Initialize CUDA kernel parameters.
+    bias_act_kernel_params p;
+    p.x     = x.data_ptr();
+    p.b     = (b.numel()) ? b.data_ptr() : NULL;
+    p.xref  = (xref.numel()) ? xref.data_ptr() : NULL;
+    p.yref  = (yref.numel()) ? yref.data_ptr() : NULL;
+    p.dy    = (dy.numel()) ? dy.data_ptr() : NULL;
+    p.y     = y.data_ptr();
+    p.grad  = grad;
+    p.act   = act;
+    p.alpha = alpha;
+    p.gain  = gain;
+    p.clamp = clamp;
+    p.sizeX = (int)x.numel();
+    p.sizeB = (int)b.numel();
+    p.stepB = (b.numel()) ? (int)x.stride(dim) : 1;
+
+    // Choose CUDA kernel.
+    void* kernel;
+    AT_DISPATCH_REDUCED_FLOATING_TYPES(x.scalar_type(), "upfirdn2d_cuda", [&]
+    {
+        kernel = choose_bias_act_kernel<scalar_t>(p);
+    });
+    TORCH_CHECK(kernel, "no CUDA kernel found for the specified activation func");
+
+    // Launch CUDA kernel.
+    p.loopX = 4;
+    int blockSize = 4 * 32;
+    int gridSize = (p.sizeX - 1) / (p.loopX * blockSize) + 1;
+    void* args[] = {&p};
+    AT_CUDA_CHECK(cudaLaunchKernel(kernel, gridSize, blockSize, args, 0, at::cuda::getCurrentCUDAStream()));
+    return y;
+}
+
+//------------------------------------------------------------------------
+
+PYBIND11_MODULE(TORCH_EXTENSION_NAME, m)
+{
+    m.def("bias_act", &bias_act);
+}
+
+//------------------------------------------------------------------------

src/vqvaes/bsqvit/stylegan_utils/ops/bias_act.cu

@@ -0,0 +1,176 @@
+// Copyright (c) 2021, NVIDIA CORPORATION.  All rights reserved.
+//
+// NVIDIA CORPORATION and its licensors retain all intellectual property
+// and proprietary rights in and to this software, related documentation
+// and any modifications thereto.  Any use, reproduction, disclosure or
+// distribution of this software and related documentation without an express
+// license agreement from NVIDIA CORPORATION is strictly prohibited.
+
+#include <ATen/ATen.h>
+#include <c10/util/Half.h>
+#include "bias_act.h"
+
+//------------------------------------------------------------------------
+// Helpers.
+
+template <class T> struct InternalType;
+template <> struct InternalType<double>         { typedef double scalar_t; };
+template <> struct InternalType<float>          { typedef float  scalar_t; };
+template <> struct InternalType<c10::Half>      { typedef float  scalar_t; };
+template <> struct InternalType<at::BFloat16>   { typedef float  scalar_t; };
+
+//------------------------------------------------------------------------
+// CUDA kernel.
+
+template <class T, int A>
+__global__ void bias_act_kernel(bias_act_kernel_params p)
+{
+    typedef typename InternalType<T>::scalar_t scalar_t;
+    int G                 = p.grad;
+    scalar_t alpha        = (scalar_t)p.alpha;
+    scalar_t gain         = (scalar_t)p.gain;
+    scalar_t clamp        = (scalar_t)p.clamp;
+    scalar_t one          = (scalar_t)1;
+    scalar_t two          = (scalar_t)2;
+    scalar_t expRange     = (scalar_t)80;
+    scalar_t halfExpRange = (scalar_t)40;
+    scalar_t seluScale    = (scalar_t)1.0507009873554804934193349852946;
+    scalar_t seluAlpha    = (scalar_t)1.6732632423543772848170429916717;
+
+    // Loop over elements.
+    int xi = blockIdx.x * p.loopX * blockDim.x + threadIdx.x;
+    for (int loopIdx = 0; loopIdx < p.loopX && xi < p.sizeX; loopIdx++, xi += blockDim.x)
+    {
+        // Load.
+        scalar_t x = (scalar_t)((const T*)p.x)[xi];
+        scalar_t b = (p.b) ? (scalar_t)((const T*)p.b)[(xi / p.stepB) % p.sizeB] : 0;
+        scalar_t xref = (p.xref) ? (scalar_t)((const T*)p.xref)[xi] : 0;
+        scalar_t yref = (p.yref) ? (scalar_t)((const T*)p.yref)[xi] : 0;
+        scalar_t dy = (p.dy) ? (scalar_t)((const T*)p.dy)[xi] : one;
+        scalar_t yy = (gain != 0) ? yref / gain : 0;
+        scalar_t y = 0;
+
+        // Apply bias.
+        ((G == 0) ? x : xref) += b;
+
+        // linear
+        if (A == 1)
+        {
+            if (G == 0) y = x;
+            if (G == 1) y = x;
+        }
+
+        // relu
+        if (A == 2)
+        {
+            if (G == 0) y = (x > 0) ? x : 0;
+            if (G == 1) y = (yy > 0) ? x : 0;
+        }
+
+        // lrelu
+        if (A == 3)
+        {
+            if (G == 0) y = (x > 0) ? x : x * alpha;
+            if (G == 1) y = (yy > 0) ? x : x * alpha;
+        }
+
+        // tanh
+        if (A == 4)
+        {
+            if (G == 0) { scalar_t c = exp(x); scalar_t d = one / c; y = (x < -expRange) ? -one : (x > expRange) ? one : (c - d) / (c + d); }
+            if (G == 1) y = x * (one - yy * yy);
+            if (G == 2) y = x * (one - yy * yy) * (-two * yy);
+        }
+
+        // sigmoid
+        if (A == 5)
+        {
+            if (G == 0) y = (x < -expRange) ? 0 : one / (exp(-x) + one);
+            if (G == 1) y = x * yy * (one - yy);
+            if (G == 2) y = x * yy * (one - yy) * (one - two * yy);
+        }
+
+        // elu
+        if (A == 6)
+        {
+            if (G == 0) y = (x >= 0) ? x : exp(x) - one;
+            if (G == 1) y = (yy >= 0) ? x : x * (yy + one);
+            if (G == 2) y = (yy >= 0) ? 0 : x * (yy + one);
+        }
+
+        // selu
+        if (A == 7)
+        {
+            if (G == 0) y = (x >= 0) ? seluScale * x : (seluScale * seluAlpha) * (exp(x) - one);
+            if (G == 1) y = (yy >= 0) ? x * seluScale : x * (yy + seluScale * seluAlpha);
+            if (G == 2) y = (yy >= 0) ? 0 : x * (yy + seluScale * seluAlpha);
+        }
+
+        // softplus
+        if (A == 8)
+        {
+            if (G == 0) y = (x > expRange) ? x : log(exp(x) + one);
+            if (G == 1) y = x * (one - exp(-yy));
+            if (G == 2) { scalar_t c = exp(-yy); y = x * c * (one - c); }
+        }
+
+        // swish
+        if (A == 9)
+        {
+            if (G == 0)
+                y = (x < -expRange) ? 0 : x / (exp(-x) + one);
+            else
+            {
+                scalar_t c = exp(xref);
+                scalar_t d = c + one;
+                if (G == 1)
+                    y = (xref > halfExpRange) ? x : x * c * (xref + d) / (d * d);
+                else
+                    y = (xref > halfExpRange) ? 0 : x * c * (xref * (two - d) + two * d) / (d * d * d);
+                yref = (xref < -expRange) ? 0 : xref / (exp(-xref) + one) * gain;
+            }
+        }
+
+        // Apply gain.
+        y *= gain * dy;
+
+        // Clamp.
+        if (clamp >= 0)
+        {
+            if (G == 0)
+                y = (y > -clamp & y < clamp) ? y : (y >= 0) ? clamp : -clamp;
+            else
+                y = (yref > -clamp & yref < clamp) ? y : 0;
+        }
+
+        // Store.
+        ((T*)p.y)[xi] = (T)y;
+    }
+}
+
+//------------------------------------------------------------------------
+// CUDA kernel selection.
+
+template <class T> void* choose_bias_act_kernel(const bias_act_kernel_params& p)
+{
+    if (p.act == 1) return (void*)bias_act_kernel<T, 1>;
+    if (p.act == 2) return (void*)bias_act_kernel<T, 2>;
+    if (p.act == 3) return (void*)bias_act_kernel<T, 3>;
+    if (p.act == 4) return (void*)bias_act_kernel<T, 4>;
+    if (p.act == 5) return (void*)bias_act_kernel<T, 5>;
+    if (p.act == 6) return (void*)bias_act_kernel<T, 6>;
+    if (p.act == 7) return (void*)bias_act_kernel<T, 7>;
+    if (p.act == 8) return (void*)bias_act_kernel<T, 8>;
+    if (p.act == 9) return (void*)bias_act_kernel<T, 9>;
+    return NULL;
+}
+
+//------------------------------------------------------------------------
+// Template specializations.
+
+template void* choose_bias_act_kernel<double>          (const bias_act_kernel_params& p);
+template void* choose_bias_act_kernel<float>           (const bias_act_kernel_params& p);
+template void* choose_bias_act_kernel<c10::Half>       (const bias_act_kernel_params& p);
+template void* choose_bias_act_kernel<at::BFloat16>    (const bias_act_kernel_params& p);
+
+//------------------------------------------------------------------------

src/vqvaes/bsqvit/stylegan_utils/ops/bias_act.h

@@ -0,0 +1,38 @@
+// Copyright (c) 2021, NVIDIA CORPORATION.  All rights reserved.
+//
+// NVIDIA CORPORATION and its licensors retain all intellectual property
+// and proprietary rights in and to this software, related documentation
+// and any modifications thereto.  Any use, reproduction, disclosure or
+// distribution of this software and related documentation without an express
+// license agreement from NVIDIA CORPORATION is strictly prohibited.
+
+//------------------------------------------------------------------------
+// CUDA kernel parameters.
+
+struct bias_act_kernel_params
+{
+    const void* x;      // [sizeX]
+    const void* b;      // [sizeB] or NULL
+    const void* xref;   // [sizeX] or NULL
+    const void* yref;   // [sizeX] or NULL
+    const void* dy;     // [sizeX] or NULL
+    void*       y;      // [sizeX]
+
+    int         grad;
+    int         act;
+    float       alpha;
+    float       gain;
+    float       clamp;
+
+    int         sizeX;
+    int         sizeB;
+    int         stepB;
+    int         loopX;
+};
+
+//------------------------------------------------------------------------
+// CUDA kernel selection.
+
+template <class T> void* choose_bias_act_kernel(const bias_act_kernel_params& p);
+
+//------------------------------------------------------------------------

src/vqvaes/bsqvit/stylegan_utils/ops/bias_act.py

@@ -0,0 +1,226 @@
+# Copyright (c) 2021, NVIDIA CORPORATION.  All rights reserved.
+#
+# NVIDIA CORPORATION and its licensors retain all intellectual property
+# and proprietary rights in and to this software, related documentation
+# and any modifications thereto.  Any use, reproduction, disclosure or
+# distribution of this software and related documentation without an express
+# license agreement from NVIDIA CORPORATION is strictly prohibited.
+
+"""Custom PyTorch ops for efficient bias and activation."""
+
+import os
+import warnings
+import numpy as np
+import torch
+import traceback
+from typing import Any
+
+from .. import custom_ops
+
+
+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]
+
+#----------------------------------------------------------------------------
+
+activation_funcs = {
+    'linear':   EasyDict(func=lambda x, **_:         x,                                          def_alpha=0,    def_gain=1,             cuda_idx=1, ref='',  has_2nd_grad=False),
+    'relu':     EasyDict(func=lambda x, **_:         torch.nn.functional.relu(x),                def_alpha=0,    def_gain=np.sqrt(2),    cuda_idx=2, ref='y', has_2nd_grad=False),
+    'lrelu':    EasyDict(func=lambda x, alpha, **_:  torch.nn.functional.leaky_relu(x, alpha),   def_alpha=0.2,  def_gain=np.sqrt(2),    cuda_idx=3, ref='y', has_2nd_grad=False),
+    'tanh':     EasyDict(func=lambda x, **_:         torch.tanh(x),                              def_alpha=0,    def_gain=1,             cuda_idx=4, ref='y', has_2nd_grad=True),
+    'sigmoid':  EasyDict(func=lambda x, **_:         torch.sigmoid(x),                           def_alpha=0,    def_gain=1,             cuda_idx=5, ref='y', has_2nd_grad=True),
+    'elu':      EasyDict(func=lambda x, **_:         torch.nn.functional.elu(x),                 def_alpha=0,    def_gain=1,             cuda_idx=6, ref='y', has_2nd_grad=True),
+    'selu':     EasyDict(func=lambda x, **_:         torch.nn.functional.selu(x),                def_alpha=0,    def_gain=1,             cuda_idx=7, ref='y', has_2nd_grad=True),
+    'softplus': EasyDict(func=lambda x, **_:         torch.nn.functional.softplus(x),            def_alpha=0,    def_gain=1,             cuda_idx=8, ref='y', has_2nd_grad=True),
+    'swish':    EasyDict(func=lambda x, **_:         torch.sigmoid(x) * x,                       def_alpha=0,    def_gain=np.sqrt(2),    cuda_idx=9, ref='x', has_2nd_grad=True),
+}
+
+#----------------------------------------------------------------------------
+
+_inited = False
+_plugin = None
+_null_tensor = torch.empty([0])
+
+def _init():
+    global _inited, _plugin
+    if not _inited:
+        _inited = True
+        sources = ['bias_act.cpp', 'bias_act.cu']
+        sources = [os.path.join(os.path.dirname(__file__), s) for s in sources]
+        try:
+            _plugin = custom_ops.get_plugin('bias_act_plugin', sources=sources, extra_cuda_cflags=['--use_fast_math'])
+        except:
+            warnings.warn('Failed to build CUDA kernels for bias_act. Falling back to slow reference implementation. Details:\n\n' + traceback.format_exc())
+    return _plugin is not None
+
+#----------------------------------------------------------------------------
+
+def bias_act(x, b=None, dim=1, act='linear', alpha=None, gain=None, clamp=None, impl='cuda'):
+    r"""Fused bias and activation function.
+
+    Adds bias `b` to activation tensor `x`, evaluates activation function `act`,
+    and scales the result by `gain`. Each of the steps is optional. In most cases,
+    the fused op is considerably more efficient than performing the same calculation
+    using standard PyTorch ops. It supports first and second order gradients,
+    but not third order gradients.
+
+    Args:
+        x:      Input activation tensor. Can be of any shape.
+        b:      Bias vector, or `None` to disable. Must be a 1D tensor of the same type
+                as `x`. The shape must be known, and it must match the dimension of `x`
+                corresponding to `dim`.
+        dim:    The dimension in `x` corresponding to the elements of `b`.
+                The value of `dim` is ignored if `b` is not specified.
+        act:    Name of the activation function to evaluate, or `"linear"` to disable.
+                Can be e.g. `"relu"`, `"lrelu"`, `"tanh"`, `"sigmoid"`, `"swish"`, etc.
+                See `activation_funcs` for a full list. `None` is not allowed.
+        alpha:  Shape parameter for the activation function, or `None` to use the default.
+        gain:   Scaling factor for the output tensor, or `None` to use default.
+                See `activation_funcs` for the default scaling of each activation function.
+                If unsure, consider specifying 1.
+        clamp:  Clamp the output values to `[-clamp, +clamp]`, or `None` to disable
+                the clamping (default).
+        impl:   Name of the implementation to use. Can be `"ref"` or `"cuda"` (default).
+
+    Returns:
+        Tensor of the same shape and datatype as `x`.
+    """
+    assert isinstance(x, torch.Tensor)
+    assert impl in ['ref', 'cuda']
+    if impl == 'cuda' and x.device.type == 'cuda' and _init():
+        return _bias_act_cuda(dim=dim, act=act, alpha=alpha, gain=gain, clamp=clamp).apply(x, b)
+    return _bias_act_ref(x=x, b=b, dim=dim, act=act, alpha=alpha, gain=gain, clamp=clamp)
+
+#----------------------------------------------------------------------------
+
+def _bias_act_ref(x, b=None, dim=1, act='linear', alpha=None, gain=None, clamp=None):
+    """Slow reference implementation of `bias_act()` using standard TensorFlow ops.
+    """
+    assert isinstance(x, torch.Tensor)
+    assert clamp is None or clamp >= 0
+    spec = activation_funcs[act]
+    alpha = float(alpha if alpha is not None else spec.def_alpha)
+    gain = float(gain if gain is not None else spec.def_gain)
+    clamp = float(clamp if clamp is not None else -1)
+
+    # Add bias.
+    if b is not None:
+        assert isinstance(b, torch.Tensor) and b.ndim == 1
+        assert 0 <= dim < x.ndim
+        assert b.shape[0] == x.shape[dim]
+        x = x + b.reshape([-1 if i == dim else 1 for i in range(x.ndim)])
+
+    # Evaluate activation function.
+    alpha = float(alpha)
+    x = spec.func(x, alpha=alpha)
+
+    # Scale by gain.
+    gain = float(gain)
+    if gain != 1:
+        x = x * gain
+
+    # Clamp.
+    if clamp >= 0:
+        x = x.clamp(-clamp, clamp) # pylint: disable=invalid-unary-operand-type
+    return x
+
+#----------------------------------------------------------------------------
+
+_bias_act_cuda_cache = dict()
+
+def _bias_act_cuda(dim=1, act='linear', alpha=None, gain=None, clamp=None):
+    """Fast CUDA implementation of `bias_act()` using custom ops.
+    """
+    # Parse arguments.
+    assert clamp is None or clamp >= 0
+    spec = activation_funcs[act]
+    alpha = float(alpha if alpha is not None else spec.def_alpha)
+    gain = float(gain if gain is not None else spec.def_gain)
+    clamp = float(clamp if clamp is not None else -1)
+
+    # Lookup from cache.
+    key = (dim, act, alpha, gain, clamp)
+    if key in _bias_act_cuda_cache:
+        return _bias_act_cuda_cache[key]
+
+    # Forward op.
+    class BiasActCuda(torch.autograd.Function):
+        @staticmethod
+        def forward(ctx, x, b): # pylint: disable=arguments-differ
+            ctx.memory_format = torch.channels_last if x.ndim > 2 and x.stride()[1] == 1 else torch.contiguous_format
+            x = x.contiguous(memory_format=ctx.memory_format)
+            b = b.contiguous() if b is not None else _null_tensor
+            y = x
+            if act != 'linear' or gain != 1 or clamp >= 0 or b is not _null_tensor:
+                y = _plugin.bias_act(x, b, _null_tensor, _null_tensor, _null_tensor, 0, dim, spec.cuda_idx, alpha, gain, clamp)
+            ctx.save_for_backward(
+                x if 'x' in spec.ref or spec.has_2nd_grad else _null_tensor,
+                b if 'x' in spec.ref or spec.has_2nd_grad else _null_tensor,
+                y if 'y' in spec.ref else _null_tensor)
+            return y
+
+        @staticmethod
+        def backward(ctx, dy): # pylint: disable=arguments-differ
+            dy = dy.contiguous(memory_format=ctx.memory_format)
+            x, b, y = ctx.saved_tensors
+            dx = None
+            db = None
+
+            if ctx.needs_input_grad[0] or ctx.needs_input_grad[1]:
+                dx = dy
+                if act != 'linear' or gain != 1 or clamp >= 0:
+                    dx = BiasActCudaGrad.apply(dy, x, b, y)
+
+            if ctx.needs_input_grad[1]:
+                db = dx.sum([i for i in range(dx.ndim) if i != dim])
+
+            return dx, db
+
+    # Backward op.
+    class BiasActCudaGrad(torch.autograd.Function):
+        @staticmethod
+        def forward(ctx, dy, x, b, y): # pylint: disable=arguments-differ
+            ctx.memory_format = torch.channels_last if dy.ndim > 2 and dy.stride()[1] == 1 else torch.contiguous_format
+            dx = _plugin.bias_act(dy, b, x, y, _null_tensor, 1, dim, spec.cuda_idx, alpha, gain, clamp)
+            ctx.save_for_backward(
+                dy if spec.has_2nd_grad else _null_tensor,
+                x, b, y)
+            return dx
+
+        @staticmethod
+        def backward(ctx, d_dx): # pylint: disable=arguments-differ
+            d_dx = d_dx.contiguous(memory_format=ctx.memory_format)
+            dy, x, b, y = ctx.saved_tensors
+            d_dy = None
+            d_x = None
+            d_b = None
+            d_y = None
+
+            if ctx.needs_input_grad[0]:
+                d_dy = BiasActCudaGrad.apply(d_dx, x, b, y)
+
+            if spec.has_2nd_grad and (ctx.needs_input_grad[1] or ctx.needs_input_grad[2]):
+                d_x = _plugin.bias_act(d_dx, b, x, y, dy, 2, dim, spec.cuda_idx, alpha, gain, clamp)
+
+            if spec.has_2nd_grad and ctx.needs_input_grad[2]:
+                d_b = d_x.sum([i for i in range(d_x.ndim) if i != dim])
+
+            return d_dy, d_x, d_b, d_y
+
+    # Add to cache.
+    _bias_act_cuda_cache[key] = BiasActCuda
+    return BiasActCuda
+
+#----------------------------------------------------------------------------

src/vqvaes/bsqvit/stylegan_utils/ops/conv2d_gradfix.py

@@ -0,0 +1,170 @@
+# Copyright (c) 2021, NVIDIA CORPORATION.  All rights reserved.
+#
+# NVIDIA CORPORATION and its licensors retain all intellectual property
+# and proprietary rights in and to this software, related documentation
+# and any modifications thereto.  Any use, reproduction, disclosure or
+# distribution of this software and related documentation without an express
+# license agreement from NVIDIA CORPORATION is strictly prohibited.
+
+"""Custom replacement for `torch.nn.functional.conv2d` that supports
+arbitrarily high order gradients with zero performance penalty."""
+
+import warnings
+import contextlib
+import torch
+
+# pylint: disable=redefined-builtin
+# pylint: disable=arguments-differ
+# pylint: disable=protected-access
+
+#----------------------------------------------------------------------------
+
+enabled = False                     # Enable the custom op by setting this to true.
+weight_gradients_disabled = False   # Forcefully disable computation of gradients with respect to the weights.
+
+@contextlib.contextmanager
+def no_weight_gradients():
+    global weight_gradients_disabled
+    old = weight_gradients_disabled
+    weight_gradients_disabled = True
+    yield
+    weight_gradients_disabled = old
+
+#----------------------------------------------------------------------------
+
+def conv2d(input, weight, bias=None, stride=1, padding=0, dilation=1, groups=1):
+    if _should_use_custom_op(input):
+        return _conv2d_gradfix(transpose=False, weight_shape=weight.shape, stride=stride, padding=padding, output_padding=0, dilation=dilation, groups=groups).apply(input, weight, bias)
+    return torch.nn.functional.conv2d(input=input, weight=weight, bias=bias, stride=stride, padding=padding, dilation=dilation, groups=groups)
+
+def conv_transpose2d(input, weight, bias=None, stride=1, padding=0, output_padding=0, groups=1, dilation=1):
+    if _should_use_custom_op(input):
+        return _conv2d_gradfix(transpose=True, weight_shape=weight.shape, stride=stride, padding=padding, output_padding=output_padding, groups=groups, dilation=dilation).apply(input, weight, bias)
+    return torch.nn.functional.conv_transpose2d(input=input, weight=weight, bias=bias, stride=stride, padding=padding, output_padding=output_padding, groups=groups, dilation=dilation)
+
+#----------------------------------------------------------------------------
+
+def _should_use_custom_op(input):
+    assert isinstance(input, torch.Tensor)
+    if (not enabled) or (not torch.backends.cudnn.enabled):
+        return False
+    if input.device.type != 'cuda':
+        return False
+    if any(torch.__version__.startswith(x) for x in ['1.7.', '1.8.', '1.9']):
+        return True
+    warnings.warn(f'conv2d_gradfix not supported on PyTorch {torch.__version__}. Falling back to torch.nn.functional.conv2d().')
+    return False
+
+def _tuple_of_ints(xs, ndim):
+    xs = tuple(xs) if isinstance(xs, (tuple, list)) else (xs,) * ndim
+    assert len(xs) == ndim
+    assert all(isinstance(x, int) for x in xs)
+    return xs
+
+#----------------------------------------------------------------------------
+
+_conv2d_gradfix_cache = dict()
+
+def _conv2d_gradfix(transpose, weight_shape, stride, padding, output_padding, dilation, groups):
+    # Parse arguments.
+    ndim = 2
+    weight_shape = tuple(weight_shape)
+    stride = _tuple_of_ints(stride, ndim)
+    padding = _tuple_of_ints(padding, ndim)
+    output_padding = _tuple_of_ints(output_padding, ndim)
+    dilation = _tuple_of_ints(dilation, ndim)
+
+    # Lookup from cache.
+    key = (transpose, weight_shape, stride, padding, output_padding, dilation, groups)
+    if key in _conv2d_gradfix_cache:
+        return _conv2d_gradfix_cache[key]
+
+    # Validate arguments.
+    assert groups >= 1
+    assert len(weight_shape) == ndim + 2
+    assert all(stride[i] >= 1 for i in range(ndim))
+    assert all(padding[i] >= 0 for i in range(ndim))
+    assert all(dilation[i] >= 0 for i in range(ndim))
+    if not transpose:
+        assert all(output_padding[i] == 0 for i in range(ndim))
+    else: # transpose
+        assert all(0 <= output_padding[i] < max(stride[i], dilation[i]) for i in range(ndim))
+
+    # Helpers.
+    common_kwargs = dict(stride=stride, padding=padding, dilation=dilation, groups=groups)
+    def calc_output_padding(input_shape, output_shape):
+        if transpose:
+            return [0, 0]
+        return [
+            input_shape[i + 2]
+            - (output_shape[i + 2] - 1) * stride[i]
+            - (1 - 2 * padding[i])
+            - dilation[i] * (weight_shape[i + 2] - 1)
+            for i in range(ndim)
+        ]
+
+    # Forward & backward.
+    class Conv2d(torch.autograd.Function):
+        @staticmethod
+        def forward(ctx, input, weight, bias):
+            assert weight.shape == weight_shape
+            if not transpose:
+                output = torch.nn.functional.conv2d(input=input, weight=weight, bias=bias, **common_kwargs)
+            else: # transpose
+                output = torch.nn.functional.conv_transpose2d(input=input, weight=weight, bias=bias, output_padding=output_padding, **common_kwargs)
+            ctx.save_for_backward(input, weight)
+            return output
+
+        @staticmethod
+        def backward(ctx, grad_output):
+            input, weight = ctx.saved_tensors
+            grad_input = None
+            grad_weight = None
+            grad_bias = None
+
+            if ctx.needs_input_grad[0]:
+                p = calc_output_padding(input_shape=input.shape, output_shape=grad_output.shape)
+                grad_input = _conv2d_gradfix(transpose=(not transpose), weight_shape=weight_shape, output_padding=p, **common_kwargs).apply(grad_output, weight, None)
+                assert grad_input.shape == input.shape
+
+            if ctx.needs_input_grad[1] and not weight_gradients_disabled:
+                grad_weight = Conv2dGradWeight.apply(grad_output, input)
+                assert grad_weight.shape == weight_shape
+
+            if ctx.needs_input_grad[2]:
+                grad_bias = grad_output.sum([0, 2, 3])
+
+            return grad_input, grad_weight, grad_bias
+
+    # Gradient with respect to the weights.
+    class Conv2dGradWeight(torch.autograd.Function):
+        @staticmethod
+        def forward(ctx, grad_output, input):
+            op = torch._C._jit_get_operation('aten::cudnn_convolution_backward_weight' if not transpose else 'aten::cudnn_convolution_transpose_backward_weight')
+            flags = [torch.backends.cudnn.benchmark, torch.backends.cudnn.deterministic, torch.backends.cudnn.allow_tf32]
+            grad_weight = op(weight_shape, grad_output, input, padding, stride, dilation, groups, *flags)
+            assert grad_weight.shape == weight_shape
+            ctx.save_for_backward(grad_output, input)
+            return grad_weight
+
+        @staticmethod
+        def backward(ctx, grad2_grad_weight):
+            grad_output, input = ctx.saved_tensors
+            grad2_grad_output = None
+            grad2_input = None
+
+            if ctx.needs_input_grad[0]:
+                grad2_grad_output = Conv2d.apply(input, grad2_grad_weight, None)
+                assert grad2_grad_output.shape == grad_output.shape
+
+            if ctx.needs_input_grad[1]:
+                p = calc_output_padding(input_shape=input.shape, output_shape=grad_output.shape)
+                grad2_input = _conv2d_gradfix(transpose=(not transpose), weight_shape=weight_shape, output_padding=p, **common_kwargs).apply(grad_output, grad2_grad_weight, None)
+                assert grad2_input.shape == input.shape
+
+            return grad2_grad_output, grad2_input
+
+    _conv2d_gradfix_cache[key] = Conv2d
+    return Conv2d
+
+#----------------------------------------------------------------------------

src/vqvaes/bsqvit/stylegan_utils/ops/conv2d_resample.py

@@ -0,0 +1,155 @@
+# Copyright (c) 2021, NVIDIA CORPORATION.  All rights reserved.
+#
+# NVIDIA CORPORATION and its licensors retain all intellectual property
+# and proprietary rights in and to this software, related documentation
+# and any modifications thereto.  Any use, reproduction, disclosure or
+# distribution of this software and related documentation without an express
+# license agreement from NVIDIA CORPORATION is strictly prohibited.
+
+"""2D convolution with optional up/downsampling."""
+
+import torch
+
+from .. import misc
+from . import conv2d_gradfix
+from . import upfirdn2d
+from .upfirdn2d import _parse_padding
+from .upfirdn2d import _get_filter_size
+
+#----------------------------------------------------------------------------
+
+def _get_weight_shape(w):
+    with misc.suppress_tracer_warnings(): # this value will be treated as a constant
+        shape = [int(sz) for sz in w.shape]
+    misc.assert_shape(w, shape)
+    return shape
+
+#----------------------------------------------------------------------------
+
+def _conv2d_wrapper(x, w, stride=1, padding=0, groups=1, transpose=False, flip_weight=True):
+    """Wrapper for the underlying `conv2d()` and `conv_transpose2d()` implementations.
+    """
+    out_channels, in_channels_per_group, kh, kw = _get_weight_shape(w)
+
+    # Flip weight if requested.
+    if not flip_weight: # conv2d() actually performs correlation (flip_weight=True) not convolution (flip_weight=False).
+        w = w.flip([2, 3])
+
+    # Workaround performance pitfall in cuDNN 8.0.5, triggered when using
+    # 1x1 kernel + memory_format=channels_last + less than 64 channels.
+    if kw == 1 and kh == 1 and stride == 1 and padding in [0, [0, 0], (0, 0)] and not transpose:
+        if x.stride()[1] == 1 and min(out_channels, in_channels_per_group) < 64:
+            if out_channels <= 4 and groups == 1:
+                in_shape = x.shape
+                x = w.squeeze(3).squeeze(2) @ x.reshape([in_shape[0], in_channels_per_group, -1])
+                x = x.reshape([in_shape[0], out_channels, in_shape[2], in_shape[3]])
+            else:
+                x = x.to(memory_format=torch.contiguous_format)
+                w = w.to(memory_format=torch.contiguous_format)
+                x = conv2d_gradfix.conv2d(x, w, groups=groups)
+            return x.to(memory_format=torch.channels_last)
+
+    # Otherwise => execute using conv2d_gradfix.
+    op = conv2d_gradfix.conv_transpose2d if transpose else conv2d_gradfix.conv2d
+    return op(x, w, stride=stride, padding=padding, groups=groups)
+
+#----------------------------------------------------------------------------
+
+def conv2d_resample(x, w, f=None, up=1, down=1, padding=0, groups=1, flip_weight=True, flip_filter=False):
+    r"""2D convolution with optional up/downsampling.
+
+    Padding is performed only once at the beginning, not between the operations.
+
+    Args:
+        x:              Input tensor of shape
+                        `[batch_size, in_channels, in_height, in_width]`.
+        w:              Weight tensor of shape
+                        `[out_channels, in_channels//groups, kernel_height, kernel_width]`.
+        f:              Low-pass filter for up/downsampling. Must be prepared beforehand by
+                        calling upfirdn2d.setup_filter(). None = identity (default).
+        up:             Integer upsampling factor (default: 1).
+        down:           Integer downsampling factor (default: 1).
+        padding:        Padding with respect to the upsampled image. Can be a single number
+                        or a list/tuple `[x, y]` or `[x_before, x_after, y_before, y_after]`
+                        (default: 0).
+        groups:         Split input channels into N groups (default: 1).
+        flip_weight:    False = convolution, True = correlation (default: True).
+        flip_filter:    False = convolution, True = correlation (default: False).
+
+    Returns:
+        Tensor of the shape `[batch_size, num_channels, out_height, out_width]`.
+    """
+    # Validate arguments.
+    assert isinstance(x, torch.Tensor) and (x.ndim == 4)
+    assert isinstance(w, torch.Tensor) and (w.ndim == 4) and (w.dtype == x.dtype)
+    assert f is None or (isinstance(f, torch.Tensor) and f.ndim in [1, 2] and f.dtype == torch.float32)
+    assert isinstance(up, int) and (up >= 1)
+    assert isinstance(down, int) and (down >= 1)
+    assert isinstance(groups, int) and (groups >= 1)
+    out_channels, in_channels_per_group, kh, kw = _get_weight_shape(w)
+    fw, fh = _get_filter_size(f)
+    px0, px1, py0, py1 = _parse_padding(padding)
+
+    # Adjust padding to account for up/downsampling.
+    if up > 1:
+        px0 += (fw + up - 1) // 2
+        px1 += (fw - up) // 2
+        py0 += (fh + up - 1) // 2
+        py1 += (fh - up) // 2
+    if down > 1:
+        px0 += (fw - down + 1) // 2
+        px1 += (fw - down) // 2
+        py0 += (fh - down + 1) // 2
+        py1 += (fh - down) // 2
+
+    # Fast path: 1x1 convolution with downsampling only => downsample first, then convolve.
+    if kw == 1 and kh == 1 and (down > 1 and up == 1):
+        x = upfirdn2d.upfirdn2d(x=x, f=f, down=down, padding=[px0,px1,py0,py1], flip_filter=flip_filter)
+        x = _conv2d_wrapper(x=x, w=w, groups=groups, flip_weight=flip_weight)
+        return x
+
+    # Fast path: 1x1 convolution with upsampling only => convolve first, then upsample.
+    if kw == 1 and kh == 1 and (up > 1 and down == 1):
+        x = _conv2d_wrapper(x=x, w=w, groups=groups, flip_weight=flip_weight)
+        x = upfirdn2d.upfirdn2d(x=x, f=f, up=up, padding=[px0,px1,py0,py1], gain=up**2, flip_filter=flip_filter)
+        return x
+
+    # Fast path: downsampling only => use strided convolution.
+    if down > 1 and up == 1:
+        x = upfirdn2d.upfirdn2d(x=x, f=f, padding=[px0,px1,py0,py1], flip_filter=flip_filter)
+        x = _conv2d_wrapper(x=x, w=w, stride=down, groups=groups, flip_weight=flip_weight)
+        return x
+
+    # Fast path: upsampling with optional downsampling => use transpose strided convolution.
+    if up > 1:
+        if groups == 1:
+            w = w.transpose(0, 1)
+        else:
+            w = w.reshape(groups, out_channels // groups, in_channels_per_group, kh, kw)
+            w = w.transpose(1, 2)
+            w = w.reshape(groups * in_channels_per_group, out_channels // groups, kh, kw)
+        px0 -= kw - 1
+        px1 -= kw - up
+        py0 -= kh - 1
+        py1 -= kh - up
+        pxt = max(min(-px0, -px1), 0)
+        pyt = max(min(-py0, -py1), 0)
+        x = _conv2d_wrapper(x=x, w=w, stride=up, padding=[pyt,pxt], groups=groups, transpose=True, flip_weight=(not flip_weight))
+        x = upfirdn2d.upfirdn2d(x=x, f=f, padding=[px0+pxt,px1+pxt,py0+pyt,py1+pyt], gain=up**2, flip_filter=flip_filter)
+        if down > 1:
+            x = upfirdn2d.upfirdn2d(x=x, f=f, down=down, flip_filter=flip_filter)
+        return x
+
+    # Fast path: no up/downsampling, padding supported by the underlying implementation => use plain conv2d.
+    if up == 1 and down == 1:
+        if px0 == px1 and py0 == py1 and px0 >= 0 and py0 >= 0:
+            return _conv2d_wrapper(x=x, w=w, padding=[py0,px0], groups=groups, flip_weight=flip_weight)
+
+    # Fallback: Generic reference implementation.
+    x = upfirdn2d.upfirdn2d(x=x, f=(f if up > 1 else None), up=up, padding=[px0,px1,py0,py1], gain=up**2, flip_filter=flip_filter)
+    x = _conv2d_wrapper(x=x, w=w, groups=groups, flip_weight=flip_weight)
+    if down > 1:
+        x = upfirdn2d.upfirdn2d(x=x, f=f, down=down, flip_filter=flip_filter)
+    return x
+
+#----------------------------------------------------------------------------

src/vqvaes/bsqvit/stylegan_utils/ops/upfirdn2d.cpp

@@ -0,0 +1,103 @@
+// Copyright (c) 2021, NVIDIA CORPORATION.  All rights reserved.
+//
+// NVIDIA CORPORATION and its licensors retain all intellectual property
+// and proprietary rights in and to this software, related documentation
+// and any modifications thereto.  Any use, reproduction, disclosure or
+// distribution of this software and related documentation without an express
+// license agreement from NVIDIA CORPORATION is strictly prohibited.
+
+#include <torch/extension.h>
+#include <ATen/cuda/CUDAContext.h>
+#include <c10/cuda/CUDAGuard.h>
+#include "upfirdn2d.h"
+
+//------------------------------------------------------------------------
+
+static torch::Tensor upfirdn2d(torch::Tensor x, torch::Tensor f, int upx, int upy, int downx, int downy, int padx0, int padx1, int pady0, int pady1, bool flip, float gain)
+{
+    // Validate arguments.
+    TORCH_CHECK(x.is_cuda(), "x must reside on CUDA device");
+    TORCH_CHECK(f.device() == x.device(), "f must reside on the same device as x");
+    TORCH_CHECK(f.dtype() == torch::kFloat, "f must be float32");
+    TORCH_CHECK(x.numel() <= INT_MAX, "x is too large");
+    TORCH_CHECK(f.numel() <= INT_MAX, "f is too large");
+    TORCH_CHECK(x.dim() == 4, "x must be rank 4");
+    TORCH_CHECK(f.dim() == 2, "f must be rank 2");
+    TORCH_CHECK(f.size(0) >= 1 && f.size(1) >= 1, "f must be at least 1x1");
+    TORCH_CHECK(upx >= 1 && upy >= 1, "upsampling factor must be at least 1");
+    TORCH_CHECK(downx >= 1 && downy >= 1, "downsampling factor must be at least 1");
+
+    // Create output tensor.
+    const at::cuda::OptionalCUDAGuard device_guard(device_of(x));
+    int outW = ((int)x.size(3) * upx + padx0 + padx1 - (int)f.size(1) + downx) / downx;
+    int outH = ((int)x.size(2) * upy + pady0 + pady1 - (int)f.size(0) + downy) / downy;
+    TORCH_CHECK(outW >= 1 && outH >= 1, "output must be at least 1x1");
+    torch::Tensor y = torch::empty({x.size(0), x.size(1), outH, outW}, x.options(), x.suggest_memory_format());
+    TORCH_CHECK(y.numel() <= INT_MAX, "output is too large");
+
+    // Initialize CUDA kernel parameters.
+    upfirdn2d_kernel_params p;
+    p.x             = x.data_ptr();
+    p.f             = f.data_ptr<float>();
+    p.y             = y.data_ptr();
+    p.up            = make_int2(upx, upy);
+    p.down          = make_int2(downx, downy);
+    p.pad0          = make_int2(padx0, pady0);
+    p.flip          = (flip) ? 1 : 0;
+    p.gain          = gain;
+    p.inSize        = make_int4((int)x.size(3), (int)x.size(2), (int)x.size(1), (int)x.size(0));
+    p.inStride      = make_int4((int)x.stride(3), (int)x.stride(2), (int)x.stride(1), (int)x.stride(0));
+    p.filterSize    = make_int2((int)f.size(1), (int)f.size(0));
+    p.filterStride  = make_int2((int)f.stride(1), (int)f.stride(0));
+    p.outSize       = make_int4((int)y.size(3), (int)y.size(2), (int)y.size(1), (int)y.size(0));
+    p.outStride     = make_int4((int)y.stride(3), (int)y.stride(2), (int)y.stride(1), (int)y.stride(0));
+    p.sizeMajor     = (p.inStride.z == 1) ? p.inSize.w : p.inSize.w * p.inSize.z;
+    p.sizeMinor     = (p.inStride.z == 1) ? p.inSize.z : 1;
+
+    // Choose CUDA kernel.
+    upfirdn2d_kernel_spec spec;
+    AT_DISPATCH_REDUCED_FLOATING_TYPES(x.scalar_type(), "upfirdn2d_cuda", [&]
+    {
+        spec = choose_upfirdn2d_kernel<scalar_t>(p);
+    });
+
+    // Set looping options.
+    p.loopMajor     = (p.sizeMajor - 1) / 16384 + 1;
+    p.loopMinor     = spec.loopMinor;
+    p.loopX         = spec.loopX;
+    p.launchMinor   = (p.sizeMinor - 1) / p.loopMinor + 1;
+    p.launchMajor   = (p.sizeMajor - 1) / p.loopMajor + 1;
+
+    // Compute grid size.
+    dim3 blockSize, gridSize;
+    if (spec.tileOutW < 0) // large
+    {
+        blockSize = dim3(4, 32, 1);
+        gridSize = dim3(
+            ((p.outSize.y - 1) / blockSize.x + 1) * p.launchMinor,
+            (p.outSize.x - 1) / (blockSize.y * p.loopX) + 1,
+            p.launchMajor);
+    }
+    else // small
+    {
+        blockSize = dim3(256, 1, 1);
+        gridSize = dim3(
+            ((p.outSize.y - 1) / spec.tileOutH + 1) * p.launchMinor,
+            (p.outSize.x - 1) / (spec.tileOutW * p.loopX) + 1,
+            p.launchMajor);
+    }
+
+    // Launch CUDA kernel.
+    void* args[] = {&p};
+    AT_CUDA_CHECK(cudaLaunchKernel(spec.kernel, gridSize, blockSize, args, 0, at::cuda::getCurrentCUDAStream()));
+    return y;
+}
+
+//------------------------------------------------------------------------
+
+PYBIND11_MODULE(TORCH_EXTENSION_NAME, m)
+{
+    m.def("upfirdn2d", &upfirdn2d);
+}
+
+//------------------------------------------------------------------------

src/vqvaes/bsqvit/stylegan_utils/ops/upfirdn2d.cu

@@ -0,0 +1,353 @@
+// Copyright (c) 2021, NVIDIA CORPORATION.  All rights reserved.
+//
+// NVIDIA CORPORATION and its licensors retain all intellectual property
+// and proprietary rights in and to this software, related documentation
+// and any modifications thereto.  Any use, reproduction, disclosure or
+// distribution of this software and related documentation without an express
+// license agreement from NVIDIA CORPORATION is strictly prohibited.
+
+#include <ATen/ATen.h>
+#include <c10/util/Half.h>
+#include "upfirdn2d.h"
+
+//------------------------------------------------------------------------
+// Helpers.
+
+template <class T> struct InternalType;
+template <> struct InternalType<double>         { typedef double scalar_t; };
+template <> struct InternalType<float>          { typedef float  scalar_t; };
+template <> struct InternalType<c10::Half>      { typedef float  scalar_t; };
+template <> struct InternalType<at::BFloat16>   { typedef float  scalar_t; };
+
+static __device__ __forceinline__ int floor_div(int a, int b)
+{
+    int t = 1 - a / b;
+    return (a + t * b) / b - t;
+}
+
+//------------------------------------------------------------------------
+// Generic CUDA implementation for large filters.
+
+template <class T> static __global__ void upfirdn2d_kernel_large(upfirdn2d_kernel_params p)
+{
+    typedef typename InternalType<T>::scalar_t scalar_t;
+
+    // Calculate thread index.
+    int minorBase = blockIdx.x * blockDim.x + threadIdx.x;
+    int outY = minorBase / p.launchMinor;
+    minorBase -= outY * p.launchMinor;
+    int outXBase = blockIdx.y * p.loopX * blockDim.y + threadIdx.y;
+    int majorBase = blockIdx.z * p.loopMajor;
+    if (outXBase >= p.outSize.x | outY >= p.outSize.y | majorBase >= p.sizeMajor)
+        return;
+
+    // Setup Y receptive field.
+    int midY = outY * p.down.y + p.up.y - 1 - p.pad0.y;
+    int inY = min(max(floor_div(midY, p.up.y), 0), p.inSize.y);
+    int h = min(max(floor_div(midY + p.filterSize.y, p.up.y), 0), p.inSize.y) - inY;
+    int filterY = midY + p.filterSize.y - (inY + 1) * p.up.y;
+    if (p.flip)
+        filterY = p.filterSize.y - 1 - filterY;
+
+    // Loop over major, minor, and X.
+    for (int majorIdx = 0, major = majorBase; majorIdx < p.loopMajor & major < p.sizeMajor; majorIdx++, major++)
+    for (int minorIdx = 0, minor = minorBase; minorIdx < p.loopMinor & minor < p.sizeMinor; minorIdx++, minor += p.launchMinor)
+    {
+        int nc = major * p.sizeMinor + minor;
+        int n = nc / p.inSize.z;
+        int c = nc - n * p.inSize.z;
+        for (int loopX = 0, outX = outXBase; loopX < p.loopX & outX < p.outSize.x; loopX++, outX += blockDim.y)
+        {
+            // Setup X receptive field.
+            int midX = outX * p.down.x + p.up.x - 1 - p.pad0.x;
+            int inX = min(max(floor_div(midX, p.up.x), 0), p.inSize.x);
+            int w = min(max(floor_div(midX + p.filterSize.x, p.up.x), 0), p.inSize.x) - inX;
+            int filterX = midX + p.filterSize.x - (inX + 1) * p.up.x;
+            if (p.flip)
+                filterX = p.filterSize.x - 1 - filterX;
+
+            // Initialize pointers.
+            const T* xp = &((const T*)p.x)[inX * p.inStride.x + inY * p.inStride.y + c * p.inStride.z + n * p.inStride.w];
+            const float* fp = &p.f[filterX * p.filterStride.x + filterY * p.filterStride.y];
+            int filterStepX = ((p.flip) ? p.up.x : -p.up.x) * p.filterStride.x;
+            int filterStepY = ((p.flip) ? p.up.y : -p.up.y) * p.filterStride.y;
+
+            // Inner loop.
+            scalar_t v = 0;
+            for (int y = 0; y < h; y++)
+            {
+                for (int x = 0; x < w; x++)
+                {
+                    v += (scalar_t)(*xp) * (scalar_t)(*fp);
+                    xp += p.inStride.x;
+                    fp += filterStepX;
+                }
+                xp += p.inStride.y - w * p.inStride.x;
+                fp += filterStepY - w * filterStepX;
+            }
+
+            // Store result.
+            v *= p.gain;
+            ((T*)p.y)[outX * p.outStride.x + outY * p.outStride.y + c * p.outStride.z + n * p.outStride.w] = (T)v;
+        }
+    }
+}
+
+//------------------------------------------------------------------------
+// Specialized CUDA implementation for small filters.
+
+template <class T, int upx, int upy, int downx, int downy, int filterW, int filterH, int tileOutW, int tileOutH, int loopMinor>
+static __global__ void upfirdn2d_kernel_small(upfirdn2d_kernel_params p)
+{
+    typedef typename InternalType<T>::scalar_t scalar_t;
+    const int tileInW = ((tileOutW - 1) * downx + filterW - 1) / upx + 1;
+    const int tileInH = ((tileOutH - 1) * downy + filterH - 1) / upy + 1;
+    __shared__ volatile scalar_t sf[filterH][filterW];
+    __shared__ volatile scalar_t sx[tileInH][tileInW][loopMinor];
+
+    // Calculate tile index.
+    int minorBase = blockIdx.x;
+    int tileOutY = minorBase / p.launchMinor;
+    minorBase -= tileOutY * p.launchMinor;
+    minorBase *= loopMinor;
+    tileOutY *= tileOutH;
+    int tileOutXBase = blockIdx.y * p.loopX * tileOutW;
+    int majorBase = blockIdx.z * p.loopMajor;
+    if (tileOutXBase >= p.outSize.x | tileOutY >= p.outSize.y | majorBase >= p.sizeMajor)
+        return;
+
+    // Load filter (flipped).
+    for (int tapIdx = threadIdx.x; tapIdx < filterH * filterW; tapIdx += blockDim.x)
+    {
+        int fy = tapIdx / filterW;
+        int fx = tapIdx - fy * filterW;
+        scalar_t v = 0;
+        if (fx < p.filterSize.x & fy < p.filterSize.y)
+        {
+            int ffx = (p.flip) ? fx : p.filterSize.x - 1 - fx;
+            int ffy = (p.flip) ? fy : p.filterSize.y - 1 - fy;
+            v = (scalar_t)p.f[ffx * p.filterStride.x + ffy * p.filterStride.y];
+        }
+        sf[fy][fx] = v;
+    }
+
+    // Loop over major and X.
+    for (int majorIdx = 0, major = majorBase; majorIdx < p.loopMajor & major < p.sizeMajor; majorIdx++, major++)
+    {
+        int baseNC = major * p.sizeMinor + minorBase;
+        int n = baseNC / p.inSize.z;
+        int baseC = baseNC - n * p.inSize.z;
+        for (int loopX = 0, tileOutX = tileOutXBase; loopX < p.loopX & tileOutX < p.outSize.x; loopX++, tileOutX += tileOutW)
+        {
+            // Load input pixels.
+            int tileMidX = tileOutX * downx + upx - 1 - p.pad0.x;
+            int tileMidY = tileOutY * downy + upy - 1 - p.pad0.y;
+            int tileInX = floor_div(tileMidX, upx);
+            int tileInY = floor_div(tileMidY, upy);
+            __syncthreads();
+            for (int inIdx = threadIdx.x; inIdx < tileInH * tileInW * loopMinor; inIdx += blockDim.x)
+            {
+                int relC = inIdx;
+                int relInX = relC / loopMinor;
+                int relInY = relInX / tileInW;
+                relC -= relInX * loopMinor;
+                relInX -= relInY * tileInW;
+                int c = baseC + relC;
+                int inX = tileInX + relInX;
+                int inY = tileInY + relInY;
+                scalar_t v = 0;
+                if (inX >= 0 & inY >= 0 & inX < p.inSize.x & inY < p.inSize.y & c < p.inSize.z)
+                    v = (scalar_t)((const T*)p.x)[inX * p.inStride.x + inY * p.inStride.y + c * p.inStride.z + n * p.inStride.w];
+                sx[relInY][relInX][relC] = v;
+            }
+
+            // Loop over output pixels.
+            __syncthreads();
+            for (int outIdx = threadIdx.x; outIdx < tileOutH * tileOutW * loopMinor; outIdx += blockDim.x)
+            {
+                int relC = outIdx;
+                int relOutX = relC / loopMinor;
+                int relOutY = relOutX / tileOutW;
+                relC -= relOutX * loopMinor;
+                relOutX -= relOutY * tileOutW;
+                int c = baseC + relC;
+                int outX = tileOutX + relOutX;
+                int outY = tileOutY + relOutY;
+
+                // Setup receptive field.
+                int midX = tileMidX + relOutX * downx;
+                int midY = tileMidY + relOutY * downy;
+                int inX = floor_div(midX, upx);
+                int inY = floor_div(midY, upy);
+                int relInX = inX - tileInX;
+                int relInY = inY - tileInY;
+                int filterX = (inX + 1) * upx - midX - 1; // flipped
+                int filterY = (inY + 1) * upy - midY - 1; // flipped
+
+                // Inner loop.
+                if (outX < p.outSize.x & outY < p.outSize.y & c < p.outSize.z)
+                {
+                    scalar_t v = 0;
+                    #pragma unroll
+                    for (int y = 0; y < filterH / upy; y++)
+                        #pragma unroll
+                        for (int x = 0; x < filterW / upx; x++)
+                            v += sx[relInY + y][relInX + x][relC] * sf[filterY + y * upy][filterX + x * upx];
+                    v *= p.gain;
+                    ((T*)p.y)[outX * p.outStride.x + outY * p.outStride.y + c * p.outStride.z + n * p.outStride.w] = (T)v;
+                }
+            }
+        }
+    }
+}
+
+//------------------------------------------------------------------------
+// CUDA kernel selection.
+
+template <class T> upfirdn2d_kernel_spec choose_upfirdn2d_kernel(const upfirdn2d_kernel_params& p)
+{
+    int s = p.inStride.z, fx = p.filterSize.x, fy = p.filterSize.y;
+
+    upfirdn2d_kernel_spec spec = {(void*)upfirdn2d_kernel_large<T>, -1,-1,1, 4}; // contiguous
+    if (s == 1)           spec = {(void*)upfirdn2d_kernel_large<T>, -1,-1,4, 1}; // channels_last
+
+    if (s != 1 && p.up.x == 1 && p.up.y == 1 && p.down.x == 1 && p.down.y == 1) // contiguous
+    {
+        if (fx <= 7  && fy <= 7 ) spec = {(void*)upfirdn2d_kernel_small<T, 1,1, 1,1, 7,7,  64,16,1>, 64,16,1, 1};
+        if (fx <= 6  && fy <= 6 ) spec = {(void*)upfirdn2d_kernel_small<T, 1,1, 1,1, 6,6,  64,16,1>, 64,16,1, 1};
+        if (fx <= 5  && fy <= 5 ) spec = {(void*)upfirdn2d_kernel_small<T, 1,1, 1,1, 5,5,  64,16,1>, 64,16,1, 1};
+        if (fx <= 4  && fy <= 4 ) spec = {(void*)upfirdn2d_kernel_small<T, 1,1, 1,1, 4,4,  64,16,1>, 64,16,1, 1};
+        if (fx <= 3  && fy <= 3 ) spec = {(void*)upfirdn2d_kernel_small<T, 1,1, 1,1, 3,3,  64,16,1>, 64,16,1, 1};
+        if (fx <= 24 && fy <= 1 ) spec = {(void*)upfirdn2d_kernel_small<T, 1,1, 1,1, 24,1, 128,8,1>, 128,8,1, 1};
+        if (fx <= 20 && fy <= 1 ) spec = {(void*)upfirdn2d_kernel_small<T, 1,1, 1,1, 20,1, 128,8,1>, 128,8,1, 1};
+        if (fx <= 16 && fy <= 1 ) spec = {(void*)upfirdn2d_kernel_small<T, 1,1, 1,1, 16,1, 128,8,1>, 128,8,1, 1};
+        if (fx <= 12 && fy <= 1 ) spec = {(void*)upfirdn2d_kernel_small<T, 1,1, 1,1, 12,1, 128,8,1>, 128,8,1, 1};
+        if (fx <= 8  && fy <= 1 ) spec = {(void*)upfirdn2d_kernel_small<T, 1,1, 1,1, 8,1,  128,8,1>, 128,8,1, 1};
+        if (fx <= 1  && fy <= 24) spec = {(void*)upfirdn2d_kernel_small<T, 1,1, 1,1, 1,24, 32,32,1>, 32,32,1, 1};
+        if (fx <= 1  && fy <= 20) spec = {(void*)upfirdn2d_kernel_small<T, 1,1, 1,1, 1,20, 32,32,1>, 32,32,1, 1};
+        if (fx <= 1  && fy <= 16) spec = {(void*)upfirdn2d_kernel_small<T, 1,1, 1,1, 1,16, 32,32,1>, 32,32,1, 1};
+        if (fx <= 1  && fy <= 12) spec = {(void*)upfirdn2d_kernel_small<T, 1,1, 1,1, 1,12, 32,32,1>, 32,32,1, 1};
+        if (fx <= 1  && fy <= 8 ) spec = {(void*)upfirdn2d_kernel_small<T, 1,1, 1,1, 1,8,  32,32,1>, 32,32,1, 1};
+    }
+    if (s == 1 && p.up.x == 1 && p.up.y == 1 && p.down.x == 1 && p.down.y == 1) // channels_last
+    {
+        if (fx <= 7  && fy <= 7 ) spec = {(void*)upfirdn2d_kernel_small<T, 1,1, 1,1, 7,7,  16,16,8>,  16,16,8,  1};
+        if (fx <= 6  && fy <= 6 ) spec = {(void*)upfirdn2d_kernel_small<T, 1,1, 1,1, 4,4,  16,16,8>,  16,16,8,  1};
+        if (fx <= 5  && fy <= 5 ) spec = {(void*)upfirdn2d_kernel_small<T, 1,1, 1,1, 4,4,  16,16,8>,  16,16,8,  1};
+        if (fx <= 4  && fy <= 4 ) spec = {(void*)upfirdn2d_kernel_small<T, 1,1, 1,1, 4,4,  16,16,8>,  16,16,8,  1};
+        if (fx <= 3  && fy <= 3 ) spec = {(void*)upfirdn2d_kernel_small<T, 1,1, 1,1, 4,4,  16,16,8>,  16,16,8,  1};
+        if (fx <= 24 && fy <= 1 ) spec = {(void*)upfirdn2d_kernel_small<T, 1,1, 1,1, 24,1, 128,1,16>, 128,1,16, 1};
+        if (fx <= 20 && fy <= 1 ) spec = {(void*)upfirdn2d_kernel_small<T, 1,1, 1,1, 20,1, 128,1,16>, 128,1,16, 1};
+        if (fx <= 16 && fy <= 1 ) spec = {(void*)upfirdn2d_kernel_small<T, 1,1, 1,1, 16,1, 128,1,16>, 128,1,16, 1};
+        if (fx <= 12 && fy <= 1 ) spec = {(void*)upfirdn2d_kernel_small<T, 1,1, 1,1, 12,1, 128,1,16>, 128,1,16, 1};
+        if (fx <= 8  && fy <= 1 ) spec = {(void*)upfirdn2d_kernel_small<T, 1,1, 1,1, 8,1,  128,1,16>, 128,1,16, 1};
+        if (fx <= 1  && fy <= 24) spec = {(void*)upfirdn2d_kernel_small<T, 1,1, 1,1, 1,24, 1,128,16>, 1,128,16, 1};
+        if (fx <= 1  && fy <= 20) spec = {(void*)upfirdn2d_kernel_small<T, 1,1, 1,1, 1,20, 1,128,16>, 1,128,16, 1};
+        if (fx <= 1  && fy <= 16) spec = {(void*)upfirdn2d_kernel_small<T, 1,1, 1,1, 1,16, 1,128,16>, 1,128,16, 1};
+        if (fx <= 1  && fy <= 12) spec = {(void*)upfirdn2d_kernel_small<T, 1,1, 1,1, 1,12, 1,128,16>, 1,128,16, 1};
+        if (fx <= 1  && fy <= 8 ) spec = {(void*)upfirdn2d_kernel_small<T, 1,1, 1,1, 1,8,  1,128,16>, 1,128,16, 1};
+    }
+    if (s != 1 && p.up.x == 2 && p.up.y == 2 && p.down.x == 1 && p.down.y == 1) // contiguous
+    {
+        if (fx <= 8  && fy <= 8 ) spec = {(void*)upfirdn2d_kernel_small<T, 2,2, 1,1, 8,8, 64,16,1>, 64,16,1, 1};
+        if (fx <= 6  && fy <= 6 ) spec = {(void*)upfirdn2d_kernel_small<T, 2,2, 1,1, 6,6, 64,16,1>, 64,16,1, 1};
+        if (fx <= 4  && fy <= 4 ) spec = {(void*)upfirdn2d_kernel_small<T, 2,2, 1,1, 4,4, 64,16,1>, 64,16,1, 1};
+        if (fx <= 2  && fy <= 2 ) spec = {(void*)upfirdn2d_kernel_small<T, 2,2, 1,1, 2,2, 64,16,1>, 64,16,1, 1};
+    }
+    if (s == 1 && p.up.x == 2 && p.up.y == 2 && p.down.x == 1 && p.down.y == 1) // channels_last
+    {
+        if (fx <= 8  && fy <= 8 ) spec = {(void*)upfirdn2d_kernel_small<T, 2,2, 1,1, 8,8, 16,16,8>, 16,16,8, 1};
+        if (fx <= 6  && fy <= 6 ) spec = {(void*)upfirdn2d_kernel_small<T, 2,2, 1,1, 6,6, 16,16,8>, 16,16,8, 1};
+        if (fx <= 4  && fy <= 4 ) spec = {(void*)upfirdn2d_kernel_small<T, 2,2, 1,1, 4,4, 16,16,8>, 16,16,8, 1};
+        if (fx <= 2  && fy <= 2 ) spec = {(void*)upfirdn2d_kernel_small<T, 2,2, 1,1, 2,2, 16,16,8>, 16,16,8, 1};
+    }
+    if (s != 1 && p.up.x == 2 && p.up.y == 1 && p.down.x == 1 && p.down.y == 1) // contiguous
+    {
+        if (fx <= 24 && fy <= 1 ) spec = {(void*)upfirdn2d_kernel_small<T, 2,1, 1,1, 24,1, 128,8,1>, 128,8,1, 1};
+        if (fx <= 20 && fy <= 1 ) spec = {(void*)upfirdn2d_kernel_small<T, 2,1, 1,1, 20,1, 128,8,1>, 128,8,1, 1};
+        if (fx <= 16 && fy <= 1 ) spec = {(void*)upfirdn2d_kernel_small<T, 2,1, 1,1, 16,1, 128,8,1>, 128,8,1, 1};
+        if (fx <= 12 && fy <= 1 ) spec = {(void*)upfirdn2d_kernel_small<T, 2,1, 1,1, 12,1, 128,8,1>, 128,8,1, 1};
+        if (fx <= 8  && fy <= 1 ) spec = {(void*)upfirdn2d_kernel_small<T, 2,1, 1,1, 8,1,  128,8,1>, 128,8,1, 1};
+    }
+    if (s == 1 && p.up.x == 2 && p.up.y == 1 && p.down.x == 1 && p.down.y == 1) // channels_last
+    {
+        if (fx <= 24 && fy <= 1 ) spec = {(void*)upfirdn2d_kernel_small<T, 2,1, 1,1, 24,1, 128,1,16>, 128,1,16, 1};
+        if (fx <= 20 && fy <= 1 ) spec = {(void*)upfirdn2d_kernel_small<T, 2,1, 1,1, 20,1, 128,1,16>, 128,1,16, 1};
+        if (fx <= 16 && fy <= 1 ) spec = {(void*)upfirdn2d_kernel_small<T, 2,1, 1,1, 16,1, 128,1,16>, 128,1,16, 1};
+        if (fx <= 12 && fy <= 1 ) spec = {(void*)upfirdn2d_kernel_small<T, 2,1, 1,1, 12,1, 128,1,16>, 128,1,16, 1};
+        if (fx <= 8  && fy <= 1 ) spec = {(void*)upfirdn2d_kernel_small<T, 2,1, 1,1, 8,1,  128,1,16>, 128,1,16, 1};
+    }
+    if (s != 1 && p.up.x == 1 && p.up.y == 2 && p.down.x == 1 && p.down.y == 1) // contiguous
+    {
+        if (fx <= 1  && fy <= 24) spec = {(void*)upfirdn2d_kernel_small<T, 1,2, 1,1, 1,24, 32,32,1>, 32,32,1, 1};
+        if (fx <= 1  && fy <= 20) spec = {(void*)upfirdn2d_kernel_small<T, 1,2, 1,1, 1,20, 32,32,1>, 32,32,1, 1};
+        if (fx <= 1  && fy <= 16) spec = {(void*)upfirdn2d_kernel_small<T, 1,2, 1,1, 1,16, 32,32,1>, 32,32,1, 1};
+        if (fx <= 1  && fy <= 12) spec = {(void*)upfirdn2d_kernel_small<T, 1,2, 1,1, 1,12, 32,32,1>, 32,32,1, 1};
+        if (fx <= 1  && fy <= 8 ) spec = {(void*)upfirdn2d_kernel_small<T, 1,2, 1,1, 1,8,  32,32,1>, 32,32,1, 1};
+    }
+    if (s == 1 && p.up.x == 1 && p.up.y == 2 && p.down.x == 1 && p.down.y == 1) // channels_last
+    {
+        if (fx <= 1  && fy <= 24) spec = {(void*)upfirdn2d_kernel_small<T, 1,2, 1,1, 1,24, 1,128,16>, 1,128,16, 1};
+        if (fx <= 1  && fy <= 20) spec = {(void*)upfirdn2d_kernel_small<T, 1,2, 1,1, 1,20, 1,128,16>, 1,128,16, 1};
+        if (fx <= 1  && fy <= 16) spec = {(void*)upfirdn2d_kernel_small<T, 1,2, 1,1, 1,16, 1,128,16>, 1,128,16, 1};
+        if (fx <= 1  && fy <= 12) spec = {(void*)upfirdn2d_kernel_small<T, 1,2, 1,1, 1,12, 1,128,16>, 1,128,16, 1};
+        if (fx <= 1  && fy <= 8 ) spec = {(void*)upfirdn2d_kernel_small<T, 1,2, 1,1, 1,8,  1,128,16>, 1,128,16, 1};
+    }
+    if (s != 1 && p.up.x == 1 && p.up.y == 1 && p.down.x == 2 && p.down.y == 2) // contiguous
+    {
+        if (fx <= 8  && fy <= 8 ) spec = {(void*)upfirdn2d_kernel_small<T, 1,1, 2,2, 8,8,  32,8,1>, 32,8,1, 1};
+        if (fx <= 6  && fy <= 6 ) spec = {(void*)upfirdn2d_kernel_small<T, 1,1, 2,2, 6,6,  32,8,1>, 32,8,1, 1};
+        if (fx <= 4  && fy <= 4 ) spec = {(void*)upfirdn2d_kernel_small<T, 1,1, 2,2, 4,4,  32,8,1>, 32,8,1, 1};
+        if (fx <= 2  && fy <= 2 ) spec = {(void*)upfirdn2d_kernel_small<T, 1,1, 2,2, 2,2,  32,8,1>, 32,8,1, 1};
+    }
+    if (s == 1 && p.up.x == 1 && p.up.y == 1 && p.down.x == 2 && p.down.y == 2) // channels_last
+    {
+        if (fx <= 8  && fy <= 8 ) spec = {(void*)upfirdn2d_kernel_small<T, 1,1, 2,2, 8,8,  8,8,8>, 8,8,8, 1};
+        if (fx <= 6  && fy <= 6 ) spec = {(void*)upfirdn2d_kernel_small<T, 1,1, 2,2, 6,6,  8,8,8>, 8,8,8, 1};
+        if (fx <= 4  && fy <= 4 ) spec = {(void*)upfirdn2d_kernel_small<T, 1,1, 2,2, 4,4,  8,8,8>, 8,8,8, 1};
+        if (fx <= 2  && fy <= 2 ) spec = {(void*)upfirdn2d_kernel_small<T, 1,1, 2,2, 2,2,  8,8,8>, 8,8,8, 1};
+    }
+    if (s != 1 && p.up.x == 1 && p.up.y == 1 && p.down.x == 2 && p.down.y == 1) // contiguous
+    {
+        if (fx <= 24 && fy <= 1 ) spec = {(void*)upfirdn2d_kernel_small<T, 1,1, 2,1, 24,1, 64,8,1>, 64,8,1, 1};
+        if (fx <= 20 && fy <= 1 ) spec = {(void*)upfirdn2d_kernel_small<T, 1,1, 2,1, 20,1, 64,8,1>, 64,8,1, 1};
+        if (fx <= 16 && fy <= 1 ) spec = {(void*)upfirdn2d_kernel_small<T, 1,1, 2,1, 16,1, 64,8,1>, 64,8,1, 1};
+        if (fx <= 12 && fy <= 1 ) spec = {(void*)upfirdn2d_kernel_small<T, 1,1, 2,1, 12,1, 64,8,1>, 64,8,1, 1};
+        if (fx <= 8  && fy <= 1 ) spec = {(void*)upfirdn2d_kernel_small<T, 1,1, 2,1, 8,1,  64,8,1>, 64,8,1, 1};
+    }
+    if (s == 1 && p.up.x == 1 && p.up.y == 1 && p.down.x == 2 && p.down.y == 1) // channels_last
+    {
+        if (fx <= 24 && fy <= 1 ) spec = {(void*)upfirdn2d_kernel_small<T, 1,1, 2,1, 24,1, 64,1,8>, 64,1,8, 1};
+        if (fx <= 20 && fy <= 1 ) spec = {(void*)upfirdn2d_kernel_small<T, 1,1, 2,1, 20,1, 64,1,8>, 64,1,8, 1};
+        if (fx <= 16 && fy <= 1 ) spec = {(void*)upfirdn2d_kernel_small<T, 1,1, 2,1, 16,1, 64,1,8>, 64,1,8, 1};
+        if (fx <= 12 && fy <= 1 ) spec = {(void*)upfirdn2d_kernel_small<T, 1,1, 2,1, 12,1, 64,1,8>, 64,1,8, 1};
+        if (fx <= 8  && fy <= 1 ) spec = {(void*)upfirdn2d_kernel_small<T, 1,1, 2,1, 8,1,  64,1,8>, 64,1,8, 1};
+    }
+    if (s != 1 && p.up.x == 1 && p.up.y == 1 && p.down.x == 1 && p.down.y == 2) // contiguous
+    {
+        if (fx <= 1  && fy <= 24) spec = {(void*)upfirdn2d_kernel_small<T, 1,1, 1,2, 1,24, 32,16,1>, 32,16,1, 1};
+        if (fx <= 1  && fy <= 20) spec = {(void*)upfirdn2d_kernel_small<T, 1,1, 1,2, 1,20, 32,16,1>, 32,16,1, 1};
+        if (fx <= 1  && fy <= 16) spec = {(void*)upfirdn2d_kernel_small<T, 1,1, 1,2, 1,16, 32,16,1>, 32,16,1, 1};
+        if (fx <= 1  && fy <= 12) spec = {(void*)upfirdn2d_kernel_small<T, 1,1, 1,2, 1,12, 32,16,1>, 32,16,1, 1};
+        if (fx <= 1  && fy <= 8 ) spec = {(void*)upfirdn2d_kernel_small<T, 1,1, 1,2, 1,8,  32,16,1>, 32,16,1, 1};
+    }
+    if (s == 1 && p.up.x == 1 && p.up.y == 1 && p.down.x == 1 && p.down.y == 2) // channels_last
+    {
+        if (fx <= 1  && fy <= 24) spec = {(void*)upfirdn2d_kernel_small<T, 1,1, 1,2, 1,24, 1,64,8>, 1,64,8, 1};
+        if (fx <= 1  && fy <= 20) spec = {(void*)upfirdn2d_kernel_small<T, 1,1, 1,2, 1,20, 1,64,8>, 1,64,8, 1};
+        if (fx <= 1  && fy <= 16) spec = {(void*)upfirdn2d_kernel_small<T, 1,1, 1,2, 1,16, 1,64,8>, 1,64,8, 1};
+        if (fx <= 1  && fy <= 12) spec = {(void*)upfirdn2d_kernel_small<T, 1,1, 1,2, 1,12, 1,64,8>, 1,64,8, 1};
+        if (fx <= 1  && fy <= 8 ) spec = {(void*)upfirdn2d_kernel_small<T, 1,1, 1,2, 1,8,  1,64,8>, 1,64,8, 1};
+    }
+    return spec;
+}
+
+//------------------------------------------------------------------------
+// Template specializations.
+
+template upfirdn2d_kernel_spec choose_upfirdn2d_kernel<double>         (const upfirdn2d_kernel_params& p);
+template upfirdn2d_kernel_spec choose_upfirdn2d_kernel<float>          (const upfirdn2d_kernel_params& p);
+template upfirdn2d_kernel_spec choose_upfirdn2d_kernel<c10::Half>      (const upfirdn2d_kernel_params& p);
+template upfirdn2d_kernel_spec choose_upfirdn2d_kernel<at::BFloat16>   (const upfirdn2d_kernel_params& p);
+
+//------------------------------------------------------------------------

src/vqvaes/bsqvit/stylegan_utils/ops/upfirdn2d.h

@@ -0,0 +1,59 @@
+// Copyright (c) 2021, NVIDIA CORPORATION.  All rights reserved.
+//
+// NVIDIA CORPORATION and its licensors retain all intellectual property
+// and proprietary rights in and to this software, related documentation
+// and any modifications thereto.  Any use, reproduction, disclosure or
+// distribution of this software and related documentation without an express
+// license agreement from NVIDIA CORPORATION is strictly prohibited.
+
+#include <cuda_runtime.h>
+
+//------------------------------------------------------------------------
+// CUDA kernel parameters.
+
+struct upfirdn2d_kernel_params
+{
+    const void*     x;
+    const float*    f;
+    void*           y;
+
+    int2            up;
+    int2            down;
+    int2            pad0;
+    int             flip;
+    float           gain;
+
+    int4            inSize;         // [width, height, channel, batch]
+    int4            inStride;
+    int2            filterSize;     // [width, height]
+    int2            filterStride;
+    int4            outSize;        // [width, height, channel, batch]
+    int4            outStride;
+    int             sizeMinor;
+    int             sizeMajor;
+
+    int             loopMinor;
+    int             loopMajor;
+    int             loopX;
+    int             launchMinor;
+    int             launchMajor;
+};
+
+//------------------------------------------------------------------------
+// CUDA kernel specialization.
+
+struct upfirdn2d_kernel_spec
+{
+    void*   kernel;
+    int     tileOutW;
+    int     tileOutH;
+    int     loopMinor;
+    int     loopX;
+};
+
+//------------------------------------------------------------------------
+// CUDA kernel selection.
+
+template <class T> upfirdn2d_kernel_spec choose_upfirdn2d_kernel(const upfirdn2d_kernel_params& p);
+
+//------------------------------------------------------------------------

src/vqvaes/bsqvit/stylegan_utils/ops/upfirdn2d.py

@@ -0,0 +1,382 @@
+# Copyright (c) 2021, NVIDIA CORPORATION.  All rights reserved.
+#
+# NVIDIA CORPORATION and its licensors retain all intellectual property
+# and proprietary rights in and to this software, related documentation
+# and any modifications thereto.  Any use, reproduction, disclosure or
+# distribution of this software and related documentation without an express
+# license agreement from NVIDIA CORPORATION is strictly prohibited.
+
+"""Custom PyTorch ops for efficient resampling of 2D images."""
+
+import os
+import warnings
+import numpy as np
+import torch
+import traceback
+
+from .. import custom_ops, misc
+from . import conv2d_gradfix
+
+#----------------------------------------------------------------------------
+
+_inited = False
+_plugin = None
+
+def _init():
+    global _inited, _plugin
+    if not _inited:
+        sources = ['upfirdn2d.cpp', 'upfirdn2d.cu']
+        sources = [os.path.join(os.path.dirname(__file__), s) for s in sources]
+        try:
+            _plugin = custom_ops.get_plugin('upfirdn2d_plugin', sources=sources, extra_cuda_cflags=['--use_fast_math'])
+        except:
+            warnings.warn('Failed to build CUDA kernels for upfirdn2d. Falling back to slow reference implementation. Details:\n\n' + traceback.format_exc())
+    return _plugin is not None
+
+def _parse_scaling(scaling):
+    if isinstance(scaling, int):
+        scaling = [scaling, scaling]
+    assert isinstance(scaling, (list, tuple))
+    assert all(isinstance(x, int) for x in scaling)
+    sx, sy = scaling
+    assert sx >= 1 and sy >= 1
+    return sx, sy
+
+def _parse_padding(padding):
+    if isinstance(padding, int):
+        padding = [padding, padding]
+    assert isinstance(padding, (list, tuple))
+    assert all(isinstance(x, int) for x in padding)
+    if len(padding) == 2:
+        padx, pady = padding
+        padding = [padx, padx, pady, pady]
+    padx0, padx1, pady0, pady1 = padding
+    return padx0, padx1, pady0, pady1
+
+def _get_filter_size(f):
+    if f is None:
+        return 1, 1
+    assert isinstance(f, torch.Tensor) and f.ndim in [1, 2]
+    fw = f.shape[-1]
+    fh = f.shape[0]
+    with misc.suppress_tracer_warnings():
+        fw = int(fw)
+        fh = int(fh)
+    misc.assert_shape(f, [fh, fw][:f.ndim])
+    assert fw >= 1 and fh >= 1
+    return fw, fh
+
+#----------------------------------------------------------------------------
+
+def setup_filter(f, device=torch.device('cpu'), normalize=True, flip_filter=False, gain=1, separable=None):
+    r"""Convenience function to setup 2D FIR filter for `upfirdn2d()`.
+
+    Args:
+        f:           Torch tensor, numpy array, or python list of the shape
+                     `[filter_height, filter_width]` (non-separable),
+                     `[filter_taps]` (separable),
+                     `[]` (impulse), or
+                     `None` (identity).
+        device:      Result device (default: cpu).
+        normalize:   Normalize the filter so that it retains the magnitude
+                     for constant input signal (DC)? (default: True).
+        flip_filter: Flip the filter? (default: False).
+        gain:        Overall scaling factor for signal magnitude (default: 1).
+        separable:   Return a separable filter? (default: select automatically).
+
+    Returns:
+        Float32 tensor of the shape
+        `[filter_height, filter_width]` (non-separable) or
+        `[filter_taps]` (separable).
+    """
+    # Validate.
+    if f is None:
+        f = 1
+    f = torch.as_tensor(f, dtype=torch.float32)
+    assert f.ndim in [0, 1, 2]
+    assert f.numel() > 0
+    if f.ndim == 0:
+        f = f[np.newaxis]
+
+    # Separable?
+    if separable is None:
+        separable = (f.ndim == 1 and f.numel() >= 8)
+    if f.ndim == 1 and not separable:
+        f = f.ger(f)
+    assert f.ndim == (1 if separable else 2)
+
+    # Apply normalize, flip, gain, and device.
+    if normalize:
+        f /= f.sum()
+    if flip_filter:
+        f = f.flip(list(range(f.ndim)))
+    f = f * (gain ** (f.ndim / 2))
+    f = f.to(device=device)
+    return f
+
+#----------------------------------------------------------------------------
+
+def upfirdn2d(x, f, up=1, down=1, padding=0, flip_filter=False, gain=1, impl='cuda'):
+    r"""Pad, upsample, filter, and downsample a batch of 2D images.
+
+    Performs the following sequence of operations for each channel:
+
+    1. Upsample the image by inserting N-1 zeros after each pixel (`up`).
+
+    2. Pad the image with the specified number of zeros on each side (`padding`).
+       Negative padding corresponds to cropping the image.
+
+    3. Convolve the image with the specified 2D FIR filter (`f`), shrinking it
+       so that the footprint of all output pixels lies within the input image.
+
+    4. Downsample the image by keeping every Nth pixel (`down`).
+
+    This sequence of operations bears close resemblance to scipy.signal.upfirdn().
+    The fused op is considerably more efficient than performing the same calculation
+    using standard PyTorch ops. It supports gradients of arbitrary order.
+
+    Args:
+        x:           Float32/float64/float16 input tensor of the shape
+                     `[batch_size, num_channels, in_height, in_width]`.
+        f:           Float32 FIR filter of the shape
+                     `[filter_height, filter_width]` (non-separable),
+                     `[filter_taps]` (separable), or
+                     `None` (identity).
+        up:          Integer upsampling factor. Can be a single int or a list/tuple
+                     `[x, y]` (default: 1).
+        down:        Integer downsampling factor. Can be a single int or a list/tuple
+                     `[x, y]` (default: 1).
+        padding:     Padding with respect to the upsampled image. Can be a single number
+                     or a list/tuple `[x, y]` or `[x_before, x_after, y_before, y_after]`
+                     (default: 0).
+        flip_filter: False = convolution, True = correlation (default: False).
+        gain:        Overall scaling factor for signal magnitude (default: 1).
+        impl:        Implementation to use. Can be `'ref'` or `'cuda'` (default: `'cuda'`).
+
+    Returns:
+        Tensor of the shape `[batch_size, num_channels, out_height, out_width]`.
+    """
+    assert isinstance(x, torch.Tensor)
+    assert impl in ['ref', 'cuda']
+    if impl == 'cuda' and x.device.type == 'cuda' and _init():
+        return _upfirdn2d_cuda(up=up, down=down, padding=padding, flip_filter=flip_filter, gain=gain).apply(x, f)
+    return _upfirdn2d_ref(x, f, up=up, down=down, padding=padding, flip_filter=flip_filter, gain=gain)
+
+#----------------------------------------------------------------------------
+
+def _upfirdn2d_ref(x, f, up=1, down=1, padding=0, flip_filter=False, gain=1):
+    """Slow reference implementation of `upfirdn2d()` using standard PyTorch ops.
+    """
+    # Validate arguments.
+    assert isinstance(x, torch.Tensor) and x.ndim == 4
+    if f is None:
+        f = torch.ones([1, 1], dtype=torch.float32, device=x.device)
+    assert isinstance(f, torch.Tensor) and f.ndim in [1, 2]
+    assert f.dtype == torch.float32 and not f.requires_grad
+    batch_size, num_channels, in_height, in_width = x.shape
+    upx, upy = _parse_scaling(up)
+    downx, downy = _parse_scaling(down)
+    padx0, padx1, pady0, pady1 = _parse_padding(padding)
+
+    # Upsample by inserting zeros.
+    x = x.reshape([batch_size, num_channels, in_height, 1, in_width, 1])
+    x = torch.nn.functional.pad(x, [0, upx - 1, 0, 0, 0, upy - 1])
+    x = x.reshape([batch_size, num_channels, in_height * upy, in_width * upx])
+
+    # Pad or crop.
+    x = torch.nn.functional.pad(x, [max(padx0, 0), max(padx1, 0), max(pady0, 0), max(pady1, 0)])
+    x = x[:, :, max(-pady0, 0) : x.shape[2] - max(-pady1, 0), max(-padx0, 0) : x.shape[3] - max(-padx1, 0)]
+
+    # Setup filter.
+    f = f * (gain ** (f.ndim / 2))
+    f = f.to(x.dtype)
+    if not flip_filter:
+        f = f.flip(list(range(f.ndim)))
+
+    # Convolve with the filter.
+    f = f[np.newaxis, np.newaxis].repeat([num_channels, 1] + [1] * f.ndim)
+    if f.ndim == 4:
+        x = conv2d_gradfix.conv2d(input=x, weight=f, groups=num_channels)
+    else:
+        x = conv2d_gradfix.conv2d(input=x, weight=f.unsqueeze(2), groups=num_channels)
+        x = conv2d_gradfix.conv2d(input=x, weight=f.unsqueeze(3), groups=num_channels)
+
+    # Downsample by throwing away pixels.
+    x = x[:, :, ::downy, ::downx]
+    return x
+
+#----------------------------------------------------------------------------
+
+_upfirdn2d_cuda_cache = dict()
+
+def _upfirdn2d_cuda(up=1, down=1, padding=0, flip_filter=False, gain=1):
+    """Fast CUDA implementation of `upfirdn2d()` using custom ops.
+    """
+    # Parse arguments.
+    upx, upy = _parse_scaling(up)
+    downx, downy = _parse_scaling(down)
+    padx0, padx1, pady0, pady1 = _parse_padding(padding)
+
+    # Lookup from cache.
+    key = (upx, upy, downx, downy, padx0, padx1, pady0, pady1, flip_filter, gain)
+    if key in _upfirdn2d_cuda_cache:
+        return _upfirdn2d_cuda_cache[key]
+
+    # Forward op.
+    class Upfirdn2dCuda(torch.autograd.Function):
+        @staticmethod
+        def forward(ctx, x, f): # pylint: disable=arguments-differ
+            assert isinstance(x, torch.Tensor) and x.ndim == 4
+            if f is None:
+                f = torch.ones([1, 1], dtype=torch.float32, device=x.device)
+            assert isinstance(f, torch.Tensor) and f.ndim in [1, 2]
+            y = x
+            if f.ndim == 2:
+                y = _plugin.upfirdn2d(y, f, upx, upy, downx, downy, padx0, padx1, pady0, pady1, flip_filter, gain)
+            else:
+                y = _plugin.upfirdn2d(y, f.unsqueeze(0), upx, 1, downx, 1, padx0, padx1, 0, 0, flip_filter, np.sqrt(gain))
+                y = _plugin.upfirdn2d(y, f.unsqueeze(1), 1, upy, 1, downy, 0, 0, pady0, pady1, flip_filter, np.sqrt(gain))
+            ctx.save_for_backward(f)
+            ctx.x_shape = x.shape
+            return y
+
+        @staticmethod
+        def backward(ctx, dy): # pylint: disable=arguments-differ
+            f, = ctx.saved_tensors
+            _, _, ih, iw = ctx.x_shape
+            _, _, oh, ow = dy.shape
+            fw, fh = _get_filter_size(f)
+            p = [
+                fw - padx0 - 1,
+                iw * upx - ow * downx + padx0 - upx + 1,
+                fh - pady0 - 1,
+                ih * upy - oh * downy + pady0 - upy + 1,
+            ]
+            dx = None
+            df = None
+
+            if ctx.needs_input_grad[0]:
+                dx = _upfirdn2d_cuda(up=down, down=up, padding=p, flip_filter=(not flip_filter), gain=gain).apply(dy, f)
+
+            assert not ctx.needs_input_grad[1]
+            return dx, df
+
+    # Add to cache.
+    _upfirdn2d_cuda_cache[key] = Upfirdn2dCuda
+    return Upfirdn2dCuda
+
+#----------------------------------------------------------------------------
+
+def filter2d(x, f, padding=0, flip_filter=False, gain=1, impl='cuda'):
+    r"""Filter a batch of 2D images using the given 2D FIR filter.
+
+    By default, the result is padded so that its shape matches the input.
+    User-specified padding is applied on top of that, with negative values
+    indicating cropping. Pixels outside the image are assumed to be zero.
+
+    Args:
+        x:           Float32/float64/float16 input tensor of the shape
+                     `[batch_size, num_channels, in_height, in_width]`.
+        f:           Float32 FIR filter of the shape
+                     `[filter_height, filter_width]` (non-separable),
+                     `[filter_taps]` (separable), or
+                     `None` (identity).
+        padding:     Padding with respect to the output. Can be a single number or a
+                     list/tuple `[x, y]` or `[x_before, x_after, y_before, y_after]`
+                     (default: 0).
+        flip_filter: False = convolution, True = correlation (default: False).
+        gain:        Overall scaling factor for signal magnitude (default: 1).
+        impl:        Implementation to use. Can be `'ref'` or `'cuda'` (default: `'cuda'`).
+
+    Returns:
+        Tensor of the shape `[batch_size, num_channels, out_height, out_width]`.
+    """
+    padx0, padx1, pady0, pady1 = _parse_padding(padding)
+    fw, fh = _get_filter_size(f)
+    p = [
+        padx0 + fw // 2,
+        padx1 + (fw - 1) // 2,
+        pady0 + fh // 2,
+        pady1 + (fh - 1) // 2,
+    ]
+    return upfirdn2d(x, f, padding=p, flip_filter=flip_filter, gain=gain, impl=impl)
+
+#----------------------------------------------------------------------------
+
+def upsample2d(x, f, up=2, padding=0, flip_filter=False, gain=1, impl='cuda'):
+    r"""Upsample a batch of 2D images using the given 2D FIR filter.
+
+    By default, the result is padded so that its shape is a multiple of the input.
+    User-specified padding is applied on top of that, with negative values
+    indicating cropping. Pixels outside the image are assumed to be zero.
+
+    Args:
+        x:           Float32/float64/float16 input tensor of the shape
+                     `[batch_size, num_channels, in_height, in_width]`.
+        f:           Float32 FIR filter of the shape
+                     `[filter_height, filter_width]` (non-separable),
+                     `[filter_taps]` (separable), or
+                     `None` (identity).
+        up:          Integer upsampling factor. Can be a single int or a list/tuple
+                     `[x, y]` (default: 1).
+        padding:     Padding with respect to the output. Can be a single number or a
+                     list/tuple `[x, y]` or `[x_before, x_after, y_before, y_after]`
+                     (default: 0).
+        flip_filter: False = convolution, True = correlation (default: False).
+        gain:        Overall scaling factor for signal magnitude (default: 1).
+        impl:        Implementation to use. Can be `'ref'` or `'cuda'` (default: `'cuda'`).
+
+    Returns:
+        Tensor of the shape `[batch_size, num_channels, out_height, out_width]`.
+    """
+    upx, upy = _parse_scaling(up)
+    padx0, padx1, pady0, pady1 = _parse_padding(padding)
+    fw, fh = _get_filter_size(f)
+    p = [
+        padx0 + (fw + upx - 1) // 2,
+        padx1 + (fw - upx) // 2,
+        pady0 + (fh + upy - 1) // 2,
+        pady1 + (fh - upy) // 2,
+    ]
+    return upfirdn2d(x, f, up=up, padding=p, flip_filter=flip_filter, gain=gain*upx*upy, impl=impl)
+
+#----------------------------------------------------------------------------
+
+def downsample2d(x, f, down=2, padding=0, flip_filter=False, gain=1, impl='cuda'):
+    r"""Downsample a batch of 2D images using the given 2D FIR filter.
+
+    By default, the result is padded so that its shape is a fraction of the input.
+    User-specified padding is applied on top of that, with negative values
+    indicating cropping. Pixels outside the image are assumed to be zero.
+
+    Args:
+        x:           Float32/float64/float16 input tensor of the shape
+                     `[batch_size, num_channels, in_height, in_width]`.
+        f:           Float32 FIR filter of the shape
+                     `[filter_height, filter_width]` (non-separable),
+                     `[filter_taps]` (separable), or
+                     `None` (identity).
+        down:        Integer downsampling factor. Can be a single int or a list/tuple
+                     `[x, y]` (default: 1).
+        padding:     Padding with respect to the input. Can be a single number or a
+                     list/tuple `[x, y]` or `[x_before, x_after, y_before, y_after]`
+                     (default: 0).
+        flip_filter: False = convolution, True = correlation (default: False).
+        gain:        Overall scaling factor for signal magnitude (default: 1).
+        impl:        Implementation to use. Can be `'ref'` or `'cuda'` (default: `'cuda'`).
+
+    Returns:
+        Tensor of the shape `[batch_size, num_channels, out_height, out_width]`.
+    """
+    downx, downy = _parse_scaling(down)
+    padx0, padx1, pady0, pady1 = _parse_padding(padding)
+    fw, fh = _get_filter_size(f)
+    p = [
+        padx0 + (fw - downx + 1) // 2,
+        padx1 + (fw - downx) // 2,
+        pady0 + (fh - downy + 1) // 2,
+        pady1 + (fh - downy) // 2,
+    ]
+    return upfirdn2d(x, f, down=down, padding=p, flip_filter=flip_filter, gain=gain, impl=impl)
+
+#----------------------------------------------------------------------------

src/vqvaes/bsqvit/transformer.py

@@ -0,0 +1,416 @@
+from collections import OrderedDict
+from typing import Callable, Optional, Union
+from einops import rearrange
+import torch
+import torch.nn as nn
+from torch.utils.checkpoint import checkpoint
+from timm.models.layers import to_2tuple
+from timm.models.layers import trunc_normal_
+from timm.models.layers import DropPath
+
+from .attention_mask import get_attention_mask
+
+
+class LayerScale(nn.Module):
+    def __init__(self, dim, init_values=1e-5, inplace=False):
+        super().__init__()
+        self.inplace = inplace
+        self.gamma = nn.Parameter(init_values * torch.ones(dim))
+
+    def forward(self, x):
+        return x.mul_(self.gamma) if self.inplace else x * self.gamma
+
+
+class ResidualAttentionBlock(nn.Module):
+    def __init__(
+            self,
+            d_model: int,
+            n_head: int,
+            mlp_ratio: float = 4.0,
+            ls_init_value: float = None,
+            drop: float = 0.,
+            attn_drop: float = 0.,
+            drop_path: float = 0.,
+            act_layer: Callable = nn.GELU,
+            norm_layer: Callable = nn.LayerNorm,
+            use_preln: bool = True,
+    ):
+        super().__init__()
+
+        self.ln_1 = norm_layer(d_model)
+        self.attn = nn.MultiheadAttention(d_model, n_head, dropout=attn_drop)
+        self.ls_1 = LayerScale(d_model, ls_init_value) if ls_init_value is not None else nn.Identity()
+
+        self.ln_2 = norm_layer(d_model)
+        mlp_width = int(d_model * mlp_ratio)
+        self.mlp = nn.Sequential(OrderedDict([
+            ("c_fc", nn.Linear(d_model, mlp_width)),
+            ("gelu", act_layer()),
+            # disable this following JAX implementation.
+            # Reference: https://github.com/google-research/magvit/blob/main/videogvt/models/simplified_bert.py#L112
+            # ("drop1", nn.Dropout(drop)),
+            ("c_proj", nn.Linear(mlp_width, d_model)),
+            ("drop2", nn.Dropout(drop)),
+        ]))
+        self.ls_2 = LayerScale(d_model, ls_init_value) if ls_init_value is not None else nn.Identity()
+
+        self.drop_path = DropPath(drop_path) if drop_path > 0. else nn.Identity()
+
+        self.use_preln = use_preln
+
+    def attention(self, x: torch.Tensor, attn_mask: Optional[torch.Tensor] = None, is_causal: bool = False):
+        attn_mask = attn_mask.to(x.dtype) if attn_mask is not None else None
+        return self.attn(x, x, x, need_weights=False, attn_mask=attn_mask, is_causal=is_causal)[0]
+
+    def checkpoint_forward(self, x: torch.Tensor, 
+                           attn_mask: Optional[torch.Tensor] = None,
+                           is_causal: bool = False):
+        state = x
+        if self.use_preln:
+            x = checkpoint(self.ln_1, x, use_reentrant=False)
+            x = self.attention(x, attn_mask, is_causal)
+            x = checkpoint(self.ls_1, x, use_reentrant=False)
+            state = state + self.drop_path(x)
+            x = checkpoint(self.ln_2, state, use_reentrant=False)
+            x = self.mlp(x)
+            x = checkpoint(self.ls_2, x, use_reentrant=False)
+            state = state + self.drop_path(x)
+        else:
+            x = self.attention(x, attn_mask, is_causal)
+            x = state + self.drop_path(x)
+            state = checkpoint(self.ln_1, x, use_reentrant=False)
+            x = self.mlp(state)
+            state = state + self.drop_path(x)
+            state = checkpoint(self.ln_2, state, use_reentrant=False)
+        return state
+
+    def forward(self, x: torch.Tensor, 
+                attn_mask: Optional[torch.Tensor] = None, is_causal: bool =False,
+                selective_checkpointing: bool = False):
+        if selective_checkpointing:
+            return self.checkpoint_forward(x, attn_mask, is_causal=is_causal)
+        if self.use_preln:
+            x = x + self.drop_path(self.ls_1(self.attention(self.ln_1(x), attn_mask=attn_mask, is_causal=is_causal)))
+            x = x + self.drop_path(self.ls_2(self.mlp(self.ln_2(x))))
+        else:
+            x = x + self.drop_path(self.attention(x, attn_mask=attn_mask, is_causal=is_causal))
+            x = self.ln_1(x)
+            x = x + self.drop_path(self.mlp(x))
+            x = self.ln_2(x)
+        return x
+
+
+class Transformer(nn.Module):
+    def __init__(self,
+                 width: int,
+                 layers: int,
+                 heads: int,
+                 mlp_ratio: float = 4.0,
+                 ls_init_value: float = None,
+                 drop: float = 0.,
+                 attn_drop: float = 0.,
+                 drop_path: float = 0.,
+                 act_layer: nn.Module = nn.GELU,
+                 norm_layer: nn.Module = nn.LayerNorm,
+                 use_preln: bool = True,
+                 ):
+        super().__init__()
+        self.width = width
+        self.layers = layers
+        self.grad_checkpointing = False
+        self.selective_checkpointing = False
+        self.grad_checkpointing_params = {'use_reentrant': False}
+        if attn_drop == 0 and drop_path == 0 and drop_path == 0:
+            self.grad_checkpointing_params.update({'preserve_rng_state': False})
+        else:
+            self.grad_checkpointing_params.update({'preserve_rng_state': True})
+
+        self.resblocks = nn.ModuleList([
+            ResidualAttentionBlock(
+                width, heads, mlp_ratio, ls_init_value=ls_init_value,
+                drop=drop, attn_drop=attn_drop, drop_path=drop_path,
+                act_layer=act_layer, norm_layer=norm_layer,
+                use_preln=use_preln)
+            for _ in range(layers)
+        ])
+
+    def forward(self, x: torch.Tensor, 
+                attn_mask: Optional[torch.Tensor] = None,
+                is_causal: bool =False):
+        for r in self.resblocks:
+            if self.training and self.grad_checkpointing and not torch.jit.is_scripting():
+                if not self.selective_checkpointing:
+                    x = checkpoint(r, x, attn_mask, is_causal=is_causal, **self.grad_checkpointing_params)
+                else:
+                    x = r(x, attn_mask=attn_mask, is_causal=is_causal, selective_checkpointing=True)
+            else:
+                x = r(x, attn_mask=attn_mask)
+        return x
+
+
+class TransformerEncoder(nn.Module):
+    def __init__(self,
+                 image_size: int,
+                 patch_size: int,
+                 width: int,
+                 layers: int,
+                 heads: int,
+                 mlp_ratio: float,
+                 num_frames: int = 1,
+                 cross_frames: bool = True,
+                 ls_init_value: float = None,
+                 drop_rate: float = 0.,
+                 attn_drop_rate: float = 0.,
+                 drop_path_rate: float = 0.,
+                 ln_pre: bool = True,
+                 ln_post: bool = True,
+                 act_layer: str = 'gelu',
+                 norm_layer: str = 'layer_norm',
+                 mask_type: Union[str, None] = 'none',
+                 mask_block_size: int = -1
+    ):
+        super().__init__()
+        self.image_size = to_2tuple(image_size)
+        self.patch_size = to_2tuple(patch_size)
+        self.grid_size = (self.image_size[0] // self.patch_size[0], self.image_size[1] // self.patch_size[1])
+        self.patches_per_frame = self.grid_size[0] * self.grid_size[1]
+        self.mask_type = mask_type
+        self.mask_block_size = mask_block_size
+
+        if act_layer.lower() == 'gelu':
+            self.act_layer = nn.GELU
+        else:
+            raise ValueError(f"Unsupported activation function: {act_layer}")
+        if norm_layer.lower() == 'layer_norm':
+            self.norm_layer = nn.LayerNorm
+        else:
+            raise ValueError(f"Unsupported normalization: {norm_layer}")
+
+        self.conv1 = nn.Linear(
+            in_features=3 * self.patch_size[0] * self.patch_size[1],
+            out_features=width,
+            bias=not ln_pre
+        )
+
+        scale = width ** -0.5
+        self.positional_embedding = nn.Parameter(scale * torch.randn(self.grid_size[0] * self.grid_size[1], width))
+        assert num_frames >= 1
+        self.num_frames = num_frames
+        self.cross_frames = cross_frames
+        if num_frames > 1 and cross_frames:
+            self.temporal_positional_embedding = nn.Parameter(torch.zeros(num_frames, width))
+        else:
+            self.temporal_positional_embedding = None
+
+        self.ln_pre = self.norm_layer(width) if ln_pre else nn.Identity()
+
+        self.transformer = Transformer(
+            width, layers, heads, mlp_ratio, ls_init_value=ls_init_value,
+            drop=drop_rate, attn_drop=attn_drop_rate, drop_path=drop_path_rate,
+            act_layer=self.act_layer, norm_layer=self.norm_layer,
+        )
+
+        self.ln_post = self.norm_layer(width)
+
+        self.init_parameters()
+
+    def init_parameters(self):
+        if self.positional_embedding is not None:
+            nn.init.normal_(self.positional_embedding, std=0.02)
+        trunc_normal_(self.conv1.weight, std=0.02)
+        for block in self.transformer.resblocks:
+            for n, p in block.named_parameters():
+                if 'weight' in n:
+                    if 'ln' not in n:
+                        trunc_normal_(p, std=0.02)
+                elif 'bias' in n:
+                    nn.init.zeros_(p)
+                else:
+                    raise NotImplementedError(f'Unknown parameters named {n}')
+
+    @torch.jit.ignore
+    def set_grad_checkpointing(self, enable=True, selective=False):
+        self.transformer.grad_checkpointing = enable
+        self.transformer.selective_checkpointing = selective
+        
+
+    def forward(self, x):
+        if self.num_frames == 1:
+            x = rearrange(
+                x, "b c (hh sh) (ww sw) -> b (hh ww) (c sh sw)",
+                sh=self.patch_size[0], sw=self.patch_size[1]
+            )
+            x = self.conv1(x)
+            x = x + self.positional_embedding.to(x.dtype)
+        elif self.cross_frames:
+            num_frames = x.shape[2]
+            assert num_frames <= self.num_frames, 'Number of frames should be less or equal to the model setting'
+            x = rearrange(
+                x, "b c t (hh sh) (ww sw) -> b (t hh ww) (c sh sw)",
+                sh=self.patch_size[0], sw=self.patch_size[1]
+            )
+            x = self.conv1(x)
+            tile_pos_embed = self.positional_embedding.repeat(num_frames, 1)
+            tile_tem_embed = self.temporal_positional_embedding[:num_frames].repeat_interleave(self.patches_per_frame, 0)
+            total_pos_embed = tile_pos_embed + tile_tem_embed
+            x = x + total_pos_embed.to(x.dtype).squeeze(0)
+        else:
+            x = rearrange(
+                x, "b c t (hh sh) (ww sw) -> (b t) (hh ww) (c sh sw)",
+                sh=self.patch_size[0], sw=self.patch_size[1]
+            )
+            x = self.conv1(x)
+            x = x + self.positional_embedding.to(x.dtype)
+
+        x = self.ln_pre(x)
+        x = x.permute(1, 0, 2)
+        block_size = self.grid_size[0] * self.grid_size[1] if self.mask_block_size <= 0 else self.mask_block_size
+        attn_mask = get_attention_mask(x.size(0), x.device, mask_type=self.mask_type, block_size=block_size)
+        x = self.transformer(x, attn_mask, is_causal=self.mask_type == 'causal')
+        x = x.permute(1, 0, 2)
+        x = self.ln_post(x)
+
+        return x
+
+
+class TransformerDecoder(nn.Module):
+    def __init__(self,
+                 image_size: int,
+                 patch_size: int,
+                 width: int,
+                 layers: int,
+                 heads: int,
+                 mlp_ratio: float,
+                 num_frames: int = 1,
+                 cross_frames: bool = True,
+                 ls_init_value: float = None,
+                 drop_rate: float = 0.,
+                 attn_drop_rate: float = 0.,
+                 drop_path_rate: float = 0.,
+                 ln_pre: bool = True,
+                 ln_post: bool = True,
+                 act_layer: str = 'gelu',
+                 norm_layer: str = 'layer_norm',
+                 use_ffn_output: bool = True,
+                 dim_ffn_output: int = 3072,
+                 logit_laplace: bool = False,
+                 mask_type: Union[str, None] = 'none',
+                 mask_block_size: int = -1
+    ):
+        super().__init__()
+        self.image_size = to_2tuple(image_size)
+        self.patch_size = to_2tuple(patch_size)
+        self.grid_size = (self.image_size[0] // self.patch_size[0], self.image_size[1] // self.patch_size[1])
+        self.patches_per_frame = self.grid_size[0] * self.grid_size[1]
+        self.mask_type = mask_type
+        self.mask_block_size = mask_block_size
+
+        if act_layer.lower() == 'gelu':
+            self.act_layer = nn.GELU
+        else:
+            raise ValueError(f"Unsupported activation function: {act_layer}")
+        if norm_layer.lower() == 'layer_norm':
+            self.norm_layer = nn.LayerNorm
+        else:
+            raise ValueError(f"Unsupported normalization: {norm_layer}")
+
+        self.use_ffn_output = use_ffn_output
+        if use_ffn_output:
+            self.ffn = nn.Sequential(
+                nn.Linear(width, dim_ffn_output),
+                nn.Tanh(),
+            )
+            self.conv_out = nn.Linear(
+            in_features=dim_ffn_output,
+            out_features=3 * self.patch_size[0] * self.patch_size[1] * (1 + logit_laplace)
+        )
+        else:
+            self.ffn = nn.Identity()
+            self.conv_out = nn.Linear(
+                in_features=width,
+                out_features=3 * self.patch_size[0] * self.patch_size[1] * (1 + logit_laplace)
+            )
+
+        scale = width ** -0.5
+        self.positional_embedding = nn.Parameter(scale * torch.randn(self.grid_size[0] * self.grid_size[1], width))
+        assert num_frames >= 1
+        self.num_frames = num_frames
+        self.cross_frames = cross_frames
+        if num_frames > 1 and cross_frames:
+            self.temporal_positional_embedding = nn.Parameter(torch.zeros(num_frames, width))
+        else:
+            self.temporal_positional_embedding = None
+
+        self.ln_pre = self.norm_layer(width) if ln_pre else nn.Identity()
+
+        self.transformer = Transformer(
+            width, layers, heads, mlp_ratio, ls_init_value=ls_init_value,
+            drop=drop_rate, attn_drop=attn_drop_rate, drop_path=drop_path_rate,
+            act_layer=self.act_layer, norm_layer=self.norm_layer,
+        )
+
+        self.ln_post = self.norm_layer(width) if ln_post else nn.Identity()
+
+        self.init_parameters()
+
+    def init_parameters(self):
+        if self.positional_embedding is not None:
+            nn.init.normal_(self.positional_embedding, std=0.02)
+
+        for block in self.transformer.resblocks:
+            for n, p in block.named_parameters():
+                if 'weight' in n:
+                    if 'ln' not in n:
+                        trunc_normal_(p, std=0.02)
+                elif 'bias' in n:
+                    nn.init.zeros_(p)
+                else:
+                    raise NotImplementedError(f'Unknown parameters named {n}')
+        if self.use_ffn_output:
+            trunc_normal_(self.ffn[0].weight, std=0.02)
+        trunc_normal_(self.conv_out.weight, std=0.02)
+
+    @torch.jit.ignore
+    def set_grad_checkpointing(self, enable=True, selective=False):
+        self.transformer.grad_checkpointing = enable
+        self.transformer.selective_checkpointing = selective
+
+    def forward(self, x):
+        if self.num_frames == 1 or not self.cross_frames:
+            x = x + self.positional_embedding.to(x.dtype)
+        else:
+            num_frames = x.shape[1] // self.patches_per_frame
+            assert num_frames <= self.num_frames, 'Number of frames should be less or equal to the model setting'
+            tile_pos_embed = self.positional_embedding.repeat(num_frames, 1)
+            tile_tem_embed = self.temporal_positional_embedding[:num_frames].repeat_interleave(self.patches_per_frame, 0)
+            total_pos_embed = tile_pos_embed + tile_tem_embed
+            x = x + total_pos_embed.to(x.dtype).squeeze(0)
+        x = self.ln_pre(x)
+        x = x.permute(1, 0, 2)
+        block_size = self.grid_size[0] * self.grid_size[1] if self.mask_block_size <= 0 else self.mask_block_size
+        attn_mask = get_attention_mask(x.size(0), x.device, mask_type=self.mask_type, block_size=block_size)
+        x = self.transformer(x, attn_mask, is_causal=self.mask_type == 'causal')
+        x = x.permute(1, 0, 2)
+        x = self.ln_post(x)
+        x = self.ffn(x)
+        x = self.conv_out(x)
+        if self.num_frames == 1:
+            x = rearrange(
+                x, "b (hh ww) (c sh sw) -> b c (hh sh) (ww sw)",
+                hh = self.grid_size[0], ww=self.grid_size[1],
+                sh=self.patch_size[0], sw=self.patch_size[1]
+            )
+        elif self.cross_frames:
+            x = rearrange(
+                x, "b (t hh ww) (c sh sw) -> b c t (hh sh) (ww sw)",
+                t = num_frames, hh = self.grid_size[0], ww=self.grid_size[1],
+                sh=self.patch_size[0], sw=self.patch_size[1]
+            )
+        else:
+            x = rearrange(
+                x, "(b t) (hh ww) (c sh sw) -> b c t (hh sh) (ww sw)",
+                t = num_frames, hh = self.grid_size[0], ww=self.grid_size[1],
+                sh=self.patch_size[0], sw=self.patch_size[1]
+            )
+
+        return x

src/vqvaes/flowmo/flowmo.py

@@ -0,0 +1,945 @@
+"""Model code for FlowMo.
+
+Sources: https://github.com/feizc/FluxMusic/blob/main/train.py
+https://github.com/black-forest-labs/flux/tree/main/src/flux
+"""
+
+import ast
+import itertools
+import math
+from dataclasses import dataclass
+from typing import List, Tuple
+
+import einops
+import torch
+from einops import rearrange, repeat
+from mup import MuReadout
+from torch import Tensor, nn
+import argparse
+import contextlib
+import copy
+import glob
+import os
+import subprocess
+import tempfile
+import time
+
+import fsspec
+import psutil
+import torch
+import torch.distributed as dist
+from mup import MuReadout, set_base_shapes
+from omegaconf import OmegaConf
+from torch.utils.data import DataLoader
+
+from .lookup_free_quantize import LFQ
+
+MUP_ENABLED = True
+
+
+def attention(q: Tensor, k: Tensor, v: Tensor, pe: Tensor) -> Tensor:
+    b, h, l, d = q.shape
+    q, k = apply_rope(q, k, pe)
+
+    if torch.__version__ == "2.0.1+cu117":  # tmp workaround
+        if d != 64:
+            print("MUP is broken in this setting! Be careful!")
+            x = torch.nn.functional.scaled_dot_product_attention(
+                q,
+                k,
+                v,
+            )
+    else:
+        x = torch.nn.functional.scaled_dot_product_attention(
+            q,
+            k,
+            v,
+            scale=8.0 / d if MUP_ENABLED else None,
+        )
+    assert x.shape == q.shape
+    x = rearrange(x, "B H L D -> B L (H D)")
+    return x
+
+
+def rope(pos: Tensor, dim: int, theta: int) -> Tensor:
+    assert dim % 2 == 0
+    scale = torch.arange(0, dim, 2, dtype=torch.float64, device=pos.device) / dim
+    omega = 1.0 / (theta**scale)
+    out = torch.einsum("...n,d->...nd", pos, omega)
+    out = torch.stack(
+        [torch.cos(out), -torch.sin(out), torch.sin(out), torch.cos(out)],
+        dim=-1,
+    )
+    out = rearrange(out, "b n d (i j) -> b n d i j", i=2, j=2)
+    return out.float()
+
+
+def apply_rope(xq: Tensor, xk: Tensor, freqs_cis: Tensor):
+    xq_ = xq.float().reshape(*xq.shape[:-1], -1, 1, 2)
+    xk_ = xk.float().reshape(*xk.shape[:-1], -1, 1, 2)
+    xq_out = freqs_cis[..., 0] * xq_[..., 0] + freqs_cis[..., 1] * xq_[..., 1]
+    xk_out = freqs_cis[..., 0] * xk_[..., 0] + freqs_cis[..., 1] * xk_[..., 1]
+
+    return xq_out.reshape(*xq.shape).type_as(xq), xk_out.reshape(*xk.shape).type_as(xk)
+
+
+def _get_diagonal_gaussian(parameters):
+    mean, logvar = torch.chunk(parameters, 2, dim=1)
+    logvar = torch.clamp(logvar, -30.0, 20.0)
+    return mean, logvar
+
+
+def _sample_diagonal_gaussian(mean, logvar):
+    std = torch.exp(0.5 * logvar)
+    x = mean + std * torch.randn(mean.shape, device=mean.device)
+    return x
+
+
+def _kl_diagonal_gaussian(mean, logvar):
+    var = torch.exp(logvar)
+    return 0.5 * torch.sum(torch.pow(mean, 2) + var - 1.0 - logvar, dim=1).mean()
+
+
+class EmbedND(nn.Module):
+    def __init__(self, dim: int, theta: int, axes_dim):
+        super().__init__()
+        self.dim = dim
+        self.theta = theta
+        self.axes_dim = axes_dim
+
+    def forward(self, ids: Tensor) -> Tensor:
+        n_axes = ids.shape[-1]
+        emb = torch.cat(
+            [rope(ids[..., i], self.axes_dim[i], self.theta) for i in range(n_axes)],
+            dim=-3,
+        )
+
+        return emb.unsqueeze(1)
+
+
+def timestep_embedding(t: Tensor, dim, max_period=10000, time_factor: float = 1000.0):
+    """
+    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.
+    """
+    t = time_factor * t
+    half = dim // 2
+    freqs = torch.exp(
+        -math.log(max_period)
+        * torch.arange(start=0, end=half, dtype=torch.float32)
+        / half
+    ).to(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)
+    if torch.is_floating_point(t):
+        embedding = embedding.to(t)
+    return embedding
+
+
+class MLPEmbedder(nn.Module):
+    def __init__(self, in_dim: int, hidden_dim: int):
+        super().__init__()
+        self.in_layer = nn.Linear(in_dim, hidden_dim, bias=True)
+        self.silu = nn.SiLU()
+        self.out_layer = nn.Linear(hidden_dim, hidden_dim, bias=True)
+
+    def forward(self, x: Tensor) -> Tensor:
+        return self.out_layer(self.silu(self.in_layer(x)))
+
+
+class RMSNorm(torch.nn.Module):
+    def __init__(self, dim: int):
+        super().__init__()
+        self.scale = nn.Parameter(torch.ones(dim))
+
+    def forward(self, x: Tensor):
+        x_dtype = x.dtype
+        x = x.float()
+        rrms = torch.rsqrt(torch.mean(x**2, dim=-1, keepdim=True) + 1e-6)
+        return (x * rrms).to(dtype=x_dtype) * self.scale
+
+
+class QKNorm(torch.nn.Module):
+    def __init__(self, dim: int):
+        super().__init__()
+        self.query_norm = RMSNorm(dim)
+        self.key_norm = RMSNorm(dim)
+
+    def forward(self, q: Tensor, k: Tensor, v: Tensor):
+        q = self.query_norm(q)
+        k = self.key_norm(k)
+        return q.to(v), k.to(v)
+
+
+class SelfAttention(nn.Module):
+    def __init__(self, dim: int, num_heads: int = 8, qkv_bias: bool = False):
+        super().__init__()
+        self.num_heads = num_heads
+        head_dim = dim // num_heads
+
+        self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias)
+        self.norm = QKNorm(head_dim)
+        self.proj = nn.Linear(dim, dim)
+
+    def forward(self, x: Tensor, pe: Tensor) -> Tensor:
+        qkv = self.qkv(x)
+        q, k, v = rearrange(qkv, "B L (K H D) -> K B H L D", K=3, H=self.num_heads)
+        q, k = self.norm(q, k, v)
+        x = attention(q, k, v, pe=pe)
+        x = self.proj(x)
+        return x
+
+
+@dataclass
+class ModulationOut:
+    shift: Tensor
+    scale: Tensor
+    gate: Tensor
+
+
+class Modulation(nn.Module):
+    def __init__(self, dim: int, double: bool):
+        super().__init__()
+        self.is_double = double
+        self.multiplier = 6 if double else 3
+        self.lin = nn.Linear(dim, self.multiplier * dim, bias=True)
+
+        self.lin.weight[dim * 2 : dim * 3].data[:] = 0.0
+        self.lin.bias[dim * 2 : dim * 3].data[:] = 0.0
+        self.lin.weight[dim * 5 : dim * 6].data[:] = 0.0
+        self.lin.bias[dim * 5 : dim * 6].data[:] = 0.0
+
+    def forward(self, vec: Tensor) -> Tuple[ModulationOut, ModulationOut]:
+        out = self.lin(nn.functional.silu(vec))[:, None, :].chunk(
+            self.multiplier, dim=-1
+        )
+        return (
+            ModulationOut(*out[:3]),
+            ModulationOut(*out[3:]) if self.is_double else None,
+        )
+
+
+class DoubleStreamBlock(nn.Module):
+    def __init__(
+        self,
+        hidden_size: int,
+        num_heads: int,
+        mlp_ratio: float,
+        qkv_bias: bool = False,
+    ):
+        super().__init__()
+
+        mlp_hidden_dim = int(hidden_size * mlp_ratio)
+        self.num_heads = num_heads
+        self.hidden_size = hidden_size
+
+        self.img_mod = Modulation(hidden_size, double=True)
+        self.img_norm1 = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6)
+        self.img_attn = SelfAttention(
+            dim=hidden_size, num_heads=num_heads, qkv_bias=qkv_bias
+        )
+        self.img_norm2 = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6)
+        self.img_mlp = nn.Sequential(
+            nn.Linear(hidden_size, mlp_hidden_dim, bias=True),
+            nn.GELU(approximate="tanh"),
+            nn.Linear(mlp_hidden_dim, hidden_size, bias=True),
+        )
+
+        self.txt_mod = Modulation(hidden_size, double=True)
+        self.txt_norm1 = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6)
+        self.txt_attn = SelfAttention(
+            dim=hidden_size, num_heads=num_heads, qkv_bias=qkv_bias
+        )
+
+        self.txt_norm2 = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6)
+        self.txt_mlp = nn.Sequential(
+            nn.Linear(hidden_size, mlp_hidden_dim, bias=True),
+            nn.GELU(approximate="tanh"),
+            nn.Linear(mlp_hidden_dim, hidden_size, bias=True),
+        )
+
+    def forward(self, img: Tensor, txt: Tensor, vec: Tensor, pe: Tensor):
+        pe_single, pe_double = pe
+        p = 1
+        if vec is None:
+            img_mod1, img_mod2 = ModulationOut(0, 1 - p, 1), ModulationOut(0, 1 - p, 1)
+            txt_mod1, txt_mod2 = ModulationOut(0, 1 - p, 1), ModulationOut(0, 1 - p, 1)
+        else:
+            img_mod1, img_mod2 = self.img_mod(vec)
+            txt_mod1, txt_mod2 = self.txt_mod(vec)
+
+        # prepare image for attention
+        img_modulated = self.img_norm1(img)
+        img_modulated = (p + img_mod1.scale) * img_modulated + img_mod1.shift
+        img_qkv = self.img_attn.qkv(img_modulated)
+        img_q, img_k, img_v = rearrange(
+            img_qkv, "B L (K H D) -> K B H L D", K=3, H=self.num_heads
+        )
+        img_q, img_k = self.img_attn.norm(img_q, img_k, img_v)
+
+        # prepare txt for attention
+        txt_modulated = self.txt_norm1(txt)
+        txt_modulated = (p + txt_mod1.scale) * txt_modulated + txt_mod1.shift
+        txt_qkv = self.txt_attn.qkv(txt_modulated)
+        txt_q, txt_k, txt_v = rearrange(
+            txt_qkv, "B L (K H D) -> K B H L D", K=3, H=self.num_heads
+        )
+        txt_q, txt_k = self.txt_attn.norm(txt_q, txt_k, txt_v)
+
+        # run actual attention
+        q = torch.cat((txt_q, img_q), dim=2)
+        k = torch.cat((txt_k, img_k), dim=2)
+        v = torch.cat((txt_v, img_v), dim=2)
+
+        attn = attention(q, k, v, pe=pe_double)
+        txt_attn, img_attn = attn[:, : txt.shape[1]], attn[:, txt.shape[1] :]
+
+        # calculate the img bloks
+        img = img + img_mod1.gate * self.img_attn.proj(img_attn)
+        img = img + img_mod2.gate * self.img_mlp(
+            (p + img_mod2.scale) * self.img_norm2(img) + img_mod2.shift
+        )
+
+        # calculate the txt bloks
+        txt = txt + txt_mod1.gate * self.txt_attn.proj(txt_attn)
+        txt = txt + txt_mod2.gate * self.txt_mlp(
+            (p + txt_mod2.scale) * self.txt_norm2(txt) + txt_mod2.shift
+        )
+        return img, txt
+
+
+class LastLayer(nn.Module):
+    def __init__(
+        self,
+        hidden_size: int,
+        patch_size: int,
+        out_channels: int,
+        readout_zero_init=False,
+    ):
+        super().__init__()
+        self.norm_final = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6)
+
+        if MUP_ENABLED:
+            self.linear = MuReadout(
+                hidden_size,
+                patch_size * patch_size * out_channels,
+                bias=True,
+                readout_zero_init=readout_zero_init,
+            )
+        else:
+            self.linear = nn.Linear(
+                hidden_size, patch_size * patch_size * out_channels, bias=True
+            )
+
+        self.adaLN_modulation = nn.Sequential(
+            nn.SiLU(), nn.Linear(hidden_size, 2 * hidden_size, bias=True)
+        )
+
+    def forward(self, x: Tensor, vec) -> Tensor:
+        if vec is None:
+            pass
+        else:
+            shift, scale = self.adaLN_modulation(vec).chunk(2, dim=1)
+            x = (1 + scale[:, None, :]) * self.norm_final(x) + shift[:, None, :]
+        x = self.norm_final(x)
+        x = self.linear(x)
+        return x
+
+
+@dataclass
+class FluxParams:
+    in_channels: int
+    patch_size: int
+    context_dim: int
+    hidden_size: int
+    mlp_ratio: float
+    num_heads: int
+    depth: int
+    axes_dim: List[int]
+    theta: int
+    qkv_bias: bool
+
+
+DIT_ZOO = dict(
+    dit_xl_4=dict(
+        hidden_size=1152,
+        mlp_ratio=4.0,
+        num_heads=16,
+        axes_dim=[8, 28, 28],
+        theta=10_000,
+        qkv_bias=True,
+    ),
+    dit_l_4=dict(
+        hidden_size=1024,
+        mlp_ratio=4.0,
+        num_heads=16,
+        axes_dim=[8, 28, 28],
+        theta=10_000,
+        qkv_bias=True,
+    ),
+    dit_b_4=dict(
+        hidden_size=768,
+        mlp_ratio=4.0,
+        num_heads=12,
+        axes_dim=[8, 28, 28],
+        theta=10_000,
+        qkv_bias=True,
+    ),
+    dit_s_4=dict(
+        hidden_size=384,
+        mlp_ratio=4.0,
+        num_heads=6,
+        axes_dim=[8, 28, 28],
+        theta=10_000,
+        qkv_bias=True,
+    ),
+    dit_mup_test=dict(
+        hidden_size=768,
+        mlp_ratio=4.0,
+        num_heads=12,
+        axes_dim=[8, 28, 28],
+        theta=10_000,
+        qkv_bias=True,
+    ),
+)
+
+
+def prepare_idxs(img, code_length, patch_size):
+    bs, c, h, w = img.shape
+
+    img_ids = torch.zeros(h // patch_size, w // patch_size, 3, device=img.device)
+    img_ids[..., 1] = (
+        img_ids[..., 1] + torch.arange(h // patch_size, device=img.device)[:, None]
+    )
+    img_ids[..., 2] = (
+        img_ids[..., 2] + torch.arange(w // patch_size, device=img.device)[None, :]
+    )
+    img_ids = repeat(img_ids, "h w c -> b (h w) c", b=bs)
+
+    txt_ids = (
+        torch.zeros((bs, code_length, 3), device=img.device)
+        + torch.arange(code_length, device=img.device)[None, :, None]
+    )
+    return img_ids, txt_ids
+
+
+class Flux(nn.Module):
+    """
+    Transformer model for flow matching on sequences.
+    """
+
+    def __init__(self, params: FluxParams, name="", lsg=False):
+        super().__init__()
+
+        self.name = name
+        self.lsg = lsg
+        self.params = params
+        self.in_channels = params.in_channels
+        self.patch_size = params.patch_size
+        self.out_channels = self.in_channels
+        if params.hidden_size % params.num_heads != 0:
+            raise ValueError(
+                f"Hidden size {params.hidden_size} must be divisible by num_heads {params.num_heads}"
+            )
+        pe_dim = params.hidden_size // params.num_heads
+        if sum(params.axes_dim) != pe_dim:
+            raise ValueError(
+                f"Got {params.axes_dim} but expected positional dim {pe_dim}"
+            )
+        self.hidden_size = params.hidden_size
+        self.num_heads = params.num_heads
+        self.pe_embedder = EmbedND(
+            dim=pe_dim, theta=params.theta, axes_dim=params.axes_dim
+        )
+
+        self.time_in = MLPEmbedder(in_dim=256, hidden_dim=self.hidden_size)
+        self.img_in = nn.Linear(self.in_channels, self.hidden_size, bias=True)
+        self.txt_in = nn.Linear(params.context_dim, self.hidden_size)
+
+        self.double_blocks = nn.ModuleList(
+            [
+                DoubleStreamBlock(
+                    self.hidden_size,
+                    self.num_heads,
+                    mlp_ratio=params.mlp_ratio,
+                    qkv_bias=params.qkv_bias,
+                )
+                for idx in range(params.depth)
+            ]
+        )
+
+        self.final_layer_img = LastLayer(
+            self.hidden_size, 1, self.out_channels, readout_zero_init=False
+        )
+        self.final_layer_txt = LastLayer(
+            self.hidden_size, 1, params.context_dim, readout_zero_init=False
+        )
+
+    def forward(
+        self,
+        img: Tensor,
+        img_ids: Tensor,
+        txt: Tensor,
+        txt_ids: Tensor,
+        timesteps: Tensor,
+    ) -> Tensor:
+        b, c, h, w = img.shape
+
+        img = rearrange(
+            img,
+            "b c (gh ph) (gw pw) -> b (gh gw) (ph pw c)",
+            ph=self.patch_size,
+            pw=self.patch_size,
+        )
+        if img.ndim != 3 or txt.ndim != 3:
+            raise ValueError("Input img and txt tensors must have 3 dimensions.")
+        img = self.img_in(img)
+
+        if timesteps is None:
+            vec = None
+        else:
+            vec = self.time_in(timestep_embedding(timesteps, 256))
+
+        txt = self.txt_in(txt)
+        pe_single = self.pe_embedder(torch.cat((txt_ids,), dim=1))
+        pe_double = self.pe_embedder(torch.cat((txt_ids, img_ids), dim=1))
+
+        for block in self.double_blocks:
+            img, txt = block(img=img, txt=txt, pe=(pe_single, pe_double), vec=vec)
+
+        img = self.final_layer_img(img, vec=vec)
+        img = rearrange(
+            img,
+            "b (gh gw) (ph pw c) -> b c (gh ph) (gw pw)",
+            ph=self.patch_size,
+            pw=self.patch_size,
+            gh=h // self.patch_size,
+            gw=w // self.patch_size,
+        )
+
+        txt = self.final_layer_txt(txt, vec=vec)
+        return img, txt, {"final_txt": txt}
+
+
+def get_weights_to_fix(model):
+    with torch.no_grad():
+        for name, module in itertools.chain(model.named_modules()):
+            if "double_blocks" in name and isinstance(module, torch.nn.Linear):
+                yield name, module.weight
+
+
+class FlowMo(nn.Module):
+    def __init__(self, width, config):
+        super().__init__()
+        code_length = config.model.code_length
+        context_dim = config.model.context_dim
+        enc_depth = config.model.enc_depth
+        dec_depth = config.model.dec_depth
+
+        patch_size = config.model.patch_size
+        self.config = config
+
+        self.image_size = config.data.image_size
+        self.patch_size = config.model.patch_size
+        self.code_length = code_length
+        self.dit_mode = "dit_b_4"
+        self.context_dim = context_dim
+        self.encoder_context_dim = context_dim * (
+            1 + (self.config.model.quantization_type == "kl")
+        )
+
+        if config.model.quantization_type == "lfq":
+            self.quantizer = LFQ(
+                codebook_size=2**self.config.model.codebook_size_for_entropy,
+                dim=self.config.model.codebook_size_for_entropy,
+                num_codebooks=1,
+                token_factorization=False,
+            )
+
+        if self.config.model.enc_mup_width is not None:
+            enc_width = self.config.model.enc_mup_width
+        else:
+            enc_width = width
+
+        encoder_params = FluxParams(
+            in_channels=3 * patch_size**2,
+            context_dim=self.encoder_context_dim,
+            patch_size=patch_size,
+            depth=enc_depth,
+            **DIT_ZOO[self.dit_mode],
+        )
+        decoder_params = FluxParams(
+            in_channels=3 * patch_size**2,
+            context_dim=context_dim + 1,
+            patch_size=patch_size,
+            depth=dec_depth,
+            **DIT_ZOO[self.dit_mode],
+        )
+
+        # width=4, dit_b_4 is the usual model
+        encoder_params.hidden_size = enc_width * (encoder_params.hidden_size // 4)
+        decoder_params.hidden_size = width * (decoder_params.hidden_size // 4)
+        encoder_params.axes_dim = [
+            (d // 4) * enc_width for d in encoder_params.axes_dim
+        ]
+        decoder_params.axes_dim = [(d // 4) * width for d in decoder_params.axes_dim]
+
+        self.encoder = Flux(encoder_params, name="encoder")
+        self.decoder = Flux(decoder_params, name="decoder")
+
+    @torch.compile
+    def encode(self, img):
+        b, c, h, w = img.shape
+
+        img_idxs, txt_idxs = prepare_idxs(img, self.code_length, self.patch_size)
+        txt = torch.zeros(
+            (b, self.code_length, self.encoder_context_dim), device=img.device
+        )
+
+        _, code, aux = self.encoder(img, img_idxs, txt, txt_idxs, timesteps=None)
+
+        return code, aux
+
+    def _decode(self, img, code, timesteps):
+        b, c, h, w = img.shape
+
+        img_idxs, txt_idxs = prepare_idxs(
+            img,
+            self.code_length,
+            self.patch_size,
+        )
+        pred, _, decode_aux = self.decoder(
+            img, img_idxs, code, txt_idxs, timesteps=timesteps
+        )
+        return pred, decode_aux
+
+    @torch.compile
+    def decode(self, *args, **kwargs):
+        return self._decode(*args, **kwargs)
+
+    @torch.compile
+    def decode_checkpointed(self, *args, **kwargs):
+        # Need to compile(checkpoint), not checkpoint(compile)
+        assert not kwargs, kwargs
+        return torch.utils.checkpoint.checkpoint(
+            self._decode,
+            *args,
+            # WARNING: Do not use_reentrant=True with compile, it will silently
+            # produce incorrect gradients!
+            use_reentrant=False,
+        )
+
+    @torch.compile
+    def _quantize(self, code):
+        """
+        Args:
+            code: [b codelength context dim]
+
+        Returns:
+            quantized code of the same shape
+        """
+        b, t, f = code.shape
+        indices = None
+        if self.config.model.quantization_type == "noop":
+            quantized = code
+            quantizer_loss = torch.tensor(0.0).to(code.device)
+        elif self.config.model.quantization_type == "kl":
+            # colocating features of same token before split is maybe slightly
+            # better?
+            mean, logvar = _get_diagonal_gaussian(
+                einops.rearrange(code, "b t f -> b (f t)")
+            )
+            code = einops.rearrange(
+                _sample_diagonal_gaussian(mean, logvar),
+                "b (f t) -> b t f",
+                f=f // 2,
+                t=t,
+            )
+            quantizer_loss = _kl_diagonal_gaussian(mean, logvar)
+        elif self.config.model.quantization_type == "lfq":
+            assert f % self.config.model.codebook_size_for_entropy == 0, f
+            code = einops.rearrange(
+                code,
+                "b t (fg fh) -> b fg (t fh)",
+                fg=self.config.model.codebook_size_for_entropy,
+            )
+
+            (quantized, entropy_aux_loss, indices), breakdown = self.quantizer(
+                code, return_loss_breakdown=True
+            )
+            assert quantized.shape == code.shape
+            quantized = einops.rearrange(quantized, "b fg (t fh) -> b t (fg fh)", t=t)
+
+            quantizer_loss = (
+                entropy_aux_loss * self.config.model.entropy_loss_weight
+                + breakdown.commitment * self.config.model.commit_loss_weight
+            )
+            code = quantized
+        else:
+            raise NotImplementedError
+        return code, indices, quantizer_loss
+
+    #     def forward(
+    #         self,
+    #         img,
+    #         noised_img,
+    #         timesteps,
+    #         enable_cfg=True,
+    #     ):
+    #         aux = {}
+    #
+    #         code, encode_aux = self.encode(img)
+    #
+    #         aux["original_code"] = code
+    #
+    #         b, t, f = code.shape
+    #
+    #         code, _, aux["quantizer_loss"] = self._quantize(code)
+    #
+    #         mask = torch.ones_like(code[..., :1])
+    #         code = torch.concatenate([code, mask], axis=-1)
+    #         code_pre_cfg = code
+    #
+    #         if self.config.model.enable_cfg and enable_cfg:
+    #             cfg_mask = (torch.rand((b,), device=code.device) > 0.1)[:, None, None]
+    #             code = code * cfg_mask
+    #
+    #         v_est, decode_aux = self.decode(noised_img, code, timesteps)
+    #         aux.update(decode_aux)
+    #
+    #         if self.config.model.posttrain_sample:
+    #             aux["posttrain_sample"] = self.reconstruct_checkpoint(code_pre_cfg)
+    #
+    #         return v_est, aux
+
+    def forward(self, img):
+        return self.reconstruct(img)
+
+    def reconstruct_checkpoint(self, code):
+        with torch.autocast(
+            "cuda",
+            dtype=torch.bfloat16,
+        ):
+            bs, *_ = code.shape
+
+            z = torch.randn((bs, 3, self.image_size, self.image_size)).cuda()
+            ts = (
+                torch.rand((bs, self.config.model.posttrain_sample_k + 1))
+                .cumsum(dim=1)
+                .cuda()
+            )
+            ts = ts - ts[:, :1]
+            ts = (ts / ts[:, -1:]).flip(dims=(1,))
+            dts = ts[:, :-1] - ts[:, 1:]
+
+            for i, (t, dt) in enumerate((zip(ts.T, dts.T))):
+                if self.config.model.posttrain_sample_enable_cfg:
+                    mask = (torch.rand((bs,), device=code.device) > 0.1)[
+                        :, None, None
+                    ].to(code.dtype)
+                    code_t = code * mask
+                else:
+                    code_t = code
+
+                vc, _ = self.decode_checkpointed(z, code_t, t)
+
+                z = z - dt[:, None, None, None] * vc
+        return z
+
+    @torch.no_grad()
+    def reconstruct(self, images, dtype=torch.bfloat16, code=None):
+        """
+        Args:
+            images in [bchw] [-1, 1]
+
+        Returns:
+            images in [bchw] [-1, 1]
+        """
+        model = self
+        config = self.config.eval.sampling
+
+        with torch.autocast(
+            "cuda",
+            dtype=dtype,
+        ):
+            bs, c, h, w = images.shape
+            if code is None:
+                x = images.cuda()
+                prequantized_code = model.encode(x)[0].cuda()
+                code, indices, _ = model._quantize(prequantized_code)
+
+            z = torch.randn((bs, 3, h, w)).cuda()
+
+            mask = torch.ones_like(code[..., :1])
+            code = torch.concatenate([code * mask, mask], axis=-1)
+
+            cfg_mask = 0.0
+            null_code = code * cfg_mask if config.cfg != 1.0 else None
+
+            samples = rf_sample(
+                model,
+                z,
+                code,
+                null_code=null_code,
+                sample_steps=config.sample_steps,
+                cfg=config.cfg,
+                schedule=config.schedule,
+            )[-1].clip(-1, 1)
+        return samples.to(torch.float32), code, prequantized_code
+
+
+def rf_loss(config, model, batch, aux_state):
+    x = batch["image"]
+    b = x.size(0)
+
+    if config.opt.schedule == "lognormal":
+        nt = torch.randn((b,)).to(x.device)
+        t = torch.sigmoid(nt)
+    elif config.opt.schedule == "fat_lognormal":
+        nt = torch.randn((b,)).to(x.device)
+        t = torch.sigmoid(nt)
+        t = torch.where(torch.rand_like(t) <= 0.9, t, torch.rand_like(t))
+    elif config.opt.schedule == "uniform":
+        t = torch.rand((b,), device=x.device)
+    elif config.opt.schedule.startswith("debug"):
+        p = float(config.opt.schedule.split("_")[1])
+        t = torch.ones((b,), device=x.device) * p
+    else:
+        raise NotImplementedError
+
+    t = t.view([b, *([1] * len(x.shape[1:]))])
+    z1 = torch.randn_like(x)
+    zt = (1 - t) * x + t * z1
+
+    zt, t = zt.to(x.dtype), t.to(x.dtype)
+
+    vtheta, aux = model(
+        img=x,
+        noised_img=zt,
+        timesteps=t.reshape((b,)),
+    )
+
+    diff = z1 - vtheta - x
+    x_pred = zt - vtheta * t
+
+    loss = ((diff) ** 2).mean(dim=list(range(1, len(x.shape))))
+    loss = loss.mean()
+
+    aux["loss_dict"] = {}
+    aux["loss_dict"]["diffusion_loss"] = loss
+    aux["loss_dict"]["quantizer_loss"] = aux["quantizer_loss"]
+
+    if config.opt.lpips_weight != 0.0:
+        aux_loss = 0.0
+        if config.model.posttrain_sample:
+            x_pred = aux["posttrain_sample"]
+
+        lpips_dist = aux_state["lpips_model"](x, x_pred)
+        lpips_dist = (config.opt.lpips_weight * lpips_dist).mean() + aux_loss
+        aux["loss_dict"]["lpips_loss"] = lpips_dist
+    else:
+        lpips_dist = 0.0
+
+    loss = loss + aux["quantizer_loss"] + lpips_dist
+    aux["loss_dict"]["total_loss"] = loss
+    return loss, aux
+
+
+def _edm_to_flow_convention(noise_level):
+    # z = x + \sigma z'
+    return noise_level / (1 + noise_level)
+
+
+def rf_sample(
+    model,
+    z,
+    code,
+    null_code=None,
+    sample_steps=25,
+    cfg=2.0,
+    schedule="linear",
+):
+    b = z.size(0)
+    if schedule == "linear":
+        ts = torch.arange(1, sample_steps + 1).flip(0) / sample_steps
+        dts = torch.ones_like(ts) * (1.0 / sample_steps)
+    elif schedule.startswith("pow"):
+        p = float(schedule.split("_")[1])
+        ts = torch.arange(0, sample_steps + 1).flip(0) ** (1 / p) / sample_steps ** (
+            1 / p
+        )
+        dts = ts[:-1] - ts[1:]
+    else:
+        raise NotImplementedError
+
+    if model.config.eval.sampling.cfg_interval is None:
+        interval = None
+    else:
+        cfg_lo, cfg_hi = ast.literal_eval(model.config.eval.sampling.cfg_interval)
+        interval = _edm_to_flow_convention(cfg_lo), _edm_to_flow_convention(cfg_hi)
+
+    images = []
+    for i, (t, dt) in enumerate((zip(ts, dts))):
+        timesteps = torch.tensor([t] * b).to(z.device)
+        vc, decode_aux = model.decode(img=z, timesteps=timesteps, code=code)
+
+        if null_code is not None and (
+            interval is None
+            or ((t.item() >= interval[0]) and (t.item() <= interval[1]))
+        ):
+            vu, _ = model.decode(img=z, timesteps=timesteps, code=null_code)
+            vc = vu + cfg * (vc - vu)
+
+        z = z - dt * vc
+        images.append(z)
+    return images
+
+
+def build_model(config):
+    with tempfile.TemporaryDirectory() as log_dir:
+        MUP_ENABLED = config.model.enable_mup
+        model_partial = FlowMo
+
+        shared_kwargs = dict(config=config)
+        model = model_partial(
+            **shared_kwargs,
+            width=config.model.mup_width,
+        ).cuda()
+
+        if config.model.enable_mup:
+            print("Mup enabled!")
+            with torch.device("cpu"):
+                base_model = model_partial(
+                    **shared_kwargs, width=config.model.mup_width
+                )
+                delta_model = model_partial(
+                    **shared_kwargs,
+                    width=(
+                        config.model.mup_width * 4 if config.model.mup_width == 1 else 1
+                    ),
+                )
+                true_model = model_partial(
+                    **shared_kwargs, width=config.model.mup_width
+                )
+
+                if torch.distributed.is_initialized():
+                    bsh_path = os.path.join(log_dir, f"{dist.get_rank()}.bsh")
+                else:
+                    bsh_path = os.path.join(log_dir, "0.bsh")
+                set_base_shapes(
+                    true_model, base_model, delta=delta_model, savefile=bsh_path
+                )
+
+            model = set_base_shapes(model, base=bsh_path)
+
+            for module in model.modules():
+                if isinstance(module, MuReadout):
+                    module.width_mult = lambda: module.weight.infshape.width_mult()
+    return model

src/vqvaes/flowmo/lookup_free_quantize.py

@@ -0,0 +1,396 @@
+"""
+Code is from https://github.com/TencentARC/SEED-Voken. Thanks!
+
+Lookup Free Quantization
+Proposed in https://arxiv.org/abs/2310.05737
+
+In the simplest setup, each dimension is quantized into {-1, 1}.
+An entropy penalty is used to encourage utilization.
+
+Refer to
+https://github.com/lucidrains/vector-quantize-pytorch/blob/master/vector_quantize_pytorch/lookup_free_quantization.py
+https://github.com/theAdamColton/ijepa-enhanced/blob/7edef5f7288ae8f537f0db8a10044a2a487f70c9/ijepa_enhanced/lfq.py
+"""
+
+from collections import namedtuple
+from math import ceil, log2
+
+import torch
+import torch.nn.functional as F
+from einops import pack, rearrange, reduce, unpack
+from torch import einsum
+from torch.nn import Module
+
+# constants
+
+LossBreakdown = namedtuple(
+    "LossBreakdown",
+    ["per_sample_entropy", "codebook_entropy", "commitment", "avg_probs"],
+)
+
+# helper functions
+
+
+def exists(v):
+    return v is not None
+
+
+def default(*args):
+    for arg in args:
+        if exists(arg):
+            return arg() if callable(arg) else arg
+    return None
+
+
+def pack_one(t, pattern):
+    return pack([t], pattern)
+
+
+def unpack_one(t, ps, pattern):
+    return unpack(t, ps, pattern)[0]
+
+
+# entropy
+
+# def log(t, eps = 1e-5):
+#     return t.clamp(min = eps).log()
+
+
+def entropy(prob):
+    return (-prob * torch.log(prob + 1e-5)).sum(dim=-1)
+
+
+# class
+
+
+def mult_along_first_dims(x, y):
+    """
+    returns x * y elementwise along the leading dimensions of y
+    """
+    ndim_to_expand = x.ndim - y.ndim
+    for _ in range(ndim_to_expand):
+        y = y.unsqueeze(-1)
+    return x * y
+
+
+def masked_mean(x, m):
+    """
+    takes the mean of the elements of x that are not masked
+    the mean is taken along the shared leading dims of m
+    equivalent to: x[m].mean(tuple(range(m.ndim)))
+
+    The benefit of using masked_mean rather than using
+    tensor indexing is that masked_mean is much faster
+    for torch-compile on batches.
+
+    The drawback is larger floating point errors
+    """
+    x = mult_along_first_dims(x, m)
+    x = x / m.sum()
+    return x.sum(tuple(range(m.ndim)))
+
+
+def entropy_loss(
+    logits,
+    mask=None,
+    # temperature=0.01,
+    sample_minimization_weight=1.0,
+    batch_maximization_weight=1.0,
+    eps=1e-5,
+):
+    """
+    Entropy loss of unnormalized logits
+
+    logits: Affinities are over the last dimension
+
+    https://github.com/google-research/magvit/blob/05e8cfd6559c47955793d70602d62a2f9b0bdef5/videogvt/train_lib/losses.py#L279
+    LANGUAGE MODEL BEATS DIFFUSION — TOKENIZER IS KEY TO VISUAL GENERATION (2024)
+    """
+    # import pdb
+    # pdb.set_trace()
+    # print(logits.shape)
+    # raise
+
+    temperature = 0.1
+    probs = F.softmax(logits / temperature, -1)
+    log_probs = F.log_softmax(logits / temperature + eps, -1)
+
+    if mask is not None:
+        # avg_probs = probs[mask].mean(tuple(range(probs.ndim - 1)))
+        # avg_probs = einx.mean("... D -> D", probs[mask])
+
+        avg_probs = masked_mean(probs, mask)
+        # avg_probs = einx.mean("... D -> D", avg_probs)
+    else:
+        avg_probs = reduce(probs, "... D -> D", "mean")
+
+    avg_entropy = -torch.sum(avg_probs * torch.log(avg_probs + eps))
+
+    sample_entropy = -torch.sum(probs * log_probs, -1)
+    if mask is not None:
+        # sample_entropy = sample_entropy[mask].mean()
+        sample_entropy = masked_mean(sample_entropy, mask).mean()
+    else:
+        sample_entropy = torch.mean(sample_entropy)
+
+    loss = (sample_minimization_weight * sample_entropy) - (
+        batch_maximization_weight * avg_entropy
+    )
+
+    return sample_entropy, avg_entropy, loss
+
+
+class LFQ(Module):
+    def __init__(
+        self,
+        *,
+        dim=None,
+        codebook_size=None,
+        num_codebooks=1,
+        sample_minimization_weight=1.0,
+        batch_maximization_weight=1.0,
+        token_factorization=False,
+        factorized_bits=[9, 9],
+    ):
+        super().__init__()
+
+        # some assert validations
+
+        assert exists(dim) or exists(
+            codebook_size
+        ), "either dim or codebook_size must be specified for LFQ"
+        assert (
+            not exists(codebook_size) or log2(codebook_size).is_integer()
+        ), f"your codebook size must be a power of 2 for lookup free quantization (suggested {2 ** ceil(log2(codebook_size))})"
+
+        self.codebook_size = default(codebook_size, lambda: 2**dim)
+        self.codebook_dim = int(log2(codebook_size))
+
+        codebook_dims = self.codebook_dim * num_codebooks
+        dim = default(dim, codebook_dims)
+
+        has_projections = dim != codebook_dims
+        self.has_projections = has_projections
+
+        self.dim = dim
+        self.codebook_dim = self.codebook_dim
+        self.num_codebooks = num_codebooks
+
+        # for entropy loss
+        self.sample_minimization_weight = sample_minimization_weight
+        self.batch_maximization_weight = batch_maximization_weight
+
+        # for no auxiliary loss, during inference
+        self.token_factorization = token_factorization
+        if not self.token_factorization:  # for first stage model
+            self.register_buffer(
+                "mask", 2 ** torch.arange(self.codebook_dim), persistent=False
+            )
+        else:
+            self.factorized_bits = factorized_bits
+            self.register_buffer(
+                "pre_mask", 2 ** torch.arange(factorized_bits[0]), persistent=False
+            )
+            self.register_buffer(
+                "post_mask", 2 ** torch.arange(factorized_bits[1]), persistent=False
+            )
+
+        self.register_buffer("zero", torch.tensor(0.0), persistent=False)
+
+        # codes
+
+        all_codes = torch.arange(codebook_size)
+        bits = self.indices_to_bits(all_codes)
+        codebook = bits * 2.0 - 1.0
+
+        self.register_buffer("codebook", codebook, persistent=False)
+
+    @property
+    def dtype(self):
+        return self.codebook.dtype
+
+    def indices_to_bits(self, x):
+        """
+        x: long tensor of indices
+
+        returns big endian bits
+        """
+        mask = 2 ** torch.arange(self.codebook_dim, device=x.device, dtype=torch.long)
+        # x is now big endian bits, the last dimension being the bits
+        x = (x.unsqueeze(-1) & mask) != 0
+        return x
+
+    def get_codebook_entry(self, x, bhwc, order):  # 0610
+        if self.token_factorization:
+            if order == "pre":
+                mask = 2 ** torch.arange(
+                    self.factorized_bits[0], device=x.device, dtype=torch.long
+                )
+            else:
+                mask = 2 ** torch.arange(
+                    self.factorized_bits[1], device=x.device, dtype=torch.long
+                )
+        else:
+            mask = 2 ** torch.arange(
+                self.codebook_dim, device=x.device, dtype=torch.long
+            )
+
+        x = (x.unsqueeze(-1) & mask) != 0
+        x = x * 2.0 - 1.0  # back to the float
+        ## scale back to the
+        b, h, w, c = bhwc
+        x = rearrange(x, "b (h w) c -> b h w c", h=h, w=w, c=c)
+        x = rearrange(x, "b h w c -> b c h w")
+        return x
+
+    def bits_to_indices(self, bits):
+        """
+        bits: bool tensor of big endian bits, where the last dimension is the bit dimension
+
+        returns indices, which are long integers from 0 to self.codebook_size
+        """
+        assert bits.shape[-1] == self.codebook_dim
+        indices = 2 ** torch.arange(
+            0,
+            self.codebook_dim,
+            1,
+            dtype=torch.long,
+            device=bits.device,
+        )
+        return (bits * indices).sum(-1)
+
+    def decode(self, x):
+        """
+        x: ... NH
+            where NH is number of codebook heads
+            A longtensor of codebook indices, containing values from
+            0 to self.codebook_size
+        """
+        x = self.indices_to_bits(x)
+        # to some sort of float
+        x = x.to(self.dtype)
+        # -1 or 1
+        x = x * 2 - 1
+        x = rearrange(x, "... NC Z-> ... (NC Z)")
+        return x
+
+    def forward(
+        self,
+        x,
+        inv_temperature=100.0,
+        return_loss_breakdown=False,
+        mask=None,
+        return_loss=True,
+    ):
+        """
+        einstein notation
+        b - batch
+        n - sequence (or flattened spatial dimensions)
+        d - feature dimension, which is also log2(codebook size)
+        c - number of codebook dim
+        """
+        # x = x.tanh() * 1.5
+
+        x = rearrange(x, "b d ... -> b ... d")
+        x, ps = pack_one(x, "b * d")
+        # split out number of codebooks
+
+        x = rearrange(x, "b n (c d) -> b n c d", c=self.num_codebooks)
+
+        codebook_value = torch.Tensor([1.0]).to(device=x.device, dtype=x.dtype)
+        quantized = torch.where(
+            x > 0, codebook_value, -codebook_value
+        )  # higher than 0 filled
+
+        # calculate indices
+        if self.token_factorization:
+            indices_pre = reduce(
+                (quantized[..., : self.factorized_bits[0]] > 0).int()
+                * self.pre_mask.int(),
+                "b n c d -> b n c",
+                "sum",
+            )
+            indices_post = reduce(
+                (quantized[..., self.factorized_bits[0] :] > 0).int()
+                * self.post_mask.int(),
+                "b n c d -> b n c",
+                "sum",
+            )
+        else:
+            # print(quantized.shape)
+            indices = reduce(
+                (quantized > 0).int() * self.mask.int(), "b n c d -> b n c", "sum"
+            )
+            # print(indices.shape)
+
+        # entropy aux loss
+
+        if self.training and return_loss:
+            logits = 2 * einsum("... i d, j d -> ... i j", x, self.codebook)
+            # the same as euclidean distance up to a constant
+            # import pdb
+            # pdb.set_trace()
+
+            per_sample_entropy, codebook_entropy, entropy_aux_loss = entropy_loss(
+                logits=logits,
+                sample_minimization_weight=self.sample_minimization_weight,
+                batch_maximization_weight=self.batch_maximization_weight,
+            )
+
+            avg_probs = self.zero
+        else:
+            # logits = 2 * einsum('... i d, j d -> ... i j', x, self.codebook)
+            # probs = F.softmax(logits / 0.01, -1)
+            # avg_probs = reduce(probs, "b n c d -> b d", "mean")
+            # avg_probs = torch.sum(avg_probs, 0) #batch dimension
+            # if not training, just return dummy 0
+            per_sample_entropy = codebook_entropy = self.zero
+            ## calculate the codebook_entropy needed for one batch evaluation
+            entropy_aux_loss = self.zero
+            avg_probs = self.zero
+
+        # commit loss
+
+        if self.training:
+            commit_loss = F.mse_loss(x, quantized.detach(), reduction="none")
+
+            if exists(mask):
+                commit_loss = commit_loss[mask]
+
+            commit_loss = commit_loss.mean()
+        else:
+            commit_loss = self.zero
+
+        # use straight-through gradients (optionally with custom activation fn) if training
+
+        quantized = x + (quantized - x).detach()  # transfer to quantized
+
+        # merge back codebook dim
+
+        quantized = rearrange(quantized, "b n c d -> b n (c d)")
+
+        # reconstitute image or video dimensions
+
+        quantized = unpack_one(quantized, ps, "b * d")
+        quantized = rearrange(quantized, "b ... d -> b d ...")
+
+        if self.token_factorization:
+            indices_pre = unpack_one(indices_pre, ps, "b * c")
+            indices_post = unpack_one(indices_post, ps, "b * c")
+            indices_pre = indices_pre.flatten()
+            indices_post = indices_post.flatten()
+            indices = (indices_pre, indices_post)
+        else:
+            # print(indices.shape, ps)
+            indices = unpack_one(indices, ps, "b * c")
+            # print(indices.shape)
+            indices = indices.flatten()
+
+        ret = (quantized, entropy_aux_loss, indices)
+
+        if not return_loss_breakdown:
+            return ret
+
+        return ret, LossBreakdown(
+            per_sample_entropy, codebook_entropy, commit_loss, avg_probs
+        )

src/vqvaes/infinity/conv.py

@@ -0,0 +1,107 @@
+import torch
+import torch.nn as nn
+from einops import rearrange
+import torch.nn.functional as F
+
+
+class Conv(nn.Module):
+    def __init__(
+        self,
+        in_channels,
+        out_channels,
+        kernel_size,
+        stride=1,
+        padding=0,
+        cnn_type="2d",
+        causal_offset=0,
+        temporal_down=False,
+    ):
+        super().__init__()
+        self.cnn_type = cnn_type
+        self.slice_seq_len = 17
+
+        if cnn_type == "2d":
+            self.conv = nn.Conv2d(
+                in_channels, out_channels, kernel_size, stride=stride, padding=padding
+            )
+        if cnn_type == "3d":
+            if temporal_down == False:
+                stride = (1, stride, stride)
+            else:
+                stride = (stride, stride, stride)
+            self.conv = nn.Conv3d(
+                in_channels, out_channels, kernel_size, stride=stride, padding=0
+            )
+            if isinstance(kernel_size, int):
+                kernel_size = (kernel_size, kernel_size, kernel_size)
+            self.padding = (
+                kernel_size[0] - 1 + causal_offset,  # Temporal causal padding
+                padding,  # Height padding
+                padding,  # Width padding
+            )
+        self.causal_offset = causal_offset
+        self.stride = stride
+        self.kernel_size = kernel_size
+
+    def forward(self, x):
+        if self.cnn_type == "2d":
+            if x.ndim == 5:
+                B, C, T, H, W = x.shape
+                x = rearrange(x, "B C T H W -> (B T) C H W")
+                x = self.conv(x)
+                x = rearrange(x, "(B T) C H W -> B C T H W", T=T)
+                return x
+            else:
+                return self.conv(x)
+        if self.cnn_type == "3d":
+            assert (
+                self.stride[0] == 1 or self.stride[0] == 2
+            ), f"only temporal stride = 1 or 2 are supported"
+            xs = []
+            for i in range(0, x.shape[2], self.slice_seq_len + self.stride[0] - 1):
+                st = i
+                en = min(i + self.slice_seq_len, x.shape[2])
+                _x = x[:, :, st:en, :, :]
+                if i == 0:
+                    _x = F.pad(
+                        _x,
+                        (
+                            self.padding[2],
+                            self.padding[2],  # Width
+                            self.padding[1],
+                            self.padding[1],  # Height
+                            self.padding[0],
+                            0,
+                        ),
+                    )  # Temporal
+                else:
+                    padding_0 = self.kernel_size[0] - 1
+                    _x = F.pad(
+                        _x,
+                        (
+                            self.padding[2],
+                            self.padding[2],  # Width
+                            self.padding[1],
+                            self.padding[1],  # Height
+                            padding_0,
+                            0,
+                        ),
+                    )  # Temporal
+                    _x[
+                        :,
+                        :,
+                        :padding_0,
+                        self.padding[1] : _x.shape[-2] - self.padding[1],
+                        self.padding[2] : _x.shape[-1] - self.padding[2],
+                    ] += x[:, :, i - padding_0 : i, :, :]
+                _x = self.conv(_x)
+                xs.append(_x)
+            try:
+                x = torch.cat(xs, dim=2)
+            except:
+                device = x.device
+                del x
+                xs = [_x.cpu().pin_memory() for _x in xs]
+                torch.cuda.empty_cache()
+                x = torch.cat([_x.cpu() for _x in xs], dim=2).to(device=device)
+            return x

src/vqvaes/infinity/dynamic_resolution.py

@@ -0,0 +1,147 @@
+import json
+import numpy as np
+import tqdm
+
+vae_stride = 16
+ratio2hws = {
+    1.000: [
+        (1, 1),
+        (2, 2),
+        (4, 4),
+        (6, 6),
+        (8, 8),
+        (12, 12),
+        (16, 16),
+        (20, 20),
+        (24, 24),
+        (32, 32),
+        (40, 40),
+        (48, 48),
+        (64, 64),
+    ],
+    1.250: [
+        (1, 1),
+        (2, 2),
+        (3, 3),
+        (5, 4),
+        (10, 8),
+        (15, 12),
+        (20, 16),
+        (25, 20),
+        (30, 24),
+        (35, 28),
+        (45, 36),
+        (55, 44),
+        (70, 56),
+    ],
+    1.333: [
+        (1, 1),
+        (2, 2),
+        (4, 3),
+        (8, 6),
+        (12, 9),
+        (16, 12),
+        (20, 15),
+        (24, 18),
+        (28, 21),
+        (36, 27),
+        (48, 36),
+        (60, 45),
+        (72, 54),
+    ],
+    1.500: [
+        (1, 1),
+        (2, 2),
+        (3, 2),
+        (6, 4),
+        (9, 6),
+        (15, 10),
+        (21, 14),
+        (27, 18),
+        (33, 22),
+        (39, 26),
+        (48, 32),
+        (63, 42),
+        (78, 52),
+    ],
+    1.750: [
+        (1, 1),
+        (2, 2),
+        (3, 3),
+        (7, 4),
+        (11, 6),
+        (14, 8),
+        (21, 12),
+        (28, 16),
+        (35, 20),
+        (42, 24),
+        (56, 32),
+        (70, 40),
+        (84, 48),
+    ],
+    2.000: [
+        (1, 1),
+        (2, 2),
+        (4, 2),
+        (6, 3),
+        (10, 5),
+        (16, 8),
+        (22, 11),
+        (30, 15),
+        (38, 19),
+        (46, 23),
+        (60, 30),
+        (74, 37),
+        (90, 45),
+    ],
+    2.500: [
+        (1, 1),
+        (2, 2),
+        (5, 2),
+        (10, 4),
+        (15, 6),
+        (20, 8),
+        (25, 10),
+        (30, 12),
+        (40, 16),
+        (50, 20),
+        (65, 26),
+        (80, 32),
+        (100, 40),
+    ],
+    3.000: [
+        (1, 1),
+        (2, 2),
+        (6, 2),
+        (9, 3),
+        (15, 5),
+        (21, 7),
+        (27, 9),
+        (36, 12),
+        (45, 15),
+        (54, 18),
+        (72, 24),
+        (90, 30),
+        (111, 37),
+    ],
+}
+full_ratio2hws = {}
+for ratio, hws in ratio2hws.items():
+    full_ratio2hws[ratio] = hws
+    full_ratio2hws[int(1 / ratio * 1000) / 1000] = [(item[1], item[0]) for item in hws]
+
+dynamic_resolution_h_w = {}
+predefined_HW_Scales_dynamic = {}
+for ratio in full_ratio2hws:
+    dynamic_resolution_h_w[ratio] = {}
+    for ind, leng in enumerate([7, 10, 13]):
+        h, w = (
+            full_ratio2hws[ratio][leng - 1][0],
+            full_ratio2hws[ratio][leng - 1][1],
+        )  # feature map size
+        pixel = (h * vae_stride, w * vae_stride)  # The original image (H, W)
+        dynamic_resolution_h_w[ratio][pixel[1]] = {
+            "pixel": pixel,
+            "scales": full_ratio2hws[ratio][:leng],
+        }  # W as key
+        predefined_HW_Scales_dynamic[(h, w)] = full_ratio2hws[ratio][:leng]

src/vqvaes/infinity/flux_vqgan.py

@@ -0,0 +1,771 @@
+import argparse
+import os
+import imageio
+import torch
+import numpy as np
+from einops import rearrange
+from torch import Tensor, nn
+import torch.nn.functional as F
+import torchvision
+from torchvision import transforms
+from safetensors.torch import load_file
+import torch.utils.checkpoint as checkpoint
+
+from .conv import Conv
+from .multiscale_bsq import MultiScaleBSQ
+
+ptdtype = {None: torch.float32, "fp32": torch.float32, "bf16": torch.bfloat16}
+
+
+class Normalize(nn.Module):
+    def __init__(self, in_channels, norm_type, norm_axis="spatial"):
+        super().__init__()
+        self.norm_axis = norm_axis
+        assert norm_type in ["group", "batch", "no"]
+        if norm_type == "group":
+            if in_channels % 32 == 0:
+                self.norm = nn.GroupNorm(
+                    num_groups=32, num_channels=in_channels, eps=1e-6, affine=True
+                )
+            elif in_channels % 24 == 0:
+                self.norm = nn.GroupNorm(
+                    num_groups=24, num_channels=in_channels, eps=1e-6, affine=True
+                )
+            else:
+                raise NotImplementedError
+        elif norm_type == "batch":
+            self.norm = nn.SyncBatchNorm(
+                in_channels, track_running_stats=False
+            )  # Runtime Error: grad inplace if set track_running_stats to True
+        elif norm_type == "no":
+            self.norm = nn.Identity()
+
+    def forward(self, x):
+        if self.norm_axis == "spatial":
+            if x.ndim == 4:
+                x = self.norm(x)
+            else:
+                B, C, T, H, W = x.shape
+                x = rearrange(x, "B C T H W -> (B T) C H W")
+                x = self.norm(x)
+                x = rearrange(x, "(B T) C H W -> B C T H W", T=T)
+        elif self.norm_axis == "spatial-temporal":
+            x = self.norm(x)
+        else:
+            raise NotImplementedError
+        return x
+
+
+def swish(x: Tensor) -> Tensor:
+    try:
+        return x * torch.sigmoid(x)
+    except:
+        device = x.device
+        x = x.cpu().pin_memory()
+        return (x * torch.sigmoid(x)).to(device=device)
+
+
+class AttnBlock(nn.Module):
+    def __init__(self, in_channels, norm_type="group", cnn_param=None):
+        super().__init__()
+        self.in_channels = in_channels
+
+        self.norm = Normalize(
+            in_channels, norm_type, norm_axis=cnn_param["cnn_norm_axis"]
+        )
+
+        self.q = Conv(in_channels, in_channels, kernel_size=1)
+        self.k = Conv(in_channels, in_channels, kernel_size=1)
+        self.v = Conv(in_channels, in_channels, kernel_size=1)
+        self.proj_out = Conv(in_channels, in_channels, kernel_size=1)
+
+    def attention(self, h_: Tensor) -> Tensor:
+        B, _, T, _, _ = h_.shape
+        h_ = self.norm(h_)
+        h_ = rearrange(h_, "B C T H W -> (B T) C H W")  # spatial attention only
+        q = self.q(h_)
+        k = self.k(h_)
+        v = self.v(h_)
+
+        b, c, h, w = q.shape
+        q = rearrange(q, "b c h w -> b 1 (h w) c").contiguous()
+        k = rearrange(k, "b c h w -> b 1 (h w) c").contiguous()
+        v = rearrange(v, "b c h w -> b 1 (h w) c").contiguous()
+        h_ = nn.functional.scaled_dot_product_attention(q, k, v)
+
+        return rearrange(h_, "(b t) 1 (h w) c -> b c t h w", h=h, w=w, c=c, b=B, t=T)
+
+    def forward(self, x: Tensor) -> Tensor:
+        return x + self.proj_out(self.attention(x))
+
+
+class ResnetBlock(nn.Module):
+    def __init__(
+        self, in_channels: int, out_channels: int, norm_type="group", cnn_param=None
+    ):
+        super().__init__()
+        self.in_channels = in_channels
+        out_channels = in_channels if out_channels is None else out_channels
+        self.out_channels = out_channels
+
+        self.norm1 = Normalize(
+            in_channels, norm_type, norm_axis=cnn_param["cnn_norm_axis"]
+        )
+        if cnn_param["res_conv_2d"] in ["half", "full"]:
+            self.conv1 = Conv(
+                in_channels,
+                out_channels,
+                kernel_size=3,
+                stride=1,
+                padding=1,
+                cnn_type="2d",
+            )
+        else:
+            self.conv1 = Conv(
+                in_channels,
+                out_channels,
+                kernel_size=3,
+                stride=1,
+                padding=1,
+                cnn_type=cnn_param["cnn_type"],
+            )
+        self.norm2 = Normalize(
+            out_channels, norm_type, norm_axis=cnn_param["cnn_norm_axis"]
+        )
+        if cnn_param["res_conv_2d"] in ["full"]:
+            self.conv2 = Conv(
+                out_channels,
+                out_channels,
+                kernel_size=3,
+                stride=1,
+                padding=1,
+                cnn_type="2d",
+            )
+        else:
+            self.conv2 = Conv(
+                out_channels,
+                out_channels,
+                kernel_size=3,
+                stride=1,
+                padding=1,
+                cnn_type=cnn_param["cnn_type"],
+            )
+        if self.in_channels != self.out_channels:
+            self.nin_shortcut = Conv(
+                in_channels, out_channels, kernel_size=1, stride=1, padding=0
+            )
+
+    def forward(self, x):
+        h = x
+        h = self.norm1(h)
+        h = swish(h)
+        h = self.conv1(h)
+
+        h = self.norm2(h)
+        h = swish(h)
+        h = self.conv2(h)
+
+        if self.in_channels != self.out_channels:
+            x = self.nin_shortcut(x)
+
+        return x + h
+
+
+class Downsample(nn.Module):
+    def __init__(
+        self, in_channels, cnn_type="2d", spatial_down=False, temporal_down=False
+    ):
+        super().__init__()
+        assert spatial_down == True
+        if cnn_type == "2d":
+            self.pad = (0, 1, 0, 1)
+        if cnn_type == "3d":
+            self.pad = (
+                0,
+                1,
+                0,
+                1,
+                0,
+                0,
+            )  # add padding to the right for h-axis and w-axis. No padding for t-axis
+        # no asymmetric padding in torch conv, must do it ourselves
+        self.conv = Conv(
+            in_channels,
+            in_channels,
+            kernel_size=3,
+            stride=2,
+            padding=0,
+            cnn_type=cnn_type,
+            temporal_down=temporal_down,
+        )
+
+    def forward(self, x: Tensor):
+        x = nn.functional.pad(x, self.pad, mode="constant", value=0)
+        x = self.conv(x)
+        return x
+
+
+class Upsample(nn.Module):
+    def __init__(
+        self,
+        in_channels,
+        cnn_type="2d",
+        spatial_up=False,
+        temporal_up=False,
+        use_pxsl=False,
+    ):
+        super().__init__()
+        if cnn_type == "2d":
+            self.scale_factor = 2
+            self.causal_offset = 0
+        else:
+            assert spatial_up == True
+            if temporal_up:
+                self.scale_factor = (2, 2, 2)
+                self.causal_offset = -1
+            else:
+                self.scale_factor = (1, 2, 2)
+                self.causal_offset = 0
+        self.use_pxsl = use_pxsl
+        if self.use_pxsl:
+            self.conv = Conv(
+                in_channels,
+                in_channels * 4,
+                kernel_size=3,
+                stride=1,
+                padding=1,
+                cnn_type=cnn_type,
+                causal_offset=self.causal_offset,
+            )
+            self.pxsl = nn.PixelShuffle(2)
+        else:
+            self.conv = Conv(
+                in_channels,
+                in_channels,
+                kernel_size=3,
+                stride=1,
+                padding=1,
+                cnn_type=cnn_type,
+                causal_offset=self.causal_offset,
+            )
+
+    def forward(self, x: Tensor):
+        if self.use_pxsl:
+            x = self.conv(x)
+            x = self.pxsl(x)
+        else:
+            try:
+                x = F.interpolate(x, scale_factor=self.scale_factor, mode="nearest")
+            except:
+                # shard across channel
+                _xs = []
+                for i in range(x.shape[1]):
+                    _x = F.interpolate(
+                        x[:, i : i + 1, ...],
+                        scale_factor=self.scale_factor,
+                        mode="nearest",
+                    )
+                    _xs.append(_x)
+                x = torch.cat(_xs, dim=1)
+            x = self.conv(x)
+        return x
+
+
+class Encoder(nn.Module):
+    def __init__(
+        self,
+        ch: int,
+        ch_mult: list[int],
+        num_res_blocks: int,
+        z_channels: int,
+        in_channels=3,
+        patch_size=8,
+        temporal_patch_size=4,
+        norm_type="group",
+        cnn_param=None,
+        use_checkpoint=False,
+        use_vae=True,
+    ):
+        super().__init__()
+        self.max_down = np.log2(patch_size)
+        self.temporal_max_down = np.log2(temporal_patch_size)
+        self.temporal_down_offset = self.max_down - self.temporal_max_down
+        self.ch = ch
+        self.num_resolutions = len(ch_mult)
+        self.num_res_blocks = num_res_blocks
+        self.in_channels = in_channels
+        self.cnn_param = cnn_param
+        self.use_checkpoint = use_checkpoint
+        # downsampling
+        # self.conv_in = Conv(in_channels, self.ch, kernel_size=3, stride=1, padding=1)
+        # cnn_param["cnn_type"] = "2d" for images, cnn_param["cnn_type"] = "3d" for videos
+        if cnn_param["conv_in_out_2d"] == "yes":  # "yes" for video
+            self.conv_in = Conv(
+                in_channels, ch, kernel_size=3, stride=1, padding=1, cnn_type="2d"
+            )
+        else:
+            self.conv_in = Conv(
+                in_channels,
+                ch,
+                kernel_size=3,
+                stride=1,
+                padding=1,
+                cnn_type=cnn_param["cnn_type"],
+            )
+
+        in_ch_mult = (1,) + tuple(ch_mult)
+        self.in_ch_mult = in_ch_mult
+        self.down = nn.ModuleList()
+        block_in = self.ch
+        for i_level in range(self.num_resolutions):
+            block = nn.ModuleList()
+            attn = nn.ModuleList()
+            block_in = ch * in_ch_mult[i_level]
+            block_out = ch * ch_mult[i_level]
+            for _ in range(self.num_res_blocks):
+                block.append(
+                    ResnetBlock(
+                        in_channels=block_in,
+                        out_channels=block_out,
+                        norm_type=norm_type,
+                        cnn_param=cnn_param,
+                    )
+                )
+                block_in = block_out
+            down = nn.Module()
+            down.block = block
+            down.attn = attn
+            # downsample, stride=1, stride=2, stride=2 for 4x8x8 Video VAE
+            spatial_down = True if i_level < self.max_down else False
+            temporal_down = (
+                True
+                if i_level < self.max_down and i_level >= self.temporal_down_offset
+                else False
+            )
+            if spatial_down or temporal_down:
+                down.downsample = Downsample(
+                    block_in,
+                    cnn_type=cnn_param["cnn_type"],
+                    spatial_down=spatial_down,
+                    temporal_down=temporal_down,
+                )
+            self.down.append(down)
+
+        # middle
+        self.mid = nn.Module()
+        self.mid.block_1 = ResnetBlock(
+            in_channels=block_in,
+            out_channels=block_in,
+            norm_type=norm_type,
+            cnn_param=cnn_param,
+        )
+        if cnn_param["cnn_attention"] == "yes":
+            self.mid.attn_1 = AttnBlock(block_in, norm_type, cnn_param=cnn_param)
+        self.mid.block_2 = ResnetBlock(
+            in_channels=block_in,
+            out_channels=block_in,
+            norm_type=norm_type,
+            cnn_param=cnn_param,
+        )
+
+        # end
+        self.norm_out = Normalize(
+            block_in, norm_type, norm_axis=cnn_param["cnn_norm_axis"]
+        )
+        if cnn_param["conv_inner_2d"] == "yes":
+            self.conv_out = Conv(
+                block_in,
+                (int(use_vae) + 1) * z_channels,
+                kernel_size=3,
+                stride=1,
+                padding=1,
+                cnn_type="2d",
+            )
+        else:
+            self.conv_out = Conv(
+                block_in,
+                (int(use_vae) + 1) * z_channels,
+                kernel_size=3,
+                stride=1,
+                padding=1,
+                cnn_type=cnn_param["cnn_type"],
+            )
+
+    def forward(self, x, return_hidden=False):
+        if not self.use_checkpoint:
+            return self._forward(x, return_hidden=return_hidden)
+        else:
+            return checkpoint.checkpoint(
+                self._forward, x, return_hidden, use_reentrant=False
+            )
+
+    def _forward(self, x: Tensor, return_hidden=False) -> Tensor:
+        # downsampling
+        h0 = self.conv_in(x)
+        hs = [h0]
+        for i_level in range(self.num_resolutions):
+            for i_block in range(self.num_res_blocks):
+                h = self.down[i_level].block[i_block](hs[-1])
+                if len(self.down[i_level].attn) > 0:
+                    h = self.down[i_level].attn[i_block](h)
+                hs.append(h)
+            if hasattr(self.down[i_level], "downsample"):
+                hs.append(self.down[i_level].downsample(hs[-1]))
+
+        # middle
+        h = hs[-1]
+        hs_mid = [h]
+        h = self.mid.block_1(h)
+        if self.cnn_param["cnn_attention"] == "yes":
+            h = self.mid.attn_1(h)
+        h = self.mid.block_2(h)
+        hs_mid.append(h)
+        # end
+        h = self.norm_out(h)
+        h = swish(h)
+        h = self.conv_out(h)
+        if return_hidden:
+            return h, hs, hs_mid
+        else:
+            return h
+
+
+class Decoder(nn.Module):
+    def __init__(
+        self,
+        ch: int,
+        ch_mult: list[int],
+        num_res_blocks: int,
+        z_channels: int,
+        out_ch=3,
+        patch_size=8,
+        temporal_patch_size=4,
+        norm_type="group",
+        cnn_param=None,
+        use_checkpoint=False,
+        use_freq_dec=False,  # use frequency features for decoder
+        use_pxsf=False,
+    ):
+        super().__init__()
+        self.max_up = np.log2(patch_size)
+        self.temporal_max_up = np.log2(temporal_patch_size)
+        self.temporal_up_offset = self.max_up - self.temporal_max_up
+        self.ch = ch
+        self.num_resolutions = len(ch_mult)
+        self.num_res_blocks = num_res_blocks
+        self.ffactor = 2 ** (self.num_resolutions - 1)
+        self.cnn_param = cnn_param
+        self.use_checkpoint = use_checkpoint
+        self.use_freq_dec = use_freq_dec
+        self.use_pxsf = use_pxsf
+
+        # compute in_ch_mult, block_in and curr_res at lowest res
+        block_in = ch * ch_mult[self.num_resolutions - 1]
+
+        # z to block_in
+        if cnn_param["conv_inner_2d"] == "yes":
+            self.conv_in = Conv(
+                z_channels, block_in, kernel_size=3, stride=1, padding=1, cnn_type="2d"
+            )
+        else:
+            self.conv_in = Conv(
+                z_channels,
+                block_in,
+                kernel_size=3,
+                stride=1,
+                padding=1,
+                cnn_type=cnn_param["cnn_type"],
+            )
+
+        # middle
+        self.mid = nn.Module()
+        self.mid.block_1 = ResnetBlock(
+            in_channels=block_in,
+            out_channels=block_in,
+            norm_type=norm_type,
+            cnn_param=cnn_param,
+        )
+        if cnn_param["cnn_attention"] == "yes":
+            self.mid.attn_1 = AttnBlock(
+                block_in, norm_type=norm_type, cnn_param=cnn_param
+            )
+        self.mid.block_2 = ResnetBlock(
+            in_channels=block_in,
+            out_channels=block_in,
+            norm_type=norm_type,
+            cnn_param=cnn_param,
+        )
+
+        # upsampling
+        self.up = nn.ModuleList()
+        for i_level in reversed(range(self.num_resolutions)):
+            block = nn.ModuleList()
+            attn = nn.ModuleList()
+            block_out = ch * ch_mult[i_level]
+            for _ in range(self.num_res_blocks + 1):
+                block.append(
+                    ResnetBlock(
+                        in_channels=block_in,
+                        out_channels=block_out,
+                        norm_type=norm_type,
+                        cnn_param=cnn_param,
+                    )
+                )
+                block_in = block_out
+            up = nn.Module()
+            up.block = block
+            up.attn = attn
+            # upsample, stride=1, stride=2, stride=2 for 4x8x8 Video VAE, offset 1 compared with encoder
+            # https://github.com/black-forest-labs/flux/blob/b4f689aaccd40de93429865793e84a734f4a6254/src/flux/modules/autoencoder.py#L228
+            spatial_up = True if 1 <= i_level <= self.max_up else False
+            temporal_up = (
+                True
+                if 1 <= i_level <= self.max_up
+                and i_level >= self.temporal_up_offset + 1
+                else False
+            )
+            if spatial_up or temporal_up:
+                up.upsample = Upsample(
+                    block_in,
+                    cnn_type=cnn_param["cnn_type"],
+                    spatial_up=spatial_up,
+                    temporal_up=temporal_up,
+                    use_pxsl=self.use_pxsf,
+                )
+            self.up.insert(0, up)  # prepend to get consistent order
+
+        # end
+        self.norm_out = Normalize(
+            block_in, norm_type, norm_axis=cnn_param["cnn_norm_axis"]
+        )
+        if cnn_param["conv_in_out_2d"] == "yes":
+            self.conv_out = Conv(
+                block_in, out_ch, kernel_size=3, stride=1, padding=1, cnn_type="2d"
+            )
+        else:
+            self.conv_out = Conv(
+                block_in,
+                out_ch,
+                kernel_size=3,
+                stride=1,
+                padding=1,
+                cnn_type=cnn_param["cnn_type"],
+            )
+
+    def forward(self, z):
+        if not self.use_checkpoint:
+            return self._forward(z)
+        else:
+            return checkpoint.checkpoint(self._forward, z, use_reentrant=False)
+
+    def _forward(self, z: Tensor) -> Tensor:
+        # z to block_in
+        h = self.conv_in(z)
+
+        # middle
+        h = self.mid.block_1(h)
+        if self.cnn_param["cnn_attention"] == "yes":
+            h = self.mid.attn_1(h)
+        h = self.mid.block_2(h)
+
+        # upsampling
+        for i_level in reversed(range(self.num_resolutions)):
+            for i_block in range(self.num_res_blocks + 1):
+                h = self.up[i_level].block[i_block](h)
+                if len(self.up[i_level].attn) > 0:
+                    h = self.up[i_level].attn[i_block](h)
+            if hasattr(self.up[i_level], "upsample"):
+                h = self.up[i_level].upsample(h)
+
+        # end
+        h = self.norm_out(h)
+        h = swish(h)
+        h = self.conv_out(h)
+        return h
+
+
+class AutoEncoder(nn.Module):
+    def __init__(self, args):
+        super().__init__()
+        self.args = args
+        cnn_param = dict(
+            cnn_type=args.cnn_type,
+            conv_in_out_2d=args.conv_in_out_2d,
+            res_conv_2d=args.res_conv_2d,
+            cnn_attention=args.cnn_attention,
+            cnn_norm_axis=args.cnn_norm_axis,
+            conv_inner_2d=args.conv_inner_2d,
+        )
+        self.encoder = Encoder(
+            ch=args.base_ch,
+            ch_mult=args.encoder_ch_mult,
+            num_res_blocks=args.num_res_blocks,
+            z_channels=args.codebook_dim,
+            patch_size=args.patch_size,
+            temporal_patch_size=args.temporal_patch_size,
+            cnn_param=cnn_param,
+            use_checkpoint=args.use_checkpoint,
+            use_vae=args.use_vae,
+        )
+        self.decoder = Decoder(
+            ch=args.base_ch,
+            ch_mult=args.decoder_ch_mult,
+            num_res_blocks=args.num_res_blocks,
+            z_channels=args.codebook_dim,
+            patch_size=args.patch_size,
+            temporal_patch_size=args.temporal_patch_size,
+            cnn_param=cnn_param,
+            use_checkpoint=args.use_checkpoint,
+            use_freq_dec=args.use_freq_dec,
+            use_pxsf=args.use_pxsf,  # pixelshuffle for upsampling
+        )
+        self.z_drop = nn.Dropout(args.z_drop)
+        self.scale_factor = 0.3611
+        self.shift_factor = 0.1159
+        self.codebook_dim = self.embed_dim = args.codebook_dim
+
+        self.gan_feat_weight = args.gan_feat_weight
+        self.video_perceptual_weight = args.video_perceptual_weight
+        self.recon_loss_type = args.recon_loss_type
+        self.l1_weight = args.l1_weight
+        self.use_vae = args.use_vae
+        self.kl_weight = args.kl_weight
+        self.lfq_weight = args.lfq_weight
+        self.image_gan_weight = args.image_gan_weight  # image GAN loss weight
+        self.video_gan_weight = args.video_gan_weight  # video GAN loss weight
+        self.perceptual_weight = args.perceptual_weight
+        self.flux_weight = args.flux_weight
+        self.cycle_weight = args.cycle_weight
+        self.cycle_feat_weight = args.cycle_feat_weight
+        self.cycle_gan_weight = args.cycle_gan_weight
+
+        self.flux_image_encoder = None
+
+        if not args.use_vae:
+            if args.quantizer_type == "MultiScaleBSQ":
+                self.quantizer = MultiScaleBSQ(
+                    dim=args.codebook_dim,  # this is the input feature dimension, defaults to log2(codebook_size) if not defined
+                    codebook_size=args.codebook_size,  # codebook size, must be a power of 2
+                    entropy_loss_weight=args.entropy_loss_weight,  # how much weight to place on entropy loss
+                    diversity_gamma=args.diversity_gamma,  # within entropy loss, how much weight to give to diversity of codes, taken from https://arxiv.org/abs/1911.05894
+                    preserve_norm=args.preserve_norm,  # preserve norm of the input for BSQ
+                    ln_before_quant=args.ln_before_quant,  # use layer norm before quantization
+                    ln_init_by_sqrt=args.ln_init_by_sqrt,  # layer norm init value 1/sqrt(d)
+                    commitment_loss_weight=args.commitment_loss_weight,  # loss weight of commitment loss
+                    new_quant=args.new_quant,
+                    use_decay_factor=args.use_decay_factor,
+                    mask_out=args.mask_out,
+                    use_stochastic_depth=args.use_stochastic_depth,
+                    drop_rate=args.drop_rate,
+                    schedule_mode=args.schedule_mode,
+                    keep_first_quant=args.keep_first_quant,
+                    keep_last_quant=args.keep_last_quant,
+                    remove_residual_detach=args.remove_residual_detach,
+                    use_out_phi=args.use_out_phi,
+                    use_out_phi_res=args.use_out_phi_res,
+                    random_flip=args.random_flip,
+                    flip_prob=args.flip_prob,
+                    flip_mode=args.flip_mode,
+                    max_flip_lvl=args.max_flip_lvl,
+                    random_flip_1lvl=args.random_flip_1lvl,
+                    flip_lvl_idx=args.flip_lvl_idx,
+                    drop_when_test=args.drop_when_test,
+                    drop_lvl_idx=args.drop_lvl_idx,
+                    drop_lvl_num=args.drop_lvl_num,
+                )
+                self.quantize = self.quantizer
+                self.vocab_size = args.codebook_size
+            else:
+                raise NotImplementedError(f"{args.quantizer_type} not supported")
+
+    def forward(self, x):
+        is_image = x.ndim == 4
+        if not is_image:
+            B, C, T, H, W = x.shape
+        else:
+            B, C, H, W = x.shape
+            T = 1
+        enc_dtype = ptdtype[self.args.encoder_dtype]
+
+        with torch.amp.autocast("cuda", dtype=enc_dtype):
+            h, hs, hs_mid = self.encoder(x, return_hidden=True)  # B C H W or B C T H W
+        hs = [_h.detach() for _h in hs]
+        hs_mid = [_h.detach() for _h in hs_mid]
+        h = h.to(dtype=torch.float32)
+        # print(z.shape)
+        # Multiscale LFQ
+        z, all_indices, _, _, all_loss, _ = self.quantizer(h)
+        x_recon = self.decoder(z)
+        vq_output = {
+            "commitment_loss": torch.mean(all_loss)
+            * self.lfq_weight,  # here commitment loss is sum of commitment loss and entropy penalty
+            "encodings": all_indices,
+        }
+        # return x_recon, vq_output
+        return x_recon, None, z
+
+    def encode_for_raw_features(
+        self, x, scale_schedule, return_residual_norm_per_scale=False
+    ):
+        is_image = x.ndim == 4
+        if not is_image:
+            B, C, T, H, W = x.shape
+        else:
+            B, C, H, W = x.shape
+            T = 1
+
+        enc_dtype = ptdtype[self.args.encoder_dtype]
+        with torch.amp.autocast("cuda", dtype=enc_dtype):
+            h, hs, hs_mid = self.encoder(x, return_hidden=True)  # B C H W or B C T H W
+
+        hs = [_h.detach() for _h in hs]
+        hs_mid = [_h.detach() for _h in hs_mid]
+        h = h.to(dtype=torch.float32)
+        return h, hs, hs_mid
+
+    def encode(self, x, scale_schedule, return_residual_norm_per_scale=False):
+        h, hs, hs_mid = self.encode_for_raw_features(
+            x, scale_schedule, return_residual_norm_per_scale
+        )
+        # Multiscale LFQ
+        (
+            z,
+            all_indices,
+            all_bit_indices,
+            residual_norm_per_scale,
+            all_loss,
+            var_input,
+        ) = self.quantizer(
+            h,
+            scale_schedule=scale_schedule,
+            return_residual_norm_per_scale=return_residual_norm_per_scale,
+        )
+        return h, z, all_indices, all_bit_indices, residual_norm_per_scale, var_input
+
+    def decode(self, z):
+        x_recon = self.decoder(z)
+        x_recon = torch.clamp(x_recon, min=-1, max=1)
+        return x_recon
+
+    def decode_from_indices(self, all_indices, scale_schedule, label_type):
+        summed_codes = 0
+        for idx_Bl in all_indices:
+            codes = self.quantizer.lfq.indices_to_codes(idx_Bl, label_type)
+            summed_codes += F.interpolate(
+                codes, size=scale_schedule[-1], mode=self.quantizer.z_interplote_up
+            )
+        assert summed_codes.shape[-3] == 1
+        x_recon = self.decoder(summed_codes.squeeze(-3))
+        x_recon = torch.clamp(x_recon, min=-1, max=1)
+        return summed_codes, x_recon
+
+    @staticmethod
+    def add_model_specific_args(parent_parser):
+        parser = argparse.ArgumentParser(parents=[parent_parser], add_help=False)
+        parser.add_argument("--flux_weight", type=float, default=0)
+        parser.add_argument("--cycle_weight", type=float, default=0)
+        parser.add_argument("--cycle_feat_weight", type=float, default=0)
+        parser.add_argument("--cycle_gan_weight", type=float, default=0)
+        parser.add_argument("--cycle_loop", type=int, default=0)
+        parser.add_argument("--z_drop", type=float, default=0.0)
+        return parser