Repository: lllyasviel/FramePack
Branch: main
Commit: 97fe5dbe06ac
Files: 19
Total size: 150.0 KB
Directory structure:
gitextract_wnt3hu2x/
├── .gitignore
├── LICENSE
├── README.md
├── demo_gradio.py
├── demo_gradio_f1.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
================================================
FILE CONTENTS
================================================
================================================
FILE: .gitignore
================================================
hf_download/
outputs/
repo/
# Byte-compiled / optimized / DLL files
__pycache__/
*.py[cod]
*$py.class
# C extensions
*.so
# Distribution / packaging
.Python
build/
develop-eggs/
dist/
downloads/
eggs/
.eggs/
lib/
lib64/
parts/
sdist/
var/
wheels/
share/python-wheels/
*.egg-info/
.installed.cfg
*.egg
MANIFEST
# PyInstaller
# Usually these files are written by a python script from a template
# before PyInstaller builds the exe, so as to inject date/other infos into it.
*.manifest
*.spec
# Installer logs
pip-log.txt
pip-delete-this-directory.txt
# Unit test / coverage reports
htmlcov/
.tox/
.nox/
.coverage
.coverage.*
.cache
nosetests.xml
coverage.xml
*.cover
*.py,cover
.hypothesis/
.pytest_cache/
cover/
# Translations
*.mo
*.pot
# Django stuff:
*.log
local_settings.py
db.sqlite3
db.sqlite3-journal
# Flask stuff:
instance/
.webassets-cache
# Scrapy stuff:
.scrapy
# Sphinx documentation
docs/_build/
# PyBuilder
.pybuilder/
target/
# Jupyter Notebook
.ipynb_checkpoints
# IPython
profile_default/
ipython_config.py
# pyenv
# For a library or package, you might want to ignore these files since the code is
# intended to run in multiple environments; otherwise, check them in:
# .python-version
# pipenv
# According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control.
# However, in case of collaboration, if having platform-specific dependencies or dependencies
# having no cross-platform support, pipenv may install dependencies that don't work, or not
# install all needed dependencies.
#Pipfile.lock
# UV
# Similar to Pipfile.lock, it is generally recommended to include uv.lock in version control.
# This is especially recommended for binary packages to ensure reproducibility, and is more
# commonly ignored for libraries.
#uv.lock
# poetry
# Similar to Pipfile.lock, it is generally recommended to include poetry.lock in version control.
# This is especially recommended for binary packages to ensure reproducibility, and is more
# commonly ignored for libraries.
# https://python-poetry.org/docs/basic-usage/#commit-your-poetrylock-file-to-version-control
#poetry.lock
# pdm
# Similar to Pipfile.lock, it is generally recommended to include pdm.lock in version control.
#pdm.lock
# pdm stores project-wide configurations in .pdm.toml, but it is recommended to not include it
# in version control.
# https://pdm.fming.dev/latest/usage/project/#working-with-version-control
.pdm.toml
.pdm-python
.pdm-build/
# PEP 582; used by e.g. github.com/David-OConnor/pyflow and github.com/pdm-project/pdm
__pypackages__/
# Celery stuff
celerybeat-schedule
celerybeat.pid
# SageMath parsed files
*.sage.py
# Environments
.env
.venv
env/
venv/
ENV/
env.bak/
venv.bak/
# Spyder project settings
.spyderproject
.spyproject
# Rope project settings
.ropeproject
# mkdocs documentation
/site
# mypy
.mypy_cache/
.dmypy.json
dmypy.json
# Pyre type checker
.pyre/
# pytype static type analyzer
.pytype/
# Cython debug symbols
cython_debug/
# PyCharm
# JetBrains specific template is maintained in a separate JetBrains.gitignore that can
# be found at https://github.com/github/gitignore/blob/main/Global/JetBrains.gitignore
# and can be added to the global gitignore or merged into this file. For a more nuclear
# option (not recommended) you can uncomment the following to ignore the entire idea folder.
.idea/
# Ruff stuff:
.ruff_cache/
# PyPI configuration file
.pypirc
================================================
FILE: LICENSE
================================================
Apache License
Version 2.0, January 2004
http://www.apache.org/licenses/
TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
1. Definitions.
"License" shall mean the terms and conditions for use, reproduction,
and distribution as defined by Sections 1 through 9 of this document.
"Licensor" shall mean the copyright owner or entity authorized by
the copyright owner that is granting the License.
"Legal Entity" shall mean the union of the acting entity and all
other entities that control, are controlled by, or are under common
control with that entity. For the purposes of this definition,
"control" means (i) the power, direct or indirect, to cause the
direction or management of such entity, whether by contract or
otherwise, or (ii) ownership of fifty percent (50%) or more of the
outstanding shares, or (iii) beneficial ownership of such entity.
"You" (or "Your") shall mean an individual or Legal Entity
exercising permissions granted by this License.
"Source" form shall mean the preferred form for making modifications,
including but not limited to software source code, documentation
source, and configuration files.
"Object" form shall mean any form resulting from mechanical
transformation or translation of a Source form, including but
not limited to compiled object code, generated documentation,
and conversions to other media types.
"Work" shall mean the work of authorship, whether in Source or
Object form, made available under the License, as indicated by a
copyright notice that is included in or attached to the work
(an example is provided in the Appendix below).
"Derivative Works" shall mean any work, whether in Source or Object
form, that is based on (or derived from) the Work and for which the
editorial revisions, annotations, elaborations, or other modifications
represent, as a whole, an original work of authorship. For the purposes
of this License, Derivative Works shall not include works that remain
separable from, or merely link (or bind by name) to the interfaces of,
the Work and Derivative Works thereof.
"Contribution" shall mean any work of authorship, including
the original version of the Work and any modifications or additions
to that Work or Derivative Works thereof, that is intentionally
submitted to Licensor for inclusion in the Work by the copyright owner
or by an individual or Legal Entity authorized to submit on behalf of
the copyright owner. For the purposes of this definition, "submitted"
means any form of electronic, verbal, or written communication sent
to the Licensor or its representatives, including but not limited to
communication on electronic mailing lists, source code control systems,
and issue tracking systems that are managed by, or on behalf of, the
Licensor for the purpose of discussing and improving the Work, but
excluding communication that is conspicuously marked or otherwise
designated in writing by the copyright owner as "Not a Contribution."
"Contributor" shall mean Licensor and any individual or Legal Entity
on behalf of whom a Contribution has been received by Licensor and
subsequently incorporated within the Work.
2. Grant of Copyright License. Subject to the terms and conditions of
this License, each Contributor hereby grants to You a perpetual,
worldwide, non-exclusive, no-charge, royalty-free, irrevocable
copyright license to reproduce, prepare Derivative Works of,
publicly display, publicly perform, sublicense, and distribute the
Work and such Derivative Works in Source or Object form.
3. Grant of Patent License. Subject to the terms and conditions of
this License, each Contributor hereby grants to You a perpetual,
worldwide, non-exclusive, no-charge, royalty-free, irrevocable
(except as stated in this section) patent license to make, have made,
use, offer to sell, sell, import, and otherwise transfer the Work,
where such license applies only to those patent claims licensable
by such Contributor that are necessarily infringed by their
Contribution(s) alone or by combination of their Contribution(s)
with the Work to which such Contribution(s) was submitted. If You
institute patent litigation against any entity (including a
cross-claim or counterclaim in a lawsuit) alleging that the Work
or a Contribution incorporated within the Work constitutes direct
or contributory patent infringement, then any patent licenses
granted to You under this License for that Work shall terminate
as of the date such litigation is filed.
4. Redistribution. You may reproduce and distribute copies of the
Work or Derivative Works thereof in any medium, with or without
modifications, and in Source or Object form, provided that You
meet the following conditions:
(a) You must give any other recipients of the Work or
Derivative Works a copy of this License; and
(b) You must cause any modified files to carry prominent notices
stating that You changed the files; and
(c) You must retain, in the Source form of any Derivative Works
that You distribute, all copyright, patent, trademark, and
attribution notices from the Source form of the Work,
excluding those notices that do not pertain to any part of
the Derivative Works; and
(d) If the Work includes a "NOTICE" text file as part of its
distribution, then any Derivative Works that You distribute must
include a readable copy of the attribution notices contained
within such NOTICE file, excluding those notices that do not
pertain to any part of the Derivative Works, in at least one
of the following places: within a NOTICE text file distributed
as part of the Derivative Works; within the Source form or
documentation, if provided along with the Derivative Works; or,
within a display generated by the Derivative Works, if and
wherever such third-party notices normally appear. The contents
of the NOTICE file are for informational purposes only and
do not modify the License. You may add Your own attribution
notices within Derivative Works that You distribute, alongside
or as an addendum to the NOTICE text from the Work, provided
that such additional attribution notices cannot be construed
as modifying the License.
You may add Your own copyright statement to Your modifications and
may provide additional or different license terms and conditions
for use, reproduction, or distribution of Your modifications, or
for any such Derivative Works as a whole, provided Your use,
reproduction, and distribution of the Work otherwise complies with
the conditions stated in this License.
5. Submission of Contributions. Unless You explicitly state otherwise,
any Contribution intentionally submitted for inclusion in the Work
by You to the Licensor shall be under the terms and conditions of
this License, without any additional terms or conditions.
Notwithstanding the above, nothing herein shall supersede or modify
the terms of any separate license agreement you may have executed
with Licensor regarding such Contributions.
6. Trademarks. This License does not grant permission to use the trade
names, trademarks, service marks, or product names of the Licensor,
except as required for reasonable and customary use in describing the
origin of the Work and reproducing the content of the NOTICE file.
7. Disclaimer of Warranty. Unless required by applicable law or
agreed to in writing, Licensor provides the Work (and each
Contributor provides its Contributions) on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
implied, including, without limitation, any warranties or conditions
of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A
PARTICULAR PURPOSE. You are solely responsible for determining the
appropriateness of using or redistributing the Work and assume any
risks associated with Your exercise of permissions under this License.
8. Limitation of Liability. In no event and under no legal theory,
whether in tort (including negligence), contract, or otherwise,
unless required by applicable law (such as deliberate and grossly
negligent acts) or agreed to in writing, shall any Contributor be
liable to You for damages, including any direct, indirect, special,
incidental, or consequential damages of any character arising as a
result of this License or out of the use or inability to use the
Work (including but not limited to damages for loss of goodwill,
work stoppage, computer failure or malfunction, or any and all
other commercial damages or losses), even if such Contributor
has been advised of the possibility of such damages.
9. Accepting Warranty or Additional Liability. While redistributing
the Work or Derivative Works thereof, You may choose to offer,
and charge a fee for, acceptance of support, warranty, indemnity,
or other liability obligations and/or rights consistent with this
License. However, in accepting such obligations, You may act only
on Your own behalf and on Your sole responsibility, not on behalf
of any other Contributor, and only if You agree to indemnify,
defend, and hold each Contributor harmless for any liability
incurred by, or claims asserted against, such Contributor by reason
of your accepting any such warranty or additional liability.
END OF TERMS AND CONDITIONS
APPENDIX: How to apply the Apache License to your work.
To apply the Apache License to your work, attach the following
boilerplate notice, with the fields enclosed by brackets "[]"
replaced with your own identifying information. (Don't include
the brackets!) The text should be enclosed in the appropriate
comment syntax for the file format. We also recommend that a
file or class name and description of purpose be included on the
same "printed page" as the copyright notice for easier
identification within third-party archives.
Copyright [yyyy] [name of copyright owner]
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
================================================
FILE: README.md
================================================
<p align="center">
<img src="https://github.com/user-attachments/assets/2cc030b4-87e1-40a0-b5bf-1b7d6b62820b" width="300">
</p>
# FramePack
Official implementation and desktop software for ["Frame Context Packing and Drift Prevention in Next-Frame-Prediction Video Diffusion Models"](https://lllyasviel.github.io/frame_pack_gitpage/).
Links: [**Paper**](https://arxiv.org/abs/2504.12626), [**Project Page**](https://lllyasviel.github.io/frame_pack_gitpage/)
FramePack is a next-frame (next-frame-section) prediction neural network structure that generates videos progressively.
FramePack compresses input contexts to a constant length so that the generation workload is invariant to video length.
FramePack can process a very large number of frames with 13B models even on laptop GPUs.
FramePack can be trained with a much larger batch size, similar to the batch size for image diffusion training.
**Video diffusion, but feels like image diffusion.**
# News
**2025 July 14:** Some pure text2video anti-drifting stress-test results of FramePack-P1 are uploaded [here,](https://lllyasviel.github.io/frame_pack_gitpage/p1/#text-to-video-stress-tests) using common prompts without any reference images.
**2025 June 26:** Some results of FramePack-P1 are uploaded [here.](https://lllyasviel.github.io/frame_pack_gitpage/p1) The FramePack-P1 will be the next version of FramePack with two designs: Planned Anti-Drifting and History Discretization.
**2025 May 03:** The FramePack-F1 is released. [Try it here.](https://github.com/lllyasviel/FramePack/discussions/459)
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.**
# Requirements
Note that this repo is a functional desktop software with minimal standalone high-quality sampling system and memory management.
**Start with this repo before you try anything else!**
Requirements:
* Nvidia GPU in RTX 30XX, 40XX, 50XX series that supports fp16 and bf16. The GTX 10XX/20XX are not tested.
* Linux or Windows operating system.
* At least 6GB GPU memory.
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.)
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)
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.
# Installation
**Windows**:
[>>> 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)
After you download, you uncompress, use `update.bat` to update, and use `run.bat` to run.
Note that running `update.bat` is important, otherwise you may be using a previous version with potential bugs unfixed.
Note that the models will be downloaded automatically. You will download more than 30GB from HuggingFace.
**Linux**:
We recommend having an independent Python 3.10.
pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu126
pip install -r requirements.txt
To start the GUI, run:
python demo_gradio.py
Note that it supports `--share`, `--port`, `--server`, and so on.
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.
For example, to install sage-attention (linux):
pip install sageattention==1.0.6
However, you are highly recommended to first try without sage-attention since it will influence results, though the influence is minimal.
# GUI
On the left you upload an image and write a prompt.
On the right are the generated videos and latent previews.
Because this is a next-frame-section prediction model, videos will be generated longer and longer.
You will see the progress bar for each section and the latent preview for the next section.
Note that the initial progress may be slower than later diffusion as the device may need some warmup.
# Sanity Check
Before trying your own inputs, we highly recommend going through the sanity check to find out if any hardware or software went wrong.
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.
## Image-to-5-seconds
Download this image:
<img src="https://github.com/user-attachments/assets/f3bc35cf-656a-4c9c-a83a-bbab24858b09" width="150">
Copy this prompt:
`The man dances energetically, leaping mid-air with fluid arm swings and quick footwork.`
Set like this:
(all default parameters, with teacache turned off)
The result will be:
<table>
<tr>
<td align="center" width="300">
<video
src="https://github.com/user-attachments/assets/bc74f039-2b14-4260-a30b-ceacf611a185"
controls
style="max-width:100%;">
</video>
</td>
</tr>
<tr>
<td align="center">
<em>Video may be compressed by GitHub</em>
</td>
</tr>
</table>
**Important Note:**
Again, this is a next-frame-section prediction model. This means you will generate videos frame-by-frame or section-by-section.
**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.
## Know the influence of TeaCache and Quantization
Download this image:
<img src="https://github.com/user-attachments/assets/42293e30-bdd4-456d-895c-8fedff71be04" width="150">
Copy this prompt:
`The girl dances gracefully, with clear movements, full of charm.`
Set like this:
Turn off teacache:
You will get this:
<table>
<tr>
<td align="center" width="300">
<video
src="https://github.com/user-attachments/assets/04ab527b-6da1-4726-9210-a8853dda5577"
controls
style="max-width:100%;">
</video>
</td>
</tr>
<tr>
<td align="center">
<em>Video may be compressed by GitHub</em>
</td>
</tr>
</table>
Now turn on teacache:
About 30% users will get this (the other 70% will get other random results depending on their hardware):
<table>
<tr>
<td align="center" width="300">
<video
src="https://github.com/user-attachments/assets/149fb486-9ccc-4a48-b1f0-326253051e9b"
controls
style="max-width:100%;">
</video>
</td>
</tr>
<tr>
<td align="center">
<em>A typical worse result.</em>
</td>
</tr>
</table>
So you can see that teacache is not really lossless and sometimes can influence the result a lot.
We recommend using teacache to try ideas and then using the full diffusion process to get high-quality results.
This recommendation also applies to sage-attention, bnb quant, gguf, etc., etc.
## Image-to-1-minute
<img src="https://github.com/user-attachments/assets/820af6ca-3c2e-4bbc-afe8-9a9be1994ff5" width="150">
`The girl dances gracefully, with clear movements, full of charm.`
Set video length to 60 seconds:
If everything is in order you will get some result like this eventually.
60s version:
<table>
<tr>
<td align="center" width="300">
<video
src="https://github.com/user-attachments/assets/c3be4bde-2e33-4fd4-b76d-289a036d3a47"
controls
style="max-width:100%;">
</video>
</td>
</tr>
<tr>
<td align="center">
<em>Video may be compressed by GitHub</em>
</td>
</tr>
</table>
6s version:
<table>
<tr>
<td align="center" width="300">
<video
src="https://github.com/user-attachments/assets/37fe2c33-cb03-41e8-acca-920ab3e34861"
controls
style="max-width:100%;">
</video>
</td>
</tr>
<tr>
<td align="center">
<em>Video may be compressed by GitHub</em>
</td>
</tr>
</table>
# More Examples
Many more examples are in [**Project Page**](https://lllyasviel.github.io/frame_pack_gitpage/).
Below are some more examples that you may be interested in reproducing.
---
<img src="https://github.com/user-attachments/assets/99f4d281-28ad-44f5-8700-aa7a4e5638fa" width="150">
`The girl dances gracefully, with clear movements, full of charm.`
<table>
<tr>
<td align="center" width="300">
<video
src="https://github.com/user-attachments/assets/cebe178a-09ce-4b7a-8f3c-060332f4dab1"
controls
style="max-width:100%;">
</video>
</td>
</tr>
<tr>
<td align="center">
<em>Video may be compressed by GitHub</em>
</td>
</tr>
</table>
---
<img src="https://github.com/user-attachments/assets/853f4f40-2956-472f-aa7a-fa50da03ed92" width="150">
`The girl suddenly took out a sign that said “cute” using right hand`
<table>
<tr>
<td align="center" width="300">
<video
src="https://github.com/user-attachments/assets/116069d2-7499-4f38-ada7-8f85517d1fbb"
controls
style="max-width:100%;">
</video>
</td>
</tr>
<tr>
<td align="center">
<em>Video may be compressed by GitHub</em>
</td>
</tr>
</table>
---
<img src="https://github.com/user-attachments/assets/6d87c53f-81b2-4108-a704-697164ae2e81" width="150">
`The girl skateboarding, repeating the endless spinning and dancing and jumping on a skateboard, with clear movements, full of charm.`
<table>
<tr>
<td align="center" width="300">
<video
src="https://github.com/user-attachments/assets/d9e3534a-eb17-4af2-a8ed-8e692e9993d2"
controls
style="max-width:100%;">
</video>
</td>
</tr>
<tr>
<td align="center">
<em>Video may be compressed by GitHub</em>
</td>
</tr>
</table>
---
<img src="https://github.com/user-attachments/assets/6e95d1a5-9674-4c9a-97a9-ddf704159b79" width="150">
`The girl dances gracefully, with clear movements, full of charm.`
<table>
<tr>
<td align="center" width="300">
<video
src="https://github.com/user-attachments/assets/e1b3279e-e30d-4d32-b55f-2fb1d37c81d2"
controls
style="max-width:100%;">
</video>
</td>
</tr>
<tr>
<td align="center">
<em>Video may be compressed by GitHub</em>
</td>
</tr>
</table>
---
<img src="https://github.com/user-attachments/assets/90fc6d7e-8f6b-4f8c-a5df-ee5b1c8b63c9" width="150">
`The man dances flamboyantly, swinging his hips and striking bold poses with dramatic flair.`
<table>
<tr>
<td align="center" width="300">
<video
src="https://github.com/user-attachments/assets/aaa4481b-7bf8-4c64-bc32-909659767115"
controls
style="max-width:100%;">
</video>
</td>
</tr>
<tr>
<td align="center">
<em>Video may be compressed by GitHub</em>
</td>
</tr>
</table>
---
<img src="https://github.com/user-attachments/assets/62ecf987-ec0c-401d-b3c9-be9ffe84ee5b" width="150">
`The woman dances elegantly among the blossoms, spinning slowly with flowing sleeves and graceful hand movements.`
<table>
<tr>
<td align="center" width="300">
<video
src="https://github.com/user-attachments/assets/f23f2f37-c9b8-45d5-a1be-7c87bd4b41cf"
controls
style="max-width:100%;">
</video>
</td>
</tr>
<tr>
<td align="center">
<em>Video may be compressed by GitHub</em>
</td>
</tr>
</table>
---
<img src="https://github.com/user-attachments/assets/4f740c1a-2d2f-40a6-9613-d6fe64c428aa" width="150">
`The young man writes intensely, flipping papers and adjusting his glasses with swift, focused movements.`
<table>
<tr>
<td align="center" width="300">
<video
src="https://github.com/user-attachments/assets/62e9910e-aea6-4b2b-9333-2e727bccfc64"
controls
style="max-width:100%;">
</video>
</td>
</tr>
<tr>
<td align="center">
<em>Video may be compressed by GitHub</em>
</td>
</tr>
</table>
---
# Prompting Guideline
Many people would ask how to write better prompts.
Below is a ChatGPT template that I personally often use to get prompts:
You are an assistant that writes short, motion-focused prompts for animating images.
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.
Larger and more dynamic motions (like dancing, jumping, running, etc.) are preferred over smaller or more subtle ones (like standing still, sitting, etc.).
Describe subject, then motion, then other things. For example: "The girl dances gracefully, with clear movements, full of charm."
If there is something that can dance (like a man, girl, robot, etc.), then prefer to describe it as dancing.
Stay in a loop: one image in, one motion prompt out. Do not explain, ask questions, or generate multiple options.
You paste the instruct to ChatGPT and then feed it an image to get prompt like this:
*The man dances powerfully, striking sharp poses and gliding smoothly across the reflective floor.*
Usually this will give you a prompt that works well.
You can also write prompts yourself. Concise prompts are usually preferred, for example:
*The girl dances gracefully, with clear movements, full of charm.*
*The man dances powerfully, with clear movements, full of energy.*
and so on.
# Cite
@inproceedings{zhang2025framepack,
title={Frame Context Packing and Drift Prevention in Next-Frame-Prediction Video Diffusion Models},
author={Lvmin Zhang and Shengqu Cai and Muyang Li and Gordon Wetzstein and Maneesh Agrawala},
booktitle={The Thirty-ninth Annual Conference on Neural Information Processing Systems},
year={2025},
}
@article{zhang2025framepackv1,
title={Packing Input Frame Contexts in Next-Frame Prediction Models for Video Generation},
author={Lvmin Zhang and Maneesh Agrawala},
journal={Arxiv},
year={2025}
}
================================================
FILE: demo_gradio.py
================================================
from diffusers_helper.hf_login import login
import os
os.environ['HF_HOME'] = os.path.abspath(os.path.realpath(os.path.join(os.path.dirname(__file__), './hf_download')))
import gradio as gr
import torch
import traceback
import einops
import safetensors.torch as sf
import numpy as np
import argparse
import math
from PIL import Image
from diffusers import AutoencoderKLHunyuanVideo
from transformers import LlamaModel, CLIPTextModel, LlamaTokenizerFast, CLIPTokenizer
from diffusers_helper.hunyuan import encode_prompt_conds, vae_decode, vae_encode, vae_decode_fake
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
from diffusers_helper.models.hunyuan_video_packed import HunyuanVideoTransformer3DModelPacked
from diffusers_helper.pipelines.k_diffusion_hunyuan import sample_hunyuan
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
from diffusers_helper.thread_utils import AsyncStream, async_run
from diffusers_helper.gradio.progress_bar import make_progress_bar_css, make_progress_bar_html
from transformers import SiglipImageProcessor, SiglipVisionModel
from diffusers_helper.clip_vision import hf_clip_vision_encode
from diffusers_helper.bucket_tools import find_nearest_bucket
parser = argparse.ArgumentParser()
parser.add_argument('--share', action='store_true')
parser.add_argument("--server", type=str, default='0.0.0.0')
parser.add_argument("--port", type=int, required=False)
parser.add_argument("--inbrowser", action='store_true')
args = parser.parse_args()
# for win desktop probably use --server 127.0.0.1 --inbrowser
# For linux server probably use --server 127.0.0.1 or do not use any cmd flags
print(args)
free_mem_gb = get_cuda_free_memory_gb(gpu)
high_vram = free_mem_gb > 60
print(f'Free VRAM {free_mem_gb} GB')
print(f'High-VRAM Mode: {high_vram}')
text_encoder = LlamaModel.from_pretrained("hunyuanvideo-community/HunyuanVideo", subfolder='text_encoder', torch_dtype=torch.float16).cpu()
text_encoder_2 = CLIPTextModel.from_pretrained("hunyuanvideo-community/HunyuanVideo", subfolder='text_encoder_2', torch_dtype=torch.float16).cpu()
tokenizer = LlamaTokenizerFast.from_pretrained("hunyuanvideo-community/HunyuanVideo", subfolder='tokenizer')
tokenizer_2 = CLIPTokenizer.from_pretrained("hunyuanvideo-community/HunyuanVideo", subfolder='tokenizer_2')
vae = AutoencoderKLHunyuanVideo.from_pretrained("hunyuanvideo-community/HunyuanVideo", subfolder='vae', torch_dtype=torch.float16).cpu()
feature_extractor = SiglipImageProcessor.from_pretrained("lllyasviel/flux_redux_bfl", subfolder='feature_extractor')
image_encoder = SiglipVisionModel.from_pretrained("lllyasviel/flux_redux_bfl", subfolder='image_encoder', torch_dtype=torch.float16).cpu()
transformer = HunyuanVideoTransformer3DModelPacked.from_pretrained('lllyasviel/FramePackI2V_HY', torch_dtype=torch.bfloat16).cpu()
vae.eval()
text_encoder.eval()
text_encoder_2.eval()
image_encoder.eval()
transformer.eval()
if not high_vram:
vae.enable_slicing()
vae.enable_tiling()
transformer.high_quality_fp32_output_for_inference = True
print('transformer.high_quality_fp32_output_for_inference = True')
transformer.to(dtype=torch.bfloat16)
vae.to(dtype=torch.float16)
image_encoder.to(dtype=torch.float16)
text_encoder.to(dtype=torch.float16)
text_encoder_2.to(dtype=torch.float16)
vae.requires_grad_(False)
text_encoder.requires_grad_(False)
text_encoder_2.requires_grad_(False)
image_encoder.requires_grad_(False)
transformer.requires_grad_(False)
if not high_vram:
# DynamicSwapInstaller is same as huggingface's enable_sequential_offload but 3x faster
DynamicSwapInstaller.install_model(transformer, device=gpu)
DynamicSwapInstaller.install_model(text_encoder, device=gpu)
else:
text_encoder.to(gpu)
text_encoder_2.to(gpu)
image_encoder.to(gpu)
vae.to(gpu)
transformer.to(gpu)
stream = AsyncStream()
outputs_folder = './outputs/'
os.makedirs(outputs_folder, exist_ok=True)
@torch.no_grad()
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):
total_latent_sections = (total_second_length * 30) / (latent_window_size * 4)
total_latent_sections = int(max(round(total_latent_sections), 1))
job_id = generate_timestamp()
stream.output_queue.push(('progress', (None, '', make_progress_bar_html(0, 'Starting ...'))))
try:
# Clean GPU
if not high_vram:
unload_complete_models(
text_encoder, text_encoder_2, image_encoder, vae, transformer
)
# Text encoding
stream.output_queue.push(('progress', (None, '', make_progress_bar_html(0, 'Text encoding ...'))))
if not high_vram:
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.
load_model_as_complete(text_encoder_2, target_device=gpu)
llama_vec, clip_l_pooler = encode_prompt_conds(prompt, text_encoder, text_encoder_2, tokenizer, tokenizer_2)
if cfg == 1:
llama_vec_n, clip_l_pooler_n = torch.zeros_like(llama_vec), torch.zeros_like(clip_l_pooler)
else:
llama_vec_n, clip_l_pooler_n = encode_prompt_conds(n_prompt, text_encoder, text_encoder_2, tokenizer, tokenizer_2)
llama_vec, llama_attention_mask = crop_or_pad_yield_mask(llama_vec, length=512)
llama_vec_n, llama_attention_mask_n = crop_or_pad_yield_mask(llama_vec_n, length=512)
# Processing input image
stream.output_queue.push(('progress', (None, '', make_progress_bar_html(0, 'Image processing ...'))))
H, W, C = input_image.shape
height, width = find_nearest_bucket(H, W, resolution=640)
input_image_np = resize_and_center_crop(input_image, target_width=width, target_height=height)
Image.fromarray(input_image_np).save(os.path.join(outputs_folder, f'{job_id}.png'))
input_image_pt = torch.from_numpy(input_image_np).float() / 127.5 - 1
input_image_pt = input_image_pt.permute(2, 0, 1)[None, :, None]
# VAE encoding
stream.output_queue.push(('progress', (None, '', make_progress_bar_html(0, 'VAE encoding ...'))))
if not high_vram:
load_model_as_complete(vae, target_device=gpu)
start_latent = vae_encode(input_image_pt, vae)
# CLIP Vision
stream.output_queue.push(('progress', (None, '', make_progress_bar_html(0, 'CLIP Vision encoding ...'))))
if not high_vram:
load_model_as_complete(image_encoder, target_device=gpu)
image_encoder_output = hf_clip_vision_encode(input_image_np, feature_extractor, image_encoder)
image_encoder_last_hidden_state = image_encoder_output.last_hidden_state
# Dtype
llama_vec = llama_vec.to(transformer.dtype)
llama_vec_n = llama_vec_n.to(transformer.dtype)
clip_l_pooler = clip_l_pooler.to(transformer.dtype)
clip_l_pooler_n = clip_l_pooler_n.to(transformer.dtype)
image_encoder_last_hidden_state = image_encoder_last_hidden_state.to(transformer.dtype)
# Sampling
stream.output_queue.push(('progress', (None, '', make_progress_bar_html(0, 'Start sampling ...'))))
rnd = torch.Generator("cpu").manual_seed(seed)
num_frames = latent_window_size * 4 - 3
history_latents = torch.zeros(size=(1, 16, 1 + 2 + 16, height // 8, width // 8), dtype=torch.float32).cpu()
history_pixels = None
total_generated_latent_frames = 0
latent_paddings = reversed(range(total_latent_sections))
if total_latent_sections > 4:
# In theory the latent_paddings should follow the above sequence, but it seems that duplicating some
# items looks better than expanding it when total_latent_sections > 4
# One can try to remove below trick and just
# use `latent_paddings = list(reversed(range(total_latent_sections)))` to compare
latent_paddings = [3] + [2] * (total_latent_sections - 3) + [1, 0]
for latent_padding in latent_paddings:
is_last_section = latent_padding == 0
latent_padding_size = latent_padding * latent_window_size
if stream.input_queue.top() == 'end':
stream.output_queue.push(('end', None))
return
print(f'latent_padding_size = {latent_padding_size}, is_last_section = {is_last_section}')
indices = torch.arange(0, sum([1, latent_padding_size, latent_window_size, 1, 2, 16])).unsqueeze(0)
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)
clean_latent_indices = torch.cat([clean_latent_indices_pre, clean_latent_indices_post], dim=1)
clean_latents_pre = start_latent.to(history_latents)
clean_latents_post, clean_latents_2x, clean_latents_4x = history_latents[:, :, :1 + 2 + 16, :, :].split([1, 2, 16], dim=2)
clean_latents = torch.cat([clean_latents_pre, clean_latents_post], dim=2)
if not high_vram:
unload_complete_models()
move_model_to_device_with_memory_preservation(transformer, target_device=gpu, preserved_memory_gb=gpu_memory_preservation)
if use_teacache:
transformer.initialize_teacache(enable_teacache=True, num_steps=steps)
else:
transformer.initialize_teacache(enable_teacache=False)
def callback(d):
preview = d['denoised']
preview = vae_decode_fake(preview)
preview = (preview * 255.0).detach().cpu().numpy().clip(0, 255).astype(np.uint8)
preview = einops.rearrange(preview, 'b c t h w -> (b h) (t w) c')
if stream.input_queue.top() == 'end':
stream.output_queue.push(('end', None))
raise KeyboardInterrupt('User ends the task.')
current_step = d['i'] + 1
percentage = int(100.0 * current_step / steps)
hint = f'Sampling {current_step}/{steps}'
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 ...'
stream.output_queue.push(('progress', (preview, desc, make_progress_bar_html(percentage, hint))))
return
generated_latents = sample_hunyuan(
transformer=transformer,
sampler='unipc',
width=width,
height=height,
frames=num_frames,
real_guidance_scale=cfg,
distilled_guidance_scale=gs,
guidance_rescale=rs,
# shift=3.0,
num_inference_steps=steps,
generator=rnd,
prompt_embeds=llama_vec,
prompt_embeds_mask=llama_attention_mask,
prompt_poolers=clip_l_pooler,
negative_prompt_embeds=llama_vec_n,
negative_prompt_embeds_mask=llama_attention_mask_n,
negative_prompt_poolers=clip_l_pooler_n,
device=gpu,
dtype=torch.bfloat16,
image_embeddings=image_encoder_last_hidden_state,
latent_indices=latent_indices,
clean_latents=clean_latents,
clean_latent_indices=clean_latent_indices,
clean_latents_2x=clean_latents_2x,
clean_latent_2x_indices=clean_latent_2x_indices,
clean_latents_4x=clean_latents_4x,
clean_latent_4x_indices=clean_latent_4x_indices,
callback=callback,
)
if is_last_section:
generated_latents = torch.cat([start_latent.to(generated_latents), generated_latents], dim=2)
total_generated_latent_frames += int(generated_latents.shape[2])
history_latents = torch.cat([generated_latents.to(history_latents), history_latents], dim=2)
if not high_vram:
offload_model_from_device_for_memory_preservation(transformer, target_device=gpu, preserved_memory_gb=8)
load_model_as_complete(vae, target_device=gpu)
real_history_latents = history_latents[:, :, :total_generated_latent_frames, :, :]
if history_pixels is None:
history_pixels = vae_decode(real_history_latents, vae).cpu()
else:
section_latent_frames = (latent_window_size * 2 + 1) if is_last_section else (latent_window_size * 2)
overlapped_frames = latent_window_size * 4 - 3
current_pixels = vae_decode(real_history_latents[:, :, :section_latent_frames], vae).cpu()
history_pixels = soft_append_bcthw(current_pixels, history_pixels, overlapped_frames)
if not high_vram:
unload_complete_models()
output_filename = os.path.join(outputs_folder, f'{job_id}_{total_generated_latent_frames}.mp4')
save_bcthw_as_mp4(history_pixels, output_filename, fps=30, crf=mp4_crf)
print(f'Decoded. Current latent shape {real_history_latents.shape}; pixel shape {history_pixels.shape}')
stream.output_queue.push(('file', output_filename))
if is_last_section:
break
except:
traceback.print_exc()
if not high_vram:
unload_complete_models(
text_encoder, text_encoder_2, image_encoder, vae, transformer
)
stream.output_queue.push(('end', None))
return
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):
global stream
assert input_image is not None, 'No input image!'
yield None, None, '', '', gr.update(interactive=False), gr.update(interactive=True)
stream = AsyncStream()
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)
output_filename = None
while True:
flag, data = stream.output_queue.next()
if flag == 'file':
output_filename = data
yield output_filename, gr.update(), gr.update(), gr.update(), gr.update(interactive=False), gr.update(interactive=True)
if flag == 'progress':
preview, desc, html = data
yield gr.update(), gr.update(visible=True, value=preview), desc, html, gr.update(interactive=False), gr.update(interactive=True)
if flag == 'end':
yield output_filename, gr.update(visible=False), gr.update(), '', gr.update(interactive=True), gr.update(interactive=False)
break
def end_process():
stream.input_queue.push('end')
quick_prompts = [
'The girl dances gracefully, with clear movements, full of charm.',
'A character doing some simple body movements.',
]
quick_prompts = [[x] for x in quick_prompts]
css = make_progress_bar_css()
block = gr.Blocks(css=css).queue()
with block:
gr.Markdown('# FramePack')
with gr.Row():
with gr.Column():
input_image = gr.Image(sources='upload', type="numpy", label="Image", height=320)
prompt = gr.Textbox(label="Prompt", value='')
example_quick_prompts = gr.Dataset(samples=quick_prompts, label='Quick List', samples_per_page=1000, components=[prompt])
example_quick_prompts.click(lambda x: x[0], inputs=[example_quick_prompts], outputs=prompt, show_progress=False, queue=False)
with gr.Row():
start_button = gr.Button(value="Start Generation")
end_button = gr.Button(value="End Generation", interactive=False)
with gr.Group():
use_teacache = gr.Checkbox(label='Use TeaCache', value=True, info='Faster speed, but often makes hands and fingers slightly worse.')
n_prompt = gr.Textbox(label="Negative Prompt", value="", visible=False) # Not used
seed = gr.Number(label="Seed", value=31337, precision=0)
total_second_length = gr.Slider(label="Total Video Length (Seconds)", minimum=1, maximum=120, value=5, step=0.1)
latent_window_size = gr.Slider(label="Latent Window Size", minimum=1, maximum=33, value=9, step=1, visible=False) # Should not change
steps = gr.Slider(label="Steps", minimum=1, maximum=100, value=25, step=1, info='Changing this value is not recommended.')
cfg = gr.Slider(label="CFG Scale", minimum=1.0, maximum=32.0, value=1.0, step=0.01, visible=False) # Should not change
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.')
rs = gr.Slider(label="CFG Re-Scale", minimum=0.0, maximum=1.0, value=0.0, step=0.01, visible=False) # Should not change
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.")
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. ")
with gr.Column():
preview_image = gr.Image(label="Next Latents", height=200, visible=False)
result_video = gr.Video(label="Finished Frames", autoplay=True, show_share_button=False, height=512, loop=True)
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.')
progress_desc = gr.Markdown('', elem_classes='no-generating-animation')
progress_bar = gr.HTML('', elem_classes='no-generating-animation')
gr.HTML('<div style="text-align:center; margin-top:20px;">Share your results and find ideas at the <a href="https://x.com/search?q=framepack&f=live" target="_blank">FramePack Twitter (X) thread</a></div>')
ips = [input_image, prompt, n_prompt, seed, total_second_length, latent_window_size, steps, cfg, gs, rs, gpu_memory_preservation, use_teacache, mp4_crf]
start_button.click(fn=process, inputs=ips, outputs=[result_video, preview_image, progress_desc, progress_bar, start_button, end_button])
end_button.click(fn=end_process)
block.launch(
server_name=args.server,
server_port=args.port,
share=args.share,
inbrowser=args.inbrowser,
)
================================================
FILE: demo_gradio_f1.py
================================================
from diffusers_helper.hf_login import login
import os
os.environ['HF_HOME'] = os.path.abspath(os.path.realpath(os.path.join(os.path.dirname(__file__), './hf_download')))
import gradio as gr
import torch
import traceback
import einops
import safetensors.torch as sf
import numpy as np
import argparse
import math
from PIL import Image
from diffusers import AutoencoderKLHunyuanVideo
from transformers import LlamaModel, CLIPTextModel, LlamaTokenizerFast, CLIPTokenizer
from diffusers_helper.hunyuan import encode_prompt_conds, vae_decode, vae_encode, vae_decode_fake
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
from diffusers_helper.models.hunyuan_video_packed import HunyuanVideoTransformer3DModelPacked
from diffusers_helper.pipelines.k_diffusion_hunyuan import sample_hunyuan
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
from diffusers_helper.thread_utils import AsyncStream, async_run
from diffusers_helper.gradio.progress_bar import make_progress_bar_css, make_progress_bar_html
from transformers import SiglipImageProcessor, SiglipVisionModel
from diffusers_helper.clip_vision import hf_clip_vision_encode
from diffusers_helper.bucket_tools import find_nearest_bucket
parser = argparse.ArgumentParser()
parser.add_argument('--share', action='store_true')
parser.add_argument("--server", type=str, default='0.0.0.0')
parser.add_argument("--port", type=int, required=False)
parser.add_argument("--inbrowser", action='store_true')
args = parser.parse_args()
# for win desktop probably use --server 127.0.0.1 --inbrowser
# For linux server probably use --server 127.0.0.1 or do not use any cmd flags
print(args)
free_mem_gb = get_cuda_free_memory_gb(gpu)
high_vram = free_mem_gb > 60
print(f'Free VRAM {free_mem_gb} GB')
print(f'High-VRAM Mode: {high_vram}')
text_encoder = LlamaModel.from_pretrained("hunyuanvideo-community/HunyuanVideo", subfolder='text_encoder', torch_dtype=torch.float16).cpu()
text_encoder_2 = CLIPTextModel.from_pretrained("hunyuanvideo-community/HunyuanVideo", subfolder='text_encoder_2', torch_dtype=torch.float16).cpu()
tokenizer = LlamaTokenizerFast.from_pretrained("hunyuanvideo-community/HunyuanVideo", subfolder='tokenizer')
tokenizer_2 = CLIPTokenizer.from_pretrained("hunyuanvideo-community/HunyuanVideo", subfolder='tokenizer_2')
vae = AutoencoderKLHunyuanVideo.from_pretrained("hunyuanvideo-community/HunyuanVideo", subfolder='vae', torch_dtype=torch.float16).cpu()
feature_extractor = SiglipImageProcessor.from_pretrained("lllyasviel/flux_redux_bfl", subfolder='feature_extractor')
image_encoder = SiglipVisionModel.from_pretrained("lllyasviel/flux_redux_bfl", subfolder='image_encoder', torch_dtype=torch.float16).cpu()
transformer = HunyuanVideoTransformer3DModelPacked.from_pretrained('lllyasviel/FramePack_F1_I2V_HY_20250503', torch_dtype=torch.bfloat16).cpu()
vae.eval()
text_encoder.eval()
text_encoder_2.eval()
image_encoder.eval()
transformer.eval()
if not high_vram:
vae.enable_slicing()
vae.enable_tiling()
transformer.high_quality_fp32_output_for_inference = True
print('transformer.high_quality_fp32_output_for_inference = True')
transformer.to(dtype=torch.bfloat16)
vae.to(dtype=torch.float16)
image_encoder.to(dtype=torch.float16)
text_encoder.to(dtype=torch.float16)
text_encoder_2.to(dtype=torch.float16)
vae.requires_grad_(False)
text_encoder.requires_grad_(False)
text_encoder_2.requires_grad_(False)
image_encoder.requires_grad_(False)
transformer.requires_grad_(False)
if not high_vram:
# DynamicSwapInstaller is same as huggingface's enable_sequential_offload but 3x faster
DynamicSwapInstaller.install_model(transformer, device=gpu)
DynamicSwapInstaller.install_model(text_encoder, device=gpu)
else:
text_encoder.to(gpu)
text_encoder_2.to(gpu)
image_encoder.to(gpu)
vae.to(gpu)
transformer.to(gpu)
stream = AsyncStream()
outputs_folder = './outputs/'
os.makedirs(outputs_folder, exist_ok=True)
@torch.no_grad()
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):
total_latent_sections = (total_second_length * 30) / (latent_window_size * 4)
total_latent_sections = int(max(round(total_latent_sections), 1))
job_id = generate_timestamp()
stream.output_queue.push(('progress', (None, '', make_progress_bar_html(0, 'Starting ...'))))
try:
# Clean GPU
if not high_vram:
unload_complete_models(
text_encoder, text_encoder_2, image_encoder, vae, transformer
)
# Text encoding
stream.output_queue.push(('progress', (None, '', make_progress_bar_html(0, 'Text encoding ...'))))
if not high_vram:
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.
load_model_as_complete(text_encoder_2, target_device=gpu)
llama_vec, clip_l_pooler = encode_prompt_conds(prompt, text_encoder, text_encoder_2, tokenizer, tokenizer_2)
if cfg == 1:
llama_vec_n, clip_l_pooler_n = torch.zeros_like(llama_vec), torch.zeros_like(clip_l_pooler)
else:
llama_vec_n, clip_l_pooler_n = encode_prompt_conds(n_prompt, text_encoder, text_encoder_2, tokenizer, tokenizer_2)
llama_vec, llama_attention_mask = crop_or_pad_yield_mask(llama_vec, length=512)
llama_vec_n, llama_attention_mask_n = crop_or_pad_yield_mask(llama_vec_n, length=512)
# Processing input image
stream.output_queue.push(('progress', (None, '', make_progress_bar_html(0, 'Image processing ...'))))
H, W, C = input_image.shape
height, width = find_nearest_bucket(H, W, resolution=640)
input_image_np = resize_and_center_crop(input_image, target_width=width, target_height=height)
Image.fromarray(input_image_np).save(os.path.join(outputs_folder, f'{job_id}.png'))
input_image_pt = torch.from_numpy(input_image_np).float() / 127.5 - 1
input_image_pt = input_image_pt.permute(2, 0, 1)[None, :, None]
# VAE encoding
stream.output_queue.push(('progress', (None, '', make_progress_bar_html(0, 'VAE encoding ...'))))
if not high_vram:
load_model_as_complete(vae, target_device=gpu)
start_latent = vae_encode(input_image_pt, vae)
# CLIP Vision
stream.output_queue.push(('progress', (None, '', make_progress_bar_html(0, 'CLIP Vision encoding ...'))))
if not high_vram:
load_model_as_complete(image_encoder, target_device=gpu)
image_encoder_output = hf_clip_vision_encode(input_image_np, feature_extractor, image_encoder)
image_encoder_last_hidden_state = image_encoder_output.last_hidden_state
# Dtype
llama_vec = llama_vec.to(transformer.dtype)
llama_vec_n = llama_vec_n.to(transformer.dtype)
clip_l_pooler = clip_l_pooler.to(transformer.dtype)
clip_l_pooler_n = clip_l_pooler_n.to(transformer.dtype)
image_encoder_last_hidden_state = image_encoder_last_hidden_state.to(transformer.dtype)
# Sampling
stream.output_queue.push(('progress', (None, '', make_progress_bar_html(0, 'Start sampling ...'))))
rnd = torch.Generator("cpu").manual_seed(seed)
history_latents = torch.zeros(size=(1, 16, 16 + 2 + 1, height // 8, width // 8), dtype=torch.float32).cpu()
history_pixels = None
history_latents = torch.cat([history_latents, start_latent.to(history_latents)], dim=2)
total_generated_latent_frames = 1
for section_index in range(total_latent_sections):
if stream.input_queue.top() == 'end':
stream.output_queue.push(('end', None))
return
print(f'section_index = {section_index}, total_latent_sections = {total_latent_sections}')
if not high_vram:
unload_complete_models()
move_model_to_device_with_memory_preservation(transformer, target_device=gpu, preserved_memory_gb=gpu_memory_preservation)
if use_teacache:
transformer.initialize_teacache(enable_teacache=True, num_steps=steps)
else:
transformer.initialize_teacache(enable_teacache=False)
def callback(d):
preview = d['denoised']
preview = vae_decode_fake(preview)
preview = (preview * 255.0).detach().cpu().numpy().clip(0, 255).astype(np.uint8)
preview = einops.rearrange(preview, 'b c t h w -> (b h) (t w) c')
if stream.input_queue.top() == 'end':
stream.output_queue.push(('end', None))
raise KeyboardInterrupt('User ends the task.')
current_step = d['i'] + 1
percentage = int(100.0 * current_step / steps)
hint = f'Sampling {current_step}/{steps}'
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 ...'
stream.output_queue.push(('progress', (preview, desc, make_progress_bar_html(percentage, hint))))
return
indices = torch.arange(0, sum([1, 16, 2, 1, latent_window_size])).unsqueeze(0)
clean_latent_indices_start, clean_latent_4x_indices, clean_latent_2x_indices, clean_latent_1x_indices, latent_indices = indices.split([1, 16, 2, 1, latent_window_size], dim=1)
clean_latent_indices = torch.cat([clean_latent_indices_start, clean_latent_1x_indices], dim=1)
clean_latents_4x, clean_latents_2x, clean_latents_1x = history_latents[:, :, -sum([16, 2, 1]):, :, :].split([16, 2, 1], dim=2)
clean_latents = torch.cat([start_latent.to(history_latents), clean_latents_1x], dim=2)
generated_latents = sample_hunyuan(
transformer=transformer,
sampler='unipc',
width=width,
height=height,
frames=latent_window_size * 4 - 3,
real_guidance_scale=cfg,
distilled_guidance_scale=gs,
guidance_rescale=rs,
# shift=3.0,
num_inference_steps=steps,
generator=rnd,
prompt_embeds=llama_vec,
prompt_embeds_mask=llama_attention_mask,
prompt_poolers=clip_l_pooler,
negative_prompt_embeds=llama_vec_n,
negative_prompt_embeds_mask=llama_attention_mask_n,
negative_prompt_poolers=clip_l_pooler_n,
device=gpu,
dtype=torch.bfloat16,
image_embeddings=image_encoder_last_hidden_state,
latent_indices=latent_indices,
clean_latents=clean_latents,
clean_latent_indices=clean_latent_indices,
clean_latents_2x=clean_latents_2x,
clean_latent_2x_indices=clean_latent_2x_indices,
clean_latents_4x=clean_latents_4x,
clean_latent_4x_indices=clean_latent_4x_indices,
callback=callback,
)
total_generated_latent_frames += int(generated_latents.shape[2])
history_latents = torch.cat([history_latents, generated_latents.to(history_latents)], dim=2)
if not high_vram:
offload_model_from_device_for_memory_preservation(transformer, target_device=gpu, preserved_memory_gb=8)
load_model_as_complete(vae, target_device=gpu)
real_history_latents = history_latents[:, :, -total_generated_latent_frames:, :, :]
if history_pixels is None:
history_pixels = vae_decode(real_history_latents, vae).cpu()
else:
section_latent_frames = latent_window_size * 2
overlapped_frames = latent_window_size * 4 - 3
current_pixels = vae_decode(real_history_latents[:, :, -section_latent_frames:], vae).cpu()
history_pixels = soft_append_bcthw(history_pixels, current_pixels, overlapped_frames)
if not high_vram:
unload_complete_models()
output_filename = os.path.join(outputs_folder, f'{job_id}_{total_generated_latent_frames}.mp4')
save_bcthw_as_mp4(history_pixels, output_filename, fps=30, crf=mp4_crf)
print(f'Decoded. Current latent shape {real_history_latents.shape}; pixel shape {history_pixels.shape}')
stream.output_queue.push(('file', output_filename))
except:
traceback.print_exc()
if not high_vram:
unload_complete_models(
text_encoder, text_encoder_2, image_encoder, vae, transformer
)
stream.output_queue.push(('end', None))
return
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):
global stream
assert input_image is not None, 'No input image!'
yield None, None, '', '', gr.update(interactive=False), gr.update(interactive=True)
stream = AsyncStream()
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)
output_filename = None
while True:
flag, data = stream.output_queue.next()
if flag == 'file':
output_filename = data
yield output_filename, gr.update(), gr.update(), gr.update(), gr.update(interactive=False), gr.update(interactive=True)
if flag == 'progress':
preview, desc, html = data
yield gr.update(), gr.update(visible=True, value=preview), desc, html, gr.update(interactive=False), gr.update(interactive=True)
if flag == 'end':
yield output_filename, gr.update(visible=False), gr.update(), '', gr.update(interactive=True), gr.update(interactive=False)
break
def end_process():
stream.input_queue.push('end')
quick_prompts = [
'The girl dances gracefully, with clear movements, full of charm.',
'A character doing some simple body movements.',
]
quick_prompts = [[x] for x in quick_prompts]
css = make_progress_bar_css()
block = gr.Blocks(css=css).queue()
with block:
gr.Markdown('# FramePack-F1')
with gr.Row():
with gr.Column():
input_image = gr.Image(sources='upload', type="numpy", label="Image", height=320)
prompt = gr.Textbox(label="Prompt", value='')
example_quick_prompts = gr.Dataset(samples=quick_prompts, label='Quick List', samples_per_page=1000, components=[prompt])
example_quick_prompts.click(lambda x: x[0], inputs=[example_quick_prompts], outputs=prompt, show_progress=False, queue=False)
with gr.Row():
start_button = gr.Button(value="Start Generation")
end_button = gr.Button(value="End Generation", interactive=False)
with gr.Group():
use_teacache = gr.Checkbox(label='Use TeaCache', value=True, info='Faster speed, but often makes hands and fingers slightly worse.')
n_prompt = gr.Textbox(label="Negative Prompt", value="", visible=False) # Not used
seed = gr.Number(label="Seed", value=31337, precision=0)
total_second_length = gr.Slider(label="Total Video Length (Seconds)", minimum=1, maximum=120, value=5, step=0.1)
latent_window_size = gr.Slider(label="Latent Window Size", minimum=1, maximum=33, value=9, step=1, visible=False) # Should not change
steps = gr.Slider(label="Steps", minimum=1, maximum=100, value=25, step=1, info='Changing this value is not recommended.')
cfg = gr.Slider(label="CFG Scale", minimum=1.0, maximum=32.0, value=1.0, step=0.01, visible=False) # Should not change
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.')
rs = gr.Slider(label="CFG Re-Scale", minimum=0.0, maximum=1.0, value=0.0, step=0.01, visible=False) # Should not change
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.")
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. ")
with gr.Column():
preview_image = gr.Image(label="Next Latents", height=200, visible=False)
result_video = gr.Video(label="Finished Frames", autoplay=True, show_share_button=False, height=512, loop=True)
progress_desc = gr.Markdown('', elem_classes='no-generating-animation')
progress_bar = gr.HTML('', elem_classes='no-generating-animation')
gr.HTML('<div style="text-align:center; margin-top:20px;">Share your results and find ideas at the <a href="https://x.com/search?q=framepack&f=live" target="_blank">FramePack Twitter (X) thread</a></div>')
ips = [input_image, prompt, n_prompt, seed, total_second_length, latent_window_size, steps, cfg, gs, rs, gpu_memory_preservation, use_teacache, mp4_crf]
start_button.click(fn=process, inputs=ips, outputs=[result_video, preview_image, progress_desc, progress_bar, start_button, end_button])
end_button.click(fn=end_process)
block.launch(
server_name=args.server,
server_port=args.port,
share=args.share,
inbrowser=args.inbrowser,
)
================================================
FILE: diffusers_helper/bucket_tools.py
================================================
bucket_options = {
640: [
(416, 960),
(448, 864),
(480, 832),
(512, 768),
(544, 704),
(576, 672),
(608, 640),
(640, 608),
(672, 576),
(704, 544),
(768, 512),
(832, 480),
(864, 448),
(960, 416),
],
}
def find_nearest_bucket(h, w, resolution=640):
min_metric = float('inf')
best_bucket = None
for (bucket_h, bucket_w) in bucket_options[resolution]:
metric = abs(h * bucket_w - w * bucket_h)
if metric <= min_metric:
min_metric = metric
best_bucket = (bucket_h, bucket_w)
return best_bucket
================================================
FILE: diffusers_helper/clip_vision.py
================================================
import numpy as np
def hf_clip_vision_encode(image, feature_extractor, image_encoder):
assert isinstance(image, np.ndarray)
assert image.ndim == 3 and image.shape[2] == 3
assert image.dtype == np.uint8
preprocessed = feature_extractor.preprocess(images=image, return_tensors="pt").to(device=image_encoder.device, dtype=image_encoder.dtype)
image_encoder_output = image_encoder(**preprocessed)
return image_encoder_output
================================================
FILE: diffusers_helper/dit_common.py
================================================
import torch
import accelerate.accelerator
from diffusers.models.normalization import RMSNorm, LayerNorm, FP32LayerNorm, AdaLayerNormContinuous
accelerate.accelerator.convert_outputs_to_fp32 = lambda x: x
def LayerNorm_forward(self, x):
return torch.nn.functional.layer_norm(x, self.normalized_shape, self.weight, self.bias, self.eps).to(x)
LayerNorm.forward = LayerNorm_forward
torch.nn.LayerNorm.forward = LayerNorm_forward
def FP32LayerNorm_forward(self, x):
origin_dtype = x.dtype
return torch.nn.functional.layer_norm(
x.float(),
self.normalized_shape,
self.weight.float() if self.weight is not None else None,
self.bias.float() if self.bias is not None else None,
self.eps,
).to(origin_dtype)
FP32LayerNorm.forward = FP32LayerNorm_forward
def RMSNorm_forward(self, hidden_states):
input_dtype = hidden_states.dtype
variance = hidden_states.to(torch.float32).pow(2).mean(-1, keepdim=True)
hidden_states = hidden_states * torch.rsqrt(variance + self.eps)
if self.weight is None:
return hidden_states.to(input_dtype)
return hidden_states.to(input_dtype) * self.weight.to(input_dtype)
RMSNorm.forward = RMSNorm_forward
def AdaLayerNormContinuous_forward(self, x, conditioning_embedding):
emb = self.linear(self.silu(conditioning_embedding))
scale, shift = emb.chunk(2, dim=1)
x = self.norm(x) * (1 + scale)[:, None, :] + shift[:, None, :]
return x
AdaLayerNormContinuous.forward = AdaLayerNormContinuous_forward
================================================
FILE: diffusers_helper/gradio/progress_bar.py
================================================
progress_html = '''
<div class="loader-container">
<div class="loader"></div>
<div class="progress-container">
<progress value="*number*" max="100"></progress>
</div>
<span>*text*</span>
</div>
'''
css = '''
.loader-container {
display: flex; /* Use flex to align items horizontally */
align-items: center; /* Center items vertically within the container */
white-space: nowrap; /* Prevent line breaks within the container */
}
.loader {
border: 8px solid #f3f3f3; /* Light grey */
border-top: 8px solid #3498db; /* Blue */
border-radius: 50%;
width: 30px;
height: 30px;
animation: spin 2s linear infinite;
}
@keyframes spin {
0% { transform: rotate(0deg); }
100% { transform: rotate(360deg); }
}
/* Style the progress bar */
progress {
appearance: none; /* Remove default styling */
height: 20px; /* Set the height of the progress bar */
border-radius: 5px; /* Round the corners of the progress bar */
background-color: #f3f3f3; /* Light grey background */
width: 100%;
vertical-align: middle !important;
}
/* Style the progress bar container */
.progress-container {
margin-left: 20px;
margin-right: 20px;
flex-grow: 1; /* Allow the progress container to take up remaining space */
}
/* Set the color of the progress bar fill */
progress::-webkit-progress-value {
background-color: #3498db; /* Blue color for the fill */
}
progress::-moz-progress-bar {
background-color: #3498db; /* Blue color for the fill in Firefox */
}
/* Style the text on the progress bar */
progress::after {
content: attr(value '%'); /* Display the progress value followed by '%' */
position: absolute;
top: 50%;
left: 50%;
transform: translate(-50%, -50%);
color: white; /* Set text color */
font-size: 14px; /* Set font size */
}
/* Style other texts */
.loader-container > span {
margin-left: 5px; /* Add spacing between the progress bar and the text */
}
.no-generating-animation > .generating {
display: none !important;
}
'''
def make_progress_bar_html(number, text):
return progress_html.replace('*number*', str(number)).replace('*text*', text)
def make_progress_bar_css():
return css
================================================
FILE: diffusers_helper/hf_login.py
================================================
import os
def login(token):
from huggingface_hub import login
import time
while True:
try:
login(token)
print('HF login ok.')
break
except Exception as e:
print(f'HF login failed: {e}. Retrying')
time.sleep(0.5)
hf_token = os.environ.get('HF_TOKEN', None)
if hf_token is not None:
login(hf_token)
================================================
FILE: diffusers_helper/hunyuan.py
================================================
import torch
from diffusers.pipelines.hunyuan_video.pipeline_hunyuan_video import DEFAULT_PROMPT_TEMPLATE
from diffusers_helper.utils import crop_or_pad_yield_mask
@torch.no_grad()
def encode_prompt_conds(prompt, text_encoder, text_encoder_2, tokenizer, tokenizer_2, max_length=256):
assert isinstance(prompt, str)
prompt = [prompt]
# LLAMA
prompt_llama = [DEFAULT_PROMPT_TEMPLATE["template"].format(p) for p in prompt]
crop_start = DEFAULT_PROMPT_TEMPLATE["crop_start"]
llama_inputs = tokenizer(
prompt_llama,
padding="max_length",
max_length=max_length + crop_start,
truncation=True,
return_tensors="pt",
return_length=False,
return_overflowing_tokens=False,
return_attention_mask=True,
)
llama_input_ids = llama_inputs.input_ids.to(text_encoder.device)
llama_attention_mask = llama_inputs.attention_mask.to(text_encoder.device)
llama_attention_length = int(llama_attention_mask.sum())
llama_outputs = text_encoder(
input_ids=llama_input_ids,
attention_mask=llama_attention_mask,
output_hidden_states=True,
)
llama_vec = llama_outputs.hidden_states[-3][:, crop_start:llama_attention_length]
# llama_vec_remaining = llama_outputs.hidden_states[-3][:, llama_attention_length:]
llama_attention_mask = llama_attention_mask[:, crop_start:llama_attention_length]
assert torch.all(llama_attention_mask.bool())
# CLIP
clip_l_input_ids = tokenizer_2(
prompt,
padding="max_length",
max_length=77,
truncation=True,
return_overflowing_tokens=False,
return_length=False,
return_tensors="pt",
).input_ids
clip_l_pooler = text_encoder_2(clip_l_input_ids.to(text_encoder_2.device), output_hidden_states=False).pooler_output
return llama_vec, clip_l_pooler
@torch.no_grad()
def vae_decode_fake(latents):
latent_rgb_factors = [
[-0.0395, -0.0331, 0.0445],
[0.0696, 0.0795, 0.0518],
[0.0135, -0.0945, -0.0282],
[0.0108, -0.0250, -0.0765],
[-0.0209, 0.0032, 0.0224],
[-0.0804, -0.0254, -0.0639],
[-0.0991, 0.0271, -0.0669],
[-0.0646, -0.0422, -0.0400],
[-0.0696, -0.0595, -0.0894],
[-0.0799, -0.0208, -0.0375],
[0.1166, 0.1627, 0.0962],
[0.1165, 0.0432, 0.0407],
[-0.2315, -0.1920, -0.1355],
[-0.0270, 0.0401, -0.0821],
[-0.0616, -0.0997, -0.0727],
[0.0249, -0.0469, -0.1703]
] # From comfyui
latent_rgb_factors_bias = [0.0259, -0.0192, -0.0761]
weight = torch.tensor(latent_rgb_factors, device=latents.device, dtype=latents.dtype).transpose(0, 1)[:, :, None, None, None]
bias = torch.tensor(latent_rgb_factors_bias, device=latents.device, dtype=latents.dtype)
images = torch.nn.functional.conv3d(latents, weight, bias=bias, stride=1, padding=0, dilation=1, groups=1)
images = images.clamp(0.0, 1.0)
return images
@torch.no_grad()
def vae_decode(latents, vae, image_mode=False):
latents = latents / vae.config.scaling_factor
if not image_mode:
image = vae.decode(latents.to(device=vae.device, dtype=vae.dtype)).sample
else:
latents = latents.to(device=vae.device, dtype=vae.dtype).unbind(2)
image = [vae.decode(l.unsqueeze(2)).sample for l in latents]
image = torch.cat(image, dim=2)
return image
@torch.no_grad()
def vae_encode(image, vae):
latents = vae.encode(image.to(device=vae.device, dtype=vae.dtype)).latent_dist.sample()
latents = latents * vae.config.scaling_factor
return latents
================================================
FILE: diffusers_helper/k_diffusion/uni_pc_fm.py
================================================
# Better Flow Matching UniPC by Lvmin Zhang
# (c) 2025
# CC BY-SA 4.0
# Attribution-ShareAlike 4.0 International Licence
import torch
from tqdm.auto import trange
def expand_dims(v, dims):
return v[(...,) + (None,) * (dims - 1)]
class FlowMatchUniPC:
def __init__(self, model, extra_args, variant='bh1'):
self.model = model
self.variant = variant
self.extra_args = extra_args
def model_fn(self, x, t):
return self.model(x, t, **self.extra_args)
def update_fn(self, x, model_prev_list, t_prev_list, t, order):
assert order <= len(model_prev_list)
dims = x.dim()
t_prev_0 = t_prev_list[-1]
lambda_prev_0 = - torch.log(t_prev_0)
lambda_t = - torch.log(t)
model_prev_0 = model_prev_list[-1]
h = lambda_t - lambda_prev_0
rks = []
D1s = []
for i in range(1, order):
t_prev_i = t_prev_list[-(i + 1)]
model_prev_i = model_prev_list[-(i + 1)]
lambda_prev_i = - torch.log(t_prev_i)
rk = ((lambda_prev_i - lambda_prev_0) / h)[0]
rks.append(rk)
D1s.append((model_prev_i - model_prev_0) / rk)
rks.append(1.)
rks = torch.tensor(rks, device=x.device)
R = []
b = []
hh = -h[0]
h_phi_1 = torch.expm1(hh)
h_phi_k = h_phi_1 / hh - 1
factorial_i = 1
if self.variant == 'bh1':
B_h = hh
elif self.variant == 'bh2':
B_h = torch.expm1(hh)
else:
raise NotImplementedError('Bad variant!')
for i in range(1, order + 1):
R.append(torch.pow(rks, i - 1))
b.append(h_phi_k * factorial_i / B_h)
factorial_i *= (i + 1)
h_phi_k = h_phi_k / hh - 1 / factorial_i
R = torch.stack(R)
b = torch.tensor(b, device=x.device)
use_predictor = len(D1s) > 0
if use_predictor:
D1s = torch.stack(D1s, dim=1)
if order == 2:
rhos_p = torch.tensor([0.5], device=b.device)
else:
rhos_p = torch.linalg.solve(R[:-1, :-1], b[:-1])
else:
D1s = None
rhos_p = None
if order == 1:
rhos_c = torch.tensor([0.5], device=b.device)
else:
rhos_c = torch.linalg.solve(R, b)
x_t_ = expand_dims(t / t_prev_0, dims) * x - expand_dims(h_phi_1, dims) * model_prev_0
if use_predictor:
pred_res = torch.tensordot(D1s, rhos_p, dims=([1], [0]))
else:
pred_res = 0
x_t = x_t_ - expand_dims(B_h, dims) * pred_res
model_t = self.model_fn(x_t, t)
if D1s is not None:
corr_res = torch.tensordot(D1s, rhos_c[:-1], dims=([1], [0]))
else:
corr_res = 0
D1_t = (model_t - model_prev_0)
x_t = x_t_ - expand_dims(B_h, dims) * (corr_res + rhos_c[-1] * D1_t)
return x_t, model_t
def sample(self, x, sigmas, callback=None, disable_pbar=False):
order = min(3, len(sigmas) - 2)
model_prev_list, t_prev_list = [], []
for i in trange(len(sigmas) - 1, disable=disable_pbar):
vec_t = sigmas[i].expand(x.shape[0])
if i == 0:
model_prev_list = [self.model_fn(x, vec_t)]
t_prev_list = [vec_t]
elif i < order:
init_order = i
x, model_x = self.update_fn(x, model_prev_list, t_prev_list, vec_t, init_order)
model_prev_list.append(model_x)
t_prev_list.append(vec_t)
else:
x, model_x = self.update_fn(x, model_prev_list, t_prev_list, vec_t, order)
model_prev_list.append(model_x)
t_prev_list.append(vec_t)
model_prev_list = model_prev_list[-order:]
t_prev_list = t_prev_list[-order:]
if callback is not None:
callback({'x': x, 'i': i, 'denoised': model_prev_list[-1]})
return model_prev_list[-1]
def sample_unipc(model, noise, sigmas, extra_args=None, callback=None, disable=False, variant='bh1'):
assert variant in ['bh1', 'bh2']
return FlowMatchUniPC(model, extra_args=extra_args, variant=variant).sample(noise, sigmas=sigmas, callback=callback, disable_pbar=disable)
================================================
FILE: diffusers_helper/k_diffusion/wrapper.py
================================================
import torch
def append_dims(x, target_dims):
return x[(...,) + (None,) * (target_dims - x.ndim)]
def rescale_noise_cfg(noise_cfg, noise_pred_text, guidance_rescale=1.0):
if guidance_rescale == 0:
return noise_cfg
std_text = noise_pred_text.std(dim=list(range(1, noise_pred_text.ndim)), keepdim=True)
std_cfg = noise_cfg.std(dim=list(range(1, noise_cfg.ndim)), keepdim=True)
noise_pred_rescaled = noise_cfg * (std_text / std_cfg)
noise_cfg = guidance_rescale * noise_pred_rescaled + (1.0 - guidance_rescale) * noise_cfg
return noise_cfg
def fm_wrapper(transformer, t_scale=1000.0):
def k_model(x, sigma, **extra_args):
dtype = extra_args['dtype']
cfg_scale = extra_args['cfg_scale']
cfg_rescale = extra_args['cfg_rescale']
concat_latent = extra_args['concat_latent']
original_dtype = x.dtype
sigma = sigma.float()
x = x.to(dtype)
timestep = (sigma * t_scale).to(dtype)
if concat_latent is None:
hidden_states = x
else:
hidden_states = torch.cat([x, concat_latent.to(x)], dim=1)
pred_positive = transformer(hidden_states=hidden_states, timestep=timestep, return_dict=False, **extra_args['positive'])[0].float()
if cfg_scale == 1.0:
pred_negative = torch.zeros_like(pred_positive)
else:
pred_negative = transformer(hidden_states=hidden_states, timestep=timestep, return_dict=False, **extra_args['negative'])[0].float()
pred_cfg = pred_negative + cfg_scale * (pred_positive - pred_negative)
pred = rescale_noise_cfg(pred_cfg, pred_positive, guidance_rescale=cfg_rescale)
x0 = x.float() - pred.float() * append_dims(sigma, x.ndim)
return x0.to(dtype=original_dtype)
return k_model
================================================
FILE: diffusers_helper/memory.py
================================================
# By lllyasviel
import torch
cpu = torch.device('cpu')
gpu = torch.device(f'cuda:{torch.cuda.current_device()}')
gpu_complete_modules = []
class DynamicSwapInstaller:
@staticmethod
def _install_module(module: torch.nn.Module, **kwargs):
original_class = module.__class__
module.__dict__['forge_backup_original_class'] = original_class
def hacked_get_attr(self, name: str):
if '_parameters' in self.__dict__:
_parameters = self.__dict__['_parameters']
if name in _parameters:
p = _parameters[name]
if p is None:
return None
if p.__class__ == torch.nn.Parameter:
return torch.nn.Parameter(p.to(**kwargs), requires_grad=p.requires_grad)
else:
return p.to(**kwargs)
if '_buffers' in self.__dict__:
_buffers = self.__dict__['_buffers']
if name in _buffers:
return _buffers[name].to(**kwargs)
return super(original_class, self).__getattr__(name)
module.__class__ = type('DynamicSwap_' + original_class.__name__, (original_class,), {
'__getattr__': hacked_get_attr,
})
return
@staticmethod
def _uninstall_module(module: torch.nn.Module):
if 'forge_backup_original_class' in module.__dict__:
module.__class__ = module.__dict__.pop('forge_backup_original_class')
return
@staticmethod
def install_model(model: torch.nn.Module, **kwargs):
for m in model.modules():
DynamicSwapInstaller._install_module(m, **kwargs)
return
@staticmethod
def uninstall_model(model: torch.nn.Module):
for m in model.modules():
DynamicSwapInstaller._uninstall_module(m)
return
def fake_diffusers_current_device(model: torch.nn.Module, target_device: torch.device):
if hasattr(model, 'scale_shift_table'):
model.scale_shift_table.data = model.scale_shift_table.data.to(target_device)
return
for k, p in model.named_modules():
if hasattr(p, 'weight'):
p.to(target_device)
return
def get_cuda_free_memory_gb(device=None):
if device is None:
device = gpu
memory_stats = torch.cuda.memory_stats(device)
bytes_active = memory_stats['active_bytes.all.current']
bytes_reserved = memory_stats['reserved_bytes.all.current']
bytes_free_cuda, _ = torch.cuda.mem_get_info(device)
bytes_inactive_reserved = bytes_reserved - bytes_active
bytes_total_available = bytes_free_cuda + bytes_inactive_reserved
return bytes_total_available / (1024 ** 3)
def move_model_to_device_with_memory_preservation(model, target_device, preserved_memory_gb=0):
print(f'Moving {model.__class__.__name__} to {target_device} with preserved memory: {preserved_memory_gb} GB')
for m in model.modules():
if get_cuda_free_memory_gb(target_device) <= preserved_memory_gb:
torch.cuda.empty_cache()
return
if hasattr(m, 'weight'):
m.to(device=target_device)
model.to(device=target_device)
torch.cuda.empty_cache()
return
def offload_model_from_device_for_memory_preservation(model, target_device, preserved_memory_gb=0):
print(f'Offloading {model.__class__.__name__} from {target_device} to preserve memory: {preserved_memory_gb} GB')
for m in model.modules():
if get_cuda_free_memory_gb(target_device) >= preserved_memory_gb:
torch.cuda.empty_cache()
return
if hasattr(m, 'weight'):
m.to(device=cpu)
model.to(device=cpu)
torch.cuda.empty_cache()
return
def unload_complete_models(*args):
for m in gpu_complete_modules + list(args):
m.to(device=cpu)
print(f'Unloaded {m.__class__.__name__} as complete.')
gpu_complete_modules.clear()
torch.cuda.empty_cache()
return
def load_model_as_complete(model, target_device, unload=True):
if unload:
unload_complete_models()
model.to(device=target_device)
print(f'Loaded {model.__class__.__name__} to {target_device} as complete.')
gpu_complete_modules.append(model)
return
================================================
FILE: diffusers_helper/models/hunyuan_video_packed.py
================================================
from typing import Any, Dict, List, Optional, Tuple, Union
import torch
import einops
import torch.nn as nn
import numpy as np
from diffusers.loaders import FromOriginalModelMixin
from diffusers.configuration_utils import ConfigMixin, register_to_config
from diffusers.loaders import PeftAdapterMixin
from diffusers.utils import logging
from diffusers.models.attention import FeedForward
from diffusers.models.attention_processor import Attention
from diffusers.models.embeddings import TimestepEmbedding, Timesteps, PixArtAlphaTextProjection
from diffusers.models.modeling_outputs import Transformer2DModelOutput
from diffusers.models.modeling_utils import ModelMixin
from diffusers_helper.dit_common import LayerNorm
from diffusers_helper.utils import zero_module
enabled_backends = []
if torch.backends.cuda.flash_sdp_enabled():
enabled_backends.append("flash")
if torch.backends.cuda.math_sdp_enabled():
enabled_backends.append("math")
if torch.backends.cuda.mem_efficient_sdp_enabled():
enabled_backends.append("mem_efficient")
if torch.backends.cuda.cudnn_sdp_enabled():
enabled_backends.append("cudnn")
print("Currently enabled native sdp backends:", enabled_backends)
try:
# raise NotImplementedError
from xformers.ops import memory_efficient_attention as xformers_attn_func
print('Xformers is installed!')
except:
print('Xformers is not installed!')
xformers_attn_func = None
try:
# raise NotImplementedError
from flash_attn import flash_attn_varlen_func, flash_attn_func
print('Flash Attn is installed!')
except:
print('Flash Attn is not installed!')
flash_attn_varlen_func = None
flash_attn_func = None
try:
# raise NotImplementedError
from sageattention import sageattn_varlen, sageattn
print('Sage Attn is installed!')
except:
print('Sage Attn is not installed!')
sageattn_varlen = None
sageattn = None
logger = logging.get_logger(__name__) # pylint: disable=invalid-name
def pad_for_3d_conv(x, kernel_size):
b, c, t, h, w = x.shape
pt, ph, pw = kernel_size
pad_t = (pt - (t % pt)) % pt
pad_h = (ph - (h % ph)) % ph
pad_w = (pw - (w % pw)) % pw
return torch.nn.functional.pad(x, (0, pad_w, 0, pad_h, 0, pad_t), mode='replicate')
def center_down_sample_3d(x, kernel_size):
# pt, ph, pw = kernel_size
# cp = (pt * ph * pw) // 2
# 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)
# xc = xp[cp]
# return xc
return torch.nn.functional.avg_pool3d(x, kernel_size, stride=kernel_size)
def get_cu_seqlens(text_mask, img_len):
batch_size = text_mask.shape[0]
text_len = text_mask.sum(dim=1)
max_len = text_mask.shape[1] + img_len
cu_seqlens = torch.zeros([2 * batch_size + 1], dtype=torch.int32, device="cuda")
for i in range(batch_size):
s = text_len[i] + img_len
s1 = i * max_len + s
s2 = (i + 1) * max_len
cu_seqlens[2 * i + 1] = s1
cu_seqlens[2 * i + 2] = s2
return cu_seqlens
def apply_rotary_emb_transposed(x, freqs_cis):
cos, sin = freqs_cis.unsqueeze(-2).chunk(2, dim=-1)
x_real, x_imag = x.unflatten(-1, (-1, 2)).unbind(-1)
x_rotated = torch.stack([-x_imag, x_real], dim=-1).flatten(3)
out = x.float() * cos + x_rotated.float() * sin
out = out.to(x)
return out
def attn_varlen_func(q, k, v, cu_seqlens_q, cu_seqlens_kv, max_seqlen_q, max_seqlen_kv):
if cu_seqlens_q is None and cu_seqlens_kv is None and max_seqlen_q is None and max_seqlen_kv is None:
if sageattn is not None:
x = sageattn(q, k, v, tensor_layout='NHD')
return x
if flash_attn_func is not None:
x = flash_attn_func(q, k, v)
return x
if xformers_attn_func is not None:
x = xformers_attn_func(q, k, v)
return x
x = torch.nn.functional.scaled_dot_product_attention(q.transpose(1, 2), k.transpose(1, 2), v.transpose(1, 2)).transpose(1, 2)
return x
B, L, H, C = q.shape
q = q.flatten(0, 1)
k = k.flatten(0, 1)
v = v.flatten(0, 1)
if sageattn_varlen is not None:
x = sageattn_varlen(q, k, v, cu_seqlens_q, cu_seqlens_kv, max_seqlen_q, max_seqlen_kv)
elif flash_attn_varlen_func is not None:
x = flash_attn_varlen_func(q, k, v, cu_seqlens_q, cu_seqlens_kv, max_seqlen_q, max_seqlen_kv)
else:
raise NotImplementedError('No Attn Installed!')
x = x.unflatten(0, (B, L))
return x
class HunyuanAttnProcessorFlashAttnDouble:
def __call__(self, attn, hidden_states, encoder_hidden_states, attention_mask, image_rotary_emb):
cu_seqlens_q, cu_seqlens_kv, max_seqlen_q, max_seqlen_kv = attention_mask
query = attn.to_q(hidden_states)
key = attn.to_k(hidden_states)
value = attn.to_v(hidden_states)
query = query.unflatten(2, (attn.heads, -1))
key = key.unflatten(2, (attn.heads, -1))
value = value.unflatten(2, (attn.heads, -1))
query = attn.norm_q(query)
key = attn.norm_k(key)
query = apply_rotary_emb_transposed(query, image_rotary_emb)
key = apply_rotary_emb_transposed(key, image_rotary_emb)
encoder_query = attn.add_q_proj(encoder_hidden_states)
encoder_key = attn.add_k_proj(encoder_hidden_states)
encoder_value = attn.add_v_proj(encoder_hidden_states)
encoder_query = encoder_query.unflatten(2, (attn.heads, -1))
encoder_key = encoder_key.unflatten(2, (attn.heads, -1))
encoder_value = encoder_value.unflatten(2, (attn.heads, -1))
encoder_query = attn.norm_added_q(encoder_query)
encoder_key = attn.norm_added_k(encoder_key)
query = torch.cat([query, encoder_query], dim=1)
key = torch.cat([key, encoder_key], dim=1)
value = torch.cat([value, encoder_value], dim=1)
hidden_states = attn_varlen_func(query, key, value, cu_seqlens_q, cu_seqlens_kv, max_seqlen_q, max_seqlen_kv)
hidden_states = hidden_states.flatten(-2)
txt_length = encoder_hidden_states.shape[1]
hidden_states, encoder_hidden_states = hidden_states[:, :-txt_length], hidden_states[:, -txt_length:]
hidden_states = attn.to_out[0](hidden_states)
hidden_states = attn.to_out[1](hidden_states)
encoder_hidden_states = attn.to_add_out(encoder_hidden_states)
return hidden_states, encoder_hidden_states
class HunyuanAttnProcessorFlashAttnSingle:
def __call__(self, attn, hidden_states, encoder_hidden_states, attention_mask, image_rotary_emb):
cu_seqlens_q, cu_seqlens_kv, max_seqlen_q, max_seqlen_kv = attention_mask
hidden_states = torch.cat([hidden_states, encoder_hidden_states], dim=1)
query = attn.to_q(hidden_states)
key = attn.to_k(hidden_states)
value = attn.to_v(hidden_states)
query = query.unflatten(2, (attn.heads, -1))
key = key.unflatten(2, (attn.heads, -1))
value = value.unflatten(2, (attn.heads, -1))
query = attn.norm_q(query)
key = attn.norm_k(key)
txt_length = encoder_hidden_states.shape[1]
query = torch.cat([apply_rotary_emb_transposed(query[:, :-txt_length], image_rotary_emb), query[:, -txt_length:]], dim=1)
key = torch.cat([apply_rotary_emb_transposed(key[:, :-txt_length], image_rotary_emb), key[:, -txt_length:]], dim=1)
hidden_states = attn_varlen_func(query, key, value, cu_seqlens_q, cu_seqlens_kv, max_seqlen_q, max_seqlen_kv)
hidden_states = hidden_states.flatten(-2)
hidden_states, encoder_hidden_states = hidden_states[:, :-txt_length], hidden_states[:, -txt_length:]
return hidden_states, encoder_hidden_states
class CombinedTimestepGuidanceTextProjEmbeddings(nn.Module):
def __init__(self, embedding_dim, pooled_projection_dim):
super().__init__()
self.time_proj = Timesteps(num_channels=256, flip_sin_to_cos=True, downscale_freq_shift=0)
self.timestep_embedder = TimestepEmbedding(in_channels=256, time_embed_dim=embedding_dim)
self.guidance_embedder = TimestepEmbedding(in_channels=256, time_embed_dim=embedding_dim)
self.text_embedder = PixArtAlphaTextProjection(pooled_projection_dim, embedding_dim, act_fn="silu")
def forward(self, timestep, guidance, pooled_projection):
timesteps_proj = self.time_proj(timestep)
timesteps_emb = self.timestep_embedder(timesteps_proj.to(dtype=pooled_projection.dtype))
guidance_proj = self.time_proj(guidance)
guidance_emb = self.guidance_embedder(guidance_proj.to(dtype=pooled_projection.dtype))
time_guidance_emb = timesteps_emb + guidance_emb
pooled_projections = self.text_embedder(pooled_projection)
conditioning = time_guidance_emb + pooled_projections
return conditioning
class CombinedTimestepTextProjEmbeddings(nn.Module):
def __init__(self, embedding_dim, pooled_projection_dim):
super().__init__()
self.time_proj = Timesteps(num_channels=256, flip_sin_to_cos=True, downscale_freq_shift=0)
self.timestep_embedder = TimestepEmbedding(in_channels=256, time_embed_dim=embedding_dim)
self.text_embedder = PixArtAlphaTextProjection(pooled_projection_dim, embedding_dim, act_fn="silu")
def forward(self, timestep, pooled_projection):
timesteps_proj = self.time_proj(timestep)
timesteps_emb = self.timestep_embedder(timesteps_proj.to(dtype=pooled_projection.dtype))
pooled_projections = self.text_embedder(pooled_projection)
conditioning = timesteps_emb + pooled_projections
return conditioning
class HunyuanVideoAdaNorm(nn.Module):
def __init__(self, in_features: int, out_features: Optional[int] = None) -> None:
super().__init__()
out_features = out_features or 2 * in_features
self.linear = nn.Linear(in_features, out_features)
self.nonlinearity = nn.SiLU()
def forward(
self, temb: torch.Tensor
) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
temb = self.linear(self.nonlinearity(temb))
gate_msa, gate_mlp = temb.chunk(2, dim=-1)
gate_msa, gate_mlp = gate_msa.unsqueeze(1), gate_mlp.unsqueeze(1)
return gate_msa, gate_mlp
class HunyuanVideoIndividualTokenRefinerBlock(nn.Module):
def __init__(
self,
num_attention_heads: int,
attention_head_dim: int,
mlp_width_ratio: str = 4.0,
mlp_drop_rate: float = 0.0,
attention_bias: bool = True,
) -> None:
super().__init__()
hidden_size = num_attention_heads * attention_head_dim
self.norm1 = LayerNorm(hidden_size, elementwise_affine=True, eps=1e-6)
self.attn = Attention(
query_dim=hidden_size,
cross_attention_dim=None,
heads=num_attention_heads,
dim_head=attention_head_dim,
bias=attention_bias,
)
self.norm2 = LayerNorm(hidden_size, elementwise_affine=True, eps=1e-6)
self.ff = FeedForward(hidden_size, mult=mlp_width_ratio, activation_fn="linear-silu", dropout=mlp_drop_rate)
self.norm_out = HunyuanVideoAdaNorm(hidden_size, 2 * hidden_size)
def forward(
self,
hidden_states: torch.Tensor,
temb: torch.Tensor,
attention_mask: Optional[torch.Tensor] = None,
) -> torch.Tensor:
norm_hidden_states = self.norm1(hidden_states)
attn_output = self.attn(
hidden_states=norm_hidden_states,
encoder_hidden_states=None,
attention_mask=attention_mask,
)
gate_msa, gate_mlp = self.norm_out(temb)
hidden_states = hidden_states + attn_output * gate_msa
ff_output = self.ff(self.norm2(hidden_states))
hidden_states = hidden_states + ff_output * gate_mlp
return hidden_states
class HunyuanVideoIndividualTokenRefiner(nn.Module):
def __init__(
self,
num_attention_heads: int,
attention_head_dim: int,
num_layers: int,
mlp_width_ratio: float = 4.0,
mlp_drop_rate: float = 0.0,
attention_bias: bool = True,
) -> None:
super().__init__()
self.refiner_blocks = nn.ModuleList(
[
HunyuanVideoIndividualTokenRefinerBlock(
num_attention_heads=num_attention_heads,
attention_head_dim=attention_head_dim,
mlp_width_ratio=mlp_width_ratio,
mlp_drop_rate=mlp_drop_rate,
attention_bias=attention_bias,
)
for _ in range(num_layers)
]
)
def forward(
self,
hidden_states: torch.Tensor,
temb: torch.Tensor,
attention_mask: Optional[torch.Tensor] = None,
) -> None:
self_attn_mask = None
if attention_mask is not None:
batch_size = attention_mask.shape[0]
seq_len = attention_mask.shape[1]
attention_mask = attention_mask.to(hidden_states.device).bool()
self_attn_mask_1 = attention_mask.view(batch_size, 1, 1, seq_len).repeat(1, 1, seq_len, 1)
self_attn_mask_2 = self_attn_mask_1.transpose(2, 3)
self_attn_mask = (self_attn_mask_1 & self_attn_mask_2).bool()
self_attn_mask[:, :, :, 0] = True
for block in self.refiner_blocks:
hidden_states = block(hidden_states, temb, self_attn_mask)
return hidden_states
class HunyuanVideoTokenRefiner(nn.Module):
def __init__(
self,
in_channels: int,
num_attention_heads: int,
attention_head_dim: int,
num_layers: int,
mlp_ratio: float = 4.0,
mlp_drop_rate: float = 0.0,
attention_bias: bool = True,
) -> None:
super().__init__()
hidden_size = num_attention_heads * attention_head_dim
self.time_text_embed = CombinedTimestepTextProjEmbeddings(
embedding_dim=hidden_size, pooled_projection_dim=in_channels
)
self.proj_in = nn.Linear(in_channels, hidden_size, bias=True)
self.token_refiner = HunyuanVideoIndividualTokenRefiner(
num_attention_heads=num_attention_heads,
attention_head_dim=attention_head_dim,
num_layers=num_layers,
mlp_width_ratio=mlp_ratio,
mlp_drop_rate=mlp_drop_rate,
attention_bias=attention_bias,
)
def forward(
self,
hidden_states: torch.Tensor,
timestep: torch.LongTensor,
attention_mask: Optional[torch.LongTensor] = None,
) -> torch.Tensor:
if attention_mask is None:
pooled_projections = hidden_states.mean(dim=1)
else:
original_dtype = hidden_states.dtype
mask_float = attention_mask.float().unsqueeze(-1)
pooled_projections = (hidden_states * mask_float).sum(dim=1) / mask_float.sum(dim=1)
pooled_projections = pooled_projections.to(original_dtype)
temb = self.time_text_embed(timestep, pooled_projections)
hidden_states = self.proj_in(hidden_states)
hidden_states = self.token_refiner(hidden_states, temb, attention_mask)
return hidden_states
class HunyuanVideoRotaryPosEmbed(nn.Module):
def __init__(self, rope_dim, theta):
super().__init__()
self.DT, self.DY, self.DX = rope_dim
self.theta = theta
@torch.no_grad()
def get_frequency(self, dim, pos):
T, H, W = pos.shape
freqs = 1.0 / (self.theta ** (torch.arange(0, dim, 2, dtype=torch.float32, device=pos.device)[: (dim // 2)] / dim))
freqs = torch.outer(freqs, pos.reshape(-1)).unflatten(-1, (T, H, W)).repeat_interleave(2, dim=0)
return freqs.cos(), freqs.sin()
@torch.no_grad()
def forward_inner(self, frame_indices, height, width, device):
GT, GY, GX = torch.meshgrid(
frame_indices.to(device=device, dtype=torch.float32),
torch.arange(0, height, device=device, dtype=torch.float32),
torch.arange(0, width, device=device, dtype=torch.float32),
indexing="ij"
)
FCT, FST = self.get_frequency(self.DT, GT)
FCY, FSY = self.get_frequency(self.DY, GY)
FCX, FSX = self.get_frequency(self.DX, GX)
result = torch.cat([FCT, FCY, FCX, FST, FSY, FSX], dim=0)
return result.to(device)
@torch.no_grad()
def forward(self, frame_indices, height, width, device):
frame_indices = frame_indices.unbind(0)
results = [self.forward_inner(f, height, width, device) for f in frame_indices]
results = torch.stack(results, dim=0)
return results
class AdaLayerNormZero(nn.Module):
def __init__(self, embedding_dim: int, norm_type="layer_norm", bias=True):
super().__init__()
self.silu = nn.SiLU()
self.linear = nn.Linear(embedding_dim, 6 * embedding_dim, bias=bias)
if norm_type == "layer_norm":
self.norm = LayerNorm(embedding_dim, elementwise_affine=False, eps=1e-6)
else:
raise ValueError(f"unknown norm_type {norm_type}")
def forward(
self,
x: torch.Tensor,
emb: Optional[torch.Tensor] = None,
) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
emb = emb.unsqueeze(-2)
emb = self.linear(self.silu(emb))
shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = emb.chunk(6, dim=-1)
x = self.norm(x) * (1 + scale_msa) + shift_msa
return x, gate_msa, shift_mlp, scale_mlp, gate_mlp
class AdaLayerNormZeroSingle(nn.Module):
def __init__(self, embedding_dim: int, norm_type="layer_norm", bias=True):
super().__init__()
self.silu = nn.SiLU()
self.linear = nn.Linear(embedding_dim, 3 * embedding_dim, bias=bias)
if norm_type == "layer_norm":
self.norm = LayerNorm(embedding_dim, elementwise_affine=False, eps=1e-6)
else:
raise ValueError(f"unknown norm_type {norm_type}")
def forward(
self,
x: torch.Tensor,
emb: Optional[torch.Tensor] = None,
) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
emb = emb.unsqueeze(-2)
emb = self.linear(self.silu(emb))
shift_msa, scale_msa, gate_msa = emb.chunk(3, dim=-1)
x = self.norm(x) * (1 + scale_msa) + shift_msa
return x, gate_msa
class AdaLayerNormContinuous(nn.Module):
def __init__(
self,
embedding_dim: int,
conditioning_embedding_dim: int,
elementwise_affine=True,
eps=1e-5,
bias=True,
norm_type="layer_norm",
):
super().__init__()
self.silu = nn.SiLU()
self.linear = nn.Linear(conditioning_embedding_dim, embedding_dim * 2, bias=bias)
if norm_type == "layer_norm":
self.norm = LayerNorm(embedding_dim, eps, elementwise_affine, bias)
else:
raise ValueError(f"unknown norm_type {norm_type}")
def forward(self, x: torch.Tensor, emb: torch.Tensor) -> torch.Tensor:
emb = emb.unsqueeze(-2)
emb = self.linear(self.silu(emb))
scale, shift = emb.chunk(2, dim=-1)
x = self.norm(x) * (1 + scale) + shift
return x
class HunyuanVideoSingleTransformerBlock(nn.Module):
def __init__(
self,
num_attention_heads: int,
attention_head_dim: int,
mlp_ratio: float = 4.0,
qk_norm: str = "rms_norm",
) -> None:
super().__init__()
hidden_size = num_attention_heads * attention_head_dim
mlp_dim = int(hidden_size * mlp_ratio)
self.attn = Attention(
query_dim=hidden_size,
cross_attention_dim=None,
dim_head=attention_head_dim,
heads=num_attention_heads,
out_dim=hidden_size,
bias=True,
processor=HunyuanAttnProcessorFlashAttnSingle(),
qk_norm=qk_norm,
eps=1e-6,
pre_only=True,
)
self.norm = AdaLayerNormZeroSingle(hidden_size, norm_type="layer_norm")
self.proj_mlp = nn.Linear(hidden_size, mlp_dim)
self.act_mlp = nn.GELU(approximate="tanh")
self.proj_out = nn.Linear(hidden_size + mlp_dim, hidden_size)
def forward(
self,
hidden_states: torch.Tensor,
encoder_hidden_states: torch.Tensor,
temb: torch.Tensor,
attention_mask: Optional[torch.Tensor] = None,
image_rotary_emb: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
) -> torch.Tensor:
text_seq_length = encoder_hidden_states.shape[1]
hidden_states = torch.cat([hidden_states, encoder_hidden_states], dim=1)
residual = hidden_states
# 1. Input normalization
norm_hidden_states, gate = self.norm(hidden_states, emb=temb)
mlp_hidden_states = self.act_mlp(self.proj_mlp(norm_hidden_states))
norm_hidden_states, norm_encoder_hidden_states = (
norm_hidden_states[:, :-text_seq_length, :],
norm_hidden_states[:, -text_seq_length:, :],
)
# 2. Attention
attn_output, context_attn_output = self.attn(
hidden_states=norm_hidden_states,
encoder_hidden_states=norm_encoder_hidden_states,
attention_mask=attention_mask,
image_rotary_emb=image_rotary_emb,
)
attn_output = torch.cat([attn_output, context_attn_output], dim=1)
# 3. Modulation and residual connection
hidden_states = torch.cat([attn_output, mlp_hidden_states], dim=2)
hidden_states = gate * self.proj_out(hidden_states)
hidden_states = hidden_states + residual
hidden_states, encoder_hidden_states = (
hidden_states[:, :-text_seq_length, :],
hidden_states[:, -text_seq_length:, :],
)
return hidden_states, encoder_hidden_states
class HunyuanVideoTransformerBlock(nn.Module):
def __init__(
self,
num_attention_heads: int,
attention_head_dim: int,
mlp_ratio: float,
qk_norm: str = "rms_norm",
) -> None:
super().__init__()
hidden_size = num_attention_heads * attention_head_dim
self.norm1 = AdaLayerNormZero(hidden_size, norm_type="layer_norm")
self.norm1_context = AdaLayerNormZero(hidden_size, norm_type="layer_norm")
self.attn = Attention(
query_dim=hidden_size,
cross_attention_dim=None,
added_kv_proj_dim=hidden_size,
dim_head=attention_head_dim,
heads=num_attention_heads,
out_dim=hidden_size,
context_pre_only=False,
bias=True,
processor=HunyuanAttnProcessorFlashAttnDouble(),
qk_norm=qk_norm,
eps=1e-6,
)
self.norm2 = LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6)
self.ff = FeedForward(hidden_size, mult=mlp_ratio, activation_fn="gelu-approximate")
self.norm2_context = LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6)
self.ff_context = FeedForward(hidden_size, mult=mlp_ratio, activation_fn="gelu-approximate")
def forward(
self,
hidden_states: torch.Tensor,
encoder_hidden_states: torch.Tensor,
temb: torch.Tensor,
attention_mask: Optional[torch.Tensor] = None,
freqs_cis: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
) -> Tuple[torch.Tensor, torch.Tensor]:
# 1. Input normalization
norm_hidden_states, gate_msa, shift_mlp, scale_mlp, gate_mlp = self.norm1(hidden_states, emb=temb)
norm_encoder_hidden_states, c_gate_msa, c_shift_mlp, c_scale_mlp, c_gate_mlp = self.norm1_context(encoder_hidden_states, emb=temb)
# 2. Joint attention
attn_output, context_attn_output = self.attn(
hidden_states=norm_hidden_states,
encoder_hidden_states=norm_encoder_hidden_states,
attention_mask=attention_mask,
image_rotary_emb=freqs_cis,
)
# 3. Modulation and residual connection
hidden_states = hidden_states + attn_output * gate_msa
encoder_hidden_states = encoder_hidden_states + context_attn_output * c_gate_msa
norm_hidden_states = self.norm2(hidden_states)
norm_encoder_hidden_states = self.norm2_context(encoder_hidden_states)
norm_hidden_states = norm_hidden_states * (1 + scale_mlp) + shift_mlp
norm_encoder_hidden_states = norm_encoder_hidden_states * (1 + c_scale_mlp) + c_shift_mlp
# 4. Feed-forward
ff_output = self.ff(norm_hidden_states)
context_ff_output = self.ff_context(norm_encoder_hidden_states)
hidden_states = hidden_states + gate_mlp * ff_output
encoder_hidden_states = encoder_hidden_states + c_gate_mlp * context_ff_output
return hidden_states, encoder_hidden_states
class ClipVisionProjection(nn.Module):
def __init__(self, in_channels, out_channels):
super().__init__()
self.up = nn.Linear(in_channels, out_channels * 3)
self.down = nn.Linear(out_channels * 3, out_channels)
def forward(self, x):
projected_x = self.down(nn.functional.silu(self.up(x)))
return projected_x
class HunyuanVideoPatchEmbed(nn.Module):
def __init__(self, patch_size, in_chans, embed_dim):
super().__init__()
self.proj = nn.Conv3d(in_chans, embed_dim, kernel_size=patch_size, stride=patch_size)
class HunyuanVideoPatchEmbedForCleanLatents(nn.Module):
def __init__(self, inner_dim):
super().__init__()
self.proj = nn.Conv3d(16, inner_dim, kernel_size=(1, 2, 2), stride=(1, 2, 2))
self.proj_2x = nn.Conv3d(16, inner_dim, kernel_size=(2, 4, 4), stride=(2, 4, 4))
self.proj_4x = nn.Conv3d(16, inner_dim, kernel_size=(4, 8, 8), stride=(4, 8, 8))
@torch.no_grad()
def initialize_weight_from_another_conv3d(self, another_layer):
weight = another_layer.weight.detach().clone()
bias = another_layer.bias.detach().clone()
sd = {
'proj.weight': weight.clone(),
'proj.bias': bias.clone(),
'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,
'proj_2x.bias': bias.clone(),
'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,
'proj_4x.bias': bias.clone(),
}
sd = {k: v.clone() for k, v in sd.items()}
self.load_state_dict(sd)
return
class HunyuanVideoTransformer3DModelPacked(ModelMixin, ConfigMixin, PeftAdapterMixin, FromOriginalModelMixin):
@register_to_config
def __init__(
self,
in_channels: int = 16,
out_channels: int = 16,
num_attention_heads: int = 24,
attention_head_dim: int = 128,
num_layers: int = 20,
num_single_layers: int = 40,
num_refiner_layers: int = 2,
mlp_ratio: float = 4.0,
patch_size: int = 2,
patch_size_t: int = 1,
qk_norm: str = "rms_norm",
guidance_embeds: bool = True,
text_embed_dim: int = 4096,
pooled_projection_dim: int = 768,
rope_theta: float = 256.0,
rope_axes_dim: Tuple[int] = (16, 56, 56),
has_image_proj=False,
image_proj_dim=1152,
has_clean_x_embedder=False,
) -> None:
super().__init__()
inner_dim = num_attention_heads * attention_head_dim
out_channels = out_channels or in_channels
# 1. Latent and condition embedders
self.x_embedder = HunyuanVideoPatchEmbed((patch_size_t, patch_size, patch_size), in_channels, inner_dim)
self.context_embedder = HunyuanVideoTokenRefiner(
text_embed_dim, num_attention_heads, attention_head_dim, num_layers=num_refiner_layers
)
self.time_text_embed = CombinedTimestepGuidanceTextProjEmbeddings(inner_dim, pooled_projection_dim)
self.clean_x_embedder = None
self.image_projection = None
# 2. RoPE
self.rope = HunyuanVideoRotaryPosEmbed(rope_axes_dim, rope_theta)
# 3. Dual stream transformer blocks
self.transformer_blocks = nn.ModuleList(
[
HunyuanVideoTransformerBlock(
num_attention_heads, attention_head_dim, mlp_ratio=mlp_ratio, qk_norm=qk_norm
)
for _ in range(num_layers)
]
)
# 4. Single stream transformer blocks
self.single_transformer_blocks = nn.ModuleList(
[
HunyuanVideoSingleTransformerBlock(
num_attention_heads, attention_head_dim, mlp_ratio=mlp_ratio, qk_norm=qk_norm
)
for _ in range(num_single_layers)
]
)
# 5. Output projection
self.norm_out = AdaLayerNormContinuous(inner_dim, inner_dim, elementwise_affine=False, eps=1e-6)
self.proj_out = nn.Linear(inner_dim, patch_size_t * patch_size * patch_size * out_channels)
self.inner_dim = inner_dim
self.use_gradient_checkpointing = False
self.enable_teacache = False
if has_image_proj:
self.install_image_projection(image_proj_dim)
if has_clean_x_embedder:
self.install_clean_x_embedder()
self.high_quality_fp32_output_for_inference = False
def install_image_projection(self, in_channels):
self.image_projection = ClipVisionProjection(in_channels=in_channels, out_channels=self.inner_dim)
self.config['has_image_proj'] = True
self.config['image_proj_dim'] = in_channels
def install_clean_x_embedder(self):
self.clean_x_embedder = HunyuanVideoPatchEmbedForCleanLatents(self.inner_dim)
self.config['has_clean_x_embedder'] = True
def enable_gradient_checkpointing(self):
self.use_gradient_checkpointing = True
print('self.use_gradient_checkpointing = True')
def disable_gradient_checkpointing(self):
self.use_gradient_checkpointing = False
print('self.use_gradient_checkpointing = False')
def initialize_teacache(self, enable_teacache=True, num_steps=25, rel_l1_thresh=0.15):
self.enable_teacache = enable_teacache
self.cnt = 0
self.num_steps = num_steps
self.rel_l1_thresh = rel_l1_thresh # 0.1 for 1.6x speedup, 0.15 for 2.1x speedup
self.accumulated_rel_l1_distance = 0
self.previous_modulated_input = None
self.previous_residual = None
self.teacache_rescale_func = np.poly1d([7.33226126e+02, -4.01131952e+02, 6.75869174e+01, -3.14987800e+00, 9.61237896e-02])
def gradient_checkpointing_method(self, block, *args):
if self.use_gradient_checkpointing:
result = torch.utils.checkpoint.checkpoint(block, *args, use_reentrant=False)
else:
result = block(*args)
return result
def process_input_hidden_states(
self,
latents, latent_indices=None,
clean_latents=None, clean_latent_indices=None,
clean_latents_2x=None, clean_latent_2x_indices=None,
clean_latents_4x=None, clean_latent_4x_indices=None
):
hidden_states = self.gradient_checkpointing_method(self.x_embedder.proj, latents)
B, C, T, H, W = hidden_states.shape
if latent_indices is None:
latent_indices = torch.arange(0, T).unsqueeze(0).expand(B, -1)
hidden_states = hidden_states.flatten(2).transpose(1, 2)
rope_freqs = self.rope(frame_indices=latent_indices, height=H, width=W, device=hidden_states.device)
rope_freqs = rope_freqs.flatten(2).transpose(1, 2)
if clean_latents is not None and clean_latent_indices is not None:
clean_latents = clean_latents.to(hidden_states)
clean_latents = self.gradient_checkpointing_method(self.clean_x_embedder.proj, clean_latents)
clean_latents = clean_latents.flatten(2).transpose(1, 2)
clean_latent_rope_freqs = self.rope(frame_indices=clean_latent_indices, height=H, width=W, device=clean_latents.device)
clean_latent_rope_freqs = clean_latent_rope_freqs.flatten(2).transpose(1, 2)
hidden_states = torch.cat([clean_latents, hidden_states], dim=1)
rope_freqs = torch.cat([clean_latent_rope_freqs, rope_freqs], dim=1)
if clean_latents_2x is not None and clean_latent_2x_indices is not None:
clean_latents_2x = clean_latents_2x.to(hidden_states)
clean_latents_2x = pad_for_3d_conv(clean_latents_2x, (2, 4, 4))
clean_latents_2x = self.gradient_checkpointing_method(self.clean_x_embedder.proj_2x, clean_latents_2x)
clean_latents_2x = clean_latents_2x.flatten(2).transpose(1, 2)
clean_latent_2x_rope_freqs = self.rope(frame_indices=clean_latent_2x_indices, height=H, width=W, device=clean_latents_2x.device)
clean_latent_2x_rope_freqs = pad_for_3d_conv(clean_latent_2x_rope_freqs, (2, 2, 2))
clean_latent_2x_rope_freqs = center_down_sample_3d(clean_latent_2x_rope_freqs, (2, 2, 2))
clean_latent_2x_rope_freqs = clean_latent_2x_rope_freqs.flatten(2).transpose(1, 2)
hidden_states = torch.cat([clean_latents_2x, hidden_states], dim=1)
rope_freqs = torch.cat([clean_latent_2x_rope_freqs, rope_freqs], dim=1)
if clean_latents_4x is not None and clean_latent_4x_indices is not None:
clean_latents_4x = clean_latents_4x.to(hidden_states)
clean_latents_4x = pad_for_3d_conv(clean_latents_4x, (4, 8, 8))
clean_latents_4x = self.gradient_checkpointing_method(self.clean_x_embedder.proj_4x, clean_latents_4x)
clean_latents_4x = clean_latents_4x.flatten(2).transpose(1, 2)
clean_latent_4x_rope_freqs = self.rope(frame_indices=clean_latent_4x_indices, height=H, width=W, device=clean_latents_4x.device)
clean_latent_4x_rope_freqs = pad_for_3d_conv(clean_latent_4x_rope_freqs, (4, 4, 4))
clean_latent_4x_rope_freqs = center_down_sample_3d(clean_latent_4x_rope_freqs, (4, 4, 4))
clean_latent_4x_rope_freqs = clean_latent_4x_rope_freqs.flatten(2).transpose(1, 2)
hidden_states = torch.cat([clean_latents_4x, hidden_states], dim=1)
rope_freqs = torch.cat([clean_latent_4x_rope_freqs, rope_freqs], dim=1)
return hidden_states, rope_freqs
def forward(
self,
hidden_states, timestep, encoder_hidden_states, encoder_attention_mask, pooled_projections, guidance,
latent_indices=None,
clean_latents=None, clean_latent_indices=None,
clean_latents_2x=None, clean_latent_2x_indices=None,
clean_latents_4x=None, clean_latent_4x_indices=None,
image_embeddings=None,
attention_kwargs=None, return_dict=True
):
if attention_kwargs is None:
attention_kwargs = {}
batch_size, num_channels, num_frames, height, width = hidden_states.shape
p, p_t = self.config['patch_size'], self.config['patch_size_t']
post_patch_num_frames = num_frames // p_t
post_patch_height = height // p
post_patch_width = width // p
original_context_length = post_patch_num_frames * post_patch_height * post_patch_width
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)
temb = self.gradient_checkpointing_method(self.time_text_embed, timestep, guidance, pooled_projections)
encoder_hidden_states = self.gradient_checkpointing_method(self.context_embedder, encoder_hidden_states, timestep, encoder_attention_mask)
if self.image_projection is not None:
assert image_embeddings is not None, 'You must use image embeddings!'
extra_encoder_hidden_states = self.gradient_checkpointing_method(self.image_projection, image_embeddings)
extra_attention_mask = torch.ones((batch_size, extra_encoder_hidden_states.shape[1]), dtype=encoder_attention_mask.dtype, device=encoder_attention_mask.device)
# must cat before (not after) encoder_hidden_states, due to attn masking
encoder_hidden_states = torch.cat([extra_encoder_hidden_states, encoder_hidden_states], dim=1)
encoder_attention_mask = torch.cat([extra_attention_mask, encoder_attention_mask], dim=1)
if batch_size == 1:
# When batch size is 1, we do not need any masks or var-len funcs since cropping is mathematically same to what we want
# If they are not same, then their impls are wrong. Ours are always the correct one.
text_len = encoder_attention_mask.sum().item()
encoder_hidden_states = encoder_hidden_states[:, :text_len]
attention_mask = None, None, None, None
else:
img_seq_len = hidden_states.shape[1]
txt_seq_len = encoder_hidden_states.shape[1]
cu_seqlens_q = get_cu_seqlens(encoder_attention_mask, img_seq_len)
cu_seqlens_kv = cu_seqlens_q
max_seqlen_q = img_seq_len + txt_seq_len
max_seqlen_kv = max_seqlen_q
attention_mask = cu_seqlens_q, cu_seqlens_kv, max_seqlen_q, max_seqlen_kv
if self.enable_teacache:
modulated_inp = self.transformer_blocks[0].norm1(hidden_states, emb=temb)[0]
if self.cnt == 0 or self.cnt == self.num_steps-1:
should_calc = True
self.accumulated_rel_l1_distance = 0
else:
curr_rel_l1 = ((modulated_inp - self.previous_modulated_input).abs().mean() / self.previous_modulated_input.abs().mean()).cpu().item()
self.accumulated_rel_l1_distance += self.teacache_rescale_func(curr_rel_l1)
should_calc = self.accumulated_rel_l1_distance >= self.rel_l1_thresh
if should_calc:
self.accumulated_rel_l1_distance = 0
self.previous_modulated_input = modulated_inp
self.cnt += 1
if self.cnt == self.num_steps:
self.cnt = 0
if not should_calc:
hidden_states = hidden_states + self.previous_residual
else:
ori_hidden_states = hidden_states.clone()
for block_id, block in enumerate(self.transformer_blocks):
hidden_states, encoder_hidden_states = self.gradient_checkpointing_method(
block,
hidden_states,
encoder_hidden_states,
temb,
attention_mask,
rope_freqs
)
for block_id, block in enumerate(self.single_transformer_blocks):
hidden_states, encoder_hidden_states = self.gradient_checkpointing_method(
block,
hidden_states,
encoder_hidden_states,
temb,
attention_mask,
rope_freqs
)
self.previous_residual = hidden_states - ori_hidden_states
else:
for block_id, block in enumerate(self.transformer_blocks):
hidden_states, encoder_hidden_states = self.gradient_checkpointing_method(
block,
hidden_states,
encoder_hidden_states,
temb,
attention_mask,
rope_freqs
)
for block_id, block in enumerate(self.single_transformer_blocks):
hidden_states, encoder_hidden_states = self.gradient_checkpointing_method(
block,
hidden_states,
encoder_hidden_states,
temb,
attention_mask,
rope_freqs
)
hidden_states = self.gradient_checkpointing_method(self.norm_out, hidden_states, temb)
hidden_states = hidden_states[:, -original_context_length:, :]
if self.high_quality_fp32_output_for_inference:
hidden_states = hidden_states.to(dtype=torch.float32)
if self.proj_out.weight.dtype != torch.float32:
self.proj_out.to(dtype=torch.float32)
hidden_states = self.gradient_checkpointing_method(self.proj_out, hidden_states)
hidden_states = einops.rearrange(hidden_states, 'b (t h w) (c pt ph pw) -> b c (t pt) (h ph) (w pw)',
t=post_patch_num_frames, h=post_patch_height, w=post_patch_width,
pt=p_t, ph=p, pw=p)
if return_dict:
return Transformer2DModelOutput(sample=hidden_states)
return hidden_states,
================================================
FILE: diffusers_helper/pipelines/k_diffusion_hunyuan.py
================================================
import torch
import math
from diffusers_helper.k_diffusion.uni_pc_fm import sample_unipc
from diffusers_helper.k_diffusion.wrapper import fm_wrapper
from diffusers_helper.utils import repeat_to_batch_size
def flux_time_shift(t, mu=1.15, sigma=1.0):
return math.exp(mu) / (math.exp(mu) + (1 / t - 1) ** sigma)
def calculate_flux_mu(context_length, x1=256, y1=0.5, x2=4096, y2=1.15, exp_max=7.0):
k = (y2 - y1) / (x2 - x1)
b = y1 - k * x1
mu = k * context_length + b
mu = min(mu, math.log(exp_max))
return mu
def get_flux_sigmas_from_mu(n, mu):
sigmas = torch.linspace(1, 0, steps=n + 1)
sigmas = flux_time_shift(sigmas, mu=mu)
return sigmas
@torch.inference_mode()
def sample_hunyuan(
transformer,
sampler='unipc',
initial_latent=None,
concat_latent=None,
strength=1.0,
width=512,
height=512,
frames=16,
real_guidance_scale=1.0,
distilled_guidance_scale=6.0,
guidance_rescale=0.0,
shift=None,
num_inference_steps=25,
batch_size=None,
generator=None,
prompt_embeds=None,
prompt_embeds_mask=None,
prompt_poolers=None,
negative_prompt_embeds=None,
negative_prompt_embeds_mask=None,
negative_prompt_poolers=None,
dtype=torch.bfloat16,
device=None,
negative_kwargs=None,
callback=None,
**kwargs,
):
device = device or transformer.device
if batch_size is None:
batch_size = int(prompt_embeds.shape[0])
latents = torch.randn((batch_size, 16, (frames + 3) // 4, height // 8, width // 8), generator=generator, device=generator.device).to(device=device, dtype=torch.float32)
B, C, T, H, W = latents.shape
seq_length = T * H * W // 4
if shift is None:
mu = calculate_flux_mu(seq_length, exp_max=7.0)
else:
mu = math.log(shift)
sigmas = get_flux_sigmas_from_mu(num_inference_steps, mu).to(device)
k_model = fm_wrapper(transformer)
if initial_latent is not None:
sigmas = sigmas * strength
first_sigma = sigmas[0].to(device=device, dtype=torch.float32)
initial_latent = initial_latent.to(device=device, dtype=torch.float32)
latents = initial_latent.float() * (1.0 - first_sigma) + latents.float() * first_sigma
if concat_latent is not None:
concat_latent = concat_latent.to(latents)
distilled_guidance = torch.tensor([distilled_guidance_scale * 1000.0] * batch_size).to(device=device, dtype=dtype)
prompt_embeds = repeat_to_batch_size(prompt_embeds, batch_size)
prompt_embeds_mask = repeat_to_batch_size(prompt_embeds_mask, batch_size)
prompt_poolers = repeat_to_batch_size(prompt_poolers, batch_size)
negative_prompt_embeds = repeat_to_batch_size(negative_prompt_embeds, batch_size)
negative_prompt_embeds_mask = repeat_to_batch_size(negative_prompt_embeds_mask, batch_size)
negative_prompt_poolers = repeat_to_batch_size(negative_prompt_poolers, batch_size)
concat_latent = repeat_to_batch_size(concat_latent, batch_size)
sampler_kwargs = dict(
dtype=dtype,
cfg_scale=real_guidance_scale,
cfg_rescale=guidance_rescale,
concat_latent=concat_latent,
positive=dict(
pooled_projections=prompt_poolers,
encoder_hidden_states=prompt_embeds,
encoder_attention_mask=prompt_embeds_mask,
guidance=distilled_guidance,
**kwargs,
),
negative=dict(
pooled_projections=negative_prompt_poolers,
encoder_hidden_states=negative_prompt_embeds,
encoder_attention_mask=negative_prompt_embeds_mask,
guidance=distilled_guidance,
**(kwargs if negative_kwargs is None else {**kwargs, **negative_kwargs}),
)
)
if sampler == 'unipc':
results = sample_unipc(k_model, latents, sigmas, extra_args=sampler_kwargs, disable=False, callback=callback)
else:
raise NotImplementedError(f'Sampler {sampler} is not supported.')
return results
================================================
FILE: diffusers_helper/thread_utils.py
================================================
import time
from threading import Thread, Lock
class Listener:
task_queue = []
lock = Lock()
thread = None
@classmethod
def _process_tasks(cls):
while True:
task = None
with cls.lock:
if cls.task_queue:
task = cls.task_queue.pop(0)
if task is None:
time.sleep(0.001)
continue
func, args, kwargs = task
try:
func(*args, **kwargs)
except Exception as e:
print(f"Error in listener thread: {e}")
@classmethod
def add_task(cls, func, *args, **kwargs):
with cls.lock:
cls.task_queue.append((func, args, kwargs))
if cls.thread is None:
cls.thread = Thread(target=cls._process_tasks, daemon=True)
cls.thread.start()
def async_run(func, *args, **kwargs):
Listener.add_task(func, *args, **kwargs)
class FIFOQueue:
def __init__(self):
self.queue = []
self.lock = Lock()
def push(self, item):
with self.lock:
self.queue.append(item)
def pop(self):
with self.lock:
if self.queue:
return self.queue.pop(0)
return None
def top(self):
with self.lock:
if self.queue:
return self.queue[0]
return None
def next(self):
while True:
with self.lock:
if self.queue:
return self.queue.pop(0)
time.sleep(0.001)
class AsyncStream:
def __init__(self):
self.input_queue = FIFOQueue()
self.output_queue = FIFOQueue()
================================================
FILE: diffusers_helper/utils.py
================================================
import os
import cv2
import json
import random
import glob
import torch
import einops
import numpy as np
import datetime
import torchvision
import safetensors.torch as sf
from PIL import Image
def min_resize(x, m):
if x.shape[0] < x.shape[1]:
s0 = m
s1 = int(float(m) / float(x.shape[0]) * float(x.shape[1]))
else:
s0 = int(float(m) / float(x.shape[1]) * float(x.shape[0]))
s1 = m
new_max = max(s1, s0)
raw_max = max(x.shape[0], x.shape[1])
if new_max < raw_max:
interpolation = cv2.INTER_AREA
else:
interpolation = cv2.INTER_LANCZOS4
y = cv2.resize(x, (s1, s0), interpolation=interpolation)
return y
def d_resize(x, y):
H, W, C = y.shape
new_min = min(H, W)
raw_min = min(x.shape[0], x.shape[1])
if new_min < raw_min:
interpolation = cv2.INTER_AREA
else:
interpolation = cv2.INTER_LANCZOS4
y = cv2.resize(x, (W, H), interpolation=interpolation)
return y
def resize_and_center_crop(image, target_width, target_height):
if target_height == image.shape[0] and target_width == image.shape[1]:
return image
pil_image = Image.fromarray(image)
original_width, original_height = pil_image.size
scale_factor = max(target_width / original_width, target_height / original_height)
resized_width = int(round(original_width * scale_factor))
resized_height = int(round(original_height * scale_factor))
resized_image = pil_image.resize((resized_width, resized_height), Image.LANCZOS)
left = (resized_width - target_width) / 2
top = (resized_height - target_height) / 2
right = (resized_width + target_width) / 2
bottom = (resized_height + target_height) / 2
cropped_image = resized_image.crop((left, top, right, bottom))
return np.array(cropped_image)
def resize_and_center_crop_pytorch(image, target_width, target_height):
B, C, H, W = image.shape
if H == target_height and W == target_width:
return image
scale_factor = max(target_width / W, target_height / H)
resized_width = int(round(W * scale_factor))
resized_height = int(round(H * scale_factor))
resized = torch.nn.functional.interpolate(image, size=(resized_height, resized_width), mode='bilinear', align_corners=False)
top = (resized_height - target_height) // 2
left = (resized_width - target_width) // 2
cropped = resized[:, :, top:top + target_height, left:left + target_width]
return cropped
def resize_without_crop(image, target_width, target_height):
if target_height == image.shape[0] and target_width == image.shape[1]:
return image
pil_image = Image.fromarray(image)
resized_image = pil_image.resize((target_width, target_height), Image.LANCZOS)
return np.array(resized_image)
def just_crop(image, w, h):
if h == image.shape[0] and w == image.shape[1]:
return image
original_height, original_width = image.shape[:2]
k = min(original_height / h, original_width / w)
new_width = int(round(w * k))
new_height = int(round(h * k))
x_start = (original_width - new_width) // 2
y_start = (original_height - new_height) // 2
cropped_image = image[y_start:y_start + new_height, x_start:x_start + new_width]
return cropped_image
def write_to_json(data, file_path):
temp_file_path = file_path + ".tmp"
with open(temp_file_path, 'wt', encoding='utf-8') as temp_file:
json.dump(data, temp_file, indent=4)
os.replace(temp_file_path, file_path)
return
def read_from_json(file_path):
with open(file_path, 'rt', encoding='utf-8') as file:
data = json.load(file)
return data
def get_active_parameters(m):
return {k: v for k, v in m.named_parameters() if v.requires_grad}
def cast_training_params(m, dtype=torch.float32):
result = {}
for n, param in m.named_parameters():
if param.requires_grad:
param.data = param.to(dtype)
result[n] = param
return result
def separate_lora_AB(parameters, B_patterns=None):
parameters_normal = {}
parameters_B = {}
if B_patterns is None:
B_patterns = ['.lora_B.', '__zero__']
for k, v in parameters.items():
if any(B_pattern in k for B_pattern in B_patterns):
parameters_B[k] = v
else:
parameters_normal[k] = v
return parameters_normal, parameters_B
def set_attr_recursive(obj, attr, value):
attrs = attr.split(".")
for name in attrs[:-1]:
obj = getattr(obj, name)
setattr(obj, attrs[-1], value)
return
def print_tensor_list_size(tensors):
total_size = 0
total_elements = 0
if isinstance(tensors, dict):
tensors = tensors.values()
for tensor in tensors:
total_size += tensor.nelement() * tensor.element_size()
total_elements += tensor.nelement()
total_size_MB = total_size / (1024 ** 2)
total_elements_B = total_elements / 1e9
print(f"Total number of tensors: {len(tensors)}")
print(f"Total size of tensors: {total_size_MB:.2f} MB")
print(f"Total number of parameters: {total_elements_B:.3f} billion")
return
@torch.no_grad()
def batch_mixture(a, b=None, probability_a=0.5, mask_a=None):
batch_size = a.size(0)
if b is None:
b = torch.zeros_like(a)
if mask_a is None:
mask_a = torch.rand(batch_size) < probability_a
mask_a = mask_a.to(a.device)
mask_a = mask_a.reshape((batch_size,) + (1,) * (a.dim() - 1))
result = torch.where(mask_a, a, b)
return result
@torch.no_grad()
def zero_module(module):
for p in module.parameters():
p.detach().zero_()
return module
@torch.no_grad()
def supress_lower_channels(m, k, alpha=0.01):
data = m.weight.data.clone()
assert int(data.shape[1]) >= k
data[:, :k] = data[:, :k] * alpha
m.weight.data = data.contiguous().clone()
return m
def freeze_module(m):
if not hasattr(m, '_forward_inside_frozen_module'):
m._forward_inside_frozen_module = m.forward
m.requires_grad_(False)
m.forward = torch.no_grad()(m.forward)
return m
def get_latest_safetensors(folder_path):
safetensors_files = glob.glob(os.path.join(folder_path, '*.safetensors'))
if not safetensors_files:
raise ValueError('No file to resume!')
latest_file = max(safetensors_files, key=os.path.getmtime)
latest_file = os.path.abspath(os.path.realpath(latest_file))
return latest_file
def generate_random_prompt_from_tags(tags_str, min_length=3, max_length=32):
tags = tags_str.split(', ')
tags = random.sample(tags, k=min(random.randint(min_length, max_length), len(tags)))
prompt = ', '.join(tags)
return prompt
def interpolate_numbers(a, b, n, round_to_int=False, gamma=1.0):
numbers = a + (b - a) * (np.linspace(0, 1, n) ** gamma)
if round_to_int:
numbers = np.round(numbers).astype(int)
return numbers.tolist()
def uniform_random_by_intervals(inclusive, exclusive, n, round_to_int=False):
edges = np.linspace(0, 1, n + 1)
points = np.random.uniform(edges[:-1], edges[1:])
numbers = inclusive + (exclusive - inclusive) * points
if round_to_int:
numbers = np.round(numbers).astype(int)
return numbers.tolist()
def soft_append_bcthw(history, current, overlap=0):
if overlap <= 0:
return torch.cat([history, current], dim=2)
assert history.shape[2] >= overlap, f"History length ({history.shape[2]}) must be >= overlap ({overlap})"
assert current.shape[2] >= overlap, f"Current length ({current.shape[2]}) must be >= overlap ({overlap})"
weights = torch.linspace(1, 0, overlap, dtype=history.dtype, device=history.device).view(1, 1, -1, 1, 1)
blended = weights * history[:, :, -overlap:] + (1 - weights) * current[:, :, :overlap]
output = torch.cat([history[:, :, :-overlap], blended, current[:, :, overlap:]], dim=2)
return output.to(history)
def save_bcthw_as_mp4(x, output_filename, fps=10, crf=0):
b, c, t, h, w = x.shape
per_row = b
for p in [6, 5, 4, 3, 2]:
if b % p == 0:
per_row = p
break
os.makedirs(os.path.dirname(os.path.abspath(os.path.realpath(output_filename))), exist_ok=True)
x = torch.clamp(x.float(), -1., 1.) * 127.5 + 127.5
x = x.detach().cpu().to(torch.uint8)
x = einops.rearrange(x, '(m n) c t h w -> t (m h) (n w) c', n=per_row)
torchvision.io.write_video(output_filename, x, fps=fps, video_codec='libx264', options={'crf': str(int(crf))})
return x
def save_bcthw_as_png(x, output_filename):
os.makedirs(os.path.dirname(os.path.abspath(os.path.realpath(output_filename))), exist_ok=True)
x = torch.clamp(x.float(), -1., 1.) * 127.5 + 127.5
x = x.detach().cpu().to(torch.uint8)
x = einops.rearrange(x, 'b c t h w -> c (b h) (t w)')
torchvision.io.write_png(x, output_filename)
return output_filename
def save_bchw_as_png(x, output_filename):
os.makedirs(os.path.dirname(os.path.abspath(os.path.realpath(output_filename))), exist_ok=True)
x = torch.clamp(x.float(), -1., 1.) * 127.5 + 127.5
x = x.detach().cpu().to(torch.uint8)
x = einops.rearrange(x, 'b c h w -> c h (b w)')
torchvision.io.write_png(x, output_filename)
return output_filename
def add_tensors_with_padding(tensor1, tensor2):
if tensor1.shape == tensor2.shape:
return tensor1 + tensor2
shape1 = tensor1.shape
shape2 = tensor2.shape
new_shape = tuple(max(s1, s2) for s1, s2 in zip(shape1, shape2))
padded_tensor1 = torch.zeros(new_shape)
padded_tensor2 = torch.zeros(new_shape)
padded_tensor1[tuple(slice(0, s) for s in shape1)] = tensor1
padded_tensor2[tuple(slice(0, s) for s in shape2)] = tensor2
result = padded_tensor1 + padded_tensor2
return result
def print_free_mem():
torch.cuda.empty_cache()
free_mem, total_mem = torch.cuda.mem_get_info(0)
free_mem_mb = free_mem / (1024 ** 2)
total_mem_mb = total_mem / (1024 ** 2)
print(f"Free memory: {free_mem_mb:.2f} MB")
print(f"Total memory: {total_mem_mb:.2f} MB")
return
def print_gpu_parameters(device, state_dict, log_count=1):
summary = {"device": device, "keys_count": len(state_dict)}
logged_params = {}
for i, (key, tensor) in enumerate(state_dict.items()):
if i >= log_count:
break
logged_params[key] = tensor.flatten()[:3].tolist()
summary["params"] = logged_params
print(str(summary))
return
def visualize_txt_as_img(width, height, text, font_path='font/DejaVuSans.ttf', size=18):
from PIL import Image, ImageDraw, ImageFont
txt = Image.new("RGB", (width, height), color="white")
draw = ImageDraw.Draw(txt)
font = ImageFont.truetype(font_path, size=size)
if text == '':
return np.array(txt)
# Split text into lines that fit within the image width
lines = []
words = text.split()
current_line = words[0]
for word in words[1:]:
line_with_word = f"{current_line} {word}"
if draw.textbbox((0, 0), line_with_word, font=font)[2] <= width:
current_line = line_with_word
else:
lines.append(current_line)
current_line = word
lines.append(current_line)
# Draw the text line by line
y = 0
line_height = draw.textbbox((0, 0), "A", font=font)[3]
for line in lines:
if y + line_height > height:
break # stop drawing if the next line will be outside the image
draw.text((0, y), line, fill="black", font=font)
y += line_height
return np.array(txt)
def blue_mark(x):
x = x.copy()
c = x[:, :, 2]
b = cv2.blur(c, (9, 9))
x[:, :, 2] = ((c - b) * 16.0 + b).clip(-1, 1)
return x
def green_mark(x):
x = x.copy()
x[:, :, 2] = -1
x[:, :, 0] = -1
return x
def frame_mark(x):
x = x.copy()
x[:64] = -1
x[-64:] = -1
x[:, :8] = 1
x[:, -8:] = 1
return x
@torch.inference_mode()
def pytorch2numpy(imgs):
results = []
for x in imgs:
y = x.movedim(0, -1)
y = y * 127.5 + 127.5
y = y.detach().float().cpu().numpy().clip(0, 255).astype(np.uint8)
results.append(y)
return results
@torch.inference_mode()
def numpy2pytorch(imgs):
h = torch.from_numpy(np.stack(imgs, axis=0)).float() / 127.5 - 1.0
h = h.movedim(-1, 1)
return h
@torch.no_grad()
def duplicate_prefix_to_suffix(x, count, zero_out=False):
if zero_out:
return torch.cat([x, torch.zeros_like(x[:count])], dim=0)
else:
return torch.cat([x, x[:count]], dim=0)
def weighted_mse(a, b, weight):
return torch.mean(weight.float() * (a.float() - b.float()) ** 2)
def clamped_linear_interpolation(x, x_min, y_min, x_max, y_max, sigma=1.0):
x = (x - x_min) / (x_max - x_min)
x = max(0.0, min(x, 1.0))
x = x ** sigma
return y_min + x * (y_max - y_min)
def expand_to_dims(x, target_dims):
return x.view(*x.shape, *([1] * max(0, target_dims - x.dim())))
def repeat_to_batch_size(tensor: torch.Tensor, batch_size: int):
if tensor is None:
return None
first_dim = tensor.shape[0]
if first_dim == batch_size:
return tensor
if batch_size % first_dim != 0:
raise ValueError(f"Cannot evenly repeat first dim {first_dim} to match batch_size {batch_size}.")
repeat_times = batch_size // first_dim
return tensor.repeat(repeat_times, *[1] * (tensor.dim() - 1))
def dim5(x):
return expand_to_dims(x, 5)
def dim4(x):
return expand_to_dims(x, 4)
def dim3(x):
return expand_to_dims(x, 3)
def crop_or_pad_yield_mask(x, length):
B, F, C = x.shape
device = x.device
dtype = x.dtype
if F < length:
y = torch.zeros((B, length, C), dtype=dtype, device=device)
mask = torch.zeros((B, length), dtype=torch.bool, device=device)
y[:, :F, :] = x
mask[:, :F] = True
return y, mask
return x[:, :length, :], torch.ones((B, length), dtype=torch.bool, device=device)
def extend_dim(x, dim, minimal_length, zero_pad=False):
original_length = int(x.shape[dim])
if original_length >= minimal_length:
return x
if zero_pad:
padding_shape = list(x.shape)
padding_shape[dim] = minimal_length - original_length
padding = torch.zeros(padding_shape, dtype=x.dtype, device=x.device)
else:
idx = (slice(None),) * dim + (slice(-1, None),) + (slice(None),) * (len(x.shape) - dim - 1)
last_element = x[idx]
padding = last_element.repeat_interleave(minimal_length - original_length, dim=dim)
return torch.cat([x, padding], dim=dim)
def lazy_positional_encoding(t, repeats=None):
if not isinstance(t, list):
t = [t]
from diffusers.models.embeddings import get_timestep_embedding
te = torch.tensor(t)
te = get_timestep_embedding(timesteps=te, embedding_dim=256, flip_sin_to_cos=True, downscale_freq_shift=0.0, scale=1.0)
if repeats is None:
return te
te = te[:, None, :].expand(-1, repeats, -1)
return te
def state_dict_offset_merge(A, B, C=None):
result = {}
keys = A.keys()
for key in keys:
A_value = A[key]
B_value = B[key].to(A_value)
if C is None:
result[key] = A_value + B_value
else:
C_value = C[key].to(A_value)
result[key] = A_value + B_value - C_value
return result
def state_dict_weighted_merge(state_dicts, weights):
if len(state_dicts) != len(weights):
raise ValueError("Number of state dictionaries must match number of weights")
if not state_dicts:
return {}
total_weight = sum(weights)
if total_weight == 0:
raise ValueError("Sum of weights cannot be zero")
normalized_weights = [w / total_weight for w in weights]
keys = state_dicts[0].keys()
result = {}
for key in keys:
result[key] = state_dicts[0][key] * normalized_weights[0]
for i in range(1, len(state_dicts)):
state_dict_value = state_dicts[i][key].to(result[key])
result[key] += state_dict_value * normalized_weights[i]
return result
def group_files_by_folder(all_files):
grouped_files = {}
for file in all_files:
folder_name = os.path.basename(os.path.dirname(file))
if folder_name not in grouped_files:
grouped_files[folder_name] = []
grouped_files[folder_name].append(file)
list_of_lists = list(grouped_files.values())
return list_of_lists
def generate_timestamp():
now = datetime.datetime.now()
timestamp = now.strftime('%y%m%d_%H%M%S')
milliseconds = f"{int(now.microsecond / 1000):03d}"
random_number = random.randint(0, 9999)
return f"{timestamp}_{milliseconds}_{random_number}"
def write_PIL_image_with_png_info(image, metadata, path):
from PIL.PngImagePlugin import PngInfo
png_info = PngInfo()
for key, value in metadata.items():
png_info.add_text(key, value)
image.save(path, "PNG", pnginfo=png_info)
return image
def torch_safe_save(content, path):
torch.save(content, path + '_tmp')
os.replace(path + '_tmp', path)
return path
def move_optimizer_to_device(optimizer, device):
for state in optimizer.state.values():
for k, v in state.items():
if isinstance(v, torch.Tensor):
state[k] = v.to(device)
================================================
FILE: requirements.txt
================================================
accelerate==1.6.0
diffusers==0.33.1
transformers==4.46.2
gradio==5.23.0
sentencepiece==0.2.0
pillow==11.1.0
av==12.1.0
numpy==1.26.2
scipy==1.12.0
requests==2.31.0
torchsde==0.2.6
einops
opencv-contrib-python
safetensors
gitextract_wnt3hu2x/ ├── .gitignore ├── LICENSE ├── README.md ├── demo_gradio.py ├── demo_gradio_f1.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
SYMBOL INDEX (173 symbols across 15 files)
FILE: demo_gradio.py
function worker (line 103) | def worker(input_image, prompt, n_prompt, seed, total_second_length, lat...
function process (line 318) | def process(input_image, prompt, n_prompt, seed, total_second_length, la...
function end_process (line 346) | def end_process():
FILE: demo_gradio_f1.py
function worker (line 103) | def worker(input_image, prompt, n_prompt, seed, total_second_length, lat...
function process (line 300) | def process(input_image, prompt, n_prompt, seed, total_second_length, la...
function end_process (line 328) | def end_process():
FILE: diffusers_helper/bucket_tools.py
function find_nearest_bucket (line 21) | def find_nearest_bucket(h, w, resolution=640):
FILE: diffusers_helper/clip_vision.py
function hf_clip_vision_encode (line 4) | def hf_clip_vision_encode(image, feature_extractor, image_encoder):
FILE: diffusers_helper/dit_common.py
function LayerNorm_forward (line 10) | def LayerNorm_forward(self, x):
function FP32LayerNorm_forward (line 18) | def FP32LayerNorm_forward(self, x):
function RMSNorm_forward (line 32) | def RMSNorm_forward(self, hidden_states):
function AdaLayerNormContinuous_forward (line 46) | def AdaLayerNormContinuous_forward(self, x, conditioning_embedding):
FILE: diffusers_helper/gradio/progress_bar.py
function make_progress_bar_html (line 81) | def make_progress_bar_html(number, text):
function make_progress_bar_css (line 85) | def make_progress_bar_css():
FILE: diffusers_helper/hf_login.py
function login (line 4) | def login(token):
FILE: diffusers_helper/hunyuan.py
function encode_prompt_conds (line 8) | def encode_prompt_conds(prompt, text_encoder, text_encoder_2, tokenizer,...
function vae_decode_fake (line 62) | def vae_decode_fake(latents):
function vae_decode (line 94) | def vae_decode(latents, vae, image_mode=False):
function vae_encode (line 108) | def vae_encode(image, vae):
FILE: diffusers_helper/k_diffusion/uni_pc_fm.py
function expand_dims (line 12) | def expand_dims(v, dims):
class FlowMatchUniPC (line 16) | class FlowMatchUniPC:
method __init__ (line 17) | def __init__(self, model, extra_args, variant='bh1'):
method model_fn (line 22) | def model_fn(self, x, t):
method update_fn (line 25) | def update_fn(self, x, model_prev_list, t_prev_list, t, order):
method sample (line 111) | def sample(self, x, sigmas, callback=None, disable_pbar=False):
function sample_unipc (line 139) | def sample_unipc(model, noise, sigmas, extra_args=None, callback=None, d...
FILE: diffusers_helper/k_diffusion/wrapper.py
function append_dims (line 4) | def append_dims(x, target_dims):
function rescale_noise_cfg (line 8) | def rescale_noise_cfg(noise_cfg, noise_pred_text, guidance_rescale=1.0):
function fm_wrapper (line 19) | def fm_wrapper(transformer, t_scale=1000.0):
FILE: diffusers_helper/memory.py
class DynamicSwapInstaller (line 12) | class DynamicSwapInstaller:
method _install_module (line 14) | def _install_module(module: torch.nn.Module, **kwargs):
method _uninstall_module (line 42) | def _uninstall_module(module: torch.nn.Module):
method install_model (line 48) | def install_model(model: torch.nn.Module, **kwargs):
method uninstall_model (line 54) | def uninstall_model(model: torch.nn.Module):
function fake_diffusers_current_device (line 60) | def fake_diffusers_current_device(model: torch.nn.Module, target_device:...
function get_cuda_free_memory_gb (line 71) | def get_cuda_free_memory_gb(device=None):
function move_model_to_device_with_memory_preservation (line 84) | def move_model_to_device_with_memory_preservation(model, target_device, ...
function offload_model_from_device_for_memory_preservation (line 100) | def offload_model_from_device_for_memory_preservation(model, target_devi...
function unload_complete_models (line 116) | def unload_complete_models(*args):
function load_model_as_complete (line 126) | def load_model_as_complete(model, target_device, unload=True):
FILE: diffusers_helper/models/hunyuan_video_packed.py
function pad_for_3d_conv (line 64) | def pad_for_3d_conv(x, kernel_size):
function center_down_sample_3d (line 73) | def center_down_sample_3d(x, kernel_size):
function get_cu_seqlens (line 82) | def get_cu_seqlens(text_mask, img_len):
function apply_rotary_emb_transposed (line 99) | def apply_rotary_emb_transposed(x, freqs_cis):
function attn_varlen_func (line 108) | def attn_varlen_func(q, k, v, cu_seqlens_q, cu_seqlens_kv, max_seqlen_q,...
class HunyuanAttnProcessorFlashAttnDouble (line 143) | class HunyuanAttnProcessorFlashAttnDouble:
method __call__ (line 144) | def __call__(self, attn, hidden_states, encoder_hidden_states, attenti...
class HunyuanAttnProcessorFlashAttnSingle (line 189) | class HunyuanAttnProcessorFlashAttnSingle:
method __call__ (line 190) | def __call__(self, attn, hidden_states, encoder_hidden_states, attenti...
class CombinedTimestepGuidanceTextProjEmbeddings (line 219) | class CombinedTimestepGuidanceTextProjEmbeddings(nn.Module):
method __init__ (line 220) | def __init__(self, embedding_dim, pooled_projection_dim):
method forward (line 228) | def forward(self, timestep, guidance, pooled_projection):
class CombinedTimestepTextProjEmbeddings (line 243) | class CombinedTimestepTextProjEmbeddings(nn.Module):
method __init__ (line 244) | def __init__(self, embedding_dim, pooled_projection_dim):
method forward (line 251) | def forward(self, timestep, pooled_projection):
class HunyuanVideoAdaNorm (line 262) | class HunyuanVideoAdaNorm(nn.Module):
method __init__ (line 263) | def __init__(self, in_features: int, out_features: Optional[int] = Non...
method forward (line 270) | def forward(
class HunyuanVideoIndividualTokenRefinerBlock (line 279) | class HunyuanVideoIndividualTokenRefinerBlock(nn.Module):
method __init__ (line 280) | def __init__(
method forward (line 306) | def forward(
class HunyuanVideoIndividualTokenRefiner (line 329) | class HunyuanVideoIndividualTokenRefiner(nn.Module):
method __init__ (line 330) | def __init__(
method forward (line 354) | def forward(
class HunyuanVideoTokenRefiner (line 376) | class HunyuanVideoTokenRefiner(nn.Module):
method __init__ (line 377) | def __init__(
method forward (line 404) | def forward(
class HunyuanVideoRotaryPosEmbed (line 425) | class HunyuanVideoRotaryPosEmbed(nn.Module):
method __init__ (line 426) | def __init__(self, rope_dim, theta):
method get_frequency (line 432) | def get_frequency(self, dim, pos):
method forward_inner (line 439) | def forward_inner(self, frame_indices, height, width, device):
method forward (line 456) | def forward(self, frame_indices, height, width, device):
class AdaLayerNormZero (line 463) | class AdaLayerNormZero(nn.Module):
method __init__ (line 464) | def __init__(self, embedding_dim: int, norm_type="layer_norm", bias=Tr...
method forward (line 473) | def forward(
class AdaLayerNormZeroSingle (line 485) | class AdaLayerNormZeroSingle(nn.Module):
method __init__ (line 486) | def __init__(self, embedding_dim: int, norm_type="layer_norm", bias=Tr...
method forward (line 496) | def forward(
class AdaLayerNormContinuous (line 508) | class AdaLayerNormContinuous(nn.Module):
method __init__ (line 509) | def __init__(
method forward (line 526) | def forward(self, x: torch.Tensor, emb: torch.Tensor) -> torch.Tensor:
class HunyuanVideoSingleTransformerBlock (line 534) | class HunyuanVideoSingleTransformerBlock(nn.Module):
method __init__ (line 535) | def __init__(
method forward (line 565) | def forward(
class HunyuanVideoTransformerBlock (line 608) | class HunyuanVideoTransformerBlock(nn.Module):
method __init__ (line 609) | def __init__(
method forward (line 643) | def forward(
class ClipVisionProjection (line 683) | class ClipVisionProjection(nn.Module):
method __init__ (line 684) | def __init__(self, in_channels, out_channels):
method forward (line 689) | def forward(self, x):
class HunyuanVideoPatchEmbed (line 694) | class HunyuanVideoPatchEmbed(nn.Module):
method __init__ (line 695) | def __init__(self, patch_size, in_chans, embed_dim):
class HunyuanVideoPatchEmbedForCleanLatents (line 700) | class HunyuanVideoPatchEmbedForCleanLatents(nn.Module):
method __init__ (line 701) | def __init__(self, inner_dim):
method initialize_weight_from_another_conv3d (line 708) | def initialize_weight_from_another_conv3d(self, another_layer):
class HunyuanVideoTransformer3DModelPacked (line 727) | class HunyuanVideoTransformer3DModelPacked(ModelMixin, ConfigMixin, Peft...
method __init__ (line 729) | def __init__(
method install_image_projection (line 805) | def install_image_projection(self, in_channels):
method install_clean_x_embedder (line 810) | def install_clean_x_embedder(self):
method enable_gradient_checkpointing (line 814) | def enable_gradient_checkpointing(self):
method disable_gradient_checkpointing (line 818) | def disable_gradient_checkpointing(self):
method initialize_teacache (line 822) | def initialize_teacache(self, enable_teacache=True, num_steps=25, rel_...
method gradient_checkpointing_method (line 832) | def gradient_checkpointing_method(self, block, *args):
method process_input_hidden_states (line 839) | def process_input_hidden_states(
method forward (line 898) | def forward(
FILE: diffusers_helper/pipelines/k_diffusion_hunyuan.py
function flux_time_shift (line 9) | def flux_time_shift(t, mu=1.15, sigma=1.0):
function calculate_flux_mu (line 13) | def calculate_flux_mu(context_length, x1=256, y1=0.5, x2=4096, y2=1.15, ...
function get_flux_sigmas_from_mu (line 21) | def get_flux_sigmas_from_mu(n, mu):
function sample_hunyuan (line 28) | def sample_hunyuan(
FILE: diffusers_helper/thread_utils.py
class Listener (line 6) | class Listener:
method _process_tasks (line 12) | def _process_tasks(cls):
method add_task (line 30) | def add_task(cls, func, *args, **kwargs):
function async_run (line 39) | def async_run(func, *args, **kwargs):
class FIFOQueue (line 43) | class FIFOQueue:
method __init__ (line 44) | def __init__(self):
method push (line 48) | def push(self, item):
method pop (line 52) | def pop(self):
method top (line 58) | def top(self):
method next (line 64) | def next(self):
class AsyncStream (line 73) | class AsyncStream:
method __init__ (line 74) | def __init__(self):
FILE: diffusers_helper/utils.py
function min_resize (line 16) | def min_resize(x, m):
function d_resize (line 33) | def d_resize(x, y):
function resize_and_center_crop (line 45) | def resize_and_center_crop(image, target_width, target_height):
function resize_and_center_crop_pytorch (line 63) | def resize_and_center_crop_pytorch(image, target_width, target_height):
function resize_without_crop (line 82) | def resize_without_crop(image, target_width, target_height):
function just_crop (line 91) | def just_crop(image, w, h):
function write_to_json (line 105) | def write_to_json(data, file_path):
function read_from_json (line 113) | def read_from_json(file_path):
function get_active_parameters (line 119) | def get_active_parameters(m):
function cast_training_params (line 123) | def cast_training_params(m, dtype=torch.float32):
function separate_lora_AB (line 132) | def separate_lora_AB(parameters, B_patterns=None):
function set_attr_recursive (line 148) | def set_attr_recursive(obj, attr, value):
function print_tensor_list_size (line 156) | def print_tensor_list_size(tensors):
function batch_mixture (line 177) | def batch_mixture(a, b=None, probability_a=0.5, mask_a=None):
function zero_module (line 193) | def zero_module(module):
function supress_lower_channels (line 200) | def supress_lower_channels(m, k, alpha=0.01):
function freeze_module (line 210) | def freeze_module(m):
function get_latest_safetensors (line 218) | def get_latest_safetensors(folder_path):
function generate_random_prompt_from_tags (line 229) | def generate_random_prompt_from_tags(tags_str, min_length=3, max_length=...
function interpolate_numbers (line 236) | def interpolate_numbers(a, b, n, round_to_int=False, gamma=1.0):
function uniform_random_by_intervals (line 243) | def uniform_random_by_intervals(inclusive, exclusive, n, round_to_int=Fa...
function soft_append_bcthw (line 252) | def soft_append_bcthw(history, current, overlap=0):
function save_bcthw_as_mp4 (line 266) | def save_bcthw_as_mp4(x, output_filename, fps=10, crf=0):
function save_bcthw_as_png (line 283) | def save_bcthw_as_png(x, output_filename):
function save_bchw_as_png (line 292) | def save_bchw_as_png(x, output_filename):
function add_tensors_with_padding (line 301) | def add_tensors_with_padding(tensor1, tensor2):
function print_free_mem (line 320) | def print_free_mem():
function print_gpu_parameters (line 330) | def print_gpu_parameters(device, state_dict, log_count=1):
function visualize_txt_as_img (line 345) | def visualize_txt_as_img(width, height, text, font_path='font/DejaVuSans...
function blue_mark (line 383) | def blue_mark(x):
function green_mark (line 391) | def green_mark(x):
function frame_mark (line 398) | def frame_mark(x):
function pytorch2numpy (line 408) | def pytorch2numpy(imgs):
function numpy2pytorch (line 419) | def numpy2pytorch(imgs):
function duplicate_prefix_to_suffix (line 426) | def duplicate_prefix_to_suffix(x, count, zero_out=False):
function weighted_mse (line 433) | def weighted_mse(a, b, weight):
function clamped_linear_interpolation (line 437) | def clamped_linear_interpolation(x, x_min, y_min, x_max, y_max, sigma=1.0):
function expand_to_dims (line 444) | def expand_to_dims(x, target_dims):
function repeat_to_batch_size (line 448) | def repeat_to_batch_size(tensor: torch.Tensor, batch_size: int):
function dim5 (line 465) | def dim5(x):
function dim4 (line 469) | def dim4(x):
function dim3 (line 473) | def dim3(x):
function crop_or_pad_yield_mask (line 477) | def crop_or_pad_yield_mask(x, length):
function extend_dim (line 492) | def extend_dim(x, dim, minimal_length, zero_pad=False):
function lazy_positional_encoding (line 510) | def lazy_positional_encoding(t, repeats=None):
function state_dict_offset_merge (line 527) | def state_dict_offset_merge(A, B, C=None):
function state_dict_weighted_merge (line 544) | def state_dict_weighted_merge(state_dicts, weights):
function group_files_by_folder (line 571) | def group_files_by_folder(all_files):
function generate_timestamp (line 584) | def generate_timestamp():
function write_PIL_image_with_png_info (line 592) | def write_PIL_image_with_png_info(image, metadata, path):
function torch_safe_save (line 603) | def torch_safe_save(content, path):
function move_optimizer_to_device (line 609) | def move_optimizer_to_device(optimizer, device):
Condensed preview — 19 files, each showing path, character count, and a content snippet. Download the .json file or copy for the full structured content (159K chars).
[
{
"path": ".gitignore",
"chars": 3471,
"preview": "hf_download/\noutputs/\nrepo/\n\n# Byte-compiled / optimized / DLL files\n__pycache__/\n*.py[cod]\n*$py.class\n\n# C extensions\n*"
},
{
"path": "LICENSE",
"chars": 11357,
"preview": " Apache License\n Version 2.0, January 2004\n "
},
{
"path": "README.md",
"chars": 16558,
"preview": "<p align=\"center\">\n <img src=\"https://github.com/user-attachments/assets/2cc030b4-87e1-40a0-b5bf-1b7d6b62820b\" width="
},
{
"path": "demo_gradio.py",
"chars": 19500,
"preview": "from diffusers_helper.hf_login import login\n\nimport os\n\nos.environ['HF_HOME'] = os.path.abspath(os.path.realpath(os.path"
},
{
"path": "demo_gradio_f1.py",
"chars": 18355,
"preview": "from diffusers_helper.hf_login import login\n\nimport os\n\nos.environ['HF_HOME'] = os.path.abspath(os.path.realpath(os.path"
},
{
"path": "diffusers_helper/bucket_tools.py",
"chars": 667,
"preview": "bucket_options = {\n 640: [\n (416, 960),\n (448, 864),\n (480, 832),\n (512, 768),\n (5"
},
{
"path": "diffusers_helper/clip_vision.py",
"chars": 449,
"preview": "import numpy as np\n\n\ndef hf_clip_vision_encode(image, feature_extractor, image_encoder):\n assert isinstance(image, np"
},
{
"path": "diffusers_helper/dit_common.py",
"chars": 1535,
"preview": "import torch\nimport accelerate.accelerator\n\nfrom diffusers.models.normalization import RMSNorm, LayerNorm, FP32LayerNorm"
},
{
"path": "diffusers_helper/gradio/progress_bar.py",
"chars": 2165,
"preview": "progress_html = '''\n<div class=\"loader-container\">\n <div class=\"loader\"></div>\n <div class=\"progress-container\">\n <"
},
{
"path": "diffusers_helper/hf_login.py",
"chars": 395,
"preview": "import os\n\n\ndef login(token):\n from huggingface_hub import login\n import time\n\n while True:\n try:\n "
},
{
"path": "diffusers_helper/hunyuan.py",
"chars": 3648,
"preview": "import torch\n\nfrom diffusers.pipelines.hunyuan_video.pipeline_hunyuan_video import DEFAULT_PROMPT_TEMPLATE\nfrom diffuser"
},
{
"path": "diffusers_helper/k_diffusion/uni_pc_fm.py",
"chars": 4369,
"preview": "# Better Flow Matching UniPC by Lvmin Zhang\n# (c) 2025\n# CC BY-SA 4.0\n# Attribution-ShareAlike 4.0 International Licence"
},
{
"path": "diffusers_helper/k_diffusion/wrapper.py",
"chars": 1821,
"preview": "import torch\n\n\ndef append_dims(x, target_dims):\n return x[(...,) + (None,) * (target_dims - x.ndim)]\n\n\ndef rescale_no"
},
{
"path": "diffusers_helper/memory.py",
"chars": 4322,
"preview": "# By lllyasviel\n\n\nimport torch\n\n\ncpu = torch.device('cpu')\ngpu = torch.device(f'cuda:{torch.cuda.current_device()}')\ngpu"
},
{
"path": "diffusers_helper/models/hunyuan_video_packed.py",
"chars": 41413,
"preview": "from typing import Any, Dict, List, Optional, Tuple, Union\n\nimport torch\nimport einops\nimport torch.nn as nn\nimport nump"
},
{
"path": "diffusers_helper/pipelines/k_diffusion_hunyuan.py",
"chars": 4118,
"preview": "import torch\nimport math\n\nfrom diffusers_helper.k_diffusion.uni_pc_fm import sample_unipc\nfrom diffusers_helper.k_diffus"
},
{
"path": "diffusers_helper/thread_utils.py",
"chars": 1754,
"preview": "import time\n\nfrom threading import Thread, Lock\n\n\nclass Listener:\n task_queue = []\n lock = Lock()\n thread = Non"
},
{
"path": "diffusers_helper/utils.py",
"chars": 17469,
"preview": "import os\nimport cv2\nimport json\nimport random\nimport glob\nimport torch\nimport einops\nimport numpy as np\nimport datetime"
},
{
"path": "requirements.txt",
"chars": 222,
"preview": "accelerate==1.6.0\ndiffusers==0.33.1\ntransformers==4.46.2\ngradio==5.23.0\nsentencepiece==0.2.0\npillow==11.1.0\nav==12.1.0\nn"
}
]
About this extraction
This page contains the full source code of the lllyasviel/FramePack GitHub repository, extracted and formatted as plain text for AI agents and large language models (LLMs). The extraction includes 19 files (150.0 KB), approximately 38.8k tokens, and a symbol index with 173 extracted functions, classes, methods, constants, and types. Use this with OpenClaw, Claude, ChatGPT, Cursor, Windsurf, or any other AI tool that accepts text input. You can copy the full output to your clipboard or download it as a .txt file.
Extracted by GitExtract — free GitHub repo to text converter for AI. Built by Nikandr Surkov.