Upload folder using huggingface_hub

.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
+assets/comparision.png filter=lfs diff=lfs merge=lfs -text
+tokenizer.json filter=lfs diff=lfs merge=lfs -text

README.md

@@ -1,3 +1,92 @@
----
-license: apache-2.0
----
+---
+license: apache-2.0
+pipeline_tag: text-to-image
+library_name: transformers
+---
+
+## NextStep-1.1
+
+[Homepage](https://stepfun.ai/research/en/nextstep-1) 
+| [GitHub](https://github.com/stepfun-ai/NextStep-1) 
+| [Paper](https://arxiv.org/abs/2508.10711) 
+
+We introduce **NextStep-1.1**, a new model represents a significant leap forward in the NextStep series. This version effectively resolves the visualization failures seen in **NextStep-1** and substantially elevates image quality through extended training and a Flow-based Reinforcement Learning (RL) post-training paradigm.
+
+<div align='center'>
+<img src="assets/comparision.png" class="interpolation-image" alt="arch." width="100%" />
+</div>
+
+## What's New in 1.1?
+
+NextStep-1.1 is not just a fine-tune; it is a re-engineered version focused on stability and high-fidelity output. Key improvements include:
+
+- RL Enhanced Visual Fidelity: Significant improvement in image texture and a substantial reduction in visual artifacts via RL, ensuring much cleaner and more professional outputs.
+
+- Technical Stability: Solves numerical instability inherent in autoregressive flow-based models.
+
+## Environment Setup
+
+To avoid potential errors when loading and running your models, we recommend using the following settings:
+
+```shell
+conda create -n nextstep python=3.11 -y
+conda activate nextstep
+
+pip install uv # optional
+
+GIT_LFS_SKIP_SMUDGE=1 git clone https://huggingface.co/stepfun-ai/NextStep-1.1-Pretrain-256px && cd NextStep-1.1-Pretrain-256px
+uv pip install -r requirements.txt
+
+hf download stepfun-ai/NextStep-1.1-Pretrain-256px "vae/checkpoint.pt" --local-dir ./
+```
+
+## Usage
+
+```python
+import torch
+from transformers import AutoTokenizer, AutoModel
+from models.gen_pipeline import NextStepPipeline
+
+HF_HUB = "stepfun-ai/NextStep-1.1-Pretrain-256px"
+
+# load model and tokenizer
+tokenizer = AutoTokenizer.from_pretrained(HF_HUB, local_files_only=True, trust_remote_code=True)
+model = AutoModel.from_pretrained(HF_HUB, local_files_only=True, trust_remote_code=True)
+pipeline = NextStepPipeline(tokenizer=tokenizer, model=model).to(device="cuda", dtype=torch.bfloat16)
+
+# set prompts
+positive_prompt = ""
+negative_prompt = "lowres, bad anatomy, bad hands, text, error, missing fingers, extra digit, fewer digits, cropped, worst quality, low quality, normal quality, jpeg artifacts, signature, watermark, username, blurry."
+example_prompt = "A REALISTIC PHOTOGRAPH OF A WALL WITH \"TOWARD AUTOREGRESSIVE IMAGE GENERATION WITH CONTINUOUS TOKENS AT SCALE\" PROMINENTLY DISPLAYED"
+
+# generate image from text
+IMG_SIZE = 256
+image = pipeline.generate_image(
+    example_prompt,
+    hw=(IMG_SIZE, IMG_SIZE),
+    num_images_per_caption=1,
+    positive_prompt=positive_prompt,
+    negative_prompt=negative_prompt,
+    cfg=7.5,
+    cfg_img=1.0,
+    cfg_schedule="constant",
+    use_norm=False,
+    num_sampling_steps=28,
+    timesteps_shift=1.0,
+    seed=3407,
+)[0]
+image.save("./assets/output.jpg")
+```
+
+## Citation
+
+If you find NextStep useful for your research and applications, please consider starring this repository and citing:
+
+```bibtex
+@article{nextstepteam2025nextstep1,
+  title={NextStep-1: Toward Autoregressive Image Generation with Continuous Tokens at Scale},
+  author={NextStep Team and Chunrui Han and Guopeng Li and Jingwei Wu and Quan Sun and Yan Cai and Yuang Peng and Zheng Ge and Deyu Zhou and Haomiao Tang and Hongyu Zhou and Kenkun Liu and Ailin Huang and Bin Wang and Changxin Miao and Deshan Sun and En Yu and Fukun Yin and Gang Yu and Hao Nie and Haoran Lv and Hanpeng Hu and Jia Wang and Jian Zhou and Jianjian Sun and Kaijun Tan and Kang An and Kangheng Lin and Liang Zhao and Mei Chen and Peng Xing and Rui Wang and Shiyu Liu and Shutao Xia and Tianhao You and Wei Ji and Xianfang Zeng and Xin Han and Xuelin Zhang and Yana Wei and Yanming Xu and Yimin Jiang and Yingming Wang and Yu Zhou and Yucheng Han and Ziyang Meng and Binxing Jiao and Daxin Jiang and Xiangyu Zhang and Yibo Zhu},
+  journal={arXiv preprint arXiv:2508.10711},
+  year={2025}
+}
+```

added_tokens.json

@@ -0,0 +1,34 @@
+{
+  "</tool_call>": 151658,
+  "<tool_call>": 151657,
+  "<|begin_of_image|>": 151667,
+  "<|begin_of_prompt_refinement|>": 151670,
+  "<|begin_of_thinking|>": 151672,
+  "<|box_end|>": 151649,
+  "<|box_start|>": 151648,
+  "<|end_of_image|>": 151668,
+  "<|end_of_prompt_refinement|>": 151671,
+  "<|end_of_thinking|>": 151673,
+  "<|beginoftext|>": 151674,
+  "<|endoftext|>": 151643,
+  "<|file_sep|>": 151664,
+  "<|fim_middle|>": 151660,
+  "<|fim_pad|>": 151662,
+  "<|fim_prefix|>": 151659,
+  "<|fim_suffix|>": 151661,
+  "<|im_end|>": 151645,
+  "<|im_start|>": 151644,
+  "<|image_area|>": 151666,
+  "<|image_pad|>": 151655,
+  "<|image_placeholder|>": 151669,
+  "<|object_ref_end|>": 151647,
+  "<|object_ref_start|>": 151646,
+  "<|quad_end|>": 151651,
+  "<|quad_start|>": 151650,
+  "<|repo_name|>": 151663,
+  "<|video_pad|>": 151656,
+  "<|vision_end|>": 151653,
+  "<|vision_pad|>": 151654,
+  "<|vision_start|>": 151652,
+  "[PAD]": 151665
+}

config.json

@@ -0,0 +1,66 @@
+{
+  "_attn_implementation_autoset": true,
+  "architectures": [
+    "LlamaForCausalLM"
+  ],
+  "auto_map":{
+    "AutoConfig": "models/config.NextStepConfig",
+    "AutoModel": "models/nextstep_model.NextStep"
+  },
+  "attention_bias": true,
+  "attention_dropout": 0.0,
+  "base_image_grid_size": 64,
+  "boi": 151667,
+  "bos_token_id": 151643,
+  "create_kwargs": {
+    "snr_type": "lognorm"
+  },
+  "eoi": 151668,
+  "eos_token_id": 151643,
+  "genloss_batch_mul": 4,
+  "genloss_depth": 12,
+  "genloss_net_arch": "mlp",
+  "genloss_num_sampling_steps": "100",
+  "genloss_type": "transport",
+  "genloss_width": 1536,
+  "head_dim": 128,
+  "hidden_act": "silu",
+  "hidden_size": 5120,
+  "image_decoder_arch": "Trans_E",
+  "image_encoder_name": null,
+  "image_feature_layer": -2,
+  "image_loss_weight": 1.0,
+  "image_placeholder_id": 151669,
+  "image_size": 64,
+  "initializer_range": 0.02,
+  "intermediate_size": 13824,
+  "lm_loss_weight": 0.01,
+  "max_position_embeddings": 131072,
+  "max_window_layers": 48,
+  "mlp_bias": false,
+  "model_type": "nextstep",
+  "noise_strength": 0.0,
+  "num_attention_heads": 40,
+  "num_channels": 16,
+  "num_hidden_layers": 48,
+  "num_key_value_heads": 8,
+  "o_attention_bias": false,
+  "pad_token_id_added": 151665,
+  "patch_size": 2,
+  "pretraining_tp": 1,
+  "rms_norm_eps": 1e-05,
+  "rope_scaling": null,
+  "rope_theta": 1000000.0,
+  "sliding_window": 131072,
+  "tie_word_embeddings": false,
+  "torch_dtype": "bfloat16",
+  "transformers_version": "4.55.0",
+  "use_2d_rope": false,
+  "use_cache": true,
+  "use_gen_pos_embed": false,
+  "use_mlp_before_lm_head": false,
+  "use_sliding_window": false,
+  "use_token_length_weight": false,
+  "vae_name_or_path": "vae/",
+  "vocab_size": 152064
+}

model.safetensors.index.json

@@ -0,0 +1,698 @@
+{
+  "metadata": {
+    "total_size": 29907628160
+  },
+  "weight_map": {
+    "embed_tokens.weight": "pytorch-model-00004-of-00004.safetensors",
+    "image_head.net.cond_embed.bias": "pytorch-model-00003-of-00004.safetensors",
+    "image_head.net.cond_embed.weight": "pytorch-model-00003-of-00004.safetensors",
+    "image_head.net.final_layer.adaLN_modulation.1.bias": "pytorch-model-00003-of-00004.safetensors",
+    "image_head.net.final_layer.adaLN_modulation.1.weight": "pytorch-model-00003-of-00004.safetensors",
+    "image_head.net.final_layer.linear.bias": "pytorch-model-00003-of-00004.safetensors",
+    "image_head.net.final_layer.linear.weight": "pytorch-model-00003-of-00004.safetensors",
+    "image_head.net.input_proj.bias": "pytorch-model-00003-of-00004.safetensors",
+    "image_head.net.input_proj.weight": "pytorch-model-00003-of-00004.safetensors",
+    "image_head.net.res_blocks.0.adaLN_modulation.1.bias": "pytorch-model-00003-of-00004.safetensors",
+    "image_head.net.res_blocks.0.adaLN_modulation.1.weight": "pytorch-model-00003-of-00004.safetensors",
+    "image_head.net.res_blocks.0.in_ln.bias": "pytorch-model-00003-of-00004.safetensors",
+    "image_head.net.res_blocks.0.in_ln.weight": "pytorch-model-00003-of-00004.safetensors",
+    "image_head.net.res_blocks.0.mlp.0.bias": "pytorch-model-00003-of-00004.safetensors",
+    "image_head.net.res_blocks.0.mlp.0.weight": "pytorch-model-00003-of-00004.safetensors",
+    "image_head.net.res_blocks.0.mlp.2.bias": "pytorch-model-00003-of-00004.safetensors",
+    "image_head.net.res_blocks.0.mlp.2.weight": "pytorch-model-00003-of-00004.safetensors",
+    "image_head.net.res_blocks.1.adaLN_modulation.1.bias": "pytorch-model-00003-of-00004.safetensors",
+    "image_head.net.res_blocks.1.adaLN_modulation.1.weight": "pytorch-model-00003-of-00004.safetensors",
+    "image_head.net.res_blocks.1.in_ln.bias": "pytorch-model-00003-of-00004.safetensors",
+    "image_head.net.res_blocks.1.in_ln.weight": "pytorch-model-00003-of-00004.safetensors",
+    "image_head.net.res_blocks.1.mlp.0.bias": "pytorch-model-00003-of-00004.safetensors",
+    "image_head.net.res_blocks.1.mlp.0.weight": "pytorch-model-00003-of-00004.safetensors",
+    "image_head.net.res_blocks.1.mlp.2.bias": "pytorch-model-00003-of-00004.safetensors",
+    "image_head.net.res_blocks.1.mlp.2.weight": "pytorch-model-00003-of-00004.safetensors",
+    "image_head.net.res_blocks.10.adaLN_modulation.1.bias": "pytorch-model-00003-of-00004.safetensors",
+    "image_head.net.res_blocks.10.adaLN_modulation.1.weight": "pytorch-model-00003-of-00004.safetensors",
+    "image_head.net.res_blocks.10.in_ln.bias": "pytorch-model-00003-of-00004.safetensors",
+    "image_head.net.res_blocks.10.in_ln.weight": "pytorch-model-00003-of-00004.safetensors",
+    "image_head.net.res_blocks.10.mlp.0.bias": "pytorch-model-00003-of-00004.safetensors",
+    "image_head.net.res_blocks.10.mlp.0.weight": "pytorch-model-00003-of-00004.safetensors",
+    "image_head.net.res_blocks.10.mlp.2.bias": "pytorch-model-00003-of-00004.safetensors",
+    "image_head.net.res_blocks.10.mlp.2.weight": "pytorch-model-00003-of-00004.safetensors",
+    "image_head.net.res_blocks.11.adaLN_modulation.1.bias": "pytorch-model-00003-of-00004.safetensors",
+    "image_head.net.res_blocks.11.adaLN_modulation.1.weight": "pytorch-model-00003-of-00004.safetensors",
+    "image_head.net.res_blocks.11.in_ln.bias": "pytorch-model-00003-of-00004.safetensors",
+    "image_head.net.res_blocks.11.in_ln.weight": "pytorch-model-00003-of-00004.safetensors",
+    "image_head.net.res_blocks.11.mlp.0.bias": "pytorch-model-00003-of-00004.safetensors",
+    "image_head.net.res_blocks.11.mlp.0.weight": "pytorch-model-00003-of-00004.safetensors",
+    "image_head.net.res_blocks.11.mlp.2.bias": "pytorch-model-00003-of-00004.safetensors",
+    "image_head.net.res_blocks.11.mlp.2.weight": "pytorch-model-00003-of-00004.safetensors",
+    "image_head.net.res_blocks.2.adaLN_modulation.1.bias": "pytorch-model-00003-of-00004.safetensors",
+    "image_head.net.res_blocks.2.adaLN_modulation.1.weight": "pytorch-model-00003-of-00004.safetensors",
+    "image_head.net.res_blocks.2.in_ln.bias": "pytorch-model-00003-of-00004.safetensors",
+    "image_head.net.res_blocks.2.in_ln.weight": "pytorch-model-00003-of-00004.safetensors",
+    "image_head.net.res_blocks.2.mlp.0.bias": "pytorch-model-00003-of-00004.safetensors",
+    "image_head.net.res_blocks.2.mlp.0.weight": "pytorch-model-00003-of-00004.safetensors",
+    "image_head.net.res_blocks.2.mlp.2.bias": "pytorch-model-00003-of-00004.safetensors",
+    "image_head.net.res_blocks.2.mlp.2.weight": "pytorch-model-00003-of-00004.safetensors",
+    "image_head.net.res_blocks.3.adaLN_modulation.1.bias": "pytorch-model-00003-of-00004.safetensors",
+    "image_head.net.res_blocks.3.adaLN_modulation.1.weight": "pytorch-model-00003-of-00004.safetensors",
+    "image_head.net.res_blocks.3.in_ln.bias": "pytorch-model-00003-of-00004.safetensors",
+    "image_head.net.res_blocks.3.in_ln.weight": "pytorch-model-00003-of-00004.safetensors",
+    "image_head.net.res_blocks.3.mlp.0.bias": "pytorch-model-00003-of-00004.safetensors",
+    "image_head.net.res_blocks.3.mlp.0.weight": "pytorch-model-00003-of-00004.safetensors",
+    "image_head.net.res_blocks.3.mlp.2.bias": "pytorch-model-00003-of-00004.safetensors",
+    "image_head.net.res_blocks.3.mlp.2.weight": "pytorch-model-00003-of-00004.safetensors",
+    "image_head.net.res_blocks.4.adaLN_modulation.1.bias": "pytorch-model-00003-of-00004.safetensors",
+    "image_head.net.res_blocks.4.adaLN_modulation.1.weight": "pytorch-model-00003-of-00004.safetensors",
+    "image_head.net.res_blocks.4.in_ln.bias": "pytorch-model-00003-of-00004.safetensors",
+    "image_head.net.res_blocks.4.in_ln.weight": "pytorch-model-00003-of-00004.safetensors",
+    "image_head.net.res_blocks.4.mlp.0.bias": "pytorch-model-00003-of-00004.safetensors",
+    "image_head.net.res_blocks.4.mlp.0.weight": "pytorch-model-00003-of-00004.safetensors",
+    "image_head.net.res_blocks.4.mlp.2.bias": "pytorch-model-00003-of-00004.safetensors",
+    "image_head.net.res_blocks.4.mlp.2.weight": "pytorch-model-00003-of-00004.safetensors",
+    "image_head.net.res_blocks.5.adaLN_modulation.1.bias": "pytorch-model-00003-of-00004.safetensors",
+    "image_head.net.res_blocks.5.adaLN_modulation.1.weight": "pytorch-model-00003-of-00004.safetensors",
+    "image_head.net.res_blocks.5.in_ln.bias": "pytorch-model-00003-of-00004.safetensors",
+    "image_head.net.res_blocks.5.in_ln.weight": "pytorch-model-00003-of-00004.safetensors",
+    "image_head.net.res_blocks.5.mlp.0.bias": "pytorch-model-00003-of-00004.safetensors",
+    "image_head.net.res_blocks.5.mlp.0.weight": "pytorch-model-00003-of-00004.safetensors",
+    "image_head.net.res_blocks.5.mlp.2.bias": "pytorch-model-00003-of-00004.safetensors",
+    "image_head.net.res_blocks.5.mlp.2.weight": "pytorch-model-00003-of-00004.safetensors",
+    "image_head.net.res_blocks.6.adaLN_modulation.1.bias": "pytorch-model-00003-of-00004.safetensors",
+    "image_head.net.res_blocks.6.adaLN_modulation.1.weight": "pytorch-model-00003-of-00004.safetensors",
+    "image_head.net.res_blocks.6.in_ln.bias": "pytorch-model-00003-of-00004.safetensors",
+    "image_head.net.res_blocks.6.in_ln.weight": "pytorch-model-00003-of-00004.safetensors",
+    "image_head.net.res_blocks.6.mlp.0.bias": "pytorch-model-00003-of-00004.safetensors",
+    "image_head.net.res_blocks.6.mlp.0.weight": "pytorch-model-00003-of-00004.safetensors",
+    "image_head.net.res_blocks.6.mlp.2.bias": "pytorch-model-00003-of-00004.safetensors",
+    "image_head.net.res_blocks.6.mlp.2.weight": "pytorch-model-00003-of-00004.safetensors",
+    "image_head.net.res_blocks.7.adaLN_modulation.1.bias": "pytorch-model-00003-of-00004.safetensors",
+    "image_head.net.res_blocks.7.adaLN_modulation.1.weight": "pytorch-model-00003-of-00004.safetensors",
+    "image_head.net.res_blocks.7.in_ln.bias": "pytorch-model-00003-of-00004.safetensors",
+    "image_head.net.res_blocks.7.in_ln.weight": "pytorch-model-00003-of-00004.safetensors",
+    "image_head.net.res_blocks.7.mlp.0.bias": "pytorch-model-00003-of-00004.safetensors",
+    "image_head.net.res_blocks.7.mlp.0.weight": "pytorch-model-00003-of-00004.safetensors",
+    "image_head.net.res_blocks.7.mlp.2.bias": "pytorch-model-00003-of-00004.safetensors",
+    "image_head.net.res_blocks.7.mlp.2.weight": "pytorch-model-00003-of-00004.safetensors",
+    "image_head.net.res_blocks.8.adaLN_modulation.1.bias": "pytorch-model-00003-of-00004.safetensors",
+    "image_head.net.res_blocks.8.adaLN_modulation.1.weight": "pytorch-model-00003-of-00004.safetensors",
+    "image_head.net.res_blocks.8.in_ln.bias": "pytorch-model-00003-of-00004.safetensors",
+    "image_head.net.res_blocks.8.in_ln.weight": "pytorch-model-00003-of-00004.safetensors",
+    "image_head.net.res_blocks.8.mlp.0.bias": "pytorch-model-00003-of-00004.safetensors",
+    "image_head.net.res_blocks.8.mlp.0.weight": "pytorch-model-00003-of-00004.safetensors",
+    "image_head.net.res_blocks.8.mlp.2.bias": "pytorch-model-00003-of-00004.safetensors",
+    "image_head.net.res_blocks.8.mlp.2.weight": "pytorch-model-00003-of-00004.safetensors",
+    "image_head.net.res_blocks.9.adaLN_modulation.1.bias": "pytorch-model-00003-of-00004.safetensors",
+    "image_head.net.res_blocks.9.adaLN_modulation.1.weight": "pytorch-model-00003-of-00004.safetensors",
+    "image_head.net.res_blocks.9.in_ln.bias": "pytorch-model-00003-of-00004.safetensors",
+    "image_head.net.res_blocks.9.in_ln.weight": "pytorch-model-00003-of-00004.safetensors",
+    "image_head.net.res_blocks.9.mlp.0.bias": "pytorch-model-00003-of-00004.safetensors",
+    "image_head.net.res_blocks.9.mlp.0.weight": "pytorch-model-00003-of-00004.safetensors",
+    "image_head.net.res_blocks.9.mlp.2.bias": "pytorch-model-00003-of-00004.safetensors",
+    "image_head.net.res_blocks.9.mlp.2.weight": "pytorch-model-00003-of-00004.safetensors",
+    "image_head.net.time_embed.mlp.0.bias": "pytorch-model-00003-of-00004.safetensors",
+    "image_head.net.time_embed.mlp.0.weight": "pytorch-model-00003-of-00004.safetensors",
+    "image_head.net.time_embed.mlp.2.bias": "pytorch-model-00003-of-00004.safetensors",
+    "image_head.net.time_embed.mlp.2.weight": "pytorch-model-00003-of-00004.safetensors",
+    "image_in_projector.bias": "pytorch-model-00003-of-00004.safetensors",
+    "image_in_projector.weight": "pytorch-model-00003-of-00004.safetensors",
+    "image_out_projector.bias": "pytorch-model-00003-of-00004.safetensors",
+    "image_out_projector.weight": "pytorch-model-00003-of-00004.safetensors",
+    "layers.0.input_layernorm.weight": "pytorch-model-00001-of-00004.safetensors",
+    "layers.0.mlp.down_proj.weight": "pytorch-model-00001-of-00004.safetensors",
+    "layers.0.mlp.gate_proj.weight": "pytorch-model-00001-of-00004.safetensors",
+    "layers.0.mlp.up_proj.weight": "pytorch-model-00001-of-00004.safetensors",
+    "layers.0.post_attention_layernorm.weight": "pytorch-model-00001-of-00004.safetensors",
+    "layers.0.self_attn.k_proj.bias": "pytorch-model-00001-of-00004.safetensors",
+    "layers.0.self_attn.k_proj.weight": "pytorch-model-00001-of-00004.safetensors",
+    "layers.0.self_attn.o_proj.weight": "pytorch-model-00001-of-00004.safetensors",
+    "layers.0.self_attn.q_proj.bias": "pytorch-model-00001-of-00004.safetensors",
+    "layers.0.self_attn.q_proj.weight": "pytorch-model-00001-of-00004.safetensors",
+    "layers.0.self_attn.v_proj.bias": "pytorch-model-00001-of-00004.safetensors",
+    "layers.0.self_attn.v_proj.weight": "pytorch-model-00001-of-00004.safetensors",
+    "layers.1.input_layernorm.weight": "pytorch-model-00001-of-00004.safetensors",
+    "layers.1.mlp.down_proj.weight": "pytorch-model-00001-of-00004.safetensors",
+    "layers.1.mlp.gate_proj.weight": "pytorch-model-00001-of-00004.safetensors",
+    "layers.1.mlp.up_proj.weight": "pytorch-model-00001-of-00004.safetensors",
+    "layers.1.post_attention_layernorm.weight": "pytorch-model-00001-of-00004.safetensors",
+    "layers.1.self_attn.k_proj.bias": "pytorch-model-00001-of-00004.safetensors",
+    "layers.1.self_attn.k_proj.weight": "pytorch-model-00001-of-00004.safetensors",
+    "layers.1.self_attn.o_proj.weight": "pytorch-model-00001-of-00004.safetensors",
+    "layers.1.self_attn.q_proj.bias": "pytorch-model-00001-of-00004.safetensors",
+    "layers.1.self_attn.q_proj.weight": "pytorch-model-00001-of-00004.safetensors",
+    "layers.1.self_attn.v_proj.bias": "pytorch-model-00001-of-00004.safetensors",
+    "layers.1.self_attn.v_proj.weight": "pytorch-model-00001-of-00004.safetensors",
+    "layers.10.input_layernorm.weight": "pytorch-model-00001-of-00004.safetensors",
+    "layers.10.mlp.down_proj.weight": "pytorch-model-00001-of-00004.safetensors",
+    "layers.10.mlp.gate_proj.weight": "pytorch-model-00001-of-00004.safetensors",
+    "layers.10.mlp.up_proj.weight": "pytorch-model-00001-of-00004.safetensors",
+    "layers.10.post_attention_layernorm.weight": "pytorch-model-00001-of-00004.safetensors",
+    "layers.10.self_attn.k_proj.bias": "pytorch-model-00001-of-00004.safetensors",
+    "layers.10.self_attn.k_proj.weight": "pytorch-model-00001-of-00004.safetensors",
+    "layers.10.self_attn.o_proj.weight": "pytorch-model-00001-of-00004.safetensors",
+    "layers.10.self_attn.q_proj.bias": "pytorch-model-00001-of-00004.safetensors",
+    "layers.10.self_attn.q_proj.weight": "pytorch-model-00001-of-00004.safetensors",
+    "layers.10.self_attn.v_proj.bias": "pytorch-model-00001-of-00004.safetensors",
+    "layers.10.self_attn.v_proj.weight": "pytorch-model-00001-of-00004.safetensors",
+    "layers.11.input_layernorm.weight": "pytorch-model-00001-of-00004.safetensors",
+    "layers.11.mlp.down_proj.weight": "pytorch-model-00001-of-00004.safetensors",
+    "layers.11.mlp.gate_proj.weight": "pytorch-model-00001-of-00004.safetensors",
+    "layers.11.mlp.up_proj.weight": "pytorch-model-00001-of-00004.safetensors",
+    "layers.11.post_attention_layernorm.weight": "pytorch-model-00001-of-00004.safetensors",
+    "layers.11.self_attn.k_proj.bias": "pytorch-model-00001-of-00004.safetensors",
+    "layers.11.self_attn.k_proj.weight": "pytorch-model-00001-of-00004.safetensors",
+    "layers.11.self_attn.o_proj.weight": "pytorch-model-00001-of-00004.safetensors",
+    "layers.11.self_attn.q_proj.bias": "pytorch-model-00001-of-00004.safetensors",
+    "layers.11.self_attn.q_proj.weight": "pytorch-model-00001-of-00004.safetensors",
+    "layers.11.self_attn.v_proj.bias": "pytorch-model-00001-of-00004.safetensors",
+    "layers.11.self_attn.v_proj.weight": "pytorch-model-00001-of-00004.safetensors",
+    "layers.12.input_layernorm.weight": "pytorch-model-00001-of-00004.safetensors",
+    "layers.12.mlp.down_proj.weight": "pytorch-model-00001-of-00004.safetensors",
+    "layers.12.mlp.gate_proj.weight": "pytorch-model-00001-of-00004.safetensors",
+    "layers.12.mlp.up_proj.weight": "pytorch-model-00001-of-00004.safetensors",
+    "layers.12.post_attention_layernorm.weight": "pytorch-model-00001-of-00004.safetensors",
+    "layers.12.self_attn.k_proj.bias": "pytorch-model-00001-of-00004.safetensors",
+    "layers.12.self_attn.k_proj.weight": "pytorch-model-00001-of-00004.safetensors",
+    "layers.12.self_attn.o_proj.weight": "pytorch-model-00001-of-00004.safetensors",
+    "layers.12.self_attn.q_proj.bias": "pytorch-model-00001-of-00004.safetensors",
+    "layers.12.self_attn.q_proj.weight": "pytorch-model-00001-of-00004.safetensors",
+    "layers.12.self_attn.v_proj.bias": "pytorch-model-00001-of-00004.safetensors",
+    "layers.12.self_attn.v_proj.weight": "pytorch-model-00001-of-00004.safetensors",
+    "layers.13.input_layernorm.weight": "pytorch-model-00001-of-00004.safetensors",
+    "layers.13.mlp.down_proj.weight": "pytorch-model-00001-of-00004.safetensors",
+    "layers.13.mlp.gate_proj.weight": "pytorch-model-00001-of-00004.safetensors",
+    "layers.13.mlp.up_proj.weight": "pytorch-model-00001-of-00004.safetensors",
+    "layers.13.post_attention_layernorm.weight": "pytorch-model-00001-of-00004.safetensors",
+    "layers.13.self_attn.k_proj.bias": "pytorch-model-00001-of-00004.safetensors",
+    "layers.13.self_attn.k_proj.weight": "pytorch-model-00001-of-00004.safetensors",
+    "layers.13.self_attn.o_proj.weight": "pytorch-model-00001-of-00004.safetensors",
+    "layers.13.self_attn.q_proj.bias": "pytorch-model-00001-of-00004.safetensors",
+    "layers.13.self_attn.q_proj.weight": "pytorch-model-00001-of-00004.safetensors",
+    "layers.13.self_attn.v_proj.bias": "pytorch-model-00001-of-00004.safetensors",
+    "layers.13.self_attn.v_proj.weight": "pytorch-model-00001-of-00004.safetensors",
+    "layers.14.input_layernorm.weight": "pytorch-model-00001-of-00004.safetensors",
+    "layers.14.mlp.down_proj.weight": "pytorch-model-00001-of-00004.safetensors",
+    "layers.14.mlp.gate_proj.weight": "pytorch-model-00001-of-00004.safetensors",
+    "layers.14.mlp.up_proj.weight": "pytorch-model-00001-of-00004.safetensors",
+    "layers.14.post_attention_layernorm.weight": "pytorch-model-00001-of-00004.safetensors",
+    "layers.14.self_attn.k_proj.bias": "pytorch-model-00001-of-00004.safetensors",
+    "layers.14.self_attn.k_proj.weight": "pytorch-model-00001-of-00004.safetensors",
+    "layers.14.self_attn.o_proj.weight": "pytorch-model-00001-of-00004.safetensors",
+    "layers.14.self_attn.q_proj.bias": "pytorch-model-00001-of-00004.safetensors",
+    "layers.14.self_attn.q_proj.weight": "pytorch-model-00001-of-00004.safetensors",
+    "layers.14.self_attn.v_proj.bias": "pytorch-model-00001-of-00004.safetensors",
+    "layers.14.self_attn.v_proj.weight": "pytorch-model-00001-of-00004.safetensors",
+    "layers.15.input_layernorm.weight": "pytorch-model-00001-of-00004.safetensors",
+    "layers.15.mlp.down_proj.weight": "pytorch-model-00001-of-00004.safetensors",
+    "layers.15.mlp.gate_proj.weight": "pytorch-model-00001-of-00004.safetensors",
+    "layers.15.mlp.up_proj.weight": "pytorch-model-00001-of-00004.safetensors",
+    "layers.15.post_attention_layernorm.weight": "pytorch-model-00001-of-00004.safetensors",
+    "layers.15.self_attn.k_proj.bias": "pytorch-model-00001-of-00004.safetensors",
+    "layers.15.self_attn.k_proj.weight": "pytorch-model-00001-of-00004.safetensors",
+    "layers.15.self_attn.o_proj.weight": "pytorch-model-00001-of-00004.safetensors",
+    "layers.15.self_attn.q_proj.bias": "pytorch-model-00001-of-00004.safetensors",
+    "layers.15.self_attn.q_proj.weight": "pytorch-model-00001-of-00004.safetensors",
+    "layers.15.self_attn.v_proj.bias": "pytorch-model-00001-of-00004.safetensors",
+    "layers.15.self_attn.v_proj.weight": "pytorch-model-00001-of-00004.safetensors",
+    "layers.16.input_layernorm.weight": "pytorch-model-00001-of-00004.safetensors",
+    "layers.16.mlp.down_proj.weight": "pytorch-model-00001-of-00004.safetensors",
+    "layers.16.mlp.gate_proj.weight": "pytorch-model-00001-of-00004.safetensors",
+    "layers.16.mlp.up_proj.weight": "pytorch-model-00001-of-00004.safetensors",
+    "layers.16.post_attention_layernorm.weight": "pytorch-model-00001-of-00004.safetensors",
+    "layers.16.self_attn.k_proj.bias": "pytorch-model-00001-of-00004.safetensors",
+    "layers.16.self_attn.k_proj.weight": "pytorch-model-00001-of-00004.safetensors",
+    "layers.16.self_attn.o_proj.weight": "pytorch-model-00001-of-00004.safetensors",
+    "layers.16.self_attn.q_proj.bias": "pytorch-model-00001-of-00004.safetensors",
+    "layers.16.self_attn.q_proj.weight": "pytorch-model-00001-of-00004.safetensors",
+    "layers.16.self_attn.v_proj.bias": "pytorch-model-00001-of-00004.safetensors",
+    "layers.16.self_attn.v_proj.weight": "pytorch-model-00001-of-00004.safetensors",
+    "layers.17.input_layernorm.weight": "pytorch-model-00001-of-00004.safetensors",
+    "layers.17.mlp.down_proj.weight": "pytorch-model-00001-of-00004.safetensors",
+    "layers.17.mlp.gate_proj.weight": "pytorch-model-00001-of-00004.safetensors",
+    "layers.17.mlp.up_proj.weight": "pytorch-model-00001-of-00004.safetensors",
+    "layers.17.post_attention_layernorm.weight": "pytorch-model-00001-of-00004.safetensors",
+    "layers.17.self_attn.k_proj.bias": "pytorch-model-00001-of-00004.safetensors",
+    "layers.17.self_attn.k_proj.weight": "pytorch-model-00001-of-00004.safetensors",
+    "layers.17.self_attn.o_proj.weight": "pytorch-model-00001-of-00004.safetensors",
+    "layers.17.self_attn.q_proj.bias": "pytorch-model-00001-of-00004.safetensors",
+    "layers.17.self_attn.q_proj.weight": "pytorch-model-00001-of-00004.safetensors",
+    "layers.17.self_attn.v_proj.bias": "pytorch-model-00001-of-00004.safetensors",
+    "layers.17.self_attn.v_proj.weight": "pytorch-model-00001-of-00004.safetensors",
+    "layers.18.input_layernorm.weight": "pytorch-model-00001-of-00004.safetensors",
+    "layers.18.mlp.down_proj.weight": "pytorch-model-00001-of-00004.safetensors",
+    "layers.18.mlp.gate_proj.weight": "pytorch-model-00001-of-00004.safetensors",
+    "layers.18.mlp.up_proj.weight": "pytorch-model-00001-of-00004.safetensors",
+    "layers.18.post_attention_layernorm.weight": "pytorch-model-00001-of-00004.safetensors",
+    "layers.18.self_attn.k_proj.bias": "pytorch-model-00001-of-00004.safetensors",
+    "layers.18.self_attn.k_proj.weight": "pytorch-model-00001-of-00004.safetensors",
+    "layers.18.self_attn.o_proj.weight": "pytorch-model-00001-of-00004.safetensors",
+    "layers.18.self_attn.q_proj.bias": "pytorch-model-00001-of-00004.safetensors",
+    "layers.18.self_attn.q_proj.weight": "pytorch-model-00001-of-00004.safetensors",
+    "layers.18.self_attn.v_proj.bias": "pytorch-model-00001-of-00004.safetensors",
+    "layers.18.self_attn.v_proj.weight": "pytorch-model-00001-of-00004.safetensors",
+    "layers.19.input_layernorm.weight": "pytorch-model-00001-of-00004.safetensors",
+    "layers.19.mlp.down_proj.weight": "pytorch-model-00001-of-00004.safetensors",
+    "layers.19.mlp.gate_proj.weight": "pytorch-model-00001-of-00004.safetensors",
+    "layers.19.mlp.up_proj.weight": "pytorch-model-00001-of-00004.safetensors",
+    "layers.19.post_attention_layernorm.weight": "pytorch-model-00001-of-00004.safetensors",
+    "layers.19.self_attn.k_proj.bias": "pytorch-model-00001-of-00004.safetensors",
+    "layers.19.self_attn.k_proj.weight": "pytorch-model-00001-of-00004.safetensors",
+    "layers.19.self_attn.o_proj.weight": "pytorch-model-00001-of-00004.safetensors",
+    "layers.19.self_attn.q_proj.bias": "pytorch-model-00001-of-00004.safetensors",
+    "layers.19.self_attn.q_proj.weight": "pytorch-model-00001-of-00004.safetensors",
+    "layers.19.self_attn.v_proj.bias": "pytorch-model-00001-of-00004.safetensors",
+    "layers.19.self_attn.v_proj.weight": "pytorch-model-00001-of-00004.safetensors",
+    "layers.2.input_layernorm.weight": "pytorch-model-00001-of-00004.safetensors",
+    "layers.2.mlp.down_proj.weight": "pytorch-model-00001-of-00004.safetensors",
+    "layers.2.mlp.gate_proj.weight": "pytorch-model-00001-of-00004.safetensors",
+    "layers.2.mlp.up_proj.weight": "pytorch-model-00001-of-00004.safetensors",
+    "layers.2.post_attention_layernorm.weight": "pytorch-model-00001-of-00004.safetensors",
+    "layers.2.self_attn.k_proj.bias": "pytorch-model-00001-of-00004.safetensors",
+    "layers.2.self_attn.k_proj.weight": "pytorch-model-00001-of-00004.safetensors",
+    "layers.2.self_attn.o_proj.weight": "pytorch-model-00001-of-00004.safetensors",
+    "layers.2.self_attn.q_proj.bias": "pytorch-model-00001-of-00004.safetensors",
+    "layers.2.self_attn.q_proj.weight": "pytorch-model-00001-of-00004.safetensors",
+    "layers.2.self_attn.v_proj.bias": "pytorch-model-00001-of-00004.safetensors",
+    "layers.2.self_attn.v_proj.weight": "pytorch-model-00001-of-00004.safetensors",
+    "layers.20.input_layernorm.weight": "pytorch-model-00001-of-00004.safetensors",
+    "layers.20.mlp.down_proj.weight": "pytorch-model-00001-of-00004.safetensors",
+    "layers.20.mlp.gate_proj.weight": "pytorch-model-00001-of-00004.safetensors",
+    "layers.20.mlp.up_proj.weight": "pytorch-model-00001-of-00004.safetensors",
+    "layers.20.post_attention_layernorm.weight": "pytorch-model-00001-of-00004.safetensors",
+    "layers.20.self_attn.k_proj.bias": "pytorch-model-00001-of-00004.safetensors",
+    "layers.20.self_attn.k_proj.weight": "pytorch-model-00001-of-00004.safetensors",
+    "layers.20.self_attn.o_proj.weight": "pytorch-model-00001-of-00004.safetensors",
+    "layers.20.self_attn.q_proj.bias": "pytorch-model-00001-of-00004.safetensors",
+    "layers.20.self_attn.q_proj.weight": "pytorch-model-00001-of-00004.safetensors",
+    "layers.20.self_attn.v_proj.bias": "pytorch-model-00001-of-00004.safetensors",
+    "layers.20.self_attn.v_proj.weight": "pytorch-model-00001-of-00004.safetensors",
+    "layers.21.input_layernorm.weight": "pytorch-model-00001-of-00004.safetensors",
+    "layers.21.mlp.down_proj.weight": "pytorch-model-00001-of-00004.safetensors",
+    "layers.21.mlp.gate_proj.weight": "pytorch-model-00001-of-00004.safetensors",
+    "layers.21.mlp.up_proj.weight": "pytorch-model-00001-of-00004.safetensors",
+    "layers.21.post_attention_layernorm.weight": "pytorch-model-00001-of-00004.safetensors",
+    "layers.21.self_attn.k_proj.bias": "pytorch-model-00001-of-00004.safetensors",
+    "layers.21.self_attn.k_proj.weight": "pytorch-model-00001-of-00004.safetensors",
+    "layers.21.self_attn.o_proj.weight": "pytorch-model-00001-of-00004.safetensors",
+    "layers.21.self_attn.q_proj.bias": "pytorch-model-00001-of-00004.safetensors",
+    "layers.21.self_attn.q_proj.weight": "pytorch-model-00001-of-00004.safetensors",
+    "layers.21.self_attn.v_proj.bias": "pytorch-model-00001-of-00004.safetensors",
+    "layers.21.self_attn.v_proj.weight": "pytorch-model-00001-of-00004.safetensors",
+    "layers.22.input_layernorm.weight": "pytorch-model-00001-of-00004.safetensors",
+    "layers.22.mlp.down_proj.weight": "pytorch-model-00001-of-00004.safetensors",
+    "layers.22.mlp.gate_proj.weight": "pytorch-model-00001-of-00004.safetensors",
+    "layers.22.mlp.up_proj.weight": "pytorch-model-00001-of-00004.safetensors",
+    "layers.22.post_attention_layernorm.weight": "pytorch-model-00001-of-00004.safetensors",
+    "layers.22.self_attn.k_proj.bias": "pytorch-model-00001-of-00004.safetensors",
+    "layers.22.self_attn.k_proj.weight": "pytorch-model-00001-of-00004.safetensors",
+    "layers.22.self_attn.o_proj.weight": "pytorch-model-00001-of-00004.safetensors",
+    "layers.22.self_attn.q_proj.bias": "pytorch-model-00001-of-00004.safetensors",
+    "layers.22.self_attn.q_proj.weight": "pytorch-model-00001-of-00004.safetensors",
+    "layers.22.self_attn.v_proj.bias": "pytorch-model-00001-of-00004.safetensors",
+    "layers.22.self_attn.v_proj.weight": "pytorch-model-00001-of-00004.safetensors",
+    "layers.23.input_layernorm.weight": "pytorch-model-00001-of-00004.safetensors",
+    "layers.23.mlp.down_proj.weight": "pytorch-model-00001-of-00004.safetensors",
+    "layers.23.mlp.gate_proj.weight": "pytorch-model-00001-of-00004.safetensors",
+    "layers.23.mlp.up_proj.weight": "pytorch-model-00001-of-00004.safetensors",
+    "layers.23.post_attention_layernorm.weight": "pytorch-model-00001-of-00004.safetensors",
+    "layers.23.self_attn.k_proj.bias": "pytorch-model-00001-of-00004.safetensors",
+    "layers.23.self_attn.k_proj.weight": "pytorch-model-00001-of-00004.safetensors",
+    "layers.23.self_attn.o_proj.weight": "pytorch-model-00001-of-00004.safetensors",
+    "layers.23.self_attn.q_proj.bias": "pytorch-model-00001-of-00004.safetensors",
+    "layers.23.self_attn.q_proj.weight": "pytorch-model-00001-of-00004.safetensors",
+    "layers.23.self_attn.v_proj.bias": "pytorch-model-00001-of-00004.safetensors",
+    "layers.23.self_attn.v_proj.weight": "pytorch-model-00001-of-00004.safetensors",
+    "layers.24.input_layernorm.weight": "pytorch-model-00001-of-00004.safetensors",
+    "layers.24.mlp.down_proj.weight": "pytorch-model-00001-of-00004.safetensors",
+    "layers.24.mlp.gate_proj.weight": "pytorch-model-00001-of-00004.safetensors",
+    "layers.24.mlp.up_proj.weight": "pytorch-model-00001-of-00004.safetensors",
+    "layers.24.post_attention_layernorm.weight": "pytorch-model-00001-of-00004.safetensors",
+    "layers.24.self_attn.k_proj.bias": "pytorch-model-00001-of-00004.safetensors",
+    "layers.24.self_attn.k_proj.weight": "pytorch-model-00001-of-00004.safetensors",
+    "layers.24.self_attn.o_proj.weight": "pytorch-model-00001-of-00004.safetensors",
+    "layers.24.self_attn.q_proj.bias": "pytorch-model-00001-of-00004.safetensors",
+    "layers.24.self_attn.q_proj.weight": "pytorch-model-00001-of-00004.safetensors",
+    "layers.24.self_attn.v_proj.bias": "pytorch-model-00001-of-00004.safetensors",
+    "layers.24.self_attn.v_proj.weight": "pytorch-model-00001-of-00004.safetensors",
+    "layers.25.input_layernorm.weight": "pytorch-model-00001-of-00004.safetensors",
+    "layers.25.mlp.down_proj.weight": "pytorch-model-00002-of-00004.safetensors",
+    "layers.25.mlp.gate_proj.weight": "pytorch-model-00002-of-00004.safetensors",
+    "layers.25.mlp.up_proj.weight": "pytorch-model-00002-of-00004.safetensors",
+    "layers.25.post_attention_layernorm.weight": "pytorch-model-00002-of-00004.safetensors",
+    "layers.25.self_attn.k_proj.bias": "pytorch-model-00002-of-00004.safetensors",
+    "layers.25.self_attn.k_proj.weight": "pytorch-model-00002-of-00004.safetensors",
+    "layers.25.self_attn.o_proj.weight": "pytorch-model-00001-of-00004.safetensors",
+    "layers.25.self_attn.q_proj.bias": "pytorch-model-00002-of-00004.safetensors",
+    "layers.25.self_attn.q_proj.weight": "pytorch-model-00002-of-00004.safetensors",
+    "layers.25.self_attn.v_proj.bias": "pytorch-model-00002-of-00004.safetensors",
+    "layers.25.self_attn.v_proj.weight": "pytorch-model-00002-of-00004.safetensors",
+    "layers.26.input_layernorm.weight": "pytorch-model-00002-of-00004.safetensors",
+    "layers.26.mlp.down_proj.weight": "pytorch-model-00002-of-00004.safetensors",
+    "layers.26.mlp.gate_proj.weight": "pytorch-model-00002-of-00004.safetensors",
+    "layers.26.mlp.up_proj.weight": "pytorch-model-00002-of-00004.safetensors",
+    "layers.26.post_attention_layernorm.weight": "pytorch-model-00002-of-00004.safetensors",
+    "layers.26.self_attn.k_proj.bias": "pytorch-model-00002-of-00004.safetensors",
+    "layers.26.self_attn.k_proj.weight": "pytorch-model-00002-of-00004.safetensors",
+    "layers.26.self_attn.o_proj.weight": "pytorch-model-00002-of-00004.safetensors",
+    "layers.26.self_attn.q_proj.bias": "pytorch-model-00002-of-00004.safetensors",
+    "layers.26.self_attn.q_proj.weight": "pytorch-model-00002-of-00004.safetensors",
+    "layers.26.self_attn.v_proj.bias": "pytorch-model-00002-of-00004.safetensors",
+    "layers.26.self_attn.v_proj.weight": "pytorch-model-00002-of-00004.safetensors",
+    "layers.27.input_layernorm.weight": "pytorch-model-00002-of-00004.safetensors",
+    "layers.27.mlp.down_proj.weight": "pytorch-model-00002-of-00004.safetensors",
+    "layers.27.mlp.gate_proj.weight": "pytorch-model-00002-of-00004.safetensors",
+    "layers.27.mlp.up_proj.weight": "pytorch-model-00002-of-00004.safetensors",
+    "layers.27.post_attention_layernorm.weight": "pytorch-model-00002-of-00004.safetensors",
+    "layers.27.self_attn.k_proj.bias": "pytorch-model-00002-of-00004.safetensors",
+    "layers.27.self_attn.k_proj.weight": "pytorch-model-00002-of-00004.safetensors",
+    "layers.27.self_attn.o_proj.weight": "pytorch-model-00002-of-00004.safetensors",
+    "layers.27.self_attn.q_proj.bias": "pytorch-model-00002-of-00004.safetensors",
+    "layers.27.self_attn.q_proj.weight": "pytorch-model-00002-of-00004.safetensors",
+    "layers.27.self_attn.v_proj.bias": "pytorch-model-00002-of-00004.safetensors",
+    "layers.27.self_attn.v_proj.weight": "pytorch-model-00002-of-00004.safetensors",
+    "layers.28.input_layernorm.weight": "pytorch-model-00002-of-00004.safetensors",
+    "layers.28.mlp.down_proj.weight": "pytorch-model-00002-of-00004.safetensors",
+    "layers.28.mlp.gate_proj.weight": "pytorch-model-00002-of-00004.safetensors",
+    "layers.28.mlp.up_proj.weight": "pytorch-model-00002-of-00004.safetensors",
+    "layers.28.post_attention_layernorm.weight": "pytorch-model-00002-of-00004.safetensors",
+    "layers.28.self_attn.k_proj.bias": "pytorch-model-00002-of-00004.safetensors",
+    "layers.28.self_attn.k_proj.weight": "pytorch-model-00002-of-00004.safetensors",
+    "layers.28.self_attn.o_proj.weight": "pytorch-model-00002-of-00004.safetensors",
+    "layers.28.self_attn.q_proj.bias": "pytorch-model-00002-of-00004.safetensors",
+    "layers.28.self_attn.q_proj.weight": "pytorch-model-00002-of-00004.safetensors",
+    "layers.28.self_attn.v_proj.bias": "pytorch-model-00002-of-00004.safetensors",
+    "layers.28.self_attn.v_proj.weight": "pytorch-model-00002-of-00004.safetensors",
+    "layers.29.input_layernorm.weight": "pytorch-model-00002-of-00004.safetensors",
+    "layers.29.mlp.down_proj.weight": "pytorch-model-00002-of-00004.safetensors",
+    "layers.29.mlp.gate_proj.weight": "pytorch-model-00002-of-00004.safetensors",
+    "layers.29.mlp.up_proj.weight": "pytorch-model-00002-of-00004.safetensors",
+    "layers.29.post_attention_layernorm.weight": "pytorch-model-00002-of-00004.safetensors",
+    "layers.29.self_attn.k_proj.bias": "pytorch-model-00002-of-00004.safetensors",
+    "layers.29.self_attn.k_proj.weight": "pytorch-model-00002-of-00004.safetensors",
+    "layers.29.self_attn.o_proj.weight": "pytorch-model-00002-of-00004.safetensors",
+    "layers.29.self_attn.q_proj.bias": "pytorch-model-00002-of-00004.safetensors",
+    "layers.29.self_attn.q_proj.weight": "pytorch-model-00002-of-00004.safetensors",
+    "layers.29.self_attn.v_proj.bias": "pytorch-model-00002-of-00004.safetensors",
+    "layers.29.self_attn.v_proj.weight": "pytorch-model-00002-of-00004.safetensors",
+    "layers.3.input_layernorm.weight": "pytorch-model-00002-of-00004.safetensors",
+    "layers.3.mlp.down_proj.weight": "pytorch-model-00002-of-00004.safetensors",
+    "layers.3.mlp.gate_proj.weight": "pytorch-model-00002-of-00004.safetensors",
+    "layers.3.mlp.up_proj.weight": "pytorch-model-00002-of-00004.safetensors",
+    "layers.3.post_attention_layernorm.weight": "pytorch-model-00002-of-00004.safetensors",
+    "layers.3.self_attn.k_proj.bias": "pytorch-model-00002-of-00004.safetensors",
+    "layers.3.self_attn.k_proj.weight": "pytorch-model-00002-of-00004.safetensors",
+    "layers.3.self_attn.o_proj.weight": "pytorch-model-00002-of-00004.safetensors",
+    "layers.3.self_attn.q_proj.bias": "pytorch-model-00002-of-00004.safetensors",
+    "layers.3.self_attn.q_proj.weight": "pytorch-model-00002-of-00004.safetensors",
+    "layers.3.self_attn.v_proj.bias": "pytorch-model-00002-of-00004.safetensors",
+    "layers.3.self_attn.v_proj.weight": "pytorch-model-00002-of-00004.safetensors",
+    "layers.30.input_layernorm.weight": "pytorch-model-00002-of-00004.safetensors",
+    "layers.30.mlp.down_proj.weight": "pytorch-model-00002-of-00004.safetensors",
+    "layers.30.mlp.gate_proj.weight": "pytorch-model-00002-of-00004.safetensors",
+    "layers.30.mlp.up_proj.weight": "pytorch-model-00002-of-00004.safetensors",
+    "layers.30.post_attention_layernorm.weight": "pytorch-model-00002-of-00004.safetensors",
+    "layers.30.self_attn.k_proj.bias": "pytorch-model-00002-of-00004.safetensors",
+    "layers.30.self_attn.k_proj.weight": "pytorch-model-00002-of-00004.safetensors",
+    "layers.30.self_attn.o_proj.weight": "pytorch-model-00002-of-00004.safetensors",
+    "layers.30.self_attn.q_proj.bias": "pytorch-model-00002-of-00004.safetensors",
+    "layers.30.self_attn.q_proj.weight": "pytorch-model-00002-of-00004.safetensors",
+    "layers.30.self_attn.v_proj.bias": "pytorch-model-00002-of-00004.safetensors",
+    "layers.30.self_attn.v_proj.weight": "pytorch-model-00002-of-00004.safetensors",
+    "layers.31.input_layernorm.weight": "pytorch-model-00002-of-00004.safetensors",
+    "layers.31.mlp.down_proj.weight": "pytorch-model-00002-of-00004.safetensors",
+    "layers.31.mlp.gate_proj.weight": "pytorch-model-00002-of-00004.safetensors",
+    "layers.31.mlp.up_proj.weight": "pytorch-model-00002-of-00004.safetensors",
+    "layers.31.post_attention_layernorm.weight": "pytorch-model-00002-of-00004.safetensors",
+    "layers.31.self_attn.k_proj.bias": "pytorch-model-00002-of-00004.safetensors",
+    "layers.31.self_attn.k_proj.weight": "pytorch-model-00002-of-00004.safetensors",
+    "layers.31.self_attn.o_proj.weight": "pytorch-model-00002-of-00004.safetensors",
+    "layers.31.self_attn.q_proj.bias": "pytorch-model-00002-of-00004.safetensors",
+    "layers.31.self_attn.q_proj.weight": "pytorch-model-00002-of-00004.safetensors",
+    "layers.31.self_attn.v_proj.bias": "pytorch-model-00002-of-00004.safetensors",
+    "layers.31.self_attn.v_proj.weight": "pytorch-model-00002-of-00004.safetensors",
+    "layers.32.input_layernorm.weight": "pytorch-model-00002-of-00004.safetensors",
+    "layers.32.mlp.down_proj.weight": "pytorch-model-00002-of-00004.safetensors",
+    "layers.32.mlp.gate_proj.weight": "pytorch-model-00002-of-00004.safetensors",
+    "layers.32.mlp.up_proj.weight": "pytorch-model-00002-of-00004.safetensors",
+    "layers.32.post_attention_layernorm.weight": "pytorch-model-00002-of-00004.safetensors",
+    "layers.32.self_attn.k_proj.bias": "pytorch-model-00002-of-00004.safetensors",
+    "layers.32.self_attn.k_proj.weight": "pytorch-model-00002-of-00004.safetensors",
+    "layers.32.self_attn.o_proj.weight": "pytorch-model-00002-of-00004.safetensors",
+    "layers.32.self_attn.q_proj.bias": "pytorch-model-00002-of-00004.safetensors",
+    "layers.32.self_attn.q_proj.weight": "pytorch-model-00002-of-00004.safetensors",
+    "layers.32.self_attn.v_proj.bias": "pytorch-model-00002-of-00004.safetensors",
+    "layers.32.self_attn.v_proj.weight": "pytorch-model-00002-of-00004.safetensors",
+    "layers.33.input_layernorm.weight": "pytorch-model-00002-of-00004.safetensors",
+    "layers.33.mlp.down_proj.weight": "pytorch-model-00002-of-00004.safetensors",
+    "layers.33.mlp.gate_proj.weight": "pytorch-model-00002-of-00004.safetensors",
+    "layers.33.mlp.up_proj.weight": "pytorch-model-00002-of-00004.safetensors",
+    "layers.33.post_attention_layernorm.weight": "pytorch-model-00002-of-00004.safetensors",
+    "layers.33.self_attn.k_proj.bias": "pytorch-model-00002-of-00004.safetensors",
+    "layers.33.self_attn.k_proj.weight": "pytorch-model-00002-of-00004.safetensors",
+    "layers.33.self_attn.o_proj.weight": "pytorch-model-00002-of-00004.safetensors",
+    "layers.33.self_attn.q_proj.bias": "pytorch-model-00002-of-00004.safetensors",
+    "layers.33.self_attn.q_proj.weight": "pytorch-model-00002-of-00004.safetensors",
+    "layers.33.self_attn.v_proj.bias": "pytorch-model-00002-of-00004.safetensors",
+    "layers.33.self_attn.v_proj.weight": "pytorch-model-00002-of-00004.safetensors",
+    "layers.34.input_layernorm.weight": "pytorch-model-00002-of-00004.safetensors",
+    "layers.34.mlp.down_proj.weight": "pytorch-model-00002-of-00004.safetensors",
+    "layers.34.mlp.gate_proj.weight": "pytorch-model-00002-of-00004.safetensors",
+    "layers.34.mlp.up_proj.weight": "pytorch-model-00002-of-00004.safetensors",
+    "layers.34.post_attention_layernorm.weight": "pytorch-model-00002-of-00004.safetensors",
+    "layers.34.self_attn.k_proj.bias": "pytorch-model-00002-of-00004.safetensors",
+    "layers.34.self_attn.k_proj.weight": "pytorch-model-00002-of-00004.safetensors",
+    "layers.34.self_attn.o_proj.weight": "pytorch-model-00002-of-00004.safetensors",
+    "layers.34.self_attn.q_proj.bias": "pytorch-model-00002-of-00004.safetensors",
+    "layers.34.self_attn.q_proj.weight": "pytorch-model-00002-of-00004.safetensors",
+    "layers.34.self_attn.v_proj.bias": "pytorch-model-00002-of-00004.safetensors",
+    "layers.34.self_attn.v_proj.weight": "pytorch-model-00002-of-00004.safetensors",
+    "layers.35.input_layernorm.weight": "pytorch-model-00002-of-00004.safetensors",
+    "layers.35.mlp.down_proj.weight": "pytorch-model-00002-of-00004.safetensors",
+    "layers.35.mlp.gate_proj.weight": "pytorch-model-00002-of-00004.safetensors",
+    "layers.35.mlp.up_proj.weight": "pytorch-model-00002-of-00004.safetensors",
+    "layers.35.post_attention_layernorm.weight": "pytorch-model-00002-of-00004.safetensors",
+    "layers.35.self_attn.k_proj.bias": "pytorch-model-00002-of-00004.safetensors",
+    "layers.35.self_attn.k_proj.weight": "pytorch-model-00002-of-00004.safetensors",
+    "layers.35.self_attn.o_proj.weight": "pytorch-model-00002-of-00004.safetensors",
+    "layers.35.self_attn.q_proj.bias": "pytorch-model-00002-of-00004.safetensors",
+    "layers.35.self_attn.q_proj.weight": "pytorch-model-00002-of-00004.safetensors",
+    "layers.35.self_attn.v_proj.bias": "pytorch-model-00002-of-00004.safetensors",
+    "layers.35.self_attn.v_proj.weight": "pytorch-model-00002-of-00004.safetensors",
+    "layers.36.input_layernorm.weight": "pytorch-model-00002-of-00004.safetensors",
+    "layers.36.mlp.down_proj.weight": "pytorch-model-00002-of-00004.safetensors",
+    "layers.36.mlp.gate_proj.weight": "pytorch-model-00002-of-00004.safetensors",
+    "layers.36.mlp.up_proj.weight": "pytorch-model-00002-of-00004.safetensors",
+    "layers.36.post_attention_layernorm.weight": "pytorch-model-00002-of-00004.safetensors",
+    "layers.36.self_attn.k_proj.bias": "pytorch-model-00002-of-00004.safetensors",
+    "layers.36.self_attn.k_proj.weight": "pytorch-model-00002-of-00004.safetensors",
+    "layers.36.self_attn.o_proj.weight": "pytorch-model-00002-of-00004.safetensors",
+    "layers.36.self_attn.q_proj.bias": "pytorch-model-00002-of-00004.safetensors",
+    "layers.36.self_attn.q_proj.weight": "pytorch-model-00002-of-00004.safetensors",
+    "layers.36.self_attn.v_proj.bias": "pytorch-model-00002-of-00004.safetensors",
+    "layers.36.self_attn.v_proj.weight": "pytorch-model-00002-of-00004.safetensors",
+    "layers.37.input_layernorm.weight": "pytorch-model-00002-of-00004.safetensors",
+    "layers.37.mlp.down_proj.weight": "pytorch-model-00002-of-00004.safetensors",
+    "layers.37.mlp.gate_proj.weight": "pytorch-model-00002-of-00004.safetensors",
+    "layers.37.mlp.up_proj.weight": "pytorch-model-00002-of-00004.safetensors",
+    "layers.37.post_attention_layernorm.weight": "pytorch-model-00002-of-00004.safetensors",
+    "layers.37.self_attn.k_proj.bias": "pytorch-model-00002-of-00004.safetensors",
+    "layers.37.self_attn.k_proj.weight": "pytorch-model-00002-of-00004.safetensors",
+    "layers.37.self_attn.o_proj.weight": "pytorch-model-00002-of-00004.safetensors",
+    "layers.37.self_attn.q_proj.bias": "pytorch-model-00002-of-00004.safetensors",
+    "layers.37.self_attn.q_proj.weight": "pytorch-model-00002-of-00004.safetensors",
+    "layers.37.self_attn.v_proj.bias": "pytorch-model-00002-of-00004.safetensors",
+    "layers.37.self_attn.v_proj.weight": "pytorch-model-00002-of-00004.safetensors",
+    "layers.38.input_layernorm.weight": "pytorch-model-00002-of-00004.safetensors",
+    "layers.38.mlp.down_proj.weight": "pytorch-model-00002-of-00004.safetensors",
+    "layers.38.mlp.gate_proj.weight": "pytorch-model-00002-of-00004.safetensors",
+    "layers.38.mlp.up_proj.weight": "pytorch-model-00002-of-00004.safetensors",
+    "layers.38.post_attention_layernorm.weight": "pytorch-model-00002-of-00004.safetensors",
+    "layers.38.self_attn.k_proj.bias": "pytorch-model-00002-of-00004.safetensors",
+    "layers.38.self_attn.k_proj.weight": "pytorch-model-00002-of-00004.safetensors",
+    "layers.38.self_attn.o_proj.weight": "pytorch-model-00002-of-00004.safetensors",
+    "layers.38.self_attn.q_proj.bias": "pytorch-model-00002-of-00004.safetensors",
+    "layers.38.self_attn.q_proj.weight": "pytorch-model-00002-of-00004.safetensors",
+    "layers.38.self_attn.v_proj.bias": "pytorch-model-00002-of-00004.safetensors",
+    "layers.38.self_attn.v_proj.weight": "pytorch-model-00002-of-00004.safetensors",
+    "layers.39.input_layernorm.weight": "pytorch-model-00002-of-00004.safetensors",
+    "layers.39.mlp.down_proj.weight": "pytorch-model-00002-of-00004.safetensors",
+    "layers.39.mlp.gate_proj.weight": "pytorch-model-00002-of-00004.safetensors",
+    "layers.39.mlp.up_proj.weight": "pytorch-model-00002-of-00004.safetensors",
+    "layers.39.post_attention_layernorm.weight": "pytorch-model-00002-of-00004.safetensors",
+    "layers.39.self_attn.k_proj.bias": "pytorch-model-00002-of-00004.safetensors",
+    "layers.39.self_attn.k_proj.weight": "pytorch-model-00002-of-00004.safetensors",
+    "layers.39.self_attn.o_proj.weight": "pytorch-model-00002-of-00004.safetensors",
+    "layers.39.self_attn.q_proj.bias": "pytorch-model-00002-of-00004.safetensors",
+    "layers.39.self_attn.q_proj.weight": "pytorch-model-00002-of-00004.safetensors",
+    "layers.39.self_attn.v_proj.bias": "pytorch-model-00002-of-00004.safetensors",
+    "layers.39.self_attn.v_proj.weight": "pytorch-model-00002-of-00004.safetensors",
+    "layers.4.input_layernorm.weight": "pytorch-model-00002-of-00004.safetensors",
+    "layers.4.mlp.down_proj.weight": "pytorch-model-00002-of-00004.safetensors",
+    "layers.4.mlp.gate_proj.weight": "pytorch-model-00002-of-00004.safetensors",
+    "layers.4.mlp.up_proj.weight": "pytorch-model-00002-of-00004.safetensors",
+    "layers.4.post_attention_layernorm.weight": "pytorch-model-00002-of-00004.safetensors",
+    "layers.4.self_attn.k_proj.bias": "pytorch-model-00002-of-00004.safetensors",
+    "layers.4.self_attn.k_proj.weight": "pytorch-model-00002-of-00004.safetensors",
+    "layers.4.self_attn.o_proj.weight": "pytorch-model-00002-of-00004.safetensors",
+    "layers.4.self_attn.q_proj.bias": "pytorch-model-00002-of-00004.safetensors",
+    "layers.4.self_attn.q_proj.weight": "pytorch-model-00002-of-00004.safetensors",
+    "layers.4.self_attn.v_proj.bias": "pytorch-model-00002-of-00004.safetensors",
+    "layers.4.self_attn.v_proj.weight": "pytorch-model-00002-of-00004.safetensors",
+    "layers.40.input_layernorm.weight": "pytorch-model-00002-of-00004.safetensors",
+    "layers.40.mlp.down_proj.weight": "pytorch-model-00002-of-00004.safetensors",
+    "layers.40.mlp.gate_proj.weight": "pytorch-model-00002-of-00004.safetensors",
+    "layers.40.mlp.up_proj.weight": "pytorch-model-00002-of-00004.safetensors",
+    "layers.40.post_attention_layernorm.weight": "pytorch-model-00002-of-00004.safetensors",
+    "layers.40.self_attn.k_proj.bias": "pytorch-model-00002-of-00004.safetensors",
+    "layers.40.self_attn.k_proj.weight": "pytorch-model-00002-of-00004.safetensors",
+    "layers.40.self_attn.o_proj.weight": "pytorch-model-00002-of-00004.safetensors",
+    "layers.40.self_attn.q_proj.bias": "pytorch-model-00002-of-00004.safetensors",
+    "layers.40.self_attn.q_proj.weight": "pytorch-model-00002-of-00004.safetensors",
+    "layers.40.self_attn.v_proj.bias": "pytorch-model-00002-of-00004.safetensors",
+    "layers.40.self_attn.v_proj.weight": "pytorch-model-00002-of-00004.safetensors",
+    "layers.41.input_layernorm.weight": "pytorch-model-00002-of-00004.safetensors",
+    "layers.41.mlp.down_proj.weight": "pytorch-model-00003-of-00004.safetensors",
+    "layers.41.mlp.gate_proj.weight": "pytorch-model-00003-of-00004.safetensors",
+    "layers.41.mlp.up_proj.weight": "pytorch-model-00003-of-00004.safetensors",
+    "layers.41.post_attention_layernorm.weight": "pytorch-model-00003-of-00004.safetensors",
+    "layers.41.self_attn.k_proj.bias": "pytorch-model-00003-of-00004.safetensors",
+    "layers.41.self_attn.k_proj.weight": "pytorch-model-00003-of-00004.safetensors",
+    "layers.41.self_attn.o_proj.weight": "pytorch-model-00002-of-00004.safetensors",
+    "layers.41.self_attn.q_proj.bias": "pytorch-model-00003-of-00004.safetensors",
+    "layers.41.self_attn.q_proj.weight": "pytorch-model-00003-of-00004.safetensors",
+    "layers.41.self_attn.v_proj.bias": "pytorch-model-00003-of-00004.safetensors",
+    "layers.41.self_attn.v_proj.weight": "pytorch-model-00003-of-00004.safetensors",
+    "layers.42.input_layernorm.weight": "pytorch-model-00003-of-00004.safetensors",
+    "layers.42.mlp.down_proj.weight": "pytorch-model-00003-of-00004.safetensors",
+    "layers.42.mlp.gate_proj.weight": "pytorch-model-00003-of-00004.safetensors",
+    "layers.42.mlp.up_proj.weight": "pytorch-model-00003-of-00004.safetensors",
+    "layers.42.post_attention_layernorm.weight": "pytorch-model-00003-of-00004.safetensors",
+    "layers.42.self_attn.k_proj.bias": "pytorch-model-00003-of-00004.safetensors",
+    "layers.42.self_attn.k_proj.weight": "pytorch-model-00003-of-00004.safetensors",
+    "layers.42.self_attn.o_proj.weight": "pytorch-model-00003-of-00004.safetensors",
+    "layers.42.self_attn.q_proj.bias": "pytorch-model-00003-of-00004.safetensors",
+    "layers.42.self_attn.q_proj.weight": "pytorch-model-00003-of-00004.safetensors",
+    "layers.42.self_attn.v_proj.bias": "pytorch-model-00003-of-00004.safetensors",
+    "layers.42.self_attn.v_proj.weight": "pytorch-model-00003-of-00004.safetensors",
+    "layers.43.input_layernorm.weight": "pytorch-model-00003-of-00004.safetensors",
+    "layers.43.mlp.down_proj.weight": "pytorch-model-00003-of-00004.safetensors",
+    "layers.43.mlp.gate_proj.weight": "pytorch-model-00003-of-00004.safetensors",
+    "layers.43.mlp.up_proj.weight": "pytorch-model-00003-of-00004.safetensors",
+    "layers.43.post_attention_layernorm.weight": "pytorch-model-00003-of-00004.safetensors",
+    "layers.43.self_attn.k_proj.bias": "pytorch-model-00003-of-00004.safetensors",
+    "layers.43.self_attn.k_proj.weight": "pytorch-model-00003-of-00004.safetensors",
+    "layers.43.self_attn.o_proj.weight": "pytorch-model-00003-of-00004.safetensors",
+    "layers.43.self_attn.q_proj.bias": "pytorch-model-00003-of-00004.safetensors",
+    "layers.43.self_attn.q_proj.weight": "pytorch-model-00003-of-00004.safetensors",
+    "layers.43.self_attn.v_proj.bias": "pytorch-model-00003-of-00004.safetensors",
+    "layers.43.self_attn.v_proj.weight": "pytorch-model-00003-of-00004.safetensors",
+    "layers.44.input_layernorm.weight": "pytorch-model-00003-of-00004.safetensors",
+    "layers.44.mlp.down_proj.weight": "pytorch-model-00003-of-00004.safetensors",
+    "layers.44.mlp.gate_proj.weight": "pytorch-model-00003-of-00004.safetensors",
+    "layers.44.mlp.up_proj.weight": "pytorch-model-00003-of-00004.safetensors",
+    "layers.44.post_attention_layernorm.weight": "pytorch-model-00003-of-00004.safetensors",
+    "layers.44.self_attn.k_proj.bias": "pytorch-model-00003-of-00004.safetensors",
+    "layers.44.self_attn.k_proj.weight": "pytorch-model-00003-of-00004.safetensors",
+    "layers.44.self_attn.o_proj.weight": "pytorch-model-00003-of-00004.safetensors",
+    "layers.44.self_attn.q_proj.bias": "pytorch-model-00003-of-00004.safetensors",
+    "layers.44.self_attn.q_proj.weight": "pytorch-model-00003-of-00004.safetensors",
+    "layers.44.self_attn.v_proj.bias": "pytorch-model-00003-of-00004.safetensors",
+    "layers.44.self_attn.v_proj.weight": "pytorch-model-00003-of-00004.safetensors",
+    "layers.45.input_layernorm.weight": "pytorch-model-00003-of-00004.safetensors",
+    "layers.45.mlp.down_proj.weight": "pytorch-model-00003-of-00004.safetensors",
+    "layers.45.mlp.gate_proj.weight": "pytorch-model-00003-of-00004.safetensors",
+    "layers.45.mlp.up_proj.weight": "pytorch-model-00003-of-00004.safetensors",
+    "layers.45.post_attention_layernorm.weight": "pytorch-model-00003-of-00004.safetensors",
+    "layers.45.self_attn.k_proj.bias": "pytorch-model-00003-of-00004.safetensors",
+    "layers.45.self_attn.k_proj.weight": "pytorch-model-00003-of-00004.safetensors",
+    "layers.45.self_attn.o_proj.weight": "pytorch-model-00003-of-00004.safetensors",
+    "layers.45.self_attn.q_proj.bias": "pytorch-model-00003-of-00004.safetensors",
+    "layers.45.self_attn.q_proj.weight": "pytorch-model-00003-of-00004.safetensors",
+    "layers.45.self_attn.v_proj.bias": "pytorch-model-00003-of-00004.safetensors",
+    "layers.45.self_attn.v_proj.weight": "pytorch-model-00003-of-00004.safetensors",
+    "layers.46.input_layernorm.weight": "pytorch-model-00003-of-00004.safetensors",
+    "layers.46.mlp.down_proj.weight": "pytorch-model-00003-of-00004.safetensors",
+    "layers.46.mlp.gate_proj.weight": "pytorch-model-00003-of-00004.safetensors",
+    "layers.46.mlp.up_proj.weight": "pytorch-model-00003-of-00004.safetensors",
+    "layers.46.post_attention_layernorm.weight": "pytorch-model-00003-of-00004.safetensors",
+    "layers.46.self_attn.k_proj.bias": "pytorch-model-00003-of-00004.safetensors",
+    "layers.46.self_attn.k_proj.weight": "pytorch-model-00003-of-00004.safetensors",
+    "layers.46.self_attn.o_proj.weight": "pytorch-model-00003-of-00004.safetensors",
+    "layers.46.self_attn.q_proj.bias": "pytorch-model-00003-of-00004.safetensors",
+    "layers.46.self_attn.q_proj.weight": "pytorch-model-00003-of-00004.safetensors",
+    "layers.46.self_attn.v_proj.bias": "pytorch-model-00003-of-00004.safetensors",
+    "layers.46.self_attn.v_proj.weight": "pytorch-model-00003-of-00004.safetensors",
+    "layers.47.input_layernorm.weight": "pytorch-model-00003-of-00004.safetensors",
+    "layers.47.mlp.down_proj.weight": "pytorch-model-00003-of-00004.safetensors",
+    "layers.47.mlp.gate_proj.weight": "pytorch-model-00003-of-00004.safetensors",
+    "layers.47.mlp.up_proj.weight": "pytorch-model-00003-of-00004.safetensors",
+    "layers.47.post_attention_layernorm.weight": "pytorch-model-00003-of-00004.safetensors",
+    "layers.47.self_attn.k_proj.bias": "pytorch-model-00003-of-00004.safetensors",
+    "layers.47.self_attn.k_proj.weight": "pytorch-model-00003-of-00004.safetensors",
+    "layers.47.self_attn.o_proj.weight": "pytorch-model-00003-of-00004.safetensors",
+    "layers.47.self_attn.q_proj.bias": "pytorch-model-00003-of-00004.safetensors",
+    "layers.47.self_attn.q_proj.weight": "pytorch-model-00003-of-00004.safetensors",
+    "layers.47.self_attn.v_proj.bias": "pytorch-model-00003-of-00004.safetensors",
+    "layers.47.self_attn.v_proj.weight": "pytorch-model-00003-of-00004.safetensors",
+    "layers.5.input_layernorm.weight": "pytorch-model-00003-of-00004.safetensors",
+    "layers.5.mlp.down_proj.weight": "pytorch-model-00003-of-00004.safetensors",
+    "layers.5.mlp.gate_proj.weight": "pytorch-model-00003-of-00004.safetensors",
+    "layers.5.mlp.up_proj.weight": "pytorch-model-00003-of-00004.safetensors",
+    "layers.5.post_attention_layernorm.weight": "pytorch-model-00003-of-00004.safetensors",
+    "layers.5.self_attn.k_proj.bias": "pytorch-model-00003-of-00004.safetensors",
+    "layers.5.self_attn.k_proj.weight": "pytorch-model-00003-of-00004.safetensors",
+    "layers.5.self_attn.o_proj.weight": "pytorch-model-00003-of-00004.safetensors",
+    "layers.5.self_attn.q_proj.bias": "pytorch-model-00003-of-00004.safetensors",
+    "layers.5.self_attn.q_proj.weight": "pytorch-model-00003-of-00004.safetensors",
+    "layers.5.self_attn.v_proj.bias": "pytorch-model-00003-of-00004.safetensors",
+    "layers.5.self_attn.v_proj.weight": "pytorch-model-00003-of-00004.safetensors",
+    "layers.6.input_layernorm.weight": "pytorch-model-00003-of-00004.safetensors",
+    "layers.6.mlp.down_proj.weight": "pytorch-model-00003-of-00004.safetensors",
+    "layers.6.mlp.gate_proj.weight": "pytorch-model-00003-of-00004.safetensors",
+    "layers.6.mlp.up_proj.weight": "pytorch-model-00003-of-00004.safetensors",
+    "layers.6.post_attention_layernorm.weight": "pytorch-model-00003-of-00004.safetensors",
+    "layers.6.self_attn.k_proj.bias": "pytorch-model-00003-of-00004.safetensors",
+    "layers.6.self_attn.k_proj.weight": "pytorch-model-00003-of-00004.safetensors",
+    "layers.6.self_attn.o_proj.weight": "pytorch-model-00003-of-00004.safetensors",
+    "layers.6.self_attn.q_proj.bias": "pytorch-model-00003-of-00004.safetensors",
+    "layers.6.self_attn.q_proj.weight": "pytorch-model-00003-of-00004.safetensors",
+    "layers.6.self_attn.v_proj.bias": "pytorch-model-00003-of-00004.safetensors",
+    "layers.6.self_attn.v_proj.weight": "pytorch-model-00003-of-00004.safetensors",
+    "layers.7.input_layernorm.weight": "pytorch-model-00003-of-00004.safetensors",
+    "layers.7.mlp.down_proj.weight": "pytorch-model-00003-of-00004.safetensors",
+    "layers.7.mlp.gate_proj.weight": "pytorch-model-00003-of-00004.safetensors",
+    "layers.7.mlp.up_proj.weight": "pytorch-model-00003-of-00004.safetensors",
+    "layers.7.post_attention_layernorm.weight": "pytorch-model-00003-of-00004.safetensors",
+    "layers.7.self_attn.k_proj.bias": "pytorch-model-00003-of-00004.safetensors",
+    "layers.7.self_attn.k_proj.weight": "pytorch-model-00003-of-00004.safetensors",
+    "layers.7.self_attn.o_proj.weight": "pytorch-model-00003-of-00004.safetensors",
+    "layers.7.self_attn.q_proj.bias": "pytorch-model-00003-of-00004.safetensors",
+    "layers.7.self_attn.q_proj.weight": "pytorch-model-00003-of-00004.safetensors",
+    "layers.7.self_attn.v_proj.bias": "pytorch-model-00003-of-00004.safetensors",
+    "layers.7.self_attn.v_proj.weight": "pytorch-model-00003-of-00004.safetensors",
+    "layers.8.input_layernorm.weight": "pytorch-model-00003-of-00004.safetensors",
+    "layers.8.mlp.down_proj.weight": "pytorch-model-00003-of-00004.safetensors",
+    "layers.8.mlp.gate_proj.weight": "pytorch-model-00003-of-00004.safetensors",
+    "layers.8.mlp.up_proj.weight": "pytorch-model-00003-of-00004.safetensors",
+    "layers.8.post_attention_layernorm.weight": "pytorch-model-00003-of-00004.safetensors",
+    "layers.8.self_attn.k_proj.bias": "pytorch-model-00003-of-00004.safetensors",
+    "layers.8.self_attn.k_proj.weight": "pytorch-model-00003-of-00004.safetensors",
+    "layers.8.self_attn.o_proj.weight": "pytorch-model-00003-of-00004.safetensors",
+    "layers.8.self_attn.q_proj.bias": "pytorch-model-00003-of-00004.safetensors",
+    "layers.8.self_attn.q_proj.weight": "pytorch-model-00003-of-00004.safetensors",
+    "layers.8.self_attn.v_proj.bias": "pytorch-model-00003-of-00004.safetensors",
+    "layers.8.self_attn.v_proj.weight": "pytorch-model-00003-of-00004.safetensors",
+    "layers.9.input_layernorm.weight": "pytorch-model-00003-of-00004.safetensors",
+    "layers.9.mlp.down_proj.weight": "pytorch-model-00003-of-00004.safetensors",
+    "layers.9.mlp.gate_proj.weight": "pytorch-model-00003-of-00004.safetensors",
+    "layers.9.mlp.up_proj.weight": "pytorch-model-00003-of-00004.safetensors",
+    "layers.9.post_attention_layernorm.weight": "pytorch-model-00003-of-00004.safetensors",
+    "layers.9.self_attn.k_proj.bias": "pytorch-model-00003-of-00004.safetensors",
+    "layers.9.self_attn.k_proj.weight": "pytorch-model-00003-of-00004.safetensors",
+    "layers.9.self_attn.o_proj.weight": "pytorch-model-00003-of-00004.safetensors",
+    "layers.9.self_attn.q_proj.bias": "pytorch-model-00003-of-00004.safetensors",
+    "layers.9.self_attn.q_proj.weight": "pytorch-model-00003-of-00004.safetensors",
+    "layers.9.self_attn.v_proj.bias": "pytorch-model-00003-of-00004.safetensors",
+    "layers.9.self_attn.v_proj.weight": "pytorch-model-00003-of-00004.safetensors",
+    "lm_head.weight": "pytorch-model-00003-of-00004.safetensors",
+    "norm.weight": "pytorch-model-00003-of-00004.safetensors"
+  }
+}

models/gen_pipeline.py

@@ -0,0 +1,631 @@
+import re
+import os
+import json
+import inspect
+import random
+from typing import Literal
+from dataclasses import dataclass, field
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import torchvision.transforms as transforms
+import numpy as np
+from PIL import Image
+from tqdm.auto import tqdm
+from loguru import logger
+
+from transformers import AutoTokenizer
+from transformers.cache_utils import Cache
+from diffusers.models.autoencoders.autoencoder_kl import AutoencoderKL as DiffusersAutoencoderKL
+from diffusers.models.autoencoders.vae import DiagonalGaussianDistribution
+from diffusers.models.modeling_outputs import AutoencoderKLOutput
+
+try:
+    from .nextstep_model import NextStep
+except ImportError:
+    from nextstep_model import NextStep
+
+DEFAULT_IMAGE_AREA_TOKEN = "<|image_area|>"
+@dataclass
+class AutoEncoderParams:
+    resolution: int = 256
+    in_channels: int = 3
+    ch: int = 128
+    out_ch: int = 3
+    ch_mult: list[int] = field(default_factory=lambda: [1, 2, 4, 4])
+    num_res_blocks: int = 2
+    z_channels: int = 16
+    scaling_factor: float = 1
+    shift_factor: float = 0
+    deterministic: bool = True
+    encoder_norm: bool = True
+    psz: int | None = 1
+
+
+class AutoencoderKL(DiffusersAutoencoderKL):
+    """
+    继承自 diffusers 的 AutoencoderKL，添加了 patchify 和 encoder_norm 功能。
+    这样可以复用 diffusers 的完整实现，只扩展必要的功能。
+    """
+
+    def __init__(self, params: AutoEncoderParams):
+        """
+        从 AutoEncoderParams 初始化模型，转换为 diffusers 格式
+        """
+        # 转换参数格式为 diffusers 的 AutoencoderKL 格式
+        down_block_types = ["DownEncoderBlock2D"] * len(params.ch_mult)
+        up_block_types = ["UpDecoderBlock2D"] * len(params.ch_mult)
+        block_out_channels = [params.ch * m for m in params.ch_mult]
+
+        # 调用父类初始化，创建 diffusers 的 encoder 和 decoder
+        super().__init__(
+            in_channels=params.in_channels,
+            out_channels=params.out_ch,
+            down_block_types=tuple(down_block_types),
+            up_block_types=tuple(up_block_types),
+            block_out_channels=tuple(block_out_channels),
+            layers_per_block=params.num_res_blocks,
+            latent_channels=params.z_channels,
+            norm_num_groups=32,
+            sample_size=params.resolution,
+            scaling_factor=params.scaling_factor,
+            shift_factor=params.shift_factor,
+            act_fn="silu",
+            mid_block_add_attention=True,
+            use_quant_conv=False,  # 旧的 VAE 没有使用 quant_conv
+            use_post_quant_conv=False,
+        )
+
+        # 保存自定义参数
+        self.params = params
+        self.encoder_norm = params.encoder_norm
+        self.psz = params.psz
+        self.deterministic = params.deterministic
+
+    def layer_norm_2d(self, input: torch.Tensor, normalized_shape: torch.Size, eps: float = 1e-6) -> torch.Tensor:
+        """Layer normalization for 2D spatial features."""
+        # input.shape = (bsz, c, h, w)
+        return F.layer_norm(
+            input.permute(0, 2, 3, 1), normalized_shape, None, None, eps
+        ).permute(0, 3, 1, 2)
+
+    def patchify(self, img: torch.Tensor):
+        """
+        img: (bsz, C, H, W)
+        x: (bsz, patch_size**2 * C, H / patch_size, W / patch_size)
+        """
+        bsz, c, h, w = img.shape
+        p = self.psz
+        h_, w_ = h // p, w // p
+
+        img = img.reshape(bsz, c, h_, p, w_, p)
+        img = torch.einsum("nchpwq->ncpqhw", img)
+        x = img.reshape(bsz, c * p**2, h_, w_)
+        return x
+
+    def unpatchify(self, x: torch.Tensor):
+        """
+        x: (bsz, patch_size**2 * C, H / patch_size, W / patch_size)
+        img: (bsz, C, H, W)
+        """
+        bsz = x.shape[0]
+        p = self.psz
+        c = self.config.latent_channels
+        h_, w_ = x.shape[2], x.shape[3]
+
+        x = x.reshape(bsz, c, p, p, h_, w_)
+        x = torch.einsum("ncpqhw->nchpwq", x)
+        img = x.reshape(bsz, c, h_ * p, w_ * p)
+        return img
+
+    def encode(self, x: torch.Tensor, return_dict: bool = True):
+        """
+        重写 encode 方法以支持 patchify 和 encoder_norm
+        """
+        # 使用父类的 encoder
+        h = self.encoder(x)
+
+        # 使用父类的 quant_conv（如果存在）
+        moments = self.quant_conv(h) if self.config.use_quant_conv else h
+
+        # 应用 patchify 和 normalization（如果启用）
+        mean, logvar = torch.chunk(moments, 2, dim=1)
+
+        if self.psz is not None:
+            mean = self.patchify(mean)
+            if self.encoder_norm:
+                mean = self.layer_norm_2d(mean, (mean.size(1),))
+            mean = self.unpatchify(mean)
+
+        moments = torch.cat([mean, logvar], dim=1).contiguous()
+        posterior = DiagonalGaussianDistribution(moments, deterministic=self.deterministic)
+
+        return (posterior,) if not return_dict else AutoencoderKLOutput(latent_dist=posterior)
+
+    @staticmethod
+    def _convert_old_keys_to_diffusers(state_dict, num_resolutions=4):
+        """
+        将旧的自定义 VAE 键名转换为 diffusers AutoencoderKL 的键名
+        旧格式: encoder.down.X.block.Y -> 新格式: encoder.down_blocks.X.resnets.Y
+
+        注意：decoder 的 up_blocks 顺序需要反转，因为旧实现使用 insert(0, ...)
+        """
+        import re
+        new_state_dict = {}
+
+        # 定义替换规则
+        replacements = [
+            # 通用替换
+            (".nin_shortcut.", ".conv_shortcut."),
+            (".norm_out.", ".conv_norm_out."),
+            # Encoder down blocks
+            (".down.", ".down_blocks."),
+            (".block.", ".resnets."),
+            (".downsample.", ".downsamplers.0."),
+            # Encoder mid blocks
+            (".mid.block_1.", ".mid_block.resnets.0."),
+            (".mid.block_2.", ".mid_block.resnets.1."),
+            (".mid.attn_1.norm.", ".mid_block.attentions.0.group_norm."),
+            (".mid.attn_1.q.", ".mid_block.attentions.0.to_q."),
+            (".mid.attn_1.k.", ".mid_block.attentions.0.to_k."),
+            (".mid.attn_1.v.", ".mid_block.attentions.0.to_v."),
+            (".mid.attn_1.proj_out.", ".mid_block.attentions.0.to_out.0."),
+            # Decoder up blocks
+            (".upsample.", ".upsamplers.0."),
+        ]
+
+        for key, value in state_dict.items():
+            new_key = key
+
+            # 跳过不需要转换的键
+            if any(skip in key for skip in ["conv_in", "conv_out"]) and ("encoder." in key or "decoder." in key):
+                new_state_dict[new_key] = value
+                continue
+
+            # Encoder 转换
+            if key.startswith("encoder."):
+                for old_pattern, new_pattern in replacements:
+                    new_key = new_key.replace(old_pattern, new_pattern)
+
+            # Decoder 转换
+            elif key.startswith("decoder."):
+                # 处理 up blocks 的索引反转
+                if ".up." in key:
+                    match = re.search(r'\.up\.(\d+)\.', key)
+                    if match:
+                        old_idx = int(match.group(1))
+                        new_idx = num_resolutions - 1 - old_idx
+                        new_key = re.sub(r'\.up\.\d+\.', f'.up_blocks.{new_idx}.', key)
+                    else:
+                        new_key = new_key.replace(".up.", ".up_blocks.")
+
+                # 应用其他替换规则
+                for old_pattern, new_pattern in replacements:
+                    new_key = new_key.replace(old_pattern, new_pattern)
+
+            # 处理 Conv2d (1x1) -> Linear 的权重形状转换
+            if "attentions" in new_key and "weight" in new_key and len(value.shape) == 4:
+                if value.shape[2] == 1 and value.shape[3] == 1:
+                    value = value.squeeze(-1).squeeze(-1)
+
+            new_state_dict[new_key] = value
+
+        return new_state_dict
+
+    @classmethod
+    def from_pretrained(cls, model_path, **kwargs):
+        """
+        从本地路径加载模型（兼容旧格式）
+        """
+        config_path = os.path.join(model_path, "config.json")
+        ckpt_path = os.path.join(model_path, "checkpoint.pt")
+
+        if not os.path.isdir(model_path) or not os.path.isfile(ckpt_path):
+            raise ValueError(f"Invalid model path: {model_path}. Missing config.json or checkpoint.pt files.")
+
+        state_dict = torch.load(ckpt_path, map_location="cpu", weights_only=True)
+
+        # 加载配置
+        config = {}
+        if os.path.isfile(config_path):
+            with open(config_path, "r") as f:
+                config = json.load(f)
+        config.update(kwargs)
+
+        # 过滤出 AutoEncoderParams 中的参数
+        param_signature = inspect.signature(AutoEncoderParams.__init__).parameters
+        valid_kwargs = {k: v for k, v in config.items() if k in param_signature}
+
+        # 记录被忽略的参数
+        ignored_params = [k for k in config.keys() if k not in param_signature]
+        if ignored_params:
+            logger.debug(f"Ignoring parameters: {ignored_params}")
+
+        params = AutoEncoderParams(**valid_kwargs)
+        model = cls(params)
+
+        # 转换旧格式的键名到 diffusers 格式
+        logger.info("Converting old VAE keys to diffusers format...")
+        state_dict = cls._convert_old_keys_to_diffusers(state_dict, num_resolutions=len(params.ch_mult))
+
+        try:
+            msg = model.load_state_dict(state_dict, strict=False)
+            logger.info(f"Loaded state_dict from {ckpt_path}")
+            if msg.missing_keys:
+                logger.warning(f"Missing keys: {msg.missing_keys}")
+            # if msg.unexpected_keys:
+            #     logger.warning(f"Unexpected keys: {msg.unexpected_keys}")
+        except Exception as e:
+            logger.error(f"Failed to load state_dict: {e}, using random initialization")
+
+        return model
+
+
+class NextStepPipeline:
+    def __init__(
+        self,
+        model_name_or_path: str | None = None,
+        vae_name_or_path: str | None = None,
+        tokenizer: AutoTokenizer | None = None,
+        model: nn.Module | None = None,
+        vae: AutoencoderKL | None = None,
+        enable_gradient_checkpointing: bool = False,
+        attn_implementation: str | None = None,  # "sdpa", "flash_attention_2", "eager"
+        device: str | None = "cuda",
+        dtype: torch.dtype | None = torch.bfloat16,
+    ):
+        self.tokenizer = AutoTokenizer.from_pretrained(
+            model_name_or_path,
+            local_files_only=True,
+            padding_side="left",
+            use_fast=True,
+        )
+        self.model: NextStep = NextStep.from_pretrained(model_name_or_path, local_files_only=True, enable_gradient_checkpointing=enable_gradient_checkpointing)
+        self.model.to(device=device,dtype=dtype)
+
+        self.tokenizer.add_eos_token = False
+        if vae_name_or_path is None:
+            vae_name_or_path = getattr(self.model.config, "vae_name_or_path", None)
+            vae_name_or_path = os.path.join(model_name_or_path, vae_name_or_path)
+        self.vae = AutoencoderKL.from_pretrained(vae_name_or_path).to(device=device,dtype=dtype)
+
+        vae_factor = 2 ** (len(self.vae.config.block_out_channels) - 1)
+        self.down_factor = vae_factor * self.model.config.latent_patch_size
+        self.shift_factor = getattr(self.vae.config, "shift_factor", 0.0)
+        self.scaling_factor = getattr(self.vae.config, "scaling_factor", 1.0)
+
+        self.boi = self.model.config.boi
+        self.eoi = self.model.config.eoi
+
+        self.image_placeholder_id = self.model.config.image_placeholder_id
+        self.pil2tensor = transforms.Compose(
+            [
+                transforms.ToTensor(),
+                transforms.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5]),
+            ]
+        )
+        self.device = self.model.device
+        self.dtype = self.model.dtype
+
+    def set_seed(self, seed: int, rank: int = 0):
+        random.seed(seed + rank)
+        np.random.seed(seed + rank)
+        torch.manual_seed(seed + rank)
+        torch.cuda.manual_seed_all(seed + rank)
+        torch.backends.cudnn.deterministic = True
+        os.environ["PYTHONHASHSEED"] = str(seed + rank)
+
+    def to_pil(self, image: torch.Tensor) -> Image.Image:
+        """Convert PyTorch tensor to PIL image"""
+        if not isinstance(image, torch.Tensor):
+            raise TypeError(f"Expected torch.Tensor, got {type(image)}")
+
+        # 标准化到[0,1]并转换为uint8
+        image = (image / 2 + 0.5).clamp(0, 1).mul(255).round().to(torch.uint8)
+        image = image.cpu().permute(1, 2, 0).numpy()
+
+        # 处理单通道图像
+        if image.shape[-1] == 1:
+            image = image[:, :, 0]
+            mode = "L"
+        else:
+            mode = "RGB"
+
+        return Image.fromarray(image, mode=mode)
+
+    def to(self, device: str | None = None, dtype: torch.dtype | None = None):
+        if device is not None:
+            self.device = device
+        if dtype is not None:
+            self.dtype = dtype
+        self.model.to(self.device, dtype=self.dtype)
+        self.vae.to(self.device, dtype=self.dtype)
+        return self
+
+    def hw2str(self, h: int, w: int) -> str:
+        return f"{h}*{w}"
+
+    def _image_str(self, hw: tuple[int, int] = (256, 256)):
+        latent_hw = (hw[0] // self.down_factor, hw[1] // self.down_factor)
+        image_ids = [self.boi] + [self.image_placeholder_id] * (latent_hw[0] * latent_hw[1]) + [self.eoi]
+        image_str = DEFAULT_IMAGE_AREA_TOKEN + self.hw2str(*latent_hw) + self.tokenizer.decode(image_ids)
+        return image_str
+
+    def _check_input(
+        self, captions: str | list[str], images: Image.Image | list[Image.Image] | None
+    ) -> tuple[list[str], list[Image.Image] | None]:
+        if not isinstance(captions, list):
+            captions = [captions]
+        if images is not None:
+            if not isinstance(images, list):
+                images = [images]
+            # 验证图像数量与<image>标记匹配
+            image_token_count = sum(len(re.findall(r"<image>", caption)) for caption in captions)
+            if image_token_count != len(images):
+                raise ValueError(f"图像数量({len(images)})与图像标记数量({image_token_count})不匹配")
+            # 替换<image>标记为图像字符串
+            hws = [(img.size[1], img.size[0]) for img in images]
+            processed_captions = []
+            image_idx = 0
+
+            for caption in captions:
+                processed_caption = caption
+                while "<image>" in processed_caption:
+                    processed_caption = processed_caption.replace("<image>", self._image_str(hws[image_idx]), 1)
+                    image_idx += 1
+                processed_captions.append(processed_caption)
+            captions = processed_captions
+        return captions, images
+
+    def _build_captions(
+        self,
+        captions: str | list[str],
+        images: list[Image.Image] | None = None,
+        num_images_per_caption: int = 1,
+        positive_prompt: str | None = None,
+        negative_prompt: str | None = None,
+        cfg: float = 1.0,
+        cfg_img: float = 1.0,
+    ):
+        # 标准化输入
+        if not isinstance(captions, list):
+            captions = [captions]
+
+        # 重复captions和images
+        captions = [caption for caption in captions for _ in range(num_images_per_caption)]
+        if images is not None:
+            images = [img for img in images for _ in range(num_images_per_caption)]
+
+        # 添加positive prompt
+        if positive_prompt:
+            captions = [f"{caption} {positive_prompt}" for caption in captions]
+
+        # 设置negative prompt默认值
+        negative_prompt = negative_prompt or ""
+
+        num_samples = len(captions)
+
+        # 简化CFG逻辑
+        if cfg != 1.0:
+            if cfg_img != 1.0:  # 使用图像和文本CFG
+                w, h = images[0].size
+                captions = captions + [self._image_str((h, w)) + negative_prompt] * num_samples
+                images = images + images
+            captions = captions + [negative_prompt] * num_samples
+
+        return captions, images
+
+    def _add_prefix_ids(self, hw: tuple[int, int], input_ids: torch.Tensor, attention_mask: torch.Tensor):
+        """添加图像区域前缀ID"""
+        prefix_str = DEFAULT_IMAGE_AREA_TOKEN + self.hw2str(hw[0] // self.down_factor, hw[1] // self.down_factor)
+        prefix_output = self.tokenizer(prefix_str, truncation=False, add_special_tokens=True, return_tensors="pt")
+
+        prefix_input_ids = prefix_output.input_ids.to(input_ids.device, dtype=input_ids.dtype)
+        prefix_attention_mask = prefix_output.attention_mask.to(attention_mask.device, dtype=attention_mask.dtype)
+
+        # 移除BOS token并添加BOI token
+        if self.tokenizer.bos_token is not None:
+            prefix_input_ids = prefix_input_ids[:, 1:]
+            prefix_attention_mask = prefix_attention_mask[:, 1:]
+
+        boi_token = prefix_input_ids.new_tensor([self.model.config.boi]).unsqueeze(0)
+        prefix_input_ids = torch.cat([prefix_input_ids, boi_token], dim=1)
+        prefix_attention_mask = torch.cat([prefix_attention_mask, prefix_attention_mask.new_ones((1, 1))], dim=1)
+
+        # 扩展到batch维度并拼接
+        bsz = input_ids.shape[0]
+        input_ids = torch.cat([input_ids, prefix_input_ids.expand(bsz, -1)], dim=1)
+        attention_mask = torch.cat([attention_mask, prefix_attention_mask.expand(bsz, -1)], dim=1)
+
+        return input_ids, attention_mask
+
+
+    def layer_norm(self, input: torch.Tensor, normalized_shape: torch.Size, eps: float = 1e-6) -> torch.Tensor:
+        """简化的layer norm实现，使用PyTorch内置函数"""
+        return F.layer_norm(input, normalized_shape, eps=eps)
+
+    @torch.no_grad()
+    def decoding(
+        self,
+        c: torch.Tensor,
+        attention_mask: torch.Tensor,
+        past_key_values: Cache,
+        max_new_len: int,
+        num_images_per_caption: int,
+        noise: torch.Tensor = None,
+        use_norm: bool = False,
+        cfg: float = 1.0,
+        cfg_img: float = 1.0,
+        cfg_schedule: Literal["linear", "constant"] = "constant",
+        timesteps_shift: float = 1.0,
+        num_sampling_steps: int = 20,
+        progress: bool = True,
+        hw: tuple[int, int] = (256, 256),
+        step: int = 0,
+        sde_solver: bool = False,
+        sde_type: str = "sde",
+    ):
+        indices = list(range(max_new_len))
+        indices = tqdm(indices, unit="tokens") if progress else indices
+        tokens = None
+        for step in indices:
+            cur_noise = None
+            if noise is not None:
+                cur_noise = noise[:,step:step+1,:][0]
+
+            # 简化CFG调度逻辑
+            if cfg_schedule == "linear":
+                tokens_len = 0 if tokens is None else tokens.shape[1]
+                cfg_iter = max(cfg / 2, 1 + (cfg - 1) * tokens_len / max_new_len)
+                cfg_img_iter = max(cfg_img / 2, 1 + (cfg_img - 1) * tokens_len / max_new_len)
+            else:  # constant or other
+                cfg_iter = cfg
+                cfg_img_iter = cfg_img
+
+            c = self.model.image_out_projector(c)
+            token_sampled = self.model.image_head.sample(
+                c=c.squeeze(1),
+                noise=cur_noise,
+                cfg=cfg_iter,
+                cfg_img=cfg_img_iter,
+                timesteps_shift=timesteps_shift,
+                num_sampling_steps=num_sampling_steps,
+                noise_repeat=num_images_per_caption,
+                sde_solver=sde_solver,
+                sde_type=sde_type,
+            )
+
+            if use_norm:
+                token_sampled = self.layer_norm(token_sampled, normalized_shape=token_sampled.size()[1:])
+            if tokens is not None:
+                tokens = torch.cat([tokens, token_sampled.unsqueeze(1)], dim=1)
+            else:
+                tokens = token_sampled.unsqueeze(1)
+
+            cur_inputs_embeds = self.model.image_in_projector(tokens[:, -1:])
+            # 简化CFG处理逻辑
+            if cfg != 1.0:
+                repeat_count = 3 if cfg_img != 1.0 else 2
+                cur_inputs_embeds = torch.cat([cur_inputs_embeds] * repeat_count, dim=0)
+
+            attention_mask = torch.cat([attention_mask, attention_mask.new_ones((attention_mask.shape[0], 1))], dim=-1)
+            outputs = self.model.forward(
+                inputs_embeds=cur_inputs_embeds,
+                attention_mask=attention_mask,
+                past_key_values=past_key_values,
+                use_cache=True,
+            )
+            past_key_values = outputs.past_key_values
+            c = outputs.last_hidden_state[:, -1:]
+
+
+        return tokens
+
+    @torch.no_grad()
+    def generate_image(
+        self,
+        captions: str | list[str],
+        images: list[Image.Image] | None = None,
+        num_images_per_caption: int = 1,
+        positive_prompt: str | None = None,
+        negative_prompt: str | None = None,
+        hw: tuple[int, int] = (256, 256),
+        use_norm: bool = False,
+        cfg: float = 1.0,
+        cfg_img: float = 1.0,
+        cfg_schedule: Literal["linear", "constant"] = "constant",
+        num_sampling_steps: int = 20,
+        timesteps_shift: float = 1.0,
+        seed: int = 42,
+        progress: bool = True,
+        sde_type: str = "sde",
+    ) -> list[Image.Image]:
+        # 0. set seed
+
+        # 1. check input
+        captions, images = self._check_input(captions, images)
+
+        # 2. build captions
+        captions, images = self._build_captions(
+            captions, images, num_images_per_caption, positive_prompt, negative_prompt, cfg, cfg_img
+        )
+
+        # 3. encode images
+        latents = None
+        if images is not None:
+            pixel_values = torch.stack([self.pil2tensor(img) for img in images]).to(self.device)
+            posterior = self.vae.encode(pixel_values.to(self.vae.dtype)).latent_dist
+            latents = (posterior.sample() - self.shift_factor) * self.scaling_factor
+
+        # 添加BOS token
+        if self.tokenizer.bos_token is not None:
+            captions = [self.tokenizer.bos_token + caption for caption in captions]
+
+        if seed is not None:
+            self.set_seed(seed)
+
+        # 4. tokenize caption & add prefix ids
+        output = self.tokenizer(
+            captions,
+            padding="longest",
+            truncation=False,
+            add_special_tokens=True,
+            return_tensors="pt",
+            padding_side="left"
+        )
+        input_ids = output.input_ids.to(self.device)
+        attention_mask = output.attention_mask.to(self.device)
+        input_ids, attention_mask = self._add_prefix_ids(hw, input_ids, attention_mask)
+
+        # 5. LLM prefill
+        max_new_len = (hw[0] // self.down_factor) * (hw[1] // self.down_factor)
+        max_cache_len = input_ids.shape[1] + max_new_len
+        # past_key_values = StaticCache(
+        #     config=self.model.config,
+        #     max_batch_size=input_ids.shape[0],
+        #     max_cache_len=max_cache_len,
+        #     device=self.device,
+        #     dtype=self.dtype,
+        # )
+        inputs_embeds = self.model.prepare_inputs_embeds(input_ids, latents)
+
+        outputs = self.model.forward(
+            inputs_embeds=inputs_embeds,
+            attention_mask=attention_mask,
+            past_key_values=None,
+            use_cache=True,
+        )
+        past_key_values = outputs.past_key_values
+        c = outputs.last_hidden_state[:, -1:]
+
+
+        # 6. decoding
+        tokens = self.decoding(
+            c=c,
+            attention_mask=attention_mask,
+            past_key_values=past_key_values,
+            max_new_len=max_new_len,
+            num_images_per_caption=num_images_per_caption,
+            use_norm=use_norm,
+            cfg=cfg,
+            cfg_img=cfg_img,
+            cfg_schedule=cfg_schedule,
+            timesteps_shift=timesteps_shift,
+            num_sampling_steps=num_sampling_steps,
+            progress=progress,
+            hw=hw,
+            sde_type=sde_type,
+        )
+
+        # 7. unpatchify
+        latents = self.model.unpatchify(tokens, h=hw[0] // self.down_factor, w=hw[1] // self.down_factor)
+        latents = (latents / self.scaling_factor) + self.shift_factor
+
+        # 8. decode latents
+        sampled_images = self.vae.decode(latents.to(self.vae.dtype)).sample
+        sampled_images = sampled_images.detach().cpu().to(torch.float32)
+        pil_images = [self.to_pil(img) for img in sampled_images]
+
+        return pil_images
+

models/nextstep_model.py

@@ -0,0 +1,538 @@
+import math
+import torch
+import torch.nn as nn
+
+from torch.utils.checkpoint import checkpoint
+
+from transformers import Qwen2Model, Qwen2Config
+
+def modulate(x, shift, scale=None):
+    """Adaptive layer normalization modulation"""
+    if shift is None:
+        return x * (1 + scale)
+    return x * (1 + scale) + shift
+
+
+def expand_t(t, x):
+    """Function to reshape time t to broadcastable dimension of x
+    Args:
+        t: [bsz,], time vector
+        x: [bsz,...], data point
+    """
+    return t.view(-1, *([1] * (x.ndim - 1)))
+
+class ResBlock(nn.Module):
+    def __init__(self, channels, mlp_ratio=1.0):
+        super().__init__()
+        self.channels = channels
+        self.intermediate_size = int(channels * mlp_ratio)
+
+        self.in_ln = nn.LayerNorm(self.channels, eps=1e-6)
+        self.mlp = nn.Sequential(
+            nn.Linear(self.channels, self.intermediate_size),
+            nn.SiLU(),
+            nn.Linear(self.intermediate_size, self.channels),
+        )
+
+        self.adaLN_modulation = nn.Sequential(nn.SiLU(), nn.Linear(channels, 3 * channels, bias=True))
+
+    def forward(self, x, y):
+        shift_mlp, scale_mlp, gate_mlp = self.adaLN_modulation(y).chunk(3, dim=-1)
+        h = modulate(self.in_ln(x), shift_mlp, scale_mlp)
+        h = self.mlp(h)
+        return x + gate_mlp * h
+
+
+class FinalLayer(nn.Module):
+    def __init__(self, model_channels, out_channels):
+        super().__init__()
+        self.norm_final = nn.LayerNorm(model_channels, elementwise_affine=False, eps=1e-6)
+        self.linear = nn.Linear(model_channels, out_channels, bias=True)
+        self.adaLN_modulation = nn.Sequential(nn.SiLU(), nn.Linear(model_channels, 2 * model_channels, bias=True))
+
+    def forward(self, x, c):
+        shift, scale = self.adaLN_modulation(c).chunk(2, dim=-1)
+        x = modulate(self.norm_final(x), shift, scale)
+        x = self.linear(x)
+        return x
+
+
+class TimestepEmbedder(nn.Module):
+    """
+    Embeds scalar timesteps into vector representations.
+
+    Note: diffusers.models.embeddings.TimestepEmbedding provides similar functionality
+    but with a different interface (requires in_channels, time_embed_dim).
+    This implementation follows the Glide-style timestep embedding.
+    """
+
+    def __init__(self, hidden_size, frequency_embedding_size=256):
+        super().__init__()
+        self.mlp = nn.Sequential(
+            nn.Linear(frequency_embedding_size, hidden_size, bias=True),
+            nn.SiLU(),
+            nn.Linear(hidden_size, hidden_size, bias=True),
+        )
+        self.frequency_embedding_size = frequency_embedding_size
+
+    @staticmethod
+    def timestep_embedding(t: torch.Tensor, dim: int, max_period: float = 10000.0):
+        """
+        Create sinusoidal timestep embeddings.
+        :param t: a 1-D Tensor of N indices, one per batch element. These may be fractional.
+        :param dim: the dimension of the output.
+        :param max_period: controls the minimum frequency of the embeddings.
+        :return: an (N, D) Tensor of positional embeddings.
+        """
+        # https://github.com/openai/glide-text2im/blob/main/glide_text2im/nn.py
+        half = dim // 2
+        freqs = torch.exp(-math.log(max_period) * torch.arange(start=0, end=half, dtype=torch.float32) / half).to(
+            device=t.device
+        )
+        args = t[:, None].float() * freqs[None]
+        embedding = torch.cat([torch.cos(args), torch.sin(args)], dim=-1)
+        if dim % 2:
+            embedding = torch.cat([embedding, torch.zeros_like(embedding[:, :1])], dim=-1)
+        return embedding
+
+    def forward(self, t):
+        t_freq = self.timestep_embedding(t, self.frequency_embedding_size)
+        t_emb = self.mlp(t_freq.to(self.mlp[0].weight.dtype))
+        return t_emb
+
+
+class SimpleMLPAdaLN(nn.Module):
+    def __init__(self, input_dim, cond_dim, dim=1536, layers=12, mlp_ratio=1.0):
+        super().__init__()
+        self.input_dim = input_dim
+        self.cond_dim = cond_dim
+        self.dim = dim
+        self.layers = layers
+        self.mlp_ratio = mlp_ratio
+
+        self.time_embed = TimestepEmbedder(dim)
+        self.cond_embed = nn.Linear(cond_dim, dim)
+        self.input_proj = nn.Linear(input_dim, dim)
+
+        res_blocks = []
+        for _ in range(layers):
+            res_blocks.append(ResBlock(dim, mlp_ratio))
+        self.res_blocks = nn.ModuleList(res_blocks)
+
+        self.final_layer = FinalLayer(dim, input_dim)
+
+        self.grad_checkpointing = False
+
+        self.initialize_weights()
+
+    def initialize_weights(self):
+        def _basic_init(module):
+            if isinstance(module, nn.Linear):
+                torch.nn.init.xavier_uniform_(module.weight)
+                if module.bias is not None:
+                    nn.init.constant_(module.bias, 0)
+
+        self.apply(_basic_init)
+
+        # Initialize timestep embedding MLP
+        nn.init.normal_(self.time_embed.mlp[0].weight, std=0.02)
+        nn.init.normal_(self.time_embed.mlp[2].weight, std=0.02)
+
+        # Zero-out adaLN modulation layers
+        for block in self.res_blocks:
+            nn.init.constant_(block.adaLN_modulation[-1].weight, 0)
+            nn.init.constant_(block.adaLN_modulation[-1].bias, 0)
+
+        # Zero-out output layers
+        nn.init.constant_(self.final_layer.adaLN_modulation[-1].weight, 0)
+        nn.init.constant_(self.final_layer.adaLN_modulation[-1].bias, 0)
+        nn.init.constant_(self.final_layer.linear.weight, 0)
+        nn.init.constant_(self.final_layer.linear.bias, 0)
+
+    def forward(self, x, t, c):
+        """
+        x.shape = (bsz, input_dim)
+        t.shape = (bsz,)
+        c.shape = (bsz, cond_dim)
+        """
+
+        x = self.input_proj(x)
+        t = self.time_embed(t)
+        c = self.cond_embed(c)
+
+        y = t + c
+
+        for block in self.res_blocks:
+            if self.grad_checkpointing and self.training:
+                x = checkpoint(block, x, y, use_reentrant=False)
+            else:
+                x = block(x, y)
+
+        return self.final_layer(x, y)
+
+
+class FlowMatchingHead(nn.Module):
+
+    def __init__(self, input_dim, cond_dim, dim=1536, layers=12, mlp_ratio=1.0):
+        super(FlowMatchingHead, self).__init__()
+        self.input_dim = input_dim
+        self.net = SimpleMLPAdaLN(input_dim=input_dim, cond_dim=cond_dim, dim=dim, layers=layers, mlp_ratio=mlp_ratio)
+
+    @property
+    def dtype(self):
+        return self.net.input_proj.weight.dtype
+
+    @property
+    def device(self):
+        return self.net.input_proj.weight.device
+
+    def get_score_from_velocity(self, velocity, x, t):
+        """Wrapper function: transfrom velocity prediction model to score
+        Args:
+            velocity: [bsz, ...] shaped tensor; velocity model output
+            x:        [bsz, ...] shaped tensor; x_t data point
+            t:        [bsz,] time tensor
+        """
+        t = expand_t(t, x)
+        alpha_t, d_alpha_t = t, 1
+        sigma_t, d_sigma_t = 1 - t, -1
+        mean = x
+        reverse_alpha_ratio = alpha_t / d_alpha_t
+        var = sigma_t**2 - reverse_alpha_ratio * d_sigma_t * sigma_t
+        score = (reverse_alpha_ratio * velocity - mean) / var
+        return score
+
+    def get_velocity_from_cfg(self, velocity, cfg, cfg_img, cfg_mul, normalize=False):
+        if cfg_mul == 2:
+            cond_v, uncond_v = torch.chunk(velocity, 2, dim=0)
+            velocity = uncond_v + cfg * (cond_v - uncond_v)
+            if normalize and cfg > 1.0:
+                norm_cond = torch.norm(cond_v, dim=-1, keepdim=True)
+                norm_velocity = torch.norm(velocity, dim=-1, keepdim=True)
+                scale_factor = norm_cond / (norm_velocity + 1e-6)
+                scale_factor = torch.clamp(scale_factor, min=0, max=1)
+                velocity = velocity * scale_factor
+        elif cfg_mul == 3:
+            cond_v, uncond_v1, uncond_v2 = torch.chunk(velocity, 3, dim=0)
+            velocity = uncond_v2 + cfg_img * (uncond_v1 - uncond_v2) + cfg * (cond_v - uncond_v1)
+            if normalize and cfg > 1.0 and cfg_img > 1.0:
+                norm_cond = torch.norm(cond_v, dim=-1, keepdim=True)
+                norm_velocity = torch.norm(velocity, dim=-1, keepdim=True)
+                scale_factor = norm_cond / (norm_velocity + 1e-6)
+                scale_factor = torch.clamp(scale_factor, min=0, max=1)
+                velocity = velocity * scale_factor
+        return velocity
+
+    def _compute_cfg_mult(self, cfg: float, cfg_img: float) -> int:
+        """计算 CFG 的倍数"""
+        cfg_mul = 1
+        if cfg > 1.0:
+            cfg_mul += 1
+        if cfg_img > 1.0:
+            cfg_mul += 1
+        return cfg_mul
+
+    def _compute_velocity_and_basics(self, x, c, ti, cfg, cfg_img, cfg_mul):
+        """计算 velocity 和基础值（cur_t, next_t, x0, x1）"""
+        combined = torch.cat([x] * cfg_mul, dim=0)
+        velocity = self.net(combined.to(c.dtype), ti.expand(c.shape[0]).to(c), c)
+        velocity = velocity.to(torch.float32)
+        velocity = self.get_velocity_from_cfg(velocity, cfg, cfg_img, cfg_mul)
+
+        cur_t = ti.view(-1, *([1] * (len(x.shape) - 1))).to(x.device)
+        x0 = x - velocity * cur_t  # noise
+        x1 = x + velocity * (1 - cur_t)  # ode image
+
+        return velocity, cur_t, x0, x1
+
+
+    @torch.no_grad()
+    def sample(
+        self,
+        c: torch.Tensor,
+        noise: torch.Tensor = None,
+        cfg: float = 1.0,
+        cfg_img: float = 1.0,
+        timesteps_shift: float = 1.0,
+        num_sampling_steps: int = 20,
+        last_step_size: float = 0.0,
+        noise_repeat: int = 1,
+        sde_solver: bool = False,
+        sde_type: str = "sde",
+        noise_level: float = 0.8,
+    ):
+        # """c.shape = (bsz, cond_dim)"""
+        cfg_mul = self._compute_cfg_mult(cfg, cfg_img)
+
+        if noise is None:
+            noise = torch.randn((c.shape[0] // cfg_mul, self.input_dim), device=c.device, dtype=c.dtype)
+
+        x = noise
+        xs = []
+
+        t0, t1 = 0, 1
+        timesteps = torch.linspace(t0, t1, num_sampling_steps + 1, device=c.device)[:-1]
+        timesteps = timesteps / (timesteps_shift - (timesteps_shift - 1) * timesteps)
+        timesteps = torch.cat([timesteps, torch.ones(1, device=c.device)])
+        # timesteps = torch.tensor([1.0000, 0.9601, 0.9133, 0.8577, 0.7904, 0.7073, 0.6022, 0.4649, 0.2780, 0.0089, 0.0000], device=c.device)
+        # timesteps = 1 - timesteps
+        sigma_max = timesteps[-2]
+        for ti, tj in zip(timesteps[:-1], timesteps[1:]):
+            velocity, cur_t, x0, x1 = self._compute_velocity_and_basics(x, c, ti, cfg, cfg_img, cfg_mul)
+            next_t = tj.view(-1, *([1] * (len(x.shape) - 1))).to(x.device)
+
+            if sde_type == "cps":
+                # Flow-CPS
+                std_dev_t = (1 - next_t)  * math.sin(noise_level * math.pi / 2) # sigma_t in paper
+                sde_noise = torch.randn((c.shape[0] // cfg_mul, self.input_dim), device=c.device, dtype=c.dtype)
+
+                x = x0 * torch.sqrt((1 - next_t)**2 - std_dev_t**2) + x1 * next_t + std_dev_t * sde_noise
+            elif sde_type == "sde":
+                sigma = 1 - cur_t
+                dt = tj - ti
+                std_dev_t = torch.sqrt(sigma / (1 - torch.where(sigma == 1, sigma_max, sigma)))*noise_level
+
+                variance_noise = torch.randn((c.shape[0] // cfg_mul, self.input_dim), device=c.device, dtype=c.dtype)
+
+                x = x0 * (1 - next_t) + x1 * next_t - std_dev_t ** 2 * dt / (2 * sigma) * x0
+                x = x + std_dev_t * torch.sqrt(dt) * variance_noise
+            else:
+                # ODE
+                x = x0 * (1 - next_t) + x1 * next_t
+
+            xs.append(x)
+
+
+        if len(xs) != num_sampling_steps:
+            raise ValueError(f"Samples ({len(xs)}) does not match the number of steps ({num_sampling_steps})")
+
+        return xs[-1].to(c.dtype)
+
+class NextStepConfig(Qwen2Config):
+    model_type = "nextstep"
+
+    def __init__(
+        self,
+        vae_name_or_path: str | None = None,
+        latent_size: int = 32,
+        latent_patch_size: int = 2,
+        latent_channels: int = 16,
+        boi: int | None = None,
+        eoi: int | None = None,
+        image_placeholder_id: int | None = None,
+        pad_token_id_added: int | None = None,
+        lm_loss_weight: float = 0.01,
+        im_loss_weight: float = 1.0,
+        fm_head_dim: int = 1536,
+        fm_head_layers: int = 12,
+        fm_head_batch_mul: int = 4,
+        o_attention_bias: bool | None = None,
+        attn_implementation: str | None = None,  # Add flash attention support
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+
+        # 图像相关参数
+        self.vae_name_or_path = vae_name_or_path
+        self.latent_size = latent_size
+        self.latent_patch_size = latent_patch_size
+        self.latent_channels = latent_channels
+
+        # 特殊token ID
+        self.boi = boi
+        self.eoi = eoi
+        self.image_placeholder_id = image_placeholder_id
+        self.pad_token_id_added = pad_token_id_added
+
+        # 损失权重
+        self.lm_loss_weight = lm_loss_weight
+        self.im_loss_weight = im_loss_weight
+
+        # Flow Matching Head参数
+        self.fm_head_dim = fm_head_dim
+        self.fm_head_layers = fm_head_layers
+        self.fm_head_batch_mul = fm_head_batch_mul
+
+        # Attention bias
+        self.o_attention_bias = o_attention_bias
+
+        # Flash attention support
+        self._attn_implementation = attn_implementation
+
+class NextStep(Qwen2Model):
+    config_class = NextStepConfig
+
+    def __init__(self, config: NextStepConfig, enable_gradient_checkpointing: bool = False):
+        super().__init__(config)
+
+        # 初始化投影器和头部
+        self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
+        token_dim = config.latent_channels * config.latent_patch_size**2
+
+        # 图像投影器
+        self.image_in_projector = self._create_linear(token_dim, config.hidden_size, config.initializer_range)
+        self.image_out_projector = self._create_linear(config.hidden_size, config.hidden_size, config.initializer_range)
+
+        # Flow Matching Head
+        self.image_head = FlowMatchingHead(
+            input_dim=token_dim,
+            cond_dim=config.hidden_size,
+            dim=config.fm_head_dim,
+            layers=config.fm_head_layers,
+        )
+
+        self.gradient_checkpointing = False
+        # 缓存第一个可训练参数，用于建立梯度连接（避免每次遍历）
+        self._first_trainable_param = None
+
+        if enable_gradient_checkpointing:
+            try:
+                self.gradient_checkpointing_enable(gradient_checkpointing_kwargs={"use_reentrant": False})
+                print("Enabled gradient checkpointing (use_reentrant=False) at init for NextStep model")
+            except Exception as e:
+                print(f"Enable gradient checkpointing failed at init: {e}")
+
+    def gradient_checkpointing_enable(self, **kwargs):
+        super().gradient_checkpointing_enable(**kwargs)
+        self.image_head.net.grad_checkpointing = True
+        self.gradient_checkpointing = True
+
+    def _get_first_trainable_param(self):
+        """获取第一个可训练参数的引用（延迟初始化，避免在 __init__ 时 LoRA 还未添加）"""
+        if self._first_trainable_param is None:
+            for param in self.parameters():
+                if param.requires_grad:
+                    self._first_trainable_param = param
+                    break
+        return self._first_trainable_param
+
+    def _ensure_gradient_connection(self, hidden_states):
+        """确保 hidden_states 有梯度连接（对于 LoRA + gradient checkpointing 很重要）"""
+        if self.training and not hidden_states.requires_grad:
+            # 通过一个可训练参数的恒等操作来建立梯度连接
+            # 这样即使基础层被冻结，梯度也能通过 LoRA 层传播
+            # 使用缓存的可训练参数，避免每次遍历
+            param = self._get_first_trainable_param()
+            if param is not None:
+                # 创建一个恒等操作，但通过可训练参数建立连接
+                # 这样 hidden_states 就能连接到计算图
+                # 使用 flatten()[0] 而不是 sum() 以减少计算开销
+                hidden_states = hidden_states + 0.0 * param.flatten()[0]
+            else:
+                # 如果没有可训练参数，直接设置 requires_grad
+                hidden_states = hidden_states.detach().requires_grad_(True)
+        return hidden_states
+
+    def _create_linear(self, in_features: int, out_features: int, std: float) -> nn.Linear:
+        """创建并初始化线性层"""
+        linear = nn.Linear(in_features, out_features)
+        linear.weight.data.normal_(mean=0.0, std=std)
+        linear.bias.data.zero_()
+        return linear
+
+    def patchify(self, img: torch.Tensor):
+        """
+        img: (bsz, C, H, W)
+        x: (bsz, H * W / patch_size**2, patch_size**2 * C)
+        """
+        bsz, c, h, w = img.shape
+        p = self.config.latent_patch_size
+        h_, w_ = h // p, w // p
+
+        img = img.reshape(bsz, c, h_, p, w_, p)
+        img = torch.einsum("nchpwq->nhwcpq", img)
+        x = img.reshape(bsz, h_ * w_, c * p**2)
+        return x
+
+    def unpatchify(self, x: torch.Tensor, h: int = None, w: int = None):
+        """
+        x: (bsz, H * W / patch_size**2, patch_size**2 * C)
+        img: (bsz, C, H, W)
+        """
+        bsz = x.shape[0]
+        p = self.config.latent_patch_size
+        c = self.config.latent_channels
+        if h is None and w is None:
+            h_ = w_ = int(x.shape[1] ** 0.5)
+        else:
+            h_, w_ = h, w
+        assert h_ * w_ == x.shape[1], f"Invalid sequence length {x.shape[1]}."
+
+        x = x.reshape(bsz, h_, w_, c, p, p)
+        x = torch.einsum("nhwcpq->nchpwq", x)
+        img = x.reshape(bsz, c, h_ * p, w_ * p)
+        return img
+
+    def prepare_inputs_embeds(self, input_ids: torch.LongTensor | None = None, latents: torch.FloatTensor | None = None):
+        """准备输入嵌入，支持文本和图像token混合"""
+        if latents is None:
+            return self.embed_tokens(input_ids)
+
+        bs, seq_length = input_ids.shape
+        inputs_embeds = torch.zeros(
+            (bs, seq_length, self.config.hidden_size),
+            device=self.embed_tokens.weight.device,
+            dtype=self.embed_tokens.weight.dtype,
+        )
+
+        im_indices = input_ids == self.config.image_placeholder_id
+        lm_indices = ~im_indices
+
+        # 处理图像latents
+        try:
+            if isinstance(latents, list):
+                tokens = torch.cat([self.patchify(latent) for latent in latents], dim=1)
+            else:
+                tokens = self.patchify(latents)
+        except Exception as e:
+            tokens = latents
+
+        image_embeds = self.image_in_projector(tokens).view(-1, self.config.hidden_size)
+        token_embeds = self.embed_tokens(input_ids[lm_indices])
+
+        inputs_embeds[im_indices] = image_embeds.to(inputs_embeds.dtype)
+        inputs_embeds[lm_indices] = token_embeds
+
+        return inputs_embeds
+
+    def forward(self, input_ids=None, attention_mask=None, position_ids=None,
+                inputs_embeds=None, past_key_values=None, use_cache=None,
+                output_attentions=None, output_hidden_states=None,
+                return_dict=None, forward_head=False,normalize=False, **kwargs):
+        """重写 forward 方法，添加梯度连接逻辑"""
+        # 如果提供了 inputs_embeds，确保它有梯度连接
+        if forward_head:
+            x, t, c, cfg, cfg_img, cfg_mul = kwargs["x"], kwargs["t"], kwargs["c"],kwargs["cfg"], kwargs["cfg_img"], kwargs["cfg_mul"]
+            return self.forward_head(x, t, c, cfg, cfg_img, cfg_mul, normalize=normalize)
+
+        if inputs_embeds is not None:
+            inputs_embeds = self._ensure_gradient_connection(inputs_embeds)
+
+        # 调用父类的 forward 方法
+        return super().forward(
+            input_ids=input_ids,
+            attention_mask=attention_mask,
+            position_ids=position_ids,
+            inputs_embeds=inputs_embeds,
+            past_key_values=past_key_values,
+            use_cache=use_cache,
+            output_attentions=output_attentions,
+            output_hidden_states=output_hidden_states,
+            return_dict=return_dict,
+            **kwargs
+        )
+
+    def forward_head(self, x, t, c, cfg, cfg_img, cfg_mul, normalize=False, **kwargs):
+        # Align tensors with the actual module dtypes so Deepspeed casting does not cause matmul mismatches.
+        projector_dtype = next(self.image_out_projector.parameters()).dtype
+        c = c.to(projector_dtype)
+        x = x.to(projector_dtype)
+        t = t.to(projector_dtype)
+
+        c = self.image_out_projector(c)
+        c = c.squeeze(1)
+
+        velocity = self.image_head.net(x, t, c)
+        velocity = velocity.to(torch.float32)
+        velocity = self.image_head.get_velocity_from_cfg(velocity, cfg, cfg_img, cfg_mul, normalize)
+        return velocity

pytorch-model-00001-of-00004.safetensors

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

pytorch-model-00002-of-00004.safetensors

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

pytorch-model-00003-of-00004.safetensors

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

pytorch-model-00004-of-00004.safetensors

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

requirements.txt

@@ -0,0 +1,14 @@
+diffusers==0.34.0
+einops==0.8.1
+gradio==5.42.0
+loguru==0.7.3
+numpy==1.26.4
+omegaconf==2.3.0
+Pillow==11.0.0
+Requests==2.32.4
+safetensors==0.5.3
+tabulate==0.9.0
+torch==2.5.1
+torchvision==0.20.1
+tqdm==4.67.1
+transformers==4.55.0

special_tokens_map.json

@@ -0,0 +1,27 @@
+{
+  "additional_special_tokens": [
+    "<|image_area|>",
+    "<|begin_of_image|>",
+    "<|end_of_image|>",
+    "<|image_placeholder|>",
+    "<|begin_of_prompt_refinement|>",
+    "<|end_of_prompt_refinement|>",
+    "<|begin_of_thinking|>",
+    "<|end_of_thinking|>",
+    "<|beginoftext|>"
+  ],
+  "eos_token": {
+    "content": "<|endoftext|>",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  },
+  "pad_token": {
+    "content": "[PAD]",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  }
+}

tokenizer.json

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

tokenizer_config.json

@@ -0,0 +1,285 @@
+{
+  "add_bos_token": false,
+  "add_prefix_space": false,
+  "added_tokens_decoder": {
+    "151643": {
+      "content": "<|endoftext|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151644": {
+      "content": "<|im_start|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151645": {
+      "content": "<|im_end|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151646": {
+      "content": "<|object_ref_start|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151647": {
+      "content": "<|object_ref_end|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151648": {
+      "content": "<|box_start|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151649": {
+      "content": "<|box_end|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151650": {
+      "content": "<|quad_start|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151651": {
+      "content": "<|quad_end|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151652": {
+      "content": "<|vision_start|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151653": {
+      "content": "<|vision_end|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151654": {
+      "content": "<|vision_pad|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151655": {
+      "content": "<|image_pad|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151656": {
+      "content": "<|video_pad|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151657": {
+      "content": "<tool_call>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151658": {
+      "content": "</tool_call>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151659": {
+      "content": "<|fim_prefix|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151660": {
+      "content": "<|fim_middle|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151661": {
+      "content": "<|fim_suffix|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151662": {
+      "content": "<|fim_pad|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151663": {
+      "content": "<|repo_name|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151664": {
+      "content": "<|file_sep|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151665": {
+      "content": "[PAD]",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151666": {
+      "content": "<|image_area|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151667": {
+      "content": "<|begin_of_image|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151668": {
+      "content": "<|end_of_image|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151669": {
+      "content": "<|image_placeholder|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151670": {
+      "content": "<|begin_of_prompt_refinement|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151671": {
+      "content": "<|end_of_prompt_refinement|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151672": {
+      "content": "<|begin_of_thinking|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151673": {
+      "content": "<|end_of_thinking|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151674": {
+      "content": "<|beginoftext|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    }
+  },
+  "additional_special_tokens": [
+    "<|image_area|>",
+    "<|begin_of_image|>",
+    "<|end_of_image|>",
+    "<|image_placeholder|>",
+    "<|begin_of_prompt_refinement|>",
+    "<|end_of_prompt_refinement|>",
+    "<|begin_of_thinking|>",
+    "<|end_of_thinking|>",
+    "<|beginoftext|>"
+  ],
+  "bos_token": null,
+  "chat_template": "{%- if tools %}\n    {{- '<|im_start|>system\\n' }}\n    {%- if messages[0]['role'] == 'system' %}\n        {{- messages[0]['content'] }}\n    {%- else %}\n        {{- 'You are a helpful assistant.' }}\n    {%- endif %}\n    {{- \"\\n\\n# Tools\\n\\nYou may call one or more functions to assist with the user query.\\n\\nYou are provided with function signatures within <tools></tools> XML tags:\\n<tools>\" }}\n    {%- for tool in tools %}\n        {{- \"\\n\" }}\n        {{- tool | tojson }}\n    {%- endfor %}\n    {{- \"\\n</tools>\\n\\nFor each function call, return a json object with function name and arguments within <tool_call></tool_call> XML tags:\\n<tool_call>\\n{\\\"name\\\": <function-name>, \\\"arguments\\\": <args-json-object>}\\n</tool_call><|im_end|>\\n\" }}\n{%- else %}\n    {%- if messages[0]['role'] == 'system' %}\n        {{- '<|im_start|>system\\n' + messages[0]['content'] + '<|im_end|>\\n' }}\n    {%- else %}\n        {{- '<|im_start|>system\\nYou are a helpful assistant.<|im_end|>\\n' }}\n    {%- endif %}\n{%- endif %}\n{%- for message in messages %}\n    {%- if (message.role == \"user\") or (message.role == \"system\" and not loop.first) or (message.role == \"assistant\" and not message.tool_calls) %}\n        {{- '<|im_start|>' + message.role + '\\n' + message.content + '<|im_end|>' + '\\n' }}\n    {%- elif message.role == \"assistant\" %}\n        {{- '<|im_start|>' + message.role }}\n        {%- if message.content %}\n            {{- '\\n' + message.content }}\n        {%- endif %}\n        {%- for tool_call in message.tool_calls %}\n            {%- if tool_call.function is defined %}\n                {%- set tool_call = tool_call.function %}\n            {%- endif %}\n            {{- '\\n<tool_call>\\n{\"name\": \"' }}\n            {{- tool_call.name }}\n            {{- '\", \"arguments\": ' }}\n            {{- tool_call.arguments | tojson }}\n            {{- '}\\n</tool_call>' }}\n        {%- endfor %}\n        {{- '<|im_end|>\\n' }}\n    {%- elif message.role == \"tool\" %}\n        {%- if (loop.index0 == 0) or (messages[loop.index0 - 1].role != \"tool\") %}\n            {{- '<|im_start|>user' }}\n        {%- endif %}\n        {{- '\\n<tool_response>\\n' }}\n        {{- message.content }}\n        {{- '\\n</tool_response>' }}\n        {%- if loop.last or (messages[loop.index0 + 1].role != \"tool\") %}\n            {{- '<|im_end|>\\n' }}\n        {%- endif %}\n    {%- endif %}\n{%- endfor %}\n{%- if add_generation_prompt %}\n    {{- '<|im_start|>assistant\\n' }}\n{%- endif %}\n",
+  "clean_up_tokenization_spaces": false,
+  "eos_token": "<|endoftext|>",
+  "errors": "replace",
+  "extra_special_tokens": {},
+  "model_max_length": 8192,
+  "pad_token": "[PAD]",
+  "padding_side": "right",
+  "split_special_tokens": false,
+  "tokenizer_class": "Qwen2Tokenizer",
+  "unk_token": null
+}

vae/checkpoint.pt

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