├── .gitignore ├── FramePack-user-DEVICE0.bat ├── FramePack-user-DEVICE1.bat ├── FramePack-user.bat ├── FramePackF1-user.bat ├── 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 ├── install-win-zluda.bat └── requirements.txt /.gitignore: -------------------------------------------------------------------------------- 1 | hf_download/ 2 | outputs/ 3 | repo/ 4 | 5 | # Byte-compiled / optimized / DLL files 6 | __pycache__/ 7 | *.py[cod] 8 | *$py.class 9 | 10 | # C extensions 11 | *.so 12 | 13 | # Distribution / packaging 14 | .Python 15 | build/ 16 | develop-eggs/ 17 | dist/ 18 | downloads/ 19 | eggs/ 20 | .eggs/ 21 | lib/ 22 | lib64/ 23 | parts/ 24 | sdist/ 25 | var/ 26 | wheels/ 27 | share/python-wheels/ 28 | *.egg-info/ 29 | .installed.cfg 30 | *.egg 31 | MANIFEST 32 | 33 | # PyInstaller 34 | # Usually these files are written by a python script from a template 35 | # before PyInstaller builds the exe, so as to inject date/other infos into it. 36 | *.manifest 37 | *.spec 38 | 39 | # Installer logs 40 | pip-log.txt 41 | pip-delete-this-directory.txt 42 | 43 | # Unit test / coverage reports 44 | htmlcov/ 45 | .tox/ 46 | .nox/ 47 | .coverage 48 | .coverage.* 49 | .cache 50 | nosetests.xml 51 | coverage.xml 52 | *.cover 53 | *.py,cover 54 | .hypothesis/ 55 | .pytest_cache/ 56 | cover/ 57 | 58 | # Translations 59 | *.mo 60 | *.pot 61 | 62 | # Django stuff: 63 | *.log 64 | local_settings.py 65 | db.sqlite3 66 | db.sqlite3-journal 67 | 68 | # Flask stuff: 69 | instance/ 70 | .webassets-cache 71 | 72 | # Scrapy stuff: 73 | .scrapy 74 | 75 | # Sphinx documentation 76 | docs/_build/ 77 | 78 | # PyBuilder 79 | .pybuilder/ 80 | target/ 81 | 82 | # Jupyter Notebook 83 | .ipynb_checkpoints 84 | 85 | # IPython 86 | profile_default/ 87 | ipython_config.py 88 | 89 | # pyenv 90 | # For a library or package, you might want to ignore these files since the code is 91 | # intended to run in multiple environments; otherwise, check them in: 92 | # .python-version 93 | 94 | # pipenv 95 | # According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control. 96 | # However, in case of collaboration, if having platform-specific dependencies or dependencies 97 | # having no cross-platform support, pipenv may install dependencies that don't work, or not 98 | # install all needed dependencies. 99 | #Pipfile.lock 100 | 101 | # UV 102 | # Similar to Pipfile.lock, it is generally recommended to include uv.lock in version control. 103 | # This is especially recommended for binary packages to ensure reproducibility, and is more 104 | # commonly ignored for libraries. 105 | #uv.lock 106 | 107 | # poetry 108 | # Similar to Pipfile.lock, it is generally recommended to include poetry.lock in version control. 109 | # This is especially recommended for binary packages to ensure reproducibility, and is more 110 | # commonly ignored for libraries. 111 | # https://python-poetry.org/docs/basic-usage/#commit-your-poetrylock-file-to-version-control 112 | #poetry.lock 113 | 114 | # pdm 115 | # Similar to Pipfile.lock, it is generally recommended to include pdm.lock in version control. 116 | #pdm.lock 117 | # pdm stores project-wide configurations in .pdm.toml, but it is recommended to not include it 118 | # in version control. 119 | # https://pdm.fming.dev/latest/usage/project/#working-with-version-control 120 | .pdm.toml 121 | .pdm-python 122 | .pdm-build/ 123 | 124 | # PEP 582; used by e.g. github.com/David-OConnor/pyflow and github.com/pdm-project/pdm 125 | __pypackages__/ 126 | 127 | # Celery stuff 128 | celerybeat-schedule 129 | celerybeat.pid 130 | 131 | # SageMath parsed files 132 | *.sage.py 133 | 134 | # Environments 135 | .env 136 | .venv 137 | env/ 138 | venv/ 139 | ENV/ 140 | env.bak/ 141 | venv.bak/ 142 | 143 | # Spyder project settings 144 | .spyderproject 145 | .spyproject 146 | 147 | # Rope project settings 148 | .ropeproject 149 | 150 | # mkdocs documentation 151 | /site 152 | 153 | # mypy 154 | .mypy_cache/ 155 | .dmypy.json 156 | dmypy.json 157 | 158 | # Pyre type checker 159 | .pyre/ 160 | 161 | # pytype static type analyzer 162 | .pytype/ 163 | 164 | # Cython debug symbols 165 | cython_debug/ 166 | 167 | # PyCharm 168 | # JetBrains specific template is maintained in a separate JetBrains.gitignore that can 169 | # be found at https://github.com/github/gitignore/blob/main/Global/JetBrains.gitignore 170 | # and can be added to the global gitignore or merged into this file. For a more nuclear 171 | # option (not recommended) you can uncomment the following to ignore the entire idea folder. 172 | .idea/ 173 | 174 | # Ruff stuff: 175 | .ruff_cache/ 176 | 177 | # PyPI configuration file 178 | .pypirc 179 | -------------------------------------------------------------------------------- /FramePack-user-DEVICE0.bat: -------------------------------------------------------------------------------- 1 | @echo off 2 | 3 | set PATH=%~dp0.zluda;%HIP_PATH%bin;%PATH% 4 | set PYTHON="%~dp0venv/Scripts/python.exe" 5 | set GIT= 6 | set VENV_DIR=./venv 7 | set COMMANDLINE_ARGS= 8 | @REM set COMMANDLINE_ARGS= --inbrowser 9 | 10 | set DISABLE_ADDMM_CUDA_LT=1 11 | set ZLUDA_COMGR_LOG_LEVEL=1 12 | 13 | @REM set INCLUDE="%~dp0venv\include";%INCLUDE% 14 | @REM set ZLUDA_NVRTC_LIB="%~dp0venv\Lib\site-packages\torch\lib\nvrtc_cuda.dll" 15 | 16 | set HIP_VISIBLE_DEVICES=0 17 | @REM set HIP_VISIBLE_DEVICES=1 18 | 19 | %PYTHON% demo_gradio.py %COMMANDLINE_ARGS% 20 | 21 | pause 22 | 23 | -------------------------------------------------------------------------------- /FramePack-user-DEVICE1.bat: -------------------------------------------------------------------------------- 1 | @echo off 2 | 3 | set PATH=%~dp0.zluda;%HIP_PATH%bin;%PATH% 4 | set PYTHON="%~dp0venv/Scripts/python.exe" 5 | set GIT= 6 | set VENV_DIR=./venv 7 | set COMMANDLINE_ARGS= 8 | @REM set COMMANDLINE_ARGS= --inbrowser 9 | 10 | set DISABLE_ADDMM_CUDA_LT=1 11 | set ZLUDA_COMGR_LOG_LEVEL=1 12 | 13 | @REM set INCLUDE="%~dp0venv\include";%INCLUDE% 14 | @REM set ZLUDA_NVRTC_LIB="%~dp0venv\Lib\site-packages\torch\lib\nvrtc_cuda.dll" 15 | 16 | @REM set HIP_VISIBLE_DEVICES=0 17 | set HIP_VISIBLE_DEVICES=1 18 | 19 | %PYTHON% demo_gradio.py %COMMANDLINE_ARGS% 20 | 21 | pause 22 | 23 | -------------------------------------------------------------------------------- /FramePack-user.bat: -------------------------------------------------------------------------------- 1 | @echo off 2 | 3 | set PATH=%~dp0.zluda;%HIP_PATH%bin;%PATH% 4 | set PYTHON="%~dp0venv/Scripts/python.exe" 5 | set GIT= 6 | set VENV_DIR=./venv 7 | set COMMANDLINE_ARGS= 8 | @REM set COMMANDLINE_ARGS= --inbrowser 9 | 10 | set DISABLE_ADDMM_CUDA_LT=1 11 | set ZLUDA_COMGR_LOG_LEVEL=1 12 | 13 | @REM set INCLUDE="%~dp0venv\include";%INCLUDE% 14 | @REM set ZLUDA_NVRTC_LIB="%~dp0venv\Lib\site-packages\torch\lib\nvrtc_cuda.dll" 15 | 16 | %PYTHON% demo_gradio.py %COMMANDLINE_ARGS% 17 | 18 | pause 19 | 20 | -------------------------------------------------------------------------------- /FramePackF1-user.bat: -------------------------------------------------------------------------------- 1 | @echo off 2 | 3 | set PATH=%~dp0.zluda;%HIP_PATH%bin;%PATH% 4 | set PYTHON="%~dp0venv/Scripts/python.exe" 5 | set GIT= 6 | set VENV_DIR=./venv 7 | set COMMANDLINE_ARGS= 8 | @REM set COMMANDLINE_ARGS= --inbrowser 9 | 10 | set DISABLE_ADDMM_CUDA_LT=1 11 | set ZLUDA_COMGR_LOG_LEVEL=1 12 | 13 | @REM set INCLUDE="%~dp0venv\include";%INCLUDE% 14 | @REM set ZLUDA_NVRTC_LIB="%~dp0venv\Lib\site-packages\torch\lib\nvrtc_cuda.dll" 15 | 16 | @REM set HIP_VISIBLE_DEVICES=0 17 | @REM set HIP_VISIBLE_DEVICES=1 18 | 19 | %PYTHON% demo_gradio_f1.py %COMMANDLINE_ARGS% 20 | 21 | pause 22 | 23 | -------------------------------------------------------------------------------- /LICENSE: -------------------------------------------------------------------------------- 1 | Apache License 2 | Version 2.0, January 2004 3 | http://www.apache.org/licenses/ 4 | 5 | TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION 6 | 7 | 1. Definitions. 8 | 9 | "License" shall mean the terms and conditions for use, reproduction, 10 | and distribution as defined by Sections 1 through 9 of this document. 11 | 12 | "Licensor" shall mean the copyright owner or entity authorized by 13 | the copyright owner that is granting the License. 14 | 15 | "Legal Entity" shall mean the union of the acting entity and all 16 | other entities that control, are controlled by, or are under common 17 | control with that entity. For the purposes of this definition, 18 | "control" means (i) the power, direct or indirect, to cause the 19 | direction or management of such entity, whether by contract or 20 | otherwise, or (ii) ownership of fifty percent (50%) or more of the 21 | outstanding shares, or (iii) beneficial ownership of such entity. 22 | 23 | "You" (or "Your") shall mean an individual or Legal Entity 24 | exercising permissions granted by this License. 25 | 26 | "Source" form shall mean the preferred form for making modifications, 27 | including but not limited to software source code, documentation 28 | source, and configuration files. 29 | 30 | "Object" form shall mean any form resulting from mechanical 31 | transformation or translation of a Source form, including but 32 | not limited to compiled object code, generated documentation, 33 | and conversions to other media types. 34 | 35 | "Work" shall mean the work of authorship, whether in Source or 36 | Object form, made available under the License, as indicated by a 37 | copyright notice that is included in or attached to the work 38 | (an example is provided in the Appendix below). 39 | 40 | "Derivative Works" shall mean any work, whether in Source or Object 41 | form, that is based on (or derived from) the Work and for which the 42 | editorial revisions, annotations, elaborations, or other modifications 43 | represent, as a whole, an original work of authorship. For the purposes 44 | of this License, Derivative Works shall not include works that remain 45 | separable from, or merely link (or bind by name) to the interfaces of, 46 | the Work and Derivative Works thereof. 47 | 48 | "Contribution" shall mean any work of authorship, including 49 | the original version of the Work and any modifications or additions 50 | to that Work or Derivative Works thereof, that is intentionally 51 | submitted to Licensor for inclusion in the Work by the copyright owner 52 | or by an individual or Legal Entity authorized to submit on behalf of 53 | the copyright owner. For the purposes of this definition, "submitted" 54 | means any form of electronic, verbal, or written communication sent 55 | to the Licensor or its representatives, including but not limited to 56 | communication on electronic mailing lists, source code control systems, 57 | and issue tracking systems that are managed by, or on behalf of, the 58 | Licensor for the purpose of discussing and improving the Work, but 59 | excluding communication that is conspicuously marked or otherwise 60 | designated in writing by the copyright owner as "Not a Contribution." 61 | 62 | "Contributor" shall mean Licensor and any individual or Legal Entity 63 | on behalf of whom a Contribution has been received by Licensor and 64 | subsequently incorporated within the Work. 65 | 66 | 2. Grant of Copyright License. Subject to the terms and conditions of 67 | this License, each Contributor hereby grants to You a perpetual, 68 | worldwide, non-exclusive, no-charge, royalty-free, irrevocable 69 | copyright license to reproduce, prepare Derivative Works of, 70 | publicly display, publicly perform, sublicense, and distribute the 71 | Work and such Derivative Works in Source or Object form. 72 | 73 | 3. Grant of Patent License. Subject to the terms and conditions of 74 | this License, each Contributor hereby grants to You a perpetual, 75 | worldwide, non-exclusive, no-charge, royalty-free, irrevocable 76 | (except as stated in this section) patent license to make, have made, 77 | use, offer to sell, sell, import, and otherwise transfer the Work, 78 | where such license applies only to those patent claims licensable 79 | by such Contributor that are necessarily infringed by their 80 | Contribution(s) alone or by combination of their Contribution(s) 81 | with the Work to which such Contribution(s) was submitted. If You 82 | institute patent litigation against any entity (including a 83 | cross-claim or counterclaim in a lawsuit) alleging that the Work 84 | or a Contribution incorporated within the Work constitutes direct 85 | or contributory patent infringement, then any patent licenses 86 | granted to You under this License for that Work shall terminate 87 | as of the date such litigation is filed. 88 | 89 | 4. Redistribution. You may reproduce and distribute copies of the 90 | Work or Derivative Works thereof in any medium, with or without 91 | modifications, and in Source or Object form, provided that You 92 | meet the following conditions: 93 | 94 | (a) You must give any other recipients of the Work or 95 | Derivative Works a copy of this License; and 96 | 97 | (b) You must cause any modified files to carry prominent notices 98 | stating that You changed the files; and 99 | 100 | (c) You must retain, in the Source form of any Derivative Works 101 | that You distribute, all copyright, patent, trademark, and 102 | attribution notices from the Source form of the Work, 103 | excluding those notices that do not pertain to any part of 104 | the Derivative Works; and 105 | 106 | (d) If the Work includes a "NOTICE" text file as part of its 107 | distribution, then any Derivative Works that You distribute must 108 | include a readable copy of the attribution notices contained 109 | within such NOTICE file, excluding those notices that do not 110 | pertain to any part of the Derivative Works, in at least one 111 | of the following places: within a NOTICE text file distributed 112 | as part of the Derivative Works; within the Source form or 113 | documentation, if provided along with the Derivative Works; or, 114 | within a display generated by the Derivative Works, if and 115 | wherever such third-party notices normally appear. The contents 116 | of the NOTICE file are for informational purposes only and 117 | do not modify the License. You may add Your own attribution 118 | notices within Derivative Works that You distribute, alongside 119 | or as an addendum to the NOTICE text from the Work, provided 120 | that such additional attribution notices cannot be construed 121 | as modifying the License. 122 | 123 | You may add Your own copyright statement to Your modifications and 124 | may provide additional or different license terms and conditions 125 | for use, reproduction, or distribution of Your modifications, or 126 | for any such Derivative Works as a whole, provided Your use, 127 | reproduction, and distribution of the Work otherwise complies with 128 | the conditions stated in this License. 129 | 130 | 5. Submission of Contributions. Unless You explicitly state otherwise, 131 | any Contribution intentionally submitted for inclusion in the Work 132 | by You to the Licensor shall be under the terms and conditions of 133 | this License, without any additional terms or conditions. 134 | Notwithstanding the above, nothing herein shall supersede or modify 135 | the terms of any separate license agreement you may have executed 136 | with Licensor regarding such Contributions. 137 | 138 | 6. Trademarks. This License does not grant permission to use the trade 139 | names, trademarks, service marks, or product names of the Licensor, 140 | except as required for reasonable and customary use in describing the 141 | origin of the Work and reproducing the content of the NOTICE file. 142 | 143 | 7. Disclaimer of Warranty. Unless required by applicable law or 144 | agreed to in writing, Licensor provides the Work (and each 145 | Contributor provides its Contributions) on an "AS IS" BASIS, 146 | WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or 147 | implied, including, without limitation, any warranties or conditions 148 | of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A 149 | PARTICULAR PURPOSE. You are solely responsible for determining the 150 | appropriateness of using or redistributing the Work and assume any 151 | risks associated with Your exercise of permissions under this License. 152 | 153 | 8. Limitation of Liability. In no event and under no legal theory, 154 | whether in tort (including negligence), contract, or otherwise, 155 | unless required by applicable law (such as deliberate and grossly 156 | negligent acts) or agreed to in writing, shall any Contributor be 157 | liable to You for damages, including any direct, indirect, special, 158 | incidental, or consequential damages of any character arising as a 159 | result of this License or out of the use or inability to use the 160 | Work (including but not limited to damages for loss of goodwill, 161 | work stoppage, computer failure or malfunction, or any and all 162 | other commercial damages or losses), even if such Contributor 163 | has been advised of the possibility of such damages. 164 | 165 | 9. Accepting Warranty or Additional Liability. While redistributing 166 | the Work or Derivative Works thereof, You may choose to offer, 167 | and charge a fee for, acceptance of support, warranty, indemnity, 168 | or other liability obligations and/or rights consistent with this 169 | License. However, in accepting such obligations, You may act only 170 | on Your own behalf and on Your sole responsibility, not on behalf 171 | of any other Contributor, and only if You agree to indemnify, 172 | defend, and hold each Contributor harmless for any liability 173 | incurred by, or claims asserted against, such Contributor by reason 174 | of your accepting any such warranty or additional liability. 175 | 176 | END OF TERMS AND CONDITIONS 177 | 178 | APPENDIX: How to apply the Apache License to your work. 179 | 180 | To apply the Apache License to your work, attach the following 181 | boilerplate notice, with the fields enclosed by brackets "[]" 182 | replaced with your own identifying information. (Don't include 183 | the brackets!) The text should be enclosed in the appropriate 184 | comment syntax for the file format. We also recommend that a 185 | file or class name and description of purpose be included on the 186 | same "printed page" as the copyright notice for easier 187 | identification within third-party archives. 188 | 189 | Copyright [yyyy] [name of copyright owner] 190 | 191 | Licensed under the Apache License, Version 2.0 (the "License"); 192 | you may not use this file except in compliance with the License. 193 | You may obtain a copy of the License at 194 | 195 | http://www.apache.org/licenses/LICENSE-2.0 196 | 197 | Unless required by applicable law or agreed to in writing, software 198 | distributed under the License is distributed on an "AS IS" BASIS, 199 | WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 200 | See the License for the specific language governing permissions and 201 | limitations under the License. 202 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 |
2 | 3 | # FramePack-ZLUDA 4 | 5 | FramePack ZLUDA 6 | 7 |
8 | 9 | --- 10 | 11 | ## attention branch 12 | attention branch is available. Flash Attention 2 built-in. 13 | 14 | https://github.com/githubcto/FramePack-ZLUDA/tree/attention 15 | 16 | ## AMD VGA Driver Adrenalin 25.5.1 is NOT COMPATIBLE with ZLUDA 3.9.4. 17 | 18 | ZLUDA v3.9.5 support Adrenalin 25.5.1. 19 | 20 | remove .zluda folder, download v3.9.5 and replace DLLs. 21 | ``` 22 | curl -s -L https://github.com/lshqqytiger/ZLUDA/releases/download/rel.5e717459179dc272b7d7d23391f0fad66c7459cf/ZLUDA-windows-rocm6-amd64.zip > zluda.zip 23 | mkdir .zluda && tar -xf zluda.zip -C .zluda --strip-components=1 24 | del zluda.zip 25 | if not exist "%~dp0venv\Lib\site-packages\torch\lib\nvrtc_cuda.dll" copy venv\Lib\site-packages\torch\lib\nvrtc64_112_0.dll venv\Lib\site-packages\torch\lib\nvrtc_cuda.dll & REM (optional) 26 | copy .zluda\cublas.dll venv\Lib\site-packages\torch\lib\cublas64_11.dll /y 27 | copy .zluda\cusparse.dll venv\Lib\site-packages\torch\lib\cusparse64_11.dll /y 28 | copy .zluda\nvrtc.dll venv\Lib\site-packages\torch\lib\nvrtc64_112_0.dll /y 29 | echo Finished 30 | ``` 31 | optional: remove zluda.db (C:\Users\userID\AppData\Local\ZLUDA\ComputeCache\zluda.db) . 32 | 33 | This remove is optional, I recommend this when driver/torch/zluda change. 34 | 35 | ## Install 36 | 37 | - install Microsoft Visual C++ Redistributable from [microsoft.com](https://learn.microsoft.com/en-us/cpp/windows/latest-supported-vc-redist?view=msvc-170) . 38 | 39 | https://aka.ms/vs/17/release/vc_redist.x86.exe and https://aka.ms/vs/17/release/vc_redist.x64.exe 40 | - install python3.10.x from [www.python.org](https://www.python.org/) . 41 | Tested Python3.10.17 and Python3.11.12 . 42 | - install [git](https://git-scm.com/). 43 | - install [AMD HIP SDK for Windows 6.2.4](https://www.amd.com/en/developer/resources/rocm-hub/hip-sdk.html) . 44 | If 5.7.1 already installed, uninstall 5.7.1 before installing 6.2.4 . 45 | - Open Command Prompt (not PowerShell), then run the following: 46 | ``` 47 | git clone https://github.com/githubcto/FramePack-ZLUDA.git 48 | cd FramePack-ZLUDA 49 | python.exe -m venv venv 50 | venv\Scripts\activate.bat 51 | python.exe -m pip install --upgrade pip 52 | pip install torch==2.7.0 torchvision torchaudio --index-url https://download.pytorch.org/whl/cu118 53 | pip install -r requirements.txt 54 | curl -s -L https://github.com/lshqqytiger/ZLUDA/releases/download/rel.5e717459179dc272b7d7d23391f0fad66c7459cf/ZLUDA-windows-rocm6-amd64.zip > zluda.zip 55 | mkdir .zluda && tar -xf zluda.zip -C .zluda --strip-components=1 56 | echo . 57 | ``` 58 | Next, replace the following DLL files in venv/Lib/site-packages/torch/lib with the ones provided in the .zluda folder: 59 | 60 | cublas.dll 61 | cusparse.dll 62 | cufft.dll 63 | cufftw.dll 64 | nvrtc.dll 65 | 66 | You can find them at venv/Lib/site-packages/torch/lib. 67 | ``` 68 | if not exist "%~dp0venv\Lib\site-packages\torch\lib\nvrtc_cuda.dll" copy venv\Lib\site-packages\torch\lib\nvrtc64_112_0.dll venv\Lib\site-packages\torch\lib\nvrtc_cuda.dll & REM (optional) 69 | copy .zluda\cublas.dll venv\Lib\site-packages\torch\lib\cublas64_11.dll /y 70 | copy .zluda\cusparse.dll venv\Lib\site-packages\torch\lib\cusparse64_11.dll /y 71 | copy .zluda\nvrtc.dll venv\Lib\site-packages\torch\lib\nvrtc64_112_0.dll /y 72 | echo Finished 73 | ``` 74 | or, 75 | 76 | [download zip](https://github.com/githubcto/FramePack-ZLUDA/archive/refs/heads/main.zip), extract, and refer to install-win-zluda.bat . 77 | 78 | 79 | 80 | ## Run 81 | FramePack-user.bat 82 | 83 | 1st run, 84 | - will download 40GB. 85 | - ZLUDA compile takes 30 minutes or more. 86 | 87 | Duriing this 30 minutes, you'll see the message 88 | 89 | "Compilation is in progress. Please wait..." every minute. 90 | 91 | If FramePack-user.bat does not work, try FramePack-user-DEVICE0.bat or FramePack-user-DEVICE1.bat . 92 | 93 | FramePackF1-user.bat use different model files, will download additional 25GB. 94 | 95 | ## Tips 96 | 1st time generate, 97 | try 98 | - Use square image. FramePack read it and resize automatically. 99 | - 1 sec 100 | - 10 steps 101 | - other values: use preset 102 | - Start Generation, see VRAM and DRAM usage. 103 | 104 | DRAM 64GB minimum. 64GB enough for linux. 96GB enough for windows. 128GB recommend. 105 | 106 | Set windowsOS page file "auto", "64GB" or more. 107 | 108 | TeaCache is fast, but output quality is not so good. Try TeaCache first and you feel good movie, then disable TeaCache and try same seed again. 109 | 110 | Saved png files can be converted mp4 movie [using ffmpeg](https://ffmpeg.org/) like this. 111 | ``` 112 | ffmpeg.exe -framerate 24 -i %4d.png -c:v libx264 -crf 23 -pix_fmt yuv420p -an out.mp4 113 | ``` 114 | 115 | ### for RDNA2 RX 6000 VRAM16GB 116 | try attention branch. Flash Attention 2 works. 117 | 118 | https://github.com/githubcto/FramePack-ZLUDA/tree/attention 119 | 120 | ## Issue 121 | Since my VGA is RX 6000, I can not verify some Attentions which RX 7000 support, for example, 122 | 123 | [Repeerc/flash-attention-v2-RDNA3-minimal](https://github.com/Repeerc/flash-attention-v2-RDNA3-minimal) 124 | 125 | ( You need modify [demo_gradio.py (around line 10th.)](https://github.com/githubcto/FramePack-ZLUDA/blob/main/demo_gradio.py#L10) from "False" to "True", maybe.) 126 | 127 | torch.backends.cuda.enable_flash_sdp(False) 128 | 129 | . 130 | 131 | So, I shall close FramePack-ZLUDA repo without any notice. 132 | 133 | Edit 2025 May 14th: made attention branch. Flash Attention 2 works, even rx 6000. 134 | 135 | ## for Linux ROCm 136 | 137 | This code may help. 138 | 139 | FramePack-ZLUDA [demo_gradio.py](https://github.com/githubcto/FramePack-ZLUDA/blob/main/demo_gradio.py#L79-L91) 140 | ``` 141 | # VAE Tiling size 142 | vae.enable_tiling( 143 | tile_sample_min_height=128, #256 144 | tile_sample_min_width=128, #256 145 | tile_sample_min_num_frames=12, #16 146 | tile_sample_stride_height=96, #292 147 | tile_sample_stride_width=96, #192 148 | tile_sample_stride_num_frames=10 #12 149 | ) 150 | ``` 151 | 152 | VAE tile size is adjustable. not only resolution but also frames. 153 | 154 | source code 155 | 156 | venv/Lib/site-packages/diffusers/models/autoencoders/autoencoder_kl_hunyuan_video.py 157 | 158 | HuggingFace diffusers [autoencoder_kl_hunyuan_video.py](https://github.com/huggingface/diffusers/blob/f00a995753732210a696de447cd0db80e181c30a/src/diffusers/models/autoencoders/autoencoder_kl_hunyuan_video.py#L717-L766) 159 | 160 | ## ChangeLog 161 | 162 | 2025 May 14th : torch2.7.0 and ZLUDA3.9.5. add attention branch. 163 | 164 | 2025 May 05th : FramePack-F1, unveil Latent Window Size slider. 165 | 166 | 2025 Apr. 26th : add FPS switch. default=24fps. QuickList2nd changed. (torch2.7.0 and ZLUDA3.9.3 works, but keep torch2.6.0 for a while). 167 | 168 | 2025 Apr. 25th : Init. ZLUDA, RESOLUTION, SAVE PNG, README. 169 | 170 | --- 171 | --- 172 | 173 |

174 | 175 |

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