moss_pht: step 12,500 EMA

Browse files

Files changed (10) hide show

.gitattributes +35 -0
LICENSE +202 -0
README.md +122 -0
diffusion_models/moss_pht_bf16.safetensors +3 -0
transformer/config.json +12 -0
transformer/diffusion_pytorch_model-00001-of-00003.safetensors +3 -0
transformer/diffusion_pytorch_model-00002-of-00003.safetensors +3 -0
transformer/diffusion_pytorch_model-00003-of-00003.safetensors +3 -0
transformer/diffusion_pytorch_model.safetensors.index.json +470 -0
transformer/pipeline_moss_image.py +440 -0

.gitattributes ADDED Viewed

	@@ -0,0 +1,35 @@

+*.7z filter=lfs diff=lfs merge=lfs -text
+*.arrow filter=lfs diff=lfs merge=lfs -text
+*.bin filter=lfs diff=lfs merge=lfs -text
+*.bz2 filter=lfs diff=lfs merge=lfs -text
+*.ckpt filter=lfs diff=lfs merge=lfs -text
+*.ftz filter=lfs diff=lfs merge=lfs -text
+*.gz filter=lfs diff=lfs merge=lfs -text
+*.h5 filter=lfs diff=lfs merge=lfs -text
+*.joblib filter=lfs diff=lfs merge=lfs -text
+*.lfs.* filter=lfs diff=lfs merge=lfs -text
+*.mlmodel filter=lfs diff=lfs merge=lfs -text
+*.model filter=lfs diff=lfs merge=lfs -text
+*.msgpack filter=lfs diff=lfs merge=lfs -text
+*.npy filter=lfs diff=lfs merge=lfs -text
+*.npz filter=lfs diff=lfs merge=lfs -text
+*.onnx filter=lfs diff=lfs merge=lfs -text
+*.ot filter=lfs diff=lfs merge=lfs -text
+*.parquet filter=lfs diff=lfs merge=lfs -text
+*.pb filter=lfs diff=lfs merge=lfs -text
+*.pickle filter=lfs diff=lfs merge=lfs -text
+*.pkl filter=lfs diff=lfs merge=lfs -text
+*.pt filter=lfs diff=lfs merge=lfs -text
+*.pth filter=lfs diff=lfs merge=lfs -text
+*.rar filter=lfs diff=lfs merge=lfs -text
+*.safetensors filter=lfs diff=lfs merge=lfs -text
+saved_model/**/* filter=lfs diff=lfs merge=lfs -text
+*.tar.* filter=lfs diff=lfs merge=lfs -text
+*.tar filter=lfs diff=lfs merge=lfs -text
+*.tflite filter=lfs diff=lfs merge=lfs -text
+*.tgz filter=lfs diff=lfs merge=lfs -text
+*.wasm filter=lfs diff=lfs merge=lfs -text
+*.xz filter=lfs diff=lfs merge=lfs -text
+*.zip filter=lfs diff=lfs merge=lfs -text
+*.zst filter=lfs diff=lfs merge=lfs -text
+*tfevents* filter=lfs diff=lfs merge=lfs -text

LICENSE ADDED Viewed

	@@ -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 ADDED Viewed

	@@ -0,0 +1,122 @@

