Upload Mage-VL: unified codec-native streaming VLM (image+video understanding + proactive gate)

.gitattributes

@@ -33,3 +33,9 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+tokenizer.json filter=lfs diff=lfs merge=lfs -text
+neural_codec/codec_tools/pipeline/__pycache__/generate_codec_patch_smart_resize.cpython-312.pyc filter=lfs diff=lfs merge=lfs -text
+examples/dog.jpg filter=lfs diff=lfs merge=lfs -text
+examples/soccer-broadcast.mp4 filter=lfs diff=lfs merge=lfs -text
+assets/mage-vl-cover.png filter=lfs diff=lfs merge=lfs -text
+assets/mage-vl-framework.png filter=lfs diff=lfs merge=lfs -text

README.md

@@ -1,3 +1,349 @@
 ---
 license: apache-2.0
+library_name: transformers
+pipeline_tag: image-text-to-text
+tags:
+- multimodal
+- vision-language-model
+- mage-vl
+- video-understanding
+- streaming
 ---
+
+<h1 align="center">Mage-VL<br><span style="font-size: 0.55em; font-weight: normal;">An Efficient Codec-Native Streaming Multimodal Foundation Model</span></h1>
+
+<p align="center">
+  <a href="https://microsoft.github.io/Mage"><img alt="Project Page" src="https://img.shields.io/badge/%F0%9F%8C%90-Project%20Page-blue" height="22" /></a>
+  <a href="https://github.com/microsoft/Mage/blob/main/assets/mage_vl_tech_report.pdf"><img alt="arXiv" src="https://img.shields.io/badge/arXiv-Mage--VL-b31b1b" height="22" /></a>
+  <a href="https://github.com/microsoft/Mage"><img src="https://img.shields.io/badge/Code-GitHub-181717?logo=github" alt="GitHub" height="22"></a>
+  <a href="https://huggingface.co/microsoft/Mage-VL"><img alt="Mage-VL" src="https://img.shields.io/badge/%F0%9F%A4%97-Mage--VL-yellow" height="22" /></a>
+  <a href="https://huggingface.co/microsoft/Mage-ViT"><img alt="Mage-ViT" src="https://img.shields.io/badge/%F0%9F%A4%97-Mage--ViT-yellow" height="22" /></a>
+  <a href="https://www.apache.org/licenses/LICENSE-2.0"><img src="https://img.shields.io/badge/License-Apache%202.0-green" alt="License: Apache 2.0" height="22"></a>
+</p>
+
+<div align="center">
+<img src="assets/mage-vl-cover.png" width="100%" alt="Mage-VL">
+</div>
+
+---
+
+**Mage-VL** is a **codec-native, proactive-streaming multimodal foundation model** for image and video understanding, whose visual encoder is trained **entirely from scratch** at a compact **4B** scale. It targets a modern *Moravec's paradox* of VLMs — strong at complex offline reasoning, yet slow and compute-heavy on simple real-time streaming perception. Instead of decoding video into uniformly-sampled frames and pushing a dense grid of patch tokens through a frozen web-pretrained ViT, Mage-VL follows the structure of modern video codecs: it separates a stream into **anchor (I) frames** and **predicted (P) frames**, keeps every anchor patch, and retains only the predicted-frame patches where the codec spends bits — the regions carrying real motion and new detail. This codec-aligned sparsity cuts visual tokens by **over 75%** while preserving spatio-temporal context, yielding **up to 3.5× wall-clock inference speedup** over uniform frame sampling.
+
+The system pairs **two components**:
+
+- **Mage-ViT** — a from-scratch *Codec-ViT* visual encoder that allocates tokens by codec-derived spatio-temporal importance, on a shared `16×16` patch grid with 3D rotary position encoding. It is **codec-agnostic**: the same interface accepts a traditional codec (H.264/AVC, HEVC/H.265) via motion vectors + residual energy, or a neural codec (DCVC-RT) via its learned rate map — no architecture or retraining change.
+- **Qwen3-4B causal decoder** — a Qwen3-4B-Instruct-2507 language backbone (the only pretrained component) that consumes Mage-ViT's variable-length token stream through a lightweight two-layer MLP projector, with a unified interface for images, short/long/ultra-long video, and streaming.
+
+On top of this pair, a **System 1 & System 2 dual-process design** adds proactive streaming inside a single model: a lightweight **cognition gate** (System 1) watches each rolling codec window and stays silent on routine content, invoking the full VLM (System 2) only when a response-worthy event completes — no multi-agent pipeline required.
+
+## ✨ Highlights
+
+- **Codec-native & from scratch.** The entire visual stack is trained from scratch — no billion-scale image-text ViT initialization. The bio-inspired predictive-patch mechanism (I/P frames at `16×16`) cuts visual-token consumption by **over 75%** (**~1/8 or less** of dense frame sampling), letting the model train on videos **8× longer** under the same budget.
+- **Codec-native speedup.** Codec tokenization sets a superior accuracy–efficiency frontier — **up to 3.5× wall-clock inference speedup** over uniform frame sampling at matched accuracy, and the fastest of all compared models on most video benchmarks (single 8×B200 node).
+- **Data-efficient tokenizer.** Trained on only **~100M unlabeled images/videos**, Mage-ViT matches or beats frontier encoders trained on billions of image-text pairs (SigLIP2 @ 10B, MoonViT @ 2B) — e.g. **99.33% on CIFAR-10** and **85.69% on ImageNet** with 256 tokens, showing web-scale pretraining is *not* essential for a strong VLM front-end.
+- **Native-resolution scaling.** Variable-resolution pretraining lets Mage-ViT improve *monotonically* with the token budget (peaking **>96.1% Food-101 / >86.3% ImageNet** at 676 tokens) where fixed-resolution encoders saturate or degrade.
+- **Matched-LLM video gains.** With the 4B Qwen3 backbone held fixed and only the ViT swapped, Mage-VL improves over Qwen3-VL-4B on **every** reported video and temporal-grounding benchmark — largest on localization-heavy tasks (**+22.5 QVHighlight**, +17.1 ActivityNet, +11.0 VSI-Bench, +24.5 VideoEval-Pro).
+- **Strong for its size.** On par with Qwen3-VL-4B on static images, and clearly ahead on video understanding and spatial intelligence (**+11.0** VSI-Bench, **+53.1** CrossPoint, **+5.2** EmbSpatial, **+22.5** QVHighlight).
+- **Proactive streaming, single model.** A frozen-backbone cognition gate delivers low-latency, event-gated commentary; it tops **TimVal / F1 / ROC-AUC / PR-AUC** on SoccerNet streaming and generalizes to real 2026 World Cup broadcasts.
+
+## 📥 Model
+
+A **single checkpoint**, `microsoft/Mage-VL`, is one unified model that **simultaneously** provides image & video understanding **and** the proactive streaming gate — the same weights answer offline image/video questions and drive event-gated commentary. It covers every Mage-VL capability: image understanding, frame-sampled video, traditional H.264/HEVC codec video, neural DCVC-RT codec video, and event-gated streaming. The repository bundles the codec processor, the neural codec package, and the proactive gate weights — no separate understanding, NVC, or streaming checkpoint is required.
+
+We additionally release **`microsoft/Mage-ViT`** — the standalone visual encoder from the two-stage, from-scratch ViT pre-training (cluster-discrimination on ~100M unlabeled image/video frames). This is the **ViT-pre-trained checkpoint only**: it has **not** gone through the joint VLM training with the language model. Use it as a data-efficient, codec-native visual encoder or as a drop-in ViT for your own multimodal training.
+
+| Model | Task | Backbone | Hugging Face |
+| :--- | :--- | :--- | :--- |
+| `Mage-VL` | image & video understanding **+** proactive streaming gate | Mage-ViT + Qwen3-4B-Instruct-2507 | [🤗 microsoft/Mage-VL](https://huggingface.co/microsoft/Mage-VL) |
+| `Mage-ViT` | codec-native visual encoder — ViT pre-training only, no VLM joint training | Codec-ViT (from scratch) | [🤗 microsoft/Mage-ViT](https://huggingface.co/microsoft/Mage-ViT) |
+
+## 🏗️ Architecture
+
+<div align="center">
+<img src="assets/mage-vl-framework.png" width="100%" alt="Mage-VL proactive streaming framework">
+<em>Proactive streaming framework — Mage-ViT incrementally encodes the continuous stream into codec-native visual features shared by the event gate and the causal decoder. The gate scores each rolling window and stays silent on routine content; when it opens, the decoder emits an event-conditioned response.</em>
+</div>
+
+**Mage-ViT** — a from-scratch Codec-ViT visual encoder. On a `16×16` patch grid it keeps all anchor (I) frame patches and only the motion-salient predicted (P) frame patches, cutting visual tokens by over 75% while a shared 3D RoPE preserves spatio-temporal positions.
+
+**Mage-VL** — a unified model where projected visual tokens and text tokens share one causal Qwen3 decoder. Still images become a single spatial block; videos become temporally-ordered codec windows. In **streaming** mode, a lightweight **cognition gate** predicts `p_speak = g(h_t)` per rolling window (over a recurrent streaming memory kept by an event-preserving feature extractor) and triggers generation when `p_speak ≥ τ`; the response is decoded by the frozen base model from a local sliding window of the most recent codec segments, and a text query can be injected at any time.
+
+**Training** — a progressive **five-stage** supervised curriculum (no preference/RL post-training) that produces one unified model:
+
+1. **Multimodal alignment via captions** — ~350M dense image captions + 4.2M short-video captions.
+2. **Instruction tuning + short temporal grounding** — ~54M image-instruction samples + 3.4M 30–180s video captions.
+3. **Temporal-horizon expansion** — medium/long video (LLaVA-Video, TimeLens, VideoChat-Flash, Molmo2) with retained image SFT.
+4. **Codec-native long-context adaptation** — 350K long videos as rolling codec windows (up to 384/768 frames).
+5. **Proactive streaming alignment** — a cognition gate fine-tuned on ~3.3M streaming samples with the visual encoder and LLM kept frozen (only the gate is trained).
+
+The five stages together produce a **single unified model**, `Mage-VL`, that handles image understanding, offline video reasoning, and proactive streaming — no separate variants are shipped.
+
+Two parts of the pipeline apply an **AI4AI** (AI-for-AI) paradigm: (1) dense recaptioning runs through an agentic closed loop where a GPT-5 rubric scorer grades captions and a Copilot coding agent co-designs the prompt *and* harness code (e.g. rendering timestamp overlays) under a human validation gate — improving every downstream OCR/doc/chart/perception benchmark and inspiring SkillOpt-Lite; and (2) Stage-3 uses AI-based diagnostics to decide which video categories, resolutions, and frame counts to train on.
+
+## 📊 Performance
+
+<details>
+<summary><b>Image understanding & spatial intelligence — click to expand</b></summary>
+
+Performance comparison across models. Mage-VL-4B and Qwen3-VL-4B use the same 4B Qwen3 LLM backbone; Phi-4-Multimodal-Instruct (Phi-4-MM, 5.6B) and Phi-4-Reasoning-Vision (Phi-4-R-V, 15B) are reported for reference. `–` = not run. **Bold** = best in row.
+
+| Benchmark | Mage-VL-4B | Qwen3-VL-4B | Phi-4-MM-5.6B | Phi-4-R-V-15B |
+| :--- | :---: | :---: | :---: | :---: |
+| *Document understanding* | | | | |
+| DocVQA-val | **95.14** | 94.69 | 92.79 | 76.20 |
+| InfoVQA-val | **80.33** | 79.50 | 71.84 | 55.41 |
+| AI2D w/ Mask | **83.16** | 81.54 | 81.83 | 82.87 |
+| ChartQA | **84.88** | 83.96 | 83.76 | 83.40 |
+| OCRBench | **81.80** | 81.60 | 81.70 | 73.90 |
+| MultiDocVQA-val | **87.46** | 87.21 | 46.84 | 58.35 |
+| ChartQAPro | **32.57** | 26.79 | 0.13 | 25.38 |
+| TextVQA-val | 77.28 | **80.55** | 39.93 | 76.06 |
+| CC-OCR Doc | 32.25 | **39.69** | 4.99 | 17.65 |
+| *General VQA* | | | | |
+| MMBench-EN-dev | 84.02 | 83.25 | 65.81 | **84.19** |
+| MMBench-CN-dev | **82.04** | 80.58 | 75.17 | 79.47 |
+| MMStar | **67.32** | 62.04 | 61.24 | 59.63 |
+| MME-Perception | **1709.54** | 1703.50 | 1409.66 | 1590.21 |
+| SeedBench (All) | **76.78** | 75.65 | 68.28 | 73.70 |
+| CV-Bench | **87.79** | 85.37 | 57.09 | 81.31 |
+| MME-RealWorld | **66.52** | 63.20 | 32.45 | 57.80 |
+| *Spatial intelligence* | | | | |
+| CV-Bench-2D | **82.13** | 81.00 | 56.12 | 80.11 |
+| CV-Bench-3D | **94.75** | 92.30 | 56.92 | 82.50 |
+| BLINK | **65.11** | 65.10 | 35.24 | 57.80 |
+| EmbSpatial | **82.67** | 77.50 | 41.51 | 72.67 |
+| CrossPoint | **80.00** | 26.90 | 12.20 | 47.73 |
+| CRPE-Relation | 76.12 | **77.70** | 34.60 | 74.46 |
+| SAT | 67.33 | **69.30** | 55.33 | 66.67 |
+
+</details>
+
+<details>
+<summary><b>Video understanding & temporal grounding — click to expand</b></summary>
+
+**Bold** = best in row.
+
+| Benchmark | Mage-VL-4B | Qwen3-VL-4B | Phi-4-MM-5.6B | Phi-4-R-V-15B |
+| :--- | :---: | :---: | :---: | :---: |
+| *Video QA* | | | | |
+| MV-Bench | 65.1 | **66.7** | 44.9 | 49.2 |
+| NextQA | **83.1** | 79.8 | 54.1 | 69.0 |
+| VideoMME | **64.0** | 59.7 | 44.7 | 55.3 |
+| LongVideoBench | **61.3** | 57.7 | 41.14 | 51.2 |
+| LVBench | **41.8** | 39.2 | 25.31 | 34.4 |
+| MLVU-dev | **68.7** | 61.5 | 44.18 | 51.8 |
+| VideoEval-Pro | **45.2** | 20.7 | 14.35 | 16.8 |
+| *Temporal grounding* | | | | |
+| Timelens-Charades | **50.7** | 43.1 | 4.09 | 20.6 |
+| Timelens-ActivityNet | **45.4** | 28.4 | 2.03 | 23.0 |
+| Timelens-QVHighlight | **57.4** | 34.9 | 2.47 | 11.6 |
+| *Spatial reasoning* | | | | |
+| VSI-Bench | **64.3** | 53.3 | 24.09 | 25.5 |
+| *Tracking (J&F)* | | | | |
+| Ref-DAVIS17 | **25.83** | 7.48 | 3.14 | 2.15 |
+| MeViS-ValidU | **22.55** | 3.16 | 10.28 | 1.53 |
+| ReasonVOS | **17.76** | 9.66 | 9.50 | 9.77 |
+| Ref-YT-VOS | **25.57** | 5.28 | 8.64 | 3.85 |
+
+</details>
+
+<details>
+<summary><b>Proactive streaming (SoccerNet) & online video (OVO-Bench) — click to expand</b></summary>
+
+**SoccerNet — response timing** (StreamMind protocol, codec-native inputs, zero-tolerance canvas matching). **Bold** = best in column.
+
+| Method | TriggerAcc | TimVal | F1 | ROC-AUC | PR-AUC |
+| :--- | :---: | :---: | :---: | :---: | :---: |
+| StreamMind | 52.18 | 47.36 | – | – | – |
+| JoyAI-VL-Interaction-9B | **97.98** | 19.25 | 3.55 | 56.26 | 1.68 |
+| **Mage-VL-4B** | 79.21 | **55.54** | **16.35** | **83.14** | **9.30** |
+
+JoyAI's high TriggerAcc comes from predicting silence almost everywhere under SoccerNet's heavy class imbalance, so it collapses on the precision-sensitive metrics; StreamMind is trained *in-distribution* on SoccerNet, whereas Mage-VL is not.
+
+**OVO-Bench** — online video understanding (SimpleStream recent-window protocol, 4 frames @ 1 fps; no streaming-specific fine-tuning). Mage-VL sets a new state-of-the-art overall score **among streaming architectures**. RT-Avg / BT-Avg are the Real-Time Visual Perception / Backward Tracing sub-task averages; Overall is their mean. **Bold** = best model per column (Human is the reference upper bound).
+
+| Model | #Frames | RT-Avg | BT-Avg | Overall |
+| :--- | :---: | :---: | :---: | :---: |
+| Human | – | 93.2 | 92.3 | 92.77 |
+| *Offline video LLMs* | | | | |
+| Qwen2.5-VL-7B | 1 fps | 59.9 | 44.7 | 52.28 |
+| LLaVA-Video-7B | 64 | 63.5 | 40.4 | 51.95 |
+| Qwen3-VL-4B | 64 | 72.8 | **53.1** | 63.00 |
+| *Online / streaming video LLMs* | | | | |
+| VideoLLM-online-8B | 2 fps | 20.8 | 17.7 | 19.26 |
+| Flash-VStream-7B | 1 fps | 28.4 | 27.4 | 27.90 |
+| Dispider-7B | 1 fps | 54.6 | 36.1 | 45.35 |
+| TimeChat-Online-7B | 1 fps | 61.9 | 41.7 | 51.80 |
+| StreamForest-7B | 1 fps | 61.2 | 52.0 | 56.60 |
+| Streamo-7B | 1 fps | 66.0 | 46.1 | 56.05 |
+| HERMES-7B<sup>†</sup> | 1 fps | 69.0 | 49.4 | 59.20 |
+| JoyAI-VL-Interaction-9B | 1 fps | 68.4 | 48.6 | 58.50 |
+| **Mage-VL-4B** | 1 fps | **79.84** | 48.15 | **64.00** |
+
+<sub><sup>†</sup> HERMES = Qwen2.5-VL-7B + HERMES (4K tokens). Baseline results and table structure follow SimpleStream.</sub>
+
+</details>
+
+## 🔬 Key Findings
+
+Beyond the model, the report distills **seven empirical findings** for efficient multimodal training:
+
+1. **Web-scale pretraining is not essential.** A from-scratch backbone on ~100M unlabeled frames matches encoders trained on billions of image-text pairs.
+2. **Variable-resolution pretraining scales monotonically.** Quality keeps improving with the visual-token budget instead of saturating/degrading like fixed-resolution encoders.
+3. **Codec-native tokenization sets a better accuracy–efficiency frontier** — up to **3.5× wall-clock inference speedup** over uniform frame sampling.
+4. **Explicit VideoQA SFT is redundant.** Dense video *captions* + standard image SFT are sufficient for strong zero-shot VideoQA.
+5. **Motion–spatial synergy.** Dynamic video training substantially improves static 2D/3D spatial reasoning.
+6. **AI4AI data pipeline.** Agentic closed-loop feedback + prompt/code co-design systematically lift caption quality and downstream scores (inspired SkillOpt-Lite).
+7. **Zero-Vision SFT for multimodal RL.** Bypassing visual SFT in favor of pure-text reasoning SFT unlocks stronger multimodal RL — a compute-efficient path.
+
+## 🚀 Quick Start
+
+A single checkpoint, `microsoft/Mage-VL`, covers every capability below.
+
+| Capability | Script | How to run |
+|---|---|---|
+| Image understanding | `inference.py` | `--mode offline --image` |
+| Frame-sampled video | `inference.py` | `--mode offline --video --video-backend frames` |
+| Traditional H.264/HEVC codec video | `inference.py` | `--mode offline --video --video-backend codec --codec-engine traditional` |
+| Neural DCVC-RT codec video | `inference.py` | `--mode offline --video --video-backend codec --codec-engine neural` |
+| Online image / video (SGLang) | `inference.py` | `--mode online … --base-url <server>` |
+| Event-gated streaming commentary | `inference_streaming.py` | in the [GitHub repo](https://github.com/microsoft/Mage/tree/main/mage_vl) |
+
+### Installation
+
+For offline Transformers inference:
+
+```bash
+pip install "transformers>=5.7" accelerate pillow torch torchvision \
+  opencv-python codec-video-prep
+```
+
+Codec-based video inference also requires `ffmpeg` and `ffprobe` on `PATH`.
+
+### Examples
+
+Two sample inputs ship with this repository:
+
+| Input | Question | Content |
+|---|---|---|
+| [`examples/dog.jpg`](examples/dog.jpg) | Describe this image in detail. | Photo of a dog sitting in front of a patterned rug |
+| [`examples/soccer-broadcast.mp4`](examples/soccer-broadcast.mp4) | Describe this video. | 30s, 960×540 football broadcast clip |
+
+### Offline inference
+
+Download [`inference.py`](inference.py). Offline mode loads the checkpoint with `AutoModelForCausalLM.from_pretrained` and supports images, frame sampling, and both codec engines:
+
+```bash
+# image
+python inference.py --mode offline --image examples/dog.jpg \
+  --question "Describe this image in detail."
+```
+
+> The image depicts a dog sitting on a patterned rug. The dog appears to be a
+> medium-sized breed with a thick, fluffy coat. Its fur is primarily white with
+> patches of black and brown. The dog's ears are perked up, and it has a calm and
+> attentive expression. [...]
+
+```bash
+# video — uniform frame sampling
+python inference.py --mode offline --video examples/soccer-broadcast.mp4 \
+  --video-backend frames --num-frames 32 \
+  --question "Describe this video."
+```
+
+> The video opens with a man in a black polo shirt, sporting a short haircut,
+> standing in a stadium. He is holding a yellow microphone with the BBC Sport
+> logo on it. The background reveals a large crowd of spectators. [...]
+
+```bash
+# video — traditional codec (HEVC/H.264)
+python inference.py --mode offline --video examples/soccer-broadcast.mp4 \
+  --video-backend codec --codec-engine traditional --num-frames 32 \
+  --question "Describe this video."
+```
+
+> The video opens with a BBC Sport broadcast, featuring a presenter in a black
+> shirt holding a yellow microphone. The background reveals a packed stadium,
+> with the scoreboard displaying "ENG 1 ARG 2 FT", indicating the final score of
+> the match. [...]
+
+```bash
+# video — neural codec (DCVC-RT)
+python inference.py --mode offline --video examples/soccer-broadcast.mp4 \
+  --video-backend codec --codec-engine neural --num-frames 32 \
+  --question "Describe this video."
+```
+
+> The video opens with a BBC Sport broadcast, featuring a presenter standing in a
+> stadium filled with spectators. The presenter, dressed in a black shirt, holds
+> a yellow BBC Sport microphone and wears a black earpiece. [...]
+
+### Online inference
+
+Online mode talks to an OpenAI-compatible SGLang server. **First** build and launch the server with the Mage-VL SGLang branch (building it needs `protobuf-compiler` and a Rust toolchain):
+
+```bash
+sudo apt-get update && sudo apt-get install -y protobuf-compiler
+curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs \
+  | sh -s -- -y --profile minimal --default-toolchain 1.90.0
+source "$HOME/.cargo/env"
+
+git clone -b feat/mage-vl https://github.com/kcz358/sglang
+cd sglang
+pip install -e 'python[all]'
+python -m sglang.launch_server \
+  --model-path microsoft/Mage-VL \
+  --trust-remote-code
+```
+
+**Then** send an image or sampled video frames to the running server:
+
+```bash
+pip install openai
+
+python inference.py --mode online --image examples/dog.jpg \
+  --question "Describe this image in detail." \
+  --base-url http://localhost:30000/v1
+
+python inference.py --mode online --video examples/soccer-broadcast.mp4 \
+  --num-frames 32 \
+  --question "Describe this video." \
+  --base-url http://localhost:30000/v1
+```
+
+Use `--model`, `--max-new-tokens`, and `--api-key` to override their defaults.
+
+### Streaming inference
+
+`streammind_gate.safetensors` in this repository holds the event gate. Streaming inference splits a video into non-overlapping segments, stays silent on routine content, and generates a caption only when a response-worthy event is detected. Run it with `inference_streaming.py` from the [GitHub repository](https://github.com/microsoft/Mage/tree/main/mage_vl):
+
+```bash
+python inference_streaming.py \
+  --video examples/soccer-broadcast.mp4 \
+  --video_backend codec \
+  --segment_sec 8
+```
+
+```text
+[t=0.0-8.0s] gate=silence (p=0.19)
+[t=8.0-16.0s] gate=response (p=0.55) -> The video features a live sports broadcast from BBC Sport, set in a large stadium filled with spectators. The broadcast focuses on a football match between England and Argentina, with the score displayed as England 1, Argentina 2. [...]
+[t=16.0-24.0s] gate=response (p=0.73) -> The video features a sports broadcast set in a large stadium filled with spectators. Four commentators are gathered around a table with a 'BBC Sport' logo, each holding a yellow microphone. [...]
+[t=24.0-30.0s] gate=silence (p=0.31)
+```
+
+The gate is trained on codec inputs, so `--video_backend codec` is the intended setting. Use `--video_backend frames` for direct frame sampling. Additional controls include `--num_frames`, `--cur_fps`, `--max_segments`, `--max_new_tokens`, `--gate_threshold`, and `--attn_impl`.
+
+## 📝 Citation
+
+```bibtex
+@article{mage2026magevl,
+  title={Mage-VL: An Efficient Codec-Native Streaming Multimodal Foundation Model},
+  author={Yang, Senqiao and Zhang, Kaichen and Jia, Zhaoyang and Guo, Jinghao and Shen, Yifei and Zhang, Xinjie and Zhang, Xiaoyi and Wang, Haoqing and Li, Xiao and An, Xiang and Xie, Yin and Liu, Zhening and Guo, Xun and Li, Jiahao and Zheng, Shicheng and Wang, Jinglu and Guo, Zongyu and Xie, Wenxuan and Zheng, Zihan and Luo, Yuxuan and Li, Bin and Lu, Yan},
+  journal={arXiv preprint},
+  year={2026}
+}
+```
+
+## 📄 License
+
+Mage-VL is released under the [Apache-2.0 License](https://www.apache.org/licenses/LICENSE-2.0).

added_tokens.json

@@ -0,0 +1,24 @@
+{
+  "</tool_call>": 151658,
+  "<tool_call>": 151657,
+  "<|box_end|>": 151649,
+  "<|box_start|>": 151648,
+  "<|endoftext|>": 151643,
+  "<|file_sep|>": 151664,
+  "<|fim_middle|>": 151660,
+  "<|fim_pad|>": 151662,
+  "<|fim_prefix|>": 151659,
+  "<|fim_suffix|>": 151661,
+  "<|im_end|>": 151645,
+  "<|im_start|>": 151644,
+  "<|image_pad|>": 151655,
+  "<|object_ref_end|>": 151647,
+  "<|object_ref_start|>": 151646,
+  "<|quad_end|>": 151651,
+  "<|quad_start|>": 151650,
+  "<|repo_name|>": 151663,
+  "<|video_pad|>": 151656,
+  "<|vision_end|>": 151653,
+  "<|vision_pad|>": 151654,
+  "<|vision_start|>": 151652
+}

chat_template.jinja

@@ -0,0 +1,7 @@
+{% set image_count = namespace(value=0) %}{% set video_count = namespace(value=0) %}{% for message in messages %}{% if loop.first and message['role'] != 'system' %}<|im_start|>system
+You are a helpful assistant.<|im_end|>
+{% endif %}<|im_start|>{{ message['role'] }}
+{% if message['content'] is string %}{{ message['content'] }}<|im_end|>
+{% else %}{% for content in message['content'] %}{% if content['type'] == 'image' or 'image' in content or 'image_url' in content %}{% set image_count.value = image_count.value + 1 %}{% if add_vision_id %}Picture {{ image_count.value }}: {% endif %}<|vision_start|><|image_pad|><|vision_end|>{% elif content['type'] == 'video' or 'video' in content %}{% set video_count.value = video_count.value + 1 %}{% if add_vision_id %}Video {{ video_count.value }}: {% endif %}<|vision_start|><|video_pad|><|vision_end|>{% elif 'text' in content %}{{ content['text'] }}{% endif %}{% endfor %}<|im_end|>
+{% endif %}{% endfor %}{% if add_generation_prompt %}<|im_start|>assistant
+{% endif %}

codec_video_processing_mage_vl.py

@@ -0,0 +1,591 @@
+"""Codec-based video preprocessing for MageVL (trust_remote_code).
+
+This module is the codec analogue of ``video_processing_magevl.py``.
+It is invoked when a user calls::
+
+    processor(messages=..., video_backend="codec", max_pixels=...)
+
+and is responsible for:
+
+  - Decoding the video and assembling canvas images via ``cv-preinfer``
+    (PyPI: ``codec-video-prep``, requires ``ffmpeg`` on PATH).
+  - Running the bundled ``Qwen2VLImageProcessor`` on those canvases with a
+    pixel budget that is *aligned* to the canvas dimensions (so the
+    smart_resize step never desynchronises ``image_grid_thw`` from the
+    codec-emitted ``src_patch_position`` array).
+  - Producing the per-patch ``patch_positions`` table that
+    ``modeling_magevl.py`` reads for the 2D-MRoPE block layout.
+
+The result is a ``BatchFeature``-shaped dict containing the same keys that
+the frame-sampling video path produces (``pixel_values`` /
+``image_grid_thw`` / ``patch_positions``), so downstream
+``modeling_magevl.py`` consumes it without changes.
+"""
+
+from __future__ import annotations
+
+import hashlib
+import json
+import os
+import shutil
+import subprocess
+import sys
+import tempfile
+import warnings
+from dataclasses import dataclass, field, fields as _dc_fields
+from pathlib import Path
+from typing import Optional
+
+try:
+    import fcntl
+except ImportError:
+    fcntl = None  # type: ignore
+
+import numpy as np
+import torch
+from PIL import Image
+
+
+VISION_START = "<|vision_start|>"
+VISION_END = "<|vision_end|>"
+IMAGE_PAD = "<|image_pad|>"
+
+
+# ----------------------------------------------------------------- config
+
+@dataclass
+class DcvcConfig:
+    """DCVC-RT neural-codec knobs (nested under ``CodecConfig.dcvc``).
+
+    Only used when ``CodecConfig.engine == "dcvc-rt"``. Checkpoints default to
+    the ``DCVC_INTRA_TAR`` / ``DCVC_INTER_TAR`` env vars (environment-specific,
+    so better left out of the model config).
+    """
+    qp: int = 42                          # quantization (0-63); scoring only
+    reset_interval: int = 64              # feature-adaptor reset period (frames)
+    intra_period: int = -1                # -1 => single I-frame at t=0, rest P
+    max_side: int = 0                     # >0 downscales frames before DCVC scoring;
+                                          # 0 => full-resolution scoring (matches the
+                                          # hevc baseline, which reads h264 bits at
+                                          # native stream resolution). Faithful but
+                                          # slower to encode.
+    seq_len_frames: int = 0               # 0 => CodecConfig.num_sampled_frames()
+    patch: int = 16                       # MUST match image processor patch_size
+    canvas_token_side: Optional[int] = None  # None => derive from max_pixels
+    # ---- readiness grouping knobs (used by the bundled process_video_bitcost_
+    # readiness pipeline). Defaults reproduce EXACTLY the offline/cluster DCVC
+    # sweep + hevc baseline generation params, so the release DCVC path yields
+    # the same assets we evaluated. (These intentionally differ from the release
+    # hevc _run_cv_preinfer knobs, e.g. max_group_frames=128 vs CodecConfig's 64.)
+    readiness_sum_threshold_mode: str = "auto"
+    readiness_coverage_bins: int = 3
+    readiness_delta_ratio: float = 0.05
+    bitcost_grid: str = "sub"             # DCVC bitmap is at 16px granularity
+    bitcost_pct: int = 99
+    decode_backsearch_max: int = 16
+    max_group_frames: int = 128
+    intra_ckpt: str = field(default_factory=lambda: os.getenv("DCVC_INTRA_TAR", ""))
+    inter_ckpt: str = field(default_factory=lambda: os.getenv("DCVC_INTER_TAR", ""))
+    device: str = field(default_factory=lambda: os.getenv("DCVC_DEVICE", "cuda:0"))
+    pkg_dir: str = ""                     # dir with dcvc_rt_engine.py + dcvc_readiness_gen.py + codec_tools/
+
+
+@dataclass
+class CodecConfig:
+    """All knobs for the codec preprocessing pipeline.
+
+    ``max_pixels`` is shared with the image_processor / video_processor pixel
+    budget. The processor sets it from the user's ``max_pixels=`` kwarg, so
+    canvas size and HF smart_resize budget stay consistent.
+    """
+
+    target_canvas: int = 32
+    group_size: int = 32
+    images_per_group: int = 4
+    patch: int = 14
+    max_pixels: int = 150000
+    min_group_frames: int = 8
+    max_group_frames: int = 64
+    spatial_mask_mode: str = "off"
+    cache_root: Path = field(default_factory=lambda: Path(
+        os.getenv(
+            "ONLINE_CODEC_CACHE_DIR",
+            os.path.join(
+                os.getenv("HF_HOME", os.path.expanduser("~/.cache/huggingface")),
+                "online_codec",
+            ),
+        )
+    ))
+    timeout_seconds: int = int(os.getenv("ONLINE_CODEC_TIMEOUT", "7200"))
+
+    # ---- codec engine selection ------------------------------------------
+    # "hevc"    : HEVC/h264 bit-cost readiness (external cv-preinfer binary; the
+    #             default). "cv-preinfer" is accepted as a legacy alias.
+    # "dcvc-rt" : DCVC-RT neural-codec bit-cost readiness (self-contained, uses
+    #             the bundled ``neural_codec/`` package + DCVC-RT checkpoints).
+    # Selectable per call via ``codec_config={"engine": "dcvc-rt", ...}``.
+    engine: str = field(default_factory=lambda: os.getenv("CODEC_ENGINE", "hevc"))
+
+    # DCVC-RT knobs, nested. In the model config / codec_config this is a plain
+    # dict (``"dcvc": {"qp": 21, ...}``); __post_init__ coerces it to DcvcConfig.
+    dcvc: "DcvcConfig" = field(default_factory=DcvcConfig)
+
+    def __post_init__(self):
+        if isinstance(self.dcvc, dict):
+            # Filter to known DcvcConfig fields — preprocessor_config.json's
+            # codec.dcvc also carries selection knobs (per_frame_cap_ratio,
+            # bottom_atten, threshold_scale, ...) read directly by the readiness
+            # pipeline via codec_dcvc_config, which are not DcvcConfig fields.
+            _known = {f.name for f in _dc_fields(DcvcConfig)}
+            self.dcvc = DcvcConfig(**{k: v for k, v in self.dcvc.items() if k in _known})
+
+    def validate(self) -> None:
+        if self.target_canvas <= 0:
+            raise ValueError("CodecConfig.target_canvas must be > 0")
+        if self.target_canvas % self.images_per_group != 0:
+            raise ValueError(
+                "CodecConfig.target_canvas must be divisible by images_per_group"
+            )
+        if self.group_size % self.images_per_group != 0:
+            raise ValueError(
+                "CodecConfig.group_size must be divisible by images_per_group"
+            )
+
+    def num_sampled_frames(self) -> int:
+        return (self.target_canvas // self.images_per_group) * self.group_size
+
+
+# ---------------------------------------------------------- text/position
+
+def _format_timestamp(seconds: float, decimals: int) -> str:
+    return f"<{seconds:.{decimals}f} seconds>"
+
+
+def convert_positions_to_block_layout(
+    positions: torch.Tensor, t: int, h: int, w: int, spatial_merge_size: int = 2,
+) -> torch.Tensor:
+    """Reorder a (T*H*W, 3) patch position table into 2D-MRoPE block layout."""
+    sms = int(spatial_merge_size)
+    if sms == 1:
+        return positions
+    total = int(t) * int(h) * int(w)
+    indices = torch.arange(total, device=positions.device).view(t, h, w)
+    h_m, w_m = int(h) // sms, int(w) // sms
+    indices = (
+        indices.view(t, h_m, sms, w_m, sms)
+        .permute(0, 1, 3, 2, 4).contiguous().view(total)
+    )
+    return positions[indices]
+
+
+def codec_positions_for_processor(
+    src_positions: np.ndarray, image_grid_thw: torch.Tensor, device: torch.device,
+) -> torch.Tensor:
+    positions = torch.from_numpy(src_positions).long().to(device)
+    expected_total = int(image_grid_thw.prod(dim=1).sum().item())
+    if expected_total != positions.shape[0]:
+        raise ValueError(
+            "codec patch position length mismatch: "
+            f"thw_total={expected_total}, positions={positions.shape[0]}"
+        )
+    chunks, offset = [], 0
+    for row in image_grid_thw:
+        t, h, w = int(row[0]), int(row[1]), int(row[2])
+        n = t * h * w
+        chunks.append(convert_positions_to_block_layout(positions[offset: offset + n], t, h, w))
+        offset += n
+    return torch.cat(chunks, dim=0)
+
+
+def _timestamp_runs(
+    patch_positions: torch.Tensor, fps: float, decimals: int, spatial_merge_size: int = 2,
+) -> list[tuple[str, int]]:
+    t_values = patch_positions[:, 0]
+    unique_t, counts = torch.unique_consecutive(t_values, return_counts=True)
+    merge_factor = int(spatial_merge_size) ** 2
+    runs = []
+    for t_val, count in zip(unique_t.tolist(), counts.tolist()):
+        if int(t_val) < 0:
+            continue
+        token_count = int(count) // merge_factor
+        if token_count <= 0:
+            continue
+        runs.append((_format_timestamp(float(t_val) / float(fps), decimals), token_count))
+    return runs
+
+
+def rewrite_text_with_codec_positions(
+    text: str, patch_positions: torch.Tensor, fps: float, decimals: int,
+) -> str:
+    """Replace the vision span in a chat-template string with codec-aware tokens."""
+    parts = []
+    for timestamp, token_count in _timestamp_runs(patch_positions, fps, decimals):
+        parts.extend([timestamp, VISION_START, IMAGE_PAD * token_count, VISION_END, "\n"])
+    vision_text = "".join(parts)
+    first_vs, last_ve = text.find(VISION_START), text.rfind(VISION_END)
+    if first_vs == -1 or last_ve == -1:
+        return text
+    tail_start = last_ve + len(VISION_END)
+    if tail_start < len(text) and text[tail_start] == "\n":
+        tail_start += 1
+    return text[:first_vs] + vision_text + text[tail_start:]
+
+
+def drop_padding_canvases(
+    images: list[Image.Image], src_positions: np.ndarray,
+) -> tuple[list[Image.Image], np.ndarray, int]:
+    """Drop fully-padding canvases (all-negative timestamps) and their patches."""
+    n_canvas = len(images)
+    if n_canvas == 0:
+        return images, src_positions, 0
+    total_patches = src_positions.shape[0]
+    if total_patches % n_canvas != 0:
+        raise ValueError(
+            f"src_positions length {total_patches} not divisible by canvas count {n_canvas}"
+        )
+    ppc = total_patches // n_canvas
+    positions = src_positions.reshape(n_canvas, ppc, 3)
+    canvas_t = positions[..., 0]
+    keep_mask = (canvas_t >= 0).any(axis=1)
+    if bool((keep_mask & ~((canvas_t >= 0).all(axis=1))).any()):
+        raise ValueError("encountered half-padding canvas; padding is expected to be canvas-granular")
+    dropped = int(n_canvas - int(keep_mask.sum()))
+    if dropped == 0:
+        return images, src_positions, 0
+    kept_images = [img for img, keep in zip(images, keep_mask.tolist()) if keep]
+    kept_positions = positions[keep_mask].reshape(-1, 3)
+    return kept_images, kept_positions, dropped
+
+
+# ------------------------------------------------------- cv-preinfer driver
+
+def _get_video_total_frames(video_url: str) -> int:
+    import cv2
+    cap = cv2.VideoCapture(video_url)
+    try:
+        total = int(cap.get(cv2.CAP_PROP_FRAME_COUNT) or 0)
+    finally:
+        cap.release()
+    return max(1, total)
+
+
+def _cache_dir_for(video_url: str, cfg: CodecConfig) -> Path:
+    raw = (
+        f"{video_url}|eng={cfg.engine}|tc={cfg.target_canvas}|gs={cfg.group_size}"
+        f"|ipg={cfg.images_per_group}|patch={cfg.patch}"
+        f"|mp={cfg.max_pixels}|mask={cfg.spatial_mask_mode}"
+    )
+    if cfg.engine == "dcvc-rt":
+        d = cfg.dcvc
+        raw += (
+            f"|dqp={d.qp}|drst={d.reset_interval}|dip={d.intra_period}"
+            f"|dms={d.max_side}|dpatch={d.patch}|dseq={d.seq_len_frames}"
+            f"|dcts={d.canvas_token_side}"
+        )
+        # The DCVC subprocess reads the selection-tuning knobs (threshold_scale,
+        # bottom_atten / bottom_band, per_frame_cap_ratio, readiness_*, random_*, ...)
+        # straight from preprocessor_config.json's codec.dcvc — not via cfg.dcvc — so
+        # fold that whole block into the key. Otherwise editing a tuning knob leaves
+        # the key unchanged and process_codec_video returns stale cached canvases.
+        try:
+            _cfg_path = os.path.join(os.path.dirname(os.path.abspath(__file__)),
+                                     "preprocessor_config.json")
+            with open(_cfg_path, "r", encoding="utf-8") as _f:
+                _dcvc_blk = (json.load(_f).get("codec", {}) or {}).get("dcvc", {}) or {}
+            raw += "|dcvcblk=" + json.dumps(_dcvc_blk, sort_keys=True, separators=(",", ":"))
+        except (OSError, ValueError):
+            pass
+    key = hashlib.md5(raw.encode()).hexdigest()
+    return cfg.cache_root / f"{Path(video_url).stem}_{key}"
+
+
+def _load_codec_result(out_dir: Path) -> dict:
+    with open(out_dir / "meta.json", "r", encoding="utf-8") as f:
+        meta = json.load(f)
+    canvas_files = meta.get("canvas_files")
+    if not canvas_files:
+        for ext in ("npy", "jpg", "png"):
+            hits = sorted(p.name for p in out_dir.glob(f"canvas_*.{ext}"))
+            if hits:
+                canvas_files = hits
+                break
+        canvas_files = canvas_files or []
+    images = []
+    for name in canvas_files:
+        fp = out_dir / name
+        if name.endswith(".npy"):
+            images.append(Image.fromarray(np.load(fp)))
+        else:
+            images.append(Image.open(fp).convert("RGB"))
+    src_positions = np.load(out_dir / "src_patch_position.npy")
+    fps = float(meta.get("fps") or 30.0)
+    return {"images": images, "src_positions": src_positions, "fps": fps,
+            "out_dir": str(out_dir), "meta": meta}
+
+
+def _run_cv_preinfer(video_url: str, out_dir: Path, cfg: CodecConfig) -> dict:
+    bin_name = os.environ.get("CV_PREINFER_BIN", "cv-preinfer")
+    if shutil.which(bin_name) is None and not os.path.isfile(bin_name):
+        raise RuntimeError(
+            f"engine='hevc' (traditional codec) needs the external '{bin_name}' binary on "
+            "PATH, which is not bundled with this repo. Install it and add it to PATH (or set "
+            "CV_PREINFER_BIN to its path), or use the neural codec (engine='dcvc-rt')."
+        )
+    tmp_dir = Path(tempfile.mkdtemp(dir=str(cfg.cache_root), prefix=f".tmp_{out_dir.name[:48]}_"))
+    num_sampled = min(cfg.num_sampled_frames(), _get_video_total_frames(video_url))
+    cmd = [
+        bin_name, "--video", video_url, "--out_dir", str(tmp_dir),
+        "--num_sampled_frames", str(num_sampled),
+        "--grouping_mode", "readiness",
+        "--group_size", str(cfg.group_size),
+        "--images_per_group", str(cfg.images_per_group),
+        "--patch", str(cfg.patch),
+        "--max_pixels", str(cfg.max_pixels),
+        "--readiness_sum_threshold", "0",
+        "--min_group_frames", str(cfg.min_group_frames),
+        "--max_group_frames", str(cfg.max_group_frames),
+        "--avoid_keyframes",
+        "--canvas_format", "jpg",
+    ]
+    try:
+        result = subprocess.run(cmd, text=True, capture_output=True, timeout=cfg.timeout_seconds)
+        if result.returncode != 0:
+            detail = (result.stderr or result.stdout)[-2000:]
+            raise RuntimeError(f"online codec failed rc={result.returncode}: {detail}")
+        if out_dir.exists():
+            shutil.rmtree(out_dir)
+        tmp_dir.rename(out_dir)
+    except Exception:
+        shutil.rmtree(tmp_dir, ignore_errors=True)
+        raise
+    return _load_codec_result(out_dir)
+
+
+# ------------------------------------------------------- dcvc-rt driver
+def _run_dcvc_rt(video_url: str, out_dir: Path, cfg: CodecConfig) -> dict:
+    """DCVC-RT analogue of ``_run_cv_preinfer``. Generates the codec asset dir
+    (canvas_*.jpg + src_patch_position.npy + meta.json) by running the SAME
+    readiness pipeline the hevc/offline path uses (``process_video_bitcost_
+    readiness``), swapping ONLY the per-frame score source from h264 block bits
+    to the DCVC-RT neural-codec bit-cost bitmap.
+
+    Faithful by construction: it shells out to the bundled
+    ``neural_codec/dcvc_readiness_gen.py`` — the exact script used for the offline /
+    cluster DCVC generation — with the same GEN_PARAMS, so the release
+    reproduces the evaluated assets. ffmpeg/ffprobe must be on PATH (same
+    requirement as the hevc path); the DCVC engine + checkpoints are loaded by
+    the subprocess.
+    """
+    d = cfg.dcvc
+    module_dir = os.path.dirname(os.path.abspath(__file__))
+    flat_pkg = os.path.join(module_dir, "neural_codec")
+    nested_pkg = os.path.join(os.path.dirname(module_dir), "neural_codec")
+    pkg = d.pkg_dir or (flat_pkg if os.path.isdir(flat_pkg) else nested_pkg)
+    gen = os.path.join(pkg, "dcvc_readiness_gen.py")
+    if not os.path.exists(gen):
+        raise RuntimeError(f"dcvc-rt: bundled generator not found at {gen}")
+    # Checkpoints default to the copies bundled in neural_codec/ (self-contained);
+    # DCVC_INTRA_TAR / DCVC_INTER_TAR still override.
+    if not d.intra_ckpt:
+        d.intra_ckpt = os.path.join(pkg, "dcvc_rt_intra.tar")
+    if not d.inter_ckpt:
+        d.inter_ckpt = os.path.join(pkg, "dcvc_rt_inter.tar")
+    if not os.path.exists(d.intra_ckpt) or not os.path.exists(d.inter_ckpt):
+        raise ValueError(
+            f"engine='dcvc-rt' checkpoints not found (intra={d.intra_ckpt!r}, "
+            f"inter={d.inter_ckpt!r}); they ship in neural_codec/ — or set "
+            "DCVC_INTRA_TAR / DCVC_INTER_TAR."
+        )
+    # The DCVC-RT source is bundled at neural_codec/DCVC and loaded by dcvc_rt_engine
+    # (no env var); sanity-check it so the subprocess fails fast with a clear message.
+    if not os.path.isdir(os.path.join(pkg, "DCVC", "src")):
+        raise ValueError(
+            f"engine='dcvc-rt' bundled DCVC source missing at "
+            f"{os.path.join(pkg, 'DCVC')!r} (expected a 'src/' dir); the "
+            "neural_codec/DCVC/ folder looks incomplete."
+        )
+
+    tmp_dir = Path(tempfile.mkdtemp(dir=str(cfg.cache_root), prefix=f".tmp_{out_dir.name[:48]}_"))
+    num_sampled = min(cfg.num_sampled_frames(), _get_video_total_frames(video_url))
+    # Readiness CLI args — mirror the offline/cluster GEN_PARAMS exactly so the
+    # generated assets match what we evaluated. Shared dims from cfg; DCVC +
+    # readiness knobs from cfg.dcvc (patch=16, max_group_frames=128, auto
+    # threshold, etc. — intentionally differing from the hevc _run_cv_preinfer).
+    cmd = [
+        sys.executable, gen,
+        "--video", video_url, "--out_dir", str(tmp_dir),
+        "--num_sampled_frames", str(num_sampled),
+        "--grouping_mode", "readiness",
+        "--readiness_sum_threshold_mode", str(d.readiness_sum_threshold_mode),
+        "--group_size", str(cfg.group_size),
+        "--images_per_group", str(cfg.images_per_group),
+        "--patch", str(d.patch),
+        "--max_pixels", str(cfg.max_pixels),
+        "--min_group_frames", str(cfg.min_group_frames),
+        "--max_group_frames", str(d.max_group_frames),
+        "--readiness_coverage_bins", str(d.readiness_coverage_bins),
+        "--readiness_delta_ratio", str(d.readiness_delta_ratio),
+        "--bitcost_grid", str(d.bitcost_grid),
+        "--bitcost_pct", str(d.bitcost_pct),
+        "--decode_backsearch_max", str(d.decode_backsearch_max),
+        "--canvas_format", "jpg",
+    ]
+    # Only DCVC engine *infra* travels via env (checkpoints / device / repo roots).
+    # Selection + scoring params (qp, reset_interval, intra_period, max_side, and
+    # the readiness knobs) are read by dcvc_readiness_gen from preprocessor_config.
+    # json's codec.dcvc via codec_dcvc_config — NOT from env.
+    env = dict(os.environ)
+    env.update({
+        "DCVC_INTRA_TAR": d.intra_ckpt, "DCVC_INTER_TAR": d.inter_ckpt,
+        "DCVC_DEVICE": d.device,
+        "DCVC_ENGINE_DIR": pkg, "DCVC_REPO_DIR": pkg,
+    })
+    try:
+        result = subprocess.run(cmd, text=True, capture_output=True,
+                                timeout=cfg.timeout_seconds, env=env)
+        if result.returncode != 0:
+            detail = (result.stderr or result.stdout)[-2000:]
+            raise RuntimeError(f"dcvc-rt gen failed rc={result.returncode}: {detail}")
+        if not (tmp_dir / "meta.json").exists():
+            detail = (result.stderr or result.stdout)[-2000:]
+            raise RuntimeError(f"dcvc-rt gen produced no meta.json: {detail}")
+        # Tag the score source into meta so downstream can distinguish it.
+        try:
+            with open(tmp_dir / "meta.json", "r", encoding="utf-8") as f:
+                meta = json.load(f)
+            meta.update({"engine": "dcvc-rt", "score_source": "dcvc-rt",
+                         "qp": d.qp, "reset_interval": d.reset_interval,
+                         "dcvc_max_side": d.max_side})
+            with open(tmp_dir / "meta.json", "w", encoding="utf-8") as f:
+                json.dump(meta, f, ensure_ascii=False, indent=2)
+        except Exception:
+            pass
+        if out_dir.exists():
+            shutil.rmtree(out_dir)
+        tmp_dir.rename(out_dir)
+    except Exception:
+        shutil.rmtree(tmp_dir, ignore_errors=True)
+        raise
+    return _load_codec_result(out_dir)
+
+
+def process_codec_video(video_url: str, cfg: CodecConfig) -> dict:
+    """Public entrypoint: video URL + config -> dict(images, src_positions, fps, ...).
+
+    Result is cached on disk under ``cfg.cache_root``; concurrent workers
+    coordinate via a flock-protected sentinel.
+
+    Soft-warning behaviour (B-mode):
+      - If the video has fewer frames than needed to fill ``target_canvas``,
+        we emit a one-time UserWarning describing the shortfall but proceed
+        normally (cv-preinfer will produce fewer canvases than requested).
+      - If the video is so short that cv-preinfer cannot form a single
+        group (``< min_group_frames``), we emit a clearer warning and let
+        cv-preinfer's own error propagate.
+    """
+    cfg.validate()
+    out_dir = _cache_dir_for(video_url, cfg)
+    if (out_dir / "meta.json").exists() and (out_dir / "src_patch_position.npy").exists():
+        return _load_codec_result(out_dir)
+
+    _maybe_warn_short_video(video_url, cfg)
+
+    cfg.cache_root.mkdir(parents=True, exist_ok=True)
+    lock_path = cfg.cache_root / f".{out_dir.name}.lock"
+    lock_fd = os.open(str(lock_path), os.O_CREAT | os.O_RDWR, 0o644)
+    try:
+        if fcntl is not None:
+            fcntl.flock(lock_fd, fcntl.LOCK_EX)
+        if (out_dir / "meta.json").exists() and (out_dir / "src_patch_position.npy").exists():
+            return _load_codec_result(out_dir)
+        if cfg.engine == "dcvc-rt":
+            return _run_dcvc_rt(video_url, out_dir, cfg)
+        if cfg.engine in ("hevc", "cv-preinfer"):
+            return _run_cv_preinfer(video_url, out_dir, cfg)
+        raise ValueError(
+            f"unknown codec engine: {cfg.engine!r} (use 'hevc' or 'dcvc-rt')"
+        )
+    finally:
+        try:
+            if fcntl is not None:
+                fcntl.flock(lock_fd, fcntl.LOCK_UN)
+        finally:
+            os.close(lock_fd)
+
+
+def _maybe_warn_short_video(video_url: str, cfg: CodecConfig) -> None:
+    """Soft-warn (B-mode) when a video is too short to fill target_canvas.
+
+    Logic:
+      * needed_frames  = num_sampled_frames() = (target_canvas/ipg)*group_size
+      * usable_frames  = min(needed_frames, total_frames)
+      * expected_canv  = (usable_frames // group_size) * images_per_group
+    If ``expected_canv < target_canvas`` we warn. If
+    ``total_frames < min_group_frames`` we warn more loudly (cv-preinfer
+    will fail downstream and that error is allowed to propagate).
+    """
+    try:
+        total_frames = _get_video_total_frames(video_url)
+    except Exception:
+        return  # don't fail on probe errors; cv-preinfer will report its own
+
+    needed = cfg.num_sampled_frames()
+    usable = min(needed, total_frames)
+    expected_canv = (usable // cfg.group_size) * cfg.images_per_group
+
+    if total_frames < cfg.min_group_frames:
+        warnings.warn(
+            f"[codec] video {video_url!r} has only {total_frames} frames "
+            f"(< min_group_frames={cfg.min_group_frames}); cv-preinfer cannot "
+            f"form even a single group and will error out. Consider lowering "
+            f"min_group_frames or using video_backend='frames' for this clip.",
+            UserWarning,
+            stacklevel=2,
+        )
+        return
+
+    if expected_canv < cfg.target_canvas:
+        warnings.warn(
+            f"[codec] video {video_url!r} has {total_frames} frames; with "
+            f"group_size={cfg.group_size}, images_per_group={cfg.images_per_group} "
+            f"this yields ~{expected_canv} canvas(es) instead of the requested "
+            f"target_canvas={cfg.target_canvas}. Inference will proceed with the "
+            f"smaller canvas count.",
+            UserWarning,
+            stacklevel=2,
+        )
+
+
+# ----------------------------------------------------- processor wiring
+
+def codec_image_processor_outputs(
+    image_processor, images: list[Image.Image], max_pixels: int,
+) -> dict:
+    """Run ``Qwen2VLImageProcessor`` on codec canvases without smart_resize-ing.
+
+    The codec emits canvases already aligned to the patch grid. To keep
+    ``image_grid_thw`` consistent with ``src_patch_position``:
+      - ``max_pixels`` is clamped up to the largest canvas (never shrinks)
+      - ``min_pixels`` is clamped down to the smallest canvas (never upscales)
+
+    Without the ``min_pixels`` clamp, ``Qwen2VLImageProcessor``'s default
+    ``min_pixels=200704`` would grow any canvas below that threshold,
+    producing extra patches and a chunk/index mismatch downstream.
+    """
+    canvas_pixels = [im.width * im.height for im in images]
+    proc_max = max(int(max_pixels), max(canvas_pixels, default=int(max_pixels)))
+    proc_min = min(canvas_pixels) if canvas_pixels else 1
+    return image_processor(
+        images=images, min_pixels=proc_min, max_pixels=proc_max, return_tensors="pt",
+    )
+
+
+__all__ = [
+    "CodecConfig",
+    "process_codec_video",
+    "drop_padding_canvases",
+    "codec_positions_for_processor",
+    "rewrite_text_with_codec_positions",
+    "codec_image_processor_outputs",
+    "VISION_START", "VISION_END", "IMAGE_PAD",
+]

config.json

@@ -0,0 +1,118 @@
+{
+  "auto_map": {
+    "AutoConfig": "configuration_mage_vl.MageVLConfig",
+    "AutoModel": "modeling_mage_vl.MageVLModel",
+    "AutoModelForImageTextToText": "modeling_mage_vl.MageVLForConditionalGeneration",
+    "AutoProcessor": "processing_mage_vl.MageVLProcessor",
+    "AutoVideoProcessor": "video_processing_mage_vl.MageVLVideoProcessor",
+    "AutoModelForCausalLM": "modeling_mage_vl.MageVLForConditionalGeneration"
+  },
+  "dtype": "bfloat16",
+  "image_token_id": 151655,
+  "model_type": "mage_vl",
+  "text_config": {
+    "_name_or_path": "/mage-vl/pretrain_models/Qwen3-4B-Instruct-2507",
+    "architectures": [
+      "MageVLForConditionalGeneration"
+    ],
+    "attention_bias": false,
+    "attention_dropout": 0.0,
+    "bos_token_id": 151643,
+    "dtype": "bfloat16",
+    "eos_token_id": 151645,
+    "head_dim": 128,
+    "hidden_act": "silu",
+    "hidden_size": 2560,
+    "initializer_range": 0.02,
+    "intermediate_size": 9728,
+    "layer_types": [
+      "full_attention",
+      "full_attention",
+      "full_attention",
+      "full_attention",
+      "full_attention",
+      "full_attention",
+      "full_attention",
+      "full_attention",
+      "full_attention",
+      "full_attention",
+      "full_attention",
+      "full_attention",
+      "full_attention",
+      "full_attention",
+      "full_attention",
+      "full_attention",
+      "full_attention",
+      "full_attention",
+      "full_attention",
+      "full_attention",
+      "full_attention",
+      "full_attention",
+      "full_attention",
+      "full_attention",
+      "full_attention",
+      "full_attention",
+      "full_attention",
+      "full_attention",
+      "full_attention",
+      "full_attention",
+      "full_attention",
+      "full_attention",
+      "full_attention",
+      "full_attention",
+      "full_attention",
+      "full_attention"
+    ],
+    "max_position_embeddings": 262144,
+    "max_window_layers": 36,
+    "model_type": "qwen3",
+    "num_attention_heads": 32,
+    "num_hidden_layers": 36,
+    "num_key_value_heads": 8,
+    "rms_norm_eps": 1e-06,
+    "rope_scaling": null,
+    "sliding_window": null,
+    "use_cache": true,
+    "use_sliding_window": false,
+    "vocab_size": 151936,
+    "rope_parameters": {
+      "rope_theta": 5000000,
+      "rope_type": "default"
+    },
+    "tie_word_embeddings": false
+  },
+  "transformers_version": "5.7.0",
+  "video_token_id": 151656,
+  "vision_config": {
+    "attention_dropout": 0.0,
+    "frame_windows_size": 4,
+    "hidden_act": "gelu",
+    "hidden_size": 1024,
+    "image_size": 448,
+    "initializer_range": 0.02,
+    "intermediate_size": 4096,
+    "layer_norm_eps": 1e-06,
+    "layer_norm_type": "layer_norm",
+    "model_type": "mage_vl_vision",
+    "num_attention_heads": 16,
+    "num_channels": 3,
+    "num_hidden_layers": 24,
+    "out_hidden_size": 2560,
+    "patch_size": 16,
+    "rope_theta": 10000.0,
+    "spatial_merge_size": 2,
+    "temporal_patch_size": 1,
+    "text_hidden_size": 2560,
+    "tokens_per_second": 1,
+    "use_head": false,
+    "max_position_embeddings": 8192
+  },
+  "vision_end_token_id": 151653,
+  "vision_start_token_id": 151652,
+  "architectures": [
+    "MageVLForConditionalGeneration"
+  ],
+  "bos_token_id": 151643,
+  "eos_token_id": 151645,
+  "tie_word_embeddings": false
+}

configuration_mage_vl.py

@@ -0,0 +1,110 @@
+from huggingface_hub.dataclasses import strict
+
+from transformers import CONFIG_MAPPING, AutoConfig
+from transformers.configuration_utils import PreTrainedConfig
+
+
+@strict
+class MageVLVisionConfig(PreTrainedConfig):
+    model_type = "mage_vl_vision"
+    base_config_key = "vision_config"
+
+    hidden_size: int = 1024
+    intermediate_size: int = 4096
+    num_hidden_layers: int = 24
+    num_attention_heads: int = 16
+    num_channels: int = 3
+    image_size: int = 448
+    patch_size: int = 14
+    hidden_act: str = "gelu"
+    layer_norm_eps: float = 1e-6
+    layer_norm_type: str = "layer_norm"
+    attention_dropout: float = 0.0
+    initializer_range: float = 0.02
+    rope_theta: float = 10000.0
+    use_head: bool = False
+    out_hidden_size: int = 1024
+    spatial_merge_size: int = 2
+    tokens_per_second: int = 1
+    frame_windows_size: int = 4
+    use_patch_position_encoding: bool = False
+    patch_position_encoding_type: str = "absolute"
+    max_position_embeddings: int = 8192
+
+
+@strict
+class MageVLConfig(PreTrainedConfig):
+    r"""
+    This is the configuration class to store the configuration of a [`MageVLModel`]. It is used to instantiate a
+    MageVLModel model according to the specified arguments, defining the model architecture. Instantiating a configuration
+    with the defaults will yield a Mage-VL configuration.
+
+    Configuration objects inherit from [`PreTrainedConfig`] and can be used to control the model outputs. Read the
+    documentation from [`PreTrainedConfig`] for more information.
+
+    Args:
+        text_config (`Union[PreTrainedConfig, dict]`, *optional*, defaults to `Qwen3Config`):
+            The config object or dictionary of the text backbone.
+        vision_config (`Union[PreTrainedConfig, dict]`, *optional*, defaults to `MageVLVisionConfig`):
+            The config object or dictionary of the vision backbone.
+        image_token_id (`int`, *optional*, defaults to 151655):
+            The image token index to encode the image prompt.
+        video_token_id (`int`, *optional*, defaults to 151656):
+            The video token index to encode the image prompt.
+        vision_start_token_id (`int`, *optional*, defaults to 151652):
+            The token index to denote start of vision input.
+        vision_end_token_id (`int`, *optional*, defaults to 151653):
+            The token index to denote end of vision input.
+    """
+
+    model_type = "mage_vl"
+    # `text_config` is resolved dynamically based on its `model_type` (defaults to `qwen3`),
+    # so we use `AutoConfig` here as a placeholder; `__post_init__` swaps it for the
+    # concrete config class via `CONFIG_MAPPING`.
+    sub_configs = {"vision_config": MageVLVisionConfig, "text_config": AutoConfig}
+    keys_to_ignore_at_inference = ["past_key_values"]
+
+    text_config: dict | PreTrainedConfig | None = None
+    vision_config: dict | PreTrainedConfig | None = None
+    image_token_id: int = 151655
+    video_token_id: int = 151656
+    vision_start_token_id: int = 151652
+    vision_end_token_id: int = 151653
+    tie_word_embeddings: bool = False
+    # Generation-related token ids are mirrored from `text_config` in `__post_init__`
+    # so downstream tools (e.g. `generate`, vLLM) that read them at the top level keep working.
+    bos_token_id: int | None = None
+    eos_token_id: int | list[int] | None = None
+    pad_token_id: int | None = None
+
+    def __post_init__(self, **kwargs):
+        # Resolve vision_config
+        if isinstance(self.vision_config, dict):
+            self.vision_config = self.sub_configs["vision_config"](**self.vision_config)
+        elif self.vision_config is None:
+            self.vision_config = self.sub_configs["vision_config"]()
+
+        # Resolve text_config dynamically via CONFIG_MAPPING (defaults to qwen3)
+        if isinstance(self.text_config, dict):
+            text_model_type = self.text_config.get("model_type", "qwen3")
+            self.text_config["model_type"] = text_model_type
+            text_config_cls = CONFIG_MAPPING[text_model_type]
+            self.sub_configs["text_config"] = text_config_cls
+            self.text_config = text_config_cls(**self.text_config)
+        elif self.text_config is None:
+            text_config_cls = CONFIG_MAPPING["qwen3"]
+            self.sub_configs["text_config"] = text_config_cls
+            self.text_config = text_config_cls()
+
+        # Mirror generation-related token ids from text_config to the top level so
+        # downstream tools (e.g. `generate`, chat templates, vLLM) that read them
+        # from the top-level config keep working.
+        for tok_key in ("bos_token_id", "eos_token_id", "pad_token_id"):
+            text_val = getattr(self.text_config, tok_key, None)
+            if text_val is not None and getattr(self, tok_key, None) is None:
+                setattr(self, tok_key, text_val)
+
+        super().__post_init__(**kwargs)
+
+
+__all__ = ["MageVLConfig", "MageVLVisionConfig"]

examples/soccer-broadcast.mp4

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

generation_config.json

@@ -0,0 +1,6 @@
+{
+  "_from_model_config": true,
+  "bos_token_id": 151643,
+  "eos_token_id": 151645,
+  "transformers_version": "4.57.3"
+}

inference.py

@@ -0,0 +1,160 @@
+#!/usr/bin/env python
+"""Run Mage-VL-Base image or video inference offline or through SGLang."""
+
+import argparse
+import base64
+import mimetypes
+from pathlib import Path
+
+
+def image_url(image: str) -> str:
+    if image.startswith(("http://", "https://", "data:")):
+        return image
+    path = Path(image)
+    mime = mimetypes.guess_type(path.name)[0] or "image/jpeg"
+    encoded = base64.b64encode(path.read_bytes()).decode("ascii")
+    return f"data:{mime};base64,{encoded}"
+
+
+def sample_video(video: str, num_frames: int):
+    import cv2
+    import numpy as np
+    from PIL import Image
+
+    capture = cv2.VideoCapture(video)
+    frame_count = int(capture.get(cv2.CAP_PROP_FRAME_COUNT))
+    if frame_count <= 0:
+        capture.release()
+        raise ValueError(f"Could not read video: {video}")
+    indices = np.linspace(0, frame_count - 1, min(num_frames, frame_count), dtype=int)
+    frames = []
+    for index in indices:
+        capture.set(cv2.CAP_PROP_POS_FRAMES, int(index))
+        ok, frame = capture.read()
+        if not ok:
+            capture.release()
+            raise ValueError(f"Could not decode frame {index} from: {video}")
+        frames.append(Image.fromarray(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)))
+    capture.release()
+    return frames
+
+
+def run_offline(args):
+    import os
+    import torch
+    from PIL import Image
+    from transformers import AutoModelForCausalLM, AutoProcessor
+
+    model_path = args.model
+    if args.video and args.video_backend == "codec" and args.codec_engine == "neural":
+        if not os.path.isdir(model_path):
+            from huggingface_hub import snapshot_download
+
+            model_path = snapshot_download(args.model)
+    processor = AutoProcessor.from_pretrained(model_path, trust_remote_code=True)
+    model = AutoModelForCausalLM.from_pretrained(
+        model_path, trust_remote_code=True, torch_dtype="auto", device_map="auto"
+    ).eval()
+    media_type = "image" if args.image else "video"
+    messages = [{"role": "user", "content": [
+        {"type": media_type}, {"type": "text", "text": args.question},
+    ]}]
+    text = processor.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
+    if args.image:
+        inputs = processor(
+            text=[text], images=[Image.open(args.image).convert("RGB")], return_tensors="pt"
+        )
+    elif args.video_backend == "codec":
+        codec_config = {
+            "engine": "hevc" if args.codec_engine == "traditional" else "dcvc-rt",
+            "target_canvas": args.num_frames,
+            "patch": 16,
+        }
+        if args.codec_engine == "neural":
+            codec_config["dcvc"] = {
+                "pkg_dir": os.path.join(model_path, "neural_codec"),
+                "device": str(model.device),
+            }
+        inputs = processor(
+            text=[text],
+            videos=[args.video],
+            video_backend="codec",
+            max_pixels=args.max_pixels,
+            codec_config=codec_config,
+            return_tensors="pt",
+            padding=True,
+        )
+    else:
+        inputs = processor(
+            text=[text],
+            videos=[sample_video(args.video, args.num_frames)],
+            return_tensors="pt",
+            padding=True,
+        )
+    inputs = {k: (v.to(model.device) if hasattr(v, "to") else v) for k, v in inputs.items()}
+    if "pixel_values" in inputs:
+        inputs["pixel_values"] = inputs["pixel_values"].to(model.dtype)
+    with torch.inference_mode():
+        output = model.generate(**inputs, max_new_tokens=args.max_new_tokens, do_sample=False)
+    answer = processor.tokenizer.decode(
+        output[0, inputs["input_ids"].shape[1]:], skip_special_tokens=True
+    )
+    print(answer.strip())
+
+
+def run_online(args):
+    from openai import OpenAI
+
+    if args.image:
+        content = [{"type": "image_url", "image_url": {"url": image_url(args.image)}}]
+    else:
+        content = [
+            {"type": "image_url", "image_url": {"url": frame_url(frame)}}
+            for frame in sample_video(args.video, args.num_frames)
+        ]
+    content.append({"type": "text", "text": args.question})
+    client = OpenAI(base_url=args.base_url, api_key=args.api_key)
+    response = client.chat.completions.create(
+        model=args.model,
+        messages=[{"role": "user", "content": content}],
+        max_tokens=args.max_new_tokens,
+    )
+    print(response.choices[0].message.content)
+
+
+def frame_url(frame) -> str:
+    import io
+
+    buffer = io.BytesIO()
+    frame.save(buffer, format="JPEG")
+    encoded = base64.b64encode(buffer.getvalue()).decode("ascii")
+    return f"data:image/jpeg;base64,{encoded}"
+
+
+def main():
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument("--mode", choices=("offline", "online"), required=True)
+    media = parser.add_mutually_exclusive_group(required=True)
+    media.add_argument("--image", help="Local image path")
+    media.add_argument("--video", help="Local video path")
+    parser.add_argument("--video-backend", choices=("frames", "codec"), default="frames")
+    parser.add_argument(
+        "--codec-engine", choices=("traditional", "neural"), default="traditional"
+    )
+    parser.add_argument("--num-frames", type=int, default=32)
+    parser.add_argument("--max-pixels", type=int, default=150000)
+    parser.add_argument("--question", default="Describe this media.")
+    parser.add_argument("--model", default="microsoft/Mage-VL")
+    parser.add_argument("--max-new-tokens", type=int, default=256)
+    parser.add_argument("--base-url", default="http://localhost:30000/v1")
+    parser.add_argument("--api-key", default="EMPTY")
+    args = parser.parse_args()
+    if args.mode == "online" and args.video_backend == "codec":
+        parser.error("online video inference supports only --video-backend frames")
+    if args.num_frames <= 0:
+        parser.error("--num-frames must be positive")
+    (run_offline if args.mode == "offline" else run_online)(args)
+
+
+if __name__ == "__main__":
+    main()

model-00001-of-00002.safetensors

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

model-00002-of-00002.safetensors

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

model.safetensors.index.json

@@ -0,0 +1,703 @@
+{
+  "metadata": {
+    "total_size": 9483587584
+  },
+  "weight_map": {
+    "model.language_model.embed_tokens.weight": "model-00001-of-00002.safetensors",
+    "model.language_model.layers.0.input_layernorm.weight": "model-00001-of-00002.safetensors",
+    "model.language_model.layers.0.mlp.down_proj.weight": "model-00001-of-00002.safetensors",
+    "model.language_model.layers.0.mlp.gate_proj.weight": "model-00001-of-00002.safetensors",
+    "model.language_model.layers.0.mlp.up_proj.weight": "model-00001-of-00002.safetensors",
+    "model.language_model.layers.0.post_attention_layernorm.weight": "model-00001-of-00002.safetensors",
+    "model.language_model.layers.0.self_attn.k_norm.weight": "model-00001-of-00002.safetensors",
+    "model.language_model.layers.0.self_attn.k_proj.weight": "model-00001-of-00002.safetensors",
+    "model.language_model.layers.0.self_attn.o_proj.weight": "model-00001-of-00002.safetensors",
+    "model.language_model.layers.0.self_attn.q_norm.weight": "model-00001-of-00002.safetensors",
+    "model.language_model.layers.0.self_attn.q_proj.weight": "model-00001-of-00002.safetensors",
+    "model.language_model.layers.0.self_attn.v_proj.weight": "model-00001-of-00002.safetensors",
+    "model.language_model.layers.1.input_layernorm.weight": "model-00001-of-00002.safetensors",
+    "model.language_model.layers.1.mlp.down_proj.weight": "model-00001-of-00002.safetensors",
+    "model.language_model.layers.1.mlp.gate_proj.weight": "model-00001-of-00002.safetensors",
+    "model.language_model.layers.1.mlp.up_proj.weight": "model-00001-of-00002.safetensors",
+    "model.language_model.layers.1.post_attention_layernorm.weight": "model-00001-of-00002.safetensors",
+    "model.language_model.layers.1.self_attn.k_norm.weight": "model-00001-of-00002.safetensors",
+    "model.language_model.layers.1.self_attn.k_proj.weight": "model-00001-of-00002.safetensors",
+    "model.language_model.layers.1.self_attn.o_proj.weight": "model-00001-of-00002.safetensors",
+    "model.language_model.layers.1.self_attn.q_norm.weight": "model-00001-of-00002.safetensors",
+    "model.language_model.layers.1.self_attn.q_proj.weight": "model-00001-of-00002.safetensors",
+    "model.language_model.layers.1.self_attn.v_proj.weight": "model-00001-of-00002.safetensors",
+    "model.language_model.layers.10.input_layernorm.weight": "model-00001-of-00002.safetensors",
+    "model.language_model.layers.10.mlp.down_proj.weight": "model-00001-of-00002.safetensors",
+    "model.language_model.layers.10.mlp.gate_proj.weight": "model-00001-of-00002.safetensors",
+    "model.language_model.layers.10.mlp.up_proj.weight": "model-00001-of-00002.safetensors",
+    "model.language_model.layers.10.post_attention_layernorm.weight": "model-00001-of-00002.safetensors",
+    "model.language_model.layers.10.self_attn.k_norm.weight": "model-00001-of-00002.safetensors",
+    "model.language_model.layers.10.self_attn.k_proj.weight": "model-00001-of-00002.safetensors",
+    "model.language_model.layers.10.self_attn.o_proj.weight": "model-00001-of-00002.safetensors",
+    "model.language_model.layers.10.self_attn.q_norm.weight": "model-00001-of-00002.safetensors",
+    "model.language_model.layers.10.self_attn.q_proj.weight": "model-00001-of-00002.safetensors",
+    "model.language_model.layers.10.self_attn.v_proj.weight": "model-00001-of-00002.safetensors",
+    "model.language_model.layers.11.input_layernorm.weight": "model-00001-of-00002.safetensors",
+    "model.language_model.layers.11.mlp.down_proj.weight": "model-00001-of-00002.safetensors",
+    "model.language_model.layers.11.mlp.gate_proj.weight": "model-00001-of-00002.safetensors",
+    "model.language_model.layers.11.mlp.up_proj.weight": "model-00001-of-00002.safetensors",
+    "model.language_model.layers.11.post_attention_layernorm.weight": "model-00001-of-00002.safetensors",
+    "model.language_model.layers.11.self_attn.k_norm.weight": "model-00001-of-00002.safetensors",
+    "model.language_model.layers.11.self_attn.k_proj.weight": "model-00001-of-00002.safetensors",
+    "model.language_model.layers.11.self_attn.o_proj.weight": "model-00001-of-00002.safetensors",
+    "model.language_model.layers.11.self_attn.q_norm.weight": "model-00001-of-00002.safetensors",
+    "model.language_model.layers.11.self_attn.q_proj.weight": "model-00001-of-00002.safetensors",
+    "model.language_model.layers.11.self_attn.v_proj.weight": "model-00001-of-00002.safetensors",
+    "model.language_model.layers.12.input_layernorm.weight": "model-00001-of-00002.safetensors",
+    "model.language_model.layers.12.mlp.down_proj.weight": "model-00001-of-00002.safetensors",
+    "model.language_model.layers.12.mlp.gate_proj.weight": "model-00001-of-00002.safetensors",
+    "model.language_model.layers.12.mlp.up_proj.weight": "model-00001-of-00002.safetensors",
+    "model.language_model.layers.12.post_attention_layernorm.weight": "model-00001-of-00002.safetensors",
+    "model.language_model.layers.12.self_attn.k_norm.weight": "model-00001-of-00002.safetensors",
+    "model.language_model.layers.12.self_attn.k_proj.weight": "model-00001-of-00002.safetensors",
+    "model.language_model.layers.12.self_attn.o_proj.weight": "model-00001-of-00002.safetensors",
+    "model.language_model.layers.12.self_attn.q_norm.weight": "model-00001-of-00002.safetensors",
+    "model.language_model.layers.12.self_attn.q_proj.weight": "model-00001-of-00002.safetensors",
+    "model.language_model.layers.12.self_attn.v_proj.weight": "model-00001-of-00002.safetensors",
+    "model.language_model.layers.13.input_layernorm.weight": "model-00001-of-00002.safetensors",
+    "model.language_model.layers.13.mlp.down_proj.weight": "model-00001-of-00002.safetensors",
+    "model.language_model.layers.13.mlp.gate_proj.weight": "model-00001-of-00002.safetensors",
+    "model.language_model.layers.13.mlp.up_proj.weight": "model-00001-of-00002.safetensors",
+    "model.language_model.layers.13.post_attention_layernorm.weight": "model-00001-of-00002.safetensors",
+    "model.language_model.layers.13.self_attn.k_norm.weight": "model-00001-of-00002.safetensors",
+    "model.language_model.layers.13.self_attn.k_proj.weight": "model-00001-of-00002.safetensors",
+    "model.language_model.layers.13.self_attn.o_proj.weight": "model-00001-of-00002.safetensors",
+    "model.language_model.layers.13.self_attn.q_norm.weight": "model-00001-of-00002.safetensors",
+    "model.language_model.layers.13.self_attn.q_proj.weight": "model-00001-of-00002.safetensors",
+    "model.language_model.layers.13.self_attn.v_proj.weight": "model-00001-of-00002.safetensors",
+    "model.language_model.layers.14.input_layernorm.weight": "model-00001-of-00002.safetensors",
+    "model.language_model.layers.14.mlp.down_proj.weight": "model-00001-of-00002.safetensors",
+    "model.language_model.layers.14.mlp.gate_proj.weight": "model-00001-of-00002.safetensors",
+    "model.language_model.layers.14.mlp.up_proj.weight": "model-00001-of-00002.safetensors",
+    "model.language_model.layers.14.post_attention_layernorm.weight": "model-00001-of-00002.safetensors",
+    "model.language_model.layers.14.self_attn.k_norm.weight": "model-00001-of-00002.safetensors",
+    "model.language_model.layers.14.self_attn.k_proj.weight": "model-00001-of-00002.safetensors",
+    "model.language_model.layers.14.self_attn.o_proj.weight": "model-00001-of-00002.safetensors",
+    "model.language_model.layers.14.self_attn.q_norm.weight": "model-00001-of-00002.safetensors",
+    "model.language_model.layers.14.self_attn.q_proj.weight": "model-00001-of-00002.safetensors",
+    "model.language_model.layers.14.self_attn.v_proj.weight": "model-00001-of-00002.safetensors",
+    "model.language_model.layers.15.input_layernorm.weight": "model-00001-of-00002.safetensors",
+    "model.language_model.layers.15.mlp.down_proj.weight": "model-00001-of-00002.safetensors",
+    "model.language_model.layers.15.mlp.gate_proj.weight": "model-00001-of-00002.safetensors",
+    "model.language_model.layers.15.mlp.up_proj.weight": "model-00001-of-00002.safetensors",
+    "model.language_model.layers.15.post_attention_layernorm.weight": "model-00001-of-00002.safetensors",
+    "model.language_model.layers.15.self_attn.k_norm.weight": "model-00001-of-00002.safetensors",
+    "model.language_model.layers.15.self_attn.k_proj.weight": "model-00001-of-00002.safetensors",
+    "model.language_model.layers.15.self_attn.o_proj.weight": "model-00001-of-00002.safetensors",
+    "model.language_model.layers.15.self_attn.q_norm.weight": "model-00001-of-00002.safetensors",
+    "model.language_model.layers.15.self_attn.q_proj.weight": "model-00001-of-00002.safetensors",
+    "model.language_model.layers.15.self_attn.v_proj.weight": "model-00001-of-00002.safetensors",
+    "model.language_model.layers.16.input_layernorm.weight": "model-00001-of-00002.safetensors",
+    "model.language_model.layers.16.mlp.down_proj.weight": "model-00001-of-00002.safetensors",
+    "model.language_model.layers.16.mlp.gate_proj.weight": "model-00001-of-00002.safetensors",
+    "model.language_model.layers.16.mlp.up_proj.weight": "model-00001-of-00002.safetensors",
+    "model.language_model.layers.16.post_attention_layernorm.weight": "model-00001-of-00002.safetensors",
+    "model.language_model.layers.16.self_attn.k_norm.weight": "model-00001-of-00002.safetensors",
+    "model.language_model.layers.16.self_attn.k_proj.weight": "model-00001-of-00002.safetensors",
+    "model.language_model.layers.16.self_attn.o_proj.weight": "model-00001-of-00002.safetensors",
+    "model.language_model.layers.16.self_attn.q_norm.weight": "model-00001-of-00002.safetensors",
+    "model.language_model.layers.16.self_attn.q_proj.weight": "model-00001-of-00002.safetensors",
+    "model.language_model.layers.16.self_attn.v_proj.weight": "model-00001-of-00002.safetensors",
+    "model.language_model.layers.17.input_layernorm.weight": "model-00001-of-00002.safetensors",
+    "model.language_model.layers.17.mlp.down_proj.weight": "model-00001-of-00002.safetensors",
+    "model.language_model.layers.17.mlp.gate_proj.weight": "model-00001-of-00002.safetensors",
+    "model.language_model.layers.17.mlp.up_proj.weight": "model-00001-of-00002.safetensors",
+    "model.language_model.layers.17.post_attention_layernorm.weight": "model-00001-of-00002.safetensors",
+    "model.language_model.layers.17.self_attn.k_norm.weight": "model-00001-of-00002.safetensors",
+    "model.language_model.layers.17.self_attn.k_proj.weight": "model-00001-of-00002.safetensors",
+    "model.language_model.layers.17.self_attn.o_proj.weight": "model-00001-of-00002.safetensors",
+    "model.language_model.layers.17.self_attn.q_norm.weight": "model-00001-of-00002.safetensors",
+    "model.language_model.layers.17.self_attn.q_proj.weight": "model-00001-of-00002.safetensors",
+    "model.language_model.layers.17.self_attn.v_proj.weight": "model-00001-of-00002.safetensors",
+    "model.language_model.layers.18.input_layernorm.weight": "model-00001-of-00002.safetensors",
+    "model.language_model.layers.18.mlp.down_proj.weight": "model-00001-of-00002.safetensors",
+    "model.language_model.layers.18.mlp.gate_proj.weight": "model-00001-of-00002.safetensors",
+    "model.language_model.layers.18.mlp.up_proj.weight": "model-00001-of-00002.safetensors",
+    "model.language_model.layers.18.post_attention_layernorm.weight": "model-00001-of-00002.safetensors",
+    "model.language_model.layers.18.self_attn.k_norm.weight": "model-00001-of-00002.safetensors",
+    "model.language_model.layers.18.self_attn.k_proj.weight": "model-00001-of-00002.safetensors",
+    "model.language_model.layers.18.self_attn.o_proj.weight": "model-00001-of-00002.safetensors",
+    "model.language_model.layers.18.self_attn.q_norm.weight": "model-00001-of-00002.safetensors",
+    "model.language_model.layers.18.self_attn.q_proj.weight": "model-00001-of-00002.safetensors",
+    "model.language_model.layers.18.self_attn.v_proj.weight": "model-00001-of-00002.safetensors",
+    "model.language_model.layers.19.input_layernorm.weight": "model-00001-of-00002.safetensors",
+    "model.language_model.layers.19.mlp.down_proj.weight": "model-00001-of-00002.safetensors",
+    "model.language_model.layers.19.mlp.gate_proj.weight": "model-00001-of-00002.safetensors",
+    "model.language_model.layers.19.mlp.up_proj.weight": "model-00001-of-00002.safetensors",
+    "model.language_model.layers.19.post_attention_layernorm.weight": "model-00001-of-00002.safetensors",
+    "model.language_model.layers.19.self_attn.k_norm.weight": "model-00001-of-00002.safetensors",
+    "model.language_model.layers.19.self_attn.k_proj.weight": "model-00001-of-00002.safetensors",
+    "model.language_model.layers.19.self_attn.o_proj.weight": "model-00001-of-00002.safetensors",
+    "model.language_model.layers.19.self_attn.q_norm.weight": "model-00001-of-00002.safetensors",
+    "model.language_model.layers.19.self_attn.q_proj.weight": "model-00001-of-00002.safetensors",
+    "model.language_model.layers.19.self_attn.v_proj.weight": "model-00001-of-00002.safetensors",
+    "model.language_model.layers.2.input_layernorm.weight": "model-00001-of-00002.safetensors",
+    "model.language_model.layers.2.mlp.down_proj.weight": "model-00001-of-00002.safetensors",
+    "model.language_model.layers.2.mlp.gate_proj.weight": "model-00001-of-00002.safetensors",
+    "model.language_model.layers.2.mlp.up_proj.weight": "model-00001-of-00002.safetensors",
+    "model.language_model.layers.2.post_attention_layernorm.weight": "model-00001-of-00002.safetensors",
+    "model.language_model.layers.2.self_attn.k_norm.weight": "model-00001-of-00002.safetensors",
+    "model.language_model.layers.2.self_attn.k_proj.weight": "model-00001-of-00002.safetensors",
+    "model.language_model.layers.2.self_attn.o_proj.weight": "model-00001-of-00002.safetensors",
+    "model.language_model.layers.2.self_attn.q_norm.weight": "model-00001-of-00002.safetensors",
+    "model.language_model.layers.2.self_attn.q_proj.weight": "model-00001-of-00002.safetensors",
+    "model.language_model.layers.2.self_attn.v_proj.weight": "model-00001-of-00002.safetensors",
+    "model.language_model.layers.20.input_layernorm.weight": "model-00001-of-00002.safetensors",
+    "model.language_model.layers.20.mlp.down_proj.weight": "model-00001-of-00002.safetensors",
+    "model.language_model.layers.20.mlp.gate_proj.weight": "model-00001-of-00002.safetensors",
+    "model.language_model.layers.20.mlp.up_proj.weight": "model-00001-of-00002.safetensors",
+    "model.language_model.layers.20.post_attention_layernorm.weight": "model-00001-of-00002.safetensors",
+    "model.language_model.layers.20.self_attn.k_norm.weight": "model-00001-of-00002.safetensors",
+    "model.language_model.layers.20.self_attn.k_proj.weight": "model-00001-of-00002.safetensors",
+    "model.language_model.layers.20.self_attn.o_proj.weight": "model-00001-of-00002.safetensors",
+    "model.language_model.layers.20.self_attn.q_norm.weight": "model-00001-of-00002.safetensors",
+    "model.language_model.layers.20.self_attn.q_proj.weight": "model-00001-of-00002.safetensors",
+    "model.language_model.layers.20.self_attn.v_proj.weight": "model-00001-of-00002.safetensors",
+    "model.language_model.layers.21.input_layernorm.weight": "model-00001-of-00002.safetensors",
+    "model.language_model.layers.21.post_attention_layernorm.weight": "model-00001-of-00002.safetensors",
+    "model.language_model.layers.21.self_attn.k_norm.weight": "model-00001-of-00002.safetensors",
+    "model.language_model.layers.21.self_attn.q_norm.weight": "model-00001-of-00002.safetensors",
+    "model.language_model.layers.22.input_layernorm.weight": "model-00001-of-00002.safetensors",
+    "model.language_model.layers.22.post_attention_layernorm.weight": "model-00001-of-00002.safetensors",
+    "model.language_model.layers.22.self_attn.k_norm.weight": "model-00001-of-00002.safetensors",
+    "model.language_model.layers.22.self_attn.q_norm.weight": "model-00001-of-00002.safetensors",
+    "model.language_model.layers.23.input_layernorm.weight": "model-00001-of-00002.safetensors",
+    "model.language_model.layers.23.post_attention_layernorm.weight": "model-00001-of-00002.safetensors",
+    "model.language_model.layers.23.self_attn.k_norm.weight": "model-00001-of-00002.safetensors",
+    "model.language_model.layers.23.self_attn.q_norm.weight": "model-00001-of-00002.safetensors",
+    "model.language_model.layers.24.input_layernorm.weight": "model-00001-of-00002.safetensors",
+    "model.language_model.layers.24.post_attention_layernorm.weight": "model-00001-of-00002.safetensors",
+    "model.language_model.layers.24.self_attn.k_norm.weight": "model-00001-of-00002.safetensors",
+    "model.language_model.layers.24.self_attn.q_norm.weight": "model-00001-of-00002.safetensors",
+    "model.language_model.layers.25.input_layernorm.weight": "model-00001-of-00002.safetensors",
+    "model.language_model.layers.25.post_attention_layernorm.weight": "model-00001-of-00002.safetensors",
+    "model.language_model.layers.25.self_attn.k_norm.weight": "model-00001-of-00002.safetensors",
+    "model.language_model.layers.25.self_attn.q_norm.weight": "model-00001-of-00002.safetensors",
+    "model.language_model.layers.26.input_layernorm.weight": "model-00001-of-00002.safetensors",
+    "model.language_model.layers.26.post_attention_layernorm.weight": "model-00001-of-00002.safetensors",
+    "model.language_model.layers.26.self_attn.k_norm.weight": "model-00001-of-00002.safetensors",
+    "model.language_model.layers.26.self_attn.q_norm.weight": "model-00001-of-00002.safetensors",
+    "model.language_model.layers.27.input_layernorm.weight": "model-00001-of-00002.safetensors",
+    "model.language_model.layers.27.post_attention_layernorm.weight": "model-00001-of-00002.safetensors",
+    "model.language_model.layers.27.self_attn.k_norm.weight": "model-00001-of-00002.safetensors",
+    "model.language_model.layers.27.self_attn.q_norm.weight": "model-00001-of-00002.safetensors",
+    "model.language_model.layers.28.input_layernorm.weight": "model-00001-of-00002.safetensors",
+    "model.language_model.layers.28.post_attention_layernorm.weight": "model-00001-of-00002.safetensors",
+    "model.language_model.layers.28.self_attn.k_norm.weight": "model-00001-of-00002.safetensors",
+    "model.language_model.layers.28.self_attn.q_norm.weight": "model-00001-of-00002.safetensors",
+    "model.language_model.layers.29.input_layernorm.weight": "model-00001-of-00002.safetensors",
+    "model.language_model.layers.29.post_attention_layernorm.weight": "model-00001-of-00002.safetensors",
+    "model.language_model.layers.29.self_attn.k_norm.weight": "model-00001-of-00002.safetensors",
+    "model.language_model.layers.29.self_attn.q_norm.weight": "model-00001-of-00002.safetensors",
+    "model.language_model.layers.3.input_layernorm.weight": "model-00001-of-00002.safetensors",
+    "model.language_model.layers.3.mlp.down_proj.weight": "model-00001-of-00002.safetensors",
+    "model.language_model.layers.3.mlp.gate_proj.weight": "model-00001-of-00002.safetensors",
+    "model.language_model.layers.3.mlp.up_proj.weight": "model-00001-of-00002.safetensors",
+    "model.language_model.layers.3.post_attention_layernorm.weight": "model-00001-of-00002.safetensors",
+    "model.language_model.layers.3.self_attn.k_norm.weight": "model-00001-of-00002.safetensors",
+    "model.language_model.layers.3.self_attn.k_proj.weight": "model-00001-of-00002.safetensors",
+    "model.language_model.layers.3.self_attn.o_proj.weight": "model-00001-of-00002.safetensors",
+    "model.language_model.layers.3.self_attn.q_norm.weight": "model-00001-of-00002.safetensors",
+    "model.language_model.layers.3.self_attn.q_proj.weight": "model-00001-of-00002.safetensors",
+    "model.language_model.layers.3.self_attn.v_proj.weight": "model-00001-of-00002.safetensors",
+    "model.language_model.layers.30.input_layernorm.weight": "model-00001-of-00002.safetensors",
+    "model.language_model.layers.30.post_attention_layernorm.weight": "model-00001-of-00002.safetensors",
+    "model.language_model.layers.30.self_attn.k_norm.weight": "model-00001-of-00002.safetensors",
+    "model.language_model.layers.30.self_attn.q_norm.weight": "model-00001-of-00002.safetensors",
+    "model.language_model.layers.31.input_layernorm.weight": "model-00001-of-00002.safetensors",
+    "model.language_model.layers.31.post_attention_layernorm.weight": "model-00001-of-00002.safetensors",
+    "model.language_model.layers.31.self_attn.k_norm.weight": "model-00001-of-00002.safetensors",
+    "model.language_model.layers.31.self_attn.q_norm.weight": "model-00001-of-00002.safetensors",
+    "model.language_model.layers.32.input_layernorm.weight": "model-00001-of-00002.safetensors",
+    "model.language_model.layers.32.post_attention_layernorm.weight": "model-00001-of-00002.safetensors",
+    "model.language_model.layers.32.self_attn.k_norm.weight": "model-00001-of-00002.safetensors",
+    "model.language_model.layers.32.self_attn.q_norm.weight": "model-00001-of-00002.safetensors",
+    "model.language_model.layers.33.input_layernorm.weight": "model-00001-of-00002.safetensors",
+    "model.language_model.layers.33.post_attention_layernorm.weight": "model-00001-of-00002.safetensors",
+    "model.language_model.layers.33.self_attn.k_norm.weight": "model-00001-of-00002.safetensors",
+    "model.language_model.layers.33.self_attn.q_norm.weight": "model-00001-of-00002.safetensors",
+    "model.language_model.layers.34.input_layernorm.weight": "model-00001-of-00002.safetensors",
+    "model.language_model.layers.34.post_attention_layernorm.weight": "model-00001-of-00002.safetensors",
+    "model.language_model.layers.34.self_attn.k_norm.weight": "model-00001-of-00002.safetensors",
+    "model.language_model.layers.34.self_attn.q_norm.weight": "model-00001-of-00002.safetensors",
+    "model.language_model.layers.35.input_layernorm.weight": "model-00001-of-00002.safetensors",
+    "model.language_model.layers.35.post_attention_layernorm.weight": "model-00001-of-00002.safetensors",
+    "model.language_model.layers.35.self_attn.k_norm.weight": "model-00001-of-00002.safetensors",
+    "model.language_model.layers.35.self_attn.q_norm.weight": "model-00001-of-00002.safetensors",
+    "model.language_model.layers.4.input_layernorm.weight": "model-00001-of-00002.safetensors",
+    "model.language_model.layers.4.mlp.down_proj.weight": "model-00001-of-00002.safetensors",
+    "model.language_model.layers.4.mlp.gate_proj.weight": "model-00001-of-00002.safetensors",
+    "model.language_model.layers.4.mlp.up_proj.weight": "model-00001-of-00002.safetensors",
+    "model.language_model.layers.4.post_attention_layernorm.weight": "model-00001-of-00002.safetensors",
+    "model.language_model.layers.4.self_attn.k_norm.weight": "model-00001-of-00002.safetensors",
+    "model.language_model.layers.4.self_attn.k_proj.weight": "model-00001-of-00002.safetensors",
+    "model.language_model.layers.4.self_attn.o_proj.weight": "model-00001-of-00002.safetensors",
+    "model.language_model.layers.4.self_attn.q_norm.weight": "model-00001-of-00002.safetensors",
+    "model.language_model.layers.4.self_attn.q_proj.weight": "model-00001-of-00002.safetensors",
+    "model.language_model.layers.4.self_attn.v_proj.weight": "model-00001-of-00002.safetensors",
+    "model.language_model.layers.5.input_layernorm.weight": "model-00001-of-00002.safetensors",
+    "model.language_model.layers.5.mlp.down_proj.weight": "model-00001-of-00002.safetensors",
+    "model.language_model.layers.5.mlp.gate_proj.weight": "model-00001-of-00002.safetensors",
+    "model.language_model.layers.5.mlp.up_proj.weight": "model-00001-of-00002.safetensors",
+    "model.language_model.layers.5.post_attention_layernorm.weight": "model-00001-of-00002.safetensors",
+    "model.language_model.layers.5.self_attn.k_norm.weight": "model-00001-of-00002.safetensors",
+    "model.language_model.layers.5.self_attn.k_proj.weight": "model-00001-of-00002.safetensors",
+    "model.language_model.layers.5.self_attn.o_proj.weight": "model-00001-of-00002.safetensors",
+    "model.language_model.layers.5.self_attn.q_norm.weight": "model-00001-of-00002.safetensors",
+    "model.language_model.layers.5.self_attn.q_proj.weight": "model-00001-of-00002.safetensors",
+    "model.language_model.layers.5.self_attn.v_proj.weight": "model-00001-of-00002.safetensors",
+    "model.language_model.layers.6.input_layernorm.weight": "model-00001-of-00002.safetensors",
+    "model.language_model.layers.6.mlp.down_proj.weight": "model-00001-of-00002.safetensors",
+    "model.language_model.layers.6.mlp.gate_proj.weight": "model-00001-of-00002.safetensors",
+    "model.language_model.layers.6.mlp.up_proj.weight": "model-00001-of-00002.safetensors",
+    "model.language_model.layers.6.post_attention_layernorm.weight": "model-00001-of-00002.safetensors",
+    "model.language_model.layers.6.self_attn.k_norm.weight": "model-00001-of-00002.safetensors",
+    "model.language_model.layers.6.self_attn.k_proj.weight": "model-00001-of-00002.safetensors",
+    "model.language_model.layers.6.self_attn.o_proj.weight": "model-00001-of-00002.safetensors",
+    "model.language_model.layers.6.self_attn.q_norm.weight": "model-00001-of-00002.safetensors",
+    "model.language_model.layers.6.self_attn.q_proj.weight": "model-00001-of-00002.safetensors",
+    "model.language_model.layers.6.self_attn.v_proj.weight": "model-00001-of-00002.safetensors",
+    "model.language_model.layers.7.input_layernorm.weight": "model-00001-of-00002.safetensors",
+    "model.language_model.layers.7.mlp.down_proj.weight": "model-00001-of-00002.safetensors",
+    "model.language_model.layers.7.mlp.gate_proj.weight": "model-00001-of-00002.safetensors",
+    "model.language_model.layers.7.mlp.up_proj.weight": "model-00001-of-00002.safetensors",
+    "model.language_model.layers.7.post_attention_layernorm.weight": "model-00001-of-00002.safetensors",
+    "model.language_model.layers.7.self_attn.k_norm.weight": "model-00001-of-00002.safetensors",
+    "model.language_model.layers.7.self_attn.k_proj.weight": "model-00001-of-00002.safetensors",
+    "model.language_model.layers.7.self_attn.o_proj.weight": "model-00001-of-00002.safetensors",
+    "model.language_model.layers.7.self_attn.q_norm.weight": "model-00001-of-00002.safetensors",
+    "model.language_model.layers.7.self_attn.q_proj.weight": "model-00001-of-00002.safetensors",
+    "model.language_model.layers.7.self_attn.v_proj.weight": "model-00001-of-00002.safetensors",
+    "model.language_model.layers.8.input_layernorm.weight": "model-00001-of-00002.safetensors",
+    "model.language_model.layers.8.mlp.down_proj.weight": "model-00001-of-00002.safetensors",
+    "model.language_model.layers.8.mlp.gate_proj.weight": "model-00001-of-00002.safetensors",
+    "model.language_model.layers.8.mlp.up_proj.weight": "model-00001-of-00002.safetensors",
+    "model.language_model.layers.8.post_attention_layernorm.weight": "model-00001-of-00002.safetensors",
+    "model.language_model.layers.8.self_attn.k_norm.weight": "model-00001-of-00002.safetensors",
+    "model.language_model.layers.8.self_attn.k_proj.weight": "model-00001-of-00002.safetensors",
+    "model.language_model.layers.8.self_attn.o_proj.weight": "model-00001-of-00002.safetensors",
+    "model.language_model.layers.8.self_attn.q_norm.weight": "model-00001-of-00002.safetensors",
+    "model.language_model.layers.8.self_attn.q_proj.weight": "model-00001-of-00002.safetensors",
+    "model.language_model.layers.8.self_attn.v_proj.weight": "model-00001-of-00002.safetensors",
+    "model.language_model.layers.9.input_layernorm.weight": "model-00001-of-00002.safetensors",
+    "model.language_model.layers.9.mlp.gate_proj.weight": "model-00001-of-00002.safetensors",
+    "model.language_model.layers.9.mlp.up_proj.weight": "model-00001-of-00002.safetensors",
+    "model.language_model.layers.9.post_attention_layernorm.weight": "model-00001-of-00002.safetensors",
+    "model.language_model.layers.9.self_attn.k_norm.weight": "model-00001-of-00002.safetensors",
+    "model.language_model.layers.9.self_attn.k_proj.weight": "model-00001-of-00002.safetensors",
+    "model.language_model.layers.9.self_attn.o_proj.weight": "model-00001-of-00002.safetensors",
+    "model.language_model.layers.9.self_attn.q_norm.weight": "model-00001-of-00002.safetensors",
+    "model.language_model.layers.9.self_attn.q_proj.weight": "model-00001-of-00002.safetensors",
+    "model.language_model.layers.9.self_attn.v_proj.weight": "model-00001-of-00002.safetensors",
+    "model.language_model.norm.weight": "model-00001-of-00002.safetensors",
+    "lm_head.weight": "model-00002-of-00002.safetensors",
+    "model.language_model.layers.21.mlp.down_proj.weight": "model-00002-of-00002.safetensors",
+    "model.language_model.layers.21.mlp.gate_proj.weight": "model-00002-of-00002.safetensors",
+    "model.language_model.layers.21.mlp.up_proj.weight": "model-00002-of-00002.safetensors",
+    "model.language_model.layers.21.self_attn.k_proj.weight": "model-00002-of-00002.safetensors",
+    "model.language_model.layers.21.self_attn.o_proj.weight": "model-00002-of-00002.safetensors",
+    "model.language_model.layers.21.self_attn.q_proj.weight": "model-00002-of-00002.safetensors",
+    "model.language_model.layers.21.self_attn.v_proj.weight": "model-00002-of-00002.safetensors",
+    "model.language_model.layers.22.mlp.down_proj.weight": "model-00002-of-00002.safetensors",
+    "model.language_model.layers.22.mlp.gate_proj.weight": "model-00002-of-00002.safetensors",
+    "model.language_model.layers.22.mlp.up_proj.weight": "model-00002-of-00002.safetensors",
+    "model.language_model.layers.22.self_attn.k_proj.weight": "model-00002-of-00002.safetensors",
+    "model.language_model.layers.22.self_attn.o_proj.weight": "model-00002-of-00002.safetensors",
+    "model.language_model.layers.22.self_attn.q_proj.weight": "model-00002-of-00002.safetensors",
+    "model.language_model.layers.22.self_attn.v_proj.weight": "model-00002-of-00002.safetensors",
+    "model.language_model.layers.23.mlp.down_proj.weight": "model-00002-of-00002.safetensors",
+    "model.language_model.layers.23.mlp.gate_proj.weight": "model-00002-of-00002.safetensors",
+    "model.language_model.layers.23.mlp.up_proj.weight": "model-00002-of-00002.safetensors",
+    "model.language_model.layers.23.self_attn.k_proj.weight": "model-00002-of-00002.safetensors",
+    "model.language_model.layers.23.self_attn.o_proj.weight": "model-00002-of-00002.safetensors",
+    "model.language_model.layers.23.self_attn.q_proj.weight": "model-00002-of-00002.safetensors",
+    "model.language_model.layers.23.self_attn.v_proj.weight": "model-00002-of-00002.safetensors",
+    "model.language_model.layers.24.mlp.down_proj.weight": "model-00002-of-00002.safetensors",
+    "model.language_model.layers.24.mlp.gate_proj.weight": "model-00002-of-00002.safetensors",
+    "model.language_model.layers.24.mlp.up_proj.weight": "model-00002-of-00002.safetensors",
+    "model.language_model.layers.24.self_attn.k_proj.weight": "model-00002-of-00002.safetensors",
+    "model.language_model.layers.24.self_attn.o_proj.weight": "model-00002-of-00002.safetensors",
+    "model.language_model.layers.24.self_attn.q_proj.weight": "model-00002-of-00002.safetensors",
+    "model.language_model.layers.24.self_attn.v_proj.weight": "model-00002-of-00002.safetensors",
+    "model.language_model.layers.25.mlp.down_proj.weight": "model-00002-of-00002.safetensors",
+    "model.language_model.layers.25.mlp.gate_proj.weight": "model-00002-of-00002.safetensors",
+    "model.language_model.layers.25.mlp.up_proj.weight": "model-00002-of-00002.safetensors",
+    "model.language_model.layers.25.self_attn.k_proj.weight": "model-00002-of-00002.safetensors",
+    "model.language_model.layers.25.self_attn.o_proj.weight": "model-00002-of-00002.safetensors",
+    "model.language_model.layers.25.self_attn.q_proj.weight": "model-00002-of-00002.safetensors",
+    "model.language_model.layers.25.self_attn.v_proj.weight": "model-00002-of-00002.safetensors",
+    "model.language_model.layers.26.mlp.down_proj.weight": "model-00002-of-00002.safetensors",
+    "model.language_model.layers.26.mlp.gate_proj.weight": "model-00002-of-00002.safetensors",
+    "model.language_model.layers.26.mlp.up_proj.weight": "model-00002-of-00002.safetensors",
+    "model.language_model.layers.26.self_attn.k_proj.weight": "model-00002-of-00002.safetensors",
+    "model.language_model.layers.26.self_attn.o_proj.weight": "model-00002-of-00002.safetensors",
+    "model.language_model.layers.26.self_attn.q_proj.weight": "model-00002-of-00002.safetensors",
+    "model.language_model.layers.26.self_attn.v_proj.weight": "model-00002-of-00002.safetensors",
+    "model.language_model.layers.27.mlp.down_proj.weight": "model-00002-of-00002.safetensors",
+    "model.language_model.layers.27.mlp.gate_proj.weight": "model-00002-of-00002.safetensors",
+    "model.language_model.layers.27.mlp.up_proj.weight": "model-00002-of-00002.safetensors",
+    "model.language_model.layers.27.self_attn.k_proj.weight": "model-00002-of-00002.safetensors",
+    "model.language_model.layers.27.self_attn.o_proj.weight": "model-00002-of-00002.safetensors",
+    "model.language_model.layers.27.self_attn.q_proj.weight": "model-00002-of-00002.safetensors",
+    "model.language_model.layers.27.self_attn.v_proj.weight": "model-00002-of-00002.safetensors",
+    "model.language_model.layers.28.mlp.down_proj.weight": "model-00002-of-00002.safetensors",
+    "model.language_model.layers.28.mlp.gate_proj.weight": "model-00002-of-00002.safetensors",
+    "model.language_model.layers.28.mlp.up_proj.weight": "model-00002-of-00002.safetensors",
+    "model.language_model.layers.28.self_attn.k_proj.weight": "model-00002-of-00002.safetensors",
+    "model.language_model.layers.28.self_attn.o_proj.weight": "model-00002-of-00002.safetensors",
+    "model.language_model.layers.28.self_attn.q_proj.weight": "model-00002-of-00002.safetensors",
+    "model.language_model.layers.28.self_attn.v_proj.weight": "model-00002-of-00002.safetensors",
+    "model.language_model.layers.29.mlp.down_proj.weight": "model-00002-of-00002.safetensors",
+    "model.language_model.layers.29.mlp.gate_proj.weight": "model-00002-of-00002.safetensors",
+    "model.language_model.layers.29.mlp.up_proj.weight": "model-00002-of-00002.safetensors",
+    "model.language_model.layers.29.self_attn.k_proj.weight": "model-00002-of-00002.safetensors",
+    "model.language_model.layers.29.self_attn.o_proj.weight": "model-00002-of-00002.safetensors",
+    "model.language_model.layers.29.self_attn.q_proj.weight": "model-00002-of-00002.safetensors",
+    "model.language_model.layers.29.self_attn.v_proj.weight": "model-00002-of-00002.safetensors",
+    "model.language_model.layers.30.mlp.down_proj.weight": "model-00002-of-00002.safetensors",
+    "model.language_model.layers.30.mlp.gate_proj.weight": "model-00002-of-00002.safetensors",
+    "model.language_model.layers.30.mlp.up_proj.weight": "model-00002-of-00002.safetensors",
+    "model.language_model.layers.30.self_attn.k_proj.weight": "model-00002-of-00002.safetensors",
+    "model.language_model.layers.30.self_attn.o_proj.weight": "model-00002-of-00002.safetensors",
+    "model.language_model.layers.30.self_attn.q_proj.weight": "model-00002-of-00002.safetensors",
+    "model.language_model.layers.30.self_attn.v_proj.weight": "model-00002-of-00002.safetensors",
+    "model.language_model.layers.31.mlp.down_proj.weight": "model-00002-of-00002.safetensors",
+    "model.language_model.layers.31.mlp.gate_proj.weight": "model-00002-of-00002.safetensors",
+    "model.language_model.layers.31.mlp.up_proj.weight": "model-00002-of-00002.safetensors",
+    "model.language_model.layers.31.self_attn.k_proj.weight": "model-00002-of-00002.safetensors",
+    "model.language_model.layers.31.self_attn.o_proj.weight": "model-00002-of-00002.safetensors",
+    "model.language_model.layers.31.self_attn.q_proj.weight": "model-00002-of-00002.safetensors",
+    "model.language_model.layers.31.self_attn.v_proj.weight": "model-00002-of-00002.safetensors",
+    "model.language_model.layers.32.mlp.down_proj.weight": "model-00002-of-00002.safetensors",
+    "model.language_model.layers.32.mlp.gate_proj.weight": "model-00002-of-00002.safetensors",
+    "model.language_model.layers.32.mlp.up_proj.weight": "model-00002-of-00002.safetensors",
+    "model.language_model.layers.32.self_attn.k_proj.weight": "model-00002-of-00002.safetensors",
+    "model.language_model.layers.32.self_attn.o_proj.weight": "model-00002-of-00002.safetensors",
+    "model.language_model.layers.32.self_attn.q_proj.weight": "model-00002-of-00002.safetensors",
+    "model.language_model.layers.32.self_attn.v_proj.weight": "model-00002-of-00002.safetensors",
+    "model.language_model.layers.33.mlp.down_proj.weight": "model-00002-of-00002.safetensors",
+    "model.language_model.layers.33.mlp.gate_proj.weight": "model-00002-of-00002.safetensors",
+    "model.language_model.layers.33.mlp.up_proj.weight": "model-00002-of-00002.safetensors",
+    "model.language_model.layers.33.self_attn.k_proj.weight": "model-00002-of-00002.safetensors",
+    "model.language_model.layers.33.self_attn.o_proj.weight": "model-00002-of-00002.safetensors",
+    "model.language_model.layers.33.self_attn.q_proj.weight": "model-00002-of-00002.safetensors",
+    "model.language_model.layers.33.self_attn.v_proj.weight": "model-00002-of-00002.safetensors",
+    "model.language_model.layers.34.mlp.down_proj.weight": "model-00002-of-00002.safetensors",
+    "model.language_model.layers.34.mlp.gate_proj.weight": "model-00002-of-00002.safetensors",
+    "model.language_model.layers.34.mlp.up_proj.weight": "model-00002-of-00002.safetensors",
+    "model.language_model.layers.34.self_attn.k_proj.weight": "model-00002-of-00002.safetensors",
+    "model.language_model.layers.34.self_attn.o_proj.weight": "model-00002-of-00002.safetensors",
+    "model.language_model.layers.34.self_attn.q_proj.weight": "model-00002-of-00002.safetensors",
+    "model.language_model.layers.34.self_attn.v_proj.weight": "model-00002-of-00002.safetensors",
+    "model.language_model.layers.35.mlp.down_proj.weight": "model-00002-of-00002.safetensors",
+    "model.language_model.layers.35.mlp.gate_proj.weight": "model-00002-of-00002.safetensors",
+    "model.language_model.layers.35.mlp.up_proj.weight": "model-00002-of-00002.safetensors",
+    "model.language_model.layers.35.self_attn.k_proj.weight": "model-00002-of-00002.safetensors",
+    "model.language_model.layers.35.self_attn.o_proj.weight": "model-00002-of-00002.safetensors",
+    "model.language_model.layers.35.self_attn.q_proj.weight": "model-00002-of-00002.safetensors",
+    "model.language_model.layers.35.self_attn.v_proj.weight": "model-00002-of-00002.safetensors",
+    "model.language_model.layers.9.mlp.down_proj.weight": "model-00002-of-00002.safetensors",
+    "model.visual.embeddings.patch_embedding.weight": "model-00002-of-00002.safetensors",
+    "model.visual.encoder.layers.0.layer_norm1.bias": "model-00002-of-00002.safetensors",
+    "model.visual.encoder.layers.0.layer_norm1.weight": "model-00002-of-00002.safetensors",
+    "model.visual.encoder.layers.0.layer_norm2.bias": "model-00002-of-00002.safetensors",
+    "model.visual.encoder.layers.0.layer_norm2.weight": "model-00002-of-00002.safetensors",
+    "model.visual.encoder.layers.0.mlp.fc1.bias": "model-00002-of-00002.safetensors",
+    "model.visual.encoder.layers.0.mlp.fc1.weight": "model-00002-of-00002.safetensors",
+    "model.visual.encoder.layers.0.mlp.fc2.bias": "model-00002-of-00002.safetensors",
+    "model.visual.encoder.layers.0.mlp.fc2.weight": "model-00002-of-00002.safetensors",
+    "model.visual.encoder.layers.0.self_attn.proj.bias": "model-00002-of-00002.safetensors",
+    "model.visual.encoder.layers.0.self_attn.proj.weight": "model-00002-of-00002.safetensors",
+    "model.visual.encoder.layers.0.self_attn.qkv.bias": "model-00002-of-00002.safetensors",
+    "model.visual.encoder.layers.0.self_attn.qkv.weight": "model-00002-of-00002.safetensors",
+    "model.visual.encoder.layers.1.layer_norm1.bias": "model-00002-of-00002.safetensors",
+    "model.visual.encoder.layers.1.layer_norm1.weight": "model-00002-of-00002.safetensors",
+    "model.visual.encoder.layers.1.layer_norm2.bias": "model-00002-of-00002.safetensors",
+    "model.visual.encoder.layers.1.layer_norm2.weight": "model-00002-of-00002.safetensors",
+    "model.visual.encoder.layers.1.mlp.fc1.bias": "model-00002-of-00002.safetensors",
+    "model.visual.encoder.layers.1.mlp.fc1.weight": "model-00002-of-00002.safetensors",
+    "model.visual.encoder.layers.1.mlp.fc2.bias": "model-00002-of-00002.safetensors",
+    "model.visual.encoder.layers.1.mlp.fc2.weight": "model-00002-of-00002.safetensors",
+    "model.visual.encoder.layers.1.self_attn.proj.bias": "model-00002-of-00002.safetensors",
+    "model.visual.encoder.layers.1.self_attn.proj.weight": "model-00002-of-00002.safetensors",
+    "model.visual.encoder.layers.1.self_attn.qkv.bias": "model-00002-of-00002.safetensors",
+    "model.visual.encoder.layers.1.self_attn.qkv.weight": "model-00002-of-00002.safetensors",
+    "model.visual.encoder.layers.10.layer_norm1.bias": "model-00002-of-00002.safetensors",
+    "model.visual.encoder.layers.10.layer_norm1.weight": "model-00002-of-00002.safetensors",
+    "model.visual.encoder.layers.10.layer_norm2.bias": "model-00002-of-00002.safetensors",
+    "model.visual.encoder.layers.10.layer_norm2.weight": "model-00002-of-00002.safetensors",
+    "model.visual.encoder.layers.10.mlp.fc1.bias": "model-00002-of-00002.safetensors",
+    "model.visual.encoder.layers.10.mlp.fc1.weight": "model-00002-of-00002.safetensors",
+    "model.visual.encoder.layers.10.mlp.fc2.bias": "model-00002-of-00002.safetensors",
+    "model.visual.encoder.layers.10.mlp.fc2.weight": "model-00002-of-00002.safetensors",
+    "model.visual.encoder.layers.10.self_attn.proj.bias": "model-00002-of-00002.safetensors",
+    "model.visual.encoder.layers.10.self_attn.proj.weight": "model-00002-of-00002.safetensors",
+    "model.visual.encoder.layers.10.self_attn.qkv.bias": "model-00002-of-00002.safetensors",
+    "model.visual.encoder.layers.10.self_attn.qkv.weight": "model-00002-of-00002.safetensors",
+    "model.visual.encoder.layers.11.layer_norm1.bias": "model-00002-of-00002.safetensors",
+    "model.visual.encoder.layers.11.layer_norm1.weight": "model-00002-of-00002.safetensors",
+    "model.visual.encoder.layers.11.layer_norm2.bias": "model-00002-of-00002.safetensors",
+    "model.visual.encoder.layers.11.layer_norm2.weight": "model-00002-of-00002.safetensors",
+    "model.visual.encoder.layers.11.mlp.fc1.bias": "model-00002-of-00002.safetensors",
+    "model.visual.encoder.layers.11.mlp.fc1.weight": "model-00002-of-00002.safetensors",
+    "model.visual.encoder.layers.11.mlp.fc2.bias": "model-00002-of-00002.safetensors",
+    "model.visual.encoder.layers.11.mlp.fc2.weight": "model-00002-of-00002.safetensors",
+    "model.visual.encoder.layers.11.self_attn.proj.bias": "model-00002-of-00002.safetensors",
+    "model.visual.encoder.layers.11.self_attn.proj.weight": "model-00002-of-00002.safetensors",
+    "model.visual.encoder.layers.11.self_attn.qkv.bias": "model-00002-of-00002.safetensors",
+    "model.visual.encoder.layers.11.self_attn.qkv.weight": "model-00002-of-00002.safetensors",
+    "model.visual.encoder.layers.12.layer_norm1.bias": "model-00002-of-00002.safetensors",
+    "model.visual.encoder.layers.12.layer_norm1.weight": "model-00002-of-00002.safetensors",
+    "model.visual.encoder.layers.12.layer_norm2.bias": "model-00002-of-00002.safetensors",
+    "model.visual.encoder.layers.12.layer_norm2.weight": "model-00002-of-00002.safetensors",
+    "model.visual.encoder.layers.12.mlp.fc1.bias": "model-00002-of-00002.safetensors",
+    "model.visual.encoder.layers.12.mlp.fc1.weight": "model-00002-of-00002.safetensors",
+    "model.visual.encoder.layers.12.mlp.fc2.bias": "model-00002-of-00002.safetensors",
+    "model.visual.encoder.layers.12.mlp.fc2.weight": "model-00002-of-00002.safetensors",
+    "model.visual.encoder.layers.12.self_attn.proj.bias": "model-00002-of-00002.safetensors",
+    "model.visual.encoder.layers.12.self_attn.proj.weight": "model-00002-of-00002.safetensors",
+    "model.visual.encoder.layers.12.self_attn.qkv.bias": "model-00002-of-00002.safetensors",
+    "model.visual.encoder.layers.12.self_attn.qkv.weight": "model-00002-of-00002.safetensors",
+    "model.visual.encoder.layers.13.layer_norm1.bias": "model-00002-of-00002.safetensors",
+    "model.visual.encoder.layers.13.layer_norm1.weight": "model-00002-of-00002.safetensors",
+    "model.visual.encoder.layers.13.layer_norm2.bias": "model-00002-of-00002.safetensors",
+    "model.visual.encoder.layers.13.layer_norm2.weight": "model-00002-of-00002.safetensors",
+    "model.visual.encoder.layers.13.mlp.fc1.bias": "model-00002-of-00002.safetensors",
+    "model.visual.encoder.layers.13.mlp.fc1.weight": "model-00002-of-00002.safetensors",
+    "model.visual.encoder.layers.13.mlp.fc2.bias": "model-00002-of-00002.safetensors",
+    "model.visual.encoder.layers.13.mlp.fc2.weight": "model-00002-of-00002.safetensors",
+    "model.visual.encoder.layers.13.self_attn.proj.bias": "model-00002-of-00002.safetensors",
+    "model.visual.encoder.layers.13.self_attn.proj.weight": "model-00002-of-00002.safetensors",
+    "model.visual.encoder.layers.13.self_attn.qkv.bias": "model-00002-of-00002.safetensors",
+    "model.visual.encoder.layers.13.self_attn.qkv.weight": "model-00002-of-00002.safetensors",
+    "model.visual.encoder.layers.14.layer_norm1.bias": "model-00002-of-00002.safetensors",
+    "model.visual.encoder.layers.14.layer_norm1.weight": "model-00002-of-00002.safetensors",
+    "model.visual.encoder.layers.14.layer_norm2.bias": "model-00002-of-00002.safetensors",
+    "model.visual.encoder.layers.14.layer_norm2.weight": "model-00002-of-00002.safetensors",
+    "model.visual.encoder.layers.14.mlp.fc1.bias": "model-00002-of-00002.safetensors",
+    "model.visual.encoder.layers.14.mlp.fc1.weight": "model-00002-of-00002.safetensors",
+    "model.visual.encoder.layers.14.mlp.fc2.bias": "model-00002-of-00002.safetensors",
+    "model.visual.encoder.layers.14.mlp.fc2.weight": "model-00002-of-00002.safetensors",
+    "model.visual.encoder.layers.14.self_attn.proj.bias": "model-00002-of-00002.safetensors",
+    "model.visual.encoder.layers.14.self_attn.proj.weight": "model-00002-of-00002.safetensors",
+    "model.visual.encoder.layers.14.self_attn.qkv.bias": "model-00002-of-00002.safetensors",
+    "model.visual.encoder.layers.14.self_attn.qkv.weight": "model-00002-of-00002.safetensors",
+    "model.visual.encoder.layers.15.layer_norm1.bias": "model-00002-of-00002.safetensors",
+    "model.visual.encoder.layers.15.layer_norm1.weight": "model-00002-of-00002.safetensors",
+    "model.visual.encoder.layers.15.layer_norm2.bias": "model-00002-of-00002.safetensors",
+    "model.visual.encoder.layers.15.layer_norm2.weight": "model-00002-of-00002.safetensors",
+    "model.visual.encoder.layers.15.mlp.fc1.bias": "model-00002-of-00002.safetensors",
+    "model.visual.encoder.layers.15.mlp.fc1.weight": "model-00002-of-00002.safetensors",
+    "model.visual.encoder.layers.15.mlp.fc2.bias": "model-00002-of-00002.safetensors",
+    "model.visual.encoder.layers.15.mlp.fc2.weight": "model-00002-of-00002.safetensors",
+    "model.visual.encoder.layers.15.self_attn.proj.bias": "model-00002-of-00002.safetensors",
+    "model.visual.encoder.layers.15.self_attn.proj.weight": "model-00002-of-00002.safetensors",
+    "model.visual.encoder.layers.15.self_attn.qkv.bias": "model-00002-of-00002.safetensors",
+    "model.visual.encoder.layers.15.self_attn.qkv.weight": "model-00002-of-00002.safetensors",
+    "model.visual.encoder.layers.16.layer_norm1.bias": "model-00002-of-00002.safetensors",
+    "model.visual.encoder.layers.16.layer_norm1.weight": "model-00002-of-00002.safetensors",
+    "model.visual.encoder.layers.16.layer_norm2.bias": "model-00002-of-00002.safetensors",
+    "model.visual.encoder.layers.16.layer_norm2.weight": "model-00002-of-00002.safetensors",
+    "model.visual.encoder.layers.16.mlp.fc1.bias": "model-00002-of-00002.safetensors",
+    "model.visual.encoder.layers.16.mlp.fc1.weight": "model-00002-of-00002.safetensors",
+    "model.visual.encoder.layers.16.mlp.fc2.bias": "model-00002-of-00002.safetensors",
+    "model.visual.encoder.layers.16.mlp.fc2.weight": "model-00002-of-00002.safetensors",
+    "model.visual.encoder.layers.16.self_attn.proj.bias": "model-00002-of-00002.safetensors",
+    "model.visual.encoder.layers.16.self_attn.proj.weight": "model-00002-of-00002.safetensors",
+    "model.visual.encoder.layers.16.self_attn.qkv.bias": "model-00002-of-00002.safetensors",
+    "model.visual.encoder.layers.16.self_attn.qkv.weight": "model-00002-of-00002.safetensors",
+    "model.visual.encoder.layers.17.layer_norm1.bias": "model-00002-of-00002.safetensors",
+    "model.visual.encoder.layers.17.layer_norm1.weight": "model-00002-of-00002.safetensors",
+    "model.visual.encoder.layers.17.layer_norm2.bias": "model-00002-of-00002.safetensors",
+    "model.visual.encoder.layers.17.layer_norm2.weight": "model-00002-of-00002.safetensors",
+    "model.visual.encoder.layers.17.mlp.fc1.bias": "model-00002-of-00002.safetensors",
+    "model.visual.encoder.layers.17.mlp.fc1.weight": "model-00002-of-00002.safetensors",
+    "model.visual.encoder.layers.17.mlp.fc2.bias": "model-00002-of-00002.safetensors",
+    "model.visual.encoder.layers.17.mlp.fc2.weight": "model-00002-of-00002.safetensors",
+    "model.visual.encoder.layers.17.self_attn.proj.bias": "model-00002-of-00002.safetensors",
+    "model.visual.encoder.layers.17.self_attn.proj.weight": "model-00002-of-00002.safetensors",
+    "model.visual.encoder.layers.17.self_attn.qkv.bias": "model-00002-of-00002.safetensors",
+    "model.visual.encoder.layers.17.self_attn.qkv.weight": "model-00002-of-00002.safetensors",
+    "model.visual.encoder.layers.18.layer_norm1.bias": "model-00002-of-00002.safetensors",
+    "model.visual.encoder.layers.18.layer_norm1.weight": "model-00002-of-00002.safetensors",
+    "model.visual.encoder.layers.18.layer_norm2.bias": "model-00002-of-00002.safetensors",
+    "model.visual.encoder.layers.18.layer_norm2.weight": "model-00002-of-00002.safetensors",
+    "model.visual.encoder.layers.18.mlp.fc1.bias": "model-00002-of-00002.safetensors",
+    "model.visual.encoder.layers.18.mlp.fc1.weight": "model-00002-of-00002.safetensors",
+    "model.visual.encoder.layers.18.mlp.fc2.bias": "model-00002-of-00002.safetensors",
+    "model.visual.encoder.layers.18.mlp.fc2.weight": "model-00002-of-00002.safetensors",
+    "model.visual.encoder.layers.18.self_attn.proj.bias": "model-00002-of-00002.safetensors",
+    "model.visual.encoder.layers.18.self_attn.proj.weight": "model-00002-of-00002.safetensors",
+    "model.visual.encoder.layers.18.self_attn.qkv.bias": "model-00002-of-00002.safetensors",
+    "model.visual.encoder.layers.18.self_attn.qkv.weight": "model-00002-of-00002.safetensors",
+    "model.visual.encoder.layers.19.layer_norm1.bias": "model-00002-of-00002.safetensors",
+    "model.visual.encoder.layers.19.layer_norm1.weight": "model-00002-of-00002.safetensors",
+    "model.visual.encoder.layers.19.layer_norm2.bias": "model-00002-of-00002.safetensors",
+    "model.visual.encoder.layers.19.layer_norm2.weight": "model-00002-of-00002.safetensors",
+    "model.visual.encoder.layers.19.mlp.fc1.bias": "model-00002-of-00002.safetensors",
+    "model.visual.encoder.layers.19.mlp.fc1.weight": "model-00002-of-00002.safetensors",
+    "model.visual.encoder.layers.19.mlp.fc2.bias": "model-00002-of-00002.safetensors",
+    "model.visual.encoder.layers.19.mlp.fc2.weight": "model-00002-of-00002.safetensors",
+    "model.visual.encoder.layers.19.self_attn.proj.bias": "model-00002-of-00002.safetensors",
+    "model.visual.encoder.layers.19.self_attn.proj.weight": "model-00002-of-00002.safetensors",
+    "model.visual.encoder.layers.19.self_attn.qkv.bias": "model-00002-of-00002.safetensors",
+    "model.visual.encoder.layers.19.self_attn.qkv.weight": "model-00002-of-00002.safetensors",
+    "model.visual.encoder.layers.2.layer_norm1.bias": "model-00002-of-00002.safetensors",
+    "model.visual.encoder.layers.2.layer_norm1.weight": "model-00002-of-00002.safetensors",
+    "model.visual.encoder.layers.2.layer_norm2.bias": "model-00002-of-00002.safetensors",
+    "model.visual.encoder.layers.2.layer_norm2.weight": "model-00002-of-00002.safetensors",
+    "model.visual.encoder.layers.2.mlp.fc1.bias": "model-00002-of-00002.safetensors",
+    "model.visual.encoder.layers.2.mlp.fc1.weight": "model-00002-of-00002.safetensors",
+    "model.visual.encoder.layers.2.mlp.fc2.bias": "model-00002-of-00002.safetensors",
+    "model.visual.encoder.layers.2.mlp.fc2.weight": "model-00002-of-00002.safetensors",
+    "model.visual.encoder.layers.2.self_attn.proj.bias": "model-00002-of-00002.safetensors",
+    "model.visual.encoder.layers.2.self_attn.proj.weight": "model-00002-of-00002.safetensors",
+    "model.visual.encoder.layers.2.self_attn.qkv.bias": "model-00002-of-00002.safetensors",
+    "model.visual.encoder.layers.2.self_attn.qkv.weight": "model-00002-of-00002.safetensors",
+    "model.visual.encoder.layers.20.layer_norm1.bias": "model-00002-of-00002.safetensors",
+    "model.visual.encoder.layers.20.layer_norm1.weight": "model-00002-of-00002.safetensors",
+    "model.visual.encoder.layers.20.layer_norm2.bias": "model-00002-of-00002.safetensors",
+    "model.visual.encoder.layers.20.layer_norm2.weight": "model-00002-of-00002.safetensors",
+    "model.visual.encoder.layers.20.mlp.fc1.bias": "model-00002-of-00002.safetensors",
+    "model.visual.encoder.layers.20.mlp.fc1.weight": "model-00002-of-00002.safetensors",
+    "model.visual.encoder.layers.20.mlp.fc2.bias": "model-00002-of-00002.safetensors",
+    "model.visual.encoder.layers.20.mlp.fc2.weight": "model-00002-of-00002.safetensors",
+    "model.visual.encoder.layers.20.self_attn.proj.bias": "model-00002-of-00002.safetensors",
+    "model.visual.encoder.layers.20.self_attn.proj.weight": "model-00002-of-00002.safetensors",
+    "model.visual.encoder.layers.20.self_attn.qkv.bias": "model-00002-of-00002.safetensors",
+    "model.visual.encoder.layers.20.self_attn.qkv.weight": "model-00002-of-00002.safetensors",
+    "model.visual.encoder.layers.21.layer_norm1.bias": "model-00002-of-00002.safetensors",
+    "model.visual.encoder.layers.21.layer_norm1.weight": "model-00002-of-00002.safetensors",
+    "model.visual.encoder.layers.21.layer_norm2.bias": "model-00002-of-00002.safetensors",
+    "model.visual.encoder.layers.21.layer_norm2.weight": "model-00002-of-00002.safetensors",
+    "model.visual.encoder.layers.21.mlp.fc1.bias": "model-00002-of-00002.safetensors",
+    "model.visual.encoder.layers.21.mlp.fc1.weight": "model-00002-of-00002.safetensors",
+    "model.visual.encoder.layers.21.mlp.fc2.bias": "model-00002-of-00002.safetensors",
+    "model.visual.encoder.layers.21.mlp.fc2.weight": "model-00002-of-00002.safetensors",
+    "model.visual.encoder.layers.21.self_attn.proj.bias": "model-00002-of-00002.safetensors",
+    "model.visual.encoder.layers.21.self_attn.proj.weight": "model-00002-of-00002.safetensors",
+    "model.visual.encoder.layers.21.self_attn.qkv.bias": "model-00002-of-00002.safetensors",
+    "model.visual.encoder.layers.21.self_attn.qkv.weight": "model-00002-of-00002.safetensors",
+    "model.visual.encoder.layers.22.layer_norm1.bias": "model-00002-of-00002.safetensors",
+    "model.visual.encoder.layers.22.layer_norm1.weight": "model-00002-of-00002.safetensors",
+    "model.visual.encoder.layers.22.layer_norm2.bias": "model-00002-of-00002.safetensors",
+    "model.visual.encoder.layers.22.layer_norm2.weight": "model-00002-of-00002.safetensors",
+    "model.visual.encoder.layers.22.mlp.fc1.bias": "model-00002-of-00002.safetensors",
+    "model.visual.encoder.layers.22.mlp.fc1.weight": "model-00002-of-00002.safetensors",
+    "model.visual.encoder.layers.22.mlp.fc2.bias": "model-00002-of-00002.safetensors",
+    "model.visual.encoder.layers.22.mlp.fc2.weight": "model-00002-of-00002.safetensors",
+    "model.visual.encoder.layers.22.self_attn.proj.bias": "model-00002-of-00002.safetensors",
+    "model.visual.encoder.layers.22.self_attn.proj.weight": "model-00002-of-00002.safetensors",
+    "model.visual.encoder.layers.22.self_attn.qkv.bias": "model-00002-of-00002.safetensors",
+    "model.visual.encoder.layers.22.self_attn.qkv.weight": "model-00002-of-00002.safetensors",
+    "model.visual.encoder.layers.23.layer_norm1.bias": "model-00002-of-00002.safetensors",
+    "model.visual.encoder.layers.23.layer_norm1.weight": "model-00002-of-00002.safetensors",
+    "model.visual.encoder.layers.23.layer_norm2.bias": "model-00002-of-00002.safetensors",
+    "model.visual.encoder.layers.23.layer_norm2.weight": "model-00002-of-00002.safetensors",
+    "model.visual.encoder.layers.23.mlp.fc1.bias": "model-00002-of-00002.safetensors",
+    "model.visual.encoder.layers.23.mlp.fc1.weight": "model-00002-of-00002.safetensors",
+    "model.visual.encoder.layers.23.mlp.fc2.bias": "model-00002-of-00002.safetensors",
+    "model.visual.encoder.layers.23.mlp.fc2.weight": "model-00002-of-00002.safetensors",
+    "model.visual.encoder.layers.23.self_attn.proj.bias": "model-00002-of-00002.safetensors",
+    "model.visual.encoder.layers.23.self_attn.proj.weight": "model-00002-of-00002.safetensors",
+    "model.visual.encoder.layers.23.self_attn.qkv.bias": "model-00002-of-00002.safetensors",
+    "model.visual.encoder.layers.23.self_attn.qkv.weight": "model-00002-of-00002.safetensors",
+    "model.visual.encoder.layers.3.layer_norm1.bias": "model-00002-of-00002.safetensors",
+    "model.visual.encoder.layers.3.layer_norm1.weight": "model-00002-of-00002.safetensors",
+    "model.visual.encoder.layers.3.layer_norm2.bias": "model-00002-of-00002.safetensors",
+    "model.visual.encoder.layers.3.layer_norm2.weight": "model-00002-of-00002.safetensors",
+    "model.visual.encoder.layers.3.mlp.fc1.bias": "model-00002-of-00002.safetensors",
+    "model.visual.encoder.layers.3.mlp.fc1.weight": "model-00002-of-00002.safetensors",
+    "model.visual.encoder.layers.3.mlp.fc2.bias": "model-00002-of-00002.safetensors",
+    "model.visual.encoder.layers.3.mlp.fc2.weight": "model-00002-of-00002.safetensors",
+    "model.visual.encoder.layers.3.self_attn.proj.bias": "model-00002-of-00002.safetensors",
+    "model.visual.encoder.layers.3.self_attn.proj.weight": "model-00002-of-00002.safetensors",
+    "model.visual.encoder.layers.3.self_attn.qkv.bias": "model-00002-of-00002.safetensors",
+    "model.visual.encoder.layers.3.self_attn.qkv.weight": "model-00002-of-00002.safetensors",
+    "model.visual.encoder.layers.4.layer_norm1.bias": "model-00002-of-00002.safetensors",
+    "model.visual.encoder.layers.4.layer_norm1.weight": "model-00002-of-00002.safetensors",
+    "model.visual.encoder.layers.4.layer_norm2.bias": "model-00002-of-00002.safetensors",
+    "model.visual.encoder.layers.4.layer_norm2.weight": "model-00002-of-00002.safetensors",
+    "model.visual.encoder.layers.4.mlp.fc1.bias": "model-00002-of-00002.safetensors",
+    "model.visual.encoder.layers.4.mlp.fc1.weight": "model-00002-of-00002.safetensors",
+    "model.visual.encoder.layers.4.mlp.fc2.bias": "model-00002-of-00002.safetensors",
+    "model.visual.encoder.layers.4.mlp.fc2.weight": "model-00002-of-00002.safetensors",
+    "model.visual.encoder.layers.4.self_attn.proj.bias": "model-00002-of-00002.safetensors",
+    "model.visual.encoder.layers.4.self_attn.proj.weight": "model-00002-of-00002.safetensors",
+    "model.visual.encoder.layers.4.self_attn.qkv.bias": "model-00002-of-00002.safetensors",
+    "model.visual.encoder.layers.4.self_attn.qkv.weight": "model-00002-of-00002.safetensors",
+    "model.visual.encoder.layers.5.layer_norm1.bias": "model-00002-of-00002.safetensors",
+    "model.visual.encoder.layers.5.layer_norm1.weight": "model-00002-of-00002.safetensors",
+    "model.visual.encoder.layers.5.layer_norm2.bias": "model-00002-of-00002.safetensors",
+    "model.visual.encoder.layers.5.layer_norm2.weight": "model-00002-of-00002.safetensors",
+    "model.visual.encoder.layers.5.mlp.fc1.bias": "model-00002-of-00002.safetensors",
+    "model.visual.encoder.layers.5.mlp.fc1.weight": "model-00002-of-00002.safetensors",
+    "model.visual.encoder.layers.5.mlp.fc2.bias": "model-00002-of-00002.safetensors",
+    "model.visual.encoder.layers.5.mlp.fc2.weight": "model-00002-of-00002.safetensors",
+    "model.visual.encoder.layers.5.self_attn.proj.bias": "model-00002-of-00002.safetensors",
+    "model.visual.encoder.layers.5.self_attn.proj.weight": "model-00002-of-00002.safetensors",
+    "model.visual.encoder.layers.5.self_attn.qkv.bias": "model-00002-of-00002.safetensors",
+    "model.visual.encoder.layers.5.self_attn.qkv.weight": "model-00002-of-00002.safetensors",
+    "model.visual.encoder.layers.6.layer_norm1.bias": "model-00002-of-00002.safetensors",
+    "model.visual.encoder.layers.6.layer_norm1.weight": "model-00002-of-00002.safetensors",
+    "model.visual.encoder.layers.6.layer_norm2.bias": "model-00002-of-00002.safetensors",
+    "model.visual.encoder.layers.6.layer_norm2.weight": "model-00002-of-00002.safetensors",
+    "model.visual.encoder.layers.6.mlp.fc1.bias": "model-00002-of-00002.safetensors",
+    "model.visual.encoder.layers.6.mlp.fc1.weight": "model-00002-of-00002.safetensors",
+    "model.visual.encoder.layers.6.mlp.fc2.bias": "model-00002-of-00002.safetensors",
+    "model.visual.encoder.layers.6.mlp.fc2.weight": "model-00002-of-00002.safetensors",
+    "model.visual.encoder.layers.6.self_attn.proj.bias": "model-00002-of-00002.safetensors",
+    "model.visual.encoder.layers.6.self_attn.proj.weight": "model-00002-of-00002.safetensors",
+    "model.visual.encoder.layers.6.self_attn.qkv.bias": "model-00002-of-00002.safetensors",
+    "model.visual.encoder.layers.6.self_attn.qkv.weight": "model-00002-of-00002.safetensors",
+    "model.visual.encoder.layers.7.layer_norm1.bias": "model-00002-of-00002.safetensors",
+    "model.visual.encoder.layers.7.layer_norm1.weight": "model-00002-of-00002.safetensors",
+    "model.visual.encoder.layers.7.layer_norm2.bias": "model-00002-of-00002.safetensors",
+    "model.visual.encoder.layers.7.layer_norm2.weight": "model-00002-of-00002.safetensors",
+    "model.visual.encoder.layers.7.mlp.fc1.bias": "model-00002-of-00002.safetensors",
+    "model.visual.encoder.layers.7.mlp.fc1.weight": "model-00002-of-00002.safetensors",
+    "model.visual.encoder.layers.7.mlp.fc2.bias": "model-00002-of-00002.safetensors",
+    "model.visual.encoder.layers.7.mlp.fc2.weight": "model-00002-of-00002.safetensors",
+    "model.visual.encoder.layers.7.self_attn.proj.bias": "model-00002-of-00002.safetensors",
+    "model.visual.encoder.layers.7.self_attn.proj.weight": "model-00002-of-00002.safetensors",
+    "model.visual.encoder.layers.7.self_attn.qkv.bias": "model-00002-of-00002.safetensors",
+    "model.visual.encoder.layers.7.self_attn.qkv.weight": "model-00002-of-00002.safetensors",
+    "model.visual.encoder.layers.8.layer_norm1.bias": "model-00002-of-00002.safetensors",
+    "model.visual.encoder.layers.8.layer_norm1.weight": "model-00002-of-00002.safetensors",
+    "model.visual.encoder.layers.8.layer_norm2.bias": "model-00002-of-00002.safetensors",
+    "model.visual.encoder.layers.8.layer_norm2.weight": "model-00002-of-00002.safetensors",
+    "model.visual.encoder.layers.8.mlp.fc1.bias": "model-00002-of-00002.safetensors",
+    "model.visual.encoder.layers.8.mlp.fc1.weight": "model-00002-of-00002.safetensors",
+    "model.visual.encoder.layers.8.mlp.fc2.bias": "model-00002-of-00002.safetensors",
+    "model.visual.encoder.layers.8.mlp.fc2.weight": "model-00002-of-00002.safetensors",
+    "model.visual.encoder.layers.8.self_attn.proj.bias": "model-00002-of-00002.safetensors",
+    "model.visual.encoder.layers.8.self_attn.proj.weight": "model-00002-of-00002.safetensors",
+    "model.visual.encoder.layers.8.self_attn.qkv.bias": "model-00002-of-00002.safetensors",
+    "model.visual.encoder.layers.8.self_attn.qkv.weight": "model-00002-of-00002.safetensors",
+    "model.visual.encoder.layers.9.layer_norm1.bias": "model-00002-of-00002.safetensors",
+    "model.visual.encoder.layers.9.layer_norm1.weight": "model-00002-of-00002.safetensors",
+    "model.visual.encoder.layers.9.layer_norm2.bias": "model-00002-of-00002.safetensors",
+    "model.visual.encoder.layers.9.layer_norm2.weight": "model-00002-of-00002.safetensors",
+    "model.visual.encoder.layers.9.mlp.fc1.bias": "model-00002-of-00002.safetensors",
+    "model.visual.encoder.layers.9.mlp.fc1.weight": "model-00002-of-00002.safetensors",
+    "model.visual.encoder.layers.9.mlp.fc2.bias": "model-00002-of-00002.safetensors",
+    "model.visual.encoder.layers.9.mlp.fc2.weight": "model-00002-of-00002.safetensors",
+    "model.visual.encoder.layers.9.self_attn.proj.bias": "model-00002-of-00002.safetensors",
+    "model.visual.encoder.layers.9.self_attn.proj.weight": "model-00002-of-00002.safetensors",
+    "model.visual.encoder.layers.9.self_attn.qkv.bias": "model-00002-of-00002.safetensors",
+    "model.visual.encoder.layers.9.self_attn.qkv.weight": "model-00002-of-00002.safetensors",
+    "model.visual.layernorm_pre.bias": "model-00002-of-00002.safetensors",
+    "model.visual.layernorm_pre.weight": "model-00002-of-00002.safetensors",
+    "model.visual.merger.ln_q.bias": "model-00002-of-00002.safetensors",
+    "model.visual.merger.ln_q.weight": "model-00002-of-00002.safetensors",
+    "model.visual.merger.mlp.0.bias": "model-00002-of-00002.safetensors",
+    "model.visual.merger.mlp.0.weight": "model-00002-of-00002.safetensors",
+    "model.visual.merger.mlp.2.bias": "model-00002-of-00002.safetensors",
+    "model.visual.merger.mlp.2.weight": "model-00002-of-00002.safetensors"
+  }
+}

modeling_mage_vl.py

@@ -0,0 +1,1709 @@
+from collections.abc import Callable
+from dataclasses import dataclass
+from typing import Any, Optional, Union
+from pathlib import Path
+
+import torch
+import torch.nn as nn
+from torch.nn import LayerNorm
+
+from transformers import AutoModel
+from transformers.cache_utils import Cache
+from transformers.generation import GenerationMixin
+from transformers.modeling_outputs import BaseModelOutput, BaseModelOutputWithPooling, ModelOutput
+from transformers.modeling_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedModel
+from transformers.models.siglip.modeling_siglip import SiglipMLP
+from transformers.processing_utils import Unpack
+from transformers.utils import (
+    TransformersKwargs,
+    auto_docstring,
+    can_return_tuple,
+    replace_return_docstrings,
+)
+from transformers.utils.generic import is_flash_attention_requested
+
+from .configuration_mage_vl import MageVLConfig, MageVLVisionConfig
+
+
+@dataclass
+@auto_docstring(
+    custom_intro="""
+    Base class for Mage-VL outputs, with hidden states and attentions.
+    """
+)
+class MageVLModelOutputWithPast(ModelOutput):
+    r"""
+    past_key_values (`Cache`, *optional*, returned when `use_cache=True` is passed or when `config.use_cache=True`):
+        It is a [`~cache_utils.Cache`] instance. For more details, see our [kv cache guide](https://huggingface.co/docs/transformers/en/kv_cache).
+
+        Contains pre-computed hidden-states (key and values in the self-attention blocks) that can be used (see
+        `past_key_values` input) to speed up sequential decoding.
+    """
+
+    last_hidden_state: Optional[torch.FloatTensor] = None
+    past_key_values: Optional[Cache] = None
+    hidden_states: Optional[tuple[torch.FloatTensor]] = None
+    attentions: Optional[tuple[torch.FloatTensor]] = None
+
+
+@dataclass
+@auto_docstring(
+    custom_intro="""
+    Base class for Mage-VL causal language model (or autoregressive) outputs.
+    """
+)
+class MageVLCausalLMOutputWithPast(ModelOutput):
+    r"""
+    loss (`torch.FloatTensor` of shape `(1,)`, *optional*, returned when `labels` is provided):
+        Language modeling loss (for next-token prediction).
+    logits (`torch.FloatTensor` of shape `(batch_size, sequence_length, config.vocab_size)`):
+        Prediction scores of the language modeling head (scores for each vocabulary token before SoftMax).
+    past_key_values (`Cache`, *optional*, returned when `use_cache=True` is passed or when `config.use_cache=True`):
+        It is a [`~cache_utils.Cache`] instance. For more details, see our [kv cache guide](https://huggingface.co/docs/transformers/en/kv_cache).
+
+        Contains pre-computed hidden-states (key and values in the self-attention blocks) that can be used (see
+        `past_key_values` input) to speed up sequential decoding.
+    """
+
+    loss: Optional[torch.FloatTensor] = None
+    logits: Optional[torch.FloatTensor] = None
+    past_key_values: Optional[Cache] = None
+    hidden_states: Optional[tuple[torch.FloatTensor]] = None
+    attentions: Optional[tuple[torch.FloatTensor]] = None
+
+
+# ---------------------------------------------------------------------------
+# Vision Rotary Embedding
+# ---------------------------------------------------------------------------
+
+
+class VisionRotaryEmbedding(nn.Module):
+    """
+    3D (T,H,W) Rotary frequency constructor with 4:6:6 split.
+    Supports both grid_thw-based and explicit position-based RoPE computation.
+    """
+
+    def __init__(self, config: MageVLVisionConfig):
+        super().__init__()
+        head_dim = config.hidden_size // config.num_attention_heads
+        base = config.rope_theta
+
+        assert head_dim % 2 == 0, "head_dim must be even for rotary."
+        assert head_dim % 16 == 0, "head_dim must be divisible by 16."
+        half = head_dim // 2
+        assert half % 16 == 0, "head_dim//2 must also be divisible by 16 to split into 4:6:6."
+
+        self.head_dim = head_dim
+        self.half = half
+        self.base = base
+
+        # 4:6:6 split for T:H:W
+        unit = half // 16
+        self.t_size = 4 * unit
+        self.h_size = 6 * unit
+        self.w_size = 6 * unit
+
+        self.register_buffer(
+            "inv_freq_t",
+            1.0 / (base ** (torch.arange(self.t_size, dtype=torch.float32) / self.t_size)),
+            persistent=False,
+        )
+        self.register_buffer(
+            "inv_freq_h",
+            1.0 / (base ** (torch.arange(self.h_size, dtype=torch.float32) / self.h_size)),
+            persistent=False,
+        )
+        self.register_buffer(
+            "inv_freq_w",
+            1.0 / (base ** (torch.arange(self.w_size, dtype=torch.float32) / self.w_size)),
+            persistent=False,
+        )
+
+    def forward(self, grid_thw: torch.Tensor) -> torch.Tensor:
+        """
+        Compute rotary position embeddings from grid_thw (Qwen2VL style).
+
+        Args:
+            grid_thw: [num_samples, 3] tensor with [t, h, w] for each sample
+
+        Returns:
+            freqs: [total_seq_len, half] tensor of position frequencies
+        """
+        device = grid_thw.device
+        inv_t = self.inv_freq_t.to(device=device)
+        inv_h = self.inv_freq_h.to(device=device)
+        inv_w = self.inv_freq_w.to(device=device)
+
+        all_freqs = []
+        for sample_thw in grid_thw:
+            t, h, w = sample_thw[0].item(), sample_thw[1].item(), sample_thw[2].item()
+
+            # Compute frequency tables
+            ft = torch.outer(torch.arange(t, device=device, dtype=torch.float32), inv_t)
+            fh = torch.outer(torch.arange(h, device=device, dtype=torch.float32), inv_h)
+            fw = torch.outer(torch.arange(w, device=device, dtype=torch.float32), inv_w)
+
+            # Build position indices for this sample
+            t_ids = torch.arange(t, device=device).repeat_interleave(h * w)
+            h_ids = torch.arange(h, device=device).repeat_interleave(w).repeat(t)
+            w_ids = torch.arange(w, device=device).repeat(h).repeat(t)
+
+            # Concatenate frequencies: [seq_len, half]
+            sample_freqs = torch.cat([ft[t_ids], fh[h_ids], fw[w_ids]], dim=-1)
+            all_freqs.append(sample_freqs)
+
+        return torch.cat(all_freqs, dim=0)
+
+    def forward_from_positions(self, patch_positions: torch.Tensor) -> torch.Tensor:
+        """
+        Compute rotary position embeddings from explicit patch positions.
+
+        Args:
+            patch_positions: [seq_len, 3] tensor with [t, h, w] positions for each patch
+
+        Returns:
+            freqs: [seq_len, half] tensor of position frequencies
+        """
+        device = patch_positions.device
+        inv_t = self.inv_freq_t.to(device=device)
+        inv_h = self.inv_freq_h.to(device=device)
+        inv_w = self.inv_freq_w.to(device=device)
+
+        t_pos = patch_positions[:, 0].float()
+        h_pos = patch_positions[:, 1].float()
+        w_pos = patch_positions[:, 2].float()
+
+        ft = torch.outer(t_pos, inv_t)
+        fh = torch.outer(h_pos, inv_h)
+        fw = torch.outer(w_pos, inv_w)
+
+        return torch.cat([ft, fh, fw], dim=-1)
+
+    def forward_with_thw(self, t: int, h: int, w: int, device=None) -> torch.Tensor:
+        """
+        Compute rotary position embeddings from explicit t, h, w dimensions.
+
+        Args:
+            t: Number of temporal frames
+            h: Number of height patches
+            w: Number of width patches
+            device: Target device
+
+        Returns:
+            freqs: [t*h*w, half] tensor of position frequencies
+        """
+        if device is None:
+            device = self.inv_freq_t.device
+
+        inv_t = self.inv_freq_t.to(device=device)
+        inv_h = self.inv_freq_h.to(device=device)
+        inv_w = self.inv_freq_w.to(device=device)
+
+        ft = torch.outer(torch.arange(t, device=device, dtype=torch.float32), inv_t)
+        fh = torch.outer(torch.arange(h, device=device, dtype=torch.float32), inv_h)
+        fw = torch.outer(torch.arange(w, device=device, dtype=torch.float32), inv_w)
+
+        t_ids = torch.arange(t, device=device).repeat_interleave(h * w)
+        h_ids = torch.arange(h, device=device).repeat_interleave(w).repeat(t)
+        w_ids = torch.arange(w, device=device).repeat(h).repeat(t)
+
+        freqs = torch.cat([ft[t_ids], fh[h_ids], fw[w_ids]], dim=-1)
+        return freqs
+
+
+# ---------------------------------------------------------------------------
+# Patch Embedding
+# ---------------------------------------------------------------------------
+
+
+class MageVLVisionEmbeddings(nn.Module):
+    """
+    Patch embedding layer that converts pre-processed patches to embeddings.
+
+    This module is designed to receive patches that have already been extracted
+    and arranged by the Qwen2VL image processor in 2x2 block spatial order.
+
+    Input format: [total_patches, num_channels, patch_size, patch_size]
+    Output format: [total_patches, embed_dim]
+    """
+
+    def __init__(self, config: MageVLVisionConfig):
+        super().__init__()
+        self.config = config
+        self.embed_dim = config.hidden_size
+        self.image_size = config.image_size
+        self.patch_size = config.patch_size
+        self.in_channels = config.num_channels
+
+        self.patch_embedding = nn.Conv2d(
+            in_channels=config.num_channels,
+            out_channels=self.embed_dim,
+            kernel_size=self.patch_size,
+            stride=self.patch_size,
+            bias=False,
+        )
+
+    def forward(self, hidden_states: torch.FloatTensor) -> torch.Tensor:
+        target_dtype = self.patch_embedding.weight.dtype
+        hidden_states = hidden_states.view(-1, self.in_channels, self.patch_size, self.patch_size)
+        hidden_states = self.patch_embedding(hidden_states.to(dtype=target_dtype)).view(-1, self.embed_dim)
+
+        return hidden_states
+
+
+# ---------------------------------------------------------------------------
+# Patch Merger
+# ---------------------------------------------------------------------------
+
+
+class MageVLVisionPatchMerger(nn.Module):
+    """
+    Patch merger that merges spatial_merge_size x spatial_merge_size patches into one.
+
+    This module is designed to work with Qwen2VL-style patch processing where patches
+    are already arranged in 2x2 block order by the image processor.
+    """
+
+    def __init__(
+        self,
+        dim: int,
+        context_dim: int,
+        spatial_merge_size: int = 2,
+        layer_norm_eps: float = 1e-05,
+        use_patch_position_encoding: bool = False,
+        patch_position_encoding_type: str = "absolute",
+        max_position_embeddings: int = 8192,
+    ) -> None:
+        super().__init__()
+        self.hidden_size = context_dim * (spatial_merge_size**2)
+        self.ln_q = LayerNorm(context_dim, eps=layer_norm_eps)
+        self.mlp = nn.Sequential(
+            nn.Linear(self.hidden_size, self.hidden_size),
+            nn.GELU(),
+            nn.Linear(self.hidden_size, dim),
+        )
+        self.spatial_merge_size = spatial_merge_size
+        self.use_patch_position_encoding = use_patch_position_encoding
+        self.patch_position_encoding_type = patch_position_encoding_type
+
+        if self.use_patch_position_encoding:
+            if self.patch_position_encoding_type != "absolute":
+                raise ValueError(
+                    f"Unknown patch_position_encoding_type: {self.patch_position_encoding_type}. "
+                    "Only 'absolute' is supported."
+                )
+            self.pos_emb_h = nn.Embedding(max_position_embeddings, dim)
+            self.pos_emb_w = nn.Embedding(max_position_embeddings, dim)
+
+    def forward(self, x: torch.Tensor, patch_positions: Optional[torch.Tensor] = None) -> torch.Tensor:
+        """
+        Merge patches from Qwen2VL-style input.
+
+        The input patches are already arranged in 2x2 block order by the image processor,
+        so we simply need to apply LayerNorm, reshape, and project through MLP.
+
+        Args:
+            x: Input tensor of shape [batch_size, seq_len, hidden_size] or [seq_len, hidden_size]
+               where seq_len = t * h * w (patches in 2x2 block order)
+
+        Returns:
+            Merged tensor of shape [batch_size, seq_len // spatial_merge_size^2, dim]
+            or [seq_len // spatial_merge_size^2, dim]
+        """
+        if patch_positions is not None and patch_positions.dim() == 3:
+            patch_positions = patch_positions.squeeze(0)
+
+        x = self.ln_q(x).view(-1, self.hidden_size)
+        x = self.mlp(x)
+
+        if self.use_patch_position_encoding and patch_positions is not None:
+            pp = patch_positions.view(-1, self.spatial_merge_size**2, 3)
+            pp = pp[:, 0, :]
+            pp = (pp // self.spatial_merge_size).long()
+
+            x = x + self.pos_emb_h(pp[:, 1]) + self.pos_emb_w(pp[:, 2])
+
+        return x
+
+
+def rotate_half(x):
+    """
+    Interleaved rotation to match Source model's implementation.
+    (x1, x2, x3, x4) -> (-x2, x1, -x4, x3)
+    """
+    x_even = x[..., ::2]
+    x_odd = x[..., 1::2]
+    return torch.stack((-x_odd, x_even), dim=-1).flatten(-2)
+
+
+def get_norm_layer(config):
+    if config.layer_norm_type == "rms_norm":
+        return nn.RMSNorm(config.hidden_size, eps=config.layer_norm_eps)
+    else:
+        return nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
+
+
+def apply_rotary_pos_emb(q, k, freqs):
+    # q, k: (B, H, L, D)
+    # freqs: (B, L, D)
+    orig_q_dtype = q.dtype
+    orig_k_dtype = k.dtype
+    q, k = q.float(), k.float()
+    # We need to broadcast freqs to match heads
+    # (B, L, D) -> (B, 1, L, D)
+    # Keep the same dtype as q, k to avoid memory doubling from float32 promotion
+    cos = freqs.cos().unsqueeze(1).float()
+    sin = freqs.sin().unsqueeze(1).float()
+
+    q_embed = (q * cos) + (rotate_half(q) * sin)
+    k_embed = (k * cos) + (rotate_half(k) * sin)
+    q_embed = q_embed.to(orig_q_dtype)
+    k_embed = k_embed.to(orig_k_dtype)
+    return q_embed, k_embed
+
+
+def eager_attention_forward(
+    module: nn.Module,
+    query: torch.Tensor,
+    key: torch.Tensor,
+    value: torch.Tensor,
+    attention_mask: Optional[torch.Tensor],
+    scaling: float,
+    dropout: float = 0.0,
+    **kwargs,
+):
+    """Eager attention; query/key/value are expected as ``(B, H, L, D)``."""
+    attn_weights = torch.matmul(query, key.transpose(2, 3)) * scaling
+    if attention_mask is not None:
+        attn_weights = attn_weights + attention_mask
+    attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to(query.dtype)
+    attn_weights = nn.functional.dropout(attn_weights, p=dropout, training=module.training)
+    attn_output = torch.matmul(attn_weights, value)
+    attn_output = attn_output.transpose(1, 2).contiguous()  # (B, L, H, D)
+    return attn_output, attn_weights
+
+
+class MageVLVisionAttention(nn.Module):
+    """
+    Multi-headed attention with RoPE support, dispatched through
+    :data:`ALL_ATTENTION_FUNCTIONS` (``eager`` / ``sdpa`` / ``flash_attention_2``)
+    based on ``config._attn_implementation``.
+    """
+
+    def __init__(self, config: MageVLVisionConfig):
+        super().__init__()
+        self.config = config
+        self.embed_dim = config.hidden_size
+        self.num_heads = config.num_attention_heads
+        self.head_dim = self.embed_dim // self.num_heads
+        if self.head_dim * self.num_heads != self.embed_dim:
+            raise ValueError(
+                f"embed_dim must be divisible by num_heads (got `embed_dim`: {self.embed_dim} and `num_heads`: {self.num_heads})."
+            )
+
+        self.num_key_value_groups = 1  # required by repeat_kv-aware eager paths
+        self.scale = self.head_dim**-0.5
+        self.scaling = self.scale  # alias expected by some attention interfaces
+        self.attention_dropout = config.attention_dropout
+        self.is_causal = False
+        self.qkv = nn.Linear(self.embed_dim, self.embed_dim * 3)
+        self.proj = nn.Linear(self.embed_dim, self.embed_dim)
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        attention_mask: Optional[torch.Tensor] = None,
+        rotary_pos_emb: Optional[torch.Tensor] = None,
+        output_attentions: bool = False,
+        cu_seqlens: Optional[torch.Tensor] = None,
+        max_seqlen: Optional[int] = None,
+        **kwargs,
+    ) -> tuple[torch.Tensor, Optional[torch.Tensor]]:
+        batch_size, q_len, _ = hidden_states.size()
+        # (B, L, 3*H*D) -> (B, L, 3, H, D) -> 3 x (B, L, H, D) -> 3 x (B, H, L, D)
+        q, k, v = (
+            self.qkv(hidden_states)
+            .reshape(batch_size, q_len, 3, self.num_heads, self.head_dim)
+            .permute(2, 0, 1, 3, 4)
+            .unbind(0)
+        )
+        query_states = q.transpose(1, 2)
+        key_states = k.transpose(1, 2)
+        value_states = v.transpose(1, 2)
+
+        if rotary_pos_emb is not None:
+            query_states, key_states = apply_rotary_pos_emb(query_states, key_states, rotary_pos_emb)
+
+        attention_interface: Callable = ALL_ATTENTION_FUNCTIONS.get_interface(
+            self.config._attn_implementation, eager_attention_forward
+        )
+        dropout = 0.0 if not self.training else self.attention_dropout
+
+        if cu_seqlens is not None and is_flash_attention_requested(self.config):
+            # Flash Attention varlen path: pass cu_seq_lens / max_length kwargs.
+            if max_seqlen is None:
+                max_seqlen = (cu_seqlens[1:] - cu_seqlens[:-1]).max()
+            attn_output, _ = attention_interface(
+                self,
+                query_states,
+                key_states,
+                value_states,
+                attention_mask=None,
+                scaling=self.scale,
+                dropout=dropout,
+                cu_seq_lens_q=cu_seqlens,
+                cu_seq_lens_k=cu_seqlens,
+                max_length_q=max_seqlen,
+                max_length_k=max_seqlen,
+                is_causal=False,
+                **kwargs,
+            )
+        elif cu_seqlens is not None:
+            # Non-FA implementations do not understand cu_seqlens directly; mirror
+            # Qwen3-VL by splitting the packed sequence into per-sample chunks
+            # along the L dim of (B, H, L, D) and running attention per chunk.
+            lengths = (cu_seqlens[1:] - cu_seqlens[:-1]).tolist()
+            splits = [torch.split(t, lengths, dim=2) for t in (query_states, key_states, value_states)]
+            attn_outputs = [
+                attention_interface(
+                    self,
+                    q_chunk,
+                    k_chunk,
+                    v_chunk,
+                    attention_mask=None,
+                    scaling=self.scale,
+                    dropout=dropout,
+                    is_causal=False,
+                    **kwargs,
+                )[0]
+                for q_chunk, k_chunk, v_chunk in zip(*splits)
+            ]
+            # interface output is (B, l_i, H, D); concat along the L axis
+            attn_output = torch.cat(attn_outputs, dim=1)
+        else:
+            attn_mask = None
+            if attention_mask is not None:
+                attn_mask = attention_mask
+                if attn_mask.dim() == 2:
+                    attn_mask = attn_mask.unsqueeze(0)
+                if attn_mask.shape[0] == 1 and batch_size > 1:
+                    attn_mask = attn_mask.expand(batch_size, -1, -1)
+                attn_mask = attn_mask.unsqueeze(1)  # (B, 1, L, L)
+            attn_output, _ = attention_interface(
+                self,
+                query_states,
+                key_states,
+                value_states,
+                attention_mask=attn_mask,
+                scaling=self.scale,
+                dropout=dropout,
+                is_causal=False,
+                **kwargs,
+            )
+
+        attn_output = attn_output.reshape(batch_size, q_len, self.embed_dim)
+        attn_output = self.proj(attn_output)
+
+        return attn_output, None
+
+
+class MageVLVisionEncoderLayer(nn.Module):
+    """Vision encoder layer with pre-norm and Flash Attention 2."""
+
+    def __init__(self, config: MageVLVisionConfig):
+        super().__init__()
+        self.embed_dim = config.hidden_size
+        self.self_attn = MageVLVisionAttention(config)
+        self.layer_norm1 = get_norm_layer(config)
+        self.mlp = SiglipMLP(config)
+        self.layer_norm2 = get_norm_layer(config)
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        attention_mask: Optional[torch.Tensor] = None,
+        rotary_pos_emb: Optional[torch.Tensor] = None,
+        output_attentions: bool = False,
+        cu_seqlens: Optional[torch.Tensor] = None,
+        max_seqlen: Optional[int] = None,
+    ) -> tuple[torch.Tensor, Optional[torch.Tensor]]:
+        residual = hidden_states
+        hidden_states = self.layer_norm1(hidden_states)
+
+        hidden_states, attn_weights = self.self_attn(
+            hidden_states=hidden_states,
+            attention_mask=attention_mask,
+            rotary_pos_emb=rotary_pos_emb,
+            output_attentions=output_attentions,
+            cu_seqlens=cu_seqlens,
+            max_seqlen=max_seqlen,
+        )
+        hidden_states = residual + hidden_states
+
+        residual = hidden_states
+        hidden_states = self.layer_norm2(hidden_states)
+        hidden_states = self.mlp(hidden_states)
+        hidden_states = residual + hidden_states
+
+        outputs = (hidden_states, attn_weights) if output_attentions else (hidden_states,)
+        return outputs
+
+
+class MageVLVisionEncoder(nn.Module):
+    def __init__(self, config: MageVLVisionConfig):
+        super().__init__()
+        self.config = config
+        self.layers = nn.ModuleList([MageVLVisionEncoderLayer(config) for _ in range(config.num_hidden_layers)])
+        # Gradient checkpointing support
+        self.gradient_checkpointing = False
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        attention_mask: Optional[torch.Tensor] = None,
+        rotary_pos_emb: Optional[torch.Tensor] = None,
+        output_attentions: bool = False,
+        output_hidden_states: bool = False,
+        return_dict: bool = True,
+        cu_seqlens: Optional[torch.Tensor] = None,
+        max_seqlen: Optional[int] = None,
+    ) -> Union[tuple, BaseModelOutput]:
+        all_hidden_states = () if output_hidden_states else None
+        all_self_attentions = () if output_attentions else None
+
+        for layer in self.layers:
+            if output_hidden_states:
+                all_hidden_states = all_hidden_states + (hidden_states,)
+
+            if self.gradient_checkpointing and self.training:
+                layer_outputs = self._gradient_checkpointing_func(
+                    layer.__call__,
+                    hidden_states,
+                    attention_mask,
+                    rotary_pos_emb,
+                    output_attentions,
+                    cu_seqlens,
+                    max_seqlen,
+                )
+            else:
+                layer_outputs = layer(
+                    hidden_states,
+                    attention_mask=attention_mask,
+                    rotary_pos_emb=rotary_pos_emb,
+                    output_attentions=output_attentions,
+                    cu_seqlens=cu_seqlens,
+                    max_seqlen=max_seqlen,
+                )
+
+            hidden_states = layer_outputs[0]
+
+            if output_attentions:
+                all_self_attentions = all_self_attentions + (layer_outputs[1],)
+
+        if output_hidden_states:
+            all_hidden_states = all_hidden_states + (hidden_states,)
+
+        if not return_dict:
+            return tuple(v for v in [hidden_states, all_hidden_states, all_self_attentions] if v is not None)
+
+        return BaseModelOutput(
+            last_hidden_state=hidden_states,
+            hidden_states=all_hidden_states,
+            attentions=all_self_attentions,
+        )
+
+
+class MageVLPreTrainedModel(PreTrainedModel):
+    config_class = MageVLConfig
+    base_model_prefix = "model"
+    input_modalities = ("image", "video", "text")
+    supports_gradient_checkpointing = True
+    _no_split_modules = ["MageVLVisionEncoderLayer", "Qwen3DecoderLayer"]
+    _skip_keys_device_placement = "past_key_values"
+    _supports_flash_attn = True
+    _supports_sdpa = True
+
+    def _init_weights(self, module):
+        super()._init_weights(module)
+        # Re-initialize VisionRotaryEmbedding inv_freq buffers.
+        # These are registered with persistent=False, so they are not in the checkpoint
+        # state_dict. When ``from_pretrained`` materializes the model from meta tensors,
+        # the values in these buffers end up uninitialized. Mirror Qwen3-VL by explicitly
+        # filling them here so RoPE produces the correct frequencies post-load.
+        if isinstance(module, VisionRotaryEmbedding):
+            base = module.base
+            with torch.no_grad():
+                inv_t = 1.0 / (base ** (torch.arange(module.t_size, dtype=torch.float32) / module.t_size))
+                inv_h = 1.0 / (base ** (torch.arange(module.h_size, dtype=torch.float32) / module.h_size))
+                inv_w = 1.0 / (base ** (torch.arange(module.w_size, dtype=torch.float32) / module.w_size))
+                module.inv_freq_t.copy_(inv_t.to(module.inv_freq_t.device))
+                module.inv_freq_h.copy_(inv_h.to(module.inv_freq_h.device))
+                module.inv_freq_w.copy_(inv_w.to(module.inv_freq_w.device))
+
+
+class Siglip2MultiheadAttentionPoolingHead(nn.Module):
+    """
+    Multi-Head Attention Pooling with a learned probe (PMA-style).
+    """
+
+    def __init__(self, config: MageVLVisionConfig):
+        super().__init__()
+        self.embed_dim = config.hidden_size
+        self.probe = nn.Parameter(torch.randn(1, 1, config.hidden_size))
+        self.attention = nn.MultiheadAttention(config.hidden_size, config.num_attention_heads, batch_first=True)
+        self.norm = nn.RMSNorm(config.hidden_size, eps=config.layer_norm_eps)
+        self.mlp = SiglipMLP(config)
+
+    def forward(self, hidden_states):
+        batch_size = hidden_states.shape[0]
+        probe = self.probe.repeat(batch_size, 1, 1)
+
+        attn_output, _ = self.attention(probe, hidden_states, hidden_states)
+
+        residual = attn_output
+        attn_output = self.norm(attn_output)
+        attn_output = residual + self.mlp(attn_output)
+
+        return attn_output[:, 0]
+
+
+# ---------------------------------------------------------------------------
+# Vision Model
+# ---------------------------------------------------------------------------
+
+
+class MageVLVisionPretrainedModel(MageVLPreTrainedModel):
+    """
+    Mage-VL Vision Model.
+
+    This vision model is designed to work with Qwen2VL-style image processing:
+        - Receives pre-processed patches in 2x2 block spatial order
+        - Applies RoPE with matching 2x2 block layout conversion
+        - Accepts explicit patch_positions for RoPE computation
+
+    Input format:
+        hidden_state: [total_patches, num_channels, patch_size, patch_size]
+        grid_thw: [num_samples, 3] with [t, h, w] for each sample
+    """
+
+    def __init__(self, config: MageVLVisionConfig):
+        super().__init__(config)
+        self.config = config
+        self.spatial_merge_size = config.spatial_merge_size
+
+        # Vision components
+        self.embeddings = MageVLVisionEmbeddings(config)
+        self.layernorm_pre = get_norm_layer(config)
+        self.encoder = MageVLVisionEncoder(config)
+        self.video_rope = VisionRotaryEmbedding(config)
+
+        if config.use_head:
+            self.layernorm_post = get_norm_layer(config)
+            self.head = Siglip2MultiheadAttentionPoolingHead(config)
+        else:
+            self.layernorm_post = None
+            self.head = None
+
+        self.merger = MageVLVisionPatchMerger(
+            dim=config.out_hidden_size,
+            context_dim=config.hidden_size,
+            spatial_merge_size=config.spatial_merge_size,
+            layer_norm_eps=config.layer_norm_eps,
+            use_patch_position_encoding=getattr(config, "use_patch_position_encoding", False),
+            patch_position_encoding_type=getattr(config, "patch_position_encoding_type", "absolute"),
+            max_position_embeddings=getattr(config, "max_position_embeddings", 8192),
+        )
+
+        self.post_init()
+
+    def _build_cu_seqlens(
+        self,
+        grid_thw: torch.Tensor,
+        total_patches: int,
+        fixed_t: Optional[int] = 4,
+        device: Optional[torch.device] = None,
+    ) -> tuple[torch.Tensor, int]:
+        if grid_thw is None or grid_thw.numel() == 0:
+            # Fallback for no grid_thw: treat as single sequence
+            return torch.tensor([0, total_patches], dtype=torch.int32, device=device), total_patches
+
+        if device is None:
+            device = grid_thw.device
+
+        cu_seqlens = [0]
+        max_seqlen = 0
+        total_entries = grid_thw.shape[0]
+        current_len = 0
+
+        # Calculate cumulative lengths: split sequences based on fixed_t if provided
+        for idx in range(total_entries):
+            t_val = grid_thw[idx, 0].item()
+            h_val = grid_thw[idx, 1].item()
+            w_val = grid_thw[idx, 2].item()
+
+            if fixed_t is not None and fixed_t > 0 and t_val > fixed_t:
+                # Split large t into chunks of fixed_t
+                num_full_windows = t_val // fixed_t
+                remainder = t_val % fixed_t
+
+                # Add full windows
+                for _ in range(num_full_windows):
+                    chunk_patches = fixed_t * int(h_val) * int(w_val)
+                    current_len += chunk_patches
+                    max_seqlen = max(max_seqlen, chunk_patches)
+                    cu_seqlens.append(current_len)
+
+                # Add remainder if any
+                if remainder > 0:
+                    chunk_patches = remainder * int(h_val) * int(w_val)
+                    current_len += chunk_patches
+                    max_seqlen = max(max_seqlen, chunk_patches)
+                    cu_seqlens.append(current_len)
+            else:
+                # Standard case: add as one chunk
+                chunk_patches = t_val * int(h_val) * int(w_val)
+                current_len += chunk_patches
+                max_seqlen = max(max_seqlen, chunk_patches)
+                cu_seqlens.append(current_len)
+
+        last_len = cu_seqlens[-1]
+        if last_len != total_patches:
+            raise ValueError(
+                "cu_seqlens calculation mismatch:\n"
+                f"- total_patches: {total_patches}\n"
+                f"- calculated total: {last_len}\n"
+                f"- grid_thw: {grid_thw}"
+            )
+
+        return torch.tensor(cu_seqlens, dtype=torch.int32, device=device), max_seqlen
+
+    def _build_block_attention_mask(
+        self,
+        grid_thw: torch.Tensor,
+        total_patches: int,
+        fixed_t: Optional[int] = 4,
+        device: Optional[torch.device] = None,
+    ) -> Optional[torch.Tensor]:
+        if grid_thw is None or grid_thw.numel() == 0:
+            return None
+
+        if device is None:
+            device = grid_thw.device
+
+        lengths = []
+        total_entries = grid_thw.shape[0]
+
+        for idx in range(total_entries):
+            t_val = grid_thw[idx, 0].item()
+            h_val = grid_thw[idx, 1].item()
+            w_val = grid_thw[idx, 2].item()
+
+            if fixed_t is not None and fixed_t > 0 and t_val > fixed_t:
+                # Split large t into chunks of fixed_t
+                num_full_windows = t_val // fixed_t
+                remainder = t_val % fixed_t
+
+                # Add full windows
+                for _ in range(num_full_windows):
+                    lengths.append(fixed_t * int(h_val) * int(w_val))
+
+                # Add remainder if any
+                if remainder > 0:
+                    lengths.append(remainder * int(h_val) * int(w_val))
+            else:
+                lengths.append(t_val * int(h_val) * int(w_val))
+
+        total_len = sum(lengths)
+        if total_len != total_patches:
+            raise ValueError(
+                "Block attention mask length mismatch:\n"
+                f"- total_patches: {total_patches}\n"
+                f"- total_len: {total_len}\n"
+                f"- grid_thw: {grid_thw}"
+            )
+
+        attn_mask = torch.ones((total_len, total_len), dtype=torch.bool, device=device)
+        start = 0
+        for size in lengths:
+            end = start + size
+            attn_mask[start:end, start:end] = False
+            start = end
+
+        return attn_mask
+
+    @replace_return_docstrings(output_type=BaseModelOutputWithPooling, config_class=MageVLVisionConfig)
+    def forward(
+        self,
+        hidden_state: torch.Tensor,
+        grid_thw: Optional[torch.Tensor] = None,
+        patch_positions: Optional[torch.Tensor] = None,
+        output_attentions: Optional[bool] = None,
+        output_hidden_states: Optional[bool] = None,
+        return_dict: Optional[bool] = None,
+        skip_merger: Optional[bool] = False,
+    ) -> Union[tuple, BaseModelOutputWithPooling]:
+        r"""
+        Forward pass for vision model.
+
+        This method accepts pre-processed patches from Qwen2VL image processor and applies
+        RoPE (Rotary Position Embedding) in 2x2 block layout to match the spatial arrangement
+        of patches.
+
+        Args:
+            hidden_state: Pre-processed patches from Qwen2VL processor.
+                Shape: [total_patches, num_channels, patch_size, patch_size]
+            grid_thw: Grid sizes tensor of shape [num_samples, 3] with [t, h, w] for each sample.
+                Required for computing RoPE and handling visible indices.
+            patch_positions: Optional explicit patch positions for RoPE computation.
+            output_attentions: Whether to return attention weights.
+            output_hidden_states: Whether to return all hidden states.
+            return_dict: Whether to return a ModelOutput instead of tuple.
+            skip_merger: If True, skip patch merger (useful for consistency checking).
+
+        Returns:
+            BaseModelOutputWithPooling with last_hidden_state containing merged features.
+        """
+        output_attentions = (
+            output_attentions if output_attentions is not None else getattr(self.config, "output_attentions", False)
+        )
+        output_hidden_states = (
+            output_hidden_states
+            if output_hidden_states is not None
+            else getattr(self.config, "output_hidden_states", False)
+        )
+        return_dict = True if return_dict is None else return_dict
+
+        # 1. Embeddings
+        # Note: embeddings returns [total_patches, embed_dim], we need to add batch dimension
+        hidden_states = self.embeddings(hidden_state)
+        if hidden_states.dim() == 2:
+            hidden_states = hidden_states.unsqueeze(0)  # [1, total_patches, embed_dim]
+        batch_size, total_patches, _ = hidden_states.shape
+
+        # 2. RoPE Construction
+        if patch_positions is not None and patch_positions.dim() == 3:
+            patch_positions = patch_positions.squeeze(0)
+        freqs_visible = self.video_rope.forward_from_positions(patch_positions)
+
+        # Concatenate D/2 + D/2 -> D for applying rope
+        freqs_visible = torch.cat([freqs_visible, freqs_visible], dim=-1)
+        if freqs_visible.dim() == 2:
+            freqs_visible = freqs_visible.unsqueeze(0)
+
+        # 3. Pre-Norm & Encoder
+        hidden_states = self.layernorm_pre(hidden_states)
+
+        cu_seqlens, max_seqlen = self._build_cu_seqlens(
+            grid_thw=grid_thw,
+            total_patches=total_patches,
+            fixed_t=getattr(self.config, "frame_windows_size", 4),
+            device=hidden_states.device,
+        )
+
+        encoder_outputs = self.encoder(
+            hidden_states,
+            attention_mask=None,
+            rotary_pos_emb=freqs_visible,
+            output_attentions=output_attentions,
+            output_hidden_states=True,  # Always get hidden states to use -2 layer
+            return_dict=True,
+            cu_seqlens=cu_seqlens,
+            max_seqlen=max_seqlen,
+        )
+
+        # Use second-to-last layer output for better feature representation
+        if encoder_outputs.hidden_states is not None and len(encoder_outputs.hidden_states) >= 2 and not skip_merger:
+            sequence_output = encoder_outputs.hidden_states[-1]
+        else:
+            sequence_output = encoder_outputs[0]
+
+        # Post-Norm
+        if self.layernorm_post is not None:
+            sequence_output = self.layernorm_post(sequence_output)
+
+        # Skip merger for consistency check with original ViT
+        if skip_merger:
+            pooled_output = None
+            if self.head is not None:
+                pooled_output = self.head(sequence_output)
+
+            if not return_dict:
+                return (sequence_output, pooled_output) + (
+                    encoder_outputs.hidden_states if output_hidden_states else None,
+                )
+            return BaseModelOutputWithPooling(
+                last_hidden_state=sequence_output,
+                pooler_output=pooled_output,
+                hidden_states=encoder_outputs.hidden_states if output_hidden_states else None,
+                attentions=encoder_outputs.attentions if output_attentions else None,
+            )
+
+        # Patch merger: input patches are already in 2x2 block order from Qwen2VL processor
+        merged_output = self.merger(sequence_output, patch_positions=patch_positions)
+
+        if not return_dict:
+            return (merged_output,) + (encoder_outputs.hidden_states if output_hidden_states else None,)
+
+        return BaseModelOutputWithPooling(
+            last_hidden_state=merged_output,
+            pooler_output=None,
+            hidden_states=encoder_outputs.hidden_states if output_hidden_states else None,
+            attentions=encoder_outputs.attentions if output_attentions else None,
+        )
+
+
+@auto_docstring
+class MageVLModel(MageVLPreTrainedModel):
+    base_model_prefix = ""
+    # Reference: fix gemma3 grad acc #37208
+    accepts_loss_kwargs = False
+    config: MageVLConfig
+    _no_split_modules = ["MageVLVisionEncoderLayer", "Qwen3DecoderLayer"]
+
+    def __init__(self, config: MageVLConfig):
+        super().__init__(config)
+        self.visual = MageVLVisionPretrainedModel._from_config(config.vision_config)
+        self.language_model = AutoModel.from_config(config.text_config)
+        self.streammind_gate = None
+        self._streammind_model_path = None
+
+        # Initialize weights and apply final processing
+        self.post_init()
+
+    def _load_streammind_gate(self):
+        if self.streammind_gate is not None:
+            return self.streammind_gate
+        from safetensors.torch import load_file
+        from .streammind_gate import StreamMindGate
+
+        if not self._streammind_model_path:
+            raise RuntimeError(
+                "StreamMind gate path is unavailable. Load the model with "
+                "MageVLStreamMindForConditionalGeneration.from_pretrained()."
+            )
+        gate = StreamMindGate(self.config.text_config.hidden_size)
+        gate_path = Path(self._streammind_model_path) / "streammind_gate.safetensors"
+        if not gate_path.is_file():
+            from huggingface_hub import hf_hub_download
+
+            gate_path = Path(
+                hf_hub_download(
+                    repo_id=self._streammind_model_path,
+                    filename="streammind_gate.safetensors",
+                )
+            )
+        state = load_file(str(gate_path))
+        gate.load_state_dict(state, strict=True)
+        gate.to(
+            device=next(self.visual.parameters()).device,
+            dtype=next(self.visual.parameters()).dtype,
+        ).eval()
+        self.streammind_gate = gate
+        return gate
+
+    def _streammind_vision_tokens(self, pixel_values, image_grid_thw, patch_positions=None):
+        pixel_values = pixel_values.type(self.visual.embeddings.patch_embedding.weight.dtype)
+        rope = self.visual.video_rope
+        saved_rope = (rope.inv_freq_t, rope.inv_freq_h, rope.inv_freq_w)
+        try:
+            # The StreamMind checkpoint was trained after the whole OV model,
+            # including non-persistent RoPE buffers, was cast to BF16.
+            rope.inv_freq_t = rope.inv_freq_t.to(pixel_values.dtype)
+            rope.inv_freq_h = rope.inv_freq_h.to(pixel_values.dtype)
+            rope.inv_freq_w = rope.inv_freq_w.to(pixel_values.dtype)
+            vision_output = self.visual(
+                pixel_values,
+                grid_thw=image_grid_thw,
+                patch_positions=patch_positions,
+            )
+        finally:
+            rope.inv_freq_t, rope.inv_freq_h, rope.inv_freq_w = saved_rope
+        merged = vision_output.last_hidden_state.reshape(-1, vision_output.last_hidden_state.shape[-1])
+        merge = self.visual.spatial_merge_size
+        time = int(image_grid_thw[:, 0].sum().item())
+        patches_per_time = int(
+            (image_grid_thw[0, 1].item() // merge)
+            * (image_grid_thw[0, 2].item() // merge)
+        )
+        vision_tokens = merged.reshape(1, time, patches_per_time, -1)
+        return vision_tokens
+
+    def streammind_gate_forward(self, pixel_values, image_grid_thw, patch_positions=None):
+        """Run the gate on one segment without changing the base LLM visual path."""
+        vision_tokens = self._streammind_vision_tokens(
+            pixel_values, image_grid_thw, patch_positions=patch_positions
+        )
+        return self._load_streammind_gate()(vision_tokens)
+
+    def streammind_gate_forward_segments(self, segments):
+        """Run EPFE continuously over a list of codec segments from one stream."""
+        tokens = [
+            self._streammind_vision_tokens(
+                segment["pixel_values"],
+                segment["image_grid_thw"],
+                patch_positions=segment.get("patch_positions"),
+            )
+            for segment in segments
+        ]
+        lengths = [token.shape[1] for token in tokens]
+        boundaries = torch.tensor(lengths).cumsum(0).tolist()
+        return self._load_streammind_gate()(
+            torch.cat(tokens, dim=1), response_positions=boundaries
+        )
+
+    def get_input_embeddings(self):
+        return self.language_model.get_input_embeddings()
+
+    def set_input_embeddings(self, value):
+        self.language_model.set_input_embeddings(value)
+
+    def set_decoder(self, decoder):
+        self.language_model = decoder
+
+    def get_decoder(self):
+        return self.language_model
+
+    def get_video_features(
+        self,
+        pixel_values_videos: torch.FloatTensor,
+        video_grid_thw: Optional[torch.LongTensor] = None,
+        patch_positions=None,
+    ):
+        """
+        Encodes videos into continuous embeddings that can be forwarded to the language model.
+
+        Args:
+            pixel_values_videos: Pre-processed patches from Qwen2VL processor.
+                `torch.FloatTensor` of shape `(total_patches, num_channels, patch_size, patch_size)`
+            video_grid_thw (`torch.LongTensor` of shape `(num_videos, 3)`, *optional*):
+                The temporal, height and width of feature shape of each video in LLM.
+        """
+        # Convert to correct dtype
+        pixel_values_videos = pixel_values_videos.type(self.visual.embeddings.patch_embedding.weight.dtype)
+
+        # Forward through vision model with grid_thw
+        vision_output = self.visual(pixel_values_videos, grid_thw=video_grid_thw, patch_positions=patch_positions)
+
+        # Extract the actual tensor from BaseModelOutputWithPooling
+        if hasattr(vision_output, "last_hidden_state"):
+            video_embeds = vision_output.last_hidden_state
+        else:
+            video_embeds = vision_output[0]  # Fallback for tuple output
+
+        # Compute split sizes from video_grid_thw or from input shape
+        if video_grid_thw is not None:
+            split_sizes = (video_grid_thw.prod(-1) // self.visual.spatial_merge_size**2).tolist()
+        else:
+            # Compute from input shape
+            batch_size = pixel_values_videos.shape[0]
+            split_sizes = [video_embeds.shape[1]] * batch_size
+
+        # Split embeddings per video
+        if len(split_sizes) > 1:
+            video_embeds = torch.split(video_embeds.view(-1, video_embeds.shape[-1]), split_sizes)
+        else:
+            video_embeds = [video_embeds.view(-1, video_embeds.shape[-1])]
+
+        return video_embeds
+
+    def get_image_features(
+        self, pixel_values, image_grid_thw: Optional[torch.LongTensor] = None, patch_positions=None
+    ):
+        """
+        Encodes images into continuous embeddings that can be forwarded to the language model.
+
+        Args:
+            pixel_values: Pre-processed patches from Qwen2VL processor.
+                - `torch.FloatTensor` of shape `(total_patches, num_channels, patch_size, patch_size)`
+            image_grid_thw (`torch.LongTensor` of shape `(num_images, 3)`, *optional*):
+                The temporal, height and width of feature shape of each image in LLM.
+        """
+        # Standard format from Qwen2VL processor
+        if pixel_values.dim() == 2:
+            # Convert to correct dtype
+            pixel_values = pixel_values.type(self.visual.embeddings.patch_embedding.weight.dtype)
+
+            # Forward through vision model with grid_thw
+            vision_output = self.visual(pixel_values, grid_thw=image_grid_thw, patch_positions=patch_positions)
+
+            # Extract the actual tensor from BaseModelOutputWithPooling
+            if hasattr(vision_output, "last_hidden_state"):
+                image_embeds = vision_output.last_hidden_state
+            else:
+                image_embeds = vision_output[0]
+
+            # Compute split sizes from grid_thw
+            if image_grid_thw is not None:
+                split_sizes = (image_grid_thw.prod(-1) // self.visual.spatial_merge_size**2).tolist()
+            else:
+                # Fallback: assume single image
+                split_sizes = [image_embeds.shape[0] if image_embeds.dim() == 2 else image_embeds.shape[1]]
+
+            # Split embeddings per image
+            image_embeds_flat = image_embeds.view(-1, image_embeds.shape[-1])
+            if len(split_sizes) > 1:
+                image_embeds = list(torch.split(image_embeds_flat, split_sizes))
+            else:
+                image_embeds = [image_embeds_flat]
+
+            return image_embeds
+        else:
+            raise ValueError(
+                f"Unsupported pixel_values shape: expected 4D tensor [total_patches, C, H, W], "
+                f"got {pixel_values.shape if hasattr(pixel_values, 'shape') else type(pixel_values)}"
+            )
+
+    def get_placeholder_mask(
+        self,
+        input_ids: torch.LongTensor,
+        inputs_embeds: torch.FloatTensor,
+        image_features: Optional[torch.FloatTensor] = None,
+        video_features: Optional[torch.FloatTensor] = None,
+    ):
+        """
+        Obtains multimodal placeholder mask from `input_ids` or `inputs_embeds`, and checks that the placeholder token count is
+        equal to the length of multimodal features. If the lengths are different, an error is raised.
+        """
+        if input_ids is None:
+            special_image_mask = inputs_embeds == self.get_input_embeddings()(
+                torch.tensor(self.config.image_token_id, dtype=torch.long, device=inputs_embeds.device)
+            )
+            special_image_mask = special_image_mask.all(-1)
+            special_video_mask = inputs_embeds == self.get_input_embeddings()(
+                torch.tensor(self.config.video_token_id, dtype=torch.long, device=inputs_embeds.device)
+            )
+            special_video_mask = special_video_mask.all(-1)
+        else:
+            special_image_mask = input_ids == self.config.image_token_id
+            special_video_mask = input_ids == self.config.video_token_id
+
+        n_image_tokens = special_image_mask.sum()
+        special_image_mask = special_image_mask.unsqueeze(-1).expand_as(inputs_embeds).to(inputs_embeds.device)
+        if image_features is not None and inputs_embeds[special_image_mask].numel() != image_features.numel():
+            raise ValueError(
+                f"Image features and image tokens do not match: tokens: {n_image_tokens}, features {image_features.shape[0]}"
+            )
+
+        n_video_tokens = special_video_mask.sum()
+        special_video_mask = special_video_mask.unsqueeze(-1).expand_as(inputs_embeds).to(inputs_embeds.device)
+        if video_features is not None and inputs_embeds[special_video_mask].numel() != video_features.numel():
+            raise ValueError(
+                f"Videos features and video tokens do not match: tokens: {n_video_tokens}, features {video_features.shape[0]}"
+            )
+
+        return special_image_mask, special_video_mask
+
+    @auto_docstring
+    def forward(
+        self,
+        input_ids: Optional[torch.LongTensor] = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        past_key_values: Optional[Cache] = None,
+        inputs_embeds: Optional[torch.FloatTensor] = None,
+        use_cache: Optional[bool] = None,
+        output_attentions: Optional[bool] = None,
+        output_hidden_states: Optional[bool] = None,
+        return_dict: Optional[bool] = None,
+        pixel_values: Optional[torch.Tensor] = None,
+        pixel_values_videos: Optional[torch.FloatTensor] = None,
+        image_grid_thw: Optional[torch.LongTensor] = None,
+        patch_positions: Optional[torch.LongTensor] = None,
+        video_grid_thw: Optional[torch.LongTensor] = None,
+        cache_position: Optional[torch.LongTensor] = None,
+        second_per_grid_ts: Optional[torch.Tensor] = None,
+        **kwargs: Unpack[TransformersKwargs],
+    ) -> Union[tuple, MageVLModelOutputWithPast]:
+        r"""
+        image_grid_thw (`torch.LongTensor` of shape `(num_images, 3)`, *optional*):
+            The temporal, height and width of feature shape of each image in LLM.
+        video_grid_thw (`torch.LongTensor` of shape `(num_videos, 3)`, *optional*):
+            The temporal, height and width of feature shape of each video in LLM.
+        patch_positions (`torch.LongTensor` of shape `(total_patches, 3)` or `(1, total_patches, 3)`, *optional*):
+            Explicit per-patch `(t, h, w)` position indices used by the vision tower to compute 3D rotary
+            position embeddings (and the optional absolute position embedding inside the patch merger).
+            `total_patches` is the sum of `t * h * w` across all images and videos in the batch, matching
+            the layout produced by the Qwen2VL-style image processor.
+        second_per_grid_ts (`torch.Tensor` of shape `(num_videos)`, *optional*):
+            The time interval (in seconds) for each grid along the temporal dimension in the 3D position IDs.
+        cache_position (`torch.LongTensor` of shape `(sequence_length)`, *optional*):
+              Indices depicting the position of the input sequence tokens in the sequence. Contrarily to
+              `position_ids`, this tensor is not affected by padding.
+
+        Note: see the top-level ``MageVLForConditionalGeneration.forward``
+        docstring; currently video flows in via the ``image_grid_thw`` / ``pixel_values``
+        alias, so ``pixel_values_videos`` / ``video_grid_thw`` /
+        ``second_per_grid_ts`` are unused at this layer.
+        """
+        output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
+        output_hidden_states = (
+            output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
+        )
+        return_dict = True if return_dict is None else return_dict
+
+        if inputs_embeds is None:
+            inputs_embeds = self.get_input_embeddings()(input_ids)
+
+        image_embeds = None
+
+        if pixel_values is not None:
+            image_embeds = self.get_image_features(pixel_values, image_grid_thw, patch_positions=patch_positions)
+
+        if image_embeds is not None:
+            image_embeds = torch.cat(image_embeds, dim=0).to(inputs_embeds.device, inputs_embeds.dtype)
+            image_mask, _ = self.get_placeholder_mask(
+                input_ids, inputs_embeds=inputs_embeds, image_features=image_embeds
+            )
+            inputs_embeds = inputs_embeds.masked_scatter(image_mask, image_embeds)
+
+        if pixel_values_videos is not None:
+            video_embeds = self.get_video_features(
+                pixel_values_videos, video_grid_thw, patch_positions=patch_positions
+            )
+            video_embeds = torch.cat(video_embeds, dim=0).to(inputs_embeds.device, inputs_embeds.dtype)
+            _, video_mask = self.get_placeholder_mask(
+                input_ids, inputs_embeds=inputs_embeds, video_features=video_embeds
+            )
+            inputs_embeds = inputs_embeds.masked_scatter(video_mask, video_embeds)
+
+        # Use simple 1D position_ids
+        if position_ids is None:
+            batch_size, seq_length, _ = inputs_embeds.shape
+            if attention_mask is not None:
+                position_ids = attention_mask.long().cumsum(-1) - 1
+                position_ids.masked_fill_(attention_mask == 0, 1)
+            else:
+                position_ids = (
+                    torch.arange(seq_length, device=inputs_embeds.device).unsqueeze(0).expand(batch_size, -1)
+                )
+
+            # Handle cache_position for generation
+            if cache_position is not None and cache_position[0] != 0:
+                position_ids = position_ids + cache_position[0]
+
+        outputs = self.language_model(
+            input_ids=None,
+            position_ids=position_ids,
+            attention_mask=attention_mask,
+            past_key_values=past_key_values,
+            inputs_embeds=inputs_embeds,
+            use_cache=use_cache,
+            output_attentions=output_attentions,
+            output_hidden_states=output_hidden_states,
+            return_dict=True,
+            cache_position=cache_position,
+            **kwargs,
+        )
+
+        output = MageVLModelOutputWithPast(
+            last_hidden_state=outputs.last_hidden_state,
+            past_key_values=outputs.past_key_values,
+            hidden_states=outputs.hidden_states,
+            attentions=outputs.attentions,
+        )
+        return output if return_dict else output.to_tuple()
+
+
+@auto_docstring
+class MageVLForConditionalGeneration(MageVLPreTrainedModel, GenerationMixin):
+    _tied_weights_keys = {"lm_head.weight": "model.language_model.embed_tokens.weight"}
+    # Reference: fix gemma3 grad acc #37208
+    accepts_loss_kwargs = False
+
+    def __init__(self, config):
+        super().__init__(config)
+        self.model = MageVLModel(config)
+        self.lm_head = nn.Linear(config.text_config.hidden_size, config.text_config.vocab_size, bias=False)
+        self.post_init()
+
+    @classmethod
+    def from_pretrained(cls, pretrained_model_name_or_path, *args, **kwargs):
+        model = super().from_pretrained(pretrained_model_name_or_path, *args, **kwargs)
+        model.model._streammind_model_path = str(pretrained_model_name_or_path)
+        return model
+
+    def get_input_embeddings(self):
+        return self.model.get_input_embeddings()
+
+    def set_input_embeddings(self, value):
+        self.model.set_input_embeddings(value)
+
+    def set_decoder(self, decoder):
+        self.model.set_decoder(decoder)
+
+    def get_decoder(self):
+        return self.model.get_decoder()
+
+    def get_video_features(
+        self,
+        pixel_values_videos: torch.FloatTensor,
+        video_grid_thw: Optional[torch.LongTensor] = None,
+        patch_positions=None,
+    ):
+        return self.model.get_video_features(pixel_values_videos, video_grid_thw, patch_positions=patch_positions)
+
+    def get_image_features(self, pixel_values: torch.FloatTensor, image_grid_thw: Optional[torch.LongTensor] = None):
+        return self.model.get_image_features(pixel_values, image_grid_thw)
+
+    # Make modules available through conditional class for BC
+    @property
+    def language_model(self):
+        return self.model.language_model
+
+    @property
+    def visual(self):
+        return self.model.visual
+
+    def streammind_gate_forward(self, pixel_values, image_grid_thw, patch_positions=None):
+        return self.model.streammind_gate_forward(
+            pixel_values,
+            image_grid_thw,
+            patch_positions=patch_positions,
+        )
+
+    def streammind_gate_forward_segments(self, segments):
+        return self.model.streammind_gate_forward_segments(segments)
+
+    @can_return_tuple
+    @auto_docstring
+    def forward(
+        self,
+        input_ids: Optional[torch.LongTensor] = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        past_key_values: Optional[Cache] = None,
+        inputs_embeds: Optional[torch.FloatTensor] = None,
+        labels: Optional[torch.LongTensor] = None,
+        use_cache: Optional[bool] = None,
+        output_attentions: Optional[bool] = None,
+        output_hidden_states: Optional[bool] = None,
+        pixel_values: Optional[torch.Tensor] = None,
+        pixel_values_videos: Optional[torch.FloatTensor] = None,
+        image_grid_thw: Optional[torch.LongTensor] = None,
+        patch_positions: Optional[torch.LongTensor] = None,
+        video_grid_thw: Optional[torch.LongTensor] = None,
+        cache_position: Optional[torch.LongTensor] = None,
+        second_per_grid_ts: Optional[torch.Tensor] = None,
+        logits_to_keep: Union[int, torch.Tensor] = 0,
+        **kwargs: Unpack[TransformersKwargs],
+    ) -> Union[tuple, MageVLCausalLMOutputWithPast]:
+        r"""
+        labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
+            Labels for computing the masked language modeling loss. Indices should either be in `[0, ...,
+            config.vocab_size]` or -100 (see `input_ids` docstring). Tokens with indices set to `-100` are ignored
+            (masked), the loss is only computed for the tokens with labels in `[0, ..., config.vocab_size]`.
+        image_grid_thw (`torch.LongTensor` of shape `(num_images, 3)`, *optional*):
+            The temporal, height and width of feature shape of each image in LLM.
+        video_grid_thw (`torch.LongTensor` of shape `(num_videos, 3)`, *optional*):
+            The temporal, height and width of feature shape of each video in LLM.
+        patch_positions (`torch.LongTensor` of shape `(total_patches, 3)` or `(1, total_patches, 3)`, *optional*):
+            Explicit per-patch `(t, h, w)` position indices used by the vision tower to compute 3D rotary
+            position embeddings (and the optional absolute position embedding inside the patch merger).
+            `total_patches` is the sum of `t * h * w` across all images and videos in the batch, matching
+            the layout produced by the Qwen2VL-style image processor.
+        second_per_grid_ts (`torch.Tensor` of shape `(num_videos)`, *optional*):
+            The time interval (in seconds) for each grid along the temporal dimension in the 3D position IDs.
+        cache_position (`torch.LongTensor` of shape `(sequence_length)`, *optional*):
+              Indices depicting the position of the input sequence tokens in the sequence. Contrarily to
+              `position_ids`, this tensor is not affected by padding.
+
+        Note (native-video alias):
+            The companion ``MageVLProcessor.__call__(videos=...)`` does NOT
+            pass ``pixel_values_videos`` / ``video_grid_thw`` / ``second_per_grid_ts``
+            to this forward. Instead it aliases the video patch tensor as
+            ``pixel_values=`` and ``image_grid_thw=``, so video inputs share the
+            same code path as multi-image inputs (Mage-ViT is purely
+            spatial; temporal information is carried by per-frame ``<X.X seconds>``
+            text tags emitted by the processor). The ``*_videos`` and
+            ``second_per_grid_ts`` kwargs are kept declared here only for API
+            completeness and future use (e.g. 3D mRoPE / ``get_rope_index``); they
+            are NOT consumed by the current Mage-ViT encoder.
+
+        Example:
+
+        ```python
+        >>> from PIL import Image
+        >>> import requests
+        >>> from transformers import AutoProcessor, MageVLForConditionalGeneration
+
+        >>> model = MageVLForConditionalGeneration.from_pretrained("microsoft/Mage-VL", trust_remote_code=True)
+        >>> processor = AutoProcessor.from_pretrained("microsoft/Mage-VL", trust_remote_code=True)
+
+        >>> messages = [
+            {
+                "role": "user",
+                "content": [
+                    {"type": "image"},
+                    {"type": "text", "text": "What is shown in this image?"},
+                ],
+            },
+        ]
+        >>> url = "https://www.ilankelman.org/stopsigns/australia.jpg"
+        >>> image = Image.open(requests.get(url, stream=True).raw)
+
+        >>> text = processor.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
+        >>> inputs = processor(text=[text], images=[image], return_tensors="pt")
+
+        >>> # Generate
+        >>> generate_ids = model.generate(inputs.input_ids, max_length=30)
+        >>> tokenizer.batch_decode(generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0]
+        "The image shows a street scene with a red stop sign in the foreground. In the background, there is a large red gate with Chinese characters ..."
+        ```"""
+        output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
+        output_hidden_states = (
+            output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
+        )
+        outputs = self.model(
+            input_ids=input_ids,
+            pixel_values=pixel_values,
+            pixel_values_videos=pixel_values_videos,
+            image_grid_thw=image_grid_thw,
+            patch_positions=patch_positions,
+            video_grid_thw=video_grid_thw,
+            second_per_grid_ts=second_per_grid_ts,
+            position_ids=position_ids,
+            attention_mask=attention_mask,
+            past_key_values=past_key_values,
+            inputs_embeds=inputs_embeds,
+            use_cache=use_cache,
+            output_attentions=output_attentions,
+            output_hidden_states=output_hidden_states,
+            return_dict=True,
+            cache_position=cache_position,
+            **kwargs,
+        )
+
+        hidden_states = outputs[0]
+
+        # Only compute necessary logits, and do not upcast them to float if we are not computing the loss
+        slice_indices = slice(-logits_to_keep, None) if isinstance(logits_to_keep, int) else logits_to_keep
+        logits = self.lm_head(hidden_states[:, slice_indices, :])
+
+        loss = None
+        if labels is not None:
+            loss = self.loss_function(
+                logits=logits, labels=labels, vocab_size=self.config.text_config.vocab_size, **kwargs
+            )
+
+        return MageVLCausalLMOutputWithPast(
+            loss=loss,
+            logits=logits,
+            past_key_values=outputs.past_key_values,
+            hidden_states=outputs.hidden_states,
+            attentions=outputs.attentions,
+        )
+
+
+
+    def prepare_inputs_for_generation(
+        self,
+        input_ids,
+        past_key_values=None,
+        attention_mask=None,
+        inputs_embeds=None,
+        cache_position=None,
+        position_ids=None,
+        use_cache=True,
+        pixel_values=None,
+        pixel_values_videos=None,
+        image_grid_thw=None,
+        patch_positions=None,
+        video_grid_thw=None,
+        second_per_grid_ts=None,
+        is_first_iteration=False,
+        **kwargs,
+    ):
+        # Overwritten -- in specific circumstances we don't want to forward image inputs to the model
+        model_inputs = super().prepare_inputs_for_generation(
+            input_ids,
+            past_key_values=past_key_values,
+            attention_mask=attention_mask,
+            inputs_embeds=inputs_embeds,
+            cache_position=cache_position,
+            position_ids=position_ids,
+            pixel_values=pixel_values,
+            pixel_values_videos=pixel_values_videos,
+            image_grid_thw=image_grid_thw,
+            video_grid_thw=video_grid_thw,
+            second_per_grid_ts=second_per_grid_ts,
+            patch_positions=patch_positions,
+            use_cache=use_cache,
+            is_first_iteration=is_first_iteration,
+            **kwargs,
+        )
+
+        # After the prefill iteration, drop image inputs so the vision tower
+        # isn't re-run on decode steps. Gating on `is_first_iteration` (the
+        # Qwen3-VL convention) is the only reliable signal in transformers
+        # 5.x: `past_key_values` is non-None even on the first call (an empty
+        # DynamicCache is created up-front by `generate`), and `cache_position`
+        # may be `None` for remote-code models.
+        if not is_first_iteration and use_cache:
+            model_inputs["pixel_values"] = None
+            model_inputs["pixel_values_videos"] = None
+
+        return model_inputs
+
+    def _get_image_nums_and_video_nums(
+        self,
+        input_ids: Optional[torch.LongTensor],
+        inputs_embeds: Optional[torch.Tensor] = None,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        """
+        Get the number of images and videos for each sample to calculate the separation length of the sample tensor.
+        These parameters are not passed through the processor to avoid unpredictable impacts from interface modifications.
+
+        Args:
+            input_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`):
+                Indices of input sequence tokens in the vocabulary.
+
+        Returns:
+            image_nums (`torch.LongTensor` of shape `(batch_size, num_images_sample)`)
+            video_nums (`torch.LongTensor` of shape `(batch_size, num_videos_sample)`)
+        """
+        image_token_id = self.config.image_token_id
+        video_token_id = self.config.video_token_id
+        vision_start_token_id = self.config.vision_start_token_id
+
+        if inputs_embeds is not None:
+            vision_start_mask = (
+                inputs_embeds
+                == self.get_input_embeddings()(
+                    torch.tensor(vision_start_token_id, dtype=torch.long, device=inputs_embeds.device)
+                )
+            )[..., 0]
+            image_mask = (
+                inputs_embeds
+                == self.get_input_embeddings()(
+                    torch.tensor(image_token_id, dtype=torch.long, device=inputs_embeds.device)
+                )
+            )[..., 0]
+            video_mask = (
+                inputs_embeds
+                == self.get_input_embeddings()(
+                    torch.tensor(video_token_id, dtype=torch.long, device=inputs_embeds.device)
+                )
+            )[..., 0]
+        else:
+            vision_start_mask = input_ids == vision_start_token_id
+            image_mask = input_ids == image_token_id
+            video_mask = input_ids == video_token_id
+
+        vision_first_mask = torch.roll(vision_start_mask, shifts=1, dims=1)
+        image_nums = torch.sum(vision_first_mask & image_mask, dim=1)
+        video_nums = torch.sum(vision_first_mask & video_mask, dim=1)
+
+        return image_nums, video_nums
+
+    def _expand_inputs_for_generation(
+        self,
+        expand_size: int = 1,
+        is_encoder_decoder: bool = False,
+        input_ids: Optional[torch.LongTensor] = None,
+        **model_kwargs,
+    ) -> tuple[torch.LongTensor, dict[str, Any]]:
+        # Overwritten -- Support for expanding tensors without a batch size dimension
+        # e.g., pixel_values, image_grid_thw, pixel_values_videos, video_grid_thw, second_per_grid_t
+        # pixel_values.shape[0] is sum(seqlen_images for samples)
+        # image_grid_thw.shape[0] is sum(num_images for samples)
+
+        if expand_size == 1:
+            return input_ids, model_kwargs
+
+        visual_keys = [
+            "pixel_values",
+            "image_grid_thw",
+            "pixel_values_videos",
+            "video_grid_thw",
+            "second_per_grid_ts",
+            "patch_positions",
+        ]
+
+        def _expand_dict_for_generation_visual(dict_to_expand):
+            image_grid_thw = model_kwargs.get("image_grid_thw", None)
+            video_grid_thw = model_kwargs.get("video_grid_thw", None)
+            image_nums, video_nums = self._get_image_nums_and_video_nums(
+                input_ids, inputs_embeds=model_kwargs.get("inputs_embeds", None)
+            )
+
+            def _repeat_interleave_samples(x, lengths, repeat_times):
+                samples = torch.split(x, lengths)
+                repeat_args = [repeat_times] + [1] * (x.dim() - 1)
+                result = torch.cat([sample.repeat(*repeat_args) for sample in samples], dim=0)
+                return result
+
+            for key in dict_to_expand:
+                if key == "pixel_values":
+                    # split images into samples
+                    samples = torch.split(image_grid_thw, list(image_nums))
+                    # compute the sequence length of images for each sample
+                    lengths = [torch.prod(sample, dim=1).sum() for sample in samples]
+                    dict_to_expand[key] = _repeat_interleave_samples(
+                        dict_to_expand[key], lengths=lengths, repeat_times=expand_size
+                    )
+                elif key == "image_grid_thw":
+                    # get the num of images for each sample
+                    lengths = list(image_nums)
+                    dict_to_expand[key] = _repeat_interleave_samples(
+                        dict_to_expand[key], lengths=lengths, repeat_times=expand_size
+                    )
+                elif key == "pixel_values_videos":
+                    samples = torch.split(video_grid_thw, list(video_nums))
+                    lengths = [torch.prod(sample, dim=1).sum() for sample in samples]
+                    dict_to_expand[key] = _repeat_interleave_samples(
+                        dict_to_expand[key], lengths=lengths, repeat_times=expand_size
+                    )
+                elif key == "video_grid_thw":
+                    lengths = list(video_nums)
+                    dict_to_expand[key] = _repeat_interleave_samples(
+                        dict_to_expand[key], lengths=lengths, repeat_times=expand_size
+                    )
+                elif key == "second_per_grid_ts":
+                    dict_to_expand[key] = _repeat_interleave_samples(
+                        dict_to_expand[key], lengths=list(video_nums), repeat_times=expand_size
+                    )
+                elif key == "patch_positions":
+                    if image_grid_thw is not None and image_grid_thw.numel() > 0 and image_nums.sum() > 0:
+                        samples = torch.split(image_grid_thw, list(image_nums))
+                        lengths = [torch.prod(sample, dim=1).sum() for sample in samples]
+                    elif video_grid_thw is not None and video_grid_thw.numel() > 0 and video_nums.sum() > 0:
+                        samples = torch.split(video_grid_thw, list(video_nums))
+                        lengths = [torch.prod(sample, dim=1).sum() for sample in samples]
+                    else:
+                        continue
+                    dict_to_expand[key] = _repeat_interleave_samples(
+                        dict_to_expand[key], lengths=lengths, repeat_times=expand_size
+                    )
+            return dict_to_expand
+
+        def _expand_dict_for_generation(dict_to_expand):
+            for key in dict_to_expand:
+                if (
+                    key != "cache_position"
+                    and dict_to_expand[key] is not None
+                    and isinstance(dict_to_expand[key], torch.Tensor)
+                    and key not in visual_keys
+                ):
+                    dict_to_expand[key] = dict_to_expand[key].repeat_interleave(expand_size, dim=0)
+            return dict_to_expand
+
+        model_kwargs = _expand_dict_for_generation_visual(model_kwargs)
+
+        if input_ids is not None:
+            input_ids = input_ids.repeat_interleave(expand_size, dim=0)
+
+        model_kwargs = _expand_dict_for_generation(model_kwargs)
+
+        if is_encoder_decoder:
+            if model_kwargs.get("encoder_outputs") is None:
+                raise ValueError("If `is_encoder_decoder` is True, make sure that `encoder_outputs` is defined.")
+            model_kwargs["encoder_outputs"] = _expand_dict_for_generation(model_kwargs["encoder_outputs"])
+
+        return input_ids, model_kwargs
+
+
+__all__ = [
+    "MageVLForConditionalGeneration",
+    "MageVLModel",
+    "MageVLPreTrainedModel",
+]

neural_codec/DCVC/LICENSE.txt

@@ -0,0 +1,21 @@
+Copyright (c) Microsoft Corporation.
+
+MIT License
+
+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.

neural_codec/DCVC/NOTICE.txt

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

neural_codec/DCVC/src/cpp/py_rans/py_rans.cpp

@@ -0,0 +1,393 @@
+// Copyright (c) Microsoft Corporation.
+// Licensed under the MIT License.
+
+#include "py_rans.h"
+
+#include <algorithm>
+#include <cmath>
+#include <future>
+#include <numeric>
+#include <vector>
+
+namespace py = pybind11;
+
+RansEncoder::RansEncoder()
+{
+    m_encoder0 = std::make_shared<RansEncoderLibMultiThread>();
+    m_encoder1 = std::make_shared<RansEncoderLibMultiThread>();
+}
+
+void RansEncoder::encode_y(const py::array_t<int16_t>& symbols, const int cdf_group_index)
+{
+    py::buffer_info symbols_buf = symbols.request();
+    int16_t* symbols_ptr = static_cast<int16_t*>(symbols_buf.ptr);
+
+    int symbolSize = static_cast<int>(symbols.size());
+    if (m_use_two_encoders) {
+        int symbolSize0 = symbolSize / 2;
+        int symbolSize1 = symbolSize - symbolSize0;
+
+        auto vec_symbols0 = std::make_shared<std::vector<int16_t>>(symbolSize0);
+        memcpy(vec_symbols0->data(), symbols_ptr, symbolSize0 * sizeof(int16_t));
+        m_encoder0->encode_y(vec_symbols0, cdf_group_index);
+        auto vec_symbols1 = std::make_shared<std::vector<int16_t>>(symbolSize1);
+        memcpy(vec_symbols1->data(), symbols_ptr + symbolSize0, symbolSize1 * sizeof(int16_t));
+        m_encoder1->encode_y(vec_symbols1, cdf_group_index);
+    } else {
+        auto vec_symbols0 = std::make_shared<std::vector<int16_t>>(symbolSize);
+        memcpy(vec_symbols0->data(), symbols_ptr, symbolSize * sizeof(int16_t));
+        m_encoder0->encode_y(vec_symbols0, cdf_group_index);
+    }
+}
+
+void RansEncoder::encode_z(const py::array_t<int8_t>& symbols, const int cdf_group_index,
+                           const int start_offset, const int per_channel_size)
+{
+    py::buffer_info symbols_buf = symbols.request();
+    int8_t* symbols_ptr = static_cast<int8_t*>(symbols_buf.ptr);
+
+    int symbolSize = static_cast<int>(symbols.size());
+    if (m_use_two_encoders) {
+        int symbolSize0 = symbolSize / 2;
+        int symbolSize1 = symbolSize - symbolSize0;
+        int channel_half = symbolSize0 / per_channel_size;
+
+        auto vec_symbols0 = std::make_shared<std::vector<int8_t>>(symbolSize0);
+        memcpy(vec_symbols0->data(), symbols_ptr, symbolSize0 * sizeof(int8_t));
+        m_encoder0->encode_z(vec_symbols0, cdf_group_index, start_offset, per_channel_size);
+        auto vec_symbols1 = std::make_shared<std::vector<int8_t>>(symbolSize1);
+        memcpy(vec_symbols1->data(), symbols_ptr + symbolSize0, symbolSize1 * sizeof(int8_t));
+        m_encoder1->encode_z(vec_symbols1, cdf_group_index, start_offset + channel_half,
+                             per_channel_size);
+    } else {
+        auto vec_symbols0 = std::make_shared<std::vector<int8_t>>(symbolSize);
+        memcpy(vec_symbols0->data(), symbols_ptr, symbolSize * sizeof(int8_t));
+        m_encoder0->encode_z(vec_symbols0, cdf_group_index, start_offset, per_channel_size);
+    }
+}
+
+int RansEncoder::add_cdf(const py::array_t<int32_t>& cdfs, const py::array_t<int32_t>& cdfs_sizes,
+                         const py::array_t<int32_t>& offsets)
+{
+    py::buffer_info cdfs_sizes_buf = cdfs_sizes.request();
+    py::buffer_info offsets_buf = offsets.request();
+    int32_t* cdfs_sizes_ptr = static_cast<int32_t*>(cdfs_sizes_buf.ptr);
+    int32_t* offsets_ptr = static_cast<int32_t*>(offsets_buf.ptr);
+
+    int cdf_num = static_cast<int>(cdfs_sizes.size());
+    auto vec_cdfs_sizes = std::make_shared<std::vector<int32_t>>(cdf_num);
+    memcpy(vec_cdfs_sizes->data(), cdfs_sizes_ptr, sizeof(int32_t) * cdf_num);
+    auto vec_offsets = std::make_shared<std::vector<int32_t>>(offsets.size());
+    memcpy(vec_offsets->data(), offsets_ptr, sizeof(int32_t) * offsets.size());
+
+    int per_vector_size = static_cast<int>(cdfs.size() / cdf_num);
+    auto vec_cdfs = std::make_shared<std::vector<std::vector<int32_t>>>(cdf_num);
+    auto cdfs_raw = cdfs.unchecked<2>();
+    for (int i = 0; i < cdf_num; i++) {
+        std::vector<int32_t> t(per_vector_size);
+        memcpy(t.data(), cdfs_raw.data(i, 0), sizeof(int32_t) * per_vector_size);
+        vec_cdfs->at(i) = t;
+    }
+
+    int cdf_idx = m_encoder0->add_cdf(vec_cdfs, vec_cdfs_sizes, vec_offsets);
+    m_encoder1->add_cdf(vec_cdfs, vec_cdfs_sizes, vec_offsets);
+    return cdf_idx;
+}
+
+void RansEncoder::empty_cdf_buffer()
+{
+    m_encoder0->empty_cdf_buffer();
+    m_encoder1->empty_cdf_buffer();
+}
+
+void RansEncoder::flush()
+{
+    m_encoder0->flush();
+    m_encoder1->flush();
+}
+
+py::array_t<uint8_t> RansEncoder::get_encoded_stream()
+{
+    if (m_use_two_encoders) {
+        auto result0 = m_encoder0->get_encoded_stream();
+        int nbytes0 = static_cast<int>(result0->size());
+        auto result1 = m_encoder1->get_encoded_stream();
+        int nbytes1 = static_cast<int>(result1->size());
+
+        int identical_bytes = 0;
+        int check_bytes = std::min(nbytes0, nbytes1);
+        check_bytes = std::min(check_bytes, 8);
+        for (int i = 0; i < check_bytes; i++) {
+            if (result0->at(nbytes0 - 1 - i) != 0) {
+                break;
+            }
+            if (result1->at(nbytes1 - 1 - i) != 0) {
+                break;
+            }
+            identical_bytes++;
+        }
+        if (identical_bytes == 0 && result0->at(nbytes0 - 1) == result1->at(nbytes1 - 1)) {
+            identical_bytes = 1;
+        }
+
+        py::array_t<uint8_t> stream(nbytes0 + nbytes1 - identical_bytes);
+        py::buffer_info stream_buf = stream.request();
+        uint8_t* stream_ptr = static_cast<uint8_t*>(stream_buf.ptr);
+
+        std::copy(result0->begin(), result0->end(), stream_ptr);
+        std::reverse_copy(result1->begin(), result1->end() - identical_bytes, stream_ptr + nbytes0);
+        return stream;
+    }
+
+    auto result0 = m_encoder0->get_encoded_stream();
+    int nbytes0 = static_cast<int>(result0->size());
+
+    py::array_t<uint8_t> stream(nbytes0);
+    py::buffer_info stream_buf = stream.request();
+    uint8_t* stream_ptr = static_cast<uint8_t*>(stream_buf.ptr);
+
+    std::copy(result0->begin(), result0->end(), stream_ptr);
+    return stream;
+}
+
+void RansEncoder::reset()
+{
+    m_encoder0->reset();
+    m_encoder1->reset();
+}
+
+void RansEncoder::set_use_two_encoders(bool b)
+{
+    m_use_two_encoders = b;
+}
+
+bool RansEncoder::get_use_two_encoders()
+{
+    return m_use_two_encoders;
+}
+
+RansDecoder::RansDecoder()
+{
+    m_decoder0 = std::make_shared<RansDecoderLibMultiThread>();
+    m_decoder1 = std::make_shared<RansDecoderLibMultiThread>();
+}
+
+void RansDecoder::set_stream(const py::array_t<uint8_t>& encoded)
+{
+    py::buffer_info encoded_buf = encoded.request();
+    const uint8_t* encoded_ptr = static_cast<uint8_t*>(encoded_buf.ptr);
+    const int encoded_size = static_cast<int>(encoded.size());
+    auto stream0 = std::make_shared<std::vector<uint8_t>>(encoded.size());
+    std::copy(encoded_ptr, encoded_ptr + encoded_size, stream0->data());
+    m_decoder0->set_stream(stream0);
+    if (m_use_two_decoders) {
+        auto stream1 = std::make_shared<std::vector<uint8_t>>(encoded.size());
+        std::reverse_copy(encoded_ptr, encoded_ptr + encoded_size, stream1->data());
+        m_decoder1->set_stream(stream1);
+    }
+}
+
+void RansDecoder::decode_y(const py::array_t<uint8_t>& indexes, const int cdf_group_index)
+{
+    py::buffer_info indexes_buf = indexes.request();
+    uint8_t* indexes_ptr = static_cast<uint8_t*>(indexes_buf.ptr);
+
+    int indexSize = static_cast<int>(indexes.size());
+    if (m_use_two_decoders) {
+        int indexSize0 = indexSize / 2;
+        int indexSize1 = indexSize - indexSize0;
+
+        auto vec_indexes0 = std::make_shared<std::vector<uint8_t>>(indexSize0);
+        std::copy(indexes_ptr, indexes_ptr + indexSize0, vec_indexes0->data());
+        m_decoder0->decode_y(vec_indexes0, cdf_group_index);
+
+        auto vec_indexes1 = std::make_shared<std::vector<uint8_t>>(indexSize1);
+        std::copy(indexes_ptr + indexSize0, indexes_ptr + indexSize, vec_indexes1->data());
+        m_decoder1->decode_y(vec_indexes1, cdf_group_index);
+    } else {
+        auto vec_indexes0 = std::make_shared<std::vector<uint8_t>>(indexSize);
+        std::copy(indexes_ptr, indexes_ptr + indexSize, vec_indexes0->data());
+        m_decoder0->decode_y(vec_indexes0, cdf_group_index);
+    }
+}
+
+py::array_t<int8_t> RansDecoder::decode_and_get_y(const py::array_t<uint8_t>& indexes,
+                                                  const int cdf_group_index)
+{
+    decode_y(indexes, cdf_group_index);
+    return get_decoded_tensor();
+}
+
+void RansDecoder::decode_z(const int total_size, const int cdf_group_index, const int start_offset,
+                           const int per_channel_size)
+{
+    if (m_use_two_decoders) {
+        int symbolSize0 = total_size / 2;
+        int symbolSize1 = total_size - symbolSize0;
+        int channel_half = symbolSize0 / per_channel_size;
+        m_decoder0->decode_z(symbolSize0, cdf_group_index, start_offset, per_channel_size);
+        m_decoder1->decode_z(symbolSize1, cdf_group_index, start_offset + channel_half,
+                             per_channel_size);
+    } else {
+        m_decoder0->decode_z(total_size, cdf_group_index, start_offset, per_channel_size);
+    }
+}
+
+py::array_t<int8_t> RansDecoder::get_decoded_tensor()
+{
+    if (m_use_two_decoders) {
+        auto result0 = m_decoder0->get_decoded_tensor();
+        const int total_size0 = static_cast<int>(result0->size());
+
+        auto result1 = m_decoder1->get_decoded_tensor();
+        const int total_size1 = static_cast<int>(result1->size());
+        py::array_t<int8_t> output(total_size0 + total_size1);
+        py::buffer_info buf = output.request();
+        int8_t* buf_ptr = static_cast<int8_t*>(buf.ptr);
+        std::copy(result0->begin(), result0->end(), buf_ptr);
+        std::copy(result1->begin(), result1->end(), buf_ptr + total_size0);
+
+        return output;
+    }
+
+    auto result0 = m_decoder0->get_decoded_tensor();
+    const int total_size0 = static_cast<int>(result0->size());
+
+    py::array_t<int8_t> output(total_size0);
+    py::buffer_info buf = output.request();
+    int8_t* buf_ptr = static_cast<int8_t*>(buf.ptr);
+    std::copy(result0->begin(), result0->end(), buf_ptr);
+
+    return output;
+}
+
+int RansDecoder::add_cdf(const py::array_t<int32_t>& cdfs, const py::array_t<int32_t>& cdfs_sizes,
+                         const py::array_t<int32_t>& offsets)
+{
+    py::buffer_info cdfs_sizes_buf = cdfs_sizes.request();
+    py::buffer_info offsets_buf = offsets.request();
+    int32_t* cdfs_sizes_ptr = static_cast<int32_t*>(cdfs_sizes_buf.ptr);
+    int32_t* offsets_ptr = static_cast<int32_t*>(offsets_buf.ptr);
+
+    int cdf_num = static_cast<int>(cdfs_sizes.size());
+    auto vec_cdfs_sizes = std::make_shared<std::vector<int32_t>>(cdf_num);
+    memcpy(vec_cdfs_sizes->data(), cdfs_sizes_ptr, sizeof(int32_t) * cdf_num);
+    auto vec_offsets = std::make_shared<std::vector<int32_t>>(offsets.size());
+    memcpy(vec_offsets->data(), offsets_ptr, sizeof(int32_t) * offsets.size());
+
+    int per_vector_size = static_cast<int>(cdfs.size() / cdf_num);
+    auto vec_cdfs = std::make_shared<std::vector<std::vector<int32_t>>>(cdf_num);
+    auto cdfs_raw = cdfs.unchecked<2>();
+    for (int i = 0; i < cdf_num; i++) {
+        std::vector<int32_t> t(per_vector_size);
+        memcpy(t.data(), cdfs_raw.data(i, 0), sizeof(int32_t) * per_vector_size);
+        vec_cdfs->at(i) = t;
+    }
+    int cdf_idx = m_decoder0->add_cdf(vec_cdfs, vec_cdfs_sizes, vec_offsets);
+    m_decoder1->add_cdf(vec_cdfs, vec_cdfs_sizes, vec_offsets);
+    return cdf_idx;
+}
+
+void RansDecoder::empty_cdf_buffer()
+{
+    m_decoder0->empty_cdf_buffer();
+    m_decoder1->empty_cdf_buffer();
+}
+
+void RansDecoder::set_use_two_decoders(bool b)
+{
+    m_use_two_decoders = b;
+}
+
+bool RansDecoder::get_use_two_decoders()
+{
+    return m_use_two_decoders;
+}
+
+std::vector<uint32_t> pmf_to_quantized_cdf(const std::vector<float>& pmf, int precision)
+{
+    /* NOTE(begaintj): ported from `ryg_rans` public implementation. Not optimal
+     * although it's only run once per model after training. See TF/compression
+     * implementation for an optimized version. */
+
+    std::vector<uint32_t> cdf(pmf.size() + 1);
+    cdf[0] = 0; /* freq 0 */
+
+    std::transform(pmf.begin(), pmf.end(), cdf.begin() + 1, [=](float p) {
+        return static_cast<uint32_t>(std::round(p * (1 << precision)) + 0.5);
+    });
+
+    const uint32_t total = std::accumulate(cdf.begin(), cdf.end(), 0);
+
+    std::transform(cdf.begin(), cdf.end(), cdf.begin(), [precision, total](uint32_t p) {
+        return static_cast<uint32_t>((((1ull << precision) * p) / total));
+    });
+
+    std::partial_sum(cdf.begin(), cdf.end(), cdf.begin());
+    cdf.back() = 1 << precision;
+
+    for (int i = 0; i < static_cast<int>(cdf.size() - 1); ++i) {
+        if (cdf[i] == cdf[i + 1]) {
+            /* Try to steal frequency from low-frequency symbols */
+            uint32_t best_freq = ~0u;
+            int best_steal = -1;
+            for (int j = 0; j < static_cast<int>(cdf.size()) - 1; ++j) {
+                uint32_t freq = cdf[j + 1] - cdf[j];
+                if (freq > 1 && freq < best_freq) {
+                    best_freq = freq;
+                    best_steal = j;
+                }
+            }
+
+            assert(best_steal != -1);
+
+            if (best_steal < i) {
+                for (int j = best_steal + 1; j <= i; ++j) {
+                    cdf[j]--;
+                }
+            } else {
+                assert(best_steal > i);
+                for (int j = i + 1; j <= best_steal; ++j) {
+                    cdf[j]++;
+                }
+            }
+        }
+    }
+
+    assert(cdf[0] == 0);
+    assert(cdf.back() == (1u << precision));
+    for (int i = 0; i < static_cast<int>(cdf.size()) - 1; ++i) {
+        assert(cdf[i + 1] > cdf[i]);
+    }
+
+    return cdf;
+}
+
+PYBIND11_MODULE(MLCodec_extensions_cpp, m)
+{
+    py::class_<RansEncoder>(m, "RansEncoder")
+        .def(py::init<>())
+        .def("encode_y", &RansEncoder::encode_y)
+        .def("encode_z", &RansEncoder::encode_z)
+        .def("flush", &RansEncoder::flush)
+        .def("get_encoded_stream", &RansEncoder::get_encoded_stream)
+        .def("reset", &RansEncoder::reset)
+        .def("add_cdf", &RansEncoder::add_cdf)
+        .def("empty_cdf_buffer", &RansEncoder::empty_cdf_buffer)
+        .def("set_use_two_encoders", &RansEncoder::set_use_two_encoders)
+        .def("get_use_two_encoders", &RansEncoder::get_use_two_encoders);
+
+    py::class_<RansDecoder>(m, "RansDecoder")
+        .def(py::init<>())
+        .def("set_stream", &RansDecoder::set_stream)
+        .def("decode_y", &RansDecoder::decode_y)
+        .def("decode_and_get_y", &RansDecoder::decode_and_get_y)
+        .def("decode_z", &RansDecoder::decode_z)
+        .def("get_decoded_tensor", &RansDecoder::get_decoded_tensor)
+        .def("add_cdf", &RansDecoder::add_cdf)
+        .def("empty_cdf_buffer", &RansDecoder::empty_cdf_buffer)
+        .def("set_use_two_decoders", &RansDecoder::set_use_two_decoders)
+        .def("get_use_two_decoders", &RansDecoder::get_use_two_decoders);
+
+    m.def("pmf_to_quantized_cdf", &pmf_to_quantized_cdf, "Return quantized CDF for a given PMF");
+}

neural_codec/DCVC/src/cpp/py_rans/py_rans.h

@@ -0,0 +1,71 @@
+// Copyright (c) Microsoft Corporation.
+// Licensed under the MIT License.
+
+#pragma once
+#include "rans.h"
+#include <memory>
+
+#include <pybind11/numpy.h>
+#include <pybind11/pybind11.h>
+#include <pybind11/stl.h>
+
+namespace py = pybind11;
+
+// the classes in this file only perform the type conversion
+// from python type (numpy) to C++ type (vector)
+class RansEncoder {
+public:
+    RansEncoder();
+
+    RansEncoder(const RansEncoder&) = delete;
+    RansEncoder(RansEncoder&&) = delete;
+    RansEncoder& operator=(const RansEncoder&) = delete;
+    RansEncoder& operator=(RansEncoder&&) = delete;
+
+    void encode_y(const py::array_t<int16_t>& symbols, const int cdf_group_index);
+    void encode_z(const py::array_t<int8_t>& symbols, const int cdf_group_index,
+                  const int start_offset, const int per_channel_size);
+    void flush();
+    py::array_t<uint8_t> get_encoded_stream();
+    void reset();
+    int add_cdf(const py::array_t<int32_t>& cdfs, const py::array_t<int32_t>& cdfs_sizes,
+                const py::array_t<int32_t>& offsets);
+    void empty_cdf_buffer();
+    void set_use_two_encoders(bool b);
+    bool get_use_two_encoders();
+
+private:
+    std::shared_ptr<RansEncoderLib> m_encoder0;
+    std::shared_ptr<RansEncoderLib> m_encoder1;
+    bool m_use_two_encoders{ false };
+};
+
+class RansDecoder {
+public:
+    RansDecoder();
+
+    RansDecoder(const RansDecoder&) = delete;
+    RansDecoder(RansDecoder&&) = delete;
+    RansDecoder& operator=(const RansDecoder&) = delete;
+    RansDecoder& operator=(RansDecoder&&) = delete;
+
+    void set_stream(const py::array_t<uint8_t>&);
+
+    void decode_y(const py::array_t<uint8_t>& indexes, const int cdf_group_index);
+    py::array_t<int8_t> decode_and_get_y(const py::array_t<uint8_t>& indexes, const int cdf_group_index);
+    void decode_z(const int total_size, const int cdf_group_index, const int start_offset,
+                  const int per_channel_size);
+    py::array_t<int8_t> get_decoded_tensor();
+    int add_cdf(const py::array_t<int32_t>& cdfs, const py::array_t<int32_t>& cdfs_sizes,
+                const py::array_t<int32_t>& offsets);
+    void empty_cdf_buffer();
+    void set_use_two_decoders(bool b);
+    bool get_use_two_decoders();
+
+private:
+    std::shared_ptr<RansDecoderLib> m_decoder0;
+    std::shared_ptr<RansDecoderLib> m_decoder1;
+    bool m_use_two_decoders{ false };
+};
+
+std::vector<uint32_t> pmf_to_quantized_cdf(const std::vector<float>& pmf, int precision);

neural_codec/DCVC/src/cpp/py_rans/rans.cpp

@@ -0,0 +1,534 @@
+/* Copyright 2020 InterDigital Communications, Inc.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+/* Rans64 extensions from:
+ * https://fgiesen.wordpress.com/2015/12/21/rans-in-practice/
+ * Unbounded range coding from:
+ * https://github.com/tensorflow/compression/blob/master/tensorflow_compression/cc/kernels/unbounded_index_range_coding_kernels.cc
+ **/
+
+#include "rans.h"
+
+#include <algorithm>
+#include <cassert>
+#include <cstring>
+
+constexpr uint16_t bypass_precision = 2; /* number of bits in bypass mode */
+constexpr uint16_t max_bypass_val = (1 << bypass_precision) - 1;
+
+inline void RansEncPutBits(RansState& r, uint8_t*& ptr, uint32_t val)
+{
+    RansAssert(bypass_precision <= 8);
+    RansAssert(val < (1u << bypass_precision));
+
+    constexpr uint32_t freq = 1 << (SCALE_BITS - bypass_precision);
+    constexpr uint32_t x_max = freq << ENC_RENORM_SHIFT_BITS;
+    while (r >= x_max) {
+        *(--ptr) = static_cast<uint8_t>(r & 0xff);
+        r >>= 8;
+    }
+
+    r = (r << bypass_precision) | val;
+}
+
+inline uint32_t RansDecGetBits(RansState& r, uint8_t*& ptr)
+{
+    uint32_t val = r & ((1u << bypass_precision) - 1);
+
+    /* Re-normalize */
+    r = r >> bypass_precision;
+    if (r < RANS_BYTE_L) {
+        r = (r << 8) | *ptr++;
+        RansAssert(r >= RANS_BYTE_L);
+    }
+
+    return val;
+}
+
+RansEncoderLib::RansEncoderLib()
+{
+    _stream = std::make_shared<std::vector<uint8_t>>();
+}
+
+int RansEncoderLib::add_cdf(const std::shared_ptr<std::vector<std::vector<int32_t>>> cdfs,
+                            const std::shared_ptr<std::vector<int32_t>> cdfs_sizes,
+                            const std::shared_ptr<std::vector<int32_t>> offsets)
+{
+
+    auto ransSymbols = std::make_shared<std::vector<std::vector<RansSymbol>>>(cdfs->size());
+    for (int i = 0; i < static_cast<int>(cdfs->size()); i++) {
+        const int32_t* cdf = cdfs->at(i).data();
+        std::vector<RansSymbol> ransSym(cdfs->at(i).size());
+        const int ransSize = static_cast<int>(ransSym.size() - 1);
+        for (int j = 0; j < ransSize; j++) {
+            ransSym[j] = RansSymbol(
+                { static_cast<uint16_t>(cdf[j]), static_cast<uint16_t>(cdf[j + 1] - cdf[j]) });
+        }
+        ransSymbols->at(i) = ransSym;
+    }
+
+    _ransSymbols.push_back(ransSymbols);
+    _cdfs_sizes.push_back(cdfs_sizes);
+    _offsets.push_back(offsets);
+    return static_cast<int>(_ransSymbols.size()) - 1;
+}
+
+void RansEncoderLib::empty_cdf_buffer()
+{
+    _ransSymbols.clear();
+    _cdfs_sizes.clear();
+    _offsets.clear();
+}
+
+FORCE_INLINE void RansEncoderLib::encode_one_symbol(uint8_t*& ptr, RansState& rans, const int32_t symbol,
+                                                    const int32_t cdf_size, const int32_t offset,
+                                                    const std::vector<RansSymbol>& ransSymbols)
+{
+    const int32_t max_value = cdf_size - 2;
+    int32_t value = symbol - offset;
+
+    uint32_t raw_val = 0;
+    if (value < 0) {
+        raw_val = -2 * value - 1;
+        value = max_value;
+    } else if (value >= max_value) {
+        raw_val = 2 * (value - max_value);
+        value = max_value;
+    }
+
+    if (value == max_value) {
+        std::vector<uint16_t> bypassBins;
+        bypassBins.reserve(20);
+        /* Determine the number of bypasses (in bypass_precision size) needed to
+         * encode the raw value. */
+        int32_t n_bypass = 0;
+        while ((raw_val >> (n_bypass * bypass_precision)) != 0) {
+            ++n_bypass;
+        }
+
+        /* Encode number of bypasses */
+        int32_t val = n_bypass;
+        while (val >= max_bypass_val) {
+            bypassBins.push_back(max_bypass_val);
+            val -= max_bypass_val;
+        }
+        bypassBins.push_back(static_cast<uint16_t>(val));
+
+        /* Encode raw value */
+        for (int32_t j = 0; j < n_bypass; ++j) {
+            const int32_t val1 = (raw_val >> (j * bypass_precision)) & max_bypass_val;
+            bypassBins.push_back(static_cast<uint16_t>(val1));
+        }
+
+        for (auto it = bypassBins.rbegin(); it < bypassBins.rend(); it++) {
+            RansEncPutBits(rans, ptr, *it);
+        }
+    }
+    RansEncPut(rans, ptr, ransSymbols[value].start, ransSymbols[value].range);
+}
+
+void RansEncoderLib::encode_y(const std::shared_ptr<std::vector<int16_t>> symbols,
+                              const int cdf_group_index)
+{
+    PendingTask p;
+    p.workType = WorkType::EncodeDecodeY;
+    p.symbols_y = symbols;
+    p.cdf_group_index = cdf_group_index;
+    m_pendingEncodingList.push_back(p);
+}
+
+void RansEncoderLib::encode_z(const std::shared_ptr<std::vector<int8_t>> symbols,
+                              const int cdf_group_index, const int start_offset,
+                              const int per_channel_size)
+{
+    PendingTask p;
+    p.workType = WorkType::EncodeDecodeZ;
+    p.symbols_z = symbols;
+    p.cdf_group_index = cdf_group_index;
+    p.start_offset = start_offset;
+    p.per_channel_size = per_channel_size;
+    m_pendingEncodingList.push_back(p);
+}
+#include <iostream>
+FORCE_INLINE void RansEncoderLib::encode_y_internal(uint8_t*& ptr, RansState& rans,
+                                                    const std::shared_ptr<std::vector<int16_t>> symbols,
+                                                    const int cdf_group_index)
+{
+    // backward loop on symbols from the end;
+    const int16_t* symbols_ptr = symbols->data();
+    const int32_t* cdfs_sizes_ptr = _cdfs_sizes[cdf_group_index]->data();
+    const int32_t* offsets_ptr = _offsets[cdf_group_index]->data();
+    const int symbol_size = static_cast<int>(symbols->size());
+
+    for (int i = symbol_size - 1; i >= 0; i--) {
+        const int32_t combined_symbol = symbols_ptr[i];
+        const int32_t cdf_idx = combined_symbol & 0xff;
+        const int32_t s = combined_symbol >> 8;
+        encode_one_symbol(ptr, rans, s, cdfs_sizes_ptr[cdf_idx], offsets_ptr[cdf_idx],
+                          _ransSymbols[cdf_group_index]->at(cdf_idx));
+    }
+}
+
+FORCE_INLINE void RansEncoderLib::encode_z_internal(uint8_t*& ptr, RansState& rans,
+                                                    const std::shared_ptr<std::vector<int8_t>> symbols,
+                                                    const int cdf_group_index, const int start_offset,
+                                                    const int per_channel_size)
+{
+    // backward loop on symbols from the end;
+    const int8_t* symbols_ptr = symbols->data();
+    const int32_t* cdfs_sizes_ptr = _cdfs_sizes[cdf_group_index]->data();
+    const int32_t* offsets_ptr = _offsets[cdf_group_index]->data();
+    const int symbol_size = static_cast<int>(symbols->size());
+
+    for (int i = symbol_size - 1; i >= 0; i--) {
+        const int32_t cdf_idx = i / per_channel_size + start_offset;
+        encode_one_symbol(ptr, rans, symbols_ptr[i], cdfs_sizes_ptr[cdf_idx], offsets_ptr[cdf_idx],
+                          _ransSymbols[cdf_group_index]->at(cdf_idx));
+    }
+}
+
+void RansEncoderLib::flush()
+{
+    RansState rans;
+    RansEncInit(rans);
+
+    int32_t total_symbol_size = 0;
+    for (auto it = m_pendingEncodingList.begin(); it != m_pendingEncodingList.end(); it++) {
+        if (it->workType == WorkType::EncodeDecodeY) {
+            total_symbol_size += static_cast<int32_t>(it->symbols_y->size());
+        } else if (it->workType == WorkType::EncodeDecodeZ) {
+            total_symbol_size += static_cast<int32_t>(it->symbols_z->size());
+        }
+    }
+
+    if (total_symbol_size == 0) {
+        _stream->resize(0);
+        return;
+    }
+
+    uint8_t* output = new uint8_t[total_symbol_size];  // too much space ?
+    uint8_t* ptrEnd = output + total_symbol_size;
+    uint8_t* ptr = ptrEnd;
+    assert(ptr != nullptr);
+
+    for (auto it = m_pendingEncodingList.rbegin(); it != m_pendingEncodingList.rend(); it++) {
+        PendingTask p = *it;
+        if (p.workType == WorkType::EncodeDecodeY) {
+            encode_y_internal(ptr, rans, p.symbols_y, p.cdf_group_index);
+        } else if (p.workType == WorkType::EncodeDecodeZ) {
+            encode_z_internal(ptr, rans, p.symbols_z, p.cdf_group_index, p.start_offset,
+                              p.per_channel_size);
+        }
+    }
+
+    RansEncFlush(rans, ptr);
+
+    const int nbytes = static_cast<int>(std::distance(ptr, ptrEnd));
+
+    _stream->resize(nbytes);
+    memcpy(_stream->data(), ptr, nbytes);
+    delete[] output;
+}
+
+std::shared_ptr<std::vector<uint8_t>> RansEncoderLib::get_encoded_stream()
+{
+    return _stream;
+}
+
+void RansEncoderLib::reset()
+{
+    m_pendingEncodingList.clear();
+    _stream->clear();
+}
+
+RansEncoderLibMultiThread::RansEncoderLibMultiThread()
+    : RansEncoderLib()
+    , m_finish(false)
+    , m_result_ready(false)
+{
+    m_thread = std::thread(&RansEncoderLibMultiThread::worker, this);
+}
+RansEncoderLibMultiThread::~RansEncoderLibMultiThread()
+{
+    {
+        std::lock_guard<std::mutex> lk(m_mutex_pending);
+        std::lock_guard<std::mutex> lk1(m_mutex_result);
+        m_finish = true;
+    }
+    m_cv_pending.notify_one();
+    m_cv_result.notify_one();
+    m_thread.join();
+}
+
+void RansEncoderLibMultiThread::flush()
+{
+    PendingTask p;
+    p.workType = WorkType::Flush;
+    {
+        std::unique_lock<std::mutex> lk(m_mutex_pending);
+        m_pending.push_back(p);
+    }
+    m_cv_pending.notify_one();
+}
+
+std::shared_ptr<std::vector<uint8_t>> RansEncoderLibMultiThread::get_encoded_stream()
+{
+    std::unique_lock<std::mutex> lk(m_mutex_result);
+    m_cv_result.wait(lk, [this] { return m_result_ready || m_finish; });
+    return RansEncoderLib::get_encoded_stream();
+}
+
+void RansEncoderLibMultiThread::reset()
+{
+    RansEncoderLib::reset();
+    std::lock_guard<std::mutex> lk(m_mutex_result);
+    m_result_ready = false;
+}
+
+void RansEncoderLibMultiThread::worker()
+{
+    while (!m_finish) {
+        std::unique_lock<std::mutex> lk(m_mutex_pending);
+        m_cv_pending.wait(lk, [this] { return m_pending.size() > 0 || m_finish; });
+        if (m_finish) {
+            lk.unlock();
+            break;
+        }
+        if (m_pending.size() == 0) {
+            lk.unlock();
+            // std::cout << "contine in worker" << std::endl;
+            continue;
+        }
+        while (m_pending.size() > 0) {
+            auto p = m_pending.front();
+            m_pending.pop_front();
+            lk.unlock();
+            if (p.workType == WorkType::Flush) {
+                RansEncoderLib::flush();
+                {
+                    std::lock_guard<std::mutex> lk_result(m_mutex_result);
+                    m_result_ready = true;
+                }
+                m_cv_result.notify_one();
+            }
+            lk.lock();
+        }
+        lk.unlock();
+    }
+}
+
+void RansDecoderLib::set_stream(const std::shared_ptr<std::vector<uint8_t>> encoded)
+{
+    _stream = encoded;
+    _ptr8 = (uint8_t*)(_stream->data());
+    RansDecInit(_rans, _ptr8);
+}
+
+int RansDecoderLib::add_cdf(const std::shared_ptr<std::vector<std::vector<int32_t>>> cdfs,
+                            const std::shared_ptr<std::vector<int32_t>> cdfs_sizes,
+                            const std::shared_ptr<std::vector<int32_t>> offsets)
+{
+    _cdfs.push_back(cdfs);
+    _cdfs_sizes.push_back(cdfs_sizes);
+    _offsets.push_back(offsets);
+    return static_cast<int>(_cdfs.size()) - 1;
+}
+
+void RansDecoderLib::empty_cdf_buffer()
+{
+    _cdfs.clear();
+    _cdfs_sizes.clear();
+    _offsets.clear();
+}
+
+FORCE_INLINE int8_t RansDecoderLib::decode_one_symbol(const int32_t* cdf, const int32_t cdf_size,
+                                                      const int32_t offset)
+{
+    const int32_t max_value = cdf_size - 2;
+    const int32_t cum_freq = static_cast<int32_t>(RansDecGet(_rans));
+
+    int s = 1;
+    while (cdf[s++] <= cum_freq) {
+    }
+    s -= 2;
+
+    RansDecAdvance(_rans, _ptr8, cdf[s], cdf[s + 1] - cdf[s]);
+
+    int32_t value = static_cast<int32_t>(s);
+
+    if (value == max_value) {
+        /* Bypass decoding mode */
+        int32_t val = RansDecGetBits(_rans, _ptr8);
+        int32_t n_bypass = val;
+
+        while (val == max_bypass_val) {
+            val = RansDecGetBits(_rans, _ptr8);
+            n_bypass += val;
+        }
+
+        int32_t raw_val = 0;
+        for (int j = 0; j < n_bypass; ++j) {
+            val = RansDecGetBits(_rans, _ptr8);
+            raw_val |= val << (j * bypass_precision);
+        }
+        value = raw_val >> 1;
+        if (raw_val & 1) {
+            value = -value - 1;
+        } else {
+            value += max_value;
+        }
+    }
+
+    return static_cast<int8_t>(value + offset);
+}
+
+void RansDecoderLib::decode_y(const std::shared_ptr<std::vector<uint8_t>> indexes,
+                              const int cdf_group_index)
+{
+    int index_size = static_cast<int>(indexes->size());
+    m_decoded = std::make_shared<std::vector<int8_t>>(index_size);
+
+    int8_t* outout_ptr = m_decoded->data();
+    const uint8_t* indexes_ptr = indexes->data();
+    const int32_t* cdfs_sizes_ptr = _cdfs_sizes[cdf_group_index]->data();
+    const int32_t* offsets_ptr = _offsets[cdf_group_index]->data();
+    const auto& cdfs = _cdfs[cdf_group_index];
+    for (int i = 0; i < index_size; ++i) {
+        const int32_t cdf_idx = indexes_ptr[i];
+        outout_ptr[i] = decode_one_symbol(cdfs->at(cdf_idx).data(), cdfs_sizes_ptr[cdf_idx],
+                                          offsets_ptr[cdf_idx]);
+    }
+}
+
+void RansDecoderLib::decode_z(const int total_size, const int cdf_group_index,
+                              const int start_offset, const int per_channel_size)
+{
+    m_decoded = std::make_shared<std::vector<int8_t>>(total_size);
+
+    int8_t* outout_ptr = m_decoded->data();
+    const int32_t* cdfs_sizes_ptr = _cdfs_sizes[cdf_group_index]->data();
+    const int32_t* offsets_ptr = _offsets[cdf_group_index]->data();
+    const auto& cdfs = _cdfs[cdf_group_index];
+    for (int i = 0; i < total_size; ++i) {
+        const int32_t cdf_idx = i / per_channel_size + start_offset;
+        outout_ptr[i] = decode_one_symbol(cdfs->at(cdf_idx).data(), cdfs_sizes_ptr[cdf_idx],
+                                          offsets_ptr[cdf_idx]);
+    }
+}
+
+std::shared_ptr<std::vector<int8_t>> RansDecoderLib::get_decoded_tensor()
+{
+    return m_decoded;
+}
+
+RansDecoderLibMultiThread::RansDecoderLibMultiThread()
+    : RansDecoderLib()
+    , m_finish(false)
+    , m_result_ready(false)
+{
+    m_thread = std::thread(&RansDecoderLibMultiThread::worker, this);
+}
+
+RansDecoderLibMultiThread::~RansDecoderLibMultiThread()
+{
+    {
+        std::lock_guard<std::mutex> lk(m_mutex_pending);
+        std::lock_guard<std::mutex> lk1(m_mutex_result);
+        m_finish = true;
+    }
+    m_cv_pending.notify_one();
+    m_cv_result.notify_one();
+    m_thread.join();
+}
+
+void RansDecoderLibMultiThread::decode_y(const std::shared_ptr<std::vector<uint8_t>> indexes,
+                                         const int cdf_group_index)
+{
+    {
+        std::lock_guard<std::mutex> lk(m_mutex_result);
+        m_result_ready = false;
+    }
+    PendingTask p;
+    p.workType = WorkType::EncodeDecodeY;
+    p.indexes = indexes;
+    p.cdf_group_index = cdf_group_index;
+    {
+        std::unique_lock<std::mutex> lk(m_mutex_pending);
+        m_pending.push_back(p);
+    }
+    m_cv_pending.notify_one();
+}
+
+void RansDecoderLibMultiThread::decode_z(const int total_size, const int cdf_group_index,
+                                         const int start_offset, const int per_channel_size)
+{
+    {
+        std::lock_guard<std::mutex> lk(m_mutex_result);
+        m_result_ready = false;
+    }
+    PendingTask p;
+    p.workType = WorkType::EncodeDecodeZ;
+    p.total_size = total_size;
+    p.cdf_group_index = cdf_group_index;
+    p.start_offset = start_offset;
+    p.per_channel_size = per_channel_size;
+    {
+        std::unique_lock<std::mutex> lk(m_mutex_pending);
+        m_pending.push_back(p);
+    }
+    m_cv_pending.notify_one();
+}
+
+std::shared_ptr<std::vector<int8_t>> RansDecoderLibMultiThread::get_decoded_tensor()
+{
+    std::unique_lock<std::mutex> lk(m_mutex_result);
+    m_cv_result.wait(lk, [this] { return m_result_ready || m_finish; });
+    return RansDecoderLib::get_decoded_tensor();
+}
+
+void RansDecoderLibMultiThread::worker()
+{
+    while (!m_finish) {
+        std::unique_lock<std::mutex> lk(m_mutex_pending);
+        m_cv_pending.wait(lk, [this] { return m_pending.size() > 0 || m_finish; });
+        if (m_finish) {
+            lk.unlock();
+            break;
+        }
+        if (m_pending.size() == 0) {
+            lk.unlock();
+            // std::cout << "contine in worker" << std::endl;
+            continue;
+        }
+        while (m_pending.size() > 0) {
+            auto p = m_pending.front();
+            m_pending.pop_front();
+            lk.unlock();
+            if (p.workType == WorkType::EncodeDecodeY) {
+                RansDecoderLib::decode_y(p.indexes, p.cdf_group_index);
+            } else if (p.workType == WorkType::EncodeDecodeZ) {
+                RansDecoderLib::decode_z(p.total_size, p.cdf_group_index, p.start_offset,
+                                         p.per_channel_size);
+            }
+            {
+                std::lock_guard<std::mutex> lk_result(m_mutex_result);
+                m_result_ready = true;
+            }
+            m_cv_result.notify_one();
+            lk.lock();
+        }
+        lk.unlock();
+    }
+}

neural_codec/DCVC/src/cpp/py_rans/rans.h

@@ -0,0 +1,201 @@
+/* Copyright 2020 InterDigital Communications, Inc.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#pragma once
+
+#include <condition_variable>
+#include <list>
+#include <thread>
+#include <vector>
+
+#ifdef __GNUC__
+    #pragma GCC diagnostic push
+    #pragma GCC diagnostic ignored "-Wpedantic"
+    #pragma GCC diagnostic ignored "-Wsign-compare"
+#endif
+
+#ifdef _MSC_VER
+    #pragma warning(disable : 4244)
+#endif
+
+#include "rans_byte.h"
+
+#ifdef _MSC_VER
+    #pragma warning(default : 4244)
+#endif
+
+#ifdef __GNUC__
+    #pragma GCC diagnostic pop
+#endif
+
+#ifdef _MSC_VER
+    #define FORCE_INLINE __forceinline
+#endif
+
+#ifdef __GNUC__
+    #define FORCE_INLINE __attribute__((always_inline)) inline
+#endif
+
+struct RansSymbol {
+    uint16_t start;
+    uint16_t range;  // range for normal coding and 0 for bypass coding
+};
+
+enum class WorkType {
+    EncodeDecodeY,
+    EncodeDecodeZ,
+    Flush,
+};
+
+struct PendingTask {
+    WorkType workType;
+    std::shared_ptr<std::vector<int16_t>> symbols_y;
+    std::shared_ptr<std::vector<int8_t>> symbols_z;
+    std::shared_ptr<std::vector<uint8_t>> indexes;
+    int total_size{ 0 };
+    int cdf_group_index{ 0 };
+    int start_offset{ 0 };
+    int per_channel_size{ 0 };
+};
+
+/* NOTE: Warning, we buffer everything for now... In case of large files we
+ * should split the bitstream into chunks... Or for a memory-bounded encoder
+ **/
+class RansEncoderLib {
+public:
+    RansEncoderLib();
+    virtual ~RansEncoderLib() = default;
+
+    RansEncoderLib(const RansEncoderLib&) = delete;
+    RansEncoderLib(RansEncoderLib&&) = delete;
+    RansEncoderLib& operator=(const RansEncoderLib&) = delete;
+    RansEncoderLib& operator=(RansEncoderLib&&) = delete;
+
+    void encode_y(const std::shared_ptr<std::vector<int16_t>> symbols, const int cdf_group_index);
+    void encode_z(const std::shared_ptr<std::vector<int8_t>> symbols, const int cdf_group_index,
+                  const int start_offset, const int per_channel_size);
+
+    FORCE_INLINE void encode_y_internal(uint8_t*& ptr, RansState& rans,
+                                        const std::shared_ptr<std::vector<int16_t>> symbols,
+                                        const int cdf_group_index);
+    FORCE_INLINE void encode_z_internal(uint8_t*& ptr, RansState& rans,
+                                        const std::shared_ptr<std::vector<int8_t>> symbols,
+                                        const int cdf_group_index, const int start_offset,
+                                        const int per_channel_size);
+    FORCE_INLINE void encode_one_symbol(uint8_t*& ptr, RansState& rans, const int32_t symbol,
+                                        const int32_t cdf_size, const int32_t offset,
+                                        const std::vector<RansSymbol>& ransSymbols);
+    virtual void flush();
+    virtual std::shared_ptr<std::vector<uint8_t>> get_encoded_stream();
+    virtual void reset();
+    virtual int add_cdf(const std::shared_ptr<std::vector<std::vector<int32_t>>> cdfs,
+                        const std::shared_ptr<std::vector<int32_t>> cdfs_sizes,
+                        const std::shared_ptr<std::vector<int32_t>> offsets);
+    virtual void empty_cdf_buffer();
+
+private:
+    std::shared_ptr<std::vector<uint8_t>> _stream;
+
+    std::vector<std::shared_ptr<std::vector<std::vector<RansSymbol>>>> _ransSymbols;
+    std::vector<std::shared_ptr<std::vector<int32_t>>> _cdfs_sizes;
+    std::vector<std::shared_ptr<std::vector<int32_t>>> _offsets;
+
+    std::list<PendingTask> m_pendingEncodingList;
+};
+
+class RansEncoderLibMultiThread : public RansEncoderLib {
+public:
+    RansEncoderLibMultiThread();
+    virtual ~RansEncoderLibMultiThread();
+    virtual void flush() override;
+    virtual std::shared_ptr<std::vector<uint8_t>> get_encoded_stream() override;
+    virtual void reset() override;
+
+    void worker();
+
+private:
+    bool m_finish;
+    bool m_result_ready;
+    std::thread m_thread;
+    std::mutex m_mutex_result;
+    std::mutex m_mutex_pending;
+    std::condition_variable m_cv_pending;
+    std::condition_variable m_cv_result;
+    std::list<PendingTask> m_pending;
+};
+
+class RansDecoderLib {
+public:
+    RansDecoderLib() {}
+    virtual ~RansDecoderLib() = default;
+
+    RansDecoderLib(const RansDecoderLib&) = delete;
+    RansDecoderLib(RansDecoderLib&&) = delete;
+    RansDecoderLib& operator=(const RansDecoderLib&) = delete;
+    RansDecoderLib& operator=(RansDecoderLib&&) = delete;
+
+    virtual void set_stream(const std::shared_ptr<std::vector<uint8_t>> encoded);
+
+    FORCE_INLINE int8_t decode_one_symbol(const int32_t* cdf, const int32_t cdf_size,
+                                          const int32_t offset);
+
+    virtual void decode_y(const std::shared_ptr<std::vector<uint8_t>> indexes,
+                          const int cdf_group_index);
+    virtual void decode_z(const int total_size, const int cdf_group_index, const int start_offset,
+                          const int per_channel_size);
+
+    virtual std::shared_ptr<std::vector<int8_t>> get_decoded_tensor();
+
+    virtual int add_cdf(const std::shared_ptr<std::vector<std::vector<int32_t>>> cdfs,
+                        const std::shared_ptr<std::vector<int32_t>> cdfs_sizes,
+                        const std::shared_ptr<std::vector<int32_t>> offsets);
+    virtual void empty_cdf_buffer();
+
+private:
+    RansState _rans;
+    uint8_t* _ptr8;
+    std::shared_ptr<std::vector<uint8_t>> _stream;
+    std::shared_ptr<std::vector<int8_t>> m_decoded;
+
+    std::vector<std::shared_ptr<std::vector<std::vector<int32_t>>>> _cdfs;
+    std::vector<std::shared_ptr<std::vector<int32_t>>> _cdfs_sizes;
+    std::vector<std::shared_ptr<std::vector<int32_t>>> _offsets;
+};
+
+class RansDecoderLibMultiThread : public RansDecoderLib {
+public:
+    RansDecoderLibMultiThread();
+    virtual ~RansDecoderLibMultiThread();
+
+    virtual void decode_y(const std::shared_ptr<std::vector<uint8_t>> indexes,
+                          const int cdf_group_index) override;
+
+    virtual void decode_z(const int total_size, const int cdf_group_index, const int start_offset,
+                          const int per_channel_size) override;
+
+    virtual std::shared_ptr<std::vector<int8_t>> get_decoded_tensor() override;
+
+    void worker();
+
+private:
+    bool m_finish;
+    bool m_result_ready;
+    std::thread m_thread;
+    std::mutex m_mutex_result;
+    std::mutex m_mutex_pending;
+    std::condition_variable m_cv_pending;
+    std::condition_variable m_cv_result;
+    std::list<PendingTask> m_pending;
+};

neural_codec/DCVC/src/cpp/py_rans/rans_byte.h

@@ -0,0 +1,141 @@
+// The code is from https://github.com/rygorous/ryg_rans
+// The original lisence is below.
+
+// To the extent possible under law, Fabian Giesen has waived all
+// copyright and related or neighboring rights to ryg_rans, as
+// per the terms of the CC0 license:
+
+//   https://creativecommons.org/publicdomain/zero/1.0
+
+// This work is published from the United States.
+
+// Simple byte-aligned rANS encoder/decoder - public domain - Fabian 'ryg'
+// Giesen 2014
+//
+// Not intended to be "industrial strength"; just meant to illustrate the
+// general idea.
+
+#pragma once
+
+#include <stdint.h>
+
+#ifdef assert
+    #define RansAssert assert
+#else
+    #define RansAssert(x)
+#endif
+
+// READ ME FIRST:
+//
+// This is designed like a typical arithmetic coder API, but there's three
+// twists you absolutely should be aware of before you start hacking:
+//
+// 1. You need to encode data in *reverse* - last symbol first. rANS works
+//    like a stack: last in, first out.
+// 2. Likewise, the encoder outputs bytes *in reverse* - that is, you give
+//    it a pointer to the *end* of your buffer (exclusive), and it will
+//    slowly move towards the beginning as more bytes are emitted.
+// 3. Unlike basically any other entropy coder implementation you might
+//    have used, you can interleave data from multiple independent rANS
+//    encoders into the same bytestream without any extra signaling;
+//    you can also just write some bytes by yourself in the middle if
+//    you want to. This is in addition to the usual arithmetic encoder
+//    property of being able to switch models on the fly. Writing raw
+//    bytes can be useful when you have some data that you know is
+//    incompressible, and is cheaper than going through the rANS encode
+//    function. Using multiple rANS coders on the same byte stream wastes
+//    a few bytes compared to using just one, but execution of two
+//    independent encoders can happen in parallel on superscalar and
+//    Out-of-Order CPUs, so this can be *much* faster in tight decoding
+//    loops.
+//
+//    This is why all the rANS functions take the write pointer as an
+//    argument instead of just storing it in some context struct.
+
+// --------------------------------------------------------------------------
+
+// L ('l' in the paper) is the lower bound of our normalization interval.
+// Between this and our byte-aligned emission, we use 31 (not 32!) bits.
+// This is done intentionally because exact reciprocals for 31-bit uints
+// fit in 32-bit uints: this permits some optimizations during encoding.
+constexpr int SCALE_BITS = 16;
+constexpr int RANS_SHIFT_BITS = 23;
+constexpr uint32_t RANS_BYTE_L = (1u << RANS_SHIFT_BITS);
+constexpr int ENC_RENORM_SHIFT_BITS = RANS_SHIFT_BITS - SCALE_BITS + 8;
+constexpr int DEC_MASK = ((1u << SCALE_BITS) - 1);
+
+// State for a rANS encoder. Yep, that's all there is to it.
+typedef uint32_t RansState;
+
+// Initialize a rANS encoder.
+static inline void RansEncInit(RansState& r)
+{
+    r = RANS_BYTE_L;
+}
+
+// Renormalize the encoder. Internal function.
+static inline void RansEncRenorm(RansState& x, uint8_t*& ptr, uint32_t freq)
+{
+    const uint32_t x_max = freq << ENC_RENORM_SHIFT_BITS;
+    while (x >= x_max) {
+        *(--ptr) = static_cast<uint8_t>(x & 0xff);
+        x >>= 8;
+    }
+}
+
+// Encodes a single symbol with range start "start" and frequency "freq".
+// All frequencies are assumed to sum to "1 << scale_bits", and the
+// resulting bytes get written to ptr (which is updated).
+//
+// NOTE: With rANS, you need to encode symbols in *reverse order*, i.e. from
+// beginning to end! Likewise, the output bytestream is written *backwards*:
+// ptr starts pointing at the end of the output buffer and keeps decrementing.
+static inline void RansEncPut(RansState& r, uint8_t*& ptr, uint32_t start, uint32_t freq)
+{
+    // renormalize
+    RansEncRenorm(r, ptr, freq);
+
+    // x = C(s,x)
+    r = ((r / freq) << SCALE_BITS) + (r % freq) + start;
+}
+
+// Flushes the rANS encoder.
+static inline void RansEncFlush(const RansState& r, uint8_t*& ptr)
+{
+    ptr -= 4;
+    ptr[0] = (uint8_t)(r >> 0);
+    ptr[1] = (uint8_t)(r >> 8);
+    ptr[2] = (uint8_t)(r >> 16);
+    ptr[3] = (uint8_t)(r >> 24);
+}
+
+// Initializes a rANS decoder.
+// Unlike the encoder, the decoder works forwards as you'd expect.
+static inline void RansDecInit(RansState& r, uint8_t*& ptr)
+{
+    r = (*ptr++) << 0;
+    r |= (*ptr++) << 8;
+    r |= (*ptr++) << 16;
+    r |= (*ptr++) << 24;
+}
+
+// Returns the current cumulative frequency (map it to a symbol yourself!)
+static inline uint32_t RansDecGet(RansState& r)
+{
+    return r & DEC_MASK;
+}
+
+// Advances in the bit stream by "popping" a single symbol with range start
+// "start" and frequency "freq". All frequencies are assumed to sum to "1 <<
+// scale_bits", and the resulting bytes get written to ptr (which is updated).
+static inline void RansDecAdvance(RansState& r, uint8_t*& ptr, uint32_t start, uint32_t freq)
+{
+
+    // s, x = D(x)
+    r = freq * (r >> SCALE_BITS) + (r & DEC_MASK) - start;
+
+    // renormalize
+    while (r < RANS_BYTE_L) {
+        r = (r << 8) | *ptr++;
+    }
+}

neural_codec/DCVC/src/cpp/setup.py

@@ -0,0 +1,31 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+
+import glob
+import sys
+from setuptools import setup
+from pybind11.setup_helpers import Pybind11Extension, build_ext
+
+
+if sys.platform == "win32":
+    extra_compile_args = ['/std:c++17', '/O2', '/W4', '/WX', '/wd4100']
+    extra_link_args = []
+else:
+    extra_compile_args = ['-std=c++17', '-O3', '-fPIC', '-Wall', '-Wextra', '-Werror']
+    extra_link_args = []
+
+
+setup(
+    name="MLCodec_extensions_cpp",
+    ext_modules=[
+        Pybind11Extension(
+            name='MLCodec_extensions_cpp',
+            sources=glob.glob('py_rans/*.cpp'),
+            extra_compile_args=extra_compile_args,
+            extra_link_args=extra_link_args,
+        ),
+    ],
+    cmdclass={"build_ext": build_ext},
+    zip_safe=False,
+    python_requires=">=3.12",
+)

neural_codec/DCVC/src/layers/cuda_inference.py

@@ -0,0 +1,208 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+
+import os
+
+import torch
+import torch.nn.functional as F
+
+
+CUSTOMIZED_CUDA_INFERENCE = False
+# Escape hatch: the native inference kernel makes vec4-aligned memory accesses
+# that can crash (illegal/misaligned address) on some frame geometries under
+# full-resolution scoring. Set DCVC_FORCE_PYTORCH=1 to force the pure-pytorch
+# fallback (slower but always correct) for such inputs.
+if os.environ.get("DCVC_FORCE_PYTORCH", "") not in ("1", "true", "True"):
+    try:
+        from inference_extensions_cuda import process_with_mask_cuda, combine_for_reading_2x_cuda, \
+            restore_y_2x_cuda, restore_y_4x_cuda, build_index_dec_cuda, \
+            round_and_to_int8_cuda, clamp_reciprocal_with_quant_cuda, bias_quant_cuda, \
+            add_and_multiply_cuda, bias_pixel_shuffle_8_cuda, replicate_pad_cuda, \
+            build_index_enc_cuda, DepthConvProxy, SubpelConv2xProxy  # noqa: F401
+        CUSTOMIZED_CUDA_INFERENCE = True
+    except Exception:  # pylint: disable=W0718
+        pass
+
+
+if not CUSTOMIZED_CUDA_INFERENCE and 'SUPPRESS_CUSTOM_KERNEL_WARNING' not in os.environ:
+    print("cannot import cuda implementation for inference, fallback to pytorch.")
+
+
+def round_and_to_int8(z):
+    if CUSTOMIZED_CUDA_INFERENCE and z.is_cuda:
+        z_int8 = round_and_to_int8_cuda(z)
+        return z, z_int8
+
+    z_hat = torch.clamp(torch.round(z), -128., 127.)
+    z_hat_write = z_hat.to(dtype=torch.int8)
+    return z_hat, z_hat_write
+
+
+def clamp_reciprocal_with_quant(q_dec, y, min_val):
+    if CUSTOMIZED_CUDA_INFERENCE and q_dec.is_cuda:
+        # q_dec is not inplace modified at decoder side
+        q_dec = clamp_reciprocal_with_quant_cuda(q_dec, y, min_val)
+        return q_dec, y
+
+    q_dec = torch.clamp_min(q_dec, min_val)
+    q_enc = torch.reciprocal(q_dec)
+    y = y * q_enc
+    return q_dec, y
+
+
+def add_and_multiply(y_hat_0, y_hat_1, q_dec):
+    if CUSTOMIZED_CUDA_INFERENCE and y_hat_0.is_cuda:
+        add_and_multiply_cuda(y_hat_0, y_hat_1, q_dec)
+        return y_hat_0
+
+    y_hat = y_hat_0 + y_hat_1
+    y_hat = y_hat * q_dec
+    return y_hat
+
+
+def process_with_mask(y, scales, means, mask, force_zero_thres):
+    if CUSTOMIZED_CUDA_INFERENCE and y.is_cuda:
+        thres = force_zero_thres if force_zero_thres is not None else -1.
+        return process_with_mask_cuda(y, scales, means, mask, thres)
+
+    scales_hat = scales * mask
+    means_hat = means * mask
+
+    y_res = (y - means_hat) * mask
+    y_q = torch.round(y_res)
+    if force_zero_thres is not None:
+        cond = scales_hat > force_zero_thres
+        y_q = y_q * cond
+    y_q = torch.clamp(y_q, -128., 127.)
+    y_hat = y_q + means_hat
+
+    return y_res, y_q, y_hat, scales_hat
+
+
+def combine_for_reading_2x(x, mask, inplace=False):
+    if CUSTOMIZED_CUDA_INFERENCE and x.is_cuda and x.is_contiguous():
+        B, C, H, W = x.shape
+        if inplace:
+            out = x[:, :C // 2, :, :]
+        else:
+            out = torch.empty((B, C // 2, H, W), dtype=x.dtype, layout=x.layout, device=x.device)
+        combine_for_reading_2x_cuda(out, x, mask)
+        return out
+
+    x = x * mask
+    x0, x1 = x.chunk(2, 1)
+    return x0 + x1
+
+
+def restore_y_2x(y, means, mask):
+    if CUSTOMIZED_CUDA_INFERENCE and y.is_cuda and y.is_contiguous():
+        out = torch.empty_like(means)
+        restore_y_2x_cuda(out, y, means, mask)
+        return out
+
+    return (torch.cat((y, y), dim=1) + means) * mask
+
+
+def restore_y_2x_with_cat_after(y, means, mask, to_cat):
+    if CUSTOMIZED_CUDA_INFERENCE and y.is_cuda and y.is_contiguous():
+        B, C1, H, W = means.shape
+        C2 = to_cat.shape[1]
+        out = torch.empty((B, C1 + C2, H, W), dtype=means.dtype, layout=means.layout,
+                          device=means.device)
+        restore_y_2x_cuda(out[:, :C1, :, :], y, means, mask)
+        out[:, C1:, :, :] = to_cat
+        return out[:, :C1, :, :], out
+
+    out = (torch.cat((y, y), dim=1) + means) * mask
+    return out, torch.cat((out, to_cat), dim=1)
+
+
+def restore_y_4x(y, means, mask):
+    if CUSTOMIZED_CUDA_INFERENCE and y.is_cuda and y.is_contiguous():
+        out = torch.empty_like(means)
+        restore_y_4x_cuda(out, y, means, mask)
+        return out
+
+    return (torch.cat((y, y, y, y), dim=1) + means) * mask
+
+
+def build_index_dec(scales, scale_min, scale_max, log_scale_min, log_step_recip, skip_thres=None):
+    if CUSTOMIZED_CUDA_INFERENCE and scales.is_cuda:
+        out = torch.empty_like(scales, dtype=torch.uint8)
+        skip_cond = None
+        if skip_thres is not None:
+            skip_cond = torch.empty_like(scales, dtype=torch.bool)
+        else:
+            skip_thres = -1.
+
+        build_index_dec_cuda(out, skip_cond, scales, scale_min, scale_max, log_scale_min,
+                             log_step_recip, skip_thres)
+        return out, skip_cond
+
+    skip_cond = None
+    scales = scales.clamp_(scale_min, scale_max)
+    indexes = (torch.log(scales) - log_scale_min) * log_step_recip
+    indexes = indexes.to(dtype=torch.uint8)
+    if skip_thres is not None:
+        skip_cond = scales > skip_thres
+    return indexes, skip_cond
+
+
+def build_index_enc(symbols, scales, scale_min, scale_max, log_scale_min,
+                    log_step_recip, skip_thres=None):
+    if CUSTOMIZED_CUDA_INFERENCE and scales.is_cuda:
+        out = torch.empty_like(scales, dtype=torch.int16)
+        skip_cond = None
+        if skip_thres is not None:
+            skip_cond = torch.empty_like(scales, dtype=torch.bool)
+        else:
+            skip_thres = -1.
+
+        build_index_enc_cuda(out, skip_cond, symbols, scales, scale_min, scale_max, log_scale_min,
+                             log_step_recip, skip_thres)
+
+        out = out[skip_cond]
+        return out
+
+    scales = scales.clamp_(scale_min, scale_max)
+    indexes = (torch.log(scales) - log_scale_min) * log_step_recip
+    indexes = indexes.to(dtype=torch.uint8)
+    symbols = symbols.to(dtype=torch.int16)
+    out = (symbols << 8) + indexes
+    out = out.to(dtype=torch.int16)
+    if skip_thres is not None:
+        skip_cond = scales > skip_thres
+        out = out[skip_cond]
+    return out
+
+
+def replicate_pad(x, pad_b, pad_r):
+    if pad_b == 0 and pad_r == 0:
+        return x
+    if CUSTOMIZED_CUDA_INFERENCE and x.is_cuda:
+        return replicate_pad_cuda(x, pad_b, pad_r)
+    return F.pad(x, (0, pad_r, 0, pad_b), mode="replicate")
+
+
+def bias_pixel_shuffle_8(x, bias):
+    if CUSTOMIZED_CUDA_INFERENCE and x.is_cuda:
+        B, C, H, W = x.shape
+        assert B == 1
+        out = torch.empty((B, 3, H * 8, W * 8), dtype=x.dtype, device=x.device, layout=x.layout)
+        bias_pixel_shuffle_8_cuda(out, x, bias, C, H * W, W, True)
+        return out
+
+    out = x + bias[None, :, None, None]
+    out = F.pixel_shuffle(out, 8)
+    out = torch.clamp(out, 0., 1.)
+    return out
+
+
+def bias_quant(x, bias, quant_step):
+    if CUSTOMIZED_CUDA_INFERENCE and x.is_cuda:
+        bias_quant_cuda(x, bias, quant_step)
+        return x
+
+    out = x + bias[None, :, None, None]
+    out = out * quant_step
+    return out

neural_codec/DCVC/src/layers/extensions/inference/bind.cpp

@@ -0,0 +1,36 @@
+// Copyright (c) Microsoft Corporation.
+// Licensed under the MIT License.
+
+#include "def.h"
+#include <torch/extension.h>
+
+PYBIND11_MODULE(TORCH_EXTENSION_NAME, m)
+{
+    m.def("process_with_mask_cuda", &process_with_mask_cuda);
+    m.def("combine_for_reading_2x_cuda", &combine_for_reading_2x_cuda);
+    m.def("restore_y_2x_cuda", &restore_y_2x_cuda);
+    m.def("restore_y_4x_cuda", &restore_y_4x_cuda);
+    m.def("build_index_dec_cuda", &build_index_dec_cuda);
+    m.def("build_index_enc_cuda", &build_index_enc_cuda);
+    m.def("bias_quant_cuda", &bias_quant_cuda);
+    m.def("round_and_to_int8_cuda", &round_and_to_int8_cuda);
+    m.def("clamp_reciprocal_with_quant_cuda", &clamp_reciprocal_with_quant_cuda);
+    m.def("add_and_multiply_cuda", &add_and_multiply_cuda);
+    m.def("bias_pixel_shuffle_8_cuda", &bias_pixel_shuffle_8_cuda);
+    m.def("replicate_pad_cuda", &replicate_pad_cuda);
+    m.def("bias_wsilu_depthwise_conv2d_cuda", &bias_wsilu_depthwise_conv2d_cuda);
+
+    py::class_<DepthConvProxy>(m, "DepthConvProxy")
+        .def(py::init<>())
+        .def("set_param", &DepthConvProxy::set_param)
+        .def("set_param_with_adaptor", &DepthConvProxy::set_param_with_adaptor)
+        .def("forward", &DepthConvProxy::forward)
+        .def("forward_with_quant_step", &DepthConvProxy::forward_with_quant_step)
+        .def("forward_with_cat", &DepthConvProxy::forward_with_cat);
+
+    py::class_<SubpelConv2xProxy>(m, "SubpelConv2xProxy")
+        .def(py::init<>())
+        .def("set_param", &SubpelConv2xProxy::set_param)
+        .def("forward", &SubpelConv2xProxy::forward)
+        .def("forward_with_cat", &SubpelConv2xProxy::forward_with_cat);
+}

neural_codec/DCVC/src/layers/extensions/inference/common.h

@@ -0,0 +1,352 @@
+// Copyright (c) Microsoft Corporation.
+// Licensed under the MIT License.
+
+#pragma once
+#include <ATen/cuda/CUDAContext.h>
+#include <cuda.h>
+#include <cuda_fp16.h>
+#include <cuda_runtime.h>
+#include <torch/extension.h>
+#include <type_traits>
+#include <utility>
+
+// Constrain the generic vec4 helper templates below to genuine 4-lane vector
+// types (float4/Half4/short4/char4/uchar4/bool4). We detect them by the
+// presence of a `.w` member: all four-component vectors have `.x/.y/.z/.w`,
+// whereas scalar library types do NOT — crucially `c10::Half` (and
+// `c10::BFloat16`) store their raw bits in a member literally named `x`, so a
+// `.x`-based probe would wrongly classify Half as a vector and make these
+// templates collide with the scalar Half operators (and with torch's own
+// operators) under newer pybind11 / C++20 headers. `.w` is unique to the real
+// vectors, so it cleanly separates them.
+namespace dcvc_detail {
+template <typename T, typename = void>
+struct is_vec4 : ::std::false_type {};
+template <typename T>
+struct is_vec4<T, ::std::void_t<decltype(::std::declval<T&>().w)>> : ::std::true_type {};
+}  // namespace dcvc_detail
+#define DCVC_VEC_ENABLE(TT) typename ::std::enable_if<dcvc_detail::is_vec4<TT>::value, int>::type = 0
+
+// T maybe vector type, and may be different from t.dtype
+template <typename T>
+struct GPUTensor1D {
+    GPUTensor1D(torch::Tensor& t) : ptr(static_cast<T*>(t.data_ptr())) {}
+    GPUTensor1D(const torch::Tensor& t) : ptr(static_cast<T*>(t.data_ptr())) {}
+    GPUTensor1D(T* t) : ptr(static_cast<T*>(t)) { assert(t == nullptr); }
+
+    __device__ T& operator[](int idx) { return ptr[idx]; }
+    __device__ T& operator[](int idx) const { return ptr[idx]; }
+
+    T* __restrict__ const ptr;
+};
+
+template <typename scalar_t>
+using Packed4DTensorAccessor32 = torch::PackedTensorAccessor32<scalar_t, 4, torch::RestrictPtrTraits>;
+template <typename scalar_t>
+using Packed1DTensorAccessor32 = torch::PackedTensorAccessor32<scalar_t, 1, torch::RestrictPtrTraits>;
+
+struct __align__(8) Half4
+{
+    c10::Half x;
+    c10::Half y;
+    c10::Half z;
+    c10::Half w;
+};
+
+struct __align__(4) bool4
+{
+    bool x;
+    bool y;
+    bool z;
+    bool w;
+};
+
+__forceinline__ __device__ float4 make_vec4(const float& x, const float& y, const float& z,
+                                            const float& w)
+{
+    return make_float4(x, y, z, w);
+}
+
+__forceinline__ __device__ Half4 make_vec4(const c10::Half& x, const c10::Half& y,
+                                           const c10::Half& z, const c10::Half& w)
+{
+    Half4 t;
+    t.x = x;
+    t.y = y;
+    t.z = z;
+    t.w = w;
+    return t;
+}
+
+__forceinline__ __device__ Half4 make_Half4(const c10::Half& x, const c10::Half& y,
+                                            const c10::Half& z, const c10::Half& w)
+{
+    Half4 t;
+    t.x = x;
+    t.y = y;
+    t.z = z;
+    t.w = w;
+    return t;
+}
+
+__forceinline__ __device__ bool4 make_vec4(const bool& x, const bool& y, const bool& z, const bool& w)
+{
+    bool4 t;
+    t.x = x;
+    t.y = y;
+    t.z = z;
+    t.w = w;
+    return t;
+}
+
+__forceinline__ __device__ c10::Half round(const c10::Half& a)
+{
+    return static_cast<c10::Half>(__half2int_rn(a));
+}
+
+template <typename T, DCVC_VEC_ENABLE(T)>
+__forceinline__ __device__ T round(const T& a)
+{
+    return make_vec4(round(a.x), round(a.y), round(a.z), round(a.w));
+}
+
+__forceinline__ __device__ int8_t to_int8(const float& a)
+{
+    return static_cast<int8_t>(a);
+}
+
+__forceinline__ __device__ int8_t to_int8(const c10::Half& a)
+{
+    return static_cast<int8_t>(a);
+}
+
+template <typename T, DCVC_VEC_ENABLE(T)>
+__forceinline__ __device__ char4 to_int8(const T& a)
+{
+    return make_char4(to_int8(a.x), to_int8(a.y), to_int8(a.z), to_int8(a.w));
+}
+
+__forceinline__ __device__ uint8_t to_uint8(const float& a)
+{
+    return static_cast<uint8_t>(a);
+}
+
+__forceinline__ __device__ uint8_t to_uint8(const c10::Half& a)
+{
+    return static_cast<uint8_t>(__half2uint_rd(a));
+}
+
+template <typename T, DCVC_VEC_ENABLE(T)>
+__forceinline__ __device__ uchar4 to_uint8(const T& a)
+{
+    return make_uchar4(to_uint8(a.x), to_uint8(a.y), to_uint8(a.z), to_uint8(a.w));
+}
+
+__forceinline__ __device__ int16_t to_int16(const float& a)
+{
+    return static_cast<int16_t>(a);
+}
+
+__forceinline__ __device__ int16_t to_int16(const c10::Half& a)
+{
+    return static_cast<int16_t>(__half2int_rd(a));
+}
+
+template <typename T, DCVC_VEC_ENABLE(T)>
+__forceinline__ __device__ short4 to_int16(const T& a)
+{
+    return make_short4(to_int16(a.x), to_int16(a.y), to_int16(a.z), to_int16(a.w));
+}
+
+__forceinline__ __device__ short4 operator<<(const short4& a, const int b)
+{
+    return make_short4(a.x << b, a.y << b, a.z << b, a.w << b);
+}
+
+__forceinline__ __device__ c10::Half min(const c10::Half& a, const c10::Half& b)
+{
+    return __hmin(a, b);
+}
+__forceinline__ __device__ c10::Half max(const c10::Half& a, const c10::Half& b)
+{
+    return __hmax(a, b);
+}
+
+// Native-half compare. Constrained to EXACT c10::Half operands (via SFINAE) so
+// it can never become a viable candidate for a *mixed* comparison such as
+// `int64_t > c10::Half` inside torch's C++20 headers (TypeSafeSignMath.h). If
+// it did, the implicit int64_t->c10::Half conversion would make it tie with the
+// built-in `<=>`/`>` operators and the build fails with "more than one operator
+// matches" (seen under torch 2.12 / CUDA 13.3, which compile with -std=c++20).
+template <typename A, typename B,
+          typename ::std::enable_if<::std::is_same<A, c10::Half>::value &&
+                                        ::std::is_same<B, c10::Half>::value,
+                                    int>::type = 0>
+__forceinline__ __device__ bool operator>(const A& a, const B& b)
+{
+    return __hgt(a, b);
+}
+
+template <typename A, typename B,
+          typename ::std::enable_if<::std::is_same<A, c10::Half>::value &&
+                                        ::std::is_same<B, c10::Half>::value,
+                                    int>::type = 0>
+__forceinline__ __device__ bool operator<(const A& a, const B& b)
+{
+    return __hlt(a, b);
+}
+
+__forceinline__ __device__ c10::Half log(const c10::Half& a)
+{
+    return hlog(a);
+}
+
+__forceinline__ __device__ short4 operator+(const short4& a, const short4& b)
+{
+    return make_short4(a.x + b.x, a.y + b.y, a.z + b.z, a.w + b.w);
+}
+
+__forceinline__ __device__ float4 operator+(const float4& a, const float4& b)
+{
+    return make_float4(a.x + b.x, a.y + b.y, a.z + b.z, a.w + b.w);
+}
+
+__forceinline__ __device__ Half4 operator+(const Half4& a, const Half4& b)
+{
+    return make_Half4(a.x + b.x, a.y + b.y, a.z + b.z, a.w + b.w);
+}
+
+template <typename T, DCVC_VEC_ENABLE(T)>
+__forceinline__ __device__ T operator-(const T& a, const T& b)
+{
+    return make_vec4(a.x - b.x, a.y - b.y, a.z - b.z, a.w - b.w);
+}
+
+__forceinline__ __device__ float4 operator-(const float4& a, const float& b)
+{
+    return make_vec4(a.x - b, a.y - b, a.z - b, a.w - b);
+}
+
+__forceinline__ __device__ Half4 operator-(const Half4& a, const c10::Half& b)
+{
+    return make_vec4(a.x - b, a.y - b, a.z - b, a.w - b);
+}
+
+__forceinline__ __device__ c10::Half operator*(const c10::Half& a, const bool b)
+{
+    return b ? a : static_cast<c10::Half>(0.f);
+}
+
+template <typename T, DCVC_VEC_ENABLE(T)>
+__forceinline__ __device__ T operator*(const T& a, const T& b)
+{
+    return make_vec4(a.x * b.x, a.y * b.y, a.z * b.z, a.w * b.w);
+}
+
+template <typename T, DCVC_VEC_ENABLE(T)>
+__forceinline__ __device__ T operator*(const T& a, const bool4& b)
+{
+    return make_vec4(a.x * b.x, a.y * b.y, a.z * b.z, a.w * b.w);
+}
+
+template <typename T1, typename T2, DCVC_VEC_ENABLE(T1)>
+__forceinline__ __device__ T1 operator*(const T1& a, const T2& b)
+{
+    return make_vec4(a.x * b, a.y * b, a.z * b, a.w * b);
+}
+
+template <typename T, DCVC_VEC_ENABLE(T)>
+__forceinline__ __device__ T max(const T& a, const T& b)
+{
+    return make_vec4(max(a.x, b.x), max(a.y, b.y), max(a.z, b.z), max(a.w, b.w));
+}
+
+template <typename T1, typename T2, DCVC_VEC_ENABLE(T1)>
+__forceinline__ __device__ T1 max(const T1& a, const T2& b)
+{
+    return make_vec4(max(a.x, b), max(a.y, b), max(a.z, b), max(a.w, b));
+}
+
+template <typename T, DCVC_VEC_ENABLE(T)>
+__forceinline__ __device__ T min(const T& a, const T& b)
+{
+    return make_vec4(min(a.x, b.x), min(a.y, b.y), min(a.z, b.z), min(a.w, b.w));
+}
+
+template <typename T1, typename T2, DCVC_VEC_ENABLE(T1)>
+__forceinline__ __device__ T1 min(const T1& a, const T2& b)
+{
+    return make_vec4(min(a.x, b), min(a.y, b), min(a.z, b), min(a.w, b));
+}
+
+template <typename T, DCVC_VEC_ENABLE(T)>
+__forceinline__ __device__ T log(const T& a)
+{
+    return make_vec4(log(a.x), log(a.y), log(a.z), log(a.w));
+}
+
+__forceinline__ __device__ float reciprocal(const float& a)
+{
+    return __frcp_rd(a);
+}
+
+__forceinline__ __device__ c10::Half reciprocal(const c10::Half& a)
+{
+    return hrcp(a);
+}
+
+template <typename T, DCVC_VEC_ENABLE(T)>
+__forceinline__ __device__ T reciprocal(const T& a)
+{
+    return make_vec4(reciprocal(a.x), reciprocal(a.y), reciprocal(a.z), reciprocal(a.w));
+}
+
+template <typename T1, typename T2, DCVC_VEC_ENABLE(T1)>
+__forceinline__ __device__ bool4 operator>(const T1& a, const T2& b)
+{
+    return make_vec4(a.x > b, a.y > b, a.z > b, a.w > b);
+}
+
+__forceinline__ __device__ float sigmoid(const float x)
+{
+    return 1.0f / (1.0f + expf(-x));
+}
+
+__forceinline__ __device__ float wsilu(const float x)
+{
+    return x * sigmoid(4.0f * x);
+}
+
+__forceinline__ __device__ c10::Half wsilu(const c10::Half x)
+{
+    return __float2half_rn(wsilu(__half2float(x)));
+}
+
+__forceinline__ __device__ float4 wsilu(float4 data)
+{
+    data.x = wsilu(data.x);
+    data.y = wsilu(data.y);
+    data.z = wsilu(data.z);
+    data.w = wsilu(data.w);
+    return data;
+}
+
+__forceinline__ __device__ Half4 wsilu(Half4 data)
+{
+    data.x = wsilu(data.x);
+    data.y = wsilu(data.y);
+    data.z = wsilu(data.z);
+    data.w = wsilu(data.w);
+    return data;
+}
+
+__forceinline__ __device__ float multiply_add(const float a, const float b, const float c)
+{
+    return __fmaf_rn(a, b, c);
+}
+
+__forceinline__ __device__ c10::Half multiply_add(const c10::Half a, const c10::Half b, const c10::Half c)
+{
+
+    return __hfma(a, b, c);
+}

neural_codec/DCVC/src/layers/extensions/inference/def.h

@@ -0,0 +1,107 @@
+// Copyright (c) Microsoft Corporation.
+// Licensed under the MIT License.
+
+#include <torch/extension.h>
+
+std::tuple<torch::Tensor, torch::Tensor, torch::Tensor, torch::Tensor>
+process_with_mask_cuda(const torch::Tensor& y, const torch::Tensor& scales, const torch::Tensor& means,
+                       const torch::Tensor& mask, const float force_zero_thres);
+
+void combine_for_reading_2x_cuda(torch::Tensor& out, const torch::Tensor& x, const torch::Tensor& mask);
+void restore_y_2x_cuda(torch::Tensor& out, const torch::Tensor& y, const torch::Tensor& means,
+                       const torch::Tensor& mask);
+void restore_y_4x_cuda(torch::Tensor& out, const torch::Tensor& y, const torch::Tensor& means,
+                       const torch::Tensor& mask);
+
+void build_index_dec_cuda(torch::Tensor& out, torch::optional<torch::Tensor>& cond_out,
+                          const torch::Tensor& scales, const float scale_min, const float scale_max,
+                          const float log_scale_min, const float log_step_recip,
+                          const float skip_thres);
+
+void build_index_enc_cuda(torch::Tensor& out, torch::optional<torch::Tensor>& cond_out,
+                          const torch::Tensor& symbols, const torch::Tensor& scales,
+                          const float scale_min, const float scale_max, const float log_scale_min,
+                          const float log_step_recip, const float skip_thres);
+
+void bias_wsilu_cuda(torch::Tensor& x, const torch::Tensor& bias);
+
+void bias_shortcut_cuda(torch::Tensor& x, const torch::Tensor& bias, const torch::Tensor& shortcut);
+void bias_shortcut_no_inplace_cuda(torch::Tensor& out, const torch::Tensor& x,
+                                   const torch::Tensor& bias, const torch::Tensor& shortcut);
+void bias_shortcut_2_cuda(torch::Tensor& x, const torch::Tensor& bias, torch::Tensor& shortcut);
+void bias_shortcut_with_quant_step_cuda(torch::Tensor& x, const torch::Tensor& bias,
+                                        const torch::Tensor& quant_step, const torch::Tensor& shortcut);
+
+void bias_quant_cuda(torch::Tensor& x, const torch::Tensor& bias, const torch::Tensor& quant_step);
+
+void bias_wsilu_chunk_add_cuda(torch::Tensor& x, const torch::Tensor& bias);
+
+void bias_pixel_shuffle_2_cuda(torch::Tensor& out, const torch::Tensor& x,
+                               const torch::Tensor& bias, const int C, const int N, const int W);
+void bias_pixel_shuffle_8_cuda(torch::Tensor& out, const torch::Tensor& x, const torch::Tensor& bias,
+                               const int C, const int N, const int W, bool clamp);
+torch::Tensor replicate_pad_cuda(const torch::Tensor& x, const int padB, const int padR);
+
+torch::Tensor round_and_to_int8_cuda(torch::Tensor& z);
+torch::Tensor clamp_reciprocal_with_quant_cuda(const torch::Tensor& q_dec, torch::Tensor& y,
+                                               const float min_val);
+void add_and_multiply_cuda(torch::Tensor& x0, const torch::Tensor& x1, const torch::Tensor q);
+
+torch::Tensor bias_wsilu_depthwise_conv2d_cuda(const torch::Tensor& x, const torch::Tensor& weight,
+                                               const torch::Tensor& bias);
+
+class DepthConvProxy {
+public:
+    DepthConvProxy() = default;
+    ~DepthConvProxy() = default;
+
+    void set_param(const torch::Tensor& dc_conv1_weight, const torch::Tensor& dc_conv1_bias,
+                   const torch::Tensor& dc_depth_conv_weight,
+                   const torch::Tensor& dc_depth_conv_bias, const torch::Tensor& dc_conv2_weight,
+                   const torch::Tensor& dc_conv2_bias, const torch::Tensor& ffn_conv1_weight,
+                   const torch::Tensor& ffn_conv1_bias, const torch::Tensor& ffn_conv2_weight,
+                   const torch::Tensor& ffn_conv2_bias, const bool shortcut);
+    void set_param_with_adaptor(
+        const torch::Tensor& dc_conv1_weight, const torch::Tensor& dc_conv1_bias,
+        const torch::Tensor& dc_depth_conv_weight, const torch::Tensor& dc_depth_conv_bias,
+        const torch::Tensor& dc_conv2_weight, const torch::Tensor& dc_conv2_bias,
+        const torch::Tensor& ffn_conv1_weight, const torch::Tensor& ffn_conv1_bias,
+        const torch::Tensor& ffn_conv2_weight, const torch::Tensor& ffn_conv2_bias,
+        const torch::Tensor& adaptor_weight, const torch::Tensor& adaptor_bias, const bool shortcut);
+    torch::Tensor forward(const torch::Tensor& x);
+    torch::Tensor forward_with_quant_step(const torch::Tensor& x, const torch::Tensor& quant_step);
+    torch::Tensor forward_with_cat(const torch::Tensor& x, const torch::Tensor& to_cat,
+                                   const bool cat_at_front);
+
+private:
+    std::tuple<torch::Tensor, torch::Tensor> forward_common(const torch::Tensor& x);
+
+private:
+    torch::Tensor _dc_conv1_weight;
+    torch::Tensor _dc_conv1_bias;
+    torch::Tensor _dc_depth_conv_weight;
+    torch::Tensor _dc_conv2_weight;
+    torch::Tensor _dc_conv2_bias;
+    torch::Tensor _ffn_conv1_weight;
+    torch::Tensor _ffn_conv1_bias;
+    torch::Tensor _ffn_conv2_weight;
+    torch::Tensor _ffn_conv2_bias;
+    bool _adaptor{ false };
+    bool _shortcut{ false };
+};
+
+class SubpelConv2xProxy {
+public:
+    SubpelConv2xProxy() = default;
+    ~SubpelConv2xProxy() = default;
+
+    void set_param(const torch::Tensor& weight, const torch::Tensor& bias, const int padding);
+    torch::Tensor forward(const torch::Tensor& x);
+    torch::Tensor forward_with_cat(const torch::Tensor& x, const torch::Tensor& to_cat,
+                                   const bool cat_at_front);
+
+private:
+    torch::Tensor _weight;
+    torch::Tensor _bias;
+    int _padding{ 0 };
+};

neural_codec/DCVC/src/layers/extensions/inference/impl.cpp

@@ -0,0 +1,167 @@
+// Copyright (c) Microsoft Corporation.
+// Licensed under the MIT License.
+
+#include "def.h"
+namespace F = torch::nn::functional;
+
+void DepthConvProxy::set_param(
+    const torch::Tensor& dc_conv1_weight, const torch::Tensor& dc_conv1_bias,
+    const torch::Tensor& dc_depth_conv_weight, const torch::Tensor& dc_depth_conv_bias,
+    const torch::Tensor& dc_conv2_weight, const torch::Tensor& dc_conv2_bias,
+    const torch::Tensor& ffn_conv1_weight, const torch::Tensor& ffn_conv1_bias,
+    const torch::Tensor& ffn_conv2_weight, const torch::Tensor& ffn_conv2_bias, const bool shortcut)
+{
+    _dc_conv1_weight = dc_conv1_weight;
+    _dc_conv1_bias = dc_conv1_bias;
+    _dc_depth_conv_weight = dc_depth_conv_weight;
+    _dc_conv2_weight = dc_conv2_weight;
+    _dc_conv2_bias = F::conv2d(dc_depth_conv_bias.reshape({ 1, -1, 1, 1 }), dc_conv2_weight);
+    _dc_conv2_bias = _dc_conv2_bias.index({ 0, torch::indexing::Slice(), 0, 0 }) + dc_conv2_bias;
+    _ffn_conv1_weight = ffn_conv1_weight;
+    _ffn_conv1_bias = ffn_conv1_bias;
+    _ffn_conv2_weight = ffn_conv2_weight;
+    _ffn_conv2_bias = ffn_conv2_bias;
+    _shortcut = shortcut;
+}
+
+void DepthConvProxy::set_param_with_adaptor(
+    const torch::Tensor& dc_conv1_weight, const torch::Tensor& dc_conv1_bias,
+    const torch::Tensor& dc_depth_conv_weight, const torch::Tensor& dc_depth_conv_bias,
+    const torch::Tensor& dc_conv2_weight, const torch::Tensor& dc_conv2_bias,
+    const torch::Tensor& ffn_conv1_weight, const torch::Tensor& ffn_conv1_bias,
+    const torch::Tensor& ffn_conv2_weight, const torch::Tensor& ffn_conv2_bias,
+    const torch::Tensor& adaptor_weight, const torch::Tensor& adaptor_bias, const bool shortcut)
+{
+    _dc_conv1_weight = F::conv2d(torch::transpose(adaptor_weight, 0, 1), dc_conv1_weight);
+    _dc_conv1_weight = torch::transpose(_dc_conv1_weight, 0, 1);
+    _dc_conv1_weight = torch::cat({ _dc_conv1_weight, adaptor_weight }, 0);
+    _dc_conv1_bias = F::conv2d(adaptor_bias.reshape({ 1, -1, 1, 1 }), dc_conv1_weight);
+    _dc_conv1_bias = _dc_conv1_bias.index({ 0, torch::indexing::Slice(), 0, 0 }) + dc_conv1_bias;
+    _dc_depth_conv_weight = dc_depth_conv_weight;
+    _dc_conv2_weight = dc_conv2_weight;
+    _dc_conv2_bias = F::conv2d(dc_depth_conv_bias.reshape({ 1, -1, 1, 1 }), dc_conv2_weight);
+    _dc_conv2_bias = _dc_conv2_bias.index({ 0, torch::indexing::Slice(), 0, 0 }) + dc_conv2_bias;
+    _dc_conv2_bias = _dc_conv2_bias + adaptor_bias;
+    _ffn_conv1_weight = ffn_conv1_weight;
+    _ffn_conv1_bias = ffn_conv1_bias;
+    _ffn_conv2_weight = ffn_conv2_weight;
+    _ffn_conv2_bias = ffn_conv2_bias;
+    _shortcut = shortcut;
+    _adaptor = true;
+}
+
+std::tuple<torch::Tensor, torch::Tensor> DepthConvProxy::forward_common(const torch::Tensor& x)
+{
+    auto identity = x;
+    // depthconv
+    torch::Tensor out;
+    if (_adaptor) {
+        // NOTE: Here we always fuse adaptor with the first conv1x1 (when even in_ch > out_ch).
+        //       It brings larger MACs, but it faster on A100 due to lower memory cost.
+        auto out_identity = F::conv2d(identity, _dc_conv1_weight);
+        auto chunks = torch::chunk(out_identity, 2, 1);
+        out = chunks[0];
+        identity = chunks[1];
+    } else {
+        out = F::conv2d(identity, _dc_conv1_weight);
+    }
+    out = bias_wsilu_depthwise_conv2d_cuda(out, _dc_depth_conv_weight, _dc_conv1_bias);
+    out = F::conv2d(out, _dc_conv2_weight);
+
+    if (_shortcut) {
+        bias_shortcut_2_cuda(out, _dc_conv2_bias, identity);
+    } else {
+        bias_shortcut_cuda(out, _dc_conv2_bias, identity);
+        identity = out;
+    }
+    // ffn
+    out = F::conv2d(out, _ffn_conv1_weight);
+    bias_wsilu_chunk_add_cuda(out, _ffn_conv1_bias);
+    out = F::conv2d(out, _ffn_conv2_weight);
+    return { out, identity };
+}
+
+torch::Tensor DepthConvProxy::forward(const torch::Tensor& x)
+{
+    auto [out, identity] = forward_common(x);
+    bias_shortcut_cuda(out, _ffn_conv2_bias, identity);
+    return out;
+}
+
+torch::Tensor DepthConvProxy::forward_with_quant_step(const torch::Tensor& x,
+                                                      const torch::Tensor& quant_step)
+{
+    auto [out, identity] = forward_common(x);
+    bias_shortcut_with_quant_step_cuda(out, _ffn_conv2_bias, quant_step, identity);
+    return out;
+}
+
+torch::Tensor DepthConvProxy::forward_with_cat(const torch::Tensor& x, const torch::Tensor& to_cat,
+                                               const bool cat_at_front)
+{
+    auto [t, identity] = forward_common(x);
+
+    auto t_shape = t.sizes();
+    auto B = t_shape[0];
+    auto C = t_shape[1];
+    auto H = t_shape[2];
+    auto W = t_shape[3];
+    auto add_ch = to_cat.sizes()[1];
+    auto out = torch::empty({ B, C + add_ch, H, W }, t.options());
+    if (cat_at_front) {
+        auto t_out = out.narrow(1, add_ch, C);
+        bias_shortcut_no_inplace_cuda(t_out, t, _ffn_conv2_bias, identity);
+        out.narrow(1, 0, add_ch) = to_cat;
+    } else {
+        auto t_out = out.narrow(1, 0, C);
+        bias_shortcut_no_inplace_cuda(t_out, t, _ffn_conv2_bias, identity);
+        out.narrow(1, C, add_ch) = to_cat;
+    }
+    return out;
+}
+
+void SubpelConv2xProxy::set_param(const torch::Tensor& weight, const torch::Tensor& bias,
+                                  const int padding)
+{
+    _weight = weight;
+    _bias = bias;
+    _padding = padding;
+}
+
+torch::Tensor SubpelConv2xProxy::forward(const torch::Tensor& x)
+{
+    auto t = F::conv2d(x, _weight, F::Conv2dFuncOptions().padding(_padding));
+    auto t_shape = t.sizes();
+    auto B = t_shape[0];
+    auto C = t_shape[1];
+    auto H = t_shape[2];
+    auto W = t_shape[3];
+    auto out = torch::empty({ B, C / 4, H * 2, W * 2 }, t.options());
+    assert(B == 1);
+    bias_pixel_shuffle_2_cuda(out, t, _bias, C, H * W, W);
+    return out;
+}
+
+torch::Tensor SubpelConv2xProxy::forward_with_cat(const torch::Tensor& x, const torch::Tensor& to_cat,
+                                                  const bool cat_at_front)
+{
+    auto t = F::conv2d(x, _weight, F::Conv2dFuncOptions().padding(_padding));
+    auto t_shape = t.sizes();
+    auto B = t_shape[0];
+    auto C = t_shape[1];
+    auto H = t_shape[2];
+    auto W = t_shape[3];
+    auto add_ch = to_cat.sizes()[1];
+    auto out = torch::empty({ B, add_ch + C / 4, H * 2, W * 2 }, t.options());
+    assert(B == 1);
+    if (cat_at_front) {
+        auto t_out = out.narrow(1, add_ch, C / 4);
+        bias_pixel_shuffle_2_cuda(t_out, t, _bias, C, H * W, W);
+        out.narrow(1, 0, add_ch) = to_cat;
+    } else {
+        auto t_out = out.narrow(1, 0, C / 4);
+        bias_pixel_shuffle_2_cuda(t_out, t, _bias, C, H * W, W);
+        out.narrow(1, C / 4, add_ch) = to_cat;
+    }
+    return out;
+}

neural_codec/DCVC/src/layers/extensions/inference/kernel.cu

@@ -0,0 +1,1150 @@
+// Copyright (c) Microsoft Corporation.
+// Licensed under the MIT License.
+
+#include <ATen/cuda/CUDAContext.h>
+#include <cstdint>
+#include <cuda.h>
+#include <cuda_fp16.h>
+#include <cuda_runtime.h>
+#include <torch/extension.h>
+
+#include "common.h"
+#include "def.h"
+#include <vector>
+
+template <typename vec_t>
+__forceinline__ __host__ bool can_vectorize(void* pointer)
+{
+    uint64_t address = reinterpret_cast<uint64_t>(pointer);
+    constexpr int vec4_alignment = std::alignment_of<vec_t>::value;
+    return address % vec4_alignment == 0;
+}
+
+template <typename vec_t>
+__forceinline__ std::tuple<dim3, dim3, at::cuda::CUDAStream, bool, bool, int, int>
+get_kernel_launch_info(const torch::Tensor& x, const int cDiv = 1, const bool allow_useVec = true)
+{
+    const torch::IntArrayRef x_shape = x.sizes();
+    const int B = x_shape[0];
+    assert(B == 1);
+    const int C = x_shape[1];
+    const int HW = x_shape[2] * x_shape[3];
+    const int N = C * HW / cDiv;
+    const int BLOCK_SIZE = 128;
+    const dim3 blockDim(BLOCK_SIZE);
+    const bool useVec = allow_useVec && N % 4 == 0 && can_vectorize<vec_t>(x.data_ptr());
+    const bool biasSafe = HW % 4 == 0;
+    const int factor = useVec ? 4 : 1;
+    const dim3 gridDim((N / factor + BLOCK_SIZE - 1) / BLOCK_SIZE);
+    return { blockDim, gridDim, at::cuda::getCurrentCUDAStream(), useVec, biasSafe, N / factor, HW };
+}
+
+template <typename vec_t>
+__forceinline__ std::tuple<dim3, dim3, at::cuda::CUDAStream, bool, int>
+get_kernel_launch_info_flatten(const torch::Tensor& x)
+{
+    const int N = x.numel();
+    const int BLOCK_SIZE = 128;
+    const dim3 blockDim(BLOCK_SIZE);
+    const bool useVec = N % 4 == 0 && can_vectorize<vec_t>(x.data_ptr());
+    const int factor = useVec ? 4 : 1;
+    const dim3 gridDim((N / factor + BLOCK_SIZE - 1) / BLOCK_SIZE);
+    return { blockDim, gridDim, at::cuda::getCurrentCUDAStream(), useVec, N / factor };
+}
+
+template <typename scalar_t, typename T, bool forceZero = false>
+__global__ void process_with_mask_kernel(GPUTensor1D<T> y_res, GPUTensor1D<T> y_q,
+                                         GPUTensor1D<T> y_hat, GPUTensor1D<T> s_hat,
+                                         const GPUTensor1D<T> y, const GPUTensor1D<T> scales,
+                                         const GPUTensor1D<T> means, const GPUTensor1D<T> mask,
+                                         const scalar_t force_zero_thres, const int N)
+{
+    const scalar_t __min_val = static_cast<scalar_t>(-128.f);
+    const scalar_t __max_val = static_cast<scalar_t>(127.f);
+    const int chw = blockIdx.x * blockDim.x + threadIdx.x;
+
+    if (chw < N) {
+        T _y = y[chw];
+        T _scale = scales[chw];
+        T _means = means[chw];
+        T _mask = mask[chw];
+
+        T _s_hat = _scale * _mask;
+        T _means_hat = _means * _mask;
+        T _y_res = (_y - _means_hat) * _mask;
+        T _y_q = round(_y_res);
+
+        if constexpr (forceZero) {
+            _y_q = _y_q * (_s_hat > force_zero_thres);
+        }
+        _y_q = max(min(_y_q, __max_val), __min_val);
+        T _y_hat = _y_q + _means_hat;
+
+        y_res[chw] = _y_res;
+        y_q[chw] = _y_q;
+        y_hat[chw] = _y_hat;
+        s_hat[chw] = _s_hat;
+    }
+}
+
+template <typename scalar_t, typename vec_t>
+__forceinline__ void
+process_with_mask_dispatcher(torch::Tensor& y_res, torch::Tensor& y_q, torch::Tensor& y_hat,
+                             torch::Tensor& s_hat, const torch::Tensor& y,
+                             const torch::Tensor& scales, const torch::Tensor& means,
+                             const torch::Tensor& mask, const float force_zero_thres)
+{
+    auto [blockDim, gridDim, stream, useVec, biasSafe, N, HW] = get_kernel_launch_info<vec_t>(y);
+    const bool force_zero = force_zero_thres > 0.f;
+
+    auto launch_kernel = [&](auto in_v) {
+        using in_t = decltype(in_v);
+        if (force_zero) {
+            process_with_mask_kernel<scalar_t, in_t, true>
+                <<<gridDim, blockDim, 0, stream>>>(y_res, y_q, y_hat, s_hat, y, scales, means, mask,
+                                                   static_cast<scalar_t>(force_zero_thres), N);
+        } else {
+            process_with_mask_kernel<scalar_t, in_t, false>
+                <<<gridDim, blockDim, 0, stream>>>(y_res, y_q, y_hat, s_hat, y, scales, means, mask,
+                                                   static_cast<scalar_t>(force_zero_thres), N);
+        }
+    };
+
+    if (useVec) {
+        launch_kernel(vec_t{});
+    } else {
+        launch_kernel(scalar_t{});
+    }
+}
+
+std::tuple<torch::Tensor, torch::Tensor, torch::Tensor, torch::Tensor>
+process_with_mask_cuda(const torch::Tensor& y, const torch::Tensor& scales, const torch::Tensor& means,
+                       const torch::Tensor& mask, const float force_zero_thres)
+{
+    auto y_res = torch::empty_like(y);
+    auto y_q = torch::empty_like(y);
+    auto y_hat = torch::empty_like(y);
+    auto s_hat = torch::empty_like(y);
+
+    if (y.dtype() == torch::kFloat32) {
+        process_with_mask_dispatcher<float, float4>(y_res, y_q, y_hat, s_hat, y, scales, means,
+                                                    mask, force_zero_thres);
+    } else if (y.dtype() == torch::kFloat16) {
+        process_with_mask_dispatcher<c10::Half, Half4>(y_res, y_q, y_hat, s_hat, y, scales, means,
+                                                       mask, force_zero_thres);
+    }
+
+    return { y_res, y_q, y_hat, s_hat };
+}
+
+template <typename T>
+__global__ void combine_for_reading_2x_kernel(GPUTensor1D<T> out, const GPUTensor1D<T> x,
+                                              const GPUTensor1D<T> mask, const int N)
+{
+    const int chw1 = blockIdx.x * blockDim.x + threadIdx.x;
+    const int chw2 = chw1 + N;
+
+    if (chw1 < N) {
+        T _s1 = x[chw1];
+        T _s2 = x[chw2];
+        T _m1 = mask[chw1];
+        T _m2 = mask[chw2];
+
+        _s1 = _s1 * _m1;
+        _s2 = _s2 * _m2;
+        out[chw1] = _s1 + _s2;
+    }
+}
+
+template <typename scalar_t, typename vec_t>
+__forceinline__ void combine_for_reading_2x_dispatcher(torch::Tensor& out, const torch::Tensor& x,
+                                                       const torch::Tensor& mask)
+{
+    auto [blockDim, gridDim, stream, useVec, biasSafe, N, HW] = get_kernel_launch_info<vec_t>(x, 2);
+    if (useVec) {
+        combine_for_reading_2x_kernel<vec_t><<<gridDim, blockDim, 0, stream>>>(out, x, mask, N);
+    } else {
+        combine_for_reading_2x_kernel<scalar_t><<<gridDim, blockDim, 0, stream>>>(out, x, mask, N);
+    }
+}
+
+void combine_for_reading_2x_cuda(torch::Tensor& out, const torch::Tensor& x, const torch::Tensor& mask)
+{
+    if (x.dtype() == torch::kFloat32) {
+        combine_for_reading_2x_dispatcher<float, float4>(out, x, mask);
+    } else if (x.dtype() == torch::kFloat16) {
+        combine_for_reading_2x_dispatcher<c10::Half, Half4>(out, x, mask);
+    }
+}
+
+template <typename T>
+__global__ void restore_y_2x_kernel(GPUTensor1D<T> out, const GPUTensor1D<T> y,
+                                    const GPUTensor1D<T> means, const GPUTensor1D<T> mask, const int N)
+{
+    const int chw1 = blockIdx.x * blockDim.x + threadIdx.x;
+    const int chw2 = chw1 + N;
+
+    if (chw1 < N) {
+        T _y = y[chw1];
+        T _means1 = means[chw1];
+        T _means2 = means[chw2];
+        T _mask1 = mask[chw1];
+        T _mask2 = mask[chw2];
+
+        _means1 = (_y + _means1) * _mask1;
+        _means2 = (_y + _means2) * _mask2;
+        out[chw1] = _means1;
+        out[chw2] = _means2;
+    }
+}
+
+template <typename scalar_t, typename vec_t>
+__forceinline__ void restore_y_2x_dispatcher(torch::Tensor& out, const torch::Tensor& y,
+                                             const torch::Tensor& means, const torch::Tensor& mask)
+{
+    auto [blockDim, gridDim, stream, useVec, biasSafe, N, HW] = get_kernel_launch_info<vec_t>(y);
+    if (useVec) {
+        restore_y_2x_kernel<vec_t><<<gridDim, blockDim, 0, stream>>>(out, y, means, mask, N);
+    } else {
+        restore_y_2x_kernel<scalar_t><<<gridDim, blockDim, 0, stream>>>(out, y, means, mask, N);
+    }
+}
+
+void restore_y_2x_cuda(torch::Tensor& out, const torch::Tensor& y, const torch::Tensor& means,
+                       const torch::Tensor& mask)
+{
+    if (y.dtype() == torch::kFloat32) {
+        restore_y_2x_dispatcher<float, float4>(out, y, means, mask);
+    } else if (y.dtype() == torch::kFloat16) {
+        restore_y_2x_dispatcher<c10::Half, Half4>(out, y, means, mask);
+    }
+}
+
+template <typename T>
+__global__ void restore_y_4x_kernel(GPUTensor1D<T> out, const GPUTensor1D<T> y,
+                                    const GPUTensor1D<T> means, const GPUTensor1D<T> mask, const int N)
+{
+    const int chw1 = blockIdx.x * blockDim.x + threadIdx.x;
+    const int chw2 = chw1 + N;
+    const int chw3 = chw2 + N;
+    const int chw4 = chw3 + N;
+
+    if (chw1 < N) {
+        T _y = y[chw1];
+        T _means1 = means[chw1];
+        T _means2 = means[chw2];
+        T _means3 = means[chw3];
+        T _means4 = means[chw4];
+        T _mask1 = mask[chw1];
+        T _mask2 = mask[chw2];
+        T _mask3 = mask[chw3];
+        T _mask4 = mask[chw4];
+
+        _means1 = (_y + _means1) * _mask1;
+        _means2 = (_y + _means2) * _mask2;
+        _means3 = (_y + _means3) * _mask3;
+        _means4 = (_y + _means4) * _mask4;
+        out[chw1] = _means1;
+        out[chw2] = _means2;
+        out[chw3] = _means3;
+        out[chw4] = _means4;
+    }
+}
+
+template <typename scalar_t, typename vec_t>
+__forceinline__ void restore_y_4x_dispatcher(torch::Tensor& out, const torch::Tensor& y,
+                                             const torch::Tensor& means, const torch::Tensor& mask)
+{
+    auto [blockDim, gridDim, stream, useVec, biasSafe, N, HW] = get_kernel_launch_info<vec_t>(y);
+    if (useVec) {
+        restore_y_4x_kernel<vec_t><<<gridDim, blockDim, 0, stream>>>(out, y, means, mask, N);
+    } else {
+        restore_y_4x_kernel<scalar_t><<<gridDim, blockDim, 0, stream>>>(out, y, means, mask, N);
+    }
+}
+
+void restore_y_4x_cuda(torch::Tensor& out, const torch::Tensor& y, const torch::Tensor& means,
+                       const torch::Tensor& mask)
+{
+    if (y.dtype() == torch::kFloat32) {
+        restore_y_4x_dispatcher<float, float4>(out, y, means, mask);
+    } else if (y.dtype() == torch::kFloat16) {
+        restore_y_4x_dispatcher<c10::Half, Half4>(out, y, means, mask);
+    }
+}
+
+template <typename T, typename scalar_t>
+__forceinline__ __device__ T scale_to_index(T scale, const scalar_t scale_min,
+                                            const scalar_t scale_max, const scalar_t log_scale_min,
+                                            const scalar_t log_step_recip)
+{
+    scale = max(scale, scale_min);
+    scale = min(scale, scale_max);
+    scale = log(scale) - log_scale_min;
+    scale = scale * log_step_recip;
+    return scale;
+}
+
+template <typename scalar_t, typename in_t, typename out_t, typename cond_out_t, bool with_cond = false>
+__global__ void build_index_dec_kernel(GPUTensor1D<out_t> out, GPUTensor1D<cond_out_t> cond_out,
+                                       const GPUTensor1D<in_t> scales, const scalar_t scale_min,
+                                       const scalar_t scale_max, const scalar_t log_scale_min,
+                                       const scalar_t log_step_recip, const scalar_t skip_thres,
+                                       const int N)
+{
+    const int n = blockIdx.x * blockDim.x + threadIdx.x;
+
+    if (n < N) {
+        in_t _scale = scales[n];
+        in_t _index = scale_to_index(_scale, scale_min, scale_max, log_scale_min, log_step_recip);
+        out[n] = to_uint8(_index);
+        if constexpr (with_cond) {
+            cond_out_t _cond = _scale > skip_thres;
+            cond_out[n] = _cond;
+        }
+    }
+}
+
+template <typename scalar_t, typename vec_t>
+__forceinline__ void
+build_index_dec_dispatcher(torch::Tensor& out, torch::optional<torch::Tensor>& cond_out,
+                           const torch::Tensor& scales, const scalar_t scale_min,
+                           const scalar_t scale_max, const scalar_t log_scale_min,
+                           const scalar_t log_step_recip, const scalar_t skip_thres)
+{
+    auto [blockDim, gridDim, stream, useVec, N] = get_kernel_launch_info_flatten<vec_t>(scales);
+    const bool with_cond = static_cast<float>(skip_thres) > 0.f;
+
+    auto launch_kernel = [&](auto in_v, auto out_v, auto cond_out_v) {
+        using in_t = decltype(in_v);
+        using out_t = decltype(out_v);
+        using cond_out_t = decltype(cond_out_v);
+        if (with_cond) {
+            build_index_dec_kernel<scalar_t, in_t, out_t, cond_out_t, true>
+                <<<gridDim, blockDim, 0, stream>>>(out, cond_out.value(), scales, scale_min, scale_max,
+                                                   log_scale_min, log_step_recip, skip_thres, N);
+        } else {
+            build_index_dec_kernel<scalar_t, in_t, out_t, cond_out_t, false>
+                <<<gridDim, blockDim, 0, stream>>>(out, nullptr, scales, scale_min, scale_max,
+                                                   log_scale_min, log_step_recip, skip_thres, N);
+        }
+    };
+
+    if (useVec) {
+        launch_kernel(vec_t{}, uchar4{}, bool4{});
+    } else {
+        launch_kernel(scalar_t{}, uint8_t{}, bool{});
+    }
+}
+
+void build_index_dec_cuda(torch::Tensor& out, torch::optional<torch::Tensor>& cond_out,
+                          const torch::Tensor& scales, const float scale_min, const float scale_max,
+                          const float log_scale_min, const float log_step_recip, const float skip_thres)
+{
+    if (scales.dtype() == torch::kFloat32) {
+        build_index_dec_dispatcher<float, float4>(out, cond_out, scales, scale_min, scale_max,
+                                                  log_scale_min, log_step_recip, skip_thres);
+    } else if (scales.dtype() == torch::kFloat16) {
+        build_index_dec_dispatcher<c10::Half, Half4>(
+            out, cond_out, scales, static_cast<c10::Half>(scale_min),
+            static_cast<c10::Half>(scale_max), static_cast<c10::Half>(log_scale_min),
+            static_cast<c10::Half>(log_step_recip), static_cast<c10::Half>(skip_thres));
+    }
+}
+
+template <typename scalar_t, typename in_t, typename out_t, typename cond_out_t, bool with_cond = false>
+__global__ void build_index_enc_kernel(GPUTensor1D<out_t> out, GPUTensor1D<cond_out_t> cond_out,
+                                       const GPUTensor1D<in_t> symbols, const GPUTensor1D<in_t> scales,
+                                       const scalar_t scale_min, const scalar_t scale_max,
+                                       const scalar_t log_scale_min, const scalar_t log_step_recip,
+                                       const scalar_t skip_thres, const int N)
+{
+    const int n = blockIdx.x * blockDim.x + threadIdx.x;
+
+    if (n < N) {
+        in_t _scale = scales[n];
+        in_t _symbol = symbols[n];
+        in_t _index = scale_to_index(_scale, scale_min, scale_max, log_scale_min, log_step_recip);
+
+        out[n] = (to_int16(_symbol) << 8) + to_int16(_index);
+        if constexpr (with_cond) {
+            cond_out_t _cond = _scale > skip_thres;
+            cond_out[n] = _cond;
+        }
+    }
+}
+
+template <typename scalar_t, typename vec_t>
+__forceinline__ void build_index_enc_dispatcher(
+    torch::Tensor& out, torch::optional<torch::Tensor>& cond_out, const torch::Tensor& symbols,
+    const torch::Tensor& scales, const scalar_t scale_min, const scalar_t scale_max,
+    const scalar_t log_scale_min, const scalar_t log_step_recip, const scalar_t skip_thres)
+{
+    auto [blockDim, gridDim, stream, useVec, N] = get_kernel_launch_info_flatten<vec_t>(scales);
+    const bool with_cond = static_cast<float>(skip_thres) > 0.f;
+
+    auto launch_kernel = [&](auto in_v, auto out_v, auto cond_out_v) {
+        using in_t = decltype(in_v);
+        using out_t = decltype(out_v);
+        using cond_out_t = decltype(cond_out_v);
+        if (with_cond) {
+            build_index_enc_kernel<scalar_t, in_t, out_t, cond_out_t, true>
+                <<<gridDim, blockDim, 0, stream>>>(out, cond_out.value(), symbols, scales,
+                                                   scale_min, scale_max, log_scale_min,
+                                                   log_step_recip, skip_thres, N);
+        } else {
+            build_index_enc_kernel<scalar_t, in_t, out_t, cond_out_t, false>
+                <<<gridDim, blockDim, 0, stream>>>(out, nullptr, symbols, scales, scale_min, scale_max,
+                                                   log_scale_min, log_step_recip, skip_thres, N);
+        }
+    };
+
+    if (useVec) {
+        launch_kernel(vec_t{}, short4{}, bool4{});
+    } else {
+        launch_kernel(scalar_t{}, int16_t{}, bool{});
+    }
+}
+
+void build_index_enc_cuda(torch::Tensor& out, torch::optional<torch::Tensor>& cond_out,
+                          const torch::Tensor& symbols, const torch::Tensor& scales,
+                          const float scale_min, const float scale_max, const float log_scale_min,
+                          const float log_step_recip, const float skip_thres)
+{
+    if (scales.dtype() == torch::kFloat32) {
+        build_index_enc_dispatcher<float, float4>(out, cond_out, symbols, scales, scale_min, scale_max,
+                                                  log_scale_min, log_step_recip, skip_thres);
+    } else if (scales.dtype() == torch::kFloat16) {
+        build_index_enc_dispatcher<c10::Half, Half4>(
+            out, cond_out, symbols, scales, static_cast<c10::Half>(scale_min),
+            static_cast<c10::Half>(scale_max), static_cast<c10::Half>(log_scale_min),
+            static_cast<c10::Half>(log_step_recip), static_cast<c10::Half>(skip_thres));
+    }
+}
+
+template <typename vec_t, typename scalar_t, bool biasSafe = false>
+__forceinline__ __device__ vec_t get_bias(const GPUTensor1D<scalar_t> bias, const int HW, const int chw)
+{
+    vec_t _bias;
+    if constexpr (sizeof(vec_t) / sizeof(scalar_t) == 4) {
+        if constexpr (biasSafe) {
+            scalar_t b = bias[(chw * 4 + 0) / HW];
+            _bias = make_vec4(b, b, b, b);
+        } else {
+            _bias = make_vec4(bias[(chw * 4 + 0) / HW], bias[(chw * 4 + 1) / HW],
+                              bias[(chw * 4 + 2) / HW], bias[(chw * 4 + 3) / HW]);
+        }
+    } else {
+        _bias = bias[(chw) / HW];
+    }
+    return _bias;
+}
+
+template <typename scalar_t, typename vec_t, bool biasSafe>
+__global__ void bias_wsilu_kernel(GPUTensor1D<vec_t> x, const GPUTensor1D<scalar_t> bias,
+                                  const int N, const int HW)
+{
+    const int chw = blockIdx.x * blockDim.x + threadIdx.x;
+
+    if (chw < N) {
+        vec_t _bias = get_bias<vec_t, scalar_t, biasSafe>(bias, HW, chw);
+        vec_t _x = x[chw];
+        _x = _x + _bias;
+        _x = wsilu(_x);
+        x[chw] = _x;
+    }
+}
+
+template <typename scalar_t, typename vec_t>
+__forceinline__ void bias_wsilu_dispatcher(torch::Tensor& x, const torch::Tensor& bias)
+{
+    auto [blockDim, gridDim, stream, useVec, biasSafe, N, HW] = get_kernel_launch_info<vec_t>(x);
+    if (useVec) {
+        if (biasSafe) {
+            bias_wsilu_kernel<scalar_t, vec_t, true><<<gridDim, blockDim, 0, stream>>>(x, bias, N, HW);
+        } else {
+            bias_wsilu_kernel<scalar_t, vec_t, false><<<gridDim, blockDim, 0, stream>>>(x, bias, N, HW);
+        }
+    } else {
+        bias_wsilu_kernel<scalar_t, scalar_t, true><<<gridDim, blockDim, 0, stream>>>(x, bias, N, HW);
+    }
+}
+
+void bias_wsilu_cuda(torch::Tensor& x, const torch::Tensor& bias)
+{
+    if (x.dtype() == torch::kFloat32) {
+        bias_wsilu_dispatcher<float, float4>(x, bias);
+    } else if (x.dtype() == torch::kFloat16) {
+        bias_wsilu_dispatcher<c10::Half, Half4>(x, bias);
+    }
+}
+
+template <typename scalar_t, typename vec_t, bool biasSafe, bool with_shortcut = true, bool with_quant = false>
+__global__ void bias_shortcut_kernel(GPUTensor1D<vec_t> x, const GPUTensor1D<scalar_t> bias,
+                                     const GPUTensor1D<scalar_t> quant_step,
+                                     const GPUTensor1D<vec_t> shortcut, const int N, const int HW)
+{
+    const int chw = blockIdx.x * blockDim.x + threadIdx.x;
+
+    if (chw < N) {
+        vec_t _x = x[chw];
+        vec_t _bias = get_bias<vec_t, scalar_t, biasSafe>(bias, HW, chw);
+        _x = _x + _bias;
+        if constexpr (with_shortcut) {
+            vec_t _s = shortcut[chw];
+            _x = _x + _s;
+        }
+        if constexpr (with_quant) {
+            vec_t _q = get_bias<vec_t, scalar_t, biasSafe>(quant_step, HW, chw);
+            _x = _x * _q;
+        }
+        x[chw] = _x;
+    }
+}
+
+template <typename scalar_t, typename vec_t, bool with_shortcut = true, bool with_quant = false>
+__forceinline__ void bias_shortcut_dispatcher(torch::Tensor& x, const torch::Tensor& bias,
+                                              const torch::Tensor& quant_step,
+                                              const torch::Tensor& shortcut)
+{
+    auto [blockDim, gridDim, stream, useVec, biasSafe, N, HW] = get_kernel_launch_info<vec_t>(x);
+    if (useVec) {
+        if (biasSafe) {
+            bias_shortcut_kernel<scalar_t, vec_t, true, with_shortcut, with_quant>
+                <<<gridDim, blockDim, 0, stream>>>(x, bias, quant_step, shortcut, N, HW);
+        } else {
+            bias_shortcut_kernel<scalar_t, vec_t, false, with_shortcut, with_quant>
+                <<<gridDim, blockDim, 0, stream>>>(x, bias, quant_step, shortcut, N, HW);
+        }
+    } else {
+        bias_shortcut_kernel<scalar_t, scalar_t, true, with_shortcut, with_quant>
+            <<<gridDim, blockDim, 0, stream>>>(x, bias, quant_step, shortcut, N, HW);
+    }
+}
+
+void bias_shortcut_cuda(torch::Tensor& x, const torch::Tensor& bias, const torch::Tensor& shortcut)
+{
+    if (x.dtype() == torch::kFloat32) {
+        bias_shortcut_dispatcher<float, float4, true, false>(x, bias, bias, shortcut);
+    } else if (x.dtype() == torch::kFloat16) {
+        bias_shortcut_dispatcher<c10::Half, Half4, true, false>(x, bias, bias, shortcut);
+    }
+}
+
+void bias_quant_cuda(torch::Tensor& x, const torch::Tensor& bias, const torch::Tensor& quant_step)
+{
+    if (x.dtype() == torch::kFloat32) {
+        bias_shortcut_dispatcher<float, float4, false, true>(x, bias, quant_step, bias);
+    } else if (x.dtype() == torch::kFloat16) {
+        bias_shortcut_dispatcher<c10::Half, Half4, false, true>(x, bias, quant_step, bias);
+    }
+}
+
+template <typename scalar_t, typename vec_t, bool biasSafe>
+__global__ void bias_shortcut_no_inplace_kernel(GPUTensor1D<vec_t> out, const GPUTensor1D<vec_t> x,
+                                                const GPUTensor1D<scalar_t> bias,
+                                                const GPUTensor1D<vec_t> shortcut, const int N,
+                                                const int HW)
+{
+    const int chw = blockIdx.x * blockDim.x + threadIdx.x;
+
+    if (chw < N) {
+        vec_t _x = x[chw];
+        vec_t _bias = get_bias<vec_t, scalar_t, biasSafe>(bias, HW, chw);
+        _x = _x + _bias;
+        vec_t _s = shortcut[chw];
+        _x = _x + _s;
+        out[chw] = _x;
+    }
+}
+
+template <typename scalar_t, typename vec_t>
+__forceinline__ void bias_shortcut_no_inplace_dispatcher(torch::Tensor& out, const torch::Tensor& x,
+                                                         const torch::Tensor& bias,
+                                                         const torch::Tensor& shortcut)
+{
+    auto [blockDim, gridDim, stream, useVec, biasSafe, N, HW] = get_kernel_launch_info<vec_t>(x);
+    if (useVec) {
+        if (biasSafe) {
+            bias_shortcut_no_inplace_kernel<scalar_t, vec_t, true>
+                <<<gridDim, blockDim, 0, stream>>>(out, x, bias, shortcut, N, HW);
+        } else {
+            bias_shortcut_no_inplace_kernel<scalar_t, vec_t, false>
+                <<<gridDim, blockDim, 0, stream>>>(out, x, bias, shortcut, N, HW);
+        }
+    } else {
+        bias_shortcut_no_inplace_kernel<scalar_t, scalar_t, true>
+            <<<gridDim, blockDim, 0, stream>>>(out, x, bias, shortcut, N, HW);
+    }
+}
+void bias_shortcut_no_inplace_cuda(torch::Tensor& out, const torch::Tensor& x,
+                                   const torch::Tensor& bias, const torch::Tensor& shortcut)
+{
+    if (x.dtype() == torch::kFloat32) {
+        bias_shortcut_no_inplace_dispatcher<float, float4>(out, x, bias, shortcut);
+    } else if (x.dtype() == torch::kFloat16) {
+        bias_shortcut_no_inplace_dispatcher<c10::Half, Half4>(out, x, bias, shortcut);
+    }
+}
+
+template <typename scalar_t, typename vec_t, bool biasSafe>
+__global__ void bias_shortcut_2_kernel(GPUTensor1D<vec_t> x, const GPUTensor1D<scalar_t> bias,
+                                       GPUTensor1D<vec_t> shortcut, const int N, const int HW)
+{
+    const int chw = blockIdx.x * blockDim.x + threadIdx.x;
+
+    if (chw < N) {
+        vec_t _x = x[chw];
+        vec_t _bias = get_bias<vec_t, scalar_t, biasSafe>(bias, HW, chw);
+        _x = _x + _bias;
+        vec_t _s = shortcut[chw];
+        _x = _x + _s;
+        x[chw] = _x;
+        shortcut[chw] = _x + _s;
+    }
+}
+
+template <typename scalar_t, typename vec_t>
+__forceinline__ void bias_shortcut_2_dispatcher(torch::Tensor& x, const torch::Tensor& bias,
+                                                torch::Tensor& shortcut)
+{
+    auto [blockDim, gridDim, stream, useVec, biasSafe, N, HW] = get_kernel_launch_info<vec_t>(x);
+    if (useVec) {
+        if (biasSafe) {
+            bias_shortcut_2_kernel<scalar_t, vec_t, true>
+                <<<gridDim, blockDim, 0, stream>>>(x, bias, shortcut, N, HW);
+        } else {
+            bias_shortcut_2_kernel<scalar_t, vec_t, false>
+                <<<gridDim, blockDim, 0, stream>>>(x, bias, shortcut, N, HW);
+        }
+    } else {
+        bias_shortcut_2_kernel<scalar_t, scalar_t, true>
+            <<<gridDim, blockDim, 0, stream>>>(x, bias, shortcut, N, HW);
+    }
+}
+
+void bias_shortcut_2_cuda(torch::Tensor& x, const torch::Tensor& bias, torch::Tensor& shortcut)
+{
+    if (x.dtype() == torch::kFloat32) {
+        bias_shortcut_2_dispatcher<float, float4>(x, bias, shortcut);
+    } else if (x.dtype() == torch::kFloat16) {
+        bias_shortcut_2_dispatcher<c10::Half, Half4>(x, bias, shortcut);
+    }
+}
+
+void bias_shortcut_with_quant_step_cuda(torch::Tensor& x, const torch::Tensor& bias,
+                                        const torch::Tensor& quant_step, const torch::Tensor& shortcut)
+{
+    if (x.dtype() == torch::kFloat32) {
+        bias_shortcut_dispatcher<float, float4, true, true>(x, bias, quant_step, shortcut);
+    } else if (x.dtype() == torch::kFloat16) {
+        bias_shortcut_dispatcher<c10::Half, Half4, true, true>(x, bias, quant_step, shortcut);
+    }
+}
+
+template <typename scalar_t, typename vec_t, bool biasSafe>
+__global__ void bias_wsilu_chunk_add_kernel(GPUTensor1D<vec_t> x, const GPUTensor1D<scalar_t> bias,
+                                            const int N, const int HW)
+{
+    const int chw1 = blockIdx.x * blockDim.x + threadIdx.x;
+    const int chw2 = chw1 + N;
+
+    if (chw1 < N) {
+        vec_t _x1 = x[chw1];
+        vec_t _bias1 = get_bias<vec_t, scalar_t, biasSafe>(bias, HW, chw1);
+        _x1 = _x1 + _bias1;
+        _x1 = wsilu(_x1);
+
+        vec_t _x2 = x[chw2];
+        vec_t _bias2 = get_bias<vec_t, scalar_t, biasSafe>(bias, HW, chw2);
+        _x2 = _x2 + _bias2;
+        _x2 = wsilu(_x2);
+
+        x[chw1] = _x1 + _x2;
+    }
+}
+
+template <typename scalar_t, typename vec_t>
+__forceinline__ void bias_wsilu_chunk_add_dispatcher(torch::Tensor& x, const torch::Tensor& bias)
+{
+    auto [blockDim, gridDim, stream, useVec, biasSafe, N, HW] = get_kernel_launch_info<vec_t>(x, 2);
+    if (useVec) {
+        if (biasSafe) {
+            bias_wsilu_chunk_add_kernel<scalar_t, vec_t, true>
+                <<<gridDim, blockDim, 0, stream>>>(x, bias, N, HW);
+        } else {
+            bias_wsilu_chunk_add_kernel<scalar_t, vec_t, false>
+                <<<gridDim, blockDim, 0, stream>>>(x, bias, N, HW);
+        }
+    } else {
+        bias_wsilu_chunk_add_kernel<scalar_t, scalar_t, true>
+            <<<gridDim, blockDim, 0, stream>>>(x, bias, N, HW);
+    }
+}
+
+void bias_wsilu_chunk_add_cuda(torch::Tensor& x, const torch::Tensor& bias)
+{
+    if (x.dtype() == torch::kFloat32) {
+        bias_wsilu_chunk_add_dispatcher<float, float4>(x, bias);
+    } else if (x.dtype() == torch::kFloat16) {
+        bias_wsilu_chunk_add_dispatcher<c10::Half, Half4>(x, bias);
+    }
+    const torch::IntArrayRef x_shape = x.sizes();
+    x = x.narrow(1, 0, x_shape[1] / 2);
+}
+
+template <typename scalar_t>
+__global__ void bias_pixel_shuffle_2_kernel(Packed4DTensorAccessor32<scalar_t> out,
+                                            const Packed4DTensorAccessor32<scalar_t> x,
+                                            const Packed1DTensorAccessor32<scalar_t> bias,
+                                            const int N, const int W)
+{
+    const int c = blockIdx.y * 4;
+    const int c1 = c / 4;
+    const int hw = blockIdx.x * blockDim.x + threadIdx.x;
+    const int h = (hw / W) * 2;
+    const int w = (hw % W) * 2;
+
+    __shared__ scalar_t _bias[4];
+    if (threadIdx.x < 4) {
+        _bias[threadIdx.x] = bias[c + threadIdx.x];
+    }
+    __syncthreads();
+
+    if (hw < N) {
+        scalar_t _bias_0 = _bias[0];
+        scalar_t _bias_1 = _bias[1];
+        scalar_t _bias_2 = _bias[2];
+        scalar_t _bias_3 = _bias[3];
+
+        scalar_t _x0 = x[0][c + 0][0][hw];
+        scalar_t _x1 = x[0][c + 1][0][hw];
+        scalar_t _x2 = x[0][c + 2][0][hw];
+        scalar_t _x3 = x[0][c + 3][0][hw];
+
+        _x0 = _x0 + _bias_0;
+        _x1 = _x1 + _bias_1;
+        _x2 = _x2 + _bias_2;
+        _x3 = _x3 + _bias_3;
+
+        out[0][c1][h + 0][w + 0] = _x0;
+        out[0][c1][h + 0][w + 1] = _x1;
+        out[0][c1][h + 1][w + 0] = _x2;
+        out[0][c1][h + 1][w + 1] = _x3;
+    }
+}
+
+template <typename scalar_t>
+__forceinline__ void bias_pixel_shuffle_2_dispatcher(torch::Tensor& out, const torch::Tensor& x,
+                                                     const torch::Tensor& bias, const int C,
+                                                     const int N, const int W)
+{
+    const int BLOCK_SIZE = 128;
+    const dim3 gridDim((N + BLOCK_SIZE - 1) / BLOCK_SIZE, C / 4);
+    const dim3 blockDim(BLOCK_SIZE);
+    auto stream = at::cuda::getCurrentCUDAStream();
+    bias_pixel_shuffle_2_kernel<scalar_t><<<gridDim, blockDim, 0, stream>>>(
+        out.packed_accessor32<scalar_t, 4, torch::RestrictPtrTraits>(),
+        x.packed_accessor32<scalar_t, 4, torch::RestrictPtrTraits>(),
+        bias.packed_accessor32<scalar_t, 1, torch::RestrictPtrTraits>(), N, W);
+}
+
+void bias_pixel_shuffle_2_cuda(torch::Tensor& out, const torch::Tensor& x,
+                               const torch::Tensor& bias, const int C, const int N, const int W)
+{
+    if (x.dtype() == torch::kFloat32) {
+        bias_pixel_shuffle_2_dispatcher<float>(out, x, bias, C, N, W);
+    } else if (x.dtype() == torch::kFloat16) {
+        bias_pixel_shuffle_2_dispatcher<at::Half>(out, x, bias, C, N, W);
+    }
+}
+
+template <typename scalar_t, bool clamp = false>
+__global__ void bias_pixel_shuffle_8_kernel(Packed4DTensorAccessor32<scalar_t> out,
+                                            const Packed4DTensorAccessor32<scalar_t> x,
+                                            const Packed1DTensorAccessor32<scalar_t> bias,
+                                            const int N, const int W)
+{
+    const int c = blockIdx.y * 64;
+    const int c1 = c / 64;
+    const int hw = blockIdx.x * blockDim.x + threadIdx.x;
+    const int h = (hw / W) * 8;
+    const int w = (hw % W) * 8;
+
+    __shared__ scalar_t _bias[64];
+    if (threadIdx.x < 64) {
+        _bias[threadIdx.x] = bias[c + threadIdx.x];
+    }
+    __syncthreads();
+
+    if (hw < N) {
+        for (int i = 0; i < 64; i++) {
+            scalar_t _x = x[0][c + i][0][hw];
+            _x = _x + _bias[i];
+            const int out_y_offset = i >> 3;
+            const int out_x_offset = i & 7;
+            if constexpr (clamp) {
+                _x = max(_x, static_cast<scalar_t>(0.f));
+                _x = min(_x, static_cast<scalar_t>(1.f));
+            }
+            out[0][c1][h + out_y_offset][w + out_x_offset] = _x;
+        }
+    }
+}
+
+template <typename scalar_t>
+__forceinline__ void bias_pixel_shuffle_8_dispatcher(torch::Tensor& out, const torch::Tensor& x,
+                                                     const torch::Tensor& bias, const int C,
+                                                     const int N, const int W, bool clamp)
+{
+    const int BLOCK_SIZE = 128;
+    const dim3 gridDim((N + BLOCK_SIZE - 1) / BLOCK_SIZE, C / 64);
+    const dim3 blockDim(BLOCK_SIZE);
+    auto stream = at::cuda::getCurrentCUDAStream();
+    if (clamp) {
+        bias_pixel_shuffle_8_kernel<scalar_t, true><<<gridDim, blockDim, 0, stream>>>(
+            out.packed_accessor32<scalar_t, 4, torch::RestrictPtrTraits>(),
+            x.packed_accessor32<scalar_t, 4, torch::RestrictPtrTraits>(),
+            bias.packed_accessor32<scalar_t, 1, torch::RestrictPtrTraits>(), N, W);
+    } else {
+        bias_pixel_shuffle_8_kernel<scalar_t, false><<<gridDim, blockDim, 0, stream>>>(
+            out.packed_accessor32<scalar_t, 4, torch::RestrictPtrTraits>(),
+            x.packed_accessor32<scalar_t, 4, torch::RestrictPtrTraits>(),
+            bias.packed_accessor32<scalar_t, 1, torch::RestrictPtrTraits>(), N, W);
+    }
+}
+
+void bias_pixel_shuffle_8_cuda(torch::Tensor& out, const torch::Tensor& x, const torch::Tensor& bias,
+                               const int C, const int N, const int W, bool clamp)
+{
+    if (x.dtype() == torch::kFloat32) {
+        bias_pixel_shuffle_8_dispatcher<float>(out, x, bias, C, N, W, clamp);
+    } else if (x.dtype() == torch::kFloat16) {
+        bias_pixel_shuffle_8_dispatcher<at::Half>(out, x, bias, C, N, W, clamp);
+    }
+}
+
+template <typename scalar_t, typename vec_t1, typename vec_t2>
+__global__ void round_and_to_int8_kernel(GPUTensor1D<vec_t1> z, GPUTensor1D<vec_t2> z_int8, const int N)
+{
+    const int chw = blockIdx.x * blockDim.x + threadIdx.x;
+
+    if (chw < N) {
+        vec_t1 _z = z[chw];
+        _z = round(_z);
+        _z = max(_z, static_cast<scalar_t>(-128.f));
+        _z = min(_z, static_cast<scalar_t>(127.f));
+        z[chw] = _z;
+        vec_t2 _z_int8 = to_int8(_z);
+        z_int8[chw] = _z_int8;
+    }
+}
+
+template <typename scalar_t, typename vec_t>
+__forceinline__ void round_and_to_int8_dispatcher(torch::Tensor& z, torch::Tensor& z_int8)
+{
+    auto [blockDim, gridDim, stream, useVec, biasSafe, N, HW] = get_kernel_launch_info<vec_t>(z);
+    if (useVec) {
+        round_and_to_int8_kernel<scalar_t, vec_t, char4>
+            <<<gridDim, blockDim, 0, stream>>>(z, z_int8, N);
+    } else {
+        round_and_to_int8_kernel<scalar_t, scalar_t, int8_t>
+            <<<gridDim, blockDim, 0, stream>>>(z, z_int8, N);
+    }
+}
+
+torch::Tensor round_and_to_int8_cuda(torch::Tensor& z)
+{
+    auto z_int8 = torch::empty_like(z, at::TensorOptions().dtype(torch::kInt8));
+    if (z.dtype() == torch::kFloat32) {
+        round_and_to_int8_dispatcher<float, float4>(z, z_int8);
+    } else if (z.dtype() == torch::kFloat16) {
+        round_and_to_int8_dispatcher<c10::Half, Half4>(z, z_int8);
+    }
+    return z_int8;
+}
+
+template <typename scalar_t, typename vec_t>
+__global__ void clamp_reciprocal_with_quant_kernel(GPUTensor1D<vec_t> q_dec_clamp,
+                                                   const GPUTensor1D<vec_t> q_dec, GPUTensor1D<vec_t> y,
+                                                   const scalar_t min_val, const int N)
+{
+    const int chw = blockIdx.x * blockDim.x + threadIdx.x;
+
+    if (chw < N) {
+        vec_t _q_dec = q_dec[chw];
+        vec_t _y = y[chw];
+        _q_dec = max(_q_dec, min_val);
+        q_dec_clamp[chw] = _q_dec;
+        vec_t _q_enc = reciprocal(_q_dec);
+        _y = _y * _q_enc;
+        y[chw] = _y;
+    }
+}
+
+template <typename scalar_t, typename vec_t>
+__forceinline__ void clamp_reciprocal_with_quant_dispatcher(torch::Tensor& q_dec_clamp,
+                                                            const torch::Tensor& q_dec,
+                                                            torch::Tensor& y, const float min_val)
+{
+    auto [blockDim, gridDim, stream, useVec, biasSafe, N, HW] = get_kernel_launch_info<vec_t>(q_dec);
+    if (useVec) {
+        clamp_reciprocal_with_quant_kernel<scalar_t, vec_t><<<gridDim, blockDim, 0, stream>>>(
+            q_dec_clamp, q_dec, y, static_cast<scalar_t>(min_val), N);
+    } else {
+        clamp_reciprocal_with_quant_kernel<scalar_t, scalar_t><<<gridDim, blockDim, 0, stream>>>(
+            q_dec_clamp, q_dec, y, static_cast<scalar_t>(min_val), N);
+    }
+}
+
+torch::Tensor clamp_reciprocal_with_quant_cuda(const torch::Tensor& q_dec, torch::Tensor& y,
+                                               const float min_val)
+{
+    auto q_dec_clamp = torch::empty_like(q_dec);
+    if (q_dec.dtype() == torch::kFloat32) {
+        clamp_reciprocal_with_quant_dispatcher<float, float4>(q_dec_clamp, q_dec, y, min_val);
+    } else if (q_dec.dtype() == torch::kFloat16) {
+        clamp_reciprocal_with_quant_dispatcher<c10::Half, Half4>(q_dec_clamp, q_dec, y, min_val);
+    }
+    return q_dec_clamp;
+}
+
+template <typename T>
+__global__ void add_and_multiply_kernel(GPUTensor1D<T> x0, const GPUTensor1D<T> x1,
+                                        const GPUTensor1D<T> q, const int N)
+{
+    const int chw = blockIdx.x * blockDim.x + threadIdx.x;
+
+    if (chw < N) {
+        T _x0 = x0[chw];
+        T _x1 = x1[chw];
+        T _q = q[chw];
+        _x0 = _x0 + _x1;
+        _x0 = _x0 * _q;
+        x0[chw] = _x0;
+    }
+}
+
+template <typename scalar_t, typename vec_t>
+__forceinline__ void add_and_multiply_dispatcher(torch::Tensor& x0, const torch::Tensor& x1,
+                                                 const torch::Tensor& q)
+{
+    auto [blockDim, gridDim, stream, useVec, biasSafe, N, HW] = get_kernel_launch_info<vec_t>(x0);
+    if (useVec) {
+        add_and_multiply_kernel<vec_t><<<gridDim, blockDim, 0, stream>>>(x0, x1, q, N);
+    } else {
+        add_and_multiply_kernel<scalar_t><<<gridDim, blockDim, 0, stream>>>(x0, x1, q, N);
+    }
+}
+
+void add_and_multiply_cuda(torch::Tensor& x0, const torch::Tensor& x1, const torch::Tensor q)
+{
+    if (x0.dtype() == torch::kFloat32) {
+        add_and_multiply_dispatcher<float, float4>(x0, x1, q);
+    } else if (x0.dtype() == torch::kFloat16) {
+        add_and_multiply_dispatcher<c10::Half, Half4>(x0, x1, q);
+    }
+}
+
+template <typename scalar_t>
+__global__ void replicate_pad_kernel(Packed4DTensorAccessor32<scalar_t> out,
+                                     const Packed4DTensorAccessor32<scalar_t> x, const int C,
+                                     const int H, const int W, const int H_padded, const int W_padded)
+{
+    const int b = blockIdx.y;
+    const int n = blockIdx.x * blockDim.x + threadIdx.x;
+
+    if (n < H_padded * W_padded) {
+        const int dst_y = n / W_padded;
+        const int dst_x = n % W_padded;
+        const int src_y = min(dst_y, H - 1);
+        const int src_x = min(dst_x, W - 1);
+        for (int i = 0; i < C; i++) {
+            scalar_t _x = x[b][i][src_y][src_x];
+            out[b][i][dst_y][dst_x] = _x;
+        }
+    }
+}
+
+template <typename scalar_t>
+__forceinline__ void replicate_pad_dispatcher(torch::Tensor& out, const torch::Tensor& x,
+                                              const int B, const int C, const int H, const int W,
+                                              const int padB, const int padR)
+{
+    const int totalOutPixel = (H + padB) * (W + padR);
+    const int BLOCK_SIZE = 128;
+    const dim3 blockDim(BLOCK_SIZE);
+    const dim3 gridDim((totalOutPixel + BLOCK_SIZE - 1) / BLOCK_SIZE, B);
+    auto stream = at::cuda::getCurrentCUDAStream();
+
+    replicate_pad_kernel<scalar_t><<<gridDim, blockDim, 0, stream>>>(
+        out.packed_accessor32<scalar_t, 4, torch::RestrictPtrTraits>(),
+        x.packed_accessor32<scalar_t, 4, torch::RestrictPtrTraits>(), C, H, W, H + padB, W + padR);
+}
+
+torch::Tensor replicate_pad_cuda(const torch::Tensor& x, const int padB, const int padR)
+{
+    const torch::IntArrayRef x_shape = x.sizes();
+    const int B = x_shape[0];
+    const int C = x_shape[1];
+    const int H = x_shape[2];
+    const int W = x_shape[3];
+    auto out = torch::empty({ B, C, H + padB, W + padR }, x.options());
+    if (x.dtype() == torch::kFloat32) {
+        replicate_pad_dispatcher<float>(out, x, B, C, H, W, padB, padR);
+    } else if (x.dtype() == torch::kFloat16) {
+        replicate_pad_dispatcher<c10::Half>(out, x, B, C, H, W, padB, padR);
+    } else if (x.dtype() == torch::kInt8) {
+        replicate_pad_dispatcher<int8_t>(out, x, B, C, H, W, padB, padR);
+    } else if (x.dtype() == torch::kInt16) {
+        replicate_pad_dispatcher<int16_t>(out, x, B, C, H, W, padB, padR);
+    }
+    return out;
+}
+
+template <typename T, typename T1, int BLOCK_SIZE, int THREAD_NUM_X, int THREAD_NUM_Y>
+__global__ void bias_wsilu_depthwise_conv2d_kernel(Packed4DTensorAccessor32<T> out,
+                                                   const Packed4DTensorAccessor32<T> x,
+                                                   const Packed4DTensorAccessor32<T> weight,
+                                                   const Packed1DTensorAccessor32<T> bias,
+                                                   const int B, const int C, const int H, const int W)
+{
+    const int b = blockIdx.z / C;
+    const int c = blockIdx.z % C;
+    const int h = blockIdx.y * BLOCK_SIZE;  // start of the block
+    const int w = blockIdx.x * BLOCK_SIZE;
+    const int THREAD_NUM = THREAD_NUM_Y * THREAD_NUM_X;
+    const int t_idx = threadIdx.y * THREAD_NUM_X + threadIdx.x;
+
+    __shared__ T1 x_shared[BLOCK_SIZE + 2][BLOCK_SIZE + 2];
+    const T1 __bias = static_cast<T1>(bias[c]);
+    T1 __weight[3][3];
+#pragma unroll
+    for (int i = 0; i < 3; i++) {
+#pragma unroll
+        for (int j = 0; j < 3; j++) {
+            __weight[i][j] = static_cast<T1>(weight[c][0][i][j]);
+        }
+    }
+
+    // load boundary padded pixels
+    const int read_times = (BLOCK_SIZE * 4 + THREAD_NUM - 1) / THREAD_NUM;
+    const int boundary_pos = BLOCK_SIZE + 1;
+
+    for (int i = 0; i < read_times; i++) {
+        int pixel_idx = i * THREAD_NUM + t_idx;
+        if (pixel_idx < BLOCK_SIZE * 2) {
+            const int y_offset = pixel_idx / 2 + 1;
+            const int x_offset = (pixel_idx & 1) * boundary_pos;
+            const int curr_y = h + y_offset - 1;
+            const int curr_x = w + x_offset - 1;
+            if (curr_y < 0 || curr_x < 0 || curr_y >= H || curr_x >= W) {
+                x_shared[y_offset][x_offset] = static_cast<T1>(0.f);
+            } else {
+                T1 x_tmp = static_cast<T1>(x[b][c][curr_y][curr_x]);
+                x_shared[y_offset][x_offset] = wsilu(x_tmp + __bias);
+            }
+        } else if (pixel_idx < BLOCK_SIZE * 4) {
+            pixel_idx -= BLOCK_SIZE * 2;
+            const int y_offset = (pixel_idx & 1) * boundary_pos;
+            const int x_offset = pixel_idx / 2 + 1;
+            const int curr_y = h + y_offset - 1;
+            const int curr_x = w + x_offset - 1;
+            if (curr_y < 0 || curr_x < 0 || curr_y >= H || curr_x >= W) {
+                x_shared[y_offset][x_offset] = static_cast<T1>(0.f);
+            } else {
+                T1 x_tmp = static_cast<T1>(x[b][c][curr_y][curr_x]);
+                x_shared[y_offset][x_offset] = wsilu(x_tmp + __bias);
+            }
+        }
+    }
+
+    // load corner 4 pixels
+    if (t_idx < 4) {
+        const int y_offset = (t_idx / 2) * boundary_pos;
+        const int x_offset = (t_idx & 1) * boundary_pos;
+        const int curr_y = h + y_offset - 1;
+        const int curr_x = w + x_offset - 1;
+        if (curr_y < 0 || curr_x < 0 || curr_y >= H || curr_x >= W) {
+            x_shared[y_offset][x_offset] = static_cast<T1>(0.f);
+        } else {
+            T1 x_tmp = static_cast<T1>(x[b][c][curr_y][curr_x]);
+            x_shared[y_offset][x_offset] = wsilu(x_tmp + __bias);
+        }
+    }
+
+    const int per_y_thread_pix_num = BLOCK_SIZE / THREAD_NUM_Y;
+    const int per_x_thread_pix_num = BLOCK_SIZE / THREAD_NUM_X;
+
+    for (int t_y = 0; t_y < per_y_thread_pix_num; t_y++) {
+        for (int t_x = 0; t_x < per_x_thread_pix_num; t_x++) {
+            const int h_offset = threadIdx.y * per_y_thread_pix_num + t_y + 1;
+            const int w_offset = threadIdx.x * per_x_thread_pix_num + t_x + 1;
+            const int curr_y = h + h_offset - 1;
+            const int curr_x = w + w_offset - 1;
+            // curr_x and curr_y cannot < 0
+            if (curr_y >= H || curr_x >= W) {
+                x_shared[h_offset][w_offset] = static_cast<T1>(0.f);
+            } else {
+                T1 x_tmp = static_cast<T1>(x[b][c][curr_y][curr_x]);
+                x_shared[h_offset][w_offset] = wsilu(x_tmp + __bias);
+            }
+        }
+    }
+    __syncthreads();
+
+    // calculation
+    for (int t_y = 0; t_y < per_y_thread_pix_num; t_y++) {
+        for (int t_x = 0; t_x < per_x_thread_pix_num; t_x++) {
+            const int h_offset = threadIdx.y * per_y_thread_pix_num + t_y;
+            const int w_offset = threadIdx.x * per_x_thread_pix_num + t_x;
+            if (h + h_offset < H && w + w_offset < W) {
+                T1 r = static_cast<T1>(0.f);
+#pragma unroll
+                for (int i = 0; i < 3; i++) {
+#pragma unroll
+                    for (int j = 0; j < 3; j++) {
+                        r = multiply_add(__weight[i][j], x_shared[h_offset + i][w_offset + j], r);
+                    }
+                }
+
+                out[b][c][h + h_offset][w + w_offset] = static_cast<T>(r);
+            }
+        }
+    }
+}
+
+torch::Tensor bias_wsilu_depthwise_conv2d_cuda(const torch::Tensor& x, const torch::Tensor& weight,
+                                               const torch::Tensor& bias)
+{
+    const torch::IntArrayRef x_shape = x.sizes();
+    const int B = x_shape[0];
+    const int C = x_shape[1];
+    const int H = x_shape[2];
+    const int W = x_shape[3];
+
+    auto out = torch::empty_like(x);
+
+    const int BLOCK_SIZE = 32;
+    const int THREAD_NUM_X = 16;
+    const int THREAD_NUM_Y = 8;
+    const dim3 gridDim((W + BLOCK_SIZE - 1) / BLOCK_SIZE, (H + BLOCK_SIZE - 1) / BLOCK_SIZE, B * C);
+    const dim3 blockDim(THREAD_NUM_X, THREAD_NUM_Y);
+    auto stream = at::cuda::getCurrentCUDAStream();
+    if (x.dtype() == torch::kFloat32) {
+        bias_wsilu_depthwise_conv2d_kernel<float, float, BLOCK_SIZE, THREAD_NUM_X, THREAD_NUM_Y>
+            <<<gridDim, blockDim, 0, stream>>>(
+                out.packed_accessor32<float, 4, torch::RestrictPtrTraits>(),
+                x.packed_accessor32<float, 4, torch::RestrictPtrTraits>(),
+                weight.packed_accessor32<float, 4, torch::RestrictPtrTraits>(),
+                bias.packed_accessor32<float, 1, torch::RestrictPtrTraits>(), B, C, H, W);
+    } else if (x.dtype() == torch::kFloat16) {
+        bias_wsilu_depthwise_conv2d_kernel<c10::Half, float, BLOCK_SIZE, THREAD_NUM_X, THREAD_NUM_Y>
+            <<<gridDim, blockDim, 0, stream>>>(
+                out.packed_accessor32<c10::Half, 4, torch::RestrictPtrTraits>(),
+                x.packed_accessor32<c10::Half, 4, torch::RestrictPtrTraits>(),
+                weight.packed_accessor32<c10::Half, 4, torch::RestrictPtrTraits>(),
+                bias.packed_accessor32<c10::Half, 1, torch::RestrictPtrTraits>(), B, C, H, W);
+    }
+    return out;
+}

neural_codec/DCVC/src/layers/extensions/inference/setup.py

@@ -0,0 +1,39 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+
+import os
+import glob
+import sys
+from setuptools import setup
+from torch.utils.cpp_extension import BuildExtension, CUDAExtension
+
+
+cxx_flags = ["-O3"]
+nvcc_flags = ["-O3", "--use_fast_math", "--extra-device-vectorization"]
+# `-arch=native` detects the building machine's GPU at compile time, which is
+# convenient interactively but (a) fails when no GPU is visible during the build
+# and (b) can conflict with the `-gencode` flags torch appends from
+# TORCH_CUDA_ARCH_LIST. So only use `native` when the caller has NOT pinned the
+# arch via TORCH_CUDA_ARCH_LIST (e.g. cluster builds set it to "8.0").
+if not os.environ.get("TORCH_CUDA_ARCH_LIST"):
+    nvcc_flags.append("-arch=native")
+if sys.platform == 'win32':
+    cxx_flags = ["/O2"]
+
+
+setup(
+    name='inference_extensions_cuda',
+    ext_modules=[
+        CUDAExtension(
+            name='inference_extensions_cuda',
+            sources=glob.glob('*.cpp') + glob.glob('*.cu'),
+            extra_compile_args={
+                "cxx": cxx_flags,
+                "nvcc": nvcc_flags,
+            },
+        ),
+    ],
+    cmdclass={
+        'build_ext': BuildExtension
+    }
+)

neural_codec/DCVC/src/layers/layers.py

@@ -0,0 +1,156 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+
+import torch
+from torch import nn
+from .cuda_inference import CUSTOMIZED_CUDA_INFERENCE
+if CUSTOMIZED_CUDA_INFERENCE:
+    from .cuda_inference import DepthConvProxy, SubpelConv2xProxy
+
+
+class WSiLU(nn.Module):
+    def __init__(self):
+        super().__init__()
+
+    def forward(self, x):
+        return torch.sigmoid(4.0 * x) * x
+
+
+class WSiLUChunkAdd(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.silu = WSiLU()
+
+    def forward(self, x):
+        x1, x2 = self.silu(x).chunk(2, 1)
+        return x1 + x2
+
+
+class SubpelConv2x(nn.Module):
+    def __init__(self, in_ch, out_ch, kernel_size, padding=0):
+        super().__init__()
+        self.conv = nn.Sequential(
+            nn.Conv2d(in_ch, out_ch * 4, kernel_size=kernel_size, padding=padding),
+            nn.PixelShuffle(2),
+        )
+        self.padding = padding
+
+        self.proxy = None
+
+    def forward(self, x, to_cat=None, cat_at_front=True):
+        if not CUSTOMIZED_CUDA_INFERENCE or not x.is_cuda:
+            return self.forward_torch(x, to_cat, cat_at_front)
+
+        return self.forward_cuda(x, to_cat, cat_at_front)
+
+    def forward_torch(self, x, to_cat=None, cat_at_front=True):
+        out = self.conv(x)
+        if to_cat is None:
+            return out
+        if cat_at_front:
+            return torch.cat((to_cat, out), dim=1)
+        return torch.cat((out, to_cat), dim=1)
+
+    def forward_cuda(self, x, to_cat=None, cat_at_front=True):
+        if self.proxy is None:
+            self.proxy = SubpelConv2xProxy()
+            self.proxy.set_param(self.conv[0].weight, self.conv[0].bias, self.padding)
+
+        if to_cat is None:
+            return self.proxy.forward(x)
+
+        return self.proxy.forward_with_cat(x, to_cat, cat_at_front)
+
+
+class DepthConvBlock(nn.Module):
+    def __init__(self, in_ch, out_ch, shortcut=False, force_adaptor=False):
+        super().__init__()
+        self.adaptor = None
+        if in_ch != out_ch or force_adaptor:
+            self.adaptor = nn.Conv2d(in_ch, out_ch, 1)
+        self.shortcut = shortcut
+        self.dc = nn.Sequential(
+            nn.Conv2d(out_ch, out_ch, 1),
+            WSiLU(),
+            nn.Conv2d(out_ch, out_ch, 3, padding=1, groups=out_ch),
+            nn.Conv2d(out_ch, out_ch, 1),
+        )
+        self.ffn = nn.Sequential(
+            nn.Conv2d(out_ch, out_ch * 4, 1),
+            WSiLUChunkAdd(),
+            nn.Conv2d(out_ch * 2, out_ch, 1),
+        )
+
+        self.proxy = None
+
+    def forward(self, x, quant_step=None, to_cat=None, cat_at_front=True):
+        if not CUSTOMIZED_CUDA_INFERENCE or not x.is_cuda:
+            return self.forward_torch(x, quant_step, to_cat, cat_at_front)
+
+        return self.forward_cuda(x, quant_step, to_cat, cat_at_front)
+
+    def forward_torch(self, x, quant_step=None, to_cat=None, cat_at_front=True):
+        if self.adaptor is not None:
+            x = self.adaptor(x)
+        out = self.dc(x) + x
+        out = self.ffn(out) + out
+        if self.shortcut:
+            out = out + x
+        if quant_step is not None:
+            out = out * quant_step
+        if to_cat is not None:
+            if cat_at_front:
+                out = torch.cat((to_cat, out), dim=1)
+            else:
+                out = torch.cat((out, to_cat), dim=1)
+        return out
+
+    def forward_cuda(self, x, quant_step=None, to_cat=None, cat_at_front=True):
+        if self.proxy is None:
+            self.proxy = DepthConvProxy()
+            if self.adaptor is not None:
+                self.proxy.set_param_with_adaptor(self.dc[0].weight, self.dc[0].bias,
+                                                  self.dc[2].weight, self.dc[2].bias,
+                                                  self.dc[3].weight, self.dc[3].bias,
+                                                  self.ffn[0].weight, self.ffn[0].bias,
+                                                  self.ffn[2].weight, self.ffn[2].bias,
+                                                  self.adaptor.weight, self.adaptor.bias,
+                                                  self.shortcut)
+            else:
+                self.proxy.set_param(self.dc[0].weight, self.dc[0].bias,
+                                     self.dc[2].weight, self.dc[2].bias,
+                                     self.dc[3].weight, self.dc[3].bias,
+                                     self.ffn[0].weight, self.ffn[0].bias,
+                                     self.ffn[2].weight, self.ffn[2].bias,
+                                     self.shortcut)
+
+        if quant_step is not None:
+            return self.proxy.forward_with_quant_step(x, quant_step)
+        if to_cat is not None:
+            return self.proxy.forward_with_cat(x, to_cat, cat_at_front)
+
+        return self.proxy.forward(x)
+
+
+class ResidualBlockWithStride2(nn.Module):
+    def __init__(self, in_ch, out_ch):
+        super().__init__()
+        self.down = nn.Conv2d(in_ch, out_ch, 2, stride=2)
+        self.conv = DepthConvBlock(out_ch, out_ch, shortcut=True)
+
+    def forward(self, x):
+        x = self.down(x)
+        out = self.conv(x)
+        return out
+
+
+class ResidualBlockUpsample(nn.Module):
+    def __init__(self, in_ch, out_ch):
+        super().__init__()
+        self.up = SubpelConv2x(in_ch, out_ch, 1)
+        self.conv = DepthConvBlock(out_ch, out_ch, shortcut=True)
+
+    def forward(self, x):
+        out = self.up(x)
+        out = self.conv(out)
+        return out

neural_codec/DCVC/src/models/common_model.py

@@ -0,0 +1,296 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+
+import torch
+from torch import nn
+
+from ..layers.cuda_inference import combine_for_reading_2x, \
+    restore_y_2x, restore_y_2x_with_cat_after, add_and_multiply, \
+    replicate_pad, restore_y_4x, clamp_reciprocal_with_quant
+from .entropy_models import BitEstimator, GaussianEncoder, EntropyCoder
+
+
+class CompressionModel(nn.Module):
+    def __init__(self, z_channel, extra_qp=0):
+        super().__init__()
+
+        self.z_channel = z_channel
+        self.entropy_coder = None
+        self.bit_estimator_z = BitEstimator(64 + extra_qp, z_channel)
+        self.gaussian_encoder = GaussianEncoder()
+
+        self.masks = {}
+        self.cuda_streams = {}
+
+    def get_cuda_stream(self, device, idx=0, priority=0):
+        key = f"{device}_{priority}_{idx}"
+        if key not in self.cuda_streams:
+            self.cuda_streams[key] = torch.cuda.Stream(device, priority=priority)
+        return self.cuda_streams[key]
+
+    @staticmethod
+    def get_qp_num():
+        return 64
+
+    @staticmethod
+    def get_padding_size(height, width, p=64):
+        new_h = (height + p - 1) // p * p
+        new_w = (width + p - 1) // p * p
+        padding_right = new_w - width
+        padding_bottom = new_h - height
+        return padding_right, padding_bottom
+
+    @staticmethod
+    def get_downsampled_shape(height, width, p):
+        new_h = (height + p - 1) // p * p
+        new_w = (width + p - 1) // p * p
+        return int(new_h / p + 0.5), int(new_w / p + 0.5)
+
+    def update(self, force_zero_thres=None):
+        self.entropy_coder = EntropyCoder()
+        self.gaussian_encoder.update(self.entropy_coder, force_zero_thres=force_zero_thres)
+        self.bit_estimator_z.update(self.entropy_coder)
+
+    def set_use_two_entropy_coders(self, use_two_entropy_coders):
+        self.entropy_coder.set_use_two_entropy_coders(use_two_entropy_coders)
+
+    def pad_for_y(self, y):
+        _, _, H, W = y.size()
+        padding_r, padding_b = self.get_padding_size(H, W, 4)
+        y_pad = replicate_pad(y, padding_b, padding_r)
+        return y_pad
+
+    def separate_prior(self, params, is_video=False):
+        if is_video:
+            quant_step, scales, means = params.chunk(3, 1)
+            quant_step = torch.clamp_min(quant_step, 0.5)
+            q_enc = 1. / quant_step
+            q_dec = quant_step
+        else:
+            q = params[:, :2, :, :]
+            q_enc, q_dec = (torch.sigmoid(q) * 1.5 + 0.5).chunk(2, 1)
+            scales, means = params[:, 2:, :, :].chunk(2, 1)
+        return q_enc, q_dec, scales, means
+
+    @staticmethod
+    def separate_prior_for_video_encoding(params, y):
+        q_dec, scales, means = params.chunk(3, 1)
+        q_dec, y = clamp_reciprocal_with_quant(q_dec, y, 0.5)
+        return y, q_dec, scales, means
+
+    @staticmethod
+    def separate_prior_for_video_decoding(params):
+        quant_step, scales, means = params.chunk(3, 1)
+        quant_step = torch.clamp_min(quant_step, 0.5)
+        return quant_step, scales, means
+
+    def process_with_mask(self, y, scales, means, mask):
+        return self.gaussian_encoder.process_with_mask(y, scales, means, mask)
+
+    @staticmethod
+    def get_one_mask(micro_mask, height, width, dtype, device):
+        mask = torch.tensor(micro_mask, dtype=dtype, device=device)
+        mask = mask.repeat((height + 1) // 2, (width + 1) // 2)
+        mask = mask[:height, :width]
+        mask = torch.unsqueeze(mask, 0)
+        mask = torch.unsqueeze(mask, 0)
+        return mask
+
+    def get_mask_4x(self, batch, channel, height, width, dtype, device):
+        curr_mask_str = f"{batch}_{channel}_{width}_{height}_4x"
+        with torch.no_grad():
+            if curr_mask_str not in self.masks:
+                assert channel % 4 == 0
+                m = torch.ones((batch, channel // 4, height, width), dtype=dtype, device=device)
+                m0 = self.get_one_mask(((1, 0), (0, 0)), height, width, dtype, device)
+                m1 = self.get_one_mask(((0, 1), (0, 0)), height, width, dtype, device)
+                m2 = self.get_one_mask(((0, 0), (1, 0)), height, width, dtype, device)
+                m3 = self.get_one_mask(((0, 0), (0, 1)), height, width, dtype, device)
+
+                mask_0 = torch.cat((m * m0, m * m1, m * m2, m * m3), dim=1)
+                mask_1 = torch.cat((m * m3, m * m2, m * m1, m * m0), dim=1)
+                mask_2 = torch.cat((m * m2, m * m3, m * m0, m * m1), dim=1)
+                mask_3 = torch.cat((m * m1, m * m0, m * m3, m * m2), dim=1)
+
+                self.masks[curr_mask_str] = [mask_0, mask_1, mask_2, mask_3]
+        return self.masks[curr_mask_str]
+
+    def get_mask_2x(self, batch, channel, height, width, dtype, device):
+        curr_mask_str = f"{batch}_{channel}_{width}_{height}_2x"
+        with torch.no_grad():
+            if curr_mask_str not in self.masks:
+                assert channel % 2 == 0
+                m = torch.ones((batch, channel // 2, height, width), dtype=dtype, device=device)
+                m0 = self.get_one_mask(((1, 0), (0, 1)), height, width, dtype, device)
+                m1 = self.get_one_mask(((0, 1), (1, 0)), height, width, dtype, device)
+
+                mask_0 = torch.cat((m * m0, m * m1), dim=1)
+                mask_1 = torch.cat((m * m1, m * m0), dim=1)
+
+                self.masks[curr_mask_str] = [mask_0, mask_1]
+        return self.masks[curr_mask_str]
+
+    @staticmethod
+    def single_part_for_writing_4x(x):
+        x0, x1, x2, x3 = x.chunk(4, 1)
+        return (x0 + x1) + (x2 + x3)
+
+    @staticmethod
+    def single_part_for_writing_2x(x):
+        x0, x1 = x.chunk(2, 1)
+        return x0 + x1
+
+    def compress_prior_2x(self, y, common_params, y_spatial_prior):
+        y, q_dec, scales, means = self.separate_prior_for_video_encoding(common_params, y)
+        dtype = y.dtype
+        device = y.device
+        B, C, H, W = y.size()
+        mask_0, mask_1 = self.get_mask_2x(B, C, H, W, dtype, device)
+
+        _, y_q_0, y_hat_0, s_hat_0 = self.process_with_mask(y, scales, means, mask_0)
+        cat_params = torch.cat((y_hat_0, common_params), dim=1)
+        scales, means = y_spatial_prior(cat_params).chunk(2, 1)
+        _, y_q_1, y_hat_1, s_hat_1 = self.process_with_mask(y, scales, means, mask_1)
+
+        y_hat = add_and_multiply(y_hat_0, y_hat_1, q_dec)
+
+        y_q_w_0 = self.single_part_for_writing_2x(y_q_0)
+        y_q_w_1 = self.single_part_for_writing_2x(y_q_1)
+        s_w_0 = self.single_part_for_writing_2x(s_hat_0)
+        s_w_1 = self.single_part_for_writing_2x(s_hat_1)
+        return y_q_w_0, y_q_w_1, s_w_0, s_w_1, y_hat
+
+    def decompress_prior_2x(self, common_params, y_spatial_prior):
+        infos = self.decompress_prior_2x_part1(common_params)
+        y_hat = self.decompress_prior_2x_part2(common_params, y_spatial_prior, infos)
+        return y_hat
+
+    def decompress_prior_2x_part1(self, common_params):
+        q_dec, scales, means = self.separate_prior_for_video_decoding(common_params)
+        dtype = means.dtype
+        device = means.device
+        B, C, H, W = means.size()
+        mask_0, mask_1 = self.get_mask_2x(B, C, H, W, dtype, device)
+
+        scales_r = combine_for_reading_2x(scales, mask_0, inplace=False)
+        indexes, skip_cond = self.gaussian_encoder.build_indexes_decoder(scales_r)
+        self.gaussian_encoder.decode_y(indexes)
+        infos = {
+            "q_dec": q_dec,
+            "mask_0": mask_0,
+            "mask_1": mask_1,
+            "means": means,
+            "scales_r": scales_r,
+            "skip_cond": skip_cond,
+            "indexes": indexes,
+        }
+        return infos
+
+    def decompress_prior_2x_part2(self, common_params, y_spatial_prior, infos):
+        dtype = common_params.dtype
+        device = common_params.device
+        y_q_r = self.gaussian_encoder.get_y(infos["scales_r"].shape,
+                                            infos["scales_r"].numel(),
+                                            dtype, device,
+                                            infos["skip_cond"], infos["indexes"])
+        y_hat_0, cat_params = restore_y_2x_with_cat_after(y_q_r, infos["means"], infos["mask_0"],
+                                                          common_params)
+        scales, means = y_spatial_prior(cat_params).chunk(2, 1)
+        scales_r = combine_for_reading_2x(scales, infos["mask_1"], inplace=True)
+        y_q_r = self.gaussian_encoder.decode_and_get_y(scales_r, dtype, device)
+        y_hat_1 = restore_y_2x(y_q_r, means, infos["mask_1"])
+
+        y_hat = add_and_multiply(y_hat_0, y_hat_1, infos["q_dec"])
+        return y_hat
+
+    def compress_prior_4x(self, y, common_params, y_spatial_prior_reduction,
+                          y_spatial_prior_adaptor_1, y_spatial_prior_adaptor_2,
+                          y_spatial_prior_adaptor_3, y_spatial_prior):
+        '''
+        y_0 means split in channel, the 0/4 quater
+        y_1 means split in channel, the 1/4 quater
+        y_2 means split in channel, the 2/4 quater
+        y_3 means split in channel, the 3/4 quater
+        y_?_0, means multiply with mask_0
+        y_?_1, means multiply with mask_1
+        y_?_2, means multiply with mask_2
+        y_?_3, means multiply with mask_3
+        '''
+        q_enc, q_dec, scales, means = self.separate_prior(common_params, False)
+        common_params = y_spatial_prior_reduction(common_params)
+        dtype = y.dtype
+        device = y.device
+        B, C, H, W = y.size()
+        mask_0, mask_1, mask_2, mask_3 = self.get_mask_4x(B, C, H, W, dtype, device)
+
+        y = y * q_enc
+
+        _, y_q_0, y_hat_0, s_hat_0 = self.process_with_mask(y, scales, means, mask_0)
+
+        y_hat_so_far = y_hat_0
+        params = torch.cat((y_hat_so_far, common_params), dim=1)
+        scales, means = y_spatial_prior(y_spatial_prior_adaptor_1(params)).chunk(2, 1)
+        _, y_q_1, y_hat_1, s_hat_1 = self.process_with_mask(y, scales, means, mask_1)
+
+        y_hat_so_far = y_hat_so_far + y_hat_1
+        params = torch.cat((y_hat_so_far, common_params), dim=1)
+        scales, means = y_spatial_prior(y_spatial_prior_adaptor_2(params)).chunk(2, 1)
+        _, y_q_2, y_hat_2, s_hat_2 = self.process_with_mask(y, scales, means, mask_2)
+
+        y_hat_so_far = y_hat_so_far + y_hat_2
+        params = torch.cat((y_hat_so_far, common_params), dim=1)
+        scales, means = y_spatial_prior(y_spatial_prior_adaptor_3(params)).chunk(2, 1)
+        _, y_q_3, y_hat_3, s_hat_3 = self.process_with_mask(y, scales, means, mask_3)
+
+        y_hat = y_hat_so_far + y_hat_3
+        y_hat = y_hat * q_dec
+
+        y_q_w_0 = self.single_part_for_writing_4x(y_q_0)
+        y_q_w_1 = self.single_part_for_writing_4x(y_q_1)
+        y_q_w_2 = self.single_part_for_writing_4x(y_q_2)
+        y_q_w_3 = self.single_part_for_writing_4x(y_q_3)
+        s_w_0 = self.single_part_for_writing_4x(s_hat_0)
+        s_w_1 = self.single_part_for_writing_4x(s_hat_1)
+        s_w_2 = self.single_part_for_writing_4x(s_hat_2)
+        s_w_3 = self.single_part_for_writing_4x(s_hat_3)
+        return y_q_w_0, y_q_w_1, y_q_w_2, y_q_w_3, s_w_0, s_w_1, s_w_2, s_w_3, y_hat
+
+    def decompress_prior_4x(self, common_params, y_spatial_prior_reduction,
+                            y_spatial_prior_adaptor_1, y_spatial_prior_adaptor_2,
+                            y_spatial_prior_adaptor_3, y_spatial_prior):
+        _, quant_step, scales, means = self.separate_prior(common_params, False)
+        common_params = y_spatial_prior_reduction(common_params)
+        dtype = means.dtype
+        device = means.device
+        B, C, H, W = means.size()
+        mask_0, mask_1, mask_2, mask_3 = self.get_mask_4x(B, C, H, W, dtype, device)
+
+        scales_r = self.single_part_for_writing_4x(scales * mask_0)
+        y_q_r = self.gaussian_encoder.decode_and_get_y(scales_r, dtype, device)
+        y_hat_curr_step = restore_y_4x(y_q_r, means, mask_0)
+        y_hat_so_far = y_hat_curr_step
+
+        params = torch.cat((y_hat_so_far, common_params), dim=1)
+        scales, means = y_spatial_prior(y_spatial_prior_adaptor_1(params)).chunk(2, 1)
+        scales_r = self.single_part_for_writing_4x(scales * mask_1)
+        y_q_r = self.gaussian_encoder.decode_and_get_y(scales_r, dtype, device)
+        y_hat_curr_step = restore_y_4x(y_q_r, means, mask_1)
+        y_hat_so_far = y_hat_so_far + y_hat_curr_step
+
+        params = torch.cat((y_hat_so_far, common_params), dim=1)
+        scales, means = y_spatial_prior(y_spatial_prior_adaptor_2(params)).chunk(2, 1)
+        scales_r = self.single_part_for_writing_4x(scales * mask_2)
+        y_q_r = self.gaussian_encoder.decode_and_get_y(scales_r, dtype, device)
+        y_hat_curr_step = restore_y_4x(y_q_r, means, mask_2)
+        y_hat_so_far = y_hat_so_far + y_hat_curr_step
+
+        params = torch.cat((y_hat_so_far, common_params), dim=1)
+        scales, means = y_spatial_prior(y_spatial_prior_adaptor_3(params)).chunk(2, 1)
+        scales_r = self.single_part_for_writing_4x(scales * mask_3)
+        y_q_r = self.gaussian_encoder.decode_and_get_y(scales_r, dtype, device)
+        y_hat_curr_step = restore_y_4x(y_q_r, means, mask_3)
+        y_hat_so_far = y_hat_so_far + y_hat_curr_step
+
+        y_hat = y_hat_so_far * quant_step
+
+        return y_hat

neural_codec/DCVC/src/models/entropy_models.py

@@ -0,0 +1,341 @@
+import math
+
+import torch
+import numpy as np
+from torch import nn
+import torch.nn.functional as F
+
+from ..layers.cuda_inference import build_index_dec, build_index_enc, process_with_mask
+
+
+class EntropyCoder():
+    def __init__(self):
+        super().__init__()
+
+        from MLCodec_extensions_cpp import RansEncoder, RansDecoder
+        self.encoder = RansEncoder()
+        self.decoder = RansDecoder()
+
+    @staticmethod
+    def pmf_to_quantized_cdf(pmf, precision=16):
+        from MLCodec_extensions_cpp import pmf_to_quantized_cdf as _pmf_to_cdf
+        cdf = _pmf_to_cdf(pmf.tolist(), precision)
+        cdf = torch.IntTensor(cdf)
+        return cdf
+
+    @staticmethod
+    def pmf_to_cdf(pmf, tail_mass, pmf_length, max_length):
+        entropy_coder_precision = 16
+        cdf = torch.zeros((len(pmf_length), max_length + 2), dtype=torch.int32)
+        for i, p in enumerate(pmf):
+            prob = torch.cat((p[: pmf_length[i]], tail_mass[i]), dim=0)
+            _cdf = EntropyCoder.pmf_to_quantized_cdf(prob, entropy_coder_precision)
+            cdf[i, : _cdf.size(0)] = _cdf
+        return cdf
+
+    def reset(self):
+        self.encoder.reset()
+
+    def add_cdf(self, cdf, cdf_length, offset):
+        enc_cdf_idx = self.encoder.add_cdf(cdf, cdf_length, offset)
+        dec_cdf_idx = self.decoder.add_cdf(cdf, cdf_length, offset)
+        assert enc_cdf_idx == dec_cdf_idx
+        return enc_cdf_idx
+
+    def encode_y(self, symbols, cdf_group_index):
+        # symbols: int16, high 8 bits: int8 symbol to be encoded; low 8 bits: uint8 index to use
+        assert symbols.dtype == torch.int16
+        self.encoder.encode_y(symbols.cpu().numpy(), cdf_group_index)
+
+    def encode_z(self, symbols, cdf_group_index, start_offset, per_channel_size):
+        self.encoder.encode_z(symbols.to(torch.int8).cpu().numpy(),
+                              cdf_group_index, start_offset, per_channel_size)
+
+    def flush(self):
+        self.encoder.flush()
+
+    def get_encoded_stream(self):
+        return self.encoder.get_encoded_stream().tobytes()
+
+    def set_stream(self, stream):
+        self.decoder.set_stream((np.frombuffer(stream, dtype=np.uint8)))
+
+    def decode_y(self, indexes, cdf_group_index):
+        self.decoder.decode_y(indexes.to(torch.uint8).cpu().numpy(), cdf_group_index)
+
+    def decode_and_get_y(self, indexes, cdf_group_index, device, dtype):
+        rv = self.decoder.decode_and_get_y(indexes.to(torch.uint8).cpu().numpy(), cdf_group_index)
+        rv = torch.as_tensor(rv)
+        return rv.to(device).to(dtype)
+
+    def decode_z(self, total_size, cdf_group_index, start_offset, per_channel_size):
+        self.decoder.decode_z(total_size, cdf_group_index, start_offset, per_channel_size)
+
+    def get_decoded_tensor(self, device, dtype, non_blocking=False):
+        rv = self.decoder.get_decoded_tensor()
+        rv = torch.as_tensor(rv)
+        return rv.to(device, non_blocking=non_blocking).to(dtype)
+
+    def set_use_two_entropy_coders(self, use_two_entropy_coders):
+        self.encoder.set_use_two_encoders(use_two_entropy_coders)
+        self.decoder.set_use_two_decoders(use_two_entropy_coders)
+
+
+class Bitparm(nn.Module):
+    def __init__(self, qp_num, channel, final=False):
+        super().__init__()
+        self.final = final
+        self.h = nn.Parameter(torch.nn.init.normal_(
+            torch.empty([qp_num, channel, 1, 1]), 0, 0.01))
+        self.b = nn.Parameter(torch.nn.init.normal_(
+            torch.empty([qp_num, channel, 1, 1]), 0, 0.01))
+        if not final:
+            self.a = nn.Parameter(torch.nn.init.normal_(
+                torch.empty([qp_num, channel, 1, 1]), 0, 0.01))
+        else:
+            self.a = None
+
+    def forward(self, x, index):
+        h = torch.index_select(self.h, 0, index)
+        b = torch.index_select(self.b, 0, index)
+        x = x * F.softplus(h) + b
+        if self.final:
+            return x
+
+        a = torch.index_select(self.a, 0, index)
+        return x + torch.tanh(x) * torch.tanh(a)
+
+
+class AEHelper():
+    def __init__(self):
+        super().__init__()
+        self.entropy_coder = None
+        self.cdf_group_index = None
+        self._offset = None
+        self._quantized_cdf = None
+        self._cdf_length = None
+
+    def set_cdf_info(self, quantized_cdf, cdf_length, offset):
+        self._quantized_cdf = quantized_cdf.cpu().numpy()
+        self._cdf_length = cdf_length.reshape(-1).int().cpu().numpy()
+        self._offset = offset.reshape(-1).int().cpu().numpy()
+
+    def get_cdf_info(self):
+        return self._quantized_cdf, \
+            self._cdf_length, \
+            self._offset
+
+
+class BitEstimator(AEHelper, nn.Module):
+    def __init__(self, qp_num, channel):
+        super().__init__()
+        self.f1 = Bitparm(qp_num, channel)
+        self.f2 = Bitparm(qp_num, channel)
+        self.f3 = Bitparm(qp_num, channel)
+        self.f4 = Bitparm(qp_num, channel, True)
+        self.qp_num = qp_num
+        self.channel = channel
+
+    def forward(self, x, index):
+        return self.get_cdf(x, index)
+
+    def get_logits_cdf(self, x, index):
+        x = self.f1(x, index)
+        x = self.f2(x, index)
+        x = self.f3(x, index)
+        x = self.f4(x, index)
+        return x
+
+    def get_cdf(self, x, index):
+        return torch.sigmoid(self.get_logits_cdf(x, index))
+
+    def update(self, entropy_coder):
+        self.entropy_coder = entropy_coder
+
+        with torch.no_grad():
+            device = next(self.parameters()).device
+            medians = torch.zeros((self.qp_num, self.channel, 1, 1), device=device)
+            index = torch.arange(self.qp_num, device=device, dtype=torch.int32)
+
+            minima = medians + 8
+            for i in range(8, 1, -1):
+                samples = torch.zeros_like(medians) - i
+                probs = self.forward(samples, index)
+                minima = torch.where(probs < torch.zeros_like(medians) + 0.0001,
+                                     torch.zeros_like(medians) + i, minima)
+
+            maxima = medians + 8
+            for i in range(8, 1, -1):
+                samples = torch.zeros_like(medians) + i
+                probs = self.forward(samples, index)
+                maxima = torch.where(probs > torch.zeros_like(medians) + 0.9999,
+                                     torch.zeros_like(medians) + i, maxima)
+
+            minima = minima.int()
+            maxima = maxima.int()
+
+            offset = -minima
+
+            pmf_start = medians - minima
+            pmf_length = maxima + minima + 1
+
+            max_length = pmf_length.max()
+            device = pmf_start.device
+            samples = torch.arange(max_length, device=device)
+
+            samples = samples[None, None, None, :] + pmf_start
+
+            half = float(0.5)
+
+            lower = self.forward(samples - half, index)
+            upper = self.forward(samples + half, index)
+            pmf = upper - lower
+
+            pmf = pmf[:, :, 0, :]
+            upper = self.forward(maxima.to(torch.float32), index)
+            tail_mass = lower[:, :, 0, :1] + (1.0 - upper[:, :, 0, -1:])
+
+            pmf = pmf.reshape([-1, max_length])
+            tail_mass = tail_mass.reshape([-1, 1])
+            pmf_length = pmf_length.reshape([-1])
+            offset = offset.reshape([-1])
+            quantized_cdf = EntropyCoder.pmf_to_cdf(pmf, tail_mass, pmf_length, max_length)
+            cdf_length = pmf_length + 2
+            self.set_cdf_info(quantized_cdf, cdf_length, offset)
+            self.cdf_group_index = self.entropy_coder.add_cdf(*self.get_cdf_info())
+
+    def build_indexes(self, size, qp):
+        B, C, H, W = size
+        indexes = torch.arange(C, dtype=torch.int).view(1, -1, 1, 1) + qp * self.channel
+        return indexes.repeat(B, 1, H, W)
+
+    def encode_z(self, x, qp):
+        _, _, H, W = x.size()
+        return self.entropy_coder.encode_z(x.reshape(-1), self.cdf_group_index, qp * self.channel,
+                                           H * W)
+
+    def decode_z(self, size, qp):
+        self.entropy_coder.decode_z(self.channel * size[0] * size[1], self.cdf_group_index,
+                                    qp * self.channel, size[0] * size[1])
+
+    def get_z(self, size, device, dtype):
+        output_size = (1, self.channel, size[0], size[1])
+        val = self.entropy_coder.get_decoded_tensor(device, dtype, non_blocking=True)
+        return val.reshape(output_size)
+
+
+class GaussianEncoder(AEHelper):
+    def __init__(self):
+        super().__init__()
+        self.scale_min = 0.11
+        self.scale_max = 16.0
+        self.scale_level = 128  # <= 256
+        self.scale_table = self.get_scale_table(self.scale_min, self.scale_max, self.scale_level)
+
+        self.log_scale_min = math.log(self.scale_min)
+        self.log_scale_max = math.log(self.scale_max)
+        self.log_scale_step = (self.log_scale_max - self.log_scale_min) / (self.scale_level - 1)
+        self.log_step_recip = 1. / self.log_scale_step
+
+        self.force_zero_thres = None
+        self.decode_index_cache = {}
+        self.decode_zeros_cache = {}
+
+    @staticmethod
+    def get_scale_table(min_val, max_val, levels):
+        return torch.exp(torch.linspace(math.log(min_val), math.log(max_val), levels))
+
+    def update(self, entropy_coder, force_zero_thres=None):
+        self.entropy_coder = entropy_coder
+        self.force_zero_thres = force_zero_thres
+
+        pmf_center = torch.zeros_like(self.scale_table) + 8
+        scales = torch.zeros_like(pmf_center) + self.scale_table
+        cdf_distribution = torch.distributions.normal.Normal(0., scales)
+        for i in range(8, 1, -1):
+            samples = torch.zeros_like(pmf_center) + i
+            probs = cdf_distribution.cdf(samples)
+            probs = torch.squeeze(probs)
+            pmf_center = torch.where(probs > torch.zeros_like(pmf_center) + 0.9999,
+                                     torch.zeros_like(pmf_center) + i, pmf_center)
+
+        pmf_center = pmf_center.int()
+        pmf_length = 2 * pmf_center + 1
+        max_length = torch.max(pmf_length).item()
+
+        device = pmf_center.device
+        samples = torch.arange(max_length, device=device) - pmf_center[:, None]
+        samples = samples.float()
+
+        scales = torch.zeros_like(samples) + self.scale_table[:, None]
+        cdf_distribution = torch.distributions.normal.Normal(0., scales)
+
+        upper = cdf_distribution.cdf(samples + 0.5)
+        lower = cdf_distribution.cdf(samples - 0.5)
+        pmf = upper - lower
+
+        tail_mass = 2 * lower[:, :1]
+
+        quantized_cdf = torch.Tensor(len(pmf_length), max_length + 2)
+        quantized_cdf = EntropyCoder.pmf_to_cdf(pmf, tail_mass, pmf_length, max_length)
+
+        self.set_cdf_info(quantized_cdf, pmf_length+2, -pmf_center)
+        self.cdf_group_index = self.entropy_coder.add_cdf(*self.get_cdf_info())
+
+    def process_with_mask(self, y, scales, means, mask):
+        return process_with_mask(y, scales, means, mask, self.force_zero_thres)
+
+    def build_indexes_decoder(self, scales):
+        scales = scales.reshape(-1)
+        indexes, skip_cond = build_index_dec(scales, self.scale_min, self.scale_max,
+                                             self.log_scale_min, self.log_step_recip,
+                                             self.force_zero_thres)
+        if self.force_zero_thres is not None:
+            indexes = indexes[skip_cond]
+        return indexes, skip_cond
+
+    def build_indexes_encoder(self, symbols, scales):
+        symbols = symbols.reshape(-1)
+        scales = scales.reshape(-1)
+        symbols = build_index_enc(symbols, scales, self.scale_min, self.scale_max,
+                                  self.log_scale_min, self.log_step_recip, self.force_zero_thres)
+        return symbols
+
+    def encode_y(self, x, scales):
+        symbols = self.build_indexes_encoder(x, scales)
+        return self.entropy_coder.encode_y(symbols, self.cdf_group_index)
+
+    def get_decode_index_cache(self, num, device):
+        if num not in self.decode_index_cache:
+            c = torch.arange(0, num, dtype=torch.int32, device=device)
+            self.decode_index_cache[num] = c
+
+        return self.decode_index_cache[num]
+
+    def get_decode_zeros_cache(self, num, device):
+        if num not in self.decode_zeros_cache:
+            c = torch.zeros(num, dtype=torch.int32, device=device)
+            self.decode_zeros_cache[num] = c
+
+        return self.decode_zeros_cache[num].clone()
+
+    def decode_and_get_y(self, scales, dtype, device):
+        indexes, skip_cond = self.build_indexes_decoder(scales)
+        self.decode_y(indexes)
+        return self.get_y(scales.shape, scales.numel(), dtype, device, skip_cond, indexes)
+
+    def decode_y(self, indexes):
+        self.entropy_coder.decode_y(indexes, self.cdf_group_index)
+
+    def get_y(self, shape, numel, dtype, device, skip_cond, indexes):
+        if len(indexes) == 0:
+            return torch.zeros(shape, dtype=dtype, device=device)
+        if skip_cond is not None:
+            curr_index = self.get_decode_index_cache(numel, device)
+            back_index = self.get_decode_zeros_cache(numel, device)
+            back_index.masked_scatter_(skip_cond, curr_index)
+        val = self.entropy_coder.get_decoded_tensor(device, dtype, non_blocking=True)
+        if skip_cond is not None:
+            y = torch.index_select(val, 0, back_index) * skip_cond
+            return y.reshape(shape)
+        return val.reshape(shape)

neural_codec/DCVC/src/models/image_model.py

@@ -0,0 +1,209 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+
+import torch
+from torch import nn
+import torch.nn.functional as F
+
+
+from .common_model import CompressionModel
+from ..layers.layers import DepthConvBlock, ResidualBlockUpsample, ResidualBlockWithStride2
+from ..layers.cuda_inference import CUSTOMIZED_CUDA_INFERENCE, round_and_to_int8
+
+g_ch_src = 3 * 8 * 8
+g_ch_enc_dec = 368
+
+
+class IntraEncoder(nn.Module):
+    def __init__(self, N):
+        super().__init__()
+
+        self.enc_1 = DepthConvBlock(g_ch_src, g_ch_enc_dec)
+        self.enc_2 = nn.Sequential(
+            DepthConvBlock(g_ch_enc_dec, g_ch_enc_dec),
+            DepthConvBlock(g_ch_enc_dec, g_ch_enc_dec),
+            DepthConvBlock(g_ch_enc_dec, g_ch_enc_dec),
+            DepthConvBlock(g_ch_enc_dec, g_ch_enc_dec),
+            DepthConvBlock(g_ch_enc_dec, g_ch_enc_dec),
+            DepthConvBlock(g_ch_enc_dec, g_ch_enc_dec),
+            nn.Conv2d(g_ch_enc_dec, N, 3, stride=2, padding=1),
+        )
+
+    def forward(self, x, quant_step):
+        out = F.pixel_unshuffle(x, 8)
+        if not CUSTOMIZED_CUDA_INFERENCE or not x.is_cuda:
+            return self.forward_torch(out, quant_step)
+
+        return self.forward_cuda(out, quant_step)
+
+    def forward_torch(self, out, quant_step):
+        out = self.enc_1(out)
+        out = out * quant_step
+        return self.enc_2(out)
+
+    def forward_cuda(self, out, quant_step):
+        out = self.enc_1(out, quant_step=quant_step)
+        return self.enc_2(out)
+
+
+class IntraDecoder(nn.Module):
+    def __init__(self, N):
+        super().__init__()
+
+        self.dec_1 = nn.Sequential(
+            ResidualBlockUpsample(N, g_ch_enc_dec),
+            DepthConvBlock(g_ch_enc_dec, g_ch_enc_dec),
+            DepthConvBlock(g_ch_enc_dec, g_ch_enc_dec),
+            DepthConvBlock(g_ch_enc_dec, g_ch_enc_dec),
+            DepthConvBlock(g_ch_enc_dec, g_ch_enc_dec),
+            DepthConvBlock(g_ch_enc_dec, g_ch_enc_dec),
+            DepthConvBlock(g_ch_enc_dec, g_ch_enc_dec),
+            DepthConvBlock(g_ch_enc_dec, g_ch_enc_dec),
+            DepthConvBlock(g_ch_enc_dec, g_ch_enc_dec),
+            DepthConvBlock(g_ch_enc_dec, g_ch_enc_dec),
+            DepthConvBlock(g_ch_enc_dec, g_ch_enc_dec),
+            DepthConvBlock(g_ch_enc_dec, g_ch_enc_dec),
+            DepthConvBlock(g_ch_enc_dec, g_ch_enc_dec),
+        )
+        self.dec_2 = DepthConvBlock(g_ch_enc_dec, g_ch_src)
+
+    def forward(self, x, quant_step):
+        if not CUSTOMIZED_CUDA_INFERENCE or not x.is_cuda:
+            return self.forward_torch(x, quant_step)
+
+        return self.forward_cuda(x, quant_step)
+
+    def forward_torch(self, x, quant_step):
+        out = self.dec_1(x)
+        out = out * quant_step
+        out = self.dec_2(out)
+        out = F.pixel_shuffle(out, 8)
+        return out
+
+    def forward_cuda(self, x, quant_step):
+        out = self.dec_1[0](x)
+        out = self.dec_1[1](out)
+        out = self.dec_1[2](out)
+        out = self.dec_1[3](out)
+        out = self.dec_1[4](out)
+        out = self.dec_1[5](out)
+        out = self.dec_1[6](out)
+        out = self.dec_1[7](out)
+        out = self.dec_1[8](out)
+        out = self.dec_1[9](out)
+        out = self.dec_1[10](out)
+        out = self.dec_1[11](out)
+        out = self.dec_1[12](out, quant_step=quant_step)
+        out = self.dec_2(out)
+        out = F.pixel_shuffle(out, 8)
+        return out
+
+
+class DMCI(CompressionModel):
+    def __init__(self, N=256, z_channel=128):
+        super().__init__(z_channel=z_channel)
+
+        self.enc = IntraEncoder(N)
+
+        self.hyper_enc = nn.Sequential(
+            DepthConvBlock(N, z_channel),
+            ResidualBlockWithStride2(z_channel, z_channel),
+            ResidualBlockWithStride2(z_channel, z_channel),
+        )
+
+        self.hyper_dec = nn.Sequential(
+            ResidualBlockUpsample(z_channel, z_channel),
+            ResidualBlockUpsample(z_channel, z_channel),
+            DepthConvBlock(z_channel, N),
+        )
+
+        self.y_prior_fusion = nn.Sequential(
+            DepthConvBlock(N, N * 2),
+            DepthConvBlock(N * 2, N * 2),
+            DepthConvBlock(N * 2, N * 2),
+            nn.Conv2d(N * 2, N * 2 + 2, 1),
+        )
+
+        self.y_spatial_prior_reduction = nn.Conv2d(N * 2 + 2, N * 1, 1)
+        self.y_spatial_prior_adaptor_1 = DepthConvBlock(N * 2, N * 2, force_adaptor=True)
+        self.y_spatial_prior_adaptor_2 = DepthConvBlock(N * 2, N * 2, force_adaptor=True)
+        self.y_spatial_prior_adaptor_3 = DepthConvBlock(N * 2, N * 2, force_adaptor=True)
+        self.y_spatial_prior = nn.Sequential(
+            DepthConvBlock(N * 2, N * 2),
+            DepthConvBlock(N * 2, N * 2),
+            DepthConvBlock(N * 2, N * 2),
+            nn.Conv2d(N * 2, N * 2, 1),
+        )
+
+        self.dec = IntraDecoder(N)
+
+        self.q_scale_enc = nn.Parameter(torch.ones((self.get_qp_num(), g_ch_enc_dec, 1, 1)))
+        self.q_scale_dec = nn.Parameter(torch.ones((self.get_qp_num(), g_ch_enc_dec, 1, 1)))
+
+    def compress(self, x, qp):
+        device = x.device
+        curr_q_enc = self.q_scale_enc[qp:qp+1, :, :, :]
+        curr_q_dec = self.q_scale_dec[qp:qp+1, :, :, :]
+
+        y = self.enc(x, curr_q_enc)
+        y_pad = self.pad_for_y(y)
+        z = self.hyper_enc(y_pad)
+        z_hat, z_hat_write = round_and_to_int8(z)
+
+        params = self.hyper_dec(z_hat)
+        params = self.y_prior_fusion(params)
+        _, _, yH, yW = y.shape
+        params = params[:, :, :yH, :yW].contiguous()
+        y_q_w_0, y_q_w_1, y_q_w_2, y_q_w_3, s_w_0, s_w_1, s_w_2, s_w_3, y_hat = \
+            self.compress_prior_4x(
+                y, params, self.y_spatial_prior_reduction,
+                self.y_spatial_prior_adaptor_1, self.y_spatial_prior_adaptor_2,
+                self.y_spatial_prior_adaptor_3, self.y_spatial_prior)
+
+        cuda_event = torch.cuda.Event()
+        cuda_event.record()
+        x_hat = self.dec(y_hat, curr_q_dec).clamp_(0, 1)
+
+        cuda_stream = self.get_cuda_stream(device=device, priority=-1)
+        with torch.cuda.stream(cuda_stream):
+            cuda_event.wait()
+            self.entropy_coder.reset()
+            self.bit_estimator_z.encode_z(z_hat_write, qp)
+            self.gaussian_encoder.encode_y(y_q_w_0, s_w_0)
+            self.gaussian_encoder.encode_y(y_q_w_1, s_w_1)
+            self.gaussian_encoder.encode_y(y_q_w_2, s_w_2)
+            self.gaussian_encoder.encode_y(y_q_w_3, s_w_3)
+            self.entropy_coder.flush()
+
+        bit_stream = self.entropy_coder.get_encoded_stream()
+
+        torch.cuda.synchronize(device=device)
+        result = {
+            "bit_stream": bit_stream,
+            "x_hat": x_hat,
+        }
+        return result
+
+    def decompress(self, bit_stream, sps, qp):
+        dtype = next(self.parameters()).dtype
+        device = next(self.parameters()).device
+        curr_q_dec = self.q_scale_dec[qp:qp+1, :, :, :]
+
+        self.entropy_coder.set_use_two_entropy_coders(sps['ec_part'] == 1)
+        self.entropy_coder.set_stream(bit_stream)
+        z_size = self.get_downsampled_shape(sps['height'], sps['width'], 64)
+        y_height, y_width = self.get_downsampled_shape(sps['height'], sps['width'], 16)
+        self.bit_estimator_z.decode_z(z_size, qp)
+        z_q = self.bit_estimator_z.get_z(z_size, device, dtype)
+        z_hat = z_q
+
+        params = self.hyper_dec(z_hat)
+        params = self.y_prior_fusion(params)
+        params = params[:, :, :y_height, :y_width].contiguous()
+        y_hat = self.decompress_prior_4x(params, self.y_spatial_prior_reduction,
+                                         self.y_spatial_prior_adaptor_1,
+                                         self.y_spatial_prior_adaptor_2,
+                                         self.y_spatial_prior_adaptor_3, self.y_spatial_prior)
+
+        x_hat = self.dec(y_hat, curr_q_dec).clamp_(0, 1)
+        return {"x_hat": x_hat}

neural_codec/DCVC/src/models/video_model.py

@@ -0,0 +1,379 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+
+import torch
+import torch.nn.functional as F
+from torch import nn
+
+from .common_model import CompressionModel
+from ..layers.layers import SubpelConv2x, DepthConvBlock, \
+    ResidualBlockUpsample, ResidualBlockWithStride2
+from ..layers.cuda_inference import CUSTOMIZED_CUDA_INFERENCE, round_and_to_int8, \
+    bias_pixel_shuffle_8, bias_quant
+
+
+qp_shift = [0, 8, 4]
+extra_qp = max(qp_shift)
+
+g_ch_src_d = 3 * 8 * 8
+g_ch_recon = 320
+g_ch_y = 128
+g_ch_z = 128
+g_ch_d = 256
+
+
+class FeatureExtractor(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.conv1 = nn.Sequential(
+            DepthConvBlock(g_ch_d, g_ch_d),
+            DepthConvBlock(g_ch_d, g_ch_d),
+        )
+        self.conv2 = nn.Sequential(
+            DepthConvBlock(g_ch_d, g_ch_d),
+            DepthConvBlock(g_ch_d, g_ch_d),
+            DepthConvBlock(g_ch_d, g_ch_d),
+            DepthConvBlock(g_ch_d, g_ch_d),
+        )
+
+    def forward(self, x, quant):
+        x1, ctx_t = self.forward_part1(x, quant)
+        ctx = self.forward_part2(x1)
+        return ctx, ctx_t
+
+    def forward_part1(self, x, quant):
+        x1 = self.conv1(x)
+        ctx_t = x1 * quant
+        return x1, ctx_t
+
+    def forward_part2(self, x1):
+        ctx = self.conv2(x1)
+        return ctx
+
+
+class Encoder(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.conv1 = nn.Conv2d(g_ch_src_d, g_ch_d, 1)
+        self.conv2 = nn.Sequential(
+            DepthConvBlock(g_ch_d * 2, g_ch_d),
+            DepthConvBlock(g_ch_d, g_ch_d),
+        )
+        self.conv3 = DepthConvBlock(g_ch_d, g_ch_d)
+        self.down = nn.Conv2d(g_ch_d, g_ch_y, 3, stride=2, padding=1)
+
+        self.fuse_conv1_flag = False
+
+    def forward(self, x, ctx, quant_step):
+        feature = F.pixel_unshuffle(x, 8)
+        if not CUSTOMIZED_CUDA_INFERENCE or not x.is_cuda:
+            return self.forward_torch(feature, ctx, quant_step)
+        return self.forward_cuda(feature, ctx, quant_step)
+
+    def forward_torch(self, feature, ctx, quant_step):
+        feature = self.conv1(feature)
+        feature = self.conv2(torch.cat((feature, ctx), dim=1))
+        feature = self.conv3(feature)
+        feature = feature * quant_step
+        feature = self.down(feature)
+        return feature
+
+    def forward_cuda(self, feature, ctx, quant_step):
+        if not self.fuse_conv1_flag:
+            fuse_weight1 = torch.matmul(
+                self.conv2[0].adaptor.weight[:, :g_ch_d, 0, 0],
+                self.conv1.weight[:, :, 0, 0]
+            )[:, :, None, None]
+            fuse_weight2 = self.conv2[0].adaptor.weight[:, g_ch_d:]
+            self.conv2[0].adaptor.bias.data = self.conv2[0].adaptor.bias + \
+                torch.matmul(self.conv2[0].adaptor.weight[:, :g_ch_d, 0, 0],
+                             self.conv1.bias[:, None])[:, 0]
+            self.conv2[0].adaptor.weight.data = torch.cat([fuse_weight1, fuse_weight2], dim=1)
+            self.fuse_conv1_flag = True
+
+        feature = self.conv2(torch.cat((feature, ctx), dim=1))
+        feature = self.conv3(feature, quant_step=quant_step)
+        feature = self.down(feature)
+        return feature
+
+
+class Decoder(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.up = SubpelConv2x(g_ch_y, g_ch_d, 3, padding=1)
+        self.conv1 = nn.Sequential(
+            DepthConvBlock(g_ch_d * 2, g_ch_d),
+            DepthConvBlock(g_ch_d, g_ch_d),
+            DepthConvBlock(g_ch_d, g_ch_d),
+        )
+        self.conv2 = nn.Conv2d(g_ch_d, g_ch_d, 1)
+
+    def forward(self, x, ctx, quant_step,):
+        if not CUSTOMIZED_CUDA_INFERENCE or not x.is_cuda:
+            return self.forward_torch(x, ctx, quant_step)
+
+        return self.forward_cuda(x, ctx, quant_step)
+
+    def forward_torch(self, x, ctx, quant_step):
+        feature = self.up(x)
+        feature = self.conv1(torch.cat((feature, ctx), dim=1))
+        feature = self.conv2(feature)
+        feature = feature * quant_step
+        return feature
+
+    def forward_cuda(self, x, ctx, quant_step):
+        feature = self.up(x, to_cat=ctx, cat_at_front=False)
+        feature = self.conv1(feature)
+        feature = F.conv2d(feature, self.conv2.weight)
+        feature = bias_quant(feature, self.conv2.bias, quant_step)
+        return feature
+
+
+class ReconGeneration(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.conv = nn.Sequential(
+            DepthConvBlock(g_ch_d,     g_ch_recon),
+            DepthConvBlock(g_ch_recon, g_ch_recon),
+            DepthConvBlock(g_ch_recon, g_ch_recon),
+            DepthConvBlock(g_ch_recon, g_ch_recon),
+        )
+        self.head = nn.Conv2d(g_ch_recon, g_ch_src_d, 1)
+
+    def forward(self, x, quant_step):
+        if not CUSTOMIZED_CUDA_INFERENCE or not x.is_cuda:
+            return self.forward_torch(x, quant_step)
+
+        return self.forward_cuda(x, quant_step)
+
+    def forward_torch(self, x, quant_step):
+        out = self.conv(x)
+        out = out * quant_step
+        out = self.head(out)
+        out = F.pixel_shuffle(out, 8)
+        out = torch.clamp(out, 0., 1.)
+        return out
+
+    def forward_cuda(self, x, quant_step):
+        out = self.conv[0](x)
+        out = self.conv[1](out)
+        out = self.conv[2](out)
+        out = self.conv[3](out, quant_step=quant_step)
+        out = F.conv2d(out, self.head.weight)
+        return bias_pixel_shuffle_8(out, self.head.bias)
+
+
+class HyperEncoder(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.conv = nn.Sequential(
+            DepthConvBlock(g_ch_y, g_ch_z),
+            ResidualBlockWithStride2(g_ch_z, g_ch_z),
+            ResidualBlockWithStride2(g_ch_z, g_ch_z),
+        )
+
+    def forward(self, x):
+        return self.conv(x)
+
+
+class HyperDecoder(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.conv = nn.Sequential(
+            ResidualBlockUpsample(g_ch_z, g_ch_z),
+            ResidualBlockUpsample(g_ch_z, g_ch_z),
+            DepthConvBlock(g_ch_z, g_ch_y),
+        )
+
+    def forward(self, x):
+        return self.conv(x)
+
+
+class PriorFusion(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.conv = nn.Sequential(
+            DepthConvBlock(g_ch_y * 3, g_ch_y * 3),
+            DepthConvBlock(g_ch_y * 3, g_ch_y * 3),
+            DepthConvBlock(g_ch_y * 3, g_ch_y * 3),
+            nn.Conv2d(g_ch_y * 3, g_ch_y * 3, 1),
+        )
+
+    def forward(self, x):
+        return self.conv(x)
+
+
+class SpatialPrior(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.conv = nn.Sequential(
+            DepthConvBlock(g_ch_y * 4, g_ch_y * 3),
+            DepthConvBlock(g_ch_y * 3, g_ch_y * 3),
+            nn.Conv2d(g_ch_y * 3, g_ch_y * 2, 1),
+        )
+
+    def forward(self, x):
+        return self.conv(x)
+
+
+class RefFrame():
+    def __init__(self):
+        self.frame = None
+        self.feature = None
+        self.poc = None
+
+
+class DMC(CompressionModel):
+    def __init__(self):
+        super().__init__(z_channel=g_ch_z, extra_qp=extra_qp)
+        self.qp_shift = qp_shift
+
+        self.feature_adaptor_i = DepthConvBlock(g_ch_src_d, g_ch_d)
+        self.feature_adaptor_p = nn.Conv2d(g_ch_d, g_ch_d, 1)
+        self.feature_extractor = FeatureExtractor()
+
+        self.encoder = Encoder()
+        self.hyper_encoder = HyperEncoder()
+        self.hyper_decoder = HyperDecoder()
+        self.temporal_prior_encoder = ResidualBlockWithStride2(g_ch_d, g_ch_y * 2)
+        self.y_prior_fusion = PriorFusion()
+        self.y_spatial_prior = SpatialPrior()
+        self.decoder = Decoder()
+        self.recon_generation_net = ReconGeneration()
+
+        self.q_encoder = nn.Parameter(torch.ones((self.get_qp_num() + extra_qp, g_ch_d, 1, 1)))
+        self.q_decoder = nn.Parameter(torch.ones((self.get_qp_num() + extra_qp, g_ch_d, 1, 1)))
+        self.q_feature = nn.Parameter(torch.ones((self.get_qp_num() + extra_qp, g_ch_d, 1, 1)))
+        self.q_recon = nn.Parameter(torch.ones((self.get_qp_num() + extra_qp, g_ch_recon, 1, 1)))
+
+        self.dpb = []
+        self.max_dpb_size = 1
+        self.curr_poc = 0
+
+    def reset_ref_feature(self):
+        if len(self.dpb) > 0:
+            self.dpb[0].feature = None
+
+    def add_ref_frame(self, feature=None, frame=None, increase_poc=True):
+        ref_frame = RefFrame()
+        ref_frame.poc = self.curr_poc
+        ref_frame.frame = frame
+        ref_frame.feature = feature
+        if len(self.dpb) >= self.max_dpb_size:
+            self.dpb.pop(-1)
+        self.dpb.insert(0, ref_frame)
+        if increase_poc:
+            self.curr_poc += 1
+
+    def clear_dpb(self):
+        self.dpb.clear()
+
+    def set_curr_poc(self, poc):
+        self.curr_poc = poc
+
+    def apply_feature_adaptor(self):
+        if self.dpb[0].feature is None:
+            return self.feature_adaptor_i(F.pixel_unshuffle(self.dpb[0].frame, 8))
+        return self.feature_adaptor_p(self.dpb[0].feature)
+
+    def res_prior_param_decoder(self, z_hat, ctx_t):
+        hierarchical_params = self.hyper_decoder(z_hat)
+        temporal_params = self.temporal_prior_encoder(ctx_t)
+        _, _, H, W = temporal_params.shape
+        hierarchical_params = hierarchical_params[:, :, :H, :W].contiguous()
+        params = self.y_prior_fusion(
+            torch.cat((hierarchical_params, temporal_params), dim=1))
+        return params
+
+    def get_recon_and_feature(self, y_hat, ctx, q_decoder, q_recon):
+        feature = self.decoder(y_hat, ctx, q_decoder)
+        x_hat = self.recon_generation_net(feature, q_recon)
+        return x_hat, feature
+
+    def prepare_feature_adaptor_i(self, last_qp):
+        if self.dpb[0].frame is None:
+            q_recon = self.q_recon[last_qp:last_qp+1, :, :, :]
+            self.dpb[0].frame = self.recon_generation_net(self.dpb[0].feature, q_recon).clamp_(0, 1)
+            self.reset_ref_feature()
+
+    def compress(self, x, qp):
+        # pic_width and pic_height may be different from x's size. x here is after padding
+        # x_hat has the same size with x
+        device = x.device
+        q_encoder = self.q_encoder[qp:qp+1, :, :, :]
+        q_decoder = self.q_decoder[qp:qp+1, :, :, :]
+        q_feature = self.q_feature[qp:qp+1, :, :, :]
+
+        feature = self.apply_feature_adaptor()
+        ctx, ctx_t = self.feature_extractor(feature, q_feature)
+        y = self.encoder(x, ctx, q_encoder)
+
+        hyper_inp = self.pad_for_y(y)
+
+        z = self.hyper_encoder(hyper_inp)
+        z_hat, z_hat_write = round_and_to_int8(z)
+        cuda_event_z_ready = torch.cuda.Event()
+        cuda_event_z_ready.record()
+        params = self.res_prior_param_decoder(z_hat, ctx_t)
+        y_q_w_0, y_q_w_1, s_w_0, s_w_1, y_hat = \
+            self.compress_prior_2x(y, params, self.y_spatial_prior)
+
+        cuda_event_y_ready = torch.cuda.Event()
+        cuda_event_y_ready.record()
+        feature = self.decoder(y_hat, ctx, q_decoder)
+
+        cuda_stream = self.get_cuda_stream(device=device, priority=-1)
+        with torch.cuda.stream(cuda_stream):
+            self.entropy_coder.reset()
+            cuda_event_z_ready.wait()
+            self.bit_estimator_z.encode_z(z_hat_write, qp)
+            cuda_event_y_ready.wait()
+            self.gaussian_encoder.encode_y(y_q_w_0, s_w_0)
+            self.gaussian_encoder.encode_y(y_q_w_1, s_w_1)
+            self.entropy_coder.flush()
+
+        bit_stream = self.entropy_coder.get_encoded_stream()
+
+        torch.cuda.synchronize(device=device)
+        self.add_ref_frame(feature, None)
+        return {
+            'bit_stream': bit_stream,
+        }
+
+    def decompress(self, bit_stream, sps, qp):
+        dtype = next(self.parameters()).dtype
+        device = next(self.parameters()).device
+        q_decoder = self.q_decoder[qp:qp+1, :, :, :]
+        q_feature = self.q_feature[qp:qp+1, :, :, :]
+        q_recon = self.q_recon[qp:qp+1, :, :, :]
+
+        self.entropy_coder.set_use_two_entropy_coders(sps['ec_part'] == 1)
+        self.entropy_coder.set_stream(bit_stream)
+        z_size = self.get_downsampled_shape(sps['height'], sps['width'], 64)
+        self.bit_estimator_z.decode_z(z_size, qp)
+
+        feature = self.apply_feature_adaptor()
+        c1, ctx_t = self.feature_extractor.forward_part1(feature, q_feature)
+
+        z_hat = self.bit_estimator_z.get_z(z_size, device, dtype)
+        params = self.res_prior_param_decoder(z_hat, ctx_t)
+        infos = self.decompress_prior_2x_part1(params)
+
+        ctx = self.feature_extractor.forward_part2(c1)
+
+        cuda_stream = self.get_cuda_stream(device=device, priority=-1)
+        with torch.cuda.stream(cuda_stream):
+            y_hat = self.decompress_prior_2x_part2(params, self.y_spatial_prior, infos)
+            cuda_event = torch.cuda.Event()
+            cuda_event.record()
+
+        cuda_event.wait()
+        x_hat, feature = self.get_recon_and_feature(y_hat, ctx, q_decoder, q_recon)
+
+        self.add_ref_frame(feature, x_hat)
+        return {
+            'x_hat': x_hat,
+        }
+
+    def shift_qp(self, qp, fa_idx):
+        return qp + self.qp_shift[fa_idx]

neural_codec/DCVC/src/utils/common.py

@@ -0,0 +1,177 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+
+import json
+import os
+from unittest.mock import patch
+
+import torch
+from torch.nn.modules.utils import consume_prefix_in_state_dict_if_present
+import numpy as np
+
+
+def str2bool(v):
+    return str(v).lower() in ("yes", "y", "true", "t", "1")
+
+
+def set_torch_env():
+    os.environ['CUBLAS_WORKSPACE_CONFIG'] = ":4096:8"
+    torch.backends.cudnn.enabled = True
+    torch.backends.cudnn.benchmark = False
+    torch.use_deterministic_algorithms(True)
+    torch.manual_seed(0)
+    torch.set_num_threads(1)
+    np.random.seed(seed=0)
+    try:
+        # require pytorch >= 2.2.0
+        torch.utils.deterministic.fill_uninitialized_memory = False
+    except Exception:  # pylint: disable=W0718
+        pass
+
+
+def create_folder(path, print_if_create=False):
+    if not os.path.exists(path):
+        os.makedirs(path)
+        if print_if_create:
+            print(f"created folder: {path}")
+
+
+def get_state_dict(ckpt_path):
+    ckpt = torch.load(ckpt_path, map_location=torch.device('cpu'), weights_only=True)
+    if "state_dict" in ckpt:
+        ckpt = ckpt['state_dict']
+    if "net" in ckpt:
+        ckpt = ckpt["net"]
+    consume_prefix_in_state_dict_if_present(ckpt, prefix="module.")
+    return ckpt
+
+
+@patch('json.encoder.c_make_encoder', None)
+def dump_json(obj, fid, float_digits=-1, **kwargs):
+    of = json.encoder._make_iterencode  # pylint: disable=W0212
+
+    def inner(*args, **kwargs):
+        args = list(args)
+        # fifth argument is float formater which we will replace
+        args[4] = lambda o: format(o, '.%df' % float_digits)
+        return of(*args, **kwargs)
+
+    with patch('json.encoder._make_iterencode', wraps=inner):
+        json.dump(obj, fid, **kwargs)
+
+
+def generate_log_json(frame_num, frame_pixel_num, test_time, frame_types, bits, psnrs, ssims,
+                      verbose=False, avg_encoding_time=None, avg_decoding_time=None):
+    include_yuv = len(psnrs[0]) > 1
+    assert not include_yuv or (len(psnrs[0]) == 4 and len(ssims[0]) == 4)
+    i_bits = 0
+    i_psnr = 0
+    i_psnr_y = 0
+    i_psnr_u = 0
+    i_psnr_v = 0
+    i_ssim = 0
+    i_ssim_y = 0
+    i_ssim_u = 0
+    i_ssim_v = 0
+    p_bits = 0
+    p_psnr = 0
+    p_psnr_y = 0
+    p_psnr_u = 0
+    p_psnr_v = 0
+    p_ssim = 0
+    p_ssim_y = 0
+    p_ssim_u = 0
+    p_ssim_v = 0
+    i_num = 0
+    p_num = 0
+    for idx in range(frame_num):
+        if frame_types[idx] == 0:
+            i_bits += bits[idx]
+            i_psnr += psnrs[idx][0]
+            i_ssim += ssims[idx][0]
+            i_num += 1
+            if include_yuv:
+                i_psnr_y += psnrs[idx][1]
+                i_psnr_u += psnrs[idx][2]
+                i_psnr_v += psnrs[idx][3]
+                i_ssim_y += ssims[idx][1]
+                i_ssim_u += ssims[idx][2]
+                i_ssim_v += ssims[idx][3]
+        else:
+            p_bits += bits[idx]
+            p_psnr += psnrs[idx][0]
+            p_ssim += ssims[idx][0]
+            p_num += 1
+            if include_yuv:
+                p_psnr_y += psnrs[idx][1]
+                p_psnr_u += psnrs[idx][2]
+                p_psnr_v += psnrs[idx][3]
+                p_ssim_y += ssims[idx][1]
+                p_ssim_u += ssims[idx][2]
+                p_ssim_v += ssims[idx][3]
+
+    log_result = {}
+    log_result['frame_pixel_num'] = frame_pixel_num
+    log_result['i_frame_num'] = i_num
+    log_result['p_frame_num'] = p_num
+    log_result['ave_i_frame_bpp'] = i_bits / i_num / frame_pixel_num
+    log_result['ave_i_frame_psnr'] = i_psnr / i_num
+    log_result['ave_i_frame_msssim'] = i_ssim / i_num
+    if include_yuv:
+        log_result['ave_i_frame_psnr_y'] = i_psnr_y / i_num
+        log_result['ave_i_frame_psnr_u'] = i_psnr_u / i_num
+        log_result['ave_i_frame_psnr_v'] = i_psnr_v / i_num
+        log_result['ave_i_frame_msssim_y'] = i_ssim_y / i_num
+        log_result['ave_i_frame_msssim_u'] = i_ssim_u / i_num
+        log_result['ave_i_frame_msssim_v'] = i_ssim_v / i_num
+    if verbose:
+        log_result['frame_bpp'] = list(np.array(bits) / frame_pixel_num)
+        log_result['frame_psnr'] = [v[0] for v in psnrs]
+        log_result['frame_msssim'] = [v[0] for v in ssims]
+        log_result['frame_type'] = frame_types
+        if include_yuv:
+            log_result['frame_psnr_y'] = [v[1] for v in psnrs]
+            log_result['frame_psnr_u'] = [v[2] for v in psnrs]
+            log_result['frame_psnr_v'] = [v[3] for v in psnrs]
+            log_result['frame_msssim_y'] = [v[1] for v in ssims]
+            log_result['frame_msssim_u'] = [v[2] for v in ssims]
+            log_result['frame_msssim_v'] = [v[3] for v in ssims]
+    log_result['test_time'] = test_time
+    if p_num > 0:
+        total_p_pixel_num = p_num * frame_pixel_num
+        log_result['ave_p_frame_bpp'] = p_bits / total_p_pixel_num
+        log_result['ave_p_frame_psnr'] = p_psnr / p_num
+        log_result['ave_p_frame_msssim'] = p_ssim / p_num
+        if include_yuv:
+            log_result['ave_p_frame_psnr_y'] = p_psnr_y / p_num
+            log_result['ave_p_frame_psnr_u'] = p_psnr_u / p_num
+            log_result['ave_p_frame_psnr_v'] = p_psnr_v / p_num
+            log_result['ave_p_frame_msssim_y'] = p_ssim_y / p_num
+            log_result['ave_p_frame_msssim_u'] = p_ssim_u / p_num
+            log_result['ave_p_frame_msssim_v'] = p_ssim_v / p_num
+    else:
+        log_result['ave_p_frame_bpp'] = 0
+        log_result['ave_p_frame_psnr'] = 0
+        log_result['ave_p_frame_msssim'] = 0
+        if include_yuv:
+            log_result['ave_p_frame_psnr_y'] = 0
+            log_result['ave_p_frame_psnr_u'] = 0
+            log_result['ave_p_frame_psnr_v'] = 0
+            log_result['ave_p_frame_msssim_y'] = 0
+            log_result['ave_p_frame_msssim_u'] = 0
+            log_result['ave_p_frame_msssim_v'] = 0
+    log_result['ave_all_frame_bpp'] = (i_bits + p_bits) / (frame_num * frame_pixel_num)
+    log_result['ave_all_frame_psnr'] = (i_psnr + p_psnr) / frame_num
+    log_result['ave_all_frame_msssim'] = (i_ssim + p_ssim) / frame_num
+    if avg_encoding_time is not None and avg_decoding_time is not None:
+        log_result['avg_frame_encoding_time'] = avg_encoding_time
+        log_result['avg_frame_decoding_time'] = avg_decoding_time
+    if include_yuv:
+        log_result['ave_all_frame_psnr_y'] = (i_psnr_y + p_psnr_y) / frame_num
+        log_result['ave_all_frame_psnr_u'] = (i_psnr_u + p_psnr_u) / frame_num
+        log_result['ave_all_frame_psnr_v'] = (i_psnr_v + p_psnr_v) / frame_num
+        log_result['ave_all_frame_msssim_y'] = (i_ssim_y + p_ssim_y) / frame_num
+        log_result['ave_all_frame_msssim_u'] = (i_ssim_u + p_ssim_u) / frame_num
+        log_result['ave_all_frame_msssim_v'] = (i_ssim_v + p_ssim_v) / frame_num
+
+    return log_result

neural_codec/DCVC/src/utils/metrics.py

@@ -0,0 +1,96 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+
+import numpy as np
+from scipy import signal
+from scipy import ndimage
+
+
+def fspecial_gauss(size, sigma):
+    x, y = np.mgrid[-size // 2 + 1:size // 2 + 1, -size // 2 + 1:size // 2 + 1]
+    g = np.exp(-((x**2 + y**2) / (2.0 * sigma**2)))
+    return g / g.sum()
+
+
+def calc_ssim(img1, img2, data_range=255):
+    img1 = img1.astype(np.float64)
+    img2 = img2.astype(np.float64)
+    size = 11
+    sigma = 1.5
+    window = fspecial_gauss(size, sigma)
+    K1 = 0.01
+    K2 = 0.03
+    C1 = (K1 * data_range)**2
+    C2 = (K2 * data_range)**2
+    mu1 = signal.fftconvolve(window, img1, mode='valid')
+    mu2 = signal.fftconvolve(window, img2, mode='valid')
+    mu1_sq = mu1 * mu1
+    mu2_sq = mu2 * mu2
+    mu1_mu2 = mu1 * mu2
+    sigma1_sq = signal.fftconvolve(window, img1 * img1, mode='valid') - mu1_sq
+    sigma2_sq = signal.fftconvolve(window, img2 * img2, mode='valid') - mu2_sq
+    sigma12 = signal.fftconvolve(window, img1 * img2, mode='valid') - mu1_mu2
+
+    return (((2 * mu1_mu2 + C1) * (2 * sigma12 + C2)) / ((mu1_sq + mu2_sq + C1) *
+                                                         (sigma1_sq + sigma2_sq + C2)),
+            (2.0 * sigma12 + C2) / (sigma1_sq + sigma2_sq + C2))
+
+
+def calc_msssim(img1, img2, data_range=255):
+    '''
+    img1 and img2 are 2D arrays
+    '''
+    level = 5
+    weight = np.array([0.0448, 0.2856, 0.3001, 0.2363, 0.1333])
+    height, width = img1.shape
+    if height < 176 or width < 176:
+        # according to HM implementation
+        level = 4
+        weight = np.array([0.0517, 0.3295, 0.3462, 0.2726])
+    if height < 88 or width < 88:
+        assert False
+    downsample_filter = np.ones((2, 2)) / 4.0
+    im1 = img1.astype(np.float64)
+    im2 = img2.astype(np.float64)
+    mssim = np.array([])
+    mcs = np.array([])
+    for _ in range(level):
+        ssim_map, cs_map = calc_ssim(im1, im2, data_range=data_range)
+        mssim = np.append(mssim, ssim_map.mean())
+        mcs = np.append(mcs, cs_map.mean())
+        filtered_im1 = ndimage.filters.convolve(im1, downsample_filter,
+                                                mode='reflect')
+        filtered_im2 = ndimage.filters.convolve(im2, downsample_filter,
+                                                mode='reflect')
+        im1 = filtered_im1[::2, ::2]
+        im2 = filtered_im2[::2, ::2]
+    return (np.prod(mcs[0:level - 1]**weight[0:level - 1]) *
+            (mssim[level - 1]**weight[level - 1]))
+
+
+def calc_msssim_rgb(img1, img2, data_range=255):
+    '''
+    img1 and img2 are arrays with 3xHxW
+    '''
+    msssim = 0
+    for i in range(3):
+        msssim += calc_msssim(img1[i, :, :], img2[i, :, :], data_range)
+    return msssim / 3
+
+
+def calc_psnr(img1, img2, data_range=255):
+    '''
+    img1 and img2 are arrays with same shape
+    '''
+    img1 = img1.astype(np.float64)
+    img2 = img2.astype(np.float64)
+    mse = np.mean(np.square(img1 - img2))
+    if np.isnan(mse) or np.isinf(mse):
+        return -999.9
+    if mse > 1e-10:
+        psnr = 10 * np.log10(data_range * data_range / mse)
+    else:
+        psnr = 999.9
+    if psnr > 99.9:
+        psnr = 99.9
+    return psnr

neural_codec/DCVC/src/utils/stream_helper.py

@@ -0,0 +1,217 @@
+# Copyright 2020 InterDigital Communications, Inc.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import enum
+import struct
+from pathlib import Path
+
+
+def filesize(filepath: str) -> int:
+    if not Path(filepath).is_file():
+        raise ValueError(f'Invalid file "{filepath}".')
+    return Path(filepath).stat().st_size
+
+
+def write_uints(fd, values, fmt=">{:d}I"):
+    fd.write(struct.pack(fmt.format(len(values)), *values))
+    return len(values) * 4
+
+
+def write_uchars(fd, values, fmt=">{:d}B"):
+    fd.write(struct.pack(fmt.format(len(values)), *values))
+    return len(values)
+
+
+def read_uints(fd, n, fmt=">{:d}I"):
+    sz = struct.calcsize("I")
+    return struct.unpack(fmt.format(n), fd.read(n * sz))
+
+
+def read_uchars(fd, n, fmt=">{:d}B"):
+    sz = struct.calcsize("B")
+    return struct.unpack(fmt.format(n), fd.read(n * sz))
+
+
+def write_bytes(fd, values, fmt=">{:d}s"):
+    if len(values) == 0:
+        return 0
+    fd.write(struct.pack(fmt.format(len(values)), values))
+    return len(values)
+
+
+def read_bytes(fd, n, fmt=">{:d}s"):
+    sz = struct.calcsize("s")
+    return struct.unpack(fmt.format(n), fd.read(n * sz))[0]
+
+
+def write_ushorts(fd, values, fmt=">{:d}H"):
+    fd.write(struct.pack(fmt.format(len(values)), *values))
+    return len(values) * 2
+
+
+def read_ushorts(fd, n, fmt=">{:d}H"):
+    sz = struct.calcsize("H")
+    return struct.unpack(fmt.format(n), fd.read(n * sz))
+
+
+def write_uint_adaptive(f, a):
+    if a < (1 << 7):
+        a0 = (a >> 0) & 0xff
+        a0 = a0 | (0x00 << 7)
+        write_uchars(f, (a0,))
+        return 1
+
+    if a < (1 << 14):
+        a0 = (a >> 0) & 0xff
+        a1 = (a >> 8) & 0xff
+        a1 = a1 | (0x02 << 6)
+        write_uchars(f, (a1, a0))
+        return 2
+
+    assert a < (1 << 30)
+    a0 = (a >> 0) & 0xff
+    a1 = (a >> 8) & 0xff
+    a2 = (a >> 16) & 0xff
+    a3 = (a >> 24) & 0xff
+    a3 = a3 | (0x03 << 6)
+    write_uchars(f, (a3, a2, a1, a0))
+    return 4
+
+
+def read_uint_adaptive(f):
+    a3 = read_uchars(f, 1)[0]
+    if (a3 >> 7) == 0:
+        return a3
+
+    a2 = read_uchars(f, 1)[0]
+
+    if (a3 >> 6) == 0x02:
+        a3 = a3 & 0x3f
+        return (a3 << 8) + a2
+    a3 = a3 & 0x3f
+    a1 = read_uchars(f, 1)[0]
+    a0 = read_uchars(f, 1)[0]
+    return (a3 << 24) + (a2 << 16) + (a1 << 8) + a0
+
+
+class NalType(enum.IntEnum):
+    NAL_SPS = 0
+    NAL_I = 1
+    NAL_P = 2
+
+
+class SPSHelper():
+    def __init__(self):
+        super().__init__()
+        self.spss = []
+
+    def get_sps_id(self, target_sps):
+        min_id = -1
+        for sps in self.spss:
+            if sps['height'] == target_sps['height'] and sps['width'] == target_sps['width'] and \
+                    sps['use_ada_i'] == target_sps['use_ada_i'] and \
+                    sps['ec_part'] == target_sps['ec_part']:
+                return sps['sps_id'], False
+            if sps['sps_id'] > min_id:
+                min_id = sps['sps_id']
+        assert min_id < 15
+        sps = target_sps.copy()
+        sps['sps_id'] = min_id + 1
+        self.spss.append(sps)
+        return sps['sps_id'], True
+
+    def add_sps_by_id(self, sps):
+        for i in range(len(self.spss)):
+            if self.spss[i]['sps_id'] == sps['sps_id']:
+                self.spss[i] = sps.copy()
+                return
+        self.spss.append(sps.copy())
+
+    def get_sps_by_id(self, sps_id):
+        for sps in self.spss:
+            if sps['sps_id'] == sps_id:
+                return sps
+        return None
+
+
+def write_sps(f, sps):
+    # nal_type(4), sps_id(4)
+    # height (variable)
+    # width (vairable)
+    # 0(6), ec_part(1) use_ada_i(1)
+    assert sps['sps_id'] < 16
+    assert sps['use_ada_i'] < 2
+    written = 0
+    flag = int((NalType.NAL_SPS << 4) + sps['sps_id'])
+    written += write_uchars(f, (flag,))
+    written += write_uint_adaptive(f, sps['height'])
+    written += write_uint_adaptive(f, sps['width'])
+    flag = (sps['ec_part'] << 2) + sps['use_ada_i']
+    written += write_uchars(f, (flag,))
+    return written
+
+
+def read_header(f):
+    header = {}
+    flag = read_uchars(f, 1)[0]
+    nal_type = flag >> 4
+    header['nal_type'] = NalType(nal_type)
+    if nal_type < 3:
+        header['sps_id'] = flag & 0x0f
+        return header
+
+    frame_num_minus1 = flag & 0x0f
+    frame_num = frame_num_minus1 + 1
+    header['frame_num'] = frame_num
+    sps_ids = []
+    for _ in range(0, frame_num, 2):
+        flag = read_uchars(f, 1)[0]
+        sps_ids.append(flag >> 4)
+        sps_ids.append(flag & 0x0f)
+    sps_ids = sps_ids[:frame_num]
+    header['sps_ids'] = sps_ids
+    return header
+
+
+def read_sps_remaining(f, sps_id):
+    sps = {}
+    sps['sps_id'] = sps_id
+    sps['height'] = read_uint_adaptive(f)
+    sps['width'] = read_uint_adaptive(f)
+    flag = read_uchars(f, 1)[0]
+    sps['ec_part'] = (flag >> 2) & 0x01
+    sps['use_ada_i'] = flag & 0x01
+    return sps
+
+
+def write_ip(f, is_i_frame, sps_id, qp, bit_stream):
+    written = 0
+    flag = (int(NalType.NAL_I if is_i_frame else NalType.NAL_P) << 4) + sps_id
+    written += write_uchars(f, (flag,))
+    assert qp < 256 and qp >= 0
+    flag = qp
+    written += write_uchars(f, (flag,))
+    # we write all the streams in the same file, thus, we need to write the per-frame length
+    # if packed independently, we do not need to write it
+    written += write_uint_adaptive(f, len(bit_stream))
+    written += write_bytes(f, bit_stream)
+    return written
+
+
+def read_ip_remaining(f):
+    flag = read_uchars(f, 1)[0]
+    qp = flag
+    stream_length = read_uint_adaptive(f)
+    bit_stream = read_bytes(f, stream_length)
+    return qp, bit_stream

neural_codec/DCVC/src/utils/transforms.py

@@ -0,0 +1,63 @@
+import numpy as np
+import scipy.ndimage
+import torch
+import torch.nn.functional as F
+
+
+YCBCR_WEIGHTS = {
+    # Spec: (K_r, K_g, K_b) with K_g = 1 - K_r - K_b
+    "ITU-R_BT.709": (0.2126, 0.7152, 0.0722)
+}
+
+
+def ycbcr420_to_444_np(y, uv, order=0, separate=False):
+    '''
+    y is 1xhxw Y float numpy array
+    uv is 2x(h/2)x(w/2) UV float numpy array
+    order: 0 nearest neighbor (default), 1: binear
+    return value is 3xhxw YCbCr float numpy array
+    '''
+    uv = scipy.ndimage.zoom(uv, (1, 2, 2), order=order)
+    if separate:
+        return y, uv
+    yuv = np.concatenate((y, uv), axis=0)
+    return yuv
+
+
+def rgb2ycbcr(rgb, is_bgr=False):
+    if is_bgr:
+        b, g, r = rgb.chunk(3, -3)
+    else:
+        r, g, b = rgb.chunk(3, -3)
+    Kr, Kg, Kb = YCBCR_WEIGHTS["ITU-R_BT.709"]
+    y = Kr * r + Kg * g + Kb * b
+    cb = 0.5 * (b - y) / (1 - Kb) + 0.5
+    cr = 0.5 * (r - y) / (1 - Kr) + 0.5
+    ycbcr = torch.cat((y, cb, cr), dim=-3)
+    ycbcr = torch.clamp(ycbcr, 0., 1.)
+    return ycbcr
+
+
+def ycbcr2rgb(ycbcr, is_bgr=False, clamp=True):
+    y, cb, cr = ycbcr.chunk(3, -3)
+    Kr, Kg, Kb = YCBCR_WEIGHTS["ITU-R_BT.709"]
+    r = y + (2 - 2 * Kr) * (cr - 0.5)
+    b = y + (2 - 2 * Kb) * (cb - 0.5)
+    g = (y - Kr * r - Kb * b) / Kg
+    if is_bgr:
+        rgb = torch.cat((b, g, r), dim=-3)
+    else:
+        rgb = torch.cat((r, g, b), dim=-3)
+    if clamp:
+        rgb = torch.clamp(rgb, 0., 1.)
+    return rgb
+
+
+def yuv_444_to_420(yuv):
+    def _downsample(tensor):
+        return F.avg_pool2d(tensor, kernel_size=2, stride=2)
+
+    y = yuv[:, :1, :, :]
+    uv = yuv[:, 1:, :, :]
+
+    return y, _downsample(uv)

neural_codec/DCVC/src/utils/video_reader.py

@@ -0,0 +1,90 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+
+import os
+
+import numpy as np
+from PIL import Image
+
+
+class PNGReader():
+    def __init__(self, src_path, width, height, start_num=1):
+        self.eof = False
+        self.src_path = src_path
+        self.width = width
+        self.height = height
+        pngs = os.listdir(self.src_path)
+        if 'im1.png' in pngs:
+            self.padding = 1
+        elif 'im00001.png' in pngs:
+            self.padding = 5
+        else:
+            raise ValueError('unknown image naming convention; please specify')
+        self.current_frame_index = start_num
+
+    def read_one_frame(self):
+        # rgb: 3xhxw uint8 numpy array
+        if self.eof:
+            return None
+
+        png_path = os.path.join(self.src_path,
+                                f"im{str(self.current_frame_index).zfill(self.padding)}.png"
+                                )
+        if not os.path.exists(png_path):
+            self.eof = True
+            return None
+
+        rgb = Image.open(png_path).convert('RGB')
+        rgb = np.asarray(rgb).astype(np.uint8).transpose(2, 0, 1)
+        _, height, width = rgb.shape
+        assert height == self.height
+        assert width == self.width
+
+        self.current_frame_index += 1
+        return rgb
+
+    def close(self):
+        self.current_frame_index = 1
+
+
+class YUV420Reader():
+    def __init__(self, src_path, width, height, skip_frame=0):
+        self.eof = False
+        if not src_path.endswith('.yuv'):
+            src_path = src_path + '.yuv'
+        self.src_path = src_path
+
+        self.y_size = width * height
+        self.y_width = width
+        self.y_height = height
+        self.uv_size = width * height // 2
+        self.uv_width = width // 2
+        self.uv_height = height // 2
+        # pylint: disable=R1732
+        self.file = open(src_path, "rb")
+        # pylint: enable=R1732
+        skipped_frame = 0
+        while not self.eof and skipped_frame < skip_frame:
+            y = self.file.read(self.y_size)
+            uv = self.file.read(self.uv_size)
+            if not y or not uv:
+                self.eof = True
+            skipped_frame += 1
+
+    def read_one_frame(self):
+        # y: 1xhxw uint8 numpy array
+        # uv: 2x(h/2)x(w/2) uint8 numpy array
+        if self.eof:
+            return None, None
+        y = self.file.read(self.y_size)
+        uv = self.file.read(self.uv_size)
+        if not y or not uv:
+            self.eof = True
+            return None, None
+        y = np.frombuffer(y, dtype=np.uint8).copy().reshape(1, self.y_height, self.y_width)
+        uv = np.frombuffer(uv, dtype=np.uint8).copy().reshape(2, self.uv_height, self.uv_width)
+
+        return y, uv
+
+    def close(self):
+        self.file.close()

neural_codec/DCVC/src/utils/video_writer.py

@@ -0,0 +1,52 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+
+import os
+
+from PIL import Image
+
+
+class PNGWriter():
+    def __init__(self, dst_path, width, height):
+        self.dst_path = dst_path
+        self.width = width
+        self.height = height
+        self.padding = 5
+        self.current_frame_index = 1
+        os.makedirs(dst_path, exist_ok=True)
+
+    def write_one_frame(self, rgb):
+        # rgb: 3xhxw uint8 numpy array
+        rgb = rgb.transpose(1, 2, 0)
+
+        png_path = os.path.join(self.dst_path,
+                                f"im{str(self.current_frame_index).zfill(self.padding)}.png"
+                                )
+        Image.fromarray(rgb).save(png_path)
+
+        self.current_frame_index += 1
+
+    def close(self):
+        self.current_frame_index = 1
+
+
+class YUV420Writer():
+    def __init__(self, dst_path, width, height):
+        if not dst_path.endswith('.yuv'):
+            dst_path = dst_path + '/out.yuv'
+        self.dst_path = dst_path
+        self.width = width
+        self.height = height
+
+        # pylint: disable=R1732
+        self.file = open(dst_path, "wb")
+        # pylint: enable=R1732
+
+    def write_one_frame(self, y, uv):
+        # y: 1xhxw uint8 numpy array
+        # uv: 2x(h/2)x(w/2) uint8 numpy array
+        self.file.write(y.tobytes())
+        self.file.write(uv.tobytes())
+
+    def close(self):
+        self.file.close()

neural_codec/README.md

@@ -0,0 +1,40 @@
+# neural_codec — DCVC-RT patch selection
+
+Internal package for **DCVC-RT** neural-codec patch selection: the codec's per-frame
+**bit-cost map** decides which video patches to feed the VLM (regions the codec spends more
+bits on — motion / new detail — are kept; predictable background is dropped), as the neural
+alternative to the traditional HEVC (`cv-preinfer`) path.
+
+> **Usage, environment setup, and the controllable `codec.dcvc` parameters are documented in
+> the top-level [`../README.md`](../README.md).** This file is an internal file reference.
+
+## Files
+
+| File | Role |
+|------|------|
+| `dcvc_rt_engine.py` | Loads DCVC-RT intra/inter nets; `DMCIBitmap` / `DMCBitmap` add `compute_bitmap` (per-frame `(H/16, W/16)` bit-cost map) via a streaming `reset_sequence` / `step` API. Loads the bundled DCVC-RT source from `DCVC/` (no env var); checkpoints default to `dcvc_rt_intra.tar` / `dcvc_rt_inter.tar` in this dir (`DCVC_INTRA_TAR` / `DCVC_INTER_TAR` to override). |
+| `codec_dcvc_config.py` | Single source of truth — reads the `codec.dcvc` block of `../processor/preprocessor_config.json`. |
+| `dcvc_readiness_gen.py` | Config-driven generator: runs the readiness pipeline (`codec_tools/`) with DCVC bit-cost as the score source. Invoked by the model's codec path (`../processor/codec_video_processing_magevl.py::_run_dcvc_rt`). |
+| `reproduce_bench.py` | Reproduce the evaluated selection for one video (config-driven; presets `cap12` / `b50` / `s95_b50`). |
+| `codec_loader.py` | Load precomputed assets → model inputs via the release codec helpers. |
+| `infer_dcvc_rt.py` | Standalone end-to-end demo (`--asset_dir` or `--video`). |
+| `precompute_dcvc_rt.py` | Standalone, CLI-flag-driven batch precompute (video(s) → assets). |
+| `canvas_assembler.py` | Top-k patch selection + canvas packing used by the standalone precompute path. |
+| `codec_tools/` | Vendored readiness pipeline (frame sampling, grouping, 2×2-block selection, canvas packing). |
+| `DCVC/` | Bundled DCVC-RT source (MIT, [microsoft/DCVC](https://github.com/microsoft/DCVC)): the `src/` package the engine imports + the CUDA-kernel source under `src/layers/extensions/inference/`. No external checkout / `DCVC_RT_ROOT` needed. |
+| `dcvc_rt_intra.tar` / `dcvc_rt_inter.tar` | DCVC-RT checkpoints. |
+
+## Notes / limitations
+
+- The bit-cost map is the summed **y-bits** estimate from DCVC-RT's Gaussian entropy model
+  (the dominant, spatially-resolved term); the small hyperprior `z` term is omitted (it is a
+  ranking signal). It is **not** run through the RANS arithmetic coder.
+- `patch=16` is mandatory — it must match the image processor
+  (`preprocessor_config.json: patch_size=16, merge_size=2`). The `codec.patch=14` field is a
+  separate cv-preinfer internal and does not apply here.
+- Canvases are square; non-16:9 videos are letterboxed, so wide videos waste some budget on
+  padding. Tune `codec.dcvc.max_pixels` if needed.
+- DCVC-RT decodes every frame `0..max(sampled)` to keep temporal references valid, so long
+  videos are slow — use `codec.dcvc.max_side` and multiple GPUs.
+- The DCVC CUDA kernels fall back to pytorch when the compiled extension is unavailable
+  (slower but numerically fine, and deterministic on the fallback path).

neural_codec/canvas_assembler.py

@@ -0,0 +1,217 @@
+"""Patch selection + canvas assembly for the DCVC-RT release backend.
+
+Turns per-frame DCVC-RT bit-cost bitmaps into the exact on-disk contract the
+release codec path (``codec_video_processing_magevl.py``) consumes:
+
+    <out_dir>/canvas_000.jpg ...      RGB canvases, all identical size
+    <out_dir>/src_patch_position.npy  int32 [total_patches, 3] = (t, h, w)
+    <out_dir>/meta.json               {"fps": ..., "canvas_files": [...]}
+
+Layout invariants required by the downstream release code
+---------------------------------------------------------
+* All canvases share one size, so ``total_patches`` is divisible by the canvas
+  count (``drop_padding_canvases``).
+* Canvas patch grid ``(Gh, Gw)`` is even in both dims (``spatial_merge_size=2``
+  block reorder in ``codec_positions_for_processor``).
+* We select at **2x2-patch (28px) block granularity** and pack blocks in
+  time-sorted order, so after the block reorder the ``t`` column forms
+  consecutive runs whose lengths are multiples of 4 -- required for
+  ``rewrite_text_with_codec_positions`` (``count // merge**2`` token counts).
+* ``t`` is the *original video frame index* so ``t / fps`` is the real
+  timestamp (same convention the HEVC/cv-preinfer contract uses).
+* We never emit fully-padding canvases; a short final canvas is repeat-filled
+  with real blocks (so ``drop_padding_canvases`` is a no-op).
+
+The importance signal (which blocks to keep) is DCVC-RT bits; everything else
+mirrors the dev codec's ``pack_topk`` selection (mandatory full first frame,
+then global top-k over the rest).
+"""
+
+from __future__ import annotations
+
+import math
+from dataclasses import dataclass
+from typing import Dict, List, Optional, Tuple
+
+import cv2
+import numpy as np
+from PIL import Image
+
+
+@dataclass
+class AssembleConfig:
+    # Must match the release image processor (preprocessor_config.json):
+    # patch_size=16, merge_size=2 -> selectable unit is 32px. (The codec.patch=14
+    # field in that config is a cv-preinfer internal; the Qwen2VLImageProcessor
+    # grid that codec_positions_for_processor reads uses 16.)
+    patch: int = 16
+    spatial_merge_size: int = 2
+    target_canvas: int = 32
+    seq_len_frames: int = 64
+    max_pixels: int = 150000
+    # Canvas is square with ``canvas_token_side`` tokens per side (a token = one
+    # 2x2 patch block = ``patch*merge`` px). If None, derived from max_pixels.
+    canvas_token_side: Optional[int] = None
+    mandatory_first_frame: bool = True
+
+    @property
+    def unit(self) -> int:
+        return int(self.patch) * int(self.spatial_merge_size)
+
+    def token_side(self) -> int:
+        if self.canvas_token_side is not None:
+            return int(self.canvas_token_side)
+        side_px = math.sqrt(float(self.max_pixels))
+        return max(2, int(round(side_px / self.unit)))
+
+
+def _resize_longer_pad_square(img: np.ndarray, side_px: int, is_map: bool = False):
+    """Resize longer side to ``side_px`` (keep aspect) and center-pad to a
+    square. Returns (square_array, pad_info). ``is_map`` uses float + linear.
+    """
+    H, W = img.shape[:2]
+    scale = float(side_px) / float(max(H, W))
+    Hn = max(1, int(round(H * scale)))
+    Wn = max(1, int(round(W * scale)))
+    interp = cv2.INTER_LINEAR
+    resized = cv2.resize(img.astype(np.float32) if is_map else img, (Wn, Hn), interpolation=interp)
+    pad_top = (side_px - Hn) // 2
+    pad_left = (side_px - Wn) // 2
+    if is_map:
+        out = np.zeros((side_px, side_px), dtype=np.float32)
+        out[pad_top:pad_top + Hn, pad_left:pad_left + Wn] = resized
+    else:
+        out = np.zeros((side_px, side_px, 3), dtype=np.uint8)
+        out[pad_top:pad_top + Hn, pad_left:pad_left + Wn] = resized
+    info = dict(scale=scale, Hn=Hn, Wn=Wn, pad_top=pad_top, pad_left=pad_left)
+    return out, info
+
+
+def _token_valid_mask(side_px: int, unit: int, info: dict) -> np.ndarray:
+    """Boolean (Ts, Ts): a token is valid if fully inside the non-padded region."""
+    ts = side_px // unit
+    top, left = info["pad_top"], info["pad_left"]
+    bottom, right = top + info["Hn"], left + info["Wn"]
+    mask = np.zeros((ts, ts), dtype=bool)
+    for r in range(ts):
+        y0, y1 = r * unit, (r + 1) * unit
+        if y0 < top or y1 > bottom:
+            continue
+        for c in range(ts):
+            x0, x1 = c * unit, (c + 1) * unit
+            if x0 < left or x1 > right:
+                continue
+            mask[r, c] = True
+    return mask
+
+
+def _pool_map_to_tokens(square_map: np.ndarray, unit: int) -> np.ndarray:
+    """Sum a (side, side) map into (Ts, Ts) token scores."""
+    side = square_map.shape[0]
+    ts = side // unit
+    m = square_map[: ts * unit, : ts * unit]
+    return m.reshape(ts, unit, ts, unit).sum(axis=(1, 3))
+
+
+def assemble_canvases(
+    sampled_frames_rgb: List[np.ndarray],
+    frame_ids: List[int],
+    bitmaps: Dict[int, np.ndarray],
+    cfg: AssembleConfig,
+) -> Tuple[List[Image.Image], np.ndarray]:
+    """Select high-bit blocks and pack them into release-contract canvases.
+
+    ``sampled_frames_rgb[i]`` is the decoded RGB frame for ``frame_ids[i]``;
+    ``bitmaps[frame_ids[i]]`` is its DCVC-RT bit map (H/16, W/16). Returns
+    (list of PIL RGB canvases, src_positions int32 [total_patches, 3]).
+    """
+    unit = cfg.unit
+    ts = cfg.token_side()
+    side_px = ts * unit
+    seq_len = len(sampled_frames_rgb)
+
+    # Per-frame square frames + token score maps + validity.
+    frames_sq: List[np.ndarray] = []
+    scores = np.full((seq_len, ts, ts), -np.inf, dtype=np.float32)
+    for i in range(seq_len):
+        fsq, info = _resize_longer_pad_square(sampled_frames_rgb[i], side_px, is_map=False)
+        frames_sq.append(fsq)
+        bm = bitmaps.get(int(frame_ids[i]))
+        if bm is None:
+            bm = np.zeros((max(1, side_px // 16), max(1, side_px // 16)), dtype=np.float32)
+        bm_sq, _ = _resize_longer_pad_square(bm, side_px, is_map=True)
+        vmask = _token_valid_mask(side_px, unit, info)
+        tok = _pool_map_to_tokens(bm_sq, unit)
+        scores[i][vmask] = tok[vmask]
+
+    tokens_per_canvas = ts * ts
+    target_tokens = int(cfg.target_canvas) * tokens_per_canvas
+
+    # --- selection: mandatory full first frame, then global top-k over rest ---
+    selected: List[Tuple[int, int, int]] = []  # (t_orig, th, tw)
+    if cfg.mandatory_first_frame and seq_len > 0:
+        for r in range(ts):
+            for c in range(ts):
+                selected.append((int(frame_ids[0]), r, c))
+
+    remaining = target_tokens - len(selected)
+    if remaining > 0 and seq_len > 1:
+        rest = scores[1:].copy()  # (seq_len-1, ts, ts)
+        flat = rest.reshape(-1)
+        finite = np.isfinite(flat)
+        n_avail = int(finite.sum())
+        k = min(remaining, n_avail)
+        if k > 0:
+            order = np.argsort(-flat, kind="stable")[:k]
+            for idx in order:
+                fi = idx // (ts * ts) + 1
+                rem = idx % (ts * ts)
+                selected.append((int(frame_ids[fi]), rem // ts, rem % ts))
+
+    if not selected:
+        selected.append((int(frame_ids[0]) if frame_ids else 0, 0, 0))
+
+    # Time-sorted packing so per-canvas ``t`` runs are contiguous.
+    selected.sort(key=lambda x: (x[0], x[1], x[2]))
+
+    # Repeat-fill to a whole number of canvases (>=1), never half-pad.
+    if len(selected) < tokens_per_canvas:
+        n_canvas = 1
+    else:
+        n_canvas = min(int(cfg.target_canvas), len(selected) // tokens_per_canvas)
+        n_canvas = max(1, n_canvas)
+    target = n_canvas * tokens_per_canvas
+    if len(selected) < target:
+        selected = selected + [selected[-1]] * (target - len(selected))
+    else:
+        selected = selected[:target]
+
+    # index sampled frame_id -> its square frame (dup ids collapse fine)
+    fid_to_sq = {int(frame_ids[i]): frames_sq[i] for i in range(seq_len)}
+
+    Gh = Gw = ts * cfg.spatial_merge_size  # canvas patch grid
+    sms = int(cfg.spatial_merge_size)
+    images: List[Image.Image] = []
+    all_positions: List[np.ndarray] = []
+
+    for ci in range(n_canvas):
+        canvas = np.zeros((side_px, side_px, 3), dtype=np.uint8)
+        positions = np.zeros((Gh * Gw, 3), dtype=np.int32)
+        chunk = selected[ci * tokens_per_canvas:(ci + 1) * tokens_per_canvas]
+        for p, (t_orig, th, tw) in enumerate(chunk):
+            rc, cc = p // ts, p % ts  # canvas token cell
+            src = fid_to_sq.get(int(t_orig))
+            if src is not None:
+                canvas[rc * unit:(rc + 1) * unit, cc * unit:(cc + 1) * unit] = \
+                    src[th * unit:(th + 1) * unit, tw * unit:(tw + 1) * unit]
+            pr, pc = sms * rc, sms * cc          # canvas patch coords (top-left)
+            sph, spw = sms * th, sms * tw        # source patch coords (top-left)
+            for dy in range(sms):
+                for dx in range(sms):
+                    row = (pr + dy) * Gw + (pc + dx)
+                    positions[row] = (int(t_orig), sph + dy, spw + dx)
+        images.append(Image.fromarray(canvas))
+        all_positions.append(positions)
+
+    src_positions = np.concatenate(all_positions, axis=0).astype(np.int32)
+    return images, src_positions

neural_codec/codec_dcvc_config.py

@@ -0,0 +1,65 @@
+"""Single source of truth for DCVC canvas-selection parameters.
+
+Reads the ``codec.dcvc`` block of ``preprocessor_config.json`` (found by walking
+up from this file to the model directory). Used by ``dcvc_readiness_gen.py`` and
+the readiness pipeline (``process_video_bitcost_readiness.py`` /
+``process_video_bitcost_mv_mask_collage.py``) so ALL selection knobs are
+controlled from the config file — NOT environment variables.
+
+Full schema (defaults = the ``b50`` benchmark config) lives in
+``preprocessor_config.json`` under ``codec.dcvc``.
+"""
+import json
+import os
+import functools
+
+# Baseline defaults (the b50 benchmark config). Used only when a key is absent
+# from preprocessor_config.json's codec.dcvc, so the file stays authoritative.
+_DEFAULTS = {
+    # DCVC engine
+    "qp": 42, "reset_interval": 64, "intra_period": -1, "max_side": 0,
+    # readiness sampling / grouping
+    "num_sampled_frames": 256, "grouping_mode": "readiness",
+    "readiness_sum_threshold_mode": "auto", "group_size": 32,
+    "images_per_group": 4, "patch": 16, "max_pixels": 150000,
+    "min_group_frames": 8, "max_group_frames": 128,
+    "readiness_coverage_bins": 3, "readiness_delta_ratio": 0.05,
+    "bitcost_grid": "sub", "bitcost_pct": 99, "decode_backsearch_max": 16,
+    "canvas_format": "jpg",
+    # selection tuning knobs (our additions)
+    "per_frame_cap_ratio": 1.2,   # spread block budget over more time frames
+    "bottom_atten": 0.5,          # attenuate bottom-edge bit-cost (de-bias overlays)
+    "bottom_band": 0.10,          # fraction of frame height treated as bottom band
+    "threshold_scale": 1.0,       # scale readiness threshold (<1 -> more canvases)
+    "random_select": False,       # random patch baseline (control)
+    "random_seed": 0,
+}
+
+
+@functools.lru_cache(maxsize=1)
+def _load_file():
+    d = os.path.dirname(os.path.abspath(__file__))
+    for _ in range(4):
+        # preprocessor_config.json may sit next to this file (bundled layout) or in
+        # a sibling ``processor/`` subdir (multi-component model repo layout).
+        for p in (os.path.join(d, "preprocessor_config.json"),
+                  os.path.join(d, "processor", "preprocessor_config.json")):
+            if os.path.exists(p):
+                try:
+                    with open(p, encoding="utf-8") as f:
+                        return json.load(f).get("codec", {}).get("dcvc", {}) or {}
+                except Exception:
+                    return {}
+        d = os.path.dirname(d)
+    return {}
+
+
+def get(key, cast=None):
+    """Return codec.dcvc[key] from preprocessor_config.json, else the b50 default.
+    ``cast`` optionally coerces (e.g. float, int, bool)."""
+    v = _load_file().get(key, _DEFAULTS.get(key))
+    if cast is not None and v is not None:
+        if cast is bool and isinstance(v, str):
+            return v.strip().lower() in ("1", "true", "yes")
+        return cast(v)
+    return v

neural_codec/codec_loader.py

@@ -0,0 +1,109 @@
+"""Load precomputed DCVC-RT canvases and turn them into model inputs.
+
+This is the inference-side counterpart of ``precompute_dcvc_rt.py``. It reuses
+the release codec helpers verbatim, so precomputed DCVC-RT assets go through the
+*exact* same downstream the HEVC ``video_backend="codec"`` path uses.
+"""
+
+from __future__ import annotations
+
+import importlib
+import importlib.util
+import json
+import os
+import sys
+from pathlib import Path
+from typing import Optional
+
+import numpy as np
+import torch
+from PIL import Image
+
+_REPO = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))  # Mage-VL-Ported/
+
+
+def _load_release_codec_module():
+    """Import ``codec_video_processing_magevl`` from the release dir."""
+    try:
+        return importlib.import_module("codec_video_processing_magevl")
+    except Exception:
+        path = os.path.join(_REPO, "processor", "codec_video_processing_magevl.py")
+        spec = importlib.util.spec_from_file_location(
+            "codec_video_processing_magevl", path
+        )
+        mod = importlib.util.module_from_spec(spec)
+        sys.modules[spec.name] = mod  # required for dataclass + future annotations
+        spec.loader.exec_module(mod)
+        return mod
+
+
+def load_precomputed(asset_dir: str) -> dict:
+    """Read ``<asset_dir>`` written by precompute -> {images, src_positions, fps}.
+
+    Mirrors ``codec_video_processing_magevl._load_codec_result``.
+    """
+    asset_dir = Path(asset_dir)
+    with open(asset_dir / "meta.json", "r", encoding="utf-8") as f:
+        meta = json.load(f)
+    canvas_files = meta.get("canvas_files")
+    if not canvas_files:
+        canvas_files = sorted(p.name for p in asset_dir.glob("canvas_*.jpg"))
+    images = [Image.open(asset_dir / n).convert("RGB") for n in canvas_files]
+    src_positions = np.load(asset_dir / "src_patch_position.npy")
+    return {
+        "images": images,
+        "src_positions": src_positions,
+        "fps": float(meta.get("fps") or 30.0),
+        "meta": meta,
+    }
+
+
+def build_inputs_from_assets(
+    processor,
+    asset_dir: str,
+    text: str,
+    max_pixels: Optional[int] = None,
+    device: Optional[torch.device] = None,
+) -> dict:
+    """Build a model-ready input dict from precomputed DCVC-RT assets + a
+    chat-templated ``text`` string (containing a ``<|vision_start|>...<|vision_end|>``
+    video span). Returns tensors ready for ``model.generate``.
+    """
+    cm = _load_release_codec_module()
+    payload = load_precomputed(asset_dir)
+
+    if max_pixels is None:
+        # Canvas budget lives in preprocessor_config.json's codec.dcvc (150000),
+        # NOT the image_processor's global max_pixels (the full-frame budget, 4M).
+        try:
+            import codec_dcvc_config as _dc
+            max_pixels = int(_dc.get("max_pixels"))
+        except Exception:
+            max_pixels = 150000
+
+    imgs, src_positions, _ = cm.drop_padding_canvases(payload["images"], payload["src_positions"])
+    if not imgs:
+        raise RuntimeError(f"no usable canvases in {asset_dir}")
+
+    image_data = cm.codec_image_processor_outputs(processor.image_processor, imgs, max_pixels=max_pixels)
+    image_grid_thw = image_data["image_grid_thw"]
+    patch_positions = cm.codec_positions_for_processor(
+        src_positions, image_grid_thw, device=image_grid_thw.device
+    )
+    rewritten = cm.rewrite_text_with_codec_positions(
+        text, patch_positions, fps=float(payload["fps"]), decimals=1
+    )
+
+    enc = processor.tokenizer(rewritten, return_tensors="pt")
+    out = {
+        "input_ids": enc["input_ids"],
+        "attention_mask": enc["attention_mask"],
+        "pixel_values": image_data["pixel_values"],
+        "image_grid_thw": image_grid_thw,
+        "patch_positions": patch_positions,
+    }
+    if device is not None:
+        for k, v in out.items():
+            if isinstance(v, torch.Tensor):
+                out[k] = v.to(device)
+    return out

neural_codec/codec_tools/codec_patch_gop/__init__.py

@@ -0,0 +1,11 @@
+#!/usr/bin/env python3
+# -*- coding: utf-8 -*-
+
+"""
+Codec-aware patch dataset generator with GOP-based sampling.
+
+This package provides tools for extracting patches from videos using
+motion vector and residual-based scoring.
+"""
+
+__version__ = "1.0.0"