Image-to-Video
Diffusers
Safetensors
MiniMax H3
English
smd
diffusion-distillation
few-step
lora
minimax
video-generation
audio-video
4-step
reference-to-video
Instructions to use zjwfufu/SMD-MiniMax-H3-Ref2VA-4step with libraries, inference providers, notebooks, and local apps. Follow these links to get started.
- Libraries
- Diffusers
How to use zjwfufu/SMD-MiniMax-H3-Ref2VA-4step with Diffusers:
pip install -U diffusers transformers accelerate
import torch from diffusers import DiffusionPipeline from diffusers.utils import load_image, export_to_video # switch to "mps" for apple devices pipe = DiffusionPipeline.from_pretrained("MiniMaxAI/MiniMax-H3", dtype=torch.bfloat16, device_map="cuda") pipe.load_lora_weights("zjwfufu/SMD-MiniMax-H3-Ref2VA-4step") prompt = "A man with short gray hair plays a red electric guitar." input_image = load_image("https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/guitar-man.png") output = pipe(image=input_image, prompt=prompt).frames[0] export_to_video(output, "output.mp4") - Inference
- Notebooks
- Google Colab
- Kaggle
- Local Apps Settings
- Draw Things
Commit ·
8f67bf6
0
Parent(s):
Initial commit
Browse filesSigned-off-by: zjwfufu <1174987765@qq.com>
- .gitattributes +41 -0
- LICENSE-code +201 -0
- README.md +86 -0
- examples/case.json +22 -0
- examples/generation.json +46 -0
- examples/layout.json +9 -0
- examples/poster.png +3 -0
- examples/ref_dog.png +3 -0
- examples/ref_woman.png +3 -0
- examples/reference_generation.json +38 -0
- examples/references.png +3 -0
- examples/sample.mp4 +3 -0
- h3_lora.py +117 -0
- infer.py +26 -0
- lora/adapter_config.json +16 -0
- lora/adapter_model.safetensors +3 -0
- requirements.txt +14 -0
- sampling.json +20 -0
- tools/__init__.py +0 -0
- tools/infer_minimax_h3.py +1299 -0
- tools/minimax_h3_transformers_compat.py +99 -0
.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
|
examples/ref_dog.png
ADDED
|
Git LFS Details
|
examples/ref_woman.png
ADDED
|
Git LFS Details
|
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
|
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
|