Duplicate from OpenMOSS-Team/MOSS-Image-Base

Co-authored-by: Xu Zhe <Phospheneser@users.noreply.huggingface.co>

.gitattributes

@@ -0,0 +1,11 @@
+*.safetensors filter=lfs diff=lfs merge=lfs -text
+*.bin filter=lfs diff=lfs merge=lfs -text
+tokenizer/tokenizer.json filter=lfs diff=lfs merge=lfs -text
+assets/hero.png filter=lfs diff=lfs merge=lfs -text
+assets/showcase/layout-bboxes.webp filter=lfs diff=lfs merge=lfs -text
+assets/showcase/text-rendering.webp filter=lfs diff=lfs merge=lfs -text
+assets/teaser.webp filter=lfs diff=lfs merge=lfs -text
+assets/comparison/base.webp filter=lfs diff=lfs merge=lfs -text
+assets/comparison/sft.webp filter=lfs diff=lfs merge=lfs -text
+assets/showcase/overview.webp filter=lfs diff=lfs merge=lfs -text
+assets/showcase/overview2.webp filter=lfs diff=lfs merge=lfs -text

.gitignore

@@ -0,0 +1,3 @@
+__pycache__/
+.cache/
+*.tmp

LICENSE

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

README.md

@@ -0,0 +1,189 @@
+---
+license: apache-2.0
+language:
+- en
+- zh
+pipeline_tag: text-to-image
+tags:
+- moss-image
+- diffusers
+- text-to-image
+---
+# MOSS-Image-Base
+
+<div align="center">
+  <p>
+    <a href="https://github.com/OpenMOSS/MOSS-Image">
+    <img class="dark:hidden mx-auto" src="assets/lockup-light.svg" alt="MOSS-Image" width="520">
+    <img class="hidden dark:block mx-auto" src="assets/lockup-dark.svg" alt="MOSS-Image" width="520">
+    </a>
+  </p>
+  <p>
+    <a href="https://huggingface.co/OpenMOSS-Team/MOSS-Image-Base"><img alt="Hugging Face · MOSS-Image-Base" src="https://img.shields.io/badge/🤗%20Hugging%20Face-MOSS--Image--Base-ffc107?style=flat-square"></a>
+    &nbsp;
+    <a href="https://huggingface.co/OpenMOSS-Team/MOSS-Image-SFT"><img alt="Hugging Face · MOSS-Image-SFT" src="https://img.shields.io/badge/🤗%20Hugging%20Face-MOSS--Image--SFT-ffc107?style=flat-square"></a>
+    &nbsp;
+    <a href="https://modelscope.cn/models/openmoss/MOSS-Image-Base"><img alt="ModelScope · MOSS-Image-Base" src="assets/modelscope-base-badge.svg"></a>
+    &nbsp;
+    <a href="https://modelscope.cn/models/openmoss/MOSS-Image-SFT"><img alt="ModelScope · MOSS-Image-SFT" src="assets/modelscope-sft-badge.svg"></a>
+    &nbsp;
+    <a href="https://github.com/OpenMOSS/MOSS-Image"><img alt="GitHub" src="https://img.shields.io/badge/GitHub-Code-181717?style=flat-square&amp;logo=github&amp;logoColor=white"></a>
+    &nbsp;
+    <a href="https://huggingface.co/spaces/OpenMOSS-Team/MOSS-Image"><img alt="Hugging Face Spaces · Demo" src="https://img.shields.io/badge/🤗%20Spaces-Demo-4f46e5?style=flat-square"></a>
+    &nbsp;
+    <a href="https://openmoss.github.io/MOSS-Image/"><img alt="Homepage" src="https://img.shields.io/badge/🌐%20Project-Homepage-2563eb?style=flat-square"></a>
+    &nbsp;
+    <a href="https://github.com/OpenMOSS/MOSS-Image/blob/main/paper/MOSS_Image_Technical_Report.pdf"><img alt="Technical Report" src="https://img.shields.io/badge/📄%20Paper-Technical_Report-b31b1b?style=flat-square"></a>
+    &nbsp;
+    <a href="https://x.com/Open_MOSS"><img alt="Follow OpenMOSS on X" src="https://img.shields.io/badge/Follow-OpenMOSS-000000?style=flat-square&amp;logo=x&amp;logoColor=white"></a>
+    &nbsp;
+    <a href="https://qm.qq.com/q/hphM7VhaDu"><img alt="MOSS-Image QQ 交流群：755535143" src="https://img.shields.io/badge/QQ-755535143-12B7F5?style=flat-square&amp;logo=qq&amp;logoColor=white"></a>
+    &nbsp;
+    <a href="https://applink.feishu.cn/client/chat/chatter/add_by_link?link_token=7a3m2007-815f-409b-9c3d-a4129f5c058e"><img alt="MOSS-Builders Feishu Community" src="assets/feishu-builders-badge.svg"></a>
+    &nbsp;
+    <a href="https://discord.gg/nXFWEyvbb"><img alt="MOSS-Image Discord Community" src="https://img.shields.io/badge/Discord-Community-5865F2?style=flat-square&amp;logo=discord&amp;logoColor=white"></a>
+  </p>
+  <p>
+    <a href="https://github.com/OpenMOSS/MOSS-Image/blob/main/USERGUIDE.md">User guide</a> &nbsp; · &nbsp;
+    <a href="https://github.com/OpenMOSS/MOSS-Image/blob/main/USERGUIDE_ZH.md">中文用户指南</a> &nbsp; · &nbsp;
+    <a href="https://github.com/OpenMOSS/MOSS-Image/blob/main/comfyui/README.md">ComfyUI</a>
+  </p>
+</div>
+
+![MOSS-Image-Base and MOSS-Image-SFT](assets/teaser.webp)
+
+MOSS-Image-Base is the base text-to-image model in the **MOSS-Image** family
+from OpenMOSS. It supports natural-language prompts, tag sequences, and JSON
+layouts for controlling text placement and object arrangement.
+
+## Base vs. SFT
+
+In this example, SFT produces clearer text grouping and a more organized whiteboard layout.
+
+| MOSS-Image-Base | MOSS-Image-SFT |
+| --- | --- |
+| ![MOSS-Image-Base](assets/comparison/base.webp) | ![MOSS-Image-SFT](assets/comparison/sft.webp) |
+
+## Quick start
+
+Use Python 3.10 or later and an NVIDIA GPU with CUDA support. Install a PyTorch
+2.9 or later build compatible with your CUDA environment, then install:
+
+```bash
+pip install "diffusers>=0.38" "transformers>=5.3" "accelerate>=1.10" "safetensors>=0.6" "Pillow>=10"
+```
+
+```python
+import torch
+from diffusers import DiffusionPipeline
+
+pipe = DiffusionPipeline.from_pretrained(
+    "OpenMOSS-Team/MOSS-Image-Base",
+    custom_pipeline="pipeline_moss_image",
+    trust_remote_code=True,
+    torch_dtype={"default": torch.bfloat16, "vae": torch.float32},
+)
+pipe.scheduler = type(pipe.scheduler).from_config(pipe.scheduler.config)
+pipe.to("cuda")
+
+prompt = '''A wide view of a refrigerator beside a kitchen counter with a fruit bowl. A small sticky note on the refrigerator door is handwritten in Chinese:
+"记得买牛奶和鸡蛋！
+牛奶两盒，鸡蛋一打。
+顺便带一袋面包。"'''
+
+image = pipe(
+    prompt,
+    height=1024,
+    width=1024,
+    num_inference_steps=50,
+    guidance_scale=3.0,
+    uncond_skip_layers=[22],
+    generator=torch.Generator("cpu").manual_seed(42),
+).images[0]
+image.save("refrigerator_note.png")
+```
+
+The model and custom pipeline code are downloaded and cached automatically. You
+can also pass a complete local model directory to `from_pretrained`.
+The transformer and text encoder use bfloat16; the VAE and initial noise use
+float32. The bundled VAE uses the FLUX.2 small decoder.
+
+For the Gradio interface, CPU offload, and compilation options, follow the
+[project quick start](https://github.com/OpenMOSS/MOSS-Image#quick-start).
+For ComfyUI, search for **MOSS-Image** in **ComfyUI Manager** to install the
+native plugin. Integration into the official ComfyUI repository is in progress.
+See the [installation and usage guide](https://github.com/OpenMOSS/MOSS-Image/blob/main/comfyui/README.md)
+and [example workflow](https://github.com/OpenMOSS/MOSS-Image/blob/main/comfyui/moss-image/example_workflows/moss-image.json).
+
+## Sampling settings
+
+| Setting | Default |
+| --- | --- |
+| Sampler | Flow-matching Euler |
+| Steps | 50 |
+| CFG | 3 |
+| Shift | 5 |
+| Layer skipping | Negative branch, layer index 22 (zero-based) |
+| Negative prompt | Empty |
+
+The default guidance uses two transformer forwards per step. Pass
+`uncond_skip_layers=[]` to disable negative-branch skipping, or set
+`negative_prompt` to provide a negative prompt. Python API image dimensions must
+be positive multiples of 16.
+
+See the [user guide](https://github.com/OpenMOSS/MOSS-Image/blob/main/USERGUIDE.md)
+for resolution, sampling, prompt, and LoRA training recommendations.
+
+## Prompt enhancer skill
+
+The [MOSS-Image Prompt Enhancer](https://github.com/OpenMOSS/MOSS-Image/blob/main/skills/moss-image-prompt-enhancer/SKILL.md) turns
+image ideas into detailed captions or spatial Layout prompts. It is designed for
+GPT-6-family agents with local file and Python tools. The bundled script converts
+pixel-coordinate drafts into compact JSON with integer centered coordinates;
+it does not call an image-generation service.
+
+Install the [skill folder](https://github.com/OpenMOSS/MOSS-Image/tree/main/skills/moss-image-prompt-enhancer) in your agent's skills
+directory, then invoke `$moss-image-prompt-enhancer` with your image idea, desired
+caption or Layout output, and any canvas dimensions. Paste the compiled result
+into MOSS-Image's positive prompt; for Layout, use the same generation dimensions.
+
+## Layout-to-image generation
+
+Pass a JSON layout as the prompt to describe the scene, rendered text, objects,
+and their positions. Layouts use `summary.caption` and `items`, with
+`bbox_format: "yxyx"` and `bbox_coord_space: "abs_centralized"`.
+
+With the Python API, serialize the layout using
+`json.dumps(layout, ensure_ascii=False)` and pass it to `pipe`. In Gradio, paste
+the JSON into the Prompt field and set the matching image dimensions.
+See the [layout format and example](https://github.com/OpenMOSS/MOSS-Image/blob/main/USERGUIDE.md#layout-prompts-experimental).
+
+Layouts can also support controlled editing and user-friendly layout tools.
+Edit text, element descriptions, and bounding boxes to adjust image content and
+composition in a structured way.
+
+![Layout control](assets/showcase/layout-bboxes.webp)
+
+## Showcase
+
+![Examples overview](assets/showcase/overview.webp)
+
+![Examples overview 2](assets/showcase/overview2.webp)
+
+## Citation
+
+```bibtex
+@misc{mossimage2026,
+  title = {MOSS-Image},
+  author = {{OpenMOSS Team}},
+  year = {2026},
+  howpublished = {\url{https://github.com/OpenMOSS/MOSS-Image}}
+}
+```
+
+## License and acknowledgements
+
+The model weights and bundled inference code are released under the
+[Apache License, Version 2.0](LICENSE).
+
+We thank the Diffusers, Qwen, and FLUX teams.

model_index.json

@@ -0,0 +1,24 @@
+{
+  "_class_name": "MOSSImagePipeline",
+  "_diffusers_version": "0.38.0",
+  "scheduler": [
+    "pipeline_moss_image",
+    "MOSSImageScheduler"
+  ],
+  "text_encoder": [
+    "transformers",
+    "Qwen3_5Model"
+  ],
+  "tokenizer": [
+    "transformers",
+    "Qwen2Tokenizer"
+  ],
+  "transformer": [
+    "pipeline_moss_image",
+    "MOSSImageTransformer"
+  ],
+  "vae": [
+    "diffusers",
+    "AutoencoderKLFlux2"
+  ]
+}

pipeline_moss_image.py

@@ -0,0 +1,440 @@
+"""MOSSImage modeling and text-to-image pipeline (no Motor/TorchTitan imports).
+
+Dependencies: torch>=2.9, diffusers>=0.38, transformers>=5.3, safetensors,
+accelerate, Pillow. Copy this file anywhere and import it normally.
+
+    import torch
+    # Import MOSSImagePipeline from this file for component-level loading.
+    pipe = MOSSImagePipeline.from_components_pretrained(
+        transformer_path="exported_transformer",  # save_pretrained directory
+        text_encoder_path="Qwen/Qwen3.5-9B",
+        vae_path="black-forest-labs/FLUX.2-klein-4B",
+    ).to("cuda")
+    pipe("A cat", generator=torch.Generator("cuda").manual_seed(42)).images[0].save("cat.png")
+
+Default guidance uses CFG 3 and skips block 22 only in the negative branch
+(zero-based main DiT block indices; two forwards per sampling step):
+    pipe("A cat", guidance_scale=3.0, uncond_skip_layers=[22])
+    pipe("A cat", uncond_skip_layers=[])  # disable negative-branch skipping
+Optional additional positive-branch SLG adds a third forward:
+    pipe("A cat", slg_layers=[24], slg_scale=3.0)
+The additional SLG layer/scale above illustrates the API, not a quality-tuned setting.
+
+For a consolidated training state dict, create MOSSImageTransformer and
+call load_training_state_dict(state), then save_pretrained("exported_transformer").
+Select raw or EMA weights explicitly before calling it; distributed checkpoint
+shards/optimizer states are not model state dicts. Defaults match the finalized dense MOSSImage model,
+not MoE, CFG-distilled, video, or reference-image variants. Prompt strings are
+already-rendered captions (dataset-specific quality/domain tags are not added).
+
+First-release defaults: bf16 DiT/text-encoder weights, fp32 RoPE frequencies
+and phase/sin/cos calculation, then RoPE cast to the video activation dtype.
+Initial Gaussian noise is sampled directly in fp32, including with bf16 weights.
+CFG and Euler latent updates use fp32; VAE decoding defaults to fp32. Scheduler
+shift defaults to 5.0. Exact numerical parity with the original fp32-master
+inference pipeline is not a release requirement.
+"""
+from __future__ import annotations
+
+import math
+from pathlib import Path
+
+import torch
+from torch import nn
+from torch.nn import functional as F
+from diffusers import AutoencoderKLFlux2, DiffusionPipeline, FlowMatchEulerDiscreteScheduler
+from diffusers.schedulers.scheduling_utils import SchedulerMixin  # Required by custom-component discovery.
+from diffusers.configuration_utils import ConfigMixin, register_to_config
+from diffusers.image_processor import VaeImageProcessor
+from diffusers.models.modeling_utils import ModelMixin
+from diffusers.models.modeling_outputs import Transformer2DModelOutput
+from diffusers.pipelines.pipeline_utils import ImagePipelineOutput
+from diffusers.utils.torch_utils import randn_tensor
+from transformers import AutoModel, AutoTokenizer
+
+
+class _RMSNorm(nn.Module):
+    def __init__(self, dim):
+        super().__init__()
+        self.weight = nn.Parameter(torch.ones(dim))
+
+    def forward(self, x):
+        y = x.float() * torch.rsqrt(x.float().square().mean(-1, keepdim=True) + 1e-6)
+        return y.to(x.dtype) * self.weight
+
+
+def _situ(hidden, gate):
+    return (4.0 * torch.tanh(gate / 4.0) * torch.sigmoid(gate)) * (25.0 * torch.tanh(hidden / 25.0))
+
+
+def _validate_skip_layers(layers, depth, name='skip_layers'):
+    layers = () if layers is None else tuple(layers)
+    if any(type(index) is not int or not 0 <= index < depth for index in layers):
+        raise ValueError(f'{name} must contain integer block indices in [0, {depth}).')
+    return frozenset(layers)
+
+
+def _rotate(x, rope):
+    sin, cos = rope.chunk(2, -1)
+    sin = torch.cat((sin, sin), -1).unsqueeze(-2)
+    cos = torch.cat((cos, cos), -1).unsqueeze(-2)
+    a, b = x.chunk(2, -1)
+    return x * cos + torch.cat((-b, a), -1) * sin
+
+
+def _attention(q, k, v, rope, heads):
+    shape = (*q.shape[:2], heads, q.shape[-1] // heads)
+    q = _rotate(q.reshape(shape), rope).transpose(1, 2)
+    k = _rotate(k.reshape(shape), rope).transpose(1, 2)
+    v = v.reshape(shape).transpose(1, 2)
+    return F.scaled_dot_product_attention(q.to(v.dtype), k.to(v.dtype), v).transpose(1, 2).flatten(2)
+
+
+class _TextBlock(nn.Module):
+    def __init__(self, dim, ffn_dim, heads):
+        super().__init__()
+        self.heads = heads
+        self.norm1, self.norm2 = _RMSNorm(dim), _RMSNorm(dim)
+        self.q_norm, self.k_norm = _RMSNorm(dim // heads), _RMSNorm(dim // heads)
+        self.qkv = nn.Linear(dim, 3 * dim)
+        self.proj = nn.Linear(dim, dim)
+        self.ffn_gate, self.ffn_up = nn.Linear(dim, ffn_dim), nn.Linear(dim, ffn_dim)
+        self.ffn_down = nn.Linear(ffn_dim, dim)
+
+    def forward(self, x, rope):
+        q, k, v = self.qkv(self.norm1(x)).chunk(3, -1)
+        q = self.q_norm(q.unflatten(-1, (self.heads, -1))).flatten(2)
+        k = self.k_norm(k.unflatten(-1, (self.heads, -1))).flatten(2)
+        x = x + self.proj(_attention(q, k, v, rope, self.heads))
+        y = self.norm2(x)
+        return x + self.ffn_down(_situ(self.ffn_up(y), self.ffn_gate(y)))
+
+
+class _TextProjection(nn.Module):
+    def __init__(self, text_dim, dim, ffn_dim, heads):
+        super().__init__()
+        self.proj = nn.Linear(text_dim, dim)
+        self.blocks = nn.ModuleList([_TextBlock(dim, ffn_dim, heads) for _ in range(2)])
+
+    def forward(self, x, rope):
+        x = self.proj(x.to(self.proj.weight.dtype))
+        for block in self.blocks:
+            x = block(x, rope)
+        return x
+
+
+class _Block(nn.Module):
+    def __init__(self, dim, ffn_dim, heads):
+        super().__init__()
+        self.heads = heads
+        self.norm1, self.norm2 = _RMSNorm(dim), _RMSNorm(dim)
+        self.q_norm, self.k_norm = _RMSNorm(dim // heads), _RMSNorm(dim // heads)
+        self.proj = nn.ModuleDict({
+            name: nn.Linear(dim, dim, bias=False) for name in ('q', 'k', 'v', 'o')
+        })
+        self.proj['ffn_up_gate'] = nn.Linear(dim, 2 * ffn_dim, bias=False)
+        self.proj['ffn_down'] = nn.Linear(ffn_dim, dim, bias=False)
+
+    def forward(self, x, rope, modulation):
+        shift_a, scale_a, gate_a, shift_f, scale_f, gate_f = modulation.unbind(1)
+        y = self.norm1(x) * (1 + scale_a[:, None]) + shift_a[:, None]
+        dtype = self.proj['q'].weight.dtype
+        y = y.to(dtype)
+        q = self.q_norm(self.proj['q'](y).unflatten(-1, (self.heads, -1))).flatten(2)
+        k = self.k_norm(self.proj['k'](y).unflatten(-1, (self.heads, -1))).flatten(2)
+        x = x + self.proj['o'](_attention(q, k, self.proj['v'](y), rope, self.heads)) * gate_a[:, None]
+        y = self.norm2(x) * (1 + scale_f[:, None]) + shift_f[:, None]
+        hidden, gate = self.proj['ffn_up_gate'](y.to(dtype)).chunk(2, -1)
+        return x + self.proj['ffn_down'](_situ(hidden, gate)) * gate_f[:, None]
+
+
+class _RoPE(nn.Module):
+    def __init__(self, head_dim, base):
+        super().__init__()
+        self.num_bands = head_dim // 8
+        self.base = base
+        self.register_buffer('bands', self._default_bands(torch.device('cpu')))
+
+    def _default_bands(self, device):
+        # Match training's CPU fp32 construction, including under meta init.
+        exponent = torch.arange(self.num_bands, device='cpu', dtype=torch.float32) / self.num_bands
+        return (1 / (self.base ** exponent)).to(device)
+
+    def _apply(self, fn, recurse=True):
+        super()._apply(fn, recurse=recurse)
+        # .float() cannot recover frequencies already rounded by .to(bf16).
+        self.bands = self._default_bands(self.bands.device)
+        return self
+
+    def _load_from_state_dict(self, *args, **kwargs):
+        super()._load_from_state_dict(*args, **kwargs)
+        # Frequencies are derived constants; retain the key for checkpoint compatibility.
+        self.bands = self._default_bands(self.bands.device)
+
+    def forward(self, height, width, text_length, device):
+        # Diffusers/Accelerate can install tensors directly, bypassing load hooks.
+        if self.bands.dtype != torch.float32:
+            self.bands = self._default_bands(device)
+        # Coordinates: (r, f, h, w); sequence order is image followed by text.
+        h, w = torch.meshgrid(torch.arange(height, device=device), torch.arange(width, device=device), indexing='ij')
+        image = torch.stack((torch.zeros_like(h), torch.full_like(h, text_length), h, w), -1).flatten(0, 1).float()
+        image[:, 2] -= (height - 1) / 2
+        image[:, 3] -= (width - 1) / 2
+        text = torch.zeros(text_length, 4, device=device, dtype=torch.float32)
+        text[:, 1] = torch.arange(text_length, device=device)
+        phase = torch.cat((image, text), 0).unsqueeze(-1) * self.bands
+        return torch.cat((phase.sin(), phase.cos()), -2).flatten(-2).unsqueeze(0)
+
+
+class MOSSImageTransformer(ModelMixin, ConfigMixin):
+    """MOSSImage DiT with training-compatible state-dict names and SDPA attention.
+
+    Each sample is evaluated at its true text length, including text refinement;
+    padded tokens cannot influence either text or joint attention.
+    """
+    @register_to_config
+    def __init__(self, dim=3072, ffn_dim=8192, num_heads=24, num_layers=42,
+                 text_dim=4096, freq_dim=256, latent_channels=128, rope_base=10000.0):
+        super().__init__()
+        if dim % num_heads or (dim // num_heads) % 8 or freq_dim % 2:
+            raise ValueError('Require dim divisible by heads, head_dim divisible by 8, and even freq_dim.')
+        self.video_embedder = nn.Linear(latent_channels, dim)
+        self.text_embedder = _TextProjection(text_dim, dim, ffn_dim, num_heads)
+        self.rope = _RoPE(dim // num_heads, rope_base)
+        self.time_embedding = nn.Sequential(nn.Linear(freq_dim, dim), nn.SiLU(), nn.Linear(dim, dim))
+        self.shared_block_modulation = nn.Linear(dim, 6 * dim)
+        self.layers = nn.ModuleList([_Block(dim, ffn_dim, num_heads) for _ in range(num_layers)])
+        self.final_norm_video = _RMSNorm(dim)
+        self.final_linear_video = nn.Linear(dim, latent_channels, bias=False)
+        self.video_head_modulation = nn.Linear(dim, 2 * dim)
+
+    def load_training_state_dict(self, state_dict, *, prefix=''):
+        """Load a consolidated raw/EMA model mapping, strictly; no silent key drops."""
+        if prefix:
+            if any(not key.startswith(prefix) for key in state_dict):
+                raise ValueError('All keys must start with the explicit prefix.')
+            state_dict = {key[len(prefix):]: value for key, value in state_dict.items()}
+        return self.load_state_dict(state_dict, strict=True)
+
+    def forward(self, hidden_states, timestep, encoder_hidden_states,
+                encoder_attention_mask=None, return_dict=True, skip_layers=None):
+        skip_layers = _validate_skip_layers(skip_layers, len(self.layers))
+        if hidden_states.ndim != 4 or hidden_states.shape[1] != self.config.latent_channels:
+            raise ValueError('hidden_states must be [B, latent_channels, H, W].')
+        batch, _, height, width = hidden_states.shape
+        if encoder_hidden_states.ndim != 3 or encoder_hidden_states.shape[0] != batch:
+            raise ValueError('encoder_hidden_states must be [B, S, text_dim].')
+        mask = encoder_attention_mask
+        if mask is None:
+            mask = torch.ones(encoder_hidden_states.shape[:2], dtype=torch.bool, device=hidden_states.device)
+        if mask.shape != encoder_hidden_states.shape[:2] or not torch.all((mask == 0) | (mask == 1)):
+            raise ValueError('encoder_attention_mask must be a binary [B, S] mask.')
+        times = torch.as_tensor(timestep, device=hidden_states.device, dtype=torch.float32).reshape(-1)
+        times = times.expand(batch)
+        # Training computes the sinusoid in float64 and feeds float32 into its MLP.
+        half = self.config.freq_dim // 2
+        freq = 10000 ** (-torch.arange(half, device=times.device, dtype=torch.float64) / half)
+        phase = times.double()[:, None] * freq
+        embedding = torch.cat((phase.cos(), phase.sin()), -1).float()
+        time = self.time_embedding(embedding.to(self.time_embedding[0].weight.dtype)).float()
+        mod = self.shared_block_modulation(time.to(self.shared_block_modulation.weight.dtype)).view(batch, 6, -1)
+        head = self.video_head_modulation(time.to(self.video_head_modulation.weight.dtype)).view(batch, 2, -1)
+        outputs = []
+        for index in range(batch):
+            context = encoder_hidden_states[index:index + 1, mask[index].bool()]
+            if context.shape[1] == 0:
+                raise ValueError('Each sample needs at least one valid text token.')
+            video = hidden_states[index:index + 1].flatten(2).transpose(1, 2)
+            video = self.video_embedder(video.to(self.video_embedder.weight.dtype).contiguous())
+            # Match the reference activation dtype under fp32 weights + autocast.
+            rope = self.rope(height, width, context.shape[1], hidden_states.device).to(video.dtype)
+            text = self.text_embedder(context, rope[:, height * width:])
+            x = torch.cat((video, text), 1)
+            for layer_index, layer in enumerate(self.layers):
+                if layer_index not in skip_layers:
+                    x = layer(x, rope, mod[index:index + 1])
+            x = self.final_norm_video(x[:, :height * width].to(self.final_norm_video.weight.dtype))
+            x = x * (1 + head[index:index + 1, 1, None]) + head[index:index + 1, 0, None]
+            x = self.final_linear_video(x.to(self.final_linear_video.weight.dtype))
+            outputs.append(x.transpose(1, 2).reshape(1, -1, height, width))
+        sample = torch.cat(outputs)
+        return Transformer2DModelOutput(sample=sample) if return_dict else (sample,)
+
+
+class MOSSImageScheduler(FlowMatchEulerDiscreteScheduler):
+    """Euler integration using the training project's V2, nearest-grid schedule."""
+    @register_to_config
+    def __init__(self, num_train_timesteps=1000, shift=5.0, training_shift=1.0):
+        super().__init__(num_train_timesteps=num_train_timesteps, shift=1.0)
+        self.register_to_config(shift=shift, training_shift=training_shift)
+
+    def set_timesteps(self, num_inference_steps=50, device=None):
+        if num_inference_steps < 1 or self.config.shift <= 0 or self.config.training_shift <= 0:
+            raise ValueError('Steps and shifts must be positive.')
+        total = self.config.num_train_timesteps
+        def grid(count, shift):
+            lower = shift / (total + shift - 1)
+            values = torch.linspace(1, lower, count)
+            return shift * values / (1 + (shift - 1) * values)
+        sigmas = grid(num_inference_steps, self.config.shift)
+        training = grid(total, self.config.training_shift).flip(0)
+        upper = torch.searchsorted(training, sigmas).clamp(max=total - 1)
+        lower = (upper - 1).clamp(min=0)
+        nearest = torch.where((sigmas - training[lower]).abs() <= (training[upper] - sigmas).abs(), lower, upper)
+        sigmas = torch.unique_consecutive(training[nearest])
+        # Our sigmas have already been shifted and snapped. Avoid a second shift.
+        self._shift = 1.0
+        super().set_timesteps(device=device, sigmas=sigmas.tolist())
+
+
+class MOSSImagePipeline(DiffusionPipeline):
+    model_cpu_offload_seq = 'text_encoder->transformer->vae'
+    _optional_components = ['text_encoder', 'tokenizer', 'vae']
+
+    def __init__(self, transformer: MOSSImageTransformer,
+                 scheduler: MOSSImageScheduler, text_encoder=None, tokenizer=None, vae=None):
+        super().__init__()
+        self.register_modules(transformer=transformer, scheduler=scheduler,
+                              text_encoder=text_encoder, tokenizer=tokenizer, vae=vae)
+        self.image_processor = VaeImageProcessor(vae_scale_factor=16)
+
+    @classmethod
+    def from_components_pretrained(cls, transformer_path, text_encoder_path, vae_path,
+                                   torch_dtype=torch.bfloat16, vae_dtype=torch.float32,
+                                   local_files_only=False):
+        common = dict(local_files_only=local_files_only)
+        transformer = MOSSImageTransformer.from_pretrained(transformer_path, torch_dtype=torch_dtype, **common)
+        tokenizer = AutoTokenizer.from_pretrained(text_encoder_path, **common)
+        text_encoder = AutoModel.from_pretrained(text_encoder_path, dtype=torch_dtype, **common).eval()
+        subfolder = None if (Path(vae_path) / 'config.json').is_file() else 'vae'
+        vae = AutoencoderKLFlux2.from_pretrained(vae_path, subfolder=subfolder, torch_dtype=vae_dtype, **common).eval()
+        return cls(transformer.eval(), MOSSImageScheduler(), text_encoder, tokenizer, vae)
+
+    @torch.no_grad()
+    def encode_prompt(self, prompt, max_sequence_length=4096):
+        if self.text_encoder is None or self.tokenizer is None:
+            raise ValueError('Text encoder/tokenizer required for string prompts.')
+        prompts = [prompt] if isinstance(prompt, str) else list(prompt)
+        if not prompts or not all(isinstance(p, str) for p in prompts):
+            raise ValueError('prompt must be a string or a nonempty list of strings.')
+        rows = []
+        # Unpadded individual forwards also avoid padding effects in Qwen GDN.
+        for value in prompts:
+            formatted = self.tokenizer.apply_chat_template(
+                [{'role': 'user', 'content': value}], tokenize=False,
+                add_generation_prompt=True, enable_thinking=True)
+            inputs = self.tokenizer(formatted, return_tensors='pt', truncation=True,
+                                    max_length=max_sequence_length).to(self._execution_device)
+            output = self.text_encoder(**inputs, use_cache=False, return_dict=True)
+            rows.append(output.last_hidden_state[0])
+        length = max(row.shape[0] for row in rows)
+        embeds = rows[0].new_zeros(len(rows), length, rows[0].shape[-1])
+        mask = torch.zeros(len(rows), length, device=embeds.device, dtype=torch.bool)
+        for index, row in enumerate(rows):
+            embeds[index, :row.shape[0]] = row
+            mask[index, :row.shape[0]] = True
+        return embeds, mask
+
+    @torch.no_grad()
+    def __call__(self, prompt=None, *, height=1024, width=1024, num_inference_steps=50,
+                 guidance_scale=3.0, negative_prompt='', num_images_per_prompt=1,
+                 generator=None, latents=None, prompt_embeds=None, prompt_attention_mask=None,
+                 negative_prompt_embeds=None, negative_prompt_attention_mask=None,
+                 max_sequence_length=4096, output_type='pil', return_dict=True,
+                 uncond_skip_layers=(22,), slg_layers=None, slg_scale=0.0,
+                 slg_start=0.0, slg_stop=1.0):
+        """Generate images with optional full-block skipping (zero-based indices).
+
+        uncond_skip_layers defaults to (22,) and replaces CFG's negative-prompt
+        prediction with a skipped-block prediction, retaining two forwards per
+        step. Pass None or an empty sequence to disable skipping. slg_layers
+        adds a skipped positive-prompt forward and the fp32 correction
+        slg_scale * (cond - cond_skip) to CFG. With guidance_scale=1, active
+        SLG uses cond_skip + slg_scale * (cond - cond_skip), matching Diffusers'
+        default SkipLayerGuidance formulation, with two forwards and no negative
+        branch. Both options can combine. slg_scale=1 then gives normal cond.
+        Both uncond skipping and SLG are active where
+        slg_start <= step_index / num_steps < slg_stop;
+        the interval follows denoising order, from initial noise to final image.
+        Outside it, inference uses normal CFG without skipping.
+        slg_scale=0 or empty slg_layers disables the additional forward.
+        """
+        uncond_skip_layers = _validate_skip_layers(
+            uncond_skip_layers, len(self.transformer.layers), 'uncond_skip_layers')
+        slg_layers = _validate_skip_layers(slg_layers, len(self.transformer.layers), 'slg_layers')
+        if not math.isfinite(slg_scale) or slg_scale < 0:
+            raise ValueError('slg_scale must be finite and nonnegative.')
+        if not 0 <= slg_start <= slg_stop <= 1:
+            raise ValueError('Require 0 <= slg_start <= slg_stop <= 1.')
+        if height <= 0 or width <= 0 or height % 16 or width % 16:
+            raise ValueError('height and width must be positive multiples of 16.')
+        if output_type not in {'pil', 'np', 'pt', 'latent'}:
+            raise ValueError('output_type must be pil, np, pt or latent.')
+        if num_images_per_prompt < 1 or not math.isfinite(guidance_scale) or guidance_scale < 0:
+            raise ValueError('Invalid image count or guidance scale.')
+        if not 1 <= max_sequence_length <= 4096:
+            raise ValueError('max_sequence_length must be between 1 and 4096.')
+        if (prompt is None) == (prompt_embeds is None):
+            raise ValueError('Provide exactly one of prompt and prompt_embeds.')
+        device = self._execution_device
+        if prompt_embeds is None:
+            prompt_embeds, prompt_attention_mask = self.encode_prompt(prompt, max_sequence_length)
+        batch = prompt_embeds.shape[0]
+        def prepare(embeds, mask):
+            if embeds.ndim != 3 or embeds.shape[0] != batch:
+                raise ValueError('Prompt embeddings must be [B, S, C] with matching batch size.')
+            if mask is None:
+                mask = torch.ones(embeds.shape[:2], dtype=torch.bool, device=embeds.device)
+            return (embeds.to(device).repeat_interleave(num_images_per_prompt, 0),
+                    mask.to(device).repeat_interleave(num_images_per_prompt, 0))
+        positive, positive_mask = prepare(prompt_embeds, prompt_attention_mask)
+        do_cfg = guidance_scale != 1.0
+        if do_cfg:
+            if negative_prompt_embeds is None:
+                negatives = [negative_prompt] * batch if isinstance(negative_prompt, str) else negative_prompt
+                if len(negatives) != batch:
+                    raise ValueError('negative_prompt must match prompt batch size.')
+                negative_prompt_embeds, negative_prompt_attention_mask = self.encode_prompt(negatives, max_sequence_length)
+            negative, negative_mask = prepare(negative_prompt_embeds, negative_prompt_attention_mask)
+        shape = (batch * num_images_per_prompt, self.transformer.config.latent_channels, height // 16, width // 16)
+        if latents is None:
+            # Sample in fp32: promoting bf16 noise afterward cannot restore its precision.
+            latents = randn_tensor(shape, generator=generator, device=device, dtype=torch.float32)
+        elif tuple(latents.shape) != shape:
+            raise ValueError(f'latents must have shape {shape}.')
+        else:
+            latents = latents.to(device=device, dtype=torch.float32)
+        self.scheduler.set_timesteps(num_inference_steps, device=device)
+        num_steps = len(self.scheduler.timesteps)
+        for step_index, timestep in enumerate(self.progress_bar(self.scheduler.timesteps)):
+            skip_active = slg_start <= step_index / num_steps < slg_stop
+            # CFG must combine fp32 predictions, even when the DiT runs in bf16.
+            # Converting only the combined velocity would preserve rounding errors.
+            prediction = self.transformer(latents, timestep, positive, positive_mask).sample.float()
+            conditional = prediction
+            if do_cfg:
+                skip_kwargs = {'skip_layers': uncond_skip_layers} if uncond_skip_layers and skip_active else {}
+                uncond = self.transformer(latents, timestep, negative, negative_mask, **skip_kwargs).sample.float()
+                prediction = uncond + guidance_scale * (prediction - uncond)
+            if slg_layers and slg_scale > 0 and skip_active:
+                skipped = self.transformer(
+                    latents, timestep, positive, positive_mask, skip_layers=slg_layers).sample.float()
+                if not do_cfg:
+                    prediction = skipped
+                prediction = prediction + slg_scale * (conditional - skipped)
+            latents = self.scheduler.step(prediction, timestep, latents, return_dict=False)[0]
+        if output_type == 'latent':
+            images = latents
+        else:
+            if self.vae is None:
+                raise ValueError('A FLUX.2 VAE is required for decoded output.')
+            z = latents.to(dtype=self.vae.dtype)
+            bn = self.vae.bn
+            eps = self.vae.config.batch_norm_eps
+            z = z * (bn.running_var + eps).sqrt().to(z).view(1, -1, 1, 1) + bn.running_mean.to(z).view(1, -1, 1, 1)
+            b, c, h, w = z.shape
+            z = z.reshape(b, c // 4, 2, 2, h, w).permute(0, 1, 4, 2, 5, 3).reshape(b, c // 4, 2 * h, 2 * w)
+            images = self.vae.decode(z, return_dict=False)[0].clamp(-1, 1)
+            images = self.image_processor.postprocess(images, output_type=output_type)
+        self.maybe_free_model_hooks()
+        return ImagePipelineOutput(images=images) if return_dict else (images,)

scheduler/pipeline_moss_image.py

@@ -0,0 +1,440 @@
+"""MOSSImage modeling and text-to-image pipeline (no Motor/TorchTitan imports).
+
+Dependencies: torch>=2.9, diffusers>=0.38, transformers>=5.3, safetensors,
+accelerate, Pillow. Copy this file anywhere and import it normally.
+
+    import torch
+    # Import MOSSImagePipeline from this file for component-level loading.
+    pipe = MOSSImagePipeline.from_components_pretrained(
+        transformer_path="exported_transformer",  # save_pretrained directory
+        text_encoder_path="Qwen/Qwen3.5-9B",
+        vae_path="black-forest-labs/FLUX.2-klein-4B",
+    ).to("cuda")
+    pipe("A cat", generator=torch.Generator("cuda").manual_seed(42)).images[0].save("cat.png")
+
+Default guidance uses CFG 3 and skips block 22 only in the negative branch
+(zero-based main DiT block indices; two forwards per sampling step):
+    pipe("A cat", guidance_scale=3.0, uncond_skip_layers=[22])
+    pipe("A cat", uncond_skip_layers=[])  # disable negative-branch skipping
+Optional additional positive-branch SLG adds a third forward:
+    pipe("A cat", slg_layers=[24], slg_scale=3.0)
+The additional SLG layer/scale above illustrates the API, not a quality-tuned setting.
+
+For a consolidated training state dict, create MOSSImageTransformer and
+call load_training_state_dict(state), then save_pretrained("exported_transformer").
+Select raw or EMA weights explicitly before calling it; distributed checkpoint
+shards/optimizer states are not model state dicts. Defaults match the finalized dense MOSSImage model,
+not MoE, CFG-distilled, video, or reference-image variants. Prompt strings are
+already-rendered captions (dataset-specific quality/domain tags are not added).
+
+First-release defaults: bf16 DiT/text-encoder weights, fp32 RoPE frequencies
+and phase/sin/cos calculation, then RoPE cast to the video activation dtype.
+Initial Gaussian noise is sampled directly in fp32, including with bf16 weights.
+CFG and Euler latent updates use fp32; VAE decoding defaults to fp32. Scheduler
+shift defaults to 5.0. Exact numerical parity with the original fp32-master
+inference pipeline is not a release requirement.
+"""
+from __future__ import annotations
+
+import math
+from pathlib import Path
+
+import torch
+from torch import nn
+from torch.nn import functional as F
+from diffusers import AutoencoderKLFlux2, DiffusionPipeline, FlowMatchEulerDiscreteScheduler
+from diffusers.schedulers.scheduling_utils import SchedulerMixin  # Required by custom-component discovery.
+from diffusers.configuration_utils import ConfigMixin, register_to_config
+from diffusers.image_processor import VaeImageProcessor
+from diffusers.models.modeling_utils import ModelMixin
+from diffusers.models.modeling_outputs import Transformer2DModelOutput
+from diffusers.pipelines.pipeline_utils import ImagePipelineOutput
+from diffusers.utils.torch_utils import randn_tensor
+from transformers import AutoModel, AutoTokenizer
+
+
+class _RMSNorm(nn.Module):
+    def __init__(self, dim):
+        super().__init__()
+        self.weight = nn.Parameter(torch.ones(dim))
+
+    def forward(self, x):
+        y = x.float() * torch.rsqrt(x.float().square().mean(-1, keepdim=True) + 1e-6)
+        return y.to(x.dtype) * self.weight
+
+
+def _situ(hidden, gate):
+    return (4.0 * torch.tanh(gate / 4.0) * torch.sigmoid(gate)) * (25.0 * torch.tanh(hidden / 25.0))
+
+
+def _validate_skip_layers(layers, depth, name='skip_layers'):
+    layers = () if layers is None else tuple(layers)
+    if any(type(index) is not int or not 0 <= index < depth for index in layers):
+        raise ValueError(f'{name} must contain integer block indices in [0, {depth}).')
+    return frozenset(layers)
+
+
+def _rotate(x, rope):
+    sin, cos = rope.chunk(2, -1)
+    sin = torch.cat((sin, sin), -1).unsqueeze(-2)
+    cos = torch.cat((cos, cos), -1).unsqueeze(-2)
+    a, b = x.chunk(2, -1)
+    return x * cos + torch.cat((-b, a), -1) * sin
+
+
+def _attention(q, k, v, rope, heads):
+    shape = (*q.shape[:2], heads, q.shape[-1] // heads)
+    q = _rotate(q.reshape(shape), rope).transpose(1, 2)
+    k = _rotate(k.reshape(shape), rope).transpose(1, 2)
+    v = v.reshape(shape).transpose(1, 2)
+    return F.scaled_dot_product_attention(q.to(v.dtype), k.to(v.dtype), v).transpose(1, 2).flatten(2)
+
+
+class _TextBlock(nn.Module):
+    def __init__(self, dim, ffn_dim, heads):
+        super().__init__()
+        self.heads = heads
+        self.norm1, self.norm2 = _RMSNorm(dim), _RMSNorm(dim)
+        self.q_norm, self.k_norm = _RMSNorm(dim // heads), _RMSNorm(dim // heads)
+        self.qkv = nn.Linear(dim, 3 * dim)
+        self.proj = nn.Linear(dim, dim)
+        self.ffn_gate, self.ffn_up = nn.Linear(dim, ffn_dim), nn.Linear(dim, ffn_dim)
+        self.ffn_down = nn.Linear(ffn_dim, dim)
+
+    def forward(self, x, rope):
+        q, k, v = self.qkv(self.norm1(x)).chunk(3, -1)
+        q = self.q_norm(q.unflatten(-1, (self.heads, -1))).flatten(2)
+        k = self.k_norm(k.unflatten(-1, (self.heads, -1))).flatten(2)
+        x = x + self.proj(_attention(q, k, v, rope, self.heads))
+        y = self.norm2(x)
+        return x + self.ffn_down(_situ(self.ffn_up(y), self.ffn_gate(y)))
+
+
+class _TextProjection(nn.Module):
+    def __init__(self, text_dim, dim, ffn_dim, heads):
+        super().__init__()
+        self.proj = nn.Linear(text_dim, dim)
+        self.blocks = nn.ModuleList([_TextBlock(dim, ffn_dim, heads) for _ in range(2)])
+
+    def forward(self, x, rope):
+        x = self.proj(x.to(self.proj.weight.dtype))
+        for block in self.blocks:
+            x = block(x, rope)
+        return x
+
+
+class _Block(nn.Module):
+    def __init__(self, dim, ffn_dim, heads):
+        super().__init__()
+        self.heads = heads
+        self.norm1, self.norm2 = _RMSNorm(dim), _RMSNorm(dim)
+        self.q_norm, self.k_norm = _RMSNorm(dim // heads), _RMSNorm(dim // heads)
+        self.proj = nn.ModuleDict({
+            name: nn.Linear(dim, dim, bias=False) for name in ('q', 'k', 'v', 'o')
+        })
+        self.proj['ffn_up_gate'] = nn.Linear(dim, 2 * ffn_dim, bias=False)
+        self.proj['ffn_down'] = nn.Linear(ffn_dim, dim, bias=False)
+
+    def forward(self, x, rope, modulation):
+        shift_a, scale_a, gate_a, shift_f, scale_f, gate_f = modulation.unbind(1)
+        y = self.norm1(x) * (1 + scale_a[:, None]) + shift_a[:, None]
+        dtype = self.proj['q'].weight.dtype
+        y = y.to(dtype)
+        q = self.q_norm(self.proj['q'](y).unflatten(-1, (self.heads, -1))).flatten(2)
+        k = self.k_norm(self.proj['k'](y).unflatten(-1, (self.heads, -1))).flatten(2)
+        x = x + self.proj['o'](_attention(q, k, self.proj['v'](y), rope, self.heads)) * gate_a[:, None]
+        y = self.norm2(x) * (1 + scale_f[:, None]) + shift_f[:, None]
+        hidden, gate = self.proj['ffn_up_gate'](y.to(dtype)).chunk(2, -1)
+        return x + self.proj['ffn_down'](_situ(hidden, gate)) * gate_f[:, None]
+
+
+class _RoPE(nn.Module):
+    def __init__(self, head_dim, base):
+        super().__init__()
+        self.num_bands = head_dim // 8
+        self.base = base
+        self.register_buffer('bands', self._default_bands(torch.device('cpu')))
+
+    def _default_bands(self, device):
+        # Match training's CPU fp32 construction, including under meta init.
+        exponent = torch.arange(self.num_bands, device='cpu', dtype=torch.float32) / self.num_bands
+        return (1 / (self.base ** exponent)).to(device)
+
+    def _apply(self, fn, recurse=True):
+        super()._apply(fn, recurse=recurse)
+        # .float() cannot recover frequencies already rounded by .to(bf16).
+        self.bands = self._default_bands(self.bands.device)
+        return self
+
+    def _load_from_state_dict(self, *args, **kwargs):
+        super()._load_from_state_dict(*args, **kwargs)
+        # Frequencies are derived constants; retain the key for checkpoint compatibility.
+        self.bands = self._default_bands(self.bands.device)
+
+    def forward(self, height, width, text_length, device):
+        # Diffusers/Accelerate can install tensors directly, bypassing load hooks.
+        if self.bands.dtype != torch.float32:
+            self.bands = self._default_bands(device)
+        # Coordinates: (r, f, h, w); sequence order is image followed by text.
+        h, w = torch.meshgrid(torch.arange(height, device=device), torch.arange(width, device=device), indexing='ij')
+        image = torch.stack((torch.zeros_like(h), torch.full_like(h, text_length), h, w), -1).flatten(0, 1).float()
+        image[:, 2] -= (height - 1) / 2
+        image[:, 3] -= (width - 1) / 2
+        text = torch.zeros(text_length, 4, device=device, dtype=torch.float32)
+        text[:, 1] = torch.arange(text_length, device=device)
+        phase = torch.cat((image, text), 0).unsqueeze(-1) * self.bands
+        return torch.cat((phase.sin(), phase.cos()), -2).flatten(-2).unsqueeze(0)
+
+
+class MOSSImageTransformer(ModelMixin, ConfigMixin):
+    """MOSSImage DiT with training-compatible state-dict names and SDPA attention.
+
+    Each sample is evaluated at its true text length, including text refinement;
+    padded tokens cannot influence either text or joint attention.
+    """
+    @register_to_config
+    def __init__(self, dim=3072, ffn_dim=8192, num_heads=24, num_layers=42,
+                 text_dim=4096, freq_dim=256, latent_channels=128, rope_base=10000.0):
+        super().__init__()
+        if dim % num_heads or (dim // num_heads) % 8 or freq_dim % 2:
+            raise ValueError('Require dim divisible by heads, head_dim divisible by 8, and even freq_dim.')
+        self.video_embedder = nn.Linear(latent_channels, dim)
+        self.text_embedder = _TextProjection(text_dim, dim, ffn_dim, num_heads)
+        self.rope = _RoPE(dim // num_heads, rope_base)
+        self.time_embedding = nn.Sequential(nn.Linear(freq_dim, dim), nn.SiLU(), nn.Linear(dim, dim))
+        self.shared_block_modulation = nn.Linear(dim, 6 * dim)
+        self.layers = nn.ModuleList([_Block(dim, ffn_dim, num_heads) for _ in range(num_layers)])
+        self.final_norm_video = _RMSNorm(dim)
+        self.final_linear_video = nn.Linear(dim, latent_channels, bias=False)
+        self.video_head_modulation = nn.Linear(dim, 2 * dim)
+
+    def load_training_state_dict(self, state_dict, *, prefix=''):
+        """Load a consolidated raw/EMA model mapping, strictly; no silent key drops."""
+        if prefix:
+            if any(not key.startswith(prefix) for key in state_dict):
+                raise ValueError('All keys must start with the explicit prefix.')
+            state_dict = {key[len(prefix):]: value for key, value in state_dict.items()}
+        return self.load_state_dict(state_dict, strict=True)
+
+    def forward(self, hidden_states, timestep, encoder_hidden_states,
+                encoder_attention_mask=None, return_dict=True, skip_layers=None):
+        skip_layers = _validate_skip_layers(skip_layers, len(self.layers))
+        if hidden_states.ndim != 4 or hidden_states.shape[1] != self.config.latent_channels:
+            raise ValueError('hidden_states must be [B, latent_channels, H, W].')
+        batch, _, height, width = hidden_states.shape
+        if encoder_hidden_states.ndim != 3 or encoder_hidden_states.shape[0] != batch:
+            raise ValueError('encoder_hidden_states must be [B, S, text_dim].')
+        mask = encoder_attention_mask
+        if mask is None:
+            mask = torch.ones(encoder_hidden_states.shape[:2], dtype=torch.bool, device=hidden_states.device)
+        if mask.shape != encoder_hidden_states.shape[:2] or not torch.all((mask == 0) | (mask == 1)):
+            raise ValueError('encoder_attention_mask must be a binary [B, S] mask.')
+        times = torch.as_tensor(timestep, device=hidden_states.device, dtype=torch.float32).reshape(-1)
+        times = times.expand(batch)
+        # Training computes the sinusoid in float64 and feeds float32 into its MLP.
+        half = self.config.freq_dim // 2
+        freq = 10000 ** (-torch.arange(half, device=times.device, dtype=torch.float64) / half)
+        phase = times.double()[:, None] * freq
+        embedding = torch.cat((phase.cos(), phase.sin()), -1).float()
+        time = self.time_embedding(embedding.to(self.time_embedding[0].weight.dtype)).float()
+        mod = self.shared_block_modulation(time.to(self.shared_block_modulation.weight.dtype)).view(batch, 6, -1)
+        head = self.video_head_modulation(time.to(self.video_head_modulation.weight.dtype)).view(batch, 2, -1)
+        outputs = []
+        for index in range(batch):
+            context = encoder_hidden_states[index:index + 1, mask[index].bool()]
+            if context.shape[1] == 0:
+                raise ValueError('Each sample needs at least one valid text token.')
+            video = hidden_states[index:index + 1].flatten(2).transpose(1, 2)
+            video = self.video_embedder(video.to(self.video_embedder.weight.dtype).contiguous())
+            # Match the reference activation dtype under fp32 weights + autocast.
+            rope = self.rope(height, width, context.shape[1], hidden_states.device).to(video.dtype)
+            text = self.text_embedder(context, rope[:, height * width:])
+            x = torch.cat((video, text), 1)
+            for layer_index, layer in enumerate(self.layers):
+                if layer_index not in skip_layers:
+                    x = layer(x, rope, mod[index:index + 1])
+            x = self.final_norm_video(x[:, :height * width].to(self.final_norm_video.weight.dtype))
+            x = x * (1 + head[index:index + 1, 1, None]) + head[index:index + 1, 0, None]
+            x = self.final_linear_video(x.to(self.final_linear_video.weight.dtype))
+            outputs.append(x.transpose(1, 2).reshape(1, -1, height, width))
+        sample = torch.cat(outputs)
+        return Transformer2DModelOutput(sample=sample) if return_dict else (sample,)
+
+
+class MOSSImageScheduler(FlowMatchEulerDiscreteScheduler):
+    """Euler integration using the training project's V2, nearest-grid schedule."""
+    @register_to_config
+    def __init__(self, num_train_timesteps=1000, shift=5.0, training_shift=1.0):
+        super().__init__(num_train_timesteps=num_train_timesteps, shift=1.0)
+        self.register_to_config(shift=shift, training_shift=training_shift)
+
+    def set_timesteps(self, num_inference_steps=50, device=None):
+        if num_inference_steps < 1 or self.config.shift <= 0 or self.config.training_shift <= 0:
+            raise ValueError('Steps and shifts must be positive.')
+        total = self.config.num_train_timesteps
+        def grid(count, shift):
+            lower = shift / (total + shift - 1)
+            values = torch.linspace(1, lower, count)
+            return shift * values / (1 + (shift - 1) * values)
+        sigmas = grid(num_inference_steps, self.config.shift)
+        training = grid(total, self.config.training_shift).flip(0)
+        upper = torch.searchsorted(training, sigmas).clamp(max=total - 1)
+        lower = (upper - 1).clamp(min=0)
+        nearest = torch.where((sigmas - training[lower]).abs() <= (training[upper] - sigmas).abs(), lower, upper)
+        sigmas = torch.unique_consecutive(training[nearest])
+        # Our sigmas have already been shifted and snapped. Avoid a second shift.
+        self._shift = 1.0
+        super().set_timesteps(device=device, sigmas=sigmas.tolist())
+
+
+class MOSSImagePipeline(DiffusionPipeline):
+    model_cpu_offload_seq = 'text_encoder->transformer->vae'
+    _optional_components = ['text_encoder', 'tokenizer', 'vae']
+
+    def __init__(self, transformer: MOSSImageTransformer,
+                 scheduler: MOSSImageScheduler, text_encoder=None, tokenizer=None, vae=None):
+        super().__init__()
+        self.register_modules(transformer=transformer, scheduler=scheduler,
+                              text_encoder=text_encoder, tokenizer=tokenizer, vae=vae)
+        self.image_processor = VaeImageProcessor(vae_scale_factor=16)
+
+    @classmethod
+    def from_components_pretrained(cls, transformer_path, text_encoder_path, vae_path,
+                                   torch_dtype=torch.bfloat16, vae_dtype=torch.float32,
+                                   local_files_only=False):
+        common = dict(local_files_only=local_files_only)
+        transformer = MOSSImageTransformer.from_pretrained(transformer_path, torch_dtype=torch_dtype, **common)
+        tokenizer = AutoTokenizer.from_pretrained(text_encoder_path, **common)
+        text_encoder = AutoModel.from_pretrained(text_encoder_path, dtype=torch_dtype, **common).eval()
+        subfolder = None if (Path(vae_path) / 'config.json').is_file() else 'vae'
+        vae = AutoencoderKLFlux2.from_pretrained(vae_path, subfolder=subfolder, torch_dtype=vae_dtype, **common).eval()
+        return cls(transformer.eval(), MOSSImageScheduler(), text_encoder, tokenizer, vae)
+
+    @torch.no_grad()
+    def encode_prompt(self, prompt, max_sequence_length=4096):
+        if self.text_encoder is None or self.tokenizer is None:
+            raise ValueError('Text encoder/tokenizer required for string prompts.')
+        prompts = [prompt] if isinstance(prompt, str) else list(prompt)
+        if not prompts or not all(isinstance(p, str) for p in prompts):
+            raise ValueError('prompt must be a string or a nonempty list of strings.')
+        rows = []
+        # Unpadded individual forwards also avoid padding effects in Qwen GDN.
+        for value in prompts:
+            formatted = self.tokenizer.apply_chat_template(
+                [{'role': 'user', 'content': value}], tokenize=False,
+                add_generation_prompt=True, enable_thinking=True)
+            inputs = self.tokenizer(formatted, return_tensors='pt', truncation=True,
+                                    max_length=max_sequence_length).to(self._execution_device)
+            output = self.text_encoder(**inputs, use_cache=False, return_dict=True)
+            rows.append(output.last_hidden_state[0])
+        length = max(row.shape[0] for row in rows)
+        embeds = rows[0].new_zeros(len(rows), length, rows[0].shape[-1])
+        mask = torch.zeros(len(rows), length, device=embeds.device, dtype=torch.bool)
+        for index, row in enumerate(rows):
+            embeds[index, :row.shape[0]] = row
+            mask[index, :row.shape[0]] = True
+        return embeds, mask
+
+    @torch.no_grad()
+    def __call__(self, prompt=None, *, height=1024, width=1024, num_inference_steps=50,
+                 guidance_scale=3.0, negative_prompt='', num_images_per_prompt=1,
+                 generator=None, latents=None, prompt_embeds=None, prompt_attention_mask=None,
+                 negative_prompt_embeds=None, negative_prompt_attention_mask=None,
+                 max_sequence_length=4096, output_type='pil', return_dict=True,
+                 uncond_skip_layers=(22,), slg_layers=None, slg_scale=0.0,
+                 slg_start=0.0, slg_stop=1.0):
+        """Generate images with optional full-block skipping (zero-based indices).
+
+        uncond_skip_layers defaults to (22,) and replaces CFG's negative-prompt
+        prediction with a skipped-block prediction, retaining two forwards per
+        step. Pass None or an empty sequence to disable skipping. slg_layers
+        adds a skipped positive-prompt forward and the fp32 correction
+        slg_scale * (cond - cond_skip) to CFG. With guidance_scale=1, active
+        SLG uses cond_skip + slg_scale * (cond - cond_skip), matching Diffusers'
+        default SkipLayerGuidance formulation, with two forwards and no negative
+        branch. Both options can combine. slg_scale=1 then gives normal cond.
+        Both uncond skipping and SLG are active where
+        slg_start <= step_index / num_steps < slg_stop;
+        the interval follows denoising order, from initial noise to final image.
+        Outside it, inference uses normal CFG without skipping.
+        slg_scale=0 or empty slg_layers disables the additional forward.
+        """
+        uncond_skip_layers = _validate_skip_layers(
+            uncond_skip_layers, len(self.transformer.layers), 'uncond_skip_layers')
+        slg_layers = _validate_skip_layers(slg_layers, len(self.transformer.layers), 'slg_layers')
+        if not math.isfinite(slg_scale) or slg_scale < 0:
+            raise ValueError('slg_scale must be finite and nonnegative.')
+        if not 0 <= slg_start <= slg_stop <= 1:
+            raise ValueError('Require 0 <= slg_start <= slg_stop <= 1.')
+        if height <= 0 or width <= 0 or height % 16 or width % 16:
+            raise ValueError('height and width must be positive multiples of 16.')
+        if output_type not in {'pil', 'np', 'pt', 'latent'}:
+            raise ValueError('output_type must be pil, np, pt or latent.')
+        if num_images_per_prompt < 1 or not math.isfinite(guidance_scale) or guidance_scale < 0:
+            raise ValueError('Invalid image count or guidance scale.')
+        if not 1 <= max_sequence_length <= 4096:
+            raise ValueError('max_sequence_length must be between 1 and 4096.')
+        if (prompt is None) == (prompt_embeds is None):
+            raise ValueError('Provide exactly one of prompt and prompt_embeds.')
+        device = self._execution_device
+        if prompt_embeds is None:
+            prompt_embeds, prompt_attention_mask = self.encode_prompt(prompt, max_sequence_length)
+        batch = prompt_embeds.shape[0]
+        def prepare(embeds, mask):
+            if embeds.ndim != 3 or embeds.shape[0] != batch:
+                raise ValueError('Prompt embeddings must be [B, S, C] with matching batch size.')
+            if mask is None:
+                mask = torch.ones(embeds.shape[:2], dtype=torch.bool, device=embeds.device)
+            return (embeds.to(device).repeat_interleave(num_images_per_prompt, 0),
+                    mask.to(device).repeat_interleave(num_images_per_prompt, 0))
+        positive, positive_mask = prepare(prompt_embeds, prompt_attention_mask)
+        do_cfg = guidance_scale != 1.0
+        if do_cfg:
+            if negative_prompt_embeds is None:
+                negatives = [negative_prompt] * batch if isinstance(negative_prompt, str) else negative_prompt
+                if len(negatives) != batch:
+                    raise ValueError('negative_prompt must match prompt batch size.')
+                negative_prompt_embeds, negative_prompt_attention_mask = self.encode_prompt(negatives, max_sequence_length)
+            negative, negative_mask = prepare(negative_prompt_embeds, negative_prompt_attention_mask)
+        shape = (batch * num_images_per_prompt, self.transformer.config.latent_channels, height // 16, width // 16)
+        if latents is None:
+            # Sample in fp32: promoting bf16 noise afterward cannot restore its precision.
+            latents = randn_tensor(shape, generator=generator, device=device, dtype=torch.float32)
+        elif tuple(latents.shape) != shape:
+            raise ValueError(f'latents must have shape {shape}.')
+        else:
+            latents = latents.to(device=device, dtype=torch.float32)
+        self.scheduler.set_timesteps(num_inference_steps, device=device)
+        num_steps = len(self.scheduler.timesteps)
+        for step_index, timestep in enumerate(self.progress_bar(self.scheduler.timesteps)):
+            skip_active = slg_start <= step_index / num_steps < slg_stop
+            # CFG must combine fp32 predictions, even when the DiT runs in bf16.
+            # Converting only the combined velocity would preserve rounding errors.
+            prediction = self.transformer(latents, timestep, positive, positive_mask).sample.float()
+            conditional = prediction
+            if do_cfg:
+                skip_kwargs = {'skip_layers': uncond_skip_layers} if uncond_skip_layers and skip_active else {}
+                uncond = self.transformer(latents, timestep, negative, negative_mask, **skip_kwargs).sample.float()
+                prediction = uncond + guidance_scale * (prediction - uncond)
+            if slg_layers and slg_scale > 0 and skip_active:
+                skipped = self.transformer(
+                    latents, timestep, positive, positive_mask, skip_layers=slg_layers).sample.float()
+                if not do_cfg:
+                    prediction = skipped
+                prediction = prediction + slg_scale * (conditional - skipped)
+            latents = self.scheduler.step(prediction, timestep, latents, return_dict=False)[0]
+        if output_type == 'latent':
+            images = latents
+        else:
+            if self.vae is None:
+                raise ValueError('A FLUX.2 VAE is required for decoded output.')
+            z = latents.to(dtype=self.vae.dtype)
+            bn = self.vae.bn
+            eps = self.vae.config.batch_norm_eps
+            z = z * (bn.running_var + eps).sqrt().to(z).view(1, -1, 1, 1) + bn.running_mean.to(z).view(1, -1, 1, 1)
+            b, c, h, w = z.shape
+            z = z.reshape(b, c // 4, 2, 2, h, w).permute(0, 1, 4, 2, 5, 3).reshape(b, c // 4, 2 * h, 2 * w)
+            images = self.vae.decode(z, return_dict=False)[0].clamp(-1, 1)
+            images = self.image_processor.postprocess(images, output_type=output_type)
+        self.maybe_free_model_hooks()
+        return ImagePipelineOutput(images=images) if return_dict else (images,)

scheduler/scheduler_config.json

@@ -0,0 +1,19 @@
+{
+  "_class_name": "MOSSImageScheduler",
+  "_diffusers_version": "0.38.0",
+  "base_image_seq_len": 256,
+  "base_shift": 0.5,
+  "invert_sigmas": false,
+  "max_image_seq_len": 4096,
+  "max_shift": 1.15,
+  "num_train_timesteps": 1000,
+  "shift": 5.0,
+  "shift_terminal": null,
+  "stochastic_sampling": false,
+  "time_shift_type": "exponential",
+  "training_shift": 1.0,
+  "use_beta_sigmas": false,
+  "use_dynamic_shifting": false,
+  "use_exponential_sigmas": false,
+  "use_karras_sigmas": false
+}

text_encoder/config.json

@@ -0,0 +1,109 @@
+{
+  "architectures": [
+    "Qwen3_5Model"
+  ],
+  "dtype": "bfloat16",
+  "image_token_id": 248056,
+  "model_type": "qwen3_5",
+  "text_config": {
+    "attention_bias": false,
+    "attention_dropout": 0.0,
+    "attn_output_gate": true,
+    "bos_token_id": null,
+    "dtype": "bfloat16",
+    "eos_token_id": 248044,
+    "full_attention_interval": 4,
+    "head_dim": 256,
+    "hidden_act": "silu",
+    "hidden_size": 4096,
+    "initializer_range": 0.02,
+    "intermediate_size": 12288,
+    "layer_types": [
+      "linear_attention",
+      "linear_attention",
+      "linear_attention",
+      "full_attention",
+      "linear_attention",
+      "linear_attention",
+      "linear_attention",
+      "full_attention",
+      "linear_attention",
+      "linear_attention",
+      "linear_attention",
+      "full_attention",
+      "linear_attention",
+      "linear_attention",
+      "linear_attention",
+      "full_attention",
+      "linear_attention",
+      "linear_attention",
+      "linear_attention",
+      "full_attention",
+      "linear_attention",
+      "linear_attention",
+      "linear_attention",
+      "full_attention",
+      "linear_attention",
+      "linear_attention",
+      "linear_attention",
+      "full_attention",
+      "linear_attention",
+      "linear_attention",
+      "linear_attention",
+      "full_attention"
+    ],
+    "linear_conv_kernel_dim": 4,
+    "linear_key_head_dim": 128,
+    "linear_num_key_heads": 16,
+    "linear_num_value_heads": 32,
+    "linear_value_head_dim": 128,
+    "mamba_ssm_dtype": "float32",
+    "max_position_embeddings": 262144,
+    "mlp_only_layers": [],
+    "model_type": "qwen3_5_text",
+    "mtp_num_hidden_layers": 1,
+    "mtp_use_dedicated_embeddings": false,
+    "num_attention_heads": 16,
+    "num_hidden_layers": 32,
+    "num_key_value_heads": 4,
+    "pad_token_id": null,
+    "partial_rotary_factor": 0.25,
+    "rms_norm_eps": 1e-06,
+    "rope_parameters": {
+      "mrope_interleaved": true,
+      "mrope_section": [
+        11,
+        11,
+        10
+      ],
+      "partial_rotary_factor": 0.25,
+      "rope_theta": 10000000,
+      "rope_type": "default"
+    },
+    "tie_word_embeddings": false,
+    "use_cache": true,
+    "vocab_size": 248320
+  },
+  "tie_word_embeddings": false,
+  "transformers_version": "5.12.1",
+  "video_token_id": 248057,
+  "vision_config": {
+    "deepstack_visual_indexes": [],
+    "depth": 27,
+    "dtype": "bfloat16",
+    "hidden_act": "gelu_pytorch_tanh",
+    "hidden_size": 1152,
+    "in_channels": 3,
+    "initializer_range": 0.02,
+    "intermediate_size": 4304,
+    "model_type": "qwen3_5_vision",
+    "num_heads": 16,
+    "num_position_embeddings": 2304,
+    "out_hidden_size": 4096,
+    "patch_size": 16,
+    "spatial_merge_size": 2,
+    "temporal_patch_size": 2
+  },
+  "vision_end_token_id": 248054,
+  "vision_start_token_id": 248053
+}

text_encoder/model-00001-of-00004.safetensors

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

text_encoder/model-00002-of-00004.safetensors

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

text_encoder/model-00003-of-00004.safetensors

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

text_encoder/model-00004-of-00004.safetensors

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

text_encoder/model.safetensors.index.json

@@ -0,0 +1,767 @@
+{
+  "metadata": {
+    "total_parameters": 8392695024,
+    "total_size": 16785390048
+  },
+  "weight_map": {
+    "language_model.embed_tokens.weight": "model-00001-of-00004.safetensors",
+    "language_model.layers.0.input_layernorm.weight": "model-00001-of-00004.safetensors",
+    "language_model.layers.0.linear_attn.A_log": "model-00001-of-00004.safetensors",
+    "language_model.layers.0.linear_attn.conv1d.weight": "model-00001-of-00004.safetensors",
+    "language_model.layers.0.linear_attn.dt_bias": "model-00001-of-00004.safetensors",
+    "language_model.layers.0.linear_attn.in_proj_a.weight": "model-00001-of-00004.safetensors",
+    "language_model.layers.0.linear_attn.in_proj_b.weight": "model-00001-of-00004.safetensors",
+    "language_model.layers.0.linear_attn.in_proj_qkv.weight": "model-00001-of-00004.safetensors",
+    "language_model.layers.0.linear_attn.in_proj_z.weight": "model-00001-of-00004.safetensors",
+    "language_model.layers.0.linear_attn.norm.weight": "model-00001-of-00004.safetensors",
+    "language_model.layers.0.linear_attn.out_proj.weight": "model-00001-of-00004.safetensors",
+    "language_model.layers.0.mlp.down_proj.weight": "model-00001-of-00004.safetensors",
+    "language_model.layers.0.mlp.gate_proj.weight": "model-00001-of-00004.safetensors",
+    "language_model.layers.0.mlp.up_proj.weight": "model-00001-of-00004.safetensors",
+    "language_model.layers.0.post_attention_layernorm.weight": "model-00001-of-00004.safetensors",
+    "language_model.layers.1.input_layernorm.weight": "model-00001-of-00004.safetensors",
+    "language_model.layers.1.linear_attn.A_log": "model-00001-of-00004.safetensors",
+    "language_model.layers.1.linear_attn.conv1d.weight": "model-00001-of-00004.safetensors",
+    "language_model.layers.1.linear_attn.dt_bias": "model-00001-of-00004.safetensors",
+    "language_model.layers.1.linear_attn.in_proj_a.weight": "model-00001-of-00004.safetensors",
+    "language_model.layers.1.linear_attn.in_proj_b.weight": "model-00001-of-00004.safetensors",
+    "language_model.layers.1.linear_attn.in_proj_qkv.weight": "model-00001-of-00004.safetensors",
+    "language_model.layers.1.linear_attn.in_proj_z.weight": "model-00001-of-00004.safetensors",
+    "language_model.layers.1.linear_attn.norm.weight": "model-00001-of-00004.safetensors",
+    "language_model.layers.1.linear_attn.out_proj.weight": "model-00001-of-00004.safetensors",
+    "language_model.layers.1.mlp.down_proj.weight": "model-00001-of-00004.safetensors",
+    "language_model.layers.1.mlp.gate_proj.weight": "model-00001-of-00004.safetensors",
+    "language_model.layers.1.mlp.up_proj.weight": "model-00001-of-00004.safetensors",
+    "language_model.layers.1.post_attention_layernorm.weight": "model-00001-of-00004.safetensors",
+    "language_model.layers.10.input_layernorm.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.10.linear_attn.A_log": "model-00002-of-00004.safetensors",
+    "language_model.layers.10.linear_attn.conv1d.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.10.linear_attn.dt_bias": "model-00002-of-00004.safetensors",
+    "language_model.layers.10.linear_attn.in_proj_a.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.10.linear_attn.in_proj_b.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.10.linear_attn.in_proj_qkv.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.10.linear_attn.in_proj_z.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.10.linear_attn.norm.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.10.linear_attn.out_proj.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.10.mlp.down_proj.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.10.mlp.gate_proj.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.10.mlp.up_proj.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.10.post_attention_layernorm.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.11.input_layernorm.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.11.mlp.down_proj.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.11.mlp.gate_proj.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.11.mlp.up_proj.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.11.post_attention_layernorm.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.11.self_attn.k_norm.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.11.self_attn.k_proj.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.11.self_attn.o_proj.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.11.self_attn.q_norm.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.11.self_attn.q_proj.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.11.self_attn.v_proj.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.12.input_layernorm.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.12.linear_attn.A_log": "model-00002-of-00004.safetensors",
+    "language_model.layers.12.linear_attn.conv1d.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.12.linear_attn.dt_bias": "model-00002-of-00004.safetensors",
+    "language_model.layers.12.linear_attn.in_proj_a.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.12.linear_attn.in_proj_b.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.12.linear_attn.in_proj_qkv.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.12.linear_attn.in_proj_z.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.12.linear_attn.norm.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.12.linear_attn.out_proj.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.12.mlp.down_proj.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.12.mlp.gate_proj.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.12.mlp.up_proj.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.12.post_attention_layernorm.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.13.input_layernorm.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.13.linear_attn.A_log": "model-00002-of-00004.safetensors",
+    "language_model.layers.13.linear_attn.conv1d.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.13.linear_attn.dt_bias": "model-00002-of-00004.safetensors",
+    "language_model.layers.13.linear_attn.in_proj_a.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.13.linear_attn.in_proj_b.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.13.linear_attn.in_proj_qkv.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.13.linear_attn.in_proj_z.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.13.linear_attn.norm.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.13.linear_attn.out_proj.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.13.mlp.down_proj.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.13.mlp.gate_proj.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.13.mlp.up_proj.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.13.post_attention_layernorm.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.14.input_layernorm.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.14.linear_attn.A_log": "model-00002-of-00004.safetensors",
+    "language_model.layers.14.linear_attn.conv1d.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.14.linear_attn.dt_bias": "model-00002-of-00004.safetensors",
+    "language_model.layers.14.linear_attn.in_proj_a.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.14.linear_attn.in_proj_b.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.14.linear_attn.in_proj_qkv.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.14.linear_attn.in_proj_z.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.14.linear_attn.norm.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.14.linear_attn.out_proj.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.14.mlp.down_proj.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.14.mlp.gate_proj.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.14.mlp.up_proj.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.14.post_attention_layernorm.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.15.input_layernorm.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.15.mlp.down_proj.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.15.mlp.gate_proj.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.15.mlp.up_proj.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.15.post_attention_layernorm.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.15.self_attn.k_norm.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.15.self_attn.k_proj.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.15.self_attn.o_proj.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.15.self_attn.q_norm.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.15.self_attn.q_proj.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.15.self_attn.v_proj.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.16.input_layernorm.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.16.linear_attn.A_log": "model-00002-of-00004.safetensors",
+    "language_model.layers.16.linear_attn.conv1d.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.16.linear_attn.dt_bias": "model-00002-of-00004.safetensors",
+    "language_model.layers.16.linear_attn.in_proj_a.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.16.linear_attn.in_proj_b.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.16.linear_attn.in_proj_qkv.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.16.linear_attn.in_proj_z.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.16.linear_attn.norm.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.16.linear_attn.out_proj.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.16.mlp.down_proj.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.16.mlp.gate_proj.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.16.mlp.up_proj.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.16.post_attention_layernorm.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.17.input_layernorm.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.17.linear_attn.A_log": "model-00002-of-00004.safetensors",
+    "language_model.layers.17.linear_attn.conv1d.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.17.linear_attn.dt_bias": "model-00002-of-00004.safetensors",
+    "language_model.layers.17.linear_attn.in_proj_a.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.17.linear_attn.in_proj_b.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.17.linear_attn.in_proj_qkv.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.17.linear_attn.in_proj_z.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.17.linear_attn.norm.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.17.linear_attn.out_proj.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.17.mlp.down_proj.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.17.mlp.gate_proj.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.17.mlp.up_proj.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.17.post_attention_layernorm.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.18.input_layernorm.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.18.linear_attn.A_log": "model-00002-of-00004.safetensors",
+    "language_model.layers.18.linear_attn.conv1d.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.18.linear_attn.dt_bias": "model-00002-of-00004.safetensors",
+    "language_model.layers.18.linear_attn.in_proj_a.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.18.linear_attn.in_proj_b.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.18.linear_attn.in_proj_qkv.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.18.linear_attn.in_proj_z.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.18.linear_attn.norm.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.18.linear_attn.out_proj.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.18.mlp.down_proj.weight": "model-00003-of-00004.safetensors",
+    "language_model.layers.18.mlp.gate_proj.weight": "model-00003-of-00004.safetensors",
+    "language_model.layers.18.mlp.up_proj.weight": "model-00003-of-00004.safetensors",
+    "language_model.layers.18.post_attention_layernorm.weight": "model-00003-of-00004.safetensors",
+    "language_model.layers.19.input_layernorm.weight": "model-00003-of-00004.safetensors",
+    "language_model.layers.19.mlp.down_proj.weight": "model-00003-of-00004.safetensors",
+    "language_model.layers.19.mlp.gate_proj.weight": "model-00003-of-00004.safetensors",
+    "language_model.layers.19.mlp.up_proj.weight": "model-00003-of-00004.safetensors",
+    "language_model.layers.19.post_attention_layernorm.weight": "model-00003-of-00004.safetensors",
+    "language_model.layers.19.self_attn.k_norm.weight": "model-00003-of-00004.safetensors",
+    "language_model.layers.19.self_attn.k_proj.weight": "model-00003-of-00004.safetensors",
+    "language_model.layers.19.self_attn.o_proj.weight": "model-00003-of-00004.safetensors",
+    "language_model.layers.19.self_attn.q_norm.weight": "model-00003-of-00004.safetensors",
+    "language_model.layers.19.self_attn.q_proj.weight": "model-00003-of-00004.safetensors",
+    "language_model.layers.19.self_attn.v_proj.weight": "model-00003-of-00004.safetensors",
+    "language_model.layers.2.input_layernorm.weight": "model-00001-of-00004.safetensors",
+    "language_model.layers.2.linear_attn.A_log": "model-00001-of-00004.safetensors",
+    "language_model.layers.2.linear_attn.conv1d.weight": "model-00001-of-00004.safetensors",
+    "language_model.layers.2.linear_attn.dt_bias": "model-00001-of-00004.safetensors",
+    "language_model.layers.2.linear_attn.in_proj_a.weight": "model-00001-of-00004.safetensors",
+    "language_model.layers.2.linear_attn.in_proj_b.weight": "model-00001-of-00004.safetensors",
+    "language_model.layers.2.linear_attn.in_proj_qkv.weight": "model-00001-of-00004.safetensors",
+    "language_model.layers.2.linear_attn.in_proj_z.weight": "model-00001-of-00004.safetensors",
+    "language_model.layers.2.linear_attn.norm.weight": "model-00001-of-00004.safetensors",
+    "language_model.layers.2.linear_attn.out_proj.weight": "model-00001-of-00004.safetensors",
+    "language_model.layers.2.mlp.down_proj.weight": "model-00001-of-00004.safetensors",
+    "language_model.layers.2.mlp.gate_proj.weight": "model-00001-of-00004.safetensors",
+    "language_model.layers.2.mlp.up_proj.weight": "model-00001-of-00004.safetensors",
+    "language_model.layers.2.post_attention_layernorm.weight": "model-00001-of-00004.safetensors",
+    "language_model.layers.20.input_layernorm.weight": "model-00003-of-00004.safetensors",
+    "language_model.layers.20.linear_attn.A_log": "model-00003-of-00004.safetensors",
+    "language_model.layers.20.linear_attn.conv1d.weight": "model-00003-of-00004.safetensors",
+    "language_model.layers.20.linear_attn.dt_bias": "model-00003-of-00004.safetensors",
+    "language_model.layers.20.linear_attn.in_proj_a.weight": "model-00003-of-00004.safetensors",
+    "language_model.layers.20.linear_attn.in_proj_b.weight": "model-00003-of-00004.safetensors",
+    "language_model.layers.20.linear_attn.in_proj_qkv.weight": "model-00003-of-00004.safetensors",
+    "language_model.layers.20.linear_attn.in_proj_z.weight": "model-00003-of-00004.safetensors",
+    "language_model.layers.20.linear_attn.norm.weight": "model-00003-of-00004.safetensors",
+    "language_model.layers.20.linear_attn.out_proj.weight": "model-00003-of-00004.safetensors",
+    "language_model.layers.20.mlp.down_proj.weight": "model-00003-of-00004.safetensors",
+    "language_model.layers.20.mlp.gate_proj.weight": "model-00003-of-00004.safetensors",
+    "language_model.layers.20.mlp.up_proj.weight": "model-00003-of-00004.safetensors",
+    "language_model.layers.20.post_attention_layernorm.weight": "model-00003-of-00004.safetensors",
+    "language_model.layers.21.input_layernorm.weight": "model-00003-of-00004.safetensors",
+    "language_model.layers.21.linear_attn.A_log": "model-00003-of-00004.safetensors",
+    "language_model.layers.21.linear_attn.conv1d.weight": "model-00003-of-00004.safetensors",
+    "language_model.layers.21.linear_attn.dt_bias": "model-00003-of-00004.safetensors",
+    "language_model.layers.21.linear_attn.in_proj_a.weight": "model-00003-of-00004.safetensors",
+    "language_model.layers.21.linear_attn.in_proj_b.weight": "model-00003-of-00004.safetensors",
+    "language_model.layers.21.linear_attn.in_proj_qkv.weight": "model-00003-of-00004.safetensors",
+    "language_model.layers.21.linear_attn.in_proj_z.weight": "model-00003-of-00004.safetensors",
+    "language_model.layers.21.linear_attn.norm.weight": "model-00003-of-00004.safetensors",
+    "language_model.layers.21.linear_attn.out_proj.weight": "model-00003-of-00004.safetensors",
+    "language_model.layers.21.mlp.down_proj.weight": "model-00003-of-00004.safetensors",
+    "language_model.layers.21.mlp.gate_proj.weight": "model-00003-of-00004.safetensors",
+    "language_model.layers.21.mlp.up_proj.weight": "model-00003-of-00004.safetensors",
+    "language_model.layers.21.post_attention_layernorm.weight": "model-00003-of-00004.safetensors",
+    "language_model.layers.22.input_layernorm.weight": "model-00003-of-00004.safetensors",
+    "language_model.layers.22.linear_attn.A_log": "model-00003-of-00004.safetensors",
+    "language_model.layers.22.linear_attn.conv1d.weight": "model-00003-of-00004.safetensors",
+    "language_model.layers.22.linear_attn.dt_bias": "model-00003-of-00004.safetensors",
+    "language_model.layers.22.linear_attn.in_proj_a.weight": "model-00003-of-00004.safetensors",
+    "language_model.layers.22.linear_attn.in_proj_b.weight": "model-00003-of-00004.safetensors",
+    "language_model.layers.22.linear_attn.in_proj_qkv.weight": "model-00003-of-00004.safetensors",
+    "language_model.layers.22.linear_attn.in_proj_z.weight": "model-00003-of-00004.safetensors",
+    "language_model.layers.22.linear_attn.norm.weight": "model-00003-of-00004.safetensors",
+    "language_model.layers.22.linear_attn.out_proj.weight": "model-00003-of-00004.safetensors",
+    "language_model.layers.22.mlp.down_proj.weight": "model-00003-of-00004.safetensors",
+    "language_model.layers.22.mlp.gate_proj.weight": "model-00003-of-00004.safetensors",
+    "language_model.layers.22.mlp.up_proj.weight": "model-00003-of-00004.safetensors",
+    "language_model.layers.22.post_attention_layernorm.weight": "model-00003-of-00004.safetensors",
+    "language_model.layers.23.input_layernorm.weight": "model-00003-of-00004.safetensors",
+    "language_model.layers.23.mlp.down_proj.weight": "model-00003-of-00004.safetensors",
+    "language_model.layers.23.mlp.gate_proj.weight": "model-00003-of-00004.safetensors",
+    "language_model.layers.23.mlp.up_proj.weight": "model-00003-of-00004.safetensors",
+    "language_model.layers.23.post_attention_layernorm.weight": "model-00003-of-00004.safetensors",
+    "language_model.layers.23.self_attn.k_norm.weight": "model-00003-of-00004.safetensors",
+    "language_model.layers.23.self_attn.k_proj.weight": "model-00003-of-00004.safetensors",
+    "language_model.layers.23.self_attn.o_proj.weight": "model-00003-of-00004.safetensors",
+    "language_model.layers.23.self_attn.q_norm.weight": "model-00003-of-00004.safetensors",
+    "language_model.layers.23.self_attn.q_proj.weight": "model-00003-of-00004.safetensors",
+    "language_model.layers.23.self_attn.v_proj.weight": "model-00003-of-00004.safetensors",
+    "language_model.layers.24.input_layernorm.weight": "model-00003-of-00004.safetensors",
+    "language_model.layers.24.linear_attn.A_log": "model-00003-of-00004.safetensors",
+    "language_model.layers.24.linear_attn.conv1d.weight": "model-00003-of-00004.safetensors",
+    "language_model.layers.24.linear_attn.dt_bias": "model-00003-of-00004.safetensors",
+    "language_model.layers.24.linear_attn.in_proj_a.weight": "model-00003-of-00004.safetensors",
+    "language_model.layers.24.linear_attn.in_proj_b.weight": "model-00003-of-00004.safetensors",
+    "language_model.layers.24.linear_attn.in_proj_qkv.weight": "model-00003-of-00004.safetensors",
+    "language_model.layers.24.linear_attn.in_proj_z.weight": "model-00003-of-00004.safetensors",
+    "language_model.layers.24.linear_attn.norm.weight": "model-00003-of-00004.safetensors",
+    "language_model.layers.24.linear_attn.out_proj.weight": "model-00003-of-00004.safetensors",
+    "language_model.layers.24.mlp.down_proj.weight": "model-00003-of-00004.safetensors",
+    "language_model.layers.24.mlp.gate_proj.weight": "model-00003-of-00004.safetensors",
+    "language_model.layers.24.mlp.up_proj.weight": "model-00003-of-00004.safetensors",
+    "language_model.layers.24.post_attention_layernorm.weight": "model-00003-of-00004.safetensors",
+    "language_model.layers.25.input_layernorm.weight": "model-00003-of-00004.safetensors",
+    "language_model.layers.25.linear_attn.A_log": "model-00003-of-00004.safetensors",
+    "language_model.layers.25.linear_attn.conv1d.weight": "model-00003-of-00004.safetensors",
+    "language_model.layers.25.linear_attn.dt_bias": "model-00003-of-00004.safetensors",
+    "language_model.layers.25.linear_attn.in_proj_a.weight": "model-00003-of-00004.safetensors",
+    "language_model.layers.25.linear_attn.in_proj_b.weight": "model-00003-of-00004.safetensors",
+    "language_model.layers.25.linear_attn.in_proj_qkv.weight": "model-00003-of-00004.safetensors",
+    "language_model.layers.25.linear_attn.in_proj_z.weight": "model-00003-of-00004.safetensors",
+    "language_model.layers.25.linear_attn.norm.weight": "model-00003-of-00004.safetensors",
+    "language_model.layers.25.linear_attn.out_proj.weight": "model-00003-of-00004.safetensors",
+    "language_model.layers.25.mlp.down_proj.weight": "model-00003-of-00004.safetensors",
+    "language_model.layers.25.mlp.gate_proj.weight": "model-00003-of-00004.safetensors",
+    "language_model.layers.25.mlp.up_proj.weight": "model-00003-of-00004.safetensors",
+    "language_model.layers.25.post_attention_layernorm.weight": "model-00003-of-00004.safetensors",
+    "language_model.layers.26.input_layernorm.weight": "model-00003-of-00004.safetensors",
+    "language_model.layers.26.linear_attn.A_log": "model-00003-of-00004.safetensors",
+    "language_model.layers.26.linear_attn.conv1d.weight": "model-00003-of-00004.safetensors",
+    "language_model.layers.26.linear_attn.dt_bias": "model-00003-of-00004.safetensors",
+    "language_model.layers.26.linear_attn.in_proj_a.weight": "model-00003-of-00004.safetensors",
+    "language_model.layers.26.linear_attn.in_proj_b.weight": "model-00003-of-00004.safetensors",
+    "language_model.layers.26.linear_attn.in_proj_qkv.weight": "model-00003-of-00004.safetensors",
+    "language_model.layers.26.linear_attn.in_proj_z.weight": "model-00003-of-00004.safetensors",
+    "language_model.layers.26.linear_attn.norm.weight": "model-00003-of-00004.safetensors",
+    "language_model.layers.26.linear_attn.out_proj.weight": "model-00003-of-00004.safetensors",
+    "language_model.layers.26.mlp.down_proj.weight": "model-00003-of-00004.safetensors",
+    "language_model.layers.26.mlp.gate_proj.weight": "model-00003-of-00004.safetensors",
+    "language_model.layers.26.mlp.up_proj.weight": "model-00003-of-00004.safetensors",
+    "language_model.layers.26.post_attention_layernorm.weight": "model-00003-of-00004.safetensors",
+    "language_model.layers.27.input_layernorm.weight": "model-00003-of-00004.safetensors",
+    "language_model.layers.27.mlp.down_proj.weight": "model-00003-of-00004.safetensors",
+    "language_model.layers.27.mlp.gate_proj.weight": "model-00003-of-00004.safetensors",
+    "language_model.layers.27.mlp.up_proj.weight": "model-00003-of-00004.safetensors",
+    "language_model.layers.27.post_attention_layernorm.weight": "model-00003-of-00004.safetensors",
+    "language_model.layers.27.self_attn.k_norm.weight": "model-00003-of-00004.safetensors",
+    "language_model.layers.27.self_attn.k_proj.weight": "model-00003-of-00004.safetensors",
+    "language_model.layers.27.self_attn.o_proj.weight": "model-00003-of-00004.safetensors",
+    "language_model.layers.27.self_attn.q_norm.weight": "model-00003-of-00004.safetensors",
+    "language_model.layers.27.self_attn.q_proj.weight": "model-00003-of-00004.safetensors",
+    "language_model.layers.27.self_attn.v_proj.weight": "model-00003-of-00004.safetensors",
+    "language_model.layers.28.input_layernorm.weight": "model-00003-of-00004.safetensors",
+    "language_model.layers.28.linear_attn.A_log": "model-00003-of-00004.safetensors",
+    "language_model.layers.28.linear_attn.conv1d.weight": "model-00003-of-00004.safetensors",
+    "language_model.layers.28.linear_attn.dt_bias": "model-00003-of-00004.safetensors",
+    "language_model.layers.28.linear_attn.in_proj_a.weight": "model-00003-of-00004.safetensors",
+    "language_model.layers.28.linear_attn.in_proj_b.weight": "model-00003-of-00004.safetensors",
+    "language_model.layers.28.linear_attn.in_proj_qkv.weight": "model-00003-of-00004.safetensors",
+    "language_model.layers.28.linear_attn.in_proj_z.weight": "model-00003-of-00004.safetensors",
+    "language_model.layers.28.linear_attn.norm.weight": "model-00003-of-00004.safetensors",
+    "language_model.layers.28.linear_attn.out_proj.weight": "model-00003-of-00004.safetensors",
+    "language_model.layers.28.mlp.down_proj.weight": "model-00003-of-00004.safetensors",
+    "language_model.layers.28.mlp.gate_proj.weight": "model-00003-of-00004.safetensors",
+    "language_model.layers.28.mlp.up_proj.weight": "model-00003-of-00004.safetensors",
+    "language_model.layers.28.post_attention_layernorm.weight": "model-00003-of-00004.safetensors",
+    "language_model.layers.29.input_layernorm.weight": "model-00003-of-00004.safetensors",
+    "language_model.layers.29.linear_attn.A_log": "model-00003-of-00004.safetensors",
+    "language_model.layers.29.linear_attn.conv1d.weight": "model-00003-of-00004.safetensors",
+    "language_model.layers.29.linear_attn.dt_bias": "model-00003-of-00004.safetensors",
+    "language_model.layers.29.linear_attn.in_proj_a.weight": "model-00003-of-00004.safetensors",
+    "language_model.layers.29.linear_attn.in_proj_b.weight": "model-00003-of-00004.safetensors",
+    "language_model.layers.29.linear_attn.in_proj_qkv.weight": "model-00003-of-00004.safetensors",
+    "language_model.layers.29.linear_attn.in_proj_z.weight": "model-00003-of-00004.safetensors",
+    "language_model.layers.29.linear_attn.norm.weight": "model-00003-of-00004.safetensors",
+    "language_model.layers.29.linear_attn.out_proj.weight": "model-00003-of-00004.safetensors",
+    "language_model.layers.29.mlp.down_proj.weight": "model-00003-of-00004.safetensors",
+    "language_model.layers.29.mlp.gate_proj.weight": "model-00003-of-00004.safetensors",
+    "language_model.layers.29.mlp.up_proj.weight": "model-00004-of-00004.safetensors",
+    "language_model.layers.29.post_attention_layernorm.weight": "model-00004-of-00004.safetensors",
+    "language_model.layers.3.input_layernorm.weight": "model-00001-of-00004.safetensors",
+    "language_model.layers.3.mlp.down_proj.weight": "model-00001-of-00004.safetensors",
+    "language_model.layers.3.mlp.gate_proj.weight": "model-00001-of-00004.safetensors",
+    "language_model.layers.3.mlp.up_proj.weight": "model-00001-of-00004.safetensors",
+    "language_model.layers.3.post_attention_layernorm.weight": "model-00001-of-00004.safetensors",
+    "language_model.layers.3.self_attn.k_norm.weight": "model-00001-of-00004.safetensors",
+    "language_model.layers.3.self_attn.k_proj.weight": "model-00001-of-00004.safetensors",
+    "language_model.layers.3.self_attn.o_proj.weight": "model-00001-of-00004.safetensors",
+    "language_model.layers.3.self_attn.q_norm.weight": "model-00001-of-00004.safetensors",
+    "language_model.layers.3.self_attn.q_proj.weight": "model-00001-of-00004.safetensors",
+    "language_model.layers.3.self_attn.v_proj.weight": "model-00001-of-00004.safetensors",
+    "language_model.layers.30.input_layernorm.weight": "model-00004-of-00004.safetensors",
+    "language_model.layers.30.linear_attn.A_log": "model-00004-of-00004.safetensors",
+    "language_model.layers.30.linear_attn.conv1d.weight": "model-00004-of-00004.safetensors",
+    "language_model.layers.30.linear_attn.dt_bias": "model-00004-of-00004.safetensors",
+    "language_model.layers.30.linear_attn.in_proj_a.weight": "model-00004-of-00004.safetensors",
+    "language_model.layers.30.linear_attn.in_proj_b.weight": "model-00004-of-00004.safetensors",
+    "language_model.layers.30.linear_attn.in_proj_qkv.weight": "model-00004-of-00004.safetensors",
+    "language_model.layers.30.linear_attn.in_proj_z.weight": "model-00004-of-00004.safetensors",
+    "language_model.layers.30.linear_attn.norm.weight": "model-00004-of-00004.safetensors",
+    "language_model.layers.30.linear_attn.out_proj.weight": "model-00004-of-00004.safetensors",
+    "language_model.layers.30.mlp.down_proj.weight": "model-00004-of-00004.safetensors",
+    "language_model.layers.30.mlp.gate_proj.weight": "model-00004-of-00004.safetensors",
+    "language_model.layers.30.mlp.up_proj.weight": "model-00004-of-00004.safetensors",
+    "language_model.layers.30.post_attention_layernorm.weight": "model-00004-of-00004.safetensors",
+    "language_model.layers.31.input_layernorm.weight": "model-00004-of-00004.safetensors",
+    "language_model.layers.31.mlp.down_proj.weight": "model-00004-of-00004.safetensors",
+    "language_model.layers.31.mlp.gate_proj.weight": "model-00004-of-00004.safetensors",
+    "language_model.layers.31.mlp.up_proj.weight": "model-00004-of-00004.safetensors",
+    "language_model.layers.31.post_attention_layernorm.weight": "model-00004-of-00004.safetensors",
+    "language_model.layers.31.self_attn.k_norm.weight": "model-00004-of-00004.safetensors",
+    "language_model.layers.31.self_attn.k_proj.weight": "model-00004-of-00004.safetensors",
+    "language_model.layers.31.self_attn.o_proj.weight": "model-00004-of-00004.safetensors",
+    "language_model.layers.31.self_attn.q_norm.weight": "model-00004-of-00004.safetensors",
+    "language_model.layers.31.self_attn.q_proj.weight": "model-00004-of-00004.safetensors",
+    "language_model.layers.31.self_attn.v_proj.weight": "model-00004-of-00004.safetensors",
+    "language_model.layers.4.input_layernorm.weight": "model-00001-of-00004.safetensors",
+    "language_model.layers.4.linear_attn.A_log": "model-00001-of-00004.safetensors",
+    "language_model.layers.4.linear_attn.conv1d.weight": "model-00001-of-00004.safetensors",
+    "language_model.layers.4.linear_attn.dt_bias": "model-00001-of-00004.safetensors",
+    "language_model.layers.4.linear_attn.in_proj_a.weight": "model-00001-of-00004.safetensors",
+    "language_model.layers.4.linear_attn.in_proj_b.weight": "model-00001-of-00004.safetensors",
+    "language_model.layers.4.linear_attn.in_proj_qkv.weight": "model-00001-of-00004.safetensors",
+    "language_model.layers.4.linear_attn.in_proj_z.weight": "model-00001-of-00004.safetensors",
+    "language_model.layers.4.linear_attn.norm.weight": "model-00001-of-00004.safetensors",
+    "language_model.layers.4.linear_attn.out_proj.weight": "model-00001-of-00004.safetensors",
+    "language_model.layers.4.mlp.down_proj.weight": "model-00001-of-00004.safetensors",
+    "language_model.layers.4.mlp.gate_proj.weight": "model-00001-of-00004.safetensors",
+    "language_model.layers.4.mlp.up_proj.weight": "model-00001-of-00004.safetensors",
+    "language_model.layers.4.post_attention_layernorm.weight": "model-00001-of-00004.safetensors",
+    "language_model.layers.5.input_layernorm.weight": "model-00001-of-00004.safetensors",
+    "language_model.layers.5.linear_attn.A_log": "model-00001-of-00004.safetensors",
+    "language_model.layers.5.linear_attn.conv1d.weight": "model-00001-of-00004.safetensors",
+    "language_model.layers.5.linear_attn.dt_bias": "model-00001-of-00004.safetensors",
+    "language_model.layers.5.linear_attn.in_proj_a.weight": "model-00001-of-00004.safetensors",
+    "language_model.layers.5.linear_attn.in_proj_b.weight": "model-00001-of-00004.safetensors",
+    "language_model.layers.5.linear_attn.in_proj_qkv.weight": "model-00001-of-00004.safetensors",
+    "language_model.layers.5.linear_attn.in_proj_z.weight": "model-00001-of-00004.safetensors",
+    "language_model.layers.5.linear_attn.norm.weight": "model-00001-of-00004.safetensors",
+    "language_model.layers.5.linear_attn.out_proj.weight": "model-00001-of-00004.safetensors",
+    "language_model.layers.5.mlp.down_proj.weight": "model-00001-of-00004.safetensors",
+    "language_model.layers.5.mlp.gate_proj.weight": "model-00001-of-00004.safetensors",
+    "language_model.layers.5.mlp.up_proj.weight": "model-00001-of-00004.safetensors",
+    "language_model.layers.5.post_attention_layernorm.weight": "model-00001-of-00004.safetensors",
+    "language_model.layers.6.input_layernorm.weight": "model-00001-of-00004.safetensors",
+    "language_model.layers.6.linear_attn.A_log": "model-00001-of-00004.safetensors",
+    "language_model.layers.6.linear_attn.conv1d.weight": "model-00001-of-00004.safetensors",
+    "language_model.layers.6.linear_attn.dt_bias": "model-00001-of-00004.safetensors",
+    "language_model.layers.6.linear_attn.in_proj_a.weight": "model-00001-of-00004.safetensors",
+    "language_model.layers.6.linear_attn.in_proj_b.weight": "model-00001-of-00004.safetensors",
+    "language_model.layers.6.linear_attn.in_proj_qkv.weight": "model-00001-of-00004.safetensors",
+    "language_model.layers.6.linear_attn.in_proj_z.weight": "model-00001-of-00004.safetensors",
+    "language_model.layers.6.linear_attn.norm.weight": "model-00001-of-00004.safetensors",
+    "language_model.layers.6.linear_attn.out_proj.weight": "model-00001-of-00004.safetensors",
+    "language_model.layers.6.mlp.down_proj.weight": "model-00001-of-00004.safetensors",
+    "language_model.layers.6.mlp.gate_proj.weight": "model-00001-of-00004.safetensors",
+    "language_model.layers.6.mlp.up_proj.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.6.post_attention_layernorm.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.7.input_layernorm.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.7.mlp.down_proj.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.7.mlp.gate_proj.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.7.mlp.up_proj.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.7.post_attention_layernorm.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.7.self_attn.k_norm.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.7.self_attn.k_proj.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.7.self_attn.o_proj.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.7.self_attn.q_norm.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.7.self_attn.q_proj.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.7.self_attn.v_proj.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.8.input_layernorm.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.8.linear_attn.A_log": "model-00002-of-00004.safetensors",
+    "language_model.layers.8.linear_attn.conv1d.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.8.linear_attn.dt_bias": "model-00002-of-00004.safetensors",
+    "language_model.layers.8.linear_attn.in_proj_a.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.8.linear_attn.in_proj_b.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.8.linear_attn.in_proj_qkv.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.8.linear_attn.in_proj_z.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.8.linear_attn.norm.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.8.linear_attn.out_proj.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.8.mlp.down_proj.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.8.mlp.gate_proj.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.8.mlp.up_proj.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.8.post_attention_layernorm.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.9.input_layernorm.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.9.linear_attn.A_log": "model-00002-of-00004.safetensors",
+    "language_model.layers.9.linear_attn.conv1d.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.9.linear_attn.dt_bias": "model-00002-of-00004.safetensors",
+    "language_model.layers.9.linear_attn.in_proj_a.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.9.linear_attn.in_proj_b.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.9.linear_attn.in_proj_qkv.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.9.linear_attn.in_proj_z.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.9.linear_attn.norm.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.9.linear_attn.out_proj.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.9.mlp.down_proj.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.9.mlp.gate_proj.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.9.mlp.up_proj.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.9.post_attention_layernorm.weight": "model-00002-of-00004.safetensors",
+    "language_model.norm.weight": "model-00004-of-00004.safetensors",
+    "visual.blocks.0.attn.proj.bias": "model-00004-of-00004.safetensors",
+    "visual.blocks.0.attn.proj.weight": "model-00004-of-00004.safetensors",
+    "visual.blocks.0.attn.qkv.bias": "model-00004-of-00004.safetensors",
+    "visual.blocks.0.attn.qkv.weight": "model-00004-of-00004.safetensors",
+    "visual.blocks.0.mlp.linear_fc1.bias": "model-00004-of-00004.safetensors",
+    "visual.blocks.0.mlp.linear_fc1.weight": "model-00004-of-00004.safetensors",
+    "visual.blocks.0.mlp.linear_fc2.bias": "model-00004-of-00004.safetensors",
+    "visual.blocks.0.mlp.linear_fc2.weight": "model-00004-of-00004.safetensors",
+    "visual.blocks.0.norm1.bias": "model-00004-of-00004.safetensors",
+    "visual.blocks.0.norm1.weight": "model-00004-of-00004.safetensors",
+    "visual.blocks.0.norm2.bias": "model-00004-of-00004.safetensors",
+    "visual.blocks.0.norm2.weight": "model-00004-of-00004.safetensors",
+    "visual.blocks.1.attn.proj.bias": "model-00004-of-00004.safetensors",
+    "visual.blocks.1.attn.proj.weight": "model-00004-of-00004.safetensors",
+    "visual.blocks.1.attn.qkv.bias": "model-00004-of-00004.safetensors",
+    "visual.blocks.1.attn.qkv.weight": "model-00004-of-00004.safetensors",
+    "visual.blocks.1.mlp.linear_fc1.bias": "model-00004-of-00004.safetensors",
+    "visual.blocks.1.mlp.linear_fc1.weight": "model-00004-of-00004.safetensors",
+    "visual.blocks.1.mlp.linear_fc2.bias": "model-00004-of-00004.safetensors",
+    "visual.blocks.1.mlp.linear_fc2.weight": "model-00004-of-00004.safetensors",
+    "visual.blocks.1.norm1.bias": "model-00004-of-00004.safetensors",
+    "visual.blocks.1.norm1.weight": "model-00004-of-00004.safetensors",
+    "visual.blocks.1.norm2.bias": "model-00004-of-00004.safetensors",
+    "visual.blocks.1.norm2.weight": "model-00004-of-00004.safetensors",
+    "visual.blocks.10.attn.proj.bias": "model-00004-of-00004.safetensors",
+    "visual.blocks.10.attn.proj.weight": "model-00004-of-00004.safetensors",
+    "visual.blocks.10.attn.qkv.bias": "model-00004-of-00004.safetensors",
+    "visual.blocks.10.attn.qkv.weight": "model-00004-of-00004.safetensors",
+    "visual.blocks.10.mlp.linear_fc1.bias": "model-00004-of-00004.safetensors",
+    "visual.blocks.10.mlp.linear_fc1.weight": "model-00004-of-00004.safetensors",
+    "visual.blocks.10.mlp.linear_fc2.bias": "model-00004-of-00004.safetensors",
+    "visual.blocks.10.mlp.linear_fc2.weight": "model-00004-of-00004.safetensors",
+    "visual.blocks.10.norm1.bias": "model-00004-of-00004.safetensors",
+    "visual.blocks.10.norm1.weight": "model-00004-of-00004.safetensors",
+    "visual.blocks.10.norm2.bias": "model-00004-of-00004.safetensors",
+    "visual.blocks.10.norm2.weight": "model-00004-of-00004.safetensors",
+    "visual.blocks.11.attn.proj.bias": "model-00004-of-00004.safetensors",
+    "visual.blocks.11.attn.proj.weight": "model-00004-of-00004.safetensors",
+    "visual.blocks.11.attn.qkv.bias": "model-00004-of-00004.safetensors",
+    "visual.blocks.11.attn.qkv.weight": "model-00004-of-00004.safetensors",
+    "visual.blocks.11.mlp.linear_fc1.bias": "model-00004-of-00004.safetensors",
+    "visual.blocks.11.mlp.linear_fc1.weight": "model-00004-of-00004.safetensors",
+    "visual.blocks.11.mlp.linear_fc2.bias": "model-00004-of-00004.safetensors",
+    "visual.blocks.11.mlp.linear_fc2.weight": "model-00004-of-00004.safetensors",
+    "visual.blocks.11.norm1.bias": "model-00004-of-00004.safetensors",
+    "visual.blocks.11.norm1.weight": "model-00004-of-00004.safetensors",
+    "visual.blocks.11.norm2.bias": "model-00004-of-00004.safetensors",
+    "visual.blocks.11.norm2.weight": "model-00004-of-00004.safetensors",
+    "visual.blocks.12.attn.proj.bias": "model-00004-of-00004.safetensors",
+    "visual.blocks.12.attn.proj.weight": "model-00004-of-00004.safetensors",
+    "visual.blocks.12.attn.qkv.bias": "model-00004-of-00004.safetensors",
+    "visual.blocks.12.attn.qkv.weight": "model-00004-of-00004.safetensors",
+    "visual.blocks.12.mlp.linear_fc1.bias": "model-00004-of-00004.safetensors",
+    "visual.blocks.12.mlp.linear_fc1.weight": "model-00004-of-00004.safetensors",
+    "visual.blocks.12.mlp.linear_fc2.bias": "model-00004-of-00004.safetensors",
+    "visual.blocks.12.mlp.linear_fc2.weight": "model-00004-of-00004.safetensors",
+    "visual.blocks.12.norm1.bias": "model-00004-of-00004.safetensors",
+    "visual.blocks.12.norm1.weight": "model-00004-of-00004.safetensors",
+    "visual.blocks.12.norm2.bias": "model-00004-of-00004.safetensors",
+    "visual.blocks.12.norm2.weight": "model-00004-of-00004.safetensors",
+    "visual.blocks.13.attn.proj.bias": "model-00004-of-00004.safetensors",
+    "visual.blocks.13.attn.proj.weight": "model-00004-of-00004.safetensors",
+    "visual.blocks.13.attn.qkv.bias": "model-00004-of-00004.safetensors",
+    "visual.blocks.13.attn.qkv.weight": "model-00004-of-00004.safetensors",
+    "visual.blocks.13.mlp.linear_fc1.bias": "model-00004-of-00004.safetensors",
+    "visual.blocks.13.mlp.linear_fc1.weight": "model-00004-of-00004.safetensors",
+    "visual.blocks.13.mlp.linear_fc2.bias": "model-00004-of-00004.safetensors",
+    "visual.blocks.13.mlp.linear_fc2.weight": "model-00004-of-00004.safetensors",
+    "visual.blocks.13.norm1.bias": "model-00004-of-00004.safetensors",
+    "visual.blocks.13.norm1.weight": "model-00004-of-00004.safetensors",
+    "visual.blocks.13.norm2.bias": "model-00004-of-00004.safetensors",
+    "visual.blocks.13.norm2.weight": "model-00004-of-00004.safetensors",
+    "visual.blocks.14.attn.proj.bias": "model-00004-of-00004.safetensors",
+    "visual.blocks.14.attn.proj.weight": "model-00004-of-00004.safetensors",
+    "visual.blocks.14.attn.qkv.bias": "model-00004-of-00004.safetensors",
+    "visual.blocks.14.attn.qkv.weight": "model-00004-of-00004.safetensors",
+    "visual.blocks.14.mlp.linear_fc1.bias": "model-00004-of-00004.safetensors",
+    "visual.blocks.14.mlp.linear_fc1.weight": "model-00004-of-00004.safetensors",
+    "visual.blocks.14.mlp.linear_fc2.bias": "model-00004-of-00004.safetensors",
+    "visual.blocks.14.mlp.linear_fc2.weight": "model-00004-of-00004.safetensors",
+    "visual.blocks.14.norm1.bias": "model-00004-of-00004.safetensors",
+    "visual.blocks.14.norm1.weight": "model-00004-of-00004.safetensors",
+    "visual.blocks.14.norm2.bias": "model-00004-of-00004.safetensors",
+    "visual.blocks.14.norm2.weight": "model-00004-of-00004.safetensors",
+    "visual.blocks.15.attn.proj.bias": "model-00004-of-00004.safetensors",
+    "visual.blocks.15.attn.proj.weight": "model-00004-of-00004.safetensors",
+    "visual.blocks.15.attn.qkv.bias": "model-00004-of-00004.safetensors",
+    "visual.blocks.15.attn.qkv.weight": "model-00004-of-00004.safetensors",
+    "visual.blocks.15.mlp.linear_fc1.bias": "model-00004-of-00004.safetensors",
+    "visual.blocks.15.mlp.linear_fc1.weight": "model-00004-of-00004.safetensors",
+    "visual.blocks.15.mlp.linear_fc2.bias": "model-00004-of-00004.safetensors",
+    "visual.blocks.15.mlp.linear_fc2.weight": "model-00004-of-00004.safetensors",
+    "visual.blocks.15.norm1.bias": "model-00004-of-00004.safetensors",
+    "visual.blocks.15.norm1.weight": "model-00004-of-00004.safetensors",
+    "visual.blocks.15.norm2.bias": "model-00004-of-00004.safetensors",
+    "visual.blocks.15.norm2.weight": "model-00004-of-00004.safetensors",
+    "visual.blocks.16.attn.proj.bias": "model-00004-of-00004.safetensors",
+    "visual.blocks.16.attn.proj.weight": "model-00004-of-00004.safetensors",
+    "visual.blocks.16.attn.qkv.bias": "model-00004-of-00004.safetensors",
+    "visual.blocks.16.attn.qkv.weight": "model-00004-of-00004.safetensors",
+    "visual.blocks.16.mlp.linear_fc1.bias": "model-00004-of-00004.safetensors",
+    "visual.blocks.16.mlp.linear_fc1.weight": "model-00004-of-00004.safetensors",
+    "visual.blocks.16.mlp.linear_fc2.bias": "model-00004-of-00004.safetensors",
+    "visual.blocks.16.mlp.linear_fc2.weight": "model-00004-of-00004.safetensors",
+    "visual.blocks.16.norm1.bias": "model-00004-of-00004.safetensors",
+    "visual.blocks.16.norm1.weight": "model-00004-of-00004.safetensors",
+    "visual.blocks.16.norm2.bias": "model-00004-of-00004.safetensors",
+    "visual.blocks.16.norm2.weight": "model-00004-of-00004.safetensors",
+    "visual.blocks.17.attn.proj.bias": "model-00004-of-00004.safetensors",
+    "visual.blocks.17.attn.proj.weight": "model-00004-of-00004.safetensors",
+    "visual.blocks.17.attn.qkv.bias": "model-00004-of-00004.safetensors",
+    "visual.blocks.17.attn.qkv.weight": "model-00004-of-00004.safetensors",
+    "visual.blocks.17.mlp.linear_fc1.bias": "model-00004-of-00004.safetensors",
+    "visual.blocks.17.mlp.linear_fc1.weight": "model-00004-of-00004.safetensors",
+    "visual.blocks.17.mlp.linear_fc2.bias": "model-00004-of-00004.safetensors",
+    "visual.blocks.17.mlp.linear_fc2.weight": "model-00004-of-00004.safetensors",
+    "visual.blocks.17.norm1.bias": "model-00004-of-00004.safetensors",
+    "visual.blocks.17.norm1.weight": "model-00004-of-00004.safetensors",
+    "visual.blocks.17.norm2.bias": "model-00004-of-00004.safetensors",
+    "visual.blocks.17.norm2.weight": "model-00004-of-00004.safetensors",
+    "visual.blocks.18.attn.proj.bias": "model-00004-of-00004.safetensors",
+    "visual.blocks.18.attn.proj.weight": "model-00004-of-00004.safetensors",
+    "visual.blocks.18.attn.qkv.bias": "model-00004-of-00004.safetensors",
+    "visual.blocks.18.attn.qkv.weight": "model-00004-of-00004.safetensors",
+    "visual.blocks.18.mlp.linear_fc1.bias": "model-00004-of-00004.safetensors",
+    "visual.blocks.18.mlp.linear_fc1.weight": "model-00004-of-00004.safetensors",
+    "visual.blocks.18.mlp.linear_fc2.bias": "model-00004-of-00004.safetensors",
+    "visual.blocks.18.mlp.linear_fc2.weight": "model-00004-of-00004.safetensors",
+    "visual.blocks.18.norm1.bias": "model-00004-of-00004.safetensors",
+    "visual.blocks.18.norm1.weight": "model-00004-of-00004.safetensors",
+    "visual.blocks.18.norm2.bias": "model-00004-of-00004.safetensors",
+    "visual.blocks.18.norm2.weight": "model-00004-of-00004.safetensors",
+    "visual.blocks.19.attn.proj.bias": "model-00004-of-00004.safetensors",
+    "visual.blocks.19.attn.proj.weight": "model-00004-of-00004.safetensors",
+    "visual.blocks.19.attn.qkv.bias": "model-00004-of-00004.safetensors",
+    "visual.blocks.19.attn.qkv.weight": "model-00004-of-00004.safetensors",
+    "visual.blocks.19.mlp.linear_fc1.bias": "model-00004-of-00004.safetensors",
+    "visual.blocks.19.mlp.linear_fc1.weight": "model-00004-of-00004.safetensors",
+    "visual.blocks.19.mlp.linear_fc2.bias": "model-00004-of-00004.safetensors",
+    "visual.blocks.19.mlp.linear_fc2.weight": "model-00004-of-00004.safetensors",
+    "visual.blocks.19.norm1.bias": "model-00004-of-00004.safetensors",
+    "visual.blocks.19.norm1.weight": "model-00004-of-00004.safetensors",
+    "visual.blocks.19.norm2.bias": "model-00004-of-00004.safetensors",
+    "visual.blocks.19.norm2.weight": "model-00004-of-00004.safetensors",
+    "visual.blocks.2.attn.proj.bias": "model-00004-of-00004.safetensors",
+    "visual.blocks.2.attn.proj.weight": "model-00004-of-00004.safetensors",
+    "visual.blocks.2.attn.qkv.bias": "model-00004-of-00004.safetensors",
+    "visual.blocks.2.attn.qkv.weight": "model-00004-of-00004.safetensors",
+    "visual.blocks.2.mlp.linear_fc1.bias": "model-00004-of-00004.safetensors",
+    "visual.blocks.2.mlp.linear_fc1.weight": "model-00004-of-00004.safetensors",
+    "visual.blocks.2.mlp.linear_fc2.bias": "model-00004-of-00004.safetensors",
+    "visual.blocks.2.mlp.linear_fc2.weight": "model-00004-of-00004.safetensors",
+    "visual.blocks.2.norm1.bias": "model-00004-of-00004.safetensors",
+    "visual.blocks.2.norm1.weight": "model-00004-of-00004.safetensors",
+    "visual.blocks.2.norm2.bias": "model-00004-of-00004.safetensors",
+    "visual.blocks.2.norm2.weight": "model-00004-of-00004.safetensors",
+    "visual.blocks.20.attn.proj.bias": "model-00004-of-00004.safetensors",
+    "visual.blocks.20.attn.proj.weight": "model-00004-of-00004.safetensors",
+    "visual.blocks.20.attn.qkv.bias": "model-00004-of-00004.safetensors",
+    "visual.blocks.20.attn.qkv.weight": "model-00004-of-00004.safetensors",
+    "visual.blocks.20.mlp.linear_fc1.bias": "model-00004-of-00004.safetensors",
+    "visual.blocks.20.mlp.linear_fc1.weight": "model-00004-of-00004.safetensors",
+    "visual.blocks.20.mlp.linear_fc2.bias": "model-00004-of-00004.safetensors",
+    "visual.blocks.20.mlp.linear_fc2.weight": "model-00004-of-00004.safetensors",
+    "visual.blocks.20.norm1.bias": "model-00004-of-00004.safetensors",
+    "visual.blocks.20.norm1.weight": "model-00004-of-00004.safetensors",
+    "visual.blocks.20.norm2.bias": "model-00004-of-00004.safetensors",
+    "visual.blocks.20.norm2.weight": "model-00004-of-00004.safetensors",
+    "visual.blocks.21.attn.proj.bias": "model-00004-of-00004.safetensors",
+    "visual.blocks.21.attn.proj.weight": "model-00004-of-00004.safetensors",
+    "visual.blocks.21.attn.qkv.bias": "model-00004-of-00004.safetensors",
+    "visual.blocks.21.attn.qkv.weight": "model-00004-of-00004.safetensors",
+    "visual.blocks.21.mlp.linear_fc1.bias": "model-00004-of-00004.safetensors",
+    "visual.blocks.21.mlp.linear_fc1.weight": "model-00004-of-00004.safetensors",
+    "visual.blocks.21.mlp.linear_fc2.bias": "model-00004-of-00004.safetensors",
+    "visual.blocks.21.mlp.linear_fc2.weight": "model-00004-of-00004.safetensors",
+    "visual.blocks.21.norm1.bias": "model-00004-of-00004.safetensors",
+    "visual.blocks.21.norm1.weight": "model-00004-of-00004.safetensors",
+    "visual.blocks.21.norm2.bias": "model-00004-of-00004.safetensors",
+    "visual.blocks.21.norm2.weight": "model-00004-of-00004.safetensors",
+    "visual.blocks.22.attn.proj.bias": "model-00004-of-00004.safetensors",
+    "visual.blocks.22.attn.proj.weight": "model-00004-of-00004.safetensors",
+    "visual.blocks.22.attn.qkv.bias": "model-00004-of-00004.safetensors",
+    "visual.blocks.22.attn.qkv.weight": "model-00004-of-00004.safetensors",
+    "visual.blocks.22.mlp.linear_fc1.bias": "model-00004-of-00004.safetensors",
+    "visual.blocks.22.mlp.linear_fc1.weight": "model-00004-of-00004.safetensors",
+    "visual.blocks.22.mlp.linear_fc2.bias": "model-00004-of-00004.safetensors",
+    "visual.blocks.22.mlp.linear_fc2.weight": "model-00004-of-00004.safetensors",
+    "visual.blocks.22.norm1.bias": "model-00004-of-00004.safetensors",
+    "visual.blocks.22.norm1.weight": "model-00004-of-00004.safetensors",
+    "visual.blocks.22.norm2.bias": "model-00004-of-00004.safetensors",
+    "visual.blocks.22.norm2.weight": "model-00004-of-00004.safetensors",
+    "visual.blocks.23.attn.proj.bias": "model-00004-of-00004.safetensors",
+    "visual.blocks.23.attn.proj.weight": "model-00004-of-00004.safetensors",
+    "visual.blocks.23.attn.qkv.bias": "model-00004-of-00004.safetensors",
+    "visual.blocks.23.attn.qkv.weight": "model-00004-of-00004.safetensors",
+    "visual.blocks.23.mlp.linear_fc1.bias": "model-00004-of-00004.safetensors",
+    "visual.blocks.23.mlp.linear_fc1.weight": "model-00004-of-00004.safetensors",
+    "visual.blocks.23.mlp.linear_fc2.bias": "model-00004-of-00004.safetensors",
+    "visual.blocks.23.mlp.linear_fc2.weight": "model-00004-of-00004.safetensors",
+    "visual.blocks.23.norm1.bias": "model-00004-of-00004.safetensors",
+    "visual.blocks.23.norm1.weight": "model-00004-of-00004.safetensors",
+    "visual.blocks.23.norm2.bias": "model-00004-of-00004.safetensors",
+    "visual.blocks.23.norm2.weight": "model-00004-of-00004.safetensors",
+    "visual.blocks.24.attn.proj.bias": "model-00004-of-00004.safetensors",
+    "visual.blocks.24.attn.proj.weight": "model-00004-of-00004.safetensors",
+    "visual.blocks.24.attn.qkv.bias": "model-00004-of-00004.safetensors",
+    "visual.blocks.24.attn.qkv.weight": "model-00004-of-00004.safetensors",
+    "visual.blocks.24.mlp.linear_fc1.bias": "model-00004-of-00004.safetensors",
+    "visual.blocks.24.mlp.linear_fc1.weight": "model-00004-of-00004.safetensors",
+    "visual.blocks.24.mlp.linear_fc2.bias": "model-00004-of-00004.safetensors",
+    "visual.blocks.24.mlp.linear_fc2.weight": "model-00004-of-00004.safetensors",
+    "visual.blocks.24.norm1.bias": "model-00004-of-00004.safetensors",
+    "visual.blocks.24.norm1.weight": "model-00004-of-00004.safetensors",
+    "visual.blocks.24.norm2.bias": "model-00004-of-00004.safetensors",
+    "visual.blocks.24.norm2.weight": "model-00004-of-00004.safetensors",
+    "visual.blocks.25.attn.proj.bias": "model-00004-of-00004.safetensors",
+    "visual.blocks.25.attn.proj.weight": "model-00004-of-00004.safetensors",
+    "visual.blocks.25.attn.qkv.bias": "model-00004-of-00004.safetensors",
+    "visual.blocks.25.attn.qkv.weight": "model-00004-of-00004.safetensors",
+    "visual.blocks.25.mlp.linear_fc1.bias": "model-00004-of-00004.safetensors",
+    "visual.blocks.25.mlp.linear_fc1.weight": "model-00004-of-00004.safetensors",
+    "visual.blocks.25.mlp.linear_fc2.bias": "model-00004-of-00004.safetensors",
+    "visual.blocks.25.mlp.linear_fc2.weight": "model-00004-of-00004.safetensors",
+    "visual.blocks.25.norm1.bias": "model-00004-of-00004.safetensors",
+    "visual.blocks.25.norm1.weight": "model-00004-of-00004.safetensors",
+    "visual.blocks.25.norm2.bias": "model-00004-of-00004.safetensors",
+    "visual.blocks.25.norm2.weight": "model-00004-of-00004.safetensors",
+    "visual.blocks.26.attn.proj.bias": "model-00004-of-00004.safetensors",
+    "visual.blocks.26.attn.proj.weight": "model-00004-of-00004.safetensors",
+    "visual.blocks.26.attn.qkv.bias": "model-00004-of-00004.safetensors",
+    "visual.blocks.26.attn.qkv.weight": "model-00004-of-00004.safetensors",
+    "visual.blocks.26.mlp.linear_fc1.bias": "model-00004-of-00004.safetensors",
+    "visual.blocks.26.mlp.linear_fc1.weight": "model-00004-of-00004.safetensors",
+    "visual.blocks.26.mlp.linear_fc2.bias": "model-00004-of-00004.safetensors",
+    "visual.blocks.26.mlp.linear_fc2.weight": "model-00004-of-00004.safetensors",
+    "visual.blocks.26.norm1.bias": "model-00004-of-00004.safetensors",
+    "visual.blocks.26.norm1.weight": "model-00004-of-00004.safetensors",
+    "visual.blocks.26.norm2.bias": "model-00004-of-00004.safetensors",
+    "visual.blocks.26.norm2.weight": "model-00004-of-00004.safetensors",
+    "visual.blocks.3.attn.proj.bias": "model-00004-of-00004.safetensors",
+    "visual.blocks.3.attn.proj.weight": "model-00004-of-00004.safetensors",
+    "visual.blocks.3.attn.qkv.bias": "model-00004-of-00004.safetensors",
+    "visual.blocks.3.attn.qkv.weight": "model-00004-of-00004.safetensors",
+    "visual.blocks.3.mlp.linear_fc1.bias": "model-00004-of-00004.safetensors",
+    "visual.blocks.3.mlp.linear_fc1.weight": "model-00004-of-00004.safetensors",
+    "visual.blocks.3.mlp.linear_fc2.bias": "model-00004-of-00004.safetensors",
+    "visual.blocks.3.mlp.linear_fc2.weight": "model-00004-of-00004.safetensors",
+    "visual.blocks.3.norm1.bias": "model-00004-of-00004.safetensors",
+    "visual.blocks.3.norm1.weight": "model-00004-of-00004.safetensors",
+    "visual.blocks.3.norm2.bias": "model-00004-of-00004.safetensors",
+    "visual.blocks.3.norm2.weight": "model-00004-of-00004.safetensors",
+    "visual.blocks.4.attn.proj.bias": "model-00004-of-00004.safetensors",
+    "visual.blocks.4.attn.proj.weight": "model-00004-of-00004.safetensors",
+    "visual.blocks.4.attn.qkv.bias": "model-00004-of-00004.safetensors",
+    "visual.blocks.4.attn.qkv.weight": "model-00004-of-00004.safetensors",
+    "visual.blocks.4.mlp.linear_fc1.bias": "model-00004-of-00004.safetensors",
+    "visual.blocks.4.mlp.linear_fc1.weight": "model-00004-of-00004.safetensors",
+    "visual.blocks.4.mlp.linear_fc2.bias": "model-00004-of-00004.safetensors",
+    "visual.blocks.4.mlp.linear_fc2.weight": "model-00004-of-00004.safetensors",
+    "visual.blocks.4.norm1.bias": "model-00004-of-00004.safetensors",
+    "visual.blocks.4.norm1.weight": "model-00004-of-00004.safetensors",
+    "visual.blocks.4.norm2.bias": "model-00004-of-00004.safetensors",
+    "visual.blocks.4.norm2.weight": "model-00004-of-00004.safetensors",
+    "visual.blocks.5.attn.proj.bias": "model-00004-of-00004.safetensors",
+    "visual.blocks.5.attn.proj.weight": "model-00004-of-00004.safetensors",
+    "visual.blocks.5.attn.qkv.bias": "model-00004-of-00004.safetensors",
+    "visual.blocks.5.attn.qkv.weight": "model-00004-of-00004.safetensors",
+    "visual.blocks.5.mlp.linear_fc1.bias": "model-00004-of-00004.safetensors",
+    "visual.blocks.5.mlp.linear_fc1.weight": "model-00004-of-00004.safetensors",
+    "visual.blocks.5.mlp.linear_fc2.bias": "model-00004-of-00004.safetensors",
+    "visual.blocks.5.mlp.linear_fc2.weight": "model-00004-of-00004.safetensors",
+    "visual.blocks.5.norm1.bias": "model-00004-of-00004.safetensors",
+    "visual.blocks.5.norm1.weight": "model-00004-of-00004.safetensors",
+    "visual.blocks.5.norm2.bias": "model-00004-of-00004.safetensors",
+    "visual.blocks.5.norm2.weight": "model-00004-of-00004.safetensors",
+    "visual.blocks.6.attn.proj.bias": "model-00004-of-00004.safetensors",
+    "visual.blocks.6.attn.proj.weight": "model-00004-of-00004.safetensors",
+    "visual.blocks.6.attn.qkv.bias": "model-00004-of-00004.safetensors",
+    "visual.blocks.6.attn.qkv.weight": "model-00004-of-00004.safetensors",
+    "visual.blocks.6.mlp.linear_fc1.bias": "model-00004-of-00004.safetensors",
+    "visual.blocks.6.mlp.linear_fc1.weight": "model-00004-of-00004.safetensors",
+    "visual.blocks.6.mlp.linear_fc2.bias": "model-00004-of-00004.safetensors",
+    "visual.blocks.6.mlp.linear_fc2.weight": "model-00004-of-00004.safetensors",
+    "visual.blocks.6.norm1.bias": "model-00004-of-00004.safetensors",
+    "visual.blocks.6.norm1.weight": "model-00004-of-00004.safetensors",
+    "visual.blocks.6.norm2.bias": "model-00004-of-00004.safetensors",
+    "visual.blocks.6.norm2.weight": "model-00004-of-00004.safetensors",
+    "visual.blocks.7.attn.proj.bias": "model-00004-of-00004.safetensors",
+    "visual.blocks.7.attn.proj.weight": "model-00004-of-00004.safetensors",
+    "visual.blocks.7.attn.qkv.bias": "model-00004-of-00004.safetensors",
+    "visual.blocks.7.attn.qkv.weight": "model-00004-of-00004.safetensors",
+    "visual.blocks.7.mlp.linear_fc1.bias": "model-00004-of-00004.safetensors",
+    "visual.blocks.7.mlp.linear_fc1.weight": "model-00004-of-00004.safetensors",
+    "visual.blocks.7.mlp.linear_fc2.bias": "model-00004-of-00004.safetensors",
+    "visual.blocks.7.mlp.linear_fc2.weight": "model-00004-of-00004.safetensors",
+    "visual.blocks.7.norm1.bias": "model-00004-of-00004.safetensors",
+    "visual.blocks.7.norm1.weight": "model-00004-of-00004.safetensors",
+    "visual.blocks.7.norm2.bias": "model-00004-of-00004.safetensors",
+    "visual.blocks.7.norm2.weight": "model-00004-of-00004.safetensors",
+    "visual.blocks.8.attn.proj.bias": "model-00004-of-00004.safetensors",
+    "visual.blocks.8.attn.proj.weight": "model-00004-of-00004.safetensors",
+    "visual.blocks.8.attn.qkv.bias": "model-00004-of-00004.safetensors",
+    "visual.blocks.8.attn.qkv.weight": "model-00004-of-00004.safetensors",
+    "visual.blocks.8.mlp.linear_fc1.bias": "model-00004-of-00004.safetensors",
+    "visual.blocks.8.mlp.linear_fc1.weight": "model-00004-of-00004.safetensors",
+    "visual.blocks.8.mlp.linear_fc2.bias": "model-00004-of-00004.safetensors",
+    "visual.blocks.8.mlp.linear_fc2.weight": "model-00004-of-00004.safetensors",
+    "visual.blocks.8.norm1.bias": "model-00004-of-00004.safetensors",
+    "visual.blocks.8.norm1.weight": "model-00004-of-00004.safetensors",
+    "visual.blocks.8.norm2.bias": "model-00004-of-00004.safetensors",
+    "visual.blocks.8.norm2.weight": "model-00004-of-00004.safetensors",
+    "visual.blocks.9.attn.proj.bias": "model-00004-of-00004.safetensors",
+    "visual.blocks.9.attn.proj.weight": "model-00004-of-00004.safetensors",
+    "visual.blocks.9.attn.qkv.bias": "model-00004-of-00004.safetensors",
+    "visual.blocks.9.attn.qkv.weight": "model-00004-of-00004.safetensors",
+    "visual.blocks.9.mlp.linear_fc1.bias": "model-00004-of-00004.safetensors",
+    "visual.blocks.9.mlp.linear_fc1.weight": "model-00004-of-00004.safetensors",
+    "visual.blocks.9.mlp.linear_fc2.bias": "model-00004-of-00004.safetensors",
+    "visual.blocks.9.mlp.linear_fc2.weight": "model-00004-of-00004.safetensors",
+    "visual.blocks.9.norm1.bias": "model-00004-of-00004.safetensors",
+    "visual.blocks.9.norm1.weight": "model-00004-of-00004.safetensors",
+    "visual.blocks.9.norm2.bias": "model-00004-of-00004.safetensors",
+    "visual.blocks.9.norm2.weight": "model-00004-of-00004.safetensors",
+    "visual.merger.linear_fc1.bias": "model-00004-of-00004.safetensors",
+    "visual.merger.linear_fc1.weight": "model-00004-of-00004.safetensors",
+    "visual.merger.linear_fc2.bias": "model-00004-of-00004.safetensors",
+    "visual.merger.linear_fc2.weight": "model-00004-of-00004.safetensors",
+    "visual.merger.norm.bias": "model-00004-of-00004.safetensors",
+    "visual.merger.norm.weight": "model-00004-of-00004.safetensors",
+    "visual.patch_embed.proj.bias": "model-00004-of-00004.safetensors",
+    "visual.patch_embed.proj.weight": "model-00004-of-00004.safetensors",
+    "visual.pos_embed.weight": "model-00004-of-00004.safetensors"
+  }
+}

tokenizer/chat_template.jinja

@@ -0,0 +1,154 @@
+{%- set image_count = namespace(value=0) %}
+{%- set video_count = namespace(value=0) %}
+{%- macro render_content(content, do_vision_count, is_system_content=false) %}
+    {%- if content is string %}
+        {{- content }}
+    {%- elif content is iterable and content is not mapping %}
+        {%- for item in content %}
+            {%- if 'image' in item or 'image_url' in item or item.type == 'image' %}
+                {%- if is_system_content %}
+                    {{- raise_exception('System message cannot contain images.') }}
+                {%- endif %}
+                {%- if do_vision_count %}
+                    {%- set image_count.value = image_count.value + 1 %}
+                {%- endif %}
+                {%- if add_vision_id %}
+                    {{- 'Picture ' ~ image_count.value ~ ': ' }}
+                {%- endif %}
+                {{- '<|vision_start|><|image_pad|><|vision_end|>' }}
+            {%- elif 'video' in item or item.type == 'video' %}
+                {%- if is_system_content %}
+                    {{- raise_exception('System message cannot contain videos.') }}
+                {%- endif %}
+                {%- if do_vision_count %}
+                    {%- set video_count.value = video_count.value + 1 %}
+                {%- endif %}
+                {%- if add_vision_id %}
+                    {{- 'Video ' ~ video_count.value ~ ': ' }}
+                {%- endif %}
+                {{- '<|vision_start|><|video_pad|><|vision_end|>' }}
+            {%- elif 'text' in item %}
+                {{- item.text }}
+            {%- else %}
+                {{- raise_exception('Unexpected item type in content.') }}
+            {%- endif %}
+        {%- endfor %}
+    {%- elif content is none or content is undefined %}
+        {{- '' }}
+    {%- else %}
+        {{- raise_exception('Unexpected content type.') }}
+    {%- endif %}
+{%- endmacro %}
+{%- if not messages %}
+    {{- raise_exception('No messages provided.') }}
+{%- endif %}
+{%- if tools and tools is iterable and tools is not mapping %}
+    {{- '<|im_start|>system\n' }}
+    {{- "# Tools\n\nYou have access to the following functions:\n\n<tools>" }}
+    {%- for tool in tools %}
+        {{- "\n" }}
+        {{- tool | tojson }}
+    {%- endfor %}
+    {{- "\n</tools>" }}
+    {{- '\n\nIf you choose to call a function ONLY reply in the following format with NO suffix:\n\n<tool_call>\n<function=example_function_name>\n<parameter=example_parameter_1>\nvalue_1\n</parameter>\n<parameter=example_parameter_2>\nThis is the value for the second parameter\nthat can span\nmultiple lines\n</parameter>\n</function>\n</tool_call>\n\n<IMPORTANT>\nReminder:\n- Function calls MUST follow the specified format: an inner <function=...></function> block must be nested within <tool_call></tool_call> XML tags\n- Required parameters MUST be specified\n- You may provide optional reasoning for your function call in natural language BEFORE the function call, but NOT after\n- If there is no function call available, answer the question like normal with your current knowledge and do not tell the user about function calls\n</IMPORTANT>' }}
+    {%- if messages[0].role == 'system' %}
+        {%- set content = render_content(messages[0].content, false, true)|trim %}
+        {%- if content %}
+            {{- '\n\n' + content }}
+        {%- endif %}
+    {%- endif %}
+    {{- '<|im_end|>\n' }}
+{%- else %}
+    {%- if messages[0].role == 'system' %}
+        {%- set content = render_content(messages[0].content, false, true)|trim %}
+        {{- '<|im_start|>system\n' + content + '<|im_end|>\n' }}
+    {%- endif %}
+{%- endif %}
+{%- set ns = namespace(multi_step_tool=true, last_query_index=messages|length - 1) %}
+{%- for message in messages[::-1] %}
+    {%- set index = (messages|length - 1) - loop.index0 %}
+    {%- if ns.multi_step_tool and message.role == "user" %}
+        {%- set content = render_content(message.content, false)|trim %}
+        {%- if not(content.startswith('<tool_response>') and content.endswith('</tool_response>')) %}
+            {%- set ns.multi_step_tool = false %}
+            {%- set ns.last_query_index = index %}
+        {%- endif %}
+    {%- endif %}
+{%- endfor %}
+{%- if ns.multi_step_tool %}
+    {{- raise_exception('No user query found in messages.') }}
+{%- endif %}
+{%- for message in messages %}
+    {%- set content = render_content(message.content, true)|trim %}
+    {%- if message.role == "system" %}
+        {%- if not loop.first %}
+            {{- raise_exception('System message must be at the beginning.') }}
+        {%- endif %}
+    {%- elif message.role == "user" %}
+        {{- '<|im_start|>' + message.role + '\n' + content + '<|im_end|>' + '\n' }}
+    {%- elif message.role == "assistant" %}
+        {%- set reasoning_content = '' %}
+        {%- if message.reasoning_content is string %}
+            {%- set reasoning_content = message.reasoning_content %}
+        {%- else %}
+            {%- if '</think>' in content %}
+                {%- set reasoning_content = content.split('</think>')[0].rstrip('\n').split('<think>')[-1].lstrip('\n') %}
+                {%- set content = content.split('</think>')[-1].lstrip('\n') %}
+            {%- endif %}
+        {%- endif %}
+        {%- set reasoning_content = reasoning_content|trim %}
+        {%- if loop.index0 > ns.last_query_index %}
+            {{- '<|im_start|>' + message.role + '\n<think>\n' + reasoning_content + '\n</think>\n\n' + content }}
+        {%- else %}
+            {{- '<|im_start|>' + message.role + '\n' + content }}
+        {%- endif %}
+        {%- if message.tool_calls and message.tool_calls is iterable and message.tool_calls is not mapping %}
+            {%- for tool_call in message.tool_calls %}
+                {%- if tool_call.function is defined %}
+                    {%- set tool_call = tool_call.function %}
+                {%- endif %}
+                {%- if loop.first %}
+                    {%- if content|trim %}
+                        {{- '\n\n<tool_call>\n<function=' + tool_call.name + '>\n' }}
+                    {%- else %}
+                        {{- '<tool_call>\n<function=' + tool_call.name + '>\n' }}
+                    {%- endif %}
+                {%- else %}
+                    {{- '\n<tool_call>\n<function=' + tool_call.name + '>\n' }}
+                {%- endif %}
+                {%- if tool_call.arguments is defined %}
+                    {%- for args_name, args_value in tool_call.arguments|items %}
+                        {{- '<parameter=' + args_name + '>\n' }}
+                        {%- set args_value = args_value | tojson | safe if args_value is mapping or (args_value is sequence and args_value is not string) else args_value | string %}
+                        {{- args_value }}
+                        {{- '\n</parameter>\n' }}
+                    {%- endfor %}
+                {%- endif %}
+                {{- '</function>\n</tool_call>' }}
+            {%- endfor %}
+        {%- endif %}
+        {{- '<|im_end|>\n' }}
+    {%- elif message.role == "tool" %}
+        {%- if loop.previtem and loop.previtem.role != "tool" %}
+            {{- '<|im_start|>user' }}
+        {%- endif %}
+        {{- '\n<tool_response>\n' }}
+        {{- content }}
+        {{- '\n</tool_response>' }}
+        {%- if not loop.last and loop.nextitem.role != "tool" %}
+            {{- '<|im_end|>\n' }}
+        {%- elif loop.last %}
+            {{- '<|im_end|>\n' }}
+        {%- endif %}
+    {%- else %}
+        {{- raise_exception('Unexpected message role.') }}
+    {%- endif %}
+{%- endfor %}
+{%- if add_generation_prompt %}
+    {{- '<|im_start|>assistant\n' }}
+    {%- if enable_thinking is defined and enable_thinking is false %}
+        {{- '<think>\n\n</think>\n\n' }}
+    {%- else %}
+        {{- '<think>\n' }}
+    {%- endif %}
+{%- endif %}

tokenizer/tokenizer.json

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

tokenizer/tokenizer_config.json

@@ -0,0 +1,32 @@
+{
+  "add_prefix_space": false,
+  "audio_bos_token": "<|audio_start|>",
+  "audio_eos_token": "<|audio_end|>",
+  "audio_token": "<|audio_pad|>",
+  "backend": "tokenizers",
+  "bos_token": null,
+  "clean_up_tokenization_spaces": false,
+  "eos_token": "<|im_end|>",
+  "errors": "replace",
+  "image_token": "<|image_pad|>",
+  "is_local": true,
+  "local_files_only": true,
+  "model_max_length": 262144,
+  "model_specific_special_tokens": {
+    "audio_bos_token": "<|audio_start|>",
+    "audio_eos_token": "<|audio_end|>",
+    "audio_token": "<|audio_pad|>",
+    "image_token": "<|image_pad|>",
+    "video_token": "<|video_pad|>",
+    "vision_bos_token": "<|vision_start|>",
+    "vision_eos_token": "<|vision_end|>"
+  },
+  "pad_token": "<|endoftext|>",
+  "pretokenize_regex": "(?i:'s|'t|'re|'ve|'m|'ll|'d)|[^\\r\\n\\p{L}\\p{N}]?[\\p{L}\\p{M}]+|\\p{N}| ?[^\\s\\p{L}\\p{M}\\p{N}]+[\\r\\n]*|\\s*[\\r\\n]+|\\s+(?!\\S)|\\s+",
+  "split_special_tokens": false,
+  "tokenizer_class": "Qwen2Tokenizer",
+  "unk_token": null,
+  "video_token": "<|video_pad|>",
+  "vision_bos_token": "<|vision_start|>",
+  "vision_eos_token": "<|vision_end|>"
+}

transformer/config.json

@@ -0,0 +1,12 @@
+{
+  "_class_name": "MOSSImageTransformer",
+  "_diffusers_version": "0.38.0",
+  "dim": 3072,
+  "ffn_dim": 8192,
+  "freq_dim": 256,
+  "latent_channels": 128,
+  "num_heads": 24,
+  "num_layers": 42,
+  "rope_base": 10000.0,
+  "text_dim": 4096
+}

transformer/diffusion_pytorch_model-00001-of-00003.safetensors

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

transformer/diffusion_pytorch_model-00002-of-00003.safetensors

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

transformer/diffusion_pytorch_model-00003-of-00003.safetensors

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

transformer/diffusion_pytorch_model.safetensors.index.json

@@ -0,0 +1,470 @@
+{
+  "metadata": {
+    "total_size": 10164617280
+  },
+  "weight_map": {
+    "final_linear_video.weight": "diffusion_pytorch_model-00003-of-00003.safetensors",
+    "final_norm_video.weight": "diffusion_pytorch_model-00003-of-00003.safetensors",
+    "layers.0.k_norm.weight": "diffusion_pytorch_model-00001-of-00003.safetensors",
+    "layers.0.norm1.weight": "diffusion_pytorch_model-00001-of-00003.safetensors",
+    "layers.0.norm2.weight": "diffusion_pytorch_model-00001-of-00003.safetensors",
+    "layers.0.proj.ffn_down.weight": "diffusion_pytorch_model-00001-of-00003.safetensors",
+    "layers.0.proj.ffn_up_gate.weight": "diffusion_pytorch_model-00001-of-00003.safetensors",
+    "layers.0.proj.k.weight": "diffusion_pytorch_model-00001-of-00003.safetensors",
+    "layers.0.proj.o.weight": "diffusion_pytorch_model-00001-of-00003.safetensors",
+    "layers.0.proj.q.weight": "diffusion_pytorch_model-00001-of-00003.safetensors",
+    "layers.0.proj.v.weight": "diffusion_pytorch_model-00001-of-00003.safetensors",
+    "layers.0.q_norm.weight": "diffusion_pytorch_model-00001-of-00003.safetensors",
+    "layers.1.k_norm.weight": "diffusion_pytorch_model-00001-of-00003.safetensors",
+    "layers.1.norm1.weight": "diffusion_pytorch_model-00001-of-00003.safetensors",
+    "layers.1.norm2.weight": "diffusion_pytorch_model-00001-of-00003.safetensors",
+    "layers.1.proj.ffn_down.weight": "diffusion_pytorch_model-00001-of-00003.safetensors",
+    "layers.1.proj.ffn_up_gate.weight": "diffusion_pytorch_model-00001-of-00003.safetensors",
+    "layers.1.proj.k.weight": "diffusion_pytorch_model-00001-of-00003.safetensors",
+    "layers.1.proj.o.weight": "diffusion_pytorch_model-00001-of-00003.safetensors",
+    "layers.1.proj.q.weight": "diffusion_pytorch_model-00001-of-00003.safetensors",
+    "layers.1.proj.v.weight": "diffusion_pytorch_model-00001-of-00003.safetensors",
+    "layers.1.q_norm.weight": "diffusion_pytorch_model-00001-of-00003.safetensors",
+    "layers.10.k_norm.weight": "diffusion_pytorch_model-00001-of-00003.safetensors",
+    "layers.10.norm1.weight": "diffusion_pytorch_model-00001-of-00003.safetensors",
+    "layers.10.norm2.weight": "diffusion_pytorch_model-00001-of-00003.safetensors",
+    "layers.10.proj.ffn_down.weight": "diffusion_pytorch_model-00001-of-00003.safetensors",
+    "layers.10.proj.ffn_up_gate.weight": "diffusion_pytorch_model-00001-of-00003.safetensors",
+    "layers.10.proj.k.weight": "diffusion_pytorch_model-00001-of-00003.safetensors",
+    "layers.10.proj.o.weight": "diffusion_pytorch_model-00001-of-00003.safetensors",
+    "layers.10.proj.q.weight": "diffusion_pytorch_model-00001-of-00003.safetensors",
+    "layers.10.proj.v.weight": "diffusion_pytorch_model-00001-of-00003.safetensors",
+    "layers.10.q_norm.weight": "diffusion_pytorch_model-00001-of-00003.safetensors",
+    "layers.11.k_norm.weight": "diffusion_pytorch_model-00001-of-00003.safetensors",
+    "layers.11.norm1.weight": "diffusion_pytorch_model-00001-of-00003.safetensors",
+    "layers.11.norm2.weight": "diffusion_pytorch_model-00001-of-00003.safetensors",
+    "layers.11.proj.ffn_down.weight": "diffusion_pytorch_model-00001-of-00003.safetensors",
+    "layers.11.proj.ffn_up_gate.weight": "diffusion_pytorch_model-00001-of-00003.safetensors",
+    "layers.11.proj.k.weight": "diffusion_pytorch_model-00001-of-00003.safetensors",
+    "layers.11.proj.o.weight": "diffusion_pytorch_model-00001-of-00003.safetensors",
+    "layers.11.proj.q.weight": "diffusion_pytorch_model-00001-of-00003.safetensors",
+    "layers.11.proj.v.weight": "diffusion_pytorch_model-00001-of-00003.safetensors",
+    "layers.11.q_norm.weight": "diffusion_pytorch_model-00001-of-00003.safetensors",
+    "layers.12.k_norm.weight": "diffusion_pytorch_model-00001-of-00003.safetensors",
+    "layers.12.norm1.weight": "diffusion_pytorch_model-00001-of-00003.safetensors",
+    "layers.12.norm2.weight": "diffusion_pytorch_model-00001-of-00003.safetensors",
+    "layers.12.proj.ffn_down.weight": "diffusion_pytorch_model-00001-of-00003.safetensors",
+    "layers.12.proj.ffn_up_gate.weight": "diffusion_pytorch_model-00001-of-00003.safetensors",
+    "layers.12.proj.k.weight": "diffusion_pytorch_model-00001-of-00003.safetensors",
+    "layers.12.proj.o.weight": "diffusion_pytorch_model-00001-of-00003.safetensors",
+    "layers.12.proj.q.weight": "diffusion_pytorch_model-00001-of-00003.safetensors",
+    "layers.12.proj.v.weight": "diffusion_pytorch_model-00001-of-00003.safetensors",
+    "layers.12.q_norm.weight": "diffusion_pytorch_model-00001-of-00003.safetensors",
+    "layers.13.k_norm.weight": "diffusion_pytorch_model-00001-of-00003.safetensors",
+    "layers.13.norm1.weight": "diffusion_pytorch_model-00001-of-00003.safetensors",
+    "layers.13.norm2.weight": "diffusion_pytorch_model-00001-of-00003.safetensors",
+    "layers.13.proj.ffn_down.weight": "diffusion_pytorch_model-00001-of-00003.safetensors",
+    "layers.13.proj.ffn_up_gate.weight": "diffusion_pytorch_model-00001-of-00003.safetensors",
+    "layers.13.proj.k.weight": "diffusion_pytorch_model-00001-of-00003.safetensors",
+    "layers.13.proj.o.weight": "diffusion_pytorch_model-00001-of-00003.safetensors",
+    "layers.13.proj.q.weight": "diffusion_pytorch_model-00001-of-00003.safetensors",
+    "layers.13.proj.v.weight": "diffusion_pytorch_model-00001-of-00003.safetensors",
+    "layers.13.q_norm.weight": "diffusion_pytorch_model-00001-of-00003.safetensors",
+    "layers.14.k_norm.weight": "diffusion_pytorch_model-00001-of-00003.safetensors",
+    "layers.14.norm1.weight": "diffusion_pytorch_model-00001-of-00003.safetensors",
+    "layers.14.norm2.weight": "diffusion_pytorch_model-00001-of-00003.safetensors",
+    "layers.14.proj.ffn_down.weight": "diffusion_pytorch_model-00001-of-00003.safetensors",
+    "layers.14.proj.ffn_up_gate.weight": "diffusion_pytorch_model-00001-of-00003.safetensors",
+    "layers.14.proj.k.weight": "diffusion_pytorch_model-00001-of-00003.safetensors",
+    "layers.14.proj.o.weight": "diffusion_pytorch_model-00001-of-00003.safetensors",
+    "layers.14.proj.q.weight": "diffusion_pytorch_model-00001-of-00003.safetensors",
+    "layers.14.proj.v.weight": "diffusion_pytorch_model-00001-of-00003.safetensors",
+    "layers.14.q_norm.weight": "diffusion_pytorch_model-00001-of-00003.safetensors",
+    "layers.15.k_norm.weight": "diffusion_pytorch_model-00001-of-00003.safetensors",
+    "layers.15.norm1.weight": "diffusion_pytorch_model-00001-of-00003.safetensors",
+    "layers.15.norm2.weight": "diffusion_pytorch_model-00001-of-00003.safetensors",
+    "layers.15.proj.ffn_down.weight": "diffusion_pytorch_model-00001-of-00003.safetensors",
+    "layers.15.proj.ffn_up_gate.weight": "diffusion_pytorch_model-00001-of-00003.safetensors",
+    "layers.15.proj.k.weight": "diffusion_pytorch_model-00001-of-00003.safetensors",
+    "layers.15.proj.o.weight": "diffusion_pytorch_model-00001-of-00003.safetensors",
+    "layers.15.proj.q.weight": "diffusion_pytorch_model-00001-of-00003.safetensors",
+    "layers.15.proj.v.weight": "diffusion_pytorch_model-00001-of-00003.safetensors",
+    "layers.15.q_norm.weight": "diffusion_pytorch_model-00001-of-00003.safetensors",
+    "layers.16.k_norm.weight": "diffusion_pytorch_model-00001-of-00003.safetensors",
+    "layers.16.norm1.weight": "diffusion_pytorch_model-00001-of-00003.safetensors",
+    "layers.16.norm2.weight": "diffusion_pytorch_model-00001-of-00003.safetensors",
+    "layers.16.proj.ffn_down.weight": "diffusion_pytorch_model-00001-of-00003.safetensors",
+    "layers.16.proj.ffn_up_gate.weight": "diffusion_pytorch_model-00001-of-00003.safetensors",
+    "layers.16.proj.k.weight": "diffusion_pytorch_model-00001-of-00003.safetensors",
+    "layers.16.proj.o.weight": "diffusion_pytorch_model-00001-of-00003.safetensors",
+    "layers.16.proj.q.weight": "diffusion_pytorch_model-00001-of-00003.safetensors",
+    "layers.16.proj.v.weight": "diffusion_pytorch_model-00001-of-00003.safetensors",
+    "layers.16.q_norm.weight": "diffusion_pytorch_model-00001-of-00003.safetensors",
+    "layers.17.k_norm.weight": "diffusion_pytorch_model-00001-of-00003.safetensors",
+    "layers.17.norm1.weight": "diffusion_pytorch_model-00001-of-00003.safetensors",
+    "layers.17.norm2.weight": "diffusion_pytorch_model-00001-of-00003.safetensors",
+    "layers.17.proj.ffn_down.weight": "diffusion_pytorch_model-00001-of-00003.safetensors",
+    "layers.17.proj.ffn_up_gate.weight": "diffusion_pytorch_model-00001-of-00003.safetensors",
+    "layers.17.proj.k.weight": "diffusion_pytorch_model-00001-of-00003.safetensors",
+    "layers.17.proj.o.weight": "diffusion_pytorch_model-00001-of-00003.safetensors",
+    "layers.17.proj.q.weight": "diffusion_pytorch_model-00001-of-00003.safetensors",
+    "layers.17.proj.v.weight": "diffusion_pytorch_model-00001-of-00003.safetensors",
+    "layers.17.q_norm.weight": "diffusion_pytorch_model-00001-of-00003.safetensors",
+    "layers.18.k_norm.weight": "diffusion_pytorch_model-00001-of-00003.safetensors",
+    "layers.18.norm1.weight": "diffusion_pytorch_model-00001-of-00003.safetensors",
+    "layers.18.norm2.weight": "diffusion_pytorch_model-00001-of-00003.safetensors",
+    "layers.18.proj.ffn_down.weight": "diffusion_pytorch_model-00001-of-00003.safetensors",
+    "layers.18.proj.ffn_up_gate.weight": "diffusion_pytorch_model-00001-of-00003.safetensors",
+    "layers.18.proj.k.weight": "diffusion_pytorch_model-00001-of-00003.safetensors",
+    "layers.18.proj.o.weight": "diffusion_pytorch_model-00001-of-00003.safetensors",
+    "layers.18.proj.q.weight": "diffusion_pytorch_model-00001-of-00003.safetensors",
+    "layers.18.proj.v.weight": "diffusion_pytorch_model-00001-of-00003.safetensors",
+    "layers.18.q_norm.weight": "diffusion_pytorch_model-00001-of-00003.safetensors",
+    "layers.19.k_norm.weight": "diffusion_pytorch_model-00001-of-00003.safetensors",
+    "layers.19.norm1.weight": "diffusion_pytorch_model-00001-of-00003.safetensors",
+    "layers.19.norm2.weight": "diffusion_pytorch_model-00001-of-00003.safetensors",
+    "layers.19.proj.ffn_down.weight": "diffusion_pytorch_model-00002-of-00003.safetensors",
+    "layers.19.proj.ffn_up_gate.weight": "diffusion_pytorch_model-00002-of-00003.safetensors",
+    "layers.19.proj.k.weight": "diffusion_pytorch_model-00001-of-00003.safetensors",
+    "layers.19.proj.o.weight": "diffusion_pytorch_model-00001-of-00003.safetensors",
+    "layers.19.proj.q.weight": "diffusion_pytorch_model-00001-of-00003.safetensors",
+    "layers.19.proj.v.weight": "diffusion_pytorch_model-00001-of-00003.safetensors",
+    "layers.19.q_norm.weight": "diffusion_pytorch_model-00001-of-00003.safetensors",
+    "layers.2.k_norm.weight": "diffusion_pytorch_model-00001-of-00003.safetensors",
+    "layers.2.norm1.weight": "diffusion_pytorch_model-00001-of-00003.safetensors",
+    "layers.2.norm2.weight": "diffusion_pytorch_model-00001-of-00003.safetensors",
+    "layers.2.proj.ffn_down.weight": "diffusion_pytorch_model-00001-of-00003.safetensors",
+    "layers.2.proj.ffn_up_gate.weight": "diffusion_pytorch_model-00001-of-00003.safetensors",
+    "layers.2.proj.k.weight": "diffusion_pytorch_model-00001-of-00003.safetensors",
+    "layers.2.proj.o.weight": "diffusion_pytorch_model-00001-of-00003.safetensors",
+    "layers.2.proj.q.weight": "diffusion_pytorch_model-00001-of-00003.safetensors",
+    "layers.2.proj.v.weight": "diffusion_pytorch_model-00001-of-00003.safetensors",
+    "layers.2.q_norm.weight": "diffusion_pytorch_model-00001-of-00003.safetensors",
+    "layers.20.k_norm.weight": "diffusion_pytorch_model-00002-of-00003.safetensors",
+    "layers.20.norm1.weight": "diffusion_pytorch_model-00002-of-00003.safetensors",
+    "layers.20.norm2.weight": "diffusion_pytorch_model-00002-of-00003.safetensors",
+    "layers.20.proj.ffn_down.weight": "diffusion_pytorch_model-00002-of-00003.safetensors",
+    "layers.20.proj.ffn_up_gate.weight": "diffusion_pytorch_model-00002-of-00003.safetensors",
+    "layers.20.proj.k.weight": "diffusion_pytorch_model-00002-of-00003.safetensors",
+    "layers.20.proj.o.weight": "diffusion_pytorch_model-00002-of-00003.safetensors",
+    "layers.20.proj.q.weight": "diffusion_pytorch_model-00002-of-00003.safetensors",
+    "layers.20.proj.v.weight": "diffusion_pytorch_model-00002-of-00003.safetensors",
+    "layers.20.q_norm.weight": "diffusion_pytorch_model-00002-of-00003.safetensors",
+    "layers.21.k_norm.weight": "diffusion_pytorch_model-00002-of-00003.safetensors",
+    "layers.21.norm1.weight": "diffusion_pytorch_model-00002-of-00003.safetensors",
+    "layers.21.norm2.weight": "diffusion_pytorch_model-00002-of-00003.safetensors",
+    "layers.21.proj.ffn_down.weight": "diffusion_pytorch_model-00002-of-00003.safetensors",
+    "layers.21.proj.ffn_up_gate.weight": "diffusion_pytorch_model-00002-of-00003.safetensors",
+    "layers.21.proj.k.weight": "diffusion_pytorch_model-00002-of-00003.safetensors",
+    "layers.21.proj.o.weight": "diffusion_pytorch_model-00002-of-00003.safetensors",
+    "layers.21.proj.q.weight": "diffusion_pytorch_model-00002-of-00003.safetensors",
+    "layers.21.proj.v.weight": "diffusion_pytorch_model-00002-of-00003.safetensors",
+    "layers.21.q_norm.weight": "diffusion_pytorch_model-00002-of-00003.safetensors",
+    "layers.22.k_norm.weight": "diffusion_pytorch_model-00002-of-00003.safetensors",
+    "layers.22.norm1.weight": "diffusion_pytorch_model-00002-of-00003.safetensors",
+    "layers.22.norm2.weight": "diffusion_pytorch_model-00002-of-00003.safetensors",
+    "layers.22.proj.ffn_down.weight": "diffusion_pytorch_model-00002-of-00003.safetensors",
+    "layers.22.proj.ffn_up_gate.weight": "diffusion_pytorch_model-00002-of-00003.safetensors",
+    "layers.22.proj.k.weight": "diffusion_pytorch_model-00002-of-00003.safetensors",
+    "layers.22.proj.o.weight": "diffusion_pytorch_model-00002-of-00003.safetensors",
+    "layers.22.proj.q.weight": "diffusion_pytorch_model-00002-of-00003.safetensors",
+    "layers.22.proj.v.weight": "diffusion_pytorch_model-00002-of-00003.safetensors",
+    "layers.22.q_norm.weight": "diffusion_pytorch_model-00002-of-00003.safetensors",
+    "layers.23.k_norm.weight": "diffusion_pytorch_model-00002-of-00003.safetensors",
+    "layers.23.norm1.weight": "diffusion_pytorch_model-00002-of-00003.safetensors",
+    "layers.23.norm2.weight": "diffusion_pytorch_model-00002-of-00003.safetensors",
+    "layers.23.proj.ffn_down.weight": "diffusion_pytorch_model-00002-of-00003.safetensors",
+    "layers.23.proj.ffn_up_gate.weight": "diffusion_pytorch_model-00002-of-00003.safetensors",
+    "layers.23.proj.k.weight": "diffusion_pytorch_model-00002-of-00003.safetensors",
+    "layers.23.proj.o.weight": "diffusion_pytorch_model-00002-of-00003.safetensors",
+    "layers.23.proj.q.weight": "diffusion_pytorch_model-00002-of-00003.safetensors",
+    "layers.23.proj.v.weight": "diffusion_pytorch_model-00002-of-00003.safetensors",
+    "layers.23.q_norm.weight": "diffusion_pytorch_model-00002-of-00003.safetensors",
+    "layers.24.k_norm.weight": "diffusion_pytorch_model-00002-of-00003.safetensors",
+    "layers.24.norm1.weight": "diffusion_pytorch_model-00002-of-00003.safetensors",
+    "layers.24.norm2.weight": "diffusion_pytorch_model-00002-of-00003.safetensors",
+    "layers.24.proj.ffn_down.weight": "diffusion_pytorch_model-00002-of-00003.safetensors",
+    "layers.24.proj.ffn_up_gate.weight": "diffusion_pytorch_model-00002-of-00003.safetensors",
+    "layers.24.proj.k.weight": "diffusion_pytorch_model-00002-of-00003.safetensors",
+    "layers.24.proj.o.weight": "diffusion_pytorch_model-00002-of-00003.safetensors",
+    "layers.24.proj.q.weight": "diffusion_pytorch_model-00002-of-00003.safetensors",
+    "layers.24.proj.v.weight": "diffusion_pytorch_model-00002-of-00003.safetensors",
+    "layers.24.q_norm.weight": "diffusion_pytorch_model-00002-of-00003.safetensors",
+    "layers.25.k_norm.weight": "diffusion_pytorch_model-00002-of-00003.safetensors",
+    "layers.25.norm1.weight": "diffusion_pytorch_model-00002-of-00003.safetensors",
+    "layers.25.norm2.weight": "diffusion_pytorch_model-00002-of-00003.safetensors",
+    "layers.25.proj.ffn_down.weight": "diffusion_pytorch_model-00002-of-00003.safetensors",
+    "layers.25.proj.ffn_up_gate.weight": "diffusion_pytorch_model-00002-of-00003.safetensors",
+    "layers.25.proj.k.weight": "diffusion_pytorch_model-00002-of-00003.safetensors",
+    "layers.25.proj.o.weight": "diffusion_pytorch_model-00002-of-00003.safetensors",
+    "layers.25.proj.q.weight": "diffusion_pytorch_model-00002-of-00003.safetensors",
+    "layers.25.proj.v.weight": "diffusion_pytorch_model-00002-of-00003.safetensors",
+    "layers.25.q_norm.weight": "diffusion_pytorch_model-00002-of-00003.safetensors",
+    "layers.26.k_norm.weight": "diffusion_pytorch_model-00002-of-00003.safetensors",
+    "layers.26.norm1.weight": "diffusion_pytorch_model-00002-of-00003.safetensors",
+    "layers.26.norm2.weight": "diffusion_pytorch_model-00002-of-00003.safetensors",
+    "layers.26.proj.ffn_down.weight": "diffusion_pytorch_model-00002-of-00003.safetensors",
+    "layers.26.proj.ffn_up_gate.weight": "diffusion_pytorch_model-00002-of-00003.safetensors",
+    "layers.26.proj.k.weight": "diffusion_pytorch_model-00002-of-00003.safetensors",
+    "layers.26.proj.o.weight": "diffusion_pytorch_model-00002-of-00003.safetensors",
+    "layers.26.proj.q.weight": "diffusion_pytorch_model-00002-of-00003.safetensors",
+    "layers.26.proj.v.weight": "diffusion_pytorch_model-00002-of-00003.safetensors",
+    "layers.26.q_norm.weight": "diffusion_pytorch_model-00002-of-00003.safetensors",
+    "layers.27.k_norm.weight": "diffusion_pytorch_model-00002-of-00003.safetensors",
+    "layers.27.norm1.weight": "diffusion_pytorch_model-00002-of-00003.safetensors",
+    "layers.27.norm2.weight": "diffusion_pytorch_model-00002-of-00003.safetensors",
+    "layers.27.proj.ffn_down.weight": "diffusion_pytorch_model-00002-of-00003.safetensors",
+    "layers.27.proj.ffn_up_gate.weight": "diffusion_pytorch_model-00002-of-00003.safetensors",
+    "layers.27.proj.k.weight": "diffusion_pytorch_model-00002-of-00003.safetensors",
+    "layers.27.proj.o.weight": "diffusion_pytorch_model-00002-of-00003.safetensors",
+    "layers.27.proj.q.weight": "diffusion_pytorch_model-00002-of-00003.safetensors",
+    "layers.27.proj.v.weight": "diffusion_pytorch_model-00002-of-00003.safetensors",
+    "layers.27.q_norm.weight": "diffusion_pytorch_model-00002-of-00003.safetensors",
+    "layers.28.k_norm.weight": "diffusion_pytorch_model-00002-of-00003.safetensors",
+    "layers.28.norm1.weight": "diffusion_pytorch_model-00002-of-00003.safetensors",
+    "layers.28.norm2.weight": "diffusion_pytorch_model-00002-of-00003.safetensors",
+    "layers.28.proj.ffn_down.weight": "diffusion_pytorch_model-00002-of-00003.safetensors",
+    "layers.28.proj.ffn_up_gate.weight": "diffusion_pytorch_model-00002-of-00003.safetensors",
+    "layers.28.proj.k.weight": "diffusion_pytorch_model-00002-of-00003.safetensors",
+    "layers.28.proj.o.weight": "diffusion_pytorch_model-00002-of-00003.safetensors",
+    "layers.28.proj.q.weight": "diffusion_pytorch_model-00002-of-00003.safetensors",
+    "layers.28.proj.v.weight": "diffusion_pytorch_model-00002-of-00003.safetensors",
+    "layers.28.q_norm.weight": "diffusion_pytorch_model-00002-of-00003.safetensors",
+    "layers.29.k_norm.weight": "diffusion_pytorch_model-00002-of-00003.safetensors",
+    "layers.29.norm1.weight": "diffusion_pytorch_model-00002-of-00003.safetensors",
+    "layers.29.norm2.weight": "diffusion_pytorch_model-00002-of-00003.safetensors",
+    "layers.29.proj.ffn_down.weight": "diffusion_pytorch_model-00002-of-00003.safetensors",
+    "layers.29.proj.ffn_up_gate.weight": "diffusion_pytorch_model-00002-of-00003.safetensors",
+    "layers.29.proj.k.weight": "diffusion_pytorch_model-00002-of-00003.safetensors",
+    "layers.29.proj.o.weight": "diffusion_pytorch_model-00002-of-00003.safetensors",
+    "layers.29.proj.q.weight": "diffusion_pytorch_model-00002-of-00003.safetensors",
+    "layers.29.proj.v.weight": "diffusion_pytorch_model-00002-of-00003.safetensors",
+    "layers.29.q_norm.weight": "diffusion_pytorch_model-00002-of-00003.safetensors",
+    "layers.3.k_norm.weight": "diffusion_pytorch_model-00001-of-00003.safetensors",
+    "layers.3.norm1.weight": "diffusion_pytorch_model-00001-of-00003.safetensors",
+    "layers.3.norm2.weight": "diffusion_pytorch_model-00001-of-00003.safetensors",
+    "layers.3.proj.ffn_down.weight": "diffusion_pytorch_model-00001-of-00003.safetensors",
+    "layers.3.proj.ffn_up_gate.weight": "diffusion_pytorch_model-00001-of-00003.safetensors",
+    "layers.3.proj.k.weight": "diffusion_pytorch_model-00001-of-00003.safetensors",
+    "layers.3.proj.o.weight": "diffusion_pytorch_model-00001-of-00003.safetensors",
+    "layers.3.proj.q.weight": "diffusion_pytorch_model-00001-of-00003.safetensors",
+    "layers.3.proj.v.weight": "diffusion_pytorch_model-00001-of-00003.safetensors",
+    "layers.3.q_norm.weight": "diffusion_pytorch_model-00001-of-00003.safetensors",
+    "layers.30.k_norm.weight": "diffusion_pytorch_model-00002-of-00003.safetensors",
+    "layers.30.norm1.weight": "diffusion_pytorch_model-00002-of-00003.safetensors",
+    "layers.30.norm2.weight": "diffusion_pytorch_model-00002-of-00003.safetensors",
+    "layers.30.proj.ffn_down.weight": "diffusion_pytorch_model-00002-of-00003.safetensors",
+    "layers.30.proj.ffn_up_gate.weight": "diffusion_pytorch_model-00002-of-00003.safetensors",
+    "layers.30.proj.k.weight": "diffusion_pytorch_model-00002-of-00003.safetensors",
+    "layers.30.proj.o.weight": "diffusion_pytorch_model-00002-of-00003.safetensors",
+    "layers.30.proj.q.weight": "diffusion_pytorch_model-00002-of-00003.safetensors",
+    "layers.30.proj.v.weight": "diffusion_pytorch_model-00002-of-00003.safetensors",
+    "layers.30.q_norm.weight": "diffusion_pytorch_model-00002-of-00003.safetensors",
+    "layers.31.k_norm.weight": "diffusion_pytorch_model-00002-of-00003.safetensors",
+    "layers.31.norm1.weight": "diffusion_pytorch_model-00002-of-00003.safetensors",
+    "layers.31.norm2.weight": "diffusion_pytorch_model-00002-of-00003.safetensors",
+    "layers.31.proj.ffn_down.weight": "diffusion_pytorch_model-00002-of-00003.safetensors",
+    "layers.31.proj.ffn_up_gate.weight": "diffusion_pytorch_model-00002-of-00003.safetensors",
+    "layers.31.proj.k.weight": "diffusion_pytorch_model-00002-of-00003.safetensors",
+    "layers.31.proj.o.weight": "diffusion_pytorch_model-00002-of-00003.safetensors",
+    "layers.31.proj.q.weight": "diffusion_pytorch_model-00002-of-00003.safetensors",
+    "layers.31.proj.v.weight": "diffusion_pytorch_model-00002-of-00003.safetensors",
+    "layers.31.q_norm.weight": "diffusion_pytorch_model-00002-of-00003.safetensors",
+    "layers.32.k_norm.weight": "diffusion_pytorch_model-00002-of-00003.safetensors",
+    "layers.32.norm1.weight": "diffusion_pytorch_model-00002-of-00003.safetensors",
+    "layers.32.norm2.weight": "diffusion_pytorch_model-00002-of-00003.safetensors",
+    "layers.32.proj.ffn_down.weight": "diffusion_pytorch_model-00002-of-00003.safetensors",
+    "layers.32.proj.ffn_up_gate.weight": "diffusion_pytorch_model-00002-of-00003.safetensors",
+    "layers.32.proj.k.weight": "diffusion_pytorch_model-00002-of-00003.safetensors",
+    "layers.32.proj.o.weight": "diffusion_pytorch_model-00002-of-00003.safetensors",
+    "layers.32.proj.q.weight": "diffusion_pytorch_model-00002-of-00003.safetensors",
+    "layers.32.proj.v.weight": "diffusion_pytorch_model-00002-of-00003.safetensors",
+    "layers.32.q_norm.weight": "diffusion_pytorch_model-00002-of-00003.safetensors",
+    "layers.33.k_norm.weight": "diffusion_pytorch_model-00002-of-00003.safetensors",
+    "layers.33.norm1.weight": "diffusion_pytorch_model-00002-of-00003.safetensors",
+    "layers.33.norm2.weight": "diffusion_pytorch_model-00002-of-00003.safetensors",
+    "layers.33.proj.ffn_down.weight": "diffusion_pytorch_model-00002-of-00003.safetensors",
+    "layers.33.proj.ffn_up_gate.weight": "diffusion_pytorch_model-00002-of-00003.safetensors",
+    "layers.33.proj.k.weight": "diffusion_pytorch_model-00002-of-00003.safetensors",
+    "layers.33.proj.o.weight": "diffusion_pytorch_model-00002-of-00003.safetensors",
+    "layers.33.proj.q.weight": "diffusion_pytorch_model-00002-of-00003.safetensors",
+    "layers.33.proj.v.weight": "diffusion_pytorch_model-00002-of-00003.safetensors",
+    "layers.33.q_norm.weight": "diffusion_pytorch_model-00002-of-00003.safetensors",
+    "layers.34.k_norm.weight": "diffusion_pytorch_model-00002-of-00003.safetensors",
+    "layers.34.norm1.weight": "diffusion_pytorch_model-00002-of-00003.safetensors",
+    "layers.34.norm2.weight": "diffusion_pytorch_model-00002-of-00003.safetensors",
+    "layers.34.proj.ffn_down.weight": "diffusion_pytorch_model-00002-of-00003.safetensors",
+    "layers.34.proj.ffn_up_gate.weight": "diffusion_pytorch_model-00002-of-00003.safetensors",
+    "layers.34.proj.k.weight": "diffusion_pytorch_model-00002-of-00003.safetensors",
+    "layers.34.proj.o.weight": "diffusion_pytorch_model-00002-of-00003.safetensors",
+    "layers.34.proj.q.weight": "diffusion_pytorch_model-00002-of-00003.safetensors",
+    "layers.34.proj.v.weight": "diffusion_pytorch_model-00002-of-00003.safetensors",
+    "layers.34.q_norm.weight": "diffusion_pytorch_model-00002-of-00003.safetensors",
+    "layers.35.k_norm.weight": "diffusion_pytorch_model-00002-of-00003.safetensors",
+    "layers.35.norm1.weight": "diffusion_pytorch_model-00002-of-00003.safetensors",
+    "layers.35.norm2.weight": "diffusion_pytorch_model-00002-of-00003.safetensors",
+    "layers.35.proj.ffn_down.weight": "diffusion_pytorch_model-00002-of-00003.safetensors",
+    "layers.35.proj.ffn_up_gate.weight": "diffusion_pytorch_model-00002-of-00003.safetensors",
+    "layers.35.proj.k.weight": "diffusion_pytorch_model-00002-of-00003.safetensors",
+    "layers.35.proj.o.weight": "diffusion_pytorch_model-00002-of-00003.safetensors",
+    "layers.35.proj.q.weight": "diffusion_pytorch_model-00002-of-00003.safetensors",
+    "layers.35.proj.v.weight": "diffusion_pytorch_model-00002-of-00003.safetensors",
+    "layers.35.q_norm.weight": "diffusion_pytorch_model-00002-of-00003.safetensors",
+    "layers.36.k_norm.weight": "diffusion_pytorch_model-00002-of-00003.safetensors",
+    "layers.36.norm1.weight": "diffusion_pytorch_model-00002-of-00003.safetensors",
+    "layers.36.norm2.weight": "diffusion_pytorch_model-00002-of-00003.safetensors",
+    "layers.36.proj.ffn_down.weight": "diffusion_pytorch_model-00002-of-00003.safetensors",
+    "layers.36.proj.ffn_up_gate.weight": "diffusion_pytorch_model-00002-of-00003.safetensors",
+    "layers.36.proj.k.weight": "diffusion_pytorch_model-00002-of-00003.safetensors",
+    "layers.36.proj.o.weight": "diffusion_pytorch_model-00002-of-00003.safetensors",
+    "layers.36.proj.q.weight": "diffusion_pytorch_model-00002-of-00003.safetensors",
+    "layers.36.proj.v.weight": "diffusion_pytorch_model-00002-of-00003.safetensors",
+    "layers.36.q_norm.weight": "diffusion_pytorch_model-00002-of-00003.safetensors",
+    "layers.37.k_norm.weight": "diffusion_pytorch_model-00002-of-00003.safetensors",
+    "layers.37.norm1.weight": "diffusion_pytorch_model-00002-of-00003.safetensors",
+    "layers.37.norm2.weight": "diffusion_pytorch_model-00002-of-00003.safetensors",
+    "layers.37.proj.ffn_down.weight": "diffusion_pytorch_model-00002-of-00003.safetensors",
+    "layers.37.proj.ffn_up_gate.weight": "diffusion_pytorch_model-00002-of-00003.safetensors",
+    "layers.37.proj.k.weight": "diffusion_pytorch_model-00002-of-00003.safetensors",
+    "layers.37.proj.o.weight": "diffusion_pytorch_model-00002-of-00003.safetensors",
+    "layers.37.proj.q.weight": "diffusion_pytorch_model-00002-of-00003.safetensors",
+    "layers.37.proj.v.weight": "diffusion_pytorch_model-00002-of-00003.safetensors",
+    "layers.37.q_norm.weight": "diffusion_pytorch_model-00002-of-00003.safetensors",
+    "layers.38.k_norm.weight": "diffusion_pytorch_model-00002-of-00003.safetensors",
+    "layers.38.norm1.weight": "diffusion_pytorch_model-00002-of-00003.safetensors",
+    "layers.38.norm2.weight": "diffusion_pytorch_model-00002-of-00003.safetensors",
+    "layers.38.proj.ffn_down.weight": "diffusion_pytorch_model-00002-of-00003.safetensors",
+    "layers.38.proj.ffn_up_gate.weight": "diffusion_pytorch_model-00002-of-00003.safetensors",
+    "layers.38.proj.k.weight": "diffusion_pytorch_model-00002-of-00003.safetensors",
+    "layers.38.proj.o.weight": "diffusion_pytorch_model-00002-of-00003.safetensors",
+    "layers.38.proj.q.weight": "diffusion_pytorch_model-00002-of-00003.safetensors",
+    "layers.38.proj.v.weight": "diffusion_pytorch_model-00002-of-00003.safetensors",
+    "layers.38.q_norm.weight": "diffusion_pytorch_model-00002-of-00003.safetensors",
+    "layers.39.k_norm.weight": "diffusion_pytorch_model-00002-of-00003.safetensors",
+    "layers.39.norm1.weight": "diffusion_pytorch_model-00002-of-00003.safetensors",
+    "layers.39.norm2.weight": "diffusion_pytorch_model-00002-of-00003.safetensors",
+    "layers.39.proj.ffn_down.weight": "diffusion_pytorch_model-00002-of-00003.safetensors",
+    "layers.39.proj.ffn_up_gate.weight": "diffusion_pytorch_model-00002-of-00003.safetensors",
+    "layers.39.proj.k.weight": "diffusion_pytorch_model-00002-of-00003.safetensors",
+    "layers.39.proj.o.weight": "diffusion_pytorch_model-00002-of-00003.safetensors",
+    "layers.39.proj.q.weight": "diffusion_pytorch_model-00002-of-00003.safetensors",
+    "layers.39.proj.v.weight": "diffusion_pytorch_model-00002-of-00003.safetensors",
+    "layers.39.q_norm.weight": "diffusion_pytorch_model-00002-of-00003.safetensors",
+    "layers.4.k_norm.weight": "diffusion_pytorch_model-00001-of-00003.safetensors",
+    "layers.4.norm1.weight": "diffusion_pytorch_model-00001-of-00003.safetensors",
+    "layers.4.norm2.weight": "diffusion_pytorch_model-00001-of-00003.safetensors",
+    "layers.4.proj.ffn_down.weight": "diffusion_pytorch_model-00001-of-00003.safetensors",
+    "layers.4.proj.ffn_up_gate.weight": "diffusion_pytorch_model-00001-of-00003.safetensors",
+    "layers.4.proj.k.weight": "diffusion_pytorch_model-00001-of-00003.safetensors",
+    "layers.4.proj.o.weight": "diffusion_pytorch_model-00001-of-00003.safetensors",
+    "layers.4.proj.q.weight": "diffusion_pytorch_model-00001-of-00003.safetensors",
+    "layers.4.proj.v.weight": "diffusion_pytorch_model-00001-of-00003.safetensors",
+    "layers.4.q_norm.weight": "diffusion_pytorch_model-00001-of-00003.safetensors",
+    "layers.40.k_norm.weight": "diffusion_pytorch_model-00002-of-00003.safetensors",
+    "layers.40.norm1.weight": "diffusion_pytorch_model-00002-of-00003.safetensors",
+    "layers.40.norm2.weight": "diffusion_pytorch_model-00002-of-00003.safetensors",
+    "layers.40.proj.ffn_down.weight": "diffusion_pytorch_model-00002-of-00003.safetensors",
+    "layers.40.proj.ffn_up_gate.weight": "diffusion_pytorch_model-00002-of-00003.safetensors",
+    "layers.40.proj.k.weight": "diffusion_pytorch_model-00002-of-00003.safetensors",
+    "layers.40.proj.o.weight": "diffusion_pytorch_model-00002-of-00003.safetensors",
+    "layers.40.proj.q.weight": "diffusion_pytorch_model-00002-of-00003.safetensors",
+    "layers.40.proj.v.weight": "diffusion_pytorch_model-00002-of-00003.safetensors",
+    "layers.40.q_norm.weight": "diffusion_pytorch_model-00002-of-00003.safetensors",
+    "layers.41.k_norm.weight": "diffusion_pytorch_model-00002-of-00003.safetensors",
+    "layers.41.norm1.weight": "diffusion_pytorch_model-00002-of-00003.safetensors",
+    "layers.41.norm2.weight": "diffusion_pytorch_model-00002-of-00003.safetensors",
+    "layers.41.proj.ffn_down.weight": "diffusion_pytorch_model-00003-of-00003.safetensors",
+    "layers.41.proj.ffn_up_gate.weight": "diffusion_pytorch_model-00003-of-00003.safetensors",
+    "layers.41.proj.k.weight": "diffusion_pytorch_model-00002-of-00003.safetensors",
+    "layers.41.proj.o.weight": "diffusion_pytorch_model-00002-of-00003.safetensors",
+    "layers.41.proj.q.weight": "diffusion_pytorch_model-00002-of-00003.safetensors",
+    "layers.41.proj.v.weight": "diffusion_pytorch_model-00002-of-00003.safetensors",
+    "layers.41.q_norm.weight": "diffusion_pytorch_model-00002-of-00003.safetensors",
+    "layers.5.k_norm.weight": "diffusion_pytorch_model-00001-of-00003.safetensors",
+    "layers.5.norm1.weight": "diffusion_pytorch_model-00001-of-00003.safetensors",
+    "layers.5.norm2.weight": "diffusion_pytorch_model-00001-of-00003.safetensors",
+    "layers.5.proj.ffn_down.weight": "diffusion_pytorch_model-00001-of-00003.safetensors",
+    "layers.5.proj.ffn_up_gate.weight": "diffusion_pytorch_model-00001-of-00003.safetensors",
+    "layers.5.proj.k.weight": "diffusion_pytorch_model-00001-of-00003.safetensors",
+    "layers.5.proj.o.weight": "diffusion_pytorch_model-00001-of-00003.safetensors",
+    "layers.5.proj.q.weight": "diffusion_pytorch_model-00001-of-00003.safetensors",
+    "layers.5.proj.v.weight": "diffusion_pytorch_model-00001-of-00003.safetensors",
+    "layers.5.q_norm.weight": "diffusion_pytorch_model-00001-of-00003.safetensors",
+    "layers.6.k_norm.weight": "diffusion_pytorch_model-00001-of-00003.safetensors",
+    "layers.6.norm1.weight": "diffusion_pytorch_model-00001-of-00003.safetensors",
+    "layers.6.norm2.weight": "diffusion_pytorch_model-00001-of-00003.safetensors",
+    "layers.6.proj.ffn_down.weight": "diffusion_pytorch_model-00001-of-00003.safetensors",
+    "layers.6.proj.ffn_up_gate.weight": "diffusion_pytorch_model-00001-of-00003.safetensors",
+    "layers.6.proj.k.weight": "diffusion_pytorch_model-00001-of-00003.safetensors",
+    "layers.6.proj.o.weight": "diffusion_pytorch_model-00001-of-00003.safetensors",
+    "layers.6.proj.q.weight": "diffusion_pytorch_model-00001-of-00003.safetensors",
+    "layers.6.proj.v.weight": "diffusion_pytorch_model-00001-of-00003.safetensors",
+    "layers.6.q_norm.weight": "diffusion_pytorch_model-00001-of-00003.safetensors",
+    "layers.7.k_norm.weight": "diffusion_pytorch_model-00001-of-00003.safetensors",
+    "layers.7.norm1.weight": "diffusion_pytorch_model-00001-of-00003.safetensors",
+    "layers.7.norm2.weight": "diffusion_pytorch_model-00001-of-00003.safetensors",
+    "layers.7.proj.ffn_down.weight": "diffusion_pytorch_model-00001-of-00003.safetensors",
+    "layers.7.proj.ffn_up_gate.weight": "diffusion_pytorch_model-00001-of-00003.safetensors",
+    "layers.7.proj.k.weight": "diffusion_pytorch_model-00001-of-00003.safetensors",
+    "layers.7.proj.o.weight": "diffusion_pytorch_model-00001-of-00003.safetensors",
+    "layers.7.proj.q.weight": "diffusion_pytorch_model-00001-of-00003.safetensors",
+    "layers.7.proj.v.weight": "diffusion_pytorch_model-00001-of-00003.safetensors",
+    "layers.7.q_norm.weight": "diffusion_pytorch_model-00001-of-00003.safetensors",
+    "layers.8.k_norm.weight": "diffusion_pytorch_model-00001-of-00003.safetensors",
+    "layers.8.norm1.weight": "diffusion_pytorch_model-00001-of-00003.safetensors",
+    "layers.8.norm2.weight": "diffusion_pytorch_model-00001-of-00003.safetensors",
+    "layers.8.proj.ffn_down.weight": "diffusion_pytorch_model-00001-of-00003.safetensors",
+    "layers.8.proj.ffn_up_gate.weight": "diffusion_pytorch_model-00001-of-00003.safetensors",
+    "layers.8.proj.k.weight": "diffusion_pytorch_model-00001-of-00003.safetensors",
+    "layers.8.proj.o.weight": "diffusion_pytorch_model-00001-of-00003.safetensors",
+    "layers.8.proj.q.weight": "diffusion_pytorch_model-00001-of-00003.safetensors",
+    "layers.8.proj.v.weight": "diffusion_pytorch_model-00001-of-00003.safetensors",
+    "layers.8.q_norm.weight": "diffusion_pytorch_model-00001-of-00003.safetensors",
+    "layers.9.k_norm.weight": "diffusion_pytorch_model-00001-of-00003.safetensors",
+    "layers.9.norm1.weight": "diffusion_pytorch_model-00001-of-00003.safetensors",
+    "layers.9.norm2.weight": "diffusion_pytorch_model-00001-of-00003.safetensors",
+    "layers.9.proj.ffn_down.weight": "diffusion_pytorch_model-00001-of-00003.safetensors",
+    "layers.9.proj.ffn_up_gate.weight": "diffusion_pytorch_model-00001-of-00003.safetensors",
+    "layers.9.proj.k.weight": "diffusion_pytorch_model-00001-of-00003.safetensors",
+    "layers.9.proj.o.weight": "diffusion_pytorch_model-00001-of-00003.safetensors",
+    "layers.9.proj.q.weight": "diffusion_pytorch_model-00001-of-00003.safetensors",
+    "layers.9.proj.v.weight": "diffusion_pytorch_model-00001-of-00003.safetensors",
+    "layers.9.q_norm.weight": "diffusion_pytorch_model-00001-of-00003.safetensors",
+    "rope.bands": "diffusion_pytorch_model-00001-of-00003.safetensors",
+    "shared_block_modulation.bias": "diffusion_pytorch_model-00001-of-00003.safetensors",
+    "shared_block_modulation.weight": "diffusion_pytorch_model-00001-of-00003.safetensors",
+    "text_embedder.blocks.0.ffn_down.bias": "diffusion_pytorch_model-00001-of-00003.safetensors",
+    "text_embedder.blocks.0.ffn_down.weight": "diffusion_pytorch_model-00001-of-00003.safetensors",
+    "text_embedder.blocks.0.ffn_gate.bias": "diffusion_pytorch_model-00001-of-00003.safetensors",
+    "text_embedder.blocks.0.ffn_gate.weight": "diffusion_pytorch_model-00001-of-00003.safetensors",
+    "text_embedder.blocks.0.ffn_up.bias": "diffusion_pytorch_model-00001-of-00003.safetensors",
+    "text_embedder.blocks.0.ffn_up.weight": "diffusion_pytorch_model-00001-of-00003.safetensors",
+    "text_embedder.blocks.0.k_norm.weight": "diffusion_pytorch_model-00001-of-00003.safetensors",
+    "text_embedder.blocks.0.norm1.weight": "diffusion_pytorch_model-00001-of-00003.safetensors",
+    "text_embedder.blocks.0.norm2.weight": "diffusion_pytorch_model-00001-of-00003.safetensors",
+    "text_embedder.blocks.0.proj.bias": "diffusion_pytorch_model-00001-of-00003.safetensors",
+    "text_embedder.blocks.0.proj.weight": "diffusion_pytorch_model-00001-of-00003.safetensors",
+    "text_embedder.blocks.0.q_norm.weight": "diffusion_pytorch_model-00001-of-00003.safetensors",
+    "text_embedder.blocks.0.qkv.bias": "diffusion_pytorch_model-00001-of-00003.safetensors",
+    "text_embedder.blocks.0.qkv.weight": "diffusion_pytorch_model-00001-of-00003.safetensors",
+    "text_embedder.blocks.1.ffn_down.bias": "diffusion_pytorch_model-00001-of-00003.safetensors",
+    "text_embedder.blocks.1.ffn_down.weight": "diffusion_pytorch_model-00001-of-00003.safetensors",
+    "text_embedder.blocks.1.ffn_gate.bias": "diffusion_pytorch_model-00001-of-00003.safetensors",
+    "text_embedder.blocks.1.ffn_gate.weight": "diffusion_pytorch_model-00001-of-00003.safetensors",
+    "text_embedder.blocks.1.ffn_up.bias": "diffusion_pytorch_model-00001-of-00003.safetensors",
+    "text_embedder.blocks.1.ffn_up.weight": "diffusion_pytorch_model-00001-of-00003.safetensors",
+    "text_embedder.blocks.1.k_norm.weight": "diffusion_pytorch_model-00001-of-00003.safetensors",
+    "text_embedder.blocks.1.norm1.weight": "diffusion_pytorch_model-00001-of-00003.safetensors",
+    "text_embedder.blocks.1.norm2.weight": "diffusion_pytorch_model-00001-of-00003.safetensors",
+    "text_embedder.blocks.1.proj.bias": "diffusion_pytorch_model-00001-of-00003.safetensors",
+    "text_embedder.blocks.1.proj.weight": "diffusion_pytorch_model-00001-of-00003.safetensors",
+    "text_embedder.blocks.1.q_norm.weight": "diffusion_pytorch_model-00001-of-00003.safetensors",
+    "text_embedder.blocks.1.qkv.bias": "diffusion_pytorch_model-00001-of-00003.safetensors",
+    "text_embedder.blocks.1.qkv.weight": "diffusion_pytorch_model-00001-of-00003.safetensors",
+    "text_embedder.proj.bias": "diffusion_pytorch_model-00001-of-00003.safetensors",
+    "text_embedder.proj.weight": "diffusion_pytorch_model-00001-of-00003.safetensors",
+    "time_embedding.0.bias": "diffusion_pytorch_model-00001-of-00003.safetensors",
+    "time_embedding.0.weight": "diffusion_pytorch_model-00001-of-00003.safetensors",
+    "time_embedding.2.bias": "diffusion_pytorch_model-00001-of-00003.safetensors",
+    "time_embedding.2.weight": "diffusion_pytorch_model-00001-of-00003.safetensors",
+    "video_embedder.bias": "diffusion_pytorch_model-00001-of-00003.safetensors",
+    "video_embedder.weight": "diffusion_pytorch_model-00001-of-00003.safetensors",
+    "video_head_modulation.bias": "diffusion_pytorch_model-00003-of-00003.safetensors",
+    "video_head_modulation.weight": "diffusion_pytorch_model-00003-of-00003.safetensors"
+  }
+}

transformer/pipeline_moss_image.py

@@ -0,0 +1,440 @@
+"""MOSSImage modeling and text-to-image pipeline (no Motor/TorchTitan imports).
+
+Dependencies: torch>=2.9, diffusers>=0.38, transformers>=5.3, safetensors,
+accelerate, Pillow. Copy this file anywhere and import it normally.
+
+    import torch
+    # Import MOSSImagePipeline from this file for component-level loading.
+    pipe = MOSSImagePipeline.from_components_pretrained(
+        transformer_path="exported_transformer",  # save_pretrained directory
+        text_encoder_path="Qwen/Qwen3.5-9B",
+        vae_path="black-forest-labs/FLUX.2-klein-4B",
+    ).to("cuda")
+    pipe("A cat", generator=torch.Generator("cuda").manual_seed(42)).images[0].save("cat.png")
+
+Default guidance uses CFG 3 and skips block 22 only in the negative branch
+(zero-based main DiT block indices; two forwards per sampling step):
+    pipe("A cat", guidance_scale=3.0, uncond_skip_layers=[22])
+    pipe("A cat", uncond_skip_layers=[])  # disable negative-branch skipping
+Optional additional positive-branch SLG adds a third forward:
+    pipe("A cat", slg_layers=[24], slg_scale=3.0)
+The additional SLG layer/scale above illustrates the API, not a quality-tuned setting.
+
+For a consolidated training state dict, create MOSSImageTransformer and
+call load_training_state_dict(state), then save_pretrained("exported_transformer").
+Select raw or EMA weights explicitly before calling it; distributed checkpoint
+shards/optimizer states are not model state dicts. Defaults match the finalized dense MOSSImage model,
+not MoE, CFG-distilled, video, or reference-image variants. Prompt strings are
+already-rendered captions (dataset-specific quality/domain tags are not added).
+
+First-release defaults: bf16 DiT/text-encoder weights, fp32 RoPE frequencies
+and phase/sin/cos calculation, then RoPE cast to the video activation dtype.
+Initial Gaussian noise is sampled directly in fp32, including with bf16 weights.
+CFG and Euler latent updates use fp32; VAE decoding defaults to fp32. Scheduler
+shift defaults to 5.0. Exact numerical parity with the original fp32-master
+inference pipeline is not a release requirement.
+"""
+from __future__ import annotations
+
+import math
+from pathlib import Path
+
+import torch
+from torch import nn
+from torch.nn import functional as F
+from diffusers import AutoencoderKLFlux2, DiffusionPipeline, FlowMatchEulerDiscreteScheduler
+from diffusers.schedulers.scheduling_utils import SchedulerMixin  # Required by custom-component discovery.
+from diffusers.configuration_utils import ConfigMixin, register_to_config
+from diffusers.image_processor import VaeImageProcessor
+from diffusers.models.modeling_utils import ModelMixin
+from diffusers.models.modeling_outputs import Transformer2DModelOutput
+from diffusers.pipelines.pipeline_utils import ImagePipelineOutput
+from diffusers.utils.torch_utils import randn_tensor
+from transformers import AutoModel, AutoTokenizer
+
+
+class _RMSNorm(nn.Module):
+    def __init__(self, dim):
+        super().__init__()
+        self.weight = nn.Parameter(torch.ones(dim))
+
+    def forward(self, x):
+        y = x.float() * torch.rsqrt(x.float().square().mean(-1, keepdim=True) + 1e-6)
+        return y.to(x.dtype) * self.weight
+
+
+def _situ(hidden, gate):
+    return (4.0 * torch.tanh(gate / 4.0) * torch.sigmoid(gate)) * (25.0 * torch.tanh(hidden / 25.0))
+
+
+def _validate_skip_layers(layers, depth, name='skip_layers'):
+    layers = () if layers is None else tuple(layers)
+    if any(type(index) is not int or not 0 <= index < depth for index in layers):
+        raise ValueError(f'{name} must contain integer block indices in [0, {depth}).')
+    return frozenset(layers)
+
+
+def _rotate(x, rope):
+    sin, cos = rope.chunk(2, -1)
+    sin = torch.cat((sin, sin), -1).unsqueeze(-2)
+    cos = torch.cat((cos, cos), -1).unsqueeze(-2)
+    a, b = x.chunk(2, -1)
+    return x * cos + torch.cat((-b, a), -1) * sin
+
+
+def _attention(q, k, v, rope, heads):
+    shape = (*q.shape[:2], heads, q.shape[-1] // heads)
+    q = _rotate(q.reshape(shape), rope).transpose(1, 2)
+    k = _rotate(k.reshape(shape), rope).transpose(1, 2)
+    v = v.reshape(shape).transpose(1, 2)
+    return F.scaled_dot_product_attention(q.to(v.dtype), k.to(v.dtype), v).transpose(1, 2).flatten(2)
+
+
+class _TextBlock(nn.Module):
+    def __init__(self, dim, ffn_dim, heads):
+        super().__init__()
+        self.heads = heads
+        self.norm1, self.norm2 = _RMSNorm(dim), _RMSNorm(dim)
+        self.q_norm, self.k_norm = _RMSNorm(dim // heads), _RMSNorm(dim // heads)
+        self.qkv = nn.Linear(dim, 3 * dim)
+        self.proj = nn.Linear(dim, dim)
+        self.ffn_gate, self.ffn_up = nn.Linear(dim, ffn_dim), nn.Linear(dim, ffn_dim)
+        self.ffn_down = nn.Linear(ffn_dim, dim)
+
+    def forward(self, x, rope):
+        q, k, v = self.qkv(self.norm1(x)).chunk(3, -1)
+        q = self.q_norm(q.unflatten(-1, (self.heads, -1))).flatten(2)
+        k = self.k_norm(k.unflatten(-1, (self.heads, -1))).flatten(2)
+        x = x + self.proj(_attention(q, k, v, rope, self.heads))
+        y = self.norm2(x)
+        return x + self.ffn_down(_situ(self.ffn_up(y), self.ffn_gate(y)))
+
+
+class _TextProjection(nn.Module):
+    def __init__(self, text_dim, dim, ffn_dim, heads):
+        super().__init__()
+        self.proj = nn.Linear(text_dim, dim)
+        self.blocks = nn.ModuleList([_TextBlock(dim, ffn_dim, heads) for _ in range(2)])
+
+    def forward(self, x, rope):
+        x = self.proj(x.to(self.proj.weight.dtype))
+        for block in self.blocks:
+            x = block(x, rope)
+        return x
+
+
+class _Block(nn.Module):
+    def __init__(self, dim, ffn_dim, heads):
+        super().__init__()
+        self.heads = heads
+        self.norm1, self.norm2 = _RMSNorm(dim), _RMSNorm(dim)
+        self.q_norm, self.k_norm = _RMSNorm(dim // heads), _RMSNorm(dim // heads)
+        self.proj = nn.ModuleDict({
+            name: nn.Linear(dim, dim, bias=False) for name in ('q', 'k', 'v', 'o')
+        })
+        self.proj['ffn_up_gate'] = nn.Linear(dim, 2 * ffn_dim, bias=False)
+        self.proj['ffn_down'] = nn.Linear(ffn_dim, dim, bias=False)
+
+    def forward(self, x, rope, modulation):
+        shift_a, scale_a, gate_a, shift_f, scale_f, gate_f = modulation.unbind(1)
+        y = self.norm1(x) * (1 + scale_a[:, None]) + shift_a[:, None]
+        dtype = self.proj['q'].weight.dtype
+        y = y.to(dtype)
+        q = self.q_norm(self.proj['q'](y).unflatten(-1, (self.heads, -1))).flatten(2)
+        k = self.k_norm(self.proj['k'](y).unflatten(-1, (self.heads, -1))).flatten(2)
+        x = x + self.proj['o'](_attention(q, k, self.proj['v'](y), rope, self.heads)) * gate_a[:, None]
+        y = self.norm2(x) * (1 + scale_f[:, None]) + shift_f[:, None]
+        hidden, gate = self.proj['ffn_up_gate'](y.to(dtype)).chunk(2, -1)
+        return x + self.proj['ffn_down'](_situ(hidden, gate)) * gate_f[:, None]
+
+
+class _RoPE(nn.Module):
+    def __init__(self, head_dim, base):
+        super().__init__()
+        self.num_bands = head_dim // 8
+        self.base = base
+        self.register_buffer('bands', self._default_bands(torch.device('cpu')))
+
+    def _default_bands(self, device):
+        # Match training's CPU fp32 construction, including under meta init.
+        exponent = torch.arange(self.num_bands, device='cpu', dtype=torch.float32) / self.num_bands
+        return (1 / (self.base ** exponent)).to(device)
+
+    def _apply(self, fn, recurse=True):
+        super()._apply(fn, recurse=recurse)
+        # .float() cannot recover frequencies already rounded by .to(bf16).
+        self.bands = self._default_bands(self.bands.device)
+        return self
+
+    def _load_from_state_dict(self, *args, **kwargs):
+        super()._load_from_state_dict(*args, **kwargs)
+        # Frequencies are derived constants; retain the key for checkpoint compatibility.
+        self.bands = self._default_bands(self.bands.device)
+
+    def forward(self, height, width, text_length, device):
+        # Diffusers/Accelerate can install tensors directly, bypassing load hooks.
+        if self.bands.dtype != torch.float32:
+            self.bands = self._default_bands(device)
+        # Coordinates: (r, f, h, w); sequence order is image followed by text.
+        h, w = torch.meshgrid(torch.arange(height, device=device), torch.arange(width, device=device), indexing='ij')
+        image = torch.stack((torch.zeros_like(h), torch.full_like(h, text_length), h, w), -1).flatten(0, 1).float()
+        image[:, 2] -= (height - 1) / 2
+        image[:, 3] -= (width - 1) / 2
+        text = torch.zeros(text_length, 4, device=device, dtype=torch.float32)
+        text[:, 1] = torch.arange(text_length, device=device)
+        phase = torch.cat((image, text), 0).unsqueeze(-1) * self.bands
+        return torch.cat((phase.sin(), phase.cos()), -2).flatten(-2).unsqueeze(0)
+
+
+class MOSSImageTransformer(ModelMixin, ConfigMixin):
+    """MOSSImage DiT with training-compatible state-dict names and SDPA attention.
+
+    Each sample is evaluated at its true text length, including text refinement;
+    padded tokens cannot influence either text or joint attention.
+    """
+    @register_to_config
+    def __init__(self, dim=3072, ffn_dim=8192, num_heads=24, num_layers=42,
+                 text_dim=4096, freq_dim=256, latent_channels=128, rope_base=10000.0):
+        super().__init__()
+        if dim % num_heads or (dim // num_heads) % 8 or freq_dim % 2:
+            raise ValueError('Require dim divisible by heads, head_dim divisible by 8, and even freq_dim.')
+        self.video_embedder = nn.Linear(latent_channels, dim)
+        self.text_embedder = _TextProjection(text_dim, dim, ffn_dim, num_heads)
+        self.rope = _RoPE(dim // num_heads, rope_base)
+        self.time_embedding = nn.Sequential(nn.Linear(freq_dim, dim), nn.SiLU(), nn.Linear(dim, dim))
+        self.shared_block_modulation = nn.Linear(dim, 6 * dim)
+        self.layers = nn.ModuleList([_Block(dim, ffn_dim, num_heads) for _ in range(num_layers)])
+        self.final_norm_video = _RMSNorm(dim)
+        self.final_linear_video = nn.Linear(dim, latent_channels, bias=False)
+        self.video_head_modulation = nn.Linear(dim, 2 * dim)
+
+    def load_training_state_dict(self, state_dict, *, prefix=''):
+        """Load a consolidated raw/EMA model mapping, strictly; no silent key drops."""
+        if prefix:
+            if any(not key.startswith(prefix) for key in state_dict):
+                raise ValueError('All keys must start with the explicit prefix.')
+            state_dict = {key[len(prefix):]: value for key, value in state_dict.items()}
+        return self.load_state_dict(state_dict, strict=True)
+
+    def forward(self, hidden_states, timestep, encoder_hidden_states,
+                encoder_attention_mask=None, return_dict=True, skip_layers=None):
+        skip_layers = _validate_skip_layers(skip_layers, len(self.layers))
+        if hidden_states.ndim != 4 or hidden_states.shape[1] != self.config.latent_channels:
+            raise ValueError('hidden_states must be [B, latent_channels, H, W].')
+        batch, _, height, width = hidden_states.shape
+        if encoder_hidden_states.ndim != 3 or encoder_hidden_states.shape[0] != batch:
+            raise ValueError('encoder_hidden_states must be [B, S, text_dim].')
+        mask = encoder_attention_mask
+        if mask is None:
+            mask = torch.ones(encoder_hidden_states.shape[:2], dtype=torch.bool, device=hidden_states.device)
+        if mask.shape != encoder_hidden_states.shape[:2] or not torch.all((mask == 0) | (mask == 1)):
+            raise ValueError('encoder_attention_mask must be a binary [B, S] mask.')
+        times = torch.as_tensor(timestep, device=hidden_states.device, dtype=torch.float32).reshape(-1)
+        times = times.expand(batch)
+        # Training computes the sinusoid in float64 and feeds float32 into its MLP.
+        half = self.config.freq_dim // 2
+        freq = 10000 ** (-torch.arange(half, device=times.device, dtype=torch.float64) / half)
+        phase = times.double()[:, None] * freq
+        embedding = torch.cat((phase.cos(), phase.sin()), -1).float()
+        time = self.time_embedding(embedding.to(self.time_embedding[0].weight.dtype)).float()
+        mod = self.shared_block_modulation(time.to(self.shared_block_modulation.weight.dtype)).view(batch, 6, -1)
+        head = self.video_head_modulation(time.to(self.video_head_modulation.weight.dtype)).view(batch, 2, -1)
+        outputs = []
+        for index in range(batch):
+            context = encoder_hidden_states[index:index + 1, mask[index].bool()]
+            if context.shape[1] == 0:
+                raise ValueError('Each sample needs at least one valid text token.')
+            video = hidden_states[index:index + 1].flatten(2).transpose(1, 2)
+            video = self.video_embedder(video.to(self.video_embedder.weight.dtype).contiguous())
+            # Match the reference activation dtype under fp32 weights + autocast.
+            rope = self.rope(height, width, context.shape[1], hidden_states.device).to(video.dtype)
+            text = self.text_embedder(context, rope[:, height * width:])
+            x = torch.cat((video, text), 1)
+            for layer_index, layer in enumerate(self.layers):
+                if layer_index not in skip_layers:
+                    x = layer(x, rope, mod[index:index + 1])
+            x = self.final_norm_video(x[:, :height * width].to(self.final_norm_video.weight.dtype))
+            x = x * (1 + head[index:index + 1, 1, None]) + head[index:index + 1, 0, None]
+            x = self.final_linear_video(x.to(self.final_linear_video.weight.dtype))
+            outputs.append(x.transpose(1, 2).reshape(1, -1, height, width))
+        sample = torch.cat(outputs)
+        return Transformer2DModelOutput(sample=sample) if return_dict else (sample,)
+
+
+class MOSSImageScheduler(FlowMatchEulerDiscreteScheduler):
+    """Euler integration using the training project's V2, nearest-grid schedule."""
+    @register_to_config
+    def __init__(self, num_train_timesteps=1000, shift=5.0, training_shift=1.0):
+        super().__init__(num_train_timesteps=num_train_timesteps, shift=1.0)
+        self.register_to_config(shift=shift, training_shift=training_shift)
+
+    def set_timesteps(self, num_inference_steps=50, device=None):
+        if num_inference_steps < 1 or self.config.shift <= 0 or self.config.training_shift <= 0:
+            raise ValueError('Steps and shifts must be positive.')
+        total = self.config.num_train_timesteps
+        def grid(count, shift):
+            lower = shift / (total + shift - 1)
+            values = torch.linspace(1, lower, count)
+            return shift * values / (1 + (shift - 1) * values)
+        sigmas = grid(num_inference_steps, self.config.shift)
+        training = grid(total, self.config.training_shift).flip(0)
+        upper = torch.searchsorted(training, sigmas).clamp(max=total - 1)
+        lower = (upper - 1).clamp(min=0)
+        nearest = torch.where((sigmas - training[lower]).abs() <= (training[upper] - sigmas).abs(), lower, upper)
+        sigmas = torch.unique_consecutive(training[nearest])
+        # Our sigmas have already been shifted and snapped. Avoid a second shift.
+        self._shift = 1.0
+        super().set_timesteps(device=device, sigmas=sigmas.tolist())
+
+
+class MOSSImagePipeline(DiffusionPipeline):
+    model_cpu_offload_seq = 'text_encoder->transformer->vae'
+    _optional_components = ['text_encoder', 'tokenizer', 'vae']
+
+    def __init__(self, transformer: MOSSImageTransformer,
+                 scheduler: MOSSImageScheduler, text_encoder=None, tokenizer=None, vae=None):
+        super().__init__()
+        self.register_modules(transformer=transformer, scheduler=scheduler,
+                              text_encoder=text_encoder, tokenizer=tokenizer, vae=vae)
+        self.image_processor = VaeImageProcessor(vae_scale_factor=16)
+
+    @classmethod
+    def from_components_pretrained(cls, transformer_path, text_encoder_path, vae_path,
+                                   torch_dtype=torch.bfloat16, vae_dtype=torch.float32,
+                                   local_files_only=False):
+        common = dict(local_files_only=local_files_only)
+        transformer = MOSSImageTransformer.from_pretrained(transformer_path, torch_dtype=torch_dtype, **common)
+        tokenizer = AutoTokenizer.from_pretrained(text_encoder_path, **common)
+        text_encoder = AutoModel.from_pretrained(text_encoder_path, dtype=torch_dtype, **common).eval()
+        subfolder = None if (Path(vae_path) / 'config.json').is_file() else 'vae'
+        vae = AutoencoderKLFlux2.from_pretrained(vae_path, subfolder=subfolder, torch_dtype=vae_dtype, **common).eval()
+        return cls(transformer.eval(), MOSSImageScheduler(), text_encoder, tokenizer, vae)
+
+    @torch.no_grad()
+    def encode_prompt(self, prompt, max_sequence_length=4096):
+        if self.text_encoder is None or self.tokenizer is None:
+            raise ValueError('Text encoder/tokenizer required for string prompts.')
+        prompts = [prompt] if isinstance(prompt, str) else list(prompt)
+        if not prompts or not all(isinstance(p, str) for p in prompts):
+            raise ValueError('prompt must be a string or a nonempty list of strings.')
+        rows = []
+        # Unpadded individual forwards also avoid padding effects in Qwen GDN.
+        for value in prompts:
+            formatted = self.tokenizer.apply_chat_template(
+                [{'role': 'user', 'content': value}], tokenize=False,
+                add_generation_prompt=True, enable_thinking=True)
+            inputs = self.tokenizer(formatted, return_tensors='pt', truncation=True,
+                                    max_length=max_sequence_length).to(self._execution_device)
+            output = self.text_encoder(**inputs, use_cache=False, return_dict=True)
+            rows.append(output.last_hidden_state[0])
+        length = max(row.shape[0] for row in rows)
+        embeds = rows[0].new_zeros(len(rows), length, rows[0].shape[-1])
+        mask = torch.zeros(len(rows), length, device=embeds.device, dtype=torch.bool)
+        for index, row in enumerate(rows):
+            embeds[index, :row.shape[0]] = row
+            mask[index, :row.shape[0]] = True
+        return embeds, mask
+
+    @torch.no_grad()
+    def __call__(self, prompt=None, *, height=1024, width=1024, num_inference_steps=50,
+                 guidance_scale=3.0, negative_prompt='', num_images_per_prompt=1,
+                 generator=None, latents=None, prompt_embeds=None, prompt_attention_mask=None,
+                 negative_prompt_embeds=None, negative_prompt_attention_mask=None,
+                 max_sequence_length=4096, output_type='pil', return_dict=True,
+                 uncond_skip_layers=(22,), slg_layers=None, slg_scale=0.0,
+                 slg_start=0.0, slg_stop=1.0):
+        """Generate images with optional full-block skipping (zero-based indices).
+
+        uncond_skip_layers defaults to (22,) and replaces CFG's negative-prompt
+        prediction with a skipped-block prediction, retaining two forwards per
+        step. Pass None or an empty sequence to disable skipping. slg_layers
+        adds a skipped positive-prompt forward and the fp32 correction
+        slg_scale * (cond - cond_skip) to CFG. With guidance_scale=1, active
+        SLG uses cond_skip + slg_scale * (cond - cond_skip), matching Diffusers'
+        default SkipLayerGuidance formulation, with two forwards and no negative
+        branch. Both options can combine. slg_scale=1 then gives normal cond.
+        Both uncond skipping and SLG are active where
+        slg_start <= step_index / num_steps < slg_stop;
+        the interval follows denoising order, from initial noise to final image.
+        Outside it, inference uses normal CFG without skipping.
+        slg_scale=0 or empty slg_layers disables the additional forward.
+        """
+        uncond_skip_layers = _validate_skip_layers(
+            uncond_skip_layers, len(self.transformer.layers), 'uncond_skip_layers')
+        slg_layers = _validate_skip_layers(slg_layers, len(self.transformer.layers), 'slg_layers')
+        if not math.isfinite(slg_scale) or slg_scale < 0:
+            raise ValueError('slg_scale must be finite and nonnegative.')
+        if not 0 <= slg_start <= slg_stop <= 1:
+            raise ValueError('Require 0 <= slg_start <= slg_stop <= 1.')
+        if height <= 0 or width <= 0 or height % 16 or width % 16:
+            raise ValueError('height and width must be positive multiples of 16.')
+        if output_type not in {'pil', 'np', 'pt', 'latent'}:
+            raise ValueError('output_type must be pil, np, pt or latent.')
+        if num_images_per_prompt < 1 or not math.isfinite(guidance_scale) or guidance_scale < 0:
+            raise ValueError('Invalid image count or guidance scale.')
+        if not 1 <= max_sequence_length <= 4096:
+            raise ValueError('max_sequence_length must be between 1 and 4096.')
+        if (prompt is None) == (prompt_embeds is None):
+            raise ValueError('Provide exactly one of prompt and prompt_embeds.')
+        device = self._execution_device
+        if prompt_embeds is None:
+            prompt_embeds, prompt_attention_mask = self.encode_prompt(prompt, max_sequence_length)
+        batch = prompt_embeds.shape[0]
+        def prepare(embeds, mask):
+            if embeds.ndim != 3 or embeds.shape[0] != batch:
+                raise ValueError('Prompt embeddings must be [B, S, C] with matching batch size.')
+            if mask is None:
+                mask = torch.ones(embeds.shape[:2], dtype=torch.bool, device=embeds.device)
+            return (embeds.to(device).repeat_interleave(num_images_per_prompt, 0),
+                    mask.to(device).repeat_interleave(num_images_per_prompt, 0))
+        positive, positive_mask = prepare(prompt_embeds, prompt_attention_mask)
+        do_cfg = guidance_scale != 1.0
+        if do_cfg:
+            if negative_prompt_embeds is None:
+                negatives = [negative_prompt] * batch if isinstance(negative_prompt, str) else negative_prompt
+                if len(negatives) != batch:
+                    raise ValueError('negative_prompt must match prompt batch size.')
+                negative_prompt_embeds, negative_prompt_attention_mask = self.encode_prompt(negatives, max_sequence_length)
+            negative, negative_mask = prepare(negative_prompt_embeds, negative_prompt_attention_mask)
+        shape = (batch * num_images_per_prompt, self.transformer.config.latent_channels, height // 16, width // 16)
+        if latents is None:
+            # Sample in fp32: promoting bf16 noise afterward cannot restore its precision.
+            latents = randn_tensor(shape, generator=generator, device=device, dtype=torch.float32)
+        elif tuple(latents.shape) != shape:
+            raise ValueError(f'latents must have shape {shape}.')
+        else:
+            latents = latents.to(device=device, dtype=torch.float32)
+        self.scheduler.set_timesteps(num_inference_steps, device=device)
+        num_steps = len(self.scheduler.timesteps)
+        for step_index, timestep in enumerate(self.progress_bar(self.scheduler.timesteps)):
+            skip_active = slg_start <= step_index / num_steps < slg_stop
+            # CFG must combine fp32 predictions, even when the DiT runs in bf16.
+            # Converting only the combined velocity would preserve rounding errors.
+            prediction = self.transformer(latents, timestep, positive, positive_mask).sample.float()
+            conditional = prediction
+            if do_cfg:
+                skip_kwargs = {'skip_layers': uncond_skip_layers} if uncond_skip_layers and skip_active else {}
+                uncond = self.transformer(latents, timestep, negative, negative_mask, **skip_kwargs).sample.float()
+                prediction = uncond + guidance_scale * (prediction - uncond)
+            if slg_layers and slg_scale > 0 and skip_active:
+                skipped = self.transformer(
+                    latents, timestep, positive, positive_mask, skip_layers=slg_layers).sample.float()
+                if not do_cfg:
+                    prediction = skipped
+                prediction = prediction + slg_scale * (conditional - skipped)
+            latents = self.scheduler.step(prediction, timestep, latents, return_dict=False)[0]
+        if output_type == 'latent':
+            images = latents
+        else:
+            if self.vae is None:
+                raise ValueError('A FLUX.2 VAE is required for decoded output.')
+            z = latents.to(dtype=self.vae.dtype)
+            bn = self.vae.bn
+            eps = self.vae.config.batch_norm_eps
+            z = z * (bn.running_var + eps).sqrt().to(z).view(1, -1, 1, 1) + bn.running_mean.to(z).view(1, -1, 1, 1)
+            b, c, h, w = z.shape
+            z = z.reshape(b, c // 4, 2, 2, h, w).permute(0, 1, 4, 2, 5, 3).reshape(b, c // 4, 2 * h, 2 * w)
+            images = self.vae.decode(z, return_dict=False)[0].clamp(-1, 1)
+            images = self.image_processor.postprocess(images, output_type=output_type)
+        self.maybe_free_model_hooks()
+        return ImagePipelineOutput(images=images) if return_dict else (images,)

vae/config.json

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

vae/diffusion_pytorch_model.safetensors

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