Load with the stock Diffusers and Transformers classes: one-hot timestep code, 51-layer text encoder

README.md

@@ -15,47 +15,38 @@ tags:
 # MiniMax-H3-Image-Pruned
 
 [MiniMax-H3-Image](https://huggingface.co/MiniMaxAI/MiniMax-H3-Image) without the weights its inference never needs:
-135.3 GB → 94.1 GB on disk (transformer, text encoder and VAE), with images bit-identical to the original on the same
-GPU. Every kept tensor is a byte copy of the original; nothing is quantized or retrained.
+135.3 GB → 95.1 GB on disk (transformer, text encoder and VAE). The stock Diffusers and Transformers classes load
+it like the original, with no custom code and no `trust_remote_code`. Text-to-image results are bit-identical to the
+original on the same GPU; edit results differ by a small margin (PSNR 43.3 dB on the test edit). Every kept tensor is a byte copy of the original;
+nothing is quantized or retrained.
 
 ![Generation time, GPU memory and checkpoint size](assets/speed_memory_size.png)
 
-![Samples: an edit and two text-to-image results, bit-identical to the original](assets/samples.jpg)
+![Samples: two text-to-image results, bit-identical to the original, and an edit](assets/samples.jpg)
 
 | | Original | Pruned |
 |---|---|---|
-| Transformer | 66.3 GB, 33.1B parameters | 40.3 GB, 20.1B parameters + 116 MB of tables |
-| Text encoder (Qwen3-VL-32B) | 66.7 GB, 64 decoder layers | 51.5 GB, 50 decoder layers |
+| Transformer | 66.3 GB, 33.1B parameters | 40.2 GB, 20.1B parameters |
+| Text encoder (Qwen3-VL-32B) | 66.7 GB, 64 decoder layers | 52.5 GB, 51 decoder layers |
 | VAE | 2.3 GB | 2.3 GB (unchanged) |
 
-## What was removed
-
-- **AdaLN projections (13.0B parameters).** Every block projects the timestep embedding to its modulation vectors.
-  The checkpoint samples on a fixed six-step grid (one output head per step), and the reference images of an edit stay
-  at t = 0.999, so these projections only ever produce a fixed set of values. They are stored as tables (116 MB)
-  computed with the original modules, and the projections, `time_embedder` and `norm_out.linear` are gone. One third
-  of every projection served the audio modality of MiniMax-H3, which the image model never uses; it is not kept.
-- **Text encoder layers 50-63 and `lm_head`.** The transformer is conditioned on `hidden_states[50]` of the Qwen3-VL
-  text encoder, the output of decoder layer 49. Nothing after it contributes.
-- **Dead rows of the last block (compute only).** After the last block only the generated image rows are read. The
-  text rows and the reference rows still provide keys and values, but their queries, attention output and feed-forward
-  are skipped.
-
 ## Usage
 
-Install the loader and Diffusers with MiniMax-H3-Image support:
-
 ```bash
-pip install "git+https://github.com/iamwavecut/MiniMax-H3-Image-Pruned" \
-  "diffusers @ git+https://github.com/huggingface/diffusers@875373ac7f5856071e345a3e394c74efccc6b9e6" \
-  "transformers>=5.22"
+pip install "diffusers @ git+https://github.com/huggingface/diffusers@875373ac7f5856071e345a3e394c74efccc6b9e6" \
+  "transformers>=5.22" accelerate torchvision
 ```
 
 ```python
 import torch
-from minimax_h3_image_pruned import load_pipeline
+from diffusers import ComponentsManager, ModularPipeline
+
+manager = ComponentsManager()
+manager.enable_auto_cpu_offload(device="cuda")  # or skip this and call pipe.to("cuda") on a GPU with 96 GB or more
+
+pipe = ModularPipeline.from_pretrained("WaveCut/MiniMax-H3-Image-Pruned", components_manager=manager)
+pipe.load_components(dtype=torch.bfloat16)
 
-pipe = load_pipeline("WaveCut/MiniMax-H3-Image-Pruned")  # model CPU offload; offload=False keeps everything on the GPU
 image = pipe(
     prompt="The image is a top-down photograph of ...",
     height=3552,
@@ -65,76 +56,92 @@ image = pipe(
 )[0]
 ```
 
-Editing: pass `image=reference` (or a list) and refer to the images as `<Picture 1>`, `<Picture 2>`, ... in the
-prompt. The prompting guidance of the [original card](https://huggingface.co/MiniMaxAI/MiniMax-H3-Image) applies
-unchanged: long, complete English descriptions; recommended sizes from 4096×4096 (1:1) to 6272×2688 (21:9).
+This is the code of the [original card](https://huggingface.co/MiniMaxAI/MiniMax-H3-Image) with another repository
+name. Editing: pass `image=reference` (or a list) and refer to the images as `<Picture 1>`, `<Picture 2>`, ... in the
+prompt. The prompting guidance of the original card applies unchanged: long, complete English descriptions;
+recommended sizes from 4096×4096 (1:1) to 6272×2688 (21:9).
 
-The transformer alone also loads through Diffusers' remote code:
+The components also load alone with the stock classes: `MiniMaxH3ImageTransformer2DModel.from_pretrained(repo,
+subfolder="transformer")` and `Qwen3VLForConditionalGeneration.from_pretrained(repo, subfolder="text_encoder")`.
 
-```python
-from diffusers import AutoModel
+## What was removed
 
-transformer = AutoModel.from_pretrained(
-    "WaveCut/MiniMax-H3-Image-Pruned", subfolder="transformer", trust_remote_code=True, torch_dtype=torch.bfloat16
-)
-```
+- **AdaLN projection inputs (13.0B parameters).** Every block projects the timestep embedding to its modulation
+  vectors, and the checkpoint samples on a fixed grid: six steps (one output head per step), plus t = 0.999 for the
+  reference images of an edit. Seven timesteps need seven input columns, so `adaln_proj.linear` and `norm_out.linear`
+  keep 8 of their 2688 input columns (`time_embed_dim: 8` in the config). The new `time_embedder` (16 hidden units)
+  turns each grid timestep into a one-hot code, and column k holds the original projection's output for timestep k.
+  The stock forward pass then returns the original modulation values exactly; the
+  [GitHub repository](https://github.com/iamwavecut/MiniMax-H3-Image-Pruned) explains the construction.
+- **Text encoder layers 51-63 and `lm_head`.** The transformer is conditioned on `hidden_states[50]` of the Qwen3-VL
+  text encoder, the output of decoder layer 49. Layer 50 stays because the pipeline requires more than 50 layers; with
+  51 layers, `hidden_states[50]` is still the output of layer 49. `lm_head` is tied to the input embedding, so loading
+  allocates nothing for it.
+
+## Text-to-image and edit
+
+The original model computes the modulation for a step with one timestep row in text-to-image and with two rows (step
+and reference) in an edit, and the matrix products round slightly differently in the two cases. The pruned
+checkpoint stores one value per timestep: the text-to-image value for the six steps and the edit value for the
+reference timestep.
 
-The text encoder needs `minimax_h3_image_pruned.adapt_text_encoder` after loading: it replaces the final norm with
-an identity, so that the last hidden state of the 50-layer stack is the conditioning. The
-[GitHub repository](https://github.com/iamwavecut/MiniMax-H3-Image-Pruned) has the build script, the verification
-script and the details.
+- Text-to-image: bit-identical to the original.
+- Edit: not bit-identical. 59,279 of the 29.1 million modulation values of the six step rows (0.20 %) differ from the original's two-row result by rounding steps. Through the 50 blocks this changes each edit step's transformer output by 2.1-2.5 % (relative L2), and the 2048×2048 test edit by PSNR 43.3 dB: mean absolute difference 0.63 of 255, 261 of 4.2 million pixels off by more than 64. In the worst region the differences follow edges; the two images look the same.
 
 ## Profiles
 
-Measured on one RTX PRO 6000 Blackwell (96 GB), bfloat16, one image after a warm-up image. Time includes prompt
-encoding, six steps and VAE decoding.
+Measured on one RTX PRO 6000 Blackwell Server Edition (96 GB, 600 W power limit), bfloat16, one image after a
+warm-up image. Time includes prompt encoding, six steps and VAE decoding.
 
 | Run | 2048×2048 | Edit, 2048×2048 | 4736×3552 | GPU memory peak, 4736×3552 |
 |---|---|---|---|---|
-| Original, model offload | 80.3 s | 83.7 s | 101.2 s | 68.2 GiB |
-| Pruned, model offload (`load_pipeline()`) | 5.8 s | 9.2 s | 26.3 s | 92.6 GiB |
-| Pruned, all on the GPU (`offload=False`) | 5.5 s | 8.9 s | 25.8 s | 92.6 GiB |
+| Original, model offload | 97.6 s | 101.6 s | 120.8 s | 68.2 GiB |
+| Pruned, model offload (code above) | 5.6 s | 9.4 s | 26.4 s | 92.7 GiB |
+| Pruned, all on the GPU (`pipe.to("cuda")`) | 5.6 s | 9.3 s | 26.7 s | 92.7 GiB |
 
 The time difference comes from memory, not from compute: the original (124 GiB of transformer and text encoder) does
-not fit a 96 GB card, so the offload manager moves both models for every image. The pruned models (85.5 GiB) stay on
-the GPU. The sampling steps themselves cost the same. On a GPU where the original fits (141 GB or more), expect similar
-times for both; on an 80 GB GPU the pruned models also take turns on the GPU, with less data to move (neither case was
-measured).
+not fit a 96 GB card, so the offload manager moves both models for every image. The pruned models (86.4 GiB)
+stay on the GPU. The sampling steps themselves cost the same. On a GPU where the original fits (141 GB or more),
+expect similar times for both; on an 80 GB GPU the pruned models also take turns on the GPU, with less data to move
+(neither case was measured).
 
 ## Verification
 
-`scripts/verify.py` in the GitHub repository ran the official Modular Diffusers pipeline and this checkpoint on the
-same prompts and seeds: two text-to-image prompts and one edit at 2048×2048, and one text-to-image prompt at 4736×3552.
+`scripts/verify.py` in the GitHub repository ran the official Modular Diffusers pipeline and this checkpoint, both
+with the stock classes, on the same prompts and seeds: two text-to-image prompts and one edit at 2048×2048, and one
+text-to-image prompt at 4736×3552.
 
-- Images: bit-identical in every run (0 differing pixels), with `skip_dead_rows` on (default) and off, with and
-  without offload.
-- Text conditioning (`hidden_states[50]`, 36 to 1757 tokens): bit-identical.
-- Transformer: the 18 recorded calls of the original pipeline, replayed through the pruned transformer, give
-  bit-identical outputs with and without `skip_dead_rows`.
+- Text-to-image images: bit-identical in every run (0 differing pixels), with model offload and with everything on
+  the GPU.
+- Text conditioning (`hidden_states[50]`, 36 to 1757 tokens, the edit included): bit-identical.
+- Transformer: the 18 recorded calls of the original pipeline, replayed through the pruned transformer. The 12
+  text-to-image calls give bit-identical outputs; the 6 edit calls differ by 2.1-2.5 % (relative L2).
+- Edit image: PSNR 43.3 dB against the original, mean absolute difference 0.63 of 255.
+- Build checks: every grid timestep gives an exact one-hot code with float32, TF32 and bfloat16 matmuls; every
+  stored column times 32 reproduces the original projection output.
 
-Software: torch 2.11.0+cu128, diffusers at `875373ac`, transformers 5.22.2, NVIDIA driver 610.57; source checkpoint
+Software: torch 2.11.0+cu128, diffusers at `875373ac`, transformers 5.22.2, NVIDIA driver 615.71.09; source checkpoint
 `MiniMaxAI/MiniMax-H3-Image@492fd8c4`.
 
 ## Limits
 
-- Only the release's six-step sampling grid: the tables hold no other timesteps, and the transformer raises an error
-  for them. The original has one output head per step of this grid and cannot change it either.
-- The tables were computed on an RTX PRO 6000 Blackwell. On another GPU type the original model can round its own AdaLN
-  outputs differently by single bfloat16 steps, while the pruned model keeps the values of the build GPU, so the two
-  can differ there by about as much as the original differs between two GPU types (not measured).
-  `minimax-h3-image-prune --skip text_encoder copies` rebuilds the transformer and its tables on your GPU from the
-  original transformer (66 GB download, tables in about 5 s).
-- Diffusers' context-parallel and tensor-parallel plans are not carried over. SGLang and ComfyUI do not load this
-  layout.
+- Only the release's sampling grid. A timestep off the grid gives no valid code, and the stock transformer does not
+  check for it: the result is wrong without an error. The original has one output head per step of this grid and
+  cannot change it either.
+- The stored values were computed on an RTX PRO 6000 Blackwell. On another GPU type the original model can round its
+  own AdaLN outputs differently by single bfloat16 steps, while this checkpoint keeps the values of the build GPU, so
+  the two can differ there by about as much as the original differs between two GPU types (not measured).
+  `minimax-h3-image-prune --skip text_encoder copies` from the GitHub repository rebuilds the transformer on your GPU.
+- SGLang and ComfyUI were not tested with this checkpoint.
 
 ## Files
 
 | Path | Content |
 |---|---|
-| `transformer/` | Pruned transformer, modulation tables, `modeling_minimax_h3_image_pruned.py` for remote code |
-| `text_encoder/` | Qwen3-VL conditioner, decoder layers 0-49, with its tokenizer files |
+| `transformer/` | Transformer with 8-column AdaLN projections and the one-hot `time_embedder`; `build_report.json` |
+| `text_encoder/` | Qwen3-VL conditioner, decoder layers 0-50, with its tokenizer files |
 | `vae/`, `tokenizer/`, `processor/`, `scheduler/` | Original files |
-| `modular_model_index.json` | Modular Diffusers index |
+| `modular_model_index.json` | Modular Diffusers index, original component classes |
 | `LICENSE`, `NOTICE` | MiniMax-H3-Image Community License and the required notice |
 | `assets/` | Card images |
 

modular_model_index.json

@@ -69,11 +69,11 @@
   ],
   "transformer": [
     "diffusers",
-    "AutoModel",
+    "MiniMaxH3ImageTransformer2DModel",
     {
       "type_hint": [
         "diffusers",
-        "AutoModel"
+        "MiniMaxH3ImageTransformer2DModel"
       ],
       "pretrained_model_name_or_path": "WaveCut/MiniMax-H3-Image-Pruned",
       "subfolder": "transformer",

text_encoder/config.json

@@ -18,7 +18,7 @@
     "max_position_embeddings": 262144,
     "model_type": "qwen3_vl_text",
     "num_attention_heads": 64,
-    "num_hidden_layers": 50,
+    "num_hidden_layers": 51,
     "num_key_value_heads": 8,
     "rms_norm_eps": 1e-06,
     "rope_scaling": {
@@ -59,5 +59,5 @@
   },
   "vision_end_token_id": 151653,
   "vision_start_token_id": 151652,
-  "_notice": "Modified by WaveCut from MiniMaxAI/MiniMax-H3-Image@492fd8c4eca666f1085f57e30738c1e0d86b3439: decoder layers 50+ and lm_head removed; the pipeline reads hidden_states[50]."
+  "_notice": "Modified by WaveCut from MiniMaxAI/MiniMax-H3-Image@492fd8c4eca666f1085f57e30738c1e0d86b3439: decoder layers 51+ and lm_head removed; the pipeline reads hidden_states[50], the output of layer 49."
 }

text_encoder/model-00001-of-00011.safetensors

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:e7a16e4ba468fa3cf0e6a61eafffebafe20aea2b15782c9bf6b39c02356ad836
-size 4932329128
+oid sha256:5f6e3eb429c6d65aceed6e5d17819fec09ff34bbb12d4a8ccae5b3777b371b54
+size 4932329152

text_encoder/model-00002-of-00011.safetensors

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:8925f84c9aee266e20fdfe9b8163ed017e8f5f12ebff14e46c6dcf07afb122ef
-size 4876001080
+oid sha256:74c0af2731f04da4ce48caade03b1a90a725bda19a615a8fd12f4cac364ee6ef
+size 4876001104

text_encoder/model-00003-of-00011.safetensors

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:f454cd71f516e7135be68f6c9987ba1d53d6c47addf7acd435917e23c1fec512
-size 4875980376
+oid sha256:fb0dbbc58d92a3d2de295e88a90fe1fbc508310d755205eca46924c1d4d5ffaf
+size 4875980400

text_encoder/model-00004-of-00011.safetensors

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:8abd419815d37b7bfc5db51ad3961e13183872ff2a6e0443cf2bf3f99ccc1037
-size 4875990768
+oid sha256:591d440d3acb23ebeecdcaddb3b89be1da4eff6c027d93a401cedf2beb03e90e
+size 4875990792

text_encoder/model-00005-of-00011.safetensors

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:f69f2d22cd93b6b4099148da9e96674c1239e04dc8411599247d986d635f2e4b
-size 4875990768
+oid sha256:c3eb9d885d04e339c4d1ac8d2b291d33b36886b06bcc01e2070c58223080cfc3
+size 4875990792

text_encoder/model-00006-of-00011.safetensors

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:fb61024f709eb8407e5d2e7dba1be6dc6f313fd300e7b5e254277cc9512c9745
-size 4875990768
+oid sha256:a1a5c7bd0643672f25171364d6dd5e3d19e8ee4558c1bc01a3ac3b617775ce96
+size 4875990792

text_encoder/model-00007-of-00011.safetensors

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:89175e8de21cd30ccca42f670a4be62a2538b2ed651c7099aeeb57914e8712e4
-size 4875990768
+oid sha256:78ecbfa2d8cb0ee44b318d52b380767823e1d36733d03861fb1ee76865c6c737
+size 4875990792

text_encoder/model-00008-of-00011.safetensors

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:ac167ff342cb3b86a71200abfb6d4a5255394dd8a7f9dfbb4bf312f6e201b606
-size 4875990768
+oid sha256:960a0adf383dc9106e68aa00f7a8f6af59cfabb23f401b81d87d078c9ee1114c
+size 4875990792

text_encoder/model-00009-of-00011.safetensors

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:7bd0b7d702ac089d6a25310565d3343c4d4a7536ffefeb0770e2daf868fed602
-size 4875990768
+oid sha256:bf739286fe73ad7550339bb899318d874b51fe3c9d801d6f751f8202357d90f1
+size 4875990792

text_encoder/model-00010-of-00011.safetensors

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:266bcfbbc17abe40b50a3f905a6f4d3df196538888839fcf78e06bd3f8bf489e
-size 4875990768
+oid sha256:e56d34af01da070b4103ba4e8c4e28a4313e59972bd99b1f124631f5659dc8e0
+size 4875990792

text_encoder/model-00011-of-00011.safetensors

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:de64e35dbeaa4869cef2d358bb960d6900f09d3236daf5b4609501c9762aa9eb
-size 2690070560
+oid sha256:87aebd21a73bf1743204f020b879e2ccd94322c2c02d2ce49bfdd9b2bfc79b7b
+size 3665268696

text_encoder/model.safetensors.index.json

@@ -1,6 +1,6 @@
 {
   "metadata": {
-    "total_size": 51506202080
+    "total_size": 52481398752
   },
   "weight_map": {
     "model.language_model.embed_tokens.weight": "model-00001-of-00011.safetensors",
@@ -510,6 +510,17 @@
     "model.language_model.layers.5.self_attn.q_norm.weight": "model-00002-of-00011.safetensors",
     "model.language_model.layers.5.self_attn.q_proj.weight": "model-00002-of-00011.safetensors",
     "model.language_model.layers.5.self_attn.v_proj.weight": "model-00002-of-00011.safetensors",
+    "model.language_model.layers.50.input_layernorm.weight": "model-00011-of-00011.safetensors",
+    "model.language_model.layers.50.mlp.down_proj.weight": "model-00011-of-00011.safetensors",
+    "model.language_model.layers.50.mlp.gate_proj.weight": "model-00011-of-00011.safetensors",
+    "model.language_model.layers.50.mlp.up_proj.weight": "model-00011-of-00011.safetensors",
+    "model.language_model.layers.50.post_attention_layernorm.weight": "model-00011-of-00011.safetensors",
+    "model.language_model.layers.50.self_attn.k_norm.weight": "model-00011-of-00011.safetensors",
+    "model.language_model.layers.50.self_attn.k_proj.weight": "model-00011-of-00011.safetensors",
+    "model.language_model.layers.50.self_attn.o_proj.weight": "model-00011-of-00011.safetensors",
+    "model.language_model.layers.50.self_attn.q_norm.weight": "model-00011-of-00011.safetensors",
+    "model.language_model.layers.50.self_attn.q_proj.weight": "model-00011-of-00011.safetensors",
+    "model.language_model.layers.50.self_attn.v_proj.weight": "model-00011-of-00011.safetensors",
     "model.language_model.layers.6.input_layernorm.weight": "model-00002-of-00011.safetensors",
     "model.language_model.layers.6.mlp.down_proj.weight": "model-00002-of-00011.safetensors",
     "model.language_model.layers.6.mlp.gate_proj.weight": "model-00002-of-00011.safetensors",

transformer/build_report.json

@@ -0,0 +1,117 @@
+{
+ "time_embed_hidden_dim": 16,
+ "time_embed_dim": 8,
+ "grid": [
+  0.0,
+  0.03247630596160889,
+  0.11245149374008179,
+  0.25,
+  0.4672260880470276,
+  0.7679911255836487,
+  0.9990000128746033
+ ],
+ "plans": [
+  {
+   "timesteps": [
+    0.0
+   ],
+   "rows": 51,
+   "identical_rows": 51,
+   "differing": 0
+  },
+  {
+   "timesteps": [
+    0.03247630596160889
+   ],
+   "rows": 51,
+   "identical_rows": 51,
+   "differing": 0
+  },
+  {
+   "timesteps": [
+    0.11245149374008179
+   ],
+   "rows": 51,
+   "identical_rows": 51,
+   "differing": 0
+  },
+  {
+   "timesteps": [
+    0.25
+   ],
+   "rows": 51,
+   "identical_rows": 51,
+   "differing": 0
+  },
+  {
+   "timesteps": [
+    0.4672260880470276
+   ],
+   "rows": 51,
+   "identical_rows": 51,
+   "differing": 0
+  },
+  {
+   "timesteps": [
+    0.7679911255836487
+   ],
+   "rows": 51,
+   "identical_rows": 51,
+   "differing": 0
+  },
+  {
+   "timesteps": [
+    0.0,
+    0.9990000128746033
+   ],
+   "rows": 102,
+   "identical_rows": 51,
+   "differing": 3064
+  },
+  {
+   "timesteps": [
+    0.03247630596160889,
+    0.9990000128746033
+   ],
+   "rows": 102,
+   "identical_rows": 51,
+   "differing": 19629
+  },
+  {
+   "timesteps": [
+    0.11245149374008179,
+    0.9990000128746033
+   ],
+   "rows": 102,
+   "identical_rows": 51,
+   "differing": 3683
+  },
+  {
+   "timesteps": [
+    0.25,
+    0.9990000128746033
+   ],
+   "rows": 102,
+   "identical_rows": 51,
+   "differing": 18412
+  },
+  {
+   "timesteps": [
+    0.4672260880470276,
+    0.9990000128746033
+   ],
+   "rows": 102,
+   "identical_rows": 51,
+   "differing": 10318
+  },
+  {
+   "timesteps": [
+    0.7679911255836487,
+    0.9990000128746033
+   ],
+   "rows": 102,
+   "identical_rows": 51,
+   "differing": 4173
+  }
+ ]
+}

transformer/config.json

@@ -1,5 +1,5 @@
 {
-  "_class_name": "MiniMaxH3ImagePrunedTransformer2DModel",
+  "_class_name": "MiniMaxH3ImageTransformer2DModel",
   "_diffusers_version": "0.42.0.dev0",
   "num_attention_heads": 56,
   "attention_head_dim": 128,
@@ -15,61 +15,22 @@
   ],
   "text_dim": 5120,
   "freq_dim": 256,
-  "time_embed_hidden_dim": 5376,
-  "time_embed_dim": 2688,
+  "time_embed_hidden_dim": 16,
+  "time_embed_dim": 8,
   "rope_freq_dim": 16,
   "rope_theta": 10000.0,
   "norm_eps": 1e-05,
   "qk_norm_eps": 1e-05,
   "final_norm_eps": 1e-05,
   "num_pdd_heads": 6,
-  "adaln_plans": [
-    [
-      0.0
-    ],
-    [
-      0.03247630596160889
-    ],
-    [
-      0.11245149374008179
-    ],
-    [
-      0.25
-    ],
-    [
-      0.4672260880470276
-    ],
-    [
-      0.7679911255836487
-    ],
-    [
-      0.0,
-      0.9990000128746033
-    ],
-    [
-      0.03247630596160889,
-      0.9990000128746033
-    ],
-    [
-      0.11245149374008179,
-      0.9990000128746033
-    ],
-    [
-      0.25,
-      0.9990000128746033
-    ],
-    [
-      0.4672260880470276,
-      0.9990000128746033
-    ],
-    [
-      0.7679911255836487,
-      0.9990000128746033
-    ]
+  "_timestep_grid": [
+    0.0,
+    0.03247630596160889,
+    0.11245149374008179,
+    0.25,
+    0.4672260880470276,
+    0.7679911255836487,
+    0.9990000128746033
   ],
-  "reference_timestep": 0.999,
-  "auto_map": {
-    "AutoModel": "modeling_minimax_h3_image_pruned.MiniMaxH3ImagePrunedTransformer2DModel"
-  },
-  "_notice": "Modified by WaveCut from MiniMaxAI/MiniMax-H3-Image@492fd8c4eca666f1085f57e30738c1e0d86b3439: AdaLN projections, time_embedder and norm_out.linear replaced by precomputed modulation tables."
+  "_notice": "Modified by WaveCut from MiniMaxAI/MiniMax-H3-Image@492fd8c4eca666f1085f57e30738c1e0d86b3439: time_embedder replaced by a one-hot code of the 7 sampling timesteps; adaln_proj.linear and norm_out.linear reduced to 8 input columns holding the original outputs for those timesteps."
 }

transformer/diffusion_pytorch_model-00001-of-00009.safetensors

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:68e5eb0788c1a321749667c0f9798290d13c68c30b48a4568a8da88d75d73344
-size 4989811000
+oid sha256:c28f1749ea59bc84ddad0934b12fe21a636c274849963ac1fb9631cb75cd9804
+size 4998732480

transformer/diffusion_pytorch_model-00002-of-00009.safetensors

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:d0893326c9b8af7eca3dbf9299bd1c0c0147de61b7f8cdbd1688e80b333cfa45
-size 4932641896
+oid sha256:2b2d8e5099a9f8ed4b8371f8562e955ada7760497bcc2076b98c70df34cb5b30
+size 4944836432

transformer/diffusion_pytorch_model-00003-of-00009.safetensors

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:577bff7b1d1496d5fdb7cffbb3b78e1289403f783d96286cca378f0c69eb4a01
-size 4932662456
+oid sha256:6ca4bdca9a6950bdf2adb6ac52cbf176ce15b84eab6fd4c6c487e11ec0f7da66
+size 4943114936

transformer/diffusion_pytorch_model-00004-of-00009.safetensors

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:33fb113392a286a33b3584ecb7c16eac7b67cdd30bd12c7cee7f4919cb72313e
-size 4778500288
+oid sha256:3875073d69d89f17be9d000ae92416f9243584a697271f37f7e176cd3ff68d58
+size 4788952768

transformer/diffusion_pytorch_model-00005-of-00009.safetensors

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:7539f61b8d26a87e824cfa22d0d995e82dd8b65ae02becddc086cc5efd51c78d
-size 4932640856
+oid sha256:79af9470894791a484073eacd12b4fdd7f4ec7279c7d6e9fe54ce6b898b3db84
+size 4943093336

transformer/diffusion_pytorch_model-00006-of-00009.safetensors

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:07a942f1af06cf73e9ea00540d50e602d374317ef6e28f93f0e4f5aff57a7e0c
-size 4932663536
+oid sha256:e8704a48eabe8c10b66e76afc5501a1ec67f798712d2133c1441e68231049263
+size 4944858080

transformer/diffusion_pytorch_model-00007-of-00009.safetensors

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:9fe08e19d2b615020668227dff95dc07a28dbf01f8cb2ca9c3bf0efe20d1b81c
-size 4778500288
+oid sha256:be102ad0f5b504dea1e1aa23645e922d9a3ff949171f4d71010e8e1d696bdae5
+size 4788952768

transformer/diffusion_pytorch_model-00008-of-00009.safetensors

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:1d157eae066c6fced587483097a263a71e8d3ea99c5d5e0fb346f6cfe9ea5a88
-size 4945030584
+oid sha256:e1637c44dd40310c7d0455a781f1c6c0b1fbb89de7f3d11b29e23f704598cccd
+size 4955483064

transformer/diffusion_pytorch_model-00009-of-00009.safetensors

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:c7ef92ac9cc735b8c8c513ce5f98e8e86cf5083deac48c385269de1be532ca87
-size 1041403736
+oid sha256:85f043e5221effcaec1f1677555cb8f364175f08d72e870b348055e20a7902dc
+size 926631400

transformer/diffusion_pytorch_model.safetensors.index.json

@@ -1,12 +1,13 @@
 {
   "metadata": {
-    "total_size": 40263785216
+    "total_size": 40234578272
   },
   "weight_map": {
     "context_embedder.bias": "diffusion_pytorch_model-00001-of-00009.safetensors",
     "context_embedder.weight": "diffusion_pytorch_model-00001-of-00009.safetensors",
+    "norm_out.linear.bias": "diffusion_pytorch_model-00001-of-00009.safetensors",
+    "norm_out.linear.weight": "diffusion_pytorch_model-00001-of-00009.safetensors",
     "norm_out.norm.weight": "diffusion_pytorch_model-00001-of-00009.safetensors",
-    "norm_out.table": "diffusion_pytorch_model-00009-of-00009.safetensors",
     "pdd_proj_out.0.bias": "diffusion_pytorch_model-00008-of-00009.safetensors",
     "pdd_proj_out.0.weight": "diffusion_pytorch_model-00008-of-00009.safetensors",
     "pdd_proj_out.1.bias": "diffusion_pytorch_model-00008-of-00009.safetensors",
@@ -21,6 +22,10 @@
     "pdd_proj_out.5.weight": "diffusion_pytorch_model-00008-of-00009.safetensors",
     "proj_in.bias": "diffusion_pytorch_model-00001-of-00009.safetensors",
     "proj_in.weight": "diffusion_pytorch_model-00001-of-00009.safetensors",
+    "time_embedder.linear_1.bias": "diffusion_pytorch_model-00001-of-00009.safetensors",
+    "time_embedder.linear_1.weight": "diffusion_pytorch_model-00001-of-00009.safetensors",
+    "time_embedder.linear_2.bias": "diffusion_pytorch_model-00001-of-00009.safetensors",
+    "time_embedder.linear_2.weight": "diffusion_pytorch_model-00001-of-00009.safetensors",
     "token_refiner.final_norm.weight": "diffusion_pytorch_model-00001-of-00009.safetensors",
     "token_refiner.refiner_blocks.0.attn.norm_k.weight": "diffusion_pytorch_model-00001-of-00009.safetensors",
     "token_refiner.refiner_blocks.0.attn.norm_q.weight": "diffusion_pytorch_model-00001-of-00009.safetensors",
@@ -42,7 +47,8 @@
     "token_refiner.refiner_blocks.1.ff.net.2.weight": "diffusion_pytorch_model-00001-of-00009.safetensors",
     "token_refiner.refiner_blocks.1.norm1.weight": "diffusion_pytorch_model-00001-of-00009.safetensors",
     "token_refiner.refiner_blocks.1.norm2.weight": "diffusion_pytorch_model-00001-of-00009.safetensors",
-    "transformer_blocks.0.adaln_proj.table": "diffusion_pytorch_model-00009-of-00009.safetensors",
+    "transformer_blocks.0.adaln_proj.linear.bias": "diffusion_pytorch_model-00001-of-00009.safetensors",
+    "transformer_blocks.0.adaln_proj.linear.weight": "diffusion_pytorch_model-00001-of-00009.safetensors",
     "transformer_blocks.0.attn.norm_k.weight": "diffusion_pytorch_model-00001-of-00009.safetensors",
     "transformer_blocks.0.attn.norm_q.weight": "diffusion_pytorch_model-00001-of-00009.safetensors",
     "transformer_blocks.0.attn.to_k.weight": "diffusion_pytorch_model-00001-of-00009.safetensors",
@@ -53,7 +59,8 @@
     "transformer_blocks.0.ff.net.2.weight": "diffusion_pytorch_model-00001-of-00009.safetensors",
     "transformer_blocks.0.norm1.weight": "diffusion_pytorch_model-00001-of-00009.safetensors",
     "transformer_blocks.0.norm2.weight": "diffusion_pytorch_model-00001-of-00009.safetensors",
-    "transformer_blocks.1.adaln_proj.table": "diffusion_pytorch_model-00009-of-00009.safetensors",
+    "transformer_blocks.1.adaln_proj.linear.bias": "diffusion_pytorch_model-00001-of-00009.safetensors",
+    "transformer_blocks.1.adaln_proj.linear.weight": "diffusion_pytorch_model-00001-of-00009.safetensors",
     "transformer_blocks.1.attn.norm_k.weight": "diffusion_pytorch_model-00001-of-00009.safetensors",
     "transformer_blocks.1.attn.norm_q.weight": "diffusion_pytorch_model-00001-of-00009.safetensors",
     "transformer_blocks.1.attn.to_k.weight": "diffusion_pytorch_model-00001-of-00009.safetensors",
@@ -64,7 +71,8 @@
     "transformer_blocks.1.ff.net.2.weight": "diffusion_pytorch_model-00001-of-00009.safetensors",
     "transformer_blocks.1.norm1.weight": "diffusion_pytorch_model-00001-of-00009.safetensors",
     "transformer_blocks.1.norm2.weight": "diffusion_pytorch_model-00001-of-00009.safetensors",
-    "transformer_blocks.10.adaln_proj.table": "diffusion_pytorch_model-00009-of-00009.safetensors",
+    "transformer_blocks.10.adaln_proj.linear.bias": "diffusion_pytorch_model-00002-of-00009.safetensors",
+    "transformer_blocks.10.adaln_proj.linear.weight": "diffusion_pytorch_model-00002-of-00009.safetensors",
     "transformer_blocks.10.attn.norm_k.weight": "diffusion_pytorch_model-00002-of-00009.safetensors",
     "transformer_blocks.10.attn.norm_q.weight": "diffusion_pytorch_model-00002-of-00009.safetensors",
     "transformer_blocks.10.attn.to_k.weight": "diffusion_pytorch_model-00002-of-00009.safetensors",
@@ -75,7 +83,8 @@
     "transformer_blocks.10.ff.net.2.weight": "diffusion_pytorch_model-00002-of-00009.safetensors",
     "transformer_blocks.10.norm1.weight": "diffusion_pytorch_model-00002-of-00009.safetensors",
     "transformer_blocks.10.norm2.weight": "diffusion_pytorch_model-00002-of-00009.safetensors",
-    "transformer_blocks.11.adaln_proj.table": "diffusion_pytorch_model-00009-of-00009.safetensors",
+    "transformer_blocks.11.adaln_proj.linear.bias": "diffusion_pytorch_model-00002-of-00009.safetensors",
+    "transformer_blocks.11.adaln_proj.linear.weight": "diffusion_pytorch_model-00002-of-00009.safetensors",
     "transformer_blocks.11.attn.norm_k.weight": "diffusion_pytorch_model-00002-of-00009.safetensors",
     "transformer_blocks.11.attn.norm_q.weight": "diffusion_pytorch_model-00002-of-00009.safetensors",
     "transformer_blocks.11.attn.to_k.weight": "diffusion_pytorch_model-00002-of-00009.safetensors",
@@ -86,7 +95,8 @@
     "transformer_blocks.11.ff.net.2.weight": "diffusion_pytorch_model-00003-of-00009.safetensors",
     "transformer_blocks.11.norm1.weight": "diffusion_pytorch_model-00003-of-00009.safetensors",
     "transformer_blocks.11.norm2.weight": "diffusion_pytorch_model-00003-of-00009.safetensors",
-    "transformer_blocks.12.adaln_proj.table": "diffusion_pytorch_model-00009-of-00009.safetensors",
+    "transformer_blocks.12.adaln_proj.linear.bias": "diffusion_pytorch_model-00003-of-00009.safetensors",
+    "transformer_blocks.12.adaln_proj.linear.weight": "diffusion_pytorch_model-00003-of-00009.safetensors",
     "transformer_blocks.12.attn.norm_k.weight": "diffusion_pytorch_model-00003-of-00009.safetensors",
     "transformer_blocks.12.attn.norm_q.weight": "diffusion_pytorch_model-00003-of-00009.safetensors",
     "transformer_blocks.12.attn.to_k.weight": "diffusion_pytorch_model-00003-of-00009.safetensors",
@@ -97,7 +107,8 @@
     "transformer_blocks.12.ff.net.2.weight": "diffusion_pytorch_model-00003-of-00009.safetensors",
     "transformer_blocks.12.norm1.weight": "diffusion_pytorch_model-00003-of-00009.safetensors",
     "transformer_blocks.12.norm2.weight": "diffusion_pytorch_model-00003-of-00009.safetensors",
-    "transformer_blocks.13.adaln_proj.table": "diffusion_pytorch_model-00009-of-00009.safetensors",
+    "transformer_blocks.13.adaln_proj.linear.bias": "diffusion_pytorch_model-00003-of-00009.safetensors",
+    "transformer_blocks.13.adaln_proj.linear.weight": "diffusion_pytorch_model-00003-of-00009.safetensors",
     "transformer_blocks.13.attn.norm_k.weight": "diffusion_pytorch_model-00003-of-00009.safetensors",
     "transformer_blocks.13.attn.norm_q.weight": "diffusion_pytorch_model-00003-of-00009.safetensors",
     "transformer_blocks.13.attn.to_k.weight": "diffusion_pytorch_model-00003-of-00009.safetensors",
@@ -108,7 +119,8 @@
     "transformer_blocks.13.ff.net.2.weight": "diffusion_pytorch_model-00003-of-00009.safetensors",
     "transformer_blocks.13.norm1.weight": "diffusion_pytorch_model-00003-of-00009.safetensors",
     "transformer_blocks.13.norm2.weight": "diffusion_pytorch_model-00003-of-00009.safetensors",
-    "transformer_blocks.14.adaln_proj.table": "diffusion_pytorch_model-00009-of-00009.safetensors",
+    "transformer_blocks.14.adaln_proj.linear.bias": "diffusion_pytorch_model-00003-of-00009.safetensors",
+    "transformer_blocks.14.adaln_proj.linear.weight": "diffusion_pytorch_model-00003-of-00009.safetensors",
     "transformer_blocks.14.attn.norm_k.weight": "diffusion_pytorch_model-00003-of-00009.safetensors",
     "transformer_blocks.14.attn.norm_q.weight": "diffusion_pytorch_model-00003-of-00009.safetensors",
     "transformer_blocks.14.attn.to_k.weight": "diffusion_pytorch_model-00003-of-00009.safetensors",
@@ -119,7 +131,8 @@
     "transformer_blocks.14.ff.net.2.weight": "diffusion_pytorch_model-00003-of-00009.safetensors",
     "transformer_blocks.14.norm1.weight": "diffusion_pytorch_model-00003-of-00009.safetensors",
     "transformer_blocks.14.norm2.weight": "diffusion_pytorch_model-00003-of-00009.safetensors",
-    "transformer_blocks.15.adaln_proj.table": "diffusion_pytorch_model-00009-of-00009.safetensors",
+    "transformer_blocks.15.adaln_proj.linear.bias": "diffusion_pytorch_model-00003-of-00009.safetensors",
+    "transformer_blocks.15.adaln_proj.linear.weight": "diffusion_pytorch_model-00003-of-00009.safetensors",
     "transformer_blocks.15.attn.norm_k.weight": "diffusion_pytorch_model-00003-of-00009.safetensors",
     "transformer_blocks.15.attn.norm_q.weight": "diffusion_pytorch_model-00003-of-00009.safetensors",
     "transformer_blocks.15.attn.to_k.weight": "diffusion_pytorch_model-00003-of-00009.safetensors",
@@ -130,7 +143,8 @@
     "transformer_blocks.15.ff.net.2.weight": "diffusion_pytorch_model-00003-of-00009.safetensors",
     "transformer_blocks.15.norm1.weight": "diffusion_pytorch_model-00003-of-00009.safetensors",
     "transformer_blocks.15.norm2.weight": "diffusion_pytorch_model-00003-of-00009.safetensors",
-    "transformer_blocks.16.adaln_proj.table": "diffusion_pytorch_model-00009-of-00009.safetensors",
+    "transformer_blocks.16.adaln_proj.linear.bias": "diffusion_pytorch_model-00003-of-00009.safetensors",
+    "transformer_blocks.16.adaln_proj.linear.weight": "diffusion_pytorch_model-00003-of-00009.safetensors",
     "transformer_blocks.16.attn.norm_k.weight": "diffusion_pytorch_model-00003-of-00009.safetensors",
     "transformer_blocks.16.attn.norm_q.weight": "diffusion_pytorch_model-00003-of-00009.safetensors",
     "transformer_blocks.16.attn.to_k.weight": "diffusion_pytorch_model-00003-of-00009.safetensors",
@@ -141,7 +155,8 @@
     "transformer_blocks.16.ff.net.2.weight": "diffusion_pytorch_model-00003-of-00009.safetensors",
     "transformer_blocks.16.norm1.weight": "diffusion_pytorch_model-00003-of-00009.safetensors",
     "transformer_blocks.16.norm2.weight": "diffusion_pytorch_model-00003-of-00009.safetensors",
-    "transformer_blocks.17.adaln_proj.table": "diffusion_pytorch_model-00009-of-00009.safetensors",
+    "transformer_blocks.17.adaln_proj.linear.bias": "diffusion_pytorch_model-00003-of-00009.safetensors",
+    "transformer_blocks.17.adaln_proj.linear.weight": "diffusion_pytorch_model-00003-of-00009.safetensors",
     "transformer_blocks.17.attn.norm_k.weight": "diffusion_pytorch_model-00003-of-00009.safetensors",
     "transformer_blocks.17.attn.norm_q.weight": "diffusion_pytorch_model-00003-of-00009.safetensors",
     "transformer_blocks.17.attn.to_k.weight": "diffusion_pytorch_model-00003-of-00009.safetensors",
@@ -152,7 +167,8 @@
     "transformer_blocks.17.ff.net.2.weight": "diffusion_pytorch_model-00004-of-00009.safetensors",
     "transformer_blocks.17.norm1.weight": "diffusion_pytorch_model-00004-of-00009.safetensors",
     "transformer_blocks.17.norm2.weight": "diffusion_pytorch_model-00004-of-00009.safetensors",
-    "transformer_blocks.18.adaln_proj.table": "diffusion_pytorch_model-00009-of-00009.safetensors",
+    "transformer_blocks.18.adaln_proj.linear.bias": "diffusion_pytorch_model-00004-of-00009.safetensors",
+    "transformer_blocks.18.adaln_proj.linear.weight": "diffusion_pytorch_model-00004-of-00009.safetensors",
     "transformer_blocks.18.attn.norm_k.weight": "diffusion_pytorch_model-00004-of-00009.safetensors",
     "transformer_blocks.18.attn.norm_q.weight": "diffusion_pytorch_model-00004-of-00009.safetensors",
     "transformer_blocks.18.attn.to_k.weight": "diffusion_pytorch_model-00004-of-00009.safetensors",
@@ -163,7 +179,8 @@
     "transformer_blocks.18.ff.net.2.weight": "diffusion_pytorch_model-00004-of-00009.safetensors",
     "transformer_blocks.18.norm1.weight": "diffusion_pytorch_model-00004-of-00009.safetensors",
     "transformer_blocks.18.norm2.weight": "diffusion_pytorch_model-00004-of-00009.safetensors",
-    "transformer_blocks.19.adaln_proj.table": "diffusion_pytorch_model-00009-of-00009.safetensors",
+    "transformer_blocks.19.adaln_proj.linear.bias": "diffusion_pytorch_model-00004-of-00009.safetensors",
+    "transformer_blocks.19.adaln_proj.linear.weight": "diffusion_pytorch_model-00004-of-00009.safetensors",
     "transformer_blocks.19.attn.norm_k.weight": "diffusion_pytorch_model-00004-of-00009.safetensors",
     "transformer_blocks.19.attn.norm_q.weight": "diffusion_pytorch_model-00004-of-00009.safetensors",
     "transformer_blocks.19.attn.to_k.weight": "diffusion_pytorch_model-00004-of-00009.safetensors",
@@ -174,7 +191,8 @@
     "transformer_blocks.19.ff.net.2.weight": "diffusion_pytorch_model-00004-of-00009.safetensors",
     "transformer_blocks.19.norm1.weight": "diffusion_pytorch_model-00004-of-00009.safetensors",
     "transformer_blocks.19.norm2.weight": "diffusion_pytorch_model-00004-of-00009.safetensors",
-    "transformer_blocks.2.adaln_proj.table": "diffusion_pytorch_model-00009-of-00009.safetensors",
+    "transformer_blocks.2.adaln_proj.linear.bias": "diffusion_pytorch_model-00001-of-00009.safetensors",
+    "transformer_blocks.2.adaln_proj.linear.weight": "diffusion_pytorch_model-00001-of-00009.safetensors",
     "transformer_blocks.2.attn.norm_k.weight": "diffusion_pytorch_model-00001-of-00009.safetensors",
     "transformer_blocks.2.attn.norm_q.weight": "diffusion_pytorch_model-00001-of-00009.safetensors",
     "transformer_blocks.2.attn.to_k.weight": "diffusion_pytorch_model-00001-of-00009.safetensors",
@@ -185,7 +203,8 @@
     "transformer_blocks.2.ff.net.2.weight": "diffusion_pytorch_model-00001-of-00009.safetensors",
     "transformer_blocks.2.norm1.weight": "diffusion_pytorch_model-00001-of-00009.safetensors",
     "transformer_blocks.2.norm2.weight": "diffusion_pytorch_model-00001-of-00009.safetensors",
-    "transformer_blocks.20.adaln_proj.table": "diffusion_pytorch_model-00009-of-00009.safetensors",
+    "transformer_blocks.20.adaln_proj.linear.bias": "diffusion_pytorch_model-00004-of-00009.safetensors",
+    "transformer_blocks.20.adaln_proj.linear.weight": "diffusion_pytorch_model-00004-of-00009.safetensors",
     "transformer_blocks.20.attn.norm_k.weight": "diffusion_pytorch_model-00004-of-00009.safetensors",
     "transformer_blocks.20.attn.norm_q.weight": "diffusion_pytorch_model-00004-of-00009.safetensors",
     "transformer_blocks.20.attn.to_k.weight": "diffusion_pytorch_model-00004-of-00009.safetensors",
@@ -196,7 +215,8 @@
     "transformer_blocks.20.ff.net.2.weight": "diffusion_pytorch_model-00004-of-00009.safetensors",
     "transformer_blocks.20.norm1.weight": "diffusion_pytorch_model-00004-of-00009.safetensors",
     "transformer_blocks.20.norm2.weight": "diffusion_pytorch_model-00004-of-00009.safetensors",
-    "transformer_blocks.21.adaln_proj.table": "diffusion_pytorch_model-00009-of-00009.safetensors",
+    "transformer_blocks.21.adaln_proj.linear.bias": "diffusion_pytorch_model-00004-of-00009.safetensors",
+    "transformer_blocks.21.adaln_proj.linear.weight": "diffusion_pytorch_model-00004-of-00009.safetensors",
     "transformer_blocks.21.attn.norm_k.weight": "diffusion_pytorch_model-00004-of-00009.safetensors",
     "transformer_blocks.21.attn.norm_q.weight": "diffusion_pytorch_model-00004-of-00009.safetensors",
     "transformer_blocks.21.attn.to_k.weight": "diffusion_pytorch_model-00004-of-00009.safetensors",
@@ -207,7 +227,8 @@
     "transformer_blocks.21.ff.net.2.weight": "diffusion_pytorch_model-00004-of-00009.safetensors",
     "transformer_blocks.21.norm1.weight": "diffusion_pytorch_model-00004-of-00009.safetensors",
     "transformer_blocks.21.norm2.weight": "diffusion_pytorch_model-00004-of-00009.safetensors",
-    "transformer_blocks.22.adaln_proj.table": "diffusion_pytorch_model-00009-of-00009.safetensors",
+    "transformer_blocks.22.adaln_proj.linear.bias": "diffusion_pytorch_model-00004-of-00009.safetensors",
+    "transformer_blocks.22.adaln_proj.linear.weight": "diffusion_pytorch_model-00004-of-00009.safetensors",
     "transformer_blocks.22.attn.norm_k.weight": "diffusion_pytorch_model-00004-of-00009.safetensors",
     "transformer_blocks.22.attn.norm_q.weight": "diffusion_pytorch_model-00004-of-00009.safetensors",
     "transformer_blocks.22.attn.to_k.weight": "diffusion_pytorch_model-00004-of-00009.safetensors",
@@ -218,7 +239,8 @@
     "transformer_blocks.22.ff.net.2.weight": "diffusion_pytorch_model-00004-of-00009.safetensors",
     "transformer_blocks.22.norm1.weight": "diffusion_pytorch_model-00004-of-00009.safetensors",
     "transformer_blocks.22.norm2.weight": "diffusion_pytorch_model-00004-of-00009.safetensors",
-    "transformer_blocks.23.adaln_proj.table": "diffusion_pytorch_model-00009-of-00009.safetensors",
+    "transformer_blocks.23.adaln_proj.linear.bias": "diffusion_pytorch_model-00004-of-00009.safetensors",
+    "transformer_blocks.23.adaln_proj.linear.weight": "diffusion_pytorch_model-00004-of-00009.safetensors",
     "transformer_blocks.23.attn.norm_k.weight": "diffusion_pytorch_model-00004-of-00009.safetensors",
     "transformer_blocks.23.attn.norm_q.weight": "diffusion_pytorch_model-00004-of-00009.safetensors",
     "transformer_blocks.23.attn.to_k.weight": "diffusion_pytorch_model-00004-of-00009.safetensors",
@@ -229,7 +251,8 @@
     "transformer_blocks.23.ff.net.2.weight": "diffusion_pytorch_model-00005-of-00009.safetensors",
     "transformer_blocks.23.norm1.weight": "diffusion_pytorch_model-00005-of-00009.safetensors",
     "transformer_blocks.23.norm2.weight": "diffusion_pytorch_model-00005-of-00009.safetensors",
-    "transformer_blocks.24.adaln_proj.table": "diffusion_pytorch_model-00009-of-00009.safetensors",
+    "transformer_blocks.24.adaln_proj.linear.bias": "diffusion_pytorch_model-00005-of-00009.safetensors",
+    "transformer_blocks.24.adaln_proj.linear.weight": "diffusion_pytorch_model-00005-of-00009.safetensors",
     "transformer_blocks.24.attn.norm_k.weight": "diffusion_pytorch_model-00005-of-00009.safetensors",
     "transformer_blocks.24.attn.norm_q.weight": "diffusion_pytorch_model-00005-of-00009.safetensors",
     "transformer_blocks.24.attn.to_k.weight": "diffusion_pytorch_model-00005-of-00009.safetensors",
@@ -240,7 +263,8 @@
     "transformer_blocks.24.ff.net.2.weight": "diffusion_pytorch_model-00005-of-00009.safetensors",
     "transformer_blocks.24.norm1.weight": "diffusion_pytorch_model-00005-of-00009.safetensors",
     "transformer_blocks.24.norm2.weight": "diffusion_pytorch_model-00005-of-00009.safetensors",
-    "transformer_blocks.25.adaln_proj.table": "diffusion_pytorch_model-00009-of-00009.safetensors",
+    "transformer_blocks.25.adaln_proj.linear.bias": "diffusion_pytorch_model-00005-of-00009.safetensors",
+    "transformer_blocks.25.adaln_proj.linear.weight": "diffusion_pytorch_model-00005-of-00009.safetensors",
     "transformer_blocks.25.attn.norm_k.weight": "diffusion_pytorch_model-00005-of-00009.safetensors",
     "transformer_blocks.25.attn.norm_q.weight": "diffusion_pytorch_model-00005-of-00009.safetensors",
     "transformer_blocks.25.attn.to_k.weight": "diffusion_pytorch_model-00005-of-00009.safetensors",
@@ -251,7 +275,8 @@
     "transformer_blocks.25.ff.net.2.weight": "diffusion_pytorch_model-00005-of-00009.safetensors",
     "transformer_blocks.25.norm1.weight": "diffusion_pytorch_model-00005-of-00009.safetensors",
     "transformer_blocks.25.norm2.weight": "diffusion_pytorch_model-00005-of-00009.safetensors",
-    "transformer_blocks.26.adaln_proj.table": "diffusion_pytorch_model-00009-of-00009.safetensors",
+    "transformer_blocks.26.adaln_proj.linear.bias": "diffusion_pytorch_model-00005-of-00009.safetensors",
+    "transformer_blocks.26.adaln_proj.linear.weight": "diffusion_pytorch_model-00005-of-00009.safetensors",
     "transformer_blocks.26.attn.norm_k.weight": "diffusion_pytorch_model-00005-of-00009.safetensors",
     "transformer_blocks.26.attn.norm_q.weight": "diffusion_pytorch_model-00005-of-00009.safetensors",
     "transformer_blocks.26.attn.to_k.weight": "diffusion_pytorch_model-00005-of-00009.safetensors",
@@ -262,7 +287,8 @@
     "transformer_blocks.26.ff.net.2.weight": "diffusion_pytorch_model-00005-of-00009.safetensors",
     "transformer_blocks.26.norm1.weight": "diffusion_pytorch_model-00005-of-00009.safetensors",
     "transformer_blocks.26.norm2.weight": "diffusion_pytorch_model-00005-of-00009.safetensors",
-    "transformer_blocks.27.adaln_proj.table": "diffusion_pytorch_model-00009-of-00009.safetensors",
+    "transformer_blocks.27.adaln_proj.linear.bias": "diffusion_pytorch_model-00005-of-00009.safetensors",
+    "transformer_blocks.27.adaln_proj.linear.weight": "diffusion_pytorch_model-00005-of-00009.safetensors",
     "transformer_blocks.27.attn.norm_k.weight": "diffusion_pytorch_model-00005-of-00009.safetensors",
     "transformer_blocks.27.attn.norm_q.weight": "diffusion_pytorch_model-00005-of-00009.safetensors",
     "transformer_blocks.27.attn.to_k.weight": "diffusion_pytorch_model-00005-of-00009.safetensors",
@@ -273,7 +299,8 @@
     "transformer_blocks.27.ff.net.2.weight": "diffusion_pytorch_model-00005-of-00009.safetensors",
     "transformer_blocks.27.norm1.weight": "diffusion_pytorch_model-00005-of-00009.safetensors",
     "transformer_blocks.27.norm2.weight": "diffusion_pytorch_model-00005-of-00009.safetensors",
-    "transformer_blocks.28.adaln_proj.table": "diffusion_pytorch_model-00009-of-00009.safetensors",
+    "transformer_blocks.28.adaln_proj.linear.bias": "diffusion_pytorch_model-00005-of-00009.safetensors",
+    "transformer_blocks.28.adaln_proj.linear.weight": "diffusion_pytorch_model-00005-of-00009.safetensors",
     "transformer_blocks.28.attn.norm_k.weight": "diffusion_pytorch_model-00005-of-00009.safetensors",
     "transformer_blocks.28.attn.norm_q.weight": "diffusion_pytorch_model-00005-of-00009.safetensors",
     "transformer_blocks.28.attn.to_k.weight": "diffusion_pytorch_model-00005-of-00009.safetensors",
@@ -284,7 +311,8 @@
     "transformer_blocks.28.ff.net.2.weight": "diffusion_pytorch_model-00005-of-00009.safetensors",
     "transformer_blocks.28.norm1.weight": "diffusion_pytorch_model-00005-of-00009.safetensors",
     "transformer_blocks.28.norm2.weight": "diffusion_pytorch_model-00005-of-00009.safetensors",
-    "transformer_blocks.29.adaln_proj.table": "diffusion_pytorch_model-00009-of-00009.safetensors",
+    "transformer_blocks.29.adaln_proj.linear.bias": "diffusion_pytorch_model-00005-of-00009.safetensors",
+    "transformer_blocks.29.adaln_proj.linear.weight": "diffusion_pytorch_model-00005-of-00009.safetensors",
     "transformer_blocks.29.attn.norm_k.weight": "diffusion_pytorch_model-00005-of-00009.safetensors",
     "transformer_blocks.29.attn.norm_q.weight": "diffusion_pytorch_model-00005-of-00009.safetensors",
     "transformer_blocks.29.attn.to_k.weight": "diffusion_pytorch_model-00005-of-00009.safetensors",
@@ -295,7 +323,8 @@
     "transformer_blocks.29.ff.net.2.weight": "diffusion_pytorch_model-00006-of-00009.safetensors",
     "transformer_blocks.29.norm1.weight": "diffusion_pytorch_model-00006-of-00009.safetensors",
     "transformer_blocks.29.norm2.weight": "diffusion_pytorch_model-00006-of-00009.safetensors",
-    "transformer_blocks.3.adaln_proj.table": "diffusion_pytorch_model-00009-of-00009.safetensors",
+    "transformer_blocks.3.adaln_proj.linear.bias": "diffusion_pytorch_model-00001-of-00009.safetensors",
+    "transformer_blocks.3.adaln_proj.linear.weight": "diffusion_pytorch_model-00001-of-00009.safetensors",
     "transformer_blocks.3.attn.norm_k.weight": "diffusion_pytorch_model-00001-of-00009.safetensors",
     "transformer_blocks.3.attn.norm_q.weight": "diffusion_pytorch_model-00001-of-00009.safetensors",
     "transformer_blocks.3.attn.to_k.weight": "diffusion_pytorch_model-00001-of-00009.safetensors",
@@ -306,7 +335,8 @@
     "transformer_blocks.3.ff.net.2.weight": "diffusion_pytorch_model-00001-of-00009.safetensors",
     "transformer_blocks.3.norm1.weight": "diffusion_pytorch_model-00001-of-00009.safetensors",
     "transformer_blocks.3.norm2.weight": "diffusion_pytorch_model-00001-of-00009.safetensors",
-    "transformer_blocks.30.adaln_proj.table": "diffusion_pytorch_model-00009-of-00009.safetensors",
+    "transformer_blocks.30.adaln_proj.linear.bias": "diffusion_pytorch_model-00006-of-00009.safetensors",
+    "transformer_blocks.30.adaln_proj.linear.weight": "diffusion_pytorch_model-00006-of-00009.safetensors",
     "transformer_blocks.30.attn.norm_k.weight": "diffusion_pytorch_model-00006-of-00009.safetensors",
     "transformer_blocks.30.attn.norm_q.weight": "diffusion_pytorch_model-00006-of-00009.safetensors",
     "transformer_blocks.30.attn.to_k.weight": "diffusion_pytorch_model-00006-of-00009.safetensors",
@@ -317,7 +347,8 @@
     "transformer_blocks.30.ff.net.2.weight": "diffusion_pytorch_model-00006-of-00009.safetensors",
     "transformer_blocks.30.norm1.weight": "diffusion_pytorch_model-00006-of-00009.safetensors",
     "transformer_blocks.30.norm2.weight": "diffusion_pytorch_model-00006-of-00009.safetensors",
-    "transformer_blocks.31.adaln_proj.table": "diffusion_pytorch_model-00009-of-00009.safetensors",
+    "transformer_blocks.31.adaln_proj.linear.bias": "diffusion_pytorch_model-00006-of-00009.safetensors",
+    "transformer_blocks.31.adaln_proj.linear.weight": "diffusion_pytorch_model-00006-of-00009.safetensors",
     "transformer_blocks.31.attn.norm_k.weight": "diffusion_pytorch_model-00006-of-00009.safetensors",
     "transformer_blocks.31.attn.norm_q.weight": "diffusion_pytorch_model-00006-of-00009.safetensors",
     "transformer_blocks.31.attn.to_k.weight": "diffusion_pytorch_model-00006-of-00009.safetensors",
@@ -328,7 +359,8 @@
     "transformer_blocks.31.ff.net.2.weight": "diffusion_pytorch_model-00006-of-00009.safetensors",
     "transformer_blocks.31.norm1.weight": "diffusion_pytorch_model-00006-of-00009.safetensors",
     "transformer_blocks.31.norm2.weight": "diffusion_pytorch_model-00006-of-00009.safetensors",
-    "transformer_blocks.32.adaln_proj.table": "diffusion_pytorch_model-00009-of-00009.safetensors",
+    "transformer_blocks.32.adaln_proj.linear.bias": "diffusion_pytorch_model-00006-of-00009.safetensors",
+    "transformer_blocks.32.adaln_proj.linear.weight": "diffusion_pytorch_model-00006-of-00009.safetensors",
     "transformer_blocks.32.attn.norm_k.weight": "diffusion_pytorch_model-00006-of-00009.safetensors",
     "transformer_blocks.32.attn.norm_q.weight": "diffusion_pytorch_model-00006-of-00009.safetensors",
     "transformer_blocks.32.attn.to_k.weight": "diffusion_pytorch_model-00006-of-00009.safetensors",
@@ -339,7 +371,8 @@
     "transformer_blocks.32.ff.net.2.weight": "diffusion_pytorch_model-00006-of-00009.safetensors",
     "transformer_blocks.32.norm1.weight": "diffusion_pytorch_model-00006-of-00009.safetensors",
     "transformer_blocks.32.norm2.weight": "diffusion_pytorch_model-00006-of-00009.safetensors",
-    "transformer_blocks.33.adaln_proj.table": "diffusion_pytorch_model-00009-of-00009.safetensors",
+    "transformer_blocks.33.adaln_proj.linear.bias": "diffusion_pytorch_model-00006-of-00009.safetensors",
+    "transformer_blocks.33.adaln_proj.linear.weight": "diffusion_pytorch_model-00006-of-00009.safetensors",
     "transformer_blocks.33.attn.norm_k.weight": "diffusion_pytorch_model-00006-of-00009.safetensors",
     "transformer_blocks.33.attn.norm_q.weight": "diffusion_pytorch_model-00006-of-00009.safetensors",
     "transformer_blocks.33.attn.to_k.weight": "diffusion_pytorch_model-00006-of-00009.safetensors",
@@ -350,7 +383,8 @@
     "transformer_blocks.33.ff.net.2.weight": "diffusion_pytorch_model-00006-of-00009.safetensors",
     "transformer_blocks.33.norm1.weight": "diffusion_pytorch_model-00006-of-00009.safetensors",
     "transformer_blocks.33.norm2.weight": "diffusion_pytorch_model-00006-of-00009.safetensors",
-    "transformer_blocks.34.adaln_proj.table": "diffusion_pytorch_model-00009-of-00009.safetensors",
+    "transformer_blocks.34.adaln_proj.linear.bias": "diffusion_pytorch_model-00006-of-00009.safetensors",
+    "transformer_blocks.34.adaln_proj.linear.weight": "diffusion_pytorch_model-00006-of-00009.safetensors",
     "transformer_blocks.34.attn.norm_k.weight": "diffusion_pytorch_model-00006-of-00009.safetensors",
     "transformer_blocks.34.attn.norm_q.weight": "diffusion_pytorch_model-00006-of-00009.safetensors",
     "transformer_blocks.34.attn.to_k.weight": "diffusion_pytorch_model-00006-of-00009.safetensors",
@@ -361,7 +395,8 @@
     "transformer_blocks.34.ff.net.2.weight": "diffusion_pytorch_model-00006-of-00009.safetensors",
     "transformer_blocks.34.norm1.weight": "diffusion_pytorch_model-00006-of-00009.safetensors",
     "transformer_blocks.34.norm2.weight": "diffusion_pytorch_model-00006-of-00009.safetensors",
-    "transformer_blocks.35.adaln_proj.table": "diffusion_pytorch_model-00009-of-00009.safetensors",
+    "transformer_blocks.35.adaln_proj.linear.bias": "diffusion_pytorch_model-00006-of-00009.safetensors",
+    "transformer_blocks.35.adaln_proj.linear.weight": "diffusion_pytorch_model-00006-of-00009.safetensors",
     "transformer_blocks.35.attn.norm_k.weight": "diffusion_pytorch_model-00006-of-00009.safetensors",
     "transformer_blocks.35.attn.norm_q.weight": "diffusion_pytorch_model-00006-of-00009.safetensors",
     "transformer_blocks.35.attn.to_k.weight": "diffusion_pytorch_model-00006-of-00009.safetensors",
@@ -372,7 +407,8 @@
     "transformer_blocks.35.ff.net.2.weight": "diffusion_pytorch_model-00006-of-00009.safetensors",
     "transformer_blocks.35.norm1.weight": "diffusion_pytorch_model-00006-of-00009.safetensors",
     "transformer_blocks.35.norm2.weight": "diffusion_pytorch_model-00006-of-00009.safetensors",
-    "transformer_blocks.36.adaln_proj.table": "diffusion_pytorch_model-00009-of-00009.safetensors",
+    "transformer_blocks.36.adaln_proj.linear.bias": "diffusion_pytorch_model-00006-of-00009.safetensors",
+    "transformer_blocks.36.adaln_proj.linear.weight": "diffusion_pytorch_model-00006-of-00009.safetensors",
     "transformer_blocks.36.attn.norm_k.weight": "diffusion_pytorch_model-00006-of-00009.safetensors",
     "transformer_blocks.36.attn.norm_q.weight": "diffusion_pytorch_model-00006-of-00009.safetensors",
     "transformer_blocks.36.attn.to_k.weight": "diffusion_pytorch_model-00006-of-00009.safetensors",
@@ -383,7 +419,8 @@
     "transformer_blocks.36.ff.net.2.weight": "diffusion_pytorch_model-00007-of-00009.safetensors",
     "transformer_blocks.36.norm1.weight": "diffusion_pytorch_model-00007-of-00009.safetensors",
     "transformer_blocks.36.norm2.weight": "diffusion_pytorch_model-00007-of-00009.safetensors",
-    "transformer_blocks.37.adaln_proj.table": "diffusion_pytorch_model-00009-of-00009.safetensors",
+    "transformer_blocks.37.adaln_proj.linear.bias": "diffusion_pytorch_model-00007-of-00009.safetensors",
+    "transformer_blocks.37.adaln_proj.linear.weight": "diffusion_pytorch_model-00007-of-00009.safetensors",
     "transformer_blocks.37.attn.norm_k.weight": "diffusion_pytorch_model-00007-of-00009.safetensors",
     "transformer_blocks.37.attn.norm_q.weight": "diffusion_pytorch_model-00007-of-00009.safetensors",
     "transformer_blocks.37.attn.to_k.weight": "diffusion_pytorch_model-00007-of-00009.safetensors",
@@ -394,7 +431,8 @@
     "transformer_blocks.37.ff.net.2.weight": "diffusion_pytorch_model-00007-of-00009.safetensors",
     "transformer_blocks.37.norm1.weight": "diffusion_pytorch_model-00007-of-00009.safetensors",
     "transformer_blocks.37.norm2.weight": "diffusion_pytorch_model-00007-of-00009.safetensors",
-    "transformer_blocks.38.adaln_proj.table": "diffusion_pytorch_model-00009-of-00009.safetensors",
+    "transformer_blocks.38.adaln_proj.linear.bias": "diffusion_pytorch_model-00007-of-00009.safetensors",
+    "transformer_blocks.38.adaln_proj.linear.weight": "diffusion_pytorch_model-00007-of-00009.safetensors",
     "transformer_blocks.38.attn.norm_k.weight": "diffusion_pytorch_model-00007-of-00009.safetensors",
     "transformer_blocks.38.attn.norm_q.weight": "diffusion_pytorch_model-00007-of-00009.safetensors",
     "transformer_blocks.38.attn.to_k.weight": "diffusion_pytorch_model-00007-of-00009.safetensors",
@@ -405,7 +443,8 @@
     "transformer_blocks.38.ff.net.2.weight": "diffusion_pytorch_model-00007-of-00009.safetensors",
     "transformer_blocks.38.norm1.weight": "diffusion_pytorch_model-00007-of-00009.safetensors",
     "transformer_blocks.38.norm2.weight": "diffusion_pytorch_model-00007-of-00009.safetensors",
-    "transformer_blocks.39.adaln_proj.table": "diffusion_pytorch_model-00009-of-00009.safetensors",
+    "transformer_blocks.39.adaln_proj.linear.bias": "diffusion_pytorch_model-00007-of-00009.safetensors",
+    "transformer_blocks.39.adaln_proj.linear.weight": "diffusion_pytorch_model-00007-of-00009.safetensors",
     "transformer_blocks.39.attn.norm_k.weight": "diffusion_pytorch_model-00007-of-00009.safetensors",
     "transformer_blocks.39.attn.norm_q.weight": "diffusion_pytorch_model-00007-of-00009.safetensors",
     "transformer_blocks.39.attn.to_k.weight": "diffusion_pytorch_model-00007-of-00009.safetensors",
@@ -416,7 +455,8 @@
     "transformer_blocks.39.ff.net.2.weight": "diffusion_pytorch_model-00007-of-00009.safetensors",
     "transformer_blocks.39.norm1.weight": "diffusion_pytorch_model-00007-of-00009.safetensors",
     "transformer_blocks.39.norm2.weight": "diffusion_pytorch_model-00007-of-00009.safetensors",
-    "transformer_blocks.4.adaln_proj.table": "diffusion_pytorch_model-00009-of-00009.safetensors",
+    "transformer_blocks.4.adaln_proj.linear.bias": "diffusion_pytorch_model-00001-of-00009.safetensors",
+    "transformer_blocks.4.adaln_proj.linear.weight": "diffusion_pytorch_model-00001-of-00009.safetensors",
     "transformer_blocks.4.attn.norm_k.weight": "diffusion_pytorch_model-00001-of-00009.safetensors",
     "transformer_blocks.4.attn.norm_q.weight": "diffusion_pytorch_model-00001-of-00009.safetensors",
     "transformer_blocks.4.attn.to_k.weight": "diffusion_pytorch_model-00001-of-00009.safetensors",
@@ -427,7 +467,8 @@
     "transformer_blocks.4.ff.net.2.weight": "diffusion_pytorch_model-00002-of-00009.safetensors",
     "transformer_blocks.4.norm1.weight": "diffusion_pytorch_model-00002-of-00009.safetensors",
     "transformer_blocks.4.norm2.weight": "diffusion_pytorch_model-00002-of-00009.safetensors",
-    "transformer_blocks.40.adaln_proj.table": "diffusion_pytorch_model-00009-of-00009.safetensors",
+    "transformer_blocks.40.adaln_proj.linear.bias": "diffusion_pytorch_model-00007-of-00009.safetensors",
+    "transformer_blocks.40.adaln_proj.linear.weight": "diffusion_pytorch_model-00007-of-00009.safetensors",
     "transformer_blocks.40.attn.norm_k.weight": "diffusion_pytorch_model-00007-of-00009.safetensors",
     "transformer_blocks.40.attn.norm_q.weight": "diffusion_pytorch_model-00007-of-00009.safetensors",
     "transformer_blocks.40.attn.to_k.weight": "diffusion_pytorch_model-00007-of-00009.safetensors",
@@ -438,7 +479,8 @@
     "transformer_blocks.40.ff.net.2.weight": "diffusion_pytorch_model-00007-of-00009.safetensors",
     "transformer_blocks.40.norm1.weight": "diffusion_pytorch_model-00007-of-00009.safetensors",
     "transformer_blocks.40.norm2.weight": "diffusion_pytorch_model-00007-of-00009.safetensors",
-    "transformer_blocks.41.adaln_proj.table": "diffusion_pytorch_model-00009-of-00009.safetensors",
+    "transformer_blocks.41.adaln_proj.linear.bias": "diffusion_pytorch_model-00007-of-00009.safetensors",
+    "transformer_blocks.41.adaln_proj.linear.weight": "diffusion_pytorch_model-00007-of-00009.safetensors",
     "transformer_blocks.41.attn.norm_k.weight": "diffusion_pytorch_model-00007-of-00009.safetensors",
     "transformer_blocks.41.attn.norm_q.weight": "diffusion_pytorch_model-00007-of-00009.safetensors",
     "transformer_blocks.41.attn.to_k.weight": "diffusion_pytorch_model-00007-of-00009.safetensors",
@@ -449,7 +491,8 @@
     "transformer_blocks.41.ff.net.2.weight": "diffusion_pytorch_model-00007-of-00009.safetensors",
     "transformer_blocks.41.norm1.weight": "diffusion_pytorch_model-00007-of-00009.safetensors",
     "transformer_blocks.41.norm2.weight": "diffusion_pytorch_model-00007-of-00009.safetensors",
-    "transformer_blocks.42.adaln_proj.table": "diffusion_pytorch_model-00009-of-00009.safetensors",
+    "transformer_blocks.42.adaln_proj.linear.bias": "diffusion_pytorch_model-00007-of-00009.safetensors",
+    "transformer_blocks.42.adaln_proj.linear.weight": "diffusion_pytorch_model-00007-of-00009.safetensors",
     "transformer_blocks.42.attn.norm_k.weight": "diffusion_pytorch_model-00007-of-00009.safetensors",
     "transformer_blocks.42.attn.norm_q.weight": "diffusion_pytorch_model-00007-of-00009.safetensors",
     "transformer_blocks.42.attn.to_k.weight": "diffusion_pytorch_model-00007-of-00009.safetensors",
@@ -460,7 +503,8 @@
     "transformer_blocks.42.ff.net.2.weight": "diffusion_pytorch_model-00008-of-00009.safetensors",
     "transformer_blocks.42.norm1.weight": "diffusion_pytorch_model-00008-of-00009.safetensors",
     "transformer_blocks.42.norm2.weight": "diffusion_pytorch_model-00008-of-00009.safetensors",
-    "transformer_blocks.43.adaln_proj.table": "diffusion_pytorch_model-00009-of-00009.safetensors",
+    "transformer_blocks.43.adaln_proj.linear.bias": "diffusion_pytorch_model-00008-of-00009.safetensors",
+    "transformer_blocks.43.adaln_proj.linear.weight": "diffusion_pytorch_model-00008-of-00009.safetensors",
     "transformer_blocks.43.attn.norm_k.weight": "diffusion_pytorch_model-00008-of-00009.safetensors",
     "transformer_blocks.43.attn.norm_q.weight": "diffusion_pytorch_model-00008-of-00009.safetensors",
     "transformer_blocks.43.attn.to_k.weight": "diffusion_pytorch_model-00008-of-00009.safetensors",
@@ -471,7 +515,8 @@
     "transformer_blocks.43.ff.net.2.weight": "diffusion_pytorch_model-00008-of-00009.safetensors",
     "transformer_blocks.43.norm1.weight": "diffusion_pytorch_model-00008-of-00009.safetensors",
     "transformer_blocks.43.norm2.weight": "diffusion_pytorch_model-00008-of-00009.safetensors",
-    "transformer_blocks.44.adaln_proj.table": "diffusion_pytorch_model-00009-of-00009.safetensors",
+    "transformer_blocks.44.adaln_proj.linear.bias": "diffusion_pytorch_model-00008-of-00009.safetensors",
+    "transformer_blocks.44.adaln_proj.linear.weight": "diffusion_pytorch_model-00008-of-00009.safetensors",
     "transformer_blocks.44.attn.norm_k.weight": "diffusion_pytorch_model-00008-of-00009.safetensors",
     "transformer_blocks.44.attn.norm_q.weight": "diffusion_pytorch_model-00008-of-00009.safetensors",
     "transformer_blocks.44.attn.to_k.weight": "diffusion_pytorch_model-00008-of-00009.safetensors",
@@ -482,7 +527,8 @@
     "transformer_blocks.44.ff.net.2.weight": "diffusion_pytorch_model-00008-of-00009.safetensors",
     "transformer_blocks.44.norm1.weight": "diffusion_pytorch_model-00008-of-00009.safetensors",
     "transformer_blocks.44.norm2.weight": "diffusion_pytorch_model-00008-of-00009.safetensors",
-    "transformer_blocks.45.adaln_proj.table": "diffusion_pytorch_model-00009-of-00009.safetensors",
+    "transformer_blocks.45.adaln_proj.linear.bias": "diffusion_pytorch_model-00008-of-00009.safetensors",
+    "transformer_blocks.45.adaln_proj.linear.weight": "diffusion_pytorch_model-00008-of-00009.safetensors",
     "transformer_blocks.45.attn.norm_k.weight": "diffusion_pytorch_model-00008-of-00009.safetensors",
     "transformer_blocks.45.attn.norm_q.weight": "diffusion_pytorch_model-00008-of-00009.safetensors",
     "transformer_blocks.45.attn.to_k.weight": "diffusion_pytorch_model-00008-of-00009.safetensors",
@@ -493,7 +539,8 @@
     "transformer_blocks.45.ff.net.2.weight": "diffusion_pytorch_model-00008-of-00009.safetensors",
     "transformer_blocks.45.norm1.weight": "diffusion_pytorch_model-00008-of-00009.safetensors",
     "transformer_blocks.45.norm2.weight": "diffusion_pytorch_model-00008-of-00009.safetensors",
-    "transformer_blocks.46.adaln_proj.table": "diffusion_pytorch_model-00009-of-00009.safetensors",
+    "transformer_blocks.46.adaln_proj.linear.bias": "diffusion_pytorch_model-00008-of-00009.safetensors",
+    "transformer_blocks.46.adaln_proj.linear.weight": "diffusion_pytorch_model-00008-of-00009.safetensors",
     "transformer_blocks.46.attn.norm_k.weight": "diffusion_pytorch_model-00008-of-00009.safetensors",
     "transformer_blocks.46.attn.norm_q.weight": "diffusion_pytorch_model-00008-of-00009.safetensors",
     "transformer_blocks.46.attn.to_k.weight": "diffusion_pytorch_model-00008-of-00009.safetensors",
@@ -504,7 +551,8 @@
     "transformer_blocks.46.ff.net.2.weight": "diffusion_pytorch_model-00008-of-00009.safetensors",
     "transformer_blocks.46.norm1.weight": "diffusion_pytorch_model-00008-of-00009.safetensors",
     "transformer_blocks.46.norm2.weight": "diffusion_pytorch_model-00008-of-00009.safetensors",
-    "transformer_blocks.47.adaln_proj.table": "diffusion_pytorch_model-00009-of-00009.safetensors",
+    "transformer_blocks.47.adaln_proj.linear.bias": "diffusion_pytorch_model-00008-of-00009.safetensors",
+    "transformer_blocks.47.adaln_proj.linear.weight": "diffusion_pytorch_model-00008-of-00009.safetensors",
     "transformer_blocks.47.attn.norm_k.weight": "diffusion_pytorch_model-00008-of-00009.safetensors",
     "transformer_blocks.47.attn.norm_q.weight": "diffusion_pytorch_model-00008-of-00009.safetensors",
     "transformer_blocks.47.attn.to_k.weight": "diffusion_pytorch_model-00008-of-00009.safetensors",
@@ -515,7 +563,8 @@
     "transformer_blocks.47.ff.net.2.weight": "diffusion_pytorch_model-00008-of-00009.safetensors",
     "transformer_blocks.47.norm1.weight": "diffusion_pytorch_model-00008-of-00009.safetensors",
     "transformer_blocks.47.norm2.weight": "diffusion_pytorch_model-00008-of-00009.safetensors",
-    "transformer_blocks.48.adaln_proj.table": "diffusion_pytorch_model-00009-of-00009.safetensors",
+    "transformer_blocks.48.adaln_proj.linear.bias": "diffusion_pytorch_model-00008-of-00009.safetensors",
+    "transformer_blocks.48.adaln_proj.linear.weight": "diffusion_pytorch_model-00008-of-00009.safetensors",
     "transformer_blocks.48.attn.norm_k.weight": "diffusion_pytorch_model-00008-of-00009.safetensors",
     "transformer_blocks.48.attn.norm_q.weight": "diffusion_pytorch_model-00008-of-00009.safetensors",
     "transformer_blocks.48.attn.to_k.weight": "diffusion_pytorch_model-00008-of-00009.safetensors",
@@ -526,7 +575,8 @@
     "transformer_blocks.48.ff.net.2.weight": "diffusion_pytorch_model-00009-of-00009.safetensors",
     "transformer_blocks.48.norm1.weight": "diffusion_pytorch_model-00009-of-00009.safetensors",
     "transformer_blocks.48.norm2.weight": "diffusion_pytorch_model-00009-of-00009.safetensors",
-    "transformer_blocks.49.adaln_proj.table": "diffusion_pytorch_model-00009-of-00009.safetensors",
+    "transformer_blocks.49.adaln_proj.linear.bias": "diffusion_pytorch_model-00009-of-00009.safetensors",
+    "transformer_blocks.49.adaln_proj.linear.weight": "diffusion_pytorch_model-00009-of-00009.safetensors",
     "transformer_blocks.49.attn.norm_k.weight": "diffusion_pytorch_model-00009-of-00009.safetensors",
     "transformer_blocks.49.attn.norm_q.weight": "diffusion_pytorch_model-00009-of-00009.safetensors",
     "transformer_blocks.49.attn.to_k.weight": "diffusion_pytorch_model-00009-of-00009.safetensors",
@@ -537,7 +587,8 @@
     "transformer_blocks.49.ff.net.2.weight": "diffusion_pytorch_model-00009-of-00009.safetensors",
     "transformer_blocks.49.norm1.weight": "diffusion_pytorch_model-00009-of-00009.safetensors",
     "transformer_blocks.49.norm2.weight": "diffusion_pytorch_model-00009-of-00009.safetensors",
-    "transformer_blocks.5.adaln_proj.table": "diffusion_pytorch_model-00009-of-00009.safetensors",
+    "transformer_blocks.5.adaln_proj.linear.bias": "diffusion_pytorch_model-00002-of-00009.safetensors",
+    "transformer_blocks.5.adaln_proj.linear.weight": "diffusion_pytorch_model-00002-of-00009.safetensors",
     "transformer_blocks.5.attn.norm_k.weight": "diffusion_pytorch_model-00002-of-00009.safetensors",
     "transformer_blocks.5.attn.norm_q.weight": "diffusion_pytorch_model-00002-of-00009.safetensors",
     "transformer_blocks.5.attn.to_k.weight": "diffusion_pytorch_model-00002-of-00009.safetensors",
@@ -548,7 +599,8 @@
     "transformer_blocks.5.ff.net.2.weight": "diffusion_pytorch_model-00002-of-00009.safetensors",
     "transformer_blocks.5.norm1.weight": "diffusion_pytorch_model-00002-of-00009.safetensors",
     "transformer_blocks.5.norm2.weight": "diffusion_pytorch_model-00002-of-00009.safetensors",
-    "transformer_blocks.6.adaln_proj.table": "diffusion_pytorch_model-00009-of-00009.safetensors",
+    "transformer_blocks.6.adaln_proj.linear.bias": "diffusion_pytorch_model-00002-of-00009.safetensors",
+    "transformer_blocks.6.adaln_proj.linear.weight": "diffusion_pytorch_model-00002-of-00009.safetensors",
     "transformer_blocks.6.attn.norm_k.weight": "diffusion_pytorch_model-00002-of-00009.safetensors",
     "transformer_blocks.6.attn.norm_q.weight": "diffusion_pytorch_model-00002-of-00009.safetensors",
     "transformer_blocks.6.attn.to_k.weight": "diffusion_pytorch_model-00002-of-00009.safetensors",
@@ -559,7 +611,8 @@
     "transformer_blocks.6.ff.net.2.weight": "diffusion_pytorch_model-00002-of-00009.safetensors",
     "transformer_blocks.6.norm1.weight": "diffusion_pytorch_model-00002-of-00009.safetensors",
     "transformer_blocks.6.norm2.weight": "diffusion_pytorch_model-00002-of-00009.safetensors",
-    "transformer_blocks.7.adaln_proj.table": "diffusion_pytorch_model-00009-of-00009.safetensors",
+    "transformer_blocks.7.adaln_proj.linear.bias": "diffusion_pytorch_model-00002-of-00009.safetensors",
+    "transformer_blocks.7.adaln_proj.linear.weight": "diffusion_pytorch_model-00002-of-00009.safetensors",
     "transformer_blocks.7.attn.norm_k.weight": "diffusion_pytorch_model-00002-of-00009.safetensors",
     "transformer_blocks.7.attn.norm_q.weight": "diffusion_pytorch_model-00002-of-00009.safetensors",
     "transformer_blocks.7.attn.to_k.weight": "diffusion_pytorch_model-00002-of-00009.safetensors",
@@ -570,7 +623,8 @@
     "transformer_blocks.7.ff.net.2.weight": "diffusion_pytorch_model-00002-of-00009.safetensors",
     "transformer_blocks.7.norm1.weight": "diffusion_pytorch_model-00002-of-00009.safetensors",
     "transformer_blocks.7.norm2.weight": "diffusion_pytorch_model-00002-of-00009.safetensors",
-    "transformer_blocks.8.adaln_proj.table": "diffusion_pytorch_model-00009-of-00009.safetensors",
+    "transformer_blocks.8.adaln_proj.linear.bias": "diffusion_pytorch_model-00002-of-00009.safetensors",
+    "transformer_blocks.8.adaln_proj.linear.weight": "diffusion_pytorch_model-00002-of-00009.safetensors",
     "transformer_blocks.8.attn.norm_k.weight": "diffusion_pytorch_model-00002-of-00009.safetensors",
     "transformer_blocks.8.attn.norm_q.weight": "diffusion_pytorch_model-00002-of-00009.safetensors",
     "transformer_blocks.8.attn.to_k.weight": "diffusion_pytorch_model-00002-of-00009.safetensors",
@@ -581,7 +635,8 @@
     "transformer_blocks.8.ff.net.2.weight": "diffusion_pytorch_model-00002-of-00009.safetensors",
     "transformer_blocks.8.norm1.weight": "diffusion_pytorch_model-00002-of-00009.safetensors",
     "transformer_blocks.8.norm2.weight": "diffusion_pytorch_model-00002-of-00009.safetensors",
-    "transformer_blocks.9.adaln_proj.table": "diffusion_pytorch_model-00009-of-00009.safetensors",
+    "transformer_blocks.9.adaln_proj.linear.bias": "diffusion_pytorch_model-00002-of-00009.safetensors",
+    "transformer_blocks.9.adaln_proj.linear.weight": "diffusion_pytorch_model-00002-of-00009.safetensors",
     "transformer_blocks.9.attn.norm_k.weight": "diffusion_pytorch_model-00002-of-00009.safetensors",
     "transformer_blocks.9.attn.norm_q.weight": "diffusion_pytorch_model-00002-of-00009.safetensors",
     "transformer_blocks.9.attn.to_k.weight": "diffusion_pytorch_model-00002-of-00009.safetensors",

transformer/modeling_minimax_h3_image_pruned.py

@@ -1,308 +0,0 @@
-# Copyright 2026 WaveCut. Derived from diffusers' `transformer_minimax_h3_image.py`
-# (Copyright 2026 The MiniMax and HuggingFace Teams), Apache License 2.0.
-"""MiniMax-H3-Image transformer with tabulated AdaLN modulation.
-
-The released checkpoint samples on one fixed grid of six timesteps (Parallel Decoding Distillation keeps one output
-head per step), and the reference rows of an edit stay at t = 0.999. The input of every AdaLN projection is a function
-of the timestep alone, so the modulation of each block takes a fixed set of values. This model stores those values as
-tables and drops what computed them:
-
-* `transformer_blocks.*.adaln_proj.linear` (13.0B parameters, 26.0 GB in bfloat16), including the audio third of
-  every projection, which the image model never addresses;
-* `time_embedder` and `norm_out.linear`.
-
-Each table row holds exactly what the original projection returns for that timestep set, computed with the same
-modules and the same matrix shapes (one row for text-to-image steps, two rows for edit steps). Everything else in the
-forward pass is the diffusers implementation unchanged.
-
-The last block also skips the rows whose output nothing reads: the text rows and the reference rows of an edit. They
-still provide keys and values to the generated rows; their queries, attention output and feed-forward are not
-computed. While `skip_dead_rows` is on (the default), the velocity returned for reference rows is zero (the pipeline
-steps the generated rows only); `skip_dead_rows = False` restores the full last block.
-"""
-
-from __future__ import annotations
-
-import numpy as np
-import torch
-from torch import nn
-
-from diffusers import MiniMaxH3ImageTransformer2DModel
-from diffusers.configuration_utils import register_to_config
-from diffusers.models.attention_dispatch import dispatch_attention_fn
-from diffusers.models.modeling_outputs import Transformer2DModelOutput
-from diffusers.models.modeling_utils import get_parameter_dtype
-from diffusers.models.transformers.transformer_minimax_h3_image import _apply_rotary_emb
-
-# Modality tags of the packed sequence: 0 image, 1 text. MiniMax-H3 also has 2 = audio, which the image model never
-# assigns, so the tables keep two rows per timestep instead of three.
-NUM_TAGS = 2
-
-
-class TabulatedModulation(nn.Module):
-    """Replaces `MiniMaxH3ImageAdaLayerNormModulation`: returns the six modulation tensors of a timestep set.
-
-    `table` holds, for every supported timestep set, the projection output with the audio rows removed: `[t0 image,
-    t0 text, t1 image, t1 text, ...]`, `6 * hidden_size` wide. The block addresses these rows with
-    `timestep_indices * NUM_TAGS + token_tags`.
-    """
-
-    def __init__(self, rows: int, hidden_size: int):
-        super().__init__()
-        self.hidden_size = hidden_size
-        self.register_buffer("table", torch.zeros(rows, 6 * hidden_size), persistent=True)
-
-    def forward(self, rows: slice) -> tuple[torch.Tensor, ...]:
-        return self.table[rows].chunk(6, dim=-1)
-
-
-class TabulatedNormOut(nn.Module):
-    """Replaces `MiniMaxH3ImageAdaLayerNormOut`: the final RMSNorm with tabulated `shift, scale` rows."""
-
-    def __init__(self, hidden_size: int, rows: int, eps: float):
-        super().__init__()
-        self.norm = nn.RMSNorm(hidden_size, eps=eps)
-        self.register_buffer("table", torch.zeros(rows, 2 * hidden_size), persistent=True)
-
-    def forward(self, hidden_states: torch.Tensor, rows: slice, timestep_indices: torch.Tensor) -> torch.Tensor:
-        shift, scale = self.table[rows].chunk(2, dim=-1)
-        hidden_states = self.norm(hidden_states)
-        return hidden_states * (1.0 + scale.index_select(0, timestep_indices)) + shift.index_select(
-            0, timestep_indices
-        )
-
-
-def _live_last_block(block, hidden_states, rows, adaln_indices, rotary_emb, live_rows):
-    """`MiniMaxH3ImageTransformerBlock.forward` for the rows in `live_rows` only.
-
-    Every row still contributes its key and value. The result is `(batch, len(live_rows), hidden)`.
-    """
-    shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = block.adaln_proj(rows)
-
-    norm_hidden_states = block.norm1(hidden_states)
-    norm_hidden_states = norm_hidden_states * (1.0 + scale_msa.index_select(0, adaln_indices)) + shift_msa.index_select(
-        0, adaln_indices
-    )
-
-    attn = block.attn
-    if attn.fused_projections:
-        query, key, value = attn.to_qkv(norm_hidden_states).chunk(3, dim=-1)
-        query = query.index_select(1, live_rows)
-    else:
-        query = attn.to_q(norm_hidden_states.index_select(1, live_rows))
-        key = attn.to_k(norm_hidden_states)
-        value = attn.to_v(norm_hidden_states)
-    query = attn.norm_q(query.unflatten(-1, (-1, attn.head_dim)))
-    key = attn.norm_k(key.unflatten(-1, (-1, attn.head_dim)))
-    value = value.unflatten(-1, (-1, attn.head_dim))
-    cos, sin = rotary_emb
-    query = _apply_rotary_emb(query, cos.index_select(0, live_rows), sin.index_select(0, live_rows))
-    key = _apply_rotary_emb(key, cos, sin)
-    processor = attn.processor
-    attn_output = dispatch_attention_fn(
-        query,
-        key,
-        value,
-        attn_mask=None,
-        dropout_p=0.0,
-        is_causal=False,
-        backend=getattr(processor, "_attention_backend", None),
-        parallel_config=getattr(processor, "_parallel_config", None),
-    )
-    attn_output = attn_output.flatten(2, 3).type_as(query)
-    attn_output = attn.to_out[1](attn.to_out[0](attn_output))
-
-    live_indices = adaln_indices.index_select(0, live_rows)
-    hidden_states = hidden_states.index_select(1, live_rows) + gate_msa.index_select(0, live_indices) * attn_output
-    norm_hidden_states = block.norm2(hidden_states)
-    norm_hidden_states = norm_hidden_states * (1.0 + scale_mlp.index_select(0, live_indices)) + shift_mlp.index_select(
-        0, live_indices
-    )
-    return hidden_states + gate_mlp.index_select(0, live_indices) * block.ff(norm_hidden_states)
-
-
-class MiniMaxH3ImagePrunedTransformer2DModel(MiniMaxH3ImageTransformer2DModel):
-    r"""
-    [`MiniMaxH3ImageTransformer2DModel`] with the AdaLN modulation read from tables.
-
-    Additional config:
-        adaln_plans (`list[list[float]]`):
-            The supported timestep sets, each the `timestep` tensor of one transformer call: one value for a
-            text-to-image step, the step value and the reference value for an edit step. Table rows follow this
-            order.
-        reference_timestep (`float`, defaults to `0.999`):
-            The timestep the reference rows of an edit are held at.
-        auto_map (`dict`, *optional*):
-            Remote-code entry for `diffusers.AutoModel`; not used by the model.
-
-    Runtime switches (attributes):
-        skip_dead_rows (`bool`, defaults to `True`): run the last block and the output head for the generated rows
-            only.
-        validate_inputs (`bool`, defaults to `True`): check every call's `token_tags` (image or text only).
-    """
-
-    # `time_embedder` is gone; the remaining float32 modules are those of the original model.
-    _keep_in_fp32_modules = ["proj_in", "pdd_proj_out", "rope"]
-    _no_split_modules = ["MiniMaxH3ImageTransformerBlock", "MiniMaxH3ImageTokenRefinerBlock", "TabulatedNormOut"]
-    _tp_plan = None
-    _cp_plan = None
-
-    @register_to_config
-    def __init__(
-        self,
-        num_attention_heads: int = 56,
-        attention_head_dim: int = 128,
-        hidden_size: int = 5376,
-        num_layers: int = 50,
-        num_refiner_layers: int = 2,
-        ffn_dim: int = 14336,
-        in_channels: int = 24,
-        patch_size: tuple[int, int, int] = (1, 2, 2),
-        text_dim: int = 5120,
-        freq_dim: int = 256,
-        time_embed_hidden_dim: int = 5376,
-        time_embed_dim: int = 2688,
-        rope_freq_dim: int = 16,
-        rope_theta: float = 10000.0,
-        norm_eps: float = 1e-5,
-        qk_norm_eps: float = 1e-5,
-        final_norm_eps: float = 1e-5,
-        num_pdd_heads: int = 6,
-        adaln_plans: list[list[float]] | None = None,
-        reference_timestep: float = 0.999,
-        auto_map: dict | None = None,
-    ) -> None:
-        # `auto_map` is read by `diffusers.AutoModel(..., trust_remote_code=True)`; it is accepted here only so that
-        # loading the config does not warn about it.
-        super().__init__(
-            num_attention_heads=num_attention_heads,
-            attention_head_dim=attention_head_dim,
-            hidden_size=hidden_size,
-            num_layers=num_layers,
-            num_refiner_layers=num_refiner_layers,
-            ffn_dim=ffn_dim,
-            in_channels=in_channels,
-            patch_size=patch_size,
-            text_dim=text_dim,
-            freq_dim=freq_dim,
-            time_embed_hidden_dim=time_embed_hidden_dim,
-            time_embed_dim=time_embed_dim,
-            rope_freq_dim=rope_freq_dim,
-            rope_theta=rope_theta,
-            norm_eps=norm_eps,
-            qk_norm_eps=qk_norm_eps,
-            final_norm_eps=final_norm_eps,
-            num_pdd_heads=num_pdd_heads,
-        )
-        if not adaln_plans:
-            raise ValueError("`adaln_plans` is empty: this checkpoint has no modulation tables.")
-
-        del self.time_proj
-        del self.time_embedder
-
-        # Timestep set -> (block table rows, output table rows, index of the reference timestep in the set or -1).
-        self._plans: dict[tuple[float, ...], tuple[slice, slice, int]] = {}
-        block_row = out_row = 0
-        reference = float(np.float32(reference_timestep))
-        for plan in adaln_plans:
-            values = tuple(float(np.float32(t)) for t in plan)
-            if list(values) != sorted(set(values)):
-                raise ValueError(f"A timestep set must be strictly increasing, got {plan}.")
-            reference_slot = values.index(reference) if reference in values else -1
-            self._plans[values] = (
-                slice(block_row, block_row + NUM_TAGS * len(values)),
-                slice(out_row, out_row + len(values)),
-                reference_slot,
-            )
-            block_row += NUM_TAGS * len(values)
-            out_row += len(values)
-
-        for block in self.transformer_blocks:
-            block.adaln_proj = TabulatedModulation(block_row, hidden_size)
-        self.norm_out = TabulatedNormOut(hidden_size, out_row, final_norm_eps)
-
-        self.skip_dead_rows = True
-        self.validate_inputs = True
-
-    def _plan_of(self, timestep: torch.Tensor) -> tuple[slice, slice, int]:
-        values = tuple(timestep.detach().to(device="cpu", dtype=torch.float32).tolist())
-        plan = self._plans.get(values)
-        if plan is None:
-            supported = ", ".join(str(list(v)) for v in self._plans)
-            raise ValueError(
-                f"Timesteps {list(values)} are not in this checkpoint's modulation tables. MiniMax-H3-Image samples on "
-                f"a fixed six-step grid; the tables hold exactly its timestep sets: {supported}."
-            )
-        return plan
-
-    def forward(
-        self,
-        hidden_states: torch.Tensor,
-        encoder_hidden_states: torch.Tensor,
-        timestep: torch.Tensor,
-        timestep_indices: torch.Tensor,
-        token_tags: torch.Tensor,
-        position_ids: torch.Tensor,
-        image_indices: torch.Tensor,
-        text_indices: torch.Tensor,
-        pdd_step_index: int,
-        return_dict: bool = True,
-    ) -> Transformer2DModelOutput | tuple[torch.Tensor]:
-        if position_ids.ndim != 2 or position_ids.shape[-1] != 3:
-            raise ValueError(f"`position_ids` must be a `(seq_len, 3)` tensor, got {list(position_ids.shape)}.")
-        sequence_length = position_ids.shape[0]
-        if token_tags.shape != (sequence_length,) or timestep_indices.shape != (sequence_length,):
-            raise ValueError(
-                "`token_tags` and `timestep_indices` must both be `(seq_len,)` tensors matching `position_ids`, got "
-                f"{list(token_tags.shape)} and {list(timestep_indices.shape)} for seq_len={sequence_length}."
-            )
-        if self.validate_inputs and int(token_tags.max()) >= NUM_TAGS:
-            raise ValueError("MiniMax-H3-Image rows are image (0) or text (1); the audio tag (2) has no table rows.")
-        block_rows, out_rows, reference_slot = self._plan_of(timestep)
-
-        rotary_emb = self.rope(position_ids)
-
-        # The packed sequence, as in `MiniMaxH3ImageTransformer2DModel.forward`.
-        image_embeds = self.proj_in(hidden_states.to(get_parameter_dtype(self.proj_in)))
-        text_embeds = self.context_embedder(encoder_hidden_states.to(get_parameter_dtype(self.context_embedder)))
-        text_embeds = self.token_refiner(text_embeds)
-
-        hidden_states = text_embeds.new_zeros((text_embeds.shape[0], sequence_length, text_embeds.shape[-1]))
-        hidden_states = hidden_states.index_copy(1, text_indices, text_embeds)
-        hidden_states = hidden_states.index_copy(1, image_indices, image_embeds.to(text_embeds.dtype))
-
-        adaln_indices = timestep_indices * NUM_TAGS + token_tags
-
-        blocks = self.transformer_blocks
-        full_blocks = blocks[:-1] if self.skip_dead_rows else blocks
-        for block in full_blocks:
-            if torch.is_grad_enabled() and self.gradient_checkpointing:
-                hidden_states = self._gradient_checkpointing_func(
-                    block, hidden_states, block_rows, adaln_indices, rotary_emb
-                )
-            else:
-                hidden_states = block(hidden_states, block_rows, adaln_indices, rotary_emb)
-
-        proj_out = self.pdd_proj_out[pdd_step_index]
-        if not self.skip_dead_rows:
-            hidden_states = self.norm_out(hidden_states, out_rows, timestep_indices).to(get_parameter_dtype(proj_out))
-            output = proj_out(hidden_states).index_select(1, image_indices)
-        else:
-            # Generated rows: the image rows that are not held at the reference timestep. Reference rows come first.
-            if reference_slot < 0:
-                live_mask, live_rows = None, image_indices
-            else:
-                live_mask = timestep_indices.index_select(0, image_indices) != reference_slot
-                live_rows = image_indices[live_mask]
-            hidden_states = _live_last_block(blocks[-1], hidden_states, block_rows, adaln_indices, rotary_emb, live_rows)
-            hidden_states = self.norm_out(hidden_states, out_rows, timestep_indices.index_select(0, live_rows))
-            live_output = proj_out(hidden_states.to(get_parameter_dtype(proj_out)))
-            if live_mask is None:
-                output = live_output
-            else:
-                output = live_output.new_zeros((live_output.shape[0], image_indices.shape[0], live_output.shape[-1]))
-                output[:, live_mask] = live_output
-
-        if not return_dict:
-            return (output,)
-        return Transformer2DModelOutput(sample=output)