init models

.gitattributes

@@ -33,3 +33,5 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+demo.png filter=lfs diff=lfs merge=lfs -text
+tokenizer/tokenizer.json filter=lfs diff=lfs merge=lfs -text

README.md

@@ -1,3 +1,170 @@
----
-license: apache-2.0
----
+<h1 align="center">LLaDA-Image</h1>
+
+<p align="center">
+  Welcome to the official repository for LLaDA-Image, a unified model for high-quality image generation and editing.
+</p>
+
+<p align="center">
+  <a href="https://github.com/inclusionAI/LLaDA-Image"><img src="https://img.shields.io/badge/GitHub-LLaDA--Image-181717?logo=github" alt="GitHub"></a>
+  <a href="https://huggingface.co/inclusionAI/LLaDA-Image"><img src="https://img.shields.io/badge/Hugging%20Face-Base-FFD21E?logo=huggingface" alt="LLaDA-Image Base on Hugging Face"></a>
+  <a href="https://huggingface.co/inclusionAI/LLaDA-Image-Turbo"><img src="https://img.shields.io/badge/Hugging%20Face-Turbo-FFD21E?logo=huggingface" alt="LLaDA-Image Turbo on Hugging Face"></a>
+  <a href="#"><img src="https://img.shields.io/badge/arXiv-Coming%20Soon-B31B1B?logo=arxiv" alt="arXiv"></a>
+</p>
+
+<p align="center">
+  <img src="demo.png" alt="LLaDA-Image showcase" width="100%">
+</p>
+
+## Introduction
+
+LLaDA-Image is a competitive 6B-parameter open-source unified image generation and editing model family. It includes **LLaDA-Image**, a 50-step Base model for high-quality text-to-image generation and instruction-guided editing, and **LLaDA-Image-Turbo**, a 4-step distilled model for fast generation and editing. Both variants support practical text-to-image generation, VQ-conditioned generation, reference-image editing, and Chinese--English text rendering.
+
+This repository provides the checkpoints and Diffusers-based inference code for the LLaDA-Image model family.
+
+## News
+
+- **TODO:** We released the LLaDA-Image Base and Turbo checkpoints together with the inference code.
+
+## Highlights
+
+- **Unified generation and editing.** A single checkpoint supports text-to-image generation and reference-preserving, instruction-guided editing without a separate editing backbone.
+- **Realistic image generation.** LLaDA-Image produces high-quality images with rich visual details, natural lighting, and coherent compositions.
+- **Image-only pre-training for visual-prior learning.** The report establishes the visual prior through image-only pre-training and mid-training before introducing paired language supervision and joint generation--editing training.
+- **Efficient inference with distilled model.** LLaDA-Image-Turbo uses Twin-DMD distillation to deliver fast image generation and editing in only 2--4 sampling steps.
+- **SOTA on Qwen-Image-Bench.** LLaDA-Image achieves state-of-the-art overall scores of 53.53 in English and 53.38 in Chinese.
+
+## Model Zoo
+
+| Model | Description | Sampling steps | Hugging Face |
+| --- | --- | ---: | --- |
+| **LLaDA-Image** | Base model for high-fidelity text-to-image generation and instruction-guided editing. | 50 | [inclusionAI/LLaDA-Image](https://huggingface.co/inclusionAI/LLaDA-Image) |
+| **LLaDA-Image-Turbo** | Distilled model for fast generation and editing. | 4 | [inclusionAI/LLaDA-Image-Turbo](https://huggingface.co/inclusionAI/LLaDA-Image-Turbo) |
+
+## Quick Start
+
+### 1. Create an environment
+
+The implementation has been used with Python 3.11, PyTorch 2.8, Transformers 4.57.6, and Diffusers 0.39.0.
+
+```bash
+git clone https://github.com/inclusionAI/LLaDA-Image.git
+cd LLaDA-Image
+
+pip install -r requirements.txt
+```
+
+The published LLaDA2 text encoder uses `veomni.ops.fused_moe_forward`. Install a compatible LLaDA2 / VeOmni runtime before running inference.
+
+### 2. Run inference
+
+The pipeline accepts a prompt and, for editing, an optional reference image.
+
+#### LLaDA-Image (Base)
+
+Use the Base checkpoint for high-fidelity generation and editing. Its recommended sampling configuration is **50 steps**.
+
+```python
+import torch
+
+from src import LLaDAImagePipeline
+
+# Load the pipeline. The model is downloaded from Hugging Face on first use.
+pipe = LLaDAImagePipeline.from_pretrained(
+    "inclusionAI/LLaDA-Image",
+    torch_dtype=torch.bfloat16,
+    device="cuda",
+)
+
+# Generate an image.
+prompt = (
+    "A cinematic photograph of a red fox standing in fresh snow, "
+    "soft winter light, detailed fur, shallow depth of field"
+)
+negative_prompt = ""
+
+image = pipe(
+    prompt=prompt,
+    negative_prompt=negative_prompt,
+    generation_mode="text",
+    height=1024,
+    width=1024,
+    num_inference_steps=50,
+    guidance_scale=5.0,
+    generator=torch.Generator("cuda").manual_seed(42),
+).images[0]
+
+image.save("llada-image-base.png")
+```
+
+#### LLaDA-Image-Turbo
+
+Use the Turbo checkpoint for fast generation and editing. Its recommended sampling configuration is **4 steps**.
+
+```python
+import torch
+
+from src import LLaDAImagePipeline
+
+# Load the distilled Turbo checkpoint.
+pipe = LLaDAImagePipeline.from_pretrained(
+    "inclusionAI/LLaDA-Image-Turbo",
+    torch_dtype=torch.bfloat16,
+    device="cuda",
+)
+
+prompt = "A quiet observatory above a sea of clouds at sunrise, golden light, wide-angle photograph"
+
+image = pipe(
+    prompt=prompt,
+    generation_mode="text",
+    height=1024,
+    width=1024,
+    num_inference_steps=4,
+    guidance_scale=1.0,
+    generator=torch.Generator("cuda").manual_seed(42),
+).images[0]
+
+image.save("llada-image-turbo.png")
+```
+
+#### Generation modes
+
+Both checkpoints support the following modes. Text and VQ-conditioned generation require height and width divisible by 16; image editing requires dimensions divisible by 32.
+
+**VQ-conditioned generation** uses the LLaDA2 model to produce image VQ tokens from the prompt, which SigVQ embeds before diffusion. Do not provide an input image in VQ mode.
+
+```python
+image = pipe(
+    prompt="A quiet observatory above a sea of clouds at sunrise",
+    generation_mode="vq",
+    height=1024,
+    width=1024,
+    num_inference_steps=50,  # Use 4 for LLaDA-Image-Turbo.
+    guidance_scale=5.0,  # Use 1.0 for few-step inference.
+    generator=torch.Generator("cuda").manual_seed(42),
+).images[0]
+```
+
+**Image editing** requires a reference image:
+
+```python
+from diffusers.utils import load_image
+
+reference_image = load_image("/path/to/input.png")
+image = pipe(
+    prompt="Turn it into a watercolor painting",
+    image=reference_image,
+    generation_mode="editing",
+    height=1024,
+    width=1024,
+    num_inference_steps=50,  # Use 4 for LLaDA-Image-Turbo.
+    guidance_scale=5.0,  # Use 1.0 for few-step inference.
+    generator=torch.Generator("cuda").manual_seed(43),
+).images[0]
+```
+
+## Citation
+
+If you find LLaDA-Image useful for your research or applications, please consider citing our work.
+
+TODO

model_index.json

@@ -0,0 +1,36 @@
+{
+  "_class_name": "LLaDAImagePipeline",
+  "_diffusers_version": "0.40.0.dev0",
+  "queryformer": [
+    "diffusers",
+    "LLaDAImageQueryFormerModel"
+  ],
+  "scheduler": [
+    "diffusers",
+    "FlowMatchEulerDiscreteScheduler"
+  ],
+  "sigvq": [
+    "diffusers",
+    "LLaDAImageSigVQModel"
+  ],
+  "text_encoder": [
+    "modeling_llada2uni_moe",
+    "LLaDA2MoeModelLM"
+  ],
+  "text_projection": [
+    "diffusers",
+    "LLaDAImageTextProjectionModel"
+  ],
+  "tokenizer": [
+    "transformers",
+    "PreTrainedTokenizerFast"
+  ],
+  "transformer": [
+    "diffusers",
+    "LLaDAImageTransformer2DModel"
+  ],
+  "vae": [
+    "diffusers",
+    "AutoencoderKLFlux2"
+  ]
+}

queryformer/config.json

@@ -0,0 +1,11 @@
+{
+  "_class_name": "LLaDAImageQueryFormerModel",
+  "_diffusers_version": "0.40.0.dev0",
+  "dropout": 0.0,
+  "hidden_size": 2048,
+  "intermediate_size": 8192,
+  "norm_eps": 1e-06,
+  "num_attention_heads": 16,
+  "num_hidden_layers": 1,
+  "num_queries": 256
+}

queryformer/diffusion_pytorch_model.safetensors

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

scheduler/scheduler_config.json

@@ -0,0 +1,18 @@
+{
+  "_class_name": "FlowMatchEulerDiscreteScheduler",
+  "_diffusers_version": "0.40.0.dev0",
+  "base_image_seq_len": 256,
+  "base_shift": 0.5,
+  "invert_sigmas": false,
+  "max_image_seq_len": 4096,
+  "max_shift": 1.15,
+  "num_train_timesteps": 1000,
+  "shift": 1.0,
+  "shift_terminal": null,
+  "stochastic_sampling": false,
+  "time_shift_type": "exponential",
+  "use_beta_sigmas": false,
+  "use_dynamic_shifting": false,
+  "use_exponential_sigmas": false,
+  "use_karras_sigmas": false
+}

sigvq/config.json

@@ -0,0 +1,17 @@
+{
+  "_class_name": "LLaDAImageSigVQModel",
+  "_diffusers_version": "0.40.0.dev0",
+  "attention_bias": true,
+  "attention_dropout": 0.0,
+  "codebook_embed_dim": 2048,
+  "codebook_size": 16384,
+  "hidden_size": 1536,
+  "image_size": 2048,
+  "in_channels": 3,
+  "intermediate_size": 6144,
+  "norm_eps": 1e-06,
+  "num_attention_heads": 16,
+  "num_hidden_layers": 40,
+  "patch_size": 16,
+  "semantic_embed_dim": 4096
+}

sigvq/diffusion_pytorch_model.safetensors

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

text_encoder/config.json

@@ -0,0 +1,59 @@
+{
+  "architectures": [
+    "LLaDA2MoeModelLM"
+  ],
+  "auto_map": {
+    "AutoConfig": "configuration_llada2uni_moe.LLaDA2MoeConfig",
+    "AutoModel": "modeling_llada2uni_moe.LLaDA2MoeModelLM",
+    "AutoModelForCausalLM": "modeling_llada2uni_moe.LLaDA2MoeModelLM"
+  },
+  "attention_dropout": 0.0,
+  "embedding_dropout": 0.0,
+  "first_k_dense_replace": 1,
+  "head_dim": 128,
+  "hidden_act": "silu",
+  "hidden_size": 2048,
+  "image_token_offset": 157184,
+  "initializer_range": 0.02,
+  "intermediate_size": 5120,
+  "max_position_embeddings": 16384,
+  "model_type": "llada2_moe_veomni",
+  "moe_intermediate_size": 512,
+  "moe_router_enable_expert_bias": true,
+  "n_group": 8,
+  "norm_topk_prob": true,
+  "num_attention_heads": 16,
+  "num_experts": 256,
+  "num_experts_per_tok": 8,
+  "num_hidden_layers": 20,
+  "num_key_value_heads": 4,
+  "num_shared_experts": 1,
+  "output_dropout": 0.0,
+  "output_router_logits": false,
+  "pad_token_id": 156892,
+  "partial_rotary_factor": 0.5,
+  "rms_norm_eps": 1e-06,
+  "rope_scaling": {
+    "mrope_section": [
+      16,
+      24,
+      24
+    ],
+    "rope_type": "default",
+    "type": "default"
+  },
+  "rope_theta": 600000,
+  "router_dtype": "fp32",
+  "routed_scaling_factor": 2.5,
+  "score_function": "sigmoid",
+  "sliding_window": null,
+  "tie_word_embeddings": false,
+  "topk_group": 4,
+  "torch_dtype": "bfloat16",
+  "transformers_version": "4.51.0",
+  "use_bias": false,
+  "use_cache": false,
+  "use_qk_norm": true,
+  "use_qkv_bias": false,
+  "vocab_size": 173568
+}

text_encoder/configuration_llada2uni_moe.py

@@ -0,0 +1,133 @@
+# Copyright 2025 Antgroup 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.
+"""LLaDA2 MoE model configuration."""
+
+from transformers.configuration_utils import PretrainedConfig
+
+
+class LLaDA2MoeConfig(PretrainedConfig):
+    r"""
+    Configuration class for the LLaDA2 MoE model.
+
+    ```python
+    >>> from configuration_llada2uni_moe import LLaDA2MoeConfig
+    >>> config = LLaDA2MoeConfig()
+    ```
+    """
+
+    # Keep the original value because it selects the fused-expert implementation.
+    model_type = "llada2_moe_veomni"
+
+    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,
+        head_dim=None,
+        hidden_act="silu",
+        use_qkv_bias=False,
+        use_qk_norm=True,
+        use_bias=True,
+        rms_norm_eps=1e-05,
+        tie_word_embeddings=False,
+        attention_dropout=0.1,
+        initializer_range=0.02,
+        max_position_embeddings=16384,
+        rope_theta=10000.0,
+        rope_parameters=None,
+        rope_scaling=None,
+        partial_rotary_factor=0.5,
+        use_cache=True,
+        sliding_window=None,
+        pad_token_id=126081,
+        # Image
+        image_token_offset=157184,
+        # MoE
+        num_experts=16,
+        num_shared_experts=0,
+        num_experts_per_tok=2,
+        n_group=8,
+        topk_group=4,
+        routed_scaling_factor=2.5,
+        moe_router_enable_expert_bias=True,
+        norm_topk_prob=True,
+        router_dtype="fp32",
+        score_function="sigmoid",
+        moe_intermediate_size=None,
+        first_k_dense_replace=0,
+        output_router_logits=False,
+        **kwargs,
+    ):
+        self.vocab_size = vocab_size
+        self.hidden_size = hidden_size
+        self.intermediate_size = intermediate_size
+        self.num_hidden_layers = num_hidden_layers
+        self.num_attention_heads = num_attention_heads
+        self.num_key_value_heads = num_key_value_heads
+        self.head_dim = head_dim or hidden_size // num_attention_heads
+        self.hidden_act = hidden_act
+        self.use_qkv_bias = use_qkv_bias
+        self.use_qk_norm = use_qk_norm
+        self.use_bias = use_bias
+        self.rms_norm_eps = rms_norm_eps
+        self.attention_dropout = attention_dropout
+        self.initializer_range = initializer_range
+        self.max_position_embeddings = max_position_embeddings
+        self.rope_theta = rope_theta
+        self.rope_scaling = rope_scaling
+        self.partial_rotary_factor = partial_rotary_factor
+        self.use_cache = use_cache
+        self.sliding_window = sliding_window
+
+        # Image token offset: VQ codebook indices are shifted by this amount in the vocabulary
+        self.image_token_offset = image_token_offset
+
+        # RoPE parameters dict — used by LLaDA2MoeRotaryEmbedding
+        if rope_parameters is None:
+            rope_parameters = {
+                "rope_type": "default",
+                "rope_theta": rope_theta,
+                "partial_rotary_factor": partial_rotary_factor,
+            }
+        self.rope_parameters = rope_parameters
+
+        # MoE
+        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.routed_scaling_factor = routed_scaling_factor
+        self.moe_router_enable_expert_bias = moe_router_enable_expert_bias
+        self.norm_topk_prob = norm_topk_prob
+        self.router_dtype = router_dtype
+        self.score_function = score_function
+        self.moe_intermediate_size = moe_intermediate_size
+        self.first_k_dense_replace = first_k_dense_replace
+        self.output_router_logits = output_router_logits
+
+        # FP8 quantization flag — set to True to use FP8Linear for experts
+        self.use_fp8_experts = kwargs.pop("use_fp8_experts", False)
+
+        super().__init__(
+            pad_token_id=pad_token_id,
+            tie_word_embeddings=tie_word_embeddings,
+            **kwargs,
+        )
+
+
+__all__ = ["LLaDA2MoeConfig"]

text_encoder/model-00001-of-00009.safetensors

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

text_encoder/model-00002-of-00009.safetensors

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

text_encoder/model-00003-of-00009.safetensors

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

text_encoder/model-00004-of-00009.safetensors

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

text_encoder/model-00005-of-00009.safetensors

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

text_encoder/model-00006-of-00009.safetensors

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

text_encoder/model-00007-of-00009.safetensors

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

text_encoder/model-00008-of-00009.safetensors

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

text_encoder/model-00009-of-00009.safetensors

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

text_encoder/model.safetensors.index.json

@@ -0,0 +1,285 @@
+{
+  "metadata": {
+    "total_size": 32645504512
+  },
+  "weight_map": {
+    "model.language_model.layers.0.attention.dense.weight": "model-00001-of-00009.safetensors",
+    "model.language_model.layers.0.attention.key_layernorm.weight": "model-00001-of-00009.safetensors",
+    "model.language_model.layers.0.attention.query_key_value.weight": "model-00001-of-00009.safetensors",
+    "model.language_model.layers.0.attention.query_layernorm.weight": "model-00001-of-00009.safetensors",
+    "model.language_model.layers.0.input_layernorm.weight": "model-00001-of-00009.safetensors",
+    "model.language_model.layers.0.mlp.down_proj.weight": "model-00001-of-00009.safetensors",
+    "model.language_model.layers.0.mlp.gate_proj.weight": "model-00001-of-00009.safetensors",
+    "model.language_model.layers.0.mlp.up_proj.weight": "model-00001-of-00009.safetensors",
+    "model.language_model.layers.0.post_attention_layernorm.weight": "model-00001-of-00009.safetensors",
+    "model.language_model.layers.1.attention.dense.weight": "model-00001-of-00009.safetensors",
+    "model.language_model.layers.1.attention.key_layernorm.weight": "model-00001-of-00009.safetensors",
+    "model.language_model.layers.1.attention.query_key_value.weight": "model-00001-of-00009.safetensors",
+    "model.language_model.layers.1.attention.query_layernorm.weight": "model-00001-of-00009.safetensors",
+    "model.language_model.layers.1.input_layernorm.weight": "model-00001-of-00009.safetensors",
+    "model.language_model.layers.1.mlp.experts.down_proj": "model-00001-of-00009.safetensors",
+    "model.language_model.layers.1.mlp.experts.gate_proj": "model-00001-of-00009.safetensors",
+    "model.language_model.layers.1.mlp.experts.up_proj": "model-00001-of-00009.safetensors",
+    "model.language_model.layers.1.mlp.gate.expert_bias": "model-00009-of-00009.safetensors",
+    "model.language_model.layers.1.mlp.gate.weight": "model-00001-of-00009.safetensors",
+    "model.language_model.layers.1.mlp.shared_experts.down_proj.weight": "model-00001-of-00009.safetensors",
+    "model.language_model.layers.1.mlp.shared_experts.gate_proj.weight": "model-00001-of-00009.safetensors",
+    "model.language_model.layers.1.mlp.shared_experts.up_proj.weight": "model-00001-of-00009.safetensors",
+    "model.language_model.layers.1.post_attention_layernorm.weight": "model-00001-of-00009.safetensors",
+    "model.language_model.layers.10.attention.dense.weight": "model-00003-of-00009.safetensors",
+    "model.language_model.layers.10.attention.key_layernorm.weight": "model-00003-of-00009.safetensors",
+    "model.language_model.layers.10.attention.query_key_value.weight": "model-00003-of-00009.safetensors",
+    "model.language_model.layers.10.attention.query_layernorm.weight": "model-00003-of-00009.safetensors",
+    "model.language_model.layers.10.input_layernorm.weight": "model-00003-of-00009.safetensors",
+    "model.language_model.layers.10.mlp.experts.down_proj": "model-00003-of-00009.safetensors",
+    "model.language_model.layers.10.mlp.experts.gate_proj": "model-00003-of-00009.safetensors",
+    "model.language_model.layers.10.mlp.experts.up_proj": "model-00003-of-00009.safetensors",
+    "model.language_model.layers.10.mlp.gate.expert_bias": "model-00009-of-00009.safetensors",
+    "model.language_model.layers.10.mlp.gate.weight": "model-00003-of-00009.safetensors",
+    "model.language_model.layers.10.mlp.shared_experts.down_proj.weight": "model-00003-of-00009.safetensors",
+    "model.language_model.layers.10.mlp.shared_experts.gate_proj.weight": "model-00003-of-00009.safetensors",
+    "model.language_model.layers.10.mlp.shared_experts.up_proj.weight": "model-00003-of-00009.safetensors",
+    "model.language_model.layers.10.post_attention_layernorm.weight": "model-00003-of-00009.safetensors",
+    "model.language_model.layers.11.attention.dense.weight": "model-00003-of-00009.safetensors",
+    "model.language_model.layers.11.attention.key_layernorm.weight": "model-00003-of-00009.safetensors",
+    "model.language_model.layers.11.attention.query_key_value.weight": "model-00003-of-00009.safetensors",
+    "model.language_model.layers.11.attention.query_layernorm.weight": "model-00003-of-00009.safetensors",
+    "model.language_model.layers.11.input_layernorm.weight": "model-00005-of-00009.safetensors",
+    "model.language_model.layers.11.mlp.experts.down_proj": "model-00006-of-00009.safetensors",
+    "model.language_model.layers.11.mlp.experts.gate_proj": "model-00003-of-00009.safetensors",
+    "model.language_model.layers.11.mlp.experts.up_proj": "model-00004-of-00009.safetensors",
+    "model.language_model.layers.11.mlp.gate.expert_bias": "model-00009-of-00009.safetensors",
+    "model.language_model.layers.11.mlp.gate.weight": "model-00006-of-00009.safetensors",
+    "model.language_model.layers.11.mlp.shared_experts.down_proj.weight": "model-00006-of-00009.safetensors",
+    "model.language_model.layers.11.mlp.shared_experts.gate_proj.weight": "model-00006-of-00009.safetensors",
+    "model.language_model.layers.11.mlp.shared_experts.up_proj.weight": "model-00006-of-00009.safetensors",
+    "model.language_model.layers.11.post_attention_layernorm.weight": "model-00006-of-00009.safetensors",
+    "model.language_model.layers.12.attention.dense.weight": "model-00006-of-00009.safetensors",
+    "model.language_model.layers.12.attention.key_layernorm.weight": "model-00006-of-00009.safetensors",
+    "model.language_model.layers.12.attention.query_key_value.weight": "model-00006-of-00009.safetensors",
+    "model.language_model.layers.12.attention.query_layernorm.weight": "model-00006-of-00009.safetensors",
+    "model.language_model.layers.12.input_layernorm.weight": "model-00006-of-00009.safetensors",
+    "model.language_model.layers.12.mlp.experts.down_proj": "model-00006-of-00009.safetensors",
+    "model.language_model.layers.12.mlp.experts.gate_proj": "model-00006-of-00009.safetensors",
+    "model.language_model.layers.12.mlp.experts.up_proj": "model-00006-of-00009.safetensors",
+    "model.language_model.layers.12.mlp.gate.expert_bias": "model-00009-of-00009.safetensors",
+    "model.language_model.layers.12.mlp.gate.weight": "model-00006-of-00009.safetensors",
+    "model.language_model.layers.12.mlp.shared_experts.down_proj.weight": "model-00006-of-00009.safetensors",
+    "model.language_model.layers.12.mlp.shared_experts.gate_proj.weight": "model-00006-of-00009.safetensors",
+    "model.language_model.layers.12.mlp.shared_experts.up_proj.weight": "model-00006-of-00009.safetensors",
+    "model.language_model.layers.12.post_attention_layernorm.weight": "model-00006-of-00009.safetensors",
+    "model.language_model.layers.13.attention.dense.weight": "model-00006-of-00009.safetensors",
+    "model.language_model.layers.13.attention.key_layernorm.weight": "model-00006-of-00009.safetensors",
+    "model.language_model.layers.13.attention.query_key_value.weight": "model-00006-of-00009.safetensors",
+    "model.language_model.layers.13.attention.query_layernorm.weight": "model-00006-of-00009.safetensors",
+    "model.language_model.layers.13.input_layernorm.weight": "model-00006-of-00009.safetensors",
+    "model.language_model.layers.13.mlp.experts.down_proj": "model-00006-of-00009.safetensors",
+    "model.language_model.layers.13.mlp.experts.gate_proj": "model-00006-of-00009.safetensors",
+    "model.language_model.layers.13.mlp.experts.up_proj": "model-00006-of-00009.safetensors",
+    "model.language_model.layers.13.mlp.gate.expert_bias": "model-00009-of-00009.safetensors",
+    "model.language_model.layers.13.mlp.gate.weight": "model-00006-of-00009.safetensors",
+    "model.language_model.layers.13.mlp.shared_experts.down_proj.weight": "model-00006-of-00009.safetensors",
+    "model.language_model.layers.13.mlp.shared_experts.gate_proj.weight": "model-00006-of-00009.safetensors",
+    "model.language_model.layers.13.mlp.shared_experts.up_proj.weight": "model-00006-of-00009.safetensors",
+    "model.language_model.layers.13.post_attention_layernorm.weight": "model-00006-of-00009.safetensors",
+    "model.language_model.layers.14.attention.dense.weight": "model-00006-of-00009.safetensors",
+    "model.language_model.layers.14.attention.key_layernorm.weight": "model-00006-of-00009.safetensors",
+    "model.language_model.layers.14.attention.query_key_value.weight": "model-00006-of-00009.safetensors",
+    "model.language_model.layers.14.attention.query_layernorm.weight": "model-00006-of-00009.safetensors",
+    "model.language_model.layers.14.input_layernorm.weight": "model-00006-of-00009.safetensors",
+    "model.language_model.layers.14.mlp.experts.down_proj": "model-00007-of-00009.safetensors",
+    "model.language_model.layers.14.mlp.experts.gate_proj": "model-00007-of-00009.safetensors",
+    "model.language_model.layers.14.mlp.experts.up_proj": "model-00007-of-00009.safetensors",
+    "model.language_model.layers.14.mlp.gate.expert_bias": "model-00009-of-00009.safetensors",
+    "model.language_model.layers.14.mlp.gate.weight": "model-00007-of-00009.safetensors",
+    "model.language_model.layers.14.mlp.shared_experts.down_proj.weight": "model-00007-of-00009.safetensors",
+    "model.language_model.layers.14.mlp.shared_experts.gate_proj.weight": "model-00007-of-00009.safetensors",
+    "model.language_model.layers.14.mlp.shared_experts.up_proj.weight": "model-00007-of-00009.safetensors",
+    "model.language_model.layers.14.post_attention_layernorm.weight": "model-00007-of-00009.safetensors",
+    "model.language_model.layers.15.attention.dense.weight": "model-00007-of-00009.safetensors",
+    "model.language_model.layers.15.attention.key_layernorm.weight": "model-00007-of-00009.safetensors",
+    "model.language_model.layers.15.attention.query_key_value.weight": "model-00007-of-00009.safetensors",
+    "model.language_model.layers.15.attention.query_layernorm.weight": "model-00007-of-00009.safetensors",
+    "model.language_model.layers.15.input_layernorm.weight": "model-00007-of-00009.safetensors",
+    "model.language_model.layers.15.mlp.experts.down_proj": "model-00007-of-00009.safetensors",
+    "model.language_model.layers.15.mlp.experts.gate_proj": "model-00007-of-00009.safetensors",
+    "model.language_model.layers.15.mlp.experts.up_proj": "model-00007-of-00009.safetensors",
+    "model.language_model.layers.15.mlp.gate.expert_bias": "model-00009-of-00009.safetensors",
+    "model.language_model.layers.15.mlp.gate.weight": "model-00007-of-00009.safetensors",
+    "model.language_model.layers.15.mlp.shared_experts.down_proj.weight": "model-00007-of-00009.safetensors",
+    "model.language_model.layers.15.mlp.shared_experts.gate_proj.weight": "model-00007-of-00009.safetensors",
+    "model.language_model.layers.15.mlp.shared_experts.up_proj.weight": "model-00007-of-00009.safetensors",
+    "model.language_model.layers.15.post_attention_layernorm.weight": "model-00007-of-00009.safetensors",
+    "model.language_model.layers.16.attention.dense.weight": "model-00007-of-00009.safetensors",
+    "model.language_model.layers.16.attention.key_layernorm.weight": "model-00007-of-00009.safetensors",
+    "model.language_model.layers.16.attention.query_key_value.weight": "model-00007-of-00009.safetensors",
+    "model.language_model.layers.16.attention.query_layernorm.weight": "model-00007-of-00009.safetensors",
+    "model.language_model.layers.16.input_layernorm.weight": "model-00007-of-00009.safetensors",
+    "model.language_model.layers.16.mlp.experts.down_proj": "model-00007-of-00009.safetensors",
+    "model.language_model.layers.16.mlp.experts.gate_proj": "model-00008-of-00009.safetensors",
+    "model.language_model.layers.16.mlp.experts.up_proj": "model-00008-of-00009.safetensors",
+    "model.language_model.layers.16.mlp.gate.expert_bias": "model-00009-of-00009.safetensors",
+    "model.language_model.layers.16.mlp.gate.weight": "model-00008-of-00009.safetensors",
+    "model.language_model.layers.16.mlp.shared_experts.down_proj.weight": "model-00008-of-00009.safetensors",
+    "model.language_model.layers.16.mlp.shared_experts.gate_proj.weight": "model-00008-of-00009.safetensors",
+    "model.language_model.layers.16.mlp.shared_experts.up_proj.weight": "model-00008-of-00009.safetensors",
+    "model.language_model.layers.16.post_attention_layernorm.weight": "model-00008-of-00009.safetensors",
+    "model.language_model.layers.17.attention.dense.weight": "model-00008-of-00009.safetensors",
+    "model.language_model.layers.17.attention.key_layernorm.weight": "model-00008-of-00009.safetensors",
+    "model.language_model.layers.17.attention.query_key_value.weight": "model-00008-of-00009.safetensors",
+    "model.language_model.layers.17.attention.query_layernorm.weight": "model-00008-of-00009.safetensors",
+    "model.language_model.layers.17.input_layernorm.weight": "model-00008-of-00009.safetensors",
+    "model.language_model.layers.17.mlp.experts.down_proj": "model-00008-of-00009.safetensors",
+    "model.language_model.layers.17.mlp.experts.gate_proj": "model-00008-of-00009.safetensors",
+    "model.language_model.layers.17.mlp.experts.up_proj": "model-00008-of-00009.safetensors",
+    "model.language_model.layers.17.mlp.gate.expert_bias": "model-00009-of-00009.safetensors",
+    "model.language_model.layers.17.mlp.gate.weight": "model-00008-of-00009.safetensors",
+    "model.language_model.layers.17.mlp.shared_experts.down_proj.weight": "model-00008-of-00009.safetensors",
+    "model.language_model.layers.17.mlp.shared_experts.gate_proj.weight": "model-00008-of-00009.safetensors",
+    "model.language_model.layers.17.mlp.shared_experts.up_proj.weight": "model-00008-of-00009.safetensors",
+    "model.language_model.layers.17.post_attention_layernorm.weight": "model-00008-of-00009.safetensors",
+    "model.language_model.layers.18.attention.dense.weight": "model-00008-of-00009.safetensors",
+    "model.language_model.layers.18.attention.key_layernorm.weight": "model-00008-of-00009.safetensors",
+    "model.language_model.layers.18.attention.query_key_value.weight": "model-00008-of-00009.safetensors",
+    "model.language_model.layers.18.attention.query_layernorm.weight": "model-00008-of-00009.safetensors",
+    "model.language_model.layers.18.input_layernorm.weight": "model-00008-of-00009.safetensors",
+    "model.language_model.layers.18.mlp.experts.down_proj": "model-00008-of-00009.safetensors",
+    "model.language_model.layers.18.mlp.experts.gate_proj": "model-00008-of-00009.safetensors",
+    "model.language_model.layers.18.mlp.experts.up_proj": "model-00009-of-00009.safetensors",
+    "model.language_model.layers.18.mlp.gate.expert_bias": "model-00009-of-00009.safetensors",
+    "model.language_model.layers.18.mlp.gate.weight": "model-00009-of-00009.safetensors",
+    "model.language_model.layers.18.mlp.shared_experts.down_proj.weight": "model-00009-of-00009.safetensors",
+    "model.language_model.layers.18.mlp.shared_experts.gate_proj.weight": "model-00009-of-00009.safetensors",
+    "model.language_model.layers.18.mlp.shared_experts.up_proj.weight": "model-00009-of-00009.safetensors",
+    "model.language_model.layers.18.post_attention_layernorm.weight": "model-00009-of-00009.safetensors",
+    "model.language_model.layers.19.attention.dense.weight": "model-00009-of-00009.safetensors",
+    "model.language_model.layers.19.attention.key_layernorm.weight": "model-00009-of-00009.safetensors",
+    "model.language_model.layers.19.attention.query_key_value.weight": "model-00009-of-00009.safetensors",
+    "model.language_model.layers.19.attention.query_layernorm.weight": "model-00009-of-00009.safetensors",
+    "model.language_model.layers.19.input_layernorm.weight": "model-00009-of-00009.safetensors",
+    "model.language_model.layers.19.mlp.experts.down_proj": "model-00009-of-00009.safetensors",
+    "model.language_model.layers.19.mlp.experts.gate_proj": "model-00009-of-00009.safetensors",
+    "model.language_model.layers.19.mlp.experts.up_proj": "model-00009-of-00009.safetensors",
+    "model.language_model.layers.19.mlp.gate.expert_bias": "model-00009-of-00009.safetensors",
+    "model.language_model.layers.19.mlp.gate.weight": "model-00009-of-00009.safetensors",
+    "model.language_model.layers.19.mlp.shared_experts.down_proj.weight": "model-00009-of-00009.safetensors",
+    "model.language_model.layers.19.mlp.shared_experts.gate_proj.weight": "model-00009-of-00009.safetensors",
+    "model.language_model.layers.19.mlp.shared_experts.up_proj.weight": "model-00009-of-00009.safetensors",
+    "model.language_model.layers.19.post_attention_layernorm.weight": "model-00009-of-00009.safetensors",
+    "model.language_model.layers.2.attention.dense.weight": "model-00001-of-00009.safetensors",
+    "model.language_model.layers.2.attention.key_layernorm.weight": "model-00001-of-00009.safetensors",
+    "model.language_model.layers.2.attention.query_key_value.weight": "model-00001-of-00009.safetensors",
+    "model.language_model.layers.2.attention.query_layernorm.weight": "model-00001-of-00009.safetensors",
+    "model.language_model.layers.2.input_layernorm.weight": "model-00001-of-00009.safetensors",
+    "model.language_model.layers.2.mlp.experts.down_proj": "model-00001-of-00009.safetensors",
+    "model.language_model.layers.2.mlp.experts.gate_proj": "model-00001-of-00009.safetensors",
+    "model.language_model.layers.2.mlp.experts.up_proj": "model-00001-of-00009.safetensors",
+    "model.language_model.layers.2.mlp.gate.expert_bias": "model-00009-of-00009.safetensors",
+    "model.language_model.layers.2.mlp.gate.weight": "model-00001-of-00009.safetensors",
+    "model.language_model.layers.2.mlp.shared_experts.down_proj.weight": "model-00001-of-00009.safetensors",
+    "model.language_model.layers.2.mlp.shared_experts.gate_proj.weight": "model-00001-of-00009.safetensors",
+    "model.language_model.layers.2.mlp.shared_experts.up_proj.weight": "model-00001-of-00009.safetensors",
+    "model.language_model.layers.2.post_attention_layernorm.weight": "model-00001-of-00009.safetensors",
+    "model.language_model.layers.3.attention.dense.weight": "model-00001-of-00009.safetensors",
+    "model.language_model.layers.3.attention.key_layernorm.weight": "model-00001-of-00009.safetensors",
+    "model.language_model.layers.3.attention.query_key_value.weight": "model-00001-of-00009.safetensors",
+    "model.language_model.layers.3.attention.query_layernorm.weight": "model-00001-of-00009.safetensors",
+    "model.language_model.layers.3.input_layernorm.weight": "model-00001-of-00009.safetensors",
+    "model.language_model.layers.3.mlp.experts.down_proj": "model-00001-of-00009.safetensors",
+    "model.language_model.layers.3.mlp.experts.gate_proj": "model-00002-of-00009.safetensors",
+    "model.language_model.layers.3.mlp.experts.up_proj": "model-00002-of-00009.safetensors",
+    "model.language_model.layers.3.mlp.gate.expert_bias": "model-00009-of-00009.safetensors",
+    "model.language_model.layers.3.mlp.gate.weight": "model-00002-of-00009.safetensors",
+    "model.language_model.layers.3.mlp.shared_experts.down_proj.weight": "model-00002-of-00009.safetensors",
+    "model.language_model.layers.3.mlp.shared_experts.gate_proj.weight": "model-00002-of-00009.safetensors",
+    "model.language_model.layers.3.mlp.shared_experts.up_proj.weight": "model-00002-of-00009.safetensors",
+    "model.language_model.layers.3.post_attention_layernorm.weight": "model-00002-of-00009.safetensors",
+    "model.language_model.layers.4.attention.dense.weight": "model-00002-of-00009.safetensors",
+    "model.language_model.layers.4.attention.key_layernorm.weight": "model-00002-of-00009.safetensors",
+    "model.language_model.layers.4.attention.query_key_value.weight": "model-00002-of-00009.safetensors",
+    "model.language_model.layers.4.attention.query_layernorm.weight": "model-00002-of-00009.safetensors",
+    "model.language_model.layers.4.input_layernorm.weight": "model-00002-of-00009.safetensors",
+    "model.language_model.layers.4.mlp.experts.down_proj": "model-00002-of-00009.safetensors",
+    "model.language_model.layers.4.mlp.experts.gate_proj": "model-00002-of-00009.safetensors",
+    "model.language_model.layers.4.mlp.experts.up_proj": "model-00002-of-00009.safetensors",
+    "model.language_model.layers.4.mlp.gate.expert_bias": "model-00009-of-00009.safetensors",
+    "model.language_model.layers.4.mlp.gate.weight": "model-00002-of-00009.safetensors",
+    "model.language_model.layers.4.mlp.shared_experts.down_proj.weight": "model-00002-of-00009.safetensors",
+    "model.language_model.layers.4.mlp.shared_experts.gate_proj.weight": "model-00002-of-00009.safetensors",
+    "model.language_model.layers.4.mlp.shared_experts.up_proj.weight": "model-00002-of-00009.safetensors",
+    "model.language_model.layers.4.post_attention_layernorm.weight": "model-00002-of-00009.safetensors",
+    "model.language_model.layers.5.attention.dense.weight": "model-00002-of-00009.safetensors",
+    "model.language_model.layers.5.attention.key_layernorm.weight": "model-00002-of-00009.safetensors",
+    "model.language_model.layers.5.attention.query_key_value.weight": "model-00002-of-00009.safetensors",
+    "model.language_model.layers.5.attention.query_layernorm.weight": "model-00002-of-00009.safetensors",
+    "model.language_model.layers.5.input_layernorm.weight": "model-00004-of-00009.safetensors",
+    "model.language_model.layers.5.mlp.experts.down_proj": "model-00004-of-00009.safetensors",
+    "model.language_model.layers.5.mlp.experts.gate_proj": "model-00002-of-00009.safetensors",
+    "model.language_model.layers.5.mlp.experts.up_proj": "model-00002-of-00009.safetensors",
+    "model.language_model.layers.5.mlp.gate.expert_bias": "model-00009-of-00009.safetensors",
+    "model.language_model.layers.5.mlp.gate.weight": "model-00004-of-00009.safetensors",
+    "model.language_model.layers.5.mlp.shared_experts.down_proj.weight": "model-00004-of-00009.safetensors",
+    "model.language_model.layers.5.mlp.shared_experts.gate_proj.weight": "model-00004-of-00009.safetensors",
+    "model.language_model.layers.5.mlp.shared_experts.up_proj.weight": "model-00004-of-00009.safetensors",
+    "model.language_model.layers.5.post_attention_layernorm.weight": "model-00004-of-00009.safetensors",
+    "model.language_model.layers.6.attention.dense.weight": "model-00004-of-00009.safetensors",
+    "model.language_model.layers.6.attention.key_layernorm.weight": "model-00004-of-00009.safetensors",
+    "model.language_model.layers.6.attention.query_key_value.weight": "model-00004-of-00009.safetensors",
+    "model.language_model.layers.6.attention.query_layernorm.weight": "model-00004-of-00009.safetensors",
+    "model.language_model.layers.6.input_layernorm.weight": "model-00004-of-00009.safetensors",
+    "model.language_model.layers.6.mlp.experts.down_proj": "model-00004-of-00009.safetensors",
+    "model.language_model.layers.6.mlp.experts.gate_proj": "model-00004-of-00009.safetensors",
+    "model.language_model.layers.6.mlp.experts.up_proj": "model-00004-of-00009.safetensors",
+    "model.language_model.layers.6.mlp.gate.expert_bias": "model-00009-of-00009.safetensors",
+    "model.language_model.layers.6.mlp.gate.weight": "model-00004-of-00009.safetensors",
+    "model.language_model.layers.6.mlp.shared_experts.down_proj.weight": "model-00004-of-00009.safetensors",
+    "model.language_model.layers.6.mlp.shared_experts.gate_proj.weight": "model-00004-of-00009.safetensors",
+    "model.language_model.layers.6.mlp.shared_experts.up_proj.weight": "model-00004-of-00009.safetensors",
+    "model.language_model.layers.6.post_attention_layernorm.weight": "model-00004-of-00009.safetensors",
+    "model.language_model.layers.7.attention.dense.weight": "model-00004-of-00009.safetensors",
+    "model.language_model.layers.7.attention.key_layernorm.weight": "model-00004-of-00009.safetensors",
+    "model.language_model.layers.7.attention.query_key_value.weight": "model-00004-of-00009.safetensors",
+    "model.language_model.layers.7.attention.query_layernorm.weight": "model-00004-of-00009.safetensors",
+    "model.language_model.layers.7.input_layernorm.weight": "model-00004-of-00009.safetensors",
+    "model.language_model.layers.7.mlp.experts.down_proj": "model-00004-of-00009.safetensors",
+    "model.language_model.layers.7.mlp.experts.gate_proj": "model-00004-of-00009.safetensors",
+    "model.language_model.layers.7.mlp.experts.up_proj": "model-00005-of-00009.safetensors",
+    "model.language_model.layers.7.mlp.gate.expert_bias": "model-00009-of-00009.safetensors",
+    "model.language_model.layers.7.mlp.gate.weight": "model-00005-of-00009.safetensors",
+    "model.language_model.layers.7.mlp.shared_experts.down_proj.weight": "model-00005-of-00009.safetensors",
+    "model.language_model.layers.7.mlp.shared_experts.gate_proj.weight": "model-00005-of-00009.safetensors",
+    "model.language_model.layers.7.mlp.shared_experts.up_proj.weight": "model-00005-of-00009.safetensors",
+    "model.language_model.layers.7.post_attention_layernorm.weight": "model-00005-of-00009.safetensors",
+    "model.language_model.layers.8.attention.dense.weight": "model-00005-of-00009.safetensors",
+    "model.language_model.layers.8.attention.key_layernorm.weight": "model-00005-of-00009.safetensors",
+    "model.language_model.layers.8.attention.query_key_value.weight": "model-00005-of-00009.safetensors",
+    "model.language_model.layers.8.attention.query_layernorm.weight": "model-00005-of-00009.safetensors",
+    "model.language_model.layers.8.input_layernorm.weight": "model-00005-of-00009.safetensors",
+    "model.language_model.layers.8.mlp.experts.down_proj": "model-00005-of-00009.safetensors",
+    "model.language_model.layers.8.mlp.experts.gate_proj": "model-00005-of-00009.safetensors",
+    "model.language_model.layers.8.mlp.experts.up_proj": "model-00005-of-00009.safetensors",
+    "model.language_model.layers.8.mlp.gate.expert_bias": "model-00009-of-00009.safetensors",
+    "model.language_model.layers.8.mlp.gate.weight": "model-00005-of-00009.safetensors",
+    "model.language_model.layers.8.mlp.shared_experts.down_proj.weight": "model-00005-of-00009.safetensors",
+    "model.language_model.layers.8.mlp.shared_experts.gate_proj.weight": "model-00005-of-00009.safetensors",
+    "model.language_model.layers.8.mlp.shared_experts.up_proj.weight": "model-00005-of-00009.safetensors",
+    "model.language_model.layers.8.post_attention_layernorm.weight": "model-00005-of-00009.safetensors",
+    "model.language_model.layers.9.attention.dense.weight": "model-00005-of-00009.safetensors",
+    "model.language_model.layers.9.attention.key_layernorm.weight": "model-00005-of-00009.safetensors",
+    "model.language_model.layers.9.attention.query_key_value.weight": "model-00005-of-00009.safetensors",
+    "model.language_model.layers.9.attention.query_layernorm.weight": "model-00005-of-00009.safetensors",
+    "model.language_model.layers.9.input_layernorm.weight": "model-00005-of-00009.safetensors",
+    "model.language_model.layers.9.mlp.experts.down_proj": "model-00005-of-00009.safetensors",
+    "model.language_model.layers.9.mlp.experts.gate_proj": "model-00005-of-00009.safetensors",
+    "model.language_model.layers.9.mlp.experts.up_proj": "model-00005-of-00009.safetensors",
+    "model.language_model.layers.9.mlp.gate.expert_bias": "model-00009-of-00009.safetensors",
+    "model.language_model.layers.9.mlp.gate.weight": "model-00005-of-00009.safetensors",
+    "model.language_model.layers.9.mlp.shared_experts.down_proj.weight": "model-00005-of-00009.safetensors",
+    "model.language_model.layers.9.mlp.shared_experts.gate_proj.weight": "model-00005-of-00009.safetensors",
+    "model.language_model.layers.9.mlp.shared_experts.up_proj.weight": "model-00005-of-00009.safetensors",
+    "model.language_model.layers.9.post_attention_layernorm.weight": "model-00005-of-00009.safetensors",
+    "model.language_model.norm.weight": "model-00009-of-00009.safetensors",
+    "model.language_model.word_embeddings.weight": "model-00003-of-00009.safetensors",
+    "model.lm_head.weight": "model-00009-of-00009.safetensors"
+  }
+}

text_encoder/modeling_llada2uni_moe.py

@@ -0,0 +1,1289 @@
+# coding=utf-8
+# Copyright 2025 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 implementation of the fused LLaDA2 MoE model."""
+
+from dataclasses import dataclass
+import math
+from typing import List, Optional, Tuple, Union
+import warnings
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from transformers.activations import ACT2FN
+from transformers.cache_utils import Cache, DynamicCache
+from transformers.generation import GenerationMixin
+from transformers.modeling_attn_mask_utils import (
+    _prepare_4d_attention_mask,
+    _prepare_4d_causal_attention_mask,
+    _prepare_4d_causal_attention_mask_for_sdpa,
+)
+from transformers.modeling_outputs import ModelOutput, MoeModelOutputWithPast
+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
+from transformers.utils import logging
+from veomni.ops import fused_moe_forward
+
+from .configuration_llada2uni_moe import LLaDA2MoeConfig
+
+logger = logging.get_logger(__name__)
+
+
+class LLaDA2MoeRMSNorm(nn.Module):
+    """RMSNorm used by the LLaDA2 model."""
+
+    def __init__(self, hidden_size, eps=1e-6):
+        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)
+
+
+# Preserve the historical spelling used by the original implementation.
+LLaDA2MoERMSNorm = LLaDA2MoeRMSNorm
+ALL_LAYERNORM_LAYERS.append(LLaDA2MoeRMSNorm)
+
+
+class LLaDA2MoePreTrainedModel(PreTrainedModel):
+    config_class = LLaDA2MoeConfig
+    base_model_prefix = "model"
+    supports_gradient_checkpointing = True
+    _no_split_modules = ["LLaDA2MoeDecoderLayer"]
+    _skip_keys_device_placement = "past_key_values"
+    _supports_flash_attn_2 = True
+    _supports_sdpa = True
+    _supports_cache_class = True
+    _supports_flash_attn = True
+    _can_compile_fullgraph = True
+    _supports_attention_backend = 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_()
+
+
+def rotate_half(hidden_states):
+    first, second = hidden_states.chunk(2, dim=-1)
+    return torch.cat((-second, first), dim=-1)
+
+
+def apply_rotary_pos_emb(query, key, cos, sin, position_ids=None, unsqueeze_dim=1):
+    """Apply RoPE to the rotary part of query and key states."""
+    cos = cos.unsqueeze(unsqueeze_dim)
+    sin = sin.unsqueeze(unsqueeze_dim)
+    rotary_dim = cos.shape[-1]
+    query_rotary, query_pass = query[..., :rotary_dim], query[..., rotary_dim:]
+    key_rotary, key_pass = key[..., :rotary_dim], key[..., rotary_dim:]
+    query_rotary = query_rotary * cos + rotate_half(query_rotary) * sin
+    key_rotary = key_rotary * cos + rotate_half(key_rotary) * sin
+    return torch.cat((query_rotary, query_pass), dim=-1), torch.cat(
+        (key_rotary, key_pass), dim=-1
+    )
+
+
+class LLaDA2MoeRotaryEmbedding(nn.Module):
+    def __init__(self, config: LLaDA2MoeConfig, 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):
+        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)
+
+
+class LLaDA2MoeMLP(nn.Module):
+    def __init__(self, config: LLaDA2MoeConfig, 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):
+        return self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x))
+
+
+class LLaDA2MoeGate(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
+
+        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.register_buffer("expert_bias", torch.zeros((self.num_experts)))
+        self.reset_parameters()
+
+    def reset_parameters(self) -> None:
+        import torch.nn.init as init
+
+        init.kaiming_uniform_(self.weight, a=math.sqrt(5))
+
+    def group_limited_topk(
+        self,
+        scores: torch.Tensor,
+    ):
+        num_tokens, _ = scores.size()
+        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)
+
+        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)
+
+        return probs, top_indices
+
+    def forward(self, hidden_states):
+        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.float()).type_as(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 LLaDA2MoeExperts(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        self.num_experts = config.num_experts
+        self.hidden_dim = config.hidden_size
+        self.intermediate_size = config.moe_intermediate_size
+        self.gate_proj = torch.nn.Parameter(
+            torch.empty(self.num_experts, self.intermediate_size, self.hidden_dim),
+            requires_grad=True,
+        )
+        self.up_proj = torch.nn.Parameter(
+            torch.empty(self.num_experts, self.intermediate_size, self.hidden_dim),
+            requires_grad=True,
+        )
+        self.down_proj = torch.nn.Parameter(
+            torch.empty(self.num_experts, self.hidden_dim, self.intermediate_size),
+            requires_grad=True,
+        )
+
+    def forward(self, hidden_states, routing_weights, selected_experts):
+        return fused_moe_forward(
+            module=self,
+            num_experts=self.num_experts,
+            routing_weights=routing_weights,
+            selected_experts=selected_experts,
+            hidden_states=hidden_states,
+            fc1_1_weight=self.gate_proj,
+            fc1_2_weight=self.up_proj,
+            fc2_weight=self.down_proj,
+        )
+
+    def reset_parameters(self):
+        """
+        Initialize the parameters of all expert networks.
+        Uses different initialization strategies for different projection layers.
+        """
+        for expert_id in range(self.num_experts):
+            nn.init.kaiming_uniform_(self.gate_proj[expert_id], a=math.sqrt(5))
+            nn.init.kaiming_uniform_(self.up_proj[expert_id], a=math.sqrt(5))
+            nn.init.xavier_uniform_(self.down_proj[expert_id])
+
+
+class LLaDA2MoeSparseMoeBlock(nn.Module):
+    """Fused routed experts plus a shared expert."""
+
+    def __init__(self, config: LLaDA2MoeConfig):
+        super().__init__()
+        self.config = config
+        self.experts = LLaDA2MoeExperts(config)
+        self.gate = LLaDA2MoeGate(config)
+        if config.num_shared_experts is not None:
+            self.shared_experts = LLaDA2MoeMLP(
+                config=config,
+                intermediate_size=config.moe_intermediate_size
+                * config.num_shared_experts,
+            )
+
+    def forward(self, hidden_states):
+        identity = hidden_states
+        bsz, seq_len, h = hidden_states.shape
+        topk_idx, topk_weight, router_logits = self.gate(hidden_states)
+        hidden_states = hidden_states.view(-1, hidden_states.shape[-1])
+        y = self.experts(
+            hidden_states, routing_weights=topk_weight, selected_experts=topk_idx
+        ).reshape(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),
+        )
+
+
+def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor:
+    """
+    This is the equivalent of torch.repeat_interleave(x, dim=1, repeats=n_rep). The hidden states go from (batch,
+    num_key_value_heads, seqlen, head_dim) to (batch, num_attention_heads, seqlen, head_dim)
+    """
+    batch, num_key_value_heads, slen, head_dim = hidden_states.shape
+    if n_rep == 1:
+        return hidden_states
+    hidden_states = hidden_states[:, :, None, :, :].expand(
+        batch, num_key_value_heads, n_rep, slen, head_dim
+    )
+    return hidden_states.reshape(batch, num_key_value_heads * n_rep, slen, head_dim)
+
+
+class LLaDA2MoeAttention(nn.Module):
+    """Multi-headed attention from 'Attention Is All You Need' paper"""
+
+    def __init__(self, config: LLaDA2MoeConfig, 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 = False
+
+        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.query_layernorm = LLaDA2MoERMSNorm(self.head_dim, eps=config.rms_norm_eps)
+        self.key_layernorm = LLaDA2MoERMSNorm(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
+        )
+
+    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)
+
+        query_states = self.query_layernorm(query_states)
+        key_states = self.key_layernorm(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_rotary_pos_emb(
+            query_states, key_states, cos, sin, position_ids
+        )
+
+        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, key_states.transpose(2, 3)
+        ) / math.sqrt(self.head_dim)
+
+        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
+
+        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
+
+        return attn_output, attn_weights, past_key_value
+
+
+class LLaDA2MoeSdpaAttention(LLaDA2MoeAttention):
+    """
+    LLaDA2Moe attention module using torch.nn.functional.scaled_dot_product_attention. This module inherits from
+    `LLaDA2MoeAttention` as the weights of the module stays untouched. The only changes are on the forward pass to adapt to
+    SDPA API.
+    """
+
+    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 output_attentions:
+            logger.warning_once(
+                "LLaDA2MoeModel is using LLaDA2MoeSdpaAttention, but `torch.nn.functional.scaled_dot_product_attention` does not support `output_attentions=True`. Falling back to the manual attention implementation, "
+                'but specifying the manual implementation will be required from Transformers version v5.0.0 onwards. This warning can be removed using the argument `attn_implementation="eager"` when loading the model.'
+            )
+            return super().forward(
+                hidden_states=hidden_states,
+                attention_mask=attention_mask,
+                position_ids=position_ids,
+                past_key_value=past_key_value,
+                output_attentions=output_attentions,
+                use_cache=use_cache,
+            )
+
+        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)
+
+        query_states = self.query_layernorm(query_states)
+        key_states = self.key_layernorm(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
+
+        query_states, key_states = apply_rotary_pos_emb(
+            query_states, key_states, cos, sin, position_ids
+        )
+
+        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)
+
+        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()}"
+                )
+
+        # SDPA with memory-efficient backend is currently (torch==2.1.2) bugged with non-contiguous inputs with custom attn_mask,
+        # Reference: https://github.com/pytorch/pytorch/issues/112577.
+        if query_states.device.type == "cuda" and attention_mask is not None:
+            query_states = query_states.contiguous()
+            key_states = key_states.contiguous()
+            value_states = value_states.contiguous()
+
+        attn_output = torch.nn.functional.scaled_dot_product_attention(
+            query_states,
+            key_states,
+            value_states,
+            attn_mask=attention_mask,
+            dropout_p=self.attention_dropout if self.training else 0.0,
+            # The q_len > 1 is necessary to match with AttentionMaskConverter.to_causal_4d that does not create a causal mask in case q_len == 1.
+            is_causal=self.is_causal and attention_mask is None and q_len > 1,
+        )
+
+        attn_output = attn_output.transpose(1, 2).contiguous()
+        attn_output = attn_output.reshape(bsz, q_len, -1)
+
+        attn_output = self.dense(attn_output)
+
+        return attn_output, None, past_key_value
+
+
+ATTENTION_CLASSES = {
+    "eager": LLaDA2MoeSdpaAttention,
+    "flash_attention_2": LLaDA2MoeSdpaAttention,
+    "sdpa": LLaDA2MoeSdpaAttention,
+}
+
+
+class LLaDA2MoeDecoderLayer(nn.Module):
+    def __init__(self, config: LLaDA2MoeConfig, 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 = (
+            LLaDA2MoeSparseMoeBlock(config)
+            if (
+                config.num_experts is not None
+                and layer_idx >= config.first_k_dense_replace
+            )
+            else LLaDA2MoeMLP(config=config, intermediate_size=config.intermediate_size)
+        )
+        self.input_layernorm = LLaDA2MoERMSNorm(
+            config.hidden_size, eps=config.rms_norm_eps
+        )
+        self.post_attention_layernorm = LLaDA2MoERMSNorm(
+            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,
+        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)
+
+        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
+
+        residual = hidden_states
+        hidden_states = self.post_attention_layernorm(hidden_states)
+        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
+
+
+def calculate_pack_position_ids(
+    input_ids: Optional[torch.Tensor] = None,
+    inputs_embeds: Optional[torch.Tensor] = None,
+    position_ids: Optional[torch.Tensor] = None,
+    past_key_values_length: int = 0,
+    cu_lengths_list: Optional[List[torch.Tensor]] = None,
+):
+    """Build continuous or per-sequence packed position IDs."""
+    if position_ids is not None:
+        return position_ids
+
+    if input_ids is not None:
+        device = input_ids.device
+        batch_size, seq_length = input_ids.shape
+    elif inputs_embeds is not None:
+        device = inputs_embeds.device
+        batch_size, seq_length, _ = inputs_embeds.shape
+    else:
+        raise ValueError("You have to specify either input_ids or inputs_embeds")
+
+    if cu_lengths_list is not None:
+        all_position_ids = []
+
+        for i in range(batch_size):
+            cu_seqlens = cu_lengths_list[i].to(device)
+            starts = cu_seqlens[:-1]
+            lengths = cu_seqlens[1:] - cu_seqlens[:-1]
+            total_len = cu_seqlens[-1].item()
+            global_positions = torch.arange(total_len, device=device, dtype=torch.long)
+            subtraction_mask = torch.repeat_interleave(starts, lengths)
+            current_pos_ids = global_positions - subtraction_mask
+            all_position_ids.append(current_pos_ids)
+
+        position_ids = torch.nn.utils.rnn.pad_sequence(
+            all_position_ids,
+            batch_first=True,
+            padding_value=0,
+        )
+
+        if position_ids.shape[1] < seq_length:
+            pad_right = seq_length - position_ids.shape[1]
+            position_ids = F.pad(position_ids, (0, pad_right), "constant", 0)
+    else:
+        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).expand(batch_size, -1)
+
+    return position_ids
+
+
+class LLaDA2MoeModel(LLaDA2MoePreTrainedModel):
+    """
+    Transformer decoder consisting of *config.num_hidden_layers* layers. Each layer is a [`LLaDA2MoeDecoderLayer`]
+    Args:
+        config: LLaDA2MoeConfig
+    """
+
+    def __init__(self, config: LLaDA2MoeConfig):
+        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(
+            [
+                LLaDA2MoeDecoderLayer(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 = LLaDA2MoERMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+        self.rotary_emb = LLaDA2MoeRotaryEmbedding(config=config)
+        self.gradient_checkpointing = False
+
+        self.post_init()
+
+    def get_input_embeddings(self):
+        return self.word_embeddings
+
+    def set_input_embeddings(self, value):
+        self.word_embeddings = value
+
+    def forward(
+        self,
+        input_ids: torch.LongTensor = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        past_key_values: Optional[List[torch.FloatTensor]] = None,
+        inputs_embeds: Optional[torch.FloatTensor] = None,
+        use_cache: Optional[bool] = None,
+        output_attentions: Optional[bool] = None,
+        output_hidden_states: Optional[bool] = None,
+        output_router_logits: Optional[bool] = None,
+        cu_lengths_list: Optional[List] = None,
+        return_dict: Optional[bool] = None,
+        **kwargs,
+    ) -> Union[Tuple, MoeModelOutputWithPast]:
+        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
+        )
+
+        if input_ids is not None and inputs_embeds is not None:
+            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:
+            position_ids = calculate_pack_position_ids(
+                input_ids,
+                inputs_embeds,
+                position_ids,
+                past_key_values_length,
+                cu_lengths_list,
+            )
+
+        if inputs_embeds is None:
+            inputs_embeds = self.word_embeddings(input_ids)
+
+        if hasattr(attention_mask, "dim") and attention_mask.dim() == 2:
+            if 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:
+                if attention_mask is not None:
+                    attention_mask = _prepare_4d_attention_mask(
+                        attention_mask, inputs_embeds.dtype
+                    )
+                else:
+                    attention_mask = _prepare_4d_causal_attention_mask(
+                        attention_mask,
+                        (batch_size, seq_length),
+                        inputs_embeds,
+                        past_key_values_length,
+                    )
+        hidden_states = inputs_embeds
+
+        position_embeddings = self.rotary_emb(hidden_states, position_ids)
+
+        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,
+                    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,
+                    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)
+
+        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,
+        )
+
+
+@dataclass
+class LLaDA2MoeCausalLMOutputWithPast(ModelOutput):
+    r"""
+    loss (`torch.FloatTensor` of shape `(1,)`, *optional*, returned when `labels` is provided):
+        Training loss.
+    logits (`torch.FloatTensor` of shape `(batch_size, sequence_length, config.vocab_size)`):
+        Prediction scores for each vocabulary token before SoftMax.
+    past_key_values (`Cache`, *optional*, returned when `use_cache=True` is passed or when `config.use_cache=True`):
+        It is a [`~cache_utils.Cache`] instance. For more details, see our [kv cache guide](https://huggingface.co/docs/transformers/en/kv_cache).
+
+        Contains pre-computed hidden-states (key and values in the self-attention blocks) that can be used (see
+        `past_key_values` input) to speed up sequential decoding.
+    rope_deltas (`torch.LongTensor` of shape `(batch_size, )`, *optional*):
+        The offset between the sequence length and rotary position indices.
+    """
+
+    loss: Optional[torch.FloatTensor] = None
+    z_loss: Optional[torch.FloatTensor] = None
+    logits: Optional[torch.FloatTensor] = None
+    past_key_values: Optional[Cache] = None
+    hidden_states: Optional[tuple[torch.FloatTensor]] = None
+    attentions: Optional[tuple[torch.FloatTensor]] = None
+    rope_deltas: Optional[torch.LongTensor] = None
+
+
+class LLaDA2MoeBackbone(nn.Module):
+    """Container for the model backbone and output head."""
+
+    def __init__(self, config: LLaDA2MoeConfig):
+        super().__init__()
+        self.language_model = LLaDA2MoeModel(config)
+        self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
+
+    def get_input_embeddings(self):
+        return self.language_model.get_input_embeddings()
+
+    def set_input_embeddings(self, value):
+        self.language_model.set_input_embeddings(value)
+
+    def forward(self, *args, **kwargs):
+        return self.language_model(*args, **kwargs)
+
+
+class LLaDA2MoeModelLM(LLaDA2MoePreTrainedModel, GenerationMixin):
+    """Fused LLaDA2 MoE model."""
+
+    accepts_loss_kwargs = False
+
+    def __init__(self, config: LLaDA2MoeConfig):
+        super().__init__(config)
+        self.model = LLaDA2MoeBackbone(config)
+        self.img_token_id = 157184
+        self.img_start_id = 157185
+        self.img_end_id = 157186
+        self.img_pad_id = 157187
+        self.post_init()
+
+    @property
+    def language_model(self):
+        return self.model.language_model
+
+    def get_input_embeddings(self):
+        return self.model.get_input_embeddings()
+
+    def set_input_embeddings(self, value):
+        self.model.set_input_embeddings(value)
+
+    def get_output_embeddings(self):
+        return self.model.lm_head
+
+    def set_output_embeddings(self, value):
+        self.model.lm_head = value
+
+    def get_decoder(self):
+        return self.model.language_model
+
+    def set_decoder(self, decoder):
+        self.model.language_model = decoder
+
+    def forward(
+        self,
+        input_ids: Optional[torch.LongTensor] = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        past_key_values: Optional[Cache] = None,
+        inputs_embeds: Optional[torch.FloatTensor] = None,
+        labels: Optional[torch.LongTensor] = None,
+        use_cache: Optional[bool] = None,
+        output_attentions: Optional[bool] = None,
+        output_router_logits: Optional[bool] = None,
+        output_hidden_states: Optional[bool] = None,
+        return_dict: Optional[bool] = None,
+        logits_to_keep: Union[int, torch.Tensor] = 0,
+        cu_lengths_list: Optional[List] = None,
+        **kwargs,
+    ) -> Union[tuple, LLaDA2MoeCausalLMOutputWithPast]:
+        return_dict = (
+            return_dict if return_dict is not None else self.config.use_return_dict
+        )
+        if inputs_embeds is None:
+            if input_ids is None:
+                raise ValueError("Provide either input_ids or inputs_embeds")
+            inputs_embeds = self.get_input_embeddings()(input_ids)
+
+        outputs = self.model(
+            input_ids=None,
+            attention_mask=attention_mask,
+            position_ids=position_ids,
+            past_key_values=past_key_values,
+            inputs_embeds=inputs_embeds,
+            use_cache=use_cache,
+            output_attentions=output_attentions,
+            output_router_logits=output_router_logits,
+            output_hidden_states=output_hidden_states,
+            return_dict=True,
+            cu_lengths_list=cu_lengths_list,
+            **kwargs,
+        )
+        hidden_states = outputs.last_hidden_state
+        indices = (
+            slice(-logits_to_keep, None)
+            if isinstance(logits_to_keep, int)
+            else logits_to_keep
+        )
+        logits = self.model.lm_head(hidden_states[:, indices, :])
+
+        loss = None
+        if labels is not None:
+            loss = F.cross_entropy(
+                logits.reshape(-1, logits.shape[-1]), labels.reshape(-1)
+            )
+
+        if not return_dict:
+            result = (
+                logits,
+                outputs.past_key_values,
+                outputs.hidden_states,
+                outputs.attentions,
+            )
+            return ((loss,) + result) if loss is not None else result
+        return LLaDA2MoeCausalLMOutputWithPast(
+            loss=loss,
+            z_loss=None,
+            logits=logits,
+            past_key_values=outputs.past_key_values,
+            hidden_states=outputs.hidden_states,
+            attentions=outputs.attentions,
+            rope_deltas=None,
+        )
+
+    @staticmethod
+    def _top_k_logits(logits, k):
+        if k is None or k <= 0:
+            return logits
+        values, _ = torch.topk(logits, min(k, logits.shape[-1]))
+        return torch.where(logits < values[..., -1, None], -torch.inf, logits)
+
+    @staticmethod
+    def _top_p_logits(logits, p):
+        if p is None or p >= 1.0:
+            return logits
+        sorted_logits, sorted_indices = torch.sort(logits, descending=True)
+        cumulative_probs = torch.cumsum(F.softmax(sorted_logits, dim=-1), dim=-1)
+        sorted_mask = cumulative_probs > p
+        sorted_mask[..., 1:] = sorted_mask[..., :-1].clone()
+        sorted_mask[..., 0] = False
+        mask = torch.zeros_like(sorted_mask).scatter(-1, sorted_indices, sorted_mask)
+        return logits.masked_fill(mask, -torch.inf)
+
+    def _sample_with_temperature_topk_topp(
+        self, logits, temperature=1.0, top_k=0, top_p=1.0
+    ):
+        original_shape = logits.shape[:-1]
+        logits = logits.reshape(-1, logits.shape[-1])
+        if (
+            temperature == 0.0
+            and (top_k in (None, 0))
+            and (top_p is None or top_p >= 1.0)
+        ):
+            probs = F.softmax(logits, dim=-1)
+            token = logits.argmax(dim=-1, keepdim=True)
+            token_prob = probs.gather(-1, token)
+            return token.view(*original_shape), token_prob.view(*original_shape)
+        if temperature > 0 and temperature != 1.0:
+            logits = logits / temperature
+        logits = self._top_k_logits(logits, top_k)
+        logits = self._top_p_logits(logits, top_p)
+        probs = F.softmax(logits, dim=-1)
+        token = torch.multinomial(probs, num_samples=1)
+        token_prob = probs.gather(-1, token)
+        return token.view(*original_shape), token_prob.view(*original_shape)
+
+    @staticmethod
+    def _get_num_transfer_tokens(block_length, steps):
+        if steps == 0:
+            return torch.empty(0, dtype=torch.int64)
+        schedule = torch.full((steps,), block_length // steps, dtype=torch.int64)
+        schedule[: block_length % steps] += 1
+        return schedule
+
+    @torch.no_grad()
+    def generate_bd_image_logic(
+        self,
+        data: Optional[dict] = None,
+        temperature: float = 0.0,
+        block_length: int = 32,
+        steps: int = 32,
+        gen_length: int = 2048,
+        top_p: Optional[float] = None,
+        top_k: Optional[int] = None,
+        eos_early_stop: bool = True,
+        minimal_topk: int = 1,
+        threshold: float = 0.95,
+        eos_id: int = 156892,
+        mask_id: int = 156895,
+        cfg_scale: float = 1.0,
+        mode: str = "eoi",
+    ):
+        """Generate discrete image tokens with the original block-diffusion logic."""
+        if data is None or "input_ids" not in data:
+            raise ValueError("data must contain input_ids")
+        steps = min(steps, gen_length // minimal_topk)
+        input_ids = data["input_ids"]
+        eoi_id = 156902
+        prompt_length = input_ids.shape[1]
+        num_blocks = (prompt_length + gen_length + block_length - 1) // block_length
+        total_length = num_blocks * block_length
+
+        block_mask = torch.tril(torch.ones(num_blocks, num_blocks, device=self.device))
+        full_attention_mask = (
+            block_mask.repeat_interleave(block_length, 0)
+            .repeat_interleave(block_length, 1)[None, None]
+            .bool()
+        )
+        position_ids = torch.arange(total_length, device=self.device).unsqueeze(0)
+        x = torch.full((1, total_length), mask_id, dtype=torch.long, device=self.device)
+        x[:, :prompt_length] = input_ids
+        prefill_blocks = prompt_length // block_length
+        schedule = self._get_num_transfer_tokens(block_length, steps)
+        use_cfg = cfg_scale != 1.0
+
+        if use_cfg:
+            uncond_ids = data.get("uncond_ids", [27, 411, 19483, 29])
+            if torch.is_tensor(uncond_ids):
+                uncond_ids = uncond_ids.flatten().tolist()
+            pad_len = prompt_length - len(uncond_ids)
+            if pad_len < 0:
+                raise ValueError(
+                    "The unconditional prompt is longer than the conditional prompt"
+                )
+            uncond_input = torch.full(
+                (1, prompt_length), mask_id, dtype=torch.long, device=self.device
+            )
+            uncond_input[0, -len(uncond_ids) :] = torch.tensor(
+                uncond_ids, device=self.device
+            )
+            uncond_attention_mask = full_attention_mask.clone()
+            uncond_attention_mask[:, :, :, :pad_len] = False
+            uncond_position_ids = torch.cat(
+                [
+                    torch.zeros(pad_len, device=self.device, dtype=torch.long),
+                    torch.arange(total_length - pad_len, device=self.device),
+                ]
+            ).unsqueeze(0)
+
+        for block_index in range(prefill_blocks, num_blocks):
+            window_end = (block_index + 1) * block_length
+            current = x[:, :window_end]
+            current_mask = full_attention_mask[:, :, :window_end, :window_end]
+            current_positions = position_ids[:, :window_end]
+
+            for step_index in range(steps):
+                active = current[:, -block_length:] == mask_id
+                if not active.any():
+                    break
+                if use_cfg:
+                    unconditional = current.clone()
+                    unconditional[:, :prompt_length] = uncond_input
+                    combined_ids = torch.cat([current, unconditional], dim=0)
+                    combined_positions = torch.cat(
+                        [current_positions, uncond_position_ids[:, :window_end]], dim=0
+                    )
+                    combined_mask = torch.cat(
+                        [
+                            current_mask,
+                            uncond_attention_mask[:, :, :window_end, :window_end],
+                        ],
+                        dim=0,
+                    )
+                    logits = self(
+                        input_ids=combined_ids,
+                        attention_mask=combined_mask,
+                        position_ids=combined_positions,
+                    ).logits
+                    conditional_logits, unconditional_logits = logits.chunk(2, dim=0)
+                    active_logits = unconditional_logits[
+                        :, -block_length:
+                    ] + cfg_scale * (
+                        conditional_logits[:, -block_length:]
+                        - unconditional_logits[:, -block_length:]
+                    )
+                else:
+                    active_logits = self(
+                        input_ids=current,
+                        attention_mask=current_mask,
+                        position_ids=current_positions,
+                    ).logits[:, -block_length:]
+
+                tokens, confidence = self._sample_with_temperature_topk_topp(
+                    active_logits, temperature=temperature, top_k=top_k, top_p=top_p
+                )
+                count = schedule[step_index].item()
+                scores = torch.where(active, confidence, -torch.inf)
+                selected = torch.zeros_like(tokens, dtype=torch.bool)
+                high_confidence = scores[0] > threshold
+                if high_confidence.sum().item() >= count:
+                    selected[0] = high_confidence
+                else:
+                    _, indices = torch.topk(
+                        scores[0], k=min(count, active.sum().item())
+                    )
+                    selected[0, indices] = True
+                current[:, -block_length:][selected] = tokens[selected]
+
+                stop_token = eoi_id if mode == "eoi" else eos_id
+                positions = (current[0, prompt_length:] == stop_token).nonzero(
+                    as_tuple=True
+                )[0]
+                if eos_early_stop and len(positions) > 0:
+                    stop_position = positions[0].item() + prompt_length
+                    if (current[0, prompt_length:stop_position] != mask_id).all():
+                        x[:, :window_end] = current
+                        return x[:, : stop_position + 1]
+
+            x[:, :window_end] = current
+        return x[:, : prompt_length + gen_length]
+
+
+__all__ = ["LLaDA2MoeModelLM", "LLaDA2MoeModel", "LLaDA2MoePreTrainedModel"]

text_projection/config.json

@@ -0,0 +1,11 @@
+{
+  "_class_name": "LLaDAImageTextProjectionModel",
+  "_diffusers_version": "0.40.0.dev0",
+  "attention_dropout": 0.0,
+  "hidden_size": 2048,
+  "intermediate_size": 8960,
+  "norm_eps": 1e-06,
+  "num_attention_heads": 32,
+  "num_hidden_layers": 6,
+  "projection_dim": 2560
+}

text_projection/diffusion_pytorch_model.safetensors

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

tokenizer/special_tokens_map.json

@@ -0,0 +1,37 @@
+{
+  "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": "<|endoftext|>",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  },
+  "mask_token": {
+    "content": "<|mask|>",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  },
+  "pad_token": {
+    "content": "<|endoftext|>",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  }
+}

tokenizer/tokenizer.json

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

transformer/config.json

@@ -0,0 +1,31 @@
+{
+  "_class_name": "LLaDAImageTransformer2DModel",
+  "_diffusers_version": "0.40.0.dev0",
+  "all_f_patch_size": [
+    1
+  ],
+  "all_patch_size": [
+    1
+  ],
+  "axes_dims": [
+    32,
+    48,
+    48
+  ],
+  "axes_lens": [
+    32768,
+    1024,
+    1024
+  ],
+  "cap_feat_dim": 2560,
+  "dim": 3840,
+  "in_channels": 128,
+  "n_heads": 30,
+  "n_layers": 30,
+  "n_refiner_layers": 2,
+  "norm_eps": 1e-05,
+  "qk_norm": true,
+  "rope_theta": 256.0,
+  "semantic_feat_dim": 4096,
+  "t_scale": 1000.0
+}

transformer/diffusion_pytorch_model-00001-of-00004.safetensors

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

transformer/diffusion_pytorch_model-00002-of-00004.safetensors

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

transformer/diffusion_pytorch_model-00003-of-00004.safetensors

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

transformer/diffusion_pytorch_model-00004-of-00004.safetensors

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

transformer/diffusion_pytorch_model.safetensors.index.json

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

vae/config.json

@@ -0,0 +1,39 @@
+{
+  "_class_name": "AutoencoderKLFlux2",
+  "_diffusers_version": "0.36.0.dev0",
+  "act_fn": "silu",
+  "batch_norm_eps": 0.0001,
+  "batch_norm_momentum": 0.1,
+  "block_out_channels": [
+    128,
+    256,
+    512,
+    512
+  ],
+  "down_block_types": [
+    "DownEncoderBlock2D",
+    "DownEncoderBlock2D",
+    "DownEncoderBlock2D",
+    "DownEncoderBlock2D"
+  ],
+  "force_upcast": true,
+  "in_channels": 3,
+  "latent_channels": 32,
+  "layers_per_block": 2,
+  "mid_block_add_attention": true,
+  "norm_num_groups": 32,
+  "out_channels": 3,
+  "patch_size": [
+    2,
+    2
+  ],
+  "sample_size": 1024,
+  "up_block_types": [
+    "UpDecoderBlock2D",
+    "UpDecoderBlock2D",
+    "UpDecoderBlock2D",
+    "UpDecoderBlock2D"
+  ],
+  "use_post_quant_conv": true,
+  "use_quant_conv": true
+}

vae/diffusion_pytorch_model.safetensors

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