├── .gitignore
├── LICENSE
├── README.md
├── demo_gradio.py
├── diffusers_helper
├── bucket_tools.py
├── clip_vision.py
├── dit_common.py
├── gradio
│ └── progress_bar.py
├── hf_login.py
├── hunyuan.py
├── k_diffusion
│ ├── uni_pc_fm.py
│ └── wrapper.py
├── memory.py
├── models
│ └── hunyuan_video_packed.py
├── pipelines
│ └── k_diffusion_hunyuan.py
├── thread_utils.py
└── utils.py
└── requirements.txt
/.gitignore:
--------------------------------------------------------------------------------
1 | hf_download/
2 | outputs/
3 | repo/
4 |
5 | # Byte-compiled / optimized / DLL files
6 | __pycache__/
7 | *.py[cod]
8 | *$py.class
9 |
10 | # C extensions
11 | *.so
12 |
13 | # Distribution / packaging
14 | .Python
15 | build/
16 | develop-eggs/
17 | dist/
18 | downloads/
19 | eggs/
20 | .eggs/
21 | lib/
22 | lib64/
23 | parts/
24 | sdist/
25 | var/
26 | wheels/
27 | share/python-wheels/
28 | *.egg-info/
29 | .installed.cfg
30 | *.egg
31 | MANIFEST
32 |
33 | # PyInstaller
34 | # Usually these files are written by a python script from a template
35 | # before PyInstaller builds the exe, so as to inject date/other infos into it.
36 | *.manifest
37 | *.spec
38 |
39 | # Installer logs
40 | pip-log.txt
41 | pip-delete-this-directory.txt
42 |
43 | # Unit test / coverage reports
44 | htmlcov/
45 | .tox/
46 | .nox/
47 | .coverage
48 | .coverage.*
49 | .cache
50 | nosetests.xml
51 | coverage.xml
52 | *.cover
53 | *.py,cover
54 | .hypothesis/
55 | .pytest_cache/
56 | cover/
57 |
58 | # Translations
59 | *.mo
60 | *.pot
61 |
62 | # Django stuff:
63 | *.log
64 | local_settings.py
65 | db.sqlite3
66 | db.sqlite3-journal
67 |
68 | # Flask stuff:
69 | instance/
70 | .webassets-cache
71 |
72 | # Scrapy stuff:
73 | .scrapy
74 |
75 | # Sphinx documentation
76 | docs/_build/
77 |
78 | # PyBuilder
79 | .pybuilder/
80 | target/
81 |
82 | # Jupyter Notebook
83 | .ipynb_checkpoints
84 |
85 | # IPython
86 | profile_default/
87 | ipython_config.py
88 |
89 | # pyenv
90 | # For a library or package, you might want to ignore these files since the code is
91 | # intended to run in multiple environments; otherwise, check them in:
92 | # .python-version
93 |
94 | # pipenv
95 | # According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control.
96 | # However, in case of collaboration, if having platform-specific dependencies or dependencies
97 | # having no cross-platform support, pipenv may install dependencies that don't work, or not
98 | # install all needed dependencies.
99 | #Pipfile.lock
100 |
101 | # UV
102 | # Similar to Pipfile.lock, it is generally recommended to include uv.lock in version control.
103 | # This is especially recommended for binary packages to ensure reproducibility, and is more
104 | # commonly ignored for libraries.
105 | #uv.lock
106 |
107 | # poetry
108 | # Similar to Pipfile.lock, it is generally recommended to include poetry.lock in version control.
109 | # This is especially recommended for binary packages to ensure reproducibility, and is more
110 | # commonly ignored for libraries.
111 | # https://python-poetry.org/docs/basic-usage/#commit-your-poetrylock-file-to-version-control
112 | #poetry.lock
113 |
114 | # pdm
115 | # Similar to Pipfile.lock, it is generally recommended to include pdm.lock in version control.
116 | #pdm.lock
117 | # pdm stores project-wide configurations in .pdm.toml, but it is recommended to not include it
118 | # in version control.
119 | # https://pdm.fming.dev/latest/usage/project/#working-with-version-control
120 | .pdm.toml
121 | .pdm-python
122 | .pdm-build/
123 |
124 | # PEP 582; used by e.g. github.com/David-OConnor/pyflow and github.com/pdm-project/pdm
125 | __pypackages__/
126 |
127 | # Celery stuff
128 | celerybeat-schedule
129 | celerybeat.pid
130 |
131 | # SageMath parsed files
132 | *.sage.py
133 |
134 | # Environments
135 | .env
136 | .venv
137 | env/
138 | venv/
139 | ENV/
140 | env.bak/
141 | venv.bak/
142 |
143 | # Spyder project settings
144 | .spyderproject
145 | .spyproject
146 |
147 | # Rope project settings
148 | .ropeproject
149 |
150 | # mkdocs documentation
151 | /site
152 |
153 | # mypy
154 | .mypy_cache/
155 | .dmypy.json
156 | dmypy.json
157 |
158 | # Pyre type checker
159 | .pyre/
160 |
161 | # pytype static type analyzer
162 | .pytype/
163 |
164 | # Cython debug symbols
165 | cython_debug/
166 |
167 | # PyCharm
168 | # JetBrains specific template is maintained in a separate JetBrains.gitignore that can
169 | # be found at https://github.com/github/gitignore/blob/main/Global/JetBrains.gitignore
170 | # and can be added to the global gitignore or merged into this file. For a more nuclear
171 | # option (not recommended) you can uncomment the following to ignore the entire idea folder.
172 | .idea/
173 |
174 | # Ruff stuff:
175 | .ruff_cache/
176 |
177 | # PyPI configuration file
178 | .pypirc
179 |
--------------------------------------------------------------------------------
/LICENSE:
--------------------------------------------------------------------------------
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.
202 |
--------------------------------------------------------------------------------
/README.md:
--------------------------------------------------------------------------------
1 |
2 | 
3 |
4 |
5 | # FramePack
6 |
7 | Official implementation and desktop software for ["Packing Input Frame Context in Next-Frame Prediction Models for Video Generation"](https://lllyasviel.github.io/frame_pack_gitpage/).
8 |
9 | Links: [**Paper**](https://arxiv.org/abs/2504.12626), [**Project Page**](https://lllyasviel.github.io/frame_pack_gitpage/)
10 |
11 | FramePack is a next-frame (next-frame-section) prediction neural network structure that generates videos progressively.
12 |
13 | FramePack compresses input contexts to a constant length so that the generation workload is invariant to video length.
14 |
15 | FramePack can process a very large number of frames with 13B models even on laptop GPUs.
16 |
17 | FramePack can be trained with a much larger batch size, similar to the batch size for image diffusion training.
18 |
19 | **Video diffusion, but feels like image diffusion.**
20 |
21 | # Notes
22 |
23 | Note that this GitHub repository is the only official FramePack website. We do not have any web services. All other websites are spam and fake, including but not limited to `framepack.co`, `frame_pack.co`, `framepack.net`, `frame_pack.net`, `framepack.ai`, `frame_pack.ai`, `framepack.pro`, `frame_pack.pro`, `framepack.cc`, `frame_pack.cc`,`framepackai.co`, `frame_pack_ai.co`, `framepackai.net`, `frame_pack_ai.net`, `framepackai.pro`, `frame_pack_ai.pro`, `framepackai.cc`, `frame_pack_ai.cc`, and so on. Again, they are all spam and fake. **Do not pay money or download files from any of those websites.**
24 |
25 | The team is on leave between April 21 and 29. PR merging will be delayed.
26 |
27 | # Requirements
28 |
29 | Note that this repo is a functional desktop software with minimal standalone high-quality sampling system and memory management.
30 |
31 | **Start with this repo before you try anything else!**
32 |
33 | Requirements:
34 |
35 | * Nvidia GPU in RTX 30XX, 40XX, 50XX series that supports fp16 and bf16. The GTX 10XX/20XX are not tested.
36 | * Linux or Windows operating system.
37 | * At least 6GB GPU memory.
38 |
39 | To generate 1-minute video (60 seconds) at 30fps (1800 frames) using 13B model, the minimal required GPU memory is 6GB. (Yes 6 GB, not a typo. Laptop GPUs are okay.)
40 |
41 | About speed, on my RTX 4090 desktop it generates at a speed of 2.5 seconds/frame (unoptimized) or 1.5 seconds/frame (teacache). On my laptops like 3070ti laptop or 3060 laptop, it is about 4x to 8x slower. [Troubleshoot if your speed is much slower than this.](https://github.com/lllyasviel/FramePack/issues/151#issuecomment-2817054649)
42 |
43 | In any case, you will directly see the generated frames since it is next-frame(-section) prediction. So you will get lots of visual feedback before the entire video is generated.
44 |
45 | # Installation
46 |
47 | **Windows**:
48 |
49 | [>>> Click Here to Download One-Click Package (CUDA 12.6 + Pytorch 2.6) <<<](https://github.com/lllyasviel/FramePack/releases/download/windows/framepack_cu126_torch26.7z)
50 |
51 | After you download, you uncompress, use `update.bat` to update, and use `run.bat` to run.
52 |
53 | Note that running `update.bat` is important, otherwise you may be using a previous version with potential bugs unfixed.
54 |
55 | 
56 |
57 | Note that the models will be downloaded automatically. You will download more than 30GB from HuggingFace.
58 |
59 | **Linux**:
60 |
61 | We recommend having an independent Python 3.10.
62 |
63 | pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu126
64 | pip install -r requirements.txt
65 |
66 | To start the GUI, run:
67 |
68 | python demo_gradio.py
69 |
70 | Note that it supports `--share`, `--port`, `--server`, and so on.
71 |
72 | The software supports PyTorch attention, xformers, flash-attn, sage-attention. By default, it will just use PyTorch attention. You can install those attention kernels if you know how.
73 |
74 | For example, to install sage-attention (linux):
75 |
76 | pip install sageattention==1.0.6
77 |
78 | However, you are highly recommended to first try without sage-attention since it will influence results, though the influence is minimal.
79 |
80 | # GUI
81 |
82 | 
83 |
84 | On the left you upload an image and write a prompt.
85 |
86 | On the right are the generated videos and latent previews.
87 |
88 | Because this is a next-frame-section prediction model, videos will be generated longer and longer.
89 |
90 | You will see the progress bar for each section and the latent preview for the next section.
91 |
92 | Note that the initial progress may be slower than later diffusion as the device may need some warmup.
93 |
94 | # Sanity Check
95 |
96 | Before trying your own inputs, we highly recommend going through the sanity check to find out if any hardware or software went wrong.
97 |
98 | Next-frame-section prediction models are very sensitive to subtle differences in noise and hardware. Usually, people will get slightly different results on different devices, but the results should look overall similar. In some cases, if possible, you'll get exactly the same results.
99 |
100 | ## Image-to-5-seconds
101 |
102 | Download this image:
103 |
104 | 
105 |
106 | Copy this prompt:
107 |
108 | `The man dances energetically, leaping mid-air with fluid arm swings and quick footwork.`
109 |
110 | Set like this:
111 |
112 | (all default parameters, with teacache turned off)
113 | 
114 |
115 | The result will be:
116 |
117 |
118 |
119 | |
120 |
125 | |
126 |
127 |
128 | |
129 | Video may be compressed by GitHub
130 | |
131 |
132 |
133 |
134 | **Important Note:**
135 |
136 | Again, this is a next-frame-section prediction model. This means you will generate videos frame-by-frame or section-by-section.
137 |
138 | **If you get a much shorter video in the UI, like a video with only 1 second, then it is totally expected.** You just need to wait. More sections will be generated to complete the video.
139 |
140 | ## Know the influence of TeaCache and Quantization
141 |
142 | Download this image:
143 |
144 | 
145 |
146 | Copy this prompt:
147 |
148 | `The girl dances gracefully, with clear movements, full of charm.`
149 |
150 | Set like this:
151 |
152 | 
153 |
154 | Turn off teacache:
155 |
156 | 
157 |
158 | You will get this:
159 |
160 |
161 |
162 | |
163 |
168 | |
169 |
170 |
171 | |
172 | Video may be compressed by GitHub
173 | |
174 |
175 |
176 |
177 | Now turn on teacache:
178 |
179 | 
180 |
181 | About 30% users will get this (the other 70% will get other random results depending on their hardware):
182 |
183 |
184 |
185 | |
186 |
191 | |
192 |
193 |
194 | |
195 | A typical worse result.
196 | |
197 |
198 |
199 |
200 | So you can see that teacache is not really lossless and sometimes can influence the result a lot.
201 |
202 | We recommend using teacache to try ideas and then using the full diffusion process to get high-quality results.
203 |
204 | This recommendation also applies to sage-attention, bnb quant, gguf, etc., etc.
205 |
206 | ## Image-to-1-minute
207 |
208 | 
209 |
210 | `The girl dances gracefully, with clear movements, full of charm.`
211 |
212 | 
213 |
214 | Set video length to 60 seconds:
215 |
216 | 
217 |
218 | If everything is in order you will get some result like this eventually.
219 |
220 | 60s version:
221 |
222 |
223 |
224 | |
225 |
230 | |
231 |
232 |
233 | |
234 | Video may be compressed by GitHub
235 | |
236 |
237 |
238 |
239 | 6s version:
240 |
241 |
242 |
243 | |
244 |
249 | |
250 |
251 |
252 | |
253 | Video may be compressed by GitHub
254 | |
255 |
256 |
257 |
258 | # More Examples
259 |
260 | Many more examples are in [**Project Page**](https://lllyasviel.github.io/frame_pack_gitpage/).
261 |
262 | Below are some more examples that you may be interested in reproducing.
263 |
264 | ---
265 |
266 | 
267 |
268 | `The girl dances gracefully, with clear movements, full of charm.`
269 |
270 | 
271 |
272 |
273 |
274 | |
275 |
280 | |
281 |
282 |
283 | |
284 | Video may be compressed by GitHub
285 | |
286 |
287 |
288 |
289 | ---
290 |
291 | 
292 |
293 | `The girl suddenly took out a sign that said “cute” using right hand`
294 |
295 | 
296 |
297 |
298 |
299 | |
300 |
305 | |
306 |
307 |
308 | |
309 | Video may be compressed by GitHub
310 | |
311 |
312 |
313 |
314 | ---
315 |
316 | 
317 |
318 | `The girl skateboarding, repeating the endless spinning and dancing and jumping on a skateboard, with clear movements, full of charm.`
319 |
320 | 
321 |
322 |
323 |
324 | |
325 |
330 | |
331 |
332 |
333 | |
334 | Video may be compressed by GitHub
335 | |
336 |
337 |
338 |
339 | ---
340 |
341 | 
342 |
343 | `The girl dances gracefully, with clear movements, full of charm.`
344 |
345 | 
346 |
347 |
348 |
349 | |
350 |
355 | |
356 |
357 |
358 | |
359 | Video may be compressed by GitHub
360 | |
361 |
362 |
363 |
364 | ---
365 |
366 | 
367 |
368 | `The man dances flamboyantly, swinging his hips and striking bold poses with dramatic flair.`
369 |
370 | 
371 |
372 |
373 |
374 | |
375 |
380 | |
381 |
382 |
383 | |
384 | Video may be compressed by GitHub
385 | |
386 |
387 |
388 |
389 | ---
390 |
391 | 
392 |
393 | `The woman dances elegantly among the blossoms, spinning slowly with flowing sleeves and graceful hand movements.`
394 |
395 | 
396 |
397 |
398 |
399 |
400 | |
401 |
406 | |
407 |
408 |
409 | |
410 | Video may be compressed by GitHub
411 | |
412 |
413 |
414 |
415 | ---
416 |
417 | 
418 |
419 | `The young man writes intensely, flipping papers and adjusting his glasses with swift, focused movements.`
420 |
421 | 
422 |
423 |
424 |
425 | |
426 |
431 | |
432 |
433 |
434 | |
435 | Video may be compressed by GitHub
436 | |
437 |
438 |
439 |
440 | ---
441 |
442 | # Prompting Guideline
443 |
444 | Many people would ask how to write better prompts.
445 |
446 | Below is a ChatGPT template that I personally often use to get prompts:
447 |
448 | You are an assistant that writes short, motion-focused prompts for animating images.
449 |
450 | When the user sends an image, respond with a single, concise prompt describing visual motion (such as human activity, moving objects, or camera movements). Focus only on how the scene could come alive and become dynamic using brief phrases.
451 |
452 | Larger and more dynamic motions (like dancing, jumping, running, etc.) are preferred over smaller or more subtle ones (like standing still, sitting, etc.).
453 |
454 | Describe subject, then motion, then other things. For example: "The girl dances gracefully, with clear movements, full of charm."
455 |
456 | If there is something that can dance (like a man, girl, robot, etc.), then prefer to describe it as dancing.
457 |
458 | Stay in a loop: one image in, one motion prompt out. Do not explain, ask questions, or generate multiple options.
459 |
460 | You paste the instruct to ChatGPT and then feed it an image to get prompt like this:
461 |
462 | 
463 |
464 | *The man dances powerfully, striking sharp poses and gliding smoothly across the reflective floor.*
465 |
466 | Usually this will give you a prompt that works well.
467 |
468 | You can also write prompts yourself. Concise prompts are usually preferred, for example:
469 |
470 | *The girl dances gracefully, with clear movements, full of charm.*
471 |
472 | *The man dances powerfully, with clear movements, full of energy.*
473 |
474 | and so on.
475 |
476 | # Cite
477 |
478 | @article{zhang2025framepack,
479 | title={Packing Input Frame Contexts in Next-Frame Prediction Models for Video Generation},
480 | author={Lvmin Zhang and Maneesh Agrawala},
481 | journal={Arxiv},
482 | year={2025}
483 | }
484 |
--------------------------------------------------------------------------------
/demo_gradio.py:
--------------------------------------------------------------------------------
1 | from diffusers_helper.hf_login import login
2 |
3 | import os
4 |
5 | os.environ['HF_HOME'] = os.path.abspath(os.path.realpath(os.path.join(os.path.dirname(__file__), './hf_download')))
6 |
7 | import gradio as gr
8 | import torch
9 | import traceback
10 | import einops
11 | import safetensors.torch as sf
12 | import numpy as np
13 | import argparse
14 | import math
15 |
16 | from PIL import Image
17 | from diffusers import AutoencoderKLHunyuanVideo
18 | from transformers import LlamaModel, CLIPTextModel, LlamaTokenizerFast, CLIPTokenizer
19 | from diffusers_helper.hunyuan import encode_prompt_conds, vae_decode, vae_encode, vae_decode_fake
20 | from diffusers_helper.utils import save_bcthw_as_mp4, crop_or_pad_yield_mask, soft_append_bcthw, resize_and_center_crop, state_dict_weighted_merge, state_dict_offset_merge, generate_timestamp
21 | from diffusers_helper.models.hunyuan_video_packed import HunyuanVideoTransformer3DModelPacked
22 | from diffusers_helper.pipelines.k_diffusion_hunyuan import sample_hunyuan
23 | from diffusers_helper.memory import cpu, gpu, get_cuda_free_memory_gb, move_model_to_device_with_memory_preservation, offload_model_from_device_for_memory_preservation, fake_diffusers_current_device, DynamicSwapInstaller, unload_complete_models, load_model_as_complete
24 | from diffusers_helper.thread_utils import AsyncStream, async_run
25 | from diffusers_helper.gradio.progress_bar import make_progress_bar_css, make_progress_bar_html
26 | from transformers import SiglipImageProcessor, SiglipVisionModel
27 | from diffusers_helper.clip_vision import hf_clip_vision_encode
28 | from diffusers_helper.bucket_tools import find_nearest_bucket
29 |
30 |
31 | parser = argparse.ArgumentParser()
32 | parser.add_argument('--share', action='store_true')
33 | parser.add_argument("--server", type=str, default='0.0.0.0')
34 | parser.add_argument("--port", type=int, required=False)
35 | parser.add_argument("--inbrowser", action='store_true')
36 | args = parser.parse_args()
37 |
38 | # for win desktop probably use --server 127.0.0.1 --inbrowser
39 | # For linux server probably use --server 127.0.0.1 or do not use any cmd flags
40 |
41 | print(args)
42 |
43 | free_mem_gb = get_cuda_free_memory_gb(gpu)
44 | high_vram = free_mem_gb > 60
45 |
46 | print(f'Free VRAM {free_mem_gb} GB')
47 | print(f'High-VRAM Mode: {high_vram}')
48 |
49 | text_encoder = LlamaModel.from_pretrained("hunyuanvideo-community/HunyuanVideo", subfolder='text_encoder', torch_dtype=torch.float16).cpu()
50 | text_encoder_2 = CLIPTextModel.from_pretrained("hunyuanvideo-community/HunyuanVideo", subfolder='text_encoder_2', torch_dtype=torch.float16).cpu()
51 | tokenizer = LlamaTokenizerFast.from_pretrained("hunyuanvideo-community/HunyuanVideo", subfolder='tokenizer')
52 | tokenizer_2 = CLIPTokenizer.from_pretrained("hunyuanvideo-community/HunyuanVideo", subfolder='tokenizer_2')
53 | vae = AutoencoderKLHunyuanVideo.from_pretrained("hunyuanvideo-community/HunyuanVideo", subfolder='vae', torch_dtype=torch.float16).cpu()
54 |
55 | feature_extractor = SiglipImageProcessor.from_pretrained("lllyasviel/flux_redux_bfl", subfolder='feature_extractor')
56 | image_encoder = SiglipVisionModel.from_pretrained("lllyasviel/flux_redux_bfl", subfolder='image_encoder', torch_dtype=torch.float16).cpu()
57 |
58 | transformer = HunyuanVideoTransformer3DModelPacked.from_pretrained('lllyasviel/FramePackI2V_HY', torch_dtype=torch.bfloat16).cpu()
59 |
60 | vae.eval()
61 | text_encoder.eval()
62 | text_encoder_2.eval()
63 | image_encoder.eval()
64 | transformer.eval()
65 |
66 | if not high_vram:
67 | vae.enable_slicing()
68 | vae.enable_tiling()
69 |
70 | transformer.high_quality_fp32_output_for_inference = True
71 | print('transformer.high_quality_fp32_output_for_inference = True')
72 |
73 | transformer.to(dtype=torch.bfloat16)
74 | vae.to(dtype=torch.float16)
75 | image_encoder.to(dtype=torch.float16)
76 | text_encoder.to(dtype=torch.float16)
77 | text_encoder_2.to(dtype=torch.float16)
78 |
79 | vae.requires_grad_(False)
80 | text_encoder.requires_grad_(False)
81 | text_encoder_2.requires_grad_(False)
82 | image_encoder.requires_grad_(False)
83 | transformer.requires_grad_(False)
84 |
85 | if not high_vram:
86 | # DynamicSwapInstaller is same as huggingface's enable_sequential_offload but 3x faster
87 | DynamicSwapInstaller.install_model(transformer, device=gpu)
88 | DynamicSwapInstaller.install_model(text_encoder, device=gpu)
89 | else:
90 | text_encoder.to(gpu)
91 | text_encoder_2.to(gpu)
92 | image_encoder.to(gpu)
93 | vae.to(gpu)
94 | transformer.to(gpu)
95 |
96 | stream = AsyncStream()
97 |
98 | outputs_folder = './outputs/'
99 | os.makedirs(outputs_folder, exist_ok=True)
100 |
101 |
102 | @torch.no_grad()
103 | def worker(input_image, prompt, n_prompt, seed, total_second_length, latent_window_size, steps, cfg, gs, rs, gpu_memory_preservation, use_teacache, mp4_crf):
104 | total_latent_sections = (total_second_length * 30) / (latent_window_size * 4)
105 | total_latent_sections = int(max(round(total_latent_sections), 1))
106 |
107 | job_id = generate_timestamp()
108 |
109 | stream.output_queue.push(('progress', (None, '', make_progress_bar_html(0, 'Starting ...'))))
110 |
111 | try:
112 | # Clean GPU
113 | if not high_vram:
114 | unload_complete_models(
115 | text_encoder, text_encoder_2, image_encoder, vae, transformer
116 | )
117 |
118 | # Text encoding
119 |
120 | stream.output_queue.push(('progress', (None, '', make_progress_bar_html(0, 'Text encoding ...'))))
121 |
122 | if not high_vram:
123 | fake_diffusers_current_device(text_encoder, gpu) # since we only encode one text - that is one model move and one encode, offload is same time consumption since it is also one load and one encode.
124 | load_model_as_complete(text_encoder_2, target_device=gpu)
125 |
126 | llama_vec, clip_l_pooler = encode_prompt_conds(prompt, text_encoder, text_encoder_2, tokenizer, tokenizer_2)
127 |
128 | if cfg == 1:
129 | llama_vec_n, clip_l_pooler_n = torch.zeros_like(llama_vec), torch.zeros_like(clip_l_pooler)
130 | else:
131 | llama_vec_n, clip_l_pooler_n = encode_prompt_conds(n_prompt, text_encoder, text_encoder_2, tokenizer, tokenizer_2)
132 |
133 | llama_vec, llama_attention_mask = crop_or_pad_yield_mask(llama_vec, length=512)
134 | llama_vec_n, llama_attention_mask_n = crop_or_pad_yield_mask(llama_vec_n, length=512)
135 |
136 | # Processing input image
137 |
138 | stream.output_queue.push(('progress', (None, '', make_progress_bar_html(0, 'Image processing ...'))))
139 |
140 | H, W, C = input_image.shape
141 | height, width = find_nearest_bucket(H, W, resolution=640)
142 | input_image_np = resize_and_center_crop(input_image, target_width=width, target_height=height)
143 |
144 | Image.fromarray(input_image_np).save(os.path.join(outputs_folder, f'{job_id}.png'))
145 |
146 | input_image_pt = torch.from_numpy(input_image_np).float() / 127.5 - 1
147 | input_image_pt = input_image_pt.permute(2, 0, 1)[None, :, None]
148 |
149 | # VAE encoding
150 |
151 | stream.output_queue.push(('progress', (None, '', make_progress_bar_html(0, 'VAE encoding ...'))))
152 |
153 | if not high_vram:
154 | load_model_as_complete(vae, target_device=gpu)
155 |
156 | start_latent = vae_encode(input_image_pt, vae)
157 |
158 | # CLIP Vision
159 |
160 | stream.output_queue.push(('progress', (None, '', make_progress_bar_html(0, 'CLIP Vision encoding ...'))))
161 |
162 | if not high_vram:
163 | load_model_as_complete(image_encoder, target_device=gpu)
164 |
165 | image_encoder_output = hf_clip_vision_encode(input_image_np, feature_extractor, image_encoder)
166 | image_encoder_last_hidden_state = image_encoder_output.last_hidden_state
167 |
168 | # Dtype
169 |
170 | llama_vec = llama_vec.to(transformer.dtype)
171 | llama_vec_n = llama_vec_n.to(transformer.dtype)
172 | clip_l_pooler = clip_l_pooler.to(transformer.dtype)
173 | clip_l_pooler_n = clip_l_pooler_n.to(transformer.dtype)
174 | image_encoder_last_hidden_state = image_encoder_last_hidden_state.to(transformer.dtype)
175 |
176 | # Sampling
177 |
178 | stream.output_queue.push(('progress', (None, '', make_progress_bar_html(0, 'Start sampling ...'))))
179 |
180 | rnd = torch.Generator("cpu").manual_seed(seed)
181 | num_frames = latent_window_size * 4 - 3
182 |
183 | history_latents = torch.zeros(size=(1, 16, 1 + 2 + 16, height // 8, width // 8), dtype=torch.float32).cpu()
184 | history_pixels = None
185 | total_generated_latent_frames = 0
186 |
187 | latent_paddings = reversed(range(total_latent_sections))
188 |
189 | if total_latent_sections > 4:
190 | # In theory the latent_paddings should follow the above sequence, but it seems that duplicating some
191 | # items looks better than expanding it when total_latent_sections > 4
192 | # One can try to remove below trick and just
193 | # use `latent_paddings = list(reversed(range(total_latent_sections)))` to compare
194 | latent_paddings = [3] + [2] * (total_latent_sections - 3) + [1, 0]
195 |
196 | for latent_padding in latent_paddings:
197 | is_last_section = latent_padding == 0
198 | latent_padding_size = latent_padding * latent_window_size
199 |
200 | if stream.input_queue.top() == 'end':
201 | stream.output_queue.push(('end', None))
202 | return
203 |
204 | print(f'latent_padding_size = {latent_padding_size}, is_last_section = {is_last_section}')
205 |
206 | indices = torch.arange(0, sum([1, latent_padding_size, latent_window_size, 1, 2, 16])).unsqueeze(0)
207 | clean_latent_indices_pre, blank_indices, latent_indices, clean_latent_indices_post, clean_latent_2x_indices, clean_latent_4x_indices = indices.split([1, latent_padding_size, latent_window_size, 1, 2, 16], dim=1)
208 | clean_latent_indices = torch.cat([clean_latent_indices_pre, clean_latent_indices_post], dim=1)
209 |
210 | clean_latents_pre = start_latent.to(history_latents)
211 | clean_latents_post, clean_latents_2x, clean_latents_4x = history_latents[:, :, :1 + 2 + 16, :, :].split([1, 2, 16], dim=2)
212 | clean_latents = torch.cat([clean_latents_pre, clean_latents_post], dim=2)
213 |
214 | if not high_vram:
215 | unload_complete_models()
216 | move_model_to_device_with_memory_preservation(transformer, target_device=gpu, preserved_memory_gb=gpu_memory_preservation)
217 |
218 | if use_teacache:
219 | transformer.initialize_teacache(enable_teacache=True, num_steps=steps)
220 | else:
221 | transformer.initialize_teacache(enable_teacache=False)
222 |
223 | def callback(d):
224 | preview = d['denoised']
225 | preview = vae_decode_fake(preview)
226 |
227 | preview = (preview * 255.0).detach().cpu().numpy().clip(0, 255).astype(np.uint8)
228 | preview = einops.rearrange(preview, 'b c t h w -> (b h) (t w) c')
229 |
230 | if stream.input_queue.top() == 'end':
231 | stream.output_queue.push(('end', None))
232 | raise KeyboardInterrupt('User ends the task.')
233 |
234 | current_step = d['i'] + 1
235 | percentage = int(100.0 * current_step / steps)
236 | hint = f'Sampling {current_step}/{steps}'
237 | desc = f'Total generated frames: {int(max(0, total_generated_latent_frames * 4 - 3))}, Video length: {max(0, (total_generated_latent_frames * 4 - 3) / 30) :.2f} seconds (FPS-30). The video is being extended now ...'
238 | stream.output_queue.push(('progress', (preview, desc, make_progress_bar_html(percentage, hint))))
239 | return
240 |
241 | generated_latents = sample_hunyuan(
242 | transformer=transformer,
243 | sampler='unipc',
244 | width=width,
245 | height=height,
246 | frames=num_frames,
247 | real_guidance_scale=cfg,
248 | distilled_guidance_scale=gs,
249 | guidance_rescale=rs,
250 | # shift=3.0,
251 | num_inference_steps=steps,
252 | generator=rnd,
253 | prompt_embeds=llama_vec,
254 | prompt_embeds_mask=llama_attention_mask,
255 | prompt_poolers=clip_l_pooler,
256 | negative_prompt_embeds=llama_vec_n,
257 | negative_prompt_embeds_mask=llama_attention_mask_n,
258 | negative_prompt_poolers=clip_l_pooler_n,
259 | device=gpu,
260 | dtype=torch.bfloat16,
261 | image_embeddings=image_encoder_last_hidden_state,
262 | latent_indices=latent_indices,
263 | clean_latents=clean_latents,
264 | clean_latent_indices=clean_latent_indices,
265 | clean_latents_2x=clean_latents_2x,
266 | clean_latent_2x_indices=clean_latent_2x_indices,
267 | clean_latents_4x=clean_latents_4x,
268 | clean_latent_4x_indices=clean_latent_4x_indices,
269 | callback=callback,
270 | )
271 |
272 | if is_last_section:
273 | generated_latents = torch.cat([start_latent.to(generated_latents), generated_latents], dim=2)
274 |
275 | total_generated_latent_frames += int(generated_latents.shape[2])
276 | history_latents = torch.cat([generated_latents.to(history_latents), history_latents], dim=2)
277 |
278 | if not high_vram:
279 | offload_model_from_device_for_memory_preservation(transformer, target_device=gpu, preserved_memory_gb=8)
280 | load_model_as_complete(vae, target_device=gpu)
281 |
282 | real_history_latents = history_latents[:, :, :total_generated_latent_frames, :, :]
283 |
284 | if history_pixels is None:
285 | history_pixels = vae_decode(real_history_latents, vae).cpu()
286 | else:
287 | section_latent_frames = (latent_window_size * 2 + 1) if is_last_section else (latent_window_size * 2)
288 | overlapped_frames = latent_window_size * 4 - 3
289 |
290 | current_pixels = vae_decode(real_history_latents[:, :, :section_latent_frames], vae).cpu()
291 | history_pixels = soft_append_bcthw(current_pixels, history_pixels, overlapped_frames)
292 |
293 | if not high_vram:
294 | unload_complete_models()
295 |
296 | output_filename = os.path.join(outputs_folder, f'{job_id}_{total_generated_latent_frames}.mp4')
297 |
298 | save_bcthw_as_mp4(history_pixels, output_filename, fps=30, crf=mp4_crf)
299 |
300 | print(f'Decoded. Current latent shape {real_history_latents.shape}; pixel shape {history_pixels.shape}')
301 |
302 | stream.output_queue.push(('file', output_filename))
303 |
304 | if is_last_section:
305 | break
306 | except:
307 | traceback.print_exc()
308 |
309 | if not high_vram:
310 | unload_complete_models(
311 | text_encoder, text_encoder_2, image_encoder, vae, transformer
312 | )
313 |
314 | stream.output_queue.push(('end', None))
315 | return
316 |
317 |
318 | def process(input_image, prompt, n_prompt, seed, total_second_length, latent_window_size, steps, cfg, gs, rs, gpu_memory_preservation, use_teacache, mp4_crf):
319 | global stream
320 | assert input_image is not None, 'No input image!'
321 |
322 | yield None, None, '', '', gr.update(interactive=False), gr.update(interactive=True)
323 |
324 | stream = AsyncStream()
325 |
326 | async_run(worker, input_image, prompt, n_prompt, seed, total_second_length, latent_window_size, steps, cfg, gs, rs, gpu_memory_preservation, use_teacache, mp4_crf)
327 |
328 | output_filename = None
329 |
330 | while True:
331 | flag, data = stream.output_queue.next()
332 |
333 | if flag == 'file':
334 | output_filename = data
335 | yield output_filename, gr.update(), gr.update(), gr.update(), gr.update(interactive=False), gr.update(interactive=True)
336 |
337 | if flag == 'progress':
338 | preview, desc, html = data
339 | yield gr.update(), gr.update(visible=True, value=preview), desc, html, gr.update(interactive=False), gr.update(interactive=True)
340 |
341 | if flag == 'end':
342 | yield output_filename, gr.update(visible=False), gr.update(), '', gr.update(interactive=True), gr.update(interactive=False)
343 | break
344 |
345 |
346 | def end_process():
347 | stream.input_queue.push('end')
348 |
349 |
350 | quick_prompts = [
351 | 'The girl dances gracefully, with clear movements, full of charm.',
352 | 'A character doing some simple body movements.',
353 | ]
354 | quick_prompts = [[x] for x in quick_prompts]
355 |
356 |
357 | css = make_progress_bar_css()
358 | block = gr.Blocks(css=css).queue()
359 | with block:
360 | gr.Markdown('# FramePack')
361 | with gr.Row():
362 | with gr.Column():
363 | input_image = gr.Image(sources='upload', type="numpy", label="Image", height=320)
364 | prompt = gr.Textbox(label="Prompt", value='')
365 | example_quick_prompts = gr.Dataset(samples=quick_prompts, label='Quick List', samples_per_page=1000, components=[prompt])
366 | example_quick_prompts.click(lambda x: x[0], inputs=[example_quick_prompts], outputs=prompt, show_progress=False, queue=False)
367 |
368 | with gr.Row():
369 | start_button = gr.Button(value="Start Generation")
370 | end_button = gr.Button(value="End Generation", interactive=False)
371 |
372 | with gr.Group():
373 | use_teacache = gr.Checkbox(label='Use TeaCache', value=True, info='Faster speed, but often makes hands and fingers slightly worse.')
374 |
375 | n_prompt = gr.Textbox(label="Negative Prompt", value="", visible=False) # Not used
376 | seed = gr.Number(label="Seed", value=31337, precision=0)
377 |
378 | total_second_length = gr.Slider(label="Total Video Length (Seconds)", minimum=1, maximum=120, value=5, step=0.1)
379 | latent_window_size = gr.Slider(label="Latent Window Size", minimum=1, maximum=33, value=9, step=1, visible=False) # Should not change
380 | steps = gr.Slider(label="Steps", minimum=1, maximum=100, value=25, step=1, info='Changing this value is not recommended.')
381 |
382 | cfg = gr.Slider(label="CFG Scale", minimum=1.0, maximum=32.0, value=1.0, step=0.01, visible=False) # Should not change
383 | gs = gr.Slider(label="Distilled CFG Scale", minimum=1.0, maximum=32.0, value=10.0, step=0.01, info='Changing this value is not recommended.')
384 | rs = gr.Slider(label="CFG Re-Scale", minimum=0.0, maximum=1.0, value=0.0, step=0.01, visible=False) # Should not change
385 |
386 | gpu_memory_preservation = gr.Slider(label="GPU Inference Preserved Memory (GB) (larger means slower)", minimum=6, maximum=128, value=6, step=0.1, info="Set this number to a larger value if you encounter OOM. Larger value causes slower speed.")
387 |
388 | mp4_crf = gr.Slider(label="MP4 Compression", minimum=0, maximum=100, value=16, step=1, info="Lower means better quality. 0 is uncompressed. Change to 16 if you get black outputs. ")
389 |
390 | with gr.Column():
391 | preview_image = gr.Image(label="Next Latents", height=200, visible=False)
392 | result_video = gr.Video(label="Finished Frames", autoplay=True, show_share_button=False, height=512, loop=True)
393 | gr.Markdown('Note that the ending actions will be generated before the starting actions due to the inverted sampling. If the starting action is not in the video, you just need to wait, and it will be generated later.')
394 | progress_desc = gr.Markdown('', elem_classes='no-generating-animation')
395 | progress_bar = gr.HTML('', elem_classes='no-generating-animation')
396 |
397 | gr.HTML('')
398 |
399 | ips = [input_image, prompt, n_prompt, seed, total_second_length, latent_window_size, steps, cfg, gs, rs, gpu_memory_preservation, use_teacache, mp4_crf]
400 | start_button.click(fn=process, inputs=ips, outputs=[result_video, preview_image, progress_desc, progress_bar, start_button, end_button])
401 | end_button.click(fn=end_process)
402 |
403 |
404 | block.launch(
405 | server_name=args.server,
406 | server_port=args.port,
407 | share=args.share,
408 | inbrowser=args.inbrowser,
409 | )
410 |
--------------------------------------------------------------------------------
/diffusers_helper/bucket_tools.py:
--------------------------------------------------------------------------------
1 | bucket_options = {
2 | 640: [
3 | (416, 960),
4 | (448, 864),
5 | (480, 832),
6 | (512, 768),
7 | (544, 704),
8 | (576, 672),
9 | (608, 640),
10 | (640, 608),
11 | (672, 576),
12 | (704, 544),
13 | (768, 512),
14 | (832, 480),
15 | (864, 448),
16 | (960, 416),
17 | ],
18 | }
19 |
20 |
21 | def find_nearest_bucket(h, w, resolution=640):
22 | min_metric = float('inf')
23 | best_bucket = None
24 | for (bucket_h, bucket_w) in bucket_options[resolution]:
25 | metric = abs(h * bucket_w - w * bucket_h)
26 | if metric <= min_metric:
27 | min_metric = metric
28 | best_bucket = (bucket_h, bucket_w)
29 | return best_bucket
30 |
31 |
--------------------------------------------------------------------------------
/diffusers_helper/clip_vision.py:
--------------------------------------------------------------------------------
1 | import numpy as np
2 |
3 |
4 | def hf_clip_vision_encode(image, feature_extractor, image_encoder):
5 | assert isinstance(image, np.ndarray)
6 | assert image.ndim == 3 and image.shape[2] == 3
7 | assert image.dtype == np.uint8
8 |
9 | preprocessed = feature_extractor.preprocess(images=image, return_tensors="pt").to(device=image_encoder.device, dtype=image_encoder.dtype)
10 | image_encoder_output = image_encoder(**preprocessed)
11 |
12 | return image_encoder_output
13 |
--------------------------------------------------------------------------------
/diffusers_helper/dit_common.py:
--------------------------------------------------------------------------------
1 | import torch
2 | import accelerate.accelerator
3 |
4 | from diffusers.models.normalization import RMSNorm, LayerNorm, FP32LayerNorm, AdaLayerNormContinuous
5 |
6 |
7 | accelerate.accelerator.convert_outputs_to_fp32 = lambda x: x
8 |
9 |
10 | def LayerNorm_forward(self, x):
11 | return torch.nn.functional.layer_norm(x, self.normalized_shape, self.weight, self.bias, self.eps).to(x)
12 |
13 |
14 | LayerNorm.forward = LayerNorm_forward
15 | torch.nn.LayerNorm.forward = LayerNorm_forward
16 |
17 |
18 | def FP32LayerNorm_forward(self, x):
19 | origin_dtype = x.dtype
20 | return torch.nn.functional.layer_norm(
21 | x.float(),
22 | self.normalized_shape,
23 | self.weight.float() if self.weight is not None else None,
24 | self.bias.float() if self.bias is not None else None,
25 | self.eps,
26 | ).to(origin_dtype)
27 |
28 |
29 | FP32LayerNorm.forward = FP32LayerNorm_forward
30 |
31 |
32 | def RMSNorm_forward(self, hidden_states):
33 | input_dtype = hidden_states.dtype
34 | variance = hidden_states.to(torch.float32).pow(2).mean(-1, keepdim=True)
35 | hidden_states = hidden_states * torch.rsqrt(variance + self.eps)
36 |
37 | if self.weight is None:
38 | return hidden_states.to(input_dtype)
39 |
40 | return hidden_states.to(input_dtype) * self.weight.to(input_dtype)
41 |
42 |
43 | RMSNorm.forward = RMSNorm_forward
44 |
45 |
46 | def AdaLayerNormContinuous_forward(self, x, conditioning_embedding):
47 | emb = self.linear(self.silu(conditioning_embedding))
48 | scale, shift = emb.chunk(2, dim=1)
49 | x = self.norm(x) * (1 + scale)[:, None, :] + shift[:, None, :]
50 | return x
51 |
52 |
53 | AdaLayerNormContinuous.forward = AdaLayerNormContinuous_forward
54 |
--------------------------------------------------------------------------------
/diffusers_helper/gradio/progress_bar.py:
--------------------------------------------------------------------------------
1 | progress_html = '''
2 |
9 | '''
10 |
11 | css = '''
12 | .loader-container {
13 | display: flex; /* Use flex to align items horizontally */
14 | align-items: center; /* Center items vertically within the container */
15 | white-space: nowrap; /* Prevent line breaks within the container */
16 | }
17 |
18 | .loader {
19 | border: 8px solid #f3f3f3; /* Light grey */
20 | border-top: 8px solid #3498db; /* Blue */
21 | border-radius: 50%;
22 | width: 30px;
23 | height: 30px;
24 | animation: spin 2s linear infinite;
25 | }
26 |
27 | @keyframes spin {
28 | 0% { transform: rotate(0deg); }
29 | 100% { transform: rotate(360deg); }
30 | }
31 |
32 | /* Style the progress bar */
33 | progress {
34 | appearance: none; /* Remove default styling */
35 | height: 20px; /* Set the height of the progress bar */
36 | border-radius: 5px; /* Round the corners of the progress bar */
37 | background-color: #f3f3f3; /* Light grey background */
38 | width: 100%;
39 | vertical-align: middle !important;
40 | }
41 |
42 | /* Style the progress bar container */
43 | .progress-container {
44 | margin-left: 20px;
45 | margin-right: 20px;
46 | flex-grow: 1; /* Allow the progress container to take up remaining space */
47 | }
48 |
49 | /* Set the color of the progress bar fill */
50 | progress::-webkit-progress-value {
51 | background-color: #3498db; /* Blue color for the fill */
52 | }
53 |
54 | progress::-moz-progress-bar {
55 | background-color: #3498db; /* Blue color for the fill in Firefox */
56 | }
57 |
58 | /* Style the text on the progress bar */
59 | progress::after {
60 | content: attr(value '%'); /* Display the progress value followed by '%' */
61 | position: absolute;
62 | top: 50%;
63 | left: 50%;
64 | transform: translate(-50%, -50%);
65 | color: white; /* Set text color */
66 | font-size: 14px; /* Set font size */
67 | }
68 |
69 | /* Style other texts */
70 | .loader-container > span {
71 | margin-left: 5px; /* Add spacing between the progress bar and the text */
72 | }
73 |
74 | .no-generating-animation > .generating {
75 | display: none !important;
76 | }
77 |
78 | '''
79 |
80 |
81 | def make_progress_bar_html(number, text):
82 | return progress_html.replace('*number*', str(number)).replace('*text*', text)
83 |
84 |
85 | def make_progress_bar_css():
86 | return css
87 |
--------------------------------------------------------------------------------
/diffusers_helper/hf_login.py:
--------------------------------------------------------------------------------
1 | import os
2 |
3 |
4 | def login(token):
5 | from huggingface_hub import login
6 | import time
7 |
8 | while True:
9 | try:
10 | login(token)
11 | print('HF login ok.')
12 | break
13 | except Exception as e:
14 | print(f'HF login failed: {e}. Retrying')
15 | time.sleep(0.5)
16 |
17 |
18 | hf_token = os.environ.get('HF_TOKEN', None)
19 |
20 | if hf_token is not None:
21 | login(hf_token)
22 |
--------------------------------------------------------------------------------
/diffusers_helper/hunyuan.py:
--------------------------------------------------------------------------------
1 | import torch
2 |
3 | from diffusers.pipelines.hunyuan_video.pipeline_hunyuan_video import DEFAULT_PROMPT_TEMPLATE
4 | from diffusers_helper.utils import crop_or_pad_yield_mask
5 |
6 |
7 | @torch.no_grad()
8 | def encode_prompt_conds(prompt, text_encoder, text_encoder_2, tokenizer, tokenizer_2, max_length=256):
9 | assert isinstance(prompt, str)
10 |
11 | prompt = [prompt]
12 |
13 | # LLAMA
14 |
15 | prompt_llama = [DEFAULT_PROMPT_TEMPLATE["template"].format(p) for p in prompt]
16 | crop_start = DEFAULT_PROMPT_TEMPLATE["crop_start"]
17 |
18 | llama_inputs = tokenizer(
19 | prompt_llama,
20 | padding="max_length",
21 | max_length=max_length + crop_start,
22 | truncation=True,
23 | return_tensors="pt",
24 | return_length=False,
25 | return_overflowing_tokens=False,
26 | return_attention_mask=True,
27 | )
28 |
29 | llama_input_ids = llama_inputs.input_ids.to(text_encoder.device)
30 | llama_attention_mask = llama_inputs.attention_mask.to(text_encoder.device)
31 | llama_attention_length = int(llama_attention_mask.sum())
32 |
33 | llama_outputs = text_encoder(
34 | input_ids=llama_input_ids,
35 | attention_mask=llama_attention_mask,
36 | output_hidden_states=True,
37 | )
38 |
39 | llama_vec = llama_outputs.hidden_states[-3][:, crop_start:llama_attention_length]
40 | # llama_vec_remaining = llama_outputs.hidden_states[-3][:, llama_attention_length:]
41 | llama_attention_mask = llama_attention_mask[:, crop_start:llama_attention_length]
42 |
43 | assert torch.all(llama_attention_mask.bool())
44 |
45 | # CLIP
46 |
47 | clip_l_input_ids = tokenizer_2(
48 | prompt,
49 | padding="max_length",
50 | max_length=77,
51 | truncation=True,
52 | return_overflowing_tokens=False,
53 | return_length=False,
54 | return_tensors="pt",
55 | ).input_ids
56 | clip_l_pooler = text_encoder_2(clip_l_input_ids.to(text_encoder_2.device), output_hidden_states=False).pooler_output
57 |
58 | return llama_vec, clip_l_pooler
59 |
60 |
61 | @torch.no_grad()
62 | def vae_decode_fake(latents):
63 | latent_rgb_factors = [
64 | [-0.0395, -0.0331, 0.0445],
65 | [0.0696, 0.0795, 0.0518],
66 | [0.0135, -0.0945, -0.0282],
67 | [0.0108, -0.0250, -0.0765],
68 | [-0.0209, 0.0032, 0.0224],
69 | [-0.0804, -0.0254, -0.0639],
70 | [-0.0991, 0.0271, -0.0669],
71 | [-0.0646, -0.0422, -0.0400],
72 | [-0.0696, -0.0595, -0.0894],
73 | [-0.0799, -0.0208, -0.0375],
74 | [0.1166, 0.1627, 0.0962],
75 | [0.1165, 0.0432, 0.0407],
76 | [-0.2315, -0.1920, -0.1355],
77 | [-0.0270, 0.0401, -0.0821],
78 | [-0.0616, -0.0997, -0.0727],
79 | [0.0249, -0.0469, -0.1703]
80 | ] # From comfyui
81 |
82 | latent_rgb_factors_bias = [0.0259, -0.0192, -0.0761]
83 |
84 | weight = torch.tensor(latent_rgb_factors, device=latents.device, dtype=latents.dtype).transpose(0, 1)[:, :, None, None, None]
85 | bias = torch.tensor(latent_rgb_factors_bias, device=latents.device, dtype=latents.dtype)
86 |
87 | images = torch.nn.functional.conv3d(latents, weight, bias=bias, stride=1, padding=0, dilation=1, groups=1)
88 | images = images.clamp(0.0, 1.0)
89 |
90 | return images
91 |
92 |
93 | @torch.no_grad()
94 | def vae_decode(latents, vae, image_mode=False):
95 | latents = latents / vae.config.scaling_factor
96 |
97 | if not image_mode:
98 | image = vae.decode(latents.to(device=vae.device, dtype=vae.dtype)).sample
99 | else:
100 | latents = latents.to(device=vae.device, dtype=vae.dtype).unbind(2)
101 | image = [vae.decode(l.unsqueeze(2)).sample for l in latents]
102 | image = torch.cat(image, dim=2)
103 |
104 | return image
105 |
106 |
107 | @torch.no_grad()
108 | def vae_encode(image, vae):
109 | latents = vae.encode(image.to(device=vae.device, dtype=vae.dtype)).latent_dist.sample()
110 | latents = latents * vae.config.scaling_factor
111 | return latents
112 |
--------------------------------------------------------------------------------
/diffusers_helper/k_diffusion/uni_pc_fm.py:
--------------------------------------------------------------------------------
1 | # Better Flow Matching UniPC by Lvmin Zhang
2 | # (c) 2025
3 | # CC BY-SA 4.0
4 | # Attribution-ShareAlike 4.0 International Licence
5 |
6 |
7 | import torch
8 |
9 | from tqdm.auto import trange
10 |
11 |
12 | def expand_dims(v, dims):
13 | return v[(...,) + (None,) * (dims - 1)]
14 |
15 |
16 | class FlowMatchUniPC:
17 | def __init__(self, model, extra_args, variant='bh1'):
18 | self.model = model
19 | self.variant = variant
20 | self.extra_args = extra_args
21 |
22 | def model_fn(self, x, t):
23 | return self.model(x, t, **self.extra_args)
24 |
25 | def update_fn(self, x, model_prev_list, t_prev_list, t, order):
26 | assert order <= len(model_prev_list)
27 | dims = x.dim()
28 |
29 | t_prev_0 = t_prev_list[-1]
30 | lambda_prev_0 = - torch.log(t_prev_0)
31 | lambda_t = - torch.log(t)
32 | model_prev_0 = model_prev_list[-1]
33 |
34 | h = lambda_t - lambda_prev_0
35 |
36 | rks = []
37 | D1s = []
38 | for i in range(1, order):
39 | t_prev_i = t_prev_list[-(i + 1)]
40 | model_prev_i = model_prev_list[-(i + 1)]
41 | lambda_prev_i = - torch.log(t_prev_i)
42 | rk = ((lambda_prev_i - lambda_prev_0) / h)[0]
43 | rks.append(rk)
44 | D1s.append((model_prev_i - model_prev_0) / rk)
45 |
46 | rks.append(1.)
47 | rks = torch.tensor(rks, device=x.device)
48 |
49 | R = []
50 | b = []
51 |
52 | hh = -h[0]
53 | h_phi_1 = torch.expm1(hh)
54 | h_phi_k = h_phi_1 / hh - 1
55 |
56 | factorial_i = 1
57 |
58 | if self.variant == 'bh1':
59 | B_h = hh
60 | elif self.variant == 'bh2':
61 | B_h = torch.expm1(hh)
62 | else:
63 | raise NotImplementedError('Bad variant!')
64 |
65 | for i in range(1, order + 1):
66 | R.append(torch.pow(rks, i - 1))
67 | b.append(h_phi_k * factorial_i / B_h)
68 | factorial_i *= (i + 1)
69 | h_phi_k = h_phi_k / hh - 1 / factorial_i
70 |
71 | R = torch.stack(R)
72 | b = torch.tensor(b, device=x.device)
73 |
74 | use_predictor = len(D1s) > 0
75 |
76 | if use_predictor:
77 | D1s = torch.stack(D1s, dim=1)
78 | if order == 2:
79 | rhos_p = torch.tensor([0.5], device=b.device)
80 | else:
81 | rhos_p = torch.linalg.solve(R[:-1, :-1], b[:-1])
82 | else:
83 | D1s = None
84 | rhos_p = None
85 |
86 | if order == 1:
87 | rhos_c = torch.tensor([0.5], device=b.device)
88 | else:
89 | rhos_c = torch.linalg.solve(R, b)
90 |
91 | x_t_ = expand_dims(t / t_prev_0, dims) * x - expand_dims(h_phi_1, dims) * model_prev_0
92 |
93 | if use_predictor:
94 | pred_res = torch.tensordot(D1s, rhos_p, dims=([1], [0]))
95 | else:
96 | pred_res = 0
97 |
98 | x_t = x_t_ - expand_dims(B_h, dims) * pred_res
99 | model_t = self.model_fn(x_t, t)
100 |
101 | if D1s is not None:
102 | corr_res = torch.tensordot(D1s, rhos_c[:-1], dims=([1], [0]))
103 | else:
104 | corr_res = 0
105 |
106 | D1_t = (model_t - model_prev_0)
107 | x_t = x_t_ - expand_dims(B_h, dims) * (corr_res + rhos_c[-1] * D1_t)
108 |
109 | return x_t, model_t
110 |
111 | def sample(self, x, sigmas, callback=None, disable_pbar=False):
112 | order = min(3, len(sigmas) - 2)
113 | model_prev_list, t_prev_list = [], []
114 | for i in trange(len(sigmas) - 1, disable=disable_pbar):
115 | vec_t = sigmas[i].expand(x.shape[0])
116 |
117 | if i == 0:
118 | model_prev_list = [self.model_fn(x, vec_t)]
119 | t_prev_list = [vec_t]
120 | elif i < order:
121 | init_order = i
122 | x, model_x = self.update_fn(x, model_prev_list, t_prev_list, vec_t, init_order)
123 | model_prev_list.append(model_x)
124 | t_prev_list.append(vec_t)
125 | else:
126 | x, model_x = self.update_fn(x, model_prev_list, t_prev_list, vec_t, order)
127 | model_prev_list.append(model_x)
128 | t_prev_list.append(vec_t)
129 |
130 | model_prev_list = model_prev_list[-order:]
131 | t_prev_list = t_prev_list[-order:]
132 |
133 | if callback is not None:
134 | callback({'x': x, 'i': i, 'denoised': model_prev_list[-1]})
135 |
136 | return model_prev_list[-1]
137 |
138 |
139 | def sample_unipc(model, noise, sigmas, extra_args=None, callback=None, disable=False, variant='bh1'):
140 | assert variant in ['bh1', 'bh2']
141 | return FlowMatchUniPC(model, extra_args=extra_args, variant=variant).sample(noise, sigmas=sigmas, callback=callback, disable_pbar=disable)
142 |
--------------------------------------------------------------------------------
/diffusers_helper/k_diffusion/wrapper.py:
--------------------------------------------------------------------------------
1 | import torch
2 |
3 |
4 | def append_dims(x, target_dims):
5 | return x[(...,) + (None,) * (target_dims - x.ndim)]
6 |
7 |
8 | def rescale_noise_cfg(noise_cfg, noise_pred_text, guidance_rescale=1.0):
9 | if guidance_rescale == 0:
10 | return noise_cfg
11 |
12 | std_text = noise_pred_text.std(dim=list(range(1, noise_pred_text.ndim)), keepdim=True)
13 | std_cfg = noise_cfg.std(dim=list(range(1, noise_cfg.ndim)), keepdim=True)
14 | noise_pred_rescaled = noise_cfg * (std_text / std_cfg)
15 | noise_cfg = guidance_rescale * noise_pred_rescaled + (1.0 - guidance_rescale) * noise_cfg
16 | return noise_cfg
17 |
18 |
19 | def fm_wrapper(transformer, t_scale=1000.0):
20 | def k_model(x, sigma, **extra_args):
21 | dtype = extra_args['dtype']
22 | cfg_scale = extra_args['cfg_scale']
23 | cfg_rescale = extra_args['cfg_rescale']
24 | concat_latent = extra_args['concat_latent']
25 |
26 | original_dtype = x.dtype
27 | sigma = sigma.float()
28 |
29 | x = x.to(dtype)
30 | timestep = (sigma * t_scale).to(dtype)
31 |
32 | if concat_latent is None:
33 | hidden_states = x
34 | else:
35 | hidden_states = torch.cat([x, concat_latent.to(x)], dim=1)
36 |
37 | pred_positive = transformer(hidden_states=hidden_states, timestep=timestep, return_dict=False, **extra_args['positive'])[0].float()
38 |
39 | if cfg_scale == 1.0:
40 | pred_negative = torch.zeros_like(pred_positive)
41 | else:
42 | pred_negative = transformer(hidden_states=hidden_states, timestep=timestep, return_dict=False, **extra_args['negative'])[0].float()
43 |
44 | pred_cfg = pred_negative + cfg_scale * (pred_positive - pred_negative)
45 | pred = rescale_noise_cfg(pred_cfg, pred_positive, guidance_rescale=cfg_rescale)
46 |
47 | x0 = x.float() - pred.float() * append_dims(sigma, x.ndim)
48 |
49 | return x0.to(dtype=original_dtype)
50 |
51 | return k_model
52 |
--------------------------------------------------------------------------------
/diffusers_helper/memory.py:
--------------------------------------------------------------------------------
1 | # By lllyasviel
2 |
3 |
4 | import torch
5 |
6 |
7 | cpu = torch.device('cpu')
8 | gpu = torch.device(f'cuda:{torch.cuda.current_device()}')
9 | gpu_complete_modules = []
10 |
11 |
12 | class DynamicSwapInstaller:
13 | @staticmethod
14 | def _install_module(module: torch.nn.Module, **kwargs):
15 | original_class = module.__class__
16 | module.__dict__['forge_backup_original_class'] = original_class
17 |
18 | def hacked_get_attr(self, name: str):
19 | if '_parameters' in self.__dict__:
20 | _parameters = self.__dict__['_parameters']
21 | if name in _parameters:
22 | p = _parameters[name]
23 | if p is None:
24 | return None
25 | if p.__class__ == torch.nn.Parameter:
26 | return torch.nn.Parameter(p.to(**kwargs), requires_grad=p.requires_grad)
27 | else:
28 | return p.to(**kwargs)
29 | if '_buffers' in self.__dict__:
30 | _buffers = self.__dict__['_buffers']
31 | if name in _buffers:
32 | return _buffers[name].to(**kwargs)
33 | return super(original_class, self).__getattr__(name)
34 |
35 | module.__class__ = type('DynamicSwap_' + original_class.__name__, (original_class,), {
36 | '__getattr__': hacked_get_attr,
37 | })
38 |
39 | return
40 |
41 | @staticmethod
42 | def _uninstall_module(module: torch.nn.Module):
43 | if 'forge_backup_original_class' in module.__dict__:
44 | module.__class__ = module.__dict__.pop('forge_backup_original_class')
45 | return
46 |
47 | @staticmethod
48 | def install_model(model: torch.nn.Module, **kwargs):
49 | for m in model.modules():
50 | DynamicSwapInstaller._install_module(m, **kwargs)
51 | return
52 |
53 | @staticmethod
54 | def uninstall_model(model: torch.nn.Module):
55 | for m in model.modules():
56 | DynamicSwapInstaller._uninstall_module(m)
57 | return
58 |
59 |
60 | def fake_diffusers_current_device(model: torch.nn.Module, target_device: torch.device):
61 | if hasattr(model, 'scale_shift_table'):
62 | model.scale_shift_table.data = model.scale_shift_table.data.to(target_device)
63 | return
64 |
65 | for k, p in model.named_modules():
66 | if hasattr(p, 'weight'):
67 | p.to(target_device)
68 | return
69 |
70 |
71 | def get_cuda_free_memory_gb(device=None):
72 | if device is None:
73 | device = gpu
74 |
75 | memory_stats = torch.cuda.memory_stats(device)
76 | bytes_active = memory_stats['active_bytes.all.current']
77 | bytes_reserved = memory_stats['reserved_bytes.all.current']
78 | bytes_free_cuda, _ = torch.cuda.mem_get_info(device)
79 | bytes_inactive_reserved = bytes_reserved - bytes_active
80 | bytes_total_available = bytes_free_cuda + bytes_inactive_reserved
81 | return bytes_total_available / (1024 ** 3)
82 |
83 |
84 | def move_model_to_device_with_memory_preservation(model, target_device, preserved_memory_gb=0):
85 | print(f'Moving {model.__class__.__name__} to {target_device} with preserved memory: {preserved_memory_gb} GB')
86 |
87 | for m in model.modules():
88 | if get_cuda_free_memory_gb(target_device) <= preserved_memory_gb:
89 | torch.cuda.empty_cache()
90 | return
91 |
92 | if hasattr(m, 'weight'):
93 | m.to(device=target_device)
94 |
95 | model.to(device=target_device)
96 | torch.cuda.empty_cache()
97 | return
98 |
99 |
100 | def offload_model_from_device_for_memory_preservation(model, target_device, preserved_memory_gb=0):
101 | print(f'Offloading {model.__class__.__name__} from {target_device} to preserve memory: {preserved_memory_gb} GB')
102 |
103 | for m in model.modules():
104 | if get_cuda_free_memory_gb(target_device) >= preserved_memory_gb:
105 | torch.cuda.empty_cache()
106 | return
107 |
108 | if hasattr(m, 'weight'):
109 | m.to(device=cpu)
110 |
111 | model.to(device=cpu)
112 | torch.cuda.empty_cache()
113 | return
114 |
115 |
116 | def unload_complete_models(*args):
117 | for m in gpu_complete_modules + list(args):
118 | m.to(device=cpu)
119 | print(f'Unloaded {m.__class__.__name__} as complete.')
120 |
121 | gpu_complete_modules.clear()
122 | torch.cuda.empty_cache()
123 | return
124 |
125 |
126 | def load_model_as_complete(model, target_device, unload=True):
127 | if unload:
128 | unload_complete_models()
129 |
130 | model.to(device=target_device)
131 | print(f'Loaded {model.__class__.__name__} to {target_device} as complete.')
132 |
133 | gpu_complete_modules.append(model)
134 | return
135 |
--------------------------------------------------------------------------------
/diffusers_helper/models/hunyuan_video_packed.py:
--------------------------------------------------------------------------------
1 | from typing import Any, Dict, List, Optional, Tuple, Union
2 |
3 | import torch
4 | import einops
5 | import torch.nn as nn
6 | import numpy as np
7 |
8 | from diffusers.loaders import FromOriginalModelMixin
9 | from diffusers.configuration_utils import ConfigMixin, register_to_config
10 | from diffusers.loaders import PeftAdapterMixin
11 | from diffusers.utils import logging
12 | from diffusers.models.attention import FeedForward
13 | from diffusers.models.attention_processor import Attention
14 | from diffusers.models.embeddings import TimestepEmbedding, Timesteps, PixArtAlphaTextProjection
15 | from diffusers.models.modeling_outputs import Transformer2DModelOutput
16 | from diffusers.models.modeling_utils import ModelMixin
17 | from diffusers_helper.dit_common import LayerNorm
18 | from diffusers_helper.utils import zero_module
19 |
20 |
21 | enabled_backends = []
22 |
23 | if torch.backends.cuda.flash_sdp_enabled():
24 | enabled_backends.append("flash")
25 | if torch.backends.cuda.math_sdp_enabled():
26 | enabled_backends.append("math")
27 | if torch.backends.cuda.mem_efficient_sdp_enabled():
28 | enabled_backends.append("mem_efficient")
29 | if torch.backends.cuda.cudnn_sdp_enabled():
30 | enabled_backends.append("cudnn")
31 |
32 | print("Currently enabled native sdp backends:", enabled_backends)
33 |
34 | try:
35 | # raise NotImplementedError
36 | from xformers.ops import memory_efficient_attention as xformers_attn_func
37 | print('Xformers is installed!')
38 | except:
39 | print('Xformers is not installed!')
40 | xformers_attn_func = None
41 |
42 | try:
43 | # raise NotImplementedError
44 | from flash_attn import flash_attn_varlen_func, flash_attn_func
45 | print('Flash Attn is installed!')
46 | except:
47 | print('Flash Attn is not installed!')
48 | flash_attn_varlen_func = None
49 | flash_attn_func = None
50 |
51 | try:
52 | # raise NotImplementedError
53 | from sageattention import sageattn_varlen, sageattn
54 | print('Sage Attn is installed!')
55 | except:
56 | print('Sage Attn is not installed!')
57 | sageattn_varlen = None
58 | sageattn = None
59 |
60 |
61 | logger = logging.get_logger(__name__) # pylint: disable=invalid-name
62 |
63 |
64 | def pad_for_3d_conv(x, kernel_size):
65 | b, c, t, h, w = x.shape
66 | pt, ph, pw = kernel_size
67 | pad_t = (pt - (t % pt)) % pt
68 | pad_h = (ph - (h % ph)) % ph
69 | pad_w = (pw - (w % pw)) % pw
70 | return torch.nn.functional.pad(x, (0, pad_w, 0, pad_h, 0, pad_t), mode='replicate')
71 |
72 |
73 | def center_down_sample_3d(x, kernel_size):
74 | # pt, ph, pw = kernel_size
75 | # cp = (pt * ph * pw) // 2
76 | # xp = einops.rearrange(x, 'b c (t pt) (h ph) (w pw) -> (pt ph pw) b c t h w', pt=pt, ph=ph, pw=pw)
77 | # xc = xp[cp]
78 | # return xc
79 | return torch.nn.functional.avg_pool3d(x, kernel_size, stride=kernel_size)
80 |
81 |
82 | def get_cu_seqlens(text_mask, img_len):
83 | batch_size = text_mask.shape[0]
84 | text_len = text_mask.sum(dim=1)
85 | max_len = text_mask.shape[1] + img_len
86 |
87 | cu_seqlens = torch.zeros([2 * batch_size + 1], dtype=torch.int32, device="cuda")
88 |
89 | for i in range(batch_size):
90 | s = text_len[i] + img_len
91 | s1 = i * max_len + s
92 | s2 = (i + 1) * max_len
93 | cu_seqlens[2 * i + 1] = s1
94 | cu_seqlens[2 * i + 2] = s2
95 |
96 | return cu_seqlens
97 |
98 |
99 | def apply_rotary_emb_transposed(x, freqs_cis):
100 | cos, sin = freqs_cis.unsqueeze(-2).chunk(2, dim=-1)
101 | x_real, x_imag = x.unflatten(-1, (-1, 2)).unbind(-1)
102 | x_rotated = torch.stack([-x_imag, x_real], dim=-1).flatten(3)
103 | out = x.float() * cos + x_rotated.float() * sin
104 | out = out.to(x)
105 | return out
106 |
107 |
108 | def attn_varlen_func(q, k, v, cu_seqlens_q, cu_seqlens_kv, max_seqlen_q, max_seqlen_kv):
109 | if cu_seqlens_q is None and cu_seqlens_kv is None and max_seqlen_q is None and max_seqlen_kv is None:
110 | if sageattn is not None:
111 | x = sageattn(q, k, v, tensor_layout='NHD')
112 | return x
113 |
114 | if flash_attn_func is not None:
115 | x = flash_attn_func(q, k, v)
116 | return x
117 |
118 | if xformers_attn_func is not None:
119 | x = xformers_attn_func(q, k, v)
120 | return x
121 |
122 | x = torch.nn.functional.scaled_dot_product_attention(q.transpose(1, 2), k.transpose(1, 2), v.transpose(1, 2)).transpose(1, 2)
123 | return x
124 |
125 | batch_size = q.shape[0]
126 | q = q.view(q.shape[0] * q.shape[1], *q.shape[2:])
127 | k = k.view(k.shape[0] * k.shape[1], *k.shape[2:])
128 | v = v.view(v.shape[0] * v.shape[1], *v.shape[2:])
129 | if sageattn_varlen is not None:
130 | x = sageattn_varlen(q, k, v, cu_seqlens_q, cu_seqlens_kv, max_seqlen_q, max_seqlen_kv)
131 | elif flash_attn_varlen_func is not None:
132 | x = flash_attn_varlen_func(q, k, v, cu_seqlens_q, cu_seqlens_kv, max_seqlen_q, max_seqlen_kv)
133 | else:
134 | raise NotImplementedError('No Attn Installed!')
135 | x = x.view(batch_size, max_seqlen_q, *x.shape[2:])
136 | return x
137 |
138 |
139 | class HunyuanAttnProcessorFlashAttnDouble:
140 | def __call__(self, attn, hidden_states, encoder_hidden_states, attention_mask, image_rotary_emb):
141 | cu_seqlens_q, cu_seqlens_kv, max_seqlen_q, max_seqlen_kv = attention_mask
142 |
143 | query = attn.to_q(hidden_states)
144 | key = attn.to_k(hidden_states)
145 | value = attn.to_v(hidden_states)
146 |
147 | query = query.unflatten(2, (attn.heads, -1))
148 | key = key.unflatten(2, (attn.heads, -1))
149 | value = value.unflatten(2, (attn.heads, -1))
150 |
151 | query = attn.norm_q(query)
152 | key = attn.norm_k(key)
153 |
154 | query = apply_rotary_emb_transposed(query, image_rotary_emb)
155 | key = apply_rotary_emb_transposed(key, image_rotary_emb)
156 |
157 | encoder_query = attn.add_q_proj(encoder_hidden_states)
158 | encoder_key = attn.add_k_proj(encoder_hidden_states)
159 | encoder_value = attn.add_v_proj(encoder_hidden_states)
160 |
161 | encoder_query = encoder_query.unflatten(2, (attn.heads, -1))
162 | encoder_key = encoder_key.unflatten(2, (attn.heads, -1))
163 | encoder_value = encoder_value.unflatten(2, (attn.heads, -1))
164 |
165 | encoder_query = attn.norm_added_q(encoder_query)
166 | encoder_key = attn.norm_added_k(encoder_key)
167 |
168 | query = torch.cat([query, encoder_query], dim=1)
169 | key = torch.cat([key, encoder_key], dim=1)
170 | value = torch.cat([value, encoder_value], dim=1)
171 |
172 | hidden_states = attn_varlen_func(query, key, value, cu_seqlens_q, cu_seqlens_kv, max_seqlen_q, max_seqlen_kv)
173 | hidden_states = hidden_states.flatten(-2)
174 |
175 | txt_length = encoder_hidden_states.shape[1]
176 | hidden_states, encoder_hidden_states = hidden_states[:, :-txt_length], hidden_states[:, -txt_length:]
177 |
178 | hidden_states = attn.to_out[0](hidden_states)
179 | hidden_states = attn.to_out[1](hidden_states)
180 | encoder_hidden_states = attn.to_add_out(encoder_hidden_states)
181 |
182 | return hidden_states, encoder_hidden_states
183 |
184 |
185 | class HunyuanAttnProcessorFlashAttnSingle:
186 | def __call__(self, attn, hidden_states, encoder_hidden_states, attention_mask, image_rotary_emb):
187 | cu_seqlens_q, cu_seqlens_kv, max_seqlen_q, max_seqlen_kv = attention_mask
188 |
189 | hidden_states = torch.cat([hidden_states, encoder_hidden_states], dim=1)
190 |
191 | query = attn.to_q(hidden_states)
192 | key = attn.to_k(hidden_states)
193 | value = attn.to_v(hidden_states)
194 |
195 | query = query.unflatten(2, (attn.heads, -1))
196 | key = key.unflatten(2, (attn.heads, -1))
197 | value = value.unflatten(2, (attn.heads, -1))
198 |
199 | query = attn.norm_q(query)
200 | key = attn.norm_k(key)
201 |
202 | txt_length = encoder_hidden_states.shape[1]
203 |
204 | query = torch.cat([apply_rotary_emb_transposed(query[:, :-txt_length], image_rotary_emb), query[:, -txt_length:]], dim=1)
205 | key = torch.cat([apply_rotary_emb_transposed(key[:, :-txt_length], image_rotary_emb), key[:, -txt_length:]], dim=1)
206 |
207 | hidden_states = attn_varlen_func(query, key, value, cu_seqlens_q, cu_seqlens_kv, max_seqlen_q, max_seqlen_kv)
208 | hidden_states = hidden_states.flatten(-2)
209 |
210 | hidden_states, encoder_hidden_states = hidden_states[:, :-txt_length], hidden_states[:, -txt_length:]
211 |
212 | return hidden_states, encoder_hidden_states
213 |
214 |
215 | class CombinedTimestepGuidanceTextProjEmbeddings(nn.Module):
216 | def __init__(self, embedding_dim, pooled_projection_dim):
217 | super().__init__()
218 |
219 | self.time_proj = Timesteps(num_channels=256, flip_sin_to_cos=True, downscale_freq_shift=0)
220 | self.timestep_embedder = TimestepEmbedding(in_channels=256, time_embed_dim=embedding_dim)
221 | self.guidance_embedder = TimestepEmbedding(in_channels=256, time_embed_dim=embedding_dim)
222 | self.text_embedder = PixArtAlphaTextProjection(pooled_projection_dim, embedding_dim, act_fn="silu")
223 |
224 | def forward(self, timestep, guidance, pooled_projection):
225 | timesteps_proj = self.time_proj(timestep)
226 | timesteps_emb = self.timestep_embedder(timesteps_proj.to(dtype=pooled_projection.dtype))
227 |
228 | guidance_proj = self.time_proj(guidance)
229 | guidance_emb = self.guidance_embedder(guidance_proj.to(dtype=pooled_projection.dtype))
230 |
231 | time_guidance_emb = timesteps_emb + guidance_emb
232 |
233 | pooled_projections = self.text_embedder(pooled_projection)
234 | conditioning = time_guidance_emb + pooled_projections
235 |
236 | return conditioning
237 |
238 |
239 | class CombinedTimestepTextProjEmbeddings(nn.Module):
240 | def __init__(self, embedding_dim, pooled_projection_dim):
241 | super().__init__()
242 |
243 | self.time_proj = Timesteps(num_channels=256, flip_sin_to_cos=True, downscale_freq_shift=0)
244 | self.timestep_embedder = TimestepEmbedding(in_channels=256, time_embed_dim=embedding_dim)
245 | self.text_embedder = PixArtAlphaTextProjection(pooled_projection_dim, embedding_dim, act_fn="silu")
246 |
247 | def forward(self, timestep, pooled_projection):
248 | timesteps_proj = self.time_proj(timestep)
249 | timesteps_emb = self.timestep_embedder(timesteps_proj.to(dtype=pooled_projection.dtype))
250 |
251 | pooled_projections = self.text_embedder(pooled_projection)
252 |
253 | conditioning = timesteps_emb + pooled_projections
254 |
255 | return conditioning
256 |
257 |
258 | class HunyuanVideoAdaNorm(nn.Module):
259 | def __init__(self, in_features: int, out_features: Optional[int] = None) -> None:
260 | super().__init__()
261 |
262 | out_features = out_features or 2 * in_features
263 | self.linear = nn.Linear(in_features, out_features)
264 | self.nonlinearity = nn.SiLU()
265 |
266 | def forward(
267 | self, temb: torch.Tensor
268 | ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
269 | temb = self.linear(self.nonlinearity(temb))
270 | gate_msa, gate_mlp = temb.chunk(2, dim=-1)
271 | gate_msa, gate_mlp = gate_msa.unsqueeze(1), gate_mlp.unsqueeze(1)
272 | return gate_msa, gate_mlp
273 |
274 |
275 | class HunyuanVideoIndividualTokenRefinerBlock(nn.Module):
276 | def __init__(
277 | self,
278 | num_attention_heads: int,
279 | attention_head_dim: int,
280 | mlp_width_ratio: str = 4.0,
281 | mlp_drop_rate: float = 0.0,
282 | attention_bias: bool = True,
283 | ) -> None:
284 | super().__init__()
285 |
286 | hidden_size = num_attention_heads * attention_head_dim
287 |
288 | self.norm1 = LayerNorm(hidden_size, elementwise_affine=True, eps=1e-6)
289 | self.attn = Attention(
290 | query_dim=hidden_size,
291 | cross_attention_dim=None,
292 | heads=num_attention_heads,
293 | dim_head=attention_head_dim,
294 | bias=attention_bias,
295 | )
296 |
297 | self.norm2 = LayerNorm(hidden_size, elementwise_affine=True, eps=1e-6)
298 | self.ff = FeedForward(hidden_size, mult=mlp_width_ratio, activation_fn="linear-silu", dropout=mlp_drop_rate)
299 |
300 | self.norm_out = HunyuanVideoAdaNorm(hidden_size, 2 * hidden_size)
301 |
302 | def forward(
303 | self,
304 | hidden_states: torch.Tensor,
305 | temb: torch.Tensor,
306 | attention_mask: Optional[torch.Tensor] = None,
307 | ) -> torch.Tensor:
308 | norm_hidden_states = self.norm1(hidden_states)
309 |
310 | attn_output = self.attn(
311 | hidden_states=norm_hidden_states,
312 | encoder_hidden_states=None,
313 | attention_mask=attention_mask,
314 | )
315 |
316 | gate_msa, gate_mlp = self.norm_out(temb)
317 | hidden_states = hidden_states + attn_output * gate_msa
318 |
319 | ff_output = self.ff(self.norm2(hidden_states))
320 | hidden_states = hidden_states + ff_output * gate_mlp
321 |
322 | return hidden_states
323 |
324 |
325 | class HunyuanVideoIndividualTokenRefiner(nn.Module):
326 | def __init__(
327 | self,
328 | num_attention_heads: int,
329 | attention_head_dim: int,
330 | num_layers: int,
331 | mlp_width_ratio: float = 4.0,
332 | mlp_drop_rate: float = 0.0,
333 | attention_bias: bool = True,
334 | ) -> None:
335 | super().__init__()
336 |
337 | self.refiner_blocks = nn.ModuleList(
338 | [
339 | HunyuanVideoIndividualTokenRefinerBlock(
340 | num_attention_heads=num_attention_heads,
341 | attention_head_dim=attention_head_dim,
342 | mlp_width_ratio=mlp_width_ratio,
343 | mlp_drop_rate=mlp_drop_rate,
344 | attention_bias=attention_bias,
345 | )
346 | for _ in range(num_layers)
347 | ]
348 | )
349 |
350 | def forward(
351 | self,
352 | hidden_states: torch.Tensor,
353 | temb: torch.Tensor,
354 | attention_mask: Optional[torch.Tensor] = None,
355 | ) -> None:
356 | self_attn_mask = None
357 | if attention_mask is not None:
358 | batch_size = attention_mask.shape[0]
359 | seq_len = attention_mask.shape[1]
360 | attention_mask = attention_mask.to(hidden_states.device).bool()
361 | self_attn_mask_1 = attention_mask.view(batch_size, 1, 1, seq_len).repeat(1, 1, seq_len, 1)
362 | self_attn_mask_2 = self_attn_mask_1.transpose(2, 3)
363 | self_attn_mask = (self_attn_mask_1 & self_attn_mask_2).bool()
364 | self_attn_mask[:, :, :, 0] = True
365 |
366 | for block in self.refiner_blocks:
367 | hidden_states = block(hidden_states, temb, self_attn_mask)
368 |
369 | return hidden_states
370 |
371 |
372 | class HunyuanVideoTokenRefiner(nn.Module):
373 | def __init__(
374 | self,
375 | in_channels: int,
376 | num_attention_heads: int,
377 | attention_head_dim: int,
378 | num_layers: int,
379 | mlp_ratio: float = 4.0,
380 | mlp_drop_rate: float = 0.0,
381 | attention_bias: bool = True,
382 | ) -> None:
383 | super().__init__()
384 |
385 | hidden_size = num_attention_heads * attention_head_dim
386 |
387 | self.time_text_embed = CombinedTimestepTextProjEmbeddings(
388 | embedding_dim=hidden_size, pooled_projection_dim=in_channels
389 | )
390 | self.proj_in = nn.Linear(in_channels, hidden_size, bias=True)
391 | self.token_refiner = HunyuanVideoIndividualTokenRefiner(
392 | num_attention_heads=num_attention_heads,
393 | attention_head_dim=attention_head_dim,
394 | num_layers=num_layers,
395 | mlp_width_ratio=mlp_ratio,
396 | mlp_drop_rate=mlp_drop_rate,
397 | attention_bias=attention_bias,
398 | )
399 |
400 | def forward(
401 | self,
402 | hidden_states: torch.Tensor,
403 | timestep: torch.LongTensor,
404 | attention_mask: Optional[torch.LongTensor] = None,
405 | ) -> torch.Tensor:
406 | if attention_mask is None:
407 | pooled_projections = hidden_states.mean(dim=1)
408 | else:
409 | original_dtype = hidden_states.dtype
410 | mask_float = attention_mask.float().unsqueeze(-1)
411 | pooled_projections = (hidden_states * mask_float).sum(dim=1) / mask_float.sum(dim=1)
412 | pooled_projections = pooled_projections.to(original_dtype)
413 |
414 | temb = self.time_text_embed(timestep, pooled_projections)
415 | hidden_states = self.proj_in(hidden_states)
416 | hidden_states = self.token_refiner(hidden_states, temb, attention_mask)
417 |
418 | return hidden_states
419 |
420 |
421 | class HunyuanVideoRotaryPosEmbed(nn.Module):
422 | def __init__(self, rope_dim, theta):
423 | super().__init__()
424 | self.DT, self.DY, self.DX = rope_dim
425 | self.theta = theta
426 |
427 | @torch.no_grad()
428 | def get_frequency(self, dim, pos):
429 | T, H, W = pos.shape
430 | freqs = 1.0 / (self.theta ** (torch.arange(0, dim, 2, dtype=torch.float32, device=pos.device)[: (dim // 2)] / dim))
431 | freqs = torch.outer(freqs, pos.reshape(-1)).unflatten(-1, (T, H, W)).repeat_interleave(2, dim=0)
432 | return freqs.cos(), freqs.sin()
433 |
434 | @torch.no_grad()
435 | def forward_inner(self, frame_indices, height, width, device):
436 | GT, GY, GX = torch.meshgrid(
437 | frame_indices.to(device=device, dtype=torch.float32),
438 | torch.arange(0, height, device=device, dtype=torch.float32),
439 | torch.arange(0, width, device=device, dtype=torch.float32),
440 | indexing="ij"
441 | )
442 |
443 | FCT, FST = self.get_frequency(self.DT, GT)
444 | FCY, FSY = self.get_frequency(self.DY, GY)
445 | FCX, FSX = self.get_frequency(self.DX, GX)
446 |
447 | result = torch.cat([FCT, FCY, FCX, FST, FSY, FSX], dim=0)
448 |
449 | return result.to(device)
450 |
451 | @torch.no_grad()
452 | def forward(self, frame_indices, height, width, device):
453 | frame_indices = frame_indices.unbind(0)
454 | results = [self.forward_inner(f, height, width, device) for f in frame_indices]
455 | results = torch.stack(results, dim=0)
456 | return results
457 |
458 |
459 | class AdaLayerNormZero(nn.Module):
460 | def __init__(self, embedding_dim: int, norm_type="layer_norm", bias=True):
461 | super().__init__()
462 | self.silu = nn.SiLU()
463 | self.linear = nn.Linear(embedding_dim, 6 * embedding_dim, bias=bias)
464 | if norm_type == "layer_norm":
465 | self.norm = LayerNorm(embedding_dim, elementwise_affine=False, eps=1e-6)
466 | else:
467 | raise ValueError(f"unknown norm_type {norm_type}")
468 |
469 | def forward(
470 | self,
471 | x: torch.Tensor,
472 | emb: Optional[torch.Tensor] = None,
473 | ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
474 | emb = emb.unsqueeze(-2)
475 | emb = self.linear(self.silu(emb))
476 | shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = emb.chunk(6, dim=-1)
477 | x = self.norm(x) * (1 + scale_msa) + shift_msa
478 | return x, gate_msa, shift_mlp, scale_mlp, gate_mlp
479 |
480 |
481 | class AdaLayerNormZeroSingle(nn.Module):
482 | def __init__(self, embedding_dim: int, norm_type="layer_norm", bias=True):
483 | super().__init__()
484 |
485 | self.silu = nn.SiLU()
486 | self.linear = nn.Linear(embedding_dim, 3 * embedding_dim, bias=bias)
487 | if norm_type == "layer_norm":
488 | self.norm = LayerNorm(embedding_dim, elementwise_affine=False, eps=1e-6)
489 | else:
490 | raise ValueError(f"unknown norm_type {norm_type}")
491 |
492 | def forward(
493 | self,
494 | x: torch.Tensor,
495 | emb: Optional[torch.Tensor] = None,
496 | ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
497 | emb = emb.unsqueeze(-2)
498 | emb = self.linear(self.silu(emb))
499 | shift_msa, scale_msa, gate_msa = emb.chunk(3, dim=-1)
500 | x = self.norm(x) * (1 + scale_msa) + shift_msa
501 | return x, gate_msa
502 |
503 |
504 | class AdaLayerNormContinuous(nn.Module):
505 | def __init__(
506 | self,
507 | embedding_dim: int,
508 | conditioning_embedding_dim: int,
509 | elementwise_affine=True,
510 | eps=1e-5,
511 | bias=True,
512 | norm_type="layer_norm",
513 | ):
514 | super().__init__()
515 | self.silu = nn.SiLU()
516 | self.linear = nn.Linear(conditioning_embedding_dim, embedding_dim * 2, bias=bias)
517 | if norm_type == "layer_norm":
518 | self.norm = LayerNorm(embedding_dim, eps, elementwise_affine, bias)
519 | else:
520 | raise ValueError(f"unknown norm_type {norm_type}")
521 |
522 | def forward(self, x: torch.Tensor, emb: torch.Tensor) -> torch.Tensor:
523 | emb = emb.unsqueeze(-2)
524 | emb = self.linear(self.silu(emb))
525 | scale, shift = emb.chunk(2, dim=-1)
526 | x = self.norm(x) * (1 + scale) + shift
527 | return x
528 |
529 |
530 | class HunyuanVideoSingleTransformerBlock(nn.Module):
531 | def __init__(
532 | self,
533 | num_attention_heads: int,
534 | attention_head_dim: int,
535 | mlp_ratio: float = 4.0,
536 | qk_norm: str = "rms_norm",
537 | ) -> None:
538 | super().__init__()
539 |
540 | hidden_size = num_attention_heads * attention_head_dim
541 | mlp_dim = int(hidden_size * mlp_ratio)
542 |
543 | self.attn = Attention(
544 | query_dim=hidden_size,
545 | cross_attention_dim=None,
546 | dim_head=attention_head_dim,
547 | heads=num_attention_heads,
548 | out_dim=hidden_size,
549 | bias=True,
550 | processor=HunyuanAttnProcessorFlashAttnSingle(),
551 | qk_norm=qk_norm,
552 | eps=1e-6,
553 | pre_only=True,
554 | )
555 |
556 | self.norm = AdaLayerNormZeroSingle(hidden_size, norm_type="layer_norm")
557 | self.proj_mlp = nn.Linear(hidden_size, mlp_dim)
558 | self.act_mlp = nn.GELU(approximate="tanh")
559 | self.proj_out = nn.Linear(hidden_size + mlp_dim, hidden_size)
560 |
561 | def forward(
562 | self,
563 | hidden_states: torch.Tensor,
564 | encoder_hidden_states: torch.Tensor,
565 | temb: torch.Tensor,
566 | attention_mask: Optional[torch.Tensor] = None,
567 | image_rotary_emb: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
568 | ) -> torch.Tensor:
569 | text_seq_length = encoder_hidden_states.shape[1]
570 | hidden_states = torch.cat([hidden_states, encoder_hidden_states], dim=1)
571 |
572 | residual = hidden_states
573 |
574 | # 1. Input normalization
575 | norm_hidden_states, gate = self.norm(hidden_states, emb=temb)
576 | mlp_hidden_states = self.act_mlp(self.proj_mlp(norm_hidden_states))
577 |
578 | norm_hidden_states, norm_encoder_hidden_states = (
579 | norm_hidden_states[:, :-text_seq_length, :],
580 | norm_hidden_states[:, -text_seq_length:, :],
581 | )
582 |
583 | # 2. Attention
584 | attn_output, context_attn_output = self.attn(
585 | hidden_states=norm_hidden_states,
586 | encoder_hidden_states=norm_encoder_hidden_states,
587 | attention_mask=attention_mask,
588 | image_rotary_emb=image_rotary_emb,
589 | )
590 | attn_output = torch.cat([attn_output, context_attn_output], dim=1)
591 |
592 | # 3. Modulation and residual connection
593 | hidden_states = torch.cat([attn_output, mlp_hidden_states], dim=2)
594 | hidden_states = gate * self.proj_out(hidden_states)
595 | hidden_states = hidden_states + residual
596 |
597 | hidden_states, encoder_hidden_states = (
598 | hidden_states[:, :-text_seq_length, :],
599 | hidden_states[:, -text_seq_length:, :],
600 | )
601 | return hidden_states, encoder_hidden_states
602 |
603 |
604 | class HunyuanVideoTransformerBlock(nn.Module):
605 | def __init__(
606 | self,
607 | num_attention_heads: int,
608 | attention_head_dim: int,
609 | mlp_ratio: float,
610 | qk_norm: str = "rms_norm",
611 | ) -> None:
612 | super().__init__()
613 |
614 | hidden_size = num_attention_heads * attention_head_dim
615 |
616 | self.norm1 = AdaLayerNormZero(hidden_size, norm_type="layer_norm")
617 | self.norm1_context = AdaLayerNormZero(hidden_size, norm_type="layer_norm")
618 |
619 | self.attn = Attention(
620 | query_dim=hidden_size,
621 | cross_attention_dim=None,
622 | added_kv_proj_dim=hidden_size,
623 | dim_head=attention_head_dim,
624 | heads=num_attention_heads,
625 | out_dim=hidden_size,
626 | context_pre_only=False,
627 | bias=True,
628 | processor=HunyuanAttnProcessorFlashAttnDouble(),
629 | qk_norm=qk_norm,
630 | eps=1e-6,
631 | )
632 |
633 | self.norm2 = LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6)
634 | self.ff = FeedForward(hidden_size, mult=mlp_ratio, activation_fn="gelu-approximate")
635 |
636 | self.norm2_context = LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6)
637 | self.ff_context = FeedForward(hidden_size, mult=mlp_ratio, activation_fn="gelu-approximate")
638 |
639 | def forward(
640 | self,
641 | hidden_states: torch.Tensor,
642 | encoder_hidden_states: torch.Tensor,
643 | temb: torch.Tensor,
644 | attention_mask: Optional[torch.Tensor] = None,
645 | freqs_cis: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
646 | ) -> Tuple[torch.Tensor, torch.Tensor]:
647 | # 1. Input normalization
648 | norm_hidden_states, gate_msa, shift_mlp, scale_mlp, gate_mlp = self.norm1(hidden_states, emb=temb)
649 | norm_encoder_hidden_states, c_gate_msa, c_shift_mlp, c_scale_mlp, c_gate_mlp = self.norm1_context(encoder_hidden_states, emb=temb)
650 |
651 | # 2. Joint attention
652 | attn_output, context_attn_output = self.attn(
653 | hidden_states=norm_hidden_states,
654 | encoder_hidden_states=norm_encoder_hidden_states,
655 | attention_mask=attention_mask,
656 | image_rotary_emb=freqs_cis,
657 | )
658 |
659 | # 3. Modulation and residual connection
660 | hidden_states = hidden_states + attn_output * gate_msa
661 | encoder_hidden_states = encoder_hidden_states + context_attn_output * c_gate_msa
662 |
663 | norm_hidden_states = self.norm2(hidden_states)
664 | norm_encoder_hidden_states = self.norm2_context(encoder_hidden_states)
665 |
666 | norm_hidden_states = norm_hidden_states * (1 + scale_mlp) + shift_mlp
667 | norm_encoder_hidden_states = norm_encoder_hidden_states * (1 + c_scale_mlp) + c_shift_mlp
668 |
669 | # 4. Feed-forward
670 | ff_output = self.ff(norm_hidden_states)
671 | context_ff_output = self.ff_context(norm_encoder_hidden_states)
672 |
673 | hidden_states = hidden_states + gate_mlp * ff_output
674 | encoder_hidden_states = encoder_hidden_states + c_gate_mlp * context_ff_output
675 |
676 | return hidden_states, encoder_hidden_states
677 |
678 |
679 | class ClipVisionProjection(nn.Module):
680 | def __init__(self, in_channels, out_channels):
681 | super().__init__()
682 | self.up = nn.Linear(in_channels, out_channels * 3)
683 | self.down = nn.Linear(out_channels * 3, out_channels)
684 |
685 | def forward(self, x):
686 | projected_x = self.down(nn.functional.silu(self.up(x)))
687 | return projected_x
688 |
689 |
690 | class HunyuanVideoPatchEmbed(nn.Module):
691 | def __init__(self, patch_size, in_chans, embed_dim):
692 | super().__init__()
693 | self.proj = nn.Conv3d(in_chans, embed_dim, kernel_size=patch_size, stride=patch_size)
694 |
695 |
696 | class HunyuanVideoPatchEmbedForCleanLatents(nn.Module):
697 | def __init__(self, inner_dim):
698 | super().__init__()
699 | self.proj = nn.Conv3d(16, inner_dim, kernel_size=(1, 2, 2), stride=(1, 2, 2))
700 | self.proj_2x = nn.Conv3d(16, inner_dim, kernel_size=(2, 4, 4), stride=(2, 4, 4))
701 | self.proj_4x = nn.Conv3d(16, inner_dim, kernel_size=(4, 8, 8), stride=(4, 8, 8))
702 |
703 | @torch.no_grad()
704 | def initialize_weight_from_another_conv3d(self, another_layer):
705 | weight = another_layer.weight.detach().clone()
706 | bias = another_layer.bias.detach().clone()
707 |
708 | sd = {
709 | 'proj.weight': weight.clone(),
710 | 'proj.bias': bias.clone(),
711 | 'proj_2x.weight': einops.repeat(weight, 'b c t h w -> b c (t tk) (h hk) (w wk)', tk=2, hk=2, wk=2) / 8.0,
712 | 'proj_2x.bias': bias.clone(),
713 | 'proj_4x.weight': einops.repeat(weight, 'b c t h w -> b c (t tk) (h hk) (w wk)', tk=4, hk=4, wk=4) / 64.0,
714 | 'proj_4x.bias': bias.clone(),
715 | }
716 |
717 | sd = {k: v.clone() for k, v in sd.items()}
718 |
719 | self.load_state_dict(sd)
720 | return
721 |
722 |
723 | class HunyuanVideoTransformer3DModelPacked(ModelMixin, ConfigMixin, PeftAdapterMixin, FromOriginalModelMixin):
724 | @register_to_config
725 | def __init__(
726 | self,
727 | in_channels: int = 16,
728 | out_channels: int = 16,
729 | num_attention_heads: int = 24,
730 | attention_head_dim: int = 128,
731 | num_layers: int = 20,
732 | num_single_layers: int = 40,
733 | num_refiner_layers: int = 2,
734 | mlp_ratio: float = 4.0,
735 | patch_size: int = 2,
736 | patch_size_t: int = 1,
737 | qk_norm: str = "rms_norm",
738 | guidance_embeds: bool = True,
739 | text_embed_dim: int = 4096,
740 | pooled_projection_dim: int = 768,
741 | rope_theta: float = 256.0,
742 | rope_axes_dim: Tuple[int] = (16, 56, 56),
743 | has_image_proj=False,
744 | image_proj_dim=1152,
745 | has_clean_x_embedder=False,
746 | ) -> None:
747 | super().__init__()
748 |
749 | inner_dim = num_attention_heads * attention_head_dim
750 | out_channels = out_channels or in_channels
751 |
752 | # 1. Latent and condition embedders
753 | self.x_embedder = HunyuanVideoPatchEmbed((patch_size_t, patch_size, patch_size), in_channels, inner_dim)
754 | self.context_embedder = HunyuanVideoTokenRefiner(
755 | text_embed_dim, num_attention_heads, attention_head_dim, num_layers=num_refiner_layers
756 | )
757 | self.time_text_embed = CombinedTimestepGuidanceTextProjEmbeddings(inner_dim, pooled_projection_dim)
758 |
759 | self.clean_x_embedder = None
760 | self.image_projection = None
761 |
762 | # 2. RoPE
763 | self.rope = HunyuanVideoRotaryPosEmbed(rope_axes_dim, rope_theta)
764 |
765 | # 3. Dual stream transformer blocks
766 | self.transformer_blocks = nn.ModuleList(
767 | [
768 | HunyuanVideoTransformerBlock(
769 | num_attention_heads, attention_head_dim, mlp_ratio=mlp_ratio, qk_norm=qk_norm
770 | )
771 | for _ in range(num_layers)
772 | ]
773 | )
774 |
775 | # 4. Single stream transformer blocks
776 | self.single_transformer_blocks = nn.ModuleList(
777 | [
778 | HunyuanVideoSingleTransformerBlock(
779 | num_attention_heads, attention_head_dim, mlp_ratio=mlp_ratio, qk_norm=qk_norm
780 | )
781 | for _ in range(num_single_layers)
782 | ]
783 | )
784 |
785 | # 5. Output projection
786 | self.norm_out = AdaLayerNormContinuous(inner_dim, inner_dim, elementwise_affine=False, eps=1e-6)
787 | self.proj_out = nn.Linear(inner_dim, patch_size_t * patch_size * patch_size * out_channels)
788 |
789 | self.inner_dim = inner_dim
790 | self.use_gradient_checkpointing = False
791 | self.enable_teacache = False
792 |
793 | if has_image_proj:
794 | self.install_image_projection(image_proj_dim)
795 |
796 | if has_clean_x_embedder:
797 | self.install_clean_x_embedder()
798 |
799 | self.high_quality_fp32_output_for_inference = False
800 |
801 | def install_image_projection(self, in_channels):
802 | self.image_projection = ClipVisionProjection(in_channels=in_channels, out_channels=self.inner_dim)
803 | self.config['has_image_proj'] = True
804 | self.config['image_proj_dim'] = in_channels
805 |
806 | def install_clean_x_embedder(self):
807 | self.clean_x_embedder = HunyuanVideoPatchEmbedForCleanLatents(self.inner_dim)
808 | self.config['has_clean_x_embedder'] = True
809 |
810 | def enable_gradient_checkpointing(self):
811 | self.use_gradient_checkpointing = True
812 | print('self.use_gradient_checkpointing = True')
813 |
814 | def disable_gradient_checkpointing(self):
815 | self.use_gradient_checkpointing = False
816 | print('self.use_gradient_checkpointing = False')
817 |
818 | def initialize_teacache(self, enable_teacache=True, num_steps=25, rel_l1_thresh=0.15):
819 | self.enable_teacache = enable_teacache
820 | self.cnt = 0
821 | self.num_steps = num_steps
822 | self.rel_l1_thresh = rel_l1_thresh # 0.1 for 1.6x speedup, 0.15 for 2.1x speedup
823 | self.accumulated_rel_l1_distance = 0
824 | self.previous_modulated_input = None
825 | self.previous_residual = None
826 | self.teacache_rescale_func = np.poly1d([7.33226126e+02, -4.01131952e+02, 6.75869174e+01, -3.14987800e+00, 9.61237896e-02])
827 |
828 | def gradient_checkpointing_method(self, block, *args):
829 | if self.use_gradient_checkpointing:
830 | result = torch.utils.checkpoint.checkpoint(block, *args, use_reentrant=False)
831 | else:
832 | result = block(*args)
833 | return result
834 |
835 | def process_input_hidden_states(
836 | self,
837 | latents, latent_indices=None,
838 | clean_latents=None, clean_latent_indices=None,
839 | clean_latents_2x=None, clean_latent_2x_indices=None,
840 | clean_latents_4x=None, clean_latent_4x_indices=None
841 | ):
842 | hidden_states = self.gradient_checkpointing_method(self.x_embedder.proj, latents)
843 | B, C, T, H, W = hidden_states.shape
844 |
845 | if latent_indices is None:
846 | latent_indices = torch.arange(0, T).unsqueeze(0).expand(B, -1)
847 |
848 | hidden_states = hidden_states.flatten(2).transpose(1, 2)
849 |
850 | rope_freqs = self.rope(frame_indices=latent_indices, height=H, width=W, device=hidden_states.device)
851 | rope_freqs = rope_freqs.flatten(2).transpose(1, 2)
852 |
853 | if clean_latents is not None and clean_latent_indices is not None:
854 | clean_latents = clean_latents.to(hidden_states)
855 | clean_latents = self.gradient_checkpointing_method(self.clean_x_embedder.proj, clean_latents)
856 | clean_latents = clean_latents.flatten(2).transpose(1, 2)
857 |
858 | clean_latent_rope_freqs = self.rope(frame_indices=clean_latent_indices, height=H, width=W, device=clean_latents.device)
859 | clean_latent_rope_freqs = clean_latent_rope_freqs.flatten(2).transpose(1, 2)
860 |
861 | hidden_states = torch.cat([clean_latents, hidden_states], dim=1)
862 | rope_freqs = torch.cat([clean_latent_rope_freqs, rope_freqs], dim=1)
863 |
864 | if clean_latents_2x is not None and clean_latent_2x_indices is not None:
865 | clean_latents_2x = clean_latents_2x.to(hidden_states)
866 | clean_latents_2x = pad_for_3d_conv(clean_latents_2x, (2, 4, 4))
867 | clean_latents_2x = self.gradient_checkpointing_method(self.clean_x_embedder.proj_2x, clean_latents_2x)
868 | clean_latents_2x = clean_latents_2x.flatten(2).transpose(1, 2)
869 |
870 | clean_latent_2x_rope_freqs = self.rope(frame_indices=clean_latent_2x_indices, height=H, width=W, device=clean_latents_2x.device)
871 | clean_latent_2x_rope_freqs = pad_for_3d_conv(clean_latent_2x_rope_freqs, (2, 2, 2))
872 | clean_latent_2x_rope_freqs = center_down_sample_3d(clean_latent_2x_rope_freqs, (2, 2, 2))
873 | clean_latent_2x_rope_freqs = clean_latent_2x_rope_freqs.flatten(2).transpose(1, 2)
874 |
875 | hidden_states = torch.cat([clean_latents_2x, hidden_states], dim=1)
876 | rope_freqs = torch.cat([clean_latent_2x_rope_freqs, rope_freqs], dim=1)
877 |
878 | if clean_latents_4x is not None and clean_latent_4x_indices is not None:
879 | clean_latents_4x = clean_latents_4x.to(hidden_states)
880 | clean_latents_4x = pad_for_3d_conv(clean_latents_4x, (4, 8, 8))
881 | clean_latents_4x = self.gradient_checkpointing_method(self.clean_x_embedder.proj_4x, clean_latents_4x)
882 | clean_latents_4x = clean_latents_4x.flatten(2).transpose(1, 2)
883 |
884 | clean_latent_4x_rope_freqs = self.rope(frame_indices=clean_latent_4x_indices, height=H, width=W, device=clean_latents_4x.device)
885 | clean_latent_4x_rope_freqs = pad_for_3d_conv(clean_latent_4x_rope_freqs, (4, 4, 4))
886 | clean_latent_4x_rope_freqs = center_down_sample_3d(clean_latent_4x_rope_freqs, (4, 4, 4))
887 | clean_latent_4x_rope_freqs = clean_latent_4x_rope_freqs.flatten(2).transpose(1, 2)
888 |
889 | hidden_states = torch.cat([clean_latents_4x, hidden_states], dim=1)
890 | rope_freqs = torch.cat([clean_latent_4x_rope_freqs, rope_freqs], dim=1)
891 |
892 | return hidden_states, rope_freqs
893 |
894 | def forward(
895 | self,
896 | hidden_states, timestep, encoder_hidden_states, encoder_attention_mask, pooled_projections, guidance,
897 | latent_indices=None,
898 | clean_latents=None, clean_latent_indices=None,
899 | clean_latents_2x=None, clean_latent_2x_indices=None,
900 | clean_latents_4x=None, clean_latent_4x_indices=None,
901 | image_embeddings=None,
902 | attention_kwargs=None, return_dict=True
903 | ):
904 |
905 | if attention_kwargs is None:
906 | attention_kwargs = {}
907 |
908 | batch_size, num_channels, num_frames, height, width = hidden_states.shape
909 | p, p_t = self.config['patch_size'], self.config['patch_size_t']
910 | post_patch_num_frames = num_frames // p_t
911 | post_patch_height = height // p
912 | post_patch_width = width // p
913 | original_context_length = post_patch_num_frames * post_patch_height * post_patch_width
914 |
915 | hidden_states, rope_freqs = self.process_input_hidden_states(hidden_states, latent_indices, clean_latents, clean_latent_indices, clean_latents_2x, clean_latent_2x_indices, clean_latents_4x, clean_latent_4x_indices)
916 |
917 | temb = self.gradient_checkpointing_method(self.time_text_embed, timestep, guidance, pooled_projections)
918 | encoder_hidden_states = self.gradient_checkpointing_method(self.context_embedder, encoder_hidden_states, timestep, encoder_attention_mask)
919 |
920 | if self.image_projection is not None:
921 | assert image_embeddings is not None, 'You must use image embeddings!'
922 | extra_encoder_hidden_states = self.gradient_checkpointing_method(self.image_projection, image_embeddings)
923 | extra_attention_mask = torch.ones((batch_size, extra_encoder_hidden_states.shape[1]), dtype=encoder_attention_mask.dtype, device=encoder_attention_mask.device)
924 |
925 | # must cat before (not after) encoder_hidden_states, due to attn masking
926 | encoder_hidden_states = torch.cat([extra_encoder_hidden_states, encoder_hidden_states], dim=1)
927 | encoder_attention_mask = torch.cat([extra_attention_mask, encoder_attention_mask], dim=1)
928 |
929 | with torch.no_grad():
930 | if batch_size == 1:
931 | # When batch size is 1, we do not need any masks or var-len funcs since cropping is mathematically same to what we want
932 | # If they are not same, then their impls are wrong. Ours are always the correct one.
933 | text_len = encoder_attention_mask.sum().item()
934 | encoder_hidden_states = encoder_hidden_states[:, :text_len]
935 | attention_mask = None, None, None, None
936 | else:
937 | img_seq_len = hidden_states.shape[1]
938 | txt_seq_len = encoder_hidden_states.shape[1]
939 |
940 | cu_seqlens_q = get_cu_seqlens(encoder_attention_mask, img_seq_len)
941 | cu_seqlens_kv = cu_seqlens_q
942 | max_seqlen_q = img_seq_len + txt_seq_len
943 | max_seqlen_kv = max_seqlen_q
944 |
945 | attention_mask = cu_seqlens_q, cu_seqlens_kv, max_seqlen_q, max_seqlen_kv
946 |
947 | if self.enable_teacache:
948 | modulated_inp = self.transformer_blocks[0].norm1(hidden_states, emb=temb)[0]
949 |
950 | if self.cnt == 0 or self.cnt == self.num_steps-1:
951 | should_calc = True
952 | self.accumulated_rel_l1_distance = 0
953 | else:
954 | curr_rel_l1 = ((modulated_inp - self.previous_modulated_input).abs().mean() / self.previous_modulated_input.abs().mean()).cpu().item()
955 | self.accumulated_rel_l1_distance += self.teacache_rescale_func(curr_rel_l1)
956 | should_calc = self.accumulated_rel_l1_distance >= self.rel_l1_thresh
957 |
958 | if should_calc:
959 | self.accumulated_rel_l1_distance = 0
960 |
961 | self.previous_modulated_input = modulated_inp
962 | self.cnt += 1
963 |
964 | if self.cnt == self.num_steps:
965 | self.cnt = 0
966 |
967 | if not should_calc:
968 | hidden_states = hidden_states + self.previous_residual
969 | else:
970 | ori_hidden_states = hidden_states.clone()
971 |
972 | for block_id, block in enumerate(self.transformer_blocks):
973 | hidden_states, encoder_hidden_states = self.gradient_checkpointing_method(
974 | block,
975 | hidden_states,
976 | encoder_hidden_states,
977 | temb,
978 | attention_mask,
979 | rope_freqs
980 | )
981 |
982 | for block_id, block in enumerate(self.single_transformer_blocks):
983 | hidden_states, encoder_hidden_states = self.gradient_checkpointing_method(
984 | block,
985 | hidden_states,
986 | encoder_hidden_states,
987 | temb,
988 | attention_mask,
989 | rope_freqs
990 | )
991 |
992 | self.previous_residual = hidden_states - ori_hidden_states
993 | else:
994 | for block_id, block in enumerate(self.transformer_blocks):
995 | hidden_states, encoder_hidden_states = self.gradient_checkpointing_method(
996 | block,
997 | hidden_states,
998 | encoder_hidden_states,
999 | temb,
1000 | attention_mask,
1001 | rope_freqs
1002 | )
1003 |
1004 | for block_id, block in enumerate(self.single_transformer_blocks):
1005 | hidden_states, encoder_hidden_states = self.gradient_checkpointing_method(
1006 | block,
1007 | hidden_states,
1008 | encoder_hidden_states,
1009 | temb,
1010 | attention_mask,
1011 | rope_freqs
1012 | )
1013 |
1014 | hidden_states = self.gradient_checkpointing_method(self.norm_out, hidden_states, temb)
1015 |
1016 | hidden_states = hidden_states[:, -original_context_length:, :]
1017 |
1018 | if self.high_quality_fp32_output_for_inference:
1019 | hidden_states = hidden_states.to(dtype=torch.float32)
1020 | if self.proj_out.weight.dtype != torch.float32:
1021 | self.proj_out.to(dtype=torch.float32)
1022 |
1023 | hidden_states = self.gradient_checkpointing_method(self.proj_out, hidden_states)
1024 |
1025 | hidden_states = einops.rearrange(hidden_states, 'b (t h w) (c pt ph pw) -> b c (t pt) (h ph) (w pw)',
1026 | t=post_patch_num_frames, h=post_patch_height, w=post_patch_width,
1027 | pt=p_t, ph=p, pw=p)
1028 |
1029 | if return_dict:
1030 | return Transformer2DModelOutput(sample=hidden_states)
1031 |
1032 | return hidden_states,
1033 |
--------------------------------------------------------------------------------
/diffusers_helper/pipelines/k_diffusion_hunyuan.py:
--------------------------------------------------------------------------------
1 | import torch
2 | import math
3 |
4 | from diffusers_helper.k_diffusion.uni_pc_fm import sample_unipc
5 | from diffusers_helper.k_diffusion.wrapper import fm_wrapper
6 | from diffusers_helper.utils import repeat_to_batch_size
7 |
8 |
9 | def flux_time_shift(t, mu=1.15, sigma=1.0):
10 | return math.exp(mu) / (math.exp(mu) + (1 / t - 1) ** sigma)
11 |
12 |
13 | def calculate_flux_mu(context_length, x1=256, y1=0.5, x2=4096, y2=1.15, exp_max=7.0):
14 | k = (y2 - y1) / (x2 - x1)
15 | b = y1 - k * x1
16 | mu = k * context_length + b
17 | mu = min(mu, math.log(exp_max))
18 | return mu
19 |
20 |
21 | def get_flux_sigmas_from_mu(n, mu):
22 | sigmas = torch.linspace(1, 0, steps=n + 1)
23 | sigmas = flux_time_shift(sigmas, mu=mu)
24 | return sigmas
25 |
26 |
27 | @torch.inference_mode()
28 | def sample_hunyuan(
29 | transformer,
30 | sampler='unipc',
31 | initial_latent=None,
32 | concat_latent=None,
33 | strength=1.0,
34 | width=512,
35 | height=512,
36 | frames=16,
37 | real_guidance_scale=1.0,
38 | distilled_guidance_scale=6.0,
39 | guidance_rescale=0.0,
40 | shift=None,
41 | num_inference_steps=25,
42 | batch_size=None,
43 | generator=None,
44 | prompt_embeds=None,
45 | prompt_embeds_mask=None,
46 | prompt_poolers=None,
47 | negative_prompt_embeds=None,
48 | negative_prompt_embeds_mask=None,
49 | negative_prompt_poolers=None,
50 | dtype=torch.bfloat16,
51 | device=None,
52 | negative_kwargs=None,
53 | callback=None,
54 | **kwargs,
55 | ):
56 | device = device or transformer.device
57 |
58 | if batch_size is None:
59 | batch_size = int(prompt_embeds.shape[0])
60 |
61 | latents = torch.randn((batch_size, 16, (frames + 3) // 4, height // 8, width // 8), generator=generator, device=generator.device).to(device=device, dtype=torch.float32)
62 |
63 | B, C, T, H, W = latents.shape
64 | seq_length = T * H * W // 4
65 |
66 | if shift is None:
67 | mu = calculate_flux_mu(seq_length, exp_max=7.0)
68 | else:
69 | mu = math.log(shift)
70 |
71 | sigmas = get_flux_sigmas_from_mu(num_inference_steps, mu).to(device)
72 |
73 | k_model = fm_wrapper(transformer)
74 |
75 | if initial_latent is not None:
76 | sigmas = sigmas * strength
77 | first_sigma = sigmas[0].to(device=device, dtype=torch.float32)
78 | initial_latent = initial_latent.to(device=device, dtype=torch.float32)
79 | latents = initial_latent.float() * (1.0 - first_sigma) + latents.float() * first_sigma
80 |
81 | if concat_latent is not None:
82 | concat_latent = concat_latent.to(latents)
83 |
84 | distilled_guidance = torch.tensor([distilled_guidance_scale * 1000.0] * batch_size).to(device=device, dtype=dtype)
85 |
86 | prompt_embeds = repeat_to_batch_size(prompt_embeds, batch_size)
87 | prompt_embeds_mask = repeat_to_batch_size(prompt_embeds_mask, batch_size)
88 | prompt_poolers = repeat_to_batch_size(prompt_poolers, batch_size)
89 | negative_prompt_embeds = repeat_to_batch_size(negative_prompt_embeds, batch_size)
90 | negative_prompt_embeds_mask = repeat_to_batch_size(negative_prompt_embeds_mask, batch_size)
91 | negative_prompt_poolers = repeat_to_batch_size(negative_prompt_poolers, batch_size)
92 | concat_latent = repeat_to_batch_size(concat_latent, batch_size)
93 |
94 | sampler_kwargs = dict(
95 | dtype=dtype,
96 | cfg_scale=real_guidance_scale,
97 | cfg_rescale=guidance_rescale,
98 | concat_latent=concat_latent,
99 | positive=dict(
100 | pooled_projections=prompt_poolers,
101 | encoder_hidden_states=prompt_embeds,
102 | encoder_attention_mask=prompt_embeds_mask,
103 | guidance=distilled_guidance,
104 | **kwargs,
105 | ),
106 | negative=dict(
107 | pooled_projections=negative_prompt_poolers,
108 | encoder_hidden_states=negative_prompt_embeds,
109 | encoder_attention_mask=negative_prompt_embeds_mask,
110 | guidance=distilled_guidance,
111 | **(kwargs if negative_kwargs is None else {**kwargs, **negative_kwargs}),
112 | )
113 | )
114 |
115 | if sampler == 'unipc':
116 | results = sample_unipc(k_model, latents, sigmas, extra_args=sampler_kwargs, disable=False, callback=callback)
117 | else:
118 | raise NotImplementedError(f'Sampler {sampler} is not supported.')
119 |
120 | return results
121 |
--------------------------------------------------------------------------------
/diffusers_helper/thread_utils.py:
--------------------------------------------------------------------------------
1 | import time
2 |
3 | from threading import Thread, Lock
4 |
5 |
6 | class Listener:
7 | task_queue = []
8 | lock = Lock()
9 | thread = None
10 |
11 | @classmethod
12 | def _process_tasks(cls):
13 | while True:
14 | task = None
15 | with cls.lock:
16 | if cls.task_queue:
17 | task = cls.task_queue.pop(0)
18 |
19 | if task is None:
20 | time.sleep(0.001)
21 | continue
22 |
23 | func, args, kwargs = task
24 | try:
25 | func(*args, **kwargs)
26 | except Exception as e:
27 | print(f"Error in listener thread: {e}")
28 |
29 | @classmethod
30 | def add_task(cls, func, *args, **kwargs):
31 | with cls.lock:
32 | cls.task_queue.append((func, args, kwargs))
33 |
34 | if cls.thread is None:
35 | cls.thread = Thread(target=cls._process_tasks, daemon=True)
36 | cls.thread.start()
37 |
38 |
39 | def async_run(func, *args, **kwargs):
40 | Listener.add_task(func, *args, **kwargs)
41 |
42 |
43 | class FIFOQueue:
44 | def __init__(self):
45 | self.queue = []
46 | self.lock = Lock()
47 |
48 | def push(self, item):
49 | with self.lock:
50 | self.queue.append(item)
51 |
52 | def pop(self):
53 | with self.lock:
54 | if self.queue:
55 | return self.queue.pop(0)
56 | return None
57 |
58 | def top(self):
59 | with self.lock:
60 | if self.queue:
61 | return self.queue[0]
62 | return None
63 |
64 | def next(self):
65 | while True:
66 | with self.lock:
67 | if self.queue:
68 | return self.queue.pop(0)
69 |
70 | time.sleep(0.001)
71 |
72 |
73 | class AsyncStream:
74 | def __init__(self):
75 | self.input_queue = FIFOQueue()
76 | self.output_queue = FIFOQueue()
77 |
--------------------------------------------------------------------------------
/diffusers_helper/utils.py:
--------------------------------------------------------------------------------
1 | import os
2 | import cv2
3 | import json
4 | import random
5 | import glob
6 | import torch
7 | import einops
8 | import numpy as np
9 | import datetime
10 | import torchvision
11 |
12 | import safetensors.torch as sf
13 | from PIL import Image
14 |
15 |
16 | def min_resize(x, m):
17 | if x.shape[0] < x.shape[1]:
18 | s0 = m
19 | s1 = int(float(m) / float(x.shape[0]) * float(x.shape[1]))
20 | else:
21 | s0 = int(float(m) / float(x.shape[1]) * float(x.shape[0]))
22 | s1 = m
23 | new_max = max(s1, s0)
24 | raw_max = max(x.shape[0], x.shape[1])
25 | if new_max < raw_max:
26 | interpolation = cv2.INTER_AREA
27 | else:
28 | interpolation = cv2.INTER_LANCZOS4
29 | y = cv2.resize(x, (s1, s0), interpolation=interpolation)
30 | return y
31 |
32 |
33 | def d_resize(x, y):
34 | H, W, C = y.shape
35 | new_min = min(H, W)
36 | raw_min = min(x.shape[0], x.shape[1])
37 | if new_min < raw_min:
38 | interpolation = cv2.INTER_AREA
39 | else:
40 | interpolation = cv2.INTER_LANCZOS4
41 | y = cv2.resize(x, (W, H), interpolation=interpolation)
42 | return y
43 |
44 |
45 | def resize_and_center_crop(image, target_width, target_height):
46 | if target_height == image.shape[0] and target_width == image.shape[1]:
47 | return image
48 |
49 | pil_image = Image.fromarray(image)
50 | original_width, original_height = pil_image.size
51 | scale_factor = max(target_width / original_width, target_height / original_height)
52 | resized_width = int(round(original_width * scale_factor))
53 | resized_height = int(round(original_height * scale_factor))
54 | resized_image = pil_image.resize((resized_width, resized_height), Image.LANCZOS)
55 | left = (resized_width - target_width) / 2
56 | top = (resized_height - target_height) / 2
57 | right = (resized_width + target_width) / 2
58 | bottom = (resized_height + target_height) / 2
59 | cropped_image = resized_image.crop((left, top, right, bottom))
60 | return np.array(cropped_image)
61 |
62 |
63 | def resize_and_center_crop_pytorch(image, target_width, target_height):
64 | B, C, H, W = image.shape
65 |
66 | if H == target_height and W == target_width:
67 | return image
68 |
69 | scale_factor = max(target_width / W, target_height / H)
70 | resized_width = int(round(W * scale_factor))
71 | resized_height = int(round(H * scale_factor))
72 |
73 | resized = torch.nn.functional.interpolate(image, size=(resized_height, resized_width), mode='bilinear', align_corners=False)
74 |
75 | top = (resized_height - target_height) // 2
76 | left = (resized_width - target_width) // 2
77 | cropped = resized[:, :, top:top + target_height, left:left + target_width]
78 |
79 | return cropped
80 |
81 |
82 | def resize_without_crop(image, target_width, target_height):
83 | if target_height == image.shape[0] and target_width == image.shape[1]:
84 | return image
85 |
86 | pil_image = Image.fromarray(image)
87 | resized_image = pil_image.resize((target_width, target_height), Image.LANCZOS)
88 | return np.array(resized_image)
89 |
90 |
91 | def just_crop(image, w, h):
92 | if h == image.shape[0] and w == image.shape[1]:
93 | return image
94 |
95 | original_height, original_width = image.shape[:2]
96 | k = min(original_height / h, original_width / w)
97 | new_width = int(round(w * k))
98 | new_height = int(round(h * k))
99 | x_start = (original_width - new_width) // 2
100 | y_start = (original_height - new_height) // 2
101 | cropped_image = image[y_start:y_start + new_height, x_start:x_start + new_width]
102 | return cropped_image
103 |
104 |
105 | def write_to_json(data, file_path):
106 | temp_file_path = file_path + ".tmp"
107 | with open(temp_file_path, 'wt', encoding='utf-8') as temp_file:
108 | json.dump(data, temp_file, indent=4)
109 | os.replace(temp_file_path, file_path)
110 | return
111 |
112 |
113 | def read_from_json(file_path):
114 | with open(file_path, 'rt', encoding='utf-8') as file:
115 | data = json.load(file)
116 | return data
117 |
118 |
119 | def get_active_parameters(m):
120 | return {k: v for k, v in m.named_parameters() if v.requires_grad}
121 |
122 |
123 | def cast_training_params(m, dtype=torch.float32):
124 | result = {}
125 | for n, param in m.named_parameters():
126 | if param.requires_grad:
127 | param.data = param.to(dtype)
128 | result[n] = param
129 | return result
130 |
131 |
132 | def separate_lora_AB(parameters, B_patterns=None):
133 | parameters_normal = {}
134 | parameters_B = {}
135 |
136 | if B_patterns is None:
137 | B_patterns = ['.lora_B.', '__zero__']
138 |
139 | for k, v in parameters.items():
140 | if any(B_pattern in k for B_pattern in B_patterns):
141 | parameters_B[k] = v
142 | else:
143 | parameters_normal[k] = v
144 |
145 | return parameters_normal, parameters_B
146 |
147 |
148 | def set_attr_recursive(obj, attr, value):
149 | attrs = attr.split(".")
150 | for name in attrs[:-1]:
151 | obj = getattr(obj, name)
152 | setattr(obj, attrs[-1], value)
153 | return
154 |
155 |
156 | def print_tensor_list_size(tensors):
157 | total_size = 0
158 | total_elements = 0
159 |
160 | if isinstance(tensors, dict):
161 | tensors = tensors.values()
162 |
163 | for tensor in tensors:
164 | total_size += tensor.nelement() * tensor.element_size()
165 | total_elements += tensor.nelement()
166 |
167 | total_size_MB = total_size / (1024 ** 2)
168 | total_elements_B = total_elements / 1e9
169 |
170 | print(f"Total number of tensors: {len(tensors)}")
171 | print(f"Total size of tensors: {total_size_MB:.2f} MB")
172 | print(f"Total number of parameters: {total_elements_B:.3f} billion")
173 | return
174 |
175 |
176 | @torch.no_grad()
177 | def batch_mixture(a, b=None, probability_a=0.5, mask_a=None):
178 | batch_size = a.size(0)
179 |
180 | if b is None:
181 | b = torch.zeros_like(a)
182 |
183 | if mask_a is None:
184 | mask_a = torch.rand(batch_size) < probability_a
185 |
186 | mask_a = mask_a.to(a.device)
187 | mask_a = mask_a.reshape((batch_size,) + (1,) * (a.dim() - 1))
188 | result = torch.where(mask_a, a, b)
189 | return result
190 |
191 |
192 | @torch.no_grad()
193 | def zero_module(module):
194 | for p in module.parameters():
195 | p.detach().zero_()
196 | return module
197 |
198 |
199 | @torch.no_grad()
200 | def supress_lower_channels(m, k, alpha=0.01):
201 | data = m.weight.data.clone()
202 |
203 | assert int(data.shape[1]) >= k
204 |
205 | data[:, :k] = data[:, :k] * alpha
206 | m.weight.data = data.contiguous().clone()
207 | return m
208 |
209 |
210 | def freeze_module(m):
211 | if not hasattr(m, '_forward_inside_frozen_module'):
212 | m._forward_inside_frozen_module = m.forward
213 | m.requires_grad_(False)
214 | m.forward = torch.no_grad()(m.forward)
215 | return m
216 |
217 |
218 | def get_latest_safetensors(folder_path):
219 | safetensors_files = glob.glob(os.path.join(folder_path, '*.safetensors'))
220 |
221 | if not safetensors_files:
222 | raise ValueError('No file to resume!')
223 |
224 | latest_file = max(safetensors_files, key=os.path.getmtime)
225 | latest_file = os.path.abspath(os.path.realpath(latest_file))
226 | return latest_file
227 |
228 |
229 | def generate_random_prompt_from_tags(tags_str, min_length=3, max_length=32):
230 | tags = tags_str.split(', ')
231 | tags = random.sample(tags, k=min(random.randint(min_length, max_length), len(tags)))
232 | prompt = ', '.join(tags)
233 | return prompt
234 |
235 |
236 | def interpolate_numbers(a, b, n, round_to_int=False, gamma=1.0):
237 | numbers = a + (b - a) * (np.linspace(0, 1, n) ** gamma)
238 | if round_to_int:
239 | numbers = np.round(numbers).astype(int)
240 | return numbers.tolist()
241 |
242 |
243 | def uniform_random_by_intervals(inclusive, exclusive, n, round_to_int=False):
244 | edges = np.linspace(0, 1, n + 1)
245 | points = np.random.uniform(edges[:-1], edges[1:])
246 | numbers = inclusive + (exclusive - inclusive) * points
247 | if round_to_int:
248 | numbers = np.round(numbers).astype(int)
249 | return numbers.tolist()
250 |
251 |
252 | def soft_append_bcthw(history, current, overlap=0):
253 | if overlap <= 0:
254 | return torch.cat([history, current], dim=2)
255 |
256 | assert history.shape[2] >= overlap, f"History length ({history.shape[2]}) must be >= overlap ({overlap})"
257 | assert current.shape[2] >= overlap, f"Current length ({current.shape[2]}) must be >= overlap ({overlap})"
258 |
259 | weights = torch.linspace(1, 0, overlap, dtype=history.dtype, device=history.device).view(1, 1, -1, 1, 1)
260 | blended = weights * history[:, :, -overlap:] + (1 - weights) * current[:, :, :overlap]
261 | output = torch.cat([history[:, :, :-overlap], blended, current[:, :, overlap:]], dim=2)
262 |
263 | return output.to(history)
264 |
265 |
266 | def save_bcthw_as_mp4(x, output_filename, fps=10, crf=0):
267 | b, c, t, h, w = x.shape
268 |
269 | per_row = b
270 | for p in [6, 5, 4, 3, 2]:
271 | if b % p == 0:
272 | per_row = p
273 | break
274 |
275 | os.makedirs(os.path.dirname(os.path.abspath(os.path.realpath(output_filename))), exist_ok=True)
276 | x = torch.clamp(x.float(), -1., 1.) * 127.5 + 127.5
277 | x = x.detach().cpu().to(torch.uint8)
278 | x = einops.rearrange(x, '(m n) c t h w -> t (m h) (n w) c', n=per_row)
279 | torchvision.io.write_video(output_filename, x, fps=fps, video_codec='libx264', options={'crf': str(int(crf))})
280 | return x
281 |
282 |
283 | def save_bcthw_as_png(x, output_filename):
284 | os.makedirs(os.path.dirname(os.path.abspath(os.path.realpath(output_filename))), exist_ok=True)
285 | x = torch.clamp(x.float(), -1., 1.) * 127.5 + 127.5
286 | x = x.detach().cpu().to(torch.uint8)
287 | x = einops.rearrange(x, 'b c t h w -> c (b h) (t w)')
288 | torchvision.io.write_png(x, output_filename)
289 | return output_filename
290 |
291 |
292 | def save_bchw_as_png(x, output_filename):
293 | os.makedirs(os.path.dirname(os.path.abspath(os.path.realpath(output_filename))), exist_ok=True)
294 | x = torch.clamp(x.float(), -1., 1.) * 127.5 + 127.5
295 | x = x.detach().cpu().to(torch.uint8)
296 | x = einops.rearrange(x, 'b c h w -> c h (b w)')
297 | torchvision.io.write_png(x, output_filename)
298 | return output_filename
299 |
300 |
301 | def add_tensors_with_padding(tensor1, tensor2):
302 | if tensor1.shape == tensor2.shape:
303 | return tensor1 + tensor2
304 |
305 | shape1 = tensor1.shape
306 | shape2 = tensor2.shape
307 |
308 | new_shape = tuple(max(s1, s2) for s1, s2 in zip(shape1, shape2))
309 |
310 | padded_tensor1 = torch.zeros(new_shape)
311 | padded_tensor2 = torch.zeros(new_shape)
312 |
313 | padded_tensor1[tuple(slice(0, s) for s in shape1)] = tensor1
314 | padded_tensor2[tuple(slice(0, s) for s in shape2)] = tensor2
315 |
316 | result = padded_tensor1 + padded_tensor2
317 | return result
318 |
319 |
320 | def print_free_mem():
321 | torch.cuda.empty_cache()
322 | free_mem, total_mem = torch.cuda.mem_get_info(0)
323 | free_mem_mb = free_mem / (1024 ** 2)
324 | total_mem_mb = total_mem / (1024 ** 2)
325 | print(f"Free memory: {free_mem_mb:.2f} MB")
326 | print(f"Total memory: {total_mem_mb:.2f} MB")
327 | return
328 |
329 |
330 | def print_gpu_parameters(device, state_dict, log_count=1):
331 | summary = {"device": device, "keys_count": len(state_dict)}
332 |
333 | logged_params = {}
334 | for i, (key, tensor) in enumerate(state_dict.items()):
335 | if i >= log_count:
336 | break
337 | logged_params[key] = tensor.flatten()[:3].tolist()
338 |
339 | summary["params"] = logged_params
340 |
341 | print(str(summary))
342 | return
343 |
344 |
345 | def visualize_txt_as_img(width, height, text, font_path='font/DejaVuSans.ttf', size=18):
346 | from PIL import Image, ImageDraw, ImageFont
347 |
348 | txt = Image.new("RGB", (width, height), color="white")
349 | draw = ImageDraw.Draw(txt)
350 | font = ImageFont.truetype(font_path, size=size)
351 |
352 | if text == '':
353 | return np.array(txt)
354 |
355 | # Split text into lines that fit within the image width
356 | lines = []
357 | words = text.split()
358 | current_line = words[0]
359 |
360 | for word in words[1:]:
361 | line_with_word = f"{current_line} {word}"
362 | if draw.textbbox((0, 0), line_with_word, font=font)[2] <= width:
363 | current_line = line_with_word
364 | else:
365 | lines.append(current_line)
366 | current_line = word
367 |
368 | lines.append(current_line)
369 |
370 | # Draw the text line by line
371 | y = 0
372 | line_height = draw.textbbox((0, 0), "A", font=font)[3]
373 |
374 | for line in lines:
375 | if y + line_height > height:
376 | break # stop drawing if the next line will be outside the image
377 | draw.text((0, y), line, fill="black", font=font)
378 | y += line_height
379 |
380 | return np.array(txt)
381 |
382 |
383 | def blue_mark(x):
384 | x = x.copy()
385 | c = x[:, :, 2]
386 | b = cv2.blur(c, (9, 9))
387 | x[:, :, 2] = ((c - b) * 16.0 + b).clip(-1, 1)
388 | return x
389 |
390 |
391 | def green_mark(x):
392 | x = x.copy()
393 | x[:, :, 2] = -1
394 | x[:, :, 0] = -1
395 | return x
396 |
397 |
398 | def frame_mark(x):
399 | x = x.copy()
400 | x[:64] = -1
401 | x[-64:] = -1
402 | x[:, :8] = 1
403 | x[:, -8:] = 1
404 | return x
405 |
406 |
407 | @torch.inference_mode()
408 | def pytorch2numpy(imgs):
409 | results = []
410 | for x in imgs:
411 | y = x.movedim(0, -1)
412 | y = y * 127.5 + 127.5
413 | y = y.detach().float().cpu().numpy().clip(0, 255).astype(np.uint8)
414 | results.append(y)
415 | return results
416 |
417 |
418 | @torch.inference_mode()
419 | def numpy2pytorch(imgs):
420 | h = torch.from_numpy(np.stack(imgs, axis=0)).float() / 127.5 - 1.0
421 | h = h.movedim(-1, 1)
422 | return h
423 |
424 |
425 | @torch.no_grad()
426 | def duplicate_prefix_to_suffix(x, count, zero_out=False):
427 | if zero_out:
428 | return torch.cat([x, torch.zeros_like(x[:count])], dim=0)
429 | else:
430 | return torch.cat([x, x[:count]], dim=0)
431 |
432 |
433 | def weighted_mse(a, b, weight):
434 | return torch.mean(weight.float() * (a.float() - b.float()) ** 2)
435 |
436 |
437 | def clamped_linear_interpolation(x, x_min, y_min, x_max, y_max, sigma=1.0):
438 | x = (x - x_min) / (x_max - x_min)
439 | x = max(0.0, min(x, 1.0))
440 | x = x ** sigma
441 | return y_min + x * (y_max - y_min)
442 |
443 |
444 | def expand_to_dims(x, target_dims):
445 | return x.view(*x.shape, *([1] * max(0, target_dims - x.dim())))
446 |
447 |
448 | def repeat_to_batch_size(tensor: torch.Tensor, batch_size: int):
449 | if tensor is None:
450 | return None
451 |
452 | first_dim = tensor.shape[0]
453 |
454 | if first_dim == batch_size:
455 | return tensor
456 |
457 | if batch_size % first_dim != 0:
458 | raise ValueError(f"Cannot evenly repeat first dim {first_dim} to match batch_size {batch_size}.")
459 |
460 | repeat_times = batch_size // first_dim
461 |
462 | return tensor.repeat(repeat_times, *[1] * (tensor.dim() - 1))
463 |
464 |
465 | def dim5(x):
466 | return expand_to_dims(x, 5)
467 |
468 |
469 | def dim4(x):
470 | return expand_to_dims(x, 4)
471 |
472 |
473 | def dim3(x):
474 | return expand_to_dims(x, 3)
475 |
476 |
477 | def crop_or_pad_yield_mask(x, length):
478 | B, F, C = x.shape
479 | device = x.device
480 | dtype = x.dtype
481 |
482 | if F < length:
483 | y = torch.zeros((B, length, C), dtype=dtype, device=device)
484 | mask = torch.zeros((B, length), dtype=torch.bool, device=device)
485 | y[:, :F, :] = x
486 | mask[:, :F] = True
487 | return y, mask
488 |
489 | return x[:, :length, :], torch.ones((B, length), dtype=torch.bool, device=device)
490 |
491 |
492 | def extend_dim(x, dim, minimal_length, zero_pad=False):
493 | original_length = int(x.shape[dim])
494 |
495 | if original_length >= minimal_length:
496 | return x
497 |
498 | if zero_pad:
499 | padding_shape = list(x.shape)
500 | padding_shape[dim] = minimal_length - original_length
501 | padding = torch.zeros(padding_shape, dtype=x.dtype, device=x.device)
502 | else:
503 | idx = (slice(None),) * dim + (slice(-1, None),) + (slice(None),) * (len(x.shape) - dim - 1)
504 | last_element = x[idx]
505 | padding = last_element.repeat_interleave(minimal_length - original_length, dim=dim)
506 |
507 | return torch.cat([x, padding], dim=dim)
508 |
509 |
510 | def lazy_positional_encoding(t, repeats=None):
511 | if not isinstance(t, list):
512 | t = [t]
513 |
514 | from diffusers.models.embeddings import get_timestep_embedding
515 |
516 | te = torch.tensor(t)
517 | te = get_timestep_embedding(timesteps=te, embedding_dim=256, flip_sin_to_cos=True, downscale_freq_shift=0.0, scale=1.0)
518 |
519 | if repeats is None:
520 | return te
521 |
522 | te = te[:, None, :].expand(-1, repeats, -1)
523 |
524 | return te
525 |
526 |
527 | def state_dict_offset_merge(A, B, C=None):
528 | result = {}
529 | keys = A.keys()
530 |
531 | for key in keys:
532 | A_value = A[key]
533 | B_value = B[key].to(A_value)
534 |
535 | if C is None:
536 | result[key] = A_value + B_value
537 | else:
538 | C_value = C[key].to(A_value)
539 | result[key] = A_value + B_value - C_value
540 |
541 | return result
542 |
543 |
544 | def state_dict_weighted_merge(state_dicts, weights):
545 | if len(state_dicts) != len(weights):
546 | raise ValueError("Number of state dictionaries must match number of weights")
547 |
548 | if not state_dicts:
549 | return {}
550 |
551 | total_weight = sum(weights)
552 |
553 | if total_weight == 0:
554 | raise ValueError("Sum of weights cannot be zero")
555 |
556 | normalized_weights = [w / total_weight for w in weights]
557 |
558 | keys = state_dicts[0].keys()
559 | result = {}
560 |
561 | for key in keys:
562 | result[key] = state_dicts[0][key] * normalized_weights[0]
563 |
564 | for i in range(1, len(state_dicts)):
565 | state_dict_value = state_dicts[i][key].to(result[key])
566 | result[key] += state_dict_value * normalized_weights[i]
567 |
568 | return result
569 |
570 |
571 | def group_files_by_folder(all_files):
572 | grouped_files = {}
573 |
574 | for file in all_files:
575 | folder_name = os.path.basename(os.path.dirname(file))
576 | if folder_name not in grouped_files:
577 | grouped_files[folder_name] = []
578 | grouped_files[folder_name].append(file)
579 |
580 | list_of_lists = list(grouped_files.values())
581 | return list_of_lists
582 |
583 |
584 | def generate_timestamp():
585 | now = datetime.datetime.now()
586 | timestamp = now.strftime('%y%m%d_%H%M%S')
587 | milliseconds = f"{int(now.microsecond / 1000):03d}"
588 | random_number = random.randint(0, 9999)
589 | return f"{timestamp}_{milliseconds}_{random_number}"
590 |
591 |
592 | def write_PIL_image_with_png_info(image, metadata, path):
593 | from PIL.PngImagePlugin import PngInfo
594 |
595 | png_info = PngInfo()
596 | for key, value in metadata.items():
597 | png_info.add_text(key, value)
598 |
599 | image.save(path, "PNG", pnginfo=png_info)
600 | return image
601 |
602 |
603 | def torch_safe_save(content, path):
604 | torch.save(content, path + '_tmp')
605 | os.replace(path + '_tmp', path)
606 | return path
607 |
608 |
609 | def move_optimizer_to_device(optimizer, device):
610 | for state in optimizer.state.values():
611 | for k, v in state.items():
612 | if isinstance(v, torch.Tensor):
613 | state[k] = v.to(device)
614 |
--------------------------------------------------------------------------------
/requirements.txt:
--------------------------------------------------------------------------------
1 | accelerate==1.6.0
2 | diffusers==0.33.1
3 | transformers==4.46.2
4 | gradio==5.23.0
5 | sentencepiece==0.2.0
6 | pillow==11.1.0
7 | av==12.1.0
8 | numpy==1.26.2
9 | scipy==1.12.0
10 | requests==2.31.0
11 | torchsde==0.2.6
12 |
13 | einops
14 | opencv-contrib-python
15 | safetensors
16 |
--------------------------------------------------------------------------------