+---
+base_model: OpenMOSS-Team/MOSS-Image-Base
+base_model_relation: finetune
+language:
+- en
+- zh
+license: apache-2.0
+pipeline_tag: text-to-image
+tags:
+- moss-image
+- diffusers
+- text-to-image
+---
+# moss_pht
+A lightweight, photo-oriented full fine-tune of
+[OpenMOSS-Team/MOSS-Image-Base](https://huggingface.co/OpenMOSS-Team/MOSS-Image-Base).
+The text encoder (Qwen3.5-9B), VAE
+(FLUX.2 small decoder), tokenizer, scheduler and inference pipeline are the base
+release's, unchanged. This repository therefore ships only the transformer, in
+two formats; load the other components from their official repositories, or
+swap in any compatible ones (for example a quantized text encoder).
+| File | Format |
+| --- | --- |
+| `transformer/` | Diffusers, the same layout as MOSS-Image-Base's `transformer/` |
+| `diffusion_models/moss_pht_bf16.safetensors` | ComfyUI single file, the same layout as Comfy-Org's `moss_image_base_bf16.safetensors` |
+## Diffusers
+```bash
+pip install "diffusers>=0.38" "transformers>=5.3" "accelerate>=1.10" "safetensors>=0.6" "Pillow>=10"
+```
+```python
+import torch
+from diffusers import AutoencoderKLFlux2
+from diffusers.utils.dynamic_modules_utils import get_class_from_dynamic_module
+from transformers import AutoModel, AutoTokenizer
+BASE = "OpenMOSS-Team/MOSS-Image-Base"
+def moss_class(name):
+    return get_class_from_dynamic_module(
+        BASE, "pipeline_moss_image.py", class_name=name, trust_remote_code=True
+    )
+pipe = moss_class("MOSSImagePipeline")(
+    transformer=moss_class("MOSSImageTransformer").from_pretrained(
+        "data-archetype/moss_pht", subfolder="transformer", torch_dtype=torch.bfloat16
+    ),
+    scheduler=moss_class("MOSSImageScheduler").from_pretrained(BASE, subfolder="scheduler"),
+    text_encoder=AutoModel.from_pretrained(BASE, subfolder="text_encoder", dtype=torch.bfloat16),
+    tokenizer=AutoTokenizer.from_pretrained(BASE, subfolder="tokenizer"),
+    vae=AutoencoderKLFlux2.from_pretrained(BASE, subfolder="vae", torch_dtype=torch.float32),
+)
+pipe.to("cuda")  # or pipe.enable_model_cpu_offload() on a smaller GPU
+image = pipe(
+    "A candid street photograph of a cyclist crossing a rain-soaked intersection at dusk",
+    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]
+```
+Each component is passed explicitly, so any of them can be replaced, for
+example `text_encoder=` with a quantized Qwen3.5-9B. Sampling settings, prompt
+formats (natural language, tags and JSON layouts) and requirements are the base
+model's; see its [model card](https://huggingface.co/OpenMOSS-Team/MOSS-Image-Base)
+and [user guide](https://github.com/OpenMOSS/MOSS-Image/blob/main/USERGUIDE.md).
+## ComfyUI
+Use ComfyUI's native MOSS-Image support with the
+[Comfy-Org/MOSS-Image](https://huggingface.co/Comfy-Org/MOSS-Image) text encoder
+and VAE; `moss_pht_bf16.safetensors` drops in for `moss_image_base_bf16.safetensors`
+in the Comfy-Org workflows.
+| File | Put in | Loader |
+| --- | --- | --- |
+| `diffusion_models/moss_pht_bf16.safetensors` (this repo) | `models/diffusion_models/` | Load Diffusion Model |
+| `text_encoders/qwen3.5_9b_bf16.safetensors` (or its int8 / w4a8 variants), from Comfy-Org/MOSS-Image | `models/text_encoders/` | Load CLIP, type `moss_image` |
+| `vae/flux2-vae_small.safetensors`, from Comfy-Org/MOSS-Image | `models/vae/` | Load VAE |
+## Recommended: a DMD LoRA
+Like the base checkpoint, generations with the standard multi-step sampler can be
+a bit unstable. Using a DMD LoRA distilled from MOSS-Image-Base is recommended: it
+sharpens the output distribution and stabilises generations.
+## Training
+| Setting | Value |
+| --- | --- |
+| Initialization | MOSS-Image-Base |
+| Trained parameters | All transformer parameters (fp32 master weights, bf16 autocast); text encoder and VAE frozen |
+| Steps | 12,500 at 64 images per step |
+| Resolution | About 1 megapixel, aspect-ratio buckets |
+| Text | Prompts up to 2,304 tokens; 10% of captions dropped to the empty prompt |
+| Objective | The base's flow-matching objective; noise levels from a truncated logit-normal (μ 1.09, σ 0.95, truncated at 0.95), plus 1% uniform and 1% pure noise |
+| Optimizer | [dionw](https://github.com/JTriggerFish/dionw): Dion (row-selected NorMuon) on weight matrices, AdamW on the rest |
+| Learning rate | 9e-6, AdamW-equivalent for every parameter |
+| Schedule | 200-step linear warmup, constant to step 10,500, then a 2,000-step 1 − √ cooldown to 0 |
+| Dion matrices | 25% of each block's rows updated per step (blocks with both sides ≥ 1024), Gram Newton–Schulz, heavy-ball momentum 0.95, NorMuon row normalization β₂ 0.975 |
+| AdamW parameters | β = (0.9, 0.975), ε 1e-8; includes the fused conditioning (AdaLN) projections |
+| Weight decay | 0.01 on Dion matrices, none on AdamW parameters |
+| Gradient clipping | Global norm 1.0 |
+| Released weights | EMA of the weights (decay 0.999, fp32, updated every step) at step 12,500, stored in bfloat16 (`rope.bands` in float32, as in the base release) |
+## License
+The weights and the bundled inference code (`transformer/pipeline_moss_image.py`,
+from OpenMOSS-Team/MOSS-Image-Base) are released under the base model's license,
+Apache 2.0 ([LICENSE](LICENSE), copied from OpenMOSS-Team/MOSS-Image-Base).
+We thank the OpenMOSS team for MOSS-Image.

diffusion_models/moss_pht_bf16.safetensors ADDED Viewed

	@@ -0,0 +1,3 @@

+version https://git-lfs.github.com/spec/v1
+oid sha256:42e929263750e59e3a4134e3a8f44140d3fd38d9863d124faf4bd2fc58730f68
+size 10164655808

transformer/config.json ADDED Viewed

	@@ -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 ADDED Viewed

	@@ -0,0 +1,3 @@

+version https://git-lfs.github.com/spec/v1
+oid sha256:19087e757d8e3e797a9bdda8bac28e84f284181e2637d59fb90c5b7d5ad603fe
+size 4991978016

transformer/diffusion_pytorch_model-00002-of-00003.safetensors ADDED Viewed

	@@ -0,0 +1,3 @@

+version https://git-lfs.github.com/spec/v1
+oid sha256:19501128f57aba7c0c2732f8e8e09e00405c57cd591a7d02f7f619a38c08076e
+size 4983137016

transformer/diffusion_pytorch_model-00003-of-00003.safetensors ADDED Viewed

	@@ -0,0 +1,3 @@

+version https://git-lfs.github.com/spec/v1
+oid sha256:fa067f79c0c2c023fe673936ed774cd91e89a18419c0983af611763a2b605b86
+size 189549184

transformer/diffusion_pytorch_model.safetensors.index.json ADDED Viewed

	@@ -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 ADDED Viewed

	@@ -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,)