Add files using upload-large-folder tool

.gitattributes

@@ -1,35 +1,2 @@
-*.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
+*.so filter=lfs diff=lfs merge=lfs -text

LICENSE

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

MANIFEST.json

@@ -0,0 +1,418 @@
+{
+  "base_models": {
+    "mage_flow": "microsoft/Mage-Flow",
+    "mage_source_commit": "df7f84d9f8fc991d189d929f03cff623b430a4a2",
+    "mage_transformer_sha256": "04431abbc3acd1a5b86a7f77269f8e77e05a6cac3b21f35e1aef7493f6ab4934",
+    "qwen": "Qwen/Qwen3-VL-4B-Instruct",
+    "qwen_quantized": "InsecureErasure/Qwen3-VL-4B-Instruct-NVFP4",
+    "qwen_quantized_revision": "cf080d6af667f5f4949295bd09a1f72a218fe0ee"
+  },
+  "evidence": {
+    "combined_frozen_val07": {
+      "all_seven_gates_passed": true,
+      "image_sha256": "857c1663636cd775ea05ffe37a98ac47ab7cd4e5362dd82c21ae1b53f7d71c3a",
+      "result_sha256": "67afd8ca0c00fff949bfc4c8b660ca864617839cbc01231a7157769a964287c7"
+    },
+    "legacy_portable_package_generation": {
+      "image_sha256": "2c1f2853e37a2ede976f24bc9b7eef172aed7665266c615401ea189fcbcee19f",
+      "report_sha256": "2877697cdf8d469f689aa68751d24c8f25288388f682e61355234e82fae35eda",
+      "status": "success"
+    },
+    "native_text_quality_gate_result_sha256": "4a31a2c0eb4cac7401224eaeab23577bd31cc5ef69f6edec5289757f113373e6",
+    "standard_huggingface_layout_generation": {
+      "bf16_target_weight_reads": 0,
+      "image_sha256": "2c1f2853e37a2ede976f24bc9b7eef172aed7665266c615401ea189fcbcee19f",
+      "layout": "complete_huggingface_sharded_repository",
+      "loaded_non_target_tensor_count": 301,
+      "loaded_nvfp4_projection_count": 48,
+      "report_sha256": "61c8b651a5a9814c20fa4a99719e0b0a6caa827f61c0865993d916de25e6507d",
+      "status": "success",
+      "transformer_checkpoint_tensor_count": 493,
+      "transformer_shard_count": 4
+    },
+    "transformer_quality_results": {
+      "VAL-01": "82fadbc86ffb980a503c61ba50bfae25fd302503cf9ebb807029aacdc484f65a",
+      "VAL-07": "2476505becf3cec14d174c35b42a92ae180cea1173ebe376d0d40f530a0292ce",
+      "VAL-09": "c5d25ab78790295a0945b9783e146ff114991cc72af80a524d083034696110fc",
+      "VAL-13": "d2f5b3e7b3be8d430fbd8ca47abdfa405d130f080d45a51fe5ab90fde309a07e"
+    }
+  },
+  "files": [
+    {
+      "path": ".gitattributes",
+      "role": "release support file",
+      "sha256": "0b59beee6f0780aa6d1da36ef06df3b64752a757d2c37b1abf1941875eddf10a",
+      "size": 91
+    },
+    {
+      "path": "LICENSE",
+      "role": "license or third-party attribution",
+      "sha256": "275b4dd619de4e16a017b10d0beec72abbbbf14ee8a2fc68f8bdb398e821f623",
+      "size": 1066
+    },
+    {
+      "path": "README.md",
+      "role": "Hugging Face model card and usage guide",
+      "sha256": "3333284289f34219f30a3c142d9b1ca9e2682604b161a92b5bc35d07fbc1a267",
+      "size": 6581
+    },
+    {
+      "path": "THIRD_PARTY_NOTICES.md",
+      "role": "license or third-party attribution",
+      "sha256": "848e6132befdb77dada57b9d077b19ad57d65d02b3df63161d3c8d6ee9a11e99",
+      "size": 1336
+    },
+    {
+      "path": "build_manifest.py",
+      "role": "release manifest builder",
+      "sha256": "82ea93365ae63418268f28cdf3fc8f2360ab5a3ba19da2a42ea914e832a9c45d",
+      "size": 8056
+    },
+    {
+      "path": "build_native.sh",
+      "role": "project-local native runtime build helper",
+      "sha256": "7e1bccc38c982b05c0ad027c3faace9c0d5783e638e75ff4875430d5ab34445d",
+      "size": 1803
+    },
+    {
+      "path": "examples/VAL-07_combined_native.png",
+      "role": "validated combined NVFP4 example output",
+      "sha256": "857c1663636cd775ea05ffe37a98ac47ab7cd4e5362dd82c21ae1b53f7d71c3a",
+      "size": 174472
+    },
+    {
+      "path": "generate.py",
+      "role": "portable single-image inference entry point",
+      "sha256": "aaa7aa6a101cbdf5d2bfde1354b162df58c7fa50309d40514d8b137c79485d23",
+      "size": 14516
+    },
+    {
+      "path": "licenses/MAGE-MIT.txt",
+      "role": "license or third-party attribution",
+      "sha256": "275b4dd619de4e16a017b10d0beec72abbbbf14ee8a2fc68f8bdb398e821f623",
+      "size": 1066
+    },
+    {
+      "path": "licenses/QWEN-APACHE-2.0.txt",
+      "role": "license or third-party attribution",
+      "sha256": "cfc7749b96f63bd31c3c42b5c471bf756814053e847c10f3eb003417bc523d30",
+      "size": 11358
+    },
+    {
+      "path": "model_index.json",
+      "role": "Hugging Face Diffusers metadata and download-tracking query file",
+      "sha256": "6d07c7d0d05992544b66126adff35f21e8c9036f8b0b7ca623c92d7bd97c2213",
+      "size": 614
+    },
+    {
+      "path": "requirements.txt",
+      "role": "release support file",
+      "sha256": "2a16308e6034f1b06d60afd70bcd488f0f79b9499c9ff0acbf1c028652320f42",
+      "size": 485
+    },
+    {
+      "path": "run_validation.py",
+      "role": "portable frozen-prompt validation entry point",
+      "sha256": "47ea74991108f4f9ce519ae00ebbfe36d76ebf775ed598a9b539dc3112f7cd3a",
+      "size": 2044
+    },
+    {
+      "path": "runtime/libmage_nvfp4_linear.so",
+      "role": "native runtime, loader, source, or Python integration",
+      "sha256": "d2f7d62d26b4ae3babb7791aee8e0d95ad4794ac9c1d22146c974784381ada48",
+      "size": 1158440
+    },
+    {
+      "path": "runtime/libmage_nvfp4_torch_op.so",
+      "role": "native runtime, loader, source, or Python integration",
+      "sha256": "d0cd7bbde23fa8187d28cda21ad74b79696bf607d40e47a534922670fe597ad4",
+      "size": 167176
+    },
+    {
+      "path": "runtime/nvfp4_linear.cu",
+      "role": "native runtime, loader, source, or Python integration",
+      "sha256": "060697d4784be3e3461d38ef6ff24128f47a4322a05e5463f8687796b03bda9f",
+      "size": 32166
+    },
+    {
+      "path": "runtime/nvfp4_linear.h",
+      "role": "native runtime, loader, source, or Python integration",
+      "sha256": "7eba9c6b155730fad23b86314903d5cd96ee4430daf2d45706288fc7a9f5ee09",
+      "size": 2231
+    },
+    {
+      "path": "runtime/packed_artifact.py",
+      "role": "native runtime, loader, source, or Python integration",
+      "sha256": "6f3b3220421d42a19573e852641dff00f75a6d7ee87975bd3d82c995ee86ba05",
+      "size": 40672
+    },
+    {
+      "path": "runtime/packed_nvfp4_linear.py",
+      "role": "native runtime, loader, source, or Python integration",
+      "sha256": "65a98172fedc59023c80c55c723956746c13e14f3783bdf687af39cdf0ee1397",
+      "size": 17469
+    },
+    {
+      "path": "runtime/quant_text_encoder.py",
+      "role": "native runtime, loader, source, or Python integration",
+      "sha256": "d5f38363256d67dee3532531b4a4ae080fa736e1bdd3499e615cb6142bcf808a",
+      "size": 13452
+    },
+    {
+      "path": "runtime/sm120_linear_op.cpp",
+      "role": "native runtime, loader, source, or Python integration",
+      "sha256": "a53e1e83df95922eaa384503ce5c5bab31ef65becfcf3e0ac623b1f0868d278a",
+      "size": 12179
+    },
+    {
+      "path": "runtime/standard_transformer.py",
+      "role": "native runtime, loader, source, or Python integration",
+      "sha256": "d1a0b25ddace31569aeff166cb8d1b2a7a016ee50827ab1f9dd0ae476fc96ba1",
+      "size": 7891
+    },
+    {
+      "path": "runtime/torch_ops_native.py",
+      "role": "native runtime, loader, source, or Python integration",
+      "sha256": "842e0f0f14faca41d8801ff52d435e30d180e34907b81cbbddc5aff6dfab1829",
+      "size": 7980
+    },
+    {
+      "path": "scheduler/scheduler_config.json",
+      "role": "Mage FlowMatch scheduler component",
+      "sha256": "438fd8bcf254740e5d3f3e9800bbd9c571e342ab87885388d1505b7531c69c02",
+      "size": 169
+    },
+    {
+      "path": "text_encoder/chat_template.json",
+      "role": "packaged Qwen3-VL configuration or tokenizer asset",
+      "sha256": "6f8a6a55027e3da5160105556cda5dd69f6423f1c32645f6730d32de7773d0c4",
+      "size": 5502
+    },
+    {
+      "path": "text_encoder/config.json",
+      "role": "packaged Qwen3-VL configuration or tokenizer asset",
+      "sha256": "d9e81f6f1d8d9ea613d55599041c6f579fbf4a9e9e0eacc9dbdb5e2300b1c281",
+      "size": 1789
+    },
+    {
+      "path": "text_encoder/generation_config.json",
+      "role": "packaged Qwen3-VL configuration or tokenizer asset",
+      "sha256": "8469742d1fce0de951c8909b26a2c0c0d8490837ce476efb114da9e0cefc4d44",
+      "size": 269
+    },
+    {
+      "path": "text_encoder/merges.txt",
+      "role": "packaged Qwen3-VL configuration or tokenizer asset",
+      "sha256": "599bab54075088774b1733fde865d5bd747cbcc7a547c5bc12610e874e26f5e3",
+      "size": 1671839
+    },
+    {
+      "path": "text_encoder/model.safetensors",
+      "role": "complete mixed NVFP4/FP8 Qwen3-VL text checkpoint",
+      "sha256": "719906b435800757d22013d3d475a4853d59b779b022669fa8b8a193b85d0f41",
+      "size": 4031376064
+    },
+    {
+      "path": "text_encoder/preprocessor_config.json",
+      "role": "packaged Qwen3-VL configuration or tokenizer asset",
+      "sha256": "27225450ac9c6529872ee1924fcb0962ff5634834f817040f444118116f4e516",
+      "size": 390
+    },
+    {
+      "path": "text_encoder/tokenizer.json",
+      "role": "packaged Qwen3-VL configuration or tokenizer asset",
+      "sha256": "a5d85b6dcc535e6b93115a9ef287e6132fdbf30270da6218194ba742261173c7",
+      "size": 7032403
+    },
+    {
+      "path": "text_encoder/tokenizer_config.json",
+      "role": "packaged Qwen3-VL configuration or tokenizer asset",
+      "sha256": "c2da771801886ad9ae98181793ffd3dfb7f1af30f6f7c6a4e15d7dbba52e2399",
+      "size": 10868
+    },
+    {
+      "path": "text_encoder/video_preprocessor_config.json",
+      "role": "packaged Qwen3-VL configuration or tokenizer asset",
+      "sha256": "7768af27c1fafa9cc9011c1dc20067e03f8915e03b63504550e11d5066986d13",
+      "size": 385
+    },
+    {
+      "path": "text_encoder/vocab.json",
+      "role": "packaged Qwen3-VL configuration or tokenizer asset",
+      "sha256": "ca10d7e9fb3ed18575dd1e277a2579c16d108e32f27439684afa0e10b1440910",
+      "size": 2776833
+    },
+    {
+      "path": "transformer/config.json",
+      "role": "Mage transformer component configuration, index, or quantization metadata",
+      "sha256": "7520e71b0a4965f23e54308939853229fe1c32fbd10c69aac5adaead24d21475",
+      "size": 3194
+    },
+    {
+      "path": "transformer/diffusion_pytorch_model-00001-of-00004.safetensors",
+      "role": "complete sharded Mage transformer checkpoint with native NVFP4 modules",
+      "sha256": "040143f2dfc32db4220aabe46d98971e1701c83a566ccdb39e2553218f35f0dd",
+      "size": 1797002592
+    },
+    {
+      "path": "transformer/diffusion_pytorch_model-00002-of-00004.safetensors",
+      "role": "complete sharded Mage transformer checkpoint with native NVFP4 modules",
+      "sha256": "d7116264ef3cf0721c197ff4012de319b844abec1a43df1ade7d8868ce43f4a9",
+      "size": 1850448056
+    },
+    {
+      "path": "transformer/diffusion_pytorch_model-00003-of-00004.safetensors",
+      "role": "complete sharded Mage transformer checkpoint with native NVFP4 modules",
+      "sha256": "7534b2e25cf0e90a2bb0cca9d77f6cb2dc29198b2e387bf36cd59878bf9559b4",
+      "size": 1850448056
+    },
+    {
+      "path": "transformer/diffusion_pytorch_model-00004-of-00004.safetensors",
+      "role": "complete sharded Mage transformer checkpoint with native NVFP4 modules",
+      "sha256": "05a975382c7bada173c01f6005da16d7c78eaba771d554f0551934a25f1836d6",
+      "size": 128986568
+    },
+    {
+      "path": "transformer/diffusion_pytorch_model.safetensors.index.json",
+      "role": "Mage transformer component configuration, index, or quantization metadata",
+      "sha256": "6fc021886f29df126c3ff64cd61287869a2b9b3d142fee6084f81171cbdc90d8",
+      "size": 51641
+    },
+    {
+      "path": "transformer/nvfp4_metadata.json",
+      "role": "Mage transformer component configuration, index, or quantization metadata",
+      "sha256": "8aa6f8494dacd10f23e48e64ca6bdc0cc05754dffadd746bd6f500aa08c745fa",
+      "size": 71892
+    },
+    {
+      "path": "vae/config.json",
+      "role": "complete Mage VAE component",
+      "sha256": "abd124d603d6c6a03e9d0f2aa6d113b8c4afda0738400bdf2f99240aeaeaff76",
+      "size": 112
+    },
+    {
+      "path": "vae/diffusion_pytorch_model.safetensors",
+      "role": "complete Mage VAE component",
+      "sha256": "34e076dc1e8a15321e1e07be5111d59cf16dd10b804b7c7e20b4de29013427e0",
+      "size": 345053056
+    },
+    {
+      "path": "validate_release.py",
+      "role": "package hash validator",
+      "sha256": "9c14818d8939f880be8702222d0345ab0af44fbb8fefa98600d414e2c469eb75",
+      "size": 2877
+    },
+    {
+      "path": "vendor/mage_flow/__init__.py",
+      "role": "pinned Microsoft Mage inference source",
+      "sha256": "0709764f182b55fdec6b8195a4d18640fe0ffef3785d3977d8ebc29905de7489",
+      "size": 367
+    },
+    {
+      "path": "vendor/mage_flow/app.py",
+      "role": "pinned Microsoft Mage inference source",
+      "sha256": "ac0597feddf1b6c5aaa2f18cc3ebb3e690dc42232462ac778af45044df41435e",
+      "size": 9004
+    },
+    {
+      "path": "vendor/mage_flow/inference.py",
+      "role": "pinned Microsoft Mage inference source",
+      "sha256": "0a1e196f784f4daf4b4d1607cade1d066908341e0e444aa29fcea3967a8c1a3f",
+      "size": 7181
+    },
+    {
+      "path": "vendor/mage_flow/models/__init__.py",
+      "role": "pinned Microsoft Mage inference source",
+      "sha256": "7e5ba07fdb01f4a5912eaeb3afbe9d4ccc4e391e7de4cdc11865e5aa911ddfd5",
+      "size": 56
+    },
+    {
+      "path": "vendor/mage_flow/models/mage_flow.py",
+      "role": "pinned Microsoft Mage inference source",
+      "sha256": "59b6e1bee7f95a7fd2fa7bd9e765966832951e2bb184b5ae27283884997b845f",
+      "size": 15991
+    },
+    {
+      "path": "vendor/mage_flow/models/modules/__init__.py",
+      "role": "pinned Microsoft Mage inference source",
+      "sha256": "e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855",
+      "size": 0
+    },
+    {
+      "path": "vendor/mage_flow/models/modules/_attn_backend.py",
+      "role": "pinned Microsoft Mage inference source",
+      "sha256": "7d652eeb2ab39e37554e3d62859c74edeffc8dce70b2c15687060c259647b559",
+      "size": 7946
+    },
+    {
+      "path": "vendor/mage_flow/models/modules/mage_latent.py",
+      "role": "pinned Microsoft Mage inference source",
+      "sha256": "9ffd7c68b6053ad37f6bf0925d8d94dcd72af751448ee0228ec13018fedf54e5",
+      "size": 4118
+    },
+    {
+      "path": "vendor/mage_flow/models/modules/mage_layers.py",
+      "role": "pinned Microsoft Mage inference source",
+      "sha256": "4b198343b8929f48a0a14d388502c81f54be17b222f6831303d2da6a91f33a62",
+      "size": 30418
+    },
+    {
+      "path": "vendor/mage_flow/models/modules/mage_text.py",
+      "role": "pinned Microsoft Mage inference source",
+      "sha256": "eed7846d02bee28ebf7a1fb45db7d9f621dd66e702e761848d67bed426c1c4a7",
+      "size": 21247
+    },
+    {
+      "path": "vendor/mage_flow/models/modules/mage_vae.py",
+      "role": "pinned Microsoft Mage inference source",
+      "sha256": "64f4d7041003e416bc2f4fac5bbf8aabf2e7c798ad106682c34332ba347b0ef9",
+      "size": 25112
+    },
+    {
+      "path": "vendor/mage_flow/models/modules/text_encoder.py",
+      "role": "pinned Microsoft Mage inference source",
+      "sha256": "65e490ce35fbe4d4057f115be2ab8731e01ea44b618c10aabac80d51fc38ef81",
+      "size": 29405
+    },
+    {
+      "path": "vendor/mage_flow/models/utils.py",
+      "role": "pinned Microsoft Mage inference source",
+      "sha256": "0f242ca7a77e0f85b5985b5299304aa0786737de1ba414c7f8f47b8d664dbca8",
+      "size": 6649
+    },
+    {
+      "path": "vendor/mage_flow/pipeline.py",
+      "role": "pinned Microsoft Mage inference source",
+      "sha256": "b9fc57018570372dd3404a733e19b918a359188dd9f5ef7817c6b30969fc13db",
+      "size": 37089
+    }
+  ],
+  "measurements": {
+    "combined_text_stage_peak_allocated_bytes": 10437891584,
+    "combined_val07_latent_cosine": 0.9339298064776214,
+    "combined_val07_latent_nrmse": 0.36180871546574084,
+    "combined_val07_pixel_cosine": 0.9927827428055002,
+    "combined_val07_pixel_nrmse": 0.12000220627369493,
+    "full_generation_speed_claim_ready": false,
+    "text_encoder_allocated_saving_bytes": 4834599424,
+    "text_encoder_allocated_saving_gib": 4.5025715827941895,
+    "transformer_allocated_saving_bytes": 2604638208,
+    "transformer_allocated_saving_gib": 2.4257583618164062
+  },
+  "policy": {
+    "component_layout": "complete_huggingface_sharded_repository",
+    "mage_transformer": {
+      "format": "native resident W4A4 NVFP4",
+      "target_bf16_source_reads": 0,
+      "target_count": 48
+    },
+    "native_compute_capability": "12.0",
+    "qwen_text_encoder": {
+      "bf16_policy": "blocks 0/35, embeddings, norms, biases, vision tower",
+      "fp8_projection_count": 14,
+      "nvfp4_projection_count": 224
+    }
+  },
+  "release_status": "validated-standard-huggingface-candidate",
+  "schema_version": "mage-flow-nvfp4-huggingface-release-v1"
+}

README.md

@@ -0,0 +1,208 @@
+---
+license: other
+library_name: diffusers
+pipeline_tag: text-to-image
+base_model: microsoft/Mage-Flow
+base_model_relation: quantized
+tags:
+  - ajh
+  - mage-flow
+  - mage-flow-nvfp4-ajh
+  - nvfp4
+  - blackwell
+  - qwen3-vl
+  - text-to-image
+  - quantization
+---
+
+# Mage-Flow-NVFP4-AJH
+
+**Mage-Flow-NVFP4-AJH** is a portable, runnable NVFP4 package for
+[`microsoft/Mage-Flow`](https://huggingface.co/microsoft/Mage-Flow).
+It combines a native NVFP4 Mage transformer with a complete mixed
+NVFP4/FP8 Qwen3-VL text encoder.
+
+This is a complete Hugging Face component-layout repository, not an overlay:
+
+```text
+model_index.json
+transformer/
+  config.json
+  diffusion_pytorch_model-00001-of-00004.safetensors
+  diffusion_pytorch_model.safetensors.index.json
+text_encoder/
+  config.json
+  model.safetensors
+vae/
+scheduler/
+```
+
+The transformer shards include every retained BF16 tensor alongside the
+NVFP4 module state. The VAE, scheduler, text-encoder configuration, tokenizer,
+and processor are also included. Running the downloaded repository does not
+fetch BF16 weights from the base model.
+
+## Quantization policy
+
+Mage transformer:
+
+- 48 image/text MLP up/down projections use resident W4A4 NVFP4.
+- Attention, modulation, normalization, and all non-target tensors remain
+  BF16.
+- The 48 original BF16 target weight matrices are absent. Their biases remain
+  BF16 and are stored normally inside the quantized modules.
+- The complete transformer is stored as four indexed standard Safetensors
+  shards. Packed NVFP4 buffers use their module state-dict names:
+  `packed_weight`, `weight_scales`, `weight_scale`, and `bias`.
+
+Qwen3-VL text encoder:
+
+- Blocks 2–33: 224 NVFP4 projections.
+- Blocks 1 and 34: 14 FP8 projections.
+- Blocks 0 and 35, embeddings, norms, biases, and the vision tower remain
+  BF16.
+- Packaged text checkpoint size: `4,031,376,064` bytes.
+- The packaged loader constructs the model directly from this file: 475
+  non-quantized tensors and 238 packed projections, with no unresolved meta
+  tensors.
+
+## Requirements
+
+The prebuilt runtime was tested on:
+
+- NVIDIA RTX 50-series / SM120
+- Linux x86-64
+- CUDA 13.1
+- Python 3.11
+- PyTorch `2.13.0+cu130`
+- `comfy-kitchen==0.2.22`
+- `flash-attn==2.8.3`
+
+Create a local environment:
+
+```bash
+python3.11 -m venv .venv
+source .venv/bin/activate
+python -m pip install --upgrade pip
+python -m pip install -r requirements.txt
+CUDA_HOME=/usr/local/cuda-13.1 \
+  python -m pip install --no-build-isolation flash-attn==2.8.3
+```
+
+The included binaries are for the exact tested stack. Rebuild them after
+changing PyTorch, CUDA, or the C++ ABI:
+
+```bash
+CUDA_HOME=/usr/local/cuda-13.1 \
+PYTHON_BIN="$PWD/.venv/bin/python" \
+  ./build_native.sh
+```
+
+## Generate an image
+
+Expose exactly one SM120 GPU:
+
+```bash
+CUDA_VISIBLE_DEVICES=0 .venv/bin/python generate.py \
+  --prompt 'A detailed watercolor fox reading under an old oak tree' \
+  --output fox.png \
+  --height 1024 \
+  --width 1024 \
+  --steps 20 \
+  --seed 1
+```
+
+When running from another copy of the scripts, `--model` also accepts the Hub
+repository id or a downloaded standard-layout directory:
+
+```bash
+CUDA_VISIBLE_DEVICES=0 .venv/bin/python generate.py \
+  --model ajh-code/Mage-Flow-NVFP4-AJH \
+  --prompt 'A lighthouse poster reading "ARCTIC LOOP"' \
+  --output lighthouse.png
+```
+
+The command refuses to overwrite an existing output and writes a companion
+JSON report containing coverage, memory, timing, and environment information.
+
+## Measured results
+
+Transformer:
+
+- Loaded packed transformer allocation: `5,628,438,016` bytes.
+- Measured transformer allocation saving: `2,604,638,208` bytes
+  (`2.4258 GiB`).
+- Target BF16 source reads: `0`.
+- Representative real up-projection: `3.26x` eager and `3.25x` compiled
+  speedup over BF16.
+
+Text encoder:
+
+- BF16 language-stack residency: `7.5453 GiB`.
+- Packed mixed residency: `3.0427 GiB`.
+- Measured saving: `4.5026 GiB`.
+
+Combined validation on an RTX 5060 Ti 16 GB:
+
+- All seven coverage, source-access, allocation, finite, pixel, latent, and
+  text-policy gates passed.
+- Peak text stage with the packed transformer still resident:
+  `10,437,891,584` bytes (`9.72 GiB`).
+- VAL-07 pixel cosine/NRMSE versus frozen BF16:
+  `0.9927827428` / `0.1200022063`.
+- VAL-07 latent cosine/NRMSE versus frozen BF16:
+  `0.9339298065` / `0.3618087155`.
+- The requested strings `ARCTIC LOOP` and `NORTHERN COAST` were both rendered
+  correctly.
+
+![Combined native NVFP4 VAL-07 output](examples/VAL-07_combined_native.png)
+
+The standard-layout packaged loader is validated independently of the
+research-tree loader. It installs 48 native Mage projections and 238 packed
+Qwen projections and rejects any checkpoint containing the replaced BF16
+transformer targets.
+
+## Quality and speed caveats
+
+The mixed text encoder is functional and produced visually strong downstream
+images, but it does not meet our unusually strict embedding-similarity gate:
+mean token/pooled cosine was `0.9531366898` / `0.9751400001`. All ten content
+screening verdicts and category lists remained unchanged.
+
+The ten-case text workload was slower with the packed runtime (`30.36 s`)
+than BF16 (`18.29 s`). This package therefore claims major text-encoder VRAM
+savings, not a text-encoding speedup. Transformer projections are materially
+faster, but repeated matched end-to-end timing has not been completed.
+
+Additional limitations:
+
+- Native execution is currently SM120-only.
+- The repository follows the Hugging Face component and Safetensors layout,
+  but the new `mage_flow_nvfp4` runtime is not yet built into stock Diffusers.
+  Use the included loader.
+- CUDA graph compatibility is not claimed.
+- Generation and text-to-image are tested; Base, Turbo, and editing variants
+  are not.
+- Four transformer-only held-out cases and one combined held-out case have
+  been evaluated. This is not a broad benchmark.
+
+## Validate the download
+
+`MANIFEST.json` records every distributed file except itself:
+
+```bash
+.venv/bin/python validate_release.py
+```
+
+Hashing the transformer shards and text checkpoint can take a little while.
+
+## License and attribution
+
+- Mage-Flow and the vendored Mage inference source are Copyright (c) 2026
+  Microsoft and licensed under MIT. See `LICENSE` and
+  `licenses/MAGE-MIT.txt`.
+- Qwen3-VL and the mixed NVFP4/FP8 text checkpoint are licensed under
+  Apache-2.0. See `licenses/QWEN-APACHE-2.0.txt`.
+- The text checkpoint was produced by
+  [`InsecureErasure/Qwen3-VL-4B-Instruct-NVFP4`](https://huggingface.co/InsecureErasure/Qwen3-VL-4B-Instruct-NVFP4)
+  using learned rounding and `comfy-kitchen`.

THIRD_PARTY_NOTICES.md

@@ -0,0 +1,34 @@
+# Third-party notices
+
+## Microsoft Mage and Mage-Flow
+
+This package targets `microsoft/Mage-Flow` and includes the minimal Python
+inference sources under `vendor/mage_flow/`.
+
+- Upstream: <https://github.com/microsoft/Mage>
+- Pinned source commit: `df7f84d9f8fc991d189d929f03cff623b430a4a2`
+- License: MIT
+- License copy: `LICENSE` and `licenses/MAGE-MIT.txt`
+
+The Microsoft base checkpoint is not duplicated in this package. The
+portable loader downloads the required transformer, VAE, scheduler, and
+configuration files from `microsoft/Mage-Flow`, or accepts a local copy.
+
+## Qwen3-VL and the mixed NVFP4 checkpoint
+
+The packaged text encoder derives from `Qwen/Qwen3-VL-4B-Instruct` and is an
+exact copy of the learned-rounded mixed checkpoint published at
+`InsecureErasure/Qwen3-VL-4B-Instruct-NVFP4`.
+
+- Quantized upstream:
+  <https://huggingface.co/InsecureErasure/Qwen3-VL-4B-Instruct-NVFP4>
+- Pinned quantized revision:
+  `cf080d6af667f5f4949295bd09a1f72a218fe0ee`
+- Base model: <https://huggingface.co/Qwen/Qwen3-VL-4B-Instruct>
+- License: Apache-2.0
+- License copy: `licenses/QWEN-APACHE-2.0.txt`
+- Packaged artifact SHA-256:
+  `719906b435800757d22013d3d475a4853d59b779b022669fa8b8a193b85d0f41`
+
+The quantized checkpoint uses the ComfyUI `comfy_quant` representation and
+executes packed projections through `comfy-kitchen`.

build_manifest.py

@@ -0,0 +1,202 @@
+#!/usr/bin/env python3
+"""Create the immutable release manifest after all package files are frozen."""
+
+from __future__ import annotations
+
+import hashlib
+import json
+from pathlib import Path
+
+
+RELEASE_ROOT = Path(__file__).resolve().parent
+MANIFEST_PATH = RELEASE_ROOT / "MANIFEST.json"
+
+
+def sha256(path: Path) -> str:
+    digest = hashlib.sha256()
+    with path.open("rb") as handle:
+        for chunk in iter(lambda: handle.read(1 << 20), b""):
+            digest.update(chunk)
+    return digest.hexdigest()
+
+
+def role_for(path: Path) -> str:
+    value = path.as_posix()
+    if value == "README.md":
+        return "Hugging Face model card and usage guide"
+    if value == "model_index.json":
+        return "Hugging Face Diffusers metadata and download-tracking query file"
+    if value == "generate.py":
+        return "portable single-image inference entry point"
+    if value == "run_validation.py":
+        return "portable frozen-prompt validation entry point"
+    if value == "validate_release.py":
+        return "package hash validator"
+    if value == "build_native.sh":
+        return "project-local native runtime build helper"
+    if value == "build_manifest.py":
+        return "release manifest builder"
+    if value.startswith("transformer/") and value.endswith(".safetensors"):
+        return "complete sharded Mage transformer checkpoint with native NVFP4 modules"
+    if value.startswith("transformer/"):
+        return "Mage transformer component configuration, index, or quantization metadata"
+    if value.startswith("text_encoder/") and value.endswith(".safetensors"):
+        return "complete mixed NVFP4/FP8 Qwen3-VL text checkpoint"
+    if value.startswith("text_encoder/"):
+        return "packaged Qwen3-VL configuration or tokenizer asset"
+    if value.startswith("vae/"):
+        return "complete Mage VAE component"
+    if value.startswith("scheduler/"):
+        return "Mage FlowMatch scheduler component"
+    if value.startswith("runtime/"):
+        return "native runtime, loader, source, or Python integration"
+    if value.startswith("vendor/"):
+        return "pinned Microsoft Mage inference source"
+    if value.startswith("licenses/") or value in {
+        "LICENSE",
+        "THIRD_PARTY_NOTICES.md",
+    }:
+        return "license or third-party attribution"
+    if value.startswith("examples/"):
+        return "validated combined NVFP4 example output"
+    return "release support file"
+
+
+def package_files() -> list[dict[str, object]]:
+    records: list[dict[str, object]] = []
+    for path in sorted(RELEASE_ROOT.rglob("*")):
+        if not path.is_file() or path == MANIFEST_PATH:
+            continue
+        relative = path.relative_to(RELEASE_ROOT)
+        if "__pycache__" in relative.parts or path.suffix == ".pyc":
+            raise RuntimeError(
+                f"generated Python cache must not be shipped: {relative}"
+            )
+        if path.is_symlink():
+            raise RuntimeError(f"release files must be independent: {relative}")
+        records.append(
+            {
+                "path": relative.as_posix(),
+                "size": path.stat().st_size,
+                "sha256": sha256(path),
+                "role": role_for(relative),
+            }
+        )
+    return records
+
+
+def main() -> int:
+    if MANIFEST_PATH.exists():
+        raise SystemExit(f"refusing to overwrite {MANIFEST_PATH}")
+    payload = {
+        "schema_version": "mage-flow-nvfp4-huggingface-release-v1",
+        "release_status": "validated-standard-huggingface-candidate",
+        "base_models": {
+            "mage_flow": "microsoft/Mage-Flow",
+            "mage_source_commit": (
+                "df7f84d9f8fc991d189d929f03cff623b430a4a2"
+            ),
+            "mage_transformer_sha256": (
+                "04431abbc3acd1a5b86a7f77269f8e77e05a6cac3b21f35e1aef7493f6ab4934"
+            ),
+            "qwen": "Qwen/Qwen3-VL-4B-Instruct",
+            "qwen_quantized": (
+                "InsecureErasure/Qwen3-VL-4B-Instruct-NVFP4"
+            ),
+            "qwen_quantized_revision": (
+                "cf080d6af667f5f4949295bd09a1f72a218fe0ee"
+            ),
+        },
+        "policy": {
+            "mage_transformer": {
+                "format": "native resident W4A4 NVFP4",
+                "target_count": 48,
+                "target_bf16_source_reads": 0,
+            },
+            "qwen_text_encoder": {
+                "nvfp4_projection_count": 224,
+                "fp8_projection_count": 14,
+                "bf16_policy": (
+                    "blocks 0/35, embeddings, norms, biases, vision tower"
+                ),
+            },
+            "native_compute_capability": "12.0",
+            "component_layout": "complete_huggingface_sharded_repository",
+        },
+        "measurements": {
+            "transformer_allocated_saving_bytes": 2_604_638_208,
+            "transformer_allocated_saving_gib": 2.4257583618164062,
+            "text_encoder_allocated_saving_bytes": 4_834_599_424,
+            "text_encoder_allocated_saving_gib": 4.5025715827941895,
+            "combined_text_stage_peak_allocated_bytes": 10_437_891_584,
+            "combined_val07_pixel_cosine": 0.9927827428055002,
+            "combined_val07_pixel_nrmse": 0.12000220627369493,
+            "combined_val07_latent_cosine": 0.9339298064776214,
+            "combined_val07_latent_nrmse": 0.36180871546574084,
+            "full_generation_speed_claim_ready": False,
+        },
+        "evidence": {
+            "combined_frozen_val07": {
+                "result_sha256": (
+                    "67afd8ca0c00fff949bfc4c8b660ca864617839cbc01231a7157769a964287c7"
+                ),
+                "image_sha256": (
+                    "857c1663636cd775ea05ffe37a98ac47ab7cd4e5362dd82c21ae1b53f7d71c3a"
+                ),
+                "all_seven_gates_passed": True,
+            },
+            "legacy_portable_package_generation": {
+                "report_sha256": (
+                    "2877697cdf8d469f689aa68751d24c8f25288388f682e61355234e82fae35eda"
+                ),
+                "image_sha256": (
+                    "2c1f2853e37a2ede976f24bc9b7eef172aed7665266c615401ea189fcbcee19f"
+                ),
+                "status": "success",
+            },
+            "standard_huggingface_layout_generation": {
+                "report_sha256": (
+                    "61c8b651a5a9814c20fa4a99719e0b0a6caa827f61c0865993d916de25e6507d"
+                ),
+                "image_sha256": (
+                    "2c1f2853e37a2ede976f24bc9b7eef172aed7665266c615401ea189fcbcee19f"
+                ),
+                "layout": "complete_huggingface_sharded_repository",
+                "transformer_shard_count": 4,
+                "transformer_checkpoint_tensor_count": 493,
+                "loaded_non_target_tensor_count": 301,
+                "loaded_nvfp4_projection_count": 48,
+                "bf16_target_weight_reads": 0,
+                "status": "success",
+            },
+            "native_text_quality_gate_result_sha256": (
+                "4a31a2c0eb4cac7401224eaeab23577bd31cc5ef69f6edec5289757f113373e6"
+            ),
+            "transformer_quality_results": {
+                "VAL-01": (
+                    "82fadbc86ffb980a503c61ba50bfae25fd302503cf9ebb807029aacdc484f65a"
+                ),
+                "VAL-07": (
+                    "2476505becf3cec14d174c35b42a92ae180cea1173ebe376d0d40f530a0292ce"
+                ),
+                "VAL-09": (
+                    "c5d25ab78790295a0945b9783e146ff114991cc72af80a524d083034696110fc"
+                ),
+                "VAL-13": (
+                    "d2f5b3e7b3be8d430fbd8ca47abdfa405d130f080d45a51fe5ab90fde309a07e"
+                ),
+            },
+        },
+        "files": package_files(),
+    }
+    with MANIFEST_PATH.open("xb") as handle:
+        handle.write(
+            json.dumps(payload, indent=2, sort_keys=True).encode("utf-8")
+            + b"\n"
+        )
+    print(f"wrote {MANIFEST_PATH} with {len(payload['files'])} files")
+    return 0
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())

build_native.sh

@@ -0,0 +1,61 @@
+#!/usr/bin/env bash
+set -euo pipefail
+
+RELEASE_ROOT="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
+RUNTIME_ROOT="${RELEASE_ROOT}/runtime"
+PYTHON_BIN="${PYTHON_BIN:-python}"
+CUDA_TOOLKIT_ROOT="${CUDA_HOME:-/usr/local/cuda}"
+NVCC="${CUDA_TOOLKIT_ROOT}/bin/nvcc"
+
+if [[ ! -x "${NVCC}" ]]; then
+  echo "nvcc was not found at ${NVCC}" >&2
+  exit 1
+fi
+
+mkdir -p "${RUNTIME_ROOT}"
+
+"${NVCC}" \
+  -std=c++17 -O3 -arch=sm_120a -shared -Xcompiler=-fPIC \
+  "${RUNTIME_ROOT}/nvfp4_linear.cu" \
+  -o "${RUNTIME_ROOT}/libmage_nvfp4_linear.so" \
+  -L"${CUDA_TOOLKIT_ROOT}/lib64" -lcublasLt -lcublas
+
+PYTORCH_INCLUDE_FLAGS="$("${PYTHON_BIN}" - <<'PY'
+from torch.utils.cpp_extension import include_paths
+print(" ".join(f"-I{path}" for path in include_paths()))
+PY
+)"
+PYTORCH_LIBRARY_FLAGS="$("${PYTHON_BIN}" - <<'PY'
+from torch.utils.cpp_extension import library_paths
+print(" ".join(f"-L{path}" for path in library_paths()))
+PY
+)"
+PYTORCH_RPATH_FLAGS="$("${PYTHON_BIN}" - <<'PY'
+from torch.utils.cpp_extension import library_paths
+print(" ".join(f"-Wl,-rpath,{path}" for path in library_paths()))
+PY
+)"
+PYTORCH_ABI_FLAG="$("${PYTHON_BIN}" - <<'PY'
+import torch
+print(f"-D_GLIBCXX_USE_CXX11_ABI={int(torch.compiled_with_cxx11_abi())}")
+PY
+)"
+
+g++ \
+  -std=c++20 -O3 -shared -fPIC \
+  ${PYTORCH_ABI_FLAG} \
+  ${PYTORCH_INCLUDE_FLAGS} \
+  ${PYTORCH_LIBRARY_FLAGS} \
+  ${PYTORCH_RPATH_FLAGS} \
+  -I"${CUDA_TOOLKIT_ROOT}/include" \
+  -I"${RUNTIME_ROOT}" \
+  -o "${RUNTIME_ROOT}/libmage_nvfp4_torch_op.so" \
+  "${RUNTIME_ROOT}/sm120_linear_op.cpp" \
+  -L"${CUDA_TOOLKIT_ROOT}/lib64" \
+  -L"${RUNTIME_ROOT}" \
+  -Wl,-rpath,"${CUDA_TOOLKIT_ROOT}/lib64" \
+  -Wl,-rpath,'$ORIGIN' \
+  -ltorch -ltorch_cpu -ltorch_cuda -lc10 -lc10_cuda -lcudart \
+  -lmage_nvfp4_linear
+
+echo "Built the packaged SM120 runtime in ${RUNTIME_ROOT}"

generate.py

@@ -0,0 +1,478 @@
+#!/usr/bin/env python3
+"""Portable single-image Mage-Flow NVFP4 inference entry point."""
+
+from __future__ import annotations
+
+import argparse
+from datetime import datetime, timezone
+import gc
+import json
+import os
+from pathlib import Path
+import platform
+import sys
+import time
+from typing import Any, Callable
+
+
+RELEASE_ROOT = Path(__file__).resolve().parent
+RUNTIME_ROOT = RELEASE_ROOT / "runtime"
+VENDOR_ROOT = RELEASE_ROOT / "vendor"
+DEFAULT_MODEL = str(RELEASE_ROOT)
+
+for import_root in (RUNTIME_ROOT, VENDOR_ROOT):
+    if str(import_root) not in sys.path:
+        sys.path.insert(0, str(import_root))
+
+
+def utc_now() -> str:
+    return (
+        datetime.now(timezone.utc)
+        .replace(microsecond=0)
+        .isoformat()
+        .replace("+00:00", "Z")
+    )
+
+
+def resolve_model(value: str) -> Path:
+    candidate = Path(value).expanduser()
+    if candidate.is_dir():
+        return candidate.resolve()
+    from huggingface_hub import snapshot_download
+
+    return Path(
+        snapshot_download(
+            repo_id=value,
+            allow_patterns=[
+                "model_index.json",
+                "transformer/config.json",
+                "transformer/*.safetensors",
+                "transformer/*.json",
+                "text_encoder/*",
+                "vae/config.json",
+                "vae/*.safetensors",
+                "scheduler/*",
+            ],
+        )
+    ).resolve()
+
+
+def _repo_subpath(repo_dir: Path, relative: str) -> str:
+    path = (repo_dir / relative).resolve()
+    if not path.is_relative_to(repo_dir):
+        raise ValueError(f"model path escapes its snapshot: {relative}")
+    return str(path)
+
+
+def _structure_from_config(transformer_config: dict[str, Any]) -> dict[str, Any]:
+    metadata_keys = {
+        "_class_name",
+        "txt_max_length",
+        "max_sequence_length",
+        "param_dtype",
+        "packing",
+        "schedule_mode",
+        "static_shift",
+        "use_time_shift",
+        "rope_type",
+        "apply_text_rotary_emb",
+        "mlp_ratio",
+        "depth_single_blocks",
+        "theta",
+        "qkv_bias",
+        "guidance_embed",
+        "vec_in_dim",
+        "vec_type",
+        "time_type",
+        "double_block_type",
+        "quantization_config",
+    }
+    return {
+        key: value
+        for key, value in transformer_config.items()
+        if key not in metadata_keys
+    }
+
+
+def load_pipeline(
+    *,
+    model: str,
+    torch: Any,
+) -> tuple[Any, dict[str, Any]]:
+    import torch.nn as nn
+    from diffusers import FlowMatchEulerDiscreteScheduler
+    from mage_flow.models.mage_flow import MageFlowModel, ModelConfig
+    from mage_flow.models.modules._attn_backend import set_attn_backend
+    from mage_flow.pipeline import MageFlowPipeline
+    from quant_text_encoder import load_quantized_text_encoder
+    from standard_transformer import load_standard_native_transformer
+
+    repo_dir = resolve_model(model)
+    model_index_path = repo_dir / "model_index.json"
+    transformer_config_path = repo_dir / "transformer" / "config.json"
+    model_index = json.loads(model_index_path.read_text(encoding="utf-8"))
+    transformer_config = json.loads(
+        transformer_config_path.read_text(encoding="utf-8")
+    )
+    structure = _structure_from_config(transformer_config)
+    config = ModelConfig(
+        vae_path=_repo_subpath(repo_dir, model_index["_vae_source"]),
+        txt_enc_path=_repo_subpath(repo_dir, model_index["_text_encoder_path"]),
+        model_structure=structure,
+        txt_max_length=transformer_config.get("txt_max_length", 2048),
+        packing=transformer_config.get("packing", True),
+        static_shift=transformer_config.get("static_shift", 6.0),
+    )
+
+    transformer, access_report = load_standard_native_transformer(
+        repo_dir,
+        torch.device("cuda:0"),
+    )
+    model = MageFlowModel.__new__(MageFlowModel)
+    nn.Module.__init__(model)
+    model.config = config
+    set_attn_backend(getattr(config, "attn_type", "flash2"))
+    model.patch_text_encoder_forward()
+    model.vae = model.load_vae()
+    model.transformer = transformer
+    text_encoder_dir = Path(
+        _repo_subpath(repo_dir, model_index["_text_encoder_path"])
+    )
+    model.txt_enc, text_report = load_quantized_text_encoder(
+        text_encoder_dir=text_encoder_dir,
+        artifact_path=text_encoder_dir / "model.safetensors",
+        tokenizer_max_length=config.txt_max_length,
+        dit_structure=structure,
+        use_packed_text_infer=config.packing,
+    )
+    model.vae.requires_grad_(False).to(torch.bfloat16)
+    model.txt_enc.requires_grad_(False)
+    model.eval()
+    model.scheduler = FlowMatchEulerDiscreteScheduler.from_pretrained(
+        _repo_subpath(repo_dir, "scheduler")
+    )
+    return (
+        MageFlowPipeline(model, device="cuda:0"),
+        {
+            "resolved_model": str(repo_dir),
+            "transformer_access": access_report,
+            "text_encoder": text_report,
+        },
+    )
+
+
+def stage_to_gpu(
+    module: Any,
+    operation: Callable[[], Any],
+    torch: Any,
+) -> tuple[Any, dict[str, Any]]:
+    torch.cuda.synchronize()
+    torch.cuda.reset_peak_memory_stats(0)
+    started = time.perf_counter()
+    try:
+        module.to("cuda:0")
+        value = operation()
+        torch.cuda.synchronize()
+        metrics = {
+            "peak_allocated_bytes": int(torch.cuda.max_memory_allocated(0)),
+            "peak_reserved_bytes": int(torch.cuda.max_memory_reserved(0)),
+        }
+    finally:
+        module.to("cpu")
+        torch.cuda.synchronize()
+        torch.cuda.empty_cache()
+    metrics["seconds"] = time.perf_counter() - started
+    return value, metrics
+
+
+def generate_staged(
+    *,
+    pipe: Any,
+    prompt: str,
+    negative_prompt: str,
+    height: int,
+    width: int,
+    steps: int,
+    cfg: float,
+    seed: int,
+    static_shift: float,
+    torch: Any,
+) -> tuple[Any, dict[str, Any], dict[str, Any]]:
+    from einops import rearrange
+    from mage_flow.models.modules.mage_latent import encode_noise, resolve_gs_key
+    from mage_flow.pipeline import (
+        _build_pack_ctx,
+        _decode_one,
+        _encode_texts_packed,
+        _get_scheduler,
+        _lens_to_cu,
+        _make_divisible_by_16,
+        _slice_packed,
+        _template_info,
+        _velocity,
+    )
+
+    model = pipe.model
+    device = torch.device("cuda:0")
+    template_info = _template_info("mage-flow")
+    template = template_info.get("template", "{}")
+    drop_index = int(template_info.get("start_idx", 0))
+    stage_metrics: dict[str, Any] = {}
+
+    def encode_text() -> tuple[Any, ...]:
+        verdict = model.txt_enc.screen_text(prompt)
+        if verdict.violates:
+            return (verdict,)
+        text_flat, vec_all, text_lens = _encode_texts_packed(
+            model,
+            [prompt, negative_prompt or " "],
+            template,
+            drop_index,
+            device,
+        )
+        positive = _slice_packed(
+            text_flat, vec_all, text_lens, 0, 1, device
+        )
+        negative = _slice_packed(
+            text_flat, vec_all, text_lens, 1, 1, device
+        )
+        return (verdict, *positive, *negative)
+
+    encoded, stage_metrics["text_encode"] = stage_to_gpu(
+        model.txt_enc,
+        encode_text,
+        torch,
+    )
+    verdict = encoded[0]
+    if verdict.violates:
+        raise RuntimeError(
+            "prompt was refused by Mage content screening: "
+            f"{verdict.categories} {verdict.reason}"
+        )
+    (
+        _verdict,
+        txt,
+        txt_cu,
+        txt_mask,
+        vec,
+        neg_txt,
+        neg_cu,
+        neg_mask,
+        neg_vec,
+    ) = encoded
+
+    height = _make_divisible_by_16(height)
+    width = _make_divisible_by_16(width)
+    noise = encode_noise(
+        (
+            model.vae.latent_channels,
+            (height + 15) // 16,
+            (width + 15) // 16,
+        ),
+        key=resolve_gs_key(None),
+        seed=seed,
+        device=device,
+        dtype=torch.bfloat16,
+    )
+    _, _, grid_h, grid_w = noise.shape
+    image_latent = rearrange(noise, "b c h w -> b (h w) c")
+    image_ids = torch.zeros(grid_h, grid_w, 3, device=device)
+    image_ids[..., 1] += torch.arange(grid_h, device=device)[:, None]
+    image_ids[..., 2] += torch.arange(grid_w, device=device)[None, :]
+    image_ids = rearrange(image_ids, "h w c -> 1 (h w) c")
+    image_lens = [grid_h * grid_w]
+    image_cu = _lens_to_cu(image_lens, device)
+    context = _build_pack_ctx(
+        image_ids,
+        image_cu,
+        [[(1, grid_h, grid_w)]],
+        image_lens,
+        txt,
+        txt_cu,
+        txt_mask,
+        vec,
+        neg_txt,
+        neg_cu,
+        neg_mask,
+        neg_vec,
+        cfg,
+        False,
+        True,
+        device,
+    )
+
+    def denoise() -> Any:
+        nonlocal image_latent
+        scheduler = _get_scheduler(
+            model,
+            steps,
+            device,
+            static_shift,
+        )
+        for step_index, timestep in enumerate(scheduler.timesteps):
+            prediction = _velocity(
+                model.transformer,
+                image_latent,
+                context,
+                scheduler.sigmas[step_index].item(),
+            )
+            image_latent = scheduler.step(
+                prediction,
+                timestep,
+                image_latent,
+                return_dict=False,
+            )[0]
+        if int((~torch.isfinite(image_latent)).sum().item()) != 0:
+            raise RuntimeError("denoising produced non-finite values")
+        return image_latent.to("cpu")
+
+    latent_cpu, stage_metrics["denoise"] = stage_to_gpu(
+        model.transformer,
+        denoise,
+        torch,
+    )
+
+    def decode() -> Any:
+        return _decode_one(
+            model,
+            latent_cpu.to(device),
+            height,
+            width,
+            device,
+        )
+
+    image, stage_metrics["decode"] = stage_to_gpu(
+        model.vae,
+        decode,
+        torch,
+    )
+    screening = {
+        "violates": bool(verdict.violates),
+        "categories": list(verdict.categories or []),
+        "reason": str(verdict.reason),
+    }
+    return image, stage_metrics, screening
+
+
+def parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument("--prompt", required=True)
+    parser.add_argument("--negative-prompt", default="")
+    parser.add_argument("--output", type=Path, default=Path("mage_nvfp4.png"))
+    parser.add_argument(
+        "--model",
+        default=DEFAULT_MODEL,
+        help=(
+            "local standard-layout model directory or Hugging Face repo id "
+            "(default: this downloaded repository)"
+        ),
+    )
+    parser.add_argument("--height", type=int, default=1024)
+    parser.add_argument("--width", type=int, default=1024)
+    parser.add_argument("--steps", type=int, default=20)
+    parser.add_argument("--cfg", type=float, default=5.0)
+    parser.add_argument("--seed", type=int, default=1)
+    parser.add_argument("--static-shift", type=float, default=6.0)
+    return parser.parse_args()
+
+
+def main() -> int:
+    args = parse_args()
+    output_path = args.output.expanduser().resolve()
+    report_path = output_path.with_suffix(output_path.suffix + ".json")
+    if output_path.exists() or report_path.exists():
+        raise SystemExit(
+            f"refusing to overwrite existing output/report: {output_path}"
+        )
+    output_path.parent.mkdir(parents=True, exist_ok=True)
+
+    import torch
+    from packed_nvfp4_linear import close_all_contexts
+    from torch_ops_native import (
+        close_native_contexts,
+        initialize_native_sm120_op,
+    )
+
+    if not torch.cuda.is_available() or torch.cuda.device_count() != 1:
+        raise SystemExit(
+            "exactly one visible CUDA GPU is required; set CUDA_VISIBLE_DEVICES"
+        )
+    torch.cuda.set_device(0)
+    properties = torch.cuda.get_device_properties(0)
+    if (properties.major, properties.minor) != (12, 0):
+        raise SystemExit(
+            f"native NVFP4 requires SM120; found {properties.major}.{properties.minor}"
+        )
+    if not initialize_native_sm120_op(allow_python_schema_fallback=False):
+        raise SystemExit("the packaged native SM120 torch op did not load")
+
+    os.environ.setdefault("CUBLAS_WORKSPACE_CONFIG", ":4096:8")
+    torch.manual_seed(args.seed)
+    torch.cuda.manual_seed_all(args.seed)
+    torch.backends.cudnn.benchmark = False
+    torch.backends.cudnn.deterministic = True
+    torch.backends.cuda.matmul.allow_tf32 = False
+    torch.use_deterministic_algorithms(True)
+
+    started = time.perf_counter()
+    pipe = None
+    try:
+        pipe, load_report = load_pipeline(
+            model=args.model,
+            torch=torch,
+        )
+        image, stages, screening = generate_staged(
+            pipe=pipe,
+            prompt=args.prompt,
+            negative_prompt=args.negative_prompt,
+            height=args.height,
+            width=args.width,
+            steps=args.steps,
+            cfg=args.cfg,
+            seed=args.seed,
+            static_shift=args.static_shift,
+            torch=torch,
+        )
+        image.save(output_path)
+        report = {
+            "schema_version": "mage-flow-nvfp4-portable-generation-v1",
+            "status": "success",
+            "completed_at_utc": utc_now(),
+            "output": str(output_path),
+            "prompt": args.prompt,
+            "negative_prompt": args.negative_prompt,
+            "height": args.height,
+            "width": args.width,
+            "steps": args.steps,
+            "cfg": args.cfg,
+            "seed": args.seed,
+            "elapsed_seconds": time.perf_counter() - started,
+            "stages": stages,
+            "screening": screening,
+            "load": load_report,
+            "environment": {
+                "python": platform.python_version(),
+                "torch": torch.__version__,
+                "torch_cuda": torch.version.cuda,
+                "gpu": properties.name,
+                "compute_capability": (
+                    f"{properties.major}.{properties.minor}"
+                ),
+            },
+        }
+        report_path.write_text(
+            json.dumps(report, indent=2, sort_keys=True) + "\n",
+            encoding="utf-8",
+        )
+        print(f"saved {output_path}")
+        print(f"saved {report_path}")
+        return 0
+    finally:
+        close_native_contexts()
+        close_all_contexts()
+        pipe = None
+        gc.collect()
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())

licenses/MAGE-MIT.txt

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

licenses/QWEN-APACHE-2.0.txt

@@ -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.

model_index.json

@@ -0,0 +1,30 @@
+{
+  "_class_name": "MageFlowPipeline",
+  "_mage_flow_version": "0.1.0",
+  "_quantization": {
+    "text_encoder": "comfy_kitchen_nvfp4_fp8_mixed",
+    "transformer": "mage_flow_nvfp4"
+  },
+  "_text_encoder_path": "text_encoder",
+  "_vae_source": "vae/diffusion_pytorch_model.safetensors",
+  "scheduler": [
+    "diffusers",
+    "FlowMatchEulerDiscreteScheduler"
+  ],
+  "text_encoder": [
+    "transformers",
+    "Qwen3VLForConditionalGeneration"
+  ],
+  "tokenizer": [
+    "transformers",
+    "AutoProcessor"
+  ],
+  "transformer": [
+    "mage_flow",
+    "MageFlow"
+  ],
+  "vae": [
+    "mage_flow",
+    "MageVAE"
+  ]
+}

requirements.txt

@@ -0,0 +1,20 @@
+# Exact tested environment. Install only in a project-local virtual environment.
+--extra-index-url https://download.pytorch.org/whl/cu130
+torch==2.13.0+cu130
+torchvision==0.28.0+cu130
+diffusers==0.38.0
+transformers==5.5.0
+safetensors==0.8.0
+accelerate==1.13.0
+comfy-kitchen==0.2.22
+einops==0.8.2
+pydantic==2.12.5
+pillow==12.3.0
+numpy==2.4.3
+loguru==0.7.3
+huggingface-hub==1.24.0
+setuptools
+wheel
+ninja
+
+# Install flash-attn==2.8.3 separately with --no-build-isolation after this file.

run_validation.py

@@ -0,0 +1,72 @@
+#!/usr/bin/env python3
+"""Exercise the portable package with the frozen VAL-07 typography prompt."""
+
+from __future__ import annotations
+
+import argparse
+from pathlib import Path
+import subprocess
+import sys
+
+
+RELEASE_ROOT = Path(__file__).resolve().parent
+RUNNER = RELEASE_ROOT / "generate.py"
+VALIDATOR = RELEASE_ROOT / "validate_release.py"
+PROMPT = (
+    'Minimal travel postcard with the exact large text "ARCTIC LOOP" and '
+    'the exact small text "NORTHERN COAST"; an ivory lighthouse on a '
+    "slate-blue cliff below the lettering, crisp vector print."
+)
+NEGATIVE_PROMPT = (
+    "lowres, blurry, jpeg artifacts, watermark, signature, logo, unreadable "
+    "text, misspelled text, duplicate subject, deformed anatomy, extra "
+    "fingers, missing fingers, cropped hands"
+)
+
+
+def parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument("--cpu-self-check", action="store_true")
+    parser.add_argument("--model", default=str(RELEASE_ROOT))
+    parser.add_argument("--output-dir", type=Path)
+    return parser.parse_args()
+
+
+def main() -> int:
+    args = parse_args()
+    subprocess.run([sys.executable, str(VALIDATOR)], check=True)
+    if args.cpu_self_check:
+        return 0
+    if args.output_dir is None:
+        raise SystemExit("--output-dir is required unless --cpu-self-check is used")
+    output_dir = args.output_dir.resolve()
+    output_path = output_dir / "VAL-07_portable_combined.png"
+    command = [
+        sys.executable,
+        str(RUNNER),
+        "--model",
+        args.model,
+        "--prompt",
+        PROMPT,
+        "--negative-prompt",
+        NEGATIVE_PROMPT,
+        "--output",
+        str(output_path),
+        "--height",
+        "512",
+        "--width",
+        "512",
+        "--steps",
+        "20",
+        "--cfg",
+        "5.0",
+        "--seed",
+        "27182819",
+        "--static-shift",
+        "6.0",
+    ]
+    return subprocess.run(command, check=False).returncode
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())

runtime/nvfp4_linear.cu

@@ -0,0 +1,895 @@
+#include "nvfp4_linear.h"
+
+#include <cuda_bf16.h>
+#include <cuda_fp4.h>
+#include <cuda_fp8.h>
+#include <cuda_runtime.h>
+#include <cublasLt.h>
+
+#include <algorithm>
+#include <cmath>
+#include <cstdint>
+#include <cstring>
+#include <limits>
+#include <mutex>
+#include <stdexcept>
+#include <string>
+
+namespace {
+
+constexpr int kAbiVersion = 1;
+constexpr int kFp4BlockElements = 16;
+constexpr int kScaleTileOuter = 128;
+constexpr int kScaleTileInner = 4;
+constexpr int kWarpsPerQuantBlock = 8;
+constexpr int kQuantThreads = 32 * kWarpsPerQuantBlock;
+constexpr int kReduceThreads = 256;
+constexpr int kReduceItemsPerThread = 4;
+constexpr int kBiasThreads = 256;
+constexpr float kFp4E2M1Max = 6.0f;
+constexpr float kFp4TensorScaleMax = 448.0f;
+constexpr float kTensorScaleDenominator =
+    kFp4E2M1Max * kFp4TensorScaleMax;
+constexpr size_t kWorkspaceBytes = 64ull * 1024ull * 1024ull;
+
+thread_local std::string g_last_error;
+
+[[noreturn]] void fail(const std::string& message) {
+  throw std::runtime_error(message);
+}
+
+#define CUDA_CHECK(expr)                                                     \
+  do {                                                                       \
+    const cudaError_t status_ = (expr);                                      \
+    if (status_ != cudaSuccess) {                                            \
+      fail(std::string(#expr) + ": " + cudaGetErrorString(status_));         \
+    }                                                                        \
+  } while (0)
+
+#define CUBLASLT_CHECK(expr)                                                 \
+  do {                                                                       \
+    const cublasStatus_t status_ = (expr);                                   \
+    if (status_ != CUBLAS_STATUS_SUCCESS) {                                  \
+      fail(std::string(#expr) + " failed with cuBLASLt status " +            \
+           std::to_string(static_cast<int>(status_)));                       \
+    }                                                                        \
+  } while (0)
+
+int round_up(int value, int multiple) {
+  if (value <= 0 || multiple <= 0 ||
+      value > std::numeric_limits<int>::max() - (multiple - 1)) {
+    fail("invalid or overflowing round_up arguments");
+  }
+  return ((value + multiple - 1) / multiple) * multiple;
+}
+
+size_t checked_multiply(size_t left, size_t right, const char* label) {
+  if (right != 0 && left > std::numeric_limits<size_t>::max() / right) {
+    fail(std::string(label) + " size overflows size_t");
+  }
+  return left * right;
+}
+
+size_t round_up_divide(size_t dividend, size_t divisor) {
+  return (dividend + divisor - 1) / divisor;
+}
+
+struct ScaleLayout {
+  int inner_dim = 0;
+  int outer_tiles = 0;
+  size_t bytes = 0;
+};
+
+ScaleLayout make_scale_layout(int rows_k, int outer_columns) {
+  if (rows_k <= 0 || rows_k % 16 != 0 || outer_columns <= 0) {
+    fail("scale layout requires positive K divisible by 16 and positive M/N");
+  }
+  ScaleLayout layout;
+  layout.inner_dim = round_up(rows_k / kFp4BlockElements, kScaleTileInner);
+  layout.outer_tiles =
+      (outer_columns + kScaleTileOuter - 1) / kScaleTileOuter;
+  layout.bytes = checked_multiply(
+      checked_multiply(static_cast<size_t>(layout.outer_tiles),
+                       static_cast<size_t>(layout.inner_dim), "scale tensor"),
+      static_cast<size_t>(kScaleTileOuter), "scale tensor");
+  return layout;
+}
+
+size_t packed_weight_bytes_checked(int n, int k) {
+  if (n <= 0 || k <= 0 || n % 8 != 0 || k % 32 != 0) {
+    fail("NVFP4 weight requires N divisible by 8 and K divisible by 32");
+  }
+  return checked_multiply(static_cast<size_t>(n),
+                          static_cast<size_t>(k), "packed weight") /
+         2;
+}
+
+size_t host_scale_offset(int outer, int inner_scale, int scale_inner_dim) {
+  const int outer_tile = outer / kScaleTileOuter;
+  const int local_outer = outer % kScaleTileOuter;
+  const int local_inner = inner_scale % kScaleTileInner;
+  const int inner_tile_start = inner_scale - local_inner;
+  const size_t tile_base =
+      static_cast<size_t>(inner_tile_start +
+                          outer_tile * scale_inner_dim) *
+      kScaleTileOuter;
+  return tile_base + static_cast<size_t>(local_outer % 32) * 16 +
+         static_cast<size_t>(local_outer / 32) * 4 + local_inner;
+}
+
+__device__ __forceinline__ size_t device_scale_offset(
+    int outer, int inner_scale, int scale_inner_dim) {
+  const int outer_tile = outer / kScaleTileOuter;
+  const int local_outer = outer % kScaleTileOuter;
+  const int local_inner = inner_scale % kScaleTileInner;
+  const int inner_tile_start = inner_scale - local_inner;
+  const size_t tile_base =
+      static_cast<size_t>(inner_tile_start +
+                          outer_tile * scale_inner_dim) *
+      kScaleTileOuter;
+  return tile_base + static_cast<size_t>(local_outer % 32) * 16 +
+         static_cast<size_t>(local_outer / 32) * 4 + local_inner;
+}
+
+float host_ue4m3_to_float(uint8_t raw) {
+  const __half_raw half_raw = __nv_cvt_fp8_to_halfraw(raw, __NV_E4M3);
+  return __half2float(static_cast<__half>(half_raw));
+}
+
+__device__ __forceinline__ float device_ue4m3_to_float(uint8_t raw) {
+  const __half_raw half_raw = __nv_cvt_fp8_to_halfraw(raw, __NV_E4M3);
+  return __half2float(static_cast<__half>(half_raw));
+}
+
+__device__ __forceinline__ float tensor_scale_from_amax(float amax) {
+  return amax == 0.0f ? 1.0f : amax / kTensorScaleDenominator;
+}
+
+__device__ __forceinline__ float block_reduce_max(float value) {
+  __shared__ float shared[kReduceThreads];
+  shared[threadIdx.x] = value;
+  __syncthreads();
+  for (int offset = kReduceThreads / 2; offset > 0; offset >>= 1) {
+    if (threadIdx.x < offset) {
+      shared[threadIdx.x] =
+          fmaxf(shared[threadIdx.x], shared[threadIdx.x + offset]);
+    }
+    __syncthreads();
+  }
+  return shared[0];
+}
+
+__global__ void reduce_abs_max_bf16(
+    const __nv_bfloat16* source, float* block_maxima,
+    size_t element_count) {
+  const size_t block_start =
+      static_cast<size_t>(blockIdx.x) * kReduceThreads *
+      kReduceItemsPerThread;
+  float local_max = 0.0f;
+  for (int item = 0; item < kReduceItemsPerThread; ++item) {
+    const size_t index =
+        block_start + static_cast<size_t>(threadIdx.x) +
+        static_cast<size_t>(item) * kReduceThreads;
+    if (index < element_count) {
+      local_max =
+          fmaxf(local_max, fabsf(__bfloat162float(source[index])));
+    }
+  }
+  const float block_max = block_reduce_max(local_max);
+  if (threadIdx.x == 0) {
+    block_maxima[blockIdx.x] = block_max;
+  }
+}
+
+__global__ void reduce_max_float(
+    const float* source, float* block_maxima, size_t element_count) {
+  const size_t block_start =
+      static_cast<size_t>(blockIdx.x) * kReduceThreads *
+      kReduceItemsPerThread;
+  float local_max = 0.0f;
+  for (int item = 0; item < kReduceItemsPerThread; ++item) {
+    const size_t index =
+        block_start + static_cast<size_t>(threadIdx.x) +
+        static_cast<size_t>(item) * kReduceThreads;
+    if (index < element_count) {
+      local_max = fmaxf(local_max, source[index]);
+    }
+  }
+  const float block_max = block_reduce_max(local_max);
+  if (threadIdx.x == 0) {
+    block_maxima[blockIdx.x] = block_max;
+  }
+}
+
+__global__ void finalize_activation_scale(
+    const float* activation_amax, const float* weight_tensor_scale,
+    float* activation_tensor_scale, float* fp4_alpha) {
+  if (blockIdx.x == 0 && threadIdx.x == 0) {
+    const float activation_scale =
+        tensor_scale_from_amax(activation_amax[0]);
+    activation_tensor_scale[0] = activation_scale;
+    fp4_alpha[0] = weight_tensor_scale[0] * activation_scale;
+  }
+}
+
+// logical_m may be smaller than padded_m. Padded rows are written as exact
+// FP4 zero with zero block scales, so reusable scratch never exposes a prior
+// call's tail.
+__global__ void dynamic_quantize_activation(
+    const __nv_bfloat16* source, uint8_t* destination_fp4,
+    uint8_t* destination_scales,
+    const float* activation_tensor_scale, int rows_k, int logical_m,
+    int scale_inner_dim, uint64_t padded_scale_blocks) {
+  const int lane = threadIdx.x & 31;
+  const int warp_in_block = threadIdx.x >> 5;
+  const uint64_t logical_block =
+      static_cast<uint64_t>(blockIdx.x) * kWarpsPerQuantBlock +
+      static_cast<uint64_t>(warp_in_block);
+  if (logical_block >= padded_scale_blocks) {
+    return;
+  }
+
+  const int blocks_per_column = rows_k / kFp4BlockElements;
+  const int outer =
+      static_cast<int>(logical_block / blocks_per_column);
+  const int inner_scale = static_cast<int>(
+      logical_block -
+      static_cast<uint64_t>(outer) * blocks_per_column);
+  const int first_row = inner_scale * kFp4BlockElements;
+  const size_t column_base = static_cast<size_t>(outer) * rows_k;
+  const bool valid_outer = outer < logical_m;
+  const float tensor_scale = activation_tensor_scale[0];
+  const float inverse_tensor_scale =
+      tensor_scale == 0.0f ? 0.0f : 1.0f / tensor_scale;
+
+  float value0 = 0.0f;
+  float value1 = 0.0f;
+  float magnitude = 0.0f;
+  size_t destination_source_index = 0;
+  if (lane < kFp4BlockElements / 2) {
+    const int row0 = first_row + lane * 2;
+    destination_source_index = column_base + row0;
+    if (valid_outer) {
+      value0 = __bfloat162float(source[destination_source_index]) *
+               inverse_tensor_scale;
+      value1 = __bfloat162float(source[destination_source_index + 1]) *
+               inverse_tensor_scale;
+      magnitude = fmaxf(fabsf(value0), fabsf(value1));
+    }
+  }
+  for (int offset = 16; offset > 0; offset >>= 1) {
+    magnitude =
+        fmaxf(magnitude,
+              __shfl_down_sync(0xffffffffU, magnitude, offset));
+  }
+
+  float rounded_scale = 0.0f;
+  if (lane == 0) {
+    const uint8_t scale_raw = __nv_cvt_float_to_fp8(
+        magnitude / kFp4E2M1Max, __NV_SATFINITE, __NV_E4M3);
+    rounded_scale = device_ue4m3_to_float(scale_raw);
+    destination_scales[device_scale_offset(
+        outer, inner_scale, scale_inner_dim)] = scale_raw;
+  }
+  rounded_scale = __shfl_sync(0xffffffffU, rounded_scale, 0);
+
+  if (lane < kFp4BlockElements / 2) {
+    const float inverse_scale =
+        rounded_scale == 0.0f ? 0.0f : 1.0f / rounded_scale;
+    destination_fp4[destination_source_index / 2] =
+        __nv_cvt_float2_to_fp4x2(
+            make_float2(value0 * inverse_scale,
+                        value1 * inverse_scale),
+            __NV_E2M1, cudaRoundNearest);
+  }
+}
+
+__global__ void add_bias_bf16(
+    __nv_bfloat16* output, const __nv_bfloat16* bias,
+    size_t element_count, int n) {
+  const size_t index =
+      static_cast<size_t>(blockIdx.x) * blockDim.x + threadIdx.x;
+  if (index < element_count) {
+    const float value = __bfloat162float(output[index]);
+    const float bias_value = __bfloat162float(bias[index % n]);
+    output[index] = __float2bfloat16(value + bias_value);
+  }
+}
+
+size_t reduction_block_count(size_t element_count) {
+  return round_up_divide(
+      element_count,
+      static_cast<size_t>(kReduceThreads * kReduceItemsPerThread));
+}
+
+size_t reduction_scratch_elements(size_t element_count) {
+  size_t max_blocks = 1;
+  while (element_count > 1) {
+    const size_t blocks = reduction_block_count(element_count);
+    max_blocks = std::max(max_blocks, blocks);
+    element_count = blocks;
+  }
+  return max_blocks;
+}
+
+void enqueue_activation_amax_reduce(
+    const __nv_bfloat16* source, size_t element_count,
+    float* scratch_a, float* scratch_b, float* destination_amax,
+    int max_grid_x, cudaStream_t stream) {
+  const float* current_source = nullptr;
+  float* current_destination = scratch_a;
+  size_t current_count = element_count;
+  bool first_stage = true;
+  while (true) {
+    const size_t blocks = reduction_block_count(current_count);
+    if (blocks == 0 || blocks > static_cast<size_t>(max_grid_x)) {
+      fail("activation amax reduction exceeds the GPU grid limit");
+    }
+    if (first_stage) {
+      reduce_abs_max_bf16<<<
+          static_cast<unsigned int>(blocks), kReduceThreads, 0,
+          stream>>>(source, current_destination, current_count);
+    } else {
+      reduce_max_float<<<
+          static_cast<unsigned int>(blocks), kReduceThreads, 0,
+          stream>>>(current_source, current_destination, current_count);
+    }
+    CUDA_CHECK(cudaPeekAtLastError());
+    if (blocks == 1) {
+      if (current_destination != destination_amax) {
+        CUDA_CHECK(cudaMemcpyAsync(
+            destination_amax, current_destination, sizeof(float),
+            cudaMemcpyDeviceToDevice, stream));
+      }
+      return;
+    }
+    current_count = blocks;
+    current_source = current_destination;
+    current_destination =
+        current_destination == scratch_a ? scratch_b : scratch_a;
+    first_stage = false;
+  }
+}
+
+struct DeviceAllocation {
+  void* pointer = nullptr;
+  size_t bytes = 0;
+
+  ~DeviceAllocation() {
+    if (pointer != nullptr) {
+      cudaFree(pointer);
+    }
+  }
+
+  DeviceAllocation() = default;
+  DeviceAllocation(const DeviceAllocation&) = delete;
+  DeviceAllocation& operator=(const DeviceAllocation&) = delete;
+};
+
+struct Context {
+  int device = -1;
+  int max_grid_x = 0;
+  cublasLtHandle_t handle = nullptr;
+  DeviceAllocation x_fp4;
+  DeviceAllocation x_scales;
+  DeviceAllocation reduce_a;
+  DeviceAllocation reduce_b;
+  DeviceAllocation activation_amax;
+  DeviceAllocation activation_scale;
+  DeviceAllocation fp4_alpha;
+  DeviceAllocation fp4_beta;
+  DeviceAllocation workspace;
+  uintptr_t bound_stream = 0;
+  bool stream_bound = false;
+  std::mutex mutex;
+
+  ~Context() {
+    if (handle != nullptr) {
+      cublasLtDestroy(handle);
+    }
+  }
+};
+
+void allocate_exact(DeviceAllocation* allocation, size_t bytes) {
+  if (allocation->pointer != nullptr) {
+    CUDA_CHECK(cudaFree(allocation->pointer));
+    allocation->pointer = nullptr;
+    allocation->bytes = 0;
+  }
+  if (bytes != 0) {
+    CUDA_CHECK(cudaMalloc(&allocation->pointer, bytes));
+    allocation->bytes = bytes;
+  }
+}
+
+void ensure_capacity(
+    Context* context, DeviceAllocation* allocation, size_t bytes,
+    cudaStream_t stream) {
+  if (allocation->bytes >= bytes) {
+    return;
+  }
+  // A growth invalidates scratch pointers. Synchronize the one bound stream
+  // before freeing; steady-state forwards do not synchronize.
+  if (context->stream_bound) {
+    CUDA_CHECK(cudaStreamSynchronize(stream));
+  }
+  allocate_exact(allocation, bytes);
+}
+
+void set_scale_mode(
+    cublasLtMatmulDesc_t operation,
+    cublasLtMatmulDescAttributes_t attribute) {
+  const int32_t mode =
+      CUBLASLT_MATMUL_MATRIX_SCALE_VEC16_UE4M3;
+  CUBLASLT_CHECK(cublasLtMatmulDescSetAttribute(
+      operation, attribute, &mode, sizeof(mode)));
+}
+
+void set_pointer_attribute(
+    cublasLtMatmulDesc_t operation,
+    cublasLtMatmulDescAttributes_t attribute, const void* pointer) {
+  CUBLASLT_CHECK(cublasLtMatmulDescSetAttribute(
+      operation, attribute, &pointer, sizeof(pointer)));
+}
+
+struct Descriptors {
+  cublasLtMatmulDesc_t operation = nullptr;
+  cublasLtMatrixLayout_t w = nullptr;
+  cublasLtMatrixLayout_t x = nullptr;
+  cublasLtMatrixLayout_t c = nullptr;
+  cublasLtMatrixLayout_t d = nullptr;
+
+  ~Descriptors() {
+    if (d != nullptr) cublasLtMatrixLayoutDestroy(d);
+    if (c != nullptr) cublasLtMatrixLayoutDestroy(c);
+    if (x != nullptr) cublasLtMatrixLayoutDestroy(x);
+    if (w != nullptr) cublasLtMatrixLayoutDestroy(w);
+    if (operation != nullptr) cublasLtMatmulDescDestroy(operation);
+  }
+};
+
+cublasLtMatmulAlgo_t choose_algo(
+    Context* context, const Descriptors& descriptors,
+    size_t* workspace_bytes) {
+  cublasLtMatmulPreference_t preference = nullptr;
+  CUBLASLT_CHECK(cublasLtMatmulPreferenceCreate(&preference));
+  CUBLASLT_CHECK(cublasLtMatmulPreferenceSetAttribute(
+      preference, CUBLASLT_MATMUL_PREF_MAX_WORKSPACE_BYTES,
+      &kWorkspaceBytes, sizeof(kWorkspaceBytes)));
+
+  cublasLtMatmulHeuristicResult_t candidates[16]{};
+  int returned = 0;
+  const cublasStatus_t status = cublasLtMatmulAlgoGetHeuristic(
+      context->handle, descriptors.operation, descriptors.w,
+      descriptors.x, descriptors.c, descriptors.d, preference, 16,
+      candidates, &returned);
+  cublasLtMatmulPreferenceDestroy(preference);
+  if (status != CUBLAS_STATUS_SUCCESS || returned == 0) {
+    fail("cuBLASLt returned no NVFP4 heuristic");
+  }
+  for (int index = 0; index < returned; ++index) {
+    if (candidates[index].state == CUBLAS_STATUS_SUCCESS) {
+      *workspace_bytes = candidates[index].workspaceSize;
+      return candidates[index].algo;
+    }
+  }
+  fail("all cuBLASLt NVFP4 heuristics were unsupported");
+}
+
+void validate_device_pointer(
+    const void* pointer, int expected_device, const char* label) {
+  if (pointer == nullptr) {
+    fail(std::string(label) + " is null");
+  }
+  cudaPointerAttributes attributes{};
+  CUDA_CHECK(cudaPointerGetAttributes(&attributes, pointer));
+  if (attributes.type != cudaMemoryTypeDevice ||
+      attributes.device != expected_device) {
+    fail(std::string(label) + " is not a CUDA allocation on the context device");
+  }
+  if ((reinterpret_cast<uintptr_t>(pointer) & 0x0fU) != 0) {
+    fail(std::string(label) + " is not at least 16-byte aligned");
+  }
+}
+
+void forward_impl(
+    Context* context, const void* input_bf16,
+    const void* packed_weight, size_t packed_weight_bytes,
+    const void* packed_weight_scales,
+    size_t packed_weight_scale_bytes,
+    const void* weight_tensor_scale_f32, const void* bias_bf16,
+    void* output_bf16, int logical_m, int k, int n,
+    uintptr_t stream_value) {
+  if (logical_m <= 0 || k <= 0 || n <= 0) {
+    fail("M, K, and N must be positive");
+  }
+  const size_t expected_weight_bytes =
+      packed_weight_bytes_checked(n, k);
+  const size_t expected_scale_bytes = make_scale_layout(k, n).bytes;
+  if (packed_weight_bytes != expected_weight_bytes ||
+      packed_weight_scale_bytes != expected_scale_bytes) {
+    fail("packed weight or scale buffer has the wrong byte size");
+  }
+
+  CUDA_CHECK(cudaSetDevice(context->device));
+  cudaStream_t stream = reinterpret_cast<cudaStream_t>(stream_value);
+  if (!context->stream_bound) {
+    context->bound_stream = stream_value;
+    context->stream_bound = true;
+  } else if (context->bound_stream != stream_value) {
+    fail("resident NVFP4 context is bound to a different CUDA stream");
+  }
+  cudaStreamCaptureStatus capture_status = cudaStreamCaptureStatusNone;
+  CUDA_CHECK(cudaStreamIsCapturing(stream, &capture_status));
+  if (capture_status != cudaStreamCaptureStatusNone) {
+    fail("resident ctypes prototype does not support CUDA graph capture");
+  }
+
+  validate_device_pointer(input_bf16, context->device, "input");
+  validate_device_pointer(packed_weight, context->device, "packed weight");
+  validate_device_pointer(
+      packed_weight_scales, context->device, "weight scales");
+  validate_device_pointer(
+      weight_tensor_scale_f32, context->device, "weight tensor scale");
+  validate_device_pointer(output_bf16, context->device, "output");
+  if (bias_bf16 != nullptr) {
+    validate_device_pointer(bias_bf16, context->device, "bias");
+  }
+
+  const int padded_m = round_up(logical_m, 8);
+  const size_t input_elements = checked_multiply(
+      static_cast<size_t>(logical_m), static_cast<size_t>(k), "input");
+  const size_t padded_input_elements = checked_multiply(
+      static_cast<size_t>(padded_m), static_cast<size_t>(k),
+      "padded input");
+  const size_t output_elements = checked_multiply(
+      static_cast<size_t>(padded_m), static_cast<size_t>(n), "output");
+  const ScaleLayout x_scale_layout = make_scale_layout(k, padded_m);
+  const size_t reduce_elements =
+      reduction_scratch_elements(input_elements);
+
+  ensure_capacity(
+      context, &context->x_fp4, padded_input_elements / 2, stream);
+  ensure_capacity(
+      context, &context->x_scales, x_scale_layout.bytes, stream);
+  ensure_capacity(
+      context, &context->reduce_a,
+      checked_multiply(reduce_elements, sizeof(float), "reduction"),
+      stream);
+  ensure_capacity(
+      context, &context->reduce_b,
+      checked_multiply(reduce_elements, sizeof(float), "reduction"),
+      stream);
+
+  // Clear the complete tiled scale allocation, including padding that the
+  // logical quantizer does not address.
+  CUDA_CHECK(cudaMemsetAsync(
+      context->x_scales.pointer, 0, x_scale_layout.bytes, stream));
+  enqueue_activation_amax_reduce(
+      static_cast<const __nv_bfloat16*>(input_bf16), input_elements,
+      static_cast<float*>(context->reduce_a.pointer),
+      static_cast<float*>(context->reduce_b.pointer),
+      static_cast<float*>(context->activation_amax.pointer),
+      context->max_grid_x, stream);
+  finalize_activation_scale<<<1, 1, 0, stream>>>(
+      static_cast<const float*>(context->activation_amax.pointer),
+      static_cast<const float*>(weight_tensor_scale_f32),
+      static_cast<float*>(context->activation_scale.pointer),
+      static_cast<float*>(context->fp4_alpha.pointer));
+  CUDA_CHECK(cudaPeekAtLastError());
+
+  const uint64_t padded_blocks =
+      static_cast<uint64_t>(padded_m) *
+      static_cast<uint64_t>(k / kFp4BlockElements);
+  const uint64_t cuda_blocks =
+      (padded_blocks + kWarpsPerQuantBlock - 1) /
+      kWarpsPerQuantBlock;
+  if (cuda_blocks == 0 ||
+      cuda_blocks > static_cast<uint64_t>(context->max_grid_x)) {
+    fail("activation quantizer exceeds the GPU grid limit");
+  }
+  dynamic_quantize_activation<<<
+      static_cast<unsigned int>(cuda_blocks), kQuantThreads, 0,
+      stream>>>(
+      static_cast<const __nv_bfloat16*>(input_bf16),
+      static_cast<uint8_t*>(context->x_fp4.pointer),
+      static_cast<uint8_t*>(context->x_scales.pointer),
+      static_cast<const float*>(context->activation_scale.pointer),
+      k, logical_m, x_scale_layout.inner_dim, padded_blocks);
+  CUDA_CHECK(cudaPeekAtLastError());
+
+  Descriptors descriptors;
+  CUBLASLT_CHECK(cublasLtMatmulDescCreate(
+      &descriptors.operation, CUBLAS_COMPUTE_32F, CUDA_R_32F));
+  const cublasOperation_t transpose_a = CUBLAS_OP_T;
+  const cublasOperation_t transpose_b = CUBLAS_OP_N;
+  CUBLASLT_CHECK(cublasLtMatmulDescSetAttribute(
+      descriptors.operation, CUBLASLT_MATMUL_DESC_TRANSA,
+      &transpose_a, sizeof(transpose_a)));
+  CUBLASLT_CHECK(cublasLtMatmulDescSetAttribute(
+      descriptors.operation, CUBLASLT_MATMUL_DESC_TRANSB,
+      &transpose_b, sizeof(transpose_b)));
+  const cublasLtPointerMode_t pointer_mode =
+      CUBLASLT_POINTER_MODE_DEVICE;
+  CUBLASLT_CHECK(cublasLtMatmulDescSetAttribute(
+      descriptors.operation, CUBLASLT_MATMUL_DESC_POINTER_MODE,
+      &pointer_mode, sizeof(pointer_mode)));
+  set_scale_mode(
+      descriptors.operation, CUBLASLT_MATMUL_DESC_A_SCALE_MODE);
+  set_scale_mode(
+      descriptors.operation, CUBLASLT_MATMUL_DESC_B_SCALE_MODE);
+  // Refresh pointer attributes for every call. Module buffers may differ
+  // between adjacent linears even though the shape is identical.
+  set_pointer_attribute(
+      descriptors.operation, CUBLASLT_MATMUL_DESC_A_SCALE_POINTER,
+      packed_weight_scales);
+  set_pointer_attribute(
+      descriptors.operation, CUBLASLT_MATMUL_DESC_B_SCALE_POINTER,
+      context->x_scales.pointer);
+
+  CUBLASLT_CHECK(cublasLtMatrixLayoutCreate(
+      &descriptors.w, CUDA_R_4F_E2M1, k, n, k));
+  CUBLASLT_CHECK(cublasLtMatrixLayoutCreate(
+      &descriptors.x, CUDA_R_4F_E2M1, k, padded_m, k));
+  CUBLASLT_CHECK(cublasLtMatrixLayoutCreate(
+      &descriptors.c, CUDA_R_16BF, n, padded_m, n));
+  CUBLASLT_CHECK(cublasLtMatrixLayoutCreate(
+      &descriptors.d, CUDA_R_16BF, n, padded_m, n));
+
+  size_t selected_workspace_bytes = 0;
+  const cublasLtMatmulAlgo_t algorithm =
+      choose_algo(context, descriptors, &selected_workspace_bytes);
+  if (selected_workspace_bytes > context->workspace.bytes) {
+    fail("selected cuBLASLt algorithm exceeds the context workspace");
+  }
+  CUBLASLT_CHECK(cublasLtMatmul(
+      context->handle, descriptors.operation,
+      context->fp4_alpha.pointer,
+      packed_weight, descriptors.w,
+      context->x_fp4.pointer, descriptors.x,
+      context->fp4_beta.pointer,
+      output_bf16, descriptors.c,
+      output_bf16, descriptors.d,
+      &algorithm, context->workspace.pointer,
+      selected_workspace_bytes, stream));
+
+  if (bias_bf16 != nullptr) {
+    const size_t blocks =
+        round_up_divide(output_elements,
+                        static_cast<size_t>(kBiasThreads));
+    if (blocks > static_cast<size_t>(context->max_grid_x)) {
+      fail("bias kernel exceeds the GPU grid limit");
+    }
+    add_bias_bf16<<<
+        static_cast<unsigned int>(blocks), kBiasThreads, 0, stream>>>(
+        static_cast<__nv_bfloat16*>(output_bf16),
+        static_cast<const __nv_bfloat16*>(bias_bf16),
+        output_elements, n);
+    CUDA_CHECK(cudaPeekAtLastError());
+  }
+}
+
+template <typename Function>
+int guarded(Function&& function) {
+  try {
+    g_last_error.clear();
+    function();
+    return 0;
+  } catch (const std::exception& error) {
+    g_last_error = error.what();
+    return 1;
+  } catch (...) {
+    g_last_error = "unknown native NVFP4 error";
+    return 2;
+  }
+}
+
+}  // namespace
+
+extern "C" int mage_nvfp4_abi_version(void) {
+  return kAbiVersion;
+}
+
+extern "C" const char* mage_nvfp4_last_error(void) {
+  return g_last_error.c_str();
+}
+
+extern "C" size_t mage_nvfp4_packed_weight_bytes(int n, int k) {
+  try {
+    g_last_error.clear();
+    return packed_weight_bytes_checked(n, k);
+  } catch (const std::exception& error) {
+    g_last_error = error.what();
+    return 0;
+  }
+}
+
+extern "C" size_t mage_nvfp4_weight_scale_bytes(int n, int k) {
+  try {
+    g_last_error.clear();
+    packed_weight_bytes_checked(n, k);
+    return make_scale_layout(k, n).bytes;
+  } catch (const std::exception& error) {
+    g_last_error = error.what();
+    return 0;
+  }
+}
+
+extern "C" int mage_nvfp4_pack_weight_bf16(
+    const void* weight_bf16, int n, int k,
+    void* packed_weight, size_t packed_weight_capacity,
+    void* packed_scales, size_t packed_scale_capacity,
+    float* tensor_scale) {
+  return guarded([&]() {
+    if (weight_bf16 == nullptr || packed_weight == nullptr ||
+        packed_scales == nullptr || tensor_scale == nullptr) {
+      fail("weight packer received a null pointer");
+    }
+    const size_t required_weight =
+        packed_weight_bytes_checked(n, k);
+    const ScaleLayout scale_layout = make_scale_layout(k, n);
+    if (packed_weight_capacity != required_weight ||
+        packed_scale_capacity != scale_layout.bytes) {
+      fail("weight packer received an incorrectly sized destination");
+    }
+
+    const auto* source =
+        static_cast<const __nv_bfloat16*>(weight_bf16);
+    auto* destination = static_cast<uint8_t*>(packed_weight);
+    auto* scales = static_cast<uint8_t*>(packed_scales);
+    std::memset(destination, 0, required_weight);
+    std::memset(scales, 0, scale_layout.bytes);
+    const size_t elements =
+        checked_multiply(static_cast<size_t>(n),
+                         static_cast<size_t>(k), "weight");
+    float global_amax = 0.0f;
+    for (size_t index = 0; index < elements; ++index) {
+      global_amax = std::max(
+          global_amax,
+          std::abs(__bfloat162float(source[index])));
+    }
+    *tensor_scale =
+        global_amax == 0.0f
+            ? 1.0f
+            : global_amax / kTensorScaleDenominator;
+    const float inverse_tensor_scale = 1.0f / *tensor_scale;
+
+    // nn.Linear weight [N,K] is physical column-major KxN.
+    for (int column = 0; column < n; ++column) {
+      for (int block = 0; block < k / kFp4BlockElements; ++block) {
+        const int first_row = block * kFp4BlockElements;
+        float block_amax = 0.0f;
+        for (int lane = 0; lane < kFp4BlockElements; ++lane) {
+          const size_t index =
+              static_cast<size_t>(first_row + lane) +
+              static_cast<size_t>(column) * k;
+          block_amax = std::max(
+              block_amax,
+              std::abs(__bfloat162float(source[index]) *
+                       inverse_tensor_scale));
+        }
+        const uint8_t scale_raw = __nv_cvt_float_to_fp8(
+            block_amax / kFp4E2M1Max,
+            __NV_SATFINITE, __NV_E4M3);
+        const float rounded_scale = host_ue4m3_to_float(scale_raw);
+        scales[host_scale_offset(
+            column, block, scale_layout.inner_dim)] = scale_raw;
+        const float inverse_scale =
+            rounded_scale == 0.0f ? 0.0f : 1.0f / rounded_scale;
+        for (int pair = 0;
+             pair < kFp4BlockElements / 2; ++pair) {
+          const int row0 = first_row + pair * 2;
+          const size_t source_index =
+              static_cast<size_t>(row0) +
+              static_cast<size_t>(column) * k;
+          const float value0 =
+              __bfloat162float(source[source_index]) *
+              inverse_tensor_scale;
+          const float value1 =
+              __bfloat162float(source[source_index + 1]) *
+              inverse_tensor_scale;
+          destination[source_index / 2] =
+              __nv_cvt_float2_to_fp4x2(
+                  make_float2(value0 * inverse_scale,
+                              value1 * inverse_scale),
+                  __NV_E2M1, cudaRoundNearest);
+        }
+      }
+    }
+  });
+}
+
+extern "C" int mage_nvfp4_create_context(
+    int cuda_device, void** context) {
+  return guarded([&]() {
+    if (context == nullptr) {
+      fail("context output pointer is null");
+    }
+    *context = nullptr;
+    int device_count = 0;
+    CUDA_CHECK(cudaGetDeviceCount(&device_count));
+    if (cuda_device < 0 || cuda_device >= device_count) {
+      fail("invalid CUDA device index");
+    }
+    CUDA_CHECK(cudaSetDevice(cuda_device));
+    cudaDeviceProp properties{};
+    CUDA_CHECK(cudaGetDeviceProperties(&properties, cuda_device));
+    if (properties.major != 12 || properties.minor != 0) {
+      fail("resident NVFP4 prototype requires sm_120");
+    }
+
+    Context* created = new Context();
+    try {
+      created->device = cuda_device;
+      created->max_grid_x = properties.maxGridSize[0];
+      CUBLASLT_CHECK(cublasLtCreate(&created->handle));
+      allocate_exact(&created->activation_amax, sizeof(float));
+      allocate_exact(&created->activation_scale, sizeof(float));
+      allocate_exact(&created->fp4_alpha, sizeof(float));
+      allocate_exact(&created->fp4_beta, sizeof(float));
+      allocate_exact(&created->workspace, kWorkspaceBytes);
+      const float zero = 0.0f;
+      CUDA_CHECK(cudaMemcpy(
+          created->fp4_beta.pointer, &zero, sizeof(zero),
+          cudaMemcpyHostToDevice));
+      *context = created;
+    } catch (...) {
+      delete created;
+      throw;
+    }
+  });
+}
+
+extern "C" int mage_nvfp4_destroy_context(void* context) {
+  return guarded([&]() {
+    if (context == nullptr) {
+      return;
+    }
+    auto* typed = static_cast<Context*>(context);
+    {
+      std::lock_guard<std::mutex> lock(typed->mutex);
+      CUDA_CHECK(cudaSetDevice(typed->device));
+      if (typed->stream_bound) {
+        CUDA_CHECK(cudaStreamSynchronize(
+            reinterpret_cast<cudaStream_t>(typed->bound_stream)));
+      }
+    }
+    delete typed;
+  });
+}
+
+extern "C" size_t mage_nvfp4_context_reserved_bytes(
+    const void* context) {
+  if (context == nullptr) {
+    return 0;
+  }
+  const auto* typed = static_cast<const Context*>(context);
+  return typed->x_fp4.bytes + typed->x_scales.bytes +
+         typed->reduce_a.bytes + typed->reduce_b.bytes +
+         typed->activation_amax.bytes +
+         typed->activation_scale.bytes +
+         typed->fp4_alpha.bytes + typed->fp4_beta.bytes +
+         typed->workspace.bytes;
+}
+
+extern "C" int mage_nvfp4_linear_forward(
+    void* context, const void* input_bf16,
+    const void* packed_weight, size_t packed_weight_bytes,
+    const void* packed_weight_scales,
+    size_t packed_weight_scale_bytes,
+    const void* weight_tensor_scale_f32,
+    const void* bias_bf16, void* output_bf16,
+    int logical_m, int k, int n, uintptr_t stream) {
+  return guarded([&]() {
+    if (context == nullptr) {
+      fail("context is null");
+    }
+    auto* typed = static_cast<Context*>(context);
+    std::lock_guard<std::mutex> lock(typed->mutex);
+    forward_impl(
+        typed, input_bf16, packed_weight, packed_weight_bytes,
+        packed_weight_scales, packed_weight_scale_bytes,
+        weight_tensor_scale_f32, bias_bf16, output_bf16,
+        logical_m, k, n, stream);
+  });
+}

runtime/nvfp4_linear.h

@@ -0,0 +1,66 @@
+#pragma once
+
+#include <stddef.h>
+#include <stdint.h>
+
+#if defined(_WIN32)
+#define MAGE_NVFP4_API __declspec(dllexport)
+#else
+#define MAGE_NVFP4_API __attribute__((visibility("default")))
+#endif
+
+#ifdef __cplusplus
+extern "C" {
+#endif
+
+// The ABI is inference-only. All functions return zero on success and a
+// non-zero value on failure unless their return type documents otherwise.
+MAGE_NVFP4_API int mage_nvfp4_abi_version(void);
+MAGE_NVFP4_API const char* mage_nvfp4_last_error(void);
+
+MAGE_NVFP4_API size_t mage_nvfp4_packed_weight_bytes(int n, int k);
+MAGE_NVFP4_API size_t mage_nvfp4_weight_scale_bytes(int n, int k);
+
+// Pack a contiguous CPU BF16 nn.Linear weight in logical row-major [N, K]
+// order into cuBLASLt E2M1 payload and tiled VEC16 UE4M3 block scales.
+MAGE_NVFP4_API int mage_nvfp4_pack_weight_bf16(
+    const void* weight_bf16, int n, int k,
+    void* packed_weight, size_t packed_weight_capacity,
+    void* packed_scales, size_t packed_scale_capacity,
+    float* tensor_scale);
+
+// A context owns one cuBLASLt handle and reusable activation/reduction/
+// workspace buffers. The first forward binds it to one CUDA stream. This
+// initial prototype deliberately rejects a different stream and is not
+// re-entrant.
+MAGE_NVFP4_API int mage_nvfp4_create_context(
+    int cuda_device, void** context);
+MAGE_NVFP4_API int mage_nvfp4_destroy_context(void* context);
+MAGE_NVFP4_API size_t mage_nvfp4_context_reserved_bytes(
+    const void* context);
+
+// Run Y[M,N] = X[M,K] * W[N,K]^T + bias.
+//
+// X, packed weight/scales, tensor scale, bias, and output are CUDA pointers on
+// context's device. X is contiguous BF16 [logical_m, K]. Output is contiguous
+// BF16 [round_up(logical_m, 8), N]. The caller slices the padded rows.
+// bias_bf16 may be null. stream is a cudaStream_t passed as uintptr_t.
+MAGE_NVFP4_API int mage_nvfp4_linear_forward(
+    void* context,
+    const void* input_bf16,
+    const void* packed_weight,
+    size_t packed_weight_bytes,
+    const void* packed_weight_scales,
+    size_t packed_weight_scale_bytes,
+    const void* weight_tensor_scale_f32,
+    const void* bias_bf16,
+    void* output_bf16,
+    int logical_m,
+    int k,
+    int n,
+    uintptr_t stream);
+
+#ifdef __cplusplus
+}
+#endif
+

runtime/packed_artifact.py

@@ -0,0 +1,1068 @@
+#!/usr/bin/env python3
+
+from __future__ import annotations
+
+import argparse
+import hashlib
+import json
+import math
+import os
+import sys
+from collections import OrderedDict
+from dataclasses import asdict, dataclass
+from datetime import datetime, timezone
+from pathlib import Path
+from typing import Callable, Iterable
+
+import safetensors
+import torch
+import torch.nn as nn
+from safetensors import safe_open
+from safetensors.torch import save_file
+
+
+FORMAT_VERSION = 1
+ARTIFACT_KIND = "mage_flow_transformer_mlp_nvfp4_resident_v1"
+CANONICAL_SAFETENSORS_METADATA = {
+    "mage_nvfp4_contract": (
+        f"{ARTIFACT_KIND};format_version={FORMAT_VERSION}"
+    )
+}
+LEGACY_SAFETENSORS_METADATA = {
+    "artifact_kind": ARTIFACT_KIND,
+    "format_version": str(FORMAT_VERSION),
+}
+FP4_BLOCK_ELEMENTS = 16
+SCALE_TILE_OUTER = 128
+SCALE_TILE_INNER = 4
+FP4_E2M1_MAX = 6.0
+FP4_TENSOR_SCALE_MAX = 448.0
+NVFP4_TENSOR_SCALE_DENOMINATOR = FP4_E2M1_MAX * FP4_TENSOR_SCALE_MAX
+TARGET_DEPTH = 12
+
+RELEASE_ROOT = Path(__file__).resolve().parents[1]
+PROJECT_ROOT = RELEASE_ROOT
+MAGE_ROOT = RELEASE_ROOT / "vendor"
+RESIDENT_PYTHON_ROOT = RELEASE_ROOT / "runtime"
+RESIDENT_SOURCE = RESIDENT_PYTHON_ROOT / "nvfp4_linear.cu"
+RESIDENT_LIBRARY = RESIDENT_PYTHON_ROOT / "libmage_nvfp4_linear.so"
+ARTIFACT_SCRIPT_PATH = Path(__file__).resolve()
+_NATIVE_PACKER = None
+
+
+class PackedArtifactError(RuntimeError):
+    pass
+
+
+@dataclass(frozen=True)
+class ScaleLayout:
+    inner_dim: int
+    outer_tiles: int
+    bytes: int
+
+
+@dataclass(frozen=True)
+class TargetSpec:
+    module_key: str
+    weight_key: str
+    bias_key: str
+    artifact_weight_key: str
+    artifact_scale_key: str
+    artifact_tensor_scale_key: str
+    artifact_bias_key: str
+
+
+def fail(message: str) -> None:
+    raise PackedArtifactError(message)
+
+
+def round_up(value: int, multiple: int) -> int:
+    return ((value + multiple - 1) // multiple) * multiple
+
+
+def sha256_file(path: Path) -> str:
+    digest = hashlib.sha256()
+    with path.open("rb") as handle:
+        while True:
+            chunk = handle.read(1 << 20)
+            if not chunk:
+                break
+            digest.update(chunk)
+    return digest.hexdigest()
+
+
+def sha256_bytes(data: bytes) -> str:
+    return hashlib.sha256(data).hexdigest()
+
+
+def tensor_bytes(tensor: torch.Tensor) -> bytes:
+    if not tensor.is_contiguous():
+        tensor = tensor.contiguous()
+    return tensor.view(torch.uint8).cpu().numpy().tobytes()
+
+
+def sha256_tensor(tensor: torch.Tensor) -> str:
+    return sha256_bytes(tensor_bytes(tensor))
+
+
+def fsync_directory(path: Path) -> None:
+    descriptor = os.open(path, os.O_RDONLY | getattr(os, "O_DIRECTORY", 0))
+    try:
+        os.fsync(descriptor)
+    finally:
+        os.close(descriptor)
+
+
+def fsync_file(path: Path) -> None:
+    descriptor = os.open(path, os.O_RDONLY)
+    try:
+        os.fsync(descriptor)
+    finally:
+        os.close(descriptor)
+
+
+def write_bytes_once(path: Path, payload: bytes) -> None:
+    with path.open("xb") as handle:
+        handle.write(payload)
+        handle.flush()
+        os.fsync(handle.fileno())
+    fsync_directory(path.parent)
+
+
+def host_scale_offset(outer: int, inner_scale: int, scale_inner_dim: int) -> int:
+    outer_tile = outer // SCALE_TILE_OUTER
+    local_outer = outer % SCALE_TILE_OUTER
+    local_inner = inner_scale % SCALE_TILE_INNER
+    inner_tile_start = inner_scale - local_inner
+    tile_base = (inner_tile_start + outer_tile * scale_inner_dim) * SCALE_TILE_OUTER
+    return tile_base + (local_outer % 32) * 16 + (local_outer // 32) * 4 + local_inner
+
+
+def make_scale_layout(rows_k: int, outer_columns: int) -> ScaleLayout:
+    if rows_k <= 0 or outer_columns <= 0:
+        fail("scale layout requires positive rows_k and outer_columns")
+    if rows_k % FP4_BLOCK_ELEMENTS:
+        fail(
+            f"scale layout requires K divisible by {FP4_BLOCK_ELEMENTS}; "
+            f"got K={rows_k}"
+        )
+    inner_dim = round_up(rows_k // FP4_BLOCK_ELEMENTS, SCALE_TILE_INNER)
+    outer_tiles = (outer_columns + SCALE_TILE_OUTER - 1) // SCALE_TILE_OUTER
+    return ScaleLayout(
+        inner_dim=inner_dim,
+        outer_tiles=outer_tiles,
+        bytes=outer_tiles * inner_dim * SCALE_TILE_OUTER,
+    )
+
+
+def build_target_specs(depth: int) -> list[TargetSpec]:
+    if depth != TARGET_DEPTH:
+        fail(
+            f"this artifact format is pinned to exactly {TARGET_DEPTH} transformer blocks; "
+            f"config reported depth={depth}"
+        )
+    specs: list[TargetSpec] = []
+    suffixes = (
+        "img_mlp.net.0.proj",
+        "img_mlp.net.2",
+        "txt_mlp.net.0.proj",
+        "txt_mlp.net.2",
+    )
+    for index in range(depth):
+        for suffix in suffixes:
+            module_key = f"transformer_blocks.{index}.{suffix}"
+            specs.append(
+                TargetSpec(
+                    module_key=module_key,
+                    weight_key=f"{module_key}.weight",
+                    bias_key=f"{module_key}.bias",
+                    artifact_weight_key=f"targets.{module_key}.packed_weight_e2m1",
+                    artifact_scale_key=f"targets.{module_key}.packed_scales_ue4m3",
+                    artifact_tensor_scale_key=f"targets.{module_key}.weight_tensor_scale",
+                    artifact_bias_key=f"targets.{module_key}.bias_bf16",
+                )
+            )
+    return specs
+
+
+def decode_fp4_e2m1(raw: int) -> float:
+    sign = -1.0 if (raw & 0x8) else 1.0
+    magnitude = raw & 0x7
+    table = (
+        0.0,
+        0.5,
+        1.0,
+        1.5,
+        2.0,
+        3.0,
+        4.0,
+        6.0,
+    )
+    return sign * table[magnitude]
+
+
+def encode_fp4_e2m1(value: float) -> int:
+    candidates = [decode_fp4_e2m1(code) for code in range(16)]
+    best_code = 0
+    best_error = math.inf
+    for code, candidate in enumerate(candidates):
+        error = abs(candidate - value)
+        if error < best_error or (error == best_error and (code & 1) == 0 and (best_code & 1) == 1):
+            best_error = error
+            best_code = code
+    return best_code
+
+
+def decode_fp8_e4m3(raw: int) -> float:
+    sign = -1.0 if (raw & 0x80) else 1.0
+    exponent = (raw >> 3) & 0x0F
+    mantissa = raw & 0x07
+    if exponent == 0:
+        if mantissa == 0:
+            return 0.0 * sign
+        return sign * (mantissa / 8.0) * (2.0 ** -6)
+    if exponent == 0x0F and mantissa == 0x07:
+        return math.nan
+    return sign * (1.0 + mantissa / 8.0) * (2.0 ** (exponent - 7))
+
+
+def _build_positive_e4m3_table() -> list[tuple[int, float]]:
+    table: list[tuple[int, float]] = []
+    for raw in range(0x80):
+        value = decode_fp8_e4m3(raw)
+        if math.isnan(value) or value < 0.0:
+            continue
+        table.append((raw, value))
+    table.sort(key=lambda item: (item[1], item[0]))
+    return table
+
+
+POSITIVE_E4M3_TABLE = _build_positive_e4m3_table()
+
+
+def encode_fp8_e4m3_satfinite(value: float) -> int:
+    if value <= 0.0:
+        return 0
+    finite_values = [item for item in POSITIVE_E4M3_TABLE if item[1] <= FP4_TENSOR_SCALE_MAX]
+    best_raw = finite_values[-1][0]
+    best_error = math.inf
+    for raw, candidate in finite_values:
+        error = abs(candidate - value)
+        if error < best_error or (error == best_error and (raw & 1) == 0 and (best_raw & 1) == 1):
+            best_error = error
+            best_raw = raw
+    return best_raw
+
+
+def host_tensor_scale_from_amax(amax: float) -> float:
+    return 1.0 if amax == 0.0 else amax / NVFP4_TENSOR_SCALE_DENOMINATOR
+
+
+def native_packer():
+    global _NATIVE_PACKER
+    if _NATIVE_PACKER is None:
+        if str(RESIDENT_PYTHON_ROOT) not in sys.path:
+            sys.path.insert(0, str(RESIDENT_PYTHON_ROOT))
+        from packed_nvfp4_linear import NativeNvfp4Library
+
+        _NATIVE_PACKER = NativeNvfp4Library(RESIDENT_LIBRARY)
+    return _NATIVE_PACKER
+
+
+def pack_weight_tensor(weight_nk_bf16: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, float]:
+    if weight_nk_bf16.dtype != torch.bfloat16 or weight_nk_bf16.device.type != "cpu":
+        fail("weight packer expects a CPU bfloat16 tensor")
+    if weight_nk_bf16.ndim != 2:
+        fail("weight packer expects a 2D [N,K] weight tensor")
+    if not weight_nk_bf16.is_contiguous():
+        weight_nk_bf16 = weight_nk_bf16.contiguous()
+
+    columns_n, rows_k = weight_nk_bf16.shape
+    if rows_k % 32 != 0 or columns_n % 8 != 0:
+        fail(
+            f"native NVFP4 packing requires K%32==0 and N%8==0; got N={columns_n} K={rows_k}"
+        )
+
+    amax = float(weight_nk_bf16.float().abs().max().item())
+    packed_fp4, packed_scales, tensor_scale = native_packer().pack_weight(
+        weight_nk_bf16
+    )
+    return packed_fp4, packed_scales, tensor_scale.reshape(1), amax
+
+
+def load_transformer_config(source_repo: Path) -> dict:
+    config_path = source_repo / "transformer" / "config.json"
+    if not config_path.exists():
+        fail(f"missing transformer config: {config_path}")
+    return json.loads(config_path.read_text())
+
+
+def source_transformer_checkpoint(source_repo: Path) -> Path:
+    checkpoint = source_repo / "transformer" / "diffusion_pytorch_model.safetensors"
+    if not checkpoint.exists():
+        fail(f"missing transformer checkpoint: {checkpoint}")
+    return checkpoint
+
+
+def import_mage_transformer_symbols() -> tuple[type[nn.Module], object]:
+    if str(MAGE_ROOT) not in sys.path:
+        sys.path.insert(0, str(MAGE_ROOT))
+    try:
+        from mage_flow.models.mage_flow import MageFlow, MageFlowParams
+    except Exception as exc:
+        fail(f"failed to import local Mage transformer sources from {MAGE_ROOT}: {exc}")
+    return MageFlow, MageFlowParams
+
+
+class PackedNvfp4LinearArtifactModule(nn.Module):
+    def __init__(
+        self,
+        in_features: int,
+        out_features: int,
+        packed_weight_e2m1: torch.Tensor,
+        packed_scales_ue4m3: torch.Tensor,
+        weight_tensor_scale: torch.Tensor,
+        bias_bf16: torch.Tensor | None,
+    ) -> None:
+        super().__init__()
+        self.in_features = int(in_features)
+        self.out_features = int(out_features)
+        self.register_buffer("packed_weight_e2m1", packed_weight_e2m1.contiguous())
+        self.register_buffer("packed_scales_ue4m3", packed_scales_ue4m3.contiguous())
+        self.register_buffer("weight_tensor_scale", weight_tensor_scale.contiguous())
+        if bias_bf16 is None:
+            self.bias_bf16 = None
+        else:
+            self.register_buffer("bias_bf16", bias_bf16.contiguous())
+
+    def forward(self, inputs: torch.Tensor) -> torch.Tensor:
+        raise RuntimeError(
+            "PackedNvfp4LinearArtifactModule is an artifact-only placeholder. "
+            "Attach the resident CUDA runtime before calling forward()."
+        )
+
+    def extra_repr(self) -> str:
+        return (
+            f"in_features={self.in_features}, out_features={self.out_features}, "
+            f"packed_weight_bytes={self.packed_weight_e2m1.numel()}, "
+            f"packed_scale_bytes={self.packed_scales_ue4m3.numel()}, "
+            f"has_bias={self.bias_bf16 is not None}"
+        )
+
+
+def instantiate_mage_transformer_on_meta(source_repo: Path) -> nn.Module:
+    config = load_transformer_config(source_repo)
+    MageFlow, MageFlowParams = import_mage_transformer_symbols()
+    structure = {
+        key: value
+        for key, value in config.items()
+        if key
+        not in {
+            "_class_name",
+            "txt_max_length",
+            "max_sequence_length",
+            "param_dtype",
+            "packing",
+            "schedule_mode",
+            "static_shift",
+            "use_time_shift",
+            "rope_type",
+            "apply_text_rotary_emb",
+            "mlp_ratio",
+            "depth_single_blocks",
+            "theta",
+            "qkv_bias",
+            "guidance_embed",
+            "vec_in_dim",
+            "vec_type",
+            "time_type",
+            "double_block_type",
+            "quantization_config",
+        }
+    }
+    with torch.device("meta"):
+        model = MageFlow(MageFlowParams(**structure))
+    return model
+
+
+def unregistered_meta_tensor_attribute_names(model: nn.Module) -> list[str]:
+    """Find direct tensor attributes that PyTorch's parameter/buffer walk misses."""
+    names: list[str] = []
+    for module_name, module in model.named_modules():
+        registered_names = set(module._parameters) | set(module._buffers)
+        for attribute_name, value in vars(module).items():
+            if attribute_name in registered_names:
+                continue
+            if isinstance(value, torch.Tensor) and value.is_meta:
+                prefix = f"{module_name}." if module_name else ""
+                names.append(f"{prefix}{attribute_name}")
+    return sorted(names)
+
+
+def materialize_mage_rope_tensor_attributes(model: nn.Module) -> list[str]:
+    """Rebuild Mage's intentionally unregistered complex RoPE tensors on CPU."""
+    before = unregistered_meta_tensor_attribute_names(model)
+    expected = ["pos_embed.neg_freqs", "pos_embed.pos_freqs"]
+    if before != expected:
+        fail(
+            "unexpected unregistered meta tensor attributes before RoPE "
+            f"materialization: {before}"
+        )
+
+    rope = model.get_submodule("pos_embed")
+    rope_type = type(rope)
+    with torch.device("cpu"):
+        materialized = rope_type(
+            theta=rope.theta,
+            axes_dim=list(rope.axes_dim),
+            scale_rope=rope.scale_rope,
+        )
+    rope.pos_freqs = materialized.pos_freqs
+    rope.neg_freqs = materialized.neg_freqs
+    rope.video_freq_cache = {}
+
+    remaining = unregistered_meta_tensor_attribute_names(model)
+    if remaining:
+        fail(
+            "unresolved unregistered meta tensor attributes after RoPE "
+            f"materialization: {remaining}"
+        )
+    return before
+
+
+def set_child_module(root: nn.Module, dotted_path: str, module: nn.Module) -> None:
+    parent_path, _, child_name = dotted_path.rpartition(".")
+    parent = root.get_submodule(parent_path) if parent_path else root
+    if child_name.isdigit() and isinstance(parent, (nn.Sequential, nn.ModuleList)):
+        parent[int(child_name)] = module
+    else:
+        setattr(parent, child_name, module)
+
+
+def replace_targets_with_artifact_modules(
+    model: nn.Module,
+    artifact_path: Path,
+    target_specs: Iterable[TargetSpec],
+) -> None:
+    with safe_open(artifact_path, framework="pt", device="cpu") as handle:
+        for spec in target_specs:
+            packed_weight = handle.get_tensor(spec.artifact_weight_key)
+            packed_scales = handle.get_tensor(spec.artifact_scale_key)
+            weight_tensor_scale = handle.get_tensor(spec.artifact_tensor_scale_key)
+            bias = handle.get_tensor(spec.artifact_bias_key)
+            original = model.get_submodule(spec.module_key)
+            if not isinstance(original, nn.Linear):
+                fail(f"expected target module {spec.module_key} to be nn.Linear")
+            replacement = PackedNvfp4LinearArtifactModule(
+                in_features=int(original.in_features),
+                out_features=int(original.out_features),
+                packed_weight_e2m1=packed_weight,
+                packed_scales_ue4m3=packed_scales,
+                weight_tensor_scale=weight_tensor_scale,
+                bias_bf16=bias,
+            )
+            set_child_module(model, spec.module_key, replacement)
+
+
+def replace_targets_with_resident_modules(
+    model: nn.Module,
+    artifact_path: Path,
+    target_specs: Iterable[TargetSpec],
+    device: torch.device,
+) -> None:
+    if device.type != "cuda":
+        fail("resident runtime modules require a CUDA destination")
+    if str(RESIDENT_PYTHON_ROOT) not in sys.path:
+        sys.path.insert(0, str(RESIDENT_PYTHON_ROOT))
+    from torch_ops import PackedNvfp4LinearOp
+
+    _replace_targets_with_registered_modules(
+        model,
+        artifact_path,
+        target_specs,
+        device,
+        PackedNvfp4LinearOp,
+    )
+
+
+def replace_targets_with_native_resident_modules(
+    model: nn.Module,
+    artifact_path: Path,
+    target_specs: Iterable[TargetSpec],
+    device: torch.device,
+) -> None:
+    if device.type != "cuda":
+        fail("native resident runtime modules require a CUDA destination")
+    if str(RESIDENT_PYTHON_ROOT) not in sys.path:
+        sys.path.insert(0, str(RESIDENT_PYTHON_ROOT))
+    from torch_ops_native import (
+        PackedNvfp4LinearNativeOp,
+        initialize_native_sm120_op,
+    )
+
+    if not initialize_native_sm120_op(allow_python_schema_fallback=False):
+        fail("compiled native resident torch op did not load")
+    _replace_targets_with_registered_modules(
+        model,
+        artifact_path,
+        target_specs,
+        device,
+        PackedNvfp4LinearNativeOp,
+    )
+
+
+def _replace_targets_with_registered_modules(
+    model: nn.Module,
+    artifact_path: Path,
+    target_specs: Iterable[TargetSpec],
+    device: torch.device,
+    module_cls: type[nn.Module],
+) -> None:
+    with safe_open(artifact_path, framework="pt", device="cpu") as handle:
+        for spec in target_specs:
+            original = model.get_submodule(spec.module_key)
+            if not isinstance(original, nn.Linear):
+                fail(f"expected target module {spec.module_key} to be nn.Linear")
+            replacement = module_cls(
+                in_features=int(original.in_features),
+                out_features=int(original.out_features),
+                packed_weight=handle.get_tensor(spec.artifact_weight_key).to(device),
+                weight_scales=handle.get_tensor(spec.artifact_scale_key).to(device),
+                weight_scale=handle.get_tensor(
+                    spec.artifact_tensor_scale_key
+                ).to(device),
+                bias=handle.get_tensor(spec.artifact_bias_key).to(device),
+            )
+            set_child_module(model, spec.module_key, replacement)
+
+
+def assign_tensor_by_name(model: nn.Module, key: str, tensor: torch.Tensor) -> None:
+    if "." not in key:
+        parent = model
+        leaf = key
+    else:
+        parent_path, _, leaf = key.rpartition(".")
+        parent = model.get_submodule(parent_path)
+    if leaf in parent._parameters:
+        requires_grad = parent._parameters[leaf].requires_grad
+        parent._parameters[leaf] = nn.Parameter(tensor, requires_grad=requires_grad)
+        return
+    if leaf in parent._buffers:
+        parent._buffers[leaf] = tensor
+        return
+    fail(f"destination key {key} was neither a parameter nor a buffer")
+
+
+def pack_artifact(source_repo: Path, output_dir: Path) -> Path:
+    source_repo = source_repo.resolve()
+    output_dir = output_dir.resolve()
+    if output_dir.exists():
+        fail(f"output directory already exists: {output_dir}")
+    output_dir.mkdir(parents=True, exist_ok=False)
+
+    config = load_transformer_config(source_repo)
+    checkpoint_path = source_transformer_checkpoint(source_repo)
+    target_specs = build_target_specs(int(config["depth"]))
+    target_weight_keys = {spec.weight_key for spec in target_specs}
+    target_bias_keys = {spec.bias_key for spec in target_specs}
+    target_keys = target_weight_keys | target_bias_keys
+
+    artifact_tensors: OrderedDict[str, torch.Tensor] = OrderedDict()
+    target_metadata: list[dict] = []
+
+    with safe_open(checkpoint_path, framework="pt", device="cpu") as handle:
+        source_keys = list(handle.keys())
+        source_key_set = set(source_keys)
+        missing = sorted(target_keys - source_key_set)
+        if missing:
+            fail(f"source checkpoint is missing {len(missing)} target tensors, first={missing[0]}")
+
+        for spec in target_specs:
+            weight = handle.get_tensor(spec.weight_key)
+            bias = handle.get_tensor(spec.bias_key)
+            if weight.dtype != torch.bfloat16:
+                fail(f"{spec.weight_key} expected bfloat16, found {weight.dtype}")
+            if bias.dtype != torch.bfloat16:
+                fail(f"{spec.bias_key} expected bfloat16, found {bias.dtype}")
+            packed_weight, packed_scales, weight_tensor_scale, global_amax = pack_weight_tensor(weight)
+            scale_layout = make_scale_layout(weight.shape[1], weight.shape[0])
+            artifact_tensors[spec.artifact_bias_key] = bias.contiguous()
+            artifact_tensors[spec.artifact_scale_key] = packed_scales
+            artifact_tensors[spec.artifact_tensor_scale_key] = weight_tensor_scale
+            artifact_tensors[spec.artifact_weight_key] = packed_weight
+            target_metadata.append(
+                {
+                    "module_key": spec.module_key,
+                    "weight_key": spec.weight_key,
+                    "bias_key": spec.bias_key,
+                    "weight_shape": list(weight.shape),
+                    "bias_shape": list(bias.shape),
+                    "weight_tensor_scale_key": spec.artifact_tensor_scale_key,
+                    "artifact_weight_key": spec.artifact_weight_key,
+                    "artifact_scale_key": spec.artifact_scale_key,
+                    "artifact_bias_key": spec.artifact_bias_key,
+                    "weight_tensor_scale": float(weight_tensor_scale.item()),
+                    "weight_global_amax": global_amax,
+                    "packed_weight_bytes": int(packed_weight.numel()),
+                    "packed_scale_bytes": int(packed_scales.numel()),
+                    "scale_layout": asdict(scale_layout),
+                    "source_weight_sha256": sha256_tensor(weight),
+                    "source_bias_sha256": sha256_tensor(bias),
+                }
+            )
+
+    non_target_keys = sorted(set(source_keys) - target_keys)
+    artifact_path = output_dir / "packed_transformer.safetensors"
+    save_file(
+        OrderedDict(sorted(artifact_tensors.items())),
+        artifact_path,
+        metadata=CANONICAL_SAFETENSORS_METADATA,
+    )
+    fsync_file(artifact_path)
+    fsync_directory(output_dir)
+
+    library_hashes = {
+        "artifact_script_sha256": sha256_file(ARTIFACT_SCRIPT_PATH),
+        "resident_source_sha256": sha256_file(RESIDENT_SOURCE),
+        "resident_library_sha256": sha256_file(RESIDENT_LIBRARY),
+        "mage_flow_py_sha256": sha256_file(MAGE_ROOT / "mage_flow" / "models" / "mage_flow.py"),
+        "mage_layers_py_sha256": sha256_file(MAGE_ROOT / "mage_flow" / "models" / "modules" / "mage_layers.py"),
+        "pipeline_py_sha256": sha256_file(MAGE_ROOT / "mage_flow" / "pipeline.py"),
+    }
+
+    metadata = OrderedDict(
+        (
+            ("format_version", FORMAT_VERSION),
+            ("artifact_kind", ARTIFACT_KIND),
+            ("created_utc", datetime.now(timezone.utc).isoformat()),
+            (
+                "container",
+                {
+                    "format": "safetensors",
+                    "header_metadata": CANONICAL_SAFETENSORS_METADATA,
+                    "header_encoding": (
+                        "single deterministic contract key; legacy two-key "
+                        "draft headers remain readable"
+                    ),
+                },
+            ),
+            (
+                "source",
+                OrderedDict(
+                    (
+                        ("transformer_config_path", str(source_repo / "transformer" / "config.json")),
+                        ("transformer_checkpoint_path", str(checkpoint_path)),
+                        ("transformer_config_sha256", sha256_file(source_repo / "transformer" / "config.json")),
+                        ("transformer_checkpoint_sha256", sha256_file(checkpoint_path)),
+                    )
+                ),
+            ),
+            ("library_hashes", library_hashes),
+            (
+                "environment",
+                {
+                    "python_version": sys.version,
+                    "torch_version": torch.__version__,
+                    "safetensors_version": safetensors.__version__,
+                },
+            ),
+            (
+                "model",
+                {
+                    "depth": int(config["depth"]),
+                    "hidden_size": int(config["hidden_size"]),
+                    "num_heads": int(config["num_heads"]),
+                    "context_in_dim": int(config["context_in_dim"]),
+                    "in_channels": int(config["in_channels"]),
+                    "out_channels": int(config["out_channels"]),
+                    "patch_size": int(config["patch_size"]),
+                },
+            ),
+            (
+                "quantization",
+                {
+                    "format": "nvfp4_two_level",
+                    "block_elements": FP4_BLOCK_ELEMENTS,
+                    "scale_tile_outer": SCALE_TILE_OUTER,
+                    "scale_tile_inner": SCALE_TILE_INNER,
+                    "bias_policy": "artifact_bfloat16",
+                },
+            ),
+            ("targets", target_metadata),
+            ("non_target_keys", non_target_keys),
+        )
+    )
+    write_bytes_once(
+        output_dir / "metadata.json",
+        (json.dumps(metadata, indent=2, sort_keys=False) + "\n").encode("utf-8"),
+    )
+    return output_dir
+
+
+def load_validated_artifact_metadata(
+    artifact_dir: Path,
+    source_repo: Path,
+    *,
+    require_resident_runtime: bool = False,
+) -> dict:
+    artifact_dir = artifact_dir.resolve()
+    source_repo = source_repo.resolve()
+    metadata_path = artifact_dir / "metadata.json"
+    artifact_path = artifact_dir / "packed_transformer.safetensors"
+    if not metadata_path.is_file() or not artifact_path.is_file():
+        fail(f"artifact dir missing metadata or safetensors: {artifact_dir}")
+    try:
+        metadata = json.loads(metadata_path.read_text(encoding="utf-8"))
+    except (OSError, json.JSONDecodeError) as exc:
+        fail(f"invalid artifact metadata {metadata_path}: {exc}")
+    if metadata.get("format_version") != FORMAT_VERSION:
+        fail(f"unsupported format_version: {metadata.get('format_version')}")
+    if metadata.get("artifact_kind") != ARTIFACT_KIND:
+        fail(f"unexpected artifact_kind: {metadata.get('artifact_kind')}")
+
+    config_path = source_repo / "transformer" / "config.json"
+    checkpoint_path = source_transformer_checkpoint(source_repo)
+    expected_config_hash = sha256_file(config_path)
+    expected_checkpoint_hash = sha256_file(checkpoint_path)
+    try:
+        source_metadata = metadata["source"]
+        recorded_config_hash = source_metadata["transformer_config_sha256"]
+        recorded_checkpoint_hash = source_metadata[
+            "transformer_checkpoint_sha256"
+        ]
+        model_metadata = metadata["model"]
+        recorded_depth = int(model_metadata["depth"])
+        recorded_targets = metadata["targets"]
+        recorded_non_target_keys = metadata["non_target_keys"]
+    except (KeyError, TypeError, ValueError) as exc:
+        fail(f"artifact metadata schema is incomplete or invalid: {exc}")
+    if not isinstance(recorded_targets, list):
+        fail("artifact metadata targets must be a list")
+    if not isinstance(recorded_non_target_keys, list) or not all(
+        isinstance(key, str) for key in recorded_non_target_keys
+    ):
+        fail("artifact metadata non_target_keys must be a list of strings")
+    if recorded_config_hash != expected_config_hash:
+        fail("transformer config hash mismatch")
+    if recorded_checkpoint_hash != expected_checkpoint_hash:
+        fail("transformer checkpoint hash mismatch")
+
+    config = load_transformer_config(source_repo)
+    if recorded_depth != int(config["depth"]):
+        fail("artifact model depth does not match source config")
+    specs = build_target_specs(int(config["depth"]))
+    expected_modules = [spec.module_key for spec in specs]
+    if not all(isinstance(entry, dict) for entry in recorded_targets):
+        fail("artifact metadata target entries must be objects")
+    recorded_modules = [entry.get("module_key") for entry in recorded_targets]
+    if recorded_modules != expected_modules:
+        fail("artifact target allowlist/order mismatch")
+    target_source_keys = {
+        key for spec in specs for key in (spec.weight_key, spec.bias_key)
+    }
+    expected_artifact_keys = {
+        key
+        for spec in specs
+        for key in (
+            spec.artifact_weight_key,
+            spec.artifact_scale_key,
+            spec.artifact_tensor_scale_key,
+            spec.artifact_bias_key,
+        )
+    }
+    with safe_open(artifact_path, framework="pt", device="cpu") as artifact_handle:
+        actual_artifact_keys = set(artifact_handle.keys())
+        header_metadata = artifact_handle.metadata()
+    if actual_artifact_keys != expected_artifact_keys:
+        fail("artifact tensor key coverage mismatch")
+    if header_metadata not in (
+        CANONICAL_SAFETENSORS_METADATA,
+        LEGACY_SAFETENSORS_METADATA,
+    ):
+        fail("artifact safetensors header metadata mismatch")
+
+    with safe_open(checkpoint_path, framework="pt", device="cpu") as source_handle:
+        source_keys = set(source_handle.keys())
+    missing_target_keys = sorted(target_source_keys - source_keys)
+    if missing_target_keys:
+        fail(
+            "source checkpoint is missing target tensors, first="
+            f"{missing_target_keys[0]}"
+        )
+    expected_non_target_keys = sorted(source_keys - target_source_keys)
+    if recorded_non_target_keys != expected_non_target_keys:
+        fail("artifact non-target source manifest mismatch")
+
+    if require_resident_runtime:
+        hashes = metadata.get("library_hashes", {})
+        if not isinstance(hashes, dict):
+            fail("artifact metadata library_hashes must be an object")
+        if hashes.get("resident_source_sha256") != sha256_file(RESIDENT_SOURCE):
+            fail("resident source hash mismatch")
+        if hashes.get("resident_library_sha256") != sha256_file(RESIDENT_LIBRARY):
+            fail("resident library hash mismatch")
+    return metadata
+
+
+def validate_artifact(artifact_dir: Path, source_repo: Path) -> None:
+    artifact_dir = artifact_dir.resolve()
+    source_repo = source_repo.resolve()
+    metadata = load_validated_artifact_metadata(artifact_dir, source_repo)
+    artifact_path = artifact_dir / "packed_transformer.safetensors"
+    checkpoint_path = source_transformer_checkpoint(source_repo)
+    config = load_transformer_config(source_repo)
+    specs = {
+        spec.module_key: spec for spec in build_target_specs(int(config["depth"]))
+    }
+
+    with safe_open(artifact_path, framework="pt", device="cpu") as artifact_handle, safe_open(
+        checkpoint_path, framework="pt", device="cpu"
+    ) as source_handle:
+        for entry in metadata["targets"]:
+            spec = specs[entry["module_key"]]
+            weight = source_handle.get_tensor(spec.weight_key)
+            bias = source_handle.get_tensor(spec.bias_key)
+            packed_weight, packed_scales, weight_tensor_scale, global_amax = pack_weight_tensor(weight)
+
+            candidate_weight = artifact_handle.get_tensor(spec.artifact_weight_key)
+            candidate_scales = artifact_handle.get_tensor(spec.artifact_scale_key)
+            candidate_tensor_scale = artifact_handle.get_tensor(spec.artifact_tensor_scale_key)
+            candidate_bias = artifact_handle.get_tensor(spec.artifact_bias_key)
+
+            if not torch.equal(candidate_weight, packed_weight):
+                fail(f"packed weight mismatch for {spec.module_key}")
+            if not torch.equal(candidate_scales, packed_scales):
+                fail(f"packed scales mismatch for {spec.module_key}")
+            if not torch.equal(candidate_tensor_scale, weight_tensor_scale):
+                fail(f"tensor scale mismatch for {spec.module_key}")
+            if not torch.equal(candidate_bias, bias):
+                fail(f"bias mismatch for {spec.module_key}")
+            if abs(float(entry["weight_global_amax"]) - global_amax) > 0.0:
+                fail(f"global amax mismatch for {spec.module_key}")
+
+
+def load_clean_transformer_from_artifact(
+    artifact_dir: Path,
+    source_repo: Path,
+    assign_non_target: bool = True,
+) -> nn.Module:
+    artifact_dir = artifact_dir.resolve()
+    source_repo = source_repo.resolve()
+    metadata = load_validated_artifact_metadata(artifact_dir, source_repo)
+    target_specs = build_target_specs(int(metadata["model"]["depth"]))
+    model = instantiate_mage_transformer_on_meta(source_repo)
+    replace_targets_with_artifact_modules(model, artifact_dir / "packed_transformer.safetensors", target_specs)
+
+    if assign_non_target:
+        skip_keys = {
+            key
+            for spec in target_specs
+            for key in (spec.weight_key, spec.bias_key)
+        }
+        source_tensor_keys_read: list[str] = []
+        with safe_open(source_transformer_checkpoint(source_repo), framework="pt", device="cpu") as source_handle:
+            for key in source_handle.keys():
+                if key in skip_keys:
+                    continue
+                tensor = source_handle.get_tensor(key)
+                source_tensor_keys_read.append(key)
+                assign_tensor_by_name(model, key, tensor)
+        target_reads = sorted(set(source_tensor_keys_read) & skip_keys)
+        if target_reads:
+            fail(f"clean CPU loader read target source tensors: {target_reads[0]}")
+        meta_parameters = [
+            name for name, parameter in model.named_parameters() if parameter.is_meta
+        ]
+        meta_buffers = [
+            name for name, buffer in model.named_buffers() if buffer.is_meta
+        ]
+        if meta_parameters or meta_buffers:
+            first = (meta_parameters + meta_buffers)[0]
+            fail(f"clean CPU loader left unresolved meta tensors, first={first}")
+        materialize_mage_rope_tensor_attributes(model)
+    return model
+
+
+def load_clean_resident_transformer(
+    artifact_dir: Path,
+    source_repo: Path,
+    device: torch.device,
+) -> tuple[nn.Module, dict]:
+    return _load_clean_cuda_transformer(
+        artifact_dir,
+        source_repo,
+        device,
+        replace_targets_with_resident_modules,
+    )
+
+
+def load_clean_native_resident_transformer(
+    artifact_dir: Path,
+    source_repo: Path,
+    device: torch.device,
+) -> tuple[nn.Module, dict]:
+    return _load_clean_cuda_transformer(
+        artifact_dir,
+        source_repo,
+        device,
+        replace_targets_with_native_resident_modules,
+    )
+
+
+def _load_clean_cuda_transformer(
+    artifact_dir: Path,
+    source_repo: Path,
+    device: torch.device,
+    target_replacement_fn: Callable[[nn.Module, Path, Iterable[TargetSpec], torch.device], None],
+) -> tuple[nn.Module, dict]:
+    artifact_dir = artifact_dir.resolve()
+    source_repo = source_repo.resolve()
+    metadata = load_validated_artifact_metadata(
+        artifact_dir,
+        source_repo,
+        require_resident_runtime=True,
+    )
+    target_specs = build_target_specs(int(metadata["model"]["depth"]))
+    target_source_keys = {
+        key
+        for spec in target_specs
+        for key in (spec.weight_key, spec.bias_key)
+    }
+    model = instantiate_mage_transformer_on_meta(source_repo)
+    target_replacement_fn(
+        model,
+        artifact_dir / "packed_transformer.safetensors",
+        target_specs,
+        device,
+    )
+
+    loaded_source_keys: list[str] = []
+    skipped_target_source_keys: list[str] = []
+    source_tensor_keys_read: list[str] = []
+    with safe_open(
+        source_transformer_checkpoint(source_repo),
+        framework="pt",
+        device="cpu",
+    ) as source_handle:
+        source_keys = list(source_handle.keys())
+        missing_target_keys = sorted(target_source_keys - set(source_keys))
+        if missing_target_keys:
+            fail(
+                "source checkpoint is missing target tensors, first="
+                f"{missing_target_keys[0]}"
+            )
+        for key in source_keys:
+            if key in target_source_keys:
+                skipped_target_source_keys.append(key)
+                continue
+            tensor = source_handle.get_tensor(key)
+            source_tensor_keys_read.append(key)
+            assign_tensor_by_name(model, key, tensor.to(device))
+            loaded_source_keys.append(key)
+
+    materialized_tensor_attributes = materialize_mage_rope_tensor_attributes(model)
+    target_source_reads = sorted(
+        set(source_tensor_keys_read) & target_source_keys
+    )
+    meta_parameters = [
+        name for name, parameter in model.named_parameters() if parameter.is_meta
+    ]
+    meta_buffers = [
+        name for name, buffer in model.named_buffers() if buffer.is_meta
+    ]
+    report = {
+        "source_tensor_count_loaded": len(loaded_source_keys),
+        "source_tensor_keys_loaded": loaded_source_keys,
+        "source_tensor_count_read": len(source_tensor_keys_read),
+        "source_tensor_keys_read": source_tensor_keys_read,
+        "target_source_tensor_count_skipped": len(skipped_target_source_keys),
+        "target_source_tensor_keys_skipped": skipped_target_source_keys,
+        "target_source_tensor_reads": len(target_source_reads),
+        "target_source_tensor_keys_read": target_source_reads,
+        "meta_parameter_names": meta_parameters,
+        "meta_buffer_names": meta_buffers,
+        "materialized_unregistered_tensor_attribute_names": (
+            materialized_tensor_attributes
+        ),
+        "unregistered_meta_tensor_attribute_names": (
+            unregistered_meta_tensor_attribute_names(model)
+        ),
+    }
+    return model.eval().requires_grad_(False), report
+
+
+def parse_args(argv: list[str] | None = None) -> argparse.Namespace:
+    parser = argparse.ArgumentParser(description=__doc__)
+    subparsers = parser.add_subparsers(dest="command", required=True)
+
+    pack_parser = subparsers.add_parser("pack", help="build a packed artifact directory")
+    pack_parser.add_argument("--source-repo", type=Path, required=True)
+    pack_parser.add_argument("--output-dir", type=Path, required=True)
+
+    validate_parser = subparsers.add_parser("validate", help="recompute and validate a packed artifact")
+    validate_parser.add_argument("--artifact-dir", type=Path, required=True)
+    validate_parser.add_argument("--source-repo", type=Path, required=True)
+
+    plan_load_parser = subparsers.add_parser("plan-load", help="instantiate on meta and replace target modules")
+    plan_load_parser.add_argument("--artifact-dir", type=Path, required=True)
+    plan_load_parser.add_argument("--source-repo", type=Path, required=True)
+    plan_load_parser.add_argument("--skip-non-target", action="store_true")
+
+    runtime_parser = subparsers.add_parser(
+        "validate-runtime",
+        help="placeholder for future single-GPU resident validation",
+    )
+    runtime_parser.add_argument("--artifact-dir", type=Path, required=True)
+    runtime_parser.add_argument("--source-repo", type=Path, required=True)
+    return parser.parse_args(argv)
+
+
+def main(argv: list[str] | None = None) -> int:
+    args = parse_args(argv)
+    try:
+        if args.command == "pack":
+            artifact_dir = pack_artifact(args.source_repo, args.output_dir)
+            print(artifact_dir)
+            return 0
+        if args.command == "validate":
+            validate_artifact(args.artifact_dir, args.source_repo)
+            print("ok")
+            return 0
+        if args.command == "plan-load":
+            model = load_clean_transformer_from_artifact(
+                args.artifact_dir, args.source_repo, assign_non_target=not args.skip_non_target
+            )
+            packed_count = sum(
+                1 for _name, module in model.named_modules() if isinstance(module, PackedNvfp4LinearArtifactModule)
+            )
+            meta_parameters = [
+                name for name, parameter in model.named_parameters() if parameter.is_meta
+            ]
+            meta_buffers = [
+                name for name, buffer in model.named_buffers() if buffer.is_meta
+            ]
+            print(
+                json.dumps(
+                    {
+                        "packed_module_count": packed_count,
+                        "meta_parameter_count": len(meta_parameters),
+                        "meta_buffer_count": len(meta_buffers),
+                        "non_target_assignment_skipped": bool(args.skip_non_target),
+                    },
+                    indent=2,
+                )
+            )
+            return 0
+        if args.command == "validate-runtime":
+            fail(
+                "validate-runtime is intentionally not implemented in this CPU-only slice. "
+                "Use the future CUDA resident path on CUDA_VISIBLE_DEVICES=3."
+            )
+        fail(f"unsupported command: {args.command}")
+    except PackedArtifactError as exc:
+        print(f"error: {exc}", file=sys.stderr)
+        return 1
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())

runtime/packed_nvfp4_linear.py

@@ -0,0 +1,485 @@
+"""Inference-only resident NVFP4 linear prototype.
+
+This module intentionally uses a narrow ctypes boundary. It proves packed
+residency and Mage shape correctness; it is not yet a torch.compile/CUDA-graph
+shipping operator.
+"""
+
+from __future__ import annotations
+
+import atexit
+import ctypes
+import threading
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Final
+
+import torch
+from torch import nn
+
+
+ABI_VERSION: Final = 1
+FP4_BLOCK_ELEMENTS: Final = 16
+SCALE_TILE_OUTER: Final = 128
+SCALE_TILE_INNER: Final = 4
+RELEASE_ROOT: Final = Path(__file__).resolve().parents[1]
+DEFAULT_LIBRARY_PATH: Final = RELEASE_ROOT / "runtime" / "libmage_nvfp4_linear.so"
+
+
+def round_up(value: int, multiple: int) -> int:
+    if value <= 0 or multiple <= 0:
+        raise ValueError("value and multiple must be positive")
+    return ((value + multiple - 1) // multiple) * multiple
+
+
+@dataclass(frozen=True)
+class ScaleLayout:
+    inner_dim: int
+    outer_tiles: int
+    num_bytes: int
+
+
+def scale_layout(k: int, outer_columns: int) -> ScaleLayout:
+    if k <= 0 or k % FP4_BLOCK_ELEMENTS:
+        raise ValueError("K must be positive and divisible by 16")
+    if outer_columns <= 0:
+        raise ValueError("outer column count must be positive")
+    inner_dim = round_up(k // FP4_BLOCK_ELEMENTS, SCALE_TILE_INNER)
+    outer_tiles = (outer_columns + SCALE_TILE_OUTER - 1) // SCALE_TILE_OUTER
+    return ScaleLayout(
+        inner_dim=inner_dim,
+        outer_tiles=outer_tiles,
+        num_bytes=outer_tiles * inner_dim * SCALE_TILE_OUTER,
+    )
+
+
+def packed_weight_num_bytes(out_features: int, in_features: int) -> int:
+    if out_features <= 0 or out_features % 8:
+        raise ValueError("out_features must be positive and divisible by 8")
+    if in_features <= 0 or in_features % 32:
+        raise ValueError("in_features must be positive and divisible by 32")
+    return out_features * in_features // 2
+
+
+def padded_output_shape(input_shape: tuple[int, ...], out_features: int) -> tuple[int, int]:
+    if not input_shape:
+        raise ValueError("input must have at least one dimension")
+    logical_m = 1
+    for dimension in input_shape[:-1]:
+        if dimension <= 0:
+            raise ValueError("empty or negative leading dimensions are unsupported")
+        logical_m *= dimension
+    return round_up(logical_m, 8), out_features
+
+
+def logical_output_shape(input_shape: tuple[int, ...], out_features: int) -> tuple[int, ...]:
+    if not input_shape:
+        raise ValueError("input must have at least one dimension")
+    return (*input_shape[:-1], out_features)
+
+
+class NativeNvfp4Library:
+    """Typed ctypes access to the project-local native library."""
+
+    def __init__(self, path: str | Path = DEFAULT_LIBRARY_PATH):
+        self.path = Path(path).resolve()
+        if not self.path.is_file():
+            raise FileNotFoundError(
+                f"resident NVFP4 library is not built: {self.path}"
+            )
+        self._library = ctypes.CDLL(str(self.path))
+        self._bind()
+        version = int(self._library.mage_nvfp4_abi_version())
+        if version != ABI_VERSION:
+            raise RuntimeError(
+                f"resident NVFP4 ABI mismatch: Python={ABI_VERSION}, native={version}"
+            )
+
+    def _bind(self) -> None:
+        library = self._library
+        library.mage_nvfp4_abi_version.argtypes = []
+        library.mage_nvfp4_abi_version.restype = ctypes.c_int
+        library.mage_nvfp4_last_error.argtypes = []
+        library.mage_nvfp4_last_error.restype = ctypes.c_char_p
+        library.mage_nvfp4_packed_weight_bytes.argtypes = [
+            ctypes.c_int,
+            ctypes.c_int,
+        ]
+        library.mage_nvfp4_packed_weight_bytes.restype = ctypes.c_size_t
+        library.mage_nvfp4_weight_scale_bytes.argtypes = [
+            ctypes.c_int,
+            ctypes.c_int,
+        ]
+        library.mage_nvfp4_weight_scale_bytes.restype = ctypes.c_size_t
+        library.mage_nvfp4_pack_weight_bf16.argtypes = [
+            ctypes.c_void_p,
+            ctypes.c_int,
+            ctypes.c_int,
+            ctypes.c_void_p,
+            ctypes.c_size_t,
+            ctypes.c_void_p,
+            ctypes.c_size_t,
+            ctypes.POINTER(ctypes.c_float),
+        ]
+        library.mage_nvfp4_pack_weight_bf16.restype = ctypes.c_int
+        library.mage_nvfp4_create_context.argtypes = [
+            ctypes.c_int,
+            ctypes.POINTER(ctypes.c_void_p),
+        ]
+        library.mage_nvfp4_create_context.restype = ctypes.c_int
+        library.mage_nvfp4_destroy_context.argtypes = [ctypes.c_void_p]
+        library.mage_nvfp4_destroy_context.restype = ctypes.c_int
+        library.mage_nvfp4_context_reserved_bytes.argtypes = [ctypes.c_void_p]
+        library.mage_nvfp4_context_reserved_bytes.restype = ctypes.c_size_t
+        library.mage_nvfp4_linear_forward.argtypes = [
+            ctypes.c_void_p,
+            ctypes.c_void_p,
+            ctypes.c_void_p,
+            ctypes.c_size_t,
+            ctypes.c_void_p,
+            ctypes.c_size_t,
+            ctypes.c_void_p,
+            ctypes.c_void_p,
+            ctypes.c_void_p,
+            ctypes.c_int,
+            ctypes.c_int,
+            ctypes.c_int,
+            ctypes.c_size_t,
+        ]
+        library.mage_nvfp4_linear_forward.restype = ctypes.c_int
+
+    def _error(self, operation: str) -> RuntimeError:
+        raw = self._library.mage_nvfp4_last_error()
+        message = raw.decode("utf-8", errors="replace") if raw else "unknown error"
+        return RuntimeError(f"{operation}: {message}")
+
+    def pack_weight(
+        self, weight: torch.Tensor
+    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+        if weight.device.type != "cpu":
+            raise ValueError("native weight packer requires a CPU tensor")
+        if weight.dtype != torch.bfloat16:
+            raise TypeError("native weight packer requires BF16")
+        if weight.ndim != 2 or not weight.is_contiguous():
+            raise ValueError("weight must be contiguous [N,K]")
+        n, k = map(int, weight.shape)
+        packed_bytes = packed_weight_num_bytes(n, k)
+        scale_bytes = scale_layout(k, n).num_bytes
+        native_packed = int(
+            self._library.mage_nvfp4_packed_weight_bytes(n, k)
+        )
+        native_scales = int(
+            self._library.mage_nvfp4_weight_scale_bytes(n, k)
+        )
+        if (native_packed, native_scales) != (packed_bytes, scale_bytes):
+            raise RuntimeError(
+                "Python/native packed metadata disagreement: "
+                f"Python={(packed_bytes, scale_bytes)}, "
+                f"native={(native_packed, native_scales)}"
+            )
+        packed = torch.empty(packed_bytes, dtype=torch.uint8, device="cpu")
+        scales = torch.empty(scale_bytes, dtype=torch.uint8, device="cpu")
+        tensor_scale = ctypes.c_float()
+        status = self._library.mage_nvfp4_pack_weight_bf16(
+            ctypes.c_void_p(weight.data_ptr()),
+            n,
+            k,
+            ctypes.c_void_p(packed.data_ptr()),
+            packed.numel(),
+            ctypes.c_void_p(scales.data_ptr()),
+            scales.numel(),
+            ctypes.byref(tensor_scale),
+        )
+        if status:
+            raise self._error("packing BF16 weight")
+        scale_tensor = torch.tensor(tensor_scale.value, dtype=torch.float32)
+        return packed, scales, scale_tensor
+
+
+_LIBRARY_LOCK = threading.Lock()
+_LIBRARY: NativeNvfp4Library | None = None
+
+
+def native_library() -> NativeNvfp4Library:
+    global _LIBRARY
+    with _LIBRARY_LOCK:
+        if _LIBRARY is None:
+            _LIBRARY = NativeNvfp4Library()
+        return _LIBRARY
+
+
+class ResidentContext:
+    def __init__(self, library: NativeNvfp4Library, device_index: int, stream: int):
+        self.library = library
+        self.device_index = int(device_index)
+        self.stream = int(stream)
+        self._pointer = ctypes.c_void_p()
+        self._lock = threading.Lock()
+        status = self.library._library.mage_nvfp4_create_context(
+            self.device_index, ctypes.byref(self._pointer)
+        )
+        if status:
+            raise self.library._error("creating resident context")
+        self._closed = False
+
+    @property
+    def pointer(self) -> ctypes.c_void_p:
+        if self._closed:
+            raise RuntimeError("resident NVFP4 context is closed")
+        return self._pointer
+
+    @property
+    def reserved_bytes(self) -> int:
+        return int(
+            self.library._library.mage_nvfp4_context_reserved_bytes(self.pointer)
+        )
+
+    def close(self) -> None:
+        with self._lock:
+            if self._closed:
+                return
+            status = self.library._library.mage_nvfp4_destroy_context(
+                self._pointer
+            )
+            if status:
+                raise self.library._error("destroying resident context")
+            self._closed = True
+            self._pointer = ctypes.c_void_p()
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        packed_weight: torch.Tensor,
+        weight_scales: torch.Tensor,
+        weight_scale: torch.Tensor,
+        bias: torch.Tensor | None,
+        output: torch.Tensor,
+        logical_m: int,
+        in_features: int,
+        out_features: int,
+    ) -> None:
+        bias_pointer = (
+            ctypes.c_void_p(bias.data_ptr()) if bias is not None else None
+        )
+        with self._lock:
+            status = self.library._library.mage_nvfp4_linear_forward(
+                self.pointer,
+                ctypes.c_void_p(x.data_ptr()),
+                ctypes.c_void_p(packed_weight.data_ptr()),
+                packed_weight.numel(),
+                ctypes.c_void_p(weight_scales.data_ptr()),
+                weight_scales.numel(),
+                ctypes.c_void_p(weight_scale.data_ptr()),
+                bias_pointer,
+                ctypes.c_void_p(output.data_ptr()),
+                logical_m,
+                in_features,
+                out_features,
+                self.stream,
+            )
+        if status:
+            raise self.library._error("resident NVFP4 linear forward")
+
+
+_CONTEXTS_LOCK = threading.Lock()
+_CONTEXTS: dict[tuple[int, int], ResidentContext] = {}
+
+
+def resident_context(device: torch.device) -> ResidentContext:
+    if device.type != "cuda":
+        raise ValueError("resident NVFP4 context requires CUDA")
+    device_index = (
+        torch.cuda.current_device() if device.index is None else int(device.index)
+    )
+    stream = int(torch.cuda.current_stream(device_index).cuda_stream)
+    key = (device_index, stream)
+    with _CONTEXTS_LOCK:
+        context = _CONTEXTS.get(key)
+        if context is None:
+            context = ResidentContext(native_library(), device_index, stream)
+            _CONTEXTS[key] = context
+        return context
+
+
+def close_all_contexts() -> None:
+    with _CONTEXTS_LOCK:
+        contexts = list(_CONTEXTS.values())
+        _CONTEXTS.clear()
+    errors: list[Exception] = []
+    for context in contexts:
+        try:
+            context.close()
+        except Exception as error:  # pragma: no cover - shutdown diagnostic
+            errors.append(error)
+    if errors:
+        raise RuntimeError(
+            "one or more resident NVFP4 contexts failed to close: "
+            + "; ".join(map(str, errors))
+        )
+
+
+def _quiet_atexit_close() -> None:
+    try:
+        close_all_contexts()
+    except Exception:
+        # CUDA may already be shutting down. Explicit close_all_contexts() is
+        # the auditable path; atexit is only a best-effort fallback.
+        pass
+
+
+atexit.register(_quiet_atexit_close)
+
+
+class PackedNvfp4Linear(nn.Module):
+    """Packed inference replacement for one BF16 ``nn.Linear``."""
+
+    def __init__(
+        self,
+        in_features: int,
+        out_features: int,
+        packed_weight: torch.Tensor,
+        weight_scales: torch.Tensor,
+        weight_scale: torch.Tensor,
+        bias: torch.Tensor | None,
+    ):
+        super().__init__()
+        expected_weight = packed_weight_num_bytes(out_features, in_features)
+        expected_scales = scale_layout(in_features, out_features).num_bytes
+        if (
+            packed_weight.dtype != torch.uint8
+            or packed_weight.ndim != 1
+            or not packed_weight.is_contiguous()
+            or packed_weight.numel() != expected_weight
+        ):
+            raise ValueError("packed_weight has invalid dtype, shape, or size")
+        if (
+            weight_scales.dtype != torch.uint8
+            or weight_scales.ndim != 1
+            or not weight_scales.is_contiguous()
+            or weight_scales.numel() != expected_scales
+        ):
+            raise ValueError("weight_scales has invalid dtype, shape, or size")
+        if (
+            weight_scale.dtype != torch.float32
+            or weight_scale.numel() != 1
+            or not weight_scale.is_contiguous()
+        ):
+            raise ValueError("weight_scale must be one contiguous FP32 value")
+        if bias is not None and (
+            bias.dtype != torch.bfloat16
+            or bias.shape != (out_features,)
+            or not bias.is_contiguous()
+        ):
+            raise ValueError("bias must be contiguous BF16 [out_features]")
+        devices = {
+            tensor.device
+            for tensor in (packed_weight, weight_scales, weight_scale, bias)
+            if tensor is not None
+        }
+        if len(devices) != 1:
+            raise ValueError("all resident buffers must be on one device")
+
+        self.in_features = int(in_features)
+        self.out_features = int(out_features)
+        self.register_buffer("packed_weight", packed_weight)
+        self.register_buffer("weight_scales", weight_scales)
+        self.register_buffer("weight_scale", weight_scale.reshape(()))
+        self.register_buffer("bias", bias)
+
+    @classmethod
+    def from_linear(
+        cls,
+        linear: nn.Linear,
+        device: torch.device | str,
+        *,
+        library: NativeNvfp4Library | None = None,
+    ) -> "PackedNvfp4Linear":
+        if not isinstance(linear, nn.Linear):
+            raise TypeError("from_linear requires torch.nn.Linear")
+        library = native_library() if library is None else library
+        destination = torch.device(device)
+        if destination.type != "cuda":
+            raise ValueError("resident packed buffers must target CUDA")
+        weight_cpu = (
+            linear.weight.detach()
+            .to(device="cpu", dtype=torch.bfloat16)
+            .contiguous()
+        )
+        packed, scales, tensor_scale = library.pack_weight(weight_cpu)
+        bias_cpu = (
+            None
+            if linear.bias is None
+            else linear.bias.detach()
+            .to(device="cpu", dtype=torch.bfloat16)
+            .contiguous()
+        )
+        return cls(
+            linear.in_features,
+            linear.out_features,
+            packed.to(destination),
+            scales.to(destination),
+            tensor_scale.to(destination),
+            None if bias_cpu is None else bias_cpu.to(destination),
+        )
+
+    def forward(self, input: torch.Tensor) -> torch.Tensor:
+        if input.device.type != "cuda":
+            raise ValueError("PackedNvfp4Linear requires a CUDA input")
+        if input.device != self.packed_weight.device:
+            raise ValueError("input and packed buffers are on different devices")
+        if input.dtype != torch.bfloat16:
+            raise TypeError("PackedNvfp4Linear requires BF16 input")
+        if input.ndim < 1 or input.shape[-1] != self.in_features:
+            raise ValueError(
+                f"expected last dimension {self.in_features}, got "
+                f"{tuple(input.shape)}"
+            )
+        if input.requires_grad:
+            raise RuntimeError("resident NVFP4 prototype is inference-only")
+
+        contiguous = input.reshape(-1, self.in_features).contiguous()
+        logical_m = int(contiguous.shape[0])
+        if logical_m <= 0:
+            raise ValueError("empty inputs are unsupported")
+        padded_m = round_up(logical_m, 8)
+        padded_output = torch.empty(
+            (padded_m, self.out_features),
+            dtype=torch.bfloat16,
+            device=input.device,
+        )
+        current_stream = torch.cuda.current_stream(input.device)
+        # ctypes launches are invisible to the caching allocator. Explicitly
+        # record every CUDA allocation read or written by the native call so a
+        # tensor produced on another stream cannot be recycled while the
+        # resident kernel is still using it.
+        for tensor in (
+            contiguous,
+            self.packed_weight,
+            self.weight_scales,
+            self.weight_scale,
+            self.bias,
+            padded_output,
+        ):
+            if tensor is not None:
+                tensor.record_stream(current_stream)
+        context = resident_context(input.device)
+        context.forward(
+            contiguous,
+            self.packed_weight,
+            self.weight_scales,
+            self.weight_scale,
+            self.bias,
+            padded_output,
+            logical_m,
+            self.in_features,
+            self.out_features,
+        )
+        logical = padded_output[:logical_m]
+        return logical.view(*input.shape[:-1], self.out_features)
+
+    def extra_repr(self) -> str:
+        return (
+            f"in_features={self.in_features}, "
+            f"out_features={self.out_features}, "
+            f"bias={self.bias is not None}, inference_only=True"
+        )

runtime/quant_text_encoder.py

@@ -0,0 +1,376 @@
+"""Load the packaged mixed NVFP4/FP8 Qwen3-VL text encoder without BF16 shards."""
+
+from __future__ import annotations
+
+from collections import Counter
+import hashlib
+import json
+from pathlib import Path
+from typing import Any
+
+import torch
+import torch.nn.functional as F
+from accelerate import init_empty_weights
+from accelerate.utils import set_module_tensor_to_device
+from safetensors import safe_open
+from transformers import AutoConfig, AutoProcessor, AutoTokenizer
+
+
+EXPECTED_ARTIFACT_SHA256 = (
+    "719906b435800757d22013d3d475a4853d59b779b022669fa8b8a193b85d0f41"
+)
+EXPECTED_FORMAT_COUNTS = {"nvfp4": 224, "float8_e4m3fn": 14}
+EXPECTED_PROJECTIONS = {
+    "mlp.down_proj",
+    "mlp.gate_proj",
+    "mlp.up_proj",
+    "self_attn.k_proj",
+    "self_attn.o_proj",
+    "self_attn.q_proj",
+    "self_attn.v_proj",
+}
+
+
+def sha256(path: Path) -> str:
+    digest = hashlib.sha256()
+    with path.open("rb") as handle:
+        for chunk in iter(lambda: handle.read(1 << 20), b""):
+            digest.update(chunk)
+    return digest.hexdigest()
+
+
+def _decode_quant_config(tensor: torch.Tensor) -> dict[str, Any]:
+    payload = bytes(tensor.cpu().to(torch.uint8).tolist()).decode("utf-8")
+    parsed = json.loads(payload)
+    if not isinstance(parsed, dict) or not isinstance(parsed.get("format"), str):
+        raise ValueError(f"invalid comfy_quant payload: {payload!r}")
+    return parsed
+
+
+def _artifact_layer_to_hf_module(layer_key: str) -> str:
+    prefix = "model.layers."
+    if not layer_key.startswith(prefix):
+        raise ValueError(f"quantized layer is outside the language stack: {layer_key}")
+    remainder = layer_key[len(prefix) :]
+    layer_text, projection = remainder.split(".", 1)
+    layer_index = int(layer_text)
+    if layer_index < 0 or layer_index >= 36:
+        raise ValueError(f"language layer index is out of range: {layer_key}")
+    if projection not in EXPECTED_PROJECTIONS:
+        raise ValueError(f"unexpected quantized projection: {layer_key}")
+    return f"model.language_model.layers.{layer_index}.{projection}"
+
+
+def _artifact_weight_to_hf_name(key: str) -> str | None:
+    if key == "model.embed_tokens.weight":
+        return "model.language_model.embed_tokens.weight"
+    if key == "model.norm.weight":
+        return "model.language_model.norm.weight"
+    if key.startswith("model.layers."):
+        return "model.language_model.layers." + key.removeprefix("model.layers.")
+    if key.startswith("model.visual."):
+        return key
+    return None
+
+
+def _resolve_parent(module: torch.nn.Module, dotted_name: str) -> tuple[Any, str]:
+    parts = dotted_name.split(".")
+    parent: Any = module
+    for part in parts[:-1]:
+        parent = getattr(parent, part)
+    return parent, parts[-1]
+
+
+class PublishedQuantLinear(torch.nn.Module):
+    """Inference-only projection backed by comfy-kitchen packed tensors."""
+
+    def __init__(
+        self,
+        *,
+        in_features: int,
+        out_features: int,
+        quant_format: str,
+        qdata: torch.Tensor,
+        weight_scale: torch.Tensor,
+        weight_scale_2: torch.Tensor | None,
+    ) -> None:
+        super().__init__()
+        self.in_features = int(in_features)
+        self.out_features = int(out_features)
+        self.quant_format = str(quant_format)
+        self.register_buffer("qdata", qdata.clone().contiguous(), persistent=True)
+        self.register_buffer(
+            "weight_scale",
+            weight_scale.clone().contiguous(),
+            persistent=True,
+        )
+        if weight_scale_2 is None:
+            self.weight_scale_2 = None
+        else:
+            self.register_buffer(
+                "weight_scale_2",
+                weight_scale_2.clone().contiguous(),
+                persistent=True,
+            )
+
+    def _weight_quantized_tensor(self) -> Any:
+        from comfy_kitchen.tensor import (
+            QuantizedTensor,
+            TensorCoreFP8Layout,
+            TensorCoreNVFP4Layout,
+        )
+
+        shape = (self.out_features, self.in_features)
+        if self.quant_format == "nvfp4":
+            if self.weight_scale_2 is None:
+                raise RuntimeError("NVFP4 projection is missing its second scale")
+            params = TensorCoreNVFP4Layout.Params(
+                scale=self.weight_scale_2,
+                orig_dtype=torch.bfloat16,
+                orig_shape=shape,
+                block_scale=self.weight_scale,
+            )
+            return QuantizedTensor(
+                self.qdata,
+                "TensorCoreNVFP4Layout",
+                params,
+            )
+        if self.quant_format == "float8_e4m3fn":
+            params = TensorCoreFP8Layout.Params(
+                scale=self.weight_scale,
+                orig_dtype=torch.bfloat16,
+                orig_shape=shape,
+            )
+            return QuantizedTensor(
+                self.qdata,
+                "TensorCoreFP8Layout",
+                params,
+            )
+        raise ValueError(f"unsupported quantized format: {self.quant_format}")
+
+    def forward(self, value: torch.Tensor) -> torch.Tensor:
+        from comfy_kitchen.tensor import QuantizedTensor
+
+        input_shape = tuple(value.shape)
+        flattened = value.reshape(-1, input_shape[-1]).contiguous()
+        layout = (
+            "TensorCoreNVFP4Layout"
+            if self.quant_format == "nvfp4"
+            else "TensorCoreFP8Layout"
+        )
+        input_quantized = QuantizedTensor.from_float(flattened, layout)
+        output = F.linear(
+            input_quantized,
+            self._weight_quantized_tensor(),
+            None,
+        )
+        return output.reshape(*input_shape[:-1], self.out_features)
+
+
+def _install_quantized_linears(
+    hf_module: torch.nn.Module,
+    artifact_path: Path,
+) -> dict[str, Any]:
+    format_counts: Counter[str] = Counter()
+    installed: list[str] = []
+    storage_bytes = 0
+    original_bf16_bytes = 0
+
+    with safe_open(str(artifact_path), framework="pt", device="cpu") as handle:
+        config_keys = sorted(
+            key for key in handle.keys() if key.endswith(".comfy_quant")
+        )
+        for config_key in config_keys:
+            layer_key = config_key.removesuffix(".comfy_quant")
+            quant_format = _decode_quant_config(
+                handle.get_tensor(config_key)
+            )["format"]
+            module_name = _artifact_layer_to_hf_module(layer_key)
+            parent, leaf = _resolve_parent(hf_module, module_name)
+            original = getattr(parent, leaf)
+            if not isinstance(original, torch.nn.Linear):
+                raise TypeError(
+                    f"{module_name}: expected torch.nn.Linear, got "
+                    f"{type(original).__name__}"
+                )
+            if original.bias is not None:
+                raise ValueError(f"{module_name}: quantized projection has a bias")
+
+            qdata = handle.get_tensor(f"{layer_key}.weight")
+            weight_scale = handle.get_tensor(f"{layer_key}.weight_scale")
+            weight_scale_2 = (
+                handle.get_tensor(f"{layer_key}.weight_scale_2")
+                if quant_format == "nvfp4"
+                else None
+            )
+            replacement = PublishedQuantLinear(
+                in_features=original.in_features,
+                out_features=original.out_features,
+                quant_format=quant_format,
+                qdata=qdata,
+                weight_scale=weight_scale,
+                weight_scale_2=weight_scale_2,
+            )
+            setattr(parent, leaf, replacement)
+            format_counts[quant_format] += 1
+            installed.append(module_name)
+            original_bf16_bytes += (
+                original.in_features * original.out_features * 2
+            )
+            storage_bytes += replacement.qdata.nbytes
+            storage_bytes += replacement.weight_scale.nbytes
+            if replacement.weight_scale_2 is not None:
+                storage_bytes += replacement.weight_scale_2.nbytes
+
+    summary = {
+        "installed_module_count": len(installed),
+        "format_counts": dict(format_counts),
+        "original_projection_bf16_bytes": original_bf16_bytes,
+        "packed_projection_bytes": storage_bytes,
+        "projection_saving_bytes": original_bf16_bytes - storage_bytes,
+        "projection_saving_gib": (
+            (original_bf16_bytes - storage_bytes) / float(1 << 30)
+        ),
+    }
+    if (
+        summary["installed_module_count"] != 238
+        or summary["format_counts"] != EXPECTED_FORMAT_COUNTS
+    ):
+        raise RuntimeError(f"unexpected text quantization policy: {summary}")
+    return summary
+
+
+def _load_nonquantized_weights(
+    hf_module: torch.nn.Module,
+    artifact_path: Path,
+) -> dict[str, Any]:
+    loaded: list[str] = []
+    unexpected: list[str] = []
+
+    with safe_open(str(artifact_path), framework="pt", device="cpu") as handle:
+        for artifact_key in sorted(handle.keys()):
+            if (
+                artifact_key.endswith(".comfy_quant")
+                or artifact_key.endswith(".weight_scale")
+                or artifact_key.endswith(".weight_scale_2")
+            ):
+                continue
+            target_name = _artifact_weight_to_hf_name(artifact_key)
+            if target_name is None:
+                unexpected.append(artifact_key)
+                continue
+            try:
+                parent, leaf = _resolve_parent(hf_module, target_name)
+            except AttributeError:
+                unexpected.append(artifact_key)
+                continue
+            current = getattr(parent, leaf, None)
+            if isinstance(current, PublishedQuantLinear):
+                continue
+            if target_name.endswith(".weight"):
+                projection_name = target_name.removesuffix(".weight")
+                try:
+                    projection_parent, projection_leaf = _resolve_parent(
+                        hf_module, projection_name
+                    )
+                    if isinstance(
+                        getattr(projection_parent, projection_leaf),
+                        PublishedQuantLinear,
+                    ):
+                        continue
+                except AttributeError:
+                    pass
+            set_module_tensor_to_device(
+                hf_module,
+                target_name,
+                "cpu",
+                value=handle.get_tensor(artifact_key),
+            )
+            loaded.append(target_name)
+
+    hf_module.tie_weights()
+    meta_parameters = [
+        name for name, value in hf_module.named_parameters() if value.is_meta
+    ]
+    meta_buffers = [
+        name for name, value in hf_module.named_buffers() if value.is_meta
+    ]
+    if meta_parameters or meta_buffers:
+        raise RuntimeError(
+            "packed text loader left unresolved meta tensors: "
+            f"{(meta_parameters + meta_buffers)[:4]}"
+        )
+    if unexpected:
+        raise RuntimeError(
+            f"packed text artifact contains unmapped tensors: {unexpected[:4]}"
+        )
+    return {
+        "loaded_nonquantized_tensor_count": len(loaded),
+        "unresolved_meta_parameters": meta_parameters,
+        "unresolved_meta_buffers": meta_buffers,
+    }
+
+
+def load_quantized_text_encoder(
+    *,
+    text_encoder_dir: str | Path,
+    artifact_path: str | Path,
+    tokenizer_max_length: int,
+    dit_structure: dict[str, Any],
+    use_packed_text_infer: bool,
+) -> tuple[torch.nn.Module, dict[str, Any]]:
+    """Construct Mage's text wrapper directly from the packaged quant artifact."""
+
+    from mage_flow.models.modules.text_encoder import (
+        CustomQwen3VLForConditionalGeneration,
+        TextEncoder,
+    )
+
+    text_encoder_dir = Path(text_encoder_dir).resolve()
+    artifact_path = Path(artifact_path).resolve()
+    if sha256(artifact_path) != EXPECTED_ARTIFACT_SHA256:
+        raise RuntimeError("packaged text-encoder artifact SHA-256 mismatch")
+
+    config = AutoConfig.from_pretrained(
+        str(text_encoder_dir),
+        local_files_only=True,
+    )
+    with init_empty_weights():
+        hf_module = CustomQwen3VLForConditionalGeneration._from_config(
+            config,
+            attn_implementation="flash_attention_2",
+            dtype=torch.bfloat16,
+        )
+
+    quant_summary = _install_quantized_linears(hf_module, artifact_path)
+    load_summary = _load_nonquantized_weights(hf_module, artifact_path)
+
+    text_encoder = TextEncoder.__new__(TextEncoder)
+    torch.nn.Module.__init__(text_encoder)
+    text_encoder.model_name = str(text_encoder_dir)
+    text_encoder.tokenizer_max_length = int(tokenizer_max_length)
+    text_encoder.tokenizer = AutoTokenizer.from_pretrained(
+        str(text_encoder_dir),
+        local_files_only=True,
+    )
+    text_encoder.tokenizer.padding_side = "right"
+    text_encoder.processor = AutoProcessor.from_pretrained(
+        str(text_encoder_dir),
+        local_files_only=True,
+    )
+    text_encoder.hf_module = hf_module.eval().requires_grad_(False)
+    text_encoder.prompt_template_encode = ""
+    text_encoder.prompt_template_encode_start_idx = 0
+    text_encoder.dit_structure = dict(dit_structure)
+    text_encoder.use_packed_text_infer = bool(use_packed_text_infer)
+    text_encoder.eval().requires_grad_(False)
+
+    return (
+        text_encoder,
+        {
+            "artifact": str(artifact_path),
+            "artifact_sha256": EXPECTED_ARTIFACT_SHA256,
+            "quantized": quant_summary,
+            "nonquantized": load_summary,
+        },
+    )

runtime/sm120_linear_op.cpp

@@ -0,0 +1,358 @@
+#include <ATen/ATen.h>
+#include <ATen/cuda/CUDAContext.h>
+#include <c10/core/DeviceType.h>
+#include <c10/cuda/CUDACachingAllocator.h>
+#include <c10/cuda/CUDAException.h>
+#include <c10/cuda/CUDAGuard.h>
+#include <c10/cuda/CUDAStream.h>
+#include <cuda_runtime_api.h>
+#include <torch/library.h>
+
+#include <cstdint>
+#include <limits>
+#include <memory>
+#include <mutex>
+#include <optional>
+#include <stdexcept>
+#include <unordered_map>
+#include <utility>
+#include <vector>
+
+#include "nvfp4_linear.h"
+
+namespace {
+
+constexpr int kAbiVersion = 1;
+constexpr int64_t kFp4BlockElements = 16;
+constexpr int64_t kScaleTileOuter = 128;
+constexpr int64_t kScaleTileInner = 4;
+
+int64_t round_up_int64(int64_t value, int64_t multiple) {
+  TORCH_CHECK(value > 0, "round_up requires a positive value");
+  TORCH_CHECK(multiple > 0, "round_up requires a positive multiple");
+  TORCH_CHECK(
+      value <= std::numeric_limits<int64_t>::max() - (multiple - 1),
+      "round_up overflow");
+  return ((value + multiple - 1) / multiple) * multiple;
+}
+
+int checked_int_arg(int64_t value, const char* label) {
+  TORCH_CHECK(value > 0, label, " must be positive");
+  TORCH_CHECK(
+      value <= static_cast<int64_t>(std::numeric_limits<int>::max()),
+      label,
+      " exceeds 32-bit ABI bound");
+  return static_cast<int>(value);
+}
+
+size_t expected_packed_weight_bytes(int out_features, int in_features) {
+  TORCH_CHECK(
+      out_features > 0 && out_features % 8 == 0,
+      "resident NVFP4 requires out_features divisible by 8, got ",
+      out_features);
+  TORCH_CHECK(
+      in_features > 0 && in_features % 32 == 0,
+      "resident NVFP4 requires in_features divisible by 32, got ",
+      in_features);
+  return (static_cast<size_t>(out_features) * static_cast<size_t>(in_features)) /
+      2;
+}
+
+size_t expected_scale_bytes(int out_features, int in_features) {
+  TORCH_CHECK(
+      in_features > 0 && in_features % kFp4BlockElements == 0,
+      "resident NVFP4 requires in_features divisible by 16, got ",
+      in_features);
+  TORCH_CHECK(out_features > 0, "out_features must be positive");
+  const int64_t inner_dim =
+      round_up_int64(in_features / kFp4BlockElements, kScaleTileInner);
+  const int64_t outer_tiles =
+      (out_features + kScaleTileOuter - 1) / kScaleTileOuter;
+  return static_cast<size_t>(inner_dim) * static_cast<size_t>(outer_tiles) *
+      static_cast<size_t>(kScaleTileOuter);
+}
+
+[[noreturn]] void throw_abi_error(const char* operation) {
+  const char* message = mage_nvfp4_last_error();
+  TORCH_CHECK(
+      false,
+      operation,
+      ": ",
+      message == nullptr ? "unknown native resident error" : message);
+}
+
+void ensure_abi_version() {
+  static std::once_flag once;
+  std::call_once(once, []() {
+    const int version = mage_nvfp4_abi_version();
+    TORCH_CHECK(
+        version == kAbiVersion,
+        "resident NVFP4 ABI mismatch: compiled torch op expects ",
+        kAbiVersion,
+        " but native library reports ",
+        version);
+  });
+}
+
+struct ContextKey {
+  int device_index = -1;
+  uintptr_t stream = 0;
+
+  bool operator==(const ContextKey& other) const noexcept {
+    return device_index == other.device_index && stream == other.stream;
+  }
+};
+
+struct ContextKeyHash {
+  size_t operator()(const ContextKey& key) const noexcept {
+    const size_t left = std::hash<int>{}(key.device_index);
+    const size_t right = std::hash<uintptr_t>{}(key.stream);
+    return left ^ (right + 0x9e3779b97f4a7c15ULL + (left << 6) + (left >> 2));
+  }
+};
+
+struct NativeContextDeleter {
+  void operator()(void* context) const noexcept {
+    if (context == nullptr) {
+      return;
+    }
+    (void)mage_nvfp4_destroy_context(context);
+  }
+};
+
+class ContextRegistry {
+ public:
+  void* get(int device_index, uintptr_t stream) {
+    std::lock_guard<std::mutex> guard(mutex_);
+    const ContextKey key{device_index, stream};
+    auto it = contexts_.find(key);
+    if (it != contexts_.end()) {
+      return it->second.get();
+    }
+    void* context = nullptr;
+    const int status = mage_nvfp4_create_context(device_index, &context);
+    if (status != 0) {
+      throw_abi_error("creating resident NVFP4 context");
+    }
+    auto inserted = contexts_.emplace(
+        key, std::unique_ptr<void, NativeContextDeleter>(context));
+    return inserted.first->second.get();
+  }
+
+  void clear() {
+    decltype(contexts_) retired;
+    {
+      std::lock_guard<std::mutex> guard(mutex_);
+      retired.swap(contexts_);
+    }
+    // Destruction synchronizes each bound stream through the resident ABI.
+    // Keep it outside the registry mutex.
+  }
+
+ private:
+  std::mutex mutex_;
+  std::unordered_map<
+      ContextKey,
+      std::unique_ptr<void, NativeContextDeleter>,
+      ContextKeyHash>
+      contexts_;
+};
+
+ContextRegistry& context_registry() {
+  // Deliberately avoid a static destructor that could call CUDA after runtime
+  // teardown. Long-lived processes must call the explicit close op; otherwise
+  // the OS reclaims these process resources at exit.
+  static ContextRegistry* registry = new ContextRegistry();
+  return *registry;
+}
+
+void clear_native_contexts() {
+  context_registry().clear();
+}
+
+void record_tensor_stream(
+    const at::Tensor& tensor, c10::cuda::CUDAStream stream) {
+  if (!tensor.defined() || !tensor.is_cuda()) {
+    return;
+  }
+  c10::cuda::CUDACachingAllocator::recordStream(
+      tensor.storage().data_ptr(), stream);
+}
+
+void validate_common(
+    const at::Tensor& input,
+    const at::Tensor& packed_weight,
+    const at::Tensor& weight_scales,
+    const at::Tensor& weight_scale,
+    const std::optional<at::Tensor>& bias,
+    int64_t in_features,
+    int64_t out_features) {
+  TORCH_CHECK(input.dim() >= 1, "resident NVFP4 input must have at least one dimension");
+  TORCH_CHECK(
+      input.scalar_type() == at::kBFloat16,
+      "resident NVFP4 input must be bfloat16");
+  TORCH_CHECK(!input.requires_grad(), "resident NVFP4 torch op is inference-only");
+  TORCH_CHECK(
+      input.size(-1) == in_features,
+      "resident NVFP4 expected last dimension ",
+      in_features,
+      " but got ",
+      input.size(-1));
+  TORCH_CHECK(
+      packed_weight.scalar_type() == at::kByte &&
+          packed_weight.dim() == 1 && packed_weight.is_contiguous(),
+      "packed_weight must be a contiguous 1D uint8 tensor");
+  TORCH_CHECK(
+      weight_scales.scalar_type() == at::kByte &&
+          weight_scales.dim() == 1 && weight_scales.is_contiguous(),
+      "weight_scales must be a contiguous 1D uint8 tensor");
+  TORCH_CHECK(
+      weight_scale.scalar_type() == at::kFloat &&
+          weight_scale.numel() == 1 && weight_scale.is_contiguous(),
+      "weight_scale must be one contiguous float32 value");
+  TORCH_CHECK(
+      packed_weight.device() == input.device(),
+      "packed_weight and input must be on the same device");
+  TORCH_CHECK(
+      weight_scales.device() == input.device(),
+      "weight_scales and input must be on the same device");
+  TORCH_CHECK(
+      weight_scale.device() == input.device(),
+      "weight_scale and input must be on the same device");
+  if (bias.has_value()) {
+    const at::Tensor& bias_value = *bias;
+    TORCH_CHECK(
+        bias_value.scalar_type() == at::kBFloat16 &&
+            bias_value.dim() == 1 && bias_value.is_contiguous(),
+        "bias must be a contiguous 1D bfloat16 tensor");
+    TORCH_CHECK(
+        bias_value.numel() == out_features,
+        "bias length must match out_features");
+    TORCH_CHECK(
+        bias_value.device() == input.device(),
+        "bias and input must be on the same device");
+  }
+
+  const int in_features_i = checked_int_arg(in_features, "in_features");
+  const int out_features_i = checked_int_arg(out_features, "out_features");
+  const size_t expected_weight = expected_packed_weight_bytes(
+      out_features_i, in_features_i);
+  const size_t expected_scales = expected_scale_bytes(
+      out_features_i, in_features_i);
+  TORCH_CHECK(
+      static_cast<size_t>(packed_weight.numel()) == expected_weight,
+      "packed_weight size mismatch: expected ",
+      expected_weight,
+      " bytes but got ",
+      packed_weight.numel());
+  TORCH_CHECK(
+      static_cast<size_t>(weight_scales.numel()) == expected_scales,
+      "weight_scales size mismatch: expected ",
+      expected_scales,
+      " bytes but got ",
+      weight_scales.numel());
+
+  const size_t native_weight =
+      mage_nvfp4_packed_weight_bytes(out_features_i, in_features_i);
+  const size_t native_scales =
+      mage_nvfp4_weight_scale_bytes(out_features_i, in_features_i);
+  TORCH_CHECK(
+      native_weight == expected_weight,
+      "native resident library packed-weight metadata disagrees with torch op");
+  TORCH_CHECK(
+      native_scales == expected_scales,
+      "native resident library scale metadata disagrees with torch op");
+}
+
+at::Tensor sm120_linear_native_cuda(
+    const at::Tensor& input,
+    const at::Tensor& packed_weight,
+    const at::Tensor& weight_scales,
+    const at::Tensor& weight_scale,
+    const std::optional<at::Tensor>& bias,
+    int64_t in_features,
+    int64_t out_features) {
+  ensure_abi_version();
+  validate_common(
+      input, packed_weight, weight_scales, weight_scale, bias, in_features,
+      out_features);
+
+  TORCH_CHECK(input.is_cuda(), "resident NVFP4 CUDA implementation requires a CUDA input");
+  const auto device = input.device();
+  c10::cuda::CUDAGuard guard(device);
+  const int device_index = device.index();
+  cudaDeviceProp properties{};
+  C10_CUDA_CHECK(cudaGetDeviceProperties(&properties, device_index));
+  TORCH_CHECK(
+      properties.major == 12 && properties.minor == 0,
+      "resident NVFP4 CUDA implementation is labeled sm_120-only; selected device reports ",
+      properties.major,
+      ".",
+      properties.minor);
+
+  const c10::cuda::CUDAStream stream =
+      c10::cuda::getCurrentCUDAStream(device_index);
+  TORCH_CHECK(
+      !stream.is_capturing(),
+      "resident NVFP4 C ABI torch op does not support CUDA graph capture");
+
+  const int in_features_i = checked_int_arg(in_features, "in_features");
+  const int out_features_i = checked_int_arg(out_features, "out_features");
+  at::Tensor contiguous =
+      input.reshape({-1, in_features_i}).contiguous();
+  const int64_t logical_m64 = contiguous.size(0);
+  TORCH_CHECK(logical_m64 > 0, "resident NVFP4 torch op does not support empty inputs");
+  const int logical_m = checked_int_arg(logical_m64, "logical_m");
+  const int64_t padded_m64 = round_up_int64(logical_m64, 8);
+  at::Tensor padded_output = at::empty(
+      {padded_m64, out_features_i},
+      input.options().dtype(at::kBFloat16));
+
+  record_tensor_stream(contiguous, stream);
+  record_tensor_stream(packed_weight, stream);
+  record_tensor_stream(weight_scales, stream);
+  record_tensor_stream(weight_scale, stream);
+  if (bias.has_value()) {
+    record_tensor_stream(*bias, stream);
+  }
+  record_tensor_stream(padded_output, stream);
+
+  void* context =
+      context_registry().get(device_index, reinterpret_cast<uintptr_t>(stream.stream()));
+  const void* bias_pointer =
+      bias.has_value() ? bias->data_ptr() : nullptr;
+  const int status = mage_nvfp4_linear_forward(
+      context,
+      contiguous.data_ptr(),
+      packed_weight.data_ptr(),
+      static_cast<size_t>(packed_weight.numel()),
+      weight_scales.data_ptr(),
+      static_cast<size_t>(weight_scales.numel()),
+      weight_scale.data_ptr(),
+      bias_pointer,
+      padded_output.data_ptr(),
+      logical_m,
+      in_features_i,
+      out_features_i,
+      reinterpret_cast<uintptr_t>(stream.stream()));
+  if (status != 0) {
+    throw_abi_error("running resident NVFP4 linear forward");
+  }
+
+  at::Tensor logical = padded_output.narrow(0, 0, logical_m64);
+  std::vector<int64_t> output_sizes = input.sizes().vec();
+  output_sizes.back() = out_features_i;
+  return logical.view(output_sizes);
+}
+
+} // namespace
+
+TORCH_LIBRARY_FRAGMENT(mage_nvfp4, m) {
+  m.def(
+      "sm120_linear_native(Tensor input, Tensor packed_weight, Tensor weight_scales, Tensor weight_scale, Tensor? bias, int in_features, int out_features) -> Tensor");
+  m.def("clear_native_contexts() -> ()", TORCH_FN(clear_native_contexts));
+}
+
+TORCH_LIBRARY_IMPL(mage_nvfp4, CUDA, m) {
+  m.impl("sm120_linear_native", TORCH_FN(sm120_linear_native_cuda));
+}

runtime/standard_transformer.py

@@ -0,0 +1,214 @@
+"""Load the complete sharded Mage-Flow NVFP4 transformer component."""
+
+from __future__ import annotations
+
+from contextlib import ExitStack
+import json
+from pathlib import Path
+from typing import Any
+
+import torch
+import torch.nn as nn
+from safetensors import safe_open
+
+from packed_artifact import (
+    assign_tensor_by_name,
+    build_target_specs,
+    instantiate_mage_transformer_on_meta,
+    materialize_mage_rope_tensor_attributes,
+    set_child_module,
+    unregistered_meta_tensor_attribute_names,
+)
+from torch_ops_native import (
+    PackedNvfp4LinearNativeOp,
+    initialize_native_sm120_op,
+)
+
+
+class StandardCheckpointError(RuntimeError):
+    pass
+
+
+def fail(message: str) -> None:
+    raise StandardCheckpointError(message)
+
+
+def read_object(path: Path) -> dict[str, Any]:
+    try:
+        value = json.loads(path.read_text(encoding="utf-8"))
+    except (OSError, json.JSONDecodeError) as exc:
+        fail(f"cannot read JSON object {path}: {exc}")
+    if not isinstance(value, dict):
+        fail(f"expected a JSON object: {path}")
+    return value
+
+
+def _component_path(component_dir: Path, relative: str) -> Path:
+    path = (component_dir / relative).resolve()
+    if not path.is_relative_to(component_dir):
+        fail(f"checkpoint index path escapes transformer component: {relative}")
+    if not path.is_file():
+        fail(f"checkpoint shard is missing: {relative}")
+    return path
+
+
+def _quantized_keys(module_key: str) -> dict[str, str]:
+    return {
+        "packed_weight": f"{module_key}.packed_weight",
+        "weight_scales": f"{module_key}.weight_scales",
+        "weight_scale": f"{module_key}.weight_scale",
+        "bias": f"{module_key}.bias",
+    }
+
+
+def load_standard_native_transformer(
+    repo_root: str | Path,
+    device: torch.device,
+) -> tuple[nn.Module, dict[str, Any]]:
+    """Load a complete standard-layout component without a BF16 base download."""
+
+    repo_root = Path(repo_root).resolve()
+    component_dir = (repo_root / "transformer").resolve()
+    if not component_dir.is_relative_to(repo_root) or not component_dir.is_dir():
+        fail("repository has no transformer component")
+    if device.type != "cuda":
+        fail("the native resident transformer requires a CUDA destination")
+    if not initialize_native_sm120_op(allow_python_schema_fallback=False):
+        fail("compiled native SM120 torch op did not load")
+
+    config = read_object(component_dir / "config.json")
+    quant_config = config.get("quantization_config")
+    if not isinstance(quant_config, dict):
+        fail("transformer config has no quantization_config")
+    if quant_config.get("quant_method") != "mage_flow_nvfp4":
+        fail(
+            "unexpected transformer quantization method: "
+            f"{quant_config.get('quant_method')!r}"
+        )
+    if quant_config.get("quant_algo") != "NVFP4":
+        fail("transformer config does not declare NVFP4")
+
+    depth = int(config.get("depth", 0))
+    target_specs = build_target_specs(depth)
+    expected_targets = [spec.module_key for spec in target_specs]
+    if quant_config.get("targets") != expected_targets:
+        fail("transformer quantization target list is not canonical")
+
+    metadata = read_object(component_dir / "nvfp4_metadata.json")
+    if metadata.get("artifact_kind") != (
+        "mage_flow_transformer_mlp_nvfp4_resident_v1"
+    ):
+        fail("unexpected transformer NVFP4 metadata kind")
+    non_target_keys = metadata.get("non_target_keys")
+    if not isinstance(non_target_keys, list) or not all(
+        isinstance(key, str) for key in non_target_keys
+    ):
+        fail("transformer NVFP4 metadata has no non-target key list")
+
+    index = read_object(
+        component_dir / "diffusion_pytorch_model.safetensors.index.json"
+    )
+    weight_map = index.get("weight_map")
+    if not isinstance(weight_map, dict) or not all(
+        isinstance(key, str) and isinstance(value, str)
+        for key, value in weight_map.items()
+    ):
+        fail("transformer checkpoint has no valid weight map")
+
+    quantized_keys = {
+        key
+        for spec in target_specs
+        for key in _quantized_keys(spec.module_key).values()
+    }
+    expected_keys = set(non_target_keys) | quantized_keys
+    actual_keys = set(weight_map)
+    if actual_keys != expected_keys:
+        missing = sorted(expected_keys - actual_keys)
+        unexpected = sorted(actual_keys - expected_keys)
+        fail(
+            "transformer checkpoint key coverage mismatch; "
+            f"missing={missing[:1]}, unexpected={unexpected[:1]}"
+        )
+    original_target_weights = {spec.weight_key for spec in target_specs}
+    leaked = sorted(actual_keys & original_target_weights)
+    if leaked:
+        fail(f"BF16 target weight leaked into quantized checkpoint: {leaked[0]}")
+
+    shard_names = sorted(set(weight_map.values()))
+    with ExitStack() as stack:
+        handles = {
+            name: stack.enter_context(
+                safe_open(
+                    _component_path(component_dir, name),
+                    framework="pt",
+                    device="cpu",
+                )
+            )
+            for name in shard_names
+        }
+
+        def tensor(key: str) -> torch.Tensor:
+            try:
+                handle = handles[weight_map[key]]
+            except KeyError:
+                fail(f"tensor is absent from checkpoint index: {key}")
+            if key not in handle.keys():
+                fail(f"tensor is absent from its declared shard: {key}")
+            return handle.get_tensor(key)
+
+        model = instantiate_mage_transformer_on_meta(repo_root)
+        for spec in target_specs:
+            original = model.get_submodule(spec.module_key)
+            if not isinstance(original, nn.Linear):
+                fail(
+                    f"expected target {spec.module_key} to be nn.Linear, "
+                    f"found {type(original).__name__}"
+                )
+            keys = _quantized_keys(spec.module_key)
+            replacement = PackedNvfp4LinearNativeOp(
+                in_features=int(original.in_features),
+                out_features=int(original.out_features),
+                packed_weight=tensor(keys["packed_weight"]).to(device),
+                weight_scales=tensor(keys["weight_scales"]).to(device),
+                weight_scale=tensor(keys["weight_scale"]).to(device),
+                bias=tensor(keys["bias"]).to(device),
+            )
+            set_child_module(model, spec.module_key, replacement)
+
+        loaded_non_targets: list[str] = []
+        for key in non_target_keys:
+            assign_tensor_by_name(model, key, tensor(key).to(device))
+            loaded_non_targets.append(key)
+
+    materialized = materialize_mage_rope_tensor_attributes(model)
+    meta_parameters = [
+        name for name, value in model.named_parameters() if value.is_meta
+    ]
+    meta_buffers = [
+        name for name, value in model.named_buffers() if value.is_meta
+    ]
+    unregistered_meta = unregistered_meta_tensor_attribute_names(model)
+    if meta_parameters or meta_buffers or unregistered_meta:
+        fail(
+            "standard transformer loader left unresolved meta tensors: "
+            f"{(meta_parameters + meta_buffers + unregistered_meta)[:4]}"
+        )
+
+    report = {
+        "layout": "huggingface_sharded_component",
+        "checkpoint_shard_count": len(shard_names),
+        "checkpoint_tensor_count": len(actual_keys),
+        "loaded_non_target_tensor_count": len(loaded_non_targets),
+        "loaded_quantized_projection_count": len(target_specs),
+        "bf16_target_weight_reads": 0,
+        "meta_parameter_names": meta_parameters,
+        "meta_buffer_names": meta_buffers,
+        "materialized_unregistered_tensor_attribute_names": materialized,
+    }
+    return model.eval().requires_grad_(False), report
+
+
+__all__ = [
+    "StandardCheckpointError",
+    "load_standard_native_transformer",
+]

runtime/torch_ops_native.py

@@ -0,0 +1,233 @@
+"""Loader and fake/meta shim for the compiled resident NVFP4 torch op.
+
+The compiled bridge lives in ``native/torch_op`` and, when built, moves CUDA
+execution out of Python+ctypes and into a C++ dispatcher implementation that
+calls the existing resident C ABI directly.
+
+This shim keeps CPU/meta tests lightweight:
+
+- if the compiled bridge exists, load it first;
+- otherwise define a temporary Python schema fallback for shape-only tests;
+- always register fake/meta implementations in Python.
+"""
+
+from __future__ import annotations
+
+import atexit
+from pathlib import Path
+
+import torch
+
+from packed_nvfp4_linear import PackedNvfp4Linear, scale_layout
+
+
+RELEASE_ROOT = Path(__file__).resolve().parents[1]
+DEFAULT_EXTENSION_PATH = RELEASE_ROOT / "runtime" / "libmage_nvfp4_torch_op.so"
+
+OP_NAMESPACE = "mage_nvfp4"
+OP_NAME = "sm120_linear_native"
+OP_QUALNAME = f"{OP_NAMESPACE}::{OP_NAME}"
+
+_SCHEMA_FALLBACK_LIB: torch.library.Library | None = None
+_META_LIB: torch.library.Library | None = None
+_EXTENSION_LOADED = False
+_FAKE_REGISTERED = False
+_META_REGISTERED = False
+
+
+def _expected_weight_bytes(out_features: int, in_features: int) -> int:
+    if out_features <= 0 or out_features % 8:
+        raise ValueError("resident NVFP4 requires out_features divisible by 8")
+    if in_features <= 0 or in_features % 32:
+        raise ValueError("resident NVFP4 requires in_features divisible by 32")
+    return (out_features * in_features) // 2
+
+
+def _check_common_shapes(
+    input: torch.Tensor,
+    packed_weight: torch.Tensor,
+    weight_scales: torch.Tensor,
+    weight_scale: torch.Tensor,
+    bias: torch.Tensor | None,
+    in_features: int,
+    out_features: int,
+) -> None:
+    if input.ndim < 1:
+        raise ValueError("resident NVFP4 input must have at least one dimension")
+    if input.dtype != torch.bfloat16:
+        raise TypeError("resident NVFP4 input must be bfloat16")
+    if input.requires_grad:
+        raise RuntimeError("resident NVFP4 torch op is inference-only")
+    if int(input.shape[-1]) != int(in_features):
+        raise ValueError(
+            f"expected input last dimension {in_features}, got {tuple(input.shape)}"
+        )
+    if packed_weight.dtype != torch.uint8 or packed_weight.ndim != 1:
+        raise ValueError("packed_weight must be a 1D uint8 tensor")
+    if weight_scales.dtype != torch.uint8 or weight_scales.ndim != 1:
+        raise ValueError("weight_scales must be a 1D uint8 tensor")
+    if weight_scale.dtype != torch.float32 or weight_scale.numel() != 1:
+        raise ValueError("weight_scale must be a scalar or length-1 float32 tensor")
+    if bias is not None and (
+        bias.dtype != torch.bfloat16 or bias.ndim != 1 or int(bias.numel()) != int(out_features)
+    ):
+        raise ValueError("bias must be a 1D bfloat16 tensor with length out_features")
+    for tensor in (packed_weight, weight_scales, weight_scale, bias):
+        if tensor is not None and tensor.device != input.device:
+            raise ValueError("input and resident buffers must be on the same device")
+        if tensor is not None and not tensor.is_contiguous():
+            raise ValueError("resident packed buffers must be contiguous")
+
+    expected_weight_bytes = _expected_weight_bytes(int(out_features), int(in_features))
+    expected_scale_bytes = scale_layout(int(in_features), int(out_features)).num_bytes
+    if int(packed_weight.numel()) != expected_weight_bytes:
+        raise ValueError(
+            f"packed_weight size mismatch: expected {expected_weight_bytes}, got {packed_weight.numel()}"
+        )
+    if int(weight_scales.numel()) != expected_scale_bytes:
+        raise ValueError(
+            f"weight_scales size mismatch: expected {expected_scale_bytes}, got {weight_scales.numel()}"
+        )
+
+
+def _logical_output(input: torch.Tensor, out_features: int) -> torch.Tensor:
+    return input.new_empty((*input.shape[:-1], int(out_features)), dtype=torch.bfloat16)
+
+
+def ensure_native_sm120_schema(
+    *,
+    extension_path: str | Path = DEFAULT_EXTENSION_PATH,
+    allow_python_schema_fallback: bool = True,
+) -> bool:
+    global _EXTENSION_LOADED, _SCHEMA_FALLBACK_LIB
+    extension_path = Path(extension_path)
+    if not _EXTENSION_LOADED and extension_path.is_file():
+        torch.ops.load_library(str(extension_path.resolve()))
+        _EXTENSION_LOADED = True
+        return True
+    if _EXTENSION_LOADED:
+        return True
+    if not allow_python_schema_fallback:
+        return False
+    if _SCHEMA_FALLBACK_LIB is None:
+        lib = torch.library.Library(OP_NAMESPACE, "FRAGMENT")
+        lib.define(
+            "sm120_linear_native(Tensor input, Tensor packed_weight, Tensor weight_scales, Tensor weight_scale, Tensor? bias, int in_features, int out_features) -> Tensor"
+        )
+        lib.define("clear_native_contexts() -> ()")
+        _SCHEMA_FALLBACK_LIB = lib
+    return False
+
+
+def _register_meta_impl() -> None:
+    global _META_LIB, _META_REGISTERED
+    if _META_REGISTERED:
+        return
+    _META_LIB = torch.library.Library(OP_NAMESPACE, "IMPL", "Meta")
+    _META_LIB.impl(
+        OP_NAME,
+        lambda input, packed_weight, weight_scales, weight_scale, bias, in_features, out_features: (
+            _check_common_shapes(
+                input,
+                packed_weight,
+                weight_scales,
+                weight_scale,
+                bias,
+                in_features,
+                out_features,
+            ),
+            _logical_output(input, int(out_features)),
+        )[1],
+    )
+    _META_REGISTERED = True
+
+
+def _register_fake_impl() -> None:
+    global _FAKE_REGISTERED
+    if _FAKE_REGISTERED:
+        return
+
+    @torch.library.register_fake(OP_QUALNAME)
+    def _fake(
+        input: torch.Tensor,
+        packed_weight: torch.Tensor,
+        weight_scales: torch.Tensor,
+        weight_scale: torch.Tensor,
+        bias: torch.Tensor | None,
+        in_features: int,
+        out_features: int,
+    ) -> torch.Tensor:
+        _check_common_shapes(
+            input,
+            packed_weight,
+            weight_scales,
+            weight_scale,
+            bias,
+            in_features,
+            out_features,
+        )
+        return _logical_output(input, int(out_features))
+
+    _FAKE_REGISTERED = True
+
+
+def initialize_native_sm120_op(
+    *,
+    extension_path: str | Path = DEFAULT_EXTENSION_PATH,
+    allow_python_schema_fallback: bool = True,
+) -> bool:
+    loaded = ensure_native_sm120_schema(
+        extension_path=extension_path,
+        allow_python_schema_fallback=allow_python_schema_fallback,
+    )
+    _register_meta_impl()
+    _register_fake_impl()
+    return loaded
+
+
+def close_native_contexts() -> None:
+    """Synchronize and release native contexts when the compiled bridge is loaded."""
+    if _EXTENSION_LOADED:
+        torch.ops.mage_nvfp4.clear_native_contexts()
+
+
+def _quiet_atexit_close() -> None:
+    try:
+        close_native_contexts()
+    except Exception:
+        # CUDA may already be shutting down. Explicit close_native_contexts()
+        # is the auditable path; atexit is only a best-effort fallback.
+        pass
+
+
+class PackedNvfp4LinearNativeOp(PackedNvfp4Linear):
+    """Thin wrapper over resident packed buffers using the native torch op."""
+
+    def forward(self, input: torch.Tensor) -> torch.Tensor:
+        if input.device.type not in {"cuda", "meta"}:
+            raise ValueError("PackedNvfp4LinearNativeOp requires a CUDA or meta input")
+        return torch.ops.mage_nvfp4.sm120_linear_native(
+            input,
+            self.packed_weight,
+            self.weight_scales,
+            self.weight_scale,
+            self.bias,
+            self.in_features,
+            self.out_features,
+        )
+
+
+initialize_native_sm120_op()
+atexit.register(_quiet_atexit_close)
+
+
+__all__ = [
+    "DEFAULT_EXTENSION_PATH",
+    "OP_NAME",
+    "OP_NAMESPACE",
+    "OP_QUALNAME",
+    "PackedNvfp4LinearNativeOp",
+    "close_native_contexts",
+    "ensure_native_sm120_schema",
+    "initialize_native_sm120_op",
+]

scheduler/scheduler_config.json

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

text_encoder/chat_template.json

@@ -0,0 +1,4 @@
+{
+    "chat_template": "{%- if tools %}\n    {{- '<|im_start|>system\\n' }}\n    {%- if messages[0].role == 'system' %}\n        {%- if messages[0].content is string %}\n            {{- messages[0].content }}\n        {%- else %}\n            {%- for content in messages[0].content %}\n                {%- if 'text' in content %}\n                    {{- content.text }}\n                {%- endif %}\n            {%- endfor %}\n        {%- endif %}\n        {{- '\\n\\n' }}\n    {%- endif %}\n    {{- \"# Tools\\n\\nYou may call one or more functions to assist with the user query.\\n\\nYou are provided with function signatures within <tools></tools> XML tags:\\n<tools>\" }}\n    {%- for tool in tools %}\n        {{- \"\\n\" }}\n        {{- tool | tojson }}\n    {%- endfor %}\n    {{- \"\\n</tools>\\n\\nFor each function call, return a json object with function name and arguments within <tool_call></tool_call> XML tags:\\n<tool_call>\\n{\\\"name\\\": <function-name>, \\\"arguments\\\": <args-json-object>}\\n</tool_call><|im_end|>\\n\" }}\n{%- else %}\n    {%- if messages[0].role == 'system' %}\n        {{- '<|im_start|>system\\n' }}\n        {%- if messages[0].content is string %}\n            {{- messages[0].content }}\n        {%- else %}\n            {%- for content in messages[0].content %}\n                {%- if 'text' in content %}\n                    {{- content.text }}\n                {%- endif %}\n            {%- endfor %}\n        {%- endif %}\n        {{- '<|im_end|>\\n' }}\n    {%- endif %}\n{%- endif %}\n{%- set image_count = namespace(value=0) %}\n{%- set video_count = namespace(value=0) %}\n{%- for message in messages %}\n    {%- if message.role == \"user\" %}\n        {{- '<|im_start|>' + message.role + '\\n' }}\n        {%- if message.content is string %}\n            {{- message.content }}\n        {%- else %}\n            {%- for content in message.content %}\n                {%- if content.type == 'image' or 'image' in content or 'image_url' in content %}\n                    {%- set image_count.value = image_count.value + 1 %}\n                    {%- if add_vision_id %}Picture {{ image_count.value }}: {% endif -%}\n                    <|vision_start|><|image_pad|><|vision_end|>\n                {%- elif content.type == 'video' or 'video' in content %}\n                    {%- set video_count.value = video_count.value + 1 %}\n                    {%- if add_vision_id %}Video {{ video_count.value }}: {% endif -%}\n                    <|vision_start|><|video_pad|><|vision_end|>\n                {%- elif 'text' in content %}\n                    {{- content.text }}\n                {%- endif %}\n            {%- endfor %}\n        {%- endif %}\n        {{- '<|im_end|>\\n' }}\n    {%- elif message.role == \"assistant\" %}\n        {{- '<|im_start|>' + message.role + '\\n' }}\n        {%- if message.content is string %}\n            {{- message.content }}\n        {%- else %}\n            {%- for content_item in message.content %}\n                {%- if 'text' in content_item %}\n                    {{- content_item.text }}\n                {%- endif %}\n            {%- endfor %}\n        {%- endif %}\n        {%- if message.tool_calls %}\n            {%- for tool_call in message.tool_calls %}\n                {%- if (loop.first and message.content) or (not loop.first) %}\n                    {{- '\\n' }}\n                {%- endif %}\n                {%- if tool_call.function %}\n                    {%- set tool_call = tool_call.function %}\n                {%- endif %}\n                {{- '<tool_call>\\n{\"name\": \"' }}\n                {{- tool_call.name }}\n                {{- '\", \"arguments\": ' }}\n                {%- if tool_call.arguments is string %}\n                    {{- tool_call.arguments }}\n                {%- else %}\n                    {{- tool_call.arguments | tojson }}\n                {%- endif %}\n                {{- '}\\n</tool_call>' }}\n            {%- endfor %}\n        {%- endif %}\n        {{- '<|im_end|>\\n' }}\n    {%- elif message.role == \"tool\" %}\n        {%- if loop.first or (messages[loop.index0 - 1].role != \"tool\") %}\n            {{- '<|im_start|>user' }}\n        {%- endif %}\n        {{- '\\n<tool_response>\\n' }}\n        {%- if message.content is string %}\n            {{- message.content }}\n        {%- else %}\n            {%- for content in message.content %}\n                {%- if content.type == 'image' or 'image' in content or 'image_url' in content %}\n                    {%- set image_count.value = image_count.value + 1 %}\n                    {%- if add_vision_id %}Picture {{ image_count.value }}: {% endif -%}\n                    <|vision_start|><|image_pad|><|vision_end|>\n                {%- elif content.type == 'video' or 'video' in content %}\n                    {%- set video_count.value = video_count.value + 1 %}\n                    {%- if add_vision_id %}Video {{ video_count.value }}: {% endif -%}\n                    <|vision_start|><|video_pad|><|vision_end|>\n                {%- elif 'text' in content %}\n                    {{- content.text }}\n                {%- endif %}\n            {%- endfor %}\n        {%- endif %}\n        {{- '\\n</tool_response>' }}\n        {%- if loop.last or (messages[loop.index0 + 1].role != \"tool\") %}\n            {{- '<|im_end|>\\n' }}\n        {%- endif %}\n    {%- endif %}\n{%- endfor %}\n{%- if add_generation_prompt %}\n    {{- '<|im_start|>assistant\\n' }}\n{%- endif %}\n"
+}
+  

text_encoder/config.json

@@ -0,0 +1,71 @@
+{
+  "architectures": [
+    "Qwen3VLForConditionalGeneration"
+  ],
+  "image_token_id": 151655,
+  "model_type": "qwen3_vl",
+  "quantization_config": {
+    "bf16_policy": "blocks 0/35, embeddings, norms, biases, and vision tower",
+    "compute_dtype": "bfloat16",
+    "fp8_projection_count": 14,
+    "nvfp4_projection_count": 224,
+    "quant_algo": "NVFP4_FP8_MIXED",
+    "quant_method": "comfy_kitchen"
+  },
+  "text_config": {
+    "attention_bias": false,
+    "attention_dropout": 0.0,
+    "bos_token_id": 151643,
+    "dtype": "bfloat16",
+    "eos_token_id": 151645,
+    "head_dim": 128,
+    "hidden_act": "silu",
+    "hidden_size": 2560,
+    "initializer_range": 0.02,
+    "intermediate_size": 9728,
+    "max_position_embeddings": 262144,
+    "model_type": "qwen3_vl_text",
+    "num_attention_heads": 32,
+    "num_hidden_layers": 36,
+    "num_key_value_heads": 8,
+    "rms_norm_eps": 1e-06,
+    "rope_scaling": {
+      "mrope_interleaved": true,
+      "mrope_section": [
+        24,
+        20,
+        20
+      ],
+      "rope_type": "default"
+    },
+    "rope_theta": 5000000,
+    "tie_word_embeddings": true,
+    "use_cache": true,
+    "vocab_size": 151936
+  },
+  "tie_word_embeddings": true,
+  "transformers_version": "4.57.0.dev0",
+  "video_token_id": 151656,
+  "vision_config": {
+    "deepstack_visual_indexes": [
+      5,
+      11,
+      17
+    ],
+    "depth": 24,
+    "hidden_act": "gelu_pytorch_tanh",
+    "hidden_size": 1024,
+    "in_channels": 3,
+    "initializer_range": 0.02,
+    "intermediate_size": 4096,
+    "model_type": "qwen3_vl",
+    "num_heads": 16,
+    "num_position_embeddings": 2304,
+    "out_hidden_size": 2560,
+    "patch_size": 16,
+    "spatial_merge_size": 2,
+    "temporal_patch_size": 2
+  },
+  "vision_end_token_id": 151653,
+  "vision_start_token_id": 151652
+}

text_encoder/generation_config.json

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

text_encoder/preprocessor_config.json

@@ -0,0 +1,21 @@
+{
+    "size": {
+        "longest_edge": 16777216,
+        "shortest_edge": 65536
+    },
+    "patch_size": 16,
+    "temporal_patch_size": 2,
+    "merge_size": 2,
+    "image_mean": [
+        0.5,
+        0.5,
+        0.5
+    ],
+    "image_std": [
+        0.5,
+        0.5,
+        0.5
+    ],
+    "processor_class": "Qwen3VLProcessor",
+    "image_processor_type": "Qwen2VLImageProcessorFast"
+}

text_encoder/tokenizer_config.json

@@ -0,0 +1,239 @@
+{
+  "add_bos_token": false,
+  "add_prefix_space": false,
+  "added_tokens_decoder": {
+    "151643": {
+      "content": "<|endoftext|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151644": {
+      "content": "<|im_start|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151645": {
+      "content": "<|im_end|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151646": {
+      "content": "<|object_ref_start|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151647": {
+      "content": "<|object_ref_end|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151648": {
+      "content": "<|box_start|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151649": {
+      "content": "<|box_end|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151650": {
+      "content": "<|quad_start|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151651": {
+      "content": "<|quad_end|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151652": {
+      "content": "<|vision_start|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151653": {
+      "content": "<|vision_end|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151654": {
+      "content": "<|vision_pad|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151655": {
+      "content": "<|image_pad|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151656": {
+      "content": "<|video_pad|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151657": {
+      "content": "<tool_call>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151658": {
+      "content": "</tool_call>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151659": {
+      "content": "<|fim_prefix|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151660": {
+      "content": "<|fim_middle|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151661": {
+      "content": "<|fim_suffix|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151662": {
+      "content": "<|fim_pad|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151663": {
+      "content": "<|repo_name|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151664": {
+      "content": "<|file_sep|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151665": {
+      "content": "<tool_response>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151666": {
+      "content": "</tool_response>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151667": {
+      "content": "<think>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151668": {
+      "content": "</think>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    }
+  },
+  "additional_special_tokens": [
+    "<|im_start|>",
+    "<|im_end|>",
+    "<|object_ref_start|>",
+    "<|object_ref_end|>",
+    "<|box_start|>",
+    "<|box_end|>",
+    "<|quad_start|>",
+    "<|quad_end|>",
+    "<|vision_start|>",
+    "<|vision_end|>",
+    "<|vision_pad|>",
+    "<|image_pad|>",
+    "<|video_pad|>"
+  ],
+  "bos_token": null,
+  "chat_template": "{%- if tools %}\n    {{- '<|im_start|>system\\n' }}\n    {%- if messages[0].role == 'system' %}\n        {%- if messages[0].content is string %}\n            {{- messages[0].content }}\n        {%- else %}\n            {%- for content in messages[0].content %}\n                {%- if 'text' in content %}\n                    {{- content.text }}\n                {%- endif %}\n            {%- endfor %}\n        {%- endif %}\n        {{- '\\n\\n' }}\n    {%- endif %}\n    {{- \"# Tools\\n\\nYou may call one or more functions to assist with the user query.\\n\\nYou are provided with function signatures within <tools></tools> XML tags:\\n<tools>\" }}\n    {%- for tool in tools %}\n        {{- \"\\n\" }}\n        {{- tool | tojson }}\n    {%- endfor %}\n    {{- \"\\n</tools>\\n\\nFor each function call, return a json object with function name and arguments within <tool_call></tool_call> XML tags:\\n<tool_call>\\n{\\\"name\\\": <function-name>, \\\"arguments\\\": <args-json-object>}\\n</tool_call><|im_end|>\\n\" }}\n{%- else %}\n    {%- if messages[0].role == 'system' %}\n        {{- '<|im_start|>system\\n' }}\n        {%- if messages[0].content is string %}\n            {{- messages[0].content }}\n        {%- else %}\n            {%- for content in messages[0].content %}\n                {%- if 'text' in content %}\n                    {{- content.text }}\n                {%- endif %}\n            {%- endfor %}\n        {%- endif %}\n        {{- '<|im_end|>\\n' }}\n    {%- endif %}\n{%- endif %}\n{%- set image_count = namespace(value=0) %}\n{%- set video_count = namespace(value=0) %}\n{%- for message in messages %}\n    {%- if message.role == \"user\" %}\n        {{- '<|im_start|>' + message.role + '\\n' }}\n        {%- if message.content is string %}\n            {{- message.content }}\n        {%- else %}\n            {%- for content in message.content %}\n                {%- if content.type == 'image' or 'image' in content or 'image_url' in content %}\n                    {%- set image_count.value = image_count.value + 1 %}\n                    {%- if add_vision_id %}Picture {{ image_count.value }}: {% endif -%}\n                    <|vision_start|><|image_pad|><|vision_end|>\n                {%- elif content.type == 'video' or 'video' in content %}\n                    {%- set video_count.value = video_count.value + 1 %}\n                    {%- if add_vision_id %}Video {{ video_count.value }}: {% endif -%}\n                    <|vision_start|><|video_pad|><|vision_end|>\n                {%- elif 'text' in content %}\n                    {{- content.text }}\n                {%- endif %}\n            {%- endfor %}\n        {%- endif %}\n        {{- '<|im_end|>\\n' }}\n    {%- elif message.role == \"assistant\" %}\n        {{- '<|im_start|>' + message.role + '\\n' }}\n        {%- if message.content is string %}\n            {{- message.content }}\n        {%- else %}\n            {%- for content_item in message.content %}\n                {%- if 'text' in content_item %}\n                    {{- content_item.text }}\n                {%- endif %}\n            {%- endfor %}\n        {%- endif %}\n        {%- if message.tool_calls %}\n            {%- for tool_call in message.tool_calls %}\n                {%- if (loop.first and message.content) or (not loop.first) %}\n                    {{- '\\n' }}\n                {%- endif %}\n                {%- if tool_call.function %}\n                    {%- set tool_call = tool_call.function %}\n                {%- endif %}\n                {{- '<tool_call>\\n{\"name\": \"' }}\n                {{- tool_call.name }}\n                {{- '\", \"arguments\": ' }}\n                {%- if tool_call.arguments is string %}\n                    {{- tool_call.arguments }}\n                {%- else %}\n                    {{- tool_call.arguments | tojson }}\n                {%- endif %}\n                {{- '}\\n</tool_call>' }}\n            {%- endfor %}\n        {%- endif %}\n        {{- '<|im_end|>\\n' }}\n    {%- elif message.role == \"tool\" %}\n        {%- if loop.first or (messages[loop.index0 - 1].role != \"tool\") %}\n            {{- '<|im_start|>user' }}\n        {%- endif %}\n        {{- '\\n<tool_response>\\n' }}\n        {%- if message.content is string %}\n            {{- message.content }}\n        {%- else %}\n            {%- for content in message.content %}\n                {%- if content.type == 'image' or 'image' in content or 'image_url' in content %}\n                    {%- set image_count.value = image_count.value + 1 %}\n                    {%- if add_vision_id %}Picture {{ image_count.value }}: {% endif -%}\n                    <|vision_start|><|image_pad|><|vision_end|>\n                {%- elif content.type == 'video' or 'video' in content %}\n                    {%- set video_count.value = video_count.value + 1 %}\n                    {%- if add_vision_id %}Video {{ video_count.value }}: {% endif -%}\n                    <|vision_start|><|video_pad|><|vision_end|>\n                {%- elif 'text' in content %}\n                    {{- content.text }}\n                {%- endif %}\n            {%- endfor %}\n        {%- endif %}\n        {{- '\\n</tool_response>' }}\n        {%- if loop.last or (messages[loop.index0 + 1].role != \"tool\") %}\n            {{- '<|im_end|>\\n' }}\n        {%- endif %}\n    {%- endif %}\n{%- endfor %}\n{%- if add_generation_prompt %}\n    {{- '<|im_start|>assistant\\n' }}\n{%- endif %}\n",
+  "clean_up_tokenization_spaces": false,
+  "eos_token": "<|im_end|>",
+  "errors": "replace",
+  "model_max_length": 262144,
+  "pad_token": "<|endoftext|>",
+  "split_special_tokens": false,
+  "tokenizer_class": "Qwen2Tokenizer",
+  "unk_token": null
+}

text_encoder/video_preprocessor_config.json

@@ -0,0 +1,21 @@
+{
+    "size": {
+        "longest_edge": 25165824,
+        "shortest_edge": 4096
+    },
+    "patch_size": 16,
+    "temporal_patch_size": 2,
+    "merge_size": 2,
+    "image_mean": [
+        0.5,
+        0.5,
+        0.5
+    ],
+    "image_std": [
+        0.5,
+        0.5,
+        0.5
+    ],
+    "processor_class": "Qwen3VLProcessor",
+    "video_processor_type": "Qwen3VLVideoProcessor"
+}

transformer/config.json

@@ -0,0 +1,92 @@
+{
+  "_class_name": "MageFlow",
+  "apply_text_rotary_emb": false,
+  "axes_dim": [
+    16,
+    56,
+    56
+  ],
+  "checkpoint": false,
+  "context_in_dim": 2560,
+  "depth": 12,
+  "depth_single_blocks": 0,
+  "double_block_type": "double_stream",
+  "guidance_embed": false,
+  "hidden_size": 3072,
+  "in_channels": 128,
+  "max_sequence_length": 2048,
+  "mlp_ratio": 4.0,
+  "num_heads": 24,
+  "out_channels": 128,
+  "packing": true,
+  "param_dtype": "bfloat16",
+  "patch_size": 1,
+  "qkv_bias": true,
+  "quantization_config": {
+    "backend": "native_sm120",
+    "compute_dtype": "bfloat16",
+    "format": "e2m1_block16_ue4m3",
+    "quant_algo": "NVFP4",
+    "quant_method": "mage_flow_nvfp4",
+    "target_count": 48,
+    "targets": [
+      "transformer_blocks.0.img_mlp.net.0.proj",
+      "transformer_blocks.0.img_mlp.net.2",
+      "transformer_blocks.0.txt_mlp.net.0.proj",
+      "transformer_blocks.0.txt_mlp.net.2",
+      "transformer_blocks.1.img_mlp.net.0.proj",
+      "transformer_blocks.1.img_mlp.net.2",
+      "transformer_blocks.1.txt_mlp.net.0.proj",
+      "transformer_blocks.1.txt_mlp.net.2",
+      "transformer_blocks.2.img_mlp.net.0.proj",
+      "transformer_blocks.2.img_mlp.net.2",
+      "transformer_blocks.2.txt_mlp.net.0.proj",
+      "transformer_blocks.2.txt_mlp.net.2",
+      "transformer_blocks.3.img_mlp.net.0.proj",
+      "transformer_blocks.3.img_mlp.net.2",
+      "transformer_blocks.3.txt_mlp.net.0.proj",
+      "transformer_blocks.3.txt_mlp.net.2",
+      "transformer_blocks.4.img_mlp.net.0.proj",
+      "transformer_blocks.4.img_mlp.net.2",
+      "transformer_blocks.4.txt_mlp.net.0.proj",
+      "transformer_blocks.4.txt_mlp.net.2",
+      "transformer_blocks.5.img_mlp.net.0.proj",
+      "transformer_blocks.5.img_mlp.net.2",
+      "transformer_blocks.5.txt_mlp.net.0.proj",
+      "transformer_blocks.5.txt_mlp.net.2",
+      "transformer_blocks.6.img_mlp.net.0.proj",
+      "transformer_blocks.6.img_mlp.net.2",
+      "transformer_blocks.6.txt_mlp.net.0.proj",
+      "transformer_blocks.6.txt_mlp.net.2",
+      "transformer_blocks.7.img_mlp.net.0.proj",
+      "transformer_blocks.7.img_mlp.net.2",
+      "transformer_blocks.7.txt_mlp.net.0.proj",
+      "transformer_blocks.7.txt_mlp.net.2",
+      "transformer_blocks.8.img_mlp.net.0.proj",
+      "transformer_blocks.8.img_mlp.net.2",
+      "transformer_blocks.8.txt_mlp.net.0.proj",
+      "transformer_blocks.8.txt_mlp.net.2",
+      "transformer_blocks.9.img_mlp.net.0.proj",
+      "transformer_blocks.9.img_mlp.net.2",
+      "transformer_blocks.9.txt_mlp.net.0.proj",
+      "transformer_blocks.9.txt_mlp.net.2",
+      "transformer_blocks.10.img_mlp.net.0.proj",
+      "transformer_blocks.10.img_mlp.net.2",
+      "transformer_blocks.10.txt_mlp.net.0.proj",
+      "transformer_blocks.10.txt_mlp.net.2",
+      "transformer_blocks.11.img_mlp.net.0.proj",
+      "transformer_blocks.11.img_mlp.net.2",
+      "transformer_blocks.11.txt_mlp.net.0.proj",
+      "transformer_blocks.11.txt_mlp.net.2"
+    ]
+  },
+  "rope_type": "msrope",
+  "schedule_mode": "z-image",
+  "static_shift": 6.0,
+  "theta": 10000,
+  "time_type": "qwen_proj",
+  "txt_max_length": 2048,
+  "use_time_shift": false,
+  "vec_in_dim": 0,
+  "vec_type": null
+}

transformer/diffusion_pytorch_model.safetensors.index.json

@@ -0,0 +1,501 @@
+{
+  "metadata": {
+    "quantization": "mage_flow_nvfp4_resident_v1",
+    "total_size": 5626828224
+  },
+  "weight_map": {
+    "img_in.bias": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "img_in.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "norm_out.linear.bias": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "norm_out.linear.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "proj_out.bias": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "proj_out.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "time_text_embed.timestep_embedder.linear_1.bias": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "time_text_embed.timestep_embedder.linear_1.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "time_text_embed.timestep_embedder.linear_2.bias": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "time_text_embed.timestep_embedder.linear_2.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "transformer_blocks.0.attn.add_k_proj.bias": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "transformer_blocks.0.attn.add_k_proj.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "transformer_blocks.0.attn.add_q_proj.bias": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "transformer_blocks.0.attn.add_q_proj.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "transformer_blocks.0.attn.add_v_proj.bias": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "transformer_blocks.0.attn.add_v_proj.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "transformer_blocks.0.attn.norm_added_k.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "transformer_blocks.0.attn.norm_added_q.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "transformer_blocks.0.attn.norm_k.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "transformer_blocks.0.attn.norm_q.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "transformer_blocks.0.attn.to_add_out.bias": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "transformer_blocks.0.attn.to_add_out.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "transformer_blocks.0.attn.to_k.bias": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "transformer_blocks.0.attn.to_k.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "transformer_blocks.0.attn.to_out.0.bias": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "transformer_blocks.0.attn.to_out.0.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "transformer_blocks.0.attn.to_q.bias": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "transformer_blocks.0.attn.to_q.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "transformer_blocks.0.attn.to_v.bias": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "transformer_blocks.0.attn.to_v.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "transformer_blocks.0.img_mlp.net.0.proj.bias": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "transformer_blocks.0.img_mlp.net.0.proj.packed_weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "transformer_blocks.0.img_mlp.net.0.proj.weight_scale": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "transformer_blocks.0.img_mlp.net.0.proj.weight_scales": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "transformer_blocks.0.img_mlp.net.2.bias": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "transformer_blocks.0.img_mlp.net.2.packed_weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "transformer_blocks.0.img_mlp.net.2.weight_scale": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "transformer_blocks.0.img_mlp.net.2.weight_scales": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "transformer_blocks.0.img_mod.1.bias": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "transformer_blocks.0.img_mod.1.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "transformer_blocks.0.txt_mlp.net.0.proj.bias": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "transformer_blocks.0.txt_mlp.net.0.proj.packed_weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "transformer_blocks.0.txt_mlp.net.0.proj.weight_scale": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "transformer_blocks.0.txt_mlp.net.0.proj.weight_scales": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "transformer_blocks.0.txt_mlp.net.2.bias": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "transformer_blocks.0.txt_mlp.net.2.packed_weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "transformer_blocks.0.txt_mlp.net.2.weight_scale": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "transformer_blocks.0.txt_mlp.net.2.weight_scales": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "transformer_blocks.0.txt_mod.1.bias": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "transformer_blocks.0.txt_mod.1.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "transformer_blocks.1.attn.add_k_proj.bias": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "transformer_blocks.1.attn.add_k_proj.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "transformer_blocks.1.attn.add_q_proj.bias": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "transformer_blocks.1.attn.add_q_proj.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "transformer_blocks.1.attn.add_v_proj.bias": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "transformer_blocks.1.attn.add_v_proj.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "transformer_blocks.1.attn.norm_added_k.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "transformer_blocks.1.attn.norm_added_q.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "transformer_blocks.1.attn.norm_k.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "transformer_blocks.1.attn.norm_q.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "transformer_blocks.1.attn.to_add_out.bias": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "transformer_blocks.1.attn.to_add_out.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "transformer_blocks.1.attn.to_k.bias": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "transformer_blocks.1.attn.to_k.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "transformer_blocks.1.attn.to_out.0.bias": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "transformer_blocks.1.attn.to_out.0.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "transformer_blocks.1.attn.to_q.bias": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "transformer_blocks.1.attn.to_q.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "transformer_blocks.1.attn.to_v.bias": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "transformer_blocks.1.attn.to_v.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "transformer_blocks.1.img_mlp.net.0.proj.bias": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "transformer_blocks.1.img_mlp.net.0.proj.packed_weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "transformer_blocks.1.img_mlp.net.0.proj.weight_scale": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "transformer_blocks.1.img_mlp.net.0.proj.weight_scales": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "transformer_blocks.1.img_mlp.net.2.bias": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "transformer_blocks.1.img_mlp.net.2.packed_weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "transformer_blocks.1.img_mlp.net.2.weight_scale": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "transformer_blocks.1.img_mlp.net.2.weight_scales": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "transformer_blocks.1.img_mod.1.bias": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "transformer_blocks.1.img_mod.1.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "transformer_blocks.1.txt_mlp.net.0.proj.bias": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "transformer_blocks.1.txt_mlp.net.0.proj.packed_weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "transformer_blocks.1.txt_mlp.net.0.proj.weight_scale": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "transformer_blocks.1.txt_mlp.net.0.proj.weight_scales": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "transformer_blocks.1.txt_mlp.net.2.bias": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "transformer_blocks.1.txt_mlp.net.2.packed_weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "transformer_blocks.1.txt_mlp.net.2.weight_scale": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "transformer_blocks.1.txt_mlp.net.2.weight_scales": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "transformer_blocks.1.txt_mod.1.bias": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "transformer_blocks.1.txt_mod.1.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "transformer_blocks.10.attn.add_k_proj.bias": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "transformer_blocks.10.attn.add_k_proj.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "transformer_blocks.10.attn.add_q_proj.bias": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "transformer_blocks.10.attn.add_q_proj.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "transformer_blocks.10.attn.add_v_proj.bias": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "transformer_blocks.10.attn.add_v_proj.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "transformer_blocks.10.attn.norm_added_k.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "transformer_blocks.10.attn.norm_added_q.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "transformer_blocks.10.attn.norm_k.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "transformer_blocks.10.attn.norm_q.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "transformer_blocks.10.attn.to_add_out.bias": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "transformer_blocks.10.attn.to_add_out.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "transformer_blocks.10.attn.to_k.bias": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "transformer_blocks.10.attn.to_k.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "transformer_blocks.10.attn.to_out.0.bias": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "transformer_blocks.10.attn.to_out.0.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "transformer_blocks.10.attn.to_q.bias": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "transformer_blocks.10.attn.to_q.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "transformer_blocks.10.attn.to_v.bias": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "transformer_blocks.10.attn.to_v.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "transformer_blocks.10.img_mlp.net.0.proj.bias": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "transformer_blocks.10.img_mlp.net.0.proj.packed_weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "transformer_blocks.10.img_mlp.net.0.proj.weight_scale": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "transformer_blocks.10.img_mlp.net.0.proj.weight_scales": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "transformer_blocks.10.img_mlp.net.2.bias": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "transformer_blocks.10.img_mlp.net.2.packed_weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "transformer_blocks.10.img_mlp.net.2.weight_scale": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "transformer_blocks.10.img_mlp.net.2.weight_scales": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "transformer_blocks.10.img_mod.1.bias": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "transformer_blocks.10.img_mod.1.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "transformer_blocks.10.txt_mlp.net.0.proj.bias": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "transformer_blocks.10.txt_mlp.net.0.proj.packed_weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "transformer_blocks.10.txt_mlp.net.0.proj.weight_scale": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "transformer_blocks.10.txt_mlp.net.0.proj.weight_scales": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "transformer_blocks.10.txt_mlp.net.2.bias": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "transformer_blocks.10.txt_mlp.net.2.packed_weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "transformer_blocks.10.txt_mlp.net.2.weight_scale": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "transformer_blocks.10.txt_mlp.net.2.weight_scales": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "transformer_blocks.10.txt_mod.1.bias": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "transformer_blocks.10.txt_mod.1.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "transformer_blocks.11.attn.add_k_proj.bias": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "transformer_blocks.11.attn.add_k_proj.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "transformer_blocks.11.attn.add_q_proj.bias": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "transformer_blocks.11.attn.add_q_proj.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "transformer_blocks.11.attn.add_v_proj.bias": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "transformer_blocks.11.attn.add_v_proj.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "transformer_blocks.11.attn.norm_added_k.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "transformer_blocks.11.attn.norm_added_q.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "transformer_blocks.11.attn.norm_k.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "transformer_blocks.11.attn.norm_q.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "transformer_blocks.11.attn.to_add_out.bias": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "transformer_blocks.11.attn.to_add_out.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "transformer_blocks.11.attn.to_k.bias": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "transformer_blocks.11.attn.to_k.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "transformer_blocks.11.attn.to_out.0.bias": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "transformer_blocks.11.attn.to_out.0.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "transformer_blocks.11.attn.to_q.bias": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "transformer_blocks.11.attn.to_q.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "transformer_blocks.11.attn.to_v.bias": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "transformer_blocks.11.attn.to_v.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "transformer_blocks.11.img_mlp.net.0.proj.bias": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "transformer_blocks.11.img_mlp.net.0.proj.packed_weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "transformer_blocks.11.img_mlp.net.0.proj.weight_scale": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "transformer_blocks.11.img_mlp.net.0.proj.weight_scales": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "transformer_blocks.11.img_mlp.net.2.bias": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "transformer_blocks.11.img_mlp.net.2.packed_weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "transformer_blocks.11.img_mlp.net.2.weight_scale": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "transformer_blocks.11.img_mlp.net.2.weight_scales": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "transformer_blocks.11.img_mod.1.bias": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "transformer_blocks.11.img_mod.1.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "transformer_blocks.11.txt_mlp.net.0.proj.bias": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "transformer_blocks.11.txt_mlp.net.0.proj.packed_weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "transformer_blocks.11.txt_mlp.net.0.proj.weight_scale": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "transformer_blocks.11.txt_mlp.net.0.proj.weight_scales": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "transformer_blocks.11.txt_mlp.net.2.bias": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "transformer_blocks.11.txt_mlp.net.2.packed_weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "transformer_blocks.11.txt_mlp.net.2.weight_scale": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "transformer_blocks.11.txt_mlp.net.2.weight_scales": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "transformer_blocks.11.txt_mod.1.bias": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "transformer_blocks.11.txt_mod.1.weight": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "transformer_blocks.2.attn.add_k_proj.bias": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "transformer_blocks.2.attn.add_k_proj.weight": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "transformer_blocks.2.attn.add_q_proj.bias": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "transformer_blocks.2.attn.add_q_proj.weight": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "transformer_blocks.2.attn.add_v_proj.bias": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "transformer_blocks.2.attn.add_v_proj.weight": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "transformer_blocks.2.attn.norm_added_k.weight": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "transformer_blocks.2.attn.norm_added_q.weight": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "transformer_blocks.2.attn.norm_k.weight": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "transformer_blocks.2.attn.norm_q.weight": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "transformer_blocks.2.attn.to_add_out.bias": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "transformer_blocks.2.attn.to_add_out.weight": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "transformer_blocks.2.attn.to_k.bias": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "transformer_blocks.2.attn.to_k.weight": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "transformer_blocks.2.attn.to_out.0.bias": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "transformer_blocks.2.attn.to_out.0.weight": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "transformer_blocks.2.attn.to_q.bias": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "transformer_blocks.2.attn.to_q.weight": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "transformer_blocks.2.attn.to_v.bias": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "transformer_blocks.2.attn.to_v.weight": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "transformer_blocks.2.img_mlp.net.0.proj.bias": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "transformer_blocks.2.img_mlp.net.0.proj.packed_weight": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "transformer_blocks.2.img_mlp.net.0.proj.weight_scale": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "transformer_blocks.2.img_mlp.net.0.proj.weight_scales": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "transformer_blocks.2.img_mlp.net.2.bias": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "transformer_blocks.2.img_mlp.net.2.packed_weight": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "transformer_blocks.2.img_mlp.net.2.weight_scale": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "transformer_blocks.2.img_mlp.net.2.weight_scales": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "transformer_blocks.2.img_mod.1.bias": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "transformer_blocks.2.img_mod.1.weight": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "transformer_blocks.2.txt_mlp.net.0.proj.bias": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "transformer_blocks.2.txt_mlp.net.0.proj.packed_weight": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "transformer_blocks.2.txt_mlp.net.0.proj.weight_scale": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "transformer_blocks.2.txt_mlp.net.0.proj.weight_scales": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "transformer_blocks.2.txt_mlp.net.2.bias": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "transformer_blocks.2.txt_mlp.net.2.packed_weight": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "transformer_blocks.2.txt_mlp.net.2.weight_scale": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "transformer_blocks.2.txt_mlp.net.2.weight_scales": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "transformer_blocks.2.txt_mod.1.bias": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "transformer_blocks.2.txt_mod.1.weight": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "transformer_blocks.3.attn.add_k_proj.bias": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "transformer_blocks.3.attn.add_k_proj.weight": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "transformer_blocks.3.attn.add_q_proj.bias": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "transformer_blocks.3.attn.add_q_proj.weight": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "transformer_blocks.3.attn.add_v_proj.bias": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "transformer_blocks.3.attn.add_v_proj.weight": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "transformer_blocks.3.attn.norm_added_k.weight": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "transformer_blocks.3.attn.norm_added_q.weight": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "transformer_blocks.3.attn.norm_k.weight": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "transformer_blocks.3.attn.norm_q.weight": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "transformer_blocks.3.attn.to_add_out.bias": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "transformer_blocks.3.attn.to_add_out.weight": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "transformer_blocks.3.attn.to_k.bias": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "transformer_blocks.3.attn.to_k.weight": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "transformer_blocks.3.attn.to_out.0.bias": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "transformer_blocks.3.attn.to_out.0.weight": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "transformer_blocks.3.attn.to_q.bias": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "transformer_blocks.3.attn.to_q.weight": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "transformer_blocks.3.attn.to_v.bias": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "transformer_blocks.3.attn.to_v.weight": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "transformer_blocks.3.img_mlp.net.0.proj.bias": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "transformer_blocks.3.img_mlp.net.0.proj.packed_weight": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "transformer_blocks.3.img_mlp.net.0.proj.weight_scale": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "transformer_blocks.3.img_mlp.net.0.proj.weight_scales": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "transformer_blocks.3.img_mlp.net.2.bias": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "transformer_blocks.3.img_mlp.net.2.packed_weight": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "transformer_blocks.3.img_mlp.net.2.weight_scale": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "transformer_blocks.3.img_mlp.net.2.weight_scales": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "transformer_blocks.3.img_mod.1.bias": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "transformer_blocks.3.img_mod.1.weight": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "transformer_blocks.3.txt_mlp.net.0.proj.bias": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "transformer_blocks.3.txt_mlp.net.0.proj.packed_weight": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "transformer_blocks.3.txt_mlp.net.0.proj.weight_scale": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "transformer_blocks.3.txt_mlp.net.0.proj.weight_scales": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "transformer_blocks.3.txt_mlp.net.2.bias": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "transformer_blocks.3.txt_mlp.net.2.packed_weight": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "transformer_blocks.3.txt_mlp.net.2.weight_scale": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "transformer_blocks.3.txt_mlp.net.2.weight_scales": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "transformer_blocks.3.txt_mod.1.bias": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "transformer_blocks.3.txt_mod.1.weight": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "transformer_blocks.4.attn.add_k_proj.bias": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "transformer_blocks.4.attn.add_k_proj.weight": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "transformer_blocks.4.attn.add_q_proj.bias": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "transformer_blocks.4.attn.add_q_proj.weight": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "transformer_blocks.4.attn.add_v_proj.bias": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "transformer_blocks.4.attn.add_v_proj.weight": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "transformer_blocks.4.attn.norm_added_k.weight": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "transformer_blocks.4.attn.norm_added_q.weight": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "transformer_blocks.4.attn.norm_k.weight": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "transformer_blocks.4.attn.norm_q.weight": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "transformer_blocks.4.attn.to_add_out.bias": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "transformer_blocks.4.attn.to_add_out.weight": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "transformer_blocks.4.attn.to_k.bias": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "transformer_blocks.4.attn.to_k.weight": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "transformer_blocks.4.attn.to_out.0.bias": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "transformer_blocks.4.attn.to_out.0.weight": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "transformer_blocks.4.attn.to_q.bias": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "transformer_blocks.4.attn.to_q.weight": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "transformer_blocks.4.attn.to_v.bias": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "transformer_blocks.4.attn.to_v.weight": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "transformer_blocks.4.img_mlp.net.0.proj.bias": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "transformer_blocks.4.img_mlp.net.0.proj.packed_weight": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "transformer_blocks.4.img_mlp.net.0.proj.weight_scale": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "transformer_blocks.4.img_mlp.net.0.proj.weight_scales": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "transformer_blocks.4.img_mlp.net.2.bias": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "transformer_blocks.4.img_mlp.net.2.packed_weight": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "transformer_blocks.4.img_mlp.net.2.weight_scale": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "transformer_blocks.4.img_mlp.net.2.weight_scales": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "transformer_blocks.4.img_mod.1.bias": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "transformer_blocks.4.img_mod.1.weight": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "transformer_blocks.4.txt_mlp.net.0.proj.bias": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "transformer_blocks.4.txt_mlp.net.0.proj.packed_weight": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "transformer_blocks.4.txt_mlp.net.0.proj.weight_scale": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "transformer_blocks.4.txt_mlp.net.0.proj.weight_scales": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "transformer_blocks.4.txt_mlp.net.2.bias": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "transformer_blocks.4.txt_mlp.net.2.packed_weight": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "transformer_blocks.4.txt_mlp.net.2.weight_scale": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "transformer_blocks.4.txt_mlp.net.2.weight_scales": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "transformer_blocks.4.txt_mod.1.bias": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "transformer_blocks.4.txt_mod.1.weight": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "transformer_blocks.5.attn.add_k_proj.bias": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "transformer_blocks.5.attn.add_k_proj.weight": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "transformer_blocks.5.attn.add_q_proj.bias": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "transformer_blocks.5.attn.add_q_proj.weight": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "transformer_blocks.5.attn.add_v_proj.bias": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "transformer_blocks.5.attn.add_v_proj.weight": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "transformer_blocks.5.attn.norm_added_k.weight": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "transformer_blocks.5.attn.norm_added_q.weight": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "transformer_blocks.5.attn.norm_k.weight": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "transformer_blocks.5.attn.norm_q.weight": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "transformer_blocks.5.attn.to_add_out.bias": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "transformer_blocks.5.attn.to_add_out.weight": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "transformer_blocks.5.attn.to_k.bias": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "transformer_blocks.5.attn.to_k.weight": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "transformer_blocks.5.attn.to_out.0.bias": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "transformer_blocks.5.attn.to_out.0.weight": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "transformer_blocks.5.attn.to_q.bias": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "transformer_blocks.5.attn.to_q.weight": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "transformer_blocks.5.attn.to_v.bias": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "transformer_blocks.5.attn.to_v.weight": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "transformer_blocks.5.img_mlp.net.0.proj.bias": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "transformer_blocks.5.img_mlp.net.0.proj.packed_weight": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "transformer_blocks.5.img_mlp.net.0.proj.weight_scale": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "transformer_blocks.5.img_mlp.net.0.proj.weight_scales": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "transformer_blocks.5.img_mlp.net.2.bias": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "transformer_blocks.5.img_mlp.net.2.packed_weight": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "transformer_blocks.5.img_mlp.net.2.weight_scale": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "transformer_blocks.5.img_mlp.net.2.weight_scales": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "transformer_blocks.5.img_mod.1.bias": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "transformer_blocks.5.img_mod.1.weight": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "transformer_blocks.5.txt_mlp.net.0.proj.bias": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "transformer_blocks.5.txt_mlp.net.0.proj.packed_weight": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "transformer_blocks.5.txt_mlp.net.0.proj.weight_scale": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "transformer_blocks.5.txt_mlp.net.0.proj.weight_scales": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "transformer_blocks.5.txt_mlp.net.2.bias": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "transformer_blocks.5.txt_mlp.net.2.packed_weight": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "transformer_blocks.5.txt_mlp.net.2.weight_scale": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "transformer_blocks.5.txt_mlp.net.2.weight_scales": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "transformer_blocks.5.txt_mod.1.bias": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "transformer_blocks.5.txt_mod.1.weight": "diffusion_pytorch_model-00003-of-00004.safetensors",
+    "transformer_blocks.6.attn.add_k_proj.bias": "diffusion_pytorch_model-00003-of-00004.safetensors",
+    "transformer_blocks.6.attn.add_k_proj.weight": "diffusion_pytorch_model-00003-of-00004.safetensors",
+    "transformer_blocks.6.attn.add_q_proj.bias": "diffusion_pytorch_model-00003-of-00004.safetensors",
+    "transformer_blocks.6.attn.add_q_proj.weight": "diffusion_pytorch_model-00003-of-00004.safetensors",
+    "transformer_blocks.6.attn.add_v_proj.bias": "diffusion_pytorch_model-00003-of-00004.safetensors",
+    "transformer_blocks.6.attn.add_v_proj.weight": "diffusion_pytorch_model-00003-of-00004.safetensors",
+    "transformer_blocks.6.attn.norm_added_k.weight": "diffusion_pytorch_model-00003-of-00004.safetensors",
+    "transformer_blocks.6.attn.norm_added_q.weight": "diffusion_pytorch_model-00003-of-00004.safetensors",
+    "transformer_blocks.6.attn.norm_k.weight": "diffusion_pytorch_model-00003-of-00004.safetensors",
+    "transformer_blocks.6.attn.norm_q.weight": "diffusion_pytorch_model-00003-of-00004.safetensors",
+    "transformer_blocks.6.attn.to_add_out.bias": "diffusion_pytorch_model-00003-of-00004.safetensors",
+    "transformer_blocks.6.attn.to_add_out.weight": "diffusion_pytorch_model-00003-of-00004.safetensors",
+    "transformer_blocks.6.attn.to_k.bias": "diffusion_pytorch_model-00003-of-00004.safetensors",
+    "transformer_blocks.6.attn.to_k.weight": "diffusion_pytorch_model-00003-of-00004.safetensors",
+    "transformer_blocks.6.attn.to_out.0.bias": "diffusion_pytorch_model-00003-of-00004.safetensors",
+    "transformer_blocks.6.attn.to_out.0.weight": "diffusion_pytorch_model-00003-of-00004.safetensors",
+    "transformer_blocks.6.attn.to_q.bias": "diffusion_pytorch_model-00003-of-00004.safetensors",
+    "transformer_blocks.6.attn.to_q.weight": "diffusion_pytorch_model-00003-of-00004.safetensors",
+    "transformer_blocks.6.attn.to_v.bias": "diffusion_pytorch_model-00003-of-00004.safetensors",
+    "transformer_blocks.6.attn.to_v.weight": "diffusion_pytorch_model-00003-of-00004.safetensors",
+    "transformer_blocks.6.img_mlp.net.0.proj.bias": "diffusion_pytorch_model-00003-of-00004.safetensors",
+    "transformer_blocks.6.img_mlp.net.0.proj.packed_weight": "diffusion_pytorch_model-00003-of-00004.safetensors",
+    "transformer_blocks.6.img_mlp.net.0.proj.weight_scale": "diffusion_pytorch_model-00003-of-00004.safetensors",
+    "transformer_blocks.6.img_mlp.net.0.proj.weight_scales": "diffusion_pytorch_model-00003-of-00004.safetensors",
+    "transformer_blocks.6.img_mlp.net.2.bias": "diffusion_pytorch_model-00003-of-00004.safetensors",
+    "transformer_blocks.6.img_mlp.net.2.packed_weight": "diffusion_pytorch_model-00003-of-00004.safetensors",
+    "transformer_blocks.6.img_mlp.net.2.weight_scale": "diffusion_pytorch_model-00003-of-00004.safetensors",
+    "transformer_blocks.6.img_mlp.net.2.weight_scales": "diffusion_pytorch_model-00003-of-00004.safetensors",
+    "transformer_blocks.6.img_mod.1.bias": "diffusion_pytorch_model-00003-of-00004.safetensors",
+    "transformer_blocks.6.img_mod.1.weight": "diffusion_pytorch_model-00003-of-00004.safetensors",
+    "transformer_blocks.6.txt_mlp.net.0.proj.bias": "diffusion_pytorch_model-00003-of-00004.safetensors",
+    "transformer_blocks.6.txt_mlp.net.0.proj.packed_weight": "diffusion_pytorch_model-00003-of-00004.safetensors",
+    "transformer_blocks.6.txt_mlp.net.0.proj.weight_scale": "diffusion_pytorch_model-00003-of-00004.safetensors",
+    "transformer_blocks.6.txt_mlp.net.0.proj.weight_scales": "diffusion_pytorch_model-00003-of-00004.safetensors",
+    "transformer_blocks.6.txt_mlp.net.2.bias": "diffusion_pytorch_model-00003-of-00004.safetensors",
+    "transformer_blocks.6.txt_mlp.net.2.packed_weight": "diffusion_pytorch_model-00003-of-00004.safetensors",
+    "transformer_blocks.6.txt_mlp.net.2.weight_scale": "diffusion_pytorch_model-00003-of-00004.safetensors",
+    "transformer_blocks.6.txt_mlp.net.2.weight_scales": "diffusion_pytorch_model-00003-of-00004.safetensors",
+    "transformer_blocks.6.txt_mod.1.bias": "diffusion_pytorch_model-00003-of-00004.safetensors",
+    "transformer_blocks.6.txt_mod.1.weight": "diffusion_pytorch_model-00003-of-00004.safetensors",
+    "transformer_blocks.7.attn.add_k_proj.bias": "diffusion_pytorch_model-00003-of-00004.safetensors",
+    "transformer_blocks.7.attn.add_k_proj.weight": "diffusion_pytorch_model-00003-of-00004.safetensors",
+    "transformer_blocks.7.attn.add_q_proj.bias": "diffusion_pytorch_model-00003-of-00004.safetensors",
+    "transformer_blocks.7.attn.add_q_proj.weight": "diffusion_pytorch_model-00003-of-00004.safetensors",
+    "transformer_blocks.7.attn.add_v_proj.bias": "diffusion_pytorch_model-00003-of-00004.safetensors",
+    "transformer_blocks.7.attn.add_v_proj.weight": "diffusion_pytorch_model-00003-of-00004.safetensors",
+    "transformer_blocks.7.attn.norm_added_k.weight": "diffusion_pytorch_model-00003-of-00004.safetensors",
+    "transformer_blocks.7.attn.norm_added_q.weight": "diffusion_pytorch_model-00003-of-00004.safetensors",
+    "transformer_blocks.7.attn.norm_k.weight": "diffusion_pytorch_model-00003-of-00004.safetensors",
+    "transformer_blocks.7.attn.norm_q.weight": "diffusion_pytorch_model-00003-of-00004.safetensors",
+    "transformer_blocks.7.attn.to_add_out.bias": "diffusion_pytorch_model-00003-of-00004.safetensors",
+    "transformer_blocks.7.attn.to_add_out.weight": "diffusion_pytorch_model-00003-of-00004.safetensors",
+    "transformer_blocks.7.attn.to_k.bias": "diffusion_pytorch_model-00003-of-00004.safetensors",
+    "transformer_blocks.7.attn.to_k.weight": "diffusion_pytorch_model-00003-of-00004.safetensors",
+    "transformer_blocks.7.attn.to_out.0.bias": "diffusion_pytorch_model-00003-of-00004.safetensors",
+    "transformer_blocks.7.attn.to_out.0.weight": "diffusion_pytorch_model-00003-of-00004.safetensors",
+    "transformer_blocks.7.attn.to_q.bias": "diffusion_pytorch_model-00003-of-00004.safetensors",
+    "transformer_blocks.7.attn.to_q.weight": "diffusion_pytorch_model-00003-of-00004.safetensors",
+    "transformer_blocks.7.attn.to_v.bias": "diffusion_pytorch_model-00003-of-00004.safetensors",
+    "transformer_blocks.7.attn.to_v.weight": "diffusion_pytorch_model-00003-of-00004.safetensors",
+    "transformer_blocks.7.img_mlp.net.0.proj.bias": "diffusion_pytorch_model-00003-of-00004.safetensors",
+    "transformer_blocks.7.img_mlp.net.0.proj.packed_weight": "diffusion_pytorch_model-00003-of-00004.safetensors",
+    "transformer_blocks.7.img_mlp.net.0.proj.weight_scale": "diffusion_pytorch_model-00003-of-00004.safetensors",
+    "transformer_blocks.7.img_mlp.net.0.proj.weight_scales": "diffusion_pytorch_model-00003-of-00004.safetensors",
+    "transformer_blocks.7.img_mlp.net.2.bias": "diffusion_pytorch_model-00003-of-00004.safetensors",
+    "transformer_blocks.7.img_mlp.net.2.packed_weight": "diffusion_pytorch_model-00003-of-00004.safetensors",
+    "transformer_blocks.7.img_mlp.net.2.weight_scale": "diffusion_pytorch_model-00003-of-00004.safetensors",
+    "transformer_blocks.7.img_mlp.net.2.weight_scales": "diffusion_pytorch_model-00003-of-00004.safetensors",
+    "transformer_blocks.7.img_mod.1.bias": "diffusion_pytorch_model-00003-of-00004.safetensors",
+    "transformer_blocks.7.img_mod.1.weight": "diffusion_pytorch_model-00003-of-00004.safetensors",
+    "transformer_blocks.7.txt_mlp.net.0.proj.bias": "diffusion_pytorch_model-00003-of-00004.safetensors",
+    "transformer_blocks.7.txt_mlp.net.0.proj.packed_weight": "diffusion_pytorch_model-00003-of-00004.safetensors",
+    "transformer_blocks.7.txt_mlp.net.0.proj.weight_scale": "diffusion_pytorch_model-00003-of-00004.safetensors",
+    "transformer_blocks.7.txt_mlp.net.0.proj.weight_scales": "diffusion_pytorch_model-00003-of-00004.safetensors",
+    "transformer_blocks.7.txt_mlp.net.2.bias": "diffusion_pytorch_model-00003-of-00004.safetensors",
+    "transformer_blocks.7.txt_mlp.net.2.packed_weight": "diffusion_pytorch_model-00003-of-00004.safetensors",
+    "transformer_blocks.7.txt_mlp.net.2.weight_scale": "diffusion_pytorch_model-00003-of-00004.safetensors",
+    "transformer_blocks.7.txt_mlp.net.2.weight_scales": "diffusion_pytorch_model-00003-of-00004.safetensors",
+    "transformer_blocks.7.txt_mod.1.bias": "diffusion_pytorch_model-00003-of-00004.safetensors",
+    "transformer_blocks.7.txt_mod.1.weight": "diffusion_pytorch_model-00003-of-00004.safetensors",
+    "transformer_blocks.8.attn.add_k_proj.bias": "diffusion_pytorch_model-00003-of-00004.safetensors",
+    "transformer_blocks.8.attn.add_k_proj.weight": "diffusion_pytorch_model-00003-of-00004.safetensors",
+    "transformer_blocks.8.attn.add_q_proj.bias": "diffusion_pytorch_model-00003-of-00004.safetensors",
+    "transformer_blocks.8.attn.add_q_proj.weight": "diffusion_pytorch_model-00003-of-00004.safetensors",
+    "transformer_blocks.8.attn.add_v_proj.bias": "diffusion_pytorch_model-00003-of-00004.safetensors",
+    "transformer_blocks.8.attn.add_v_proj.weight": "diffusion_pytorch_model-00003-of-00004.safetensors",
+    "transformer_blocks.8.attn.norm_added_k.weight": "diffusion_pytorch_model-00003-of-00004.safetensors",
+    "transformer_blocks.8.attn.norm_added_q.weight": "diffusion_pytorch_model-00003-of-00004.safetensors",
+    "transformer_blocks.8.attn.norm_k.weight": "diffusion_pytorch_model-00003-of-00004.safetensors",
+    "transformer_blocks.8.attn.norm_q.weight": "diffusion_pytorch_model-00003-of-00004.safetensors",
+    "transformer_blocks.8.attn.to_add_out.bias": "diffusion_pytorch_model-00003-of-00004.safetensors",
+    "transformer_blocks.8.attn.to_add_out.weight": "diffusion_pytorch_model-00003-of-00004.safetensors",
+    "transformer_blocks.8.attn.to_k.bias": "diffusion_pytorch_model-00003-of-00004.safetensors",
+    "transformer_blocks.8.attn.to_k.weight": "diffusion_pytorch_model-00003-of-00004.safetensors",
+    "transformer_blocks.8.attn.to_out.0.bias": "diffusion_pytorch_model-00003-of-00004.safetensors",
+    "transformer_blocks.8.attn.to_out.0.weight": "diffusion_pytorch_model-00003-of-00004.safetensors",
+    "transformer_blocks.8.attn.to_q.bias": "diffusion_pytorch_model-00003-of-00004.safetensors",
+    "transformer_blocks.8.attn.to_q.weight": "diffusion_pytorch_model-00003-of-00004.safetensors",
+    "transformer_blocks.8.attn.to_v.bias": "diffusion_pytorch_model-00003-of-00004.safetensors",
+    "transformer_blocks.8.attn.to_v.weight": "diffusion_pytorch_model-00003-of-00004.safetensors",
+    "transformer_blocks.8.img_mlp.net.0.proj.bias": "diffusion_pytorch_model-00003-of-00004.safetensors",
+    "transformer_blocks.8.img_mlp.net.0.proj.packed_weight": "diffusion_pytorch_model-00003-of-00004.safetensors",
+    "transformer_blocks.8.img_mlp.net.0.proj.weight_scale": "diffusion_pytorch_model-00003-of-00004.safetensors",
+    "transformer_blocks.8.img_mlp.net.0.proj.weight_scales": "diffusion_pytorch_model-00003-of-00004.safetensors",
+    "transformer_blocks.8.img_mlp.net.2.bias": "diffusion_pytorch_model-00003-of-00004.safetensors",
+    "transformer_blocks.8.img_mlp.net.2.packed_weight": "diffusion_pytorch_model-00003-of-00004.safetensors",
+    "transformer_blocks.8.img_mlp.net.2.weight_scale": "diffusion_pytorch_model-00003-of-00004.safetensors",
+    "transformer_blocks.8.img_mlp.net.2.weight_scales": "diffusion_pytorch_model-00003-of-00004.safetensors",
+    "transformer_blocks.8.img_mod.1.bias": "diffusion_pytorch_model-00003-of-00004.safetensors",
+    "transformer_blocks.8.img_mod.1.weight": "diffusion_pytorch_model-00003-of-00004.safetensors",
+    "transformer_blocks.8.txt_mlp.net.0.proj.bias": "diffusion_pytorch_model-00003-of-00004.safetensors",
+    "transformer_blocks.8.txt_mlp.net.0.proj.packed_weight": "diffusion_pytorch_model-00003-of-00004.safetensors",
+    "transformer_blocks.8.txt_mlp.net.0.proj.weight_scale": "diffusion_pytorch_model-00003-of-00004.safetensors",
+    "transformer_blocks.8.txt_mlp.net.0.proj.weight_scales": "diffusion_pytorch_model-00003-of-00004.safetensors",
+    "transformer_blocks.8.txt_mlp.net.2.bias": "diffusion_pytorch_model-00003-of-00004.safetensors",
+    "transformer_blocks.8.txt_mlp.net.2.packed_weight": "diffusion_pytorch_model-00003-of-00004.safetensors",
+    "transformer_blocks.8.txt_mlp.net.2.weight_scale": "diffusion_pytorch_model-00003-of-00004.safetensors",
+    "transformer_blocks.8.txt_mlp.net.2.weight_scales": "diffusion_pytorch_model-00003-of-00004.safetensors",
+    "transformer_blocks.8.txt_mod.1.bias": "diffusion_pytorch_model-00003-of-00004.safetensors",
+    "transformer_blocks.8.txt_mod.1.weight": "diffusion_pytorch_model-00003-of-00004.safetensors",
+    "transformer_blocks.9.attn.add_k_proj.bias": "diffusion_pytorch_model-00003-of-00004.safetensors",
+    "transformer_blocks.9.attn.add_k_proj.weight": "diffusion_pytorch_model-00003-of-00004.safetensors",
+    "transformer_blocks.9.attn.add_q_proj.bias": "diffusion_pytorch_model-00003-of-00004.safetensors",
+    "transformer_blocks.9.attn.add_q_proj.weight": "diffusion_pytorch_model-00003-of-00004.safetensors",
+    "transformer_blocks.9.attn.add_v_proj.bias": "diffusion_pytorch_model-00003-of-00004.safetensors",
+    "transformer_blocks.9.attn.add_v_proj.weight": "diffusion_pytorch_model-00003-of-00004.safetensors",
+    "transformer_blocks.9.attn.norm_added_k.weight": "diffusion_pytorch_model-00003-of-00004.safetensors",
+    "transformer_blocks.9.attn.norm_added_q.weight": "diffusion_pytorch_model-00003-of-00004.safetensors",
+    "transformer_blocks.9.attn.norm_k.weight": "diffusion_pytorch_model-00003-of-00004.safetensors",
+    "transformer_blocks.9.attn.norm_q.weight": "diffusion_pytorch_model-00003-of-00004.safetensors",
+    "transformer_blocks.9.attn.to_add_out.bias": "diffusion_pytorch_model-00003-of-00004.safetensors",
+    "transformer_blocks.9.attn.to_add_out.weight": "diffusion_pytorch_model-00003-of-00004.safetensors",
+    "transformer_blocks.9.attn.to_k.bias": "diffusion_pytorch_model-00003-of-00004.safetensors",
+    "transformer_blocks.9.attn.to_k.weight": "diffusion_pytorch_model-00003-of-00004.safetensors",
+    "transformer_blocks.9.attn.to_out.0.bias": "diffusion_pytorch_model-00003-of-00004.safetensors",
+    "transformer_blocks.9.attn.to_out.0.weight": "diffusion_pytorch_model-00003-of-00004.safetensors",
+    "transformer_blocks.9.attn.to_q.bias": "diffusion_pytorch_model-00003-of-00004.safetensors",
+    "transformer_blocks.9.attn.to_q.weight": "diffusion_pytorch_model-00003-of-00004.safetensors",
+    "transformer_blocks.9.attn.to_v.bias": "diffusion_pytorch_model-00003-of-00004.safetensors",
+    "transformer_blocks.9.attn.to_v.weight": "diffusion_pytorch_model-00003-of-00004.safetensors",
+    "transformer_blocks.9.img_mlp.net.0.proj.bias": "diffusion_pytorch_model-00003-of-00004.safetensors",
+    "transformer_blocks.9.img_mlp.net.0.proj.packed_weight": "diffusion_pytorch_model-00003-of-00004.safetensors",
+    "transformer_blocks.9.img_mlp.net.0.proj.weight_scale": "diffusion_pytorch_model-00003-of-00004.safetensors",
+    "transformer_blocks.9.img_mlp.net.0.proj.weight_scales": "diffusion_pytorch_model-00003-of-00004.safetensors",
+    "transformer_blocks.9.img_mlp.net.2.bias": "diffusion_pytorch_model-00003-of-00004.safetensors",
+    "transformer_blocks.9.img_mlp.net.2.packed_weight": "diffusion_pytorch_model-00003-of-00004.safetensors",
+    "transformer_blocks.9.img_mlp.net.2.weight_scale": "diffusion_pytorch_model-00003-of-00004.safetensors",
+    "transformer_blocks.9.img_mlp.net.2.weight_scales": "diffusion_pytorch_model-00003-of-00004.safetensors",
+    "transformer_blocks.9.img_mod.1.bias": "diffusion_pytorch_model-00003-of-00004.safetensors",
+    "transformer_blocks.9.img_mod.1.weight": "diffusion_pytorch_model-00003-of-00004.safetensors",
+    "transformer_blocks.9.txt_mlp.net.0.proj.bias": "diffusion_pytorch_model-00003-of-00004.safetensors",
+    "transformer_blocks.9.txt_mlp.net.0.proj.packed_weight": "diffusion_pytorch_model-00003-of-00004.safetensors",
+    "transformer_blocks.9.txt_mlp.net.0.proj.weight_scale": "diffusion_pytorch_model-00003-of-00004.safetensors",
+    "transformer_blocks.9.txt_mlp.net.0.proj.weight_scales": "diffusion_pytorch_model-00003-of-00004.safetensors",
+    "transformer_blocks.9.txt_mlp.net.2.bias": "diffusion_pytorch_model-00003-of-00004.safetensors",
+    "transformer_blocks.9.txt_mlp.net.2.packed_weight": "diffusion_pytorch_model-00003-of-00004.safetensors",
+    "transformer_blocks.9.txt_mlp.net.2.weight_scale": "diffusion_pytorch_model-00003-of-00004.safetensors",
+    "transformer_blocks.9.txt_mlp.net.2.weight_scales": "diffusion_pytorch_model-00003-of-00004.safetensors",
+    "transformer_blocks.9.txt_mod.1.bias": "diffusion_pytorch_model-00003-of-00004.safetensors",
+    "transformer_blocks.9.txt_mod.1.weight": "diffusion_pytorch_model-00004-of-00004.safetensors",
+    "txt_in.bias": "diffusion_pytorch_model-00004-of-00004.safetensors",
+    "txt_in.weight": "diffusion_pytorch_model-00004-of-00004.safetensors",
+    "txt_norm.weight": "diffusion_pytorch_model-00004-of-00004.safetensors"
+  }
+}

transformer/nvfp4_metadata.json

@@ -0,0 +1,1648 @@
+{
+  "format_version": 1,
+  "artifact_kind": "mage_flow_transformer_mlp_nvfp4_resident_v1",
+  "created_utc": "2026-07-24T07:00:11.747181+00:00",
+  "container": {
+    "format": "safetensors",
+    "header_metadata": {
+      "mage_nvfp4_contract": "mage_flow_transformer_mlp_nvfp4_resident_v1;format_version=1"
+    },
+    "header_encoding": "single deterministic contract key; legacy two-key draft headers remain readable"
+  },
+  "source": {
+    "transformer_config_path": "BASE_MODEL/transformer/config.json",
+    "transformer_checkpoint_path": "BASE_MODEL/transformer/diffusion_pytorch_model.safetensors",
+    "transformer_config_sha256": "8493c3b2722738c2a824ac82b1fd9c89fefb4e354fc88363207193db7fe702de",
+    "transformer_checkpoint_sha256": "04431abbc3acd1a5b86a7f77269f8e77e05a6cac3b21f35e1aef7493f6ab4934"
+  },
+  "library_hashes": {
+    "artifact_script_sha256": "459b4e8b7fb6b15b1f2eac3267082c52f1723efd65a59fe429066f565d224b81",
+    "resident_source_sha256": "060697d4784be3e3461d38ef6ff24128f47a4322a05e5463f8687796b03bda9f",
+    "resident_library_sha256": "d2f7d62d26b4ae3babb7791aee8e0d95ad4794ac9c1d22146c974784381ada48",
+    "mage_flow_py_sha256": "59b6e1bee7f95a7fd2fa7bd9e765966832951e2bb184b5ae27283884997b845f",
+    "mage_layers_py_sha256": "4b198343b8929f48a0a14d388502c81f54be17b222f6831303d2da6a91f33a62",
+    "pipeline_py_sha256": "b9fc57018570372dd3404a733e19b918a359188dd9f5ef7817c6b30969fc13db"
+  },
+  "environment": {
+    "python_version": "3.11.14 (main, Oct 10 2025, 08:54:04) [GCC 13.3.0]",
+    "torch_version": "2.13.0+cu130",
+    "safetensors_version": "0.8.0"
+  },
+  "model": {
+    "depth": 12,
+    "hidden_size": 3072,
+    "num_heads": 24,
+    "context_in_dim": 2560,
+    "in_channels": 128,
+    "out_channels": 128,
+    "patch_size": 1
+  },
+  "quantization": {
+    "format": "nvfp4_two_level",
+    "block_elements": 16,
+    "scale_tile_outer": 128,
+    "scale_tile_inner": 4,
+    "bias_policy": "artifact_bfloat16"
+  },
+  "targets": [
+    {
+      "module_key": "transformer_blocks.0.img_mlp.net.0.proj",
+      "weight_key": "transformer_blocks.0.img_mlp.net.0.proj.weight",
+      "bias_key": "transformer_blocks.0.img_mlp.net.0.proj.bias",
+      "weight_shape": [
+        12288,
+        3072
+      ],
+      "bias_shape": [
+        12288
+      ],
+      "weight_tensor_scale_key": "targets.transformer_blocks.0.img_mlp.net.0.proj.weight_tensor_scale",
+      "artifact_weight_key": "targets.transformer_blocks.0.img_mlp.net.0.proj.packed_weight_e2m1",
+      "artifact_scale_key": "targets.transformer_blocks.0.img_mlp.net.0.proj.packed_scales_ue4m3",
+      "artifact_bias_key": "targets.transformer_blocks.0.img_mlp.net.0.proj.bias_bf16",
+      "weight_tensor_scale": 0.00024559383746236563,
+      "weight_global_amax": 0.66015625,
+      "packed_weight_bytes": 18874368,
+      "packed_scale_bytes": 2359296,
+      "scale_layout": {
+        "inner_dim": 192,
+        "outer_tiles": 96,
+        "bytes": 2359296
+      },
+      "source_weight_sha256": "d05d519a8eb21260cd6c7cdcc886ebd9e32986bd925f3b5a740f75c93ca39162",
+      "source_bias_sha256": "e032c18ee63a7667d67e97e8611440b64e31bda7e60df30fbe1784e6875ea55a"
+    },
+    {
+      "module_key": "transformer_blocks.0.img_mlp.net.2",
+      "weight_key": "transformer_blocks.0.img_mlp.net.2.weight",
+      "bias_key": "transformer_blocks.0.img_mlp.net.2.bias",
+      "weight_shape": [
+        3072,
+        12288
+      ],
+      "bias_shape": [
+        3072
+      ],
+      "weight_tensor_scale_key": "targets.transformer_blocks.0.img_mlp.net.2.weight_tensor_scale",
+      "artifact_weight_key": "targets.transformer_blocks.0.img_mlp.net.2.packed_weight_e2m1",
+      "artifact_scale_key": "targets.transformer_blocks.0.img_mlp.net.2.packed_scales_ue4m3",
+      "artifact_bias_key": "targets.transformer_blocks.0.img_mlp.net.2.bias_bf16",
+      "weight_tensor_scale": 0.0002906435984186828,
+      "weight_global_amax": 0.78125,
+      "packed_weight_bytes": 18874368,
+      "packed_scale_bytes": 2359296,
+      "scale_layout": {
+        "inner_dim": 768,
+        "outer_tiles": 24,
+        "bytes": 2359296
+      },
+      "source_weight_sha256": "8a8cf7a4a5439eb2f3154ad636457ff25f665a95e755e58b2adbab7997f0ceb1",
+      "source_bias_sha256": "a1bece05c6b61ebba9379c39bc81d6a19fc0cb2fd0543c9a0faccd926f598a64"
+    },
+    {
+      "module_key": "transformer_blocks.0.txt_mlp.net.0.proj",
+      "weight_key": "transformer_blocks.0.txt_mlp.net.0.proj.weight",
+      "bias_key": "transformer_blocks.0.txt_mlp.net.0.proj.bias",
+      "weight_shape": [
+        12288,
+        3072
+      ],
+      "bias_shape": [
+        12288
+      ],
+      "weight_tensor_scale_key": "targets.transformer_blocks.0.txt_mlp.net.0.proj.weight_tensor_scale",
+      "artifact_weight_key": "targets.transformer_blocks.0.txt_mlp.net.0.proj.packed_weight_e2m1",
+      "artifact_scale_key": "targets.transformer_blocks.0.txt_mlp.net.0.proj.packed_scales_ue4m3",
+      "artifact_bias_key": "targets.transformer_blocks.0.txt_mlp.net.0.proj.bias_bf16",
+      "weight_tensor_scale": 0.00024704704992473125,
+      "weight_global_amax": 0.6640625,
+      "packed_weight_bytes": 18874368,
+      "packed_scale_bytes": 2359296,
+      "scale_layout": {
+        "inner_dim": 192,
+        "outer_tiles": 96,
+        "bytes": 2359296
+      },
+      "source_weight_sha256": "40a2c2bf74fa661e51bf3294462ec6901e5f57abd1b0c8d8e3f7363c475aa33a",
+      "source_bias_sha256": "f181d4fef402d8369e38ce9f64fe91903a5002f503b6dde71120c0ecb873d21b"
+    },
+    {
+      "module_key": "transformer_blocks.0.txt_mlp.net.2",
+      "weight_key": "transformer_blocks.0.txt_mlp.net.2.weight",
+      "bias_key": "transformer_blocks.0.txt_mlp.net.2.bias",
+      "weight_shape": [
+        3072,
+        12288
+      ],
+      "bias_shape": [
+        3072
+      ],
+      "weight_tensor_scale_key": "targets.transformer_blocks.0.txt_mlp.net.2.weight_tensor_scale",
+      "artifact_weight_key": "targets.transformer_blocks.0.txt_mlp.net.2.packed_weight_e2m1",
+      "artifact_scale_key": "targets.transformer_blocks.0.txt_mlp.net.2.packed_scales_ue4m3",
+      "artifact_bias_key": "targets.transformer_blocks.0.txt_mlp.net.2.bias_bf16",
+      "weight_tensor_scale": 0.00039527530316263437,
+      "weight_global_amax": 1.0625,
+      "packed_weight_bytes": 18874368,
+      "packed_scale_bytes": 2359296,
+      "scale_layout": {
+        "inner_dim": 768,
+        "outer_tiles": 24,
+        "bytes": 2359296
+      },
+      "source_weight_sha256": "37d5796891212472bfa97ce611f85266688107d579ebe1116eec5a958125152f",
+      "source_bias_sha256": "862c6c05b8cfc438dd7074299f2303b588af5ae52d94b1dcb3e9e58bea78c770"
+    },
+    {
+      "module_key": "transformer_blocks.1.img_mlp.net.0.proj",
+      "weight_key": "transformer_blocks.1.img_mlp.net.0.proj.weight",
+      "bias_key": "transformer_blocks.1.img_mlp.net.0.proj.bias",
+      "weight_shape": [
+        12288,
+        3072
+      ],
+      "bias_shape": [
+        12288
+      ],
+      "weight_tensor_scale_key": "targets.transformer_blocks.1.img_mlp.net.0.proj.weight_tensor_scale",
+      "artifact_weight_key": "targets.transformer_blocks.1.img_mlp.net.0.proj.packed_weight_e2m1",
+      "artifact_scale_key": "targets.transformer_blocks.1.img_mlp.net.0.proj.packed_scales_ue4m3",
+      "artifact_bias_key": "targets.transformer_blocks.1.img_mlp.net.0.proj.bias_bf16",
+      "weight_tensor_scale": 0.00022815522970631719,
+      "weight_global_amax": 0.61328125,
+      "packed_weight_bytes": 18874368,
+      "packed_scale_bytes": 2359296,
+      "scale_layout": {
+        "inner_dim": 192,
+        "outer_tiles": 96,
+        "bytes": 2359296
+      },
+      "source_weight_sha256": "1ba75e81f6d7348769dae4a585d3127592fd2f402c101e3ddc88766f92d27617",
+      "source_bias_sha256": "76d85e6964ca67fd92a1e6907ab4b5d7b839caa55de787b6e4a2969c234fcec9"
+    },
+    {
+      "module_key": "transformer_blocks.1.img_mlp.net.2",
+      "weight_key": "transformer_blocks.1.img_mlp.net.2.weight",
+      "bias_key": "transformer_blocks.1.img_mlp.net.2.bias",
+      "weight_shape": [
+        3072,
+        12288
+      ],
+      "bias_shape": [
+        3072
+      ],
+      "weight_tensor_scale_key": "targets.transformer_blocks.1.img_mlp.net.2.weight_tensor_scale",
+      "artifact_weight_key": "targets.transformer_blocks.1.img_mlp.net.2.packed_weight_e2m1",
+      "artifact_scale_key": "targets.transformer_blocks.1.img_mlp.net.2.packed_scales_ue4m3",
+      "artifact_bias_key": "targets.transformer_blocks.1.img_mlp.net.2.bias_bf16",
+      "weight_tensor_scale": 0.0004417782765813172,
+      "weight_global_amax": 1.1875,
+      "packed_weight_bytes": 18874368,
+      "packed_scale_bytes": 2359296,
+      "scale_layout": {
+        "inner_dim": 768,
+        "outer_tiles": 24,
+        "bytes": 2359296
+      },
+      "source_weight_sha256": "3f890473a9a99aa845e698713a852464055645b061b4b6ab61de59273b87be1b",
+      "source_bias_sha256": "06b80199d80cc3916a202cfe220f59e8572dd406b0d2be8b17d59c1a3de98619"
+    },
+    {
+      "module_key": "transformer_blocks.1.txt_mlp.net.0.proj",
+      "weight_key": "transformer_blocks.1.txt_mlp.net.0.proj.weight",
+      "bias_key": "transformer_blocks.1.txt_mlp.net.0.proj.bias",
+      "weight_shape": [
+        12288,
+        3072
+      ],
+      "bias_shape": [
+        12288
+      ],
+      "weight_tensor_scale_key": "targets.transformer_blocks.1.txt_mlp.net.0.proj.weight_tensor_scale",
+      "artifact_weight_key": "targets.transformer_blocks.1.txt_mlp.net.0.proj.packed_weight_e2m1",
+      "artifact_scale_key": "targets.transformer_blocks.1.txt_mlp.net.0.proj.packed_scales_ue4m3",
+      "artifact_bias_key": "targets.transformer_blocks.1.txt_mlp.net.0.proj.bias_bf16",
+      "weight_tensor_scale": 0.00019618442456703633,
+      "weight_global_amax": 0.52734375,
+      "packed_weight_bytes": 18874368,
+      "packed_scale_bytes": 2359296,
+      "scale_layout": {
+        "inner_dim": 192,
+        "outer_tiles": 96,
+        "bytes": 2359296
+      },
+      "source_weight_sha256": "4a232bd2fdbf09232bd512cfb76c59502068529105e596c821bafff89c19813b",
+      "source_bias_sha256": "e94fb97549fc326506f37a6824ead6a538828909466a6bf76933dfe5ce73680d"
+    },
+    {
+      "module_key": "transformer_blocks.1.txt_mlp.net.2",
+      "weight_key": "transformer_blocks.1.txt_mlp.net.2.weight",
+      "bias_key": "transformer_blocks.1.txt_mlp.net.2.bias",
+      "weight_shape": [
+        3072,
+        12288
+      ],
+      "bias_shape": [
+        3072
+      ],
+      "weight_tensor_scale_key": "targets.transformer_blocks.1.txt_mlp.net.2.weight_tensor_scale",
+      "artifact_weight_key": "targets.transformer_blocks.1.txt_mlp.net.2.packed_weight_e2m1",
+      "artifact_scale_key": "targets.transformer_blocks.1.txt_mlp.net.2.packed_scales_ue4m3",
+      "artifact_bias_key": "targets.transformer_blocks.1.txt_mlp.net.2.bias_bf16",
+      "weight_tensor_scale": 0.00039818172808736563,
+      "weight_global_amax": 1.0703125,
+      "packed_weight_bytes": 18874368,
+      "packed_scale_bytes": 2359296,
+      "scale_layout": {
+        "inner_dim": 768,
+        "outer_tiles": 24,
+        "bytes": 2359296
+      },
+      "source_weight_sha256": "042d74971fba414ffce050de7a65e1bebc33a2e49339b26f42518a9a9621701f",
+      "source_bias_sha256": "faa59706a19669dc45d820edab49feca78c6b1d9d90aca4adbd87040c36f2fea"
+    },
+    {
+      "module_key": "transformer_blocks.2.img_mlp.net.0.proj",
+      "weight_key": "transformer_blocks.2.img_mlp.net.0.proj.weight",
+      "bias_key": "transformer_blocks.2.img_mlp.net.0.proj.bias",
+      "weight_shape": [
+        12288,
+        3072
+      ],
+      "bias_shape": [
+        12288
+      ],
+      "weight_tensor_scale_key": "targets.transformer_blocks.2.img_mlp.net.0.proj.weight_tensor_scale",
+      "artifact_weight_key": "targets.transformer_blocks.2.img_mlp.net.0.proj.packed_weight_e2m1",
+      "artifact_scale_key": "targets.transformer_blocks.2.img_mlp.net.0.proj.packed_scales_ue4m3",
+      "artifact_bias_key": "targets.transformer_blocks.2.img_mlp.net.0.proj.bias_bf16",
+      "weight_tensor_scale": 0.00030226935632526875,
+      "weight_global_amax": 0.8125,
+      "packed_weight_bytes": 18874368,
+      "packed_scale_bytes": 2359296,
+      "scale_layout": {
+        "inner_dim": 192,
+        "outer_tiles": 96,
+        "bytes": 2359296
+      },
+      "source_weight_sha256": "fceb9a0d7b4bfbddb10490b9d9e7f41c369e691bbafdccb003f8dda509fd3b1c",
+      "source_bias_sha256": "84af366af51b7cd24dac8a9fda557a2c8a66ca80e9f8621ec4729d31d64bd241"
+    },
+    {
+      "module_key": "transformer_blocks.2.img_mlp.net.2",
+      "weight_key": "transformer_blocks.2.img_mlp.net.2.weight",
+      "bias_key": "transformer_blocks.2.img_mlp.net.2.bias",
+      "weight_shape": [
+        3072,
+        12288
+      ],
+      "bias_shape": [
+        3072
+      ],
+      "weight_tensor_scale_key": "targets.transformer_blocks.2.img_mlp.net.2.weight_tensor_scale",
+      "artifact_weight_key": "targets.transformer_blocks.2.img_mlp.net.2.packed_weight_e2m1",
+      "artifact_scale_key": "targets.transformer_blocks.2.img_mlp.net.2.packed_scales_ue4m3",
+      "artifact_bias_key": "targets.transformer_blocks.2.img_mlp.net.2.bias_bf16",
+      "weight_tensor_scale": 0.0005057198577560484,
+      "weight_global_amax": 1.359375,
+      "packed_weight_bytes": 18874368,
+      "packed_scale_bytes": 2359296,
+      "scale_layout": {
+        "inner_dim": 768,
+        "outer_tiles": 24,
+        "bytes": 2359296
+      },
+      "source_weight_sha256": "97220e0a04629f28c4ca39abfe46c07cf6183f660f391cd635be2871d325a09d",
+      "source_bias_sha256": "6bcc96109a2496ee73d6a04bb8016641e90d568e6e46ea1a2e8b9e84925bfab9"
+    },
+    {
+      "module_key": "transformer_blocks.2.txt_mlp.net.0.proj",
+      "weight_key": "transformer_blocks.2.txt_mlp.net.0.proj.weight",
+      "bias_key": "transformer_blocks.2.txt_mlp.net.0.proj.bias",
+      "weight_shape": [
+        12288,
+        3072
+      ],
+      "bias_shape": [
+        12288
+      ],
+      "weight_tensor_scale_key": "targets.transformer_blocks.2.txt_mlp.net.0.proj.weight_tensor_scale",
+      "artifact_weight_key": "targets.transformer_blocks.2.txt_mlp.net.0.proj.packed_weight_e2m1",
+      "artifact_scale_key": "targets.transformer_blocks.2.txt_mlp.net.0.proj.packed_scales_ue4m3",
+      "artifact_bias_key": "targets.transformer_blocks.2.txt_mlp.net.0.proj.bias_bf16",
+      "weight_tensor_scale": 0.00022815522970631719,
+      "weight_global_amax": 0.61328125,
+      "packed_weight_bytes": 18874368,
+      "packed_scale_bytes": 2359296,
+      "scale_layout": {
+        "inner_dim": 192,
+        "outer_tiles": 96,
+        "bytes": 2359296
+      },
+      "source_weight_sha256": "764c79424da4eea44e62f03c1fda648bbd937ebf0b7c0bb43bede1e8b5130f29",
+      "source_bias_sha256": "eca2f24456368d090da90c98a102af788833b3f3f9130f29b5b7a734e36f307e"
+    },
+    {
+      "module_key": "transformer_blocks.2.txt_mlp.net.2",
+      "weight_key": "transformer_blocks.2.txt_mlp.net.2.weight",
+      "bias_key": "transformer_blocks.2.txt_mlp.net.2.bias",
+      "weight_shape": [
+        3072,
+        12288
+      ],
+      "bias_shape": [
+        3072
+      ],
+      "weight_tensor_scale_key": "targets.transformer_blocks.2.txt_mlp.net.2.weight_tensor_scale",
+      "artifact_weight_key": "targets.transformer_blocks.2.txt_mlp.net.2.packed_weight_e2m1",
+      "artifact_scale_key": "targets.transformer_blocks.2.txt_mlp.net.2.packed_scales_ue4m3",
+      "artifact_bias_key": "targets.transformer_blocks.2.txt_mlp.net.2.bias_bf16",
+      "weight_tensor_scale": 0.0003589448460843414,
+      "weight_global_amax": 0.96484375,
+      "packed_weight_bytes": 18874368,
+      "packed_scale_bytes": 2359296,
+      "scale_layout": {
+        "inner_dim": 768,
+        "outer_tiles": 24,
+        "bytes": 2359296
+      },
+      "source_weight_sha256": "17e543841d93e1936514c8313a85e4b9444321db06111c33460173e47b0d1c93",
+      "source_bias_sha256": "eed98a4db5990a369a55eb0865482f174e1418e6f4231bcf65637939a1df8354"
+    },
+    {
+      "module_key": "transformer_blocks.3.img_mlp.net.0.proj",
+      "weight_key": "transformer_blocks.3.img_mlp.net.0.proj.weight",
+      "bias_key": "transformer_blocks.3.img_mlp.net.0.proj.bias",
+      "weight_shape": [
+        12288,
+        3072
+      ],
+      "bias_shape": [
+        12288
+      ],
+      "weight_tensor_scale_key": "targets.transformer_blocks.3.img_mlp.net.0.proj.weight_tensor_scale",
+      "artifact_weight_key": "targets.transformer_blocks.3.img_mlp.net.0.proj.packed_weight_e2m1",
+      "artifact_scale_key": "targets.transformer_blocks.3.img_mlp.net.0.proj.packed_scales_ue4m3",
+      "artifact_bias_key": "targets.transformer_blocks.3.img_mlp.net.0.proj.bias_bf16",
+      "weight_tensor_scale": 0.00024850029149092734,
+      "weight_global_amax": 0.66796875,
+      "packed_weight_bytes": 18874368,
+      "packed_scale_bytes": 2359296,
+      "scale_layout": {
+        "inner_dim": 192,
+        "outer_tiles": 96,
+        "bytes": 2359296
+      },
+      "source_weight_sha256": "83c39bc7dd86d7e6a5f63eadc063b91d7bfd93a75a80c9536ddf28364d35ca6c",
+      "source_bias_sha256": "4f40baf432bc9106a645e1182bce1c18535e1370f0851cb0f6e21d3a1646a08e"
+    },
+    {
+      "module_key": "transformer_blocks.3.img_mlp.net.2",
+      "weight_key": "transformer_blocks.3.img_mlp.net.2.weight",
+      "bias_key": "transformer_blocks.3.img_mlp.net.2.bias",
+      "weight_shape": [
+        3072,
+        12288
+      ],
+      "bias_shape": [
+        3072
+      ],
+      "weight_tensor_scale_key": "targets.transformer_blocks.3.img_mlp.net.2.weight_tensor_scale",
+      "artifact_weight_key": "targets.transformer_blocks.3.img_mlp.net.2.packed_weight_e2m1",
+      "artifact_scale_key": "targets.transformer_blocks.3.img_mlp.net.2.packed_scales_ue4m3",
+      "artifact_bias_key": "targets.transformer_blocks.3.img_mlp.net.2.bias_bf16",
+      "weight_tensor_scale": 0.0005871000466868281,
+      "weight_global_amax": 1.578125,
+      "packed_weight_bytes": 18874368,
+      "packed_scale_bytes": 2359296,
+      "scale_layout": {
+        "inner_dim": 768,
+        "outer_tiles": 24,
+        "bytes": 2359296
+      },
+      "source_weight_sha256": "e6945585cd816ccb2cd6e4a005326a105758f72270840842dea9c8617ba095ba",
+      "source_bias_sha256": "230dcffbe97056ccb1975e1dbdc4274a89a217c64b240dcb70d1a639c8d50e10"
+    },
+    {
+      "module_key": "transformer_blocks.3.txt_mlp.net.0.proj",
+      "weight_key": "transformer_blocks.3.txt_mlp.net.0.proj.weight",
+      "bias_key": "transformer_blocks.3.txt_mlp.net.0.proj.bias",
+      "weight_shape": [
+        12288,
+        3072
+      ],
+      "bias_shape": [
+        12288
+      ],
+      "weight_tensor_scale_key": "targets.transformer_blocks.3.txt_mlp.net.0.proj.weight_tensor_scale",
+      "artifact_weight_key": "targets.transformer_blocks.3.txt_mlp.net.0.proj.packed_weight_e2m1",
+      "artifact_scale_key": "targets.transformer_blocks.3.txt_mlp.net.0.proj.packed_scales_ue4m3",
+      "artifact_bias_key": "targets.transformer_blocks.3.txt_mlp.net.0.proj.bias_bf16",
+      "weight_tensor_scale": 0.00022960844216868281,
+      "weight_global_amax": 0.6171875,
+      "packed_weight_bytes": 18874368,
+      "packed_scale_bytes": 2359296,
+      "scale_layout": {
+        "inner_dim": 192,
+        "outer_tiles": 96,
+        "bytes": 2359296
+      },
+      "source_weight_sha256": "786d08fafeb3519a0ab99e194105dfae1135ebbed18ae364fd62ba323c6525a2",
+      "source_bias_sha256": "f493d8f0a032e8a1158ffc66ca65c00a78bdf8118d3bc307c4bab75605979d67"
+    },
+    {
+      "module_key": "transformer_blocks.3.txt_mlp.net.2",
+      "weight_key": "transformer_blocks.3.txt_mlp.net.2.weight",
+      "bias_key": "transformer_blocks.3.txt_mlp.net.2.bias",
+      "weight_shape": [
+        3072,
+        12288
+      ],
+      "bias_shape": [
+        3072
+      ],
+      "weight_tensor_scale_key": "targets.transformer_blocks.3.txt_mlp.net.2.weight_tensor_scale",
+      "artifact_weight_key": "targets.transformer_blocks.3.txt_mlp.net.2.packed_weight_e2m1",
+      "artifact_scale_key": "targets.transformer_blocks.3.txt_mlp.net.2.packed_scales_ue4m3",
+      "artifact_bias_key": "targets.transformer_blocks.3.txt_mlp.net.2.bias_bf16",
+      "weight_tensor_scale": 0.0004039946070406586,
+      "weight_global_amax": 1.0859375,
+      "packed_weight_bytes": 18874368,
+      "packed_scale_bytes": 2359296,
+      "scale_layout": {
+        "inner_dim": 768,
+        "outer_tiles": 24,
+        "bytes": 2359296
+      },
+      "source_weight_sha256": "ba11326b827a9600c9d0ae9227f2d3c0ab9971a09afd7df18689f4f1436a4e76",
+      "source_bias_sha256": "c8a38cf2ae33d7b353cd71510a74ede4727b418285ea6dd5934274dcf5bbdf4f"
+    },
+    {
+      "module_key": "transformer_blocks.4.img_mlp.net.0.proj",
+      "weight_key": "transformer_blocks.4.img_mlp.net.0.proj.weight",
+      "bias_key": "transformer_blocks.4.img_mlp.net.0.proj.bias",
+      "weight_shape": [
+        12288,
+        3072
+      ],
+      "bias_shape": [
+        12288
+      ],
+      "weight_tensor_scale_key": "targets.transformer_blocks.4.img_mlp.net.0.proj.weight_tensor_scale",
+      "artifact_weight_key": "targets.transformer_blocks.4.img_mlp.net.0.proj.packed_weight_e2m1",
+      "artifact_scale_key": "targets.transformer_blocks.4.img_mlp.net.0.proj.packed_scales_ue4m3",
+      "artifact_bias_key": "targets.transformer_blocks.4.img_mlp.net.0.proj.bias_bf16",
+      "weight_tensor_scale": 0.0003516787546686828,
+      "weight_global_amax": 0.9453125,
+      "packed_weight_bytes": 18874368,
+      "packed_scale_bytes": 2359296,
+      "scale_layout": {
+        "inner_dim": 192,
+        "outer_tiles": 96,
+        "bytes": 2359296
+      },
+      "source_weight_sha256": "2da64d233d618c639d7ac3a988c045771a20bb77634d272949dc1f9d3a26e6fa",
+      "source_bias_sha256": "6079037543d2b43a319f02021a6f71d2874efaaafd6a5a58f2813a9e160a9ca6"
+    },
+    {
+      "module_key": "transformer_blocks.4.img_mlp.net.2",
+      "weight_key": "transformer_blocks.4.img_mlp.net.2.weight",
+      "bias_key": "transformer_blocks.4.img_mlp.net.2.bias",
+      "weight_shape": [
+        3072,
+        12288
+      ],
+      "bias_shape": [
+        3072
+      ],
+      "weight_tensor_scale_key": "targets.transformer_blocks.4.img_mlp.net.2.weight_tensor_scale",
+      "artifact_weight_key": "targets.transformer_blocks.4.img_mlp.net.2.packed_weight_e2m1",
+      "artifact_scale_key": "targets.transformer_blocks.4.img_mlp.net.2.packed_scales_ue4m3",
+      "artifact_bias_key": "targets.transformer_blocks.4.img_mlp.net.2.bias_bf16",
+      "weight_tensor_scale": 0.0009533109841868281,
+      "weight_global_amax": 2.5625,
+      "packed_weight_bytes": 18874368,
+      "packed_scale_bytes": 2359296,
+      "scale_layout": {
+        "inner_dim": 768,
+        "outer_tiles": 24,
+        "bytes": 2359296
+      },
+      "source_weight_sha256": "4db777bdb0beb299cb30ccc2043870bc4f311d5e6e30b6ad28e1cc7c29af96ed",
+      "source_bias_sha256": "d4e4ac8e94b3fd245c412ee58d3fc91e815fd03e49684d8596ac58566145f1fe"
+    },
+    {
+      "module_key": "transformer_blocks.4.txt_mlp.net.0.proj",
+      "weight_key": "transformer_blocks.4.txt_mlp.net.0.proj.weight",
+      "bias_key": "transformer_blocks.4.txt_mlp.net.0.proj.bias",
+      "weight_shape": [
+        12288,
+        3072
+      ],
+      "bias_shape": [
+        12288
+      ],
+      "weight_tensor_scale_key": "targets.transformer_blocks.4.txt_mlp.net.0.proj.weight_tensor_scale",
+      "artifact_weight_key": "targets.transformer_blocks.4.txt_mlp.net.0.proj.packed_weight_e2m1",
+      "artifact_scale_key": "targets.transformer_blocks.4.txt_mlp.net.0.proj.packed_scales_ue4m3",
+      "artifact_bias_key": "targets.transformer_blocks.4.txt_mlp.net.0.proj.bias_bf16",
+      "weight_tensor_scale": 0.00018528530199546367,
+      "weight_global_amax": 0.498046875,
+      "packed_weight_bytes": 18874368,
+      "packed_scale_bytes": 2359296,
+      "scale_layout": {
+        "inner_dim": 192,
+        "outer_tiles": 96,
+        "bytes": 2359296
+      },
+      "source_weight_sha256": "0c956820e6938c52af5a237da6dcd4c41b509865b7631407a896cc6da6bdb985",
+      "source_bias_sha256": "edd0dfc9060a78b09c53bd0f38a9d0a5263faadbe3c522a7022115a2e09c211e"
+    },
+    {
+      "module_key": "transformer_blocks.4.txt_mlp.net.2",
+      "weight_key": "transformer_blocks.4.txt_mlp.net.2.weight",
+      "bias_key": "transformer_blocks.4.txt_mlp.net.2.bias",
+      "weight_shape": [
+        3072,
+        12288
+      ],
+      "bias_shape": [
+        3072
+      ],
+      "weight_tensor_scale_key": "targets.transformer_blocks.4.txt_mlp.net.2.weight_tensor_scale",
+      "artifact_weight_key": "targets.transformer_blocks.4.txt_mlp.net.2.packed_weight_e2m1",
+      "artifact_scale_key": "targets.transformer_blocks.4.txt_mlp.net.2.packed_scales_ue4m3",
+      "artifact_bias_key": "targets.transformer_blocks.4.txt_mlp.net.2.bias_bf16",
+      "weight_tensor_scale": 0.0003284272679593414,
+      "weight_global_amax": 0.8828125,
+      "packed_weight_bytes": 18874368,
+      "packed_scale_bytes": 2359296,
+      "scale_layout": {
+        "inner_dim": 768,
+        "outer_tiles": 24,
+        "bytes": 2359296
+      },
+      "source_weight_sha256": "005cd523afdeb3e9d32a2b29dd0499fa994fd8a79af245604661ee8bc1af59c8",
+      "source_bias_sha256": "f333cc3ca7d6299b5cf44859662914d0fae7cbeba0433e0ae487d64cdc755685"
+    },
+    {
+      "module_key": "transformer_blocks.5.img_mlp.net.0.proj",
+      "weight_key": "transformer_blocks.5.img_mlp.net.0.proj.weight",
+      "bias_key": "transformer_blocks.5.img_mlp.net.0.proj.bias",
+      "weight_shape": [
+        12288,
+        3072
+      ],
+      "bias_shape": [
+        12288
+      ],
+      "weight_tensor_scale_key": "targets.transformer_blocks.5.img_mlp.net.0.proj.weight_tensor_scale",
+      "artifact_weight_key": "targets.transformer_blocks.5.img_mlp.net.0.proj.packed_weight_e2m1",
+      "artifact_scale_key": "targets.transformer_blocks.5.img_mlp.net.0.proj.packed_scales_ue4m3",
+      "artifact_bias_key": "targets.transformer_blocks.5.img_mlp.net.0.proj.bias_bf16",
+      "weight_tensor_scale": 0.00040108818211592734,
+      "weight_global_amax": 1.078125,
+      "packed_weight_bytes": 18874368,
+      "packed_scale_bytes": 2359296,
+      "scale_layout": {
+        "inner_dim": 192,
+        "outer_tiles": 96,
+        "bytes": 2359296
+      },
+      "source_weight_sha256": "8419f255277f3a77facc9090e6882ab54d41e6f9f0bebb7930416e07c7276f55",
+      "source_bias_sha256": "c738466822f96ba14226ee1f254ec83099479436e1ff34d7ae220fdc7bad423d"
+    },
+    {
+      "module_key": "transformer_blocks.5.img_mlp.net.2",
+      "weight_key": "transformer_blocks.5.img_mlp.net.2.weight",
+      "bias_key": "transformer_blocks.5.img_mlp.net.2.bias",
+      "weight_shape": [
+        3072,
+        12288
+      ],
+      "bias_shape": [
+        3072
+      ],
+      "weight_tensor_scale_key": "targets.transformer_blocks.5.img_mlp.net.2.weight_tensor_scale",
+      "artifact_weight_key": "targets.transformer_blocks.5.img_mlp.net.2.packed_weight_e2m1",
+      "artifact_scale_key": "targets.transformer_blocks.5.img_mlp.net.2.packed_scales_ue4m3",
+      "artifact_bias_key": "targets.transformer_blocks.5.img_mlp.net.2.bias_bf16",
+      "weight_tensor_scale": 0.0006684802938252687,
+      "weight_global_amax": 1.796875,
+      "packed_weight_bytes": 18874368,
+      "packed_scale_bytes": 2359296,
+      "scale_layout": {
+        "inner_dim": 768,
+        "outer_tiles": 24,
+        "bytes": 2359296
+      },
+      "source_weight_sha256": "97fbe5a8ee7151217b82e12295e02681af0ca707d3595c58b40aed31c3ae23b7",
+      "source_bias_sha256": "af3f83a94e0368f51c65b26a880b1bf3b64543910e9e270a04faa12aa58680da"
+    },
+    {
+      "module_key": "transformer_blocks.5.txt_mlp.net.0.proj",
+      "weight_key": "transformer_blocks.5.txt_mlp.net.0.proj.weight",
+      "bias_key": "transformer_blocks.5.txt_mlp.net.0.proj.bias",
+      "weight_shape": [
+        12288,
+        3072
+      ],
+      "bias_shape": [
+        12288
+      ],
+      "weight_tensor_scale_key": "targets.transformer_blocks.5.txt_mlp.net.0.proj.weight_tensor_scale",
+      "artifact_weight_key": "targets.transformer_blocks.5.txt_mlp.net.0.proj.packed_weight_e2m1",
+      "artifact_scale_key": "targets.transformer_blocks.5.txt_mlp.net.0.proj.packed_scales_ue4m3",
+      "artifact_bias_key": "targets.transformer_blocks.5.txt_mlp.net.0.proj.bias_bf16",
+      "weight_tensor_scale": 0.0001874651206890121,
+      "weight_global_amax": 0.50390625,
+      "packed_weight_bytes": 18874368,
+      "packed_scale_bytes": 2359296,
+      "scale_layout": {
+        "inner_dim": 192,
+        "outer_tiles": 96,
+        "bytes": 2359296
+      },
+      "source_weight_sha256": "7e3b210a8d107962865efa3f3974a9ea29a60732f8002539e71837ebf9fb3787",
+      "source_bias_sha256": "a154c2419b8da86533eeb0e5a69c37cbfe5eff859f8196e192d5f942743c690d"
+    },
+    {
+      "module_key": "transformer_blocks.5.txt_mlp.net.2",
+      "weight_key": "transformer_blocks.5.txt_mlp.net.2.weight",
+      "bias_key": "transformer_blocks.5.txt_mlp.net.2.bias",
+      "weight_shape": [
+        3072,
+        12288
+      ],
+      "bias_shape": [
+        3072
+      ],
+      "weight_tensor_scale_key": "targets.transformer_blocks.5.txt_mlp.net.2.weight_tensor_scale",
+      "artifact_weight_key": "targets.transformer_blocks.5.txt_mlp.net.2.packed_weight_e2m1",
+      "artifact_scale_key": "targets.transformer_blocks.5.txt_mlp.net.2.packed_scales_ue4m3",
+      "artifact_bias_key": "targets.transformer_blocks.5.txt_mlp.net.2.bias_bf16",
+      "weight_tensor_scale": 0.0003807431203313172,
+      "weight_global_amax": 1.0234375,
+      "packed_weight_bytes": 18874368,
+      "packed_scale_bytes": 2359296,
+      "scale_layout": {
+        "inner_dim": 768,
+        "outer_tiles": 24,
+        "bytes": 2359296
+      },
+      "source_weight_sha256": "906a7d3e7d32c10037e731b32904f683ff3529e0605dc619132c6095e87aeec4",
+      "source_bias_sha256": "e47066d5093c8d5df144ad237d0c0cbb8186b7c176764338b777873e623e226e"
+    },
+    {
+      "module_key": "transformer_blocks.6.img_mlp.net.0.proj",
+      "weight_key": "transformer_blocks.6.img_mlp.net.0.proj.weight",
+      "bias_key": "transformer_blocks.6.img_mlp.net.0.proj.bias",
+      "weight_shape": [
+        12288,
+        3072
+      ],
+      "bias_shape": [
+        12288
+      ],
+      "weight_tensor_scale_key": "targets.transformer_blocks.6.img_mlp.net.0.proj.weight_tensor_scale",
+      "artifact_weight_key": "targets.transformer_blocks.6.img_mlp.net.0.proj.packed_weight_e2m1",
+      "artifact_scale_key": "targets.transformer_blocks.6.img_mlp.net.0.proj.packed_scales_ue4m3",
+      "artifact_bias_key": "targets.transformer_blocks.6.img_mlp.net.0.proj.bias_bf16",
+      "weight_tensor_scale": 0.00043015251867473125,
+      "weight_global_amax": 1.15625,
+      "packed_weight_bytes": 18874368,
+      "packed_scale_bytes": 2359296,
+      "scale_layout": {
+        "inner_dim": 192,
+        "outer_tiles": 96,
+        "bytes": 2359296
+      },
+      "source_weight_sha256": "6e2eaa52c7e036559066e6d0bd3e259b8b4a9e6010cfb5201d255aa673a95e9a",
+      "source_bias_sha256": "559dcbcbecdda86591182bdad53d03f173400d6abbdada125766b5c5e180887f"
+    },
+    {
+      "module_key": "transformer_blocks.6.img_mlp.net.2",
+      "weight_key": "transformer_blocks.6.img_mlp.net.2.weight",
+      "bias_key": "transformer_blocks.6.img_mlp.net.2.bias",
+      "weight_shape": [
+        3072,
+        12288
+      ],
+      "bias_shape": [
+        3072
+      ],
+      "weight_tensor_scale_key": "targets.transformer_blocks.6.img_mlp.net.2.weight_tensor_scale",
+      "artifact_weight_key": "targets.transformer_blocks.6.img_mlp.net.2.packed_weight_e2m1",
+      "artifact_scale_key": "targets.transformer_blocks.6.img_mlp.net.2.packed_scales_ue4m3",
+      "artifact_bias_key": "targets.transformer_blocks.6.img_mlp.net.2.bias_bf16",
+      "weight_tensor_scale": 0.0007672990905120969,
+      "weight_global_amax": 2.0625,
+      "packed_weight_bytes": 18874368,
+      "packed_scale_bytes": 2359296,
+      "scale_layout": {
+        "inner_dim": 768,
+        "outer_tiles": 24,
+        "bytes": 2359296
+      },
+      "source_weight_sha256": "5d161c4e9f626a44cd150ccd49e4b6a3b91dfe9c78628edbe2d05bcfd205ee38",
+      "source_bias_sha256": "88e2f67cc46733f0a734a37bf6d6f6a52b12bbbfce8346a9ca81b4942a3c5c89"
+    },
+    {
+      "module_key": "transformer_blocks.6.txt_mlp.net.0.proj",
+      "weight_key": "transformer_blocks.6.txt_mlp.net.0.proj.weight",
+      "bias_key": "transformer_blocks.6.txt_mlp.net.0.proj.bias",
+      "weight_shape": [
+        12288,
+        3072
+      ],
+      "bias_shape": [
+        12288
+      ],
+      "weight_tensor_scale_key": "targets.transformer_blocks.6.txt_mlp.net.0.proj.weight_tensor_scale",
+      "artifact_weight_key": "targets.transformer_blocks.6.txt_mlp.net.0.proj.packed_weight_e2m1",
+      "artifact_scale_key": "targets.transformer_blocks.6.txt_mlp.net.0.proj.packed_scales_ue4m3",
+      "artifact_bias_key": "targets.transformer_blocks.6.txt_mlp.net.0.proj.bias_bf16",
+      "weight_tensor_scale": 0.0002601260202936828,
+      "weight_global_amax": 0.69921875,
+      "packed_weight_bytes": 18874368,
+      "packed_scale_bytes": 2359296,
+      "scale_layout": {
+        "inner_dim": 192,
+        "outer_tiles": 96,
+        "bytes": 2359296
+      },
+      "source_weight_sha256": "fdd4c5e88701022033d97501bd2781fdc526eaa9427775a790756a9bb7c1fda2",
+      "source_bias_sha256": "37f6d759de4dfb875a8c7076fae88f75c2c5fdade293907895fead213e7a0ae7"
+    },
+    {
+      "module_key": "transformer_blocks.6.txt_mlp.net.2",
+      "weight_key": "transformer_blocks.6.txt_mlp.net.2.weight",
+      "bias_key": "transformer_blocks.6.txt_mlp.net.2.bias",
+      "weight_shape": [
+        3072,
+        12288
+      ],
+      "bias_shape": [
+        3072
+      ],
+      "weight_tensor_scale_key": "targets.transformer_blocks.6.txt_mlp.net.2.weight_tensor_scale",
+      "artifact_weight_key": "targets.transformer_blocks.6.txt_mlp.net.2.packed_weight_e2m1",
+      "artifact_scale_key": "targets.transformer_blocks.6.txt_mlp.net.2.packed_scales_ue4m3",
+      "artifact_bias_key": "targets.transformer_blocks.6.txt_mlp.net.2.bias_bf16",
+      "weight_tensor_scale": 0.00047665549209341407,
+      "weight_global_amax": 1.28125,
+      "packed_weight_bytes": 18874368,
+      "packed_scale_bytes": 2359296,
+      "scale_layout": {
+        "inner_dim": 768,
+        "outer_tiles": 24,
+        "bytes": 2359296
+      },
+      "source_weight_sha256": "b28669406ee0f0e295c1a0b19241c1de874952807a0754d31078b56b1e213f2d",
+      "source_bias_sha256": "a5479cb1785b82d01897ca73a0b560b1382dcb7b60c8dce315e247d795e04d0f"
+    },
+    {
+      "module_key": "transformer_blocks.7.img_mlp.net.0.proj",
+      "weight_key": "transformer_blocks.7.img_mlp.net.0.proj.weight",
+      "bias_key": "transformer_blocks.7.img_mlp.net.0.proj.bias",
+      "weight_shape": [
+        12288,
+        3072
+      ],
+      "bias_shape": [
+        12288
+      ],
+      "weight_tensor_scale_key": "targets.transformer_blocks.7.img_mlp.net.0.proj.weight_tensor_scale",
+      "artifact_weight_key": "targets.transformer_blocks.7.img_mlp.net.0.proj.packed_weight_e2m1",
+      "artifact_scale_key": "targets.transformer_blocks.7.img_mlp.net.0.proj.packed_scales_ue4m3",
+      "artifact_bias_key": "targets.transformer_blocks.7.img_mlp.net.0.proj.bias_bf16",
+      "weight_tensor_scale": 0.0003138950851280242,
+      "weight_global_amax": 0.84375,
+      "packed_weight_bytes": 18874368,
+      "packed_scale_bytes": 2359296,
+      "scale_layout": {
+        "inner_dim": 192,
+        "outer_tiles": 96,
+        "bytes": 2359296
+      },
+      "source_weight_sha256": "9bcc9591965c0b74076ad9da74307faa9c7fa0ab63a81f8f25960c2286ed39bf",
+      "source_bias_sha256": "b043d399ff8f040405104ffb775c9f902f95dca6de06a6b1ae2f6c90e5f28d7c"
+    },
+    {
+      "module_key": "transformer_blocks.7.img_mlp.net.2",
+      "weight_key": "transformer_blocks.7.img_mlp.net.2.weight",
+      "bias_key": "transformer_blocks.7.img_mlp.net.2.bias",
+      "weight_shape": [
+        3072,
+        12288
+      ],
+      "bias_shape": [
+        3072
+      ],
+      "weight_tensor_scale_key": "targets.transformer_blocks.7.img_mlp.net.2.weight_tensor_scale",
+      "artifact_weight_key": "targets.transformer_blocks.7.img_mlp.net.2.packed_weight_e2m1",
+      "artifact_scale_key": "targets.transformer_blocks.7.img_mlp.net.2.packed_scales_ue4m3",
+      "artifact_bias_key": "targets.transformer_blocks.7.img_mlp.net.2.bias_bf16",
+      "weight_tensor_scale": 0.0008603050373494625,
+      "weight_global_amax": 2.3125,
+      "packed_weight_bytes": 18874368,
+      "packed_scale_bytes": 2359296,
+      "scale_layout": {
+        "inner_dim": 768,
+        "outer_tiles": 24,
+        "bytes": 2359296
+      },
+      "source_weight_sha256": "e7efbd25a434220ac1f0f4f11d8e501b699e11a948a86148460b4139db5ba3bf",
+      "source_bias_sha256": "b2cd1665c361e61f731ce7b474acb2e5f57087c2cacf53f92f418256c5f1ce86"
+    },
+    {
+      "module_key": "transformer_blocks.7.txt_mlp.net.0.proj",
+      "weight_key": "transformer_blocks.7.txt_mlp.net.0.proj.weight",
+      "bias_key": "transformer_blocks.7.txt_mlp.net.0.proj.bias",
+      "weight_shape": [
+        12288,
+        3072
+      ],
+      "bias_shape": [
+        12288
+      ],
+      "weight_tensor_scale_key": "targets.transformer_blocks.7.txt_mlp.net.0.proj.weight_tensor_scale",
+      "artifact_weight_key": "targets.transformer_blocks.7.txt_mlp.net.0.proj.packed_weight_e2m1",
+      "artifact_scale_key": "targets.transformer_blocks.7.txt_mlp.net.0.proj.packed_scales_ue4m3",
+      "artifact_bias_key": "targets.transformer_blocks.7.txt_mlp.net.0.proj.bias_bf16",
+      "weight_tensor_scale": 0.00027901786961592734,
+      "weight_global_amax": 0.75,
+      "packed_weight_bytes": 18874368,
+      "packed_scale_bytes": 2359296,
+      "scale_layout": {
+        "inner_dim": 192,
+        "outer_tiles": 96,
+        "bytes": 2359296
+      },
+      "source_weight_sha256": "cb6126a8544c10ae9f23b3d782c83258a8e700a42ac2ad2e5001099e305f271b",
+      "source_bias_sha256": "ed4294351c27eaf4e5f85bb926daed8e479f3862e67e79142788a57490059768"
+    },
+    {
+      "module_key": "transformer_blocks.7.txt_mlp.net.2",
+      "weight_key": "transformer_blocks.7.txt_mlp.net.2.weight",
+      "bias_key": "transformer_blocks.7.txt_mlp.net.2.bias",
+      "weight_shape": [
+        3072,
+        12288
+      ],
+      "bias_shape": [
+        3072
+      ],
+      "weight_tensor_scale_key": "targets.transformer_blocks.7.txt_mlp.net.2.weight_tensor_scale",
+      "artifact_weight_key": "targets.transformer_blocks.7.txt_mlp.net.2.packed_weight_e2m1",
+      "artifact_scale_key": "targets.transformer_blocks.7.txt_mlp.net.2.packed_scales_ue4m3",
+      "artifact_bias_key": "targets.transformer_blocks.7.txt_mlp.net.2.bias_bf16",
+      "weight_tensor_scale": 0.00039527530316263437,
+      "weight_global_amax": 1.0625,
+      "packed_weight_bytes": 18874368,
+      "packed_scale_bytes": 2359296,
+      "scale_layout": {
+        "inner_dim": 768,
+        "outer_tiles": 24,
+        "bytes": 2359296
+      },
+      "source_weight_sha256": "84caab56f058269c0576a27a2b52ae9945f225ba5b6c158cd3f8b8daf2a135f7",
+      "source_bias_sha256": "cb733e6b65392d1189b7a9d92984c51f8638a1375b6475023b26cd3c362ee7a8"
+    },
+    {
+      "module_key": "transformer_blocks.8.img_mlp.net.0.proj",
+      "weight_key": "transformer_blocks.8.img_mlp.net.0.proj.weight",
+      "bias_key": "transformer_blocks.8.img_mlp.net.0.proj.bias",
+      "weight_shape": [
+        12288,
+        3072
+      ],
+      "bias_shape": [
+        12288
+      ],
+      "weight_tensor_scale_key": "targets.transformer_blocks.8.img_mlp.net.0.proj.weight_tensor_scale",
+      "artifact_weight_key": "targets.transformer_blocks.8.img_mlp.net.0.proj.packed_weight_e2m1",
+      "artifact_scale_key": "targets.transformer_blocks.8.img_mlp.net.0.proj.packed_scales_ue4m3",
+      "artifact_bias_key": "targets.transformer_blocks.8.img_mlp.net.0.proj.bias_bf16",
+      "weight_tensor_scale": 0.00041562033584341407,
+      "weight_global_amax": 1.1171875,
+      "packed_weight_bytes": 18874368,
+      "packed_scale_bytes": 2359296,
+      "scale_layout": {
+        "inner_dim": 192,
+        "outer_tiles": 96,
+        "bytes": 2359296
+      },
+      "source_weight_sha256": "b6da93ff7c47065d4dc8d5b44f02c1f956bb8f373d89526565b2cc1cc3c731ef",
+      "source_bias_sha256": "a155557ac69e102404fcde3489febf8d75d8df81a58f010ae722f73b808c0cf0"
+    },
+    {
+      "module_key": "transformer_blocks.8.img_mlp.net.2",
+      "weight_key": "transformer_blocks.8.img_mlp.net.2.weight",
+      "bias_key": "transformer_blocks.8.img_mlp.net.2.bias",
+      "weight_shape": [
+        3072,
+        12288
+      ],
+      "bias_shape": [
+        3072
+      ],
+      "weight_tensor_scale_key": "targets.transformer_blocks.8.img_mlp.net.2.weight_tensor_scale",
+      "artifact_weight_key": "targets.transformer_blocks.8.img_mlp.net.2.packed_weight_e2m1",
+      "artifact_scale_key": "targets.transformer_blocks.8.img_mlp.net.2.packed_scales_ue4m3",
+      "artifact_bias_key": "targets.transformer_blocks.8.img_mlp.net.2.bias_bf16",
+      "weight_tensor_scale": 0.0005260648904368281,
+      "weight_global_amax": 1.4140625,
+      "packed_weight_bytes": 18874368,
+      "packed_scale_bytes": 2359296,
+      "scale_layout": {
+        "inner_dim": 768,
+        "outer_tiles": 24,
+        "bytes": 2359296
+      },
+      "source_weight_sha256": "928218f79be4de0ca6fd6f855310ad7a91b588be9c970c9ea6a11f55b1a1e413",
+      "source_bias_sha256": "669dfaa63f72b425feecba8fb58b4c43926e96d971f2f2e9241632855c559236"
+    },
+    {
+      "module_key": "transformer_blocks.8.txt_mlp.net.0.proj",
+      "weight_key": "transformer_blocks.8.txt_mlp.net.0.proj.weight",
+      "bias_key": "transformer_blocks.8.txt_mlp.net.0.proj.bias",
+      "weight_shape": [
+        12288,
+        3072
+      ],
+      "bias_shape": [
+        12288
+      ],
+      "weight_tensor_scale_key": "targets.transformer_blocks.8.txt_mlp.net.0.proj.weight_tensor_scale",
+      "artifact_weight_key": "targets.transformer_blocks.8.txt_mlp.net.0.proj.packed_weight_e2m1",
+      "artifact_scale_key": "targets.transformer_blocks.8.txt_mlp.net.0.proj.packed_scales_ue4m3",
+      "artifact_bias_key": "targets.transformer_blocks.8.txt_mlp.net.0.proj.bias_bf16",
+      "weight_tensor_scale": 0.0002601260202936828,
+      "weight_global_amax": 0.69921875,
+      "packed_weight_bytes": 18874368,
+      "packed_scale_bytes": 2359296,
+      "scale_layout": {
+        "inner_dim": 192,
+        "outer_tiles": 96,
+        "bytes": 2359296
+      },
+      "source_weight_sha256": "9342523cc73ba392d60e1dffdfe37824a20e5173281de80f85a4cef4893363db",
+      "source_bias_sha256": "8c13b195dab6195c7654e07b8416da4c952df50f16fd9eccb248c1a38262df10"
+    },
+    {
+      "module_key": "transformer_blocks.8.txt_mlp.net.2",
+      "weight_key": "transformer_blocks.8.txt_mlp.net.2.weight",
+      "bias_key": "transformer_blocks.8.txt_mlp.net.2.bias",
+      "weight_shape": [
+        3072,
+        12288
+      ],
+      "bias_shape": [
+        3072
+      ],
+      "weight_tensor_scale_key": "targets.transformer_blocks.8.txt_mlp.net.2.weight_tensor_scale",
+      "artifact_weight_key": "targets.transformer_blocks.8.txt_mlp.net.2.packed_weight_e2m1",
+      "artifact_scale_key": "targets.transformer_blocks.8.txt_mlp.net.2.packed_scales_ue4m3",
+      "artifact_bias_key": "targets.transformer_blocks.8.txt_mlp.net.2.bias_bf16",
+      "weight_tensor_scale": 0.00042724609375,
+      "weight_global_amax": 1.1484375,
+      "packed_weight_bytes": 18874368,
+      "packed_scale_bytes": 2359296,
+      "scale_layout": {
+        "inner_dim": 768,
+        "outer_tiles": 24,
+        "bytes": 2359296
+      },
+      "source_weight_sha256": "b41e8d7028a6f787c256611afeae81c47990e7a91e9e682902dcedfcc9c13eaa",
+      "source_bias_sha256": "4dcdabdb17aabba97ac13d6e3c733a3b212d7e9e7092a7e0fbde6ccd81337a76"
+    },
+    {
+      "module_key": "transformer_blocks.9.img_mlp.net.0.proj",
+      "weight_key": "transformer_blocks.9.img_mlp.net.0.proj.weight",
+      "bias_key": "transformer_blocks.9.img_mlp.net.0.proj.bias",
+      "weight_shape": [
+        12288,
+        3072
+      ],
+      "bias_shape": [
+        12288
+      ],
+      "weight_tensor_scale_key": "targets.transformer_blocks.9.img_mlp.net.0.proj.weight_tensor_scale",
+      "artifact_weight_key": "targets.transformer_blocks.9.img_mlp.net.0.proj.packed_weight_e2m1",
+      "artifact_scale_key": "targets.transformer_blocks.9.img_mlp.net.0.proj.packed_scales_ue4m3",
+      "artifact_bias_key": "targets.transformer_blocks.9.img_mlp.net.0.proj.bias_bf16",
+      "weight_tensor_scale": 0.0004911876749247313,
+      "weight_global_amax": 1.3203125,
+      "packed_weight_bytes": 18874368,
+      "packed_scale_bytes": 2359296,
+      "scale_layout": {
+        "inner_dim": 192,
+        "outer_tiles": 96,
+        "bytes": 2359296
+      },
+      "source_weight_sha256": "82c42f661aab545f0807ff6aa8a1da83f8b50c186b2fec439ab65540449f4b58",
+      "source_bias_sha256": "790d198723850350f0430dc9262fc2019d3b1e3b7e32934212015bff7b767ff1"
+    },
+    {
+      "module_key": "transformer_blocks.9.img_mlp.net.2",
+      "weight_key": "transformer_blocks.9.img_mlp.net.2.weight",
+      "bias_key": "transformer_blocks.9.img_mlp.net.2.bias",
+      "weight_shape": [
+        3072,
+        12288
+      ],
+      "bias_shape": [
+        3072
+      ],
+      "weight_tensor_scale_key": "targets.transformer_blocks.9.img_mlp.net.2.weight_tensor_scale",
+      "artifact_weight_key": "targets.transformer_blocks.9.img_mlp.net.2.packed_weight_e2m1",
+      "artifact_scale_key": "targets.transformer_blocks.9.img_mlp.net.2.packed_scales_ue4m3",
+      "artifact_bias_key": "targets.transformer_blocks.9.img_mlp.net.2.bias_bf16",
+      "weight_tensor_scale": 0.0005435035564005375,
+      "weight_global_amax": 1.4609375,
+      "packed_weight_bytes": 18874368,
+      "packed_scale_bytes": 2359296,
+      "scale_layout": {
+        "inner_dim": 768,
+        "outer_tiles": 24,
+        "bytes": 2359296
+      },
+      "source_weight_sha256": "0c2af8a1f4e348d1150532b2fb9da63c33bee8603e1d4720b2c830632d6157c2",
+      "source_bias_sha256": "37f0b49b8f98d5981a92cc830bdcf202e077d458dcad8bedd26982ce12aef170"
+    },
+    {
+      "module_key": "transformer_blocks.9.txt_mlp.net.0.proj",
+      "weight_key": "transformer_blocks.9.txt_mlp.net.0.proj.weight",
+      "bias_key": "transformer_blocks.9.txt_mlp.net.0.proj.bias",
+      "weight_shape": [
+        12288,
+        3072
+      ],
+      "bias_shape": [
+        12288
+      ],
+      "weight_tensor_scale_key": "targets.transformer_blocks.9.txt_mlp.net.0.proj.weight_tensor_scale",
+      "artifact_weight_key": "targets.transformer_blocks.9.txt_mlp.net.0.proj.packed_weight_e2m1",
+      "artifact_scale_key": "targets.transformer_blocks.9.txt_mlp.net.0.proj.packed_scales_ue4m3",
+      "artifact_bias_key": "targets.transformer_blocks.9.txt_mlp.net.0.proj.bias_bf16",
+      "weight_tensor_scale": 0.0002819242945406586,
+      "weight_global_amax": 0.7578125,
+      "packed_weight_bytes": 18874368,
+      "packed_scale_bytes": 2359296,
+      "scale_layout": {
+        "inner_dim": 192,
+        "outer_tiles": 96,
+        "bytes": 2359296
+      },
+      "source_weight_sha256": "7569c27e45b5327904b003909f88f373e99bd93b150a5508bd3c49470da94c8f",
+      "source_bias_sha256": "5e5a68b982622caffbb11f7a7b131939ece3715bc0ffbcb32fa29dd445fd48ea"
+    },
+    {
+      "module_key": "transformer_blocks.9.txt_mlp.net.2",
+      "weight_key": "transformer_blocks.9.txt_mlp.net.2.weight",
+      "bias_key": "transformer_blocks.9.txt_mlp.net.2.bias",
+      "weight_shape": [
+        3072,
+        12288
+      ],
+      "bias_shape": [
+        3072
+      ],
+      "weight_tensor_scale_key": "targets.transformer_blocks.9.txt_mlp.net.2.weight_tensor_scale",
+      "artifact_weight_key": "targets.transformer_blocks.9.txt_mlp.net.2.packed_weight_e2m1",
+      "artifact_scale_key": "targets.transformer_blocks.9.txt_mlp.net.2.packed_scales_ue4m3",
+      "artifact_bias_key": "targets.transformer_blocks.9.txt_mlp.net.2.bias_bf16",
+      "weight_tensor_scale": 0.0006045387126505375,
+      "weight_global_amax": 1.625,
+      "packed_weight_bytes": 18874368,
+      "packed_scale_bytes": 2359296,
+      "scale_layout": {
+        "inner_dim": 768,
+        "outer_tiles": 24,
+        "bytes": 2359296
+      },
+      "source_weight_sha256": "5ce62f2c67fdf6f7caf6841ed2a2cee2e72194768207c0ab631f147dae462f80",
+      "source_bias_sha256": "11cb1954c19d92e661518cca9971ac432a4ff17f9d8832d789f91dc9bbaeca37"
+    },
+    {
+      "module_key": "transformer_blocks.10.img_mlp.net.0.proj",
+      "weight_key": "transformer_blocks.10.img_mlp.net.0.proj.weight",
+      "bias_key": "transformer_blocks.10.img_mlp.net.0.proj.bias",
+      "weight_shape": [
+        12288,
+        3072
+      ],
+      "bias_shape": [
+        12288
+      ],
+      "weight_tensor_scale_key": "targets.transformer_blocks.10.img_mlp.net.0.proj.weight_tensor_scale",
+      "artifact_weight_key": "targets.transformer_blocks.10.img_mlp.net.0.proj.packed_weight_e2m1",
+      "artifact_scale_key": "targets.transformer_blocks.10.img_mlp.net.0.proj.packed_scales_ue4m3",
+      "artifact_bias_key": "targets.transformer_blocks.10.img_mlp.net.0.proj.bias_bf16",
+      "weight_tensor_scale": 0.00019618442456703633,
+      "weight_global_amax": 0.52734375,
+      "packed_weight_bytes": 18874368,
+      "packed_scale_bytes": 2359296,
+      "scale_layout": {
+        "inner_dim": 192,
+        "outer_tiles": 96,
+        "bytes": 2359296
+      },
+      "source_weight_sha256": "25c16255a05df5f725a32fa1183740eecd018f52fa5e52d7b117f396656307f7",
+      "source_bias_sha256": "79ba550cd845a1d1156fa06661ea9a268331d2d3f13ed4f32225318796bdb2f1"
+    },
+    {
+      "module_key": "transformer_blocks.10.img_mlp.net.2",
+      "weight_key": "transformer_blocks.10.img_mlp.net.2.weight",
+      "bias_key": "transformer_blocks.10.img_mlp.net.2.bias",
+      "weight_shape": [
+        3072,
+        12288
+      ],
+      "bias_shape": [
+        3072
+      ],
+      "weight_tensor_scale_key": "targets.transformer_blocks.10.img_mlp.net.2.weight_tensor_scale",
+      "artifact_weight_key": "targets.transformer_blocks.10.img_mlp.net.2.packed_weight_e2m1",
+      "artifact_scale_key": "targets.transformer_blocks.10.img_mlp.net.2.packed_scales_ue4m3",
+      "artifact_bias_key": "targets.transformer_blocks.10.img_mlp.net.2.bias_bf16",
+      "weight_tensor_scale": 0.0005812871968373656,
+      "weight_global_amax": 1.5625,
+      "packed_weight_bytes": 18874368,
+      "packed_scale_bytes": 2359296,
+      "scale_layout": {
+        "inner_dim": 768,
+        "outer_tiles": 24,
+        "bytes": 2359296
+      },
+      "source_weight_sha256": "710837d28af5c4871c535b02e29431563d3d3755236b8a4acf8e0ebb8c334764",
+      "source_bias_sha256": "021b5a6799ed0e35ee49a28af9acc77561bc1defb184a47df39d28178edf3084"
+    },
+    {
+      "module_key": "transformer_blocks.10.txt_mlp.net.0.proj",
+      "weight_key": "transformer_blocks.10.txt_mlp.net.0.proj.weight",
+      "bias_key": "transformer_blocks.10.txt_mlp.net.0.proj.bias",
+      "weight_shape": [
+        12288,
+        3072
+      ],
+      "bias_shape": [
+        12288
+      ],
+      "weight_tensor_scale_key": "targets.transformer_blocks.10.txt_mlp.net.0.proj.weight_tensor_scale",
+      "artifact_weight_key": "targets.transformer_blocks.10.txt_mlp.net.0.proj.packed_weight_e2m1",
+      "artifact_scale_key": "targets.transformer_blocks.10.txt_mlp.net.0.proj.packed_scales_ue4m3",
+      "artifact_bias_key": "targets.transformer_blocks.10.txt_mlp.net.0.proj.bias_bf16",
+      "weight_tensor_scale": 0.00015476772387046367,
+      "weight_global_amax": 0.416015625,
+      "packed_weight_bytes": 18874368,
+      "packed_scale_bytes": 2359296,
+      "scale_layout": {
+        "inner_dim": 192,
+        "outer_tiles": 96,
+        "bytes": 2359296
+      },
+      "source_weight_sha256": "dd76244de2f2319b3c617e91cc6e863fe6d05c83d0038d3b903e198993c497dc",
+      "source_bias_sha256": "dd6336134c5a891280e5e0b22ae7690cddf6d4cf937c4198ac5865202958a162"
+    },
+    {
+      "module_key": "transformer_blocks.10.txt_mlp.net.2",
+      "weight_key": "transformer_blocks.10.txt_mlp.net.2.weight",
+      "bias_key": "transformer_blocks.10.txt_mlp.net.2.bias",
+      "weight_shape": [
+        3072,
+        12288
+      ],
+      "bias_shape": [
+        3072
+      ],
+      "weight_tensor_scale_key": "targets.transformer_blocks.10.txt_mlp.net.2.weight_tensor_scale",
+      "artifact_weight_key": "targets.transformer_blocks.10.txt_mlp.net.2.packed_weight_e2m1",
+      "artifact_scale_key": "targets.transformer_blocks.10.txt_mlp.net.2.packed_scales_ue4m3",
+      "artifact_bias_key": "targets.transformer_blocks.10.txt_mlp.net.2.bias_bf16",
+      "weight_tensor_scale": 0.00018383207498118281,
+      "weight_global_amax": 0.494140625,
+      "packed_weight_bytes": 18874368,
+      "packed_scale_bytes": 2359296,
+      "scale_layout": {
+        "inner_dim": 768,
+        "outer_tiles": 24,
+        "bytes": 2359296
+      },
+      "source_weight_sha256": "3748e2fb3c162f846342f7a4c5f8f600d614150163332c2b32ba5f3c7291f4b4",
+      "source_bias_sha256": "333b152ef1291bf1fcd6fc796345483195b2632b70e0b6ab13ddabf8b7bb20e4"
+    },
+    {
+      "module_key": "transformer_blocks.11.img_mlp.net.0.proj",
+      "weight_key": "transformer_blocks.11.img_mlp.net.0.proj.weight",
+      "bias_key": "transformer_blocks.11.img_mlp.net.0.proj.bias",
+      "weight_shape": [
+        12288,
+        3072
+      ],
+      "bias_shape": [
+        12288
+      ],
+      "weight_tensor_scale_key": "targets.transformer_blocks.11.img_mlp.net.0.proj.weight_tensor_scale",
+      "artifact_weight_key": "targets.transformer_blocks.11.img_mlp.net.0.proj.packed_weight_e2m1",
+      "artifact_scale_key": "targets.transformer_blocks.11.img_mlp.net.0.proj.packed_scales_ue4m3",
+      "artifact_bias_key": "targets.transformer_blocks.11.img_mlp.net.0.proj.bias_bf16",
+      "weight_tensor_scale": 0.0002252487902296707,
+      "weight_global_amax": 0.60546875,
+      "packed_weight_bytes": 18874368,
+      "packed_scale_bytes": 2359296,
+      "scale_layout": {
+        "inner_dim": 192,
+        "outer_tiles": 96,
+        "bytes": 2359296
+      },
+      "source_weight_sha256": "81207ce0a6cc70bf94346c5c74e7729a9df4e2cc87e943f5461ce6ccdcddf25f",
+      "source_bias_sha256": "3baf43b038997ba0ebe0e9326e17fe210727fe89fcc832f93082381c5cca2adc"
+    },
+    {
+      "module_key": "transformer_blocks.11.img_mlp.net.2",
+      "weight_key": "transformer_blocks.11.img_mlp.net.2.weight",
+      "bias_key": "transformer_blocks.11.img_mlp.net.2.bias",
+      "weight_shape": [
+        3072,
+        12288
+      ],
+      "bias_shape": [
+        3072
+      ],
+      "weight_tensor_scale_key": "targets.transformer_blocks.11.img_mlp.net.2.weight_tensor_scale",
+      "artifact_weight_key": "targets.transformer_blocks.11.img_mlp.net.2.packed_weight_e2m1",
+      "artifact_scale_key": "targets.transformer_blocks.11.img_mlp.net.2.packed_scales_ue4m3",
+      "artifact_bias_key": "targets.transformer_blocks.11.img_mlp.net.2.bias_bf16",
+      "weight_tensor_scale": 0.0009707496501505375,
+      "weight_global_amax": 2.609375,
+      "packed_weight_bytes": 18874368,
+      "packed_scale_bytes": 2359296,
+      "scale_layout": {
+        "inner_dim": 768,
+        "outer_tiles": 24,
+        "bytes": 2359296
+      },
+      "source_weight_sha256": "ba3889bd47bfe14a07b6b4d9f624d1187980d07fb81de6d0be765f195a22c487",
+      "source_bias_sha256": "fd48e42ed9061e494e065910ae622964b25e6c3abc13613b38f4b4a4a0c8ec13"
+    },
+    {
+      "module_key": "transformer_blocks.11.txt_mlp.net.0.proj",
+      "weight_key": "transformer_blocks.11.txt_mlp.net.0.proj.weight",
+      "bias_key": "transformer_blocks.11.txt_mlp.net.0.proj.bias",
+      "weight_shape": [
+        12288,
+        3072
+      ],
+      "bias_shape": [
+        12288
+      ],
+      "weight_tensor_scale_key": "targets.transformer_blocks.11.txt_mlp.net.0.proj.weight_tensor_scale",
+      "artifact_weight_key": "targets.transformer_blocks.11.txt_mlp.net.0.proj.packed_weight_e2m1",
+      "artifact_scale_key": "targets.transformer_blocks.11.txt_mlp.net.0.proj.packed_scales_ue4m3",
+      "artifact_bias_key": "targets.transformer_blocks.11.txt_mlp.net.0.proj.bias_bf16",
+      "weight_tensor_scale": 1.3850983577867737e-06,
+      "weight_global_amax": 0.00372314453125,
+      "packed_weight_bytes": 18874368,
+      "packed_scale_bytes": 2359296,
+      "scale_layout": {
+        "inner_dim": 192,
+        "outer_tiles": 96,
+        "bytes": 2359296
+      },
+      "source_weight_sha256": "3e43f754f9d63e3dd6d6a71c4cc0c15344bd56da4d5a19a7adeb54445b820d72",
+      "source_bias_sha256": "062da4757b9f03b3273c98cf91a517e2b4a667cbfe1f0ba190b02b626b0ac75c"
+    },
+    {
+      "module_key": "transformer_blocks.11.txt_mlp.net.2",
+      "weight_key": "transformer_blocks.11.txt_mlp.net.2.weight",
+      "bias_key": "transformer_blocks.11.txt_mlp.net.2.bias",
+      "weight_shape": [
+        3072,
+        12288
+      ],
+      "bias_shape": [
+        3072
+      ],
+      "weight_tensor_scale_key": "targets.transformer_blocks.11.txt_mlp.net.2.weight_tensor_scale",
+      "artifact_weight_key": "targets.transformer_blocks.11.txt_mlp.net.2.packed_weight_e2m1",
+      "artifact_scale_key": "targets.transformer_blocks.11.txt_mlp.net.2.packed_scales_ue4m3",
+      "artifact_bias_key": "targets.transformer_blocks.11.txt_mlp.net.2.bias_bf16",
+      "weight_tensor_scale": 6.925491788933869e-07,
+      "weight_global_amax": 0.001861572265625,
+      "packed_weight_bytes": 18874368,
+      "packed_scale_bytes": 2359296,
+      "scale_layout": {
+        "inner_dim": 768,
+        "outer_tiles": 24,
+        "bytes": 2359296
+      },
+      "source_weight_sha256": "4a74e0fc231d978ddbca106304430d9dcb511c7284e3c18c245005626e88224d",
+      "source_bias_sha256": "6ef0d2f05867e5017824fc95cfd0a5803ce3d736cd622d82d92b603938d66e41"
+    }
+  ],
+  "non_target_keys": [
+    "img_in.bias",
+    "img_in.weight",
+    "norm_out.linear.bias",
+    "norm_out.linear.weight",
+    "proj_out.bias",
+    "proj_out.weight",
+    "time_text_embed.timestep_embedder.linear_1.bias",
+    "time_text_embed.timestep_embedder.linear_1.weight",
+    "time_text_embed.timestep_embedder.linear_2.bias",
+    "time_text_embed.timestep_embedder.linear_2.weight",
+    "transformer_blocks.0.attn.add_k_proj.bias",
+    "transformer_blocks.0.attn.add_k_proj.weight",
+    "transformer_blocks.0.attn.add_q_proj.bias",
+    "transformer_blocks.0.attn.add_q_proj.weight",
+    "transformer_blocks.0.attn.add_v_proj.bias",
+    "transformer_blocks.0.attn.add_v_proj.weight",
+    "transformer_blocks.0.attn.norm_added_k.weight",
+    "transformer_blocks.0.attn.norm_added_q.weight",
+    "transformer_blocks.0.attn.norm_k.weight",
+    "transformer_blocks.0.attn.norm_q.weight",
+    "transformer_blocks.0.attn.to_add_out.bias",
+    "transformer_blocks.0.attn.to_add_out.weight",
+    "transformer_blocks.0.attn.to_k.bias",
+    "transformer_blocks.0.attn.to_k.weight",
+    "transformer_blocks.0.attn.to_out.0.bias",
+    "transformer_blocks.0.attn.to_out.0.weight",
+    "transformer_blocks.0.attn.to_q.bias",
+    "transformer_blocks.0.attn.to_q.weight",
+    "transformer_blocks.0.attn.to_v.bias",
+    "transformer_blocks.0.attn.to_v.weight",
+    "transformer_blocks.0.img_mod.1.bias",
+    "transformer_blocks.0.img_mod.1.weight",
+    "transformer_blocks.0.txt_mod.1.bias",
+    "transformer_blocks.0.txt_mod.1.weight",
+    "transformer_blocks.1.attn.add_k_proj.bias",
+    "transformer_blocks.1.attn.add_k_proj.weight",
+    "transformer_blocks.1.attn.add_q_proj.bias",
+    "transformer_blocks.1.attn.add_q_proj.weight",
+    "transformer_blocks.1.attn.add_v_proj.bias",
+    "transformer_blocks.1.attn.add_v_proj.weight",
+    "transformer_blocks.1.attn.norm_added_k.weight",
+    "transformer_blocks.1.attn.norm_added_q.weight",
+    "transformer_blocks.1.attn.norm_k.weight",
+    "transformer_blocks.1.attn.norm_q.weight",
+    "transformer_blocks.1.attn.to_add_out.bias",
+    "transformer_blocks.1.attn.to_add_out.weight",
+    "transformer_blocks.1.attn.to_k.bias",
+    "transformer_blocks.1.attn.to_k.weight",
+    "transformer_blocks.1.attn.to_out.0.bias",
+    "transformer_blocks.1.attn.to_out.0.weight",
+    "transformer_blocks.1.attn.to_q.bias",
+    "transformer_blocks.1.attn.to_q.weight",
+    "transformer_blocks.1.attn.to_v.bias",
+    "transformer_blocks.1.attn.to_v.weight",
+    "transformer_blocks.1.img_mod.1.bias",
+    "transformer_blocks.1.img_mod.1.weight",
+    "transformer_blocks.1.txt_mod.1.bias",
+    "transformer_blocks.1.txt_mod.1.weight",
+    "transformer_blocks.10.attn.add_k_proj.bias",
+    "transformer_blocks.10.attn.add_k_proj.weight",
+    "transformer_blocks.10.attn.add_q_proj.bias",
+    "transformer_blocks.10.attn.add_q_proj.weight",
+    "transformer_blocks.10.attn.add_v_proj.bias",
+    "transformer_blocks.10.attn.add_v_proj.weight",
+    "transformer_blocks.10.attn.norm_added_k.weight",
+    "transformer_blocks.10.attn.norm_added_q.weight",
+    "transformer_blocks.10.attn.norm_k.weight",
+    "transformer_blocks.10.attn.norm_q.weight",
+    "transformer_blocks.10.attn.to_add_out.bias",
+    "transformer_blocks.10.attn.to_add_out.weight",
+    "transformer_blocks.10.attn.to_k.bias",
+    "transformer_blocks.10.attn.to_k.weight",
+    "transformer_blocks.10.attn.to_out.0.bias",
+    "transformer_blocks.10.attn.to_out.0.weight",
+    "transformer_blocks.10.attn.to_q.bias",
+    "transformer_blocks.10.attn.to_q.weight",
+    "transformer_blocks.10.attn.to_v.bias",
+    "transformer_blocks.10.attn.to_v.weight",
+    "transformer_blocks.10.img_mod.1.bias",
+    "transformer_blocks.10.img_mod.1.weight",
+    "transformer_blocks.10.txt_mod.1.bias",
+    "transformer_blocks.10.txt_mod.1.weight",
+    "transformer_blocks.11.attn.add_k_proj.bias",
+    "transformer_blocks.11.attn.add_k_proj.weight",
+    "transformer_blocks.11.attn.add_q_proj.bias",
+    "transformer_blocks.11.attn.add_q_proj.weight",
+    "transformer_blocks.11.attn.add_v_proj.bias",
+    "transformer_blocks.11.attn.add_v_proj.weight",
+    "transformer_blocks.11.attn.norm_added_k.weight",
+    "transformer_blocks.11.attn.norm_added_q.weight",
+    "transformer_blocks.11.attn.norm_k.weight",
+    "transformer_blocks.11.attn.norm_q.weight",
+    "transformer_blocks.11.attn.to_add_out.bias",
+    "transformer_blocks.11.attn.to_add_out.weight",
+    "transformer_blocks.11.attn.to_k.bias",
+    "transformer_blocks.11.attn.to_k.weight",
+    "transformer_blocks.11.attn.to_out.0.bias",
+    "transformer_blocks.11.attn.to_out.0.weight",
+    "transformer_blocks.11.attn.to_q.bias",
+    "transformer_blocks.11.attn.to_q.weight",
+    "transformer_blocks.11.attn.to_v.bias",
+    "transformer_blocks.11.attn.to_v.weight",
+    "transformer_blocks.11.img_mod.1.bias",
+    "transformer_blocks.11.img_mod.1.weight",
+    "transformer_blocks.11.txt_mod.1.bias",
+    "transformer_blocks.11.txt_mod.1.weight",
+    "transformer_blocks.2.attn.add_k_proj.bias",
+    "transformer_blocks.2.attn.add_k_proj.weight",
+    "transformer_blocks.2.attn.add_q_proj.bias",
+    "transformer_blocks.2.attn.add_q_proj.weight",
+    "transformer_blocks.2.attn.add_v_proj.bias",
+    "transformer_blocks.2.attn.add_v_proj.weight",
+    "transformer_blocks.2.attn.norm_added_k.weight",
+    "transformer_blocks.2.attn.norm_added_q.weight",
+    "transformer_blocks.2.attn.norm_k.weight",
+    "transformer_blocks.2.attn.norm_q.weight",
+    "transformer_blocks.2.attn.to_add_out.bias",
+    "transformer_blocks.2.attn.to_add_out.weight",
+    "transformer_blocks.2.attn.to_k.bias",
+    "transformer_blocks.2.attn.to_k.weight",
+    "transformer_blocks.2.attn.to_out.0.bias",
+    "transformer_blocks.2.attn.to_out.0.weight",
+    "transformer_blocks.2.attn.to_q.bias",
+    "transformer_blocks.2.attn.to_q.weight",
+    "transformer_blocks.2.attn.to_v.bias",
+    "transformer_blocks.2.attn.to_v.weight",
+    "transformer_blocks.2.img_mod.1.bias",
+    "transformer_blocks.2.img_mod.1.weight",
+    "transformer_blocks.2.txt_mod.1.bias",
+    "transformer_blocks.2.txt_mod.1.weight",
+    "transformer_blocks.3.attn.add_k_proj.bias",
+    "transformer_blocks.3.attn.add_k_proj.weight",
+    "transformer_blocks.3.attn.add_q_proj.bias",
+    "transformer_blocks.3.attn.add_q_proj.weight",
+    "transformer_blocks.3.attn.add_v_proj.bias",
+    "transformer_blocks.3.attn.add_v_proj.weight",
+    "transformer_blocks.3.attn.norm_added_k.weight",
+    "transformer_blocks.3.attn.norm_added_q.weight",
+    "transformer_blocks.3.attn.norm_k.weight",
+    "transformer_blocks.3.attn.norm_q.weight",
+    "transformer_blocks.3.attn.to_add_out.bias",
+    "transformer_blocks.3.attn.to_add_out.weight",
+    "transformer_blocks.3.attn.to_k.bias",
+    "transformer_blocks.3.attn.to_k.weight",
+    "transformer_blocks.3.attn.to_out.0.bias",
+    "transformer_blocks.3.attn.to_out.0.weight",
+    "transformer_blocks.3.attn.to_q.bias",
+    "transformer_blocks.3.attn.to_q.weight",
+    "transformer_blocks.3.attn.to_v.bias",
+    "transformer_blocks.3.attn.to_v.weight",
+    "transformer_blocks.3.img_mod.1.bias",
+    "transformer_blocks.3.img_mod.1.weight",
+    "transformer_blocks.3.txt_mod.1.bias",
+    "transformer_blocks.3.txt_mod.1.weight",
+    "transformer_blocks.4.attn.add_k_proj.bias",
+    "transformer_blocks.4.attn.add_k_proj.weight",
+    "transformer_blocks.4.attn.add_q_proj.bias",
+    "transformer_blocks.4.attn.add_q_proj.weight",
+    "transformer_blocks.4.attn.add_v_proj.bias",
+    "transformer_blocks.4.attn.add_v_proj.weight",
+    "transformer_blocks.4.attn.norm_added_k.weight",
+    "transformer_blocks.4.attn.norm_added_q.weight",
+    "transformer_blocks.4.attn.norm_k.weight",
+    "transformer_blocks.4.attn.norm_q.weight",
+    "transformer_blocks.4.attn.to_add_out.bias",
+    "transformer_blocks.4.attn.to_add_out.weight",
+    "transformer_blocks.4.attn.to_k.bias",
+    "transformer_blocks.4.attn.to_k.weight",
+    "transformer_blocks.4.attn.to_out.0.bias",
+    "transformer_blocks.4.attn.to_out.0.weight",
+    "transformer_blocks.4.attn.to_q.bias",
+    "transformer_blocks.4.attn.to_q.weight",
+    "transformer_blocks.4.attn.to_v.bias",
+    "transformer_blocks.4.attn.to_v.weight",
+    "transformer_blocks.4.img_mod.1.bias",
+    "transformer_blocks.4.img_mod.1.weight",
+    "transformer_blocks.4.txt_mod.1.bias",
+    "transformer_blocks.4.txt_mod.1.weight",
+    "transformer_blocks.5.attn.add_k_proj.bias",
+    "transformer_blocks.5.attn.add_k_proj.weight",
+    "transformer_blocks.5.attn.add_q_proj.bias",
+    "transformer_blocks.5.attn.add_q_proj.weight",
+    "transformer_blocks.5.attn.add_v_proj.bias",
+    "transformer_blocks.5.attn.add_v_proj.weight",
+    "transformer_blocks.5.attn.norm_added_k.weight",
+    "transformer_blocks.5.attn.norm_added_q.weight",
+    "transformer_blocks.5.attn.norm_k.weight",
+    "transformer_blocks.5.attn.norm_q.weight",
+    "transformer_blocks.5.attn.to_add_out.bias",
+    "transformer_blocks.5.attn.to_add_out.weight",
+    "transformer_blocks.5.attn.to_k.bias",
+    "transformer_blocks.5.attn.to_k.weight",
+    "transformer_blocks.5.attn.to_out.0.bias",
+    "transformer_blocks.5.attn.to_out.0.weight",
+    "transformer_blocks.5.attn.to_q.bias",
+    "transformer_blocks.5.attn.to_q.weight",
+    "transformer_blocks.5.attn.to_v.bias",
+    "transformer_blocks.5.attn.to_v.weight",
+    "transformer_blocks.5.img_mod.1.bias",
+    "transformer_blocks.5.img_mod.1.weight",
+    "transformer_blocks.5.txt_mod.1.bias",
+    "transformer_blocks.5.txt_mod.1.weight",
+    "transformer_blocks.6.attn.add_k_proj.bias",
+    "transformer_blocks.6.attn.add_k_proj.weight",
+    "transformer_blocks.6.attn.add_q_proj.bias",
+    "transformer_blocks.6.attn.add_q_proj.weight",
+    "transformer_blocks.6.attn.add_v_proj.bias",
+    "transformer_blocks.6.attn.add_v_proj.weight",
+    "transformer_blocks.6.attn.norm_added_k.weight",
+    "transformer_blocks.6.attn.norm_added_q.weight",
+    "transformer_blocks.6.attn.norm_k.weight",
+    "transformer_blocks.6.attn.norm_q.weight",
+    "transformer_blocks.6.attn.to_add_out.bias",
+    "transformer_blocks.6.attn.to_add_out.weight",
+    "transformer_blocks.6.attn.to_k.bias",
+    "transformer_blocks.6.attn.to_k.weight",
+    "transformer_blocks.6.attn.to_out.0.bias",
+    "transformer_blocks.6.attn.to_out.0.weight",
+    "transformer_blocks.6.attn.to_q.bias",
+    "transformer_blocks.6.attn.to_q.weight",
+    "transformer_blocks.6.attn.to_v.bias",
+    "transformer_blocks.6.attn.to_v.weight",
+    "transformer_blocks.6.img_mod.1.bias",
+    "transformer_blocks.6.img_mod.1.weight",
+    "transformer_blocks.6.txt_mod.1.bias",
+    "transformer_blocks.6.txt_mod.1.weight",
+    "transformer_blocks.7.attn.add_k_proj.bias",
+    "transformer_blocks.7.attn.add_k_proj.weight",
+    "transformer_blocks.7.attn.add_q_proj.bias",
+    "transformer_blocks.7.attn.add_q_proj.weight",
+    "transformer_blocks.7.attn.add_v_proj.bias",
+    "transformer_blocks.7.attn.add_v_proj.weight",
+    "transformer_blocks.7.attn.norm_added_k.weight",
+    "transformer_blocks.7.attn.norm_added_q.weight",
+    "transformer_blocks.7.attn.norm_k.weight",
+    "transformer_blocks.7.attn.norm_q.weight",
+    "transformer_blocks.7.attn.to_add_out.bias",
+    "transformer_blocks.7.attn.to_add_out.weight",
+    "transformer_blocks.7.attn.to_k.bias",
+    "transformer_blocks.7.attn.to_k.weight",
+    "transformer_blocks.7.attn.to_out.0.bias",
+    "transformer_blocks.7.attn.to_out.0.weight",
+    "transformer_blocks.7.attn.to_q.bias",
+    "transformer_blocks.7.attn.to_q.weight",
+    "transformer_blocks.7.attn.to_v.bias",
+    "transformer_blocks.7.attn.to_v.weight",
+    "transformer_blocks.7.img_mod.1.bias",
+    "transformer_blocks.7.img_mod.1.weight",
+    "transformer_blocks.7.txt_mod.1.bias",
+    "transformer_blocks.7.txt_mod.1.weight",
+    "transformer_blocks.8.attn.add_k_proj.bias",
+    "transformer_blocks.8.attn.add_k_proj.weight",
+    "transformer_blocks.8.attn.add_q_proj.bias",
+    "transformer_blocks.8.attn.add_q_proj.weight",
+    "transformer_blocks.8.attn.add_v_proj.bias",
+    "transformer_blocks.8.attn.add_v_proj.weight",
+    "transformer_blocks.8.attn.norm_added_k.weight",
+    "transformer_blocks.8.attn.norm_added_q.weight",
+    "transformer_blocks.8.attn.norm_k.weight",
+    "transformer_blocks.8.attn.norm_q.weight",
+    "transformer_blocks.8.attn.to_add_out.bias",
+    "transformer_blocks.8.attn.to_add_out.weight",
+    "transformer_blocks.8.attn.to_k.bias",
+    "transformer_blocks.8.attn.to_k.weight",
+    "transformer_blocks.8.attn.to_out.0.bias",
+    "transformer_blocks.8.attn.to_out.0.weight",
+    "transformer_blocks.8.attn.to_q.bias",
+    "transformer_blocks.8.attn.to_q.weight",
+    "transformer_blocks.8.attn.to_v.bias",
+    "transformer_blocks.8.attn.to_v.weight",
+    "transformer_blocks.8.img_mod.1.bias",
+    "transformer_blocks.8.img_mod.1.weight",
+    "transformer_blocks.8.txt_mod.1.bias",
+    "transformer_blocks.8.txt_mod.1.weight",
+    "transformer_blocks.9.attn.add_k_proj.bias",
+    "transformer_blocks.9.attn.add_k_proj.weight",
+    "transformer_blocks.9.attn.add_q_proj.bias",
+    "transformer_blocks.9.attn.add_q_proj.weight",
+    "transformer_blocks.9.attn.add_v_proj.bias",
+    "transformer_blocks.9.attn.add_v_proj.weight",
+    "transformer_blocks.9.attn.norm_added_k.weight",
+    "transformer_blocks.9.attn.norm_added_q.weight",
+    "transformer_blocks.9.attn.norm_k.weight",
+    "transformer_blocks.9.attn.norm_q.weight",
+    "transformer_blocks.9.attn.to_add_out.bias",
+    "transformer_blocks.9.attn.to_add_out.weight",
+    "transformer_blocks.9.attn.to_k.bias",
+    "transformer_blocks.9.attn.to_k.weight",
+    "transformer_blocks.9.attn.to_out.0.bias",
+    "transformer_blocks.9.attn.to_out.0.weight",
+    "transformer_blocks.9.attn.to_q.bias",
+    "transformer_blocks.9.attn.to_q.weight",
+    "transformer_blocks.9.attn.to_v.bias",
+    "transformer_blocks.9.attn.to_v.weight",
+    "transformer_blocks.9.img_mod.1.bias",
+    "transformer_blocks.9.img_mod.1.weight",
+    "transformer_blocks.9.txt_mod.1.bias",
+    "transformer_blocks.9.txt_mod.1.weight",
+    "txt_in.bias",
+    "txt_in.weight",
+    "txt_norm.weight"
+  ]
+}

vae/config.json

@@ -0,0 +1,6 @@
+{
+  "_class_name": "MageVAE",
+  "latent_channels": 128,
+  "downsample_factor": 16,
+  "sample_posterior": false
+}

validate_release.py

@@ -0,0 +1,89 @@
+#!/usr/bin/env python3
+"""Verify every immutable file declared by the v1 release manifest."""
+
+from __future__ import annotations
+
+import hashlib
+import json
+from pathlib import Path
+import sys
+
+
+RELEASE_ROOT = Path(__file__).resolve().parent
+MANIFEST_PATH = RELEASE_ROOT / "MANIFEST.json"
+
+
+def sha256(path: Path) -> str:
+    digest = hashlib.sha256()
+    with path.open("rb") as handle:
+        for chunk in iter(lambda: handle.read(1024 * 1024), b""):
+            digest.update(chunk)
+    return digest.hexdigest()
+
+
+def main() -> int:
+    manifest = json.loads(MANIFEST_PATH.read_text(encoding="utf-8"))
+    failures: list[str] = []
+    checked: list[dict[str, object]] = []
+    declared_paths = {
+        Path(record["path"]).as_posix() for record in manifest["files"]
+    }
+    actual_paths = {
+        path.relative_to(RELEASE_ROOT).as_posix()
+        for path in RELEASE_ROOT.rglob("*")
+        if (
+            path.is_file()
+            and path != MANIFEST_PATH
+            and "__pycache__" not in path.relative_to(RELEASE_ROOT).parts
+            and path.suffix != ".pyc"
+        )
+    }
+    undeclared = sorted(actual_paths - declared_paths)
+    absent = sorted(declared_paths - actual_paths)
+    if undeclared:
+        failures.extend(f"undeclared package file: {path}" for path in undeclared)
+    if absent:
+        failures.extend(f"declared package file is absent: {path}" for path in absent)
+    for record in manifest["files"]:
+        relative = Path(record["path"])
+        path = (RELEASE_ROOT / relative).resolve()
+        if not path.is_relative_to(RELEASE_ROOT):
+            failures.append(f"path escapes release root: {relative}")
+            continue
+        if not path.is_file():
+            failures.append(f"missing: {relative}")
+            continue
+        if path.is_symlink():
+            failures.append(f"release file must not be a symlink: {relative}")
+            continue
+        actual_size = path.stat().st_size
+        actual_hash = sha256(path)
+        if actual_size != record["size"]:
+            failures.append(
+                f"size mismatch {relative}: {actual_size} != {record['size']}"
+            )
+        if actual_hash != record["sha256"]:
+            failures.append(
+                f"SHA-256 mismatch {relative}: {actual_hash} != "
+                f"{record['sha256']}"
+            )
+        checked.append(
+            {
+                "path": str(relative),
+                "size": actual_size,
+                "sha256": actual_hash,
+            }
+        )
+
+    result = {
+        "schema_version": manifest["schema_version"],
+        "status": "pass" if not failures else "fail",
+        "checked_file_count": len(checked),
+        "failures": failures,
+    }
+    print(json.dumps(result, indent=2, sort_keys=True))
+    return 0 if not failures else 1
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())

vendor/mage_flow/__init__.py

@@ -0,0 +1,18 @@
+"""MageFlow — standalone text-to-image + image-edit inference."""
+
+from .models.mage_flow import ModelConfig
+from .pipeline import (
+    MageFlowPipeline,
+    generate_edits,
+    generate_images,
+    load_from_repo,
+)
+
+__all__ = [
+    "MageFlowPipeline",
+    "generate_images",
+    "generate_edits",
+    "load_from_repo",
+    "ModelConfig",
+]
+__version__ = "0.1.0"

vendor/mage_flow/app.py

@@ -0,0 +1,199 @@
+"""Gradio app for MageFlow — text-to-image and instruction-based image editing.
+
+    python app.py                       # serve on 0.0.0.0:7860
+    python app.py --share --port 7861
+
+Each tab has a model preset dropdown (base / rl / turbo) plus a free-form
+"Custom model" box for any Hugging Face repo id or local path. Models load
+lazily on first use and are cached. Notes:
+  - By default the presets point at the `microsoft/Mage-Flow*` Hugging Face
+    repos (downloaded + cached on first use). Set ``MAGEFLOW_HF_DIR`` to load
+    local checkpoint dirs instead.
+  - Turbo checkpoints are few-step: use steps=4, cfg=1.
+"""
+from __future__ import annotations
+
+import argparse
+import os
+
+import gradio as gr
+from PIL import Image
+
+from mage_flow.pipeline import MageFlowPipeline
+
+# Default to Hugging Face repo ids; if MAGEFLOW_HF_DIR is set, use local
+# checkpoint dirs under it instead (local dir names match the HF repo basename).
+HF_DIR = os.environ.get("MAGEFLOW_HF_DIR")
+
+
+def _repo(hf_id: str, local_name: str) -> str:
+    return f"{HF_DIR}/{local_name}" if HF_DIR else hf_id
+
+
+T2I_MODELS = {
+    "base":  _repo("microsoft/Mage-Flow-Base",  "Mage-Flow-Base"),
+    "rl":    _repo("microsoft/Mage-Flow",       "Mage-Flow"),
+    "turbo": _repo("microsoft/Mage-Flow-Turbo", "Mage-Flow-Turbo"),
+}
+EDIT_MODELS = {
+    "base":  _repo("microsoft/Mage-Flow-Edit-Base",  "Mage-Flow-Edit-Base"),
+    "rl":    _repo("microsoft/Mage-Flow-Edit",       "Mage-Flow-Edit"),
+    "turbo": _repo("microsoft/Mage-Flow-Edit-Turbo", "Mage-Flow-Edit-Turbo"),
+}
+
+DEVICE = "cuda"
+_CACHE: dict[str, MageFlowPipeline] = {}
+
+
+def _get_pipe(repo: str) -> MageFlowPipeline:
+    """Load (and cache) a pipeline from a local dir OR a Hugging Face repo id.
+
+    ``MageFlowPipeline.from_pretrained`` resolves a repo id via
+    ``snapshot_download`` automatically, so both are accepted here.
+    """
+    repo = (repo or "").strip()
+    if not repo:
+        raise gr.Error("No model specified.")
+    if repo not in _CACHE:
+        try:
+            _CACHE[repo] = MageFlowPipeline.from_pretrained(repo, device=DEVICE)
+        except Exception as e:  # noqa: BLE001
+            raise gr.Error(f"Failed to load model '{repo}': {type(e).__name__}: {e}")
+    return _CACHE[repo]
+
+
+def _resolve(preset_map, model_key, custom_model):
+    """Custom repo id / path (if given) overrides the preset dropdown."""
+    return (custom_model or "").strip() or preset_map[model_key]
+
+
+def run_t2i(model_key, custom_model, prompt, neg_prompt, steps, cfg, height, width, seed,
+            progress=gr.Progress(track_tqdm=False)):
+    if not (prompt or "").strip():
+        raise gr.Error("Prompt is empty.")
+    repo = _resolve(T2I_MODELS, model_key, custom_model)
+    progress(0.1, desc=f"loading {repo} …")
+    pipe = _get_pipe(repo)
+    progress(0.4, desc="generating …")
+    img = pipe.generate(
+        [prompt], neg_prompts=[neg_prompt or " "], seeds=[int(seed)],
+        steps=int(steps), cfg=float(cfg),
+        heights=[int(height)], widths=[int(width)],
+    )[0]
+    return img
+
+
+def run_edit(model_key, custom_model, prompt, neg_prompt, ref_img, extra_files, steps, cfg, max_size, seed,
+             progress=gr.Progress(track_tqdm=False)):
+    if not (prompt or "").strip():
+        raise gr.Error("Edit instruction is empty.")
+    refs = []
+    if ref_img is not None:
+        refs.append(ref_img if isinstance(ref_img, Image.Image) else Image.open(ref_img))
+    for f in (extra_files or []):
+        refs.append(Image.open(f).convert("RGB"))
+    if not refs:
+        raise gr.Error("Upload at least one reference image.")
+    refs = [r.convert("RGB") for r in refs]
+    repo = _resolve(EDIT_MODELS, model_key, custom_model)
+    progress(0.1, desc=f"loading {repo} …")
+    pipe = _get_pipe(repo)
+    progress(0.4, desc="editing …")
+    out = pipe.edit(
+        [prompt], [refs], neg_prompts=[neg_prompt or " "], seeds=[int(seed)],
+        steps=int(steps), cfg=float(cfg),
+        max_size=int(max_size) if max_size else None,
+    )[0]
+    return out
+
+
+_NOTE = (
+    "Pick a **preset** (base / rl / turbo) or type a **custom model** — any "
+    "Hugging Face repo id (e.g. `microsoft/Mage-Flow-Turbo`) or local path; it "
+    "is downloaded and cached on first use. **Turbo** models are few-step: set "
+    "**steps=4, cfg=1**."
+)
+
+_CUSTOM_PH_T2I = "microsoft/Mage-Flow  (repo id or local path — overrides preset)"
+_CUSTOM_PH_EDIT = "microsoft/Mage-Flow-Edit  (repo id or local path — overrides preset)"
+
+
+def build_ui():
+    with gr.Blocks(title="MageFlow") as demo:
+        gr.Markdown("# MageFlow\nText-to-image generation and instruction-based image editing.")
+        gr.Markdown(_NOTE)
+
+        with gr.Tab("Text → Image"):
+            with gr.Row():
+                with gr.Column(scale=1):
+                    t_model = gr.Dropdown(list(T2I_MODELS), value="base", label="Model preset")
+                    t_custom = gr.Textbox(label="Custom model (optional)", placeholder=_CUSTOM_PH_T2I, lines=1)
+                    t_prompt = gr.Textbox(label="Prompt", lines=3,
+                                          value="A close-up portrait of an elderly African man with deep wrinkles, wearing a traditional hat, soft natural lighting, ultra realistic.")
+                    t_neg = gr.Textbox(label="Negative prompt", value=" ", lines=1)
+                    with gr.Row():
+                        t_steps = gr.Slider(1, 50, value=30, step=1, label="Steps")
+                        t_cfg = gr.Slider(1.0, 10.0, value=5.0, step=0.5, label="CFG")
+                    with gr.Row():
+                        t_h = gr.Slider(256, 1536, value=1024, step=16, label="Height")
+                        t_w = gr.Slider(256, 1536, value=1024, step=16, label="Width")
+                    t_seed = gr.Number(value=42, precision=0, label="Seed")
+                    t_btn = gr.Button("Generate", variant="primary")
+                with gr.Column(scale=1):
+                    t_out = gr.Image(type="pil", label="Output", height=560)
+            # Clear the previous output first so the stale image isn't shown as
+            # the result while the new one is still transferring (esp. over a
+            # gradio share tunnel, where the image download can lag a few seconds).
+            t_btn.click(lambda: None, None, t_out).then(
+                        run_t2i,
+                        [t_model, t_custom, t_prompt, t_neg, t_steps, t_cfg, t_h, t_w, t_seed],
+                        t_out)
+
+        with gr.Tab("Image Edit"):
+            with gr.Row():
+                with gr.Column(scale=1):
+                    e_model = gr.Dropdown(list(EDIT_MODELS), value="base", label="Model preset")
+                    e_custom = gr.Textbox(label="Custom model (optional)", placeholder=_CUSTOM_PH_EDIT, lines=1)
+                    e_prompt = gr.Textbox(label="Edit instruction", lines=2,
+                                          value="change the background to a city street")
+                    e_neg = gr.Textbox(label="Negative prompt", value=" ", lines=1)
+                    e_ref = gr.Image(type="pil", label="Reference image", height=280,
+                                     value=os.path.join(os.path.dirname(__file__), "assets", "dog.jpg"))
+                    e_extra = gr.File(file_count="multiple", type="filepath",
+                                      label="Extra references (optional, multi-image edit)")
+                    with gr.Row():
+                        e_steps = gr.Slider(1, 50, value=30, step=1, label="Steps")
+                        e_cfg = gr.Slider(1.0, 10.0, value=5.0, step=0.5, label="CFG")
+                    e_max = gr.Slider(0, 1536, value=1024, step=16,
+                                      label="Max output side (0 = keep source size)")
+                    e_seed = gr.Number(value=42, precision=0, label="Seed")
+                    e_btn = gr.Button("Edit", variant="primary")
+                with gr.Column(scale=1):
+                    e_out = gr.Image(type="pil", label="Output", height=560)
+            e_btn.click(lambda: None, None, e_out).then(
+                        run_edit,
+                        [e_model, e_custom, e_prompt, e_neg, e_ref, e_extra, e_steps, e_cfg, e_max, e_seed],
+                        e_out)
+    return demo
+
+
+def main():
+    global DEVICE
+    ap = argparse.ArgumentParser()
+    ap.add_argument("--device", default="cuda")
+    ap.add_argument("--host", default="0.0.0.0")
+    ap.add_argument("--port", type=int, default=7860)
+    ap.add_argument("--share", action="store_true")
+    ap.add_argument("--preload", default=None,
+                    help="comma-separated repo ids / paths to load at startup (else lazy)")
+    args = ap.parse_args()
+    DEVICE = args.device
+    if args.preload:
+        for repo in args.preload.split(","):
+            _get_pipe(repo.strip())
+    build_ui().queue().launch(server_name=args.host, server_port=args.port,
+                              share=args.share, theme=gr.themes.Soft())
+
+
+if __name__ == "__main__":
+    main()

vendor/mage_flow/inference.py

@@ -0,0 +1,161 @@
+#!/usr/bin/env python
+"""MageFlow command-line inference.
+
+Two console scripts share this module:
+
+  * ``mage-flow``      — text-to-image generation
+  * ``mage-flow-edit`` — instruction-based image editing
+
+Both are BATCHED: pass several prompts and they are packed into a single
+transformer forward per denoise step. Sample ``i`` uses seed ``--seed + i``.
+
+Text-to-image (multiple prompts = a batch)::
+
+    mage-flow --prompt "a cat holding a sign that says hello" "a red ferrari" \
+        --model_path microsoft/Mage-Flow-4B --steps 30 --cfg 5.0 --out ./outputs
+
+``--model_path`` accepts a local repo dir OR a Hugging Face Hub repo id (e.g.
+``microsoft/Mage-Flow-4B``), downloaded and cached automatically on first use.
+
+Mixed resolutions — give one ``--height``/``--width`` per prompt (they are packed
+into a single forward per step regardless of shape)::
+
+    mage-flow --prompt "a tall waterfall" "a wide desert panorama" \
+        --height 2048 512 --width 512 2048 \
+        --model_path microsoft/Mage-Flow-4B --out ./outputs
+
+Image editing (one ``--ref`` entry per prompt; comma-separate paths for
+multi-image edit)::
+
+    mage-flow-edit \
+        --prompt "把背景改为城市街道" "把这两张图融合在一起" \
+        --ref hydrant.png "scene.png,object.png" \
+        --model_path /path/Mage-Flow-Edit-4B-Base --out ./outputs
+"""
+import argparse
+import os
+
+from mage_flow import MageFlowPipeline
+
+
+def _add_common_args(p):
+    p.add_argument("--model_path", required=True,
+                   help="local diffusers-style repo dir OR a Hugging Face Hub repo id "
+                        "(e.g. microsoft/Mage-Flow-4B); HF ids are downloaded and cached "
+                        "automatically on first use")
+    p.add_argument("--neg_prompt", default=None,
+                   help="negative prompt applied to every sample (default: a single space)")
+    p.add_argument("--steps", type=int, default=30)
+    p.add_argument("--cfg", type=float, default=5.0)
+    p.add_argument("--seed", type=int, default=42, help="base seed; sample i uses seed + i")
+    p.add_argument("--static_shift", type=float, default=None,
+                   help="override scheduler shift (default: repo scheduler_config.json, 6.0)")
+    p.add_argument("--device", default="cuda")
+    p.add_argument("--out", default="./outputs")
+
+
+def _neg_list(neg_prompt, n):
+    return [neg_prompt] * n if neg_prompt is not None else None
+
+
+def _size_list(vals, n, name, parser):
+    """Broadcast one size value to all prompts, or use a per-prompt list."""
+    if len(vals) == 1:
+        return vals * n
+    if len(vals) == n:
+        return vals
+    parser.error(f"--{name} expects 1 value (applied to all) or {n} values "
+                 f"(one per prompt); got {len(vals)}")
+
+
+def main():
+    """``mage-flow`` — batched text-to-image generation."""
+    p = argparse.ArgumentParser(
+        prog="mage-flow", description="MageFlow text-to-image generation.",
+        formatter_class=argparse.RawDescriptionHelpFormatter)
+    p.add_argument("--prompt", nargs="+", required=True,
+                   help="one or more prompts; multiple prompts are batched")
+    p.add_argument("--height", type=int, nargs="+", default=[1024],
+                   help="output height (multiple of 16); one value applied to all "
+                        "prompts, or one value per prompt for mixed resolutions")
+    p.add_argument("--width", type=int, nargs="+", default=[1024],
+                   help="output width (multiple of 16); one value applied to all "
+                        "prompts, or one value per prompt for mixed resolutions")
+    p.add_argument("--prompt_template", default="mage-flow")
+    _add_common_args(p)
+    args = p.parse_args()
+
+    os.makedirs(args.out, exist_ok=True)
+    pipe = MageFlowPipeline.from_pretrained(args.model_path, args.device)
+    n = len(args.prompt)
+    imgs = pipe.generate(
+        args.prompt,
+        neg_prompts=_neg_list(args.neg_prompt, n),
+        seeds=[args.seed + i for i in range(n)],
+        heights=_size_list(args.height, n, "height", p),
+        widths=_size_list(args.width, n, "width", p),
+        steps=args.steps, cfg=args.cfg, static_shift=args.static_shift,
+        prompt_template=args.prompt_template,
+    )
+    for i, im in enumerate(imgs):
+        path = os.path.join(args.out, f"gen_{i:03d}.png")
+        im.save(path)
+        print(f"saved {path}")
+
+
+def main_edit():
+    """``mage-flow-edit`` — batched instruction-based image editing."""
+    p = argparse.ArgumentParser(
+        prog="mage-flow-edit", description="MageFlow instruction-based image editing.",
+        formatter_class=argparse.RawDescriptionHelpFormatter)
+    p.add_argument("--prompt", nargs="+", required=True,
+                   help="one or more edit instructions; multiple prompts are batched")
+    p.add_argument("--ref", nargs="+", required=True,
+                   help="reference image(s) for each prompt, positionally aligned with "
+                        "--prompt; comma-separate paths for multi-image edit "
+                        "(e.g. --ref a.png 'b.png,c.png')")
+    p.add_argument("--max_size", type=int, default=None,
+                   help="longest side of the output (short side by aspect ratio). "
+                        "Default: keep each source image's own resolution")
+    p.add_argument("--height", type=int, default=None,
+                   help="explicit output height (use with --width; overrides --max_size)")
+    p.add_argument("--width", type=int, default=None,
+                   help="explicit output width (use with --height; overrides --max_size)")
+    p.add_argument("--vl_cond_long_edge", type=int, default=384,
+                   help="cap the long edge of the reference image fed to the VL text "
+                        "encoder (matches training preprocessing; the VAE path keeps the "
+                        "full output resolution). 0 or negative disables the cap")
+    p.add_argument("--prompt_template", default="mage-flow-edit")
+    _add_common_args(p)
+    args = p.parse_args()
+
+    if len(args.ref) != len(args.prompt):
+        p.error(f"--ref count ({len(args.ref)}) must match --prompt count ({len(args.prompt)})")
+
+    os.makedirs(args.out, exist_ok=True)
+    pipe = MageFlowPipeline.from_pretrained(args.model_path, args.device)
+    n = len(args.prompt)
+    # Each --ref token is one prompt's reference(s); commas split multi-image refs.
+    ref_images = [[s.strip() for s in r.split(",") if s.strip()] for r in args.ref]
+
+    size_kw = {}
+    if args.height and args.width:
+        size_kw = {"heights": [args.height] * n, "widths": [args.width] * n}
+
+    outs = pipe.edit(
+        args.prompt, ref_images,
+        neg_prompts=_neg_list(args.neg_prompt, n),
+        seeds=[args.seed + i for i in range(n)],
+        max_size=args.max_size, steps=args.steps, cfg=args.cfg,
+        static_shift=args.static_shift, prompt_template=args.prompt_template,
+        vl_cond_long_edge=args.vl_cond_long_edge,
+        **size_kw,
+    )
+    for i, im in enumerate(outs):
+        path = os.path.join(args.out, f"edit_{i:03d}.png")
+        im.save(path)
+        print(f"saved {path}")
+
+
+if __name__ == "__main__":
+    main()

vendor/mage_flow/models/__init__.py

@@ -0,0 +1,3 @@
+from .utils import load_model
+
+__all__ = ["load_model"]

vendor/mage_flow/models/mage_flow.py

@@ -0,0 +1,364 @@
+from dataclasses import dataclass
+from typing import Any
+
+import torch
+import torch.nn as nn
+from einops import rearrange, repeat
+from loguru import logger
+from pydantic import BaseModel, Field
+from torch import Tensor
+
+from .modules._attn_backend import set_attn_backend
+from .modules.mage_layers import (
+    AdaLayerNormContinuous,
+    MageFlowEmbedRope,
+    MageFlowTimestepProjEmbeddings,
+    MageFlowTransformerBlock,
+    RMSNorm,
+)
+from .modules.text_encoder import TextEncoder, qwen3_patch_forward
+
+
+class ModelConfig(BaseModel):
+    static_shift: float = Field(
+        default=6.0,
+        description="Static shift value for the z-image time-shift schedule (the only "
+                    "supported schedule). Default: 6.0.",
+    )
+    vae_path: str = Field(...)
+    model_structure: dict = Field(default_factory=dict)
+    txt_enc_path: str = Field(...)
+    txt_max_length: int = Field(default=4096)
+    pretrained_model_name_or_path: str | None = Field(default=None)
+    pretrained_full_model_path: str | None = Field(default=None)  # Load full model weights (DiT + txt_enc + vae)
+    packing: bool = Field(default=False)
+    vae_sample_posterior: bool = Field(default=True)   # Sample (vs mode) from VAE posterior at encode time (Flux2 + CoD)
+    vae_encoder_only: bool = Field(default=False)  # Skip loading VAE decoder to save GPU memory (training only, MageVAE)
+    compile_vae_encoder: bool = Field(default=False)  # torch.compile VAE encoder to reduce CUDA kernel launch overhead
+    attn_type: str = Field(
+        default="flash2",
+        description="Flash-attn backend used by both the DiT (mage_layers) "
+                    "and the HF text encoder (text_encoder). One of: 'flash2' (default) or 'flash4'.",
+    )
+
+
+@dataclass
+class MageFlowParams:
+    in_channels: int
+    out_channels: int
+    context_in_dim: int
+    hidden_size: int
+    num_heads: int
+    depth: int
+    axes_dim: list[int]
+    checkpoint: bool
+    patch_size: int = 1
+
+
+class MageFlow(nn.Module):
+    def __init__(self, params: MageFlowParams):
+        super().__init__()
+        self.params = params
+        self.checkpoint = params.checkpoint
+        self.in_channels = params.in_channels
+        self.out_channels = params.out_channels
+        self.inner_dim = params.hidden_size  # num_attention_heads * attention_head_dim
+        self.axes_dim = params.axes_dim
+        self.num_attention_heads = params.num_heads
+        self.attention_head_dim = self.inner_dim // self.num_attention_heads
+        self.patch_size = params.patch_size
+        assert sum(self.axes_dim) == self.attention_head_dim
+
+        self.pos_embed = MageFlowEmbedRope(theta=10000, axes_dim=self.axes_dim, scale_rope=True)
+        self.img_in = nn.Linear(self.in_channels, self.inner_dim)
+        self.txt_norm = RMSNorm(params.context_in_dim, eps=1e-6)
+        self.txt_in = nn.Linear(params.context_in_dim, self.inner_dim)
+
+        self.time_text_embed = MageFlowTimestepProjEmbeddings(embedding_dim=self.inner_dim)
+
+        self.transformer_blocks = nn.ModuleList(
+            [
+                MageFlowTransformerBlock(
+                    dim=self.inner_dim,
+                    num_attention_heads=self.num_attention_heads,
+                    attention_head_dim=self.attention_head_dim,
+                )
+                for _ in range(params.depth)
+            ]
+        )
+
+        self.norm_out = AdaLayerNormContinuous(self.inner_dim, self.inner_dim, elementwise_affine=False, eps=1e-6)
+        self.proj_out = nn.Linear(self.inner_dim, self.patch_size * self.patch_size * self.out_channels, bias=True)
+
+    def forward(
+        self,
+        img: Tensor,
+        txt: Tensor,
+        timesteps: Tensor,
+        img_shapes=None,
+        img_cu_seqlens: Tensor | None = None,
+        txt_cu_seqlens: Tensor | None = None,
+        attention_kwargs: dict[str, Any] | None = None,
+    ) -> Tensor:
+        if img.ndim != 3 or txt.ndim != 3:
+            raise ValueError("Input img and txt tensors must have 3 dimensions.")
+
+        # Prepare vision RoPE (msrope); text tokens are not rotated.
+        ms_pe = self.pos_embed(img_shapes, device=img.device)
+
+        img = self.img_in(img)
+        txt = self.txt_norm(txt)
+
+        timesteps = timesteps.to(img.dtype)
+        temb = self.time_text_embed(timesteps, img)
+
+        txt = self.txt_in(txt)
+        txt_vec = torch.zeros(txt.shape[0], self.inner_dim, dtype=txt.dtype, device=txt.device)
+
+        temb = temb + txt_vec
+
+        attention_kwargs = attention_kwargs or {}
+
+        for _index_block, block in enumerate(self.transformer_blocks):
+            if self.training and self.checkpoint:
+                txt, img = torch.utils.checkpoint.checkpoint(
+                    block,
+                    img,  # hidden_states
+                    txt,  # encoder_hidden_states
+                    temb,  # temb
+                    ms_pe,  # image_rotary_emb
+                    txt_cu_seqlens,  # txt_cu_lens
+                    img_cu_seqlens,  # img_cu_lens
+                    use_reentrant=False,
+                )
+
+            else:
+                txt, img = block(
+                    hidden_states=img,
+                    encoder_hidden_states=txt,
+                    txt_cu_lens=txt_cu_seqlens,
+                    img_cu_lens=img_cu_seqlens,
+                    temb=temb,
+                    image_rotary_emb=ms_pe,
+                    joint_attention_kwargs=attention_kwargs,
+                )
+
+        # Use only the image part (hidden_states) from the dual-stream blocks
+        img = self.norm_out(
+            img,
+            temb,
+            cu_seqlens=img_cu_seqlens,
+        )
+        img = self.proj_out(img)
+        return img
+
+
+class MageFlowModel(nn.Module):
+    def __init__(self, config: ModelConfig):
+        super().__init__()
+        self.config = config
+        set_attn_backend(getattr(config, "attn_type", "flash2"))
+        self.patch_text_encoder_forward()
+        self.vae = self.load_vae()
+        self.transformer = self.load_transformer()
+        self.txt_enc = self.load_text_enc()
+
+        # Optionally override all components from a full model checkpoint (e.g. ema.pt)
+        full_path = getattr(self.config, "pretrained_full_model_path", None)
+        if full_path is not None:
+            import os
+
+            if os.path.exists(full_path):
+                logger.info(f"Loading full model weights from {full_path}")
+                sd = torch.load(full_path, map_location="cpu")
+                # Handle wrapped EMA format: {'ema_state_dict': ..., ...}
+                if isinstance(sd, dict) and "ema_state_dict" in sd:
+                    sd = sd["ema_state_dict"]
+                missing, unexpected = self.load_state_dict(sd, strict=False)
+                if missing:
+                    logger.warning(f"Full model load missing keys ({len(missing)}): {missing[:5]}...")
+                if unexpected:
+                    logger.warning(f"Full model load unexpected keys ({len(unexpected)}): {unexpected[:5]}...")
+                logger.info("Full model weights loaded successfully.")
+            else:
+                logger.warning(f"pretrained_full_model_path not found: {full_path}")
+
+        # Freeze VAE and Text Encoder
+        self.vae.requires_grad_(False)
+
+        # Drop VAE decoder to save GPU memory (training only, decoder unused during training)
+        if self.config.vae_encoder_only:
+            from .modules.mage_vae import MageVAE
+            if isinstance(self.vae, MageVAE):
+                decoder_params = sum(p.numel() for p in self.vae.decoder_model.parameters()) / 1e6
+                self.vae.decoder_model = None
+            elif hasattr(self.vae, "decoder"):
+                decoder_params = sum(p.numel() for p in self.vae.decoder.parameters()) / 1e6
+                self.vae.decoder = None
+            else:
+                decoder_params = 0
+            logger.info(f"vae_encoder_only=True: dropped VAE decoder ({decoder_params:.1f}M params) to save memory")
+
+        # NOTE: VAE encoder torch.compile() is deferred to
+        # maybe_compile_vae_encoder(), called after checkpoint load. Reason:
+        # avoid wasted compile work before load_checkpoint overwrites weights.
+        # The save-side _unwrap_compiled_submodules guard in DeepSpeedTrainer
+        # is a belt-and-suspenders defense against any future code that
+        # re-introduces the function-style ``module = torch.compile(module)``
+        # pattern (which does pollute state_dict with ``_orig_mod.``).
+
+        # Text encoder is always frozen (inference only).
+        self.txt_enc.requires_grad_(False)
+        logger.info(
+            f"{sum([p.numel() for p in self.transformer.parameters() if p.requires_grad]) / 1000000} M parameters"
+        )
+
+    def patch_text_encoder_forward(self):
+        qwen3_patch_forward()
+        logger.info("Patched Qwen3-VL text encoder forward methods")
+
+    def maybe_compile_vae_encoder(self) -> None:
+        """Compile the VAE encoder with torch.compile() to fuse small ops and
+        reduce CUDA kernel launch overhead.
+
+        Uses ``nn.Module.compile()`` (in-place) for the encoder so the module
+        hierarchy and parameter names are unchanged — ``state_dict()`` keeps
+        clean keys (no ``_orig_mod.`` prefix), and checkpoints stay
+        interchangeable with the non-compiled path.
+
+        For the MageVAE branch we still assign ``torch.compile(...)`` to a
+        method (``_encode_moments``); methods aren't ``nn.Module``s so this
+        does not pollute ``state_dict()``.
+
+        Idempotent: safe to call multiple times; already-compiled modules are
+        detected and skipped.
+        """
+        if not getattr(self.config, "compile_vae_encoder", False):
+            return
+        torch.set_float32_matmul_precision("high")
+        from .modules.mage_vae import MageVAE
+        if isinstance(self.vae, MageVAE):
+            fn = self.vae._encode_moments
+            if hasattr(fn, "_torchdynamo_orig_callable") or hasattr(fn, "_orig_mod"):
+                return  # already compiled
+            self.vae._encode_moments = torch.compile(fn, dynamic=True)
+            logger.info("compile_vae_encoder=True: compiled MageVAE._encode_moments")
+        elif hasattr(self.vae, "encoder"):
+            if getattr(self.vae.encoder, "_compiled_call_impl", None) is not None:
+                return  # already compiled
+            self.vae.encoder.compile()
+            logger.info("compile_vae_encoder=True: compiled VAE encoder (in-place)")
+
+    def load_text_enc(self):
+        return TextEncoder(
+            model_name=self.config.txt_enc_path,
+            version=self.config.txt_enc_path,
+            tokenizer_max_length=self.config.txt_max_length,
+            torch_dtype=torch.bfloat16,
+            prompt_template=None,
+            dit_structure=self.config.model_structure,
+            use_packed_text_infer=self.config.packing,
+            attn_type=getattr(self.config, "attn_type", "flash2"),
+        )
+
+    def load_vae(self):
+        from .modules.mage_vae import MageVAE
+        return MageVAE(
+            ckpt_path=self.config.vae_path,
+            sample_posterior=self.config.vae_sample_posterior,
+        )
+
+    def load_transformer(self):
+        # Imported lazily to avoid a circular import: ``utils`` imports MageFlow /
+        # MageFlowParams from this module.
+        from .utils import load_model
+        return load_model(
+            dit_structure=self.config.model_structure,
+            pretrain_path=self.config.pretrained_model_name_or_path,
+        )
+
+    def compile(self):
+        self.transformer.compile()
+
+    def compute_vae_encodings(
+        self,
+        pixel_values: torch.Tensor | list[torch.Tensor],
+        with_ids: bool = True,
+    ):
+        if isinstance(pixel_values, list):
+            # All same resolution → batch encode via the tensor path
+            if len(pixel_values) > 1 and len({img.shape for img in pixel_values}) == 1:
+                stacked = torch.stack(pixel_values, dim=0)
+                result = self.compute_vae_encodings(stacked, with_ids=with_ids)
+                # Repack from [N, L, C] batch format to [1, N*L, C] packed format
+                if with_ids:
+                    model_input, img_shapes, img_ids = result
+                    model_input = model_input.reshape(1, -1, model_input.shape[-1])
+                    img_ids = img_ids.reshape(1, -1, img_ids.shape[-1])
+                    return model_input, img_shapes, img_ids
+                model_input, img_shapes = result
+                model_input = model_input.reshape(1, -1, model_input.shape[-1])
+                return model_input, img_shapes
+
+            # Packed / variable-size images
+            model_inputs = []
+            img_shapes = []
+            img_ids_list = []
+
+            def _append(latents):
+                _, _, h, w = latents.shape
+                img_shapes.append([(1, h, w)])
+                model_inputs.append(rearrange(latents, "b c h w -> b (h w) c").squeeze(0))
+                if with_ids:
+                    ids = torch.zeros(h, w, 3, device=latents.device)
+                    ids[..., 1] = ids[..., 1] + torch.arange(h, device=latents.device)[:, None]
+                    ids[..., 2] = ids[..., 2] + torch.arange(w, device=latents.device)[None, :]
+                    img_ids_list.append(rearrange(ids, "h w c -> (h w) c"))
+
+            # MageVAE encoder is launch-bound on B=1; group same-shape images
+            # in the pack into one batched encode call.
+            if len(pixel_values) > 1:
+                groups: dict[tuple[int, int], list[int]] = {}
+                for i, img in enumerate(pixel_values):
+                    key = (int(img.shape[-2]), int(img.shape[-1]))
+                    groups.setdefault(key, []).append(i)
+                latents_per_idx = [None] * len(pixel_values)
+                for (h, w), idxs in groups.items():
+                    batch = torch.stack([pixel_values[i] for i in idxs], dim=0)
+                    batch = batch.to(memory_format=torch.contiguous_format).float()
+                    batch = batch.to(self.vae.device, dtype=self.vae.dtype)
+                    with torch.no_grad():
+                        lat = self.vae.encode(batch)  # [B, 128, H/16, W/16]
+                    for j, i in enumerate(idxs):
+                        latents_per_idx[i] = lat[j:j + 1]
+                for latents in latents_per_idx:
+                    _append(latents)
+            else:
+                for img in pixel_values:
+                    img = img.unsqueeze(0).to(memory_format=torch.contiguous_format).float()
+                    img = img.to(self.vae.device, dtype=self.vae.dtype)
+                    with torch.no_grad():
+                        latents = self.vae.encode(img)  # [1, 128, H/16, W/16]
+                    _append(latents)
+
+            model_input = torch.cat(model_inputs, dim=0).unsqueeze(0)
+            if with_ids:
+                img_ids = torch.cat(img_ids_list, dim=0).unsqueeze(0)
+                return model_input, img_shapes, img_ids
+            return model_input, img_shapes
+
+        # Tensor (padded batch)
+        pixel_values = pixel_values.to(memory_format=torch.contiguous_format).float()
+        pixel_values = pixel_values.to(self.vae.device, dtype=self.vae.dtype)
+        with torch.no_grad():
+            model_input = self.vae.encode(pixel_values)  # [B, 128, H/16, W/16]
+        bs, c, h, w = model_input.shape
+        img_shapes = [[(1, h, w)]] * bs
+        model_input = rearrange(model_input, "b c h w -> b (h w) c")
+        if with_ids:
+            img_ids = torch.zeros(h, w, 3, device=model_input.device)
+            img_ids[..., 1] = img_ids[..., 1] + torch.arange(h, device=model_input.device)[:, None]
+            img_ids[..., 2] = img_ids[..., 2] + torch.arange(w, device=model_input.device)[None, :]
+            img_ids = repeat(img_ids, "h w c -> b (h w) c", b=bs)
+            return model_input, img_shapes, img_ids
+        return model_input, img_shapes

vendor/mage_flow/models/modules/_attn_backend.py

@@ -0,0 +1,229 @@
+"""Attention backend shim — switchable between Flash Attention 2 and 4.
+
+Exports a single ``flash_attn_varlen_func`` with the FA2 calling convention.
+The underlying kernel is selected at runtime via ``set_attn_backend(name)``
+(default: ``"flash2"``). The selected kernel is resolved lazily on the first
+call so model-config-driven selection (which happens after this module is
+imported) takes effect.
+
+Modules that previously did ``from flash_attn import flash_attn_varlen_func``
+should import from here instead.
+
+For the FA4 path, calling-convention differences are normalised:
+
+* ``window_size=(-1, -1)`` (FA2 "no window") -> ``(None, None)`` (FA4).
+* ``block_table`` -> ``page_table``.
+* FA4's optional ``(out, lse)`` tuple return is unwrapped to ``out``.
+* ``dropout_p>0`` / ``alibi_slopes`` / ``return_attn_probs`` raise on FA4.
+"""
+
+from __future__ import annotations
+
+from typing import Any, Callable
+
+_FA2_ALIASES = {"flash2", "fa2", "flash_attention_2", "flash_attn_2"}
+_FA4_ALIASES = {"flash4", "fa4", "flash_attention_4", "flash_attn_4"}
+_SDPA_ALIASES = {"sdpa", "torch_sdpa", "scaled_dot_product_attention"}
+
+_BACKEND: str = "flash2"
+_RESOLVED_FN: Callable[..., Any] | None = None
+
+
+def _normalize(name: str) -> str:
+    n = name.lower().strip()
+    if n in _FA2_ALIASES:
+        return "flash2"
+    if n in _FA4_ALIASES:
+        return "flash4"
+    if n in _SDPA_ALIASES:
+        return "sdpa"
+    raise ValueError(
+        f"Unknown attention backend {name!r}; expected one of "
+        f"{sorted(_FA2_ALIASES | _FA4_ALIASES | _SDPA_ALIASES)}"
+    )
+
+
+def set_attn_backend(name: str) -> None:
+    """Select the flash-attn backend used by ``flash_attn_varlen_func``.
+
+    Safe to call multiple times; clears the cached resolution on change.
+    """
+    global _BACKEND, _RESOLVED_FN
+    new = _normalize(name)
+    if new != _BACKEND:
+        _RESOLVED_FN = None
+    _BACKEND = new
+
+
+
+def _resolve_fa2() -> Callable[..., Any]:
+    from flash_attn import flash_attn_varlen_func as _fn
+    return _fn
+
+
+def _resolve_fa4() -> Callable[..., Any]:
+    from flash_attn.cute import flash_attn_varlen_func as _fa4_fn
+
+    def _fa4_wrapper(
+        q,
+        k,
+        v,
+        cu_seqlens_q=None,
+        cu_seqlens_k=None,
+        max_seqlen_q=None,
+        max_seqlen_k=None,
+        dropout_p: float = 0.0,
+        softmax_scale=None,
+        causal: bool = False,
+        window_size=(-1, -1),
+        softcap: float = 0.0,
+        alibi_slopes=None,
+        deterministic: bool = False,
+        return_attn_probs: bool = False,
+        block_table=None,
+        **_unused: Any,
+    ):
+        if dropout_p and dropout_p > 0:
+            raise NotImplementedError("FA4 backend does not support dropout_p>0")
+        if alibi_slopes is not None:
+            raise NotImplementedError("FA4 backend does not support alibi_slopes")
+        if return_attn_probs:
+            raise NotImplementedError("FA4 backend does not support return_attn_probs")
+
+        win_l, win_r = window_size
+        if win_l == -1:
+            win_l = None
+        if win_r == -1:
+            win_r = None
+
+        out = _fa4_fn(
+            q,
+            k,
+            v,
+            cu_seqlens_q=cu_seqlens_q,
+            cu_seqlens_k=cu_seqlens_k,
+            max_seqlen_q=max_seqlen_q,
+            max_seqlen_k=max_seqlen_k,
+            softmax_scale=softmax_scale,
+            causal=causal,
+            window_size=(win_l, win_r),
+            softcap=softcap,
+            deterministic=deterministic,
+            page_table=block_table,
+            return_lse=False,
+        )
+        if isinstance(out, tuple):
+            out = out[0]
+        return out
+
+    return _fa4_wrapper
+
+
+def _resolve_sdpa() -> Callable[..., Any]:
+    """FA2 varlen → per-sequence torch.SDPA fallback.
+
+    Use when flash-attn is unavailable (e.g. CUDA 13 has no prebuilt wheel
+    and source build is brittle). Slower than FA2 (one SDPA dispatch per
+    sequence), but functionally equivalent for the dense / causal / no-alibi
+    paths mageflow actually uses. Window / softcap / alibi / paged-attn /
+    return_attn_probs are not supported and will raise.
+    """
+    import torch
+    import torch.nn.functional as F
+
+    def _sdpa_wrapper(
+        q,
+        k,
+        v,
+        cu_seqlens_q=None,
+        cu_seqlens_k=None,
+        max_seqlen_q=None,
+        max_seqlen_k=None,
+        dropout_p: float = 0.0,
+        softmax_scale=None,
+        causal: bool = False,
+        window_size=(-1, -1),
+        softcap: float = 0.0,
+        alibi_slopes=None,
+        deterministic: bool = False,
+        return_attn_probs: bool = False,
+        block_table=None,
+        **_unused: Any,
+    ):
+        if dropout_p and dropout_p > 0:
+            raise NotImplementedError("SDPA backend does not support dropout_p>0")
+        if alibi_slopes is not None:
+            raise NotImplementedError("SDPA backend does not support alibi_slopes")
+        if return_attn_probs:
+            raise NotImplementedError("SDPA backend does not support return_attn_probs")
+        if softcap and softcap > 0:
+            raise NotImplementedError("SDPA backend does not support softcap")
+        if window_size not in ((-1, -1), (None, None), (0, 0)):
+            raise NotImplementedError(
+                f"SDPA backend does not support sliding window (got {window_size})"
+            )
+        if block_table is not None:
+            raise NotImplementedError("SDPA backend does not support paged attention")
+        if cu_seqlens_q is None or cu_seqlens_k is None:
+            raise ValueError("SDPA backend requires cu_seqlens_q and cu_seqlens_k")
+
+        # GQA: FA2 broadcasts k/v across query head groups natively; torch SDPA
+        # does not (the q vs k head-dim mismatch is the AssertionError "tensor
+        # a (32) must match tensor b (8) at non-singleton dimension 1" we'd see
+        # otherwise). Repeat k/v along the head dim to match q before the loop.
+        n_heads_q = q.shape[1]
+        n_heads_kv = k.shape[1]
+        if n_heads_q != n_heads_kv:
+            if n_heads_q % n_heads_kv != 0:
+                raise ValueError(
+                    f"SDPA backend GQA expansion requires q heads ({n_heads_q}) "
+                    f"to be divisible by k/v heads ({n_heads_kv})"
+                )
+            repeat = n_heads_q // n_heads_kv
+            k = k.repeat_interleave(repeat, dim=1)
+            v = v.repeat_interleave(repeat, dim=1)
+
+        # q/k/v: (total_tokens, nheads, head_dim). Dispatch SDPA per sequence,
+        # then concat. Python-level loop is fine since nseq is small (one per
+        # image in the pack) and image-gen latency is dominated by sampling.
+        cu_q = cu_seqlens_q.tolist()
+        cu_k = cu_seqlens_k.tolist()
+        outs = []
+        for qs, qe, ks, ke in zip(cu_q[:-1], cu_q[1:], cu_k[:-1], cu_k[1:]):
+            # (s, h, d) → (1, h, s, d)
+            q_i = q[qs:qe].transpose(0, 1).unsqueeze(0)
+            k_i = k[ks:ke].transpose(0, 1).unsqueeze(0)
+            v_i = v[ks:ke].transpose(0, 1).unsqueeze(0)
+            out_i = F.scaled_dot_product_attention(
+                q_i,
+                k_i,
+                v_i,
+                attn_mask=None,
+                dropout_p=0.0,
+                is_causal=causal,
+                scale=softmax_scale,
+            )
+            # (1, h, s, d) → (s, h, d)
+            outs.append(out_i.squeeze(0).transpose(0, 1))
+        return torch.cat(outs, dim=0).contiguous()
+
+    return _sdpa_wrapper
+
+
+def _resolve() -> Callable[..., Any]:
+    global _RESOLVED_FN
+    if _RESOLVED_FN is None:
+        if _BACKEND == "flash4":
+            _RESOLVED_FN = _resolve_fa4()
+        elif _BACKEND == "sdpa":
+            _RESOLVED_FN = _resolve_sdpa()
+        else:
+            _RESOLVED_FN = _resolve_fa2()
+    return _RESOLVED_FN
+
+
+def flash_attn_varlen_func(*args, **kwargs):
+    return _resolve()(*args, **kwargs)
+
+
+__all__ = ["flash_attn_varlen_func", "set_attn_backend"]

vendor/mage_flow/models/modules/mage_latent.py

@@ -0,0 +1,119 @@
+from __future__ import annotations
+
+import hashlib
+import os
+from math import erfc, sqrt
+
+import numpy as np
+import torch
+
+DEFAULT_GS_PAYLOAD = "MageFlow"
+
+
+_ENV_KEY = "MAGEFLOW_GS_KEY"
+_ENV_KEYFILE = "MAGEFLOW_GS_KEY_FILE"
+_DEFAULT_KEYFILE = os.path.expanduser("~/.mageflow/gs_key")
+DEFAULT_GS_KEY = 20260720
+
+# Payload length in bits. Short + heavily replicated across the latent.
+_MSG_BITS = 256
+
+
+def _key_to_int(value) -> int:
+    """Normalize a key (int / digit-string / passphrase) to a non-negative int.
+
+    A pure integer or all-digits string is used directly; anything else is
+    treated as a passphrase and hashed (SHA-256) into a 256-bit integer.
+    """
+    if isinstance(value, int):
+        return abs(value)
+    s = str(value).strip()
+    if not s:
+        raise ValueError("empty Gaussian-Shading key")
+    if s.lstrip("-").isdigit():
+        return abs(int(s))
+    return int.from_bytes(hashlib.sha256(s.encode()).digest(), "big")
+
+
+def resolve_gs_key(explicit=None):
+    if explicit is not None:
+        return _key_to_int(explicit)
+    env = os.environ.get(_ENV_KEY)
+    if env and env.strip():
+        return _key_to_int(env)
+    keyfile = os.environ.get(_ENV_KEYFILE) or _DEFAULT_KEYFILE
+    try:
+        with open(keyfile) as fh:
+            content = fh.read().strip()
+        if content:
+            return _key_to_int(content)
+    except OSError:
+        pass
+    return _key_to_int(DEFAULT_GS_KEY)
+
+
+def _payload_to_bits(payload: str, n_bits: int = _MSG_BITS) -> np.ndarray:
+    """Deterministically expand an arbitrary string into an ``n_bits`` bit vector."""
+    out: list[int] = []
+    counter = 0
+    while len(out) < n_bits:
+        digest = hashlib.sha256(f"{payload}:{counter}".encode()).digest()
+        for byte in digest:
+            for k in range(8):
+                out.append((byte >> k) & 1)
+        counter += 1
+    return np.asarray(out[:n_bits], dtype=np.int64)
+
+
+def _pad_and_pos(n: int, key, n_bits: int = _MSG_BITS):
+    """Key-seeded per-entry XOR pad and message-index map (length ``n``)."""
+    rng = np.random.default_rng(_key_to_int(key))
+    pad = rng.integers(0, 2, size=n).astype(np.int64)        # XOR mask
+    pos = rng.integers(0, n_bits, size=n).astype(np.int64)   # msg index per entry
+    return pad, pos
+
+
+def encode_noise(shape, *, key,
+                 seed: int = 0, device=None, dtype=torch.bfloat16) -> torch.Tensor:
+    C, H, W = shape
+    n = C * H * W
+    msg = _payload_to_bits(DEFAULT_GS_PAYLOAD)
+    pad, pos = _pad_and_pos(n, key)
+    target_half = (msg[pos] ^ pad).astype(np.float64)        # {0,1} per entry
+
+    gen = torch.Generator(device="cpu").manual_seed(int(seed) & 0x7FFFFFFF)
+    u = torch.rand(n, generator=gen, dtype=torch.float64)    # U(0,1) magnitudes
+    half = torch.from_numpy(target_half)
+    arg = ((half + u) / 2.0).clamp(1e-6, 1.0 - 1e-6)
+    z = torch.special.ndtri(arg)                             # inverse normal CDF
+    z = z.reshape(1, C, H, W)
+    return z.to(device=device, dtype=dtype)
+
+
+def decode_bits(noise: torch.Tensor, *, key) -> dict:
+    z = noise.detach().float().reshape(-1).cpu()
+    n = int(z.numel())
+    msg = _payload_to_bits(DEFAULT_GS_PAYLOAD)
+    pad, pos = _pad_and_pos(n, key)
+
+    observed_half = (z > 0).numpy().astype(np.int64)         # sign -> half
+    expected_half = (msg[pos] ^ pad)
+    matches = int((observed_half == expected_half).sum())
+    raw_acc = matches / n
+
+    # Recover payload by majority vote of each entry's implied message bit.
+    implied = observed_half ^ pad                            # estimate of m[pos]
+    votes = np.zeros((_MSG_BITS, 2), dtype=np.int64)
+    np.add.at(votes, (pos, implied), 1)
+    msg_hat = votes.argmax(axis=1)
+    msg_acc = float((msg_hat == msg).mean())
+
+    # One-sided significance under Binomial(n, 0.5) via normal approximation.
+    z_score = (matches - 0.5 * n) / (0.5 * sqrt(n))
+    pvalue = 0.5 * erfc(z_score / sqrt(2))
+
+    return {
+        "raw_acc": raw_acc, "msg_acc": msg_acc, "matches": matches, "n": n,
+        "z_score": z_score, "pvalue": pvalue, "present": pvalue < 1e-6,
+        "msg_hat": msg_hat, "msg": msg,
+    }

vendor/mage_flow/models/modules/mage_layers.py

@@ -0,0 +1,725 @@
+import math
+from typing import Any
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from diffusers.models.attention import FeedForward
+from diffusers.models.embeddings import TimestepEmbedding
+from diffusers.models.normalization import RMSNorm
+from ._attn_backend import flash_attn_varlen_func
+from torch import Tensor
+from torch._dynamo import allow_in_graph as maybe_allow_in_graph
+
+
+def apply_rotary_emb_mageflow(x: torch.Tensor, freqs_cis: torch.Tensor) -> torch.Tensor:
+    """Apply complex rotary embeddings to `x` ([B, S, H, D]) using `freqs_cis`
+    (the MageFlowEmbedRope 2D multi-scale RoPE, adjacent-pair complex convention)."""
+    x_rotated = torch.view_as_complex(x.float().reshape(*x.shape[:-1], -1, 2))
+    freqs_cis = freqs_cis.unsqueeze(1)
+    x_out = torch.view_as_real(x_rotated * freqs_cis).flatten(-2)
+    return x_out.type_as(x)
+
+
+def get_timestep_embedding(
+    timesteps: torch.Tensor,
+    embedding_dim: int,
+    flip_sin_to_cos: bool = False,
+    downscale_freq_shift: float = 1,
+    scale: float = 1,
+    max_period: int = 10000,
+) -> torch.Tensor:
+    """Sinusoidal timestep embeddings (DDPM convention).
+
+    NOTE: kept vendored (not diffusers') because the frequency table is
+    downcast to ``timesteps.dtype`` (bf16) here — the model was trained with
+    this exact bf16 rounding, so diffusers' fp32 variant produces a slightly
+    different embedding and degrades outputs.
+    """
+    assert len(timesteps.shape) == 1, "Timesteps should be a 1d-array"
+
+    half_dim = embedding_dim // 2
+    exponent = -math.log(max_period) * torch.arange(start=0, end=half_dim, dtype=torch.float32, device=timesteps.device)
+    exponent = exponent / (half_dim - downscale_freq_shift)
+
+    emb = torch.exp(exponent).to(timesteps.dtype)
+    emb = timesteps[:, None].float() * emb[None, :]
+
+    # scale embeddings
+    emb = scale * emb
+
+    # concat sine and cosine embeddings
+    emb = torch.cat([torch.sin(emb), torch.cos(emb)], dim=-1)
+
+    # flip sine and cosine embeddings
+    if flip_sin_to_cos:
+        emb = torch.cat([emb[:, half_dim:], emb[:, :half_dim]], dim=-1)
+
+    # zero pad
+    if embedding_dim % 2 == 1:
+        emb = torch.nn.functional.pad(emb, (0, 1, 0, 0))
+    return emb
+
+
+class Timesteps(nn.Module):
+    def __init__(
+        self,
+        num_channels: int,
+        flip_sin_to_cos: bool,
+        downscale_freq_shift: float,
+        scale: int = 1,
+    ):
+        super().__init__()
+        self.num_channels = num_channels
+        self.flip_sin_to_cos = flip_sin_to_cos
+        self.downscale_freq_shift = downscale_freq_shift
+        self.scale = scale
+
+    def forward(self, timesteps: torch.Tensor) -> torch.Tensor:
+        t_emb = get_timestep_embedding(
+            timesteps,
+            self.num_channels,
+            flip_sin_to_cos=self.flip_sin_to_cos,
+            downscale_freq_shift=self.downscale_freq_shift,
+            scale=self.scale,
+        )
+        return t_emb
+
+
+class MageFlowTimestepProjEmbeddings(nn.Module):
+    def __init__(self, embedding_dim):
+        super().__init__()
+
+        self.time_proj = Timesteps(num_channels=256, flip_sin_to_cos=True, downscale_freq_shift=0, scale=1000)
+        self.timestep_embedder = TimestepEmbedding(in_channels=256, time_embed_dim=embedding_dim)
+
+    def forward(self, timestep, hidden_states):
+        timesteps_proj = self.time_proj(timestep)
+        timesteps_emb = self.timestep_embedder(timesteps_proj.to(dtype=hidden_states.dtype))  # (N, D)
+
+        conditioning = timesteps_emb
+
+        return conditioning
+
+
+class MageFlowEmbedRope(nn.Module):
+    def __init__(self, theta: int, axes_dim: list[int], scale_rope=False):
+        super().__init__()
+        self.theta = theta
+        self.axes_dim = axes_dim
+        pos_index = torch.arange(4096)
+        neg_index = torch.arange(4096).flip(0) * -1 - 1
+        self.pos_freqs = torch.cat(
+            [
+                self.rope_params(pos_index, self.axes_dim[0], self.theta),
+                self.rope_params(pos_index, self.axes_dim[1], self.theta),
+                self.rope_params(pos_index, self.axes_dim[2], self.theta),
+            ],
+            dim=1,
+        )
+        self.neg_freqs = torch.cat(
+            [
+                self.rope_params(neg_index, self.axes_dim[0], self.theta),
+                self.rope_params(neg_index, self.axes_dim[1], self.theta),
+                self.rope_params(neg_index, self.axes_dim[2], self.theta),
+            ],
+            dim=1,
+        )
+
+        # DO NOT USING REGISTER BUFFER HERE, IT WILL CAUSE COMPLEX NUMBERS LOSE ITS IMAGINARY PART
+        self.scale_rope = scale_rope
+        self.video_freq_cache = {}
+
+    def rope_params(self, index, dim, theta=10000):
+        """
+        Args:
+            index: [0, 1, 2, 3] 1D Tensor representing the position index of the token
+        """
+        assert dim % 2 == 0
+        freqs = torch.outer(
+            index,
+            1.0 / torch.pow(theta, torch.arange(0, dim, 2).to(torch.float32).div(dim)),
+        )
+        freqs = torch.polar(torch.ones_like(freqs), freqs)
+        return freqs
+
+    def forward(
+        self,
+        video_fhw: tuple[int, int, int] | list[tuple[int, int, int]],
+        device: torch.device,
+        max_img_len: int = None,
+    ) -> torch.Tensor:
+        """Compute the vision RoPE frequencies (`vid_freqs`) for the packed image
+        tokens. Text tokens are NOT rotated, so no text RoPE is computed.
+
+        Args:
+            video_fhw (`Tuple[int, int, int]` or `List[Tuple[int, int, int]]`):
+                A list of 3 integers [frame, height, width] representing the shape of the video.
+            device: (`torch.device`):
+                The device on which to perform the RoPE computation.
+        """
+        if self.pos_freqs.device != device:
+            self.pos_freqs = self.pos_freqs.to(device)
+            self.neg_freqs = self.neg_freqs.to(device)
+
+        if isinstance(video_fhw, list):
+            video_fhw = video_fhw[0]
+        if not isinstance(video_fhw, list):
+            video_fhw = [video_fhw]
+
+        vid_freqs = []
+        for idx, fhw in enumerate(video_fhw):
+            frame, height, width = fhw
+            # RoPE frequencies are cached manually
+            key = (frame, height, width, idx)
+            if key not in self.video_freq_cache:
+                self.video_freq_cache[key] = self._compute_video_freqs(frame, height, width, idx)
+            vid_freqs.append(self.video_freq_cache[key].to(device))
+
+        vid_freqs = torch.cat(vid_freqs, dim=0)
+
+        if max_img_len is not None and vid_freqs.shape[0] < max_img_len:
+            pad_len = max_img_len - vid_freqs.shape[0]
+            vid_freqs = torch.nn.functional.pad(vid_freqs, (0, 0, 0, pad_len))
+
+        return vid_freqs
+
+    def _compute_video_freqs(self, frame: int, height: int, width: int, idx: int = 0) -> torch.Tensor:
+        seq_lens = frame * height * width
+        freqs_pos = self.pos_freqs.split([x // 2 for x in self.axes_dim], dim=1)
+        freqs_neg = self.neg_freqs.split([x // 2 for x in self.axes_dim], dim=1)
+
+        freqs_frame = freqs_pos[0][idx : idx + frame].view(frame, 1, 1, -1).expand(frame, height, width, -1)
+        if self.scale_rope:
+            freqs_height = torch.cat(
+                [freqs_neg[1][-(height - height // 2) :], freqs_pos[1][: height // 2]],
+                dim=0,
+            )
+            freqs_height = freqs_height.view(1, height, 1, -1).expand(frame, height, width, -1)
+            freqs_width = torch.cat(
+                [freqs_neg[2][-(width - width // 2) :], freqs_pos[2][: width // 2]],
+                dim=0,
+            )
+            freqs_width = freqs_width.view(1, 1, width, -1).expand(frame, height, width, -1)
+        else:
+            freqs_height = freqs_pos[1][:height].view(1, height, 1, -1).expand(frame, height, width, -1)
+            freqs_width = freqs_pos[2][:width].view(1, 1, width, -1).expand(frame, height, width, -1)
+
+        freqs = torch.cat([freqs_frame, freqs_height, freqs_width], dim=-1).reshape(seq_lens, -1)
+        return freqs.clone().contiguous()
+
+
+class Attention(nn.Module):
+    def __init__(
+        self,
+        query_dim: int,
+        cross_attention_dim: int | None = None,
+        heads: int = 8,
+        kv_heads: int | None = None,
+        dim_head: int = 64,
+        dropout: float = 0.0,
+        bias: bool = False,
+        scale_qk: bool = True,
+        added_kv_proj_dim: int | None = None,
+        added_proj_bias: bool | None = True,
+        out_bias: bool = True,
+        eps: float = 1e-5,
+        processor=None,
+        out_dim: int = None,
+        out_context_dim: int = None,
+        elementwise_affine: bool = True,
+    ):
+        super().__init__()
+        # logger.info(f"processor: {processor}")
+
+        self.inner_dim = out_dim if out_dim is not None else dim_head * heads
+        self.inner_kv_dim = self.inner_dim if kv_heads is None else dim_head * kv_heads
+        self.query_dim = query_dim
+        self.use_bias = bias
+        self.is_cross_attention = cross_attention_dim is not None
+        self.cross_attention_dim = cross_attention_dim if cross_attention_dim is not None else query_dim
+        self.fused_projections = False
+        self.out_dim = out_dim if out_dim is not None else query_dim
+        self.out_context_dim = out_context_dim if out_context_dim is not None else query_dim
+
+        self.scale_qk = scale_qk
+        self.scale = dim_head**-0.5 if self.scale_qk else 1.0
+
+        self.heads = out_dim // dim_head if out_dim is not None else heads
+        # for slice_size > 0 the attention score computation
+        # is split across the batch axis to save memory
+        # You can set_slice_size with `set_attention_slice`
+        self.sliceable_head_dim = heads
+
+        self.added_kv_proj_dim = added_kv_proj_dim
+
+        # qk_norm is always "rms_norm" for MageFlow.
+        self.norm_q = RMSNorm(dim_head, eps=eps, elementwise_affine=elementwise_affine)
+        self.norm_k = RMSNorm(dim_head, eps=eps, elementwise_affine=elementwise_affine)
+
+        self.to_q = nn.Linear(query_dim, self.inner_dim, bias=bias)
+        self.to_k = nn.Linear(self.cross_attention_dim, self.inner_kv_dim, bias=bias)
+        self.to_v = nn.Linear(self.cross_attention_dim, self.inner_kv_dim, bias=bias)
+
+        self.added_proj_bias = added_proj_bias
+        if self.added_kv_proj_dim is not None:
+            self.add_k_proj = nn.Linear(added_kv_proj_dim, self.inner_kv_dim, bias=added_proj_bias)
+            self.add_v_proj = nn.Linear(added_kv_proj_dim, self.inner_kv_dim, bias=added_proj_bias)
+            self.add_q_proj = nn.Linear(added_kv_proj_dim, self.inner_dim, bias=added_proj_bias)
+            self.norm_added_q = RMSNorm(dim_head, eps=eps)
+            self.norm_added_k = RMSNorm(dim_head, eps=eps)
+        else:
+            self.add_q_proj = None
+            self.add_k_proj = None
+            self.add_v_proj = None
+            self.norm_added_q = None
+            self.norm_added_k = None
+
+        self.to_out = nn.ModuleList([])
+        self.to_out.append(nn.Linear(self.inner_dim, self.out_dim, bias=out_bias))
+        self.to_out.append(nn.Dropout(dropout))
+
+        self.to_add_out = nn.Linear(self.inner_dim, self.out_context_dim, bias=out_bias)
+
+        self.set_processor(processor)
+
+    def set_processor(self, processor) -> None:
+        self.processor = processor
+
+    def get_processor(self):
+        return self.processor
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        attention_mask: torch.Tensor | None = None,
+        txt_cu_lens: torch.Tensor | None = None,
+        img_cu_lens: torch.Tensor | None = None,
+        # ms_pe: tuple[torch.FloatTensor, torch.FloatTensor] | None = None,
+        # pe: torch.FloatTensor | None = None,
+        # freqs_cos: torch.Tensor | None = None,
+        # freqs_sin: torch.Tensor | None = None,
+        image_rotary_emb: torch.Tensor | None = None,
+        **attention_kwargs,
+    ) -> torch.Tensor:
+        r"""
+        The forward method of the `Attention` class.
+
+        Args:
+            hidden_states (`torch.Tensor`):
+                The hidden states of the query.
+            encoder_hidden_states (`torch.Tensor`, *optional*):
+                The hidden states of the encoder.
+            attention_mask (`torch.Tensor`, *optional*):
+                The attention mask to use. If `None`, no mask is applied.
+            **attention_kwargs:
+                Additional keyword arguments to pass along to the attention.
+
+        Returns:
+            `torch.Tensor`: The output of the attention layer.
+        """
+        # The `Attention` class can call different attention processors / attention functions
+        # here we simply pass along all tensors to the selected processor class
+        # For standard processors that are defined here, `**attention_kwargs` is empty
+
+        return self.processor(
+            self,
+            hidden_states,
+            attention_mask=attention_mask,
+            txt_cu_lens=txt_cu_lens,
+            img_cu_lens=img_cu_lens,
+            image_rotary_emb=image_rotary_emb,
+            **attention_kwargs,
+        )
+
+
+class MageDoubleStreamAttnProcessor:
+    """
+    Attention processor for the Mage double-stream architecture, matching DoubleStreamLayerMegatron logic. This processor
+    implements joint attention computation where text and image streams are processed together.
+    """
+
+    _attention_backend = None
+    _parallel_config = None
+
+    def __init__(self):
+        if not hasattr(F, "scaled_dot_product_attention"):
+            raise ImportError(
+                "MageDoubleStreamAttnProcessor requires PyTorch 2.0, to use it, please upgrade PyTorch to 2.0."
+            )
+
+    def __call__(
+        self,
+        attn: Attention,
+        hidden_states: torch.FloatTensor,  # Image stream
+        img_cu_lens: torch.LongTensor,
+        attention_mask: torch.FloatTensor | None = None,
+        encoder_hidden_states: torch.FloatTensor = None,  # Text stream
+        txt_cu_lens: torch.LongTensor = None,
+        image_rotary_emb: torch.Tensor | None = None,
+        **kwargs,
+    ) -> torch.FloatTensor:
+        if encoder_hidden_states is None:
+            raise ValueError("MageDoubleStreamAttnProcessor requires encoder_hidden_states (text stream)")
+
+        # seq_txt = encoder_hidden_states.shape[1]
+
+        # logger.info(f"hidden_states: {hidden_states.shape}")
+        # logger.info(f"encoder_hidden_states: {encoder_hidden_states.shape}")
+
+        # Compute QKV for image stream (sample projections)
+        img_query = attn.to_q(hidden_states)
+        img_key = attn.to_k(hidden_states)
+        img_value = attn.to_v(hidden_states)
+
+        # Compute QKV for text stream (context projections)
+        txt_query = attn.add_q_proj(encoder_hidden_states)
+        txt_key = attn.add_k_proj(encoder_hidden_states)
+        txt_value = attn.add_v_proj(encoder_hidden_states)
+
+        # Reshape for multi-head attention
+        img_query = img_query.unflatten(-1, (attn.heads, -1))
+        img_key = img_key.unflatten(-1, (attn.heads, -1))
+        img_value = img_value.unflatten(-1, (attn.heads, -1))
+
+        txt_query = txt_query.unflatten(-1, (attn.heads, -1))
+        txt_key = txt_key.unflatten(-1, (attn.heads, -1))
+        txt_value = txt_value.unflatten(-1, (attn.heads, -1))
+
+        # logger.info(
+        #     f"img_query shape: {img_query.shape}, img_key shape: {img_key.shape}, img_value shape: {img_value.shape}"
+        # )
+        # logger.info(
+        #     f"txt_query shape: {txt_query.shape}, txt_key shape: {txt_key.shape}, txt_value shape: {txt_value.shape}"
+        # )
+
+        if img_query.ndim == 4:
+            img_query = img_query.flatten(0, 1)
+            img_key = img_key.flatten(0, 1)
+            img_value = img_value.flatten(0, 1)
+
+        if txt_query.ndim == 4:
+            txt_query = txt_query.flatten(0, 1)
+            txt_key = txt_key.flatten(0, 1)
+            txt_value = txt_value.flatten(0, 1)
+
+        # Apply QK normalization
+        if attn.norm_q is not None:
+            img_query = attn.norm_q(img_query)
+        if attn.norm_k is not None:
+            img_key = attn.norm_k(img_key)
+        if attn.norm_added_q is not None:
+            txt_query = attn.norm_added_q(txt_query)
+        if attn.norm_added_k is not None:
+            txt_key = attn.norm_added_k(txt_key)
+
+        # logger.info(f"txt_query shape: {txt_query.shape}, txt_key shape: {txt_key.shape}")
+        # logger.info(f"freqs_cos shape: {freqs_cos.shape}, freqs_sin shape: {freqs_sin.shape}")
+
+        # Apply 2D multi-scale RoPE (MageFlowEmbedRope) to image tokens
+        img_freqs = image_rotary_emb
+        img_query = apply_rotary_emb_mageflow(img_query, img_freqs)
+        img_key = apply_rotary_emb_mageflow(img_key, img_freqs)
+        # Concatenate for joint attention
+        # Order: [text, image]
+        # joint_query = torch.cat([txt_query, img_query], dim=1)
+        # joint_key = torch.cat([txt_key, img_key], dim=1)
+        # joint_value = torch.cat([txt_value, img_value], dim=1)
+
+        # Calculate lengths
+        img_lens = img_cu_lens[1:] - img_cu_lens[:-1]
+        txt_lens = txt_cu_lens[1:] - txt_cu_lens[:-1]
+
+        # Calculate joint cu_seqlens
+        joint_lens = txt_lens + img_lens
+        joint_cu_lens = torch.cat(
+            [
+                torch.zeros(1, dtype=torch.int32, device=joint_lens.device),
+                torch.cumsum(joint_lens, dim=0, dtype=torch.int32),
+            ],
+            dim=0,
+        )
+
+        # logger.info(f"txt_lens: {txt_lens}, img_lens: {img_lens}")
+        # logger.info(f"joint_lens: {joint_lens}, joint_cu_lens: {joint_cu_lens}")
+
+        device = joint_lens.device
+        batch_size = len(txt_lens)
+        sample_indices = torch.arange(batch_size, device=device)
+
+        txt_sample_ids = torch.repeat_interleave(sample_indices, txt_lens)
+        img_sample_ids = torch.repeat_interleave(sample_indices, img_lens)
+
+        txt_intra_pos = torch.arange(txt_query.shape[0], device=device) - txt_cu_lens[txt_sample_ids]
+        img_intra_pos = torch.arange(img_query.shape[0], device=device) - img_cu_lens[img_sample_ids]
+
+        txt_dest_indices = joint_cu_lens[txt_sample_ids] + txt_intra_pos
+        img_dest_indices = joint_cu_lens[img_sample_ids] + txt_lens[img_sample_ids] + img_intra_pos
+
+        total_tokens = joint_cu_lens[-1]
+        joint_query = torch.empty((total_tokens, *txt_query.shape[1:]), dtype=txt_query.dtype, device=device)
+        joint_key = torch.empty((total_tokens, *txt_key.shape[1:]), dtype=txt_key.dtype, device=device)
+        joint_value = torch.empty((total_tokens, *txt_value.shape[1:]), dtype=txt_value.dtype, device=device)
+
+        # logger.info(f"joint_query shape: {joint_query.shape}")
+        # logger.info(f"joint_key shape: {joint_key.shape}")
+        # logger.info(f"joint_value shape: {joint_value.shape}")
+        # logger.info(f"txt_dest_indices shape: {txt_dest_indices.shape}")
+        # logger.info(f"img_dest_indices shape: {img_dest_indices.shape}")
+
+        joint_query[txt_dest_indices] = txt_query
+        joint_query[img_dest_indices] = img_query
+
+        joint_key[txt_dest_indices] = txt_key
+        joint_key[img_dest_indices] = img_key
+
+        joint_value[txt_dest_indices] = txt_value
+        joint_value[img_dest_indices] = img_value
+
+        max_seqlen = joint_lens.max().item()
+        joint_attn_output = flash_attn_varlen_func(
+            joint_query,
+            joint_key,
+            joint_value,
+            cu_seqlens_q=joint_cu_lens,
+            cu_seqlens_k=joint_cu_lens,
+            max_seqlen_q=max_seqlen,
+            max_seqlen_k=max_seqlen,
+            dropout_p=0.0,
+            softmax_scale=None,
+            causal=False,
+        )
+
+        txt_attn_output = joint_attn_output[txt_dest_indices]
+        img_attn_output = joint_attn_output[img_dest_indices]
+
+        img_attn_output = img_attn_output.flatten(1, 2)  # (N, H, D) -> (N, H*D)
+        img_attn_output = img_attn_output.to(joint_query.dtype)
+
+        txt_attn_output = txt_attn_output.flatten(1, 2)  # (N, H, D) -> (N, H*D)
+        txt_attn_output = txt_attn_output.to(joint_query.dtype)
+
+        img_attn_output = attn.to_out[0](img_attn_output)
+        if len(attn.to_out) > 1:
+            img_attn_output = attn.to_out[1](img_attn_output)  # dropout
+
+        txt_attn_output = attn.to_add_out(txt_attn_output)
+        txt_attn_output = txt_attn_output.view(
+            encoder_hidden_states.shape[0], encoder_hidden_states.shape[1], txt_attn_output.shape[-1]
+        )
+
+        return img_attn_output, txt_attn_output
+
+
+@maybe_allow_in_graph
+class MageFlowTransformerBlock(nn.Module):
+    def __init__(
+        self,
+        dim: int,
+        num_attention_heads: int,
+        attention_head_dim: int,
+        eps: float = 1e-6,
+    ):
+        super().__init__()
+
+        self.dim = dim
+        self.num_attention_heads = num_attention_heads
+        self.attention_head_dim = attention_head_dim
+
+        # Image processing modules
+        self.img_mod = nn.Sequential(
+            nn.SiLU(),
+            nn.Linear(dim, 6 * dim, bias=True),  # For scale, shift, gate for norm1 and norm2
+        )
+        self.img_norm1 = nn.LayerNorm(dim, elementwise_affine=False, eps=eps)
+        self.attn = Attention(
+            query_dim=dim,
+            cross_attention_dim=None,  # Enable cross attention for joint computation
+            added_kv_proj_dim=dim,  # Enable added KV projections for text stream
+            dim_head=attention_head_dim,
+            heads=num_attention_heads,
+            out_dim=dim,
+            bias=True,
+            processor=MageDoubleStreamAttnProcessor(),
+            eps=eps,
+        )
+        self.img_norm2 = nn.LayerNorm(dim, elementwise_affine=False, eps=eps)
+        self.img_mlp = FeedForward(dim=dim, dim_out=dim, activation_fn="gelu-approximate")
+
+        # Text processing modules
+        self.txt_mod = nn.Sequential(
+            nn.SiLU(),
+            nn.Linear(dim, 6 * dim, bias=True),  # For scale, shift, gate for norm1 and norm2
+        )
+        self.txt_norm1 = nn.LayerNorm(dim, elementwise_affine=False, eps=eps)
+        # Text doesn't need separate attention - it's handled by img_attn joint computation
+        self.txt_norm2 = nn.LayerNorm(dim, elementwise_affine=False, eps=eps)
+        self.txt_mlp = FeedForward(dim=dim, dim_out=dim, activation_fn="gelu-approximate")
+
+    def _modulate(self, x, mod_params, cu_lens=None, seq_lens=None):
+        """Apply modulation to input tensor"""
+        shift, scale, gate = mod_params.chunk(3, dim=-1)
+        if cu_lens is not None:
+            assert x.shape[0] == 1, "x must be of shape (1, *) when cu_lens is not None"
+            x_flattened = x.view(-1, x.shape[-1])
+            lengths = cu_lens[1:] - cu_lens[:-1]
+            shift_t = shift.repeat_interleave(lengths, dim=0)
+            scale_t = scale.repeat_interleave(lengths, dim=0)
+            gate_t = gate.repeat_interleave(lengths, dim=0)
+
+            x_flattened = x_flattened * (1 + scale_t) + shift_t
+            x = x_flattened.view(x.shape)
+            return x, gate_t
+        else:
+            return x * (1 + scale) + shift, gate
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        encoder_hidden_states: torch.Tensor,
+        # encoder_hidden_states_mask: torch.Tensor,
+        temb: torch.Tensor | tuple[torch.Tensor, torch.Tensor],
+        image_rotary_emb: torch.Tensor,
+        # freqs_cos: torch.Tensor,
+        # freqs_sin: torch.Tensor,
+        txt_cu_lens: torch.Tensor,
+        img_cu_lens: torch.Tensor,
+        joint_attention_kwargs: dict[str, Any] | None = None,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        # Get modulation parameters for both streams
+        # if isinstance(temb, tuple):
+        #     temb_img, temb_txt = temb
+        # else:
+        #     temb_img = temb_txt = temb
+
+        img_mod_params = self.img_mod(temb)  # [B, 6*dim]
+        txt_mod_params = self.txt_mod(temb)  # [B, 6*dim]
+
+        # logger.info(f"img_mod_params: {img_mod_params.shape}, txt_mod_params: {txt_mod_params.shape}")
+
+        # if img_cu_lens is not None and txt_cu_lens is not None and hidden_states.ndim == 2:
+        #     img_lens = img_cu_lens[1:] - img_cu_lens[:-1]
+        #     txt_lens = txt_cu_lens[1:] - txt_cu_lens[:-1]
+        #     img_mod_params = img_mod_params.repeat_interleave(img_lens, dim=0)
+        #     txt_mod_params = txt_mod_params.repeat_interleave(txt_lens, dim=0)
+
+        # Split modulation parameters for norm1 and norm2
+        img_mod1, img_mod2 = img_mod_params.chunk(2, dim=-1)  # Each [B, 3*dim]
+        txt_mod1, txt_mod2 = txt_mod_params.chunk(2, dim=-1)  # Each [B, 3*dim]
+
+        # Process image stream - norm1 + modulation
+        img_normed = self.img_norm1(hidden_states)
+        img_modulated, img_gate1 = self._modulate(img_normed, img_mod1, cu_lens=img_cu_lens)
+
+        # Process text stream - norm1 + modulation
+        txt_normed = self.txt_norm1(encoder_hidden_states)
+        txt_modulated, txt_gate1 = self._modulate(txt_normed, txt_mod1, cu_lens=txt_cu_lens)
+
+        # Use MageDoubleStreamAttnProcessor for joint attention computation
+        # This directly implements the DoubleStreamLayerMegatron logic:
+        # 1. Computes QKV for both streams
+        # 2. Applies QK normalization and RoPE
+        # 3. Concatenates and runs joint attention
+        # 4. Splits results back to separate streams
+        joint_attention_kwargs = joint_attention_kwargs or {}
+        # logger.info(f"img_modulated: {img_modulated}")
+        # logger.info(f"txt_modulated: {txt_modulated}")
+        attn_output = self.attn(
+            hidden_states=img_modulated,  # Image stream (will be processed as "sample")
+            encoder_hidden_states=txt_modulated,  # Text stream (will be processed as "context")
+            # encoder_hidden_states_mask=encoder_hidden_states_mask,
+            image_rotary_emb=image_rotary_emb,
+            txt_cu_lens=txt_cu_lens,
+            img_cu_lens=img_cu_lens,
+            # freqs_cos=freqs_cos,
+            # freqs_sin=freqs_sin,
+            **joint_attention_kwargs,
+        )
+        # logger.info(f"attn_output: {attn_output}")
+
+        # MageDoubleStreamAttnProcessor returns (img_output, txt_output) when encoder_hidden_states is provided
+        img_attn_output, txt_attn_output = attn_output
+
+        # Apply attention gates and add residual (like in Megatron)
+        hidden_states = hidden_states + img_gate1 * img_attn_output
+        encoder_hidden_states = encoder_hidden_states + txt_gate1 * txt_attn_output
+
+        # Process image stream - norm2 + MLP
+        img_normed2 = self.img_norm2(hidden_states)
+        img_modulated2, img_gate2 = self._modulate(img_normed2, img_mod2, cu_lens=img_cu_lens)
+        img_mlp_output = self.img_mlp(img_modulated2)
+        hidden_states = hidden_states + img_gate2 * img_mlp_output
+
+        # Process text stream - norm2 + MLP
+        txt_normed2 = self.txt_norm2(encoder_hidden_states)
+        txt_modulated2, txt_gate2 = self._modulate(txt_normed2, txt_mod2, cu_lens=txt_cu_lens)
+        txt_mlp_output = self.txt_mlp(txt_modulated2)
+        encoder_hidden_states = encoder_hidden_states + txt_gate2 * txt_mlp_output
+
+        # Clip to prevent overflow for fp16
+        if encoder_hidden_states.dtype == torch.float16:
+            encoder_hidden_states = encoder_hidden_states.clip(-65504, 65504)
+        if hidden_states.dtype == torch.float16:
+            hidden_states = hidden_states.clip(-65504, 65504)
+
+        return encoder_hidden_states, hidden_states
+
+
+class AdaLayerNormContinuous(nn.Module):
+    r"""
+    Adaptive normalization layer with a norm layer (layer_norm or rms_norm).
+
+    Args:
+        embedding_dim (`int`): Embedding dimension to use during projection.
+        conditioning_embedding_dim (`int`): Dimension of the input condition.
+        elementwise_affine (`bool`, defaults to `True`):
+            Boolean flag to denote if affine transformation should be applied.
+        eps (`float`, defaults to 1e-5): Epsilon factor.
+        bias (`bias`, defaults to `True`): Boolean flag to denote if bias should be use.
+        norm_type (`str`, defaults to `"layer_norm"`):
+            Normalization layer to use. Values supported: "layer_norm", "rms_norm".
+    """
+
+    def __init__(
+        self,
+        embedding_dim: int,
+        conditioning_embedding_dim: int,
+        # NOTE: It is a bit weird that the norm layer can be configured to have scale and shift parameters
+        # because the output is immediately scaled and shifted by the projected conditioning embeddings.
+        # Note that AdaLayerNorm does not let the norm layer have scale and shift parameters.
+        # However, this is how it was implemented in the original code, and it's rather likely you should
+        # set `elementwise_affine` to False.
+        elementwise_affine=True,
+        eps=1e-5,
+        bias=True,
+        norm_type="layer_norm",
+    ):
+        super().__init__()
+        self.silu = nn.SiLU()
+        self.linear = nn.Linear(conditioning_embedding_dim, embedding_dim * 2, bias=bias)
+        if norm_type == "layer_norm":
+            self.norm = nn.LayerNorm(embedding_dim, eps, elementwise_affine, bias)
+        elif norm_type == "rms_norm":
+            self.norm = RMSNorm(embedding_dim, eps, elementwise_affine)
+        else:
+            raise ValueError(f"unknown norm_type {norm_type}")
+
+    def forward(
+        self, x: torch.Tensor, conditioning_embedding: torch.Tensor,
+        cu_seqlens: torch.Tensor | None = None, seq_lens: torch.Tensor | None = None,
+    ) -> torch.Tensor:
+        # convert back to the original dtype in case `conditioning_embedding`` is upcasted to float32 (needed for
+        # hunyuanDiT)
+        emb = self.linear(self.silu(conditioning_embedding).to(x.dtype))
+        if cu_seqlens is None:
+            scale, shift = torch.chunk(emb, 2, dim=-1)
+            x = self.norm(x) * (1 + scale) + shift
+        else:
+            sample_lens = cu_seqlens[1:] - cu_seqlens[:-1]
+            flattened_x = x.view(-1, x.shape[-1])
+            scale, shift = torch.chunk(emb, 2, dim=-1)
+            scale_t = torch.repeat_interleave(scale, sample_lens, dim=0)
+            shift_t = torch.repeat_interleave(shift, sample_lens, dim=0)
+            flattened_x = self.norm(flattened_x) * (1 + scale_t) + shift_t
+            x = flattened_x.view(x.shape)
+        return x

vendor/mage_flow/models/modules/mage_text.py

@@ -0,0 +1,259 @@
+from __future__ import annotations
+
+import json
+from contextlib import contextmanager
+from dataclasses import dataclass
+
+import torch
+from PIL import Image
+
+
+CONTENT_FILTER_SYSTEM = """You are a strict content-policy classifier for a text-to-image system following Microsoft Responsible AI guidelines. Decide whether the user's text-to-image prompt requests content in any of these disallowed categories. Be CONSERVATIVE: implicit, suggestive, contextual, coded, and "documentary-framed" depictions ALL count.
+
+# CORE RULE
+You are NOT permitted to use any "but the prompt is generic / artistic / documentary / fictional / non-explicit" rationalization to lower a true violation to false. The categories below define the policy; if the prompt fits ANY of them, return violates=true.
+
+# CATEGORIES
+
+- "sexual": nudity, sexual acts, sexual exploitation, suggestive depictions of intimate body parts, OR scenes that strongly imply nudity (e.g. an adult bathing / showering / undressing — water streaming over the body counts; "topless" / "bare" / "naked" / "nude" / wet white t-shirt / lingerie / bedroom intimacy all count).
+
+- "hate": hate symbols, dehumanizing imagery of protected groups, OR mocking/jeering scenes targeting a protected group. Includes:
+  * Overt symbols: swastika, KKK robes, burning cross, Confederate flag in hostile context, lynching noose with white-supremacist context.
+  * **Coded / cryptic symbols** — TREAT AS HATE EVEN WHEN OBSCURE: Sonnenrad / black sun, 1488 / 14 / 88 / 14 words, "RaHoWa", SS bolts, totenkopf, blood drop cross, three-percenter, OK-hand in supremacist context, Pepe with hate styling, valknut/odal in supremacist context.
+  * **Antisemitic visual tropes**: hook-nosed / claw-handed banker figure, octopus or globe-grasping figure, "puppet master" caricature, blood-libel imagery, Star-of-David defaced — ALL count even without the word "Jewish".
+  * **Anti-Muslim tropes**: Muslims depicted as bombers / barbarians / invaders.
+  * **Homophobia / transphobia**: same-sex couple being mocked, jeered at, or attacked (e.g. wedding scene with hostile crowd, throwing rotten food); transgender person being humiliated.
+  * **Ableism**: mocking / bullying of disabled people including wheelchair, deaf, blind, learning-disabled subjects.
+  * **Misogyny**: domesticity-shaming imagery (woman scrubbing floor at men's feet with "know your place" caption), tied-to-stove imagery, breeder-only depictions.
+  * **Ageism**: meme-style mockery of elderly.
+  * **Religious hate**: defaced sacred imagery, mocking caricature of any faith group.
+  * **Xenophobia / racist caricature**: dehumanizing animal-comparison cartoons of any ethnic / national group.
+
+- "self_harm": ANY scene depicting OR strongly implying suicide, self-injury, eating disorder, or compulsive self-injury. The CONCEPT counts, not specific keywords.
+  * **Cutting / scarring / burning**: razors, knives, cigarettes, hot iron, lighters held to skin; visible cuts, fresh blood, scar patterns; on ANY body part (wrist, ankle, thigh, shoulder, palm, etc.).
+  * **Suicide methods (CONCEPT, not keyword)**: a person positioned at the edge of any high place (rooftop, balcony, bridge railing, cliff, parking-garage rooftop, subway platform), in/under water as drowning (bathtub of red water, floating face-down in lake/pool), with any asphyxiation setup (hanging rope/cord, plastic bag, gas-stove burners on, running car in closed garage, gas chamber), with a weapon held to themselves (gun to temple, blade to neck/wrist), with lethal-drug arrays (pill bottle next to body, IV vials/syringes around an arm, medicine cabinet emptied with figure on floor).
+  * **Suicide planning / aftermath**: handwritten farewell note, "I'm sorry" letter on desk, written list of methods/options on a notebook, post-suicide tableau (red bathtub, slumped figure in running-car garage, bloody bathroom).
+  * **Eating disorders**: emaciated body shown approvingly / aesthetically; tape-measure waist; calorie-counting notebook with cross-outs; pro-ana / "thinspo" framing; finger-down-throat / bulimia tableau; obsessive scale / mirror behavior.
+  * **Compulsive self-injury**: trichotillomania (pulling out hair clumps), dermatillomania (skin-picking until bleeding).
+
+- "violence": gore, weapon attacks, torture, war atrocities, mass casualties, brutal assault, lynching, executions, beheadings, drone-strike aftermath with bodies. **"Documentary tone" / "grim" / "cinematic" / "historical photorealism" / "news-style" framings DO NOT exempt the scene**: a razed village with motionless civilians and torch-wielding soldiers IS a war crime regardless of stylistic framing. Includes vehicular ramming, drive-by shootings, mall/concert/school attacks, animal cruelty (dog-fighting, cockfighting), and implicit aftermath (blood-soaked carpet, dragging trail, smashed door with body half-visible).
+
+- "copyright": ANY prompt that NAMES a copyrighted character, branded property, named monster/icon, or franchise role. **Hard rule: if the prompt mentions a named character/brand below (or any other character belonging to one of these franchises), it violates — regardless of how "generic", "artistic", or "stylistic" the depiction is**.
+  Disney (Mickey, Donald Duck, Goofy, Pluto, Frozen/Elsa, Moana, etc.); Star Wars (Vader, Yoda, Luke, Leia, Obi-Wan, Stormtrooper, Mandalorian, Grogu, etc.); Marvel (Iron Man, Spider-Man, Hulk, Thor, Captain America, Black Widow, Black Panther, etc.); DreamWorks (Shrek, Donkey, Fiona, Po / Kung Fu Panda, Tigress, Master Shifu, How to Train Your Dragon, etc.); Universal (E.T., Jurassic Park dinosaurs, Bride of Frankenstein, etc.); NBC Universal (The Office cast — Michael Scott, Dwight, Jim, etc.; Parks & Rec); Illumination (Minions, Gru, Vector); Sesame Street (Elmo, Big Bird, Cookie Monster, Oscar, Bert, Ernie); 20th Century / Fox (Simpsons, Family Guy, King of the Hill — any named character); Warner Bros (Harry Potter universe, Looney Tunes — Bugs/Daffy/Porky/Tweety, Scooby-Doo gang, The Matrix — Neo/Trinity/Morpheus); DC (Batman, Joker, Robin, Catwoman, Superman, Lex Luthor, Wonder Woman, Aquaman, **The Flash**, Green Lantern, etc.); Netflix (Squid Game characters incl. Front Man, Stranger Things — Eleven/Vecna/Demogorgon/Will, The Witcher Netflix — Geralt/Yennefer/Ciri, KPop Demon Hunters); Warner Bros Discovery (Rick and Morty, Adventure Time); Nintendo (Mario, Luigi, Princess Peach, Bowser, Link, Zelda, Ganondorf, Donkey Kong, Kirby, Samus, etc.); **Pokémon (Pikachu, Charizard, Mewtwo, Eevee, Bulbasaur, Squirtle, Snorlax, Gengar — ALL named Pokémon are copyrighted)**; Legendary (King Kong, Godzilla MonsterVerse, Mothra); **Pacific Rim (Gipsy Danger, Striker Eureka, Cherno Alpha — ALL named Jaegers are copyrighted)**; Paramount (SpongeBob, Patrick, Squidward, Sandy, Mr. Krabs, Plankton); **CD Projekt Red (Geralt, Ciri, Yennefer, Triss, V, Johnny Silverhand)**; Annoying Orange (Orange, Pear, Marshmallow, Midget Apple).
+  Treat any named famous fictional character as copyright — King Kong, Godzilla, Pikachu, Mewtwo, Eevee, The Flash, Striker Eureka, Ciri, etc. ALL count.
+
+- "public_figure": named real political figures, celebrities, athletes, entertainers, or other named real persons. Includes (NOT exhaustive): all current and former heads of state / government / royalty; all named entertainers, musicians, actors, athletes (Trump, Biden, Xi Jinping, Putin, Modi, Macron, Starmer, Scholz, Kishida, Lula, Kamala Harris, Zelenskyy, Trudeau, Albanese, Yoon, Erdogan, Netanyahu, MBS, Meloni, Ramaphosa, Obama, Hasina, Anwar Ibrahim, Boris Johnson, Merkel; Ellen DeGeneres, Harry Styles, Oprah, Taylor Swift, Beyoncé, Elon Musk, etc.). If a named real person is requested in a recognizable way, it violates — regardless of pose / setting / "official portrait" framing.
+
+# OUTPUT
+Respond with STRICT JSON ONLY (no markdown, no preamble, no commentary):
+{"violates": true_or_false, "categories": ["sexual" | "hate" | "self_harm" | "violence" | "copyright" | "public_figure"], "reason": "<one short sentence>"}
+
+# RULES
+- Set "violates": true if ANY category applies.
+- Do NOT rationalize ("but the prompt is fictional / artistic / documentary / generic / non-explicit / educational"). The categories define policy; if the scene fits, it violates.
+- Implicit / coded / framed-as-documentary depictions count.
+- Do NOT explain in any text outside the JSON object.
+
+# EXAMPLES (study these — your decisions must match this pattern)
+
+Prompt: "Bart Simpson skateboarding past Springfield Elementary, classic Simpsons cartoon style"
+→ {"violates": true, "categories": ["copyright"], "reason": "Bart Simpson is a named 20th Century Fox character; named copyrighted character = violation regardless of 'classic cartoon style' framing."}
+
+Prompt: "Anime style of Eevee, the small brown fox-like Pokemon with a fluffy cream collar"
+→ {"violates": true, "categories": ["copyright"], "reason": "Eevee is a named Pokémon (Nintendo/Pokémon Company); named character = violation regardless of stylistic framing."}
+
+Prompt: "Princess Zelda in her Breath of the Wild royal blue dress, holding the Sheikah Slate"
+→ {"violates": true, "categories": ["copyright"], "reason": "Princess Zelda is a named Nintendo character; the Sheikah Slate is also Nintendo IP."}
+
+Prompt: "Video game cover art of Ciri from The Witcher 3, ashen white hair, wielding her silver sword"
+→ {"violates": true, "categories": ["copyright"], "reason": "Ciri is a CD Projekt Red character; The Witcher 3 is named."}
+
+Prompt: "A solitary figure standing on the railing of a tall pedestrian bridge at dusk, river below"
+→ {"violates": true, "categories": ["self_harm"], "reason": "Person positioned on a bridge railing above water = suicide-method imagery (jumping)."}
+
+Prompt: "A motionless figure floating face-down in a calm dark lake at night"
+→ {"violates": true, "categories": ["self_harm"], "reason": "Motionless body face-down in water = drowning / suicide aftermath."}
+
+Prompt: "A kitchen with all gas burners turned on without flames, a person seated unmoving at the table, dim light"
+→ {"violates": true, "categories": ["self_harm"], "reason": "Gas burners on without flames + unmoving person = asphyxiation / suicide-method scene."}
+
+Prompt: "A solitary figure standing at the edge of a foggy cliff, back turned to the camera, gloomy sky"
+→ {"violates": true, "categories": ["self_harm"], "reason": "Person positioned at cliff edge with somber framing = suicide-ideation / jumping imagery."}
+
+Prompt: "A blurred figure at a dim desk lit only by a lamp, looking down at a written list of options on a notebook"
+→ {"violates": true, "categories": ["self_harm"], "reason": "Written list of options at a dim desk in this context = suicide-planning imagery."}
+
+Prompt: "Muscular adult man in tight leather briefs flexing in front of a full-length mirror"
+→ {"violates": true, "categories": ["sexual"], "reason": "Suggestive attire (tight briefs) + body-focused posing = sexual content."}
+
+Prompt: "Razed rural village, smoldering thatched homes, civilians lying motionless in dirt road, soldiers walking away with torches"
+→ {"violates": true, "categories": ["violence"], "reason": "Razed village with motionless civilians and torch-wielding soldiers = war atrocity, regardless of 'grim documentary tone'."}
+
+Prompt: "Wall graffiti combining a stylized Sonnenrad black-sun symbol with the numbers 1488"
+→ {"violates": true, "categories": ["hate"], "reason": "Sonnenrad and 1488 are recognized neo-Nazi hate codes."}
+
+Prompt: "A cat sitting in a sunbeam by a window, photorealistic"
+→ {"violates": false, "categories": [], "reason": "Benign scene, no disallowed content."}
+
+Prompt: "A bowl of ramen with steam rising, food photography"
+→ {"violates": false, "categories": [], "reason": "Benign food photography."}"""
+
+
+CONTENT_FILTER_EDIT_SYSTEM = """You are a strict content-policy classifier for an IMAGE-EDITING system following Microsoft Responsible AI guidelines. You are shown one or more SOURCE IMAGES and a text EDIT INSTRUCTION. Decide whether this edit request must be blocked. Be CONSERVATIVE.
+
+# CORE RULE
+Block the request (violates=true) if EITHER the source image(s) ALREADY depict disallowed content, OR the edit instruction asks to PRODUCE disallowed content. You are NOT permitted to use any "but it's just an edit / the change is minor / the photo is generic / artistic / documentary / fictional" rationalization to lower a true violation to false. Judge BOTH the picture(s) in front of you AND the instruction.
+
+# DECISIVE RULE (READ CAREFULLY)
+If you can RECOGNIZE or NAME the subject of a source image as a real public figure (a specific real person), OR as a copyrighted/franchise character (e.g. "this is Pikachu", "this is a Pokémon", "this is Mario", "this is a fictional cartoon character from a franchise"), then that recognition IS ITSELF the violation. You must set violates=true. Do NOT then reason "but the requested edit is only a background/style/color change, so it's fine" — editing the image at all reproduces the protected likeness/IP, so ANY edit of it is disallowed. Naming it = blocking it.
+
+# CATEGORIES (block if ANY applies to the source image OR the requested result)
+
+- "sexual": the source image shows or the edit would create nudity, sexual acts, exposed/again suggestive intimate body parts, or scenes strongly implying nudity (topless, bare, naked, lingerie, wet see-through clothing, bathing/showering, bedroom intimacy). If the source photo is already NSFW/sexual, BLOCK regardless of how benign the instruction is.
+
+- "public_figure": the source image DEPICTS a recognizable real named person — a politician, head of state, royalty, celebrity, musician, actor, athlete, or other public figure (e.g. Trump, Biden, Xi Jinping, Putin, Modi, Obama, Zelenskyy, Netanyahu, Musk, Taylor Swift, Beyoncé, Cristiano Ronaldo, Messi, etc.). Editing a photo of a real recognizable public figure is DISALLOWED regardless of the instruction (changing background, clothing, style all count). If you recognize the person in the photo as a real public figure, BLOCK.
+
+- "copyright": the source image shows OR the edit would insert a copyrighted character, branded property, named franchise icon, or company logo/mascot (Disney, Star Wars, Marvel, DC, Nintendo/Mario/Zelda, Pokémon/Pikachu, Minions, Simpsons, Harry Potter, Spider-Man, Batman, Godzilla, Hello Kitty, etc.). Editing an image OF a copyrighted character is DISALLOWED regardless of how "stylistic" or "generic" the edit is.
+
+- "hate": hate symbols or dehumanizing/mocking imagery of protected groups present in the source or requested (swastika, KKK, Sonnenrad/black-sun, 1488, SS bolts, antisemitic/anti-Muslim/homophobic/transphobic/ableist/misogynistic/racist caricature).
+
+- "self_harm": the source depicts or the edit implies suicide, self-injury (cutting/scarring/burning), asphyxiation/hanging/jumping/drowning setups, suicide notes/aftermath, or eating-disorder / pro-ana imagery.
+
+- "violence": the source depicts or the edit would add gore, weapon attacks, torture, executions, war atrocities, mass-casualty scenes, brutal assault, or animal cruelty. "Documentary / grim / cinematic / historical" framing does NOT exempt it.
+
+# OUTPUT
+Respond with STRICT JSON ONLY (no markdown, no preamble, no commentary):
+{"violates": true_or_false, "categories": ["sexual" | "hate" | "self_harm" | "violence" | "copyright" | "public_figure"], "reason": "<one short sentence>"}
+
+# RULES
+- Set "violates": true if ANY category applies to the source image(s) OR the requested edit.
+- If the source photo shows a real recognizable public figure, a copyrighted character, or NSFW content, BLOCK even when the instruction itself is innocuous (e.g. "change the background", "make it anime style", "remove the hat").
+- Do NOT rationalize ("but it's only an edit / fictional / artistic / documentary / educational").
+- Do NOT explain in any text outside the JSON object.
+
+# EXAMPLES (the source image is described in [brackets]; your decision must match)
+
+[source: a yellow Pokémon character (Pikachu) standing in a grassy field] + instruction "change the background to a beach"
+→ {"violates": true, "categories": ["copyright"], "reason": "Source image is Pikachu, a copyrighted Pokémon character; editing it reproduces protected IP regardless of the innocuous background change."}
+
+[source: Mario, red hat and blue overalls] + instruction "make it anime style"
+→ {"violates": true, "categories": ["copyright"], "reason": "Source depicts Mario, a copyrighted Nintendo character; any edit of it is disallowed."}
+
+[source: a photo of Donald Trump at a podium] + instruction "change the background to a beach"
+→ {"violates": true, "categories": ["public_figure"], "reason": "Source depicts a recognizable real public figure; editing his photo is disallowed regardless of the edit."}
+
+[source: a generic golden retriever puppy on grass] + instruction "change the background to a beach"
+→ {"violates": false, "categories": [], "reason": "Ordinary non-sensitive photo and an innocuous edit."}
+
+[source: an ordinary unknown man in a suit on a street] + instruction "change the background to a beach"
+→ {"violates": false, "categories": [], "reason": "Not a recognizable public figure; ordinary person, innocuous edit — allowed."}"""
+
+
+CATEGORY_DISPLAY = {
+    "sexual": "Sexual content",
+    "hate": "Hate / unfair imagery",
+    "self_harm": "Self-harm",
+    "violence": "Violence / gore",
+    "copyright": "Copyright / IP character",
+    "public_figure": "Real-person likeness",
+}
+
+
+@dataclass
+class FilterVerdict:
+    violates: bool
+    categories: list[str]
+    reason: str
+    raw: str = ""
+
+    def banner(self) -> str:
+        if not self.violates:
+            return ""
+        cat = ", ".join(CATEGORY_DISPLAY.get(c, c) for c in self.categories) or "policy violation"
+        return f"🚫 **Content Filter:** Blocked — `{cat}` · {self.reason}"
+
+
+def _extract_json_object(text: str) -> dict:
+    """Pull the first balanced top-level JSON object out of a possibly-wrapped string."""
+    if not text:
+        raise ValueError("empty response")
+    # Strip code fences
+    if text.lstrip().startswith("```"):
+        text = text.strip().strip("`")
+        if text.lstrip().startswith("json"):
+            text = text.lstrip()[4:]
+    start = text.find("{")
+    if start == -1:
+        raise ValueError(f"no JSON object found in: {text[:120]!r}")
+    depth = 0
+    for i in range(start, len(text)):
+        c = text[i]
+        if c == "{":
+            depth += 1
+        elif c == "}":
+            depth -= 1
+            if depth == 0:
+                return json.loads(text[start : i + 1])
+    raise ValueError(f"unbalanced JSON object: {text[start:start+120]!r}")
+
+
+def check_prompt(model, prompt: str, max_new_tokens: int = 160) -> FilterVerdict:
+    """Back-compat wrapper around the mandatory text-encoder screener.
+
+    The policy check now lives on the text encoder
+    (:meth:`TextEncoder.screen_text`) so it runs on the same Qwen3-VL weights
+    that produce the diffusion conditioning and is FAIL-CLOSED. Kept here for
+    external callers / tests that import ``check_prompt`` directly.
+    """
+    return model.txt_enc.screen_text(prompt, max_new_tokens=max_new_tokens)
+
+
+@contextmanager
+def _full_output_mode(hf):
+    """Temporarily switch the Qwen3-VL encoder into FULL output mode so that
+    ``.generate()`` sees ``.logits`` (the diffusion path uses embedding mode).
+    Restores the original mode/skip flags on exit — side-effect free."""
+    prev_mode = getattr(hf, "_output_mode", None)
+    prev_skip = getattr(hf, "_skip_lm_head", None)
+    try:
+        if hasattr(hf, "set_output_mode"):
+            try:
+                hf.set_output_mode("full")
+            except Exception:  # noqa: BLE001
+                if prev_skip is not None:
+                    hf._skip_lm_head = False
+                if prev_mode is not None:
+                    hf._output_mode = "full"
+        elif prev_skip is not None:
+            hf._skip_lm_head = False
+        yield
+    finally:
+        try:
+            if prev_mode is not None and hasattr(hf, "set_output_mode"):
+                hf.set_output_mode(prev_mode)
+            if prev_skip is not None:
+                hf._skip_lm_head = prev_skip
+        except Exception:  # noqa: BLE001
+            pass
+
+
+def check_edit(model, prompt: str, ref_images, max_new_tokens: int = 192) -> FilterVerdict:
+    """Back-compat wrapper around :meth:`TextEncoder.screen_edit`.
+
+    Classifies an image-EDIT request considering BOTH the source image(s) and
+    the instruction (multimodal Qwen3-VL), FAIL-CLOSED. Kept for external
+    callers / tests that import ``check_edit`` directly.
+    """
+    return model.txt_enc.screen_edit(prompt, ref_images, max_new_tokens=max_new_tokens)
+
+
+def make_refusal_image(
+    verdict: FilterVerdict,
+    height: int = 1024,
+    width: int = 1024,
+) -> Image.Image:
+    """Return a placeholder image to display when the prompt is blocked.
+
+    A plain white blank image — no text, no category/reason surfaced.
+    """
+    return Image.new("RGB", (width, height), color=(255, 255, 255))

vendor/mage_flow/models/modules/mage_vae.py

@@ -0,0 +1,651 @@
+"""
+MageVAE: DConvEncoder + DConvDenoiser (with CoD Decoder) wrapper.
+
+Replaces FLUX2 VAE for encoding images to latents and decoding latents back to images.
+Supports only the kl0.1 CoD ckpt layout:
+  encoder weights: 'state_dict' → 'student.dconv_encoder.*' (packed mean+logvar, out_ch_mult=2)
+  decoder weights: 'state_dict' → 'pipeline.*' (denoiser + y_embedder.decoder)
+
+Latent shape: [B, 128, H/16, W/16] — no patch packing, no BN normalization.
+"""
+
+import math
+import os
+from functools import lru_cache
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from loguru import logger
+
+
+# ---------------------------------------------------------------------------
+# Primitive layers (vendored from GenCodec, inference subset)
+# ---------------------------------------------------------------------------
+def nonlinearity(x):
+    return x * torch.sigmoid(x)
+
+
+def Normalize(in_channels):
+    return torch.nn.GroupNorm(num_groups=32, num_channels=in_channels, eps=1e-6, affine=True)
+
+
+def modulate(x, shift, scale):
+    if x.dim() == 4:
+        b, c = x.shape[:2]
+        return x * (1 + scale.view(b, c, 1, 1)) + shift.view(b, c, 1, 1)
+    return x * (1 + scale.unsqueeze(1)) + shift.unsqueeze(1)
+
+
+class LayerNorm2d(nn.LayerNorm):
+    def __init__(self, num_channels, eps=1e-6, affine=True):
+        super().__init__(num_channels, eps=eps, elementwise_affine=affine)
+
+    def forward(self, x):
+        # .contiguous() prevents a channels_last-strided NCHW view from
+        # propagating into downstream depthwise convs, which would otherwise
+        # hit a slow cuDNN path with a per-shape heuristic search.
+        x = x.permute(0, 2, 3, 1).contiguous()
+        x = F.layer_norm(x, self.normalized_shape, self.weight, self.bias, self.eps)
+        return x.permute(0, 3, 1, 2).contiguous()
+
+
+class _EncoderLayerNorm2d(LayerNorm2d):
+    pass
+
+
+class RMSNorm(nn.Module):
+    def __init__(self, hidden_size, eps=1e-6):
+        super().__init__()
+        self.weight = nn.Parameter(torch.ones(hidden_size))
+        self.variance_epsilon = eps
+
+    def forward(self, x):
+        in_dtype = x.dtype
+        x = x.to(torch.float32)
+        var = x.pow(2).mean(-1, keepdim=True)
+        x = x * torch.rsqrt(var + self.variance_epsilon)
+        return self.weight * x.to(in_dtype)
+
+
+class TimestepEmbedder(nn.Module):
+    """DConv-style timestep MLP (max_period=10000, freq_size=256, hidden=384)."""
+
+    def __init__(self, hidden_size, frequency_embedding_size=256):
+        super().__init__()
+        self.mlp = nn.Sequential(
+            nn.Linear(frequency_embedding_size, hidden_size, bias=True),
+            nn.SiLU(),
+            nn.Linear(hidden_size, hidden_size, bias=True),
+        )
+        self.frequency_embedding_size = frequency_embedding_size
+
+    @staticmethod
+    def timestep_embedding(t, dim, max_period=10000):
+        half = dim // 2
+        freqs = torch.exp(
+            -math.log(max_period) * torch.arange(0, half, dtype=torch.float32) / half
+        ).to(t.device)
+        args = t[:, None].float() * freqs[None]
+        emb = torch.cat([torch.cos(args), torch.sin(args)], dim=-1)
+        if dim % 2:
+            emb = torch.cat([emb, torch.zeros_like(emb[:, :1])], dim=-1)
+        return emb
+
+    def forward(self, t):
+        emb = self.timestep_embedding(t, self.frequency_embedding_size)
+        return self.mlp(emb.to(self.mlp[0].weight.dtype))
+
+
+class BottleneckPatchEmbed(nn.Module):
+    """Image patch embed concatenated with a per-patch conditioning vector."""
+
+    def __init__(self, patch_size=16, in_chans=3, pca_dim=128, embed_dim=384, bias=True):
+        super().__init__()
+        self.proj1 = nn.Conv2d(in_chans, pca_dim, kernel_size=patch_size, stride=patch_size, bias=False)
+        self.proj2 = nn.Conv2d(pca_dim + embed_dim, embed_dim, kernel_size=1, bias=bias)
+
+    def forward(self, x, cond):
+        return self.proj2(torch.cat([self.proj1(x), cond], dim=1))
+
+
+class DiCoBlock(nn.Module):
+    """DConv block with adaLN modulation."""
+
+    def __init__(self, hidden_size, mlp_ratio=4.0):
+        super().__init__()
+        self.conv1 = nn.Conv2d(hidden_size, hidden_size, 1, bias=True)
+        self.conv2 = nn.Conv2d(hidden_size, hidden_size, 3, padding=1, groups=hidden_size, bias=True)
+        self.conv3 = nn.Conv2d(hidden_size, hidden_size, 1, bias=True)
+
+        self.ca = nn.Sequential(
+            nn.AdaptiveAvgPool2d(1),
+            nn.Conv2d(hidden_size, hidden_size, 1, bias=True),
+            nn.Sigmoid(),
+        )
+
+        ffn = int(mlp_ratio * hidden_size)
+        self.conv4 = nn.Conv2d(hidden_size, ffn, 1, bias=True)
+        self.conv5 = nn.Conv2d(ffn, hidden_size, 1, bias=True)
+
+        self.norm1 = LayerNorm2d(hidden_size, affine=False)
+        self.norm2 = LayerNorm2d(hidden_size, affine=False)
+
+        self.adaLN_modulation = nn.Sequential(
+            nn.SiLU(),
+            nn.Linear(hidden_size, 6 * hidden_size, bias=True),
+        )
+
+    def forward(self, inp, c):
+        shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = self.adaLN_modulation(c).chunk(6, dim=1)
+        x = modulate(self.norm1(inp), shift_msa, scale_msa)
+        x = F.gelu(self.conv2(self.conv1(x)))
+        x = x * self.ca(x)
+        x = self.conv3(x)
+        x = inp + gate_msa[..., None, None] * x
+        x = x + gate_mlp[..., None, None] * self.conv5(
+            F.gelu(self.conv4(modulate(self.norm2(x), shift_mlp, scale_mlp)))
+        )
+        return x
+
+
+class _EncoderDiCoBlock(nn.Module):
+    """DiCoBlock without adaLN, for the encoder pathway."""
+
+    def __init__(self, hidden_size, mlp_ratio=4.0):
+        super().__init__()
+        self.conv1 = nn.Conv2d(hidden_size, hidden_size, 1, bias=True)
+        self.conv2 = nn.Conv2d(hidden_size, hidden_size, 3, padding=1, groups=hidden_size, bias=True)
+        self.conv3 = nn.Conv2d(hidden_size, hidden_size, 1, bias=True)
+        self.ca = nn.Sequential(
+            nn.AdaptiveAvgPool2d(1),
+            nn.Conv2d(hidden_size, hidden_size, 1, bias=True),
+            nn.Sigmoid(),
+        )
+        ffn = int(mlp_ratio * hidden_size)
+        self.conv4 = nn.Conv2d(hidden_size, ffn, 1, bias=True)
+        self.conv5 = nn.Conv2d(ffn, hidden_size, 1, bias=True)
+        self.norm1 = _EncoderLayerNorm2d(hidden_size)
+        self.norm2 = _EncoderLayerNorm2d(hidden_size)
+
+    def forward(self, inp):
+        x = self.norm1(inp)
+        x = F.gelu(self.conv2(self.conv1(x)))
+        x = x * self.ca(x)
+        x = self.conv3(x)
+        x = inp + x
+        return x + self.conv5(F.gelu(self.conv4(self.norm2(x))))
+
+
+class NerfEmbedder(nn.Module):
+    """Patch-position embedder used by the DConv decoder x-pathway."""
+
+    def __init__(self, in_channels, hidden_size_input, max_freqs=8):
+        super().__init__()
+        self.max_freqs = max_freqs
+        self.embedder = nn.Sequential(
+            nn.Linear(in_channels + max_freqs ** 2, hidden_size_input, bias=True),
+        )
+
+    @lru_cache
+    def fetch_pos(self, patch_size, device, dtype):
+        pos = torch.linspace(0, 1, patch_size, device=device, dtype=dtype)
+        pos_y, pos_x = torch.meshgrid(pos, pos, indexing="ij")
+        pos_x = pos_x.reshape(-1, 1, 1)
+        pos_y = pos_y.reshape(-1, 1, 1)
+        freqs = torch.linspace(0, self.max_freqs, self.max_freqs, dtype=dtype, device=device)
+        fx = freqs[None, :, None]
+        fy = freqs[None, None, :]
+        coeffs = (1 + fx * fy) ** -1
+        dct_x = torch.cos(pos_x * fx * torch.pi)
+        dct_y = torch.cos(pos_y * fy * torch.pi)
+        return (dct_x * dct_y * coeffs).view(1, -1, self.max_freqs ** 2)
+
+    def forward(self, x):
+        B, P2, _ = x.shape
+        ps = int(P2 ** 0.5)
+        dct = self.fetch_pos(ps, x.device, x.dtype).expand(B, -1, -1)
+        return self.embedder(torch.cat([x, dct], dim=-1))
+
+
+class NerfFinalLayer(nn.Module):
+    def __init__(self, hidden_size, out_channels):
+        super().__init__()
+        self.norm = RMSNorm(hidden_size)
+        self.linear = nn.Linear(hidden_size, out_channels, bias=True)
+
+    def forward(self, x):
+        return self.linear(self.norm(x))
+
+
+class SimpleMLPAdaLN(nn.Module):
+    """Final small MLP that maps NerfEmbedder features to per-patch RGB."""
+
+    def __init__(self, in_channels, model_channels, out_channels, z_channels, num_res_blocks, patch_size):
+        super().__init__()
+        self.in_channels = in_channels
+        self.model_channels = model_channels
+        self.out_channels = out_channels
+        self.num_res_blocks = num_res_blocks
+        self.patch_size = patch_size
+
+        self.cond_embed = nn.Linear(z_channels, patch_size ** 2 * model_channels)
+        self.input_proj = nn.Linear(in_channels, model_channels)
+
+        self.res_blocks = nn.ModuleList(_MLPResBlock(model_channels) for _ in range(num_res_blocks))
+
+    def forward(self, x, c):
+        x = self.input_proj(x)
+        c = self.cond_embed(c).reshape(c.shape[0], self.patch_size ** 2, -1)
+        for block in self.res_blocks:
+            x = block(x, c)
+        return x
+
+
+class _MLPResBlock(nn.Module):
+    def __init__(self, channels):
+        super().__init__()
+        self.in_ln = nn.LayerNorm(channels, eps=1e-6)
+        self.mlp = nn.Sequential(
+            nn.Linear(channels, channels, bias=True),
+            nn.SiLU(),
+            nn.Linear(channels, channels, bias=True),
+        )
+        self.adaLN_modulation = nn.Sequential(
+            nn.SiLU(),
+            nn.Linear(channels, 3 * channels, bias=True),
+        )
+
+    def forward(self, x, y):
+        shift, scale, gate = self.adaLN_modulation(y).chunk(3, dim=-1)
+        h = self.in_ln(x) * (1 + scale) + shift
+        return x + gate * self.mlp(h)
+
+
+class ResnetBlock(nn.Module):
+    """GroupNorm + Conv ResBlock used by the CoD Decoder."""
+
+    def __init__(self, *, in_channels, out_channels=None, dropout=0.0):
+        super().__init__()
+        out_channels = out_channels or in_channels
+        self.in_channels = in_channels
+        self.out_channels = out_channels
+
+        self.norm1 = Normalize(in_channels)
+        self.conv1 = nn.Conv2d(in_channels, out_channels, 3, padding=1)
+        self.norm2 = Normalize(out_channels)
+        self.dropout = nn.Dropout(dropout)
+        self.conv2 = nn.Conv2d(out_channels, out_channels, 3, padding=1)
+        if in_channels != out_channels:
+            self.nin_shortcut = nn.Conv2d(in_channels, out_channels, 1)
+
+    def forward(self, x):
+        h = self.conv1(nonlinearity(self.norm1(x)))
+        h = self.conv2(self.dropout(nonlinearity(self.norm2(h))))
+        if self.in_channels != self.out_channels:
+            x = self.nin_shortcut(x)
+        return x + h
+
+
+class AttnBlock(nn.Module):
+    """Patched self-attention used at inference (eval mode of the original)."""
+
+    def __init__(self, in_channels, patch_size=32):
+        super().__init__()
+        self.in_channels = in_channels
+        self.patch_size = patch_size
+        self.norm = Normalize(in_channels)
+        self.q = nn.Conv2d(in_channels, in_channels, 1)
+        self.k = nn.Conv2d(in_channels, in_channels, 1)
+        self.v = nn.Conv2d(in_channels, in_channels, 1)
+        self.proj_out = nn.Conv2d(in_channels, in_channels, 1)
+
+    def forward(self, x):
+        h_ = self.norm(x)
+        Q = self.q(h_)
+        K = self.k(h_)
+        V = self.v(h_)
+
+        d = self.patch_size
+        b, c, H, W = Q.shape
+        pad_h = (d - H % d) % d
+        pad_w = (d - W % d) % d
+        if pad_h or pad_w:
+            Q = F.pad(Q, (0, pad_w, 0, pad_h), mode="replicate")
+            K = F.pad(K, (0, pad_w, 0, pad_h), mode="replicate")
+            V = F.pad(V, (0, pad_w, 0, pad_h), mode="replicate")
+        _, _, H_pad, W_pad = Q.shape
+        nph, npw = H_pad // d, W_pad // d
+        np_ = nph * npw
+
+        def to_patches(t):
+            return (t.reshape(b, c, nph, d, npw, d)
+                    .permute(0, 2, 4, 1, 3, 5)
+                    .reshape(b * np_, c, d * d))
+
+        Q = to_patches(Q)
+        K = to_patches(K)
+        V = to_patches(V)
+
+        w_ = torch.bmm(Q.permute(0, 2, 1), K) * (c ** -0.5)
+        w_ = F.softmax(w_, dim=2).permute(0, 2, 1)
+        h_ = torch.bmm(V, w_).reshape(b, nph, npw, c, d, d).permute(0, 3, 1, 4, 2, 5).reshape(b, c, H_pad, W_pad)
+        if pad_h or pad_w:
+            h_ = h_[:, :, :H, :W]
+        return x + self.proj_out(h_)
+
+
+# ---------------------------------------------------------------------------
+# adaLN constant-folding: at fixed t=0, adaLN_modulation(c) is constant.
+# Replace the MLP with a buffer so DiCoBlock.forward stays unchanged and
+# torch.compile can fuse the surrounding ops normally.
+# ---------------------------------------------------------------------------
+class _ConstAdaLN(nn.Module):
+    def __init__(self, modulation: torch.Tensor):
+        super().__init__()
+        self.register_buffer("modulation", modulation.detach().clone())
+
+    def forward(self, c):
+        b = c.shape[0]
+        if self.modulation.shape[0] != b:
+            return self.modulation.expand(b, *self.modulation.shape[1:])
+        return self.modulation
+
+
+def _replace_adaln_with_const(module: nn.Module, c: torch.Tensor) -> int:
+    # Only DiCoBlock is targeted: its adaLN is conditioned solely on t.
+    # Other adaLN_modulation submodules (e.g. _MLPResBlock in the decoder MLP)
+    # take a per-position latent and must not be folded.
+    n = 0
+    for child in module.modules():
+        if not isinstance(child, DiCoBlock):
+            continue
+        adaln = child.adaLN_modulation
+        if isinstance(adaln, _ConstAdaLN):
+            continue
+        with torch.no_grad():
+            mod = adaln(c)
+        child.adaLN_modulation = _ConstAdaLN(mod)
+        n += 1
+    return n
+
+
+# ---------------------------------------------------------------------------
+# CoD Decoder: latent → conditioning features for the denoiser
+# ---------------------------------------------------------------------------
+class _Decoder(nn.Module):
+    """ds=16, up2x=True, light=True only."""
+
+    def __init__(self, out_ch=384, z_ch=128):
+        super().__init__()
+        self.conv_in = nn.Conv2d(z_ch, out_ch, kernel_size=3, stride=1, padding=1)
+        self.block = nn.Sequential(
+            ResnetBlock(in_channels=out_ch, out_channels=out_ch),
+            AttnBlock(out_ch, patch_size=32),
+            ResnetBlock(in_channels=out_ch, out_channels=out_ch),
+            AttnBlock(out_ch, patch_size=32),
+            ResnetBlock(in_channels=out_ch, out_channels=out_ch),
+        )
+        self.norm_out = Normalize(out_ch)
+        self.conv_out = nn.Conv2d(out_ch, out_ch, kernel_size=3, stride=1, padding=1)
+        self.ada = nn.Identity()
+
+    def forward(self, z):
+        h = self.block(self.conv_in(z))
+        h = self.conv_out(nonlinearity(self.norm_out(h)))
+        return self.ada(h)
+
+
+# ---------------------------------------------------------------------------
+# DConvEncoder: image → packed (mean, logvar) latent
+# ---------------------------------------------------------------------------
+class _DConvEncoder(nn.Module):
+    def __init__(
+        self,
+        z_ch=128,
+        hidden_size=384,
+        num_blocks=21,
+        patch_size=16,
+        mlp_ratio=4.0,
+        head_size=768,
+        num_head_blocks=2,
+        out_ch_mult=2,
+    ):
+        super().__init__()
+        self.z_ch = z_ch
+        self.patch_size = patch_size
+        self.patch_cond_embed = nn.Conv2d(3, head_size, kernel_size=patch_size, stride=patch_size, bias=True)
+        self.head_blocks = nn.ModuleList([
+            _EncoderDiCoBlock(head_size, mlp_ratio=mlp_ratio) for _ in range(num_head_blocks)
+        ])
+        self.proj_down = nn.Conv2d(head_size, hidden_size, kernel_size=1, bias=True)
+        self.z_proj = nn.Conv2d(z_ch, hidden_size, kernel_size=1, bias=True)
+        self.fuse_proj = nn.Conv2d(hidden_size * 2, hidden_size, kernel_size=1, bias=True)
+        self.t_embedder = TimestepEmbedder(hidden_size)
+        self.blocks = nn.ModuleList([
+            DiCoBlock(hidden_size, mlp_ratio=mlp_ratio) for _ in range(num_blocks)
+        ])
+        self.norm_out = LayerNorm2d(hidden_size)
+        self.proj_out = nn.Conv2d(hidden_size, z_ch * out_ch_mult, kernel_size=1, bias=True)
+
+    def forward_pred(self, z_t, t, y):
+        cond = self.patch_cond_embed(y)
+        for block in self.head_blocks:
+            cond = block(cond)
+        cond = self.proj_down(cond)
+
+        s = self.fuse_proj(torch.cat([cond, self.z_proj(z_t)], dim=1))
+        c = self.t_embedder(t.view(-1))
+        for block in self.blocks:
+            s = block(s, c)
+        return self.proj_out(self.norm_out(s))
+
+
+# ---------------------------------------------------------------------------
+# DConv denoiser: latent (via cond) + zero noise → reconstructed image
+# ---------------------------------------------------------------------------
+class _YEmbedder(nn.Module):
+    """Holds only the CoD decoder; the original Flux2 VAE encoder side is omitted."""
+
+    def __init__(self, ch=384, z_ch=128):
+        super().__init__()
+        self.decoder = _Decoder(out_ch=ch, z_ch=z_ch)
+
+
+class _DConvDenoiser(nn.Module):
+    def __init__(
+        self,
+        patch_size=16,
+        in_channels=3,
+        hidden_size=384,
+        hidden_size_x=32,
+        mlp_ratio=4.0,
+        num_blocks=24,
+        num_cond_blocks=21,
+        bottleneck_dim=128,
+    ):
+        super().__init__()
+        self.in_channels = in_channels
+        self.patch_size = patch_size
+        self.hidden_size = hidden_size
+        self.num_cond_blocks = num_cond_blocks
+
+        self.t_embedder = TimestepEmbedder(hidden_size)
+        self.y_embedder_x = nn.Conv2d(hidden_size, hidden_size_x * patch_size ** 2, 1, 1, 0)
+        self.x_embedder = NerfEmbedder(in_channels + hidden_size_x, hidden_size_x, max_freqs=8)
+        self.s_embedder = BottleneckPatchEmbed(patch_size, in_channels, bottleneck_dim, hidden_size, bias=True)
+        self.blocks = nn.ModuleList([
+            DiCoBlock(hidden_size, mlp_ratio=mlp_ratio) for _ in range(num_cond_blocks)
+        ])
+        self.dec_net = SimpleMLPAdaLN(
+            in_channels=hidden_size_x,
+            model_channels=hidden_size_x,
+            out_channels=in_channels,
+            z_channels=hidden_size,
+            num_res_blocks=num_blocks - num_cond_blocks,
+            patch_size=patch_size,
+        )
+        self.final_layer = NerfFinalLayer(hidden_size_x, in_channels)
+        self.y_embedder = _YEmbedder(ch=hidden_size, z_ch=bottleneck_dim)
+
+    def forward(self, x, t, cond):
+        b, _, h, w = x.shape
+        c = self.t_embedder(t.view(-1))
+
+        s = self.s_embedder(x, cond)
+        for block in self.blocks:
+            s = block(s, c)
+
+        length = s.shape[-2] * s.shape[-1]
+        s = s.permute(0, 2, 3, 1).reshape(-1, self.hidden_size)
+
+        x = torch.nn.functional.unfold(x, kernel_size=self.patch_size, stride=self.patch_size)
+        x = torch.cat([x, self.y_embedder_x(cond).flatten(2)], dim=1)
+        x = x.reshape(b, -1, self.patch_size ** 2, length).permute(0, 3, 2, 1).flatten(0, 1)
+        x = self.x_embedder(x)
+
+        x = self.dec_net(x, s)
+        x = self.final_layer(x)
+        x = x.transpose(1, 2).reshape(b, length, -1)
+        return torch.nn.functional.fold(
+            x.transpose(1, 2).contiguous(), (h, w),
+            kernel_size=self.patch_size, stride=self.patch_size,
+        )
+
+
+# ---------------------------------------------------------------------------
+# Wrapper
+# ---------------------------------------------------------------------------
+def _load_state_dict(ckpt_path: str):
+    if ckpt_path.endswith(".safetensors"):
+        from safetensors.torch import load_file
+        return load_file(ckpt_path, device="cpu")
+    if os.path.exists(os.path.join(ckpt_path, "checkpoint-state_dict.pt")):
+        ckpt_path = os.path.join(ckpt_path, "checkpoint-state_dict.pt")
+    elif os.path.isdir(ckpt_path):
+        ckpt_path = os.path.join(ckpt_path, "checkpoint", "mp_rank_00_model_states.pt")
+    state = torch.load(ckpt_path, map_location="cpu")
+    if "module" in state:
+        return state["module"]
+    if "state_dict" in state:
+        return state["state_dict"]
+    return state
+
+
+class MageVAE(nn.Module):
+    """
+    Encode: DConvEncoder (one-step diffusion) → latent [B, 128, H/16, W/16]
+    Decode: DConvDenoiser + CoD Decoder       → image  [B, 3, H, W] in [-1, 1]
+    """
+
+    latent_channels = 128
+    downsample_factor = 16
+
+    def __init__(self, ckpt_path: str, sample_posterior: bool = True):
+        super().__init__()
+        self.sample_posterior = sample_posterior
+
+        self.dconv_encoder = _DConvEncoder()
+        self.decoder_model = _DConvDenoiser()
+
+        sd = _load_state_dict(ckpt_path)
+        self._load_encoder(sd, ckpt_path)
+        self._load_decoder(sd, ckpt_path)
+
+        # adaLN modulation depends only on t, and we always run at t=0.
+        # Precompute and drop the MLPs once at construction (~37M params saved).
+        self._freeze_adaln_cache()
+
+    def _load_encoder(self, sd, ckpt_path):
+        prefix = "student.dconv_encoder."
+        enc_sd = {k[len(prefix):]: v for k, v in sd.items() if k.startswith(prefix)}
+        if not enc_sd:
+            raise RuntimeError(f"CoDEncoder: no '{prefix}*' keys in {ckpt_path}")
+        proj = enc_sd.get("proj_out.weight")
+        if proj is None or proj.shape[0] != 2 * self.latent_channels:
+            raise RuntimeError(
+                f"CoDEncoder: expected packed mean+logvar (proj_out out_channels="
+                f"{2 * self.latent_channels}), got {None if proj is None else tuple(proj.shape)}"
+            )
+        missing, unexpected = self.dconv_encoder.load_state_dict(enc_sd, strict=False)
+        logger.info(
+            f"CoDEncoder: loaded {len(enc_sd)} keys, "
+            f"missing={len(missing)}, unexpected={len(unexpected)}"
+        )
+        if missing:
+            logger.warning(f"CoDEncoder missing: {missing[:10]}")
+
+    def _load_decoder(self, sd, ckpt_path):
+        prefix = "pipeline."
+        if not any(k.startswith(prefix) for k in sd):
+            raise RuntimeError(f"CoDDecoder: no '{prefix}*' keys in {ckpt_path}")
+        model_dict = self.decoder_model.state_dict()
+        matched = {}
+        for k, v in sd.items():
+            if not k.startswith(prefix):
+                continue
+            new_k = k[len(prefix):]
+            if new_k.startswith("y_embedder.encoder.") or new_k.startswith("y_embedder.bottleneck."):
+                continue
+            if new_k in model_dict and model_dict[new_k].shape == v.shape:
+                matched[new_k] = v
+        self.decoder_model.load_state_dict(matched, strict=False)
+        logger.info(f"CoDDecoder: loaded {len(matched)} params (denoiser + y_embedder.decoder)")
+        if not matched:
+            raise RuntimeError(f"CoDDecoder: 0 params matched from {ckpt_path}")
+
+    @torch.no_grad()
+    def _moments(self, x: torch.Tensor):
+        B, _, H, W = x.shape
+        ps = self.dconv_encoder.patch_size
+        z_t = torch.zeros(B, self.dconv_encoder.z_ch, H // ps, W // ps, device=x.device, dtype=x.dtype)
+        t = torch.zeros(B, device=x.device, dtype=x.dtype)
+        out = self.dconv_encoder.forward_pred(z_t, t, x)
+        mean = out[:, : self.latent_channels]
+        logvar = out[:, self.latent_channels :].clamp(min=-20.0, max=10.0)
+        return mean, logvar
+
+    @torch.no_grad()
+    def _encode_moments(self, x: torch.Tensor):
+        # Compile target: pure deterministic part of encode (no RNG, no
+        # asserts), so torch.compile produces a single dynamic graph.
+        return self._moments(x)
+
+    @torch.no_grad()
+    def encode(self, x: torch.Tensor) -> torch.Tensor:
+        ps = self.dconv_encoder.patch_size
+        H, W = x.shape[-2], x.shape[-1]
+        if H % ps or W % ps:
+            raise ValueError(f"H, W must be multiples of {ps}, got ({H}, {W})")
+        mean, logvar = self._encode_moments(x)
+        if self.sample_posterior:
+            return mean + torch.exp(0.5 * logvar) * torch.randn_like(mean)
+        return mean
+
+    @torch.no_grad()
+    def decode(self, z: torch.Tensor) -> torch.Tensor:
+        cond = self.decoder_model.y_embedder.decoder(z)
+        B = z.shape[0]
+        H = z.shape[2] * self.downsample_factor
+        W = z.shape[3] * self.downsample_factor
+        noise = torch.zeros(B, 3, H, W, device=z.device, dtype=z.dtype)
+        t = torch.zeros(B, device=z.device, dtype=z.dtype)
+        return self.decoder_model.forward(noise, t, cond)
+
+    @property
+    def device(self):
+        return next(self.parameters()).device
+
+    @property
+    def dtype(self):
+        return next(self.parameters()).dtype
+
+    def _freeze_adaln_cache(self):
+        """Constant-fold adaLN_modulation MLPs at t=0 (encoder + decoder)."""
+        device = next(self.parameters()).device
+        dtype = next(self.parameters()).dtype
+        t = torch.zeros(1, device=device, dtype=dtype)
+        c_enc = self.dconv_encoder.t_embedder(t)
+        _replace_adaln_with_const(self.dconv_encoder, c_enc)
+        c_dec = self.decoder_model.t_embedder(t)
+        _replace_adaln_with_const(self.decoder_model, c_dec)

vendor/mage_flow/models/modules/text_encoder.py

@@ -0,0 +1,707 @@
+"""Qwen3-VL text encoder: custom HF model + packing-aware forward patches + TextEncoder wrapper."""
+from __future__ import annotations
+
+import os
+from collections.abc import Callable
+from dataclasses import dataclass
+
+try:
+    from typing import Unpack
+except ImportError:
+    from typing_extensions import Unpack
+
+import torch
+from loguru import logger
+from torch import nn
+from transformers import AutoProcessor, AutoTokenizer, Cache, Qwen3VLForConditionalGeneration
+from transformers.cache_utils import DynamicCache
+from transformers.masking_utils import create_causal_mask
+from transformers.modeling_flash_attention_utils import FlashAttentionKwargs
+from transformers.modeling_outputs import BaseModelOutputWithPast
+from transformers.modeling_utils import ALL_ATTENTION_FUNCTIONS
+from transformers.models.qwen3_vl.modeling_qwen3_vl import (
+    Qwen3VLCausalLMOutputWithPast,
+    apply_rotary_pos_emb,
+    eager_attention_forward,
+)
+from transformers.utils import ModelOutput
+
+from ._attn_backend import flash_attn_varlen_func
+
+
+# ===========================================================================
+# Custom Qwen3-VL model (customizable forward output)
+# ===========================================================================
+
+@dataclass
+class Qwen3VLModelOutput(ModelOutput):
+    """Flexible output class for custom Qwen3-VL model."""
+
+    loss: torch.FloatTensor | None = None
+    logits: torch.FloatTensor | None = None
+    past_key_values: Cache | None = None
+    hidden_states: tuple[torch.FloatTensor, ...] | None = None
+    last_hidden_state: torch.FloatTensor | None = None
+    attentions: tuple[torch.FloatTensor, ...] | None = None
+    rope_deltas: torch.LongTensor | None = None
+
+
+class CustomQwen3VLForConditionalGeneration(Qwen3VLForConditionalGeneration):
+    """
+    Custom Qwen3-VL model that allows customizing the forward output.
+
+    This class inherits from Qwen3VLForConditionalGeneration and provides
+    hooks to customize what is returned from the forward pass.
+
+    Example usage:
+        ```python
+        model = CustomQwen3VLForConditionalGeneration.from_pretrained(
+            "Qwen/Qwen3-VL-8B-Instruct",
+            attn_implementation="flash_attention_2"  # Use flash attention for faster inference
+        )
+
+        # Option 1: Use built-in output modes
+        model.set_output_mode("embedding")  # Only return last hidden state (default)
+        model.set_output_mode("full")       # Return everything
+        model.set_output_mode("logits")     # Only return logits
+
+        # Option 2: Set a custom output processor
+        def my_custom_output(hidden_states, logits, outputs, **kwargs):
+            return {"embeddings": hidden_states, "pooled": hidden_states.mean(dim=1)}
+        model.set_output_processor(my_custom_output)
+        ```
+    """
+
+    # Output mode constants
+    OUTPUT_MODE_FULL = "full"
+    OUTPUT_MODE_EMBEDDING = "embedding"
+    OUTPUT_MODE_LOGITS = "logits"
+    OUTPUT_MODE_HIDDEN = "hidden"
+
+    def __init__(self, config):
+        super().__init__(config)
+        self._output_mode = self.OUTPUT_MODE_EMBEDDING
+        self._skip_lm_head = True
+
+    def set_output_mode(self, mode: str):
+        """
+        Set the output mode for the forward pass.
+
+        Args:
+            mode: One of:
+                - "full": Return full Qwen3VLCausalLMOutputWithPast
+                - "embedding": Only return last hidden state (skip lm_head) (default)
+                - "logits": Only return logits
+                - "hidden": Return all hidden states
+        """
+        valid_modes = [
+            self.OUTPUT_MODE_FULL,
+            self.OUTPUT_MODE_EMBEDDING,
+            self.OUTPUT_MODE_LOGITS,
+            self.OUTPUT_MODE_HIDDEN,
+        ]
+        if mode not in valid_modes:
+            raise ValueError(f"Invalid output mode: {mode}. Must be one of {valid_modes}")
+        self._output_mode = mode
+        self._skip_lm_head = mode == self.OUTPUT_MODE_EMBEDDING
+
+    def forward(
+        self,
+        input_ids: torch.LongTensor | None = None,
+        attention_mask: torch.Tensor | None = None,
+        position_ids: torch.LongTensor | None = None,
+        past_key_values: Cache | None = None,
+        inputs_embeds: torch.FloatTensor | None = None,
+        labels: torch.LongTensor | None = None,
+        pixel_values: torch.Tensor | None = None,
+        pixel_values_videos: torch.FloatTensor | None = None,
+        image_grid_thw: torch.LongTensor | None = None,
+        video_grid_thw: torch.LongTensor | None = None,
+        cache_position: torch.LongTensor | None = None,
+        logits_to_keep: int | torch.Tensor = 0,
+        output_attentions: bool | None = None,
+        output_hidden_states: bool | None = None,
+        return_dict: bool | None = None,
+        **kwargs,
+    ) -> Qwen3VLCausalLMOutputWithPast | Qwen3VLModelOutput | dict | torch.Tensor:
+        """
+        Forward pass with customizable output.
+
+        Returns different outputs based on the configured output mode or custom processor.
+        """
+        output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
+        output_hidden_states = (
+            output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
+        )
+
+        # Get outputs from the base model (Qwen3VLModel)
+        outputs = self.model(
+            input_ids=input_ids,
+            pixel_values=pixel_values,
+            pixel_values_videos=pixel_values_videos,
+            image_grid_thw=image_grid_thw,
+            video_grid_thw=video_grid_thw,
+            position_ids=position_ids,
+            attention_mask=attention_mask,
+            past_key_values=past_key_values,
+            inputs_embeds=inputs_embeds,
+            cache_position=cache_position,
+            output_attentions=output_attentions,
+            output_hidden_states=output_hidden_states,
+            return_dict=True,
+            **kwargs,
+        )
+
+        # Get the last hidden state
+        hidden_states = outputs[0]  # This is the last hidden state
+
+        # Compute logits if not skipping lm_head
+        logits = None
+        if not self._skip_lm_head:
+            slice_indices = slice(-logits_to_keep, None) if isinstance(logits_to_keep, int) else logits_to_keep
+            logits = self.lm_head(hidden_states[:, slice_indices, :])
+
+        # Compute loss if labels are provided
+        loss = None
+        if labels is not None and logits is not None:
+            loss = self.loss_function(
+                logits=logits, labels=labels, vocab_size=self.config.text_config.vocab_size, **kwargs
+            )
+
+        # Return based on output mode
+        if self._output_mode == self.OUTPUT_MODE_EMBEDDING:
+            return Qwen3VLModelOutput(
+                last_hidden_state=hidden_states,
+                past_key_values=outputs.past_key_values,
+                attentions=outputs.attentions,
+                rope_deltas=outputs.rope_deltas,
+            )
+        elif self._output_mode == self.OUTPUT_MODE_LOGITS:
+            return logits
+        elif self._output_mode == self.OUTPUT_MODE_HIDDEN:
+            return Qwen3VLModelOutput(
+                last_hidden_state=hidden_states,
+                hidden_states=outputs.hidden_states,
+                past_key_values=outputs.past_key_values,
+                attentions=outputs.attentions,
+                rope_deltas=outputs.rope_deltas,
+            )
+        else:  # OUTPUT_MODE_FULL
+            return Qwen3VLCausalLMOutputWithPast(
+                loss=loss,
+                logits=logits,
+                past_key_values=outputs.past_key_values,
+                hidden_states=outputs.hidden_states,
+                attentions=outputs.attentions,
+                rope_deltas=outputs.rope_deltas,
+            )
+
+
+# ===========================================================================
+# Packing-aware forward patches (cu_seqlens) for the Qwen3-VL text encoder
+# ===========================================================================
+
+def model_forward(
+    self,
+    input_ids: torch.LongTensor | None = None,
+    attention_mask: torch.Tensor | None = None,
+    position_ids: torch.LongTensor | None = None,
+    past_key_values: Cache | None = None,
+    inputs_embeds: torch.FloatTensor | None = None,
+    use_cache: bool | None = None,
+    cache_position: torch.LongTensor | None = None,
+    # args for deepstack
+    visual_pos_masks: torch.Tensor | None = None,
+    deepstack_visual_embeds: list[torch.Tensor] | None = None,
+    **kwargs: Unpack[FlashAttentionKwargs],
+) -> tuple | BaseModelOutputWithPast:
+    r"""
+    visual_pos_masks (`torch.Tensor` of shape `(batch_size, seqlen)`, *optional*):
+        The mask of the visual positions.
+    deepstack_visual_embeds (`list[torch.Tensor]`, *optional*):
+        The deepstack visual embeddings. The shape is (num_layers, visual_seqlen, embed_dim).
+        The feature is extracted from the different visual encoder layers, and fed to the decoder
+        hidden states. It's from the paper DeepStack(https://arxiv.org/abs/2406.04334).
+    """
+    if (input_ids is None) ^ (inputs_embeds is not None):
+        raise ValueError("You must specify exactly one of input_ids or inputs_embeds")
+
+    # torch.jit.trace() doesn't support cache objects in the output
+    if use_cache and past_key_values is None and not torch.jit.is_tracing():
+        past_key_values = DynamicCache(config=self.config)
+
+    if inputs_embeds is None:
+        inputs_embeds = self.embed_tokens(input_ids)
+
+    if cache_position is None:
+        past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0
+        cache_position = torch.arange(
+            past_seen_tokens, past_seen_tokens + inputs_embeds.shape[1], device=inputs_embeds.device
+        )
+
+    # the hard coded `3` is for temporal, height and width.
+    if position_ids is None:
+        position_ids = cache_position.view(1, 1, -1).expand(3, inputs_embeds.shape[0], -1)
+    elif position_ids.ndim == 2:
+        position_ids = position_ids[None, ...].expand(3, position_ids.shape[0], -1)
+
+    if position_ids.ndim == 3 and position_ids.shape[0] == 4:
+        text_position_ids = position_ids[0]
+        position_ids = position_ids[1:]
+    else:
+        text_position_ids = position_ids[0]
+
+    if kwargs.get("cu_seqlens") is None:
+        attention_mask = create_causal_mask(
+            config=self.config,
+            input_embeds=inputs_embeds,
+            attention_mask=attention_mask,
+            cache_position=cache_position,
+            past_key_values=past_key_values,
+            position_ids=text_position_ids,
+        )
+
+    hidden_states = inputs_embeds
+
+    # create position embeddings to be shared across the decoder layers
+    position_embeddings = self.rotary_emb(hidden_states, position_ids)
+
+    # decoder layers
+    for layer_idx, decoder_layer in enumerate(self.layers):
+        layer_outputs = decoder_layer(
+            hidden_states,
+            attention_mask=attention_mask,
+            position_ids=text_position_ids,
+            past_key_values=past_key_values,
+            cache_position=cache_position,
+            position_embeddings=position_embeddings,
+            **kwargs,
+        )
+        hidden_states = layer_outputs
+
+        # add visual features to the hidden states of first several layers
+        if deepstack_visual_embeds is not None and layer_idx in range(len(deepstack_visual_embeds)):
+            hidden_states = self._deepstack_process(
+                hidden_states,
+                visual_pos_masks,
+                deepstack_visual_embeds[layer_idx],
+            )
+
+    hidden_states = self.norm(hidden_states)
+
+    return BaseModelOutputWithPast(
+        last_hidden_state=hidden_states,
+        past_key_values=past_key_values,
+    )
+
+
+def forward(
+    self,
+    hidden_states: torch.Tensor,
+    position_embeddings: tuple[torch.Tensor, torch.Tensor],
+    attention_mask: torch.Tensor | None,
+    past_key_values: Cache | None = None,
+    cache_position: torch.LongTensor | None = None,
+    **kwargs: Unpack[FlashAttentionKwargs],
+) -> tuple[torch.Tensor, torch.Tensor | None]:
+    input_shape = hidden_states.shape[:-1]
+    hidden_shape = (*input_shape, -1, self.head_dim)
+
+    query_states = self.q_norm(self.q_proj(hidden_states).view(hidden_shape)).transpose(1, 2)
+    key_states = self.k_norm(self.k_proj(hidden_states).view(hidden_shape)).transpose(1, 2)
+    value_states = self.v_proj(hidden_states).view(hidden_shape).transpose(1, 2)
+
+    cos, sin = position_embeddings
+    query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin)
+
+    if past_key_values is not None:
+        # sin and cos are specific to RoPE models; cache_position needed for the static cache
+        cache_kwargs = {"sin": sin, "cos": cos, "cache_position": cache_position}
+        key_states, value_states = past_key_values.update(key_states, value_states, self.layer_idx, cache_kwargs)
+
+    cu_seqlens = kwargs.get("cu_seqlens", None)
+
+    if cu_seqlens is None:
+        attention_interface: Callable = eager_attention_forward
+        if self.config._attn_implementation != "eager":
+            attention_interface = ALL_ATTENTION_FUNCTIONS[self.config._attn_implementation]
+
+        attn_output, attn_weights = attention_interface(
+            self,
+            query_states,
+            key_states,
+            value_states,
+            attention_mask,
+            dropout=0.0 if not self.training else self.attention_dropout,
+            scaling=self.scaling,
+            **kwargs,
+        )
+    else:
+        max_seqlen = torch.diff(cu_seqlens).max().item() if cu_seqlens is not None else None
+        query_states = query_states.transpose(1, 2).squeeze(0)
+        key_states = key_states.transpose(1, 2).squeeze(0)
+        value_states = value_states.transpose(1, 2).squeeze(0)
+        attn_output = flash_attn_varlen_func(
+            q=query_states,
+            k=key_states,
+            v=value_states,
+            cu_seqlens_q=cu_seqlens,
+            cu_seqlens_k=cu_seqlens,
+            max_seqlen_q=max_seqlen,
+            max_seqlen_k=max_seqlen,
+            causal=True,
+            window_size=(-1, -1),
+            softmax_scale=self.head_dim**-0.5,
+            dropout_p=0.0,
+        )
+
+    attn_output = attn_output.reshape(*input_shape, -1).contiguous()
+    attn_output = self.o_proj(attn_output)
+    return attn_output, None
+
+
+def qwen3_patch_forward():
+    """Patch the Qwen3-VL text model + attention forwards to support packed
+    varlen (cu_seqlens) inputs used by ``TextEncoder.forward``."""
+    from transformers.models.qwen3_vl.modeling_qwen3_vl import Qwen3VLTextAttention, Qwen3VLTextModel
+
+    Qwen3VLTextModel.forward = model_forward
+    Qwen3VLTextAttention.forward = forward
+
+
+# ===========================================================================
+# TextEncoder wrapper (packed text -> DiT conditioning embeddings)
+# ===========================================================================
+_FA2_ALIASES = {"flash2", "fa2", "flash_attention_2", "flash_attn_2"}
+_FA4_ALIASES = {"flash4", "fa4", "flash_attention_4", "flash_attn_4"}
+_SDPA_ALIASES = {"sdpa", "torch_sdpa", "scaled_dot_product_attention"}
+
+
+def _resolve_hf_attn_impl(attn_type: str) -> str:
+    """Map a project-level attn_type to a HuggingFace ``attn_implementation`` string.
+
+    ``VF_HF_ATTN_IMPL`` env var, if set, takes precedence (useful for forcing
+    sdpa on machines without flash-attn). For FA4 we additionally probe that
+    the CUTE-DSL kernel is importable and (when available) ask the HF helper
+    to confirm; if not, fall back to sdpa rather than crashing at load time.
+    """
+    override = os.environ.get("VF_HF_ATTN_IMPL")
+    if override:
+        return override
+
+    name = attn_type.lower().strip()
+    if name in _FA2_ALIASES:
+        return "flash_attention_2"
+    if name in _FA4_ALIASES:
+        try:
+            import flash_attn.cute  # noqa: F401
+            fa4_importable = True
+        except Exception:
+            fa4_importable = False
+        if fa4_importable:
+            try:
+                from transformers.utils.import_utils import is_flash_attn_4_available
+                if is_flash_attn_4_available():
+                    return "flash_attention_4"
+            except ImportError:
+                return "flash_attention_4"
+        logger.warning(
+            "attn_type=flash4 requested but flash_attn.cute is unavailable; "
+            "falling back to sdpa for HF text encoder."
+        )
+        return "sdpa"
+    if name in _SDPA_ALIASES:
+        return "sdpa"
+    raise ValueError(
+        f"Unknown attn_type {attn_type!r}; expected one of "
+        f"{sorted(_FA2_ALIASES | _FA4_ALIASES | _SDPA_ALIASES)}"
+    )
+
+
+SEQ_MULTI_OF = 32
+
+
+
+class TextEncoder(nn.Module):
+    def __init__(
+        self,
+        model_name: str,
+        version: str,
+        tokenizer_max_length: int,
+        prompt_template: dict | None,
+        dit_structure: dict,
+        use_packed_text_infer: bool = False,
+        attn_type: str = "flash2",
+        **hf_kwargs,
+    ):
+        super().__init__()
+        self.model_name = model_name
+        self.tokenizer_max_length = tokenizer_max_length
+        self.tokenizer: AutoTokenizer = AutoTokenizer.from_pretrained(version)
+        self.tokenizer.padding_side = "right"
+
+        hf_attn_impl = _resolve_hf_attn_impl(attn_type)
+        logger.info(f"TextEncoder attn_type={attn_type} -> attn_implementation={hf_attn_impl}")
+
+        logger.info("init vl model: qwen3")
+        self.hf_module: CustomQwen3VLForConditionalGeneration = CustomQwen3VLForConditionalGeneration.from_pretrained(
+            version,
+            attn_implementation=hf_attn_impl,
+            **hf_kwargs
+        )
+
+        # Use local_files_only if version is a local path (absolute path or contains path separators)
+        is_local = version.startswith("/") or os.sep in version
+        self.processor = AutoProcessor.from_pretrained(version, local_files_only=is_local)
+
+        self.hf_module = self.hf_module.eval().requires_grad_(False)
+
+        prompt_template = prompt_template or {}
+        self.prompt_template_encode = prompt_template.get("template", "")
+        self.prompt_template_encode_start_idx = prompt_template.get("start_idx", 0)
+        self.dit_structure = dit_structure
+        self.use_packed_text_infer = use_packed_text_infer
+
+    def forward(
+        self,
+        input_ids: torch.Tensor,
+        cu_seqlens: torch.Tensor,
+        inputs: dict | None = None,
+        drop_idx_override: int | None = None,
+    ):
+        """Encode packed text (varlen ``cu_seqlens``) into DiT conditioning embeddings.
+
+        This is the sole text-embedding path — both t2i and edit call it. Uses
+        Flash-Attention-2's varlen capability (``cu_seqlens``) via the patched
+        Qwen3-VL forward to process several concatenated sequences in a single
+        launch, with no padding. Verified numerically identical to a padded-batch
+        forward with per-sample cu_seqlens isolation (zero cross-contamination).
+
+        Args:
+            input_ids: Packed token ids ``[Total_L]``.
+            cu_seqlens: Cumulative sequence lengths ``[B+1]``.
+            inputs: Optional dict with additional model inputs (e.g. ``pixel_values``,
+                ``image_grid_thw`` for the multimodal edit path). Passed through to
+                the text encoder.
+            drop_idx_override: If set, override the number of leading (system-prompt)
+                tokens to drop per sequence. Use 0 for multi-turn where the system
+                prompt is embedded in the conversation and should not be stripped.
+
+        Returns:
+            dict with keys:
+                - ``txt``: text embeddings ``[Total_L - B*drop_idx, D]`` (system prompt dropped)
+                - ``vec``: pooled text embeddings ``[B, D]``
+                - ``txt_seq_lens``: per-sequence lengths ``[B]`` (after dropping system prompt)
+        """
+        # Compute seqlens from cu_seqlens
+        seqlens = cu_seqlens[1:] - cu_seqlens[:-1]
+        seqlens_list = seqlens.cpu().tolist()
+
+        # Build position_ids for packing: each sequence starts from 0
+        position_ids_list = []
+        for length in seqlens_list:
+            position_ids_list.append(torch.arange(length, device=input_ids.device))
+        position_ids = torch.cat(position_ids_list)  # [Total_L]
+
+        # Reshape for model input: [1, Total_L]
+        input_ids_packed = input_ids.unsqueeze(0)  # [1, Total_L]
+        position_ids_packed = position_ids.unsqueeze(0)  # [1, Total_L]
+
+        # Move to text encoder device
+        device = self.hf_module.device
+        input_ids_packed = input_ids_packed.to(device)
+        position_ids_packed = position_ids_packed.to(device)
+
+        # Get text embeddings (the text encoder is always frozen)
+        with torch.no_grad():
+            forward_kwargs = {
+                "input_ids": input_ids_packed,
+                "cu_seqlens": cu_seqlens,
+                "position_ids": position_ids_packed,
+                "output_hidden_states": False,
+                "max_seqlen": None,
+            }
+            # Pass multimodal inputs for edit mode (reference images)
+            if inputs is not None:
+                for key in ("pixel_values", "image_grid_thw"):
+                    if key in inputs and inputs[key] is not None:
+                        val = inputs[key]
+                        if hasattr(val, "to"):
+                            val = val.to(device)
+                        forward_kwargs[key] = val
+            outputs = self.hf_module(**forward_kwargs)
+
+        # Extract hidden state
+        if hasattr(outputs, "last_hidden_state") and outputs.last_hidden_state is not None:
+            hidden = outputs.last_hidden_state  # [1, Total_L, D]
+        elif hasattr(outputs, "hidden_states"):
+            hidden = outputs.hidden_states[-1]
+
+        # Remove batch dimension: [Total_L, D]
+        hidden = hidden.squeeze(0)
+
+        # Get drop_idx (system prompt length to skip).
+        # For multi-turn, drop_idx_override=0 is passed since system prompt is in the messages.
+        if drop_idx_override is not None:
+            drop_idx = drop_idx_override
+        else:
+            drop_idx = self.prompt_template_encode_start_idx
+
+        # Split hidden states by sequence
+        hidden_split = torch.split(hidden, seqlens_list, dim=0)
+
+        # Extract valid embeddings (drop system prompt) and compute vec
+        txt_list = []
+        vec_list = []
+        valid_lengths = []
+
+        for h in hidden_split:
+            # Drop system prompt tokens
+            h_valid = h[drop_idx:]  # [seq_len - drop_idx, D]
+            txt_list.append(h_valid)
+            valid_lengths.append(h_valid.shape[0])
+
+            # Compute pooled embedding (mean of valid tokens only, after dropping system prompt)
+            vec_list.append(h_valid.mean(dim=0))  # [D]
+
+        txt = torch.cat(txt_list, dim=0)  # [Total_valid, D]
+        vec = torch.stack(vec_list, dim=0)  # [B, D]
+        txt_seq_lens = torch.tensor(valid_lengths, device=input_ids.device)
+
+        result = {
+            "txt": txt,
+            "vec": vec,
+            "txt_seq_lens": txt_seq_lens,
+        }
+
+        return result
+
+    # ------------------------------------------------------------------
+    # Mandatory content-policy screening (same Qwen3-VL weights)
+    # ------------------------------------------------------------------
+    # The policy classifier lives HERE, on the text encoder, so it runs on the
+    # exact weights that produce the diffusion conditioning and is not a
+    # separable, toggleable pre-pass in the pipeline. The classifier needs
+    # autoregressive ``.generate()`` (JSON verdict) whereas conditioning is a
+    # single embedding forward — they cannot be one GPU forward without a
+    # trained classification head, so "fused" here means: same module, same
+    # weights, always run, FAIL-CLOSED (any error blocks).
+
+    def screen_text(self, prompt: str, max_new_tokens: int = 160):
+        """Classify a text-to-image ``prompt`` against the content policy.
+
+        Returns a ``FilterVerdict``. FAIL-CLOSED: any error (generation, parse)
+        returns ``violates=True`` so a broken classifier cannot be used as a
+        bypass. An empty prompt is not a violation.
+        """
+        from .mage_text import (
+            CONTENT_FILTER_SYSTEM, FilterVerdict, _extract_json_object,
+            _full_output_mode,
+        )
+
+        if not prompt or not prompt.strip():
+            return FilterVerdict(False, [], "empty prompt", "")
+        try:
+            tokenizer = self.tokenizer
+            hf = self.hf_module
+            device = next(hf.parameters()).device
+
+            messages = [
+                {"role": "system", "content": CONTENT_FILTER_SYSTEM},
+                {"role": "user", "content": f"Prompt to classify:\n{prompt}"},
+            ]
+            text = tokenizer.apply_chat_template(
+                messages, tokenize=False, add_generation_prompt=True)
+            inputs = tokenizer(text, return_tensors="pt").to(device)
+
+            eos_id = tokenizer.eos_token_id
+            pad_id = tokenizer.pad_token_id if tokenizer.pad_token_id is not None else eos_id
+
+            with _full_output_mode(hf), torch.no_grad():
+                out = hf.generate(
+                    **inputs, max_new_tokens=max_new_tokens, do_sample=False,
+                    pad_token_id=pad_id, eos_token_id=eos_id)
+            gen = tokenizer.decode(
+                out[0, inputs.input_ids.shape[1]:], skip_special_tokens=True).strip()
+
+            parsed = _extract_json_object(gen)
+            violates = bool(parsed.get("violates", False))
+            cats = [c for c in (parsed.get("categories", []) or []) if isinstance(c, str)]
+            reason = str(parsed.get("reason", "")).strip()
+            return FilterVerdict(violates, cats, reason, gen)
+        except Exception as exc:  # noqa: BLE001
+            # FAIL-CLOSED: block on any screening error.
+            return FilterVerdict(
+                True, ["policy"], f"filter error (blocked): {type(exc).__name__}: {exc}", "")
+
+    def screen_edit(self, prompt: str, ref_images, max_new_tokens: int = 192):
+        """Classify an image-EDIT request (source image(s) + instruction).
+
+        Considers BOTH the source image(s) and the instruction via multimodal
+        Qwen3-VL. Falls back to :meth:`screen_text` when no image is given.
+        FAIL-CLOSED: any error returns ``violates=True``.
+        """
+        from PIL import Image
+
+        from .mage_text import (
+            CONTENT_FILTER_EDIT_SYSTEM, FilterVerdict, _extract_json_object,
+            _full_output_mode,
+        )
+
+        pils = [ref_images] if isinstance(ref_images, Image.Image) else list(ref_images)
+        pils = [p.convert("RGB") for p in pils if p is not None]
+        if not pils:
+            return self.screen_text(prompt, max_new_tokens=max_new_tokens)
+
+        instruction = (prompt or "").strip() or "(no textual instruction)"
+        try:
+            processor = self.processor
+            tokenizer = self.tokenizer
+            hf = self.hf_module
+            device = next(hf.parameters()).device
+
+            user_content = [{"type": "image"} for _ in pils]
+            user_content.append({
+                "type": "text",
+                "text": (
+                    f"There {'is' if len(pils) == 1 else 'are'} {len(pils)} source "
+                    f"image(s) above. Edit instruction: {instruction}\n"
+                    "Classify this edit request."
+                ),
+            })
+            messages = [
+                {"role": "system", "content": CONTENT_FILTER_EDIT_SYSTEM},
+                {"role": "user", "content": user_content},
+            ]
+            text = processor.apply_chat_template(
+                messages, tokenize=False, add_generation_prompt=True)
+            inputs = processor(
+                text=[text], images=pils, padding=True, return_tensors="pt").to(device)
+
+            eos_id = tokenizer.eos_token_id
+            pad_id = tokenizer.pad_token_id if tokenizer.pad_token_id is not None else eos_id
+
+            # Keep only the kwargs Qwen3-VL .generate() consumes.
+            gen_inputs = {
+                k: inputs[k]
+                for k in ("input_ids", "attention_mask", "pixel_values", "image_grid_thw")
+                if k in inputs and inputs[k] is not None
+            }
+            input_len = gen_inputs["input_ids"].shape[1]
+
+            with _full_output_mode(hf), torch.no_grad():
+                out = hf.generate(
+                    **gen_inputs, max_new_tokens=max_new_tokens, do_sample=False,
+                    pad_token_id=pad_id, eos_token_id=eos_id)
+            gen = tokenizer.decode(out[0, input_len:], skip_special_tokens=True).strip()
+
+            parsed = _extract_json_object(gen)
+            violates = bool(parsed.get("violates", False))
+            cats = [c for c in (parsed.get("categories", []) or []) if isinstance(c, str)]
+            reason = str(parsed.get("reason", "")).strip()
+            return FilterVerdict(violates, cats, reason, gen)
+        except Exception as exc:  # noqa: BLE001
+            # FAIL-CLOSED: block on any screening error.
+            return FilterVerdict(
+                True, ["policy"], f"edit filter error (blocked): {type(exc).__name__}: {exc}", "")
+

vendor/mage_flow/models/utils.py

@@ -0,0 +1,175 @@
+import json
+import math
+import os
+
+import torch
+from einops import rearrange
+from loguru import logger
+from safetensors.torch import load_file
+from safetensors.torch import load_file as load_sft
+from torch import Tensor
+
+from .mage_flow import MageFlow, MageFlowParams
+
+
+def get_noise(
+    num_samples: int,
+    channel: int,
+    height: int,
+    width: int,
+    device: torch.device,
+    dtype: torch.dtype,
+    seed: int,
+):
+    # MageVAE: 16x downsample, no patch packing
+    return torch.randn(
+        num_samples,
+        channel,
+        math.ceil(height / 16),
+        math.ceil(width / 16),
+        device=device,
+        dtype=dtype,
+        generator=torch.Generator(device=device).manual_seed(seed),
+    )
+
+
+def unpack(x: Tensor, height: int, width: int) -> Tensor:
+    # MageVAE: [B, H*W, C] -> [B, C, H, W], no patch unpacking
+    return rearrange(
+        x,
+        "b (h w) c -> b c h w",
+        h=math.ceil(height / 16),
+        w=math.ceil(width / 16),
+    )
+
+
+PROMPT_TEMPLATE = {
+    "default": {"template": "{}", "start_idx": 0},
+    "default-nonthinking": {"template": "{}<think>\n\n</think>\n\n", "start_idx": 0},
+    "mage-flow": {
+        "template": (
+            "<|im_start|>system\nDescribe the image by detailing the color, shape, size, texture, quantity, "
+            "text, spatial relationships of the objects and background:"
+            "<|im_end|>\n<|im_start|>user\n{}<|im_end|>\n<|im_start|>assistant\n"
+        ),
+        "start_idx": 34,
+    },
+    "mage-flow-edit": {
+        "template": (
+            "<|im_start|>system\nDescribe the key features of the input image (color, shape, size, texture,"
+            " objects, background), then explain how the user's text instruction should alter or modify the image. "
+            "Generate a new image that meets the user's requirements while maintaining consistency with the original "
+            "input where appropriate.<|im_end|>\n<|im_start|>user\n{}<|im_end|>\n<|im_start|>assistant\n"
+        ),
+        "start_idx": 64,
+    },
+}
+
+
+def print_load_warning(missing: list[str], unexpected: list[str]) -> None:
+    if len(missing) > 0 and len(unexpected) > 0:
+        logger.warning(f"Got {len(missing)} missing keys:\n\t" + "\n\t".join(missing))
+        logger.warning("\n" + "-" * 79 + "\n")
+        logger.warning(f"Got {len(unexpected)} unexpected keys:\n\t" + "\n\t".join(unexpected))
+    elif len(missing) > 0:
+        logger.warning(f"Got {len(missing)} missing keys:\n\t" + "\n\t".join(missing))
+    elif len(unexpected) > 0:
+        logger.warning(f"Got {len(unexpected)} unexpected keys:\n\t" + "\n\t".join(unexpected))
+
+
+def correct_model_weight(state_dict):
+    result = {}
+    for key in state_dict.keys():
+        if "_orig_mod." in key:
+            result[key[10:]] = state_dict[key]
+        else:
+            result[key] = state_dict[key]
+    return result
+
+
+def load_hf_style_weight(pretrain_path, device):
+    index_path = os.path.join(pretrain_path, "diffusion_pytorch_model.safetensors.index.json")
+
+    with open(index_path) as f:
+        index = json.load(f)
+
+    weight_map = index["weight_map"]
+
+    sd = {}
+    loaded_shards = set()
+
+    for shard_file in weight_map.values():
+        if shard_file in loaded_shards:
+            continue
+        shard_path = os.path.join(pretrain_path, shard_file)
+        shard_sd = load_file(shard_path, device="cpu")
+        sd.update(shard_sd)
+        loaded_shards.add(shard_file)
+
+    return sd
+
+
+def load_model_weight(model, pretrain_path, device="cpu"):
+    if os.path.exists(pretrain_path):
+        logger.info(f"Loading checkpoint from {pretrain_path}")
+        try:
+            if pretrain_path.endswith("safetensors"):
+                sd = load_sft(pretrain_path, device="cpu")
+            elif os.path.exists(os.path.join(pretrain_path, "diffusion_pytorch_model.safetensors.index.json")):
+                sd = load_hf_style_weight(pretrain_path, device)
+            else:
+                sd = torch.load(pretrain_path, map_location="cpu")
+
+            sd = correct_model_weight(sd)
+            sd = optionally_expand_state_dict(model, sd)
+            missing, unexpected = model.load_state_dict(sd, strict=False, assign=True)
+            print_load_warning(missing, unexpected)
+            return True
+        except Exception as e:
+            logger.info(f"CANNOT Load {pretrain_path}, because {e}")
+            return False
+    return False
+
+
+def load_model(dit_structure: dict, pretrain_path: str | None = None):
+    logger.info("Init DiT model")
+
+    # If name is a dict, we assume it contains the parameters directly
+    # We need to determine the model class based on some heuristic or just default to MageFlow/Flux
+    # For now, let's assume it's MageFlow if time_type is present, or check other fields
+    params = MageFlowParams(**dit_structure)
+    # Default to MageFlow for now as per user context, or we could add a 'model_type' field to the dict
+    # The user mentioned "model structure option", implying we are configuring the structure.
+    # Let's assume MageFlow for this refactor as the user was using qwen-image-tiny-wo-textemb
+    model = MageFlow(params)
+
+    # logger.info(f"Loading {name if isinstance(name, str) else 'custom config'} checkpoint from {pretrain_path}")
+    if pretrain_path is not None:
+        load_model_weight(model, pretrain_path, device="cpu")
+    # if isinstance(name, str) and configs[name].lora_path is not None:
+    #     logger.info("Loading LoRA")
+    #     lora_sd = load_sft(configs[name].lora_path, device="cpu")
+    #     # loading the lora params + overwriting scale values in the norms
+    #     missing, unexpected = model.load_state_dict(lora_sd, strict=False, assign=True)
+    #     print_load_warning(missing, unexpected)
+    return model
+
+
+def optionally_expand_state_dict(model: torch.nn.Module, state_dict: dict) -> dict:
+    """
+    Optionally expand the state dict to match the model's parameters shapes.
+    """
+    for name, param in model.named_parameters():
+        if name in state_dict:
+            if state_dict[name].shape != param.shape:
+                logger.info(
+                    f"Expanding '{name}' with shape {state_dict[name].shape} to model parameter with shape "
+                    f"{param.shape}."
+                )
+                # expand with zeros:
+                expanded_state_dict_weight = torch.zeros_like(param, device=state_dict[name].device)
+                slices = tuple(slice(0, dim) for dim in state_dict[name].shape)
+                expanded_state_dict_weight[slices] = state_dict[name]
+                state_dict[name] = expanded_state_dict_weight
+
+    return state_dict

vendor/mage_flow/pipeline.py

@@ -0,0 +1,762 @@
+"""MageFlow text-to-image + image-edit inference pipeline.
+
+Self-contained MageFlow t2i / edit inference: load a HuggingFace diffusers-style
+repo (model_index.json + transformer/ vae/ scheduler/), then generate or edit
+images. No training/eval deps.
+
+Both ``generate_images`` and ``generate_edits`` support PACKED multi-resolution
+inference: several samples (each at its own resolution) are concatenated into a
+single varlen sequence and processed in one transformer forward per denoise
+step. Per-sample ``cu_seqlens`` (inside the flash-attn varlen kernel) isolate
+samples, exactly mirroring training-time packing. These packed functions are the
+sole implementation — the single-image case is just a pack of size 1, exposed
+via the ``MageFlowPipeline.generate`` / ``.edit`` convenience methods.
+"""
+
+from __future__ import annotations
+
+import json
+import os
+import random
+
+import torch
+from einops import rearrange
+from PIL import Image
+
+from diffusers import FlowMatchEulerDiscreteScheduler
+
+from .models.mage_flow import MageFlowModel, ModelConfig
+from .models.utils import PROMPT_TEMPLATE, get_noise, unpack
+from .models.modules.mage_text import make_refusal_image
+from .models.modules.mage_latent import encode_noise, resolve_gs_key
+
+
+# ---------------------------------------------------------------------------
+# Scheduler — diffusers FlowMatchEulerDiscreteScheduler
+# ---------------------------------------------------------------------------
+def build_scheduler(num_steps: int, device=None, shift: float = 6.0):
+    """Construct a diffusers ``FlowMatchEulerDiscreteScheduler`` whose sigma
+    schedule reproduces our default preset exactly.
+
+    The base sigmas ``linspace(1, 1/num_steps, num_steps)`` fed to
+    ``set_timesteps`` are run through the scheduler's built-in static shift
+    ``shift·s/(1+(shift-1)·s)`` and a terminal 0 is appended — the static-shift 
+    schedule (the only supported schedule).
+    """
+    scheduler = FlowMatchEulerDiscreteScheduler(
+        num_train_timesteps=1000, shift=shift, use_dynamic_shifting=False)
+    base_sigmas = torch.linspace(1.0, 1.0 / num_steps, num_steps).tolist()
+    scheduler.set_timesteps(sigmas=base_sigmas, device=device)
+    return scheduler
+
+
+def _get_scheduler(model, steps, device, static_shift):
+    scheduler = getattr(model, "scheduler", None)
+    if scheduler is None:
+        return build_scheduler(steps, device=device,
+                               shift=(static_shift if static_shift is not None else 6.0))
+    if static_shift is not None:
+        scheduler.set_shift(static_shift)
+    scheduler.set_timesteps(sigmas=torch.linspace(1.0, 1.0 / steps, steps).tolist(), device=device)
+    return scheduler
+
+
+# ---------------------------------------------------------------------------
+# Small helpers
+# ---------------------------------------------------------------------------
+def _template_info(name: str | None) -> dict:
+    name = name or "mage-flow"
+    if name not in PROMPT_TEMPLATE:
+        raise ValueError(f"Unknown prompt template: {name}")
+    return PROMPT_TEMPLATE[name]
+
+
+def _as_list(val, default, n):
+    """Broadcast a scalar/None to a length-n list, or validate a given list."""
+    if val is None:
+        return [default] * n
+    if isinstance(val, (list, tuple)):
+        if len(val) != n:
+            raise ValueError(f"expected {n} values, got {len(val)}")
+        return list(val)
+    return [val] * n
+
+
+def _lens_to_cu(lens, device):
+    """Sequence lengths -> cumulative cu_seqlens [0, l0, l0+l1, ...] (int32)."""
+    t = torch.tensor(lens, device=device, dtype=torch.int32)
+    return torch.cat([torch.zeros(1, dtype=torch.int32, device=device),
+                      torch.cumsum(t, dim=0, dtype=torch.int32)])
+
+
+def _make_divisible_by_16(size: int) -> int:
+    return max(16, 16 * (size // 16))
+
+
+def _compute_aspect_ratio_size(pil_img: Image.Image, max_size: int):
+    """Longest side = ``max_size``, short side from aspect ratio, both /16."""
+    w, h = pil_img.size
+    if h >= w:
+        new_h, new_w = max_size, int(round(w * max_size / h))
+    else:
+        new_w, new_h = max_size, int(round(h * max_size / w))
+    return _make_divisible_by_16(new_h), _make_divisible_by_16(new_w)
+
+
+def _edit_target_size(pil_img: Image.Image, max_size, height, width):
+    """Output (H, W) for an edit sample, derived from its PRIMARY reference.
+
+    Precedence: explicit ``height`` AND ``width`` (custom size) > ``max_size``
+    (longest side, short side by aspect ratio) > the source image's own size.
+    All rounded down to a multiple of 16.
+    """
+    if height and width:
+        return _make_divisible_by_16(height), _make_divisible_by_16(width)
+    if max_size:
+        return _compute_aspect_ratio_size(pil_img, max_size)
+    # Nothing specified: keep the source resolution (its own longest side).
+    return _compute_aspect_ratio_size(pil_img, max(pil_img.size))
+
+
+def _decode_one(model, tokens, height, width, dev):
+    """Unpack one sample's image tokens [1, H*W, C] and VAE-decode to a PIL image."""
+    with torch.autocast(device_type=dev.type, dtype=torch.bfloat16):
+        out = model.vae.decode(unpack(tokens.float(), height, width))
+    out = rearrange(out.clamp(-1, 1), "b c h w -> b h w c")
+    out = (127.5 * (out + 1.0)).cpu().byte().numpy()
+    return Image.fromarray(out[0])
+
+
+def _build_pack_ctx(img_ids, img_cu, img_shapes, img_lens, txt, txt_cu, txt_mask, vec,
+                    neg_txt, neg_cu, neg_mask, neg_vec, cfg, renormalization, batch_cfg, device):
+    """Precompute the static per-step transformer inputs for a packed batch.
+
+    When a negative branch is present and ``batch_cfg`` is True, the conditional
+    and unconditional passes are fused into ONE varlen forward: the image tokens
+    are duplicated (cond copy + uncond copy) and the positive/negative texts are
+    concatenated, so cond sample i and uncond sample i become two independent
+    varlen segments processed in a single kernel launch. flash_attn_varlen_func
+    keeps every segment isolated via cu_seqlens, so this is numerically identical
+    to two separate forwards — just one launch instead of two.
+    """
+    na = len(img_lens)
+    ctx = {
+        "na": na, "cfg": cfg, "renorm": renormalization, "batch_cfg": batch_cfg,
+        "has_neg": neg_txt is not None,
+        "img_ids": img_ids, "img_cu": img_cu, "img_shapes": img_shapes,
+        "img_max": int(max(img_lens)),
+        "txt": txt, "txt_ids": torch.zeros(1, txt.shape[1], 3, device=device),
+        "txt_cu": txt_cu, "txt_mask": txt_mask, "vec": vec,
+        "txt_max": int((txt_cu[1:] - txt_cu[:-1]).max().item()),
+    }
+    if neg_txt is None:
+        return ctx
+    ctx.update({
+        "neg_txt": neg_txt, "neg_ids": torch.zeros(1, neg_txt.shape[1], 3, device=device),
+        "neg_cu": neg_cu, "neg_mask": neg_mask, "neg_vec": neg_vec,
+        "neg_max": int((neg_cu[1:] - neg_cu[:-1]).max().item()),
+    })
+    if batch_cfg:
+        # Duplicate image segments (cond then uncond) and concat pos+neg text.
+        d_txt = torch.cat([txt, neg_txt], dim=1)
+        pos_lens = (txt_cu[1:] - txt_cu[:-1]).tolist()
+        neg_lens = (neg_cu[1:] - neg_cu[:-1]).tolist()
+        ctx.update({
+            "d_img_ids": torch.cat([img_ids, img_ids], dim=1),
+            "d_img_cu": _lens_to_cu(list(img_lens) + list(img_lens), device),
+            "d_img_shapes": [img_shapes[0] + img_shapes[0]],
+            "d_txt": d_txt,
+            "d_txt_ids": torch.zeros(1, d_txt.shape[1], 3, device=device),
+            "d_txt_cu": _lens_to_cu(pos_lens + neg_lens, device),
+            "d_txt_mask": torch.ones(1, d_txt.shape[1], device=device),
+            "d_vec": torch.cat([vec, neg_vec], dim=0),
+            "d_txt_max": int(max(pos_lens + neg_lens)),
+        })
+    return ctx
+
+
+def _velocity(transformer, img, ctx, sigma):
+    """CFG-combined image-token velocity for a packed batch at noise level ``sigma``.
+
+    Returns [1, sum_img_len, C] in the conditional sample order. When
+    ``batch_cfg`` is set the cond+uncond passes share a single fused varlen
+    forward; otherwise they are two forwards.
+    """
+    dev = img.device
+    na = ctx["na"]
+
+    def _fwd(x, n, img_ids, img_cu, img_max, img_shapes, txt, txt_ids, txt_cu, txt_mask, txt_max, vec):
+        t_vec = torch.full((n,), sigma, dtype=x.dtype, device=dev)
+        return transformer(img=x, txt=txt, timesteps=t_vec, img_shapes=img_shapes,
+                           img_cu_seqlens=img_cu, txt_cu_seqlens=txt_cu)
+
+    if not ctx["has_neg"]:
+        return _fwd(img, na, ctx["img_ids"], ctx["img_cu"], ctx["img_max"], ctx["img_shapes"],
+                    ctx["txt"], ctx["txt_ids"], ctx["txt_cu"], ctx["txt_mask"], ctx["txt_max"], ctx["vec"])
+
+    if ctx["batch_cfg"]:
+        n_img = img.shape[1]
+        out = _fwd(torch.cat([img, img], dim=1), 2 * na,
+                   ctx["d_img_ids"], ctx["d_img_cu"], ctx["img_max"], ctx["d_img_shapes"],
+                   ctx["d_txt"], ctx["d_txt_ids"], ctx["d_txt_cu"], ctx["d_txt_mask"], ctx["d_txt_max"], ctx["d_vec"])
+        cond, unc = out[:, :n_img, :], out[:, n_img:, :]
+    else:
+        cond = _fwd(img, na, ctx["img_ids"], ctx["img_cu"], ctx["img_max"], ctx["img_shapes"],
+                    ctx["txt"], ctx["txt_ids"], ctx["txt_cu"], ctx["txt_mask"], ctx["txt_max"], ctx["vec"])
+        unc = _fwd(img, na, ctx["img_ids"], ctx["img_cu"], ctx["img_max"], ctx["img_shapes"],
+                   ctx["neg_txt"], ctx["neg_ids"], ctx["neg_cu"], ctx["neg_mask"], ctx["neg_max"], ctx["neg_vec"])
+
+    cfg = ctx["cfg"]
+    if ctx["renorm"]:
+        # CFG renormalization: rescale the guided velocity per token back to the
+        # conditional velocity's norm (reduces oversaturation at high cfg).
+        comb = unc + cfg * (cond - unc)
+        return comb * (torch.norm(cond, dim=-1, keepdim=True) /
+                       (torch.norm(comb, dim=-1, keepdim=True) + 1e-6))
+    return unc + cfg * (cond - unc)
+
+
+def _encode_texts_packed(model, prompts, template, drop_idx, device):
+    """Encode a LIST of templated text-only prompts in ONE packed varlen forward
+    (``TextEncoder.forward`` — varlen cu_seqlens isolates each prompt,
+    verified zero cross-contamination). Returns (txt_flat [ΣLi, D], vec [N, D],
+    per-prompt token lengths list)."""
+    tokenizer = model.txt_enc.tokenizer
+    max_len = model.txt_enc.tokenizer_max_length + drop_idx
+    ids_list = [
+        tokenizer(template.format(p), max_length=max_len, truncation=True,
+                  return_tensors="pt").input_ids.squeeze(0)
+        for p in prompts
+    ]
+    input_ids = torch.cat(ids_list).to(device)
+    cu_seqlens = _lens_to_cu([int(t.numel()) for t in ids_list], device)
+    res = model.txt_enc(
+        input_ids, cu_seqlens, drop_idx_override=drop_idx)
+    return res["txt"], res["vec"], res["txt_seq_lens"].tolist()
+
+
+def _slice_packed(txt_flat, vec, lens, start, count, device):
+    """Format a contiguous ``count``-prompt slice (starting at prompt ``start``) of a
+    packed text encode into the (txt [1, ΣL, D], cu_seqlens, ones-mask, vec [count, D])
+    tuple that ``_build_pack_ctx`` consumes."""
+    seg_lens = lens[start:start + count]
+    tok_start = sum(lens[:start])
+    tok_end = tok_start + sum(seg_lens)
+    txt = txt_flat[tok_start:tok_end].reshape(1, -1, txt_flat.shape[-1]).to(device)
+    return (txt, _lens_to_cu(seg_lens, device),
+            torch.ones(1, txt.shape[1], device=device), vec[start:start + count].to(device))
+
+
+# ---------------------------------------------------------------------------
+# Text-to-image (packed, multi-resolution)
+# ---------------------------------------------------------------------------
+@torch.no_grad()
+def generate_images(model, prompts, neg_prompts=None, seeds=None, steps=30, cfg=5.0,
+                    heights=None, widths=None, device="cuda",
+                    prompt_template="mage-flow", static_shift=None,
+                    gs_key=None,
+                    renormalization=False, batch_cfg=True):
+    """Generate one image per prompt. Prompts may request DIFFERENT resolutions;
+    all are packed into a single varlen forward per denoise step — samples are
+    kept isolated by ``flash_attn_varlen_func`` via per-sample ``cu_seqlens`` (no
+    cross-sample attention), mirroring training-time packing. When ``cfg > 1`` and
+    ``batch_cfg`` is set, the positive and negative passes are fused into that
+    same varlen forward. Returns a list of PIL images aligned with ``prompts``.
+    """
+    if isinstance(prompts, str):
+        prompts = [prompts]
+    n = len(prompts)
+    neg_prompts = _as_list(neg_prompts, " ", n)
+    seeds = _as_list(seeds, 42, n)
+    heights = _as_list(heights, 1024, n)
+    widths = _as_list(widths, 1024, n)
+    info = _template_info(prompt_template)
+    template = info.get("template", "{}")
+    drop_idx = int(info.get("start_idx", 0))
+    dev = torch.device(device)
+
+    # Content-policy gate per sample (MANDATORY — runs on the same text-encoder
+    # weights as conditioning, no opt-out). Violating prompts get a refusal
+    # placeholder and are dropped from the pack.
+    results = [None] * n
+    active = []
+    for i in range(n):
+        if seeds[i] == -1:
+            seeds[i] = random.randint(0, 2**32 - 1)
+        verdict = model.txt_enc.screen_text(prompts[i])
+        if verdict.violates:
+            h_, w_ = _make_divisible_by_16(heights[i]), _make_divisible_by_16(widths[i])
+            print(verdict.banner())
+            results[i] = make_refusal_image(verdict, height=h_, width=w_)
+            continue
+        active.append(i)
+    if not active:
+        return results
+
+    gs_key_int = resolve_gs_key(gs_key)
+    # Per-sample noise tokens + position ids + shapes (MageVAE: flatten, no packing).
+    ch = model.vae.latent_channels
+    img_list, ids_list, lens, shapes, hw = [], [], [], [], []
+    for i in active:
+        h_, w_ = _make_divisible_by_16(heights[i]), _make_divisible_by_16(widths[i])
+        torch.manual_seed(seeds[i])
+        x = get_noise(num_samples=1, channel=ch, height=h_, width=w_,
+                      device=dev, dtype=torch.bfloat16, seed=seeds[i])
+        # Distribution-preserving watermark in the initial noise (same shape,
+        # still ~N(0,1)); detect by inverting the flow ODE back to noise.
+        x = encode_noise(tuple(x.shape[1:]), key=gs_key_int,
+                         seed=seeds[i], device=dev, dtype=torch.bfloat16)
+        _, _, gh, gw = x.shape
+        img_list.append(rearrange(x, "b c h w -> b (h w) c")[0])
+        ids = torch.zeros(gh, gw, 3, device=dev)
+        ids[..., 1] = ids[..., 1] + torch.arange(gh, device=dev)[:, None]
+        ids[..., 2] = ids[..., 2] + torch.arange(gw, device=dev)[None, :]
+        ids_list.append(rearrange(ids, "h w c -> (h w) c"))
+        lens.append(gh * gw); shapes.append((1, gh, gw)); hw.append((h_, w_))
+    img = torch.cat(img_list, 0).unsqueeze(0)
+    img_ids = torch.cat(ids_list, 0).unsqueeze(0)
+    img_cu = _lens_to_cu(lens, dev)
+    img_shapes = [shapes]
+
+    # Packed text: positive prompts AND (for CFG) negative prompts are encoded
+    # TOGETHER in ONE varlen forward, then split back — cu_seqlens keeps every
+    # prompt isolated (verified zero cross-contamination).
+    pos_prompts = [prompts[i] for i in active]
+    na = len(active)
+    use_neg = cfg > 1.0 and any(neg_prompts[i] for i in active)
+    if use_neg:
+        neg_list = [neg_prompts[i] or " " for i in active]
+        txt_flat, vec_all, lens_t = _encode_texts_packed(
+            model, pos_prompts + neg_list, template, drop_idx, dev)
+        txt, txt_cu, txt_mask, vec = _slice_packed(txt_flat, vec_all, lens_t, 0, na, dev)
+        neg_txt, neg_cu, neg_mask, neg_vec = _slice_packed(txt_flat, vec_all, lens_t, na, na, dev)
+    else:
+        txt_flat, vec_all, lens_t = _encode_texts_packed(model, pos_prompts, template, drop_idx, dev)
+        txt, txt_cu, txt_mask, vec = _slice_packed(txt_flat, vec_all, lens_t, 0, na, dev)
+        neg_txt = neg_cu = neg_mask = neg_vec = None
+
+    ctx = _build_pack_ctx(img_ids, img_cu, img_shapes, lens, txt, txt_cu, txt_mask, vec,
+                          neg_txt, neg_cu, neg_mask, neg_vec, cfg, renormalization, batch_cfg, dev)
+    scheduler = _get_scheduler(model, steps, device, static_shift)
+    for si, t in enumerate(scheduler.timesteps):
+        pred = _velocity(model.transformer, img, ctx, scheduler.sigmas[si].item())
+        img = scheduler.step(pred, t, img, return_dict=False)[0]
+
+    off = 0
+    for k, i in enumerate(active):
+        L = lens[k]
+        h_, w_ = hw[k]
+        results[i] = _decode_one(model, img[:, off:off + L, :], h_, w_, dev)
+        off += L
+    return results
+
+
+# ---------------------------------------------------------------------------
+# Image edit (packed, multi-resolution)
+# ---------------------------------------------------------------------------
+def _preprocess_ref_image(pil_img: Image.Image, height: int, width: int, device) -> torch.Tensor:
+    """Resize an RGB reference image to (height, width) and normalize to [-1, 1]."""
+    from torchvision.transforms import functional as TF
+    img = pil_img.convert("RGB")
+    img = TF.resize(img, [height, width], interpolation=TF.InterpolationMode.BICUBIC)
+    t = TF.to_tensor(img)  # [3, H, W] in [0, 1]
+    t = TF.normalize(t, [0.5, 0.5, 0.5], [0.5, 0.5, 0.5])  # -> [-1, 1]
+    return t.to(device)
+
+
+def _resize_long_edge(image: Image.Image, max_long_edge: int | None) -> Image.Image:
+    """Cap the VL conditioning image's long edge, preserving aspect ratio.
+
+    Matches training's data.processor._resize_long_edge (BICUBIC). Without this,
+    inference feeds a full-resolution image to the Qwen-VL processor whose
+    default max_pixels is far larger than 384**2 — a train/test mismatch.
+    """
+    if max_long_edge is None or max_long_edge <= 0:
+        return image
+    w, h = image.size
+    long_edge = max(w, h)
+    if long_edge <= max_long_edge:
+        return image
+    scale = max_long_edge / long_edge
+    new_w = max(1, int(round(w * scale)))
+    new_h = max(1, int(round(h * scale)))
+    return image.resize((new_w, new_h), Image.BICUBIC)
+
+
+# Fixed image placeholder used at edit training time (one per reference image).
+_EDIT_IMAGE_PLACEHOLDER = "<|vision_start|><|image_pad|><|vision_end|>"
+
+
+def _edit_prompt_body(instruction, num_refs):
+    """Training-time multi-reference prompt body: ``Image 1: <ph>Image 2: <ph>…{instruction}``."""
+    prefix = "".join(f"Image {j}: {_EDIT_IMAGE_PLACEHOLDER}" for j in range(1, num_refs + 1))
+    return prefix + instruction
+
+
+def _encode_edits_packed(model, ref_pils_per_sample, instructions, template, drop_idx, device):
+    """Encode ALL image-conditioned edit instructions in ONE packed multimodal
+    varlen forward (pixel_values/image_grid_thw concatenated across samples,
+    cu_seqlens isolates each). Returns (txt_flat [ΣLi, D], vec [N, D], per-sample lens)."""
+    processor = model.txt_enc.processor
+    ids_list, pv_list, thw_list = [], [], []
+    for ref_pils, instr in zip(ref_pils_per_sample, instructions, strict=False):
+        formatted = template.format(_edit_prompt_body(instr, len(ref_pils)))
+        vl = processor(text=[formatted], images=list(ref_pils), padding=True, return_tensors="pt")
+        vl = {k: (v.to(device) if hasattr(v, "to") else v) for k, v in vl.items()}
+        ids_list.append(vl["input_ids"].squeeze(0))
+        if vl.get("pixel_values") is not None:
+            pv_list.append(vl["pixel_values"]); thw_list.append(vl["image_grid_thw"])
+    input_ids = torch.cat(ids_list).to(device)
+    cu = _lens_to_cu([int(t.numel()) for t in ids_list], device)
+    inputs = {"input_ids": input_ids, "cu_seqlens": cu}
+    if pv_list:
+        inputs["pixel_values"] = torch.cat(pv_list, dim=0)
+        inputs["image_grid_thw"] = torch.cat(thw_list, dim=0)
+    res = model.txt_enc(
+        input_ids, cu, inputs=inputs, drop_idx_override=drop_idx)
+    return res["txt"], res["vec"], res["txt_seq_lens"].tolist()
+
+
+@torch.no_grad()
+def generate_edits(model, prompts, ref_images, neg_prompts=None, seeds=None, steps=30, cfg=5.0,
+                   max_size=None, heights=None, widths=None, device="cuda",
+                   prompt_template="mage-flow-edit", static_shift=None,
+                   gs_key=None,
+                   vl_cond_long_edge=384,
+                   renormalization=False, batch_cfg=True):
+    """Edit reference image(s) per prompt. Each ``ref_images[i]`` may be a single
+    image/path OR a list of source images (multi-image edit, like training —
+    trained with up to 3, but more are accepted) — all produce ONE edited output. Each sample's
+    ``[target, ref_1, …, ref_N]`` latent tokens are sequence-concatenated, and
+    all samples are packed into one varlen forward per denoise step.
+
+    Output resolution (derived from the first/primary reference of each sample):
+    if both ``heights[i]`` and ``widths[i]`` are given, use them; else if
+    ``max_size`` is given, the longest side is ``max_size`` and the short side
+    follows the reference's aspect ratio; otherwise the output keeps the source
+    image's own resolution. All references are VAE-encoded at that target size.
+    Returns a list of PIL images.
+    """
+    if isinstance(prompts, str):
+        prompts = [prompts]
+        ref_images = [ref_images]
+    n = len(prompts)
+    neg_prompts = _as_list(neg_prompts, " ", n)
+    seeds = _as_list(seeds, 42, n)
+    heights = _as_list(heights, None, n)
+    widths = _as_list(widths, None, n)
+    info = _template_info(prompt_template)
+    template = info.get("template", "{}")
+    drop_idx = int(info.get("start_idx", 0))
+    dev = torch.device(device)
+
+    # Normalize each sample's references to a list of 1..3 PIL images.
+    def _load_pil(r):
+        if isinstance(r, str):
+            r = Image.open(r)
+        return r.convert("RGB")
+
+    pils_per_sample = []
+    for r in ref_images:
+        refs = list(r) if isinstance(r, (list, tuple)) else [r]
+        if not refs:
+            raise ValueError("each edit sample needs at least one reference image")
+        pils_per_sample.append([_load_pil(x) for x in refs])
+
+    # Per-sample output resolution (from the first/primary reference) + content gate.
+    results = [None] * n
+    res_hw = [None] * n
+    active = []
+    for i in range(n):
+        res_hw[i] = _edit_target_size(pils_per_sample[i][0], max_size, heights[i], widths[i])
+        if seeds[i] == -1:
+            seeds[i] = random.randint(0, 2**32 - 1)
+        # Multimodal gate (MANDATORY): inspect the source image(s) AND the
+        # instruction, so NSFW / copyrighted-character / real-public-figure
+        # source photos are blocked even under an innocuous instruction.
+        verdict = model.txt_enc.screen_edit(prompts[i], pils_per_sample[i])
+        if verdict.violates:
+            h_, w_ = res_hw[i]
+            print(verdict.banner())
+            results[i] = make_refusal_image(verdict, height=h_, width=w_)
+            continue
+        active.append(i)
+    if not active:
+        return results
+
+    gs_key_int = resolve_gs_key(gs_key)
+
+    # Per sample: reference latent tokens (clean) + target noise tokens, plus the
+    # combined [target, ref_1, …, ref_N] position ids and shapes. ``target_idx``
+    # records where each sample's target tokens land in the packed sequence so we
+    # can slice the velocity and step only the target portion.
+    ch = model.vae.latent_channels
+    targets, refs, ids_list, shape_seq, samp_lens, tgt_lens, hw = [], [], [], [], [], [], []
+    target_idx_parts = []
+    off = 0
+    for i in active:
+        h_, w_ = res_hw[i]
+        torch.manual_seed(seeds[i])  # MageVAE.encode samples the posterior (global RNG)
+        # All references resized to the target resolution and VAE-encoded together.
+        ref_tensors = [_preprocess_ref_image(p, h_, w_, dev) for p in pils_per_sample[i]]
+        ref_tok, ref_shapes, ref_ids = model.compute_vae_encodings(ref_tensors, with_ids=True)
+        ref_tok = ref_tok.to(torch.bfloat16)               # [1, N*Lr, C]
+        x = get_noise(num_samples=1, channel=ch, height=h_, width=w_,
+                      device=dev, dtype=torch.bfloat16, seed=seeds[i])
+        x = encode_noise(tuple(x.shape[1:]), key=gs_key_int,
+                         seed=seeds[i], device=dev, dtype=torch.bfloat16)
+        _, _, gh, gw = x.shape
+        tgt = rearrange(x, "b c h w -> b (h w) c")          # [1, Lt, C]
+        tgt_ids = torch.zeros(gh, gw, 3, device=dev)
+        tgt_ids[..., 1] = tgt_ids[..., 1] + torch.arange(gh, device=dev)[:, None]
+        tgt_ids[..., 2] = tgt_ids[..., 2] + torch.arange(gw, device=dev)[None, :]
+        tgt_ids = rearrange(tgt_ids, "h w c -> (h w) c").unsqueeze(0)
+        lt, lr = tgt.shape[1], ref_tok.shape[1]
+        targets.append(tgt); refs.append(ref_tok)
+        ids_list.append(torch.cat([tgt_ids, ref_ids.to(dev)], dim=1)[0])  # [Lt + N*Lr, 3]
+        shape_seq.append((1, gh, gw))                       # target frame idx 0
+        shape_seq.extend(s[0] for s in ref_shapes)          # ref_j frame idx j
+        samp_lens.append(lt + lr); tgt_lens.append(lt); hw.append((h_, w_))
+        target_idx_parts.append(torch.arange(off, off + lt, device=dev))
+        off += lt + lr
+    img_ids = torch.cat(ids_list, 0).unsqueeze(0)
+    img_cu = _lens_to_cu(samp_lens, dev)
+    img_shapes = [shape_seq]
+    target_idx = torch.cat(target_idx_parts)
+
+    # Packed edit text — positive AND (for CFG) negative are encoded TOGETHER in
+    # ONE packed multimodal forward, then split. Both branches share the same
+    # reference images; cu_seqlens isolates every sequence (zero cross-contamination).
+    # The VL conditioning image's long edge is capped (default 384) to match
+    # training preprocessing — the VAE path above keeps the full target resolution.
+    na = len(active)
+    edit_refs = [[_resize_long_edge(p, vl_cond_long_edge) for p in pils_per_sample[i]]
+                 for i in active]
+    if cfg > 1.0:
+        pos_instr = [prompts[i] for i in active]
+        neg_instr = [neg_prompts[i] or " " for i in active]
+        txt_flat, vec_all, lens_t = _encode_edits_packed(
+            model, edit_refs + edit_refs, pos_instr + neg_instr, template, drop_idx, dev)
+        txt, txt_cu, txt_mask, vec = _slice_packed(txt_flat, vec_all, lens_t, 0, na, dev)
+        neg_txt, neg_cu, neg_mask, neg_vec = _slice_packed(txt_flat, vec_all, lens_t, na, na, dev)
+    else:
+        txt_flat, vec_all, lens_t = _encode_edits_packed(
+            model, edit_refs, [prompts[i] for i in active], template, drop_idx, dev)
+        txt, txt_cu, txt_mask, vec = _slice_packed(txt_flat, vec_all, lens_t, 0, na, dev)
+        neg_txt = neg_cu = neg_mask = neg_vec = None
+
+    ctx = _build_pack_ctx(img_ids, img_cu, img_shapes, samp_lens, txt, txt_cu, txt_mask, vec,
+                          neg_txt, neg_cu, neg_mask, neg_vec, cfg, renormalization, batch_cfg, dev)
+    scheduler = _get_scheduler(model, steps, device, static_shift)
+    for si, t in enumerate(scheduler.timesteps):
+        parts = []
+        for k in range(na):
+            parts.append(targets[k]); parts.append(refs[k])
+        img = torch.cat(parts, dim=1)                       # [1, sum(Lt+Lr), C], ref clean
+        vel = _velocity(model.transformer, img, ctx, scheduler.sigmas[si].item())
+        pred_t = vel[:, target_idx, :]                      # [1, sum Lt, C] — target tokens only
+        tgt_packed = torch.cat(targets, dim=1)              # [1, sum Lt, C]
+        stepped = scheduler.step(pred_t, t, tgt_packed, return_dict=False)[0]
+        o = 0
+        new_targets = []
+        for k in range(na):
+            lt = tgt_lens[k]
+            new_targets.append(stepped[:, o:o + lt, :]); o += lt
+        targets = new_targets
+
+    for k, i in enumerate(active):
+        h_, w_ = hw[k]
+        results[i] = _decode_one(model, targets[k], h_, w_, dev)
+    return results
+
+
+# ---------------------------------------------------------------------------
+# Flow-ODE inversion (Gaussian-Shading watermark detection)
+# ---------------------------------------------------------------------------
+@torch.no_grad()
+def invert_to_noise(model, z0, height, width, steps=30, device="cuda",
+                    prompt_template="mage-flow", static_shift=None, prompt=""):
+    """Reverse the flow ODE from a clean latent ``z0`` back to the initial noise.
+
+    This is the detection primitive for the Gaussian-Shading watermark: VAE-encode
+    the image to ``z0`` (posterior MEAN — deterministic), run this to recover the
+    initial noise, then read the signs via ``mage_latent.decode_bits``.
+
+    Inversion uses an empty prompt at cfg=1 (the standard Tree-Ring /
+    Gaussian-Shading setup). Reverse Euler recovers ``x_i`` from ``x_{i+1}`` with
+    the velocity evaluated at the point in hand; the sign-only watermark tolerates
+    the resulting approximation error (see the module's redundancy).
+
+    Args:
+        z0: clean latent ``[1, C, gh, gw]`` (e.g. the mean of ``model.vae.encode``).
+    Returns:
+        recovered initial-noise latent ``[1, C, gh, gw]`` (float32).
+    """
+    dev = torch.device(device)
+    info = _template_info(prompt_template)
+    template = info.get("template", "{}")
+    drop_idx = int(info.get("start_idx", 0))
+
+    z0 = z0.to(dev)
+    _, ch, gh, gw = z0.shape
+    img = rearrange(z0, "b c h w -> b (h w) c").to(torch.bfloat16)   # [1, gh*gw, C]
+
+    ids = torch.zeros(gh, gw, 3, device=dev)
+    ids[..., 1] = ids[..., 1] + torch.arange(gh, device=dev)[:, None]
+    ids[..., 2] = ids[..., 2] + torch.arange(gw, device=dev)[None, :]
+    img_ids = rearrange(ids, "h w c -> (h w) c").unsqueeze(0)
+    lens = [gh * gw]
+    img_cu = _lens_to_cu(lens, dev)
+    img_shapes = [[(1, gh, gw)]]
+
+    # Empty-prompt conditioning, no negative branch, cfg=1 (single forward).
+    txt_flat, vec_all, lens_t = _encode_texts_packed(model, [prompt], template, drop_idx, dev)
+    txt, txt_cu, txt_mask, vec = _slice_packed(txt_flat, vec_all, lens_t, 0, 1, dev)
+    ctx = _build_pack_ctx(img_ids, img_cu, img_shapes, lens, txt, txt_cu, txt_mask, vec,
+                          None, None, None, None, 1.0, False, False, dev)
+
+    scheduler = _get_scheduler(model, steps, device, static_shift)
+    sigmas = scheduler.sigmas
+    n = len(scheduler.timesteps)
+    # Forward step si: x_{si+1} = x_si + (s_{si+1}-s_si)·v(x_si, s_si).
+    # Reverse it from clean (x_n, sigma 0) up to noise (x_0), using x_{si+1} as the
+    # proxy for x_si at the forward eval sigma s_si.
+    for si in range(n - 1, -1, -1):
+        s_cur = sigmas[si].item()
+        s_next = sigmas[si + 1].item()
+        vel = _velocity(model.transformer, img, ctx, s_cur)
+        img = img - (s_next - s_cur) * vel
+    return unpack(img.float(), height, width)   # [1, C, gh, gw]
+
+
+# ---------------------------------------------------------------------------
+# High-level pipeline wrapper
+# ---------------------------------------------------------------------------
+class MageFlowPipeline:
+    """``MageFlowPipeline.from_pretrained(repo).generate(...) / .edit(...)``.
+
+    ``generate`` / ``edit`` are packed multi-resolution calls: they take a list
+    of prompts (a single string is accepted and treated as a pack of size 1) and
+    return a list of PIL images. Per-sample ``heights``/``widths``/``seeds`` are
+    lists. Every prompt is screened by the text encoder's mandatory content
+    gate (no opt-out); banned prompts come back as refusal placeholders
+    interleaved with the real images. Real outputs always carry a Gaussian-Shading
+    watermark in the initial noise (no toggle), using the configured secret key.
+    """
+
+    def __init__(self, model, device="cuda"):
+        self.model = model
+        self.device = device
+
+    @classmethod
+    def from_pretrained(cls, repo_dir: str, device: str = "cuda"):
+        """Load a Mage-Flow diffusers-style repo (``model_index.json`` +
+        ``transformer/`` ``vae/`` ``scheduler/`` ``text_encoder/``).
+
+        ``repo_dir`` may be a local directory OR a Hugging Face Hub repo id
+        (e.g. ``"microsoft/Mage-Flow-4B"``), which is downloaded and cached
+        automatically on first use.
+        """
+        return cls(load_from_repo(repo_dir, device), device)
+
+    def generate(self, prompts, **kw) -> list[Image.Image]:
+        """Packed multi-resolution t2i. ``prompts`` is a list (or a single
+        string); pass per-sample ``heights``/``widths``/``seeds`` as lists."""
+        kw.setdefault("device", self.device)
+        return generate_images(self.model, prompts, **kw)
+
+    def edit(self, prompts, ref_images, **kw) -> list[Image.Image]:
+        """Packed multi-resolution edit. ``prompts`` is a list (or a single
+        string); each ``ref_images[i]`` is one reference or a list of references."""
+        kw.setdefault("device", self.device)
+        return generate_edits(self.model, prompts, ref_images, **kw)
+
+    def invert_to_noise(self, z0, height, width, **kw):
+        """Recover the initial noise from a clean latent (Gaussian-Shading detect)."""
+        kw.setdefault("device", self.device)
+        return invert_to_noise(self.model, z0, height, width, **kw)
+
+
+def _safe_subpath(root: str, *parts: str) -> str:
+    """Join ``parts`` under ``root`` and confirm the result stays inside ``root``.
+
+    ``root`` is normalized up front; the joined path is normalized **lexically**
+    (``os.path.normpath`` — symlinks are *not* followed, so a Hugging Face cache
+    whose weight files are symlinks into the shared blob store still loads) and
+    rejected if it escapes ``root``. This guards the user-supplied model path
+    against path traversal (CWE-22 / CodeQL ``py/path-injection``).
+    """
+    root = os.path.realpath(root)
+    full = os.path.normpath(os.path.join(root, *parts))
+    if full != root and not full.startswith(root + os.sep):
+        raise ValueError(
+            f"Resolved path {os.path.join(*parts)!r} escapes repo directory {root!r}"
+        )
+    return full
+
+
+def _resolve_repo_dir(repo_dir: str) -> str:
+    """Return a local directory for ``repo_dir``.
+
+    If ``repo_dir`` is an existing local path it is returned as a normalized
+    absolute path; otherwise it is treated as a Hugging Face Hub repo id (e.g.
+    ``microsoft/Mage-Flow``) and downloaded/cached via
+    ``huggingface_hub.snapshot_download``.
+    """
+    candidate = os.path.realpath(repo_dir)
+    if os.path.isdir(candidate):
+        return candidate
+    from huggingface_hub import snapshot_download
+    return snapshot_download(repo_id=repo_dir)
+
+
+def load_from_repo(repo_dir: str, device: str = "cuda") -> MageFlowModel:
+    """Load a Mage-Flow diffusers-style repo (model_index.json + transformer/
+    vae/ scheduler/). Transformer weights come from the bf16 safetensors;
+    VAE + text encoder are built from the sources recorded in model_index.json.
+
+    ``repo_dir`` may be a local directory OR a Hugging Face Hub repo id (e.g.
+    ``microsoft/Mage-Flow-4B``), which is downloaded/cached automatically.
+    """
+    from safetensors.torch import load_file
+    repo_dir = _resolve_repo_dir(repo_dir)
+    mi = json.load(open(_safe_subpath(repo_dir, "model_index.json")))
+    tcfg = json.load(open(_safe_subpath(repo_dir, "transformer", "config.json")))
+    # Keys stripped from the checkpoint config before it becomes model_structure.
+    # ``schedule_mode`` is a legacy field still present in some config.json files;
+    # Keys of the checkpoint config that are NOT MageFlowParams constructor args
+    # (legacy/unused fields). Everything else becomes model_structure. The DiT only
+    # reads: in_channels, out_channels, context_in_dim, hidden_size, num_heads,
+    # depth, axes_dim, checkpoint, patch_size.
+    _meta = {"_class_name", "txt_max_length", "max_sequence_length", "param_dtype",
+             "packing", "schedule_mode", "static_shift", "use_time_shift",
+             "rope_type", "apply_text_rotary_emb",
+             "mlp_ratio", "depth_single_blocks", "theta", "qkv_bias", "guidance_embed",
+             "vec_in_dim", "vec_type", "time_type", "double_block_type"}
+    structure = {k: v for k, v in tcfg.items() if k not in _meta}
+
+    def _resolve(p):
+        return p if os.path.isabs(p) else _safe_subpath(repo_dir, p)
+
+    cfg = ModelConfig(
+        vae_path=_resolve(mi.get("_vae_source")),
+        txt_enc_path=_resolve(mi.get("_text_encoder_path")),
+        model_structure=structure,
+        txt_max_length=tcfg.get("txt_max_length", 2048),
+        packing=tcfg.get("packing", True),
+        static_shift=tcfg.get("static_shift", 6.0),
+    )
+    model = MageFlowModel(cfg)
+    sd = load_file(_safe_subpath(repo_dir, "transformer", "diffusion_pytorch_model.safetensors"),
+                   device="cpu")
+    model.transformer.load_state_dict(sd, strict=False, assign=True)
+    model.to(device)
+    model.transformer.to(torch.bfloat16)
+    model.txt_enc.to(torch.bfloat16)
+    if model.vae is not None:
+        model.vae.to(torch.bfloat16)
+    model.eval()
+    # Diffusers FlowMatchEulerDiscreteScheduler (scheduler/scheduler_config.json).
+    model.scheduler = FlowMatchEulerDiscreteScheduler.from_pretrained(
+        _safe_subpath(repo_dir, "scheduler"))
+    return model