Add files using upload-large-folder tool

LICENSE

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

README.md

@@ -0,0 +1,239 @@
+---
+license: mit
+base_model: inclusionAI/Ming-Image-0.1-Design
+base_model_relation: quantized
+pipeline_tag: text-to-image
+tags:
+- text-to-image
+- rocm
+- amd
+- strix-halo
+- gfx1151
+- int8
+- design
+- rgba
+---
+
+# Ming-Image-0.1-Design — ROCm build (AMD Strix Halo, gfx1151) · INT8 MLLM · paired with Ling-3.0-flash-VL
+
+[inclusionAI/Ming-Image-0.1-Design](https://huggingface.co/inclusionAI/Ming-Image-0.1-Design) (text-to-image for
+UI, posters and infographics, RGBA output) made to run on AMD ROCm, with the routed experts of its 17B-parameter
+MoE language model stored as weight-only INT8, and wired to the prompt enhancer its model card names: **Ling-3.0-flash-VL**
+(served from our [ROCmFP4 build](https://huggingface.co/kingjones777/Ling-3.0-flash-VL-MTP-ROCmFP4-GGUF)).
+
+Everything below was measured on one AMD Ryzen AI Max+ 395 (Radeon 8060S, gfx1151) — see
+[Reproduction](#reproduction). Nothing here was run on CUDA.
+
+## Results (1024 × 1024, 12 steps, cfg 1.0, seed 42, four prompts)
+
+| prompt (1024², 12 steps) | BF16 · upstream code | BF16 · contiguous attention | INT8 · contiguous attention | INT8 · fast attention |
+|---|---:|---:|---:|---:|
+| four-seasons cabin (upstream's example prompt) | 339.9 s | 211.1 s | 181.1 s | 77.8 s |
+| banking-app screen (Ling rewrite) | 501.8 s | 270.0 s | 259.7 s | 88.5 s |
+| jazz-night poster (Ling rewrite) | 355.5 s | 221.4 s | 219.9 s | 79.4 s |
+| water-cycle infographic (Ling rewrite) | 606.7 s | 312.0 s | 250.0 s | 100.3 s |
+| **mean** | 451.0 s | 253.6 s | 227.7 s | 86.5 s |
+| speed-up vs upstream code | 1.00× | 1.78× | 1.98× † | 5.21× |
+
+| configuration | PyTorch peak allocated | Ming's own GTT peak | box |
+|---|---:|---:|---|
+| BF16 · upstream code, all components resident | 58.7 GiB | 74.5 GiB | Ling stopped |
+| BF16 · contiguous attention | 58.9 GiB | 74.5 GiB | Ling stopped |
+| BF16 · contiguous · `--release-mllm-after-conditioning` (cabin) | 47.2 GiB | — (baseline not settled) | Ling stopped |
+| **INT8 · contiguous · `--release-mllm-after-conditioning`** | **33.6 GiB** | **35.8 GiB** | Ling-3.0-flash-VL resident (67.5 GiB) → box peak 103.3 GiB of 128 GiB GTT |
+| INT8 · fast attention · release | 33.6 GiB | 35.1 GiB | Ling resident |
+
+- **Contiguous attention is the big one:** 1.78× faster end to end, and the images are **byte-identical** to the upstream code path (4/4 prompts plus the `--release` run, compared with `cmp`).
+- † **INT8** runs the same BF16 diffusion transformer, so it does not change the step time: its 227.7 s mean vs 253.6 s for BF16 with the same attention reflects when the runs happened (later, with Ling resident), not the quantization. What INT8 does cost is the conditioning pass, where the weights are dequantized on the fly: 2.4–4.7 s vs 1.1–2.0 s in BF16. INT8 exists for memory: with `--release-mllm-after-conditioning` Ming needs 35.8 GiB at its peak, which fits next to the resident Ling-3.0-flash-VL (67.5 GiB) on a 128 GiB box — BF16 with the upstream code needed 74.5 GiB and did not. Loading takes 30–100 s; the slow end is a cold page cache (first load after the files were written or after Ling was restarted).
+- **Fast attention** (`--attention-bf16-reduction`) is opt-in: 5.21× vs upstream, at the fidelity cost shown below.
+- The very first image on a fresh box is slower once: the reference cabin took 608.4 s the first time and 339.9 s warm (MIOpen tunes the VAE's 3-D convolutions and caches the result in `~/.cache/miopen`).
+
+### Fidelity against the BF16 original
+
+Inference on this box is **deterministic**: the same prompt and seed produced a byte-identical PNG twice, in BF16
+and in INT8, so every difference below is caused by the quantization (or by the fast-attention option), not by
+run-to-run noise.
+
+| prompt | INT8 SSIM | INT8 PSNR | cond cos (query / VLM tokens) | cond rel L2 (query / VLM) | INT8 fast SSIM | INT8 fast PSNR |
+|---|---:|---:|---:|---:|---:|---:|
+| four-seasons cabin (upstream's example prompt) | 0.9345 | 22.90 dB | 0.98169 / 0.99672 | 0.1931 / 0.0810 | 0.9150 | 22.13 dB |
+| banking-app screen (Ling rewrite) | 0.9666 | 23.76 dB | 0.98638 / 0.99623 | 0.1655 / 0.0868 | 0.9624 | 22.75 dB |
+| jazz-night poster (Ling rewrite) | 0.9450 | 23.31 dB | 0.98874 / 0.99587 | 0.1504 / 0.0909 | 0.9456 | 23.55 dB |
+| water-cycle infographic (Ling rewrite) | 0.8486 | 17.49 dB | 0.99061 / 0.99500 | 0.1378 / 0.1000 | 0.8439 | 17.42 dB |
+
+`cos` / `rel L2` compare the conditioning tensors the diffusion transformer receives (the MLLM output — the only
+part that is quantized). SSIM is windowed 7×7 on luminance; PSNR over RGB.
+
+## Run it
+
+```bash
+hf download kingjones777/Ming-Image-0.1-Design-ROCm-INT8 --local-dir Ming-Image-0.1-Design-ROCm-INT8
+cd Ming-Image-0.1-Design-ROCm-INT8
+# caption -> Ling-3.0-flash-VL rewrites it into the layered JSON prompt -> Ming renders it (one image per process)
+bash code/generate_paired.sh --model . --base-url http://127.0.0.1:8090/v1 \
+    "A minimalist concert poster for a jazz night called \"Blue Hour\"" --resolution 1024 --output-dir out \
+    -- --device-map none --release-mllm-after-conditioning        # add --attention-bf16-reduction for fast mode
+```
+
+`PYTHON=/path/to/python` selects the interpreter (it needs a ROCm build of PyTorch and
+[`code/requirements-rocm.txt`](code/requirements-rocm.txt)); `--pe-model` names the chat model your endpoint serves.
+This is the command shape our end-to-end run used (see [Samples](#samples)); 1024² is the resolution we measured —
+`infer.py`'s own default is 2048², which we did not run.
+
+`generate_paired.sh` calls `pe_ling.py` (standard library only) against an OpenAI-compatible endpoint
+(default `http://127.0.0.1:8090/v1`, our Ling-3.0-flash-VL llama-server seat), validates the rewrite against the
+schema the upstream system prompt demands (one retry with the errors), then runs `infer.py`. You can also pass
+your own JSON prompt straight to `infer.py --prompt prompt.json`.
+
+The code is in [`code/`](code/): upstream `inclusionAI/Ming-Image` at `62c6072` plus the changes below
+([`code/rocm.patch`](code/rocm.patch)).
+
+## What changed for ROCm, and why
+
+1. **`transformer_engine` removed.** The vision tower imported NVIDIA Transformer Engine (CUDA-only) for a
+   single `te.RMSNorm`; it is now a plain RMSNorm with the same parameter name (`weight`), so all 65 vision norm
+   tensors in the checkpoint load unchanged. A dead import in `modeling_bailing_moe_v2.py` went too.
+2. **`--attn-implementation eager` now reaches the towers.** `BailingMM2Config` declared no `sub_configs`, so
+   transformers never copied the chosen attention implementation into the vision and language configs; their
+   `flash_attention_2` defaults raised `ImportError` at model construction on any box without flash-attn. Declaring
+   `sub_configs` fixes it. (Upstream's `--validate-only` cannot catch this — it never builds the model.)
+3. **Contiguous attention inputs.** The diffusion transformer handed PyTorch's attention permuted views of
+   `[B, L, H, D]` tensors. On gfx1151 the only working SDPA kernel is the math one, and with those strides its
+   fp32 GEMMs fall onto an 8×8×8 macro-tile: 644 ms per attention call at the cabin prompt's length
+   (L = 5,759) versus 363 ms for the same tensors made contiguous, bit-identical output. The profiler put
+   SDPA-math at 87.6% of all GPU time before the fix. `diffusion/transformer.py` now passes
+   contiguous tensors whenever diffusers' default native backend is active.
+4. **INT8 MLLM — routed experts only.** Weight-only, per-output-channel symmetric INT8 (fp32 scales) for the
+   256 routed experts of the language model's 19 MoE layers — 14,592 Linears, which hold
+   almost all of its weights. Everything that runs on every token stays byte-identical BF16: attention
+   (`query_key_value`, `dense`), the shared experts, layer 0's dense MLP, the routers (`gate`, `image_gate`,
+   `audio_gate`), embeddings, `lm_head`, norms and the vision tower (691 tensors).
+   `mllm/` goes from 34.00 GB to 18.76 GB. Relative weight error
+   ‖W − Ŵ‖/‖W‖: mean 0.00833, p99 0.01035, max
+   0.01412. `quant/quantize_stream.py` writes it shard by shard without building the model
+   (6.7 min; peak RSS 17.84 GiB, sampled on an earlier run of the same
+   tool); `quant/load_int8.py` builds the model on the meta device and loads the INT8 shards straight onto the GPU, so
+   BF16 weights for the quantized layers never exist in memory. The scales stay fp32 through `.to(bfloat16)`.
+   *Measured alternative:* quantizing the always-on Linears too saves another 0.30 GB but raised the mean
+   conditioning error (rel L2, VLM tokens) from 0.090 to 0.108 and changed the cabin prompt's
+   surround from scenery to white (SSIM 0.699); mean SSIM over the four prompts 0.877 vs 0.924.
+5. **Connector stored as bf16.** It shipped as float32 (6.17 GB) but `infer.py` always loads it as bf16; storing
+   it pre-rounded halves the download and every tensor equals `fp32.to(bfloat16)` exactly (all 338 checked).
+6. **`--release-mllm-after-conditioning`** (opt-in, one image per process): the language model, vision tower and
+   connector are only needed for the ~1–2 s conditioning pass, so they are freed before the 12 diffusion steps.
+   PyTorch peak allocated on the cabin prompt (BF16): 55.6 → 47.2 GiB, byte-identical image.
+7. **`--attention-bf16-reduction`** (opt-in): lets the math SDPA kernel stay in bf16 instead of upcasting to fp32 —
+   363 → 115 ms per call, attention error vs an fp32 reference 1.658e-03 → 5.032e-03
+   (rel L2, random inputs). Image-level cost is in the fidelity table.
+
+## Things that do not work on gfx1151 (measured, so you don't have to)
+
+- **Only the math SDPA kernel runs.** AOTriton's efficient and flash kernels are unavailable by default
+  (`UNAVAILABLE: No available kernel. Aborting execution.`).
+- **Do not set `TORCH_ROCM_AOTRITON_ENABLE_EXPERIMENTAL=1`**, even though PyTorch's warning suggests it: the shipped
+  `amd-gfx11xx` AOTriton kernels then get selected and fail (`UNAVAILABLE: HIP error: invalid argument`), and a full
+  generation crashes with the same `HIP error: invalid argument`.
+- **The first image on a fresh box is slower once.** The same cabin prompt took 608.4 s the
+  first time and 339.9 s warm; MIOpen tunes the VAE's 3-D convolutions on first use and its cache
+  (`~/.cache/miopen`) grew during that first run.
+- **Inference is deterministic here, BF16 and INT8:** same prompt + seed → byte-identical PNG (both checked), which
+  is what makes the fidelity numbers above exact rather than statistical.
+
+## Pairing with Ling-3.0-flash-VL
+
+Ming-Image's text-to-image quality depends on a structured, Figma-like JSON prompt (canvas settings +
+positioned layers with colours). Upstream ships the rewriter system prompt (`assets/t2i_rewriter_system_prompt.txt`) and names Ling-3.0-flash-VL as the model to run it.
+`pe_ling.py` sends that prompt verbatim plus your caption to any OpenAI-compatible endpoint, extracts the JSON, validates it, and retries once with the validation errors if needed.
+
+Measured against our Ling-3.0-flash-VL seat (llama-server, `ling-3.0-flash-vl-mtp-halo-STRIX_LEAN`, same box) — all three valid on the first attempt:
+
+| caption | rewrite time | layers |
+|---|---:|---:|
+| banking-app screen (Ling rewrite) | 45.9 s | 7 |
+| jazz-night poster (Ling rewrite) | 34.8 s | 8 |
+| water-cycle infographic (Ling rewrite) | 49.8 s | 14 |
+
+The exact captions and the JSON Ling returned are in [`samples/`](samples/).
+
+## Samples
+
+INT8 build, `--release-mllm-after-conditioning`, 1024², 12 steps, cfg 1.0, seed 42 — exactly the runs in
+the tables above. Full-resolution RGBA PNGs and the JSON prompts are in `samples/`.
+
+**cabin upstream** — upstream's own structured prompt (`assets/t2i_four_seasons_cabin_prompt.json`), no rewrite · [prompt JSON](samples/cabin_upstream.json)
+
+![cabin_upstream](samples/cabin_upstream.png)
+
+**ui banking** — caption → Ling-3.0-flash-VL: *A mobile banking app home screen: a balance card at the top, a recent transactions list, quick-action buttons for send, pay and top up, and a bottom navigation bar. Clean modern fintech style.* · [prompt JSON](samples/ui_banking.json)
+
+![ui_banking](samples/ui_banking.png)
+
+**poster jazz** — caption → Ling-3.0-flash-VL: *A minimalist concert poster for a jazz night called "Blue Hour" on Friday, October 3, 8 PM at The Lantern Room, with a saxophone silhouette over a deep blue gradient.* · [prompt JSON](samples/poster_jazz.json)
+
+![poster_jazz](samples/poster_jazz.png)
+
+**info water** — caption → Ling-3.0-flash-VL: *An infographic that explains the water cycle in four labeled stages - evaporation, condensation, precipitation, collection - with arrows and simple flat icons.* · [prompt JSON](samples/info_water.json)
+
+![info_water](samples/info_water.png)
+
+### End to end on the box
+
+`ming-paired` (our box's wrapper around `generate_paired.sh`, not part of this repo) with Ling-3.0-flash-VL and two other model servers running: it paused 2 of them for the render and restarted them afterwards (halo-bonsai8b, halo-bonsai). Ling rewrote the caption in 48.3 s (9 layers); the whole command took 323 s.
+
+*A landing page hero section for a coffee subscription service called Daily Grind: the headline Fresh beans every Monday, a short tagline, a Start your subscription button, and a photo of latte art on the right. Warm earthy palette.* · [prompt JSON](samples/e2e_daily_grind.json)
+
+![end to end](samples/e2e_daily_grind.png)
+
+## Files
+
+| path | size | what |
+|---|---:|---|
+| `mllm/` | 18.78 GB | language model + vision tower; the 14,592 routed-expert Linears are **INT8** (`int8_manifest.json` lists them) |
+| `transformer/` | 12.31 GB | diffusion transformer, BF16, unchanged |
+| `connector/` | 3.09 GB | Qwen2 connector, **stored as bf16** (upstream ships fp32; runtime identical) |
+| `vae/` | 253.82 MB | VAE (4-channel RGBA), unchanged |
+| `mlp/` | 124.84 MB | conditioning MLP, unchanged |
+| `scheduler/` | 173 B | flow-matching scheduler config, unchanged |
+| `code/` | 2.31 MB | patched inference code + `rocm.patch` + tools |
+| `samples/` | 4.56 MB | the sample images and prompts shown above |
+| **total** | **34.56 GB** | upstream: 52.88 GB |
+
+## Reproduction
+
+```
+upstream : inclusionAI/Ming-Image-0.1-Design  revision 1cd7fac3b0dcb54196fe2cd12b80da09edf8fcf4
+code     : inclusionAI/Ming-Image @ 62c6072e1ff15af83f7c4963a0a1954c1424e80e + code/rocm.patch  (branch rocm-halo @ f986f7a (upstream inclusionAI/Ming-Image 62c6072 + 3 commits); rocm.patch sha256 03fe16f8ed566bf99ec654caa687c7e659f489df6d077cfb7fc7aee134ebc777)
+python   : 3.13.5 · torch 2.10.0 (HIP 7.13.99004) · transformers 4.57.1 · diffusers 0.36.0 · accelerate 1.13.0 · safetensors 0.8.0
+box      : amd-halo · AMD RYZEN AI MAX+ 395 w/ Radeon 8060S · 125 GiB RAM visible, 128 GiB GTT · ROCm 7.13.0 · kernel 6.18.35+rex+2-amd64
+power    : platform_profile=balanced · governor=powersave · GPU 83–107 W at 100% busy (step probe, 0.5 s samples)
+env      : nothing set; in particular TORCH_ROCM_AOTRITON_ENABLE_EXPERIMENTAL is NOT set (it crashes gfx1151)
+date     : 2026-09-22 to 2026-09-23 (runs crossed midnight, America/Chicago)
+```
+
+Build the INT8 package from the upstream download:
+
+```bash
+python code/quant/quantize_stream.py <upstream>/mllm <package>/mllm \
+    --exclude '\.attention\.|\.shared_experts\.|^model\.model\.layers\.0\.mlp\.'   # routed experts only
+python code/tools/convert_connector.py <upstream>/connector <package>/connector # fp32 -> bf16, proven exact
+# transformer/ vae/ mlp/ scheduler/ LICENSE are the upstream files, unchanged
+python code/tools/verify_package.py <upstream> <package>                        # the checks behind this card
+```
+
+Measure (one process per image, the pairing configuration):
+
+```bash
+cd code && PYTHONPATH=. python tools/ming_bench.py --prompts <prompt.json> --out <dir> -- \
+    --model <package> --task text-to-image --resolution 1024 --device-map none \
+    --attn-implementation eager --release-mllm-after-conditioning        # add --attention-bf16-reduction for fast
+python tools/fidelity_compare.py <bf16-dir> <int8-dir> --json fidelity.json
+```
+
+BF16 references ran with Ling stopped (the box cannot hold both); every INT8 run had Ling-3.0-flash-VL resident.
+GPU memory is the amdgpu `mem_info_gtt_used` peak sampled every second, minus a baseline taken after GTT settled.
+
+## License
+
+MIT, same as the original. Model weights, architecture and inference code © 2026 inclusionAI
+([LICENSE](LICENSE)). The INT8 quantization, the ROCm changes and the pairing scripts are ours; they are listed in
+this card and in `code/rocm.patch`.

SHA256SUMS

@@ -0,0 +1,94 @@
+b772c54f1f150a680e7c083a92431c518cf62821a407f6cccdf53f6a07bcf702  LICENSE
+b1a5d04726679cfa37262199478de1c0c1735b2d5ac34ede7ac49c72fd35e10a  README.md
+561cdddbafb1cbabe417d2ad2e90ec1dd28a7c131aaead2cc9ded59e9e27b685  code/.gitignore
+b772c54f1f150a680e7c083a92431c518cf62821a407f6cccdf53f6a07bcf702  code/LICENSE
+bef9e77db8ade9d625b80869a43319beb9e8353927ea279d095f69a2fd7c9dcc  code/README.md
+7416e29e85a7997d7550ed455226473a0a95217faa3523aa9ab2125055fca4ad  code/REVISION
+a763ccc5a39740b19fc610581f14041914739feeb83c28b712a72a8a896cc978  code/assets/layer_decompose_5layers.txt
+e445a68e4969c1f2734f00a4223559cf58904c2299019346816571db9dee35c3  code/assets/layer_samples/card_making_decomposition.png
+a52347559312b7ddbf435d042bdc32c83efe8750d3071580e6a8eceee2412112  code/assets/layer_samples/card_making_input.png
+3e06063e76ffd32b80ad6989352f47006fd7f936eb808ff563ec2392898c1884  code/assets/layer_samples/card_making_prompt.txt
+387468da6db92aeb773e1776434a5c9b723de7d2f059e3304f4de3a9c862eaf6  code/assets/t2i_four_seasons_cabin_prompt.json
+62b0c7c0380fdccf34abf490d1da6ad971a2e63d82f99144c5cd4c9a8c80480c  code/assets/t2i_rewriter_system_prompt.txt
+897fd9b0541efcc15e5734c9afd943c37a7b29b37d465556a43b10f5815d05cd  code/bailingmm_utils.py
+7d49dc16d2c31c05ea25e76089c635a9b56c629a6a416c88d8fba4e05a13c0ab  code/configuration_bailing_moe_v2.py
+ac329fb7c8959b16b68b774693672d1f991d7fe4489d484dc72b4e389c10706f  code/configuration_bailingmm2.py
+e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855  code/diffusion/__init__.py
+4013c3e09e3684d2ca8bbf1f98144a8ebab2a9408641831b1661d89b4a0aa106  code/diffusion/autoencoder_kl_qwenimage.py
+a31f006c123febdcfdfeb92b1b55f2716b154a98e3c3e13bd7bf825148e9901e  code/diffusion/generator.py
+511f76e9292f82076200b4c493a481782df9a8d3eea38ecc5c453b75d3faebe0  code/diffusion/padding.py
+ab42fc223ddff73f854d8b6595e1207ab184248d32151adbd847c0e9974f2343  code/diffusion/pipeline.py
+287522eea36b6538e5e5dd38a9cee6a21d47053f9868d79457083cae281f0293  code/diffusion/transformer.py
+b373bf064d60b7dd493b44d8e54fa1d939fc134fc63766c3f9e9d7d405cf31d6  code/examples/profiles/generation_edit.json
+51f773286518ce8ca9034fb12bf7b527ce52a4fcc09f975e9df04d968276a244  code/examples/profiles/layer_decompose.json
+99e45db68ad2a119d07feb5efe2da3f788ea33ad5997b7a0c8406c9a6e00d5e8  code/generate_paired.sh
+b54ff75e3452de090d29c0f1bc0ce082c510a77405808ea6a59aebef6a691142  code/image_processing_bailingmm2.py
+e39ee57066c183f8411ae3038f8ac491a82a59bde8a9e75ac41f4ed2bb368f92  code/infer.py
+84f407d5e323c7261792492ed400bb23e3a189cb3edfa017e6c854263a8ee490  code/inference_profile.py
+02bae9078a230d2414861368894eee027403f324a94b95b07272d79a75a0bf88  code/mllm_device_map.py
+9da007026e814d17b1bbb8577a03fe70bc4ab6f43854bf6b7f5fba5f846fd429  code/modeling_bailing_moe_v2.py
+9816b917c9661020160f1cd6f4b895f8a9c6dfa42a163db7751f9aad29bab00b  code/modeling_bailingmm2.py
+81b0be70fbe32ca6ac34110d58653a18462c0c73bf797579e8e6be56472be156  code/modeling_utils.py
+aa4d34184c0d513de0ce6948bc0c0d17195b521ecfa1de606eb7f437dab4739b  code/pe_ling.py
+96988388c47935b2f54eeb14ca10dab16044dd7dcf0ca67d2c284c46e49f4861  code/processing_bailingmm2.py
+ba6414475da49f4696e414eb48b34da71108519bf5b538bd51acd780da8f3f5e  code/quant/__init__.py
+05be56bc8b26fe63825b0303e9f830a898e95c21b23847cbe7eb8e40616dcc41  code/quant/int8_linear.py
+d80aa13e407519078059f847ce5f0c5117ce67606a2766be65809e3a44999b53  code/quant/load_int8.py
+a0466c47611eed8ed3fdd627f1c0cc00ed9ee6760407b3cf4011e24c6c420870  code/quant/quantize_stream.py
+3eb9c6eab13676931f1a60856249e0cdebb80e4d55b8d4772c4fe862c69bc50b  code/quant/test_int8.py
+08bc15e548fc07d55626b2707b69253c7cb9e69ffa490ba00d215b6740126ba5  code/qwen2_5_vit.py
+022e5312f4c0eb2be01c03abbdceb905e0fb9d7005473b51d337f4b755c6d807  code/requirements-rocm.txt
+f82a235b7e7da97b59d4121473e2b002bfef9960344050913e8c75f21e6d1c6e  code/requirements.txt
+03fe16f8ed566bf99ec654caa687c7e659f489df6d077cfb7fc7aee134ebc777  code/rocm.patch
+e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855  code/tests/__init__.py
+d5d02c0618be0e4e13dc74f14bc8376396d0a9ed3c06152d3a78ed35485de019  code/tests/assets/smoke_input.png
+0cc02c16a9f34b3700bef5c1e04997b71e3cf89696c1ffe6daf13151dd5c4451  code/tests/test_infer_cli.py
+1d9f5d94101fadb72416320bf073284e5cfc4b30e3ad539b606ff4485c83dfa9  code/tests/test_inference_profile.py
+35edb0b9a378cd7c7be0e0488b9f6ae66ac01b15cab65a71df6b993c21f17af0  code/tests/test_inference_smoke.py
+21648462c682f9cd03eddb178907104ffa1e0b6a4e79fb760f755385aea4abf8  code/tests/test_mllm_device_map.py
+4f2d6347f360399392257caa22da310d2a170f570eb5da1dff4cde63feee3a62  code/tests/test_padding.py
+9c9edaf036544e209c9b6195233fcb3b0baa3752dc1f59923b058bb654ffc5b8  code/tests/test_runtime_precision.py
+719b7dbdb29f18d74f0a7fae070383d2d612173f8c5c6d42d5a0496f3071ae98  code/tokenization_bailing.py
+4f9040a6b8b7917860c4df4ecc9dd56fafead0643458221f6b0179255f03597b  code/tools/convert_connector.py
+47057cbfc5224951617b55c2e6600867145d57c0571c62857f01f1990b96f9f2  code/tools/fidelity_compare.py
+a247b28869d96e3323f47b573fcb0d432af86f67e41db29a37dfbff92ffe1bcf  code/tools/ming_bench.py
+f5d885aefc9ec44476e194d32bf614a58d0a20eff15dea3e005cb79994e043ba  code/tools/sdpa_layout.py
+1c54fd63d74608da8f61ef3012cf9420627b7d9d01b337f45653d08a2abe11cf  code/tools/step_probe.py
+7f4ddb78131e9de22c30e2522cf1b667f7be5b2aeceafbc63cf67d3dcc94e00c  code/tools/verify_package.py
+cc743612957ccbb78d9812e060e87fe096ea1c0f2ad63519b5eb49f249724f95  connector/config.json
+cd5866254bd02b52464b75af06818afa882298a64d30a3819343a4fa4ed16ca9  connector/generation_config.json
+1522edb909db45f90d3628aa5cb03805fba81aca1950dfe9385b057c6b636a47  connector/model.safetensors
+f797c39adcd1624e5bfc94f60d2a86095dc40978f8cf230fcd11ada877c597ea  mllm/config.json
+1472f23b9a1bc48af3298875a075fa667ffc2c147c33acad6c3ae6ba77bc13f7  mllm/int8_manifest.json
+fd19afe5acf56675a154a2ea5f79a7844514d7500dd4efe6b7e771524fdf9c52  mllm/model-00001-of-00004.safetensors
+a1f3472a28da6607565cd2e31324c25b17499b61620666d090d3f6a434ebdfa7  mllm/model-00002-of-00004.safetensors
+d9be6de14c1f0fabaab1941b895bee49e97a015a469f7da31e622f6f10ba368e  mllm/model-00003-of-00004.safetensors
+c770d343ff03238ec3f3a6c2cb594c7efe286eca5a790612bbe72c4e5155783c  mllm/model-00004-of-00004.safetensors
+7dcaa09cd8129511cf675dc5ad832ebdff251f4a3e155f4ebeac272977d4fa30  mllm/model.safetensors.index.json
+e63066dde0329df1642c8dd9a269d613307804084d12e0787ccc887839339091  mllm/preprocessor_config.json
+8c383d444e75b12f2643b4695bece7f3dd1d43725d8a26deeb019aa40a8f891d  mllm/special_tokens_map.json
+e7ff01708d504f7bf4dbf7f5815adde57bab9a40e7f563ab6ad1acace4464917  mllm/tokenizer.json
+861c20bea04c2057362f226658ad9209978ff294b8a95591a0d32f69dbd3fc28  mllm/tokenizer_config.json
+682c34c6c39173b396aa24f444098e495c7b731247d935552dc7d6714154adce  mlp/config.json
+53e47c1ec942749f07025a41be148f5e8364d8587aa2d954969716abe521b77b  mlp/model.safetensors
+387468da6db92aeb773e1776434a5c9b723de7d2f059e3304f4de3a9c862eaf6  samples/cabin_upstream.json
+1c9aaa45455f78c289bc41d8aa175cc4e01c46b1c57e109be043b9f78de57bc1  samples/cabin_upstream.png
+cf9acb660827aad5de550c74cc825508f1b46db59dab42af07f4c0f82830f303  samples/e2e_daily_grind.caption.txt
+f5c561cce8c455622a8342331b0eaa57d4e276041266ba77dee63332b0b904da  samples/e2e_daily_grind.json
+654d8f20a74959a987db7ea306c03618762f07a103fefdef39305818db2e91e8  samples/e2e_daily_grind.png
+e86c214427617b69d1ac4bf126120edbb75c1a6e70fe37669d908997abc8c4f4  samples/info_water.json
+e9906c15e21da79428bd3dd0e8dbb37eca28d61ffbef49aa3b6801ade1f4d8e6  samples/info_water.png
+864b07f559fad0bbc38b5dd3ff831debf35a4fe82a33ab39f8d82fdbff8625f8  samples/poster_jazz.json
+1e7083ba5d40fb97b4269148b127f780a21aae6c7ebc534adb3c1ffc6fd7b3b5  samples/poster_jazz.png
+6c9ddda781fa033068da8dcc414996f6d592efae322d8b2d45b24d157c34215e  samples/ui_banking.json
+264207a5c73e24dd8292a85a0ad1fb4216dec03998ac98e81da7b85b27754e5b  samples/ui_banking.png
+79bd98cddf1a60d365f3072d713ffad832239f57c837bba0e9da3cff40bfafc5  scheduler/scheduler_config.json
+9a64223dd9858f5a1120f4593534173b57741bb8ec4bac53f1636058c820abcd  transformer/config.json
+1488ae867d45a4c184fa0df9fde68d8f7bd9c2dea52b611b77fc6bbd685d9d70  transformer/diffusion_pytorch_model-00001-of-00005.safetensors
+6f9441923d41408c74bd3761a34eb4db2c583fc08d5785a2c7f2260c8e17089e  transformer/diffusion_pytorch_model-00002-of-00005.safetensors
+45cbd9a557c20c7ae5d676049b301f4856046be2f802f3e9e7de477c77a35b8c  transformer/diffusion_pytorch_model-00003-of-00005.safetensors
+73c70763356a57d5d3cd04e91e714cb0b86081a3e922b594659b58f2ab49bc90  transformer/diffusion_pytorch_model-00004-of-00005.safetensors
+0c178789512c7f5d45a7b0c1c561f868704cd0fa77d4cb32f8d7ea06b2707781  transformer/diffusion_pytorch_model-00005-of-00005.safetensors
+21f362456d5f8394b516215521d7de8d44f8353e170043f82d902913eef5a2cc  transformer/diffusion_pytorch_model.safetensors.index.json
+8db67ee46bb5810ed819d698a8acfc95936d157fe68b66c9ed16b2d5ea04413e  vae/config.json
+06520463778e64dca1039c7447890065ee220bc408d15412182e5c3e06f304f1  vae/diffusion_pytorch_model.safetensors

code/.gitignore

@@ -0,0 +1,10 @@
+__pycache__/
+*.pyc
+*.pyo
+*.egg-info/
+.venv/
+venv/
+env/
+outputs/
+cmd.log
+.DS_Store

code/LICENSE

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

code/README.md

@@ -0,0 +1,279 @@
+# Ming Image 0.1 Design
+
+Ming-Image-0.1-Design is an open-source series for visual-design generation
+and editable layer decomposition.
+
+The series includes two 6B-parameter models:
+
+- [Ming-Image-0.1-Design](https://huggingface.co/inclusionAI/Ming-Image-0.1-Design)
+  generates complete visual designs for UI, infographics, posters, and
+  text-rich compositions.
+- [Ming-Image-0.1-Design-Layer](https://huggingface.co/inclusionAI/Ming-Image-0.1-Design-Layer)
+  decomposes flattened design images into independently editable transparent
+  layers.
+
+![Ming-Image-0.1-Design on the Artificial Analysis UI/UX Design leaderboard](assets/ming-image-design-ui-ux-leaderboard.webp)
+
+## Gallery
+
+### Text-to-image
+
+![Text-to-image showcase](assets/model_cards/design_showcase.webp)
+
+### Transparent-background text-to-image
+
+![Transparent-background text-to-image showcase](assets/model_cards/design_transparency_showcase.webp)
+
+### Layer decomposition
+
+![Six-layer card decomposition showcase](assets/model_cards/layer_showcase.webp)
+
+![Layer-decomposition gallery](assets/model_cards/layer_gallery.webp)
+
+![Layer-decomposition results on the Crello test set](assets/model_cards/layer_performance.webp)
+
+## Requirements
+
+- Python 3.10 or newer;
+- CUDA-capable PyTorch for full model inference;
+- a local checkpoint directory or a Hugging Face Hub repository ID.
+
+Install the runtime dependencies in a clean environment:
+
+```bash
+pip install -r requirements.txt
+```
+
+`flash_attention_2` is optional. The CLI defaults to `eager` attention because
+the LLM only implements eager and FlashAttention 2 attention classes; the
+diffusion transformer always uses PyTorch SDPA internally. Selecting
+`--attn-implementation sdpa` fails closed at load time.
+
+## Inference
+
+The `--model` argument accepts either a local directory or an HF Hub repo ID.
+Hub models are resolved to one immutable local snapshot before any component is
+loaded. Use `--revision` to pin a branch, tag or commit. The examples use the
+published Hub IDs; replace them with local checkpoint directories for offline
+inference.
+
+Sampling defaults and public resolution buckets are task-specific:
+
+| Task | Steps | CFG | Resolution buckets | Default / recommended |
+| --- | ---: | ---: | --- | ---: |
+| Text-to-image | 12 | 1.0 | 1024, 2048 | 2048 |
+| Layer decomposition | 12 | 2.0 | 512, 1024 | 1024 |
+
+Pass `--steps` or `--cfg` to override either value explicitly.
+`--resolution` is optional. A supplied positive integer snaps to the nearest
+bucket supported by the selected task, with ties going to the smaller bucket.
+For faster layer decomposition, explicitly pass `--resolution 512`; use 1024
+for the recommended output quality. Text-to-image output is square at the
+selected bucket. Layer decomposition preserves the reference image's aspect
+ratio while selecting a predefined working size from the effective 1024 or
+512 bucket.
+
+The default and minimum validated deployment is **one GPU with at least 80 GiB
+of memory**, running in BF16. This configuration passes both model families
+end-to-end, and all demos below use it by default.
+
+```bash
+export CUDA_VISIBLE_DEVICES=0
+```
+
+Select `--attn-implementation flash_attention_2` when FlashAttention 2 is
+installed; the portable CLI default remains `eager`.
+
+For both tasks, `--prompt` accepts either raw text or a path to a prompt file.
+Use `--validate-only` to check the checkpoint contract and task combination
+without loading model weights:
+
+```bash
+python infer.py --model inclusionAI/Ming-Image-0.1-Design --task text-to-image \
+  --prompt "A red circle on a white background" --validate-only
+```
+
+### Text-to-image demo
+
+This demo renders a fixed cabin across spring, summer, autumn, and winter. It
+uses the [exact structured JSON prompt](assets/t2i_four_seasons_cabin_prompt.json)
+without duplicating that long input in the README.
+
+```bash
+python infer.py \
+  --model inclusionAI/Ming-Image-0.1-Design \
+  --task text-to-image \
+  --prompt assets/t2i_four_seasons_cabin_prompt.json \
+  --attn-implementation flash_attention_2 \
+  --resolution 2048 \
+  --output-dir outputs/t2i
+```
+
+### Text-to-image prompt rewriting
+
+An instruction-following VLM can turn a short caption into a precise,
+layout-structured JSON description: Figma-style layers ordered back to front,
+with exact coordinates, hierarchy, color specs, and every rendered string
+quoted verbatim and owned exactly once. Unlike layer decomposition, the
+text-to-image rewriter describes the complete 1:1 canvas.
+
+This rewrite is a pre-processing step *outside* `infer.py`. Prompt enhancement
+(PE) can use `Ling-3.0-flash-VL` or `qwen3.8-27B`; run it first, then pass its
+output to `--prompt` as raw text or via an asset file.
+
+The released system prompt is stored in
+[`assets/t2i_rewriter_system_prompt.txt`](assets/t2i_rewriter_system_prompt.txt)
+and reproduced below:
+
+```text
+You are a senior visual designer and image-prompt engineer. Expand the user's request into one precise, high-resolution Figma-style caption. Return only one JSON object.
+
+Use exactly two top-level keys. `canvas_settings` contains exactly `aspect_ratio`, `ambient_lighting`, and `image_style`. `layers` lists visible groups from background to topmost overlay. Every layer contains exactly `description`, `coordinates`, `hierarchy_and_relation`, and `color_specs`; `color_specs` is an array of hex colors.
+
+`coordinates` MUST be one string, never an object or array, in exactly this form: `"cx: 0.500, cy: 0.500, w: 1.000, h: 1.000"`. Values are normalized; each bbox encloses its complete owned object and stays inside the canvas.
+
+A layer is one selectable visible semantic group: background, full person, coherent object, panel, card, row, or text block. Prefer the fewest groups that preserve the layout. Keep people and objects intact. Never create invisible parents, guides, placeholders, empty layers, duplicate summaries, or multiple owners for one element.
+
+Preserve every user-supplied rendered string character-for-character and as one contiguous string. Unless multiple visible copies are requested, it must occur exactly once across all `description` fields and zero times in `hierarchy_and_relation`. Quote it only where describing its visible rendering; refer to the related subject elsewhere with unquoted semantic wording. Enumerate intended copy, invent extra copy sparingly, and never hide content behind "other text", "remaining labels", or "etc."
+
+Describe concrete composition, typography, materials, texture, lighting, pose, and camera treatment without literary filler. Use `hierarchy_and_relation` only for ownership, alignment, containment, stacking, and occlusion.
+
+Infer structured layouts first. Use one complete layer per card and state its row and column. A compact secondary table may be one layer only if every header and cell is listed; otherwise use a visible shared frame when present, one complete header, and one complete layer per body row, binding values to columns and stating blanks. Enumerate sequences, schedules, spans, gaps, and vacant tracks in visual order. Do not mistake ordinary alignment for a table.
+
+Silently verify schema, string coordinates, Z-order, exact-text counts, geometry, bbox validity, and completeness.
+```
+
+### Transparent-background generation tip
+
+To generate an image with an alpha channel, choose exactly one of the following
+fixed phrases and place it at the beginning of the prompt. Do not combine
+multiple prefixes.
+
+- `带透明通道，4通道RGBA图像`
+- `透明背景，alpha通道，无底图`
+- `抠图素材，背景alpha=0`
+- `孤立主体，透明PNG图层`
+- `不要白底，不要棋盘格，只要透明通道`
+- `RGBA, 4-channel, transparent background`
+- `isolated subject, alpha matte, no background`
+- `cutout PNG, alpha=0 outside the object`
+- `transparent canvas, not white, not checkerboard`
+- `production RGBA layer for compositing`
+
+### Layer-decomposition demo
+
+This example separates the flattened card design into six transparent RGBA
+layers. See the [source input](assets/layer_samples/card_making_input.png) and
+the [exact six-layer specification](assets/layer_samples/card_making_prompt.txt).
+
+```bash
+python infer.py \
+  --model inclusionAI/Ming-Image-0.1-Design-Layer \
+  --task layer-decompose \
+  --input-image assets/layer_samples/card_making_input.png \
+  --prompt assets/layer_samples/card_making_prompt.txt \
+  --attn-implementation flash_attention_2 \
+  --resolution 1024 \
+  --output-dir outputs/layers
+```
+
+The layer count is parsed from a "Decompose this image into N layers" or
+"Number of layers: N" specification in the prompt. When `--prompt` is omitted,
+`--num-layers N` generates the default request `Decompose this image into N
+layers.`
+
+The layer model returns the requested layers plus one leading
+composite/full-canvas image. The CLI skips that first image and saves the
+standalone layers as `layer_01.png`, `layer_02.png`, and so on.
+
+### Layer-decomposition prompt rewriting
+
+Layer decomposition is driven by an explicit per-layer specification, not a
+free-form caption. The reference pipeline first runs a prompt enhancer: an
+instruction-following VLM rewrites the user's rough layer plan into a precise
+decomposition — concrete colors/positions/shapes per layer, real text quoted
+verbatim on its own front layer, a text-supporting card/panel/badge/banner as
+its own layer directly behind the text, the main subject as its own layer, and
+the background last absorbing the remaining supports and surfaces.
+
+This rewrite is a pre-processing step *outside* `infer.py`: run the enhancer
+first, then pass its output to `--prompt` (as raw text or via an asset file).
+The enhancer's output is exactly the format this CLI parses — it regenerates the
+"Decompose this image into N layers" / "Number of layers: N" spec, so the layer
+count flows through the same `--prompt` parsing path described above.
+
+Prompt enhancement (PE) can use `Ling-3.0-flash-VL` or `qwen3.8-27B` as the
+instruction-following VLM. The guided prompt below is part of the released
+pipeline and is kept here for reproducibility:
+
+```text
+GUIDED_PROMPT = """You are a graphic-design layer-decomposition expert. You are given ONE flattened design image and a ROUGH layer plan from the user. Rewrite the rough plan into a precise layer decomposition that matches the image.
+
+User's rough layer plan:
+{spec}
+
+Guidelines:
+- Follow the user's plan EXACTLY: use the same number of layers and the same per-layer role/meaning, in the same order. Layer 1 is the FRONT-most (topmost); the last layer is the background/environment. Stacking the layers back-to-front must reproduce the image.
+- For each layer, write a concrete one-or-two-sentence description grounded in the image: real colors, positions and shapes.
+- TEXT goes in the front layer(s); quote any real text VERBATIM in double quotes and keep its original language (do not translate).
+- A text-supporting CARD / PANEL / BADGE / BANNER is its OWN layer directly behind the text — do not merge it into the text layer or into the background.
+- The MAIN SUBJECT (hero product/photo/illustration) is its own layer.
+- The BACKGROUND/ENVIRONMENT is the LAST layer and ABSORBS supporting props and surfaces under the subject (tables, boards, plates, floors, shadows, gradients, patterns) — these are NOT separate layers.
+
+Write the description DIRECTLY about the content; do NOT mention "image" or narrate your reasoning. Output EXACTLY in this format and NOTHING else (N = the number of layers in the user's plan):
+
+Decompose this image into N layers with the following specifications:
+
+Number of layers: N
+Layer 1: <front-most layer>
+Layer 2: <...>
+Layer N: <background/environment layer>"""
+```
+
+## Deployment
+
+We recommend the following inference frameworks to serve the model:
+
+- vLLM-Omni: see the [recipes](https://github.com/vllm-project/vllm-omni/blob/main/recipes/inclusionAI/Ming-Image.md)
+  and [installation guide](https://docs.vllm.ai/projects/vllm-omni/en/latest/getting_started/quickstart/).
+
+## Verification
+
+Fast contract tests do not require model weights:
+
+```bash
+python -m unittest -v tests.test_inference_profile
+python -m unittest -v tests.test_padding
+```
+
+Full inference is a GPU smoke test and requires both checkpoint families. At a
+minimum, validate the four-seasons-cabin text-to-image case and the six-layer
+card-making decomposition case with fixed seeds. The smoke test defaults to
+the single-GPU placement and FlashAttention 2 configuration shown above:
+
+```bash
+export CUDA_VISIBLE_DEVICES=0
+export MING_GENERATION_MODEL=inclusionAI/Ming-Image-0.1-Design
+export MING_LAYER_MODEL=inclusionAI/Ming-Image-0.1-Design-Layer
+export MING_SMOKE_OUTPUT_DIR=/path/to/persistent/results
+python -m unittest -v tests.test_inference_smoke.InferenceSmokeTest.test_two_step_showcase_smoke
+```
+
+## Checkpoint contract
+
+Each checkpoint declares its runtime capability in `transformer/config.json`:
+
+| Family | `alignment_padding_mode` | `multi_frame_output` |
+| --- | --- | ---: |
+| Text-to-image | `"zero_masked"` | `false` |
+| Layer decomposition | `"learned"` | `true` |
+
+Both fields must be present together, and the VAE component must be the
+4-channel `AutoencoderKLQwenImage` (argmax reference encoding). Missing,
+partial, or unknown values are errors; the loader never infers padding
+behavior from a directory name or silently falls back to another mode.
+
+Legacy packages that predate the component metadata are loaded strictly from
+a root `inference_profile.json` during the compatibility window; when both
+exist, the component configs are authoritative and any disagreement with the
+legacy file is an error.

code/REVISION

@@ -0,0 +1 @@
+branch rocm-halo @ f986f7a (upstream inclusionAI/Ming-Image 62c6072 + 3 commits); rocm.patch sha256 03fe16f8ed566bf99ec654caa687c7e659f489df6d077cfb7fc7aee134ebc777

code/bailingmm_utils.py

@@ -0,0 +1,504 @@
+import base64
+import logging
+import math
+import os
+from io import BytesIO
+from tqdm.contrib.concurrent import thread_map
+
+import numpy as np
+
+import requests
+import torch
+
+from PIL import Image
+from typing import Union, Tuple, List
+
+def fetch_video(ele):
+    raise ValueError(
+        "video input is not supported by Ming Image inference; "
+        "the public API accepts image and text input only"
+    )
+
+logger = logging.getLogger(__name__)
+
+IMAGE_FACTOR = 28
+MIN_PIXELS = 4 * 28 * 28
+MAX_PIXELS = 1024 * 28 * 28
+MAX_RATIO = 200
+
+VideoInput = Union[
+    List["Image.Image"],
+    "np.ndarray",
+    "torch.Tensor",
+    List["np.ndarray"],
+    List["torch.Tensor"],
+    List[List["Image.Image"]],
+    List[List["np.ndarray"]],
+    List[List["torch.Tensor"]],
+]
+
+
+def round_by_factor(number: int, factor: int) -> int:
+    """Returns the closest integer to 'number' that is divisible by 'factor'."""
+    return round(number / factor) * factor
+
+def ceil_by_factor(number: int, factor: int) -> int:
+    """Returns the smallest integer greater than or equal to 'number' that is divisible by 'factor'."""
+    return math.ceil(number / factor) * factor
+
+def floor_by_factor(number: int, factor: int) -> int:
+    """Returns the largest integer less than or equal to 'number' that is divisible by 'factor'."""
+    return math.floor(number / factor) * factor
+
+def is_image(image_file):
+    if isinstance(image_file, str) and (image_file.startswith("base64,") or image_file.lower().endswith(
+            ('.bmp', '.dib', '.png', '.jpg', '.jpeg', '.pbm', '.pgm', '.ppm', '.tif', '.tiff'))):
+        return True
+    elif isinstance(image_file, Image.Image):
+        return True
+    else:
+        return False
+
+def is_video(video_file):
+    if isinstance(video_file, str) and video_file.lower().endswith(
+            ('.mp4', '.mkv', '.avi', '.wmv', '.iso', ".webm")):
+        return True
+    else:
+        return False
+
+def is_audio(audio_file):
+    if isinstance(audio_file, str) and audio_file.lower().endswith(
+            (".wav", ".mp3", ".aac", ".flac", ".alac", ".m4a", ".ogg", ".wma", ".aiff", ".amr", ".au")):
+        return True
+    else:
+        return False
+
+def smart_resize(
+    height: int, width: int, factor: int = IMAGE_FACTOR, min_pixels: int = MIN_PIXELS, max_pixels: int = MAX_PIXELS
+) -> tuple[int, int]:
+    """
+    Rescales the image so that the following conditions are met:
+
+    1. Both dimensions (height and width) are divisible by 'factor'.
+
+    2. The total number of pixels is within the range ['min_pixels', 'max_pixels'].
+
+    3. The aspect ratio of the image is maintained as closely as possible.
+    """
+    if max(height, width) / min(height, width) > MAX_RATIO:
+        raise ValueError(
+            f"absolute aspect ratio must be smaller than {MAX_RATIO}, got {max(height, width) / min(height, width)}"
+        )
+    h_bar = max(factor, round_by_factor(height, factor))
+    w_bar = max(factor, round_by_factor(width, factor))
+    if h_bar * w_bar > max_pixels:
+        beta = math.sqrt((height * width) / max_pixels)
+        h_bar = floor_by_factor(height / beta, factor)
+        w_bar = floor_by_factor(width / beta, factor)
+    elif h_bar * w_bar < min_pixels:
+        beta = math.sqrt(min_pixels / (height * width))
+        h_bar = ceil_by_factor(height * beta, factor)
+        w_bar = ceil_by_factor(width * beta, factor)
+    return h_bar, w_bar
+
+def fetch_image(ele: dict[str, str | Image.Image], size_factor: int = IMAGE_FACTOR) -> Image.Image:
+    if "image" in ele:
+        image = ele["image"]
+    else:
+        image = ele["image_url"]
+    image_obj = None
+    if isinstance(image, Image.Image):
+        image_obj = image
+    elif image.startswith("http://") or image.startswith("https://"):
+        image_obj = Image.open(requests.get(image, stream=True).raw)
+    elif image.startswith("file://"):
+        image_obj = Image.open(image[7:])
+    elif image.startswith("data:image"):
+        if "base64," in image:
+            _, base64_data = image.split("base64,", 1)
+            data = base64.b64decode(base64_data)
+            image_obj = Image.open(BytesIO(data))
+    else:
+        image_obj = Image.open(image)
+    if image_obj is None:
+        raise ValueError(f"Unrecognized image input, support local path, http url, base64 and PIL.Image, got {image}")
+    image = image_obj.convert("RGB")
+    ## resize
+    if "resized_height" in ele and "resized_width" in ele:
+        resized_height, resized_width = smart_resize(
+            ele["resized_height"],
+            ele["resized_width"],
+            factor=size_factor,
+        )
+    else:
+        width, height = image.size
+        min_pixels = ele.get("min_pixels", MIN_PIXELS)
+        max_pixels = ele.get("max_pixels", MAX_PIXELS)
+        resized_height, resized_width = smart_resize(
+            height,
+            width,
+            factor=size_factor,
+            min_pixels=min_pixels,
+            max_pixels=max_pixels,
+        )
+    image = image.resize((resized_width, resized_height))
+
+    return image
+
+def fetch_image_wo_resize(ele: dict[str, str | Image.Image], size_factor: int = IMAGE_FACTOR) -> Image.Image:
+    if "image" in ele:
+        image = ele["image"]
+    else:
+        image = ele["image_url"]
+    image_obj = None
+    if isinstance(image, Image.Image):
+        image_obj = image
+    elif image.startswith("http://") or image.startswith("https://"):
+        image_obj = Image.open(requests.get(image, stream=True).raw)
+    elif image.startswith("file://"):
+        image_obj = Image.open(image[7:])
+    elif image.startswith("data:image"):
+        if "base64," in image:
+            _, base64_data = image.split("base64,", 1)
+            data = base64.b64decode(base64_data)
+            image_obj = Image.open(BytesIO(data))
+    else:
+        image_obj = Image.open(image)
+    if image_obj is None:
+        raise ValueError(f"Unrecognized image input, support local path, http url, base64 and PIL.Image, got {image}")
+    
+    #image = image_obj.convert("RGB")
+
+    return image_obj
+
+def fetch_audio(ele: dict[str, str | torch.Tensor], return_tensor="pt") -> Tuple[Union[torch.Tensor, np.ndarray], int]:
+    import torchaudio
+
+    if "audio" in ele:
+        audio = ele["audio"]
+    else:
+        audio = ele["audio_url"]
+
+    if isinstance(audio, torch.Tensor):
+        waveform = audio
+        sample_rate: int = ele.get("sample_rate", 16000)
+    elif audio.startswith("http://") or audio.startswith("https://"):
+        audio_file = BytesIO(requests.get(audio, stream=True).content)
+        waveform, sample_rate = torchaudio.load(audio_file)
+    elif audio.startswith("file://"):
+        waveform, sample_rate = torchaudio.load(audio[7:])
+    else:
+        waveform, sample_rate = torchaudio.load(audio)
+    if return_tensor == "pt":
+        return waveform, sample_rate
+    else:
+        return waveform.numpy(), sample_rate
+
+def extract_vision_info(conversations: list[dict] | list[list[dict]]) -> list[dict]:
+    vision_infos = []
+    if isinstance(conversations[0], dict):
+        conversations = [conversations]
+    for conversation in conversations:
+        for message in conversation:
+            if isinstance(message["content"], list):
+                for ele in message["content"]:
+                    if (
+                            "image" in ele
+                            or "image_url" in ele
+                            or "video" in ele
+                            or "video_url" in ele
+                            or "audio" in ele
+                            or "audio_url" in ele
+                            or ele["type"] in ["image", "image_url", "video", "video_url", "audio", "audio_url"]
+                    ):
+                        vision_infos.append(ele)
+    return vision_infos
+
+
+def process_reference_vision_info(
+    conversations: list[dict] | list[list[dict]],
+) -> list[Image.Image] | None:
+    vision_infos = extract_vision_info(conversations)
+    ## Read images
+    image_inputs = []
+
+    def inner_process_func(vision_info):
+        if "image" in vision_info or "image_url" in vision_info:
+            res_list = []
+            if "image" in vision_info and isinstance(vision_info["image"], (tuple, list)):
+                for i in range(len(vision_info["image"])):
+                    res_list.append(fetch_image_wo_resize({"type": "image", "image": vision_info["image"][i]}))
+            elif "image_url" in vision_info and vision_info["image_url"].get("url", None) is not None:
+                vision_info["image_url"] = vision_info["image_url"].get("url")
+                res_list.extend([fetch_image_wo_resize(vision_info)])
+            else:
+                res_list.extend([fetch_image_wo_resize(vision_info)])
+            return {'image_inputs':res_list}
+        else:
+            return None
+
+    vision_infos_reslist = thread_map(inner_process_func, vision_infos, disable=True)
+    for res in vision_infos_reslist:
+        if res is None:
+            raise ValueError("image, image_url, video, video_url, audio or audio_url should in content.")
+        elif 'image_inputs' in res:
+            image_inputs.extend(res['image_inputs'])
+
+    if len(image_inputs) > 1: # multi-image input keeps only the first image as the VAE reference
+        image_inputs = [image_inputs[0]]
+
+    if len(image_inputs) == 0:
+        image_inputs = None
+
+    return image_inputs
+
+
+def process_vision_info(
+    conversations: list[dict] | list[list[dict]],
+) -> tuple[list[Image.Image] | None, list[torch.Tensor | list[Image.Image]] | None, list[
+    torch.Tensor | list[np.ndarray]] | None]:
+    vision_infos = extract_vision_info(conversations)
+    ## Read images, videos or audios
+    image_inputs = []
+    video_inputs = []
+    audio_inputs = []
+
+    def inner_process_func(vision_info):
+        if "image" in vision_info or "image_url" in vision_info:
+            res_list = []
+            if "image" in vision_info and isinstance(vision_info["image"], (tuple, list)):
+                for i in range(len(vision_info["image"])):
+                    res_list.append(fetch_image({"type": "image", "image": vision_info["image"][i]}))
+            elif "image_url" in vision_info and vision_info["image_url"].get("url", None) is not None:
+                vision_info["image_url"] = vision_info["image_url"].get("url")
+                res_list.extend([fetch_image(vision_info)])
+            else:
+                res_list.extend([fetch_image(vision_info)])
+            return {'image_inputs':res_list}
+
+        elif "video" in vision_info or "video_url" in vision_info:
+            if "video_url" in vision_info and vision_info["video_url"].get("url", None) is not None:
+                data_value = vision_info["video_url"].get("url")
+            elif "video" in vision_info and not os.path.isdir(vision_info['video']):
+                data_value = vision_info['video']
+            else:
+                data_value = [os.path.join(vision_info['video'], frame) for frame in sorted(os.listdir(vision_info['video']))]
+            vision_info['video']=data_value
+            return {"video_inputs": [fetch_video(vision_info)]}
+
+        elif "audio" in vision_info or "audio_url" in vision_info:
+            if "audio" in vision_info and isinstance(vision_info["audio"], (tuple, list)):
+                return {"audio_inputs":[fetch_audio(info) for info in vision_info["audio"]]}
+            elif "audio_url" in vision_info and vision_info["audio_url"].get("url", None) is not None:
+                vision_info["audio_url"] = vision_info["audio_url"].get("url")
+                return {"audio_inputs":[fetch_audio(vision_info)]}
+            else:
+                return {"audio_inputs":[fetch_audio(vision_info)]}
+        else:
+            return None
+
+    vision_infos_reslist = thread_map(inner_process_func, vision_infos, disable=True)
+    for res in vision_infos_reslist:
+        if res is None:
+            raise ValueError("image, image_url, video, video_url, audio or audio_url should in content.")
+        elif 'image_inputs' in res:
+            image_inputs.extend(res['image_inputs'])
+        elif 'video_inputs' in res:
+            video_inputs.extend(res['video_inputs'])
+        elif 'audio_inputs' in res:
+            audio_inputs.extend(res['audio_inputs'])
+
+    if len(image_inputs) == 0:
+        image_inputs = None
+    if len(video_inputs) == 0:
+        video_inputs = None
+    if len(audio_inputs) == 0:
+        audio_inputs = None
+    return image_inputs, video_inputs, audio_inputs
+
+
+def get_closest_ratio(height: float, width: float, aspect_ratios: dict):
+    aspect_ratio = height / width
+    closest_ratio = min(aspect_ratios.keys(), key=lambda ratio: abs(float(ratio) - aspect_ratio))
+    return aspect_ratios[closest_ratio], float(closest_ratio)
+
+def process_ratio(ori_h, ori_w, highres=512):
+    ASPECT_RATIO_512 = {
+        "0.25": [256, 1024], "0.26": [256, 992], "0.27": [256, 960], "0.28": [256, 928],
+        "0.32": [288, 896], "0.33": [288, 864], "0.35": [288, 832], "0.4": [320, 800],
+        "0.42": [320, 768], "0.48": [352, 736], "0.5": [352, 704], "0.52": [352, 672],
+        "0.5455": [384, 704], "0.57": [384, 672], "0.6": [384, 640], "0.65": [416, 640],
+        "0.68": [416, 608], "0.72": [416, 576], "0.78": [448, 576],
+        "0.82": [448, 544], "0.88": [480, 544], "0.94": [480, 512],
+        "1.0": [512, 512], "1.07": [512, 480], "1.13": [544, 480], "1.21": [544, 448],
+        "1.29": [576, 448], "1.38": [576, 416],
+        "1.46": [608, 416], "1.5385": [640, 416], "1.67": [640, 384], "1.75": [672, 384],
+        "1.8333": [704, 384], "2.0": [704, 352], "2.09": [736, 352], "2.4": [768, 320],
+        "2.5": [800, 320], "2.89": [832, 288], "3.0": [864, 288], "3.11": [896, 288],
+        "3.62": [928, 256], "3.75": [960, 256], "3.88": [992, 256], "4.0": [1024, 256],
+    }
+    ASPECT_RATIO_1024 = {
+        "0.25": [512, 2048], "0.26": [512, 1984], "0.27": [512, 1920], "0.28": [512, 1856],
+        "0.32": [576, 1792], "0.33": [576, 1728], "0.35": [576, 1664], "0.4": [640, 1600],
+        "0.42": [640, 1536], "0.48": [704, 1472], "0.5": [704, 1408], "0.52": [704, 1344],
+        "0.5581": [768, 1376], "0.5625": [720, 1280], "0.5647": [768, 1360], "0.57": [768, 1344],
+        "0.6": [768, 1280], "0.622": [816, 1312], "0.625": [800, 1280], "0.65": [832, 1280],
+        "0.6582": [832, 1264], "0.6667": [832, 1248], "0.6709": [848, 1264], "0.68": [832, 1216],
+        "0.7013": [864, 1232], "0.72": [832, 1152], "0.7467": [896, 1200], "0.75": [864, 1152],
+        "0.7568": [896, 1184], "0.78": [896, 1152], "0.8": [896, 1120], "0.8056": [928, 1152],
+        "0.82": [896, 1088], "0.88": [960, 1088], "0.94": [960, 1024], "0.9846": [1024, 1040],
+        "1.0": [1024, 1024], "1.07": [1024, 960], "1.13": [1088, 960], "1.21": [1088, 896],
+        "1.2414": [1152, 928], "1.25": [1120, 896], "1.2807": [1168, 912], "1.29": [1152, 896],
+        "1.3333": [1152, 864], "1.3393": [1200, 896], "1.38": [1152, 832], "1.46": [1216, 832],
+        "1.4906": [1264, 848], "1.5": [1248, 832], "1.6": [1280, 800], "1.67": [1280, 768],
+        "1.75": [1344, 768], "1.7708": [1360, 768], "1.7778": [1280, 720], "2.0": [1408, 704],
+        "2.09": [1472, 704], "2.4": [1536, 640], "2.5": [1600, 640], "2.89": [1664, 576],
+        "3.0": [1728, 576], "3.11": [1792, 576], "3.62": [1856, 512], "3.75": [1920, 512],
+        "3.88": [1984, 512], "4.0": [2048, 512],
+    }
+
+    ASPECT_RATIO_672 = {
+        "0.28": [352, 1280], "0.32": [384, 1184], "0.38": [416, 1088], "0.44": [448, 1024],
+        "0.52": [480, 928], "0.5636": [496, 880], "0.57": [512, 896], "0.65": [544, 832],
+        "0.6667": [544, 816], "0.75": [576, 768], "0.8": [576, 720], "0.83": [608, 736],
+        "0.91": [640, 704], "1.00": [672, 672], "1.10": [704, 640], "1.21": [736, 608],
+        "1.25": [720, 576], "1.33": [768, 576], "1.39": [800, 576], "1.5": [816, 544],
+        "1.53": [832, 544], "1.69": [864, 512], "1.75": [896, 512], "1.7742": [880, 496],
+        "1.93": [928, 480], "2.00": [960, 480],
+        "2.21": [992, 448], "2.29": [1024, 448], "2.54": [1056, 416], "2.62": [1088, 416],
+        "2.69": [1120, 416], "3.00": [1152, 384], "3.08": [1184, 384], "3.17": [1216, 384],
+        "3.55": [1248, 352], "3.64": [1280, 352],
+    }
+
+    ASPECT_RATIO_2048 = {
+        "0.25": [1024, 4096], "0.26": [1024, 3968], "0.27": [1024, 3840], "0.28": [1024, 3712],
+        "0.32": [1152, 3584], "0.33": [1152, 3456], "0.35": [1152, 3328], "0.4": [1280, 3200],
+        "0.42": [1280, 3072], "0.48": [1408, 2944], "0.5": [1408, 2816], "0.52": [1408, 2688],
+        "0.5625": [1440, 2560], "0.57": [1536, 2688], "0.6": [1536, 2560], "0.6667": [1664, 2496],
+        "0.68": [1664, 2432], "0.72": [1664, 2304], "0.75": [1824, 2432], "0.78": [1792, 2304],
+        "0.7917": [1824, 2304], "0.8": [1792, 2240], "0.82": [1792, 2176], "0.88": [1920, 2176],
+        "0.94": [1920, 2048],
+        "1.0": [2048, 2048], "1.07": [2048, 1920], "1.13": [2176, 1920], "1.21": [2176, 1792],
+        "1.25": [2240, 1792], "1.2632": [2304, 1824], "1.29": [2304, 1792],
+        "1.3333": [2432, 1824], "1.38": [2304, 1664],
+        "1.46": [2432, 1664], "1.5": [2496, 1664], "1.67": [2560, 1536], "1.75": [2688, 1536],
+        "1.7778": [2560, 1440], "2.0": [2816, 1408], "2.09": [2944, 1408], "2.4": [3072, 1280],
+        "2.5": [3200, 1280], "2.89": [3328, 1152], "3.0": [3456, 1152], "3.11": [3584, 1152],
+        "3.62": [3712, 1024], "3.75": [3840, 1024], "3.88": [3968, 1024], "4.0": [4096, 1024],
+        "0.2941": [1120, 3808], "0.3043": [1120, 3680], "0.3679": [1248, 3392],
+        "0.3846": [1280, 3328], "0.433": [1344, 3104], "0.4574": [1376, 3008],
+        "0.4681": [1408, 3008], "0.5465": [1504, 2752], "0.6296": [1632, 2592],
+        "0.6582": [1664, 2528], "0.7432": [1760, 2368], "0.8551": [1888, 2208],
+        "0.9692": [2016, 2080], "1.0317": [2080, 2016], "1.1695": [2208, 1888],
+        "1.3455": [2368, 1760], "1.5192": [2528, 1664], "1.5882": [2592, 1632],
+        "1.8298": [2752, 1504], "1.913": [2816, 1472], "2.1364": [3008, 1408],
+        "2.186": [3008, 1376], "2.3095": [3104, 1344], "2.6": [3328, 1280],
+        "2.7179": [3392, 1248], "3.2857": [3680, 1120], "3.4": [3808, 1120],
+    }
+
+    aspect_ratio_dict = {
+        512 : ASPECT_RATIO_512,
+        672 : ASPECT_RATIO_672,
+        1024 : ASPECT_RATIO_1024,
+        2048 : ASPECT_RATIO_2048,
+    }
+
+    if highres is None or highres is False:
+        highres = 512
+    elif highres is True:
+        highres = 1024
+
+    aspect_ratio = aspect_ratio_dict[min([i for i in aspect_ratio_dict], key=lambda x: abs(x - highres))]
+
+    closest_size, _ = get_closest_ratio(ori_h, ori_w, aspect_ratios=aspect_ratio)
+    closest_size = list(map(lambda x: int(x), closest_size))
+    if closest_size[0] / ori_h > closest_size[1] / ori_w:
+        resize_size = closest_size[0], int(ori_w * closest_size[0] / ori_h)
+    else:
+        resize_size = int(ori_h * closest_size[1] / ori_w), closest_size[1]
+    return closest_size, resize_size
+
+
+def find_first_index_of_consecutive_ones(lst):
+    """
+    Given a list of 0s and 1s, return the index of the first 1 of each
+    consecutive run of 1s.
+
+    Args:
+        lst (list): list of 0s and 1s
+
+    Returns:
+        list: indices of the first 1 of each consecutive run of 1s
+    """
+    result = []
+    i = 0
+    n = len(lst)
+    
+    while i < n:
+        if lst[i] == 1:
+            # find the start of a consecutive run of 1s
+            result.append(i)
+            # skip the remainder of the consecutive run of 1s
+            while i < n and lst[i] == 1:
+                i += 1
+        else:
+            i += 1
+    
+    return result
+
+def merge_consecutive_ones(lst, n):
+    """
+    Given a list of 0s and 1s, merge every n consecutive 1s of each run
+    (length >= 1) into a single 1. Each run must have a length divisible by n.
+    The relative order of 0s and 1s is preserved.
+
+    Args:
+        lst: list of 0s and 1s
+        n: positive integer merge unit size
+
+    Returns:
+        list: the merged list
+    """
+    assert isinstance(lst, list), "input must be a list"
+    assert isinstance(n, int) and n > 0, "n must be a positive integer"
+    
+    # iterate over the list, extract runs of 1s, and verify each run length is divisible by n
+    i = 0
+    while i < len(lst):
+        if lst[i] == 1:
+            count = 0
+            start = i
+            # count the run of consecutive 1s
+            while i < len(lst) and lst[i] == 1:
+                count += 1
+                i += 1
+            # every run of 1s must be divisible by n
+            assert count % n == 0, f"run of 1s at index {start} has length {count}, not divisible by n={n}"
+        else:
+            i += 1
+
+    # build the new list by merging groups
+    result = []
+    i = 0
+    while i < len(lst):
+        if lst[i] == 0:
+            result.append(0)
+            i += 1
+        else:
+            # process a run of 1s
+            count = 0
+            while i < len(lst) and lst[i] == 1:
+                count += 1
+                i += 1
+            # merge every n 1s into a single 1
+            result.extend([1] * (count // n))
+    
+    return result
+
+def get_default_image_gen_hw(image_gen_highres, image_gen_aspect_ratio):
+    if image_gen_aspect_ratio is None:
+        image_gen_aspect_ratio = 1.0
+
+    closest_size, _ = process_ratio(ori_h=512, ori_w=int(512.0 * image_gen_aspect_ratio), highres=image_gen_highres)
+    h, w = closest_size
+    return h, w

code/configuration_bailing_moe_v2.py

@@ -0,0 +1,85 @@
+"""Bailing MoE model configuration"""
+from transformers.configuration_utils import PretrainedConfig
+
+class BailingMoeV2Config(PretrainedConfig):
+    model_type = "bailing_moe_v2"
+    def __init__(
+        self,
+        vocab_size=30592,
+        hidden_size=1024,
+        intermediate_size=None,
+        num_hidden_layers=24,
+        num_attention_heads=16,
+        num_key_value_heads=0,
+        hidden_act="silu",
+        use_qkv_bias=False,  # bailing only
+        use_qk_norm=False,
+        use_bias=True,  # bailing only
+        rms_norm_eps=1e-05,
+        norm_head=False,  # bailing only
+        tie_word_embeddings=False,  # PretrainedConfig key, here change default value.
+        embedding_dropout=0.1,
+        attention_dropout=0.1,
+        output_dropout=0.1,
+        initializer_range=0.02,
+        max_position_embeddings=16384,
+        rope_theta=10000.0,
+        use_cache=True,
+        use_sliding_window=False,
+        sliding_window=4096,
+        max_window_layers=28,
+        rope_scaling=None,
+        pad_token_id=126081,
+        num_experts=16,
+        num_shared_experts=0,
+        num_experts_per_tok=2,
+        n_group=8,
+        topk_group=4,
+        routed_scaling_factor=2.5,
+        moe_intermediate_size=None,
+        first_k_dense_replace=0,
+        head_dim=None,
+        output_router_logits=False,
+        partial_rotary_factor=0.5,
+        router_type="topN",
+        _attn_implementation="flash_attention_2",
+        **kwargs,
+    ):
+        self.num_hidden_layers = num_hidden_layers
+        self.vocab_size = vocab_size
+        self.hidden_size = hidden_size
+        self.intermediate_size = intermediate_size
+        self.num_attention_heads = num_attention_heads
+        self.num_key_value_heads = num_key_value_heads
+        self.hidden_act = hidden_act
+        self.use_qkv_bias = use_qkv_bias
+        self.use_bias = use_bias
+        self.norm_head = norm_head
+        self.rms_norm_eps = rms_norm_eps
+        self.embedding_dropout = embedding_dropout
+        self.attention_dropout = attention_dropout
+        self.output_dropout = output_dropout
+        self.initializer_range = initializer_range
+        self.max_position_embeddings = max_position_embeddings
+        self.rope_theta = rope_theta
+        self.use_cache = use_cache
+        self.use_sliding_window = use_sliding_window
+        self.sliding_window = sliding_window
+        self.max_window_layers = max_window_layers
+        self.head_dim = head_dim or self.hidden_size // self.num_attention_heads
+        self.rope_scaling = rope_scaling
+        # MoE configs
+        self.num_experts = num_experts
+        self.num_shared_experts = num_shared_experts
+        self.num_experts_per_tok = num_experts_per_tok
+        self.n_group = n_group
+        self.topk_group = topk_group
+        self.moe_intermediate_size = moe_intermediate_size
+        self.first_k_dense_replace = first_k_dense_replace
+        self.output_router_logits = output_router_logits
+        self.routed_scaling_factor = routed_scaling_factor
+        self.partial_rotary_factor = partial_rotary_factor
+        self.router_type = router_type
+        super().__init__(pad_token_id=pad_token_id, tie_word_embeddings=tie_word_embeddings, **kwargs)
+        self._attn_implementation = _attn_implementation
+

code/configuration_bailingmm2.py

@@ -0,0 +1,43 @@
+# coding=utf-8
+# Copyright 2024 ANT Group and the HuggingFace Inc. team. All rights reserved.
+#
+# 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.
+
+from transformers import PretrainedConfig
+from qwen2_5_vit import Qwen2_5_VLVisionConfig
+from configuration_bailing_moe_v2 import BailingMoeV2Config
+
+
+class BailingMM2Config(PretrainedConfig):
+    model_type = "bailingmm_moe_v2_lite"
+    # Declared so transformers' `_attn_implementation` setter recurses into both towers.
+    # Without it an explicit attn_implementation (e.g. "eager" on ROCm, which has no
+    # flash-attn) never reaches them, and their "flash_attention_2" defaults raise at
+    # model construction.
+    sub_configs = {"vision_config": Qwen2_5_VLVisionConfig, "llm_config": BailingMoeV2Config}
+
+    def __init__(
+        self,
+        mlp_depth=1,
+        llm_config: BailingMoeV2Config = None,
+        vision_config: Qwen2_5_VLVisionConfig = None,
+        audio_config=None,
+        **kwargs
+    ):
+        if audio_config is not None:
+            raise ValueError("audio_config is not supported by Ming Image inference")
+        self.audio_config = None
+        self.vision_config = Qwen2_5_VLVisionConfig(**vision_config) if isinstance(vision_config, dict) else vision_config
+        self.llm_config = BailingMoeV2Config(**llm_config) if isinstance(llm_config, dict) else llm_config
+        self.mlp_depth = mlp_depth
+        super().__init__(**kwargs)

code/examples/profiles/generation_edit.json

@@ -0,0 +1,8 @@
+{
+  "schema_version": 1,
+  "inference_profile": "generation_edit",
+  "alignment_padding_mode": "zero_masked",
+  "multi_frame_output": false,
+  "vae_input_channels": 4,
+  "vae_sample_mode": "argmax"
+}

code/examples/profiles/layer_decompose.json

@@ -0,0 +1,8 @@
+{
+  "schema_version": 1,
+  "inference_profile": "layer_decompose",
+  "alignment_padding_mode": "learned",
+  "multi_frame_output": true,
+  "vae_input_channels": 4,
+  "vae_sample_mode": "argmax"
+}

code/generate_paired.sh

@@ -0,0 +1,162 @@
+#!/usr/bin/env bash
+# Paired pipeline: Ling-3.0-flash-VL prompt enhancement -> Ming-Image text-to-image.
+#
+#   Stage 1  pe_ling.py  caption -> validated structured JSON prompt
+#                        (system prompt: assets/t2i_rewriter_system_prompt.txt)
+#   Stage 2  infer.py    --task text-to-image --prompt <json file> -> PNG
+#                        (infer.py reads --prompt as a file when the path exists)
+#
+# Artifacts land in --output-dir: enhanced_prompt.json (overwritten per run)
+# plus the PNG(s) infer.py writes (image_00.png for text-to-image).
+# Fails loudly at every stage (set -Eeuo pipefail + ERR trap + stage checks).
+set -Eeuo pipefail
+trap 'printf "generate_paired: FAILED at line %d (exit %d)\n" "$LINENO" "$?" >&2' ERR
+
+SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd)"
+PYTHON="${PYTHON:-python3}"
+
+# Local llama-server seat serving Ling-3.0-flash-VL on the target box.
+DEFAULT_BASE_URL="http://127.0.0.1:8090/v1"
+DEFAULT_PE_MODEL="ling-3.0-flash-vl-mtp-halo-STRIX_LEAN"
+
+usage() {
+  cat <<'EOF'
+Usage: generate_paired.sh --model DIR_OR_REPO CAPTION [options] [-- EXTRA_INFER_ARGS...]
+
+Enhances CAPTION with Ling-3.0-flash-VL (pe_ling.py), validates the structured
+JSON rewrite, then renders it with infer.py --task text-to-image.
+
+Required:
+  CAPTION                free-form design caption (positional)
+  --model DIR_OR_REPO    Ming checkpoint directory or HF repo id
+                         (may also be set via the MING_MODEL environment variable)
+
+Passthrough to infer.py (all optional; infer.py defaults in parentheses):
+  --resolution N         resolution bucket, 1024 or 2048 for text-to-image;
+                         other positive values snap to the nearest bucket (2048)
+  --seed N               generation seed (42)
+  --steps N              diffusion steps (12)
+  --                     everything after this is passed to infer.py verbatim
+                         (e.g. -- --validate-only --dtype float16)
+
+Prompt-enhancement endpoint:
+  --base-url URL         OpenAI-compatible base URL (http://127.0.0.1:8090/v1)
+  --pe-model ID          chat model id served there
+                         (ling-3.0-flash-vl-mtp-halo-STRIX_LEAN)
+  LITELLM_API_KEY env    exported key is sent as a Bearer token (for a gated
+                         OpenAI-compatible gateway such as LiteLLM)
+
+Other:
+  --output-dir DIR       artifact directory (outputs/paired)
+  -h, --help             this help
+
+Examples:
+  ./generate_paired.sh --model /models/Ming-Image-0.1-Design \
+      "espresso machine product poster, warm morning light" --resolution 2048
+
+  LITELLM_API_KEY=sk-... ./generate_paired.sh \
+      --base-url http://<gateway-host>:4000/v1 --pe-model <gateway-model-name> \
+      --model /models/Ming-Image-0.1-Design "a caption" --seed 7
+EOF
+}
+
+die() {
+  printf 'generate_paired: %s\n' "$*" >&2
+  exit 1
+}
+
+model="${MING_MODEL:-}"
+base_url="$DEFAULT_BASE_URL"
+pe_model="$DEFAULT_PE_MODEL"
+output_dir="outputs/paired"
+resolution=""
+seed=""
+steps=""
+caption=""
+extra_infer_args=()
+
+while [[ $# -gt 0 ]]; do
+  case "$1" in
+    --model)      [[ $# -ge 2 ]] || die "--model requires a value";      model="$2";      shift 2 ;;
+    --base-url)   [[ $# -ge 2 ]] || die "--base-url requires a value";   base_url="$2";   shift 2 ;;
+    --pe-model)   [[ $# -ge 2 ]] || die "--pe-model requires a value";   pe_model="$2";   shift 2 ;;
+    --output-dir) [[ $# -ge 2 ]] || die "--output-dir requires a value"; output_dir="$2"; shift 2 ;;
+    --resolution) [[ $# -ge 2 ]] || die "--resolution requires a value"; resolution="$2"; shift 2 ;;
+    --seed)       [[ $# -ge 2 ]] || die "--seed requires a value";       seed="$2";       shift 2 ;;
+    --steps)      [[ $# -ge 2 ]] || die "--steps requires a value";      steps="$2";      shift 2 ;;
+    -h|--help)    usage; exit 0 ;;
+    --)           shift; extra_infer_args+=("$@"); break ;;
+    -*)           usage >&2; die "unknown option: $1" ;;
+    *)
+      if [[ -n "$caption" ]]; then
+        usage >&2
+        die "unexpected extra argument: $1 (CAPTION was already given)"
+      fi
+      caption="$1"
+      shift
+      ;;
+  esac
+done
+
+[[ -n "$caption" ]] || { usage >&2; die "CAPTION is required"; }
+[[ -n "$model" ]]   || { usage >&2; die "--model DIR_OR_REPO is required (or set MING_MODEL)"; }
+if [[ -n "$resolution" && ! "$resolution" =~ ^[0-9]+$ ]]; then
+  die "--resolution must be a positive integer, got: $resolution"
+fi
+if [[ -n "$seed" && ! "$seed" =~ ^-?[0-9]+$ ]]; then
+  die "--seed must be an integer, got: $seed"
+fi
+if [[ -n "$steps" && ! "$steps" =~ ^[0-9]+$ ]]; then
+  die "--steps must be a positive integer, got: $steps"
+fi
+[[ -f "$SCRIPT_DIR/pe_ling.py" ]] || die "missing stage-1 script: $SCRIPT_DIR/pe_ling.py"
+[[ -f "$SCRIPT_DIR/infer.py" ]]   || die "missing stage-2 script: $SCRIPT_DIR/infer.py"
+command -v "$PYTHON" >/dev/null 2>&1 || die "python interpreter not found: $PYTHON (override with PYTHON=...)"
+
+mkdir -p -- "$output_dir" || die "cannot create output directory: $output_dir"
+prompt_json="$output_dir/enhanced_prompt.json"
+
+printf '== stage 1/2: prompt enhancement (pe_ling.py, model %s @ %s)\n' "$pe_model" "$base_url" >&2
+"$PYTHON" "$SCRIPT_DIR/pe_ling.py" "$caption" \
+  --out "$prompt_json" \
+  --base-url "$base_url" \
+  --model "$pe_model"
+[[ -s "$prompt_json" ]] || die "prompt enhancement produced no prompt file: $prompt_json"
+
+printf '== stage 2/2: Ming-Image text-to-image (infer.py, model %s)\n' "$model" >&2
+infer_args=(
+  --model "$model"
+  --task text-to-image
+  --prompt "$prompt_json"
+  --output-dir "$output_dir"
+)
+if [[ -n "$resolution" ]]; then infer_args+=(--resolution "$resolution"); fi
+if [[ -n "$seed" ]];       then infer_args+=(--seed "$seed"); fi
+if [[ -n "$steps" ]];      then infer_args+=(--steps "$steps"); fi
+if [[ ${#extra_infer_args[@]} -gt 0 ]]; then infer_args+=("${extra_infer_args[@]}"); fi
+validate_only=0
+for arg in ${extra_infer_args[@]+"${extra_infer_args[@]}"}; do
+  if [[ "$arg" == "--validate-only" ]]; then validate_only=1; fi
+done
+"$PYTHON" "$SCRIPT_DIR/infer.py" "${infer_args[@]}"
+
+if [[ "$validate_only" -eq 1 ]]; then
+  printf 'generate_paired: --validate-only dry run, no PNG expected; enhanced prompt: %s\n' \
+    "$prompt_json" >&2
+  exit 0
+fi
+
+# infer.py exits non-zero on failure (set -e above); additionally verify the
+# promised PNG artifacts actually exist so a silent no-write still fails
+# loudly. -newer pins the check to THIS run: stage 2 always writes its PNG
+# after stage 1 wrote enhanced_prompt.json, so stale PNGs do not satisfy it.
+pngs=()
+while IFS= read -r png; do
+  pngs+=("$png")
+done < <(find "$output_dir" -maxdepth 1 -name '*.png' -type f -newer "$prompt_json" | sort)
+if [[ ${#pngs[@]} -eq 0 ]]; then
+  die "infer.py exited 0 but wrote no PNG under $output_dir in this run"
+fi
+printf 'generate_paired: enhanced prompt: %s\n' "$prompt_json" >&2
+printf 'generate_paired: %d PNG(s):\n' "${#pngs[@]}" >&2
+printf '%s\n' "${pngs[@]}"

code/image_processing_bailingmm2.py

@@ -0,0 +1,451 @@
+# coding=utf-8
+# Copyright 2024 The Qwen team, Alibaba Group and the HuggingFace Inc. team. All rights reserved.
+#
+# This code is based on EleutherAI's GPT-NeoX library and the GPT-NeoX
+# and OPT implementations in this library. It has been modified from its
+# original forms to accommodate minor architectural differences compared
+# to GPT-NeoX and OPT used by the Meta AI team that trained the model.
+#
+# 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.
+"""Image processor class for BailingMM"""
+
+import math
+from typing import Dict, List, Optional, Union
+
+import numpy as np
+
+from transformers.image_processing_utils import BaseImageProcessor, BatchFeature
+from transformers.image_transforms import (
+    convert_to_rgb,
+    resize,
+    to_channel_dimension_format,
+)
+from bailingmm_utils import VideoInput
+from transformers.image_utils import (
+    OPENAI_CLIP_MEAN,
+    OPENAI_CLIP_STD,
+    ChannelDimension,
+    ImageInput,
+    PILImageResampling,
+    get_image_size,
+    infer_channel_dimension_format,
+    is_scaled_image,
+    is_valid_image,
+    make_list_of_images,
+    to_numpy_array,
+    valid_images,
+    validate_preprocess_arguments,
+)
+from transformers.utils import TensorType, is_vision_available, logging
+logger = logging.get_logger(__name__)
+
+if is_vision_available():
+    from PIL import Image
+
+def make_batched_images(images) -> List[List[ImageInput]]:
+    """
+    Accepts images in list or nested list format, and makes a list of images for preprocessing.
+
+    Args:
+        images (`Union[List[List[ImageInput]], List[ImageInput], ImageInput]`):
+            The input image.
+
+    Returns:
+        list: A list of images.
+    """
+    if isinstance(images, (list, tuple)) and isinstance(images[0], (list, tuple)) and is_valid_image(images[0][0]):
+        return [img for img_list in images for img in img_list]
+
+    elif isinstance(images, (list, tuple)) and is_valid_image(images[0]):
+        return images
+
+    elif is_valid_image(images):
+        return [images]
+
+    raise ValueError(f"Could not make batched images from {images}")
+
+# Copied from transformers.models.llava_next_video.image_processing_llava_next_video.make_batched_videos
+def make_batched_videos(videos) -> List[VideoInput]:
+    if isinstance(videos, (list, tuple)) and isinstance(videos[0], (list, tuple)) and is_valid_image(videos[0][0]):
+        return videos
+
+    elif isinstance(videos, (list, tuple)) and is_valid_image(videos[0]):
+        if isinstance(videos[0], Image.Image):
+            return [videos]
+        elif len(videos[0].shape) == 4:
+            return [list(video) for video in videos]
+
+    elif is_valid_image(videos) and len(videos.shape) == 4:
+        return [list(videos)]
+
+    raise ValueError(f"Could not make batched video from {videos}")
+
+def smart_resize(
+    height: int, width: int, factor: int = 28, min_pixels: int = 56 * 56, max_pixels: int = 14 * 14 * 4 * 1280
+):
+    """Rescales the image so that the following conditions are met:
+
+    1. Both dimensions (height and width) are divisible by 'factor'.
+
+    2. The total number of pixels is within the range ['min_pixels', 'max_pixels'].
+
+    3. The aspect ratio of the image is maintained as closely as possible.
+
+    """
+    if height < factor or width < factor:
+        raise ValueError(f"height:{height} or width:{width} must be larger than factor:{factor}")
+    elif max(height, width) / min(height, width) > 200:
+        raise ValueError(
+            f"absolute aspect ratio must be smaller than 200, got {max(height, width) / min(height, width)}"
+        )
+    h_bar = round(height / factor) * factor
+    w_bar = round(width / factor) * factor
+    if h_bar * w_bar > max_pixels:
+        beta = math.sqrt((height * width) / max_pixels)
+        h_bar = math.floor(height / beta / factor) * factor
+        w_bar = math.floor(width / beta / factor) * factor
+    elif h_bar * w_bar < min_pixels:
+        beta = math.sqrt(min_pixels / (height * width))
+        h_bar = math.ceil(height * beta / factor) * factor
+        w_bar = math.ceil(width * beta / factor) * factor
+    return h_bar, w_bar
+
+class BailingMM2ImageProcessor(BaseImageProcessor):
+    r"""
+    Constructs a BailingMM2 image processor that dynamically resizes images based on the original images.
+
+    Args:
+        do_resize (`bool`, *optional*, defaults to `True`):
+            Whether to resize the image's (height, width) dimensions.
+        resample (`PILImageResampling`, *optional*, defaults to `Resampling.BICUBIC`):
+            Resampling filter to use when resizing the image.
+        do_rescale (`bool`, *optional*, defaults to `True`):
+            Whether to rescale the image by the specified scale `rescale_factor`.
+        rescale_factor (`int` or `float`, *optional*, defaults to `1/255`):
+            Scale factor to use if rescaling the image.
+        do_normalize (`bool`, *optional*, defaults to `True`):
+            Whether to normalize the image.
+        image_mean (`float` or `List[float]`, *optional*, defaults to `[0.48145466, 0.4578275, 0.40821073]`):
+            Mean to use if normalizing the image. This is a float or list of floats for each channel in the image.
+        image_std (`float` or `List[float]`, *optional*, defaults to `[0.26862954, 0.26130258, 0.27577711]`):
+            Standard deviation to use if normalizing the image. This is a float or list of floats for each channel in the image.
+        do_convert_rgb (`bool`, *optional*, defaults to `True`):
+            Whether to convert the image to RGB.
+        min_pixels (`int`, *optional*, defaults to `56 * 56`):
+            The min pixels of the image to resize the image.
+        max_pixels (`int`, *optional*, defaults to `28 * 28 * 1280`):
+            The max pixels of the image to resize the image.
+        patch_size (`int`, *optional*, defaults to 14):
+            The spacial patch size of the vision encoder.
+        temporal_patch_size (`int`, *optional*, defaults to 2):
+            The temporal patch size of the vision encoder.
+        merge_size (`int`, *optional*, defaults to 2):
+            The merge size of the vision encoder to llm encoder.
+    """
+
+    model_input_names = ["pixel_values", "image_grid_thw", "pixel_values_videos", "video_grid_thw"]
+
+    def __init__(
+        self,
+        do_resize: bool = True,
+        resample: PILImageResampling = PILImageResampling.BICUBIC,
+        do_rescale: bool = True,
+        rescale_factor: Union[int, float] = 1 / 255,
+        do_normalize: bool = True,
+        image_mean: Optional[Union[float, List[float]]] = None,
+        image_std: Optional[Union[float, List[float]]] = None,
+        do_convert_rgb: bool = True,
+        min_pixels: int = 78400,
+        max_pixels: int = 2007040,
+        patch_size: int = 14,
+        temporal_patch_size: int = 2,
+        merge_size: int = 2,
+        **kwargs,
+    ) -> None:
+        super().__init__(**kwargs)
+        self.do_resize = do_resize
+        self.resample = resample
+        self.do_rescale = do_rescale
+        self.rescale_factor = rescale_factor
+        self.do_normalize = do_normalize
+        self.image_mean = image_mean if image_mean is not None else OPENAI_CLIP_MEAN
+        self.image_std = image_std if image_std is not None else OPENAI_CLIP_STD
+        self.min_pixels = min_pixels
+        self.max_pixels = max_pixels
+        self.patch_size = patch_size
+        self.temporal_patch_size = temporal_patch_size
+        self.merge_size = merge_size
+        self.size = {"min_pixels": min_pixels, "max_pixels": max_pixels}
+        self.do_convert_rgb = do_convert_rgb
+
+    def _preprocess(
+        self,
+        images: Union[ImageInput, VideoInput],
+        do_resize: bool = None,
+        resample: PILImageResampling = None,
+        do_rescale: bool = None,
+        rescale_factor: float = None,
+        do_normalize: bool = None,
+        image_mean: Optional[Union[float, List[float]]] = None,
+        image_std: Optional[Union[float, List[float]]] = None,
+        do_convert_rgb: bool = None,
+        data_format: Optional[ChannelDimension] = ChannelDimension.FIRST,
+        input_data_format: Optional[Union[str, ChannelDimension]] = None,
+    ):
+        """
+        Preprocess an image or batch of images. Copy of the `preprocess` method from `CLIPImageProcessor`.
+
+        Args:
+            images (`ImageInput`):
+                Image or batch of images to preprocess. Expects pixel values ranging from 0 to 255. If pixel values range from 0 to 1, set `do_rescale=False`.
+            vision_info (`List[Dict]`, *optional*):
+                Optional list of dictionaries containing additional information about vision inputs.
+            do_resize (`bool`, *optional*, defaults to `self.do_resize`):
+                Whether to resize the image.
+            resample (`PILImageResampling`, *optional*, defaults to `self.resample`):
+                Resampling filter to use if resizing the image. This can be one of the `PILImageResampling` enums.
+            do_rescale (`bool`, *optional*, defaults to `self.do_rescale`):
+                Whether to rescale the image.
+            rescale_factor (`float`, *optional*, defaults to `self.rescale_factor`):
+                Scale factor to use if rescaling the image.
+            do_normalize (`bool`, *optional*, defaults to `self.do_normalize`):
+                Whether to normalize the image.
+            image_mean (`float` or `List[float]`, *optional*, defaults to `self.image_mean`):
+                Mean to use if normalizing the image. Can be a float or a list of floats corresponding to the number of channels in the image.
+            image_std (`float` or `List[float]`, *optional*, defaults to `self.image_std`):
+                Standard deviation to use if normalizing the image. Can be a float or a list of floats corresponding to the number of channels in the image.
+            do_convert_rgb (`bool`, *optional*, defaults to `self.do_convert_rgb`):
+                Whether to convert the image to RGB.
+            data_format (`ChannelDimension`, *optional*, defaults to `ChannelDimension.FIRST`):
+                The channel dimension format for the output image. Can be one of:
+                - `"channels_first"` or `ChannelDimension.FIRST`: image in (num_channels, height, width) format.
+                - `"channels_last"` or `ChannelDimension.LAST`: image in (height, width, num_channels) format.
+                - Unset: Use the channel dimension format of the input image.
+            input_data_format (`ChannelDimension` or `str`, *optional*):
+                The channel dimension format for the input image. Can be one of:
+                - `"channels_first"` or `ChannelDimension.FIRST`: image in (num_channels, height, width) format.
+                - `"channels_last"` or `ChannelDimension.LAST`: image in (height, width, num_channels) format.
+                - `"none"` or `ChannelDimension.NONE`: image in (height, width) format.   - `"none"` or `ChannelDimension.NONE`: image in (height, width) format.
+        """
+        images = make_list_of_images(images)
+
+        if do_convert_rgb:
+            images = [convert_to_rgb(image) for image in images]
+
+        # All transformations expect numpy arrays.
+        images = [to_numpy_array(image) for image in images]
+
+        if is_scaled_image(images[0]) and do_rescale:
+            logger.warning_once(
+                "It looks like you are trying to rescale already rescaled images. If the input"
+                " images have pixel values between 0 and 1, set `do_rescale=False` to avoid rescaling them again."
+            )
+        if input_data_format is None:
+            # We assume that all images have the same channel dimension format.
+            input_data_format = infer_channel_dimension_format(images[0])
+
+        height, width = get_image_size(images[0], channel_dim=input_data_format)
+        resized_height, resized_width = height, width
+        processed_images = []
+        for image in images:
+            if do_resize:
+                resized_height, resized_width = smart_resize(
+                    height,
+                    width,
+                    factor=self.patch_size * self.merge_size,
+                    min_pixels=self.min_pixels,
+                    max_pixels=self.max_pixels,
+                )
+                image = resize(
+                    image, size=(resized_height, resized_width), resample=resample, input_data_format=input_data_format
+                )
+
+            if do_rescale:
+                image = self.rescale(image, scale=rescale_factor, input_data_format=input_data_format)
+
+            if do_normalize:
+                image = self.normalize(
+                    image=image, mean=image_mean, std=image_std, input_data_format=input_data_format
+                )
+
+            image = to_channel_dimension_format(image, data_format, input_channel_dim=input_data_format)
+            processed_images.append(image)
+
+        patches = np.array(processed_images)
+        if data_format == ChannelDimension.LAST:
+            patches = patches.transpose(0, 3, 1, 2)
+        if patches.shape[0] == 1:
+            patches = np.tile(patches, (self.temporal_patch_size, 1, 1, 1))
+        channel = patches.shape[1]
+        grid_t = patches.shape[0] // self.temporal_patch_size
+        grid_h, grid_w = resized_height // self.patch_size, resized_width // self.patch_size
+        patches = patches.reshape(
+            grid_t,
+            self.temporal_patch_size,
+            channel,
+            grid_h // self.merge_size,
+            self.merge_size,
+            self.patch_size,
+            grid_w // self.merge_size,
+            self.merge_size,
+            self.patch_size,
+        )
+        patches = patches.transpose(0, 3, 6, 4, 7, 2, 1, 5, 8)
+        flatten_patches = patches.reshape(
+            grid_t * grid_h * grid_w, channel * self.temporal_patch_size * self.patch_size * self.patch_size
+        )
+
+        return flatten_patches, (grid_t, grid_h, grid_w)
+
+    def preprocess(
+        self,
+        images: ImageInput,
+        videos: VideoInput = None,
+        do_resize: bool = None,
+        size: Dict[str, int] = None,
+        resample: PILImageResampling = None,
+        do_rescale: bool = None,
+        rescale_factor: float = None,
+        do_normalize: bool = None,
+        image_mean: Optional[Union[float, List[float]]] = None,
+        image_std: Optional[Union[float, List[float]]] = None,
+        do_convert_rgb: bool = None,
+        return_tensors: Optional[Union[str, TensorType]] = None,
+        data_format: Optional[ChannelDimension] = ChannelDimension.FIRST,
+        input_data_format: Optional[Union[str, ChannelDimension]] = None,
+    ):
+        """
+        Args:
+            images (`ImageInput`):
+                Image to preprocess. Expects a single or batch of images with pixel values ranging from 0 to 255. If
+                passing in images with pixel values between 0 and 1, set `do_rescale=False`.
+            videos (`VideoInput`):
+                Video to preprocess. Expects a single or batch of videos with pixel values ranging from 0 to 255. If
+                passing in videos with pixel values between 0 and 1, set `do_rescale=False`.
+            do_resize (`bool`, *optional*, defaults to `self.do_resize`):
+                Whether to resize the image.
+            size (`Dict[str, int]`, *optional*, defaults to `self.size`):
+                Size of the image after resizing. Shortest edge of the image is resized to size["shortest_edge"], with
+                the longest edge resized to keep the input aspect ratio.
+            resample (`int`, *optional*, defaults to `self.resample`):
+                Resampling filter to use if resizing the image. This can be one of the enum `PILImageResampling`. Only
+                has an effect if `do_resize` is set to `True`.
+            do_rescale (`bool`, *optional*, defaults to `self.do_rescale`):
+                Whether to rescale the image.
+            rescale_factor (`float`, *optional*, defaults to `self.rescale_factor`):
+                Rescale factor to rescale the image by if `do_rescale` is set to `True`.
+            do_normalize (`bool`, *optional*, defaults to `self.do_normalize`):
+                Whether to normalize the image.
+            image_mean (`float` or `List[float]`, *optional*, defaults to `self.image_mean`):
+                Image mean to use for normalization. Only has an effect if `do_normalize` is set to `True`.
+            image_std (`float` or `List[float]`, *optional*, defaults to `self.image_std`):
+                Image standard deviation to use for normalization. Only has an effect if `do_normalize` is set to
+                `True`.
+            do_convert_rgb (`bool`, *optional*, defaults to `self.do_convert_rgb`):
+                Whether to convert the image to RGB.
+            return_tensors (`str` or `TensorType`, *optional*):
+                The type of tensors to return. Can be one of:
+                - Unset: Return a list of `np.ndarray`.
+                - `TensorType.TENSORFLOW` or `'tf'`: Return a batch of type `tf.Tensor`.
+                - `TensorType.PYTORCH` or `'pt'`: Return a batch of type `torch.Tensor`.
+                - `TensorType.NUMPY` or `'np'`: Return a batch of type `np.ndarray`.
+                - `TensorType.JAX` or `'jax'`: Return a batch of type `jax.numpy.ndarray`.
+            data_format (`ChannelDimension` or `str`, *optional*, defaults to `ChannelDimension.FIRST`):
+                The channel dimension format for the output image. Can be one of:
+                - `"channels_first"` or `ChannelDimension.FIRST`: image in (num_channels, height, width) format.
+                - `"channels_last"` or `ChannelDimension.LAST`: image in (height, width, num_channels) format.
+                - Unset: Use the channel dimension format of the input image.
+            input_data_format (`ChannelDimension` or `str`, *optional*):
+                The channel dimension format for the input image. If unset, the channel dimension format is inferred
+                from the input image. Can be one of:
+                - `"channels_first"` or `ChannelDimension.FIRST`: image in (num_channels, height, width) format.
+                - `"channels_last"` or `ChannelDimension.LAST`: image in (height, width, num_channels) format.
+                - `"none"` or `ChannelDimension.NONE`: image in (height, width) format.
+
+        """
+        do_resize = do_resize if do_resize is not None else self.do_resize
+        size = size if size is not None else self.size
+        resample = resample if resample is not None else self.resample
+        do_rescale = do_rescale if do_rescale is not None else self.do_rescale
+        rescale_factor = rescale_factor if rescale_factor is not None else self.rescale_factor
+        do_normalize = do_normalize if do_normalize is not None else self.do_normalize
+        image_mean = image_mean if image_mean is not None else self.image_mean
+        image_std = image_std if image_std is not None else self.image_std
+        do_convert_rgb = do_convert_rgb if do_convert_rgb is not None else self.do_convert_rgb
+
+        if images is not None:
+            images = make_batched_images(images)
+        if videos is not None:
+            videos = make_batched_videos(videos)
+
+        if images is not None and not valid_images(images):
+            raise ValueError(
+                "Invalid image type. Must be of type PIL.Image.Image, numpy.ndarray, "
+                "torch.Tensor, tf.Tensor or jax.ndarray."
+            )
+
+        validate_preprocess_arguments(
+            rescale_factor=rescale_factor,
+            do_normalize=do_normalize,
+            image_mean=image_mean,
+            image_std=image_std,
+            do_resize=do_resize,
+            size=size,
+            resample=resample,
+        )
+
+        if images is not None:
+            pixel_values, vision_grid_thws = [], []
+            for image in images:
+                patches, image_grid_thw = self._preprocess(
+                    image,
+                    do_resize=do_resize,
+                    resample=resample,
+                    do_rescale=do_rescale,
+                    rescale_factor=rescale_factor,
+                    do_normalize=do_normalize,
+                    image_mean=image_mean,
+                    image_std=image_std,
+                    data_format=data_format,
+                    do_convert_rgb=do_convert_rgb,
+                    input_data_format=input_data_format,
+                )
+                pixel_values.extend(patches)
+                vision_grid_thws.append(image_grid_thw)
+            pixel_values = np.array(pixel_values)
+            vision_grid_thws = np.array(vision_grid_thws)
+            data = {"pixel_values": pixel_values, "image_grid_thw": vision_grid_thws}
+
+        if videos is not None:
+            pixel_values, vision_grid_thws = [], []
+            for images in videos:
+                patches, video_grid_thw = self._preprocess(
+                    images,
+                    do_resize=do_resize,
+                    resample=resample,
+                    do_rescale=do_rescale,
+                    rescale_factor=rescale_factor,
+                    do_normalize=do_normalize,
+                    image_mean=image_mean,
+                    image_std=image_std,
+                    data_format=data_format,
+                    do_convert_rgb=do_convert_rgb,
+                    input_data_format=input_data_format,
+                )
+                pixel_values.extend(patches)
+                vision_grid_thws.append(video_grid_thw)
+            pixel_values = np.array(pixel_values)
+            vision_grid_thws = np.array(vision_grid_thws)
+            data = {"pixel_values_videos": pixel_values, "video_grid_thw": vision_grid_thws}
+
+        return BatchFeature(data=data, tensor_type=return_tensors)

code/infer.py

@@ -0,0 +1,468 @@
+#!/usr/bin/env python3
+"""Unified Hugging Face inference entry point for Ming image checkpoints."""
+
+from __future__ import annotations
+
+import argparse
+import json
+import re
+import sys
+from pathlib import Path
+from typing import Iterable, List
+
+from inference_profile import (
+    VALID_TASKS,
+    load_checkpoint_capabilities,
+    resolve_model_directory,
+)
+from mllm_device_map import (
+    build_mllm_device_plan,
+    load_mllm_num_hidden_layers,
+    validate_loaded_layer_devices,
+)
+CODE_DIRECTORY = Path(__file__).resolve().parent
+
+TASK_RESOLUTION_BUCKETS = {
+    "text-to-image": (1024, 2048),
+    "image-edit": (1024,),
+    "layer-decompose": (512, 1024),
+}
+TASK_DEFAULT_RESOLUTIONS = {
+    "text-to-image": 2048,
+    "image-edit": 1024,
+    "layer-decompose": 1024,
+}
+
+
+def parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser(
+        description=(
+            "Run text-to-image, image editing, or layer decomposition with an "
+            "explicit checkpoint capability profile."
+        )
+    )
+    parser.add_argument("--model", required=True, help="Local model directory or HF Hub repo ID")
+    parser.add_argument("--task", required=True, choices=sorted(VALID_TASKS))
+    parser.add_argument("--prompt", help="Generation/edit prompt; optional for layer decomposition")
+    parser.add_argument("--input-image", type=Path, help="Required for edit and layer decomposition")
+    parser.add_argument("--num-layers", type=int, default=1)
+    parser.add_argument("--output-dir", type=Path, default=Path("outputs"))
+    parser.add_argument(
+        "--resolution",
+        type=int,
+        help=(
+            "Requested resolution bucket. Defaults: text-to-image 2048, "
+            "image-edit 1024, layer-decompose 1024. Requests snap to the "
+            "nearest bucket supported by the selected task."
+        ),
+    )
+    parser.add_argument(
+        "--steps",
+        type=int,
+        help="Override the checkpoint-family default (generation: 12, layers: 12)",
+    )
+    parser.add_argument("--seed", type=int, default=42)
+    parser.add_argument(
+        "--cfg",
+        type=float,
+        help="Override the checkpoint-family default (generation: 1.0, layers: 2.0)",
+    )
+    parser.add_argument("--dtype", choices=("bfloat16", "float16", "float32"), default="bfloat16")
+    parser.add_argument(
+        "--attn-implementation",
+        choices=("sdpa", "flash_attention_2", "eager"),
+        # The BailingMoeV2 LLM only implements eager and flash_attention_2
+        # attention classes; selecting "sdpa" fails closed at load time.
+        default="eager",
+    )
+    parser.add_argument("--device", default="cuda:0")
+    parser.add_argument(
+        "--device-map",
+        choices=("balanced", "none", "auto"),
+        default="balanced",
+        help="balanced reserves GPU 0 for fixed image modules and shards the MLLM",
+    )
+    parser.add_argument(
+        "--num-gpus",
+        type=int,
+        default=1,
+        help="required visible GPU count for balanced placement (any positive integer)",
+    )
+    parser.add_argument(
+        "--processor",
+        help="Optional processor data directory; defaults to <checkpoint>/mllm",
+    )
+    parser.add_argument("--revision", help="HF Hub model revision")
+    parser.add_argument("--cache-dir", type=Path)
+    parser.add_argument("--local-files-only", action="store_true")
+    parser.add_argument(
+        "--validate-only",
+        action="store_true",
+        help="Validate model profile and task arguments without loading weights",
+    )
+    parser.add_argument(
+        "--attention-bf16-reduction",
+        action="store_true",
+        help=(
+            "Let PyTorch's math attention kernel (the only SDPA kernel that runs on ROCm gfx1151) "
+            "stay in bf16 instead of upcasting to fp32: faster, less precise."
+        ),
+    )
+    parser.add_argument(
+        "--release-mllm-after-conditioning",
+        action="store_true",
+        help=(
+            "Free the MLLM, vision tower and connector as soon as the conditioning is computed, "
+            "before the diffusion steps. Lowers peak memory; one image per process."
+        ),
+    )
+    return parser.parse_args()
+
+
+def parse_num_layers(text: str) -> int:
+    decompose_match = re.search(
+        r"decompose this image into\s+(\d+)\s+layers?", text.lower()
+    )
+    if decompose_match:
+        return int(decompose_match.group(1))
+    for line in text.splitlines():
+        s = line.strip().lower()
+        if s.startswith("number of layers:"):
+            try:
+                return int(s.split(":", 1)[1].strip())
+            except ValueError:
+                pass
+    return 5
+
+
+def resolve_task_resolution(task: str, requested: int | None) -> int:
+    """Resolve a user request to the nearest supported task-level bucket."""
+    try:
+        buckets = TASK_RESOLUTION_BUCKETS[task]
+    except KeyError as exc:
+        raise ValueError(f"unsupported task for resolution policy: {task!r}") from exc
+    if requested is None:
+        return TASK_DEFAULT_RESOLUTIONS[task]
+    if isinstance(requested, bool) or not isinstance(requested, int) or requested <= 0:
+        raise ValueError("--resolution must be a positive integer")
+    return min(buckets, key=lambda value: (abs(value - requested), value))
+
+
+def _load_prompt(prompt: str) -> str:
+    prompt_path = Path(prompt)
+    try:
+        is_file = prompt_path.is_file()
+    except OSError:
+        # Long literal prompts can be invalid filesystem paths. In that case,
+        # keep treating the argument as prompt text.
+        return prompt
+    if is_file:
+        return prompt_path.read_text(encoding="utf-8")
+    return prompt
+
+
+def _dtype(name: str):
+    import torch
+
+    return {
+        "bfloat16": torch.bfloat16,
+        "float16": torch.float16,
+        "float32": torch.float32,
+    }[name]
+
+
+def _build_messages(task: str, prompt: str, input_image: Path | None):
+    content = []
+    if input_image is not None:
+        content.append({"type": "image", "image": str(input_image)})
+    content.append({"type": "text", "text": prompt})
+    return [{"role": "HUMAN", "content": content}]
+
+
+def _normalize_outputs(output) -> List[Image.Image]:
+    from PIL import Image
+
+    if isinstance(output, Image.Image):
+        return [output]
+    if isinstance(output, (list, tuple)) and all(
+        isinstance(item, Image.Image) for item in output
+    ):
+        return list(output)
+    raise TypeError(f"model returned unsupported output type: {type(output)!r}")
+
+
+def _model_input_device(model):
+    try:
+        return model.device
+    except AttributeError:
+        return next(parameter.device for parameter in model.parameters() if not parameter.is_meta)
+
+
+def _validate_balanced_placement(model, plan, torch) -> None:
+    layer_devices = []
+    for index, layer in enumerate(model.model.model.layers):
+        devices = {parameter.device for parameter in layer.parameters()}
+        if len(devices) != 1:
+            raise RuntimeError(f"MLLM layer {index} spans devices {sorted(map(str, devices))}")
+        device = devices.pop()
+        if device.type != "cuda" or device.index is None:
+            raise RuntimeError(f"MLLM layer {index} loaded on {device}, expected CUDA")
+        layer_devices.append(device.index)
+    validate_loaded_layer_devices(layer_devices, plan)
+
+    fixed_modules = {
+        "vision": model.vision,
+        "linear_proj": model.linear_proj,
+        "connector": getattr(model, "connector", None),
+        "proj_in": getattr(model, "proj_in", None),
+        "proj_out": getattr(model, "proj_out", None),
+        "proj_directvlm": getattr(model, "proj_directvlm", None),
+        "query_tokens": getattr(model, "query_tokens_dict", None),
+        "diffusion": getattr(model, "diffusion_loss", None),
+    }
+    expected = {torch.device("cuda:0")}
+    for name, module in fixed_modules.items():
+        if module is None:
+            continue
+        devices = {parameter.device for parameter in module.parameters()}
+        if devices and devices != expected:
+            raise RuntimeError(
+                f"fixed module {name} loaded on {sorted(map(str, devices))}, expected cuda:0"
+            )
+
+
+def _save_outputs(
+    images: Iterable, output_dir: Path, task: str
+) -> List[Path]:
+    output_dir.mkdir(parents=True, exist_ok=True)
+    prefix = "layer" if task == "layer-decompose" else "image"
+    skip_first = task == "layer-decompose"
+    paths = []
+    for index, image in enumerate(images):
+        if skip_first and index == 0:
+            continue
+        output_path = output_dir / f"{prefix}_{index:02d}.png"
+        image.save(output_path)
+        paths.append(output_path)
+    return paths
+
+
+def main() -> None:
+    args = parse_args()
+    model_directory = resolve_model_directory(
+        args.model,
+        revision=args.revision,
+        cache_dir=args.cache_dir,
+        local_files_only=args.local_files_only,
+    )
+    profile = load_checkpoint_capabilities(model_directory)
+
+    has_reference_image = args.input_image is not None
+    profile.validate_task(
+        args.task,
+        has_reference_image=has_reference_image,
+        num_layers=args.num_layers,
+    )
+    effective_resolution = resolve_task_resolution(args.task, args.resolution)
+    if args.resolution is not None and args.resolution != effective_resolution:
+        print(
+            f"resolution {args.resolution} snapped to {effective_resolution} "
+            f"for task {args.task}",
+            file=sys.stderr,
+        )
+    sampling = profile.resolve_sampling_parameters(steps=args.steps, cfg=args.cfg)
+    if args.input_image is not None and not args.input_image.is_file():
+        raise FileNotFoundError(f"input image does not exist: {args.input_image}")
+    if args.task != "layer-decompose" and not args.prompt:
+        raise ValueError(f"--prompt is required for {args.task}")
+
+    if args.prompt is not None:
+        prompt = _load_prompt(args.prompt)
+    else:
+        prompt = f"Decompose this image into {args.num_layers} layers."
+    num_layers = parse_num_layers(prompt) if args.task == "layer-decompose" else args.num_layers
+    if args.validate_only:
+        print(
+            json.dumps(
+                {
+                    "model": str(model_directory),
+                    "task": args.task,
+                    "profile": profile.__dict__,
+                    "sampling": sampling.__dict__,
+                    "resolution": {
+                        "requested": args.resolution,
+                        "effective": effective_resolution,
+                    },
+                },
+                indent=2,
+            )
+        )
+        return
+
+    model, processor = load_model_and_processor(model_directory, args)
+    images = run_generation(
+        model,
+        processor,
+        profile,
+        task=args.task,
+        prompt=prompt,
+        input_image=args.input_image,
+        resolution=effective_resolution,
+        sampling=sampling,
+        seed=args.seed,
+        num_layers=num_layers,
+        dtype=_dtype(args.dtype),
+    )
+
+    output_paths = _save_outputs(images, args.output_dir, args.task)
+    print(json.dumps({"outputs": [str(path.resolve()) for path in output_paths]}, indent=2))
+
+
+def load_model_and_processor(model_directory: Path, args):
+    """Load the model and processor once; shared by the CLI and batch tools."""
+    import torch
+
+    from modeling_bailingmm2 import BailingMM2NativeForConditionalGeneration
+    from processing_bailingmm2 import load_bailingmm2_processor
+
+    # Processor/tokenizer data lives in the package's mllm/ component; the
+    # Python implementations stay in this repository (AutoProcessor would
+    # require them inside the data directory). --processor overrides the data
+    # directory only.
+    processor_directory = (
+        Path(args.processor).expanduser().resolve()
+        if args.processor
+        else model_directory / "mllm"
+    )
+    processor = load_bailingmm2_processor(processor_directory)
+
+    if getattr(args, "attention_bf16_reduction", False):
+        # The math SDPA kernel upcasts bf16 inputs to fp32 by default; this keeps it in bf16.
+        torch.backends.cuda.allow_fp16_bf16_reduction_math_sdp(True)
+
+    dtype = _dtype(args.dtype)
+    load_kwargs = {
+        "torch_dtype": dtype,
+        "attn_implementation": args.attn_implementation,
+        "load_image_gen": True,
+        "image_gen_device": args.device,
+    }
+    device_plan = None
+    if args.device_map == "balanced":
+        visible_gpus = torch.cuda.device_count()
+        if visible_gpus != args.num_gpus:
+            raise RuntimeError(
+                f"balanced placement requires {args.num_gpus} visible GPUs, got {visible_gpus}; "
+                "set CUDA_VISIBLE_DEVICES before starting Python"
+            )
+        num_hidden_layers = load_mllm_num_hidden_layers(model_directory / "mllm")
+        device_plan = build_mllm_device_plan(num_hidden_layers, visible_gpus)
+        load_kwargs["device_map"] = device_plan.device_map
+        print(
+            f"MLLM device plan: {visible_gpus} GPUs, layer_counts={device_plan.layer_counts}"
+        )
+    elif args.device_map == "auto":
+        load_kwargs["device_map"] = "auto"
+
+    model = BailingMM2NativeForConditionalGeneration.from_pretrained(
+        str(model_directory), **load_kwargs
+    )
+    if args.device_map == "none":
+        model = model.to(device=args.device, dtype=dtype)
+    elif device_plan is not None:
+        _validate_balanced_placement(model, device_plan, torch)
+    if getattr(args, "release_mllm_after_conditioning", False):
+        _release_mllm_before_sampling(model)
+    return model, processor
+
+
+def _release_mllm_before_sampling(model) -> None:
+    """Free the MLLM-side modules once the conditioning exists (--release-mllm-after-conditioning).
+
+    Wraps the diffusion sampler: by the time it is called the conditioning tensors are computed,
+    so the language model, vision tower and connector are moved to the meta device (releasing
+    their memory) before the diffusion steps start. The model cannot generate again afterwards.
+    """
+    import gc
+
+    import torch
+
+    original_sample = model.diffusion_loss.sample
+
+    def sample_after_release(*args, **kwargs):
+        for name in ("model", "vision", "linear_proj", "connector"):
+            module = getattr(model, name, None)
+            if module is not None:
+                module.to("meta")
+        gc.collect()
+        if torch.cuda.is_available():
+            torch.cuda.empty_cache()
+        return original_sample(*args, **kwargs)
+
+    model.diffusion_loss.sample = sample_after_release
+
+
+def run_generation(
+    model,
+    processor,
+    profile,
+    *,
+    task: str,
+    prompt: str,
+    input_image,
+    resolution: int,
+    sampling,
+    seed: int,
+    num_layers: int,
+    dtype,
+) -> List:
+    """Run one inference with an already-loaded model and processor."""
+    import torch
+    from PIL import Image
+
+    resolution = resolve_task_resolution(task, resolution)
+
+    messages = _build_messages(task, prompt, input_image)
+    text = processor.apply_chat_template(messages, add_generation_prompt=True)
+    image_inputs, video_inputs, _ = processor.process_vision_info(messages)
+
+    reference_image = None
+    if input_image is not None:
+        reference_mode = "RGB" if profile.vae_input_channels == 3 else "RGBA"
+        reference_image = Image.open(input_image).convert(reference_mode)
+
+    inputs = processor(
+        text=[text],
+        images=image_inputs,
+        videos=video_inputs,
+        return_tensors="pt",
+        image_gen_highres=resolution,
+        image_gen_ref_images=reference_image,
+        image_gen_input_channels=profile.vae_input_channels,
+    )
+    input_device = _model_input_device(model)
+    inputs = inputs.to(input_device)
+    for key, value in inputs.items():
+        if isinstance(value, torch.Tensor) and torch.is_floating_point(value):
+            inputs[key] = value.to(dtype=dtype)
+
+    output = model.generate(
+        **inputs,
+        image_gen=True,
+        image_gen_task=task,
+        image_gen_seed=seed,
+        image_gen_steps=sampling.steps,
+        image_gen_cfg=sampling.cfg,
+        num_frames_per_prompt=num_layers+1 if task == "layer-decompose" else num_layers,
+    )
+    images = _normalize_outputs(output)
+    expected_outputs = num_layers + 1 if task == "layer-decompose" else 1
+    if len(images) != expected_outputs:
+        raise RuntimeError(
+            f"expected {expected_outputs} output image(s), got {len(images)}"
+        )
+    return images
+
+
+if __name__ == "__main__":
+    main()

code/inference_profile.py

@@ -0,0 +1,420 @@
+"""Checkpoint capability contract for public Ming image inference.
+
+New packages declare their capability in ``transformer/config.json`` via the
+``alignment_padding_mode`` / ``multi_frame_output`` pair; the VAE contract is
+derived from ``vae/config.json``. Legacy packages that predate the component
+metadata are loaded strictly from the root ``inference_profile.json`` during
+the compatibility window. The contract is intentionally strict: selecting
+behavior from directory names, missing state-dict keys, or task names can
+silently load the wrong padding semantics and produce degraded images.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+import json
+import math
+from pathlib import Path
+from typing import Any, Mapping, Optional, Union
+
+
+PROFILE_FILENAME = "inference_profile.json"
+PROFILE_SCHEMA_VERSION = 1
+
+GENERATION_PROFILE = "generation_edit"
+LAYER_PROFILE = "layer_decompose"
+VALID_PROFILES = {GENERATION_PROFILE, LAYER_PROFILE}
+
+LEARNED_PADDING = "learned"
+ZERO_MASKED_PADDING = "zero_masked"
+VALID_PADDING_MODES = {LEARNED_PADDING, ZERO_MASKED_PADDING}
+
+VALID_VAE_SAMPLE_MODES = {"sample", "argmax"}
+VALID_TASKS = {"text-to-image", "image-edit", "layer-decompose"}
+
+REQUIRED_PROFILE_FIELDS = {
+    "schema_version",
+    "inference_profile",
+    "alignment_padding_mode",
+    "multi_frame_output",
+    "vae_input_channels",
+    "vae_sample_mode",
+}
+
+
+class InferenceProfileError(ValueError):
+    """Raised when a checkpoint does not satisfy the inference contract."""
+
+
+@dataclass(frozen=True)
+class SamplingParameters:
+    steps: int
+    cfg: float
+
+
+DEFAULT_SAMPLING_PARAMETERS = {
+    GENERATION_PROFILE: SamplingParameters(steps=12, cfg=1.0),
+    LAYER_PROFILE: SamplingParameters(steps=12, cfg=2.0),
+}
+
+
+@dataclass(frozen=True)
+class InferenceProfile:
+    schema_version: int
+    inference_profile: str
+    alignment_padding_mode: str
+    multi_frame_output: bool
+    vae_input_channels: int
+    vae_sample_mode: str
+
+    @classmethod
+    def from_dict(cls, raw: Mapping[str, Any]) -> "InferenceProfile":
+        missing = sorted(REQUIRED_PROFILE_FIELDS.difference(raw))
+        if missing:
+            raise InferenceProfileError(
+                "checkpoint inference profile is missing required fields: "
+                + ", ".join(missing)
+            )
+
+        unknown = sorted(set(raw).difference(REQUIRED_PROFILE_FIELDS))
+        if unknown:
+            raise InferenceProfileError(
+                "checkpoint inference profile contains unsupported fields: "
+                + ", ".join(unknown)
+            )
+
+        if type(raw["schema_version"]) is not int:
+            raise InferenceProfileError("schema_version must be an integer")
+        if type(raw["multi_frame_output"]) is not bool:
+            raise InferenceProfileError("multi_frame_output must be a boolean")
+        if type(raw["vae_input_channels"]) is not int:
+            raise InferenceProfileError("vae_input_channels must be an integer")
+
+        profile = cls(**{key: raw[key] for key in REQUIRED_PROFILE_FIELDS})
+        profile.validate()
+        return profile
+
+    def validate(self) -> None:
+        if self.schema_version != PROFILE_SCHEMA_VERSION:
+            raise InferenceProfileError(
+                f"unsupported profile schema_version={self.schema_version}; "
+                f"expected {PROFILE_SCHEMA_VERSION}"
+            )
+        if self.inference_profile not in VALID_PROFILES:
+            raise InferenceProfileError(
+                f"inference_profile must be one of {sorted(VALID_PROFILES)}, "
+                f"got {self.inference_profile!r}"
+            )
+        if self.alignment_padding_mode not in VALID_PADDING_MODES:
+            raise InferenceProfileError(
+                "alignment_padding_mode must be one of "
+                f"{sorted(VALID_PADDING_MODES)}, got {self.alignment_padding_mode!r}"
+            )
+        if self.vae_input_channels not in (3, 4):
+            raise InferenceProfileError(
+                f"vae_input_channels must be 3 or 4, got {self.vae_input_channels}"
+            )
+        if self.vae_sample_mode not in VALID_VAE_SAMPLE_MODES:
+            raise InferenceProfileError(
+                f"vae_sample_mode must be one of {sorted(VALID_VAE_SAMPLE_MODES)}, "
+                f"got {self.vae_sample_mode!r}"
+            )
+
+        if self.inference_profile == GENERATION_PROFILE:
+            if self.alignment_padding_mode != ZERO_MASKED_PADDING:
+                raise InferenceProfileError(
+                    "generation_edit checkpoints must use zero_masked alignment padding"
+                )
+            if self.multi_frame_output:
+                raise InferenceProfileError(
+                    "generation_edit checkpoints must set multi_frame_output=false"
+                )
+            if self.vae_input_channels != 4:
+                raise InferenceProfileError(
+                    "generation_edit checkpoints must declare vae_input_channels=4"
+                )
+            if self.vae_sample_mode != "argmax":
+                raise InferenceProfileError(
+                    "generation_edit checkpoints must declare vae_sample_mode='argmax'"
+                )
+        else:
+            if self.alignment_padding_mode != LEARNED_PADDING:
+                raise InferenceProfileError(
+                    "layer_decompose checkpoints must use learned alignment padding"
+                )
+            if not self.multi_frame_output:
+                raise InferenceProfileError(
+                    "layer_decompose checkpoints must set multi_frame_output=true"
+                )
+            if self.vae_input_channels != 4:
+                raise InferenceProfileError(
+                    "layer_decompose checkpoints must declare vae_input_channels=4"
+                )
+            if self.vae_sample_mode != "argmax":
+                raise InferenceProfileError(
+                    "layer_decompose checkpoints must declare vae_sample_mode='argmax'"
+                )
+
+    def validate_task(
+        self,
+        task: str,
+        *,
+        has_reference_image: bool,
+        num_layers: int = 1,
+    ) -> None:
+        if task not in VALID_TASKS:
+            raise InferenceProfileError(
+                f"task must be one of {sorted(VALID_TASKS)}, got {task!r}"
+            )
+        if num_layers < 1:
+            raise InferenceProfileError("num_layers must be at least 1")
+
+        if task == "layer-decompose":
+            if self.inference_profile != LAYER_PROFILE:
+                raise InferenceProfileError(
+                    "layer-decompose requires a layer_decompose checkpoint"
+                )
+            if not has_reference_image:
+                raise InferenceProfileError(
+                    "layer-decompose requires an input reference image"
+                )
+            return
+
+        if self.inference_profile != GENERATION_PROFILE:
+            raise InferenceProfileError(
+                f"{task} requires a generation_edit checkpoint"
+            )
+        if num_layers != 1:
+            raise InferenceProfileError(
+                f"{task} does not support num_layers={num_layers}; expected 1"
+            )
+        if task == "text-to-image" and has_reference_image:
+            raise InferenceProfileError("text-to-image does not accept a reference image")
+        if task == "image-edit" and not has_reference_image:
+            raise InferenceProfileError("image-edit requires a reference image")
+
+    def resolve_sampling_parameters(
+        self,
+        *,
+        steps: Optional[int] = None,
+        cfg: Optional[float] = None,
+    ) -> SamplingParameters:
+        defaults = DEFAULT_SAMPLING_PARAMETERS[self.inference_profile]
+        resolved_steps = defaults.steps if steps is None else steps
+        resolved_cfg = defaults.cfg if cfg is None else cfg
+
+        if type(resolved_steps) is not int or resolved_steps < 1:
+            raise InferenceProfileError("sampling steps must be an integer >= 1")
+        if (
+            isinstance(resolved_cfg, bool)
+            or not isinstance(resolved_cfg, (int, float))
+            or not math.isfinite(float(resolved_cfg))
+            or resolved_cfg < 0
+        ):
+            raise InferenceProfileError("CFG must be a finite number >= 0")
+
+        return SamplingParameters(
+            steps=resolved_steps,
+            cfg=float(resolved_cfg),
+        )
+
+
+def load_inference_profile(model_directory: Union[str, Path]) -> InferenceProfile:
+    """Legacy parser: strict root ``inference_profile.json`` (see module doc)."""
+
+    model_directory = Path(model_directory)
+    profile_path = model_directory / PROFILE_FILENAME
+    if not profile_path.is_file():
+        raise InferenceProfileError(
+            f"checkpoint must contain {PROFILE_FILENAME}: {profile_path}"
+        )
+    try:
+        raw = json.loads(profile_path.read_text(encoding="utf-8"))
+    except (OSError, json.JSONDecodeError) as error:
+        raise InferenceProfileError(
+            f"failed to read checkpoint inference profile {profile_path}: {error}"
+        ) from error
+    if not isinstance(raw, dict):
+        raise InferenceProfileError("checkpoint inference profile must be a JSON object")
+    return InferenceProfile.from_dict(raw)
+
+
+TRANSFORMER_CONFIG_FILENAME = "transformer/config.json"
+VAE_CONFIG_FILENAME = "vae/config.json"
+CAPABILITY_FIELDS = ("alignment_padding_mode", "multi_frame_output")
+QWEN_VAE_CLASS_NAME = "AutoencoderKLQwenImage"
+
+# The only valid capability pairs and the runtime family they select.
+CAPABILITY_PROFILES = {
+    (ZERO_MASKED_PADDING, False): GENERATION_PROFILE,
+    (LEARNED_PADDING, True): LAYER_PROFILE,
+}
+
+
+def _read_component_config(model_directory: Path, relative: str) -> Mapping[str, Any]:
+    config_path = model_directory / relative
+    if not config_path.is_file():
+        raise InferenceProfileError(
+            f"checkpoint is missing component config {relative}: {config_path}"
+        )
+    try:
+        raw = json.loads(config_path.read_text(encoding="utf-8"))
+    except (OSError, json.JSONDecodeError) as error:
+        raise InferenceProfileError(
+            f"failed to read component config {config_path}: {error}"
+        ) from error
+    if not isinstance(raw, dict):
+        raise InferenceProfileError(f"component config {relative} must be a JSON object")
+    return raw
+
+
+def _derive_vae_contract(model_directory: Path) -> tuple[int, str]:
+    """(vae_input_channels, vae_sample_mode) from the VAE component config."""
+
+    config = _read_component_config(model_directory, VAE_CONFIG_FILENAME)
+    class_name = config.get("_class_name")
+    if class_name != QWEN_VAE_CLASS_NAME:
+        raise InferenceProfileError(
+            f"unsupported VAE contract: {VAE_CONFIG_FILENAME} _class_name must be "
+            f"{QWEN_VAE_CLASS_NAME!r} for argmax reference encoding, got "
+            f"{class_name!r}"
+        )
+    declared = [
+        config[key]
+        for key in ("input_channels", "in_channels")
+        if key in config
+    ]
+    if not declared:
+        raise InferenceProfileError(
+            f"{VAE_CONFIG_FILENAME} must declare input_channels or in_channels"
+        )
+    if len(set(declared)) != 1:
+        raise InferenceProfileError(
+            f"{VAE_CONFIG_FILENAME} input_channels and in_channels disagree: "
+            f"{declared}"
+        )
+    channels = declared[0]
+    if type(channels) is not int or channels != 4:
+        raise InferenceProfileError(
+            "the supported public families require a 4-channel VAE, got "
+            f"{channels!r} in {VAE_CONFIG_FILENAME}"
+        )
+    return channels, "argmax"
+
+
+def _profile_from_components(
+    model_directory: Path,
+    alignment_padding_mode: Any,
+    multi_frame_output: Any,
+) -> InferenceProfile:
+    if type(alignment_padding_mode) is not str:
+        raise InferenceProfileError(
+            f"{TRANSFORMER_CONFIG_FILENAME} alignment_padding_mode must be a "
+            f"string, got {alignment_padding_mode!r}"
+        )
+    if type(multi_frame_output) is not bool:
+        raise InferenceProfileError(
+            f"{TRANSFORMER_CONFIG_FILENAME} multi_frame_output must be a "
+            f"boolean, got {multi_frame_output!r}"
+        )
+    pair = (alignment_padding_mode, multi_frame_output)
+    if pair not in CAPABILITY_PROFILES:
+        raise InferenceProfileError(
+            f"unsupported capability pair {pair!r} in "
+            f"{TRANSFORMER_CONFIG_FILENAME}; expected one of "
+            f"{sorted(CAPABILITY_PROFILES)}"
+        )
+    channels, sample_mode = _derive_vae_contract(model_directory)
+    profile = InferenceProfile(
+        schema_version=PROFILE_SCHEMA_VERSION,
+        inference_profile=CAPABILITY_PROFILES[pair],
+        alignment_padding_mode=alignment_padding_mode,
+        multi_frame_output=multi_frame_output,
+        vae_input_channels=channels,
+        vae_sample_mode=sample_mode,
+    )
+    profile.validate()
+    return profile
+
+
+def load_checkpoint_capabilities(
+    model_directory: Union[str, Path],
+) -> InferenceProfile:
+    """Derive the runtime capability from component configs.
+
+    New packages declare ``alignment_padding_mode`` and
+    ``multi_frame_output`` in ``transformer/config.json``; legacy packages
+    without them fall back to the strict root ``inference_profile.json``.
+    A missing ``transformer/config.json`` means both fields are absent, so
+    the legacy path applies. A partial pair is a hard error. When both the
+    component fields and the legacy file exist, the component metadata is
+    authoritative and the legacy file may only agree with it.
+    """
+
+    model_directory = Path(model_directory)
+    transformer_path = model_directory / TRANSFORMER_CONFIG_FILENAME
+    transformer_config = (
+        _read_component_config(model_directory, TRANSFORMER_CONFIG_FILENAME)
+        if transformer_path.is_file()
+        else {}
+    )
+    present = [field for field in CAPABILITY_FIELDS if field in transformer_config]
+    if len(present) == 1:
+        raise InferenceProfileError(
+            f"{TRANSFORMER_CONFIG_FILENAME} carries only {present[0]!r}; "
+            "alignment_padding_mode and multi_frame_output must be declared "
+            "together"
+        )
+    if not present:
+        return load_inference_profile(model_directory)
+
+    profile = _profile_from_components(
+        model_directory,
+        transformer_config["alignment_padding_mode"],
+        transformer_config["multi_frame_output"],
+    )
+    legacy_path = model_directory / PROFILE_FILENAME
+    if legacy_path.is_file():
+        legacy = load_inference_profile(model_directory)
+        if legacy != profile:
+            raise InferenceProfileError(
+                f"component configs disagree with legacy {PROFILE_FILENAME}: "
+                f"transformer/vae derive {profile!r} but the profile declares "
+                f"{legacy!r}"
+            )
+    return profile
+
+
+def resolve_model_directory(
+    model_name_or_path: Union[str, Path],
+    *,
+    revision: Optional[str] = None,
+    cache_dir: Optional[Union[str, Path]] = None,
+    local_files_only: bool = False,
+    token: Optional[Union[str, bool]] = None,
+) -> Path:
+    """Resolve a local directory or materialize a Hugging Face Hub snapshot."""
+
+    candidate = Path(model_name_or_path).expanduser()
+    if candidate.is_dir():
+        return candidate.resolve()
+    if candidate.exists():
+        raise ValueError(f"model path must be a directory: {candidate}")
+    if candidate.is_absolute():
+        raise FileNotFoundError(f"local model directory does not exist: {candidate}")
+
+    try:
+        from huggingface_hub import snapshot_download
+    except ImportError as error:
+        raise RuntimeError(
+            "huggingface_hub is required when --model is a Hub repository ID"
+        ) from error
+
+    snapshot_path = snapshot_download(
+        repo_id=str(model_name_or_path),
+        revision=revision,
+        cache_dir=str(cache_dir) if cache_dir is not None else None,
+        local_files_only=local_files_only,
+        token=token,
+    )
+    return Path(snapshot_path).resolve()

code/mllm_device_map.py

@@ -0,0 +1,146 @@
+"""Memory-aware device placement for the MLLM inference frontend."""
+
+from __future__ import annotations
+
+import json
+import math
+from dataclasses import dataclass
+from pathlib import Path
+
+
+DEFAULT_GPU0_RESERVED_LAYER_EQUIVALENTS = 4
+
+
+class MLLMDeviceMapError(ValueError):
+    """Raised when a requested MLLM layout is incomplete or unsafe."""
+
+
+@dataclass(frozen=True)
+class MLLMDevicePlan:
+    num_hidden_layers: int
+    n_gpu: int
+    gpu0_reserved_layer_equivalents: int
+    layer_counts: tuple[int, ...]
+    layer_devices: tuple[int, ...]
+    device_map: dict[str, int]
+
+
+def load_mllm_num_hidden_layers(model_directory: str | Path) -> int:
+    config_path = Path(model_directory) / "config.json"
+    try:
+        with config_path.open(encoding="utf-8") as handle:
+            config = json.load(handle)
+    except (OSError, json.JSONDecodeError) as exc:
+        raise MLLMDeviceMapError(f"cannot read MLLM config {config_path}: {exc}") from None
+
+    llm_config = config.get("llm_config")
+    nested = llm_config.get("num_hidden_layers") if isinstance(llm_config, dict) else None
+    root = config.get("num_hidden_layers")
+    if nested is not None and root is not None and nested != root:
+        raise MLLMDeviceMapError(
+            "ambiguous decoder depth: "
+            f"llm_config.num_hidden_layers={nested!r}, num_hidden_layers={root!r}"
+        )
+    value = nested if nested is not None else root
+    if isinstance(value, bool) or not isinstance(value, int) or value <= 0:
+        raise MLLMDeviceMapError(
+            f"missing or invalid MLLM num_hidden_layers in {config_path}: {value!r}"
+        )
+    return value
+
+
+def allocate_mllm_layer_counts(
+    num_hidden_layers: int,
+    n_gpu: int,
+    *,
+    gpu0_reserved_layer_equivalents: int = DEFAULT_GPU0_RESERVED_LAYER_EQUIVALENTS,
+) -> tuple[int, ...]:
+    if (
+        isinstance(n_gpu, bool)
+        or not isinstance(n_gpu, int)
+        or n_gpu <= 0
+    ):
+        raise MLLMDeviceMapError(f"n_gpu must be a positive integer, got {n_gpu!r}")
+    if isinstance(num_hidden_layers, bool) or not isinstance(num_hidden_layers, int) or num_hidden_layers <= 0:
+        raise MLLMDeviceMapError(f"num_hidden_layers must be positive, got {num_hidden_layers!r}")
+    if (
+        isinstance(gpu0_reserved_layer_equivalents, bool)
+        or not isinstance(gpu0_reserved_layer_equivalents, int)
+        or gpu0_reserved_layer_equivalents < 0
+    ):
+        raise MLLMDeviceMapError(
+            "gpu0_reserved_layer_equivalents must be a non-negative integer"
+        )
+
+    if n_gpu == 1:
+        # Single-GPU plan: every decoder layer and every fixed image module
+        # shares logical device 0, so nothing is kept sharded.
+        return (num_hidden_layers,)
+
+    effective_target = max(
+        gpu0_reserved_layer_equivalents,
+        math.ceil((num_hidden_layers + gpu0_reserved_layer_equivalents) / n_gpu),
+    )
+    gpu0_layers = min(
+        num_hidden_layers,
+        max(0, effective_target - gpu0_reserved_layer_equivalents),
+    )
+    base, remainder = divmod(num_hidden_layers - gpu0_layers, n_gpu - 1)
+    other_counts = [base] * (n_gpu - 1)
+    for index in range(len(other_counts) - remainder, len(other_counts)):
+        other_counts[index] += 1
+    counts = (gpu0_layers, *other_counts)
+    if sum(counts) != num_hidden_layers:
+        raise AssertionError(f"invalid internal MLLM allocation: {counts}")
+    return counts
+
+
+def build_mllm_device_plan(
+    num_hidden_layers: int,
+    n_gpu: int,
+    *,
+    gpu0_reserved_layer_equivalents: int = DEFAULT_GPU0_RESERVED_LAYER_EQUIVALENTS,
+) -> MLLMDevicePlan:
+    counts = allocate_mllm_layer_counts(
+        num_hidden_layers,
+        n_gpu,
+        gpu0_reserved_layer_equivalents=gpu0_reserved_layer_equivalents,
+    )
+    layer_devices = tuple(
+        device for device, count in enumerate(counts) for _ in range(count)
+    )
+    device_map = {
+        f"model.model.layers.{layer_index}": device
+        for layer_index, device in enumerate(layer_devices)
+    }
+    # Image conditioning and diffusion modules are attached after the base
+    # checkpoint load and are intentionally placed on logical CUDA device 0.
+    device_map.update(
+        {
+            "vision": 0,
+            "linear_proj": 0,
+            "model.model.word_embeddings.weight": 0,
+            "model.model.norm.weight": 0,
+            "model.lm_head.weight": 0,
+            "model.model.norm": 0,
+        }
+    )
+    return MLLMDevicePlan(
+        num_hidden_layers=num_hidden_layers,
+        n_gpu=n_gpu,
+        gpu0_reserved_layer_equivalents=gpu0_reserved_layer_equivalents,
+        layer_counts=counts,
+        layer_devices=layer_devices,
+        device_map=device_map,
+    )
+
+
+def validate_loaded_layer_devices(
+    actual_devices: list[int], plan: MLLMDevicePlan
+) -> None:
+    actual = tuple(actual_devices)
+    if actual != plan.layer_devices:
+        raise MLLMDeviceMapError(
+            "loaded MLLM layer placement disagrees with the requested plan: "
+            f"actual={actual}, expected={plan.layer_devices}"
+        )

code/modeling_bailing_moe_v2.py

@@ -0,0 +1,2031 @@
+# coding=utf-8
+# Copyright 2023 Antgroup and The HuggingFace Inc. team. All rights reserved.
+#
+# This code is based on EleutherAI's GPT-NeoX library and the GPT-NeoX
+# and OPT implementations in this library. It has been modified from its
+# original forms to accommodate minor architectural differences compared
+# to GPT-NeoX and OPT used by the Meta AI team that trained the model.
+#
+# 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.
+"""PyTorch BailingMoE model."""
+import math
+import warnings
+import os
+from typing import List, Optional, Tuple, Union
+
+import torch
+import torch.nn.functional as F
+import torch.utils.checkpoint
+from torch import nn
+from torch.nn import CrossEntropyLoss
+from transformers.activations import ACT2FN
+from transformers.cache_utils import Cache, DynamicCache
+from transformers.modeling_attn_mask_utils import (
+    AttentionMaskConverter,
+    _prepare_4d_attention_mask,
+    _prepare_4d_causal_attention_mask,
+    _prepare_4d_causal_attention_mask_for_sdpa,
+)
+from transformers.modeling_outputs import (
+    MoeModelOutputWithPast,
+    MoeCausalLMOutputWithPast,
+)
+from transformers.modeling_rope_utils import ROPE_INIT_FUNCTIONS, dynamic_rope_update
+from transformers.modeling_utils import PreTrainedModel
+from transformers.pytorch_utils import ALL_LAYERNORM_LAYERS, is_torch_greater_or_equal_than_1_13
+from transformers.utils import (
+    add_start_docstrings,
+    add_start_docstrings_to_model_forward,
+    is_flash_attn_2_available,
+    is_flash_attn_greater_or_equal_2_10,
+    logging,
+    replace_return_docstrings,
+)
+from transformers.utils.import_utils import is_torch_fx_available
+from configuration_bailing_moe_v2 import BailingMoeV2Config
+from transformers.generation.utils import GenerationMixin
+from modeling_utils import patch_continuous_features, build_modality_mask
+
+
+
+if is_flash_attn_2_available():
+    from flash_attn import flash_attn_func, flash_attn_varlen_func
+    from flash_attn.bert_padding import index_first_axis, pad_input, unpad_input  # noqa
+
+
+# This makes `_prepare_4d_causal_attention_mask` a leaf function in the FX graph.
+# It means that the function will not be traced through and simply appear as a node in the graph.
+if is_torch_fx_available():
+    if not is_torch_greater_or_equal_than_1_13:
+        import torch.fx
+
+    _prepare_4d_causal_attention_mask = torch.fx.wrap(_prepare_4d_causal_attention_mask)
+
+
+logger = logging.get_logger(__name__)
+
+_CONFIG_FOR_DOC = "BailingMoeV2Config"
+
+
+def _get_unpad_data(attention_mask):
+    seqlens_in_batch = attention_mask.sum(dim=-1, dtype=torch.int32)
+    indices = torch.nonzero(attention_mask.flatten(), as_tuple=False).flatten()
+    max_seqlen_in_batch = seqlens_in_batch.max().item()
+    cu_seqlens = F.pad(torch.cumsum(seqlens_in_batch, dim=0, dtype=torch.torch.int32), (1, 0))
+    return (
+        indices,
+        cu_seqlens,
+        max_seqlen_in_batch,
+    )
+
+
+def _expand_mask(mask: torch.Tensor, dtype: torch.dtype, tgt_len: Optional[int] = None):
+    warnings.warn(
+        "Calling `transformers.models.BailingMoeV2.modeling_BailingMoeV2._prepare_4d_attention_mask` is deprecated and will be removed in v4.37. Use `transformers.modeling_attn_mask_utils._prepare_4d_attention_mask"
+    )
+    return _prepare_4d_attention_mask(mask=mask, dtype=dtype, tgt_len=tgt_len)
+
+
+def _make_causal_mask(
+    input_ids_shape: torch.Size, dtype: torch.dtype, device: torch.device, past_key_values_length: int = 0
+):
+    warnings.warn(
+        "Calling `transformers.models.BailingMoeV2.modeling_BailingMoeV2._make_causal_mask` is deprecated and will be removed in v4.37. Use `transformers.models.BailingMoeV2.modeling_BailingMoeV2.AttentionMaskConverter._make_causal_mask"
+    )
+    return AttentionMaskConverter._make_causal_mask(
+        input_ids_shape=input_ids_shape, dtype=dtype, device=device, past_key_values_length=past_key_values_length
+    )
+
+
+class BailingMoeV2RMSNorm(nn.Module):
+    def __init__(self, hidden_size, eps=1e-6):
+        """
+        BailingMoeV2RMSNorm is equivalent to T5LayerNorm
+        """
+        super().__init__()
+        self.weight = nn.Parameter(torch.ones(hidden_size))
+        self.variance_epsilon = eps
+
+    def forward(self, hidden_states):
+        input_dtype = hidden_states.dtype
+        hidden_states = hidden_states.to(torch.float32)
+        variance = hidden_states.pow(2).mean(-1, keepdim=True)
+        hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon)
+        return self.weight * hidden_states.to(input_dtype)
+
+    def reset_parameters(self):
+        nn.init.ones_(self.weight)  # explicit reset for enabling fsdp
+
+
+ALL_LAYERNORM_LAYERS.append(BailingMoeV2RMSNorm)
+
+
+class BailingMoeV2RotaryEmbedding(nn.Module):
+    def __init__(self, config: BailingMoeV2Config, device=None):
+        super().__init__()
+        # BC: "rope_type" was originally "type"
+        if hasattr(config, "rope_scaling") and config.rope_scaling is not None:
+            self.rope_type = config.rope_scaling.get("rope_type", config.rope_scaling.get("type"))
+        else:
+            self.rope_type = "default"
+        self.max_seq_len_cached = config.max_position_embeddings
+        self.original_max_seq_len = config.max_position_embeddings
+
+        self.config = config
+        self.rope_init_fn = ROPE_INIT_FUNCTIONS[self.rope_type]
+
+        inv_freq, self.attention_scaling = self.rope_init_fn(self.config, device)
+        self.register_buffer("inv_freq", inv_freq, persistent=False)
+        self.original_inv_freq = self.inv_freq
+
+    @torch.no_grad()
+    @dynamic_rope_update  # power user: used with advanced RoPE types (e.g. dynamic rope)
+    def forward(self, x, position_ids):
+        position_ids = position_ids[:, 0, :] 
+        inv_freq_expanded = self.inv_freq[None, :, None].float().expand(position_ids.shape[0], -1, 1).to(x.device)
+        position_ids_expanded = position_ids[:, None, :].float()
+
+
+        device_type = x.device.type if isinstance(x.device.type, str) and x.device.type != "mps" else "cpu"
+        with torch.autocast(device_type=device_type, enabled=False):  # Force float32
+            freqs = (inv_freq_expanded.float() @ position_ids_expanded.float()).transpose(1, 2)
+            emb = torch.cat((freqs, freqs), dim=-1)
+            cos = emb.cos() * self.attention_scaling
+            sin = emb.sin() * self.attention_scaling
+
+        return cos.to(dtype=x.dtype), sin.to(dtype=x.dtype)
+
+    def reset_parameters(self) -> None:
+        new_inv_freq, self.attention_scaling = self.rope_init_fn(self.config, None)
+        self.inv_freq.copy_(new_inv_freq)
+    
+    
+class BailingMoeV2RotaryEmbedding3D(nn.Module):
+    def __init__(self, config: BailingMoeV2Config, device=None):
+        super().__init__()
+        # BC: "rope_type" was originally "type"
+        self.rope_init_type = "default"
+        self.rope_type = config.rope_scaling.get("rope_type", config.rope_scaling.get("type"))
+        self.max_seq_len_cached = config.max_position_embeddings
+        self.original_max_seq_len = config.max_position_embeddings
+
+        self.config = config
+        self.rope_init_fn = ROPE_INIT_FUNCTIONS[self.rope_init_type]
+
+        inv_freq, self.attention_scaling = self.rope_init_fn(self.config, device)
+        self.register_buffer("inv_freq", inv_freq, persistent=False)
+        self.original_inv_freq = self.inv_freq
+        self.config = config
+
+    @torch.no_grad()
+    @dynamic_rope_update  # power user: used with advanced RoPE types (e.g. dynamic rope)
+    def forward(self, x, position_ids):
+        if self.rope_type == "3D" or self.rope_type == "video_rope":
+            inv_freq_expanded = self.inv_freq[None, None, :, None].float().expand(3, position_ids.shape[1], -1, 1).to(x.device)
+            position_ids_expanded = position_ids[:, :, None, :].float()
+        else:
+            inv_freq_expanded = self.inv_freq[None, :, None].float().expand(position_ids.shape[0], -1, 1).to(x.device)
+            position_ids_expanded = position_ids[:, None, :].float()
+
+        device_type = x.device.type if isinstance(x.device.type, str) and x.device.type != "mps" else "cpu"
+        with torch.autocast(device_type=device_type, enabled=False):  # Force float32
+            if self.rope_type == "3D" or self.rope_type == "video_rope":
+                freqs = (inv_freq_expanded.float() @ position_ids_expanded.float()).transpose(2, 3)
+            else:
+                freqs = (inv_freq_expanded.float() @ position_ids_expanded.float()).transpose(1, 2)
+            emb = torch.cat((freqs, freqs), dim=-1)
+            cos = emb.cos() * self.attention_scaling
+            sin = emb.sin() * self.attention_scaling
+
+        return cos.to(dtype=x.dtype), sin.to(dtype=x.dtype)
+
+    def reset_parameters(self) -> None:
+        new_inv_freq, self.attention_scaling = self.rope_init_fn(self.config, None)
+        self.inv_freq.copy_(new_inv_freq)
+
+
+# Copied from transformers.models.llama.modeling_llama.rotate_half
+def rotate_half(x):
+    """Rotates half the hidden dims of the input."""
+    x1 = x[..., : x.shape[-1] // 2]
+    x2 = x[..., x.shape[-1] // 2 :]
+    return torch.cat((-x2, x1), dim=-1)
+
+
+def apply_3d_rotary_pos_emb(
+    q, 
+    k, 
+    cos, 
+    sin, 
+    mrope_section=[8, 12, 12], 
+    unsqueeze_dim=1, 
+    rope_type="m_rope", 
+    rotary_half=True
+):
+    """Applies Rotary Position Embedding with Multimodal Sections to the query and key tensors (https://qwenlm.github.io/blog/qwen2-vl/).
+    Explanation:
+        Multimodal 3D rotary position embedding is an extension to 1D rotary position embedding. The input embedding
+        sequence contains vision (images / videos) embedding and text embedding or just contains text embedding. For
+        vision embedding part, we apply rotary position embedding on temporal, height and width dimension separately.
+        Here we split the channel dimension to 3 chunks for the temporal, height and width rotary position embedding.
+        For text embedding part, we just apply 1D rotary position embedding. The three rotary position index (temporal,
+        height and width) of text embedding is always the same, so the text embedding rotary position embedding has no
+        difference with modern LLMs.
+    Args:
+        q (`torch.Tensor`): The query tensor.
+        k (`torch.Tensor`): The key tensor.
+        cos (`torch.Tensor`): The cosine part of the rotary embedding.
+        sin (`torch.Tensor`): The sine part of the rotary embedding.
+        mrope_section(`List(int)`):
+            Multimodal rope section is for channel dimension of temporal, height and width in rope calculation.
+        unsqueeze_dim (`int`, *optional*, defaults to 1):
+            The 'unsqueeze_dim' argument specifies the dimension along which to unsqueeze cos[position_ids] and
+            sin[position_ids] so that they can be properly broadcasted to the dimensions of q and k. For example, note
+            that cos[position_ids] and sin[position_ids] have the shape [batch_size, seq_len, head_dim]. Then, if q and
+            k have the shape [batch_size, heads, seq_len, head_dim], then setting unsqueeze_dim=1 makes
+            cos[position_ids] and sin[position_ids] broadcastable to the shapes of q and k. Similarly, if q and k have
+            the shape [batch_size, seq_len, heads, head_dim], then set unsqueeze_dim=2.
+        rope_type (`str`, *optional*, defaults to "m_rope"):
+        rotary_half (`bool`, *optional*, defaults to `False`): Keep half or full tensor for later concatenation
+    Returns:
+        `tuple(torch.Tensor)` comprising the query and key tensors rotated using the Rotary Position Embedding.
+    """
+    if rope_type == "3D":  # rename rope_type
+        rope_type = "m_rope"
+
+    if rope_type == "m_rope":
+        mrope_section = mrope_section * 2
+        cos = torch.cat([m[i % 3] for i, m in enumerate(cos.split(mrope_section, dim=-1))], dim=-1).unsqueeze(unsqueeze_dim).to(q.device)
+        sin = torch.cat([m[i % 3] for i, m in enumerate(sin.split(mrope_section, dim=-1))], dim=-1).unsqueeze(unsqueeze_dim).to(q.device)
+    elif rope_type == "video_rope":
+        mrope_section = list(mrope_section)
+        mrope_section = [mrope_section[0], mrope_section[1] + mrope_section[2]]
+        mrope_section = mrope_section * 2
+        # adjust t last -> (48, 16, 48, 16)
+        mrope_section = mrope_section[::-1]
+        index = 0
+        result_cos = []
+        result_sin = []
+        # get x1, y1, x2, y2, ..., t1, t2, ...
+        for i, section in enumerate(mrope_section):
+            if i % 2 == 0:
+                for j in range(section):
+                    row = 1 if j % 2 == 0 else 2
+                    result_cos.append(cos[row, ..., index: index + 1])
+                    result_sin.append(sin[row, ..., index: index + 1])
+                    index += 1
+            else:
+                result_cos.append(cos[0, ..., index:index + section])
+                result_sin.append(sin[0, ..., index:index + section])
+                index += section
+        cos, sin = torch.cat(result_cos, dim=-1).unsqueeze(dim=unsqueeze_dim).to(q.device), torch.cat(result_sin, dim=-1).unsqueeze(dim=unsqueeze_dim).to(q.device)
+    else:  # vanilla rope for llm
+        cos = cos.unsqueeze(unsqueeze_dim).to(q.device)
+        sin = sin.unsqueeze(unsqueeze_dim).to(q.device)
+
+    if rotary_half:
+        rotary_dim = cos.shape[-1]
+        q_rot, q_pass = q[..., :rotary_dim], q[..., rotary_dim:]
+        k_rot, k_pass = k[..., :rotary_dim], k[..., rotary_dim:]
+
+        # Apply rotary embeddings on the first half
+        q_embed = (q_rot * cos) + (rotate_half(q_rot) * sin)
+        k_embed = (k_rot * cos) + (rotate_half(k_rot) * sin)
+
+        # Concatenate back to full shape
+        q_embed = torch.cat([q_embed, q_pass], dim=-1)
+        k_embed = torch.cat([k_embed, k_pass], dim=-1)
+    else:
+        q_embed = (q * cos) + (rotate_half(q) * sin)
+        k_embed = (k * cos) + (rotate_half(k) * sin)
+    return q_embed, k_embed
+
+def get_t_scale_rope_index(
+    config,
+    input_ids: torch.LongTensor,
+    image_grid_thw: Optional[torch.LongTensor] = None,
+    video_grid_thw: Optional[torch.LongTensor] = None,
+    attention_mask: Optional[torch.Tensor] = None,
+    scale_factor: float = 1.0,
+    second_per_grid_ts: Optional[torch.Tensor] = None,
+) -> Tuple[torch.Tensor, torch.Tensor]:
+    spatial_merge_size = config.spatial_merge_size
+    image_token_id = config.image_patch_token
+    video_token_id = config.video_patch_token
+    image_start_token_id = config.image_start_token
+    video_start_token_id = config.video_start_token
+    use_abs_time_pos = second_per_grid_ts is not None
+
+    mrope_position_deltas = []
+    if image_grid_thw is not None or video_grid_thw is not None:
+        total_input_ids = input_ids
+        if attention_mask is None:
+            attention_mask = torch.ones_like(total_input_ids)
+        position_ids = torch.ones(
+            3,
+            input_ids.shape[0],
+            input_ids.shape[1],
+            dtype=input_ids.dtype,
+            device=input_ids.device,
+        )
+
+        image_index, video_index = 0, 0
+        attention_mask = attention_mask.to(total_input_ids.device)
+
+        for i, input_ids in enumerate(total_input_ids):
+            if attention_mask is not None:
+                input_ids = input_ids[attention_mask[i] == 1]
+            image_nums, video_nums = 0, 0
+            if image_grid_thw is not None:
+                vision_start_indices = torch.argwhere(input_ids == image_start_token_id).squeeze(1)
+                vision_tokens = input_ids[vision_start_indices + 1]
+                image_nums = (vision_tokens == image_token_id).sum()
+            if video_grid_thw is not None:
+                vision_start_indices = torch.argwhere(input_ids == video_start_token_id).squeeze(1)
+                vision_tokens = input_ids[vision_start_indices + 1]
+                video_nums = (vision_tokens == video_token_id).sum()
+
+            input_tokens = input_ids.tolist()
+            llm_pos_ids_list: list = []
+            st = 0
+            remain_images, remain_videos = image_nums, video_nums
+            for _ in range(image_nums + video_nums):
+                if image_token_id in input_tokens and remain_images > 0:
+                    ed_image = input_tokens.index(image_token_id, st)
+                else:
+                    ed_image = len(input_tokens) + 1
+                if video_token_id in input_tokens and remain_videos > 0:
+                    ed_video = input_tokens.index(video_token_id, st)
+                else:
+                    ed_video = len(input_tokens) + 1
+                if ed_image < ed_video:
+                    t, h, w = (
+                        image_grid_thw[image_index][0],
+                        image_grid_thw[image_index][1],
+                        image_grid_thw[image_index][2],
+                    )
+                    second_per_grid_t = 0
+                    image_index += 1
+                    remain_images -= 1
+                    ed = ed_image
+                else:
+                    t, h, w = (
+                        video_grid_thw[video_index][0],
+                        video_grid_thw[video_index][1],
+                        video_grid_thw[video_index][2],
+                    )
+                    if second_per_grid_ts is not None:
+                        second_per_grid_t = second_per_grid_ts[video_index]
+                    else:
+                        second_per_grid_t = 1.0
+                    video_index += 1
+                    remain_videos -= 1
+                    ed = ed_video
+                llm_grid_t, llm_grid_h, llm_grid_w = (
+                    t.item(),
+                    h.item() // spatial_merge_size,
+                    w.item() // spatial_merge_size,
+                )
+                text_len = ed - st
+
+                st_idx = llm_pos_ids_list[-1][0].max() + 1 if len(llm_pos_ids_list) > 0 else 0
+                llm_pos_ids_list.append(torch.arange(text_len).view(1, -1).expand(3, -1) + st_idx)
+
+                # body-diagonal symmetry
+                t_index = torch.arange(llm_grid_t).view(-1, 1).expand(
+                    -1, llm_grid_h * llm_grid_w).flatten()
+                h_index = torch.arange(llm_grid_h).view(1, -1, 1).expand(
+                    llm_grid_t, -1, llm_grid_w).flatten() - (llm_grid_h - 1) // 2
+                w_index = torch.arange(llm_grid_w).view(1, 1, -1).expand(
+                    llm_grid_t, llm_grid_h, -1).flatten() - (llm_grid_w - 1) // 2
+
+                # time dim adjust step size
+                if use_abs_time_pos:
+                    t_index = t_index * second_per_grid_t * scale_factor
+                else:
+                    t_index = t_index * scale_factor
+                t_index = t_index + text_len + st_idx
+
+                h_index = h_index + t_index
+                w_index = w_index + t_index
+                llm_pos_ids_list.append(
+                    torch.stack([t_index, h_index, w_index]))
+                st = ed + llm_grid_t * llm_grid_h * llm_grid_w
+
+            if st < len(input_tokens):
+                # next text token near last video token position = last t + 1
+                st_idx = llm_pos_ids_list[-1][0].max() + 1 if len(llm_pos_ids_list) > 0 else 0
+                text_len = len(input_tokens) - st
+                llm_pos_ids_list.append(torch.arange(text_len).view(1, -1).expand(3, -1) + st_idx)
+
+            llm_positions = torch.cat(llm_pos_ids_list, dim=1).reshape(3, -1)
+            llm_positions = llm_positions.to(dtype=position_ids.dtype, device=position_ids.device)
+            position_ids[..., i, attention_mask[i] == 1] = llm_positions
+            # generate first token = last t + 1
+            mrope_position_deltas.append(llm_positions[0].max() + 1 - len(total_input_ids[i]))
+        mrope_position_deltas = torch.tensor(mrope_position_deltas, device=input_ids.device).unsqueeze(1)
+    else:
+        if attention_mask is not None:
+            position_ids = attention_mask.long().cumsum(-1) - 1
+            position_ids.masked_fill_(attention_mask == 0, 1)
+            position_ids = position_ids.unsqueeze(0).expand(3, -1, -1).to(input_ids.device)
+            max_position_ids = position_ids.max(0, keepdim=False)[0].max(-1, keepdim=True)[0]
+            mrope_position_deltas = max_position_ids + 1 - attention_mask.shape[-1]
+        else:
+            position_ids = (
+                torch.arange(input_ids.shape[1], device=input_ids.device)
+                .view(1, 1, -1)
+                .expand(3, input_ids.shape[0], -1)
+            )
+            mrope_position_deltas = torch.zeros(
+                [input_ids.shape[0], 1],
+                device=input_ids.device,
+                dtype=input_ids.dtype,
+            )
+
+    return position_ids, mrope_position_deltas
+
+def get_rope_index(
+    config,
+    input_ids: Optional[torch.LongTensor] = None,
+    image_grid_thw: Optional[torch.LongTensor] = None,
+    video_grid_thw: Optional[torch.LongTensor] = None,
+    attention_mask: Optional[torch.Tensor] = None,
+    second_per_grid_ts: Optional[torch.Tensor] = None,
+) -> Tuple[torch.Tensor, torch.Tensor]:
+        """
+        Calculate the 3D rope index based on image and video's temporal, height and width in LLM.
+
+        Explanation:
+            Each embedding sequence contains vision embedding and text embedding or just contains text embedding.
+
+            For pure text embedding sequence, the rotary position embedding has no difference with modern LLMs.
+            Examples:
+                input_ids: [T T T T T], here T is for text.
+                temporal position_ids: [0, 1, 2, 3, 4]
+                height position_ids: [0, 1, 2, 3, 4]
+                width position_ids: [0, 1, 2, 3, 4]
+
+            For vision and text embedding sequence, we calculate 3D rotary position embedding for vision part
+            and 1D rotary position embeddin for text part.
+            Examples:
+                Temporal (Time): 3 patches, representing different segments of the video in time.
+                Height: 2 patches, dividing each frame vertically.
+                Width: 2 patches, dividing each frame horizontally.
+                We also have some important parameters:
+                fps (Frames Per Second): The video's frame rate, set to 1. This means one frame is processed each second.
+                tokens_per_second: This is a crucial parameter. It dictates how many "time-steps" or "temporal tokens" are conceptually packed into a one-second interval of the video. In this case, we have 25 tokens per second. So each second of the video will be represented with 25 separate time points. It essentially defines the temporal granularity.
+                temporal_patch_size: The number of frames that compose one temporal patch. Here, it's 2 frames.
+                interval: The step size for the temporal position IDs, calculated as tokens_per_second * temporal_patch_size / fps. In this case, 25 * 2 / 1 = 50. This means that each temporal patch will be have a difference of 50 in the temporal position IDs.
+                input_ids: [V V V V V V V V V V V V T T T T T], here V is for vision.
+                vision temporal position_ids: [0, 0, 0, 0, 50, 50, 50, 50, 100, 100, 100, 100]
+                vision height position_ids: [0, 0, 1, 1, 0, 0, 1, 1, 0, 0, 1, 1]
+                vision width position_ids: [0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1]
+                text temporal position_ids: [101, 102, 103, 104, 105]
+                text height position_ids: [101, 102, 103, 104, 105]
+                text width position_ids: [101, 102, 103, 104, 105]
+                Here we calculate the text start position_ids as the max vision position_ids plus 1.
+
+        Args:
+            input_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`):
+                Indices of input sequence tokens in the vocabulary. Padding will be ignored by default should you provide
+                it.
+            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.
+            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.
+            attention_mask (`torch.Tensor` of shape `(batch_size, sequence_length)`, *optional*):
+                Mask to avoid performing attention on padding token indices. Mask values selected in `[0, 1]`:
+
+                - 1 for tokens that are **not masked**,
+                - 0 for tokens that are **masked**.
+
+        Returns:
+            position_ids (`torch.LongTensor` of shape `(3, batch_size, sequence_length)`)
+            mrope_position_deltas (`torch.Tensor` of shape `(batch_size)`)
+        """
+        spatial_merge_size = config.spatial_merge_size
+        image_token_id = config.image_patch_token
+        video_token_id = config.video_patch_token
+        image_start_token_id = config.image_start_token
+        video_start_token_id = config.video_start_token
+
+        use_abs_time_pos = second_per_grid_ts is not None
+
+        mrope_position_deltas = []
+        if image_grid_thw is not None or video_grid_thw is not None:
+            total_input_ids = input_ids
+            if attention_mask is None:
+                attention_mask = torch.ones_like(total_input_ids)
+            position_ids = torch.ones(
+                3,
+                input_ids.shape[0],
+                input_ids.shape[1],
+                dtype=input_ids.dtype,
+                device=input_ids.device,
+            )
+            image_index, video_index = 0, 0
+            attention_mask = attention_mask.to(total_input_ids.device)
+            for i, input_ids in enumerate(total_input_ids):
+                input_ids = input_ids[attention_mask[i] == 1]
+                image_nums, video_nums = 0, 0
+                if image_grid_thw is not None:
+                    vision_start_indices = torch.argwhere(input_ids == image_start_token_id).squeeze(1)
+                    vision_tokens = input_ids[vision_start_indices + 1]
+                    image_nums = (vision_tokens == image_token_id).sum()
+                if video_grid_thw is not None:
+                    vision_start_indices = torch.argwhere(input_ids == video_start_token_id).squeeze(1)
+                    vision_tokens = input_ids[vision_start_indices + 1]
+                    video_nums = (vision_tokens == video_token_id).sum()
+
+                input_tokens = input_ids.tolist()
+                llm_pos_ids_list: list = []
+                st = 0
+                remain_images, remain_videos = image_nums, video_nums
+                for _ in range(image_nums + video_nums):
+                    if image_token_id in input_tokens and remain_images > 0:
+                        ed_image = input_tokens.index(image_token_id, st)
+                    else:
+                        ed_image = len(input_tokens) + 1
+                    if video_token_id in input_tokens and remain_videos > 0:
+                        ed_video = input_tokens.index(video_token_id, st)
+                    else:
+                        ed_video = len(input_tokens) + 1
+                    if ed_image < ed_video:
+                        t, h, w = (
+                            image_grid_thw[image_index][0],
+                            image_grid_thw[image_index][1],
+                            image_grid_thw[image_index][2],
+                        )
+                        second_per_grid_t = 0
+                        image_index += 1
+                        remain_images -= 1
+                        ed = ed_image
+
+                    else:
+                        t, h, w = (
+                            video_grid_thw[video_index][0],
+                            video_grid_thw[video_index][1],
+                            video_grid_thw[video_index][2],
+                        )
+                        if second_per_grid_ts is not None:
+                            second_per_grid_t = second_per_grid_ts[video_index]
+                        else:
+                            second_per_grid_t = 1.0
+                        video_index += 1
+                        remain_videos -= 1
+                        ed = ed_video
+                    llm_grid_t, llm_grid_h, llm_grid_w = (
+                        t.item(),
+                        h.item() // spatial_merge_size,
+                        w.item() // spatial_merge_size,
+                    )
+                    text_len = ed - st
+
+                    st_idx = llm_pos_ids_list[-1].max() + 1 if len(llm_pos_ids_list) > 0 else 0
+                    llm_pos_ids_list.append(torch.arange(text_len).view(1, -1).expand(3, -1) + st_idx)
+
+                    range_tensor = torch.arange(llm_grid_t).view(-1, 1)
+                    expanded_range = range_tensor.expand(-1, llm_grid_h * llm_grid_w)
+                    if use_abs_time_pos:
+                        time_tensor = expanded_range * second_per_grid_t * config.tokens_per_second
+                        time_tensor_long = time_tensor.long()
+                    else:
+                        time_tensor_long = expanded_range.long()
+                    t_index = time_tensor_long.flatten()
+
+                    h_index = torch.arange(llm_grid_h).view(1, -1, 1).expand(llm_grid_t, -1, llm_grid_w).flatten()
+                    w_index = torch.arange(llm_grid_w).view(1, 1, -1).expand(llm_grid_t, llm_grid_h, -1).flatten()
+                    llm_pos_ids_list.append(torch.stack([t_index, h_index, w_index]) + text_len + st_idx)
+                    st = ed + llm_grid_t * llm_grid_h * llm_grid_w
+
+                if st < len(input_tokens):
+                    st_idx = llm_pos_ids_list[-1].max() + 1 if len(llm_pos_ids_list) > 0 else 0
+                    text_len = len(input_tokens) - st
+                    llm_pos_ids_list.append(torch.arange(text_len).view(1, -1).expand(3, -1) + st_idx)
+
+                llm_positions = torch.cat(llm_pos_ids_list, dim=1).reshape(3, -1)
+                position_ids[..., i, attention_mask[i] == 1] = llm_positions.to(position_ids.device)
+                mrope_position_deltas.append(llm_positions.max() + 1 - len(total_input_ids[i]))
+            mrope_position_deltas = torch.tensor(mrope_position_deltas, device=input_ids.device).unsqueeze(1)
+        else:
+            if attention_mask is not None:
+                position_ids = attention_mask.long().cumsum(-1) - 1
+                position_ids.masked_fill_(attention_mask == 0, 1)
+                position_ids = position_ids.unsqueeze(0).expand(3, -1, -1).to(input_ids.device)
+                max_position_ids = position_ids.max(0, keepdim=False)[0].max(-1, keepdim=True)[0]
+                mrope_position_deltas = max_position_ids + 1 - attention_mask.shape[-1]
+            else:
+                position_ids = (
+                    torch.arange(input_ids.shape[1], device=input_ids.device)
+                    .view(1, 1, -1)
+                    .expand(3, input_ids.shape[0], -1)
+                )
+                mrope_position_deltas = torch.zeros(
+                    [input_ids.shape[0], 1],
+                    device=input_ids.device,
+                    dtype=input_ids.dtype,
+                )
+
+        return position_ids, mrope_position_deltas
+
+
+class BailingMoeV2MLP(nn.Module):
+    def __init__(self, config: BailingMoeV2Config, intermediate_size: int):
+        super().__init__()
+        self.config = config
+        self.hidden_size = config.hidden_size
+        self.intermediate_size = intermediate_size
+
+        self.gate_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False)
+        self.up_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False)
+        self.down_proj = nn.Linear(self.intermediate_size, self.hidden_size, bias=False)
+        self.act_fn = ACT2FN[config.hidden_act]
+
+    def forward(self, x, image_mask=None, audio_mask=None):
+        return self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x))
+
+
+class BailingMoeV2Gate(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        self.config = config
+        self.top_k = config.num_experts_per_tok
+        self.num_experts = config.num_experts
+
+        self.n_group = config.n_group
+        self.topk_group = config.topk_group
+
+        # topk selection algorithm
+        self.gating_dim = config.hidden_size
+        self.weight = nn.Parameter(torch.empty((self.num_experts, self.gating_dim)))
+        self.routed_scaling_factor = config.routed_scaling_factor
+        self.bias_update_coeff = 0.001
+
+        self.expert_bias = torch.nn.Parameter(torch.zeros(self.num_experts), requires_grad=False)
+        self.reset_parameters()
+
+    def reset_parameters(self) -> None:
+        import torch.nn.init as init
+
+        init.kaiming_uniform_(self.weight, a=math.sqrt(5))
+
+    def update_bias(self, distribution: torch.Tensor):
+        with torch.no_grad():
+            delta_bias = (distribution.mean() - distribution).sign()
+        self.expert_bias.data = self.expert_bias.data + self.bias_update_coeff * delta_bias
+
+    def group_limited_topk(
+        self,
+        scores: torch.Tensor,
+    ):
+        num_tokens, _ = scores.size()
+        # Organize the experts into groups
+        group_scores = scores.view(num_tokens, self.n_group, -1).topk(2, dim=-1)[0].sum(dim=-1)
+        group_idx = torch.topk(group_scores, k=self.topk_group, dim=-1, sorted=False)[1]
+        group_mask = torch.zeros_like(group_scores)
+        group_mask.scatter_(1, group_idx, 1)
+
+        # Mask the experts based on selection groups
+        score_mask = (
+            group_mask.unsqueeze(-1)
+            .expand(num_tokens, self.n_group, self.num_experts // self.n_group)
+            .reshape(num_tokens, -1)
+        )
+
+        masked_scores = scores.masked_fill(~score_mask.bool(), float('-inf'))
+        probs, top_indices = torch.topk(masked_scores, k=self.top_k, dim=-1, sorted=False)
+
+        return probs, top_indices
+
+    def forward(self, hidden_states):
+        # compute gating score
+        hidden_states = hidden_states.view(-1, hidden_states.shape[-1])
+        logits = F.linear(hidden_states.type(torch.float32), self.weight.type(torch.float32))
+
+        scores = torch.sigmoid(logits)
+
+        scores_for_routing = scores + self.expert_bias
+        _, topk_idx = self.group_limited_topk(scores_for_routing)
+
+        scores = torch.gather(scores, dim=1, index=topk_idx).type_as(logits)
+
+        topk_weight = scores / (scores.sum(dim=-1, keepdim=True) + 1e-20) if self.top_k > 1 else scores
+        topk_weight = topk_weight * self.routed_scaling_factor
+
+        return topk_idx, topk_weight, logits
+
+
+class BailingMoeV2SparseMoeBlock(nn.Module):
+    """
+    A mixed expert module containing shared experts.
+    """
+
+    def __init__(self, config: BailingMoeV2Config):
+        super().__init__()
+        self.config = config
+        self.num_experts_per_tok = config.num_experts_per_tok
+        self._setup_experts()
+        self.router_type = config.router_type
+        if self.router_type == "topN":
+            logger.info(f"use topN Router")
+            self.gate = BailingMoeV2Gate(config)
+        elif self.router_type == "MultiRouter":
+            logger.info(f"use MultiRouter")
+            self.gate = BailingMoeV2Gate(config)
+            self.image_gate = BailingMoeV2Gate(config)
+            self.audio_gate = BailingMoeV2Gate(config)
+        if config.num_shared_experts is not None:
+            self.shared_experts = BailingMoeV2MLP(
+                config=config, intermediate_size=config.moe_intermediate_size * config.num_shared_experts
+            )
+
+    def _setup_experts(self):
+        self.experts = nn.ModuleList(
+            [
+                BailingMoeV2MLP(config=self.config, intermediate_size=self.config.moe_intermediate_size)
+                for _ in range(self.config.num_experts)
+            ]
+        )
+
+    def forward(self, hidden_states, image_mask, audio_mask):
+        identity = hidden_states
+        bsz, seq_len, h = hidden_states.shape
+        if self.router_type == "MultiRouter":
+            if image_mask is not None:
+                if len(image_mask.shape) == 3:
+                    assert image_mask.shape[-1] == 1
+                elif len(image_mask.shape) == 2:
+                    assert image_mask.shape == hidden_states.shape[:2]
+                    image_mask = image_mask.unsqueeze(-1)
+                else:
+                    raise ValueError(
+                        "unexpected image mask shape: "
+                        f"{tuple(image_mask.shape)}"
+                    )
+            if audio_mask is not None:
+                if len(audio_mask.shape) == 3:
+                    assert audio_mask.shape[-1] == 1
+                elif len(audio_mask.shape) == 2:
+                    assert audio_mask.shape == hidden_states.shape[:2]
+                    audio_mask = audio_mask.unsqueeze(-1)
+                else:
+                    raise ValueError(
+                        "unexpected audio mask shape: "
+                        f"{tuple(audio_mask.shape)}"
+                    )
+            if image_mask is not None and audio_mask is not None:
+                assert torch.logical_and(image_mask, audio_mask).sum() == 0
+            
+            image_topk_idx, image_topk_weight, image_router_logits = self.image_gate(hidden_states)
+            audio_topk_idx, audio_topk_weight, audio_router_logits = self.audio_gate(hidden_states)
+            topk_idx, topk_weight, router_logits = self.gate(hidden_states)
+
+            if image_mask is not None:
+                image_mask = image_mask.view(-1, 1)
+                topk_idx = image_topk_idx * image_mask + topk_idx * torch.logical_not(image_mask)
+                topk_weight = image_topk_weight * image_mask + topk_weight * torch.logical_not(image_mask)
+                router_logits = image_router_logits * image_mask + router_logits * torch.logical_not(image_mask)
+            if audio_mask is not None:
+                audio_mask = audio_mask.view(-1, 1)
+                audio_mask = audio_mask.to(router_logits.device)
+                topk_idx = audio_topk_idx * audio_mask + topk_idx * torch.logical_not(audio_mask)
+                topk_weight = audio_topk_weight * audio_mask + topk_weight * torch.logical_not(audio_mask)
+                router_logits = audio_router_logits * audio_mask + router_logits * torch.logical_not(audio_mask)
+        else:
+            topk_idx, topk_weight, router_logits = self.gate(hidden_states)
+
+        hidden_states = hidden_states.view(-1, hidden_states.shape[-1])
+        flat_topk_idx = topk_idx.view(-1)
+        y = self.moe_infer(hidden_states, topk_idx, topk_weight).view(bsz, seq_len, h)
+        if self.config.num_shared_experts is not None:
+            y = y + self.shared_experts(identity)
+        return y, (router_logits.view(bsz, seq_len, -1), topk_idx.view(bsz, seq_len, -1))
+
+    @torch.no_grad()
+    def moe_infer(self, x, topk_ids, topk_weight):
+        cnts = topk_ids.new_zeros((topk_ids.shape[0], len(self.experts)))
+        cnts.scatter_(1, topk_ids, 1)
+        tokens_per_expert = cnts.sum(dim=0)
+        idxs = topk_ids.view(-1).argsort()
+        sorted_tokens = x[idxs // topk_ids.shape[1]]
+        sorted_tokens_shape = sorted_tokens.shape
+        tokens_per_expert = tokens_per_expert.cpu().numpy()
+        outputs = []
+        start_idx = 0
+        for i, num_tokens in enumerate(tokens_per_expert):
+            end_idx = start_idx + num_tokens
+            if num_tokens == 0:
+                continue
+            expert = self.experts[i]
+            tokens_for_this_expert = sorted_tokens[start_idx:end_idx]
+            expert_out = expert(tokens_for_this_expert)
+            outputs.append(expert_out.to(x.device))
+            start_idx = end_idx
+
+        outs = torch.cat(outputs, dim=0) if len(outputs) else sorted_tokens.new_empty(0)
+        new_x = torch.empty_like(outs)
+        new_x[idxs] = outs
+        final_out = (
+            new_x.view(*topk_ids.shape, -1)
+            .type(topk_weight.dtype)
+            .mul_(topk_weight.unsqueeze(dim=-1))
+            .sum(dim=1)
+            .type(new_x.dtype)
+        )
+        return final_out
+
+
+# Copied from transformers.models.llama.modeling_llama.repeat_kv
+def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor:
+    """
+    This is the equivalent of torch.repeat_interleave(x, dim=1, repeats=n_rep). The hidden states go from (batch,
+    num_key_value_heads, seqlen, head_dim) to (batch, num_attention_heads, seqlen, head_dim)
+    """
+    batch, num_key_value_heads, slen, head_dim = hidden_states.shape
+    if n_rep == 1:
+        return hidden_states
+    hidden_states = hidden_states[:, :, None, :, :].expand(batch, num_key_value_heads, n_rep, slen, head_dim)
+    return hidden_states.reshape(batch, num_key_value_heads * n_rep, slen, head_dim)
+
+
+# Copied from transformers.models.llama.modeling_llama.LlamaAttention with Llama->BailingMoeV2
+class BailingMoeV2Attention(nn.Module):
+    """Multi-headed attention from 'Attention Is All You Need' paper"""
+
+    def __init__(self, config: BailingMoeV2Config, layer_idx: Optional[int] = None):
+        super().__init__()
+        self.config = config
+        self.layer_idx = layer_idx
+        if layer_idx is None:
+            logger.warning_once(
+                f"Instantiating {self.__class__.__name__} without passing `layer_idx` is not recommended and will "
+                "to errors during the forward call, if caching is used. Please make sure to provide a `layer_idx` "
+                "when creating this class."
+            )
+
+        self.attention_dropout = config.attention_dropout
+        self.hidden_size = config.hidden_size
+        self.num_heads = config.num_attention_heads
+        self.head_dim = config.head_dim or self.hidden_size // self.num_heads
+        partial_rotary_factor = config.partial_rotary_factor if hasattr(config, "partial_rotary_factor") else 1.0
+        self.rope_dim = int(self.head_dim * partial_rotary_factor)
+        self.num_key_value_heads = config.num_key_value_heads
+        self.num_key_value_groups = self.num_heads // self.num_key_value_heads
+        self.max_position_embeddings = config.max_position_embeddings
+        self.rope_theta = config.rope_theta
+        self.is_causal = True
+
+        self.query_key_value = nn.Linear(
+            self.hidden_size,
+            (self.num_heads + 2 * self.num_key_value_heads) * self.head_dim,
+            bias=config.use_qkv_bias,
+        )
+
+        self.q_norm = BailingMoeV2RMSNorm(self.head_dim, eps=config.rms_norm_eps)
+        self.k_norm = BailingMoeV2RMSNorm(self.head_dim, eps=config.rms_norm_eps)
+        self.dense = nn.Linear(self.num_heads * self.head_dim, self.hidden_size, bias=config.use_bias)
+
+        if self.config.rope_scaling is not None:
+            self.rope_scaling = {"type": "mrope", "mrope_section": [8, 12, 12]}
+            self.mrope_section = self.rope_scaling["mrope_section"]
+
+    def _shape(self, tensor: torch.Tensor, seq_len: int, bsz: int):
+        return tensor.view(bsz, seq_len, self.num_heads, self.head_dim).transpose(1, 2).contiguous()
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        attention_mask: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        past_key_value: Optional[Cache] = None,
+        output_attentions: bool = False,
+        use_cache: bool = False,
+        position_embeddings: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,  # necessary, but kept here for BC
+        **kwargs,
+    ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]:
+        if "padding_mask" in kwargs:
+            warnings.warn(
+                "Passing `padding_mask` is deprecated and will be removed in v4.37. Please make sure use `attention_mask` instead.`"
+            )
+
+        bsz, q_len, _ = hidden_states.size()
+
+        qkv = self.query_key_value(hidden_states)
+        qkv = qkv.view(bsz, q_len, self.num_heads + 2 * self.num_key_value_heads, self.head_dim)
+
+        query_states, key_states, value_states = qkv.split(
+            [self.num_heads, self.num_key_value_heads, self.num_key_value_heads], dim=-2
+        )
+        query_states = query_states.transpose(1, 2)
+        key_states = key_states.transpose(1, 2)
+        value_states = value_states.transpose(1, 2)
+
+        if not self.training:
+            query_states = query_states.contiguous()
+            key_states = key_states.contiguous()
+            value_states = value_states.contiguous()
+            del qkv
+
+        query_states = self.q_norm(query_states)
+        key_states = self.k_norm(key_states)
+
+        kv_seq_len = key_states.shape[-2]
+        if past_key_value is not None:
+            if self.layer_idx is None:
+                raise ValueError(
+                    f"The cache structure has changed since version v4.36. If you are using {self.__class__.__name__} "
+                    "for auto-regressive decoding with k/v caching, please make sure to initialize the attention class "
+                    "with a layer index."
+                )
+            kv_seq_len += past_key_value.get_usable_length(kv_seq_len, self.layer_idx)
+        cos, sin = position_embeddings
+        query_states, key_states = apply_3d_rotary_pos_emb(
+            query_states, 
+            key_states, 
+            cos, 
+            sin, 
+            mrope_section=self.rope_scaling["mrope_section"], 
+            rope_type=self.config.rope_scaling["type"]
+        )
+        
+        if past_key_value is not None:
+            cache_kwargs = {"sin": sin, "cos": cos}  # Specific to RoPE models
+            key_states, value_states = past_key_value.update(key_states, value_states, self.layer_idx, cache_kwargs)
+
+        key_states = repeat_kv(key_states, self.num_key_value_groups)
+        value_states = repeat_kv(value_states, self.num_key_value_groups)
+
+        attn_weights = torch.matmul(query_states / math.sqrt(self.head_dim), key_states.transpose(2, 3))
+
+        if attn_weights.size() != (bsz, self.num_heads, q_len, kv_seq_len):
+            raise ValueError(
+                f"Attention weights should be of size {(bsz, self.num_heads, q_len, kv_seq_len)}, but is"
+                f" {attn_weights.size()}"
+            )
+
+        if attention_mask is not None:
+            if attention_mask.size() != (bsz, 1, q_len, kv_seq_len):
+                raise ValueError(
+                    f"Attention mask should be of size {(bsz, 1, q_len, kv_seq_len)}, but is {attention_mask.size()}"
+                )
+            attn_weights = attn_weights + attention_mask
+
+        # upcast attention to fp32
+        attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to(query_states.dtype)
+        attn_weights = nn.functional.dropout(attn_weights, p=self.attention_dropout, training=self.training)
+        attn_output = torch.matmul(attn_weights, value_states)
+
+        if attn_output.size() != (bsz, self.num_heads, q_len, self.head_dim):
+            raise ValueError(
+                f"`attn_output` should be of size {(bsz, self.num_heads, q_len, self.head_dim)}, but is"
+                f" {attn_output.size()}"
+            )
+
+        attn_output = attn_output.transpose(1, 2).contiguous()
+
+        attn_output = attn_output.reshape(bsz, q_len, -1)
+
+        attn_output = self.dense(attn_output)
+
+        if not output_attentions:
+            attn_weights = None
+        
+        if not self.training:
+            del query_states, key_states, value_states
+
+        return attn_output, attn_weights, past_key_value
+
+
+# Copied from transformers.models.llama.modeling_llama.LlamaFlashAttention2 with Llama->BailingMoeV2
+class BailingMoeV2FlashAttention2(BailingMoeV2Attention):
+    """
+    BailingMoeV2 flash attention module. This module inherits from `BailingMoeV2Attention` as the weights of the module stays
+    untouched. The only required change would be on the forward pass where it needs to correctly call the public API of
+    flash attention and deal with padding tokens in case the input contains any of them.
+    """
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+
+        # TODO: Should be removed once Flash Attention for RoCm is bumped to 2.1.
+        # flash_attn<2.1 generates top-left aligned causal mask, while what is needed here is bottom-right alignement, that was made default for flash_attn>=2.1. This attribute is used to handle this difference. Reference: https://github.com/Dao-AILab/flash-attention/releases/tag/v2.1.0.
+        # Beware that with flash_attn<2.1, using q_seqlen != k_seqlen (except for the case q_seqlen == 1) produces a wrong mask (top-left).
+        self._flash_attn_uses_top_left_mask = not is_flash_attn_greater_or_equal_2_10()
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        attention_mask: Optional[torch.LongTensor] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        past_key_value: Optional[Cache] = None,
+        output_attentions: bool = False,
+        use_cache: bool = False,
+        position_embeddings: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,  # necessary, but kept here for BC
+        **kwargs,
+    ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]:
+        # BailingMoeV2FlashAttention2 attention does not support output_attentions
+        if "padding_mask" in kwargs:
+            warnings.warn(
+                "Passing `padding_mask` is deprecated and will be removed in v4.37. Please make sure use `attention_mask` instead.`"
+            )
+
+            # overwrite attention_mask with padding_mask
+            attention_mask = kwargs.pop("padding_mask")
+
+        output_attentions = False
+
+        bsz, q_len, _ = hidden_states.size()
+
+        # Flash attention requires the input to have the shape
+        # batch_size x seq_length x head_dim x hidden_dim
+        # therefore we just need to keep the original shape
+
+        qkv = self.query_key_value(hidden_states)
+        qkv = qkv.view(bsz, q_len, self.num_heads + 2 * self.num_key_value_heads, self.head_dim)
+
+        query_states, key_states, value_states = qkv.split(
+            [self.num_heads, self.num_key_value_heads, self.num_key_value_heads], dim=-2
+        )
+        query_states = query_states.transpose(1, 2)
+        key_states = key_states.transpose(1, 2)
+        value_states = value_states.transpose(1, 2)
+
+        if not self.training:
+            query_states = query_states.contiguous()
+            key_states = key_states.contiguous()
+            value_states = value_states.contiguous()
+            del qkv
+
+        query_states = self.q_norm(query_states)
+        key_states = self.k_norm(key_states)
+
+        kv_seq_len = key_states.shape[-2]
+        if past_key_value is not None:
+            kv_seq_len += past_key_value.get_usable_length(kv_seq_len, self.layer_idx)
+        cos, sin = position_embeddings
+        if self.config.rope_scaling is not None:
+            rope_type = self.config.rope_scaling.get("rope_type", self.config.rope_scaling.get("type"))
+        else:
+            rope_type = "default"
+        query_states, key_states = apply_3d_rotary_pos_emb(
+            query_states,
+            key_states,
+            cos,
+            sin,
+            rope_type=rope_type
+        )
+
+        if past_key_value is not None:
+            cache_kwargs = {"sin": sin, "cos": cos}  # Specific to RoPE models
+            key_states, value_states = past_key_value.update(key_states, value_states, self.layer_idx, cache_kwargs)
+
+        # TODO: These transpose are quite inefficient but Flash Attention requires the layout [batch_size, sequence_length, num_heads, head_dim]. We would need to refactor the KV cache
+        # to be able to avoid many of these transpose/reshape/view.
+        query_states = query_states.transpose(1, 2)
+        key_states = key_states.transpose(1, 2)
+        value_states = value_states.transpose(1, 2)
+
+        dropout_rate = self.attention_dropout if self.training else 0.0
+
+        # In PEFT, usually we cast the layer norms in float32 for training stability reasons
+        # therefore the input hidden states gets silently cast in float32. Hence, we need
+        # cast them back in the correct dtype just to be sure everything works as expected.
+        # This might slow down training & inference so it is recommended to not cast the LayerNorms
+        # in fp32. (BailingMoeV2RMSNorm handles it correctly)
+
+        input_dtype = query_states.dtype
+        if input_dtype == torch.float32:
+            # Handle the case where the model is quantized
+            if hasattr(self.config, "_pre_quantization_dtype"):
+                target_dtype = self.config._pre_quantization_dtype
+            elif torch.is_autocast_enabled():
+                target_dtype = torch.get_autocast_gpu_dtype()
+            else:
+                target_dtype = self.query_key_value.weight.dtype
+
+            logger.warning_once(
+                f"The input hidden states seems to be silently casted in float32, this might be related to"
+                f" the fact you have upcasted embedding or layer norm layers in float32. We will cast back the input in"
+                f" {target_dtype}."
+            )
+
+            query_states = query_states.to(target_dtype)
+            key_states = key_states.to(target_dtype)
+            value_states = value_states.to(target_dtype)
+
+        attn_output = self._flash_attention_forward(
+            query_states, key_states, value_states, attention_mask, q_len, dropout=dropout_rate
+        )
+
+        attn_output = attn_output.reshape(bsz, q_len, -1).contiguous()
+        attn_output = self.dense(attn_output)
+
+        if not output_attentions:
+            attn_weights = None
+
+        if not self.training:
+            del query_states, key_states, value_states
+
+        return attn_output, attn_weights, past_key_value
+
+    def _flash_attention_forward(
+        self, query_states, key_states, value_states, attention_mask, query_length, dropout=0.0, softmax_scale=None
+    ):
+        """
+        Calls the forward method of Flash Attention - if the input hidden states contain at least one padding token
+        first unpad the input, then computes the attention scores and pad the final attention scores.
+
+        Args:
+            query_states (`torch.Tensor`):
+                Input query states to be passed to Flash Attention API
+            key_states (`torch.Tensor`):
+                Input key states to be passed to Flash Attention API
+            value_states (`torch.Tensor`):
+                Input value states to be passed to Flash Attention API
+            attention_mask (`torch.Tensor`):
+                The padding mask - corresponds to a tensor of size `(batch_size, seq_len)` where 0 stands for the
+                position of padding tokens and 1 for the position of non-padding tokens.
+            dropout (`int`, *optional*):
+                Attention dropout
+            softmax_scale (`float`, *optional*):
+                The scaling of QK^T before applying softmax. Default to 1 / sqrt(head_dim)
+            query_length (`int`):
+                The length of the query sequence in terms of tokens. This represents the number of tokens in the
+                `query_states` tensor along the sequence dimension. It is used to determine the effective sequence
+                length for attention computations.
+        """
+        if not self._flash_attn_uses_top_left_mask:
+            causal = self.is_causal
+        else:
+            # TODO: Remove the `query_length != 1` check once Flash Attention for RoCm is bumped to 2.1. For details, please see the comment in BailingMoeV2FlashAttention2 __init__.
+            causal = self.is_causal and query_length != 1
+
+        # Contains at least one padding token in the sequence
+        if attention_mask is not None:
+            batch_size = query_states.shape[0]
+            query_states, key_states, value_states, indices_q, cu_seq_lens, max_seq_lens = self._upad_input(
+                query_states, key_states, value_states, attention_mask, query_length
+            )
+
+            cu_seqlens_q, cu_seqlens_k = cu_seq_lens
+            max_seqlen_in_batch_q, max_seqlen_in_batch_k = max_seq_lens
+
+            attn_output_unpad = flash_attn_varlen_func(
+                query_states,
+                key_states,
+                value_states,
+                cu_seqlens_q=cu_seqlens_q,
+                cu_seqlens_k=cu_seqlens_k,
+                max_seqlen_q=max_seqlen_in_batch_q,
+                max_seqlen_k=max_seqlen_in_batch_k,
+                dropout_p=dropout,
+                softmax_scale=softmax_scale,
+                causal=causal,
+            )
+
+            attn_output = pad_input(attn_output_unpad, indices_q, batch_size, query_length)
+        else:
+            attn_output = flash_attn_func(
+                query_states, key_states, value_states, dropout, softmax_scale=softmax_scale, causal=causal
+            )
+
+        return attn_output
+
+    def _upad_input(self, query_layer, key_layer, value_layer, attention_mask, query_length):
+        indices_k, cu_seqlens_k, max_seqlen_in_batch_k = _get_unpad_data(attention_mask)
+        batch_size, kv_seq_len, num_key_value_heads, head_dim = key_layer.shape
+
+        key_layer = index_first_axis(
+            key_layer.reshape(batch_size * kv_seq_len, num_key_value_heads, head_dim), indices_k
+        )
+        value_layer = index_first_axis(
+            value_layer.reshape(batch_size * kv_seq_len, num_key_value_heads, head_dim), indices_k
+        )
+        if query_length == kv_seq_len:
+            query_layer = index_first_axis(
+                query_layer.reshape(batch_size * kv_seq_len, self.num_heads, head_dim), indices_k
+            )
+            cu_seqlens_q = cu_seqlens_k
+            max_seqlen_in_batch_q = max_seqlen_in_batch_k
+            indices_q = indices_k
+        elif query_length == 1:
+            max_seqlen_in_batch_q = 1
+            cu_seqlens_q = torch.arange(
+                batch_size + 1, dtype=torch.int32, device=query_layer.device
+            )  # There is a memcpy here, that is very bad.
+            indices_q = cu_seqlens_q[:-1]
+            query_layer = query_layer.squeeze(1)
+        else:
+            # The -q_len: slice assumes left padding.
+            attention_mask = attention_mask[:, -query_length:]
+            query_layer, indices_q, cu_seqlens_q, max_seqlen_in_batch_q = unpad_input(query_layer, attention_mask)
+
+        return (
+            query_layer,
+            key_layer,
+            value_layer,
+            indices_q,
+            (cu_seqlens_q, cu_seqlens_k),
+            (max_seqlen_in_batch_q, max_seqlen_in_batch_k),
+        )
+
+
+ATTENTION_CLASSES = {
+    "eager": BailingMoeV2Attention,
+    "flash_attention_2": BailingMoeV2FlashAttention2,
+}
+
+
+class BailingMoeV2DecoderLayer(nn.Module):
+    def __init__(self, config: BailingMoeV2Config, layer_idx: int):
+        super().__init__()
+        self.hidden_size = config.hidden_size
+
+        self.attention = ATTENTION_CLASSES[config._attn_implementation](config=config, layer_idx=layer_idx)
+
+        self.mlp = (
+            BailingMoeV2SparseMoeBlock(config)
+            if (config.num_experts is not None and layer_idx >= config.first_k_dense_replace)
+            else BailingMoeV2MLP(config=config, intermediate_size=config.intermediate_size)
+        )
+        self.input_layernorm = BailingMoeV2RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+        self.post_attention_layernorm = BailingMoeV2RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        attention_mask: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        image_mask: Optional[torch.Tensor] = None,
+        audio_mask: Optional[torch.Tensor] = None,
+        past_key_value: Optional[Tuple[torch.Tensor]] = None,
+        output_attentions: Optional[bool] = False,
+        output_router_logits: Optional[bool] = False,
+        use_cache: Optional[bool] = False,
+        position_embeddings: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,  # necessary, but kept here for BC
+        **kwargs,
+    ) -> Tuple[torch.FloatTensor, Optional[Tuple[torch.FloatTensor, torch.FloatTensor]]]:
+        """
+        Args:
+            hidden_states (`torch.FloatTensor`): input to the layer of shape `(batch, seq_len, embed_dim)`
+            attention_mask (`torch.FloatTensor`, *optional*):
+                attention mask of size `(batch_size, sequence_length)` if flash attention is used or `(batch_size, 1,
+                query_sequence_length, key_sequence_length)` if default attention is used.
+            position_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
+                Indices of positions of each input sequence tokens in the position embeddings. Selected in the range `[0,
+                config.n_positions - 1]`.
+            past_key_value (`Tuple(torch.FloatTensor)`, *optional*):
+                cached past key and value projection states
+            output_attentions (`bool`, *optional*):
+                Whether to return the attentions tensors of all attention layers. See `attentions` under
+                returned tensors for more detail.
+            output_router_logits (`bool`, *optional*):
+                Whether or not to return the logits of all the routers. They are useful for computing the router loss,
+                and should not be returned during inference.
+            use_cache (`bool`, *optional*):
+                If set to `True`, `past_key_values` key value states are returned and can be used to speed up decoding
+                (see `past_key_values`).
+        """
+        if "padding_mask" in kwargs:
+            warnings.warn(
+                "Passing `padding_mask` is deprecated and will be removed in v4.37. Please make sure use `attention_mask` instead.`"
+            )
+        residual = hidden_states
+
+        hidden_states = self.input_layernorm(hidden_states)
+
+        # Self Attention
+        hidden_states, self_attn_weights, present_key_value = self.attention(
+            hidden_states=hidden_states,
+            attention_mask=attention_mask,
+            position_ids=position_ids,
+            past_key_value=past_key_value,
+            output_attentions=output_attentions,
+            position_embeddings=position_embeddings,
+            use_cache=use_cache,
+        )
+        hidden_states = residual + hidden_states
+
+        # Fully Connected
+        residual = hidden_states
+        hidden_states = self.post_attention_layernorm(hidden_states)
+        hidden_states = self.mlp(hidden_states, image_mask, audio_mask)
+        #hidden_states = self.mlp(hidden_states)
+        if isinstance(hidden_states, tuple):
+            hidden_states, router_logits = hidden_states
+        else:
+            router_logits = None
+        hidden_states = residual + hidden_states.to(residual.device)
+
+        outputs = (hidden_states,)
+
+        if output_attentions:
+            outputs += (self_attn_weights,)
+
+        if use_cache:
+            outputs += (present_key_value,)
+
+        if output_router_logits:
+            outputs += (router_logits,)
+
+        return outputs
+
+
+BAILINGMOEV2_START_DOCSTRING = r"""
+    This model inherits from [`PreTrainedModel`]. Check the superclass documentation for the generic methods the
+    library implements for all its model (such as downloading or saving, resizing the input embeddings, pruning heads
+    etc.)
+
+    This model is also a PyTorch [torch.nn.Module](https://pytorch.org/docs/stable/nn.html#torch.nn.Module) subclass.
+    Use it as a regular PyTorch Module and refer to the PyTorch documentation for all matter related to general usage
+    and behavior.
+
+    Parameters:
+        config ([`BailingMoeV2Config`]):
+            Model configuration class with all the parameters of the model. Initializing with a config file does not
+            load the weights associated with the model, only the configuration. Check out the
+            [`~PreTrainedModel.from_pretrained`] method to load the model weights.
+"""
+
+
+@add_start_docstrings(
+    "The bare BailingMoeV2 Model outputting raw hidden-states without any specific head on top.",
+    BAILINGMOEV2_START_DOCSTRING,
+)
+class BailingMoeV2PreTrainedModel(PreTrainedModel):
+    config_class = BailingMoeV2Config
+    base_model_prefix = "model"
+    supports_gradient_checkpointing = True
+    _no_split_modules = ["BailingMoeV2DecoderLayer"]
+    _skip_keys_device_placement = "past_key_values"
+    _supports_flash_attn_2 = True
+    _supports_sdpa = True
+    _supports_cache_class = True
+
+    def _init_weights(self, module):
+        std = self.config.initializer_range
+        if isinstance(module, nn.Linear):
+            module.weight.data.normal_(mean=0.0, std=std)
+            if module.bias is not None:
+                module.bias.data.zero_()
+        elif isinstance(module, nn.Embedding):
+            module.weight.data.normal_(mean=0.0, std=std)
+            if module.padding_idx is not None:
+                module.weight.data[module.padding_idx].zero_()
+
+
+BAILINGMOEV2_INPUTS_DOCSTRING = r"""
+    Args:
+        input_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`):
+            Indices of input sequence tokens in the vocabulary. Padding will be ignored by default should you provide
+            it.
+
+            Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.encode`] and
+            [`PreTrainedTokenizer.__call__`] for details.
+
+            [What are input IDs?](../glossary#input-ids)
+        attention_mask (`torch.Tensor` of shape `(batch_size, sequence_length)`, *optional*):
+            Mask to avoid performing attention on padding token indices. Mask values selected in `[0, 1]`:
+
+            - 1 for tokens that are **not masked**,
+            - 0 for tokens that are **masked**.
+
+            [What are attention masks?](../glossary#attention-mask)
+
+            Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.encode`] and
+            [`PreTrainedTokenizer.__call__`] for details.
+
+            If `past_key_values` is used, optionally only the last `input_ids` have to be input (see
+            `past_key_values`).
+
+            If you want to change padding behavior, you should read [`modeling_opt._prepare_decoder_attention_mask`]
+            and modify to your needs. See diagram 1 in [the paper](https://arxiv.org/abs/1910.13461) for more
+            information on the default strategy.
+
+            - 1 indicates the head is **not masked**,
+            - 0 indicates the head is **masked**.
+        position_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
+            Indices of positions of each input sequence tokens in the position embeddings. Selected in the range `[0,
+            config.n_positions - 1]`.
+
+            [What are position IDs?](../glossary#position-ids)
+        past_key_values (`Cache` or `tuple(tuple(torch.FloatTensor))`, *optional*):
+            Pre-computed hidden-states (key and values in the self-attention blocks and in the cross-attention
+            blocks) that can be used to speed up sequential decoding. This typically consists in the `past_key_values`
+            returned by the model at a previous stage of decoding, when `use_cache=True` or `config.use_cache=True`.
+
+            Two formats are allowed:
+            - a [`~cache_utils.Cache`] instance;
+            - Tuple of `tuple(torch.FloatTensor)` of length `config.n_layers`, with each tuple having 2 tensors of
+            shape `(batch_size, num_heads, sequence_length, embed_size_per_head)`). This is also known as the legacy
+            cache format.
+
+            The model will output the same cache format that is fed as input. If no `past_key_values` are passed, the
+            legacy cache format will be returned.
+
+            If `past_key_values` are used, the user can optionally input only the last `input_ids` (those that don't
+            have their past key value states given to this model) of shape `(batch_size, 1)` instead of all `input_ids`
+            of shape `(batch_size, sequence_length)`.
+        inputs_embeds (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`, *optional*):
+            Optionally, instead of passing `input_ids` you can choose to directly pass an embedded representation. This
+            is useful if you want more control over how to convert `input_ids` indices into associated vectors than the
+            model's internal embedding lookup matrix.
+        use_cache (`bool`, *optional*):
+            If set to `True`, `past_key_values` key value states are returned and can be used to speed up decoding (see
+            `past_key_values`).
+        output_attentions (`bool`, *optional*):
+            Whether or not to return the attentions tensors of all attention layers. See `attentions` under returned
+            tensors for more detail.
+        output_hidden_states (`bool`, *optional*):
+            Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors for
+            more detail.
+        return_dict (`bool`, *optional*):
+            Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple.
+"""
+
+
+@add_start_docstrings(
+    "The bare BailingMoeV2 Model outputting raw hidden-states without any specific head on top.",
+    BAILINGMOEV2_START_DOCSTRING,
+)
+class BailingMoeV2Model(BailingMoeV2PreTrainedModel):
+    """
+    Transformer decoder consisting of *config.num_hidden_layers* layers. Each layer is a [`BailingMoeV2DecoderLayer`]
+
+    Args:
+        config: BailingMoeV2Config
+    """
+
+    def __init__(self, config: BailingMoeV2Config):
+        super().__init__(config)
+        self.padding_idx = config.pad_token_id
+        self.vocab_size = config.vocab_size
+
+        self.word_embeddings = nn.Embedding(config.vocab_size, config.hidden_size, self.padding_idx)
+        self.layers = nn.ModuleList(
+            [BailingMoeV2DecoderLayer(config, layer_idx) for layer_idx in range(config.num_hidden_layers)]
+        )
+        self._use_sdpa = config._attn_implementation == "sdpa"
+        self._use_flash_attention_2 = config._attn_implementation == "flash_attention_2"
+        self.norm = BailingMoeV2RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+        if self.config.rope_scaling is not None:
+            self.rotary_emb = BailingMoeV2RotaryEmbedding3D(config=config)
+        else:
+            self.rotary_emb = BailingMoeV2RotaryEmbedding(config=config)
+        self.gradient_checkpointing = False
+        config.spatial_merge_size = 2 
+        config.tokens_per_second = 2
+        self.rope_deltas = None
+        # Initialize weights and apply final processing
+        self.post_init()
+
+    def get_input_embeddings(self):
+        return self.word_embeddings
+
+    def set_input_embeddings(self, value):
+        self.word_embeddings = value
+ 
+    def prompt_wrap_vision(self, input_ids, inputs_embeds, vision_embeds, vision_token_id):
+        if vision_embeds is None or input_ids is None:
+            return inputs_embeds
+
+        if len(vision_embeds.shape) == 3:
+            vision_embeds = vision_embeds.reshape(-1, vision_embeds.shape[-1])
+
+        n_image_tokens = (input_ids == vision_token_id).sum().item()
+        n_image_features = vision_embeds.shape[0]
+
+        if n_image_tokens != n_image_features:
+            raise ValueError(
+                f"Image features and image tokens do not match: tokens: {n_image_tokens}, features {n_image_features}"
+            )
+        vision_mask = (input_ids == vision_token_id).unsqueeze(-1).expand_as(inputs_embeds).to(inputs_embeds.device)
+        
+        image_embeds = vision_embeds.to(inputs_embeds.device, inputs_embeds.dtype)
+        inputs_embeds = inputs_embeds.masked_scatter(vision_mask, image_embeds)
+
+        return inputs_embeds
+
+    def prompt_wrap_audio(self, input_ids, inputs_embeds, audio_embeds, audio_embeds_lengths, placeholder_audio_loc_lens):
+        inputs_embeds = patch_continuous_features(
+           input_embeddings=inputs_embeds, placeholder_loc_lens=placeholder_audio_loc_lens,
+           encoded_feats=audio_embeds, encoded_feat_lens=audio_embeds_lengths,
+        )
+        router_mask_audio = build_modality_mask(placeholder_audio_loc_lens, inputs_embeds.shape[:-1])
+        router_mask_audio = router_mask_audio.to(inputs_embeds.device)
+        return inputs_embeds, router_mask_audio
+     
+    def prompt_wrap_navit(self, input_ids, config, query_embeds_image=None, query_embeds_video=None, query_embeds_audio=None,
+        query_embeds_audio_lengths=None, placeholder_audio_loc_lens=None, target_embeds=None):
+        inputs_embeds = self.word_embeddings(input_ids)
+        vision_mask = None
+        audio_mask = None
+        if query_embeds_image is None and query_embeds_video is None and query_embeds_audio is None and target_embeds is None:
+            return inputs_embeds, vision_mask, audio_mask
+
+        if query_embeds_image is not None:
+            inputs_embeds = self.prompt_wrap_vision(input_ids, inputs_embeds, query_embeds_image, config.image_patch_token)
+        if query_embeds_video is not None:
+            inputs_embeds = self.prompt_wrap_vision(input_ids, inputs_embeds, query_embeds_video, config.video_patch_token)
+
+        image_mask = input_ids == config.image_patch_token
+        video_mask = input_ids == config.video_patch_token
+        vision_mask = (image_mask + video_mask) > 0
+        vision_mask = vision_mask.unsqueeze(-1).to(input_ids.device)
+
+        if query_embeds_audio is not None:
+            inputs_embeds, audio_mask = self.prompt_wrap_audio(
+                input_ids, inputs_embeds, query_embeds_audio, query_embeds_audio_lengths, placeholder_audio_loc_lens,
+            )
+
+        return inputs_embeds, vision_mask, audio_mask
+
+
+    @add_start_docstrings_to_model_forward(BAILINGMOEV2_INPUTS_DOCSTRING)
+    def forward(
+        self,
+        input_ids: torch.LongTensor = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        query_embeds_image: Optional[torch.Tensor] = None,
+        query_embeds_video: Optional[torch.Tensor] = None,
+        query_embeds_audio: Optional[torch.Tensor] = None,
+        query_embeds_audio_lengths: Optional[torch.Tensor] = None,
+        placeholder_audio_loc_lens: Optional[torch.Tensor] = None,
+        target_embeds: Optional[torch.Tensor] = None, 
+        past_key_values: Optional[List[torch.FloatTensor]] = None,
+        inputs_embeds: Optional[torch.FloatTensor] = None,
+        image_grid_thw: Optional[torch.Tensor] = None,
+        image_grid_thw_video: Optional[torch.Tensor] = None, 
+        use_cache: Optional[bool] = None,
+        output_attentions: Optional[bool] = None,
+        output_hidden_states: Optional[bool] = None,
+        output_router_logits: Optional[bool] = None,
+        return_dict: Optional[bool] = None,
+        second_per_grid_ts: Optional[torch.Tensor] = None,
+        image_mask=None,
+        audio_mask=None,
+        **kwargs,
+    ) -> Union[Tuple, MoeModelOutputWithPast]:
+
+        if inputs_embeds is not None:
+            words_embeddings = inputs_embeds
+            input_shape = inputs_embeds.size()[:2]
+        else:
+            if (query_embeds_image is None and query_embeds_video is None and query_embeds_audio is None and target_embeds is None) or input_ids.size(1) == 1:
+                words_embeddings = self.word_embeddings(input_ids.clip(0, self.word_embeddings.weight.shape[0] - 1))
+                input_shape = input_ids.size()
+                image_mask = None
+                audio_mask = None
+            else:
+                words_embeddings, image_mask, audio_mask = self.prompt_wrap_navit(
+                    input_ids.clip(0, self.word_embeddings.weight.shape[0] - 1), self.config, query_embeds_image, query_embeds_video, query_embeds_audio,
+                    query_embeds_audio_lengths, placeholder_audio_loc_lens, target_embeds,  # noqa
+                )
+
+                input_shape = words_embeddings.size()[:2]
+        embeddings = words_embeddings
+
+        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
+        )
+        output_router_logits = (
+            output_router_logits if output_router_logits is not None else self.config.output_router_logits
+        )
+        use_cache = use_cache if use_cache is not None else self.config.use_cache
+
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+
+        # retrieve input_ids and inputs_embeds
+        if input_ids is not None and inputs_embeds is not None:
+            assert input_ids.size(1) == inputs_embeds.size(1), "{} vs {}".format(
+                input_ids.size,
+                inputs_embeds.size,
+            )
+            batch_size, seq_length = inputs_embeds.shape[:2]
+            #raise ValueError("You cannot specify both input_ids and inputs_embeds at the same time")
+        elif input_ids is not None:
+            batch_size, seq_length = input_ids.shape[:2]
+        elif inputs_embeds is not None:
+            batch_size, seq_length = inputs_embeds.shape[:2]
+        else:
+            raise ValueError("You have to specify either input_ids or inputs_embeds")
+
+        if self.gradient_checkpointing and self.training:
+            if use_cache:
+                logger.warning_once(
+                    "`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`transformers."
+                )
+                use_cache = False
+
+        past_key_values_length = 0
+        if use_cache:
+            use_legacy_cache = not isinstance(past_key_values, Cache)
+            if use_legacy_cache:
+                past_key_values = DynamicCache.from_legacy_cache(past_key_values)
+            past_key_values_length = past_key_values.get_usable_length(seq_length)
+
+        if position_ids is None:
+            device = input_ids.device if input_ids is not None else inputs_embeds.device
+            position_ids = torch.arange(
+                past_key_values_length, seq_length + past_key_values_length, dtype=torch.long, device=device
+            )
+            position_ids = position_ids.unsqueeze(0)
+
+        if self.rotary_emb.rope_type == "video_rope" and input_ids.size(1) != 1:
+            position_ids, rope_deltas = get_t_scale_rope_index(
+                self.config,
+                input_ids,
+                image_grid_thw,
+                image_grid_thw_video,
+                attention_mask,
+                scale_factor=2.0,
+                second_per_grid_ts=second_per_grid_ts,
+            )
+            self.rope_deltas = rope_deltas
+        elif self.rotary_emb.rope_type == "3D" and input_ids.size(1) != 1:
+            position_ids, rope_deltas = get_rope_index(
+                self.config,
+                input_ids,
+                image_grid_thw,
+                image_grid_thw_video,
+                attention_mask,
+                second_per_grid_ts,
+            )
+            self.rope_deltas = rope_deltas
+        elif self.rotary_emb.rope_type == "3D" or self.rotary_emb.rope_type == "video_rope": # decode stage
+            batch_size, seq_length = input_ids.shape
+            if past_key_values and self.rope_deltas:
+                delta = past_key_values[0][1].shape[2] + self.rope_deltas
+            elif past_key_values:
+                delta = torch.tensor(past_key_values[0][1].shape[2])
+            else:
+                delta = torch.tensor(0)
+            delta = delta.to(input_ids.device)
+            position_ids = torch.arange(seq_length, device=input_ids.device)
+            position_ids = position_ids.view(1, -1).expand(batch_size, -1)
+            position_ids = position_ids.add(delta)
+            position_ids = position_ids.unsqueeze(0).expand(3, -1, -1)
+
+        inputs_embeds = embeddings
+        if self._use_flash_attention_2:
+            # 2d mask is passed through the layers
+            attention_mask = attention_mask if (attention_mask is not None and 0 in attention_mask) else None
+        elif self._use_sdpa and not output_attentions:
+            # output_attentions=True can not be supported when using SDPA, and we fall back on
+            # the manual implementation that requires a 4D causal mask in all cases.
+            attention_mask = _prepare_4d_causal_attention_mask_for_sdpa(
+                attention_mask,
+                (batch_size, seq_length),
+                inputs_embeds,
+                past_key_values_length,
+            )
+        else:
+            # 4d mask is passed through the layers
+            attention_mask = _prepare_4d_causal_attention_mask(
+                attention_mask, (batch_size, seq_length), inputs_embeds, past_key_values_length
+            )
+
+        # embed positions
+        hidden_states = inputs_embeds
+
+        # create position embeddings to be shared across the decoder layers
+        position_embeddings = self.rotary_emb(hidden_states, position_ids)
+
+        # decoder layers
+        all_hidden_states = () if output_hidden_states else None
+        all_self_attns = () if output_attentions else None
+        all_router_logits = () if output_router_logits else None
+        next_decoder_cache = None
+
+        for decoder_layer in self.layers:
+            if output_hidden_states:
+                all_hidden_states += (hidden_states,)
+
+            if self.gradient_checkpointing and self.training:
+                layer_outputs = self._gradient_checkpointing_func(
+                    decoder_layer.__call__,
+                    hidden_states,
+                    attention_mask,
+                    position_ids,
+                    image_mask,
+                    audio_mask,
+                    past_key_values,
+                    output_attentions,
+                    output_router_logits,
+                    use_cache,
+                    position_embeddings,
+                )
+            else:
+                layer_outputs = decoder_layer(
+                    hidden_states,
+                    attention_mask=attention_mask,
+                    position_ids=position_ids,
+                    image_mask=image_mask,
+                    audio_mask=audio_mask,
+                    past_key_value=past_key_values,
+                    output_attentions=output_attentions,
+                    output_router_logits=output_router_logits,
+                    use_cache=use_cache,
+                    position_embeddings=position_embeddings,
+                )
+            hidden_states = layer_outputs[0]
+
+            if use_cache:
+                next_decoder_cache = layer_outputs[2 if output_attentions else 1]
+
+            if output_attentions:
+                all_self_attns += (layer_outputs[1],)
+
+            if output_router_logits and layer_outputs[-1] is not None:
+                all_router_logits += (layer_outputs[-1],)
+
+        hidden_states = self.norm(hidden_states)
+
+        # add hidden states from the last decoder layer
+        if output_hidden_states:
+            all_hidden_states += (hidden_states,)
+
+        next_cache = None
+        if use_cache:
+            next_cache = next_decoder_cache.to_legacy_cache() if use_legacy_cache else next_decoder_cache
+        if not return_dict:
+            return tuple(
+                v
+                for v in [hidden_states, next_cache, all_hidden_states, all_self_attns, all_router_logits]
+                if v is not None
+            )
+        return MoeModelOutputWithPast(
+            last_hidden_state=hidden_states,
+            past_key_values=next_cache,
+            hidden_states=all_hidden_states,
+            attentions=all_self_attns,
+            router_logits=all_router_logits,
+        )
+
+
+class BailingMoeV2ForCausalLM(BailingMoeV2PreTrainedModel, GenerationMixin):
+    _tied_weights_keys = ["lm_head.weight"]
+
+    def __init__(self, config: BailingMoeV2Config):
+        super().__init__(config)
+        self.model = BailingMoeV2Model(config)
+        self.vocab_size = config.vocab_size
+        self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
+
+        # Initialize weights and apply final processing
+        self.post_init()
+
+    def get_input_embeddings(self):
+        return self.model.word_embeddings
+
+    def set_input_embeddings(self, value):
+        self.model.word_embeddings = value
+
+    def get_output_embeddings(self):
+        return self.lm_head
+
+    def set_output_embeddings(self, new_embeddings):
+        self.lm_head = new_embeddings
+
+    def set_decoder(self, decoder):
+        self.model = decoder
+
+    def get_decoder(self):
+        return self.model
+
+    @add_start_docstrings_to_model_forward(BAILINGMOEV2_INPUTS_DOCSTRING)
+    @replace_return_docstrings(output_type=MoeCausalLMOutputWithPast, config_class=_CONFIG_FOR_DOC)
+    def forward(
+        self,
+        input_ids: torch.LongTensor = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        query_embeds_image: Optional[torch.Tensor] = None,
+        query_embeds_video: Optional[torch.Tensor] = None,
+        query_embeds_audio: Optional[torch.Tensor] = None,
+        query_embeds_audio_lengths: Optional[torch.Tensor] = None,
+        placeholder_audio_loc_lens: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        past_key_values: Optional[List[torch.FloatTensor]] = None,
+        inputs_embeds: Optional[torch.FloatTensor] = None,
+        image_grid_thw: Optional[torch.Tensor] = None,
+        image_grid_thw_video: Optional[torch.Tensor] = None,
+        labels: Optional[torch.LongTensor] = None,
+        use_cache: Optional[bool] = None,
+        output_attentions: Optional[bool] = None,
+        output_hidden_states: Optional[bool] = None,
+        output_router_logits: Optional[bool] = None,
+        return_dict: Optional[bool] = None,
+        second_per_grid_ts: Optional[torch.Tensor] = None,
+        num_logits_to_keep: Optional[int] = 0,
+        image_mask=None,
+        audio_mask=None,
+        **kwargs,
+    ) -> Union[Tuple, MoeCausalLMOutputWithPast]:
+        r"""
+        Args:
+            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]`.
+
+        Returns:
+
+        Example:
+
+        ```python
+        >>> from transformers import AutoTokenizer
+
+        >>> model = BailingMoeV2ForCausalLM.from_pretrained(PATH_TO_CONVERTED_WEIGHTS)
+        >>> tokenizer = AutoTokenizer.from_pretrained(PATH_TO_CONVERTED_TOKENIZER)
+
+        >>> prompt = "Hey, are you conscious? Can you talk to me?"
+        >>> inputs = tokenizer(prompt, 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]
+        "Hey, are you conscious? Can you talk to me?\nI'm not conscious, but I can talk to you."
+        ```"""
+        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
+        )
+        output_router_logits = (
+            output_router_logits if output_router_logits is not None else self.config.output_router_logits
+        )
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+        # decoder outputs consists of (dec_features, layer_state, dec_hidden, dec_attn)
+        outputs = self.model(
+            input_ids=input_ids,
+            attention_mask=attention_mask,
+            position_ids=position_ids,
+            query_embeds_image=query_embeds_image,
+            query_embeds_video=query_embeds_video,
+            query_embeds_audio=query_embeds_audio,
+            query_embeds_audio_lengths=query_embeds_audio_lengths,
+            placeholder_audio_loc_lens=placeholder_audio_loc_lens,
+            past_key_values=past_key_values,
+            inputs_embeds=inputs_embeds,
+            image_grid_thw=image_grid_thw,
+            image_grid_thw_video=image_grid_thw_video,
+            use_cache=use_cache,
+            output_attentions=output_attentions,
+            output_hidden_states=output_hidden_states,
+            output_router_logits=output_router_logits,
+            return_dict=return_dict,
+            second_per_grid_ts=second_per_grid_ts,
+            image_mask=image_mask,
+            audio_mask=audio_mask,
+            **kwargs,
+        )
+
+        hidden_states = outputs[0]
+
+        logits = self.lm_head(hidden_states[:, -num_logits_to_keep:, :])
+        # logits = logits.float()
+
+        loss = None
+        aux_loss = None
+
+        if labels is not None:
+            # Shift so that tokens < n predict n
+            shift_logits = logits[..., :-1, :].contiguous()
+            shift_labels = labels[..., 1:].contiguous()
+            # Flatten the tokens
+            loss_fct = CrossEntropyLoss(inplace_backward=True)
+            shift_logits = shift_logits.view(-1, self.config.vocab_size)
+            shift_labels = shift_labels.view(-1)
+            # Enable model parallelism
+            shift_labels = shift_labels.to(shift_logits.device)
+            loss = loss_fct(shift_logits, shift_labels)
+
+        if not return_dict:
+            output = (logits,) + outputs[1:]
+            if output_router_logits:
+                output = (aux_loss,) + output
+            return (loss,) + output if loss is not None else output
+
+        return MoeCausalLMOutputWithPast(
+            loss=loss,
+            aux_loss=aux_loss,
+            logits=logits,
+            past_key_values=outputs.past_key_values,
+            hidden_states=outputs.hidden_states,
+            attentions=outputs.attentions,
+            router_logits=outputs.router_logits,
+        )
+
+    def prepare_inputs_for_generation(
+        self,
+        input_ids,
+        query_embeds_image: torch.Tensor = None,
+        query_embeds_video: torch.Tensor = None,
+        query_embeds_audio: torch.Tensor = None,
+        query_embeds_audio_lengths: torch.Tensor = None,
+        placeholder_audio_loc_lens: torch.Tensor = None,
+        image_mask=None,
+        audio_mask=None,
+        past_key_values=None,
+        attention_mask=None,
+        inputs_embeds=None,
+        cache_position=None,
+        position_ids=None,
+        use_cache=True,
+        image_grid_thw=None,
+        image_grid_thw_video=None,
+        second_per_grid_ts=None,
+        is_audio_generation_mode=False,
+        num_logits_to_keep=0,
+        **kwargs,
+    ):
+        if past_key_values is not None:
+            if isinstance(past_key_values, Cache):
+                cache_length = past_key_values.get_seq_length()
+                past_length = past_key_values.seen_tokens
+                max_cache_length = (
+                    past_key_values.get_max_length()
+                    if hasattr(past_key_values, "get_max_length")
+                    else past_key_values.get_max_cache_shape()
+                )
+            else:
+                cache_length = past_length = past_key_values[0][0].shape[2]
+                max_cache_length = None
+
+            # Keep only the unprocessed tokens:
+            # 1 - If the length of the attention_mask exceeds the length of input_ids, then we are in a setting where
+            # some of the inputs are exclusivelly passed as part of the cache (e.g. when passing input_embeds as input)
+            if attention_mask is not None and attention_mask.shape[1] > input_ids.shape[1]:
+                input_ids = input_ids[:, -(attention_mask.shape[1] - past_length) :]
+            # 2 - If the past_length is smaller than input_ids', then input_ids holds all input tokens. We can discard
+            # input_ids based on the past_length.
+            elif past_length < input_ids.shape[1]:
+                input_ids = input_ids[:, past_length:]
+            # 3 - Otherwise (past_length >= input_ids.shape[1]), let's assume input_ids only has unprocessed tokens.
+
+            # If we are about to go beyond the maximum cache length, we need to crop the input attention mask.
+            if (
+                max_cache_length is not None
+                and attention_mask is not None
+                and cache_length + input_ids.shape[1] > max_cache_length
+            ):
+                attention_mask = attention_mask[:, -max_cache_length:]
+
+        position_ids = kwargs.get("position_ids", None)
+        if attention_mask is not None and position_ids is None:
+            # create position_ids on the fly for batch generation
+            position_ids = attention_mask.long().cumsum(-1) - 1
+            position_ids.masked_fill_(attention_mask == 0, 1)
+            if past_key_values:
+                position_ids = position_ids[:, -input_ids.shape[1] :]
+            position_ids = position_ids.unsqueeze(0)
+
+        # if `inputs_embeds` are passed, we only want to use them in the 1st generation step
+        if inputs_embeds is not None and past_key_values is None:
+            model_inputs = {"inputs_embeds": inputs_embeds, "input_ids": None}
+        else:
+            model_inputs = {"input_ids": input_ids, "inputs_embeds": None}
+
+        model_inputs.update(
+            {
+                "query_embeds_image": query_embeds_image,
+                "query_embeds_video": query_embeds_video,
+                "query_embeds_audio": query_embeds_audio,
+                "position_ids": position_ids,
+                "past_key_values": past_key_values,
+                "use_cache": use_cache,
+                "attention_mask": attention_mask,
+                "image_grid_thw": image_grid_thw,
+                "image_grid_thw_video": image_grid_thw_video,
+                "image_mask": image_mask,
+                "audio_mask": audio_mask,
+                "query_embeds_audio_lengths": query_embeds_audio_lengths,
+                "placeholder_audio_loc_lens": placeholder_audio_loc_lens,
+                "num_logits_to_keep": num_logits_to_keep,
+            }
+        )
+        return model_inputs
+
+    @staticmethod
+    def _reorder_cache(past_key_values, beam_idx):
+        reordered_past = ()
+        for layer_past in past_key_values:
+            reordered_past += (
+                tuple(past_state.index_select(0, beam_idx.to(past_state.device)) for past_state in layer_past),
+            )
+        return reordered_past
+

code/modeling_bailingmm2.py

@@ -0,0 +1,814 @@
+#!/usr/bin/env python3
+# coding=utf-8
+# Copyright (c) Ant Group. All rights reserved.
+
+from typing import List, Optional, Union
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from PIL import Image
+
+from diffusers.models.normalization import RMSNorm
+from transformers import PreTrainedModel
+from transformers.utils import logging
+from configuration_bailingmm2 import BailingMM2Config
+from modeling_bailing_moe_v2 import BailingMoeV2ForCausalLM
+from bailingmm_utils import process_ratio, find_first_index_of_consecutive_ones, merge_consecutive_ones
+from inference_profile import load_checkpoint_capabilities, resolve_model_directory
+import os
+from copy import deepcopy
+
+# vision encoder
+from qwen2_5_vit import Qwen2_5_VisionTransformer
+
+logger = logging.get_logger(__name__)
+
+_CONFIG_FOR_DOC = "BailingMM2Config"
+
+
+class BailingMM2NativeForConditionalGeneration(PreTrainedModel):
+    config_class = BailingMM2Config
+    base_model_prefix = "model"
+    supports_gradient_checkpointing = True
+    _skip_keys_device_placement = "past_key_values"
+    _supports_flash_attn_2 = True
+
+    def __init__(
+        self,
+        config: BailingMM2Config,
+        empty_load=False, 
+    ):
+        super().__init__(config)
+        self.config: BailingMM2Config = config
+        self.vision = None
+
+        self.llm_dytpe = torch.bfloat16
+
+        if empty_load:
+            self.model = None
+            return
+
+        if self.config.vision_config:
+            self.vision = Qwen2_5_VisionTransformer(self.config.vision_config)
+
+        self.model = BailingMoeV2ForCausalLM(self.config.llm_config)
+
+        mlp_modules_img = [nn.Linear(self.vision.image_emb_dim, self.model.config.hidden_size)]
+        for _ in range(1, self.config.mlp_depth):
+            mlp_modules_img.append(nn.GELU())
+            mlp_modules_img.append(nn.Linear(self.model.config.hidden_size, self.model.config.hidden_size))
+        self.linear_proj = nn.Sequential(*mlp_modules_img)
+
+        self.post_init()
+
+
+    def extract_image_feature(self, pixel_values, grid_thw):
+        with torch.cuda.amp.autocast(dtype=torch.bfloat16):
+            image_embeds = self.vision(pixel_values, grid_thw=grid_thw)
+        image_embeds = self.linear_proj(image_embeds)
+        image_embeds = F.normalize(image_embeds, dim=-1)
+        return image_embeds
+
+
+    @torch.no_grad()
+    def generate(
+        self,
+        input_ids: torch.LongTensor = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        use_cache: Optional[bool] = None,
+        pixel_values: Optional[torch.FloatTensor] = None,
+        pixel_values_videos: Optional[torch.FloatTensor] = None,
+        audio_feats: Optional[torch.FloatTensor] = None,
+        audio_feats_lengths: Optional[torch.LongTensor] = None,
+        audio_placeholder_loc_lens: Optional[torch.LongTensor] = None,
+        image_grid_thw: Optional[torch.LongTensor] = None,
+        video_grid_thw: Optional[torch.LongTensor] = None,
+        past_key_values: Optional[List[torch.Tensor]] = None,
+        num_logits_to_keep: Optional[int] = 0,
+        image_gen: Optional[bool] = False,
+        image_gen_pixel_values_reference: Optional[torch.FloatTensor] = None,
+        image_gen_negative_input_ids: Optional[torch.LongTensor] = None,
+        image_gen_negative_attention_mask: Optional[torch.Tensor] = None,
+        image_gen_steps: Optional[int] = None,
+        image_gen_seed: Optional[int] = None,
+        image_gen_cfg: Optional[float] = None,
+        image_gen_image_cfg: Optional[float] = 1.0,
+        image_gen_cfg_mode: Optional[int] = 1,
+        image_gen_height: Optional[int] = None,
+        image_gen_width: Optional[int] = None,
+        image_gen_llm_hidden_states:  Optional[torch.LongTensor] = None,
+        image_gen_negative_llm_hidden_states:  Optional[torch.LongTensor] = None,
+        image_gen_text: Optional[list] = None,
+        image_gen_highres = 512,
+        image_gen_only_extract_hidden_states = False,
+        image_gen_condition_embeds=None,
+        image_gen_negative_condition_embeds=None,
+        image_gen_condition_embeds_2=None,
+        image_gen_negative_condition_embeds_2=None,
+        image_gen_return_batch=False,
+        image_gen_task=None,
+        num_frames_per_prompt=1,
+        **generate_kwargs,
+    ):
+        if audio_feats is not None or audio_feats_lengths is not None:
+            raise ValueError("audio input is not supported by Ming Image inference")
+        if image_gen_image_cfg not in (None, 1.0):
+            raise ValueError(
+                "image_gen_image_cfg is not supported by this inference path; "
+                "guidance is controlled by image_gen_cfg only"
+            )
+        image_embeds, video_embeds, audio_embeds, audio_embeds_lengths = None, None, None, None
+
+        if image_gen:
+            if not hasattr(self, "inference_profile"):
+                raise RuntimeError(
+                    "image modules were loaded without a checkpoint inference profile"
+                )
+            self.inference_profile.validate_task(
+                image_gen_task,
+                has_reference_image=image_gen_pixel_values_reference is not None,
+                num_layers=num_frames_per_prompt,
+            )
+            sampling = self.inference_profile.resolve_sampling_parameters(
+                steps=image_gen_steps,
+                cfg=image_gen_cfg,
+            )
+            image_gen_steps = sampling.steps
+            image_gen_cfg = sampling.cfg
+            if image_gen_pixel_values_reference is not None:
+                input_channels = image_gen_pixel_values_reference.shape[1]
+                expected_channels = self.inference_profile.vae_input_channels
+                if input_channels % expected_channels != 0:
+                    raise ValueError(
+                        "reference image channels do not match the checkpoint "
+                        f"VAE contract: input={input_channels}, expected a "
+                        f"multiple of {expected_channels}"
+                    )
+            condition_embeds, negative_condition_embeds = None, None
+            condition_embeds_2, negative_condition_embeds_2 = None, None
+            if (image_gen_condition_embeds is not None) or (image_gen_condition_embeds_2 is not None):
+                if image_gen_condition_embeds is not None:
+                    condition_embeds = image_gen_condition_embeds
+                    negative_condition_embeds = condition_embeds * 0.0 if image_gen_negative_condition_embeds is None else image_gen_negative_condition_embeds
+                
+                if image_gen_condition_embeds_2 is not None:
+                    condition_embeds_2 = image_gen_condition_embeds_2
+                    negative_condition_embeds_2 = condition_embeds_2 * 0.0 if image_gen_negative_condition_embeds_2 is None else image_gen_negative_condition_embeds_2
+
+            else:
+                if image_gen_llm_hidden_states is None:
+                    assert self.model is not None
+                    assert self.vision is not None
+                    if pixel_values is not None:
+                        image_embeds = self.extract_image_feature(pixel_values, grid_thw=image_grid_thw)
+                
+                assert self.loaded_image_gen_modules is True, "please add `load_image_gen=True` in from_pretrained() method"
+                assert position_ids is None
+
+                
+                condition_embeds, condition_embeds_2 = self.get_condition_embeds_for_image_gen(
+                    input_ids=input_ids, 
+                    attention_mask=attention_mask,
+                    image_embeds=image_embeds, 
+                    position_ids=position_ids,
+                    use_cache=use_cache,
+                    image_grid_thw=image_grid_thw,
+                    llm_hidden_states=image_gen_llm_hidden_states,
+                )
+                if condition_embeds is not None:
+                    negative_condition_embeds = condition_embeds * 0.0
+                
+                if condition_embeds_2 is not None:
+                    negative_condition_embeds_2 = condition_embeds_2 * 0.0
+
+                # negative prompt feature is deprecated
+                # negative_condition_embeds = self.get_learnable_token_embeds_for_image_gen(
+                #     input_ids=image_gen_negative_input_ids, 
+                #     attention_mask=image_gen_negative_attention_mask,
+                #     image_embeds=image_embeds, 
+                #     position_ids=position_ids,
+                #     use_cache=use_cache,
+                #     image_grid_thw=image_grid_thw,
+                #     llm_hidden_states=image_gen_negative_llm_hidden_states,
+                # ) if ((image_gen_negative_input_ids is not None) or (image_gen_negative_llm_hidden_states is not None)) else condition_embeds * 0.0
+
+            
+
+                if image_gen_only_extract_hidden_states:
+                    return condition_embeds, negative_condition_embeds, condition_embeds_2, negative_condition_embeds_2
+            
+            assert (condition_embeds is not None) or (condition_embeds_2 is not None)
+            if (condition_embeds is not None) and (condition_embeds_2 is not None):
+                assert condition_embeds.shape[0] == condition_embeds_2.shape[0]
+
+            bsz = condition_embeds.shape[0] if condition_embeds is not None else condition_embeds_2.shape[0]
+
+            if image_gen_height is None or image_gen_width is None:
+                if isinstance(image_gen_highres, int):
+                    image_gen_height, image_gen_width = [image_gen_highres] * bsz, [image_gen_highres] * bsz
+                elif image_gen_highres is True:
+                    image_gen_height, image_gen_width = [1024] * bsz, [1024] * bsz
+                else:
+                    image_gen_height, image_gen_width = [512] * bsz, [512] * bsz
+            elif isinstance(image_gen_height, torch.Tensor) or isinstance(image_gen_width, torch.Tensor):
+                assert isinstance(image_gen_height, torch.Tensor), image_gen_height
+                assert isinstance(image_gen_width, torch.Tensor), image_gen_width
+                image_gen_height = image_gen_height.cpu().tolist()
+                image_gen_width = image_gen_width.cpu().tolist()
+                assert len(image_gen_height) == bsz
+                assert len(image_gen_width)  == bsz
+            elif isinstance(image_gen_height, int) or isinstance(image_gen_width, int):
+                assert isinstance(image_gen_height, int), image_gen_height
+                assert isinstance(image_gen_width, int), image_gen_width
+                image_gen_height = [image_gen_height] * bsz
+                image_gen_width = [image_gen_width] * bsz
+            else:
+                assert isinstance(image_gen_height, list), image_gen_height
+                assert isinstance(image_gen_width, list), image_gen_width
+                assert len(image_gen_height) == bsz
+                assert len(image_gen_width)  == bsz
+
+
+            image_gen_height_diffusion_list = []
+            image_gen_width_diffusion_list = []
+            image_gen_output_resize_height = []
+            image_gen_output_resize_width = []
+            for height, width in zip(image_gen_height, image_gen_width):
+                closest_size, resize_size = process_ratio(ori_h=height, ori_w=width, highres=image_gen_highres)
+                height, width = closest_size
+                image_gen_height_diffusion_list.append(height)
+                image_gen_width_diffusion_list.append(width)
+                height, width = resize_size
+                image_gen_output_resize_height.append(height)
+                image_gen_output_resize_width.append(width)
+
+            image_gen_height = image_gen_height_diffusion_list[0]
+            assert all([i == image_gen_height for i in image_gen_height_diffusion_list])
+            image_gen_width = image_gen_width_diffusion_list[0]
+            assert all([i == image_gen_width for i in image_gen_width_diffusion_list])
+
+            if image_gen_pixel_values_reference is not None:
+                assert (image_gen_height, image_gen_width) == (image_gen_pixel_values_reference.shape[-2], image_gen_pixel_values_reference.shape[-1])
+
+            if image_gen_seed is None or image_gen_seed < 0:
+                from datetime import datetime
+                image_gen_seed = datetime.now().microsecond % 1000
+                
+            sample_kwargs = {
+                "steps": image_gen_steps,
+                "seed": image_gen_seed,
+                "cfg": image_gen_cfg,
+                "height": image_gen_height,
+                "width": image_gen_width,
+                "cfg_mode": image_gen_cfg_mode,
+                "ref_x": image_gen_pixel_values_reference,
+                "encoder_hidden_states": condition_embeds,
+                "directvlm_hidden_states": condition_embeds_2,
+                "num_frames_per_prompt": num_frames_per_prompt,
+            }
+
+            image = self.diffusion_loss.sample(
+                **sample_kwargs,
+            )
+            if image_gen_task == "layer-decompose":
+                output_size = (
+                    image_gen_output_resize_width[0],
+                    image_gen_output_resize_height[0],
+                )
+                image = [item.resize(output_size, Image.LANCZOS) for item in image]
+            else:
+                image = [
+                    item.resize((width, height), Image.LANCZOS)
+                    for item, width, height in zip(
+                        image,
+                        image_gen_output_resize_width,
+                        image_gen_output_resize_height,
+                    )
+                ]
+
+            if (
+                image_gen_task != "layer-decompose"
+                and not image_gen_return_batch
+                and len(image) == 1
+            ):
+                image = image[0]
+            
+            return image
+
+        if pixel_values is not None:
+            image_embeds = self.extract_image_feature(pixel_values, grid_thw=image_grid_thw)
+        if pixel_values_videos is not None:
+            video_embeds = self.extract_image_feature(pixel_values_videos, grid_thw=video_grid_thw)
+        with torch.cuda.amp.autocast(dtype=torch.bfloat16):
+            outputs = self.model.generate(
+                input_ids=input_ids,
+                query_embeds_image=image_embeds,
+                query_embeds_video=video_embeds,
+                query_embeds_audio=audio_embeds,
+                query_embeds_audio_lengths=audio_embeds_lengths,
+                placeholder_audio_loc_lens=audio_placeholder_loc_lens,
+                image_grid_thw=image_grid_thw,
+                image_grid_thw_video=video_grid_thw,
+                attention_mask=attention_mask,
+                position_ids=position_ids,
+                past_key_values=past_key_values,
+                use_cache=use_cache,
+                num_logits_to_keep=num_logits_to_keep,
+                **generate_kwargs,
+            )
+        return outputs
+
+    def load_image_gen_modules(self, inference_model_path, torch_dtype=torch.float32, load_image_gen_diffusion=True, load_image_gen_others=True, device=None):
+        inference_model_path = str(resolve_model_directory(inference_model_path))
+        if os.path.exists(os.path.join(inference_model_path, "byt5")):
+            raise ValueError(
+                "Ming Image inference does not support a byt5 component; "
+                "the public checkpoint layout has no byt5/ directory."
+            )
+        self.inference_profile = load_checkpoint_capabilities(inference_model_path)
+        if device is not None:
+            device = torch.device(device)
+        elif self.model is not None:
+            device = self.model.device
+        else:
+            device = torch.device(torch.cuda.current_device())
+        logger.info(f"load_image_gen_modules device={device}")
+        from transformers import AutoModelForCausalLM
+        from safetensors.torch import load_file
+        temp_state_dict = load_file(
+            os.path.join(inference_model_path, "mlp", "model.safetensors")
+        )
+        with open(os.path.join(inference_model_path, 'mlp', 'config.json'), 'r') as f:
+            import json
+            metax_config = json.load(f)
+            diffusion_c_input_dim = metax_config.get("diffusion_c_input_dim", 2048)
+            self.img_gen_scales = metax_config.get("img_gen_scales", [4, 8, 16])
+            self.connector_norm = metax_config.get("connector_norm", True)
+            self.use_vlm_directvlm_condition = metax_config.get(
+                "use_vlm_directvlm_condition", False
+            )
+            self.use_learnable_token_condition = metax_config.get(
+                "use_learnable_token_condition", True
+            )
+            self.selected_hidden_states_layers = metax_config.get(
+                "selected_hidden_states_layers"
+            )
+            self.diffusion_inner_dim = metax_config.get("diffusion_inner_dim")
+
+        if load_image_gen_others:
+            self.connector = None    
+            self.query_tokens_dict = nn.ParameterDict()
+            # cumulative token index across the scales
+            self.scale_indices = []
+            current_idx = 0
+            for scale in self.img_gen_scales:                    
+                num_tokens = scale * scale
+                scale_name = f"{scale}x{scale}"
+                #weights = temp_state_dict[f"query_tokens_dict.{scale_name}"]
+                self.query_tokens_dict[scale_name] = nn.Parameter(
+                    torch.nn.functional.normalize(torch.randn(num_tokens, self.config.llm_config.hidden_size), dim=-1)
+                )
+                current_idx += scale * scale
+                self.scale_indices.append(current_idx)
+
+            self.query_tokens_dict.to(torch_dtype).to(device)
+
+            if self.use_learnable_token_condition:
+                modified_state_dict_query_tokens = {
+                    f"{scale}x{scale}": temp_state_dict[f"query_tokens_dict.{scale}x{scale}"]
+                    for scale in self.img_gen_scales   
+                }
+            
+                self.query_tokens_dict.load_state_dict(modified_state_dict_query_tokens, strict=True)
+                
+                # self.norm_query_embeds = True
+                # load connector
+                self.connector = AutoModelForCausalLM.from_pretrained(inference_model_path, subfolder='connector', torch_dtype=torch_dtype)
+                for layer in self.connector.model.layers:
+                    layer.self_attn.is_causal = False
+                self.connector.to(device)
+                
+                
+                self.proj_in = nn.Linear(self.config.llm_config.hidden_size, self.connector.config.hidden_size)
+                self.proj_out = nn.Linear(self.connector.config.hidden_size, diffusion_c_input_dim)
+                
+                modified_state_dict_in = {
+                    'weight': temp_state_dict['proj_in.weight'],
+                    'bias': temp_state_dict['proj_in.bias']
+                }
+                self.proj_in.load_state_dict(modified_state_dict_in, strict=True)
+                modified_state_dict_out = {
+                    'weight': temp_state_dict['proj_out.weight'],
+                    'bias': temp_state_dict['proj_out.bias']
+                }
+                self.proj_out.load_state_dict(modified_state_dict_out, strict=True)
+                self.proj_in.to(device=device, dtype=torch_dtype)
+                self.proj_out.to(device=device, dtype=torch_dtype)
+
+            self.proj_directvlm = None
+            if self.use_vlm_directvlm_condition:
+                directvlm_dim = self.model.config.hidden_size
+                if self.selected_hidden_states_layers is not None:
+                    directvlm_dim = directvlm_dim * len(self.selected_hidden_states_layers)
+
+                self.proj_directvlm = nn.Sequential(RMSNorm(directvlm_dim, eps=1e-5), nn.Linear(directvlm_dim, self.diffusion_inner_dim, bias=True))
+                
+                modified_state_dict_directvlm = {
+                    '0.weight': temp_state_dict["proj_directvlm.0.weight"],
+                    '1.weight': temp_state_dict["proj_directvlm.1.weight"],
+                    '1.bias': temp_state_dict["proj_directvlm.1.bias"],
+                }
+                self.proj_directvlm.load_state_dict(modified_state_dict_directvlm, strict=True)
+                self.proj_directvlm.to(device=device, dtype=torch_dtype)
+
+        if load_image_gen_diffusion:
+            diffusion_mlp_state_dict = {
+                key[len("mlp.") :] : temp_state_dict[key]
+                for key in temp_state_dict if key.startswith("mlp.")
+            }
+            from diffusion.generator import ImageGenerator
+
+            self.diffusion_loss = ImageGenerator(
+                model_path=inference_model_path,
+                scheduler_path=inference_model_path,
+                vision_dim=diffusion_c_input_dim,
+                mlp_state_dict=diffusion_mlp_state_dict,
+                torch_dtype=torch_dtype,
+                device=device,
+                use_identity_mlp=metax_config.get("use_identity_mlp", False),
+                text_encoder_norm=metax_config.get("text_encoder_norm", False),
+                inference_profile=self.inference_profile,
+            )
+            self.diffusion_loss.to(device)
+        self.loaded_image_gen_modules = True
+    @classmethod
+    def _from_int8_checkpoint(cls, vlm_directory, device, **kwargs):
+        """Load an mllm/ component written by quant/quantize_stream.py (weight-only int8).
+
+        The model is built with its parameters on the meta device, the Linear modules listed
+        in int8_manifest.json become Int8Linear shells, and every stored tensor is loaded
+        straight onto `device`, so BF16 weights for the quantized modules never exist in memory.
+        """
+        from accelerate import init_empty_weights
+        from quant.load_int8 import load_int8_mllm_
+
+        device_map = kwargs.pop("device_map", None)
+        if device_map is not None:
+            # infer.py's default "balanced" plan on a single-GPU box maps every module to GPU 0;
+            # that is honoured. Splitting the int8 model across devices is not supported.
+            targets = set(device_map.values()) if isinstance(device_map, dict) else {device_map}
+            if len(targets) != 1 or not isinstance(next(iter(targets)), int):
+                raise ValueError(
+                    "the int8 mllm checkpoint loads onto a single GPU; device_map targets "
+                    f"{sorted(map(str, targets))} (use --device-map none)"
+                )
+            device = torch.device("cuda", next(iter(targets)))
+        supported = ("torch_dtype", "dtype", "attn_implementation")
+        unsupported = sorted(key for key in kwargs if key not in supported)
+        if unsupported:
+            raise ValueError(
+                f"the int8 mllm checkpoint loads onto a single device; unsupported arguments: {unsupported}"
+            )
+        device = torch.device(device) if device is not None else torch.device("cpu")
+        if device.type == "cuda" and device.index is None:
+            device = torch.device("cuda", torch.cuda.current_device())
+        config = BailingMM2Config.from_pretrained(vlm_directory)
+        with init_empty_weights():
+            model = cls._from_config(config, **kwargs)
+        report = load_int8_mllm_(model, vlm_directory, device)
+        # Buffers built in __init__ (rotary inv_freq) are not stored in the checkpoint; they follow
+        # the weights. Int8Linear keeps its fp32 scales through this and any later dtype cast.
+        model.to(device)
+        logger.info(f"int8 mllm loaded from {vlm_directory}: {report}")
+        model.tie_weights()
+        model.eval()
+        return model
+
+    @classmethod
+    def from_pretrained(
+        cls,
+        pretrained_model_name_or_path: Optional[Union[str, os.PathLike]],
+        *model_args,
+        **kwargs,
+    ):
+        load_image_gen = False
+        if "load_image_gen" in kwargs:
+            load_image_gen = kwargs["load_image_gen"]
+            del kwargs["load_image_gen"]
+        load_image_gen_diffusion = True
+        if "load_image_gen_diffusion" in kwargs:
+            load_image_gen_diffusion = kwargs["load_image_gen_diffusion"]
+            del kwargs["load_image_gen_diffusion"]
+        
+        load_image_gen_others = True
+        if "load_image_gen_others" in kwargs:
+            load_image_gen_others = kwargs["load_image_gen_others"]
+            del kwargs["load_image_gen_others"]
+        load_vlm = True
+        if "load_vlm" in kwargs:
+            load_vlm = kwargs["load_vlm"]
+            del kwargs["load_vlm"]
+        image_gen_device = kwargs.pop("image_gen_device", None)
+        vlm_directory = pretrained_model_name_or_path
+        if load_image_gen:
+            pretrained_model_name_or_path = str(
+                resolve_model_directory(
+                    pretrained_model_name_or_path,
+                    revision=kwargs.get("revision"),
+                    cache_dir=kwargs.get("cache_dir"),
+                    local_files_only=kwargs.get("local_files_only", False),
+                    token=kwargs.get("token"),
+                )
+            )
+            # The package root keeps connector/mlp/transformer/vae/scheduler;
+            # the MLLM itself (config, weights, tokenizer data) lives in mllm/.
+            vlm_directory = os.path.join(pretrained_model_name_or_path, "mllm")
+            if not os.path.isdir(vlm_directory):
+                raise FileNotFoundError(
+                    "checkpoint package is missing the mllm/ component: "
+                    f"{vlm_directory}. Migrate the package to the component "
+                    "layout before loading."
+                )
+        if load_vlm and os.path.exists(os.path.join(vlm_directory, "int8_manifest.json")):
+            model = cls._from_int8_checkpoint(vlm_directory, image_gen_device, **kwargs)
+        elif load_vlm:
+            model = super().from_pretrained(
+                vlm_directory,
+                *model_args,
+                **kwargs,
+            )
+        else:
+            model = cls(
+                BailingMM2Config.from_dict(BailingMM2Config.get_config_dict(vlm_directory)[0]),
+                empty_load=True,
+            )
+        if load_image_gen:
+            model.load_image_gen_modules(
+                pretrained_model_name_or_path, 
+                torch_dtype=kwargs["torch_dtype"] if "torch_dtype" in kwargs else torch.float32,
+                load_image_gen_diffusion=load_image_gen_diffusion,
+                load_image_gen_others=load_image_gen_others,
+                device=image_gen_device,
+            )
+        return model
+    
+    def append_input_ids_with_multiscale_learnable_tokens(
+        self,
+        text_ids,
+        attention_mask,
+        scales,
+        start_token_id,
+        end_token_id,
+        patch_token_id,
+    ):
+        default_scaled_tokens = []
+        default_scaled_attn_masks = []
+        default_gen_masks = []
+        for scale in scales:
+            default_scaled_tokens.append(start_token_id)
+            default_scaled_tokens.extend([patch_token_id for _ in range(scale * scale)])
+            default_scaled_tokens.append(end_token_id)
+            default_scaled_attn_masks.extend([1 for _ in range(scale * scale + 2)])
+            default_gen_masks.append(0)
+            default_gen_masks.extend([1 for _ in range(scale * scale)])
+            default_gen_masks.append(0)
+        
+        text_ids_list = text_ids.cpu().tolist()
+        attention_mask_list = attention_mask.cpu().tolist()
+
+        new_text_ids_list = []
+        new_attention_mask_list = []
+        gen_mask_list = []
+        new_labels_list = []
+        for text_ids_one_batch, attention_mask_one_batch in zip(
+            text_ids_list, attention_mask_list
+        ):
+            assert len(text_ids_one_batch) == len(attention_mask_one_batch)
+
+            padding_start = 0
+            for idx, value in enumerate(attention_mask_one_batch):
+                if value == 0:
+                    break
+                
+                padding_start += 1
+
+            new_text_ids_list.append(text_ids_one_batch[:padding_start] + deepcopy(default_scaled_tokens) + text_ids_one_batch[padding_start:])
+            new_labels_list.append([ -100 for _ in range(padding_start)] + [1 for _ in range(len(default_scaled_tokens))] + [-100 for _ in range(len(text_ids_one_batch[padding_start:]))] )
+
+            new_attention_mask_list.append(attention_mask_one_batch[:padding_start] + deepcopy(default_scaled_attn_masks) + attention_mask_one_batch[padding_start:])
+            gen_mask_list.append(
+                [0 for _ in range(len(attention_mask_one_batch[:padding_start]))] + \
+                deepcopy(default_gen_masks) + \
+                [0 for _ in range(len(attention_mask_one_batch[padding_start:]))]
+            )
+
+        text_ids_append_lq = torch.tensor(new_text_ids_list, dtype=text_ids.dtype).to(text_ids.device)
+        attention_mask_append_lq = torch.tensor(new_attention_mask_list, dtype=attention_mask.dtype).to(attention_mask.device)
+        gen_mask = torch.tensor(gen_mask_list, dtype=attention_mask.dtype).to(attention_mask.device)
+        labels = torch.tensor(new_labels_list, dtype=text_ids.dtype).to(text_ids.device)
+
+        assert attention_mask_append_lq.shape == text_ids_append_lq.shape
+        assert labels.shape == text_ids_append_lq.shape
+        assert gen_mask.shape == text_ids_append_lq.shape
+        return text_ids_append_lq, labels, attention_mask_append_lq, gen_mask
+
+    def appand_learnable_tokens(
+        self,
+        text_ids,
+        gen_mask,
+        image_embeds,
+        image_grid_thw,
+        patch_token_id,
+    ):
+        query_tokens_embeds = torch.cat(
+            [self.query_tokens_dict[f"{scale}x{scale}"] for scale in self.img_gen_scales], 
+            dim=0,
+        )
+        if image_embeds is not None:
+            query_tokens_embeds = query_tokens_embeds.to(image_embeds.dtype).to(image_embeds.device)
+
+        assert text_ids.shape == gen_mask.shape
+        text_ids_aslist = text_ids.cpu().view(-1).tolist()
+        gen_mask_aslist = gen_mask.cpu().view(-1).tolist()
+        is_patch_list = [1 if i == patch_token_id else 0 for i in text_ids_aslist]
+        idxes_start_of_patch = find_first_index_of_consecutive_ones(is_patch_list)
+        isgen_indicators = merge_consecutive_ones([1 if gen_mask_aslist[i] else 0 for i in idxes_start_of_patch], len(self.img_gen_scales))
+        if any([i == 0 for i in isgen_indicators]):
+            assert image_grid_thw is not None
+            assert image_grid_thw.ndim == 2
+            assert image_embeds is not None
+            assert image_embeds.ndim == 2
+
+        new_image_grid_thw = []
+        new_image_embeds = []
+        cum_image_token = 0
+        cnt_input_image = 0
+
+        for is_gen in isgen_indicators:
+            if is_gen:
+                for scale in self.img_gen_scales:
+                    new_image_grid_thw.append([1, 2, scale * scale * 2])
+                
+                new_image_embeds.append(query_tokens_embeds)
+            else:
+                thw = image_grid_thw[cnt_input_image].tolist()
+                assert thw[0] == 1
+                assert thw[1] % 2 == 0 # h
+                assert thw[2] % 2 == 0 # w
+                n_image_token = (thw[1] // 2) * (thw[2] // 2)
+                image_embed_one = image_embeds[cum_image_token : cum_image_token + n_image_token, :]
+                new_image_embeds.append(image_embed_one)
+                new_image_grid_thw.append(thw)
+                cnt_input_image += 1
+                cum_image_token += n_image_token
+
+        if image_grid_thw is not None:
+            assert cnt_input_image == image_grid_thw.shape[0]
+            assert cum_image_token == image_embeds.shape[0]
+        else:
+            assert cnt_input_image == 0
+            assert cum_image_token == 0
+
+        new_image_grid_thw = torch.tensor(new_image_grid_thw, dtype=text_ids.dtype).to(text_ids.device)
+        new_image_embeds = torch.cat(new_image_embeds, dim=0).to(text_ids.device)
+
+        total_patch_token = 0
+        for bid in range(new_image_grid_thw.shape[0]):
+            thw = new_image_grid_thw[bid].tolist()
+            assert thw[0] == 1
+            assert thw[1] % 2 == 0
+            assert thw[2] % 2 == 0
+            patch_h = thw[1] // 2
+            patch_w = thw[2] // 2
+            n_patch_token = patch_h * patch_w
+            total_patch_token += n_patch_token
+        
+        # if torch.distributed.get_rank() == 0:
+        #     embed()
+        # torch.distributed.barrier()
+        
+        assert total_patch_token == new_image_embeds.shape[0], f"{total_patch_token}, vs. {new_image_embeds.shape}"
+        
+        return new_image_grid_thw, new_image_embeds
+
+    def get_condition_embeds_for_image_gen(
+        self,
+        input_ids, 
+        attention_mask,
+        image_embeds, 
+        position_ids,
+        use_cache,
+        image_grid_thw,
+        llm_hidden_states,
+    ):
+        input_ids, labels, attention_mask, gen_mask = self.append_input_ids_with_multiscale_learnable_tokens(
+            input_ids,
+            attention_mask,
+            self.img_gen_scales,
+            self.config.llm_config.image_patch_token + 1,
+            self.config.llm_config.image_patch_token + 2,
+            self.config.llm_config.image_patch_token,
+        )
+        
+        if llm_hidden_states is None:
+            image_grid_thw, image_embeds = self.appand_learnable_tokens(
+                input_ids,
+                gen_mask,
+                image_embeds,
+                image_grid_thw,
+                self.config.llm_config.image_patch_token,
+            )
+            
+            with torch.cuda.amp.autocast(dtype=torch.bfloat16):
+                if image_embeds is None or input_ids.size(1) == 1:
+                    words_embeddings = self.model.get_input_embeddings()(input_ids.clip(0, self.model.get_input_embeddings().weight.shape[0] - 1))
+                    image_mask = None
+                    audio_mask = None
+                else:
+                    words_embeddings, image_mask, audio_mask = self.model.model.prompt_wrap_navit(
+                        input_ids=input_ids.clip(0, self.model.get_input_embeddings().weight.shape[0] - 1),
+                        config=self.model.model.config, 
+                        query_embeds_image=image_embeds, 
+                    )
+                
+                assert input_ids.size(1) == words_embeddings.size(1), "{} vs {}".format(
+                    input_ids.size,
+                    words_embeddings.size,
+                )
+
+                # if torch.distributed.get_rank() == 3:
+                #     embed()
+                # torch.distributed.barrier()
+
+                outputs = self.model.forward(
+                    input_ids=input_ids,
+                    attention_mask=attention_mask,
+                    position_ids=position_ids,
+                    past_key_values=None,
+                    inputs_embeds=words_embeddings,
+                    image_grid_thw=image_grid_thw,
+                    use_cache=False,
+                    image_mask=image_mask,
+                    audio_mask=None,
+                    output_hidden_states=True,
+                )
+                hidden_states = outputs.hidden_states[-1]
+        else:
+            hidden_states = llm_hidden_states
+
+        directvlm_hidden_states = None
+        if self.use_vlm_directvlm_condition:
+            # use hidden states
+            use_input_mask = torch.lt(labels, 0).int().to(attention_mask.dtype) * attention_mask
+            assert use_input_mask.ndim == 2
+            directvlm_max_valid_ind = use_input_mask.cumsum(-1).argmax(-1).max().item() + 1
+            #directvlm_max_valid_ind = min(directvlm_max_valid_ind, self.max_vlm_directvlm_length)
+            use_input_mask = use_input_mask[:, :directvlm_max_valid_ind]
+
+            if self.selected_hidden_states_layers is not None:
+                directvlm_hidden_states = torch.cat([
+                    outputs.hidden_states[layer_i].to(labels.device)[:, :directvlm_max_valid_ind, :] * use_input_mask.unsqueeze(-1)
+                    for layer_i in self.selected_hidden_states_layers
+                ], dim=-1)
+            else:
+                directvlm_hidden_states = outputs.hidden_states[-1].to(labels.device)[:, :directvlm_max_valid_ind, :] * use_input_mask.unsqueeze(-1)
+            
+            directvlm_hidden_states = directvlm_hidden_states.detach()
+            directvlm_hidden_states = self.proj_directvlm(directvlm_hidden_states)
+
+        scale_embeds = None
+        if self.use_learnable_token_condition:
+            with torch.cuda.amp.autocast(dtype=torch.bfloat16):
+                gen_mask = gen_mask.unsqueeze(-1).expand(gen_mask.shape[0], gen_mask.shape[1], hidden_states.shape[-1]).to(hidden_states.device).bool()
+                hidden_states_gen = torch.masked_select(hidden_states, gen_mask).view(hidden_states.shape[0], -1, hidden_states.shape[-1])
+                # split hidden_states into per-scale representations
+                scale_start_idxes = [0] + self.scale_indices[:-1]
+                scale_end_idxes = self.scale_indices
+                assert scale_end_idxes[-1] == hidden_states_gen.shape[1]
+                
+                scale, scale_start_idx, scale_end_idx = [
+                    i for i in zip(self.img_gen_scales, scale_start_idxes, scale_end_idxes)
+                ][-1]
+                
+                scale_hidden = hidden_states_gen[:, scale_start_idx : scale_end_idx, :]
+                scale_embeds = self.proj_in(scale_hidden)
+                seq_shape = scale_embeds.shape
+                with torch.cuda.amp.autocast(dtype=torch.bfloat16):
+                    scale_embeds = self.connector(
+                        inputs_embeds=scale_embeds, 
+                        attention_mask=torch.ones(seq_shape[0],1,seq_shape[1],seq_shape[1]).to(scale_embeds.device), 
+                        output_hidden_states=True
+                    ).hidden_states[-1]
+                    
+                scale_embeds = self.proj_out(scale_embeds)
+                # normalize
+                if self.connector_norm:
+                    scale_embeds = torch.nn.functional.normalize(scale_embeds, dim=-1)
+
+        return scale_embeds, directvlm_hidden_states
+
+__all__ = [
+    "BailingMM2NativeForConditionalGeneration"
+]

code/modeling_utils.py

@@ -0,0 +1,191 @@
+import torch
+import torch.nn as nn
+from typing import Optional
+
+import logging
+
+logger = logging.getLogger(__name__)
+
+class Transpose(nn.Module):
+    def __init__(self, dim0: int, dim1: int):
+        super().__init__()
+        self.dim0 = dim0
+        self.dim1 = dim1
+
+    def forward(self, x):
+        return x.transpose(self.dim0, self.dim1)
+
+def patch_continuous_features(
+    input_embeddings: torch.Tensor,
+    placeholder_loc_lens: torch.Tensor,
+    encoded_feats: torch.Tensor,
+    encoded_feat_lens: torch.Tensor,
+):
+    """
+    Patch continuous features into input embeddings, while keeping a valid gradient flow.
+
+    input_embeddings: torch.Tensor, size = [B, C?, T, D]
+    placeholder_loc_lens: torch.LongTensor, size = [B, N, 2]
+        Each 2-tuple represents (start, length) of a placeholder.
+    encoded_feats: torch.Tensor, size = [B, L1 + L2 + ... + LN, ...]
+    encoded_feat_lens: torch.LongTensor, size = [B, N]
+
+    Example ('X' for patch placeholder tokens):
+    Inputs:
+        input_embeddings = [[1, 2, 3, X, X, X, 4, 5, 6, X, X, X, 7, 8]]
+        placeholder_loc_lens = [[[3, 3]], [[9, 3]]]
+        encoded_feats = [[A, A, A, B, B]]
+        encoded_feat_lens = [[3], [2]]
+    Outputs:
+        embeddings = [[1, 2, 3, A, A, A, 4, 5, 6, B, B, X, 7, 8]]
+    """
+    batch_size = input_embeddings.size(0)
+    for i in range(batch_size):
+        audio_feat_start = 0
+        for j in range(placeholder_loc_lens.shape[1]):
+            placeholder_start: int = int(placeholder_loc_lens[i, j, 0].item())
+            placeholder_len: int = int(placeholder_loc_lens[i, j, 1].item())
+            if placeholder_len <= 0:
+                break
+            feat_len = int(encoded_feat_lens[i, j].item())
+            real_feat_len = feat_len
+            if feat_len > placeholder_len:
+                logger.warning(
+                    f"Feature length ({feat_len}) > placeholder length ({placeholder_len}). "
+                    "This is not expected; please check estimate_audio_feature_length(). "
+                    "We truncate the feature to avoid errors."
+                )
+                feat_len = placeholder_len
+            target_len = min(feat_len, placeholder_len)
+            input_embeddings[i, placeholder_start:placeholder_start + target_len] = encoded_feats[i, audio_feat_start:audio_feat_start + target_len]
+            audio_feat_start += real_feat_len
+    return input_embeddings
+
+def build_modality_mask(placeholder_loc_lens: torch.Tensor, shape: torch.Size):
+    mask = torch.zeros(shape, dtype=torch.bool)
+    for i in range(placeholder_loc_lens.shape[0]):
+        for j in range(placeholder_loc_lens.shape[1]):
+            start: int = int(placeholder_loc_lens[i, j, 0].item())
+            length: int = int(placeholder_loc_lens[i, j, 1].item())
+            if length <= 0:
+                break
+            mask[i, start:start + length] = True
+    return mask
+
+def encode_audio_segments(
+    encoder,
+    proj_layer,
+    wav_feats=None,
+    wav_feats_lengths=None,
+    waveforms=None,
+    waveforms_lengths=None,
+    use_waveform=False,
+    audio_config=None,
+):
+    """
+    Apply audio encoder to input audio features in wrapped format.
+    See the documentation of unwrap_feats() for details about 'wrapped format'.
+    """
+
+    # Forward audio encoder.
+    if use_waveform:
+        assert waveforms is not None and waveforms_lengths is not None
+        # Unwrap the waveforms so each waveform is placed at an independent row.
+        waveform_segs_batch, waveform_seg_lengths = unwrap_feats(waveforms, waveforms_lengths)
+        audio_feats_seg, audio_feat_seg_lengths = encoder(waveform_segs_batch, waveform_seg_lengths)[:2]
+    else:
+        assert wav_feats is not None and wav_feats_lengths is not None
+        # Unwrap the features so the feature of each waveform is placed at an independent row.
+        feat_segs_batch, feat_seg_lengths = unwrap_feats(wav_feats, wav_feats_lengths)
+        # for whisper encoder
+        # feat_segs_batch: [B, T, n_mels]
+        # feat_seg_lengths: [B]
+        audio_feats_seg = encoder(feat_segs_batch)
+        audio_feats_seg_proj = proj_layer(audio_feats_seg.transpose(-1, -2)).transpose(-1, -2)
+        feat_seg_lengths = feat_seg_lengths.to(feat_segs_batch.device)
+        # whisper encoder conv
+        audio_feat_seg_lengths = (feat_seg_lengths - 3 + 2 * 1) // 2 + 1
+        # project layer conv
+        audio_feat_seg_lengths = (audio_feat_seg_lengths - audio_config.ds_kernel_size + 2 *
+                                  (audio_config.ds_kernel_size//2)) // audio_config.ds_stride + 1
+
+    # Wrap the features so the 1st dim represents batch_size.
+    input_lengths = waveforms_lengths if use_waveform else wav_feats_lengths
+    assert input_lengths is not None
+    audio_feats, _, audio_feats_lengths = wrap_feats(audio_feats_seg, input_lengths, audio_feat_seg_lengths)
+    audio_feats_proj, _, audio_feats_lengths2 = wrap_feats(audio_feats_seg_proj, input_lengths, audio_feat_seg_lengths)
+    assert torch.all(audio_feats_lengths == audio_feats_lengths2), f"{audio_feats_lengths}, {audio_feats_lengths2}"
+
+    return audio_feats_proj, audio_feats, audio_feats_lengths
+
+def unwrap_feats(feats: torch.Tensor, feats_lengths: torch.Tensor):
+    """
+    The input feats are in the "wrapped" format, which means that features from (at most) N audios are concatenated
+    as a single sample feats[i]. In this case, each row of feats_lengths contains the lengths of the concatenated
+    feature. This function unwraps the features.
+    For samples with less than N segments, one should pad feats_lengths with 0. The result will contain valid
+    segments only.
+
+    feats: torch.Tensor, size = [B, L1 + L2 + ... + LN, ...]
+    feats_lengths: torch.LongTensor, size = [B, N]
+
+    Example ('X' for padding):
+    Inputs:
+        feats = [[A, A, A, A, X],
+                 [B, B, C, C, C]]
+        feats_lengths = [[4, 0],
+                         [2, 3]]
+    Outputs:
+        feat_segs = [[A, A, A, A],
+                     [B, B, X, X],
+                     [C, C, C, X]]
+        feat_seg_lengths = [4, 2, 3]
+    """
+    feat_segs = []
+    feat_seg_lengths = []
+    for i in range(feats_lengths.shape[0]):
+        feat_index = 0
+        for j in range(feats_lengths.shape[1]):
+            feat_len = feats_lengths[i, j].item()
+            if feat_len == 0: break
+            feat_segs.append(feats[i, feat_index:feat_index + feat_len])
+            feat_seg_lengths.append(feat_len)
+            feat_index += feat_len
+    feat_segs_batch = torch.nn.utils.rnn.pad_sequence(feat_segs, True).to(feats.device)
+    feat_seg_lengths = torch.tensor(feat_seg_lengths, dtype=torch.long, device=feats.device)
+    return feat_segs_batch, feat_seg_lengths
+
+def wrap_feats(feat_segs: torch.Tensor, feats_lengths: torch.Tensor, feats_seg_lengths: Optional[torch.Tensor] = None):
+    """
+    Wrap segmented features back to the wrapped format.
+    This function is the inverse operation of unwrap_feats(). See its documentation for details.
+    Note that the feats_lengths value does not matter a lot. We only check the location of the first 0 to determine the
+    number of feature segments.
+    """
+    feat_idx = 0
+    feats_buffer = []
+    feats_locs_buffer = []
+    feats_lengths_buffer = []
+    for i in range(feats_lengths.shape[0]):
+        feat_buffer = []
+        feat_locs_buffer = []
+        feat_lengths_buffer = []
+        feat_total_len = 0
+        for j in range(feats_lengths.shape[1]):
+            feat_len = feats_lengths[i, j].item()
+            if feat_len == 0:
+                break
+            if feats_seg_lengths is not None:
+                feat_len = feats_seg_lengths[feat_idx].item()
+            feat_buffer.append(feat_segs[feat_idx, :feat_len])
+            feat_locs_buffer.append(feat_total_len)
+            feat_lengths_buffer.append(feat_len)
+            feat_idx += 1
+            feat_total_len += feat_len
+        feats_buffer.append(torch.cat(feat_buffer))
+        feats_locs_buffer.append(torch.tensor(feat_locs_buffer, dtype=torch.long))
+        feats_lengths_buffer.append(torch.tensor(feat_lengths_buffer, dtype=torch.long))
+    feats = torch.nn.utils.rnn.pad_sequence(feats_buffer, True).to(feat_segs.device)
+    feats_locs = torch.nn.utils.rnn.pad_sequence(feats_locs_buffer, True).to(feats_lengths.device)
+    feats_new_lengths = torch.nn.utils.rnn.pad_sequence(feats_lengths_buffer, True).to(feats_lengths.device)
+    return feats, feats_locs, feats_new_lengths

code/pe_ling.py

@@ -0,0 +1,446 @@
+#!/usr/bin/env python3
+"""Prompt enhancement (PE) for Ming-Image text-to-image via a Ling-3.0-flash-VL seat.
+
+Per the README, PE is a pre-processing step *outside* ``infer.py``: an
+instruction-following VLM rewrites a short caption into the structured
+Figma-style JSON prompt that the text-to-image pipeline consumes, and the
+result is passed to ``infer.py --prompt`` as raw text or via a file.
+
+This module drives any OpenAI-compatible ``/chat/completions`` endpoint using
+only the standard library (``urllib``): by default the local llama-server seat
+serving Ling-3.0-flash-VL, optionally the LiteLLM lab gateway (Bearer auth via
+``--api-key`` or the ``LITELLM_API_KEY`` environment variable). The rewriter
+system prompt is read verbatim from ``assets/t2i_rewriter_system_prompt.txt``.
+
+The reply is parsed robustly (```json fences and surrounding prose are
+tolerated), then validated against the schema the system prompt demands. On a
+parse or validation failure the request is retried exactly once with the
+errors appended to the user turn; if that still fails, PromptEnhancementError
+is raised with the errors. Invalid JSON is never passed through silently.
+
+CLI:
+    python pe_ling.py "a caption" --out prompt.json \
+        [--base-url http://127.0.0.1:8090/v1] \
+        [--model ling-3.0-flash-vl-mtp-halo-STRIX_LEAN]
+"""
+
+from __future__ import annotations
+
+import argparse
+import json
+import os
+import re
+import sys
+import time
+import urllib.error
+import urllib.request
+from pathlib import Path
+from typing import Any, Dict, List, Optional, Tuple
+
+CODE_DIRECTORY = Path(__file__).resolve().parent
+SYSTEM_PROMPT_PATH = CODE_DIRECTORY / "assets" / "t2i_rewriter_system_prompt.txt"
+
+# The Ling-3.0-flash-VL seat already served on the target box (llama-server,
+# OpenAI-compatible, thinking disabled); both endpoints speak the same
+# /chat/completions protocol.
+DEFAULT_BASE_URL = "http://127.0.0.1:8090/v1"
+DEFAULT_MODEL = "ling-3.0-flash-vl-mtp-halo-STRIX_LEAN"
+API_KEY_ENV = "LITELLM_API_KEY"
+
+# Low temperature: the rewrite is a deterministic schema transformation, not
+# creative sampling.
+DEFAULT_TEMPERATURE = 0.2
+# The upstream example rewrite (assets/t2i_four_seasons_cabin_prompt.json) is
+# ~5 KB (~2k tokens); dense multi-layer infographic rewrites run several times
+# longer, so leave generous headroom for a complete JSON object.
+DEFAULT_MAX_TOKENS = 16384
+# A multi-thousand-token completion on the local seat can take minutes.
+DEFAULT_TIMEOUT_SECONDS = 600.0
+
+REPAIR_INSTRUCTION = "Return only the corrected JSON object: no prose, no code fences."
+
+CANVAS_SETTINGS_KEYS = ("aspect_ratio", "ambient_lighting", "image_style")
+LAYER_KEYS = ("description", "coordinates", "hierarchy_and_relation", "color_specs")
+COORDINATE_FIELDS = ("cx", "cy", "w", "h")
+
+# `coordinates` must be ONE string of the form
+# "cx: 0.500, cy: 0.500, w: 1.000, h: 1.000". The upstream example also uses
+# bare integers ("h: 1"), so accept any decimal spelling and enforce the
+# [0, 1] range on the parsed value. Whitespace around ':' and ',' is
+# tolerated; the key order is fixed.
+_COORDINATE_NUMBER = r"[-+]?(?:\d+(?:\.\d*)?|\.\d+)"
+COORDINATES_RE = re.compile(
+    rf"^\s*cx:\s*(?P<cx>{_COORDINATE_NUMBER})\s*,"
+    rf"\s*cy:\s*(?P<cy>{_COORDINATE_NUMBER})\s*,"
+    rf"\s*w:\s*(?P<w>{_COORDINATE_NUMBER})\s*,"
+    rf"\s*h:\s*(?P<h>{_COORDINATE_NUMBER})\s*$"
+)
+
+# Hex colors: #RGB, #RGBA, #RRGGBB, #RRGGBBAA (the upstream example uses
+# #RRGGBB; the alpha forms keep RGBA-design outputs from failing validation).
+HEX_COLOR_RE = re.compile(
+    r"^#(?:[0-9a-fA-F]{3}|[0-9a-fA-F]{4}|[0-9a-fA-F]{6}|[0-9a-fA-F]{8})$"
+)
+
+
+class PromptEnhancementError(RuntimeError):
+    """PE failed: transport/protocol error, or schema failure after the retry."""
+
+    def __init__(
+        self,
+        message: str,
+        errors: Optional[List[str]] = None,
+        reply: Optional[str] = None,
+    ):
+        super().__init__(message)
+        self.errors = list(errors or [])
+        self.reply = reply
+
+
+def load_system_prompt(path: Path = SYSTEM_PROMPT_PATH) -> str:
+    """Return the released rewriter system prompt, verbatim."""
+    return path.read_text(encoding="utf-8")
+
+
+def extract_json_object(text: str) -> Dict[str, Any]:
+    """Return the first complete top-level JSON object found in ``text``.
+
+    Models sometimes wrap JSON in ```json fences or add prose around it.
+    Scanning every ``{`` position with ``JSONDecoder.raw_decode`` (which
+    decodes a document at an offset and ignores trailing data) recovers the
+    object in all of those shapes. Raises ValueError when no complete JSON
+    object is present, e.g. a reply truncated mid-object.
+    """
+    decoder = json.JSONDecoder()
+    position = text.find("{")
+    while position != -1:
+        try:
+            document, _ = decoder.raw_decode(text, position)
+        except ValueError:
+            position = text.find("{", position + 1)
+            continue
+        return document
+    snippet = text.strip()
+    if len(snippet) > 300:
+        snippet = snippet[:300] + "..."
+    raise ValueError(
+        f"reply contains no complete top-level JSON object "
+        f"({len(text)} characters); starts with: {snippet!r}"
+    )
+
+
+def _check_exact_keys(
+    mapping: Dict[str, Any], expected: Tuple[str, ...], path: str, errors: List[str]
+) -> None:
+    missing = [key for key in expected if key not in mapping]
+    unexpected = [key for key in mapping if key not in expected]
+    if missing:
+        errors.append(f"{path}: missing required key(s): {', '.join(missing)}")
+    if unexpected:
+        errors.append(
+            f"{path}: unexpected key(s): {', '.join(unexpected)} "
+            f"(exactly {', '.join(expected)} are required)"
+        )
+
+
+def _check_non_empty_string(value: Any, path: str, errors: List[str]) -> None:
+    if not isinstance(value, str):
+        errors.append(f"{path}: expected a string, got {type(value).__name__}")
+    elif not value.strip():
+        errors.append(f"{path}: string is empty")
+
+
+def _check_coordinates(value: Any, path: str, errors: List[str]) -> None:
+    if not isinstance(value, str):
+        errors.append(
+            f"{path}: must be ONE string of the form "
+            f"'cx: 0.500, cy: 0.500, w: 1.000, h: 1.000', got {type(value).__name__}"
+        )
+        return
+    match = COORDINATES_RE.match(value)
+    if match is None:
+        errors.append(
+            f"{path}: {value!r} is not of the form "
+            f"'cx: 0.500, cy: 0.500, w: 1.000, h: 1.000'"
+        )
+        return
+    for field in COORDINATE_FIELDS:
+        number = float(match.group(field))
+        if not 0.0 <= number <= 1.0:
+            errors.append(f"{path}: {field}={match.group(field)} is outside [0, 1]")
+
+
+def _check_color_specs(value: Any, path: str, errors: List[str]) -> None:
+    if not isinstance(value, list):
+        errors.append(
+            f"{path}: expected a list of hex colors, got {type(value).__name__}"
+        )
+        return
+    for index, color in enumerate(value):
+        if not isinstance(color, str) or HEX_COLOR_RE.match(color) is None:
+            errors.append(
+                f"{path}[{index}]: {color!r} is not a hex color "
+                f"(expected #RGB, #RGBA, #RRGGBB, or #RRGGBBAA)"
+            )
+
+
+def validate_enhanced_prompt(document: Any) -> List[str]:
+    """Return schema errors for a rewritten prompt; an empty list means valid.
+
+    Schema demanded by assets/t2i_rewriter_system_prompt.txt: exactly two
+    top-level keys ``canvas_settings`` (exactly ``aspect_ratio``,
+    ``ambient_lighting``, ``image_style``) and ``layers`` (each layer exactly
+    ``description``, ``coordinates``, ``hierarchy_and_relation``,
+    ``color_specs``); ``coordinates`` is a string "cx: 0.500, cy: 0.500,
+    w: 1.000, h: 1.000" with values in [0, 1]; ``color_specs`` is a list of
+    hex colors. ``layers`` must hold at least one visible layer -- an empty
+    list means the rewrite failed even though it is type-correct.
+    """
+    if not isinstance(document, dict):
+        return [f"top level: expected a JSON object, got {type(document).__name__}"]
+    errors: List[str] = []
+    _check_exact_keys(document, ("canvas_settings", "layers"), "top level", errors)
+
+    if "canvas_settings" in document:
+        canvas = document["canvas_settings"]
+        if not isinstance(canvas, dict):
+            errors.append(
+                f"canvas_settings: expected a JSON object, got {type(canvas).__name__}"
+            )
+        else:
+            _check_exact_keys(canvas, CANVAS_SETTINGS_KEYS, "canvas_settings", errors)
+            for key in CANVAS_SETTINGS_KEYS:
+                if key in canvas:
+                    _check_non_empty_string(
+                        canvas[key], f"canvas_settings.{key}", errors
+                    )
+
+    if "layers" in document:
+        layers = document["layers"]
+        if not isinstance(layers, list):
+            errors.append(f"layers: expected a list, got {type(layers).__name__}")
+        elif not layers:
+            errors.append("layers: expected at least one visible layer")
+        else:
+            for index, layer in enumerate(layers):
+                path = f"layers[{index}]"
+                if not isinstance(layer, dict):
+                    errors.append(
+                        f"{path}: expected a JSON object, got {type(layer).__name__}"
+                    )
+                    continue
+                _check_exact_keys(layer, LAYER_KEYS, path, errors)
+                for key in ("description", "hierarchy_and_relation"):
+                    if key in layer:
+                        _check_non_empty_string(layer[key], f"{path}.{key}", errors)
+                if "coordinates" in layer:
+                    _check_coordinates(
+                        layer["coordinates"], f"{path}.coordinates", errors
+                    )
+                if "color_specs" in layer:
+                    _check_color_specs(layer["color_specs"], f"{path}.color_specs", errors)
+    return errors
+
+
+def _chat_completion(
+    base_url: str,
+    model: str,
+    messages: List[Dict[str, str]],
+    *,
+    temperature: float,
+    max_tokens: int,
+    api_key: Optional[str],
+    timeout: float,
+) -> Tuple[str, Optional[str]]:
+    """POST one chat completion; return (content, finish_reason)."""
+    url = base_url.rstrip("/") + "/chat/completions"
+    payload = json.dumps(
+        {
+            "model": model,
+            "messages": messages,
+            "temperature": temperature,
+            "max_tokens": max_tokens,
+            "stream": False,
+        }
+    ).encode("utf-8")
+    headers = {"Content-Type": "application/json"}
+    if api_key:
+        headers["Authorization"] = f"Bearer {api_key}"
+    request = urllib.request.Request(url, data=payload, headers=headers, method="POST")
+    try:
+        with urllib.request.urlopen(request, timeout=timeout) as response:
+            body = response.read().decode("utf-8", errors="replace")
+    except urllib.error.HTTPError as error:
+        detail = error.read().decode("utf-8", errors="replace")
+        raise PromptEnhancementError(
+            f"HTTP {error.code} from {url}: {detail[:2000]}"
+        ) from error
+    except urllib.error.URLError as error:
+        raise PromptEnhancementError(f"cannot reach {url}: {error.reason}") from error
+    except OSError as error:  # includes socket timeouts during the read
+        raise PromptEnhancementError(f"request to {url} failed: {error}") from error
+
+    try:
+        envelope = json.loads(body)
+        choice = envelope["choices"][0]
+        content = choice["message"]["content"]
+    except (json.JSONDecodeError, KeyError, IndexError, TypeError) as error:
+        raise PromptEnhancementError(
+            f"malformed chat completion response from {url}: {body[:500]}"
+        ) from error
+    finish_reason = choice.get("finish_reason")
+    if not isinstance(content, str) or not content.strip():
+        raise PromptEnhancementError(
+            f"empty completion content from {url} (finish_reason={finish_reason!r})"
+        )
+    return content, finish_reason
+
+
+def enhance(
+    caption: str,
+    base_url: str,
+    model: str,
+    api_key: Optional[str] = None,
+    timeout: float = DEFAULT_TIMEOUT_SECONDS,
+    temperature: float = DEFAULT_TEMPERATURE,
+    max_tokens: int = DEFAULT_MAX_TOKENS,
+) -> Dict[str, Any]:
+    """Return the validated structured rewrite of ``caption``.
+
+    Sends the verbatim rewriter system prompt plus the caption to
+    ``{base_url}/chat/completions``. On a parse or schema failure, retries
+    exactly once with the validation errors appended to the user turn; if
+    that also fails, raises PromptEnhancementError carrying the errors.
+    """
+    system_prompt = load_system_prompt()
+    messages = [
+        {"role": "system", "content": system_prompt},
+        {"role": "user", "content": caption},
+    ]
+    request_kwargs = {
+        "temperature": temperature,
+        "max_tokens": max_tokens,
+        "api_key": api_key,
+        "timeout": timeout,
+    }
+    errors: List[str] = []
+    content = ""
+    for attempt in (1, 2):
+        content, finish_reason = _chat_completion(
+            base_url, model, messages, **request_kwargs
+        )
+        document: Optional[Dict[str, Any]] = None
+        try:
+            document = extract_json_object(content)
+        except ValueError as error:
+            errors = [str(error)]
+        if document is not None:
+            errors = validate_enhanced_prompt(document)
+        if not errors:
+            assert document is not None  # errors empty implies extraction succeeded
+            return document
+        if finish_reason == "length":
+            errors.append(
+                "the reply was cut off (finish_reason='length'): the complete "
+                f"JSON object must fit within max_tokens={max_tokens}"
+            )
+        print(f"pe_ling: attempt {attempt}/2 failed validation:", file=sys.stderr)
+        for error in errors:
+            print(f"pe_ling:   - {error}", file=sys.stderr)
+        if attempt == 1:
+            retry_content = (
+                f"{caption}\n\n"
+                "Your previous reply failed schema validation:\n"
+                + "".join(f"- {error}\n" for error in errors)
+                + "\n"
+                + REPAIR_INSTRUCTION
+            )
+            messages = [
+                {"role": "system", "content": system_prompt},
+                {"role": "user", "content": retry_content},
+            ]
+    raise PromptEnhancementError(
+        "prompt enhancement failed schema validation after 2 attempts:\n"
+        + "".join(f"  - {error}\n" for error in errors).rstrip(),
+        errors=errors,
+        reply=content,
+    )
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser(
+        description=(
+            "Enhance a Ming-Image text-to-image caption into the validated "
+            "structured JSON prompt via an OpenAI-compatible Ling-3.0-flash-VL "
+            "endpoint."
+        )
+    )
+    parser.add_argument("caption", help="free-form design caption to enhance")
+    parser.add_argument(
+        "--out",
+        type=Path,
+        help="write the validated JSON here (default: stdout, summary on stderr)",
+    )
+    parser.add_argument(
+        "--base-url",
+        default=DEFAULT_BASE_URL,
+        help=f"OpenAI-compatible base URL (default: {DEFAULT_BASE_URL})",
+    )
+    parser.add_argument(
+        "--model",
+        default=DEFAULT_MODEL,
+        help=f"chat model id served at the endpoint (default: {DEFAULT_MODEL})",
+    )
+    parser.add_argument(
+        "--api-key",
+        default=os.environ.get(API_KEY_ENV),
+        help=f"Bearer token for gated endpoints; defaults to ${API_KEY_ENV} when set",
+    )
+    parser.add_argument(
+        "--timeout",
+        type=float,
+        default=DEFAULT_TIMEOUT_SECONDS,
+        help=f"per-request timeout in seconds (default: {DEFAULT_TIMEOUT_SECONDS})",
+    )
+    parser.add_argument(
+        "--temperature",
+        type=float,
+        default=DEFAULT_TEMPERATURE,
+        help=f"sampling temperature (default: {DEFAULT_TEMPERATURE})",
+    )
+    parser.add_argument(
+        "--max-tokens",
+        type=int,
+        default=DEFAULT_MAX_TOKENS,
+        help=f"completion token budget (default: {DEFAULT_MAX_TOKENS})",
+    )
+    args = parser.parse_args()
+
+    started = time.perf_counter()
+    try:
+        document = enhance(
+            args.caption,
+            args.base_url,
+            args.model,
+            api_key=args.api_key,
+            timeout=args.timeout,
+            temperature=args.temperature,
+            max_tokens=args.max_tokens,
+        )
+    except PromptEnhancementError as error:
+        print(f"pe_ling: {error}", file=sys.stderr)
+        raise SystemExit(1)
+    elapsed = time.perf_counter() - started
+    layer_count = len(document["layers"])
+    payload = json.dumps(document, indent=2, ensure_ascii=False) + "\n"
+    if args.out is not None:
+        args.out.parent.mkdir(parents=True, exist_ok=True)
+        args.out.write_text(payload, encoding="utf-8")
+        print(f"pe_ling: {elapsed:.1f}s, {layer_count} layer(s) -> {args.out}")
+    else:
+        sys.stdout.write(payload)
+        print(f"pe_ling: {elapsed:.1f}s, {layer_count} layer(s)", file=sys.stderr)
+
+
+if __name__ == "__main__":
+    main()

code/processing_bailingmm2.py

@@ -0,0 +1,512 @@
+# coding=utf-8
+# Copyright 2024 The HuggingFace Inc. team.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Processor class for BailingMM2."""
+
+import sys
+from typing import List, Union, Dict, Optional
+
+import torch
+import PIL
+from PIL import Image
+
+if sys.version_info >= (3, 11):
+    from typing import Unpack
+else:
+    from typing_extensions import Unpack
+
+from transformers.feature_extraction_utils import BatchFeature
+from transformers.image_utils import ImageInput
+from transformers.processing_utils import (
+    ProcessingKwargs,
+    ProcessorMixin,
+)
+from transformers.tokenization_utils_base import PreTokenizedInput, TextInput
+
+from bailingmm_utils import process_vision_info, VideoInput, process_ratio, process_reference_vision_info, get_default_image_gen_hw
+import torchvision
+import math
+
+DEFAULT_IMAGE_PATCH_TOKEN = "<imagePatch>"
+DEFAULT_IM_START_TOKEN = "<image>"
+DEFAULT_IM_END_TOKEN = "</image>"
+DEFAULT_VID_START_TOKEN = "<video>"
+DEFAULT_VID_END_TOKEN = "</video>"
+DEFAULT_GEN_IMAGE_PATCH_TOKEN = "<gen_imagePatch>"
+DEFAULT_GEN_IM_START_TOKEN = "<gen_image>"
+DEFAULT_GEN_IM_END_TOKEN = "</gen_image>"
+PLACEHOLDER_IMAGE_TOKEN_IN_TEXT = "<imageHere>"
+DEFAULT_END_OF_CHUNK_TOKEN = "<end_of_chunk>"
+
+DEFAULT_FRAME_PATCH_TOKEN = "<framePatch>"
+DEFAULT_TEXT_TOKEN = '<text>'
+DEFAULT_ASR_TOKEN = '<asr>'
+DEFAULT_TTS_TOKEN = '<tts>'
+
+USER_PREFIX = "<role>HUMAN</role>"
+ASSISTANT_PREFIX = "<role>ASSISTANT</role>"
+
+SYSTEM_PROMPT_LINGV2_FLASH_NOTHINK = "<role>SYSTEM</role>你是一个友好的AI助手。\n\ndetailed thinking off"
+SYSTEM_PROMPT_LINGV2_FLASH_THINK = "<role>SYSTEM</role>你是一个友好的AI助手。\n\ndetailed thinking on"
+
+def check_single_quotes(s):
+    count = s.count("'")
+    if count % 2 != 0:
+        return False
+
+    positions = [i for i, char in enumerate(s) if char == "'"]
+    for i in range(0, len(positions), 2):
+        start = positions[i]
+        end = positions[i+1]
+        substr = s[start+1:end]
+        chinese_count = 0
+        for char in substr:
+            if '\u4e00' <= char <= '\u9fff':
+                chinese_count += 1
+        other_count = len(substr) - chinese_count
+        total = 3 * chinese_count + other_count
+        if total >= 20:
+            return False
+    return True
+
+def get_text_from_prompt(prompt):
+    if "'" in prompt and check_single_quotes(prompt):
+        prompt = prompt.replace("'", '"')
+
+    patterns = [r'\"(.*?)\"', r'‘(.*?)’', r'“(.*?)”']
+
+        
+    import re
+    texts = []
+    patterns = [r'\"(.*?)\"', r'‘(.*?)’', r'“(.*?)”']
+    for pattern in patterns:
+        texts.extend(re.findall(pattern, prompt))
+
+    if len(texts) == 1:
+        assert texts[0] in prompt
+        is_remove = False
+        remove_keywords = ["remove", "delete", "erase"]
+        text_start = min([j for j in [prompt.find(i) for i in ['"', '‘', '“']] if j >= 0])
+        for kw in remove_keywords:
+            if kw in prompt.lower():
+                if prompt.lower().find(kw) < text_start:
+                    is_remove = True
+                    break
+        
+        if is_remove:
+            texts = []
+    
+    text = " ".join(texts[-1:])
+    if len(text) > 0:
+        text = f'Text "{text}"'
+        text += ". "
+
+    return text
+
+def crop_to_aspect_max(img: Image.Image, target_ratio: float) -> Image.Image:
+    """
+    Center-crop a PIL.Image to the largest area fitting the target aspect ratio
+    (width/height) without resizing. Uses torchvision CenterCrop.
+
+    Args:
+        img: PIL.Image.Image input image
+        target_ratio: float target aspect ratio (width/height), must be positive
+
+    Returns:
+        the center-cropped PIL.Image.Image
+    """
+    if not isinstance(img, Image.Image):
+        raise TypeError("img must be a PIL.Image.Image")
+    if not math.isfinite(target_ratio) or target_ratio <= 0:
+        raise ValueError("target_ratio must be a positive, finite number")
+
+    W, H = img.size
+    if W <= 0 or H <= 0:
+        raise ValueError("image size is invalid")
+
+    orig_ratio = W / H
+
+    if orig_ratio >= target_ratio:
+        # image is wider than the target: use full height, crop left/right
+        new_h = H
+        new_w = int(math.floor(target_ratio * H))
+        new_w = max(1, min(new_w, W))  # guard against extreme ratios producing invalid sizes
+    else:
+        # image is narrower than the target: use full width, crop top/bottom
+        new_w = W
+        new_h = int(math.floor(W / target_ratio))
+        new_h = max(1, min(new_h, H))
+
+    crop = torchvision.transforms.CenterCrop((new_h, new_w))  # size is (h, w)
+    return crop(img)
+
+def transform_reference_images(
+    images,
+    image_gen_aspect_ratio=None,
+    image_gen_resolution=512,
+    image_gen_input_channels=None,
+):
+    if image_gen_input_channels not in (3, 4):
+        raise ValueError(
+            "image_gen_input_channels must be explicitly set to 3 or 4 "
+            "from the checkpoint capability contract"
+        )
+    image_mode = "RGB" if image_gen_input_channels == 3 else "RGBA"
+    images = [image.convert(image_mode) for image in images]
+
+    ref_pil = images[0]
+    if image_gen_aspect_ratio is not None:
+        ref_pil = crop_to_aspect_max(ref_pil, image_gen_aspect_ratio)
+
+    ori_h = ref_pil.size[1]
+    ori_w = ref_pil.size[0]
+    closest_size, _ = process_ratio(ori_h=ori_h, ori_w=ori_w, highres=image_gen_resolution)
+    
+    ref_pils = [torchvision.transforms.functional.resize(i, closest_size, interpolation=torchvision.transforms.InterpolationMode.BILINEAR) for i in images]
+    
+    ref_tensor = torch.cat([
+        ((torchvision.transforms.functional.to_tensor(i) - 0.5) * 2.0).unsqueeze(0)
+        for i in ref_pils
+    ], dim=0)
+
+    return ref_tensor, ref_pil.size[1], ref_pil.size[0]
+
+class BailingMM2ProcessorKwargs(ProcessingKwargs, total=False):
+    # see processing_utils.ProcessingKwargs documentation for usage.
+    _defaults = {
+        "text_kwargs": {"padding": True, "padding_side": "right"},
+        "image_kwargs": {},
+        "video_kwargs": {},
+    }
+
+class BailingMM2Processor(ProcessorMixin):
+    r"""
+    Constructs a BailingMM2 processor which wraps a bailingmm2 image processor, bailing audio processor and a LLaMa tokenizer into a single processor.
+    Args:
+        image_processor ([`BailingMM2ImageProcessor`], *optional*):
+            The image processor is a required input.
+        tokenizer ([`LlamaTokenizerFast`], *optional*):
+            The tokenizer is a required input.
+        chat_template (`str`, *optional*): A Jinja template which will be used to convert lists of messages
+            in a chat into a tokenizable string.
+        image_token (`str`, *optional*, defaults to `"<image>"`):
+            Special token used to denote image location.
+        video_token (`str`, *optional*, defaults to `"<video>"`):
+            Special token used to denote video location.
+    """
+
+    attributes = ["image_processor", "tokenizer"]
+    optional_attributes = ["chat_template"]
+
+    image_processor_class = "AutoImageProcessor"
+    tokenizer_class = "AutoTokenizer"
+
+    valid_kwargs = [
+        "chat_template",
+        "num_image_tokens",
+        "image_token",
+        "video_token",
+    ]
+
+    def __init__(
+        self,
+        image_processor=None,
+        tokenizer=None,
+        chat_template=None,
+        image_token="<image>",
+        video_token="<video>",
+        **kwargs: Unpack[BailingMM2ProcessorKwargs],
+    ):
+        self.image_token = image_token
+        self.video_token = video_token
+        if chat_template is None:
+            chat_template = tokenizer.chat_template
+
+        self.gen_terminator = [tokenizer.eos_token_id]
+        super().__init__(image_processor, tokenizer, chat_template=chat_template)
+
+    def __call__(
+        self,
+        images: ImageInput = None,
+        videos: VideoInput = None,
+        text: Union[TextInput, PreTokenizedInput, List[TextInput], List[PreTokenizedInput]] = None,
+        image_gen_highres = 512,
+        image_gen_aspect_ratio = None,
+        image_gen_ref_images: Union["PIL.Image.Image", list["PIL.Image.Image"]] = None,
+        image_gen_input_channels = None,
+        **kwargs,
+    ) -> BatchFeature:
+        """
+        Main method to prepare for the model one or several sequences(s) and image(s). This method forwards the `text`
+        and `kwargs` arguments to LlamaTokenizerFast's [`~LlamaTokenizerFast.__call__`] if `text` is not `None` to encode
+        the text. To prepare the image(s), this method forwards the `images` and `kwrags` arguments to
+        LlavaNextImageProcessor's [`~LlavaNextImageProcessor.__call__`] if `images` is not `None`. Please refer to the doctsring
+        of the above two methods for more information.
+
+        Args:
+            images (`PIL.Image.Image`, `np.ndarray`, `torch.Tensor`, `List[PIL.Image.Image]`, `List[np.ndarray]`, `List[torch.Tensor]`):
+                The image or batch of images to be prepared. Each image can be a PIL image, NumPy array or torch Tensor.
+                tensor. Both channels-first and channels-last formats are supported.
+            videos (`np.ndarray`, `torch.Tensor`, `List[np.ndarray]`, `List[torch.Tensor]`):
+                The image or batch of videos to be prepared. Each video can be a 4D NumPy array or torch Tensor.
+            audios (`Tuple[torch.Tensor, int]`, `List[Tuple[torch.Tensor, int]]`):
+                The sequence or batch of audios to be prepared. Each audio can be a 1D torch Tensor (with its sampling rate).
+            text (`str`, `List[str]`, `List[List[str]]`):
+                The sequence or batch of sequences to be encoded. Each sequence can be a string or a list of strings
+                (pretokenized string). If the sequences are provided as a list of strings (pretokenized), you must set
+                `is_split_into_words=True` (to lift the ambiguity with a batch of sequences).
+
+        Returns:
+            [`BatchFeature`]: A [`BatchFeature`] with the following fields:
+
+            - **input_ids** -- List of token ids to be fed to a model. Returned when `text` is not `None`.
+            - **attention_mask** -- List of indices specifying which tokens should be attended to by the model (when
+              `return_attention_mask=True` or if *"attention_mask"* is in `self.model_input_names` and if `text` is not
+              `None`).
+            - **pixel_values** -- Pixel values to be fed to a model. Returned when `images` is not `None`.
+            - **image_num_patches** -- Patch number to be fed to a model. Returned when `images` is not `None`.
+            - **image_sizes** -- Size of each image that will be used to unpad an image. Returned when `images` is not `None`.
+            - **pixel_values_videos** -- Pixel values of a video input to be fed to a model. Returned when `videos` is not `None`.
+            - **pixel_values_audios** -- Pixel values of an audio input to be fed to a model. Returned when `audios` is not `None`.
+
+        """
+        output_kwargs = self._merge_kwargs(
+            BailingMM2ProcessorKwargs,
+            tokenizer_init_kwargs=self.tokenizer.init_kwargs,
+            **kwargs,
+        )
+        if isinstance(text, str):
+            text = [text]
+        elif not isinstance(text, list) and not isinstance(text[0], str):
+            raise ValueError("Invalid input text. Please provide a string, or a list of strings")
+
+        image_inputs = {}
+        video_inputs = {}
+        image_gen_inputs = {}
+
+        text_in_text = [get_text_from_prompt(i) for i in text]
+
+
+        default_image_gen_height, default_image_gen_width = get_default_image_gen_hw(image_gen_highres, image_gen_aspect_ratio)
+        
+        image_gen_inputs.update({
+            "image_gen_text": text_in_text,
+            "image_gen_highres": image_gen_highres,
+            "image_gen_height": torch.LongTensor([default_image_gen_height] * len(text)),
+            "image_gen_width": torch.LongTensor([default_image_gen_width] * len(text)),
+        })
+
+        if images is not None:
+            image_inputs = self.image_processor(images=images, videos=None, **output_kwargs["images_kwargs"])
+            image_grid_thw = image_inputs["image_grid_thw"]
+            
+            text = self._expand_image_tokens(text, image_grid_thw)
+
+            # image_gen_pixel_values_reference, image_gen_height, image_gen_width = None, 512, 512
+            if image_gen_ref_images is not None:
+                if isinstance(image_gen_ref_images, PIL.Image.Image):
+                    image_gen_ref_images = [image_gen_ref_images]
+                elif not isinstance(image_gen_ref_images, list) and not isinstance(image_gen_ref_images[0], PIL.Image.Image):
+                    raise ValueError("Invalid input image_gen_ref_images. Please provide a PIL.Image.Image, or a list of PIL.Image.Image")
+
+                assert len(image_gen_ref_images) == len(text) # same batch_size 
+
+                image_gen_pixel_values_reference, image_gen_height_list, image_gen_width_list = transform_reference_images(
+                    image_gen_ref_images,
+                    image_gen_aspect_ratio,
+                    image_gen_highres,
+                    image_gen_input_channels,
+                )
+            
+                image_gen_inputs.update({
+                    "image_gen_pixel_values_reference": image_gen_pixel_values_reference,
+                    "image_gen_height": torch.LongTensor([image_gen_height_list] * len(text)),
+                    "image_gen_width": torch.LongTensor([image_gen_width_list] * len(text)),
+                    #"image_gen_height": torch.LongTensor([ori_h]),
+                    #"image_gen_width": torch.LongTensor([ori_w]),
+                })
+
+        if videos is not None:
+            video_inputs = self.image_processor(images=None, videos=videos, do_resize=False, **output_kwargs["videos_kwargs"])
+            video_grid_thw = video_inputs["video_grid_thw"]
+            text = self._expand_video_tokens(text, video_grid_thw)
+
+        # Padding side can be in TextKwargs but is not accepted by the tokenizer
+        _ = output_kwargs["text_kwargs"].pop("padding_side", None)
+        text_inputs = self.tokenizer(text, **output_kwargs["text_kwargs"])
+
+        return BatchFeature(data={**text_inputs, **image_inputs, **video_inputs, **image_gen_inputs})
+
+    def apply_system_template(self, sys_prompt_exp=None, use_cot_system_prompt=False):
+        if use_cot_system_prompt:
+            sys_prompt = SYSTEM_PROMPT_LINGV2_FLASH_THINK
+        else:
+            sys_prompt = SYSTEM_PROMPT_LINGV2_FLASH_NOTHINK
+        if sys_prompt_exp is not None:
+            sys_prompt = sys_prompt.replace("你是一个友好的AI助手。", sys_prompt_exp)
+
+        return sys_prompt
+
+    def apply_chat_template(
+        self,
+        conversation: Union[List[Dict[str, str]]],
+        sys_prompt_exp: Optional[str] = None,
+        use_cot_system_prompt: Optional[bool] = False,
+        **kwargs,
+    ) -> str:
+        """
+        Similar to the `apply_chat_template` method on tokenizers, this method applies a Jinja template to input
+        conversations to turn them into a single tokenizable string.
+
+        Args:
+            conversation (`List[Dict, str, str]`):
+                The conversation to format.
+            sys_prompt_exp (`Optional[str]`, *optional*):
+                The system prompt. If not provided, the processor's sysyetm template is used.
+            **kwargs:
+                Additional keyword arguments
+        """
+        text = ""
+        sys_prompt = self.apply_system_template(sys_prompt_exp, use_cot_system_prompt)
+        text = sys_prompt + self.tokenizer.eos_token
+
+        for idx, message in enumerate(conversation):
+            assert message["role"] in ["HUMAN", "ASSISTANT"]
+            if idx == len(conversation) - 1:
+                assert message["role"] == "HUMAN"
+
+            if message["role"] == "HUMAN":
+                text += USER_PREFIX
+            elif message["role"] == "ASSISTANT":
+                text += ASSISTANT_PREFIX
+
+            image_counts = str(message["content"]).count("<image>")
+            video_counts = str(message["content"]).count("<video>")
+
+            for content in message["content"]:
+                if content["type"] == "image":
+                    num_images = 1 if isinstance(content["image"], (str, Image.Image)) else len(content["image"])
+                    if image_counts < num_images:
+                        image_placeholder = "<IMAGE>\n" * (num_images - image_counts)
+                        text += image_placeholder.rstrip("\n")
+                # only one video supported now
+                elif content["type"] == "video":
+                    assert video_counts <= 1, "Video count must be at most 1!"
+                    if video_counts == 0:
+                        text += "<VIDEO>"
+                elif content["type"] == "audio":
+                    raise ValueError("audio input is not supported by Ming Image inference")
+                elif content["type"] == "text":
+                    text += content['text']
+            text += self.tokenizer.eos_token
+        text += ASSISTANT_PREFIX
+        
+        return text
+
+    def process_vision_info(
+        self,
+        conversations,
+    ):
+        return process_vision_info(conversations)
+
+    def process_reference_vision_info(
+        self,
+        conversations,
+    ):
+        return process_reference_vision_info(conversations)
+
+    def _expand_image_tokens(
+        self,
+        text: List[TextInput],
+        image_grid_thw: Union[List[int], int],
+        special_token: str = "<IMAGE>",
+    ):
+        prompt_strings = []
+        image_index = 0
+        num_query_token = torch.prod(image_grid_thw, dim=1) // 4
+        for sample in text:
+            num_images = sample.count(special_token)
+            if num_images > 0:
+                for i in range(image_index, num_images + image_index):
+                    img_text = DEFAULT_IM_START_TOKEN + num_query_token[i] * DEFAULT_IMAGE_PATCH_TOKEN + DEFAULT_IM_END_TOKEN + "\n"
+                    sample = sample.replace(special_token, img_text, 1)
+            image_index += num_images
+            prompt_strings.append(sample)
+        text = [sample for sample in prompt_strings]
+        return text
+
+    def _expand_video_tokens(
+        self,
+        text: List[TextInput],
+        video_grid_thw: Union[List[int], int],
+        special_token: str = "<VIDEO>",
+    ):
+        prompt_strings = []
+        video_index = 0
+        num_query_token = torch.prod(video_grid_thw, dim=1) // 4
+        for sample in text:
+            num_videos = sample.count(special_token)
+            if num_videos > 0:
+                for i in range(video_index, num_videos + video_index):
+                    video_text = num_query_token[i] * DEFAULT_FRAME_PATCH_TOKEN
+                    video_text = DEFAULT_VID_START_TOKEN + video_text + DEFAULT_VID_END_TOKEN + "\n"
+                    sample = sample.replace(special_token, video_text, 1)
+            video_index += num_videos
+            prompt_strings.append(sample)
+        text = [sample for sample in prompt_strings]
+        return text
+
+    # Copied from transformers.models.clip.processing_clip.CLIPProcessor.batch_decode with CLIP->Llama
+    def batch_decode(self, *args, **kwargs):
+        """
+        This method forwards all its arguments to LlamaTokenizerFast's [`~PreTrainedTokenizer.batch_decode`]. Please
+        refer to the docstring of this method for more information.
+        """
+        return self.tokenizer.batch_decode(*args, **kwargs)
+
+    # Copied from transformers.models.clip.processing_clip.CLIPProcessor.decode with CLIP->Llama
+    def decode(self, *args, **kwargs):
+        """
+        This method forwards all its arguments to LlamaTokenizerFast's [`~PreTrainedTokenizer.decode`]. Please refer to
+        the docstring of this method for more information.
+        """
+        return self.tokenizer.decode(*args, **kwargs)
+
+    @property
+    def model_input_names(self):
+        tokenizer_input_names = self.tokenizer.model_input_names
+        image_processor_input_names = self.image_processor.model_input_names
+        return list(
+            dict.fromkeys(
+                tokenizer_input_names + image_processor_input_names))
+
+
+def load_bailingmm2_processor(data_directory):
+    """Build the processor from data files with this repository's classes.
+
+    Checkpoint component directories (``mllm/``) carry only data files
+    (``preprocessor_config.json``, ``tokenizer_config.json``,
+    ``special_tokens_map.json``, ``tokenizer.json``). Loading them through
+    ``AutoProcessor`` with ``trust_remote_code=True`` fails because
+    Transformers requires the Python implementation inside the loaded
+    directory, while the implementation intentionally lives only in this
+    repository. Construct the components explicitly instead.
+    """
+    from image_processing_bailingmm2 import BailingMM2ImageProcessor
+    from tokenization_bailing import BailingTokenizer
+
+    data_directory = str(data_directory)
+    tokenizer = BailingTokenizer.from_pretrained(data_directory)
+    image_processor = BailingMM2ImageProcessor.from_pretrained(data_directory)
+    return BailingMM2Processor(image_processor=image_processor, tokenizer=tokenizer)

code/quant/__init__.py

@@ -0,0 +1 @@
+"""Weight-only INT8 for the Ming-Image MLLM: quantize_stream.py writes it, load_int8.py loads it."""

code/quant/int8_linear.py

@@ -0,0 +1,221 @@
+"""Weight-only symmetric per-output-channel INT8 linear.
+
+Scales stay float32 across dtype casts. ``module.to(dtype=torch.bfloat16)``
+(and ``.bfloat16()`` / ``.half()`` / ``.to(device, dtype)``) must not touch them;
+device moves still do. The int8 weight codes are likewise dtype-stable.
+"""
+
+from __future__ import annotations
+
+import torch
+import torch.nn.functional as F
+from torch import nn
+
+# Leaf names of Linear modules whose 2-D weights are quantized.
+# Exact match: the routers are `gate` / `image_gate` / `audio_gate`, NOT `gate_proj`.
+QUANT_LEAVES = frozenset(
+    {"query_key_value", "dense", "gate_proj", "up_proj", "down_proj"}
+)
+
+QUANT_RULE = (
+    "Quantize ONLY 2-D .weight tensors under model.model.layers. whose owning "
+    "module's leaf name is exactly one of query_key_value, dense, gate_proj, "
+    "up_proj, down_proj. Everything else stays byte-identical BF16: embeddings, "
+    "lm_head, all norms, the vision tower, linear_proj, and the three routers "
+    "(modules named gate, image_gate, audio_gate — leaf match, not a substring). "
+    "Per-output-channel symmetric: scale = absmax/127, "
+    "q = clamp(round(w/scale), -127, 127). All-zero rows: scale 1.0, q 0."
+)
+
+# Real checkpoint keys look like `model.model.layers.N...`. A bare
+# `layers.N...` name is the same stack with the root prefix omitted (tests).
+_DECODER_LAYER_PREFIXES = ((), ("model", "model"))
+
+
+def _weight_leaf(tensor_name: str) -> str | None:
+    """Owning module's leaf name if `tensor_name` ends in `.weight`, else None."""
+    if not isinstance(tensor_name, str) or not tensor_name.endswith(".weight"):
+        return None
+    module = tensor_name[: -len(".weight")]
+    if not module:
+        return None
+    return module.rsplit(".", 1)[-1]
+
+
+def _under_decoder_layers(tensor_name: str) -> bool:
+    """True when the tensor lives under the MLLM decoder `model.model.layers` stack.
+
+    `layers` must be its own path component, followed by a layer index. The
+    components before it must be empty or end in `model.model` — so a vision
+    tower that happens to contain the substring "layers" is not selected, and
+    `gate` is never selected just because `gate_proj` contains those letters.
+    """
+    parts = tensor_name.split(".")
+    for i, part in enumerate(parts):
+        if part != "layers":
+            continue
+        if i + 1 >= len(parts) or not parts[i + 1].isdigit():
+            continue
+        prefix = tuple(parts[:i])
+        if prefix in _DECODER_LAYER_PREFIXES:
+            return True
+        if len(prefix) >= 2 and prefix[-2:] == ("model", "model"):
+            return True
+    return False
+
+
+def quant_rule_leaf(tensor_name: str) -> str | None:
+    """Leaf name if the name matches the quantize rule, ignoring rank.
+
+    Returns None when the tensor is not a candidate. A candidate whose rank is
+    not 2 is a hard error for the stream (see quantize_stream); ``is_quantizable``
+    itself returns False for that case.
+    """
+    leaf = _weight_leaf(tensor_name)
+    if leaf not in QUANT_LEAVES:
+        return None
+    if not _under_decoder_layers(tensor_name):
+        return None
+    return leaf
+
+
+def is_quantizable(tensor_name: str, shape) -> bool:
+    """True only for 2-D quantize-rule weights. See ``QUANT_RULE``."""
+    if quant_rule_leaf(tensor_name) is None:
+        return False
+    try:
+        rank = len(shape)
+    except TypeError:
+        return False
+    return rank == 2
+
+
+def quantize_weight(weight: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
+    """Per-output-channel symmetric int8.
+
+    ``scale = absmax(row) / 127``, ``q = clamp(round(w / scale), -127, 127)``.
+    An all-zero row gets scale 1.0 and q 0 (no div-by-zero, no NaN/Inf).
+    """
+    if weight.ndim != 2:
+        raise ValueError(
+            f"quantize_weight expects a 2-D weight, got shape {tuple(weight.shape)}"
+        )
+    wf = weight.detach().to(dtype=torch.float32)
+    absmax = wf.abs().amax(dim=1)
+    scale = absmax / 127.0
+    zero = scale == 0
+    # All-zero rows would divide by 0. Force scale 1 and q 0 instead of NaN.
+    scale = torch.where(zero, torch.ones_like(scale), scale)
+    q = torch.round(wf / scale[:, None]).clamp(-127, 127).to(dtype=torch.int8)
+    q = torch.where(zero[:, None], torch.zeros_like(q), q)
+    return q.contiguous(), scale.to(dtype=torch.float32).contiguous()
+
+
+def _scale_name(weight_name: str) -> str:
+    if not weight_name.endswith(".weight"):
+        raise ValueError(f"not a weight tensor name: {weight_name}")
+    return weight_name[: -len("weight")] + "scale"
+
+
+class Int8Linear(nn.Module):
+    """``F.linear`` on a weight dequantized from int8 + per-row float32 scale.
+
+    ``weight`` is int8 ``[out, in]``, ``scale`` is float32 ``[out]``, ``bias``
+    (optional) keeps the source dtype. All three are buffers.
+    """
+
+    def __init__(self, weight: torch.Tensor, scale: torch.Tensor, bias: torch.Tensor | None):
+        super().__init__()
+        if weight.dtype != torch.int8 or weight.ndim != 2:
+            raise ValueError(
+                f"weight must be int8 [out, in], got dtype={weight.dtype} shape={tuple(weight.shape)}"
+            )
+        if scale.dtype != torch.float32 or tuple(scale.shape) != (weight.shape[0],):
+            raise ValueError(
+                f"scale must be float32 [{weight.shape[0]}], got dtype={scale.dtype} shape={tuple(scale.shape)}"
+            )
+        if bias is not None:
+            if bias.ndim != 1 or bias.shape[0] != weight.shape[0]:
+                raise ValueError(
+                    f"bias must be [{weight.shape[0]}], got shape={tuple(bias.shape)}"
+                )
+        self.in_features = int(weight.shape[1])
+        self.out_features = int(weight.shape[0])
+        self.register_buffer("weight", weight)
+        self.register_buffer("scale", scale)
+        self.register_buffer("bias", bias)
+
+    def _apply(self, fn, *args, **kwargs):
+        # Pull dtype-stable buffers out before Module._apply. Putting them back
+        # with only a device move (never fn's dtype cast) keeps scale float32
+        # and weight int8. Bias is left in the dict so it follows the cast.
+        saved: dict[str, torch.Tensor] = {}
+        for name in ("weight", "scale"):
+            buf = self._buffers.get(name, None)
+            if buf is not None:
+                saved[name] = buf
+                self._buffers[name] = None
+        try:
+            out = super()._apply(fn, *args, **kwargs)
+        finally:
+            for name, buf in saved.items():
+                self._buffers[name] = _move_device_keep_dtype(buf, fn)
+        return out
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        # One dequant in fp32, one cast to the activation dtype, then linear.
+        w = (self.weight.float() * self.scale[:, None]).to(dtype=x.dtype)
+        return F.linear(x, w, self.bias)
+
+    @classmethod
+    def from_linear(cls, linear: nn.Linear) -> "Int8Linear":
+        if not isinstance(linear, nn.Linear):
+            raise TypeError(f"from_linear expects nn.Linear, got {type(linear).__name__}")
+        q, scale = quantize_weight(linear.weight.data)
+        if linear.bias is None:
+            bias = None
+        else:
+            bias = linear.bias.detach().clone()
+        return cls(q, scale, bias)
+
+    @classmethod
+    def shell(
+        cls,
+        in_features: int,
+        out_features: int,
+        bias: bool,
+        bias_dtype: torch.dtype,
+        device,
+    ) -> "Int8Linear":
+        """Empty buffers (for ``meta``). Does not read or write any weight values."""
+        dev = torch.device(device) if not isinstance(device, torch.device) else device
+        weight = torch.empty((out_features, in_features), dtype=torch.int8, device=dev)
+        scale = torch.empty((out_features,), dtype=torch.float32, device=dev)
+        if bias:
+            bias_t: torch.Tensor | None = torch.empty(
+                (out_features,), dtype=bias_dtype, device=dev
+            )
+        else:
+            bias_t = None
+        return cls(weight, scale, bias_t)
+
+    def extra_repr(self) -> str:
+        return (
+            f"in_features={self.in_features}, out_features={self.out_features}, "
+            f"bias={self.bias is not None}"
+        )
+
+
+def _move_device_keep_dtype(buf: torch.Tensor, fn) -> torch.Tensor:
+    """Apply only the device change implied by ``fn``, preserving ``buf``'s dtype and values.
+
+    Probed with a 0-element tensor so a dtype cast cannot round the real scale.
+    """
+    try:
+        probe = torch.empty((), dtype=buf.dtype, device=buf.device)
+        moved = fn(probe)
+    except Exception:
+        return buf
+    if not torch.is_tensor(moved) or moved.device == buf.device:
+        return buf
+    return buf.to(device=moved.device)

code/quant/quantize_stream.py

@@ -0,0 +1,556 @@
+"""Stream a Ming MLLM directory to weight-only INT8 shards.
+
+Never builds the model: it buffers at most one output shard (<= 5 GB) of tensors at a time. Measured on
+the real 34.0 GB checkpoint (AMD Strix Halo, 2026-09-23): 266 s wall, peak RSS 17.8 GiB.
+CLI: ``python quantize_stream.py SRC_MLLM_DIR DST_DIR [--exclude MODULE_REGEX]`` (matching modules stay BF16).
+"""
+
+from __future__ import annotations
+
+import json
+import math
+import re
+import os
+import shutil
+import sys
+from dataclasses import dataclass
+from pathlib import Path
+
+import torch
+from safetensors import safe_open
+from safetensors.torch import save_file
+
+try:  # imported as the `quant` package
+    from .int8_linear import QUANT_RULE, quant_rule_leaf, quantize_weight
+except ImportError:  # run as a script: python quant/quantize_stream.py SRC DST
+    from int8_linear import QUANT_RULE, quant_rule_leaf, quantize_weight
+
+# Decimal GB, same unit Hugging Face uses for max_shard_size="5GB".
+MAX_SHARD_BYTES = 5 * 10**9
+
+_DTYPE_BYTES = {
+    "BOOL": 1,
+    "U8": 1,
+    "I8": 1,
+    "F8_E4M3": 1,
+    "F8_E5M2": 1,
+    "F8_E8M0": 1,
+    "U16": 2,
+    "I16": 2,
+    "F16": 2,
+    "BF16": 2,
+    "U32": 4,
+    "I32": 4,
+    "F32": 4,
+    "U64": 8,
+    "I64": 8,
+    "F64": 8,
+}
+
+INDEX_NAME = "model.safetensors.index.json"
+MANIFEST_NAME = "int8_manifest.json"
+
+
+class QuantizeError(Exception):
+    """User-facing checkpoint error. main() prints it and returns 1."""
+
+
+def _die(msg: str) -> None:
+    raise QuantizeError(msg)
+
+
+def _normalize_dtype(dtype_name) -> str:
+    text = str(dtype_name).upper()
+    if "." in text:
+        text = text.rsplit(".", 1)[-1]
+    aliases = {
+        "BFLOAT16": "BF16",
+        "FLOAT16": "F16",
+        "FLOAT32": "F32",
+        "FLOAT64": "F64",
+        "FLOAT8_E4M3FN": "F8_E4M3",
+        "FLOAT8_E5M2": "F8_E5M2",
+        "INT8": "I8",
+        "INT16": "I16",
+        "INT32": "I32",
+        "INT64": "I64",
+        "UINT8": "U8",
+    }
+    return aliases.get(text, text)
+
+
+def _dtype_nbytes(dtype_name: str) -> int:
+    try:
+        return _DTYPE_BYTES[dtype_name]
+    except KeyError:
+        _die(f"unsupported safetensors dtype {dtype_name!r}")
+        raise  # unreachable; satisfies type checkers
+
+
+def _numel(shape: tuple[int, ...]) -> int:
+    n = 1
+    for d in shape:
+        n *= int(d)
+    return n
+
+
+def _load_index(path: Path) -> dict:
+    if not path.is_file():
+        _die(f"missing index: {path}")
+
+    def _pairs(pairs):
+        keys = [k for k, _ in pairs]
+        dupes = sorted({k for k in keys if keys.count(k) > 1})
+        if dupes:
+            _die(f"duplicate key(s) in {path}: {dupes}")
+        return dict(pairs)
+
+    try:
+        raw = path.read_text(encoding="utf-8")
+        index = json.loads(raw, object_pairs_hook=_pairs)
+    except QuantizeError:
+        raise
+    except (OSError, json.JSONDecodeError) as exc:
+        _die(f"cannot read index {path}: {exc}")
+    if not isinstance(index, dict) or not isinstance(index.get("weight_map"), dict):
+        _die(f"index {path} has no weight_map object")
+    if not index["weight_map"]:
+        _die(f"index {path} weight_map is empty")
+    return index
+
+
+def _check_dst_clean(dst: Path) -> None:
+    if not dst.exists():
+        return
+    if not dst.is_dir():
+        _die(f"destination is not a directory: {dst}")
+    found = sorted(p.relative_to(dst).as_posix() for p in dst.rglob("*.safetensors"))
+    if found:
+        _die(f"destination already contains safetensors: {found}")
+
+
+def _reject_nested(src: Path, dst: Path) -> None:
+    src_r = src.resolve()
+    dst_r = dst.resolve()
+    if src_r == dst_r or src_r in dst_r.parents or dst_r in src_r.parents:
+        _die(f"SRC and DST must be distinct and not nested: {src} vs {dst}")
+
+
+def _shard_path(src: Path, shard_name: str) -> Path:
+    rel = Path(shard_name)
+    if rel.is_absolute() or ".." in rel.parts:
+        _die(f"unsafe shard path in index: {shard_name}")
+    path = src / rel
+    if not path.is_file():
+        _die(f"index lists missing shard: {shard_name}")
+    return path
+
+
+@dataclass
+class Item:
+    src_shard: str
+    name: str
+    kind: str  # "copy" or "quant"
+    shape: tuple[int, ...]
+    src_dtype: str
+    src_bytes: int
+    out_bytes: int
+    group: int = -1
+
+
+def _scale_name(weight_name: str) -> str:
+    return weight_name[: -len("weight")] + "scale"
+
+
+def _plan(src: Path, index: dict, exclude: str | None = None) -> list[Item]:
+    """Metadata-only pass. Reads shapes and dtypes, not tensor bodies."""
+    weight_map: dict[str, str] = index["weight_map"]
+    shard_order: list[str] = []
+    seen_shards: set[str] = set()
+    for shard in weight_map.values():
+        if shard not in seen_shards:
+            seen_shards.add(shard)
+            shard_order.append(shard)
+
+    index_names_by_shard: dict[str, set[str]] = {s: set() for s in shard_order}
+    for name, shard in weight_map.items():
+        if shard not in index_names_by_shard:
+            _die(f"weight_map value {shard!r} for {name} was not collected")
+        index_names_by_shard[shard].add(name)
+
+    items: list[Item] = []
+    seen_names: dict[str, str] = {}
+    for shard in shard_order:
+        path = _shard_path(src, shard)
+        with safe_open(str(path), framework="pt", device="cpu") as handle:
+            file_names = list(handle.keys())
+            file_set = set(file_names)
+            if len(file_set) != len(file_names):
+                _die(f"shard {shard} header lists a tensor name twice")
+            missing = sorted(index_names_by_shard[shard] - file_set)
+            extra = sorted(file_set - index_names_by_shard[shard])
+            if missing:
+                _die(f"index lists tensors missing from {shard}: {missing}")
+            if extra:
+                _die(f"{shard} contains tensors absent from the index: {extra}")
+            for name in file_names:
+                if name in seen_names:
+                    _die(
+                        f"tensor name appears twice: {name} "
+                        f"({seen_names[name]} and {shard})"
+                    )
+                seen_names[name] = shard
+                sl = handle.get_slice(name)
+                if not hasattr(sl, "get_dtype") or not hasattr(sl, "get_shape"):
+                    _die(
+                        "safetensors safe_open slice is missing get_shape/get_dtype; "
+                        "cannot plan shards without loading tensor bodies"
+                    )
+                shape = tuple(int(d) for d in sl.get_shape())
+                dtype_name = _normalize_dtype(sl.get_dtype())
+                src_bytes = _numel(shape) * _dtype_nbytes(dtype_name)
+                leaf = quant_rule_leaf(name)
+                if leaf is not None and exclude and re.search(exclude, name[: -len(".weight")]):
+                    leaf = None  # kept BF16 by --exclude
+                if leaf is not None and len(shape) != 2:
+                    _die(
+                        f"tensor {name} matches the quantize rule but is not 2-D "
+                        f"(shape={list(shape)}, dtype={dtype_name})"
+                    )
+                if leaf is not None:
+                    out_bytes = _numel(shape) * 1 + shape[0] * 4  # int8 weight + fp32 scale
+                    items.append(
+                        Item(shard, name, "quant", shape, dtype_name, src_bytes, out_bytes)
+                    )
+                else:
+                    items.append(
+                        Item(shard, name, "copy", shape, dtype_name, src_bytes, src_bytes)
+                    )
+
+    index_names = set(weight_map)
+    planned = {it.name for it in items}
+    if planned != index_names:
+        _die(
+            "index / shard mismatch after scan: "
+            f"only_in_index={sorted(index_names - planned)[:8]} "
+            f"only_in_shards={sorted(planned - index_names)[:8]}"
+        )
+
+    produced = set(planned)
+    for it in items:
+        if it.kind != "quant":
+            continue
+        sname = _scale_name(it.name)
+        if sname in produced:
+            _die(f"scale name collides with an existing tensor: {sname}")
+        produced.add(sname)
+    return items
+
+
+def _assign_groups(items: list[Item], max_shard_bytes: int) -> list[list[Item]]:
+    if max_shard_bytes <= 0:
+        _die(f"max_shard_bytes must be positive, got {max_shard_bytes}")
+    groups: list[list[Item]] = []
+    cur: list[Item] = []
+    cur_bytes = 0
+    for it in items:
+        if cur and cur_bytes + it.out_bytes > max_shard_bytes:
+            groups.append(cur)
+            cur = []
+            cur_bytes = 0
+        if cur_bytes == 0 and it.out_bytes > max_shard_bytes:
+            print(
+                f"warning: {it.name} contributes {it.out_bytes} bytes, "
+                f"over the {max_shard_bytes}-byte shard target; writing it alone",
+                file=sys.stderr,
+                flush=True,
+            )
+        it.group = len(groups)
+        cur.append(it)
+        cur_bytes += it.out_bytes
+    if cur:
+        groups.append(cur)
+    return groups
+
+
+def _relative_frobenius(weight: torch.Tensor, q: torch.Tensor, scale: torch.Tensor) -> float:
+    w = weight.detach().to(dtype=torch.float64)
+    deq = q.detach().to(dtype=torch.float64) * scale.detach().to(dtype=torch.float64)[:, None]
+    denom = torch.linalg.matrix_norm(w, ord="fro")
+    numer = torch.linalg.matrix_norm(w - deq, ord="fro")
+    d = denom.item()
+    n = numer.item()
+    if d == 0.0:
+        return 0.0 if n == 0.0 else math.inf
+    return n / d
+
+
+def _percentile_linear(values: list[float], pct: float) -> float:
+    """NumPy-style linear percentile. Empty → 0."""
+    if not values:
+        return 0.0
+    ordered = sorted(values)
+    if len(ordered) == 1:
+        return ordered[0]
+    rank = (len(ordered) - 1) * (pct / 100.0)
+    lo = math.floor(rank)
+    hi = math.ceil(rank)
+    if lo == hi:
+        return ordered[lo]
+    w = rank - lo
+    return ordered[lo] * (1.0 - w) + ordered[hi] * w
+
+
+def _copy_sidecars(src: Path, dst: Path) -> list[str]:
+    copied: list[str] = []
+    for dirpath, _dirnames, filenames in os.walk(src):
+        rel = Path(dirpath).relative_to(src)
+        out_dir = dst / rel
+        out_dir.mkdir(parents=True, exist_ok=True)
+        for filename in filenames:
+            if filename.endswith(".safetensors"):
+                continue
+            if filename == INDEX_NAME and rel == Path("."):
+                continue
+            src_file = Path(dirpath) / filename
+            dst_file = out_dir / filename
+            shutil.copy2(src_file, dst_file)
+            copied.append((rel / filename).as_posix())
+    return copied
+
+
+def _write_shards(
+    src: Path,
+    dst: Path,
+    items: list[Item],
+    groups: list[list[Item]],
+) -> tuple[dict[str, str], int, int, list[tuple[str, float]], list[Path]]:
+    n_out = len(groups)
+    weight_map: dict[str, str] = {}
+    bytes_in = 0
+    bytes_out = 0
+    errors: list[tuple[str, float]] = []
+    written: list[Path] = []
+
+    n_src = len({it.src_shard for it in items})
+    src_seen = 0
+    open_name: str | None = None
+    handle = None
+    buf: dict[str, torch.Tensor] = {}
+    buf_q = 0
+    buf_c = 0
+    current_group = 0
+
+    def flush() -> None:
+        nonlocal buf, buf_q, buf_c, current_group
+        if not buf:
+            return
+        fname = f"model-{current_group + 1:05d}-of-{n_out:05d}.safetensors"
+        path = dst / fname
+        for key, tensor in buf.items():
+            if not tensor.is_contiguous():
+                buf[key] = tensor.contiguous()
+        save_file(buf, str(path))
+        shard_bytes = 0
+        for key, tensor in buf.items():
+            weight_map[key] = fname
+            shard_bytes += tensor.numel() * tensor.element_size()
+        written.append(path)
+        print(
+            f"wrote {fname}: tensors={len(buf)} quantized={buf_q} copied={buf_c} "
+            f"bytes={shard_bytes}",
+            flush=True,
+        )
+        buf = {}
+        buf_q = 0
+        buf_c = 0
+        current_group += 1
+
+    try:
+        for it in items:
+            if it.src_shard != open_name:
+                if handle is not None:
+                    handle.__exit__(None, None, None)
+                    handle = None
+                path = _shard_path(src, it.src_shard)
+                handle = safe_open(str(path), framework="pt", device="cpu")
+                handle.__enter__()
+                open_name = it.src_shard
+                src_seen += 1
+                n_here = sum(1 for x in items if x.src_shard == it.src_shard)
+                print(
+                    f"reading source shard {src_seen}/{n_src} {it.src_shard} ({n_here} tensors)",
+                    flush=True,
+                )
+            assert handle is not None
+            tensor = handle.get_tensor(it.name)
+            got = tensor.numel() * tensor.element_size()
+            if got != it.src_bytes:
+                _die(
+                    f"{it.name} byte size {got} != planned {it.src_bytes} "
+                    f"(dtype={tensor.dtype}, shape={tuple(tensor.shape)})"
+                )
+            bytes_in += got
+            if it.kind == "quant":
+                if not tensor.is_floating_point():
+                    _die(
+                        f"{it.name} matches the quantize rule but dtype is {tensor.dtype}, "
+                        "expected a floating dtype"
+                    )
+                if tuple(tensor.shape) != it.shape:
+                    _die(f"{it.name} shape changed between passes: {tuple(tensor.shape)} vs {it.shape}")
+                q, scale = quantize_weight(tensor)
+                err = _relative_frobenius(tensor, q, scale)
+                if math.isnan(err) or math.isinf(err):
+                    _die(f"non-finite relative error for {it.name}: {err}")
+                errors.append((it.name, err))
+                del tensor
+                sname = _scale_name(it.name)
+                buf[it.name] = q
+                buf[sname] = scale
+                produced = q.numel() * q.element_size() + scale.numel() * scale.element_size()
+                if produced != it.out_bytes:
+                    _die(f"{it.name} output bytes {produced} != planned {it.out_bytes}")
+                buf_q += 1
+            else:
+                if not tensor.is_contiguous():
+                    tensor = tensor.contiguous()
+                buf[it.name] = tensor
+                buf_c += 1
+            bytes_out += it.out_bytes
+            # Flush when this item closes its planned output shard.
+            group_items = groups[it.group]
+            if it is group_items[-1]:
+                flush()
+    finally:
+        if handle is not None:
+            handle.__exit__(None, None, None)
+
+    if buf:
+        _die("internal error: output buffer not flushed")
+    if current_group != n_out:
+        _die(f"internal error: wrote {current_group} shards, planned {n_out}")
+    return weight_map, bytes_in, bytes_out, errors, written
+
+
+def _summary(
+    errors: list[tuple[str, float]],
+    n_quant: int,
+    n_copy: int,
+    bytes_in: int,
+    bytes_out: int,
+) -> dict:
+    vals = [e for _, e in errors]
+    if errors:
+        worst_name, worst_err = min(
+            errors,
+            key=lambda pair: (-pair[1], pair[0]),
+        )
+    else:
+        worst_name, worst_err = None, 0.0
+    mean = (sum(vals) / len(vals)) if vals else 0.0
+    return {
+        "tensors_quantized": n_quant,
+        "tensors_copied": n_copy,
+        "bytes_in": bytes_in,
+        "bytes_out": bytes_out,
+        "mean_relative_error": mean,
+        "p99_relative_error": _percentile_linear(vals, 99.0),
+        "max_relative_error": worst_err if vals else 0.0,
+        "worst_tensor": worst_name,
+    }
+
+
+def run(src: Path, dst: Path, max_shard_bytes: int = MAX_SHARD_BYTES, exclude: str | None = None) -> dict:
+    src = src.resolve()
+    dst = dst.resolve()
+    if not src.is_dir():
+        _die(f"SRC is not a directory: {src}")
+    _reject_nested(src, dst)
+    _check_dst_clean(dst)
+    index = _load_index(src / INDEX_NAME)
+    items = _plan(src, index, exclude)
+    groups = _assign_groups(items, max_shard_bytes)
+    dst.mkdir(parents=True, exist_ok=True)
+
+    written: list[Path] = []
+    try:
+        weight_map, bytes_in, bytes_out, errors, written = _write_shards(src, dst, items, groups)
+        copied = _copy_sidecars(src, dst)
+        n_quant = sum(1 for it in items if it.kind == "quant")
+        n_copy = sum(1 for it in items if it.kind == "copy")
+        measured = _summary(errors, n_quant, n_copy, bytes_in, bytes_out)
+        if measured["bytes_in"] != bytes_in or measured["bytes_out"] != bytes_out:
+            _die("internal error: summary byte counters diverged")
+        # Recompute the on-disk total from the tensors we recorded. weight_map
+        # values are what we just saved; bytes_out is that sum.
+        out_index = {"metadata": {"total_size": bytes_out}, "weight_map": weight_map}
+        (dst / INDEX_NAME).write_text(
+            json.dumps(out_index, indent=2) + "\n", encoding="utf-8"
+        )
+        modules = sorted(
+            it.name[: -len(".weight")] for it in items if it.kind == "quant"
+        )
+        manifest = {
+            "format": "ming-int8-wo-v1",
+            "scheme": "weight-only int8, per-output-channel symmetric, fp32 scales",
+            "rule": QUANT_RULE + (f" Additionally kept BF16: modules matching /{exclude}/." if exclude else ""),
+            "exclude": exclude,
+            "quantized_modules": modules,
+            "source_total_size": bytes_in,
+            "total_size": bytes_out,
+            "measured": measured,
+        }
+        (dst / MANIFEST_NAME).write_text(
+            json.dumps(manifest, indent=2, allow_nan=False) + "\n", encoding="utf-8"
+        )
+    except Exception:
+        for path in written:
+            try:
+                path.unlink()
+            except OSError:
+                pass
+        raise
+
+    print(f"copied {len(copied)} non-safetensors file(s)", flush=True)
+    m = measured
+    print(
+        "summary: "
+        f"quantized={m['tensors_quantized']} copied={m['tensors_copied']} "
+        f"bytes_in={m['bytes_in']} bytes_out={m['bytes_out']} "
+        f"mean_rel={m['mean_relative_error']:.8g} "
+        f"p99_rel={m['p99_relative_error']:.8g} "
+        f"max_rel={m['max_relative_error']:.8g} "
+        f"worst={m['worst_tensor']}",
+        flush=True,
+    )
+    return manifest
+
+
+def main(argv: list[str] | None = None) -> int:
+    args = list(sys.argv[1:] if argv is None else argv)
+    exclude = None
+    if "--exclude" in args:
+        i = args.index("--exclude")
+        if i + 1 >= len(args):
+            print("--exclude needs a regex", file=sys.stderr)
+            return 2
+        exclude = args[i + 1]
+        re.compile(exclude)
+        del args[i : i + 2]
+    if len(args) != 2:
+        print(
+            "usage: python quantize_stream.py SRC_MLLM_DIR DST_DIR [--exclude MODULE_REGEX]",
+            file=sys.stderr,
+        )
+        return 2
+    try:
+        run(Path(args[0]), Path(args[1]), max_shard_bytes=MAX_SHARD_BYTES, exclude=exclude)
+    except QuantizeError as exc:
+        print(f"error: {exc}", file=sys.stderr)
+        return 1
+    return 0
+
+
+if __name__ == "__main__":
+    sys.exit(main())

code/qwen2_5_vit.py

@@ -0,0 +1,508 @@
+# coding=utf-8
+# Copyright 2025 The Qwen Team and The HuggingFace Inc. team. All rights reserved.
+#
+# This code is based on EleutherAI's GPT-NeoX library and the GPT-NeoX
+# and OPT implementations in this library. It has been modified from its
+# original forms to accommodate minor architectural differences compared
+# to GPT-NeoX and OPT used by the Meta AI team that trained the model.
+#
+# 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.
+"""PyTorch Qwen2_5_ViT model."""
+
+import math
+
+import os
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from transformers.activations import ACT2FN
+from transformers.modeling_utils import PreTrainedModel
+from transformers.utils import (
+    is_flash_attn_2_available,
+    logging,
+)
+
+from typing import Union
+
+from transformers.configuration_utils import PretrainedConfig
+
+if is_flash_attn_2_available():
+    from flash_attn import flash_attn_varlen_func
+    from flash_attn.layers.rotary import apply_rotary_emb
+
+else:
+    flash_attn_varlen_func = None
+    apply_rotary_emb = None
+
+logger = logging.get_logger(__name__)
+
+class Qwen2_5_VLVisionConfig(PretrainedConfig):
+    model_type = "qwen2_5_vit"
+
+    def __init__(
+        self,
+        depth=32,
+        hidden_size=3584,
+        hidden_act="silu",
+        intermediate_size=3420,
+        num_heads=16,
+        in_channels=3,
+        patch_size=14,
+        spatial_merge_size=2,
+        temporal_patch_size=2,
+        tokens_per_second=4,
+        window_size=112,
+        out_hidden_size=3584,
+        fullatt_block_indexes=[7, 15, 23, 31],
+        _attn_implementation="flash_attention_2",
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+        self.depth = depth
+        self.hidden_size = hidden_size
+        self.hidden_act = hidden_act
+        self.intermediate_size = intermediate_size
+        self.num_heads = num_heads
+        self.in_channels = in_channels
+        self.patch_size = patch_size
+        self.spatial_merge_size = spatial_merge_size
+        self.temporal_patch_size = temporal_patch_size
+        self.tokens_per_second = tokens_per_second
+        self.window_size = window_size
+        self.fullatt_block_indexes = fullatt_block_indexes
+        self.out_hidden_size = out_hidden_size
+        self._attn_implementation = _attn_implementation
+
+    @classmethod
+    def from_pretrained(cls, pretrained_model_name_or_path: Union[str, os.PathLike], **kwargs) -> "PretrainedConfig":
+        cls._set_token_in_kwargs(kwargs)
+
+        config_dict, kwargs = cls.get_config_dict(pretrained_model_name_or_path, **kwargs)
+
+        if 'vision_config' in config_dict:
+            config_dict = config_dict['vision_config']
+
+        if "model_type" in config_dict and hasattr(cls, "model_type") and config_dict["model_type"] != cls.model_type:
+            logger.warning(
+                f"You are using a model of type {config_dict['model_type']} to instantiate a model of type "
+                f"{cls.model_type}. This is not supported for all configurations of models and can yield errors."
+            )
+
+        return cls.from_dict(config_dict, **kwargs)
+
+class Qwen2_5_VLMLP(nn.Module):
+    def __init__(self, config, bias: bool = False):
+        super().__init__()
+        self.hidden_size = config.hidden_size
+        self.intermediate_size = config.intermediate_size
+        self.gate_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=bias)
+        self.up_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=bias)
+        self.down_proj = nn.Linear(self.intermediate_size, self.hidden_size, bias=bias)
+        self.act_fn = ACT2FN[config.hidden_act]
+
+    def forward(self, hidden_state):
+        return self.down_proj(self.act_fn(self.gate_proj(hidden_state)) * self.up_proj(hidden_state))
+
+class Qwen2_5_VisionPatchEmbed(nn.Module):
+    def __init__(
+        self,
+        patch_size: int = 14,
+        temporal_patch_size: int = 2,
+        in_channels: int = 3,
+        embed_dim: int = 1152,
+    ) -> None:
+        super().__init__()
+        self.patch_size = patch_size
+        self.temporal_patch_size = temporal_patch_size
+        self.in_channels = in_channels
+        self.embed_dim = embed_dim
+
+        kernel_size = [temporal_patch_size, patch_size, patch_size]
+        self.proj = nn.Conv3d(in_channels, embed_dim, kernel_size=kernel_size, stride=kernel_size, bias=False)
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        target_dtype = self.proj.weight.dtype
+        hidden_states = hidden_states.view(
+            -1, self.in_channels, self.temporal_patch_size, self.patch_size, self.patch_size
+        )
+        hidden_states = self.proj(hidden_states.to(dtype=target_dtype)).view(-1, self.embed_dim)
+        return hidden_states
+
+class Qwen2_5_VisionRotaryEmbedding(nn.Module):
+    def __init__(self, dim: int, theta: float = 10000.0) -> None:
+        super().__init__()
+        inv_freq = 1.0 / (theta ** (torch.arange(0, dim, 2, dtype=torch.float) / dim))
+        self.register_buffer("inv_freq", inv_freq, persistent=False)
+        self.theta = theta
+        self.dim = dim
+
+    def forward(self, seqlen: int) -> torch.Tensor:
+        seq = torch.arange(seqlen, device=self.inv_freq.device, dtype=self.inv_freq.dtype)
+        freqs = torch.outer(seq, self.inv_freq)
+        return freqs
+
+    def reset_parameters(self):
+        # recompute inv_freq (for dynamic adjustment)
+        new_inv_freq = 1.0 / (self.theta ** (torch.arange(0, self.dim, 2, dtype=torch.float) / self.dim))
+        self.inv_freq.copy_(new_inv_freq)
+
+class Qwen2RMSNorm(nn.Module):
+    def __init__(self, hidden_size, eps=1e-6):
+        """
+        Qwen2RMSNorm is equivalent to T5LayerNorm.
+
+        Replaces transformer_engine.pytorch.RMSNorm: ROCm has no transformer-engine.
+        te.RMSNorm defaults (zero_centered_gamma=False) are standard RMSNorm, and the
+        checkpoint stores this affine as `weight`.
+        """
+        super().__init__()
+        self.weight = nn.Parameter(torch.ones(hidden_size))
+        self.variance_epsilon = eps
+
+    def forward(self, hidden_states):
+        input_dtype = hidden_states.dtype
+        hidden_states = hidden_states.to(torch.float32)
+        variance = hidden_states.pow(2).mean(-1, keepdim=True)
+        hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon)
+        return self.weight * hidden_states.to(input_dtype)
+
+    def extra_repr(self):
+        return f"{tuple(self.weight.shape)}, eps={self.variance_epsilon}"
+
+class Qwen2_5_VLPatchMerger(nn.Module):
+    def __init__(self, dim: int, context_dim: int, spatial_merge_size: int = 2) -> None:
+        super().__init__()
+        self.hidden_size = context_dim * (spatial_merge_size ** 2)
+        self.ln_q = Qwen2RMSNorm(context_dim, eps=1e-6)
+        self.mlp = nn.Sequential(
+            nn.Linear(self.hidden_size, self.hidden_size),
+            nn.GELU(),
+            nn.Linear(self.hidden_size, dim),
+        )
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        x = self.mlp(self.ln_q(x).view(-1, self.hidden_size))
+        return x
+
+def apply_rotary_pos_emb_flashatt(tensor: torch.Tensor, freqs: torch.Tensor) -> torch.Tensor:
+    tensor_ = tensor.float()
+    cos = freqs.cos().float()
+    sin = freqs.sin().float()
+    output = apply_rotary_emb(tensor_, cos, sin).type_as(tensor)
+    return output
+
+class Qwen2_5_VLVisionFlashAttention2(nn.Module):
+    def __init__(self, dim: int, num_heads: int = 16) -> None:
+        super().__init__()
+        self.num_heads = num_heads
+        self.qkv = nn.Linear(dim, dim * 3, bias=True)
+        self.proj = nn.Linear(dim, dim)
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        cu_seqlens: torch.Tensor,
+        rotary_pos_emb: torch.Tensor = None,
+    ) -> torch.Tensor:
+        seq_length = hidden_states.shape[0]
+        q, k, v = self.qkv(hidden_states).reshape(seq_length, 3, self.num_heads, -1).permute(1, 0, 2, 3).unbind(0)
+        q = apply_rotary_pos_emb_flashatt(q.unsqueeze(0), rotary_pos_emb).squeeze(0)
+        k = apply_rotary_pos_emb_flashatt(k.unsqueeze(0), rotary_pos_emb).squeeze(0)
+
+        max_seqlen = (cu_seqlens[1:] - cu_seqlens[:-1]).max().item()
+        attn_output = flash_attn_varlen_func(q, k, v, cu_seqlens, cu_seqlens, max_seqlen, max_seqlen).reshape(
+            seq_length, -1
+        )
+        attn_output = self.proj(attn_output)
+        return attn_output
+
+def rotate_half(x):
+    """Rotates half the hidden dims of the input."""
+    x1 = x[..., : x.shape[-1] // 2]
+    x2 = x[..., x.shape[-1] // 2:]
+    return torch.cat((-x2, x1), dim=-1)
+
+def apply_rotary_pos_emb_vision(tensor: torch.Tensor, freqs: torch.Tensor) -> torch.Tensor:
+    orig_dtype = tensor.dtype
+    tensor = tensor.float()
+    cos = freqs.cos()
+    sin = freqs.sin()
+    cos = cos.unsqueeze(1).repeat(1, 1, 2).unsqueeze(0).float()
+    sin = sin.unsqueeze(1).repeat(1, 1, 2).unsqueeze(0).float()
+    output = (tensor * cos) + (rotate_half(tensor) * sin)
+    output = output.to(orig_dtype)
+    return output
+
+class Qwen2_5_VLVisionAttention(nn.Module):
+    def __init__(self, dim: int, num_heads: int = 16) -> None:
+        super().__init__()
+        self.num_heads = num_heads
+        self.head_dim = dim // num_heads
+        self.qkv = nn.Linear(dim, dim * 3, bias=True)
+        self.proj = nn.Linear(dim, dim)
+
+    def forward(
+        self, hidden_states: torch.Tensor, cu_seqlens: torch.Tensor, rotary_pos_emb: torch.Tensor = None
+    ) -> torch.Tensor:
+        seq_length = hidden_states.shape[0]
+        q, k, v = self.qkv(hidden_states).reshape(seq_length, 3, self.num_heads, -1).permute(1, 0, 2, 3).unbind(0)
+        q = apply_rotary_pos_emb_vision(q.unsqueeze(0), rotary_pos_emb).squeeze(0)
+        k = apply_rotary_pos_emb_vision(k.unsqueeze(0), rotary_pos_emb).squeeze(0)
+
+        attention_mask = torch.full(
+            [1, seq_length, seq_length], torch.finfo(q.dtype).min, device=q.device, dtype=q.dtype
+        )
+        for i in range(1, len(cu_seqlens)):
+            attention_mask[..., cu_seqlens[i - 1]: cu_seqlens[i], cu_seqlens[i - 1]: cu_seqlens[i]] = 0
+
+        q = q.transpose(0, 1)
+        k = k.transpose(0, 1)
+        v = v.transpose(0, 1)
+        attn_weights = torch.matmul(q, k.transpose(1, 2)) / math.sqrt(self.head_dim)
+        attn_weights = attn_weights + attention_mask
+        attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to(q.dtype)
+        attn_output = torch.matmul(attn_weights, v)
+        attn_output = attn_output.transpose(0, 1)
+        attn_output = attn_output.reshape(seq_length, -1)
+        attn_output = self.proj(attn_output)
+        return attn_output
+
+class Qwen2_5_VLVisionSdpaAttention(nn.Module):
+    def __init__(self, dim: int, num_heads: int = 16) -> None:
+        super().__init__()
+        self.num_heads = num_heads
+        self.qkv = nn.Linear(dim, dim * 3, bias=True)
+        self.proj = nn.Linear(dim, dim)
+
+    def forward(
+        self, hidden_states: torch.Tensor, cu_seqlens: torch.Tensor, rotary_pos_emb: torch.Tensor = None
+    ) -> torch.Tensor:
+        seq_length = hidden_states.shape[0]
+        q, k, v = self.qkv(hidden_states).reshape(seq_length, 3, self.num_heads, -1).permute(1, 0, 2, 3).unbind(0)
+        q = apply_rotary_pos_emb_vision(q.unsqueeze(0), rotary_pos_emb).squeeze(0)
+        k = apply_rotary_pos_emb_vision(k.unsqueeze(0), rotary_pos_emb).squeeze(0)
+
+        attention_mask = torch.zeros([1, seq_length, seq_length], device=q.device, dtype=torch.bool)
+        for i in range(1, len(cu_seqlens)):
+            attention_mask[..., cu_seqlens[i - 1]: cu_seqlens[i], cu_seqlens[i - 1]: cu_seqlens[i]] = True
+        q = q.transpose(0, 1)
+        k = k.transpose(0, 1)
+        v = v.transpose(0, 1)
+        attn_output = F.scaled_dot_product_attention(q, k, v, attention_mask, dropout_p=0.0)
+        attn_output = attn_output.transpose(0, 1)
+        attn_output = attn_output.reshape(seq_length, -1)
+        attn_output = self.proj(attn_output)
+        return attn_output
+
+QWEN2_5_VL_VISION_ATTENTION_CLASSES = {
+    "eager": Qwen2_5_VLVisionAttention,
+    "flash_attention_2": Qwen2_5_VLVisionFlashAttention2,
+    "sdpa": Qwen2_5_VLVisionSdpaAttention,
+}
+
+class Qwen2_5_VLVisionBlock(nn.Module):
+    def __init__(self, config, attn_implementation: str = "sdpa") -> None:
+        super().__init__()
+        self.norm1 = Qwen2RMSNorm(config.hidden_size, eps=1e-6)
+        self.norm2 = Qwen2RMSNorm(config.hidden_size, eps=1e-6)
+        self.attn = QWEN2_5_VL_VISION_ATTENTION_CLASSES[attn_implementation](
+            config.hidden_size, num_heads=config.num_heads
+        )
+        self.mlp = Qwen2_5_VLMLP(config, bias=True)
+
+    def forward(self, hidden_states, cu_seqlens, rotary_pos_emb) -> torch.Tensor:
+        hidden_states = hidden_states + self.attn(
+            self.norm1(hidden_states),
+            cu_seqlens=cu_seqlens,
+            rotary_pos_emb=rotary_pos_emb,
+        )
+        hidden_states = hidden_states + self.mlp(self.norm2(hidden_states))
+        return hidden_states
+
+class Qwen2_5_VisionTransformer(PreTrainedModel):
+    config_class = Qwen2_5_VLVisionConfig
+    _no_split_modules = ["Qwen2_5_VLVisionBlock"]
+    _supports_flash_attn_2 = True
+    _supports_sdpa = True
+
+    def __init__(self, config, *inputs, **kwargs) -> None:
+        super().__init__(config, *inputs, **kwargs)
+        self.spatial_merge_size = config.spatial_merge_size
+        self.patch_size = config.patch_size
+        self.fullatt_block_indexes = config.fullatt_block_indexes
+        self.window_size = config.window_size
+        self.spatial_merge_unit = self.spatial_merge_size * self.spatial_merge_size
+
+        self.patch_embed = Qwen2_5_VisionPatchEmbed(
+            patch_size=config.patch_size,
+            temporal_patch_size=config.temporal_patch_size,
+            in_channels=config.in_channels,
+            embed_dim=config.hidden_size,
+        )
+
+        head_dim = config.hidden_size // config.num_heads
+        self.rotary_pos_emb = Qwen2_5_VisionRotaryEmbedding(head_dim // 2)
+
+        self.blocks = nn.ModuleList(
+            [Qwen2_5_VLVisionBlock(config, config._attn_implementation) for _ in range(config.depth)]
+        )
+        self.merger = Qwen2_5_VLPatchMerger(
+            dim=config.out_hidden_size,
+            context_dim=config.hidden_size,
+            spatial_merge_size=config.spatial_merge_size,
+        )
+        self.gradient_checkpointing = False
+        self.image_emb_dim = 8192
+
+    def get_dtype(self) -> torch.dtype:
+        return self.blocks[0].mlp.down_proj.weight.dtype
+
+    def rot_pos_emb(self, grid_thw):
+        pos_ids = []
+        for t, h, w in grid_thw:
+            hpos_ids = torch.arange(h).unsqueeze(1).expand(-1, w)
+            hpos_ids = hpos_ids.reshape(
+                h // self.spatial_merge_size,
+                self.spatial_merge_size,
+                w // self.spatial_merge_size,
+                self.spatial_merge_size,
+            )
+            hpos_ids = hpos_ids.permute(0, 2, 1, 3)
+            hpos_ids = hpos_ids.flatten()
+
+            wpos_ids = torch.arange(w).unsqueeze(0).expand(h, -1)
+            wpos_ids = wpos_ids.reshape(
+                h // self.spatial_merge_size,
+                self.spatial_merge_size,
+                w // self.spatial_merge_size,
+                self.spatial_merge_size,
+            )
+            wpos_ids = wpos_ids.permute(0, 2, 1, 3)
+            wpos_ids = wpos_ids.flatten()
+            pos_ids.append(torch.stack([hpos_ids, wpos_ids], dim=-1).repeat(t, 1))
+        pos_ids = torch.cat(pos_ids, dim=0)
+        max_grid_size = grid_thw[:, 1:].max()
+        rotary_pos_emb_full = self.rotary_pos_emb(max_grid_size)
+        rotary_pos_emb = rotary_pos_emb_full[pos_ids].flatten(1)
+        return rotary_pos_emb
+
+    def get_window_index(self, grid_thw):
+        window_index: list = []
+        cu_window_seqlens: list = [0]
+        window_index_id = 0
+        vit_merger_window_size = self.window_size // self.spatial_merge_size // self.patch_size
+
+        for grid_t, grid_h, grid_w in grid_thw:
+            llm_grid_h, llm_grid_w = (
+                grid_h // self.spatial_merge_size,
+                grid_w // self.spatial_merge_size,
+            )
+            index = torch.arange(grid_t * llm_grid_h * llm_grid_w).reshape(grid_t, llm_grid_h, llm_grid_w)
+            pad_h = vit_merger_window_size - llm_grid_h % vit_merger_window_size
+            pad_w = vit_merger_window_size - llm_grid_w % vit_merger_window_size
+            num_windows_h = (llm_grid_h + pad_h) // vit_merger_window_size
+            num_windows_w = (llm_grid_w + pad_w) // vit_merger_window_size
+            index_padded = F.pad(index, (0, pad_w, 0, pad_h), "constant", -100)
+            index_padded = index_padded.reshape(
+                grid_t,
+                num_windows_h,
+                vit_merger_window_size,
+                num_windows_w,
+                vit_merger_window_size,
+            )
+            index_padded = index_padded.permute(0, 1, 3, 2, 4).reshape(
+                grid_t,
+                num_windows_h * num_windows_w,
+                vit_merger_window_size,
+                vit_merger_window_size,
+            )
+            seqlens = (index_padded != -100).sum([2, 3]).reshape(-1)
+            index_padded = index_padded.reshape(-1)
+            index_new = index_padded[index_padded != -100]
+            window_index.append(index_new + window_index_id)
+            cu_seqlens_tmp = seqlens.cumsum(0) * self.spatial_merge_unit + cu_window_seqlens[-1]
+            cu_window_seqlens.extend(cu_seqlens_tmp.tolist())
+            window_index_id += (grid_t * llm_grid_h * llm_grid_w).item()
+        window_index = torch.cat(window_index, dim=0)
+
+        return window_index, cu_window_seqlens
+
+    def forward(self, hidden_states: torch.Tensor, grid_thw: torch.Tensor, is_list=False, **kwargs) -> torch.Tensor:
+        """
+        Args:
+            hidden_states (`torch.Tensor` of shape `(batch_size, seq_len, hidden_size)`):
+                The final hidden states of the model.
+            grid_thw (`torch.Tensor` of shape `(num_images_or_videos, 3)`):
+                The temporal, height and width of feature shape of each image in LLM.
+
+        Returns:
+            `torch.Tensor`: hidden_states.
+        """
+        hidden_states = hidden_states.type(self.get_dtype())
+        if is_list:
+            image_grid_thw_array = []
+            for image_grid_thw_ in grid_thw:
+                image_grid_thw_array.extend(image_grid_thw_)
+            grid_thw = torch.tensor(image_grid_thw_array, device=hidden_states.device)
+
+        hidden_states = self.patch_embed(hidden_states)
+        rotary_pos_emb = self.rot_pos_emb(grid_thw)
+        window_index, cu_window_seqlens = self.get_window_index(grid_thw)
+        cu_window_seqlens = torch.tensor(
+            cu_window_seqlens,
+            device=hidden_states.device,
+            dtype=grid_thw.dtype if torch.jit.is_tracing() else torch.int32,
+        )
+        cu_window_seqlens = torch.unique_consecutive(cu_window_seqlens)
+
+        seq_len, _ = hidden_states.size()
+        hidden_states = hidden_states.reshape(seq_len // self.spatial_merge_unit, self.spatial_merge_unit, -1)
+        hidden_states = hidden_states[window_index, :, :]
+        hidden_states = hidden_states.reshape(seq_len, -1)
+        rotary_pos_emb = rotary_pos_emb.reshape(seq_len // self.spatial_merge_unit, self.spatial_merge_unit, -1)
+        rotary_pos_emb = rotary_pos_emb[window_index, :, :]
+        rotary_pos_emb = rotary_pos_emb.reshape(seq_len, -1)
+
+        cu_seqlens = torch.repeat_interleave(grid_thw[:, 1] * grid_thw[:, 2], grid_thw[:, 0]).cumsum(
+            dim=0,
+            # Select dtype based on the following factors:
+            #  - FA2 requires that cu_seqlens_q must have dtype int32
+            #  - torch.onnx.export requires that cu_seqlens_q must have same dtype as grid_thw
+            # See https://github.com/huggingface/transformers/pull/34852 for more information
+            dtype=grid_thw.dtype if torch.jit.is_tracing() else torch.int32,
+        )
+        cu_seqlens = F.pad(cu_seqlens, (1, 0), value=0)
+
+        for layer_num, blk in enumerate(self.blocks):
+            if layer_num in self.fullatt_block_indexes:
+                cu_seqlens_now = cu_seqlens
+            else:
+                cu_seqlens_now = cu_window_seqlens
+            if self.gradient_checkpointing and self.training:
+                hidden_states = self._gradient_checkpointing_func(
+                    blk.__call__, hidden_states, cu_seqlens_now, rotary_pos_emb
+                )
+            else:
+                hidden_states = blk(
+                    hidden_states,
+                    cu_seqlens=cu_seqlens_now,
+                    rotary_pos_emb=rotary_pos_emb,
+                )
+
+        hidden_states = self.merger(hidden_states)
+
+        reverse_indices = torch.argsort(window_index)
+        hidden_states = hidden_states[reverse_indices, :]
+
+        return hidden_states

code/requirements-rocm.txt

@@ -0,0 +1,36 @@
+# ROCm port of requirements.txt for AMD gfx1151 (ROCm 7.13).
+# Omitted on purpose — do not add them back:
+#   torch, torchvision: the target interpreter already has a working ROCm
+#     build (torch 2.10.0, torch.version.hip 7.13.99004). Reinstalling the
+#     upstream CUDA pins would replace it.
+#   transformer-engine: NVIDIA CUDA-only; ROCm has no TE. Qwen2RMSNorm in
+#     qwen2_5_vit.py is pure PyTorch, and the unused TE import is gone.
+# Create the venv with system site packages so that ROCm torch is inherited:
+#   python3 -m venv --system-site-packages .venv
+#   .venv/bin/pip install -r requirements-rocm.txt
+#
+# Also omitted / relaxed versus upstream, because the ROCm interpreter is Python 3.13
+# and its torch is built against numpy 2.x:
+#   numpy      upstream 1.23.1 has no Python 3.13 wheels, and downgrading would
+#              break the inherited torch. Inherit the system numpy (validated 2.2.4).
+#   Pillow     upstream 10.4.0 has no Python 3.13 wheels. Inherit (validated 11.1.0).
+#   safetensors  inherit the system build (validated 0.8.0).
+#
+# Validated on halo (gfx1151, ROCm 7.13, Python 3.13.5) on 2026-09-22:
+#   torch 2.10.0 (hip 7.13.99004) | numpy 2.2.4 | Pillow 11.1.0 | safetensors 0.8.0
+#   transformers 4.57.1 | diffusers 0.36.0 | accelerate 1.13.0 | tokenizers 0.22.2
+#   huggingface-hub 0.34.0 | peft 0.17.0
+transformers==4.57.1
+diffusers==0.36.0
+accelerate==1.13.0
+tokenizers==0.22.2
+huggingface-hub==0.34.0
+peft==0.17.0
+requests==2.32.3
+tqdm==4.67.1
+typing-extensions==4.15.0
+
+# Optional FlashAttention 2 backend (validated: flash-attn==2.7.3). The CLI
+# default is eager attention; --attn-implementation flash_attention_2 needs
+# this package.
+# flash-attn==2.7.3

code/requirements.txt

@@ -0,0 +1,23 @@
+# Dependency set validated end-to-end on 2026-09-16. Exact pins are
+# intentional: this is the environment known to run both published checkpoint
+# families, not a broader compatibility matrix.
+torch==2.4.0
+torchvision==0.19.0
+transformers==4.57.1
+diffusers==0.36.0
+accelerate==1.13.0
+transformer-engine[pytorch]==1.11.0
+safetensors==0.7.0
+tokenizers==0.22.2
+huggingface-hub==0.34.0
+peft==0.17.0
+numpy==1.23.1
+Pillow==10.4.0
+requests==2.32.3
+tqdm==4.67.1
+typing-extensions==4.15.0
+
+# Optional FlashAttention 2 backend (validated: flash-attn==2.7.3). The CLI
+# default is eager attention; --attn-implementation flash_attention_2 needs
+# this package.
+# flash-attn==2.7.3

code/tokenization_bailing.py

@@ -0,0 +1,1024 @@
+#!/usr/bin/env python3
+# coding=utf-8
+# Copyright (c) Ant Group. All rights reserved.
+
+import itertools
+from typing import Any, Dict, List, Optional, Union
+
+import torch
+from transformers import PreTrainedTokenizerFast
+from transformers.tokenization_utils_base import AddedToken, BatchEncoding
+from transformers.utils import TensorType, logging
+
+logger = logging.get_logger(__name__)
+
+
+def is_system(msg):
+    return msg['role'].lower() == 'system'
+
+
+def is_user(msg):
+    return msg['role'].lower() in ['human', 'user']
+
+
+def is_assistant(msg):
+    return msg['role'].lower() == 'assistant'
+
+
+def _convert_to_conversation(query, system=None):
+    conversation = []
+    if system:
+        conversation.append({"role": "SYSTEM", "content": system})
+    if isinstance(query, str):
+        conversation.append({"role": "HUMAN", "content": query})
+    elif isinstance(query, List):
+        conversation.extend(query)
+    elif isinstance(query, Dict):
+        if "messages" in query:
+            conversation.extend(query["messages"])
+            if "system_message" in query and len(conversation) > 0 and not is_system(conversation[0]):
+                conversation.insert(0, {"role": "SYSTEM", "content": query["system_message"]})
+        else:
+            conversation.append(query)
+    return conversation
+
+
+class BailingTokenizer(PreTrainedTokenizerFast):
+    is_bailing_tokenizer = True
+    model_input_names = ["input_ids", "attention_mask"]
+    slow_tokenizer_class = None
+
+    # add gmask_token
+    SPECIAL_TOKENS_ATTRIBUTES = [
+        "bos_token",
+        "eos_token",
+        "unk_token",
+        "sep_token",
+        "pad_token",
+        "cls_token",
+        "mask_token",
+        "gmask_token",
+        "additional_special_tokens",
+    ]
+
+    def __init__(
+        self,
+        vocab_file=None,
+        merges_file=None,
+        tokenizer_file=None,
+        clean_up_tokenization_spaces=False,
+        bos_token="<|startoftext|>",
+        eos_token="<|endoftext|>",
+        cls_token="[CLS]",
+        pad_token="<|endoftext|>",
+        gmask_token="[gMASK]",
+        add_bos_token=False,
+        add_eos_token=False,
+        **kwargs,
+    ):
+        self.add_bos_token = add_bos_token
+
+        self._gmask_token = (
+            AddedToken(gmask_token, lstrip=False, rstrip=False, normalized=False)
+            if isinstance(gmask_token, str)
+            else gmask_token
+        )
+
+        self._sop_token = (
+            AddedToken(bos_token, lstrip=False, rstrip=False, normalized=False)
+            if isinstance(bos_token, str)
+            else bos_token
+        )
+
+        self._eop_token = (
+            AddedToken(eos_token, lstrip=False, rstrip=False, normalized=False)
+            if isinstance(eos_token, str)
+            else eos_token
+        )
+
+        super().__init__(
+            vocab_file=vocab_file,
+            merges_file=merges_file,
+            tokenizer_file=tokenizer_file,
+            clean_up_tokenization_spaces=clean_up_tokenization_spaces,
+            bos_token=bos_token,
+            eos_token=eos_token,
+            cls_token=cls_token,
+            pad_token=pad_token,
+            gmask_token=gmask_token,
+            add_bos_token=add_bos_token,
+            add_eos_token=add_eos_token,
+            **kwargs,
+        )
+
+        self.check_special_tokens()
+
+    def check_special_tokens(self):
+        '''
+        eos_token, cls_token, mask_token
+        special tokens should init, check special token is not None
+        '''
+        for name, special_token in zip(
+            ['eos', 'bos', 'cls', 'gmask'],
+            [self.eos_token, self.bos_token, self.cls_token, self.gmask_token],
+        ):
+            assert special_token is not None, f'should init special token [{name}] in tokenizer_config.json'
+
+    @property
+    def gmask_token(self) -> Optional[str]:
+        if self._gmask_token is None:
+            if self.verbose:
+                logger.error("Using gmask_token, but it is not set yet.")
+            return None
+        return str(self._gmask_token)
+
+    @gmask_token.setter
+    def gmask_token(self, value):
+        if not isinstance(value, (str, AddedToken)) and value is not None:
+            raise ValueError("Cannot set a non-string value as the gmask token")
+        self._gmask_token = value
+
+    @property
+    def gmask_token_id(self) -> Optional[int]:
+        if self._gmask_token is None:
+            return None
+        return self.convert_tokens_to_ids(self.gmask_token)
+
+    @property
+    def sop_token(self) -> Optional[str]:
+        if self._sop_token is None:
+            if self.verbose:
+                logger.error("Using sop_token, but it is not set yet.")
+            return None
+        return str(self._sop_token)
+
+    @sop_token.setter
+    def sop_token(self, value):
+        if not isinstance(value, (str, AddedToken)) and value is not None:
+            raise ValueError("Cannot set a non-string value as the sop token")
+        self._sop_token = value
+
+    @property
+    def sop_token_id(self) -> Optional[int]:
+        if self._sop_token is None:
+            return None
+        return self.convert_tokens_to_ids(self.sop_token)
+
+    @property
+    def eop_token(self) -> Optional[str]:
+        if self._eop_token is None:
+            if self.verbose:
+                logger.error("Using eop_token, but it is not set yet.")
+            return None
+        return str(self._eop_token)
+
+    @eop_token.setter
+    def eop_token(self, value):
+        if not isinstance(value, (str, AddedToken)) and value is not None:
+            raise ValueError("Cannot set a non-string value as the eop token")
+        self._eop_token = value
+
+    @property
+    def eop_token_id(self) -> Optional[int]:
+        if self._eop_token is None:
+            return None
+        return self.convert_tokens_to_ids(self.eop_token)
+
+    @property
+    def vocab_size(self):
+        return len(self.get_vocab())
+
+    def _chat_from_json(self, chat, chat_format="antglm_chat", system=None):
+        raise NotImplementedError(
+            "the legacy antglm_chat rendering path was removed; a checkpoint "
+            "tokenizer must provide a chat_template in tokenizer_config.json"
+        )
+
+    def apply_chat_template(
+        self,
+        conversation: Union[List[Dict[str, str]], List[List[Dict[str, str]]]],
+        tools: Optional[List[Dict]] = None,
+        documents: Optional[List[Dict[str, str]]] = None,
+        chat_template: Optional[str] = None,
+        add_generation_prompt: bool = False,
+        system: str = None,  # only used for legacy chatml
+        tokenize=False,
+        padding: bool = False,
+        truncation: bool = False,
+        max_length: Optional[int] = None,
+        return_tensors: Optional[Union[str, TensorType]] = None,
+        return_dict: bool = False,
+        return_assistant_tokens_mask: bool = False,
+        tokenizer_kwargs: Optional[Dict[str, Any]] = None,
+        **kwargs,
+    ):
+        if hasattr(self, "chat_template") and self.chat_template:
+            if isinstance(conversation, Dict) and "messages" in conversation:
+                conversation = conversation["messages"]
+            # use transformers built-in method
+            return super().apply_chat_template(
+                conversation=conversation,
+                tools=tools,
+                documents=documents,
+                chat_template=chat_template,
+                add_generation_prompt=add_generation_prompt,
+                tokenize=tokenize,
+                padding=padding,
+                truncation=truncation,
+                return_tensors=return_tensors,
+                return_dict=return_dict,
+                return_assistant_tokens_mask=return_assistant_tokens_mask,
+                tokenizer_kwargs=tokenizer_kwargs,
+            )
+
+        # The legacy antglm_chat rendering path was removed. A checkpoint
+        # tokenizer must provide a chat_template in tokenizer_config.json.
+        raise ValueError(
+            "BailingMM2 inference requires a chat_template in the checkpoint "
+            "tokenizer_config.json"
+        )
+
+    def _build_position_ids(
+        self,
+        mask_pos: int,
+        bos_pos: int,
+        max_output_length: int,
+        rotary_type: Optional[str] = "none",
+        **kwargs,
+    ) -> List[List[int]]:
+        window_size = kwargs.get("window_size", 1024) - 1
+        block_position_ids = [0] * bos_pos
+
+        # location of the mask, used to construct output position ids later
+        if "1d" in rotary_type:
+            position_ids = list(range(bos_pos)) + list(range(mask_pos + 1, mask_pos + max_output_length + 2))
+            block_position_ids = block_position_ids + list(range(1, max_output_length + 2))
+        elif "2d" in rotary_type:
+            # a bos_id is appended to input_ids
+            position_ids = list(range(bos_pos))
+            position_ids = position_ids + [mask_pos] * (1 + max_output_length)
+            block_position_ids = block_position_ids + list(range(1, max_output_length + 2))
+        else:
+            # build position ids
+            position_ids = []
+            repeat_times = bos_pos // window_size
+            for _ in range(repeat_times):
+                position_ids += list(range(window_size))
+            position_ids += list(range(bos_pos - window_size * repeat_times))
+            # need consider additional bos_id after input_ids
+            mask_pos = position_ids[-1]
+            position_ids += [mask_pos] * (max_output_length + 1)
+
+            block_repeat_times = max_output_length // (window_size - 1)
+            additional_block_position_ids = []
+            for _ in range(block_repeat_times):
+                additional_block_position_ids += list(range(1, window_size))
+            additional_block_position_ids += list(
+                range(1, max_output_length + 2 - (window_size - 1) * block_repeat_times)
+            )
+            block_position_ids = block_position_ids + additional_block_position_ids
+
+        position_ids = [position_ids, block_position_ids]
+        return position_ids
+
+    def _build_inputs_for_generation(
+        self,
+        input_ids: List[int],
+        max_input_length=None,
+        left_truncate=True,
+        max_output_length=1024,
+        rotary_type="none",
+        unidirectional_attention: bool = True,
+        attention_dtype=None,
+        **kwargs,
+    ):
+        if max_input_length and len(input_ids) > max_input_length:
+            if left_truncate:
+                input_ids = input_ids[-max_input_length:]
+            else:
+                input_ids = input_ids[:max_input_length]
+
+        is_left_padding = input_ids[0] == self.eos_token_id
+        if not unidirectional_attention:
+            if input_ids[0] != self.cls_token_id:
+                input_ids = [self.cls_token_id] + input_ids
+
+            if self.gmask_token_id not in set(input_ids):
+                input_ids = input_ids + [self.gmask_token_id]
+
+            mask_pos = input_ids.index(self.gmask_token_id)
+            sep = len(input_ids)
+        else:
+            if self.add_bos_token:
+                input_ids = input_ids + [self.bos_token_id]
+                if self.eos_token_id in input_ids:
+                    mask_pos = input_ids.index(self.eos_token_id) - 1
+                else:
+                    mask_pos = len(input_ids) - 1
+                sep = len(input_ids) - 1
+            else:
+                sep = len(input_ids)
+                if self.eos_token_id in input_ids:
+                    if is_left_padding:
+                        ori_input_ids = input_ids
+                        input_ids = input_ids[::-1]
+                    mask_pos = input_ids.index(self.eos_token_id) - 1
+                    mask_pos = max(0, mask_pos)  # for empty sequence
+                    if is_left_padding:
+                        input_ids = ori_input_ids
+                        mask_pos = sep - 1 - mask_pos  # the first non-eos token
+
+                else:
+                    mask_pos = len(input_ids) - 1
+
+        position_ids = self._build_position_ids(mask_pos, sep, max_output_length, rotary_type, **kwargs)
+
+        if is_left_padding:
+            position_ids[0] = [max(0, i - mask_pos) for i in range(len(position_ids[0]))]
+
+        # a bos_id is appended to input_ids
+        total_length = sep + max_output_length
+        if self.add_bos_token:
+            total_length += 1
+
+        def build_mask_matrix(seq_length, sep, mask_pos, unidirectional_attention):
+            # use bool attention masks for long sequences to save memory
+            if unidirectional_attention:
+                attention_mask = torch.ones([seq_length, seq_length], dtype=attention_dtype)
+                attention_mask = torch.tril(attention_mask)
+                if is_left_padding:
+                    attention_mask[:, :mask_pos] = 0
+                else:
+                    attention_mask[:, mask_pos + 1 : sep] = 0
+            else:
+                attention_mask = torch.zeros([seq_length, seq_length], dtype=attention_dtype)
+                attention_mask[:, : mask_pos + 1] = 1
+                for i in range(sep, total_length):
+                    attention_mask[i, sep : i + 1] = 1
+            return attention_mask
+
+        if self.add_bos_token:
+            attention_mask = build_mask_matrix(total_length, sep + 1, mask_pos, unidirectional_attention)
+        else:
+            attention_mask = build_mask_matrix(total_length, sep, mask_pos, unidirectional_attention)
+        attention_mask = torch.unsqueeze(attention_mask, dim=0)
+        attention_mask = torch.unsqueeze(attention_mask, dim=1)
+        if attention_dtype is None:
+            attention_mask = attention_mask.long()
+        inputs = {
+            "input_ids": torch.Tensor([input_ids]).long(),
+            "position_ids": torch.Tensor([position_ids]).long(),
+            "attention_mask": attention_mask,
+        }
+        return BatchEncoding(inputs)
+
+    def build_inputs_for_generation(
+        self,
+        input_ids: Union[List[int], List[List[int]], torch.Tensor],
+        max_input_length=None,
+        left_truncate=True,
+        max_output_length=1024,
+        rotary_type="1d",
+        unidirectional_attention=True,
+        attention_dtype=None,
+        **kwargs,
+    ):
+        if isinstance(input_ids, torch.Tensor):
+            input_ids = input_ids.tolist()
+
+        if isinstance(input_ids[0], list):
+            input_ids_list = []
+            position_ids_list = []
+            attention_mask_list = []
+            for _input_ids in input_ids:
+                inputs = self._build_inputs_for_generation(
+                    _input_ids,
+                    max_input_length=max_input_length,
+                    left_truncate=left_truncate,
+                    max_output_length=max_output_length,
+                    rotary_type=rotary_type,
+                    unidirectional_attention=unidirectional_attention,
+                    attention_dtype=attention_dtype,
+                    **kwargs,
+                )
+                input_ids_list.append(inputs['input_ids'])
+                position_ids_list.append(inputs['position_ids'])
+                attention_mask_list.append(inputs["attention_mask"])
+
+            max_ids_length = max([input.size(1) for input in input_ids_list])
+
+            for i in range(len(input_ids)):
+                cur_ids_length = input_ids_list[i].size(1)
+                if cur_ids_length < max_ids_length:
+                    # pad input ids
+                    pad_input_ids = input_ids_list[i].new_zeros((1, max_ids_length - cur_ids_length))
+                    input_ids_list[i] = torch.cat([pad_input_ids, input_ids_list[i]], dim=-1)
+
+                    # pad postition ids with left pad
+                    # 0, 1, 2, 3, 4 ... -> 0, ..., 0, 1, 2, 3, 4, ...
+                    pad_position_ids = input_ids_list[i].new_zeros((1, 2, max_ids_length - cur_ids_length))
+                    position_ids_list[i] = torch.cat([pad_position_ids, position_ids_list[i]], dim=-1)
+
+                    # pad generation attention mask with left and bottom pad
+                    new_attention_mask = input_ids_list[i].new_zeros(
+                        1,
+                        1,
+                        max_ids_length + max_output_length,
+                        max_ids_length + max_output_length,
+                    )
+                    new_attention_mask[
+                        :,
+                        :,
+                        max_ids_length - cur_ids_length :,
+                        max_ids_length - cur_ids_length :,
+                    ] = attention_mask_list[i]
+                    attention_mask_list[i] = new_attention_mask.contiguous()
+
+            input_ids_list = torch.cat(input_ids_list, dim=0)
+            position_ids_list = torch.cat(position_ids_list, dim=0)
+            attention_mask_list = torch.cat(attention_mask_list, dim=0)
+
+            inputs = {
+                "input_ids": input_ids_list,
+                "position_ids": position_ids_list,
+                "attention_mask": attention_mask_list,
+            }
+
+            return BatchEncoding(inputs)
+        else:
+            return self._build_inputs_for_generation(
+                input_ids,
+                max_input_length=max_input_length,
+                left_truncate=left_truncate,
+                max_output_length=max_output_length,
+                rotary_type=rotary_type,
+                unidirectional_attention=unidirectional_attention,
+                **kwargs,
+            )
+
+    def _build_inputs_for_train(
+        self,
+        inputs: Union[str, List[str]],
+        outputs: Union[str, List[str]],
+        new_conversation_offset: List[int] = None,
+        max_length: int = 2048,
+        rotary_type: str = "1d",
+        left_truncate: bool = True,
+        unidirectional_attention: bool = True,
+        isolation_position_ids: bool = False,
+        padding: bool = True,
+        use_fa2: bool = True,
+        use_packed: bool = True,
+        use_baichuan_packed: bool = False,
+        skip_truncated_turn: bool = False,
+        return_attention_mask: bool = True,
+    ):
+        r"""
+        Build tensor input for model training. If inputs and outputs are list, will pack them.
+
+        Args:
+            inputs (str, List[str], List[Dict], List[List[Dict]]): the input prompts.
+            outputs (str, List[str]): the output responses.
+            max_length (int, Optional): the maximum length of the final input ids for training. Default: 2048
+            rotary_type (str, Optional): the rotary type of position embedding. Default: 1d
+            left_truncate (bool, Optional): whether truncate the inputs from left. Default: True
+            use_fa2 (bool, Optional): whether to build attention mask under flash attention 2.
+            new_conversation_offset (List[int], Optional): marks turns that start a brand-new conversation; [0, 1] means inputs[0]/outputs[0] is one turn, inputs[1]/outputs[1] another.
+        """
+        if use_packed and use_baichuan_packed and unidirectional_attention:
+            return self._build_baichuan_inputs_for_train(
+                inputs,
+                outputs,
+                new_conversation_offset,
+                max_length,
+                rotary_type,
+                left_truncate,
+                skip_truncated_turn,
+                use_fa2,
+                padding,
+            )
+        if isinstance(inputs, str):
+            inputs = [inputs]
+        if isinstance(outputs, str):
+            outputs = [outputs]
+
+        assert len(inputs) == len(outputs)
+
+        input_ids = [self(item)['input_ids'] for item in inputs]
+        output_ids = [self(item)['input_ids'] for item in outputs]
+
+        packed_input_ids = []
+        packed_output_ids = []
+        if new_conversation_offset is None:
+            new_conversation_offset = list(range(0, len(inputs)))
+        assert 0 in new_conversation_offset, f"0 must be present, check new_conversation_offset: {new_conversation_offset}"
+        current_len = 0
+
+        for idx, (input, output) in enumerate(zip(input_ids, output_ids)):
+            num_special_tokens = 0
+            if not unidirectional_attention:
+                if idx in new_conversation_offset:
+                    # cls and gmask
+                    num_special_tokens += 2
+                else:
+                    # only gmask
+                    num_special_tokens += 1
+            else:
+                # sop and eos
+                if self.add_bos_token:
+                    num_special_tokens += 2
+                else:
+                    num_special_tokens += 1
+
+            # truncate
+            if len(input) + len(output) + current_len > max_length - num_special_tokens:
+                if not use_packed or use_fa2 and unidirectional_attention:
+                    attention_mask = torch.tensor(0)
+                elif use_fa2:
+                    attention_mask = -1 * torch.ones([2, max_length])
+                else:
+                    attention_mask = torch.tril(torch.ones([max_length, max_length]))
+                # return an empty sample that does not participate in training
+                default_return = {
+                    'input_ids': (torch.ones(max_length) * self.eos_token_id).long(),
+                    'position_ids': torch.zeros(2, max_length).long(),
+                    'attention_mask': (attention_mask.long()),
+                    'labels': (torch.ones(max_length) * -100).long(),
+                }
+                # if no truncation is needed, return directly
+                if skip_truncated_turn:
+                    if current_len == 0:
+                        return default_return
+                    else:
+                        break
+                left_len = max_length - num_special_tokens - current_len
+                # truncate only the prompt when truncation is required
+                if left_len - len(output) > 0:
+                    if left_truncate:
+                        input = input[-(left_len - len(output)) :]
+                    else:
+                        input = input[: left_len - len(output)]
+                else:
+                    # the response exceeds left_len, return directly
+                    if current_len == 0:
+                        return default_return
+                    else:
+                        break
+            if unidirectional_attention:
+                packed_input_ids.append(list(input))
+            else:
+                if num_special_tokens == 4:
+                    packed_input_ids.append([self.cls_token_id] + list(input) + [self.gmask_token_id])
+                else:
+                    packed_input_ids.append(list(input) + [self.gmask_token_id])
+
+            packed_output_ids.append(list(output) + [self.eos_token_id])
+            current_len += len(input) + len(output) + num_special_tokens
+
+        assert current_len <= max_length
+
+        if use_packed:
+            # pack mode
+            def build_mask_matrix(seq_length, sep):
+                # https://github.com/pytorch/pytorch/issues/101932, fix triu/tril bf16 support
+                m = torch.ones((1, seq_length, seq_length))
+                mask = torch.arange(1, m.shape[-1] + 1).reshape(1, -1, 1).to(m.device)
+                ids = torch.arange(1, m.shape[-1] + 1).reshape(1, 1, -1).expand(1, m.shape[-1], -1).to(m.device)
+                m = (ids <= mask).type_as(m)
+
+                m[0, :, : int(sep)] = 1
+                m = m.squeeze(0)
+                return m
+
+            tokens = []
+            attention_mask_list = []
+            input_length_list = []
+            position_id_list = []
+            block_position_id_list = []
+            for input, output in zip(packed_input_ids, packed_output_ids):
+                if self.add_bos_token:
+                    data = input + [self.sop_token_id] + output
+                    mask_pos = len(input) - 1
+                else:
+                    data = input + output
+                    mask_pos = len(input) - 2
+                if return_attention_mask:
+                    if unidirectional_attention:
+                        attention_mask = build_mask_matrix(len(data), 0)
+                    else:
+                        attention_mask = build_mask_matrix(len(data), len(input))
+                    attention_mask = attention_mask.squeeze((0, 1))
+
+                    attention_mask_list.append(attention_mask)
+                input_length_list.append(len(input))
+                tokens += data
+
+                sop_pos = mask_pos + 1
+                position_ids, block_position_ids = self._build_position_ids(
+                    mask_pos=mask_pos, bos_pos=sop_pos, max_output_length=len(output), rotary_type=rotary_type
+                )
+
+                position_id_list.append(position_ids)
+                block_position_id_list.append(block_position_ids)
+
+            labels = []
+            for i in range(len(packed_input_ids)):
+                if self.add_bos_token:
+                    labels += [-100] * len(packed_input_ids[i]) + packed_output_ids[i] + [-100]
+                else:
+                    labels += [-100] * (len(packed_input_ids[i]) - 1) + packed_output_ids[i] + [-100]
+
+            total_len = 0
+            if use_fa2:
+                pack_attention_mask = -1 * torch.ones([2, current_len])
+            else:
+                pack_attention_mask = torch.tril(torch.ones([current_len, current_len]))
+
+            pack_position_ids = []
+            pack_block_position_ids = []
+            total_len = 0
+            max_index = 0
+            for i in range(len(position_id_list)):
+
+                if use_fa2:
+                    pack_attention_mask[0][i] = total_len
+                    pack_attention_mask[1][i] = total_len + input_length_list[i]
+                else:
+                    pack_attention_mask[
+                        total_len : total_len + attention_mask.shape[0],
+                        total_len : total_len + attention_mask.shape[0],
+                    ] = attention_mask
+                position_ids = [pid + max_index for pid in position_id_list[i]]
+                block_position_ids = block_position_id_list[i]
+                pack_position_ids.extend(position_ids)
+                pack_block_position_ids.extend(block_position_ids)
+                if not isolation_position_ids:
+                    max_index = pack_position_ids[-1] + 1
+                total_len += len(position_id_list[i])
+            position_ids = [pack_position_ids, pack_block_position_ids]
+        else:
+            # single-input mode
+            # in true multi-turn, one sample can span several turns; find the end position of the first turn
+            if len(new_conversation_offset) > 1:
+                end_idx = new_conversation_offset[1]
+            else:
+                end_idx = 1
+            input, output = list(itertools.chain(*packed_input_ids[:end_idx])), list(
+                itertools.chain(*packed_output_ids[:end_idx])
+            )
+            if self.add_bos_token:
+                tokens = input + [self.sop_token_id] + output
+            else:
+                tokens = input + output
+
+            if self.add_bos_token:
+                labels = [-100] * len(input) + output + [-100]
+                position_ids = self._build_position_ids(
+                    mask_pos=len(input) - 1, bos_pos=len(input), max_output_length=len(output), rotary_type=rotary_type
+                )
+            else:
+                labels = [-100] * (len(input) - 1) + output + [-100]
+                position_ids = self._build_position_ids(
+                    mask_pos=len(input) - 2,
+                    bos_pos=len(input) - 1,
+                    max_output_length=len(output),
+                    rotary_type=rotary_type,
+                )
+            attention_mask = len(input)
+        assert current_len == len(tokens)
+
+        # pad up to the maximum length
+        if max_length > 0 and len(tokens) < max_length and padding:
+            pad_length = max_length - len(tokens)
+            tokens += [self.pad_token_id] * pad_length
+            labels.extend([-100] * pad_length)
+            position_ids[0] += [0] * pad_length
+            position_ids[1] += [0] * pad_length
+
+            if use_packed:
+                if use_fa2:
+                    new_attention_mask = -1 * torch.ones([2, max_length])
+                    new_attention_mask[:, :current_len] = pack_attention_mask
+                else:
+                    new_attention_mask = torch.tril(torch.ones([max_length, max_length]))
+                    new_attention_mask[:current_len, :current_len] = pack_attention_mask
+                pack_attention_mask = new_attention_mask.contiguous()
+
+        assert len(tokens) == len(labels)
+
+        if max_length > 0 and padding:
+            assert len(tokens) == max_length
+
+        if use_fa2 and unidirectional_attention:
+            # pack_attention_mask = torch.zeros([1], dtype=torch.long)
+            pack_attention_mask = torch.tensor(0)
+
+        if use_packed:
+            if not use_fa2:
+                attention_mask = pack_attention_mask.unsqueeze(0).long()
+            else:
+                attention_mask = pack_attention_mask
+        else:
+            attention_mask = torch.tensor(attention_mask).long()
+        return {
+            'input_ids': torch.tensor(tokens).long(),
+            'position_ids': torch.tensor(position_ids).long(),
+            'attention_mask': attention_mask,
+            'labels': torch.tensor(labels).long(),
+        }
+
+    def _build_baichuan_inputs_for_train(
+        self,
+        inputs: Union[str, List[str]],
+        outputs: Union[str, List[str]],
+        new_conversation_offset: List[int] = None,
+        max_length: int = 2048,
+        rotary_type: str = "1d",
+        left_truncate: bool = True,
+        skip_truncated_turn: bool = True,
+        use_fa2: bool = True,
+        padding: bool = True,
+    ):
+        '''
+        input:  <role> HUMAN </role> u1 <role>  ASSISTANT </role> a11 a12            <role> HUMAN </role> u2 <role> ASSISTANT </role> a21 a22           <|endoftext|> <role> HUMAN </role> u1 <role>  ASSISTANT </role> a11 a12            <role> HUMAN </role> u2 <role> ASSISTANT </role> a21 a22           <|endoftext|>
+        output: x      x     x       x  x       x         a11     a12 <|endoftext|>  x      x     x       x  x      x         a21     a22 <|endoftext|> x             x      x     x       x  x       x         a11     a12 <|endoftext|>  x      x     x       x  x      x         a21     a22 <|endoftext|> x
+        Only for true multi-turn + pack training of a unidirectional model; requires use_true_multiturn=True
+        '''
+        if isinstance(inputs, str):
+            inputs = [inputs]
+        if isinstance(outputs, str):
+            outputs = [outputs]
+        assert len(inputs) == len(outputs)
+
+        input_ids = [self(item)['input_ids'] for item in inputs]
+        output_ids = [self(item)['input_ids'] for item in outputs]
+
+        packed_input_ids = []
+        packed_output_ids = []
+
+        if new_conversation_offset is None:
+            new_conversation_offset = list(range(0, len(inputs)))
+        assert 0 in new_conversation_offset, f"0 must be present, check new_conversation_offset: {new_conversation_offset}"
+        current_len = 0
+
+        for idx, (input, output) in enumerate(zip(input_ids, output_ids)):
+            num_special_tokens = 0
+            if idx != 0 and idx in new_conversation_offset:
+                # append eos to input_ids; skip it only for the first sample
+                num_special_tokens += 1
+
+            # truncate
+            if len(input) + len(output) + current_len > max_length - num_special_tokens:
+                if use_fa2:
+                    attention_mask = torch.tensor(0)
+                else:
+                    attention_mask = torch.tril(torch.ones([max_length, max_length]))
+                # return an empty sample that does not participate in training
+                default_return = {
+                    'input_ids': (torch.ones(max_length) * self.eos_token_id).long(),
+                    'position_ids': torch.zeros(2, max_length).long(),
+                    'attention_mask': (attention_mask.long()),
+                    'labels': (torch.ones(max_length) * -100).long(),
+                }
+
+                # if no truncation is needed, return directly
+                if skip_truncated_turn:
+                    if current_len == 0:
+                        return default_return
+                    else:
+                        break
+                left_len = max_length - num_special_tokens - current_len
+                # truncate only the prompt when truncation is required
+                if left_len - len(output) > 0:
+                    if left_truncate:
+                        input = input[-(left_len - len(output)) :]
+                    else:
+                        input = input[: left_len - len(output)]
+                else:
+                    # the response exceeds left_len, return directly
+                    if current_len == 0:
+                        return default_return
+                    else:
+                        break
+            # input_ids are concatenated here
+            if num_special_tokens == 1:
+                packed_input_ids.append([self.eos_token_id] + list(input))
+            else:
+                packed_input_ids.append(list(input))
+            packed_output_ids.append(list(output))
+            current_len += len(input) + len(output) + num_special_tokens
+        assert current_len <= max_length
+
+        def build_mask_matrix(seq_length, sep):
+            # https://github.com/pytorch/pytorch/issues/101932, fix triu/tril bf16 support
+            m = torch.ones((1, seq_length, seq_length))
+            mask = torch.arange(1, m.shape[-1] + 1).reshape(1, -1, 1).to(m.device)
+            ids = torch.arange(1, m.shape[-1] + 1).reshape(1, 1, -1).expand(1, m.shape[-1], -1).to(m.device)
+            m = (ids <= mask).type_as(m)
+
+            m[0, :, : int(sep)] = 1
+            m = m.squeeze(0)
+            return m
+
+        tokens = []
+        attention_mask_list = []
+        position_id_list = []
+        block_position_id_list = []
+        token_lens = []
+        for input, output in zip(packed_input_ids, packed_output_ids):
+            data = input + output
+            if not use_fa2:
+                attention_mask = build_mask_matrix(len(data), 0)
+                attention_mask_list.append(attention_mask)
+            tokens += data
+            token_lens.append(len(data))
+
+            position_ids, block_position_ids = self._build_position_ids(
+                mask_pos=len(input) - 2, bos_pos=len(input) - 1, max_output_length=len(output), rotary_type=rotary_type
+            )
+
+            position_id_list.append(position_ids)
+            block_position_id_list.append(block_position_ids)
+
+        labels = []
+        for i in range(len(packed_input_ids)):
+            labels += [-100] * (len(packed_input_ids[i]) - 1) + packed_output_ids[i] + [self.eos_token_id]
+
+        total_len = 0
+        if use_fa2:
+            pack_attention_mask = torch.Tensor([[0], [1]])
+        else:
+            pack_attention_mask = torch.tril(torch.ones([max_length, max_length]))
+
+        pack_position_ids = []
+        pack_block_position_ids = []
+        total_len = 0
+        max_index = 0
+        for i in range(len(token_lens)):
+            if not use_fa2:
+                attention_mask = attention_mask_list[i]
+                pack_attention_mask[
+                    total_len : total_len + attention_mask.shape[0], total_len : total_len + attention_mask.shape[0]
+                ] = attention_mask
+            position_ids = [pid + max_index for pid in position_id_list[i]]
+            block_position_ids = block_position_id_list[i]
+            pack_position_ids.extend(position_ids)
+            pack_block_position_ids.extend(block_position_ids)
+            max_index = pack_position_ids[-1] + 1
+            total_len += token_lens[i]
+        position_ids = [pack_position_ids, pack_block_position_ids]
+
+        if max_length > 0 and len(tokens) < max_length and padding:
+            pad_length = max_length - len(tokens)
+            tokens += [self.pad_token_id] * pad_length
+            labels.extend([-100] * pad_length)
+            position_ids[0] += [0] * pad_length
+            position_ids[1] += [0] * pad_length
+
+        assert len(tokens) == len(labels)
+
+        if not use_fa2:
+            attention_mask = pack_attention_mask.unsqueeze(0).long()
+        else:
+            attention_mask = torch.tensor(0)
+        return {
+            'input_ids': torch.tensor(tokens).long(),
+            'position_ids': torch.tensor(position_ids).long(),
+            'attention_mask': attention_mask,
+            'labels': torch.tensor(labels).long(),
+        }
+
+    def build_inputs_for_train(
+        self,
+        data: Union[Dict, List[Dict]],
+        new_conversation_offset: List[int] = None,
+        chat_format="antglm_chat",
+        is_chat_format=True,  # whether the input is already formatted as chat text
+        use_true_multiturn=False,
+        max_length: int = 2048,
+        rotary_type: str = "1d",
+        left_truncate: bool = True,
+        unidirectional_attention: bool = True,
+        isolation_position_ids: bool = False,
+        padding: bool = True,
+        use_fa2: bool = True,
+        use_packed: bool = True,
+        use_baichuan_packed: bool = False,
+        skip_truncated_turn: bool = False,
+        return_attention_mask: bool = True,
+    ):
+        r"""
+        Build tensor input for model training. If inputs and outputs are list, will pack them.
+
+        Args:
+            inputs (str, List[str], List[Dict], List[List[Dict]]): the input prompts.
+            outputs (str, List[str]): the output responses.
+            new_conversation_offset (List[int]): the offset index of the new conversation turn.
+            is_chat_format (bool): whether the input is already chatml format
+            max_length (int, Optional): the maximum length of the final input ids for training. Default: 2048
+            rotary_type (str, Optional): the rotary type of position embedding. Default: 1d
+            left_truncate (bool, Optional): whether truncate the inputs from left. Default: True
+            use_fa2 (bool, Optional): whether to build attention mask under flash attention 2.
+        """
+        if isinstance(data, List):
+            # chatml list
+            _inputs = []
+            _outputs = []
+            new_conversation_offset = []
+            for _input in data:
+                if use_true_multiturn:
+                    chat = self._chat_from_json(_input, chat_format=chat_format)
+                    chat_data = chat.prompt_pack
+                    new_conversation_offset.append(len(_inputs))
+                    _inputs.extend(chat_data['input'])
+                    _outputs.extend(chat_data['output'])
+                else:
+                    _conversation = _convert_to_conversation(_input)
+                    assert is_assistant(_conversation[-1])
+
+                    _inputs.append(
+                        self.apply_chat_template(_conversation[:-1], tokenize=False, add_generation_prompt=True)
+                    )
+                    _outputs.append(_conversation[-1]['content'])
+
+            return self._build_inputs_for_train(
+                inputs=_inputs,
+                outputs=_outputs,
+                new_conversation_offset=new_conversation_offset,
+                max_length=max_length,
+                rotary_type=rotary_type,
+                left_truncate=left_truncate,
+                unidirectional_attention=unidirectional_attention,
+                isolation_position_ids=isolation_position_ids,
+                padding=padding,
+                use_fa2=use_fa2,
+                use_packed=use_packed,
+                use_baichuan_packed=use_baichuan_packed,
+                skip_truncated_turn=skip_truncated_turn,
+                return_attention_mask=return_attention_mask,
+            )
+        elif isinstance(data, Dict):
+            if 'messages' in data:
+                # chatml format
+                if use_true_multiturn:
+                    chat = self._chat_from_json(data, chat_format=chat_format)
+                    chat_data = chat.prompt_pack
+                else:
+                    _conversation = _convert_to_conversation(data)
+                    assert is_assistant(_conversation[-1])
+
+                    chat_data = {
+                        "input": self.apply_chat_template(
+                            _conversation[:-1], tokenize=False, add_generation_prompt=True
+                        ),
+                        "output": _conversation[-1]['content'],
+                    }
+
+                return self._build_inputs_for_train(
+                    inputs=chat_data['input'],
+                    outputs=chat_data['output'],
+                    max_length=max_length,
+                    rotary_type=rotary_type,
+                    left_truncate=left_truncate,
+                    unidirectional_attention=unidirectional_attention,
+                    isolation_position_ids=isolation_position_ids,
+                    padding=padding,
+                    use_fa2=use_fa2,
+                    use_packed=use_packed,
+                    use_baichuan_packed=use_baichuan_packed,
+                    skip_truncated_turn=skip_truncated_turn,
+                    return_attention_mask=return_attention_mask,
+                )
+            else:
+                inputs = data['input']
+                outputs = data['output']
+
+                if isinstance(inputs, str):
+                    inputs = [inputs]
+                if isinstance(outputs, str):
+                    outputs = [outputs]
+
+                if not is_chat_format and chat_format:
+                    inputs = [
+                        self.apply_chat_template(
+                            [{"role": "HUMAN", "content": item}], tokenize=False, chat_format=chat_format
+                        )
+                        for item in inputs
+                    ]
+
+                return self._build_inputs_for_train(
+                    inputs=inputs,
+                    outputs=outputs,
+                    new_conversation_offset=new_conversation_offset,
+                    max_length=max_length,
+                    rotary_type=rotary_type,
+                    left_truncate=left_truncate,
+                    unidirectional_attention=unidirectional_attention,
+                    isolation_position_ids=isolation_position_ids,
+                    padding=padding,
+                    use_fa2=use_fa2,
+                    use_packed=use_packed,
+                    use_baichuan_packed=use_baichuan_packed,
+                    skip_truncated_turn=skip_truncated_turn,
+                    return_attention_mask=return_attention_mask,
+                )

connector/config.json

@@ -0,0 +1,58 @@
+{
+  "architectures": [
+    "Qwen2ForCausalLM"
+  ],
+  "attention_dropout": 0.0,
+  "bos_token_id": 151643,
+  "dtype": "bfloat16",
+  "eos_token_id": 151645,
+  "hidden_act": "silu",
+  "hidden_size": 1536,
+  "initializer_range": 0.02,
+  "intermediate_size": 8960,
+  "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"
+  ],
+  "max_position_embeddings": 32768,
+  "max_window_layers": 21,
+  "model_type": "qwen2",
+  "num_attention_heads": 12,
+  "num_hidden_layers": 28,
+  "num_key_value_heads": 2,
+  "rms_norm_eps": 1e-06,
+  "rope_scaling": null,
+  "rope_theta": 1000000.0,
+  "sliding_window": null,
+  "tie_word_embeddings": true,
+  "transformers_version": "4.57.1",
+  "use_cache": true,
+  "use_sliding_window": false,
+  "vocab_size": 151936
+}

connector/generation_config.json

@@ -0,0 +1,14 @@
+{
+  "bos_token_id": 151643,
+  "do_sample": true,
+  "eos_token_id": [
+    151645,
+    151643
+  ],
+  "pad_token_id": 151643,
+  "repetition_penalty": 1.1,
+  "temperature": 0.7,
+  "top_k": 20,
+  "top_p": 0.8,
+  "transformers_version": "4.57.1"
+}

mllm/config.json

@@ -0,0 +1,210 @@
+{
+  "architectures": [
+    "BailingMM2NativeForConditionalGeneration"
+  ],
+  "audio_config": null,
+  "llm_config": {
+    "_name_or_path": "",
+    "add_cross_attention": false,
+    "architectures": [
+      "BailingMoeV2ForCausalLM"
+    ],
+    "attention_dropout": 0.0,
+    "bad_words_ids": null,
+    "begin_suppress_tokens": null,
+    "bos_token_id": null,
+    "chunk_size_feed_forward": 0,
+    "cross_attention_hidden_size": null,
+    "decoder_start_token_id": null,
+    "diversity_penalty": 0.0,
+    "do_sample": false,
+    "early_stopping": false,
+    "embedding_dropout": 0.0,
+    "encoder_no_repeat_ngram_size": 0,
+    "eos_token_id": 156895,
+    "exponential_decay_length_penalty": null,
+    "finetuning_task": null,
+    "first_k_dense_replace": 1,
+    "forced_bos_token_id": null,
+    "forced_eos_token_id": null,
+    "head_dim": 128,
+    "hidden_act": "silu",
+    "hidden_size": 2048,
+    "id2label": {
+      "0": "LABEL_0",
+      "1": "LABEL_1"
+    },
+    "image_end_token": 157159,
+    "image_patch_token": 157157,
+    "image_start_token": 157158,
+    "video_end_token": 157161,
+    "video_patch_token": 157175,
+    "video_start_token": 157160,
+    "initializer_range": 0.006,
+    "intermediate_size": 5120,
+    "is_decoder": false,
+    "is_encoder_decoder": false,
+    "label2id": {
+      "LABEL_0": 0,
+      "LABEL_1": 1
+    },
+    "length_penalty": 1.0,
+    "max_length": 20,
+    "max_position_embeddings": 32768,
+    "max_window_layers": 28,
+    "min_length": 0,
+    "model_type": "bailing_moe_v2",
+    "moe_intermediate_size": 512,
+    "moe_router_topk_scaling_factor": 2.5,
+    "n_group": 8,
+    "no_repeat_ngram_size": 0,
+    "norm_head": false,
+    "norm_softmax": false,
+    "norm_topk_prob": true,
+    "num_attention_heads": 16,
+    "num_beam_groups": 1,
+    "num_beams": 1,
+    "num_experts": 256,
+    "num_experts_per_tok": 8,
+    "num_hidden_layers": 20,
+    "num_key_value_heads": 4,
+    "num_return_sequences": 1,
+    "num_shared_experts": 1,
+    "output_attentions": false,
+    "output_dropout": 0.0,
+    "output_hidden_states": false,
+    "output_router_logits": false,
+    "output_scores": false,
+    "pad_token_id": 156892,
+    "partial_rotary_factor": 0.5,
+    "prefix": null,
+    "pretraining_tp": 1,
+    "problem_type": null,
+    "pruned_heads": {},
+    "remove_invalid_values": false,
+    "repetition_penalty": 1.0,
+    "return_dict": true,
+    "return_dict_in_generate": false,
+    "rms_norm_eps": 1e-06,
+    "rope_scaling": {
+      "factor": null,
+      "type": "video_rope"
+    },
+    "rope_theta": 600000,
+    "routed_scaling_factor": 2.5,
+    "router_type": "MultiRouter",
+    "sep_token_id": null,
+    "sliding_window": 4096,
+    "spatial_merge_size": 2,
+    "suppress_tokens": null,
+    "task_specific_params": null,
+    "temperature": 1.0,
+    "tf_legacy_loss": false,
+    "tie_encoder_decoder": false,
+    "tie_word_embeddings": false,
+    "tokenizer_class": null,
+    "tokens_per_second": 2,
+    "top_k": 50,
+    "top_p": 1.0,
+    "topk_group": 4,
+    "torch_dtype": "bfloat16",
+    "torchscript": false,
+    "typical_p": 1.0,
+    "use_bfloat16": false,
+    "use_bias": false,
+    "use_cache": true,
+    "use_expert_bias": true,
+    "use_qkv_bias": false,
+    "use_sliding_window": false,
+    "vocab_size": 157184
+  },
+  "mlp_depth": 2,
+  "model_type": "bailingmm_moe_v2_lite",
+  "torch_dtype": "bfloat16",
+  "transformers_version": "4.53.0.dev0",
+  "vision_config": {
+    "_name_or_path": "",
+    "add_cross_attention": false,
+    "architectures": [
+      "Qwen2_5_VisionTransformer"
+    ],
+    "bad_words_ids": null,
+    "begin_suppress_tokens": null,
+    "bos_token_id": null,
+    "chunk_size_feed_forward": 0,
+    "cross_attention_hidden_size": null,
+    "decoder_start_token_id": null,
+    "depth": 32,
+    "diversity_penalty": 0.0,
+    "do_sample": false,
+    "early_stopping": false,
+    "encoder_no_repeat_ngram_size": 0,
+    "eos_token_id": null,
+    "exponential_decay_length_penalty": null,
+    "finetuning_task": null,
+    "forced_bos_token_id": null,
+    "forced_eos_token_id": null,
+    "fullatt_block_indexes": [
+      7,
+      15,
+      23,
+      31
+    ],
+    "hidden_act": "silu",
+    "hidden_size": 1280,
+    "id2label": {
+      "0": "LABEL_0",
+      "1": "LABEL_1"
+    },
+    "in_channels": 3,
+    "in_chans": 3,
+    "intermediate_size": 3456,
+    "is_decoder": false,
+    "is_encoder_decoder": false,
+    "label2id": {
+      "LABEL_0": 0,
+      "LABEL_1": 1
+    },
+    "length_penalty": 1.0,
+    "max_length": 20,
+    "min_length": 0,
+    "model_type": "qwen2_5_vit",
+    "no_repeat_ngram_size": 0,
+    "num_beam_groups": 1,
+    "num_beams": 1,
+    "num_heads": 16,
+    "num_return_sequences": 1,
+    "out_hidden_size": 8192,
+    "output_attentions": false,
+    "output_hidden_states": false,
+    "output_scores": false,
+    "pad_token_id": null,
+    "patch_size": 14,
+    "prefix": null,
+    "problem_type": null,
+    "pruned_heads": {},
+    "remove_invalid_values": false,
+    "repetition_penalty": 1.0,
+    "return_dict": true,
+    "return_dict_in_generate": false,
+    "sep_token_id": null,
+    "spatial_merge_size": 2,
+    "spatial_patch_size": 14,
+    "suppress_tokens": null,
+    "task_specific_params": null,
+    "temperature": 1.0,
+    "temporal_patch_size": 2,
+    "tf_legacy_loss": false,
+    "tie_encoder_decoder": false,
+    "tie_word_embeddings": true,
+    "tokenizer_class": null,
+    "tokens_per_second": 2,
+    "top_k": 50,
+    "top_p": 1.0,
+    "torch_dtype": "bfloat16",
+    "torchscript": false,
+    "typical_p": 1.0,
+    "use_bfloat16": false,
+    "window_size": 112
+  }
+}

mllm/preprocessor_config.json

@@ -0,0 +1,23 @@
+{
+  "min_pixels": 451584,
+  "max_pixels": 451584,
+  "patch_size": 14,
+  "temporal_patch_size": 2,
+  "merge_size": 2,
+  "image_mean": [
+    0.48145466,
+    0.4578275,
+    0.40821073
+  ],
+  "image_std": [
+    0.26862954,
+    0.26130258,
+    0.27577711
+  ],
+  "image_token": "<image>",
+  "video_token": "<video>",
+  "image_processor_type": "BailingMM2ImageProcessor",
+  "return_attention_mask": true,
+  "padding_side": "right",
+  "padding_value": 0.0
+}

mllm/special_tokens_map.json

@@ -0,0 +1,44 @@
+{
+  "additional_special_tokens": [
+    "<|arithmetic_end|>",
+    "<role>",
+    "<|number_end|>",
+    "<|arithmetic_start|>",
+    "</role>",
+    "<|number_start|>",
+    "<|role_end|>",
+    "<tool_call>",
+    "</tool_call>",
+    "<tool_response>",
+    "</tool_response>"
+  ],
+  "bos_token": {
+    "content": "<|startoftext|>",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  },
+  "cls_token": {
+    "content": "[CLS]",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  },
+  "eos_token": {
+    "content": "<|role_end|>",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  },
+  "gmask_token": {
+    "content": "[gMASK]",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  },
+  "pad_token": "<|role_end|>"
+}

mllm/tokenizer_config.json

@@ -0,0 +1,2334 @@
+{
+  "add_bos_token": false,
+  "add_eos_token": false,
+  "added_tokens_decoder": {
+    "156891": {
+      "content": "<|startoftext|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "156892": {
+      "content": "<|endoftext|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "156893": {
+      "content": "[CLS]",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "156894": {
+      "content": "[gMASK]",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "156895": {
+      "content": "<|role_end|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "156896": {
+      "content": "<tool_call>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "156897": {
+      "content": "</tool_call>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "156898": {
+      "content": "<tool_response>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "156899": {
+      "content": "</tool_response>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "156900": {
+      "content": "<|reserved_token_5|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "156901": {
+      "content": "<|reserved_token_6|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "156902": {
+      "content": "<|reserved_token_7|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "156903": {
+      "content": "<|reserved_token_8|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "156904": {
+      "content": "<|reserved_token_9|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "156905": {
+      "content": "<|reserved_token_10|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "156906": {
+      "content": "<|reserved_token_11|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "156907": {
+      "content": "<|reserved_token_12|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "156908": {
+      "content": "<|reserved_token_13|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "156909": {
+      "content": "<|reserved_token_14|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "156910": {
+      "content": "<|reserved_token_15|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "156911": {
+      "content": "<|reserved_token_16|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "156912": {
+      "content": "<|reserved_token_17|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "156913": {
+      "content": "<|reserved_token_18|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "156914": {
+      "content": "<|reserved_token_19|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "156915": {
+      "content": "<|reserved_token_20|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "156916": {
+      "content": "<|reserved_token_21|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "156917": {
+      "content": "<|reserved_token_22|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "156918": {
+      "content": "<|reserved_token_23|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "156919": {
+      "content": "<|reserved_token_24|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "156920": {
+      "content": "<|reserved_token_25|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "156921": {
+      "content": "<|reserved_token_26|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "156922": {
+      "content": "<|reserved_token_27|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "156923": {
+      "content": "<|reserved_token_28|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "156924": {
+      "content": "<|reserved_token_29|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "156925": {
+      "content": "<|reserved_token_30|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "156926": {
+      "content": "<|reserved_token_31|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "156927": {
+      "content": "<|reserved_token_32|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "156928": {
+      "content": "<|reserved_token_33|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "156929": {
+      "content": "<|reserved_token_34|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "156930": {
+      "content": "<|reserved_token_35|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "156931": {
+      "content": "<|reserved_token_36|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "156932": {
+      "content": "<|reserved_token_37|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "156933": {
+      "content": "<|reserved_token_38|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "156934": {
+      "content": "<|reserved_token_39|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "156935": {
+      "content": "<|reserved_token_40|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "156936": {
+      "content": "<|reserved_token_41|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "156937": {
+      "content": "<|reserved_token_42|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "156938": {
+      "content": "<|reserved_token_43|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "156939": {
+      "content": "<|reserved_token_44|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "156940": {
+      "content": "<|reserved_token_45|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "156941": {
+      "content": "<|reserved_token_46|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "156942": {
+      "content": "<|reserved_token_47|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "156943": {
+      "content": "<|reserved_token_48|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "156944": {
+      "content": "<|reserved_token_49|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "156945": {
+      "content": "<|reserved_token_50|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "156946": {
+      "content": "<|reserved_token_51|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "156947": {
+      "content": "<|reserved_token_52|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "156948": {
+      "content": "<|reserved_token_53|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "156949": {
+      "content": "<|reserved_token_54|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "156950": {
+      "content": "<|reserved_token_55|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "156951": {
+      "content": "<|reserved_token_56|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "156952": {
+      "content": "<|reserved_token_57|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "156953": {
+      "content": "<|reserved_token_58|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "156954": {
+      "content": "<|reserved_token_59|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "156955": {
+      "content": "<|reserved_token_60|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "156956": {
+      "content": "<|reserved_token_61|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "156957": {
+      "content": "<|reserved_token_62|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "156958": {
+      "content": "<|reserved_token_63|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "156959": {
+      "content": "<|reserved_token_64|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "156960": {
+      "content": "<|reserved_token_65|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "156961": {
+      "content": "<|reserved_token_66|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "156962": {
+      "content": "<|reserved_token_67|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "156963": {
+      "content": "<|reserved_token_68|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "156964": {
+      "content": "<|reserved_token_69|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "156965": {
+      "content": "<|reserved_token_70|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "156966": {
+      "content": "<|reserved_token_71|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "156967": {
+      "content": "<|reserved_token_72|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "156968": {
+      "content": "<|reserved_token_73|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "156969": {
+      "content": "<|reserved_token_74|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "156970": {
+      "content": "<|reserved_token_75|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "156971": {
+      "content": "<|reserved_token_76|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "156972": {
+      "content": "<|reserved_token_77|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "156973": {
+      "content": "<|reserved_token_78|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "156974": {
+      "content": "<|reserved_token_79|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "156975": {
+      "content": "<|reserved_token_80|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "156976": {
+      "content": "<|reserved_token_81|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "156977": {
+      "content": "<|reserved_token_82|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "156978": {
+      "content": "<|reserved_token_83|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "156979": {
+      "content": "<|reserved_token_84|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "156980": {
+      "content": "<|reserved_token_85|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "156981": {
+      "content": "<|reserved_token_86|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "156982": {
+      "content": "<|reserved_token_87|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "156983": {
+      "content": "<|reserved_token_88|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "156984": {
+      "content": "<|reserved_token_89|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "156985": {
+      "content": "<|reserved_token_90|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "156986": {
+      "content": "<|reserved_token_91|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "156987": {
+      "content": "<|reserved_token_92|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "156988": {
+      "content": "<|reserved_token_93|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "156989": {
+      "content": "<|reserved_token_94|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "156990": {
+      "content": "<|reserved_token_95|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "156991": {
+      "content": "<|reserved_token_96|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "156992": {
+      "content": "<|reserved_token_97|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "156993": {
+      "content": "<|reserved_token_98|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "156994": {
+      "content": "<|reserved_token_99|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "156995": {
+      "content": "<|reserved_token_100|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "156996": {
+      "content": "<|reserved_token_101|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "156997": {
+      "content": "<|reserved_token_102|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "156998": {
+      "content": "<|reserved_token_103|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "156999": {
+      "content": "<|reserved_token_104|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "157000": {
+      "content": "<|reserved_token_105|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "157001": {
+      "content": "<|reserved_token_106|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "157002": {
+      "content": "<|reserved_token_107|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "157003": {
+      "content": "<|reserved_token_108|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "157004": {
+      "content": "<|reserved_token_109|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "157005": {
+      "content": "<|reserved_token_110|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "157006": {
+      "content": "<|reserved_token_111|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "157007": {
+      "content": "<|reserved_token_112|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "157008": {
+      "content": "<|reserved_token_113|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "157009": {
+      "content": "<|reserved_token_114|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "157010": {
+      "content": "<|reserved_token_115|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "157011": {
+      "content": "<|reserved_token_116|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "157012": {
+      "content": "<|reserved_token_117|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "157013": {
+      "content": "<|reserved_token_118|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "157014": {
+      "content": "<|reserved_token_119|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "157015": {
+      "content": "<|reserved_token_120|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "157016": {
+      "content": "<|reserved_token_121|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "157017": {
+      "content": "<|reserved_token_122|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "157018": {
+      "content": "<|reserved_token_123|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "157019": {
+      "content": "<|reserved_token_124|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "157020": {
+      "content": "<|reserved_token_125|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "157021": {
+      "content": "<|reserved_token_126|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "157022": {
+      "content": "<|reserved_token_127|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "157023": {
+      "content": "<|reserved_token_128|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "157024": {
+      "content": "<|reserved_token_129|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "157025": {
+      "content": "<|reserved_token_130|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "157026": {
+      "content": "<|reserved_token_131|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "157027": {
+      "content": "<|reserved_token_132|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "157028": {
+      "content": "<|reserved_token_133|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "157029": {
+      "content": "<|reserved_token_134|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "157030": {
+      "content": "<|reserved_token_135|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "157031": {
+      "content": "<|reserved_token_136|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "157032": {
+      "content": "<|reserved_token_137|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "157033": {
+      "content": "<|reserved_token_138|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "157034": {
+      "content": "<|reserved_token_139|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "157035": {
+      "content": "<|reserved_token_140|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "157036": {
+      "content": "<|reserved_token_141|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "157037": {
+      "content": "<|reserved_token_142|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "157038": {
+      "content": "<|reserved_token_143|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "157039": {
+      "content": "<|reserved_token_144|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "157040": {
+      "content": "<|reserved_token_145|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "157041": {
+      "content": "<|reserved_token_146|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "157042": {
+      "content": "<|reserved_token_147|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "157043": {
+      "content": "<|reserved_token_148|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "157044": {
+      "content": "<|reserved_token_149|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "157045": {
+      "content": "<|reserved_token_150|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "157046": {
+      "content": "<|reserved_token_151|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "157047": {
+      "content": "<|reserved_token_152|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "157048": {
+      "content": "<|reserved_token_153|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "157049": {
+      "content": "<|reserved_token_154|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "157050": {
+      "content": "<|reserved_token_155|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "157051": {
+      "content": "<|reserved_token_156|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "157052": {
+      "content": "<|reserved_token_157|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "157053": {
+      "content": "<|reserved_token_158|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "157054": {
+      "content": "<|reserved_token_159|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "157055": {
+      "content": "<|reserved_token_160|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "157056": {
+      "content": "<|reserved_token_161|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "157057": {
+      "content": "<|reserved_token_162|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "157058": {
+      "content": "<|reserved_token_163|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "157059": {
+      "content": "<|reserved_token_164|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "157060": {
+      "content": "<|reserved_token_165|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "157061": {
+      "content": "<|reserved_token_166|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "157062": {
+      "content": "<|reserved_token_167|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "157063": {
+      "content": "<|reserved_token_168|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "157064": {
+      "content": "<|reserved_token_169|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "157065": {
+      "content": "<|reserved_token_170|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "157066": {
+      "content": "<|reserved_token_171|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "157067": {
+      "content": "<|reserved_token_172|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "157068": {
+      "content": "<|reserved_token_173|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "157069": {
+      "content": "<|reserved_token_174|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "157070": {
+      "content": "<|reserved_token_175|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "157071": {
+      "content": "<|reserved_token_176|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "157072": {
+      "content": "<|reserved_token_177|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "157073": {
+      "content": "<|reserved_token_178|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "157074": {
+      "content": "<|reserved_token_179|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "157075": {
+      "content": "<|reserved_token_180|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "157076": {
+      "content": "<|reserved_token_181|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "157077": {
+      "content": "<|reserved_token_182|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "157078": {
+      "content": "<|reserved_token_183|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "157079": {
+      "content": "<|reserved_token_184|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "157080": {
+      "content": "<|reserved_token_185|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "157081": {
+      "content": "<|reserved_token_186|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "157082": {
+      "content": "<|reserved_token_187|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "157083": {
+      "content": "<|reserved_token_188|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "157084": {
+      "content": "<|reserved_token_189|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "157085": {
+      "content": "<|reserved_token_190|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "157086": {
+      "content": "<|reserved_token_191|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "157087": {
+      "content": "<|reserved_token_192|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "157088": {
+      "content": "<|reserved_token_193|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "157089": {
+      "content": "<|reserved_token_194|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "157090": {
+      "content": "<|reserved_token_195|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "157091": {
+      "content": "<|reserved_token_196|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "157092": {
+      "content": "<|reserved_token_197|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "157093": {
+      "content": "<|reserved_token_198|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "157094": {
+      "content": "<|reserved_token_199|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "157095": {
+      "content": "<|reserved_token_200|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "157096": {
+      "content": "<|reserved_token_201|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "157097": {
+      "content": "<|reserved_token_202|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "157098": {
+      "content": "<|reserved_token_203|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "157099": {
+      "content": "<|reserved_token_204|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "157100": {
+      "content": "<|reserved_token_205|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "157101": {
+      "content": "<|reserved_token_206|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "157102": {
+      "content": "<|reserved_token_207|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "157103": {
+      "content": "<|reserved_token_208|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "157104": {
+      "content": "<|reserved_token_209|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "157105": {
+      "content": "<|reserved_token_210|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "157106": {
+      "content": "<|reserved_token_211|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "157107": {
+      "content": "<|reserved_token_212|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "157108": {
+      "content": "<|reserved_token_213|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "157109": {
+      "content": "<|reserved_token_214|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "157110": {
+      "content": "<|reserved_token_215|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "157111": {
+      "content": "<|reserved_token_216|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "157112": {
+      "content": "<|reserved_token_217|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "157113": {
+      "content": "<|reserved_token_218|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "157114": {
+      "content": "<|reserved_token_219|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "157115": {
+      "content": "<|reserved_token_220|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "157116": {
+      "content": "<|reserved_token_221|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "157117": {
+      "content": "<|reserved_token_222|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "157118": {
+      "content": "<|reserved_token_223|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "157119": {
+      "content": "<|reserved_token_224|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "157120": {
+      "content": "<|reserved_token_225|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "157121": {
+      "content": "<|reserved_token_226|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "157122": {
+      "content": "<|reserved_token_227|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "157123": {
+      "content": "<|reserved_token_228|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "157124": {
+      "content": "<|reserved_token_229|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "157125": {
+      "content": "<|reserved_token_230|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "157126": {
+      "content": "<|reserved_token_231|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "157127": {
+      "content": "<|reserved_token_232|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "157128": {
+      "content": "<|reserved_token_233|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "157129": {
+      "content": "<|reserved_token_234|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "157130": {
+      "content": "<|reserved_token_235|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "157131": {
+      "content": "<|reserved_token_236|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "157132": {
+      "content": "<|reserved_token_237|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "157133": {
+      "content": "<|reserved_token_238|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "157134": {
+      "content": "<|reserved_token_239|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "157135": {
+      "content": "<|reserved_token_240|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "157136": {
+      "content": "<|reserved_token_241|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "157137": {
+      "content": "<|reserved_token_242|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "157138": {
+      "content": "<|reserved_token_243|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "157139": {
+      "content": "<|reserved_token_244|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "157140": {
+      "content": "<|reserved_token_245|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "157141": {
+      "content": "<|reserved_token_246|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "157142": {
+      "content": "<|reserved_token_247|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "157143": {
+      "content": "<|reserved_token_248|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "157144": {
+      "content": "<|reserved_token_249|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "157145": {
+      "content": "<|reserved_token_250|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "157146": {
+      "content": "<|reserved_token_251|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "157147": {
+      "content": "<|reserved_token_252|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "157148": {
+      "content": "<|reserved_token_253|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "157149": {
+      "content": "<|reserved_token_254|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "157150": {
+      "content": "<|reserved_token_255|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "157151": {
+      "content": "<role>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "157152": {
+      "content": "</role>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "157153": {
+      "content": "<|arithmetic_end|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "157154": {
+      "content": "<|number_end|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "157155": {
+      "content": "<|arithmetic_start|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "157156": {
+      "content": "<|number_start|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "157157": {
+      "content": "<imagePatch>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "157158": {
+      "content": "<image>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "157159": {
+      "content": "</image>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "157160": {
+      "content": "<video>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "157161": {
+      "content": "</video>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "157162": {
+      "content": "<gen_imagePatch>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "157163": {
+      "content": "<gen_image>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "157164": {
+      "content": "</gen_image>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "157165": {
+      "content": "<imageHere>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "157166": {
+      "content": "<end_of_chunk>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "157167": {
+      "content": "<end_of_audio>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "157168": {
+      "content": "<audioPatch>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "157169": {
+      "content": "<audio>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "157170": {
+      "content": "</audio>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "157171": {
+      "content": "<gen_audioPatch>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "157172": {
+      "content": "<gen_audio>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "157173": {
+      "content": "</gen_audio>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "157174": {
+      "content": "<audioHere>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "157175": {
+      "content": "<framePatch>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "157176": {
+      "content": "<text>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "157177": {
+      "content": "<asr>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "157178": {
+      "content": "<tts>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    }
+  },
+  "additional_special_tokens": [
+    "<|arithmetic_end|>",
+    "<role>",
+    "<|number_end|>",
+    "<|arithmetic_start|>",
+    "</role>",
+    "<|number_start|>",
+    "<|role_end|>",
+    "<tool_call>",
+    "</tool_call>",
+    "<tool_response>",
+    "</tool_response>"
+  ],
+  "bos_token": "<|startoftext|>",
+  "chat_template": "{% set thinking_option = 'off' %}\n{{- '<role>SYSTEM</role>' }}\n{%- if messages[0].role == 'system' %}\n    {{- messages[0].content + '\\n\\n' }}\n{%- endif %}\n{%- if tools %}\n    {{- \"# Tools\\n\\nYou may call one or more functions to assist with the user query.\\n\\nYou are provided with function signatures within <tools></tools> XML tags:\\n<tools>\" }}\n    {%- for tool in tools %}\n        {{- \"\\n\" }}\n        {{- tool | tojson }}\n    {%- endfor %}\n    {{- \"\\n</tools>\\n\\nFor each function call, return a json object with function name and arguments within <tool_call></tool_call> XML tags:\\n<tool_call>\\n{\\\"name\\\": <function-name>, \\\"arguments\\\": <args-json-object>}\\n</tool_call>\\n\" }}\n{%- endif %}\n{{- 'detailed thinking ' + thinking_option + '<|role_end|>' }}\n{%- set ns = namespace(multi_step_tool=true, last_query_index=messages|length - 1) %}\n{%- for message in messages[::-1] %}\n    {%- set index = (messages|length - 1) - loop.index0 %}\n    {%- if ns.multi_step_tool and message.role == \"user\" and message.content is string and not(message.content.startswith('<tool_response>') and message.content.endswith('</tool_response>')) %}\n        {%- set ns.multi_step_tool = false %}\n        {%- set ns.last_query_index = index %}\n    {%- endif %}\n{%- endfor %}\n{%- for message in messages %}\n    {%- if message.content is string %}\n        {%- set content = message.content %}\n    {%- else %}\n        {%- set content = '' %}\n    {%- endif %}\n    {%- if message.role == \"user\" %}\n        {{- '<role>HUMAN</role>' + message.content + '<|role_end|>' }}\n    {%- elif message.role == \"system\" and not loop.first %}\n        {{- '<role>SYSTEM</role>' + message.content + '<|role_end|>' }}\n    {%- elif message.role == \"assistant\" %}\n        {%- set reasoning_content = '' %}\n        {%- if message.reasoning_content is string %}\n            {%- set reasoning_content = message.reasoning_content %}\n        {%- else %}\n            {%- if '</think>' in content %}\n                {%- set reasoning_content = content.split('</think>')[0].rstrip('\\n').split('<think>')[-1].lstrip('\\n') %}\n                {%- set content = content.split('</think>')[-1].lstrip('\\n') %}\n            {%- endif %}\n        {%- endif %}\n        {%- if loop.index0 > ns.last_query_index %}\n            {%- if reasoning_content %}\n                {{- '<role>ASSISTANT</role>' + '\\n<think>\\n' + reasoning_content.strip('\\n') + '\\n</think>\\n\\n' + content.lstrip('\\n') }}\n            {%- else %}\n                {{- '<role>ASSISTANT</role>' + content }}\n            {%- endif %}\n        {%- else %}\n            {{- '<role>ASSISTANT</role>' + content }}\n        {%- endif %}\n        {%- if message.tool_calls %}\n            {%- for tool_call in message.tool_calls %}\n                {%- if (loop.first and content) or (not loop.first) %}\n                    {{- '\\n' }}\n                {%- endif %}\n                {%- if tool_call.function %}\n                    {%- set tool_call = tool_call.function %}\n                {%- endif %}\n                {{- '<tool_call>\\n{\"name\": \"' }}\n                {{- tool_call.name }}\n                {{- '\", \"arguments\": ' }}\n                {%- if tool_call.arguments is string %}\n                    {{- tool_call.arguments }}\n                {%- else %}\n                    {{- tool_call.arguments | tojson }}\n                {%- endif %}\n                {{- '}\\n</tool_call>' }}\n            {%- endfor %}\n        {%- endif %}\n        {{- '<|role_end|>' }}\n    {%- elif message.role == \"tool\" %}\n        {%- if loop.first or (messages[loop.index0 - 1].role != \"tool\") %}\n            {{- '<role>OBSERVATION</role>' }}\n        {%- endif %}\n        {{- '\\n<tool_response>\\n' }}\n        {{- content }}\n        {{- '\\n</tool_response>' }}\n        {%- if loop.last or (messages[loop.index0 + 1].role != \"tool\") %}\n            {{- '<|role_end|>' }}\n        {%- endif %}\n    {%- endif %}\n{%- endfor %}\n{%- if add_generation_prompt %}\n    {{- '<role>ASSISTANT</role>' }}\n{%- endif %}",
+  "clean_up_tokenization_spaces": false,
+  "cls_token": "[CLS]",
+  "eos_token": "<|role_end|>",
+  "extra_special_tokens": {},
+  "fast_tokenizer": true,
+  "gmask_token": "[gMASK]",
+  "merges_file": null,
+  "model_max_length": 1000000000000000019884624838656,
+  "pad_token": "<|role_end|>"
+}

mlp/config.json

@@ -0,0 +1,17 @@
+{
+  "use_identity_mlp": true,
+  "use_vlm_directvlm_condition": true,
+  "use_learnable_token_condition": true,
+  "diffusion_c_input_dim": 2560,
+  "text_encoder_norm": false,
+  "connector_norm": false,
+  "img_gen_scales": [
+    16
+  ],
+  "selected_hidden_states_layers": [
+    5,
+    12,
+    20
+  ],
+  "diffusion_inner_dim": 3840
+}

samples/cabin_upstream.json

@@ -0,0 +1,33 @@
+{
+  "canvas_settings": {
+    "aspect_ratio": "1:1, 2048 × 2048 px",
+    "ambient_lighting": "Four season-specific lighting conditions unified by one fixed eye-level sun direction, consistent horizon at 48% canvas height, matching atmospheric depth, and natural cinematic exposure.",
+    "image_style": "High-resolution cinematic photorealistic environmental concept sheet; four equal edge-to-edge vertical panels with crisp shared boundaries and no gutters, frames, extra title, or decorative inter-panel objects. Fixed 35 mm lens, eye-level three-quarter front-right camera angle, identical perspective and cabin registration in every panel. Each panel uses restrained bold uppercase condensed sans-serif labeling at the bottom."
+  },
+  "layers": [
+    {
+      "description": "Leftmost spring panel. The same small one-story rectangular log cabin appears at a fixed local position and scale: horizontal dark-brown logs, pale stone foundation, centered vertical-plank front door, two square four-pane front windows, one matching right-side window, and a charcoal standing-seam gabled roof with identical ridge height and overhang. Fresh pale-green grass, budding branches, white and pink blossoms, rain-darkened soil, shallow puddles, wet roof highlights, fine recent-rain droplets, and soft mist fill the landscape beneath an overcast clearing sky. Render the exact label \"SPRING\" centered near the bottom in bold condensed uppercase white sans-serif with a subtle dark shadow.",
+      "coordinates": "cx: 0.125, cy: 0.500, w: 0.250, h: 1",
+      "hierarchy_and_relation": "Complete first panel; aligned to the shared horizon and fixed cabin registration, flush with the canvas left edge and directly abutting the second panel without a gutter.",
+      "color_specs": ["#A9C7D2", "#D8E4DE", "#7FAF69", "#B7D58B", "#F2D7DE", "#F4F1E8", "#5A3B28", "#2C3032", "#7B7468", "#FFFFFF"]
+    },
+    {
+      "description": "Second summer panel. Repeat the cabin with exactly the same architecture, camera angle, local placement, dimensions, centered plank door, front and side windows, stone foundation, roof pitch, ridge height, and overhang. Surround it with lush meadow grass, dense layered deciduous foliage, full tree canopies, small sunlit shrubs, and dry natural ground; use a bright clear sky, strong warm sunlight from the same direction, crisp leaf highlights, controlled shadows, and mild atmospheric haze. Render the exact label \"SUMMER\" centered near the bottom in bold condensed uppercase white sans-serif with a subtle dark shadow.",
+      "coordinates": "cx: 0.375, cy: 0.500, w: 0.250, h: 1",
+      "hierarchy_and_relation": "Complete second panel; aligned to the shared horizon and fixed cabin registration, directly abutting the first and third panels without gutters.",
+      "color_specs": ["#4FA9D8", "#BDE8F2", "#26733E", "#4F9A45", "#79B84A", "#D8C66A", "#5A3B28", "#2C3032", "#7B7468", "#FFFFFF"]
+    },
+    {
+      "description": "Third autumn panel. Repeat the cabin with exactly the same architecture, camera angle, local placement, dimensions, centered plank door, front and side windows, stone foundation, roof pitch, ridge height, and overhang. Cover the landscape with amber grass, rust-orange and golden deciduous foliage, scattered fallen leaves, and a few leaves drifting naturally in the air. Use warm late-afternoon sunlight from the same direction, long soft shadows, copper rim light on logs and roof, a pale warm sky, and subtle atmospheric depth. Render the exact label \"AUTUMN\" centered near the bottom in bold condensed uppercase white sans-serif with a subtle dark shadow.",
+      "coordinates": "cx: 0.625, cy: 0.500, w: 0.250, h: 1",
+      "hierarchy_and_relation": "Complete third panel; aligned to the shared horizon and fixed cabin registration, directly abutting the second and fourth panels without gutters.",
+      "color_specs": ["#E6B06A", "#F0D6A2", "#D66A2C", "#A84324", "#E2A62B", "#80613B", "#5A3B28", "#2C3032", "#7B7468", "#FFFFFF"]
+    },
+    {
+      "description": "Rightmost winter panel. Repeat the cabin with exactly the same architecture, camera angle, local placement, dimensions, centered plank door, front and side windows, stone foundation, roof pitch, ridge height, and overhang. Blanket the ground and roof with natural snow accumulation while keeping doors and windows readable; add bare deciduous branches, snow-laden low shrubs, faint footprints away from the cabin, and sparse fine snowfall. Use cool blue ambient light, a pale clouded sky, soft low-contrast shadows from the same direction, icy edge highlights, and slight winter haze. Render the exact label \"WINTER\" centered near the bottom in bold condensed uppercase white sans-serif with a subtle dark shadow.",
+      "coordinates": "cx: 0.875, cy: 0.500, w: 0.250, h: 1",
+      "hierarchy_and_relation": "Complete fourth panel; aligned to the shared horizon and fixed cabin registration, directly abutting the third panel without a gutter and flush with the canvas right edge.",
+      "color_specs": ["#AFC8DE", "#DCE8F1", "#F4F7F8", "#C6D8E5", "#8198AA", "#6D5949", "#5A3B28", "#2C3032", "#7B7468", "#FFFFFF"]
+    }
+  ]
+}

samples/e2e_daily_grind.caption.txt

@@ -0,0 +1 @@
+A landing page hero section for a coffee subscription service called Daily Grind: the headline Fresh beans every Monday, a short tagline, a Start your subscription button, and a photo of latte art on the right. Warm earthy palette.

samples/e2e_daily_grind.json

@@ -0,0 +1,102 @@
+{
+  "canvas_settings": {
+    "aspect_ratio": "1:1",
+    "ambient_lighting": "Warm, softly diffused natural light from the upper left, producing gentle cream highlights, muted brown shadows, and a cozy editorial café atmosphere.",
+    "image_style": "Polished square landing-page hero with a warm minimalist editorial layout; the left side is clean graphic space and the right side is a softly blurred lifestyle photograph."
+  },
+  "layers": [
+    {
+      "description": "Full square hero background divided visually into a warm off-white paper-like left field and a softly defocused café-photo field on the right. The left field has subtle beige grain and a faint horizontal tonal transition near the lower third; the right field shows blurred amber-brown wood, pale reflections, and indistinct dark plant silhouettes.",
+      "coordinates": "cx: 0.500, cy: 0.500, w: 1.000, h: 1.000",
+      "hierarchy_and_relation": "Base layer beneath all typography, controls, icons, and foreground drink elements; the left and right fields meet near the vertical center without a hard border.",
+      "color_specs": [
+        "#F7F1E8",
+        "#E9DDCB",
+        "#C8955D",
+        "#6A442D",
+        "#2F2118"
+      ]
+    },
+    {
+      "description": "Upper-left brand lockup: a small dark-brown outlined coffee-bean icon tilted diagonally, followed by the widely tracked uppercase wordmark 'DAILY GRIND'. The icon and wordmark share one horizontal row, with the wordmark aligned to the icon's right and vertically centered against it. The lettering is compact, bold, sans-serif, and aligned horizontally with generous surrounding whitespace.",
+      "coordinates": "cx: 0.151, cy: 0.082, w: 0.205, h: 0.045",
+      "hierarchy_and_relation": "Sits directly over the off-white background at the top-left and remains visually separate from the headline below.",
+      "color_specs": [
+        "#3B2418",
+        "#F7F1E8"
+      ]
+    },
+    {
+      "description": "Large two-line hero headline in a heavy rounded sans-serif. The first line reads 'Fresh beans' and the second line reads 'every Monday'. The two rendered-text occurrences form a left-aligned vertical stack, with the second row directly beneath the first and both sharing the same left edge. Both rows are left aligned, dark espresso brown, tightly stacked, and occupy most of the left half.",
+      "coordinates": "cx: 0.276, cy: 0.350, w: 0.458, h: 0.210",
+      "hierarchy_and_relation": "Primary typographic focal point on the left, positioned below the brand and above the supporting copy; it remains unobstructed against the pale background.",
+      "color_specs": [
+        "#352117",
+        "#F7F1E8"
+      ]
+    },
+    {
+      "description": "Two-line supporting statement in a smaller medium-weight sans-serif: first line 'Handpicked roasts delivered weekly.', second line 'Start your week with better coffee.'. The two rendered-text occurrences are stacked as consecutive left-aligned rows beneath the headline, with the second row directly below the first. The text is left aligned, muted warm brown, and evenly spaced.",
+      "coordinates": "cx: 0.224, cy: 0.505, w: 0.348, h: 0.068",
+      "hierarchy_and_relation": "Placed beneath the headline and above the call-to-action, with strong contrast against the off-white field.",
+      "color_specs": [
+        "#6F503D",
+        "#F7F1E8"
+      ]
+    },
+    {
+      "description": "Rounded rectangular call-to-action button with a solid burnt-orange fill. Centered inside is the white bold sans-serif label 'Start your subscription', followed by a thin white right-pointing arrow icon. The label and arrow share one horizontal row, with the arrow positioned to the label's right and both vertically centered within the button. The button has softly rounded corners and a subtle warm shadow.",
+      "coordinates": "cx: 0.202, cy: 0.612, w: 0.299, h: 0.064",
+      "hierarchy_and_relation": "Floats over the left background below the supporting statement; its saturated fill creates the strongest interactive accent on the page.",
+      "color_specs": [
+        "#C9682F",
+        "#FFFFFF",
+        "#A95225"
+      ]
+    },
+    {
+      "description": "Three compact benefit items arranged in one horizontal row beneath the button. Each item pairs a thin brown outline icon with two stacked text lines: a coffee-bean icon above 'Ethically' and 'sourced'; a leaf icon above 'Small-batch' and 'roasted'; and a calendar icon above 'Delivered' and 'every Monday'. Within each item, the two rendered-text occurrences form a left-aligned two-row column beside its icon; the three icon-and-text groups proceed left to right on a shared horizontal row. Headings are darker and slightly bolder than the smaller muted sublines.",
+      "coordinates": "cx: 0.237, cy: 0.720, w: 0.375, h: 0.060",
+      "hierarchy_and_relation": "Anchored along the lower-left area, aligned beneath the call-to-action and separated from the photograph by open negative space.",
+      "color_specs": [
+        "#5A3827",
+        "#806451",
+        "#F7F1E8"
+      ]
+    },
+    {
+      "description": "Large shallow ceramic latte cup viewed from a slightly elevated front angle, positioned on the right. The cup has a thick cream rim, a rounded body with warm beige speckling, and a broad curved handle extending to the right. It rests on a dark brown wooden saucer or tabletop with visible grain and soft highlights.",
+      "coordinates": "cx: 0.735, cy: 0.666, w: 0.491, h: 0.430",
+      "hierarchy_and_relation": "Foreground photographic subject overlapping the blurred café background; the cup occludes the wood surface and supports the latte-art surface above it.",
+      "color_specs": [
+        "#F3E7D4",
+        "#D8B992",
+        "#B98255",
+        "#5A321F",
+        "#2B1A12"
+      ]
+    },
+    {
+      "description": "Creamy latte surface inside the cup, filling the upper bowl with a pale beige foam. Symmetrical rosette latte art is formed by crisp white milk foam and darker caramel-brown coffee channels, with a pointed leaf-like center and soft blurred foam edges.",
+      "coordinates": "cx: 0.724, cy: 0.535, w: 0.315, h: 0.225",
+      "hierarchy_and_relation": "Contained within and partially occluded by the cup rim; the bright foam contrasts against the darker coffee and draws attention toward the upper-right.",
+      "color_specs": [
+        "#F6E8D2",
+        "#D9B27D",
+        "#A9683D",
+        "#7B4328"
+      ]
+    },
+    {
+      "description": "Foreground coffee beans scattered across the lower-right wooden surface, including several sharply focused oval beans near the bottom edge and softer beans behind them. The beans are roasted chestnut brown with glossy highlights and natural central creases.",
+      "coordinates": "cx: 0.858, cy: 0.842, w: 0.282, h: 0.210",
+      "hierarchy_and_relation": "Topmost photographic foreground detail, overlapping the wooden surface and partially approaching the cup without covering its main body.",
+      "color_specs": [
+        "#4A2818",
+        "#6E3D23",
+        "#8B502D",
+        "#C18A5B"
+      ]
+    }
+  ]
+}

samples/info_water.json

@@ -0,0 +1,141 @@
+{
+  "canvas_settings": {
+    "aspect_ratio": "1:1",
+    "ambient_lighting": "Bright, even white studio illumination with soft pastel gradients and no cast shadows.",
+    "image_style": "Clean flat-vector educational infographic with rounded geometry, thin white outlines, and a centered four-stage cycle."
+  },
+  "layers": [
+    {
+      "description": "Solid white square background with generous margins around a centered water-cycle infographic.",
+      "coordinates": "cx: 0.500, cy: 0.500, w: 1.000, h: 1.000",
+      "hierarchy_and_relation": "Base layer beneath every illustration, arrow, panel, and text element.",
+      "color_specs": [
+        "#FFFFFF"
+      ]
+    },
+    {
+      "description": "Large circular ocean-evaporation scene: a blue sun with short radial yellow rays rises over layered turquoise and deep-blue water. White-capped waves fill the lower half; a pale blue dashed vertical path rises from the water toward the top of the circle.",
+      "coordinates": "cx: 0.247, cy: 0.258, w: 0.333, h: 0.333",
+      "hierarchy_and_relation": "Occupies the upper-left cycle position; its top edge is partially occluded by the adjacent upward arrow, and it sits behind the stage label below.",
+      "color_specs": [
+        "#29A9E2",
+        "#7ED6F4",
+        "#087FC1",
+        "#075E91",
+        "#FFFFFF",
+        "#FFD547"
+      ]
+    },
+    {
+      "description": "Large circular cloud-formation scene: a pale blue sky contains several white and light-gray puffy clouds, with a darker gray rain cloud near the lower center. A pale blue dashed vertical path descends from the upper cloud area.",
+      "coordinates": "cx: 0.750, cy: 0.258, w: 0.333, h: 0.333",
+      "hierarchy_and_relation": "Occupies the upper-right cycle position; its left edge is partially occluded by the incoming arrow from the evaporation stage and its lower edge is partially occluded by the downward arrow.",
+      "color_specs": [
+        "#B9E5F5",
+        "#FFFFFF",
+        "#D9D9D9",
+        "#858585",
+        "#29A9E2"
+      ]
+    },
+    {
+      "description": "Large circular rainfall scene: dark blue-gray storm clouds and falling blue raindrops fill the upper portion above a vivid turquoise-blue water surface with small white wave marks.",
+      "coordinates": "cx: 0.750, cy: 0.742, w: 0.333, h: 0.333",
+      "hierarchy_and_relation": "Occupies the lower-right cycle position; receives the downward arrow from the cloud stage and sends a curved arrow toward the collection stage.",
+      "color_specs": [
+        "#555F68",
+        "#747D85",
+        "#29A9E2",
+        "#087FC1",
+        "#FFFFFF"
+      ]
+    },
+    {
+      "description": "Large circular land-collection scene: green hills and grassy banks surround a blue lake or river. Dark evergreen trees stand on the left and right banks, pale gray mountains recede in the distance, and a small white cloud floats at upper left.",
+      "coordinates": "cx: 0.247, cy: 0.742, w: 0.333, h: 0.333",
+      "hierarchy_and_relation": "Occupies the lower-left cycle position; receives the curved arrow from rainfall and sends a long curved arrow upward toward evaporation.",
+      "color_specs": [
+        "#238F38",
+        "#56A94B",
+        "#087FC1",
+        "#29A9E2",
+        "#A9C7D8",
+        "#FFFFFF",
+        "#4F7E35"
+      ]
+    },
+    {
+      "description": "Thick curved directional arrows with rounded shafts and triangular heads: one rises from the upper-left circle toward the upper-right circle, one descends from the upper-right circle toward the lower-right circle, one curves from the lower-right circle toward the lower-left circle, and one curves from the lower-left circle toward the upper-left circle. Each arrow is medium blue with a white outline.",
+      "coordinates": "cx: 0.500, cy: 0.500, w: 0.744, h: 0.744",
+      "hierarchy_and_relation": "Overlays the gaps between the four circular scenes, visually linking them into a clockwise cycle while remaining behind the stage labels.",
+      "color_specs": [
+        "#29A9E2",
+        "#FFFFFF"
+      ]
+    },
+    {
+      "description": "Bold dark-blue uppercase stage heading 'EVAPORATION', centered above the upper-left circular scene. The heading forms the first row of a centered two-row label group, stacked directly above its explanatory copy.",
+      "coordinates": "cx: 0.247, cy: 0.070, w: 0.245, h: 0.031",
+      "hierarchy_and_relation": "Topmost label for the evaporation stage, positioned above and not touching the circular water illustration.",
+      "color_specs": [
+        "#075E91"
+      ]
+    },
+    {
+      "description": "Two centered lines of small dark-gray explanatory copy beneath the upper-left heading: first line 'The sun heats up water', second line 'and it turns into vapor.'. The two occurrences form a compact stacked pair on separate rows, share a common center axis, and sit directly beneath the heading.",
+      "coordinates": "cx: 0.247, cy: 0.111, w: 0.190, h: 0.043",
+      "hierarchy_and_relation": "Secondary text block directly below the evaporation heading and above the ocean scene.",
+      "color_specs": [
+        "#333333"
+      ]
+    },
+    {
+      "description": "Bold dark-blue uppercase stage heading 'CONDENSATION', centered above the upper-right circular scene. The heading forms the first row of a centered two-row label group, stacked directly above its explanatory copy.",
+      "coordinates": "cx: 0.750, cy: 0.070, w: 0.245, h: 0.031",
+      "hierarchy_and_relation": "Topmost label for the condensation stage, positioned above the cloud illustration.",
+      "color_specs": [
+        "#075E91"
+      ]
+    },
+    {
+      "description": "Two centered lines of small dark-gray explanatory copy beneath the upper-right heading: first line 'Water vapor cools', second line 'and forms clouds.'. The two occurrences form a compact stacked pair on separate rows, share a common center axis, and sit directly beneath the heading.",
+      "coordinates": "cx: 0.750, cy: 0.111, w: 0.170, h: 0.043",
+      "hierarchy_and_relation": "Secondary text block directly below the condensation heading and above the cloud scene.",
+      "color_specs": [
+        "#333333"
+      ]
+    },
+    {
+      "description": "Bold dark-blue uppercase stage heading 'PRECIPITATION', centered beside the lower-right circular scene. The heading forms the first row of a centered two-row label group, stacked directly above its explanatory copy.",
+      "coordinates": "cx: 0.750, cy: 0.918, w: 0.245, h: 0.031",
+      "hierarchy_and_relation": "Topmost label below the rainfall illustration, aligned to its horizontal center.",
+      "color_specs": [
+        "#075E91"
+      ]
+    },
+    {
+      "description": "Two centered lines of small dark-gray explanatory copy beneath the lower-right heading: first line 'Water falls as rain,', second line 'snow, or hail.'. The two occurrences form a compact stacked pair on separate rows, share a common center axis, and sit directly beneath the heading.",
+      "coordinates": "cx: 0.750, cy: 0.959, w: 0.170, h: 0.043",
+      "hierarchy_and_relation": "Secondary text block directly below the precipitation heading and above the bottom canvas margin.",
+      "color_specs": [
+        "#333333"
+      ]
+    },
+    {
+      "description": "Bold dark-blue uppercase stage heading 'COLLECTION', centered above the lower-left circular scene. The heading forms the first row of a centered two-row label group, stacked directly above its explanatory copy.",
+      "coordinates": "cx: 0.247, cy: 0.918, w: 0.200, h: 0.031",
+      "hierarchy_and_relation": "Topmost label below the land-and-water illustration, aligned to its horizontal center.",
+      "color_specs": [
+        "#075E91"
+      ]
+    },
+    {
+      "description": "Two centered lines of small dark-gray explanatory copy beneath the lower-left heading: first line 'Water collects in rivers, lakes,', second line 'and oceans.'. The two occurrences form a compact stacked pair on separate rows, share a common center axis, and sit directly beneath the heading.",
+      "coordinates": "cx: 0.247, cy: 0.959, w: 0.220, h: 0.043",
+      "hierarchy_and_relation": "Secondary text block directly below the collection heading and above the bottom canvas margin.",
+      "color_specs": [
+        "#333333"
+      ]
+    }
+  ]
+}

samples/poster_jazz.json

@@ -0,0 +1,86 @@
+{
+  "canvas_settings": {
+    "aspect_ratio": "1:1",
+    "ambient_lighting": "Low-key nocturnal lighting with a soft cyan-blue glow rising from the lower center, fading into near-black navy at the top; the saxophone silhouette is rim-lit by this cool backlight.",
+    "image_style": "Minimalist square concert-poster design with smooth digital gradients, crisp centered typography, and a single high-contrast photographic-style instrument silhouette."
+  },
+  "layers": [
+    {
+      "description": "Full-square background formed by a smooth deep navy-to-black vertical gradient, subtly brighter blue around the lower center and nearly black along the upper edge; no visible texture beyond the soft tonal falloff.",
+      "coordinates": "cx: 0.500, cy: 0.500, w: 1.000, h: 1.000",
+      "hierarchy_and_relation": "Base layer beneath all typography, ornament, and instrument imagery; its lower-center glow silhouettes the saxophone.",
+      "color_specs": [
+        "#020A18",
+        "#06152A",
+        "#0A2848",
+        "#123E63"
+      ]
+    },
+    {
+      "description": "Large centered event title 'Blue Hour' in an elegant high-contrast serif typeface, pale blue-white, with very large letterforms and generous spacing; the title occupies the upper-middle area on a single centered baseline and is horizontally centered over the saxophone and lower information block.",
+      "coordinates": "cx: 0.500, cy: 0.250, w: 0.760, h: 0.145",
+      "hierarchy_and_relation": "Topmost typography over the dark gradient, centered above the saxophone and visually dominant over all other text.",
+      "color_specs": [
+        "#D8E8F7",
+        "#BFD8EF"
+      ]
+    },
+    {
+      "description": "Small uppercase subtitle 'A JAZZ NIGHT' centered below the main title, set in widely tracked sans-serif letters; a thin horizontal rule appears on each side, aligned to the subtitle baseline and extending outward symmetrically toward the left and right margins.",
+      "coordinates": "cx: 0.500, cy: 0.345, w: 0.520, h: 0.030",
+      "hierarchy_and_relation": "Placed directly beneath the title and above the saxophone, with the flanking rules framing the subtitle without touching the letters.",
+      "color_specs": [
+        "#8FB4D2",
+        "#5F88AA"
+      ]
+    },
+    {
+      "description": "Vertical saxophone silhouette in dark navy-black, shown in side profile facing left with the bell flaring at the lower left; curved neck, mouthpiece, body keys, and bell are visible, with a narrow cool-blue highlight tracing the right edge and lower rim.",
+      "coordinates": "cx: 0.493, cy: 0.574, w: 0.285, h: 0.445",
+      "hierarchy_and_relation": "Central photographic-style object over the gradient; it sits below the title block and behind the lower event-information typography, with its lower portion partially occluded by the date row.",
+      "color_specs": [
+        "#020B18",
+        "#06152A",
+        "#0B2742",
+        "#2E6F9D"
+      ]
+    },
+    {
+      "description": "Thin circular blue outline behind the saxophone, centered around the instrument body; the ring is delicate, mostly visible along the left and lower arc, and partially hidden by the saxophone silhouette.",
+      "coordinates": "cx: 0.499, cy: 0.574, w: 0.350, h: 0.350",
+      "hierarchy_and_relation": "Decorative midground element between the gradient background and saxophone; the saxophone occludes portions of the ring while the ring remains visible around the left and bottom edges.",
+      "color_specs": [
+        "#174A73",
+        "#0D3153"
+      ]
+    },
+    {
+      "description": "Centered date line 'Friday, October 3' in a clean light sans-serif typeface, pale blue-gray, positioned below the saxophone bell and body.",
+      "coordinates": "cx: 0.500, cy: 0.786, w: 0.430, h: 0.040",
+      "hierarchy_and_relation": "Foreground text below the instrument, centered over the lower gradient and above the time-and-venue row.",
+      "color_specs": [
+        "#C7D9EA",
+        "#9FB9D0"
+      ]
+    },
+    {
+      "description": "Centered lower information row with '8 PM' at left and 'The Lantern Room' at right, both in small widely tracked uppercase sans-serif lettering; a thin horizontal rule separates this row from the date above and aligns with the left and right text endpoints.",
+      "coordinates": "cx: 0.500, cy: 0.835, w: 0.430, h: 0.055",
+      "hierarchy_and_relation": "Foreground information layer beneath the date line; the two text occurrences share a baseline and are separated by open space, while the rule spans behind them horizontally.",
+      "color_specs": [
+        "#A8C2D9",
+        "#6F93B1",
+        "#315776"
+      ]
+    },
+    {
+      "description": "Small bottom decorative mark centered near the lower edge: a thin vertical stem with a tiny blue dot above it and a short horizontal base line.",
+      "coordinates": "cx: 0.500, cy: 0.923, w: 0.060, h: 0.055",
+      "hierarchy_and_relation": "Topmost minimal ornament below the event information, isolated against the dark background and aligned to the poster center axis.",
+      "color_specs": [
+        "#1E5A86",
+        "#0B2D4A"
+      ]
+    }
+  ]
+}

samples/ui_banking.json

@@ -0,0 +1,94 @@
+{
+  "canvas_settings": {
+    "aspect_ratio": "1:1",
+    "ambient_lighting": "soft cool daylight UI lighting with pale blue-white background glow and subtle card shadows",
+    "image_style": "clean modern fintech mobile app UI, rounded white cards, soft drop shadows, sans-serif typography, flat vector icons"
+  },
+  "layers": [
+    {
+      "description": "Full-screen mobile app background with a very pale blue-to-white gradient, slightly brighter at the top and lower edges, smooth glassy fintech aesthetic with no visible texture beyond soft gradient shading.",
+      "coordinates": "cx: 0.500, cy: 0.500, w: 1.000, h: 1.000",
+      "hierarchy_and_relation": "Z-index 0 background layer; all cards, text, icons, and navigation elements sit above it with soft shadows separating them from the pale surface.",
+      "color_specs": [
+        "#F3F8FF",
+        "#EAF3FF",
+        "#FFFFFF"
+      ]
+    },
+    {
+      "description": "Top status bar with dark time text '9:41' at the upper left, and compact black cellular signal bars, Wi-Fi arcs, and a rounded battery outline with dark fill at the upper right. The time and the right-side icon cluster sit on the same horizontal status-bar row, aligned near the top edge with wide empty space between them.",
+      "coordinates": "cx: 0.500, cy: 0.030, w: 0.850, h: 0.035",
+      "hierarchy_and_relation": "Sits above the background and above the app header, visually aligned to the top safe area without overlapping the balance card.",
+      "color_specs": [
+        "#111827",
+        "#000000",
+        "#FFFFFF",
+        "#D1D5DB"
+      ]
+    },
+    {
+      "description": "Header area with bold dark greeting text 'Good morning, Alex' near the upper left, smaller muted subtitle 'Here's your financial overview' directly underneath, and a circular white notification button on the upper right containing a thin outlined bell icon with a small blue unread dot. The greeting and subtitle form a left-aligned two-line stack, while the notification button is horizontally aligned with the greeting block and separated to the far right.",
+      "coordinates": "cx: 0.500, cy: 0.095, w: 0.850, h: 0.075",
+      "hierarchy_and_relation": "Placed below the status bar and above the balance card; the notification button floats on the same horizontal band and casts a faint shadow over the background.",
+      "color_specs": [
+        "#111827",
+        "#6B7280",
+        "#FFFFFF",
+        "#2F80ED",
+        "#E5E7EB"
+      ]
+    },
+    {
+      "description": "Large rounded balance card with deep blue gradient fill, subtle curved wave decoration in the lower right, and a white outlined eye icon button near the top right. Left side contains small label 'Total Balance', oversized white balance amount '$8,450.75', and smaller white currency label 'USD'. Right lower area contains a white rounded square icon tile with a blue wallet/card symbol. The label, amount, and currency label are stacked in a left-aligned column inside the card, while the eye button aligns near the card's top-right corner and the wallet icon tile anchors the lower-right area.",
+      "coordinates": "cx: 0.500, cy: 0.250, w: 0.850, h: 0.225",
+      "hierarchy_and_relation": "Main hero card sits above the background and below the header; its shadow and saturated blue fill make it the dominant visual block, with the wallet icon tile layered on top of the card gradient.",
+      "color_specs": [
+        "#0F62FE",
+        "#1B4FD8",
+        "#FFFFFF",
+        "#EAF2FF",
+        "#1E5BDB"
+      ]
+    },
+    {
+      "description": "White rounded recent transactions panel with title 'Recent Transactions' at the top left and blue link 'View all' at the top right. First transaction row shows a green circular food icon with a white shopping bag outline, merchant name 'Whole Foods Market', date line 'May 20, 2024', right-aligned negative amount '-$85.42', and category 'Groceries'. Second row shows a blue circular ride icon with a white car outline, merchant name 'Uber Ride', date line 'May 20, 2024', right-aligned negative amount '-$18.75', and category 'Transport'. Third row shows a purple circular salary icon with a white wallet outline, merchant name 'Salary Deposit', date line 'May 19, 2024', right-aligned positive amount '+$2,500.00', and category 'Income'. Thin light gray dividers separate rows. The panel title and link share the top header row, and each transaction row uses a consistent left-to-right layout of icon, two-line merchant/date text block, and right-aligned amount/category stack.",
+      "coordinates": "cx: 0.500, cy: 0.520, w: 0.850, h: 0.285",
+      "hierarchy_and_relation": "Panel sits below the balance card with a soft shadow; each icon, text block, amount, and category is contained within the white card, with dividers layered behind the row content.",
+      "color_specs": [
+        "#FFFFFF",
+        "#111827",
+        "#6B7280",
+        "#2F80ED",
+        "#22C55E",
+        "#8B5CF6",
+        "#E5E7EB"
+      ]
+    },
+    {
+      "description": "White rounded quick-actions panel with heading 'Quick Actions' at the top left. Three evenly spaced action columns appear below: left column has a blue circular icon with a white paper-plane/send symbol and label 'Send' with sublabel 'Transfer money'; middle column has a green circular icon with a white payment-card symbol and label 'Pay' with sublabel 'Pay bills & more'; right column has a purple circular icon with a white plus symbol and label 'Top Up' with sublabel 'Add funds'. The heading sits above a single row of three equal action columns, each column stacking its circular icon above the bold label and smaller sublabel, all centered within its column.",
+      "coordinates": "cx: 0.500, cy: 0.765, w: 0.850, h: 0.170",
+      "hierarchy_and_relation": "Positioned below the transactions panel and above the bottom navigation; the three action groups are children of the white rounded card and are separated by spacing rather than visible dividers.",
+      "color_specs": [
+        "#FFFFFF",
+        "#111827",
+        "#6B7280",
+        "#2F80ED",
+        "#22C55E",
+        "#8B5CF6",
+        "#F3F4F6"
+      ]
+    },
+    {
+      "description": "Bottom navigation bar as a wide white rounded rectangle with four evenly spaced tabs. Left tab is active with a blue home icon and blue label 'Home'. Second tab has a gray vertical card/statement icon and label 'Cards'. Third tab has a gray pie-chart icon and label 'Insights'. Fourth tab has a gray person outline icon and label 'Profile'. The four tabs form a single evenly spaced horizontal row, with each icon centered above its label and all labels aligned along the same baseline.",
+      "coordinates": "cx: 0.500, cy: 0.915, w: 0.850, h: 0.095",
+      "hierarchy_and_relation": "Topmost persistent navigation layer sits above the background and below no other content; it overlaps the lower margin with a soft shadow and keeps all tab labels aligned along the same baseline.",
+      "color_specs": [
+        "#FFFFFF",
+        "#2F80ED",
+        "#6B7280",
+        "#E5E7EB",
+        "#111827"
+      ]
+    }
+  ]
+}

scheduler/scheduler_config.json

@@ -0,0 +1,7 @@
+{
+  "_class_name": "FlowMatchEulerDiscreteScheduler",
+  "_diffusers_version": "0.37.0.dev0",
+  "num_train_timesteps": 1000,
+  "use_dynamic_shifting": false,
+  "shift": 6.0
+}

transformer/config.json

@@ -0,0 +1,35 @@
+{
+  "_class_name": "DiffusionTransformer",
+  "_diffusers_version": "0.36.0",
+  "_name_or_path": "",
+  "all_f_patch_size": [
+    1
+  ],
+  "all_patch_size": [
+    2
+  ],
+  "axes_dims": [
+    32,
+    48,
+    48
+  ],
+  "axes_lens": [
+    20480,
+    512,
+    512
+  ],
+  "cap_feat_dim": 2560,
+  "dim": 3840,
+  "in_channels": 16,
+  "n_heads": 30,
+  "n_kv_heads": 30,
+  "n_layers": 30,
+  "n_refiner_layers": 2,
+  "norm_eps": 1e-05,
+  "qk_norm": true,
+  "rope_theta": 256.0,
+  "siglip_feat_dim": null,
+  "t_scale": 1000.0,
+  "alignment_padding_mode": "zero_masked",
+  "multi_frame_output": false
+}

transformer/diffusion_pytorch_model.safetensors.index.json

@@ -0,0 +1,526 @@
+{
+  "metadata": {
+    "total_size": 12309802112
+  },
+  "weight_map": {
+    "all_final_layer.2-1.adaLN_modulation.1.bias": "diffusion_pytorch_model-00001-of-00005.safetensors",
+    "all_final_layer.2-1.adaLN_modulation.1.weight": "diffusion_pytorch_model-00001-of-00005.safetensors",
+    "all_final_layer.2-1.linear.bias": "diffusion_pytorch_model-00001-of-00005.safetensors",
+    "all_final_layer.2-1.linear.weight": "diffusion_pytorch_model-00001-of-00005.safetensors",
+    "all_x_embedder.2-1.bias": "diffusion_pytorch_model-00001-of-00005.safetensors",
+    "all_x_embedder.2-1.weight": "diffusion_pytorch_model-00001-of-00005.safetensors",
+    "cap_embedder.0.weight": "diffusion_pytorch_model-00001-of-00005.safetensors",
+    "cap_embedder.1.bias": "diffusion_pytorch_model-00001-of-00005.safetensors",
+    "cap_embedder.1.weight": "diffusion_pytorch_model-00001-of-00005.safetensors",
+    "context_refiner.0.attention.norm_k.weight": "diffusion_pytorch_model-00001-of-00005.safetensors",
+    "context_refiner.0.attention.norm_q.weight": "diffusion_pytorch_model-00001-of-00005.safetensors",
+    "context_refiner.0.attention.to_k.weight": "diffusion_pytorch_model-00001-of-00005.safetensors",
+    "context_refiner.0.attention.to_out.0.weight": "diffusion_pytorch_model-00001-of-00005.safetensors",
+    "context_refiner.0.attention.to_q.weight": "diffusion_pytorch_model-00001-of-00005.safetensors",
+    "context_refiner.0.attention.to_v.weight": "diffusion_pytorch_model-00001-of-00005.safetensors",
+    "context_refiner.0.attention_norm1.weight": "diffusion_pytorch_model-00001-of-00005.safetensors",
+    "context_refiner.0.attention_norm2.weight": "diffusion_pytorch_model-00001-of-00005.safetensors",
+    "context_refiner.0.feed_forward.w1.weight": "diffusion_pytorch_model-00001-of-00005.safetensors",
+    "context_refiner.0.feed_forward.w2.weight": "diffusion_pytorch_model-00001-of-00005.safetensors",
+    "context_refiner.0.feed_forward.w3.weight": "diffusion_pytorch_model-00001-of-00005.safetensors",
+    "context_refiner.0.ffn_norm1.weight": "diffusion_pytorch_model-00001-of-00005.safetensors",
+    "context_refiner.0.ffn_norm2.weight": "diffusion_pytorch_model-00001-of-00005.safetensors",
+    "context_refiner.1.attention.norm_k.weight": "diffusion_pytorch_model-00001-of-00005.safetensors",
+    "context_refiner.1.attention.norm_q.weight": "diffusion_pytorch_model-00001-of-00005.safetensors",
+    "context_refiner.1.attention.to_k.weight": "diffusion_pytorch_model-00001-of-00005.safetensors",
+    "context_refiner.1.attention.to_out.0.weight": "diffusion_pytorch_model-00001-of-00005.safetensors",
+    "context_refiner.1.attention.to_q.weight": "diffusion_pytorch_model-00001-of-00005.safetensors",
+    "context_refiner.1.attention.to_v.weight": "diffusion_pytorch_model-00001-of-00005.safetensors",
+    "context_refiner.1.attention_norm1.weight": "diffusion_pytorch_model-00001-of-00005.safetensors",
+    "context_refiner.1.attention_norm2.weight": "diffusion_pytorch_model-00001-of-00005.safetensors",
+    "context_refiner.1.feed_forward.w1.weight": "diffusion_pytorch_model-00001-of-00005.safetensors",
+    "context_refiner.1.feed_forward.w2.weight": "diffusion_pytorch_model-00001-of-00005.safetensors",
+    "context_refiner.1.feed_forward.w3.weight": "diffusion_pytorch_model-00001-of-00005.safetensors",
+    "context_refiner.1.ffn_norm1.weight": "diffusion_pytorch_model-00001-of-00005.safetensors",
+    "context_refiner.1.ffn_norm2.weight": "diffusion_pytorch_model-00001-of-00005.safetensors",
+    "layers.0.adaLN_modulation.0.bias": "diffusion_pytorch_model-00001-of-00005.safetensors",
+    "layers.0.adaLN_modulation.0.weight": "diffusion_pytorch_model-00001-of-00005.safetensors",
+    "layers.0.attention.norm_k.weight": "diffusion_pytorch_model-00001-of-00005.safetensors",
+    "layers.0.attention.norm_q.weight": "diffusion_pytorch_model-00001-of-00005.safetensors",
+    "layers.0.attention.to_k.weight": "diffusion_pytorch_model-00001-of-00005.safetensors",
+    "layers.0.attention.to_out.0.weight": "diffusion_pytorch_model-00001-of-00005.safetensors",
+    "layers.0.attention.to_q.weight": "diffusion_pytorch_model-00001-of-00005.safetensors",
+    "layers.0.attention.to_v.weight": "diffusion_pytorch_model-00001-of-00005.safetensors",
+    "layers.0.attention_norm1.weight": "diffusion_pytorch_model-00001-of-00005.safetensors",
+    "layers.0.attention_norm2.weight": "diffusion_pytorch_model-00001-of-00005.safetensors",
+    "layers.0.feed_forward.w1.weight": "diffusion_pytorch_model-00001-of-00005.safetensors",
+    "layers.0.feed_forward.w2.weight": "diffusion_pytorch_model-00001-of-00005.safetensors",
+    "layers.0.feed_forward.w3.weight": "diffusion_pytorch_model-00001-of-00005.safetensors",
+    "layers.0.ffn_norm1.weight": "diffusion_pytorch_model-00001-of-00005.safetensors",
+    "layers.0.ffn_norm2.weight": "diffusion_pytorch_model-00001-of-00005.safetensors",
+    "layers.1.adaLN_modulation.0.bias": "diffusion_pytorch_model-00001-of-00005.safetensors",
+    "layers.1.adaLN_modulation.0.weight": "diffusion_pytorch_model-00001-of-00005.safetensors",
+    "layers.1.attention.norm_k.weight": "diffusion_pytorch_model-00001-of-00005.safetensors",
+    "layers.1.attention.norm_q.weight": "diffusion_pytorch_model-00001-of-00005.safetensors",
+    "layers.1.attention.to_k.weight": "diffusion_pytorch_model-00001-of-00005.safetensors",
+    "layers.1.attention.to_out.0.weight": "diffusion_pytorch_model-00001-of-00005.safetensors",
+    "layers.1.attention.to_q.weight": "diffusion_pytorch_model-00001-of-00005.safetensors",
+    "layers.1.attention.to_v.weight": "diffusion_pytorch_model-00001-of-00005.safetensors",
+    "layers.1.attention_norm1.weight": "diffusion_pytorch_model-00001-of-00005.safetensors",
+    "layers.1.attention_norm2.weight": "diffusion_pytorch_model-00001-of-00005.safetensors",
+    "layers.1.feed_forward.w1.weight": "diffusion_pytorch_model-00001-of-00005.safetensors",
+    "layers.1.feed_forward.w2.weight": "diffusion_pytorch_model-00001-of-00005.safetensors",
+    "layers.1.feed_forward.w3.weight": "diffusion_pytorch_model-00001-of-00005.safetensors",
+    "layers.1.ffn_norm1.weight": "diffusion_pytorch_model-00001-of-00005.safetensors",
+    "layers.1.ffn_norm2.weight": "diffusion_pytorch_model-00001-of-00005.safetensors",
+    "layers.10.adaLN_modulation.0.bias": "diffusion_pytorch_model-00002-of-00005.safetensors",
+    "layers.10.adaLN_modulation.0.weight": "diffusion_pytorch_model-00002-of-00005.safetensors",
+    "layers.10.attention.norm_k.weight": "diffusion_pytorch_model-00002-of-00005.safetensors",
+    "layers.10.attention.norm_q.weight": "diffusion_pytorch_model-00002-of-00005.safetensors",
+    "layers.10.attention.to_k.weight": "diffusion_pytorch_model-00002-of-00005.safetensors",
+    "layers.10.attention.to_out.0.weight": "diffusion_pytorch_model-00002-of-00005.safetensors",
+    "layers.10.attention.to_q.weight": "diffusion_pytorch_model-00002-of-00005.safetensors",
+    "layers.10.attention.to_v.weight": "diffusion_pytorch_model-00002-of-00005.safetensors",
+    "layers.10.attention_norm1.weight": "diffusion_pytorch_model-00002-of-00005.safetensors",
+    "layers.10.attention_norm2.weight": "diffusion_pytorch_model-00002-of-00005.safetensors",
+    "layers.10.feed_forward.w1.weight": "diffusion_pytorch_model-00002-of-00005.safetensors",
+    "layers.10.feed_forward.w2.weight": "diffusion_pytorch_model-00002-of-00005.safetensors",
+    "layers.10.feed_forward.w3.weight": "diffusion_pytorch_model-00002-of-00005.safetensors",
+    "layers.10.ffn_norm1.weight": "diffusion_pytorch_model-00002-of-00005.safetensors",
+    "layers.10.ffn_norm2.weight": "diffusion_pytorch_model-00002-of-00005.safetensors",
+    "layers.11.adaLN_modulation.0.bias": "diffusion_pytorch_model-00002-of-00005.safetensors",
+    "layers.11.adaLN_modulation.0.weight": "diffusion_pytorch_model-00002-of-00005.safetensors",
+    "layers.11.attention.norm_k.weight": "diffusion_pytorch_model-00002-of-00005.safetensors",
+    "layers.11.attention.norm_q.weight": "diffusion_pytorch_model-00002-of-00005.safetensors",
+    "layers.11.attention.to_k.weight": "diffusion_pytorch_model-00002-of-00005.safetensors",
+    "layers.11.attention.to_out.0.weight": "diffusion_pytorch_model-00002-of-00005.safetensors",
+    "layers.11.attention.to_q.weight": "diffusion_pytorch_model-00002-of-00005.safetensors",
+    "layers.11.attention.to_v.weight": "diffusion_pytorch_model-00002-of-00005.safetensors",
+    "layers.11.attention_norm1.weight": "diffusion_pytorch_model-00002-of-00005.safetensors",
+    "layers.11.attention_norm2.weight": "diffusion_pytorch_model-00002-of-00005.safetensors",
+    "layers.11.feed_forward.w1.weight": "diffusion_pytorch_model-00002-of-00005.safetensors",
+    "layers.11.feed_forward.w2.weight": "diffusion_pytorch_model-00002-of-00005.safetensors",
+    "layers.11.feed_forward.w3.weight": "diffusion_pytorch_model-00002-of-00005.safetensors",
+    "layers.11.ffn_norm1.weight": "diffusion_pytorch_model-00002-of-00005.safetensors",
+    "layers.11.ffn_norm2.weight": "diffusion_pytorch_model-00002-of-00005.safetensors",
+    "layers.12.adaLN_modulation.0.bias": "diffusion_pytorch_model-00003-of-00005.safetensors",
+    "layers.12.adaLN_modulation.0.weight": "diffusion_pytorch_model-00003-of-00005.safetensors",
+    "layers.12.attention.norm_k.weight": "diffusion_pytorch_model-00002-of-00005.safetensors",
+    "layers.12.attention.norm_q.weight": "diffusion_pytorch_model-00002-of-00005.safetensors",
+    "layers.12.attention.to_k.weight": "diffusion_pytorch_model-00002-of-00005.safetensors",
+    "layers.12.attention.to_out.0.weight": "diffusion_pytorch_model-00002-of-00005.safetensors",
+    "layers.12.attention.to_q.weight": "diffusion_pytorch_model-00002-of-00005.safetensors",
+    "layers.12.attention.to_v.weight": "diffusion_pytorch_model-00002-of-00005.safetensors",
+    "layers.12.attention_norm1.weight": "diffusion_pytorch_model-00003-of-00005.safetensors",
+    "layers.12.attention_norm2.weight": "diffusion_pytorch_model-00003-of-00005.safetensors",
+    "layers.12.feed_forward.w1.weight": "diffusion_pytorch_model-00003-of-00005.safetensors",
+    "layers.12.feed_forward.w2.weight": "diffusion_pytorch_model-00003-of-00005.safetensors",
+    "layers.12.feed_forward.w3.weight": "diffusion_pytorch_model-00003-of-00005.safetensors",
+    "layers.12.ffn_norm1.weight": "diffusion_pytorch_model-00003-of-00005.safetensors",
+    "layers.12.ffn_norm2.weight": "diffusion_pytorch_model-00003-of-00005.safetensors",
+    "layers.13.adaLN_modulation.0.bias": "diffusion_pytorch_model-00003-of-00005.safetensors",
+    "layers.13.adaLN_modulation.0.weight": "diffusion_pytorch_model-00003-of-00005.safetensors",
+    "layers.13.attention.norm_k.weight": "diffusion_pytorch_model-00003-of-00005.safetensors",
+    "layers.13.attention.norm_q.weight": "diffusion_pytorch_model-00003-of-00005.safetensors",
+    "layers.13.attention.to_k.weight": "diffusion_pytorch_model-00003-of-00005.safetensors",
+    "layers.13.attention.to_out.0.weight": "diffusion_pytorch_model-00003-of-00005.safetensors",
+    "layers.13.attention.to_q.weight": "diffusion_pytorch_model-00003-of-00005.safetensors",
+    "layers.13.attention.to_v.weight": "diffusion_pytorch_model-00003-of-00005.safetensors",
+    "layers.13.attention_norm1.weight": "diffusion_pytorch_model-00003-of-00005.safetensors",
+    "layers.13.attention_norm2.weight": "diffusion_pytorch_model-00003-of-00005.safetensors",
+    "layers.13.feed_forward.w1.weight": "diffusion_pytorch_model-00003-of-00005.safetensors",
+    "layers.13.feed_forward.w2.weight": "diffusion_pytorch_model-00003-of-00005.safetensors",
+    "layers.13.feed_forward.w3.weight": "diffusion_pytorch_model-00003-of-00005.safetensors",
+    "layers.13.ffn_norm1.weight": "diffusion_pytorch_model-00003-of-00005.safetensors",
+    "layers.13.ffn_norm2.weight": "diffusion_pytorch_model-00003-of-00005.safetensors",
+    "layers.14.adaLN_modulation.0.bias": "diffusion_pytorch_model-00003-of-00005.safetensors",
+    "layers.14.adaLN_modulation.0.weight": "diffusion_pytorch_model-00003-of-00005.safetensors",
+    "layers.14.attention.norm_k.weight": "diffusion_pytorch_model-00003-of-00005.safetensors",
+    "layers.14.attention.norm_q.weight": "diffusion_pytorch_model-00003-of-00005.safetensors",
+    "layers.14.attention.to_k.weight": "diffusion_pytorch_model-00003-of-00005.safetensors",
+    "layers.14.attention.to_out.0.weight": "diffusion_pytorch_model-00003-of-00005.safetensors",
+    "layers.14.attention.to_q.weight": "diffusion_pytorch_model-00003-of-00005.safetensors",
+    "layers.14.attention.to_v.weight": "diffusion_pytorch_model-00003-of-00005.safetensors",
+    "layers.14.attention_norm1.weight": "diffusion_pytorch_model-00003-of-00005.safetensors",
+    "layers.14.attention_norm2.weight": "diffusion_pytorch_model-00003-of-00005.safetensors",
+    "layers.14.feed_forward.w1.weight": "diffusion_pytorch_model-00003-of-00005.safetensors",
+    "layers.14.feed_forward.w2.weight": "diffusion_pytorch_model-00003-of-00005.safetensors",
+    "layers.14.feed_forward.w3.weight": "diffusion_pytorch_model-00003-of-00005.safetensors",
+    "layers.14.ffn_norm1.weight": "diffusion_pytorch_model-00003-of-00005.safetensors",
+    "layers.14.ffn_norm2.weight": "diffusion_pytorch_model-00003-of-00005.safetensors",
+    "layers.15.adaLN_modulation.0.bias": "diffusion_pytorch_model-00003-of-00005.safetensors",
+    "layers.15.adaLN_modulation.0.weight": "diffusion_pytorch_model-00003-of-00005.safetensors",
+    "layers.15.attention.norm_k.weight": "diffusion_pytorch_model-00003-of-00005.safetensors",
+    "layers.15.attention.norm_q.weight": "diffusion_pytorch_model-00003-of-00005.safetensors",
+    "layers.15.attention.to_k.weight": "diffusion_pytorch_model-00003-of-00005.safetensors",
+    "layers.15.attention.to_out.0.weight": "diffusion_pytorch_model-00003-of-00005.safetensors",
+    "layers.15.attention.to_q.weight": "diffusion_pytorch_model-00003-of-00005.safetensors",
+    "layers.15.attention.to_v.weight": "diffusion_pytorch_model-00003-of-00005.safetensors",
+    "layers.15.attention_norm1.weight": "diffusion_pytorch_model-00003-of-00005.safetensors",
+    "layers.15.attention_norm2.weight": "diffusion_pytorch_model-00003-of-00005.safetensors",
+    "layers.15.feed_forward.w1.weight": "diffusion_pytorch_model-00003-of-00005.safetensors",
+    "layers.15.feed_forward.w2.weight": "diffusion_pytorch_model-00003-of-00005.safetensors",
+    "layers.15.feed_forward.w3.weight": "diffusion_pytorch_model-00003-of-00005.safetensors",
+    "layers.15.ffn_norm1.weight": "diffusion_pytorch_model-00003-of-00005.safetensors",
+    "layers.15.ffn_norm2.weight": "diffusion_pytorch_model-00003-of-00005.safetensors",
+    "layers.16.adaLN_modulation.0.bias": "diffusion_pytorch_model-00003-of-00005.safetensors",
+    "layers.16.adaLN_modulation.0.weight": "diffusion_pytorch_model-00003-of-00005.safetensors",
+    "layers.16.attention.norm_k.weight": "diffusion_pytorch_model-00003-of-00005.safetensors",
+    "layers.16.attention.norm_q.weight": "diffusion_pytorch_model-00003-of-00005.safetensors",
+    "layers.16.attention.to_k.weight": "diffusion_pytorch_model-00003-of-00005.safetensors",
+    "layers.16.attention.to_out.0.weight": "diffusion_pytorch_model-00003-of-00005.safetensors",
+    "layers.16.attention.to_q.weight": "diffusion_pytorch_model-00003-of-00005.safetensors",
+    "layers.16.attention.to_v.weight": "diffusion_pytorch_model-00003-of-00005.safetensors",
+    "layers.16.attention_norm1.weight": "diffusion_pytorch_model-00003-of-00005.safetensors",
+    "layers.16.attention_norm2.weight": "diffusion_pytorch_model-00003-of-00005.safetensors",
+    "layers.16.feed_forward.w1.weight": "diffusion_pytorch_model-00003-of-00005.safetensors",
+    "layers.16.feed_forward.w2.weight": "diffusion_pytorch_model-00003-of-00005.safetensors",
+    "layers.16.feed_forward.w3.weight": "diffusion_pytorch_model-00003-of-00005.safetensors",
+    "layers.16.ffn_norm1.weight": "diffusion_pytorch_model-00003-of-00005.safetensors",
+    "layers.16.ffn_norm2.weight": "diffusion_pytorch_model-00003-of-00005.safetensors",
+    "layers.17.adaLN_modulation.0.bias": "diffusion_pytorch_model-00003-of-00005.safetensors",
+    "layers.17.adaLN_modulation.0.weight": "diffusion_pytorch_model-00003-of-00005.safetensors",
+    "layers.17.attention.norm_k.weight": "diffusion_pytorch_model-00003-of-00005.safetensors",
+    "layers.17.attention.norm_q.weight": "diffusion_pytorch_model-00003-of-00005.safetensors",
+    "layers.17.attention.to_k.weight": "diffusion_pytorch_model-00003-of-00005.safetensors",
+    "layers.17.attention.to_out.0.weight": "diffusion_pytorch_model-00003-of-00005.safetensors",
+    "layers.17.attention.to_q.weight": "diffusion_pytorch_model-00003-of-00005.safetensors",
+    "layers.17.attention.to_v.weight": "diffusion_pytorch_model-00003-of-00005.safetensors",
+    "layers.17.attention_norm1.weight": "diffusion_pytorch_model-00003-of-00005.safetensors",
+    "layers.17.attention_norm2.weight": "diffusion_pytorch_model-00003-of-00005.safetensors",
+    "layers.17.feed_forward.w1.weight": "diffusion_pytorch_model-00003-of-00005.safetensors",
+    "layers.17.feed_forward.w2.weight": "diffusion_pytorch_model-00003-of-00005.safetensors",
+    "layers.17.feed_forward.w3.weight": "diffusion_pytorch_model-00003-of-00005.safetensors",
+    "layers.17.ffn_norm1.weight": "diffusion_pytorch_model-00003-of-00005.safetensors",
+    "layers.17.ffn_norm2.weight": "diffusion_pytorch_model-00003-of-00005.safetensors",
+    "layers.18.adaLN_modulation.0.bias": "diffusion_pytorch_model-00003-of-00005.safetensors",
+    "layers.18.adaLN_modulation.0.weight": "diffusion_pytorch_model-00003-of-00005.safetensors",
+    "layers.18.attention.norm_k.weight": "diffusion_pytorch_model-00003-of-00005.safetensors",
+    "layers.18.attention.norm_q.weight": "diffusion_pytorch_model-00003-of-00005.safetensors",
+    "layers.18.attention.to_k.weight": "diffusion_pytorch_model-00003-of-00005.safetensors",
+    "layers.18.attention.to_out.0.weight": "diffusion_pytorch_model-00003-of-00005.safetensors",
+    "layers.18.attention.to_q.weight": "diffusion_pytorch_model-00003-of-00005.safetensors",
+    "layers.18.attention.to_v.weight": "diffusion_pytorch_model-00003-of-00005.safetensors",
+    "layers.18.attention_norm1.weight": "diffusion_pytorch_model-00003-of-00005.safetensors",
+    "layers.18.attention_norm2.weight": "diffusion_pytorch_model-00003-of-00005.safetensors",
+    "layers.18.feed_forward.w1.weight": "diffusion_pytorch_model-00003-of-00005.safetensors",
+    "layers.18.feed_forward.w2.weight": "diffusion_pytorch_model-00003-of-00005.safetensors",
+    "layers.18.feed_forward.w3.weight": "diffusion_pytorch_model-00003-of-00005.safetensors",
+    "layers.18.ffn_norm1.weight": "diffusion_pytorch_model-00003-of-00005.safetensors",
+    "layers.18.ffn_norm2.weight": "diffusion_pytorch_model-00003-of-00005.safetensors",
+    "layers.19.adaLN_modulation.0.bias": "diffusion_pytorch_model-00003-of-00005.safetensors",
+    "layers.19.adaLN_modulation.0.weight": "diffusion_pytorch_model-00003-of-00005.safetensors",
+    "layers.19.attention.norm_k.weight": "diffusion_pytorch_model-00003-of-00005.safetensors",
+    "layers.19.attention.norm_q.weight": "diffusion_pytorch_model-00003-of-00005.safetensors",
+    "layers.19.attention.to_k.weight": "diffusion_pytorch_model-00003-of-00005.safetensors",
+    "layers.19.attention.to_out.0.weight": "diffusion_pytorch_model-00003-of-00005.safetensors",
+    "layers.19.attention.to_q.weight": "diffusion_pytorch_model-00003-of-00005.safetensors",
+    "layers.19.attention.to_v.weight": "diffusion_pytorch_model-00003-of-00005.safetensors",
+    "layers.19.attention_norm1.weight": "diffusion_pytorch_model-00003-of-00005.safetensors",
+    "layers.19.attention_norm2.weight": "diffusion_pytorch_model-00003-of-00005.safetensors",
+    "layers.19.feed_forward.w1.weight": "diffusion_pytorch_model-00003-of-00005.safetensors",
+    "layers.19.feed_forward.w2.weight": "diffusion_pytorch_model-00003-of-00005.safetensors",
+    "layers.19.feed_forward.w3.weight": "diffusion_pytorch_model-00003-of-00005.safetensors",
+    "layers.19.ffn_norm1.weight": "diffusion_pytorch_model-00003-of-00005.safetensors",
+    "layers.19.ffn_norm2.weight": "diffusion_pytorch_model-00003-of-00005.safetensors",
+    "layers.2.adaLN_modulation.0.bias": "diffusion_pytorch_model-00001-of-00005.safetensors",
+    "layers.2.adaLN_modulation.0.weight": "diffusion_pytorch_model-00001-of-00005.safetensors",
+    "layers.2.attention.norm_k.weight": "diffusion_pytorch_model-00001-of-00005.safetensors",
+    "layers.2.attention.norm_q.weight": "diffusion_pytorch_model-00001-of-00005.safetensors",
+    "layers.2.attention.to_k.weight": "diffusion_pytorch_model-00001-of-00005.safetensors",
+    "layers.2.attention.to_out.0.weight": "diffusion_pytorch_model-00001-of-00005.safetensors",
+    "layers.2.attention.to_q.weight": "diffusion_pytorch_model-00001-of-00005.safetensors",
+    "layers.2.attention.to_v.weight": "diffusion_pytorch_model-00001-of-00005.safetensors",
+    "layers.2.attention_norm1.weight": "diffusion_pytorch_model-00001-of-00005.safetensors",
+    "layers.2.attention_norm2.weight": "diffusion_pytorch_model-00001-of-00005.safetensors",
+    "layers.2.feed_forward.w1.weight": "diffusion_pytorch_model-00001-of-00005.safetensors",
+    "layers.2.feed_forward.w2.weight": "diffusion_pytorch_model-00001-of-00005.safetensors",
+    "layers.2.feed_forward.w3.weight": "diffusion_pytorch_model-00001-of-00005.safetensors",
+    "layers.2.ffn_norm1.weight": "diffusion_pytorch_model-00001-of-00005.safetensors",
+    "layers.2.ffn_norm2.weight": "diffusion_pytorch_model-00001-of-00005.safetensors",
+    "layers.20.adaLN_modulation.0.bias": "diffusion_pytorch_model-00004-of-00005.safetensors",
+    "layers.20.adaLN_modulation.0.weight": "diffusion_pytorch_model-00004-of-00005.safetensors",
+    "layers.20.attention.norm_k.weight": "diffusion_pytorch_model-00003-of-00005.safetensors",
+    "layers.20.attention.norm_q.weight": "diffusion_pytorch_model-00003-of-00005.safetensors",
+    "layers.20.attention.to_k.weight": "diffusion_pytorch_model-00003-of-00005.safetensors",
+    "layers.20.attention.to_out.0.weight": "diffusion_pytorch_model-00003-of-00005.safetensors",
+    "layers.20.attention.to_q.weight": "diffusion_pytorch_model-00003-of-00005.safetensors",
+    "layers.20.attention.to_v.weight": "diffusion_pytorch_model-00003-of-00005.safetensors",
+    "layers.20.attention_norm1.weight": "diffusion_pytorch_model-00004-of-00005.safetensors",
+    "layers.20.attention_norm2.weight": "diffusion_pytorch_model-00004-of-00005.safetensors",
+    "layers.20.feed_forward.w1.weight": "diffusion_pytorch_model-00003-of-00005.safetensors",
+    "layers.20.feed_forward.w2.weight": "diffusion_pytorch_model-00004-of-00005.safetensors",
+    "layers.20.feed_forward.w3.weight": "diffusion_pytorch_model-00004-of-00005.safetensors",
+    "layers.20.ffn_norm1.weight": "diffusion_pytorch_model-00004-of-00005.safetensors",
+    "layers.20.ffn_norm2.weight": "diffusion_pytorch_model-00004-of-00005.safetensors",
+    "layers.21.adaLN_modulation.0.bias": "diffusion_pytorch_model-00004-of-00005.safetensors",
+    "layers.21.adaLN_modulation.0.weight": "diffusion_pytorch_model-00004-of-00005.safetensors",
+    "layers.21.attention.norm_k.weight": "diffusion_pytorch_model-00004-of-00005.safetensors",
+    "layers.21.attention.norm_q.weight": "diffusion_pytorch_model-00004-of-00005.safetensors",
+    "layers.21.attention.to_k.weight": "diffusion_pytorch_model-00004-of-00005.safetensors",
+    "layers.21.attention.to_out.0.weight": "diffusion_pytorch_model-00004-of-00005.safetensors",
+    "layers.21.attention.to_q.weight": "diffusion_pytorch_model-00004-of-00005.safetensors",
+    "layers.21.attention.to_v.weight": "diffusion_pytorch_model-00004-of-00005.safetensors",
+    "layers.21.attention_norm1.weight": "diffusion_pytorch_model-00004-of-00005.safetensors",
+    "layers.21.attention_norm2.weight": "diffusion_pytorch_model-00004-of-00005.safetensors",
+    "layers.21.feed_forward.w1.weight": "diffusion_pytorch_model-00004-of-00005.safetensors",
+    "layers.21.feed_forward.w2.weight": "diffusion_pytorch_model-00004-of-00005.safetensors",
+    "layers.21.feed_forward.w3.weight": "diffusion_pytorch_model-00004-of-00005.safetensors",
+    "layers.21.ffn_norm1.weight": "diffusion_pytorch_model-00004-of-00005.safetensors",
+    "layers.21.ffn_norm2.weight": "diffusion_pytorch_model-00004-of-00005.safetensors",
+    "layers.22.adaLN_modulation.0.bias": "diffusion_pytorch_model-00004-of-00005.safetensors",
+    "layers.22.adaLN_modulation.0.weight": "diffusion_pytorch_model-00004-of-00005.safetensors",
+    "layers.22.attention.norm_k.weight": "diffusion_pytorch_model-00004-of-00005.safetensors",
+    "layers.22.attention.norm_q.weight": "diffusion_pytorch_model-00004-of-00005.safetensors",
+    "layers.22.attention.to_k.weight": "diffusion_pytorch_model-00004-of-00005.safetensors",
+    "layers.22.attention.to_out.0.weight": "diffusion_pytorch_model-00004-of-00005.safetensors",
+    "layers.22.attention.to_q.weight": "diffusion_pytorch_model-00004-of-00005.safetensors",
+    "layers.22.attention.to_v.weight": "diffusion_pytorch_model-00004-of-00005.safetensors",
+    "layers.22.attention_norm1.weight": "diffusion_pytorch_model-00004-of-00005.safetensors",
+    "layers.22.attention_norm2.weight": "diffusion_pytorch_model-00004-of-00005.safetensors",
+    "layers.22.feed_forward.w1.weight": "diffusion_pytorch_model-00004-of-00005.safetensors",
+    "layers.22.feed_forward.w2.weight": "diffusion_pytorch_model-00004-of-00005.safetensors",
+    "layers.22.feed_forward.w3.weight": "diffusion_pytorch_model-00004-of-00005.safetensors",
+    "layers.22.ffn_norm1.weight": "diffusion_pytorch_model-00004-of-00005.safetensors",
+    "layers.22.ffn_norm2.weight": "diffusion_pytorch_model-00004-of-00005.safetensors",
+    "layers.23.adaLN_modulation.0.bias": "diffusion_pytorch_model-00004-of-00005.safetensors",
+    "layers.23.adaLN_modulation.0.weight": "diffusion_pytorch_model-00004-of-00005.safetensors",
+    "layers.23.attention.norm_k.weight": "diffusion_pytorch_model-00004-of-00005.safetensors",
+    "layers.23.attention.norm_q.weight": "diffusion_pytorch_model-00004-of-00005.safetensors",
+    "layers.23.attention.to_k.weight": "diffusion_pytorch_model-00004-of-00005.safetensors",
+    "layers.23.attention.to_out.0.weight": "diffusion_pytorch_model-00004-of-00005.safetensors",
+    "layers.23.attention.to_q.weight": "diffusion_pytorch_model-00004-of-00005.safetensors",
+    "layers.23.attention.to_v.weight": "diffusion_pytorch_model-00004-of-00005.safetensors",
+    "layers.23.attention_norm1.weight": "diffusion_pytorch_model-00004-of-00005.safetensors",
+    "layers.23.attention_norm2.weight": "diffusion_pytorch_model-00004-of-00005.safetensors",
+    "layers.23.feed_forward.w1.weight": "diffusion_pytorch_model-00004-of-00005.safetensors",
+    "layers.23.feed_forward.w2.weight": "diffusion_pytorch_model-00004-of-00005.safetensors",
+    "layers.23.feed_forward.w3.weight": "diffusion_pytorch_model-00004-of-00005.safetensors",
+    "layers.23.ffn_norm1.weight": "diffusion_pytorch_model-00004-of-00005.safetensors",
+    "layers.23.ffn_norm2.weight": "diffusion_pytorch_model-00004-of-00005.safetensors",
+    "layers.24.adaLN_modulation.0.bias": "diffusion_pytorch_model-00004-of-00005.safetensors",
+    "layers.24.adaLN_modulation.0.weight": "diffusion_pytorch_model-00004-of-00005.safetensors",
+    "layers.24.attention.norm_k.weight": "diffusion_pytorch_model-00004-of-00005.safetensors",
+    "layers.24.attention.norm_q.weight": "diffusion_pytorch_model-00004-of-00005.safetensors",
+    "layers.24.attention.to_k.weight": "diffusion_pytorch_model-00004-of-00005.safetensors",
+    "layers.24.attention.to_out.0.weight": "diffusion_pytorch_model-00004-of-00005.safetensors",
+    "layers.24.attention.to_q.weight": "diffusion_pytorch_model-00004-of-00005.safetensors",
+    "layers.24.attention.to_v.weight": "diffusion_pytorch_model-00004-of-00005.safetensors",
+    "layers.24.attention_norm1.weight": "diffusion_pytorch_model-00004-of-00005.safetensors",
+    "layers.24.attention_norm2.weight": "diffusion_pytorch_model-00004-of-00005.safetensors",
+    "layers.24.feed_forward.w1.weight": "diffusion_pytorch_model-00004-of-00005.safetensors",
+    "layers.24.feed_forward.w2.weight": "diffusion_pytorch_model-00004-of-00005.safetensors",
+    "layers.24.feed_forward.w3.weight": "diffusion_pytorch_model-00004-of-00005.safetensors",
+    "layers.24.ffn_norm1.weight": "diffusion_pytorch_model-00004-of-00005.safetensors",
+    "layers.24.ffn_norm2.weight": "diffusion_pytorch_model-00004-of-00005.safetensors",
+    "layers.25.adaLN_modulation.0.bias": "diffusion_pytorch_model-00004-of-00005.safetensors",
+    "layers.25.adaLN_modulation.0.weight": "diffusion_pytorch_model-00004-of-00005.safetensors",
+    "layers.25.attention.norm_k.weight": "diffusion_pytorch_model-00004-of-00005.safetensors",
+    "layers.25.attention.norm_q.weight": "diffusion_pytorch_model-00004-of-00005.safetensors",
+    "layers.25.attention.to_k.weight": "diffusion_pytorch_model-00004-of-00005.safetensors",
+    "layers.25.attention.to_out.0.weight": "diffusion_pytorch_model-00004-of-00005.safetensors",
+    "layers.25.attention.to_q.weight": "diffusion_pytorch_model-00004-of-00005.safetensors",
+    "layers.25.attention.to_v.weight": "diffusion_pytorch_model-00004-of-00005.safetensors",
+    "layers.25.attention_norm1.weight": "diffusion_pytorch_model-00004-of-00005.safetensors",
+    "layers.25.attention_norm2.weight": "diffusion_pytorch_model-00004-of-00005.safetensors",
+    "layers.25.feed_forward.w1.weight": "diffusion_pytorch_model-00004-of-00005.safetensors",
+    "layers.25.feed_forward.w2.weight": "diffusion_pytorch_model-00004-of-00005.safetensors",
+    "layers.25.feed_forward.w3.weight": "diffusion_pytorch_model-00004-of-00005.safetensors",
+    "layers.25.ffn_norm1.weight": "diffusion_pytorch_model-00004-of-00005.safetensors",
+    "layers.25.ffn_norm2.weight": "diffusion_pytorch_model-00004-of-00005.safetensors",
+    "layers.26.adaLN_modulation.0.bias": "diffusion_pytorch_model-00004-of-00005.safetensors",
+    "layers.26.adaLN_modulation.0.weight": "diffusion_pytorch_model-00004-of-00005.safetensors",
+    "layers.26.attention.norm_k.weight": "diffusion_pytorch_model-00004-of-00005.safetensors",
+    "layers.26.attention.norm_q.weight": "diffusion_pytorch_model-00004-of-00005.safetensors",
+    "layers.26.attention.to_k.weight": "diffusion_pytorch_model-00004-of-00005.safetensors",
+    "layers.26.attention.to_out.0.weight": "diffusion_pytorch_model-00004-of-00005.safetensors",
+    "layers.26.attention.to_q.weight": "diffusion_pytorch_model-00004-of-00005.safetensors",
+    "layers.26.attention.to_v.weight": "diffusion_pytorch_model-00004-of-00005.safetensors",
+    "layers.26.attention_norm1.weight": "diffusion_pytorch_model-00004-of-00005.safetensors",
+    "layers.26.attention_norm2.weight": "diffusion_pytorch_model-00004-of-00005.safetensors",
+    "layers.26.feed_forward.w1.weight": "diffusion_pytorch_model-00004-of-00005.safetensors",
+    "layers.26.feed_forward.w2.weight": "diffusion_pytorch_model-00004-of-00005.safetensors",
+    "layers.26.feed_forward.w3.weight": "diffusion_pytorch_model-00004-of-00005.safetensors",
+    "layers.26.ffn_norm1.weight": "diffusion_pytorch_model-00004-of-00005.safetensors",
+    "layers.26.ffn_norm2.weight": "diffusion_pytorch_model-00004-of-00005.safetensors",
+    "layers.27.adaLN_modulation.0.bias": "diffusion_pytorch_model-00004-of-00005.safetensors",
+    "layers.27.adaLN_modulation.0.weight": "diffusion_pytorch_model-00004-of-00005.safetensors",
+    "layers.27.attention.norm_k.weight": "diffusion_pytorch_model-00004-of-00005.safetensors",
+    "layers.27.attention.norm_q.weight": "diffusion_pytorch_model-00004-of-00005.safetensors",
+    "layers.27.attention.to_k.weight": "diffusion_pytorch_model-00004-of-00005.safetensors",
+    "layers.27.attention.to_out.0.weight": "diffusion_pytorch_model-00004-of-00005.safetensors",
+    "layers.27.attention.to_q.weight": "diffusion_pytorch_model-00004-of-00005.safetensors",
+    "layers.27.attention.to_v.weight": "diffusion_pytorch_model-00004-of-00005.safetensors",
+    "layers.27.attention_norm1.weight": "diffusion_pytorch_model-00004-of-00005.safetensors",
+    "layers.27.attention_norm2.weight": "diffusion_pytorch_model-00004-of-00005.safetensors",
+    "layers.27.feed_forward.w1.weight": "diffusion_pytorch_model-00004-of-00005.safetensors",
+    "layers.27.feed_forward.w2.weight": "diffusion_pytorch_model-00004-of-00005.safetensors",
+    "layers.27.feed_forward.w3.weight": "diffusion_pytorch_model-00004-of-00005.safetensors",
+    "layers.27.ffn_norm1.weight": "diffusion_pytorch_model-00004-of-00005.safetensors",
+    "layers.27.ffn_norm2.weight": "diffusion_pytorch_model-00004-of-00005.safetensors",
+    "layers.28.adaLN_modulation.0.bias": "diffusion_pytorch_model-00005-of-00005.safetensors",
+    "layers.28.adaLN_modulation.0.weight": "diffusion_pytorch_model-00005-of-00005.safetensors",
+    "layers.28.attention.norm_k.weight": "diffusion_pytorch_model-00004-of-00005.safetensors",
+    "layers.28.attention.norm_q.weight": "diffusion_pytorch_model-00004-of-00005.safetensors",
+    "layers.28.attention.to_k.weight": "diffusion_pytorch_model-00004-of-00005.safetensors",
+    "layers.28.attention.to_out.0.weight": "diffusion_pytorch_model-00004-of-00005.safetensors",
+    "layers.28.attention.to_q.weight": "diffusion_pytorch_model-00004-of-00005.safetensors",
+    "layers.28.attention.to_v.weight": "diffusion_pytorch_model-00004-of-00005.safetensors",
+    "layers.28.attention_norm1.weight": "diffusion_pytorch_model-00005-of-00005.safetensors",
+    "layers.28.attention_norm2.weight": "diffusion_pytorch_model-00005-of-00005.safetensors",
+    "layers.28.feed_forward.w1.weight": "diffusion_pytorch_model-00004-of-00005.safetensors",
+    "layers.28.feed_forward.w2.weight": "diffusion_pytorch_model-00004-of-00005.safetensors",
+    "layers.28.feed_forward.w3.weight": "diffusion_pytorch_model-00005-of-00005.safetensors",
+    "layers.28.ffn_norm1.weight": "diffusion_pytorch_model-00005-of-00005.safetensors",
+    "layers.28.ffn_norm2.weight": "diffusion_pytorch_model-00005-of-00005.safetensors",
+    "layers.29.adaLN_modulation.0.bias": "diffusion_pytorch_model-00005-of-00005.safetensors",
+    "layers.29.adaLN_modulation.0.weight": "diffusion_pytorch_model-00005-of-00005.safetensors",
+    "layers.29.attention.norm_k.weight": "diffusion_pytorch_model-00005-of-00005.safetensors",
+    "layers.29.attention.norm_q.weight": "diffusion_pytorch_model-00005-of-00005.safetensors",
+    "layers.29.attention.to_k.weight": "diffusion_pytorch_model-00005-of-00005.safetensors",
+    "layers.29.attention.to_out.0.weight": "diffusion_pytorch_model-00005-of-00005.safetensors",
+    "layers.29.attention.to_q.weight": "diffusion_pytorch_model-00005-of-00005.safetensors",
+    "layers.29.attention.to_v.weight": "diffusion_pytorch_model-00005-of-00005.safetensors",
+    "layers.29.attention_norm1.weight": "diffusion_pytorch_model-00005-of-00005.safetensors",
+    "layers.29.attention_norm2.weight": "diffusion_pytorch_model-00005-of-00005.safetensors",
+    "layers.29.feed_forward.w1.weight": "diffusion_pytorch_model-00005-of-00005.safetensors",
+    "layers.29.feed_forward.w2.weight": "diffusion_pytorch_model-00005-of-00005.safetensors",
+    "layers.29.feed_forward.w3.weight": "diffusion_pytorch_model-00005-of-00005.safetensors",
+    "layers.29.ffn_norm1.weight": "diffusion_pytorch_model-00005-of-00005.safetensors",
+    "layers.29.ffn_norm2.weight": "diffusion_pytorch_model-00005-of-00005.safetensors",
+    "layers.3.adaLN_modulation.0.bias": "diffusion_pytorch_model-00001-of-00005.safetensors",
+    "layers.3.adaLN_modulation.0.weight": "diffusion_pytorch_model-00001-of-00005.safetensors",
+    "layers.3.attention.norm_k.weight": "diffusion_pytorch_model-00001-of-00005.safetensors",
+    "layers.3.attention.norm_q.weight": "diffusion_pytorch_model-00001-of-00005.safetensors",
+    "layers.3.attention.to_k.weight": "diffusion_pytorch_model-00001-of-00005.safetensors",
+    "layers.3.attention.to_out.0.weight": "diffusion_pytorch_model-00001-of-00005.safetensors",
+    "layers.3.attention.to_q.weight": "diffusion_pytorch_model-00001-of-00005.safetensors",
+    "layers.3.attention.to_v.weight": "diffusion_pytorch_model-00001-of-00005.safetensors",
+    "layers.3.attention_norm1.weight": "diffusion_pytorch_model-00001-of-00005.safetensors",
+    "layers.3.attention_norm2.weight": "diffusion_pytorch_model-00001-of-00005.safetensors",
+    "layers.3.feed_forward.w1.weight": "diffusion_pytorch_model-00001-of-00005.safetensors",
+    "layers.3.feed_forward.w2.weight": "diffusion_pytorch_model-00001-of-00005.safetensors",
+    "layers.3.feed_forward.w3.weight": "diffusion_pytorch_model-00001-of-00005.safetensors",
+    "layers.3.ffn_norm1.weight": "diffusion_pytorch_model-00001-of-00005.safetensors",
+    "layers.3.ffn_norm2.weight": "diffusion_pytorch_model-00001-of-00005.safetensors",
+    "layers.4.adaLN_modulation.0.bias": "diffusion_pytorch_model-00002-of-00005.safetensors",
+    "layers.4.adaLN_modulation.0.weight": "diffusion_pytorch_model-00002-of-00005.safetensors",
+    "layers.4.attention.norm_k.weight": "diffusion_pytorch_model-00001-of-00005.safetensors",
+    "layers.4.attention.norm_q.weight": "diffusion_pytorch_model-00001-of-00005.safetensors",
+    "layers.4.attention.to_k.weight": "diffusion_pytorch_model-00001-of-00005.safetensors",
+    "layers.4.attention.to_out.0.weight": "diffusion_pytorch_model-00002-of-00005.safetensors",
+    "layers.4.attention.to_q.weight": "diffusion_pytorch_model-00001-of-00005.safetensors",
+    "layers.4.attention.to_v.weight": "diffusion_pytorch_model-00001-of-00005.safetensors",
+    "layers.4.attention_norm1.weight": "diffusion_pytorch_model-00002-of-00005.safetensors",
+    "layers.4.attention_norm2.weight": "diffusion_pytorch_model-00002-of-00005.safetensors",
+    "layers.4.feed_forward.w1.weight": "diffusion_pytorch_model-00002-of-00005.safetensors",
+    "layers.4.feed_forward.w2.weight": "diffusion_pytorch_model-00002-of-00005.safetensors",
+    "layers.4.feed_forward.w3.weight": "diffusion_pytorch_model-00002-of-00005.safetensors",
+    "layers.4.ffn_norm1.weight": "diffusion_pytorch_model-00002-of-00005.safetensors",
+    "layers.4.ffn_norm2.weight": "diffusion_pytorch_model-00002-of-00005.safetensors",
+    "layers.5.adaLN_modulation.0.bias": "diffusion_pytorch_model-00002-of-00005.safetensors",
+    "layers.5.adaLN_modulation.0.weight": "diffusion_pytorch_model-00002-of-00005.safetensors",
+    "layers.5.attention.norm_k.weight": "diffusion_pytorch_model-00002-of-00005.safetensors",
+    "layers.5.attention.norm_q.weight": "diffusion_pytorch_model-00002-of-00005.safetensors",
+    "layers.5.attention.to_k.weight": "diffusion_pytorch_model-00002-of-00005.safetensors",
+    "layers.5.attention.to_out.0.weight": "diffusion_pytorch_model-00002-of-00005.safetensors",
+    "layers.5.attention.to_q.weight": "diffusion_pytorch_model-00002-of-00005.safetensors",
+    "layers.5.attention.to_v.weight": "diffusion_pytorch_model-00002-of-00005.safetensors",
+    "layers.5.attention_norm1.weight": "diffusion_pytorch_model-00002-of-00005.safetensors",
+    "layers.5.attention_norm2.weight": "diffusion_pytorch_model-00002-of-00005.safetensors",
+    "layers.5.feed_forward.w1.weight": "diffusion_pytorch_model-00002-of-00005.safetensors",
+    "layers.5.feed_forward.w2.weight": "diffusion_pytorch_model-00002-of-00005.safetensors",
+    "layers.5.feed_forward.w3.weight": "diffusion_pytorch_model-00002-of-00005.safetensors",
+    "layers.5.ffn_norm1.weight": "diffusion_pytorch_model-00002-of-00005.safetensors",
+    "layers.5.ffn_norm2.weight": "diffusion_pytorch_model-00002-of-00005.safetensors",
+    "layers.6.adaLN_modulation.0.bias": "diffusion_pytorch_model-00002-of-00005.safetensors",
+    "layers.6.adaLN_modulation.0.weight": "diffusion_pytorch_model-00002-of-00005.safetensors",
+    "layers.6.attention.norm_k.weight": "diffusion_pytorch_model-00002-of-00005.safetensors",
+    "layers.6.attention.norm_q.weight": "diffusion_pytorch_model-00002-of-00005.safetensors",
+    "layers.6.attention.to_k.weight": "diffusion_pytorch_model-00002-of-00005.safetensors",
+    "layers.6.attention.to_out.0.weight": "diffusion_pytorch_model-00002-of-00005.safetensors",
+    "layers.6.attention.to_q.weight": "diffusion_pytorch_model-00002-of-00005.safetensors",
+    "layers.6.attention.to_v.weight": "diffusion_pytorch_model-00002-of-00005.safetensors",
+    "layers.6.attention_norm1.weight": "diffusion_pytorch_model-00002-of-00005.safetensors",
+    "layers.6.attention_norm2.weight": "diffusion_pytorch_model-00002-of-00005.safetensors",
+    "layers.6.feed_forward.w1.weight": "diffusion_pytorch_model-00002-of-00005.safetensors",
+    "layers.6.feed_forward.w2.weight": "diffusion_pytorch_model-00002-of-00005.safetensors",
+    "layers.6.feed_forward.w3.weight": "diffusion_pytorch_model-00002-of-00005.safetensors",
+    "layers.6.ffn_norm1.weight": "diffusion_pytorch_model-00002-of-00005.safetensors",
+    "layers.6.ffn_norm2.weight": "diffusion_pytorch_model-00002-of-00005.safetensors",
+    "layers.7.adaLN_modulation.0.bias": "diffusion_pytorch_model-00002-of-00005.safetensors",
+    "layers.7.adaLN_modulation.0.weight": "diffusion_pytorch_model-00002-of-00005.safetensors",
+    "layers.7.attention.norm_k.weight": "diffusion_pytorch_model-00002-of-00005.safetensors",
+    "layers.7.attention.norm_q.weight": "diffusion_pytorch_model-00002-of-00005.safetensors",
+    "layers.7.attention.to_k.weight": "diffusion_pytorch_model-00002-of-00005.safetensors",
+    "layers.7.attention.to_out.0.weight": "diffusion_pytorch_model-00002-of-00005.safetensors",
+    "layers.7.attention.to_q.weight": "diffusion_pytorch_model-00002-of-00005.safetensors",
+    "layers.7.attention.to_v.weight": "diffusion_pytorch_model-00002-of-00005.safetensors",
+    "layers.7.attention_norm1.weight": "diffusion_pytorch_model-00002-of-00005.safetensors",
+    "layers.7.attention_norm2.weight": "diffusion_pytorch_model-00002-of-00005.safetensors",
+    "layers.7.feed_forward.w1.weight": "diffusion_pytorch_model-00002-of-00005.safetensors",
+    "layers.7.feed_forward.w2.weight": "diffusion_pytorch_model-00002-of-00005.safetensors",
+    "layers.7.feed_forward.w3.weight": "diffusion_pytorch_model-00002-of-00005.safetensors",
+    "layers.7.ffn_norm1.weight": "diffusion_pytorch_model-00002-of-00005.safetensors",
+    "layers.7.ffn_norm2.weight": "diffusion_pytorch_model-00002-of-00005.safetensors",
+    "layers.8.adaLN_modulation.0.bias": "diffusion_pytorch_model-00002-of-00005.safetensors",
+    "layers.8.adaLN_modulation.0.weight": "diffusion_pytorch_model-00002-of-00005.safetensors",
+    "layers.8.attention.norm_k.weight": "diffusion_pytorch_model-00002-of-00005.safetensors",
+    "layers.8.attention.norm_q.weight": "diffusion_pytorch_model-00002-of-00005.safetensors",
+    "layers.8.attention.to_k.weight": "diffusion_pytorch_model-00002-of-00005.safetensors",
+    "layers.8.attention.to_out.0.weight": "diffusion_pytorch_model-00002-of-00005.safetensors",
+    "layers.8.attention.to_q.weight": "diffusion_pytorch_model-00002-of-00005.safetensors",
+    "layers.8.attention.to_v.weight": "diffusion_pytorch_model-00002-of-00005.safetensors",
+    "layers.8.attention_norm1.weight": "diffusion_pytorch_model-00002-of-00005.safetensors",
+    "layers.8.attention_norm2.weight": "diffusion_pytorch_model-00002-of-00005.safetensors",
+    "layers.8.feed_forward.w1.weight": "diffusion_pytorch_model-00002-of-00005.safetensors",
+    "layers.8.feed_forward.w2.weight": "diffusion_pytorch_model-00002-of-00005.safetensors",
+    "layers.8.feed_forward.w3.weight": "diffusion_pytorch_model-00002-of-00005.safetensors",
+    "layers.8.ffn_norm1.weight": "diffusion_pytorch_model-00002-of-00005.safetensors",
+    "layers.8.ffn_norm2.weight": "diffusion_pytorch_model-00002-of-00005.safetensors",
+    "layers.9.adaLN_modulation.0.bias": "diffusion_pytorch_model-00002-of-00005.safetensors",
+    "layers.9.adaLN_modulation.0.weight": "diffusion_pytorch_model-00002-of-00005.safetensors",
+    "layers.9.attention.norm_k.weight": "diffusion_pytorch_model-00002-of-00005.safetensors",
+    "layers.9.attention.norm_q.weight": "diffusion_pytorch_model-00002-of-00005.safetensors",
+    "layers.9.attention.to_k.weight": "diffusion_pytorch_model-00002-of-00005.safetensors",
+    "layers.9.attention.to_out.0.weight": "diffusion_pytorch_model-00002-of-00005.safetensors",
+    "layers.9.attention.to_q.weight": "diffusion_pytorch_model-00002-of-00005.safetensors",
+    "layers.9.attention.to_v.weight": "diffusion_pytorch_model-00002-of-00005.safetensors",
+    "layers.9.attention_norm1.weight": "diffusion_pytorch_model-00002-of-00005.safetensors",
+    "layers.9.attention_norm2.weight": "diffusion_pytorch_model-00002-of-00005.safetensors",
+    "layers.9.feed_forward.w1.weight": "diffusion_pytorch_model-00002-of-00005.safetensors",
+    "layers.9.feed_forward.w2.weight": "diffusion_pytorch_model-00002-of-00005.safetensors",
+    "layers.9.feed_forward.w3.weight": "diffusion_pytorch_model-00002-of-00005.safetensors",
+    "layers.9.ffn_norm1.weight": "diffusion_pytorch_model-00002-of-00005.safetensors",
+    "layers.9.ffn_norm2.weight": "diffusion_pytorch_model-00002-of-00005.safetensors",
+    "noise_refiner.0.adaLN_modulation.0.bias": "diffusion_pytorch_model-00001-of-00005.safetensors",
+    "noise_refiner.0.adaLN_modulation.0.weight": "diffusion_pytorch_model-00001-of-00005.safetensors",
+    "noise_refiner.0.attention.norm_k.weight": "diffusion_pytorch_model-00001-of-00005.safetensors",
+    "noise_refiner.0.attention.norm_q.weight": "diffusion_pytorch_model-00001-of-00005.safetensors",
+    "noise_refiner.0.attention.to_k.weight": "diffusion_pytorch_model-00001-of-00005.safetensors",
+    "noise_refiner.0.attention.to_out.0.weight": "diffusion_pytorch_model-00001-of-00005.safetensors",
+    "noise_refiner.0.attention.to_q.weight": "diffusion_pytorch_model-00001-of-00005.safetensors",
+    "noise_refiner.0.attention.to_v.weight": "diffusion_pytorch_model-00001-of-00005.safetensors",
+    "noise_refiner.0.attention_norm1.weight": "diffusion_pytorch_model-00001-of-00005.safetensors",
+    "noise_refiner.0.attention_norm2.weight": "diffusion_pytorch_model-00001-of-00005.safetensors",
+    "noise_refiner.0.feed_forward.w1.weight": "diffusion_pytorch_model-00001-of-00005.safetensors",
+    "noise_refiner.0.feed_forward.w2.weight": "diffusion_pytorch_model-00001-of-00005.safetensors",
+    "noise_refiner.0.feed_forward.w3.weight": "diffusion_pytorch_model-00001-of-00005.safetensors",
+    "noise_refiner.0.ffn_norm1.weight": "diffusion_pytorch_model-00001-of-00005.safetensors",
+    "noise_refiner.0.ffn_norm2.weight": "diffusion_pytorch_model-00001-of-00005.safetensors",
+    "noise_refiner.1.adaLN_modulation.0.bias": "diffusion_pytorch_model-00001-of-00005.safetensors",
+    "noise_refiner.1.adaLN_modulation.0.weight": "diffusion_pytorch_model-00001-of-00005.safetensors",
+    "noise_refiner.1.attention.norm_k.weight": "diffusion_pytorch_model-00001-of-00005.safetensors",
+    "noise_refiner.1.attention.norm_q.weight": "diffusion_pytorch_model-00001-of-00005.safetensors",
+    "noise_refiner.1.attention.to_k.weight": "diffusion_pytorch_model-00001-of-00005.safetensors",
+    "noise_refiner.1.attention.to_out.0.weight": "diffusion_pytorch_model-00001-of-00005.safetensors",
+    "noise_refiner.1.attention.to_q.weight": "diffusion_pytorch_model-00001-of-00005.safetensors",
+    "noise_refiner.1.attention.to_v.weight": "diffusion_pytorch_model-00001-of-00005.safetensors",
+    "noise_refiner.1.attention_norm1.weight": "diffusion_pytorch_model-00001-of-00005.safetensors",
+    "noise_refiner.1.attention_norm2.weight": "diffusion_pytorch_model-00001-of-00005.safetensors",
+    "noise_refiner.1.feed_forward.w1.weight": "diffusion_pytorch_model-00001-of-00005.safetensors",
+    "noise_refiner.1.feed_forward.w2.weight": "diffusion_pytorch_model-00001-of-00005.safetensors",
+    "noise_refiner.1.feed_forward.w3.weight": "diffusion_pytorch_model-00001-of-00005.safetensors",
+    "noise_refiner.1.ffn_norm1.weight": "diffusion_pytorch_model-00001-of-00005.safetensors",
+    "noise_refiner.1.ffn_norm2.weight": "diffusion_pytorch_model-00001-of-00005.safetensors",
+    "t_embedder.mlp.0.bias": "diffusion_pytorch_model-00001-of-00005.safetensors",
+    "t_embedder.mlp.0.weight": "diffusion_pytorch_model-00001-of-00005.safetensors",
+    "t_embedder.mlp.2.bias": "diffusion_pytorch_model-00001-of-00005.safetensors",
+    "t_embedder.mlp.2.weight": "diffusion_pytorch_model-00001-of-00005.safetensors"
+  }
+}

vae/config.json

@@ -0,0 +1,25 @@
+{
+  "_class_name": "AutoencoderKLQwenImage",
+  "_diffusers_version": "0.36.0.dev0",
+  "attn_scales": [],
+  "base_dim": 96,
+  "dim_mult": [
+    1,
+    2,
+    4,
+    4
+  ],
+  "dropout": 0.0,
+  "input_channels": 4,
+  "latents_mean": null,
+  "latents_std": null,
+  "scaling_factor": 8.0064,
+  "shift_factor": 0.0,
+  "num_res_blocks": 2,
+  "temperal_downsample": [
+    false,
+    true,
+    true
+  ],
+  "z_dim": 16
+}