zjwfufu commited on
Commit
8f67bf6
·
0 Parent(s):

Initial commit

Browse files

Signed-off-by: zjwfufu <1174987765@qq.com>

.gitattributes ADDED
@@ -0,0 +1,41 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ *.7z filter=lfs diff=lfs merge=lfs -text
2
+ *.arrow filter=lfs diff=lfs merge=lfs -text
3
+ *.bin filter=lfs diff=lfs merge=lfs -text
4
+ *.bz2 filter=lfs diff=lfs merge=lfs -text
5
+ *.ckpt filter=lfs diff=lfs merge=lfs -text
6
+ *.ftz filter=lfs diff=lfs merge=lfs -text
7
+ *.gz filter=lfs diff=lfs merge=lfs -text
8
+ *.h5 filter=lfs diff=lfs merge=lfs -text
9
+ *.joblib filter=lfs diff=lfs merge=lfs -text
10
+ *.lfs.* filter=lfs diff=lfs merge=lfs -text
11
+ *.mlmodel filter=lfs diff=lfs merge=lfs -text
12
+ *.model filter=lfs diff=lfs merge=lfs -text
13
+ *.msgpack filter=lfs diff=lfs merge=lfs -text
14
+ *.npy filter=lfs diff=lfs merge=lfs -text
15
+ *.npz filter=lfs diff=lfs merge=lfs -text
16
+ *.onnx filter=lfs diff=lfs merge=lfs -text
17
+ *.ot filter=lfs diff=lfs merge=lfs -text
18
+ *.parquet filter=lfs diff=lfs merge=lfs -text
19
+ *.pb filter=lfs diff=lfs merge=lfs -text
20
+ *.pickle filter=lfs diff=lfs merge=lfs -text
21
+ *.pkl filter=lfs diff=lfs merge=lfs -text
22
+ *.pt filter=lfs diff=lfs merge=lfs -text
23
+ *.pth filter=lfs diff=lfs merge=lfs -text
24
+ *.rar filter=lfs diff=lfs merge=lfs -text
25
+ *.safetensors filter=lfs diff=lfs merge=lfs -text
26
+ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
27
+ *.tar.* filter=lfs diff=lfs merge=lfs -text
28
+ *.tar filter=lfs diff=lfs merge=lfs -text
29
+ *.tflite filter=lfs diff=lfs merge=lfs -text
30
+ *.tgz filter=lfs diff=lfs merge=lfs -text
31
+ *.wasm filter=lfs diff=lfs merge=lfs -text
32
+ *.xz filter=lfs diff=lfs merge=lfs -text
33
+ *.zip filter=lfs diff=lfs merge=lfs -text
34
+ *.zst filter=lfs diff=lfs merge=lfs -text
35
+ *tfevents* filter=lfs diff=lfs merge=lfs -text
36
+ examples/input.png filter=lfs diff=lfs merge=lfs -text
37
+ examples/sample.mp4 filter=lfs diff=lfs merge=lfs -text
38
+ examples/poster.png filter=lfs diff=lfs merge=lfs -text
39
+ examples/ref_dog.png filter=lfs diff=lfs merge=lfs -text
40
+ examples/ref_woman.png filter=lfs diff=lfs merge=lfs -text
41
+ examples/references.png filter=lfs diff=lfs merge=lfs -text
LICENSE-code ADDED
@@ -0,0 +1,201 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ Apache License
2
+ Version 2.0, January 2004
3
+ http://www.apache.org/licenses/
4
+
5
+ TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
6
+
7
+ 1. Definitions.
8
+
9
+ "License" shall mean the terms and conditions for use, reproduction,
10
+ and distribution as defined by Sections 1 through 9 of this document.
11
+
12
+ "Licensor" shall mean the copyright owner or entity authorized by
13
+ the copyright owner that is granting the License.
14
+
15
+ "Legal Entity" shall mean the union of the acting entity and all
16
+ other entities that control, are controlled by, or are under common
17
+ control with that entity. For the purposes of this definition,
18
+ "control" means (i) the power, direct or indirect, to cause the
19
+ direction or management of such entity, whether by contract or
20
+ otherwise, or (ii) ownership of fifty percent (50%) or more of the
21
+ outstanding shares, or (iii) beneficial ownership of such entity.
22
+
23
+ "You" (or "Your") shall mean an individual or Legal Entity
24
+ exercising permissions granted by this License.
25
+
26
+ "Source" form shall mean the preferred form for making modifications,
27
+ including but not limited to software source code, documentation
28
+ source, and configuration files.
29
+
30
+ "Object" form shall mean any form resulting from mechanical
31
+ transformation or translation of a Source form, including but
32
+ not limited to compiled object code, generated documentation,
33
+ and conversions to other media types.
34
+
35
+ "Work" shall mean the work of authorship, whether in Source or
36
+ Object form, made available under the License, as indicated by a
37
+ copyright notice that is included in or attached to the work
38
+ (an example is provided in the Appendix below).
39
+
40
+ "Derivative Works" shall mean any work, whether in Source or Object
41
+ form, that is based on (or derived from) the Work and for which the
42
+ editorial revisions, annotations, elaborations, or other modifications
43
+ represent, as a whole, an original work of authorship. For the purposes
44
+ of this License, Derivative Works shall not include works that remain
45
+ separable from, or merely link (or bind by name) to the interfaces of,
46
+ the Work and Derivative Works thereof.
47
+
48
+ "Contribution" shall mean any work of authorship, including
49
+ the original version of the Work and any modifications or additions
50
+ to that Work or Derivative Works thereof, that is intentionally
51
+ submitted to Licensor for inclusion in the Work by the copyright owner
52
+ or by an individual or Legal Entity authorized to submit on behalf of
53
+ the copyright owner. For the purposes of this definition, "submitted"
54
+ means any form of electronic, verbal, or written communication sent
55
+ to the Licensor or its representatives, including but not limited to
56
+ communication on electronic mailing lists, source code control systems,
57
+ and issue tracking systems that are managed by, or on behalf of, the
58
+ Licensor for the purpose of discussing and improving the Work, but
59
+ excluding communication that is conspicuously marked or otherwise
60
+ designated in writing by the copyright owner as "Not a Contribution."
61
+
62
+ "Contributor" shall mean Licensor and any individual or Legal Entity
63
+ on behalf of whom a Contribution has been received by Licensor and
64
+ subsequently incorporated within the Work.
65
+
66
+ 2. Grant of Copyright License. Subject to the terms and conditions of
67
+ this License, each Contributor hereby grants to You a perpetual,
68
+ worldwide, non-exclusive, no-charge, royalty-free, irrevocable
69
+ copyright license to reproduce, prepare Derivative Works of,
70
+ publicly display, publicly perform, sublicense, and distribute the
71
+ Work and such Derivative Works in Source or Object form.
72
+
73
+ 3. Grant of Patent License. Subject to the terms and conditions of
74
+ this License, each Contributor hereby grants to You a perpetual,
75
+ worldwide, non-exclusive, no-charge, royalty-free, irrevocable
76
+ (except as stated in this section) patent license to make, have made,
77
+ use, offer to sell, sell, import, and otherwise transfer the Work,
78
+ where such license applies only to those patent claims licensable
79
+ by such Contributor that are necessarily infringed by their
80
+ Contribution(s) alone or by combination of their Contribution(s)
81
+ with the Work to which such Contribution(s) was submitted. If You
82
+ institute patent litigation against any entity (including a
83
+ cross-claim or counterclaim in a lawsuit) alleging that the Work
84
+ or a Contribution incorporated within the Work constitutes direct
85
+ or contributory patent infringement, then any patent licenses
86
+ granted to You under this License for that Work shall terminate
87
+ as of the date such litigation is filed.
88
+
89
+ 4. Redistribution. You may reproduce and distribute copies of the
90
+ Work or Derivative Works thereof in any medium, with or without
91
+ modifications, and in Source or Object form, provided that You
92
+ meet the following conditions:
93
+
94
+ (a) You must give any other recipients of the Work or
95
+ Derivative Works a copy of this License; and
96
+
97
+ (b) You must cause any modified files to carry prominent notices
98
+ stating that You changed the files; and
99
+
100
+ (c) You must retain, in the Source form of any Derivative Works
101
+ that You distribute, all copyright, patent, trademark, and
102
+ attribution notices from the Source form of the Work,
103
+ excluding those notices that do not pertain to any part of
104
+ the Derivative Works; and
105
+
106
+ (d) If the Work includes a "NOTICE" text file as part of its
107
+ distribution, then any Derivative Works that You distribute must
108
+ include a readable copy of the attribution notices contained
109
+ within such NOTICE file, excluding those notices that do not
110
+ pertain to any part of the Derivative Works, in at least one
111
+ of the following places: within a NOTICE text file distributed
112
+ as part of the Derivative Works; within the Source form or
113
+ documentation, if provided along with the Derivative Works; or,
114
+ within a display generated by the Derivative Works, if and
115
+ wherever such third-party notices normally appear. The contents
116
+ of the NOTICE file are for informational purposes only and
117
+ do not modify the License. You may add Your own attribution
118
+ notices within Derivative Works that You distribute, alongside
119
+ or as an addendum to the NOTICE text from the Work, provided
120
+ that such additional attribution notices cannot be construed
121
+ as modifying the License.
122
+
123
+ You may add Your own copyright statement to Your modifications and
124
+ may provide additional or different license terms and conditions
125
+ for use, reproduction, or distribution of Your modifications, or
126
+ for any such Derivative Works as a whole, provided Your use,
127
+ reproduction, and distribution of the Work otherwise complies with
128
+ the conditions stated in this License.
129
+
130
+ 5. Submission of Contributions. Unless You explicitly state otherwise,
131
+ any Contribution intentionally submitted for inclusion in the Work
132
+ by You to the Licensor shall be under the terms and conditions of
133
+ this License, without any additional terms or conditions.
134
+ Notwithstanding the above, nothing herein shall supersede or modify
135
+ the terms of any separate license agreement you may have executed
136
+ with Licensor regarding such Contributions.
137
+
138
+ 6. Trademarks. This License does not grant permission to use the trade
139
+ names, trademarks, service marks, or product names of the Licensor,
140
+ except as required for reasonable and customary use in describing the
141
+ origin of the Work and reproducing the content of the NOTICE file.
142
+
143
+ 7. Disclaimer of Warranty. Unless required by applicable law or
144
+ agreed to in writing, Licensor provides the Work (and each
145
+ Contributor provides its Contributions) on an "AS IS" BASIS,
146
+ WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
147
+ implied, including, without limitation, any warranties or conditions
148
+ of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A
149
+ PARTICULAR PURPOSE. You are solely responsible for determining the
150
+ appropriateness of using or redistributing the Work and assume any
151
+ risks associated with Your exercise of permissions under this License.
152
+
153
+ 8. Limitation of Liability. In no event and under no legal theory,
154
+ whether in tort (including negligence), contract, or otherwise,
155
+ unless required by applicable law (such as deliberate and grossly
156
+ negligent acts) or agreed to in writing, shall any Contributor be
157
+ liable to You for damages, including any direct, indirect, special,
158
+ incidental, or consequential damages of any character arising as a
159
+ result of this License or out of the use or inability to use the
160
+ Work (including but not limited to damages for loss of goodwill,
161
+ work stoppage, computer failure or malfunction, or any and all
162
+ other commercial damages or losses), even if such Contributor
163
+ has been advised of the possibility of such damages.
164
+
165
+ 9. Accepting Warranty or Additional Liability. While redistributing
166
+ the Work or Derivative Works thereof, You may choose to offer,
167
+ and charge a fee for, acceptance of support, warranty, indemnity,
168
+ or other liability obligations and/or rights consistent with this
169
+ License. However, in accepting such obligations, You may act only
170
+ on Your own behalf and on Your sole responsibility, not on behalf
171
+ of any other Contributor, and only if You agree to indemnify,
172
+ defend, and hold each Contributor harmless for any liability
173
+ incurred by, or claims asserted against, such Contributor by reason
174
+ of your accepting any such warranty or additional liability.
175
+
176
+ END OF TERMS AND CONDITIONS
177
+
178
+ APPENDIX: How to apply the Apache License to your work.
179
+
180
+ To apply the Apache License to your work, attach the following
181
+ boilerplate notice, with the fields enclosed by brackets "[]"
182
+ replaced with your own identifying information. (Don't include
183
+ the brackets!) The text should be enclosed in the appropriate
184
+ comment syntax for the file format. We also recommend that a
185
+ file or class name and description of purpose be included on the
186
+ same "printed page" as the copyright notice for easier
187
+ identification within third-party archives.
188
+
189
+ Copyright [yyyy] [name of copyright owner]
190
+
191
+ Licensed under the Apache License, Version 2.0 (the "License");
192
+ you may not use this file except in compliance with the License.
193
+ You may obtain a copy of the License at
194
+
195
+ http://www.apache.org/licenses/LICENSE-2.0
196
+
197
+ Unless required by applicable law or agreed to in writing, software
198
+ distributed under the License is distributed on an "AS IS" BASIS,
199
+ WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
200
+ See the License for the specific language governing permissions and
201
+ limitations under the License.
README.md ADDED
@@ -0,0 +1,86 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ ---
2
+ language:
3
+ - en
4
+ tags:
5
+ - smd
6
+ - diffusion-distillation
7
+ - few-step
8
+ - lora
9
+ base_model:
10
+ - MiniMaxAI/MiniMax-H3
11
+ base_model_relation: adapter
12
+ pipeline_tag: image-to-video
13
+ license: other
14
+ license_name: minimax-h3-community-license-agreement
15
+ license_link: https://huggingface.co/MiniMaxAI/MiniMax-H3/blob/main/LICENSE
16
+ ---
17
+
18
+ # SMD-MiniMax-H3-Ref2VA-4step
19
+
20
+ **4-step SMD for image-guided video and audio generation.** The student adapters, standalone inference code and reproducible examples are included. Upstream base components are downloaded separately.
21
+
22
+ [Training code](https://github.com/zjwfufu/SMD) · [SMD collection](https://huggingface.co/collections/zjwfufu/smd-6ac5f4b53c26548a75672189)
23
+
24
+ ## Example
25
+
26
+ A woman and a Samoyed, supplied as separate references, share a playful moment on a snowy mountain overlook.
27
+
28
+ <table>
29
+ <tr><th width="50%" align="center">Reference subjects</th><th width="50%" align="center">Generated video · 4 steps</th></tr>
30
+ <tr>
31
+ <td width="50%" align="center" valign="middle"><img src="https://huggingface.co/zjwfufu/SMD-MiniMax-H3-Ref2VA-4step/resolve/main/examples/references.png" alt="Person reference and dog reference" width="100%"></td>
32
+ <td width="50%" align="center" valign="middle"><video controls playsinline poster="https://huggingface.co/zjwfufu/SMD-MiniMax-H3-Ref2VA-4step/resolve/main/examples/poster.png" src="https://huggingface.co/zjwfufu/SMD-MiniMax-H3-Ref2VA-4step/resolve/main/examples/sample.mp4" width="100%"></video></td>
33
+ </tr>
34
+ </table>
35
+
36
+ [Person reference](examples/ref_woman.png) · [Dog reference](examples/ref_dog.png)
37
+
38
+ The references define the subjects' appearance; the model generates their new poses, shared setting, interaction and sound. They are not first-frame anchors.
39
+
40
+ [Download video with audio](examples/sample.mp4)
41
+
42
+ **Seed:** 42. **Size:** 1344 × 768. **Frames:** 124 at 24 fps. **Steps:** 4.
43
+
44
+ [Full prompt and case settings](examples/case.json) · [Generation receipt](examples/generation.json)
45
+
46
+ The separate reference images were generated with [SMD-FLUX.2-Klein-4B](https://huggingface.co/zjwfufu/SMD-FLUX.2-Klein-4B), using the `teacher-9b` variant. [Reference generation settings](examples/reference_generation.json).
47
+
48
+ ## Models
49
+
50
+ | Variant | Weights | Steps | Size |
51
+ |---|---|---:|---:|
52
+ | default | LoRA | 4 | 1.38 GB |
53
+
54
+ ## Run
55
+
56
+ For a native ComfyUI workflow on a single GPU, use the corresponding [SMD-MiniMax-H3-Ref2VA-4step-ComfyUI](https://huggingface.co/zjwfufu/SMD-MiniMax-H3-Ref2VA-4step-ComfyUI). It includes a Compact profile and a BF16 transformer profile. The Python runner below uses two GPUs.
57
+
58
+ Base: [MiniMaxAI/MiniMax-H3](https://huggingface.co/MiniMaxAI/MiniMax-H3).
59
+
60
+ Use Linux, Python 3.12, Git and two NVIDIA GPUs with CUDA support. The requirements install Diffusers from a pinned Git revision. The included inference code does not import the SMD training repository.
61
+
62
+ ```bash
63
+ python3.12 -m venv .venv
64
+ source .venv/bin/activate
65
+ python -m pip install huggingface-hub
66
+ # Sign in to Hugging Face.
67
+ hf auth login
68
+ hf download zjwfufu/SMD-MiniMax-H3-Ref2VA-4step --local-dir SMD-MiniMax-H3-Ref2VA-4step
69
+ cd SMD-MiniMax-H3-Ref2VA-4step
70
+ python -m pip install -r requirements.txt
71
+ python infer.py --output outputs
72
+ ```
73
+
74
+ The runner requires two CUDA GPUs and launches two processes with FSDP2. It was tested on **2 × A100 80 GB**, with 433 GiB of host RAM. It downloads only the selected transformer plus shared encoders, decoders and processor files to the Hugging Face cache. The selected base components occupy approximately **144 GB (134 GiB) per workflow**; keeping both workflows reuses the shared components and requires approximately **210 GB (196 GiB)** in total. Loading and fusing these weights can take longer than the sampling steps. `--base` selects a complete local upstream snapshot, and `--case` selects another case JSON.
75
+
76
+ ## Sampling
77
+
78
+ The released sampler uses raw final-segment scale **0.5**. The raw grid is built before the model-specific time shift. The included runner follows `sampling.json`; changing only an upstream pipeline's step count does not reproduce the student.
79
+
80
+ Video shift is 12; audio shift is 3. LoRA scale is 1.0. The default case uses 124 frames at 24 fps and joint video/audio generation.
81
+
82
+ The supplied loader validates and applies all LoRA matrices and saved replacement tensors, with the recorded merge precision. Omitting replacement tensors or using a different base can change the result.
83
+
84
+ ## License and attribution
85
+
86
+ Base components retain the terms described in the [MiniMaxAI/MiniMax-H3 model card](https://huggingface.co/MiniMaxAI/MiniMax-H3). Original code notices and bundled licenses are retained. Adapter use remains subject to the applicable base-model terms.
examples/case.json ADDED
@@ -0,0 +1,22 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "examples": [
3
+ {
4
+ "task": "ref2va",
5
+ "prompt": "subject_definitions:\n<Subject 1> is the adult woman in <Picture 1>, with a short dark bob, mustard yellow raincoat, dark trousers and brown boots. This image supplies her identity and clothing only.\n<Subject 2> is the white Samoyed in <Picture 2>, with fluffy fur, pointed ears, dark nose and a red bandana. This image supplies the dog's appearance only.\n\nsummary:\n[reference generation] <Subject 1> and <Subject 2> share a playful moment on a new snowy mountain overlook. Their separate studio backgrounds and standing poses do not define the target composition.\n\nretention_analysis:\n<Subject 1> (appears in [Shot 1]): fully_preserved - retain her face, short dark hair, yellow raincoat, trousers and boots while placing her in the new outdoor scene.\n<Subject 2> (appears in [Shot 1]): fully_preserved - retain the Samoyed's fluffy white fur, pointed ears, dark nose and red bandana while generating a new pose and action.\n\ndetailed_description:\nThe target video is a live-action outdoor travel scene with crisp winter daylight and gentle natural contrast.\n[Shot 1] A medium-wide view opens on a snowy mountain overlook. <Subject 1>, the adult woman with the short dark bob and mustard yellow raincoat from <Picture 1>, kneels on the left side of the frame beside <Subject 2>, the white Samoyed with the red bandana from <Picture 2>. Both subjects are already together in this new environment when the shot begins. The woman's brown boots press shallow tracks into the snow. The dog sits on the right, its thick white coat clearly separated from the snow by soft gray shadows and warm sunlight along its fur. Behind them, a low wooden trail marker has the three large dark painted letters \"SMD\" across its face. The marker is unobstructed near the left edge of the frame, with a distant blue mountain ridge and scattered dark fir trees beyond it. The camera holds a static shot at the dog's eye level throughout the five-second scene. During the first second, the dog looks toward the woman with its ears raised. She smiles and slowly brings her right hand to the top of its head. In the middle of the shot, she gently strokes the fur once from the crown toward the back of its neck. The hand stays above the dog's head and moves with light contact rather than gripping. The dog tilts its muzzle slightly toward her, then gives one short friendly bark with a visible opening and closing of its mouth. The woman responds with a quiet natural laugh and keeps her hand resting lightly on its neck. A faint breeze moves loose strands of her dark hair and lifts one corner of the dog's red bandana. At the end, both remain beside each other in the same positions while the dog turns its gaze back toward the camera. Snow, trees and the trail marker remain stationary; the only background motion is a little powder blowing close to the ground. There is no studio backdrop or compositional match to either source photograph.\n\noverall_soundscape:\nA light open-air breeze continues beneath the dog's single friendly bark. The woman's brief soft laugh follows the bark, accompanied by faint clothing rustle and a small crunch of snow beneath her knee.\n\nnon_diegetic_music:\nN/A",
6
+ "seed": 42,
7
+ "width": 1344,
8
+ "height": 768,
9
+ "num_frames": 124,
10
+ "references": [
11
+ {
12
+ "kind": "image",
13
+ "path": "ref_woman.png"
14
+ },
15
+ {
16
+ "kind": "image",
17
+ "path": "ref_dog.png"
18
+ }
19
+ ]
20
+ }
21
+ ]
22
+ }
examples/generation.json ADDED
@@ -0,0 +1,46 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "task": "ref2va",
3
+ "prompt": "subject_definitions:\n<Subject 1> is the adult woman in <Picture 1>, with a short dark bob, mustard yellow raincoat, dark trousers and brown boots. This image supplies her identity and clothing only.\n<Subject 2> is the white Samoyed in <Picture 2>, with fluffy fur, pointed ears, dark nose and a red bandana. This image supplies the dog's appearance only.\n\nsummary:\n[reference generation] <Subject 1> and <Subject 2> share a playful moment on a new snowy mountain overlook. Their separate studio backgrounds and standing poses do not define the target composition.\n\nretention_analysis:\n<Subject 1> (appears in [Shot 1]): fully_preserved - retain her face, short dark hair, yellow raincoat, trousers and boots while placing her in the new outdoor scene.\n<Subject 2> (appears in [Shot 1]): fully_preserved - retain the Samoyed's fluffy white fur, pointed ears, dark nose and red bandana while generating a new pose and action.\n\ndetailed_description:\nThe target video is a live-action outdoor travel scene with crisp winter daylight and gentle natural contrast.\n[Shot 1] A medium-wide view opens on a snowy mountain overlook. <Subject 1>, the adult woman with the short dark bob and mustard yellow raincoat from <Picture 1>, kneels on the left side of the frame beside <Subject 2>, the white Samoyed with the red bandana from <Picture 2>. Both subjects are already together in this new environment when the shot begins. The woman's brown boots press shallow tracks into the snow. The dog sits on the right, its thick white coat clearly separated from the snow by soft gray shadows and warm sunlight along its fur. Behind them, a low wooden trail marker has the three large dark painted letters \"SMD\" across its face. The marker is unobstructed near the left edge of the frame, with a distant blue mountain ridge and scattered dark fir trees beyond it. The camera holds a static shot at the dog's eye level throughout the five-second scene. During the first second, the dog looks toward the woman with its ears raised. She smiles and slowly brings her right hand to the top of its head. In the middle of the shot, she gently strokes the fur once from the crown toward the back of its neck. The hand stays above the dog's head and moves with light contact rather than gripping. The dog tilts its muzzle slightly toward her, then gives one short friendly bark with a visible opening and closing of its mouth. The woman responds with a quiet natural laugh and keeps her hand resting lightly on its neck. A faint breeze moves loose strands of her dark hair and lifts one corner of the dog's red bandana. At the end, both remain beside each other in the same positions while the dog turns its gaze back toward the camera. Snow, trees and the trail marker remain stationary; the only background motion is a little powder blowing close to the ground. There is no studio backdrop or compositional match to either source photograph.\n\noverall_soundscape:\nA light open-air breeze continues beneath the dog's single friendly bark. The woman's brief soft laugh follows the bark, accompanied by faint clothing rustle and a small crunch of snow beneath her knee.\n\nnon_diegetic_music:\nN/A",
4
+ "seed": 42,
5
+ "width": 1344,
6
+ "height": 768,
7
+ "num_frames": 124,
8
+ "fps": 24,
9
+ "nfe": 4,
10
+ "ff_chunk_size": 1024,
11
+ "qk_norm_chunk_size": 2048,
12
+ "low_peak_rope": true,
13
+ "explicit_transformer_offload": true,
14
+ "final_step_size_scale": 0.5,
15
+ "raw_sigmas": [
16
+ 1.0,
17
+ 0.7142857142857143,
18
+ 0.4285714285714286,
19
+ 0.1428571428571429,
20
+ 0.0
21
+ ],
22
+ "video_shift": 12.0,
23
+ "audio_shift": 3.0,
24
+ "base_model_id": "MiniMaxAI/MiniMax-H3",
25
+ "model_revision": "939557dc319dd91227e30195a763f272ba7f8765",
26
+ "lora_sha256": "b8cd339e167ab29e112d4f8934d3ac7f20b7e3269895fb149761bd959336bdd7",
27
+ "lora_scale": 1.0,
28
+ "lora_fused": true,
29
+ "reference_image_short_edge": 512,
30
+ "output": "sample.mp4",
31
+ "output_sha256": "6d9fda48b9b5b5cc331b22f5aa4994c12d20a3a110759f75f111561e70a8beb6",
32
+ "diffusers_revision": "4e0466f3e5260f0d78b5e2b68ffbf27d819cc6db",
33
+ "hardware": "2 x NVIDIA A100 80GB",
34
+ "reference_inputs": [
35
+ {
36
+ "kind": "image",
37
+ "path": "ref_woman.png",
38
+ "sha256": "08bdadae2f0ce2cb39b98153a6173ec360378b73b389276e2acfaac7f37971b3"
39
+ },
40
+ {
41
+ "kind": "image",
42
+ "path": "ref_dog.png",
43
+ "sha256": "6b16db993ff16833d48012234d811ae9712a358da1c8b3ac6bc05124ddc49c71"
44
+ }
45
+ ]
46
+ }
examples/layout.json ADDED
@@ -0,0 +1,9 @@
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "operation": "Proportional resize and padding for display; no cropping or retouching.",
3
+ "preview": "references.png",
4
+ "preview_sha256": "41655fdce97dc1fbcb1187f4b6209136a10d42e7440f7fef20f2767e09399d44",
5
+ "sources": {
6
+ "ref_woman.png": "08bdadae2f0ce2cb39b98153a6173ec360378b73b389276e2acfaac7f37971b3",
7
+ "ref_dog.png": "6b16db993ff16833d48012234d811ae9712a358da1c8b3ac6bc05124ddc49c71"
8
+ }
9
+ }
examples/poster.png ADDED

Git LFS Details

  • SHA256: b3453e3650ec4ed37e8891b4a80bc5c1b2185b3aa99ca9c53ff8a489780e546d
  • Pointer size: 131 Bytes
  • Size of remote file: 730 kB
examples/ref_dog.png ADDED

Git LFS Details

  • SHA256: 6b16db993ff16833d48012234d811ae9712a358da1c8b3ac6bc05124ddc49c71
  • Pointer size: 131 Bytes
  • Size of remote file: 811 kB
examples/ref_woman.png ADDED

Git LFS Details

  • SHA256: 08bdadae2f0ce2cb39b98153a6173ec360378b73b389276e2acfaac7f37971b3
  • Pointer size: 131 Bytes
  • Size of remote file: 799 kB
examples/reference_generation.json ADDED
@@ -0,0 +1,38 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ [
2
+ {
3
+ "prompt": "Full-length studio photograph of an adult woman with a short dark bob, wearing a mustard yellow raincoat, dark trousers and brown boots. Relaxed standing pose with empty hands. Plain light gray background, soft even lighting.",
4
+ "seed": 1234,
5
+ "height": 1024,
6
+ "width": 1024,
7
+ "sigmas": [
8
+ 0.999,
9
+ 0.9320267290061409,
10
+ 0.8044392117499536,
11
+ 0.4775642604809712,
12
+ 0.0
13
+ ],
14
+ "output": "ref_woman.png",
15
+ "output_sha256": "08bdadae2f0ce2cb39b98153a6173ec360378b73b389276e2acfaac7f37971b3",
16
+ "model_id": "zjwfufu/SMD-FLUX.2-Klein-4B",
17
+ "variant": "teacher-9b",
18
+ "steps": 4
19
+ },
20
+ {
21
+ "prompt": "Full-body studio photograph of a fluffy white Samoyed dog with a red bandana, standing in a relaxed pose, mouth slightly open. Isolated on a plain light gray background, soft even light.",
22
+ "seed": 1234,
23
+ "height": 1024,
24
+ "width": 1024,
25
+ "sigmas": [
26
+ 0.999,
27
+ 0.9320267290061409,
28
+ 0.8044392117499536,
29
+ 0.4775642604809712,
30
+ 0.0
31
+ ],
32
+ "output": "ref_dog.png",
33
+ "output_sha256": "6b16db993ff16833d48012234d811ae9712a358da1c8b3ac6bc05124ddc49c71",
34
+ "model_id": "zjwfufu/SMD-FLUX.2-Klein-4B",
35
+ "variant": "teacher-9b",
36
+ "steps": 4
37
+ }
38
+ ]
examples/references.png ADDED

Git LFS Details

  • SHA256: 41655fdce97dc1fbcb1187f4b6209136a10d42e7440f7fef20f2767e09399d44
  • Pointer size: 131 Bytes
  • Size of remote file: 681 kB
examples/sample.mp4 ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:6d9fda48b9b5b5cc331b22f5aa4994c12d20a3a110759f75f111561e70a8beb6
3
+ size 688999
h3_lora.py ADDED
@@ -0,0 +1,117 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # SPDX-License-Identifier: Apache-2.0
2
+ """Strict inference-only loading for the released H3 adapters."""
3
+ from __future__ import annotations
4
+ from collections.abc import Mapping
5
+ from pathlib import Path
6
+ import torch
7
+ from peft import LoraConfig
8
+ from safetensors import safe_open
9
+ from safetensors.torch import load_file
10
+
11
+
12
+ H3_LORA_TARGET_MODULES = (
13
+ "to_q",
14
+ "to_k",
15
+ "to_v",
16
+ "to_out.0",
17
+ "ff.net.0.proj",
18
+ "ff.net.2",
19
+ )
20
+
21
+ LORA_A_SUFFIX = ".lora_A.default.weight"
22
+
23
+ LORA_B_SUFFIX = ".lora_B.default.weight"
24
+
25
+ def validate_h3_lora_state_dict(state_dict: Mapping[str, torch.Tensor]) -> int:
26
+ lora_a: dict[str, torch.Tensor] = {}
27
+ lora_b: dict[str, torch.Tensor] = {}
28
+ for key, tensor in state_dict.items():
29
+ if key.endswith(LORA_A_SUFFIX):
30
+ lora_a[key[: -len(LORA_A_SUFFIX)]] = tensor
31
+ elif key.endswith(LORA_B_SUFFIX):
32
+ lora_b[key[: -len(LORA_B_SUFFIX)]] = tensor
33
+ else:
34
+ raise ValueError(f"Unsupported MiniMax H3 LoRA tensor key: {key!r}.")
35
+ if tensor.dtype != torch.bfloat16:
36
+ raise ValueError(f"MiniMax H3 LoRA tensor {key!r} must be BF16, got {tensor.dtype}.")
37
+ if tensor.ndim != 2 or not torch.isfinite(tensor).all():
38
+ raise ValueError(f"MiniMax H3 LoRA tensor {key!r} must be a finite matrix.")
39
+
40
+ if not lora_a:
41
+ raise ValueError("MiniMax H3 LoRA state dict contains no adapter tensors.")
42
+ missing_a = sorted(lora_b.keys() - lora_a.keys())
43
+ missing_b = sorted(lora_a.keys() - lora_b.keys())
44
+ if missing_a or missing_b:
45
+ raise ValueError(f"Unpaired MiniMax H3 LoRA tensors: missing A={missing_a[:3]}, missing B={missing_b[:3]}.")
46
+
47
+ ranks = set()
48
+ for module_name, a_tensor in lora_a.items():
49
+ b_tensor = lora_b[module_name]
50
+ if a_tensor.shape[0] != b_tensor.shape[1]:
51
+ raise ValueError(
52
+ f"MiniMax H3 LoRA rank mismatch for {module_name}: A{tuple(a_tensor.shape)} B{tuple(b_tensor.shape)}."
53
+ )
54
+ if not module_name.endswith(H3_LORA_TARGET_MODULES):
55
+ raise ValueError(f"Unsupported MiniMax H3 LoRA target: {module_name!r}.")
56
+ ranks.add(int(a_tensor.shape[0]))
57
+ if len(ranks) != 1:
58
+ raise ValueError(f"MiniMax H3 LoRA checkpoint has mixed ranks: {sorted(ranks)}.")
59
+ return ranks.pop()
60
+
61
+ def load_h3_lora_file(path: str | Path) -> tuple[dict[str, torch.Tensor], dict[str, str], int, int]:
62
+ path = Path(path).expanduser().resolve()
63
+ with safe_open(str(path), framework="pt", device="cpu") as handle:
64
+ metadata = handle.metadata() or {}
65
+ state_dict = load_file(str(path), device="cpu")
66
+ rank = validate_h3_lora_state_dict(state_dict)
67
+ alpha = int(metadata.get("alpha", rank))
68
+ return state_dict, metadata, rank, alpha
69
+
70
+ def apply_h3_lora(
71
+ transformer: torch.nn.Module,
72
+ path: str | Path,
73
+ *,
74
+ scale: float = 1.0,
75
+ fuse: bool = False,
76
+ ) -> tuple[int, int]:
77
+ state_dict, _, rank, alpha = load_h3_lora_file(path)
78
+ transformer.add_adapter(
79
+ LoraConfig(
80
+ r=rank,
81
+ lora_alpha=alpha,
82
+ init_lora_weights=False,
83
+ target_modules=list(H3_LORA_TARGET_MODULES),
84
+ use_rslora=False,
85
+ )
86
+ )
87
+ adapter_parameters = {
88
+ name: parameter
89
+ for name, parameter in transformer.named_parameters()
90
+ if ".lora_A." in name or ".lora_B." in name
91
+ }
92
+ missing = sorted(adapter_parameters.keys() - state_dict.keys())
93
+ unexpected = sorted(state_dict.keys() - adapter_parameters.keys())
94
+ mismatches = [
95
+ (name, tuple(state_dict[name].shape), tuple(parameter.shape))
96
+ for name, parameter in adapter_parameters.items()
97
+ if name in state_dict and state_dict[name].shape != parameter.shape
98
+ ]
99
+ if missing or unexpected or mismatches:
100
+ raise ValueError(
101
+ "MiniMax H3 LoRA is incompatible with the transformer: "
102
+ f"missing={missing[:3]}, unexpected={unexpected[:3]}, mismatches={mismatches[:3]}."
103
+ )
104
+ incompatible = transformer.load_state_dict(state_dict, strict=False)
105
+ missing_lora = [key for key in incompatible.missing_keys if ".lora_A." in key or ".lora_B." in key]
106
+ if incompatible.unexpected_keys or missing_lora:
107
+ raise RuntimeError(
108
+ "MiniMax H3 LoRA loading did not consume all adapter tensors: "
109
+ f"missing={missing_lora[:3]}, unexpected={incompatible.unexpected_keys[:3]}."
110
+ )
111
+ transformer.set_adapters("default", weights=float(scale))
112
+ if fuse:
113
+ transformer.fuse_lora(lora_scale=1.0, safe_fusing=True, adapter_names=["default"])
114
+ transformer.unload_lora()
115
+ transformer.requires_grad_(False)
116
+ transformer.eval()
117
+ return rank, alpha
infer.py ADDED
@@ -0,0 +1,26 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # SPDX-License-Identifier: Apache-2.0
2
+ """Download the pinned upstream components, then run the bundled H3 example."""
3
+ import argparse,json,os,subprocess,sys
4
+ from pathlib import Path
5
+ from huggingface_hub import snapshot_download
6
+
7
+ def main():
8
+ root=Path(__file__).resolve().parent
9
+ parser=argparse.ArgumentParser(description=__doc__)
10
+ parser.add_argument('--output',type=Path,default=Path('outputs'))
11
+ parser.add_argument('--case',type=Path,default=root/'examples/case.json')
12
+ parser.add_argument('--base',type=Path,help='Optional complete local upstream snapshot')
13
+ args=parser.parse_args()
14
+ args.case=args.case.resolve()
15
+ if not args.case.is_file():parser.error(f'Case file does not exist or is not a regular file: {args.case}')
16
+ spec=json.loads((root/'sampling.json').read_text());sampler=spec['sampler']
17
+ components=['model_index.json','scheduler/*','audio_scheduler/*','vae/*','audio_vae/*','text_encoder/*','tokenizer/*','processor/*',sampler['transformer_subfolder']+'/*']
18
+ base=str(args.base) if args.base else snapshot_download(spec['model_id'],revision=spec['model_revision'],allow_patterns=components)
19
+ command=[sys.executable,'-m','torch.distributed.run','--standalone','--nproc_per_node=2',str(root/'tools/infer_minimax_h3.py'),
20
+ '--jobs-json',str(args.case.resolve()),'--model-id',base,'--output-dir',str(args.output.resolve()),
21
+ '--lora-path',str(root/spec['weights']),'--fuse-lora','--inference-steps',str(sampler['student_sample_steps']),
22
+ '--final-step-size-scale','0.5','--video-shift','12','--audio-shift','3','--fsdp2','--no-cpu-offload',
23
+ '--ff-chunk-size','1024','--qk-norm-chunk-size','2048','--no-resume']
24
+ if sampler['condition_mode']=='ref2va':command+=['--reference-image-short-edge','512']
25
+ subprocess.run(command,check=True,cwd=root,env=os.environ.copy())
26
+ if __name__=='__main__':main()
lora/adapter_config.json ADDED
@@ -0,0 +1,16 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "format": "h3_peft_state_dict",
3
+ "r": 128,
4
+ "lora_alpha": 128,
5
+ "base_model": "MiniMaxAI/MiniMax-H3",
6
+ "base_model_revision": "939557dc319dd91227e30195a763f272ba7f8765",
7
+ "transformer_subfolder": "transformer_ref",
8
+ "target_modules": [
9
+ "to_q",
10
+ "to_k",
11
+ "to_v",
12
+ "to_out.0",
13
+ "ff.net.0.proj",
14
+ "ff.net.2"
15
+ ]
16
+ }
lora/adapter_model.safetensors ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:b8cd339e167ab29e112d4f8934d3ac7f20b7e3269895fb149761bd959336bdd7
3
+ size 1383677792
requirements.txt ADDED
@@ -0,0 +1,14 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ torch==2.7.1
2
+ torchvision==0.22.1
3
+ transformers==5.17.0
4
+ peft==0.21.0
5
+ accelerate==1.15.0
6
+ safetensors==0.9.0rc0
7
+ huggingface-hub==1.32.0
8
+ av==18.1.0
9
+ imageio-ffmpeg==0.6.0
10
+ sentencepiece
11
+ protobuf
12
+ librosa
13
+ soundfile
14
+ diffusers @ git+https://github.com/huggingface/diffusers.git@4e0466f3e5260f0d78b5e2b68ffbf27d819cc6db
sampling.json ADDED
@@ -0,0 +1,20 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "schema_version": 1,
3
+ "model_id": "MiniMaxAI/MiniMax-H3",
4
+ "model_revision": "939557dc319dd91227e30195a763f272ba7f8765",
5
+ "weights": "lora/adapter_model.safetensors",
6
+ "precision": "bfloat16",
7
+ "sampler": {
8
+ "type": "joint_video_audio_euler",
9
+ "student_sample_steps": 4,
10
+ "final_step_size_scale": 0.5,
11
+ "video_shift": 12.0,
12
+ "audio_shift": 3.0,
13
+ "condition_mode": "ref2va",
14
+ "transformer_subfolder": "transformer_ref",
15
+ "height": 768,
16
+ "width": 1344,
17
+ "num_frames": 124,
18
+ "fps": 24
19
+ }
20
+ }
tools/__init__.py ADDED
File without changes
tools/infer_minimax_h3.py ADDED
@@ -0,0 +1,1299 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ #!/usr/bin/env python3
2
+ # SPDX-License-Identifier: Apache-2.0
3
+
4
+ """Run reproducible MiniMax H3 T2VA or Ref2VA inference with optional PEFT LoRA and FSDP2."""
5
+
6
+ from __future__ import annotations
7
+
8
+ import argparse
9
+ import gc
10
+ import hashlib
11
+ import json
12
+ import os
13
+ import shutil
14
+ import sys
15
+ import tempfile
16
+ from dataclasses import dataclass, replace
17
+ from pathlib import Path
18
+ from types import MethodType
19
+ from typing import Any
20
+
21
+ import torch
22
+ import torch.distributed as dist
23
+
24
+ _REPO_ROOT = Path(__file__).resolve().parents[1]
25
+ if str(_REPO_ROOT) not in sys.path:
26
+ sys.path.insert(0, str(_REPO_ROOT))
27
+
28
+ from safetensors.torch import load_file
29
+
30
+ from h3_lora import apply_h3_lora
31
+ from tools.minimax_h3_transformers_compat import (
32
+ ensure_qwen3vl_hidden_state_capture,
33
+ ensure_qwen3vl_mm_token_type_ids,
34
+ )
35
+
36
+ FPS = 24
37
+ DEFAULT_HEIGHT = 768
38
+ DEFAULT_WIDTH = 1344
39
+ DEFAULT_FRAMES = 124
40
+ DEFAULT_NFE = 4
41
+ SUPPORTED_SPATIAL_GEOMETRIES = ((1344, 768), (768, 1344), (1024, 1024), (1792, 576))
42
+ REF2VA_SUPPORTED_SPATIAL_GEOMETRIES = (
43
+ (1152, 512),
44
+ (544, 544),
45
+ (960, 544),
46
+ (768, 576),
47
+ (1024, 576),
48
+ (1152, 640),
49
+ (1536, 672),
50
+ (1280, 704),
51
+ (576, 768),
52
+ (768, 768),
53
+ (1024, 768),
54
+ (1344, 768),
55
+ (544, 960),
56
+ (768, 1024),
57
+ (640, 1152),
58
+ (768, 1344),
59
+ )
60
+ DEFAULT_BASE_MODEL_ID = "MiniMaxAI/MiniMax-H3"
61
+ DEFAULT_MODEL_REVISION = "939557dc319dd91227e30195a763f272ba7f8765"
62
+ DEFAULT_STAGING_PROTOCOL_VERSION = "node-local-mp4-copy-sha256-v1"
63
+
64
+
65
+ @dataclass(frozen=True)
66
+ class ReferenceInput:
67
+ kind: str
68
+ path: Path
69
+ name: str
70
+ sha256: str
71
+ bytes: int
72
+
73
+ def identity(self) -> dict[str, Any]:
74
+ return {
75
+ "kind": self.kind,
76
+ "name": self.name,
77
+ "sha256": self.sha256,
78
+ "bytes": self.bytes,
79
+ }
80
+
81
+
82
+ @dataclass(frozen=True)
83
+ class GenerationJob:
84
+ index: int
85
+ prompt: str
86
+ seed: int
87
+ task: str = "t2va"
88
+ references: tuple[ReferenceInput, ...] = ()
89
+ is_padding: bool = False
90
+ case_name: str = ""
91
+ source_request: str = ""
92
+ item_id: str = ""
93
+ category: str = ""
94
+ width: int | None = None
95
+ height: int | None = None
96
+ num_frames: int | None = None
97
+
98
+
99
+ @dataclass(frozen=True)
100
+ class DistributedContext:
101
+ enabled: bool
102
+ rank: int
103
+ local_rank: int
104
+ world_size: int
105
+
106
+ @property
107
+ def device(self) -> torch.device:
108
+ return torch.device("cuda", self.local_rank) if self.enabled else torch.device("cuda")
109
+
110
+
111
+ def _sha256(path: Path) -> str:
112
+ digest = hashlib.sha256()
113
+ with path.open("rb") as handle:
114
+ for block in iter(lambda: handle.read(8 * 1024 * 1024), b""):
115
+ digest.update(block)
116
+ return digest.hexdigest()
117
+
118
+
119
+ def raw_sigma_grid(num_steps: int, final_step_size_scale: float) -> list[float]:
120
+ if num_steps < 1:
121
+ raise ValueError("num_steps must be positive.")
122
+ if final_step_size_scale <= 0:
123
+ raise ValueError("final_step_size_scale must be positive.")
124
+ base_segment = 1.0 / (num_steps - 1 + final_step_size_scale)
125
+ sigmas = [1.0]
126
+ for step in range(num_steps):
127
+ segment = base_segment * (final_step_size_scale if step == num_steps - 1 else 1.0)
128
+ sigmas.append(sigmas[-1] - segment)
129
+ sigmas[-1] = 0.0
130
+ return sigmas
131
+
132
+
133
+ def shifted_sigma_grid(raw_sigmas: list[float], shift: float) -> list[float]:
134
+ return [shift * sigma / (1.0 + (shift - 1.0) * sigma) for sigma in raw_sigmas]
135
+
136
+
137
+ def install_explicit_schedule(scheduler: Any, sigmas: list[float]) -> None:
138
+ original_set_timesteps = scheduler.set_timesteps
139
+
140
+ def set_timesteps(_num_inference_steps=None, device=None, **kwargs):
141
+ explicit = kwargs.pop("sigmas", None)
142
+ if kwargs:
143
+ raise TypeError(f"Unsupported scheduler arguments: {sorted(kwargs)}")
144
+ original_set_timesteps(device=device, sigmas=sigmas if explicit is None else explicit)
145
+
146
+ scheduler.set_timesteps = set_timesteps
147
+
148
+
149
+ def load_jobs(path: Path, default_seed: int) -> list[GenerationJob]:
150
+ payload = json.loads(path.read_text(encoding="utf-8"))
151
+ examples = payload.get("examples")
152
+ if examples is None and isinstance(payload.get("records"), list):
153
+ examples = [
154
+ {
155
+ "task": "t2va",
156
+ "prompt": record.get("prompt"),
157
+ "seed": int(record.get("prompt_index", index)) * 1000,
158
+ }
159
+ for index, record in enumerate(payload["records"])
160
+ ]
161
+ if not isinstance(examples, list) or not examples:
162
+ raise ValueError(f"{path} must contain a non-empty examples or records list.")
163
+ jobs = []
164
+ for index, example in enumerate(examples):
165
+ if not isinstance(example, dict):
166
+ raise TypeError(f"Example {index} must be an object.")
167
+ task = str(example.get("task", "t2va")).lower()
168
+ if task not in {"t2va", "ref2va"}:
169
+ raise ValueError(f"Example {index} has unsupported task {task!r}; expected t2va or ref2va.")
170
+ prompt = example.get("prompt")
171
+ if not isinstance(prompt, str) or not prompt.strip():
172
+ raise ValueError(f"Example {index} requires a non-empty prompt.")
173
+ seed = example.get("seed", default_seed + index)
174
+ if isinstance(seed, bool) or not isinstance(seed, int) or seed < 0:
175
+ raise ValueError(f"Example {index} seed must be a non-negative integer.")
176
+ case_name = example.get("case_name", "")
177
+ source_request = example.get("source_request", "")
178
+ item_id = example.get("item_id", "")
179
+ category = example.get("category", "")
180
+ if not all(isinstance(value, str) for value in (case_name, source_request, item_id, category)):
181
+ raise TypeError(f"Example {index} identity fields must be strings.")
182
+ width = example.get("width")
183
+ height = example.get("height")
184
+ num_frames = example.get("num_frames")
185
+ for name, value in (("width", width), ("height", height), ("num_frames", num_frames)):
186
+ if value is not None and (isinstance(value, bool) or not isinstance(value, int) or value <= 0):
187
+ raise ValueError(f"Example {index} {name} must be a positive integer.")
188
+ if width is not None and width % 32:
189
+ raise ValueError(f"Example {index} width must be divisible by 32.")
190
+ if height is not None and height % 32:
191
+ raise ValueError(f"Example {index} height must be divisible by 32.")
192
+ if num_frames is not None and (num_frames < 5 or (num_frames - 5) % 17):
193
+ raise ValueError(f"Example {index} num_frames must satisfy 17*n+5.")
194
+ raw_references = example.get("references", [])
195
+ if not isinstance(raw_references, list):
196
+ raise TypeError(f"Example {index} references must be a list.")
197
+ if task == "t2va" and raw_references:
198
+ raise ValueError(f"Example {index} T2VA job must not declare references.")
199
+ if task == "ref2va" and not raw_references:
200
+ raise ValueError(f"Example {index} Ref2VA job requires at least one reference.")
201
+ references = []
202
+ reference_counts = {"image": 0, "video": 0, "audio": 0}
203
+ for reference_index, raw_reference in enumerate(raw_references):
204
+ if not isinstance(raw_reference, dict):
205
+ raise TypeError(f"Example {index} reference {reference_index} must be an object.")
206
+ kind = str(raw_reference.get("kind", "")).lower()
207
+ if kind not in reference_counts:
208
+ raise ValueError(
209
+ f"Example {index} reference {reference_index} has unsupported kind {kind!r}; "
210
+ "expected image, video, or audio."
211
+ )
212
+ raw_path = raw_reference.get("path")
213
+ if not isinstance(raw_path, str) or not raw_path.strip():
214
+ raise ValueError(f"Example {index} reference {reference_index} requires a non-empty path.")
215
+ reference_path = Path(raw_path).expanduser()
216
+ if not reference_path.is_absolute():
217
+ reference_path = path.parent / reference_path
218
+ reference_path = reference_path.resolve()
219
+ if not reference_path.is_file() or reference_path.stat().st_size <= 0:
220
+ raise FileNotFoundError(f"Example {index} reference is missing or empty: {reference_path}")
221
+ name = raw_reference.get("name", raw_path)
222
+ if not isinstance(name, str) or not name.strip():
223
+ raise ValueError(f"Example {index} reference {reference_index} name must be a non-empty string.")
224
+ reference_sha256 = _sha256(reference_path)
225
+ expected_sha256 = raw_reference.get("sha256")
226
+ if expected_sha256 is not None and expected_sha256 != reference_sha256:
227
+ raise ValueError(
228
+ f"Example {index} reference {reference_index} SHA256 mismatch: "
229
+ f"expected {expected_sha256!r}, got {reference_sha256!r}."
230
+ )
231
+ reference_bytes = reference_path.stat().st_size
232
+ expected_bytes = raw_reference.get("bytes")
233
+ if expected_bytes is not None and expected_bytes != reference_bytes:
234
+ raise ValueError(
235
+ f"Example {index} reference {reference_index} byte-size mismatch: "
236
+ f"expected {expected_bytes!r}, got {reference_bytes!r}."
237
+ )
238
+ reference_counts[kind] += 1
239
+ references.append(
240
+ ReferenceInput(
241
+ kind=kind,
242
+ path=reference_path,
243
+ name=name,
244
+ sha256=reference_sha256,
245
+ bytes=reference_bytes,
246
+ )
247
+ )
248
+ if len(references) > 12 or reference_counts["image"] > 9 or reference_counts["video"] > 3:
249
+ raise ValueError(f"Example {index} exceeds the MiniMax H3 Ref2VA visual-reference limits.")
250
+ if reference_counts["audio"] > 3:
251
+ raise ValueError(f"Example {index} exceeds the MiniMax H3 Ref2VA audio-reference limit.")
252
+ if reference_counts["audio"] and not (reference_counts["image"] or reference_counts["video"]):
253
+ raise ValueError(f"Example {index} audio references require at least one image or video reference.")
254
+ jobs.append(
255
+ GenerationJob(
256
+ index=index,
257
+ prompt=prompt,
258
+ seed=seed,
259
+ task=task,
260
+ references=tuple(references),
261
+ case_name=case_name or item_id,
262
+ source_request=source_request,
263
+ item_id=item_id,
264
+ category=category,
265
+ width=width,
266
+ height=height,
267
+ num_frames=num_frames,
268
+ )
269
+ )
270
+ return jobs
271
+
272
+
273
+ def distributed_context(enabled: bool) -> DistributedContext:
274
+ world_size = int(os.environ.get("WORLD_SIZE", "1"))
275
+ if not enabled:
276
+ if world_size != 1:
277
+ raise RuntimeError("Multi-process launch requires --fsdp2.")
278
+ if not torch.cuda.is_available():
279
+ raise RuntimeError("MiniMax H3 inference requires CUDA.")
280
+ return DistributedContext(False, 0, 0, 1)
281
+ for name in ("RANK", "LOCAL_RANK", "WORLD_SIZE"):
282
+ if name not in os.environ:
283
+ raise RuntimeError(f"--fsdp2 requires torchrun environment variable {name}.")
284
+ context = DistributedContext(
285
+ True,
286
+ int(os.environ["RANK"]),
287
+ int(os.environ["LOCAL_RANK"]),
288
+ int(os.environ["WORLD_SIZE"]),
289
+ )
290
+ if not 0 <= context.rank < context.world_size:
291
+ raise ValueError(f"Invalid distributed rank {context.rank}/{context.world_size}.")
292
+ return context
293
+
294
+
295
+ def assign_jobs(jobs: list[GenerationJob], context: DistributedContext) -> list[GenerationJob]:
296
+ if not context.enabled:
297
+ return jobs
298
+ if not jobs:
299
+ return []
300
+ buckets = [
301
+ [job for job in jobs if job.index % context.world_size == rank] for rank in range(context.world_size)
302
+ ]
303
+ jobs_per_rank = max(len(bucket) for bucket in buckets)
304
+ assigned = list(buckets[context.rank])
305
+ for offset in range(jobs_per_rank - len(assigned)):
306
+ source = jobs[offset % len(jobs)]
307
+ assigned.append(replace(source, is_padding=True))
308
+ return assigned
309
+
310
+
311
+ def select_job_shard(jobs: list[GenerationJob], shard_index: int, shard_count: int) -> list[GenerationJob]:
312
+ if shard_count < 1:
313
+ raise ValueError("job shard count must be positive.")
314
+ if not 0 <= shard_index < shard_count:
315
+ raise ValueError(f"job shard index must be in [0, {shard_count}).")
316
+ return [job for job in jobs if job.index % shard_count == shard_index]
317
+
318
+
319
+ def load_job_references(job: GenerationJob) -> list[Any]:
320
+ if job.task != "ref2va":
321
+ return []
322
+ from diffusers.modular_pipelines.minimax_h3.references import (
323
+ MiniMaxH3AudioReference,
324
+ MiniMaxH3ImageReference,
325
+ MiniMaxH3VideoReference,
326
+ )
327
+
328
+ classes = {
329
+ "image": MiniMaxH3ImageReference,
330
+ "video": MiniMaxH3VideoReference,
331
+ "audio": MiniMaxH3AudioReference,
332
+ }
333
+ return [classes[reference.kind].from_file(reference.path) for reference in job.references]
334
+
335
+
336
+ def validate_reference_image_policy(jobs: list[GenerationJob], reference_image_short_edge: int | None) -> None:
337
+ has_images = any(reference.kind == "image" for job in jobs for reference in job.references)
338
+ if not has_images:
339
+ return
340
+ if reference_image_short_edge is None:
341
+ raise ValueError("Ref2VA image-reference jobs require explicit --reference-image-short-edge.")
342
+ if reference_image_short_edge <= 0 or reference_image_short_edge % 32:
343
+ raise ValueError("--reference-image-short-edge must be a positive multiple of 32.")
344
+
345
+
346
+ def validate_resolved_geometry(
347
+ width: int,
348
+ height: int,
349
+ num_frames: int,
350
+ *,
351
+ allow_variable_spatial_geometry: bool,
352
+ task: str = "t2va",
353
+ ) -> None:
354
+ if width <= 0 or height <= 0 or width % 32 or height % 32 or num_frames < 5 or (num_frames - 5) % 17:
355
+ raise ValueError(f"Invalid resolved geometry: {width}x{height}x{num_frames}.")
356
+ supported = REF2VA_SUPPORTED_SPATIAL_GEOMETRIES if task == "ref2va" else SUPPORTED_SPATIAL_GEOMETRIES
357
+ if (width, height) not in supported:
358
+ geometry_label = "MiniMax H3 spatial geometry" if task == "t2va" else "MiniMax H3 Ref2VA spatial geometry"
359
+ raise ValueError(f"Unsupported {geometry_label}: {width}x{height}; supported={list(supported)}.")
360
+ if (width, height) != (DEFAULT_WIDTH, DEFAULT_HEIGHT) and not allow_variable_spatial_geometry:
361
+ raise ValueError(
362
+ f"Non-default MiniMax H3 geometry {width}x{height} requires --allow-variable-spatial-geometry."
363
+ )
364
+
365
+
366
+ def initialize_fsdp2(context: DistributedContext):
367
+ if not context.enabled:
368
+ return None
369
+ if not torch.cuda.is_available() or context.local_rank >= torch.cuda.device_count():
370
+ raise RuntimeError(f"CUDA local rank {context.local_rank} is unavailable.")
371
+ torch.cuda.set_device(context.local_rank)
372
+ dist.init_process_group(backend="nccl", device_id=context.device)
373
+ from torch.distributed.device_mesh import init_device_mesh
374
+
375
+ return init_device_mesh("cuda", (context.world_size,), mesh_dim_names=("fsdp",))
376
+
377
+
378
+ def shard_h3_transformer(transformer: torch.nn.Module, mesh) -> None:
379
+ from torch.distributed.fsdp import fully_shard
380
+
381
+ refiner_blocks = transformer.token_refiner.refiner_blocks
382
+ for block in refiner_blocks:
383
+ fully_shard(block, mesh=mesh, reshard_after_forward=True)
384
+ fully_shard(transformer.token_refiner, mesh=mesh, reshard_after_forward=True)
385
+ for block in transformer.transformer_blocks:
386
+ fully_shard(block, mesh=mesh, reshard_after_forward=True)
387
+ for name in (
388
+ "proj_in",
389
+ "audio_proj_in",
390
+ "context_embedder",
391
+ "time_embedder",
392
+ "norm_out",
393
+ "proj_out",
394
+ "audio_proj_out",
395
+ ):
396
+ fully_shard(getattr(transformer, name), mesh=mesh, reshard_after_forward=True)
397
+ fully_shard(transformer, mesh=mesh, reshard_after_forward=True)
398
+
399
+
400
+ def shard_h3_text_encoder(text_encoder: torch.nn.Module, mesh) -> None:
401
+ from torch.distributed.fsdp import CPUOffloadPolicy, fully_shard
402
+
403
+ # Keep the large conditioner shards off GPU while the denoiser and VAEs
404
+ # are resident. Forward still computes on the FSDP CUDA mesh, layer by layer.
405
+ policy = CPUOffloadPolicy(pin_memory=False)
406
+ model = text_encoder.model
407
+ for layer in model.language_model.layers:
408
+ fully_shard(layer, mesh=mesh, reshard_after_forward=True, offload_policy=policy)
409
+ fully_shard(model, mesh=mesh, reshard_after_forward=True, offload_policy=policy)
410
+
411
+ def reshard_root(module, _inputs, output):
412
+ module.reshard()
413
+ return output
414
+
415
+ model.register_forward_hook(reshard_root)
416
+
417
+
418
+ def enable_h3_ff_chunking(transformer: torch.nn.Module, chunk_size: int) -> int:
419
+ if chunk_size <= 0:
420
+ return 0
421
+ wrapped = 0
422
+ for block in transformer.transformer_blocks:
423
+ feed_forward = block.ff
424
+ if getattr(feed_forward, "_fastgen_chunk_size", None) == chunk_size:
425
+ continue
426
+ original_forward = feed_forward.forward
427
+
428
+ def chunked_forward(self, hidden_states, *args, _forward=original_forward, **kwargs):
429
+ if hidden_states.ndim < 3 or hidden_states.shape[1] <= self._fastgen_chunk_size:
430
+ return _forward(hidden_states, *args, **kwargs)
431
+ return torch.cat(
432
+ [_forward(chunk, *args, **kwargs) for chunk in hidden_states.split(self._fastgen_chunk_size, dim=1)],
433
+ dim=1,
434
+ )
435
+
436
+ feed_forward._fastgen_chunk_size = int(chunk_size)
437
+ feed_forward.forward = MethodType(chunked_forward, feed_forward)
438
+ wrapped += 1
439
+ return wrapped
440
+
441
+
442
+ def low_peak_h3_rotary(
443
+ hidden_states: torch.Tensor,
444
+ cos: torch.Tensor,
445
+ sin: torch.Tensor,
446
+ ) -> torch.Tensor:
447
+ rotary_dim = cos.shape[-1]
448
+ if rotary_dim > hidden_states.shape[-1] or rotary_dim % 2:
449
+ raise ValueError(f"Invalid MiniMax H3 rotary dim {rotary_dim} for {hidden_states.shape[-1]}.")
450
+ half = rotary_dim // 2
451
+ rotary = hidden_states[..., :rotary_dim]
452
+ cos = cos.to(hidden_states.dtype)[None, :, None, :half]
453
+ sin = sin.to(hidden_states.dtype)[None, :, None, :half]
454
+ for start in range(0, half, 4):
455
+ end = min(start + 4, half)
456
+ first = rotary[..., start:end].clone()
457
+ second = rotary[..., half + start : half + end].clone()
458
+ cos_chunk = cos[..., start:end]
459
+ sin_chunk = sin[..., start:end]
460
+ rotary[..., start:end] = first * cos_chunk - second * sin_chunk
461
+ rotary[..., half + start : half + end] = second * cos_chunk + first * sin_chunk
462
+ return hidden_states.contiguous()
463
+
464
+
465
+ def enable_h3_low_peak_rope() -> None:
466
+ from diffusers.models.transformers import transformer_minimax_h3
467
+
468
+ transformer_minimax_h3._apply_rotary_emb = low_peak_h3_rotary
469
+
470
+
471
+ def low_peak_h3_rms_norm_inplace(
472
+ norm: torch.nn.Module,
473
+ hidden_states: torch.Tensor,
474
+ chunk_size: int,
475
+ ) -> torch.Tensor:
476
+ if chunk_size <= 0 or hidden_states.shape[1] <= chunk_size:
477
+ return norm(hidden_states)
478
+ if torch.is_grad_enabled() and hidden_states.requires_grad:
479
+ raise RuntimeError("In-place MiniMax H3 Q/K RMSNorm chunking is inference-only.")
480
+ for start in range(0, hidden_states.shape[1], chunk_size):
481
+ end = min(start + chunk_size, hidden_states.shape[1])
482
+ normalized = norm(hidden_states[:, start:end])
483
+ hidden_states[:, start:end].copy_(normalized)
484
+ return hidden_states
485
+
486
+
487
+ class _LowPeakH3QKNorm(torch.nn.Module):
488
+ def __init__(self, norm: torch.nn.Module, chunk_size: int) -> None:
489
+ super().__init__()
490
+ self.norm = norm
491
+ self.chunk_size = int(chunk_size)
492
+
493
+ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
494
+ return low_peak_h3_rms_norm_inplace(self.norm, hidden_states, self.chunk_size)
495
+
496
+
497
+ def enable_h3_qk_norm_chunking(transformer: torch.nn.Module, chunk_size: int) -> int:
498
+ if chunk_size <= 0:
499
+ return 0
500
+ wrapped = 0
501
+ for block in transformer.transformer_blocks:
502
+ attention = block.attn
503
+ for name in ("norm_q", "norm_k"):
504
+ norm = getattr(attention, name)
505
+ if isinstance(norm, _LowPeakH3QKNorm):
506
+ norm.chunk_size = int(chunk_size)
507
+ continue
508
+ setattr(attention, name, _LowPeakH3QKNorm(norm, chunk_size))
509
+ wrapped += 1
510
+ return wrapped
511
+
512
+
513
+ def offload_h3_transformer(transformer: torch.nn.Module, *, group_offload: bool = False) -> None:
514
+ if group_offload:
515
+ from diffusers.hooks.group_offloading import _is_group_offload_enabled
516
+
517
+ if not _is_group_offload_enabled(transformer):
518
+ raise RuntimeError("MiniMax H3 group offload was requested but is not installed.")
519
+ gc.collect()
520
+ torch.cuda.empty_cache()
521
+ return
522
+ hook = getattr(transformer, "_hf_hook", None)
523
+ if hook is not None and hasattr(hook, "init_hook"):
524
+ hook.init_hook(transformer)
525
+ else:
526
+ transformer.to("cpu")
527
+ device = next(transformer.parameters()).device
528
+ if device.type != "cpu":
529
+ raise RuntimeError(f"MiniMax H3 Transformer did not offload before VAE decode: {device}")
530
+ gc.collect()
531
+ torch.cuda.empty_cache()
532
+
533
+
534
+ def move_unsharded_components(pipe, transformer: torch.nn.Module, device: torch.device) -> None:
535
+ sharded = {id(transformer)}
536
+ if getattr(pipe, "text_encoder", None) is not None:
537
+ sharded.add(id(pipe.text_encoder))
538
+ for component in pipe.components.values():
539
+ if isinstance(component, torch.nn.Module) and id(component) not in sharded:
540
+ component.to(device)
541
+
542
+
543
+ def cached_h3_blocks(*, group_offload: bool = False, execution_device: torch.device | None = None):
544
+ from diffusers.modular_pipelines.minimax_h3.before_denoise import (
545
+ MiniMaxH3NoKeyframeAnchorsStep,
546
+ MiniMaxH3PrepareLatentsStep,
547
+ MiniMaxH3PrepareLayoutStep,
548
+ MiniMaxH3SetTimestepsStep,
549
+ )
550
+ from diffusers.modular_pipelines.minimax_h3.decoders import (
551
+ MiniMaxH3AfterDenoiseStep,
552
+ MiniMaxH3AudioDecodeStep,
553
+ MiniMaxH3VideoDecodeStep,
554
+ )
555
+ from diffusers.modular_pipelines.minimax_h3.denoise import MiniMaxH3DenoiseStep
556
+ from diffusers.modular_pipelines.modular_pipeline import (
557
+ ModularPipelineBlocks,
558
+ SequentialPipelineBlocks,
559
+ )
560
+
561
+ class PrepareDenoiseDevicesStep(ModularPipelineBlocks):
562
+ @torch.no_grad()
563
+ def __call__(self, components, state):
564
+ if group_offload:
565
+ components.vae.to("cpu")
566
+ components.audio_vae.to("cpu")
567
+ gc.collect()
568
+ torch.cuda.empty_cache()
569
+ return components, state
570
+
571
+ class OffloadTransformerStep(ModularPipelineBlocks):
572
+ @torch.no_grad()
573
+ def __call__(self, components, state):
574
+ offload_h3_transformer(components.transformer, group_offload=group_offload)
575
+ if group_offload:
576
+ if execution_device is None:
577
+ raise RuntimeError("MiniMax H3 group offload requires an execution device.")
578
+ components.vae.to(execution_device)
579
+ components.audio_vae.to(execution_device)
580
+ return components, state
581
+
582
+ return SequentialPipelineBlocks.from_blocks_dict(
583
+ {
584
+ "prepare_denoise_devices": PrepareDenoiseDevicesStep(),
585
+ "no_keyframe_anchors": MiniMaxH3NoKeyframeAnchorsStep(),
586
+ "prepare_layout": MiniMaxH3PrepareLayoutStep(),
587
+ "prepare_latents": MiniMaxH3PrepareLatentsStep(),
588
+ "set_timesteps": MiniMaxH3SetTimestepsStep(),
589
+ "denoise": MiniMaxH3DenoiseStep(),
590
+ "after_denoise": MiniMaxH3AfterDenoiseStep(),
591
+ "offload_transformer": OffloadTransformerStep(),
592
+ "decode_video": MiniMaxH3VideoDecodeStep(),
593
+ "decode_audio": MiniMaxH3AudioDecodeStep(),
594
+ }
595
+ )
596
+
597
+
598
+ def load_pipeline(args: argparse.Namespace, context: DistributedContext, mesh, task: str):
599
+ enable_h3_low_peak_rope()
600
+ text_encoder_sharded = False
601
+ component_roots = {"default": args.model_id}
602
+ if args.reference_model_root:
603
+ component_roots["transformer_ref"] = str(args.reference_model_root)
604
+ try:
605
+ from diffusers import ComponentsManager, ModularPipeline
606
+ from diffusers.modular_pipelines.minimax_h3.modular_pipeline import MiniMaxH3ModularPipeline
607
+ except ImportError as exc:
608
+ raise RuntimeError(
609
+ "MiniMax H3 inference requires a current Diffusers build with ModularPipeline support; "
610
+ "the training environment's pinned Diffusers 0.38 is insufficient."
611
+ ) from exc
612
+
613
+ if task == "t2va" and args.prompt_cache is not None:
614
+
615
+ class CachedMiniMaxH3ModularPipeline(MiniMaxH3ModularPipeline):
616
+ @property
617
+ def max_duration(self):
618
+ if args.allow_aligned_15s:
619
+ return 362 / FPS
620
+ return super().max_duration
621
+
622
+ manager = None
623
+ if not context.enabled and not args.no_cpu_offload and args.group_offload_blocks == 0:
624
+ manager = ComponentsManager()
625
+ manager.enable_auto_cpu_offload(
626
+ device=args.device,
627
+ memory_reserve_margin=args.memory_reserve_margin,
628
+ )
629
+ pipe = CachedMiniMaxH3ModularPipeline(
630
+ blocks=cached_h3_blocks(
631
+ group_offload=args.group_offload_blocks > 0,
632
+ execution_device=context.device,
633
+ ),
634
+ pretrained_model_name_or_path=args.model_id,
635
+ components_manager=manager,
636
+ )
637
+ pipe.load_components(
638
+ names=["scheduler", "audio_scheduler", "transformer", "vae", "audio_vae"],
639
+ dtype=torch.bfloat16,
640
+ pretrained_model_name_or_path=component_roots,
641
+ )
642
+ elif context.enabled or args.no_cpu_offload:
643
+ if context.enabled:
644
+ class DistributedH3Pipeline(MiniMaxH3ModularPipeline):
645
+ @property
646
+ def _execution_device(self):
647
+ # CPU-resident FSDP shards must not make the pipeline select
648
+ # CPU for noise, reference inputs or scheduler tensors.
649
+ return context.device
650
+
651
+ pipe = DistributedH3Pipeline.from_pretrained(args.model_id, workflow=task)
652
+ else:
653
+ pipe = ModularPipeline.from_pretrained(args.model_id, workflow=task)
654
+ # Load and shard the encoder before allocating the denoiser. Two ranks
655
+ # must not each hold both full 62-GiB modules in host memory at once.
656
+ if context.enabled and args.prompt_cache is None:
657
+ pipe.load_components(names=["text_encoder"], dtype=torch.bfloat16,
658
+ pretrained_model_name_or_path=component_roots)
659
+ pipe.text_encoder.requires_grad_(False).eval()
660
+ shard_h3_text_encoder(pipe.text_encoder, mesh)
661
+ text_encoder_sharded = True
662
+ gc.collect()
663
+ remaining = ["scheduler", "audio_scheduler", "vae", "audio_vae",
664
+ "transformer_ref" if task == "ref2va" else "transformer"]
665
+ if args.prompt_cache is None:
666
+ remaining += ["tokenizer", "processor"]
667
+ if not text_encoder_sharded:
668
+ remaining.append("text_encoder")
669
+ pipe.load_components(
670
+ names=remaining,
671
+ dtype=torch.bfloat16,
672
+ pretrained_model_name_or_path=component_roots,
673
+ )
674
+ else:
675
+ manager = ComponentsManager()
676
+ manager.enable_auto_cpu_offload(
677
+ device=args.device,
678
+ memory_reserve_margin=args.memory_reserve_margin,
679
+ )
680
+ pipe = ModularPipeline.from_pretrained(
681
+ args.model_id,
682
+ components_manager=manager,
683
+ workflow=task,
684
+ )
685
+ pipe.load_components(
686
+ dtype=torch.bfloat16,
687
+ pretrained_model_name_or_path=component_roots,
688
+ )
689
+ transformer_name = "transformer_ref" if task == "ref2va" else "transformer"
690
+ transformer = getattr(pipe, transformer_name, None)
691
+ if args.reference_image_short_edge is not None:
692
+ pipe.register_to_config(reference_image_short_edge=args.reference_image_short_edge)
693
+ required = [
694
+ "vae",
695
+ "scheduler",
696
+ "audio_scheduler",
697
+ "audio_vae",
698
+ transformer_name,
699
+ ]
700
+ if args.prompt_cache is None:
701
+ required.extend(["text_encoder", "tokenizer", "processor"])
702
+ missing = [name for name in required if getattr(pipe, name, None) is None]
703
+ if missing:
704
+ raise RuntimeError(f"MiniMax H3 pipeline is missing components: {missing}.")
705
+ pipe.scheduler.set_shift(args.video_shift)
706
+ pipe.audio_scheduler.set_shift(args.audio_shift)
707
+ raw_sigmas = raw_sigma_grid(args.inference_steps, args.final_step_size_scale)
708
+ install_explicit_schedule(pipe.scheduler, shifted_sigma_grid(raw_sigmas, args.video_shift))
709
+ install_explicit_schedule(pipe.audio_scheduler, shifted_sigma_grid(raw_sigmas, args.audio_shift))
710
+
711
+ if args.lora_path is not None:
712
+ rank, alpha = apply_h3_lora(
713
+ transformer,
714
+ args.lora_path,
715
+ scale=args.lora_scale,
716
+ fuse=args.fuse_lora,
717
+ )
718
+ print(
719
+ f"[rank {context.rank}] loaded LoRA path={args.lora_path} rank={rank} alpha={alpha} "
720
+ f"scale={args.lora_scale:g} fused={args.fuse_lora}",
721
+ flush=True,
722
+ )
723
+ chunked_blocks = enable_h3_ff_chunking(transformer, args.ff_chunk_size)
724
+ print(
725
+ f"[rank {context.rank}] FF chunking size={args.ff_chunk_size} blocks={chunked_blocks}",
726
+ flush=True,
727
+ )
728
+ chunked_qk_norms = enable_h3_qk_norm_chunking(transformer, args.qk_norm_chunk_size)
729
+ print(
730
+ f"[rank {context.rank}] Q/K norm chunking size={args.qk_norm_chunk_size} norms={chunked_qk_norms}",
731
+ flush=True,
732
+ )
733
+ transformer.requires_grad_(False).eval()
734
+ if args.group_offload_blocks > 0:
735
+ if context.enabled:
736
+ raise ValueError("MiniMax H3 group offload is only supported by the single-GPU cached-condition path.")
737
+ if task != "t2va" or args.prompt_cache is None:
738
+ raise ValueError("MiniMax H3 group offload currently requires cached-condition T2VA inference.")
739
+ transformer.enable_group_offload(
740
+ onload_device=context.device,
741
+ offload_device=torch.device("cpu"),
742
+ offload_type="block_level",
743
+ num_blocks_per_group=args.group_offload_blocks,
744
+ use_stream=args.group_offload_stream,
745
+ )
746
+ print(
747
+ f"[rank {context.rank}] group offload blocks={args.group_offload_blocks} "
748
+ f"stream={args.group_offload_stream}",
749
+ flush=True,
750
+ )
751
+ if args.attention_backend:
752
+ transformer.set_attention_backend(args.attention_backend)
753
+
754
+ if context.enabled:
755
+ if args.prompt_cache is None and not text_encoder_sharded:
756
+ pipe.text_encoder.requires_grad_(False).eval()
757
+ shard_h3_text_encoder(pipe.text_encoder, mesh)
758
+ shard_h3_transformer(transformer, mesh)
759
+ move_unsharded_components(pipe, transformer, context.device)
760
+ print(f"[rank {context.rank}/{context.world_size}] FSDP2 ready on {context.device}", flush=True)
761
+ elif args.no_cpu_offload:
762
+ pipe.to(args.device)
763
+ return pipe
764
+
765
+
766
+ def save_result(result: dict[str, Any], path: Path) -> None:
767
+ from diffusers.utils.export_utils import encode_video
768
+
769
+ audio = result["audio"][0]
770
+ if not isinstance(audio, torch.Tensor):
771
+ audio = torch.as_tensor(audio)
772
+ encode_video(
773
+ result["videos"][0],
774
+ fps=FPS,
775
+ output_path=str(path),
776
+ audio=audio.detach(),
777
+ audio_sample_rate=int(result["sampling_rate"]),
778
+ )
779
+
780
+
781
+ def is_structurally_valid_mp4(path: Path) -> bool:
782
+ try:
783
+ if not path.is_file() or path.stat().st_size < 12:
784
+ return False
785
+ with path.open("rb") as handle:
786
+ return b"ftyp" in handle.read(64)
787
+ except OSError:
788
+ return False
789
+
790
+
791
+ def publish_result(result: dict[str, Any], output: Path, staging_dir: Path | None) -> tuple[str, str]:
792
+ output.parent.mkdir(parents=True, exist_ok=True)
793
+ if staging_dir is None:
794
+ staged = output.with_name(f".{output.name}.{os.getpid()}.partial.mp4")
795
+ publish_mode = "adjacent-atomic-replace"
796
+ else:
797
+ staging_dir.mkdir(parents=True, exist_ok=True)
798
+ staged = staging_dir / f"{output.stem}.{os.getpid()}.partial.mp4"
799
+ publish_mode = DEFAULT_STAGING_PROTOCOL_VERSION
800
+ staged.unlink(missing_ok=True)
801
+ save_result(result, staged)
802
+ if not is_structurally_valid_mp4(staged):
803
+ raise RuntimeError(f"MiniMax H3 encoder produced an invalid MP4: {staged}")
804
+ staged_sha256 = _sha256(staged)
805
+
806
+ if staging_dir is None:
807
+ staged.replace(output)
808
+ else:
809
+ shared_partial = output.with_name(f".{output.name}.{os.getpid()}.copying.partial.mp4")
810
+ shared_partial.unlink(missing_ok=True)
811
+ try:
812
+ shutil.copyfile(staged, shared_partial)
813
+ if not is_structurally_valid_mp4(shared_partial) or _sha256(shared_partial) != staged_sha256:
814
+ raise RuntimeError(f"MiniMax H3 shared MP4 copy failed structural/hash validation: {shared_partial}")
815
+ shared_partial.replace(output)
816
+ finally:
817
+ shared_partial.unlink(missing_ok=True)
818
+ staged.unlink(missing_ok=True)
819
+ if not is_structurally_valid_mp4(output) or _sha256(output) != staged_sha256:
820
+ raise RuntimeError(f"MiniMax H3 published MP4 failed structural/hash validation: {output}")
821
+ return staged_sha256, publish_mode
822
+
823
+
824
+ def cleanup_result_partials(output: Path) -> None:
825
+ if not output.parent.is_dir():
826
+ return
827
+ for pattern in (
828
+ f".{output.name}.*.partial.mp4",
829
+ f".{output.name}.*.copying.partial.mp4",
830
+ ):
831
+ for path in output.parent.glob(pattern):
832
+ path.unlink(missing_ok=True)
833
+
834
+
835
+ def result_path(output_dir: Path, job: GenerationJob, rank: int) -> Path:
836
+ return output_dir / f"{job.index:04d}_{job.task}_seed{job.seed}_rank{rank:04d}.mp4"
837
+
838
+
839
+ def _atomic_write_json(path: Path, payload: dict[str, Any]) -> None:
840
+ path.parent.mkdir(parents=True, exist_ok=True)
841
+ with tempfile.NamedTemporaryFile(
842
+ mode="w",
843
+ dir=path.parent,
844
+ prefix=f".{path.name}.",
845
+ suffix=".tmp",
846
+ encoding="utf-8",
847
+ delete=False,
848
+ ) as handle:
849
+ staged = Path(handle.name)
850
+ json.dump(payload, handle, indent=2)
851
+ handle.write("\n")
852
+ handle.flush()
853
+ os.fsync(handle.fileno())
854
+ staged.replace(path)
855
+
856
+
857
+ def load_cached_condition(prompt_cache: Path, job: GenerationJob) -> tuple[torch.Tensor, torch.Tensor, str]:
858
+ condition_path = prompt_cache / f"{job.index:04d}.safetensors"
859
+ metadata_path = prompt_cache / f"{job.index:04d}.json"
860
+ if not condition_path.is_file() or not metadata_path.is_file():
861
+ raise FileNotFoundError(f"Missing cached MiniMax H3 condition for job {job.index}: {condition_path}")
862
+ metadata = json.loads(metadata_path.read_text(encoding="utf-8"))
863
+ prompt_sha256 = hashlib.sha256(job.prompt.encode("utf-8")).hexdigest()
864
+ if metadata.get("complete") is not True or metadata.get("prompt_sha256") != prompt_sha256:
865
+ raise ValueError(f"MiniMax H3 cached condition does not match job {job.index}.")
866
+ tensors = load_file(str(condition_path), device="cpu")
867
+ embeds = tensors.get("prompt_embeds")
868
+ tags = tensors.get("text_token_tags")
869
+ if (
870
+ embeds is None
871
+ or tags is None
872
+ or embeds.dtype != torch.bfloat16
873
+ or embeds.ndim != 3
874
+ or embeds.shape[0] != 1
875
+ or embeds.shape[-1] != 5120
876
+ or tags.dtype != torch.int64
877
+ or tags.shape != (embeds.shape[1],)
878
+ or not torch.isfinite(embeds).all()
879
+ ):
880
+ raise ValueError(f"Invalid cached MiniMax H3 condition for job {job.index}.")
881
+ return embeds, tags, _sha256(condition_path)
882
+
883
+
884
+ def output_is_complete(
885
+ output: Path,
886
+ job: GenerationJob,
887
+ *,
888
+ lora_sha256: str | None,
889
+ condition_sha256: str | None,
890
+ inference_steps: int,
891
+ ff_chunk_size: int,
892
+ lora_fused: bool,
893
+ width: int,
894
+ height: int,
895
+ num_frames: int,
896
+ group_offload_blocks: int,
897
+ group_offload_stream: bool,
898
+ allow_aligned_15s: bool,
899
+ final_step_size_scale: float,
900
+ video_shift: float,
901
+ audio_shift: float,
902
+ allow_variable_spatial_geometry: bool | None = None,
903
+ base_model_id: str | None = None,
904
+ model_revision: str | None = None,
905
+ env_profile_version: str | None = None,
906
+ generation_protocol_version: str | None = None,
907
+ staging_protocol_version: str | None = None,
908
+ prompt_manifest_sha256: str | None = None,
909
+ prompt_cache_audit_sha256: str | None = None,
910
+ publish_mode: str | None = None,
911
+ lora_scale: float | None = None,
912
+ task: str | None = None,
913
+ reference_inputs: list[dict[str, Any]] | None = None,
914
+ reference_inputs_audit_sha256: str | None = None,
915
+ reference_image_short_edge: int | None = None,
916
+ ) -> bool:
917
+ metadata_path = output.with_suffix(".json")
918
+ if not is_structurally_valid_mp4(output) or not metadata_path.is_file():
919
+ return False
920
+ try:
921
+ metadata = json.loads(metadata_path.read_text(encoding="utf-8"))
922
+ except (OSError, json.JSONDecodeError):
923
+ return False
924
+ expected = {
925
+ "complete": True,
926
+ "prompt": job.prompt,
927
+ "seed": job.seed,
928
+ "nfe": inference_steps,
929
+ "ff_chunk_size": ff_chunk_size,
930
+ "low_peak_rope": True,
931
+ "explicit_transformer_offload": True,
932
+ "final_step_size_scale": final_step_size_scale,
933
+ "video_shift": video_shift,
934
+ "audio_shift": audio_shift,
935
+ "lora_sha256": lora_sha256,
936
+ "lora_fused": lora_fused,
937
+ "condition_sha256": condition_sha256,
938
+ "width": width,
939
+ "height": height,
940
+ "num_frames": num_frames,
941
+ "group_offload_blocks": group_offload_blocks,
942
+ "group_offload_stream": group_offload_stream,
943
+ "allow_aligned_15s": allow_aligned_15s,
944
+ }
945
+ optional_expected = {
946
+ "base_model_id": base_model_id,
947
+ "model_revision": model_revision,
948
+ "env_profile_version": env_profile_version,
949
+ "generation_protocol_version": generation_protocol_version,
950
+ "staging_protocol_version": staging_protocol_version,
951
+ "prompt_manifest_sha256": prompt_manifest_sha256,
952
+ "prompt_cache_audit_sha256": prompt_cache_audit_sha256,
953
+ "publish_mode": publish_mode,
954
+ "lora_scale": lora_scale,
955
+ "allow_variable_spatial_geometry": allow_variable_spatial_geometry,
956
+ "task": task,
957
+ "reference_inputs": reference_inputs,
958
+ "reference_inputs_audit_sha256": reference_inputs_audit_sha256,
959
+ "reference_image_short_edge": reference_image_short_edge,
960
+ }
961
+ expected.update({key: value for key, value in optional_expected.items() if value is not None})
962
+ if any(metadata.get(key) != value for key, value in expected.items()):
963
+ return False
964
+ return metadata.get("output_sha256") == _sha256(output)
965
+
966
+
967
+ def parse_args() -> argparse.Namespace:
968
+ parser = argparse.ArgumentParser()
969
+ parser.add_argument("--jobs-json", type=Path, required=True)
970
+ parser.add_argument("--model-id", required=True)
971
+ parser.add_argument("--base-model-id", default=DEFAULT_BASE_MODEL_ID)
972
+ parser.add_argument("--model-revision", default=DEFAULT_MODEL_REVISION)
973
+ parser.add_argument("--output-dir", type=Path, required=True)
974
+ parser.add_argument("--reference-model-root", type=Path)
975
+ parser.add_argument("--lora-path", type=Path)
976
+ parser.add_argument("--lora-scale", type=float, default=1.0)
977
+ parser.add_argument("--ff-chunk-size", type=int, default=2048)
978
+ parser.add_argument("--qk-norm-chunk-size", type=int, default=0)
979
+ parser.add_argument("--prompt-cache", type=Path)
980
+ parser.add_argument("--fuse-lora", action="store_true")
981
+ parser.add_argument("--seed", type=int, default=0)
982
+ parser.add_argument("--height", type=int, default=DEFAULT_HEIGHT)
983
+ parser.add_argument("--width", type=int, default=DEFAULT_WIDTH)
984
+ parser.add_argument("--num-frames", type=int, default=DEFAULT_FRAMES)
985
+ parser.add_argument("--inference-steps", type=int, default=DEFAULT_NFE)
986
+ parser.add_argument("--final-step-size-scale", type=float, default=0.5)
987
+ parser.add_argument("--video-shift", type=float, default=12.0)
988
+ parser.add_argument("--audio-shift", type=float, default=3.0)
989
+ parser.add_argument("--device", default="cuda")
990
+ parser.add_argument("--memory-reserve-margin", default="12GB")
991
+ parser.add_argument("--no-cpu-offload", action="store_true")
992
+ parser.add_argument("--fsdp2", action="store_true")
993
+ parser.add_argument("--attention-backend")
994
+ parser.add_argument("--group-offload-blocks", type=int, default=0)
995
+ parser.add_argument("--group-offload-stream", action="store_true")
996
+ parser.add_argument("--allow-aligned-15s", action="store_true")
997
+ parser.add_argument("--allow-variable-spatial-geometry", action=argparse.BooleanOptionalAction, default=False)
998
+ parser.add_argument("--staging-dir", type=Path)
999
+ parser.add_argument("--env-profile-version")
1000
+ parser.add_argument("--generation-protocol-version")
1001
+ parser.add_argument("--staging-protocol-version", default=DEFAULT_STAGING_PROTOCOL_VERSION)
1002
+ parser.add_argument("--prompt-manifest-sha256")
1003
+ parser.add_argument("--prompt-cache-audit-sha256")
1004
+ parser.add_argument("--reference-inputs-audit-sha256")
1005
+ parser.add_argument("--reference-image-short-edge", type=int)
1006
+ parser.add_argument("--job-shard-count", type=int, default=1)
1007
+ parser.add_argument("--job-shard-index", type=int, default=0)
1008
+ parser.add_argument("--resume", action=argparse.BooleanOptionalAction, default=True)
1009
+ parser.add_argument("--dry-run", action="store_true")
1010
+ return parser.parse_args()
1011
+
1012
+
1013
+ def main() -> None:
1014
+ args = parse_args()
1015
+ ensure_qwen3vl_mm_token_type_ids()
1016
+ ensure_qwen3vl_hidden_state_capture()
1017
+ if args.height <= 0 or args.width <= 0 or args.height % 32 or args.width % 32:
1018
+ raise ValueError("--height and --width must be positive multiples of 32.")
1019
+ if args.num_frames < 5 or (args.num_frames - 5) % 17:
1020
+ raise ValueError("--num-frames must satisfy 17*n+5.")
1021
+ if args.inference_steps < 1:
1022
+ raise ValueError("--inference-steps must be positive.")
1023
+ if args.final_step_size_scale <= 0:
1024
+ raise ValueError("--final-step-size-scale must be positive.")
1025
+ if args.video_shift <= 0 or args.audio_shift <= 0:
1026
+ raise ValueError("--video-shift and --audio-shift must be positive.")
1027
+ if args.lora_scale < 0:
1028
+ raise ValueError("--lora-scale must be non-negative.")
1029
+ if args.ff_chunk_size < 0:
1030
+ raise ValueError("--ff-chunk-size must be non-negative.")
1031
+ if args.qk_norm_chunk_size < 0:
1032
+ raise ValueError("--qk-norm-chunk-size must be non-negative.")
1033
+ if args.group_offload_blocks < 0:
1034
+ raise ValueError("--group-offload-blocks must be non-negative.")
1035
+ if args.job_shard_count < 1:
1036
+ raise ValueError("--job-shard-count must be positive.")
1037
+ if not 0 <= args.job_shard_index < args.job_shard_count:
1038
+ raise ValueError("--job-shard-index must be smaller than --job-shard-count.")
1039
+ args.jobs_json = args.jobs_json.expanduser().resolve()
1040
+ args.output_dir = args.output_dir.expanduser().resolve()
1041
+ if args.staging_dir is not None:
1042
+ args.staging_dir = args.staging_dir.expanduser().resolve()
1043
+ args.model_id = (
1044
+ os.fspath(Path(args.model_id).expanduser().resolve()) if Path(args.model_id).exists() else args.model_id
1045
+ )
1046
+ if args.prompt_cache is not None:
1047
+ args.prompt_cache = args.prompt_cache.expanduser().resolve()
1048
+ if not args.prompt_cache.is_dir():
1049
+ raise FileNotFoundError(f"Prompt cache does not exist: {args.prompt_cache}")
1050
+ if args.lora_path is not None:
1051
+ args.lora_path = args.lora_path.expanduser().resolve()
1052
+ if not args.lora_path.is_file():
1053
+ raise FileNotFoundError(f"LoRA file does not exist: {args.lora_path}.")
1054
+ lora_sha256 = None if args.lora_path is None else _sha256(args.lora_path)
1055
+
1056
+ context = distributed_context(args.fsdp2)
1057
+ jobs = load_jobs(args.jobs_json, args.seed)
1058
+ tasks = {job.task for job in jobs}
1059
+ if len(tasks) != 1:
1060
+ raise ValueError(f"All MiniMax H3 jobs in one launch must use the same task, got {sorted(tasks)}.")
1061
+ task = tasks.pop()
1062
+ if task == "ref2va" and args.prompt_cache is not None:
1063
+ raise ValueError(
1064
+ "Ref2VA prompt presentations depend on the references and cannot use the T2VA text-only cache."
1065
+ )
1066
+ validate_reference_image_policy(jobs, args.reference_image_short_edge)
1067
+ for job in jobs:
1068
+ width = job.width or args.width
1069
+ height = job.height or args.height
1070
+ num_frames = job.num_frames or args.num_frames
1071
+ try:
1072
+ validate_resolved_geometry(
1073
+ width,
1074
+ height,
1075
+ num_frames,
1076
+ allow_variable_spatial_geometry=args.allow_variable_spatial_geometry,
1077
+ task=job.task,
1078
+ )
1079
+ except ValueError as error:
1080
+ raise ValueError(f"Invalid resolved geometry for job {job.index}: {error}") from error
1081
+ total_jobs = len(jobs)
1082
+ jobs = select_job_shard(jobs, args.job_shard_index, args.job_shard_count)
1083
+ if not jobs:
1084
+ print(
1085
+ f"jobs=0/{total_jobs} job_shard={args.job_shard_index}/{args.job_shard_count}; "
1086
+ "empty original-index shard requires no model load.",
1087
+ flush=True,
1088
+ )
1089
+ return
1090
+ args.output_dir.mkdir(parents=True, exist_ok=True)
1091
+ if context.enabled and task == "ref2va" and args.resume:
1092
+ pending_jobs = []
1093
+ for job in jobs:
1094
+ output_rank = job.index % context.world_size
1095
+ if not output_is_complete(
1096
+ result_path(args.output_dir, job, output_rank),
1097
+ job,
1098
+ lora_sha256=lora_sha256,
1099
+ condition_sha256=None,
1100
+ inference_steps=args.inference_steps,
1101
+ ff_chunk_size=args.ff_chunk_size,
1102
+ lora_fused=args.fuse_lora,
1103
+ width=job.width or args.width,
1104
+ height=job.height or args.height,
1105
+ num_frames=job.num_frames or args.num_frames,
1106
+ group_offload_blocks=0,
1107
+ group_offload_stream=False,
1108
+ allow_aligned_15s=args.allow_aligned_15s,
1109
+ allow_variable_spatial_geometry=args.allow_variable_spatial_geometry,
1110
+ final_step_size_scale=args.final_step_size_scale,
1111
+ video_shift=args.video_shift,
1112
+ audio_shift=args.audio_shift,
1113
+ base_model_id=args.base_model_id,
1114
+ model_revision=args.model_revision,
1115
+ env_profile_version=args.env_profile_version,
1116
+ generation_protocol_version=args.generation_protocol_version,
1117
+ staging_protocol_version=args.staging_protocol_version,
1118
+ prompt_manifest_sha256=args.prompt_manifest_sha256,
1119
+ prompt_cache_audit_sha256=None,
1120
+ publish_mode=(
1121
+ args.staging_protocol_version if args.staging_dir is not None else "adjacent-atomic-replace"
1122
+ ),
1123
+ lora_scale=args.lora_scale,
1124
+ task=job.task,
1125
+ reference_inputs=[reference.identity() for reference in job.references],
1126
+ reference_inputs_audit_sha256=args.reference_inputs_audit_sha256,
1127
+ reference_image_short_edge=args.reference_image_short_edge,
1128
+ ):
1129
+ pending_jobs.append(job)
1130
+ if context.rank == 0:
1131
+ print(f"distributed resume pending={len(pending_jobs)}/{len(jobs)}", flush=True)
1132
+ jobs = pending_jobs
1133
+ if not jobs:
1134
+ return
1135
+ assigned = assign_jobs(jobs, context)
1136
+ if context.rank == 0:
1137
+ print(
1138
+ f"jobs={len(jobs)}/{total_jobs} job_shard={args.job_shard_index}/{args.job_shard_count} "
1139
+ f"world_size={context.world_size} default_geometry="
1140
+ f"{args.width}x{args.height}x{args.num_frames} NFE={args.inference_steps} "
1141
+ f"final_scale={args.final_step_size_scale:g} shifts={args.video_shift:g}/{args.audio_shift:g}",
1142
+ flush=True,
1143
+ )
1144
+ for job in assigned:
1145
+ print(
1146
+ f"[rank {context.rank}] index={job.index} seed={job.seed} padding={job.is_padding} "
1147
+ f"geometry={job.width or args.width}x{job.height or args.height}x{job.num_frames or args.num_frames} "
1148
+ f"prompt={job.prompt!r}",
1149
+ flush=True,
1150
+ )
1151
+ if args.dry_run:
1152
+ return
1153
+
1154
+ condition_hashes: dict[int, str | None] = {}
1155
+ if args.prompt_cache is not None:
1156
+ for job in assigned:
1157
+ _, _, condition_hashes[job.index] = load_cached_condition(args.prompt_cache, job)
1158
+ for job in assigned:
1159
+ if not job.is_padding:
1160
+ cleanup_result_partials(result_path(args.output_dir, job, context.rank))
1161
+ if not context.enabled and args.resume:
1162
+ assigned = [
1163
+ job
1164
+ for job in assigned
1165
+ if not output_is_complete(
1166
+ result_path(args.output_dir, job, context.rank),
1167
+ job,
1168
+ lora_sha256=lora_sha256,
1169
+ condition_sha256=condition_hashes.get(job.index),
1170
+ inference_steps=args.inference_steps,
1171
+ ff_chunk_size=args.ff_chunk_size,
1172
+ lora_fused=args.fuse_lora,
1173
+ width=job.width or args.width,
1174
+ height=job.height or args.height,
1175
+ num_frames=job.num_frames or args.num_frames,
1176
+ group_offload_blocks=args.group_offload_blocks,
1177
+ group_offload_stream=args.group_offload_stream,
1178
+ allow_aligned_15s=args.allow_aligned_15s,
1179
+ allow_variable_spatial_geometry=args.allow_variable_spatial_geometry,
1180
+ final_step_size_scale=args.final_step_size_scale,
1181
+ video_shift=args.video_shift,
1182
+ audio_shift=args.audio_shift,
1183
+ base_model_id=args.base_model_id,
1184
+ model_revision=args.model_revision,
1185
+ env_profile_version=args.env_profile_version,
1186
+ generation_protocol_version=args.generation_protocol_version,
1187
+ staging_protocol_version=args.staging_protocol_version,
1188
+ prompt_manifest_sha256=args.prompt_manifest_sha256,
1189
+ prompt_cache_audit_sha256=args.prompt_cache_audit_sha256,
1190
+ publish_mode=(
1191
+ args.staging_protocol_version if args.staging_dir is not None else "adjacent-atomic-replace"
1192
+ ),
1193
+ lora_scale=args.lora_scale,
1194
+ task=job.task,
1195
+ reference_inputs=[reference.identity() for reference in job.references],
1196
+ reference_inputs_audit_sha256=args.reference_inputs_audit_sha256,
1197
+ reference_image_short_edge=args.reference_image_short_edge,
1198
+ )
1199
+ ]
1200
+ print(f"[rank {context.rank}] resume pending={len(assigned)}/{len(jobs)}", flush=True)
1201
+ if not assigned:
1202
+ return
1203
+ mesh = initialize_fsdp2(context)
1204
+ try:
1205
+ pipe = load_pipeline(args, context, mesh, task)
1206
+ scheduler_grid_points = args.inference_steps + 1
1207
+ for job in assigned:
1208
+ job_width = job.width or args.width
1209
+ job_height = job.height or args.height
1210
+ job_num_frames = job.num_frames or args.num_frames
1211
+ generator = torch.Generator(device=context.device).manual_seed(job.seed)
1212
+ condition_kwargs = {}
1213
+ if args.prompt_cache is not None:
1214
+ prompt_embeds, text_token_tags, condition_hash = load_cached_condition(args.prompt_cache, job)
1215
+ condition_kwargs = {
1216
+ "prompt_embeds": prompt_embeds,
1217
+ "text_token_tags": text_token_tags,
1218
+ }
1219
+ else:
1220
+ condition_kwargs = {"prompt": job.prompt}
1221
+ condition_hash = None
1222
+ if job.task == "ref2va":
1223
+ condition_kwargs["references"] = load_job_references(job)
1224
+ with torch.inference_mode():
1225
+ result = pipe(
1226
+ **condition_kwargs,
1227
+ height=job_height,
1228
+ width=job_width,
1229
+ num_frames=job_num_frames,
1230
+ num_inference_steps=scheduler_grid_points,
1231
+ generator=generator,
1232
+ output_type="np",
1233
+ output=["videos", "audio", "sampling_rate"],
1234
+ )
1235
+ if not job.is_padding:
1236
+ output = result_path(args.output_dir, job, context.rank)
1237
+ output_sha256, publish_mode = publish_result(result, output, args.staging_dir)
1238
+ metadata = {
1239
+ "schema_version": 1,
1240
+ "complete": True,
1241
+ "task": job.task,
1242
+ "case_name": job.case_name,
1243
+ "source_request": job.source_request,
1244
+ "item_id": job.item_id,
1245
+ "category": job.category,
1246
+ "prompt": job.prompt,
1247
+ "seed": job.seed,
1248
+ "width": job_width,
1249
+ "height": job_height,
1250
+ "num_frames": job_num_frames,
1251
+ "fps": FPS,
1252
+ "nfe": args.inference_steps,
1253
+ "ff_chunk_size": args.ff_chunk_size,
1254
+ "qk_norm_chunk_size": args.qk_norm_chunk_size,
1255
+ "low_peak_rope": True,
1256
+ "explicit_transformer_offload": True,
1257
+ "scheduler_grid_points": scheduler_grid_points,
1258
+ "final_step_size_scale": args.final_step_size_scale,
1259
+ "raw_sigmas": raw_sigma_grid(args.inference_steps, args.final_step_size_scale),
1260
+ "video_shift": args.video_shift,
1261
+ "audio_shift": args.audio_shift,
1262
+ "model_id": args.model_id,
1263
+ "base_model_id": args.base_model_id,
1264
+ "model_revision": args.model_revision,
1265
+ "lora_path": None if args.lora_path is None else os.fspath(args.lora_path),
1266
+ "lora_sha256": lora_sha256,
1267
+ "lora_scale": args.lora_scale,
1268
+ "lora_fused": args.fuse_lora,
1269
+ "group_offload_blocks": args.group_offload_blocks,
1270
+ "group_offload_stream": args.group_offload_stream,
1271
+ "allow_aligned_15s": args.allow_aligned_15s,
1272
+ "allow_variable_spatial_geometry": args.allow_variable_spatial_geometry,
1273
+ "condition_sha256": condition_hash,
1274
+ "env_profile_version": args.env_profile_version,
1275
+ "generation_protocol_version": args.generation_protocol_version,
1276
+ "staging_protocol_version": args.staging_protocol_version,
1277
+ "publish_mode": publish_mode,
1278
+ "prompt_manifest_sha256": args.prompt_manifest_sha256,
1279
+ "prompt_cache_audit_sha256": args.prompt_cache_audit_sha256,
1280
+ "reference_inputs": [reference.identity() for reference in job.references],
1281
+ "reference_inputs_audit_sha256": args.reference_inputs_audit_sha256,
1282
+ "reference_image_short_edge": args.reference_image_short_edge,
1283
+ "output": output.name,
1284
+ "output_sha256": output_sha256,
1285
+ }
1286
+ _atomic_write_json(output.with_suffix(".json"), metadata)
1287
+ print(f"[rank {context.rank}] wrote {output}", flush=True)
1288
+ del result, generator
1289
+ gc.collect()
1290
+ torch.cuda.empty_cache()
1291
+ if context.enabled:
1292
+ dist.barrier()
1293
+ finally:
1294
+ if dist.is_initialized():
1295
+ dist.destroy_process_group()
1296
+
1297
+
1298
+ if __name__ == "__main__":
1299
+ main()
tools/minimax_h3_transformers_compat.py ADDED
@@ -0,0 +1,99 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # SPDX-License-Identifier: Apache-2.0
2
+
3
+ """Compatibility for the Qwen3-VL processor used by pinned MiniMax H3 jobs."""
4
+
5
+ from __future__ import annotations
6
+
7
+ from typing import Any
8
+
9
+ import numpy as np
10
+ import torch
11
+ from transformers import Qwen3VLProcessor
12
+
13
+
14
+ def _token_ids(processor: Any, plural: str, singular: str) -> list[int]:
15
+ values = getattr(processor, plural, None)
16
+ if values is not None:
17
+ return [int(value) for value in values if value is not None]
18
+ value = getattr(processor, singular, None)
19
+ return [] if value is None else [int(value)]
20
+
21
+
22
+ def _create_mm_token_type_ids(processor: Any, input_ids: list[list[int]]) -> list[list[int]]:
23
+ image_ids = _token_ids(processor, "image_token_ids", "image_token_id")
24
+ video_ids = _token_ids(processor, "video_token_ids", "video_token_id")
25
+ result = []
26
+ for sequence in input_ids:
27
+ token_ids = np.asarray(sequence)
28
+ token_types = np.zeros_like(token_ids)
29
+ token_types[np.isin(token_ids, image_ids)] = 1
30
+ token_types[np.isin(token_ids, video_ids)] = 2
31
+ result.append(token_types.tolist())
32
+ return result
33
+
34
+
35
+ def ensure_qwen3vl_mm_token_type_ids() -> None:
36
+ if not hasattr(Qwen3VLProcessor, "create_mm_token_type_ids"):
37
+ Qwen3VLProcessor.create_mm_token_type_ids = _create_mm_token_type_ids
38
+
39
+
40
+ def get_qwen3vl_prompt_embeds_with_layer_capture(
41
+ text_encoder,
42
+ processor,
43
+ token_ids: list[int],
44
+ vision_inputs: dict | None = None,
45
+ text_encoder_layer: int = 50,
46
+ device: torch.device | None = None,
47
+ dtype: torch.dtype | None = None,
48
+ ) -> torch.Tensor:
49
+ """Capture Qwen3-VL's pre-norm intermediate state without Transformers output hooks."""
50
+
51
+ num_layers = text_encoder.config.text_config.num_hidden_layers
52
+ if not 0 < text_encoder_layer < num_layers:
53
+ raise ValueError(f"MiniMax H3 requires an intermediate layer in (0, {num_layers}), got {text_encoder_layer}.")
54
+ input_ids = torch.tensor([token_ids], dtype=torch.long, device=device)
55
+ mm_token_type_ids = torch.tensor(
56
+ processor.create_mm_token_type_ids([token_ids]),
57
+ dtype=torch.long,
58
+ device=device,
59
+ )
60
+ vision_kwargs = {}
61
+ for name, value in (vision_inputs or {}).items():
62
+ vision_kwargs[name] = value.to(device, text_encoder.dtype) if name.startswith("pixel_") else value.to(device)
63
+
64
+ captured: dict[str, torch.Tensor] = {}
65
+
66
+ def capture_layer(_module, _inputs, output):
67
+ hidden_states = output[0] if isinstance(output, (tuple, list)) else output
68
+ if not isinstance(hidden_states, torch.Tensor):
69
+ raise TypeError(f"Qwen3-VL layer returned unsupported output type {type(hidden_states)!r}.")
70
+ captured["hidden_states"] = hidden_states
71
+
72
+ layer = text_encoder.model.language_model.layers[text_encoder_layer - 1]
73
+ handle = layer.register_forward_hook(capture_layer)
74
+ try:
75
+ hook = getattr(text_encoder, "_hf_hook", None)
76
+ if hook is not None and hasattr(hook, "pre_forward"):
77
+ hook.pre_forward(text_encoder)
78
+ text_encoder.model(
79
+ input_ids=input_ids,
80
+ attention_mask=torch.ones_like(input_ids),
81
+ mm_token_type_ids=mm_token_type_ids,
82
+ use_cache=False,
83
+ output_hidden_states=False,
84
+ **vision_kwargs,
85
+ )
86
+ finally:
87
+ handle.remove()
88
+ hidden_states = captured.get("hidden_states")
89
+ if hidden_states is None:
90
+ raise RuntimeError(f"Qwen3-VL decoder layer {text_encoder_layer - 1} did not produce a captured state.")
91
+ return hidden_states.to(device=device, dtype=dtype)
92
+
93
+
94
+ def ensure_qwen3vl_hidden_state_capture() -> None:
95
+ """Install an FSDP-safe MiniMax H3 layer-50 capture helper into Diffusers."""
96
+
97
+ from diffusers.modular_pipelines.minimax_h3 import encoders
98
+
99
+ encoders.get_qwen3vl_prompt_embeds = get_qwen3vl_prompt_embeds_with_layer_capture