[update] app.py

.gitignore

@@ -0,0 +1,30 @@
+checkpoints
+__pycache__
+imagenet_tiny
+*.log
+*.ckpt
+data
+wandb
+.vscode
+heun-steps*
+log.txt
+imagenet_openai_ref_images
+*.jpg
+*.npz
+src
+scan
+*.txt
+*.env
+*.pem, 
+secrets.yaml
+scripts
+scripts/*/local*.sh
+checkpoints_vision
+tmp_videos
+weights
+.DS_Store
+local
+*.mp4
+tmp
+.git_bk
+demo

LICENSE

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

README.md

@@ -0,0 +1,389 @@
+# GRN: Generative Refinement Networks
+
+[![arXiv](https://img.shields.io/badge/arXiv%20paper-2604.13030-b31b1b.svg)](https://arxiv.org/abs/2604.13030)
+[![Homepage](https://img.shields.io/badge/🏠%20Homepage-GRN-green.svg)](https://bytedance.github.io/GRN/)
+[![Models](https://img.shields.io/badge/🤗%20Hugging%20Face-Models-blue.svg)](https://huggingface.co/bytedance-research/GRN)
+[![Demo](https://img.shields.io/badge/🤗%20Hugging%20Face-Demo-yellow.svg)](https://huggingface.co/spaces/hanjian/GRN)
+[![License](https://img.shields.io/badge/License-MIT-blue.svg)](LICENSE)
+[![GitHub stars](https://img.shields.io/github/stars/bytedance/GRN?style=social)](https://github.com/bytedance/GRN)
+
+---
+
+## 🔥 Updates!!
+* June 3, 2026: 🍉 A toy image-video dataset is provided for GRN-T2I/GRN-T2V training and fine-tuning.
+* May 23, 2026: 🌺 We release the training and evaluation code for HBQ tokenizer, enjoy~
+* April 14, 2026: 🤗 Paper and code release
+
+## 📋 Table of Contents
+
+- [🌟 Introduction](#-introduction)
+- [✨ Gallery](#-gallery)
+- [🚀 Demo](#-demo)
+- [📦 Model Zoo](#-model-zoo)
+- [🛠️ Installation](#️-installation)
+- [🖼️ Class-to-Image](#️-class-to-image)
+  - [Dataset](#dataset)
+  - [Training](#training)
+  - [Evaluation](#evaluation)
+- [🎨 Text-to-Image](#-text-to-image)
+  - [Data](#data)
+  - [Train](#train)
+  - [Inference](#inference)
+- [🎬 Text-to-Video](#-text-to-video)
+  - [Data](#data-1)
+  - [Train](#train-1)
+  - [Inference](#inference-1)
+- [📦 HBQ Tokenizer](#-hbq-tokenizer)
+  - [Data](#data-2)
+  - [Training](#training-1)
+  - [Evaluation](#evaluation-1)
+- [📧 Contact](#-contact)
+- [🤗 Acknowledgements](#-acknowledgements)
+- [📝 Citation](#-citation)
+
+---
+
+## 🌟 Introduction
+
+This is the official implementation of the paper **Generative Refinement Networks for Visual Synthesis**. Neither diffusion nor autoregressive — GRN is a third way. 🧠 Refines globally like an artist. ⚡ Generates adaptively by complexity. 🏆 New SOTA across image & video. The visual generation paradigm just got rewritten.
+
+Diffusion models dominate visual generation but they allocate uniform computational effort to samples with varying levels of complexity. Autoregressive (AR) models are complexity-aware, as evidenced by their variable likelihoods, but suffer from lossy tokenization and error accumulation.
+
+We introduce **Generative Refinement Networks (GRN)**, a new visual synthesis paradigm that addresses these issues:
+- **Near-lossless tokenization** via Hierarchical Binary Quantization (HBQ)
+- **Global refinement mechanism** that progressively perfects outputs like a human artist
+- **Entropy-guided sampling** for complexity-aware, adaptive-step generation
+
+GRN achieves state-of-the-art results on ImageNet reconstruction and class-conditional generation, and scales effectively to text-to-image and text-to-video tasks.
+
+---
+
+<figure align="center">
+  <figcaption><strong><em>Generative Refinement Framework</em></strong></figcaption>
+  <img src="demo/framework.jpg" width="100%" alt="Framework">
+</figure>
+
+<p align="center">
+Starting from a random token map, GRN randomly selects more predictions at each step and refines all input tokens. For example, compared to the second step, the third step filled six new tokens (<span style="color: rgb(220, 120, 117);">pink</span>), kept two tokens (<span style="color: rgb(88, 160, 227);">blue</span>), erased two tokens (<span style="color: rgb(240, 180, 40);">yellow</span>), and left six tokens blank (<span style="color: rgb(128, 138, 151);">gray</span>).
+</p>
+
+---
+
+## ✨ Gallery
+
+### GRN-8B Text-to-Video Examples
+
+<div align="center">
+  <table style="border-spacing: 6px; margin: auto;">
+    <tr>
+      <td style="padding: 2px;"><video src="https://github.com/user-attachments/assets/6ce844dc-3185-4239-bcf1-d72ff20a3031" width="33%" autoplay muted loop playsinline></video></td>
+      <td style="padding: 2px;"><video src="https://github.com/user-attachments/assets/1697066e-f00f-4e23-a55c-c6af5948c4af" width="33%" autoplay muted loop playsinline></video></td>
+      <td style="padding: 2px;"><video src="https://github.com/user-attachments/assets/1023ea4d-d814-4be1-95f2-b1623de0f6bd" width="33%" autoplay muted loop playsinline></video></td>
+    </tr>
+    <tr>
+      <td style="padding: 2px;"><video src="https://github.com/user-attachments/assets/6244dae4-f480-408a-ac3d-19e4d1ef0a2d" width="33%" autoplay muted loop playsinline></video></td>
+      <td style="padding: 2px;"><video src="https://github.com/user-attachments/assets/5aefc8d2-bc99-48e4-bd1c-9b3077c9c35e" width="33%" autoplay muted loop playsinline></video></td>
+      <td style="padding: 2px;"><video src="https://github.com/user-attachments/assets/014e8bb4-04a7-4fa4-a597-d0dfbcc23e02" width="33%" autoplay muted loop playsinline></video></td>
+    </tr>
+    <tr>
+      <td style="padding: 2px;"><video src="https://github.com/user-attachments/assets/6bde1f2e-cebe-4f47-9eac-4fe817c3ebc7" width="33%" autoplay muted loop playsinline></video></td>
+      <td style="padding: 2px;"><video src="https://github.com/user-attachments/assets/b9957300-fa98-411c-83d5-f972621245ad" width="33%" autoplay muted loop playsinline></video></td>
+      <td style="padding: 2px;"><video src="https://github.com/user-attachments/assets/d07cef92-3eec-4e7c-93da-f6c6a8dc1658" width="33%" autoplay muted loop playsinline></video></td>
+    </tr>
+  </table>
+</div>
+
+---
+
+### GRN-8B Image-to-Video Examples
+
+<div align="center">
+  <table style="border-spacing: 6px; margin: auto;">
+    <tr>
+      <td style="padding: 2px;"><video src="https://github.com/user-attachments/assets/527f94b0-4b04-4cbb-a86d-9f5ae05fab67" width="33%" autoplay muted loop playsinline></video></td>
+      <td style="padding: 2px;"><video src="https://github.com/user-attachments/assets/0b63a9ed-2940-402f-8339-db0d05a09525" width="33%" autoplay muted loop playsinline></video></td>
+      <td style="padding: 2px;"><video src="https://github.com/user-attachments/assets/8d108a5f-1414-43ca-af6b-8862640741e5" width="33%" autoplay muted loop playsinline></video></td>
+    </tr>
+    <tr>
+      <td style="padding: 2px;"><video src="https://github.com/user-attachments/assets/64cd45a9-0c2f-4926-bcc0-b8a0a939ae54" width="33%" autoplay muted loop playsinline></video></td>
+      <td style="padding: 2px;"><video src="https://github.com/user-attachments/assets/6c31c9e5-0742-4416-925c-16c39bc5a03a" width="33%" autoplay muted loop playsinline></video></td>
+      <td style="padding: 2px;"><video src="https://github.com/user-attachments/assets/4e966b46-6107-4ffe-a24b-37dc3c8461dd" width="33%" autoplay muted loop playsinline></video></td>
+    </tr>
+    <tr>
+      <td style="padding: 2px;"><video src="https://github.com/user-attachments/assets/56ce2dc1-3b64-4493-ab27-b2ba273c64ef" width="33%" autoplay muted loop playsinline></video></td>
+      <td style="padding: 2px;"><video src="https://github.com/user-attachments/assets/ef45a3f4-8fb2-4bb5-885e-19645e5a0fb5" width="33%" autoplay muted loop playsinline></video></td>
+      <td style="padding: 2px;"><video src="https://github.com/user-attachments/assets/98e3fbb0-9a54-49e6-8cec-96a42d0634e6" width="33%" autoplay muted loop playsinline></video></td>
+    </tr>
+  </table>
+</div>
+
+### GRN-2B Class-to-Image Examples
+<figure align="center">
+  <!-- <figcaption><strong><em>GRN-2B Class-to-Image Examples</em></strong></figcaption> -->
+  <img src="demo/c2i_examples.jpg" width="100%" alt="Class-to-Image Examples">
+</figure>
+
+### GRN-2B Text-to-Image Examples
+<figure align="center">
+  <!-- <figcaption><strong><em>GRN-2B Text-to-Image Examples</em></strong></figcaption> -->
+  <img src="demo/t2i_examples.jpg" width="100%" alt="Text-to-Image Examples">
+</figure>
+
+---
+
+## 🚀 Demo
+
+### 🖼️ Text-to-Image
+Try our interactive Text-to-Image demo on 🤗 Hugging Face Space:
+
+**[GRN T2I Demo](https://huggingface.co/spaces/hanjian/GRN)**
+
+Experience the power of Generative Refinement Networks firsthand by generating images from text prompts directly in your browser!
+
+---
+
+### 🎬 Text-to-Video
+Try our interactive Text-to-Video demo on Discord:
+
+[![Discord](https://img.shields.io/badge/Discord-Join%20Server-5865F2?style=for-the-badge&logo=discord&logoColor=white)](http://opensource.bytedance.com/discord/invite)
+
+
+<figure align="center">
+  <figcaption><strong><em>T2V Demo on Discord</em></strong></figcaption>
+  <img src="demo/t2v_demo.png" width="100%" alt="T2V Demo">
+</figure>
+
+---
+
+## 📦 Model Zoo
+
+| Model | Checkpoints |
+|-------|:-----------:|
+| **Tokenizers** | ✅ [ImageNet Tokenizer](https://huggingface.co/bytedance-research/GRN/blob/main/HBQ_image_tokenizer_16dim_M4.ckpt)<br>✅ [Joint Image/Video Tokenizer](https://huggingface.co/bytedance-research/GRN/blob/main/HBQ_tokenizer_64dim_M4.ckpt) |
+| **GRN_ind_C2I** | ✅ [B](https://huggingface.co/bytedance-research/GRN/blob/main/GRN_ind_B_ep599.pth)<br>⬜ L (TBD)<br>⬜ H (TBD)<br>⬜ G (TBD) |
+| **GRN_bit_T2I** | ✅ [GRN_T2I](https://huggingface.co/bytedance-research/GRN/blob/main/GRN_T2I_2B.pth) |
+| **GRN_bit_T2V** | ✅ [GRN_T2V](https://huggingface.co/bytedance-research/GRN/blob/main/GRN_T2V_2B.pth) |
+
+---
+
+## 🛠️ Installation
+
+### Step 1: Clone the repository
+```bash
+git clone https://github.com/bytedance/GRN
+cd GRN
+```
+
+### Step 2: Create conda environment
+A suitable [conda](https://conda.io/) environment named `GRN` can be created and activated with:
+```bash
+conda env create -f environment.yaml
+conda activate GRN
+```
+
+### Troubleshooting
+If you get `undefined symbol: iJIT_NotifyEvent` when importing `torch`, simply:
+```bash
+pip uninstall torch
+pip install torch==2.5.1 --index-url https://download.pytorch.org/whl/cu124
+```
+Check this [issue](https://github.com/conda/conda/issues/13812#issuecomment-2071445372) for more details.
+
+---
+
+## 🖼️ Class-to-Image
+
+### Dataset
+Download [ImageNet](http://image-net.org/download) dataset, and place it in your `IMAGENET_PATH`.
+
+### Training
+
+All training scripts are located in `scripts/c2i/`. We suggest using 8x80GB GPUs for most models.
+
+| Model | Training Script | GPUs Required |
+|-------|:-------------:|:-------------:|
+| GRN_ind_B | `bash scripts/c2i/train_GRN_ind_B.sh` | 8x80GB |
+| GRN_bit_B | `bash scripts/c2i/train_GRN_bit_B.sh` | 8x80GB |
+| GRN_ind_L | `bash scripts/c2i/train_GRN_ind_L.sh` | 8x80GB |
+| GRN_ind_H | `bash scripts/c2i/train_GRN_ind_H.sh` | 16x80GB |
+| GRN_ind_G | `bash scripts/c2i/train_GRN_ind_G.sh` | 32x80GB |
+
+### Evaluation
+
+PyTorch pre-trained models are available [here](https://huggingface.co/bytedance-research/GRN/tree/main).
+
+All evaluation scripts are located in `scripts/c2i/`. We suggest using 8x80GB vRAM GPUs.
+
+| Model | Evaluation Script |
+|-------|:--------------:|
+| GRN_ind_B | `bash scripts/c2i/eval_GRN_ind_B.sh` |
+| GRN_bit_B | `bash scripts/c2i/eval_GRN_bit_B.sh` |
+| GRN_ind_L | `bash scripts/c2i/eval_GRN_ind_L.sh` |
+| GRN_ind_H | `bash scripts/c2i/eval_GRN_ind_H.sh` |
+| GRN_ind_G | `bash scripts/c2i/eval_GRN_ind_G.sh` |
+
+We use [torch-fidelity](https://github.com/LTH14/torch-fidelity) to evaluate FID and IS against a reference image folder or statistics. We use the JiT's pre-computed reference stats under `grn/utils_c2i/fid_stats`.
+
+---
+
+## 🎨 Text-to-Image
+### Data
+Refer to `data/toy_data/jsonls/000001/0001_0008_000000100.jsonl`
+```
+{"image_path": "[image_path_1]", "long_caption": "xxx", "long_caption_type": "caption-InternVL2.0", "text": "", "short_caption_type": "blip2_caption", "width": 1080, "height": 1920}
+{"image_path": "[image_path_2]", "long_caption": "xxx", "long_caption_type": "caption-InternVL2.0", "text": "", "short_caption_type": "blip2_caption", "width": 1080, "height": 1920}
+...
+```
+
+### Train
+Run `bash scripts/train_GRN_ind_t2i.sh`
+
+### Inference
+
+You can simply run `python3 t2i_infer.py` or use the following code:
+
+```python
+from PIL import Image
+from grn_pipeline import GRNPipeline
+
+# Load pipeline
+pipeline = GRNPipeline.from_pretrained(
+    hf_repo_id='bytedance-research/GRN',
+    task='T2I',
+    pn='1M', 
+    device='cpu',
+).to('cuda')
+
+# Generate one image
+result = pipeline(
+    prompt="A cute cat playing in the garden",
+    guidance_scale=3.0,
+    temperature=1.1,
+    complexity_aware_Tmin=10,
+    complexity_aware_Tmax=50,
+    complexity_aware_k = 0,
+    complexity_aware_b = 50,
+    complexity_aware_wp = 5,
+    snr_shift = 1.,
+    h_div_w=1.,
+    content_type='image',
+    seed=42,
+)
+image = result.images[0]
+image.save('./generated_image.jpg')
+```
+
+---
+
+## 🎬 Text-to-Video
+### Data
+Refer to `data/toy_data/jsonls/000001/0001_0008_000000100.jsonl`
+```
+{"video_path": "[video_path_1]", "begin_frame_id": xxx, "end_frame_id": xxx, "quality_prompt": "There is text in the video.", "fps": 25.0, "duration": 3.88, "width": 1280, "height": 720, "caption": [{"type": "short", "content": "[short_caption]"}, {"type": "medium", "content": "[medium_caption]"}, {"type": "long", "content": "[long_caption]"}]}
+{"video_path": "[video_path_1]", "begin_frame_id": xxx, "end_frame_id": xxx, "quality_prompt": "The quality is very high!", "fps": 25.0, "duration": 3.88, "width": 1280, "height": 720, "caption": [{"type": "short", "content": "[short_caption]"}, {"type": "medium", "content": "[medium_caption]"}, {"type": "long", "content": "[long_caption]"}]}
+...
+```
+
+### Train
+Run `bash scripts/train_GRN_ind_t2v.sh`
+
+### Inference
+
+You can simply run `python3 t2v_infer.py` or use the following code:
+
+```python
+from grn_pipeline import GRNPipeline
+
+# Load pipeline
+pipeline = GRNPipeline.from_pretrained(
+    hf_repo_id='bytedance-research/GRN', 
+    task='T2V', 
+    pn='0.41M', 
+    device='cpu'
+).to('cuda')
+
+# Generate one video
+result = pipeline(
+    prompt="Two women demonstrate a makeup product, applying it with a sponge while smiling and engaging with the camera in a bright, clean setting.",
+    guidance_scale=4.0,
+    temperature=1.0,
+    complexity_aware_Tmin=10,
+    complexity_aware_Tmax=50,
+    complexity_aware_k = 0,
+    complexity_aware_b = 50,
+    complexity_aware_wp = 5,
+    snr_shift = 1.,
+    h_div_w=9/16,
+    duration=2.,
+    content_type='video',
+    seed=42,
+)
+video_file = result.videos[0]
+```
+
+---
+
+## 📦 HBQ Tokenizer
+
+### Data
+Image Dataset, e.g., data_root/username/labels/imagenet/train.txt:
+```
+[image_1_full_path]
+[image_2_full_path]
+[image_3_full_path]
+...
+```
+
+Video Dataset, e.g., data_root/username/labels_hanjian/high-quality-video/horizontal_videos.txt
+```
+[video_1_full_path]
+[video_2_full_path]
+[video_3_full_path]
+...
+```
+
+### Training
+For example, set `latent_channels=16/64` and `quant_method=hierarchical_binary_quant_round_4` in `scripts/hbq_tokenizer_train.sh`, then run:
+```bash
+cd grn/tokenizer
+bash scripts/hbq_tokenizer_train.sh
+```
+
+### Evaluation
+For example, set `latent_channels=16/64` and `quant_method=hierarchical_binary_quant_round_4` in `scripts/hbq_tokenizer_train.sh`, then run:
+```bash
+cd grn/tokenizer
+bash scripts/hbq_tokenizer_eval.sh
+```
+
+---
+
+## 📧 Contact
+
+If you are interested in scaling GRN for image generation / image editing / video generation / video editing / unified model directions, please feel free to reach out!
+
+**📧 Email:** [hanjian.thu123@bytedance.com](mailto:hanjian.thu123@bytedance.com)
+
+---
+
+## 🤗 Acknowledgements
+
+- Thanks to [JiT](https://github.com/LTH14/JiT), [Infinity](https://github.com/FoundationVision/Infinity) and [InfinityStar](https://github.com/FoundationVision/InfinityStar) for their wonderful work and codebase!
+
+---
+
+## 📝 Citation
+
+If you find our work useful, please consider citing:
+
+```bibtex
+@misc{han2026grn,
+      title={Generative Refinement Networks for Visual Synthesis}, 
+      author={Jian Han and Jinlai Liu and Jiahuan Wang and Bingyue Peng and Zehuan Yuan},
+      year={2026},
+      eprint={2604.13030},
+      archivePrefix={arXiv},
+      primaryClass={cs.CV},
+      url={https://arxiv.org/abs/2604.13030}, 
+}
+```

app.py

@@ -0,0 +1,132 @@
+import os
+import sys
+import traceback
+import torch
+import gradio as gr
+import spaces
+
+sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
+
+from grn_pipeline import GRNPipeline
+
+# Global pipeline
+pipe = None
+device = "cuda" if torch.cuda.is_available() else "cpu"
+
+def load_pipeline():
+    global pipe
+    print("Loading GRN pipeline...")
+    # 从 Hugging Face Hub 下载权重
+    pipe = GRNPipeline.from_pretrained(
+        hf_repo_id='bytedance-research/GRN',
+        task='T2I',
+        pn='1M', 
+        model='GRN2b',
+        use_slow_attn=True,
+        device=device,
+    )
+    print("Pipeline loaded successfully!")
+    return pipe
+
+# @spaces.GPU #[uncomment to use ZeroGPU]
+@spaces.GPU(duration=40)
+def generate(prompt, content_type="image", guidance_scale=3.0, temperature=1.0, seed=42, width=1024, height=1024):
+    global pipe
+    if pipe is None:
+        try:
+            pipe = load_pipeline()
+        except Exception as e:
+            print(f"Error loading pipeline: {e}")
+            traceback.print_exc()
+            return f"Error loading pipeline: {e}\n\n{traceback.format_exc()}"
+    
+    try:
+        result = pipe(
+            prompt="<T2I>"+prompt,
+            guidance_scale=guidance_scale,
+            temperature=temperature,
+            complexity_aware_Tmin=10,
+            complexity_aware_Tmax=50,
+            complexity_aware_k = 0,
+            complexity_aware_b = 50,
+            complexity_aware_wp = 5,
+            snr_shift = 1.,
+            h_div_w=1.,
+            content_type=content_type,
+            seed=seed,
+            width=width,
+            height=height
+        )
+        
+        if content_type == "image" and hasattr(result, 'images'):
+            return result.images[0]
+        elif content_type == "video" and hasattr(result, 'videos'):
+            return result.videos[0]
+        return f"Error: Invalid result from pipeline"
+    except Exception as e:
+        print(f"Error generating content: {e}")
+        traceback.print_exc()
+        return f"Error generating content: {e}\n\n{traceback.format_exc()}"
+
+def create_demo():
+    with gr.Blocks(title="GRN: Generative Refinement Networks", theme=gr.themes.Soft()) as demo:
+        gr.Markdown("# GRN: Generative Refinement Networks")
+        gr.Markdown("Text-to-Image generation using GRN")
+        
+        with gr.Row():
+            with gr.Column():
+                prompt_input = gr.Textbox(
+                    label="Text Prompt",
+                    placeholder="Enter your prompt here...",
+                    value="A cute cat playing in the garden"
+                )
+                
+                content_type = gr.Radio(
+                    choices=["image"], # , "video"
+                    value="image",
+                    label="Content Type"
+                )
+                
+                with gr.Accordion("Settings", open=True):
+                    guidance_scale = gr.Slider(minimum=0, maximum=10, value=3.0, label="Guidance Scale")
+                    temperature = gr.Slider(minimum=0.1, maximum=1.5, value=1.1, label="Temperature")
+                    seed = gr.Number(value=42, label="Seed", precision=0)
+                    width = gr.Number(value=1024, label="Width", precision=0)
+                    height = gr.Number(value=1024, label="Height", precision=0)
+                
+                generate_btn = gr.Button("Generate", variant="primary")
+            
+            with gr.Column():
+                output = gr.Gallery(label="Output", show_label=True, elem_id="gallery", columns=1, height="auto", preview=True, object_fit="contain")
+        
+        def generate_and_display(prompt, content_type, guidance_scale, temperature, seed, width, height):
+            result = generate(prompt, content_type, guidance_scale, temperature, seed, width, height)
+            if result:
+                return [result]
+            return []
+        
+        generate_btn.click(
+            fn=generate_and_display,
+            inputs=[prompt_input, content_type, guidance_scale, temperature, seed, width, height],
+            outputs=output
+        )
+        
+        gr.Examples(
+            examples=[
+                ["A majestic lion standing on a cliff at sunset", "image", 3.0, 1.0, 42, 1024, 1024],
+            ],
+            inputs=[prompt_input, content_type, guidance_scale, temperature, seed, width, height],
+            cache_examples=False
+        )
+    
+    return demo
+
+if __name__ == "__main__":
+    try:
+        load_pipeline()
+    except Exception as e:
+        print(f"Error loading pipeline: {e}")
+        traceback.print_exc()
+    
+    demo = create_demo()
+    demo.launch()

c2i_train_infer.py

@@ -0,0 +1,378 @@
+import argparse
+import datetime
+import numpy as np
+import os
+import time
+import functools
+from pathlib import Path
+
+import torch
+import torch.backends.cudnn as cudnn
+from torch.utils.tensorboard import SummaryWriter
+import torchvision.transforms as transforms
+import torchvision.datasets as datasets
+from torch.distributed.device_mesh import init_device_mesh
+from torch.distributed.fsdp import (
+    FullyShardedDataParallel as FSDP,
+    MixedPrecision,
+    BackwardPrefetch,
+    ShardingStrategy,
+    FullStateDictConfig,
+    StateDictType,
+)
+from torch.distributed.fsdp.wrap import (
+    transformer_auto_wrap_policy,
+    enable_wrap,
+    wrap,
+)
+
+from grn.utils_c2i.crop import center_crop_arr
+import grn.utils_c2i.misc as misc
+
+import copy
+from grn.utils_c2i.engine import train_one_epoch, evaluate
+from grn.utils import wandb_utils as wandb_utils
+
+from grn.utils_c2i.denoiser import Denoiser
+from grn.models.grn_c2i import GRNblock
+
+
+def get_args_parser():
+    parser = argparse.ArgumentParser('GRN', add_help=False)
+
+    # architecture
+    parser.add_argument('--model', default='GRN_B', type=str, metavar='MODEL',
+                        help='Name of the model to train')
+    parser.add_argument('--img_size', default=256, type=int, help='Image size')
+    parser.add_argument('--attn_dropout', type=float, default=0.0, help='Attention dropout rate')
+    parser.add_argument('--proj_dropout', type=float, default=0.0, help='Projection dropout rate')
+
+    # training
+    parser.add_argument('--epochs', default=200, type=int)
+    parser.add_argument('--warmup_epochs', type=int, default=5, metavar='N',
+                        help='Epochs to warm up LR')
+    parser.add_argument('--batch_size', default=128, type=int,
+                        help='Batch size per GPU (effective batch size = batch_size * # GPUs)')
+    parser.add_argument('--lr', type=float, default=None, metavar='LR',
+                        help='Learning rate (absolute)')
+    parser.add_argument('--blr', type=float, default=5e-5, metavar='LR',
+                        help='Base learning rate: absolute_lr = base_lr * total_batch_size / 256')
+    parser.add_argument('--min_lr', type=float, default=0., metavar='LR',
+                        help='Minimum LR for cyclic schedulers that hit 0')
+    parser.add_argument('--lr_schedule', type=str, default='constant',
+                        help='Learning rate schedule')
+    parser.add_argument('--weight_decay', type=float, default=0.0,
+                        help='Weight decay (default: 0.0)')
+    parser.add_argument('--ema_decay1', type=float, default=0.9999,
+                        help='The first ema to track. Use the first ema for sampling by default.')
+    parser.add_argument('--ema_decay2', type=float, default=0.9996,
+                        help='The second ema to track')
+    parser.add_argument('--P_mean', default=-0.8, type=float)
+    parser.add_argument('--P_std', default=0.8, type=float)
+    parser.add_argument('--noise_scale', default=1.0, type=float)
+    parser.add_argument('--t_eps', default=5e-2, type=float)
+    parser.add_argument('--label_drop_prob', default=0.1, type=float)
+
+    parser.add_argument('--seed', default=0, type=int)
+    parser.add_argument('--start_epoch', default=0, type=int, metavar='N',
+                        help='Starting epoch')
+    parser.add_argument('--num_workers', default=12, type=int)
+    parser.add_argument('--pin_mem', action='store_true',
+                        help='Pin CPU memory in DataLoader for faster GPU transfers')
+    parser.add_argument('--no_pin_mem', action='store_false', dest='pin_mem')
+    parser.set_defaults(pin_mem=True)
+
+    # sampling
+    parser.add_argument('--sampling_method', default='heun', type=str,
+                        help='ODE samping method')
+    parser.add_argument('--num_sampling_steps', default=50, type=int,
+                        help='Sampling steps')
+    parser.add_argument('--cfg', default=1.0, type=float,
+                        help='Classifier-free guidance factor')
+    parser.add_argument('--interval_min', default=0.0, type=float,
+                        help='CFG interval min')
+    parser.add_argument('--interval_max', default=1.0, type=float,
+                        help='CFG interval max')
+    parser.add_argument('--num_images', default=50000, type=int,
+                        help='Number of images to generate')
+    parser.add_argument('--eval_freq', type=int, default=40,
+                        help='Frequency (in epochs) for evaluation')
+    parser.add_argument('--online_eval', type=int, default=0, choices=[0,1],
+                        help='Whether to evaluate the model online')
+    parser.add_argument('--evaluate_gen', action='store_true')
+    parser.add_argument('--gen_bsz', type=int, default=256,
+                        help='Generation batch size')
+
+    # dataset
+    parser.add_argument('--data_path', default='./data/imagenet', type=str,
+                        help='Path to the dataset')
+    parser.add_argument('--class_num', default=1000, type=int)
+
+    # checkpointing
+    parser.add_argument('--output_dir', default='./output_dir',
+                        help='Directory to save outputs (empty for no saving)')
+    parser.add_argument('--resume', default='',
+                        help='Folder that contains checkpoint to resume from')
+    parser.add_argument('--save_last_freq', type=int, default=5,
+                        help='Frequency (in epochs) to save checkpoints')
+    parser.add_argument('--log_freq', default=100, type=int)
+    parser.add_argument('--device', default='cuda',
+                        help='Device to use for training/testing')
+
+    # distributed training
+    parser.add_argument('--world_size', default=1, type=int,
+                        help='Number of distributed processes')
+    parser.add_argument('--local_rank', default=-1, type=int)
+    parser.add_argument('--dist_on_itp', action='store_true')
+    parser.add_argument('--dist_url', default='env://',
+                        help='URL used to set up distributed training')
+    parser.add_argument('--hbq_round', default=4, type=int,)
+    parser.add_argument('--in_channels', default=3, type=int,)
+    parser.add_argument('--method', default='GRN_ind', type=str, choices=['GRN_ind', 'GRN_bit'])
+    parser.add_argument('--vae_path', default='', type=str,)
+    parser.add_argument('--tau', default=1.0, type=float,)
+    parser.add_argument('--wandb', default=1, type=int, choices=[0,1])
+    parser.add_argument('--generation_dir', default='/tmp', type=str)
+    parser.add_argument('--clip_grad_norm', default=1., type=float)
+    parser.add_argument('--use_fsdp_train', default=0, type=int, choices=[0, 1])
+    parser.add_argument('--delete_images', default=1, type=int, choices=[0, 1])
+    parser.add_argument('--use_confidence_sampling', default=0, type=int, choices=[0, 1])
+    parser.add_argument('--inner_shard_degree', default=8, type=int)
+    parser.add_argument('--patch_size', default=1, type=int)
+    parser.add_argument('--convert_type', default='', type=str)
+    parser.add_argument('--mask_group_size', default=-1, type=int)
+    parser.add_argument('--grn_shift_factor', default=1., type=float)
+    parser.add_argument('--use_focal_loss', default=0, type=int, choices=[0, 1])
+    return parser
+
+
+def main(args):
+    misc.init_distributed_mode(args)
+    print('Job directory:', os.path.dirname(os.path.realpath(__file__)))
+    print("Arguments:\n{}".format(args).replace(', ', ',\n'))
+
+    device = torch.device(args.device)
+
+    # Set seeds for reproducibility
+    seed = args.seed + misc.get_rank()
+    torch.manual_seed(seed)
+    np.random.seed(seed)
+
+    cudnn.benchmark = True
+
+    num_tasks = misc.get_world_size()
+    global_rank = misc.get_rank()
+
+    # Set up TensorBoard logging (only on main process)
+    if global_rank == 0 and args.output_dir is not None:
+        os.makedirs(args.output_dir, exist_ok=True)
+        log_writer = SummaryWriter(log_dir=args.output_dir)
+        if args.wandb:
+            entity = os.environ["EXP_NAME"]
+            project = os.environ["PROJECT"]
+            wandb_utils.wandb.init(project=project, name=entity, config={})
+    else:
+        log_writer = None
+
+    # Data augmentation transforms
+    transform_train = transforms.Compose([
+        transforms.Lambda(lambda img: center_crop_arr(img, args.img_size)),
+        transforms.RandomHorizontalFlip(),
+        transforms.PILToTensor()
+    ])
+
+    dataset_train = datasets.ImageFolder(os.path.join(args.data_path, 'train'), transform=transform_train)
+    print(dataset_train)
+
+    sampler_train = torch.utils.data.DistributedSampler(
+        dataset_train, num_replicas=num_tasks, rank=global_rank, shuffle=True
+    )
+    print("Sampler_train =", sampler_train)
+
+    data_loader_train = torch.utils.data.DataLoader(
+        dataset_train, sampler=sampler_train,
+        batch_size=args.batch_size,
+        num_workers=args.num_workers,
+        pin_memory=args.pin_mem,
+        drop_last=True
+    )
+
+    torch._dynamo.config.cache_size_limit = 128
+    torch._dynamo.config.optimize_ddp = False
+
+    # Create denoiser
+    model = Denoiser(args)
+
+    # ininitalize vae
+    from grn.models.hbq_tokenizer import HBQ_Tokenizer
+    vae = HBQ_Tokenizer(args=args, latent_channels=16, encoder_out_type='feature_tanh')
+    vae.eval()
+    vae = vae.to('cuda')
+    for param in vae.parameters():
+        param.requires_grad = False
+    state_dict = torch.load(args.vae_path, map_location='cuda')
+    if 'ema' in state_dict:
+        print(f'Load ema vae weights')
+        state_dict = state_dict['ema']
+    else:
+        print(f'Load non ema vae weights')
+        state_dict = state_dict['vae']
+    print('Load vae: ', vae.load_state_dict(state_dict, assign=True))
+
+    print("Model =", model)
+    n_params = sum(p.numel() for p in model.parameters() if p.requires_grad)
+    print("Number of trainable parameters: {:.6f}M".format(n_params / 1e6))
+
+    model.to(device)
+
+    eff_batch_size = args.batch_size * misc.get_world_size()
+    if args.lr is None:  # only base_lr (blr) is specified
+        args.lr = args.blr * eff_batch_size / 256
+
+    print("Base lr: {:.2e}".format(args.lr * 256 / eff_batch_size))
+    print("Actual lr: {:.2e}".format(args.lr))
+    print("Effective batch size: %d" % eff_batch_size)
+
+    if args.use_fsdp_train:
+        auto_wrap_policy = functools.partial(
+            transformer_auto_wrap_policy,
+            transformer_layer_cls={GRNblock},
+        )
+        if args.inner_shard_degree > 0:
+            sharding_strategy = ShardingStrategy.HYBRID_SHARD
+            world_size = misc.get_world_size()
+            assert world_size % args.inner_shard_degree == 0
+            assert args.inner_shard_degree > 1 and args.inner_shard_degree <= world_size
+            device_mesh = init_device_mesh('cuda', (world_size // args.inner_shard_degree, args.inner_shard_degree))
+        else:
+            sharding_strategy = ShardingStrategy.FULL_SHARD
+            device_mesh = None
+        model = FSDP(
+            model,
+            auto_wrap_policy=auto_wrap_policy,
+            mixed_precision=MixedPrecision(
+                param_dtype=torch.bfloat16, 
+                reduce_dtype=torch.bfloat16, 
+                buffer_dtype=torch.bfloat16
+            ),
+            device_id=torch.cuda.current_device(),
+            sharding_strategy=sharding_strategy, 
+            use_orig_params=True,
+            device_mesh=device_mesh,
+        )
+        model_without_ddp = model
+    else:
+        model = torch.nn.parallel.DistributedDataParallel(model, device_ids=[args.gpu])
+        model_without_ddp = model.module
+
+    # Set up optimizer with weight decay adjustment for bias and norm layers
+    param_groups = misc.add_weight_decay(model_without_ddp, args.weight_decay)
+    optimizer = torch.optim.AdamW(param_groups, lr=args.lr, betas=(0.9, 0.95))
+    print(optimizer)
+
+    # Resume from checkpoint if provided
+    # checkpoint_path = os.path.join(args.resume, "checkpoint-last.pth") if args.resume else None
+    checkpoint_path = args.resume if args.resume else None
+    if checkpoint_path and os.path.exists(checkpoint_path):
+        checkpoint = torch.load(checkpoint_path, map_location='cpu')
+        model_without_ddp.load_state_dict(checkpoint['model'])
+
+        if args.use_fsdp_train:
+            # For FSDP, load EMA state dict into model temporarily to set ema_params
+            model_without_ddp.load_state_dict(checkpoint['model_ema1'])
+            model_without_ddp.module.ema_params1 = [p.detach().clone() for p in model_without_ddp.parameters()]
+            
+            model_without_ddp.load_state_dict(checkpoint['model_ema2'])
+            model_without_ddp.module.ema_params2 = [p.detach().clone() for p in model_without_ddp.parameters()]
+            
+            # Restore model
+            model_without_ddp.load_state_dict(checkpoint['model'])
+        else:
+            ema_state_dict1 = checkpoint['model_ema1']
+            ema_state_dict2 = checkpoint['model_ema2']
+            model_without_ddp.ema_params1 = [ema_state_dict1[name].cuda() for name, _ in model_without_ddp.named_parameters()]
+            model_without_ddp.ema_params2 = [ema_state_dict2[name].cuda() for name, _ in model_without_ddp.named_parameters()]
+        
+        print("Resumed checkpoint from", args.resume)
+
+        try:
+            if 'optimizer' in checkpoint and 'epoch' in checkpoint:
+                if args.use_fsdp_train:
+                    opt_state = FSDP.optim_state_dict_to_load(
+                        model_without_ddp, optimizer, checkpoint['optimizer']
+                    )
+                    optimizer.load_state_dict(opt_state)
+                else:
+                    optimizer.load_state_dict(checkpoint['optimizer'])
+                print("Loaded optimizer & scaler state!")
+        except:
+            print("Failed to load optimizer & scaler state! Just load checkpoint.")
+        args.start_epoch = checkpoint['epoch'] + 1
+        del checkpoint
+    else:
+        if args.use_fsdp_train:
+            model_without_ddp.module.ema_params1 = [p.detach().clone() for p in model_without_ddp.parameters()]
+            model_without_ddp.module.ema_params2 = [p.detach().clone() for p in model_without_ddp.parameters()]
+        else:
+            model_without_ddp.ema_params1 = [p.detach().clone() for p in model_without_ddp.parameters()]
+            model_without_ddp.ema_params2 = [p.detach().clone() for p in model_without_ddp.parameters()]
+        print("Training from scratch")
+
+    # Evaluate generation
+    if args.evaluate_gen:
+        print("Evaluating checkpoint at {} epoch".format(args.start_epoch))
+        with torch.random.fork_rng():
+            torch.manual_seed(seed)
+            with torch.no_grad():
+                evaluate(model_without_ddp, args, args.start_epoch, batch_size=args.gen_bsz, log_writer=log_writer, vae=vae)
+        return
+
+    # Training loop
+    print(f"Start training for {args.epochs} epochs")
+    start_time = time.time()
+    for epoch in range(args.start_epoch, args.epochs):
+        if args.distributed:
+            data_loader_train.sampler.set_epoch(epoch)
+
+        train_one_epoch(model, model_without_ddp, data_loader_train, optimizer, device, epoch, log_writer=log_writer, args=args, vae=vae)
+
+        # Save checkpoint periodically
+        if epoch % args.save_last_freq == 0 or epoch + 1 == args.epochs:
+            if misc.is_main_process():
+                from grn.utils.safe_rm import safe_remove
+                safe_remove(f'{args.output_dir}/checkpoint-tmp_*.pth', args.output_dir)
+            misc.save_model(
+                args=args,
+                model_without_ddp=model_without_ddp,
+                optimizer=optimizer,
+                epoch=epoch,
+                epoch_name=f"tmp_{epoch}"
+            )
+
+        if epoch % 100 == 0 and epoch > 0:
+            misc.save_model(
+                args=args,
+                model_without_ddp=model_without_ddp,
+                optimizer=optimizer,
+                epoch=epoch
+            )
+
+        # Perform online evaluation at specified intervals
+        if args.online_eval and (epoch % args.eval_freq == 0 or epoch + 1 == args.epochs):
+            torch.cuda.empty_cache()
+            with torch.no_grad():
+                evaluate(model_without_ddp, args, epoch, batch_size=args.gen_bsz, log_writer=log_writer, vae=vae)
+            torch.cuda.empty_cache()
+
+        if misc.is_main_process() and log_writer is not None:
+            log_writer.flush()
+
+    total_time = time.time() - start_time
+    total_time_str = str(datetime.timedelta(seconds=int(total_time)))
+    print('Training time:', total_time_str)
+
+
+if __name__ == '__main__':
+    args = get_args_parser().parse_args()
+    Path(args.output_dir).mkdir(parents=True, exist_ok=True)
+    main(args)

environment.yaml

@@ -0,0 +1,19 @@
+name: grn
+channels:
+  - pytorch
+  - defaults
+  - nvidia
+dependencies:
+  - python=3.10
+  - pip=22.3
+  - pytorch-cuda=12.4
+  - pytorch=2.5.1
+  - torchvision=0.20.1
+  - numpy=1.22
+  - pip:
+    - opencv-python==4.11.0.86
+    - timm==0.9.12
+    - tensorboard==2.10.0
+    - scipy==1.9.1
+    - einops==0.8.1
+    - gdown==5.2.0

evaluation/gen_eval/_base_/datasets/coco_panoptic.py

@@ -0,0 +1,59 @@
+# dataset settings
+dataset_type = 'CocoPanopticDataset'
+data_root = 'data/coco/'
+img_norm_cfg = dict(
+    mean=[123.675, 116.28, 103.53], std=[58.395, 57.12, 57.375], to_rgb=True)
+train_pipeline = [
+    dict(type='LoadImageFromFile'),
+    dict(
+        type='LoadPanopticAnnotations',
+        with_bbox=True,
+        with_mask=True,
+        with_seg=True),
+    dict(type='Resize', img_scale=(1333, 800), keep_ratio=True),
+    dict(type='RandomFlip', flip_ratio=0.5),
+    dict(type='Normalize', **img_norm_cfg),
+    dict(type='Pad', size_divisor=32),
+    dict(type='SegRescale', scale_factor=1 / 4),
+    dict(type='DefaultFormatBundle'),
+    dict(
+        type='Collect',
+        keys=['img', 'gt_bboxes', 'gt_labels', 'gt_masks', 'gt_semantic_seg']),
+]
+test_pipeline = [
+    dict(type='LoadImageFromFile'),
+    dict(
+        type='MultiScaleFlipAug',
+        img_scale=(1333, 800),
+        flip=False,
+        transforms=[
+            dict(type='Resize', keep_ratio=True),
+            dict(type='RandomFlip'),
+            dict(type='Normalize', **img_norm_cfg),
+            dict(type='Pad', size_divisor=32),
+            dict(type='ImageToTensor', keys=['img']),
+            dict(type='Collect', keys=['img']),
+        ])
+]
+data = dict(
+    samples_per_gpu=2,
+    workers_per_gpu=2,
+    train=dict(
+        type=dataset_type,
+        ann_file=data_root + 'annotations/panoptic_train2017.json',
+        img_prefix=data_root + 'train2017/',
+        seg_prefix=data_root + 'annotations/panoptic_train2017/',
+        pipeline=train_pipeline),
+    val=dict(
+        type=dataset_type,
+        ann_file=data_root + 'annotations/panoptic_val2017.json',
+        img_prefix=data_root + 'val2017/',
+        seg_prefix=data_root + 'annotations/panoptic_val2017/',
+        pipeline=test_pipeline),
+    test=dict(
+        type=dataset_type,
+        ann_file=data_root + 'annotations/panoptic_val2017.json',
+        img_prefix=data_root + 'val2017/',
+        seg_prefix=data_root + 'annotations/panoptic_val2017/',
+        pipeline=test_pipeline))
+evaluation = dict(interval=1, metric=['PQ'])

evaluation/gen_eval/_base_/default_runtime.py

@@ -0,0 +1,27 @@
+checkpoint_config = dict(interval=1)
+# yapf:disable
+log_config = dict(
+    interval=50,
+    hooks=[
+        dict(type='TextLoggerHook'),
+        # dict(type='TensorboardLoggerHook')
+    ])
+# yapf:enable
+custom_hooks = [dict(type='NumClassCheckHook')]
+
+dist_params = dict(backend='nccl')
+log_level = 'INFO'
+load_from = None
+resume_from = None
+workflow = [('train', 1)]
+
+# disable opencv multithreading to avoid system being overloaded
+opencv_num_threads = 0
+# set multi-process start method as `fork` to speed up the training
+mp_start_method = 'fork'
+
+# Default setting for scaling LR automatically
+#   - `enable` means enable scaling LR automatically
+#       or not by default.
+#   - `base_batch_size` = (8 GPUs) x (2 samples per GPU).
+auto_scale_lr = dict(enable=False, base_batch_size=16)

evaluation/gen_eval/evaluate_images.py

@@ -0,0 +1,298 @@
+"""
+Evaluate generated images using Mask2Former (or other object detector model)
+"""
+
+import argparse
+import json
+import os
+import re
+import sys
+import time
+import os.path as osp
+
+import warnings
+warnings.filterwarnings("ignore")
+
+import numpy as np
+import tqdm
+import pandas as pd
+from PIL import Image, ImageOps
+import torch
+import mmdet
+from mmdet.apis import inference_detector, init_detector
+
+import open_clip
+from clip_benchmark.metrics import zeroshot_classification as zsc
+zsc.tqdm = lambda it, *args, **kwargs: it
+
+# Get directory path
+
+def parse_args():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("imagedir", type=str)
+    parser.add_argument("--outfile", type=str, default="results.jsonl")
+    parser.add_argument("--model-config", type=str, default="")
+    parser.add_argument("--model-path", type=str, default="")
+    # Other arguments
+    parser.add_argument("--options", nargs="*", type=str, default=[])
+    args = parser.parse_args()
+    args.options = dict(opt.split("=", 1) for opt in args.options)
+    if args.model_config is None:
+        args.model_config = os.path.join(
+            os.path.dirname(mmdet.__file__),
+            "../configs/mask2former/mask2former_swin-s-p4-w7-224_lsj_8x2_50e_coco.py"
+        )
+    return args
+
+DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
+assert DEVICE == "cuda"
+
+def timed(fn):
+    def wrapper(*args, **kwargs):
+        startt = time.time()
+        result = fn(*args, **kwargs)
+        endt = time.time()
+        print(f'Function {fn.__name__!r} executed in {endt - startt:.3f}s', file=sys.stderr)
+        return result
+    return wrapper
+
+# Load models
+
+@timed
+def load_models(args):
+    CONFIG_PATH = osp.abspath(args.model_config)
+    OBJECT_DETECTOR = args.options.get('model', "mask2former_swin-s-p4-w7-224_lsj_8x2_50e_coco")
+    CKPT_PATH = os.path.join(args.model_path, f"{OBJECT_DETECTOR}.pth")
+    object_detector = init_detector(CONFIG_PATH, CKPT_PATH, device=DEVICE)
+
+    clip_arch = args.options.get('clip_model', "ViT-L-14")
+    clip_model, _, transform = open_clip.create_model_and_transforms(clip_arch, pretrained="openai", device=DEVICE)
+    tokenizer = open_clip.get_tokenizer(clip_arch)
+
+    with open(os.path.join(os.path.dirname(__file__), "object_names.txt")) as cls_file:
+        classnames = [line.strip() for line in cls_file]
+
+    return object_detector, (clip_model, transform, tokenizer), classnames
+
+
+COLORS = ["red", "orange", "yellow", "green", "blue", "purple", "pink", "brown", "black", "white"]
+COLOR_CLASSIFIERS = {}
+
+# Evaluation parts
+
+class ImageCrops(torch.utils.data.Dataset):
+    def __init__(self, image: Image.Image, objects):
+        self._image = image.convert("RGB")
+        bgcolor = args.options.get('bgcolor', "#999")
+        if bgcolor == "original":
+            self._blank = self._image.copy()
+        else:
+            self._blank = Image.new("RGB", image.size, color=bgcolor)
+        self._objects = objects
+
+    def __len__(self):
+        return len(self._objects)
+
+    def __getitem__(self, index):
+        box, mask = self._objects[index]
+        if mask is not None:
+            assert tuple(self._image.size[::-1]) == tuple(mask.shape), (index, self._image.size[::-1], mask.shape)
+            image = Image.composite(self._image, self._blank, Image.fromarray(mask))
+        else:
+            image = self._image
+        if args.options.get('crop', '1') == '1':
+            image = image.crop(box[:4])
+        # if args.save:
+        #     base_count = len(os.listdir(args.save))
+        #     image.save(os.path.join(args.save, f"cropped_{base_count:05}.png"))
+        return (transform(image), 0)
+
+
+def color_classification(image, bboxes, classname):
+    if classname not in COLOR_CLASSIFIERS:
+        COLOR_CLASSIFIERS[classname] = zsc.zero_shot_classifier(
+            clip_model, tokenizer, COLORS,
+            [
+                f"a photo of a {{c}} {classname}",
+                f"a photo of a {{c}}-colored {classname}",
+                f"a photo of a {{c}} object"
+            ],
+            DEVICE
+        )
+    clf = COLOR_CLASSIFIERS[classname]
+    dataloader = torch.utils.data.DataLoader(
+        ImageCrops(image, bboxes),
+        batch_size=16, num_workers=4
+    )
+    with torch.no_grad():
+        pred, _ = zsc.run_classification(clip_model, clf, dataloader, DEVICE)
+        return [COLORS[index.item()] for index in pred.argmax(1)]
+
+
+def compute_iou(box_a, box_b):
+    area_fn = lambda box: max(box[2] - box[0] + 1, 0) * max(box[3] - box[1] + 1, 0)
+    i_area = area_fn([
+        max(box_a[0], box_b[0]), max(box_a[1], box_b[1]),
+        min(box_a[2], box_b[2]), min(box_a[3], box_b[3])
+    ])
+    u_area = area_fn(box_a) + area_fn(box_b) - i_area
+    return i_area / u_area if u_area else 0
+
+
+def relative_position(obj_a, obj_b):
+    """Give position of A relative to B, factoring in object dimensions"""
+    boxes = np.array([obj_a[0], obj_b[0]])[:, :4].reshape(2, 2, 2)
+    center_a, center_b = boxes.mean(axis=-2)
+    dim_a, dim_b = np.abs(np.diff(boxes, axis=-2))[..., 0, :]
+    offset = center_a - center_b
+    #
+    revised_offset = np.maximum(np.abs(offset) - POSITION_THRESHOLD * (dim_a + dim_b), 0) * np.sign(offset)
+    if np.all(np.abs(revised_offset) < 1e-3):
+        return set()
+    #
+    dx, dy = revised_offset / np.linalg.norm(offset)
+    relations = set()
+    if dx < -0.5: relations.add("left of")
+    if dx > 0.5: relations.add("right of")
+    if dy < -0.5: relations.add("above")
+    if dy > 0.5: relations.add("below")
+    return relations
+
+
+def evaluate(image, objects, metadata):
+    """
+    Evaluate given image using detected objects on the global metadata specifications.
+    Assumptions:
+    * Metadata combines 'include' clauses with AND, and 'exclude' clauses with OR
+    * All clauses are independent, i.e., duplicating a clause has no effect on the correctness
+    * CHANGED: Color and position will only be evaluated on the most confidently predicted objects;
+        therefore, objects are expected to appear in sorted order
+    """
+    correct = True
+    reason = []
+    matched_groups = []
+    # Check for expected objects
+    for req in metadata.get('include', []):
+        classname = req['class']
+        matched = True
+        found_objects = objects.get(classname, [])[:req['count']]
+        if len(found_objects) < req['count']:
+            correct = matched = False
+            reason.append(f"expected {classname}>={req['count']}, found {len(found_objects)}")
+        else:
+            if 'color' in req:
+                # Color check
+                colors = color_classification(image, found_objects, classname)
+                if colors.count(req['color']) < req['count']:
+                    correct = matched = False
+                    reason.append(
+                        f"expected {req['color']} {classname}>={req['count']}, found " +
+                        f"{colors.count(req['color'])} {req['color']}; and " +
+                        ", ".join(f"{colors.count(c)} {c}" for c in COLORS if c in colors)
+                    )
+            if 'position' in req and matched:
+                # Relative position check
+                expected_rel, target_group = req['position']
+                if matched_groups[target_group] is None:
+                    correct = matched = False
+                    reason.append(f"no target for {classname} to be {expected_rel}")
+                else:
+                    for obj in found_objects:
+                        for target_obj in matched_groups[target_group]:
+                            true_rels = relative_position(obj, target_obj)
+                            if expected_rel not in true_rels:
+                                correct = matched = False
+                                reason.append(
+                                    f"expected {classname} {expected_rel} target, found " +
+                                    f"{' and '.join(true_rels)} target"
+                                )
+                                break
+                        if not matched:
+                            break
+        if matched:
+            matched_groups.append(found_objects)
+        else:
+            matched_groups.append(None)
+    # Check for non-expected objects
+    for req in metadata.get('exclude', []):
+        classname = req['class']
+        if len(objects.get(classname, [])) >= req['count']:
+            correct = False
+            reason.append(f"expected {classname}<{req['count']}, found {len(objects[classname])}")
+    return correct, "\n".join(reason)
+
+
+def evaluate_image(filepath, metadata):
+    result = inference_detector(object_detector, filepath)
+    bbox = result[0] if isinstance(result, tuple) else result
+    segm = result[1] if isinstance(result, tuple) and len(result) > 1 else None
+    image = ImageOps.exif_transpose(Image.open(filepath))
+    detected = {}
+    # Determine bounding boxes to keep
+    confidence_threshold = THRESHOLD if metadata['tag'] != "counting" else COUNTING_THRESHOLD
+    for index, classname in enumerate(classnames):
+        ordering = np.argsort(bbox[index][:, 4])[::-1]
+        ordering = ordering[bbox[index][ordering, 4] > confidence_threshold] # Threshold
+        ordering = ordering[:MAX_OBJECTS].tolist() # Limit number of detected objects per class
+        detected[classname] = []
+        while ordering:
+            max_obj = ordering.pop(0)
+            detected[classname].append((bbox[index][max_obj], None if segm is None else segm[index][max_obj]))
+            ordering = [
+                obj for obj in ordering
+                if NMS_THRESHOLD == 1 or compute_iou(bbox[index][max_obj], bbox[index][obj]) < NMS_THRESHOLD
+            ]
+        if not detected[classname]:
+            del detected[classname]
+    # Evaluate
+    is_correct, reason = evaluate(image, detected, metadata)
+    return {
+        'filename': filepath,
+        'tag': metadata['tag'],
+        'prompt': metadata['prompt'],
+        'correct': is_correct,
+        'reason': reason,
+        'metadata': json.dumps(metadata),
+        'details': json.dumps({
+            key: [box.tolist() for box, _ in value]
+            for key, value in detected.items()
+        })
+    }
+
+
+def main(args):
+    full_results = []
+    pbar = tqdm.tqdm(total=len(os.listdir(args.imagedir)))
+    for subfolder in os.listdir(args.imagedir):
+        pbar.update(1)
+        folderpath = os.path.join(args.imagedir, subfolder)
+        if not os.path.isdir(folderpath) or not subfolder.isdigit():
+            print('skip 269')
+            continue
+        with open(os.path.join(folderpath, "metadata.jsonl")) as fp:
+            metadata = json.load(fp)
+        # Evaluate each image
+        for imagename in os.listdir(os.path.join(folderpath, "samples")):
+            imagepath = os.path.join(folderpath, "samples", imagename)
+            if not os.path.isfile(imagepath) or not (re.match(r"\d+\.png", imagename) or re.match(r"\d+\.jpg", imagename)):
+                print('skip 276')
+                continue
+            result = evaluate_image(imagepath, metadata)
+            full_results.append(result)
+    # Save results
+    if os.path.dirname(args.outfile):
+        os.makedirs(os.path.dirname(args.outfile), exist_ok=True)
+    with open(args.outfile, "w") as fp:
+        pd.DataFrame(full_results).to_json(fp, orient="records", lines=True)
+
+
+if __name__ == "__main__":
+    args = parse_args()
+    object_detector, (clip_model, transform, tokenizer), classnames = load_models(args)
+    THRESHOLD = float(args.options.get('threshold', 0.3))
+    COUNTING_THRESHOLD = float(args.options.get('counting_threshold', 0.9))
+    MAX_OBJECTS = int(args.options.get('max_objects', 16))
+    NMS_THRESHOLD = float(args.options.get('max_overlap', 1.0))
+    POSITION_THRESHOLD = float(args.options.get('position_threshold', 0.1))
+
+    main(args)

evaluation/gen_eval/mask2former/mask2former_r50_lsj_8x2_50e_coco-panoptic.py

@@ -0,0 +1,253 @@
+_base_ = [
+    '../_base_/datasets/coco_panoptic.py', '../_base_/default_runtime.py'
+]
+num_things_classes = 80
+num_stuff_classes = 53
+num_classes = num_things_classes + num_stuff_classes
+model = dict(
+    type='Mask2Former',
+    backbone=dict(
+        type='ResNet',
+        depth=50,
+        num_stages=4,
+        out_indices=(0, 1, 2, 3),
+        frozen_stages=-1,
+        norm_cfg=dict(type='BN', requires_grad=False),
+        norm_eval=True,
+        style='pytorch',
+        init_cfg=dict(type='Pretrained', checkpoint='torchvision://resnet50')),
+    panoptic_head=dict(
+        type='Mask2FormerHead',
+        in_channels=[256, 512, 1024, 2048],  # pass to pixel_decoder inside
+        strides=[4, 8, 16, 32],
+        feat_channels=256,
+        out_channels=256,
+        num_things_classes=num_things_classes,
+        num_stuff_classes=num_stuff_classes,
+        num_queries=100,
+        num_transformer_feat_level=3,
+        pixel_decoder=dict(
+            type='MSDeformAttnPixelDecoder',
+            num_outs=3,
+            norm_cfg=dict(type='GN', num_groups=32),
+            act_cfg=dict(type='ReLU'),
+            encoder=dict(
+                type='DetrTransformerEncoder',
+                num_layers=6,
+                transformerlayers=dict(
+                    type='BaseTransformerLayer',
+                    attn_cfgs=dict(
+                        type='MultiScaleDeformableAttention',
+                        embed_dims=256,
+                        num_heads=8,
+                        num_levels=3,
+                        num_points=4,
+                        im2col_step=64,
+                        dropout=0.0,
+                        batch_first=False,
+                        norm_cfg=None,
+                        init_cfg=None),
+                    ffn_cfgs=dict(
+                        type='FFN',
+                        embed_dims=256,
+                        feedforward_channels=1024,
+                        num_fcs=2,
+                        ffn_drop=0.0,
+                        act_cfg=dict(type='ReLU', inplace=True)),
+                    operation_order=('self_attn', 'norm', 'ffn', 'norm')),
+                init_cfg=None),
+            positional_encoding=dict(
+                type='SinePositionalEncoding', num_feats=128, normalize=True),
+            init_cfg=None),
+        enforce_decoder_input_project=False,
+        positional_encoding=dict(
+            type='SinePositionalEncoding', num_feats=128, normalize=True),
+        transformer_decoder=dict(
+            type='DetrTransformerDecoder',
+            return_intermediate=True,
+            num_layers=9,
+            transformerlayers=dict(
+                type='DetrTransformerDecoderLayer',
+                attn_cfgs=dict(
+                    type='MultiheadAttention',
+                    embed_dims=256,
+                    num_heads=8,
+                    attn_drop=0.0,
+                    proj_drop=0.0,
+                    dropout_layer=None,
+                    batch_first=False),
+                ffn_cfgs=dict(
+                    embed_dims=256,
+                    feedforward_channels=2048,
+                    num_fcs=2,
+                    act_cfg=dict(type='ReLU', inplace=True),
+                    ffn_drop=0.0,
+                    dropout_layer=None,
+                    add_identity=True),
+                feedforward_channels=2048,
+                operation_order=('cross_attn', 'norm', 'self_attn', 'norm',
+                                 'ffn', 'norm')),
+            init_cfg=None),
+        loss_cls=dict(
+            type='CrossEntropyLoss',
+            use_sigmoid=False,
+            loss_weight=2.0,
+            reduction='mean',
+            class_weight=[1.0] * num_classes + [0.1]),
+        loss_mask=dict(
+            type='CrossEntropyLoss',
+            use_sigmoid=True,
+            reduction='mean',
+            loss_weight=5.0),
+        loss_dice=dict(
+            type='DiceLoss',
+            use_sigmoid=True,
+            activate=True,
+            reduction='mean',
+            naive_dice=True,
+            eps=1.0,
+            loss_weight=5.0)),
+    panoptic_fusion_head=dict(
+        type='MaskFormerFusionHead',
+        num_things_classes=num_things_classes,
+        num_stuff_classes=num_stuff_classes,
+        loss_panoptic=None,
+        init_cfg=None),
+    train_cfg=dict(
+        num_points=12544,
+        oversample_ratio=3.0,
+        importance_sample_ratio=0.75,
+        assigner=dict(
+            type='MaskHungarianAssigner',
+            cls_cost=dict(type='ClassificationCost', weight=2.0),
+            mask_cost=dict(
+                type='CrossEntropyLossCost', weight=5.0, use_sigmoid=True),
+            dice_cost=dict(
+                type='DiceCost', weight=5.0, pred_act=True, eps=1.0)),
+        sampler=dict(type='MaskPseudoSampler')),
+    test_cfg=dict(
+        panoptic_on=True,
+        # For now, the dataset does not support
+        # evaluating semantic segmentation metric.
+        semantic_on=False,
+        instance_on=True,
+        # max_per_image is for instance segmentation.
+        max_per_image=100,
+        iou_thr=0.8,
+        # In Mask2Former's panoptic postprocessing,
+        # it will filter mask area where score is less than 0.5 .
+        filter_low_score=True),
+    init_cfg=None)
+
+# dataset settings
+image_size = (1024, 1024)
+img_norm_cfg = dict(
+    mean=[123.675, 116.28, 103.53], std=[58.395, 57.12, 57.375], to_rgb=True)
+train_pipeline = [
+    dict(type='LoadImageFromFile', to_float32=True),
+    dict(
+        type='LoadPanopticAnnotations',
+        with_bbox=True,
+        with_mask=True,
+        with_seg=True),
+    dict(type='RandomFlip', flip_ratio=0.5),
+    # large scale jittering
+    dict(
+        type='Resize',
+        img_scale=image_size,
+        ratio_range=(0.1, 2.0),
+        multiscale_mode='range',
+        keep_ratio=True),
+    dict(
+        type='RandomCrop',
+        crop_size=image_size,
+        crop_type='absolute',
+        recompute_bbox=True,
+        allow_negative_crop=True),
+    dict(type='Normalize', **img_norm_cfg),
+    dict(type='Pad', size=image_size),
+    dict(type='DefaultFormatBundle', img_to_float=True),
+    dict(
+        type='Collect',
+        keys=['img', 'gt_bboxes', 'gt_labels', 'gt_masks', 'gt_semantic_seg']),
+]
+test_pipeline = [
+    dict(type='LoadImageFromFile'),
+    dict(
+        type='MultiScaleFlipAug',
+        img_scale=(1333, 800),
+        flip=False,
+        transforms=[
+            dict(type='Resize', keep_ratio=True),
+            dict(type='RandomFlip'),
+            dict(type='Normalize', **img_norm_cfg),
+            dict(type='Pad', size_divisor=32),
+            dict(type='ImageToTensor', keys=['img']),
+            dict(type='Collect', keys=['img']),
+        ])
+]
+data_root = 'data/coco/'
+data = dict(
+    samples_per_gpu=2,
+    workers_per_gpu=2,
+    train=dict(pipeline=train_pipeline),
+    val=dict(
+        pipeline=test_pipeline,
+        ins_ann_file=data_root + 'annotations/instances_val2017.json',
+    ),
+    test=dict(
+        pipeline=test_pipeline,
+        ins_ann_file=data_root + 'annotations/instances_val2017.json',
+    ))
+
+embed_multi = dict(lr_mult=1.0, decay_mult=0.0)
+# optimizer
+optimizer = dict(
+    type='AdamW',
+    lr=0.0001,
+    weight_decay=0.05,
+    eps=1e-8,
+    betas=(0.9, 0.999),
+    paramwise_cfg=dict(
+        custom_keys={
+            'backbone': dict(lr_mult=0.1, decay_mult=1.0),
+            'query_embed': embed_multi,
+            'query_feat': embed_multi,
+            'level_embed': embed_multi,
+        },
+        norm_decay_mult=0.0))
+optimizer_config = dict(grad_clip=dict(max_norm=0.01, norm_type=2))
+
+# learning policy
+lr_config = dict(
+    policy='step',
+    gamma=0.1,
+    by_epoch=False,
+    step=[327778, 355092],
+    warmup='linear',
+    warmup_by_epoch=False,
+    warmup_ratio=1.0,  # no warmup
+    warmup_iters=10)
+
+max_iters = 368750
+runner = dict(type='IterBasedRunner', max_iters=max_iters)
+
+log_config = dict(
+    interval=50,
+    hooks=[
+        dict(type='TextLoggerHook', by_epoch=False),
+        dict(type='TensorboardLoggerHook', by_epoch=False)
+    ])
+interval = 5000
+workflow = [('train', interval)]
+checkpoint_config = dict(
+    by_epoch=False, interval=interval, save_last=True, max_keep_ckpts=3)
+
+# Before 365001th iteration, we do evaluation every 5000 iterations.
+# After 365000th iteration, we do evaluation every 368750 iterations,
+# which means that we do evaluation at the end of training.
+dynamic_intervals = [(max_iters // interval * interval + 1, max_iters)]
+evaluation = dict(
+    interval=interval,
+    dynamic_intervals=dynamic_intervals,
+    metric=['PQ', 'bbox', 'segm'])

evaluation/gen_eval/mask2former/mask2former_r50_lsj_8x2_50e_coco.py

@@ -0,0 +1,79 @@
+_base_ = ['./mask2former_r50_lsj_8x2_50e_coco-panoptic.py']
+num_things_classes = 80
+num_stuff_classes = 0
+num_classes = num_things_classes + num_stuff_classes
+model = dict(
+    panoptic_head=dict(
+        num_things_classes=num_things_classes,
+        num_stuff_classes=num_stuff_classes,
+        loss_cls=dict(class_weight=[1.0] * num_classes + [0.1])),
+    panoptic_fusion_head=dict(
+        num_things_classes=num_things_classes,
+        num_stuff_classes=num_stuff_classes),
+    test_cfg=dict(panoptic_on=False))
+
+# dataset settings
+image_size = (1024, 1024)
+img_norm_cfg = dict(
+    mean=[123.675, 116.28, 103.53], std=[58.395, 57.12, 57.375], to_rgb=True)
+pad_cfg = dict(img=(128, 128, 128), masks=0, seg=255)
+train_pipeline = [
+    dict(type='LoadImageFromFile', to_float32=True),
+    dict(type='LoadAnnotations', with_bbox=True, with_mask=True),
+    dict(type='RandomFlip', flip_ratio=0.5),
+    # large scale jittering
+    dict(
+        type='Resize',
+        img_scale=image_size,
+        ratio_range=(0.1, 2.0),
+        multiscale_mode='range',
+        keep_ratio=True),
+    dict(
+        type='RandomCrop',
+        crop_size=image_size,
+        crop_type='absolute',
+        recompute_bbox=True,
+        allow_negative_crop=True),
+    dict(type='FilterAnnotations', min_gt_bbox_wh=(1e-5, 1e-5), by_mask=True),
+    dict(type='Pad', size=image_size, pad_val=pad_cfg),
+    dict(type='Normalize', **img_norm_cfg),
+    dict(type='DefaultFormatBundle', img_to_float=True),
+    dict(type='Collect', keys=['img', 'gt_bboxes', 'gt_labels', 'gt_masks']),
+]
+test_pipeline = [
+    dict(type='LoadImageFromFile'),
+    dict(
+        type='MultiScaleFlipAug',
+        img_scale=(1333, 800),
+        flip=False,
+        transforms=[
+            dict(type='Resize', keep_ratio=True),
+            dict(type='RandomFlip'),
+            dict(type='Pad', size_divisor=32, pad_val=pad_cfg),
+            dict(type='Normalize', **img_norm_cfg),
+            dict(type='ImageToTensor', keys=['img']),
+            dict(type='Collect', keys=['img']),
+        ])
+]
+dataset_type = 'CocoDataset'
+data_root = 'data/coco/'
+data = dict(
+    _delete_=True,
+    samples_per_gpu=2,
+    workers_per_gpu=2,
+    train=dict(
+        type=dataset_type,
+        ann_file=data_root + 'annotations/instances_train2017.json',
+        img_prefix=data_root + 'train2017/',
+        pipeline=train_pipeline),
+    val=dict(
+        type=dataset_type,
+        ann_file=data_root + 'annotations/instances_val2017.json',
+        img_prefix=data_root + 'val2017/',
+        pipeline=test_pipeline),
+    test=dict(
+        type=dataset_type,
+        ann_file=data_root + 'annotations/instances_val2017.json',
+        img_prefix=data_root + 'val2017/',
+        pipeline=test_pipeline))
+evaluation = dict(metric=['bbox', 'segm'])

evaluation/gen_eval/mask2former/mask2former_swin-s-p4-w7-224_lsj_8x2_50e_coco.py

@@ -0,0 +1,37 @@
+_base_ = ['./mask2former_swin-t-p4-w7-224_lsj_8x2_50e_coco.py']
+pretrained = 'https://github.com/SwinTransformer/storage/releases/download/v1.0.0/swin_small_patch4_window7_224.pth'  # noqa
+
+depths = [2, 2, 18, 2]
+model = dict(
+    backbone=dict(
+        depths=depths, init_cfg=dict(type='Pretrained',
+                                     checkpoint=pretrained)))
+
+# set all layers in backbone to lr_mult=0.1
+# set all norm layers, position_embeding,
+# query_embeding, level_embeding to decay_multi=0.0
+backbone_norm_multi = dict(lr_mult=0.1, decay_mult=0.0)
+backbone_embed_multi = dict(lr_mult=0.1, decay_mult=0.0)
+embed_multi = dict(lr_mult=1.0, decay_mult=0.0)
+custom_keys = {
+    'backbone': dict(lr_mult=0.1, decay_mult=1.0),
+    'backbone.patch_embed.norm': backbone_norm_multi,
+    'backbone.norm': backbone_norm_multi,
+    'absolute_pos_embed': backbone_embed_multi,
+    'relative_position_bias_table': backbone_embed_multi,
+    'query_embed': embed_multi,
+    'query_feat': embed_multi,
+    'level_embed': embed_multi
+}
+custom_keys.update({
+    f'backbone.stages.{stage_id}.blocks.{block_id}.norm': backbone_norm_multi
+    for stage_id, num_blocks in enumerate(depths)
+    for block_id in range(num_blocks)
+})
+custom_keys.update({
+    f'backbone.stages.{stage_id}.downsample.norm': backbone_norm_multi
+    for stage_id in range(len(depths) - 1)
+})
+# optimizer
+optimizer = dict(
+    paramwise_cfg=dict(custom_keys=custom_keys, norm_decay_mult=0.0))

evaluation/gen_eval/mask2former/mask2former_swin-t-p4-w7-224_lsj_8x2_50e_coco.py

@@ -0,0 +1,61 @@
+_base_ = ['./mask2former_r50_lsj_8x2_50e_coco.py']
+pretrained = 'https://github.com/SwinTransformer/storage/releases/download/v1.0.0/swin_tiny_patch4_window7_224.pth'  # noqa
+depths = [2, 2, 6, 2]
+model = dict(
+    type='Mask2Former',
+    backbone=dict(
+        _delete_=True,
+        type='SwinTransformer',
+        embed_dims=96,
+        depths=depths,
+        num_heads=[3, 6, 12, 24],
+        window_size=7,
+        mlp_ratio=4,
+        qkv_bias=True,
+        qk_scale=None,
+        drop_rate=0.,
+        attn_drop_rate=0.,
+        drop_path_rate=0.3,
+        patch_norm=True,
+        out_indices=(0, 1, 2, 3),
+        with_cp=False,
+        convert_weights=True,
+        frozen_stages=-1,
+        init_cfg=dict(type='Pretrained', checkpoint=pretrained)),
+    panoptic_head=dict(
+        type='Mask2FormerHead', in_channels=[96, 192, 384, 768]),
+    init_cfg=None)
+
+# set all layers in backbone to lr_mult=0.1
+# set all norm layers, position_embeding,
+# query_embeding, level_embeding to decay_multi=0.0
+backbone_norm_multi = dict(lr_mult=0.1, decay_mult=0.0)
+backbone_embed_multi = dict(lr_mult=0.1, decay_mult=0.0)
+embed_multi = dict(lr_mult=1.0, decay_mult=0.0)
+custom_keys = {
+    'backbone': dict(lr_mult=0.1, decay_mult=1.0),
+    'backbone.patch_embed.norm': backbone_norm_multi,
+    'backbone.norm': backbone_norm_multi,
+    'absolute_pos_embed': backbone_embed_multi,
+    'relative_position_bias_table': backbone_embed_multi,
+    'query_embed': embed_multi,
+    'query_feat': embed_multi,
+    'level_embed': embed_multi
+}
+custom_keys.update({
+    f'backbone.stages.{stage_id}.blocks.{block_id}.norm': backbone_norm_multi
+    for stage_id, num_blocks in enumerate(depths)
+    for block_id in range(num_blocks)
+})
+custom_keys.update({
+    f'backbone.stages.{stage_id}.downsample.norm': backbone_norm_multi
+    for stage_id in range(len(depths) - 1)
+})
+# optimizer
+optimizer = dict(
+    type='AdamW',
+    lr=0.0001,
+    weight_decay=0.05,
+    eps=1e-8,
+    betas=(0.9, 0.999),
+    paramwise_cfg=dict(custom_keys=custom_keys, norm_decay_mult=0.0))

evaluation/gen_eval/prompts/create_prompts.py

@@ -0,0 +1,183 @@
+"""
+Generate prompts for evaluation
+"""
+
+import argparse
+import json
+import os
+import yaml
+
+import numpy as np
+
+# Load classnames
+
+with open("object_names.txt") as cls_file:
+    classnames = [line.strip() for line in cls_file]
+
+# Proper a vs an
+
+def with_article(name: str):
+    if name[0] in "aeiou":
+        return f"an {name}"
+    return f"a {name}"
+
+# Proper plural
+
+def make_plural(name: str):
+    if name[-1] in "s":
+        return f"{name}es"
+    return f"{name}s"
+
+# Generates single object samples
+
+def generate_single_object_sample(rng: np.random.Generator, size: int = None):
+    TAG = "single_object"
+    if size > len(classnames):
+        size = len(classnames)
+        print(f"Not enough distinct classes, generating only {size} samples")
+    return_scalar = size is None
+    size = size or 1
+    idxs = rng.choice(len(classnames), size=size, replace=False)
+    samples = [dict(
+        tag=TAG,
+        include=[
+            {"class": classnames[idx], "count": 1}
+        ],
+        prompt=f"a photo of {with_article(classnames[idx])}"
+    ) for idx in idxs]
+    if return_scalar:
+        return samples[0]
+    return samples
+
+# Generate two object samples
+
+def generate_two_object_sample(rng: np.random.Generator):
+    TAG = "two_object"
+    idx_a, idx_b = rng.choice(len(classnames), size=2, replace=False)
+    return dict(
+        tag=TAG,
+        include=[
+            {"class": classnames[idx_a], "count": 1},
+            {"class": classnames[idx_b], "count": 1}
+        ],
+        prompt=f"a photo of {with_article(classnames[idx_a])} and {with_article(classnames[idx_b])}"
+    )
+
+# Generate counting samples
+
+numbers = ["zero", "one", "two", "three", "four", "five", "six", "seven", "eight", "nine", "ten"]
+
+def generate_counting_sample(rng: np.random.Generator, max_count=4):
+    TAG = "counting"
+    idx = rng.choice(len(classnames))
+    num = int(rng.integers(2, max_count, endpoint=True))
+    return dict(
+        tag=TAG,
+        include=[
+            {"class": classnames[idx], "count": num}
+        ],
+        exclude=[
+            {"class": classnames[idx], "count": num + 1}
+        ],
+        prompt=f"a photo of {numbers[num]} {make_plural(classnames[idx])}"
+    )
+
+# Generate color samples
+
+colors = ["red", "orange", "yellow", "green", "blue", "purple", "pink", "brown", "black", "white"]
+
+def generate_color_sample(rng: np.random.Generator):
+    TAG = "colors"
+    idx = rng.choice(len(classnames) - 1) + 1
+    idx = (idx + classnames.index("person")) % len(classnames) # No "[COLOR] person" prompts
+    color = colors[rng.choice(len(colors))]
+    return dict(
+        tag=TAG,
+        include=[
+            {"class": classnames[idx], "count": 1, "color": color}
+        ],
+        prompt=f"a photo of {with_article(color)} {classnames[idx]}"
+    )
+
+# Generate position samples
+
+positions = ["left of", "right of", "above", "below"]
+
+def generate_position_sample(rng: np.random.Generator):
+    TAG = "position"
+    idx_a, idx_b = rng.choice(len(classnames), size=2, replace=False)
+    position = positions[rng.choice(len(positions))]
+    return dict(
+        tag=TAG,
+        include=[
+            {"class": classnames[idx_b], "count": 1},
+            {"class": classnames[idx_a], "count": 1, "position": (position, 0)}
+        ],
+        prompt=f"a photo of {with_article(classnames[idx_a])} {position} {with_article(classnames[idx_b])}"
+    )
+
+# Generate color attribution samples
+
+def generate_color_attribution_sample(rng: np.random.Generator):
+    TAG = "color_attr"
+    idxs = rng.choice(len(classnames) - 1, size=2, replace=False) + 1
+    idx_a, idx_b = (idxs + classnames.index("person")) % len(classnames) # No "[COLOR] person" prompts
+    cidx_a, cidx_b = rng.choice(len(colors), size=2, replace=False)
+    return dict(
+        tag=TAG,
+        include=[
+            {"class": classnames[idx_a], "count": 1, "color": colors[cidx_a]},
+            {"class": classnames[idx_b], "count": 1, "color": colors[cidx_b]}
+        ],
+        prompt=f"a photo of {with_article(colors[cidx_a])} {classnames[idx_a]} and {with_article(colors[cidx_b])} {classnames[idx_b]}"
+    )
+
+
+# Generate evaluation suite
+
+def generate_suite(rng: np.random.Generator, n: int = 100, output_path: str = ""):
+    samples = []
+    # Generate single object samples for all COCO classnames
+    samples.extend(generate_single_object_sample(rng, size=len(classnames)))
+    # Generate two object samples (~100)
+    for _ in range(n):
+        samples.append(generate_two_object_sample(rng))
+    # Generate counting samples
+    for _ in range(n):
+        samples.append(generate_counting_sample(rng, max_count=4))
+    # Generate color samples
+    for _ in range(n):
+        samples.append(generate_color_sample(rng))
+    # Generate position samples
+    for _ in range(n):
+        samples.append(generate_position_sample(rng))
+    # Generate color attribution samples
+    for _ in range(n):
+        samples.append(generate_color_attribution_sample(rng))
+    # De-duplicate
+    unique_samples, used_samples = [], set()
+    for sample in samples:
+        sample_text = yaml.safe_dump(sample)
+        if sample_text not in used_samples:
+            unique_samples.append(sample)
+            used_samples.add(sample_text)
+
+    # Write to files
+    os.makedirs(output_path, exist_ok=True)
+    with open(os.path.join(output_path, "generation_prompts.txt"), "w") as fp:
+        for sample in unique_samples:
+            print(sample['prompt'], file=fp)
+    with open(os.path.join(output_path, "evaluation_metadata.jsonl"), "w") as fp:
+        for sample in unique_samples:
+            print(json.dumps(sample), file=fp)
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--seed", type=int, default=43, help="generation seed (default: 43)")
+    parser.add_argument("--num-prompts", "-n", type=int, default=100, help="number of prompts per task (default: 100)")
+    parser.add_argument("--output-path", "-o", type=str, default="prompts", help="output folder for prompts and metadata (default: 'prompts/')")
+    args = parser.parse_args()
+    rng = np.random.default_rng(args.seed)
+    generate_suite(rng, args.num_prompts, args.output_path)
+

evaluation/gen_eval/summary_scores.py

@@ -0,0 +1,45 @@
+# Get results of evaluation
+
+import argparse
+import os
+
+import numpy as np
+import pandas as pd
+
+
+parser = argparse.ArgumentParser()
+parser.add_argument("filename", type=str)
+args = parser.parse_args()
+
+# Load classnames
+
+with open(os.path.join(os.path.dirname(__file__), "object_names.txt")) as cls_file:
+    classnames = [line.strip() for line in cls_file]
+    cls_to_idx = {"_".join(cls.split()):idx for idx, cls in enumerate(classnames)}
+
+# Load results
+
+df = pd.read_json(args.filename, orient="records", lines=True)
+
+# Measure overall success
+
+print("Summary")
+print("=======")
+print(f"Total images: {len(df)}")
+print(f"Total prompts: {len(df.groupby('metadata'))}")
+print(f"% correct images: {df['correct'].mean():.2%}")
+print(f"% correct prompts: {df.groupby('metadata')['correct'].any().mean():.2%}")
+print()
+
+# By group
+
+task_scores = []
+
+print("Task breakdown")
+print("==============")
+for tag, task_df in df.groupby('tag', sort=False):
+    task_scores.append(task_df['correct'].mean())
+    print(f"{tag:<16} = {task_df['correct'].mean():.2%} ({task_df['correct'].sum()} / {len(task_df)})")
+print()
+
+print(f"Overall score (avg. over tasks): {np.mean(task_scores):.5f}")

grn/dataset/build.py

@@ -0,0 +1,150 @@
+import datetime
+import os
+import os.path as osp
+import random
+import subprocess
+from functools import partial
+from typing import Optional
+import time
+
+import pytz
+
+from grn.dataset.dataset_joint_vi import JointViDataset
+from grn.utils_t2iv.sequence_parallel import SequenceParallelManager as sp_manager
+
+try:
+    from grp import getgrgid
+    from pwd import getpwuid
+except:
+    pass
+import PIL.Image as PImage
+from PIL import ImageFile
+import numpy as np
+from torchvision.transforms import transforms
+from torchvision.transforms.functional import resize, to_tensor
+import torch.distributed as tdist
+
+from torchvision.transforms import InterpolationMode
+bicubic = InterpolationMode.BICUBIC
+lanczos = InterpolationMode.LANCZOS
+PImage.MAX_IMAGE_PIXELS = (1024 * 1024 * 1024 // 4 // 3) * 5
+ImageFile.LOAD_TRUNCATED_IMAGES = False
+
+
+def time_str(fmt='[%m-%d %H:%M:%S]'):
+    return datetime.datetime.now(tz=pytz.timezone('Asia/Shanghai')).strftime(fmt)
+
+
+def normalize_01_into_pm1(x):  # normalize x from [0, 1] to [-1, 1] by (x*2) - 1
+    return x.add(x).add_(-1)
+
+
+def denormalize_pm1_into_01(x):  # denormalize x from [-1, 1] to [0, 1]
+    return x.add(1).mul_(0.5)
+
+
+def center_crop_arr(pil_image, image_size):
+    """
+    Center cropping implementation from ADM.
+    https://github.com/openai/guided-diffusion/blob/8fb3ad9197f16bbc40620447b2742e13458d2831/guided_diffusion/image_datasets.py#L126
+    """
+    while min(*pil_image.size) >= 2 * image_size:
+        pil_image = pil_image.resize(
+            tuple(x // 2 for x in pil_image.size), resample=PImage.BOX
+        )
+    
+    scale = image_size / min(*pil_image.size)
+    pil_image = pil_image.resize(
+        tuple(round(x * scale) for x in pil_image.size), resample=PImage.LANCZOS
+    )
+    
+    arr = np.array(pil_image)
+    crop_y = (arr.shape[0] - image_size) // 2
+    crop_x = (arr.shape[1] - image_size) // 2
+    return PImage.fromarray(arr[crop_y: crop_y + image_size, crop_x: crop_x + image_size])
+
+
+class RandomResize:
+    def __init__(self, mid_reso, final_reso, interpolation):
+        ub = max(round((mid_reso + (mid_reso-final_reso) / 8) / 4) * 4, mid_reso)
+        self.reso_lb, self.reso_ub = final_reso, ub
+        self.interpolation = interpolation
+    
+    def __call__(self, img):
+        return resize(img, size=random.randint(self.reso_lb, self.reso_ub), interpolation=self.interpolation)
+    
+    def __repr__(self):
+        return f'RandomResize(reso=({self.reso_lb}, {self.reso_ub}), interpolation={self.interpolation})'
+
+
+def print_aug(transform, label):
+    print(f'Transform {label} = ')
+    if hasattr(transform, 'transforms'):
+        for t in transform.transforms:
+            print(t)
+    else:
+        print(transform)
+    print('---------------------------\n')
+
+
+def build_joint_dataset(
+    args,
+    meta_folders: str,
+    meta_folder_repeats: str,
+    max_caption_len: int,
+    short_prob=0.2,
+    load_vae_instead_of_image=False
+):
+    return JointViDataset(
+        meta_folders=meta_folders, 
+        meta_folder_repeats=meta_folder_repeats,
+        max_caption_len=max_caption_len,
+        short_prob=short_prob, 
+        load_vae_instead_of_image=load_vae_instead_of_image, 
+        video_fps=args.video_fps,
+        num_frames=args.video_frames,
+        online_t5=args.online_t5,
+        num_replicas=sp_manager.get_sp_group_nums() if sp_manager.sp_on() else tdist.get_world_size(), # 1,
+        rank = sp_manager.get_sp_group_rank() if sp_manager.sp_on() else tdist.get_rank(),
+        dataloader_workers=args.workers,
+        enable_dynamic_length_prompt=args.enable_dynamic_length_prompt,
+        hdfs_mode=args.hdfs_mode,
+        dynamic_scale_schedule=args.dynamic_scale_schedule,
+        seed=args.seed,
+        other_args=args,
+    ) 
+
+
+def pil_load(path: str, proposal_size):
+    with open(path, 'rb') as f:
+        img: PImage.Image = PImage.open(f)
+        w: int = img.width
+        h: int = img.height
+        sh: int = min(h, w)
+        if sh > proposal_size:
+            ratio: float = proposal_size / sh
+            w = round(ratio * w)
+            h = round(ratio * h)
+        img.draft('RGB', (w, h))
+        img = img.convert('RGB')
+    return img
+
+
+def rewrite(im: PImage, file: str, info: str):
+    kw = dict(quality=100)
+    if file.lower().endswith('.tif') or file.lower().endswith('.tiff'):
+        kw['compression'] = 'none'
+    elif file.lower().endswith('.webp'):
+        kw['lossless'] = True
+    
+    st = os.stat(file)
+    uname = getpwuid(st.st_uid).pw_name
+    gname = getgrgid(st.st_gid).gr_name
+    mode = oct(st.st_mode)[-3:]
+    
+    local_file = osp.basename(file)
+    im.save(local_file, **kw)
+    print(f'************* <REWRITE: {info}> *************  @  {file}')
+    subprocess.call(['sudo', 'mv', local_file, file])
+    subprocess.call(['sudo', 'chown', f'{uname}:{gname}', file])
+    subprocess.call(['sudo', 'chmod', str(mode), file])

grn/dataset/dataset_joint_vi.py

@@ -0,0 +1,687 @@
+import glob
+import os
+import pickle
+import random
+import re
+import time
+from functools import partial
+from os import path as osp
+from typing import List, Tuple, Union
+import json
+import itertools
+import hashlib
+import copy
+import collections
+import math
+
+import tqdm
+import numpy as np
+import torch
+import pandas as pd
+from decord import VideoReader
+from PIL import Image as PImage
+from torch.nn import functional as F
+from torchvision.transforms.functional import to_tensor, hflip
+from torchvision.transforms import transforms, InterpolationMode
+from torch.utils.data import Dataset, DataLoader
+import torch.distributed as tdist
+from PIL import Image
+os.environ["TOKENIZERS_PARALLELISM"] = "false"
+
+from grn.schedules.dynamic_resolution import get_dynamic_resolution_meta
+from grn.utils.video_decoder import EncodedVideoDecord
+from grn.utils.compress_tokens import load_packed_tensor
+from transformers import AutoTokenizer
+
+def transform(pil_img, tgt_h, tgt_w):
+    width, height = pil_img.size
+    if width / height <= tgt_w / tgt_h:
+        resized_width = tgt_w
+        resized_height = int(tgt_w / (width / height))
+    else:
+        resized_height = tgt_h
+        resized_width = int((width / height) * tgt_h)
+    pil_img = pil_img.resize((resized_width, resized_height), resample=PImage.LANCZOS)
+    # crop the center out
+    arr = np.array(pil_img)
+    crop_y = (arr.shape[0] - tgt_h) // 2
+    crop_x = (arr.shape[1] - tgt_w) // 2
+    im = to_tensor(arr[crop_y: crop_y + tgt_h, crop_x: crop_x + tgt_w])
+    return im.add(im).add_(-1)
+
+def normalize(x):  # normalize x from [0, 1] to [-1, 1] by (x*2) - 1
+    return x.add(x).add_(-1)
+
+def get_prompt_id(prompt):
+    md5 = hashlib.md5()
+    md5.update(prompt.encode('utf-8'))
+    prompt_id = md5.hexdigest()
+    return prompt_id
+
+def prepend_motion_score(prompt, motion_score):
+    return f'<<<motion_score: {round(motion_score):.1f}>>> {prompt}'
+
+class VideoReaderWrapper(VideoReader):
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.seek(0)
+    def __getitem__(self, key):
+        frames = super().__getitem__(key)
+        self.seek(0)
+        return frames
+
+class JointViDataset(Dataset):
+    def __init__(
+        self,
+        meta_folders: str = '',
+        meta_folder_repeats: str = '',
+        buffersize: int = 1000000 * 300,
+        seed: int = 0,
+        pn: str = '',
+        video_fps: int = 1,
+        num_replicas: int = 1,
+        rank: int = 0,
+        dataloader_workers: int = 2,
+        enable_dynamic_length_prompt: bool = True,
+        shuffle: bool = True,
+        short_prob: float = 0.2,
+        verbose=False,
+        temp_dir= "/dev/shm",
+        hdfs_mode='read',
+        other_args=None,
+        **kwargs,
+    ):
+        self.meta_folders = json.loads(meta_folders)
+        self.meta_folder_repeats = json.loads(meta_folder_repeats)
+        self.meta_folder_identifiers = json.loads(other_args.meta_folder_identifiers)
+        self.pn_list = json.loads(other_args.pn_list)
+        self.pn_probs = json.loads(other_args.pn_probs)
+        assert len(self.meta_folders) == len(self.meta_folder_repeats), f'{len(self.meta_folders)} != {len(self.meta_folder_repeats)}'
+        assert len(self.meta_folders) == len(self.meta_folder_identifiers), f'{len(self.meta_folders)} != {len(self.meta_folder_identifiers)}'
+        self.verbose = verbose
+        self.buffer_size = buffersize
+        self.num_replicas = num_replicas
+        self.rank = rank
+        self.worker_id = 0
+        self.global_worker_id = 0
+        self.short_prob = short_prob
+        self.dataloader_workers = max(1, dataloader_workers)
+        self.shuffle = shuffle
+        self.global_workers = self.num_replicas * self.dataloader_workers
+        self.seed = seed
+        self.text_tokenizer = other_args.text_tokenizer
+        self.feature_extraction = other_args.only_images4extract_feats # < 0 # no sequence packing, for feature extraction
+        self.epoch_generator = None
+        self.epoch_rank_generator = None
+        self.other_args = other_args
+        self.pair_input = other_args.pair_input
+        self.drop_long_video = other_args.drop_long_video
+        self.enable_dynamic_length_prompt = enable_dynamic_length_prompt
+        self.set_epoch_generator(other_args.epoch)
+        self.temporal_compress_rate = other_args.temporal_compress_rate
+        self.dynamic_resolution_h_w, self.h_div_w_templates = get_dynamic_resolution_meta(other_args.dynamic_scale_schedule, other_args.train_h_div_w_list, other_args.video_frames) # here video_frames is the max video frames
+        self.video_fps = video_fps
+        self.min_training_duration = (other_args.min_video_frames-1) // self.video_fps
+        self.max_training_duration = (other_args.video_frames-1) // self.video_fps
+        self.c2i = self.other_args.add_class_token > 0
+        if self.c2i:
+            self.c2i_transform = transforms.Compose([
+                transforms.RandomHorizontalFlip(),
+                transforms.Resize(288, interpolation=InterpolationMode.LANCZOS), # transforms.Resize: resize the shorter edge to mid_reso
+                transforms.RandomCrop((256, 256)),
+                transforms.ToTensor(), 
+                normalize,
+            ])
+        else:
+            self.c2i_transform = None
+        self.print(f"{self.rank=} dataset {self.seed=}, {self.h_div_w_templates=} {self.min_training_duration=} {self.max_training_duration=}, cache_check_mode={self.other_args.cache_check_mode}")
+        self.token_cache_dir = other_args.token_cache_dir
+        self.use_vae_token_cache = other_args.use_vae_token_cache
+        self.allow_online_vae_feature_extraction = other_args.allow_online_vae_feature_extraction
+        self.use_text_token_cache = other_args.use_text_token_cache
+        self.max_video_frames = other_args.video_frames
+        self.cached_video_frames = other_args.cached_video_frames # cached max video frames
+        self.down_size_limit = other_args.down_size_limit
+        self.video_caption_type = other_args.video_caption_type
+        self.train_max_token_len = other_args.train_max_token_len
+        self.duration_resolution = other_args.duration_resolution
+        self.device = other_args.device
+        print(f'self.down_size_limit: {self.down_size_limit}')
+        self.hdfs_mode = hdfs_mode
+        self.max_text_len = other_args.tlen
+        self.temp_dir = temp_dir.rstrip("/")
+        self.mapped_duration2metas, self.mapped_duration2freqs = self.get_mapped_duration2metas()
+        self.batches = self.form_batches(self.mapped_duration2metas)
+        print(f'{num_replicas=}, {rank=}, {dataloader_workers=}, {len(self.batches)=}, {self.drop_long_video=} {self.max_text_len=} self.batches[:10]={self.batches[:10]}')
+
+    def print(self, string):
+        if self.feature_extraction:
+            print(string)
+        else:
+            print(string, force=True)
+    
+    def get_captions_lens(self, captions):
+        if self.other_args.text_tokenizer_type == 'flan_t5':
+            tokens = self.other_args.text_tokenizer(text=captions, max_length=self.other_args.text_tokenizer.model_max_length, padding='max_length', truncation=True, return_tensors='pt')
+            mask = tokens.attention_mask.cuda(non_blocking=True)
+            lens: List[int] = mask.sum(dim=-1).tolist()
+        else: # umt5-xxl
+            ids, mask = self.other_args.text_tokenizer( captions, return_mask=True, add_special_tokens=True)
+            lens = mask.gt(0).sum(dim=1).tolist()
+        return lens
+        
+    def get_video_caption(self, meta, mapped_duration):
+        if 'tarsier2_caption' not in meta:
+            caption = self.epoch_rank_generator.choice(meta['caption'])['content']
+        else:
+            caption_type = 'tarsier2_caption'
+            if ('MiniCPM_V_2_6_caption' in meta) and meta['MiniCPM_V_2_6_caption']:
+                caption_type = self.epoch_rank_generator.choice(['tarsier2_caption', 'MiniCPM_V_2_6_caption'])
+            caption = meta[caption_type]
+        if self.enable_dynamic_length_prompt and (self.epoch_rank_generator.random() < self.other_args.short_cap_prob):
+            caption = self.random_drop_sentences(caption, min_sentences=2)
+        if 'quality_prompt' in meta:
+            caption = caption + ' ' + meta['quality_prompt']
+        if meta['first_frame_condition']:
+            caption = '<I2V>' + caption
+        else:
+            caption = '<T2V>' + caption
+        assert caption
+        return caption
+    
+    def get_image_caption(self, meta):
+        caption = meta['long_caption']
+        if not meta['long_caption']:
+            caption = meta['text']
+        else:
+            if self.epoch_rank_generator.random() < self.other_args.short_cap_prob:
+                if meta['text']:
+                    caption = meta['text']
+                elif ('InternVL' in meta['long_caption_type']):
+                    caption = self.random_drop_sentences(meta['long_caption'], min_sentences=2)
+        caption = '<T2I>' + caption
+        assert caption
+        return caption
+
+    def sample_pn(self, meta, pn_list, pn_probs):
+        if ('height' in meta) and ('width' in meta):
+            real_pn = meta['height'] * meta['width'] / 1000000
+            valid_pn_list, valid_pn_probs = [], []
+            for pn, pn_prob in zip(pn_list, pn_probs):
+                if real_pn > 0.7 * float(pn[:-1]):
+                    valid_pn_list.append(pn)
+                    valid_pn_probs.append(pn_prob)
+            if not len(valid_pn_list):
+                valid_pn_list = pn_list[:1]
+                valid_pn_probs = pn_probs[:1]
+        else:
+            valid_pn_list = pn_list
+            valid_pn_probs = pn_probs
+        return self.epoch_rank_generator.choice(valid_pn_list, p=valid_pn_probs)
+
+    def get_mapped_duration2metas(self):
+        part_filepaths = []
+        part_file2identifier = {}
+        for meta_folder, meta_folder_repeat, meta_folder_identifier in zip(self.meta_folders, self.meta_folder_repeats, self.meta_folder_identifiers):
+            tmp_part_filepaths = sorted(glob.glob(osp.join(meta_folder, '*/*.jsonl')))
+            self.epoch_generator.shuffle(tmp_part_filepaths)
+            if meta_folder_repeat > 1:
+                tmp_part_filepaths = tmp_part_filepaths * int(np.ceil(meta_folder_repeat))
+            tmp_part_filepaths = tmp_part_filepaths[:int(len(tmp_part_filepaths)*meta_folder_repeat)]
+            for file in tmp_part_filepaths: part_file2identifier[file] = meta_folder_identifier
+            part_filepaths.extend(tmp_part_filepaths)
+        self.epoch_generator.shuffle(part_filepaths)
+        self.print(f'{self.rank=} jsonls sample: {part_filepaths[:4]}')
+        if self.num_replicas > 1:
+            part_filepaths = part_filepaths[self.rank::self.num_replicas]
+        
+        mapped_duration2metas = {}
+        pbar = tqdm.tqdm(total=len(part_filepaths))
+        total, corrupt = 0, 0
+        stop_read = False
+        rough_h_div_w = self.h_div_w_templates[np.argmin(np.abs((9/16-self.h_div_w_templates)))]
+        for part_filepath in part_filepaths:
+            file_quality_prompt = part_file2identifier[part_filepath]
+            if stop_read:
+                break
+            pbar.update(1)
+            try:
+                with open(part_filepath, encoding='utf-8') as f:
+                    lines = f.readlines()
+            except Exception as e:
+                print(f'{part_filepath=} Error: {e}')
+                lines = []
+            for line in lines:
+                total += 1
+                try:
+                    meta = json.loads(line)
+                except Exception as e:
+                    print(e)
+                    corrupt += 1
+                    print(e, corrupt, total, corrupt/total)
+                    continue
+                if file_quality_prompt: # override quality prompt
+                    meta['quality_prompt'] = file_quality_prompt
+                if ('height' in meta) and ('width' in meta):
+                    cur_h_div_w_template = self.h_div_w_templates[np.argmin(np.abs((meta['height']/meta['width']-self.h_div_w_templates)))]
+                else:
+                    cur_h_div_w_template = rough_h_div_w
+                if 'h_div_w' in meta:
+                    del meta['h_div_w']
+                meta['first_frame_condition'] = False
+                meta['pn'] = self.sample_pn(meta, self.pn_list, self.pn_probs)
+                if 'video_path' in meta:
+                    if self.epoch_rank_generator.random() < self.other_args.i2v_ratio:
+                        meta['first_frame_condition'] = True
+                    begin_frame_id, end_frame_id, fps = meta['begin_frame_id'], meta['end_frame_id'], meta['fps']
+                    real_duration = (end_frame_id - begin_frame_id) / fps
+                    mapped_duration = int(np.round(real_duration / self.duration_resolution)) * self.duration_resolution
+                    if mapped_duration < self.min_training_duration:
+                        continue
+                    if mapped_duration > self.max_training_duration:
+                        if self.drop_long_video:
+                            continue
+                        else:
+                            mapped_duration = self.max_training_duration
+                    if self.other_args.use_clipwise_caption:
+                        meta['caption'] = [
+                            meta['caption-InternVL2.0'],
+                            self.get_video_caption(meta, mapped_duration),
+                        ]
+                    else:
+                        meta['caption'] = [self.get_video_caption(meta, mapped_duration)]
+                    sample_frames = int(mapped_duration * self.video_fps + 1)
+                    pt = (sample_frames-1) // self.temporal_compress_rate + 1
+                    scale_schedule = self.dynamic_resolution_h_w[cur_h_div_w_template][meta['pn']]['pt2scale_schedule'][pt]
+                    meta['sample_frames'] = sample_frames
+                elif 'image_path' in meta:
+                    mapped_duration = -1
+                    scale_schedule = self.dynamic_resolution_h_w[cur_h_div_w_template][meta['pn']]['pt2scale_schedule'][1]
+                    meta['caption'] = [self.get_image_caption(meta)]
+                # random set caption to "" for classifier-free guidance
+                # refer to: https://github.com/PixArt-alpha/PixArt-alpha/blob/master/train_scripts/train_diffusers.py#L67
+                for caption_ind in range(len(meta['caption'])):
+                    if self.epoch_rank_generator.random() < self.other_args.drop_condition_prob:
+                        meta['caption'][caption_ind] = ""
+                if mapped_duration not in mapped_duration2metas:
+                    mapped_duration2metas[mapped_duration] = []
+                
+                # get cum_text_visual_tokens
+                cum_visual_tokens = []
+                preserve_scale_inds = {}
+                assert len(scale_schedule) == len(self.other_args.video_scale_probs), f'{len(scale_schedule)=} {len(self.other_args.video_scale_probs)=}'
+                for scale_ind, scale in enumerate(scale_schedule):
+                    if self.epoch_rank_generator.random() < self.other_args.video_scale_probs[scale_ind]:
+                        preserve_scale_inds[scale_ind] = True
+                        tokens_this_scale = np.array(scale).prod(-1) + self.other_args.add_scale_token
+                        cum_visual_tokens.append(tokens_this_scale)
+                cum_visual_tokens = np.array(cum_visual_tokens).cumsum()
+                meta['cum_text_visual_tokens'] = cum_visual_tokens
+                meta['preserve_scale_inds'] = preserve_scale_inds
+                
+                if self.other_args.cache_check_mode == 1: # check at the begining
+                    if self.exists_cache_file(meta):
+                        mapped_duration2metas[mapped_duration].append(meta)
+                elif self.other_args.cache_check_mode == -1: # select unexist, used for token cache
+                    if not self.exists_cache_file(meta):
+                        mapped_duration2metas[mapped_duration].append(meta)
+                else:
+                    mapped_duration2metas[mapped_duration].append(meta)
+                
+                total_metas = sum([len(item) for item in mapped_duration2metas.values()])
+                if (self.other_args.restrict_data_size > 0) and (total_metas > self.other_args.restrict_data_size / self.num_replicas):
+                    stop_read = True
+                    break
+                
+        # set mapped_duration2freqs
+        mapped_duration2freqs = {}
+        for mapped_duration in sorted(mapped_duration2metas.keys()):
+            mapped_duration2freqs[mapped_duration] = len(mapped_duration2metas[mapped_duration])
+
+        for mapped_duration in mapped_duration2metas.keys():
+            freqs = mapped_duration2freqs[mapped_duration]
+            assert len(mapped_duration2metas[mapped_duration]) >= freqs
+            self.epoch_rank_generator.shuffle(mapped_duration2metas[mapped_duration])
+            mapped_duration2metas[mapped_duration] = mapped_duration2metas[mapped_duration][:freqs]
+            # append text tokens
+            skip_count_text_token = self.other_args.skip_count_text_token or self.other_args.add_class_token > 0
+            mapped_duration2metas[mapped_duration] = self.append_text_tokens(mapped_duration2metas[mapped_duration], skip_count_text_token=skip_count_text_token)
+
+        total_metas = sum([len(item) for item in mapped_duration2metas.values()])
+        for mapped_duration in sorted(mapped_duration2freqs.keys()):
+            freq = mapped_duration2freqs[mapped_duration]
+            proportion = freq / total_metas * 100
+            print(f'{mapped_duration=}, {freq=}, {proportion=:.1f}%')
+        return mapped_duration2metas, mapped_duration2freqs
+
+    def append_text_tokens(self, metas, skip_count_text_token=False, bucket_size=100):
+        t1 = time.time()
+        pbar = tqdm.tqdm(total=len(metas) // bucket_size + 1, desc='append text tokens')
+        valid_metas = []
+        for bucket_id in range(len(metas) // bucket_size + 1):
+            pbar.update(1)
+            start = bucket_id * bucket_size
+            end = min(start + bucket_size, len(metas))
+            if start >= end:
+                break
+            captions = []
+            caps_per_meta = []
+            for i in range(start, end):
+                captions.extend(metas[i]['caption'])
+                caps_per_meta.append(len(metas[i]['caption']))
+            assert len(captions), f'{len(captions)=}'
+            if skip_count_text_token:
+                lens = [0 for _ in range(len(captions))]
+            else:
+                lens = self.get_captions_lens(captions)
+                lens = np.clip(np.array(lens), a_min=0, a_max=self.max_text_len)
+            ptr = 0
+            for i in range(start, end):
+                text_tokens = sum(lens[ptr:ptr+caps_per_meta[i-start]])
+                ptr += caps_per_meta[i-start]
+                metas[i]['text_tokens'] = text_tokens
+                metas[i]['cum_text_visual_tokens'] = metas[i]['cum_text_visual_tokens'] + metas[i]['text_tokens']
+                metas[i]['text_visual_tokens'] = metas[i]['cum_text_visual_tokens'][-1]
+                if metas[i]['text_visual_tokens'] <= self.train_max_token_len * self.other_args.dense_ratio4seqpack:
+                    valid_metas.append(metas[i])
+        t2 = time.time()
+        print(f'append text tokens: {t2-t1:.1f}s')
+        return valid_metas
+
+    def exists_cache_file(self, meta):
+        pn = meta['pn']
+        if 'image_path' in meta:
+            return osp.exists(self.get_image_cache_file(meta['image_path'], pn))
+        else:
+            if '/vdataset/clip' in meta['video_path']: # clip
+                cache_file = self.get_video_cache_file(meta['video_path'], 0, meta['end_frame_id']-meta['begin_frame_id'], self.video_fps, pn)
+            else:
+                cache_file = self.get_video_cache_file(meta['video_path'], meta['begin_frame_id'], meta['end_frame_id'], self.video_fps, pn)
+            return osp.exists(cache_file)
+    
+    def form_batches(self, mapped_duration2metas):
+        examples = []
+        for mapped_duration in sorted(mapped_duration2metas.keys()):
+            for example_ind in range(len(mapped_duration2metas[mapped_duration])):
+                text_visual_tokens = mapped_duration2metas[mapped_duration][example_ind]['text_visual_tokens']
+                examples.append((mapped_duration, example_ind, text_visual_tokens))
+        examples = sorted(examples, key=lambda x: -x[2])
+        max_text_visual_tokens = examples[0][2] if len(examples) else 0
+        assert self.train_max_token_len >= max_text_visual_tokens, f'{self.train_max_token_len=} should >= {max_text_visual_tokens=}'
+        self.print(f'{self.rank=} {self.mapped_duration2freqs=} form_batches details: {self.rank=} examples={examples[:20]}')
+        
+        st = time.time()
+        if self.feature_extraction or self.pair_input: # no sequence packing, for feature extraction or dpo training
+            batches = [[item[:2]] for item in examples]
+        else:
+            batches = []
+            left_ptr, right_ptr = 0, len(examples)-1
+            while left_ptr <= right_ptr:
+                tokens_remain = self.train_max_token_len
+                tmp_batch = []
+                while left_ptr <= right_ptr and (tokens_remain - examples[left_ptr][2] >= 0):
+                    tokens_remain = tokens_remain - examples[left_ptr][2]
+                    tmp_batch.append((left_ptr, examples[left_ptr][2]))
+                    left_ptr += 1
+                while left_ptr <= right_ptr and (tokens_remain - examples[right_ptr][2] >= 0):
+                    tokens_remain = tokens_remain - examples[right_ptr][2]
+                    tmp_batch.append((right_ptr, examples[right_ptr][2]))
+                    right_ptr -= 1
+                if len(tmp_batch):
+                    tmp_batch = sorted(tmp_batch, key=lambda x: -x[1])
+                    # total_tokens = sum(item[1] for item in tmp_batch)
+                    # if total_tokens / self.train_max_token_len < 0.7:
+                    #     import pdb; pdb.set_trace()
+                    batches.append([examples[ptr][:2] for (ptr, _) in tmp_batch])
+                    if len(batches) % 1000 == 0:
+                        print(f'form {len(batches)} batches, len(metas)={len(examples)}')
+        print(f'[data preprocess] form_batches done, got {len(batches)} batches, cost {time.time()-st:.2f}s')
+        self.epoch_rank_generator.shuffle(batches)
+        print(f'[data preprocess] shuffle batches done')
+        batch_num = len(batches)
+        try:
+            if self.num_replicas > 1:
+                batch_num = torch.tensor([batch_num], device=self.device)
+                if tdist.is_initialized():
+                    tdist.all_reduce(batch_num, op=tdist.ReduceOp.MIN)
+                batch_num = batch_num.item()
+        except Exception as e:
+            print(e)
+        batches = batches[:batch_num]
+        print(f'[data preprocess] aligned batch number among gpus, got {batch_num} batches')
+        return batches
+    
+    def set_epoch_generator(self, epoch):
+        self.epoch = epoch
+        self.epoch_generator = np.random.default_rng(self.seed + self.epoch)
+        self.epoch_rank_generator = np.random.default_rng(self.seed + self.epoch + self.rank)
+
+    def __getitem__(self, batch_ind_ptr):
+        try:
+            batch_info = self.batches[batch_ind_ptr%len(self.batches)]
+            batch_data = []
+            for (mapped_duration, example_ind) in batch_info:
+                ret = False
+                repeat_times = 0
+                mapped_duration_metas = self.mapped_duration2metas[mapped_duration]
+                while not ret:
+                    example_ind = example_ind % len(mapped_duration_metas)
+                    meta = mapped_duration_metas[example_ind]
+                    if 'video_path' in meta:
+                        if self.pair_input:
+                            ret, model_input = self.prepare_pair_video_input(meta)
+                        else:
+                            ret, model_input = self.prepare_video_input(meta)
+                    elif 'image_path' in meta:
+                        if self.pair_input:
+                            ret, model_input = self.prepare_pair_image_input(meta)
+                        else:
+                            ret, model_input = self.prepare_image_input(meta)
+                    if ret:
+                        if self.pair_input:
+                            batch_data.extend(model_input)
+                        else:
+                            batch_data.append(model_input)
+                    else: # Handle corrupt example in a batch, just try to read the next one
+                        example_ind = example_ind + 1
+                        repeat_times += 1
+                        if repeat_times % 20 == 0: # Too many corrupt files, switch to another batch
+                            self.print(f'Caution! I have repeat {repeat_times} times to read a video/image, but still failed to read it. {example_ind=} {meta=}')
+                            return self.__getitem__(batch_ind_ptr+1)
+            
+            images, raw_features_bcthw, feature_cache_files4images  = [], [], []
+            text_feature_cache_files = []
+            addition_pn_images = {}
+            batch_data4images, batch_data4raw_features = [], []
+            for item in batch_data:
+                if item['raw_features_cthw'] is None:
+                    images.append(item['img_T3HW'].permute(1,0,2,3)) # # tchw -> cthw
+                    for key in item:
+                        if key.startswith('img_T3HW_'):
+                            if key not in addition_pn_images:
+                                addition_pn_images[key] = []
+                            addition_pn_images[key].append(item[key].permute(1,0,2,3))
+                    feature_cache_files4images.append(item['feature_cache_file'])
+                    batch_data4images.append(item)
+                else:
+                    raw_features_bcthw.append(item['raw_features_cthw'])
+                    batch_data4raw_features.append(item)
+            batch_data4images_raw_features = batch_data4images + batch_data4raw_features
+            captions = [item['text_input'] for item in batch_data4images_raw_features]
+            text_feature_cache_files = [item['text_feature_cache_file'] for item in batch_data4images_raw_features]
+            meta_list = [item['meta'] for item in batch_data4images_raw_features]
+            return {
+                'captions': captions, 
+                'images': images, 
+                'addition_pn_images': addition_pn_images,
+                'feature_cache_files4images': feature_cache_files4images,
+                'raw_features_bcthw': raw_features_bcthw, 
+                'text_cond_tuple': None,
+                'text_feature_cache_files': text_feature_cache_files,
+                'meta_list': meta_list,
+                'media': 'videos',
+            }
+        except Exception as e:
+            print(f'get item error: {e}')
+            return self.__getitem__(batch_ind_ptr+1)
+
+
+    def prepare_image_input(self, info) -> Tuple:
+        try:
+            img_path, text_input = osp.abspath(info['image_path']), info['caption']
+            img_T3HW, raw_features_cthw, feature_cache_file, text_features_lenxdim, text_feature_cache_file = [None] * 5
+            if self.use_vae_token_cache:
+                feature_cache_file = self.get_image_cache_file(img_path, info['pn'])
+                if osp.exists(feature_cache_file):
+                    try:
+                        raw_features_cthw = self.load_visual_token(feature_cache_file)
+                    except Exception as e:
+                        print(f'load cache file error: {e}')
+                        os.remove(feature_cache_file)
+                if raw_features_cthw is None and (not self.allow_online_vae_feature_extraction):
+                    return False, None
+            if raw_features_cthw is None:
+                with open(img_path, 'rb') as f:
+                    img: PImage.Image = PImage.open(f)
+                    w, h = img.size
+                    h_div_w = h / w
+                    h_div_w_template = self.h_div_w_templates[np.argmin(np.abs((h_div_w-self.h_div_w_templates)))]
+                    tgt_h, tgt_w = self.dynamic_resolution_h_w[h_div_w_template][info['pn']]['pixel']
+                    img = img.convert('RGB')
+                    if self.c2i:
+                        img_T3HW = self.c2i_transform(img)
+                    else:
+                        img_T3HW = transform(img, tgt_h, tgt_w)
+                    img_T3HW = img_T3HW.unsqueeze(0)
+                    assert img_T3HW.shape == (1, 3, tgt_h, tgt_w)
+            data_item = {
+                'text_input': text_input,
+                'img_T3HW': img_T3HW,
+                'raw_features_cthw': raw_features_cthw,
+                'feature_cache_file': feature_cache_file,
+                'text_features_lenxdim': text_features_lenxdim,
+                'text_feature_cache_file': text_feature_cache_file,
+                'meta': info,
+            }
+            return True, data_item
+        except Exception as e:
+            print(f'prepare_image_input error: {e}')
+            return False, None
+
+    def prepare_pair_image_input(self, info) -> Tuple:
+        pass
+        
+    def prepare_pair_video_input(self, info) -> Tuple:
+        win_flag, win_data_item = self.prepare_video_input(copy.deepcopy(info))
+        
+        info['video_path'] = info['lose_video_path']
+        lose_flag, lose_data_item = self.prepare_video_input(info)
+
+        flag = win_flag and lose_flag
+        return flag, [win_data_item, lose_data_item]
+
+    def load_visual_token(self, feature_cache_file):
+        raw_features_cthw = load_packed_tensor(feature_cache_file)
+        from grn.utils_t2iv.hbq_util_t2iv import bit_label2raw_feature
+        raw_features_cthw = bit_label2raw_feature(raw_features_cthw.unsqueeze(0), self.other_args.hbq_round)[0]
+        return raw_features_cthw
+    
+    def prepare_video_input(self, info) -> Tuple:
+        filename, begin_frame_id, end_frame_id = (
+            info["video_path"],
+            info["begin_frame_id"],
+            info["end_frame_id"],
+        )
+
+        try:
+            img_T3HW, raw_features_cthw, feature_cache_file, text_features_lenxdim, text_feature_cache_file = None, None, None, None, None
+            img_T3HW_4additional_pn = {}
+            text_input = info['caption']
+            if '/vdataset/clip' in filename: # clip
+                begin_frame_id, end_frame_id = 0, end_frame_id - begin_frame_id
+            sample_frames = info['sample_frames']
+            tmp_local_path = ''
+            if self.use_vae_token_cache:
+                feature_cache_file = self.get_video_cache_file(info["video_path"], begin_frame_id, end_frame_id, self.video_fps, info['pn'])
+                if osp.exists(feature_cache_file):
+                    try:
+                        pt = (sample_frames-1) // self.temporal_compress_rate + 1
+                        raw_features_cthw = self.load_visual_token(feature_cache_file)
+                        assert raw_features_cthw.shape[1] >= pt, f'raw_features_cthw.shape[1] >= pt: {raw_features_cthw.shape[1]} vs {pt}'
+                        if raw_features_cthw.shape[1] > pt:
+                            raw_features_cthw = raw_features_cthw[:,:pt]
+                    except Exception as e:
+                        self.print(f'load video cache file error: {e}')
+                        os.remove(feature_cache_file)
+                        raw_features_cthw = None
+                if raw_features_cthw is None and (not self.allow_online_vae_feature_extraction):
+                    return False, None
+            pn_list = [info['pn']]
+            if raw_features_cthw is None:
+                tmp_local_path = info["video_path"]
+                if not osp.exists(tmp_local_path):
+                    return False, None
+                video = EncodedVideoDecord(tmp_local_path, os.path.basename(tmp_local_path), num_threads=0)
+                start_interval = max(0, begin_frame_id / video._fps)
+                end_interval = start_interval+(sample_frames-1)/self.video_fps
+                assert end_interval <= video.duration + 0.2, f'{end_interval=}, but {video.duration=}' # 0.2s margin
+                end_interval = min(end_interval, video.duration)
+                raw_video, _ = video.get_clip(start_interval, end_interval, sample_frames) # rgb order
+                h, w, _ = raw_video[0].shape
+                h_div_w = h / w
+                h_div_w_template = self.h_div_w_templates[np.argmin(np.abs((h_div_w-self.h_div_w_templates)))]
+                tgt_h, tgt_w = self.dynamic_resolution_h_w[h_div_w_template][info['pn']]['pixel']
+                    
+                for pn in pn_list:
+                    img_T3HW = [transform(Image.fromarray(frame).convert("RGB"), tgt_h, tgt_w) for frame in raw_video]
+                    img_T3HW = torch.stack(img_T3HW, 0)
+                    img_T3HW_4additional_pn[pn] = img_T3HW
+                del video
+                assert img_T3HW.shape[-3:] == (3, tgt_h, tgt_w)
+            data_item = {
+                'text_input': text_input,
+                'img_T3HW': img_T3HW_4additional_pn.get(info['pn'], None),
+                'raw_features_cthw': raw_features_cthw,
+                'feature_cache_file': feature_cache_file,
+                'text_features_lenxdim': text_features_lenxdim,
+                'text_feature_cache_file': text_feature_cache_file,
+                'meta': info,
+            }
+            for pn in pn_list[1:]:
+                data_item.update({f'img_T3HW_{pn}': img_T3HW_4additional_pn.get(pn, None)})
+            return True, data_item
+        except Exception as e:
+            self.print(f'prepare_video_input error: {e}, info: {info}')
+            return False, None
+
+        
+    @staticmethod
+    def collate_function(batch, online_t5: bool = False) -> None:
+        pass
+    
+    def random_drop_sentences(self, caption, min_sentences):
+        elems = [item for item in caption.split('.') if item]
+        if len(elems) <= min_sentences:
+            return caption
+        sentences = self.epoch_rank_generator.integers(min_sentences, len(elems)+1)
+        return '.'.join(elems[:sentences]) + '.'
+
+    def __len__(self):
+        return len(self.batches) * self.other_args.loop_data_per_epoch
+
+    def get_image_cache_file(self, image_path, pn):
+        elems = image_path.split('/')
+        elems = [item for item in elems if item]
+        filename, ext = osp.splitext(elems[-1])
+        filename = get_prompt_id(filename)
+        save_filepath = osp.join(self.token_cache_dir, f'images_pn_{pn}', '/'.join(elems[4:-1]), f'{filename}.npz')
+        return save_filepath
+
+    def get_video_cache_file(self, video_path, begin_frame_id, end_frame_id, video_fps, pn):
+        elems = video_path.split('/')
+        elems = [item for item in elems if item]
+        filename, ext = osp.splitext(elems[-1])
+        filename = get_prompt_id(filename)
+        save_filepath = osp.join(self.token_cache_dir, f'pn_{pn}_sample_fps_{video_fps}', '/'.join(elems[4:-1]), f'{filename}_sf_{begin_frame_id}_ef_{end_frame_id}.npz')
+        return save_filepath
+    

grn/models/basic.py

@@ -0,0 +1,256 @@
+import math
+import os
+from functools import partial
+from typing import Optional, Tuple, Union
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import numpy as np
+from torch.utils.checkpoint import checkpoint
+from torch.nn.functional import scaled_dot_product_attention as slow_attn    # q, k, v: BHLc
+
+from grn.models.rope import apply_rotary_emb
+from grn.utils_t2iv.sequence_parallel import sp_all_to_all, SequenceParallelManager as sp_manager
+
+try:
+    from flash_attn.cute import flash_attn_varlen_func
+except:
+    from flash_attn import flash_attn_varlen_func
+                
+# Import flash_attn's fused ops
+try:
+    from flash_attn.ops.rms_norm import rms_norm as rms_norm_impl
+except ImportError:
+    def rms_norm_impl(x, weight, epsilon):
+        return (x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True).add_(epsilon))) * weight
+
+def merge_states(states, splits, cfg):
+    """
+    pick key and value states for flash_attn_varlen_func
+    Args:
+        states: list of states
+        splits: list of split sizes
+        cfg: bool, use cfg or not"""
+    if cfg:
+        cond_len, uncond_len = 0, 0
+        cond_states, uncond_states = [], []
+        for stat_, split_ in zip(states, splits):
+            cond, uncond = torch.split(stat_, split_, dim=2)
+            cond_states.append(cond)
+            uncond_states.append(uncond)
+            cond_len += cond.shape[2]
+            uncond_len += uncond.shape[2]
+        return cond_states + uncond_states, [cond_len, uncond_len]
+    else:
+        cond_len = 0
+        for stat_ in states:
+            cond_len += stat_.shape[2]
+        return states, [cond_len]
+
+class FastRMSNorm(nn.Module):
+    def __init__(self, C, eps=1e-6, elementwise_affine=True):
+        super().__init__()
+        self.C = C
+        self.eps = eps
+        self.elementwise_affine = elementwise_affine
+        if self.elementwise_affine:
+            self.weight = nn.Parameter(torch.ones(C))
+        else:
+            self.register_buffer('weight', torch.ones(C))
+    
+    def forward(self, x):
+        src_type = x.dtype
+        return rms_norm_impl(x.float(), self.weight, epsilon=self.eps).to(src_type)
+    
+    def extra_repr(self) -> str:
+        return f'C={self.C}, eps={self.eps:g}, elementwise_affine={self.elementwise_affine}'
+
+
+class WanLayerNorm(nn.LayerNorm):
+
+    def __init__(self, dim, eps=1e-6, elementwise_affine=False):
+        super().__init__(dim, elementwise_affine=elementwise_affine, eps=eps)
+
+    def forward(self, x):
+        r"""
+        Args:
+            x(Tensor): Shape [B, L, C]
+        """
+        return super().forward(x.float()).type_as(x)
+
+
+class Qwen3MLP(nn.Module):
+    def __init__(self, hidden_size, intermediate_size):
+        super().__init__()
+        self.hidden_size = hidden_size
+        self.intermediate_size = intermediate_size
+        self.gate_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False)
+        self.up_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False)
+        self.down_proj = nn.Linear(self.intermediate_size, self.hidden_size, bias=False)
+        self.act_fn = nn.SiLU()
+
+    def forward(self, x):
+        down_proj = self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x))
+        return down_proj
+
+def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor:
+    """
+    This is the equivalent of torch.repeat_interleave(x, dim=1, repeats=n_rep). The hidden states go from (batch,
+    num_key_value_heads, seqlen, head_dim) to (batch, num_attention_heads, seqlen, head_dim)
+    """
+    batch, num_key_value_heads, slen, head_dim = hidden_states.shape
+    if n_rep == 1:
+        return hidden_states
+    hidden_states = hidden_states[:, :, None, :, :].expand(batch, num_key_value_heads, n_rep, slen, head_dim)
+    return hidden_states.reshape(batch, num_key_value_heads * n_rep, slen, head_dim)
+
+class SelfAttention(nn.Module):
+    def __init__(
+        self, embed_dim=768, num_heads=12, num_key_value_heads=-1,
+        use_flex_attn=False, qwen_qkvo_bias=False, **kwargs,
+    ):
+        """
+        :param embed_dim: model's width
+        :param num_heads: num heads of multi-head attention
+        :param proj_drop: always 0 for testing
+        :param tau: always 1
+        :param cos_attn: always True: during attention, q and k will be L2-normalized and scaled by a head-wise learnable parameter self.scale_mul_1H11
+        """
+        super().__init__()
+        assert embed_dim % num_heads == 0
+        assert num_key_value_heads == -1 or num_heads % num_key_value_heads == 0
+        
+        self.num_heads, self.head_dim = num_heads, embed_dim // num_heads
+        self.num_key_value_heads = num_key_value_heads if num_key_value_heads > 0 else num_heads
+        self.q_proj = nn.Linear(embed_dim, self.num_heads*self.head_dim, bias=qwen_qkvo_bias)
+        self.k_proj = nn.Linear(embed_dim, self.num_key_value_heads*self.head_dim, bias=qwen_qkvo_bias)
+        self.v_proj = nn.Linear(embed_dim, self.num_key_value_heads*self.head_dim, bias=qwen_qkvo_bias)
+        self.o_proj = nn.Linear(self.num_heads*self.head_dim, embed_dim, bias=qwen_qkvo_bias)
+        self.q_norm = FastRMSNorm(self.head_dim)
+        self.k_norm = FastRMSNorm(self.head_dim)
+        self.num_key_value_groups = self.num_heads // self.num_key_value_heads
+        self.scale = self.head_dim**-0.5
+        
+        self.caching = False    # kv caching: only used during inference
+        self.cached_k = {}    # kv caching: only used during inference
+        self.cached_v = {}    # kv caching: only used during inference
+        self.cached_split_cond_uncond = {} # only used during inference
+
+        self.use_flex_attn = use_flex_attn
+    
+    def kv_caching(self, enable: bool): # kv caching: only used during inference
+        self.caching = enable
+        self.cached_k = {}
+        self.cached_v = {}
+        self.cached_split_cond_uncond = {}
+
+    # NOTE: attn_bias_or_two_vector is None during inference
+    def forward(self, x, cu_seqlens, max_seqlen, attn_bias_or_two_vector: Union[torch.Tensor, Tuple[torch.IntTensor, torch.IntTensor]], attn_fn=None, rope2d_freqs_grid=[], scale_ind=0, context_info=None, last_diffusion_step=True, ref_text_scale_inds=[], use_cfg=False, split_cond_uncond=[], **kwargs):
+        # x: fp32
+        B, L, C = x.shape
+        hidden_states = x
+        input_shape = hidden_states.shape[:-1]
+        hidden_shape = (*input_shape, -1, self.head_dim)
+        query_states = self.q_norm(self.q_proj(hidden_states).view(hidden_shape)).contiguous()# batch, slen, heads, head_dim
+        key_states = self.k_norm(self.k_proj(hidden_states).view(hidden_shape)).contiguous() # batch, slen, num_key_value_heads, head_dim
+        value_states = self.v_proj(hidden_states).view(hidden_shape).contiguous() # batch, slen, num_key_value_heads, head_dim
+
+        if sp_manager.sp_on():
+            # Headnum need to be sharded and L needs to be gathered
+            # [B, H, raw_L/sp, C] --> [B, H/sp, raw_L, C]
+            sdim = 1
+            gdim = 2
+            L = L * sp_manager.get_sp_size()
+            C = C // sp_manager.get_sp_size()
+            query_states = sp_all_to_all(query_states, sdim, gdim)
+            key_states = sp_all_to_all(key_states, sdim, gdim)
+            value_states = sp_all_to_all(value_states, sdim, gdim)
+
+        query_states, key_states = apply_rotary_emb(query_states, key_states, rope2d_freqs_grid)
+        key_states = repeat_kv(key_states, self.num_key_value_groups)
+        value_states = repeat_kv(value_states, self.num_key_value_groups)
+
+        if attn_bias_or_two_vector is None:
+            # fa4, flash_attn_func input/output should be (batch_size, seqlen, nheads, headdim)
+            from flash_attn.cute import flash_attn_varlen_func
+            attn_output = flash_attn_varlen_func(
+                q = query_states.squeeze(0),
+                k = key_states.squeeze(0),
+                v = value_states.squeeze(0),
+                cu_seqlens_q=cu_seqlens,
+                cu_seqlens_k=cu_seqlens,
+                max_seqlen_q=max_seqlen,
+                max_seqlen_k=max_seqlen,
+                softmax_scale=self.scale,
+            )
+            attn_output = attn_output[0].reshape(B, L, C).contiguous()
+        else:
+            # slow attn
+            attn_output = slow_attn(query=query_states.transpose(1, 2), key=key_states.transpose(1, 2), value=value_states.transpose(1, 2), scale=self.scale, attn_mask=attn_bias_or_two_vector, dropout_p=0).transpose(1, 2).reshape(B, L, C)
+
+        if sp_manager.sp_on():
+            # [B, raw_L, C/sp] --> [B, raw_L/sp, C]
+            sdim = 1
+            gdim = 2
+            attn_output = sp_all_to_all(attn_output, sdim, gdim)
+
+        attn_output = self.o_proj(attn_output)
+
+        return attn_output
+    
+class SelfAttnBlock(nn.Module):
+    def __init__(
+        self, embed_dim, num_heads, num_key_value_heads, mlp_ratio=4.,
+        use_flex_attn=False,
+        qwen_qkvo_bias=False, use_ada_layer_norm=False, **kwargs,
+    ):
+        super(SelfAttnBlock, self).__init__()
+        self.C = embed_dim
+        self.attn = SelfAttention(
+            embed_dim=embed_dim, num_heads=num_heads, num_key_value_heads=num_key_value_heads,
+            use_flex_attn=use_flex_attn, qwen_qkvo_bias=qwen_qkvo_bias, **kwargs,
+        )
+        self.mlp = Qwen3MLP(hidden_size=embed_dim, intermediate_size=round(embed_dim * mlp_ratio / 256) * 256)
+        self.use_ada_layer_norm = use_ada_layer_norm
+        if self.use_ada_layer_norm:
+            self.modulation = nn.Parameter(torch.randn(1, 6, embed_dim) / embed_dim**0.5)
+            self.input_layernorm = WanLayerNorm(embed_dim)
+            self.post_attention_layernorm = WanLayerNorm(embed_dim)
+        else:
+            self.input_layernorm = FastRMSNorm(embed_dim)
+            self.post_attention_layernorm = FastRMSNorm(embed_dim)
+        
+    # NOTE: attn_bias_or_two_vector is None during inference
+    def forward(self, x, cu_seqlens, max_seqlen, e0, attn_bias_or_two_vector, attn_fn=None, rope2d_freqs_grid=[], scale_ind=0, context_info=None, last_diffusion_step=True, ref_text_scale_inds=[], use_cfg=False, split_cond_uncond=[], **kwargs):
+        # x: [B,L,C]
+        # e0: [B, L, 6, C]
+        if self.use_ada_layer_norm:
+            assert e0.dtype == torch.float32
+            e = e0
+            with torch.amp.autocast('cuda', dtype=torch.float32):
+                e = (self.modulation.unsqueeze(0) + e).chunk(6, dim=2)
+            residual = x
+            hidden_states = x
+            hidden_states = self.input_layernorm(hidden_states).float() * (1 + e[1].squeeze(2)) + e[0].squeeze(2)
+            hidden_states = self.attn(hidden_states, cu_seqlens, max_seqlen, attn_bias_or_two_vector, attn_fn, rope2d_freqs_grid, scale_ind, context_info, last_diffusion_step, ref_text_scale_inds, use_cfg, split_cond_uncond, **kwargs)
+            with torch.amp.autocast('cuda', dtype=torch.float32):
+                hidden_states = residual + hidden_states * e[2].squeeze(2)
+            # Fully Connected
+            residual = hidden_states
+            hidden_states = self.post_attention_layernorm(hidden_states).float() * (1 + e[4].squeeze(2)) + e[3].squeeze(2)
+            hidden_states = self.mlp(hidden_states)
+            with torch.amp.autocast('cuda', dtype=torch.float32):
+                hidden_states = residual + hidden_states * e[5].squeeze(2)
+        else:
+            residual = x
+            hidden_states = x
+            hidden_states = self.input_layernorm(hidden_states)
+            hidden_states = self.attn(hidden_states, cu_seqlens, max_seqlen, attn_bias_or_two_vector, attn_fn, rope2d_freqs_grid, scale_ind, context_info, last_diffusion_step, ref_text_scale_inds, use_cfg, split_cond_uncond, **kwargs)
+            hidden_states = residual + hidden_states
+            # Fully Connected
+            residual = hidden_states
+            hidden_states = self.post_attention_layernorm(hidden_states)
+            hidden_states = self.mlp(hidden_states)
+            hidden_states = residual + hidden_states
+        return hidden_states

grn/models/ema.py

@@ -0,0 +1,23 @@
+import copy
+import torch
+from collections import OrderedDict
+
+
+def get_ema_model(model):
+    ema_model = copy.deepcopy(model)
+    ema_model.eval()
+    for param in ema_model.parameters():
+        param.requires_grad = False
+    return ema_model
+
+@torch.no_grad()
+def update_ema(ema_model, model, decay):
+    """
+    Step the EMA model towards the current model.
+    """
+    ema_params = OrderedDict(ema_model.named_parameters())
+    model_params = OrderedDict(model.named_parameters())
+
+    for name, param in model_params.items():
+        # TODO: Consider applying only to params that require_grad to avoid small numerical changes of pos_embed
+        ema_params[name].mul_(decay).add_(param.data, alpha=1 - decay)

grn/models/flex_attn_mask.py

@@ -0,0 +1,67 @@
+from functools import partial
+import torch
+import numpy as np
+import torch.nn as nn
+import torch.nn.functional as F
+from torch.nn.attention.flex_attention import flex_attention, create_block_mask
+
+
+def _length_to_offsets(lengths, device):
+    offsets = [0]
+    offsets.extend(lengths)
+    offsets = torch.tensor(offsets, device=device, dtype=torch.int32)
+    offsets = torch.cumsum(offsets, dim=-1)
+    return offsets
+
+def _offsets_to_doc_ids_tensor(offsets):
+    device = offsets.device
+    counts = offsets[1:] - offsets[:-1]
+    visual = torch.repeat_interleave(torch.arange(len(counts), device=device, dtype=torch.int32), counts)
+    return visual
+
+def _generate_overall_mask(offsets, querysid_refsid):
+    document_id = _offsets_to_doc_ids_tensor(offsets) # to scale_ind
+    def overall_mask(b, h, q_idx, kv_idx):
+        querysid = document_id[q_idx]
+        kv_sid = document_id[kv_idx]
+        return querysid_refsid[querysid][kv_sid]
+    return overall_mask
+
+def causal(b, h, q_idx, kv_idx):
+    return q_idx >= kv_idx
+
+def build_flex_attn_func(
+        flex_attention,
+        seq_l,
+        prefix_lens,
+        args,
+        device,
+        batch_size,
+        heads,
+        pad_seq_len,
+        sequece_packing_scales,
+        super_scale_lengths,
+        super_querysid_super_refsid,
+):
+    """
+    Build a flex attn function for a given scale schedule.
+    Args:
+        flex_attention: compiled flex attention
+        seq_l: seq length
+        prefix_lens: valid text prefix lens, [bs]
+        args: arguments
+        device: device
+        batch_size: batch size
+        heads: heads
+        pad_seq_len: pad_seq_len
+        sequece_packing_scales: list of scale schedule
+        querysid_refsid: list of scale_pack_info
+    Returns:
+        attn_fn: flex attn function
+    """
+    assert sum(super_scale_lengths) == seq_l, f'{sum(super_scale_lengths)}!= {seq_l}'
+    offsets = _length_to_offsets(super_scale_lengths, device=device)
+    mask_mod = _generate_overall_mask(offsets, super_querysid_super_refsid)
+    block_mask = create_block_mask(mask_mod, B = batch_size, H = heads, Q_LEN = seq_l, KV_LEN = seq_l, device = device, _compile = True)
+    attn_fn = partial(flex_attention, block_mask=block_mask)
+    return attn_fn

grn/models/fused_op.py

@@ -0,0 +1,27 @@
+import gc
+from copy import deepcopy
+from typing import Union
+
+import torch
+from torch import nn as nn
+from torch.nn import functional as F
+
+
+@torch.compile(fullgraph=True)
+def fused_rms_norm(x: torch.Tensor, weight: nn.Parameter, eps: float):
+    x = x.float()
+    return (x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True).add_(eps))) * weight
+
+
+@torch.compile(fullgraph=True)
+def fused_ada_layer_norm(C: int, eps: float, x: torch.Tensor, scale: torch.Tensor, shift: torch.Tensor):
+    x = x.float()
+    x = F.layer_norm(input=x, normalized_shape=(C,), weight=None, bias=None, eps=eps)
+    return x.mul(scale.add(1)).add_(shift)
+
+
+@torch.compile(fullgraph=True)
+def fused_ada_rms_norm(C: int, eps: float, x: torch.Tensor, scale: torch.Tensor, shift: torch.Tensor):
+    x = x.float()
+    x = (x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True).add_(eps)))
+    return x.mul(scale.add(1)).add_(shift)

grn/models/grn.py

@@ -0,0 +1,754 @@
+import json
+import math
+import time
+from contextlib import nullcontext
+from functools import partial
+from typing import Any, Dict, List, Optional, Tuple, Union
+
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import torch.utils.checkpoint
+import tqdm
+from timm.models import register_model
+
+import grn.utils_t2iv.dist as dist
+from grn.models.basic import FastRMSNorm, SelfAttnBlock
+from grn.models.rope import precompute_rope3d_freqs_grid
+from grn.schedules.dynamic_resolution import get_dynamic_resolution_meta
+from grn.utils_t2iv.dist import for_visualize
+from grn.utils_t2iv.hbq_util_t2iv import multiclass_labels2onehot_input
+from grn.utils_t2iv.sequence_parallel import SequenceParallelManager as sp_manager
+from grn.utils_t2iv.sequence_parallel import sp_gather_sequence_by_dim, sp_split_sequence_by_dim
+
+
+class MultipleLayers(nn.Module):
+    """A sequential container for a chunk of multiple transformer blocks."""
+
+    def __init__(self, layers: List[nn.Module], num_blocks: int, start_index: int):
+        super().__init__()
+        self.module = nn.ModuleList([
+            layers[i] for i in range(start_index, start_index + num_blocks)
+        ])
+
+    def forward(
+        self, x, cu_seqlens, max_seqlen, e0: Optional[torch.Tensor], 
+        attn_bias_or_two_vector: Optional[Any], attn_fn: Optional[Any] = None, 
+        checkpointing_full_block: bool = False, rope2d_freqs_grid: Optional[torch.Tensor] = None, 
+        scale_ind: Optional[Any] = None, context_info: Optional[Any] = None, 
+        last_diffusion_step: bool = True, ref_text_scale_inds: Optional[List[Any]] = None, 
+        use_cfg: bool = False, split_cond_uncond: Optional[List[Any]] = None
+    ) -> torch.Tensor:
+        h = x
+        for m in self.module:
+            if checkpointing_full_block:
+                h = torch.utils.checkpoint.checkpoint(
+                    m, h, cu_seqlens, max_seqlen, e0, attn_bias_or_two_vector, attn_fn, 
+                    rope2d_freqs_grid, scale_ind, context_info, 
+                    last_diffusion_step, ref_text_scale_inds, 
+                    use_cfg, split_cond_uncond, use_reentrant=False
+                )
+            else:
+                h = m(
+                    h, cu_seqlens, max_seqlen, e0, attn_bias_or_two_vector, attn_fn, 
+                    rope2d_freqs_grid, scale_ind, context_info, 
+                    last_diffusion_step, ref_text_scale_inds, 
+                    use_cfg, split_cond_uncond
+                )
+        return h
+
+
+def sinusoidal_embedding_1d(dim: int, position: torch.Tensor) -> torch.Tensor:
+    """
+    Generate 1D sinusoidal embeddings.
+    
+    Args:
+        dim (int): Embedding dimension (must be even).
+        position (torch.Tensor): Position tensor of shape [B, L].
+        
+    Returns:
+        torch.Tensor: Embeddings of shape [B, L, dim].
+    """
+    if dim % 2 != 0:
+        raise ValueError(f"Embedding dimension must be even, got {dim}")
+        
+    half = dim // 2
+    b, l = position.shape
+    position = position.reshape(-1).type(torch.float64)
+
+    sinusoid = torch.outer(
+        position, 
+        torch.pow(10000, -torch.arange(half).to(position).div(half))
+    )
+    x = torch.cat([torch.cos(sinusoid), torch.sin(sinusoid)], dim=1)
+    return x.reshape(b, l, dim)
+
+
+class TimestepEmbedder(nn.Module):
+    """Embeds scalar timesteps into vector representations."""
+
+    def __init__(self, hidden_size: int, frequency_embedding_size: int = 256):
+        super().__init__()
+        self.mlp = nn.Sequential(
+            nn.Linear(frequency_embedding_size, hidden_size, bias=True),
+            nn.SiLU(),
+            nn.Linear(hidden_size, hidden_size, bias=True),
+        )
+        self.frequency_embedding_size = frequency_embedding_size
+
+    @staticmethod
+    def timestep_embedding(t: torch.Tensor, dim: int, max_period: int = 10000) -> torch.Tensor:
+        """Create sinusoidal timestep embeddings."""
+        half = dim // 2
+        freqs = torch.exp(
+            -math.log(max_period) * torch.arange(start=0, end=half, dtype=torch.float32) / half
+        ).to(device=t.device)
+        args = t[:, None].float() * freqs[None]
+        embedding = torch.cat([torch.cos(args), torch.sin(args)], dim=-1)
+        if dim % 2:
+            embedding = torch.cat([embedding, torch.zeros_like(embedding[:, :1])], dim=-1)
+        return embedding
+
+    def forward(self, t: torch.Tensor) -> torch.Tensor:
+        t_freq = self.timestep_embedding(t, self.frequency_embedding_size)
+        return self.mlp(t_freq)
+
+
+def bld_to_bthwd(item: torch.Tensor, patch_time: int, patch_height: int, patch_width: int, apply_spatial_patchify: bool = False) -> torch.Tensor:
+    """Reshape a sequence tensor to a spatial tensor."""
+    batch_size = item.shape[0]
+    return item.reshape(batch_size, patch_time, patch_height, patch_width, -1)
+
+
+def build_attn_mask(seqlens, device):
+    attn_mask = torch.zeros((1, 1, sum(seqlens), sum(seqlens)), dtype=torch.bool, device=device)
+    q_start = 0
+    for i in range(len(seqlens)):
+        q_len = seqlens[i]
+        q_end = q_start + q_len
+        attn_mask[:, :, q_start:q_end, q_start:q_end] = True
+        q_start = q_end
+    return attn_mask
+
+
+class FsqHead(nn.Module):
+    """Classification head for Finite Scalar Quantization (FSQ)."""
+
+    def __init__(self, hidden_dim: int, fsq_dim: int, fsq_lvl: int, use_ada_layer_norm: bool, eps: float = 1e-6):
+        super().__init__()
+        self.proj = nn.Linear(hidden_dim, fsq_dim * fsq_lvl)
+        self.norm = FastRMSNorm(hidden_dim)
+
+    def forward(self, x: torch.Tensor, e: Optional[torch.Tensor] = None) -> torch.Tensor:
+        with torch.amp.autocast('cuda', dtype=torch.float32):
+            return self.proj(self.norm(x))
+
+
+class GRN(nn.Module):
+    def __init__(
+        self,
+        vae_local: Any,
+        arch: str = 'var',
+        qwen_qkvo_bias: bool = False,
+        text_channels: int = 0,
+        text_maxlen: int = 0,
+        embed_dim: int = 1024,
+        depth: int = 16,
+        num_key_value_heads: int = -1,
+        num_heads: int = 16,
+        mlp_ratio: float = 4.0,
+        drop_path_rate: float = 0.0,
+        norm_eps: float = 1e-6,
+        block_chunks: int = 1,
+        checkpointing: Optional[str] = None,
+        pad_to_multiplier: int = 0,
+        use_flex_attn: bool = False,
+        num_of_label_value: int = 2,
+        rope2d_normalized_by_hw: int = 0,
+        pn: Optional[str] = None,
+        video_frames: int = 1,
+        always_training_scales: int = 20,
+        apply_spatial_patchify: int = 0,
+        inference_mode: bool = False,
+        other_args: Optional[Any] = None,
+        **kwargs: Any,
+    ):
+        super().__init__()
+        # 1. Model Configuration
+        self.embed_dim = embed_dim
+        self.depth = depth
+        self.num_heads = num_heads
+        self.arch = arch
+        self.mlp_ratio = mlp_ratio
+        self.norm_eps = norm_eps
+        self.drop_path_rate = drop_path_rate
+        self.use_flex_attn = use_flex_attn
+        self.checkpointing = checkpointing
+        self.inference_mode = inference_mode
+        self.other_args = other_args
+
+        # 2. Embedding & Scale Configuration
+        self.vae_embed_dim = vae_local.codebook_dim
+        self.apply_spatial_patchify = apply_spatial_patchify
+        self.text_channels = text_channels
+        self.text_maxlen = text_maxlen
+        self.is_text_to_image = text_channels != 0
+
+        classifier_head_dim = other_args.detail_scale_dim
+        classifier_head_lvl = other_args.detail_num_lvl
+        hbq_round = other_args.hbq_round
+
+        if other_args.refine_mode in ['ar_discrete_GRN_ind']:
+            self.visual_embedding_in_dim = vae_local.codebook_dim * (2**hbq_round)
+            classifier_head_dim = vae_local.codebook_dim
+        elif other_args.refine_mode in ['ar_discrete_GRN_bit']:
+            self.visual_embedding_in_dim = hbq_round * vae_local.codebook_dim * 2
+            classifier_head_dim = hbq_round * vae_local.codebook_dim
+        else:
+            self.visual_embedding_in_dim = vae_local.codebook_dim
+
+        if self.apply_spatial_patchify:
+            self.visual_embedding_in_dim *= 4
+
+        # 3. Dynamic Resolution & Video Specifics
+        self.video_frames = video_frames
+        self.always_training_scales = always_training_scales
+        self.num_of_label_value = num_of_label_value
+        self.rope2d_normalized_by_hw = rope2d_normalized_by_hw
+
+        self.dynamic_resolution_h_w, self.h_div_w_templates = get_dynamic_resolution_meta(
+            other_args.dynamic_scale_schedule, other_args.train_h_div_w_list, other_args.video_frames
+        )
+        self.train_h_div_w_list = self.h_div_w_templates
+        print(f"train_h_div_w_list: {self.train_h_div_w_list}")
+
+        # 4. Utilities
+        self.entrophy_statistics = []
+        self.top_p, self.top_k = 1.0, 100
+        self.rng = torch.Generator(device=dist.get_device())
+        self.maybe_record_function = nullcontext
+        self.infer_ts = None
+
+        # 5. Model Components (Projections, Embeddings)
+        self.norm0_cond = nn.Identity()
+        self.text_proj = nn.Linear(self.text_channels, self.embed_dim)
+
+        if self.other_args.use_ada_layer_norm:
+            self.scale_or_time_dim = 256
+            self.scale_or_time_embedding = nn.Sequential(
+                nn.Linear(self.scale_or_time_dim, self.embed_dim), nn.SiLU(), nn.Linear(self.embed_dim, self.embed_dim),
+            )
+            self.scale_or_time_projection = nn.Sequential(nn.SiLU(), nn.Linear(self.embed_dim, self.embed_dim * 6))
+
+        tmp_h_div_w_template = self.train_h_div_w_list[0]
+
+        # RoPE grid initialization
+        with torch.amp.autocast('cuda', dtype=torch.float32):
+            self.rope2d_freqs_grid = precompute_rope3d_freqs_grid(
+                dim=self.embed_dim // self.num_heads,
+                rope2d_normalized_by_hw=self.rope2d_normalized_by_hw,
+                activated_h_div_w_templates=self.train_h_div_w_list,
+                max_scales=1010, # never used
+                max_frames=int(self.video_frames / other_args.temporal_compress_rate + 1),
+                max_height=1800 // 8, 
+                max_width=1800 // 8,
+                text_maxlen=self.text_maxlen,
+                args=other_args,
+            )
+
+        self.word_embed = nn.Linear(self.visual_embedding_in_dim, self.embed_dim)
+        self.head = FsqHead(
+            hidden_dim=self.embed_dim, 
+            fsq_dim=classifier_head_dim, 
+            fsq_lvl=classifier_head_lvl, 
+            use_ada_layer_norm=other_args.use_ada_layer_norm,
+        )
+
+        if other_args.add_scale_token > 0:
+            self.pt_embedder = TimestepEmbedder(self.embed_dim)
+
+        # 6. Transformer Blocks
+        self.attn_fn_compile_dict = {}
+        self.unregistered_blocks = []
+        for block_idx in range(depth):
+            block = SelfAttnBlock(
+                embed_dim=self.embed_dim, 
+                num_heads=num_heads, 
+                num_key_value_heads=num_key_value_heads, 
+                mlp_ratio=mlp_ratio,
+                use_flex_attn=use_flex_attn,
+                qwen_qkvo_bias=qwen_qkvo_bias,
+                use_ada_layer_norm=other_args.use_ada_layer_norm,
+            )
+            self.unregistered_blocks.append(block)
+
+        self.num_block_chunks = block_chunks or 1
+        self.num_blocks_in_a_chunk = depth // self.num_block_chunks
+        assert self.num_blocks_in_a_chunk * self.num_block_chunks == depth, "Depth must be divisible by block_chunks"
+        
+        self.block_chunks = nn.ModuleList([
+            MultipleLayers(self.unregistered_blocks, self.num_blocks_in_a_chunk, i * self.num_blocks_in_a_chunk)
+            for i in range(self.num_block_chunks)
+        ])
+
+        print(f"    [Model Config] embed_dim={embed_dim}, num_heads={num_heads}, depth={depth}, "
+              f"mlp_ratio={mlp_ratio}, num_blocks_in_a_chunk={self.num_blocks_in_a_chunk}")
+        print(f"    drop_path_rate={drop_path_rate:g}", end='\n\n', flush=True)
+        
+    def get_loss_acc(
+        self,
+        hidden_states: torch.Tensor,
+        hidden_states_mask: Optional[torch.Tensor],
+        e: Optional[torch.Tensor],
+        sequence_packing_scales: List[List[Tuple[int, int, int]]],
+        gt: List[torch.Tensor],
+        other_info_by_scale: List[Dict[str, Any]],
+        return_last_hidden_states: bool
+    ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+        """
+        Calculate loss and accuracy for the predicted logits.
+        
+        Args:
+            hidden_states: shaped (B, L, C)
+            hidden_states_mask: Optional mask for hidden states
+            e: scale or time embeddings
+            sequence_packing_scales: List of scales for sequence packing
+            gt: Ground truth labels
+            other_info_by_scale: Meta information for each scale
+            return_last_hidden_states: Whether to return the last hidden states
+            
+        Returns:
+            Tuple of (logits_norm, loss_list, acc_list)
+        """
+        logits_norm = []
+        logits_full = self.head(hidden_states, e)
+        global_token_ptr, global_scale_ptr = 0, 0
+        loss_list, acc_list = [], []
+
+        for pack_scales in sequence_packing_scales:
+            for pt, ph, pw in pack_scales:
+                mul_pt_ph_pw = pt * ph * pw
+                cur_bits = other_info_by_scale[global_scale_ptr]['cur_bits']
+                cur_lvl = other_info_by_scale[global_scale_ptr]['cur_lvl']
+                predict_tokens = other_info_by_scale[global_scale_ptr]['predict_tokens']
+                all_tokens = other_info_by_scale[global_scale_ptr]['all_tokens']
+                logits = logits_full[:, global_token_ptr:global_token_ptr + predict_tokens]
+                logits = logits.reshape(hidden_states.shape[0], mul_pt_ph_pw, cur_bits, cur_lvl)
+                logits = logits.permute(0, 3, 1, 2) # [1, num_of_label_value, mul_pt_ph_pw, d]
+                
+                logits_norm.append(logits.abs().mean())
+                
+                # gt[global_scale_ptr]: [1, mul_pt_ph_pw, d]
+                loss_this_scale = F.cross_entropy(logits, gt[global_scale_ptr], reduction='none')[0] # [mul_pt_ph_pw, d]
+                acc_this_scale = (logits.argmax(1) == gt[global_scale_ptr]).float()[0] # [mul_pt_ph_pw, d]
+                
+                loss_list.append(loss_this_scale.mean(-1))
+                acc_list.append(acc_this_scale.mean(-1))
+                
+                global_scale_ptr += 1
+                global_token_ptr += all_tokens
+                
+        loss_tensor = torch.cat(loss_list) if loss_list else torch.tensor([], device=hidden_states.device)
+        acc_tensor = torch.cat(acc_list) if acc_list else torch.tensor([], device=hidden_states.device)
+        logits_norm_tensor = torch.stack(logits_norm).mean() if logits_norm else torch.tensor(0.0, device=hidden_states.device)
+        
+        return logits_norm_tensor, loss_tensor, acc_tensor
+    
+    def get_logits_during_infer(self, hidden_states: torch.Tensor, e: Optional[torch.Tensor] = None) -> torch.Tensor:
+        """Get logits during inference."""
+        return self.head(hidden_states.float(), e)
+    
+    def forward(
+        self,
+        label_B_or_BLT: Union[torch.LongTensor, Tuple[torch.FloatTensor, torch.IntTensor, int]],
+        x_BLC: torch.Tensor,
+        visual_rope_cache: Optional[List[torch.Tensor]] = None,
+        sequece_packing_scales: Optional[List[List[Tuple[int, int, int]]]] = None,
+        super_scale_lengths: Optional[List[int]] = None,
+        other_info_by_scale: Optional[List[Dict[str, Any]]] = None,
+        gt_BL: Optional[List[torch.Tensor]] = None,
+        x_BLC_mask: Optional[torch.Tensor] = None,
+        scale_or_time_ids: Optional[torch.Tensor] = None,
+        return_last_hidden_states: bool = False,
+        **kwargs: Any,
+    ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, float]:
+        """
+        Forward pass for the GRN model.
+        
+        Args:
+            label_B_or_BLT: Text conditions or labels
+            x_BLC: Input sequence hidden states
+            visual_rope_cache: Cache for visual RoPE embeddings
+            sequece_packing_scales: Scales for sequence packing
+            super_scale_lengths: Lengths of super scales
+            other_info_by_scale: Meta info for scales
+            gt_BL: Ground truth
+            x_BLC_mask: Mask for input sequence
+            scale_or_time_ids: IDs for scale or time embeddings
+            return_last_hidden_states: Whether to return last hidden states
+            
+        Returns:
+            Tuple of (logits_norm, loss_list, acc_list, valid_sequence_ratio)
+        """
+        device = x_BLC[0].device
+
+        # [1. get input sequence x_BLC]
+        # word embedding
+        sub_L_list = [item.shape[1] for item in x_BLC]
+        cat_x_BLC = torch.cat(x_BLC, dim=1)
+        with torch.amp.autocast('cuda', dtype=torch.float32):
+            cat_x_BLC = self.word_embed(cat_x_BLC.float())
+        x_BLC = list(torch.split(cat_x_BLC, sub_L_list, dim=1))
+
+        # text tokens embedding
+        kv_compact, lens, cu_seqlens_k, max_seqlen_k, _ = label_B_or_BLT
+        with torch.amp.autocast('cuda', dtype=torch.float32):
+            kv_compact = self.text_proj(kv_compact).contiguous() # [sum(lens), C]
+        kv_compact_splits = torch.split(kv_compact, lens, dim=0)
+
+        # scale tokens embedding
+        scale_token_ids = torch.tensor([info["scale_token_id"] for info in other_info_by_scale], device=device)
+        with torch.amp.autocast("cuda", dtype=torch.float32):
+            pt_tokens = self.pt_embedder((scale_token_ids)) # [num_scales, C]
+        
+        # construct final X_BLC input, [visual token, text token, scale token]
+        x_BLC_lists = []
+        for i in range(len(x_BLC)):
+            x_BLC_lists.extend([x_BLC[i], kv_compact_splits[i].unsqueeze(0), pt_tokens[i][None, None]])
+        x_BLC = torch.cat(x_BLC_lists, dim=1)
+
+        valid_sequence_ratio = x_BLC.shape[1] / self.other_args.train_max_token_len
+        attn_fn, attn_bias_or_two_vector = None, None
+
+        # calculate finalrope cache, [visual token, text token, scale token]
+        self.rope2d_freqs_grid['freqs_text'] = self.rope2d_freqs_grid['freqs_text'].to(x_BLC.device)
+        rope_cache_list = []
+        for i in range(len(visual_rope_cache)):
+            rope_cache_list.append(visual_rope_cache[i])
+            rope_cache_list.append(self.rope2d_freqs_grid['freqs_text'][:,:,:,:,:lens[i]])
+            rope_cache_list.append(self.rope2d_freqs_grid['freqs_text'][:,:,:,:,512:512+self.other_args.add_scale_token])
+        rope_cache = torch.cat(rope_cache_list, dim=4) # (2, 1, 1, 1, seq_len, head_dim / 2)
+        assert rope_cache.shape[4] == x_BLC.shape[1], f'{rope_cache.shape[4]} != {x_BLC.shape[1]}'
+        rope_cache = rope_cache[:,0].permute(0, 1, 3, 2, 4) # (2, 1, 1, 1, seq_len, head_dim / 2) -> (2, 1, 1, seq_len, head_dim / 2) -> (2, 1, seq_len, 1, head_dim / 2)
+
+        # calculate time or scale embeddings
+        if self.other_args.use_ada_layer_norm:
+            with torch.amp.autocast('cuda', dtype=torch.float32):
+                e = self.scale_or_time_embedding(sinusoidal_embedding_1d(self.scale_or_time_dim, scale_or_time_ids).float()) # [1, visual_seq_len,] -> [1, visual_seq_len, 256] -> [1, visual_seq_len, C]
+                if e.shape[1] < x_BLC.shape[1]:
+                    e = F.pad(e, (0,0,0,x_BLC.shape[1]-e.shape[1]), 'constant', 0.) # [1, visual_seq_len, C] -> [1, L, C]
+                e0 = self.scale_or_time_projection(e).unflatten(2, (6, self.C)) # [1, L, C] -> [1, L, 6C] -> [1, L, 6, C]
+                assert e.dtype == torch.float32 and e0.dtype == torch.float32
+        else:
+            e, e0 = None, None
+        
+        # [2. block loop]
+        checkpointing_full_block = self.checkpointing == 'full-block' and self.training
+
+        if sp_manager.sp_on():
+            # [B, raw_L, C] --> [B, raw_L/sp_size, C]
+            x_BLC = sp_split_sequence_by_dim(x_BLC, 1)
+
+        cu_seqlens = torch.tensor([0]+super_scale_lengths, device=device).cumsum(-1).to(torch.int32)
+        max_seqlen = max(super_scale_lengths)
+        for i, chunk in enumerate(self.block_chunks): # this path
+            x_BLC = chunk(x=x_BLC, cu_seqlens=cu_seqlens, max_seqlen=max_seqlen, e0=e0, attn_bias_or_two_vector=attn_bias_or_two_vector, attn_fn=attn_fn, checkpointing_full_block=checkpointing_full_block, rope2d_freqs_grid=rope_cache)
+
+        if sp_manager.sp_on():
+            # [B, raw_L/sp_size, C] --> [B, raw_L, C]
+            x_BLC = sp_gather_sequence_by_dim(x_BLC, 1)
+
+        # [3. unpad the seqlen dim, and then get logits]
+        logits_norm, loss_list, acc_list = self.get_loss_acc(x_BLC, x_BLC_mask, e, sequece_packing_scales, gt_BL, other_info_by_scale, return_last_hidden_states)
+        return logits_norm, loss_list, acc_list, valid_sequence_ratio
+
+    def prepare_text_conditions(
+        self,
+        label_B_or_BLT: Tuple[torch.Tensor, ...],
+        negative_label_B_or_BLT: Optional[Tuple[torch.Tensor, ...]],
+        use_cfg: bool = False,
+    ) -> Tuple[torch.Tensor, List[int]]:
+        """Prepare text conditions for inference."""
+        kv_compact, lens, cu_seqlens_k, max_seqlen_k = label_B_or_BLT
+        if use_cfg:
+            kv_compact_un, lens_un, cu_seqlens_k_un, max_seqlen_k_un = negative_label_B_or_BLT
+            kv_compact = torch.cat((kv_compact, kv_compact_un), dim=0)
+            cu_seqlens_k = torch.cat((cu_seqlens_k, cu_seqlens_k_un[1:] + cu_seqlens_k[-1]), dim=0)
+            max_seqlen_k = max(max_seqlen_k, max_seqlen_k_un)
+            lens = lens + lens_un
+        kv_compact = self.text_proj(kv_compact).contiguous()
+        return kv_compact, lens
+    
+    def embeds_codes2input(self, last_stage: torch.Tensor) -> torch.Tensor:
+        """Embed discrete codes into continuous input representations."""
+        last_stage = last_stage.reshape(*last_stage.shape[:2], -1) # [B, d, t*h*w] or [B, 4d, t*h*w]
+        last_stage = torch.permute(last_stage, [0, 2, 1]) # [B, t*h*w, d] or [B, t*h*w, 4d]
+        last_stage = self.word_embed(last_stage) # norm0_ve is Identity
+        return last_stage
+    
+    @torch.no_grad()
+    def autoregressive_infer(
+        self,
+        vae: Optional[Any] = None,
+        scale_schedule: Optional[List[Tuple[int, int, int]]] = None,
+        label_B_or_BLT: Optional[List[Tuple[torch.Tensor, ...]]] = None,
+        negative_label_B_or_BLT: Optional[List[Tuple[torch.Tensor, ...]]] = None,
+        g_seed: Optional[int] = None,
+        cfg_list: Optional[List[float]] = None,
+        tau_list: Optional[List[float]] = None,
+        gt_leak: int = 0,
+        args: Optional[Any] = None,
+        get_visual_rope_embeds: Optional[Any] = None,
+        noise_list: Optional[List[torch.Tensor]] = None,
+        uncond_class_token_id: int = 1000,
+        first_frame_condition: bool = False,
+        **kwargs: Any,
+    ):
+        """Autoregressive inference loop for the GRN model."""
+        if cfg_list is None: cfg_list = []
+        if tau_list is None: tau_list = []
+        
+        from grn.schedules.global_refine import shift_pt
+        
+        rng = None
+        assert len(cfg_list) >= len(scale_schedule), "Not enough CFG values for scales"
+        assert len(tau_list) >= len(scale_schedule), "Not enough tau values for scales"
+
+        ret, idx_Bl_list = [], []  # current length, list of reconstructed images
+        for b in self.unregistered_blocks: b.attn.kv_caching(True)
+        total_steps = args.max_infer_steps
+        pbar = tqdm.tqdm(total=total_steps)
+        block_chunks = self.block_chunks if self.num_block_chunks > 1 else self.blocks
+        use_cfg = True
+        cfg_interval = float(args.cfg_type.split('_')[-1])
+        full_pt, ph, pw = scale_schedule[0]
+        if first_frame_condition:
+            pt = full_pt - 1
+            visual_rope_cache = get_visual_rope_embeds(self.rope2d_freqs_grid, (pt, ph, pw), 'cuda', args.mapped_h_div_w_template, t_offset=1)
+        else:
+            pt = full_pt
+            visual_rope_cache = get_visual_rope_embeds(self.rope2d_freqs_grid, (pt, ph, pw), 'cuda', args.mapped_h_div_w_template, t_offset=0)
+
+        # text tokens forward
+        self.rope2d_freqs_grid['freqs_text'] = self.rope2d_freqs_grid['freqs_text'].to('cuda')
+        prefix_tokens, lens = self.prepare_text_conditions(label_B_or_BLT[0], negative_label_B_or_BLT, use_cfg)
+        device = prefix_tokens.device
+        infer_device, infer_dtype = prefix_tokens.device, prefix_tokens.dtype
+        prefix_tokens = torch.split(prefix_tokens, lens, dim=0)
+        rope_cache_text_cond = self.rope2d_freqs_grid['freqs_text'][:,:,:,:,:lens[0]]
+        rope_cache_text_uncond = self.rope2d_freqs_grid['freqs_text'][:,:,:,:,:lens[1]]
+
+        if args.refine_mode in ['ar_discrete_GRN_bit']:
+            classes = 2
+            labels_shape = (1,args.detail_scale_dim*args.hbq_round,pt,ph,pw)
+        elif args.refine_mode in ['ar_discrete_GRN_index']:
+            classes = 2**args.hbq_round
+            labels_shape = (1,args.detail_scale_dim,pt,ph,pw)
+
+        mul_pt_ph_pw = pt * ph * pw
+        repeat_idx = -1
+        scale_token_rope_cache = self.rope2d_freqs_grid['freqs_text'][:,:,:,:,512:512+args.add_scale_token]
+        if noise_list is not None:
+            absolute_gt_labels = noise_list[0].to('cuda').permute(0,2,3,4,1) # [B,d,t,h,w] -> [B,t,h,w,d]
+        assert len(scale_schedule) == 1
+        if first_frame_condition:
+            first_frame_labels = noise_list[0][:,:,:1] # [B,d,1,h,w]
+            first_frame_tokens_cond = self.embeds_codes2input(multiclass_labels2onehot_input(first_frame_labels, classes))
+            fist_frame_rope_cache = get_visual_rope_embeds(self.rope2d_freqs_grid, (1, ph, pw), device, args.mapped_h_div_w_template, t_offset=0)
+            visual_rope_cache = torch.cat((visual_rope_cache, fist_frame_rope_cache), dim=4)
+            tmp_seqlens = [mul_pt_ph_pw + ph * pw + lens[0] + args.add_scale_token, mul_pt_ph_pw + lens[1] + ph * pw + args.add_scale_token]
+        else:
+            tmp_seqlens = [mul_pt_ph_pw+lens[0]+args.add_scale_token, mul_pt_ph_pw+lens[1]+args.add_scale_token]
+
+        # [visual tokens, text tokens, pt tokens]
+        rope_cache = torch.cat([visual_rope_cache, rope_cache_text_cond, scale_token_rope_cache, visual_rope_cache, rope_cache_text_uncond, scale_token_rope_cache], dim=4) # (2, 1, 1, 1, seq_len, dim / 2)
+        rope_cache = rope_cache[:,0].permute(0, 1, 3, 2, 4) # (2, 1, 1, 1, seq_len, dim / 2) -> (2, 1, 1, seq_len, dim / 2) -> (2, 1, seq_len, 1, dim / 2)
+
+        cu_seqlens = torch.tensor([0]+tmp_seqlens, device=device).cumsum(-1).to(torch.int32)
+        max_seqlen = max(tmp_seqlens)
+
+        pure_rand_labels = torch.randint(low=0, high=classes, size=labels_shape, device=infer_device, dtype=infer_dtype)
+        mixed_xt = pure_rand_labels
+        next_pt = 0.
+        attn_mask = build_attn_mask(tmp_seqlens, device) if args.use_slow_attn else None
+        for cur_inner_round_si in range(args.max_infer_steps):
+            cur_pt = next_pt
+            is_last_step = np.abs(cur_pt - 1) < 0.02
+            if cur_inner_round_si == 0:
+                self.entrophy_statistics.append([])
+            repeat_idx += 1 # index scale tokens, very important
+            cfg = cfg_list[0] if cur_pt >= cfg_interval else 1.0
+            last_stage = self.embeds_codes2input(multiclass_labels2onehot_input(mixed_xt, classes))
+            pt_tokens = self.pt_embedder(torch.tensor([cur_pt], device=device)).unsqueeze(0)
+            # [visual tokens, text tokens, pt tokens]
+            if first_frame_condition:
+                last_stage_cond = torch.cat((last_stage, first_frame_tokens_cond, prefix_tokens[0].unsqueeze(0), pt_tokens), dim=1)
+                last_stage_uncond = torch.cat((last_stage, first_frame_tokens_cond, prefix_tokens[1].unsqueeze(0), pt_tokens), dim=1)
+            else:
+                last_stage_cond = torch.cat((last_stage, prefix_tokens[0].unsqueeze(0), pt_tokens), dim=1)
+                last_stage_uncond = torch.cat((last_stage, prefix_tokens[1].unsqueeze(0), pt_tokens), dim=1)
+            last_stage = torch.cat([last_stage_cond, last_stage_uncond], dim=1)
+
+            e, e0 = None, None
+            last_diffusion_step = False
+            for block_idx, b in enumerate(block_chunks):
+                last_stage = b(x=last_stage, cu_seqlens=cu_seqlens, max_seqlen=max_seqlen, e0=e0, attn_bias_or_two_vector=attn_mask, attn_fn=None, rope2d_freqs_grid=rope_cache, last_diffusion_step=last_diffusion_step)
+            logits = self.get_logits_during_infer(last_stage, e=e)
+            tmp_bs, tmp_seq_len = logits.shape[:2]
+            logits = logits.reshape(tmp_bs, tmp_seq_len, -1, args.detail_num_lvl) # [B,thw+...,d,2]
+            pred_cond_logits = logits[:,:mul_pt_ph_pw] # [B,thw,d,2]
+            pred_uncond_logits = logits[:,tmp_seqlens[0]:tmp_seqlens[0]+mul_pt_ph_pw] # [B,thw,d,2]
+            pred_cond_probs = pred_cond_logits.softmax(-1) # [B,thw,d,2]
+            categories = pred_cond_logits.shape[-1]
+            entrophy = (-pred_cond_probs * torch.log2(pred_cond_probs)).sum(-1).mean().item() / np.log2(categories)
+
+            pt_unshift = (cur_inner_round_si + 1) / (args.complexity_aware_Tmax - 1)
+            pt_shift = shift_pt(min(1., pt_unshift), args.snr_shift)
+            next_pt = 1 - np.cos(np.pi/2*pt_shift)
+            next_pt = next_pt * 0.999
+
+            pred_cond_labels = torch.argmax(pred_cond_probs, dim=-1) # [B,thw,d]
+            pred_cond_labels = bld_to_bthwd(pred_cond_labels, pt, ph, pw)
+            if cfg != 1:
+                pred_cfg_logits = pred_uncond_logits + cfg * (pred_cond_logits - pred_uncond_logits)
+            else:
+                pred_cfg_logits = pred_cond_logits
+            pred_cfg_logits = pred_cfg_logits.mul(1/tau_list[0]) # [B,thw,d,2]
+            pred_cfg_probs = pred_cfg_logits.softmax(dim=-1) # [B,thw,d,2]
+            pred_cfg_labels = torch.argmax(pred_cfg_probs, dim=-1) # [B,thw,d]
+            pred_cfg_labels = bld_to_bthwd(pred_cfg_labels, pt, ph, pw) # [B,t,h,w,d]
+            pred_sample_labels = torch.multinomial(pred_cfg_probs.view(-1, args.detail_num_lvl), num_samples=1, replacement=True, generator=rng).view(tmp_bs, mul_pt_ph_pw, -1) # [B, thw,d]
+            pred_sample_probs = torch.gather(pred_cfg_probs, dim=3, index=pred_sample_labels.unsqueeze(-1)).squeeze(-1) # [B,thw,d]
+            pred_sample_probs = bld_to_bthwd(pred_sample_probs, pt, ph, pw) # [B,t,h,w,d]
+            pred_sample_labels = bld_to_bthwd(pred_sample_labels, pt, ph, pw) # [B,t,h,w,d]
+
+            assume_flip_ratio = (1 - cur_pt) / args.detail_num_lvl * 100. # different ratio between prediciton and input
+            pred_zero_ratio = (pred_cond_labels == 0).sum() / pred_cond_labels.numel() * 100.
+            pred_one_ratio = (pred_cond_labels == 1).sum() / pred_cond_labels.numel() * 100.
+            mixed_xt_Bthwd_01 = mixed_xt.clone().permute(0,2,3,4,1)
+            mixed_xt_Bthwd_01[mixed_xt_Bthwd_01<0] = 0
+            pred_cond_flip_ratio = (pred_cond_labels != mixed_xt_Bthwd_01).sum() / pred_cond_labels.numel() * 100.
+            pred_cfg_flip_ratio = (pred_cfg_labels != mixed_xt_Bthwd_01).sum() / pred_cfg_labels.numel() * 100.
+            pred_sample_flip_ratio = (pred_sample_labels != mixed_xt_Bthwd_01).sum() / pred_sample_labels.numel() * 100.
+            self.entrophy_statistics[-1].append({
+                'cur_inner_round_si': cur_inner_round_si,
+                'cur_pt': cur_pt,
+                # 'cur_tau': cur_tau,
+                # 'cur_cfg': cur_cfg,
+                'entrophy': entrophy,
+                'assume_flip_ratio': assume_flip_ratio,
+                'pred_cond_flip_ratio': pred_cond_flip_ratio.item(),
+                'pred_cfg_flip_ratio': pred_cfg_flip_ratio.item(),
+                'pred_sample_flip_ratio': pred_sample_flip_ratio.item(),
+                'pred_zero_ratio': pred_zero_ratio.item(),
+                'pred_one_ratio': pred_one_ratio.item(),
+                'meta': args.meta,
+            })
+            print(f'{repeat_idx=} {cur_inner_round_si=} {cur_pt=:.3f} {pred_sample_labels.shape=}')
+            print(f'{assume_flip_ratio=:.2f}% {pred_cond_flip_ratio=:.2f}% {pred_cfg_flip_ratio=:.2f}% {pred_sample_flip_ratio=:.2f}%')
+            if repeat_idx < gt_leak:
+                gt_labels = absolute_gt_labels
+                gt_flip_ratio = (gt_labels != mixed_xt_Bthwd_01).sum() / gt_labels.numel() * 100.
+                gt_flip_ratio = gt_flip_ratio.item()
+                pred_cond_acc = (gt_labels==pred_cond_labels).to(float).mean().item()
+                pred_cfg_acc = (gt_labels==pred_cfg_labels).to(float).mean().item()
+                pred_sample_acc = (gt_labels==pred_sample_labels).to(float).mean().item()
+                print(f'{repeat_idx=} {entrophy=:.4f} {pred_cond_acc=:.4f} {pred_cfg_acc=:.4f} {pred_sample_acc=:.4f}')
+                self.entrophy_statistics[-1][-1].update({
+                    'gt_flip_ratio': gt_flip_ratio,
+                    'pred_cond_acc': pred_cond_acc,
+                    'pred_cfg_acc': pred_cfg_acc,
+                    'pred_sample_acc': pred_sample_acc,
+                })
+                pred_sample_labels = gt_labels
+            
+            pred_sample_labels = pred_sample_labels.permute(0,4,1,2,3) # [B,t,h,w,d] -> [B,d,t,h,w]
+            pred_sample_probs = pred_sample_probs.permute(0,4,1,2,3) # [B,t,h,w,d] -> [B,d,t,h,w]
+            use_predict_mask = torch.rand(pred_sample_labels.shape, device=device) < next_pt
+            mixed_xt = torch.where(use_predict_mask, pred_sample_labels, pure_rand_labels)
+            next_pt = use_predict_mask.float().mean().item()
+            pbar.update(1)
+            if is_last_step: break
+
+        if first_frame_condition:
+            pred_sample_labels = torch.cat((first_frame_labels, pred_sample_labels), dim=2)
+         
+        if args.refine_mode == 'ar_discrete_GRN_ind':
+            from grn.utils_t2iv.hbq_util_t2iv import index_label2quant_features
+            approx_signal = index_label2quant_features(pred_sample_labels, hbq_round=args.hbq_round)
+        elif args.refine_mode == 'ar_discrete_GRN_bit':
+            from grn.utils_t2iv.hbq_util_t2iv import bit_label2raw_feature
+            approx_signal = bit_label2raw_feature(pred_sample_labels, hbq_round=args.hbq_round) # [B, hbq_round_mul_d, t, h, w] -> [B,d,t,h,w]
+        for b in self.unregistered_blocks: b.attn.kv_caching(False)
+        img = self.summed_codes2images(vae, approx_signal)
+        return ret, idx_Bl_list, img
+
+    def summed_codes2images(self, vae: Any, summed_codes: torch.Tensor) -> torch.Tensor:
+        """Decode summed codes into images using the VAE."""
+        t1 = time.time()
+        img = vae.decode(summed_codes, slice=True)
+        img = (img + 1) / 2
+        img = torch.clamp(img, 0, 1)
+        img = img.permute(0, 2, 3, 4, 1) # [bs, 3, t, h, w] -> [bs, t, h, w, 3]
+        img = img.mul_(255).to(torch.uint8).flip(dims=(4,))
+        print(f"Decode takes {time.time() - t1:.1f}s")
+        return img # bgr order
+
+    @for_visualize
+    def vis_key_params(self, ep: int) -> None:
+        return
+    
+    def load_state_dict(self, state_dict: Dict[str, Any], strict: bool = False, assign: bool = False) -> Any:       
+        return super().load_state_dict(state_dict=state_dict, strict=strict, assign=assign)
+    
+    def special_init(self, **kwargs: Any) -> None:
+        """Apply special initialization to specific layers."""
+        std = 0.02
+        for name, module in self.named_modules():
+            if isinstance(module, nn.Linear):
+                module.weight.data.normal_(mean=0.0, std=std)
+                if module.bias is not None:
+                    module.bias.data.zero_()
+            elif isinstance(module, nn.Embedding):
+                module.weight.data.normal_(mean=0.0, std=std)
+                if module.padding_idx is not None:
+                    module.weight.data[module.padding_idx].zero_()
+    
+    def extra_repr(self) -> str:
+        return f'drop_path_rate={self.drop_path_rate}'
+    
+    def get_layer_id_and_scale_exp(self, para_name: str) -> Any:
+        raise NotImplementedError
+
+TIMM_KEYS = {'img_size', 'pretrained', 'pretrained_cfg', 'pretrained_cfg_overlay', 'global_pool'}
+
+@register_model
+def GRN0b(depth: int = 4, block_chunks: int = 2, embed_dim: int = 512, num_heads: int = 4, num_key_value_heads: int = 4, drop_path_rate: float = 0.0, **kwargs: Any) -> GRN: 
+    return GRN(
+        arch='qwen',
+        qwen_qkvo_bias=False,
+        depth=depth,
+        block_chunks=block_chunks,
+        embed_dim=embed_dim,
+        num_heads=num_heads,
+        num_key_value_heads=num_key_value_heads,
+        mlp_ratio=3.55,
+        drop_path_rate=drop_path_rate,
+        **{k: v for k, v in kwargs.items() if k not in TIMM_KEYS}
+    )
+
+@register_model
+def GRN2b(depth: int = 28, block_chunks: int = 7, embed_dim: int = 2304, num_heads: int = 18, num_key_value_heads: int = 18, drop_path_rate: float = 0.0, **kwargs: Any) -> GRN: 
+    return GRN(
+        arch='qwen',
+        qwen_qkvo_bias=False,
+        depth=depth,
+        block_chunks=block_chunks,
+        embed_dim=embed_dim,
+        num_heads=num_heads,
+        num_key_value_heads=num_key_value_heads,
+        mlp_ratio=3.55,
+        drop_path_rate=drop_path_rate,
+        **{k: v for k, v in kwargs.items() if k not in TIMM_KEYS}
+    )

grn/models/grn_c2i.py

@@ -0,0 +1,399 @@
+# --------------------------------------------------------
+# References:
+# SiT: https://github.com/willisma/SiT
+# Lightning-DiT: https://github.com/hustvl/LightningDiT
+# --------------------------------------------------------
+import torch
+import torch.nn as nn
+import numpy as np
+import math
+import torch.nn.functional as F
+from grn.utils_c2i.model_util import VisionRotaryEmbeddingFast, get_2d_sincos_pos_embed, RMSNorm
+from grn.utils_c2i.hbq_util_c2i import multiclass_labels2onehot_input
+
+def modulate(x, shift, scale):
+    return x * (1 + scale.unsqueeze(1)) + shift.unsqueeze(1)
+
+
+class BottleneckPatchEmbed(nn.Module):
+    """ Image to Patch Embedding
+    """
+    def __init__(self, img_size=224, patch_size=16, in_chans=3, pca_dim=768, embed_dim=768, bias=True):
+        super().__init__()
+        img_size = (img_size, img_size)
+        patch_size = (patch_size, patch_size)
+        num_patches = (img_size[1] // patch_size[1]) * (img_size[0] // patch_size[0])
+        self.img_size = img_size
+        self.patch_size = patch_size
+        self.num_patches = num_patches
+
+        self.proj1 = nn.Conv2d(in_chans, pca_dim, kernel_size=patch_size, stride=patch_size, bias=False)
+        self.proj2 = nn.Conv2d(pca_dim, embed_dim, kernel_size=1, stride=1, bias=bias)
+
+    def forward(self, x):
+        B, C, H, W = x.shape
+        assert H == self.img_size[0] and W == self.img_size[1], \
+            f"Input image size ({H}*{W}) doesn't match model ({self.img_size[0]}*{self.img_size[1]})."
+        x = self.proj2(self.proj1(x)).flatten(2).transpose(1, 2)
+        return x
+
+
+class TimestepEmbedder(nn.Module):
+    """
+    Embeds scalar timesteps into vector representations.
+    """
+    def __init__(self, hidden_size, frequency_embedding_size=256):
+        super().__init__()
+        self.mlp = nn.Sequential(
+            nn.Linear(frequency_embedding_size, hidden_size, bias=True),
+            nn.SiLU(),
+            nn.Linear(hidden_size, hidden_size, bias=True),
+        )
+        self.frequency_embedding_size = frequency_embedding_size
+
+    @staticmethod
+    def timestep_embedding(t, dim, max_period=10000):
+        """
+        Create sinusoidal timestep embeddings.
+        :param t: a 1-D Tensor of N indices, one per batch element.
+                          These may be fractional.
+        :param dim: the dimension of the output.
+        :param max_period: controls the minimum frequency of the embeddings.
+        :return: an (N, D) Tensor of positional embeddings.
+        """
+        # https://github.com/openai/glide-text2im/blob/main/glide_text2im/nn.py
+        half = dim // 2
+        freqs = torch.exp(
+            -math.log(max_period) * torch.arange(start=0, end=half, dtype=torch.float32) / half
+        ).to(device=t.device)
+        args = t[:, None].float() * freqs[None]
+        embedding = torch.cat([torch.cos(args), torch.sin(args)], dim=-1)
+        if dim % 2:
+            embedding = torch.cat([embedding, torch.zeros_like(embedding[:, :1])], dim=-1)
+        return embedding
+
+    def forward(self, t):
+        t_freq = self.timestep_embedding(t, self.frequency_embedding_size)
+        t_emb = self.mlp(t_freq)
+        return t_emb
+
+
+class LabelEmbedder(nn.Module):
+    """
+    Embeds class labels into vector representations. Also handles label dropout for classifier-free guidance.
+    """
+    def __init__(self, num_classes, hidden_size):
+        super().__init__()
+        self.embedding_table = nn.Embedding(num_classes + 1, hidden_size)
+        self.num_classes = num_classes
+
+    def forward(self, labels):
+        embeddings = self.embedding_table(labels)
+        return embeddings
+
+
+def scaled_dot_product_attention(query, key, value, dropout_p=0.0) -> torch.Tensor:
+    L, S = query.size(-2), key.size(-2)
+    scale_factor = 1 / math.sqrt(query.size(-1))
+    attn_bias = torch.zeros(query.size(0), 1, L, S, dtype=query.dtype).cuda()
+
+    with torch.cuda.amp.autocast(enabled=False):
+        attn_weight = query.float() @ key.float().transpose(-2, -1) * scale_factor
+    attn_weight += attn_bias
+    attn_weight = torch.softmax(attn_weight, dim=-1)
+    attn_weight = torch.dropout(attn_weight, dropout_p, train=True)
+    return attn_weight @ value
+
+
+class Attention(nn.Module):
+    def __init__(self, dim, num_heads=8, qkv_bias=True, qk_norm=True, attn_drop=0., proj_drop=0.):
+        super().__init__()
+        self.num_heads = num_heads
+        head_dim = dim // num_heads
+
+        self.q_norm = RMSNorm(head_dim) if qk_norm else nn.Identity()
+        self.k_norm = RMSNorm(head_dim) if qk_norm else nn.Identity()
+
+        self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias)
+        self.attn_drop = nn.Dropout(attn_drop)
+        self.proj = nn.Linear(dim, dim)
+        self.proj_drop = nn.Dropout(proj_drop)
+
+    def forward(self, x, rope):
+        B, N, C = x.shape
+        qkv = self.qkv(x).reshape(B, N, 3, self.num_heads, C // self.num_heads).permute(2, 0, 3, 1, 4)
+        q, k, v = qkv[0], qkv[1], qkv[2]   # make torchscript happy (cannot use tensor as tuple)
+
+        q = self.q_norm(q)
+        k = self.k_norm(k)
+
+        q = rope(q)
+        k = rope(k)
+
+        x = scaled_dot_product_attention(q, k, v, dropout_p=self.attn_drop.p if self.training else 0.)
+
+        x = x.transpose(1, 2).reshape(B, N, C)
+
+        x = self.proj(x)
+        x = self.proj_drop(x)
+        return x
+
+
+class SwiGLUFFN(nn.Module):
+    def __init__(
+        self,
+        dim: int,
+        hidden_dim: int,
+        drop=0.0,
+        bias=True
+    ) -> None:
+        super().__init__()
+        hidden_dim = int(hidden_dim * 2 / 3)
+        self.w12 = nn.Linear(dim, 2 * hidden_dim, bias=bias)
+        self.w3 = nn.Linear(hidden_dim, dim, bias=bias)
+        self.ffn_dropout = nn.Dropout(drop)
+
+    def forward(self, x):
+        x12 = self.w12(x)
+        x1, x2 = x12.chunk(2, dim=-1)
+        hidden = F.silu(x1) * x2
+        return self.w3(self.ffn_dropout(hidden))
+
+
+class FinalLayer(nn.Module):
+    def __init__(self, hidden_size, patch_size, out_channels):
+        super().__init__()
+        self.norm_final = RMSNorm(hidden_size)
+        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)
+        )
+
+    @torch.compile
+    def forward(self, x, c):
+        shift, scale = self.adaLN_modulation(c).chunk(2, dim=1)
+        x = modulate(self.norm_final(x), shift, scale)
+        x = self.linear(x)
+        return x
+
+
+class GRNblock(nn.Module):
+    def __init__(self, hidden_size, num_heads, mlp_ratio=4.0, attn_drop=0.0, proj_drop=0.0):
+        super().__init__()
+        self.norm1 = RMSNorm(hidden_size, eps=1e-6)
+        self.attn = Attention(hidden_size, num_heads=num_heads, qkv_bias=True, qk_norm=True,
+                              attn_drop=attn_drop, proj_drop=proj_drop)
+        self.norm2 = RMSNorm(hidden_size, eps=1e-6)
+        mlp_hidden_dim = int(hidden_size * mlp_ratio)
+        self.mlp = SwiGLUFFN(hidden_size, mlp_hidden_dim, drop=proj_drop)
+        self.adaLN_modulation = nn.Sequential(
+            nn.SiLU(),
+            nn.Linear(hidden_size, 6 * hidden_size, bias=True)
+        )
+
+    @torch.compile
+    def forward(self, x,  c, feat_rope=None):
+        shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = self.adaLN_modulation(c).chunk(6, dim=-1)
+        x = x + gate_msa.unsqueeze(1) * self.attn(modulate(self.norm1(x), shift_msa, scale_msa), rope=feat_rope)
+        x = x + gate_mlp.unsqueeze(1) * self.mlp(modulate(self.norm2(x), shift_mlp, scale_mlp))
+        return x
+
+
+class GRN(nn.Module):
+    """
+    GRN image Transformer.
+    """
+    def __init__(
+        self,
+        input_size=256,
+        patch_size=16,
+        in_channels=3,
+        hidden_size=1024,
+        depth=24,
+        num_heads=16,
+        mlp_ratio=4.0,
+        attn_drop=0.0,
+        proj_drop=0.0,
+        num_classes=1000,
+        bottleneck_dim=128,
+        in_context_len=32,
+        in_context_start=8,
+        args=None,
+    ):
+        super().__init__()
+        self.in_channels = in_channels
+        self.out_channels = in_channels
+        self.patch_size = patch_size
+        self.num_heads = num_heads
+        self.hidden_size = hidden_size
+        self.input_size = input_size
+        self.in_context_len = in_context_len
+        self.in_context_start = in_context_start
+        self.num_classes = num_classes
+        self.args=args
+
+        # time and class embed
+        self.t_embedder = TimestepEmbedder(hidden_size)
+        self.y_embedder = LabelEmbedder(num_classes, hidden_size)
+
+        # linear embed
+        self.x_embedder = BottleneckPatchEmbed(input_size, patch_size, in_channels, bottleneck_dim, hidden_size, bias=True)
+
+        # use fixed sin-cos embedding
+        num_patches = self.x_embedder.num_patches
+        self.pos_embed = nn.Parameter(torch.zeros(1, num_patches, hidden_size), requires_grad=False)
+
+        # in-context cls token
+        if self.in_context_len > 0:
+            self.in_context_posemb = nn.Parameter(torch.zeros(1, self.in_context_len, hidden_size), requires_grad=True)
+            torch.nn.init.normal_(self.in_context_posemb, std=.02)
+
+        # rope
+        half_head_dim = hidden_size // num_heads // 2
+        hw_seq_len = input_size // patch_size
+        self.feat_rope = VisionRotaryEmbeddingFast(
+            dim=half_head_dim,
+            pt_seq_len=hw_seq_len,
+            num_cls_token=0
+        )
+        self.feat_rope_incontext = VisionRotaryEmbeddingFast(
+            dim=half_head_dim,
+            pt_seq_len=hw_seq_len,
+            num_cls_token=self.in_context_len
+        )
+
+        # transformer
+        self.blocks = nn.ModuleList([
+            GRNblock(hidden_size, num_heads, mlp_ratio=mlp_ratio,
+                     attn_drop=attn_drop if (depth // 4 * 3 > i >= depth // 4) else 0.0,
+                     proj_drop=proj_drop if (depth // 4 * 3 > i >= depth // 4) else 0.0)
+            for i in range(depth)
+        ])
+
+        # linear predict
+        self.final_layer = FinalLayer(hidden_size, patch_size, self.out_channels)
+
+        self.initialize_weights()
+
+    def initialize_weights(self):
+        # Initialize transformer layers:
+        def _basic_init(module):
+            if isinstance(module, nn.Linear):
+                torch.nn.init.xavier_uniform_(module.weight)
+                if module.bias is not None:
+                    nn.init.constant_(module.bias, 0)
+        self.apply(_basic_init)
+
+        # Initialize (and freeze) pos_embed by sin-cos embedding:
+        pos_embed = get_2d_sincos_pos_embed(self.pos_embed.shape[-1], int(self.x_embedder.num_patches ** 0.5))
+        self.pos_embed.data.copy_(torch.from_numpy(pos_embed).float().unsqueeze(0))
+
+        # Initialize patch_embed like nn.Linear (instead of nn.Conv2d):
+        w1 = self.x_embedder.proj1.weight.data
+        nn.init.xavier_uniform_(w1.view([w1.shape[0], -1]))
+        w2 = self.x_embedder.proj2.weight.data
+        nn.init.xavier_uniform_(w2.view([w2.shape[0], -1]))
+        nn.init.constant_(self.x_embedder.proj2.bias, 0)
+
+        # Initialize label embedding table:
+        nn.init.normal_(self.y_embedder.embedding_table.weight, std=0.02)
+
+        nn.init.normal_(self.t_embedder.mlp[0].weight, std=0.02)
+        nn.init.normal_(self.t_embedder.mlp[2].weight, std=0.02)
+
+        # Zero-out adaLN modulation layers:
+        for block in self.blocks:
+            nn.init.constant_(block.adaLN_modulation[-1].weight, 0)
+            nn.init.constant_(block.adaLN_modulation[-1].bias, 0)
+
+        # Zero-out output layers:
+        nn.init.constant_(self.final_layer.adaLN_modulation[-1].weight, 0)
+        nn.init.constant_(self.final_layer.adaLN_modulation[-1].bias, 0)
+
+        nn.init.constant_(self.final_layer.linear.weight, 0)
+        nn.init.constant_(self.final_layer.linear.bias, 0)
+
+    def unpatchify(self, x, p):
+        """
+        x: (N, T, patch_size**2 * C)
+        imgs: (N, H, W, C)
+        """
+        c = self.out_channels
+        h = w = int(x.shape[1] ** 0.5)
+        assert h * w == x.shape[1]
+
+        x = x.reshape(shape=(x.shape[0], h, w, p, p, c))
+        x = torch.einsum('nhwpqc->nchpwq', x)
+        imgs = x.reshape(shape=(x.shape[0], c, h * p, h * p))
+        return imgs
+
+    def forward(self, x, t, y):
+        """
+        x: (N, C, H, W)
+        t: (N,)
+        y: (N,)
+        """
+        if self.args.method in ['GRN_ind']:
+            x = multiclass_labels2onehot_input(x, 2**self.args.hbq_round)
+        elif self.args.method in ['GRN_bit']:
+            x = multiclass_labels2onehot_input(x, 2)
+        
+        # class and time embeddings
+        t_emb = self.t_embedder(t)
+        y_emb = self.y_embedder(y)
+        c = t_emb + y_emb
+
+        # forward
+        x = self.x_embedder(x)
+        x += self.pos_embed
+
+        for i, block in enumerate(self.blocks):
+            # in-context
+            if self.in_context_len > 0 and i == self.in_context_start:
+                in_context_tokens = y_emb.unsqueeze(1).repeat(1, self.in_context_len, 1)
+                in_context_tokens += self.in_context_posemb
+                x = torch.cat([in_context_tokens, x], dim=1)
+            x = block(x, c, self.feat_rope if i < self.in_context_start else self.feat_rope_incontext)
+
+        x = x[:, self.in_context_len:]
+
+        with torch.amp.autocast('cuda', dtype=torch.float32):
+            x = self.final_layer(x, c)
+        
+        if self.args.method == 'GRN_ind':
+            B, h_mul_w, classes_mul_d = x.shape
+            classes = 2**self.args.hbq_round
+            h = w = int(np.round(math.sqrt(h_mul_w)))
+            output = x.reshape(B, h, w, classes, classes_mul_d//classes) # [B, h, w, classes, d]
+            output = output.permute(0, 3, 4, 1, 2) # [B, classes, d, h, w]
+        elif self.args.method == 'GRN_bit':
+            B, h_mul_w, classes_mul_d = x.shape
+            h = w = int(np.round(math.sqrt(h_mul_w)))
+            output = x.reshape(B, h, w, 2, classes_mul_d//2) # [B, h, w, 2, d]
+            output = output.permute(0, 3, 4, 1, 2) # [B, 2, d, h, w]
+        return output
+
+
+def GRN_B(**kwargs):
+    return GRN(depth=12, hidden_size=768, num_heads=12,
+               bottleneck_dim=128, in_context_len=32, in_context_start=4, patch_size=1, **kwargs)
+
+def GRN_L(**kwargs):
+    return GRN(depth=24, hidden_size=1024, num_heads=16,
+               bottleneck_dim=128, in_context_len=32, in_context_start=8, patch_size=1, **kwargs)
+
+def GRN_H(**kwargs):
+    return GRN(depth=32, hidden_size=1280, num_heads=16,
+               bottleneck_dim=256, in_context_len=32, in_context_start=10, patch_size=1, **kwargs)
+
+def GRN_G(**kwargs):
+    return GRN(depth=40, hidden_size=1664, num_heads=16,
+               bottleneck_dim=256, in_context_len=32, in_context_start=10, patch_size=1, **kwargs)
+
+GRN_models = {
+    'GRN_B': GRN_B,
+    'GRN_L': GRN_L,
+    'GRN_H': GRN_H,
+    'GRN_G': GRN_G,
+}

grn/models/hbq_tokenizer.py

@@ -0,0 +1,932 @@
+import logging
+import os
+import os.path as osp
+
+import torch
+import torch.cuda.amp as amp
+import torch.nn as nn
+import torch.nn.functional as F
+import numpy as np
+from einops import rearrange
+
+
+CACHE_T = 2
+
+
+class CausalConv3d(nn.Conv3d):
+    """
+    Causal 3d convolusion.
+    """
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self._padding = (
+            self.padding[2],
+            self.padding[2],
+            self.padding[1],
+            self.padding[1],
+            2 * self.padding[0],
+            0,
+        )
+        self.padding = (0, 0, 0)
+
+    def forward(self, x, cache_x=None):
+        padding = list(self._padding)
+        if cache_x is not None and self._padding[4] > 0:
+            cache_x = cache_x.to(x.device)
+            x = torch.cat([cache_x, x], dim=2)
+            padding[4] -= cache_x.shape[2]
+        x = F.pad(x, padding)
+
+        return super().forward(x)
+
+
+class RMS_norm(nn.Module):
+
+    def __init__(self, dim, channel_first=True, images=True, bias=False):
+        super().__init__()
+        broadcastable_dims = (1, 1, 1) if not images else (1, 1)
+        shape = (dim, *broadcastable_dims) if channel_first else (dim,)
+
+        self.channel_first = channel_first
+        self.scale = dim**0.5
+        self.gamma = nn.Parameter(torch.ones(shape))
+        self.bias = nn.Parameter(torch.zeros(shape)) if bias else 0.0
+
+    def forward(self, x):
+        return (F.normalize(x, dim=(1 if self.channel_first else -1)) *
+                self.scale * self.gamma + self.bias)
+
+
+class Upsample(nn.Upsample):
+
+    def forward(self, x):
+        """
+        Fix bfloat16 support for nearest neighbor interpolation.
+        """
+        return super().forward(x.float()).type_as(x)
+
+
+class Resample(nn.Module):
+
+    def __init__(self, dim, mode):
+        assert mode in (
+            "none",
+            "upsample2d",
+            "upsample3d",
+            "downsample2d",
+            "downsample3d",
+        )
+        super().__init__()
+        self.dim = dim
+        self.mode = mode
+
+        # layers
+        if mode == "upsample2d":
+            self.resample = nn.Sequential(
+                Upsample(scale_factor=(2.0, 2.0), mode="nearest-exact"),
+                nn.Conv2d(dim, dim, 3, padding=1),
+            )
+        elif mode == "upsample3d":
+            self.resample = nn.Sequential(
+                Upsample(scale_factor=(2.0, 2.0), mode="nearest-exact"),
+                nn.Conv2d(dim, dim, 3, padding=1),
+                # nn.Conv2d(dim, dim//2, 3, padding=1)
+            )
+            self.time_conv = CausalConv3d(
+                dim, dim * 2, (3, 1, 1), padding=(1, 0, 0))
+        elif mode == "downsample2d":
+            self.resample = nn.Sequential(
+                nn.ZeroPad2d((0, 1, 0, 1)),
+                nn.Conv2d(dim, dim, 3, stride=(2, 2)))
+        elif mode == "downsample3d":
+            self.resample = nn.Sequential(
+                nn.ZeroPad2d((0, 1, 0, 1)),
+                nn.Conv2d(dim, dim, 3, stride=(2, 2)))
+            self.time_conv = CausalConv3d(
+                dim, dim, (3, 1, 1), stride=(2, 1, 1), padding=(0, 0, 0))
+        else:
+            self.resample = nn.Identity()
+
+    def forward(self, x, feat_cache=None, feat_idx=[0]):
+        b, c, t, h, w = x.size()
+        if self.mode == "upsample3d":
+            if feat_cache is not None:
+                idx = feat_idx[0]
+                if feat_cache[idx] is None:
+                    feat_cache[idx] = "Rep"
+                    feat_idx[0] += 1
+                else:
+                    cache_x = x[:, :, -CACHE_T:, :, :].clone()
+                    if (cache_x.shape[2] < 2 and feat_cache[idx] is not None and
+                            feat_cache[idx] != "Rep"):
+                        # cache last frame of last two chunk
+                        cache_x = torch.cat(
+                            [
+                                feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to(
+                                    cache_x.device),
+                                cache_x,
+                            ],
+                            dim=2,
+                        )
+                    if (cache_x.shape[2] < 2 and feat_cache[idx] is not None and
+                            feat_cache[idx] == "Rep"):
+                        cache_x = torch.cat(
+                            [
+                                torch.zeros_like(cache_x).to(cache_x.device),
+                                cache_x
+                            ],
+                            dim=2,
+                        )
+                    if feat_cache[idx] == "Rep":
+                        x = self.time_conv(x)
+                    else:
+                        x = self.time_conv(x, feat_cache[idx])
+                    feat_cache[idx] = cache_x
+                    feat_idx[0] += 1
+                    x = x.reshape(b, 2, c, t, h, w)
+                    x = torch.stack((x[:, 0, :, :, :, :], x[:, 1, :, :, :, :]),
+                                    3)
+                    x = x.reshape(b, c, t * 2, h, w)
+        t = x.shape[2]
+        x = rearrange(x, "b c t h w -> (b t) c h w")
+        x = self.resample(x)
+        x = rearrange(x, "(b t) c h w -> b c t h w", t=t) # this 4 lines do spatial down / up sample
+
+        if self.mode == "downsample3d":
+            if feat_cache is not None:
+                idx = feat_idx[0]
+                if feat_cache[idx] is None:
+                    feat_cache[idx] = x.clone()
+                    feat_idx[0] += 1
+                else:
+                    cache_x = x[:, :, -1:, :, :].clone()
+                    x = self.time_conv(
+                        torch.cat([feat_cache[idx][:, :, -1:, :, :], x], 2))
+                    feat_cache[idx] = cache_x
+                    feat_idx[0] += 1
+        return x
+
+    def init_weight(self, conv):
+        conv_weight = conv.weight.detach().clone()
+        nn.init.zeros_(conv_weight)
+        c1, c2, t, h, w = conv_weight.size()
+        one_matrix = torch.eye(c1, c2)
+        init_matrix = one_matrix
+        nn.init.zeros_(conv_weight)
+        conv_weight.data[:, :, 1, 0, 0] = init_matrix  # * 0.5
+        conv.weight = nn.Parameter(conv_weight)
+        nn.init.zeros_(conv.bias.data)
+
+    def init_weight2(self, conv):
+        conv_weight = conv.weight.data.detach().clone()
+        nn.init.zeros_(conv_weight)
+        c1, c2, t, h, w = conv_weight.size()
+        init_matrix = torch.eye(c1 // 2, c2)
+        conv_weight[:c1 // 2, :, -1, 0, 0] = init_matrix
+        conv_weight[c1 // 2:, :, -1, 0, 0] = init_matrix
+        conv.weight = nn.Parameter(conv_weight)
+        nn.init.zeros_(conv.bias.data)
+
+
+class ResidualBlock(nn.Module):
+
+    def __init__(self, in_dim, out_dim, dropout=0.0):
+        super().__init__()
+        self.in_dim = in_dim
+        self.out_dim = out_dim
+
+        # layers
+        self.residual = nn.Sequential(
+            RMS_norm(in_dim, images=False),
+            nn.SiLU(),
+            CausalConv3d(in_dim, out_dim, 3, padding=1),
+            RMS_norm(out_dim, images=False),
+            nn.SiLU(),
+            nn.Dropout(dropout),
+            CausalConv3d(out_dim, out_dim, 3, padding=1),
+        )
+        self.shortcut = (
+            CausalConv3d(in_dim, out_dim, 1)
+            if in_dim != out_dim else nn.Identity())
+
+    def forward(self, x, feat_cache=None, feat_idx=[0]):
+        h = self.shortcut(x)
+        for layer in self.residual:
+            if isinstance(layer, CausalConv3d) and feat_cache is not None:
+                idx = feat_idx[0]
+                cache_x = x[:, :, -CACHE_T:, :, :].clone()
+                if cache_x.shape[2] < 2 and feat_cache[idx] is not None:
+                    # cache last frame of last two chunk
+                    cache_x = torch.cat(
+                        [
+                            feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to(
+                                cache_x.device),
+                            cache_x,
+                        ],
+                        dim=2,
+                    )
+                x = layer(x, feat_cache[idx])
+                feat_cache[idx] = cache_x
+                feat_idx[0] += 1
+            else:
+                x = layer(x)
+        return x + h
+
+
+class AttentionBlock(nn.Module):
+    """
+    Causal self-attention with a single head.
+    """
+
+    def __init__(self, dim):
+        super().__init__()
+        self.dim = dim
+
+        # layers
+        self.norm = RMS_norm(dim)
+        self.to_qkv = nn.Conv2d(dim, dim * 3, 1)
+        self.proj = nn.Conv2d(dim, dim, 1)
+
+        # zero out the last layer params
+        nn.init.zeros_(self.proj.weight)
+
+    def forward(self, x):
+        identity = x
+        b, c, t, h, w = x.size()
+        x = rearrange(x, "b c t h w -> (b t) c h w")
+        x = self.norm(x)
+        # compute query, key, value
+        q, k, v = (
+            self.to_qkv(x).reshape(b * t, 1, c * 3,
+                                   -1).permute(0, 1, 3,
+                                               2).contiguous().chunk(3, dim=-1))
+
+        # apply attention
+        x = F.scaled_dot_product_attention(
+            q,
+            k,
+            v,
+        )
+        x = x.squeeze(1).permute(0, 2, 1).reshape(b * t, c, h, w)
+
+        # output
+        x = self.proj(x)
+        x = rearrange(x, "(b t) c h w-> b c t h w", t=t)
+        return x + identity
+
+
+def patchify(x, patch_size):
+    if patch_size == 1:
+        return x
+    if x.dim() == 4:
+        x = rearrange(
+            x, "b c (h q) (w r) -> b (c r q) h w", q=patch_size, r=patch_size)
+    elif x.dim() == 5:
+        x = rearrange(
+            x,
+            "b c f (h q) (w r) -> b (c r q) f h w",
+            q=patch_size,
+            r=patch_size,
+        )
+    else:
+        raise ValueError(f"Invalid input shape: {x.shape}")
+
+    return x
+
+
+def unpatchify(x, patch_size):
+    if patch_size == 1:
+        return x
+
+    if x.dim() == 4:
+        x = rearrange(
+            x, "b (c r q) h w -> b c (h q) (w r)", q=patch_size, r=patch_size)
+    elif x.dim() == 5:
+        x = rearrange(
+            x,
+            "b (c r q) f h w -> b c f (h q) (w r)",
+            q=patch_size,
+            r=patch_size,
+        )
+    return x
+
+
+class AvgDown3D(nn.Module):
+
+    def __init__(
+        self,
+        in_channels,
+        out_channels,
+        factor_t,
+        factor_s=1,
+    ):
+        super().__init__()
+        self.in_channels = in_channels
+        self.out_channels = out_channels
+        self.factor_t = factor_t
+        self.factor_s = factor_s
+        self.factor = self.factor_t * self.factor_s * self.factor_s
+
+        assert in_channels * self.factor % out_channels == 0
+        self.group_size = in_channels * self.factor // out_channels
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        pad_t = (self.factor_t - x.shape[2] % self.factor_t) % self.factor_t
+        pad = (0, 0, 0, 0, pad_t, 0)
+        x = F.pad(x, pad)
+        B, C, T, H, W = x.shape
+        x = x.view(
+            B,
+            C,
+            T // self.factor_t,
+            self.factor_t,
+            H // self.factor_s,
+            self.factor_s,
+            W // self.factor_s,
+            self.factor_s,
+        )
+        x = x.permute(0, 1, 3, 5, 7, 2, 4, 6).contiguous()
+        x = x.view(
+            B,
+            C * self.factor,
+            T // self.factor_t,
+            H // self.factor_s,
+            W // self.factor_s,
+        )
+        x = x.view(
+            B,
+            self.out_channels,
+            self.group_size,
+            T // self.factor_t,
+            H // self.factor_s,
+            W // self.factor_s,
+        )
+        x = x.mean(dim=2)
+        return x
+
+
+class DupUp3D(nn.Module):
+
+    def __init__(
+        self,
+        in_channels: int,
+        out_channels: int,
+        factor_t,
+        factor_s=1,
+    ):
+        super().__init__()
+        self.in_channels = in_channels
+        self.out_channels = out_channels
+
+        self.factor_t = factor_t
+        self.factor_s = factor_s
+        self.factor = self.factor_t * self.factor_s * self.factor_s
+
+        assert out_channels * self.factor % in_channels == 0
+        self.repeats = out_channels * self.factor // in_channels
+
+    def forward(self, x: torch.Tensor, first_chunk=False) -> torch.Tensor:
+        x = x.repeat_interleave(self.repeats, dim=1)
+        x = x.view(
+            x.size(0),
+            self.out_channels,
+            self.factor_t,
+            self.factor_s,
+            self.factor_s,
+            x.size(2),
+            x.size(3),
+            x.size(4),
+        )
+        x = x.permute(0, 1, 5, 2, 6, 3, 7, 4).contiguous()
+        x = x.view(
+            x.size(0),
+            self.out_channels,
+            x.size(2) * self.factor_t,
+            x.size(4) * self.factor_s,
+            x.size(6) * self.factor_s,
+        )
+        if first_chunk:
+            x = x[:, :, self.factor_t - 1:, :, :]
+        return x
+
+
+class Down_ResidualBlock(nn.Module):
+
+    def __init__(self,
+                 in_dim,
+                 out_dim,
+                 dropout,
+                 mult,
+                 temperal_downsample=False,
+                 down_flag=False):
+        super().__init__()
+
+        # Shortcut path with downsample
+        self.avg_shortcut = AvgDown3D(
+            in_dim,
+            out_dim,
+            factor_t=2 if temperal_downsample else 1,
+            factor_s=2 if down_flag else 1,
+        )
+
+        # Main path with residual blocks and downsample
+        downsamples = []
+        for _ in range(mult): # mult=2, two block
+            downsamples.append(ResidualBlock(in_dim, out_dim, dropout))
+            in_dim = out_dim
+
+        # Add the final downsample block
+        if down_flag:
+            mode = "downsample3d" if temperal_downsample else "downsample2d"
+            downsamples.append(Resample(out_dim, mode=mode))
+
+        self.downsamples = nn.Sequential(*downsamples)
+
+    def forward(self, x, feat_cache=None, feat_idx=[0]):
+        x_copy = x.clone()
+        for module in self.downsamples:
+            x = module(x, feat_cache, feat_idx)
+
+        return x + self.avg_shortcut(x_copy)
+
+
+class Up_ResidualBlock(nn.Module):
+
+    def __init__(self,
+                 in_dim,
+                 out_dim,
+                 dropout,
+                 mult,
+                 temperal_upsample=False,
+                 up_flag=False):
+        super().__init__()
+        # Shortcut path with upsample
+        if up_flag:
+            self.avg_shortcut = DupUp3D(
+                in_dim,
+                out_dim,
+                factor_t=2 if temperal_upsample else 1,
+                factor_s=2 if up_flag else 1,
+            )
+        else:
+            self.avg_shortcut = None
+
+        # Main path with residual blocks and upsample
+        upsamples = []
+        for _ in range(mult):
+            upsamples.append(ResidualBlock(in_dim, out_dim, dropout))
+            in_dim = out_dim
+
+        # Add the final upsample block
+        if up_flag:
+            mode = "upsample3d" if temperal_upsample else "upsample2d"
+            upsamples.append(Resample(out_dim, mode=mode))
+
+        self.upsamples = nn.Sequential(*upsamples)
+
+    def forward(self, x, feat_cache=None, feat_idx=[0], first_chunk=False):
+        x_main = x.clone()
+        for module in self.upsamples:
+            x_main = module(x_main, feat_cache, feat_idx)
+        if self.avg_shortcut is not None:
+            x_shortcut = self.avg_shortcut(x, first_chunk)
+            return x_main + x_shortcut
+        else:
+            return x_main
+
+
+class Encoder3d(nn.Module):
+
+    def __init__(
+        self,
+        dim=128,
+        z_dim=4,
+        dim_mult=[1, 2, 4, 4],
+        num_res_blocks=2,
+        attn_scales=[],
+        temperal_downsample=[True, True, False],
+        dropout=0.0,
+    ):
+        super().__init__()
+        self.dim = dim
+        self.z_dim = z_dim
+        self.dim_mult = dim_mult
+        self.num_res_blocks = num_res_blocks
+        self.attn_scales = attn_scales
+        self.temperal_downsample = temperal_downsample
+
+        # dimensions
+        dims = [dim * u for u in [1] + dim_mult] # [1,2,4,4] -> [1,1,2,4,4] -> [128,128,256,512,512]
+
+        # init block
+        self.conv1 = CausalConv3d(12, dims[0], 3, padding=1)
+
+        # downsample blocks
+        downsamples = []
+        for i, (in_dim, out_dim) in enumerate(zip(dims[:-1], dims[1:])):
+            t_down_flag = (
+                temperal_downsample[i]
+                if i < len(temperal_downsample) else False)
+            downsamples.append(
+                Down_ResidualBlock(
+                    in_dim=in_dim,
+                    out_dim=out_dim,
+                    dropout=dropout,
+                    mult=num_res_blocks,
+                    temperal_downsample=t_down_flag,
+                    down_flag=i != len(dim_mult) - 1,
+                ))
+        self.downsamples = nn.Sequential(*downsamples)
+
+        # middle blocks
+        self.middle = nn.Sequential(
+            ResidualBlock(out_dim, out_dim, dropout),
+            AttentionBlock(out_dim),
+            ResidualBlock(out_dim, out_dim, dropout),
+        )
+
+        # # output blocks
+        self.head = nn.Sequential(
+            RMS_norm(out_dim, images=False),
+            nn.SiLU(),
+            CausalConv3d(out_dim, z_dim, 3, padding=1),
+        )
+
+    def forward(self, x, feat_cache=None, feat_idx=[0]):
+
+        if feat_cache is not None:
+            idx = feat_idx[0]
+            cache_x = x[:, :, -CACHE_T:, :, :].clone()
+            if cache_x.shape[2] < 2 and feat_cache[idx] is not None:
+                cache_x = torch.cat(
+                    [
+                        feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to(
+                            cache_x.device),
+                        cache_x,
+                    ],
+                    dim=2,
+                )
+            x = self.conv1(x, feat_cache[idx])
+            feat_cache[idx] = cache_x
+            feat_idx[0] += 1
+        else:
+            x = self.conv1(x)
+
+        ## downsamples
+        for layer in self.downsamples:
+            if feat_cache is not None:
+                x = layer(x, feat_cache, feat_idx)
+            else:
+                x = layer(x)
+
+        ## middle
+        for layer in self.middle:
+            if isinstance(layer, ResidualBlock) and feat_cache is not None:
+                x = layer(x, feat_cache, feat_idx)
+            else:
+                x = layer(x)
+
+        ## head
+        for layer in self.head:
+            if isinstance(layer, CausalConv3d) and feat_cache is not None:
+                idx = feat_idx[0]
+                cache_x = x[:, :, -CACHE_T:, :, :].clone()
+                if cache_x.shape[2] < 2 and feat_cache[idx] is not None:
+                    cache_x = torch.cat(
+                        [
+                            feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to(
+                                cache_x.device),
+                            cache_x,
+                        ],
+                        dim=2,
+                    )
+                x = layer(x, feat_cache[idx])
+                feat_cache[idx] = cache_x
+                feat_idx[0] += 1
+            else:
+                x = layer(x)
+
+        return x
+
+
+class Decoder3d(nn.Module):
+
+    def __init__(
+        self,
+        dim=128,
+        z_dim=4,
+        dim_mult=[1, 2, 4, 4],
+        num_res_blocks=2,
+        attn_scales=[],
+        temperal_upsample=[False, True, True],
+        dropout=0.0,
+    ):
+        super().__init__()
+        self.dim = dim
+        self.z_dim = z_dim
+        self.dim_mult = dim_mult
+        self.num_res_blocks = num_res_blocks
+        self.attn_scales = attn_scales
+        self.temperal_upsample = temperal_upsample
+
+        # dimensions
+        dims = [dim * u for u in [dim_mult[-1]] + dim_mult[::-1]]
+        # scale = 1.0 / 2**(len(dim_mult) - 2)
+        # init block
+        self.conv1 = CausalConv3d(z_dim, dims[0], 3, padding=1)
+
+        # middle blocks
+        self.middle = nn.Sequential(
+            ResidualBlock(dims[0], dims[0], dropout),
+            AttentionBlock(dims[0]),
+            ResidualBlock(dims[0], dims[0], dropout),
+        )
+
+        # upsample blocks
+        upsamples = []
+        for i, (in_dim, out_dim) in enumerate(zip(dims[:-1], dims[1:])):
+            t_up_flag = temperal_upsample[i] if i < len(
+                temperal_upsample) else False
+            upsamples.append(
+                Up_ResidualBlock(
+                    in_dim=in_dim,
+                    out_dim=out_dim,
+                    dropout=dropout,
+                    mult=num_res_blocks + 1,
+                    temperal_upsample=t_up_flag,
+                    up_flag=i != len(dim_mult) - 1,
+                ))
+        self.upsamples = nn.Sequential(*upsamples)
+
+        # output blocks
+        self.head = nn.Sequential(
+            RMS_norm(out_dim, images=False),
+            nn.SiLU(),
+            CausalConv3d(out_dim, 12, 3, padding=1),
+        )
+
+    def forward(self, x, feat_cache=None, feat_idx=[0], first_chunk=False):
+        if feat_cache is not None:
+            idx = feat_idx[0]
+            cache_x = x[:, :, -CACHE_T:, :, :].clone()
+            if cache_x.shape[2] < 2 and feat_cache[idx] is not None:
+                cache_x = torch.cat(
+                    [
+                        feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to(
+                            cache_x.device),
+                        cache_x,
+                    ],
+                    dim=2,
+                )
+            x = self.conv1(x, feat_cache[idx])
+            feat_cache[idx] = cache_x
+            feat_idx[0] += 1
+        else:
+            x = self.conv1(x)
+
+        for layer in self.middle:
+            if isinstance(layer, ResidualBlock) and feat_cache is not None:
+                x = layer(x, feat_cache, feat_idx)
+            else:
+                x = layer(x)
+
+        ## upsamples
+        for layer in self.upsamples:
+            if feat_cache is not None:
+                x = layer(x, feat_cache, feat_idx, first_chunk)
+            else:
+                x = layer(x)
+
+        ## head
+        for layer in self.head:
+            if isinstance(layer, CausalConv3d) and feat_cache is not None:
+                idx = feat_idx[0]
+                cache_x = x[:, :, -CACHE_T:, :, :].clone()
+                if cache_x.shape[2] < 2 and feat_cache[idx] is not None:
+                    cache_x = torch.cat(
+                        [
+                            feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to(
+                                cache_x.device),
+                            cache_x,
+                        ],
+                        dim=2,
+                    )
+                x = layer(x, feat_cache[idx])
+                feat_cache[idx] = cache_x
+                feat_idx[0] += 1
+            else:
+                x = layer(x)
+        return x
+
+
+def count_conv3d(model):
+    count = 0
+    for m in model.modules():
+        if isinstance(m, CausalConv3d):
+            count += 1
+    return count
+
+
+class WanVAE_(nn.Module):
+
+    def __init__(
+        self,
+        dim=160,
+        dec_dim=256,
+        z_dim=16,
+        dim_mult=[1, 2, 4, 4],
+        num_res_blocks=2,
+        attn_scales=[],
+        temperal_downsample=[True, True, False],
+        dropout=0.0,
+    ):
+        super().__init__()
+        self.dim = dim
+        self.z_dim = z_dim
+        self.dim_mult = dim_mult
+        self.num_res_blocks = num_res_blocks
+        self.attn_scales = attn_scales
+        self.temperal_downsample = temperal_downsample
+        self.temperal_upsample = temperal_downsample[::-1]
+
+        # modules
+        self.encoder = Encoder3d(
+            dim,
+            z_dim * 2,
+            dim_mult,
+            num_res_blocks,
+            attn_scales,
+            self.temperal_downsample,
+            dropout,
+        )
+        self.conv1 = CausalConv3d(z_dim * 2, z_dim * 2, 1)
+        self.conv2 = CausalConv3d(z_dim, z_dim, 1)
+        self.decoder = Decoder3d(
+            dec_dim,
+            z_dim,
+            dim_mult,
+            num_res_blocks,
+            attn_scales,
+            self.temperal_upsample,
+            dropout,
+        )
+
+    def forward(self, x, scale=[0, 1]):
+        mu = self.encode(x, scale)
+        x_recon = self.decode(mu, scale)
+        return x_recon, mu
+
+    def encode(self, x, scale=None):
+        self.clear_cache()
+        x = patchify(x, patch_size=2)
+        t = x.shape[2]
+        iter_ = 1 + (t - 1) // 4
+        for i in range(iter_):
+            self._enc_conv_idx = [0]
+            if i == 0:
+                out = self.encoder(
+                    x[:, :, :1, :, :],
+                    feat_cache=self._enc_feat_map,
+                    feat_idx=self._enc_conv_idx,
+                )
+            else:
+                out_ = self.encoder(
+                    x[:, :, 1 + 4 * (i - 1):1 + 4 * i, :, :],
+                    feat_cache=self._enc_feat_map,
+                    feat_idx=self._enc_conv_idx,
+                )
+                out = torch.cat([out, out_], 2)
+        mu, log_var = self.conv1(out).chunk(2, dim=1)
+        if scale is not None:
+            if isinstance(scale[0], torch.Tensor):
+                mu = (mu - scale[0].view(1, self.z_dim, 1, 1, 1)) * scale[1].view(
+                    1, self.z_dim, 1, 1, 1)
+            else:
+                mu = (mu - scale[0]) * scale[1]
+        self.clear_cache()
+        if self.encoder_out_type == 'feature_tanh':
+            mu = torch.tanh(mu)
+        elif self.encoder_out_type == 'feature':
+            pass
+        else:
+            raise ValueError(f'{self.encoder_out_type=} is not supported!')
+        return mu
+
+    def decode(self, z, scale=None, **kwargs):
+        self.clear_cache()
+        if scale is not None:
+            if isinstance(scale[0], torch.Tensor):
+                z = z / scale[1].view(1, self.z_dim, 1, 1, 1) + scale[0].view(
+                    1, self.z_dim, 1, 1, 1)
+            else:
+                z = z / scale[1] + scale[0]
+        x = self.conv2(z)
+        iter_ = z.shape[2]
+        for i in range(iter_):
+            self._conv_idx = [0]
+            if i == 0:
+                out = self.decoder(
+                    x[:, :, i:i + 1, :, :],
+                    feat_cache=self._feat_map,
+                    feat_idx=self._conv_idx,
+                    first_chunk=True,
+                )
+            else:
+                out_ = self.decoder(
+                    x[:, :, i:i + 1, :, :],
+                    feat_cache=self._feat_map,
+                    feat_idx=self._conv_idx,
+                )
+                out = torch.cat([out, out_], 2)
+        out = unpatchify(out, patch_size=2)
+        self.clear_cache()
+        return out
+
+    def reparameterize(self, mu, log_var):
+        std = torch.exp(0.5 * log_var)
+        eps = torch.randn_like(std)
+        return eps * std + mu
+
+    def sample(self, imgs, deterministic=False):
+        import pdb; pdb.set_trace()
+        mu, log_var = self.encode(imgs)
+        if deterministic:
+            return mu
+        std = torch.exp(0.5 * log_var.clamp(-30.0, 20.0))
+        return mu + std * torch.randn_like(std)
+
+    def clear_cache(self):
+        self._conv_num = count_conv3d(self.decoder)
+        self._conv_idx = [0]
+        self._feat_map = [None] * self._conv_num
+        # cache encode
+        self._enc_conv_num = count_conv3d(self.encoder)
+        self._enc_conv_idx = [0]
+        self._enc_feat_map = [None] * self._enc_conv_num
+
+class HBQ_Tokenizer(WanVAE_):
+    def __init__(
+        self,
+        args,
+        dim=160,
+        dec_dim=256,
+        latent_channels=16,
+        dim_mult=[1, 2, 4, 4],
+        num_res_blocks=2,
+        temperal_downsample=[0,1,1],
+        dropout=0.,
+        encoder_out_type='',
+    ):
+        super().__init__(
+            dim=dim, 
+            dec_dim=dec_dim, 
+            z_dim=latent_channels, 
+            dim_mult=dim_mult, 
+            num_res_blocks=num_res_blocks, 
+            temperal_downsample=temperal_downsample, 
+            dropout=dropout,
+        )
+        self.other_args = args
+        self.codebook_dim = latent_channels
+        self.encoder_out_type = encoder_out_type
+
+    def encode_for_raw_features(
+        self, x: torch.Tensor, 
+        **kwargs,
+    ):
+        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
+            x = x.unsqueeze(2)
+        with torch.amp.autocast("cuda", dtype=torch.float):
+            z = self.encode(x)
+        return [z], None, None
+
+def _video_vae(pretrained_path=None, z_dim=16, dim=160, device="cpu", **kwargs):
+    # params
+    cfg = dict(
+        dim=dim,
+        z_dim=z_dim,
+        dim_mult=[1, 2, 4, 4],
+        num_res_blocks=2,
+        attn_scales=[],
+        temperal_downsample=[True, True, True],
+        dropout=0.0,
+    )
+    cfg.update(**kwargs)
+
+    # init model
+    with torch.device("meta"):
+        model = WanVAE_(**cfg)
+
+    # load checkpoint
+    logging.info(f"loading {pretrained_path}")
+    model.load_state_dict(
+        torch.load(pretrained_path, map_location=device), assign=True)
+
+    return model
+

grn/models/init_param.py

@@ -0,0 +1,33 @@
+import torch.nn as nn
+
+
+def init_weights(model: nn.Module, conv_std_or_gain: float = 0.02, other_std: float = 0.02):
+    """
+    :param model: the model to be inited
+    :param conv_std_or_gain: how to init every conv layer `m`
+        > 0: nn.init.trunc_normal_(m.weight.data, std=conv_std_or_gain)
+        < 0: nn.init.xavier_normal_(m.weight.data, gain=-conv_std_or_gain)
+    :param other_std: how to init every linear layer or embedding layer
+        use nn.init.trunc_normal_(m.weight.data, std=other_std)
+    """
+    skip = abs(conv_std_or_gain) > 10
+    if skip: return
+    print(f'[init_weights] {type(model).__name__} with {"std" if conv_std_or_gain > 0 else "gain"}={abs(conv_std_or_gain):g}')
+    for m in model.modules():
+        if isinstance(m, nn.Linear):
+            nn.init.trunc_normal_(m.weight.data, std=other_std)
+            if m.bias is not None:
+                nn.init.constant_(m.bias.data, 0.)
+        elif isinstance(m, nn.Embedding):
+            nn.init.trunc_normal_(m.weight.data, std=other_std)
+            if m.padding_idx is not None:
+                m.weight.data[m.padding_idx].zero_()
+        elif isinstance(m, (nn.Conv1d, nn.Conv2d, nn.ConvTranspose1d, nn.ConvTranspose2d)):
+            nn.init.trunc_normal_(m.weight.data, std=conv_std_or_gain) if conv_std_or_gain > 0 else nn.init.xavier_normal_(m.weight.data, gain=-conv_std_or_gain)   # todo: StyleSwin: (..., gain=.02)
+            if hasattr(m, 'bias') and m.bias is not None:
+                nn.init.constant_(m.bias.data, 0.)
+        elif isinstance(m, (nn.LayerNorm, nn.BatchNorm1d, nn.BatchNorm2d, nn.BatchNorm3d, nn.SyncBatchNorm, nn.GroupNorm, nn.InstanceNorm1d, nn.InstanceNorm2d, nn.InstanceNorm3d)):
+            if m.bias is not None:
+                nn.init.constant_(m.bias.data, 0.)
+            if m.weight is not None:
+                nn.init.constant_(m.weight.data, 1.)

grn/models/rope.py

@@ -0,0 +1,191 @@
+import math
+import os
+from functools import partial
+from typing import Optional, Tuple, Union
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import numpy as np
+from timm.models.layers import DropPath, drop_path
+from torch.utils.checkpoint import checkpoint
+
+
+def precompute_rope2d_freqs_grid(dim, dynamic_resolution_h_w, rope2d_normalized_by_hw, pad_to_multiplier=1, max_height=2048 // 16, max_width=2048 // 16, base=10000.0, device=None, scaling_factor=1.0, activated_h_div_w_templates=[]):
+    # split the dimension into half, one for x and one for y
+    half_dim = dim // 2
+    inv_freq = 1.0 / (base ** (torch.arange(0, half_dim, 2, dtype=torch.int64).float().to(device) / half_dim)) # namely theta, 1 / (10000^(i/half_dim)), i=0,2,..., half_dim-2
+    t_height = torch.arange(max_height, device=device, dtype=torch.int64).type_as(inv_freq)
+    t_width = torch.arange(max_width, device=device, dtype=torch.int64).type_as(inv_freq)
+    t_height = t_height / scaling_factor
+    freqs_height = torch.outer(t_height, inv_freq)  # (max_height, dim / (1 for 1d, 2 for 2d, 3 for 3d) / 2), namely y*theta
+    t_width = t_width / scaling_factor
+    freqs_width = torch.outer(t_width, inv_freq)  # (max_width, dim / (1 for 1d, 2 for 2d, 3 for 3d) / 2), namely x*theta
+    freqs_grid_map = torch.concat([
+        freqs_height[:, None, :].expand(-1, max_width, -1), # (max_height, max_width, dim / (1 for 1d, 2 for 2d, 3 for 3d) / 2)
+        freqs_width[None, :, :].expand(max_height, -1, -1), # (max_height, max_width, dim / (1 for 1d, 2 for 2d, 3 for 3d) / 2)
+    ], dim=-1)  # (max_height, max_width, dim / (1 for 1d, 2 for 2d, 3 for 3d))
+    freqs_grid_map = torch.stack([torch.cos(freqs_grid_map), torch.sin(freqs_grid_map)], dim=0)
+    # (2, max_height, max_width, dim / (1 for 1d, 2 for 2d, 3 for 3d))
+
+    rope2d_freqs_grid = {}
+    for h_div_w in activated_h_div_w_templates:
+        assert h_div_w in dynamic_resolution_h_w, f'Unknown h_div_w: {h_div_w}'
+        scale_schedule = dynamic_resolution_h_w[h_div_w]['1M']['image_scales']
+        _, ph, pw = scale_schedule[-1]
+        max_edge_length = freqs_grid_map.shape[1]
+        if ph >= pw:
+            uph, upw = max_edge_length, int(max_edge_length / ph * pw)
+        else:
+            uph, upw = int(max_edge_length / pw * ph), max_edge_length
+        rope_cache_list = []
+        for (_, ph, pw) in scale_schedule:
+            ph_mul_pw = ph * pw
+            if rope2d_normalized_by_hw == 1: # downsample
+                rope_cache = F.interpolate(freqs_grid_map[:, :uph, :upw, :].permute([0,3,1,2]), size=(ph, pw), mode='bilinear', align_corners=True)
+                rope_cache = rope_cache.permute([0,2,3,1]) # (2, ph, pw, half_head_dim)
+            elif rope2d_normalized_by_hw == 2: # star stylee
+                _, uph, upw = scale_schedule[-1]
+                indices = torch.stack([
+                    (torch.arange(ph) * (uph / ph)).reshape(ph, 1).expand(ph, pw),
+                    (torch.arange(pw) * (upw / pw)).reshape(1, pw).expand(ph, pw),
+                ], dim=-1).round().int() # (ph, pw, 2)
+                indices = indices.reshape(-1, 2) # (ph*pw, 2)
+                rope_cache = freqs_grid_map[:, indices[:,0], indices[:,1], :] # (2, ph*pw, half_head_dim)
+                rope_cache = rope_cache.reshape(2, ph, pw, -1)
+            elif rope2d_normalized_by_hw == 0:
+                rope_cache = freqs_grid_map[:, :ph, :pw, :] # (2, ph, pw, half_head_dim)
+            else:
+                raise ValueError(f'Unknown rope2d_normalized_by_hw: {rope2d_normalized_by_hw}')
+            rope_cache_list.append(rope_cache.reshape(2, ph_mul_pw, -1))
+        cat_rope_cache = torch.cat(rope_cache_list, 1) # (2, seq_len, half_head_dim)
+        if cat_rope_cache.shape[1] % pad_to_multiplier:
+            pad = torch.zeros(2, pad_to_multiplier - cat_rope_cache.shape[1] % pad_to_multiplier, half_dim)
+            cat_rope_cache = torch.cat([cat_rope_cache, pad], dim=1)
+        cat_rope_cache = cat_rope_cache[:,None,None,None] # (2, 1, 1, 1, seq_len, half_dim)
+        for pn in dynamic_resolution_h_w[h_div_w]:
+            scale_schedule = dynamic_resolution_h_w[h_div_w][pn]['image_scales']
+            tmp_scale_schedule = [(1, h, w) for _, h, w in scale_schedule]
+            rope2d_freqs_grid[str(tuple(tmp_scale_schedule))] = cat_rope_cache
+    return rope2d_freqs_grid
+
+
+def precompute_rope3d_freqs_grid(
+        dim, 
+        rope2d_normalized_by_hw, 
+        max_frames=128, 
+        max_height=2048 // 8, 
+        max_width=2048 // 8, 
+        base=10000.0, 
+        device=None, 
+        activated_h_div_w_templates=[], 
+        text_maxlen=0, 
+        pn=None, 
+        args=None,
+        **kwargs,
+):
+    # split the dimension into three parts, one for x, one for y, and one for t
+    print(f'[precompute_rope4d_freqs_grid: 3d]: start')
+    assert dim % 2 == 0, f'Only support dim % 2 == 0, but got dim={dim}'
+    dim_div_2 = dim // 2
+    num_of_freqs_former = dim_div_2 // 3
+    preserve_1d_length = 600
+    num_of_freqs_last = dim_div_2 - num_of_freqs_former * 2 # in some cases, dim_div_2 % 3 != 0. here tackle with these cases
+    inv_freq_former = 1.0 / (base ** (torch.arange(num_of_freqs_former, dtype=torch.int64).float().to(device) / num_of_freqs_former)) # namely theta, 1 / (10000^(i/dim_div_3)), i=0,2,..., dim_div_3-2, totally dim_div_3 / 2 elems
+    inv_freq_last = 1.0 / (base ** (torch.arange(num_of_freqs_last, dtype=torch.int64).float().to(device) / num_of_freqs_last))
+    t_frames = torch.arange(preserve_1d_length+max_frames, device=device, dtype=torch.int64).type_as(inv_freq_former)
+    t_height = torch.arange(max_height, device=device, dtype=torch.int64).type_as(inv_freq_former)
+    t_width = torch.arange(max_width, device=device, dtype=torch.int64).type_as(inv_freq_former)
+    freqs_frames = torch.outer(t_frames, inv_freq_former)  # (max_frames, (dim_div_2 / 3)), namely x*theta
+    freqs_height = torch.outer(t_height, inv_freq_former)  # (max_height, (dim_div_2 / 3), namely y*theta
+    freqs_width = torch.outer(t_width, inv_freq_last)  # (max_width, (dim_div_2 / 3)), namely x*theta
+    freqs_frames = torch.stack([torch.cos(freqs_frames), torch.sin(freqs_frames)], dim=0)
+    freqs_height = torch.stack([torch.cos(freqs_height), torch.sin(freqs_height)], dim=0)
+    freqs_width = torch.stack([torch.cos(freqs_width), torch.sin(freqs_width)], dim=0)
+    tm = preserve_1d_length
+    rope_text_embeds = torch.cat([
+        freqs_frames[   :,   :tm,  None,   None,   :].expand(-1, -1, -1, -1, -1),
+        freqs_height[   :,  None,    :1,   None,   :].expand(-1, tm, -1, -1, -1),
+        freqs_width[   :,  None,  None,     :1,   :].expand(-1, tm, -1, -1, -1),
+    ], dim=-1)  # (2, tm, 1, 1, dim_div_2)
+    rope_text_embeds = rope_text_embeds.reshape(2, 1, 1, 1, tm, dim_div_2)
+    rope2d_freqs_grid = {}
+    rope2d_freqs_grid['freqs_text'] = rope_text_embeds # (2, 1, 1, 1, preserve_1d_length, dim / 2)
+    rope2d_freqs_grid['freqs_frames'] = freqs_frames[:, tm:] # (2, max_frames, ceil(dim_div_2 / 4))
+    rope2d_freqs_grid['freqs_height'] = freqs_height # (2, max_height, ceil(dim_div_2 / 4))
+    rope2d_freqs_grid['freqs_width'] = freqs_width # (2, max_width, ceil(dim_div_2 / 4))
+    return rope2d_freqs_grid
+
+
+def precompute_rope4d_freqs_grid(
+        dim, 
+        rope2d_normalized_by_hw, 
+        max_scales=128, 
+        max_frames=128, 
+        max_height=2048 // 8, 
+        max_width=2048 // 8, 
+        base=10000.0, 
+        device=None, 
+        activated_h_div_w_templates=[], 
+        text_maxlen=0, 
+        pn=None, 
+        args=None,
+        **kwargs,
+):
+    # split the dimension into three parts, one for x, one for y, and one for t
+    print(f'[precompute_rope4d_freqs_grid: 4d]: start')
+    assert dim % 2 == 0, f'Only support dim % 2 == 0, but got dim={dim}'
+    dim_div_2 = dim // 2
+    num_of_freqs = int(np.ceil(dim_div_2 / 4))
+    inv_freq = 1.0 / (base ** (torch.arange(num_of_freqs, dtype=torch.int64).float().to(device) / num_of_freqs)) # namely theta, 1 / (10000^(i/dim_div_4)), i=0,2,..., dim_div_4-2, totally dim_div_4 / 2 elems
+    t_scales = torch.arange(text_maxlen+max_scales, device=device, dtype=torch.int64).type_as(inv_freq)
+    t_frames = torch.arange(max_frames, device=device, dtype=torch.int64).type_as(inv_freq)
+    t_height = torch.arange(max_height, device=device, dtype=torch.int64).type_as(inv_freq)
+    t_width = torch.arange(max_width, device=device, dtype=torch.int64).type_as(inv_freq)
+    freqs_scales = torch.outer(t_scales, inv_freq)  # (text_maxlen+max_scales, ceil(dim_div_2 / 4)), namely x*theta
+    freqs_frames = torch.outer(t_frames, inv_freq)  # (max_frames, ceil(dim_div_2 / 4)), namely x*theta
+    freqs_height = torch.outer(t_height, inv_freq)  # (max_height, ceil(dim_div_2 / 4)), namely y*theta
+    freqs_width = torch.outer(t_width, inv_freq)  # (max_width, ceil(dim_div_2 / 4)), namely x*theta
+    assert num_of_freqs*4==dim_div_2
+    freqs_scales = torch.stack([torch.cos(freqs_scales), torch.sin(freqs_scales)], dim=0)
+    freqs_frames = torch.stack([torch.cos(freqs_frames), torch.sin(freqs_frames)], dim=0)
+    freqs_height = torch.stack([torch.cos(freqs_height), torch.sin(freqs_height)], dim=0)
+    freqs_width = torch.stack([torch.cos(freqs_width), torch.sin(freqs_width)], dim=0)
+    tm = text_maxlen
+    rope_text_embeds = torch.cat([
+        freqs_scales[   :,   :tm,  None,   None,   None,   :].expand(-1, -1, -1, -1, -1, -1),
+        freqs_frames[   :,  None,    :1,   None,   None,   :].expand(-1, tm, -1, -1, -1, -1),
+        freqs_height[   :,  None,  None,     :1,   None,   :].expand(-1, tm, -1, -1, -1, -1),
+        freqs_width[    :,  None,  None,   None,     :1,   :].expand(-1, tm, -1, -1, -1, -1),
+    ], dim=-1)  # (2, tm, 1, 1, 1, dim_div_2)
+    rope_text_embeds = rope_text_embeds.reshape(2, 1, 1, 1, tm, dim_div_2)
+    rope2d_freqs_grid = {}
+    rope2d_freqs_grid['freqs_text'] = rope_text_embeds # (2, 1, 1, 1, text_maxlen, dim / 2)
+    rope2d_freqs_grid['freqs_scales'] = freqs_scales[:, tm:] # (2, max_scales, ceil(dim_div_2 / 4))
+    rope2d_freqs_grid['freqs_frames'] = freqs_frames # (2, max_frames, ceil(dim_div_2 / 4))
+    rope2d_freqs_grid['freqs_height'] = freqs_height # (2, max_height, ceil(dim_div_2 / 4))
+    rope2d_freqs_grid['freqs_width'] = freqs_width # (2, max_width, ceil(dim_div_2 / 4))
+    return rope2d_freqs_grid
+
+def apply_rotary_emb(q, k, rope_cache):
+    device_type = q.device.type
+    device_type = device_type if isinstance(device_type, str) and device_type != "mps" else "cpu"
+    qk = [q, k]
+    rope_cache = rope_cache[:,0]
+    with torch.autocast(device_type=device_type, enabled=False):
+        for i in range(2):
+            qk[i] = qk[i].reshape(*qk[i].shape[:-1], -1, 2)
+            tmp1 = qk[i][..., 1] * rope_cache[1]
+            tmp2 = qk[i][..., 0] * rope_cache[1]
+            qk[i][..., 0].mul_(rope_cache[0]).sub_(tmp1)
+            qk[i][..., 1].mul_(rope_cache[0]).add_(tmp2)
+            qk[i] = qk[i].reshape(*qk[i].shape[:-2], -1)
+        q, k = qk
+        # qk = qk.reshape(*qk.shape[:-1], -1, 2) #(2, batch_size, heads, seq_len, half_head_dim, 2)
+        # qk = torch.stack([
+        #     qk[...,0] * rope_cache[0] - qk[...,1] * rope_cache[1],
+        #     qk[...,0] * rope_cache[1] + qk[...,1] * rope_cache[0],
+        # ], dim=-1) # (2, batch_size, heads, seq_len, half_head_dim, 2), here stack + reshape should not be concate
+        # qk = qk.reshape(*qk.shape[:-2], -1) #(2, batch_size, heads, seq_len, head_dim)
+        # q, k = qk.unbind(dim=0) # (batch_size, heads, seq_len, head_dim)
+    return q, k

grn/models/umt5/fsdp.py

@@ -0,0 +1,42 @@
+import gc
+from functools import partial
+
+import torch
+from torch.distributed.fsdp import FullyShardedDataParallel as FSDP
+from torch.distributed.fsdp import MixedPrecision, ShardingStrategy
+from torch.distributed.fsdp.wrap import lambda_auto_wrap_policy
+from torch.distributed.utils import _free_storage
+
+
+def shard_model(
+    model,
+    device_id,
+    param_dtype=torch.bfloat16,
+    reduce_dtype=torch.float32,
+    buffer_dtype=torch.float32,
+    process_group=None,
+    sharding_strategy=ShardingStrategy.FULL_SHARD,
+    sync_module_states=True,
+):
+    model = FSDP(
+        module=model,
+        process_group=process_group,
+        sharding_strategy=sharding_strategy,
+        auto_wrap_policy=partial(
+            lambda_auto_wrap_policy, lambda_fn=lambda m: m in model.blocks),
+        mixed_precision=MixedPrecision(
+            param_dtype=param_dtype,
+            reduce_dtype=reduce_dtype,
+            buffer_dtype=buffer_dtype),
+        device_id=device_id,
+        sync_module_states=sync_module_states)
+    return model
+
+
+def free_model(model):
+    for m in model.modules():
+        if isinstance(m, FSDP):
+            _free_storage(m._handle.flat_param.data)
+    del model
+    gc.collect()
+    torch.cuda.empty_cache()

grn/models/umt5/t5.py

@@ -0,0 +1,514 @@
+import logging
+import math
+from functools import partial
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from grn.models.umt5.fsdp import shard_model
+from grn.models.umt5.umt5_tokenizers import HuggingfaceTokenizer
+
+__all__ = [
+    'T5Model',
+    'T5Encoder',
+    'T5Decoder',
+    'T5EncoderModel',
+]
+
+
+def fp16_clamp(x):
+    if x.dtype == torch.float16 and torch.isinf(x).any():
+        clamp = torch.finfo(x.dtype).max - 1000
+        x = torch.clamp(x, min=-clamp, max=clamp)
+    return x
+
+
+def init_weights(m):
+    if isinstance(m, T5LayerNorm):
+        nn.init.ones_(m.weight)
+    elif isinstance(m, T5Model):
+        nn.init.normal_(m.token_embedding.weight, std=1.0)
+    elif isinstance(m, T5FeedForward):
+        nn.init.normal_(m.gate[0].weight, std=m.dim**-0.5)
+        nn.init.normal_(m.fc1.weight, std=m.dim**-0.5)
+        nn.init.normal_(m.fc2.weight, std=m.dim_ffn**-0.5)
+    elif isinstance(m, T5Attention):
+        nn.init.normal_(m.q.weight, std=(m.dim * m.dim_attn)**-0.5)
+        nn.init.normal_(m.k.weight, std=m.dim**-0.5)
+        nn.init.normal_(m.v.weight, std=m.dim**-0.5)
+        nn.init.normal_(m.o.weight, std=(m.num_heads * m.dim_attn)**-0.5)
+    elif isinstance(m, T5RelativeEmbedding):
+        nn.init.normal_(
+            m.embedding.weight, std=(2 * m.num_buckets * m.num_heads)**-0.5)
+
+
+class GELU(nn.Module):
+
+    def forward(self, x):
+        return 0.5 * x * (1.0 + torch.tanh(
+            math.sqrt(2.0 / math.pi) * (x + 0.044715 * torch.pow(x, 3.0))))
+
+
+class T5LayerNorm(nn.Module):
+
+    def __init__(self, dim, eps=1e-6):
+        super(T5LayerNorm, self).__init__()
+        self.dim = dim
+        self.eps = eps
+        self.weight = nn.Parameter(torch.ones(dim))
+
+    def forward(self, x):
+        x = x * torch.rsqrt(x.float().pow(2).mean(dim=-1, keepdim=True) +
+                            self.eps)
+        if self.weight.dtype in [torch.float16, torch.bfloat16]:
+            x = x.type_as(self.weight)
+        return self.weight * x
+
+
+class T5Attention(nn.Module):
+
+    def __init__(self, dim, dim_attn, num_heads, dropout=0.1):
+        assert dim_attn % num_heads == 0
+        super(T5Attention, self).__init__()
+        self.dim = dim
+        self.dim_attn = dim_attn
+        self.num_heads = num_heads
+        self.head_dim = dim_attn // num_heads
+
+        # layers
+        self.q = nn.Linear(dim, dim_attn, bias=False)
+        self.k = nn.Linear(dim, dim_attn, bias=False)
+        self.v = nn.Linear(dim, dim_attn, bias=False)
+        self.o = nn.Linear(dim_attn, dim, bias=False)
+        self.dropout = nn.Dropout(dropout)
+
+    def forward(self, x, context=None, mask=None, pos_bias=None):
+        """
+        x:          [B, L1, C].
+        context:    [B, L2, C] or None.
+        mask:       [B, L2] or [B, L1, L2] or None.
+        """
+        # check inputs
+        context = x if context is None else context
+        b, n, c = x.size(0), self.num_heads, self.head_dim
+
+        # compute query, key, value
+        q = self.q(x).view(b, -1, n, c)
+        k = self.k(context).view(b, -1, n, c)
+        v = self.v(context).view(b, -1, n, c)
+
+        # attention bias
+        attn_bias = x.new_zeros(b, n, q.size(1), k.size(1))
+        if pos_bias is not None:
+            attn_bias += pos_bias
+        if mask is not None:
+            assert mask.ndim in [2, 3]
+            mask = mask.view(b, 1, 1,
+                             -1) if mask.ndim == 2 else mask.unsqueeze(1)
+            attn_bias.masked_fill_(mask == 0, torch.finfo(x.dtype).min)
+
+        # compute attention (T5 does not use scaling)
+        attn = torch.einsum('binc,bjnc->bnij', q, k) + attn_bias
+        attn = F.softmax(attn.float(), dim=-1).type_as(attn)
+        x = torch.einsum('bnij,bjnc->binc', attn, v)
+
+        # output
+        x = x.reshape(b, -1, n * c)
+        x = self.o(x)
+        x = self.dropout(x)
+        return x
+
+
+class T5FeedForward(nn.Module):
+
+    def __init__(self, dim, dim_ffn, dropout=0.1):
+        super(T5FeedForward, self).__init__()
+        self.dim = dim
+        self.dim_ffn = dim_ffn
+
+        # layers
+        self.gate = nn.Sequential(nn.Linear(dim, dim_ffn, bias=False), GELU())
+        self.fc1 = nn.Linear(dim, dim_ffn, bias=False)
+        self.fc2 = nn.Linear(dim_ffn, dim, bias=False)
+        self.dropout = nn.Dropout(dropout)
+
+    def forward(self, x):
+        x = self.fc1(x) * self.gate(x)
+        x = self.dropout(x)
+        x = self.fc2(x)
+        x = self.dropout(x)
+        return x
+
+
+class T5SelfAttention(nn.Module):
+
+    def __init__(self,
+                 dim,
+                 dim_attn,
+                 dim_ffn,
+                 num_heads,
+                 num_buckets,
+                 shared_pos=True,
+                 dropout=0.1):
+        super(T5SelfAttention, self).__init__()
+        self.dim = dim
+        self.dim_attn = dim_attn
+        self.dim_ffn = dim_ffn
+        self.num_heads = num_heads
+        self.num_buckets = num_buckets
+        self.shared_pos = shared_pos
+
+        # layers
+        self.norm1 = T5LayerNorm(dim)
+        self.attn = T5Attention(dim, dim_attn, num_heads, dropout)
+        self.norm2 = T5LayerNorm(dim)
+        self.ffn = T5FeedForward(dim, dim_ffn, dropout)
+        self.pos_embedding = None if shared_pos else T5RelativeEmbedding(
+            num_buckets, num_heads, bidirectional=True)
+
+    def forward(self, x, mask=None, pos_bias=None):
+        e = pos_bias if self.shared_pos else self.pos_embedding(
+            x.size(1), x.size(1))
+        x = fp16_clamp(x + self.attn(self.norm1(x), mask=mask, pos_bias=e))
+        x = fp16_clamp(x + self.ffn(self.norm2(x)))
+        return x
+
+
+class T5CrossAttention(nn.Module):
+
+    def __init__(self,
+                 dim,
+                 dim_attn,
+                 dim_ffn,
+                 num_heads,
+                 num_buckets,
+                 shared_pos=True,
+                 dropout=0.1):
+        super(T5CrossAttention, self).__init__()
+        self.dim = dim
+        self.dim_attn = dim_attn
+        self.dim_ffn = dim_ffn
+        self.num_heads = num_heads
+        self.num_buckets = num_buckets
+        self.shared_pos = shared_pos
+
+        # layers
+        self.norm1 = T5LayerNorm(dim)
+        self.self_attn = T5Attention(dim, dim_attn, num_heads, dropout)
+        self.norm2 = T5LayerNorm(dim)
+        self.cross_attn = T5Attention(dim, dim_attn, num_heads, dropout)
+        self.norm3 = T5LayerNorm(dim)
+        self.ffn = T5FeedForward(dim, dim_ffn, dropout)
+        self.pos_embedding = None if shared_pos else T5RelativeEmbedding(
+            num_buckets, num_heads, bidirectional=False)
+
+    def forward(self,
+                x,
+                mask=None,
+                encoder_states=None,
+                encoder_mask=None,
+                pos_bias=None):
+        e = pos_bias if self.shared_pos else self.pos_embedding(
+            x.size(1), x.size(1))
+        x = fp16_clamp(x + self.self_attn(self.norm1(x), mask=mask, pos_bias=e))
+        x = fp16_clamp(x + self.cross_attn(
+            self.norm2(x), context=encoder_states, mask=encoder_mask))
+        x = fp16_clamp(x + self.ffn(self.norm3(x)))
+        return x
+
+
+class T5RelativeEmbedding(nn.Module):
+
+    def __init__(self, num_buckets, num_heads, bidirectional, max_dist=128):
+        super(T5RelativeEmbedding, self).__init__()
+        self.num_buckets = num_buckets
+        self.num_heads = num_heads
+        self.bidirectional = bidirectional
+        self.max_dist = max_dist
+
+        # layers
+        self.embedding = nn.Embedding(num_buckets, num_heads)
+
+    def forward(self, lq, lk):
+        device = self.embedding.weight.device
+        # rel_pos = torch.arange(lk).unsqueeze(0).to(device) - \
+        #     torch.arange(lq).unsqueeze(1).to(device)
+        rel_pos = torch.arange(lk, device=device).unsqueeze(0) - \
+            torch.arange(lq, device=device).unsqueeze(1)
+        rel_pos = self._relative_position_bucket(rel_pos)
+        rel_pos_embeds = self.embedding(rel_pos)
+        rel_pos_embeds = rel_pos_embeds.permute(2, 0, 1).unsqueeze(
+            0)  # [1, N, Lq, Lk]
+        return rel_pos_embeds.contiguous()
+
+    def _relative_position_bucket(self, rel_pos):
+        # preprocess
+        if self.bidirectional:
+            num_buckets = self.num_buckets // 2
+            rel_buckets = (rel_pos > 0).long() * num_buckets
+            rel_pos = torch.abs(rel_pos)
+        else:
+            num_buckets = self.num_buckets
+            rel_buckets = 0
+            rel_pos = -torch.min(rel_pos, torch.zeros_like(rel_pos))
+
+        # embeddings for small and large positions
+        max_exact = num_buckets // 2
+        rel_pos_large = max_exact + (torch.log(rel_pos.float() / max_exact) /
+                                     math.log(self.max_dist / max_exact) *
+                                     (num_buckets - max_exact)).long()
+        rel_pos_large = torch.min(
+            rel_pos_large, torch.full_like(rel_pos_large, num_buckets - 1))
+        rel_buckets += torch.where(rel_pos < max_exact, rel_pos, rel_pos_large)
+        return rel_buckets
+
+
+class T5Encoder(nn.Module):
+
+    def __init__(self,
+                 vocab,
+                 dim,
+                 dim_attn,
+                 dim_ffn,
+                 num_heads,
+                 num_layers,
+                 num_buckets,
+                 shared_pos=True,
+                 dropout=0.1):
+        super(T5Encoder, self).__init__()
+        self.dim = dim
+        self.dim_attn = dim_attn
+        self.dim_ffn = dim_ffn
+        self.num_heads = num_heads
+        self.num_layers = num_layers
+        self.num_buckets = num_buckets
+        self.shared_pos = shared_pos
+
+        # layers
+        self.token_embedding = vocab if isinstance(vocab, nn.Embedding) \
+            else nn.Embedding(vocab, dim)
+        self.pos_embedding = T5RelativeEmbedding(
+            num_buckets, num_heads, bidirectional=True) if shared_pos else None
+        self.dropout = nn.Dropout(dropout)
+        self.blocks = nn.ModuleList([
+            T5SelfAttention(dim, dim_attn, dim_ffn, num_heads, num_buckets,
+                            shared_pos, dropout) for _ in range(num_layers)
+        ])
+        self.norm = T5LayerNorm(dim)
+
+        # initialize weights
+        self.apply(init_weights)
+
+    def forward(self, ids, mask=None):
+        x = self.token_embedding(ids)
+        x = self.dropout(x)
+        e = self.pos_embedding(x.size(1),
+                               x.size(1)) if self.shared_pos else None
+        for block in self.blocks:
+            x = block(x, mask, pos_bias=e)
+        x = self.norm(x)
+        x = self.dropout(x)
+        return x
+
+
+class T5Decoder(nn.Module):
+
+    def __init__(self,
+                 vocab,
+                 dim,
+                 dim_attn,
+                 dim_ffn,
+                 num_heads,
+                 num_layers,
+                 num_buckets,
+                 shared_pos=True,
+                 dropout=0.1):
+        super(T5Decoder, self).__init__()
+        self.dim = dim
+        self.dim_attn = dim_attn
+        self.dim_ffn = dim_ffn
+        self.num_heads = num_heads
+        self.num_layers = num_layers
+        self.num_buckets = num_buckets
+        self.shared_pos = shared_pos
+
+        # layers
+        self.token_embedding = vocab if isinstance(vocab, nn.Embedding) \
+            else nn.Embedding(vocab, dim)
+        self.pos_embedding = T5RelativeEmbedding(
+            num_buckets, num_heads, bidirectional=False) if shared_pos else None
+        self.dropout = nn.Dropout(dropout)
+        self.blocks = nn.ModuleList([
+            T5CrossAttention(dim, dim_attn, dim_ffn, num_heads, num_buckets,
+                             shared_pos, dropout) for _ in range(num_layers)
+        ])
+        self.norm = T5LayerNorm(dim)
+
+        # initialize weights
+        self.apply(init_weights)
+
+    def forward(self, ids, mask=None, encoder_states=None, encoder_mask=None):
+        b, s = ids.size()
+
+        # causal mask
+        if mask is None:
+            mask = torch.tril(torch.ones(1, s, s).to(ids.device))
+        elif mask.ndim == 2:
+            mask = torch.tril(mask.unsqueeze(1).expand(-1, s, -1))
+
+        # layers
+        x = self.token_embedding(ids)
+        x = self.dropout(x)
+        e = self.pos_embedding(x.size(1),
+                               x.size(1)) if self.shared_pos else None
+        for block in self.blocks:
+            x = block(x, mask, encoder_states, encoder_mask, pos_bias=e)
+        x = self.norm(x)
+        x = self.dropout(x)
+        return x
+
+
+class T5Model(nn.Module):
+
+    def __init__(self,
+                 vocab_size,
+                 dim,
+                 dim_attn,
+                 dim_ffn,
+                 num_heads,
+                 encoder_layers,
+                 decoder_layers,
+                 num_buckets,
+                 shared_pos=True,
+                 dropout=0.1):
+        super(T5Model, self).__init__()
+        self.vocab_size = vocab_size
+        self.dim = dim
+        self.dim_attn = dim_attn
+        self.dim_ffn = dim_ffn
+        self.num_heads = num_heads
+        self.encoder_layers = encoder_layers
+        self.decoder_layers = decoder_layers
+        self.num_buckets = num_buckets
+
+        # layers
+        self.token_embedding = nn.Embedding(vocab_size, dim)
+        self.encoder = T5Encoder(self.token_embedding, dim, dim_attn, dim_ffn,
+                                 num_heads, encoder_layers, num_buckets,
+                                 shared_pos, dropout)
+        self.decoder = T5Decoder(self.token_embedding, dim, dim_attn, dim_ffn,
+                                 num_heads, decoder_layers, num_buckets,
+                                 shared_pos, dropout)
+        self.head = nn.Linear(dim, vocab_size, bias=False)
+
+        # initialize weights
+        self.apply(init_weights)
+
+    def forward(self, encoder_ids, encoder_mask, decoder_ids, decoder_mask):
+        x = self.encoder(encoder_ids, encoder_mask)
+        x = self.decoder(decoder_ids, decoder_mask, x, encoder_mask)
+        x = self.head(x)
+        return x
+
+
+def _t5(name,
+        encoder_only=False,
+        decoder_only=False,
+        return_tokenizer=False,
+        tokenizer_kwargs={},
+        dtype=torch.float32,
+        device='cpu',
+        **kwargs):
+    # sanity check
+    assert not (encoder_only and decoder_only)
+
+    # params
+    if encoder_only:
+        model_cls = T5Encoder
+        kwargs['vocab'] = kwargs.pop('vocab_size')
+        kwargs['num_layers'] = kwargs.pop('encoder_layers')
+        _ = kwargs.pop('decoder_layers')
+    elif decoder_only:
+        model_cls = T5Decoder
+        kwargs['vocab'] = kwargs.pop('vocab_size')
+        kwargs['num_layers'] = kwargs.pop('decoder_layers')
+        _ = kwargs.pop('encoder_layers')
+    else:
+        model_cls = T5Model
+
+    # init model
+    with torch.device(device):
+        model = model_cls(**kwargs)
+
+    # set device
+    model = model.to(dtype=dtype, device=device)
+
+    # init tokenizer
+    if return_tokenizer:
+        from .tokenizers import HuggingfaceTokenizer
+        tokenizer = HuggingfaceTokenizer(f'google/{name}', **tokenizer_kwargs)
+        return model, tokenizer
+    else:
+        return model
+
+
+def umt5_xxl(**kwargs):
+    cfg = dict(
+        vocab_size=256384,
+        dim=4096,
+        dim_attn=4096,
+        dim_ffn=10240,
+        num_heads=64,
+        encoder_layers=24,
+        decoder_layers=24,
+        num_buckets=32,
+        shared_pos=False,
+        dropout=0.1)
+    cfg.update(**kwargs)
+    return _t5('umt5-xxl', **cfg)
+
+
+class T5EncoderModel:
+
+    def __init__(
+        self,
+        text_len,
+        dtype=torch.bfloat16,
+        device=torch.cuda.current_device(),
+        checkpoint_path=None,
+        tokenizer_path=None,
+        enable_fsdp=False,
+    ):
+        self.text_len = text_len
+        self.dtype = dtype
+        self.device = device
+        self.checkpoint_path = checkpoint_path
+        self.tokenizer_path = tokenizer_path
+
+        # init model
+        model = umt5_xxl(
+            encoder_only=True,
+            return_tokenizer=False,
+            dtype=dtype,
+            device=device).eval().requires_grad_(False)
+        logging.info(f'loading {checkpoint_path}')
+        model.load_state_dict(torch.load(checkpoint_path, map_location='cpu'))
+        self.model = model
+        if enable_fsdp:
+            shard_fn = partial(shard_model, device_id=device)
+            self.model = shard_fn(self.model, sync_module_states=False)
+        else:
+            self.model.to(self.device)
+        # init tokenizer
+        self.tokenizer = HuggingfaceTokenizer(
+            name=tokenizer_path, seq_len=text_len, clean='whitespace')
+
+    def __call__(self, texts, device):
+        ids, mask = self.tokenizer(
+            texts, return_mask=True, add_special_tokens=True)
+        ids = ids.to(device)
+        mask = mask.to(device)
+        seq_lens = mask.gt(0).sum(dim=1).long()
+        context = self.model(ids, mask)
+        return [u[:v] for u, v in zip(context, seq_lens)]

grn/models/umt5/umt5_tokenizers.py

@@ -0,0 +1,81 @@
+import html
+import string
+
+import ftfy
+import regex as re
+from transformers import AutoTokenizer
+
+__all__ = ['HuggingfaceTokenizer']
+
+
+def basic_clean(text):
+    text = ftfy.fix_text(text)
+    text = html.unescape(html.unescape(text))
+    return text.strip()
+
+
+def whitespace_clean(text):
+    text = re.sub(r'\s+', ' ', text)
+    text = text.strip()
+    return text
+
+
+def canonicalize(text, keep_punctuation_exact_string=None):
+    text = text.replace('_', ' ')
+    if keep_punctuation_exact_string:
+        text = keep_punctuation_exact_string.join(
+            part.translate(str.maketrans('', '', string.punctuation))
+            for part in text.split(keep_punctuation_exact_string))
+    else:
+        text = text.translate(str.maketrans('', '', string.punctuation))
+    text = text.lower()
+    text = re.sub(r'\s+', ' ', text)
+    return text.strip()
+
+
+class HuggingfaceTokenizer:
+
+    def __init__(self, name, seq_len=None, clean=None, **kwargs):
+        assert clean in (None, 'whitespace', 'lower', 'canonicalize')
+        self.name = name
+        self.seq_len = seq_len
+        self.clean = clean
+
+        # init tokenizer
+        self.tokenizer = AutoTokenizer.from_pretrained(name, **kwargs)
+        self.vocab_size = self.tokenizer.vocab_size
+
+    def __call__(self, sequence, **kwargs):
+        return_mask = kwargs.pop('return_mask', False)
+
+        # arguments
+        _kwargs = {'return_tensors': 'pt'}
+        if self.seq_len is not None:
+            _kwargs.update({
+                'padding': 'max_length',
+                'truncation': True,
+                'max_length': self.seq_len
+            })
+        _kwargs.update(**kwargs)
+
+        # tokenization
+        if isinstance(sequence, str):
+            sequence = [sequence]
+        if self.clean:
+            sequence = [self._clean(u) for u in sequence]
+        ids = self.tokenizer(sequence, **_kwargs)
+
+        # output
+        if return_mask:
+            return ids.input_ids, ids.attention_mask
+        else:
+            return ids.input_ids
+
+    def _clean(self, text):
+        if self.clean == 'whitespace':
+            text = whitespace_clean(basic_clean(text))
+        elif self.clean == 'lower':
+            text = whitespace_clean(basic_clean(text)).lower()
+        elif self.clean == 'canonicalize':
+            text = canonicalize(basic_clean(text))
+        return text

grn/schedules/__init__.py

@@ -0,0 +1,6 @@
+def get_encode_decode_func(dynamic_scale_schedule):
+    if 'GRN_vae_stride16' in dynamic_scale_schedule:
+        from grn.schedules.global_refine import video_encode, video_decode, get_visual_rope_embeds, get_scale_pack_info
+    else:
+        raise NotImplementedError(f'{dynamic_scale_schedule} is unsupported')
+    return video_encode, video_decode, get_visual_rope_embeds, get_scale_pack_info

grn/schedules/dynamic_resolution.py

@@ -0,0 +1,99 @@
+import json
+import math
+import copy
+
+import tqdm
+import numpy as np
+
+
+def get_first_full_spatial_size_scale_index(vae_scale_schedule):
+    for si, (pt, ph, pw) in enumerate(vae_scale_schedule):
+        if vae_scale_schedule[si][-2:] == vae_scale_schedule[-1][-2:]:
+            return si
+
+def get_full_spatial_size_scale_indices(vae_scale_schedule):
+    full_spatial_size_scale_indices = []
+    for si, (pt, ph, pw) in enumerate(vae_scale_schedule):
+        if vae_scale_schedule[si][-2:] == vae_scale_schedule[-1][-2:]:
+            full_spatial_size_scale_indices.append(si)
+    return full_spatial_size_scale_indices
+
+def get_ratio2hws_pixels2scales(dynamic_scale_schedule, train_h_div_w_list, video_frames):
+    compressed_frames = video_frames // 4 + 1
+    if dynamic_scale_schedule in ['GRN_vae_stride16']:
+        assert type(train_h_div_w_list) is str
+        train_h_div_w_list = json.loads(train_h_div_w_list)
+        if len(train_h_div_w_list) == 0:
+            train_h_div_w_list = [3/1, 5/2, 2/1, 16/9, 3/2, 4/3, 116/100, 1, 100/116, 3/4, 2/3, 9/16, 1/2, 2/5, 1/3]
+        vae_stride = 16
+        if 'vae_stride' in dynamic_scale_schedule:
+            vae_stride = int(dynamic_scale_schedule.split('vae_stride')[-1])
+        dynamic_resolution_h_w = {}
+        for h_div_w in train_h_div_w_list:
+            ratio = int(h_div_w*1000)/1000
+            dynamic_resolution_h_w[ratio] = {}
+            for pn in ['0.06M', '0.25M', '0.41M', '0.92M', '1M', '2M']:
+                if pn == '0.06M': # 256x256, 192p
+                    scale = 8
+                elif pn == '0.25M': # 512x512, 384p
+                    scale = 16
+                elif pn == '0.41M': # 640x640, 480p
+                    scale = 20
+                elif pn == '0.92M': # 960x960, 720p
+                    scale = 30
+                elif pn == '1M': # 1024x1024, 768p
+                    scale = 32
+                elif pn == '2M': # 1440x1440, 1080p
+                    scale = 45
+                if vae_stride == 16:
+                    scale = scale * 2
+                elif vae_stride == 32:
+                    scale = scale * 1
+                else:
+                    raise ValueError(f'vae_stride {vae_stride} is not supported')
+                area = scale * scale
+                pw_float = math.sqrt(area / h_div_w)
+                ph_float = pw_float * h_div_w
+                ph, pw = int(np.round(ph_float)), int(np.round(pw_float))
+                scales = [(ph,pw)]
+                pixel = (scales[-1][0] * vae_stride, scales[-1][1] * vae_stride)
+                dynamic_resolution_h_w[ratio][pn] = {
+                    'pixel': pixel,
+                    'scales': scales
+                }
+        for ratio in dynamic_resolution_h_w:
+            for pn in dynamic_resolution_h_w[ratio]:
+                base_scale_schedule = dynamic_resolution_h_w[ratio][pn]['scales']
+                scales_in_one_clip = len(base_scale_schedule)
+                dynamic_resolution_h_w[ratio][pn]['pt2scale_schedule'] = {}
+                for pt in range(1, compressed_frames+1, 1):
+                    dynamic_resolution_h_w[ratio][pn]['pt2scale_schedule'][pt] = [(pt, h, w) for h, w in base_scale_schedule]
+                dynamic_resolution_h_w[ratio][pn]['image_scales'] = scales_in_one_clip
+                dynamic_resolution_h_w[ratio][pn]['scales_in_one_clip'] = scales_in_one_clip
+                dynamic_resolution_h_w[ratio][pn]['max_video_scales'] = len(dynamic_resolution_h_w[ratio][pn]['pt2scale_schedule'][compressed_frames])
+                del dynamic_resolution_h_w[ratio][pn]['scales']
+    else:
+        raise ValueError(f'dynamic_scale_schedule={dynamic_scale_schedule} not implemented')
+    return dynamic_resolution_h_w
+
+def get_dynamic_resolution_meta(dynamic_scale_schedule, train_h_div_w_list, video_frames):
+    dynamic_resolution_h_w = get_ratio2hws_pixels2scales(dynamic_scale_schedule, train_h_div_w_list, video_frames)
+    h_div_w_templates = []
+    for h_div_w in dynamic_resolution_h_w.keys():
+        h_div_w_templates.append(h_div_w)
+    h_div_w_templates = np.array(h_div_w_templates)
+    return dynamic_resolution_h_w, h_div_w_templates
+
+def get_h_div_w_template2indices(h_div_w_list, h_div_w_templates):
+    indices = list(range(len(h_div_w_list)))
+    h_div_w_template2indices = {}
+    pbar = tqdm.tqdm(total=len(indices), desc='get_h_div_w_template2indices...')
+    for h_div_w, index in zip(h_div_w_list, indices):
+        pbar.update(1)
+        nearest_h_div_w_template_ = h_div_w_templates[np.argmin(np.abs(h_div_w-h_div_w_templates))]
+        if nearest_h_div_w_template_ not in h_div_w_template2indices:
+            h_div_w_template2indices[nearest_h_div_w_template_] = []
+        h_div_w_template2indices[nearest_h_div_w_template_].append(index)
+    for h_div_w_template_, sub_indices in h_div_w_template2indices.items():
+        h_div_w_template2indices[h_div_w_template_] = np.array(sub_indices)
+    return h_div_w_template2indices

grn/schedules/global_refine.py

@@ -0,0 +1,220 @@
+import os
+import json
+import math
+import bisect
+
+import numpy as np
+import torch
+import torch.nn.functional as F
+
+from grn.utils_t2iv.hbq_util_t2iv import multiclass_labels2onehot_input
+
+def get_scale_pack_info(**kwargs):
+    return
+
+def flatten_two_level_list(two_level_list):
+    flatten_list = []
+    for item in two_level_list:
+        flatten_list.extend(item)
+    return flatten_list
+
+def shift_pt(pt, alpha):
+    """shift pt (signal ratio) to lower one, recommand alpha=sqrt(height*width/256/256)"""
+    if alpha > 1000:
+        alpha = alpha - 1000
+    noise_pt = 1 - pt
+    noise_pt = alpha * noise_pt / (1+(alpha-1)*noise_pt) # shift noise_pt to higer one
+    pt = 1 - noise_pt
+    return pt
+
+def video_encode(
+    vae,
+    inp_B3HW,
+    vae_features=None,
+    device='cuda',
+    args=None,
+    infer_mode=False,
+    rope2d_freqs_grid=None,
+    dynamic_resolution_h_w=None,
+    tokens_remain=9999999,
+    text_lens=[],
+    caption_nums=[],
+    rank_vary_generator=None,
+    vis_verbose=False,
+    meta_list=None,
+    **kwargs,
+):
+    if rank_vary_generator is not None:
+        numpy_generator = rank_vary_generator['numpy_generator']
+        torch_cuda_generator = rank_vary_generator['torch_cuda_generator']
+    else:
+        numpy_generator = np.random.default_rng()
+        torch_cuda_generator = torch.Generator(device='cuda')
+    
+    if vae_features is None:
+        raw_features, _, _ = vae.encode_for_raw_features(inp_B3HW, scale_schedule=None, slice=True)
+        raw_features_list = [raw_features]
+        x_recon_raw = vae.decode(raw_features[-1], slice=True)
+        x_recon_raw = torch.clamp(x_recon_raw, min=-1, max=1)
+        print(f'raw_features[-1].shape: {raw_features[-1].shape}')
+    else:
+        raw_features_list = vae_features
+    # raw_features_list: list of [1,d,t,h,w]:
+    # import pdb; pdb.set_trace()
+    gt_all_bit_indices = []
+    pred_all_bit_indices = []
+    var_input_list = []
+    sequece_packing_scales = [] # with trunk
+    h_div_w_template_list = np.array(list(dynamic_resolution_h_w.keys()))
+    visual_rope_cache_list = []
+    other_info_by_scale = []
+    scale_lengths = []
+    with torch.amp.autocast('cuda', enabled = False):
+        for example_ind, raw_features in enumerate(raw_features_list):
+            meta = meta_list[example_ind]
+            gt_all_bit_indices.append([])
+            pred_all_bit_indices.append([])
+            var_input_list.append([])
+            visual_rope_cache_list.append([])
+            other_info_by_scale.append([])
+            B, C, T, H, W = raw_features[-1].shape
+            h_div_w = H / W
+            mapped_h_div_w_template = h_div_w_template_list[np.argmin(np.abs(h_div_w-h_div_w_template_list))]
+            pn = meta['pn']
+            if meta['first_frame_condition']:
+                scale_schedule = dynamic_resolution_h_w[mapped_h_div_w_template][pn]['pt2scale_schedule'][T-1]
+            else:
+                scale_schedule = dynamic_resolution_h_w[mapped_h_div_w_template][pn]['pt2scale_schedule'][T]
+            if not infer_mode:
+                next_tokens_remain = tokens_remain - T * H * W - args.add_scale_token - text_lens[example_ind]
+                if next_tokens_remain < 0:
+                    break
+                tokens_remain = next_tokens_remain
+                scale_lengths.append(T * H * W + text_lens[example_ind] + args.add_scale_token)
+            preserve_scale_schedule = []
+            preserve_scale_schedule.append(scale_schedule[0])
+            target = raw_features[0]
+            if not infer_mode and args.log_norm_sigma > 0:
+                spt = torch.sigmoid(torch.randn(1, generator=torch_cuda_generator, device=target.device) * args.log_norm_sigma + args.log_norm_mean).item()
+                spt = shift_pt(spt, args.alpha)
+            else:
+                spt = shift_pt(numpy_generator.random(), args.alpha)
+            
+            if args.refine_mode in ['ar_discrete_GRN_ind']:
+                from grn.utils_t2iv.hbq_util_t2iv import raw_feature2index_label
+                labels = raw_feature2index_label(target, hbq_round=args.hbq_round) # [B, hbq_round * d, t, h, w]
+                classes = 2**args.hbq_round
+            elif args.refine_mode in ['ar_discrete_GRN_bit']:
+                from grn.utils_t2iv.hbq_util_t2iv import raw_feature2bit_label
+                labels = raw_feature2bit_label(target, hbq_round=args.hbq_round) # [B, hbq_round * d, t, h, w]
+                classes = 2
+                
+            random_labels = torch.randint(0, classes, size=labels.shape, generator=torch_cuda_generator, device=labels.device, dtype=labels.dtype) # random 0 or 1 labels
+            random_mask = torch.rand(size=labels.shape, generator=torch_cuda_generator, device=labels.device, dtype=target.dtype) < spt
+            mixed_xt = torch.where(random_mask, labels, random_labels) # [B, hbq_round * d, t, h, w]
+            precise_spt = random_mask.float().mean()
+            wandb_plot_index = min(9, int(precise_spt / 0.1)) # 0~9
+            
+            # get visual rope
+            if not infer_mode:
+                if meta['first_frame_condition']:
+                    visual_rope_cache_list[-1].append(get_visual_rope_embeds(rope2d_freqs_grid, scale_schedule[0], device, mapped_h_div_w_template, t_offset=1)) # (2, 1, 1, 1, pt*ph*pw, dim_div_2)
+                    visual_rope_cache_list[-1].append(get_visual_rope_embeds(rope2d_freqs_grid, (1, H, W), device, mapped_h_div_w_template, t_offset=0))
+                    visual_rope_cache_list[-1] = [torch.cat(visual_rope_cache_list[-1], dim=-2)]
+                else:
+                    visual_rope_cache_list[-1].append(get_visual_rope_embeds(rope2d_freqs_grid, scale_schedule[0], device, mapped_h_div_w_template, t_offset=0)) # (2, 1, 1, 1, pt*ph*pw, dim_div_2)
+
+            # get visual tokens
+            visual_token_dim = mixed_xt.shape[1] * classes
+            
+            if not infer_mode and meta['first_frame_condition']:
+                first_frame_labels = labels[:,:,:1] # [B, hbq_round * d, 1, h, w]
+                first_frame_tokens = multiclass_labels2onehot_input(first_frame_labels, classes).reshape(1, visual_token_dim, -1).permute(0, 2, 1)
+                cur_visual_tokens = multiclass_labels2onehot_input(mixed_xt[:,:,1:], classes).reshape(1, visual_token_dim, -1).permute(0, 2, 1)
+                cur_visual_tokens = torch.cat((cur_visual_tokens, first_frame_tokens), dim=1)
+                indices = labels[:,:,1:]
+            else:
+                cur_visual_tokens = multiclass_labels2onehot_input(mixed_xt, classes).reshape(1, visual_token_dim, -1).permute(0, 2, 1)
+                indices = labels
+            indices = indices.type(torch.long).permute(0,2,3,4,1) # [B,d,t,h,w] -> [B,t,h,w,d]
+            gt_all_bit_indices[-1].append(indices)
+            var_input_list[-1].append(cur_visual_tokens)
+            other_info_by_scale[-1].append(
+                {
+                    'largest_scale': scale_schedule[-1], 
+                    'wandb_plot_index': wandb_plot_index, 
+                    'cur_bits': indices.shape[-1],
+                    'cur_lvl': args.detail_num_lvl,
+                    'scale_token_id': precise_spt,
+                    'predict_tokens': np.prod(scale_schedule[0]),
+                    'all_tokens': scale_lengths[-1] if len(scale_lengths) else -1,
+                    'first_frame_condition': meta['first_frame_condition'],
+                }
+            )
+            sequece_packing_scales.append(preserve_scale_schedule)
+
+    gt_all_bit_indices = flatten_two_level_list(gt_all_bit_indices)
+    pred_all_bit_indices = flatten_two_level_list(pred_all_bit_indices)
+    var_input_list = flatten_two_level_list(var_input_list)
+    visual_rope_cache_list = flatten_two_level_list(visual_rope_cache_list)
+    other_info_by_scale = flatten_two_level_list(other_info_by_scale)
+
+    if infer_mode:
+        return [labels, target], x_recon_raw, [target], None, None, None
+        
+    gt_ms_idx_Bl = []
+    for item in gt_all_bit_indices:
+        _, tt, hh, ww, dd = item.shape
+        item = item.reshape(B, tt*hh*ww, dd)
+        gt_ms_idx_Bl.append(item)
+    gt_BLC = gt_ms_idx_Bl # torch.cat(gt_ms_idx_Bl, 1).contiguous().type(torch.long)
+    x_BLC = var_input_list
+    x_BLC_mask = None
+    scale_or_time_ids = None
+    return x_BLC, x_BLC_mask, scale_or_time_ids, gt_BLC, pred_all_bit_indices, visual_rope_cache_list, sequece_packing_scales, scale_lengths, other_info_by_scale
+
+def video_decode(
+    vae,
+    all_indices,
+    scale_schedule,
+    label_type,
+    args=None,
+    noise_list=None,
+    trunc_scales=-1,
+    **kwargs,
+):  
+    if trunc_scales < 0:
+        summed_codes = all_indices[-1]
+    else:
+        summed_codes = all_indices[trunc_scales-1]
+    x_recon = vae.decode(summed_codes, slice=True)
+    x_recon = torch.clamp(x_recon, min=-1, max=1)
+    x_recon_256 = None
+    return x_recon, x_recon_256
+
+def get_visual_rope_embeds(rope2d_freqs_grid, scale_schedule, device=None, mapped_h_div_w_template=None, t_offset=0):
+    # freqs_frames: (2, max_frames, dim_div_2 / 3)
+    rope2d_freqs_grid['freqs_frames'] = rope2d_freqs_grid['freqs_frames'].to(device)
+    rope2d_freqs_grid['freqs_height'] = rope2d_freqs_grid['freqs_height'].to(device)
+    rope2d_freqs_grid['freqs_width'] = rope2d_freqs_grid['freqs_width'].to(device)
+    max_height = rope2d_freqs_grid['freqs_height'].shape[1]
+    max_width = rope2d_freqs_grid['freqs_width'].shape[1]
+    extreme_h_div_w = 3
+    assert mapped_h_div_w_template <= extreme_h_div_w
+    extreme_h = max_height
+    extreme_w = extreme_h / extreme_h_div_w
+    upw = np.sqrt(extreme_h * extreme_w / mapped_h_div_w_template)
+    uph = mapped_h_div_w_template * upw
+    uph, upw = int(uph), int(upw)
+    pt, ph, pw = scale_schedule
+    assert ph <= uph and pw <= upw
+    f_frames = rope2d_freqs_grid['freqs_frames'][:, t_offset:t_offset+pt]
+    f_height = rope2d_freqs_grid['freqs_height'][:, (torch.arange(ph) * (uph / ph)).round().int()]
+    f_width = rope2d_freqs_grid['freqs_width'][:, (torch.arange(pw) * (upw / pw)).round().int()]
+    rope_embeds = torch.cat([
+        f_frames[   :,     :,  None,   None,   :].expand(-1, -1, ph, pw, -1),
+        f_height[   :,  None,      :,  None,   :].expand(-1,  pt,-1, pw, -1),
+        f_width[   :,  None,   None,      :,   :].expand(-1,  pt,ph, -1, -1),
+    ], dim=-1)  # (2, pt, ph, pw, dim_div_2)
+    rope_embeds = rope_embeds.reshape(2, 1, 1, 1, pt*ph*pw, -1)  # (2, 1, 1, 1, pt*ph*pw, dim_div_2)
+    return rope_embeds

grn/tokenizer/.gitignore

@@ -0,0 +1,25 @@
+**/__pycache__/
+lightning_logs/
+.ipynb_checkpoints/
+*.egg-info
+.pyc
+results*
+cmp_results*
+logs*
+recon.*
+*.pt
+*.npy
+dataset/
+**/span.log
+cli/batch_post.py
+model_arch*.txt
+*.mp4
+*.png
+labels
+video_vae_results
+video_vae_results_bk
+results
+video_vae_results_full
+bashrc_gpu_worker
+hj_video_vae_results
+wandb

grn/tokenizer/sample.py

@@ -0,0 +1,694 @@
+import os
+import tqdm
+import json
+import re
+import torch
+import torch.nn.functional as F
+import argparse
+import time
+import datetime
+import numpy as np
+import hashlib
+import random
+import torch.nn as nn
+from torchvision.models.inception import inception_v3
+from torch.profiler import record_function as torch_record_function
+from contextlib import nullcontext
+import lpips
+import cv2
+from einops import rearrange
+from tqdm import tqdm
+from PIL import Image
+import os.path as osp
+Image.MAX_IMAGE_PIXELS = None
+
+from videovae.modules.commitments import DiagonalGaussianDistribution
+
+import torch.distributed as dist
+from torch.multiprocessing import spawn
+from torch.nn.parallel import DistributedDataParallel as DDP
+
+import imageio
+import random
+from skimage.metrics import peak_signal_noise_ratio as psnr_loss
+from skimage.metrics import structural_similarity as ssim_loss
+
+from videovae.data import VideoData
+from videovae.utils.misc import save_video_grid, shift_dim, data_prefix_manager, rearranged_forward, seed_everything
+from videovae.utils.init_models import init_cnn_from_image, load_cnn
+from videovae.utils.arguments import MainArgs, add_model_specific_args, init_resolution
+from videovae.evaluation import get_fvd_logits, frechet_distance, load_fvd_model
+from videovae.evaluation import calculate_frechet_distance
+from videovae.evaluation import InceptionV3
+from videovae.evaluation import calculate_fvd, calculate_lpips, calculate_psnr, calculate_ssim
+
+torch.set_num_threads(32)
+os.environ["NCCL_DEBUG"] = "WARN"
+os.environ['PYTORCH_CUDA_ALLOC_CONF'] = 'expandable_segments:True'
+
+
+def calculate_batch_codebook_usage_percentage(batch_encoding_indices,n_codes):
+    if isinstance(batch_encoding_indices, list):
+        all_indices = []
+        for one_encoding_indices in batch_encoding_indices:
+            all_indices.append(one_encoding_indices.flatten())
+        all_indices = torch.cat(all_indices, dim=0)
+    else:
+        # Flatten the batch of encoding indices into a single 1D tensor
+        all_indices = batch_encoding_indices.flatten()
+    all_indices = all_indices.detach().cpu()
+    
+    # Obtain the total number of encoding indices in the batch to calculate percentages
+    total_indices = all_indices.numel()
+    
+    # Initialize a tensor to store the percentage usage of each code
+    codebook_usage = torch.zeros(n_codes, dtype=torch.long)
+    
+    # Count the number of occurrences of each index and get their frequency as percentages
+    unique_indices, counts = torch.unique(all_indices, return_counts=True)
+    
+    # Populate the corresponding percentages in the codebook_usage_percentage tensor
+    codebook_usage[unique_indices.long()] = counts
+    
+    return codebook_usage
+
+
+def disabled_train(self, mode=True):
+    """Overwrite model.train with this function to make sure train/eval mode
+    does not change anymore."""
+    return self
+
+def default_parse_args():
+    parser = argparse.ArgumentParser()
+    parser.add_argument('--vqgan_ckpt', type=str, default=None)
+    parser.add_argument('--sd_ckpt', type=str, default=None)
+    parser.add_argument('--use_frames', type=int, default=None)
+    parser.add_argument('--inference_type', type=str, choices=["image", "video", "video_concat"])
+    parser.add_argument('--save_prediction', action='store_true')
+    parser.add_argument('--save_dir',  type=str, default="results")
+    parser.add_argument('--intermediate_tensor', action='store_true')
+    parser.add_argument('--save_z', action='store_true')
+    parser.add_argument('--save_frames', action='store_true')
+    parser.add_argument('--image_recon4video', action='store_true')
+    parser.add_argument('--junke_old', action='store_true')
+    parser.add_argument('--cal_norm', action='store_true')
+    parser.add_argument('--save_samples', type=str, default=None)
+    parser.add_argument('--device', type=str, default="cuda", choices=["cpu", "cuda"])
+    parser.add_argument('--noise_scale', type=float, default=0.0)
+    parser = MainArgs.add_main_args(parser)
+    parser = VideoData.add_data_specific_args(parser)
+    args, unknown = parser.parse_known_args()
+    args, parser, vae_model = add_model_specific_args(args, parser)
+    args = parser.parse_args()
+    return args, vae_model
+
+
+def setup(rank, world_size):
+    os.environ['MASTER_ADDR'] = 'localhost'
+    os.environ['MASTER_PORT'] = str(12355+int(time.time())%1000)
+    # dist.init_process_group("nccl", rank=rank, world_size=world_size)
+    dist.init_process_group("nccl", rank=rank, world_size=world_size, timeout=datetime.timedelta(seconds=30 * 60))
+
+def cleanup():
+    dist.destroy_process_group()
+
+def main():
+    args, vae_model = default_parse_args()
+    assert len(args.dataset_list) == 1
+
+    # init data_prefix_manager
+    data_prefix_manager.set_data_root(args.data_root, username=args.username)
+    args.default_root_dir = data_prefix_manager(args.default_root_dir)
+    os.makedirs(args.default_root_dir, exist_ok=True)
+    print(args.default_root_dir)
+
+    # init intermediate_tensor_dir
+    if args.intermediate_tensor:
+        random.seed(time.time())
+        random_folder_name = hashlib.sha256(str(random.random()).encode('utf-8')).hexdigest()[:16]
+        args.intermediate_tensor_dir = os.path.join(args.default_root_dir, random_folder_name)
+        print(f"save temporal tensor to {args.intermediate_tensor_dir}")
+    
+    seed_everything(seed=0, allow_tf32=True) # ALERT: allow_tf32=True may cause accumulate error in conv3d forward > 
+
+    # init resolution
+    args.resolution = init_resolution(args.resolution, len(args.dataset_list))
+ 
+    # init profiler
+    def trace_handler(p):
+        p.export_chrome_trace(os.path.join(args.default_root_dir, f"trace_step_{p.step_num}_rank_{0}.json"))
+
+    tp = None
+    if args.turn_on_profiler:
+        tp = torch.profiler.profile(
+            activities=[
+                torch.profiler.ProfilerActivity.CPU,
+                torch.profiler.ProfilerActivity.CUDA,
+            ],
+            schedule=torch.profiler.schedule(
+                wait=args.profiler_scheduler_wait_steps,
+                warmup=3,
+                active=2,
+                repeat=1,
+            ),
+            with_stack=True,
+            record_shapes=True,
+            profile_memory=True,
+            on_trace_ready=trace_handler
+        )
+        tp.start()
+        record_function = torch_record_function
+    else:
+        record_function = nullcontext
+
+
+    vae = None
+    use_vae = None
+    num_codes = None
+    if args.vqgan_ckpt:
+        args.vqgan_ckpt = data_prefix_manager(args.vqgan_ckpt)
+    if args.tokenizer in ["hbq_tokenizer"]:
+        vae = vae_model(args)
+        state_dict = torch.load(args.vqgan_ckpt, map_location=torch.device("cpu"), weights_only=True)
+        new_state_dict = {}
+        for key in ['vae', 'ema']:
+            if (key not in state_dict) or (not state_dict[key]):
+                continue
+            if 'quantizer.scale_learnable_parameters' in state_dict[key]:
+                if len(state_dict[key]['quantizer.scale_learnable_parameters']) == 1:
+                    state_dict[key]['quantizer.scale_learnable_parameters'] = state_dict[key]['quantizer.scale_learnable_parameters'].expand(4)
+                state_dict[key]['scale_learnable_parameters'] = state_dict[key]['quantizer.scale_learnable_parameters']
+                del state_dict[key]['quantizer.scale_learnable_parameters']
+            if 'z_mean' in state_dict[key]:
+                if state_dict[key]['z_mean'].shape != vae.z_mean.shape:
+                    del state_dict[key]['z_mean']
+                    del state_dict[key]['z_std']
+            new_state_dict[key] = state_dict[key]
+            slim_model_path = args.vqgan_ckpt.replace('/checkpoints/', f'/slim_{key}/')
+            if not osp.exists(slim_model_path):
+                os.makedirs(os.path.dirname(slim_model_path), exist_ok=True)
+                torch.save({key: state_dict[key]}, slim_model_path)
+                print(f'save to {slim_model_path}')
+
+        if args.ema == "yes":
+            print("testing ema weights")
+            print(vae.load_state_dict(new_state_dict["ema"], strict=False))
+        else:
+            print("testing non ema weights")
+            print(vae.load_state_dict(new_state_dict["vae"], strict=False))
+        for name, param in vae.named_parameters():
+            if name.startswith("scale_learnable_"):
+                try:
+                    print(f"{name}: {param[:32,0,0].cpu().detach().reshape(-1).tolist()}")
+                except:
+                    print(f"{name}: {param[:32].cpu().detach().reshape(-1).tolist()}")
+        for name, param in vae.named_buffers():
+            if name.startswith("scale_learnable_"):
+                try:
+                    print(f"{name}: {param[:32,0,0].cpu().detach().reshape(-1).tolist()}")
+                except:
+                    print(f"{name}: {param[:32].cpu().detach().reshape(-1).tolist()}")
+            if ("scale_wise_std_" in name) or ("scale_wise_mean_" in name):
+                print(f"{name}: {param[:32,0,0].cpu().detach().reshape(-1).tolist()}")
+            if ('signal_' in name):
+                print(f"{name}: {param.cpu().detach().reshape(-1).tolist()}")
+        if args.tokenizer != 'hbq_tokenizer':
+            vae.enable_slicing()
+        # vae.enable_tiling()
+    else:
+        raise NotImplementedError
+
+    if args.inference_type == "video":
+        def extract_results(return_dict, world_size):
+            real_embeddings, fake_embeddings, all_real_videos, all_fake_videos, zs = [], [], [], [], []
+            if args.intermediate_tensor:
+                for rank in range(world_size):
+                    real_embeddings.append(return_dict[rank]['real_embeddings'])
+                    fake_embeddings.append(return_dict[rank]['fake_embeddings'])
+                    all_real_videos += return_dict[rank]['all_real_videos']
+                    all_fake_videos += return_dict[rank]['all_fake_videos']
+                    zs.append(return_dict[rank]['zs'])
+                real_embeddings = torch.cat(real_embeddings, 0).to('cuda:0')
+                fake_embeddings = torch.cat(fake_embeddings, 0).to('cuda:0')
+                zs = torch.cat(zs, 0).to('cuda:0')
+            else:
+                for rank in range(world_size):
+                    real_embeddings.append(return_dict[rank]['real_embeddings'])
+                    fake_embeddings.append(return_dict[rank]['fake_embeddings'])
+                    all_real_videos.append(return_dict[rank]['all_real_videos'])
+                    all_fake_videos.append(return_dict[rank]['all_fake_videos'])
+                    zs.append(return_dict[rank]['zs'])
+                real_embeddings = torch.cat(real_embeddings, 0).to('cuda:0')
+                fake_embeddings = torch.cat(fake_embeddings, 0).to('cuda:0')
+                all_real_videos = torch.cat(all_real_videos, 0)
+                all_fake_videos = torch.cat(all_fake_videos, 0)
+                zs = torch.cat(zs, 0).to('cuda:0')
+            return real_embeddings, fake_embeddings, all_real_videos, all_fake_videos, zs
+
+        def inference(mean=None, std=None, noise_scale=0):
+            world_size = torch.cuda.device_count()
+            manager = torch.multiprocessing.Manager()
+            return_dict = manager.dict()
+            ### multi-process
+            # try:
+            #     spawn(inference_DDP, args=(world_size, args, vae_model, vae, record_function, tp, use_vae, num_codes, return_dict, mean, std, noise_scale), nprocs=world_size, join=True)
+            # except Exception as e:
+            #     print(f"Error during spawn {e}")
+
+            ## single process
+            world_size = 1
+            inference_DDP(0, world_size, args, vae_model, vae, record_function, tp, use_vae, num_codes, return_dict, mean=mean, std=std, noise_scale=noise_scale)
+            
+            real_embeddings, fake_embeddings, all_real_videos, all_fake_videos, zs = extract_results(return_dict, world_size)
+            return real_embeddings, fake_embeddings, all_real_videos, all_fake_videos, zs
+        
+        def cal_std(zs):
+            dims_to_reduce = [i for i in range(zs.dim()) if i != 1]
+            total_std = zs.std().item()
+            _mean = zs.mean(dim=dims_to_reduce)
+            _std = zs.std(dim=dims_to_reduce)
+            return total_std, _mean, _std
+        
+        real_embeddings, fake_embeddings, all_real_videos, all_fake_videos, zs = inference()
+        if args.noise_scale > 0:
+            total_std, _mean, _std = cal_std(zs)
+            real_embeddings, fake_embeddings, all_real_videos, all_fake_videos, zs = inference(mean=_mean, std=_std, noise_scale=args.noise_scale)
+        
+        if args.save_samples:
+            torch.save(zs.cpu(), args.save_samples)
+
+        if args.cal_norm:
+            total_std, _mean, _std = cal_std(zs)
+            print(f"{total_std = } {_mean = } {_std = }")
+            if args.save_prediction:
+                fname = os.path.join(args.save_dir, args.dataset_list[0], "gt_recon", "mean_std.pth")
+                torch.save({'_mean': _mean, '_std': _std}, fname)
+
+        result_str = video_eval(real_embeddings, fake_embeddings, all_real_videos, all_fake_videos)
+    else:
+        world_size = 1 if args.debug else torch.cuda.device_count()
+        manager = torch.multiprocessing.Manager()
+        return_dict = manager.dict()
+
+        if args.debug:
+            inference_eval(0, world_size, args, vae_model, vae, record_function, use_vae, num_codes, return_dict)
+        else:
+            spawn(inference_eval, args=(world_size, args, vae_model, vae, record_function, use_vae, num_codes, return_dict), nprocs=world_size, join=True)
+
+        pred_xs, pred_recs, lpips_alex, lpips_vgg, ssim_value, psnr_value, num_iter, total_usage, total_usage_bit, total_num_token, all_bit_indices_cat = [], [], 0, 0, 0, 0, 0, 0, 0, 0, []
+        for rank in range(world_size):
+            pred_xs.append(return_dict[rank]['pred_xs'])
+            pred_recs.append(return_dict[rank]['pred_recs'])
+            lpips_alex += return_dict[rank]['lpips_alex']
+            lpips_vgg += return_dict[rank]['lpips_vgg']
+            ssim_value += return_dict[rank]['ssim_value']
+            psnr_value += return_dict[rank]['psnr_value']
+            num_iter += return_dict[rank]['num_iter']
+            total_usage += return_dict[rank]['total_usage']
+        pred_xs = np.concatenate(pred_xs, 0)
+        pred_recs = np.concatenate(pred_recs, 0)
+
+        result_str = image_eval(pred_xs, pred_recs, lpips_alex, lpips_vgg, ssim_value, psnr_value, num_iter, total_usage, num_codes, total_usage_bit, total_num_token)
+        # result_str = inference_eval(args, vae_model, vae, record_function, use_vae, num_codes) 
+
+    print(f"noise scale = {args.noise_scale}")
+    print(result_str)
+    # save result_str to exp_dir
+    basename = os.path.basename(args.vqgan_ckpt)
+    match = re.search(r'model_step_(\d+)\.ckpt', basename)
+    iter_num = match.group(1) if match else None
+    data_prefix_manager.set_data_root(args.data_root, username=args.username)
+    ckpt_dir = os.path.dirname(data_prefix_manager(args.vqgan_ckpt))
+    use_frames = args.use_frames if args.use_frames else args.sequence_length
+    save_dir = os.path.join(ckpt_dir, "evaluation", args.dataset_list[0], f"{args.resolution[0][0]}_{args.resolution[0][1]}", f"{use_frames}")
+    os.makedirs(save_dir, exist_ok=True)
+    ema_suffix = "_ema" if args.ema == "yes" else ""
+    result_name = os.path.join(save_dir, f"result_{iter_num}{ema_suffix}.txt")
+    if (not args.save_prediction) and (args.noise_scale == 0):
+        with open(result_name, "w") as f:
+            f.write(result_str)
+    # print('Usage = %.2f'%((total_usage > 0.).sum() / num_codes))
+    if args.intermediate_tensor:
+        os.system(f"rm -rf {args.intermediate_tensor_dir}")
+
+def add_noise(z, mean, std, noise_scale):
+    if noise_scale > 0:
+        mean = mean.view(1, mean.shape[0], 1, 1, 1).to(z.device)
+        std = std.view(1, std.shape[0], 1, 1, 1).to(z.device)
+        z = (z - mean) / std
+        noise = torch.randn(z.size()).to(z.device)
+        z = (z + noise * noise_scale) * std + mean
+    return z
+
+def inference_DDP(rank, world_size, args, vae_model, vae, record_function, tp, use_vae, num_codes, return_dict, mean=None, std=None, noise_scale=0):
+    setup(rank, world_size)
+    # init data_prefix_manager
+    data_prefix_manager.set_data_root(args.data_root, username=args.username)
+
+    for param in vae.parameters():
+        param.requires_grad = False
+    vae = vae.eval()
+    vae = vae.to(f"cuda:{rank}")
+    # vae = torch.compile(vae)
+
+    save_dir = os.path.join(args.save_dir, args.dataset_list[0])
+    print('generating and saving video to %s...'%save_dir)
+    os.makedirs(save_dir, exist_ok=True)
+
+    data = VideoData(args)
+    loader = data.val_dataloader()
+
+    i3d = load_fvd_model(f"cuda:{rank}")
+
+    os.makedirs(os.path.join(save_dir, "gt"), exist_ok=True)
+    os.makedirs(os.path.join(save_dir, "recons"), exist_ok=True)
+
+    zs = []
+    real_embeddings = []
+    fake_embeddings = []
+
+    all_real_videos = []
+    all_fake_videos = []
+
+    num_videos = len(loader)
+    loader_iter = iter(loader)
+    progress_bar = tqdm(total=num_videos, desc=f"Testing {num_videos} batches")
+    for batch_idx in range(num_videos):
+        if args.turn_on_profiler and tp:
+            tp.step()
+        batch = next(loader_iter)
+        with torch.no_grad():
+            input_ = batch['video'] # B C T H W        
+            B = input_.shape[0]
+            if args.tokenizer in ["hbq_tokenizer"]:
+                input_ = input_.to(f"cuda:{rank}").to(torch.bfloat16)
+                with torch.amp.autocast("cuda", dtype=torch.bfloat16):
+                    x_raw, x_recons, z = vae(input_, 0, is_train=False)
+                batch['video'] = x_raw.to('cpu').to(torch.float32)
+                x_recons = x_recons.to(torch.float32)
+            else:
+                raise NotImplementedError
+                
+            if args.tokenizer in ["icvivit", "sd"]:
+                x_recons = rearrange(x_recons, "(b t) c h w -> b c t h w", b=B)
+
+            real_videos = torch.clamp(batch['video'] / 2 + 0.5, 0, 1)
+            if args.junke_old:
+                fake_videos = torch.clamp(x_recons.detach().cpu() + 0.5, 0, 1)
+            else:
+                fake_videos = torch.clamp(x_recons.detach().cpu() / 2 + 0.5, 0, 1)
+            
+            use_frames = args.use_frames if args.use_frames else args.sequence_length
+            if args.intermediate_tensor:
+                folder_name = os.path.join(args.intermediate_tensor_dir, f"{rank}_{batch_idx}")
+                os.makedirs(folder_name, exist_ok=True)
+                real_file = os.path.join(folder_name, "real_videos.pt")
+                fake_file = os.path.join(folder_name, "fake_videos.pt")
+                real_videos = real_videos[:,:,:use_frames,...]
+                fake_videos = fake_videos[:,:,:use_frames,...]
+                torch.save(real_videos.permute(0, 2, 1, 3, 4).squeeze(0), real_file)
+                torch.save(fake_videos.permute(0, 2, 1, 3, 4).squeeze(0), fake_file)
+                all_real_videos.append(real_file)
+                all_fake_videos.append(fake_file)
+            else:
+                real_videos = real_videos[:,:,:use_frames,...]
+                fake_videos = fake_videos[:,:,:use_frames,...]
+                all_real_videos.append(real_videos.clone())
+                all_fake_videos.append(fake_videos.clone())
+            if args.cal_norm or args.save_samples or args.noise_scale > 0:
+                zs.append(z)
+            real_embedding = get_fvd_logits(shift_dim(real_videos * 255, 1, -1).byte().data.numpy(), i3d=i3d, device=f"cuda:{rank}").cpu()
+            real_embeddings.append(real_embedding)
+            fake_embedding = get_fvd_logits(shift_dim(fake_videos * 255, 1, -1).byte().data.numpy(), i3d=i3d, device=f"cuda:{rank}").cpu()
+            fake_embeddings.append(fake_embedding)
+
+        if args.tokenizer in ['cvivit', "icvivit"] and not use_vae:
+            batch_codebook_usage = vq_output["batch_usage"]
+            total_usage += batch_codebook_usage
+
+        if args.save_prediction:
+            video = torch.cat([real_videos[:,:,:fake_videos.shape[2],:,:], fake_videos], dim=-1)
+            b, c, t, h, w = video.shape
+            video = video.permute(0, 2, 3, 4, 1).contiguous()
+            video = (video.squeeze().detach().cpu().numpy() * 255).astype('uint8')
+            os.makedirs(os.path.join(save_dir, "gt_recon"), exist_ok=True)
+            this_filename = batch["path"][0].split('/')[-1]
+            fname = os.path.join(save_dir, "gt_recon", this_filename)
+            import imageio
+            imageio.mimsave(fname, video, fps=15)
+        
+        if args.save_z:
+            os.makedirs(os.path.join(save_dir, "gt_recon"), exist_ok=True)
+            this_filename = batch["path"][0].split('/')[-1].split(".")[0]
+            fname = os.path.join(save_dir, "gt_recon", this_filename+".pt")
+            torch.save(z, fname)
+        
+        if args.save_frames:
+            
+            def convert_to_uint8(image):
+                return (image.detach().cpu().numpy() * 255).astype(np.uint8)
+
+            # artifact_grid_size = 32
+            assert real_videos.shape == fake_videos.shape, f"shape of gt and predicted videos are not equal"
+            assert real_videos.shape[0] == fake_videos.shape[0] == 1, f"batch size must be 1, real_videos {real_videos.shape[0]}, fake_videos {fake_videos.shape[0]}"
+            _real_videos = real_videos.squeeze(0)
+            _fake_videos = fake_videos.squeeze(0)
+            # h, w = real_videos.shape[-2:]
+            # assert (h % artifact_grid_size == 0) and (w % artifact_grid_size == 0), f"height and width of video must be divisible by {artifact_grid_size}"
+            
+            frame_num = _real_videos.shape[1]
+            for frame_idx in range(frame_num):
+                real_image = _real_videos[:,frame_idx,:,:]
+                fake_image = _fake_videos[:,frame_idx,:,:]
+            #     most_different_top_left, max_difference = find_most_different_patch(real_image, fake_image, artifact_grid_size)
+                real_image_uint8 = convert_to_uint8(real_image)
+                predicted_image_uint8 = convert_to_uint8(fake_image)
+
+                real_image_bgr = cv2.cvtColor(real_image_uint8.transpose(1, 2, 0), cv2.COLOR_RGB2BGR)
+                predicted_image_bgr = cv2.cvtColor(predicted_image_uint8.transpose(1, 2, 0), cv2.COLOR_RGB2BGR)
+
+                concatenated_image = np.concatenate((real_image_bgr, predicted_image_bgr), axis=1)
+                fname = os.path.join(save_dir, "gt_recon", f"{this_filename}_{frame_idx}.png")
+                cv2.imwrite(fname, concatenated_image)
+        progress_bar.update(1)
+
+    real_embeddings = torch.cat(real_embeddings, 0)
+    fake_embeddings = torch.cat(fake_embeddings, 0)
+    zs = torch.cat(zs, 0) if len(zs) > 0 else torch.tensor([])
+    if args.intermediate_tensor:
+        temp_dict = {
+            'real_embeddings':real_embeddings.cpu(),
+            'fake_embeddings':fake_embeddings.cpu(),
+            'all_real_videos':all_real_videos,
+            'all_fake_videos':all_fake_videos,
+            'zs': zs.cpu(),
+        }
+    else:
+        all_real_videos = torch.cat(all_real_videos, 0).permute(0, 2, 1, 3, 4)
+        all_fake_videos = torch.cat(all_fake_videos, 0).permute(0, 2, 1, 3, 4)
+        temp_dict = {
+            'real_embeddings':real_embeddings.cpu(),
+            'fake_embeddings':fake_embeddings.cpu(),
+            'all_real_videos':all_real_videos.cpu(),
+            'all_fake_videos':all_fake_videos.cpu(),
+            'zs': zs.cpu(),
+        }
+    # if dist.is_initialized():
+    #     dist.barrier()
+    return_dict[rank] = temp_dict
+    cleanup()
+
+            
+def video_eval(real_embeddings, fake_embeddings, all_real_videos, all_fake_videos):
+    fake_embeddings = fake_embeddings.to(torch.float64)
+    real_embeddings = real_embeddings.to(torch.float64)
+    FVD = frechet_distance(fake_embeddings, real_embeddings)
+    print(f"FVD: {FVD}")  # can't wait to see this number :)
+    del real_embeddings, fake_embeddings
+
+    lpips = calculate_lpips(all_real_videos, all_fake_videos, device="cuda")["value"].values()
+    psnr = calculate_psnr(all_real_videos, all_fake_videos)["value"].values()
+    ssim = calculate_ssim(all_real_videos, all_fake_videos)["value"].values()
+    lpips = np.mean(np.stack(list(lpips)))
+    ssim = np.mean(np.stack(list(ssim)))
+    psnr = np.mean(np.stack(list(psnr)))
+
+    result_str = f"""
+    FVD = {FVD:.4f}
+    LPIPS = {lpips:.4f}
+    SSIM = {ssim:.4f}
+    PSNR = {psnr:.3f}
+    """
+    return result_str
+
+def inference_eval(rank, world_size, args, vae_model, vae, record_function, use_vae, num_codes, return_dict):
+    # Don't remove this setup!!! dist.init_process_group is important for building loader (data.distributed.DistributedSampler)
+    setup(rank, world_size) 
+    # init data_prefix_manager
+    data_prefix_manager.set_data_root(args.data_root, username=args.username)
+
+    device = torch.device(f"cuda:{rank}")
+
+    for param in vae.parameters():
+        param.requires_grad = False
+    vae.to(device).eval()
+
+    save_dir = os.path.join(args.save_dir, args.dataset_list[0])
+    print('generating and saving video to %s...'%save_dir)
+    os.makedirs(save_dir, exist_ok=True)
+
+    data = VideoData(args)
+
+    loader = data.val_dataloader()
+
+    dims = 2048
+    block_idx = InceptionV3.BLOCK_INDEX_BY_DIM[dims]
+    inception_model = InceptionV3([block_idx]).to(device)
+    inception_model.eval()
+
+    loader_iter = iter(loader)
+
+    pred_xs = []
+    pred_recs = []
+    # LPIPS score related
+    loss_fn_alex = lpips.LPIPS(net='alex').to(device)  # best forward scores
+    loss_fn_vgg = lpips.LPIPS(net='vgg').to(device)   # closer to "traditional" perceptual loss, when used for optimization
+    lpips_alex = 0.0
+    lpips_vgg = 0.0
+
+    # SSIM score related
+    ssim_value = 0.0
+
+    # PSNR score related
+    psnr_value = 0.0
+    
+    num_images = len(loader)
+    print(f"Testing {num_images} files")
+    num_iter = 0
+
+    total_usage = 0.0
+    total_usage_bit = 0.0
+    total_num_token = 0
+    for batch_idx in tqdm(range(num_images)):
+        batch = next(loader_iter)
+            
+        with torch.no_grad():
+            x = batch['video']
+            if args.tokenizer in ["hbq_tokenizer"]:
+                x_raw, x_recons, z = vae(x.to(device), 0, is_train=False)
+                x_recons = x_recons.squeeze(-3).cpu()
+            else:
+                raise NotImplementedError
+
+            if args.image_recon4video:
+                # convert back to image format
+                x = x.squeeze(2)
+                x_recons = x_recons.squeeze(2)
+        
+        if args.tokenizer in ["cvivit", "icvivit"] and not use_vae:
+        
+            # encoding_indices = vq_output["encodings"].detach().cpu()
+            code_counts = calculate_batch_codebook_usage_percentage(vq_output["encodings"], num_codes)
+            total_counts += code_counts
+
+            batch_codebook_usage = vq_output["batch_usage"]
+            total_usage += batch_codebook_usage
+        
+        paths = batch["path"]
+        assert len(paths) == x.shape[0]
+
+        for p, input_ori, recon_ori in zip(paths, x, x_recons):
+            if os.path.isabs(p):
+                p = "/".join(p.split("/")[6:])
+            assert not os.path.isabs(p), f"{p} should not be abspath"
+            path = os.path.join(save_dir, "input_recon", os.path.basename(p))
+            os.makedirs(os.path.split(path)[0], exist_ok=True)
+            
+            input_ori = input_ori.unsqueeze(0).to(device)
+            input_ = (input_ori + 1) / 2 # [0, 1]
+            
+            pred_x = inception_model(input_)[0]
+            pred_x = pred_x.squeeze(3).squeeze(2).cpu().numpy()
+
+            recon_ori = recon_ori.unsqueeze(0).to(device)
+            recon_ = (recon_ori + 1) / 2 # [0, 1]
+            # recon_ = recon_.permute(1, 2, 0).detach().cpu()
+            with torch.no_grad():
+                pred_rec = inception_model(recon_)[0]
+            pred_rec = pred_rec.squeeze(3).squeeze(2).cpu().numpy()
+            if args.save_prediction:
+                if input_.dim() == 4:
+                    input_image = input_.squeeze(0)
+                if recon_.dim() == 4:
+                    recon_image = recon_.squeeze(0)
+                input_recon = torch.cat([input_image, recon_image], dim=-1)
+                input_recon = Image.fromarray((torch.clamp(input_recon.permute(1, 2, 0).detach().cpu(), 0, 1).numpy() * 255).astype(np.uint8))
+                input_recon.save(path)
+
+            pred_xs.append(pred_x)
+            pred_recs.append(pred_rec)
+
+            # calculate lpips
+            with torch.no_grad():
+                lpips_alex += loss_fn_alex(input_ori, recon_ori).sum() # [-1, 1]
+                lpips_vgg += loss_fn_vgg(input_ori, recon_ori).sum() # [-1, 1]
+
+            #calculate PSNR and SSIM
+            rgb_restored = (recon_ * 255.0).permute(0, 2, 3, 1).to("cpu", dtype=torch.uint8).numpy()
+            rgb_gt = (input_ * 255.0).permute(0, 2, 3, 1).to("cpu", dtype=torch.uint8).numpy()
+            rgb_restored = rgb_restored.astype(np.float32) / 255.
+            rgb_gt = rgb_gt.astype(np.float32) / 255.
+            ssim_temp = 0
+            psnr_temp = 0
+            B, _, _, _ = rgb_restored.shape
+            for i in range(B):
+                rgb_restored_s, rgb_gt_s = rgb_restored[i], rgb_gt[i]
+                with torch.no_grad():
+                    ssim_temp += ssim_loss(rgb_restored_s, rgb_gt_s, data_range=1.0, channel_axis=-1)
+                    psnr_temp += psnr_loss(rgb_gt, rgb_restored)
+            ssim_value += ssim_temp / B
+            psnr_value += psnr_temp / B
+            num_iter += 1
+        
+    pred_xs = np.concatenate(pred_xs, axis=0)
+    pred_recs = np.concatenate(pred_recs, axis=0)
+    temp_dict = {
+        'pred_xs':pred_xs,
+        'pred_recs':pred_recs,
+        'lpips_alex':lpips_alex.cpu(),
+        'lpips_vgg':lpips_vgg.cpu(),
+        'ssim_value': ssim_value,
+        'psnr_value': psnr_value,
+        'num_iter': num_iter,
+        'total_usage': total_usage,
+        'total_usage_bit': total_usage_bit,
+        'total_num_token': total_num_token,
+    }
+    return_dict[rank] = temp_dict
+
+    # if dist.is_initialized():
+    #     dist.barrier()
+    cleanup()
+
+def image_eval(pred_xs, pred_recs, lpips_alex, lpips_vgg, ssim_value, psnr_value, num_iter, total_usage, num_codes, total_usage_bit, total_num_token):
+    mu_x = np.mean(pred_xs, axis=0)
+    sigma_x = np.cov(pred_xs, rowvar=False)
+    mu_rec = np.mean(pred_recs, axis=0)
+    sigma_rec = np.cov(pred_recs, rowvar=False)
+    
+    fid_value = calculate_frechet_distance(mu_x, sigma_x, mu_rec, sigma_rec)
+    lpips_alex_value = lpips_alex / num_iter
+    lpips_vgg_value = lpips_vgg / num_iter
+    ssim_value = ssim_value / num_iter
+    psnr_value = psnr_value / num_iter
+
+    result_str = f"""
+    FID = {fid_value:.4f}
+    LPIPS_VGG: {lpips_vgg_value.item():.4f}
+    LPIPS_ALEX: {lpips_alex_value.item():.4f}
+    SSIM: {ssim_value:.4f}
+    PSNR: {psnr_value:.3f}
+    """
+    return result_str
+if __name__ == '__main__':
+    main()

grn/tokenizer/train.py

@@ -0,0 +1,561 @@
+from json import load
+import os
+import argparse
+import math
+import glob
+import time
+import logging
+from distutils.util import strtobool
+from copy import deepcopy
+import gc
+gc.disable()
+import os.path as osp
+
+import torch
+import torch.nn.functional as F
+import torch.optim as optim
+import torch.distributed as dist
+from torch.profiler import record_function as torch_record_function
+from contextlib import nullcontext
+from torch.nn.parallel import DistributedDataParallel as DDP
+from safetensors.torch import load_file
+from torch.distributed.fsdp import FullyShardedDataParallel as FSDP
+from torch.distributed.fsdp import StateDictType, FullStateDictConfig
+
+from videovae.utils.misc import data_prefix_manager, COLOR_BLUE, COLOR_RESET, is_torch_optim_sch
+from videovae.utils.distributed import init_distributed_mode, reduce_losses, average_losses, _FSDP
+from videovae.utils.ema import update_ema, requires_grad
+
+from videovae.models.discriminator import ImageDiscriminator, VideoDiscriminator
+from videovae.data import VideoData
+from videovae.modules import build_lpips_model
+from videovae.modules.loss import get_disc_loss, adopt_weight
+from videovae.utils.misc import get_last_ckpt, seed_everything, print_gpu_usage, print_model_summary, version_checker
+from videovae.utils.init_models import init_vae_only, init_vit_from_image, resume_from_ckpt, init_cnn_from_image, load_cnn
+from videovae.utils.nan_detector import NanDetector
+from videovae.utils.arguments import MainArgs, add_model_specific_args, init_args, format_args
+from videovae.utils.scheduler import get_lambda
+from videovae.utils.mfu import register_mfu_hook, get_mfu, get_tflops, get_tflops_dict
+
+from videovae.utils.context_parallel import ContextParallelUtils as cp
+
+def save_model(fsdp_model, rank, model_path, global_step):
+    # FSDP推荐用 state_dict_type=FULL_STATE_DICT 来保存
+    with FSDP.state_dict_type(fsdp_model, StateDictType.FULL_STATE_DICT, FullStateDictConfig(offload_to_cpu=True, rank0_only=True)):
+        state_dict = fsdp_model.state_dict()
+        os.makedirs(os.path.dirname(model_path), exist_ok=True)
+        torch.save({'vae': state_dict, 'step': global_step}, model_path)
+        print(f"模型已保存到 {model_path}")
+
+# enable_timeline_sdk = strtobool(os.getenv("GenAI_USE_TIMELINE_SDK", "0"))
+enable_timeline_sdk = False
+
+def split_to_ranks(x):
+    bs = x.shape[0]
+    cp_size = cp.get_cp_size()
+    if cp_size > 1 and bs % cp_size == 0:
+        cp_rank = cp.get_cp_rank()
+        return x.chunk(cp_size, dim=0)[cp_rank]
+    else:
+        return x
+
+if enable_timeline_sdk:
+    try:
+        import bytedance.ndtimeline as ndtimeline
+    except ImportError:
+        print(f"import vescale.ndtimeline failed, skipped")
+        enable_timeline_sdk = False
+
+def init_data_scheduler(video_ranks_ratio: float = -1.0, cp_size: int = 1):
+    if video_ranks_ratio < 0:
+        return None,None
+
+    cp_size = max(1, cp_size)
+
+    rank = torch.distributed.get_rank()
+    world_size = torch.distributed.get_world_size()
+
+    video_ranks = list(range(int((world_size * video_ranks_ratio) // cp_size) * cp_size)) # align to cp_size for video
+    image_ranks = list(range(len(video_ranks), world_size))
+
+    print(f"[info] video_ranks: {video_ranks}, image_ranks: {image_ranks}")
+
+    if rank in image_ranks:
+        group = torch.distributed.new_group(image_ranks)
+        dataset_type_on_this_rank = "image"
+    else:
+        group = torch.distributed.new_group(video_ranks)
+        dataset_type_on_this_rank = "video"
+
+    return group, dataset_type_on_this_rank
+
+def main():
+    parser = argparse.ArgumentParser()
+    parser = MainArgs.add_main_args(parser)
+    parser = VideoData.add_data_specific_args(parser)
+    args, unknown = parser.parse_known_args()
+    args, parser, vae_model = add_model_specific_args(args, parser)
+    args = parser.parse_args()
+
+    args = init_args(args) # post process args
+
+    # init data_prefix_manager
+    data_prefix_manager.set_data_root(args.data_root, username=args.username)
+    args.default_root_dir = data_prefix_manager(args.default_root_dir)
+
+    # Setup DDP:
+    init_distributed_mode(args)
+    rank = dist.get_rank()
+    world_size = dist.get_world_size()
+    device = rank % torch.cuda.device_count()
+    seed_everything(args.seed)
+    torch.cuda.set_device(device)
+
+    # init context parallel
+    cp_cfg = {"cp_size": args.context_parallel_size}
+    cp.initialize_context_parallel(cp_cfg)
+
+    ds_group, ds_type = init_data_scheduler(args.video_ranks_ratio, cp_size = args.context_parallel_size)
+
+    # Setup an experiment folder:
+    checkpoint_dir = f"{args.default_root_dir}/checkpoints"  # Stores saved model checkpoints
+    os.makedirs(checkpoint_dir, exist_ok=True)
+    if rank == 0:
+        script_str = format_args(args)
+        with open(os.path.join(args.default_root_dir, "script.sh"), "w") as f:
+            f.write(script_str)
+        print(f"{COLOR_BLUE}Experiment directory created at {args.default_root_dir}{COLOR_RESET}")
+
+        import wandb
+        wandb_project = "HBQ_Tokenizer"
+        wandb.init(
+            project=wandb_project,
+            name=os.path.basename(os.path.normpath(args.default_root_dir)),
+            dir=args.default_root_dir,
+            config=args,
+            mode="offline" if args.debug else "online"
+        )
+    
+    # init model    
+    vae = vae_model(args).to(device)
+    if rank == 0:
+        model_arch_save_path = os.path.join(args.default_root_dir, "model_arch.txt")
+        print(f"{COLOR_BLUE}Logging model architecture at {model_arch_save_path}{COLOR_RESET}")
+        with open(model_arch_save_path, "w") as f:
+            f.write(str(vae))
+    image_disc = ImageDiscriminator(args).to(device)
+    video_disc = VideoDiscriminator(args).to(device)
+
+    # init optimizers and schedulers
+    if args.optim_type == "Adam":
+        vae_optim = torch.optim.Adam
+    elif args.optim_type == "AdamW":
+        vae_optim = torch.optim.AdamW
+    if args.disc_optim_type is None:
+        disc_optim = vae_optim
+    elif args.disc_optim_type == "rmsprop":
+        disc_optim = torch.optim.RMSprop
+
+    def get_param_groups(model):
+        decay = []
+        no_decay = []
+        for name, param in model.named_parameters():
+            if param.requires_grad:
+                if len(param.shape) == 1 or name.endswith(".bias") or ('scale_learnable_parameters' in name):
+                    no_decay.append(param)
+                    print(f'disable weight deacy for {name}')
+                else:
+                    decay.append(param)
+        optimizer_grouped_parameters = [
+            {'params': decay, 'weight_decay': 0.01},
+            {'params': no_decay, 'weight_decay': 0.0}
+        ]
+        return optimizer_grouped_parameters
+
+    opt_vae = vae_optim(get_param_groups(vae), lr=args.lr, betas=(args.beta1, args.beta2))
+    if disc_optim == torch.optim.RMSprop:
+        opt_image_disc = disc_optim(image_disc.parameters(), lr=args.lr * args.dis_lr_multiplier)
+        opt_video_disc = disc_optim(video_disc.parameters(), lr=args.lr * args.dis_lr_multiplier)
+    else:
+        opt_image_disc = disc_optim(image_disc.parameters(), lr=args.lr * args.dis_lr_multiplier, betas=(args.beta1, args.beta2))
+        opt_video_disc = disc_optim(video_disc.parameters(), lr=args.lr * args.dis_lr_multiplier, betas=(args.beta1, args.beta2))
+
+    if args.scheduler == "no":
+        sch_vae, sch_image_disc, sch_video_disc = None, None, None
+    else:
+        lr_lambda = get_lambda(args)
+        sch_vae = optim.lr_scheduler.LambdaLR(opt_vae, lr_lambda)
+        sch_image_disc = optim.lr_scheduler.LambdaLR(opt_image_disc, lr_lambda)
+        sch_video_disc = optim.lr_scheduler.LambdaLR(opt_video_disc, lr_lambda)
+
+    ### ema
+    ema = None
+    if args.ema == "yes":
+        ema = deepcopy(vae).to(device)  # Create an EMA of the model for use after training
+        requires_grad(ema, False)
+        print(f"EMA Parameters: {sum(p.numel() for p in ema.parameters()):,}")
+        update_ema(ema, vae, decay=0)  # Ensure EMA is initialized with synced weights
+        ema.eval()  # EMA model should always be in eval mode
+
+    model_optims = {
+        "vae" : vae,
+        "image_disc" : image_disc,
+        "video_disc" : video_disc,
+        "opt_vae" : opt_vae,
+        "opt_image_disc" : opt_image_disc,
+        "opt_video_disc" : opt_video_disc,
+        "sch_vae" : sch_vae,
+        "sch_image_disc" : sch_image_disc,
+        "sch_video_disc" : sch_video_disc,
+        "ema": ema, 
+    }
+
+    ### Resume from checkpoint in default_root_dir or load pretrained weights if specified
+    ckpt_path = None
+    assert not args.default_root_dir is None # required argument
+    ckpt_path = get_last_ckpt(args.default_root_dir)
+    init_step = 0
+    if ckpt_path:
+        print(f"Resuming from {ckpt_path}")
+        state_dict = torch.load(ckpt_path, map_location="cpu")
+        model_optims, init_step = resume_from_ckpt(state_dict, model_optims, load_optims=args.zero<=0, remove_disc=args.remove_disc, ckpt_path=ckpt_path, args=args)
+    elif args.pretrained is not None:
+        args.pretrained = data_prefix_manager(args.pretrained)
+        # read weight
+        state_dict = torch.load(args.pretrained, map_location="cpu", weights_only=True)
+        
+        if args.pretrained_ema == "yes":
+            state_dict["vae"] = state_dict["ema"] # replace vae weights with ema weight
+        # load model
+        if args.pretrained_mode == "weights":
+            model_optims, _ = resume_from_ckpt(state_dict, model_optims, load_optims=False, remove_disc=args.remove_disc, remove_enlarge_factors=args.remove_enlarge_factors, args=args) # load all models and optims
+            del state_dict
+        else:
+            raise NotImplementedError
+        print(f"Successfully loaded ckpt {args.pretrained}, pretrained_mode {args.pretrained_mode}")
+
+    # init dataloader
+    data = VideoData(args, ds_group = ds_group, ds_type = ds_type)
+    dataloaders = data.train_dataloader()
+    dataloader_iters = [iter(loader) for loader in dataloaders]
+    ### init epoch in resuming
+    dataloader_init_epoch = (
+        init_step if init_step > 0 # in case of resuming
+        else args.dataloader_init_epoch if args.dataloader_init_epoch > 0 # in case of fintuning
+        else 0
+    )
+    data_epochs = [dataloader_init_epoch for _ in dataloaders]
+    for idx in range(len(dataloaders)):
+        print(f"Reset the {idx}th dataloader as epoch {data_epochs[idx]}")
+        if hasattr(dataloaders[idx], "sampler"):
+            dataloaders[idx].sampler.set_epoch(data_epochs[idx])
+        else:
+            raise NotImplementedError
+
+    ### torch.compile after loading all weights
+    print_model_summary([vae, image_disc, video_disc])
+    if args.zero > 0:
+        from torch.distributed.fsdp import (
+            FullyShardedDataParallel as FSDP,
+            ShardingStrategy,
+            MixedPrecision,
+        )
+        def my_policy(
+            module: torch.nn.Module,
+            recurse: bool,
+            **kwargs,
+        ) -> bool:
+            return True
+        auto_wrap_policy = my_policy
+        vae = FSDP(
+            vae, 
+            device_id=device,
+            sharding_strategy=ShardingStrategy.FULL_SHARD, 
+            mixed_precision=None,
+            auto_wrap_policy=auto_wrap_policy, 
+            use_orig_params=True, 
+            sync_module_states=True, 
+            limit_all_gathers=True,
+            device_mesh=None,
+        ).to(device)
+        # vae = _FSDP(vae, device, args.zero)
+    else:
+        vae = DDP(vae.to(device), device_ids=[args.gpu], bucket_cap_mb=args.bucket_cap_mb, find_unused_parameters=True)
+
+    image_disc = DDP(image_disc.to(device), device_ids=[args.gpu], bucket_cap_mb=args.bucket_cap_mb)
+    video_disc = DDP(video_disc.to(device), device_ids=[args.gpu], bucket_cap_mb=args.bucket_cap_mb)
+
+    image_perceptual_model, video_perceptual_model = build_lpips_model(args)
+    image_perceptual_model = image_perceptual_model.to(device)
+    video_perceptual_model = video_perceptual_model.to(device)
+
+    if args.compile == "yes":
+        if args.vf_weight > 0 and args.vf_weight_approx < 0:
+            torch._functorch.config.donated_buffer = False # This backward function was compiled with non-empty donated buffers which requires create_graph=False and retain_graph=False
+        torch._dynamo.config.cache_size_limit = 256
+        torch._dynamo.config.accumulated_cache_size_limit = 4096
+        torch._dynamo.config.automatic_dynamic_shapes = True
+        torch._dynamo.config.suppress_errors = False
+        torch._dynamo.config.optimize_ddp = False if args.use_checkpoint else True
+
+        for k in model_optims:
+            if k != "ema" and model_optims[k] and not is_torch_optim_sch(model_optims[k]):
+                print(f"compiling model {k}")
+                if k == "vae":
+                    model_optims[k].encoder.compile()#options={'fx_graph_cache':True})
+                    model_optims[k].decoder.compile()#options={'fx_graph_cache':True})
+                else:
+                    model_optims[k].compile()#options={'fx_graph_cache':True})
+        print(f"Successfully compiled all models")
+
+    disc_loss = get_disc_loss(args.disc_loss_type)
+    
+    if enable_timeline_sdk:
+        version_checker("2.0.0", "3.0.0")
+        bmq_cluster = os.getenv('CUDA_TIMER_STREAM_KAFKA_CLUSTER', 'bmq_bigbang_3rd')
+        bmq_topic = os.getenv('CUDA_TIMER_STREAM_KAFKA_TOPIC', 'megatron_cuda_timer_tracing_original')
+        ndtimeline.init_ndtimers(
+            mode="fsdp",
+            mesh_shape=(world_size,),
+            world_size=world_size,
+            enable_streamer=True,
+            post_handlers=[ndtimeline.handlers.MQNDHandler(mq_sinks=[ndtimeline.handlers.format_mq_sink(bmq_cluster, bmq_topic)])],
+        )
+        ndtimeline.set_global_step(init_step)
+        print(f"init timeline successfully")
+
+    # init profiler
+    def trace_handler(p):
+        p.export_chrome_trace(os.path.join(args.default_root_dir, f"trace_step_{p.step_num}_rank_{rank}.json.gz"))
+
+    if args.turn_on_profiler:
+        print(f"start to init profiler")
+        tp = torch.profiler.profile(
+            activities=[
+                torch.profiler.ProfilerActivity.CPU,
+                torch.profiler.ProfilerActivity.CUDA,
+            ],
+            schedule=torch.profiler.schedule(
+                wait=args.profiler_scheduler_wait_steps,
+                warmup=3,
+                active=2,
+                repeat=1,
+            ),
+            with_stack=True,
+            record_shapes=True,
+            profile_memory=True,
+            on_trace_ready=trace_handler
+        )
+        tp.start()
+        record_function = torch_record_function
+        print(f"finish to init profiler")
+    else:
+        record_function = nullcontext
+
+    start_time = time.time()
+    cnt = 0
+    debug_root_dir = data_prefix_manager(f"debug/{cnt}")
+    while os.path.exists(debug_root_dir):
+        cnt += 1
+        debug_root_dir = data_prefix_manager(f"debug/{cnt}")
+    os.makedirs(debug_root_dir, exist_ok=True)
+    if cp.get_cp_rank() == 0:
+        cp_group = cp.get_cp_group_rank()
+        debug_f = open(os.path.join(debug_root_dir, f"rank_{rank}_cp_group_{cp_group}.txt"), "w")
+    else:
+        debug_f = None
+    for global_step in range(init_step, args.max_steps):
+        # if args.turn_on_profiler and tp:
+        #     tp.step()
+        loss_dicts = []
+        
+        if global_step == args.discriminator_iter_start - args.disc_pretrain_iter:
+            logging.info(f"discriminator begins pretraining ")
+        if global_step == args.discriminator_iter_start:
+            log_str = "add GAN loss into training"
+            if args.disc_pretrain_iter > 0:
+                log_str += ", discriminator ends pretraining"
+            logging.info(log_str)
+        
+        for idx in range(len(dataloader_iters)):
+            try:
+                _batch = next(dataloader_iters[idx])
+            except StopIteration:
+                data_epochs[idx] += 1
+                print(f"Reset the {idx}th dataloader as epoch {data_epochs[idx]}")
+                dataloaders[idx].sampler.set_epoch(data_epochs[idx])
+                dataloader_iters[idx] = iter(dataloaders[idx]) # update dataloader iter
+                _batch = next(dataloader_iters[idx])
+            except Exception as e:
+                raise e 
+            x = _batch["video"]
+
+            _type = _batch["type"][0]
+
+            if _type == "image" and ds_type is None:
+                x = split_to_ranks(x)
+
+            disc_factor = 1.
+            with NanDetector(vae) if args.enable_nan_detector else nullcontext():
+                with record_function("vae"):
+                    if _type == "image":
+                        x, x_recon, flat_frames, flat_frames_recon, vae_loss_dict, vae_log_dict = vae(x, disc_factor, image_disc=image_disc, image_perceptual_model=image_perceptual_model)
+                    elif _type == "video":
+                        if debug_f is not None:
+                            debug_f.write(f'step {idx}, {_batch["path"]}\n')
+                        x, x_recon, flat_frames, flat_frames_recon, vae_loss_dict, vae_log_dict = vae(
+                            x, disc_factor, 
+                            image_disc=image_disc, video_disc=video_disc, 
+                            image_perceptual_model=image_perceptual_model, 
+                            video_perceptual_model=video_perceptual_model,
+                        )
+                    g_loss = sum(vae_loss_dict.values())
+                    opt_vae.zero_grad()
+                    g_loss.backward()
+                    # print_gpu_usage("vae")
+                    if args.max_grad_norm > 0:
+                        torch.nn.utils.clip_grad_norm_(vae.parameters(), args.max_grad_norm)
+                    
+                    # from pnp.utils import detect_anomalous_params
+                    # detect_anomalous_params(g_loss, vae)
+
+                    opt_vae.step()
+                    opt_vae.zero_grad() # free memory
+                    if args.ema == "yes":
+                        update_ema(ema, vae.module)
+                        
+                with record_function("disc"):
+                    disc_loss_dict = {}
+                    # args.discriminator_iter_start=-1 args.disc_pretrain_iter=0
+                    discloss = d_image_loss = d_video_loss = torch.tensor(0.).to(x.device)
+                    ### enable pool warmup
+                    for disc_step in range(args.disc_optim_steps):
+                        require_optim = False
+                        if _type == "image":
+                            if args.image_disc_weight > 0:
+                                require_optim = True
+                                logits_image_real = image_disc(x, pool_name="real")
+                                logits_image_fake = image_disc(x_recon.detach(), pool_name="fake")
+                                d_image_loss = disc_loss(logits_image_real, logits_image_fake)
+                                discloss = d_image_loss * args.image_disc_weight
+                                disc_loss_dict["train/logits_image_real"] = logits_image_real.mean().detach()
+                                disc_loss_dict["train/logits_image_fake"] = logits_image_fake.mean().detach()
+                                disc_loss_dict["train/d_image_loss"] = discloss.detach()
+                            opt_discs, sch_discs = [opt_image_disc], [sch_image_disc]
+                        elif _type == "video":
+                            if args.image_disc_weight > 0 and args.gan_image4video == "yes":
+                                require_optim = True
+                                logits_image_real = image_disc(flat_frames.detach(), pool_name="real")
+                                logits_image_fake = image_disc(flat_frames_recon.detach(), pool_name="fake")
+                                d_image_loss = disc_loss(logits_image_real, logits_image_fake)
+                                disc_loss_dict["train/logits_image_real"] = logits_image_real.mean().detach()
+                                disc_loss_dict["train/logits_image_fake"] = logits_image_fake.mean().detach()
+                                disc_loss_dict["train/d_image_loss"] = (d_image_loss * args.image_disc_weight).detach()
+                            if args.video_disc_weight > 0:
+                                require_optim = True
+                                logits_video_real = video_disc(x.detach(), pool_name="real")
+                                logits_video_fake = video_disc(x_recon.detach(), pool_name="fake")
+                                d_video_loss = disc_loss(logits_video_real, logits_video_fake)
+                                disc_loss_dict["train/logits_video_real"] = logits_video_real.mean().detach()
+                                disc_loss_dict["train/logits_video_fake"] = logits_video_fake.mean().detach()
+                                disc_loss_dict["train/d_video_loss"] = (d_video_loss * args.video_disc_weight).detach()
+                            discloss = d_image_loss * args.image_disc_weight + d_video_loss * args.video_disc_weight
+                            opt_discs, sch_discs = [opt_image_disc, opt_video_disc], [sch_image_disc, sch_video_disc]
+                        discloss = disc_factor * discloss
+
+                        if require_optim:
+                            for opt_disc in opt_discs:
+                                opt_disc.zero_grad()
+                            discloss.backward()
+                            # print_gpu_usage("disc")
+                            if args.max_grad_norm_disc > 0:
+                                torch.nn.utils.clip_grad_norm_(image_disc.parameters(), args.max_grad_norm_disc)
+                                torch.nn.utils.clip_grad_norm_(video_disc.parameters(), args.max_grad_norm_disc)
+
+                            for opt_disc in opt_discs:
+                                opt_disc.step()
+                            for opt_disc in opt_discs:
+                                opt_disc.zero_grad() # free memory
+
+            with record_function("loss"):
+                loss_dict = {**vae_loss_dict, **disc_loss_dict, **vae_log_dict}
+                if (global_step+1) % args.log_every == 0:
+                    reduced_loss_dict = reduce_losses(loss_dict)
+                else:
+                    reduced_loss_dict = {}
+                loss_dicts.append(reduced_loss_dict)
+        
+        # update scheduler   
+        if not sch_vae is None:
+            sch_vae.step()
+        for sch_disc in sch_discs:
+            if not sch_disc is None:
+                sch_disc.step()
+
+        # if enable_timeline_sdk:
+            # ndtimeline.inc_step()
+
+        if (global_step+1) % args.log_every == 0:
+            avg_loss_dict = average_losses(loss_dicts)
+            torch.cuda.synchronize()
+            end_time = time.time()
+            iter_speed = (end_time - start_time) / args.log_every
+
+            if args.mfu_logging == "yes":
+                tflops = get_tflops() / args.log_every
+                tflops_log_str = f"tflops={tflops:.1f}, "
+                tflops_dict = get_tflops_dict(args.log_every)
+                tflops_dict_log_str = f"tflops_Dict={tflops_dict}, "
+                mfu = get_mfu(iter_speed) / args.log_every
+                mfu_log_str = f"mfu={mfu:.3f}, "
+            else:
+                tflops_log_str = ""
+                tflops_dict_log_str = ""
+                mfu_log_str = ""
+
+            if rank == 0:
+                avg_loss_dict["lr"] = opt_vae.param_groups[0]['lr']
+                for key, value in avg_loss_dict.items():
+                    wandb.log({key: value}, step=global_step)
+                
+                recons_loss_sum, video_perceptual_loss_sum = 0., 0.
+                for key in avg_loss_dict:
+                    if 'recon_loss' in key:
+                        recons_loss_sum += avg_loss_dict[key]
+                    if 'video_perceptual_loss' in key:
+                        video_perceptual_loss_sum += avg_loss_dict[key]
+                
+                print(f'global_step={global_step}, recon_loss={recons_loss_sum:.4f}, ' \
+                            f'video_perceptual_loss={video_perceptual_loss_sum:.4f}, ' \
+                            f'iter_speed={iter_speed:.2f}s, ' \
+                            f'{mfu_log_str}' \
+                            f'{tflops_log_str}' \
+                            f'{tflops_dict_log_str}' \
+                            )
+            start_time = time.time()
+            if enable_timeline_sdk:
+                ndtimeline.flush()
+        
+        if (global_step+1) % args.ckpt_every == 0 and global_step != init_step:
+            checkpoint_path = os.path.join(checkpoint_dir, f'model_step_{global_step}.ckpt')
+            if args.zero > 0:
+                save_model(vae, rank, checkpoint_path, global_step)
+            else:
+                if rank == 0:
+                    save_dict = {}
+                    for k in model_optims:
+                        model = model_optims[k]
+                        save_dict[k] = None if model is None \
+                            else model.module.state_dict() if hasattr(model, "module") \
+                            else model.state_dict()
+                    torch.save({
+                        'step': global_step,
+                        **save_dict,
+                    }, checkpoint_path)
+                    print(f'Checkpoint saved at step {global_step}')
+        
+        if (global_step+1) % args.manual_gc_interval == 0:
+            gc.collect()
+    
+if __name__ == '__main__':
+    main()

grn/tokenizer/videovae/evaluation/__init__.py

@@ -0,0 +1,8 @@
+from .common_metrics_on_video_quality.calculate_fvd import calculate_fvd
+from .common_metrics_on_video_quality.calculate_lpips import calculate_lpips
+from .common_metrics_on_video_quality.calculate_psnr import calculate_psnr
+from .common_metrics_on_video_quality.calculate_ssim import calculate_ssim
+
+from .fvd import get_fvd_logits, frechet_distance, load_fvd_model
+from .fid import calculate_frechet_distance
+from .inception import InceptionV3

grn/tokenizer/videovae/evaluation/common_metrics_on_video_quality/.gitignore

@@ -0,0 +1 @@
+__pycache__

grn/tokenizer/videovae/evaluation/common_metrics_on_video_quality/calculate_fvd.py

@@ -0,0 +1,85 @@
+import numpy as np
+import torch
+from tqdm import tqdm
+
+def trans(x):
+    # if greyscale images add channel
+    if x.shape[-3] == 1:
+        x = x.repeat(1, 1, 3, 1, 1)
+
+    # permute BTCHW -> BCTHW
+    x = x.permute(0, 2, 1, 3, 4) 
+
+    return x
+
+def calculate_fvd(videos1, videos2, device, method='styleganv'):
+
+    if method == 'styleganv':
+        from .fvd.styleganv.fvd import get_fvd_feats, frechet_distance, load_i3d_pretrained
+    elif method == 'videogpt':
+        from .fvd.videogpt.fvd import load_i3d_pretrained
+        from .fvd.videogpt.fvd import get_fvd_logits as get_fvd_feats
+        from .fvd.videogpt.fvd import frechet_distance
+
+    print("calculate_fvd...")
+
+    # videos [batch_size, timestamps, channel, h, w]
+    
+    assert videos1.shape == videos2.shape
+
+    i3d = load_i3d_pretrained(device=device)
+    fvd_results = []
+
+    # support grayscale input, if grayscale -> channel*3
+    # BTCHW -> BCTHW
+    # videos -> [batch_size, channel, timestamps, h, w]
+
+    videos1 = trans(videos1)
+    videos2 = trans(videos2)
+
+    fvd_results = {}
+
+    # for calculate FVD, each clip_timestamp must >= 10
+    for clip_timestamp in tqdm(range(10, videos1.shape[-3]+1)):
+       
+        # get a video clip
+        # videos_clip [batch_size, channel, timestamps[:clip], h, w]
+        videos_clip1 = videos1[:, :, : clip_timestamp]
+        videos_clip2 = videos2[:, :, : clip_timestamp]
+
+        # get FVD features
+        feats1 = get_fvd_feats(videos_clip1, i3d=i3d, device=device)
+        feats2 = get_fvd_feats(videos_clip2, i3d=i3d, device=device)
+      
+        # calculate FVD when timestamps[:clip]
+        fvd_results[clip_timestamp] = frechet_distance(feats1, feats2)
+
+    result = {
+        "value": fvd_results,
+        "video_setting": videos1.shape,
+        "video_setting_name": "batch_size, channel, time, heigth, width",
+    }
+
+    return result
+
+# test code / using example
+
+def main():
+    NUMBER_OF_VIDEOS = 8
+    VIDEO_LENGTH = 50
+    CHANNEL = 3
+    SIZE = 64
+    videos1 = torch.zeros(NUMBER_OF_VIDEOS, VIDEO_LENGTH, CHANNEL, SIZE, SIZE, requires_grad=False)
+    videos2 = torch.ones(NUMBER_OF_VIDEOS, VIDEO_LENGTH, CHANNEL, SIZE, SIZE, requires_grad=False)
+    device = torch.device("cuda")
+    # device = torch.device("cpu")
+
+    import json
+    result = calculate_fvd(videos1, videos2, device, method='videogpt')
+    print(json.dumps(result, indent=4))
+
+    result = calculate_fvd(videos1, videos2, device, method='styleganv')
+    print(json.dumps(result, indent=4))
+
+if __name__ == "__main__":
+    main()

grn/tokenizer/videovae/evaluation/common_metrics_on_video_quality/calculate_lpips.py

@@ -0,0 +1,72 @@
+import numpy as np
+import torch
+from tqdm import tqdm
+import math
+
+import torch
+import lpips
+from .utils import build_dataloader
+
+spatial = True         # Return a spatial map of perceptual distance.
+
+# Linearly calibrated models (LPIPS)
+loss_fn = lpips.LPIPS(net='vgg', spatial=spatial) # Can also set net = 'squeeze' or 'vgg'
+# loss_fn = lpips.LPIPS(net='alex', spatial=spatial, lpips=False) # Can also set net = 'squeeze' or 'vgg'
+
+def calculate_lpips(videos1, videos2, device):
+    # image should be RGB, IMPORTANT: normalized to [-1,1]
+    print("calculate_lpips...")
+
+    lpips_results = {}
+    dataloader1, dataloader2 = build_dataloader(videos1, videos2)
+    for video1, video2 in tqdm(zip(dataloader1, dataloader2), total=len(dataloader1)):
+        # get a video [timestamps, channel, h, w]
+        
+        assert video1.shape == video2.shape
+        video1 = video1.squeeze(0) * 2 - 1
+        video2 = video2.squeeze(0) * 2 - 1
+
+        for clip_timestamp in range(len(video1)):
+            # get a img
+            # img [timestamps[x], channel, h, w]
+            # img [channel, h, w] tensor
+
+            img1 = video1[clip_timestamp].unsqueeze(0).to(device)
+            img2 = video2[clip_timestamp].unsqueeze(0).to(device)
+            
+            loss_fn.to(device)
+
+            # calculate lpips of a video
+            value = np.array(loss_fn.forward(img1, img2).mean().detach().cpu().tolist())
+            if clip_timestamp not in lpips_results:
+                lpips_results[clip_timestamp] = value
+            else:
+                lpips_results[clip_timestamp] += value
+
+    for clip_timestamp in range(len(video1)):
+        lpips_results[clip_timestamp] /= len(dataloader1)
+
+    result = {
+        "value": lpips_results,
+    }
+
+    return result
+
+# test code / using example
+
+def main():
+    NUMBER_OF_VIDEOS = 8
+    VIDEO_LENGTH = 50
+    CHANNEL = 3
+    SIZE = 64
+    videos1 = torch.zeros(NUMBER_OF_VIDEOS, VIDEO_LENGTH, CHANNEL, SIZE, SIZE, requires_grad=False)
+    videos2 = torch.ones(NUMBER_OF_VIDEOS, VIDEO_LENGTH, CHANNEL, SIZE, SIZE, requires_grad=False)
+    device = torch.device("cuda")
+    # device = torch.device("cpu")
+
+    import json
+    result = calculate_lpips(videos1, videos2, device)
+    print(json.dumps(result, indent=4))
+
+if __name__ == "__main__":
+    main()

grn/tokenizer/videovae/evaluation/common_metrics_on_video_quality/calculate_psnr.py

@@ -0,0 +1,83 @@
+import numpy as np
+import torch
+from tqdm import tqdm
+import math
+from multiprocessing import Pool
+from .utils import build_dataloader
+
+def img_psnr(img1, img2):
+    # [0,1]
+    # compute mse
+    # mse = np.mean((img1-img2)**2)
+    mse = np.mean((img1 / 1.0 - img2 / 1.0) ** 2)
+    # compute psnr
+    if mse < 1e-10:
+        return 100
+    psnr = 20 * math.log10(1 / math.sqrt(mse))
+    return psnr
+
+def trans(x):
+    return x
+
+def process_video(video_pair):
+    video1, video2 = video_pair
+    psnr_results_of_a_video = []
+    for clip_timestamp in range(len(video1)):
+        # get a img
+        # img [timestamps[x], channel, h, w]
+        # img [channel, h, w] numpy
+
+        img1 = video1[clip_timestamp].numpy()
+        img2 = video2[clip_timestamp].numpy()
+        
+        # calculate psnr of a video
+        psnr_results_of_a_video.append(img_psnr(img1, img2))
+    
+    return psnr_results_of_a_video
+    
+
+def calculate_psnr(videos1, videos2):
+    print("calculate_psnr...")
+
+    # videos [batch_size, timestamps, channel, h, w]
+    dataloader1, dataloader2 = build_dataloader(videos1, videos2)
+    psnr_results = []
+    for video1, video2 in tqdm(zip(dataloader1, dataloader2), total=len(dataloader1)):
+        video1 = video1.squeeze(0)
+        video2 = video2.squeeze(0)
+        video_pair = (video1, video2)
+        result = process_video(video_pair)
+        psnr_results.append(result)
+
+    psnr_results = np.array(psnr_results)
+    
+    psnr = {}
+    psnr_std = {}
+
+    for clip_timestamp in range(len(video1)):
+        psnr[clip_timestamp] = np.mean(psnr_results[:,clip_timestamp])
+        psnr_std[clip_timestamp] = np.std(psnr_results[:,clip_timestamp])
+
+    result = {
+        "value": psnr,
+        "value_std": psnr_std,
+    }
+
+    return result
+
+# test code / using example
+
+def main():
+    NUMBER_OF_VIDEOS = 8
+    VIDEO_LENGTH = 50
+    CHANNEL = 3
+    SIZE = 64
+    videos1 = torch.zeros(NUMBER_OF_VIDEOS, VIDEO_LENGTH, CHANNEL, SIZE, SIZE, requires_grad=False)
+    videos2 = torch.zeros(NUMBER_OF_VIDEOS, VIDEO_LENGTH, CHANNEL, SIZE, SIZE, requires_grad=False)
+
+    import json
+    result = calculate_psnr(videos1, videos2)
+    print(json.dumps(result, indent=4))
+
+if __name__ == "__main__":
+    main()

grn/tokenizer/videovae/evaluation/common_metrics_on_video_quality/calculate_ssim.py

@@ -0,0 +1,138 @@
+import numpy as np
+import torch
+from tqdm import tqdm
+import cv2
+from multiprocessing import Pool
+import torch
+import torch.nn.functional as F
+from .utils import build_dataloader
+ 
+def ssim(img1, img2):
+    _img1, _img2 = img1, img2
+
+    ### original implementation
+    # C1 = 0.01 ** 2
+    # C2 = 0.03 ** 2
+    # img1 = img1.astype(np.float64)
+    # img2 = img2.astype(np.float64)
+    # kernel = cv2.getGaussianKernel(11, 1.5)
+    # window = np.outer(kernel, kernel.transpose())
+    # mu1 = cv2.filter2D(img1, -1, window)[5:-5, 5:-5]  # valid
+    # mu2 = cv2.filter2D(img2, -1, window)[5:-5, 5:-5]
+    # mu1_sq = mu1 ** 2
+    # mu2_sq = mu2 ** 2
+    # mu1_mu2 = mu1 * mu2
+    # sigma1_sq = cv2.filter2D(img1 ** 2, -1, window)[5:-5, 5:-5] - mu1_sq
+    # sigma2_sq = cv2.filter2D(img2 ** 2, -1, window)[5:-5, 5:-5] - mu2_sq
+    # sigma12 = cv2.filter2D(img1 * img2, -1, window)[5:-5, 5:-5] - mu1_mu2
+    # ssim_map = ((2 * mu1_mu2 + C1) * (2 * sigma12 + C2)) / ((mu1_sq + mu2_sq + C1) *
+    #                                                         (sigma1_sq + sigma2_sq + C2))
+    # return ssim_map.mean()
+    # res1 = ssim_map.mean()
+
+    ### accelerated implementation
+    img1, img2 = torch.from_numpy(_img1), torch.from_numpy(_img2)
+    C1 = 0.01 ** 2
+    C2 = 0.03 ** 2
+    
+    # Ensure data is float and move to GPU
+    img1 = img1.to(torch.float64).cuda()
+    img2 = img2.to(torch.float64).cuda()
+    
+    # Gaussian kernel
+    kernel = torch.tensor(cv2.getGaussianKernel(11, 1.5)).to(torch.float64).cuda()
+    window = kernel @ kernel.t()
+    window = window.unsqueeze(0).unsqueeze(0).cuda()
+    
+    mu1 = F.conv2d(img1.unsqueeze(0), window, padding=0, groups=1)
+    mu2 = F.conv2d(img2.unsqueeze(0), window, padding=0, groups=1)
+    mu1_sq = mu1.pow(2)
+    mu2_sq = mu2.pow(2)
+    mu1_mu2 = mu1 * mu2
+    sigma1_sq = F.conv2d(img1.unsqueeze(0) ** 2, window, padding=0, groups=1) - mu1_sq
+    sigma2_sq = F.conv2d(img2.unsqueeze(0) ** 2, window, padding=0, groups=1) - mu2_sq
+    sigma12 = F.conv2d(img1.unsqueeze(0) * img2.unsqueeze(0), window, padding=0, groups=1) - mu1_mu2
+    
+    ssim_map = ((2 * mu1_mu2 + C1) * (2 * sigma12 + C2)) / ((mu1_sq + mu2_sq + C1) *
+                                                            (sigma1_sq + sigma2_sq + C2))
+    res2 = ssim_map.mean().item()
+    # print(res1-res2)
+    return res2
+    # return ssim_map.mean().item()
+ 
+ 
+def calculate_ssim_function(img1, img2):
+    # [0,1]
+    # ssim is the only metric extremely sensitive to gray being compared to b/w 
+    if not img1.shape == img2.shape:
+        raise ValueError('Input images must have the same dimensions.')
+    if img1.ndim == 2:
+        return ssim(img1, img2)
+    elif img1.ndim == 3:
+        if img1.shape[0] == 3:
+            ssims = []
+            for i in range(3):
+                ssims.append(ssim(img1[i], img2[i]))
+            return np.array(ssims).mean()                   
+        elif img1.shape[0] == 1:
+            return ssim(np.squeeze(img1), np.squeeze(img2))
+    else:
+        raise ValueError('Wrong input image dimensions.')
+
+def trans(x):
+    return x
+
+def process_video(video_pair):
+    video1, video2 = video_pair
+    ssim_results_of_a_video = []
+    for clip_timestamp in range(len(video1)):
+        img1 = video1[clip_timestamp].numpy()
+        img2 = video2[clip_timestamp].numpy()
+        ssim_results_of_a_video.append(calculate_ssim_function(img1, img2))
+    return ssim_results_of_a_video
+
+def calculate_ssim(videos1, videos2):
+    print("calculate_ssim...")
+
+    ssim_results = []
+    dataloader1, dataloader2 = build_dataloader(videos1, videos2)
+    for video1, video2 in tqdm(zip(dataloader1, dataloader2), total=len(dataloader1)):
+        video1 = video1.squeeze(0)
+        video2 = video2.squeeze(0)
+        video_pair = (video1, video2)
+        result = process_video(video_pair)
+        ssim_results.append(result)
+
+    ssim_results = np.array(ssim_results)
+
+    ssim = {}
+    ssim_std = {}
+
+    for clip_timestamp in range(len(video1)):
+        ssim[clip_timestamp] = np.mean(ssim_results[:,clip_timestamp])
+        ssim_std[clip_timestamp] = np.std(ssim_results[:,clip_timestamp])
+
+    result = {
+        "value": ssim,
+        "value_std": ssim_std,
+    }
+
+    return result
+
+# test code / using example
+
+def main():
+    NUMBER_OF_VIDEOS = 8
+    VIDEO_LENGTH = 50
+    CHANNEL = 3
+    SIZE = 64
+    videos1 = torch.zeros(NUMBER_OF_VIDEOS, VIDEO_LENGTH, CHANNEL, SIZE, SIZE, requires_grad=False)
+    videos2 = torch.zeros(NUMBER_OF_VIDEOS, VIDEO_LENGTH, CHANNEL, SIZE, SIZE, requires_grad=False)
+    device = torch.device("cuda")
+
+    import json
+    result = calculate_ssim(videos1, videos2)
+    print(json.dumps(result, indent=4))
+
+if __name__ == "__main__":
+    main()

grn/tokenizer/videovae/evaluation/common_metrics_on_video_quality/fvd/styleganv/fvd.py

@@ -0,0 +1,90 @@
+import torch
+import os
+import math
+import torch.nn.functional as F
+
+# https://github.com/universome/fvd-comparison
+
+
+def load_i3d_pretrained(device=torch.device('cpu')):
+    i3D_WEIGHTS_URL = "https://www.dropbox.com/s/ge9e5ujwgetktms/i3d_torchscript.pt"
+    filepath = os.path.join(os.path.dirname(os.path.abspath(__file__)), 'i3d_torchscript.pt')
+    print(filepath)
+    if not os.path.exists(filepath):
+        print(f"preparing for download {i3D_WEIGHTS_URL}, you can download it by yourself.")
+        os.system(f"wget {i3D_WEIGHTS_URL} -O {filepath}")
+    i3d = torch.jit.load(filepath).eval().to(device)
+    i3d = torch.nn.DataParallel(i3d)
+    return i3d
+    
+
+def get_feats(videos, detector, device, bs=10):
+    # videos : torch.tensor BCTHW [0, 1]
+    detector_kwargs = dict(rescale=False, resize=False, return_features=True) # Return raw features before the softmax layer.
+    feats = np.empty((0, 400))
+    with torch.no_grad():
+        for i in range((len(videos)-1)//bs + 1):
+            feats = np.vstack([feats, detector(torch.stack([preprocess_single(video) for video in videos[i*bs:(i+1)*bs]]).to(device), **detector_kwargs).detach().cpu().numpy()])
+    return feats
+
+
+def get_fvd_feats(videos, i3d, device, bs=10):
+    # videos in [0, 1] as torch tensor BCTHW
+    # videos = [preprocess_single(video) for video in videos]
+    embeddings = get_feats(videos, i3d, device, bs)
+    return embeddings
+
+
+def preprocess_single(video, resolution=224, sequence_length=None):
+    # video: CTHW, [0, 1]
+    c, t, h, w = video.shape
+
+    # temporal crop
+    if sequence_length is not None:
+        assert sequence_length <= t
+        video = video[:, :sequence_length]
+
+    # scale shorter side to resolution
+    scale = resolution / min(h, w)
+    if h < w:
+        target_size = (resolution, math.ceil(w * scale))
+    else:
+        target_size = (math.ceil(h * scale), resolution)
+    video = F.interpolate(video, size=target_size, mode='bilinear', align_corners=False)
+
+    # center crop
+    c, t, h, w = video.shape
+    w_start = (w - resolution) // 2
+    h_start = (h - resolution) // 2
+    video = video[:, :, h_start:h_start + resolution, w_start:w_start + resolution]
+
+    # [0, 1] -> [-1, 1]
+    video = (video - 0.5) * 2
+
+    return video.contiguous()
+
+
+"""
+Copy-pasted from https://github.com/cvpr2022-stylegan-v/stylegan-v/blob/main/src/metrics/frechet_video_distance.py
+"""
+from typing import Tuple
+from scipy.linalg import sqrtm
+import numpy as np
+
+
+def compute_stats(feats: np.ndarray) -> Tuple[np.ndarray, np.ndarray]:
+    mu = feats.mean(axis=0) # [d]
+    sigma = np.cov(feats, rowvar=False) # [d, d]
+    return mu, sigma
+
+
+def frechet_distance(feats_fake: np.ndarray, feats_real: np.ndarray) -> float:
+    mu_gen, sigma_gen = compute_stats(feats_fake)
+    mu_real, sigma_real = compute_stats(feats_real)
+    m = np.square(mu_gen - mu_real).sum()
+    if feats_fake.shape[0]>1:
+        s, _ = sqrtm(np.dot(sigma_gen, sigma_real), disp=False) # pylint: disable=no-member
+        fid = np.real(m + np.trace(sigma_gen + sigma_real - s * 2))
+    else:
+        fid = np.real(m)
+    return float(fid)

grn/tokenizer/videovae/evaluation/common_metrics_on_video_quality/fvd/videogpt/fvd.py

@@ -0,0 +1,137 @@
+import torch
+import os
+import math
+import torch.nn.functional as F
+import numpy as np
+import einops
+
+def load_i3d_pretrained(device=torch.device('cpu')):
+    i3D_WEIGHTS_URL = "https://onedrive.live.com/download?cid=78EEF3EB6AE7DBCB&resid=78EEF3EB6AE7DBCB%21199&authkey=AApKdFHPXzWLNyI"
+    filepath = os.path.join(os.path.dirname(os.path.abspath(__file__)), 'i3d_pretrained_400.pt')
+    print(filepath)
+    if not os.path.exists(filepath):
+        print(f"preparing for download {i3D_WEIGHTS_URL}, you can download it by yourself.")
+        os.system(f"wget {i3D_WEIGHTS_URL} -O {filepath}")
+    from .pytorch_i3d import InceptionI3d
+    i3d = InceptionI3d(400, in_channels=3).eval().to(device)
+    i3d.load_state_dict(torch.load(filepath, map_location=device, weights_only=True))
+    i3d = torch.nn.DataParallel(i3d)
+    return i3d
+
+def preprocess_single(video, resolution, sequence_length=None):
+    # video: THWC, {0, ..., 255}
+    video = video.permute(0, 3, 1, 2).float() / 255. # TCHW
+    t, c, h, w = video.shape
+
+    # temporal crop
+    if sequence_length is not None:
+        assert sequence_length <= t
+        video = video[:sequence_length]
+
+    # scale shorter side to resolution
+    scale = resolution / min(h, w)
+    if h < w:
+        target_size = (resolution, math.ceil(w * scale))
+    else:
+        target_size = (math.ceil(h * scale), resolution)
+    video = F.interpolate(video, size=target_size, mode='bilinear',
+                          align_corners=False)
+
+    # center crop
+    t, c, h, w = video.shape
+    w_start = (w - resolution) // 2
+    h_start = (h - resolution) // 2
+    video = video[:, :, h_start:h_start + resolution, w_start:w_start + resolution]
+    video = video.permute(1, 0, 2, 3).contiguous() # CTHW
+
+    video -= 0.5
+
+    return video
+
+def preprocess(videos, target_resolution=224):
+    # we should tras videos in [0-1] [b c t h w] as th.float 
+    # -> videos in {0, ..., 255} [b t h w c] as np.uint8 array
+    videos = einops.rearrange(videos, 'b c t h w -> b t h w c')
+    videos = (videos*255).numpy().astype(np.uint8)
+
+    b, t, h, w, c = videos.shape
+    videos = torch.from_numpy(videos)
+    videos = torch.stack([preprocess_single(video, target_resolution) for video in videos])
+    return videos * 2 # [-0.5, 0.5] -> [-1, 1]
+
+def get_fvd_logits(videos, i3d, device, bs=10):
+    videos = preprocess(videos)
+    embeddings = get_logits(i3d, videos, device, bs=10)
+    return embeddings
+
+# https://github.com/tensorflow/gan/blob/de4b8da3853058ea380a6152bd3bd454013bf619/tensorflow_gan/python/eval/classifier_metrics.py#L161
+def _symmetric_matrix_square_root(mat, eps=1e-10):
+    u, s, v = torch.svd(mat)
+    si = torch.where(s < eps, s, torch.sqrt(s))
+    return torch.matmul(torch.matmul(u, torch.diag(si)), v.t())
+
+# https://github.com/tensorflow/gan/blob/de4b8da3853058ea380a6152bd3bd454013bf619/tensorflow_gan/python/eval/classifier_metrics.py#L400
+def trace_sqrt_product(sigma, sigma_v):
+    sqrt_sigma = _symmetric_matrix_square_root(sigma)
+    sqrt_a_sigmav_a = torch.matmul(sqrt_sigma, torch.matmul(sigma_v, sqrt_sigma))
+    return torch.trace(_symmetric_matrix_square_root(sqrt_a_sigmav_a))
+
+# https://discuss.pytorch.org/t/covariance-and-gradient-support/16217/2
+def cov(m, rowvar=False):
+    '''Estimate a covariance matrix given data.
+
+    Covariance indicates the level to which two variables vary together.
+    If we examine N-dimensional samples, `X = [x_1, x_2, ... x_N]^T`,
+    then the covariance matrix element `C_{ij}` is the covariance of
+    `x_i` and `x_j`. The element `C_{ii}` is the variance of `x_i`.
+
+    Args:
+        m: A 1-D or 2-D array containing multiple variables and observations.
+            Each row of `m` represents a variable, and each column a single
+            observation of all those variables.
+        rowvar: If `rowvar` is True, then each row represents a
+            variable, with observations in the columns. Otherwise, the
+            relationship is transposed: each column represents a variable,
+            while the rows contain observations.
+
+    Returns:
+        The covariance matrix of the variables.
+    '''
+    if m.dim() > 2:
+        raise ValueError('m has more than 2 dimensions')
+    if m.dim() < 2:
+        m = m.view(1, -1)
+    if not rowvar and m.size(0) != 1:
+        m = m.t()
+
+    fact = 1.0 / (m.size(1) - 1) # unbiased estimate
+    m -= torch.mean(m, dim=1, keepdim=True)
+    mt = m.t()  # if complex: mt = m.t().conj()
+    return fact * m.matmul(mt).squeeze()
+
+
+def frechet_distance(x1, x2):
+    x1 = x1.flatten(start_dim=1)
+    x2 = x2.flatten(start_dim=1)
+    m, m_w = x1.mean(dim=0), x2.mean(dim=0)
+    sigma, sigma_w = cov(x1, rowvar=False), cov(x2, rowvar=False)
+    mean = torch.sum((m - m_w) ** 2)
+    if x1.shape[0]>1:
+        sqrt_trace_component = trace_sqrt_product(sigma, sigma_w)
+        trace = torch.trace(sigma + sigma_w) - 2.0 * sqrt_trace_component
+        fd = trace + mean
+    else:
+        fd = np.real(mean)
+    return float(fd)
+
+
+def get_logits(i3d, videos, device, bs=10):
+    # assert videos.shape[0] % 16 == 0
+    with torch.no_grad():
+        logits = []
+        for i in range(0, videos.shape[0], bs):
+            batch = videos[i:i + bs].to(device)
+            # logits.append(i3d.module.extract_features(batch)) # wrong
+            logits.append(i3d(batch)) # right
+        logits = torch.cat(logits, dim=0)
+        return logits

grn/tokenizer/videovae/evaluation/common_metrics_on_video_quality/fvd/videogpt/pytorch_i3d.py

@@ -0,0 +1,322 @@
+# Original code from https://github.com/piergiaj/pytorch-i3d
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import numpy as np
+
+class MaxPool3dSamePadding(nn.MaxPool3d):
+
+    def compute_pad(self, dim, s):
+        if s % self.stride[dim] == 0:
+            return max(self.kernel_size[dim] - self.stride[dim], 0)
+        else:
+            return max(self.kernel_size[dim] - (s % self.stride[dim]), 0)
+
+    def forward(self, x):
+        # compute 'same' padding
+        (batch, channel, t, h, w) = x.size()
+        out_t = np.ceil(float(t) / float(self.stride[0]))
+        out_h = np.ceil(float(h) / float(self.stride[1]))
+        out_w = np.ceil(float(w) / float(self.stride[2]))
+        pad_t = self.compute_pad(0, t)
+        pad_h = self.compute_pad(1, h)
+        pad_w = self.compute_pad(2, w)
+
+        pad_t_f = pad_t // 2
+        pad_t_b = pad_t - pad_t_f
+        pad_h_f = pad_h // 2
+        pad_h_b = pad_h - pad_h_f
+        pad_w_f = pad_w // 2
+        pad_w_b = pad_w - pad_w_f
+
+        pad = (pad_w_f, pad_w_b, pad_h_f, pad_h_b, pad_t_f, pad_t_b)
+        x = F.pad(x, pad)
+        return super(MaxPool3dSamePadding, self).forward(x)
+
+
+class Unit3D(nn.Module):
+
+    def __init__(self, in_channels,
+                 output_channels,
+                 kernel_shape=(1, 1, 1),
+                 stride=(1, 1, 1),
+                 padding=0,
+                 activation_fn=F.relu,
+                 use_batch_norm=True,
+                 use_bias=False,
+                 name='unit_3d'):
+
+        """Initializes Unit3D module."""
+        super(Unit3D, self).__init__()
+
+        self._output_channels = output_channels
+        self._kernel_shape = kernel_shape
+        self._stride = stride
+        self._use_batch_norm = use_batch_norm
+        self._activation_fn = activation_fn
+        self._use_bias = use_bias
+        self.name = name
+        self.padding = padding
+
+        self.conv3d = nn.Conv3d(in_channels=in_channels,
+                                out_channels=self._output_channels,
+                                kernel_size=self._kernel_shape,
+                                stride=self._stride,
+                                padding=0, # we always want padding to be 0 here. We will dynamically pad based on input size in forward function
+                                bias=self._use_bias)
+
+        if self._use_batch_norm:
+            self.bn = nn.BatchNorm3d(self._output_channels, eps=1e-5, momentum=0.001)
+
+    def compute_pad(self, dim, s):
+        if s % self._stride[dim] == 0:
+            return max(self._kernel_shape[dim] - self._stride[dim], 0)
+        else:
+            return max(self._kernel_shape[dim] - (s % self._stride[dim]), 0)
+
+
+    def forward(self, x):
+        # compute 'same' padding
+        (batch, channel, t, h, w) = x.size()
+        out_t = np.ceil(float(t) / float(self._stride[0]))
+        out_h = np.ceil(float(h) / float(self._stride[1]))
+        out_w = np.ceil(float(w) / float(self._stride[2]))
+        pad_t = self.compute_pad(0, t)
+        pad_h = self.compute_pad(1, h)
+        pad_w = self.compute_pad(2, w)
+
+        pad_t_f = pad_t // 2
+        pad_t_b = pad_t - pad_t_f
+        pad_h_f = pad_h // 2
+        pad_h_b = pad_h - pad_h_f
+        pad_w_f = pad_w // 2
+        pad_w_b = pad_w - pad_w_f
+
+        pad = (pad_w_f, pad_w_b, pad_h_f, pad_h_b, pad_t_f, pad_t_b)
+        x = F.pad(x, pad)
+
+        x = self.conv3d(x)
+        if self._use_batch_norm:
+            x = self.bn(x)
+        if self._activation_fn is not None:
+            x = self._activation_fn(x)
+        return x
+
+
+
+class InceptionModule(nn.Module):
+    def __init__(self, in_channels, out_channels, name):
+        super(InceptionModule, self).__init__()
+
+        self.b0 = Unit3D(in_channels=in_channels, output_channels=out_channels[0], kernel_shape=[1, 1, 1], padding=0,
+                         name=name+'/Branch_0/Conv3d_0a_1x1')
+        self.b1a = Unit3D(in_channels=in_channels, output_channels=out_channels[1], kernel_shape=[1, 1, 1], padding=0,
+                          name=name+'/Branch_1/Conv3d_0a_1x1')
+        self.b1b = Unit3D(in_channels=out_channels[1], output_channels=out_channels[2], kernel_shape=[3, 3, 3],
+                          name=name+'/Branch_1/Conv3d_0b_3x3')
+        self.b2a = Unit3D(in_channels=in_channels, output_channels=out_channels[3], kernel_shape=[1, 1, 1], padding=0,
+                          name=name+'/Branch_2/Conv3d_0a_1x1')
+        self.b2b = Unit3D(in_channels=out_channels[3], output_channels=out_channels[4], kernel_shape=[3, 3, 3],
+                          name=name+'/Branch_2/Conv3d_0b_3x3')
+        self.b3a = MaxPool3dSamePadding(kernel_size=[3, 3, 3],
+                                stride=(1, 1, 1), padding=0)
+        self.b3b = Unit3D(in_channels=in_channels, output_channels=out_channels[5], kernel_shape=[1, 1, 1], padding=0,
+                          name=name+'/Branch_3/Conv3d_0b_1x1')
+        self.name = name
+
+    def forward(self, x):
+        b0 = self.b0(x)
+        b1 = self.b1b(self.b1a(x))
+        b2 = self.b2b(self.b2a(x))
+        b3 = self.b3b(self.b3a(x))
+        return torch.cat([b0,b1,b2,b3], dim=1)
+
+
+class InceptionI3d(nn.Module):
+    """Inception-v1 I3D architecture.
+    The model is introduced in:
+        Quo Vadis, Action Recognition? A New Model and the Kinetics Dataset
+        Joao Carreira, Andrew Zisserman
+        https://arxiv.org/pdf/1705.07750v1.pdf.
+    See also the Inception architecture, introduced in:
+        Going deeper with convolutions
+        Christian Szegedy, Wei Liu, Yangqing Jia, Pierre Sermanet, Scott Reed,
+        Dragomir Anguelov, Dumitru Erhan, Vincent Vanhoucke, Andrew Rabinovich.
+        http://arxiv.org/pdf/1409.4842v1.pdf.
+    """
+
+    # Endpoints of the model in order. During construction, all the endpoints up
+    # to a designated `final_endpoint` are returned in a dictionary as the
+    # second return value.
+    VALID_ENDPOINTS = (
+        'Conv3d_1a_7x7',
+        'MaxPool3d_2a_3x3',
+        'Conv3d_2b_1x1',
+        'Conv3d_2c_3x3',
+        'MaxPool3d_3a_3x3',
+        'Mixed_3b',
+        'Mixed_3c',
+        'MaxPool3d_4a_3x3',
+        'Mixed_4b',
+        'Mixed_4c',
+        'Mixed_4d',
+        'Mixed_4e',
+        'Mixed_4f',
+        'MaxPool3d_5a_2x2',
+        'Mixed_5b',
+        'Mixed_5c',
+        'Logits',
+        'Predictions',
+    )
+
+    def __init__(self, num_classes=400, spatial_squeeze=True,
+                 final_endpoint='Logits', name='inception_i3d', in_channels=3, dropout_keep_prob=0.5):
+        """Initializes I3D model instance.
+        Args:
+          num_classes: The number of outputs in the logit layer (default 400, which
+              matches the Kinetics dataset).
+          spatial_squeeze: Whether to squeeze the spatial dimensions for the logits
+              before returning (default True).
+          final_endpoint: The model contains many possible endpoints.
+              `final_endpoint` specifies the last endpoint for the model to be built
+              up to. In addition to the output at `final_endpoint`, all the outputs
+              at endpoints up to `final_endpoint` will also be returned, in a
+              dictionary. `final_endpoint` must be one of
+              InceptionI3d.VALID_ENDPOINTS (default 'Logits').
+          name: A string (optional). The name of this module.
+        Raises:
+          ValueError: if `final_endpoint` is not recognized.
+        """
+
+        if final_endpoint not in self.VALID_ENDPOINTS:
+            raise ValueError('Unknown final endpoint %s' % final_endpoint)
+
+        super(InceptionI3d, self).__init__()
+        self._num_classes = num_classes
+        self._spatial_squeeze = spatial_squeeze
+        self._final_endpoint = final_endpoint
+        self.logits = None
+
+        if self._final_endpoint not in self.VALID_ENDPOINTS:
+            raise ValueError('Unknown final endpoint %s' % self._final_endpoint)
+
+        self.end_points = {}
+        end_point = 'Conv3d_1a_7x7'
+        self.end_points[end_point] = Unit3D(in_channels=in_channels, output_channels=64, kernel_shape=[7, 7, 7],
+                                            stride=(2, 2, 2), padding=(3,3,3),  name=name+end_point)
+        if self._final_endpoint == end_point: return
+
+        end_point = 'MaxPool3d_2a_3x3'
+        self.end_points[end_point] = MaxPool3dSamePadding(kernel_size=[1, 3, 3], stride=(1, 2, 2),
+                                                             padding=0)
+        if self._final_endpoint == end_point: return
+
+        end_point = 'Conv3d_2b_1x1'
+        self.end_points[end_point] = Unit3D(in_channels=64, output_channels=64, kernel_shape=[1, 1, 1], padding=0,
+                                       name=name+end_point)
+        if self._final_endpoint == end_point: return
+
+        end_point = 'Conv3d_2c_3x3'
+        self.end_points[end_point] = Unit3D(in_channels=64, output_channels=192, kernel_shape=[3, 3, 3], padding=1,
+                                       name=name+end_point)
+        if self._final_endpoint == end_point: return
+
+        end_point = 'MaxPool3d_3a_3x3'
+        self.end_points[end_point] = MaxPool3dSamePadding(kernel_size=[1, 3, 3], stride=(1, 2, 2),
+                                                             padding=0)
+        if self._final_endpoint == end_point: return
+
+        end_point = 'Mixed_3b'
+        self.end_points[end_point] = InceptionModule(192, [64,96,128,16,32,32], name+end_point)
+        if self._final_endpoint == end_point: return
+
+        end_point = 'Mixed_3c'
+        self.end_points[end_point] = InceptionModule(256, [128,128,192,32,96,64], name+end_point)
+        if self._final_endpoint == end_point: return
+
+        end_point = 'MaxPool3d_4a_3x3'
+        self.end_points[end_point] = MaxPool3dSamePadding(kernel_size=[3, 3, 3], stride=(2, 2, 2),
+                                                             padding=0)
+        if self._final_endpoint == end_point: return
+
+        end_point = 'Mixed_4b'
+        self.end_points[end_point] = InceptionModule(128+192+96+64, [192,96,208,16,48,64], name+end_point)
+        if self._final_endpoint == end_point: return
+
+        end_point = 'Mixed_4c'
+        self.end_points[end_point] = InceptionModule(192+208+48+64, [160,112,224,24,64,64], name+end_point)
+        if self._final_endpoint == end_point: return
+
+        end_point = 'Mixed_4d'
+        self.end_points[end_point] = InceptionModule(160+224+64+64, [128,128,256,24,64,64], name+end_point)
+        if self._final_endpoint == end_point: return
+
+        end_point = 'Mixed_4e'
+        self.end_points[end_point] = InceptionModule(128+256+64+64, [112,144,288,32,64,64], name+end_point)
+        if self._final_endpoint == end_point: return
+
+        end_point = 'Mixed_4f'
+        self.end_points[end_point] = InceptionModule(112+288+64+64, [256,160,320,32,128,128], name+end_point)
+        if self._final_endpoint == end_point: return
+
+        end_point = 'MaxPool3d_5a_2x2'
+        self.end_points[end_point] = MaxPool3dSamePadding(kernel_size=[2, 2, 2], stride=(2, 2, 2),
+                                                             padding=0)
+        if self._final_endpoint == end_point: return
+
+        end_point = 'Mixed_5b'
+        self.end_points[end_point] = InceptionModule(256+320+128+128, [256,160,320,32,128,128], name+end_point)
+        if self._final_endpoint == end_point: return
+
+        end_point = 'Mixed_5c'
+        self.end_points[end_point] = InceptionModule(256+320+128+128, [384,192,384,48,128,128], name+end_point)
+        if self._final_endpoint == end_point: return
+
+        end_point = 'Logits'
+        self.avg_pool = nn.AvgPool3d(kernel_size=[2, 7, 7],
+                                     stride=(1, 1, 1))
+        self.dropout = nn.Dropout(dropout_keep_prob)
+        self.logits = Unit3D(in_channels=384+384+128+128, output_channels=self._num_classes,
+                             kernel_shape=[1, 1, 1],
+                             padding=0,
+                             activation_fn=None,
+                             use_batch_norm=False,
+                             use_bias=True,
+                             name='logits')
+
+        self.build()
+
+
+    def replace_logits(self, num_classes):
+        self._num_classes = num_classes
+        self.logits = Unit3D(in_channels=384+384+128+128, output_channels=self._num_classes,
+                             kernel_shape=[1, 1, 1],
+                             padding=0,
+                             activation_fn=None,
+                             use_batch_norm=False,
+                             use_bias=True,
+                             name='logits')
+
+
+    def build(self):
+        for k in self.end_points.keys():
+            self.add_module(k, self.end_points[k])
+
+    def forward(self, x):
+        for end_point in self.VALID_ENDPOINTS:
+            if end_point in self.end_points:
+                x = self._modules[end_point](x) # use _modules to work with dataparallel
+
+        x = self.logits(self.dropout(self.avg_pool(x)))
+        if self._spatial_squeeze:
+            logits = x.squeeze(3).squeeze(3)
+        logits = logits.mean(dim=2)
+        # logits is batch X time X classes, which is what we want to work with
+        return logits
+
+
+    def extract_features(self, x):
+        for end_point in self.VALID_ENDPOINTS:
+            if end_point in self.end_points:
+                x = self._modules[end_point](x)
+        return self.avg_pool(x)

grn/tokenizer/videovae/evaluation/common_metrics_on_video_quality/utils.py

@@ -0,0 +1,27 @@
+import torch
+from torch.utils.data import Dataset, DataLoader
+
+class VideoDataset(Dataset):
+    def __init__(self, videos):
+        self.videos = videos
+
+    def __len__(self):
+        return len(self.videos) if type(self.videos) == list else self.videos.shape[0]
+
+    def __getitem__(self, idx):
+        video = self.videos[idx]
+        if isinstance(video, str):
+            video = torch.load(video, weights_only=True)
+        return video
+
+def build_dataloader(videos1, videos2):
+    dataset1 = VideoDataset(videos1)
+    dataset2 = VideoDataset(videos2)
+    assert len(dataset1) == len(dataset2)
+
+    dataloader1 = DataLoader(dataset1, batch_size=1, num_workers=8, shuffle=False)
+    dataloader2 = DataLoader(dataset2, batch_size=1, num_workers=8, shuffle=False)
+
+    return dataloader1, dataloader2
+
+

grn/tokenizer/videovae/evaluation/fid.py

@@ -0,0 +1,62 @@
+import numpy as np
+from scipy import linalg
+
+
+def calculate_frechet_distance(mu1, sigma1, mu2, sigma2, eps=1e-6):
+    """Numpy implementation of the Frechet Distance.
+    The Frechet distance between two multivariate Gaussians X_1 ~ N(mu_1, C_1)
+    and X_2 ~ N(mu_2, C_2) is
+            d^2 = ||mu_1 - mu_2||^2 + Tr(C_1 + C_2 - 2*sqrt(C_1*C_2)).
+
+    Stable version by Dougal J. Sutherland.
+
+    Params:
+    -- mu1   : Numpy array containing the activations of a layer of the
+               inception net (like returned by the function 'get_predictions')
+               for generated samples.
+    -- mu2   : The sample mean over activations, precalculated on an
+               representative data set.
+    -- sigma1: The covariance matrix over activations for generated samples.
+    -- sigma2: The covariance matrix over activations, precalculated on an
+               representative data set.
+
+    Returns:
+    --   : The Frechet Distance.
+    """
+
+    mu1 = np.atleast_1d(mu1)
+    mu2 = np.atleast_1d(mu2)
+
+    sigma1 = np.atleast_2d(sigma1)
+    sigma2 = np.atleast_2d(sigma2)
+
+    assert (
+        mu1.shape == mu2.shape
+    ), "Training and test mean vectors have different lengths"
+    assert (
+        sigma1.shape == sigma2.shape
+    ), "Training and test covariances have different dimensions"
+
+    diff = mu1 - mu2
+
+    # Product might be almost singular
+    covmean, _ = linalg.sqrtm(sigma1.dot(sigma2), disp=False)
+    if not np.isfinite(covmean).all():
+        msg = (
+            "fid calculation produces singular product; "
+            "adding %s to diagonal of cov estimates"
+        ) % eps
+        print(msg)
+        offset = np.eye(sigma1.shape[0]) * eps
+        covmean = linalg.sqrtm((sigma1 + offset).dot(sigma2 + offset))
+
+    # Numerical error might give slight imaginary component
+    if np.iscomplexobj(covmean):
+        if not np.allclose(np.diagonal(covmean).imag, 0, atol=1e-3):
+            m = np.max(np.abs(covmean.imag))
+            raise ValueError("Imaginary component {}".format(m))
+        covmean = covmean.real
+
+    tr_covmean = np.trace(covmean)
+
+    return diff.dot(diff) + np.trace(sigma1) + np.trace(sigma2) - 2 * tr_covmean

grn/tokenizer/videovae/evaluation/fvd.py

@@ -0,0 +1,150 @@
+import argparse
+from email.policy import strict
+import numpy as np
+
+import torch
+import torch.nn.functional as F
+import torch.utils.data as data
+
+from .pytorch_i3d import InceptionI3d
+import os
+from videovae.utils.misc import data_prefix_manager
+
+from sklearn.metrics.pairwise import polynomial_kernel
+
+MAX_BATCH = 16
+FVD_SAMPLE_SIZE = 2048
+TARGET_RESOLUTION = (224, 224)
+
+def preprocess(videos, target_resolution):
+    # videos in {0, ..., 255} as np.uint8 array
+    b, t, h, w, c = videos.shape
+    all_frames = torch.FloatTensor(videos).flatten(end_dim=1) # (b * t, h, w, c)
+    
+    all_frames = all_frames.permute(0, 3, 1, 2).contiguous() # (b * t, c, h, w)
+    resized_videos = F.interpolate(all_frames, size=target_resolution,
+                                   mode='bilinear', align_corners=False)
+    resized_videos = resized_videos.view(b, t, c, *target_resolution)
+    output_videos = resized_videos.transpose(1, 2).contiguous() # (b, c, t, *)
+    scaled_videos = 2. * output_videos / 255. - 1 # [-1, 1]
+    return scaled_videos
+
+def get_fvd_logits(videos, i3d, device):
+    videos = preprocess(videos, TARGET_RESOLUTION)
+    embeddings = get_logits(i3d, videos, device)
+    return embeddings
+
+def load_fvd_model(device="cpu"):
+    i3d = InceptionI3d(400, in_channels=3).to(device)
+    i3d_path = data_prefix_manager('checkpoints/i3d_pretrained_400.pt')
+    i3d.load_state_dict(torch.load(i3d_path, map_location=device, weights_only=True))
+    i3d.eval()
+    return i3d
+
+
+def load_i3d_perceptual():
+    i3d = InceptionI3d(400, in_channels=3)
+    current_dir = os.path.dirname(os.path.abspath(__file__))
+    i3d_path = os.path.join(current_dir, 'i3d_pretrained_400.pt')
+    i3d.load_state_dict(torch.load(i3d_path, map_location=torch.device("cpu"), weights_only=True), strict=False)
+    for param in i3d.parameters():
+        param.requires_grad = False
+    
+    return i3d
+
+# https://github.com/tensorflow/gan/blob/de4b8da3853058ea380a6152bd3bd454013bf619/tensorflow_gan/python/eval/classifier_metrics.py#L161
+def _symmetric_matrix_square_root(mat, eps=1e-10):
+    u, s, v = torch.svd(mat)
+    si = torch.where(s < eps, s, torch.sqrt(s))
+    return torch.matmul(torch.matmul(u, torch.diag(si)), v.t())
+
+# https://github.com/tensorflow/gan/blob/de4b8da3853058ea380a6152bd3bd454013bf619/tensorflow_gan/python/eval/classifier_metrics.py#L400
+def trace_sqrt_product(sigma, sigma_v):
+    sqrt_sigma = _symmetric_matrix_square_root(sigma)
+    sqrt_a_sigmav_a = torch.matmul(sqrt_sigma, torch.matmul(sigma_v, sqrt_sigma))
+    return torch.trace(_symmetric_matrix_square_root(sqrt_a_sigmav_a))
+
+# https://discuss.pytorch.org/t/covariance-and-gradient-support/16217/2
+def cov(m, rowvar=False):
+    '''Estimate a covariance matrix given data.
+
+    Covariance indicates the level to which two variables vary together.
+    If we examine N-dimensional samples, `X = [x_1, x_2, ... x_N]^T`,
+    then the covariance matrix element `C_{ij}` is the covariance of
+    `x_i` and `x_j`. The element `C_{ii}` is the variance of `x_i`.
+
+    Args:
+        m: A 1-D or 2-D array containing multiple variables and observations.
+            Each row of `m` represents a variable, and each column a single
+            observation of all those variables.
+        rowvar: If `rowvar` is True, then each row represents a
+            variable, with observations in the columns. Otherwise, the
+            relationship is transposed: each column represents a variable,
+            while the rows contain observations.
+
+    Returns:
+        The covariance matrix of the variables.
+    '''
+    if m.dim() > 2:
+        raise ValueError('m has more than 2 dimensions')
+    if m.dim() < 2:
+        m = m.view(1, -1)
+    if not rowvar and m.size(0) != 1:
+        m = m.t()
+
+    fact = 1.0 / (m.size(1) - 1) # unbiased estimate
+    m_center = m - torch.mean(m, dim=1, keepdim=True)
+    mt = m_center.t()  # if complex: mt = m.t().conj()
+    return fact * m_center.matmul(mt).squeeze()
+
+
+def frechet_distance(x1, x2):
+    x1 = x1.flatten(start_dim=1)
+    x2 = x2.flatten(start_dim=1)
+    m, m_w = x1.mean(dim=0), x2.mean(dim=0)
+    sigma, sigma_w = cov(x1, rowvar=False), cov(x2, rowvar=False)
+
+    sqrt_trace_component = trace_sqrt_product(sigma, sigma_w)
+    trace = torch.trace(sigma + sigma_w) - 2.0 * sqrt_trace_component
+
+    mean = torch.sum((m - m_w) ** 2)
+    fd = trace + mean
+    return fd
+
+
+def polynomial_mmd(X, Y):
+    m = X.shape[0]
+    n = Y.shape[0]
+    # compute kernels
+    K_XX = polynomial_kernel(X)
+    K_YY = polynomial_kernel(Y)
+    K_XY = polynomial_kernel(X, Y)
+    # compute mmd distance
+    K_XX_sum = (K_XX.sum() - np.diagonal(K_XX).sum()) / (m * (m - 1))
+    K_YY_sum = (K_YY.sum() - np.diagonal(K_YY).sum()) / (n * (n - 1))
+    K_XY_sum = K_XY.sum() / (m * n)
+    mmd = K_XX_sum + K_YY_sum - 2 * K_XY_sum
+    return mmd
+
+
+
+def get_logits(i3d, videos, device):
+    # assert videos.shape[0] % MAX_BATCH == 0
+    with torch.no_grad():
+        logits = []
+        for i in range(0, videos.shape[0], MAX_BATCH):
+            batch = videos[i:i + MAX_BATCH].to(device)
+            logits.append(i3d(batch))
+        logits = torch.cat(logits, dim=0)
+        return logits
+
+
+def compute_fvd(real, samples, i3d, device=torch.device('cpu')):
+    # i3d.to(device)
+    # real, samples are (N, T, H, W, C) numpy arrays in np.uint8
+    real, samples = preprocess(real, (224, 224)), preprocess(samples, (224, 224))
+    first_embed = get_logits(i3d, real, device)
+    second_embed = get_logits(i3d, samples, device)
+
+    return frechet_distance(first_embed, second_embed)
+

grn/tokenizer/videovae/evaluation/inception.py

@@ -0,0 +1,370 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torchvision import models
+from scipy import linalg
+import numpy as np
+
+try:
+    from torchvision.models.utils import load_state_dict_from_url
+except ImportError:
+    from torch.utils.model_zoo import load_url as load_state_dict_from_url
+
+# Inception weights ported to Pytorch from
+# http://download.tensorflow.org/models/image/imagenet/inception-2015-12-05.tgz
+FID_WEIGHTS_URL = 'https://github.com/mseitzer/pytorch-fid/releases/download/fid_weights/pt_inception-2015-12-05-6726825d.pth'
+
+FID_WEIGHTS_PATH = "../../pretrained/inception/pt_inception-2015-12-05-6726825d.pth"
+
+def calculate_frechet_distance(mu1, sigma1, mu2, sigma2, eps=1e-6):
+    """Numpy implementation of the Frechet Distance.
+    The Frechet distance between two multivariate Gaussians X_1 ~ N(mu_1, C_1)
+    and X_2 ~ N(mu_2, C_2) is
+            d^2 = ||mu_1 - mu_2||^2 + Tr(C_1 + C_2 - 2*sqrt(C_1*C_2)).
+
+    Stable version by Dougal J. Sutherland.
+
+    Params:
+    -- mu1   : Numpy array containing the activations of a layer of the
+               inception net (like returned by the function 'get_predictions')
+               for generated samples.
+    -- mu2   : The sample mean over activations, precalculated on an
+               representative data set.
+    -- sigma1: The covariance matrix over activations for generated samples.
+    -- sigma2: The covariance matrix over activations, precalculated on an
+               representative data set.
+
+    Returns:
+    --   : The Frechet Distance.
+    """
+
+    mu1 = np.atleast_1d(mu1)
+    mu2 = np.atleast_1d(mu2)
+
+    sigma1 = np.atleast_2d(sigma1)
+    sigma2 = np.atleast_2d(sigma2)
+
+    assert mu1.shape == mu2.shape, \
+        'Training and test mean vectors have different lengths'
+    assert sigma1.shape == sigma2.shape, \
+        'Training and test covariances have different dimensions'
+
+    diff = mu1 - mu2
+
+    # Product might be almost singular
+    covmean, _ = linalg.sqrtm(sigma1.dot(sigma2), disp=False)
+    if not np.isfinite(covmean).all():
+        msg = ('fid calculation produces singular product; '
+               'adding %s to diagonal of cov estimates') % eps
+        print(msg)
+        offset = np.eye(sigma1.shape[0]) * eps
+        covmean = linalg.sqrtm((sigma1 + offset).dot(sigma2 + offset))
+
+    # Numerical error might give slight imaginary component
+    if np.iscomplexobj(covmean):
+        if not np.allclose(np.diagonal(covmean).imag, 0, atol=1e-3):
+            m = np.max(np.abs(covmean.imag))
+            raise ValueError('Imaginary component {}'.format(m))
+        covmean = covmean.real
+
+    tr_covmean = np.trace(covmean)
+
+    return (diff.dot(diff) + np.trace(sigma1) +
+            np.trace(sigma2) - 2 * tr_covmean)
+
+class InceptionV3(nn.Module):
+    """Pretrained InceptionV3 network returning feature maps"""
+
+    # Index of default block of inception to return,
+    # corresponds to output of final average pooling
+    DEFAULT_BLOCK_INDEX = 3
+
+    # Maps feature dimensionality to their output blocks indices
+    BLOCK_INDEX_BY_DIM = {
+        64: 0,   # First max pooling features
+        192: 1,  # Second max pooling featurs
+        768: 2,  # Pre-aux classifier features
+        2048: 3  # Final average pooling features
+    }
+
+    def __init__(self,
+                 output_blocks=[DEFAULT_BLOCK_INDEX],
+                 resize_input=True,
+                 normalize_input=True,
+                 requires_grad=False,
+                 use_fid_inception=True):
+        """Build pretrained InceptionV3
+
+        Parameters
+        ----------
+        output_blocks : list of int
+            Indices of blocks to return features of. Possible values are:
+                - 0: corresponds to output of first max pooling
+                - 1: corresponds to output of second max pooling
+                - 2: corresponds to output which is fed to aux classifier
+                - 3: corresponds to output of final average pooling
+        resize_input : bool
+            If true, bilinearly resizes input to width and height 299 before
+            feeding input to model. As the network without fully connected
+            layers is fully convolutional, it should be able to handle inputs
+            of arbitrary size, so resizing might not be strictly needed
+        normalize_input : bool
+            If true, scales the input from range (0, 1) to the range the
+            pretrained Inception network expects, namely (-1, 1)
+        requires_grad : bool
+            If true, parameters of the model require gradients. Possibly useful
+            for finetuning the network
+        use_fid_inception : bool
+            If true, uses the pretrained Inception model used in Tensorflow's
+            FID implementation. If false, uses the pretrained Inception model
+            available in torchvision. The FID Inception model has different
+            weights and a slightly different structure from torchvision's
+            Inception model. If you want to compute FID scores, you are
+            strongly advised to set this parameter to true to get comparable
+            results.
+        """
+        super(InceptionV3, self).__init__()
+
+        self.resize_input = resize_input
+        self.normalize_input = normalize_input
+        self.output_blocks = sorted(output_blocks)
+        self.last_needed_block = max(output_blocks)
+
+        assert self.last_needed_block <= 3, \
+            'Last possible output block index is 3'
+
+        self.blocks = nn.ModuleList()
+
+        if use_fid_inception:
+            inception = fid_inception_v3()
+        else:
+            inception = models.inception_v3(pretrained=True)
+
+        # Block 0: input to maxpool1
+        block0 = [
+            inception.Conv2d_1a_3x3,
+            inception.Conv2d_2a_3x3,
+            inception.Conv2d_2b_3x3,
+            nn.MaxPool2d(kernel_size=3, stride=2)
+        ]
+        self.blocks.append(nn.Sequential(*block0))
+
+        # Block 1: maxpool1 to maxpool2
+        if self.last_needed_block >= 1:
+            block1 = [
+                inception.Conv2d_3b_1x1,
+                inception.Conv2d_4a_3x3,
+                nn.MaxPool2d(kernel_size=3, stride=2)
+            ]
+            self.blocks.append(nn.Sequential(*block1))
+
+        # Block 2: maxpool2 to aux classifier
+        if self.last_needed_block >= 2:
+            block2 = [
+                inception.Mixed_5b,
+                inception.Mixed_5c,
+                inception.Mixed_5d,
+                inception.Mixed_6a,
+                inception.Mixed_6b,
+                inception.Mixed_6c,
+                inception.Mixed_6d,
+                inception.Mixed_6e,
+            ]
+            self.blocks.append(nn.Sequential(*block2))
+
+        # Block 3: aux classifier to final avgpool
+        if self.last_needed_block >= 3:
+            block3 = [
+                inception.Mixed_7a,
+                inception.Mixed_7b,
+                inception.Mixed_7c,
+                nn.AdaptiveAvgPool2d(output_size=(1, 1))
+            ]
+            self.blocks.append(nn.Sequential(*block3))
+
+        for param in self.parameters():
+            param.requires_grad = requires_grad
+
+    def forward(self, inp):
+        """Get Inception feature maps
+
+        Parameters
+        ----------
+        inp : torch.autograd.Variable
+            Input tensor of shape Bx3xHxW. Values are expected to be in
+            range (0, 1)
+
+        Returns
+        -------
+        List of torch.autograd.Variable, corresponding to the selected output
+        block, sorted ascending by index
+        """
+        outp = []
+        x = inp
+
+        if self.resize_input:
+            x = F.interpolate(x,
+                              size=(299, 299),
+                              mode='bilinear',
+                              align_corners=False)
+
+        if self.normalize_input:
+            x = 2 * x - 1  # Scale from range (0, 1) to range (-1, 1)
+
+        for idx, block in enumerate(self.blocks):
+            x = block(x)
+            if idx in self.output_blocks:
+                outp.append(x)
+
+            if idx == self.last_needed_block:
+                break
+
+        return outp
+
+
+def fid_inception_v3():
+    """Build pretrained Inception model for FID computation
+
+    The Inception model for FID computation uses a different set of weights
+    and has a slightly different structure than torchvision's Inception.
+
+    This method first constructs torchvision's Inception and then patches the
+    necessary parts that are different in the FID Inception model.
+    """
+    inception = models.inception_v3(num_classes=1008,
+                                    aux_logits=False,
+                                    pretrained=False)
+    inception.Mixed_5b = FIDInceptionA(192, pool_features=32)
+    inception.Mixed_5c = FIDInceptionA(256, pool_features=64)
+    inception.Mixed_5d = FIDInceptionA(288, pool_features=64)
+    inception.Mixed_6b = FIDInceptionC(768, channels_7x7=128)
+    inception.Mixed_6c = FIDInceptionC(768, channels_7x7=160)
+    inception.Mixed_6d = FIDInceptionC(768, channels_7x7=160)
+    inception.Mixed_6e = FIDInceptionC(768, channels_7x7=192)
+    inception.Mixed_7b = FIDInceptionE_1(1280)
+    inception.Mixed_7c = FIDInceptionE_2(2048)
+
+    state_dict = load_state_dict_from_url(FID_WEIGHTS_URL, progress=True)
+    # state_dict = torch.load(FID_WEIGHTS_PATH)
+    inception.load_state_dict(state_dict)
+    return inception
+
+
+class FIDInceptionA(models.inception.InceptionA):
+    """InceptionA block patched for FID computation"""
+    def __init__(self, in_channels, pool_features):
+        super(FIDInceptionA, self).__init__(in_channels, pool_features)
+
+    def forward(self, x):
+        branch1x1 = self.branch1x1(x)
+
+        branch5x5 = self.branch5x5_1(x)
+        branch5x5 = self.branch5x5_2(branch5x5)
+
+        branch3x3dbl = self.branch3x3dbl_1(x)
+        branch3x3dbl = self.branch3x3dbl_2(branch3x3dbl)
+        branch3x3dbl = self.branch3x3dbl_3(branch3x3dbl)
+
+        # Patch: Tensorflow's average pool does not use the padded zero's in
+        # its average calculation
+        branch_pool = F.avg_pool2d(x, kernel_size=3, stride=1, padding=1,
+                                   count_include_pad=False)
+        branch_pool = self.branch_pool(branch_pool)
+
+        outputs = [branch1x1, branch5x5, branch3x3dbl, branch_pool]
+        return torch.cat(outputs, 1)
+
+
+class FIDInceptionC(models.inception.InceptionC):
+    """InceptionC block patched for FID computation"""
+    def __init__(self, in_channels, channels_7x7):
+        super(FIDInceptionC, self).__init__(in_channels, channels_7x7)
+
+    def forward(self, x):
+        branch1x1 = self.branch1x1(x)
+
+        branch7x7 = self.branch7x7_1(x)
+        branch7x7 = self.branch7x7_2(branch7x7)
+        branch7x7 = self.branch7x7_3(branch7x7)
+
+        branch7x7dbl = self.branch7x7dbl_1(x)
+        branch7x7dbl = self.branch7x7dbl_2(branch7x7dbl)
+        branch7x7dbl = self.branch7x7dbl_3(branch7x7dbl)
+        branch7x7dbl = self.branch7x7dbl_4(branch7x7dbl)
+        branch7x7dbl = self.branch7x7dbl_5(branch7x7dbl)
+
+        # Patch: Tensorflow's average pool does not use the padded zero's in
+        # its average calculation
+        branch_pool = F.avg_pool2d(x, kernel_size=3, stride=1, padding=1,
+                                   count_include_pad=False)
+        branch_pool = self.branch_pool(branch_pool)
+
+        outputs = [branch1x1, branch7x7, branch7x7dbl, branch_pool]
+        return torch.cat(outputs, 1)
+
+
+class FIDInceptionE_1(models.inception.InceptionE):
+    """First InceptionE block patched for FID computation"""
+    def __init__(self, in_channels):
+        super(FIDInceptionE_1, self).__init__(in_channels)
+
+    def forward(self, x):
+        branch1x1 = self.branch1x1(x)
+
+        branch3x3 = self.branch3x3_1(x)
+        branch3x3 = [
+            self.branch3x3_2a(branch3x3),
+            self.branch3x3_2b(branch3x3),
+        ]
+        branch3x3 = torch.cat(branch3x3, 1)
+
+        branch3x3dbl = self.branch3x3dbl_1(x)
+        branch3x3dbl = self.branch3x3dbl_2(branch3x3dbl)
+        branch3x3dbl = [
+            self.branch3x3dbl_3a(branch3x3dbl),
+            self.branch3x3dbl_3b(branch3x3dbl),
+        ]
+        branch3x3dbl = torch.cat(branch3x3dbl, 1)
+
+        # Patch: Tensorflow's average pool does not use the padded zero's in
+        # its average calculation
+        branch_pool = F.avg_pool2d(x, kernel_size=3, stride=1, padding=1,
+                                   count_include_pad=False)
+        branch_pool = self.branch_pool(branch_pool)
+
+        outputs = [branch1x1, branch3x3, branch3x3dbl, branch_pool]
+        return torch.cat(outputs, 1)
+
+
+class FIDInceptionE_2(models.inception.InceptionE):
+    """Second InceptionE block patched for FID computation"""
+    def __init__(self, in_channels):
+        super(FIDInceptionE_2, self).__init__(in_channels)
+
+    def forward(self, x):
+        branch1x1 = self.branch1x1(x)
+
+        branch3x3 = self.branch3x3_1(x)
+        branch3x3 = [
+            self.branch3x3_2a(branch3x3),
+            self.branch3x3_2b(branch3x3),
+        ]
+        branch3x3 = torch.cat(branch3x3, 1)
+
+        branch3x3dbl = self.branch3x3dbl_1(x)
+        branch3x3dbl = self.branch3x3dbl_2(branch3x3dbl)
+        branch3x3dbl = [
+            self.branch3x3dbl_3a(branch3x3dbl),
+            self.branch3x3dbl_3b(branch3x3dbl),
+        ]
+        branch3x3dbl = torch.cat(branch3x3dbl, 1)
+
+        # Patch: The FID Inception model uses max pooling instead of average
+        # pooling. This is likely an error in this specific Inception
+        # implementation, as other Inception models use average pooling here
+        # (which matches the description in the paper).
+        branch_pool = F.max_pool2d(x, kernel_size=3, stride=1, padding=1)
+        branch_pool = self.branch_pool(branch_pool)
+
+        outputs = [branch1x1, branch3x3, branch3x3dbl, branch_pool]
+        return torch.cat(outputs, 1)