diff --git a/.gitattributes b/.gitattributes
index a6344aac8c09253b3b630fb776ae94478aa0275b..751ca67f6a019df773400b6667857aafd622950b 100644
--- a/.gitattributes
+++ b/.gitattributes
@@ -33,3 +33,4 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
*.zip filter=lfs diff=lfs merge=lfs -text
*.zst filter=lfs diff=lfs merge=lfs -text
*tfevents* filter=lfs diff=lfs merge=lfs -text
+.github/workflows/logo.gif filter=lfs diff=lfs merge=lfs -text
diff --git a/.github/workflows/logo.gif b/.github/workflows/logo.gif
new file mode 100644
index 0000000000000000000000000000000000000000..ef5717efc17bbb2018a37a530a9d7a09e86277d9
--- /dev/null
+++ b/.github/workflows/logo.gif
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:36a7627b7f0f0a508ec64aba72e5d95d38dfe7958bd8cf42d2a63f6ac2641529
+size 149067
diff --git a/.github/workflows/publish.yaml b/.github/workflows/publish.yaml
new file mode 100644
index 0000000000000000000000000000000000000000..31e8947290b23601a51e105fbe22bc31dee4693e
--- /dev/null
+++ b/.github/workflows/publish.yaml
@@ -0,0 +1,29 @@
+name: release
+
+on:
+ push:
+ tags:
+ - 'v**'
+
+concurrency:
+ group: ${{ github.workflow }}-${{ github.ref }}-publish
+ cancel-in-progress: true
+
+jobs:
+ build-n-publish:
+ runs-on: ubuntu-20.04
+ #if: startsWith(github.event.ref, 'refs/tags')
+ steps:
+ - uses: actions/checkout@v2
+ - name: Set up Python 3.10
+ uses: actions/setup-python@v2
+ with:
+ python-version: '3.10'
+ - name: Install wheel
+ run: pip install wheel==0.44.0 && pip install -r requirements.txt
+ - name: Build DiffSynth
+ run: python -m build
+ - name: Publish package to PyPI
+ run: |
+ pip install twine
+ twine upload dist/* --skip-existing -u __token__ -p ${{ secrets.PYPI_API_TOKEN }}
diff --git a/.gitignore b/.gitignore
new file mode 100644
index 0000000000000000000000000000000000000000..7cce2c12dcecb4ddd9ca29f2b20a58fc395b6e07
--- /dev/null
+++ b/.gitignore
@@ -0,0 +1,178 @@
+/data
+/models
+/scripts
+/diffusers
+/.vscode
+/.opencode
+*.pkl
+*.safetensors
+*.pth
+*.ckpt
+*.pt
+*.bin
+*.DS_Store
+*.msc
+*.mv
+log*.txt
+AGENTS.md
+
+# Byte-compiled / optimized / DLL files
+__pycache__/
+*.py[cod]
+*$py.class
+
+# C extensions
+*.so
+
+# Distribution / packaging
+.Python
+build/
+develop-eggs/
+dist/
+downloads/
+eggs/
+.eggs/
+lib/
+lib64/
+parts/
+sdist/
+var/
+wheels/
+share/python-wheels/
+*.egg-info/
+.installed.cfg
+*.egg
+MANIFEST
+
+# PyInstaller
+# Usually these files are written by a python script from a template
+# before PyInstaller builds the exe, so as to inject date/other infos into it.
+*.manifest
+*.spec
+
+# Installer logs
+pip-log.txt
+pip-delete-this-directory.txt
+
+# Unit test / coverage reports
+htmlcov/
+.tox/
+.nox/
+.coverage
+.coverage.*
+.cache
+nosetests.xml
+coverage.xml
+*.cover
+*.py,cover
+.hypothesis/
+.pytest_cache/
+cover/
+
+# Translations
+*.mo
+*.pot
+
+# Django stuff:
+*.log
+local_settings.py
+db.sqlite3
+db.sqlite3-journal
+
+# Flask stuff:
+instance/
+.webassets-cache
+
+# Scrapy stuff:
+.scrapy
+
+# Sphinx documentation
+docs/_build/
+
+# PyBuilder
+.pybuilder/
+target/
+
+# Jupyter Notebook
+.ipynb_checkpoints
+
+# IPython
+profile_default/
+ipython_config.py
+
+# pyenv
+# For a library or package, you might want to ignore these files since the code is
+# intended to run in multiple environments; otherwise, check them in:
+# .python-version
+
+# pipenv
+# According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control.
+# However, in case of collaboration, if having platform-specific dependencies or dependencies
+# having no cross-platform support, pipenv may install dependencies that don't work, or not
+# install all needed dependencies.
+#Pipfile.lock
+
+# poetry
+# Similar to Pipfile.lock, it is generally recommended to include poetry.lock in version control.
+# This is especially recommended for binary packages to ensure reproducibility, and is more
+# commonly ignored for libraries.
+# https://python-poetry.org/docs/basic-usage/#commit-your-poetrylock-file-to-version-control
+#poetry.lock
+
+# pdm
+# Similar to Pipfile.lock, it is generally recommended to include pdm.lock in version control.
+#pdm.lock
+# pdm stores project-wide configurations in .pdm.toml, but it is recommended to not include it
+# in version control.
+# https://pdm.fming.dev/#use-with-ide
+.pdm.toml
+
+# PEP 582; used by e.g. github.com/David-OConnor/pyflow and github.com/pdm-project/pdm
+__pypackages__/
+
+# Celery stuff
+celerybeat-schedule
+celerybeat.pid
+
+# SageMath parsed files
+*.sage.py
+
+# Environments
+.env
+.venv
+env/
+venv/
+ENV/
+env.bak/
+venv.bak/
+
+# Spyder project settings
+.spyderproject
+.spyproject
+
+# Rope project settings
+.ropeproject
+
+# mkdocs documentation
+/site
+
+# mypy
+.mypy_cache/
+.dmypy.json
+dmypy.json
+
+# Pyre type checker
+.pyre/
+
+# pytype static type analyzer
+.pytype/
+
+# Cython debug symbols
+cython_debug/
+
+# PyCharm
+# JetBrains specific template is maintained in a separate JetBrains.gitignore that can
+# be found at https://github.com/github/gitignore/blob/main/Global/JetBrains.gitignore
+# and can be added to the global gitignore or merged into this file. For a more nuclear
+# option (not recommended) you can uncomment the following to ignore the entire idea folder.
+#.idea/
\ No newline at end of file
diff --git a/LICENSE b/LICENSE
new file mode 100644
index 0000000000000000000000000000000000000000..0e5a49f5f70af9e9d37278f72315d1b1afd34895
--- /dev/null
+++ b/LICENSE
@@ -0,0 +1,201 @@
+ Apache License
+ Version 2.0, January 2004
+ http://www.apache.org/licenses/
+
+ TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
+
+ 1. Definitions.
+
+ "License" shall mean the terms and conditions for use, reproduction,
+ and distribution as defined by Sections 1 through 9 of this document.
+
+ "Licensor" shall mean the copyright owner or entity authorized by
+ the copyright owner that is granting the License.
+
+ "Legal Entity" shall mean the union of the acting entity and all
+ other entities that control, are controlled by, or are under common
+ control with that entity. For the purposes of this definition,
+ "control" means (i) the power, direct or indirect, to cause the
+ direction or management of such entity, whether by contract or
+ otherwise, or (ii) ownership of fifty percent (50%) or more of the
+ outstanding shares, or (iii) beneficial ownership of such entity.
+
+ "You" (or "Your") shall mean an individual or Legal Entity
+ exercising permissions granted by this License.
+
+ "Source" form shall mean the preferred form for making modifications,
+ including but not limited to software source code, documentation
+ source, and configuration files.
+
+ "Object" form shall mean any form resulting from mechanical
+ transformation or translation of a Source form, including but
+ not limited to compiled object code, generated documentation,
+ and conversions to other media types.
+
+ "Work" shall mean the work of authorship, whether in Source or
+ Object form, made available under the License, as indicated by a
+ copyright notice that is included in or attached to the work
+ (an example is provided in the Appendix below).
+
+ "Derivative Works" shall mean any work, whether in Source or Object
+ form, that is based on (or derived from) the Work and for which the
+ editorial revisions, annotations, elaborations, or other modifications
+ represent, as a whole, an original work of authorship. For the purposes
+ of this License, Derivative Works shall not include works that remain
+ separable from, or merely link (or bind by name) to the interfaces of,
+ the Work and Derivative Works thereof.
+
+ "Contribution" shall mean any work of authorship, including
+ the original version of the Work and any modifications or additions
+ to that Work or Derivative Works thereof, that is intentionally
+ submitted to Licensor for inclusion in the Work by the copyright owner
+ or by an individual or Legal Entity authorized to submit on behalf of
+ the copyright owner. For the purposes of this definition, "submitted"
+ means any form of electronic, verbal, or written communication sent
+ to the Licensor or its representatives, including but not limited to
+ communication on electronic mailing lists, source code control systems,
+ and issue tracking systems that are managed by, or on behalf of, the
+ Licensor for the purpose of discussing and improving the Work, but
+ excluding communication that is conspicuously marked or otherwise
+ designated in writing by the copyright owner as "Not a Contribution."
+
+ "Contributor" shall mean Licensor and any individual or Legal Entity
+ on behalf of whom a Contribution has been received by Licensor and
+ subsequently incorporated within the Work.
+
+ 2. Grant of Copyright License. Subject to the terms and conditions of
+ this License, each Contributor hereby grants to You a perpetual,
+ worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+ copyright license to reproduce, prepare Derivative Works of,
+ publicly display, publicly perform, sublicense, and distribute the
+ Work and such Derivative Works in Source or Object form.
+
+ 3. Grant of Patent License. Subject to the terms and conditions of
+ this License, each Contributor hereby grants to You a perpetual,
+ worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+ (except as stated in this section) patent license to make, have made,
+ use, offer to sell, sell, import, and otherwise transfer the Work,
+ where such license applies only to those patent claims licensable
+ by such Contributor that are necessarily infringed by their
+ Contribution(s) alone or by combination of their Contribution(s)
+ with the Work to which such Contribution(s) was submitted. If You
+ institute patent litigation against any entity (including a
+ cross-claim or counterclaim in a lawsuit) alleging that the Work
+ or a Contribution incorporated within the Work constitutes direct
+ or contributory patent infringement, then any patent licenses
+ granted to You under this License for that Work shall terminate
+ as of the date such litigation is filed.
+
+ 4. Redistribution. You may reproduce and distribute copies of the
+ Work or Derivative Works thereof in any medium, with or without
+ modifications, and in Source or Object form, provided that You
+ meet the following conditions:
+
+ (a) You must give any other recipients of the Work or
+ Derivative Works a copy of this License; and
+
+ (b) You must cause any modified files to carry prominent notices
+ stating that You changed the files; and
+
+ (c) You must retain, in the Source form of any Derivative Works
+ that You distribute, all copyright, patent, trademark, and
+ attribution notices from the Source form of the Work,
+ excluding those notices that do not pertain to any part of
+ the Derivative Works; and
+
+ (d) If the Work includes a "NOTICE" text file as part of its
+ distribution, then any Derivative Works that You distribute must
+ include a readable copy of the attribution notices contained
+ within such NOTICE file, excluding those notices that do not
+ pertain to any part of the Derivative Works, in at least one
+ of the following places: within a NOTICE text file distributed
+ as part of the Derivative Works; within the Source form or
+ documentation, if provided along with the Derivative Works; or,
+ within a display generated by the Derivative Works, if and
+ wherever such third-party notices normally appear. The contents
+ of the NOTICE file are for informational purposes only and
+ do not modify the License. You may add Your own attribution
+ notices within Derivative Works that You distribute, alongside
+ or as an addendum to the NOTICE text from the Work, provided
+ that such additional attribution notices cannot be construed
+ as modifying the License.
+
+ You may add Your own copyright statement to Your modifications and
+ may provide additional or different license terms and conditions
+ for use, reproduction, or distribution of Your modifications, or
+ for any such Derivative Works as a whole, provided Your use,
+ reproduction, and distribution of the Work otherwise complies with
+ the conditions stated in this License.
+
+ 5. Submission of Contributions. Unless You explicitly state otherwise,
+ any Contribution intentionally submitted for inclusion in the Work
+ by You to the Licensor shall be under the terms and conditions of
+ this License, without any additional terms or conditions.
+ Notwithstanding the above, nothing herein shall supersede or modify
+ the terms of any separate license agreement you may have executed
+ with Licensor regarding such Contributions.
+
+ 6. Trademarks. This License does not grant permission to use the trade
+ names, trademarks, service marks, or product names of the Licensor,
+ except as required for reasonable and customary use in describing the
+ origin of the Work and reproducing the content of the NOTICE file.
+
+ 7. Disclaimer of Warranty. Unless required by applicable law or
+ agreed to in writing, Licensor provides the Work (and each
+ Contributor provides its Contributions) on an "AS IS" BASIS,
+ WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
+ implied, including, without limitation, any warranties or conditions
+ of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A
+ PARTICULAR PURPOSE. You are solely responsible for determining the
+ appropriateness of using or redistributing the Work and assume any
+ risks associated with Your exercise of permissions under this License.
+
+ 8. Limitation of Liability. In no event and under no legal theory,
+ whether in tort (including negligence), contract, or otherwise,
+ unless required by applicable law (such as deliberate and grossly
+ negligent acts) or agreed to in writing, shall any Contributor be
+ liable to You for damages, including any direct, indirect, special,
+ incidental, or consequential damages of any character arising as a
+ result of this License or out of the use or inability to use the
+ Work (including but not limited to damages for loss of goodwill,
+ work stoppage, computer failure or malfunction, or any and all
+ other commercial damages or losses), even if such Contributor
+ has been advised of the possibility of such damages.
+
+ 9. Accepting Warranty or Additional Liability. While redistributing
+ the Work or Derivative Works thereof, You may choose to offer,
+ and charge a fee for, acceptance of support, warranty, indemnity,
+ or other liability obligations and/or rights consistent with this
+ License. However, in accepting such obligations, You may act only
+ on Your own behalf and on Your sole responsibility, not on behalf
+ of any other Contributor, and only if You agree to indemnify,
+ defend, and hold each Contributor harmless for any liability
+ incurred by, or claims asserted against, such Contributor by reason
+ of your accepting any such warranty or additional liability.
+
+ END OF TERMS AND CONDITIONS
+
+ APPENDIX: How to apply the Apache License to your work.
+
+ To apply the Apache License to your work, attach the following
+ boilerplate notice, with the fields enclosed by brackets "[]"
+ replaced with your own identifying information. (Don't include
+ the brackets!) The text should be enclosed in the appropriate
+ comment syntax for the file format. We also recommend that a
+ file or class name and description of purpose be included on the
+ same "printed page" as the copyright notice for easier
+ identification within third-party archives.
+
+ Copyright [2023] [Zhongjie Duan]
+
+ Licensed under the Apache License, Version 2.0 (the "License");
+ you may not use this file except in compliance with the License.
+ You may obtain a copy of the License at
+
+ http://www.apache.org/licenses/LICENSE-2.0
+
+ Unless required by applicable law or agreed to in writing, software
+ distributed under the License is distributed on an "AS IS" BASIS,
+ WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ See the License for the specific language governing permissions and
+ limitations under the License.
diff --git a/README.md b/README.md
new file mode 100644
index 0000000000000000000000000000000000000000..38515ca72c20bc8ebdf86b81c8de1af935173bb1
--- /dev/null
+++ b/README.md
@@ -0,0 +1,758 @@
+# DiffSynth-Studio
+
+
+
+[](https://pypi.org/project/DiffSynth/)
+[](https://github.com/modelscope/DiffSynth-Studio/blob/master/LICENSE)
+[](https://github.com/modelscope/DiffSynth-Studio/issues)
+[](https://GitHub.com/modelscope/DiffSynth-Studio/pull/)
+[](https://GitHub.com/modelscope/DiffSynth-Studio/commit/)
+[](https://discord.gg/Mm9suEeUDc)
+
+[切换到中文版](./README_zh.md)
+
+## Introduction
+
+Welcome to the magical world of Diffusion models! DiffSynth-Studio is an open-source Diffusion model engine developed and maintained by the [ModelScope Community](https://www.modelscope.cn/) team. We hope to foster technological innovation through framework construction, aggregate the power of the open-source community, and explore the rich capabilities of generative model technology!
+
+Framework features:
+
+* [Model Support](#all-supported-models): The framework integrates mainstream open-source Diffusion models, covering image generation, video generation, audio generation, and image quality metrics models.
+* [VRAM Management](https://diffsynth-studio-doc.readthedocs.io/en/latest/Pipeline_Usage/VRAM_management.html): Dynamically schedules model parameters across disk, memory, and VRAM, allowing consumer-grade GPUs with low VRAM to run inference with large models.
+* [Parameter Quantization](https://diffsynth-studio-doc.readthedocs.io/en/latest/Pipeline_Usage/Quantization.html): Converts model parameters to quantized precisions such as NF4 and INT8, significantly reducing the VRAM requirements of model inference and LoRA training.
+* [Arbitrary Training](https://diffsynth-studio-doc.readthedocs.io/en/latest/Pipeline_Usage/Model_Training.html): Almost every model that supports inference also supports training, whether base models, LoRAs, or any Adapter models with additional inputs.
+* [Split Training](https://diffsynth-studio-doc.readthedocs.io/en/latest/Training/Split_Training.html): Uses a computational graph inference engine to track every variable in the Pipeline, splitting the training process into two stages for efficient training.
+
+References:
+
+* Developer documentation (for humans): [中文版](https://diffsynth-studio-doc.readthedocs.io/zh-cn/latest/README.html)、[English version](https://diffsynth-studio-doc.readthedocs.io/en/latest/README.html)
+* Agent Skills (for AI): [DiffSynth-Studio Model Integration Skills](https://www.modelscope.cn/collections/DiffSynth-Studio/DiffSynth-Studio-Model-Integration-Skills)
+
+See also:
+
+* [DiffSynth-WebUI](https://github.com/modelscope/DiffSynth-WebUI): A lightweight LoRA training tool built on DiffSynth-Studio, enabling LoRA training of models on consumer-grade GPUs.
+* [ModelScope AIGC Zone (for Chinese users)](https://modelscope.cn/aigc/home): Productized features powered by DiffSynth-Studio as the core inference and training engine; experience the rich potential of the open-source model ecosystem.
+* [ModelScope Civision (for global users)](https://modelscope.ai/civision/home): Unlock the vast potential of the open-source model ecosystem through productized capabilities powered by DiffSynth-Studio.
+
+## Update History
+
+> DiffSynth-Studio has undergone major version updates, and some old features are no longer maintained. If you need to use old features, please switch to the [last historical version](https://github.com/modelscope/DiffSynth-Studio/tree/afd101f3452c9ecae0c87b79adfa2e22d65ffdc3) before the major version update.
+
+> Currently, the development personnel of this project are limited, with most of the work handled by [Artiprocher](https://github.com/Artiprocher) and [mi804](https://github.com/mi804). Therefore, the progress of new feature development will be relatively slow, and the speed of responding to and resolving issues is limited. We apologize for this and ask developers to understand.
+
+- **August 31, 2026** We have integrated [Qwen-Video-Edit](https://modelscope.cn/models/yunpeng1998/Qwen-Video-Edit), a video editing model developed by open-source community contributor [yunpeng1998](https://github.com/yunpeng1998) based on the image editing model Qwen-Image-Edit. This serves as an excellent example of exploring and expanding model capabilities.
+
+- **August 25, 2026** We have open-sourced [DiffSynth-WebUI](https://github.com/modelscope/DiffSynth-WebUI), enabling one-click private deployment of LoRA training services. Combined with the model quantization feature, you can train large models even with consumer-grade GPUs.
+
+- **August 19, 2026** We have released the model quantization feature. It provides a unified `QuantizeConfig` entry point supporting multiple quantization backends including bitsandbytes, torchao, and comfy-kitchen, with capabilities such as online quantization, loading pre-quantized weights, mixed quantization, saving quantized models, and quantization + LoRA training. For details, please refer to the [documentation](/docs/en/Pipeline_Usage/Quantization.md).
+
+- **August 17, 2026** MiniMax-Music3 open-sourced, welcome a new member to the audio model family! Support includes text-to-music generation and low VRAM inference. For details, please refer to the [documentation](/docs/en/Model_Details/MiniMax-Music3.md) and [example code](/examples/minimax_music3/).
+
+- **August 7, 2026** We add support for Wan-Animate-2 in the Wan series. Given a reference image and a driving video, it makes the reference character perform the motions in the driving video, generating high-quality character animation, with both standard and distilled variants. For details, please refer to the [documentation](/docs/en/Model_Details/Wan.md) and [example code](/examples/wanvideo/).
+
+- **August 3, 2026** MiniMax-H3 open-sourced, welcome a new member to the video model family! Support includes text-to-video-audio generation, keyframe-guided generation, reference-driven generation, low VRAM inference, and NF4-quantized inference. For details, please refer to the [documentation](/docs/en/Model_Details/MiniMax-H3.md) and [example code](/examples/minimax_h3/).
+
+- **July 28, 2026** LingBot-Video open-sourced, welcome a new member to the video model family! This release includes two variants, Dense-1.3B and MoE-30B-A3B (30B total parameters, ~3B active per token), both supporting text-to-video, image-to-video and text-to-image generation, low VRAM inference, and LoRA / full training capabilities. For details, please refer to the [documentation](/docs/en/Model_Details/LingBot-Video.md) and [example code](/examples/lingbot_video/). Huge thanks to [NancyFyong](https://github.com/NancyFyong) for contributing the integration of this model!
+
+- **July 21, 2026** We have open-sourced [DiffSynth-Studio Model Integration Skills](https://www.modelscope.cn/collections/DiffSynth-Studio/DiffSynth-Studio-Model-Integration-Skills). This is a composable collection of Agent Skills that automates the entire workflow of integrating external diffusion models into DiffSynth-Studio, significantly improving the standardization and efficiency of model integration. Get started with the [example](https://www.modelscope.cn/skills/DiffSynth-Studio/diffsynth-integrator/file/view/master/example.md?status=1)!
+
+
+More
+
+- **June 29, 2026** Boogu-Image open-sourced. Support includes text-to-image generation, image editing, low VRAM inference, and training capabilities. For details, please refer to the [documentation](/docs/en/Model_Details/Boogu-Image.md) and [example code](/examples/boogu_image/).
+
+- **June 24, 2026** Krea-2 is now open-source, and we have provided full support. For more details, please refer to the [documentation](/docs/en/Model_Details/Krea-2.md) and [example code](/examples/krea2/).
+
+- **June 16, 2026**: We have added a new Template model for ACE-Step: [vocals2music](https://www.modelscope.cn/models/DiffSynth-Studio/acestep15xlsft-vocals2music). For more details, please refer to the [documentation](/docs/en/Model_Details/ACE-Step.md) and [example code](/examples/ace_step/).
+
+- **June 15, 2026** We have open-sourced Image-to-LoRA V2, compressing the hours-long training process for image style LoRAs into a single model inference step, thereby exploring a new paradigm for LoRA model training. The [technical report](https://arxiv.org/abs/2606.13809) has been released. This release includes three models:
+ * [DiffSynth-Studio/ZImage-i2L-v2](https://modelscope.cn/models/DiffSynth-Studio/ZImage-i2L-v2): Adapted for the Z-Image model
+ * [DiffSynth-Studio/KleinBase4B-i2L-v2](https://modelscope.cn/models/DiffSynth-Studio/KleinBase4B-i2L-v2): Adapted for the FLUX.2-klein-base-4B model
+ * [DiffSynth-Studio/HidreamO1-i2L-v2](https://modelscope.cn/models/DiffSynth-Studio/HidreamO1-i2L-v2): Adapted for the Hidream-O1-Image model
+
+- **June 5, 2026** Ideogram 4 open-sourced. Support includes text-to-image inference. For details, please refer to the [documentation](/docs/en/Model_Details/Ideogram-4.md) and [example code](/examples/ideogram4/).
+
+- **May 21, 2026**: Added support for image quality metrics models, including FID, CLIP, Aesthetic, PickScore, ImageReward, HPSv2, and HPSv3. For details, refer to the [documentation](/docs/en/Model_Details/Image-Quality-Metrics.md) and [example code](/examples/image_quality_metric/).
+
+- **May 18, 2026** Added **CPU Offload Training** support. By moving model weights layer-by-layer between CPU and GPU, it significantly reduces GPU VRAM usage during training, enabling LoRA training of large models even on consumer-grade GPUs, compatible with all models. Simply add `--enable_model_cpu_offload` to your training command to enable (currently supports single-GPU training only). For details, see the [documentation](/docs/en/Training/Offload_Training.md).
+
+- **May 14, 2026** HiDream-O1-Image open-sourced, welcome a new member to the image model family! Support includes text-to-image generation, image editing, low VRAM inference, and training capabilities. For details, please refer to the [documentation](/docs/en/Model_Details/HiDream-O1-Image.md) and [example code](/examples/hidream_o1_image/).
+
+- **April 28, 2026** We released Diffusion Templates, a plugin framework designed for Diffusion models that significantly lowers the barrier to training controllable generative models. Let's explore this cutting-edge technology together!
+ * Open-source code: [DiffSynth-Studio](https://github.com/modelscope/DiffSynth-Studio)
+ * Technical report: [arXiv](https://arxiv.org/abs/2604.24351)
+ * Project homepage: [GitHub](https://modelscope.github.io/diffusion-templates-web/)
+ * Documentation: [English Version](https://diffsynth-studio-doc.readthedocs.io/en/latest/Diffusion_Templates/Introducing_Diffusion_Templates.html) | [Chinese Version](https://diffsynth-studio-doc.readthedocs.io/zh-cn/latest/Diffusion_Templates/Introducing_Diffusion_Templates.html)
+ * Online demo: [ModelScope](https://modelscope.cn/studios/DiffSynth-Studio/Diffusion-Templates)
+ * Model collections: [ModelScope](https://modelscope.cn/collections/DiffSynth-Studio/KleinBase4B-Templates) | [ModelScope International](https://modelscope.ai/collections/DiffSynth-Studio/KleinBase4B-Templates) | [HuggingFace](https://huggingface.co/collections/DiffSynth-Studio/kleinbase4b-templates)
+ * Datasets: [ModelScope](https://modelscope.cn/collections/DiffSynth-Studio/ImagePulseV2) | [ModelScope International](https://modelscope.ai/collections/DiffSynth-Studio/ImagePulseV2) | [HuggingFace](https://huggingface.co/collections/DiffSynth-Studio/imagepulsev2)
+
+- **April 27, 2026** We support ACE-Step-1.5! Support includes text-to-music generation, low VRAM inference, and LoRA training capabilities. For details, please refer to the [documentation](/docs/en/Model_Details/ACE-Step.md) and [example code](/examples/ace_step/).
+
+- **April 27, 2026**: We have reinstated support for the Stable Diffusion v1.5 and SDXL models, providing academic research support exclusively for these two model types.
+
+- **April 14, 2026** JoyAI-Image open-sourced, welcome a new member to the image editing model family! Support includes instruction-guided image editing, low VRAM inference, and training capabilities. For details, please refer to the [documentation](/docs/en/Model_Details/JoyAI-Image.md) and [example code](/examples/joyai_image/).
+
+
+- **March 19, 2026**: Added support for [openmoss/MOVA-720p](https://modelscope.cn/models/openmoss/MOVA-720p) and [openmoss/MOVA-360p](https://modelscope.cn/models/openmoss/MOVA-360p) models, including training and inference capabilities. [Documentation](/docs/en/Model_Details/Wan.md) and [example code](/examples/mova/) are now available.
+
+- **March 12, 2026**: We have added support for the [LTX-2.3](https://modelscope.cn/models/Lightricks/LTX-2.3) audio-video generation model. The features includes text-to-audio/video, image-to-audio/video, IC-LoRA control, audio-to-video, and audio-video inpainting. We have supported the complete inference and training functionalities. For details, please refer to the [documentation](/docs/en/Model_Details/LTX-2.md) and [code](/examples/ltx2/).
+
+- **March 3, 2026**: We released the [DiffSynth-Studio/Qwen-Image-Layered-Control-V2](https://www.modelscope.cn/models/DiffSynth-Studio/Qwen-Image-Layered-Control-V2) model, which is an updated version of Qwen-Image-Layered-Control. In addition to the originally supported text-guided functionality, it adds brush-controlled layer separation capabilities.
+
+- **March 2, 2026** Added support for [Anima](https://modelscope.cn/models/circlestone-labs/Anima). For details, please refer to the [documentation](docs/en/Model_Details/Anima.md). This is an interesting anime-style image generation model. We look forward to its future updates.
+
+- **February 26, 2026** Added full and lora training support for the LTX-2 audio-video generation model. See the [documentation](/docs/en/Model_Details/LTX-2.md) for details.
+
+- **February 10, 2026** Added inference support for the LTX-2 audio-video generation model. See the [documentation](/docs/en/Model_Details/LTX-2.md) for details. Support for model training will be implemented in the future.
+
+- **February 2, 2026** The first document of the Research Tutorial series is now available, guiding you through training a small 0.1B text-to-image model from scratch. For details, see the [documentation](/docs/en/Research_Tutorial/train_from_scratch.md) and [model](https://modelscope.cn/models/DiffSynth-Studio/AAAMyModel). We hope DiffSynth-Studio can evolve into a more powerful training framework for Diffusion models.
+
+- **January 27, 2026**: [Z-Image](https://modelscope.cn/models/Tongyi-MAI/Z-Image) is released, and our [Z-Image-i2L](https://www.modelscope.cn/models/DiffSynth-Studio/Z-Image-i2L) model is released concurrently. You can use it in [ModelScope Studios](https://modelscope.cn/studios/DiffSynth-Studio/Z-Image-i2L). For details, see the [documentation](/docs/zh/Model_Details/Z-Image.md).
+
+- **January 19, 2026**: Added support for [FLUX.2-klein-4B](https://modelscope.cn/models/black-forest-labs/FLUX.2-klein-4B) and [FLUX.2-klein-9B](https://modelscope.cn/models/black-forest-labs/FLUX.2-klein-9B) models, including training and inference capabilities. [Documentation](/docs/en/Model_Details/FLUX2.md) and [example code](/examples/flux2/) are now available.
+
+- **January 12, 2026**: We trained and open-sourced a text-guided image layer separation model ([Model Link](https://modelscope.cn/models/DiffSynth-Studio/Qwen-Image-Layered-Control)). Given an input image and a textual description, the model isolates the image layer corresponding to the described content. For more details, please refer to our blog post ([Chinese version](https://modelscope.cn/learn/4938), [English version](https://huggingface.co/blog/kelseye/qwen-image-layered-control)).
+
+- **December 24, 2025**: Based on Qwen-Image-Edit-2511, we trained an In-Context Editing LoRA model ([Model Link](https://modelscope.cn/models/DiffSynth-Studio/Qwen-Image-Edit-2511-ICEdit-LoRA)). This model takes three images as input (Image A, Image B, and Image C), and automatically analyzes the transformation from Image A to Image B, then applies the same transformation to Image C to generate Image D. For more details, please refer to our blog post ([Chinese version](https://mp.weixin.qq.com/s/41aEiN3lXKGCJs1-we4Q2g), [English version](https://huggingface.co/blog/kelseye/qwen-image-edit-2511-icedit-lora)).
+
+- **December 9, 2025** We release a wild model based on DiffSynth-Studio 2.0: [Qwen-Image-i2L](https://www.modelscope.cn/models/DiffSynth-Studio/Qwen-Image-i2L) (Image-to-LoRA). This model takes an image as input and outputs a LoRA. Although this version still has significant room for improvement in terms of generalization, detail preservation, and other aspects, we are open-sourcing these models to inspire more innovative research. For more details, please refer to our [blog](https://huggingface.co/blog/kelseye/qwen-image-i2l).
+
+- **December 4, 2025** DiffSynth-Studio 2.0 released! Many new features online
+ - [Documentation](/docs/en/README.md) online: Our documentation is still continuously being optimized and updated
+ - [VRAM Management](/docs/en/Pipeline_Usage/VRAM_management.md) module upgraded, supporting layer-level disk offload, releasing both memory and VRAM simultaneously
+ - New model support
+ - Z-Image Turbo: [Model](https://www.modelscope.ai/models/Tongyi-MAI/Z-Image-Turbo), [Documentation](/docs/en/Model_Details/Z-Image.md), [Code](/examples/z_image/)
+ - FLUX.2-dev: [Model](https://www.modelscope.cn/models/black-forest-labs/FLUX.2-dev), [Documentation](/docs/en/Model_Details/FLUX2.md), [Code](/examples/flux2/)
+ - Training framework upgrade
+ - [Split Training](/docs/en/Training/Split_Training.md): Supports automatically splitting the training process into two stages: data processing and training (even for training ControlNet or any other model). Computations that do not require gradient backpropagation, such as text encoding and VAE encoding, are performed during the data processing stage, while other computations are handled during the training stage. Faster speed, less VRAM requirement.
+ - [Differential LoRA Training](/docs/en/Training/Differential_LoRA.md): This is a training technique we used in [ArtAug](https://www.modelscope.cn/models/DiffSynth-Studio/ArtAug-lora-FLUX.1dev-v1), now available for LoRA training of any model.
+ - [FP8 Training](/docs/en/Training/FP8_Precision.md): FP8 can be applied to any non-training model during training, i.e., models with gradients turned off or gradients that only affect LoRA weights.
+
+- **November 4, 2025** Supported the [ByteDance/Video-As-Prompt-Wan2.1-14B](https://modelscope.cn/models/ByteDance/Video-As-Prompt-Wan2.1-14B) model, which is trained based on Wan 2.1 and supports generating corresponding actions based on reference videos.
+
+- **October 30, 2025** Supported the [meituan-longcat/LongCat-Video](https://www.modelscope.cn/models/meituan-longcat/LongCat-Video) model, which supports text-to-video, image-to-video, and video continuation. This model uses the Wan framework for inference and training in this project.
+
+- **October 27, 2025** Supported the [krea/krea-realtime-video](https://www.modelscope.cn/models/krea/krea-realtime-video) model, adding another member to the Wan model ecosystem.
+
+- **September 23, 2025** [DiffSynth-Studio/Qwen-Image-EliGen-Poster](https://www.modelscope.cn/models/DiffSynth-Studio/Qwen-Image-EliGen-Poster) released! This model was jointly developed and open-sourced by us and Taobao Experience Design Team. Built upon Qwen-Image, the model is specifically designed for e-commerce poster scenarios, supporting precise partition layout control. Please refer to [our sample code](./examples/qwen_image/model_inference/Qwen-Image-EliGen-Poster.py).
+
+- **September 9, 2025** Our training framework supports various training modes. Currently adapted for Qwen-Image, in addition to the standard SFT training mode, Direct Distill is now supported. Please refer to [our sample code](./examples/qwen_image/model_training/lora/Qwen-Image-Distill-LoRA.sh). This feature is experimental, and we will continue to improve it to support more comprehensive model training functions.
+
+- **August 28, 2025** We support Wan2.2-S2V, an audio-driven cinematic video generation model. See [./examples/wanvideo/](./examples/wanvideo/).
+
+- **August 21, 2025** [DiffSynth-Studio/Qwen-Image-EliGen-V2](https://www.modelscope.cn/models/DiffSynth-Studio/Qwen-Image-EliGen-V2) released! Compared to the V1 version, the training dataset has been changed to [Qwen-Image-Self-Generated-Dataset](https://www.modelscope.cn/datasets/DiffSynth-Studio/Qwen-Image-Self-Generated-Dataset), so the generated images better conform to Qwen-Image's own image distribution and style. Please refer to [our sample code](./examples/qwen_image/model_inference_low_vram/Qwen-Image-EliGen-V2.py).
+
+- **August 21, 2025** We open-sourced the [DiffSynth-Studio/Qwen-Image-In-Context-Control-Union](https://www.modelscope.cn/models/DiffSynth-Studio/Qwen-Image-In-Context-Control-Union) structural control LoRA model, adopting the In Context technical route, supporting multiple categories of structural control conditions, including canny, depth, lineart, softedge, normal, and openpose. Please refer to [our sample code](./examples/qwen_image/model_inference/Qwen-Image-In-Context-Control-Union.py).
+
+- **August 20, 2025** We open-sourced the [DiffSynth-Studio/Qwen-Image-Edit-Lowres-Fix](https://www.modelscope.cn/models/DiffSynth-Studio/Qwen-Image-Edit-Lowres-Fix) model, improving the editing effect of Qwen-Image-Edit on low-resolution image inputs. Please refer to [our sample code](./examples/qwen_image/model_inference/Qwen-Image-Edit-Lowres-Fix.py)
+
+- **August 19, 2025** Qwen-Image-Edit open-sourced, welcome a new member to the image editing model family!
+
+- **August 18, 2025** We trained and open-sourced the Qwen-Image inpainting ControlNet model [DiffSynth-Studio/Qwen-Image-Blockwise-ControlNet-Inpaint](https://www.modelscope.cn/models/DiffSynth-Studio/Qwen-Image-Blockwise-ControlNet-Inpaint). The model structure adopts a lightweight design. Please refer to [our sample code](./examples/qwen_image/model_inference/Qwen-Image-Blockwise-ControlNet-Inpaint.py).
+
+- **August 15, 2025** We open-sourced the [Qwen-Image-Self-Generated-Dataset](https://www.modelscope.cn/datasets/DiffSynth-Studio/Qwen-Image-Self-Generated-Dataset) dataset. This is an image dataset generated using the Qwen-Image model, containing 160,000 `1024 x 1024` images. It includes general, English text rendering, and Chinese text rendering subsets. We provide annotations for image descriptions, entities, and structural control images for each image. Developers can use this dataset to train Qwen-Image models' ControlNet and EliGen models. We aim to promote technological development through open-sourcing!
+
+- **August 13, 2025** We trained and open-sourced the Qwen-Image ControlNet model [DiffSynth-Studio/Qwen-Image-Blockwise-ControlNet-Depth](https://modelscope.cn/models/DiffSynth-Studio/Qwen-Image-Blockwise-ControlNet-Depth). The model structure adopts a lightweight design. Please refer to [our sample code](./examples/qwen_image/model_inference/Qwen-Image-Blockwise-ControlNet-Depth.py).
+
+- **August 12, 2025** We trained and open-sourced the Qwen-Image ControlNet model [DiffSynth-Studio/Qwen-Image-Blockwise-ControlNet-Canny](https://modelscope.cn/models/DiffSynth-Studio/Qwen-Image-Blockwise-ControlNet-Canny). The model structure adopts a lightweight design. Please refer to [our sample code](./examples/qwen_image/model_inference/Qwen-Image-Blockwise-ControlNet-Canny.py).
+
+- **August 11, 2025** We open-sourced the distilled acceleration model [DiffSynth-Studio/Qwen-Image-Distill-LoRA](https://www.modelscope.cn/models/DiffSynth-Studio/Qwen-Image-Distill-LoRA) for Qwen-Image, following the same training process as [DiffSynth-Studio/Qwen-Image-Distill-Full](https://www.modelscope.cn/models/DiffSynth-Studio/Qwen-Image-Distill-Full), but the model structure has been modified to LoRA, thus being better compatible with other open-source ecosystem models.
+
+- **August 7, 2025** We open-sourced the entity control LoRA model [DiffSynth-Studio/Qwen-Image-EliGen](https://www.modelscope.cn/models/DiffSynth-Studio/Qwen-Image-EliGen) for Qwen-Image. Qwen-Image-EliGen can achieve entity-level controlled text-to-image generation. Technical details can be found in [the paper](https://arxiv.org/abs/2501.01097). Training dataset: [EliGenTrainSet](https://www.modelscope.cn/datasets/DiffSynth-Studio/EliGenTrainSet).
+
+- **August 5, 2025** We open-sourced the distilled acceleration model [DiffSynth-Studio/Qwen-Image-Distill-Full](https://www.modelscope.cn/models/DiffSynth-Studio/Qwen-Image-Distill-Full) for Qwen-Image, achieving approximately 5x acceleration.
+
+- **August 4, 2025** Qwen-Image open-sourced, welcome a new member to the image generation model family!
+
+- **August 1, 2025** [FLUX.1-Krea-dev](https://www.modelscope.cn/models/black-forest-labs/FLUX.1-Krea-dev) open-sourced, a text-to-image model focused on aesthetic photography. We provided comprehensive support in a timely manner, including low VRAM layer-by-layer offload, LoRA training, and full training. For more details, please refer to [./examples/flux/](./examples/flux/).
+
+- **July 28, 2025** Wan 2.2 open-sourced. We provided comprehensive support in a timely manner, including low VRAM layer-by-layer offload, FP8 quantization, sequence parallelism, LoRA training, and full training. For more details, please refer to [./examples/wanvideo/](./examples/wanvideo/).
+
+- **July 11, 2025** We propose Nexus-Gen, a unified framework that combines the language reasoning capabilities of Large Language Models (LLMs) with the image generation capabilities of diffusion models. This framework supports seamless image understanding, generation, and editing tasks.
+ - Paper: [Nexus-Gen: Unified Image Understanding, Generation, and Editing via Prefilled Autoregression in Shared Embedding Space](https://arxiv.org/pdf/2504.21356)
+ - GitHub Repository: https://github.com/modelscope/Nexus-Gen
+ - Model: [ModelScope](https://www.modelscope.cn/models/DiffSynth-Studio/Nexus-GenV2), [HuggingFace](https://huggingface.co/modelscope/Nexus-GenV2)
+ - Training Dataset: [ModelScope Dataset](https://www.modelscope.cn/datasets/DiffSynth-Studio/Nexus-Gen-Training-Dataset)
+ - Online Experience: [ModelScope Nexus-Gen Studio](https://www.modelscope.cn/studios/DiffSynth-Studio/Nexus-Gen)
+
+- **June 15, 2025** ModelScope's official evaluation framework [EvalScope](https://github.com/modelscope/evalscope) now supports text-to-image generation evaluation. Please refer to the [best practices](https://evalscope.readthedocs.io/zh-cn/latest/best_practice/t2i_eval.html) guide to try it out.
+
+- **March 25, 2025** Our new open-source project [DiffSynth-Engine](https://github.com/modelscope/DiffSynth-Engine) is now open-sourced! Focused on stable model deployment, targeting industry, providing better engineering support, higher computational performance, and more stable features.
+
+- **March 31, 2025** We support InfiniteYou, a face feature preservation method for FLUX. More details can be found in [./examples/InfiniteYou/](./examples/InfiniteYou/).
+
+- **March 13, 2025** We support HunyuanVideo-I2V, the image-to-video generation version of Tencent's open-source HunyuanVideo. More details can be found in [./examples/HunyuanVideo/](./examples/HunyuanVideo/).
+
+- **February 25, 2025** We support Wan-Video, a series of state-of-the-art video synthesis models open-sourced by Alibaba. See [./examples/wanvideo/](./examples/wanvideo/).
+
+- **February 17, 2025** We support [StepVideo](https://modelscope.cn/models/stepfun-ai/stepvideo-t2v/summary)! Advanced video synthesis model! See [./examples/stepvideo](./examples/stepvideo/).
+
+- **December 31, 2024** We propose EliGen, a new framework for entity-level controlled text-to-image generation, supplemented with an inpainting fusion pipeline, extending its capabilities to image inpainting tasks. EliGen can seamlessly integrate existing community models such as IP-Adapter and In-Context LoRA, enhancing their versatility. For more details, see [./examples/EntityControl](./examples/EntityControl/).
+ - Paper: [EliGen: Entity-Level Controlled Image Generation with Regional Attention](https://arxiv.org/abs/2501.01097)
+ - Model: [ModelScope](https://www.modelscope.cn/models/DiffSynth-Studio/Eligen), [HuggingFace](https://huggingface.co/modelscope/EliGen)
+ - Online Experience: [ModelScope EliGen Studio](https://www.modelscope.cn/studios/DiffSynth-Studio/EliGen)
+ - Training Dataset: [EliGen Train Set](https://www.modelscope.cn/datasets/DiffSynth-Studio/EliGenTrainSet)
+
+- **December 19, 2024** We implemented advanced VRAM management for HunyuanVideo, enabling video generation with resolutions of 129x720x1280 on 24GB VRAM or 129x512x384 on just 6GB VRAM. More details can be found in [./examples/HunyuanVideo/](./examples/HunyuanVideo/).
+
+- **December 18, 2024** We propose ArtAug, a method to improve text-to-image models through synthesis-understanding interaction. We trained an ArtAug enhancement module for FLUX.1-dev in LoRA format. This model incorporates the aesthetic understanding of Qwen2-VL-72B into FLUX.1-dev, thereby improving the quality of generated images.
+ - Paper: https://arxiv.org/abs/2412.12888
+ - Example: https://github.com/modelscope/DiffSynth-Studio/tree/main/examples/ArtAug
+ - Model: [ModelScope](https://www.modelscope.cn/models/DiffSynth-Studio/ArtAug-lora-FLUX.1dev-v1), [HuggingFace](https://huggingface.co/ECNU-CILab/ArtAug-lora-FLUX.1dev-v1)
+ - Demo: [ModelScope](https://modelscope.cn/aigc/imageGeneration?tab=advanced&versionId=7228&modelType=LoRA&sdVersion=FLUX_1&modelUrl=modelscope%3A%2F%2FDiffSynth-Studio%2FArtAug-lora-FLUX.1dev-v1%3Frevision%3Dv1.0), HuggingFace (coming soon)
+
+- **October 25, 2024** We provide extensive FLUX ControlNet support. This project supports many different ControlNet models and can be freely combined, even if their structures are different. Additionally, ControlNet models are compatible with high-resolution optimization and partition control technologies, enabling very powerful controllable image generation. See [`./examples/ControlNet/`](./examples/ControlNet/).
+
+- **October 8, 2024** We released extended LoRAs based on CogVideoX-5B and ExVideo. You can download this model from [ModelScope](https://modelscope.cn/models/ECNU-CILab/ExVideo-CogVideoX-LoRA-129f-v1) or [HuggingFace](https://huggingface.co/ECNU-CILab/ExVideo-CogVideoX-LoRA-129f-v1).
+
+- **August 22, 2024** This project now supports CogVideoX-5B. See [here](/examples/video_synthesis/). We provide several interesting features for this text-to-video model, including:
+ - Text-to-video
+ - Video editing
+ - Self super-resolution
+ - Video interpolation
+
+- **August 22, 2024** We implemented an interesting brush feature that supports all text-to-image models. Now you can create stunning images with the assistance of AI using the brush!
+ - Use it in our [WebUI](#usage-in-webui).
+
+- **August 21, 2024** DiffSynth-Studio now supports FLUX.
+ - Enable CFG and high-resolution inpainting to improve visual quality. See [here](/examples/image_synthesis/README.md)
+ - LoRA, ControlNet, and other addon models will be released soon.
+
+- **June 21, 2024** We propose ExVideo, a post-training fine-tuning technique aimed at enhancing the capabilities of video generation models. We extended Stable Video Diffusion to achieve long video generation of up to 128 frames.
+ - [Project Page](https://ecnu-cilab.github.io/ExVideoProjectPage/)
+ - Source code has been released in this repository. See [`examples/ExVideo`](./examples/ExVideo/).
+ - Model has been released at [HuggingFace](https://huggingface.co/ECNU-CILab/ExVideo-SVD-128f-v1) and [ModelScope](https://modelscope.cn/models/ECNU-CILab/ExVideo-SVD-128f-v1).
+ - Technical report has been released at [arXiv](https://arxiv.org/abs/2406.14130).
+ - You can try ExVideo in this [demo](https://huggingface.co/spaces/modelscope/ExVideo-SVD-128f-v1)!
+
+- **June 13, 2024** DiffSynth Studio has migrated to ModelScope. The development team has also transitioned from "me" to "us". Of course, I will still participate in subsequent development and maintenance work.
+
+- **January 29, 2024** We propose Diffutoon, an excellent cartoon coloring solution.
+ - [Project Page](https://ecnu-cilab.github.io/DiffutoonProjectPage/)
+ - Source code has been released in this project.
+ - Technical report (IJCAI 2024) has been released at [arXiv](https://arxiv.org/abs/2401.16224).
+
+- **December 8, 2023** We decided to initiate a new project aimed at unleashing the potential of diffusion models, especially in video synthesis. The development work of this project officially began.
+
+- **November 15, 2023** We propose FastBlend, a powerful video deflickering algorithm.
+ - sd-webui extension has been released at [GitHub](https://github.com/Artiprocher/sd-webui-fastblend).
+ - Demonstration videos have been showcased on Bilibili, including three tasks:
+ - [Video Deflickering](https://www.bilibili.com/video/BV1d94y1W7PE)
+ - [Video Interpolation](https://www.bilibili.com/video/BV1Lw411m71p)
+ - [Image-Driven Video Rendering](https://www.bilibili.com/video/BV1RB4y1Z7LF)
+ - Technical report has been released at [arXiv](https://arxiv.org/abs/2311.09265).
+ - Unofficial ComfyUI extensions developed by other users have been released at [GitHub](https://github.com/AInseven/ComfyUI-fastblend).
+
+- **October 1, 2023** We released an early version of the project named FastSDXL. This was an initial attempt to build a diffusion engine.
+ - Source code has been released at [GitHub](https://github.com/Artiprocher/FastSDXL).
+ - FastSDXL includes a trainable OLSS scheduler to improve efficiency.
+ - The original repository of OLSS is located [here](https://github.com/alibaba/EasyNLP/tree/master/diffusion/olss_scheduler).
+ - Technical report (CIKM 2023) has been released at [arXiv](https://arxiv.org/abs/2305.14677).
+ - Demonstration video has been released at [Bilibili](https://www.bilibili.com/video/BV1w8411y7uj).
+ - Since OLSS requires additional training, we did not implement it in this project.
+
+- **August 29, 2023** We propose DiffSynth, a video synthesis framework.
+ - [Project Page](https://ecnu-cilab.github.io/DiffSynth.github.io/).
+ - Source code has been released at [EasyNLP](https://github.com/alibaba/EasyNLP/tree/master/diffusion/DiffSynth).
+ - Technical report (ECML PKDD 2024) has been released at [arXiv](https://arxiv.org/abs/2308.03463).
+
+
+
+## Installation
+
+Install from source (recommended):
+
+```
+git clone https://github.com/modelscope/DiffSynth-Studio.git
+cd DiffSynth-Studio
+pip install -e .
+```
+
+For more installation methods and instructions for non-NVIDIA GPUs, please refer to the [Installation Guide](/docs/en/Pipeline_Usage/Setup.md).
+
+
+Download Source Configuration
+
+> Before model inference and training, you can configure the model download source and other options through [environment variables](/docs/en/Pipeline_Usage/Environment_Variables.md).
+>
+> This project downloads models from [ModelScope](https://modelscope.cn/) by default. For users outside China, you can download models from the [ModelScope International](https://modelscope.ai) site with the following configuration:
+>
+> ```shell
+> export MODELSCOPE_ENDPOINT=https://modelscope.ai
+> ```
+>
+> To download models from [HuggingFace](https://huggingface.co/), please modify the [environment variables](/docs/en/Pipeline_Usage/Environment_Variables.md). Note that model IDs may differ across platforms:
+>
+> ```shell
+> export DIFFSYNTH_DOWNLOAD_SOURCE="huggingface"
+> ```
+
+
+
+## Basic Framework
+
+DiffSynth-Studio redesigns the inference and training pipelines for mainstream Diffusion models (including FLUX, Wan, etc.), enabling efficient memory management and flexible model training.
+
+Quick start: experience popular and the latest models:
+
+| Architecture | Model ID | Inference | Low VRAM Inference | Full Training | Validation After Full Training | LoRA Training | Validation After LoRA Training |
+|-|-|-|-|-|-|-|-|
+| MiniMax-H3 | [MiniMax/MiniMax-H3: FL2VA](https://www.modelscope.cn/models/MiniMax/MiniMax-H3) | [code](/examples/minimax_h3/model_inference/MiniMax-H3-FL2VA.py) | [code](/examples/minimax_h3/model_inference_low_vram/MiniMax-H3-FL2VA.py) | [code](/examples/minimax_h3/model_training/full/MiniMax-H3-FL2VA.sh) | [code](/examples/minimax_h3/model_training/validate_full/MiniMax-H3-FL2VA.py) | [code](/examples/minimax_h3/model_training/lora/MiniMax-H3-FL2VA.sh) | [code](/examples/minimax_h3/model_training/validate_lora/MiniMax-H3-FL2VA.py) |
+| MiniMax-H3 | [DiffSynth-Studio/MiniMax-H3-NF4: FL2VA pruned](https://www.modelscope.cn/models/DiffSynth-Studio/MiniMax-H3-NF4) | [code](/examples/minimax_h3/model_inference/MiniMax-H3-NF4-Pruned-FL2VA.py) | [code](/examples/minimax_h3/model_inference_low_vram/MiniMax-H3-NF4-Pruned-FL2VA.py) | - | - | [code](/examples/minimax_h3/model_training/lora/MiniMax-H3-NF4-Pruned-FL2VA.sh) | [code](/examples/minimax_h3/model_training/validate_lora/MiniMax-H3-NF4-Pruned-FL2VA.py) |
+| ACE-Step | [ACE-Step/acestep-v15-xl-sft](https://www.modelscope.cn/models/ACE-Step/acestep-v15-xl-sft) | [code](/examples/ace_step/model_inference/acestep-v15-xl-sft.py) | [code](/examples/ace_step/model_inference_low_vram/acestep-v15-xl-sft.py) | [code](/examples/ace_step/model_training/full/acestep-v15-xl-sft.sh) | [code](/examples/ace_step/model_training/validate_full/acestep-v15-xl-sft.py) | [code](/examples/ace_step/model_training/lora/acestep-v15-xl-sft.sh) | [code](/examples/ace_step/model_training/validate_lora/acestep-v15-xl-sft.py) |
+| Z-Image | [Tongyi-MAI/Z-Image-Turbo](https://www.modelscope.cn/models/Tongyi-MAI/Z-Image-Turbo) | [code](/examples/z_image/model_inference/Z-Image-Turbo.py) | [code](/examples/z_image/model_inference_low_vram/Z-Image-Turbo.py) | [code](/examples/z_image/model_training/full/Z-Image-Turbo.sh) | [code](/examples/z_image/model_training/validate_full/Z-Image-Turbo.py) | [code](/examples/z_image/model_training/lora/Z-Image-Turbo.sh) | [code](/examples/z_image/model_training/validate_lora/Z-Image-Turbo.py) |
+| Krea-2 | [krea/Krea-2-Raw](https://www.modelscope.cn/models/krea/Krea-2-Raw) | [code](/examples/krea2/model_inference/Krea-2-Raw.py) | [code](/examples/krea2/model_inference_low_vram/Krea-2-Raw.py) | [code](/examples/krea2/model_training/full/Krea-2-Raw.sh) | [code](/examples/krea2/model_training/validate_full/Krea-2-Raw.py) | [code](/examples/krea2/model_training/lora/Krea-2-Raw.sh) | [code](/examples/krea2/model_training/validate_lora/Krea-2-Raw.py) |
+| Krea-2 | [krea/Krea-2-Turbo](https://www.modelscope.cn/models/krea/Krea-2-Turbo) | [code](/examples/krea2/model_inference/Krea-2-Turbo.py) | [code](/examples/krea2/model_inference_low_vram/Krea-2-Turbo.py) | [code](/examples/krea2/model_training/full/Krea-2-Turbo.sh) | [code](/examples/krea2/model_training/validate_full/Krea-2-Turbo.py) | [code](/examples/krea2/model_training/lora/Krea-2-Turbo.sh) | [code](/examples/krea2/model_training/validate_lora/Krea-2-Turbo.py) |
+
+Model overview:
+
+- Image generation
+ - Boogu-Image: [Documentation](https://diffsynth-studio-doc.readthedocs.io/en/latest/Model_Details/Boogu-Image.html), [Example code](/examples/boogu_image/)
+ - Krea-2: [Documentation](https://diffsynth-studio-doc.readthedocs.io/en/latest/Model_Details/Krea-2.html), [Example code](/examples/krea2/)
+ - Ideogram 4: [Documentation](https://diffsynth-studio-doc.readthedocs.io/en/latest/Model_Details/Ideogram-4.html), [Example code](/examples/ideogram4/)
+ - HiDream-O1-Image: [Documentation](https://diffsynth-studio-doc.readthedocs.io/en/latest/Model_Details/HiDream-O1-Image.html), [Example code](/examples/hidream_o1_image/)
+ - JoyAI-Image: [Documentation](https://diffsynth-studio-doc.readthedocs.io/en/latest/Model_Details/JoyAI-Image.html), [Example code](/examples/joyai_image/)
+ - ERNIE-Image: [Documentation](https://diffsynth-studio-doc.readthedocs.io/en/latest/Model_Details/ERNIE-Image.html), [Example code](/examples/ernie_image/)
+ - FLUX.2: [Documentation](https://diffsynth-studio-doc.readthedocs.io/en/latest/Model_Details/FLUX2.html), [Example code](/examples/flux2/)
+ - Z-Image: [Documentation](https://diffsynth-studio-doc.readthedocs.io/en/latest/Model_Details/Z-Image.html), [Example code](/examples/z_image/)
+ - Anima: [Documentation](https://diffsynth-studio-doc.readthedocs.io/en/latest/Model_Details/Anima.html), [Example code](/examples/anima/)
+ - Qwen-Image: [Documentation](https://diffsynth-studio-doc.readthedocs.io/en/latest/Model_Details/Qwen-Image.html), [Example code](/examples/qwen_image/)
+ - FLUX.1: [Documentation](https://diffsynth-studio-doc.readthedocs.io/en/latest/Model_Details/FLUX.html), [Example code](/examples/flux/)
+ - Stable Diffusion XL: [Documentation](https://diffsynth-studio-doc.readthedocs.io/en/latest/Model_Details/Stable-Diffusion-XL.html), [Example code](/examples/stable_diffusion_xl/)
+ - Stable Diffusion: [Documentation](https://diffsynth-studio-doc.readthedocs.io/en/latest/Model_Details/Stable-Diffusion.html), [Example code](/examples/stable_diffusion/)
+- Video generation
+ - MiniMax-H3: [Documentation](https://diffsynth-studio-doc.readthedocs.io/en/latest/Model_Details/MiniMax-H3.html), [Example code](/examples/minimax_h3/)
+ - LingBot-Video: [Documentation](https://diffsynth-studio-doc.readthedocs.io/en/latest/Model_Details/LingBot-Video.html), [Example code](/examples/lingbot_video/)
+ - LTX-2: [Documentation](https://diffsynth-studio-doc.readthedocs.io/en/latest/Model_Details/LTX-2.html), [Example code](/examples/ltx2/)
+ - Wan: [Documentation](https://diffsynth-studio-doc.readthedocs.io/en/latest/Model_Details/Wan.html), [Example code](/examples/wanvideo/)
+- Audio generation
+ - MiniMax-Music3: [Documentation](https://diffsynth-studio-doc.readthedocs.io/en/latest/Model_Details/MiniMax-Music3.html), [Example code](/examples/minimax_music3/)
+ - ACE-Step: [Documentation](https://diffsynth-studio-doc.readthedocs.io/en/latest/Model_Details/ACE-Step.html), [Example code](/examples/ace_step/)
+- Evaluation models: [Documentation](https://diffsynth-studio-doc.readthedocs.io/en/latest/Model_Details/Image-Quality-Metrics.html), [Example code](/examples/image_quality_metric/)
+
+[View all supported models](#all-supported-models)
+
+## Innovative Achievements
+
+We believe that a well-developed open-source code framework can lower the threshold for technical exploration. We have achieved many interesting technologies based on this codebase. Perhaps you also have many wild ideas, and with DiffSynth-Studio, you can quickly realize these ideas.
+
+
+
+TreeAdapter: A Model System Built from Structured LoRAs
+
+> A model system composed of 10,000+ LoRAs challenges the precise generation of tens of thousands of rare species.
+
+* Paper: [TreeAdapter: Hierarchical Taxonomy-Guided Adapter Composition for Fine-Grained Species Image Generation](https://arxiv.org/abs/2607.24215)
+* Model: [ModelScope](https://modelscope.cn/models/DiffSynth-Studio/TreeAdapter-KleinBase4B)
+
+
+
+
+
+
+
+
+Image-to-LoRA: Compressing Model Training into Model Inference
+
+> True Meta Learning: feed a dataset of images into one end of the model, and the trained LoRA model comes out the other end.
+
+* Paper: [Compressing Image Style Training into a Single Model Forward](https://arxiv.org/abs/2606.13809)
+* Model:
+ * [DiffSynth-Studio/ZImage-i2L-v2](https://modelscope.cn/models/DiffSynth-Studio/ZImage-i2L-v2): Adapted for the Z-Image model
+ * [DiffSynth-Studio/KleinBase4B-i2L-v2](https://modelscope.cn/models/DiffSynth-Studio/KleinBase4B-i2L-v2): Adapted for the FLUX.2-klein-base-4B model
+ * [DiffSynth-Studio/HidreamO1-i2L-v2](https://modelscope.cn/models/DiffSynth-Studio/HidreamO1-i2L-v2): Adapted for the Hidream-O1-Image model
+
+|Input example 1|Output example 1|Input example 2|Output example 2|
+|-|-|-|-|
+|||||
+
+
+
+
+
+
+Diffusion-Templates: A Plugin-Based Controllable Generation Framework
+
+> One framework that turns every controllable generation capability into a plugin, allowing multiple models to combine and emerge with rich generation capabilities.
+
+* Paper: [Diffusion Templates: A Unified Plugin Framework for Controllable Diffusion](https://arxiv.org/abs/2604.24351)
+* Project homepage: [GitHub](https://modelscope.github.io/diffusion-templates-web/)
+* Documentation: [English Version](https://diffsynth-studio-doc.readthedocs.io/en/latest/Diffusion_Templates/Introducing_Diffusion_Templates.html) | [Chinese Version](https://diffsynth-studio-doc.readthedocs.io/zh-cn/latest/Diffusion_Templates/Introducing_Diffusion_Templates.html)
+* Online demo: [ModelScope](https://modelscope.cn/studios/DiffSynth-Studio/Diffusion-Templates)
+* Model collections: [ModelScope](https://modelscope.cn/collections/DiffSynth-Studio/KleinBase4B-Templates) | [ModelScope International](https://modelscope.ai/collections/DiffSynth-Studio/KleinBase4B-Templates) | [HuggingFace](https://huggingface.co/collections/DiffSynth-Studio/kleinbase4b-templates)
+* Datasets: [ModelScope](https://modelscope.cn/collections/DiffSynth-Studio/ImagePulseV2) | [ModelScope International](https://modelscope.cn/collections/DiffSynth-Studio/ImagePulseV2) | [HuggingFace](https://huggingface.co/collections/DiffSynth-Studio/imagepulsev2)
+
+|Reference image|Local editing|Style transfer|Sharpness enhancement|
+|-|-|-|-|
+|||||
+
+
+
+
+
+
+
+Spectral Evolution Search: Efficient Inference-Time Scaling for Reward-Aligned Image Generation
+
+> Trade inference time for higher quality of generated content.
+
+- Paper: [Spectral Evolution Search: Efficient Inference-Time Scaling for Reward-Aligned Image Generation
+](https://arxiv.org/abs/2602.03208)
+- Sample Code: [/docs/en/Research_Tutorial/inference_time_scaling.md](/docs/en/Research_Tutorial/inference_time_scaling.md)
+
+|FLUX.1-dev|FLUX.1-dev + SES|Qwen-Image|Qwen-Image + SES|
+|-|-|-|-|
+|||||
+
+
+
+
+
+
+VIRAL: Visual In-Context Reasoning via Analogy in Diffusion Transformers
+
+> Transform image 3 into image 4 based on the change from image 1 to image 2 — the emergent capability of image editing models.
+
+- Paper: [VIRAL: Visual In-Context Reasoning via Analogy in Diffusion Transformers
+](https://arxiv.org/abs/2602.03210)
+- Sample code: [/examples/qwen_image/model_inference/Qwen-Image-Edit-2511-ICEdit.py](/examples/qwen_image/model_inference/Qwen-Image-Edit-2511-ICEdit.py)
+- Model: [ModelScope](https://www.modelscope.cn/models/DiffSynth-Studio/Qwen-Image-Edit-2511-ICEdit-LoRA)
+
+|Example 1|Example 2|Query|Output|
+|-|-|-|-|
+|||||
+
+
+
+
+
+
+AttriCtrl: Attribute Intensity Control for Image Generation Models
+
+> Numerical attributes can also precisely control image generation models.
+
+- Paper: [AttriCtrl: Fine-Grained Control of Aesthetic Attribute Intensity in Diffusion Models](https://arxiv.org/abs/2508.02151)
+- Sample Code: [/examples/flux/model_inference/FLUX.1-dev-AttriCtrl.py](/examples/flux/model_inference/FLUX.1-dev-AttriCtrl.py)
+- Model: [ModelScope](https://www.modelscope.cn/models/DiffSynth-Studio/AttriCtrl-FLUX.1-Dev)
+
+|brightness scale = 0.1|brightness scale = 0.3|brightness scale = 0.5|brightness scale = 0.7|brightness scale = 0.9|
+|-|-|-|-|-|
+||||||
+
+
+
+
+
+
+AutoLoRA: Automated LoRA Retrieval and Fusion
+
+> LoRA is a product that unifies needs and solutions — how can we make better use of these LoRAs?
+
+- Paper: [AutoLoRA: Automatic LoRA Retrieval and Fine-Grained Gated Fusion for Text-to-Image Generation](https://arxiv.org/abs/2508.02107)
+- Sample Code: [/examples/flux/model_inference/FLUX.1-dev-LoRA-Fusion.py](/examples/flux/model_inference/FLUX.1-dev-LoRA-Fusion.py)
+- Model: [ModelScope](https://www.modelscope.cn/models/DiffSynth-Studio/LoRAFusion-preview-FLUX.1-dev)
+
+||[LoRA 1](https://modelscope.cn/models/cancel13/cxsk)|[LoRA 2](https://modelscope.cn/models/wy413928499/xuancai2)|[LoRA 3](https://modelscope.cn/models/DiffSynth-Studio/ArtAug-lora-FLUX.1dev-v1)|[LoRA 4](https://modelscope.cn/models/hongyanbujian/JPL)|
+|-|-|-|-|-|
+|[LoRA 1](https://modelscope.cn/models/cancel13/cxsk) |||||
+|[LoRA 2](https://modelscope.cn/models/wy413928499/xuancai2) |||||
+|[LoRA 3](https://modelscope.cn/models/DiffSynth-Studio/ArtAug-lora-FLUX.1dev-v1) |||||
+|[LoRA 4](https://modelscope.cn/models/hongyanbujian/JPL) |||||
+
+
+
+
+
+
+Nexus-Gen: Unified Architecture for Image Understanding, Generation, and Editing
+
+> What happens when a single model combines image understanding, generation, and editing capabilities?
+
+- Detailed Page: https://github.com/modelscope/Nexus-Gen
+- Paper: [Nexus-Gen: Unified Image Understanding, Generation, and Editing via Prefilled Autoregression in Shared Embedding Space](https://arxiv.org/pdf/2504.21356)
+- Model: [ModelScope](https://www.modelscope.cn/models/DiffSynth-Studio/Nexus-GenV2), [HuggingFace](https://huggingface.co/modelscope/Nexus-GenV2)
+- Dataset: [ModelScope Dataset](https://www.modelscope.cn/datasets/DiffSynth-Studio/Nexus-Gen-Training-Dataset)
+- Online Experience: [ModelScope Nexus-Gen Studio](https://www.modelscope.cn/studios/DiffSynth-Studio/Nexus-Gen)
+
+
+
+
+
+
+
+
+ArtAug: Aesthetic Enhancement for Image Generation Models
+
+> A single LoRA that significantly enhances detail and aesthetics.
+
+- Detailed Page: [./examples/ArtAug/](./examples/ArtAug/)
+- Paper: [ArtAug: Enhancing Text-to-Image Generation through Synthesis-Understanding Interaction](https://arxiv.org/abs/2412.12888)
+- Model: [ModelScope](https://www.modelscope.cn/models/DiffSynth-Studio/ArtAug-lora-FLUX.1dev-v1), [HuggingFace](https://huggingface.co/ECNU-CILab/ArtAug-lora-FLUX.1dev-v1)
+- Online Experience: [ModelScope AIGC Tab](https://www.modelscope.cn/aigc/imageGeneration?tab=advanced&versionId=7228&modelType=LoRA&sdVersion=FLUX_1&modelUrl=modelscope%3A%2F%2FDiffSynth-Studio%2FArtAug-lora-FLUX.1dev-v1%3Frevision%3Dv1.0)
+
+|FLUX.1-dev|FLUX.1-dev + ArtAug LoRA|
+|-|-|
+|||
+
+
+
+
+
+
+EliGen: Precise Image Partition Control
+
+> How can region-based layers control the position of content in an image?
+
+- Paper: [EliGen: Entity-Level Controlled Image Generation with Regional Attention](https://arxiv.org/abs/2501.01097)
+- Sample Code: [/examples/flux/model_inference/FLUX.1-dev-EliGen.py](/examples/flux/model_inference/FLUX.1-dev-EliGen.py)
+- Model: [ModelScope](https://www.modelscope.cn/models/DiffSynth-Studio/Eligen), [HuggingFace](https://huggingface.co/modelscope/EliGen)
+- Online Experience: [ModelScope EliGen Studio](https://www.modelscope.cn/studios/DiffSynth-Studio/EliGen)
+- Dataset: [EliGen Train Set](https://www.modelscope.cn/datasets/DiffSynth-Studio/EliGenTrainSet)
+
+|Entity Control Region|Generated Image|
+|-|-|
+|||
+
+
+
+
+
+
+ExVideo: Extended Training for Video Generation Models
+
+> If a video generation model can only generate 25 frames, how can we make it generate longer videos?
+
+- Project Page: [Project Page](https://ecnu-cilab.github.io/ExVideoProjectPage/)
+- Paper: [ExVideo: Extending Video Diffusion Models via Parameter-Efficient Post-Tuning](https://arxiv.org/abs/2406.14130)
+- Sample Code: Please refer to the [older version](https://github.com/modelscope/DiffSynth-Studio/tree/afd101f3452c9ecae0c87b79adfa2e22d65ffdc3/examples/ExVideo)
+- Model: [ModelScope](https://modelscope.cn/models/ECNU-CILab/ExVideo-SVD-128f-v1), [HuggingFace](https://huggingface.co/ECNU-CILab/ExVideo-SVD-128f-v1)
+
+https://github.com/modelscope/DiffSynth-Studio/assets/35051019/d97f6aa9-8064-4b5b-9d49-ed6001bb9acc
+
+
+
+
+
+
+Diffutoon: High-Resolution Anime-Style Video Rendering
+
+> I don't care what you say, I just love anime!
+
+- Project Page: [Project Page](https://ecnu-cilab.github.io/DiffutoonProjectPage/)
+- Paper: [Diffutoon: High-Resolution Editable Toon Shading via Diffusion Models](https://arxiv.org/abs/2401.16224)
+- Sample Code: Please refer to the [older version](https://github.com/modelscope/DiffSynth-Studio/tree/afd101f3452c9ecae0c87b79adfa2e22d65ffdc3/examples/Diffutoon)
+
+The Diffutoon examples are from the pre-2.0 codebase and are not included in the current `main` branch. Use the older source tree above with the corresponding older version of DiffSynth-Studio.
+
+https://github.com/Artiprocher/DiffSynth-Studio/assets/35051019/b54c05c5-d747-4709-be5e-b39af82404dd
+
+
+
+
+
+
+DiffSynth: The Original Version of This Project
+
+> In the era before video generation models, how could image generation models be used to process videos?
+
+- Project Page: [Project Page](https://ecnu-cilab.github.io/DiffSynth.github.io/)
+- Paper: [DiffSynth: Latent In-Iteration Deflickering for Realistic Video Synthesis](https://arxiv.org/abs/2308.03463)
+- Sample Code: Please refer to the [older version](https://github.com/modelscope/DiffSynth-Studio/tree/afd101f3452c9ecae0c87b79adfa2e22d65ffdc3/examples/diffsynth)
+
+https://github.com/Artiprocher/DiffSynth-Studio/assets/35051019/59fb2f7b-8de0-4481-b79f-0c3a7361a1ea
+
+
+
+## Contact Us
+
+|Discord：https://discord.gg/Mm9suEeUDc|
+|-|
+|
|
+
+
+## All Supported Models
+
+| Architecture | Model ID | Inference | Low VRAM Inference | Full Training | Validation After Full Training | LoRA Training | Validation After LoRA Training |
+|-|-|-|-|-|-|-|-|
+| MiniMax-Music3 | [MiniMax/MiniMax-Music3](https://www.modelscope.cn/models/MiniMax/MiniMax-Music3) | [code](/examples/minimax_music3/model_inference/MiniMax-Music3.py) | [code](/examples/minimax_music3/model_inference_low_vram/MiniMax-Music3.py) | — | — | — | — |
+| MiniMax-H3 | [MiniMax/MiniMax-H3: FL2VA](https://www.modelscope.cn/models/MiniMax/MiniMax-H3) | [code](/examples/minimax_h3/model_inference/MiniMax-H3-FL2VA.py) | [code](/examples/minimax_h3/model_inference_low_vram/MiniMax-H3-FL2VA.py) | [code](/examples/minimax_h3/model_training/full/MiniMax-H3-FL2VA.sh) | [code](/examples/minimax_h3/model_training/validate_full/MiniMax-H3-FL2VA.py) | [code](/examples/minimax_h3/model_training/lora/MiniMax-H3-FL2VA.sh) | [code](/examples/minimax_h3/model_training/validate_lora/MiniMax-H3-FL2VA.py) |
+| MiniMax-H3 | [MiniMax/MiniMax-H3: Ref2VA](https://www.modelscope.cn/models/MiniMax/MiniMax-H3) | [code](/examples/minimax_h3/model_inference/MiniMax-H3-Ref2VA.py) | [code](/examples/minimax_h3/model_inference_low_vram/MiniMax-H3-Ref2VA.py) | [code](/examples/minimax_h3/model_training/full/MiniMax-H3-Ref2VA.sh) | [code](/examples/minimax_h3/model_training/validate_full/MiniMax-H3-Ref2VA.py) | [code](/examples/minimax_h3/model_training/lora/MiniMax-H3-Ref2VA.sh) | [code](/examples/minimax_h3/model_training/validate_lora/MiniMax-H3-Ref2VA.py) |
+| MiniMax-H3 | [MiniMax/MiniMax-H3: Retake](https://www.modelscope.cn/models/MiniMax/MiniMax-H3) | [code](/examples/minimax_h3/model_inference/MiniMax-H3-Retake.py) | [code](/examples/minimax_h3/model_inference_low_vram/MiniMax-H3-Retake.py) | - | - | - | - |
+| MiniMax-H3 | [DiffSynth-Studio/MiniMax-H3-NF4: FL2VA](https://www.modelscope.cn/models/DiffSynth-Studio/MiniMax-H3-NF4) | [code](/examples/minimax_h3/model_inference/MiniMax-H3-NF4-FL2VA.py) | [code](/examples/minimax_h3/model_inference_low_vram/MiniMax-H3-NF4-FL2VA.py) | - | - | [code](/examples/minimax_h3/model_training/lora/MiniMax-H3-NF4-FL2VA.sh) | [code](/examples/minimax_h3/model_training/validate_lora/MiniMax-H3-NF4-FL2VA.py) |
+| MiniMax-H3 | [DiffSynth-Studio/MiniMax-H3-NF4: Ref2VA](https://www.modelscope.cn/models/DiffSynth-Studio/MiniMax-H3-NF4) | [code](/examples/minimax_h3/model_inference/MiniMax-H3-NF4-Ref2VA.py) | [code](/examples/minimax_h3/model_inference_low_vram/MiniMax-H3-NF4-Ref2VA.py) | - | - | [code](/examples/minimax_h3/model_training/lora/MiniMax-H3-NF4-Ref2VA.sh) | [code](/examples/minimax_h3/model_training/validate_lora/MiniMax-H3-NF4-Ref2VA.py) |
+| MiniMax-H3 | [Comfy-Org/MiniMax-H3: FL2VA pruned](https://www.modelscope.cn/models/Comfy-Org/MiniMax-H3) | [code](/examples/minimax_h3/model_inference/MiniMax-H3-Pruned-FL2VA.py) | [code](/examples/minimax_h3/model_inference_low_vram/MiniMax-H3-Pruned-FL2VA.py) | [code](/examples/minimax_h3/model_training/full/MiniMax-H3-Pruned-FL2VA.sh) | [code](/examples/minimax_h3/model_training/validate_full/MiniMax-H3-Pruned-FL2VA.py) | [code](/examples/minimax_h3/model_training/lora/MiniMax-H3-Pruned-FL2VA.sh) | [code](/examples/minimax_h3/model_training/validate_lora/MiniMax-H3-Pruned-FL2VA.py) |
+| MiniMax-H3 | [Comfy-Org/MiniMax-H3: Ref2VA pruned](https://www.modelscope.cn/models/Comfy-Org/MiniMax-H3) | [code](/examples/minimax_h3/model_inference/MiniMax-H3-Pruned-Ref2VA.py) | [code](/examples/minimax_h3/model_inference_low_vram/MiniMax-H3-Pruned-Ref2VA.py) | [code](/examples/minimax_h3/model_training/full/MiniMax-H3-Pruned-Ref2VA.sh) | [code](/examples/minimax_h3/model_training/validate_full/MiniMax-H3-Pruned-Ref2VA.py) | [code](/examples/minimax_h3/model_training/lora/MiniMax-H3-Pruned-Ref2VA.sh) | [code](/examples/minimax_h3/model_training/validate_lora/MiniMax-H3-Pruned-Ref2VA.py) |
+| MiniMax-H3 | [DiffSynth-Studio/MiniMax-H3-NF4: FL2VA pruned](https://www.modelscope.cn/models/DiffSynth-Studio/MiniMax-H3-NF4) | [code](/examples/minimax_h3/model_inference/MiniMax-H3-NF4-Pruned-FL2VA.py) | [code](/examples/minimax_h3/model_inference_low_vram/MiniMax-H3-NF4-Pruned-FL2VA.py) | - | - | [code](/examples/minimax_h3/model_training/lora/MiniMax-H3-NF4-Pruned-FL2VA.sh) | [code](/examples/minimax_h3/model_training/validate_lora/MiniMax-H3-NF4-Pruned-FL2VA.py) |
+| MiniMax-H3 | [DiffSynth-Studio/MiniMax-H3-NF4: Ref2VA pruned](https://www.modelscope.cn/models/DiffSynth-Studio/MiniMax-H3-NF4) | [code](/examples/minimax_h3/model_inference/MiniMax-H3-NF4-Pruned-Ref2VA.py) | [code](/examples/minimax_h3/model_inference_low_vram/MiniMax-H3-NF4-Pruned-Ref2VA.py) | - | - | [code](/examples/minimax_h3/model_training/lora/MiniMax-H3-NF4-Pruned-Ref2VA.sh) | [code](/examples/minimax_h3/model_training/validate_lora/MiniMax-H3-NF4-Pruned-Ref2VA.py) |
+| MiniMax-H3 | [Comfy-Org/MiniMax-H3: FL2VA int8_convrot](https://www.modelscope.cn/models/Comfy-Org/MiniMax-H3) | [code](/examples/minimax_h3/model_inference/MiniMax-H3-Int8-ConvRot-FL2VA.py) | [code](/examples/minimax_h3/model_inference_low_vram/MiniMax-H3-Int8-ConvRot-FL2VA.py) | - | - | [code](/examples/minimax_h3/model_training/lora/MiniMax-H3-Int8-ConvRot-FL2VA.sh) | [code](/examples/minimax_h3/model_training/validate_lora/MiniMax-H3-Int8-ConvRot-FL2VA.py) |
+| MiniMax-H3 | [Comfy-Org/MiniMax-H3: Ref2VA int8_convrot](https://www.modelscope.cn/models/Comfy-Org/MiniMax-H3) | [code](/examples/minimax_h3/model_inference/MiniMax-H3-Int8-ConvRot-Ref2VA.py) | [code](/examples/minimax_h3/model_inference_low_vram/MiniMax-H3-Int8-ConvRot-Ref2VA.py) | - | - | [code](/examples/minimax_h3/model_training/lora/MiniMax-H3-Int8-ConvRot-Ref2VA.sh) | [code](/examples/minimax_h3/model_training/validate_lora/MiniMax-H3-Int8-ConvRot-Ref2VA.py) |
+| MiniMax-H3 | [Comfy-Org/MiniMax-H3: FL2VA pruned int8_convrot](https://www.modelscope.cn/models/Comfy-Org/MiniMax-H3) | [code](/examples/minimax_h3/model_inference/MiniMax-H3-Int8-ConvRot-Pruned-FL2VA.py) | [code](/examples/minimax_h3/model_inference_low_vram/MiniMax-H3-Int8-ConvRot-Pruned-FL2VA.py) | - | - | [code](/examples/minimax_h3/model_training/lora/MiniMax-H3-Int8-ConvRot-Pruned-FL2VA.sh) | [code](/examples/minimax_h3/model_training/validate_lora/MiniMax-H3-Int8-ConvRot-Pruned-FL2VA.py) |
+| MiniMax-H3 | [Comfy-Org/MiniMax-H3: Ref2VA pruned int8_convrot](https://www.modelscope.cn/models/Comfy-Org/MiniMax-H3) | [code](/examples/minimax_h3/model_inference/MiniMax-H3-Int8-ConvRot-Pruned-Ref2VA.py) | [code](/examples/minimax_h3/model_inference_low_vram/MiniMax-H3-Int8-ConvRot-Pruned-Ref2VA.py) | - | - | [code](/examples/minimax_h3/model_training/lora/MiniMax-H3-Int8-ConvRot-Pruned-Ref2VA.sh) | [code](/examples/minimax_h3/model_training/validate_lora/MiniMax-H3-Int8-ConvRot-Pruned-Ref2VA.py) |
+| MiniMax-H3 | [Comfy-Org/MiniMax-H3: FL2VA pruned fp8](https://www.modelscope.cn/models/Comfy-Org/MiniMax-H3) | [code](/examples/minimax_h3/model_inference/MiniMax-H3-FP8-Pruned-FL2VA.py) | [code](/examples/minimax_h3/model_inference_low_vram/MiniMax-H3-FP8-Pruned-FL2VA.py) | - | - | [code](/examples/minimax_h3/model_training/lora/MiniMax-H3-FP8-Pruned-FL2VA.sh) | [code](/examples/minimax_h3/model_training/validate_lora/MiniMax-H3-FP8-Pruned-FL2VA.py) |
+| MiniMax-H3 | [Comfy-Org/MiniMax-H3: Ref2VA pruned fp8](https://www.modelscope.cn/models/Comfy-Org/MiniMax-H3) | [code](/examples/minimax_h3/model_inference/MiniMax-H3-FP8-Pruned-Ref2VA.py) | [code](/examples/minimax_h3/model_inference_low_vram/MiniMax-H3-FP8-Pruned-Ref2VA.py) | - | - | [code](/examples/minimax_h3/model_training/lora/MiniMax-H3-FP8-Pruned-Ref2VA.sh) | [code](/examples/minimax_h3/model_training/validate_lora/MiniMax-H3-FP8-Pruned-Ref2VA.py) |
+| MiniMax-H3 | [lightx2v/Minimax-h3-Turbo: FL2VA 4steps](https://www.modelscope.cn/models/lightx2v/Minimax-h3-Turbo) | [code](https://github.com/modelscope/DiffSynth-Studio/blob/main/examples/minimax_h3/model_inference/MiniMax-H3-FL2VA-Turbo.py) | [code](https://github.com/modelscope/DiffSynth-Studio/blob/main/examples/minimax_h3/model_inference_low_vram/MiniMax-H3-FL2VA-Turbo.py) | - | - | - | - |
+| MiniMax-H3 | [DiffSynth-Studio/MiniMax-H3-Text-Embeddings](https://www.modelscope.cn/models/DiffSynth-Studio/MiniMax-H3-Text-Embeddings) | [code](/examples/minimax_h3/model_inference/MiniMax-H3-Text-Embeddings.py) | [code](/examples/minimax_h3/model_inference_low_vram/MiniMax-H3-Text-Embeddings.py) | [code](/examples/minimax_h3/model_training/full/MiniMax-H3-Text-Embeddings.sh) | [code](/examples/minimax_h3/model_training/validate_full/MiniMax-H3-Text-Embeddings.py) | - | - |
+| LingBot-Video | [Robbyant/lingbot-video-dense-1.3b: T2V](https://modelscope.cn/models/Robbyant/lingbot-video-dense-1.3b) | [code](/examples/lingbot_video/model_inference/lingbot-video-dense-1.3b_t2v.py) | [code](/examples/lingbot_video/model_inference_low_vram/lingbot-video-dense-1.3b_t2v.py) | [code](/examples/lingbot_video/model_training/full/lingbot-video-dense-1.3b_t2v.sh) | [code](/examples/lingbot_video/model_training/validate_full/lingbot-video-dense-1.3b_t2v.py) | [code](/examples/lingbot_video/model_training/lora/lingbot-video-dense-1.3b_t2v.sh) | [code](/examples/lingbot_video/model_training/validate_lora/lingbot-video-dense-1.3b_t2v.py) |
+| LingBot-Video | [Robbyant/lingbot-video-dense-1.3b: TI2V](https://modelscope.cn/models/Robbyant/lingbot-video-dense-1.3b) | [code](/examples/lingbot_video/model_inference/lingbot-video-dense-1.3b_ti2v.py) | [code](/examples/lingbot_video/model_inference_low_vram/lingbot-video-dense-1.3b_ti2v.py) | [code](/examples/lingbot_video/model_training/full/lingbot-video-dense-1.3b_ti2v.sh) | [code](/examples/lingbot_video/model_training/validate_full/lingbot-video-dense-1.3b_ti2v.py) | [code](/examples/lingbot_video/model_training/lora/lingbot-video-dense-1.3b_ti2v.sh) | [code](/examples/lingbot_video/model_training/validate_lora/lingbot-video-dense-1.3b_ti2v.py) |
+| LingBot-Video | [Robbyant/lingbot-video-dense-1.3b: T2I](https://modelscope.cn/models/Robbyant/lingbot-video-dense-1.3b) | [code](/examples/lingbot_video/model_inference/lingbot-video-dense-1.3b_t2i.py) | [code](/examples/lingbot_video/model_inference_low_vram/lingbot-video-dense-1.3b_t2i.py) | - | - | - | - |
+| LingBot-Video | [Robbyant/lingbot-video-moe-30b-a3b: T2V](https://modelscope.cn/models/Robbyant/lingbot-video-moe-30b-a3b) | [code](/examples/lingbot_video/model_inference/lingbot-video-moe-30b-a3b_t2v.py) | [code](/examples/lingbot_video/model_inference_low_vram/lingbot-video-moe-30b-a3b_t2v.py) | [code](/examples/lingbot_video/model_training/full/lingbot-video-moe-30b-a3b_t2v.sh) | [code](/examples/lingbot_video/model_training/validate_full/lingbot-video-moe-30b-a3b_t2v.py) | [code](/examples/lingbot_video/model_training/lora/lingbot-video-moe-30b-a3b_t2v.sh) | [code](/examples/lingbot_video/model_training/validate_lora/lingbot-video-moe-30b-a3b_t2v.py) |
+| LingBot-Video | [Robbyant/lingbot-video-moe-30b-a3b: TI2V](https://modelscope.cn/models/Robbyant/lingbot-video-moe-30b-a3b) | [code](/examples/lingbot_video/model_inference/lingbot-video-moe-30b-a3b_ti2v.py) | [code](/examples/lingbot_video/model_inference_low_vram/lingbot-video-moe-30b-a3b_ti2v.py) | [code](/examples/lingbot_video/model_training/full/lingbot-video-moe-30b-a3b_ti2v.sh) | [code](/examples/lingbot_video/model_training/validate_full/lingbot-video-moe-30b-a3b_ti2v.py) | [code](/examples/lingbot_video/model_training/lora/lingbot-video-moe-30b-a3b_ti2v.sh) | [code](/examples/lingbot_video/model_training/validate_lora/lingbot-video-moe-30b-a3b_ti2v.py) |
+| LingBot-Video | [Robbyant/lingbot-video-moe-30b-a3b: T2I](https://modelscope.cn/models/Robbyant/lingbot-video-moe-30b-a3b) | [code](/examples/lingbot_video/model_inference/lingbot-video-moe-30b-a3b_t2i.py) | [code](/examples/lingbot_video/model_inference_low_vram/lingbot-video-moe-30b-a3b_t2i.py) | - | - | - | - |
+| LingBot-Video | [Robbyant/lingbot-video-moe-30b-a3b: T2V + Refinement](https://modelscope.cn/models/Robbyant/lingbot-video-moe-30b-a3b) | [code](/examples/lingbot_video/model_inference/lingbot-video-moe-30b-a3b_t2v_refiner.py) | [code](/examples/lingbot_video/model_inference_low_vram/lingbot-video-moe-30b-a3b_t2v_refiner.py) | - | - | - | - |
+| LingBot-Video | [Robbyant/lingbot-video-moe-30b-a3b: TI2V + Refinement](https://modelscope.cn/models/Robbyant/lingbot-video-moe-30b-a3b) | [code](/examples/lingbot_video/model_inference/lingbot-video-moe-30b-a3b_ti2v_refiner.py) | [code](/examples/lingbot_video/model_inference_low_vram/lingbot-video-moe-30b-a3b_ti2v_refiner.py) | - | - | - | - |
+| ACE-Step | [ACE-Step/Ace-Step1.5](https://www.modelscope.cn/models/ACE-Step/Ace-Step1.5) | [code](/examples/ace_step/model_inference/Ace-Step1.5.py) | [code](/examples/ace_step/model_inference_low_vram/Ace-Step1.5.py) | [code](/examples/ace_step/model_training/full/Ace-Step1.5.sh) | [code](/examples/ace_step/model_training/validate_full/Ace-Step1.5.py) | [code](/examples/ace_step/model_training/lora/Ace-Step1.5.sh) | [code](/examples/ace_step/model_training/validate_lora/Ace-Step1.5.py) |
+| ACE-Step | [ACE-Step/acestep-v15-turbo-shift1](https://www.modelscope.cn/models/ACE-Step/acestep-v15-turbo-shift1) | [code](/examples/ace_step/model_inference/acestep-v15-turbo-shift1.py) | [code](/examples/ace_step/model_inference_low_vram/acestep-v15-turbo-shift1.py) | [code](/examples/ace_step/model_training/full/acestep-v15-turbo-shift1.sh) | [code](/examples/ace_step/model_training/validate_full/acestep-v15-turbo-shift1.py) | [code](/examples/ace_step/model_training/lora/acestep-v15-turbo-shift1.sh) | [code](/examples/ace_step/model_training/validate_lora/acestep-v15-turbo-shift1.py) |
+| ACE-Step | [ACE-Step/acestep-v15-turbo-shift3](https://www.modelscope.cn/models/ACE-Step/acestep-v15-turbo-shift3) | [code](/examples/ace_step/model_inference/acestep-v15-turbo-shift3.py) | [code](/examples/ace_step/model_inference_low_vram/acestep-v15-turbo-shift3.py) | [code](/examples/ace_step/model_training/full/acestep-v15-turbo-shift3.sh) | [code](/examples/ace_step/model_training/validate_full/acestep-v15-turbo-shift3.py) | [code](/examples/ace_step/model_training/lora/acestep-v15-turbo-shift3.sh) | [code](/examples/ace_step/model_training/validate_lora/acestep-v15-turbo-shift3.py) |
+| ACE-Step | [ACE-Step/acestep-v15-turbo-continuous](https://www.modelscope.cn/models/ACE-Step/acestep-v15-turbo-continuous) | [code](/examples/ace_step/model_inference/acestep-v15-turbo-continuous.py) | [code](/examples/ace_step/model_inference_low_vram/acestep-v15-turbo-continuous.py) | [code](/examples/ace_step/model_training/full/acestep-v15-turbo-continuous.sh) | [code](/examples/ace_step/model_training/validate_full/acestep-v15-turbo-continuous.py) | [code](/examples/ace_step/model_training/lora/acestep-v15-turbo-continuous.sh) | [code](/examples/ace_step/model_training/validate_lora/acestep-v15-turbo-continuous.py) |
+| ACE-Step | [ACE-Step/acestep-v15-base](https://www.modelscope.cn/models/ACE-Step/acestep-v15-base) | [code](/examples/ace_step/model_inference/acestep-v15-base.py) | [code](/examples/ace_step/model_inference_low_vram/acestep-v15-base.py) | [code](/examples/ace_step/model_training/full/acestep-v15-base.sh) | [code](/examples/ace_step/model_training/validate_full/acestep-v15-base.py) | [code](/examples/ace_step/model_training/lora/acestep-v15-base.sh) | [code](/examples/ace_step/model_training/validate_lora/acestep-v15-base.py) |
+| ACE-Step | [ACE-Step/acestep-v15-base: CoverTask](https://www.modelscope.cn/models/ACE-Step/acestep-v15-base) | [code](/examples/ace_step/model_inference/acestep-v15-base-CoverTask.py) | [code](/examples/ace_step/model_inference_low_vram/acestep-v15-base-CoverTask.py) | — | — | — | — |
+| ACE-Step | [ACE-Step/acestep-v15-base: RepaintTask](https://www.modelscope.cn/models/ACE-Step/acestep-v15-base) | [code](/examples/ace_step/model_inference/acestep-v15-base-RepaintTask.py) | [code](/examples/ace_step/model_inference_low_vram/acestep-v15-base-RepaintTask.py) | — | — | — | — |
+| ACE-Step | [ACE-Step/acestep-v15-sft](https://www.modelscope.cn/models/ACE-Step/acestep-v15-sft) | [code](/examples/ace_step/model_inference/acestep-v15-sft.py) | [code](/examples/ace_step/model_inference_low_vram/acestep-v15-sft.py) | [code](/examples/ace_step/model_training/full/acestep-v15-sft.sh) | [code](/examples/ace_step/model_training/validate_full/acestep-v15-sft.py) | [code](/examples/ace_step/model_training/lora/acestep-v15-sft.sh) | [code](/examples/ace_step/model_training/validate_lora/acestep-v15-sft.py) |
+| ACE-Step | [ACE-Step/acestep-v15-xl-base](https://www.modelscope.cn/models/ACE-Step/acestep-v15-xl-base) | [code](/examples/ace_step/model_inference/acestep-v15-xl-base.py) | [code](/examples/ace_step/model_inference_low_vram/acestep-v15-xl-base.py) | [code](/examples/ace_step/model_training/full/acestep-v15-xl-base.sh) | [code](/examples/ace_step/model_training/validate_full/acestep-v15-xl-base.py) | [code](/examples/ace_step/model_training/lora/acestep-v15-xl-base.sh) | [code](/examples/ace_step/model_training/validate_lora/acestep-v15-xl-base.py) |
+| ACE-Step | [ACE-Step/acestep-v15-xl-sft](https://www.modelscope.cn/models/ACE-Step/acestep-v15-xl-sft) | [code](/examples/ace_step/model_inference/acestep-v15-xl-sft.py) | [code](/examples/ace_step/model_inference_low_vram/acestep-v15-xl-sft.py) | [code](/examples/ace_step/model_training/full/acestep-v15-xl-sft.sh) | [code](/examples/ace_step/model_training/validate_full/acestep-v15-xl-sft.py) | [code](/examples/ace_step/model_training/lora/acestep-v15-xl-sft.sh) | [code](/examples/ace_step/model_training/validate_lora/acestep-v15-xl-sft.py) |
+| ACE-Step | [ACE-Step/acestep-v15-xl-turbo](https://www.modelscope.cn/models/ACE-Step/acestep-v15-xl-turbo) | [code](/examples/ace_step/model_inference/acestep-v15-xl-turbo.py) | [code](/examples/ace_step/model_inference_low_vram/acestep-v15-xl-turbo.py) | [code](/examples/ace_step/model_training/full/acestep-v15-xl-turbo.sh) | [code](/examples/ace_step/model_training/validate_full/acestep-v15-xl-turbo.py) | [code](/examples/ace_step/model_training/lora/acestep-v15-xl-turbo.sh) | [code](/examples/ace_step/model_training/validate_lora/acestep-v15-xl-turbo.py) |
+| ACE-Step | [DiffSynth-Studio/acestep15xlsft-lora-music](https://www.modelscope.cn/models/DiffSynth-Studio/acestep15xlsft-lora-music) | [code](/examples/ace_step/model_inference/acestep15xlsft-vocals2music.py) | [code](/examples/ace_step/model_inference_low_vram/acestep15xlsft-vocals2music.py) | [code](/examples/ace_step/model_training/full/acestep15xlsft-vocals2music.sh) | [code](/examples/ace_step/model_training/validate_full/acestep15xlsft-vocals2music.py) | - | - |
+| Boogu-Image | [Boogu/Boogu-Image-0.1-Base](https://modelscope.cn/models/Boogu/Boogu-Image-0.1-Base) | [code](/examples/boogu_image/model_inference/Boogu-Image-0.1-Base.py) | [code](/examples/boogu_image/model_inference_low_vram/Boogu-Image-0.1-Base.py) | [code](/examples/boogu_image/model_training/full/Boogu-Image-0.1-Base.sh) | [code](/examples/boogu_image/model_training/validate_full/Boogu-Image-0.1-Base.py) | [code](/examples/boogu_image/model_training/lora/Boogu-Image-0.1-Base.sh) | [code](/examples/boogu_image/model_training/validate_lora/Boogu-Image-0.1-Base.py) |
+| Boogu-Image | [Boogu/Boogu-Image-0.1-Turbo](https://modelscope.cn/models/Boogu/Boogu-Image-0.1-Turbo) | [code](/examples/boogu_image/model_inference/Boogu-Image-0.1-Turbo.py) | [code](/examples/boogu_image/model_inference_low_vram/Boogu-Image-0.1-Turbo.py) | [code](/examples/boogu_image/model_training/full/Boogu-Image-0.1-Turbo.sh) | [code](/examples/boogu_image/model_training/validate_full/Boogu-Image-0.1-Turbo.py) | [code](/examples/boogu_image/model_training/lora/Boogu-Image-0.1-Turbo.sh) | [code](/examples/boogu_image/model_training/validate_lora/Boogu-Image-0.1-Turbo.py) |
+| Boogu-Image | [Boogu/Boogu-Image-0.1-Edit](https://modelscope.cn/models/Boogu/Boogu-Image-0.1-Edit) | [code](/examples/boogu_image/model_inference/Boogu-Image-0.1-Edit.py) | [code](/examples/boogu_image/model_inference_low_vram/Boogu-Image-0.1-Edit.py) | [code](/examples/boogu_image/model_training/full/Boogu-Image-0.1-Edit.sh) | [code](/examples/boogu_image/model_training/validate_full/Boogu-Image-0.1-Edit.py) | [code](/examples/boogu_image/model_training/lora/Boogu-Image-0.1-Edit.sh) | [code](/examples/boogu_image/model_training/validate_lora/Boogu-Image-0.1-Edit.py) |
+| Krea-2 | [krea/Krea-2-Raw](https://www.modelscope.cn/models/krea/Krea-2-Raw) | [code](/examples/krea2/model_inference/Krea-2-Raw.py) | [code](/examples/krea2/model_inference_low_vram/Krea-2-Raw.py) | [code](/examples/krea2/model_training/full/Krea-2-Raw.sh) | [code](/examples/krea2/model_training/validate_full/Krea-2-Raw.py) | [code](/examples/krea2/model_training/lora/Krea-2-Raw.sh) | [code](/examples/krea2/model_training/validate_lora/Krea-2-Raw.py) |
+| Krea-2 | [krea/Krea-2-Turbo](https://www.modelscope.cn/models/krea/Krea-2-Turbo) | [code](/examples/krea2/model_inference/Krea-2-Turbo.py) | [code](/examples/krea2/model_inference_low_vram/Krea-2-Turbo.py) | [code](/examples/krea2/model_training/full/Krea-2-Turbo.sh) | [code](/examples/krea2/model_training/validate_full/Krea-2-Turbo.py) | [code](/examples/krea2/model_training/lora/Krea-2-Turbo.sh) | [code](/examples/krea2/model_training/validate_lora/Krea-2-Turbo.py) |
+| Ideogram 4 | [ideogram-ai/ideogram-4-fp8](https://www.modelscope.cn/models/ideogram-ai/ideogram-4-fp8) | [code](/examples/ideogram4/model_inference/ideogram-4-fp8.py) | - | - | - | - | - |
+| Ideogram 4 | [DiffSynth-Studio/ideogram-4-bf16-repackage](https://www.modelscope.cn/models/DiffSynth-Studio/ideogram-4-bf16-repackage) | [code](/examples/ideogram4/model_inference/ideogram-4-bf16-repackage.py) | [code](/examples/ideogram4/model_inference_low_vram/ideogram-4-bf16-repackage.py) | [code](/examples/ideogram4/model_training/full/Ideogram-4-bf16-repackage.sh) | - | [code](/examples/ideogram4/model_training/lora/Ideogram-4-bf16-repackage.sh) | [code](/examples/ideogram4/model_training/validate_lora/Ideogram-4-bf16-repackage.py) |
+| HiDream-O1-Image | [HiDream-ai/HiDream-O1-Image](https://modelscope.cn/models/HiDream-ai/HiDream-O1-Image) | [code](/examples/hidream_o1_image/model_inference/HiDream-O1-Image.py) | [code](/examples/hidream_o1_image/model_inference_low_vram/HiDream-O1-Image.py) | [code](/examples/hidream_o1_image/model_training/full/HiDream-O1-Image.sh) | [code](/examples/hidream_o1_image/model_training/validate_full/HiDream-O1-Image.py) | [code](/examples/hidream_o1_image/model_training/lora/HiDream-O1-Image.sh) | [code](/examples/hidream_o1_image/model_training/validate_lora/HiDream-O1-Image.py) |
+| HiDream-O1-Image | [HiDream-ai/HiDream-O1-Image-Dev](https://modelscope.cn/models/HiDream-ai/HiDream-O1-Image-Dev) | [code](/examples/hidream_o1_image/model_inference/HiDream-O1-Image-Dev.py) | [code](/examples/hidream_o1_image/model_inference_low_vram/HiDream-O1-Image-Dev.py) | [code](/examples/hidream_o1_image/model_training/full/HiDream-O1-Image-Dev.sh) | [code](/examples/hidream_o1_image/model_training/validate_full/HiDream-O1-Image-Dev.py) | [code](/examples/hidream_o1_image/model_training/lora/HiDream-O1-Image-Dev.sh) | [code](/examples/hidream_o1_image/model_training/validate_lora/HiDream-O1-Image-Dev.py) |
+| HiDream-O1-Image | [DiffSynth-Studio/HidreamO1-i2L-v2](https://www.modelscope.cn/models/DiffSynth-Studio/HidreamO1-i2L-v2) | [code](/examples/hidream_o1_image/model_inference/HidreamO1-i2L-v2.py) | [code](/examples/hidream_o1_image/model_inference_low_vram/HidreamO1-i2L-v2.py) | [code](/examples/hidream_o1_image/model_training/full/HidreamO1-i2L-v2.sh) | [code](/examples/hidream_o1_image/model_training/validate_full/HidreamO1-i2L-v2.py) | - | - |
+| JoyAI-Image | [jd-opensource/JoyAI-Image-Edit](https://modelscope.cn/models/jd-opensource/JoyAI-Image-Edit) | [code](/examples/joyai_image/model_inference/JoyAI-Image-Edit.py) | [code](/examples/joyai_image/model_inference_low_vram/JoyAI-Image-Edit.py) | [code](/examples/joyai_image/model_training/full/JoyAI-Image-Edit.sh) | [code](/examples/joyai_image/model_training/validate_full/JoyAI-Image-Edit.py) | [code](/examples/joyai_image/model_training/lora/JoyAI-Image-Edit.sh) | [code](/examples/joyai_image/model_training/validate_lora/JoyAI-Image-Edit.py) |
+| ERNIE-Image | [PaddlePaddle/ERNIE-Image](https://www.modelscope.cn/models/PaddlePaddle/ERNIE-Image) | [code](/examples/ernie_image/model_inference/ERNIE-Image.py) | [code](/examples/ernie_image/model_inference_low_vram/ERNIE-Image.py) | [code](/examples/ernie_image/model_training/full/ERNIE-Image.sh) | [code](/examples/ernie_image/model_training/validate_full/ERNIE-Image.py) | [code](/examples/ernie_image/model_training/lora/ERNIE-Image.sh) | [code](/examples/ernie_image/model_training/validate_lora/ERNIE-Image.py) |
+| ERNIE-Image | [PaddlePaddle/ERNIE-Image-Turbo](https://www.modelscope.cn/models/PaddlePaddle/ERNIE-Image-Turbo) | [code](/examples/ernie_image/model_inference/ERNIE-Image-Turbo.py) | [code](/examples/ernie_image/model_inference_low_vram/ERNIE-Image-Turbo.py) | — | — | — | — |
+| LTX-2 | [jd-opensource/JoyAI-Echo](https://modelscope.cn/models/jd-opensource/JoyAI-Echo) | [code](/examples/ltx2/model_inference/JoyAI-Echo-T2AV.py) | [code](/examples/ltx2/model_inference_low_vram/JoyAI-Echo-T2AV.py) | [code](/examples/ltx2/model_training/full/JoyAI-Echo-T2AV-splited.sh) | [code](/examples/ltx2/model_training/validate_full/JoyAI-Echo-T2AV.py) | [code](/examples/ltx2/model_training/lora/JoyAI-Echo-T2AV-splited.sh) | [code](/examples/ltx2/model_training/validate_lora/JoyAI-Echo-T2AV.py) |
+| LTX-2 | [Lightricks/LTX-2.3: OneStagePipeline-I2AV](https://www.modelscope.cn/models/Lightricks/LTX-2.3) | [code](/examples/ltx2/model_inference/LTX-2.3-I2AV-OneStage.py) | [code](/examples/ltx2/model_inference_low_vram/LTX-2.3-I2AV-OneStage.py) | [code](/examples/ltx2/model_training/full/LTX-2.3-I2AV-splited.sh) | [code](/examples/ltx2/model_training/validate_full/LTX-2.3-I2AV.py) | [code](/examples/ltx2/model_training/lora/LTX-2.3-I2AV-splited.sh) | [code](/examples/ltx2/model_training/validate_lora/LTX-2.3-I2AV.py) |
+| LTX-2 | [Lightricks/LTX-2.3: TwoStagePipeline-I2AV](https://www.modelscope.cn/models/Lightricks/LTX-2.3) | [code](/examples/ltx2/model_inference/LTX-2.3-I2AV-TwoStage.py) | [code](/examples/ltx2/model_inference_low_vram/LTX-2.3-I2AV-TwoStage.py) | - | - | - | - |
+| LTX-2 | [Lightricks/LTX-2.3: DistilledPipeline-I2AV](https://www.modelscope.cn/models/Lightricks/LTX-2.3) | [code](/examples/ltx2/model_inference/LTX-2.3-I2AV-DistilledPipeline.py) | [code](/examples/ltx2/model_inference_low_vram/LTX-2.3-I2AV-DistilledPipeline.py) | - | - | - | - |
+| LTX-2 | [Lightricks/LTX-2.3: OneStagePipeline-T2AV](https://www.modelscope.cn/models/Lightricks/LTX-2.3) | [code](/examples/ltx2/model_inference/LTX-2.3-T2AV-OneStage.py) | [code](/examples/ltx2/model_inference_low_vram/LTX-2.3-T2AV-OneStage.py) | [code](/examples/ltx2/model_training/full/LTX-2.3-T2AV-splited.sh) | [code](/examples/ltx2/model_training/validate_full/LTX-2.3-T2AV.py) | [code](/examples/ltx2/model_training/lora/LTX-2.3-T2AV-splited.sh) | [code](/examples/ltx2/model_training/validate_lora/LTX-2.3-T2AV.py) |
+| LTX-2 | [Lightricks/LTX-2.3: TwoStagePipeline-T2AV](https://www.modelscope.cn/models/Lightricks/LTX-2.3) | [code](/examples/ltx2/model_inference/LTX-2.3-T2AV-TwoStage.py) | [code](/examples/ltx2/model_inference_low_vram/LTX-2.3-T2AV-TwoStage.py) | - | - | - | - |
+| LTX-2 | [Lightricks/LTX-2.3: DistilledPipeline-T2AV](https://www.modelscope.cn/models/Lightricks/LTX-2.3) | [code](/examples/ltx2/model_inference/LTX-2.3-T2AV-DistilledPipeline.py) | [code](/examples/ltx2/model_inference_low_vram/LTX-2.3-T2AV-DistilledPipeline.py) | - | - | - | - |
+| LTX-2 | [Lightricks/LTX-2.3: A2V](https://www.modelscope.cn/models/Lightricks/LTX-2.3) | [code](/examples/ltx2/model_inference/LTX-2.3-A2V-TwoStage.py) | [code](/examples/ltx2/model_inference_low_vram/LTX-2.3-A2V-TwoStage.py) | - | - | - | - |
+| LTX-2 | [Lightricks/LTX-2.3: Retake](https://www.modelscope.cn/models/Lightricks/LTX-2.3) | [code](/examples/ltx2/model_inference/LTX-2.3-T2AV-TwoStage-Retake.py) | [code](/examples/ltx2/model_inference_low_vram/LTX-2.3-T2AV-TwoStage-Retake.py) | - | - | - | - |
+| LTX-2 | [Lightricks/LTX-2.3-22b-IC-LoRA-Union-Control](https://www.modelscope.cn/models/Lightricks/LTX-2.3-22b-IC-LoRA-Union-Control) | [code](/examples/ltx2/model_inference/LTX-2.3-T2AV-IC-LoRA-Union-Control.py) | [code](/examples/ltx2/model_inference_low_vram/LTX-2.3-T2AV-IC-LoRA-Union-Control.py) | - | - | [code](/examples/ltx2/model_training/lora/LTX-2.3-T2AV-IC-LoRA-splited.sh) | [code](/examples/ltx2/model_training/validate_lora/LTX-2.3-T2AV-IC-LoRA.py) |
+| LTX-2 | [Lightricks/LTX-2.3-22b-IC-LoRA-Motion-Track-Control](https://www.modelscope.cn/models/Lightricks/LTX-2.3-22b-IC-LoRA-Motion-Track-Control) | [code](/examples/ltx2/model_inference/LTX-2.3-T2AV-IC-LoRA-Motion-Track-Control.py) | [code](/examples/ltx2/model_inference_low_vram/LTX-2.3-T2AV-IC-LoRA-Motion-Track-Control.py) | - | - | [code](/examples/ltx2/model_training/lora/LTX-2.3-T2AV-IC-LoRA-splited.sh) | [code](/examples/ltx2/model_training/validate_lora/LTX-2.3-T2AV-IC-LoRA.py) |
+| LTX-2 | [Lightricks/LTX-2: OneStagePipeline-T2AV](https://www.modelscope.cn/models/Lightricks/LTX-2) | [code](/examples/ltx2/model_inference/LTX-2-T2AV-OneStage.py) | [code](/examples/ltx2/model_inference_low_vram/LTX-2-T2AV-OneStage.py) | [code](/examples/ltx2/model_training/full/LTX-2-T2AV-splited.sh) | [code](/examples/ltx2/model_training/validate_full/LTX-2-T2AV.py) | [code](/examples/ltx2/model_training/lora/LTX-2-T2AV-splited.sh) | [code](/examples/ltx2/model_training/validate_lora/LTX-2-T2AV.py) |
+| LTX-2 | [Lightricks/LTX-2-19b-IC-LoRA-Union-Control](https://www.modelscope.cn/models/Lightricks/LTX-2-19b-IC-LoRA-Union-Control) | [code](/examples/ltx2/model_inference/LTX-2-T2AV-IC-LoRA-Union-Control.py) | [code](/examples/ltx2/model_inference_low_vram/LTX-2-T2AV-IC-LoRA-Union-Control.py) | - | - | [code](/examples/ltx2/model_training/lora/LTX-2-T2AV-IC-LoRA-splited.sh) | [code](/examples/ltx2/model_training/validate_lora/LTX-2-T2AV-IC-LoRA.py) |
+| LTX-2 | [Lightricks/LTX-2-19b-IC-LoRA-Detailer](https://www.modelscope.cn/models/Lightricks/LTX-2-19b-IC-LoRA-Detailer) | [code](/examples/ltx2/model_inference/LTX-2-T2AV-IC-LoRA-Detailer.py) | [code](/examples/ltx2/model_inference_low_vram/LTX-2-T2AV-IC-LoRA-Detailer.py) | - | - | [code](/examples/ltx2/model_training/lora/LTX-2-T2AV-IC-LoRA-splited.sh) | [code](/examples/ltx2/model_training/validate_lora/LTX-2-T2AV-IC-LoRA.py) |
+| LTX-2 | [Lightricks/LTX-2: TwoStagePipeline-T2AV](https://www.modelscope.cn/models/Lightricks/LTX-2) | [code](/examples/ltx2/model_inference/LTX-2-T2AV-TwoStage.py) | [code](/examples/ltx2/model_inference_low_vram/LTX-2-T2AV-TwoStage.py) | - | - | - | - |
+| LTX-2 | [Lightricks/LTX-2: DistilledPipeline-T2AV](https://www.modelscope.cn/models/Lightricks/LTX-2) | [code](/examples/ltx2/model_inference/LTX-2-T2AV-DistilledPipeline.py) | [code](/examples/ltx2/model_inference_low_vram/LTX-2-T2AV-DistilledPipeline.py) | - | - | - | - |
+| LTX-2 | [Lightricks/LTX-2: OneStagePipeline-I2AV](https://www.modelscope.cn/models/Lightricks/LTX-2) | [code](/examples/ltx2/model_inference/LTX-2-I2AV-OneStage.py) | [code](/examples/ltx2/model_inference_low_vram/LTX-2-I2AV-OneStage.py) | - | - | - | - |
+| LTX-2 | [Lightricks/LTX-2: TwoStagePipeline-I2AV](https://www.modelscope.cn/models/Lightricks/LTX-2) | [code](/examples/ltx2/model_inference/LTX-2-I2AV-TwoStage.py) | [code](/examples/ltx2/model_inference_low_vram/LTX-2-I2AV-TwoStage.py) | - | - | - | - |
+| LTX-2 | [Lightricks/LTX-2: DistilledPipeline-I2AV](https://www.modelscope.cn/models/Lightricks/LTX-2) | [code](/examples/ltx2/model_inference/LTX-2-I2AV-DistilledPipeline.py) | [code](/examples/ltx2/model_inference_low_vram/LTX-2-I2AV-DistilledPipeline.py) | - | - | - | - |
+| LTX-2 | [Lightricks/LTX-2-19b-LoRA-Camera-Control-Dolly-In](https://www.modelscope.cn/models/Lightricks/LTX-2-19b-LoRA-Camera-Control-Dolly-In) | [code](/examples/ltx2/model_inference/LTX-2-T2AV-Camera-Control-Dolly-In.py) | [code](/examples/ltx2/model_inference_low_vram/LTX-2-T2AV-Camera-Control-Dolly-In.py) | - | - | - | - |
+| LTX-2 | [Lightricks/LTX-2-19b-LoRA-Camera-Control-Dolly-Out](https://www.modelscope.cn/models/Lightricks/LTX-2-19b-LoRA-Camera-Control-Dolly-Out) | [code](/examples/ltx2/model_inference/LTX-2-T2AV-Camera-Control-Dolly-Out.py) | [code](/examples/ltx2/model_inference_low_vram/LTX-2-T2AV-Camera-Control-Dolly-Out.py) | - | - | - | - |
+| LTX-2 | [Lightricks/LTX-2-19b-LoRA-Camera-Control-Dolly-Left](https://www.modelscope.cn/models/Lightricks/LTX-2-19b-LoRA-Camera-Control-Dolly-Left) | [code](/examples/ltx2/model_inference/LTX-2-T2AV-Camera-Control-Dolly-Left.py) | [code](/examples/ltx2/model_inference_low_vram/LTX-2-T2AV-Camera-Control-Dolly-Left.py) | - | - | - | - |
+| LTX-2 | [Lightricks/LTX-2-19b-LoRA-Camera-Control-Dolly-Right](https://www.modelscope.cn/models/Lightricks/LTX-2-19b-LoRA-Camera-Control-Dolly-Right) | [code](/examples/ltx2/model_inference/LTX-2-T2AV-Camera-Control-Dolly-Right.py) | [code](/examples/ltx2/model_inference_low_vram/LTX-2-T2AV-Camera-Control-Dolly-Right.py) | - | - | - | - |
+| LTX-2 | [Lightricks/LTX-2-19b-LoRA-Camera-Control-Jib-Up](https://www.modelscope.cn/models/Lightricks/LTX-2-19b-LoRA-Camera-Control-Jib-Up) | [code](/examples/ltx2/model_inference/LTX-2-T2AV-Camera-Control-Jib-Up.py) | [code](/examples/ltx2/model_inference_low_vram/LTX-2-T2AV-Camera-Control-Jib-Up.py) | - | - | - | - |
+| LTX-2 | [Lightricks/LTX-2-19b-LoRA-Camera-Control-Jib-Down](https://www.modelscope.cn/models/Lightricks/LTX-2-19b-LoRA-Camera-Control-Jib-Down) | [code](/examples/ltx2/model_inference/LTX-2-T2AV-Camera-Control-Jib-Down.py) | [code](/examples/ltx2/model_inference_low_vram/LTX-2-T2AV-Camera-Control-Jib-Down.py) | - | - | - | - |
+| LTX-2 | [Lightricks/LTX-2-19b-LoRA-Camera-Control-Static](https://www.modelscope.cn/models/Lightricks/LTX-2-19b-LoRA-Camera-Control-Static) | [code](/examples/ltx2/model_inference/LTX-2-T2AV-Camera-Control-Static.py) | [code](/examples/ltx2/model_inference_low_vram/LTX-2-T2AV-Camera-Control-Static.py) | - | - | - | - |
+| FLUX.2 | [black-forest-labs/FLUX.2-dev](https://www.modelscope.cn/models/black-forest-labs/FLUX.2-dev) | [code](/examples/flux2/model_inference/FLUX.2-dev.py) | [code](/examples/flux2/model_inference_low_vram/FLUX.2-dev.py) | - | - | [code](/examples/flux2/model_training/lora/FLUX.2-dev.sh) | [code](/examples/flux2/model_training/validate_lora/FLUX.2-dev.py) |
+| FLUX.2 | [black-forest-labs/FLUX.2-klein-4B](https://www.modelscope.cn/models/black-forest-labs/FLUX.2-klein-4B) | [code](/examples/flux2/model_inference/FLUX.2-klein-4B.py) | [code](/examples/flux2/model_inference_low_vram/FLUX.2-klein-4B.py) | [code](/examples/flux2/model_training/full/FLUX.2-klein-4B.sh) | [code](/examples/flux2/model_training/validate_full/FLUX.2-klein-4B.py) | [code](/examples/flux2/model_training/lora/FLUX.2-klein-4B.sh) | [code](/examples/flux2/model_training/validate_lora/FLUX.2-klein-4B.py) |
+| FLUX.2 | [black-forest-labs/FLUX.2-klein-9B](https://www.modelscope.cn/models/black-forest-labs/FLUX.2-klein-9B) | [code](/examples/flux2/model_inference/FLUX.2-klein-9B.py) | [code](/examples/flux2/model_inference_low_vram/FLUX.2-klein-9B.py) | [code](/examples/flux2/model_training/full/FLUX.2-klein-9B.sh) | [code](/examples/flux2/model_training/validate_full/FLUX.2-klein-9B.py) | [code](/examples/flux2/model_training/lora/FLUX.2-klein-9B.sh) | [code](/examples/flux2/model_training/validate_lora/FLUX.2-klein-9B.py) |
+| FLUX.2 | [black-forest-labs/FLUX.2-klein-base-4B](https://www.modelscope.cn/models/black-forest-labs/FLUX.2-klein-base-4B) | [code](/examples/flux2/model_inference/FLUX.2-klein-base-4B.py) | [code](/examples/flux2/model_inference_low_vram/FLUX.2-klein-base-4B.py) | [code](/examples/flux2/model_training/full/FLUX.2-klein-base-4B.sh) | [code](/examples/flux2/model_training/validate_full/FLUX.2-klein-base-4B.py) | [code](/examples/flux2/model_training/lora/FLUX.2-klein-base-4B.sh) | [code](/examples/flux2/model_training/validate_lora/FLUX.2-klein-base-4B.py) |
+| FLUX.2 | [black-forest-labs/FLUX.2-klein-base-9B](https://www.modelscope.cn/models/black-forest-labs/FLUX.2-klein-base-9B) | [code](/examples/flux2/model_inference/FLUX.2-klein-base-9B.py) | [code](/examples/flux2/model_inference_low_vram/FLUX.2-klein-base-9B.py) | [code](/examples/flux2/model_training/full/FLUX.2-klein-base-9B.sh) | [code](/examples/flux2/model_training/validate_full/FLUX.2-klein-base-9B.py) | [code](/examples/flux2/model_training/lora/FLUX.2-klein-base-9B.sh) | [code](/examples/flux2/model_training/validate_lora/FLUX.2-klein-base-9B.py) |
+| FLUX.2 | [DiffSynth-Studio/Template-KleinBase4B-Aesthetic](https://www.modelscope.cn/models/DiffSynth-Studio/Template-KleinBase4B-Aesthetic) | [code](/examples/flux2/model_inference/Template-KleinBase4B-Aesthetic.py) | [code](/examples/flux2/model_inference_low_vram/Template-KleinBase4B-Aesthetic.py) | [code](/examples/flux2/model_training/full/Template-KleinBase4B-Aesthetic.sh) | [code](/examples/flux2/model_training/validate_full/Template-KleinBase4B-Aesthetic.py) | - | - |
+| FLUX.2 | [DiffSynth-Studio/Template-KleinBase4B-Brightness](https://www.modelscope.cn/models/DiffSynth-Studio/Template-KleinBase4B-Brightness) | [code](/examples/flux2/model_inference/Template-KleinBase4B-Brightness.py) | [code](/examples/flux2/model_inference_low_vram/Template-KleinBase4B-Brightness.py) | [code](/examples/flux2/model_training/full/Template-KleinBase4B-Brightness.sh) | [code](/examples/flux2/model_training/validate_full/Template-KleinBase4B-Brightness.py) | - | - |
+| FLUX.2 | [DiffSynth-Studio/Template-KleinBase4B-Age](https://www.modelscope.cn/models/DiffSynth-Studio/Template-KleinBase4B-Age) | [code](/examples/flux2/model_inference/Template-KleinBase4B-Age.py) | [code](/examples/flux2/model_inference_low_vram/Template-KleinBase4B-Age.py) | [code](/examples/flux2/model_training/full/Template-KleinBase4B-Age.sh) | [code](/examples/flux2/model_training/validate_full/Template-KleinBase4B-Age.py) | - | - |
+| FLUX.2 | [DiffSynth-Studio/Template-KleinBase4B-ControlNet](https://www.modelscope.cn/models/DiffSynth-Studio/Template-KleinBase4B-ControlNet) | [code](/examples/flux2/model_inference/Template-KleinBase4B-ControlNet.py) | [code](/examples/flux2/model_inference_low_vram/Template-KleinBase4B-ControlNet.py) | [code](/examples/flux2/model_training/full/Template-KleinBase4B-ControlNet.sh) | [code](/examples/flux2/model_training/validate_full/Template-KleinBase4B-ControlNet.py) | - | - |
+| FLUX.2 | [DiffSynth-Studio/Template-KleinBase4B-Edit](https://www.modelscope.cn/models/DiffSynth-Studio/Template-KleinBase4B-Edit) | [code](/examples/flux2/model_inference/Template-KleinBase4B-Edit.py) | [code](/examples/flux2/model_inference_low_vram/Template-KleinBase4B-Edit.py) | [code](/examples/flux2/model_training/full/Template-KleinBase4B-Edit.sh) | [code](/examples/flux2/model_training/validate_full/Template-KleinBase4B-Edit.py) | - | - |
+| FLUX.2 | [DiffSynth-Studio/Template-KleinBase4B-Inpaint](https://www.modelscope.cn/models/DiffSynth-Studio/Template-KleinBase4B-Inpaint) | [code](/examples/flux2/model_inference/Template-KleinBase4B-Inpaint.py) | [code](/examples/flux2/model_inference_low_vram/Template-KleinBase4B-Inpaint.py) | [code](/examples/flux2/model_training/full/Template-KleinBase4B-Inpaint.sh) | [code](/examples/flux2/model_training/validate_full/Template-KleinBase4B-Inpaint.py) | - | - |
+| FLUX.2 | [DiffSynth-Studio/Template-KleinBase4B-PandaMeme](https://www.modelscope.cn/models/DiffSynth-Studio/Template-KleinBase4B-PandaMeme) | [code](/examples/flux2/model_inference/Template-KleinBase4B-PandaMeme.py) | [code](/examples/flux2/model_inference_low_vram/Template-KleinBase4B-PandaMeme.py) | [code](/examples/flux2/model_training/full/Template-KleinBase4B-PandaMeme.sh) | [code](/examples/flux2/model_training/validate_full/Template-KleinBase4B-PandaMeme.py) | - | - |
+| FLUX.2 | [DiffSynth-Studio/Template-KleinBase4B-Sharpness](https://www.modelscope.cn/models/DiffSynth-Studio/Template-KleinBase4B-Sharpness) | [code](/examples/flux2/model_inference/Template-KleinBase4B-Sharpness.py) | [code](/examples/flux2/model_inference_low_vram/Template-KleinBase4B-Sharpness.py) | [code](/examples/flux2/model_training/full/Template-KleinBase4B-Sharpness.sh) | [code](/examples/flux2/model_training/validate_full/Template-KleinBase4B-Sharpness.py) | - | - |
+| FLUX.2 | [DiffSynth-Studio/Template-KleinBase4B-SoftRGB](https://www.modelscope.cn/models/DiffSynth-Studio/Template-KleinBase4B-SoftRGB) | [code](/examples/flux2/model_inference/Template-KleinBase4B-SoftRGB.py) | [code](/examples/flux2/model_inference_low_vram/Template-KleinBase4B-SoftRGB.py) | [code](/examples/flux2/model_training/full/Template-KleinBase4B-SoftRGB.sh) | [code](/examples/flux2/model_training/validate_full/Template-KleinBase4B-SoftRGB.py) | - | - |
+| FLUX.2 | [DiffSynth-Studio/Template-KleinBase4B-Upscaler](https://www.modelscope.cn/models/DiffSynth-Studio/Template-KleinBase4B-Upscaler) | [code](/examples/flux2/model_inference/Template-KleinBase4B-Upscaler.py) | [code](/examples/flux2/model_inference_low_vram/Template-KleinBase4B-Upscaler.py) | [code](/examples/flux2/model_training/full/Template-KleinBase4B-Upscaler.sh) | [code](/examples/flux2/model_training/validate_full/Template-KleinBase4B-Upscaler.py) | - | - |
+| FLUX.2 | [DiffSynth-Studio/Template-KleinBase4B-ContentRef](https://www.modelscope.cn/models/DiffSynth-Studio/Template-KleinBase4B-ContentRef) | [code](/examples/flux2/model_inference/Template-KleinBase4B-ContentRef.py) | [code](/examples/flux2/model_inference_low_vram/Template-KleinBase4B-ContentRef.py) | [code](/examples/flux2/model_training/full/Template-KleinBase4B-ContentRef.sh) | [code](/examples/flux2/model_training/validate_full/Template-KleinBase4B-ContentRef.py) | - | - |
+| FLUX.2 | [DiffSynth-Studio/KleinBase4B-i2L-v2](https://www.modelscope.cn/models/DiffSynth-Studio/KleinBase4B-i2L-v2) | [code](/examples/flux2/model_inference/KleinBase4B-i2L-v2.py) | [code](/examples/flux2/model_inference_low_vram/KleinBase4B-i2L-v2.py) | [code](/examples/flux2/model_training/full/KleinBase4B-i2L-v2.sh) | [code](/examples/flux2/model_training/validate_full/KleinBase4B-i2L-v2.py) | - | - |
+| Z-Image | [Tongyi-MAI/Z-Image](https://www.modelscope.cn/models/Tongyi-MAI/Z-Image) | [code](/examples/z_image/model_inference/Z-Image.py) | [code](/examples/z_image/model_inference_low_vram/Z-Image.py) | [code](/examples/z_image/model_training/full/Z-Image.sh) | [code](/examples/z_image/model_training/validate_full/Z-Image.py) | [code](/examples/z_image/model_training/lora/Z-Image.sh) | [code](/examples/z_image/model_training/validate_lora/Z-Image.py) |
+| Z-Image | [DiffSynth-Studio/Z-Image-i2L](https://www.modelscope.cn/models/DiffSynth-Studio/Z-Image-i2L) | [code](/examples/z_image/model_inference/Z-Image-i2L.py) | [code](/examples/z_image/model_inference_low_vram/Z-Image-i2L.py) | - | - | - | - |
+| Z-Image | [Tongyi-MAI/Z-Image-Turbo](https://www.modelscope.cn/models/Tongyi-MAI/Z-Image-Turbo) | [code](/examples/z_image/model_inference/Z-Image-Turbo.py) | [code](/examples/z_image/model_inference_low_vram/Z-Image-Turbo.py) | [code](/examples/z_image/model_training/full/Z-Image-Turbo.sh) | [code](/examples/z_image/model_training/validate_full/Z-Image-Turbo.py) | [code](/examples/z_image/model_training/lora/Z-Image-Turbo.sh) | [code](/examples/z_image/model_training/validate_lora/Z-Image-Turbo.py) |
+| Z-Image | [PAI/Z-Image-Turbo-Fun-Controlnet-Union-2.1](https://www.modelscope.cn/models/PAI/Z-Image-Turbo-Fun-Controlnet-Union-2.1) | [code](/examples/z_image/model_inference/Z-Image-Turbo-Fun-Controlnet-Union-2.1.py) | [code](/examples/z_image/model_inference_low_vram/Z-Image-Turbo-Fun-Controlnet-Union-2.1.py) | [code](/examples/z_image/model_training/full/Z-Image-Turbo-Fun-Controlnet-Union-2.1.sh) | [code](/examples/z_image/model_training/validate_full/Z-Image-Turbo-Fun-Controlnet-Union-2.1.py) | [code](/examples/z_image/model_training/lora/Z-Image-Turbo-Fun-Controlnet-Union-2.1.sh) | [code](/examples/z_image/model_training/validate_lora/Z-Image-Turbo-Fun-Controlnet-Union-2.1.py) |
+| Z-Image | [PAI/Z-Image-Turbo-Fun-Controlnet-Union-2.1-8steps](https://www.modelscope.cn/models/PAI/Z-Image-Turbo-Fun-Controlnet-Union-2.1) | [code](/examples/z_image/model_inference/Z-Image-Turbo-Fun-Controlnet-Union-2.1-8steps.py) | [code](/examples/z_image/model_inference_low_vram/Z-Image-Turbo-Fun-Controlnet-Union-2.1-8steps.py) | [code](/examples/z_image/model_training/full/Z-Image-Turbo-Fun-Controlnet-Union-2.1-8steps.sh) | [code](/examples/z_image/model_training/validate_full/Z-Image-Turbo-Fun-Controlnet-Union-2.1-8steps.py) | [code](/examples/z_image/model_training/lora/Z-Image-Turbo-Fun-Controlnet-Union-2.1-8steps.sh) | [code](/examples/z_image/model_training/validate_lora/Z-Image-Turbo-Fun-Controlnet-Union-2.1-8steps.py) |
+| Z-Image | [PAI/Z-Image-Turbo-Fun-Controlnet-Tile-2.1-8steps](https://www.modelscope.cn/models/PAI/Z-Image-Turbo-Fun-Controlnet-Union-2.1) | [code](/examples/z_image/model_inference/Z-Image-Turbo-Fun-Controlnet-Tile-2.1-8steps.py) | [code](/examples/z_image/model_inference_low_vram/Z-Image-Turbo-Fun-Controlnet-Tile-2.1-8steps.py) | [code](/examples/z_image/model_training/full/Z-Image-Turbo-Fun-Controlnet-Tile-2.1-8steps.sh) | [code](/examples/z_image/model_training/validate_full/Z-Image-Turbo-Fun-Controlnet-Tile-2.1-8steps.py) | [code](/examples/z_image/model_training/lora/Z-Image-Turbo-Fun-Controlnet-Tile-2.1-8steps.sh) | [code](/examples/z_image/model_training/validate_lora/Z-Image-Turbo-Fun-Controlnet-Tile-2.1-8steps.py) |
+| Z-Image | [DiffSynth-Studio/ZImage-i2L-v2](https://www.modelscope.cn/models/DiffSynth-Studio/ZImage-i2L-v2) | [code](/examples/z_image/model_inference/ZImage-i2L-v2.py) | [code](/examples/z_image/model_inference_low_vram/ZImage-i2L-v2.py) | [code](/examples/z_image/model_training/full/ZImage-i2L-v2.sh) | [code](/examples/z_image/model_training/validate_full/ZImage-i2L-v2.py) | - | - |
+| Anima | [circlestone-labs/Anima](https://www.modelscope.cn/models/circlestone-labs/Anima) | [code](/examples/anima/model_inference/anima-preview.py) | [code](/examples/anima/model_inference_low_vram/anima-preview.py) | [code](/examples/anima/model_training/full/anima-preview.sh) | [code](/examples/anima/model_training/validate_full/anima-preview.py) | [code](/examples/anima/model_training/lora/anima-preview.sh) | [code](/examples/anima/model_training/validate_lora/anima-preview.py) |
+| Qwen-Image | [Qwen/Qwen-Image](https://www.modelscope.cn/models/Qwen/Qwen-Image) | [code](/examples/qwen_image/model_inference/Qwen-Image.py) | [code](/examples/qwen_image/model_inference_low_vram/Qwen-Image.py) | [code](/examples/qwen_image/model_training/full/Qwen-Image.sh) | [code](/examples/qwen_image/model_training/validate_full/Qwen-Image.py) | [code](/examples/qwen_image/model_training/lora/Qwen-Image.sh) | [code](/examples/qwen_image/model_training/validate_lora/Qwen-Image.py) |
+| Qwen-Image | [Qwen/Qwen-Image-2512](https://www.modelscope.cn/models/Qwen/Qwen-Image-2512) | [code](/examples/qwen_image/model_inference/Qwen-Image-2512.py) | [code](/examples/qwen_image/model_inference_low_vram/Qwen-Image-2512.py) | [code](/examples/qwen_image/model_training/full/Qwen-Image-2512.sh) | [code](/examples/qwen_image/model_training/validate_full/Qwen-Image-2512.py) | [code](/examples/qwen_image/model_training/lora/Qwen-Image-2512.sh) | [code](/examples/qwen_image/model_training/validate_lora/Qwen-Image-2512.py) |
+| Qwen-Image | [Qwen/Qwen-Image-Edit](https://www.modelscope.cn/models/Qwen/Qwen-Image-Edit) | [code](/examples/qwen_image/model_inference/Qwen-Image-Edit.py) | [code](/examples/qwen_image/model_inference_low_vram/Qwen-Image-Edit.py) | [code](/examples/qwen_image/model_training/full/Qwen-Image-Edit.sh) | [code](/examples/qwen_image/model_training/validate_full/Qwen-Image-Edit.py) | [code](/examples/qwen_image/model_training/lora/Qwen-Image-Edit.sh) | [code](/examples/qwen_image/model_training/validate_lora/Qwen-Image-Edit.py) |
+| Qwen-Image | [Qwen/Qwen-Image-Edit-2509](https://www.modelscope.cn/models/Qwen/Qwen-Image-Edit-2509) | [code](/examples/qwen_image/model_inference/Qwen-Image-Edit-2509.py) | [code](/examples/qwen_image/model_inference_low_vram/Qwen-Image-Edit-2509.py) | [code](/examples/qwen_image/model_training/full/Qwen-Image-Edit-2509.sh) | [code](/examples/qwen_image/model_training/validate_full/Qwen-Image-Edit-2509.py) | [code](/examples/qwen_image/model_training/lora/Qwen-Image-Edit-2509.sh) | [code](/examples/qwen_image/model_training/validate_lora/Qwen-Image-Edit-2509.py) |
+| Qwen-Image | [Qwen/Qwen-Image-Edit-2511](https://www.modelscope.cn/models/Qwen/Qwen-Image-Edit-2511) | [code](/examples/qwen_image/model_inference/Qwen-Image-Edit-2511.py) | [code](/examples/qwen_image/model_inference_low_vram/Qwen-Image-Edit-2511.py) | [code](/examples/qwen_image/model_training/full/Qwen-Image-Edit-2511.sh) | [code](/examples/qwen_image/model_training/validate_full/Qwen-Image-Edit-2511.py) | [code](/examples/qwen_image/model_training/lora/Qwen-Image-Edit-2511.sh) | [code](/examples/qwen_image/model_training/validate_lora/Qwen-Image-Edit-2511.py) |
+| Qwen-Image | [FireRedTeam/FireRed-Image-Edit-1.0](https://www.modelscope.cn/models/FireRedTeam/FireRed-Image-Edit-1.0) | [code](/examples/qwen_image/model_inference/FireRed-Image-Edit-1.0.py) | [code](/examples/qwen_image/model_inference_low_vram/FireRed-Image-Edit-1.0.py) | [code](/examples/qwen_image/model_training/full/FireRed-Image-Edit-1.0.sh) | [code](/examples/qwen_image/model_training/validate_full/FireRed-Image-Edit-1.0.py) | [code](/examples/qwen_image/model_training/lora/FireRed-Image-Edit-1.0.sh) | [code](/examples/qwen_image/model_training/validate_lora/FireRed-Image-Edit-1.0.py) |
+| Qwen-Image | [FireRedTeam/FireRed-Image-Edit-1.1](https://www.modelscope.cn/models/FireRedTeam/FireRed-Image-Edit-1.1) | [code](/examples/qwen_image/model_inference/FireRed-Image-Edit-1.1.py) | [code](/examples/qwen_image/model_inference_low_vram/FireRed-Image-Edit-1.1.py) | [code](/examples/qwen_image/model_training/full/FireRed-Image-Edit-1.1.sh) | [code](/examples/qwen_image/model_training/validate_full/FireRed-Image-Edit-1.1.py) | [code](/examples/qwen_image/model_training/lora/FireRed-Image-Edit-1.1.sh) | [code](/examples/qwen_image/model_training/validate_lora/FireRed-Image-Edit-1.1.py) |
+| Qwen-Image | [lightx2v/Qwen-Image-Edit-2511-Lightning](https://modelscope.cn/models/lightx2v/Qwen-Image-Edit-2511-Lightning) | [code](/examples/qwen_image/model_inference/Qwen-Image-Edit-2511-Lightning.py) | [code](/examples/qwen_image/model_inference_low_vram/Qwen-Image-Edit-2511-Lightning.py) | - | - | - | - |
+| Qwen-Image | [Qwen/Qwen-Image-Layered](https://www.modelscope.cn/models/Qwen/Qwen-Image-Layered) | [code](/examples/qwen_image/model_inference/Qwen-Image-Layered.py) | [code](/examples/qwen_image/model_inference_low_vram/Qwen-Image-Layered.py) | [code](/examples/qwen_image/model_training/full/Qwen-Image-Layered.sh) | [code](/examples/qwen_image/model_training/validate_full/Qwen-Image-Layered.py) | [code](/examples/qwen_image/model_training/lora/Qwen-Image-Layered.sh) | [code](/examples/qwen_image/model_training/validate_lora/Qwen-Image-Layered.py) |
+| Qwen-Image | [DiffSynth-Studio/Qwen-Image-Layered-Control](https://www.modelscope.cn/models/DiffSynth-Studio/Qwen-Image-Layered-Control) | [code](/examples/qwen_image/model_inference/Qwen-Image-Layered-Control.py) | [code](/examples/qwen_image/model_inference_low_vram/Qwen-Image-Layered-Control.py) | [code](/examples/qwen_image/model_training/full/Qwen-Image-Layered-Control.sh) | [code](/examples/qwen_image/model_training/validate_full/Qwen-Image-Layered-Control.py) | [code](/examples/qwen_image/model_training/lora/Qwen-Image-Layered-Control.sh) | [code](/examples/qwen_image/model_training/validate_lora/Qwen-Image-Layered-Control.py) |
+| Qwen-Image | [DiffSynth-Studio/Qwen-Image-Layered-Control-V2](https://www.modelscope.cn/models/DiffSynth-Studio/Qwen-Image-Layered-Control-V2) | [code](/examples/qwen_image/model_inference/Qwen-Image-Layered-Control-V2.py) | [code](/examples/qwen_image/model_inference_low_vram/Qwen-Image-Layered-Control-V2.py) | - | - | [code](/examples/qwen_image/model_training/lora/Qwen-Image-Layered-Control-V2.sh) | [code](/examples/qwen_image/model_training/validate_lora/Qwen-Image-Layered-Control-V2.py) |
+| Qwen-Image | [DiffSynth-Studio/Qwen-Image-EliGen](https://www.modelscope.cn/models/DiffSynth-Studio/Qwen-Image-EliGen) | [code](/examples/qwen_image/model_inference/Qwen-Image-EliGen.py) | [code](/examples/qwen_image/model_inference_low_vram/Qwen-Image-EliGen.py) | - | - | [code](/examples/qwen_image/model_training/lora/Qwen-Image-EliGen.sh) | [code](/examples/qwen_image/model_training/validate_lora/Qwen-Image-EliGen.py) |
+| Qwen-Image | [DiffSynth-Studio/Qwen-Image-EliGen-V2](https://www.modelscope.cn/models/DiffSynth-Studio/Qwen-Image-EliGen-V2) | [code](/examples/qwen_image/model_inference/Qwen-Image-EliGen-V2.py) | [code](/examples/qwen_image/model_inference_low_vram/Qwen-Image-EliGen-V2.py) | - | - | [code](/examples/qwen_image/model_training/lora/Qwen-Image-EliGen.sh) | [code](/examples/qwen_image/model_training/validate_lora/Qwen-Image-EliGen.py) |
+| Qwen-Image | [DiffSynth-Studio/Qwen-Image-EliGen-Poster](https://www.modelscope.cn/models/DiffSynth-Studio/Qwen-Image-EliGen-Poster) | [code](/examples/qwen_image/model_inference/Qwen-Image-EliGen-Poster.py) | [code](/examples/qwen_image/model_inference_low_vram/Qwen-Image-EliGen-Poster.py) | - | - | [code](/examples/qwen_image/model_training/lora/Qwen-Image-EliGen-Poster.sh) | [code](/examples/qwen_image/model_training/validate_lora/Qwen-Image-EliGen-Poster.py) |
+| Qwen-Image | [DiffSynth-Studio/Qwen-Image-Distill-Full](https://www.modelscope.cn/models/DiffSynth-Studio/Qwen-Image-Distill-Full) | [code](/examples/qwen_image/model_inference/Qwen-Image-Distill-Full.py) | [code](/examples/qwen_image/model_inference_low_vram/Qwen-Image-Distill-Full.py) | [code](/examples/qwen_image/model_training/full/Qwen-Image-Distill-Full.sh) | [code](/examples/qwen_image/model_training/validate_full/Qwen-Image-Distill-Full.py) | [code](/examples/qwen_image/model_training/lora/Qwen-Image-Distill-Full.sh) | [code](/examples/qwen_image/model_training/validate_lora/Qwen-Image-Distill-Full.py) |
+| Qwen-Image | [DiffSynth-Studio/Qwen-Image-Distill-LoRA](https://www.modelscope.cn/models/DiffSynth-Studio/Qwen-Image-Distill-LoRA) | [code](/examples/qwen_image/model_inference/Qwen-Image-Distill-LoRA.py) | [code](/examples/qwen_image/model_inference_low_vram/Qwen-Image-Distill-LoRA.py) | - | - | [code](/examples/qwen_image/model_training/lora/Qwen-Image-Distill-LoRA.sh) | [code](/examples/qwen_image/model_training/validate_lora/Qwen-Image-Distill-LoRA.py) |
+| Qwen-Image | [DiffSynth-Studio/Qwen-Image-Blockwise-ControlNet-Canny](https://modelscope.cn/models/DiffSynth-Studio/Qwen-Image-Blockwise-ControlNet-Canny) | [code](/examples/qwen_image/model_inference/Qwen-Image-Blockwise-ControlNet-Canny.py) | [code](/examples/qwen_image/model_inference_low_vram/Qwen-Image-Blockwise-ControlNet-Canny.py) | [code](/examples/qwen_image/model_training/full/Qwen-Image-Blockwise-ControlNet-Canny.sh) | [code](/examples/qwen_image/model_training/validate_full/Qwen-Image-Blockwise-ControlNet-Canny.py) | [code](/examples/qwen_image/model_training/lora/Qwen-Image-Blockwise-ControlNet-Canny.sh) | [code](/examples/qwen_image/model_training/validate_lora/Qwen-Image-Blockwise-ControlNet-Canny.py) |
+| Qwen-Image | [DiffSynth-Studio/Qwen-Image-Blockwise-ControlNet-Depth](https://modelscope.cn/models/DiffSynth-Studio/Qwen-Image-Blockwise-ControlNet-Depth) | [code](/examples/qwen_image/model_inference/Qwen-Image-Blockwise-ControlNet-Depth.py) | [code](/examples/qwen_image/model_inference_low_vram/Qwen-Image-Blockwise-ControlNet-Depth.py) | [code](/examples/qwen_image/model_training/full/Qwen-Image-Blockwise-ControlNet-Depth.sh) | [code](/examples/qwen_image/model_training/validate_full/Qwen-Image-Blockwise-ControlNet-Depth.py) | [code](/examples/qwen_image/model_training/lora/Qwen-Image-Blockwise-ControlNet-Depth.sh) | [code](/examples/qwen_image/model_training/validate_lora/Qwen-Image-Blockwise-ControlNet-Depth.py) |
+| Qwen-Image | [DiffSynth-Studio/Qwen-Image-Blockwise-ControlNet-Inpaint](https://modelscope.cn/models/DiffSynth-Studio/Qwen-Image-Blockwise-ControlNet-Inpaint) | [code](/examples/qwen_image/model_inference/Qwen-Image-Blockwise-ControlNet-Inpaint.py) | [code](/examples/qwen_image/model_inference_low_vram/Qwen-Image-Blockwise-ControlNet-Inpaint.py) | [code](/examples/qwen_image/model_training/full/Qwen-Image-Blockwise-ControlNet-Inpaint.sh) | [code](/examples/qwen_image/model_training/validate_full/Qwen-Image-Blockwise-ControlNet-Inpaint.py) | [code](/examples/qwen_image/model_training/lora/Qwen-Image-Blockwise-ControlNet-Inpaint.sh) | [code](/examples/qwen_image/model_training/validate_lora/Qwen-Image-Blockwise-ControlNet-Inpaint.py) |
+| Qwen-Image | [DiffSynth-Studio/Qwen-Image-In-Context-Control-Union](https://www.modelscope.cn/models/DiffSynth-Studio/Qwen-Image-In-Context-Control-Union) | [code](/examples/qwen_image/model_inference/Qwen-Image-In-Context-Control-Union.py) | [code](/examples/qwen_image/model_inference_low_vram/Qwen-Image-In-Context-Control-Union.py) | - | - | [code](/examples/qwen_image/model_training/lora/Qwen-Image-In-Context-Control-Union.sh) | [code](/examples/qwen_image/model_training/validate_lora/Qwen-Image-In-Context-Control-Union.py) |
+| Qwen-Image | [DiffSynth-Studio/Qwen-Image-Edit-Lowres-Fix](https://www.modelscope.cn/models/DiffSynth-Studio/Qwen-Image-Edit-Lowres-Fix) | [code](/examples/qwen_image/model_inference/Qwen-Image-Edit-Lowres-Fix.py) | [code](/examples/qwen_image/model_inference_low_vram/Qwen-Image-Edit-Lowres-Fix.py) | - | - | - | - |
+| Qwen-Image | [DiffSynth-Studio/Qwen-Image-i2L](https://www.modelscope.cn/models/DiffSynth-Studio/Qwen-Image-i2L) | [code](/examples/qwen_image/model_inference/Qwen-Image-i2L.py) | [code](/examples/qwen_image/model_inference_low_vram/Qwen-Image-i2L.py) | - | - | - | - |
+| Qwen-Video-Edit | [yunpeng1998/Qwen-Video-Edit](https://www.modelscope.cn/models/yunpeng1998/Qwen-Video-Edit) | [code](/examples/qwen_video_edit/model_inference/Qwen-Video-Edit.py) | [code](/examples/qwen_video_edit/model_inference_low_vram/Qwen-Video-Edit.py) | [code](/examples/qwen_video_edit/model_training/full/Qwen-Video-Edit.sh) | [code](/examples/qwen_video_edit/model_training/validate_full/Qwen-Video-Edit.py) | [code](/examples/qwen_video_edit/model_training/lora/Qwen-Video-Edit.sh) | [code](/examples/qwen_video_edit/model_training/validate_lora/Qwen-Video-Edit.py) |
+| Wan | [Wan-AI/Wan2.1-T2V-1.3B](https://modelscope.cn/models/Wan-AI/Wan2.1-T2V-1.3B) | [code](/examples/wanvideo/model_inference/Wan2.1-T2V-1.3B.py) | [code](/examples/wanvideo/model_inference_low_vram/Wan2.1-T2V-1.3B.py) | [code](/examples/wanvideo/model_training/full/Wan2.1-T2V-1.3B.sh) | [code](/examples/wanvideo/model_training/validate_full/Wan2.1-T2V-1.3B.py) | [code](/examples/wanvideo/model_training/lora/Wan2.1-T2V-1.3B.sh) | [code](/examples/wanvideo/model_training/validate_lora/Wan2.1-T2V-1.3B.py) |
+| Wan | [Wan-AI/Wan2.1-T2V-14B](https://modelscope.cn/models/Wan-AI/Wan2.1-T2V-14B) | [code](/examples/wanvideo/model_inference/Wan2.1-T2V-14B.py) | [code](/examples/wanvideo/model_inference_low_vram/Wan2.1-T2V-14B.py) | [code](/examples/wanvideo/model_training/full/Wan2.1-T2V-14B.sh) | [code](/examples/wanvideo/model_training/validate_full/Wan2.1-T2V-14B.py) | [code](/examples/wanvideo/model_training/lora/Wan2.1-T2V-14B.sh) | [code](/examples/wanvideo/model_training/validate_lora/Wan2.1-T2V-14B.py) |
+| Wan | [Wan-AI/Wan2.1-I2V-14B-480P](https://modelscope.cn/models/Wan-AI/Wan2.1-I2V-14B-480P) | [code](/examples/wanvideo/model_inference/Wan2.1-I2V-14B-480P.py) | [code](/examples/wanvideo/model_inference_low_vram/Wan2.1-I2V-14B-480P.py) | [code](/examples/wanvideo/model_training/full/Wan2.1-I2V-14B-480P.sh) | [code](/examples/wanvideo/model_training/validate_full/Wan2.1-I2V-14B-480P.py) | [code](/examples/wanvideo/model_training/lora/Wan2.1-I2V-14B-480P.sh) | [code](/examples/wanvideo/model_training/validate_lora/Wan2.1-I2V-14B-480P.py) |
+| Wan | [Wan-AI/Wan2.1-I2V-14B-720P](https://modelscope.cn/models/Wan-AI/Wan2.1-I2V-14B-720P) | [code](/examples/wanvideo/model_inference/Wan2.1-I2V-14B-720P.py) | [code](/examples/wanvideo/model_inference_low_vram/Wan2.1-I2V-14B-720P.py) | [code](/examples/wanvideo/model_training/full/Wan2.1-I2V-14B-720P.sh) | [code](/examples/wanvideo/model_training/validate_full/Wan2.1-I2V-14B-720P.py) | [code](/examples/wanvideo/model_training/lora/Wan2.1-I2V-14B-720P.sh) | [code](/examples/wanvideo/model_training/validate_lora/Wan2.1-I2V-14B-720P.py) |
+| Wan | [Wan-AI/Wan2.1-FLF2V-14B-720P](https://modelscope.cn/models/Wan-AI/Wan2.1-FLF2V-14B-720P) | [code](/examples/wanvideo/model_inference/Wan2.1-FLF2V-14B-720P.py) | [code](/examples/wanvideo/model_inference_low_vram/Wan2.1-FLF2V-14B-720P.py) | [code](/examples/wanvideo/model_training/full/Wan2.1-FLF2V-14B-720P.sh) | [code](/examples/wanvideo/model_training/validate_full/Wan2.1-FLF2V-14B-720P.py) | [code](/examples/wanvideo/model_training/lora/Wan2.1-FLF2V-14B-720P.sh) | [code](/examples/wanvideo/model_training/validate_lora/Wan2.1-FLF2V-14B-720P.py) |
+| Wan | [iic/VACE-Wan2.1-1.3B-Preview](https://modelscope.cn/models/iic/VACE-Wan2.1-1.3B-Preview) | [code](/examples/wanvideo/model_inference/Wan2.1-VACE-1.3B-Preview.py) | [code](/examples/wanvideo/model_inference_low_vram/Wan2.1-VACE-1.3B-Preview.py) | [code](/examples/wanvideo/model_training/full/Wan2.1-VACE-1.3B-Preview.sh) | [code](/examples/wanvideo/model_training/validate_full/Wan2.1-VACE-1.3B-Preview.py) | [code](/examples/wanvideo/model_training/lora/Wan2.1-VACE-1.3B-Preview.sh) | [code](/examples/wanvideo/model_training/validate_lora/Wan2.1-VACE-1.3B-Preview.py) |
+| Wan | [Wan-AI/Wan2.1-VACE-1.3B](https://modelscope.cn/models/Wan-AI/Wan2.1-VACE-1.3B) | [code](/examples/wanvideo/model_inference/Wan2.1-VACE-1.3B.py) | [code](/examples/wanvideo/model_inference_low_vram/Wan2.1-VACE-1.3B.py) | [code](/examples/wanvideo/model_training/full/Wan2.1-VACE-1.3B.sh) | [code](/examples/wanvideo/model_training/validate_full/Wan2.1-VACE-1.3B.py) | [code](/examples/wanvideo/model_training/lora/Wan2.1-VACE-1.3B.sh) | [code](/examples/wanvideo/model_training/validate_lora/Wan2.1-VACE-1.3B.py) |
+| Wan | [Wan-AI/Wan2.1-VACE-14B](https://modelscope.cn/models/Wan-AI/Wan2.1-VACE-14B) | [code](/examples/wanvideo/model_inference/Wan2.1-VACE-14B.py) | [code](/examples/wanvideo/model_inference_low_vram/Wan2.1-VACE-14B.py) | [code](/examples/wanvideo/model_training/full/Wan2.1-VACE-14B.sh) | [code](/examples/wanvideo/model_training/validate_full/Wan2.1-VACE-14B.py) | [code](/examples/wanvideo/model_training/lora/Wan2.1-VACE-14B.sh) | [code](/examples/wanvideo/model_training/validate_lora/Wan2.1-VACE-14B.py) |
+| Wan | [PAI/Wan2.1-Fun-1.3B-InP](https://modelscope.cn/models/PAI/Wan2.1-Fun-1.3B-InP) | [code](/examples/wanvideo/model_inference/Wan2.1-Fun-1.3B-InP.py) | [code](/examples/wanvideo/model_inference_low_vram/Wan2.1-Fun-1.3B-InP.py) | [code](/examples/wanvideo/model_training/full/Wan2.1-Fun-1.3B-InP.sh) | [code](/examples/wanvideo/model_training/validate_full/Wan2.1-Fun-1.3B-InP.py) | [code](/examples/wanvideo/model_training/lora/Wan2.1-Fun-1.3B-InP.sh) | [code](/examples/wanvideo/model_training/validate_lora/Wan2.1-Fun-1.3B-InP.py) |
+| Wan | [PAI/Wan2.1-Fun-1.3B-Control](https://modelscope.cn/models/PAI/Wan2.1-Fun-1.3B-Control) | [code](/examples/wanvideo/model_inference/Wan2.1-Fun-1.3B-Control.py) | [code](/examples/wanvideo/model_inference_low_vram/Wan2.1-Fun-1.3B-Control.py) | [code](/examples/wanvideo/model_training/full/Wan2.1-Fun-1.3B-Control.sh) | [code](/examples/wanvideo/model_training/validate_full/Wan2.1-Fun-1.3B-Control.py) | [code](/examples/wanvideo/model_training/lora/Wan2.1-Fun-1.3B-Control.sh) | [code](/examples/wanvideo/model_training/validate_lora/Wan2.1-Fun-1.3B-Control.py) |
+| Wan | [PAI/Wan2.1-Fun-14B-InP](https://modelscope.cn/models/PAI/Wan2.1-Fun-14B-InP) | [code](/examples/wanvideo/model_inference/Wan2.1-Fun-14B-InP.py) | [code](/examples/wanvideo/model_inference_low_vram/Wan2.1-Fun-14B-InP.py) | [code](/examples/wanvideo/model_training/full/Wan2.1-Fun-14B-InP.sh) | [code](/examples/wanvideo/model_training/validate_full/Wan2.1-Fun-14B-InP.py) | [code](/examples/wanvideo/model_training/lora/Wan2.1-Fun-14B-InP.sh) | [code](/examples/wanvideo/model_training/validate_lora/Wan2.1-Fun-14B-InP.py) |
+| Wan | [PAI/Wan2.1-Fun-14B-Control](https://modelscope.cn/models/PAI/Wan2.1-Fun-14B-Control) | [code](/examples/wanvideo/model_inference/Wan2.1-Fun-14B-Control.py) | [code](/examples/wanvideo/model_inference_low_vram/Wan2.1-Fun-14B-Control.py) | [code](/examples/wanvideo/model_training/full/Wan2.1-Fun-14B-Control.sh) | [code](/examples/wanvideo/model_training/validate_full/Wan2.1-Fun-14B-Control.py) | [code](/examples/wanvideo/model_training/lora/Wan2.1-Fun-14B-Control.sh) | [code](/examples/wanvideo/model_training/validate_lora/Wan2.1-Fun-14B-Control.py) |
+| Wan | [PAI/Wan2.1-Fun-V1.1-1.3B-Control](https://modelscope.cn/models/PAI/Wan2.1-Fun-V1.1-1.3B-Control) | [code](/examples/wanvideo/model_inference/Wan2.1-Fun-V1.1-1.3B-Control.py) | [code](/examples/wanvideo/model_inference_low_vram/Wan2.1-Fun-V1.1-1.3B-Control.py) | [code](/examples/wanvideo/model_training/full/Wan2.1-Fun-V1.1-1.3B-Control.sh) | [code](/examples/wanvideo/model_training/validate_full/Wan2.1-Fun-V1.1-1.3B-Control.py) | [code](/examples/wanvideo/model_training/lora/Wan2.1-Fun-V1.1-1.3B-Control.sh) | [code](/examples/wanvideo/model_training/validate_lora/Wan2.1-Fun-V1.1-1.3B-Control.py) |
+| Wan | [PAI/Wan2.1-Fun-V1.1-14B-Control](https://modelscope.cn/models/PAI/Wan2.1-Fun-V1.1-14B-Control) | [code](/examples/wanvideo/model_inference/Wan2.1-Fun-V1.1-14B-Control.py) | [code](/examples/wanvideo/model_inference_low_vram/Wan2.1-Fun-V1.1-14B-Control.py) | [code](/examples/wanvideo/model_training/full/Wan2.1-Fun-V1.1-14B-Control.sh) | [code](/examples/wanvideo/model_training/validate_full/Wan2.1-Fun-V1.1-14B-Control.py) | [code](/examples/wanvideo/model_training/lora/Wan2.1-Fun-V1.1-14B-Control.sh) | [code](/examples/wanvideo/model_training/validate_lora/Wan2.1-Fun-V1.1-14B-Control.py) |
+| Wan | [PAI/Wan2.1-Fun-V1.1-1.3B-InP](https://modelscope.cn/models/PAI/Wan2.1-Fun-V1.1-1.3B-InP) | [code](/examples/wanvideo/model_inference/Wan2.1-Fun-V1.1-1.3B-InP.py) | [code](/examples/wanvideo/model_inference_low_vram/Wan2.1-Fun-V1.1-1.3B-InP.py) | [code](/examples/wanvideo/model_training/full/Wan2.1-Fun-V1.1-1.3B-InP.sh) | [code](/examples/wanvideo/model_training/validate_full/Wan2.1-Fun-V1.1-1.3B-InP.py) | [code](/examples/wanvideo/model_training/lora/Wan2.1-Fun-V1.1-1.3B-InP.sh) | [code](/examples/wanvideo/model_training/validate_lora/Wan2.1-Fun-V1.1-1.3B-InP.py) |
+| Wan | [PAI/Wan2.1-Fun-V1.1-14B-InP](https://modelscope.cn/models/PAI/Wan2.1-Fun-V1.1-14B-InP) | [code](/examples/wanvideo/model_inference/Wan2.1-Fun-V1.1-14B-InP.py) | [code](/examples/wanvideo/model_inference_low_vram/Wan2.1-Fun-V1.1-14B-InP.py) | [code](/examples/wanvideo/model_training/full/Wan2.1-Fun-V1.1-14B-InP.sh) | [code](/examples/wanvideo/model_training/validate_full/Wan2.1-Fun-V1.1-14B-InP.py) | [code](/examples/wanvideo/model_training/lora/Wan2.1-Fun-V1.1-14B-InP.sh) | [code](/examples/wanvideo/model_training/validate_lora/Wan2.1-Fun-V1.1-14B-InP.py) |
+| Wan | [PAI/Wan2.1-Fun-V1.1-1.3B-Control-Camera](https://modelscope.cn/models/PAI/Wan2.1-Fun-V1.1-1.3B-Control-Camera) | [code](/examples/wanvideo/model_inference/Wan2.1-Fun-V1.1-1.3B-Control-Camera.py) | [code](/examples/wanvideo/model_inference_low_vram/Wan2.1-Fun-V1.1-1.3B-Control-Camera.py) | [code](/examples/wanvideo/model_training/full/Wan2.1-Fun-V1.1-1.3B-Control-Camera.sh) | [code](/examples/wanvideo/model_training/validate_full/Wan2.1-Fun-V1.1-1.3B-Control-Camera.py) | [code](/examples/wanvideo/model_training/lora/Wan2.1-Fun-V1.1-1.3B-Control-Camera.sh) | [code](/examples/wanvideo/model_training/validate_lora/Wan2.1-Fun-V1.1-1.3B-Control-Camera.py) |
+| Wan | [PAI/Wan2.1-Fun-V1.1-14B-Control-Camera](https://modelscope.cn/models/PAI/Wan2.1-Fun-V1.1-14B-Control-Camera) | [code](/examples/wanvideo/model_inference/Wan2.1-Fun-V1.1-14B-Control-Camera.py) | [code](/examples/wanvideo/model_inference_low_vram/Wan2.1-Fun-V1.1-14B-Control-Camera.py) | [code](/examples/wanvideo/model_training/full/Wan2.1-Fun-V1.1-14B-Control-Camera.sh) | [code](/examples/wanvideo/model_training/validate_full/Wan2.1-Fun-V1.1-14B-Control-Camera.py) | [code](/examples/wanvideo/model_training/lora/Wan2.1-Fun-V1.1-14B-Control-Camera.sh) | [code](/examples/wanvideo/model_training/validate_lora/Wan2.1-Fun-V1.1-14B-Control-Camera.py) |
+| Wan | [DiffSynth-Studio/Wan2.1-1.3b-speedcontrol-v1](https://modelscope.cn/models/DiffSynth-Studio/Wan2.1-1.3b-speedcontrol-v1) | [code](/examples/wanvideo/model_inference/Wan2.1-1.3b-speedcontrol-v1.py) | [code](/examples/wanvideo/model_inference_low_vram/Wan2.1-1.3b-speedcontrol-v1.py) | [code](/examples/wanvideo/model_training/full/Wan2.1-1.3b-speedcontrol-v1.sh) | [code](/examples/wanvideo/model_training/validate_full/Wan2.1-1.3b-speedcontrol-v1.py) | [code](/examples/wanvideo/model_training/lora/Wan2.1-1.3b-speedcontrol-v1.sh) | [code](/examples/wanvideo/model_training/validate_lora/Wan2.1-1.3b-speedcontrol-v1.py) |
+| Wan | [krea/krea-realtime-video](https://www.modelscope.cn/models/krea/krea-realtime-video) | [code](/examples/wanvideo/model_inference/krea-realtime-video.py) | [code](/examples/wanvideo/model_inference_low_vram/krea-realtime-video.py) | [code](/examples/wanvideo/model_training/full/krea-realtime-video.sh) | [code](/examples/wanvideo/model_training/validate_full/krea-realtime-video.py) | [code](/examples/wanvideo/model_training/lora/krea-realtime-video.sh) | [code](/examples/wanvideo/model_training/validate_lora/krea-realtime-video.py) |
+| Wan | [meituan-longcat/LongCat-Video](https://www.modelscope.cn/models/meituan-longcat/LongCat-Video) | [code](/examples/wanvideo/model_inference/LongCat-Video.py) | [code](/examples/wanvideo/model_inference_low_vram/LongCat-Video.py) | [code](/examples/wanvideo/model_training/full/LongCat-Video.sh) | [code](/examples/wanvideo/model_training/validate_full/LongCat-Video.py) | [code](/examples/wanvideo/model_training/lora/LongCat-Video.sh) | [code](/examples/wanvideo/model_training/validate_lora/LongCat-Video.py) |
+| Wan | [ByteDance/Video-As-Prompt-Wan2.1-14B](https://modelscope.cn/models/ByteDance/Video-As-Prompt-Wan2.1-14B) | [code](/examples/wanvideo/model_inference/Video-As-Prompt-Wan2.1-14B.py) | [code](/examples/wanvideo/model_inference_low_vram/Video-As-Prompt-Wan2.1-14B.py) | [code](/examples/wanvideo/model_training/full/Video-As-Prompt-Wan2.1-14B.sh) | [code](/examples/wanvideo/model_training/validate_full/Video-As-Prompt-Wan2.1-14B.py) | [code](/examples/wanvideo/model_training/lora/Video-As-Prompt-Wan2.1-14B.sh) | [code](/examples/wanvideo/model_training/validate_lora/Video-As-Prompt-Wan2.1-14B.py) |
+| Wan | [Wan-AI/Wan2.2-T2V-A14B](https://modelscope.cn/models/Wan-AI/Wan2.2-T2V-A14B) | [code](/examples/wanvideo/model_inference/Wan2.2-T2V-A14B.py) | [code](/examples/wanvideo/model_inference_low_vram/Wan2.2-T2V-A14B.py) | [code](/examples/wanvideo/model_training/full/Wan2.2-T2V-A14B.sh) | [code](/examples/wanvideo/model_training/validate_full/Wan2.2-T2V-A14B.py) | [code](/examples/wanvideo/model_training/lora/Wan2.2-T2V-A14B.sh) | [code](/examples/wanvideo/model_training/validate_lora/Wan2.2-T2V-A14B.py) |
+| Wan | [Wan-AI/Wan2.2-I2V-A14B](https://modelscope.cn/models/Wan-AI/Wan2.2-I2V-A14B) | [code](/examples/wanvideo/model_inference/Wan2.2-I2V-A14B.py) | [code](/examples/wanvideo/model_inference_low_vram/Wan2.2-I2V-A14B.py) | [code](/examples/wanvideo/model_training/full/Wan2.2-I2V-A14B.sh) | [code](/examples/wanvideo/model_training/validate_full/Wan2.2-I2V-A14B.py) | [code](/examples/wanvideo/model_training/lora/Wan2.2-I2V-A14B.sh) | [code](/examples/wanvideo/model_training/validate_lora/Wan2.2-I2V-A14B.py) |
+| Wan | [Wan-AI/Wan2.2-TI2V-5B](https://modelscope.cn/models/Wan-AI/Wan2.2-TI2V-5B) | [code](/examples/wanvideo/model_inference/Wan2.2-TI2V-5B.py) | [code](/examples/wanvideo/model_inference_low_vram/Wan2.2-TI2V-5B.py) | [code](/examples/wanvideo/model_training/full/Wan2.2-TI2V-5B.sh) | [code](/examples/wanvideo/model_training/validate_full/Wan2.2-TI2V-5B.py) | [code](/examples/wanvideo/model_training/lora/Wan2.2-TI2V-5B.sh) | [code](/examples/wanvideo/model_training/validate_lora/Wan2.2-TI2V-5B.py) |
+| Wan | [Wan-AI/Wan2.2-Animate-14B](https://www.modelscope.cn/models/Wan-AI/Wan2.2-Animate-14B) | [code](/examples/wanvideo/model_inference/Wan2.2-Animate-14B.py) | [code](/examples/wanvideo/model_inference_low_vram/Wan2.2-Animate-14B.py) | [code](/examples/wanvideo/model_training/full/Wan2.2-Animate-14B.sh) | [code](/examples/wanvideo/model_training/validate_full/Wan2.2-Animate-14B.py) | [code](/examples/wanvideo/model_training/lora/Wan2.2-Animate-14B.sh) | [code](/examples/wanvideo/model_training/validate_lora/Wan2.2-Animate-14B.py) |
+| Wan | [Wan-AI/Wan2.2-Animate-2-14B](https://www.modelscope.cn/models/Wan-AI/Wan2.2-Animate-2-14B) | [code](/examples/wanvideo/model_inference/Wan2.2-Animate-2-14B.py) | [code](/examples/wanvideo/model_inference_low_vram/Wan2.2-Animate-2-14B.py) | [code](/examples/wanvideo/model_training/full/Wan2.2-Animate-2-14B.sh) | [code](/examples/wanvideo/model_training/validate_full/Wan2.2-Animate-2-14B.py) | [code](/examples/wanvideo/model_training/lora/Wan2.2-Animate-2-14B.sh) | [code](/examples/wanvideo/model_training/validate_lora/Wan2.2-Animate-2-14B.py) |
+| Wan | [Wan-AI/Wan2.2-Animate-2-14B: Distilled](https://www.modelscope.cn/models/Wan-AI/Wan2.2-Animate-2-14B) | [code](/examples/wanvideo/model_inference/Wan2.2-Animate-2-14B-Distilled.py) | [code](/examples/wanvideo/model_inference_low_vram/Wan2.2-Animate-2-14B-Distilled.py) | [code](/examples/wanvideo/model_training/full/Wan2.2-Animate-2-14B-Distilled.sh) | [code](/examples/wanvideo/model_training/validate_full/Wan2.2-Animate-2-14B-Distilled.py) | [code](/examples/wanvideo/model_training/lora/Wan2.2-Animate-2-14B-Distilled.sh) | [code](/examples/wanvideo/model_training/validate_lora/Wan2.2-Animate-2-14B-Distilled.py) |
+| Wan | [Wan-AI/Wan2.2-S2V-14B](https://www.modelscope.cn/models/Wan-AI/Wan2.2-S2V-14B) | [code](/examples/wanvideo/model_inference/Wan2.2-S2V-14B_multi_clips.py) | [code](/examples/wanvideo/model_inference_low_vram/Wan2.2-S2V-14B_multi_clips.py) | [code](/examples/wanvideo/model_training/full/Wan2.2-S2V-14B.sh) | [code](/examples/wanvideo/model_training/validate_full/Wan2.2-S2V-14B.py) | [code](/examples/wanvideo/model_training/lora/Wan2.2-S2V-14B.sh) | [code](/examples/wanvideo/model_training/validate_lora/Wan2.2-S2V-14B.py) |
+| Wan | [PAI/Wan2.2-VACE-Fun-A14B](https://www.modelscope.cn/models/PAI/Wan2.2-VACE-Fun-A14B) | [code](/examples/wanvideo/model_inference/Wan2.2-VACE-Fun-A14B.py) | [code](/examples/wanvideo/model_inference_low_vram/Wan2.2-VACE-Fun-A14B.py) | [code](/examples/wanvideo/model_training/full/Wan2.2-VACE-Fun-A14B.sh) | [code](/examples/wanvideo/model_training/validate_full/Wan2.2-VACE-Fun-A14B.py) | [code](/examples/wanvideo/model_training/lora/Wan2.2-VACE-Fun-A14B.sh) | [code](/examples/wanvideo/model_training/validate_lora/Wan2.2-VACE-Fun-A14B.py) |
+| Wan | [PAI/Wan2.2-Fun-A14B-InP](https://modelscope.cn/models/PAI/Wan2.2-Fun-A14B-InP) | [code](/examples/wanvideo/model_inference/Wan2.2-Fun-A14B-InP.py) | [code](/examples/wanvideo/model_inference_low_vram/Wan2.2-Fun-A14B-InP.py) | [code](/examples/wanvideo/model_training/full/Wan2.2-Fun-A14B-InP.sh) | [code](/examples/wanvideo/model_training/validate_full/Wan2.2-Fun-A14B-InP.py) | [code](/examples/wanvideo/model_training/lora/Wan2.2-Fun-A14B-InP.sh) | [code](/examples/wanvideo/model_training/validate_lora/Wan2.2-Fun-A14B-InP.py) |
+| Wan | [PAI/Wan2.2-Fun-A14B-Control](https://modelscope.cn/models/PAI/Wan2.2-Fun-A14B-Control) | [code](/examples/wanvideo/model_inference/Wan2.2-Fun-A14B-Control.py) | [code](/examples/wanvideo/model_inference_low_vram/Wan2.2-Fun-A14B-Control.py) | [code](/examples/wanvideo/model_training/full/Wan2.2-Fun-A14B-Control.sh) | [code](/examples/wanvideo/model_training/validate_full/Wan2.2-Fun-A14B-Control.py) | [code](/examples/wanvideo/model_training/lora/Wan2.2-Fun-A14B-Control.sh) | [code](/examples/wanvideo/model_training/validate_lora/Wan2.2-Fun-A14B-Control.py) |
+| Wan | [PAI/Wan2.2-Fun-A14B-Control-Camera](https://modelscope.cn/models/PAI/Wan2.2-Fun-A14B-Control-Camera) | [code](/examples/wanvideo/model_inference/Wan2.2-Fun-A14B-Control-Camera.py) | [code](/examples/wanvideo/model_inference_low_vram/Wan2.2-Fun-A14B-Control-Camera.py) | [code](/examples/wanvideo/model_training/full/Wan2.2-Fun-A14B-Control-Camera.sh) | [code](/examples/wanvideo/model_training/validate_full/Wan2.2-Fun-A14B-Control-Camera.py) | [code](/examples/wanvideo/model_training/lora/Wan2.2-Fun-A14B-Control-Camera.sh) | [code](/examples/wanvideo/model_training/validate_lora/Wan2.2-Fun-A14B-Control-Camera.py) |
+| Wan | [openmoss/MOVA-360p](https://modelscope.cn/models/openmoss/MOVA-360p) | [code](/examples/mova/model_inference/MOVA-360p-I2AV.py) | [code](/examples/mova/model_inference_low_vram/MOVA-360p-I2AV.py) | [code](/examples/mova/model_training/full/MOVA-360P-I2AV.sh) | [code](/examples/mova/model_training/validate_full/MOVA-360p-I2AV.py) | [code](/examples/mova/model_training/lora/MOVA-360P-I2AV.sh) | [code](/examples/mova/model_training/validate_lora/MOVA-360p-I2AV.py) |
+| Wan | [openmoss/MOVA-720p](https://modelscope.cn/models/openmoss/MOVA-720p) | [code](/examples/mova/model_inference/MOVA-720p-I2AV.py) | [code](/examples/mova/model_inference_low_vram/MOVA-720p-I2AV.py) | [code](/examples/mova/model_training/full/MOVA-720P-I2AV.sh) | [code](/examples/mova/model_training/validate_full/MOVA-720p-I2AV.py) | [code](/examples/mova/model_training/lora/MOVA-720P-I2AV.sh) | [code](/examples/mova/model_training/validate_lora/MOVA-720p-I2AV.py) |
+| Wan | [Wan-AI/Wan-Dancer-14B (global model)](https://modelscope.cn/models/Wan-AI/Wan-Dancer-14B) | [code](/examples/wanvideo/model_inference/Wan-Dancer-14B-global.py) | [code](/examples/wanvideo/model_inference_low_vram/Wan-Dancer-14B-global.py) | [code](/examples/wanvideo/model_training/full/Wan-Dancer-14B-global.sh) | [code](/examples/wanvideo/model_training/validate_full/Wan-Dancer-14B-global.py) | [code](/examples/wanvideo/model_training/lora/Wan-Dancer-14B-global.sh) | [code](/examples/wanvideo/model_training/validate_lora/Wan-Dancer-14B-global.py) |
+| Wan | [Wan-AI/Wan-Dancer-14B (local model)](https://modelscope.cn/models/Wan-AI/Wan-Dancer-14B) | [code](/examples/wanvideo/model_inference/Wan-Dancer-14B-local.py) | [code](/examples/wanvideo/model_inference_low_vram/Wan-Dancer-14B-local.py) | [code](/examples/wanvideo/model_training/full/Wan-Dancer-14B-local.sh) | [code](/examples/wanvideo/model_training/validate_full/Wan-Dancer-14B-local.py) | [code](/examples/wanvideo/model_training/lora/Wan-Dancer-14B-local.sh) | [code](/examples/wanvideo/model_training/validate_lora/Wan-Dancer-14B-local.py) |
+| FLUX.1 | [black-forest-labs/FLUX.1-dev](https://www.modelscope.cn/models/black-forest-labs/FLUX.1-dev) | [code](/examples/flux/model_inference/FLUX.1-dev.py) | [code](/examples/flux/model_inference_low_vram/FLUX.1-dev.py) | [code](/examples/flux/model_training/full/FLUX.1-dev.sh) | [code](/examples/flux/model_training/validate_full/FLUX.1-dev.py) | [code](/examples/flux/model_training/lora/FLUX.1-dev.sh) | [code](/examples/flux/model_training/validate_lora/FLUX.1-dev.py) |
+| FLUX.1 | [black-forest-labs/FLUX.1-Krea-dev](https://www.modelscope.cn/models/black-forest-labs/FLUX.1-Krea-dev) | [code](/examples/flux/model_inference/FLUX.1-Krea-dev.py) | [code](/examples/flux/model_inference_low_vram/FLUX.1-Krea-dev.py) | [code](/examples/flux/model_training/full/FLUX.1-Krea-dev.sh) | [code](/examples/flux/model_training/validate_full/FLUX.1-Krea-dev.py) | [code](/examples/flux/model_training/lora/FLUX.1-Krea-dev.sh) | [code](/examples/flux/model_training/validate_lora/FLUX.1-Krea-dev.py) |
+| FLUX.1 | [black-forest-labs/FLUX.1-Kontext-dev](https://www.modelscope.cn/models/black-forest-labs/FLUX.1-Kontext-dev) | [code](/examples/flux/model_inference/FLUX.1-Kontext-dev.py) | [code](/examples/flux/model_inference_low_vram/FLUX.1-Kontext-dev.py) | [code](/examples/flux/model_training/full/FLUX.1-Kontext-dev.sh) | [code](/examples/flux/model_training/validate_full/FLUX.1-Kontext-dev.py) | [code](/examples/flux/model_training/lora/FLUX.1-Kontext-dev.sh) | [code](/examples/flux/model_training/validate_lora/FLUX.1-Kontext-dev.py) |
+| FLUX.1 | [black-forest-labs/FLUX.1-Fill-dev](https://www.modelscope.cn/models/black-forest-labs/FLUX.1-Fill-dev) | [code](/examples/flux/model_inference/FLUX.1-Fill-dev.py) | [code](/examples/flux/model_inference_low_vram/FLUX.1-Fill-dev.py) | [code](/examples/flux/model_training/full/FLUX.1-Fill-dev.sh) | [code](/examples/flux/model_training/validate_full/FLUX.1-Fill-dev.py) | [code](/examples/flux/model_training/lora/FLUX.1-Fill-dev.sh) | [code](/examples/flux/model_training/validate_lora/FLUX.1-Fill-dev.py) |
+| FLUX.1 | [black-forest-labs/FLUX.1-Redux-dev](https://www.modelscope.cn/models/black-forest-labs/FLUX.1-Redux-dev) | [code](/examples/flux/model_inference/FLUX.1-Redux-dev.py) | [code](/examples/flux/model_inference_low_vram/FLUX.1-Redux-dev.py) | [code](/examples/flux/model_training/full/FLUX.1-Redux-dev.sh) | [code](/examples/flux/model_training/validate_full/FLUX.1-Redux-dev.py) | [code](/examples/flux/model_training/lora/FLUX.1-Redux-dev.sh) | [code](/examples/flux/model_training/validate_lora/FLUX.1-Redux-dev.py) |
+| FLUX.1 | [HuanJue/Insert-Anything](https://www.modelscope.cn/models/HuanJue/Insert-Anything) | [code](/examples/flux/model_inference/Insert-Anything.py) | [code](/examples/flux/model_inference_low_vram/Insert-Anything.py) | - | - | [code](/examples/flux/model_training/lora/Insert-Anything.sh) | [code](/examples/flux/model_training/validate_lora/Insert-Anything.py) |
+| FLUX.1 | [alimama-creative/FLUX.1-dev-Controlnet-Inpainting-Beta](https://www.modelscope.cn/models/alimama-creative/FLUX.1-dev-Controlnet-Inpainting-Beta) | [code](/examples/flux/model_inference/FLUX.1-dev-Controlnet-Inpainting-Beta.py) | [code](/examples/flux/model_inference_low_vram/FLUX.1-dev-Controlnet-Inpainting-Beta.py) | [code](/examples/flux/model_training/full/FLUX.1-dev-Controlnet-Inpainting-Beta.sh) | [code](/examples/flux/model_training/validate_full/FLUX.1-dev-Controlnet-Inpainting-Beta.py) | [code](/examples/flux/model_training/lora/FLUX.1-dev-Controlnet-Inpainting-Beta.sh) | [code](/examples/flux/model_training/validate_lora/FLUX.1-dev-Controlnet-Inpainting-Beta.py) |
+| FLUX.1 | [InstantX/FLUX.1-dev-Controlnet-Union-alpha](https://www.modelscope.cn/models/InstantX/FLUX.1-dev-Controlnet-Union-alpha) | [code](/examples/flux/model_inference/FLUX.1-dev-Controlnet-Union-alpha.py) | [code](/examples/flux/model_inference_low_vram/FLUX.1-dev-Controlnet-Union-alpha.py) | [code](/examples/flux/model_training/full/FLUX.1-dev-Controlnet-Union-alpha.sh) | [code](/examples/flux/model_training/validate_full/FLUX.1-dev-Controlnet-Union-alpha.py) | [code](/examples/flux/model_training/lora/FLUX.1-dev-Controlnet-Union-alpha.sh) | [code](/examples/flux/model_training/validate_lora/FLUX.1-dev-Controlnet-Union-alpha.py) |
+| FLUX.1 | [jasperai/Flux.1-dev-Controlnet-Upscaler](https://www.modelscope.cn/models/jasperai/Flux.1-dev-Controlnet-Upscaler) | [code](/examples/flux/model_inference/FLUX.1-dev-Controlnet-Upscaler.py) | [code](/examples/flux/model_inference_low_vram/FLUX.1-dev-Controlnet-Upscaler.py) | [code](/examples/flux/model_training/full/FLUX.1-dev-Controlnet-Upscaler.sh) | [code](/examples/flux/model_training/validate_full/FLUX.1-dev-Controlnet-Upscaler.py) | [code](/examples/flux/model_training/lora/FLUX.1-dev-Controlnet-Upscaler.sh) | [code](/examples/flux/model_training/validate_lora/FLUX.1-dev-Controlnet-Upscaler.py) |
+| FLUX.1 | [InstantX/FLUX.1-dev-IP-Adapter](https://www.modelscope.cn/models/InstantX/FLUX.1-dev-IP-Adapter) | [code](/examples/flux/model_inference/FLUX.1-dev-IP-Adapter.py) | [code](/examples/flux/model_inference_low_vram/FLUX.1-dev-IP-Adapter.py) | [code](/examples/flux/model_training/full/FLUX.1-dev-IP-Adapter.sh) | [code](/examples/flux/model_training/validate_full/FLUX.1-dev-IP-Adapter.py) | [code](/examples/flux/model_training/lora/FLUX.1-dev-IP-Adapter.sh) | [code](/examples/flux/model_training/validate_lora/FLUX.1-dev-IP-Adapter.py) |
+| FLUX.1 | [ByteDance/InfiniteYou](https://www.modelscope.cn/models/ByteDance/InfiniteYou) | [code](/examples/flux/model_inference/FLUX.1-dev-InfiniteYou.py) | [code](/examples/flux/model_inference_low_vram/FLUX.1-dev-InfiniteYou.py) | [code](/examples/flux/model_training/full/FLUX.1-dev-InfiniteYou.sh) | [code](/examples/flux/model_training/validate_full/FLUX.1-dev-InfiniteYou.py) | [code](/examples/flux/model_training/lora/FLUX.1-dev-InfiniteYou.sh) | [code](/examples/flux/model_training/validate_lora/FLUX.1-dev-InfiniteYou.py) |
+| FLUX.1 | [DiffSynth-Studio/Eligen](https://www.modelscope.cn/models/DiffSynth-Studio/Eligen) | [code](/examples/flux/model_inference/FLUX.1-dev-EliGen.py) | [code](/examples/flux/model_inference_low_vram/FLUX.1-dev-EliGen.py) | - | - | [code](/examples/flux/model_training/lora/FLUX.1-dev-EliGen.sh) | [code](/examples/flux/model_training/validate_lora/FLUX.1-dev-EliGen.py) |
+| FLUX.1 | [DiffSynth-Studio/LoRA-Encoder-FLUX.1-Dev](https://www.modelscope.cn/models/DiffSynth-Studio/LoRA-Encoder-FLUX.1-Dev) | [code](/examples/flux/model_inference/FLUX.1-dev-LoRA-Encoder.py) | [code](/examples/flux/model_inference_low_vram/FLUX.1-dev-LoRA-Encoder.py) | [code](/examples/flux/model_training/full/FLUX.1-dev-LoRA-Encoder.sh) | [code](/examples/flux/model_training/validate_full/FLUX.1-dev-LoRA-Encoder.py) | - | - |
+| FLUX.1 | [DiffSynth-Studio/LoRAFusion-preview-FLUX.1-dev](https://modelscope.cn/models/DiffSynth-Studio/LoRAFusion-preview-FLUX.1-dev) | [code](/examples/flux/model_inference/FLUX.1-dev-LoRA-Fusion.py) | - | - | - | - | - |
+| FLUX.1 | [stepfun-ai/Step1X-Edit](https://www.modelscope.cn/models/stepfun-ai/Step1X-Edit) | [code](/examples/flux/model_inference/Step1X-Edit.py) | [code](/examples/flux/model_inference_low_vram/Step1X-Edit.py) | [code](/examples/flux/model_training/full/Step1X-Edit.sh) | [code](/examples/flux/model_training/validate_full/Step1X-Edit.py) | [code](/examples/flux/model_training/lora/Step1X-Edit.sh) | [code](/examples/flux/model_training/validate_lora/Step1X-Edit.py) |
+| FLUX.1 | [ostris/Flex.2-preview](https://www.modelscope.cn/models/ostris/Flex.2-preview) | [code](/examples/flux/model_inference/FLEX.2-preview.py) | [code](/examples/flux/model_inference_low_vram/FLEX.2-preview.py) | [code](/examples/flux/model_training/full/FLEX.2-preview.sh) | [code](/examples/flux/model_training/validate_full/FLEX.2-preview.py) | [code](/examples/flux/model_training/lora/FLEX.2-preview.sh) | [code](/examples/flux/model_training/validate_lora/FLEX.2-preview.py) |
+| FLUX.1 | [DiffSynth-Studio/Nexus-GenV2](https://www.modelscope.cn/models/DiffSynth-Studio/Nexus-GenV2) | [code](/examples/flux/model_inference/Nexus-Gen-Editing.py) | [code](/examples/flux/model_inference_low_vram/Nexus-Gen-Editing.py) | [code](/examples/flux/model_training/full/Nexus-Gen.sh) | [code](/examples/flux/model_training/validate_full/Nexus-Gen.py) | [code](/examples/flux/model_training/lora/Nexus-Gen.sh) | [code](/examples/flux/model_training/validate_lora/Nexus-Gen.py) |
+| Stable Diffusion XL | [stabilityai/stable-diffusion-xl-base-1.0](https://www.modelscope.cn/models/stabilityai/stable-diffusion-xl-base-1.0) | [code](/examples/stable_diffusion_xl/model_inference/stable-diffusion-xl-base-1.0.py) | [code](/examples/stable_diffusion_xl/model_inference_low_vram/stable-diffusion-xl-base-1.0.py) | [code](/examples/stable_diffusion_xl/model_training/full/stable-diffusion-xl-base-1.0.sh) | [code](/examples/stable_diffusion_xl/model_training/validate_full/stable-diffusion-xl-base-1.0.py) | [code](/examples/stable_diffusion_xl/model_training/lora/stable-diffusion-xl-base-1.0.sh) | [code](/examples/stable_diffusion_xl/model_training/validate_lora/stable-diffusion-xl-base-1.0.py) |
+| Stable Diffusion | [AI-ModelScope/stable-diffusion-v1-5](https://www.modelscope.cn/models/AI-ModelScope/stable-diffusion-v1-5) | [code](/examples/stable_diffusion/model_inference/stable-diffusion-v1-5.py) | [code](/examples/stable_diffusion/model_inference_low_vram/stable-diffusion-v1-5.py) | [code](/examples/stable_diffusion/model_training/full/stable-diffusion-v1-5.sh) | [code](/examples/stable_diffusion/model_training/validate_full/stable-diffusion-v1-5.py) | [code](/examples/stable_diffusion/model_training/lora/stable-diffusion-v1-5.sh) | [code](/examples/stable_diffusion/model_training/validate_lora/stable-diffusion-v1-5.py) |
diff --git a/README_zh.md b/README_zh.md
new file mode 100644
index 0000000000000000000000000000000000000000..5b7c2e6bcfacff596d0d491d746a1da3f792dfa9
--- /dev/null
+++ b/README_zh.md
@@ -0,0 +1,729 @@
+# DiffSynth-Studio
+
+
+
+[](https://pypi.org/project/DiffSynth/)
+[](https://github.com/modelscope/DiffSynth-Studio/blob/master/LICENSE)
+[](https://github.com/modelscope/DiffSynth-Studio/issues)
+[](https://GitHub.com/modelscope/DiffSynth-Studio/pull/)
+[](https://GitHub.com/modelscope/DiffSynth-Studio/commit/)
+[](https://discord.gg/Mm9suEeUDc)
+
+[Switch to English](./README.md)
+
+## 简介
+
+欢迎来到 Diffusion 模型的魔法世界！DiffSynth-Studio 是由[魔搭社区](https://www.modelscope.cn/)团队开发和维护的开源 Diffusion 模型引擎。我们期望以框架建设孵化技术创新，凝聚开源社区的力量，探索生成式模型技术的丰富能力！
+
+框架功能：
+
+* [模型支持](#全部支持的模型)：框架集成了主流的开源 Diffusion 模型，涵盖图像生成、视频生成、音频生成，以及图像指标模型。
+* [显存管理](https://diffsynth-studio-doc.readthedocs.io/zh-cn/latest/Pipeline_Usage/VRAM_management.html)：在硬盘、内存、显存之间动态调度模型参数，让低显存的消费级 GPU 也能运行大模型的推理。
+* [参数量化](https://diffsynth-studio-doc.readthedocs.io/zh-cn/latest/Pipeline_Usage/Quantization.html)：将模型参数转换为 NF4、INT8 等量化精度，大幅减少模型推理和 LoRA 模型训练的显存需求。
+* [任意训练](https://diffsynth-studio-doc.readthedocs.io/zh-cn/latest/Pipeline_Usage/Model_Training.html)：几乎所有支持推理的模型都支持训练，无论是基础模型、LoRA，还是任何带有额外输入的 Adapter 模型。
+* [拆分训练](https://diffsynth-studio-doc.readthedocs.io/zh-cn/latest/Training/Split_Training.html)：使用计算图推理引擎追踪 Pipeline 中每个变量，将训练过程拆分为两阶段，高效训练。
+
+参考资料：
+
+* 开发者文档 (面向人类)：[中文版](https://diffsynth-studio-doc.readthedocs.io/zh-cn/latest/README.html)、[English version](https://diffsynth-studio-doc.readthedocs.io/en/latest/README.html)
+* Agent Skills (面向 AI)：[DiffSynth-Studio Model Integration Skills](https://www.modelscope.cn/collections/DiffSynth-Studio/DiffSynth-Studio-Model-Integration-Skills)
+
+查看更多：
+
+* [DiffSynth-WebUI](https://github.com/modelscope/DiffSynth-WebUI): 基于 DiffSynth-Studio 构建的轻量化 LoRA 训练工具，可在消费级 GPU 上训练 LoRA 模型。
+* [魔搭社区 AIGC 专区 (面向中国用户)](https://modelscope.cn/aigc/home): 以 DiffSynth-Studio 为核心推理和训练引擎的产品化功能，体验开源模型生态的丰富潜力。
+* [ModelScope Civision (for global users)](https://modelscope.ai/civision/home): Unlock the vast potential of the open-source model ecosystem through productized capabilities powered by DiffSynth-Studio.
+
+## 更新历史
+
+> DiffSynth-Studio 经历了大版本更新，部分旧功能已停止维护，如需使用旧版功能，请切换到大版本更新前的[最后一个历史版本](https://github.com/modelscope/DiffSynth-Studio/tree/afd101f3452c9ecae0c87b79adfa2e22d65ffdc3)。
+
+> 目前本项目的开发人员有限，大部分工作由 [Artiprocher](https://github.com/Artiprocher) 和 [mi804](https://github.com/mi804) 负责，因此新功能的开发进展会比较缓慢，issue 的回复和解决速度有限，我们对此感到非常抱歉，请各位开发者理解。
+
+- **2026年8月31日** 我们接入了 [Qwen-Video-Edit](https://modelscope.cn/models/yunpeng1998/Qwen-Video-Edit)，这是开源社区用户 [yunpeng1998](https://github.com/yunpeng1998) 基于图像编辑模型 Qwen-Image-Edit 训练的视频编辑模型，是探索模型能力拓展的优秀案例。
+
+- **2026年8月25日** 我们开源了 [DiffSynth-WebUI](https://github.com/modelscope/DiffSynth-WebUI)，一键私有化部署 LoRA 训练服务，与模型参数量化功能结合，用消费级 GPU 也能训练超大模型。
+
+- **2026年8月19日** 我们发布了模型量化功能。它提供了统一的 `QuantizeConfig` 入口，支持 bitsandbytes、torchao、comfy-kitchen 等多种量化后端，具备在线量化、加载预量化权重、混合量化、保存量化模型、量化 + LoRA 训练等能力。详情请参考[文档](/docs/zh/Pipeline_Usage/Quantization.md)。
+
+- **2026年8月17日** MiniMax-Music3 开源，欢迎加入音频生成模型家族！支持文生音乐推理和低显存推理能力。详情请参考[文档](/docs/zh/Model_Details/MiniMax-Music3.md)和[示例代码](/examples/minimax_music3/)。
+
+- **2026年8月7日** 我们为 Wan 系列新增了 Wan-Animate-2，输入一张参考图和一段驱动视频，即可让参考角色演绎驱动视频中的动作，生成高质量角色动画，包含标准与蒸馏两个变体。详情请参考[文档](/docs/zh/Model_Details/Wan.md)和[示例代码](/examples/wanvideo/)。
+
+- **2026年8月3日** MiniMax-H3 开源，欢迎加入视频生成模型家族！支持文生音视频、首尾帧引导生成、参考驱动生成、低显存推理以及 NF4 量化推理。详情请参考[文档](/docs/zh/Model_Details/MiniMax-H3.md)和[示例代码](/examples/minimax_h3/)。
+
+- **2026年7月28日** LingBot-Video 开源，欢迎加入视频生成模型家族！本次开源包含 Dense-1.3B 和 MoE-30B-A3B 两个版本（MoE 版总参数量 30B、每个 token 激活约 3B），均支持文生视频、图生视频、文生图推理、低显存推理以及 LoRA / 全量训练能力。详情请参考[文档](/docs/zh/Model_Details/LingBot-Video.md)和[示例代码](/examples/lingbot_video/)。特别感谢 [NancyFyong](https://github.com/NancyFyong) 贡献了本模型的接入代码！
+
+- **2026年7月21日** 我们开源了 [DiffSynth-Studio Model Integration Skills](https://www.modelscope.cn/collections/DiffSynth-Studio/DiffSynth-Studio-Model-Integration-Skills)。这是一套可组合的 Agent Skill 合集，将外部扩散模型接入 DiffSynth-Studio 的全流程自动化，大幅提升模型接入标准化程度与效率。从[使用示例](https://www.modelscope.cn/skills/DiffSynth-Studio/diffsynth-integrator/file/view/master/example.md?status=1)开始体验吧！
+
+
+更多
+
+- **2026年6月29日** Boogu-Image 开源，已支持文生图推理、图像编辑、低显存推理和训练能力。详情请参考[文档](/docs/zh/Model_Details/Boogu-Image.md)和[示例代码](/examples/boogu_image/)。
+
+- **2026年6月24日** Krea-2 开源，我们已提供全面支持。详情请参考[文档](/docs/zh/Model_Details/Krea-2.md)和[示例代码](/examples/krea2/)。
+
+- **2026年6月16日** 我们为 ACE-Step 新增了 Template 模型：[vocals2music](https://www.modelscope.cn/models/DiffSynth-Studio/acestep15xlsft-vocals2music)。详情请参考[文档](/docs/zh/Model_Details/ACE-Step.md)和[示例代码](/examples/ace_step/)。
+
+- **2026年6月15日** 我们开源了 Image-to-LoRA V2，将动辄数小时的图像风格 LoRA 训练压缩到一次模型推理中，探索 LoRA 模型训练的新方式。[技术报告](https://arxiv.org/abs/2606.13809)已公开，本次开源包括三个模型：
+ * [DiffSynth-Studio/ZImage-i2L-v2](https://modelscope.cn/models/DiffSynth-Studio/ZImage-i2L-v2)：适配模型 Z-Image
+ * [DiffSynth-Studio/KleinBase4B-i2L-v2](https://modelscope.cn/models/DiffSynth-Studio/KleinBase4B-i2L-v2)：适配模型 FLUX.2-klein-base-4B
+ * [DiffSynth-Studio/HidreamO1-i2L-v2](https://modelscope.cn/models/DiffSynth-Studio/HidreamO1-i2L-v2)：适配模型 Hidream-O1-Image
+
+- **2026年6月5日** Ideogram 4 开源，已支持文生图推理。详情请参考[文档](/docs/zh/Model_Details/Ideogram-4.md)和[示例代码](/examples/ideogram4/)。
+
+- **2026年5月21日** 新增图像质量评估模型的支持，包括 FID、CLIP、Aesthetic、PickScore、ImageReward、HPSv2、HPSv3，详情请参考[文档](/docs/zh/Model_Details/Image-Quality-Metrics.md)和[示例代码](/examples/image_quality_metric/)
+
+- **2026年5月18日** 新增 **CPU Offload Training** 功能，通过将模型权重逐层在 CPU 与 GPU 之间搬运，大幅降低训练时的 GPU 显存占用，让消费级显卡也能进行大模型 LoRA 训练，适配所有模型。只需在训练命令中添加 `--enable_model_cpu_offload` 即可启用（当前仅支持单卡训练）。详情请参考[文档](/docs/zh/Training/Offload_Training.md)。
+
+- **2026年5月14日** HiDream-O1-Image 开源，欢迎加入图像生成模型家族！支持文生图推理、图像编辑推理、低显存推理和训练能力。详情请参考[文档](/docs/zh/Model_Details/HiDream-O1-Image.md)和[示例代码](/examples/hidream_o1_image/)。
+
+- **2026年4月28日** 我们发布了 Diffusion Templates，面向 Diffusion 模型的插件框架，大幅降低了可控生成模型的训练门槛，一起来探索新奇的技术吧！
+ * 开源代码：[DiffSynth-Studio](https://github.com/modelscope/DiffSynth-Studio)
+ * 技术报告：[arXiv](https://arxiv.org/abs/2604.24351)
+ * 项目主页：[GitHub](https://modelscope.github.io/diffusion-templates-web/)
+ * 文档参考：[English Version](https://diffsynth-studio-doc.readthedocs.io/en/latest/Diffusion_Templates/Introducing_Diffusion_Templates.html)、[中文版](https://diffsynth-studio-doc.readthedocs.io/zh-cn/latest/Diffusion_Templates/Introducing_Diffusion_Templates.html)
+ * 在线体验：[魔搭社区创空间](https://modelscope.cn/studios/DiffSynth-Studio/Diffusion-Templates)
+ * 模型集：[ModelScope](https://modelscope.cn/collections/DiffSynth-Studio/KleinBase4B-Templates)、[ModelScope 国际站](https://modelscope.ai/collections/DiffSynth-Studio/KleinBase4B-Templates)、[HuggingFace](https://huggingface.co/collections/DiffSynth-Studio/kleinbase4b-templates)
+ * 数据集：[ModelScope](https://modelscope.cn/collections/DiffSynth-Studio/ImagePulseV2)、[ModelScope 国际站](https://modelscope.cn/collections/DiffSynth-Studio/ImagePulseV2)、[HuggingFace](https://huggingface.co/collections/DiffSynth-Studio/imagepulsev2)
+- **2026年4月27日** 我们支持了 ACE-Step-1.5！包括文生音乐推理、低显存推理和 LoRA 训练能力。详情请参考[文档](/docs/zh/Model_Details/ACE-Step.md)和[示例代码](/examples/ace_step/)。
+
+- **2026年4月27日** 我们重新支持了模型 Stable Diffusion v1.5 和 SDXL，仅对这两类模型提供学术科研支持。
+
+- **2026年4月14日** JoyAI-Image 开源，欢迎加入图像编辑模型家族！支持指令引导的图像编辑推理、低显存推理和训练能力。详情请参考[文档](/docs/zh/Model_Details/JoyAI-Image.md)和[示例代码](/examples/joyai_image/)。
+
+- **2026年3月19日** 新增对 [openmoss/MOVA-720p](https://modelscope.cn/models/openmoss/MOVA-720p) 和 [openmoss/MOVA-360p](https://modelscope.cn/models/openmoss/MOVA-360p) 模型的支持，包括完整的训练和推理功能。[文档](/docs/zh/Model_Details/Wan.md)和[示例代码](/examples/mova/)现已可用。
+
+- **2026年3月12日** 我们新增了 [LTX-2.3](https://modelscope.cn/models/Lightricks/LTX-2.3) 音视频生成模型的支持，模型支持的功能包括文生音视频、图生音视频、IC-LoRA控制、音频生视频、音视频局部Inpainting，框架支持完整的推理和训练功能。详细信息请参考 [文档](/docs/zh/Model_Details/LTX-2.md) 和 [示例代码](/examples/ltx2/)。
+
+- **2026年3月3日** 我们发布了 [DiffSynth-Studio/Qwen-Image-Layered-Control-V2](https://www.modelscope.cn/models/DiffSynth-Studio/Qwen-Image-Layered-Control-V2) 模型，这是 Qwen-Image-Layered-Control 的更新版本。除了原本就支持的文本引导功能，新增了画笔控制的图层拆分能力。
+
+- **2026年3月2日** 新增对[Anima](https://modelscope.cn/models/circlestone-labs/Anima)的支持，详见[文档](docs/zh/Model_Details/Anima.md)。这是一个有趣的动漫风格图像生成模型，我们期待其后续的模型更新。
+
+- **2026年2月26日** 新增对[LTX-2](https://www.modelscope.cn/models/Lightricks/LTX-2)音视频生成模型全量微调与LoRA训练支持，详见[文档](docs/zh/Model_Details/LTX-2.md)。
+
+- **2026年2月10日** 新增对[LTX-2](https://www.modelscope.cn/models/Lightricks/LTX-2)音视频生成模型的推理支持，详见[文档](docs/zh/Model_Details/LTX-2.md)，后续将推进模型训练的支持。
+
+- **2026年2月2日** Research Tutorial 的第一篇文档上线，带你从零开始训练一个 0.1B 的小型文生图模型，详见[文档](/docs/zh/Research_Tutorial/train_from_scratch.md)、[模型](https://modelscope.cn/models/DiffSynth-Studio/AAAMyModel)，我们希望 DiffSynth-Studio 能够成为一个更强大的 Diffusion 模型训练框架。
+
+- **2026年1月27日** [Z-Image](https://modelscope.cn/models/Tongyi-MAI/Z-Image) 发布，我们的 [Z-Image-i2L](https://www.modelscope.cn/models/DiffSynth-Studio/Z-Image-i2L) 模型同步发布，在[魔搭创空间](https://modelscope.cn/studios/DiffSynth-Studio/Z-Image-i2L)可直接体验，详见[文档](/docs/zh/Model_Details/Z-Image.md)。
+
+- **2026年1月19日** 新增对 [FLUX.2-klein-4B](https://modelscope.cn/models/black-forest-labs/FLUX.2-klein-4B) 和 [FLUX.2-klein-9B](https://modelscope.cn/models/black-forest-labs/FLUX.2-klein-9B) 模型的支持，包括完整的训练和推理功能。[文档](/docs/zh/Model_Details/FLUX2.md)和[示例代码](/examples/flux2/)现已可用。
+
+- **2026年1月12日** 我们训练并开源了一个文本引导的图层拆分模型（[模型链接](https://modelscope.cn/models/DiffSynth-Studio/Qwen-Image-Layered-Control)），这一模型输入一张图与一段文本描述，模型会将图像中与文本描述相关的图层拆分出来。更多细节请阅读我们的 blog（[中文版](https://modelscope.cn/learn/4938)、[英文版](https://huggingface.co/blog/kelseye/qwen-image-layered-control)）。
+
+- **2025年12月24日** 我们基于 Qwen-Image-Edit-2511 训练了一个 In-Context Editing LoRA 模型（[模型链接](https://modelscope.cn/models/DiffSynth-Studio/Qwen-Image-Edit-2511-ICEdit-LoRA)），这个模型可以输入三张图：图A、图B、图C，模型会自行分析图A到图B的变化，并将这样的变化应用到图C，生成图D。更多细节请阅读我们的 blog（[中文版](https://mp.weixin.qq.com/s/41aEiN3lXKGCJs1-we4Q2g)、[英文版](https://huggingface.co/blog/kelseye/qwen-image-edit-2511-icedit-lora)）。
+
+- **2025年12月9日** 我们基于 DiffSynth-Studio 2.0 训练了一个疯狂的模型：[Qwen-Image-i2L](https://www.modelscope.cn/models/DiffSynth-Studio/Qwen-Image-i2L)（Image to LoRA）。这一模型以图像为输入，以 LoRA 为输出。尽管这个版本的模型在泛化能力、细节保持能力等方面还有很大改进空间，我们将这些模型开源，以启发更多创新性的研究工作。更多细节，请参考我们的 [blog](https://huggingface.co/blog/kelseye/qwen-image-i2l)。
+
+- **2025年12月4日** DiffSynth-Studio 2.0 发布！众多新功能上线
+ - [文档](/docs/zh/README.md)上线：我们的文档还在持续优化更新中
+ - [显存管理](/docs/zh/Pipeline_Usage/VRAM_management.md)模块升级，支持 Layer 级别的 Disk Offload，同时释放内存与显存
+ - 新模型支持
+ - Z-Image Turbo: [模型](https://www.modelscope.ai/models/Tongyi-MAI/Z-Image-Turbo)、[文档](/docs/zh/Model_Details/Z-Image.md)、[代码](/examples/z_image/)
+ - FLUX.2-dev: [模型](https://www.modelscope.cn/models/black-forest-labs/FLUX.2-dev)、[文档](/docs/zh/Model_Details/FLUX2.md)、[代码](/examples/flux2/)
+ - 训练框架升级
+ - [拆分训练](/docs/zh/Training/Split_Training.md)：支持自动化地将训练过程拆分为数据处理和训练两阶段（即使训练的是 ControlNet 或其他任意模型），在数据处理阶段进行文本编码、VAE 编码等不需要梯度回传的计算，在训练阶段处理其他计算。速度更快，显存需求更少。
+ - [差分 LoRA 训练](/docs/zh/Training/Differential_LoRA.md)：这是我们曾在 [ArtAug](https://www.modelscope.cn/models/DiffSynth-Studio/ArtAug-lora-FLUX.1dev-v1) 中使用的训练技术，目前已可用于任意模型的 LoRA 训练。
+ - [FP8 训练](/docs/zh/Training/FP8_Precision.md)：FP8 在训练中支持应用到任意非训练模型，即梯度关闭或者梯度仅影响 LoRA 权重的模型。
+- **2025年11月4日** 支持了 [ByteDance/Video-As-Prompt-Wan2.1-14B](https://modelscope.cn/models/ByteDance/Video-As-Prompt-Wan2.1-14B) 模型，该模型基于 Wan 2.1 训练，支持根据参考视频生成相应的动作。
+
+- **2025年10月30日** 支持了 [meituan-longcat/LongCat-Video](https://www.modelscope.cn/models/meituan-longcat/LongCat-Video) 模型，该模型支持文生视频、图生视频、视频续写。这个模型在本项目中沿用 Wan 的框架进行推理和训练。
+
+- **2025年10月27日** 支持了 [krea/krea-realtime-video](https://www.modelscope.cn/models/krea/krea-realtime-video) 模型，Wan 模型生态再添一员。
+
+- **2025年9月23日** [DiffSynth-Studio/Qwen-Image-EliGen-Poster](https://www.modelscope.cn/models/DiffSynth-Studio/Qwen-Image-EliGen-Poster) 发布！本模型由我们与淘天体验设计团队联合研发并开源。模型基于 Qwen-Image 构建，专为电商海报场景设计，支持精确的分区布局控制。 请参考[我们的示例代码](./examples/qwen_image/model_inference/Qwen-Image-EliGen-Poster.py)。
+
+- **2025年9月9日** 我们的训练框架支持了多种训练模式，目前已适配 Qwen-Image，除标准 SFT 训练模式外，已支持 Direct Distill，请参考[我们的示例代码](./examples/qwen_image/model_training/lora/Qwen-Image-Distill-LoRA.sh)。这项功能是实验性的，我们将会继续完善已支持更全面的模型训练功能。
+
+- **2025年8月28日** 我们支持了Wan2.2-S2V，一个音频驱动的电影级视频生成模型。请参见[./examples/wanvideo/](./examples/wanvideo/)。
+
+- **2025年8月21日** [DiffSynth-Studio/Qwen-Image-EliGen-V2](https://www.modelscope.cn/models/DiffSynth-Studio/Qwen-Image-EliGen-V2) 发布！相比于 V1 版本，训练数据集变为 [Qwen-Image-Self-Generated-Dataset](https://www.modelscope.cn/datasets/DiffSynth-Studio/Qwen-Image-Self-Generated-Dataset)，因此，生成的图像更符合 Qwen-Image 本身的图像分布和风格。 请参考[我们的示例代码](./examples/qwen_image/model_inference_low_vram/Qwen-Image-EliGen-V2.py)。
+
+- **2025年8月21日** 我们开源了 [DiffSynth-Studio/Qwen-Image-In-Context-Control-Union](https://www.modelscope.cn/models/DiffSynth-Studio/Qwen-Image-In-Context-Control-Union) 结构控制 LoRA 模型，采用 In Context 的技术路线，支持多种类别的结构控制条件，包括 canny, depth, lineart, softedge, normal, openpose。 请参考[我们的示例代码](./examples/qwen_image/model_inference/Qwen-Image-In-Context-Control-Union.py)。
+
+- **2025年8月20日** 我们开源了 [DiffSynth-Studio/Qwen-Image-Edit-Lowres-Fix](https://www.modelscope.cn/models/DiffSynth-Studio/Qwen-Image-Edit-Lowres-Fix) 模型，提升了 Qwen-Image-Edit 对低分辨率图像输入的编辑效果。请参考[我们的示例代码](./examples/qwen_image/model_inference/Qwen-Image-Edit-Lowres-Fix.py)
+
+- **2025年8月19日** Qwen-Image-Edit 开源，欢迎图像编辑模型新成员！
+
+- **2025年8月18日** 我们训练并开源了 Qwen-Image 的图像重绘 ControlNet 模型 [DiffSynth-Studio/Qwen-Image-Blockwise-ControlNet-Inpaint](https://www.modelscope.cn/models/DiffSynth-Studio/Qwen-Image-Blockwise-ControlNet-Inpaint)，模型结构采用了轻量化的设计，请参考[我们的示例代码](./examples/qwen_image/model_inference/Qwen-Image-Blockwise-ControlNet-Inpaint.py)。
+
+- **2025年8月15日** 我们开源了 [Qwen-Image-Self-Generated-Dataset](https://www.modelscope.cn/datasets/DiffSynth-Studio/Qwen-Image-Self-Generated-Dataset) 数据集。这是一个使用 Qwen-Image 模型生成的图像数据集，共包含 160,000 张`1024 x 1024`图像。它包括通用、英文文本渲染和中文文本渲染子集。我们为每张图像提供了图像描述、实体和结构控制图像的标注。开发者可以使用这个数据集来训练 Qwen-Image 模型的 ControlNet 和 EliGen 等模型，我们旨在通过开源推动技术发展！
+
+- **2025年8月13日** 我们训练并开源了 Qwen-Image 的 ControlNet 模型 [DiffSynth-Studio/Qwen-Image-Blockwise-ControlNet-Depth](https://modelscope.cn/models/DiffSynth-Studio/Qwen-Image-Blockwise-ControlNet-Depth)，模型结构采用了轻量化的设计，请参考[我们的示例代码](./examples/qwen_image/model_inference/Qwen-Image-Blockwise-ControlNet-Depth.py)。
+
+- **2025年8月12日** 我们训练并开源了 Qwen-Image 的 ControlNet 模型 [DiffSynth-Studio/Qwen-Image-Blockwise-ControlNet-Canny](https://modelscope.cn/models/DiffSynth-Studio/Qwen-Image-Blockwise-ControlNet-Canny)，模型结构采用了轻量化的设计，请参考[我们的示例代码](./examples/qwen_image/model_inference/Qwen-Image-Blockwise-ControlNet-Canny.py)。
+
+- **2025年8月11日** 我们开源了 Qwen-Image 的蒸馏加速模型 [DiffSynth-Studio/Qwen-Image-Distill-LoRA](https://www.modelscope.cn/models/DiffSynth-Studio/Qwen-Image-Distill-LoRA)，沿用了与 [DiffSynth-Studio/Qwen-Image-Distill-Full](https://www.modelscope.cn/models/DiffSynth-Studio/Qwen-Image-Distill-Full) 相同的训练流程，但模型结构修改为了 LoRA，因此能够更好地与其他开源生态模型兼容。
+
+- **2025年8月7日** 我们开源了 Qwen-Image 的实体控制 LoRA 模型 [DiffSynth-Studio/Qwen-Image-EliGen](https://www.modelscope.cn/models/DiffSynth-Studio/Qwen-Image-EliGen)。Qwen-Image-EliGen 能够实现实体级可控的文生图。技术细节请参见[论文](https://arxiv.org/abs/2501.01097)。训练数据集：[EliGenTrainSet](https://www.modelscope.cn/datasets/DiffSynth-Studio/EliGenTrainSet)。
+
+- **2025年8月5日** 我们开源了 Qwen-Image 的蒸馏加速模型 [DiffSynth-Studio/Qwen-Image-Distill-Full](https://www.modelscope.cn/models/DiffSynth-Studio/Qwen-Image-Distill-Full)，实现了约 5 倍加速。
+
+- **2025年8月4日** Qwen-Image 开源，欢迎图像生成模型家族新成员！
+
+- **2025年8月1日** [FLUX.1-Krea-dev](https://www.modelscope.cn/models/black-forest-labs/FLUX.1-Krea-dev) 开源，这是一个专注于美学摄影的文生图模型。我们第一时间提供了全方位支持，包括低显存逐层 offload、LoRA 训练、全量训练。详细信息请参考 [./examples/flux/](./examples/flux/)。
+
+- **2025年7月28日** Wan 2.2 开源，我们第一时间提供了全方位支持，包括低显存逐层 offload、FP8 量化、序列并行、LoRA 训练、全量训练。详细信息请参考 [./examples/wanvideo/](./examples/wanvideo/)。
+
+- **2025年7月11日** 我们提出 Nexus-Gen，一个将大语言模型（LLM）的语言推理能力与扩散模型的图像生成能力相结合的统一框架。该框架支持无缝的图像理解、生成和编辑任务。
+ - 论文: [Nexus-Gen: Unified Image Understanding, Generation, and Editing via Prefilled Autoregression in Shared Embedding Space](https://arxiv.org/pdf/2504.21356)
+ - Github 仓库: https://github.com/modelscope/Nexus-Gen
+ - 模型: [ModelScope](https://www.modelscope.cn/models/DiffSynth-Studio/Nexus-GenV2), [HuggingFace](https://huggingface.co/modelscope/Nexus-GenV2)
+ - 训练数据集: [ModelScope Dataset](https://www.modelscope.cn/datasets/DiffSynth-Studio/Nexus-Gen-Training-Dataset)
+ - 在线体验: [ModelScope Nexus-Gen Studio](https://www.modelscope.cn/studios/DiffSynth-Studio/Nexus-Gen)
+- **2025年6月15日** ModelScope 官方评测框架 [EvalScope](https://github.com/modelscope/evalscope) 现已支持文生图生成评测。请参考[最佳实践](https://evalscope.readthedocs.io/zh-cn/latest/best_practice/t2i_eval.html)指南进行尝试。
+
+- **2025年3月25日** 我们的新开源项目 [DiffSynth-Engine](https://github.com/modelscope/DiffSynth-Engine) 现已开源！专注于稳定的模型部署，面向工业界，提供更好的工程支持、更高的计算性能和更稳定的功能。
+
+- **2025年3月31日** 我们支持 InfiniteYou，一种用于 FLUX 的人脸特征保留方法。更多细节请参考 [./examples/InfiniteYou/](./examples/InfiniteYou/)。
+
+- **2025年3月13日** 我们支持 HunyuanVideo-I2V，即腾讯开源的 HunyuanVideo 的图像到视频生成版本。更多细节请参考 [./examples/HunyuanVideo/](./examples/HunyuanVideo/)。
+
+- **2025年2月25日** 我们支持 Wan-Video，这是阿里巴巴开源的一系列最先进的视频合成模型。详见 [./examples/wanvideo/](./examples/wanvideo/)。
+
+- **2025年2月17日** 我们支持 [StepVideo](https://modelscope.cn/models/stepfun-ai/stepvideo-t2v/summary)！先进的视频合成模型！详见 [./examples/stepvideo](./examples/stepvideo/)。
+
+- **2024年12月31日** 我们提出 EliGen，一种用于精确实体级别控制的文本到图像生成的新框架，并辅以修复融合管道，将其能力扩展到图像修复任务。EliGen 可以无缝集成现有的社区模型，如 IP-Adapter 和 In-Context LoRA，提升其通用性。更多详情，请见 [./examples/EntityControl](./examples/EntityControl/)。
+ - 论文: [EliGen: Entity-Level Controlled Image Generation with Regional Attention](https://arxiv.org/abs/2501.01097)
+ - 模型: [ModelScope](https://www.modelscope.cn/models/DiffSynth-Studio/Eligen), [HuggingFace](https://huggingface.co/modelscope/EliGen)
+ - 在线体验: [ModelScope EliGen Studio](https://www.modelscope.cn/studios/DiffSynth-Studio/EliGen)
+ - 训练数据集: [EliGen Train Set](https://www.modelscope.cn/datasets/DiffSynth-Studio/EliGenTrainSet)
+- **2024年12月19日** 我们为 HunyuanVideo 实现了高级显存管理，使得在 24GB 显存下可以生成分辨率为 129x720x1280 的视频，或在仅 6GB 显存下生成分辨率为 129x512x384 的视频。更多细节请参考 [./examples/HunyuanVideo/](./examples/HunyuanVideo/)。
+
+- **2024年12月18日** 我们提出 ArtAug，一种通过合成-理解交互来改进文生图模型的方法。我们以 LoRA 格式为 FLUX.1-dev 训练了一个 ArtAug 增强模块。该模型将 Qwen2-VL-72B 的美学理解融入 FLUX.1-dev，从而提升了生成图像的质量。
+ - 论文: https://arxiv.org/abs/2412.12888
+ - 示例: https://github.com/modelscope/DiffSynth-Studio/tree/main/examples/ArtAug
+ - 模型: [ModelScope](https://www.modelscope.cn/models/DiffSynth-Studio/ArtAug-lora-FLUX.1dev-v1), [HuggingFace](https://huggingface.co/ECNU-CILab/ArtAug-lora-FLUX.1dev-v1)
+ - 演示: [ModelScope](https://modelscope.cn/aigc/imageGeneration?tab=advanced&versionId=7228&modelType=LoRA&sdVersion=FLUX_1&modelUrl=modelscope%3A%2F%2FDiffSynth-Studio%2FArtAug-lora-FLUX.1dev-v1%3Frevision%3Dv1.0), HuggingFace (即将上线)
+- **2024年10月25日** 我们提供了广泛的 FLUX ControlNet 支持。该项目支持许多不同的 ControlNet 模型，并且可以自由组合，即使它们的结构不同。此外，ControlNet 模型兼容高分辨率优化和分区控制技术，能够实现非常强大的可控图像生成。详见 [`./examples/ControlNet/`](./examples/ControlNet/)。
+
+- **2024年10月8日** 我们发布了基于 CogVideoX-5B 和 ExVideo 的扩展 LoRA。您可以从 [ModelScope](https://modelscope.cn/models/ECNU-CILab/ExVideo-CogVideoX-LoRA-129f-v1) 或 [HuggingFace](https://huggingface.co/ECNU-CILab/ExVideo-CogVideoX-LoRA-129f-v1) 下载此模型。
+
+- **2024年8月22日** 本项目现已支持 CogVideoX-5B。详见 [此处](/examples/video_synthesis/)。我们为这个文生视频模型提供了几个有趣的功能，包括：
+ - 文本到视频
+ - 视频编辑
+ - 自我超分
+ - 视频插帧
+- **2024年8月22日** 我们实现了一个有趣的画笔功能，支持所有文生图模型。现在，您可以在 AI 的辅助下使用画笔创作惊艳的图像了！
+ - 在我们的 [WebUI](#usage-in-webui) 中使用它。
+- **2024年8月21日** DiffSynth-Studio 现已支持 FLUX。
+ - 启用 CFG 和高分辨率修复以提升视觉质量。详见 [此处](/examples/image_synthesis/README.md)
+ - LoRA、ControlNet 和其他附加模型将很快推出。
+- **2024年6月21日** 我们提出 ExVideo，一种旨在增强视频生成模型能力的后训练微调技术。我们将 Stable Video Diffusion 进行了扩展，实现了长达 128 帧的长视频生成。
+ - [项目页面](https://ecnu-cilab.github.io/ExVideoProjectPage/)
+ - 源代码已在此仓库中发布。详见 [`examples/ExVideo`](./examples/ExVideo/)。
+ - 模型已发布于 [HuggingFace](https://huggingface.co/ECNU-CILab/ExVideo-SVD-128f-v1) 和 [ModelScope](https://modelscope.cn/models/ECNU-CILab/ExVideo-SVD-128f-v1)。
+ - 技术报告已发布于 [arXiv](https://arxiv.org/abs/2406.14130)。
+ - 您可以在此 [演示](https://huggingface.co/spaces/modelscope/ExVideo-SVD-128f-v1) 中试用 ExVideo！
+- **2024年6月13日** DiffSynth Studio 已迁移至 ModelScope。开发团队也从“我”转变为“我们”。当然，我仍会参与后续的开发和维护工作。
+
+- **2024年1月29日** 我们提出 Diffutoon，这是一个出色的卡通着色解决方案。
+ - [项目页面](https://ecnu-cilab.github.io/DiffutoonProjectPage/)
+ - 源代码已在此项目中发布。
+ - 技术报告（IJCAI 2024）已发布于 [arXiv](https://arxiv.org/abs/2401.16224)。
+- **2023年12月8日** 我们决定启动一个新项目，旨在释放扩散模型的潜力，尤其是在视频合成方面。该项目的开发工作正式开始。
+
+- **2023年11月15日** 我们提出 FastBlend，一种强大的视频去闪烁算法。
+ - sd-webui 扩展已发布于 [GitHub](https://github.com/Artiprocher/sd-webui-fastblend)。
+ - 演示视频已在 Bilibili 上展示，包含三个任务：
+ - [视频去闪烁](https://www.bilibili.com/video/BV1d94y1W7PE)
+ - [视频插帧](https://www.bilibili.com/video/BV1Lw411m71p)
+ - [图像驱动的视频渲染](https://www.bilibili.com/video/BV1RB4y1Z7LF)
+ - 技术报告已发布于 [arXiv](https://arxiv.org/abs/2311.09265)。
+ - 其他用户开发的非官方 ComfyUI 扩展已发布于 [GitHub](https://github.com/AInseven/ComfyUI-fastblend)。
+- **2023年10月1日** 我们发布了该项目的早期版本，名为 FastSDXL。这是构建一个扩散引擎的初步尝试。
+ - 源代码已发布于 [GitHub](https://github.com/Artiprocher/FastSDXL)。
+ - FastSDXL 包含一个可训练的 OLSS 调度器，以提高效率。
+ - OLSS 的原始仓库位于 [此处](https://github.com/alibaba/EasyNLP/tree/master/diffusion/olss_scheduler)。
+ - 技术报告（CIKM 2023）已发布于 [arXiv](https://arxiv.org/abs/2305.14677)。
+ - 演示视频已发布于 [Bilibili](https://www.bilibili.com/video/BV1w8411y7uj)。
+ - 由于 OLSS 需要额外训练，我们未在本项目中实现它。
+- **2023年8月29日** 我们提出 DiffSynth，一个视频合成框架。
+ - [项目页面](https://ecnu-cilab.github.io/DiffSynth.github.io/)。
+ - 源代码已发布在 [EasyNLP](https://github.com/alibaba/EasyNLP/tree/master/diffusion/DiffSynth)。
+ - 技术报告（ECML PKDD 2024）已发布于 [arXiv](https://arxiv.org/abs/2308.03463)。
+
+
+
+## 安装
+
+从源码安装（推荐）：
+
+```
+git clone https://github.com/modelscope/DiffSynth-Studio.git
+cd DiffSynth-Studio
+pip install -e .
+```
+
+更多安装方式，以及非 NVIDIA GPU 的安装，请参考[安装文档](/docs/zh/Pipeline_Usage/Setup.md)。
+
+
+下载源配置
+
+> 在进行模型推理和训练前，可通过[环境变量](/docs/zh/Pipeline_Usage/Environment_Variables.md)配置模型下载源等。
+>
+> 本项目默认从[魔搭社区](https://modelscope.cn/)下载模型。对于非中国区域的用户，可以通过以下配置从[魔搭社区国际站](https://modelscope.ai)下载模型：
+>
+> ```shell
+> export MODELSCOPE_ENDPOINT=https://modelscope.ai
+> ```
+>
+> 如需从 [HuggingFace](https://huggingface.co/) 下载模型，请修改[环境变量](/docs/zh/Pipeline_Usage/Environment_Variables.md)，注意不同模型平台上的模型 ID 可能不同：
+>
+> ```shell
+> export DIFFSYNTH_DOWNLOAD_SOURCE="huggingface"
+> ```
+
+
+
+## 基础框架
+
+DiffSynth-Studio 作为基础的模型框架，为主流 Diffusion 模型重新设计了推理和训练流水线，能够实现高效的显存管理、灵活的模型训练。
+
+快速开始，体验热门和最新模型：
+
+| 基础架构 | 模型 ID | 推理 | 低显存推理 | 全量训练 | 全量训练后验证 | LoRA 训练 | LoRA 训练后验证 |
+|-|-|-|-|-|-|-|-|
+| MiniMax-H3 | [MiniMax/MiniMax-H3: FL2VA](https://www.modelscope.cn/models/MiniMax/MiniMax-H3) | [code](/examples/minimax_h3/model_inference/MiniMax-H3-FL2VA.py) | [code](/examples/minimax_h3/model_inference_low_vram/MiniMax-H3-FL2VA.py) | [code](/examples/minimax_h3/model_training/full/MiniMax-H3-FL2VA.sh) | [code](/examples/minimax_h3/model_training/validate_full/MiniMax-H3-FL2VA.py) | [code](/examples/minimax_h3/model_training/lora/MiniMax-H3-FL2VA.sh) | [code](/examples/minimax_h3/model_training/validate_lora/MiniMax-H3-FL2VA.py) |
+| MiniMax-H3 | [DiffSynth-Studio/MiniMax-H3-NF4: FL2VA pruned](https://www.modelscope.cn/models/DiffSynth-Studio/MiniMax-H3-NF4) | [code](/examples/minimax_h3/model_inference/MiniMax-H3-NF4-Pruned-FL2VA.py) | [code](/examples/minimax_h3/model_inference_low_vram/MiniMax-H3-NF4-Pruned-FL2VA.py) | - | - | [code](/examples/minimax_h3/model_training/lora/MiniMax-H3-NF4-Pruned-FL2VA.sh) | [code](/examples/minimax_h3/model_training/validate_lora/MiniMax-H3-NF4-Pruned-FL2VA.py) |
+| ACE-Step | [ACE-Step/acestep-v15-xl-sft](https://www.modelscope.cn/models/ACE-Step/acestep-v15-xl-sft) | [code](/examples/ace_step/model_inference/acestep-v15-xl-sft.py) | [code](/examples/ace_step/model_inference_low_vram/acestep-v15-xl-sft.py) | [code](/examples/ace_step/model_training/full/acestep-v15-xl-sft.sh) | [code](/examples/ace_step/model_training/validate_full/acestep-v15-xl-sft.py) | [code](/examples/ace_step/model_training/lora/acestep-v15-xl-sft.sh) | [code](/examples/ace_step/model_training/validate_lora/acestep-v15-xl-sft.py) |
+| Z-Image | [Tongyi-MAI/Z-Image-Turbo](https://www.modelscope.cn/models/Tongyi-MAI/Z-Image-Turbo) | [code](/examples/z_image/model_inference/Z-Image-Turbo.py) | [code](/examples/z_image/model_inference_low_vram/Z-Image-Turbo.py) | [code](/examples/z_image/model_training/full/Z-Image-Turbo.sh) | [code](/examples/z_image/model_training/validate_full/Z-Image-Turbo.py) | [code](/examples/z_image/model_training/lora/Z-Image-Turbo.sh) | [code](/examples/z_image/model_training/validate_lora/Z-Image-Turbo.py) |
+| Krea-2 | [krea/Krea-2-Raw](https://www.modelscope.cn/models/krea/Krea-2-Raw) | [code](/examples/krea2/model_inference/Krea-2-Raw.py) | [code](/examples/krea2/model_inference_low_vram/Krea-2-Raw.py) | [code](/examples/krea2/model_training/full/Krea-2-Raw.sh) | [code](/examples/krea2/model_training/validate_full/Krea-2-Raw.py) | [code](/examples/krea2/model_training/lora/Krea-2-Raw.sh) | [code](/examples/krea2/model_training/validate_lora/Krea-2-Raw.py) |
+| Krea-2 | [krea/Krea-2-Turbo](https://www.modelscope.cn/models/krea/Krea-2-Turbo) | [code](/examples/krea2/model_inference/Krea-2-Turbo.py) | [code](/examples/krea2/model_inference_low_vram/Krea-2-Turbo.py) | [code](/examples/krea2/model_training/full/Krea-2-Turbo.sh) | [code](/examples/krea2/model_training/validate_full/Krea-2-Turbo.py) | [code](/examples/krea2/model_training/lora/Krea-2-Turbo.sh) | [code](/examples/krea2/model_training/validate_lora/Krea-2-Turbo.py) |
+
+模型一览：
+
+- 图像生成
+ - Boogu-Image：[文档](https://diffsynth-studio-doc.readthedocs.io/zh-cn/latest/Model_Details/Boogu-Image.html)、[样例代码](/examples/boogu_image/)
+ - Krea-2：[文档](https://diffsynth-studio-doc.readthedocs.io/zh-cn/latest/Model_Details/Krea-2.html)、[样例代码](/examples/krea2/)
+ - Ideogram 4：[文档](https://diffsynth-studio-doc.readthedocs.io/zh-cn/latest/Model_Details/Ideogram-4.html)、[样例代码](/examples/ideogram4/)
+ - HiDream-O1-Image：[文档](https://diffsynth-studio-doc.readthedocs.io/zh-cn/latest/Model_Details/HiDream-O1-Image.html)、[样例代码](/examples/hidream_o1_image/)
+ - JoyAI-Image：[文档](https://diffsynth-studio-doc.readthedocs.io/zh-cn/latest/Model_Details/JoyAI-Image.html)、[样例代码](/examples/joyai_image/)
+ - ERNIE-Image：[文档](https://diffsynth-studio-doc.readthedocs.io/zh-cn/latest/Model_Details/ERNIE-Image.html)、[样例代码](/examples/ernie_image/)
+ - FLUX.2：[文档](https://diffsynth-studio-doc.readthedocs.io/zh-cn/latest/Model_Details/FLUX2.html)、[样例代码](/examples/flux2/)
+ - Z-Image：[文档](https://diffsynth-studio-doc.readthedocs.io/zh-cn/latest/Model_Details/Z-Image.html)、[样例代码](/examples/z_image/)
+ - Anima：[文档](https://diffsynth-studio-doc.readthedocs.io/zh-cn/latest/Model_Details/Anima.html)、[样例代码](/examples/anima/)
+ - Qwen-Image：[文档](https://diffsynth-studio-doc.readthedocs.io/zh-cn/latest/Model_Details/Qwen-Image.html)、[样例代码](/examples/qwen_image/)
+ - FLUX.1：[文档](https://diffsynth-studio-doc.readthedocs.io/zh-cn/latest/Model_Details/FLUX.html)、[样例代码](/examples/flux/)
+ - Stable Diffusion XL：[文档](https://diffsynth-studio-doc.readthedocs.io/zh-cn/latest/Model_Details/Stable-Diffusion-XL.html)、[样例代码](/examples/stable_diffusion_xl/)
+ - Stable Diffusion：[文档](https://diffsynth-studio-doc.readthedocs.io/zh-cn/latest/Model_Details/Stable-Diffusion.html)、[样例代码](/examples/stable_diffusion/)
+- 视频生成
+ - MiniMax-H3：[文档](https://diffsynth-studio-doc.readthedocs.io/zh-cn/latest/Model_Details/MiniMax-H3.html)、[样例代码](/examples/minimax_h3/)
+ - LingBot-Video：[文档](https://diffsynth-studio-doc.readthedocs.io/zh-cn/latest/Model_Details/LingBot-Video.html)、[样例代码](/examples/lingbot_video/)
+ - LTX-2：[文档](https://diffsynth-studio-doc.readthedocs.io/zh-cn/latest/Model_Details/LTX-2.html)、[样例代码](/examples/ltx2/)
+ - Wan：[文档](https://diffsynth-studio-doc.readthedocs.io/zh-cn/latest/Model_Details/Wan.html)、[样例代码](/examples/wanvideo/)
+- 音频生成
+ - MiniMax-Music3：[文档](https://diffsynth-studio-doc.readthedocs.io/zh-cn/latest/Model_Details/MiniMax-Music3.html)、[样例代码](/examples/minimax_music3/)
+ - ACE-Step：[文档](https://diffsynth-studio-doc.readthedocs.io/zh-cn/latest/Model_Details/ACE-Step.html)、[样例代码](/examples/ace_step/)
+- 评测模型：[文档](https://diffsynth-studio-doc.readthedocs.io/zh-cn/latest/Model_Details/Image-Quality-Metrics.html)、[样例代码](/examples/image_quality_metric/)
+
+[查看全部支持的模型](#全部支持的模型)
+
+## 创新成果
+
+我们相信，一个完善的开源代码框架能够降低技术探索的门槛，我们基于这个代码库搞出了不少有意思的技术。或许你也有许多天马行空的构想，借助 DiffSynth-Studio，你可以快速实现这些想法。
+
+
+TreeAdapter: 由结构化 LoRA 组建的模型系统
+
+> 用 1w+ LoRA 组成的模型系统挑战上万个稀有物种的精准生成。
+
+* 论文：[TreeAdapter: Hierarchical Taxonomy-Guided Adapter Composition for Fine-Grained Species Image Generation](https://arxiv.org/abs/2607.24215)
+* 模型：[ModelScope](https://modelscope.cn/models/DiffSynth-Studio/TreeAdapter-KleinBase4B)
+
+
+
+
+
+
+Image-to-LoRA: 把模型训练压缩到模型推理中
+
+> 真正的 Meta Learning：模型一端输入图像数据集，另一端直接输出训出来的 LoRA 模型。
+
+* 论文：[Compressing Image Style Training into a Single Model Forward](https://arxiv.org/abs/2606.13809)
+* 模型：
+ * [DiffSynth-Studio/ZImage-i2L-v2](https://modelscope.cn/models/DiffSynth-Studio/ZImage-i2L-v2)：适配模型 Z-Image
+ * [DiffSynth-Studio/KleinBase4B-i2L-v2](https://modelscope.cn/models/DiffSynth-Studio/KleinBase4B-i2L-v2)：适配模型 FLUX.2-klein-base-4B
+ * [DiffSynth-Studio/HidreamO1-i2L-v2](https://modelscope.cn/models/DiffSynth-Studio/HidreamO1-i2L-v2)：适配模型 Hidream-O1-Image
+
+|输入样例1|输出样例1|输入样例2|输出样例2|
+|-|-|-|-|
+|||||
+
+
+
+
+Diffusion-Templates: 插件化的可控生成框架
+
+> 一个框架，让每一种可控生成能力都成为插件，让多个模型组合涌现出丰富生成能力。
+
+* 论文：[Diffusion Templates: A Unified Plugin Framework for Controllable Diffusion](https://arxiv.org/abs/2604.24351)
+* 项目主页：[GitHub](https://modelscope.github.io/diffusion-templates-web/)
+* 文档参考：[English Version](https://diffsynth-studio-doc.readthedocs.io/en/latest/Diffusion_Templates/Introducing_Diffusion_Templates.html)、[中文版](https://diffsynth-studio-doc.readthedocs.io/zh-cn/latest/Diffusion_Templates/Introducing_Diffusion_Templates.html)
+* 在线体验：[魔搭社区创空间](https://modelscope.cn/studios/DiffSynth-Studio/Diffusion-Templates)
+* 模型集：[ModelScope](https://modelscope.cn/collections/DiffSynth-Studio/KleinBase4B-Templates)、[ModelScope 国际站](https://modelscope.ai/collections/DiffSynth-Studio/KleinBase4B-Templates)、[HuggingFace](https://huggingface.co/collections/DiffSynth-Studio/kleinbase4b-templates)
+* 数据集：[ModelScope](https://modelscope.cn/collections/DiffSynth-Studio/ImagePulseV2)、[ModelScope 国际站](https://modelscope.cn/collections/DiffSynth-Studio/ImagePulseV2)、[HuggingFace](https://huggingface.co/collections/DiffSynth-Studio/imagepulsev2)
+
+|参考图|局部编辑|风格迁移|清晰度增强|
+|-|-|-|-|
+|||||
+
+
+
+
+SES: 用于奖励对齐图像生成的高效推理阶段缩放
+
+> 以推理时间为代价，提高生成内容的质量。
+
+- 论文：[Spectral Evolution Search: Efficient Inference-Time Scaling for Reward-Aligned Image Generation](https://arxiv.org/abs/2602.03208)
+- 代码样例：[/docs/en/Research_Tutorial/inference_time_scaling.md](/docs/en/Research_Tutorial/inference_time_scaling.md)
+
+|FLUX.1-dev|FLUX.1-dev + SES|Qwen-Image|Qwen-Image + SES|
+|-|-|-|-|
+|||||
+
+
+
+
+VIRAL：基于DiT模型的类比视觉上下文推理
+
+> 根据图1到图2的变化把图3转化为图4，图像编辑模型的能力涌现。
+
+- 论文：[VIRAL: Visual In-Context Reasoning via Analogy in Diffusion Transformers
+](https://arxiv.org/abs/2602.03210)
+- 代码样例：[/examples/qwen_image/model_inference/Qwen-Image-Edit-2511-ICEdit.py](/examples/qwen_image/model_inference/Qwen-Image-Edit-2511-ICEdit.py)
+- 模型：[ModelScope](https://www.modelscope.cn/models/DiffSynth-Studio/Qwen-Image-Edit-2511-ICEdit-LoRA)
+
+|Example 1|Example 2|Query|Output|
+|-|-|-|-|
+|||||
+
+
+
+
+AttriCtrl: 图像生成模型的属性强度控制
+
+> 用数值属性也能精准地控制图像生成模型。
+
+- 论文：[AttriCtrl: Fine-Grained Control of Aesthetic Attribute Intensity in Diffusion Models
+](https://arxiv.org/abs/2508.02151)
+- 代码样例：[/examples/flux/model_inference/FLUX.1-dev-AttriCtrl.py](/examples/flux/model_inference/FLUX.1-dev-AttriCtrl.py)
+- 模型：[ModelScope](https://www.modelscope.cn/models/DiffSynth-Studio/AttriCtrl-FLUX.1-Dev)
+
+|brightness scale = 0.1|brightness scale = 0.3|brightness scale = 0.5|brightness scale = 0.7|brightness scale = 0.9|
+|-|-|-|-|-|
+||||||
+
+
+
+
+AutoLoRA: 自动化的 LoRA 检索和融合
+
+> LoRA 是需求和解决方案一体化的产物，如何更好地利用这些 LoRA？
+
+- 论文：[AutoLoRA: Automatic LoRA Retrieval and Fine-Grained Gated Fusion for Text-to-Image Generation
+](https://arxiv.org/abs/2508.02107)
+- 代码样例：[/examples/flux/model_inference/FLUX.1-dev-LoRA-Fusion.py](/examples/flux/model_inference/FLUX.1-dev-LoRA-Fusion.py)
+- 模型：[ModelScope](https://www.modelscope.cn/models/DiffSynth-Studio/LoRAFusion-preview-FLUX.1-dev)
+
+||[LoRA 1](https://modelscope.cn/models/cancel13/cxsk)|[LoRA 2](https://modelscope.cn/models/wy413928499/xuancai2)|[LoRA 3](https://modelscope.cn/models/DiffSynth-Studio/ArtAug-lora-FLUX.1dev-v1)|[LoRA 4](https://modelscope.cn/models/hongyanbujian/JPL)|
+|-|-|-|-|-|
+|[LoRA 1](https://modelscope.cn/models/cancel13/cxsk) |||||
+|[LoRA 2](https://modelscope.cn/models/wy413928499/xuancai2) |||||
+|[LoRA 3](https://modelscope.cn/models/DiffSynth-Studio/ArtAug-lora-FLUX.1dev-v1) |||||
+|[LoRA 4](https://modelscope.cn/models/hongyanbujian/JPL) |||||
+
+
+
+
+Nexus-Gen: 统一架构的图像理解、生成、编辑
+
+> 如果一个模型集齐图像理解、图像生成、图像编辑能力，会发生什么？
+
+- 详细页面：https://github.com/modelscope/Nexus-Gen
+- 论文：[Nexus-Gen: Unified Image Understanding, Generation, and Editing via Prefilled Autoregression in Shared Embedding Space](https://arxiv.org/pdf/2504.21356)
+- 模型：[ModelScope](https://www.modelscope.cn/models/DiffSynth-Studio/Nexus-GenV2), [HuggingFace](https://huggingface.co/modelscope/Nexus-GenV2)
+- 数据集：[ModelScope Dataset](https://www.modelscope.cn/datasets/DiffSynth-Studio/Nexus-Gen-Training-Dataset)
+- 在线体验：[ModelScope Nexus-Gen Studio](https://www.modelscope.cn/studios/DiffSynth-Studio/Nexus-Gen)
+
+
+
+
+
+
+ArtAug: 图像生成模型的美学提升
+
+> 一个 LoRA，大幅提升细节和美感。
+
+- 详细页面：[./examples/ArtAug/](./examples/ArtAug/)
+- 论文：[ArtAug: Enhancing Text-to-Image Generation through Synthesis-Understanding Interaction](https://arxiv.org/abs/2412.12888)
+- 模型：[ModelScope](https://www.modelscope.cn/models/DiffSynth-Studio/ArtAug-lora-FLUX.1dev-v1), [HuggingFace](https://huggingface.co/ECNU-CILab/ArtAug-lora-FLUX.1dev-v1)
+- 在线体验：[ModelScope AIGC Tab](https://www.modelscope.cn/aigc/imageGeneration?tab=advanced&versionId=7228&modelType=LoRA&sdVersion=FLUX_1&modelUrl=modelscope%3A%2F%2FDiffSynth-Studio%2FArtAug-lora-FLUX.1dev-v1%3Frevision%3Dv1.0)
+
+|FLUX.1-dev|FLUX.1-dev + ArtAug LoRA|
+|-|-|
+|||
+
+
+
+
+EliGen: 精准的图像分区控制
+
+> 如何用分区图层控制画面内容的位置？
+
+- 论文：[EliGen: Entity-Level Controlled Image Generation with Regional Attention](https://arxiv.org/abs/2501.01097)
+- 代码样例：[/examples/flux/model_inference/FLUX.1-dev-EliGen.py](/examples/flux/model_inference/FLUX.1-dev-EliGen.py)
+- 模型：[ModelScope](https://www.modelscope.cn/models/DiffSynth-Studio/Eligen), [HuggingFace](https://huggingface.co/modelscope/EliGen)
+- 在线体验：[ModelScope EliGen Studio](https://www.modelscope.cn/studios/DiffSynth-Studio/EliGen)
+- 数据集：[EliGen Train Set](https://www.modelscope.cn/datasets/DiffSynth-Studio/EliGenTrainSet)
+
+|实体控制区域|生成图像|
+|-|-|
+|||
+
+
+
+
+ExVideo: 视频生成模型的扩展训练
+
+> 如果视频生成模型只能生成 25 帧，如何才能让它生成更长视频？
+
+- 项目页面：[Project Page](https://ecnu-cilab.github.io/ExVideoProjectPage/)
+- 论文：[ExVideo: Extending Video Diffusion Models via Parameter-Efficient Post-Tuning](https://arxiv.org/abs/2406.14130)
+- 代码样例：请前往[旧版本](https://github.com/modelscope/DiffSynth-Studio/tree/afd101f3452c9ecae0c87b79adfa2e22d65ffdc3/examples/ExVideo)查看
+- 模型：[ModelScope](https://modelscope.cn/models/ECNU-CILab/ExVideo-SVD-128f-v1), [HuggingFace](https://huggingface.co/ECNU-CILab/ExVideo-SVD-128f-v1)
+
+https://github.com/modelscope/DiffSynth-Studio/assets/35051019/d97f6aa9-8064-4b5b-9d49-ed6001bb9acc
+
+
+
+
+Diffutoon: 高分辨率动漫风格视频渲染
+
+> 我管你这那的，我就是喜欢二次元！
+
+- 项目页面：[Project Page](https://ecnu-cilab.github.io/DiffutoonProjectPage/)
+- 论文：[Diffutoon: High-Resolution Editable Toon Shading via Diffusion Models](https://arxiv.org/abs/2401.16224)
+- 代码样例：请前往[旧版本](https://github.com/modelscope/DiffSynth-Studio/tree/afd101f3452c9ecae0c87b79adfa2e22d65ffdc3/examples/Diffutoon)查看
+
+注意：Diffutoon 示例来自 2.0 之前的代码库，目前不包含在 `main` 分支中。请将上面的旧版本代码与对应版本的 DiffSynth-Studio 一起使用。
+
+https://github.com/Artiprocher/DiffSynth-Studio/assets/35051019/b54c05c5-d747-4709-be5e-b39af82404dd
+
+
+
+
+DiffSynth: 本项目的初代版本
+
+> 在没有视频生成模型的时代，如何利用图像生成模型处理视频？
+
+- 项目页面：[Project Page](https://ecnu-cilab.github.io/DiffSynth.github.io/)
+- 论文：[DiffSynth: Latent In-Iteration Deflickering for Realistic Video Synthesis](https://arxiv.org/abs/2308.03463)
+- 代码样例：请前往[旧版本](https://github.com/modelscope/DiffSynth-Studio/tree/afd101f3452c9ecae0c87b79adfa2e22d65ffdc3/examples/diffsynth)查看
+
+https://github.com/Artiprocher/DiffSynth-Studio/assets/35051019/59fb2f7b-8de0-4481-b79f-0c3a7361a1ea
+
+
+
+## 联系我们
+
+|Discord：https://discord.gg/Mm9suEeUDc|
+|-|
+||
+
+## 全部支持的模型
+
+| 基础架构 | 模型 ID | 推理 | 低显存推理 | 全量训练 | 全量训练后验证 | LoRA 训练 | LoRA 训练后验证 |
+|-|-|-|-|-|-|-|-|
+| MiniMax-Music3 | [MiniMax/MiniMax-Music3](https://www.modelscope.cn/models/MiniMax/MiniMax-Music3) | [code](/examples/minimax_music3/model_inference/MiniMax-Music3.py) | [code](/examples/minimax_music3/model_inference_low_vram/MiniMax-Music3.py) | — | — | — | — |
+| MiniMax-H3 | [MiniMax/MiniMax-H3: FL2VA](https://www.modelscope.cn/models/MiniMax/MiniMax-H3) | [code](/examples/minimax_h3/model_inference/MiniMax-H3-FL2VA.py) | [code](/examples/minimax_h3/model_inference_low_vram/MiniMax-H3-FL2VA.py) | [code](/examples/minimax_h3/model_training/full/MiniMax-H3-FL2VA.sh) | [code](/examples/minimax_h3/model_training/validate_full/MiniMax-H3-FL2VA.py) | [code](/examples/minimax_h3/model_training/lora/MiniMax-H3-FL2VA.sh) | [code](/examples/minimax_h3/model_training/validate_lora/MiniMax-H3-FL2VA.py) |
+| MiniMax-H3 | [MiniMax/MiniMax-H3: Ref2VA](https://www.modelscope.cn/models/MiniMax/MiniMax-H3) | [code](/examples/minimax_h3/model_inference/MiniMax-H3-Ref2VA.py) | [code](/examples/minimax_h3/model_inference_low_vram/MiniMax-H3-Ref2VA.py) | [code](/examples/minimax_h3/model_training/full/MiniMax-H3-Ref2VA.sh) | [code](/examples/minimax_h3/model_training/validate_full/MiniMax-H3-Ref2VA.py) | [code](/examples/minimax_h3/model_training/lora/MiniMax-H3-Ref2VA.sh) | [code](/examples/minimax_h3/model_training/validate_lora/MiniMax-H3-Ref2VA.py) |
+| MiniMax-H3 | [MiniMax/MiniMax-H3: Retake](https://www.modelscope.cn/models/MiniMax/MiniMax-H3) | [code](/examples/minimax_h3/model_inference/MiniMax-H3-Retake.py) | [code](/examples/minimax_h3/model_inference_low_vram/MiniMax-H3-Retake.py) | - | - | - | - |
+| MiniMax-H3 | [DiffSynth-Studio/MiniMax-H3-NF4: FL2VA](https://www.modelscope.cn/models/DiffSynth-Studio/MiniMax-H3-NF4) | [code](/examples/minimax_h3/model_inference/MiniMax-H3-NF4-FL2VA.py) | [code](/examples/minimax_h3/model_inference_low_vram/MiniMax-H3-NF4-FL2VA.py) | - | - | [code](/examples/minimax_h3/model_training/lora/MiniMax-H3-NF4-FL2VA.sh) | [code](/examples/minimax_h3/model_training/validate_lora/MiniMax-H3-NF4-FL2VA.py) |
+| MiniMax-H3 | [DiffSynth-Studio/MiniMax-H3-NF4: Ref2VA](https://www.modelscope.cn/models/DiffSynth-Studio/MiniMax-H3-NF4) | [code](/examples/minimax_h3/model_inference/MiniMax-H3-NF4-Ref2VA.py) | [code](/examples/minimax_h3/model_inference_low_vram/MiniMax-H3-NF4-Ref2VA.py) | - | - | [code](/examples/minimax_h3/model_training/lora/MiniMax-H3-NF4-Ref2VA.sh) | [code](/examples/minimax_h3/model_training/validate_lora/MiniMax-H3-NF4-Ref2VA.py) |
+| MiniMax-H3 | [Comfy-Org/MiniMax-H3: FL2VA pruned](https://www.modelscope.cn/models/Comfy-Org/MiniMax-H3) | [code](/examples/minimax_h3/model_inference/MiniMax-H3-Pruned-FL2VA.py) | [code](/examples/minimax_h3/model_inference_low_vram/MiniMax-H3-Pruned-FL2VA.py) | [code](/examples/minimax_h3/model_training/full/MiniMax-H3-Pruned-FL2VA.sh) | [code](/examples/minimax_h3/model_training/validate_full/MiniMax-H3-Pruned-FL2VA.py) | [code](/examples/minimax_h3/model_training/lora/MiniMax-H3-Pruned-FL2VA.sh) | [code](/examples/minimax_h3/model_training/validate_lora/MiniMax-H3-Pruned-FL2VA.py) |
+| MiniMax-H3 | [Comfy-Org/MiniMax-H3: Ref2VA pruned](https://www.modelscope.cn/models/Comfy-Org/MiniMax-H3) | [code](/examples/minimax_h3/model_inference/MiniMax-H3-Pruned-Ref2VA.py) | [code](/examples/minimax_h3/model_inference_low_vram/MiniMax-H3-Pruned-Ref2VA.py) | [code](/examples/minimax_h3/model_training/full/MiniMax-H3-Pruned-Ref2VA.sh) | [code](/examples/minimax_h3/model_training/validate_full/MiniMax-H3-Pruned-Ref2VA.py) | [code](/examples/minimax_h3/model_training/lora/MiniMax-H3-Pruned-Ref2VA.sh) | [code](/examples/minimax_h3/model_training/validate_lora/MiniMax-H3-Pruned-Ref2VA.py) |
+| MiniMax-H3 | [DiffSynth-Studio/MiniMax-H3-NF4: FL2VA pruned](https://www.modelscope.cn/models/DiffSynth-Studio/MiniMax-H3-NF4) | [code](/examples/minimax_h3/model_inference/MiniMax-H3-NF4-Pruned-FL2VA.py) | [code](/examples/minimax_h3/model_inference_low_vram/MiniMax-H3-NF4-Pruned-FL2VA.py) | - | - | [code](/examples/minimax_h3/model_training/lora/MiniMax-H3-NF4-Pruned-FL2VA.sh) | [code](/examples/minimax_h3/model_training/validate_lora/MiniMax-H3-NF4-Pruned-FL2VA.py) |
+| MiniMax-H3 | [DiffSynth-Studio/MiniMax-H3-NF4: Ref2VA pruned](https://www.modelscope.cn/models/DiffSynth-Studio/MiniMax-H3-NF4) | [code](/examples/minimax_h3/model_inference/MiniMax-H3-NF4-Pruned-Ref2VA.py) | [code](/examples/minimax_h3/model_inference_low_vram/MiniMax-H3-NF4-Pruned-Ref2VA.py) | - | - | [code](/examples/minimax_h3/model_training/lora/MiniMax-H3-NF4-Pruned-Ref2VA.sh) | [code](/examples/minimax_h3/model_training/validate_lora/MiniMax-H3-NF4-Pruned-Ref2VA.py) |
+| MiniMax-H3 | [Comfy-Org/MiniMax-H3: FL2VA int8_convrot](https://www.modelscope.cn/models/Comfy-Org/MiniMax-H3) | [code](/examples/minimax_h3/model_inference/MiniMax-H3-Int8-ConvRot-FL2VA.py) | [code](/examples/minimax_h3/model_inference_low_vram/MiniMax-H3-Int8-ConvRot-FL2VA.py) | - | - | [code](/examples/minimax_h3/model_training/lora/MiniMax-H3-Int8-ConvRot-FL2VA.sh) | [code](/examples/minimax_h3/model_training/validate_lora/MiniMax-H3-Int8-ConvRot-FL2VA.py) |
+| MiniMax-H3 | [Comfy-Org/MiniMax-H3: Ref2VA int8_convrot](https://www.modelscope.cn/models/Comfy-Org/MiniMax-H3) | [code](/examples/minimax_h3/model_inference/MiniMax-H3-Int8-ConvRot-Ref2VA.py) | [code](/examples/minimax_h3/model_inference_low_vram/MiniMax-H3-Int8-ConvRot-Ref2VA.py) | - | - | [code](/examples/minimax_h3/model_training/lora/MiniMax-H3-Int8-ConvRot-Ref2VA.sh) | [code](/examples/minimax_h3/model_training/validate_lora/MiniMax-H3-Int8-ConvRot-Ref2VA.py) |
+| MiniMax-H3 | [Comfy-Org/MiniMax-H3: FL2VA pruned int8_convrot](https://www.modelscope.cn/models/Comfy-Org/MiniMax-H3) | [code](/examples/minimax_h3/model_inference/MiniMax-H3-Int8-ConvRot-Pruned-FL2VA.py) | [code](/examples/minimax_h3/model_inference_low_vram/MiniMax-H3-Int8-ConvRot-Pruned-FL2VA.py) | - | - | [code](/examples/minimax_h3/model_training/lora/MiniMax-H3-Int8-ConvRot-Pruned-FL2VA.sh) | [code](/examples/minimax_h3/model_training/validate_lora/MiniMax-H3-Int8-ConvRot-Pruned-FL2VA.py) |
+| MiniMax-H3 | [Comfy-Org/MiniMax-H3: Ref2VA pruned int8_convrot](https://www.modelscope.cn/models/Comfy-Org/MiniMax-H3) | [code](/examples/minimax_h3/model_inference/MiniMax-H3-Int8-ConvRot-Pruned-Ref2VA.py) | [code](/examples/minimax_h3/model_inference_low_vram/MiniMax-H3-Int8-ConvRot-Pruned-Ref2VA.py) | - | - | [code](/examples/minimax_h3/model_training/lora/MiniMax-H3-Int8-ConvRot-Pruned-Ref2VA.sh) | [code](/examples/minimax_h3/model_training/validate_lora/MiniMax-H3-Int8-ConvRot-Pruned-Ref2VA.py) |
+| MiniMax-H3 | [Comfy-Org/MiniMax-H3: FL2VA pruned fp8](https://www.modelscope.cn/models/Comfy-Org/MiniMax-H3) | [code](/examples/minimax_h3/model_inference/MiniMax-H3-FP8-Pruned-FL2VA.py) | [code](/examples/minimax_h3/model_inference_low_vram/MiniMax-H3-FP8-Pruned-FL2VA.py) | - | - | [code](/examples/minimax_h3/model_training/lora/MiniMax-H3-FP8-Pruned-FL2VA.sh) | [code](/examples/minimax_h3/model_training/validate_lora/MiniMax-H3-FP8-Pruned-FL2VA.py) |
+| MiniMax-H3 | [Comfy-Org/MiniMax-H3: Ref2VA pruned fp8](https://www.modelscope.cn/models/Comfy-Org/MiniMax-H3) | [code](/examples/minimax_h3/model_inference/MiniMax-H3-FP8-Pruned-Ref2VA.py) | [code](/examples/minimax_h3/model_inference_low_vram/MiniMax-H3-FP8-Pruned-Ref2VA.py) | - | - | [code](/examples/minimax_h3/model_training/lora/MiniMax-H3-FP8-Pruned-Ref2VA.sh) | [code](/examples/minimax_h3/model_training/validate_lora/MiniMax-H3-FP8-Pruned-Ref2VA.py) |
+| MiniMax-H3 | [lightx2v/Minimax-h3-Turbo: FL2VA 4steps](https://www.modelscope.cn/models/lightx2v/Minimax-h3-Turbo) | [code](https://github.com/modelscope/DiffSynth-Studio/blob/main/examples/minimax_h3/model_inference/MiniMax-H3-FL2VA-Turbo.py) | [code](https://github.com/modelscope/DiffSynth-Studio/blob/main/examples/minimax_h3/model_inference_low_vram/MiniMax-H3-FL2VA-Turbo.py) | - | - | - | - |
+| MiniMax-H3 | [DiffSynth-Studio/MiniMax-H3-Text-Embeddings](https://www.modelscope.cn/models/DiffSynth-Studio/MiniMax-H3-Text-Embeddings) | [code](/examples/minimax_h3/model_inference/MiniMax-H3-Text-Embeddings.py) | [code](/examples/minimax_h3/model_inference_low_vram/MiniMax-H3-Text-Embeddings.py) | [code](/examples/minimax_h3/model_training/full/MiniMax-H3-Text-Embeddings.sh) | [code](/examples/minimax_h3/model_training/validate_full/MiniMax-H3-Text-Embeddings.py) | - | - |
+| LingBot-Video | [Robbyant/lingbot-video-dense-1.3b: T2V](https://modelscope.cn/models/Robbyant/lingbot-video-dense-1.3b) | [code](/examples/lingbot_video/model_inference/lingbot-video-dense-1.3b_t2v.py) | [code](/examples/lingbot_video/model_inference_low_vram/lingbot-video-dense-1.3b_t2v.py) | [code](/examples/lingbot_video/model_training/full/lingbot-video-dense-1.3b_t2v.sh) | [code](/examples/lingbot_video/model_training/validate_full/lingbot-video-dense-1.3b_t2v.py) | [code](/examples/lingbot_video/model_training/lora/lingbot-video-dense-1.3b_t2v.sh) | [code](/examples/lingbot_video/model_training/validate_lora/lingbot-video-dense-1.3b_t2v.py) |
+| LingBot-Video | [Robbyant/lingbot-video-dense-1.3b: TI2V](https://modelscope.cn/models/Robbyant/lingbot-video-dense-1.3b) | [code](/examples/lingbot_video/model_inference/lingbot-video-dense-1.3b_ti2v.py) | [code](/examples/lingbot_video/model_inference_low_vram/lingbot-video-dense-1.3b_ti2v.py) | [code](/examples/lingbot_video/model_training/full/lingbot-video-dense-1.3b_ti2v.sh) | [code](/examples/lingbot_video/model_training/validate_full/lingbot-video-dense-1.3b_ti2v.py) | [code](/examples/lingbot_video/model_training/lora/lingbot-video-dense-1.3b_ti2v.sh) | [code](/examples/lingbot_video/model_training/validate_lora/lingbot-video-dense-1.3b_ti2v.py) |
+| LingBot-Video | [Robbyant/lingbot-video-dense-1.3b: T2I](https://modelscope.cn/models/Robbyant/lingbot-video-dense-1.3b) | [code](/examples/lingbot_video/model_inference/lingbot-video-dense-1.3b_t2i.py) | [code](/examples/lingbot_video/model_inference_low_vram/lingbot-video-dense-1.3b_t2i.py) | - | - | - | - |
+| LingBot-Video | [Robbyant/lingbot-video-moe-30b-a3b: T2V](https://modelscope.cn/models/Robbyant/lingbot-video-moe-30b-a3b) | [code](/examples/lingbot_video/model_inference/lingbot-video-moe-30b-a3b_t2v.py) | [code](/examples/lingbot_video/model_inference_low_vram/lingbot-video-moe-30b-a3b_t2v.py) | [code](/examples/lingbot_video/model_training/full/lingbot-video-moe-30b-a3b_t2v.sh) | [code](/examples/lingbot_video/model_training/validate_full/lingbot-video-moe-30b-a3b_t2v.py) | [code](/examples/lingbot_video/model_training/lora/lingbot-video-moe-30b-a3b_t2v.sh) | [code](/examples/lingbot_video/model_training/validate_lora/lingbot-video-moe-30b-a3b_t2v.py) |
+| LingBot-Video | [Robbyant/lingbot-video-moe-30b-a3b: TI2V](https://modelscope.cn/models/Robbyant/lingbot-video-moe-30b-a3b) | [code](/examples/lingbot_video/model_inference/lingbot-video-moe-30b-a3b_ti2v.py) | [code](/examples/lingbot_video/model_inference_low_vram/lingbot-video-moe-30b-a3b_ti2v.py) | [code](/examples/lingbot_video/model_training/full/lingbot-video-moe-30b-a3b_ti2v.sh) | [code](/examples/lingbot_video/model_training/validate_full/lingbot-video-moe-30b-a3b_ti2v.py) | [code](/examples/lingbot_video/model_training/lora/lingbot-video-moe-30b-a3b_ti2v.sh) | [code](/examples/lingbot_video/model_training/validate_lora/lingbot-video-moe-30b-a3b_ti2v.py) |
+| LingBot-Video | [Robbyant/lingbot-video-moe-30b-a3b: T2I](https://modelscope.cn/models/Robbyant/lingbot-video-moe-30b-a3b) | [code](/examples/lingbot_video/model_inference/lingbot-video-moe-30b-a3b_t2i.py) | [code](/examples/lingbot_video/model_inference_low_vram/lingbot-video-moe-30b-a3b_t2i.py) | - | - | - | - |
+| LingBot-Video | [Robbyant/lingbot-video-moe-30b-a3b: T2V + Refinement](https://modelscope.cn/models/Robbyant/lingbot-video-moe-30b-a3b) | [code](/examples/lingbot_video/model_inference/lingbot-video-moe-30b-a3b_t2v_refiner.py) | [code](/examples/lingbot_video/model_inference_low_vram/lingbot-video-moe-30b-a3b_t2v_refiner.py) | - | - | - | - |
+| LingBot-Video | [Robbyant/lingbot-video-moe-30b-a3b: TI2V + Refinement](https://modelscope.cn/models/Robbyant/lingbot-video-moe-30b-a3b) | [code](/examples/lingbot_video/model_inference/lingbot-video-moe-30b-a3b_ti2v_refiner.py) | [code](/examples/lingbot_video/model_inference_low_vram/lingbot-video-moe-30b-a3b_ti2v_refiner.py) | - | - | - | - |
+| ACE-Step | [ACE-Step/Ace-Step1.5](https://www.modelscope.cn/models/ACE-Step/Ace-Step1.5) | [code](/examples/ace_step/model_inference/Ace-Step1.5.py) | [code](/examples/ace_step/model_inference_low_vram/Ace-Step1.5.py) | [code](/examples/ace_step/model_training/full/Ace-Step1.5.sh) | [code](/examples/ace_step/model_training/validate_full/Ace-Step1.5.py) | [code](/examples/ace_step/model_training/lora/Ace-Step1.5.sh) | [code](/examples/ace_step/model_training/validate_lora/Ace-Step1.5.py) |
+| ACE-Step | [ACE-Step/acestep-v15-turbo-shift1](https://www.modelscope.cn/models/ACE-Step/acestep-v15-turbo-shift1) | [code](/examples/ace_step/model_inference/acestep-v15-turbo-shift1.py) | [code](/examples/ace_step/model_inference_low_vram/acestep-v15-turbo-shift1.py) | [code](/examples/ace_step/model_training/full/acestep-v15-turbo-shift1.sh) | [code](/examples/ace_step/model_training/validate_full/acestep-v15-turbo-shift1.py) | [code](/examples/ace_step/model_training/lora/acestep-v15-turbo-shift1.sh) | [code](/examples/ace_step/model_training/validate_lora/acestep-v15-turbo-shift1.py) |
+| ACE-Step | [ACE-Step/acestep-v15-turbo-shift3](https://www.modelscope.cn/models/ACE-Step/acestep-v15-turbo-shift3) | [code](/examples/ace_step/model_inference/acestep-v15-turbo-shift3.py) | [code](/examples/ace_step/model_inference_low_vram/acestep-v15-turbo-shift3.py) | [code](/examples/ace_step/model_training/full/acestep-v15-turbo-shift3.sh) | [code](/examples/ace_step/model_training/validate_full/acestep-v15-turbo-shift3.py) | [code](/examples/ace_step/model_training/lora/acestep-v15-turbo-shift3.sh) | [code](/examples/ace_step/model_training/validate_lora/acestep-v15-turbo-shift3.py) |
+| ACE-Step | [ACE-Step/acestep-v15-turbo-continuous](https://www.modelscope.cn/models/ACE-Step/acestep-v15-turbo-continuous) | [code](/examples/ace_step/model_inference/acestep-v15-turbo-continuous.py) | [code](/examples/ace_step/model_inference_low_vram/acestep-v15-turbo-continuous.py) | [code](/examples/ace_step/model_training/full/acestep-v15-turbo-continuous.sh) | [code](/examples/ace_step/model_training/validate_full/acestep-v15-turbo-continuous.py) | [code](/examples/ace_step/model_training/lora/acestep-v15-turbo-continuous.sh) | [code](/examples/ace_step/model_training/validate_lora/acestep-v15-turbo-continuous.py) |
+| ACE-Step | [ACE-Step/acestep-v15-base](https://www.modelscope.cn/models/ACE-Step/acestep-v15-base) | [code](/examples/ace_step/model_inference/acestep-v15-base.py) | [code](/examples/ace_step/model_inference_low_vram/acestep-v15-base.py) | [code](/examples/ace_step/model_training/full/acestep-v15-base.sh) | [code](/examples/ace_step/model_training/validate_full/acestep-v15-base.py) | [code](/examples/ace_step/model_training/lora/acestep-v15-base.sh) | [code](/examples/ace_step/model_training/validate_lora/acestep-v15-base.py) |
+| ACE-Step | [ACE-Step/acestep-v15-base: CoverTask](https://www.modelscope.cn/models/ACE-Step/acestep-v15-base) | [code](/examples/ace_step/model_inference/acestep-v15-base-CoverTask.py) | [code](/examples/ace_step/model_inference_low_vram/acestep-v15-base-CoverTask.py) | — | — | — | — |
+| ACE-Step | [ACE-Step/acestep-v15-base: RepaintTask](https://www.modelscope.cn/models/ACE-Step/acestep-v15-base) | [code](/examples/ace_step/model_inference/acestep-v15-base-RepaintTask.py) | [code](/examples/ace_step/model_inference_low_vram/acestep-v15-base-RepaintTask.py) | — | — | — | — |
+| ACE-Step | [ACE-Step/acestep-v15-sft](https://www.modelscope.cn/models/ACE-Step/acestep-v15-sft) | [code](/examples/ace_step/model_inference/acestep-v15-sft.py) | [code](/examples/ace_step/model_inference_low_vram/acestep-v15-sft.py) | [code](/examples/ace_step/model_training/full/acestep-v15-sft.sh) | [code](/examples/ace_step/model_training/validate_full/acestep-v15-sft.py) | [code](/examples/ace_step/model_training/lora/acestep-v15-sft.sh) | [code](/examples/ace_step/model_training/validate_lora/acestep-v15-sft.py) |
+| ACE-Step | [ACE-Step/acestep-v15-xl-base](https://www.modelscope.cn/models/ACE-Step/acestep-v15-xl-base) | [code](/examples/ace_step/model_inference/acestep-v15-xl-base.py) | [code](/examples/ace_step/model_inference_low_vram/acestep-v15-xl-base.py) | [code](/examples/ace_step/model_training/full/acestep-v15-xl-base.sh) | [code](/examples/ace_step/model_training/validate_full/acestep-v15-xl-base.py) | [code](/examples/ace_step/model_training/lora/acestep-v15-xl-base.sh) | [code](/examples/ace_step/model_training/validate_lora/acestep-v15-xl-base.py) |
+| ACE-Step | [ACE-Step/acestep-v15-xl-sft](https://www.modelscope.cn/models/ACE-Step/acestep-v15-xl-sft) | [code](/examples/ace_step/model_inference/acestep-v15-xl-sft.py) | [code](/examples/ace_step/model_inference_low_vram/acestep-v15-xl-sft.py) | [code](/examples/ace_step/model_training/full/acestep-v15-xl-sft.sh) | [code](/examples/ace_step/model_training/validate_full/acestep-v15-xl-sft.py) | [code](/examples/ace_step/model_training/lora/acestep-v15-xl-sft.sh) | [code](/examples/ace_step/model_training/validate_lora/acestep-v15-xl-sft.py) |
+| ACE-Step | [ACE-Step/acestep-v15-xl-turbo](https://www.modelscope.cn/models/ACE-Step/acestep-v15-xl-turbo) | [code](/examples/ace_step/model_inference/acestep-v15-xl-turbo.py) | [code](/examples/ace_step/model_inference_low_vram/acestep-v15-xl-turbo.py) | [code](/examples/ace_step/model_training/full/acestep-v15-xl-turbo.sh) | [code](/examples/ace_step/model_training/validate_full/acestep-v15-xl-turbo.py) | [code](/examples/ace_step/model_training/lora/acestep-v15-xl-turbo.sh) | [code](/examples/ace_step/model_training/validate_lora/acestep-v15-xl-turbo.py) |
+| ACE-Step | [DiffSynth-Studio/acestep15xlsft-lora-music](https://www.modelscope.cn/models/DiffSynth-Studio/acestep15xlsft-lora-music) | [code](/examples/ace_step/model_inference/acestep15xlsft-vocals2music.py) | [code](/examples/ace_step/model_inference_low_vram/acestep15xlsft-vocals2music.py) | [code](/examples/ace_step/model_training/full/acestep15xlsft-vocals2music.sh) | [code](/examples/ace_step/model_training/validate_full/acestep15xlsft-vocals2music.py) | - | - |
+| Boogu-Image | [Boogu/Boogu-Image-0.1-Base](https://modelscope.cn/models/Boogu/Boogu-Image-0.1-Base) | [code](/examples/boogu_image/model_inference/Boogu-Image-0.1-Base.py) | [code](/examples/boogu_image/model_inference_low_vram/Boogu-Image-0.1-Base.py) | [code](/examples/boogu_image/model_training/full/Boogu-Image-0.1-Base.sh) | [code](/examples/boogu_image/model_training/validate_full/Boogu-Image-0.1-Base.py) | [code](/examples/boogu_image/model_training/lora/Boogu-Image-0.1-Base.sh) | [code](/examples/boogu_image/model_training/validate_lora/Boogu-Image-0.1-Base.py) |
+| Boogu-Image | [Boogu/Boogu-Image-0.1-Turbo](https://modelscope.cn/models/Boogu/Boogu-Image-0.1-Turbo) | [code](/examples/boogu_image/model_inference/Boogu-Image-0.1-Turbo.py) | [code](/examples/boogu_image/model_inference_low_vram/Boogu-Image-0.1-Turbo.py) | [code](/examples/boogu_image/model_training/full/Boogu-Image-0.1-Turbo.sh) | [code](/examples/boogu_image/model_training/validate_full/Boogu-Image-0.1-Turbo.py) | [code](/examples/boogu_image/model_training/lora/Boogu-Image-0.1-Turbo.sh) | [code](/examples/boogu_image/model_training/validate_lora/Boogu-Image-0.1-Turbo.py) |
+| Boogu-Image | [Boogu/Boogu-Image-0.1-Edit](https://modelscope.cn/models/Boogu/Boogu-Image-0.1-Edit) | [code](/examples/boogu_image/model_inference/Boogu-Image-0.1-Edit.py) | [code](/examples/boogu_image/model_inference_low_vram/Boogu-Image-0.1-Edit.py) | [code](/examples/boogu_image/model_training/full/Boogu-Image-0.1-Edit.sh) | [code](/examples/boogu_image/model_training/validate_full/Boogu-Image-0.1-Edit.py) | [code](/examples/boogu_image/model_training/lora/Boogu-Image-0.1-Edit.sh) | [code](/examples/boogu_image/model_training/validate_lora/Boogu-Image-0.1-Edit.py) |
+| Krea-2 | [krea/Krea-2-Raw](https://www.modelscope.cn/models/krea/Krea-2-Raw) | [code](/examples/krea2/model_inference/Krea-2-Raw.py) | [code](/examples/krea2/model_inference_low_vram/Krea-2-Raw.py) | [code](/examples/krea2/model_training/full/Krea-2-Raw.sh) | [code](/examples/krea2/model_training/validate_full/Krea-2-Raw.py) | [code](/examples/krea2/model_training/lora/Krea-2-Raw.sh) | [code](/examples/krea2/model_training/validate_lora/Krea-2-Raw.py) |
+| Krea-2 | [krea/Krea-2-Turbo](https://www.modelscope.cn/models/krea/Krea-2-Turbo) | [code](/examples/krea2/model_inference/Krea-2-Turbo.py) | [code](/examples/krea2/model_inference_low_vram/Krea-2-Turbo.py) | [code](/examples/krea2/model_training/full/Krea-2-Turbo.sh) | [code](/examples/krea2/model_training/validate_full/Krea-2-Turbo.py) | [code](/examples/krea2/model_training/lora/Krea-2-Turbo.sh) | [code](/examples/krea2/model_training/validate_lora/Krea-2-Turbo.py) |
+| Ideogram 4 | [ideogram-ai/ideogram-4-fp8](https://www.modelscope.cn/models/ideogram-ai/ideogram-4-fp8) | [code](/examples/ideogram4/model_inference/ideogram-4-fp8.py) | - | - | - | - | - |
+| Ideogram 4 | [DiffSynth-Studio/ideogram-4-bf16-repackage](https://www.modelscope.cn/models/DiffSynth-Studio/ideogram-4-bf16-repackage) | [code](/examples/ideogram4/model_inference/ideogram-4-bf16-repackage.py) | [code](/examples/ideogram4/model_inference_low_vram/ideogram-4-bf16-repackage.py) | [code](/examples/ideogram4/model_training/full/Ideogram-4-bf16-repackage.sh) | - | [code](/examples/ideogram4/model_training/lora/Ideogram-4-bf16-repackage.sh) | [code](/examples/ideogram4/model_training/validate_lora/Ideogram-4-bf16-repackage.py) |
+| HiDream-O1-Image | [HiDream-ai/HiDream-O1-Image](https://modelscope.cn/models/HiDream-ai/HiDream-O1-Image) | [code](/examples/hidream_o1_image/model_inference/HiDream-O1-Image.py) | [code](/examples/hidream_o1_image/model_inference_low_vram/HiDream-O1-Image.py) | [code](/examples/hidream_o1_image/model_training/full/HiDream-O1-Image.sh) | [code](/examples/hidream_o1_image/model_training/validate_full/HiDream-O1-Image.py) | [code](/examples/hidream_o1_image/model_training/lora/HiDream-O1-Image.sh) | [code](/examples/hidream_o1_image/model_training/validate_lora/HiDream-O1-Image.py) |
+| HiDream-O1-Image | [HiDream-ai/HiDream-O1-Image-Dev](https://modelscope.cn/models/HiDream-ai/HiDream-O1-Image-Dev) | [code](/examples/hidream_o1_image/model_inference/HiDream-O1-Image-Dev.py) | [code](/examples/hidream_o1_image/model_inference_low_vram/HiDream-O1-Image-Dev.py) | [code](/examples/hidream_o1_image/model_training/full/HiDream-O1-Image-Dev.sh) | [code](/examples/hidream_o1_image/model_training/validate_full/HiDream-O1-Image-Dev.py) | [code](/examples/hidream_o1_image/model_training/lora/HiDream-O1-Image-Dev.sh) | [code](/examples/hidream_o1_image/model_training/validate_lora/HiDream-O1-Image-Dev.py) |
+| HiDream-O1-Image | [DiffSynth-Studio/HidreamO1-i2L-v2](https://www.modelscope.cn/models/DiffSynth-Studio/HidreamO1-i2L-v2) | [code](/examples/hidream_o1_image/model_inference/HidreamO1-i2L-v2.py) | [code](/examples/hidream_o1_image/model_inference_low_vram/HidreamO1-i2L-v2.py) | [code](/examples/hidream_o1_image/model_training/full/HidreamO1-i2L-v2.sh) | [code](/examples/hidream_o1_image/model_training/validate_full/HidreamO1-i2L-v2.py) | - | - |
+| JoyAI-Image | [jd-opensource/JoyAI-Image-Edit](https://modelscope.cn/models/jd-opensource/JoyAI-Image-Edit) | [code](/examples/joyai_image/model_inference/JoyAI-Image-Edit.py) | [code](/examples/joyai_image/model_inference_low_vram/JoyAI-Image-Edit.py) | [code](/examples/joyai_image/model_training/full/JoyAI-Image-Edit.sh) | [code](/examples/joyai_image/model_training/validate_full/JoyAI-Image-Edit.py) | [code](/examples/joyai_image/model_training/lora/JoyAI-Image-Edit.sh) | [code](/examples/joyai_image/model_training/validate_lora/JoyAI-Image-Edit.py) |
+| ERNIE-Image | [PaddlePaddle/ERNIE-Image](https://www.modelscope.cn/models/PaddlePaddle/ERNIE-Image) | [code](/examples/ernie_image/model_inference/ERNIE-Image.py) | [code](/examples/ernie_image/model_inference_low_vram/ERNIE-Image.py) | [code](/examples/ernie_image/model_training/full/ERNIE-Image.sh) | [code](/examples/ernie_image/model_training/validate_full/ERNIE-Image.py) | [code](/examples/ernie_image/model_training/lora/ERNIE-Image.sh) | [code](/examples/ernie_image/model_training/validate_lora/ERNIE-Image.py) |
+| ERNIE-Image | [PaddlePaddle/ERNIE-Image-Turbo](https://www.modelscope.cn/models/PaddlePaddle/ERNIE-Image-Turbo) | [code](/examples/ernie_image/model_inference/ERNIE-Image-Turbo.py) | [code](/examples/ernie_image/model_inference_low_vram/ERNIE-Image-Turbo.py) | — | — | — | — |
+| LTX-2 | [jd-opensource/JoyAI-Echo](https://modelscope.cn/models/jd-opensource/JoyAI-Echo) | [code](/examples/ltx2/model_inference/JoyAI-Echo-T2AV.py) | [code](/examples/ltx2/model_inference_low_vram/JoyAI-Echo-T2AV.py) | [code](/examples/ltx2/model_training/full/JoyAI-Echo-T2AV-splited.sh) | [code](/examples/ltx2/model_training/validate_full/JoyAI-Echo-T2AV.py) | [code](/examples/ltx2/model_training/lora/JoyAI-Echo-T2AV-splited.sh) | [code](/examples/ltx2/model_training/validate_lora/JoyAI-Echo-T2AV.py) |
+| LTX-2 | [Lightricks/LTX-2.3: OneStagePipeline-I2AV](https://www.modelscope.cn/models/Lightricks/LTX-2.3) | [code](/examples/ltx2/model_inference/LTX-2.3-I2AV-OneStage.py) | [code](/examples/ltx2/model_inference_low_vram/LTX-2.3-I2AV-OneStage.py) | [code](/examples/ltx2/model_training/full/LTX-2.3-I2AV-splited.sh) | [code](/examples/ltx2/model_training/validate_full/LTX-2.3-I2AV.py) | [code](/examples/ltx2/model_training/lora/LTX-2.3-I2AV-splited.sh) | [code](/examples/ltx2/model_training/validate_lora/LTX-2.3-I2AV.py) |
+| LTX-2 | [Lightricks/LTX-2.3: TwoStagePipeline-I2AV](https://www.modelscope.cn/models/Lightricks/LTX-2.3) | [code](/examples/ltx2/model_inference/LTX-2.3-I2AV-TwoStage.py) | [code](/examples/ltx2/model_inference_low_vram/LTX-2.3-I2AV-TwoStage.py) | - | - | - | - |
+| LTX-2 | [Lightricks/LTX-2.3: DistilledPipeline-I2AV](https://www.modelscope.cn/models/Lightricks/LTX-2.3) | [code](/examples/ltx2/model_inference/LTX-2.3-I2AV-DistilledPipeline.py) | [code](/examples/ltx2/model_inference_low_vram/LTX-2.3-I2AV-DistilledPipeline.py) | - | - | - | - |
+| LTX-2 | [Lightricks/LTX-2.3: OneStagePipeline-T2AV](https://www.modelscope.cn/models/Lightricks/LTX-2.3) | [code](/examples/ltx2/model_inference/LTX-2.3-T2AV-OneStage.py) | [code](/examples/ltx2/model_inference_low_vram/LTX-2.3-T2AV-OneStage.py) | [code](/examples/ltx2/model_training/full/LTX-2.3-T2AV-splited.sh) | [code](/examples/ltx2/model_training/validate_full/LTX-2.3-T2AV.py) | [code](/examples/ltx2/model_training/lora/LTX-2.3-T2AV-splited.sh) | [code](/examples/ltx2/model_training/validate_lora/LTX-2.3-T2AV.py) |
+| LTX-2 | [Lightricks/LTX-2.3: TwoStagePipeline-T2AV](https://www.modelscope.cn/models/Lightricks/LTX-2.3) | [code](/examples/ltx2/model_inference/LTX-2.3-T2AV-TwoStage.py) | [code](/examples/ltx2/model_inference_low_vram/LTX-2.3-T2AV-TwoStage.py) | - | - | - | - |
+| LTX-2 | [Lightricks/LTX-2.3: DistilledPipeline-T2AV](https://www.modelscope.cn/models/Lightricks/LTX-2.3) | [code](/examples/ltx2/model_inference/LTX-2.3-T2AV-DistilledPipeline.py) | [code](/examples/ltx2/model_inference_low_vram/LTX-2.3-T2AV-DistilledPipeline.py) | - | - | - | - |
+| LTX-2 | [Lightricks/LTX-2.3: A2V](https://www.modelscope.cn/models/Lightricks/LTX-2.3) | [code](/examples/ltx2/model_inference/LTX-2.3-A2V-TwoStage.py) | [code](/examples/ltx2/model_inference_low_vram/LTX-2.3-A2V-TwoStage.py) | - | - | - | - |
+| LTX-2 | [Lightricks/LTX-2.3: Retake](https://www.modelscope.cn/models/Lightricks/LTX-2.3) | [code](/examples/ltx2/model_inference/LTX-2.3-T2AV-TwoStage-Retake.py) | [code](/examples/ltx2/model_inference_low_vram/LTX-2.3-T2AV-TwoStage-Retake.py) | - | - | - | - |
+| LTX-2 | [Lightricks/LTX-2.3-22b-IC-LoRA-Union-Control](https://www.modelscope.cn/models/Lightricks/LTX-2.3-22b-IC-LoRA-Union-Control) | [code](/examples/ltx2/model_inference/LTX-2.3-T2AV-IC-LoRA-Union-Control.py) | [code](/examples/ltx2/model_inference_low_vram/LTX-2.3-T2AV-IC-LoRA-Union-Control.py) | - | - | [code](/examples/ltx2/model_training/lora/LTX-2.3-T2AV-IC-LoRA-splited.sh) | [code](/examples/ltx2/model_training/validate_lora/LTX-2.3-T2AV-IC-LoRA.py) |
+| LTX-2 | [Lightricks/LTX-2.3-22b-IC-LoRA-Motion-Track-Control](https://www.modelscope.cn/models/Lightricks/LTX-2.3-22b-IC-LoRA-Motion-Track-Control) | [code](/examples/ltx2/model_inference/LTX-2.3-T2AV-IC-LoRA-Motion-Track-Control.py) | [code](/examples/ltx2/model_inference_low_vram/LTX-2.3-T2AV-IC-LoRA-Motion-Track-Control.py) | - | - | [code](/examples/ltx2/model_training/lora/LTX-2.3-T2AV-IC-LoRA-splited.sh) | [code](/examples/ltx2/model_training/validate_lora/LTX-2.3-T2AV-IC-LoRA.py) |
+| LTX-2 | [Lightricks/LTX-2: OneStagePipeline-T2AV](https://www.modelscope.cn/models/Lightricks/LTX-2) | [code](/examples/ltx2/model_inference/LTX-2-T2AV-OneStage.py) | [code](/examples/ltx2/model_inference_low_vram/LTX-2-T2AV-OneStage.py) | [code](/examples/ltx2/model_training/full/LTX-2-T2AV-splited.sh) | [code](/examples/ltx2/model_training/validate_full/LTX-2-T2AV.py) | [code](/examples/ltx2/model_training/lora/LTX-2-T2AV-splited.sh) | [code](/examples/ltx2/model_training/validate_lora/LTX-2-T2AV.py) |
+| LTX-2 | [Lightricks/LTX-2-19b-IC-LoRA-Union-Control](https://www.modelscope.cn/models/Lightricks/LTX-2-19b-IC-LoRA-Union-Control) | [code](/examples/ltx2/model_inference/LTX-2-T2AV-IC-LoRA-Union-Control.py) | [code](/examples/ltx2/model_inference_low_vram/LTX-2-T2AV-IC-LoRA-Union-Control.py) | - | - | [code](/examples/ltx2/model_training/lora/LTX-2-T2AV-IC-LoRA-splited.sh) | [code](/examples/ltx2/model_training/validate_lora/LTX-2-T2AV-IC-LoRA.py) |
+| LTX-2 | [Lightricks/LTX-2-19b-IC-LoRA-Detailer](https://www.modelscope.cn/models/Lightricks/LTX-2-19b-IC-LoRA-Detailer) | [code](/examples/ltx2/model_inference/LTX-2-T2AV-IC-LoRA-Detailer.py) | [code](/examples/ltx2/model_inference_low_vram/LTX-2-T2AV-IC-LoRA-Detailer.py) | - | - | [code](/examples/ltx2/model_training/lora/LTX-2-T2AV-IC-LoRA-splited.sh) | [code](/examples/ltx2/model_training/validate_lora/LTX-2-T2AV-IC-LoRA.py) |
+| LTX-2 | [Lightricks/LTX-2: TwoStagePipeline-T2AV](https://www.modelscope.cn/models/Lightricks/LTX-2) | [code](/examples/ltx2/model_inference/LTX-2-T2AV-TwoStage.py) | [code](/examples/ltx2/model_inference_low_vram/LTX-2-T2AV-TwoStage.py) | - | - | - | - |
+| LTX-2 | [Lightricks/LTX-2: DistilledPipeline-T2AV](https://www.modelscope.cn/models/Lightricks/LTX-2) | [code](/examples/ltx2/model_inference/LTX-2-T2AV-DistilledPipeline.py) | [code](/examples/ltx2/model_inference_low_vram/LTX-2-T2AV-DistilledPipeline.py) | - | - | - | - |
+| LTX-2 | [Lightricks/LTX-2: OneStagePipeline-I2AV](https://www.modelscope.cn/models/Lightricks/LTX-2) | [code](/examples/ltx2/model_inference/LTX-2-I2AV-OneStage.py) | [code](/examples/ltx2/model_inference_low_vram/LTX-2-I2AV-OneStage.py) | - | - | - | - |
+| LTX-2 | [Lightricks/LTX-2: TwoStagePipeline-I2AV](https://www.modelscope.cn/models/Lightricks/LTX-2) | [code](/examples/ltx2/model_inference/LTX-2-I2AV-TwoStage.py) | [code](/examples/ltx2/model_inference_low_vram/LTX-2-I2AV-TwoStage.py) | - | - | - | - |
+| LTX-2 | [Lightricks/LTX-2: DistilledPipeline-I2AV](https://www.modelscope.cn/models/Lightricks/LTX-2) | [code](/examples/ltx2/model_inference/LTX-2-I2AV-DistilledPipeline.py) | [code](/examples/ltx2/model_inference_low_vram/LTX-2-I2AV-DistilledPipeline.py) | - | - | - | - |
+| LTX-2 | [Lightricks/LTX-2-19b-LoRA-Camera-Control-Dolly-In](https://www.modelscope.cn/models/Lightricks/LTX-2-19b-LoRA-Camera-Control-Dolly-In) | [code](/examples/ltx2/model_inference/LTX-2-T2AV-Camera-Control-Dolly-In.py) | [code](/examples/ltx2/model_inference_low_vram/LTX-2-T2AV-Camera-Control-Dolly-In.py) | - | - | - | - |
+| LTX-2 | [Lightricks/LTX-2-19b-LoRA-Camera-Control-Dolly-Out](https://www.modelscope.cn/models/Lightricks/LTX-2-19b-LoRA-Camera-Control-Dolly-Out) | [code](/examples/ltx2/model_inference/LTX-2-T2AV-Camera-Control-Dolly-Out.py) | [code](/examples/ltx2/model_inference_low_vram/LTX-2-T2AV-Camera-Control-Dolly-Out.py) | - | - | - | - |
+| LTX-2 | [Lightricks/LTX-2-19b-LoRA-Camera-Control-Dolly-Left](https://www.modelscope.cn/models/Lightricks/LTX-2-19b-LoRA-Camera-Control-Dolly-Left) | [code](/examples/ltx2/model_inference/LTX-2-T2AV-Camera-Control-Dolly-Left.py) | [code](/examples/ltx2/model_inference_low_vram/LTX-2-T2AV-Camera-Control-Dolly-Left.py) | - | - | - | - |
+| LTX-2 | [Lightricks/LTX-2-19b-LoRA-Camera-Control-Dolly-Right](https://www.modelscope.cn/models/Lightricks/LTX-2-19b-LoRA-Camera-Control-Dolly-Right) | [code](/examples/ltx2/model_inference/LTX-2-T2AV-Camera-Control-Dolly-Right.py) | [code](/examples/ltx2/model_inference_low_vram/LTX-2-T2AV-Camera-Control-Dolly-Right.py) | - | - | - | - |
+| LTX-2 | [Lightricks/LTX-2-19b-LoRA-Camera-Control-Jib-Up](https://www.modelscope.cn/models/Lightricks/LTX-2-19b-LoRA-Camera-Control-Jib-Up) | [code](/examples/ltx2/model_inference/LTX-2-T2AV-Camera-Control-Jib-Up.py) | [code](/examples/ltx2/model_inference_low_vram/LTX-2-T2AV-Camera-Control-Jib-Up.py) | - | - | - | - |
+| LTX-2 | [Lightricks/LTX-2-19b-LoRA-Camera-Control-Jib-Down](https://www.modelscope.cn/models/Lightricks/LTX-2-19b-LoRA-Camera-Control-Jib-Down) | [code](/examples/ltx2/model_inference/LTX-2-T2AV-Camera-Control-Jib-Down.py) | [code](/examples/ltx2/model_inference_low_vram/LTX-2-T2AV-Camera-Control-Jib-Down.py) | - | - | - | - |
+| LTX-2 | [Lightricks/LTX-2-19b-LoRA-Camera-Control-Static](https://www.modelscope.cn/models/Lightricks/LTX-2-19b-LoRA-Camera-Control-Static) | [code](/examples/ltx2/model_inference/LTX-2-T2AV-Camera-Control-Static.py) | [code](/examples/ltx2/model_inference_low_vram/LTX-2-T2AV-Camera-Control-Static.py) | - | - | - | - |
+| FLUX.2 | [black-forest-labs/FLUX.2-dev](https://www.modelscope.cn/models/black-forest-labs/FLUX.2-dev) | [code](/examples/flux2/model_inference/FLUX.2-dev.py) | [code](/examples/flux2/model_inference_low_vram/FLUX.2-dev.py) | - | - | [code](/examples/flux2/model_training/lora/FLUX.2-dev.sh) | [code](/examples/flux2/model_training/validate_lora/FLUX.2-dev.py) |
+| FLUX.2 | [black-forest-labs/FLUX.2-klein-4B](https://www.modelscope.cn/models/black-forest-labs/FLUX.2-klein-4B) | [code](/examples/flux2/model_inference/FLUX.2-klein-4B.py) | [code](/examples/flux2/model_inference_low_vram/FLUX.2-klein-4B.py) | [code](/examples/flux2/model_training/full/FLUX.2-klein-4B.sh) | [code](/examples/flux2/model_training/validate_full/FLUX.2-klein-4B.py) | [code](/examples/flux2/model_training/lora/FLUX.2-klein-4B.sh) | [code](/examples/flux2/model_training/validate_lora/FLUX.2-klein-4B.py) |
+| FLUX.2 | [black-forest-labs/FLUX.2-klein-9B](https://www.modelscope.cn/models/black-forest-labs/FLUX.2-klein-9B) | [code](/examples/flux2/model_inference/FLUX.2-klein-9B.py) | [code](/examples/flux2/model_inference_low_vram/FLUX.2-klein-9B.py) | [code](/examples/flux2/model_training/full/FLUX.2-klein-9B.sh) | [code](/examples/flux2/model_training/validate_full/FLUX.2-klein-9B.py) | [code](/examples/flux2/model_training/lora/FLUX.2-klein-9B.sh) | [code](/examples/flux2/model_training/validate_lora/FLUX.2-klein-9B.py) |
+| FLUX.2 | [black-forest-labs/FLUX.2-klein-base-4B](https://www.modelscope.cn/models/black-forest-labs/FLUX.2-klein-base-4B) | [code](/examples/flux2/model_inference/FLUX.2-klein-base-4B.py) | [code](/examples/flux2/model_inference_low_vram/FLUX.2-klein-base-4B.py) | [code](/examples/flux2/model_training/full/FLUX.2-klein-base-4B.sh) | [code](/examples/flux2/model_training/validate_full/FLUX.2-klein-base-4B.py) | [code](/examples/flux2/model_training/lora/FLUX.2-klein-base-4B.sh) | [code](/examples/flux2/model_training/validate_lora/FLUX.2-klein-base-4B.py) |
+| FLUX.2 | [black-forest-labs/FLUX.2-klein-base-9B](https://www.modelscope.cn/models/black-forest-labs/FLUX.2-klein-base-9B) | [code](/examples/flux2/model_inference/FLUX.2-klein-base-9B.py) | [code](/examples/flux2/model_inference_low_vram/FLUX.2-klein-base-9B.py) | [code](/examples/flux2/model_training/full/FLUX.2-klein-base-9B.sh) | [code](/examples/flux2/model_training/validate_full/FLUX.2-klein-base-9B.py) | [code](/examples/flux2/model_training/lora/FLUX.2-klein-base-9B.sh) | [code](/examples/flux2/model_training/validate_lora/FLUX.2-klein-base-9B.py) |
+| FLUX.2 | [DiffSynth-Studio/Template-KleinBase4B-Aesthetic](https://www.modelscope.cn/models/DiffSynth-Studio/Template-KleinBase4B-Aesthetic) | [code](/examples/flux2/model_inference/Template-KleinBase4B-Aesthetic.py) | [code](/examples/flux2/model_inference_low_vram/Template-KleinBase4B-Aesthetic.py) | [code](/examples/flux2/model_training/full/Template-KleinBase4B-Aesthetic.sh) | [code](/examples/flux2/model_training/validate_full/Template-KleinBase4B-Aesthetic.py) | - | - |
+| FLUX.2 | [DiffSynth-Studio/Template-KleinBase4B-Brightness](https://www.modelscope.cn/models/DiffSynth-Studio/Template-KleinBase4B-Brightness) | [code](/examples/flux2/model_inference/Template-KleinBase4B-Brightness.py) | [code](/examples/flux2/model_inference_low_vram/Template-KleinBase4B-Brightness.py) | [code](/examples/flux2/model_training/full/Template-KleinBase4B-Brightness.sh) | [code](/examples/flux2/model_training/validate_full/Template-KleinBase4B-Brightness.py) | - | - |
+| FLUX.2 | [DiffSynth-Studio/Template-KleinBase4B-Age](https://www.modelscope.cn/models/DiffSynth-Studio/Template-KleinBase4B-Age) | [code](/examples/flux2/model_inference/Template-KleinBase4B-Age.py) | [code](/examples/flux2/model_inference_low_vram/Template-KleinBase4B-Age.py) | [code](/examples/flux2/model_training/full/Template-KleinBase4B-Age.sh) | [code](/examples/flux2/model_training/validate_full/Template-KleinBase4B-Age.py) | - | - |
+| FLUX.2 | [DiffSynth-Studio/Template-KleinBase4B-ControlNet](https://www.modelscope.cn/models/DiffSynth-Studio/Template-KleinBase4B-ControlNet) | [code](/examples/flux2/model_inference/Template-KleinBase4B-ControlNet.py) | [code](/examples/flux2/model_inference_low_vram/Template-KleinBase4B-ControlNet.py) | [code](/examples/flux2/model_training/full/Template-KleinBase4B-ControlNet.sh) | [code](/examples/flux2/model_training/validate_full/Template-KleinBase4B-ControlNet.py) | - | - |
+| FLUX.2 | [DiffSynth-Studio/Template-KleinBase4B-Edit](https://www.modelscope.cn/models/DiffSynth-Studio/Template-KleinBase4B-Edit) | [code](/examples/flux2/model_inference/Template-KleinBase4B-Edit.py) | [code](/examples/flux2/model_inference_low_vram/Template-KleinBase4B-Edit.py) | [code](/examples/flux2/model_training/full/Template-KleinBase4B-Edit.sh) | [code](/examples/flux2/model_training/validate_full/Template-KleinBase4B-Edit.py) | - | - |
+| FLUX.2 | [DiffSynth-Studio/Template-KleinBase4B-Inpaint](https://www.modelscope.cn/models/DiffSynth-Studio/Template-KleinBase4B-Inpaint) | [code](/examples/flux2/model_inference/Template-KleinBase4B-Inpaint.py) | [code](/examples/flux2/model_inference_low_vram/Template-KleinBase4B-Inpaint.py) | [code](/examples/flux2/model_training/full/Template-KleinBase4B-Inpaint.sh) | [code](/examples/flux2/model_training/validate_full/Template-KleinBase4B-Inpaint.py) | - | - |
+| FLUX.2 | [DiffSynth-Studio/Template-KleinBase4B-PandaMeme](https://www.modelscope.cn/models/DiffSynth-Studio/Template-KleinBase4B-PandaMeme) | [code](/examples/flux2/model_inference/Template-KleinBase4B-PandaMeme.py) | [code](/examples/flux2/model_inference_low_vram/Template-KleinBase4B-PandaMeme.py) | [code](/examples/flux2/model_training/full/Template-KleinBase4B-PandaMeme.sh) | [code](/examples/flux2/model_training/validate_full/Template-KleinBase4B-PandaMeme.py) | - | - |
+| FLUX.2 | [DiffSynth-Studio/Template-KleinBase4B-Sharpness](https://www.modelscope.cn/models/DiffSynth-Studio/Template-KleinBase4B-Sharpness) | [code](/examples/flux2/model_inference/Template-KleinBase4B-Sharpness.py) | [code](/examples/flux2/model_inference_low_vram/Template-KleinBase4B-Sharpness.py) | [code](/examples/flux2/model_training/full/Template-KleinBase4B-Sharpness.sh) | [code](/examples/flux2/model_training/validate_full/Template-KleinBase4B-Sharpness.py) | - | - |
+| FLUX.2 | [DiffSynth-Studio/Template-KleinBase4B-SoftRGB](https://www.modelscope.cn/models/DiffSynth-Studio/Template-KleinBase4B-SoftRGB) | [code](/examples/flux2/model_inference/Template-KleinBase4B-SoftRGB.py) | [code](/examples/flux2/model_inference_low_vram/Template-KleinBase4B-SoftRGB.py) | [code](/examples/flux2/model_training/full/Template-KleinBase4B-SoftRGB.sh) | [code](/examples/flux2/model_training/validate_full/Template-KleinBase4B-SoftRGB.py) | - | - |
+| FLUX.2 | [DiffSynth-Studio/Template-KleinBase4B-Upscaler](https://www.modelscope.cn/models/DiffSynth-Studio/Template-KleinBase4B-Upscaler) | [code](/examples/flux2/model_inference/Template-KleinBase4B-Upscaler.py) | [code](/examples/flux2/model_inference_low_vram/Template-KleinBase4B-Upscaler.py) | [code](/examples/flux2/model_training/full/Template-KleinBase4B-Upscaler.sh) | [code](/examples/flux2/model_training/validate_full/Template-KleinBase4B-Upscaler.py) | - | - |
+| FLUX.2 | [DiffSynth-Studio/Template-KleinBase4B-ContentRef](https://www.modelscope.cn/models/DiffSynth-Studio/Template-KleinBase4B-ContentRef) | [code](/examples/flux2/model_inference/Template-KleinBase4B-ContentRef.py) | [code](/examples/flux2/model_inference_low_vram/Template-KleinBase4B-ContentRef.py) | [code](/examples/flux2/model_training/full/Template-KleinBase4B-ContentRef.sh) | [code](/examples/flux2/model_training/validate_full/Template-KleinBase4B-ContentRef.py) | - | - |
+| FLUX.2 | [DiffSynth-Studio/KleinBase4B-i2L-v2](https://www.modelscope.cn/models/DiffSynth-Studio/KleinBase4B-i2L-v2) | [code](/examples/flux2/model_inference/KleinBase4B-i2L-v2.py) | [code](/examples/flux2/model_inference_low_vram/KleinBase4B-i2L-v2.py) | [code](/examples/flux2/model_training/full/KleinBase4B-i2L-v2.sh) | [code](/examples/flux2/model_training/validate_full/KleinBase4B-i2L-v2.py) | - | - |
+| Z-Image | [Tongyi-MAI/Z-Image](https://www.modelscope.cn/models/Tongyi-MAI/Z-Image) | [code](/examples/z_image/model_inference/Z-Image.py) | [code](/examples/z_image/model_inference_low_vram/Z-Image.py) | [code](/examples/z_image/model_training/full/Z-Image.sh) | [code](/examples/z_image/model_training/validate_full/Z-Image.py) | [code](/examples/z_image/model_training/lora/Z-Image.sh) | [code](/examples/z_image/model_training/validate_lora/Z-Image.py) |
+| Z-Image | [DiffSynth-Studio/Z-Image-i2L](https://www.modelscope.cn/models/DiffSynth-Studio/Z-Image-i2L) | [code](/examples/z_image/model_inference/Z-Image-i2L.py) | [code](/examples/z_image/model_inference_low_vram/Z-Image-i2L.py) | - | - | - | - |
+| Z-Image | [Tongyi-MAI/Z-Image-Turbo](https://www.modelscope.cn/models/Tongyi-MAI/Z-Image-Turbo) | [code](/examples/z_image/model_inference/Z-Image-Turbo.py) | [code](/examples/z_image/model_inference_low_vram/Z-Image-Turbo.py) | [code](/examples/z_image/model_training/full/Z-Image-Turbo.sh) | [code](/examples/z_image/model_training/validate_full/Z-Image-Turbo.py) | [code](/examples/z_image/model_training/lora/Z-Image-Turbo.sh) | [code](/examples/z_image/model_training/validate_lora/Z-Image-Turbo.py) |
+| Z-Image | [PAI/Z-Image-Turbo-Fun-Controlnet-Union-2.1](https://www.modelscope.cn/models/PAI/Z-Image-Turbo-Fun-Controlnet-Union-2.1) | [code](/examples/z_image/model_inference/Z-Image-Turbo-Fun-Controlnet-Union-2.1.py) | [code](/examples/z_image/model_inference_low_vram/Z-Image-Turbo-Fun-Controlnet-Union-2.1.py) | [code](/examples/z_image/model_training/full/Z-Image-Turbo-Fun-Controlnet-Union-2.1.sh) | [code](/examples/z_image/model_training/validate_full/Z-Image-Turbo-Fun-Controlnet-Union-2.1.py) | [code](/examples/z_image/model_training/lora/Z-Image-Turbo-Fun-Controlnet-Union-2.1.sh) | [code](/examples/z_image/model_training/validate_lora/Z-Image-Turbo-Fun-Controlnet-Union-2.1.py) |
+| Z-Image | [PAI/Z-Image-Turbo-Fun-Controlnet-Union-2.1-8steps](https://www.modelscope.cn/models/PAI/Z-Image-Turbo-Fun-Controlnet-Union-2.1) | [code](/examples/z_image/model_inference/Z-Image-Turbo-Fun-Controlnet-Union-2.1-8steps.py) | [code](/examples/z_image/model_inference_low_vram/Z-Image-Turbo-Fun-Controlnet-Union-2.1-8steps.py) | [code](/examples/z_image/model_training/full/Z-Image-Turbo-Fun-Controlnet-Union-2.1-8steps.sh) | [code](/examples/z_image/model_training/validate_full/Z-Image-Turbo-Fun-Controlnet-Union-2.1-8steps.py) | [code](/examples/z_image/model_training/lora/Z-Image-Turbo-Fun-Controlnet-Union-2.1-8steps.sh) | [code](/examples/z_image/model_training/validate_lora/Z-Image-Turbo-Fun-Controlnet-Union-2.1-8steps.py) |
+| Z-Image | [PAI/Z-Image-Turbo-Fun-Controlnet-Tile-2.1-8steps](https://www.modelscope.cn/models/PAI/Z-Image-Turbo-Fun-Controlnet-Union-2.1) | [code](/examples/z_image/model_inference/Z-Image-Turbo-Fun-Controlnet-Tile-2.1-8steps.py) | [code](/examples/z_image/model_inference_low_vram/Z-Image-Turbo-Fun-Controlnet-Tile-2.1-8steps.py) | [code](/examples/z_image/model_training/full/Z-Image-Turbo-Fun-Controlnet-Tile-2.1-8steps.sh) | [code](/examples/z_image/model_training/validate_full/Z-Image-Turbo-Fun-Controlnet-Tile-2.1-8steps.py) | [code](/examples/z_image/model_training/lora/Z-Image-Turbo-Fun-Controlnet-Tile-2.1-8steps.sh) | [code](/examples/z_image/model_training/validate_lora/Z-Image-Turbo-Fun-Controlnet-Tile-2.1-8steps.py) |
+| Z-Image | [DiffSynth-Studio/ZImage-i2L-v2](https://www.modelscope.cn/models/DiffSynth-Studio/ZImage-i2L-v2) | [code](/examples/z_image/model_inference/ZImage-i2L-v2.py) | [code](/examples/z_image/model_inference_low_vram/ZImage-i2L-v2.py) | [code](/examples/z_image/model_training/full/ZImage-i2L-v2.sh) | [code](/examples/z_image/model_training/validate_full/ZImage-i2L-v2.py) | - | - |
+| Anima | [circlestone-labs/Anima](https://www.modelscope.cn/models/circlestone-labs/Anima) | [code](/examples/anima/model_inference/anima-preview.py) | [code](/examples/anima/model_inference_low_vram/anima-preview.py) | [code](/examples/anima/model_training/full/anima-preview.sh) | [code](/examples/anima/model_training/validate_full/anima-preview.py) | [code](/examples/anima/model_training/lora/anima-preview.sh) | [code](/examples/anima/model_training/validate_lora/anima-preview.py) |
+| Qwen-Image | [Qwen/Qwen-Image](https://www.modelscope.cn/models/Qwen/Qwen-Image) | [code](/examples/qwen_image/model_inference/Qwen-Image.py) | [code](/examples/qwen_image/model_inference_low_vram/Qwen-Image.py) | [code](/examples/qwen_image/model_training/full/Qwen-Image.sh) | [code](/examples/qwen_image/model_training/validate_full/Qwen-Image.py) | [code](/examples/qwen_image/model_training/lora/Qwen-Image.sh) | [code](/examples/qwen_image/model_training/validate_lora/Qwen-Image.py) |
+| Qwen-Image | [Qwen/Qwen-Image-2512](https://www.modelscope.cn/models/Qwen/Qwen-Image-2512) | [code](/examples/qwen_image/model_inference/Qwen-Image-2512.py) | [code](/examples/qwen_image/model_inference_low_vram/Qwen-Image-2512.py) | [code](/examples/qwen_image/model_training/full/Qwen-Image-2512.sh) | [code](/examples/qwen_image/model_training/validate_full/Qwen-Image-2512.py) | [code](/examples/qwen_image/model_training/lora/Qwen-Image-2512.sh) | [code](/examples/qwen_image/model_training/validate_lora/Qwen-Image-2512.py) |
+| Qwen-Image | [Qwen/Qwen-Image-Edit](https://www.modelscope.cn/models/Qwen/Qwen-Image-Edit) | [code](/examples/qwen_image/model_inference/Qwen-Image-Edit.py) | [code](/examples/qwen_image/model_inference_low_vram/Qwen-Image-Edit.py) | [code](/examples/qwen_image/model_training/full/Qwen-Image-Edit.sh) | [code](/examples/qwen_image/model_training/validate_full/Qwen-Image-Edit.py) | [code](/examples/qwen_image/model_training/lora/Qwen-Image-Edit.sh) | [code](/examples/qwen_image/model_training/validate_lora/Qwen-Image-Edit.py) |
+| Qwen-Image | [Qwen/Qwen-Image-Edit-2509](https://www.modelscope.cn/models/Qwen/Qwen-Image-Edit-2509) | [code](/examples/qwen_image/model_inference/Qwen-Image-Edit-2509.py) | [code](/examples/qwen_image/model_inference_low_vram/Qwen-Image-Edit-2509.py) | [code](/examples/qwen_image/model_training/full/Qwen-Image-Edit-2509.sh) | [code](/examples/qwen_image/model_training/validate_full/Qwen-Image-Edit-2509.py) | [code](/examples/qwen_image/model_training/lora/Qwen-Image-Edit-2509.sh) | [code](/examples/qwen_image/model_training/validate_lora/Qwen-Image-Edit-2509.py) |
+| Qwen-Image | [Qwen/Qwen-Image-Edit-2511](https://www.modelscope.cn/models/Qwen/Qwen-Image-Edit-2511) | [code](/examples/qwen_image/model_inference/Qwen-Image-Edit-2511.py) | [code](/examples/qwen_image/model_inference_low_vram/Qwen-Image-Edit-2511.py) | [code](/examples/qwen_image/model_training/full/Qwen-Image-Edit-2511.sh) | [code](/examples/qwen_image/model_training/validate_full/Qwen-Image-Edit-2511.py) | [code](/examples/qwen_image/model_training/lora/Qwen-Image-Edit-2511.sh) | [code](/examples/qwen_image/model_training/validate_lora/Qwen-Image-Edit-2511.py) |
+| Qwen-Image | [FireRedTeam/FireRed-Image-Edit-1.0](https://www.modelscope.cn/models/FireRedTeam/FireRed-Image-Edit-1.0) | [code](/examples/qwen_image/model_inference/FireRed-Image-Edit-1.0.py) | [code](/examples/qwen_image/model_inference_low_vram/FireRed-Image-Edit-1.0.py) | [code](/examples/qwen_image/model_training/full/FireRed-Image-Edit-1.0.sh) | [code](/examples/qwen_image/model_training/validate_full/FireRed-Image-Edit-1.0.py) | [code](/examples/qwen_image/model_training/lora/FireRed-Image-Edit-1.0.sh) | [code](/examples/qwen_image/model_training/validate_lora/FireRed-Image-Edit-1.0.py) |
+| Qwen-Image | [FireRedTeam/FireRed-Image-Edit-1.1](https://www.modelscope.cn/models/FireRedTeam/FireRed-Image-Edit-1.1) | [code](/examples/qwen_image/model_inference/FireRed-Image-Edit-1.1.py) | [code](/examples/qwen_image/model_inference_low_vram/FireRed-Image-Edit-1.1.py) | [code](/examples/qwen_image/model_training/full/FireRed-Image-Edit-1.1.sh) | [code](/examples/qwen_image/model_training/validate_full/FireRed-Image-Edit-1.1.py) | [code](/examples/qwen_image/model_training/lora/FireRed-Image-Edit-1.1.sh) | [code](/examples/qwen_image/model_training/validate_lora/FireRed-Image-Edit-1.1.py) |
+| Qwen-Image | [lightx2v/Qwen-Image-Edit-2511-Lightning](https://modelscope.cn/models/lightx2v/Qwen-Image-Edit-2511-Lightning) | [code](/examples/qwen_image/model_inference/Qwen-Image-Edit-2511-Lightning.py) | [code](/examples/qwen_image/model_inference_low_vram/Qwen-Image-Edit-2511-Lightning.py) | - | - | - | - |
+| Qwen-Image | [Qwen/Qwen-Image-Layered](https://www.modelscope.cn/models/Qwen/Qwen-Image-Layered) | [code](/examples/qwen_image/model_inference/Qwen-Image-Layered.py) | [code](/examples/qwen_image/model_inference_low_vram/Qwen-Image-Layered.py) | [code](/examples/qwen_image/model_training/full/Qwen-Image-Layered.sh) | [code](/examples/qwen_image/model_training/validate_full/Qwen-Image-Layered.py) | [code](/examples/qwen_image/model_training/lora/Qwen-Image-Layered.sh) | [code](/examples/qwen_image/model_training/validate_lora/Qwen-Image-Layered.py) |
+| Qwen-Image | [DiffSynth-Studio/Qwen-Image-Layered-Control](https://www.modelscope.cn/models/DiffSynth-Studio/Qwen-Image-Layered-Control) | [code](/examples/qwen_image/model_inference/Qwen-Image-Layered-Control.py) | [code](/examples/qwen_image/model_inference_low_vram/Qwen-Image-Layered-Control.py) | [code](/examples/qwen_image/model_training/full/Qwen-Image-Layered-Control.sh) | [code](/examples/qwen_image/model_training/validate_full/Qwen-Image-Layered-Control.py) | [code](/examples/qwen_image/model_training/lora/Qwen-Image-Layered-Control.sh) | [code](/examples/qwen_image/model_training/validate_lora/Qwen-Image-Layered-Control.py) |
+| Qwen-Image | [DiffSynth-Studio/Qwen-Image-Layered-Control-V2](https://www.modelscope.cn/models/DiffSynth-Studio/Qwen-Image-Layered-Control-V2) | [code](/examples/qwen_image/model_inference/Qwen-Image-Layered-Control-V2.py) | [code](/examples/qwen_image/model_inference_low_vram/Qwen-Image-Layered-Control-V2.py) | - | - | [code](/examples/qwen_image/model_training/lora/Qwen-Image-Layered-Control-V2.sh) | [code](/examples/qwen_image/model_training/validate_lora/Qwen-Image-Layered-Control-V2.py) |
+| Qwen-Image | [DiffSynth-Studio/Qwen-Image-EliGen](https://www.modelscope.cn/models/DiffSynth-Studio/Qwen-Image-EliGen) | [code](/examples/qwen_image/model_inference/Qwen-Image-EliGen.py) | [code](/examples/qwen_image/model_inference_low_vram/Qwen-Image-EliGen.py) | - | - | [code](/examples/qwen_image/model_training/lora/Qwen-Image-EliGen.sh) | [code](/examples/qwen_image/model_training/validate_lora/Qwen-Image-EliGen.py) |
+| Qwen-Image | [DiffSynth-Studio/Qwen-Image-EliGen-V2](https://www.modelscope.cn/models/DiffSynth-Studio/Qwen-Image-EliGen-V2) | [code](/examples/qwen_image/model_inference/Qwen-Image-EliGen-V2.py) | [code](/examples/qwen_image/model_inference_low_vram/Qwen-Image-EliGen-V2.py) | - | - | [code](/examples/qwen_image/model_training/lora/Qwen-Image-EliGen.sh) | [code](/examples/qwen_image/model_training/validate_lora/Qwen-Image-EliGen.py) |
+| Qwen-Image | [DiffSynth-Studio/Qwen-Image-EliGen-Poster](https://www.modelscope.cn/models/DiffSynth-Studio/Qwen-Image-EliGen-Poster) | [code](/examples/qwen_image/model_inference/Qwen-Image-EliGen-Poster.py) | [code](/examples/qwen_image/model_inference_low_vram/Qwen-Image-EliGen-Poster.py) | - | - | [code](/examples/qwen_image/model_training/lora/Qwen-Image-EliGen-Poster.sh) | [code](/examples/qwen_image/model_training/validate_lora/Qwen-Image-EliGen-Poster.py) |
+| Qwen-Image | [DiffSynth-Studio/Qwen-Image-Distill-Full](https://www.modelscope.cn/models/DiffSynth-Studio/Qwen-Image-Distill-Full) | [code](/examples/qwen_image/model_inference/Qwen-Image-Distill-Full.py) | [code](/examples/qwen_image/model_inference_low_vram/Qwen-Image-Distill-Full.py) | [code](/examples/qwen_image/model_training/full/Qwen-Image-Distill-Full.sh) | [code](/examples/qwen_image/model_training/validate_full/Qwen-Image-Distill-Full.py) | [code](/examples/qwen_image/model_training/lora/Qwen-Image-Distill-Full.sh) | [code](/examples/qwen_image/model_training/validate_lora/Qwen-Image-Distill-Full.py) |
+| Qwen-Image | [DiffSynth-Studio/Qwen-Image-Distill-LoRA](https://www.modelscope.cn/models/DiffSynth-Studio/Qwen-Image-Distill-LoRA) | [code](/examples/qwen_image/model_inference/Qwen-Image-Distill-LoRA.py) | [code](/examples/qwen_image/model_inference_low_vram/Qwen-Image-Distill-LoRA.py) | - | - | [code](/examples/qwen_image/model_training/lora/Qwen-Image-Distill-LoRA.sh) | [code](/examples/qwen_image/model_training/validate_lora/Qwen-Image-Distill-LoRA.py) |
+| Qwen-Image | [DiffSynth-Studio/Qwen-Image-Blockwise-ControlNet-Canny](https://modelscope.cn/models/DiffSynth-Studio/Qwen-Image-Blockwise-ControlNet-Canny) | [code](/examples/qwen_image/model_inference/Qwen-Image-Blockwise-ControlNet-Canny.py) | [code](/examples/qwen_image/model_inference_low_vram/Qwen-Image-Blockwise-ControlNet-Canny.py) | [code](/examples/qwen_image/model_training/full/Qwen-Image-Blockwise-ControlNet-Canny.sh) | [code](/examples/qwen_image/model_training/validate_full/Qwen-Image-Blockwise-ControlNet-Canny.py) | [code](/examples/qwen_image/model_training/lora/Qwen-Image-Blockwise-ControlNet-Canny.sh) | [code](/examples/qwen_image/model_training/validate_lora/Qwen-Image-Blockwise-ControlNet-Canny.py) |
+| Qwen-Image | [DiffSynth-Studio/Qwen-Image-Blockwise-ControlNet-Depth](https://modelscope.cn/models/DiffSynth-Studio/Qwen-Image-Blockwise-ControlNet-Depth) | [code](/examples/qwen_image/model_inference/Qwen-Image-Blockwise-ControlNet-Depth.py) | [code](/examples/qwen_image/model_inference_low_vram/Qwen-Image-Blockwise-ControlNet-Depth.py) | [code](/examples/qwen_image/model_training/full/Qwen-Image-Blockwise-ControlNet-Depth.sh) | [code](/examples/qwen_image/model_training/validate_full/Qwen-Image-Blockwise-ControlNet-Depth.py) | [code](/examples/qwen_image/model_training/lora/Qwen-Image-Blockwise-ControlNet-Depth.sh) | [code](/examples/qwen_image/model_training/validate_lora/Qwen-Image-Blockwise-ControlNet-Depth.py) |
+| Qwen-Image | [DiffSynth-Studio/Qwen-Image-Blockwise-ControlNet-Inpaint](https://modelscope.cn/models/DiffSynth-Studio/Qwen-Image-Blockwise-ControlNet-Inpaint) | [code](/examples/qwen_image/model_inference/Qwen-Image-Blockwise-ControlNet-Inpaint.py) | [code](/examples/qwen_image/model_inference_low_vram/Qwen-Image-Blockwise-ControlNet-Inpaint.py) | [code](/examples/qwen_image/model_training/full/Qwen-Image-Blockwise-ControlNet-Inpaint.sh) | [code](/examples/qwen_image/model_training/validate_full/Qwen-Image-Blockwise-ControlNet-Inpaint.py) | [code](/examples/qwen_image/model_training/lora/Qwen-Image-Blockwise-ControlNet-Inpaint.sh) | [code](/examples/qwen_image/model_training/validate_lora/Qwen-Image-Blockwise-ControlNet-Inpaint.py) |
+| Qwen-Image | [DiffSynth-Studio/Qwen-Image-In-Context-Control-Union](https://www.modelscope.cn/models/DiffSynth-Studio/Qwen-Image-In-Context-Control-Union) | [code](/examples/qwen_image/model_inference/Qwen-Image-In-Context-Control-Union.py) | [code](/examples/qwen_image/model_inference_low_vram/Qwen-Image-In-Context-Control-Union.py) | - | - | [code](/examples/qwen_image/model_training/lora/Qwen-Image-In-Context-Control-Union.sh) | [code](/examples/qwen_image/model_training/validate_lora/Qwen-Image-In-Context-Control-Union.py) |
+| Qwen-Image | [DiffSynth-Studio/Qwen-Image-Edit-Lowres-Fix](https://www.modelscope.cn/models/DiffSynth-Studio/Qwen-Image-Edit-Lowres-Fix) | [code](/examples/qwen_image/model_inference/Qwen-Image-Edit-Lowres-Fix.py) | [code](/examples/qwen_image/model_inference_low_vram/Qwen-Image-Edit-Lowres-Fix.py) | - | - | - | - |
+| Qwen-Image | [DiffSynth-Studio/Qwen-Image-i2L](https://www.modelscope.cn/models/DiffSynth-Studio/Qwen-Image-i2L) | [code](/examples/qwen_image/model_inference/Qwen-Image-i2L.py) | [code](/examples/qwen_image/model_inference_low_vram/Qwen-Image-i2L.py) | - | - | - | - |
+| Qwen-Video-Edit | [yunpeng1998/Qwen-Video-Edit](https://www.modelscope.cn/models/yunpeng1998/Qwen-Video-Edit) | [code](/examples/qwen_video_edit/model_inference/Qwen-Video-Edit.py) | [code](/examples/qwen_video_edit/model_inference_low_vram/Qwen-Video-Edit.py) | [code](/examples/qwen_video_edit/model_training/full/Qwen-Video-Edit.sh) | [code](/examples/qwen_video_edit/model_training/validate_full/Qwen-Video-Edit.py) | [code](/examples/qwen_video_edit/model_training/lora/Qwen-Video-Edit.sh) | [code](/examples/qwen_video_edit/model_training/validate_lora/Qwen-Video-Edit.py) |
+| Wan | [Wan-AI/Wan2.1-T2V-1.3B](https://modelscope.cn/models/Wan-AI/Wan2.1-T2V-1.3B) | [code](/examples/wanvideo/model_inference/Wan2.1-T2V-1.3B.py) | [code](/examples/wanvideo/model_inference_low_vram/Wan2.1-T2V-1.3B.py) | [code](/examples/wanvideo/model_training/full/Wan2.1-T2V-1.3B.sh) | [code](/examples/wanvideo/model_training/validate_full/Wan2.1-T2V-1.3B.py) | [code](/examples/wanvideo/model_training/lora/Wan2.1-T2V-1.3B.sh) | [code](/examples/wanvideo/model_training/validate_lora/Wan2.1-T2V-1.3B.py) |
+| Wan | [Wan-AI/Wan2.1-T2V-14B](https://modelscope.cn/models/Wan-AI/Wan2.1-T2V-14B) | [code](/examples/wanvideo/model_inference/Wan2.1-T2V-14B.py) | [code](/examples/wanvideo/model_inference_low_vram/Wan2.1-T2V-14B.py) | [code](/examples/wanvideo/model_training/full/Wan2.1-T2V-14B.sh) | [code](/examples/wanvideo/model_training/validate_full/Wan2.1-T2V-14B.py) | [code](/examples/wanvideo/model_training/lora/Wan2.1-T2V-14B.sh) | [code](/examples/wanvideo/model_training/validate_lora/Wan2.1-T2V-14B.py) |
+| Wan | [Wan-AI/Wan2.1-I2V-14B-480P](https://modelscope.cn/models/Wan-AI/Wan2.1-I2V-14B-480P) | [code](/examples/wanvideo/model_inference/Wan2.1-I2V-14B-480P.py) | [code](/examples/wanvideo/model_inference_low_vram/Wan2.1-I2V-14B-480P.py) | [code](/examples/wanvideo/model_training/full/Wan2.1-I2V-14B-480P.sh) | [code](/examples/wanvideo/model_training/validate_full/Wan2.1-I2V-14B-480P.py) | [code](/examples/wanvideo/model_training/lora/Wan2.1-I2V-14B-480P.sh) | [code](/examples/wanvideo/model_training/validate_lora/Wan2.1-I2V-14B-480P.py) |
+| Wan | [Wan-AI/Wan2.1-I2V-14B-720P](https://modelscope.cn/models/Wan-AI/Wan2.1-I2V-14B-720P) | [code](/examples/wanvideo/model_inference/Wan2.1-I2V-14B-720P.py) | [code](/examples/wanvideo/model_inference_low_vram/Wan2.1-I2V-14B-720P.py) | [code](/examples/wanvideo/model_training/full/Wan2.1-I2V-14B-720P.sh) | [code](/examples/wanvideo/model_training/validate_full/Wan2.1-I2V-14B-720P.py) | [code](/examples/wanvideo/model_training/lora/Wan2.1-I2V-14B-720P.sh) | [code](/examples/wanvideo/model_training/validate_lora/Wan2.1-I2V-14B-720P.py) |
+| Wan | [Wan-AI/Wan2.1-FLF2V-14B-720P](https://modelscope.cn/models/Wan-AI/Wan2.1-FLF2V-14B-720P) | [code](/examples/wanvideo/model_inference/Wan2.1-FLF2V-14B-720P.py) | [code](/examples/wanvideo/model_inference_low_vram/Wan2.1-FLF2V-14B-720P.py) | [code](/examples/wanvideo/model_training/full/Wan2.1-FLF2V-14B-720P.sh) | [code](/examples/wanvideo/model_training/validate_full/Wan2.1-FLF2V-14B-720P.py) | [code](/examples/wanvideo/model_training/lora/Wan2.1-FLF2V-14B-720P.sh) | [code](/examples/wanvideo/model_training/validate_lora/Wan2.1-FLF2V-14B-720P.py) |
+| Wan | [iic/VACE-Wan2.1-1.3B-Preview](https://modelscope.cn/models/iic/VACE-Wan2.1-1.3B-Preview) | [code](/examples/wanvideo/model_inference/Wan2.1-VACE-1.3B-Preview.py) | [code](/examples/wanvideo/model_inference_low_vram/Wan2.1-VACE-1.3B-Preview.py) | [code](/examples/wanvideo/model_training/full/Wan2.1-VACE-1.3B-Preview.sh) | [code](/examples/wanvideo/model_training/validate_full/Wan2.1-VACE-1.3B-Preview.py) | [code](/examples/wanvideo/model_training/lora/Wan2.1-VACE-1.3B-Preview.sh) | [code](/examples/wanvideo/model_training/validate_lora/Wan2.1-VACE-1.3B-Preview.py) |
+| Wan | [Wan-AI/Wan2.1-VACE-1.3B](https://modelscope.cn/models/Wan-AI/Wan2.1-VACE-1.3B) | [code](/examples/wanvideo/model_inference/Wan2.1-VACE-1.3B.py) | [code](/examples/wanvideo/model_inference_low_vram/Wan2.1-VACE-1.3B.py) | [code](/examples/wanvideo/model_training/full/Wan2.1-VACE-1.3B.sh) | [code](/examples/wanvideo/model_training/validate_full/Wan2.1-VACE-1.3B.py) | [code](/examples/wanvideo/model_training/lora/Wan2.1-VACE-1.3B.sh) | [code](/examples/wanvideo/model_training/validate_lora/Wan2.1-VACE-1.3B.py) |
+| Wan | [Wan-AI/Wan2.1-VACE-14B](https://modelscope.cn/models/Wan-AI/Wan2.1-VACE-14B) | [code](/examples/wanvideo/model_inference/Wan2.1-VACE-14B.py) | [code](/examples/wanvideo/model_inference_low_vram/Wan2.1-VACE-14B.py) | [code](/examples/wanvideo/model_training/full/Wan2.1-VACE-14B.sh) | [code](/examples/wanvideo/model_training/validate_full/Wan2.1-VACE-14B.py) | [code](/examples/wanvideo/model_training/lora/Wan2.1-VACE-14B.sh) | [code](/examples/wanvideo/model_training/validate_lora/Wan2.1-VACE-14B.py) |
+| Wan | [PAI/Wan2.1-Fun-1.3B-InP](https://modelscope.cn/models/PAI/Wan2.1-Fun-1.3B-InP) | [code](/examples/wanvideo/model_inference/Wan2.1-Fun-1.3B-InP.py) | [code](/examples/wanvideo/model_inference_low_vram/Wan2.1-Fun-1.3B-InP.py) | [code](/examples/wanvideo/model_training/full/Wan2.1-Fun-1.3B-InP.sh) | [code](/examples/wanvideo/model_training/validate_full/Wan2.1-Fun-1.3B-InP.py) | [code](/examples/wanvideo/model_training/lora/Wan2.1-Fun-1.3B-InP.sh) | [code](/examples/wanvideo/model_training/validate_lora/Wan2.1-Fun-1.3B-InP.py) |
+| Wan | [PAI/Wan2.1-Fun-1.3B-Control](https://modelscope.cn/models/PAI/Wan2.1-Fun-1.3B-Control) | [code](/examples/wanvideo/model_inference/Wan2.1-Fun-1.3B-Control.py) | [code](/examples/wanvideo/model_inference_low_vram/Wan2.1-Fun-1.3B-Control.py) | [code](/examples/wanvideo/model_training/full/Wan2.1-Fun-1.3B-Control.sh) | [code](/examples/wanvideo/model_training/validate_full/Wan2.1-Fun-1.3B-Control.py) | [code](/examples/wanvideo/model_training/lora/Wan2.1-Fun-1.3B-Control.sh) | [code](/examples/wanvideo/model_training/validate_lora/Wan2.1-Fun-1.3B-Control.py) |
+| Wan | [PAI/Wan2.1-Fun-14B-InP](https://modelscope.cn/models/PAI/Wan2.1-Fun-14B-InP) | [code](/examples/wanvideo/model_inference/Wan2.1-Fun-14B-InP.py) | [code](/examples/wanvideo/model_inference_low_vram/Wan2.1-Fun-14B-InP.py) | [code](/examples/wanvideo/model_training/full/Wan2.1-Fun-14B-InP.sh) | [code](/examples/wanvideo/model_training/validate_full/Wan2.1-Fun-14B-InP.py) | [code](/examples/wanvideo/model_training/lora/Wan2.1-Fun-14B-InP.sh) | [code](/examples/wanvideo/model_training/validate_lora/Wan2.1-Fun-14B-InP.py) |
+| Wan | [PAI/Wan2.1-Fun-14B-Control](https://modelscope.cn/models/PAI/Wan2.1-Fun-14B-Control) | [code](/examples/wanvideo/model_inference/Wan2.1-Fun-14B-Control.py) | [code](/examples/wanvideo/model_inference_low_vram/Wan2.1-Fun-14B-Control.py) | [code](/examples/wanvideo/model_training/full/Wan2.1-Fun-14B-Control.sh) | [code](/examples/wanvideo/model_training/validate_full/Wan2.1-Fun-14B-Control.py) | [code](/examples/wanvideo/model_training/lora/Wan2.1-Fun-14B-Control.sh) | [code](/examples/wanvideo/model_training/validate_lora/Wan2.1-Fun-14B-Control.py) |
+| Wan | [PAI/Wan2.1-Fun-V1.1-1.3B-Control](https://modelscope.cn/models/PAI/Wan2.1-Fun-V1.1-1.3B-Control) | [code](/examples/wanvideo/model_inference/Wan2.1-Fun-V1.1-1.3B-Control.py) | [code](/examples/wanvideo/model_inference_low_vram/Wan2.1-Fun-V1.1-1.3B-Control.py) | [code](/examples/wanvideo/model_training/full/Wan2.1-Fun-V1.1-1.3B-Control.sh) | [code](/examples/wanvideo/model_training/validate_full/Wan2.1-Fun-V1.1-1.3B-Control.py) | [code](/examples/wanvideo/model_training/lora/Wan2.1-Fun-V1.1-1.3B-Control.sh) | [code](/examples/wanvideo/model_training/validate_lora/Wan2.1-Fun-V1.1-1.3B-Control.py) |
+| Wan | [PAI/Wan2.1-Fun-V1.1-14B-Control](https://modelscope.cn/models/PAI/Wan2.1-Fun-V1.1-14B-Control) | [code](/examples/wanvideo/model_inference/Wan2.1-Fun-V1.1-14B-Control.py) | [code](/examples/wanvideo/model_inference_low_vram/Wan2.1-Fun-V1.1-14B-Control.py) | [code](/examples/wanvideo/model_training/full/Wan2.1-Fun-V1.1-14B-Control.sh) | [code](/examples/wanvideo/model_training/validate_full/Wan2.1-Fun-V1.1-14B-Control.py) | [code](/examples/wanvideo/model_training/lora/Wan2.1-Fun-V1.1-14B-Control.sh) | [code](/examples/wanvideo/model_training/validate_lora/Wan2.1-Fun-V1.1-14B-Control.py) |
+| Wan | [PAI/Wan2.1-Fun-V1.1-1.3B-InP](https://modelscope.cn/models/PAI/Wan2.1-Fun-V1.1-1.3B-InP) | [code](/examples/wanvideo/model_inference/Wan2.1-Fun-V1.1-1.3B-InP.py) | [code](/examples/wanvideo/model_inference_low_vram/Wan2.1-Fun-V1.1-1.3B-InP.py) | [code](/examples/wanvideo/model_training/full/Wan2.1-Fun-V1.1-1.3B-InP.sh) | [code](/examples/wanvideo/model_training/validate_full/Wan2.1-Fun-V1.1-1.3B-InP.py) | [code](/examples/wanvideo/model_training/lora/Wan2.1-Fun-V1.1-1.3B-InP.sh) | [code](/examples/wanvideo/model_training/validate_lora/Wan2.1-Fun-V1.1-1.3B-InP.py) |
+| Wan | [PAI/Wan2.1-Fun-V1.1-14B-InP](https://modelscope.cn/models/PAI/Wan2.1-Fun-V1.1-14B-InP) | [code](/examples/wanvideo/model_inference/Wan2.1-Fun-V1.1-14B-InP.py) | [code](/examples/wanvideo/model_inference_low_vram/Wan2.1-Fun-V1.1-14B-InP.py) | [code](/examples/wanvideo/model_training/full/Wan2.1-Fun-V1.1-14B-InP.sh) | [code](/examples/wanvideo/model_training/validate_full/Wan2.1-Fun-V1.1-14B-InP.py) | [code](/examples/wanvideo/model_training/lora/Wan2.1-Fun-V1.1-14B-InP.sh) | [code](/examples/wanvideo/model_training/validate_lora/Wan2.1-Fun-V1.1-14B-InP.py) |
+| Wan | [PAI/Wan2.1-Fun-V1.1-1.3B-Control-Camera](https://modelscope.cn/models/PAI/Wan2.1-Fun-V1.1-1.3B-Control-Camera) | [code](/examples/wanvideo/model_inference/Wan2.1-Fun-V1.1-1.3B-Control-Camera.py) | [code](/examples/wanvideo/model_inference_low_vram/Wan2.1-Fun-V1.1-1.3B-Control-Camera.py) | [code](/examples/wanvideo/model_training/full/Wan2.1-Fun-V1.1-1.3B-Control-Camera.sh) | [code](/examples/wanvideo/model_training/validate_full/Wan2.1-Fun-V1.1-1.3B-Control-Camera.py) | [code](/examples/wanvideo/model_training/lora/Wan2.1-Fun-V1.1-1.3B-Control-Camera.sh) | [code](/examples/wanvideo/model_training/validate_lora/Wan2.1-Fun-V1.1-1.3B-Control-Camera.py) |
+| Wan | [PAI/Wan2.1-Fun-V1.1-14B-Control-Camera](https://modelscope.cn/models/PAI/Wan2.1-Fun-V1.1-14B-Control-Camera) | [code](/examples/wanvideo/model_inference/Wan2.1-Fun-V1.1-14B-Control-Camera.py) | [code](/examples/wanvideo/model_inference_low_vram/Wan2.1-Fun-V1.1-14B-Control-Camera.py) | [code](/examples/wanvideo/model_training/full/Wan2.1-Fun-V1.1-14B-Control-Camera.sh) | [code](/examples/wanvideo/model_training/validate_full/Wan2.1-Fun-V1.1-14B-Control-Camera.py) | [code](/examples/wanvideo/model_training/lora/Wan2.1-Fun-V1.1-14B-Control-Camera.sh) | [code](/examples/wanvideo/model_training/validate_lora/Wan2.1-Fun-V1.1-14B-Control-Camera.py) |
+| Wan | [DiffSynth-Studio/Wan2.1-1.3b-speedcontrol-v1](https://modelscope.cn/models/DiffSynth-Studio/Wan2.1-1.3b-speedcontrol-v1) | [code](/examples/wanvideo/model_inference/Wan2.1-1.3b-speedcontrol-v1.py) | [code](/examples/wanvideo/model_inference_low_vram/Wan2.1-1.3b-speedcontrol-v1.py) | [code](/examples/wanvideo/model_training/full/Wan2.1-1.3b-speedcontrol-v1.sh) | [code](/examples/wanvideo/model_training/validate_full/Wan2.1-1.3b-speedcontrol-v1.py) | [code](/examples/wanvideo/model_training/lora/Wan2.1-1.3b-speedcontrol-v1.sh) | [code](/examples/wanvideo/model_training/validate_lora/Wan2.1-1.3b-speedcontrol-v1.py) |
+| Wan | [krea/krea-realtime-video](https://www.modelscope.cn/models/krea/krea-realtime-video) | [code](/examples/wanvideo/model_inference/krea-realtime-video.py) | [code](/examples/wanvideo/model_inference_low_vram/krea-realtime-video.py) | [code](/examples/wanvideo/model_training/full/krea-realtime-video.sh) | [code](/examples/wanvideo/model_training/validate_full/krea-realtime-video.py) | [code](/examples/wanvideo/model_training/lora/krea-realtime-video.sh) | [code](/examples/wanvideo/model_training/validate_lora/krea-realtime-video.py) |
+| Wan | [meituan-longcat/LongCat-Video](https://www.modelscope.cn/models/meituan-longcat/LongCat-Video) | [code](/examples/wanvideo/model_inference/LongCat-Video.py) | [code](/examples/wanvideo/model_inference_low_vram/LongCat-Video.py) | [code](/examples/wanvideo/model_training/full/LongCat-Video.sh) | [code](/examples/wanvideo/model_training/validate_full/LongCat-Video.py) | [code](/examples/wanvideo/model_training/lora/LongCat-Video.sh) | [code](/examples/wanvideo/model_training/validate_lora/LongCat-Video.py) |
+| Wan | [ByteDance/Video-As-Prompt-Wan2.1-14B](https://modelscope.cn/models/ByteDance/Video-As-Prompt-Wan2.1-14B) | [code](/examples/wanvideo/model_inference/Video-As-Prompt-Wan2.1-14B.py) | [code](/examples/wanvideo/model_inference_low_vram/Video-As-Prompt-Wan2.1-14B.py) | [code](/examples/wanvideo/model_training/full/Video-As-Prompt-Wan2.1-14B.sh) | [code](/examples/wanvideo/model_training/validate_full/Video-As-Prompt-Wan2.1-14B.py) | [code](/examples/wanvideo/model_training/lora/Video-As-Prompt-Wan2.1-14B.sh) | [code](/examples/wanvideo/model_training/validate_lora/Video-As-Prompt-Wan2.1-14B.py) |
+| Wan | [Wan-AI/Wan2.2-T2V-A14B](https://modelscope.cn/models/Wan-AI/Wan2.2-T2V-A14B) | [code](/examples/wanvideo/model_inference/Wan2.2-T2V-A14B.py) | [code](/examples/wanvideo/model_inference_low_vram/Wan2.2-T2V-A14B.py) | [code](/examples/wanvideo/model_training/full/Wan2.2-T2V-A14B.sh) | [code](/examples/wanvideo/model_training/validate_full/Wan2.2-T2V-A14B.py) | [code](/examples/wanvideo/model_training/lora/Wan2.2-T2V-A14B.sh) | [code](/examples/wanvideo/model_training/validate_lora/Wan2.2-T2V-A14B.py) |
+| Wan | [Wan-AI/Wan2.2-I2V-A14B](https://modelscope.cn/models/Wan-AI/Wan2.2-I2V-A14B) | [code](/examples/wanvideo/model_inference/Wan2.2-I2V-A14B.py) | [code](/examples/wanvideo/model_inference_low_vram/Wan2.2-I2V-A14B.py) | [code](/examples/wanvideo/model_training/full/Wan2.2-I2V-A14B.sh) | [code](/examples/wanvideo/model_training/validate_full/Wan2.2-I2V-A14B.py) | [code](/examples/wanvideo/model_training/lora/Wan2.2-I2V-A14B.sh) | [code](/examples/wanvideo/model_training/validate_lora/Wan2.2-I2V-A14B.py) |
+| Wan | [Wan-AI/Wan2.2-TI2V-5B](https://modelscope.cn/models/Wan-AI/Wan2.2-TI2V-5B) | [code](/examples/wanvideo/model_inference/Wan2.2-TI2V-5B.py) | [code](/examples/wanvideo/model_inference_low_vram/Wan2.2-TI2V-5B.py) | [code](/examples/wanvideo/model_training/full/Wan2.2-TI2V-5B.sh) | [code](/examples/wanvideo/model_training/validate_full/Wan2.2-TI2V-5B.py) | [code](/examples/wanvideo/model_training/lora/Wan2.2-TI2V-5B.sh) | [code](/examples/wanvideo/model_training/validate_lora/Wan2.2-TI2V-5B.py) |
+| Wan | [Wan-AI/Wan2.2-Animate-14B](https://www.modelscope.cn/models/Wan-AI/Wan2.2-Animate-14B) | [code](/examples/wanvideo/model_inference/Wan2.2-Animate-14B.py) | [code](/examples/wanvideo/model_inference_low_vram/Wan2.2-Animate-14B.py) | [code](/examples/wanvideo/model_training/full/Wan2.2-Animate-14B.sh) | [code](/examples/wanvideo/model_training/validate_full/Wan2.2-Animate-14B.py) | [code](/examples/wanvideo/model_training/lora/Wan2.2-Animate-14B.sh) | [code](/examples/wanvideo/model_training/validate_lora/Wan2.2-Animate-14B.py) |
+| Wan | [Wan-AI/Wan2.2-Animate-2-14B](https://www.modelscope.cn/models/Wan-AI/Wan2.2-Animate-2-14B) | [code](/examples/wanvideo/model_inference/Wan2.2-Animate-2-14B.py) | [code](/examples/wanvideo/model_inference_low_vram/Wan2.2-Animate-2-14B.py) | [code](/examples/wanvideo/model_training/full/Wan2.2-Animate-2-14B.sh) | [code](/examples/wanvideo/model_training/validate_full/Wan2.2-Animate-2-14B.py) | [code](/examples/wanvideo/model_training/lora/Wan2.2-Animate-2-14B.sh) | [code](/examples/wanvideo/model_training/validate_lora/Wan2.2-Animate-2-14B.py) |
+| Wan | [Wan-AI/Wan2.2-Animate-2-14B: Distilled](https://www.modelscope.cn/models/Wan-AI/Wan2.2-Animate-2-14B) | [code](/examples/wanvideo/model_inference/Wan2.2-Animate-2-14B-Distilled.py) | [code](/examples/wanvideo/model_inference_low_vram/Wan2.2-Animate-2-14B-Distilled.py) | [code](/examples/wanvideo/model_training/full/Wan2.2-Animate-2-14B-Distilled.sh) | [code](/examples/wanvideo/model_training/validate_full/Wan2.2-Animate-2-14B-Distilled.py) | [code](/examples/wanvideo/model_training/lora/Wan2.2-Animate-2-14B-Distilled.sh) | [code](/examples/wanvideo/model_training/validate_lora/Wan2.2-Animate-2-14B-Distilled.py) |
+| Wan | [Wan-AI/Wan2.2-S2V-14B](https://www.modelscope.cn/models/Wan-AI/Wan2.2-S2V-14B) | [code](/examples/wanvideo/model_inference/Wan2.2-S2V-14B_multi_clips.py) | [code](/examples/wanvideo/model_inference_low_vram/Wan2.2-S2V-14B_multi_clips.py) | [code](/examples/wanvideo/model_training/full/Wan2.2-S2V-14B.sh) | [code](/examples/wanvideo/model_training/validate_full/Wan2.2-S2V-14B.py) | [code](/examples/wanvideo/model_training/lora/Wan2.2-S2V-14B.sh) | [code](/examples/wanvideo/model_training/validate_lora/Wan2.2-S2V-14B.py) |
+| Wan | [PAI/Wan2.2-VACE-Fun-A14B](https://www.modelscope.cn/models/PAI/Wan2.2-VACE-Fun-A14B) | [code](/examples/wanvideo/model_inference/Wan2.2-VACE-Fun-A14B.py) | [code](/examples/wanvideo/model_inference_low_vram/Wan2.2-VACE-Fun-A14B.py) | [code](/examples/wanvideo/model_training/full/Wan2.2-VACE-Fun-A14B.sh) | [code](/examples/wanvideo/model_training/validate_full/Wan2.2-VACE-Fun-A14B.py) | [code](/examples/wanvideo/model_training/lora/Wan2.2-VACE-Fun-A14B.sh) | [code](/examples/wanvideo/model_training/validate_lora/Wan2.2-VACE-Fun-A14B.py) |
+| Wan | [PAI/Wan2.2-Fun-A14B-InP](https://modelscope.cn/models/PAI/Wan2.2-Fun-A14B-InP) | [code](/examples/wanvideo/model_inference/Wan2.2-Fun-A14B-InP.py) | [code](/examples/wanvideo/model_inference_low_vram/Wan2.2-Fun-A14B-InP.py) | [code](/examples/wanvideo/model_training/full/Wan2.2-Fun-A14B-InP.sh) | [code](/examples/wanvideo/model_training/validate_full/Wan2.2-Fun-A14B-InP.py) | [code](/examples/wanvideo/model_training/lora/Wan2.2-Fun-A14B-InP.sh) | [code](/examples/wanvideo/model_training/validate_lora/Wan2.2-Fun-A14B-InP.py) |
+| Wan | [PAI/Wan2.2-Fun-A14B-Control](https://modelscope.cn/models/PAI/Wan2.2-Fun-A14B-Control) | [code](/examples/wanvideo/model_inference/Wan2.2-Fun-A14B-Control.py) | [code](/examples/wanvideo/model_inference_low_vram/Wan2.2-Fun-A14B-Control.py) | [code](/examples/wanvideo/model_training/full/Wan2.2-Fun-A14B-Control.sh) | [code](/examples/wanvideo/model_training/validate_full/Wan2.2-Fun-A14B-Control.py) | [code](/examples/wanvideo/model_training/lora/Wan2.2-Fun-A14B-Control.sh) | [code](/examples/wanvideo/model_training/validate_lora/Wan2.2-Fun-A14B-Control.py) |
+| Wan | [PAI/Wan2.2-Fun-A14B-Control-Camera](https://modelscope.cn/models/PAI/Wan2.2-Fun-A14B-Control-Camera) | [code](/examples/wanvideo/model_inference/Wan2.2-Fun-A14B-Control-Camera.py) | [code](/examples/wanvideo/model_inference_low_vram/Wan2.2-Fun-A14B-Control-Camera.py) | [code](/examples/wanvideo/model_training/full/Wan2.2-Fun-A14B-Control-Camera.sh) | [code](/examples/wanvideo/model_training/validate_full/Wan2.2-Fun-A14B-Control-Camera.py) | [code](/examples/wanvideo/model_training/lora/Wan2.2-Fun-A14B-Control-Camera.sh) | [code](/examples/wanvideo/model_training/validate_lora/Wan2.2-Fun-A14B-Control-Camera.py) |
+| Wan | [openmoss/MOVA-360p](https://modelscope.cn/models/openmoss/MOVA-360p) | [code](/examples/mova/model_inference/MOVA-360p-I2AV.py) | [code](/examples/mova/model_inference_low_vram/MOVA-360p-I2AV.py) | [code](/examples/mova/model_training/full/MOVA-360P-I2AV.sh) | [code](/examples/mova/model_training/validate_full/MOVA-360p-I2AV.py) | [code](/examples/mova/model_training/lora/MOVA-360P-I2AV.sh) | [code](/examples/mova/model_training/validate_lora/MOVA-360p-I2AV.py) |
+| Wan | [openmoss/MOVA-720p](https://modelscope.cn/models/openmoss/MOVA-720p) | [code](/examples/mova/model_inference/MOVA-720p-I2AV.py) | [code](/examples/mova/model_inference_low_vram/MOVA-720p-I2AV.py) | [code](/examples/mova/model_training/full/MOVA-720P-I2AV.sh) | [code](/examples/mova/model_training/validate_full/MOVA-720p-I2AV.py) | [code](/examples/mova/model_training/lora/MOVA-720P-I2AV.sh) | [code](/examples/mova/model_training/validate_lora/MOVA-720p-I2AV.py) |
+| Wan | [Wan-AI/Wan-Dancer-14B (global model)](https://modelscope.cn/models/Wan-AI/Wan-Dancer-14B) | [code](/examples/wanvideo/model_inference/Wan-Dancer-14B-global.py) | [code](/examples/wanvideo/model_inference_low_vram/Wan-Dancer-14B-global.py) | [code](/examples/wanvideo/model_training/full/Wan-Dancer-14B-global.sh) | [code](/examples/wanvideo/model_training/validate_full/Wan-Dancer-14B-global.py) | [code](/examples/wanvideo/model_training/lora/Wan-Dancer-14B-global.sh) | [code](/examples/wanvideo/model_training/validate_lora/Wan-Dancer-14B-global.py) |
+| Wan | [Wan-AI/Wan-Dancer-14B (local model)](https://modelscope.cn/models/Wan-AI/Wan-Dancer-14B) | [code](/examples/wanvideo/model_inference/Wan-Dancer-14B-local.py) | [code](/examples/wanvideo/model_inference_low_vram/Wan-Dancer-14B-local.py) | [code](/examples/wanvideo/model_training/full/Wan-Dancer-14B-local.sh) | [code](/examples/wanvideo/model_training/validate_full/Wan-Dancer-14B-local.py) | [code](/examples/wanvideo/model_training/lora/Wan-Dancer-14B-local.sh) | [code](/examples/wanvideo/model_training/validate_lora/Wan-Dancer-14B-local.py) |
+| FLUX.1 | [black-forest-labs/FLUX.1-dev](https://www.modelscope.cn/models/black-forest-labs/FLUX.1-dev) | [code](/examples/flux/model_inference/FLUX.1-dev.py) | [code](/examples/flux/model_inference_low_vram/FLUX.1-dev.py) | [code](/examples/flux/model_training/full/FLUX.1-dev.sh) | [code](/examples/flux/model_training/validate_full/FLUX.1-dev.py) | [code](/examples/flux/model_training/lora/FLUX.1-dev.sh) | [code](/examples/flux/model_training/validate_lora/FLUX.1-dev.py) |
+| FLUX.1 | [black-forest-labs/FLUX.1-Krea-dev](https://www.modelscope.cn/models/black-forest-labs/FLUX.1-Krea-dev) | [code](/examples/flux/model_inference/FLUX.1-Krea-dev.py) | [code](/examples/flux/model_inference_low_vram/FLUX.1-Krea-dev.py) | [code](/examples/flux/model_training/full/FLUX.1-Krea-dev.sh) | [code](/examples/flux/model_training/validate_full/FLUX.1-Krea-dev.py) | [code](/examples/flux/model_training/lora/FLUX.1-Krea-dev.sh) | [code](/examples/flux/model_training/validate_lora/FLUX.1-Krea-dev.py) |
+| FLUX.1 | [black-forest-labs/FLUX.1-Kontext-dev](https://www.modelscope.cn/models/black-forest-labs/FLUX.1-Kontext-dev) | [code](/examples/flux/model_inference/FLUX.1-Kontext-dev.py) | [code](/examples/flux/model_inference_low_vram/FLUX.1-Kontext-dev.py) | [code](/examples/flux/model_training/full/FLUX.1-Kontext-dev.sh) | [code](/examples/flux/model_training/validate_full/FLUX.1-Kontext-dev.py) | [code](/examples/flux/model_training/lora/FLUX.1-Kontext-dev.sh) | [code](/examples/flux/model_training/validate_lora/FLUX.1-Kontext-dev.py) |
+| FLUX.1 | [black-forest-labs/FLUX.1-Fill-dev](https://www.modelscope.cn/models/black-forest-labs/FLUX.1-Fill-dev) | [code](/examples/flux/model_inference/FLUX.1-Fill-dev.py) | [code](/examples/flux/model_inference_low_vram/FLUX.1-Fill-dev.py) | [code](/examples/flux/model_training/full/FLUX.1-Fill-dev.sh) | [code](/examples/flux/model_training/validate_full/FLUX.1-Fill-dev.py) | [code](/examples/flux/model_training/lora/FLUX.1-Fill-dev.sh) | [code](/examples/flux/model_training/validate_lora/FLUX.1-Fill-dev.py) |
+| FLUX.1 | [black-forest-labs/FLUX.1-Redux-dev](https://www.modelscope.cn/models/black-forest-labs/FLUX.1-Redux-dev) | [code](/examples/flux/model_inference/FLUX.1-Redux-dev.py) | [code](/examples/flux/model_inference_low_vram/FLUX.1-Redux-dev.py) | [code](/examples/flux/model_training/full/FLUX.1-Redux-dev.sh) | [code](/examples/flux/model_training/validate_full/FLUX.1-Redux-dev.py) | [code](/examples/flux/model_training/lora/FLUX.1-Redux-dev.sh) | [code](/examples/flux/model_training/validate_lora/FLUX.1-Redux-dev.py) |
+| FLUX.1 | [HuanJue/Insert-Anything](https://www.modelscope.cn/models/HuanJue/Insert-Anything) | [code](/examples/flux/model_inference/Insert-Anything.py) | [code](/examples/flux/model_inference_low_vram/Insert-Anything.py) | - | - | [code](/examples/flux/model_training/lora/Insert-Anything.sh) | [code](/examples/flux/model_training/validate_lora/Insert-Anything.py) |
+| FLUX.1 | [alimama-creative/FLUX.1-dev-Controlnet-Inpainting-Beta](https://www.modelscope.cn/models/alimama-creative/FLUX.1-dev-Controlnet-Inpainting-Beta) | [code](/examples/flux/model_inference/FLUX.1-dev-Controlnet-Inpainting-Beta.py) | [code](/examples/flux/model_inference_low_vram/FLUX.1-dev-Controlnet-Inpainting-Beta.py) | [code](/examples/flux/model_training/full/FLUX.1-dev-Controlnet-Inpainting-Beta.sh) | [code](/examples/flux/model_training/validate_full/FLUX.1-dev-Controlnet-Inpainting-Beta.py) | [code](/examples/flux/model_training/lora/FLUX.1-dev-Controlnet-Inpainting-Beta.sh) | [code](/examples/flux/model_training/validate_lora/FLUX.1-dev-Controlnet-Inpainting-Beta.py) |
+| FLUX.1 | [InstantX/FLUX.1-dev-Controlnet-Union-alpha](https://www.modelscope.cn/models/InstantX/FLUX.1-dev-Controlnet-Union-alpha) | [code](/examples/flux/model_inference/FLUX.1-dev-Controlnet-Union-alpha.py) | [code](/examples/flux/model_inference_low_vram/FLUX.1-dev-Controlnet-Union-alpha.py) | [code](/examples/flux/model_training/full/FLUX.1-dev-Controlnet-Union-alpha.sh) | [code](/examples/flux/model_training/validate_full/FLUX.1-dev-Controlnet-Union-alpha.py) | [code](/examples/flux/model_training/lora/FLUX.1-dev-Controlnet-Union-alpha.sh) | [code](/examples/flux/model_training/validate_lora/FLUX.1-dev-Controlnet-Union-alpha.py) |
+| FLUX.1 | [jasperai/Flux.1-dev-Controlnet-Upscaler](https://www.modelscope.cn/models/jasperai/Flux.1-dev-Controlnet-Upscaler) | [code](/examples/flux/model_inference/FLUX.1-dev-Controlnet-Upscaler.py) | [code](/examples/flux/model_inference_low_vram/FLUX.1-dev-Controlnet-Upscaler.py) | [code](/examples/flux/model_training/full/FLUX.1-dev-Controlnet-Upscaler.sh) | [code](/examples/flux/model_training/validate_full/FLUX.1-dev-Controlnet-Upscaler.py) | [code](/examples/flux/model_training/lora/FLUX.1-dev-Controlnet-Upscaler.sh) | [code](/examples/flux/model_training/validate_lora/FLUX.1-dev-Controlnet-Upscaler.py) |
+| FLUX.1 | [InstantX/FLUX.1-dev-IP-Adapter](https://www.modelscope.cn/models/InstantX/FLUX.1-dev-IP-Adapter) | [code](/examples/flux/model_inference/FLUX.1-dev-IP-Adapter.py) | [code](/examples/flux/model_inference_low_vram/FLUX.1-dev-IP-Adapter.py) | [code](/examples/flux/model_training/full/FLUX.1-dev-IP-Adapter.sh) | [code](/examples/flux/model_training/validate_full/FLUX.1-dev-IP-Adapter.py) | [code](/examples/flux/model_training/lora/FLUX.1-dev-IP-Adapter.sh) | [code](/examples/flux/model_training/validate_lora/FLUX.1-dev-IP-Adapter.py) |
+| FLUX.1 | [ByteDance/InfiniteYou](https://www.modelscope.cn/models/ByteDance/InfiniteYou) | [code](/examples/flux/model_inference/FLUX.1-dev-InfiniteYou.py) | [code](/examples/flux/model_inference_low_vram/FLUX.1-dev-InfiniteYou.py) | [code](/examples/flux/model_training/full/FLUX.1-dev-InfiniteYou.sh) | [code](/examples/flux/model_training/validate_full/FLUX.1-dev-InfiniteYou.py) | [code](/examples/flux/model_training/lora/FLUX.1-dev-InfiniteYou.sh) | [code](/examples/flux/model_training/validate_lora/FLUX.1-dev-InfiniteYou.py) |
+| FLUX.1 | [DiffSynth-Studio/Eligen](https://www.modelscope.cn/models/DiffSynth-Studio/Eligen) | [code](/examples/flux/model_inference/FLUX.1-dev-EliGen.py) | [code](/examples/flux/model_inference_low_vram/FLUX.1-dev-EliGen.py) | - | - | [code](/examples/flux/model_training/lora/FLUX.1-dev-EliGen.sh) | [code](/examples/flux/model_training/validate_lora/FLUX.1-dev-EliGen.py) |
+| FLUX.1 | [DiffSynth-Studio/LoRA-Encoder-FLUX.1-Dev](https://www.modelscope.cn/models/DiffSynth-Studio/LoRA-Encoder-FLUX.1-Dev) | [code](/examples/flux/model_inference/FLUX.1-dev-LoRA-Encoder.py) | [code](/examples/flux/model_inference_low_vram/FLUX.1-dev-LoRA-Encoder.py) | [code](/examples/flux/model_training/full/FLUX.1-dev-LoRA-Encoder.sh) | [code](/examples/flux/model_training/validate_full/FLUX.1-dev-LoRA-Encoder.py) | - | - |
+| FLUX.1 | [DiffSynth-Studio/LoRAFusion-preview-FLUX.1-dev](https://modelscope.cn/models/DiffSynth-Studio/LoRAFusion-preview-FLUX.1-dev) | [code](/examples/flux/model_inference/FLUX.1-dev-LoRA-Fusion.py) | - | - | - | - | - |
+| FLUX.1 | [stepfun-ai/Step1X-Edit](https://www.modelscope.cn/models/stepfun-ai/Step1X-Edit) | [code](/examples/flux/model_inference/Step1X-Edit.py) | [code](/examples/flux/model_inference_low_vram/Step1X-Edit.py) | [code](/examples/flux/model_training/full/Step1X-Edit.sh) | [code](/examples/flux/model_training/validate_full/Step1X-Edit.py) | [code](/examples/flux/model_training/lora/Step1X-Edit.sh) | [code](/examples/flux/model_training/validate_lora/Step1X-Edit.py) |
+| FLUX.1 | [ostris/Flex.2-preview](https://www.modelscope.cn/models/ostris/Flex.2-preview) | [code](/examples/flux/model_inference/FLEX.2-preview.py) | [code](/examples/flux/model_inference_low_vram/FLEX.2-preview.py) | [code](/examples/flux/model_training/full/FLEX.2-preview.sh) | [code](/examples/flux/model_training/validate_full/FLEX.2-preview.py) | [code](/examples/flux/model_training/lora/FLEX.2-preview.sh) | [code](/examples/flux/model_training/validate_lora/FLEX.2-preview.py) |
+| FLUX.1 | [DiffSynth-Studio/Nexus-GenV2](https://www.modelscope.cn/models/DiffSynth-Studio/Nexus-GenV2) | [code](/examples/flux/model_inference/Nexus-Gen-Editing.py) | [code](/examples/flux/model_inference_low_vram/Nexus-Gen-Editing.py) | [code](/examples/flux/model_training/full/Nexus-Gen.sh) | [code](/examples/flux/model_training/validate_full/Nexus-Gen.py) | [code](/examples/flux/model_training/lora/Nexus-Gen.sh) | [code](/examples/flux/model_training/validate_lora/Nexus-Gen.py) |
+| Stable Diffusion XL | [stabilityai/stable-diffusion-xl-base-1.0](https://www.modelscope.cn/models/stabilityai/stable-diffusion-xl-base-1.0) | [code](/examples/stable_diffusion_xl/model_inference/stable-diffusion-xl-base-1.0.py) | [code](/examples/stable_diffusion_xl/model_inference_low_vram/stable-diffusion-xl-base-1.0.py) | [code](/examples/stable_diffusion_xl/model_training/full/stable-diffusion-xl-base-1.0.sh) | [code](/examples/stable_diffusion_xl/model_training/validate_full/stable-diffusion-xl-base-1.0.py) | [code](/examples/stable_diffusion_xl/model_training/lora/stable-diffusion-xl-base-1.0.sh) | [code](/examples/stable_diffusion_xl/model_training/validate_lora/stable-diffusion-xl-base-1.0.py) |
+| Stable Diffusion | [AI-ModelScope/stable-diffusion-v1-5](https://www.modelscope.cn/models/AI-ModelScope/stable-diffusion-v1-5) | [code](/examples/stable_diffusion/model_inference/stable-diffusion-v1-5.py) | [code](/examples/stable_diffusion/model_inference_low_vram/stable-diffusion-v1-5.py) | [code](/examples/stable_diffusion/model_training/full/stable-diffusion-v1-5.sh) | [code](/examples/stable_diffusion/model_training/validate_full/stable-diffusion-v1-5.py) | [code](/examples/stable_diffusion/model_training/lora/stable-diffusion-v1-5.sh) | [code](/examples/stable_diffusion/model_training/validate_lora/stable-diffusion-v1-5.py) |
+| - | [PickScore](https://github.com/yuvalkirstain/pickscore) | [code](/examples/image_quality_metric/pickscore.py) | - | - | - | - | - |
+| - | [ImageReward](https://github.com/zai-org/ImageReward) | [code](/examples/image_quality_metric/image_reward.py) | - | - | - | - | - |
+| - | [HPSv2](https://github.com/tgxs002/HPSv2) | [code](/examples/image_quality_metric/hpsv2.py) | - | - | - | - | - |
+| - | [HPSv3](https://github.com/MizzenAI/HPSv3) | [code](/examples/image_quality_metric/hpsv3.py) | - | - | - | - | - |
+| - | [CLIP Score](https://github.com/openai/CLIP) | [code](/examples/image_quality_metric/clipscore.py) | - | - | - | - | - |
+| - | [UnifiedReward 2.0](https://github.com/THUDM/UnifiedReward) | [code](/examples/image_quality_metric/unified_reward_2.py) | - | - | - | - | - |
+| - | [Qwen-Image-Bench](https://github.com/QwenLM/Qwen-Image-Bench) | [code](/examples/image_quality_metric/qwen_image_bench.py) | - | - | - | - | - |
+| - | [UnifiedReward Edit](https://github.com/THUDM/UnifiedReward) | [code](/examples/image_quality_metric/unified_reward_edit.py) | - | - | - | - | - |
+| - | [Aesthetic](https://github.com/christophschuhmann/improved-aesthetic-predictor) | [code](/examples/image_quality_metric/aesthetic.py) | - | - | - | - | - |
+| - | [FID](https://github.com/mseitzer/pytorch-fid) | [code](/examples/image_quality_metric/fid.py) | - | - | - | - | - |
diff --git a/diffsynth/__init__.py b/diffsynth/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..bb67a43fa4e5791ab58e7e40260bc3df8b6bc7cc
--- /dev/null
+++ b/diffsynth/__init__.py
@@ -0,0 +1 @@
+from .core import *
diff --git a/diffsynth/configs/__init__.py b/diffsynth/configs/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..7ad5b7322143d89beb1dbec64d505c5c8b1ecf61
--- /dev/null
+++ b/diffsynth/configs/__init__.py
@@ -0,0 +1,2 @@
+from .model_configs import MODEL_CONFIGS
+from .vram_management_module_maps import VRAM_MANAGEMENT_MODULE_MAPS, VERSION_CHECKER_MAPS
diff --git a/diffsynth/configs/model_configs.py b/diffsynth/configs/model_configs.py
new file mode 100644
index 0000000000000000000000000000000000000000..3de21f9bf35f9de0c21842e3d26b49b61fe14003
--- /dev/null
+++ b/diffsynth/configs/model_configs.py
@@ -0,0 +1,1569 @@
+qwen_image_series = [
+ {
+ # Example: ModelConfig(model_id="Qwen/Qwen-Image", origin_file_pattern="text_encoder/model*.safetensors")
+ "model_hash": "0319a1cb19835fb510907dd3367c95ff",
+ "model_name": "qwen_image_dit",
+ "model_class": "diffsynth.models.qwen_image_dit.QwenImageDiT",
+ },
+ {
+ # Example: ModelConfig(model_id="Qwen/Qwen-Image", origin_file_pattern="transformer/diffusion_pytorch_model*.safetensors")
+ "model_hash": "8004730443f55db63092006dd9f7110e",
+ "model_name": "qwen_image_text_encoder",
+ "model_class": "diffsynth.models.qwen_image_text_encoder.QwenImageTextEncoder",
+ "state_dict_converter": "diffsynth.utils.state_dict_converters.qwen_image_text_encoder.QwenImageTextEncoderStateDictConverter",
+ },
+ {
+ # Example: ModelConfig(model_id="Qwen/Qwen-Image", origin_file_pattern="vae/diffusion_pytorch_model.safetensors")
+ "model_hash": "ed4ea5824d55ec3107b09815e318123a",
+ "model_name": "qwen_image_vae",
+ "model_class": "diffsynth.models.qwen_image_vae.QwenImageVAE",
+ },
+ {
+ # Example: ModelConfig(model_id="DiffSynth-Studio/Qwen-Image-Blockwise-ControlNet-Depth", origin_file_pattern="model.safetensors")
+ "model_hash": "073bce9cf969e317e5662cd570c3e79c",
+ "model_name": "qwen_image_blockwise_controlnet",
+ "model_class": "diffsynth.models.qwen_image_controlnet.QwenImageBlockWiseControlNet",
+ },
+ {
+ # Example: ModelConfig(model_id="DiffSynth-Studio/Qwen-Image-Blockwise-ControlNet-Inpaint", origin_file_pattern="model.safetensors")
+ "model_hash": "a9e54e480a628f0b956a688a81c33bab",
+ "model_name": "qwen_image_blockwise_controlnet",
+ "model_class": "diffsynth.models.qwen_image_controlnet.QwenImageBlockWiseControlNet",
+ "extra_kwargs": {"additional_in_dim": 4},
+ },
+ {
+ # Example: ModelConfig(model_id="DiffSynth-Studio/General-Image-Encoders", origin_file_pattern="SigLIP2-G384/model.safetensors")
+ "model_hash": "469c78b61e3e31bc9eec0d0af3d3f2f8",
+ "model_name": "siglip2_image_encoder",
+ "model_class": "diffsynth.models.siglip2_image_encoder.Siglip2ImageEncoder",
+ },
+ {
+ # Example: ModelConfig(model_id="DiffSynth-Studio/General-Image-Encoders", origin_file_pattern="DINOv3-7B/model.safetensors")
+ "model_hash": "5722b5c873720009de96422993b15682",
+ "model_name": "dinov3_image_encoder",
+ "model_class": "diffsynth.models.dinov3_image_encoder.DINOv3ImageEncoder",
+ "state_dict_converter": "diffsynth.utils.state_dict_converters.dino_v3.DINOv3StateDictConverter",
+ },
+ {
+ # Example:
+ "model_hash": "a166c33455cdbd89c0888a3645ca5c0f",
+ "model_name": "qwen_image_image2lora_coarse",
+ "model_class": "diffsynth.models.qwen_image_image2lora.QwenImageImage2LoRAModel",
+ },
+ {
+ # Example:
+ "model_hash": "a5476e691767a4da6d3a6634a10f7408",
+ "model_name": "qwen_image_image2lora_fine",
+ "model_class": "diffsynth.models.qwen_image_image2lora.QwenImageImage2LoRAModel",
+ "extra_kwargs": {"residual_length": 37*37+7, "residual_mid_dim": 64}
+ },
+ {
+ # Example:
+ "model_hash": "0aad514690602ecaff932c701cb4b0bb",
+ "model_name": "qwen_image_image2lora_style",
+ "model_class": "diffsynth.models.qwen_image_image2lora.QwenImageImage2LoRAModel",
+ "extra_kwargs": {"compress_dim": 64, "use_residual": False}
+ },
+ {
+ # Example: ModelConfig(model_id="Qwen/Qwen-Image-Layered", origin_file_pattern="transformer/diffusion_pytorch_model*.safetensors")
+ "model_hash": "8dc8cda05de16c73afa755e2c1ce2839",
+ "model_name": "qwen_image_dit",
+ "model_class": "diffsynth.models.qwen_image_dit.QwenImageDiT",
+ "extra_kwargs": {"use_layer3d_rope": True, "use_additional_t_cond": True}
+ },
+ {
+ # Example: ModelConfig(model_id="Qwen/Qwen-Image-Layered", origin_file_pattern="vae/diffusion_pytorch_model.safetensors")
+ "model_hash": "44b39ddc499e027cfb24f7878d7416b9",
+ "model_name": "qwen_image_vae",
+ "model_class": "diffsynth.models.qwen_image_vae.QwenImageVAE",
+ "extra_kwargs": {"image_channels": 4}
+ },
+]
+
+qwen_video_edit_series = [
+ {
+ # Example: ModelConfig(model_id="yunpeng1998/Qwen-Video-Edit", origin_file_pattern="360P/step-30000.safetensors")
+ "model_hash": "8ae0ca4ab286d00197f08986c7fbbade",
+ "model_name": "qwen_video_edit_dit",
+ "model_class": "diffsynth.models.qwen_image_dit.QwenImageDiT",
+ "state_dict_converter": "diffsynth.utils.state_dict_converters.qwen_video_edit.QwenVideoEditDiTStateDictConverter",
+ },
+ {
+ # Example: ModelConfig(model_id="yunpeng1998/Qwen-Video-Edit", origin_file_pattern="360P/step-30000.safetensors")
+ "model_hash": "8ae0ca4ab286d00197f08986c7fbbade",
+ "model_name": "qwen_video_edit_adapter",
+ "model_class": "diffsynth.models.qwen_video_edit_dit.QwenVideoEditAdapter",
+ "state_dict_converter": "diffsynth.utils.state_dict_converters.qwen_video_edit.QwenVideoEditAdapterStateDictConverter",
+ },
+]
+
+wan_series = [
+ {
+ # Example: ModelConfig(model_id="Wan-AI/Wan2.2-Animate-2-14B", origin_file_pattern="wan_animate_2/wan_animate_2_bf16.safetensors")
+ "model_hash": "4536c21ad8740ba78367af4216ae85bf",
+ "model_name": "wan_video_dit",
+ "model_class": "diffsynth.models.wan_animate_2_dit.WanAnimate2Transformer",
+ "extra_kwargs": {},
+ },
+ {
+ # Example: ModelConfig(model_id="krea/krea-realtime-video", origin_file_pattern="krea-realtime-video-14b.safetensors")
+ "model_hash": "5ec04e02b42d2580483ad69f4e76346a",
+ "model_name": "wan_video_dit",
+ "model_class": "diffsynth.models.wan_video_dit.WanModel",
+ "extra_kwargs": {'has_image_input': False, 'patch_size': [1, 2, 2], 'in_dim': 16, 'dim': 5120, 'ffn_dim': 13824, 'freq_dim': 256, 'text_dim': 4096, 'out_dim': 16, 'num_heads': 40, 'num_layers': 40, 'eps': 1e-06},
+ "state_dict_converter": "diffsynth.utils.state_dict_converters.wan_video_dit.WanVideoDiTStateDictConverter",
+ },
+ {
+ # Example: ModelConfig(model_id="Wan-AI/Wan2.1-T2V-14B", origin_file_pattern="models_t5_umt5-xxl-enc-bf16.pth")
+ "model_hash": "9c8818c2cbea55eca56c7b447df170da",
+ "model_name": "wan_video_text_encoder",
+ "model_class": "diffsynth.models.wan_video_text_encoder.WanTextEncoder",
+ },
+ {
+ # Example: ModelConfig(model_id="Wan-AI/Wan2.1-T2V-14B", origin_file_pattern="Wan2.1_VAE.pth")
+ "model_hash": "ccc42284ea13e1ad04693284c7a09be6",
+ "model_name": "wan_video_vae",
+ "model_class": "diffsynth.models.wan_video_vae.WanVideoVAE",
+ "state_dict_converter": "diffsynth.utils.state_dict_converters.wan_video_vae.WanVideoVAEStateDictConverter",
+ },
+ {
+ # Example: ModelConfig(model_id="meituan-longcat/LongCat-Video", origin_file_pattern="dit/diffusion_pytorch_model*.safetensors")
+ "model_hash": "8b27900f680d7251ce44e2dc8ae1ffef",
+ "model_name": "wan_video_dit",
+ "model_class": "diffsynth.models.longcat_video_dit.LongCatVideoTransformer3DModel",
+ },
+ {
+ # Example: ModelConfig(model_id="ByteDance/Video-As-Prompt-Wan2.1-14B", origin_file_pattern="transformer/diffusion_pytorch_model*.safetensors")
+ "model_hash": "5f90e66a0672219f12d9a626c8c21f61",
+ "model_name": "wan_video_dit",
+ "model_class": "diffsynth.models.wan_video_dit.WanModel",
+ "extra_kwargs": {'has_image_input': True, 'patch_size': [1, 2, 2], 'in_dim': 36, 'dim': 5120, 'ffn_dim': 13824, 'freq_dim': 256, 'text_dim': 4096, 'out_dim': 16, 'num_heads': 40, 'num_layers': 40, 'eps': 1e-06},
+ "state_dict_converter": "diffsynth.utils.state_dict_converters.wan_video_dit.WanVideoDiTFromDiffusers"
+ },
+ {
+ # Example: ModelConfig(model_id="ByteDance/Video-As-Prompt-Wan2.1-14B", origin_file_pattern="transformer/diffusion_pytorch_model*.safetensors")
+ "model_hash": "5f90e66a0672219f12d9a626c8c21f61",
+ "model_name": "wan_video_vap",
+ "model_class": "diffsynth.models.wan_video_mot.MotWanModel",
+ "state_dict_converter": "diffsynth.utils.state_dict_converters.wan_video_mot.WanVideoMotStateDictConverter"
+ },
+ {
+ # Example: ModelConfig(model_id="Wan-AI/Wan2.1-I2V-14B-480P", origin_file_pattern="models_clip_open-clip-xlm-roberta-large-vit-huge-14.pth")
+ "model_hash": "5941c53e207d62f20f9025686193c40b",
+ "model_name": "wan_video_image_encoder",
+ "model_class": "diffsynth.models.wan_video_image_encoder.WanImageEncoder",
+ "state_dict_converter": "diffsynth.utils.state_dict_converters.wan_video_image_encoder.WanImageEncoderStateDictConverter"
+ },
+ {
+ # Example: ModelConfig(model_id="DiffSynth-Studio/Wan2.1-1.3b-speedcontrol-v1", origin_file_pattern="model.safetensors")
+ "model_hash": "dbd5ec76bbf977983f972c151d545389",
+ "model_name": "wan_video_motion_controller",
+ "model_class": "diffsynth.models.wan_video_motion_controller.WanMotionControllerModel",
+ },
+ {
+ # Example: ModelConfig(model_id="Wan-AI/Wan2.1-T2V-1.3B", origin_file_pattern="diffusion_pytorch_model*.safetensors")
+ "model_hash": "9269f8db9040a9d860eaca435be61814",
+ "model_name": "wan_video_dit",
+ "model_class": "diffsynth.models.wan_video_dit.WanModel",
+ "extra_kwargs": {'has_image_input': False, 'patch_size': [1, 2, 2], 'in_dim': 16, 'dim': 1536, 'ffn_dim': 8960, 'freq_dim': 256, 'text_dim': 4096, 'out_dim': 16, 'num_heads': 12, 'num_layers': 30, 'eps': 1e-06}
+ },
+ {
+ # Example: ModelConfig(model_id="Wan-AI/Wan2.1-FLF2V-14B-720P", origin_file_pattern="diffusion_pytorch_model*.safetensors")
+ "model_hash": "3ef3b1f8e1dab83d5b71fd7b617f859f",
+ "model_name": "wan_video_dit",
+ "model_class": "diffsynth.models.wan_video_dit.WanModel",
+ "extra_kwargs": {'has_image_input': True, 'patch_size': [1, 2, 2], 'in_dim': 36, 'dim': 5120, 'ffn_dim': 13824, 'freq_dim': 256, 'text_dim': 4096, 'out_dim': 16, 'num_heads': 40, 'num_layers': 40, 'eps': 1e-06, 'has_image_pos_emb': True}
+ },
+ {
+ # Example: ModelConfig(model_id="PAI/Wan2.1-Fun-1.3B-Control", origin_file_pattern="diffusion_pytorch_model*.safetensors")
+ "model_hash": "349723183fc063b2bfc10bb2835cf677",
+ "model_name": "wan_video_dit",
+ "model_class": "diffsynth.models.wan_video_dit.WanModel",
+ "extra_kwargs": {'has_image_input': True, 'patch_size': [1, 2, 2], 'in_dim': 48, 'dim': 1536, 'ffn_dim': 8960, 'freq_dim': 256, 'text_dim': 4096, 'out_dim': 16, 'num_heads': 12, 'num_layers': 30, 'eps': 1e-06}
+ },
+ {
+ # Example: ModelConfig(model_id="PAI/Wan2.1-Fun-1.3B-InP", origin_file_pattern="diffusion_pytorch_model*.safetensors")
+ "model_hash": "6d6ccde6845b95ad9114ab993d917893",
+ "model_name": "wan_video_dit",
+ "model_class": "diffsynth.models.wan_video_dit.WanModel",
+ "extra_kwargs": {'has_image_input': True, 'patch_size': [1, 2, 2], 'in_dim': 36, 'dim': 1536, 'ffn_dim': 8960, 'freq_dim': 256, 'text_dim': 4096, 'out_dim': 16, 'num_heads': 12, 'num_layers': 30, 'eps': 1e-06}
+ },
+ {
+ # Example: ModelConfig(model_id="PAI/Wan2.1-Fun-14B-Control", origin_file_pattern="diffusion_pytorch_model*.safetensors")
+ "model_hash": "efa44cddf936c70abd0ea28b6cbe946c",
+ "model_name": "wan_video_dit",
+ "model_class": "diffsynth.models.wan_video_dit.WanModel",
+ "extra_kwargs": {'has_image_input': True, 'patch_size': [1, 2, 2], 'in_dim': 48, 'dim': 5120, 'ffn_dim': 13824, 'freq_dim': 256, 'text_dim': 4096, 'out_dim': 16, 'num_heads': 40, 'num_layers': 40, 'eps': 1e-06}
+ },
+ {
+ # Example: ModelConfig(model_id="PAI/Wan2.1-Fun-14B-InP", origin_file_pattern="diffusion_pytorch_model*.safetensors")
+ "model_hash": "6bfcfb3b342cb286ce886889d519a77e",
+ "model_name": "wan_video_dit",
+ "model_class": "diffsynth.models.wan_video_dit.WanModel",
+ "extra_kwargs": {'has_image_input': True, 'patch_size': [1, 2, 2], 'in_dim': 36, 'dim': 5120, 'ffn_dim': 13824, 'freq_dim': 256, 'text_dim': 4096, 'out_dim': 16, 'num_heads': 40, 'num_layers': 40, 'eps': 1e-06}
+ },
+ {
+ # Example: ModelConfig(model_id="PAI/Wan2.1-Fun-V1.1-1.3B-Control-Camera", origin_file_pattern="diffusion_pytorch_model*.safetensors")
+ "model_hash": "ac6a5aa74f4a0aab6f64eb9a72f19901",
+ "model_name": "wan_video_dit",
+ "model_class": "diffsynth.models.wan_video_dit.WanModel",
+ "extra_kwargs": {'has_image_input': True, 'patch_size': [1, 2, 2], 'in_dim': 32, 'dim': 1536, 'ffn_dim': 8960, 'freq_dim': 256, 'text_dim': 4096, 'out_dim': 16, 'num_heads': 12, 'num_layers': 30, 'eps': 1e-06, 'has_ref_conv': False, 'add_control_adapter': True, 'in_dim_control_adapter': 24}
+ },
+ {
+ # Example: ModelConfig(model_id="PAI/Wan2.1-Fun-V1.1-1.3B-Control", origin_file_pattern="diffusion_pytorch_model*.safetensors")
+ "model_hash": "70ddad9d3a133785da5ea371aae09504",
+ "model_name": "wan_video_dit",
+ "model_class": "diffsynth.models.wan_video_dit.WanModel",
+ "extra_kwargs": {'has_image_input': True, 'patch_size': [1, 2, 2], 'in_dim': 48, 'dim': 1536, 'ffn_dim': 8960, 'freq_dim': 256, 'text_dim': 4096, 'out_dim': 16, 'num_heads': 12, 'num_layers': 30, 'eps': 1e-06, 'has_ref_conv': True}
+ },
+ {
+ # Example: ModelConfig(model_id="PAI/Wan2.1-Fun-V1.1-14B-Control-Camera", origin_file_pattern="diffusion_pytorch_model*.safetensors")
+ "model_hash": "b61c605c2adbd23124d152ed28e049ae",
+ "model_name": "wan_video_dit",
+ "model_class": "diffsynth.models.wan_video_dit.WanModel",
+ "extra_kwargs": {'has_image_input': True, 'patch_size': [1, 2, 2], 'in_dim': 32, 'dim': 5120, 'ffn_dim': 13824, 'freq_dim': 256, 'text_dim': 4096, 'out_dim': 16, 'num_heads': 40, 'num_layers': 40, 'eps': 1e-06, 'has_ref_conv': False, 'add_control_adapter': True, 'in_dim_control_adapter': 24}
+ },
+ {
+ # Example: ModelConfig(model_id="PAI/Wan2.1-Fun-V1.1-14B-Control", origin_file_pattern="diffusion_pytorch_model*.safetensors")
+ "model_hash": "26bde73488a92e64cc20b0a7485b9e5b",
+ "model_name": "wan_video_dit",
+ "model_class": "diffsynth.models.wan_video_dit.WanModel",
+ "extra_kwargs": {'has_image_input': True, 'patch_size': [1, 2, 2], 'in_dim': 48, 'dim': 5120, 'ffn_dim': 13824, 'freq_dim': 256, 'text_dim': 4096, 'out_dim': 16, 'num_heads': 40, 'num_layers': 40, 'eps': 1e-06, 'has_ref_conv': True}
+ },
+ {
+ # Example: ModelConfig(model_id="Wan-AI/Wan2.1-T2V-14B", origin_file_pattern="diffusion_pytorch_model*.safetensors")
+ "model_hash": "aafcfd9672c3a2456dc46e1cb6e52c70",
+ "model_name": "wan_video_dit",
+ "model_class": "diffsynth.models.wan_video_dit.WanModel",
+ "extra_kwargs": {'has_image_input': False, 'patch_size': [1, 2, 2], 'in_dim': 16, 'dim': 5120, 'ffn_dim': 13824, 'freq_dim': 256, 'text_dim': 4096, 'out_dim': 16, 'num_heads': 40, 'num_layers': 40, 'eps': 1e-06}
+ },
+ {
+ # Example: ModelConfig(model_id="iic/VACE-Wan2.1-1.3B-Preview", origin_file_pattern="diffusion_pytorch_model*.safetensors")
+ "model_hash": "a61453409b67cd3246cf0c3bebad47ba",
+ "model_name": "wan_video_dit",
+ "model_class": "diffsynth.models.wan_video_dit.WanModel",
+ "extra_kwargs": {'has_image_input': False, 'patch_size': [1, 2, 2], 'in_dim': 16, 'dim': 1536, 'ffn_dim': 8960, 'freq_dim': 256, 'text_dim': 4096, 'out_dim': 16, 'num_heads': 12, 'num_layers': 30, 'eps': 1e-06},
+ "state_dict_converter": "diffsynth.utils.state_dict_converters.wan_video_dit.WanVideoDiTStateDictConverter",
+ },
+ {
+ # Example: ModelConfig(model_id="iic/VACE-Wan2.1-1.3B-Preview", origin_file_pattern="diffusion_pytorch_model*.safetensors")
+ "model_hash": "a61453409b67cd3246cf0c3bebad47ba",
+ "model_name": "wan_video_vace",
+ "model_class": "diffsynth.models.wan_video_vace.VaceWanModel",
+ "state_dict_converter": "diffsynth.utils.state_dict_converters.wan_video_vace.VaceWanModelDictConverter"
+ },
+ {
+ # Example: ModelConfig(model_id="Wan-AI/Wan2.1-VACE-14B", origin_file_pattern="diffusion_pytorch_model*.safetensors")
+ "model_hash": "7a513e1f257a861512b1afd387a8ecd9",
+ "model_name": "wan_video_dit",
+ "model_class": "diffsynth.models.wan_video_dit.WanModel",
+ "extra_kwargs": {'has_image_input': False, 'patch_size': [1, 2, 2], 'in_dim': 16, 'dim': 5120, 'ffn_dim': 13824, 'freq_dim': 256, 'text_dim': 4096, 'out_dim': 16, 'num_heads': 40, 'num_layers': 40, 'eps': 1e-06},
+ "state_dict_converter": "diffsynth.utils.state_dict_converters.wan_video_dit.WanVideoDiTStateDictConverter",
+ },
+ {
+ # Example: ModelConfig(model_id="Wan-AI/Wan2.1-VACE-14B", origin_file_pattern="diffusion_pytorch_model*.safetensors")
+ "model_hash": "7a513e1f257a861512b1afd387a8ecd9",
+ "model_name": "wan_video_vace",
+ "model_class": "diffsynth.models.wan_video_vace.VaceWanModel",
+ "extra_kwargs": {'vace_layers': (0, 5, 10, 15, 20, 25, 30, 35), 'vace_in_dim': 96, 'patch_size': (1, 2, 2), 'has_image_input': False, 'dim': 5120, 'num_heads': 40, 'ffn_dim': 13824, 'eps': 1e-06},
+ "state_dict_converter": "diffsynth.utils.state_dict_converters.wan_video_vace.VaceWanModelDictConverter"
+ },
+ {
+ # Example: ModelConfig(model_id="Wan-AI/Wan2.2-Animate-14B", origin_file_pattern="diffusion_pytorch_model*.safetensors")
+ "model_hash": "31fa352acb8a1b1d33cd8764273d80a2",
+ "model_name": "wan_video_dit",
+ "model_class": "diffsynth.models.wan_video_dit.WanModel",
+ "extra_kwargs": {'has_image_input': True, 'patch_size': [1, 2, 2], 'in_dim': 36, 'dim': 5120, 'ffn_dim': 13824, 'freq_dim': 256, 'text_dim': 4096, 'out_dim': 16, 'num_heads': 40, 'num_layers': 40, 'eps': 1e-06},
+ "state_dict_converter": "diffsynth.utils.state_dict_converters.wan_video_dit.WanVideoDiTStateDictConverter"
+ },
+ {
+ # Example: ModelConfig(model_id="Wan-AI/Wan2.2-Animate-14B", origin_file_pattern="diffusion_pytorch_model*.safetensors")
+ "model_hash": "31fa352acb8a1b1d33cd8764273d80a2",
+ "model_name": "wan_video_animate_adapter",
+ "model_class": "diffsynth.models.wan_video_animate_adapter.WanAnimateAdapter",
+ "state_dict_converter": "diffsynth.utils.state_dict_converters.wan_video_animate_adapter.WanAnimateAdapterStateDictConverter"
+ },
+ {
+ # Example: ModelConfig(model_id="PAI/Wan2.2-Fun-A14B-Control-Camera", origin_file_pattern="high_noise_model/diffusion_pytorch_model*.safetensors")
+ "model_hash": "47dbeab5e560db3180adf51dc0232fb1",
+ "model_name": "wan_video_dit",
+ "model_class": "diffsynth.models.wan_video_dit.WanModel",
+ "extra_kwargs": {'has_image_input': False, 'patch_size': [1, 2, 2], 'in_dim': 36, 'dim': 5120, 'ffn_dim': 13824, 'freq_dim': 256, 'text_dim': 4096, 'out_dim': 16, 'num_heads': 40, 'num_layers': 40, 'eps': 1e-06, 'has_ref_conv': False, 'add_control_adapter': True, 'in_dim_control_adapter': 24, 'require_clip_embedding': False}
+ },
+ {
+ # Example: ModelConfig(model_id="PAI/Wan2.2-Fun-A14B-Control", origin_file_pattern="high_noise_model/diffusion_pytorch_model*.safetensors")
+ "model_hash": "2267d489f0ceb9f21836532952852ee5",
+ "model_name": "wan_video_dit",
+ "model_class": "diffsynth.models.wan_video_dit.WanModel",
+ "extra_kwargs": {'has_image_input': False, 'patch_size': [1, 2, 2], 'in_dim': 52, 'dim': 5120, 'ffn_dim': 13824, 'freq_dim': 256, 'text_dim': 4096, 'out_dim': 16, 'num_heads': 40, 'num_layers': 40, 'eps': 1e-06, 'has_ref_conv': True, 'require_clip_embedding': False},
+ },
+ {
+ # Example: ModelConfig(model_id="Wan-AI/Wan2.2-I2V-A14B", origin_file_pattern="high_noise_model/diffusion_pytorch_model*.safetensors")
+ "model_hash": "5b013604280dd715f8457c6ed6d6a626",
+ "model_name": "wan_video_dit",
+ "model_class": "diffsynth.models.wan_video_dit.WanModel",
+ "extra_kwargs": {'has_image_input': False, 'patch_size': [1, 2, 2], 'in_dim': 36, 'dim': 5120, 'ffn_dim': 13824, 'freq_dim': 256, 'text_dim': 4096, 'out_dim': 16, 'num_heads': 40, 'num_layers': 40, 'eps': 1e-06, 'require_clip_embedding': False}
+ },
+ {
+ # Example: ModelConfig(model_id="Wan-AI/Wan2.2-S2V-14B", origin_file_pattern="diffusion_pytorch_model*.safetensors")
+ "model_hash": "966cffdcc52f9c46c391768b27637614",
+ "model_name": "wan_video_dit",
+ "model_class": "diffsynth.models.wan_video_dit_s2v.WanS2VModel",
+ "extra_kwargs": {'dim': 5120, 'in_dim': 16, 'ffn_dim': 13824, 'out_dim': 16, 'text_dim': 4096, 'freq_dim': 256, 'eps': 1e-06, 'patch_size': (1, 2, 2), 'num_heads': 40, 'num_layers': 40, 'cond_dim': 16, 'audio_dim': 1024, 'num_audio_token': 4}
+ },
+ {
+ # Example: ModelConfig(model_id="Wan-AI/Wan2.2-TI2V-5B", origin_file_pattern="diffusion_pytorch_model*.safetensors")
+ "model_hash": "1f5ab7703c6fc803fdded85ff040c316",
+ "model_name": "wan_video_dit",
+ "model_class": "diffsynth.models.wan_video_dit.WanModel",
+ "extra_kwargs": {'has_image_input': False, 'patch_size': [1, 2, 2], 'in_dim': 48, 'dim': 3072, 'ffn_dim': 14336, 'freq_dim': 256, 'text_dim': 4096, 'out_dim': 48, 'num_heads': 24, 'num_layers': 30, 'eps': 1e-06, 'seperated_timestep': True, 'require_clip_embedding': False, 'require_vae_embedding': False, 'fuse_vae_embedding_in_latents': True}
+ },
+ {
+ # Example: ModelConfig(model_id="Wan-AI/Wan2.2-TI2V-5B", origin_file_pattern="Wan2.2_VAE.pth")
+ "model_hash": "e1de6c02cdac79f8b739f4d3698cd216",
+ "model_name": "wan_video_vae",
+ "model_class": "diffsynth.models.wan_video_vae.WanVideoVAE38",
+ "state_dict_converter": "diffsynth.utils.state_dict_converters.wan_video_vae.WanVideoVAEStateDictConverter",
+ },
+ {
+ # Example: ModelConfig(model_id="Wan-AI/Wan2.2-S2V-14B", origin_file_pattern="wav2vec2-large-xlsr-53-english/model.safetensors")
+ "model_hash": "06be60f3a4526586d8431cd038a71486",
+ "model_name": "wans2v_audio_encoder",
+ "model_class": "diffsynth.models.wav2vec.WanS2VAudioEncoder",
+ "state_dict_converter": "diffsynth.utils.state_dict_converters.wans2v_audio_encoder.WanS2VAudioEncoderStateDictConverter",
+ },
+ {
+ # Example: ModelConfig(model_id="Wan-AI/Wan-Dancer-14B", origin_file_pattern="global_model.safetensors")
+ "model_hash": "eb18873fc0ba77b541eb7b62dbcd2059",
+ "model_name": "wan_video_dit",
+ "model_class": "diffsynth.models.wan_video_dit.WanModel",
+ "extra_kwargs": {'has_image_input': True, 'patch_size': [1, 2, 2], 'in_dim': 36, 'dim': 5120, 'ffn_dim': 13824, 'freq_dim': 256, 'text_dim': 4096, 'out_dim': 16, 'num_heads': 40, 'num_layers': 40, 'eps': 1e-06, 'wantodance_enable_music_inject': True, 'wantodance_music_inject_layers': [0, 4, 8, 12, 16, 20, 24, 27], 'wantodance_enable_refimage': True, 'has_ref_conv': True, 'wantodance_enable_refface': False, 'wantodance_enable_global': True, 'wantodance_enable_dynamicfps': True, 'wantodance_enable_unimodel': True}
+ },
+]
+
+flux_series = [
+ {
+ # Example: ModelConfig(model_id="black-forest-labs/FLUX.1-dev", origin_file_pattern="flux1-dev.safetensors")
+ "model_hash": "a29710fea6dddb0314663ee823598e50",
+ "model_name": "flux_dit",
+ "model_class": "diffsynth.models.flux_dit.FluxDiT",
+ "state_dict_converter": "diffsynth.utils.state_dict_converters.flux_dit.FluxDiTStateDictConverter",
+ },
+ {
+ # Supported due to historical reasons.
+ "model_hash": "605c56eab23e9e2af863ad8f0813a25d",
+ "model_name": "flux_dit",
+ "model_class": "diffsynth.models.flux_dit.FluxDiT",
+ "state_dict_converter": "diffsynth.utils.state_dict_converters.flux_dit.FluxDiTStateDictConverterFromDiffusers",
+ },
+ {
+ # Example: ModelConfig(model_id="black-forest-labs/FLUX.1-Fill-dev", origin_file_pattern="flux1-fill-dev.safetensors")
+ "model_hash": "f876d2a6fb247fb7da0f84fc7fb8a823",
+ "model_name": "flux_dit",
+ "model_class": "diffsynth.models.flux_dit.FluxDiT",
+ "extra_kwargs": {"input_dim": 384},
+ "state_dict_converter": "diffsynth.utils.state_dict_converters.flux_dit.FluxDiTStateDictConverter",
+ },
+ {
+ # Example: ModelConfig(model_id="black-forest-labs/FLUX.1-dev", origin_file_pattern="text_encoder/model.safetensors")
+ "model_hash": "94eefa3dac9cec93cb1ebaf1747d7b78",
+ "model_name": "flux_text_encoder_clip",
+ "model_class": "diffsynth.models.flux_text_encoder_clip.FluxTextEncoderClip",
+ "state_dict_converter": "diffsynth.utils.state_dict_converters.flux_text_encoder_clip.FluxTextEncoderClipStateDictConverter",
+ },
+ {
+ # Example: ModelConfig(model_id="black-forest-labs/FLUX.1-dev", origin_file_pattern="text_encoder_2/*.safetensors")
+ "model_hash": "22540b49eaedbc2f2784b2091a234c7c",
+ "model_name": "flux_text_encoder_t5",
+ "model_class": "diffsynth.models.flux_text_encoder_t5.FluxTextEncoderT5",
+ "state_dict_converter": "diffsynth.utils.state_dict_converters.flux_text_encoder_t5.FluxTextEncoderT5StateDictConverter",
+ },
+ {
+ # Example: ModelConfig(model_id="black-forest-labs/FLUX.1-dev", origin_file_pattern="ae.safetensors")
+ "model_hash": "21ea55f476dfc4fd135587abb59dfe5d",
+ "model_name": "flux_vae_encoder",
+ "model_class": "diffsynth.models.flux_vae.FluxVAEEncoder",
+ "state_dict_converter": "diffsynth.utils.state_dict_converters.flux_vae.FluxVAEEncoderStateDictConverter",
+ },
+ {
+ # Example: ModelConfig(model_id="black-forest-labs/FLUX.1-dev", origin_file_pattern="ae.safetensors")
+ "model_hash": "21ea55f476dfc4fd135587abb59dfe5d",
+ "model_name": "flux_vae_decoder",
+ "model_class": "diffsynth.models.flux_vae.FluxVAEDecoder",
+ "state_dict_converter": "diffsynth.utils.state_dict_converters.flux_vae.FluxVAEDecoderStateDictConverter",
+ },
+ {
+ # Example: ModelConfig(model_id="ostris/Flex.2-preview", origin_file_pattern="Flex.2-preview.safetensors")
+ "model_hash": "d02f41c13549fa5093d3521f62a5570a",
+ "model_name": "flux_dit",
+ "model_class": "diffsynth.models.flux_dit.FluxDiT",
+ "extra_kwargs": {'input_dim': 196, 'num_blocks': 8},
+ "state_dict_converter": "diffsynth.utils.state_dict_converters.flux_dit.FluxDiTStateDictConverter",
+ },
+ {
+ # Example: ModelConfig(model_id="DiffSynth-Studio/AttriCtrl-FLUX.1-Dev", origin_file_pattern="models/brightness.safetensors")
+ "model_hash": "0629116fce1472503a66992f96f3eb1a",
+ "model_name": "flux_value_controller",
+ "model_class": "diffsynth.models.flux_value_control.SingleValueEncoder",
+ },
+ {
+ # Example: ModelConfig(model_id="alimama-creative/FLUX.1-dev-Controlnet-Inpainting-Beta", origin_file_pattern="diffusion_pytorch_model.safetensors")
+ "model_hash": "52357cb26250681367488a8954c271e8",
+ "model_name": "flux_controlnet",
+ "model_class": "diffsynth.models.flux_controlnet.FluxControlNet",
+ "state_dict_converter": "diffsynth.utils.state_dict_converters.flux_controlnet.FluxControlNetStateDictConverter",
+ "extra_kwargs": {"num_joint_blocks": 6, "num_single_blocks": 0, "additional_input_dim": 4},
+ },
+ {
+ # Example: ModelConfig(model_id="InstantX/FLUX.1-dev-Controlnet-Union-alpha", origin_file_pattern="diffusion_pytorch_model.safetensors")
+ "model_hash": "78d18b9101345ff695f312e7e62538c0",
+ "model_name": "flux_controlnet",
+ "model_class": "diffsynth.models.flux_controlnet.FluxControlNet",
+ "state_dict_converter": "diffsynth.utils.state_dict_converters.flux_controlnet.FluxControlNetStateDictConverter",
+ "extra_kwargs": {"num_mode": 10, "mode_dict": {"canny": 0, "tile": 1, "depth": 2, "blur": 3, "pose": 4, "gray": 5, "lq": 6}},
+ },
+ {
+ # Example: ModelConfig(model_id="jasperai/Flux.1-dev-Controlnet-Upscaler", origin_file_pattern="diffusion_pytorch_model.safetensors")
+ "model_hash": "b001c89139b5f053c715fe772362dd2a",
+ "model_name": "flux_controlnet",
+ "model_class": "diffsynth.models.flux_controlnet.FluxControlNet",
+ "state_dict_converter": "diffsynth.utils.state_dict_converters.flux_controlnet.FluxControlNetStateDictConverter",
+ "extra_kwargs": {"num_single_blocks": 0},
+ },
+ {
+ # Example: ModelConfig(model_id="ByteDance/InfiniteYou", origin_file_pattern="infu_flux_v1.0/aes_stage2/image_proj_model.bin")
+ "model_hash": "c07c0f04f5ff55e86b4e937c7a40d481",
+ "model_name": "infiniteyou_image_projector",
+ "model_class": "diffsynth.models.flux_infiniteyou.InfiniteYouImageProjector",
+ "state_dict_converter": "diffsynth.utils.state_dict_converters.flux_infiniteyou.FluxInfiniteYouImageProjectorStateDictConverter",
+ },
+ {
+ # Example: ModelConfig(model_id="ByteDance/InfiniteYou", origin_file_pattern="infu_flux_v1.0/aes_stage2/InfuseNetModel/*.safetensors")
+ "model_hash": "7f9583eb8ba86642abb9a21a4b2c9e16",
+ "model_name": "flux_controlnet",
+ "model_class": "diffsynth.models.flux_controlnet.FluxControlNet",
+ "state_dict_converter": "diffsynth.utils.state_dict_converters.flux_controlnet.FluxControlNetStateDictConverter",
+ "extra_kwargs": {"num_joint_blocks": 4, "num_single_blocks": 10},
+ },
+ {
+ # Example: ModelConfig(model_id="DiffSynth-Studio/LoRA-Encoder-FLUX.1-Dev", origin_file_pattern="model.safetensors")
+ "model_hash": "77c2e4dd2440269eb33bfaa0d004f6ab",
+ "model_name": "flux_lora_encoder",
+ "model_class": "diffsynth.models.flux_lora_encoder.FluxLoRAEncoder",
+ },
+ {
+ # Example: ModelConfig(model_id="DiffSynth-Studio/LoRAFusion-preview-FLUX.1-dev", origin_file_pattern="model.safetensors")
+ "model_hash": "30143afb2dea73d1ac580e0787628f8c",
+ "model_name": "flux_lora_patcher",
+ "model_class": "diffsynth.models.flux_lora_patcher.FluxLoraPatcher",
+ },
+ {
+ # Example: ModelConfig(model_id="DiffSynth-Studio/Nexus-GenV2", origin_file_pattern="model*.safetensors")
+ "model_hash": "2bd19e845116e4f875a0a048e27fc219",
+ "model_name": "nexus_gen_llm",
+ "model_class": "diffsynth.models.nexus_gen.NexusGenAutoregressiveModel",
+ "state_dict_converter": "diffsynth.utils.state_dict_converters.nexus_gen.NexusGenAutoregressiveModelStateDictConverter",
+ },
+ {
+ # Example: ModelConfig(model_id="DiffSynth-Studio/Nexus-GenV2", origin_file_pattern="edit_decoder.bin")
+ "model_hash": "63c969fd37cce769a90aa781fbff5f81",
+ "model_name": "nexus_gen_editing_adapter",
+ "model_class": "diffsynth.models.nexus_gen_projector.NexusGenImageEmbeddingMerger",
+ "state_dict_converter": "diffsynth.utils.state_dict_converters.nexus_gen_projector.NexusGenMergerStateDictConverter",
+ },
+ {
+ # Example: ModelConfig(model_id="DiffSynth-Studio/Nexus-GenV2", origin_file_pattern="edit_decoder.bin")
+ "model_hash": "63c969fd37cce769a90aa781fbff5f81",
+ "model_name": "flux_dit",
+ "model_class": "diffsynth.models.flux_dit.FluxDiT",
+ "state_dict_converter": "diffsynth.utils.state_dict_converters.flux_dit.FluxDiTStateDictConverter",
+ },
+ {
+ # Example: ModelConfig(model_id="DiffSynth-Studio/Nexus-GenV2", origin_file_pattern="generation_decoder.bin")
+ "model_hash": "3e6c61b0f9471135fc9c6d6a98e98b6d",
+ "model_name": "nexus_gen_generation_adapter",
+ "model_class": "diffsynth.models.nexus_gen_projector.NexusGenAdapter",
+ "state_dict_converter": "diffsynth.utils.state_dict_converters.nexus_gen_projector.NexusGenAdapterStateDictConverter",
+ },
+ {
+ # Example: ModelConfig(model_id="DiffSynth-Studio/Nexus-GenV2", origin_file_pattern="generation_decoder.bin")
+ "model_hash": "3e6c61b0f9471135fc9c6d6a98e98b6d",
+ "model_name": "flux_dit",
+ "model_class": "diffsynth.models.flux_dit.FluxDiT",
+ "state_dict_converter": "diffsynth.utils.state_dict_converters.flux_dit.FluxDiTStateDictConverter",
+ },
+ {
+ # Example: ModelConfig(model_id="InstantX/FLUX.1-dev-IP-Adapter", origin_file_pattern="ip-adapter.bin")
+ "model_hash": "4daaa66cc656a8fe369908693dad0a35",
+ "model_name": "flux_ipadapter",
+ "model_class": "diffsynth.models.flux_ipadapter.FluxIpAdapter",
+ "state_dict_converter": "diffsynth.utils.state_dict_converters.flux_ipadapter.FluxIpAdapterStateDictConverter",
+ },
+ {
+ # Example: ModelConfig(model_id="google/siglip-so400m-patch14-384", origin_file_pattern="model.safetensors")
+ "model_hash": "04d8c1e20a1f1b25f7434f111992a33f",
+ "model_name": "siglip_vision_model",
+ "model_class": "diffsynth.models.flux_ipadapter.SiglipVisionModelSO400M",
+ "state_dict_converter": "diffsynth.utils.state_dict_converters.flux_ipadapter.SiglipStateDictConverter",
+ },
+ {
+ # Example: ModelConfig(model_id="black-forest-labs/FLUX.1-Redux-dev", origin_file_pattern="image_encoder/model.safetensors")
+ "model_hash": "6bd5343828144a2135d783519c665e2c",
+ "model_name": "flux_redux_image_encoder",
+ "model_class": "diffsynth.models.flux_ipadapter.SiglipVisionModelSO400M",
+ "state_dict_converter": "diffsynth.utils.state_dict_converters.flux_ipadapter.SiglipReduxStateDictConverter",
+ },
+ {
+ # Example: ModelConfig(model_id="black-forest-labs/FLUX.1-Redux-dev", origin_file_pattern="image_embedder/diffusion_pytorch_model.safetensors")
+ # 2-layer MLP projecting SigLIP tokens (1152) to the 4096-dim FLUX text
+ # space; keys already match FluxReduxImageEncoder, so no converter needed.
+ "model_hash": "d8ea4618391460bc831fddaae5fb5288",
+ "model_name": "flux_redux_image_embedder",
+ "model_class": "diffsynth.models.flux_redux.FluxReduxImageEncoder",
+ },
+ {
+ # Example: ModelConfig(model_id="stepfun-ai/Step1X-Edit", origin_file_pattern="step1x-edit-i1258.safetensors"),
+ "model_hash": "d30fb9e02b1dbf4e509142f05cf7dd50",
+ "model_name": "step1x_connector",
+ "model_class": "diffsynth.models.step1x_connector.Qwen2Connector",
+ "state_dict_converter": "diffsynth.utils.state_dict_converters.step1x_connector.Qwen2ConnectorStateDictConverter",
+ },
+ {
+ # Example: ModelConfig(model_id="stepfun-ai/Step1X-Edit", origin_file_pattern="step1x-edit-i1258.safetensors"),
+ "model_hash": "d30fb9e02b1dbf4e509142f05cf7dd50",
+ "model_name": "flux_dit",
+ "model_class": "diffsynth.models.flux_dit.FluxDiT",
+ "state_dict_converter": "diffsynth.utils.state_dict_converters.flux_dit.FluxDiTStateDictConverter",
+ "extra_kwargs": {"disable_guidance_embedder": True},
+ },
+ {
+ # Example: ModelConfig(model_id="MAILAND/majicflus_v1", origin_file_pattern="majicflus_v134.safetensors")
+ "model_hash": "3394f306c4cbf04334b712bf5aaed95f",
+ "model_name": "flux_dit",
+ "model_class": "diffsynth.models.flux_dit.FluxDiT",
+ "state_dict_converter": "diffsynth.utils.state_dict_converters.flux_dit.FluxDiTStateDictConverter",
+ },
+]
+
+flux2_series = [
+ {
+ # Example: ModelConfig(model_id="black-forest-labs/FLUX.2-dev", origin_file_pattern="text_encoder/*.safetensors")
+ "model_hash": "28fca3d8e5bf2a2d1271748a773f6757",
+ "model_name": "flux2_text_encoder",
+ "model_class": "diffsynth.models.flux2_text_encoder.Flux2TextEncoder",
+ "state_dict_converter": "diffsynth.utils.state_dict_converters.flux2_text_encoder.Flux2TextEncoderStateDictConverter",
+ },
+ {
+ # Example: ModelConfig(model_id="black-forest-labs/FLUX.2-dev", origin_file_pattern="transformer/*.safetensors")
+ "model_hash": "d38e1d5c5aec3b0a11e79327ac6e3b0f",
+ "model_name": "flux2_dit",
+ "model_class": "diffsynth.models.flux2_dit.Flux2DiT",
+ },
+ {
+ # Example: ModelConfig(model_id="black-forest-labs/FLUX.2-dev", origin_file_pattern="vae/diffusion_pytorch_model.safetensors")
+ "model_hash": "c54288e3ee12ca215898840682337b95",
+ "model_name": "flux2_vae",
+ "model_class": "diffsynth.models.flux2_vae.Flux2VAE",
+ },
+ {
+ # Example: ModelConfig(model_id="black-forest-labs/FLUX.2-klein-4B", origin_file_pattern="transformer/*.safetensors")
+ "model_hash": "3bde7b817fec8143028b6825a63180df",
+ "model_name": "flux2_dit",
+ "model_class": "diffsynth.models.flux2_dit.Flux2DiT",
+ "extra_kwargs": {"guidance_embeds": False, "joint_attention_dim": 7680, "num_attention_heads": 24, "num_layers": 5, "num_single_layers": 20}
+ },
+ {
+ # Example: ModelConfig(model_id="black-forest-labs/FLUX.2-klein-9B", origin_file_pattern="text_encoder/*.safetensors")
+ "model_hash": "9195f3ea256fcd0ae6d929c203470754",
+ "model_name": "z_image_text_encoder",
+ "model_class": "diffsynth.models.z_image_text_encoder.ZImageTextEncoder",
+ "extra_kwargs": {"model_size": "8B"},
+ "state_dict_converter": "diffsynth.utils.state_dict_converters.z_image_text_encoder.ZImageTextEncoderStateDictConverter",
+ },
+ {
+ # Example: ModelConfig(model_id="black-forest-labs/FLUX.2-klein-9B", origin_file_pattern="transformer/*.safetensors")
+ "model_hash": "39c6fc48f07bebecedbbaa971ff466c8",
+ "model_name": "flux2_dit",
+ "model_class": "diffsynth.models.flux2_dit.Flux2DiT",
+ "extra_kwargs": {"guidance_embeds": False, "joint_attention_dim": 12288, "num_attention_heads": 32, "num_layers": 8, "num_single_layers": 24}
+ },
+]
+
+ernie_image_series = [
+ {
+ # Example: ModelConfig(model_id="PaddlePaddle/ERNIE-Image", origin_file_pattern="transformer/diffusion_pytorch_model*.safetensors")
+ "model_hash": "584c13713849f1af4e03d5f1858b8b7b",
+ "model_name": "ernie_image_dit",
+ "model_class": "diffsynth.models.ernie_image_dit.ErnieImageDiT",
+ },
+ {
+ # Example: ModelConfig(model_id="PaddlePaddle/ERNIE-Image", origin_file_pattern="text_encoder/model.safetensors")
+ "model_hash": "404ed9f40796a38dd34c1620f1920207",
+ "model_name": "ernie_image_text_encoder",
+ "model_class": "diffsynth.models.ernie_image_text_encoder.ErnieImageTextEncoder",
+ "state_dict_converter": "diffsynth.utils.state_dict_converters.ernie_image_text_encoder.ErnieImageTextEncoderStateDictConverter",
+ },
+]
+
+z_image_series = [
+ {
+ # Example: ModelConfig(model_id="Tongyi-MAI/Z-Image-Turbo", origin_file_pattern="transformer/*.safetensors")
+ "model_hash": "fc3a8a1247fe185ce116ccbe0e426c28",
+ "model_name": "z_image_dit",
+ "model_class": "diffsynth.models.z_image_dit.ZImageDiT",
+ },
+ {
+ # Example: ModelConfig(model_id="Tongyi-MAI/Z-Image-Turbo", origin_file_pattern="text_encoder/*.safetensors")
+ "model_hash": "0f050f62a88876fea6eae0a18dac5a2e",
+ "model_name": "z_image_text_encoder",
+ "model_class": "diffsynth.models.z_image_text_encoder.ZImageTextEncoder",
+ },
+ {
+ # Example: ModelConfig(model_id="Tongyi-MAI/Z-Image-Turbo", origin_file_pattern="vae/vae/diffusion_pytorch_model.safetensors")
+ "model_hash": "1aafa3cc91716fb6b300cc1cd51b85a3",
+ "model_name": "flux_vae_encoder",
+ "model_class": "diffsynth.models.flux_vae.FluxVAEEncoder",
+ "state_dict_converter": "diffsynth.utils.state_dict_converters.flux_vae.FluxVAEEncoderStateDictConverterDiffusers",
+ "extra_kwargs": {"use_conv_attention": False},
+ },
+ {
+ # Example: ModelConfig(model_id="Tongyi-MAI/Z-Image-Turbo", origin_file_pattern="vae/vae/diffusion_pytorch_model.safetensors")
+ "model_hash": "1aafa3cc91716fb6b300cc1cd51b85a3",
+ "model_name": "flux_vae_decoder",
+ "model_class": "diffsynth.models.flux_vae.FluxVAEDecoder",
+ "state_dict_converter": "diffsynth.utils.state_dict_converters.flux_vae.FluxVAEDecoderStateDictConverterDiffusers",
+ "extra_kwargs": {"use_conv_attention": False},
+ },
+ {
+ # Example: ModelConfig(model_id="Tongyi-MAI/Z-Image-Omni-Base", origin_file_pattern="transformer/*.safetensors")
+ "model_hash": "aa3563718e5c3ecde3dfbb020ca61180",
+ "model_name": "z_image_dit",
+ "model_class": "diffsynth.models.z_image_dit.ZImageDiT",
+ "extra_kwargs": {"siglip_feat_dim": 1152},
+ },
+ {
+ # Example: ModelConfig(model_id="Tongyi-MAI/Z-Image-Omni-Base", origin_file_pattern="siglip/model.safetensors")
+ "model_hash": "89d48e420f45cff95115a9f3e698d44a",
+ "model_name": "siglip_vision_model_428m",
+ "model_class": "diffsynth.models.siglip2_image_encoder.Siglip2ImageEncoder428M",
+ },
+ {
+ # Example: ModelConfig(model_id="PAI/Z-Image-Turbo-Fun-Controlnet-Union-2.1", origin_file_pattern="Z-Image-Turbo-Fun-Controlnet-Union-2.1-8steps.safetensors")
+ "model_hash": "1677708d40029ab380a95f6c731a57d7",
+ "model_name": "z_image_controlnet",
+ "model_class": "diffsynth.models.z_image_controlnet.ZImageControlNet",
+ },
+ {
+ # Example: ???
+ "model_hash": "9510cb8cd1dd34ee0e4f111c24905510",
+ "model_name": "z_image_image2lora_style",
+ "model_class": "diffsynth.models.z_image_image2lora.ZImageImage2LoRAModel",
+ "extra_kwargs": {"compress_dim": 128},
+ },
+ {
+ # Example: ModelConfig(model_id="Qwen/Qwen3-0.6B", origin_file_pattern="model.safetensors")
+ "model_hash": "1392adecee344136041e70553f875f31",
+ "model_name": "z_image_text_encoder",
+ "model_class": "diffsynth.models.z_image_text_encoder.ZImageTextEncoder",
+ "extra_kwargs": {"model_size": "0.6B"},
+ "state_dict_converter": "diffsynth.utils.state_dict_converters.z_image_text_encoder.ZImageTextEncoderStateDictConverter",
+ },
+ {
+ # To ensure compatibility with the `model.diffusion_model` prefix introduced by other frameworks.
+ "model_hash": "8cf241a0d32f93d5de368502a086852f",
+ "model_name": "z_image_dit",
+ "model_class": "diffsynth.models.z_image_dit.ZImageDiT",
+ "state_dict_converter": "diffsynth.utils.state_dict_converters.z_image_dit.ZImageDiTStateDictConverter",
+ },
+]
+"""
+Offical model repo: https://www.modelscope.cn/models/Lightricks/LTX-2
+Repackaged model repo: https://www.modelscope.cn/models/DiffSynth-Studio/LTX-2-Repackage
+For base models of LTX-2, offical checkpoint (with model config ModelConfig(model_id="Lightricks/LTX-2", origin_file_pattern="ltx-2-19b-dev.safetensors"))
+and repackaged checkpoints (with model config ModelConfig(model_id="DiffSynth-Studio/LTX-2-Repackage", origin_file_pattern="*.safetensors")) are both supported.
+We have repackeged the official checkpoints in DiffSynth-Studio/LTX-2-Repackage repo to support separate loading of different submodules,
+and avoid redundant memory usage when users only want to use part of the model.
+"""
+ltx2_series = [
+ {
+ # Example: ModelConfig(model_id="Lightricks/LTX-2", origin_file_pattern="ltx-2-19b-dev.safetensors")
+ "model_hash": "aca7b0bbf8415e9c98360750268915fc",
+ "model_name": "ltx2_dit",
+ "model_class": "diffsynth.models.ltx2_dit.LTXModel",
+ "state_dict_converter": "diffsynth.utils.state_dict_converters.ltx2_dit.LTXModelStateDictConverter",
+ },
+ {
+ # Example: ModelConfig(model_id="DiffSynth-Studio/LTX-2-Repackage", origin_file_pattern="transformer.safetensors")
+ "model_hash": "c567aaa37d5ed7454c73aa6024458661",
+ "model_name": "ltx2_dit",
+ "model_class": "diffsynth.models.ltx2_dit.LTXModel",
+ "state_dict_converter": "diffsynth.utils.state_dict_converters.ltx2_dit.LTXModelStateDictConverter",
+ },
+ {
+ # Example: ModelConfig(model_id="Lightricks/LTX-2", origin_file_pattern="ltx-2-19b-dev.safetensors")
+ "model_hash": "aca7b0bbf8415e9c98360750268915fc",
+ "model_name": "ltx2_video_vae_encoder",
+ "model_class": "diffsynth.models.ltx2_video_vae.LTX2VideoEncoder",
+ "state_dict_converter": "diffsynth.utils.state_dict_converters.ltx2_video_vae.LTX2VideoEncoderStateDictConverter",
+ },
+ {
+ # Example: ModelConfig(model_id="DiffSynth-Studio/LTX-2-Repackage", origin_file_pattern="video_vae_encoder.safetensors")
+ "model_hash": "7f7e904a53260ec0351b05f32153754b",
+ "model_name": "ltx2_video_vae_encoder",
+ "model_class": "diffsynth.models.ltx2_video_vae.LTX2VideoEncoder",
+ "state_dict_converter": "diffsynth.utils.state_dict_converters.ltx2_video_vae.LTX2VideoEncoderStateDictConverter",
+ },
+ {
+ # Example: ModelConfig(model_id="Lightricks/LTX-2", origin_file_pattern="ltx-2-19b-dev.safetensors")
+ "model_hash": "aca7b0bbf8415e9c98360750268915fc",
+ "model_name": "ltx2_video_vae_decoder",
+ "model_class": "diffsynth.models.ltx2_video_vae.LTX2VideoDecoder",
+ "state_dict_converter": "diffsynth.utils.state_dict_converters.ltx2_video_vae.LTX2VideoDecoderStateDictConverter",
+ },
+ {
+ # Example: ModelConfig(model_id="DiffSynth-Studio/LTX-2-Repackage", origin_file_pattern="video_vae_decoder.safetensors")
+ "model_hash": "dc6029ca2825147872b45e35a2dc3a97",
+ "model_name": "ltx2_video_vae_decoder",
+ "model_class": "diffsynth.models.ltx2_video_vae.LTX2VideoDecoder",
+ "state_dict_converter": "diffsynth.utils.state_dict_converters.ltx2_video_vae.LTX2VideoDecoderStateDictConverter",
+ },
+ {
+ # Example: ModelConfig(model_id="Lightricks/LTX-2", origin_file_pattern="ltx-2-19b-dev.safetensors")
+ "model_hash": "aca7b0bbf8415e9c98360750268915fc",
+ "model_name": "ltx2_audio_vae_decoder",
+ "model_class": "diffsynth.models.ltx2_audio_vae.LTX2AudioDecoder",
+ "state_dict_converter": "diffsynth.utils.state_dict_converters.ltx2_audio_vae.LTX2AudioDecoderStateDictConverter",
+ },
+ {
+ # Example: ModelConfig(model_id="DiffSynth-Studio/LTX-2-Repackage", origin_file_pattern="audio_vae_decoder.safetensors")
+ "model_hash": "7d7823dde8f1ea0b50fb07ac329dd4cb",
+ "model_name": "ltx2_audio_vae_decoder",
+ "model_class": "diffsynth.models.ltx2_audio_vae.LTX2AudioDecoder",
+ "state_dict_converter": "diffsynth.utils.state_dict_converters.ltx2_audio_vae.LTX2AudioDecoderStateDictConverter",
+ },
+ {
+ # Example: ModelConfig(model_id="Lightricks/LTX-2", origin_file_pattern="ltx-2-19b-dev.safetensors")
+ "model_hash": "aca7b0bbf8415e9c98360750268915fc",
+ "model_name": "ltx2_audio_vocoder",
+ "model_class": "diffsynth.models.ltx2_audio_vae.LTX2Vocoder",
+ "state_dict_converter": "diffsynth.utils.state_dict_converters.ltx2_audio_vae.LTX2VocoderStateDictConverter",
+ },
+ {
+ # Example: ModelConfig(model_id="DiffSynth-Studio/LTX-2-Repackage", origin_file_pattern="audio_vocoder.safetensors")
+ "model_hash": "f471360f6b24bef702ab73133d9f8bb9",
+ "model_name": "ltx2_audio_vocoder",
+ "model_class": "diffsynth.models.ltx2_audio_vae.LTX2Vocoder",
+ "state_dict_converter": "diffsynth.utils.state_dict_converters.ltx2_audio_vae.LTX2VocoderStateDictConverter",
+ },
+ {
+ # Example: ModelConfig(model_id="Lightricks/LTX-2", origin_file_pattern="ltx-2-19b-dev.safetensors")
+ "model_hash": "aca7b0bbf8415e9c98360750268915fc",
+ "model_name": "ltx2_audio_vae_encoder",
+ "model_class": "diffsynth.models.ltx2_audio_vae.LTX2AudioEncoder",
+ "state_dict_converter": "diffsynth.utils.state_dict_converters.ltx2_audio_vae.LTX2AudioEncoderStateDictConverter",
+ },
+ {
+ # Example: ModelConfig(model_id="DiffSynth-Studio/LTX-2-Repackage", origin_file_pattern="audio_vae_encoder.safetensors")
+ "model_hash": "29338f3b95e7e312a3460a482e4f4554",
+ "model_name": "ltx2_audio_vae_encoder",
+ "model_class": "diffsynth.models.ltx2_audio_vae.LTX2AudioEncoder",
+ "state_dict_converter": "diffsynth.utils.state_dict_converters.ltx2_audio_vae.LTX2AudioEncoderStateDictConverter",
+ },
+ {
+ # Example: ModelConfig(model_id="Lightricks/LTX-2", origin_file_pattern="ltx-2-19b-dev.safetensors")
+ "model_hash": "aca7b0bbf8415e9c98360750268915fc",
+ "model_name": "ltx2_text_encoder_post_modules",
+ "model_class": "diffsynth.models.ltx2_text_encoder.LTX2TextEncoderPostModules",
+ "state_dict_converter": "diffsynth.utils.state_dict_converters.ltx2_text_encoder.LTX2TextEncoderPostModulesStateDictConverter",
+ },
+ {
+ # Example: ModelConfig(model_id="DiffSynth-Studio/LTX-2-Repackage", origin_file_pattern="text_encoder_post_modules.safetensors")
+ "model_hash": "981629689c8be92a712ab3c5eb4fc3f6",
+ "model_name": "ltx2_text_encoder_post_modules",
+ "model_class": "diffsynth.models.ltx2_text_encoder.LTX2TextEncoderPostModules",
+ "state_dict_converter": "diffsynth.utils.state_dict_converters.ltx2_text_encoder.LTX2TextEncoderPostModulesStateDictConverter",
+ },
+ {
+ # Example: ModelConfig(model_id="google/gemma-3-12b-it-qat-q4_0-unquantized", origin_file_pattern="model-*.safetensors")
+ "model_hash": "33917f31c4a79196171154cca39f165e",
+ "model_name": "ltx2_text_encoder",
+ "model_class": "diffsynth.models.ltx2_text_encoder.LTX2TextEncoder",
+ "state_dict_converter": "diffsynth.utils.state_dict_converters.ltx2_text_encoder.LTX2TextEncoderStateDictConverter",
+ },
+ {
+ # Example: ModelConfig(model_id="Lightricks/LTX-2", origin_file_pattern="ltx-2-19b-dev.safetensors")
+ "model_hash": "c79c458c6e99e0e14d47e676761732d2",
+ "model_name": "ltx2_latent_upsampler",
+ "model_class": "diffsynth.models.ltx2_upsampler.LTX2LatentUpsampler",
+ },
+ {
+ # Example: ModelConfig(model_id="Lightricks/LTX-2.3", origin_file_pattern="ltx-2.3-22b-dev.safetensors")
+ "model_hash": "f3a83ecf3995dcc4fae2d27e08ad5767",
+ "model_name": "ltx2_dit",
+ "model_class": "diffsynth.models.ltx2_dit.LTXModel",
+ "extra_kwargs": {"apply_gated_attention": True, "cross_attention_adaln": True, "caption_channels": None},
+ "state_dict_converter": "diffsynth.utils.state_dict_converters.ltx2_dit.LTXModelStateDictConverter",
+ },
+ {
+ # Example: ModelConfig(model_id="Lightricks/LTX-2.3", origin_file_pattern="ltx-2.3-22b-dev.safetensors")
+ "model_hash": "f3a83ecf3995dcc4fae2d27e08ad5767",
+ "model_name": "ltx2_video_vae_encoder",
+ "model_class": "diffsynth.models.ltx2_video_vae.LTX2VideoEncoder",
+ "extra_kwargs": {"encoder_version": "ltx-2.3"},
+ "state_dict_converter": "diffsynth.utils.state_dict_converters.ltx2_video_vae.LTX2VideoEncoderStateDictConverter",
+ },
+ {
+ # Example: ModelConfig(model_id="Lightricks/LTX-2.3", origin_file_pattern="ltx-2.3-22b-dev.safetensors")
+ "model_hash": "f3a83ecf3995dcc4fae2d27e08ad5767",
+ "model_name": "ltx2_video_vae_decoder",
+ "model_class": "diffsynth.models.ltx2_video_vae.LTX2VideoDecoder",
+ "extra_kwargs": {"decoder_version": "ltx-2.3"},
+ "state_dict_converter": "diffsynth.utils.state_dict_converters.ltx2_video_vae.LTX2VideoDecoderStateDictConverter",
+ },
+ {
+ # Example: ModelConfig(model_id="Lightricks/LTX-2.3", origin_file_pattern="ltx-2.3-22b-dev.safetensors")
+ "model_hash": "f3a83ecf3995dcc4fae2d27e08ad5767",
+ "model_name": "ltx2_audio_vae_decoder",
+ "model_class": "diffsynth.models.ltx2_audio_vae.LTX2AudioDecoder",
+ "state_dict_converter": "diffsynth.utils.state_dict_converters.ltx2_audio_vae.LTX2AudioDecoderStateDictConverter",
+ },
+ {
+ # Example: ModelConfig(model_id="Lightricks/LTX-2.3", origin_file_pattern="ltx-2.3-22b-dev.safetensors")
+ "model_hash": "f3a83ecf3995dcc4fae2d27e08ad5767",
+ "model_name": "ltx2_audio_vocoder",
+ "model_class": "diffsynth.models.ltx2_audio_vae.LTX2VocoderWithBWE",
+ "state_dict_converter": "diffsynth.utils.state_dict_converters.ltx2_audio_vae.LTX2VocoderStateDictConverter",
+ },
+ {
+ # Example: ModelConfig(model_id="Lightricks/LTX-2.3", origin_file_pattern="ltx-2.3-22b-dev.safetensors")
+ "model_hash": "f3a83ecf3995dcc4fae2d27e08ad5767",
+ "model_name": "ltx2_audio_vae_encoder",
+ "model_class": "diffsynth.models.ltx2_audio_vae.LTX2AudioEncoder",
+ "state_dict_converter": "diffsynth.utils.state_dict_converters.ltx2_audio_vae.LTX2AudioEncoderStateDictConverter",
+ },
+ {
+ # Example: ModelConfig(model_id="Lightricks/LTX-2.3", origin_file_pattern="ltx-2.3-22b-dev.safetensors")
+ "model_hash": "f3a83ecf3995dcc4fae2d27e08ad5767",
+ "model_name": "ltx2_text_encoder_post_modules",
+ "model_class": "diffsynth.models.ltx2_text_encoder.LTX2TextEncoderPostModules",
+ "extra_kwargs": {"separated_audio_video": True, "embedding_dim_gemma": 3840, "num_layers_gemma": 49, "video_attention_heads": 32, "video_attention_head_dim": 128, "audio_attention_heads": 32, "audio_attention_head_dim": 64, "num_connector_layers": 8, "apply_gated_attention": True},
+ "state_dict_converter": "diffsynth.utils.state_dict_converters.ltx2_text_encoder.LTX2TextEncoderPostModulesStateDictConverter",
+ },
+ {
+ # Example: ModelConfig(model_id="Lightricks/LTX-2.3", origin_file_pattern="ltx-2.3-spatial-upscaler-x2-1.0.safetensors")
+ "model_hash": "aed408774d694a2452f69936c32febb5",
+ "model_name": "ltx2_latent_upsampler",
+ "model_class": "diffsynth.models.ltx2_upsampler.LTX2LatentUpsampler",
+ "extra_kwargs": {"rational_resampler": False},
+ },
+ {
+ # Example: ModelConfig(model_id="DiffSynth-Studio/LTX-2.3-Repackage", origin_file_pattern="transformer.safetensors")
+ "model_hash": "1c55afad76ed33c112a2978550b524d1",
+ "model_name": "ltx2_dit",
+ "model_class": "diffsynth.models.ltx2_dit.LTXModel",
+ "extra_kwargs": {"apply_gated_attention": True, "cross_attention_adaln": True, "caption_channels": None},
+ "state_dict_converter": "diffsynth.utils.state_dict_converters.ltx2_dit.LTXModelStateDictConverter",
+ },
+ {
+ # Example: ModelConfig(model_id="DiffSynth-Studio/LTX-2.3-Repackage", origin_file_pattern="video_vae_encoder.safetensors")
+ "model_hash": "eecdc07c2ec30863b8a2b8b2134036cf",
+ "model_name": "ltx2_video_vae_encoder",
+ "model_class": "diffsynth.models.ltx2_video_vae.LTX2VideoEncoder",
+ "extra_kwargs": {"encoder_version": "ltx-2.3"},
+ "state_dict_converter": "diffsynth.utils.state_dict_converters.ltx2_video_vae.LTX2VideoEncoderStateDictConverter",
+ },
+ {
+ # Example: ModelConfig(model_id="DiffSynth-Studio/LTX-2.3-Repackage", origin_file_pattern="video_vae_decoder.safetensors")
+ "model_hash": "deda2f542e17ee25bc8c38fd605316ea",
+ "model_name": "ltx2_video_vae_decoder",
+ "model_class": "diffsynth.models.ltx2_video_vae.LTX2VideoDecoder",
+ "extra_kwargs": {"decoder_version": "ltx-2.3"},
+ "state_dict_converter": "diffsynth.utils.state_dict_converters.ltx2_video_vae.LTX2VideoDecoderStateDictConverter",
+ },
+ {
+ # Example: ModelConfig(model_id="DiffSynth-Studio/LTX-2.3-Repackage", origin_file_pattern="audio_vocoder.safetensors")
+ "model_hash": "cd436c99e69ec5c80f050f0944f02a15",
+ "model_name": "ltx2_audio_vocoder",
+ "model_class": "diffsynth.models.ltx2_audio_vae.LTX2VocoderWithBWE",
+ "state_dict_converter": "diffsynth.utils.state_dict_converters.ltx2_audio_vae.LTX2VocoderStateDictConverter",
+ },
+ {
+ # Example: ModelConfig(model_id="DiffSynth-Studio/LTX-2.3-Repackage", origin_file_pattern="text_encoder_post_modules.safetensors")
+ "model_hash": "05da2aab1c4b061f72c426311c165a43",
+ "model_name": "ltx2_text_encoder_post_modules",
+ "model_class": "diffsynth.models.ltx2_text_encoder.LTX2TextEncoderPostModules",
+ "extra_kwargs": {"separated_audio_video": True, "embedding_dim_gemma": 3840, "num_layers_gemma": 49, "video_attention_heads": 32, "video_attention_head_dim": 128, "audio_attention_heads": 32, "audio_attention_head_dim": 64, "num_connector_layers": 8, "apply_gated_attention": True},
+ "state_dict_converter": "diffsynth.utils.state_dict_converters.ltx2_text_encoder.LTX2TextEncoderPostModulesStateDictConverter",
+ },
+]
+anima_series = [
+ {
+ # Example: ModelConfig(model_id="circlestone-labs/Anima", origin_file_pattern="split_files/vae/qwen_image_vae.safetensors")
+ "model_hash": "a9995952c2d8e63cf82e115005eb61b9",
+ "model_name": "z_image_text_encoder",
+ "model_class": "diffsynth.models.z_image_text_encoder.ZImageTextEncoder",
+ "extra_kwargs": {"model_size": "0.6B"},
+ },
+ {
+ # Example: ModelConfig(model_id="circlestone-labs/Anima", origin_file_pattern="split_files/diffusion_models/anima-preview.safetensors")
+ "model_hash": "417673936471e79e31ed4d186d7a3f4a",
+ "model_name": "anima_dit",
+ "model_class": "diffsynth.models.anima_dit.AnimaDiT",
+ "state_dict_converter": "diffsynth.utils.state_dict_converters.anima_dit.AnimaDiTStateDictConverter",
+ }
+]
+
+mova_series = [
+ # Example: ModelConfig(model_id="openmoss/MOVA-720p", origin_file_pattern="audio_dit/diffusion_pytorch_model.safetensors")
+ {
+ "model_hash": "8c57e12790e2c45a64817e0ce28cde2f",
+ "model_name": "mova_audio_dit",
+ "model_class": "diffsynth.models.mova_audio_dit.MovaAudioDit",
+ "extra_kwargs": {'has_image_input': False, 'patch_size': [1], 'in_dim': 128, 'dim': 1536, 'ffn_dim': 8960, 'freq_dim': 256, 'text_dim': 4096, 'out_dim': 128, 'num_heads': 12, 'num_layers': 30, 'eps': 1e-06}
+ },
+ # Example: ModelConfig(model_id="openmoss/MOVA-720p", origin_file_pattern="audio_vae/diffusion_pytorch_model.safetensors")
+ {
+ "model_hash": "418517fb2b4e919d2cac8f314fcf82ac",
+ "model_name": "mova_audio_vae",
+ "model_class": "diffsynth.models.mova_audio_vae.DacVAE",
+ },
+ # Example: ModelConfig(model_id="openmoss/MOVA-720p", origin_file_pattern="dual_tower_bridge/diffusion_pytorch_model.safetensors")
+ {
+ "model_hash": "d1139dbbc8b4ab53cf4b4243d57bbceb",
+ "model_name": "mova_dual_tower_bridge",
+ "model_class": "diffsynth.models.mova_dual_tower_bridge.DualTowerConditionalBridge",
+ },
+]
+stable_diffusion_xl_series = [
+ {
+ # Example: ModelConfig(model_id="stabilityai/stable-diffusion-xl-base-1.0", origin_file_pattern="unet/diffusion_pytorch_model.safetensors")
+ "model_hash": "142b114f67f5ab3a6d83fb5788f12ded",
+ "model_name": "stable_diffusion_xl_unet",
+ "model_class": "diffsynth.models.stable_diffusion_xl_unet.SDXLUNet2DConditionModel",
+ "extra_kwargs": {"attention_head_dim": [5, 10, 20], "transformer_layers_per_block": [1, 2, 10], "use_linear_projection": True, "addition_embed_type": "text_time", "addition_time_embed_dim": 256, "projection_class_embeddings_input_dim": 2816},
+ },
+ {
+ # Example: ModelConfig(model_id="stabilityai/stable-diffusion-xl-base-1.0", origin_file_pattern="text_encoder_2/model.safetensors")
+ "model_hash": "98cc34ccc5b54ae0e56bdea8688dcd5a",
+ "model_name": "stable_diffusion_xl_text_encoder",
+ "model_class": "diffsynth.models.stable_diffusion_xl_text_encoder.SDXLTextEncoder2",
+ "state_dict_converter": "diffsynth.utils.state_dict_converters.stable_diffusion_xl_text_encoder.SDXLTextEncoder2StateDictConverter",
+ },
+ {
+ # Example: ModelConfig(model_id="stabilityai/stable-diffusion-xl-base-1.0", origin_file_pattern="text_encoder/model.safetensors")
+ "model_hash": "94eefa3dac9cec93cb1ebaf1747d7b78",
+ "model_name": "stable_diffusion_text_encoder",
+ "model_class": "diffsynth.models.stable_diffusion_text_encoder.SDTextEncoder",
+ "state_dict_converter": "diffsynth.utils.state_dict_converters.stable_diffusion_text_encoder.SDTextEncoderStateDictConverter",
+ },
+ {
+ # Example: ModelConfig(model_id="stabilityai/stable-diffusion-xl-base-1.0", origin_file_pattern="vae/diffusion_pytorch_model.safetensors")
+ "model_hash": "13115dd45a6e1c39860f91ab073b8a78",
+ "model_name": "stable_diffusion_xl_vae",
+ "model_class": "diffsynth.models.stable_diffusion_vae.StableDiffusionVAE",
+ "state_dict_converter": "diffsynth.utils.state_dict_converters.stable_diffusion_vae.SDVAEStateDictConverter",
+ "extra_kwargs": {"scaling_factor": 0.13025, "sample_size": 1024, "force_upcast": True},
+ },
+ {
+ # Example: ModelConfig(model_id="stabilityai/stable-diffusion-xl-base-1.0", origin_file_pattern="sd_xl_base_1.0.safetensors")
+ "model_hash": "4cf64a799d04260df438c6f33c9a047e",
+ "model_name": "stable_diffusion_xl_unet",
+ "model_class": "diffsynth.models.stable_diffusion_xl_unet.SDXLUNet2DConditionModel",
+ "extra_kwargs": {"attention_head_dim": [5, 10, 20], "transformer_layers_per_block": [1, 2, 10], "use_linear_projection": True, "addition_embed_type": "text_time", "addition_time_embed_dim": 256, "projection_class_embeddings_input_dim": 2816},
+ "state_dict_converter": "diffsynth.utils.state_dict_converters.sdxl.SDXLUNetStateDictConverter_Original2Diffusers",
+ },
+ {
+ # Example: ModelConfig(model_id="stabilityai/stable-diffusion-xl-base-1.0", origin_file_pattern="sd_xl_base_1.0.safetensors")
+ "model_hash": "4cf64a799d04260df438c6f33c9a047e",
+ "model_name": "stable_diffusion_text_encoder",
+ "model_class": "diffsynth.models.stable_diffusion_text_encoder.SDTextEncoder",
+ "state_dict_converter": "diffsynth.utils.state_dict_converters.sdxl_text_encoder.SDXLTextEncoderStateDictConverter_Original2Diffusers",
+ },
+ {
+ # Example: ModelConfig(model_id="stabilityai/stable-diffusion-xl-base-1.0", origin_file_pattern="sd_xl_base_1.0.safetensors")
+ "model_hash": "4cf64a799d04260df438c6f33c9a047e",
+ "model_name": "stable_diffusion_xl_text_encoder",
+ "model_class": "diffsynth.models.stable_diffusion_xl_text_encoder.SDXLTextEncoder2",
+ "state_dict_converter": "diffsynth.utils.state_dict_converters.sdxl_text_encoder_2.SDXLTextEncoder2StateDictConverter_Original2Diffusers",
+ },
+ {
+ # Example: ModelConfig(model_id="stabilityai/stable-diffusion-xl-base-1.0", origin_file_pattern="sd_xl_base_1.0.safetensors")
+ "model_hash": "4cf64a799d04260df438c6f33c9a047e",
+ "model_name": "stable_diffusion_xl_vae",
+ "model_class": "diffsynth.models.stable_diffusion_vae.StableDiffusionVAE",
+ "state_dict_converter": "diffsynth.utils.state_dict_converters.sdxl_vae.SDXLVAEStateDictConverter_Original2Diffusers",
+ "extra_kwargs": {"scaling_factor": 0.13025, "sample_size": 1024, "force_upcast": True},
+ },
+ {
+ # Example: Non-mainstream format 1
+ "model_hash": "d34167ecd389f17d58f6c907b65173fa",
+ "model_name": "stable_diffusion_xl_unet",
+ "model_class": "diffsynth.models.stable_diffusion_xl_unet.SDXLUNet2DConditionModel",
+ "extra_kwargs": {"attention_head_dim": [5, 10, 20], "transformer_layers_per_block": [1, 2, 10], "use_linear_projection": True, "addition_embed_type": "text_time", "addition_time_embed_dim": 256, "projection_class_embeddings_input_dim": 2816},
+ "state_dict_converter": "diffsynth.utils.state_dict_converters.sdxl.SDXLUNetStateDictConverter_Original2Diffusers",
+ },
+ {
+ # Example: Non-mainstream format 1
+ "model_hash": "d34167ecd389f17d58f6c907b65173fa",
+ "model_name": "stable_diffusion_text_encoder",
+ "model_class": "diffsynth.models.stable_diffusion_text_encoder.SDTextEncoder",
+ "state_dict_converter": "diffsynth.utils.state_dict_converters.sdxl_text_encoder.SDXLTextEncoderStateDictConverter_Original2Diffusers",
+ },
+ {
+ # Example: Non-mainstream format 1
+ "model_hash": "d34167ecd389f17d58f6c907b65173fa",
+ "model_name": "stable_diffusion_xl_text_encoder",
+ "model_class": "diffsynth.models.stable_diffusion_xl_text_encoder.SDXLTextEncoder2",
+ "state_dict_converter": "diffsynth.utils.state_dict_converters.sdxl_text_encoder_2.SDXLTextEncoder2StateDictConverter_Original2Diffusers",
+ },
+ {
+ # Example: Non-mainstream format 1
+ "model_hash": "d34167ecd389f17d58f6c907b65173fa",
+ "model_name": "stable_diffusion_xl_vae",
+ "model_class": "diffsynth.models.stable_diffusion_vae.StableDiffusionVAE",
+ "state_dict_converter": "diffsynth.utils.state_dict_converters.sdxl_vae.SDXLVAEStateDictConverter_Original2Diffusers",
+ "extra_kwargs": {"scaling_factor": 0.13025, "sample_size": 1024, "force_upcast": True},
+ },
+ {
+ # Example: Non-mainstream format 2
+ "model_hash": "c906695c0d35ec8547ebdd70d2e6a19b",
+ "model_name": "stable_diffusion_xl_unet",
+ "model_class": "diffsynth.models.stable_diffusion_xl_unet.SDXLUNet2DConditionModel",
+ "extra_kwargs": {"attention_head_dim": [5, 10, 20], "transformer_layers_per_block": [1, 2, 10], "use_linear_projection": True, "addition_embed_type": "text_time", "addition_time_embed_dim": 256, "projection_class_embeddings_input_dim": 2816},
+ "state_dict_converter": "diffsynth.utils.state_dict_converters.sdxl.SDXLUNetStateDictConverter_Original2Diffusers",
+ },
+ {
+ # Example: Non-mainstream format 2
+ "model_hash": "c906695c0d35ec8547ebdd70d2e6a19b",
+ "model_name": "stable_diffusion_text_encoder",
+ "model_class": "diffsynth.models.stable_diffusion_text_encoder.SDTextEncoder",
+ "state_dict_converter": "diffsynth.utils.state_dict_converters.sdxl_text_encoder.SDXLTextEncoderStateDictConverter_Original2Diffusers",
+ },
+ {
+ # Example: Non-mainstream format 2
+ "model_hash": "c906695c0d35ec8547ebdd70d2e6a19b",
+ "model_name": "stable_diffusion_xl_text_encoder",
+ "model_class": "diffsynth.models.stable_diffusion_xl_text_encoder.SDXLTextEncoder2",
+ "state_dict_converter": "diffsynth.utils.state_dict_converters.sdxl_text_encoder_2.SDXLTextEncoder2StateDictConverter_Original2Diffusers",
+ },
+ {
+ # Example: Non-mainstream format 2
+ "model_hash": "c906695c0d35ec8547ebdd70d2e6a19b",
+ "model_name": "stable_diffusion_xl_vae",
+ "model_class": "diffsynth.models.stable_diffusion_vae.StableDiffusionVAE",
+ "state_dict_converter": "diffsynth.utils.state_dict_converters.sdxl_vae.SDXLVAEStateDictConverter_Original2Diffusers",
+ "extra_kwargs": {"scaling_factor": 0.13025, "sample_size": 1024, "force_upcast": True},
+ },
+ {
+ # Example: Non-mainstream format 3
+ "model_hash": "1a9744a6f0c0e54ef40805c9d2b79ec6",
+ "model_name": "stable_diffusion_xl_unet",
+ "model_class": "diffsynth.models.stable_diffusion_xl_unet.SDXLUNet2DConditionModel",
+ "extra_kwargs": {"attention_head_dim": [5, 10, 20], "transformer_layers_per_block": [1, 2, 10], "use_linear_projection": True, "addition_embed_type": "text_time", "addition_time_embed_dim": 256, "projection_class_embeddings_input_dim": 2816},
+ "state_dict_converter": "diffsynth.utils.state_dict_converters.sdxl.SDXLUNetStateDictConverter_Original2Diffusers",
+ },
+ {
+ # Example: Non-mainstream format 3
+ "model_hash": "1a9744a6f0c0e54ef40805c9d2b79ec6",
+ "model_name": "stable_diffusion_text_encoder",
+ "model_class": "diffsynth.models.stable_diffusion_text_encoder.SDTextEncoder",
+ "state_dict_converter": "diffsynth.utils.state_dict_converters.sdxl_text_encoder.SDXLTextEncoderStateDictConverter_Original2Diffusers",
+ },
+ {
+ # Example: Non-mainstream format 3
+ "model_hash": "1a9744a6f0c0e54ef40805c9d2b79ec6",
+ "model_name": "stable_diffusion_xl_text_encoder",
+ "model_class": "diffsynth.models.stable_diffusion_xl_text_encoder.SDXLTextEncoder2",
+ "state_dict_converter": "diffsynth.utils.state_dict_converters.sdxl_text_encoder_2.SDXLTextEncoder2StateDictConverter_Original2Diffusers",
+ },
+ {
+ # Example: Non-mainstream format 3
+ "model_hash": "1a9744a6f0c0e54ef40805c9d2b79ec6",
+ "model_name": "stable_diffusion_xl_vae",
+ "model_class": "diffsynth.models.stable_diffusion_vae.StableDiffusionVAE",
+ "state_dict_converter": "diffsynth.utils.state_dict_converters.sdxl_vae.SDXLVAEStateDictConverter_Original2Diffusers",
+ "extra_kwargs": {"scaling_factor": 0.13025, "sample_size": 1024, "force_upcast": True},
+ },
+]
+
+stable_diffusion_series = [
+ {
+ # Example: ModelConfig(model_id="AI-ModelScope/stable-diffusion-v1-5", origin_file_pattern="text_encoder/model.safetensors")
+ "model_hash": "ffd1737ae9df7fd43f5fbed653bdad67",
+ "model_name": "stable_diffusion_text_encoder",
+ "model_class": "diffsynth.models.stable_diffusion_text_encoder.SDTextEncoder",
+ "state_dict_converter": "diffsynth.utils.state_dict_converters.stable_diffusion_text_encoder.SDTextEncoderStateDictConverter",
+ },
+ {
+ # Example: ModelConfig(model_id="AI-ModelScope/stable-diffusion-v1-5", origin_file_pattern="vae/diffusion_pytorch_model.safetensors")
+ "model_hash": "f86d5683ed32433be8ca69969c67ba69",
+ "model_name": "stable_diffusion_vae",
+ "model_class": "diffsynth.models.stable_diffusion_vae.StableDiffusionVAE",
+ "state_dict_converter": "diffsynth.utils.state_dict_converters.stable_diffusion_vae.SDVAEStateDictConverter",
+ },
+ {
+ # Example: ModelConfig(model_id="AI-ModelScope/stable-diffusion-v1-5", origin_file_pattern="unet/diffusion_pytorch_model.safetensors")
+ "model_hash": "025a4b86a84829399d89f613e580757b",
+ "model_name": "stable_diffusion_unet",
+ "model_class": "diffsynth.models.stable_diffusion_unet.UNet2DConditionModel",
+ },
+]
+
+joyai_image_series = [
+ {
+ # Example: ModelConfig(model_id="jd-opensource/JoyAI-Image-Edit", origin_file_pattern="transformer/transformer.pth")
+ "model_hash": "56592ddfd7d0249d3aa527d24161a863",
+ "model_name": "joyai_image_dit",
+ "model_class": "diffsynth.models.joyai_image_dit.JoyAIImageDiT",
+ },
+ {
+ # Example: ModelConfig(model_id="jd-opensource/JoyAI-Image-Edit", origin_file_pattern="JoyAI-Image-Und/model-*.safetensors")
+ # Example: ModelConfig(model_id="DiffSynth-Studio/ImageMetrics", origin_file_pattern="UnifiedReward-Edit-qwen3vl-8b/model-*.safetensors")
+ "model_hash": "2d11bf14bba8b4e87477c8199a895403",
+ "model_name": "joyai_image_text_encoder",
+ "model_class": "diffsynth.models.joyai_image_text_encoder.JoyAIImageTextEncoder",
+ "state_dict_converter": "diffsynth.utils.state_dict_converters.joyai_image_text_encoder.JoyAIImageTextEncoderStateDictConverter",
+ },
+]
+
+boogu_image_series = [
+ {
+ # Example: ModelConfig(model_id="Boogu/Boogu-Image-0.1-Base", origin_file_pattern="transformer/diffusion_pytorch_model-*.safetensors")
+ "model_hash": "8e29d2d28b83530320ae1d2c5eb323b3",
+ "model_name": "boogu_image_dit",
+ "model_class": "diffsynth.models.boogu_image_dit.BooguImageDiT",
+ },
+]
+
+ace_step_series = [
+ {
+ # Example: ModelConfig(model_id="ACE-Step/Ace-Step1.5", origin_file_pattern="acestep-v15-turbo/model.safetensors")
+ "model_hash": "ba29d8bddbb6ace65675f6a757a13c00",
+ "model_name": "ace_step_dit",
+ "model_class": "diffsynth.models.ace_step_dit.AceStepDiTModel",
+ "state_dict_converter": "diffsynth.utils.state_dict_converters.ace_step_dit.AceStepDiTModelStateDictConverter",
+ },
+ {
+ # Example: ModelConfig(model_id="ACE-Step/acestep-v15-xl-base", origin_file_pattern="model-*.safetensors")
+ "model_hash": "3a28a410c2246f125153ef792d8bc828",
+ "model_name": "ace_step_dit",
+ "model_class": "diffsynth.models.ace_step_dit.AceStepDiTModel",
+ "state_dict_converter": "diffsynth.utils.state_dict_converters.ace_step_dit.AceStepDiTModelStateDictConverter",
+ "extra_kwargs": {
+ "hidden_size": 2560,
+ "intermediate_size": 9728,
+ "num_hidden_layers": 32,
+ "num_attention_heads": 32,
+ "num_key_value_heads": 8,
+ "head_dim": 128,
+ "encoder_hidden_size": 2048,
+ "layer_types": ["sliding_attention", "full_attention"] * 16,
+ },
+ },
+ {
+ # Example: ModelConfig(model_id="ACE-Step/Ace-Step1.5", origin_file_pattern="acestep-v15-turbo/model.safetensors")
+ "model_hash": "ba29d8bddbb6ace65675f6a757a13c00",
+ "model_name": "ace_step_conditioner",
+ "model_class": "diffsynth.models.ace_step_conditioner.AceStepConditionEncoder",
+ "state_dict_converter": "diffsynth.utils.state_dict_converters.ace_step_conditioner.AceStepConditionEncoderStateDictConverter",
+ },
+ {
+ # Example: ModelConfig(model_id="ACE-Step/acestep-v15-xl-base", origin_file_pattern="model-*.safetensors")
+ "model_hash": "3a28a410c2246f125153ef792d8bc828",
+ "model_name": "ace_step_conditioner",
+ "model_class": "diffsynth.models.ace_step_conditioner.AceStepConditionEncoder",
+ "state_dict_converter": "diffsynth.utils.state_dict_converters.ace_step_conditioner.AceStepConditionEncoderStateDictConverter",
+ },
+ {
+ # Example: ModelConfig(model_id="ACE-Step/Ace-Step1.5", origin_file_pattern="Qwen3-Embedding-0.6B/model.safetensors")
+ "model_hash": "3509bea17b0e8cffc3dd4a15cc7899d0",
+ "model_name": "ace_step_text_encoder",
+ "model_class": "diffsynth.models.ace_step_text_encoder.AceStepTextEncoder",
+ "state_dict_converter": "diffsynth.utils.state_dict_converters.ace_step_text_encoder.AceStepTextEncoderStateDictConverter",
+ },
+ {
+ # Example: ModelConfig(model_id="ACE-Step/Ace-Step1.5", origin_file_pattern="vae/diffusion_pytorch_model.safetensors")
+ "model_hash": "51420834e54474986a7f4be0e4d6f687",
+ "model_name": "ace_step_vae",
+ "model_class": "diffsynth.models.ace_step_vae.AceStepVAE",
+ },
+ {
+ # Example: ModelConfig(model_id="ACE-Step/Ace-Step1.5", origin_file_pattern="acestep-v15-turbo/model.safetensors")
+ "model_hash": "ba29d8bddbb6ace65675f6a757a13c00",
+ "model_name": "ace_step_tokenizer",
+ "model_class": "diffsynth.models.ace_step_tokenizer.AceStepTokenizer",
+ "state_dict_converter": "diffsynth.utils.state_dict_converters.ace_step_tokenizer.AceStepTokenizerStateDictConverter",
+ },
+ {
+ # Example: ModelConfig(model_id="ACE-Step/acestep-v15-xl-base", origin_file_pattern="model-*.safetensors")
+ "model_hash": "3a28a410c2246f125153ef792d8bc828",
+ "model_name": "ace_step_tokenizer",
+ "model_class": "diffsynth.models.ace_step_tokenizer.AceStepTokenizer",
+ "state_dict_converter": "diffsynth.utils.state_dict_converters.ace_step_tokenizer.AceStepTokenizerStateDictConverter",
+ },
+ {
+ # Example: ???
+ "model_hash": "ff74b1806e6a0b52e7bbd1d3df2d26d1",
+ "model_name": "demucs",
+ "model_class": "diffsynth.models.demucs.HTDemucs",
+ },
+]
+
+image_metrics_series = [
+ {
+ # Example: ModelConfig(model_id="DiffSynth-Studio/ImageMetrics", origin_file_pattern="PickScore/model.safetensors")
+ # Example: ModelConfig(model_id="DiffSynth-Studio/ImageMetrics", origin_file_pattern="CLIP-ViT-H-14-laion2B-s32B-b79K/model.safetensors")
+ "model_hash": "b5e2c0bfcbf4085ccdb2feb8f0ba408a",
+ "model_name": "image_metrics_clip_hf",
+ "model_class": "diffsynth.models.clip.ImageMetricsCLIPModel",
+ "state_dict_converter": "diffsynth.utils.state_dict_converters.image_metrics.ImageMetricsCLIPStateDictConverter",
+ },
+ {
+ # Example: ModelConfig(model_id="DiffSynth-Studio/ImageMetrics", origin_file_pattern="HPSv2/model.safetensors")
+ "model_hash": "f79e72cec8ae5a540cff0304bfb21b00",
+ "model_name": "image_metrics_hpsv2",
+ "model_class": "diffsynth.models.clip.ImageMetricsCLIPModel",
+ "state_dict_converter": "diffsynth.utils.state_dict_converters.image_metrics.ImageMetricsOpenCLIPStateDictConverter",
+ },
+ {
+ # Example: ModelConfig(model_id="DiffSynth-Studio/ImageMetrics", origin_file_pattern="HPSv3/model.safetensors")
+ "model_hash": "5655d9cde15b759cfeefe7432d7a912c",
+ "model_name": "image_metrics_hpsv3",
+ "model_class": "diffsynth.models.hpsv3.HPSv3Qwen2VLRewardModel",
+ "state_dict_converter": "diffsynth.utils.state_dict_converters.image_metrics.ImageMetricsHPSv3StateDictConverter",
+ "extra_kwargs": {"vocab_size": 151658, "output_dim": 2, "reward_token": "special", "rm_head_type": "ranknet"},
+ },
+ {
+ # Example: ModelConfig(model_id="DiffSynth-Studio/ImageMetrics", origin_file_pattern="ImageReward/model.safetensors")
+ "model_hash": "b3cc8e10b76ca98cde653daa5cf63139",
+ "model_name": "image_metrics_image_reward",
+ "model_class": "diffsynth.models.image_reward.ImageRewardModel",
+ "state_dict_converter": "diffsynth.utils.state_dict_converters.image_metrics.ImageMetricsImageRewardStateDictConverter",
+ },
+ {
+ # Example: ModelConfig(model_id="DiffSynth-Studio/ImageMetrics", origin_file_pattern="Aesthetic/model.safetensors")
+ "model_hash": "306981222ec94302794e07cf676c84cc",
+ "model_name": "image_metrics_aesthetic",
+ "model_class": "diffsynth.models.aesthetic.AestheticModel",
+ "state_dict_converter": "diffsynth.utils.state_dict_converters.image_metrics.ImageMetricsAestheticStateDictConverter",
+ },
+ {
+ # Example: ModelConfig(model_id="DiffSynth-Studio/ImageMetrics", origin_file_pattern="FID/model.safetensors")
+ "model_hash": "d4e9549be726259b444d1f62db4ce413",
+ "model_name": "image_metrics_fid_inception",
+ "model_class": "diffsynth.models.fid.FIDInceptionModel",
+ "state_dict_converter": "diffsynth.utils.state_dict_converters.image_metrics.ImageMetricsFIDStateDictConverter",
+ },
+ {
+ # Example: ModelConfig(model_id="DiffSynth-Studio/ImageMetrics", origin_file_pattern="BioCLIPv2/open_clip_model.safetensors")
+ "model_hash": "3a020a3e47afb7c5e21c52f2d0692c09",
+ "model_name": "image_metrics_bioclip_v2",
+ "model_class": "diffsynth.models.bioclip.BioCLIPv2Model",
+ "state_dict_converter": "diffsynth.utils.state_dict_converters.image_metrics.ImageMetricsOpenCLIPStateDictConverter",
+ },
+ {
+ # Example: ModelConfig(model_id="DiffSynth-Studio/ImageMetrics", origin_file_pattern="LPIPS/alexnet.safetensors")
+ "model_hash": "08a75c660c9b2e775c530a0955857f1f",
+ "model_name": "image_metrics_lpips_alex",
+ "model_class": "diffsynth.models.lpips.LPIPSModel",
+ "extra_kwargs": {"net": "alex"},
+ },
+ {
+ # Example: ModelConfig(model_id="DiffSynth-Studio/ImageMetrics", origin_file_pattern="LPIPS/vgg.safetensors")
+ "model_hash": "5740953aaa8aba2ecd9b9c23da813591",
+ "model_name": "image_metrics_lpips_vgg",
+ "model_class": "diffsynth.models.lpips.LPIPSModel",
+ "extra_kwargs": {"net": "vgg"},
+ },
+ {
+ # Example: ModelConfig(model_id="DiffSynth-Studio/ImageMetrics", origin_file_pattern="LPIPS/squeezenet.safetensors")
+ "model_hash": "ff994b70a30599287a332105396d5004",
+ "model_name": "image_metrics_lpips_squeeze",
+ "model_class": "diffsynth.models.lpips.LPIPSModel",
+ "extra_kwargs": {"net": "squeeze"},
+ },
+ {
+ # Example: ModelConfig(model_id="DiffSynth-Studio/ImageMetrics", origin_file_pattern="UnifiedReward-2.0-qwen35-9b/model-*.safetensors")
+ "model_hash": "f9786d06eca5c0f1ece89843b2c4cc66",
+ "model_name": "image_metrics_unified_reward_2",
+ "model_class": "diffsynth.models.unified_reward_2.UnifiedReward2Qwen35ForConditionalGeneration",
+ "state_dict_converter": "diffsynth.utils.state_dict_converters.image_metrics.ImageMetricsUnifiedRewardStateDictConverter",
+ "extra_kwargs": {"variant": "qwen35_9b"},
+ },
+ {
+ # Example: ModelConfig(model_id="DiffSynth-Studio/ImageMetrics", origin_file_pattern="Qwen-Image-Bench/model-*.safetensors")
+ "model_hash": "ff4ad0463675e96738483611f6dd551b",
+ "model_name": "image_metrics_qwen_image_bench",
+ "model_class": "diffsynth.models.qwen_image_bench.QwenImageBenchQwen35ForConditionalGeneration",
+ "state_dict_converter": "diffsynth.utils.state_dict_converters.image_metrics.ImageMetricsUnifiedRewardStateDictConverter",
+ "extra_kwargs": {"variant": "qwen35"},
+ },
+]
+
+hidream_o1_image_series = [
+ {
+ # Example: ModelConfig(model_id="HiDream-ai/HiDream-O1-Image", origin_file_pattern="model-*.safetensors")
+ "model_hash": "58a7c1073d79556bfc61e05e6061b771",
+ "model_name": "hidream_o1_image_dit",
+ "model_class": "diffsynth.models.hidream_o1_image_dit.HiDreamO1ImageModel",
+ },
+]
+
+ideogram4_series = [
+ {
+ # Example: ModelConfig(model_id="ideogram-ai/ideogram-4-fp8", origin_file_pattern="transformer/diffusion_pytorch_model.safetensors")
+ "model_hash": "6f56a1d28667f2ff98e1c79af88a7516",
+ "model_name": "ideogram4_dit",
+ "model_class": "diffsynth.models.ideogram4_dit.Ideogram4DiT",
+ # The checkpoint already holds fp8 weights; the backend that reads its layout is
+ # registered next to the model, in `diffsynth.models.ideogram4_dit`.
+ "quant_config": {"method": "ideogram4_fp8", "load_prequantized": True},
+ "extra_kwargs": {"config": {"emb_dim": 4608, "num_layers": 34, "num_heads": 18, "intermediate_size": 12288, "adanln_dim": 512, "in_channels": 128, "llm_features_dim": 53248, "rope_theta": 5000000, "mrope_section": [24, 20, 20], "norm_eps": 1e-05}},
+ },
+ {
+ # Example: ModelConfig(model_id="ideogram-ai/ideogram-4-fp8", origin_file_pattern="text_encoder/model.safetensors")
+ "model_hash": "6d72a86d1027baff87e2cf8fc523aab1",
+ "model_name": "ideogram4_text_encoder",
+ "model_class": "diffsynth.models.ideogram4_text_encoder.Ideogram4TextEncoder",
+ "state_dict_converter": "diffsynth.utils.state_dict_converters.ideogram4_text_encoder.Ideogram4TextEncoderStateDictConverter",
+ # Quantized just like the DiT of the same release, and read by the same backend.
+ "quant_config": {"method": "ideogram4_fp8", "load_prequantized": True},
+ },
+ {
+ # Example: ModelConfig(model_id="ideogram-ai/ideogram-4-nf4", origin_file_pattern="transformer/diffusion_pytorch_model.safetensors")
+ # The same file also serves as `unconditional_transformer/diffusion_pytorch_model.safetensors`.
+ "model_hash": "61962d4376a8957f96d807047bc6f780",
+ "model_name": "ideogram4_dit",
+ "model_class": "diffsynth.models.ideogram4_dit.Ideogram4DiT",
+ # bnb's own 4bit layout, written without double quantization.
+ "quant_config": {"method": "bitsandbytes_nf4", "load_prequantized": True, "backend_config_kwargs": {"compress_statistics": False}},
+ "extra_kwargs": {"config": {"emb_dim": 4608, "num_layers": 34, "num_heads": 18, "intermediate_size": 12288, "adanln_dim": 512, "in_channels": 128, "llm_features_dim": 53248, "rope_theta": 5000000, "mrope_section": [24, 20, 20], "norm_eps": 1e-05}},
+ },
+ {
+ # Example: ModelConfig(model_id="ideogram-ai/ideogram-4-nf4", origin_file_pattern="text_encoder/model.safetensors")
+ "model_hash": "c14b8f1328a87ffdba145431277d133a",
+ "model_name": "ideogram4_text_encoder",
+ "model_class": "diffsynth.models.ideogram4_text_encoder.Ideogram4TextEncoder",
+ "state_dict_converter": "diffsynth.utils.state_dict_converters.ideogram4_text_encoder.Ideogram4TextEncoderStateDictConverter",
+ "quant_config": {"method": "bitsandbytes_nf4", "load_prequantized": True, "backend_config_kwargs": {"compress_statistics": False}},
+ },
+ {
+ # Example: ModelConfig(model_id="DiffSynth-Studio/ideogram-4-bf16-repackage", origin_file_pattern="transformer/diffusion_pytorch_model.safetensors")
+ "model_hash": "291b300b11c8c8e11978bd85a9c5f80c",
+ "model_name": "ideogram4_dit",
+ "model_class": "diffsynth.models.ideogram4_dit.Ideogram4DiT",
+ "extra_kwargs": {"config": {"emb_dim": 4608, "num_layers": 34, "num_heads": 18, "intermediate_size": 12288, "adanln_dim": 512, "in_channels": 128, "llm_features_dim": 53248, "rope_theta": 5000000, "mrope_section": [24, 20, 20], "norm_eps": 1e-05}},
+ },
+ {
+ # Example: ModelConfig(model_id="ideogram-ai/ideogram-4-fp8", origin_file_pattern="text_encoder/model.safetensors")
+ "model_hash": "6a269892c0757aacd46bd41b8d5a7aef",
+ "model_name": "ideogram4_text_encoder",
+ "model_class": "diffsynth.models.ideogram4_text_encoder.Ideogram4TextEncoder",
+ "state_dict_converter": "diffsynth.utils.state_dict_converters.ideogram4_text_encoder.Ideogram4TextEncoderStateDictConverter",
+ },
+]
+
+krea2_series = [
+ {
+ # Example: ModelConfig(model_id="AI-ModelScope/Krea-2", origin_file_pattern="raw.safetensors"),
+ "model_hash": "4be51f0ecc9b9dfc9825abb3f7a0b4ae",
+ "model_name": "krea2_dit",
+ "model_class": "diffsynth.models.krea2_dit.SingleStreamDiT",
+ "state_dict_converter": "diffsynth.utils.state_dict_converters.krea2_dit.Krea2DiTStateDictConverter",
+ },
+ {
+ # Example: ModelConfig(model_id="krea/Krea-2-Raw", origin_file_pattern="raw.safetensors")
+ "model_hash": "2a05e69084a6d608516ee47c8ad03334",
+ "model_name": "krea2_dit",
+ "model_class": "diffsynth.models.krea2_dit.SingleStreamDiT",
+ },
+ {
+ # Example: ModelConfig(model_id="krea/Krea-2-Raw", origin_file_pattern="transformer/*.safetensors")
+ "model_hash": "cddb3118d9f8250bf7b08c331c5c55aa",
+ "model_name": "krea2_dit",
+ "model_class": "diffsynth.models.krea2_dit.SingleStreamDiT",
+ "state_dict_converter": "diffsynth.utils.state_dict_converters.krea2_dit.Krea2DiTStateDictConverter_Diffusers",
+ },
+ {
+ # Example: ModelConfig(model_id="Qwen/Qwen3-VL-4B-Instruct", origin_file_pattern="*.safetensors")
+ "model_hash": "b8750c24f732c87797f551196c4cef78",
+ "model_name": "krea2_text_encoder",
+ "model_class": "diffsynth.models.krea2_text_encoder.Krea2TextEncoder",
+ "state_dict_converter": "diffsynth.utils.state_dict_converters.krea2_text_encoder.Krea2TextEncoderStateDictConverter",
+ },
+]
+
+lingbot_video_series = [
+ {
+ # Example: ModelConfig(model_id="Robbyant/lingbot-video-dense-1.3b", origin_file_pattern="transformer/diffusion_pytorch_model.safetensors")
+ "model_hash": "2bcf511fe5e0000519394d242b4d8abd",
+ "model_name": "lingbot_video_dit",
+ "model_class": "diffsynth.models.lingbot_video_dit.LingBotVideoDiT",
+ "state_dict_converter": "diffsynth.utils.state_dict_converters.lingbot_video_dit.LingBotVideoDiTStateDictConverter",
+ },
+ {
+ # Example: ModelConfig(model_id="Robbyant/lingbot-video-moe-30b-a3b", origin_file_pattern="transformer/diffusion_pytorch_model*.safetensors")
+ "model_hash": "65b83aa625cd362ff5ff3409fb367a6f",
+ "model_name": "lingbot_video_dit",
+ "model_class": "diffsynth.models.lingbot_video_dit.LingBotVideoDiT",
+ "state_dict_converter": "diffsynth.utils.state_dict_converters.lingbot_video_dit.LingBotVideoDiTStateDictConverter",
+ "extra_kwargs": {'depth': 48, 'axes_lens': (4096, 512, 512), 'num_experts': 128, 'moe_intermediate_size': 768, 'n_group': 4, 'topk_group': 2, 'n_shared_experts': 1, 'routed_scaling_factor': 2.5},
+ },
+]
+
+minimax_h3_series = [
+ {
+ # Example: ModelConfig(model_id="MiniMax/MiniMax-H3-FL2VA", origin_file_pattern="audio_vae/model.safetensors")
+ "model_hash": "db383f1c8960837b94059f7722e6cb11",
+ "model_name": "minimax_h3_audio_vae",
+ "model_class": "diffsynth.models.minimax_h3_audio_vae.MiniMaxH3AudioVAE",
+ "state_dict_converter": "diffsynth.utils.state_dict_converters.minimax_h3_audio_vae.MiniMaxH3AudioVAEStateDictConverter"
+ },
+ {
+ # Example: ModelConfig(model_id="MiniMax/MiniMax-H3-FL2VA", origin_file_pattern="text_encoder/model*.safetensors")
+ "model_hash": "8398104725e58fa8287c163c5c332686",
+ "model_name": "minimax_h3_text_encoder",
+ "model_class": "diffsynth.models.minimax_h3_text_encoder.MiniMaxH3TextEncoder",
+ "state_dict_converter": "diffsynth.utils.state_dict_converters.minimax_h3_text_encoder.MiniMaxH3TextEncoderStateDictConverter",
+ },
+ {
+ # Example: ModelConfig(model_id="DiffSynth-Studio/MiniMax-H3-NF4", origin_file_pattern="minimax-h3-text-encoder-nf4.safetensors")
+ "model_hash": "297933c3a2b0fc4d4dfee30e34c566b8",
+ "model_name": "minimax_h3_text_encoder",
+ "model_class": "diffsynth.models.minimax_h3_text_encoder.MiniMaxH3TextEncoder",
+ "state_dict_converter": "diffsynth.utils.state_dict_converters.minimax_h3_text_encoder.MiniMaxH3TextEncoderStateDictConverter",
+ "quant_config": {"method": "bitsandbytes_nf4", "load_prequantized": True, "exclude_modules": ["qkv", "proj", "linear_fc1", "linear_fc2"]},
+ },
+ {
+ # Example: ModelConfig(model_id="Comfy-Org/MiniMax-H3", origin_file_pattern="text_encoders/qwen3vl_32b_minimax_h3_int8_convrot.safetensors")
+ "model_hash": "d6473c909b72b73afc4fa00c273d1da2",
+ "model_name": "minimax_h3_text_encoder",
+ "model_class": "diffsynth.models.minimax_h3_text_encoder.MiniMaxH3TextEncoder",
+ "state_dict_converter": "diffsynth.utils.state_dict_converters.minimax_h3_text_encoder.MiniMaxH3TextEncoderComfyOrgStateDictConverter",
+ "quant_config": {"method": "comfy_kitchen_int8_w8a8", "load_prequantized": True, "exclude_modules": ["qkv", "proj", "linear_fc1", "linear_fc2"]},
+ },
+ {
+ # Example: ModelConfig(model_id="MiniMax/MiniMax-H3", origin_file_pattern="FL2VA/transformer/model*.safetensors")
+ "model_hash": "db0197b6919425a5c7102c54e73affc1",
+ "model_name": "minimax_h3_dit",
+ "model_class": "diffsynth.models.minimax_h3_dit.MiniMaxH3DiT",
+ },
+ {
+ # Example: ModelConfig(model_id="Comfy-Org/MiniMax-H3", origin_file_pattern="diffusion_models/minimax_h3_fl2va_int8_convrot.safetensors")
+ "model_hash": "fb6c399412920c6b817d5c1f43ec62cd",
+ "model_name": "minimax_h3_dit",
+ "model_class": "diffsynth.models.minimax_h3_dit_comfy.MiniMaxH3DiTComfy",
+ "quant_config": {"method": "comfy_kitchen_int8_w8a8", "load_prequantized": True, "exclude_modules": ["video_patch_proj", "audio_patch_proj", "condition_proj", "time_embedder.proj_in", "time_embedder.proj_out", "final_layer.adaln_proj.linear", "final_layer.video_out", "final_layer.audio_out", "token_refiner.blocks.0.attn.qkv_proj", "token_refiner.blocks.0.attn.out_proj", "token_refiner.blocks.0.mlp.fc1", "token_refiner.blocks.0.mlp.fc2", "token_refiner.blocks.1.attn.qkv_proj", "token_refiner.blocks.1.attn.out_proj", "token_refiner.blocks.1.mlp.fc1", "token_refiner.blocks.1.mlp.fc2"]},
+ },
+ {
+ # Example: ModelConfig(model_id="Comfy-Org/MiniMax-H3", origin_file_pattern="diffusion_models/minimax_h3_fl2va_pruned_int8_convrot.safetensors")
+ "model_hash": "ac9fdcd56900c9a1a5ef21eaacf6cb76",
+ "model_name": "minimax_h3_dit",
+ "model_class": "diffsynth.models.minimax_h3_dit_comfy.MiniMaxH3DiTComfyPruned",
+ "quant_config": {"method": "comfy_kitchen_int8_w8a8", "load_prequantized": True, "exclude_modules": ["video_patch_proj", "audio_patch_proj", "condition_proj", "adaln_proj.linear", "final_layer.video_out", "final_layer.audio_out", "token_refiner.blocks.0.attn.qkv_proj", "token_refiner.blocks.0.attn.out_proj", "token_refiner.blocks.0.mlp.fc1", "token_refiner.blocks.0.mlp.fc2", "token_refiner.blocks.1.attn.qkv_proj", "token_refiner.blocks.1.attn.out_proj", "token_refiner.blocks.1.mlp.fc1", "token_refiner.blocks.1.mlp.fc2"]},
+ },
+ {
+ # Example: ModelConfig(model_id="Comfy-Org/MiniMax-H3", origin_file_pattern="diffusion_models/minimax_h3_fl2va_pruned_fp8_scaled.safetensors")
+ "model_hash": "7e410be9acd325ef8db3b8bed3d814cd",
+ "model_name": "minimax_h3_dit",
+ "model_class": "diffsynth.models.minimax_h3_dit_comfy.MiniMaxH3DiTComfyPruned",
+ "quant_config": {"method": "comfy_kitchen_fp8_w8a8", "load_prequantized": True, "exclude_modules": ["video_patch_proj", "audio_patch_proj", "condition_proj", "adaln_proj.linear", "final_layer.video_out", "final_layer.audio_out", "token_refiner.blocks.0.attn.qkv_proj", "token_refiner.blocks.0.attn.out_proj", "token_refiner.blocks.0.mlp.fc1", "token_refiner.blocks.0.mlp.fc2", "token_refiner.blocks.1.attn.qkv_proj", "token_refiner.blocks.1.attn.out_proj", "token_refiner.blocks.1.mlp.fc1", "token_refiner.blocks.1.mlp.fc2"]},
+ },
+ {
+ # Example: ModelConfig(model_id="Comfy-Org/MiniMax-H3", origin_file_pattern="diffusion_models/minimax_h3_fl2va_pruned_bf16.safetensors")
+ "model_hash": "b4e0df87da0c079979ae857edbcef49f",
+ "model_name": "minimax_h3_dit",
+ "model_class": "diffsynth.models.minimax_h3_dit_comfy.MiniMaxH3DiTComfyPruned",
+ },
+ {
+ # Example: ModelConfig(model_id="DiffSynth-Studio/MiniMax-H3-NF4", origin_file_pattern="minimax-h3-fl2va-nf4.safetensors")
+ "model_hash": "4b27efacefbc4d8670e0d7b876699648",
+ "model_name": "minimax_h3_dit",
+ "model_class": "diffsynth.models.minimax_h3_dit.MiniMaxH3DiT",
+ "quant_config": {"method": "bitsandbytes_nf4", "load_prequantized": True, "exclude_modules": ["time_embedder.proj_in", "time_embedder.proj_out", "video_patch_proj", "audio_patch_proj", "condition_proj", "final_layer.video_out", "final_layer.audio_out"]},
+ },
+ {
+ # Example: ModelConfig(model_id="DiffSynth-Studio/MiniMax-H3-NF4", origin_file_pattern="minimax-h3-ref2va-nf4.safetensors")
+ "model_hash": "ae44c4b18f5eea35d533bf9ca7b529ae",
+ "model_name": "minimax_h3_dit",
+ "model_class": "diffsynth.models.minimax_h3_dit.MiniMaxH3DiT",
+ "quant_config": {"method": "bitsandbytes_nf4", "load_prequantized": True, "exclude_modules": ["time_embedder.proj_in", "time_embedder.proj_out", "video_patch_proj", "audio_patch_proj", "condition_proj", "final_layer.video_out", "final_layer.audio_out"]},
+ },
+ {
+ # Example: ModelConfig(model_id="DiffSynth-Studio/MiniMax-H3-NF4", origin_file_pattern="minimax-h3-fl2va-pruned-nf4.safetensors")
+ "model_hash": "9b8f1cbc6fb5fcd2e2c02b6691480427",
+ "model_name": "minimax_h3_dit",
+ "model_class": "diffsynth.models.minimax_h3_dit_comfy.MiniMaxH3DiTComfyPruned",
+ "quant_config": {"method": "bitsandbytes_nf4", "load_prequantized": True, "exclude_modules": ["time_embedder.proj_in", "time_embedder.proj_out", "video_patch_proj", "audio_patch_proj", "condition_proj", "final_layer.video_out", "final_layer.audio_out", "adaln_proj.linear"]},
+ },
+ {
+ # Example: ModelConfig(model_id="DiffSynth-Studio/MiniMax-H3-NF4", origin_file_pattern="minimax-h3-ref2va-pruned-nf4.safetensors")
+ "model_hash": "072bd13242d6d67df79a154e7c4274b3",
+ "model_name": "minimax_h3_dit",
+ "model_class": "diffsynth.models.minimax_h3_dit_comfy.MiniMaxH3DiTComfyPruned",
+ "quant_config": {"method": "bitsandbytes_nf4", "load_prequantized": True, "exclude_modules": ["time_embedder.proj_in", "time_embedder.proj_out", "video_patch_proj", "audio_patch_proj", "condition_proj", "final_layer.video_out", "final_layer.audio_out", "adaln_proj.linear"]},
+ },
+ {
+ # Example: ModelConfig(model_id="MiniMax/MiniMax-H3-FL2VA", origin_file_pattern="video_vae/source/model.safetensors")
+ "model_hash": "24b80900992e2024fab17c991c57da23",
+ "model_name": "minimax_h3_video_vae",
+ "model_class": "diffsynth.models.minimax_h3_video_vae.MiniMaxH3VideoVAE",
+ "state_dict_converter": "diffsynth.utils.state_dict_converters.minimax_h3_video_vae.MiniMaxH3VideoVAEStateDictConverter",
+ },
+ {
+ # Example: ModelConfig(model_id="DiffSynth-Studio/MiniMax-H3-NF4", origin_file_pattern="video_vae_nf4.safetensors")
+ "model_hash": "204819338f4fa12936684eeda37fdf71",
+ "model_name": "minimax_h3_video_vae",
+ "model_class": "diffsynth.models.minimax_h3_video_vae.MiniMaxH3VideoVAE",
+ "quant_config": {"method": "bitsandbytes_nf4", "load_prequantized": True, "exclude_modules": ["decoder.x_embedder"]},
+ "state_dict_converter": "diffsynth.utils.state_dict_converters.minimax_h3_video_vae.MiniMaxH3VideoVAEStateDictConverter",
+ },
+ {
+ # Example: ModelConfig(model_id="DiffSynth-Studio/MiniMax-H3-NF4", origin_file_pattern="audio_vae_nf4.safetensors")
+ "model_hash": "925ae768f8f4a9daa902b789de2ad2dc",
+ "model_name": "minimax_h3_audio_vae",
+ "model_class": "diffsynth.models.minimax_h3_audio_vae.MiniMaxH3AudioVAE",
+ "quant_config": {"method": "bitsandbytes_nf4", "load_prequantized": True},
+ "state_dict_converter": "diffsynth.utils.state_dict_converters.minimax_h3_audio_vae.MiniMaxH3AudioVAEStateDictConverter"
+ },
+]
+
+minimax_music3_series = [
+ {
+ # Example: ModelConfig(model_id="MiniMax/MiniMax-Music3", origin_file_pattern="transformer/diffusion_pytorch_model*.safetensors")
+ "model_hash": "8d1f1aba6a5801220ccc3a531865f694",
+ "model_name": "minimax_music3_dit",
+ "model_class": "diffsynth.models.minimax_music3_dit.MiniMaxMusic3DiT",
+ },
+ {
+ # Example: ModelConfig(model_id="MiniMax/MiniMax-Music3", origin_file_pattern="condition_encoder/diffusion_pytorch_model.safetensors")
+ "model_hash": "1370ef6377206e5ceaa78b125b2d4e1f",
+ "model_name": "minimax_music3_condition_encoder",
+ "model_class": "diffsynth.models.minimax_music3_condition_encoder.MiniMaxMusic3ConditionEncoder",
+ },
+ {
+ # Example: ModelConfig(model_id="MiniMax/MiniMax-Music3", origin_file_pattern="vocoder/diffusion_pytorch_model.safetensors")
+ "model_hash": "b6eded627c26c85b39804628ba4e5d1d",
+ "model_name": "minimax_music3_vocoder",
+ "model_class": "diffsynth.models.minimax_music3_vocoder.MiniMaxMusic3Vocoder",
+ },
+ {
+ # Example: ModelConfig(model_id="MiniMax/MiniMax-Music3", origin_file_pattern="rvq_depth_decoder/diffusion_pytorch_model.safetensors")
+ "model_hash": "5e571ded532a131c02e9d8256cb3fded",
+ "model_name": "minimax_music3_rvq_depth_decoder",
+ "model_class": "diffsynth.models.minimax_music3_rvq_depth_decoder.MiniMaxMusic3RVQDepthDecoder",
+ },
+ {
+ # Example: ModelConfig(model_id="MiniMax/MiniMax-Music3", origin_file_pattern="language_model/model*.safetensors")
+ "model_hash": "474b2197ab43bd2d4b69746506a952ef",
+ "model_name": "minimax_music3_text_encoder",
+ "model_class": "diffsynth.models.minimax_music3_text_encoder.MiniMaxMusic3TextEncoder",
+ "state_dict_converter": "diffsynth.utils.state_dict_converters.minimax_music3_text_encoder.MiniMaxMusic3TextEncoderStateDictConverter",
+ },
+]
+
+MODEL_CONFIGS = (
+ stable_diffusion_xl_series + stable_diffusion_series + qwen_image_series + qwen_video_edit_series + wan_series + flux_series + flux2_series + ernie_image_series
+ + z_image_series + ltx2_series + anima_series + mova_series + joyai_image_series + boogu_image_series + ace_step_series + hidream_o1_image_series
+ + image_metrics_series + ideogram4_series + krea2_series + lingbot_video_series + minimax_h3_series + minimax_music3_series
+)
diff --git a/diffsynth/configs/vram_management_module_maps.py b/diffsynth/configs/vram_management_module_maps.py
new file mode 100644
index 0000000000000000000000000000000000000000..427a02eafaf9f8fccdf54dd4d39cb82d1d216f28
--- /dev/null
+++ b/diffsynth/configs/vram_management_module_maps.py
@@ -0,0 +1,517 @@
+flux_general_vram_config = {
+ "torch.nn.Linear": "diffsynth.core.vram.layers.AutoWrappedLinear",
+ "torch.nn.Embedding": "diffsynth.core.vram.layers.AutoWrappedModule",
+ "torch.nn.LayerNorm": "diffsynth.core.vram.layers.AutoWrappedModule",
+ "torch.nn.Conv2d": "diffsynth.core.vram.layers.AutoWrappedModule",
+ "torch.nn.GroupNorm": "diffsynth.core.vram.layers.AutoWrappedModule",
+ "diffsynth.models.general_modules.RMSNorm": "diffsynth.core.vram.layers.AutoWrappedModule",
+ "diffsynth.models.flux_lora_encoder.LoRALayerBlock": "diffsynth.core.vram.layers.AutoWrappedModule",
+ "diffsynth.models.flux_lora_patcher.LoraMerger": "diffsynth.core.vram.layers.AutoWrappedModule",
+}
+
+VRAM_MANAGEMENT_MODULE_MAPS = {
+ "diffsynth.models.minimax_h3_dit.MiniMaxH3DiT": {
+ "torch.nn.Linear": "diffsynth.core.vram.layers.AutoWrappedLinear",
+ "torch.nn.RMSNorm": "diffsynth.core.vram.layers.AutoWrappedModule",
+ },
+ "diffsynth.models.minimax_h3_dit_comfy.MiniMaxH3DiTComfy": {
+ "torch.nn.Linear": "diffsynth.core.vram.layers.AutoWrappedLinear",
+ "torch.nn.RMSNorm": "diffsynth.core.vram.layers.AutoWrappedModule",
+ },
+ "diffsynth.models.minimax_h3_dit_comfy.MiniMaxH3DiTComfyPruned": {
+ "torch.nn.Linear": "diffsynth.core.vram.layers.AutoWrappedLinear",
+ "torch.nn.RMSNorm": "diffsynth.core.vram.layers.AutoWrappedModule",
+ },
+ "diffsynth.models.minimax_h3_text_encoder.MiniMaxH3TextEncoder": {
+ "torch.nn.Linear": "diffsynth.core.vram.layers.AutoWrappedLinear",
+ "torch.nn.Embedding": "diffsynth.core.vram.layers.AutoWrappedModule",
+ "torch.nn.LayerNorm": "diffsynth.core.vram.layers.AutoWrappedModule",
+ "torch.nn.Conv2d": "diffsynth.core.vram.layers.AutoWrappedModule",
+ "torch.nn.Conv3d": "diffsynth.core.vram.layers.AutoWrappedModule",
+ "transformers.models.qwen3_vl.modeling_qwen3_vl.Qwen3VLTextRMSNorm": "diffsynth.core.vram.layers.AutoWrappedModule",
+ "transformers.models.qwen3_vl.modeling_qwen3_vl.Qwen3VLTextRotaryEmbedding": "diffsynth.core.vram.layers.AutoWrappedModule",
+ "transformers.models.qwen3_vl.modeling_qwen3_vl.Qwen3VLVisionModel": "diffsynth.core.vram.layers.AutoWrappedModule",
+ "transformers.models.qwen3_vl.modeling_qwen3_vl.Qwen3VLVisionPatchEmbed": "diffsynth.core.vram.layers.AutoWrappedModule",
+ "transformers.models.qwen3_vl.modeling_qwen3_vl.Qwen3VLVisionRotaryEmbedding": "diffsynth.core.vram.layers.AutoWrappedModule",
+ },
+ "diffsynth.models.minimax_h3_video_vae.MiniMaxH3VideoVAE": {
+ "diffsynth.models.minimax_h3_video_vae.WarpedTensor": "diffsynth.core.vram.layers.AutoWrappedModule",
+ "torch.nn.Linear": "diffsynth.core.vram.layers.AutoWrappedLinear",
+ "torch.nn.Embedding": "diffsynth.core.vram.layers.AutoWrappedModule",
+ "torch.nn.LayerNorm": "diffsynth.core.vram.layers.AutoWrappedModule",
+ "torch.nn.RMSNorm": "diffsynth.core.vram.layers.AutoWrappedModule",
+ "torch.nn.Conv3d": "diffsynth.core.vram.layers.AutoWrappedModule",
+ "diffsynth.models.minimax_h3_video_vae.BaseConv3d": "diffsynth.core.vram.layers.AutoWrappedModule",
+ "diffsynth.models.minimax_h3_video_vae.TemporalIsolatedSpatialGroupNorm": "diffsynth.core.vram.layers.AutoWrappedModule",
+ },
+ "diffsynth.models.minimax_h3_audio_vae.MiniMaxH3AudioVAE": {
+ "diffsynth.models.minimax_h3_audio_vae.WarpedTensor": "diffsynth.core.vram.layers.AutoWrappedModule",
+ "torch.nn.Linear": "diffsynth.core.vram.layers.AutoWrappedLinear",
+ "torch.nn.LayerNorm": "diffsynth.core.vram.layers.AutoWrappedModule",
+ "torch.nn.RMSNorm": "diffsynth.core.vram.layers.AutoWrappedModule",
+ "torch.nn.Conv1d": "diffsynth.core.vram.layers.AutoWrappedModule",
+ "torch.nn.Conv3d": "diffsynth.core.vram.layers.AutoWrappedModule",
+ "diffsynth.models.minimax_h3_audio_vae.Snake": "diffsynth.core.vram.layers.AutoWrappedModule",
+ "diffsynth.models.minimax_h3_audio_vae.WeightNormedConv1d": "diffsynth.core.vram.layers.AutoWrappedModule",
+ "diffsynth.models.minimax_h3_audio_vae.WeightNormedConvTranspose1d": "diffsynth.core.vram.layers.AutoWrappedModule",
+ "diffsynth.models.minimax_h3_audio_vae.UpSample1d": "diffsynth.core.vram.layers.AutoWrappedModule",
+ "diffsynth.models.minimax_h3_audio_vae.LowPassFilter1d": "diffsynth.core.vram.layers.AutoWrappedModule",
+ },
+ "diffsynth.models.qwen_image_dit.QwenImageDiT": {
+ "diffsynth.models.qwen_image_dit.RMSNorm": "diffsynth.core.vram.layers.AutoWrappedModule",
+ "torch.nn.Linear": "diffsynth.core.vram.layers.AutoWrappedLinear",
+ "torch.nn.Embedding": "diffsynth.core.vram.layers.AutoWrappedModule",
+ },
+ "diffsynth.models.qwen_video_edit_dit.QwenVideoEditAdapter": {
+ "torch.nn.Linear": "diffsynth.core.vram.layers.AutoWrappedLinear",
+ "torch.nn.Conv3d": "diffsynth.core.vram.layers.AutoWrappedModule",
+ },
+ "diffsynth.models.qwen_image_text_encoder.QwenImageTextEncoder": {
+ "torch.nn.Linear": "diffsynth.core.vram.layers.AutoWrappedLinear",
+ "torch.nn.Embedding": "diffsynth.core.vram.layers.AutoWrappedModule",
+ "transformers.models.qwen2_5_vl.modeling_qwen2_5_vl.Qwen2_5_VLRotaryEmbedding": "diffsynth.core.vram.layers.AutoWrappedModule",
+ "transformers.models.qwen2_5_vl.modeling_qwen2_5_vl.Qwen2RMSNorm": "diffsynth.core.vram.layers.AutoWrappedModule",
+ "transformers.models.qwen2_5_vl.modeling_qwen2_5_vl.Qwen2_5_VisionPatchEmbed": "diffsynth.core.vram.layers.AutoWrappedModule",
+ "transformers.models.qwen2_5_vl.modeling_qwen2_5_vl.Qwen2_5_VisionRotaryEmbedding": "diffsynth.core.vram.layers.AutoWrappedModule",
+ },
+ "diffsynth.models.qwen_image_vae.QwenImageVAE": {
+ "torch.nn.Linear": "diffsynth.core.vram.layers.AutoWrappedLinear",
+ "torch.nn.Conv3d": "diffsynth.core.vram.layers.AutoWrappedModule",
+ "torch.nn.Conv2d": "diffsynth.core.vram.layers.AutoWrappedModule",
+ "diffsynth.models.qwen_image_vae.QwenImageRMS_norm": "diffsynth.core.vram.layers.AutoWrappedModule",
+ },
+ "diffsynth.models.qwen_image_controlnet.BlockWiseControlBlock": {
+ "diffsynth.models.qwen_image_dit.RMSNorm": "diffsynth.core.vram.layers.AutoWrappedModule",
+ "torch.nn.Linear": "diffsynth.core.vram.layers.AutoWrappedLinear",
+ },
+ "diffsynth.models.siglip2_image_encoder.Siglip2ImageEncoder": {
+ "transformers.models.siglip.modeling_siglip.SiglipVisionEmbeddings": "diffsynth.core.vram.layers.AutoWrappedModule",
+ "transformers.models.siglip.modeling_siglip.SiglipMultiheadAttentionPoolingHead": "diffsynth.core.vram.layers.AutoWrappedModule",
+ "torch.nn.Conv2d": "diffsynth.core.vram.layers.AutoWrappedModule",
+ "torch.nn.Embedding": "diffsynth.core.vram.layers.AutoWrappedModule",
+ "torch.nn.LayerNorm": "diffsynth.core.vram.layers.AutoWrappedModule",
+ "torch.nn.Linear": "diffsynth.core.vram.layers.AutoWrappedLinear",
+ },
+ "diffsynth.models.dinov3_image_encoder.DINOv3ImageEncoder": {
+ "transformers.models.dinov3_vit.modeling_dinov3_vit.DINOv3ViTLayerScale": "diffsynth.core.vram.layers.AutoWrappedModule",
+ "transformers.models.dinov3_vit.modeling_dinov3_vit.DINOv3ViTRopePositionEmbedding": "diffsynth.core.vram.layers.AutoWrappedModule",
+ "transformers.models.dinov3_vit.modeling_dinov3_vit.DINOv3ViTEmbeddings": "diffsynth.core.vram.layers.AutoWrappedModule",
+ "torch.nn.Conv2d": "diffsynth.core.vram.layers.AutoWrappedModule",
+ "torch.nn.LayerNorm": "diffsynth.core.vram.layers.AutoWrappedModule",
+ "torch.nn.Linear": "diffsynth.core.vram.layers.AutoWrappedLinear",
+ },
+ "diffsynth.models.qwen_image_image2lora.QwenImageImage2LoRAModel": {
+ "torch.nn.Linear": "diffsynth.core.vram.layers.AutoWrappedLinear",
+ },
+ "diffsynth.models.wan_video_animate_adapter.WanAnimateAdapter": {
+ "diffsynth.models.wan_video_animate_adapter.FaceEncoder": "diffsynth.core.vram.layers.AutoWrappedModule",
+ "diffsynth.models.wan_video_animate_adapter.EqualLinear": "diffsynth.core.vram.layers.AutoWrappedModule",
+ "diffsynth.models.wan_video_animate_adapter.ConvLayer": "diffsynth.core.vram.layers.AutoWrappedModule",
+ "diffsynth.models.wan_video_animate_adapter.FusedLeakyReLU": "diffsynth.core.vram.layers.AutoWrappedModule",
+ "diffsynth.models.wan_video_animate_adapter.RMSNorm": "diffsynth.core.vram.layers.AutoWrappedModule",
+ "torch.nn.Linear": "diffsynth.core.vram.layers.AutoWrappedLinear",
+ "torch.nn.LayerNorm": "diffsynth.core.vram.layers.AutoWrappedModule",
+ "torch.nn.Conv1d": "diffsynth.core.vram.layers.AutoWrappedModule",
+ "torch.nn.Conv2d": "diffsynth.core.vram.layers.AutoWrappedModule",
+ "torch.nn.Conv3d": "diffsynth.core.vram.layers.AutoWrappedModule",
+ },
+ "diffsynth.models.wan_video_dit_s2v.WanS2VModel": {
+ "diffsynth.models.wan_video_dit.Head": "diffsynth.core.vram.layers.AutoWrappedModule",
+ "diffsynth.models.wan_video_dit_s2v.WanS2VDiTBlock": "diffsynth.core.vram.layers.AutoWrappedModule",
+ "diffsynth.models.wan_video_dit_s2v.CausalAudioEncoder": "diffsynth.core.vram.layers.AutoWrappedModule",
+ "torch.nn.Embedding": "diffsynth.core.vram.layers.AutoWrappedModule",
+ "torch.nn.Linear": "diffsynth.core.vram.layers.AutoWrappedLinear",
+ "torch.nn.Conv3d": "diffsynth.core.vram.layers.AutoWrappedModule",
+ "torch.nn.LayerNorm": "diffsynth.core.vram.layers.AutoWrappedModule",
+ "diffsynth.models.wan_video_dit.RMSNorm": "diffsynth.core.vram.layers.AutoWrappedModule",
+ "torch.nn.Conv2d": "diffsynth.core.vram.layers.AutoWrappedModule",
+ },
+ "diffsynth.models.wan_animate_2_dit.WanAnimate2Transformer": {
+ "diffsynth.models.wan_video_dit.MLP": "diffsynth.core.vram.layers.AutoWrappedModule",
+ "diffsynth.models.wan_animate_2_dit.AttentionBlock": "diffsynth.core.vram.layers.AutoWrappedNonRecurseModule",
+ "diffsynth.models.wan_animate_2_dit.Head": "diffsynth.core.vram.layers.AutoWrappedModule",
+ "torch.nn.Linear": "diffsynth.core.vram.layers.AutoWrappedLinear",
+ "torch.nn.Conv3d": "diffsynth.core.vram.layers.AutoWrappedModule",
+ "torch.nn.LayerNorm": "diffsynth.core.vram.layers.AutoWrappedModule",
+ "diffsynth.models.wan_video_dit.RMSNorm": "diffsynth.core.vram.layers.AutoWrappedModule",
+ },
+ "diffsynth.models.wan_video_dit.WanModel": {
+ "diffsynth.models.wan_video_dit.MLP": "diffsynth.core.vram.layers.AutoWrappedModule",
+ "diffsynth.models.wan_video_dit.DiTBlock": "diffsynth.core.vram.layers.AutoWrappedNonRecurseModule",
+ "diffsynth.models.wan_video_dit.Head": "diffsynth.core.vram.layers.AutoWrappedModule",
+ "torch.nn.Linear": "diffsynth.core.vram.layers.AutoWrappedLinear",
+ "torch.nn.Conv3d": "diffsynth.core.vram.layers.AutoWrappedModule",
+ "torch.nn.LayerNorm": "diffsynth.core.vram.layers.AutoWrappedModule",
+ "diffsynth.models.wan_video_dit.RMSNorm": "diffsynth.core.vram.layers.AutoWrappedModule",
+ "torch.nn.Conv2d": "diffsynth.core.vram.layers.AutoWrappedModule",
+ },
+ "diffsynth.models.wan_video_image_encoder.WanImageEncoder": {
+ "diffsynth.models.wan_video_image_encoder.VisionTransformer": "diffsynth.core.vram.layers.AutoWrappedModule",
+ "torch.nn.Linear": "diffsynth.core.vram.layers.AutoWrappedLinear",
+ "torch.nn.Conv2d": "diffsynth.core.vram.layers.AutoWrappedModule",
+ "torch.nn.LayerNorm": "diffsynth.core.vram.layers.AutoWrappedModule",
+ },
+ "diffsynth.models.wan_video_mot.MotWanModel": {
+ "diffsynth.models.wan_video_mot.MotWanAttentionBlock": "diffsynth.core.vram.layers.AutoWrappedModule",
+ "torch.nn.Conv3d": "diffsynth.core.vram.layers.AutoWrappedModule",
+ "torch.nn.Linear": "diffsynth.core.vram.layers.AutoWrappedLinear",
+ "torch.nn.LayerNorm": "diffsynth.core.vram.layers.AutoWrappedModule",
+ },
+ "diffsynth.models.wan_video_motion_controller.WanMotionControllerModel": {
+ "torch.nn.Linear": "diffsynth.core.vram.layers.AutoWrappedLinear",
+ },
+ "diffsynth.models.wan_video_text_encoder.WanTextEncoder": {
+ "torch.nn.Linear": "diffsynth.core.vram.layers.AutoWrappedLinear",
+ "torch.nn.Embedding": "diffsynth.core.vram.layers.AutoWrappedModule",
+ "diffsynth.models.wan_video_text_encoder.T5RelativeEmbedding": "diffsynth.core.vram.layers.AutoWrappedModule",
+ "diffsynth.models.wan_video_text_encoder.T5LayerNorm": "diffsynth.core.vram.layers.AutoWrappedModule",
+ },
+ "diffsynth.models.wan_video_vace.VaceWanModel": {
+ "diffsynth.models.wan_video_dit.DiTBlock": "diffsynth.core.vram.layers.AutoWrappedModule",
+ "torch.nn.Linear": "diffsynth.core.vram.layers.AutoWrappedLinear",
+ "torch.nn.Conv3d": "diffsynth.core.vram.layers.AutoWrappedModule",
+ "torch.nn.LayerNorm": "diffsynth.core.vram.layers.AutoWrappedModule",
+ "diffsynth.models.wan_video_dit.RMSNorm": "diffsynth.core.vram.layers.AutoWrappedModule",
+ },
+ "diffsynth.models.wan_video_vae.WanVideoVAE": {
+ "torch.nn.Linear": "diffsynth.core.vram.layers.AutoWrappedLinear",
+ "torch.nn.Conv2d": "diffsynth.core.vram.layers.AutoWrappedModule",
+ "diffsynth.models.wan_video_vae.RMS_norm": "diffsynth.core.vram.layers.AutoWrappedModule",
+ "diffsynth.models.wan_video_vae.CausalConv3d": "diffsynth.core.vram.layers.AutoWrappedModule",
+ "diffsynth.models.wan_video_vae.Upsample": "diffsynth.core.vram.layers.AutoWrappedModule",
+ "torch.nn.SiLU": "diffsynth.core.vram.layers.AutoWrappedModule",
+ "torch.nn.Dropout": "diffsynth.core.vram.layers.AutoWrappedModule",
+ },
+ "diffsynth.models.wan_video_vae.WanVideoVAE38": {
+ "torch.nn.Linear": "diffsynth.core.vram.layers.AutoWrappedLinear",
+ "torch.nn.Conv2d": "diffsynth.core.vram.layers.AutoWrappedModule",
+ "diffsynth.models.wan_video_vae.RMS_norm": "diffsynth.core.vram.layers.AutoWrappedModule",
+ "diffsynth.models.wan_video_vae.CausalConv3d": "diffsynth.core.vram.layers.AutoWrappedModule",
+ "diffsynth.models.wan_video_vae.Upsample": "diffsynth.core.vram.layers.AutoWrappedModule",
+ "torch.nn.SiLU": "diffsynth.core.vram.layers.AutoWrappedModule",
+ "torch.nn.Dropout": "diffsynth.core.vram.layers.AutoWrappedModule",
+ },
+ "diffsynth.models.wav2vec.WanS2VAudioEncoder": {
+ "torch.nn.Linear": "diffsynth.core.vram.layers.AutoWrappedLinear",
+ "torch.nn.LayerNorm": "diffsynth.core.vram.layers.AutoWrappedModule",
+ "torch.nn.Conv1d": "diffsynth.core.vram.layers.AutoWrappedModule",
+ },
+ "diffsynth.models.longcat_video_dit.LongCatVideoTransformer3DModel": {
+ "torch.nn.Linear": "diffsynth.core.vram.layers.AutoWrappedLinear",
+ "torch.nn.Conv3d": "diffsynth.core.vram.layers.AutoWrappedModule",
+ "diffsynth.models.longcat_video_dit.RMSNorm_FP32": "diffsynth.core.vram.layers.AutoWrappedModule",
+ "diffsynth.models.longcat_video_dit.LayerNorm_FP32": "diffsynth.core.vram.layers.AutoWrappedModule",
+ },
+ "diffsynth.models.flux_dit.FluxDiT": {
+ "torch.nn.Linear": "diffsynth.core.vram.layers.AutoWrappedLinear",
+ "diffsynth.models.flux_dit.RMSNorm": "diffsynth.core.vram.layers.AutoWrappedModule",
+ },
+ "diffsynth.models.flux_text_encoder_clip.FluxTextEncoderClip": flux_general_vram_config,
+ "diffsynth.models.flux_vae.FluxVAEEncoder": flux_general_vram_config,
+ "diffsynth.models.flux_vae.FluxVAEDecoder": flux_general_vram_config,
+ "diffsynth.models.flux_controlnet.FluxControlNet": flux_general_vram_config,
+ "diffsynth.models.flux_infiniteyou.InfiniteYouImageProjector": flux_general_vram_config,
+ "diffsynth.models.flux_ipadapter.FluxIpAdapter": flux_general_vram_config,
+ "diffsynth.models.flux_lora_patcher.FluxLoraPatcher": flux_general_vram_config,
+ "diffsynth.models.step1x_connector.Qwen2Connector": flux_general_vram_config,
+ "diffsynth.models.flux_lora_encoder.FluxLoRAEncoder": flux_general_vram_config,
+ "diffsynth.models.flux_text_encoder_t5.FluxTextEncoderT5": {
+ "torch.nn.Linear": "diffsynth.core.vram.layers.AutoWrappedLinear",
+ "torch.nn.Embedding": "diffsynth.core.vram.layers.AutoWrappedModule",
+ "transformers.models.t5.modeling_t5.T5LayerNorm": "diffsynth.core.vram.layers.AutoWrappedModule",
+ "transformers.models.t5.modeling_t5.T5DenseActDense": "diffsynth.core.vram.layers.AutoWrappedModule",
+ "transformers.models.t5.modeling_t5.T5DenseGatedActDense": "diffsynth.core.vram.layers.AutoWrappedModule",
+ },
+ "diffsynth.models.flux_ipadapter.SiglipVisionModelSO400M": {
+ "transformers.models.siglip.modeling_siglip.SiglipVisionEmbeddings": "diffsynth.core.vram.layers.AutoWrappedModule",
+ "transformers.models.siglip.modeling_siglip.SiglipEncoder": "diffsynth.core.vram.layers.AutoWrappedModule",
+ "transformers.models.siglip.modeling_siglip.SiglipMultiheadAttentionPoolingHead": "diffsynth.core.vram.layers.AutoWrappedModule",
+ "torch.nn.MultiheadAttention": "diffsynth.core.vram.layers.AutoWrappedModule",
+ "torch.nn.Linear": "diffsynth.core.vram.layers.AutoWrappedLinear",
+ "torch.nn.LayerNorm": "diffsynth.core.vram.layers.AutoWrappedModule",
+ },
+ "diffsynth.models.flux2_dit.Flux2DiT": {
+ "torch.nn.Linear": "diffsynth.core.vram.layers.AutoWrappedLinear",
+ "torch.nn.LayerNorm": "diffsynth.core.vram.layers.AutoWrappedModule",
+ "torch.nn.RMSNorm": "diffsynth.core.vram.layers.AutoWrappedModule",
+ },
+ "diffsynth.models.flux2_text_encoder.Flux2TextEncoder": {
+ "torch.nn.Linear": "diffsynth.core.vram.layers.AutoWrappedLinear",
+ "torch.nn.Conv2d": "diffsynth.core.vram.layers.AutoWrappedModule",
+ "torch.nn.Embedding": "diffsynth.core.vram.layers.AutoWrappedModule",
+ "transformers.models.mistral.modeling_mistral.MistralRMSNorm": "diffsynth.core.vram.layers.AutoWrappedModule",
+ },
+ "diffsynth.models.flux2_vae.Flux2VAE": {
+ "torch.nn.Linear": "diffsynth.core.vram.layers.AutoWrappedLinear",
+ "torch.nn.Conv2d": "diffsynth.core.vram.layers.AutoWrappedModule",
+ "torch.nn.GroupNorm": "diffsynth.core.vram.layers.AutoWrappedModule",
+ },
+ "diffsynth.models.z_image_text_encoder.ZImageTextEncoder": {
+ "torch.nn.Linear": "diffsynth.core.vram.layers.AutoWrappedLinear",
+ "transformers.models.qwen3.modeling_qwen3.Qwen3RMSNorm": "diffsynth.core.vram.layers.AutoWrappedModule",
+ "torch.nn.Embedding": "diffsynth.core.vram.layers.AutoWrappedModule",
+ },
+ "diffsynth.models.z_image_dit.ZImageDiT": {
+ "torch.nn.Linear": "diffsynth.core.vram.layers.AutoWrappedLinear",
+ "diffsynth.models.z_image_dit.RMSNorm": "diffsynth.core.vram.layers.AutoWrappedModule",
+ },
+ "diffsynth.models.z_image_controlnet.ZImageControlNet": {
+ "torch.nn.Linear": "diffsynth.core.vram.layers.AutoWrappedLinear",
+ "diffsynth.models.z_image_dit.RMSNorm": "diffsynth.core.vram.layers.AutoWrappedModule",
+ },
+ "diffsynth.models.z_image_image2lora.ZImageImage2LoRAModel": {
+ "torch.nn.Linear": "diffsynth.core.vram.layers.AutoWrappedLinear",
+ },
+ "diffsynth.models.siglip2_image_encoder.Siglip2ImageEncoder428M": {
+ "transformers.models.siglip2.modeling_siglip2.Siglip2VisionEmbeddings": "diffsynth.core.vram.layers.AutoWrappedModule",
+ "transformers.models.siglip2.modeling_siglip2.Siglip2MultiheadAttentionPoolingHead": "diffsynth.core.vram.layers.AutoWrappedModule",
+ "torch.nn.Conv2d": "diffsynth.core.vram.layers.AutoWrappedModule",
+ "torch.nn.Embedding": "diffsynth.core.vram.layers.AutoWrappedModule",
+ "torch.nn.LayerNorm": "diffsynth.core.vram.layers.AutoWrappedModule",
+ "torch.nn.Linear": "diffsynth.core.vram.layers.AutoWrappedLinear",
+ },
+ "diffsynth.models.ltx2_dit.LTXModel": {
+ "torch.nn.Linear": "diffsynth.core.vram.layers.AutoWrappedLinear",
+ "torch.nn.RMSNorm": "diffsynth.core.vram.layers.AutoWrappedModule",
+ },
+ "diffsynth.models.ltx2_upsampler.LTX2LatentUpsampler": {
+ "torch.nn.Conv2d": "diffsynth.core.vram.layers.AutoWrappedModule",
+ "torch.nn.Conv3d": "diffsynth.core.vram.layers.AutoWrappedModule",
+ "torch.nn.GroupNorm": "diffsynth.core.vram.layers.AutoWrappedModule",
+ },
+ "diffsynth.models.ltx2_video_vae.LTX2VideoEncoder": {
+ "torch.nn.Conv3d": "diffsynth.core.vram.layers.AutoWrappedModule",
+ },
+ "diffsynth.models.ltx2_video_vae.LTX2VideoDecoder": {
+ "torch.nn.Conv3d": "diffsynth.core.vram.layers.AutoWrappedModule",
+ },
+ "diffsynth.models.ltx2_audio_vae.LTX2AudioDecoder": {
+ "torch.nn.Conv2d": "diffsynth.core.vram.layers.AutoWrappedModule",
+ },
+ "diffsynth.models.ltx2_audio_vae.LTX2Vocoder": {
+ "torch.nn.Conv1d": "diffsynth.core.vram.layers.AutoWrappedModule",
+ "torch.nn.ConvTranspose1d": "diffsynth.core.vram.layers.AutoWrappedModule",
+ },
+ "diffsynth.models.ltx2_text_encoder.LTX2TextEncoderPostModules": {
+ "torch.nn.Linear": "diffsynth.core.vram.layers.AutoWrappedLinear",
+ "torch.nn.RMSNorm": "diffsynth.core.vram.layers.AutoWrappedModule",
+ "diffsynth.models.ltx2_text_encoder.Embeddings1DConnector": "diffsynth.core.vram.layers.AutoWrappedModule",
+ },
+ "diffsynth.models.ltx2_text_encoder.LTX2TextEncoder": {
+ "torch.nn.Linear": "diffsynth.core.vram.layers.AutoWrappedLinear",
+ "transformers.models.gemma3.modeling_gemma3.Gemma3MultiModalProjector": "diffsynth.core.vram.layers.AutoWrappedModule",
+ "transformers.models.gemma3.modeling_gemma3.Gemma3RMSNorm": "diffsynth.core.vram.layers.AutoWrappedModule",
+ "transformers.models.gemma3.modeling_gemma3.Gemma3TextScaledWordEmbedding": "diffsynth.core.vram.layers.AutoWrappedModule",
+ },
+ "diffsynth.models.anima_dit.AnimaDiT": {
+ "torch.nn.Linear": "diffsynth.core.vram.layers.AutoWrappedLinear",
+ "torch.nn.LayerNorm": "diffsynth.core.vram.layers.AutoWrappedModule",
+ "torch.nn.RMSNorm": "diffsynth.core.vram.layers.AutoWrappedModule",
+ "torch.nn.Embedding": "diffsynth.core.vram.layers.AutoWrappedModule",
+ },
+ "diffsynth.models.mova_audio_dit.MovaAudioDit": {
+ "diffsynth.models.wan_video_dit.DiTBlock": "diffsynth.core.vram.layers.AutoWrappedNonRecurseModule",
+ "diffsynth.models.wan_video_dit.Head": "diffsynth.core.vram.layers.AutoWrappedModule",
+ "torch.nn.Linear": "diffsynth.core.vram.layers.AutoWrappedLinear",
+ "torch.nn.Conv1d": "diffsynth.core.vram.layers.AutoWrappedModule",
+ "torch.nn.LayerNorm": "diffsynth.core.vram.layers.AutoWrappedModule",
+ "diffsynth.models.wan_video_dit.RMSNorm": "diffsynth.core.vram.layers.AutoWrappedModule",
+ },
+ "diffsynth.models.mova_dual_tower_bridge.DualTowerConditionalBridge": {
+ "torch.nn.Linear": "diffsynth.core.vram.layers.AutoWrappedLinear",
+ "torch.nn.LayerNorm": "diffsynth.core.vram.layers.AutoWrappedModule",
+ "diffsynth.models.wan_video_dit.RMSNorm": "diffsynth.core.vram.layers.AutoWrappedModule",
+ },
+ "diffsynth.models.mova_audio_vae.DacVAE": {
+ "diffsynth.models.mova_audio_vae.Snake1d": "diffsynth.core.vram.layers.AutoWrappedModule",
+ "torch.nn.Conv1d": "diffsynth.core.vram.layers.AutoWrappedModule",
+ "torch.nn.ConvTranspose1d": "diffsynth.core.vram.layers.AutoWrappedModule",
+ },
+ "diffsynth.models.ernie_image_dit.ErnieImageDiT": {
+ "diffsynth.models.ernie_image_dit.ErnieImageRMSNorm": "diffsynth.core.vram.layers.AutoWrappedModule",
+ "torch.nn.Linear": "diffsynth.core.vram.layers.AutoWrappedLinear",
+ "torch.nn.Conv2d": "diffsynth.core.vram.layers.AutoWrappedModule",
+ "torch.nn.LayerNorm": "diffsynth.core.vram.layers.AutoWrappedModule",
+ "torch.nn.RMSNorm": "diffsynth.core.vram.layers.AutoWrappedModule",
+ },
+ "diffsynth.models.ernie_image_text_encoder.ErnieImageTextEncoder": {
+ "torch.nn.Linear": "diffsynth.core.vram.layers.AutoWrappedLinear",
+ "torch.nn.Embedding": "diffsynth.core.vram.layers.AutoWrappedModule",
+ "transformers.models.ministral3.modeling_ministral3.Ministral3RMSNorm": "diffsynth.core.vram.layers.AutoWrappedModule",
+ },
+ "diffsynth.models.joyai_image_dit.Transformer3DModel": {
+ "diffsynth.models.joyai_image_dit.RMSNorm": "diffsynth.core.vram.layers.AutoWrappedModule",
+ "diffsynth.models.joyai_image_dit.ModulateWan": "diffsynth.core.vram.layers.AutoWrappedModule",
+ "torch.nn.Linear": "diffsynth.core.vram.layers.AutoWrappedLinear",
+ "torch.nn.Conv3d": "diffsynth.core.vram.layers.AutoWrappedModule",
+ "torch.nn.LayerNorm": "diffsynth.core.vram.layers.AutoWrappedModule",
+ },
+ "diffsynth.models.joyai_image_text_encoder.JoyAIImageTextEncoder": {
+ "torch.nn.Linear": "diffsynth.core.vram.layers.AutoWrappedLinear",
+ "torch.nn.Embedding": "diffsynth.core.vram.layers.AutoWrappedModule",
+ "torch.nn.LayerNorm": "diffsynth.core.vram.layers.AutoWrappedModule",
+ "torch.nn.Conv3d": "diffsynth.core.vram.layers.AutoWrappedModule",
+ "transformers.models.qwen3_vl.modeling_qwen3_vl.Qwen3VLVisionModel": "diffsynth.core.vram.layers.AutoWrappedModule",
+ "transformers.models.qwen3_vl.modeling_qwen3_vl.Qwen3VLTextRMSNorm": "diffsynth.core.vram.layers.AutoWrappedModule",
+ "transformers.models.qwen3_vl.modeling_qwen3_vl.Qwen3VLTextRotaryEmbedding": "diffsynth.core.vram.layers.AutoWrappedModule",
+ },
+ "diffsynth.models.boogu_image_dit.BooguImageDiT": {
+ "torch.nn.Linear": "diffsynth.core.vram.layers.AutoWrappedLinear",
+ "torch.nn.RMSNorm": "diffsynth.core.vram.layers.AutoWrappedModule",
+ "torch.nn.LayerNorm": "diffsynth.core.vram.layers.AutoWrappedModule",
+ "torch.nn.Embedding": "diffsynth.core.vram.layers.AutoWrappedModule",
+ },
+ "diffsynth.models.ace_step_dit.AceStepDiTModel": {
+ "diffsynth.models.ace_step_dit.AceStepDiTLayer": "diffsynth.core.vram.layers.AutoWrappedNonRecurseModule",
+ "torch.nn.Linear": "diffsynth.core.vram.layers.AutoWrappedLinear",
+ "torch.nn.Conv1d": "diffsynth.core.vram.layers.AutoWrappedModule",
+ "torch.nn.ConvTranspose1d": "diffsynth.core.vram.layers.AutoWrappedModule",
+ "transformers.models.qwen3.modeling_qwen3.Qwen3RMSNorm": "diffsynth.core.vram.layers.AutoWrappedModule",
+ "transformers.models.qwen3.modeling_qwen3.Qwen3RotaryEmbedding": "diffsynth.core.vram.layers.AutoWrappedModule",
+ },
+ "diffsynth.models.ace_step_conditioner.AceStepConditionEncoder": {
+ "torch.nn.Linear": "diffsynth.core.vram.layers.AutoWrappedLinear",
+ "torch.nn.Embedding": "diffsynth.core.vram.layers.AutoWrappedModule",
+ "transformers.models.qwen3.modeling_qwen3.Qwen3RMSNorm": "diffsynth.core.vram.layers.AutoWrappedModule",
+ "transformers.models.qwen3.modeling_qwen3.Qwen3RotaryEmbedding": "diffsynth.core.vram.layers.AutoWrappedModule",
+ },
+ "diffsynth.models.ace_step_text_encoder.AceStepTextEncoder": {
+ "torch.nn.Linear": "diffsynth.core.vram.layers.AutoWrappedLinear",
+ "torch.nn.Embedding": "diffsynth.core.vram.layers.AutoWrappedModule",
+ "transformers.models.qwen3.modeling_qwen3.Qwen3RMSNorm": "diffsynth.core.vram.layers.AutoWrappedModule",
+ "transformers.models.qwen3.modeling_qwen3.Qwen3RotaryEmbedding": "diffsynth.core.vram.layers.AutoWrappedModule",
+ },
+ "diffsynth.models.ace_step_vae.AceStepVAE": {
+ "torch.nn.Linear": "diffsynth.core.vram.layers.AutoWrappedLinear",
+ "torch.nn.Conv1d": "diffsynth.core.vram.layers.AutoWrappedModule",
+ "torch.nn.ConvTranspose1d": "diffsynth.core.vram.layers.AutoWrappedModule",
+ "diffsynth.models.ace_step_vae.Snake1d": "diffsynth.core.vram.layers.AutoWrappedModule",
+ },
+ "diffsynth.models.ace_step_tokenizer.AceStepTokenizer": {
+ "torch.nn.Linear": "diffsynth.core.vram.layers.AutoWrappedLinear",
+ "torch.nn.Embedding": "diffsynth.core.vram.layers.AutoWrappedModule",
+ "diffsynth.models.ace_step_residual_fsq.ResidualFSQ": "diffsynth.core.vram.layers.AutoWrappedModule",
+ "transformers.models.qwen3.modeling_qwen3.Qwen3RMSNorm": "diffsynth.core.vram.layers.AutoWrappedModule",
+ "transformers.models.qwen3.modeling_qwen3.Qwen3MLP": "diffsynth.core.vram.layers.AutoWrappedModule",
+ "transformers.models.qwen3.modeling_qwen3.Qwen3RotaryEmbedding": "diffsynth.core.vram.layers.AutoWrappedModule",
+ },
+ "diffsynth.models.stable_diffusion_unet.UNet2DConditionModel": {
+ "torch.nn.Linear": "diffsynth.core.vram.layers.AutoWrappedLinear",
+ "torch.nn.Conv2d": "diffsynth.core.vram.layers.AutoWrappedModule",
+ "torch.nn.GroupNorm": "diffsynth.core.vram.layers.AutoWrappedModule",
+ "torch.nn.LayerNorm": "diffsynth.core.vram.layers.AutoWrappedModule",
+ "torch.nn.SiLU": "diffsynth.core.vram.layers.AutoWrappedModule",
+ "torch.nn.Dropout": "diffsynth.core.vram.layers.AutoWrappedModule",
+ },
+ "diffsynth.models.stable_diffusion_vae.StableDiffusionVAE": {
+ "torch.nn.Linear": "diffsynth.core.vram.layers.AutoWrappedLinear",
+ "torch.nn.Conv2d": "diffsynth.core.vram.layers.AutoWrappedModule",
+ "torch.nn.GroupNorm": "diffsynth.core.vram.layers.AutoWrappedModule",
+ "torch.nn.SiLU": "diffsynth.core.vram.layers.AutoWrappedModule",
+ "torch.nn.Dropout": "diffsynth.core.vram.layers.AutoWrappedModule",
+ "diffsynth.models.stable_diffusion_vae.Upsample2D": "diffsynth.core.vram.layers.AutoWrappedModule",
+ "diffsynth.models.stable_diffusion_vae.Downsample2D": "diffsynth.core.vram.layers.AutoWrappedModule",
+ },
+ "diffsynth.models.stable_diffusion_text_encoder.SDTextEncoder": {
+ "torch.nn.Linear": "diffsynth.core.vram.layers.AutoWrappedLinear",
+ "torch.nn.Embedding": "diffsynth.core.vram.layers.AutoWrappedModule",
+ "torch.nn.LayerNorm": "diffsynth.core.vram.layers.AutoWrappedModule",
+ "diffsynth.models.stable_diffusion_text_encoder.CLIPEncoderLayer": "diffsynth.core.vram.layers.AutoWrappedModule",
+ "diffsynth.models.stable_diffusion_text_encoder.CLIPAttention": "diffsynth.core.vram.layers.AutoWrappedModule",
+ },
+ "diffsynth.models.stable_diffusion_xl_unet.SDXLUNet2DConditionModel": {
+ "torch.nn.Linear": "diffsynth.core.vram.layers.AutoWrappedLinear",
+ "torch.nn.Conv2d": "diffsynth.core.vram.layers.AutoWrappedModule",
+ "torch.nn.GroupNorm": "diffsynth.core.vram.layers.AutoWrappedModule",
+ "torch.nn.LayerNorm": "diffsynth.core.vram.layers.AutoWrappedModule",
+ "torch.nn.SiLU": "diffsynth.core.vram.layers.AutoWrappedModule",
+ "torch.nn.Dropout": "diffsynth.core.vram.layers.AutoWrappedModule",
+ },
+ "diffsynth.models.stable_diffusion_xl_text_encoder.SDXLTextEncoder2": {
+ "torch.nn.Linear": "diffsynth.core.vram.layers.AutoWrappedLinear",
+ "torch.nn.Embedding": "diffsynth.core.vram.layers.AutoWrappedModule",
+ "torch.nn.LayerNorm": "diffsynth.core.vram.layers.AutoWrappedModule",
+ "diffsynth.models.stable_diffusion_text_encoder.CLIPEncoderLayer": "diffsynth.core.vram.layers.AutoWrappedModule",
+ "diffsynth.models.stable_diffusion_text_encoder.CLIPAttention": "diffsynth.core.vram.layers.AutoWrappedModule",
+ "diffsynth.models.stable_diffusion_xl_text_encoder.CLIPTextModelWithProjection": "diffsynth.core.vram.layers.AutoWrappedModule",
+ },
+ "diffsynth.models.hidream_o1_image_dit.HiDreamO1ImageModel": {
+ "torch.nn.Linear": "diffsynth.core.vram.layers.AutoWrappedLinear",
+ "torch.nn.Embedding": "diffsynth.core.vram.layers.AutoWrappedModule",
+ "torch.nn.Conv3d": "diffsynth.core.vram.layers.AutoWrappedModule",
+ "torch.nn.LayerNorm": "diffsynth.core.vram.layers.AutoWrappedModule",
+ "diffsynth.models.hidream_o1_image_dit.Qwen3VLTextRMSNorm": "diffsynth.core.vram.layers.AutoWrappedModule",
+ "diffsynth.models.hidream_o1_image_dit.Qwen3VLVisionModel": "diffsynth.core.vram.layers.AutoWrappedModule",
+ },
+ "diffsynth.models.ideogram4_dit.Ideogram4DiT": {
+ "torch.nn.Linear": "diffsynth.core.vram.layers.AutoWrappedLinear",
+ "diffsynth.models.ideogram4_dit.Ideogram4RMSNorm": "diffsynth.core.vram.layers.AutoWrappedModule",
+ "torch.nn.Embedding": "diffsynth.core.vram.layers.AutoWrappedModule",
+ "torch.nn.LayerNorm": "diffsynth.core.vram.layers.AutoWrappedModule",
+ },
+ "diffsynth.models.ideogram4_text_encoder.Ideogram4TextEncoder": {
+ "torch.nn.Linear": "diffsynth.core.vram.layers.AutoWrappedLinear",
+ "torch.nn.Embedding": "diffsynth.core.vram.layers.AutoWrappedModule",
+ "torch.nn.LayerNorm": "diffsynth.core.vram.layers.AutoWrappedModule",
+ "transformers.models.qwen3_vl.modeling_qwen3_vl.Qwen3VLTextRotaryEmbedding": "diffsynth.core.vram.layers.AutoWrappedModule",
+ "transformers.models.qwen3_vl.modeling_qwen3_vl.Qwen3VLTextRMSNorm": "diffsynth.core.vram.layers.AutoWrappedModule",
+ },
+ "diffsynth.models.krea2_dit.SingleStreamDiT": {
+ "diffsynth.models.krea2_dit.RMSNorm": "diffsynth.core.vram.layers.AutoWrappedModule",
+ "torch.nn.Linear": "diffsynth.core.vram.layers.AutoWrappedLinear",
+ "diffsynth.models.krea2_dit.DoubleSharedModulation": "diffsynth.core.vram.layers.AutoWrappedModule",
+ "diffsynth.models.krea2_dit.SimpleModulation": "diffsynth.core.vram.layers.AutoWrappedModule",
+ },
+ "diffsynth.models.krea2_text_encoder.Krea2TextEncoder": {
+ "torch.nn.Linear": "diffsynth.core.vram.layers.AutoWrappedLinear",
+ "torch.nn.Embedding": "diffsynth.core.vram.layers.AutoWrappedModule",
+ "transformers.models.qwen3_vl.modeling_qwen3_vl.Qwen3VLTextRotaryEmbedding": "diffsynth.core.vram.layers.AutoWrappedModule",
+ "transformers.models.qwen3_vl.modeling_qwen3_vl.Qwen3VLTextRMSNorm": "diffsynth.core.vram.layers.AutoWrappedModule",
+ "torch.nn.LayerNorm": "diffsynth.core.vram.layers.AutoWrappedModule",
+ "transformers.models.qwen3_vl.modeling_qwen3_vl.Qwen3VLVisionPatchEmbed": "diffsynth.core.vram.layers.AutoWrappedModule",
+ "transformers.models.qwen3_vl.modeling_qwen3_vl.Qwen3VLVisionRotaryEmbedding": "diffsynth.core.vram.layers.AutoWrappedModule",
+ },
+ "diffsynth.models.lingbot_video_dit.LingBotVideoDiT": {
+ "diffsynth.models.lingbot_video_dit.LingBotVideoBlock": "diffsynth.core.vram.layers.AutoWrappedNonRecurseModule",
+ "diffsynth.models.lingbot_video_dit.LingBotVideoRouter": "diffsynth.core.vram.layers.AutoWrappedNonRecurseModule",
+ "diffsynth.models.lingbot_video_dit.LingBotVideoGroupedExperts": "diffsynth.core.vram.layers.AutoWrappedNonRecurseModule",
+ "diffsynth.models.lingbot_video_dit.LingBotVideoRMSNorm": "diffsynth.core.vram.layers.AutoWrappedModule",
+ "torch.nn.Linear": "diffsynth.core.vram.layers.AutoWrappedLinear",
+ },
+ "diffsynth.models.minimax_music3_dit.MiniMaxMusic3DiT": {
+ "torch.nn.Linear": "diffsynth.core.vram.layers.AutoWrappedLinear",
+ "torch.nn.Conv1d": "diffsynth.core.vram.layers.AutoWrappedModule",
+ "torch.nn.LayerNorm": "diffsynth.core.vram.layers.AutoWrappedModule",
+ "diffsynth.models.minimax_music3_dit.MiniMaxMusic3FourierEmbedding": "diffsynth.core.vram.layers.AutoWrappedModule",
+ },
+ "diffsynth.models.minimax_music3_rvq_depth_decoder.MiniMaxMusic3RVQDepthDecoder": {
+ "torch.nn.Linear": "diffsynth.core.vram.layers.AutoWrappedLinear",
+ "torch.nn.Embedding": "diffsynth.core.vram.layers.AutoWrappedModule",
+ "diffsynth.models.minimax_music3_rvq_depth_decoder.MiniMaxMusic3RMSNorm": "diffsynth.core.vram.layers.AutoWrappedModule",
+ },
+ "diffsynth.models.minimax_music3_vocoder.MiniMaxMusic3Vocoder": {
+ "torch.nn.Conv1d": "diffsynth.core.vram.layers.AutoWrappedModule",
+ "diffsynth.models.minimax_music3_vocoder.MiniMaxMusic3WeightNormConv1d": "diffsynth.core.vram.layers.AutoWrappedModule",
+ "diffsynth.models.minimax_music3_vocoder.MiniMaxMusic3WeightNormConvTranspose1d": "diffsynth.core.vram.layers.AutoWrappedModule",
+ "diffsynth.models.minimax_music3_vocoder.MiniMaxMusic3Snake1d": "diffsynth.core.vram.layers.AutoWrappedModule",
+ },
+ "diffsynth.models.minimax_music3_text_encoder.MiniMaxMusic3TextEncoder": {
+ "torch.nn.Linear": "diffsynth.core.vram.layers.AutoWrappedLinear",
+ "torch.nn.Embedding": "diffsynth.core.vram.layers.AutoWrappedModule",
+ "transformers.models.qwen3.modeling_qwen3.Qwen3RMSNorm": "diffsynth.core.vram.layers.AutoWrappedModule",
+ },
+}
+
+def QwenImageTextEncoder_Module_Map_Updater():
+ current = VRAM_MANAGEMENT_MODULE_MAPS["diffsynth.models.qwen_image_text_encoder.QwenImageTextEncoder"]
+ from packaging import version
+ import transformers
+ if version.parse(transformers.__version__) >= version.parse("5.2.0"):
+ # The Qwen2RMSNorm in transformers 5.2.0+ has been renamed to Qwen2_5_VLRMSNorm, so we need to update the module map accordingly
+ current.pop("transformers.models.qwen2_5_vl.modeling_qwen2_5_vl.Qwen2RMSNorm", None)
+ current["transformers.models.qwen2_5_vl.modeling_qwen2_5_vl.Qwen2_5_VLRMSNorm"] = "diffsynth.core.vram.layers.AutoWrappedModule"
+ return current
+
+VERSION_CHECKER_MAPS = {
+ "diffsynth.models.qwen_image_text_encoder.QwenImageTextEncoder": QwenImageTextEncoder_Module_Map_Updater,
+}
\ No newline at end of file
diff --git a/diffsynth/core/__init__.py b/diffsynth/core/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..8b1747471525a8c0d6c17131afd05778faa9bc2a
--- /dev/null
+++ b/diffsynth/core/__init__.py
@@ -0,0 +1,8 @@
+from .attention import *
+from .data import *
+from .gradient import *
+from .quant import *
+from .loader import *
+from .vram import *
+from .device import *
+from .offload_training import *
diff --git a/diffsynth/core/attention/__init__.py b/diffsynth/core/attention/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..45cf8a4382397aa4ff6558f37191c726d821b95a
--- /dev/null
+++ b/diffsynth/core/attention/__init__.py
@@ -0,0 +1 @@
+from .attention import attention_forward
diff --git a/diffsynth/core/attention/attention.py b/diffsynth/core/attention/attention.py
new file mode 100644
index 0000000000000000000000000000000000000000..334f57d870818effba9560ebbc23fa0243eec7d8
--- /dev/null
+++ b/diffsynth/core/attention/attention.py
@@ -0,0 +1,255 @@
+import torch, os, inspect
+from einops import rearrange, repeat
+
+if os.environ.get("DIFFSYNTH_FLASH_ATTN_KERNEL_REPO_ID") is not None:
+ try:
+ from kernels import get_kernel
+ kernel_module = get_kernel(os.environ.get("DIFFSYNTH_FLASH_ATTN_KERNEL_REPO_ID"), version=int(os.environ.get("DIFFSYNTH_FLASH_ATTN_KERNEL_VERSION")))
+ customized_flash_attn_func = kernel_module.flash_attn_func
+ CUSTOMIZED_FA_KERNEL_AVAILABLE = True
+ except Exception:
+ CUSTOMIZED_FA_KERNEL_AVAILABLE = False
+else:
+ CUSTOMIZED_FA_KERNEL_AVAILABLE = False
+
+try:
+ from flash_attn.cute import flash_attn_func as flash_attn_func_cute
+ FLASH_ATTN_4_AVAILABLE = True
+except Exception:
+ FLASH_ATTN_4_AVAILABLE = False
+
+try:
+ import flash_attn_interface
+ FLASH_ATTN_3_AVAILABLE = True
+except Exception:
+ FLASH_ATTN_3_AVAILABLE = False
+
+try:
+ import flash_attn
+ FLASH_ATTN_2_AVAILABLE = True
+except Exception:
+ FLASH_ATTN_2_AVAILABLE = False
+
+try:
+ from sageattention import sageattn
+ SAGE_ATTN_AVAILABLE = True
+except Exception:
+ SAGE_ATTN_AVAILABLE = False
+
+try:
+ import xformers.ops as xops
+ XFORMERS_AVAILABLE = True
+except Exception:
+ XFORMERS_AVAILABLE = False
+
+try:
+ from torch.nn.attention.flex_attention import flex_attention as flex_attention_func
+ flex_attention_func = torch.compile(flex_attention_func, dynamic=False, mode="max-autotune-no-cudagraphs", fullgraph=True, backend="inductor")
+ FLEX_ATTN_AVAILABLE = True
+except Exception:
+ FLEX_ATTN_AVAILABLE = False
+
+try:
+ if "enable_gqa" in inspect.signature(torch.nn.functional.scaled_dot_product_attention).parameters:
+ TORCH_SUPPORT_GQA = True
+ else:
+ TORCH_SUPPORT_GQA = False
+except Exception:
+ TORCH_SUPPORT_GQA = False
+
+
+def initialize_attention_priority():
+ if os.environ.get('DIFFSYNTH_ATTENTION_IMPLEMENTATION') is not None:
+ return os.environ.get('DIFFSYNTH_ATTENTION_IMPLEMENTATION').lower()
+ elif CUSTOMIZED_FA_KERNEL_AVAILABLE:
+ return "customized_fa_kernel"
+ elif FLASH_ATTN_4_AVAILABLE:
+ return "flash_attention_4"
+ elif FLASH_ATTN_3_AVAILABLE:
+ return "flash_attention_3"
+ elif FLASH_ATTN_2_AVAILABLE:
+ return "flash_attention_2"
+ elif SAGE_ATTN_AVAILABLE:
+ return "sage_attention"
+ elif XFORMERS_AVAILABLE:
+ return "xformers"
+ else:
+ return "torch"
+
+
+ATTENTION_IMPLEMENTATION = initialize_attention_priority()
+
+
+def rearrange_qkv(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, q_pattern="b n s d", k_pattern="b n s d", v_pattern="b n s d", required_in_pattern="b n s d", dims=None):
+ dims = {} if dims is None else dims
+ if q_pattern != required_in_pattern:
+ q = rearrange(q, f"{q_pattern} -> {required_in_pattern}", **dims)
+ if k_pattern != required_in_pattern:
+ k = rearrange(k, f"{k_pattern} -> {required_in_pattern}", **dims)
+ if v_pattern != required_in_pattern:
+ v = rearrange(v, f"{v_pattern} -> {required_in_pattern}", **dims)
+ return q, k, v
+
+
+def rearrange_out(out: torch.Tensor, out_pattern="b n s d", required_out_pattern="b n s d", dims=None):
+ dims = {} if dims is None else dims
+ if out_pattern != required_out_pattern:
+ out = rearrange(out, f"{required_out_pattern} -> {out_pattern}", **dims)
+ return out
+
+
+def torch_sdpa(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, q_pattern="b n s d", k_pattern="b n s d", v_pattern="b n s d", out_pattern="b n s d", dims=None, attn_mask=None, scale=None, is_causal=False):
+ required_in_pattern, required_out_pattern= "b n s d", "b n s d"
+ q, k, v = rearrange_qkv(q, k, v, q_pattern, k_pattern, v_pattern, required_in_pattern, dims)
+ if q.shape[1] != k.shape[1] or q.shape[1] != v.shape[1]:
+ # Grouped Query Attention
+ if TORCH_SUPPORT_GQA:
+ out = torch.nn.functional.scaled_dot_product_attention(q, k, v, attn_mask, scale=scale, is_causal=is_causal, enable_gqa=True)
+ else:
+ # In low-version torch, `enable_gqa` is not supported.
+ k = repeat(k, "b n s d -> b (n m) s d", m=q.shape[1]//k.shape[1])
+ v = repeat(v, "b n s d -> b (n m) s d", m=q.shape[1]//v.shape[1])
+ out = torch.nn.functional.scaled_dot_product_attention(q, k, v, attn_mask, scale=scale, is_causal=is_causal)
+ else:
+ out = torch.nn.functional.scaled_dot_product_attention(q, k, v, attn_mask, scale=scale, is_causal=is_causal)
+ out = rearrange_out(out, out_pattern, required_out_pattern, dims)
+ return out
+
+
+def torch_sdpa_sliding_window(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, sliding_window: int, q_pattern="b n s d", k_pattern="b n s d", v_pattern="b n s d", out_pattern="b n s d", dims=None, scale=None):
+ required_in_pattern, required_out_pattern = "b n s d", "b n s d"
+ q, k, v = rearrange_qkv(q, k, v, q_pattern, k_pattern, v_pattern, required_in_pattern, dims)
+
+ B, N, S, D = q.shape
+ W = sliding_window
+ chunk_size = W
+ num_chunks = (S + chunk_size - 1) // chunk_size
+
+ output = torch.empty_like(q)
+ dtype = q.dtype
+ device = q.device
+ min_val = torch.finfo(dtype).min
+
+ for i in range(num_chunks):
+ q_start = i * chunk_size
+ q_end = min(q_start + chunk_size, S)
+ actual_chunk_size = q_end - q_start
+
+ kv_start = max(0, q_start - W)
+ kv_end = min(S, q_end + W)
+
+ q_chunk = q[:, :, q_start:q_end, :]
+ k_chunk = k[:, :, kv_start:kv_end, :]
+ v_chunk = v[:, :, kv_start:kv_end, :]
+
+ q_indices = torch.arange(q_start, q_end, device=device)
+ k_indices = torch.arange(kv_start, kv_end, device=device)
+ diff = q_indices.unsqueeze(1) - k_indices.unsqueeze(0)
+ valid = diff.abs() <= W
+
+ local_mask = torch.zeros(actual_chunk_size, kv_end - kv_start, dtype=dtype, device=device)
+ local_mask.masked_fill_(~valid, min_val)
+ local_mask = local_mask.unsqueeze(0).unsqueeze(0)
+
+ out_chunk = torch_sdpa(
+ q_chunk, k_chunk, v_chunk, attn_mask=local_mask, scale=scale
+ )
+ output[:, :, q_start:q_end, :] = out_chunk
+
+ output = rearrange_out(output, out_pattern, required_out_pattern, dims)
+ return output
+
+
+def flash_attention_4(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, q_pattern="b n s d", k_pattern="b n s d", v_pattern="b n s d", out_pattern="b n s d", dims=None, scale=None, is_causal=False, window_size=None):
+ required_in_pattern, required_out_pattern= "b s n d", "b s n d"
+ q, k, v = rearrange_qkv(q, k, v, q_pattern, k_pattern, v_pattern, required_in_pattern, dims)
+ window_size = (window_size, window_size) if window_size is not None else (-1, -1)
+ out = flash_attn_func_cute(q, k, v, softmax_scale=scale, causal=is_causal, window_size=window_size)
+ if isinstance(out, tuple):
+ out = out[0]
+ out = rearrange_out(out, out_pattern, required_out_pattern, dims)
+ return out
+
+
+def flash_attention_3(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, q_pattern="b n s d", k_pattern="b n s d", v_pattern="b n s d", out_pattern="b n s d", dims=None, scale=None, is_causal=False, window_size=None):
+ required_in_pattern, required_out_pattern= "b s n d", "b s n d"
+ q, k, v = rearrange_qkv(q, k, v, q_pattern, k_pattern, v_pattern, required_in_pattern, dims)
+ window_size = (window_size, window_size) if window_size is not None else (-1, -1)
+ out = flash_attn_interface.flash_attn_func(q, k, v, softmax_scale=scale, causal=is_causal, window_size=window_size)
+ if isinstance(out, tuple):
+ out = out[0]
+ out = rearrange_out(out, out_pattern, required_out_pattern, dims)
+ return out
+
+
+def flash_attention_2(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, q_pattern="b n s d", k_pattern="b n s d", v_pattern="b n s d", out_pattern="b n s d", dims=None, scale=None, is_causal=False, window_size=None):
+ required_in_pattern, required_out_pattern= "b s n d", "b s n d"
+ q, k, v = rearrange_qkv(q, k, v, q_pattern, k_pattern, v_pattern, required_in_pattern, dims)
+ window_size = (window_size, window_size) if window_size is not None else (-1, -1)
+ out = flash_attn.flash_attn_func(q, k, v, softmax_scale=scale, causal=is_causal, window_size=window_size)
+ out = rearrange_out(out, out_pattern, required_out_pattern, dims)
+ return out
+
+
+def flash_attention_customized_kernel(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, q_pattern="b n s d", k_pattern="b n s d", v_pattern="b n s d", out_pattern="b n s d", dims=None, scale=None, is_causal=False, window_size=None):
+ required_in_pattern, required_out_pattern = "b s n d", "b s n d"
+ q, k, v = rearrange_qkv(q, k, v, q_pattern, k_pattern, v_pattern, required_in_pattern, dims)
+ window_size = (window_size, window_size) if window_size is not None else (-1, -1)
+ out = customized_flash_attn_func(q, k, v, softmax_scale=scale, causal=is_causal, window_size=window_size)
+ if isinstance(out, tuple):
+ out = out[0]
+ out = rearrange_out(out, out_pattern, required_out_pattern, dims)
+ return out
+
+
+def sage_attention(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, q_pattern="b n s d", k_pattern="b n s d", v_pattern="b n s d", out_pattern="b n s d", dims=None, scale=None):
+ required_in_pattern, required_out_pattern= "b n s d", "b n s d"
+ q, k, v = rearrange_qkv(q, k, v, q_pattern, k_pattern, v_pattern, required_in_pattern, dims)
+ out = sageattn(q, k, v, sm_scale=scale)
+ out = rearrange_out(out, out_pattern, required_out_pattern, dims)
+ return out
+
+
+def xformers_attention(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, q_pattern="b n s d", k_pattern="b n s d", v_pattern="b n s d", out_pattern="b n s d", dims=None, scale=None):
+ required_in_pattern, required_out_pattern= "b s n d", "b s n d"
+ q, k, v = rearrange_qkv(q, k, v, q_pattern, k_pattern, v_pattern, required_in_pattern, dims)
+ out = xops.memory_efficient_attention(q, k, v, scale=scale)
+ out = rearrange_out(out, out_pattern, required_out_pattern, dims)
+ return out
+
+
+def flex_attention(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, q_pattern="b n s d", k_pattern="b n s d", v_pattern="b n s d", out_pattern="b n s d", dims=None, attn_mask=None, scale=None, score_mod=None):
+ assert FLEX_ATTN_AVAILABLE, "Flex Attention is not available. Please upgrade torch to 2.5.0 or later."
+ required_in_pattern, required_out_pattern = "b n s d", "b n s d"
+ q, k, v = rearrange_qkv(q, k, v, q_pattern, k_pattern, v_pattern, required_in_pattern, dims)
+ out = flex_attention_func(query=q, key=k, value=v, block_mask=attn_mask, scale=scale, score_mod=score_mod)
+ out = rearrange_out(out, out_pattern, required_out_pattern, dims)
+ return out
+
+
+def attention_forward(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, q_pattern="b n s d", k_pattern="b n s d", v_pattern="b n s d", out_pattern="b n s d", dims=None, attn_mask=None, scale=None, is_causal=False, compatibility_mode=False, window_size=None, use_flex=False, score_mod=None):
+ if compatibility_mode or (attn_mask is not None) or ATTENTION_IMPLEMENTATION == "torch":
+ if use_flex or score_mod is not None:
+ return flex_attention(q, k, v, q_pattern, k_pattern, v_pattern, out_pattern, dims, attn_mask=attn_mask, scale=scale, score_mod=score_mod)
+ elif window_size is None:
+ return torch_sdpa(q, k, v, q_pattern, k_pattern, v_pattern, out_pattern, dims, attn_mask=attn_mask, scale=scale, is_causal=is_causal)
+ else:
+ # Sliding Window Attention is not compatible with `is_causal` and `attn_mask`.
+ assert is_causal == False and attn_mask is None
+ return torch_sdpa_sliding_window(q, k, v, window_size, q_pattern, k_pattern, v_pattern, out_pattern, dims, scale=scale)
+ elif ATTENTION_IMPLEMENTATION == "customized_fa_kernel":
+ return flash_attention_customized_kernel(q, k, v, q_pattern, k_pattern, v_pattern, out_pattern, dims, scale=scale, is_causal=is_causal, window_size=window_size)
+ elif ATTENTION_IMPLEMENTATION == "flash_attention_4":
+ return flash_attention_4(q, k, v, q_pattern, k_pattern, v_pattern, out_pattern, dims, scale=scale, is_causal=is_causal, window_size=window_size)
+ elif ATTENTION_IMPLEMENTATION == "flash_attention_3":
+ return flash_attention_3(q, k, v, q_pattern, k_pattern, v_pattern, out_pattern, dims, scale=scale, is_causal=is_causal, window_size=window_size)
+ elif ATTENTION_IMPLEMENTATION == "flash_attention_2":
+ return flash_attention_2(q, k, v, q_pattern, k_pattern, v_pattern, out_pattern, dims, scale=scale, is_causal=is_causal, window_size=window_size)
+ elif ATTENTION_IMPLEMENTATION == "sage_attention":
+ if window_size is not None or is_causal: return attention_forward(q, k, v, q_pattern, k_pattern, v_pattern, out_pattern, dims, attn_mask, scale, is_causal, compatibility_mode=True, window_size=window_size)
+ return sage_attention(q, k, v, q_pattern, k_pattern, v_pattern, out_pattern, dims, scale=scale)
+ elif ATTENTION_IMPLEMENTATION == "xformers":
+ if window_size is not None or is_causal: return attention_forward(q, k, v, q_pattern, k_pattern, v_pattern, out_pattern, dims, attn_mask, scale, is_causal, compatibility_mode=True, window_size=window_size)
+ return xformers_attention(q, k, v, q_pattern, k_pattern, v_pattern, out_pattern, dims, scale=scale)
+ else:
+ raise NotImplementedError("No available attention implementation.")
diff --git a/diffsynth/core/data/__init__.py b/diffsynth/core/data/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..d494a277d81eeb2a9575155eb983d8bc3879590a
--- /dev/null
+++ b/diffsynth/core/data/__init__.py
@@ -0,0 +1 @@
+from .unified_dataset import UnifiedDataset
diff --git a/diffsynth/core/data/operators.py b/diffsynth/core/data/operators.py
new file mode 100644
index 0000000000000000000000000000000000000000..218a6d5398dacd8ff1e08cabf65910d61281c9f8
--- /dev/null
+++ b/diffsynth/core/data/operators.py
@@ -0,0 +1,341 @@
+import math, warnings, time
+import torch, torchvision, imageio, os
+import imageio.v3 as iio
+import urllib.request
+from io import BytesIO
+from PIL import Image
+from einops import repeat
+
+
+def _is_url(path):
+ return isinstance(path, str) and path.startswith(("http://", "https://"))
+
+
+def _load_image_from_url(url, timeout=30, max_retries=3):
+ request = urllib.request.Request(url, headers={"User-Agent": "Mozilla/5.0 (compatible; DiffSynth-Studio)"})
+ last_error = None
+ for attempt in range(max_retries):
+ try:
+ with urllib.request.urlopen(request, timeout=timeout) as response:
+ image = Image.open(BytesIO(response.read()))
+ image.load()
+ return image
+ except Exception as e:
+ last_error = e
+ if attempt < max_retries - 1:
+ time.sleep(2 ** attempt)
+ raise RuntimeError(f"Failed to load image from URL '{url}' after {max_retries} attempts: {last_error}") from last_error
+
+
+class DataProcessingPipeline:
+ def __init__(self, operators=None):
+ self.operators: list[DataProcessingOperator] = [] if operators is None else operators
+
+ def __call__(self, data):
+ for operator in self.operators:
+ data = operator(data)
+ return data
+
+ def __rshift__(self, pipe):
+ if isinstance(pipe, DataProcessingOperator):
+ pipe = DataProcessingPipeline([pipe])
+ return DataProcessingPipeline(self.operators + pipe.operators)
+
+
+class DataProcessingOperator:
+ def __call__(self, data):
+ raise NotImplementedError("DataProcessingOperator cannot be called directly.")
+
+ def __rshift__(self, pipe):
+ if isinstance(pipe, DataProcessingOperator):
+ pipe = DataProcessingPipeline([pipe])
+ return DataProcessingPipeline([self]).__rshift__(pipe)
+
+
+class DataProcessingOperatorRaw(DataProcessingOperator):
+ def __call__(self, data):
+ return data
+
+
+class ToInt(DataProcessingOperator):
+ def __call__(self, data):
+ return int(data)
+
+
+class ToFloat(DataProcessingOperator):
+ def __call__(self, data):
+ return float(data)
+
+
+class ToStr(DataProcessingOperator):
+ def __init__(self, none_value=""):
+ self.none_value = none_value
+
+ def __call__(self, data):
+ if data is None: data = self.none_value
+ return str(data)
+
+
+class LoadImage(DataProcessingOperator):
+ def __init__(self, convert_RGB=True, convert_RGBA=False, url_timeout=30, url_max_retries=3):
+ self.convert_RGB = convert_RGB
+ self.convert_RGBA = convert_RGBA
+ self.url_timeout = url_timeout
+ self.url_max_retries = url_max_retries
+
+ def __call__(self, data: str):
+ if _is_url(data):
+ image = _load_image_from_url(data, timeout=self.url_timeout, max_retries=self.url_max_retries)
+ else:
+ image = Image.open(data)
+ if self.convert_RGB: image = image.convert("RGB")
+ if self.convert_RGBA: image = image.convert("RGBA")
+ return image
+
+
+class ImageCropAndResize(DataProcessingOperator):
+ def __init__(self, height=None, width=None, max_pixels=None, height_division_factor=1, width_division_factor=1):
+ self.height = height
+ self.width = width
+ self.max_pixels = max_pixels
+ self.height_division_factor = height_division_factor
+ self.width_division_factor = width_division_factor
+
+ def crop_and_resize(self, image, target_height, target_width):
+ width, height = image.size
+ scale = max(target_width / width, target_height / height)
+ image = torchvision.transforms.functional.resize(
+ image,
+ (round(height*scale), round(width*scale)),
+ interpolation=torchvision.transforms.InterpolationMode.BILINEAR
+ )
+ image = torchvision.transforms.functional.center_crop(image, (target_height, target_width))
+ return image
+
+ def get_height_width(self, image):
+ if self.height is None or self.width is None:
+ width, height = image.size
+ if width * height > self.max_pixels:
+ scale = (width * height / self.max_pixels) ** 0.5
+ height, width = int(height / scale), int(width / scale)
+ height = height // self.height_division_factor * self.height_division_factor
+ width = width // self.width_division_factor * self.width_division_factor
+ else:
+ height, width = self.height, self.width
+ return height, width
+
+ def __call__(self, data: Image.Image):
+ image = self.crop_and_resize(data, *self.get_height_width(data))
+ return image
+
+
+class ToList(DataProcessingOperator):
+ def __call__(self, data):
+ return [data]
+
+
+class FrameSamplerByRateMixin:
+ def __init__(self, num_frames=81, time_division_factor=4, time_division_remainder=1, frame_rate=24, fix_frame_rate=False):
+ self.num_frames = num_frames
+ self.time_division_factor = time_division_factor
+ self.time_division_remainder = time_division_remainder
+ self.frame_rate = frame_rate
+ self.fix_frame_rate = fix_frame_rate
+
+ def get_reader(self, data: str):
+ return imageio.get_reader(data)
+
+ def get_available_num_frames(self, reader):
+ if not self.fix_frame_rate:
+ return reader.count_frames()
+ meta_data = reader.get_meta_data()
+ total_original_frames = int(reader.count_frames())
+ duration = meta_data["duration"] if "duration" in meta_data else total_original_frames / meta_data['fps']
+ total_available_frames = math.floor(duration * self.frame_rate)
+ return int(total_available_frames)
+
+ def get_num_frames(self, reader):
+ num_frames = self.num_frames
+ total_frames = self.get_available_num_frames(reader)
+ if int(total_frames) < num_frames:
+ num_frames = total_frames
+ while num_frames > 1 and num_frames % self.time_division_factor != self.time_division_remainder:
+ num_frames -= 1
+ return num_frames
+
+ def map_single_frame_id(self, new_sequence_id: int, raw_frame_rate: float, total_raw_frames: int) -> int:
+ if not self.fix_frame_rate:
+ return new_sequence_id
+ target_time_in_seconds = new_sequence_id / self.frame_rate
+ raw_frame_index_float = target_time_in_seconds * raw_frame_rate
+ frame_id = int(round(raw_frame_index_float))
+ frame_id = min(frame_id, total_raw_frames - 1)
+ return frame_id
+
+
+class LoadVideo(DataProcessingOperator, FrameSamplerByRateMixin):
+ def __init__(self, num_frames=81, time_division_factor=4, time_division_remainder=1, frame_processor=lambda x: x, frame_rate=24, fix_frame_rate=False):
+ FrameSamplerByRateMixin.__init__(self, num_frames, time_division_factor, time_division_remainder, frame_rate, fix_frame_rate)
+ # frame_processor is build in the video loader for high efficiency.
+ self.frame_processor = frame_processor
+
+ def __call__(self, data: str):
+ reader = self.get_reader(data)
+ raw_frame_rate = reader.get_meta_data()['fps']
+ num_frames = self.get_num_frames(reader)
+ total_raw_frames = reader.count_frames()
+ frames = []
+ for frame_id in range(num_frames):
+ frame_id = self.map_single_frame_id(frame_id, raw_frame_rate, total_raw_frames)
+ frame = reader.get_data(frame_id)
+ frame = Image.fromarray(frame)
+ frame = self.frame_processor(frame)
+ frames.append(frame)
+ reader.close()
+ return frames
+
+
+class SequencialProcess(DataProcessingOperator):
+ def __init__(self, operator=lambda x: x):
+ self.operator = operator
+
+ def __call__(self, data):
+ return [self.operator(i) for i in data]
+
+
+class LoadGIF(DataProcessingOperator):
+ def __init__(self, num_frames=81, time_division_factor=4, time_division_remainder=1, frame_processor=lambda x: x):
+ self.num_frames = num_frames
+ self.time_division_factor = time_division_factor
+ self.time_division_remainder = time_division_remainder
+ # frame_processor is build in the video loader for high efficiency.
+ self.frame_processor = frame_processor
+
+ def get_num_frames(self, path):
+ num_frames = self.num_frames
+ images = iio.imread(path, mode="RGB")
+ if len(images) < num_frames:
+ num_frames = len(images)
+ while num_frames > 1 and num_frames % self.time_division_factor != self.time_division_remainder:
+ num_frames -= 1
+ return num_frames
+
+ def __call__(self, data: str):
+ num_frames = self.get_num_frames(data)
+ frames = []
+ images = iio.imread(data, mode="RGB")
+ for img in images:
+ frame = Image.fromarray(img)
+ frame = self.frame_processor(frame)
+ frames.append(frame)
+ if len(frames) >= num_frames:
+ break
+ return frames
+
+
+class RouteByExtensionName(DataProcessingOperator):
+ def __init__(self, operator_map):
+ self.operator_map = operator_map
+
+ def __call__(self, data: str):
+ file_ext_name = data.split(".")[-1].lower()
+ for ext_names, operator in self.operator_map:
+ if ext_names is None or file_ext_name in ext_names:
+ return operator(data)
+ raise ValueError(f"Unsupported file: {data}")
+
+
+class RouteByType(DataProcessingOperator):
+ def __init__(self, operator_map):
+ self.operator_map = operator_map
+
+ def __call__(self, data):
+ for dtype, operator in self.operator_map:
+ if dtype is None or isinstance(data, dtype):
+ return operator(data)
+ raise ValueError(f"Unsupported data: {data}")
+
+
+class LoadTorchPickle(DataProcessingOperator):
+ def __init__(self, map_location="cpu"):
+ self.map_location = map_location
+
+ def __call__(self, data):
+ return torch.load(data, map_location=self.map_location, weights_only=False)
+
+
+class ToAbsolutePath(DataProcessingOperator):
+ def __init__(self, base_path=""):
+ self.base_path = base_path
+
+ def __call__(self, data):
+ # A remote URL is already absolute; joining it with base_path would corrupt it.
+ if _is_url(data):
+ return data
+ return os.path.join(self.base_path, data)
+
+
+class LoadAudio(DataProcessingOperator):
+ def __init__(self, sr=16000):
+ self.sr = sr
+ import librosa
+ self.audio_loader = librosa.load
+
+ def __call__(self, data: str):
+ input_audio, sample_rate = self.audio_loader(data, sr=self.sr)
+ return input_audio
+
+
+class LoadAudioWithTorchaudio(DataProcessingOperator, FrameSamplerByRateMixin):
+
+ def __init__(self, num_frames=121, time_division_factor=8, time_division_remainder=1, frame_rate=24, fix_frame_rate=True):
+ FrameSamplerByRateMixin.__init__(self, num_frames, time_division_factor, time_division_remainder, frame_rate, fix_frame_rate)
+
+ def __call__(self, data: str):
+ try:
+ reader = self.get_reader(data)
+ num_frames = self.get_num_frames(reader)
+ duration = num_frames / self.frame_rate
+ reader.close()
+ from ...utils.data.audio import read_audio
+ waveform, sample_rate = read_audio(data, duration=duration)
+ target_samples = int(duration * sample_rate)
+ current_samples = waveform.shape[-1]
+ if current_samples > target_samples:
+ waveform = waveform[..., :target_samples]
+ elif current_samples < target_samples:
+ padding = target_samples - current_samples
+ waveform = torch.nn.functional.pad(waveform, (0, padding))
+ return waveform, sample_rate
+ except Exception as e:
+ warnings.warn(f"Cannot load audio in {data} due to {e}. The audio will be `None`.")
+ return None
+
+
+class LoadPureAudioWithTorchaudio(DataProcessingOperator):
+
+ def __init__(self, target_sample_rate=None, max_audio_duration=None, padding=False, channels=2):
+ self.target_sample_rate = target_sample_rate
+ self.max_audio_duration = max_audio_duration
+ self.resample = True if target_sample_rate is not None else False
+ self.padding = padding
+ from diffsynth.utils.data.audio import read_audio
+ self.audio_loader = read_audio
+
+ def __call__(self, data: str):
+ try:
+ waveform, sample_rate = self.audio_loader(data, resample=self.resample, resample_rate=self.target_sample_rate)
+ if self.max_audio_duration is not None:
+ target_samples = int(self.max_audio_duration * sample_rate)
+ current_samples = waveform.shape[-1]
+ if current_samples > target_samples:
+ waveform = waveform[..., :target_samples]
+ elif current_samples < target_samples and self.padding:
+ padding = target_samples - current_samples
+ waveform = torch.nn.functional.pad(waveform, (0, padding))
+ if waveform.shape[0] == 1:
+ waveform = repeat(waveform, "C L -> (N C) L", N=2)
+ return waveform, sample_rate
+ except Exception as e:
+ print(f"Cannot load audio in {data} due to {e}. The audio will be `None`.")
+ return None
diff --git a/diffsynth/core/data/unified_dataset.py b/diffsynth/core/data/unified_dataset.py
new file mode 100644
index 0000000000000000000000000000000000000000..71ddef99fbcac074f5025358678452504aa6ecb0
--- /dev/null
+++ b/diffsynth/core/data/unified_dataset.py
@@ -0,0 +1,140 @@
+from .operators import *
+import torch, json, pandas
+
+
+class UnifiedDataset(torch.utils.data.Dataset):
+ def __init__(
+ self,
+ base_path=None, metadata_path=None,
+ repeat=1,
+ data_file_keys=tuple(),
+ main_data_operator=lambda x: x,
+ special_operator_map=None,
+ max_data_items=None,
+ cached_data_file_list=None,
+ ):
+ self.base_path = base_path
+ self.metadata_path = metadata_path
+ self.repeat = repeat
+ self.data_file_keys = data_file_keys
+ self.main_data_operator = main_data_operator
+ self.cached_data_operator = LoadTorchPickle()
+ self.special_operator_map = {} if special_operator_map is None else special_operator_map
+ self.max_data_items = max_data_items
+ self.data = []
+ self.cached_data = []
+ self.load_from_cache = metadata_path is None
+ self.load_metadata(metadata_path, cached_data_file_list=cached_data_file_list)
+
+ @staticmethod
+ def default_image_operator(
+ base_path="",
+ max_pixels=1920*1080, height=None, width=None,
+ height_division_factor=16, width_division_factor=16,
+ ):
+ return RouteByType(operator_map=[
+ (str, ToAbsolutePath(base_path) >> LoadImage() >> ImageCropAndResize(height, width, max_pixels, height_division_factor, width_division_factor)),
+ (list, SequencialProcess(ToAbsolutePath(base_path) >> LoadImage() >> ImageCropAndResize(height, width, max_pixels, height_division_factor, width_division_factor))),
+ ])
+
+ @staticmethod
+ def default_video_operator(
+ base_path="",
+ max_pixels=1920*1080, height=None, width=None,
+ height_division_factor=16, width_division_factor=16,
+ num_frames=81, time_division_factor=4, time_division_remainder=1,
+ frame_rate=24, fix_frame_rate=False,
+ ):
+ return RouteByType(operator_map=[
+ (str, ToAbsolutePath(base_path) >> RouteByExtensionName(operator_map=[
+ (("jpg", "jpeg", "png", "webp", "bmp"), LoadImage() >> ImageCropAndResize(height, width, max_pixels, height_division_factor, width_division_factor) >> ToList()),
+ (("gif",), LoadGIF(
+ num_frames, time_division_factor, time_division_remainder,
+ frame_processor=ImageCropAndResize(height, width, max_pixels, height_division_factor, width_division_factor),
+ )),
+ (("mp4", "avi", "mov", "wmv", "mkv", "flv", "webm"), LoadVideo(
+ num_frames, time_division_factor, time_division_remainder,
+ frame_processor=ImageCropAndResize(height, width, max_pixels, height_division_factor, width_division_factor),
+ frame_rate=frame_rate, fix_frame_rate=fix_frame_rate,
+ )),
+ ])),
+ ])
+
+ def search_for_cached_data_files(self, path):
+ for file_name in os.listdir(path):
+ subpath = os.path.join(path, file_name)
+ if os.path.isdir(subpath):
+ self.search_for_cached_data_files(subpath)
+ elif subpath.endswith(".pth"):
+ self.cached_data.append(subpath)
+
+ def load_cached_data_file_list(self, file_list_path):
+ with open(file_list_path, "r", encoding="utf-8") as file_list:
+ for line_number, line in enumerate(file_list, 1):
+ path = line.strip()
+ if not path:
+ continue
+ if not os.path.isabs(path):
+ path = os.path.join(self.base_path, path)
+ path = os.path.abspath(path)
+ if not path.endswith(".pth"):
+ raise ValueError(f"{file_list_path}:{line_number}: cached file must end in .pth: {path}")
+ if not os.path.isfile(path):
+ raise FileNotFoundError(f"{file_list_path}:{line_number}: cached file does not exist: {path}")
+ self.cached_data.append(path)
+
+ def load_metadata(self, metadata_path, cached_data_file_list=None):
+ if cached_data_file_list is not None:
+ if metadata_path is not None:
+ raise ValueError("cached_data_file_list and metadata_path are mutually exclusive")
+ print(f"Loading cached data paths from `{cached_data_file_list}`.")
+ self.load_cached_data_file_list(cached_data_file_list)
+ print(f"{len(self.cached_data)} listed cached data files found.")
+ elif metadata_path is None:
+ print("No metadata_path. Searching for cached data files.")
+ self.search_for_cached_data_files(self.base_path)
+ print(f"{len(self.cached_data)} cached data files found.")
+ elif metadata_path.endswith(".json"):
+ with open(metadata_path, "r") as f:
+ metadata = json.load(f)
+ self.data = metadata
+ elif metadata_path.endswith(".jsonl"):
+ metadata = []
+ with open(metadata_path, 'r') as f:
+ for line in f:
+ metadata.append(json.loads(line.strip()))
+ self.data = metadata
+ else:
+ metadata = pandas.read_csv(metadata_path)
+ self.data = [metadata.iloc[i].to_dict() for i in range(len(metadata))]
+
+ def __getitem__(self, data_id):
+ if self.load_from_cache:
+ data = self.cached_data[data_id % len(self.cached_data)]
+ data = self.cached_data_operator(data)
+ else:
+ data = self.data[data_id % len(self.data)].copy()
+ for key in self.data_file_keys:
+ if key in data:
+ if key in self.special_operator_map:
+ data[key] = self.special_operator_map[key](data[key])
+ elif key in self.data_file_keys:
+ data[key] = self.main_data_operator(data[key])
+ return data
+
+ def __len__(self):
+ if self.max_data_items is not None:
+ return self.max_data_items
+ elif self.load_from_cache:
+ return len(self.cached_data) * self.repeat
+ else:
+ return len(self.data) * self.repeat
+
+ def check_data_equal(self, data1, data2):
+ # Debug only
+ if len(data1) != len(data2):
+ return False
+ for k in data1:
+ if data1[k] != data2[k]:
+ return False
+ return True
diff --git a/diffsynth/core/device/__init__.py b/diffsynth/core/device/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..889d6823adb203630d2173e8ab25fcc350b32ce4
--- /dev/null
+++ b/diffsynth/core/device/__init__.py
@@ -0,0 +1,2 @@
+from .npu_compatible_device import parse_device_type, parse_nccl_backend, get_available_device_type, get_device_name
+from .npu_compatible_device import IS_NPU_AVAILABLE, IS_CUDA_AVAILABLE
diff --git a/diffsynth/core/device/npu_compatible_device.py b/diffsynth/core/device/npu_compatible_device.py
new file mode 100644
index 0000000000000000000000000000000000000000..b672f9c9116552519581b3cc066be4f88955dc96
--- /dev/null
+++ b/diffsynth/core/device/npu_compatible_device.py
@@ -0,0 +1,109 @@
+import importlib
+import torch
+from typing import Any
+
+
+def is_torch_npu_available():
+ return importlib.util.find_spec("torch_npu") is not None
+
+
+IS_CUDA_AVAILABLE = torch.cuda.is_available()
+IS_NPU_AVAILABLE = is_torch_npu_available() and torch.npu.is_available()
+
+if IS_NPU_AVAILABLE:
+ import torch_npu
+
+ torch.npu.config.allow_internal_format = False
+
+
+def get_device_type() -> str:
+ """Get device type based on current machine, currently only support CPU, CUDA, NPU."""
+ if IS_CUDA_AVAILABLE:
+ device = "cuda"
+ elif IS_NPU_AVAILABLE:
+ device = "npu"
+ else:
+ device = "cpu"
+
+ return device
+
+
+def get_torch_device() -> Any:
+ """Get torch attribute based on device type, e.g. torch.cuda or torch.npu"""
+ device_name = get_device_type()
+
+ try:
+ return getattr(torch, device_name)
+ except AttributeError:
+ print(f"Device namespace '{device_name}' not found in torch, try to load 'torch.cuda'.")
+ return torch.cuda
+
+
+def get_device_id() -> int:
+ """Get current device id based on device type."""
+ return get_torch_device().current_device()
+
+
+def get_device_name() -> str:
+ """Get current device name based on device type."""
+ return f"{get_device_type()}:{get_device_id()}"
+
+
+def synchronize() -> None:
+ """Execute torch synchronize operation."""
+ get_torch_device().synchronize()
+
+
+def empty_cache() -> None:
+ """Execute torch empty cache operation."""
+ get_torch_device().empty_cache()
+
+
+def get_nccl_backend() -> str:
+ """Return distributed communication backend type based on device type."""
+ if IS_CUDA_AVAILABLE:
+ return "nccl"
+ elif IS_NPU_AVAILABLE:
+ return "hccl"
+ else:
+ raise RuntimeError(f"No available distributed communication backend found on device type {get_device_type()}.")
+
+
+def enable_high_precision_for_bf16():
+ """
+ Set high accumulation dtype for matmul and reduction.
+ """
+ if IS_CUDA_AVAILABLE:
+ torch.backends.cuda.matmul.allow_tf32 = False
+ torch.backends.cuda.matmul.allow_bf16_reduced_precision_reduction = False
+
+ if IS_NPU_AVAILABLE:
+ torch.npu.matmul.allow_tf32 = False
+ torch.npu.matmul.allow_bf16_reduced_precision_reduction = False
+
+
+def parse_device_type(device):
+ if isinstance(device, torch.device):
+ return device.type
+ elif isinstance(device, str):
+ # `torch.device("npu")` raises when torch_npu is not installed, so keep npu ahead of the generic parse.
+ if device.startswith("npu"):
+ return "npu"
+ try:
+ return torch.device(device).type
+ except RuntimeError:
+ # Strings that are not device names at all
+ return "cpu"
+
+
+def parse_nccl_backend(device_type):
+ if device_type == "cuda":
+ return "nccl"
+ elif device_type == "npu":
+ return "hccl"
+ else:
+ raise RuntimeError(f"No available distributed communication backend found on device type {device_type}.")
+
+
+def get_available_device_type():
+ return get_device_type()
diff --git a/diffsynth/core/gradient/__init__.py b/diffsynth/core/gradient/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..57914792a78ec32f69c3c99ae37535598efc8d52
--- /dev/null
+++ b/diffsynth/core/gradient/__init__.py
@@ -0,0 +1 @@
+from .gradient_checkpoint import gradient_checkpoint_forward
diff --git a/diffsynth/core/gradient/gradient_checkpoint.py b/diffsynth/core/gradient/gradient_checkpoint.py
new file mode 100644
index 0000000000000000000000000000000000000000..495827f809d3deab94775de80e2e574fbc3dfaa1
--- /dev/null
+++ b/diffsynth/core/gradient/gradient_checkpoint.py
@@ -0,0 +1,65 @@
+import torch
+
+
+try:
+ import deepspeed
+ _HAS_DEEPSPEED = True
+except ModuleNotFoundError:
+ _HAS_DEEPSPEED = False
+
+
+def create_custom_forward(module):
+ def custom_forward(*inputs, **kwargs):
+ return module(*inputs, **kwargs)
+ return custom_forward
+
+
+def create_custom_forward_use_reentrant(module):
+ def custom_forward(*inputs):
+ return module(*inputs)
+ return custom_forward
+
+
+def judge_args_requires_grad(*args):
+ for arg in args:
+ if isinstance(arg, torch.Tensor) and arg.requires_grad:
+ return True
+ return False
+
+
+def gradient_checkpoint_forward(
+ model,
+ use_gradient_checkpointing,
+ use_gradient_checkpointing_offload,
+ *args,
+ **kwargs,
+):
+ if use_gradient_checkpointing and _HAS_DEEPSPEED and deepspeed.checkpointing.is_configured():
+ all_args = args + tuple(kwargs.values())
+ if not judge_args_requires_grad(*all_args):
+ # get the first grad_enabled tensor from un_checkpointed forward
+ model_output = model(*args, **kwargs)
+ else:
+ model_output = deepspeed.checkpointing.checkpoint(
+ create_custom_forward_use_reentrant(model),
+ *all_args,
+ )
+ return model_output
+ if use_gradient_checkpointing_offload:
+ with torch.autograd.graph.save_on_cpu():
+ model_output = torch.utils.checkpoint.checkpoint(
+ create_custom_forward(model),
+ *args,
+ **kwargs,
+ use_reentrant=False,
+ )
+ elif use_gradient_checkpointing:
+ model_output = torch.utils.checkpoint.checkpoint(
+ create_custom_forward(model),
+ *args,
+ **kwargs,
+ use_reentrant=False,
+ )
+ else:
+ model_output = model(*args, **kwargs)
+ return model_output
diff --git a/diffsynth/core/loader/__init__.py b/diffsynth/core/loader/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..1f56d814bae40436f66bca583e33d180d6e11247
--- /dev/null
+++ b/diffsynth/core/loader/__init__.py
@@ -0,0 +1,3 @@
+from .file import load_state_dict, hash_state_dict_keys, hash_model_file
+from .model import load_model, load_model_with_disk_offload
+from .config import ModelConfig
diff --git a/diffsynth/core/loader/config.py b/diffsynth/core/loader/config.py
new file mode 100644
index 0000000000000000000000000000000000000000..aecb00f2af9240147b3b36f7b29e364c049daf26
--- /dev/null
+++ b/diffsynth/core/loader/config.py
@@ -0,0 +1,177 @@
+import torch, glob, os, sys
+from typing import Optional, Union, Dict
+from dataclasses import dataclass
+from modelscope import snapshot_download
+from huggingface_hub import snapshot_download as hf_snapshot_download
+from typing import Optional
+from ..quant import QuantizeConfig, MixedQuantizeConfig
+
+_download_tips_printed = False
+
+if sys.platform.startswith('win'):
+ DOWNLOADING_TIPS = """
+┌──────────────────────────────────────────────────────────────────────┐
+│ DiffSynth-Studio Model Downloader Configuration: │
+│ │
+│ [0] Download from https://modelscope.cn/ │
+│ (default behavior) │
+│ [1] Download from https://modelscope.ai/ │
+│ (enabled via `$env:MODELSCOPE_ENDPOINT="https://modelscope.ai"`) │
+│ [2] Download from https://huggingface.co/ │
+│ (enabled via `$env:DIFFSYNTH_DOWNLOAD_SOURCE="HuggingFace"`) │
+│ [3] Skip download and load only pre-downloaded model files │
+│ (enabled via `$env:DIFFSYNTH_SKIP_DOWNLOAD="True"`) │
+└──────────────────────────────────────────────────────────────────────┘
+""".strip()
+else:
+ DOWNLOADING_TIPS = """
+┌──────────────────────────────────────────────────────────────────────┐
+│ DiffSynth-Studio Model Downloader Configuration: │
+│ │
+│ [0] Download from https://modelscope.cn/ │
+│ (default behavior) │
+│ [1] Download from https://modelscope.ai/ │
+│ (enabled via `export MODELSCOPE_ENDPOINT=https://modelscope.ai`) │
+│ [2] Download from https://huggingface.co/ │
+│ (enabled via `export DIFFSYNTH_DOWNLOAD_SOURCE=HuggingFace`) │
+│ [3] Skip download and load only pre-downloaded model files │
+│ (enabled via `export DIFFSYNTH_SKIP_DOWNLOAD=True`) │
+└──────────────────────────────────────────────────────────────────────┘
+""".strip()
+
+
+@dataclass
+class ModelConfig:
+ path: Union[str, list[str]] = None
+ model_id: str = None
+ origin_file_pattern: Union[str, list[str]] = None
+ download_source: str = None
+ local_model_path: str = None
+ skip_download: bool = None
+ offload_device: Optional[Union[str, torch.device]] = None
+ offload_dtype: Optional[torch.dtype] = None
+ onload_device: Optional[Union[str, torch.device]] = None
+ onload_dtype: Optional[torch.dtype] = None
+ preparing_device: Optional[Union[str, torch.device]] = None
+ preparing_dtype: Optional[torch.dtype] = None
+ computation_device: Optional[Union[str, torch.device]] = None
+ computation_dtype: Optional[torch.dtype] = None
+ clear_parameters: bool = False
+ state_dict: Dict[str, torch.Tensor] = None
+ quantize: Optional[Union[QuantizeConfig, MixedQuantizeConfig]] = None
+
+ def check_input(self):
+ if self.path is None and self.model_id is None:
+ raise ValueError(f"""No valid model files. Please use `ModelConfig(path="xxx")` or `ModelConfig(model_id="xxx/yyy", origin_file_pattern="zzz")`. `skip_download=True` only supports the first one.""")
+
+ def parse_original_file_pattern(self):
+ if self.origin_file_pattern in [None, "", "./"]:
+ return "*"
+ elif self.origin_file_pattern.endswith("/"):
+ return self.origin_file_pattern + "*"
+ else:
+ return self.origin_file_pattern
+
+ def parse_download_source(self):
+ if self.download_source is None:
+ if os.environ.get('DIFFSYNTH_DOWNLOAD_SOURCE') is not None:
+ return os.environ.get('DIFFSYNTH_DOWNLOAD_SOURCE').lower()
+ else:
+ return "modelscope"
+ else:
+ return self.download_source.lower()
+
+ def parse_skip_download(self):
+ if self.skip_download is None:
+ if os.environ.get('DIFFSYNTH_SKIP_DOWNLOAD') is not None:
+ if os.environ.get('DIFFSYNTH_SKIP_DOWNLOAD').lower() == "true":
+ return True
+ elif os.environ.get('DIFFSYNTH_SKIP_DOWNLOAD').lower() == "false":
+ return False
+ else:
+ return False
+ else:
+ return self.skip_download
+
+ def download(self):
+ origin_file_pattern = self.parse_original_file_pattern()
+ downloaded_files = glob.glob(origin_file_pattern, root_dir=os.path.join(self.local_model_path, self.model_id))
+ download_source = self.parse_download_source()
+ if download_source.lower() == "modelscope":
+ snapshot_download(
+ self.model_id,
+ local_dir=os.path.join(self.local_model_path, self.model_id),
+ allow_file_pattern=origin_file_pattern,
+ ignore_file_pattern=downloaded_files,
+ local_files_only=False
+ )
+ elif download_source.lower() == "huggingface":
+ hf_snapshot_download(
+ self.model_id,
+ local_dir=os.path.join(self.local_model_path, self.model_id),
+ allow_patterns=origin_file_pattern,
+ ignore_patterns=downloaded_files,
+ local_files_only=False
+ )
+ else:
+ raise ValueError("`download_source` should be `modelscope` or `huggingface`.")
+
+ def require_downloading(self):
+ if self.path is not None:
+ return False
+ skip_download = self.parse_skip_download()
+ return not skip_download
+
+ def reset_local_model_path(self):
+ if os.environ.get('DIFFSYNTH_MODEL_BASE_PATH') is not None:
+ self.local_model_path = os.environ.get('DIFFSYNTH_MODEL_BASE_PATH')
+ elif self.local_model_path is None:
+ self.local_model_path = "./models"
+
+ def check_download_source(self):
+ download_source = self.parse_download_source().lower()
+ if os.environ.get('DIFFSYNTH_SKIP_DOWNLOAD', "").lower() == "true":
+ behavior = 3
+ elif download_source == "modelscope":
+ if "modelscope.ai" in os.environ.get("MODELSCOPE_ENDPOINT", "") or "modelscope.ai" in os.environ.get("MODELSCOPE_DOMAIN", ""):
+ behavior = 1
+ else:
+ behavior = 0
+ else:
+ behavior = 2
+ tips = DOWNLOADING_TIPS
+ for i in range(4): tips = tips.replace(f"[{i}]", ["[ ]", "[√]"][i==behavior])
+ global _download_tips_printed
+ if not _download_tips_printed:
+ is_main_process = True
+ if torch.distributed.is_initialized():
+ is_main_process = torch.distributed.get_rank() == 0
+ if is_main_process:
+ print(tips)
+ _download_tips_printed = True
+
+ def download_if_necessary(self):
+ self.check_download_source()
+ self.check_input()
+ self.reset_local_model_path()
+ if self.require_downloading():
+ self.download()
+ if self.path is None:
+ if self.origin_file_pattern in [None, "", "./"]:
+ self.path = os.path.join(self.local_model_path, self.model_id)
+ else:
+ self.path = glob.glob(os.path.join(self.local_model_path, self.model_id, self.origin_file_pattern))
+ if isinstance(self.path, list) and len(self.path) == 1:
+ self.path = self.path[0]
+
+ def vram_config(self):
+ return {
+ "offload_device": self.offload_device,
+ "offload_dtype": self.offload_dtype,
+ "onload_device": self.onload_device,
+ "onload_dtype": self.onload_dtype,
+ "preparing_device": self.preparing_device,
+ "preparing_dtype": self.preparing_dtype,
+ "computation_device": self.computation_device,
+ "computation_dtype": self.computation_dtype,
+ }
diff --git a/diffsynth/core/loader/file.py b/diffsynth/core/loader/file.py
new file mode 100644
index 0000000000000000000000000000000000000000..b922f6d1a60f29a722801e6ddf361d76a7a2ecb8
--- /dev/null
+++ b/diffsynth/core/loader/file.py
@@ -0,0 +1,151 @@
+from safetensors import safe_open
+import torch, hashlib, json
+
+
+def load_state_dict(file_path, torch_dtype=None, device="cpu", pin_memory=False, verbose=0):
+ if isinstance(file_path, list):
+ state_dict = {}
+ for file_path_ in file_path:
+ state_dict.update(load_state_dict(file_path_, torch_dtype, device, pin_memory=pin_memory, verbose=verbose))
+ else:
+ if verbose >= 1:
+ print(f"Loading file [started]: {file_path}")
+ if file_path.endswith(".safetensors"):
+ state_dict = load_state_dict_from_safetensors(file_path, torch_dtype=torch_dtype, device=device)
+ else:
+ state_dict = load_state_dict_from_bin(file_path, torch_dtype=torch_dtype, device=device)
+ # If load state dict in CPU memory, `pin_memory=True` will make `model.to("cuda")` faster.
+ if pin_memory:
+ for i in state_dict:
+ state_dict[i] = state_dict[i].pin_memory()
+ if verbose >= 1:
+ print(f"Loading file [done]: {file_path}")
+ return state_dict
+
+
+def load_state_dict_from_safetensors(file_path, torch_dtype=None, device="cpu"):
+ state_dict = {}
+ with safe_open(file_path, framework="pt", device=str(device)) as f:
+ for k in f.keys():
+ state_dict[k] = f.get_tensor(k)
+ if torch_dtype is not None:
+ state_dict[k] = state_dict[k].to(torch_dtype)
+ return state_dict
+
+
+def load_metadata_from_safetensors(file_path):
+ if isinstance(file_path, list):
+ metadata = {}
+ for file_path_ in file_path:
+ for key, value in load_metadata_from_safetensors(file_path_).items():
+ metadata[key] = merge_metadata_value(metadata[key], value) if key in metadata else value
+ return metadata
+ with safe_open(file_path, framework="pt", device="cpu") as f:
+ return f.metadata() or {}
+
+
+def merge_metadata_value(old_value, new_value):
+ try:
+ old_items, new_items = json.loads(old_value), json.loads(new_value)
+ except (json.JSONDecodeError, TypeError):
+ return new_value
+ if isinstance(old_items, list) and isinstance(new_items, list):
+ return json.dumps(old_items + [i for i in new_items if i not in old_items])
+ return new_value
+
+
+def load_state_dict_from_bin(file_path, torch_dtype=None, device="cpu"):
+ state_dict = torch.load(file_path, map_location=device, weights_only=True)
+ if len(state_dict) == 1:
+ if "state_dict" in state_dict:
+ state_dict = state_dict["state_dict"]
+ elif "module" in state_dict:
+ state_dict = state_dict["module"]
+ elif "model_state" in state_dict:
+ state_dict = state_dict["model_state"]
+ if torch_dtype is not None:
+ for i in state_dict:
+ if isinstance(state_dict[i], torch.Tensor):
+ state_dict[i] = state_dict[i].to(torch_dtype)
+ return state_dict
+
+
+def convert_state_dict_keys_to_single_str(state_dict, with_shape=True):
+ keys = []
+ for key, value in state_dict.items():
+ if isinstance(key, str):
+ if isinstance(value, torch.Tensor):
+ if with_shape:
+ shape = "_".join(map(str, list(value.shape)))
+ keys.append(key + ":" + shape)
+ keys.append(key)
+ elif isinstance(value, dict):
+ keys.append(key + "|" + convert_state_dict_keys_to_single_str(value, with_shape=with_shape))
+ keys.sort()
+ keys_str = ",".join(keys)
+ return keys_str
+
+
+def hash_state_dict_keys(state_dict, with_shape=True):
+ keys_str = convert_state_dict_keys_to_single_str(state_dict, with_shape=with_shape)
+ keys_str = keys_str.encode(encoding="UTF-8")
+ return hashlib.md5(keys_str).hexdigest()
+
+
+def load_keys_dict(file_path):
+ if isinstance(file_path, list):
+ state_dict = {}
+ for file_path_ in file_path:
+ state_dict.update(load_keys_dict(file_path_))
+ return state_dict
+ if file_path.endswith(".safetensors"):
+ return load_keys_dict_from_safetensors(file_path)
+ else:
+ return load_keys_dict_from_bin(file_path)
+
+
+def load_keys_dict_from_safetensors(file_path):
+ keys_dict = {}
+ with safe_open(file_path, framework="pt", device="cpu") as f:
+ for k in f.keys():
+ keys_dict[k] = f.get_slice(k).get_shape()
+ return keys_dict
+
+
+def convert_state_dict_to_keys_dict(state_dict):
+ keys_dict = {}
+ for k, v in state_dict.items():
+ if isinstance(v, torch.Tensor):
+ keys_dict[k] = list(v.shape)
+ else:
+ keys_dict[k] = convert_state_dict_to_keys_dict(v)
+ return keys_dict
+
+
+def load_keys_dict_from_bin(file_path):
+ state_dict = load_state_dict_from_bin(file_path)
+ keys_dict = convert_state_dict_to_keys_dict(state_dict)
+ return keys_dict
+
+
+def convert_keys_dict_to_single_str(state_dict, with_shape=True):
+ keys = []
+ for key, value in state_dict.items():
+ if isinstance(key, str):
+ if isinstance(value, dict):
+ keys.append(key + "|" + convert_keys_dict_to_single_str(value, with_shape=with_shape))
+ else:
+ if with_shape:
+ shape = "_".join(map(str, list(value)))
+ keys.append(key + ":" + shape)
+ keys.append(key)
+ keys.sort()
+ keys_str = ",".join(keys)
+ return keys_str
+
+
+def hash_model_file(path, with_shape=True):
+ keys_dict = load_keys_dict(path)
+ keys_str = convert_keys_dict_to_single_str(keys_dict, with_shape=with_shape)
+ keys_str = keys_str.encode(encoding="UTF-8")
+ return hashlib.md5(keys_str).hexdigest()
diff --git a/diffsynth/core/loader/model.py b/diffsynth/core/loader/model.py
new file mode 100644
index 0000000000000000000000000000000000000000..340e87aa15b89f3af0144a6333d66aa96c91712e
--- /dev/null
+++ b/diffsynth/core/loader/model.py
@@ -0,0 +1,167 @@
+from ..vram.initialization import skip_model_initialization
+from ..vram.disk_map import DiskMap
+from ..vram.layers import enable_vram_management
+from .file import load_state_dict, load_metadata_from_safetensors
+import torch
+from contextlib import contextmanager
+from transformers.integrations import is_deepspeed_zero3_enabled
+from transformers.utils import ContextManagers
+
+
+def load_model(model_class, path, config=None, torch_dtype=torch.bfloat16, device="cpu", state_dict_converter=None, use_disk_map=False, module_map=None, vram_config=None, vram_limit=None, state_dict=None, quantize=None):
+ config = {} if config is None else config
+ with ContextManagers(get_init_context(torch_dtype=torch_dtype, device=device)):
+ model = model_class(**config)
+ # What is `module_map`?
+ # This is a module mapping table for VRAM management.
+ if module_map is not None and quantize is None:
+ devices = [vram_config["offload_device"], vram_config["onload_device"], vram_config["preparing_device"], vram_config["computation_device"]]
+ device = [d for d in devices if d != "disk"][0]
+ dtypes = [vram_config["offload_dtype"], vram_config["onload_dtype"], vram_config["preparing_dtype"], vram_config["computation_dtype"]]
+ dtype = [d for d in dtypes if d != "disk"][0]
+ if vram_config["offload_device"] != "disk":
+ if state_dict is None: state_dict = DiskMap(path, device, torch_dtype=dtype)
+ if state_dict_converter is not None:
+ state_dict = state_dict_converter(state_dict)
+ else:
+ state_dict = {i: state_dict[i] for i in state_dict}
+ model.load_state_dict(state_dict, assign=True)
+ model = enable_vram_management(model, module_map, vram_config=vram_config, disk_map=None, vram_limit=vram_limit)
+ else:
+ disk_map = DiskMap(path, device, state_dict_converter=state_dict_converter)
+ model = enable_vram_management(model, module_map, vram_config=vram_config, disk_map=disk_map, vram_limit=vram_limit)
+ elif quantize is not None and module_map is not None:
+ if "disk" in vram_config.values():
+ if not quantize.load_prequantized:
+ raise ValueError("Disk offload with quantization is only supported for pre-quantized checkpoints (load_prequantized=True).")
+ devices = [vram_config[k] for k in ("offload_device", "onload_device", "preparing_device", "computation_device")]
+ load_device = [d for d in devices if d != "disk"][0]
+ disk_map = DiskMap(path, load_device, torch_dtype=None, state_dict_converter=state_dict_converter)
+ metadata = load_metadata_from_safetensors(path)
+ model = quantize.prepare_for_prequantized_load(model, compute_dtype=vram_config["computation_dtype"])
+ model = enable_vram_management(model, module_map, vram_config=vram_config, disk_map=disk_map, vram_limit=vram_limit, quantize=quantize, metadata=metadata)
+ else:
+ offload_device = vram_config["offload_device"]
+ computation_device = vram_config["computation_device"]
+ computation_dtype = vram_config["computation_dtype"]
+ load_dtype = None if quantize.load_prequantized else computation_dtype
+ if state_dict is None: state_dict = DiskMap(path, offload_device, torch_dtype=load_dtype)
+ if state_dict_converter is not None:
+ state_dict = state_dict_converter(state_dict)
+ else:
+ state_dict = {i: state_dict[i] for i in state_dict}
+
+ if quantize.load_prequantized:
+ model = quantize.prepare_for_prequantized_load(model, compute_dtype=computation_dtype)
+ state_dict = quantize.unflatten_state_dict(state_dict, load_metadata_from_safetensors(path))
+
+ model.load_state_dict(state_dict, assign=True)
+ state_dict = None
+
+ model = quantize.quantize_model(model, compute_device=computation_device, model_device=offload_device)
+ model = quantize.dequantize_model(model, compute_dtype=computation_dtype, compute_device=computation_device, model_device=offload_device)
+ model = model.to(dtype=computation_dtype, device=offload_device)
+ model = enable_vram_management(model, module_map, vram_config=vram_config, disk_map=None, vram_limit=vram_limit, quantize=quantize)
+ elif quantize is not None:
+ # Weight-only quantization (see `diffsynth.core.quant`), isolated from the normal path below.
+ if quantize.load_prequantized:
+ load_device, load_dtype = device, None
+ else:
+ load_device, load_dtype = "cpu", torch_dtype
+
+ if state_dict is not None:
+ pass
+ elif use_disk_map:
+ state_dict = DiskMap(path, load_device, torch_dtype=load_dtype)
+ else:
+ state_dict = load_state_dict(path, load_dtype, load_device)
+
+ if state_dict_converter is not None:
+ state_dict = state_dict_converter(state_dict)
+ else:
+ state_dict = {i: state_dict[i] for i in state_dict}
+
+ if quantize.load_prequantized:
+ model = quantize.prepare_for_prequantized_load(model, compute_dtype=torch_dtype or torch.bfloat16)
+ state_dict = quantize.unflatten_state_dict(state_dict, load_metadata_from_safetensors(path))
+
+ model.load_state_dict(state_dict, assign=True)
+ model = quantize.quantize_model(model, compute_device=device, model_device=device)
+ model = quantize.dequantize_model(model, compute_dtype=torch_dtype or torch.bfloat16)
+ model = model.to(dtype=torch_dtype, device=device)
+ else:
+ # Why do we use `DiskMap`?
+ # Sometimes a model file contains multiple models,
+ # and DiskMap can load only the parameters of a single model,
+ # avoiding the need to load all parameters in the file.
+ if state_dict is not None:
+ pass
+ elif use_disk_map:
+ state_dict = DiskMap(path, device, torch_dtype=torch_dtype)
+ else:
+ state_dict = load_state_dict(path, torch_dtype, device)
+ # Why do we use `state_dict_converter`?
+ # Some models are saved in complex formats,
+ # and we need to convert the state dict into the appropriate format.
+ if state_dict_converter is not None:
+ state_dict = state_dict_converter(state_dict)
+ else:
+ state_dict = {i: state_dict[i] for i in state_dict}
+ # Why does DeepSpeed ZeRO Stage 3 need to be handled separately?
+ # Because at this stage, model parameters are partitioned across multiple GPUs.
+ # Loading them directly could lead to excessive GPU memory consumption.
+ if is_deepspeed_zero3_enabled():
+ from transformers.integrations.deepspeed import _load_state_dict_into_zero3_model
+ _load_state_dict_into_zero3_model(model, state_dict)
+ else:
+ model.load_state_dict(state_dict, assign=True)
+ # Why do we call `to()`?
+ # Because some models override the behavior of `to()`,
+ # especially those from libraries like Transformers.
+ model = model.to(dtype=torch_dtype, device=device)
+ if quantize is not None:
+ # Downstream steps (e.g. LoRA hot-loading) need the config to handle the quantized layers.
+ model.quantize_config = quantize
+ if hasattr(model, "eval"):
+ model = model.eval()
+ return model
+
+
+def load_model_with_disk_offload(model_class, path, config=None, torch_dtype=torch.bfloat16, device="cpu", state_dict_converter=None, module_map=None):
+ if isinstance(path, str):
+ path = [path]
+ config = {} if config is None else config
+ with skip_model_initialization():
+ model = model_class(**config)
+ if hasattr(model, "eval"):
+ model = model.eval()
+ disk_map = DiskMap(path, device, state_dict_converter=state_dict_converter)
+ vram_config = {
+ "offload_dtype": "disk",
+ "offload_device": "disk",
+ "onload_dtype": "disk",
+ "onload_device": "disk",
+ "preparing_dtype": torch.float8_e4m3fn,
+ "preparing_device": device,
+ "computation_dtype": torch_dtype,
+ "computation_device": device,
+ }
+ enable_vram_management(model, module_map, vram_config=vram_config, disk_map=disk_map, vram_limit=80)
+ return model
+
+
+def get_init_context(torch_dtype, device):
+ if is_deepspeed_zero3_enabled():
+ from transformers.modeling_utils import set_zero3_state
+ import deepspeed
+ # Why do we use "deepspeed.zero.Init"?
+ # Weight segmentation of the model can be performed on the CPU side
+ # and loading the segmented weights onto the computing card
+ init_contexts = [deepspeed.zero.Init(remote_device=device, dtype=torch_dtype), set_zero3_state()]
+ else:
+ # Why do we use `skip_model_initialization`?
+ # It skips the random initialization of model parameters,
+ # thereby speeding up model loading and avoiding excessive memory usage.
+ init_contexts = [skip_model_initialization()]
+
+ return init_contexts
diff --git a/diffsynth/core/npu_patch/npu_fused_operator.py b/diffsynth/core/npu_patch/npu_fused_operator.py
new file mode 100644
index 0000000000000000000000000000000000000000..71660414518ab90e59657813cc927a7e9368faeb
--- /dev/null
+++ b/diffsynth/core/npu_patch/npu_fused_operator.py
@@ -0,0 +1,30 @@
+import torch
+from ..device.npu_compatible_device import get_device_type
+try:
+ import torch_npu
+except:
+ pass
+
+
+def rms_norm_forward_npu(self, hidden_states):
+ "npu rms fused operator for RMSNorm.forward from diffsynth\models\general_modules.py"
+ if hidden_states.dtype != self.weight.dtype:
+ hidden_states = hidden_states.to(self.weight.dtype)
+ return torch_npu.npu_rms_norm(hidden_states, self.weight, self.eps)[0]
+
+
+def rms_norm_forward_transformers_npu(self, hidden_states):
+ "npu rms fused operator for transformers"
+ if hidden_states.dtype != self.weight.dtype:
+ hidden_states = hidden_states.to(self.weight.dtype)
+ return torch_npu.npu_rms_norm(hidden_states, self.weight, self.variance_epsilon)[0]
+
+
+def rotary_emb_Zimage_npu(self, x_in: torch.Tensor, freqs_cis: torch.Tensor):
+ "npu rope fused operator for Zimage"
+ with torch.amp.autocast(get_device_type(), enabled=False):
+ freqs_cis = freqs_cis.unsqueeze(2)
+ cos, sin = torch.chunk(torch.view_as_real(freqs_cis), 2, dim=-1)
+ cos = cos.expand(-1, -1, -1, -1, 2).flatten(-2)
+ sin = sin.expand(-1, -1, -1, -1, 2).flatten(-2)
+ return torch_npu.npu_rotary_mul(x_in, cos, sin, rotary_mode="interleave").to(x_in)
\ No newline at end of file
diff --git a/diffsynth/core/offload_training/__init__.py b/diffsynth/core/offload_training/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..0a2408e512c2256288bdc9a0235a68810d9f18fe
--- /dev/null
+++ b/diffsynth/core/offload_training/__init__.py
@@ -0,0 +1 @@
+from .manager import OffloadTrainingManager
diff --git a/diffsynth/core/offload_training/manager.py b/diffsynth/core/offload_training/manager.py
new file mode 100644
index 0000000000000000000000000000000000000000..d6c38a21a27f01cf9a604ee18c78093087984624
--- /dev/null
+++ b/diffsynth/core/offload_training/manager.py
@@ -0,0 +1,177 @@
+"""
+Layer offloading for training — hook-based CPU offload.
+
+Hook lifecycle per module:
+
+No checkpointing:
+ forward_pre(load→GPU) → forward() → forward_hook(offload)
+ backward_pre(load→GPU) → backward() → backward_hook(offload)
+
+With checkpointing (use_reentrant=False):
+ First forward:
+ forward_pre(load→GPU) → forward() → forward_hook(offload, mark in_recompute)
+ Recomputing forward (during backward):
+ forward_pre(load→GPU) → forward() → forward_hook(in_recompute=True → keep GPU)
+ Backward:
+ backward_pre(load→GPU) → backward() → backward_hook(offload)
+"""
+import torch
+import torch.nn as nn
+import warnings
+from .offloader import StaticParamOffloader, TrainableParamOffloader, AlwaysOnGPUParamOffloader, BufferOffloader
+from .memory_buffer import PinnedArenaPool, BaseBufferPool
+warnings.filterwarnings("ignore", message="Full backward hook is firing when gradients are computed with respect to module outputs")
+
+
+def has_parameters(module: nn.Module) -> bool:
+ return len(list(module.parameters())) > 0
+
+def count_parameters(module: nn.Module) -> int:
+ return sum(p.numel() for p in module.parameters())
+
+def is_leaf_module(module: nn.Module) -> bool:
+ return len(list(module.children())) == 0
+
+
+class UnitWiseParamManager:
+ def __init__(self, model: nn.Module, target_device: torch.device, enable_optimizer_cpu_offload: bool = False, params: list = None, buffers: list = None, memory_buffer: BaseBufferPool = None):
+ self.model = model
+ self.target_device = target_device
+ self.param_offloaders = {}
+ for param in (model.parameters() if params is None else params):
+ if not param.requires_grad:
+ self.param_offloaders[id(param)] = StaticParamOffloader(param, target_device, memory_buffer=memory_buffer)
+ else:
+ if enable_optimizer_cpu_offload:
+ self.param_offloaders[id(param)] = TrainableParamOffloader(param, target_device)
+ else:
+ self.param_offloaders[id(param)] = AlwaysOnGPUParamOffloader(param, target_device)
+ if buffers is not None and len(buffers) > 0:
+ for mod, buf_name, buf in buffers:
+ self.param_offloaders[id(buf)] = BufferOffloader(mod, buf_name, buf, target_device, memory_buffer=memory_buffer)
+
+ def move_gradients_to_cpu(self):
+ for offloader in self.param_offloaders.values():
+ offloader.offload_grad()
+
+ def onload_module(self, module: nn.Module):
+ for param in module.parameters(recurse=False):
+ if id(param) in self.param_offloaders:
+ self.param_offloaders[id(param)].onload()
+ for name, buf in module.named_buffers(recurse=False):
+ if id(buf) in self.param_offloaders:
+ self.param_offloaders[id(buf)].onload()
+
+ def offload_module(self, module: nn.Module):
+ for param in module.parameters(recurse=False):
+ if id(param) in self.param_offloaders:
+ self.param_offloaders[id(param)].offload()
+ for name, buf in module.named_buffers(recurse=False):
+ if id(buf) in self.param_offloaders:
+ self.param_offloaders[id(buf)].offload()
+
+
+class UnitWiseHookManager:
+ def __init__(self, model: nn.Module, target_device: torch.device, enable_optimizer_cpu_offload: bool = False,
+ params: list = None, buffers: list = None, memory_buffer: BaseBufferPool = None):
+ self.param_manager = UnitWiseParamManager(model, target_device, enable_optimizer_cpu_offload, params=params, buffers=buffers, memory_buffer=memory_buffer)
+ self._in_recompute: set = set()
+ self._register_hooks(model)
+
+ def _register_hooks(self, module: nn.Module):
+ def forward_pre_hook(mod, args):
+ self.param_manager.onload_module(mod)
+
+ def forward_hook(mod, args, output):
+ if mod in self._in_recompute:
+ return
+ self._in_recompute.add(mod)
+ self.param_manager.offload_module(mod)
+
+ def backward_pre_hook(mod, grad_output):
+ self.param_manager.onload_module(mod)
+
+ def backward_hook(mod, grad_input, grad_output):
+ self.param_manager.offload_module(mod)
+
+ module.register_forward_pre_hook(forward_pre_hook)
+ module.register_forward_hook(forward_hook)
+ module.register_full_backward_pre_hook(backward_pre_hook)
+ if is_leaf_module(module):
+ module.register_full_backward_hook(backward_hook)
+ else:
+ # Parent module backward_hook fires before child backward completes.
+ # Register on leaf children instead.
+ sub_modules = [m for m in module.modules() if is_leaf_module(m) and has_parameters(m)]
+ for sub_mod in sub_modules:
+ sub_mod.register_full_backward_hook(backward_hook)
+
+ def after_backward(self):
+ self._in_recompute.clear()
+ self.param_manager.move_gradients_to_cpu()
+
+ @property
+ def managed_param_ids(self):
+ return set(self.param_manager.param_offloaders.keys())
+
+
+class OffloadTrainingManager:
+ def __init__(self, model: nn.Module, target_device: torch.device, enable_optimizer_cpu_offload: bool = False, cpu_offload_split_threshold: int = None):
+ self.model = model
+ self.target_device = target_device
+ self.enable_optimizer_cpu_offload = enable_optimizer_cpu_offload
+ cpu_offload_split_threshold = cpu_offload_split_threshold * 1024 * 1024 if cpu_offload_split_threshold is not None else None
+ self._register_units(model, target_device, enable_optimizer_cpu_offload, cpu_offload_split_threshold)
+
+ def _register_units(self, model: nn.Module, target_device: torch.device, enable_optimizer_cpu_offload: bool, cpu_offload_split_threshold: int = None):
+ self.memory_buffer = PinnedArenaPool.from_model(model)
+ units = self._find_units_recursive(model, cpu_offload_split_threshold)
+ self.units = [UnitWiseHookManager(u, target_device, enable_optimizer_cpu_offload, memory_buffer=self.memory_buffer) for u in units]
+
+ managed_param_ids = set().union(*[unit.managed_param_ids for unit in self.units])
+ orphan_params, orphan_buffers = self._find_orphan_params_and_buffers(model, managed_param_ids)
+ for orphan_module in set(orphan_params.keys()) | set(orphan_buffers.keys()):
+ params = orphan_params.get(orphan_module, [])
+ buffers = orphan_buffers.get(orphan_module, [])
+ self.units.append(UnitWiseHookManager(orphan_module, target_device, enable_optimizer_cpu_offload, params=params, buffers=buffers, memory_buffer=self.memory_buffer))
+
+ def _find_orphan_params_and_buffers(self, model: nn.Module, managed_param_ids: set):
+ orphan_params_by_module = {}
+ for _, mod in model.named_modules():
+ for param in mod.parameters(recurse=False):
+ if id(param) not in managed_param_ids:
+ orphan_params_by_module.setdefault(mod, []).append(param)
+ # Collect orphan buffers grouped by owner module
+ orphan_buffers_by_module = {}
+ for _, mod in model.named_modules():
+ for name, buf in mod.named_buffers(recurse=False):
+ orphan_buffers_by_module.setdefault(mod, []).append((mod, name, buf))
+
+ return orphan_params_by_module, orphan_buffers_by_module
+
+ def _find_units_recursive(self, module: nn.Module, cpu_offload_split_threshold: int = None) -> list:
+ if cpu_offload_split_threshold is None:
+ return [m for m in module.modules() if is_leaf_module(m) and has_parameters(m)]
+ if self._should_force_recurse(module, cpu_offload_split_threshold):
+ units = []
+ for child in module.children():
+ units.extend(self._find_units_recursive(child, cpu_offload_split_threshold))
+ return units
+ return [module]
+
+ def _should_force_recurse(self, module: nn.Module, cpu_offload_split_threshold: int = None) -> bool:
+ if is_leaf_module(module):
+ return False
+ if (
+ count_parameters(module) > cpu_offload_split_threshold
+ or ('forward' not in type(module).__dict__)
+ or (hasattr(module, 'encode') and hasattr(module, 'decode'))
+ ):
+ return True
+ return False
+
+ # run after backward() and before optimizer.step()
+ def after_backward(self):
+ for unit in self.units:
+ unit.after_backward()
+ torch.cuda.synchronize()
diff --git a/diffsynth/core/offload_training/memory_buffer.py b/diffsynth/core/offload_training/memory_buffer.py
new file mode 100644
index 0000000000000000000000000000000000000000..cef0a39b928a079bd0b55720440b5114e9b8278a
--- /dev/null
+++ b/diffsynth/core/offload_training/memory_buffer.py
@@ -0,0 +1,136 @@
+import torch
+
+ALIGNMENT = 64
+
+def _align_up(x: int, alignment: int = ALIGNMENT) -> int:
+ return (x + alignment - 1) // alignment * alignment
+
+def _next_power_of_two(x: int) -> int:
+ """
+ Smallest power of two >= x.
+ For power-of-two x=2^k: (x-1) has bit_length=k, so 1< int:
+ """Largest power-of-two <= x."""
+ return 1 if x <= 1 else 1 << (x.bit_length() - 1)
+
+def _tensor_storage_size(tensor: torch.Tensor) -> int:
+ return _align_up(tensor.numel() * tensor.element_size())
+
+
+class BaseBufferPool:
+ """Naive per-tensor pin_memory allocation. No pre-allocation, no memory saving."""
+
+ def allocate_like(self, tensor: torch.Tensor) -> torch.Tensor:
+ return tensor.pin_memory()
+
+ @classmethod
+ def from_model(cls, model: torch.nn.Module, **kwargs):
+ return cls()
+
+
+class PinnedBuffer:
+ """Single pinned uint8 buffer with bump-pointer allocation. Lazy: actual memory allocated on first allocate_like."""
+
+ def __init__(self, size: int):
+ self._size = size
+ self._buf: torch.Tensor | None = None
+ self._offset = 0
+
+ def _ensure_allocated(self):
+ if self._buf is None:
+ self._buf = torch.empty(self._size, dtype=torch.uint8, device="cpu", pin_memory=True)
+
+ @property
+ def capacity(self) -> int:
+ return self._size
+
+ @property
+ def remaining(self) -> int:
+ return self._size if self._buf is None else self._buf.numel() - self._offset
+
+ @property
+ def used(self) -> int:
+ return self._offset
+
+ @classmethod
+ def from_tensor(cls, tensor: torch.Tensor, min_size: int = 1 * 1024**3):
+ size = max(_tensor_storage_size(tensor) + ALIGNMENT, min_size)
+ return cls(_next_power_of_two(size))
+
+ def allocate_like(self, tensor: torch.Tensor, *, copy: bool = True, non_blocking: bool = False) -> torch.Tensor | None:
+ """Try to allocate a view for tensor. Returns None if not enough space."""
+ num_bytes = tensor.numel() * tensor.element_size()
+ if num_bytes > self.remaining:
+ return None
+ self._ensure_allocated()
+ view = self._buf.narrow(0, self._offset, num_bytes).view(tensor.dtype).reshape(tuple(tensor.shape))
+ if copy:
+ view.copy_(tensor, non_blocking=bool(non_blocking and tensor.device.type == "cuda"))
+ self._offset = _align_up(self._offset + num_bytes)
+ return view
+
+
+class PinnedArenaPool(BaseBufferPool):
+ """Pinned arena pool — pre-allocate pinned memory, avoid per-tensor cudaHostAlloc overhead.
+
+ Pool strategy:
+ 1. Sizing: from_model() scans all non-trainable params + buffers, sums their aligned sizes
+ as total_bytes. max_chunk_size is raised to fit the largest single tensor.
+ 2. Decomposition: total_bytes is split into power-of-two chunks (min_chunk_size ~ max_chunk_size).
+ Each chunk becomes one PinnedBuffer (lazy — actual pin_memory on first use).
+ 3. Allocation: allocate_like() sequentially probes each buffer for space (first-fit).
+ Each PinnedBuffer uses bump-pointer with ALIGNMENT padding between tensors.
+ 4. Growth: if all existing buffers are full, _grow() appends a new PinnedBuffer sized
+ to the requesting tensor (at least min_chunk_size, power-of-two rounded).
+ 5. Fallback: on any exception, falls back to per-tensor pin_memory().
+ """
+
+ def __init__(self, total_bytes: int, min_chunk_size: int = 1 * 1024**3, max_chunk_size: int = 4 * 1024**3):
+ self.min_chunk_size = _next_power_of_two(int(min_chunk_size))
+ self.max_chunk_size = _next_power_of_two(int(max_chunk_size))
+ self.min_chunk_size = min(self.min_chunk_size, self.max_chunk_size)
+ self._buffers = [PinnedBuffer(s) for s in self._decompose(total_bytes, self.min_chunk_size, self.max_chunk_size)]
+
+ @classmethod
+ def from_model(cls, model: torch.nn.Module, min_chunk_size: int = 1 * 1024**3, max_chunk_size: int = 4 * 1024**3):
+ """Size pool for all non-trainable params + buffers."""
+ tensors = [p for p in model.parameters() if not p.requires_grad] + list(model.buffers())
+ total = sum(_tensor_storage_size(t) for t in tensors)
+ max_tensor_size = max((_tensor_storage_size(t) for t in tensors), default=0)
+ max_chunk_size = _next_power_of_two(max_tensor_size) if max_tensor_size > max_chunk_size else max_chunk_size
+ return cls(total, min_chunk_size=min_chunk_size, max_chunk_size=max_chunk_size)
+
+ @staticmethod
+ def _decompose(total_bytes: int, min_chunk_size: int, max_chunk_size: int) -> list:
+ """Decompose total_bytes into power-of-two chunks capped by min/max."""
+ if total_bytes <= 0:
+ return []
+ chunks, remaining = [], total_bytes
+ while remaining > 0:
+ chunk = max(min(_prev_power_of_two(remaining), max_chunk_size), min_chunk_size)
+ chunks.append(chunk)
+ remaining -= chunk
+ chunks.sort(reverse=True)
+ return chunks
+
+ def _grow(self, tensor: torch.Tensor):
+ self._buffers.append(PinnedBuffer.from_tensor(tensor, min_size=self.min_chunk_size))
+
+ def allocate_like(self, tensor: torch.Tensor, *, copy: bool = True, require_contiguous: bool = True, non_blocking: bool = False) -> torch.Tensor:
+ """Allocate a pinned view. Falls back to per-tensor pin_memory on failure."""
+ src = tensor.detach()
+ if require_contiguous and not src.is_contiguous():
+ src = src.contiguous()
+ try:
+ for buf in self._buffers:
+ view = buf.allocate_like(src, copy=copy, non_blocking=non_blocking)
+ if view is not None:
+ return view
+ self._grow(src)
+ return self._buffers[-1].allocate_like(src, copy=copy, non_blocking=non_blocking)
+ except Exception:
+ return src.pin_memory()
diff --git a/diffsynth/core/offload_training/offloader.py b/diffsynth/core/offload_training/offloader.py
new file mode 100644
index 0000000000000000000000000000000000000000..0bea513c820229f9074b1a0077b4555f5f51958b
--- /dev/null
+++ b/diffsynth/core/offload_training/offloader.py
@@ -0,0 +1,71 @@
+import torch
+from .memory_buffer import BaseBufferPool
+
+
+class OffloaderMixin:
+ def onload(self):
+ pass
+
+ def offload(self):
+ pass
+
+ def offload_grad(self):
+ pass
+
+
+class BaseParamOffloader(OffloaderMixin):
+ def __init__(self, param: torch.nn.Parameter, target_device: torch.device):
+ self.param = param
+ self.target_device = target_device
+
+
+class StaticParamOffloader(BaseParamOffloader):
+ def __init__(self, param: torch.nn.Parameter, target_device: torch.device, memory_buffer: BaseBufferPool = None):
+ super().__init__(param, target_device)
+ cpu_data = param.data.cpu().detach().contiguous()
+ self.cpu_copy = memory_buffer.allocate_like(cpu_data) if memory_buffer is not None else cpu_data.pin_memory()
+ self._placeholder = torch.empty(0, device=target_device, dtype=param.dtype)
+ param.data = self._placeholder
+
+ def onload(self):
+ self.param.data = self.cpu_copy.to(self.target_device, non_blocking=True)
+
+ def offload(self):
+ self.param.data = self._placeholder
+
+
+class TrainableParamOffloader(BaseParamOffloader):
+ def __init__(self, param: torch.nn.Parameter, target_device: torch.device):
+ super().__init__(param, target_device)
+ assert param.requires_grad, "TrainableParamOffloader can only be used with trainable parameters"
+
+ def onload(self):
+ self.param.data = self.param.data.to(self.target_device, non_blocking=True)
+
+ def offload(self):
+ self.param.data = self.param.data.to('cpu', non_blocking=True)
+
+ def offload_grad(self):
+ if self.param.grad is not None:
+ self.param.grad = self.param.grad.to('cpu', non_blocking=True)
+
+
+class AlwaysOnGPUParamOffloader(BaseParamOffloader):
+ def __init__(self, param, target_device):
+ super().__init__(param, target_device)
+ self.param.data = self.param.data.to(self.target_device)
+
+
+class BufferOffloader(OffloaderMixin):
+ def __init__(self, module: torch.nn.Module, buf_name: str, buf: torch.Tensor, target_device: torch.device, memory_buffer: BaseBufferPool = None):
+ self.module = module
+ self.buf_name = buf_name
+ self.target_device = target_device
+ cpu_data = buf.data.cpu().contiguous()
+ self.cpu_copy = memory_buffer.allocate_like(cpu_data) if memory_buffer is not None else cpu_data.pin_memory()
+
+ def onload(self):
+ self.module._buffers[self.buf_name] = self.cpu_copy.to(self.target_device, non_blocking=True)
+
+ def offload(self):
+ self.module._buffers[self.buf_name] = self.cpu_copy
diff --git a/diffsynth/core/quant/__init__.py b/diffsynth/core/quant/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..9167e6b01850740a235bbfcb8be1dbcdf90c4cc6
--- /dev/null
+++ b/diffsynth/core/quant/__init__.py
@@ -0,0 +1,3 @@
+from .base import QuantBackend, BackendConfig, QUANT_BACKENDS, register_quant_backend, check_differentiable, check_backend_contract
+from .config import QuantizeConfig, MixedQuantizeConfig, QuantMethodSpec, QUANT_METHODS, register_quant_method, describe_quant_method
+from . import backends
diff --git a/diffsynth/core/quant/backends/__init__.py b/diffsynth/core/quant/backends/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..41c6477fc79c21a583f12b6bf25a07402a092217
--- /dev/null
+++ b/diffsynth/core/quant/backends/__init__.py
@@ -0,0 +1,28 @@
+import importlib
+
+_LAZY_BACKENDS = {
+ "bitsandbytes": ".bitsandbytes",
+ "torchao": ".torchao",
+ "comfy_kitchen": ".comfy_kitchen",
+}
+_loaded = set()
+
+
+def load_backend(name):
+ module = _LAZY_BACKENDS.get(name)
+ if module is not None and name not in _loaded:
+ importlib.import_module(module, __name__)
+ _loaded.add(name)
+
+
+def load_all_backends():
+ for name in _LAZY_BACKENDS:
+ load_backend(name)
+
+
+def load_backend_for_method(method):
+ from ..config import QUANT_METHODS
+ for name in _LAZY_BACKENDS:
+ load_backend(name)
+ if method in QUANT_METHODS:
+ return
diff --git a/diffsynth/core/quant/backends/bitsandbytes.py b/diffsynth/core/quant/backends/bitsandbytes.py
new file mode 100644
index 0000000000000000000000000000000000000000..b08a54672331e61799e020788df07c81039e1fe9
--- /dev/null
+++ b/diffsynth/core/quant/backends/bitsandbytes.py
@@ -0,0 +1,155 @@
+from dataclasses import dataclass, field
+
+import torch
+from ..base import QuantBackend, BackendConfig, register_quant_backend
+from ..config import register_quant_method
+
+try:
+ import bitsandbytes as bnb
+ BITSANDBYTES_AVAILABLE = True
+except ImportError:
+ BITSANDBYTES_AVAILABLE = False
+
+
+if BITSANDBYTES_AVAILABLE:
+ class BitsAndBytesLinear4bit(bnb.nn.Linear4bit):
+ def _load_from_state_dict(self, state_dict, prefix, local_metadata, strict, missing_keys, unexpected_keys, error_msgs):
+ local_names = set()
+ for name, current in list(self._parameters.items()) + list(self._buffers.items()):
+ if current is None:
+ continue
+ local_names.add(name)
+ key = prefix + name
+ if key in state_dict:
+ value = state_dict[key]
+ if name in self._parameters:
+ self._parameters[name] = value if isinstance(value, torch.nn.Parameter) else torch.nn.Parameter(value, requires_grad=False)
+ else:
+ self._buffers[name] = value
+ elif strict:
+ missing_keys.append(key)
+ if strict:
+ for key in state_dict:
+ if key.startswith(prefix) and key[len(prefix):].split(".", 1)[0] not in local_names:
+ unexpected_keys.append(key)
+
+
+@register_quant_backend("bitsandbytes")
+class BitsAndBytesQuantBackend(QuantBackend):
+ """Adapter over `bitsandbytes.nn.Linear4bit` (NF4 / FP4, weight-only)."""
+
+ project_url = "https://github.com/bitsandbytes-foundation/bitsandbytes"
+
+ def validate_environment(self):
+ if not BITSANDBYTES_AVAILABLE:
+ raise ImportError(
+ "bitsandbytes is required for bitsandbytes quantization methods. "
+ 'Please install it via `pip install bitsandbytes` or `pip install "diffsynth[quant]"`.'
+ )
+
+ def capabilities(self):
+ return {
+ "is_serializable": True,
+ "is_differentiable": True,
+ "is_compileable": False,
+ "requires_calibration": False,
+ }
+
+ def quantized_linear_classes(self):
+ return (BitsAndBytesLinear4bit,)
+
+ def create_quantized_linear(self, linear, compute_device=None, model_device=None):
+ """The `meta` shell avoids allocating an fp weight; bnb quantizes while `Params4bit` moves to `compute_device`."""
+ cfg = self.config
+ compute_dtype = linear.weight.dtype
+ if compute_device is None:
+ compute_device = linear.weight.device
+
+ with torch.device("meta"):
+ quant_linear = BitsAndBytesLinear4bit(
+ linear.in_features,
+ linear.out_features,
+ bias=linear.bias is not None,
+ compute_dtype=compute_dtype,
+ compress_statistics=cfg.compress_statistics,
+ quant_type=cfg.quant_type,
+ quant_storage=cfg.quant_storage,
+ )
+ weight = bnb.nn.Params4bit(
+ linear.weight.data.contiguous(),
+ requires_grad=False,
+ compress_statistics=cfg.compress_statistics,
+ blocksize=cfg.blocksize,
+ quant_type=cfg.quant_type,
+ quant_storage=cfg.quant_storage,
+ )
+ quant_linear.weight = weight.to(compute_device)
+ if linear.bias is not None:
+ quant_linear.bias = torch.nn.Parameter(linear.bias.data.to(device=compute_device), requires_grad=False)
+ return quant_linear.to(device=model_device)
+
+ def create_quantized_linear_shell(self, linear, compute_dtype):
+ with torch.device("meta"):
+ shell = BitsAndBytesLinear4bit(
+ linear.in_features,
+ linear.out_features,
+ bias=linear.bias is not None,
+ compute_dtype=compute_dtype,
+ compress_statistics=self.config.compress_statistics,
+ quant_type=self.config.quant_type,
+ quant_storage=self.config.quant_storage,
+ )
+ return shell
+
+ def unflatten_state_dict(self, state_dict, metadata):
+ rebuilt = {}
+ quant_state = {}
+ for key, value in state_dict.items():
+ weight_key = key.rsplit(".weight.", 1)[0] + ".weight" if ".weight." in key else None
+ if weight_key is not None:
+ quant_state.setdefault(weight_key, {})[key] = value
+ else:
+ rebuilt[key] = value
+ for weight_key, stats in quant_state.items():
+ if weight_key not in rebuilt:
+ raise ValueError(f"Found quant state for `{weight_key}` but no packed weight.")
+ rebuilt[weight_key] = bnb.nn.Params4bit.from_prequantized(
+ data=rebuilt[weight_key], quantized_stats=stats, requires_grad=False, device=rebuilt[weight_key].device,
+ )
+ return rebuilt
+
+ def dequantize_to_linear(self, module, compute_dtype, compute_device=None, model_device=None):
+ if compute_device is not None:
+ module = module.to(device=compute_device)
+ weight = module.weight
+ fp_weight = bnb.functional.dequantize_4bit(weight.data, weight.quant_state).to(compute_dtype)
+ linear = torch.nn.Linear(module.in_features, module.out_features, bias=module.bias is not None, device="meta")
+ linear.weight = torch.nn.Parameter(fp_weight, requires_grad=False)
+ if module.bias is not None:
+ linear.bias = torch.nn.Parameter(module.bias.data.to(dtype=compute_dtype, device=fp_weight.device), requires_grad=False)
+ if model_device is not None:
+ linear = linear.to(device=model_device)
+ return linear
+
+
+@dataclass
+class BitsAndBytes4bitConfig(BackendConfig):
+ """Shared 4bit knobs; `quant_type` is pinned by the per-method subclass."""
+
+ compress_statistics: bool = True
+ blocksize: int = None # None -> bnb default (64)
+ quant_storage: torch.dtype = torch.uint8
+
+
+@dataclass
+class BitsAndBytesNF4Config(BitsAndBytes4bitConfig):
+ quant_type: str = field(init=False, default="nf4")
+
+
+@dataclass
+class BitsAndBytesFP4Config(BitsAndBytes4bitConfig):
+ quant_type: str = field(init=False, default="fp4")
+
+
+register_quant_method("bitsandbytes_nf4", "bitsandbytes", BitsAndBytesNF4Config.from_kwargs, label="4bit, nf4, weight-only")
+register_quant_method("bitsandbytes_fp4", "bitsandbytes", BitsAndBytesFP4Config.from_kwargs, label="4bit, fp4, weight-only")
diff --git a/diffsynth/core/quant/backends/comfy_kitchen.py b/diffsynth/core/quant/backends/comfy_kitchen.py
new file mode 100644
index 0000000000000000000000000000000000000000..56b11fe92754d9059b657ac321b5afe5271059ca
--- /dev/null
+++ b/diffsynth/core/quant/backends/comfy_kitchen.py
@@ -0,0 +1,318 @@
+import copy
+import json
+from dataclasses import dataclass, field
+
+import torch
+from ..base import QuantBackend, BackendConfig, register_quant_backend
+from ..config import register_quant_method
+
+try:
+ import comfy_kitchen as ck
+ from comfy_kitchen.tensor import QuantizedTensor, get_layout_class
+ COMFY_KITCHEN_AVAILABLE = True
+except ImportError:
+ COMFY_KITCHEN_AVAILABLE = False
+
+
+COMFY_QUANT_KEY = "comfy_quant"
+SCALE_KEY = "weight_scale"
+
+
+class ComfyKitchenLinear(torch.nn.Linear):
+ """`nn.Linear` whose weight is a comfy-kitchen `QuantizedTensor`."""
+
+ def __init__(self, in_features, out_features, bias, *, layout, compute_dtype):
+ with torch.device("meta"):
+ super().__init__(in_features, out_features, bias=bias, dtype=compute_dtype)
+ self.layout = layout
+ self.compute_dtype = compute_dtype
+ self.weight.requires_grad_(False)
+ if self.bias is not None:
+ self.bias.requires_grad_(False)
+
+ def forward(self, x):
+ if not x.is_contiguous():
+ x = x.contiguous()
+ return super().forward(x)
+
+ def __deepcopy__(self, memo):
+ clone = type(self)(
+ self.in_features, self.out_features, bias=self.bias is not None,
+ layout=self.layout, compute_dtype=self.compute_dtype,
+ )
+ weight = self.weight
+ params = getattr(weight, "_params", None)
+ if params is None:
+ cloned_weight = weight.detach().clone()
+ else:
+ cloned_weight = QuantizedTensor(weight._qdata.clone(), weight._layout_cls, copy.deepcopy(params))
+ clone.weight = torch.nn.Parameter(cloned_weight, requires_grad=False)
+ if self.bias is not None:
+ clone.bias = torch.nn.Parameter(self.bias.detach().clone(), requires_grad=False)
+ memo[id(self)] = clone
+ return clone
+
+
+def _fp8_linear(x_2d, weight, bias, layout, input_scale):
+ """Quantize a 2D activation per tensor, then let comfy-kitchen's dispatch pick the fp8 matmul."""
+ q_x = QuantizedTensor.from_float(x_2d, layout, scale=input_scale)
+ return torch.nn.functional.linear(q_x, weight, bias)
+
+
+class _FP8LinearFunction(torch.autograd.Function):
+
+ @staticmethod
+ def forward(ctx, x, weight, bias, layout, input_scale):
+ out = _fp8_linear(x.detach().reshape(-1, x.shape[-1]), weight, bias, layout, input_scale)
+ ctx.save_for_backward(weight)
+ ctx.x_shape = x.shape
+ return out.reshape(*x.shape[:-1], out.shape[-1])
+
+ @staticmethod
+ @torch.autograd.function.once_differentiable
+ def backward(ctx, grad_output):
+ weight, = ctx.saved_tensors
+ grad_2d = grad_output.reshape(-1, grad_output.shape[-1])
+ grad_x = torch.mm(grad_2d, weight.dequantize().to(grad_2d.dtype)).reshape(ctx.x_shape)
+ grad_bias = grad_2d.sum(dim=0) if ctx.needs_input_grad[2] else None
+ return grad_x, None, grad_bias, None, None
+
+
+class ComfyKitchenFP8Linear(ComfyKitchenLinear):
+
+ def __init__(self, in_features, out_features, bias, *, layout, compute_dtype):
+ super().__init__(in_features, out_features, bias, layout=layout, compute_dtype=compute_dtype)
+ self.register_buffer("input_scale", None)
+
+ def _load_from_state_dict(self, state_dict, prefix, local_metadata, strict, missing_keys, unexpected_keys, error_msgs):
+ scale = state_dict.pop(prefix + "input_scale", None)
+ super()._load_from_state_dict(state_dict, prefix, local_metadata, strict, missing_keys, unexpected_keys, error_msgs)
+ if scale is not None:
+ self.input_scale = scale
+
+ def forward(self, x):
+ if not x.is_contiguous():
+ x = x.contiguous()
+ scale = self.input_scale
+ if scale is not None:
+ scale = scale.to(dtype=torch.float32, device=x.device)
+ if x.requires_grad:
+ return _FP8LinearFunction.apply(x, self.weight, self.bias, self.layout, scale)
+ out = _fp8_linear(x.reshape(-1, x.shape[-1]), self.weight, self.bias, self.layout, scale)
+ return out.reshape(*x.shape[:-1], out.shape[-1])
+
+ def __deepcopy__(self, memo):
+ clone = super().__deepcopy__(memo)
+ if self.input_scale is not None:
+ clone.input_scale = self.input_scale.detach().clone()
+ return clone
+
+
+class ComfyQuantFormat:
+ """Everything specific to one weight format that ComfyUI writes, one hook per direction."""
+
+ marker_name = None # the `format` string ComfyUI writes into the marker
+ layout_cls = None # name of the comfy-kitchen layout class
+ linear_cls = ComfyKitchenLinear # the `nn.Linear` subclass that runs this format
+ sibling_keys = () # per-layer tensors to drop, so `load_state_dict` never sees them
+
+ def params_from_marker(self, marker, qdata, layer_name):
+ """`Params` fields beyond `scale`, `orig_dtype` and `orig_shape`, read off the marker."""
+ return {}
+
+ def from_float_kwargs(self, in_features, config):
+ """Keywords for `QuantizedTensor.from_float` when quantizing an fp weight online."""
+ return {}
+
+ def marker_from_params(self, params):
+ """Marker fields to write back, so a re-saved checkpoint reloads identically."""
+ return {}
+
+
+class TensorWiseInt8Format(ComfyQuantFormat):
+
+ marker_name = "int8_tensorwise"
+ layout_cls = "TensorWiseINT8Layout"
+ DEFAULT_GROUPSIZE = 256
+
+ def params_from_marker(self, marker, qdata, layer_name):
+ nested = marker.get("params", {})
+ convrot = bool(marker.get("convrot", nested.get("convrot", False)))
+ groupsize = int(marker.get("convrot_groupsize", nested.get("convrot_groupsize", self.DEFAULT_GROUPSIZE)))
+ in_features = qdata.shape[1]
+ if convrot and in_features % groupsize:
+ raise ValueError(
+ f"The marker of `{layer_name}` asks for convrot with groupsize {groupsize}, "
+ f"which does not divide in_features={in_features}."
+ )
+ return {"is_weight": True, "convrot": convrot, "convrot_groupsize": groupsize}
+
+ def from_float_kwargs(self, in_features, config):
+ groupsize = config.convrot_groupsize
+ kwargs = {"is_weight": True, "per_channel": config.per_channel, "convrot": False, "convrot_groupsize": groupsize}
+ if not config.convrot:
+ return kwargs
+ if groupsize < 4 or groupsize & (groupsize - 1) or (groupsize.bit_length() - 1) % 2:
+ raise ValueError(f"convrot_groupsize must be a power of four >= 4, got {groupsize}.")
+ while groupsize >= 4 and in_features % groupsize:
+ groupsize //= 4
+ if groupsize >= 4:
+ kwargs.update(convrot=True, convrot_groupsize=groupsize)
+ return kwargs
+
+ def marker_from_params(self, params):
+ marker = {"convrot": bool(params.convrot)}
+ if params.convrot:
+ marker["convrot_groupsize"] = int(params.convrot_groupsize)
+ return marker
+
+
+class Float8E4M3Format(ComfyQuantFormat):
+
+ marker_name = "float8_e4m3fn"
+ layout_cls = "TensorCoreFP8Layout"
+ linear_cls = ComfyKitchenFP8Linear
+
+
+FORMATS = {fmt.marker_name: fmt for fmt in (TensorWiseInt8Format(), Float8E4M3Format())}
+
+
+def _parse_marker(marker_tensor, layer_name):
+ try:
+ marker = json.loads(bytes(marker_tensor.to(torch.uint8).tolist()).decode())
+ except Exception as error:
+ raise ValueError(f"`{layer_name}.{COMFY_QUANT_KEY}` is not valid JSON: {error}")
+ marker_name = marker.get("format")
+ if marker_name is None:
+ raise ValueError(f"`{layer_name}.{COMFY_QUANT_KEY}` has no `format` field.")
+ if marker_name not in FORMATS:
+ raise ValueError(f"`{layer_name}` uses format `{marker_name}`, not in {sorted(FORMATS)}.")
+ return FORMATS[marker_name], marker
+
+
+@register_quant_backend("comfy_kitchen")
+class ComfyKitchenQuantBackend(QuantBackend):
+ """Adapter over comfy-kitchen's `QuantizedTensor`, the quantized-weight format ComfyUI writes."""
+
+ project_url = "https://github.com/Comfy-Org/comfy-kitchen"
+
+ def validate_environment(self):
+ if not COMFY_KITCHEN_AVAILABLE:
+ raise ImportError(
+ "comfy-kitchen is required for comfy_kitchen quantization methods. "
+ 'Please install it via `pip install comfy-kitchen` or `pip install "diffsynth[quant]"`.'
+ )
+ if torch.version.cuda is None or tuple(int(v) for v in torch.version.cuda.split(".")[:2]) < (13, 0):
+ ck.registry.disable("cuda")
+
+ def capabilities(self):
+ return {
+ "is_serializable": True,
+ "is_differentiable": True,
+ "is_compileable": False,
+ "requires_calibration": False,
+ }
+
+ def quantized_linear_classes(self):
+ return tuple(dict.fromkeys(fmt.linear_cls for fmt in FORMATS.values()))
+
+ def create_quantized_linear(self, linear, compute_device=None, model_device=None):
+ quant_format = self._active_format()
+ if compute_device is not None:
+ linear = linear.to(device=compute_device)
+ weight = QuantizedTensor.from_float(
+ linear.weight.data, quant_format.layout_cls,
+ **quant_format.from_float_kwargs(linear.in_features, self.config),
+ )
+ quant_linear = self._build_linear(linear, linear.weight.dtype)
+ quant_linear.weight = torch.nn.Parameter(weight, requires_grad=False)
+ if linear.bias is not None:
+ quant_linear.bias = torch.nn.Parameter(linear.bias.data, requires_grad=False)
+ return quant_linear if model_device is None else quant_linear.to(device=model_device)
+
+ def create_quantized_linear_shell(self, linear, compute_dtype):
+ return self._build_linear(linear, compute_dtype)
+
+ def unflatten_state_dict(self, state_dict, metadata):
+ rebuilt = dict(state_dict)
+ suffix = f".{COMFY_QUANT_KEY}"
+ for layer_name in [key[: -len(suffix)] for key in state_dict if key.endswith(suffix)]:
+ quant_format, marker = _parse_marker(rebuilt.pop(layer_name + suffix), layer_name)
+ qdata = rebuilt.get(f"{layer_name}.weight")
+ scale = rebuilt.pop(f"{layer_name}.{SCALE_KEY}", None)
+ for key, value in (("weight", qdata), (SCALE_KEY, scale)):
+ if value is None:
+ raise ValueError(f"`{layer_name}` has a {COMFY_QUANT_KEY} marker but no `{key}`.")
+ for key in quant_format.sibling_keys:
+ rebuilt.pop(f"{layer_name}.{key}", None)
+ params = get_layout_class(quant_format.layout_cls).Params(
+ scale=scale,
+ orig_dtype=self.config.orig_dtype,
+ orig_shape=tuple(qdata.shape),
+ **quant_format.params_from_marker(marker, qdata, layer_name),
+ )
+ rebuilt[f"{layer_name}.weight"] = QuantizedTensor(qdata, quant_format.layout_cls, params)
+ return rebuilt
+
+ def flatten_state_dict(self, state_dict):
+ flattened = {}
+ for key, value in state_dict.items():
+ params = getattr(value, "_params", None)
+ if not key.endswith(".weight") or params is None:
+ flattened[key] = value
+ continue
+ layer_name = key[: -len(".weight")]
+ quant_format = next((f for f in FORMATS.values() if f.layout_cls == value._layout_cls), None)
+ if quant_format is None:
+ raise ValueError(f"`{layer_name}` uses layout `{value._layout_cls}`, which no format claims.")
+ flattened.update(value.state_dict(prefix=key))
+ marker = {"format": quant_format.marker_name, **quant_format.marker_from_params(params)}
+ flattened[f"{layer_name}.{COMFY_QUANT_KEY}"] = torch.frombuffer(
+ bytearray(json.dumps(marker).encode()), dtype=torch.uint8
+ )
+ return flattened, {}
+
+ def dequantize_to_linear(self, module, compute_dtype, compute_device=None, model_device=None):
+ if compute_device is not None:
+ module = module.to(device=compute_device)
+ fp_weight = module.weight.dequantize().to(compute_dtype)
+ linear = torch.nn.Linear(module.in_features, module.out_features, bias=module.bias is not None, device="meta")
+ linear.weight = torch.nn.Parameter(fp_weight, requires_grad=False)
+ if module.bias is not None:
+ linear.bias = torch.nn.Parameter(module.bias.data.to(dtype=compute_dtype, device=fp_weight.device), requires_grad=False)
+ return linear if model_device is None else linear.to(device=model_device)
+
+ def _active_format(self):
+ marker_name = self.config.format
+ if marker_name not in FORMATS:
+ raise ValueError(f"Unsupported format `{marker_name}`, not in {sorted(FORMATS)}.")
+ return FORMATS[marker_name]
+
+ def _build_linear(self, linear, compute_dtype):
+ quant_format = self._active_format()
+ return quant_format.linear_cls(
+ linear.in_features,
+ linear.out_features,
+ bias=linear.bias is not None,
+ layout=quant_format.layout_cls,
+ compute_dtype=compute_dtype,
+ )
+
+
+@dataclass
+class ComfyKitchenInt8Config(BackendConfig):
+ format: str = field(init=False, default=TensorWiseInt8Format.marker_name) # pinned by the method
+ per_channel: bool = True
+ convrot: bool = True
+ convrot_groupsize: int = TensorWiseInt8Format.DEFAULT_GROUPSIZE
+ orig_dtype: torch.dtype = torch.bfloat16
+
+
+@dataclass
+class ComfyKitchenFp8Config(BackendConfig):
+ format: str = field(init=False, default=Float8E4M3Format.marker_name) # pinned by the method
+ orig_dtype: torch.dtype = torch.bfloat16
+
+
+register_quant_method("comfy_kitchen_int8_w8a8", "comfy_kitchen", ComfyKitchenInt8Config.from_kwargs, label="W8A8, int8 weight + int8 dynamic activation (ComfyUI int8_tensorwise)")
+register_quant_method("comfy_kitchen_fp8_w8a8", "comfy_kitchen", ComfyKitchenFp8Config.from_kwargs, label="W8A8, fp8 E4M3 weight + fp8 activation, static input_scale when present (ComfyUI float8_e4m3fn)")
diff --git a/diffsynth/core/quant/backends/torchao.py b/diffsynth/core/quant/backends/torchao.py
new file mode 100644
index 0000000000000000000000000000000000000000..f05edcffaaa82c85432efa803e11732b5479feb8
--- /dev/null
+++ b/diffsynth/core/quant/backends/torchao.py
@@ -0,0 +1,162 @@
+import json
+
+import torch
+from ..base import QuantBackend, register_quant_backend
+from ..config import register_quant_method
+
+try:
+ from torchao.quantization import (
+ quantize_, Int8WeightOnlyConfig, Int4WeightOnlyConfig, Float8WeightOnlyConfig,
+ Int8DynamicActivationInt8WeightConfig, Float8DynamicActivationFloat8WeightConfig,
+ Float8DynamicActivationInt4WeightConfig,
+ )
+ from torchao.prototype.mx_formats.inference_workflow import (
+ MXDynamicActivationMXWeightConfig, NVFP4DynamicActivationNVFP4WeightConfig, NVFP4WeightOnlyConfig,
+ )
+ from torchao.prototype.safetensors.safetensors_support import flatten_tensor_state_dict, unflatten_tensor_state_dict
+ from torchao.prototype.safetensors.safetensors_utils import is_metadata_torchao
+ TORCHAO_AVAILABLE = True
+except ImportError:
+ TORCHAO_AVAILABLE = False
+
+
+class TorchaoLinear(torch.nn.Linear):
+ """Marker class for torchao-quantized Linears."""
+
+
+@register_quant_backend("torchao")
+class TorchaoQuantBackend(QuantBackend):
+ """Adapter over torchao `quantize_` (weight-only configs); the quantization lives in the weight tensor subclass, not the module class."""
+
+ project_url = "https://github.com/pytorch/ao"
+
+ def validate_environment(self):
+ if not TORCHAO_AVAILABLE:
+ raise ImportError(
+ "torchao is required for torchao quantization methods. "
+ 'Please install it via `pip install torchao` or `pip install "diffsynth[quant]"`.'
+ )
+
+ _DIFFERENTIABLE_CONFIGS = frozenset({
+ "Int8WeightOnlyConfig",
+ "Float8WeightOnlyConfig",
+ "NVFP4WeightOnlyConfig",
+ })
+
+ def capabilities(self):
+ config_name = type(self.config).__name__
+ return {
+ "is_serializable": True,
+ "is_differentiable": config_name in self._DIFFERENTIABLE_CONFIGS,
+ "is_compileable": True,
+ "requires_calibration": False,
+ }
+
+ def quantized_linear_classes(self):
+ return (TorchaoLinear,)
+
+ def create_quantized_linear(self, linear, compute_device=None, model_device=None):
+ linear.requires_grad_(False)
+ if compute_device is not None:
+ linear = linear.to(device=compute_device)
+ quant_linear = TorchaoLinear(linear.in_features, linear.out_features, bias=linear.bias is not None, device="meta")
+ quant_linear.weight = linear.weight
+ quant_linear.bias = linear.bias
+ quantize_(quant_linear, self.config)
+ if model_device is not None:
+ quant_linear = quant_linear.to(device=model_device)
+ return quant_linear
+
+ def create_quantized_linear_shell(self, linear, compute_dtype):
+ return TorchaoLinear(linear.in_features, linear.out_features, bias=linear.bias is not None, device="meta")
+
+ def flatten_state_dict(self, state_dict):
+ return flatten_tensor_state_dict(state_dict)
+
+ def unflatten_state_dict(self, state_dict, metadata):
+ if not is_metadata_torchao(metadata):
+ raise ValueError("This checkpoint carries no torchao metadata.")
+ tensor_names = json.loads(metadata["tensor_names"])
+ root_names = [name for name in tensor_names if "." not in name]
+ if root_names:
+ metadata = {**metadata, "tensor_names": json.dumps([name for name in tensor_names if "." in name])}
+ rebuilt, _ = unflatten_tensor_state_dict(state_dict, metadata)
+ for name in root_names:
+ if name in state_dict:
+ rebuilt[name] = state_dict[name]
+ return rebuilt
+
+ @staticmethod
+ def _to_plain_weight(weight, compute_dtype):
+ try:
+ return weight.dequantize().to(compute_dtype)
+ except (NotImplementedError, AttributeError):
+ eye = torch.eye(weight.shape[1], dtype=compute_dtype, device=weight.device)
+ return torch.nn.functional.linear(eye, weight).t().contiguous()
+
+ def dequantize_to_linear(self, module, compute_dtype, compute_device=None, model_device=None):
+ if compute_device is not None:
+ module = module.to(device=compute_device)
+ fp_weight = self._to_plain_weight(module.weight, compute_dtype)
+ linear = torch.nn.Linear(module.in_features, module.out_features, bias=module.bias is not None, device="meta")
+ linear.weight = torch.nn.Parameter(fp_weight, requires_grad=False)
+ if module.bias is not None:
+ linear.bias = torch.nn.Parameter(module.bias.data.to(dtype=compute_dtype, device=fp_weight.device), requires_grad=False)
+ return linear if model_device is None else linear.to(device=model_device)
+
+
+def _int8_weight_only(backend_config_kwargs):
+ backend_config_kwargs.setdefault("version", 2)
+ return Int8WeightOnlyConfig(**backend_config_kwargs)
+
+
+def _int4_weight_only(backend_config_kwargs):
+ return Int4WeightOnlyConfig(**backend_config_kwargs)
+
+
+def _fp8_weight_only(backend_config_kwargs):
+ return Float8WeightOnlyConfig(**backend_config_kwargs)
+
+
+def _int8_dynamic_activation_int8_weight(backend_config_kwargs):
+ backend_config_kwargs.setdefault("version", 2)
+ return Int8DynamicActivationInt8WeightConfig(**backend_config_kwargs)
+
+
+def _float8_dynamic_activation_float8_weight(backend_config_kwargs):
+ return Float8DynamicActivationFloat8WeightConfig(**backend_config_kwargs)
+
+
+def _float8_dynamic_activation_int4_weight(backend_config_kwargs):
+ return Float8DynamicActivationInt4WeightConfig(**backend_config_kwargs)
+
+
+def _mx_dynamic_activation_mx_weight(elem_dtype):
+ def factory(backend_config_kwargs):
+ dtype = getattr(torch, elem_dtype, None)
+ if dtype is None:
+ raise ImportError(f"This torch build has no `torch.{elem_dtype}`, required by the MX config.")
+ kwargs = {"block_size": 32, "activation_dtype": dtype, "weight_dtype": dtype}
+ kwargs.update(backend_config_kwargs)
+ return MXDynamicActivationMXWeightConfig(**kwargs)
+ return factory
+
+
+def _nvfp4_dynamic_activation_nvfp4_weight(backend_config_kwargs):
+ return NVFP4DynamicActivationNVFP4WeightConfig(**backend_config_kwargs)
+
+
+def _nvfp4_weight_only(backend_config_kwargs):
+ return NVFP4WeightOnlyConfig(**backend_config_kwargs)
+
+
+register_quant_method("torchao_int8_w8a16", "torchao", _int8_weight_only, label="8bit, int8, weight-only")
+register_quant_method("torchao_int4_w4a16", "torchao", _int4_weight_only, label="W4A16, int4 weight-only")
+register_quant_method("torchao_fp8_w8a16", "torchao", _fp8_weight_only, label="8bit, fp8, weight-only")
+register_quant_method("torchao_int8_w8a8", "torchao", _int8_dynamic_activation_int8_weight, label="W8A8, int8 weight + int8 dynamic activation")
+register_quant_method("torchao_fp8_w8a8", "torchao", _float8_dynamic_activation_float8_weight, label="W8A8, fp8 weight + fp8 dynamic activation")
+register_quant_method("torchao_int4_w4a8", "torchao", _float8_dynamic_activation_int4_weight, label="W4A8, int4 weight + fp8 dynamic activation")
+register_quant_method("torchao_mxfp8_w8a8", "torchao", _mx_dynamic_activation_mx_weight("float8_e4m3fn"), label="W8A8, MXFP8 microscaling weight + activation")
+register_quant_method("torchao_mxfp4_w4a4", "torchao", _mx_dynamic_activation_mx_weight("float4_e2m1fn_x2"), label="W4A4, MXFP4 microscaling weight + activation")
+register_quant_method("torchao_nvfp4_w4a4", "torchao", _nvfp4_dynamic_activation_nvfp4_weight, label="W4A4, NVFP4 weight + activation")
+register_quant_method("torchao_nvfp4_w4a16", "torchao", _nvfp4_weight_only, label="W4A16, NVFP4 weight-only")
diff --git a/diffsynth/core/quant/base.py b/diffsynth/core/quant/base.py
new file mode 100644
index 0000000000000000000000000000000000000000..dfd0b1398f472c785d836d46590d7b209595af3d
--- /dev/null
+++ b/diffsynth/core/quant/base.py
@@ -0,0 +1,243 @@
+from abc import ABC
+from dataclasses import dataclass, fields
+import torch
+
+
+QUANT_BACKENDS = {}
+
+
+@dataclass
+class BackendConfig:
+ """Base for a backend's typed config. User-tunable knobs are ordinary init fields;
+ method-pinned values use `field(init=False, ...)` so they cannot be overridden."""
+
+ @classmethod
+ def from_kwargs(cls, kwargs):
+ accepted = {f.name for f in fields(cls) if f.init}
+ unknown = set(kwargs) - accepted
+ if unknown:
+ raise ValueError(
+ f"`backend_config_kwargs` {sorted(unknown)} are not accepted by `{cls.__name__}`. Accepted: {sorted(accepted) or '(none)'}."
+ )
+ return cls(**kwargs)
+
+
+class QuantBackend(ABC):
+ """
+ Adapter between the framework and a quantization library (bnb / torchao / custom).
+ Subclasses are registered in `QUANT_BACKENDS` ({name -> class}) and instantiated by
+ `QuantizeConfig` with the method's backend config. A backend operates on single
+ Linears; model-level traversal and replacement is done by the `QuantizeConfig` methods.
+
+ The quantized Linear produced by a backend must satisfy:
+ (a) It is an `nn.Linear` drop-in: `forward(x)` internally performs dequant + matmul.
+ (b) `.to(...)` moves devices but never re-types the packed weight / quant state:
+ a dtype cast (`.to(dtype)` / `.half()` / `.float()`) must leave their storage
+ format and values intact.
+ (c) `state_dict()` / `load_state_dict(assign=True)` round-trips (optionally via
+ `flatten_state_dict` / `unflatten_state_dict`).
+ (d) (Training branch only) `forward` is differentiable w.r.t. its input, so
+ gradients can pass through frozen quantized layers to reach LoRA branches.
+ Declared statically by `capabilities()["is_differentiable"]` and verifiable
+ at runtime by calling `check_differentiable` on a quantized Linear the
+ backend produced.
+ """
+
+ name: str = ""
+ project_url: str = ""
+
+ def __init__(self, config=None):
+ self.config = config
+
+ def capabilities(self) -> dict:
+ return {
+ "is_serializable": False,
+ "is_differentiable": True,
+ "is_compileable": False,
+ "requires_calibration": False,
+ }
+
+ def validate_environment(self):
+ return
+
+ def announce_environment(self):
+ """Point at the library that owns this backend's hardware support, instead of restating it."""
+ if self.project_url:
+ print(f"Quantization backend `{self.name}`: for hardware compatibility issues see {self.project_url}")
+
+ def create_quantized_linear(self, linear: torch.nn.Linear, compute_device=None, model_device=None) -> torch.nn.Module:
+ raise NotImplementedError(
+ f"Backend `{self.name}` cannot quantize an fp model online. Use a method whose "
+ "backend supports it, or load an already quantized checkpoint."
+ )
+
+ def create_quantized_linear_shell(self, linear: torch.nn.Linear, compute_dtype: torch.dtype) -> torch.nn.Module:
+ raise NotImplementedError(
+ f"Backend `{self.name}` cannot load pre-quantized checkpoints. Use "
+ "`load_prequantized=False` to quantize an fp model online instead."
+ )
+
+ def dequantize_to_linear(self, module: torch.nn.Module, compute_dtype: torch.dtype, compute_device=None, model_device=None) -> torch.nn.Linear:
+ raise NotImplementedError(
+ f"Backend `{self.name}` cannot dequantize back to `nn.Linear`, so "
+ '`mode="dequant_once"` is unavailable.'
+ )
+
+ def quantized_linear_classes(self) -> tuple:
+ """The Linear classes this backend produces. MUST be subclass of `torch.nn.Linear`."""
+ raise NotImplementedError(
+ f"Backend `{self.name}` must declare the Linear classes it produces."
+ )
+
+ def is_quantized_linear(self, module) -> bool:
+ return isinstance(module, self.quantized_linear_classes())
+
+ def flatten_state_dict(self, state_dict: dict):
+ self._require_serializable()
+ return state_dict, {}
+
+ def unflatten_state_dict(self, state_dict: dict, metadata: dict):
+ self._require_serializable()
+ return state_dict
+
+ def _require_serializable(self):
+ if not self.capabilities().get("is_serializable", False):
+ raise NotImplementedError(
+ f"Backend `{self.name}` declares `is_serializable=False`, so its quantized state "
+ "dict cannot be flattened or rebuilt. Override these methods if it actually can."
+ )
+
+def register_quant_backend(name):
+ def decorator(cls):
+ cls.name = name
+ QUANT_BACKENDS[name] = cls
+ return cls
+ return decorator
+
+
+def check_differentiable(module: torch.nn.Module, example_input: torch.Tensor = None, verbose: bool = True) -> bool:
+ """
+ Check whether gradients pass through `module` w.r.t. its input: run a real
+ backward pass from the output (`torch.autograd.grad`) and verify a finite
+ gradient arrives at the input. This is what LoRA training requires from
+ frozen (e.g. quantized) layers. The module is cast to bfloat16 in place and
+ probed with a bfloat16 input; if `example_input` is None, a random one is
+ built for modules exposing `in_features`.
+
+ Example (probing a torchao-quantized Linear):
+
+ import torch
+ from diffsynth.core.quant import check_differentiable
+ from torchao.quantization import quantize_, Int8WeightOnlyConfig
+
+ linear = torch.nn.Linear(1024, 1024, dtype=torch.bfloat16, device="cuda")
+ quantize_(linear, Int8WeightOnlyConfig(version=2))
+ check_differentiable(linear)
+ """
+ def report(result, detail):
+ if verbose:
+ print(f"check_differentiable ({type(module).__name__}): {'OK' if result else 'FAIL'} -- {detail}")
+ return result
+
+ try:
+ module = module.to(torch.bfloat16)
+ if example_input is None:
+ if not hasattr(module, "in_features"):
+ raise ValueError("`example_input` is required for modules without `in_features`.")
+ device = next((t.device for t in list(module.parameters()) + list(module.buffers())), torch.device("cpu"))
+ example_input = torch.randn(4, module.in_features, device=device)
+ x = example_input.detach().to(torch.bfloat16).requires_grad_(True)
+ y = module(x)
+ if not y.requires_grad:
+ return report(False, "the output does not require grad, so no autograd graph was recorded")
+ input_grad = torch.autograd.grad(y, x, grad_outputs=torch.randn_like(y), allow_unused=True)[0]
+ except Exception as error:
+ return report(False, f"{type(error).__name__}: {error}")
+ if input_grad is None:
+ return report(False, "backward finished but no gradient reached the input")
+ if not torch.isfinite(input_grad.float()).all():
+ return report(False, "the input gradient contains non-finite values")
+ return report(True, "gradients pass through the module to its input")
+
+
+def check_backend_contract(backend, in_features: int = 512, out_features: int = 512,
+ compute_dtype: torch.dtype = torch.bfloat16,
+ compute_device: str = "cuda", verbose: bool = True) -> bool:
+ """
+ Verify a backend satisfies the quantized-Linear contract: it declares its classes,
+ both factory methods return instances of them, and every declared class subclasses
+ `torch.nn.Linear` so LoRA target detection and VRAM management can see it. The
+ checkpoint key patterns are checked against the keys the backend actually writes,
+ since a pattern list that misses a scale makes disk offload load corrupt layers
+ without raising.
+
+ Example:
+
+ from diffsynth.core.quant import QUANT_BACKENDS, QUANT_METHODS, check_backend_contract
+ spec = QUANT_METHODS["bitsandbytes_nf4"]
+ check_backend_contract(QUANT_BACKENDS[spec.backend](spec.config_factory({})))
+ """
+ failures = []
+ backend.validate_environment()
+
+ def check(condition, detail):
+ if not condition:
+ failures.append(detail)
+ if verbose:
+ print(f" [{'PASS' if condition else 'FAIL'}] {detail}")
+ return condition
+
+ if verbose:
+ print(f"check_backend_contract ({backend.name}):")
+ try:
+ classes = backend.quantized_linear_classes()
+ except NotImplementedError as error:
+ if verbose:
+ print(f" [FAIL] quantized_linear_classes() is not implemented: {error}")
+ return False
+ if not check(len(classes) > 0, f"quantized_linear_classes() is non-empty: {[cls.__name__ for cls in classes]}"):
+ return False
+ for cls in classes:
+ check(issubclass(cls, torch.nn.Linear), f"{cls.__name__} subclasses torch.nn.Linear")
+
+ plain = torch.nn.Linear(in_features, out_features, bias=True, dtype=compute_dtype, device=compute_device)
+ plain.requires_grad_(False)
+ check(not backend.is_quantized_linear(plain), "a plain nn.Linear is not reported as quantized")
+
+ try:
+ shell = backend.create_quantized_linear_shell(plain, compute_dtype)
+ except NotImplementedError:
+ shell = None
+ if verbose:
+ print(" [SKIP] create_quantized_linear_shell() is unsupported by this backend")
+ if shell is not None:
+ check(isinstance(shell, classes), f"create_quantized_linear_shell() returns a declared class, got {type(shell).__name__}")
+ check(backend.is_quantized_linear(shell), "the shell is recognized before load_state_dict (disk offload routing)")
+
+ try:
+ quantized = backend.create_quantized_linear(plain, compute_device=compute_device)
+ except NotImplementedError:
+ quantized = None
+ if verbose:
+ print(" [SKIP] create_quantized_linear() is unsupported by this backend")
+ if quantized is not None:
+ check(isinstance(quantized, classes), f"create_quantized_linear() returns a declared class, got {type(quantized).__name__}")
+
+ saved = quantized if quantized is not None else shell
+ if saved is None:
+ if verbose:
+ print(" [SKIP] neither factory method is supported, so the stored keys cannot be checked")
+ else:
+ state_dict = {f"proj.{key}": value for key, value in saved.state_dict().items()}
+ if backend.capabilities().get("is_serializable", False):
+ try:
+ state_dict = backend.flatten_state_dict(state_dict)[0]
+ except Exception as error:
+ if verbose:
+ print(f" [SKIP] flatten_state_dict() failed, falling back to the raw state dict keys: {type(error).__name__}: {error}")
+ uncovered = sorted(key for key in state_dict if not key.startswith("proj."))
+ check(len(uncovered) == 0, f"every stored key lives under the layer name; uncovered: {uncovered}")
+
+ if verbose:
+ print(f" => {'OK' if not failures else str(len(failures)) + ' FAILED'}")
+ return not failures
diff --git a/diffsynth/core/quant/config.py b/diffsynth/core/quant/config.py
new file mode 100644
index 0000000000000000000000000000000000000000..85ac865721b32d5456be009ba7e4e3974f86bbaf
--- /dev/null
+++ b/diffsynth/core/quant/config.py
@@ -0,0 +1,570 @@
+from dataclasses import dataclass, field, fields
+from typing import Any, Callable, Optional
+import torch
+from .base import QUANT_BACKENDS
+
+
+QUANT_METHODS = {}
+
+
+@dataclass
+class QuantMethodSpec:
+ backend: str
+ config_factory: Callable[[dict], Any]
+ label: str = ""
+
+
+def register_quant_method(name, backend, config_factory, label=""):
+ QUANT_METHODS[name] = QuantMethodSpec(backend=backend, config_factory=config_factory, label=label)
+
+
+def describe_quant_method(name):
+ """Print a method's backend, label, and the accepted `backend_config_kwargs` with defaults."""
+ from . import backends
+ backends.load_all_backends()
+ if name not in QUANT_METHODS:
+ raise ValueError(f"Unknown quantization method: {name}. Available methods:\n{_available_methods()}")
+ spec = QUANT_METHODS[name]
+ lines = [f"method: {name}", f"backend: {spec.backend}", f"detail: {spec.label}"]
+ try:
+ config = spec.config_factory({})
+ except Exception as error:
+ lines.append(f"backend config: could not be built with default kwargs: {error}")
+ print("\n".join(lines))
+ return
+ cls = type(config)
+ lines.append(f"backend config: {cls.__module__}.{cls.__qualname__}")
+ all_fields = fields(config)
+ user_fields = [f for f in all_fields if f.init]
+ pinned_fields = [f for f in all_fields if not f.init]
+ if user_fields:
+ lines.append("backend_config_kwargs (user-tunable):")
+ width = max(len(f.name) for f in user_fields)
+ lines += [f" {f.name:<{width}} = {getattr(config, f.name)!r}" for f in user_fields]
+ if pinned_fields:
+ lines.append("pinned by method (not overridable):")
+ lines += [f" {f.name} = {getattr(config, f.name)!r}" for f in pinned_fields]
+ print("\n".join(lines))
+
+
+def _available_methods():
+ from . import backends
+ backends.load_all_backends()
+ methods = sorted(QUANT_METHODS.items())
+ if len(methods) == 0:
+ return " (no method registered)"
+ name_width = max(len(name) for name, _ in methods)
+ backend_width = max(len(spec.backend) for _, spec in methods)
+ lines = [f"method: {name:<{name_width}} backend: {spec.backend:<{backend_width}} detail: {spec.label}"
+ for name, spec in methods]
+ width = max(len(line) for line in lines)
+ body = "\n".join(f"│ {line:<{width}} │" for line in lines)
+ return f"┌{'─' * (width + 2)}┐\n{body}\n└{'─' * (width + 2)}┘"
+
+
+def _format_error_report(reports):
+ if not reports:
+ return "Quantization error report: no targeted nn.Linear layers found."
+ methods = sorted({report["method"] for report in reports})
+ lines = [f"Quantization error report ({' + '.join(methods)}, {len(reports)} layers, sorted by relative_error, worst first):"]
+ mixed = len(methods) > 1
+ name_width = max(len(report["name"]) for report in reports)
+ header = f" {'relative_error':>14} {'max_abs':>12} {'shape':<16} {'name':<{name_width}}"
+ if mixed:
+ header += " method"
+ lines.append(header)
+ for report in reports:
+ line = (
+ f" {report['relative_error'] * 100:>13.2f}% {report['max_abs']:>12.3e}"
+ f" {str(report['shape']):<16} {report['name']:<{name_width}}"
+ )
+ if mixed:
+ line += f" {report['method']}"
+ lines.append(line)
+ return "\n".join(lines)
+
+
+@dataclass
+class QuantizeConfig:
+ """
+ Quantization config and entry points for any `nn.Module`.
+
+ Online (dynamic) quantization -- quantize an fp model whose weights are loaded:
+
+ cfg = QuantizeConfig(method="bitsandbytes_nf4")
+ model.load_state_dict(fp_state_dict)
+ cfg.quantize_model(model, compute_device="cuda")
+
+ Loading a pre-quantized checkpoint -- swap in quantized shells before loading:
+
+ cfg = QuantizeConfig(method="bitsandbytes_nf4", load_prequantized=True)
+ cfg.prepare_for_prequantized_load(model, compute_dtype=torch.bfloat16)
+ state_dict = cfg.unflatten_state_dict(state_dict, metadata)
+ model.load_state_dict(state_dict, assign=True)
+
+ Dequantize-once -- after either flow above, back to plain fp Linears that carry
+ the quantization error:
+
+ cfg.dequantize_model(model, compute_dtype=torch.bfloat16)
+
+ Fields:
+ method: name from QUANT_METHODS, determines backend + scheme + backend config.
+ mode: "dynamic" keeps the backend-native quantized Linear, which dequantizes
+ its weight on every forward; "dequant_once" dequantizes to a plain fp
+ `nn.Linear` right after the weights are quantized or loaded.
+ target_modules / exclude_modules: list of module names; a layer is matched if
+ its full dotted name equals an entry, or ends with "." + entry
+ (e.g. "img_mod.1" matches "transformer_blocks.0.img_mod.1").
+ backend_config_kwargs: passed through to the method's backend config factory
+ (e.g. nf4 blocksize); `describe_quant_method(name)` prints what a method
+ accepts.
+ load_prequantized: the checkpoint already holds quantized weights; load them
+ directly instead of quantizing fp weights online.
+ """
+ method: str = None
+ mode: str = "dynamic"
+ target_modules: Optional[list] = None
+ exclude_modules: Optional[list] = None
+ backend_config_kwargs: dict = field(default_factory=dict)
+ load_prequantized: bool = False
+
+ def __post_init__(self):
+ if self.method is None:
+ raise ValueError(f"`QuantizeConfig.method` is required. Available methods:\n{_available_methods()}")
+ if self.mode not in ("dynamic", "dequant_once"):
+ raise ValueError(f"`QuantizeConfig.mode` should be `dynamic` or `dequant_once`, but got `{self.mode}`.")
+ for field_name in ("target_modules", "exclude_modules"):
+ value = getattr(self, field_name)
+ if value is not None and not isinstance(value, list):
+ raise ValueError(f"`QuantizeConfig.{field_name}` should be a list of module names, but got `{type(value).__name__}`.")
+ self.backend = self._build_backend()
+
+ def _build_backend(self):
+ from . import backends
+ backends.load_backend_for_method(self.method)
+ if self.method not in QUANT_METHODS:
+ raise ValueError(f"Unknown quantization method: {self.method}. Available methods:\n{_available_methods()}")
+ spec = QUANT_METHODS[self.method]
+ if spec.backend not in QUANT_BACKENDS:
+ raise ValueError(f"Quantization backend `{spec.backend}` (required by method `{self.method}`) is not registered.")
+ backend = QUANT_BACKENDS[spec.backend]()
+ backend.validate_environment()
+ backend.announce_environment()
+ backend.config = spec.config_factory(dict(self.backend_config_kwargs))
+ return backend
+
+ def quantize_model(self, model: torch.nn.Module, compute_device=None, model_device=None):
+ """
+ Quantize the targeted `nn.Linear` layers of an fp model in place, keeping each
+ layer's existing dtype. Call AFTER `load_state_dict`. Does nothing when
+ `load_prequantized` is set, since such a checkpoint is already quantized.
+
+ Parameters:
+ model: the fp model whose targeted layers are replaced by quantized ones.
+ compute_device: where quantization runs; `None` quantizes each layer in place.
+ model_device: where each quantized layer is stored afterwards; `None` leaves it
+ on `compute_device`. Combining an fp model on the CPU with
+ `compute_device="cuda", model_device="cpu"` streams the work layer by
+ layer, so the accelerator only ever holds one layer at a time.
+ """
+ if self.load_prequantized:
+ return model
+
+ def quantize(linear):
+ return self.backend.create_quantized_linear(linear, compute_device=compute_device, model_device=model_device)
+
+ replaced = self._replace_target_linears(model, quantize)
+ print(f"{len(replaced)} nn.Linear layers quantized (method: {self.method}).")
+ return model
+
+ def dequantize_model(self, model: torch.nn.Module, compute_dtype: torch.dtype = torch.bfloat16, compute_device=None, model_device=None):
+ """
+ Replace every quantized Linear in the model by a plain fp `nn.Linear`.
+ Does nothing unless `mode` is `"dequant_once"`.
+
+ Parameters:
+ model: the quantized model to restore in place.
+ compute_dtype: dtype of the restored fp weights.
+ compute_device: where dequantization runs; `None` restores each layer in place.
+ model_device: where each restored layer is stored afterwards; `None` leaves it
+ on `compute_device`.
+ """
+ if self.mode != "dequant_once":
+ return model
+ self._dequantize_linears(model, compute_dtype, compute_device, model_device)
+ return model
+
+ def is_quantized_linear(self, module) -> bool:
+ """Whether `module` is one of this config's backend-native quantized Linears."""
+ return self.backend.is_quantized_linear(module)
+
+ def build_quantized_shell(self, module, compute_dtype: torch.dtype, **kwargs):
+ """
+ Build an empty quantized Linear matching `module`, used to release a layer's
+ weights while keeping it routable, and to stage a transient copy on the
+ computation device.
+
+ Parameters:
+ module: the quantized layer whose shape and bias presence are mirrored.
+ compute_dtype: dtype the shell dequantizes to at forward time.
+ """
+ return self.backend.create_quantized_linear_shell(module, compute_dtype)
+
+ def prepare_for_prequantized_load(self, model: torch.nn.Module, compute_dtype: torch.dtype = torch.bfloat16):
+ """
+ Replace the targeted `nn.Linear` layers with empty quantized Linears matching a
+ pre-quantized checkpoint. Call BEFORE `load_state_dict(assign=True)`.
+
+ Parameters:
+ model: the freshly constructed model whose targeted layers become shells.
+ compute_dtype: dtype the quantized layers dequantize to at forward time.
+ """
+ def build_shell(linear):
+ return self.backend.create_quantized_linear_shell(linear, compute_dtype)
+
+ replaced = self._replace_target_linears(model, build_shell)
+ print(f"{len(replaced)} nn.Linear layers replaced for loading the pre-quantized checkpoint (method: {self.method}).")
+ return model
+
+ def unflatten_state_dict(self, state_dict: dict, metadata: dict):
+ """
+ Rebuild composite quantized tensors from a flat (safetensors) state dict, so the
+ result can be given to `load_state_dict(assign=True)`.
+
+ Parameters:
+ state_dict: the flat tensors read from the checkpoint.
+ metadata: the string-only metadata stored alongside them.
+ """
+ return self.backend.unflatten_state_dict(state_dict, metadata)
+
+ def flatten_state_dict(self, state_dict: dict):
+ """
+ Flatten a quantized model's state dict into plain tensors and string-only metadata.
+ Returns `(tensors, metadata)`, ready for
+ `safetensors.torch.save_file(tensors, path, metadata=metadata)`.
+
+ Parameters:
+ state_dict: the quantized model's state dict, holding composite tensors.
+ """
+ if not self.backend.capabilities().get("is_serializable", False):
+ raise NotImplementedError(
+ f"Backend `{self.backend.name}` (method `{self.method}`) does not declare "
+ 'serialization support (`capabilities()["is_serializable"]`), so its quantized '
+ "state dict cannot be flattened for saving."
+ )
+ tensors, metadata = self.backend.flatten_state_dict(state_dict)
+ tensors = {key: value.contiguous() for key, value in tensors.items()}
+ metadata = {"format": "pt", **{key: value if isinstance(value, str) else str(value) for key, value in metadata.items()}}
+ return tensors, metadata
+
+ def measure_quantization_error(self, model: torch.nn.Module, compute_device="cuda", verbose=True) -> list:
+ """
+ Measure the weight error this config's quantization introduces.
+
+ Parameters:
+ model: the fp model to analyze.
+ compute_device: where the quantize-dequantize round trip runs.
+ verbose: print the per-layer report, sorted by `relative_error`, worst first.
+
+ Returns:
+ A list of per-layer dicts: `name`, `method`, `shape`, `relative_error` and `max_abs` (worst-element absolute error).
+ """
+ reports = [
+ self._measure_layer_error(full_name, module, compute_device)
+ for full_name, module in model.named_modules()
+ if full_name != "" and self._should_quantize(full_name, module)
+ ]
+ reports.sort(key=lambda report: report["relative_error"], reverse=True)
+ if verbose:
+ print(_format_error_report(reports))
+ return reports
+
+ def _measure_layer_error(self, full_name, module, compute_device):
+ fp_weight = module.weight.detach().to(device=compute_device, dtype=torch.float32)
+ linear = torch.nn.Linear(
+ module.in_features, module.out_features,
+ bias=module.bias is not None, dtype=module.weight.dtype, device=compute_device,
+ )
+ linear.weight.data.copy_(module.weight.detach())
+ if module.bias is not None:
+ linear.bias.data.copy_(module.bias.detach().to(device=compute_device))
+ quantized = self.backend.create_quantized_linear(linear, compute_device=compute_device, model_device=compute_device)
+ restored = self.backend.dequantize_to_linear(quantized, compute_dtype=torch.float32, compute_device=compute_device, model_device=compute_device)
+ quant_weight = restored.weight.detach()
+ error = quant_weight - fp_weight
+ report = {
+ "name": full_name,
+ "method": self.method,
+ "shape": tuple(fp_weight.shape),
+ "relative_error": (torch.linalg.vector_norm(error) / torch.linalg.vector_norm(fp_weight)).item(),
+ "max_abs": error.abs().max().item(),
+ }
+ del fp_weight, linear, quantized, restored, quant_weight, error
+ return report
+
+ @staticmethod
+ def _name_matches(full_name, patterns):
+ if patterns is None:
+ return False
+ if full_name in patterns:
+ return True
+ return any(full_name.endswith(f".{pattern}") for pattern in patterns)
+
+ def _should_quantize(self, full_name, module):
+ if not isinstance(module, torch.nn.Linear):
+ return False
+ if self.is_quantized_linear(module):
+ return False
+ if self.target_modules is not None and not self._name_matches(full_name, self.target_modules):
+ return False
+ if self._name_matches(full_name, self.exclude_modules):
+ return False
+ return True
+
+ def _replace_target_linears(self, model, transform):
+ replaced = [
+ full_name
+ for full_name, module in model.named_modules()
+ if full_name != "" and self._should_quantize(full_name, module)
+ ]
+ for full_name in replaced:
+ parent_name, _, leaf_name = full_name.rpartition(".")
+ parent = model.get_submodule(parent_name) if parent_name else model
+ setattr(parent, leaf_name, transform(getattr(parent, leaf_name)))
+ return replaced
+
+ def _dequantize_linears(self, model, compute_dtype, compute_device=None, model_device=None):
+ restored = [
+ full_name
+ for full_name, module in model.named_modules()
+ if full_name != "" and self.backend.is_quantized_linear(module)
+ ]
+ for full_name in restored:
+ parent_name, _, leaf_name = full_name.rpartition(".")
+ parent = model.get_submodule(parent_name) if parent_name else model
+ restored_linear = self.backend.dequantize_to_linear(
+ getattr(parent, leaf_name), compute_dtype, compute_device=compute_device, model_device=model_device
+ )
+ setattr(parent, leaf_name, restored_linear)
+ return restored
+
+
+@dataclass
+class MixedQuantizeConfig:
+ """
+ Combines multiple `QuantizeConfig` passes, one per disjoint layer set, behind the
+ same interface as a single `QuantizeConfig`. Mixed quantization assigns different
+ methods to different layers, e.g. nf4 for attention/MLP and int8 for the
+ quantization-sensitive modulation layers:
+
+ mod_layers = ["img_mod.1", "txt_mod.1", ...]
+ cfg = MixedQuantizeConfig(configs=[
+ QuantizeConfig(method="bitsandbytes_nf4", exclude_modules=mod_layers),
+ QuantizeConfig(method="torchao_int8_w8a16", target_modules=mod_layers),
+ ])
+ cfg.quantize_model(model, compute_device="cuda")
+
+ The passes run in `configs` order. Their matched layer sets must be pairwise
+ disjoint; `quantize_model` and `prepare_for_prequantized_load` verify this on the
+ fp model and raise before touching it.
+
+ Fields:
+ configs: list of `QuantizeConfig`, one per layer set; all must share the same
+ `mode`, and their `load_prequantized` must stay False.
+ load_prequantized: set on this wrapper (not on the children) to load a mixed
+ pre-quantized checkpoint.
+ """
+ configs: list = None
+ load_prequantized: bool = False
+
+ def __post_init__(self):
+ if not self.configs:
+ raise ValueError("`MixedQuantizeConfig.configs` requires at least one `QuantizeConfig`.")
+ for config in self.configs:
+ if not isinstance(config, QuantizeConfig):
+ raise ValueError(f"`MixedQuantizeConfig.configs` should contain `QuantizeConfig` instances, but got `{type(config).__name__}`.")
+ if config.load_prequantized:
+ raise ValueError(
+ "Set `load_prequantized=True` on `MixedQuantizeConfig` itself, not on its "
+ "child configs; the children only describe method + layer set."
+ )
+ modes = {config.mode for config in self.configs}
+ if len(modes) > 1:
+ raise ValueError(f"All configs in `MixedQuantizeConfig` should share the same `mode`, but got {sorted(modes)}.")
+ self._ownership = {}
+
+ @property
+ def method(self):
+ return " + ".join(config.method for config in self.configs)
+
+ @property
+ def mode(self):
+ return self.configs[0].mode
+
+ def quantize_model(self, model: torch.nn.Module, compute_device=None, model_device=None):
+ """
+ Run every config's quantization pass in order, after verifying their layer sets
+ are disjoint. Does nothing when `load_prequantized` is set.
+
+ Parameters:
+ model: the fp model whose targeted layers are replaced by quantized ones.
+ compute_device: where quantization runs; `None` quantizes each layer in place.
+ model_device: where each quantized layer is stored afterwards; `None` leaves it
+ on `compute_device`.
+ """
+ if self.load_prequantized:
+ return model
+ self._build_ownership(model)
+ for config in self.configs:
+ config.quantize_model(model, compute_device=compute_device, model_device=model_device)
+ return model
+
+ def dequantize_model(self, model: torch.nn.Module, compute_dtype: torch.dtype = torch.bfloat16, compute_device=None, model_device=None):
+ """
+ Replace every quantized Linear of every config's backend by a plain fp `nn.Linear`.
+ Does nothing unless `mode` is `"dequant_once"`.
+
+ Parameters:
+ model: the quantized model to restore in place.
+ compute_dtype: dtype of the restored fp weights.
+ compute_device: where dequantization runs; `None` restores each layer in place.
+ model_device: where each restored layer is stored afterwards; `None` leaves it
+ on `compute_device`.
+ """
+ if self.mode != "dequant_once":
+ return model
+ for config in self.configs:
+ config.dequantize_model(model, compute_dtype=compute_dtype, compute_device=compute_device, model_device=model_device)
+ return model
+
+ def is_quantized_linear(self, module) -> bool:
+ """Whether `module` is a quantized Linear of any config's backend."""
+ return any(config.is_quantized_linear(module) for config in self.configs)
+
+ def build_quantized_shell(self, module, compute_dtype: torch.dtype, layer_name: str = None):
+ """
+ Build an empty quantized Linear for `module` using the config that owns it.
+
+ Parameters:
+ module: the quantized layer whose shape and bias presence are mirrored.
+ compute_dtype: dtype the shell dequantizes to at forward time.
+ layer_name: its dotted name inside the model; required to pick the right config
+ when several configs share one backend, since their quantized Linears are
+ then the same class.
+ """
+ return self._owning_config(module, layer_name).build_quantized_shell(module, compute_dtype)
+
+ def _owning_config(self, module, layer_name=None):
+ config = self._ownership.get(layer_name)
+ if config is not None:
+ return config
+ for config in self.configs:
+ if config.is_quantized_linear(module):
+ return config
+ raise ValueError(
+ f"`{type(module).__name__}` is not a quantized Linear of any config in this "
+ f"`MixedQuantizeConfig` (methods: {self.method})."
+ )
+
+ def prepare_for_prequantized_load(self, model: torch.nn.Module, compute_dtype: torch.dtype = torch.bfloat16):
+ """
+ Replace each config's targeted Linears by that backend's shells, after verifying
+ the layer sets are disjoint.
+
+ Parameters:
+ model: the freshly constructed model whose targeted layers become shells.
+ compute_dtype: dtype the quantized layers dequantize to at forward time.
+ """
+ self._build_ownership(model)
+ for config in self.configs:
+ config.prepare_for_prequantized_load(model, compute_dtype=compute_dtype)
+ return model
+
+ def unflatten_state_dict(self, state_dict: dict, metadata: dict):
+ """
+ Rebuild composite quantized tensors from a flat (safetensors) state dict.
+ """
+ if not self._ownership:
+ for config in self._distinct_backend_configs():
+ state_dict = config.unflatten_state_dict(state_dict, metadata)
+ return state_dict
+ rebuilt = {}
+ for config, keys in self._grouped_by_backend(state_dict):
+ rebuilt.update(config.unflatten_state_dict(keys, metadata) if config else keys)
+ return rebuilt
+
+ def flatten_state_dict(self, state_dict: dict):
+ """
+ Flatten the mixed model's state dict into plain tensors and string-only metadata.
+ Returns (state_dict, metadata), ready for `safetensors.torch.save_file(tensors,
+ path, metadata=metadata)`.
+ """
+ for config in self.configs:
+ if not config.backend.capabilities().get("is_serializable", False):
+ raise NotImplementedError(
+ f"Backend `{config.backend.name}` (method `{config.method}`) does not declare "
+ 'serialization support (`capabilities()["is_serializable"]`), so the mixed '
+ "state dict cannot be flattened for saving."
+ )
+ merged_metadata = {}
+ if not self._ownership:
+ flattened = state_dict
+ for config in self._distinct_backend_configs():
+ flattened, metadata = config.backend.flatten_state_dict(flattened)
+ merged_metadata.update(metadata)
+ else:
+ flattened = {}
+ for config, keys in self._grouped_by_backend(state_dict):
+ if config is None:
+ flattened.update(keys)
+ continue
+ tensors, metadata = config.backend.flatten_state_dict(keys)
+ flattened.update(tensors)
+ merged_metadata.update(metadata)
+ tensors = {key: value.contiguous() for key, value in flattened.items()}
+ metadata = {"format": "pt", **{key: value if isinstance(value, str) else str(value) for key, value in merged_metadata.items()}}
+ return tensors, metadata
+
+ def measure_quantization_error(self, model: torch.nn.Module, compute_device="cuda", verbose=True) -> list:
+ """
+ Measure the weight error of the mixed quantization.
+ """
+ reports = []
+ for config in self.configs:
+ reports.extend(config.measure_quantization_error(model, compute_device=compute_device, verbose=False))
+ reports.sort(key=lambda report: report["relative_error"], reverse=True)
+ if verbose:
+ print(_format_error_report(reports))
+ return reports
+
+ def _distinct_backend_configs(self):
+ seen = set()
+ distinct = []
+ for config in self.configs:
+ if config.backend.name not in seen:
+ seen.add(config.backend.name)
+ distinct.append(config)
+ return distinct
+
+ def _grouped_by_backend(self, state_dict):
+ groups = {}
+ for key, value in state_dict.items():
+ parts = key.split(".")
+ config = next((owner for end in range(len(parts) - 1, 0, -1)
+ if (owner := self._ownership.get(".".join(parts[:end]))) is not None), None)
+ groups.setdefault(config.backend.name if config else None, (config, {}))[1][key] = value
+ return groups.values()
+
+ def _build_ownership(self, model):
+ ownership = {}
+ for index, config in enumerate(self.configs):
+ names = [name for name, module in model.named_modules() if name and config._should_quantize(name, module)]
+ overlap = sorted(set(names) & ownership.keys())
+ if overlap:
+ owner = self.configs.index(ownership[overlap[0]])
+ raise ValueError(
+ f"Configs {owner} (`{self.configs[owner].method}`) and {index} (`{config.method}`) "
+ f"in `MixedQuantizeConfig` both match {len(overlap)} layers (e.g. {', '.join(overlap[:5])}). "
+ )
+ ownership.update({name: config for name in names})
+ self._ownership.update(ownership)
diff --git a/diffsynth/core/vram/__init__.py b/diffsynth/core/vram/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..32763bb9b4abfa4d5b2617827661c520d7e9fcae
--- /dev/null
+++ b/diffsynth/core/vram/__init__.py
@@ -0,0 +1,2 @@
+from .initialization import skip_model_initialization
+from .layers import *
diff --git a/diffsynth/core/vram/disk_map.py b/diffsynth/core/vram/disk_map.py
new file mode 100644
index 0000000000000000000000000000000000000000..70730061f8797906f8728722478f392b31478417
--- /dev/null
+++ b/diffsynth/core/vram/disk_map.py
@@ -0,0 +1,93 @@
+from safetensors import safe_open
+import torch, os
+
+
+class SafetensorsCompatibleTensor:
+ def __init__(self, tensor):
+ self.tensor = tensor
+
+ def get_shape(self):
+ return list(self.tensor.shape)
+
+
+class SafetensorsCompatibleBinaryLoader:
+ def __init__(self, path, device):
+ print("Detected non-safetensors files, which may cause slower loading. It's recommended to convert it to a safetensors file.")
+ self.state_dict = torch.load(path, weights_only=True, map_location=device)
+
+ def keys(self):
+ return self.state_dict.keys()
+
+ def get_tensor(self, name):
+ return self.state_dict[name]
+
+ def get_slice(self, name):
+ return SafetensorsCompatibleTensor(self.state_dict[name])
+
+
+class DiskMap:
+
+ def __init__(self, path, device, torch_dtype=None, state_dict_converter=None, buffer_size=10**9):
+ self.path = path if isinstance(path, list) else [path]
+ self.device = device
+ self.torch_dtype = torch_dtype
+ if os.environ.get('DIFFSYNTH_DISK_MAP_BUFFER_SIZE') is not None:
+ self.buffer_size = int(os.environ.get('DIFFSYNTH_DISK_MAP_BUFFER_SIZE'))
+ else:
+ self.buffer_size = buffer_size
+ self.files = []
+ self.flush_files()
+ self.name_map = {}
+ for file_id, file in enumerate(self.files):
+ for name in file.keys():
+ self.name_map[name] = file_id
+ self.rename_dict = self.fetch_rename_dict(state_dict_converter)
+
+ def flush_files(self):
+ if len(self.files) == 0:
+ for path in self.path:
+ if path.endswith(".safetensors"):
+ self.files.append(safe_open(path, framework="pt", device=str(self.device)))
+ else:
+ self.files.append(SafetensorsCompatibleBinaryLoader(path, device=self.device))
+ else:
+ for i, path in enumerate(self.path):
+ if path.endswith(".safetensors"):
+ self.files[i] = safe_open(path, framework="pt", device=str(self.device))
+ self.num_params = 0
+
+ def __getitem__(self, name):
+ if self.rename_dict is not None: name = self.rename_dict[name]
+ file_id = self.name_map[name]
+ param = self.files[file_id].get_tensor(name)
+ if self.torch_dtype is not None and isinstance(param, torch.Tensor):
+ param = param.to(self.torch_dtype)
+ if isinstance(param, torch.Tensor) and param.device.type == "cpu":
+ param = param.clone()
+ if isinstance(param, torch.Tensor):
+ self.num_params += param.numel()
+ if self.num_params > self.buffer_size:
+ self.flush_files()
+ return param
+
+ def fetch_rename_dict(self, state_dict_converter):
+ if state_dict_converter is None:
+ return None
+ state_dict = {}
+ for file in self.files:
+ for name in file.keys():
+ state_dict[name] = name
+ state_dict = state_dict_converter(state_dict)
+ return state_dict
+
+ def __iter__(self):
+ if self.rename_dict is not None:
+ return self.rename_dict.__iter__()
+ else:
+ return self.name_map.__iter__()
+
+ def __contains__(self, x):
+ if self.rename_dict is not None:
+ return x in self.rename_dict
+ else:
+ return x in self.name_map
diff --git a/diffsynth/core/vram/initialization.py b/diffsynth/core/vram/initialization.py
new file mode 100644
index 0000000000000000000000000000000000000000..bff2498b526638bfdd1c114c78aa0b98c251a47d
--- /dev/null
+++ b/diffsynth/core/vram/initialization.py
@@ -0,0 +1,21 @@
+import torch
+from contextlib import contextmanager
+
+
+@contextmanager
+def skip_model_initialization(device=torch.device("meta")):
+
+ def register_empty_parameter(module, name, param):
+ old_register_parameter(module, name, param)
+ if param is not None:
+ param_cls = type(module._parameters[name])
+ kwargs = module._parameters[name].__dict__
+ kwargs["requires_grad"] = param.requires_grad
+ module._parameters[name] = param_cls(module._parameters[name].to(device), **kwargs)
+
+ old_register_parameter = torch.nn.Module.register_parameter
+ torch.nn.Module.register_parameter = register_empty_parameter
+ try:
+ yield
+ finally:
+ torch.nn.Module.register_parameter = old_register_parameter
diff --git a/diffsynth/core/vram/layers.py b/diffsynth/core/vram/layers.py
new file mode 100644
index 0000000000000000000000000000000000000000..716d1d09b0eca7998e2de07c7372565757632141
--- /dev/null
+++ b/diffsynth/core/vram/layers.py
@@ -0,0 +1,620 @@
+import torch, copy
+from typing import Union
+from .initialization import skip_model_initialization
+from .disk_map import DiskMap
+from ..device import parse_device_type, get_device_name, IS_NPU_AVAILABLE
+
+
+class AutoTorchModule(torch.nn.Module):
+
+ def __init__(
+ self,
+ offload_dtype: torch.dtype = None,
+ offload_device: Union[str, torch.device] = None,
+ onload_dtype: torch.dtype = None,
+ onload_device: Union[str, torch.device] = None,
+ preparing_dtype: torch.dtype = None,
+ preparing_device: Union[str, torch.device] = None,
+ computation_dtype: torch.dtype = None,
+ computation_device: Union[str, torch.device] = None,
+ vram_limit: float = None,
+ ):
+ super().__init__()
+ self.set_dtype_and_device(
+ offload_dtype,
+ offload_device,
+ onload_dtype,
+ onload_device,
+ preparing_dtype,
+ preparing_device,
+ computation_dtype,
+ computation_device,
+ vram_limit,
+ )
+ self.state = 0
+ self.name = ""
+ self.computation_device_type = parse_device_type(self.computation_device)
+
+ def set_dtype_and_device(
+ self,
+ offload_dtype: torch.dtype = None,
+ offload_device: Union[str, torch.device] = None,
+ onload_dtype: torch.dtype = None,
+ onload_device: Union[str, torch.device] = None,
+ preparing_dtype: torch.dtype = None,
+ preparing_device: Union[str, torch.device] = None,
+ computation_dtype: torch.dtype = None,
+ computation_device: Union[str, torch.device] = None,
+ vram_limit: float = None,
+ ):
+ self.offload_dtype = offload_dtype or computation_dtype
+ self.offload_device = offload_device or ("disk" if self.offload_dtype == "disk" else computation_device)
+ self.onload_dtype = onload_dtype or computation_dtype
+ self.onload_device = onload_device or ("disk" if self.onload_dtype == "disk" else computation_device)
+ self.preparing_dtype = preparing_dtype or computation_dtype
+ self.preparing_device = preparing_device or ("disk" if self.preparing_dtype == "disk" else computation_device)
+ self.computation_dtype = computation_dtype
+ self.computation_device = computation_device
+ self.vram_limit = vram_limit
+
+ def cast_to(self, weight, dtype, device):
+ r = torch.empty_like(weight, dtype=dtype, device=device)
+ r.copy_(weight)
+ return r
+
+ def check_free_vram(self):
+ if self.computation_device_type == "mps":
+ allocated = torch.mps.current_allocated_memory()
+ recommended_max = torch.mps.recommended_max_memory()
+ gpu_mem_state = (recommended_max - allocated, recommended_max)
+ else:
+ device = self.computation_device if not IS_NPU_AVAILABLE else get_device_name()
+ gpu_mem_state = getattr(torch, self.computation_device_type).mem_get_info(device)
+ used_memory = (gpu_mem_state[1] - gpu_mem_state[0]) / (1024**3)
+ return used_memory < self.vram_limit
+
+ def offload(self):
+ if self.state != 0:
+ self.to(dtype=self.offload_dtype, device=self.offload_device)
+ self.state = 0
+
+ def onload(self):
+ if self.state != 1:
+ self.to(dtype=self.onload_dtype, device=self.onload_device)
+ self.state = 1
+
+ def param_name(self, name):
+ if self.name == "":
+ return name
+ else:
+ return self.name + "." + name
+
+
+class AutoWrappedModule(AutoTorchModule):
+
+ def __init__(
+ self,
+ module: torch.nn.Module,
+ offload_dtype: torch.dtype = None,
+ offload_device: Union[str, torch.device] = None,
+ onload_dtype: torch.dtype = None,
+ onload_device: Union[str, torch.device] = None,
+ preparing_dtype: torch.dtype = None,
+ preparing_device: Union[str, torch.device] = None,
+ computation_dtype: torch.dtype = None,
+ computation_device: Union[str, torch.device] = None,
+ vram_limit: float = None,
+ name: str = "",
+ disk_map: DiskMap = None,
+ **kwargs
+ ):
+ super().__init__(
+ offload_dtype,
+ offload_device,
+ onload_dtype,
+ onload_device,
+ preparing_dtype,
+ preparing_device,
+ computation_dtype,
+ computation_device,
+ vram_limit,
+ )
+ self.module = module
+ if offload_dtype == "disk":
+ self.name = name
+ self.disk_map = disk_map
+ self.required_params = [name for name, _ in self.module.named_parameters()]
+ self.disk_offload = True
+ else:
+ self.disk_offload = False
+
+ def load_from_disk(self, torch_dtype, device, copy_module=False):
+ if copy_module:
+ module = copy.deepcopy(self.module)
+ else:
+ module = self.module
+ state_dict = {}
+ for name in self.required_params:
+ param = self.disk_map[self.param_name(name)]
+ param = param.to(dtype=torch_dtype, device=device)
+ state_dict[name] = param
+ module.load_state_dict(state_dict, assign=True)
+ module.to(dtype=torch_dtype, device=device)
+ return module
+
+ def offload_to_disk(self, model: torch.nn.Module):
+ for buf in model.buffers():
+ # If there are some parameters are registed in buffers (not in state dict),
+ # We cannot offload the model.
+ for children in model.children():
+ self.offload_to_disk(children)
+ break
+ else:
+ model.to("meta")
+
+ def offload(self):
+ # offload / onload / preparing -> offload
+ if self.state != 0:
+ if self.disk_offload:
+ self.offload_to_disk(self.module)
+ else:
+ self.to(dtype=self.offload_dtype, device=self.offload_device)
+ self.state = 0
+
+ def onload(self):
+ # offload / onload / preparing -> onload
+ if self.state < 1:
+ if self.disk_offload and self.onload_device != "disk" and self.offload_device == "disk":
+ self.load_from_disk(self.onload_dtype, self.onload_device)
+ elif self.onload_device != "disk":
+ self.to(dtype=self.onload_dtype, device=self.onload_device)
+ self.state = 1
+
+ def preparing(self):
+ # onload / preparing -> preparing
+ if self.state != 2:
+ if self.disk_offload and self.preparing_device != "disk" and self.onload_device == "disk":
+ self.load_from_disk(self.preparing_dtype, self.preparing_device)
+ elif self.preparing_device != "disk":
+ self.to(dtype=self.preparing_dtype, device=self.preparing_device)
+ self.state = 2
+
+ def cast_to(self, module, dtype, device):
+ return copy.deepcopy(module).to(dtype=dtype, device=device)
+
+ def computation(self):
+ # onload / preparing -> computation (temporary)
+ if self.state == 2:
+ torch_dtype, device = self.preparing_dtype, self.preparing_device
+ else:
+ torch_dtype, device = self.onload_dtype, self.onload_device
+ if torch_dtype == self.computation_dtype and device == self.computation_device:
+ module = self.module
+ elif self.disk_offload and device == "disk":
+ module = self.load_from_disk(self.computation_dtype, self.computation_device, copy_module=True)
+ else:
+ module = self.cast_to(self.module, dtype=self.computation_dtype, device=self.computation_device)
+ return module
+
+ def forward(self, *args, **kwargs):
+ if self.state == 1 and (self.vram_limit is None or self.check_free_vram()):
+ self.preparing()
+ module = self.computation()
+ return module(*args, **kwargs)
+
+ def __getattr__(self, name):
+ if name in self.__dict__ or name == "module":
+ return super().__getattr__(name)
+ else:
+ return getattr(self.module, name)
+
+
+class AutoWrappedNonRecurseModule(AutoWrappedModule):
+
+ def __init__(
+ self,
+ module: torch.nn.Module,
+ offload_dtype: torch.dtype = None,
+ offload_device: Union[str, torch.device] = None,
+ onload_dtype: torch.dtype = None,
+ onload_device: Union[str, torch.device] = None,
+ preparing_dtype: torch.dtype = None,
+ preparing_device: Union[str, torch.device] = None,
+ computation_dtype: torch.dtype = None,
+ computation_device: Union[str, torch.device] = None,
+ vram_limit: float = None,
+ name: str = "",
+ disk_map: DiskMap = None,
+ **kwargs
+ ):
+ super().__init__(
+ module,
+ offload_dtype,
+ offload_device,
+ onload_dtype,
+ onload_device,
+ preparing_dtype,
+ preparing_device,
+ computation_dtype,
+ computation_device,
+ vram_limit,
+ name,
+ disk_map,
+ **kwargs
+ )
+ if self.disk_offload:
+ self.required_params = [name for name, _ in self.module.named_parameters(recurse=False)]
+
+ def load_from_disk(self, torch_dtype, device, copy_module=False):
+ if copy_module:
+ module = copy.deepcopy(self.module)
+ else:
+ module = self.module
+ state_dict = {}
+ for name in self.required_params:
+ param = self.disk_map[self.param_name(name)]
+ param = param.to(dtype=torch_dtype, device=device)
+ state_dict[name] = param
+ module.load_state_dict(state_dict, assign=True, strict=False)
+ return module
+
+ def offload_to_disk(self, model: torch.nn.Module):
+ for name in self.required_params:
+ getattr(self, name).to("meta")
+
+ def cast_to(self, module, dtype, device):
+ # Parameter casting is implemented in the model architecture.
+ return module
+
+ def __getattr__(self, name):
+ if name in self.__dict__ or name == "module":
+ return super().__getattr__(name)
+ else:
+ return getattr(self.module, name)
+
+
+class LoRAHotLoadMixin:
+ def init_lora_hotload(self):
+ self.lora_A_weights = []
+ self.lora_B_weights = []
+ self.lora_merger = None
+
+ def lora_forward(self, x, out):
+ if self.lora_merger is None:
+ for lora_A, lora_B in zip(self.lora_A_weights, self.lora_B_weights):
+ out = out + x @ lora_A.T.to(device=x.device, dtype=x.dtype) @ lora_B.T.to(device=x.device, dtype=x.dtype)
+ else:
+ lora_output = []
+ for lora_A, lora_B in zip(self.lora_A_weights, self.lora_B_weights):
+ lora_output.append(x @ lora_A.T @ lora_B.T)
+ lora_output = torch.stack(lora_output)
+ out = self.lora_merger(out, lora_output)
+ return out
+
+
+class AutoWrappedLinear(torch.nn.Linear, AutoTorchModule, LoRAHotLoadMixin):
+ def __init__(
+ self,
+ module: torch.nn.Linear,
+ offload_dtype: torch.dtype = None,
+ offload_device: Union[str, torch.device] = None,
+ onload_dtype: torch.dtype = None,
+ onload_device: Union[str, torch.device] = None,
+ preparing_dtype: torch.dtype = None,
+ preparing_device: Union[str, torch.device] = None,
+ computation_dtype: torch.dtype = None,
+ computation_device: Union[str, torch.device] = None,
+ vram_limit: float = None,
+ name: str = "",
+ disk_map: DiskMap = None,
+ **kwargs
+ ):
+ with skip_model_initialization():
+ super().__init__(
+ in_features=module.in_features,
+ out_features=module.out_features,
+ bias=module.bias is not None,
+ )
+ self.set_dtype_and_device(
+ offload_dtype,
+ offload_device,
+ onload_dtype,
+ onload_device,
+ preparing_dtype,
+ preparing_device,
+ computation_dtype,
+ computation_device,
+ vram_limit,
+ )
+ self.weight = module.weight
+ self.bias = module.bias
+ self.state = 0
+ self.name = name
+ self.init_lora_hotload()
+ self.enable_fp8 = computation_dtype in [torch.float8_e4m3fn, torch.float8_e4m3fnuz]
+ self.computation_device_type = parse_device_type(self.computation_device)
+
+ if offload_dtype == "disk":
+ self.disk_map = disk_map
+ self.disk_offload = True
+ else:
+ self.disk_offload = False
+
+ def fp8_linear(
+ self,
+ input: torch.Tensor,
+ weight: torch.Tensor,
+ bias: torch.Tensor = None,
+ ) -> torch.Tensor:
+ device = input.device
+ origin_dtype = input.dtype
+ origin_shape = input.shape
+ input = input.reshape(-1, origin_shape[-1])
+
+ x_max = torch.max(torch.abs(input), dim=-1, keepdim=True).values
+ # Use the maximum finite value defined by the selected FP8 dtype.
+ fp8_max = torch.finfo(self.computation_dtype).max
+ scale_a = torch.clamp(x_max / fp8_max, min=1.0).float().to(device=device)
+ scale_b = torch.ones((weight.shape[0], 1)).to(device=device)
+ input = input / (scale_a + 1e-8)
+ input = input.to(self.computation_dtype)
+ weight = weight.to(self.computation_dtype)
+ if bias is not None:
+ bias = bias.to(torch.bfloat16)
+
+ result = torch._scaled_mm(
+ input.to(self.computation_dtype),
+ weight.T,
+ scale_a=scale_a,
+ scale_b=scale_b.T,
+ bias=bias,
+ out_dtype=origin_dtype,
+ )
+ new_shape = origin_shape[:-1] + result.shape[-1:]
+ result = result.reshape(new_shape)
+ return result
+
+ def load_from_disk(self, torch_dtype, device, assign=True):
+ weight = self.disk_map[self.name + ".weight"].to(dtype=torch_dtype, device=device)
+ bias = None if self.bias is None else self.disk_map[self.name + ".bias"].to(dtype=torch_dtype, device=device)
+ if assign:
+ state_dict = {"weight": weight}
+ if bias is not None: state_dict["bias"] = bias
+ self.load_state_dict(state_dict, assign=True)
+ return weight, bias
+
+ def offload(self):
+ # offload / onload / preparing -> offload
+ if self.state != 0:
+ if self.disk_offload:
+ self.to("meta")
+ else:
+ self.to(dtype=self.offload_dtype, device=self.offload_device)
+ self.state = 0
+
+ def onload(self):
+ # offload / onload / preparing -> onload
+ if self.state < 1:
+ if self.disk_offload and self.onload_device != "disk" and self.offload_device == "disk":
+ self.load_from_disk(self.onload_dtype, self.onload_device)
+ elif self.onload_device != "disk":
+ self.to(dtype=self.onload_dtype, device=self.onload_device)
+ self.state = 1
+
+ def preparing(self):
+ # onload / preparing -> preparing
+ if self.state != 2:
+ if self.disk_offload and self.preparing_device != "disk" and self.onload_device == "disk":
+ self.load_from_disk(self.preparing_dtype, self.preparing_device)
+ elif self.preparing_device != "disk":
+ self.to(dtype=self.preparing_dtype, device=self.preparing_device)
+ self.state = 2
+
+ def computation(self):
+ # onload / preparing -> computation (temporary)
+ if self.state == 2:
+ torch_dtype, device = self.preparing_dtype, self.preparing_device
+ else:
+ torch_dtype, device = self.onload_dtype, self.onload_device
+ if torch_dtype == self.computation_dtype and device == self.computation_device:
+ weight, bias = self.weight, self.bias
+ elif self.disk_offload and device == "disk":
+ weight, bias = self.load_from_disk(self.computation_dtype, self.computation_device, assign=False)
+ else:
+ weight = self.cast_to(self.weight, self.computation_dtype, self.computation_device)
+ bias = None if self.bias is None else self.cast_to(self.bias, self.computation_dtype, self.computation_device)
+ return weight, bias
+
+ def linear_forward(self, x, weight, bias):
+ if self.enable_fp8:
+ out = self.fp8_linear(x, weight, bias)
+ else:
+ out = torch.nn.functional.linear(x, weight, bias)
+ return out
+
+ def forward(self, x, *args, **kwargs):
+ if self.state == 1 and (self.vram_limit is None or self.check_free_vram()):
+ self.preparing()
+ weight, bias = self.computation()
+ out = self.linear_forward(x, weight, bias)
+ if len(self.lora_A_weights) > 0:
+ out = self.lora_forward(x, out)
+ return out
+
+
+class AutoWrappedQuantizedModule(torch.nn.Linear, AutoTorchModule, LoRAHotLoadMixin):
+ def __init__(
+ self,
+ module: torch.nn.Module,
+ offload_dtype: torch.dtype = None,
+ offload_device: Union[str, torch.device] = None,
+ onload_dtype: torch.dtype = None,
+ onload_device: Union[str, torch.device] = None,
+ preparing_dtype: torch.dtype = None,
+ preparing_device: Union[str, torch.device] = None,
+ computation_dtype: torch.dtype = None,
+ computation_device: Union[str, torch.device] = None,
+ vram_limit: float = None,
+ name: str = "",
+ disk_map: DiskMap = None,
+ quantize=None,
+ metadata: dict = None,
+ **kwargs
+ ):
+ disk_offload = "disk" in (offload_dtype, offload_device, onload_dtype, onload_device)
+ if disk_offload and (disk_map is None or quantize is None):
+ raise ValueError(
+ "Disk offload for quantized layers requires both `disk_map` and `quantize`, "
+ "so each layer can rebuild its packed weight and quant state lazily."
+ )
+ AutoTorchModule.__init__(
+ self,
+ offload_dtype,
+ offload_device,
+ onload_dtype,
+ onload_device,
+ preparing_dtype,
+ preparing_device,
+ computation_dtype,
+ computation_device,
+ vram_limit,
+ )
+ self.in_features = module.in_features
+ self.out_features = module.out_features
+ self.module = module
+ self.name = name
+ self.disk_offload = disk_offload
+ self.disk_map = disk_map
+ self.quantize = quantize
+ self.metadata = metadata
+ self._required_keys = None
+ self.init_lora_hotload()
+
+ def _disk_required_keys(self):
+ if self._required_keys is None:
+ self._required_keys = [key for key in self.disk_map if key.startswith(self.name + ".")]
+ return self._required_keys
+
+ def _load_from_disk(self, device, target=None):
+ module = self.module if target is None else target
+ prefix = self.name + "."
+ subdict = {key: self.disk_map[key] for key in self._disk_required_keys()}
+ rebuilt = self.quantize.unflatten_state_dict(subdict, self.metadata or {})
+ state = {key[len(prefix):]: value for key, value in rebuilt.items()}
+ module.load_state_dict(state, assign=True)
+ module.to(device=device)
+ return module
+
+ def _module_device(self):
+ tensor = next(self.module.parameters(), None)
+ if tensor is None:
+ tensor = next(self.module.buffers(), None)
+ return tensor.device
+
+ def offload(self):
+ if self.state != 0:
+ if self.disk_offload:
+ self.module = self.quantize.build_quantized_shell(self.module, self.computation_dtype, layer_name=self.name)
+ else:
+ self.module.to(device=self.offload_device)
+ self.state = 0
+
+ def onload(self):
+ if self.state < 1:
+ if self.disk_offload and self.onload_device != "disk" and self.offload_device == "disk":
+ self._load_from_disk(self.onload_device)
+ elif self.onload_device != "disk":
+ self.module.to(device=self.onload_device)
+ self.state = 1
+
+ def preparing(self):
+ if self.state != 2:
+ if self.disk_offload and self.preparing_device != "disk" and self.onload_device == "disk":
+ self._load_from_disk(self.preparing_device)
+ elif self.preparing_device != "disk":
+ self.module.to(device=self.preparing_device)
+ self.state = 2
+
+ def computation_module(self):
+ device = self.preparing_device if self.state == 2 else self.onload_device
+ if device == self.computation_device:
+ return self.module
+ if self.disk_offload and device == "disk":
+ transient = self.quantize.build_quantized_shell(self.module, self.computation_dtype, layer_name=self.name)
+ return self._load_from_disk(self.computation_device, target=transient)
+ return copy.deepcopy(self.module).to(device=self.computation_device)
+
+ def forward(self, x, *args, **kwargs):
+ if self.state == 1 and (self.vram_limit is None or self.check_free_vram()):
+ self.preparing()
+ module = self.computation_module()
+ out = module(x, *args, **kwargs)
+ if len(self.lora_A_weights) > 0:
+ out = self.lora_forward(x, out)
+ return out
+
+ def __getattr__(self, name):
+ if name in self.__dict__ or name == "module":
+ return super().__getattr__(name)
+ else:
+ return getattr(self.module, name)
+
+
+def _materialize_root_params_from_disk(model: torch.nn.Module, disk_map: DiskMap):
+ for name, param in list(model.named_parameters(recurse=False)):
+ if param is None or not param.is_meta or name not in disk_map:
+ continue
+ model.register_parameter(name, torch.nn.Parameter(disk_map[name], requires_grad=param.requires_grad))
+ for name, buffer in list(model.named_buffers(recurse=False)):
+ if buffer is None or name not in disk_map:
+ continue
+ model.register_buffer(name, disk_map[name], persistent=True)
+
+
+def enable_vram_management_recursively(model: torch.nn.Module, module_map: dict, vram_config: dict, vram_limit=None, name_prefix="", disk_map=None, quantize=None, **kwargs):
+ if isinstance(model, AutoWrappedNonRecurseModule):
+ model = model.module
+ for name, module in model.named_children():
+ layer_name = name if name_prefix == "" else name_prefix + "." + name
+ if quantize is not None and quantize.is_quantized_linear(module):
+ module_ = AutoWrappedQuantizedModule(module, **vram_config, vram_limit=vram_limit, name=layer_name, disk_map=disk_map, quantize=quantize, **kwargs)
+ setattr(model, name, module_)
+ continue
+ for source_module, target_module in module_map.items():
+ if isinstance(module, source_module):
+ module_ = target_module(module, **vram_config, vram_limit=vram_limit, name=layer_name, disk_map=disk_map, **kwargs)
+ if isinstance(module_, AutoWrappedNonRecurseModule):
+ enable_vram_management_recursively(module_, module_map, vram_config, vram_limit=vram_limit, name_prefix=layer_name, disk_map=disk_map, quantize=quantize, **kwargs)
+ setattr(model, name, module_)
+ break
+ else:
+ enable_vram_management_recursively(module, module_map, vram_config, vram_limit=vram_limit, name_prefix=layer_name, disk_map=disk_map, quantize=quantize, **kwargs)
+
+
+def fill_vram_config(model, vram_config):
+ vram_config_ = vram_config.copy()
+ vram_config_["onload_dtype"] = vram_config["computation_dtype"]
+ vram_config_["onload_device"] = vram_config["computation_device"]
+ vram_config_["preparing_dtype"] = vram_config["computation_dtype"]
+ vram_config_["preparing_device"] = vram_config["computation_device"]
+ for k in vram_config:
+ if vram_config[k] != vram_config_[k]:
+ print(f"No fine-grained VRAM configuration is provided for {model.__class__.__name__}. [`onload`, `preparing`, `computation`] will be the same state. `vram_config` is set to {vram_config_}")
+ break
+ return vram_config_
+
+
+def enable_vram_management(model: torch.nn.Module, module_map: dict, vram_config: dict, vram_limit=None, disk_map=None, quantize=None, **kwargs):
+ for source_module, target_module in module_map.items():
+ # If no fine-grained VRAM configuration is provided, the entire model will be managed uniformly.
+ if isinstance(model, source_module):
+ vram_config = fill_vram_config(model, vram_config)
+ model = target_module(model, **vram_config, vram_limit=vram_limit, disk_map=disk_map, **kwargs)
+ break
+ else:
+ enable_vram_management_recursively(model, module_map, vram_config, vram_limit=vram_limit, disk_map=disk_map, quantize=quantize, **kwargs)
+ if disk_map is not None:
+ _materialize_root_params_from_disk(model, disk_map)
+ # `vram_management_enabled` is a flag that allows the pipeline to determine whether VRAM management is enabled.
+ model.vram_management_enabled = True
+ return model
diff --git a/diffsynth/diffusion/__init__.py b/diffsynth/diffusion/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..d285482ded3ac8ba6406c77078f9ad8f70db7272
--- /dev/null
+++ b/diffsynth/diffusion/__init__.py
@@ -0,0 +1,7 @@
+from .flow_match import FlowMatchScheduler, HiDreamO1FlashScheduler
+from .training_module import DiffusionTrainingModule
+from .logger import ModelLogger
+from .runner import launch_training_task, launch_data_process_task
+from .parsers import *
+from .loss import *
+from .dmd2 import *
diff --git a/diffsynth/diffusion/base_pipeline.py b/diffsynth/diffusion/base_pipeline.py
new file mode 100644
index 0000000000000000000000000000000000000000..5af48a11085e90ab14254f43a094a5c52012361f
--- /dev/null
+++ b/diffsynth/diffusion/base_pipeline.py
@@ -0,0 +1,542 @@
+from PIL import Image
+import torch
+import numpy as np
+from einops import repeat, reduce
+from typing import Union
+from ..core import AutoTorchModule, AutoWrappedLinear, LoRAHotLoadMixin, load_state_dict, ModelConfig, parse_device_type, enable_vram_management
+from ..core.device.npu_compatible_device import get_device_type
+from ..utils.lora import GeneralLoRALoader
+from ..models.model_loader import ModelPool
+from ..utils.controlnet import ControlNetInput
+from ..core.device import get_device_name, IS_NPU_AVAILABLE
+from .template import load_template_model, load_template_data_processor
+
+
+class PipelineUnit:
+ def __init__(
+ self,
+ seperate_cfg: bool = False,
+ take_over: bool = False,
+ input_params: tuple[str] = None,
+ output_params: tuple[str] = None,
+ input_params_posi: dict[str, str] = None,
+ input_params_nega: dict[str, str] = None,
+ onload_model_names: tuple[str] = None
+ ):
+ self.seperate_cfg = seperate_cfg
+ self.take_over = take_over
+ self.input_params = input_params
+ self.output_params = output_params
+ self.input_params_posi = input_params_posi
+ self.input_params_nega = input_params_nega
+ self.onload_model_names = onload_model_names
+
+ def fetch_input_params(self):
+ params = []
+ if self.input_params is not None:
+ for param in self.input_params:
+ params.append(param)
+ if self.input_params_posi is not None:
+ for _, param in self.input_params_posi.items():
+ params.append(param)
+ if self.input_params_nega is not None:
+ for _, param in self.input_params_nega.items():
+ params.append(param)
+ params = sorted(list(set(params)))
+ return params
+
+ def fetch_output_params(self):
+ params = []
+ if self.output_params is not None:
+ for param in self.output_params:
+ params.append(param)
+ return params
+
+ def process(self, pipe, **kwargs) -> dict:
+ return {}
+
+ def post_process(self, pipe, **kwargs) -> dict:
+ return {}
+
+
+class BasePipeline(torch.nn.Module):
+
+ def __init__(
+ self,
+ device=get_device_type(), torch_dtype=torch.float16,
+ height_division_factor=64, width_division_factor=64,
+ time_division_factor=None, time_division_remainder=None,
+ ):
+ super().__init__()
+ # The device and torch_dtype is used for the storage of intermediate variables, not models.
+ self.device = device
+ self.torch_dtype = torch_dtype
+ self.device_type = parse_device_type(device)
+ # The following parameters are used for shape check.
+ self.height_division_factor = height_division_factor
+ self.width_division_factor = width_division_factor
+ self.time_division_factor = time_division_factor
+ self.time_division_remainder = time_division_remainder
+ # VRAM management
+ self.vram_management_enabled = False
+ # Pipeline Unit Runner
+ self.unit_runner = PipelineUnitRunner()
+ # LoRA Loader
+ self.lora_loader = GeneralLoRALoader
+
+
+ def to(self, *args, **kwargs):
+ device, dtype, non_blocking, convert_to_format = torch._C._nn._parse_to(*args, **kwargs)
+ if device is not None:
+ self.device = device
+ if dtype is not None:
+ self.torch_dtype = dtype
+ super().to(*args, **kwargs)
+ return self
+
+
+ def check_resize_height_width(self, height, width, num_frames=None, verbose=1):
+ # Shape check
+ if height % self.height_division_factor != 0:
+ height = (height + self.height_division_factor - 1) // self.height_division_factor * self.height_division_factor
+ if verbose > 0:
+ print(f"height % {self.height_division_factor} != 0. We round it up to {height}.")
+ if width % self.width_division_factor != 0:
+ width = (width + self.width_division_factor - 1) // self.width_division_factor * self.width_division_factor
+ if verbose > 0:
+ print(f"width % {self.width_division_factor} != 0. We round it up to {width}.")
+ if num_frames is None:
+ return height, width
+ else:
+ if num_frames % self.time_division_factor != self.time_division_remainder:
+ num_frames = (num_frames - self.time_division_remainder + self.time_division_factor - 1) // self.time_division_factor * self.time_division_factor + self.time_division_remainder
+ if verbose > 0:
+ print(f"num_frames % {self.time_division_factor} != {self.time_division_remainder}. We round it up to {num_frames}.")
+ return height, width, num_frames
+
+
+ def preprocess_image(self, image, torch_dtype=None, device=None, pattern="B C H W", min_value=-1, max_value=1):
+ # Transform a PIL.Image to torch.Tensor
+ image = torch.Tensor(np.array(image, dtype=np.float32))
+ image = image.to(dtype=torch_dtype or self.torch_dtype, device=device or self.device)
+ image = image * ((max_value - min_value) / 255) + min_value
+ image = repeat(image, f"H W C -> {pattern}", **({"B": 1} if "B" in pattern else {}))
+ return image
+
+
+ def preprocess_video(self, video, torch_dtype=None, device=None, pattern="B C T H W", min_value=-1, max_value=1):
+ # Transform a list of PIL.Image to torch.Tensor
+ video = [self.preprocess_image(image, torch_dtype=torch_dtype, device=device, min_value=min_value, max_value=max_value) for image in video]
+ video = torch.stack(video, dim=pattern.index("T") // 2)
+ return video
+
+
+ def vae_output_to_image(self, vae_output, pattern="B C H W", min_value=-1, max_value=1):
+ # Transform a torch.Tensor to PIL.Image
+ if pattern != "H W C":
+ vae_output = reduce(vae_output, f"{pattern} -> H W C", reduction="mean")
+ image = ((vae_output - min_value) * (255 / (max_value - min_value))).clip(0, 255)
+ image = image.to(device="cpu", dtype=torch.uint8)
+ image = Image.fromarray(image.numpy())
+ return image
+
+
+ def vae_output_to_video(self, vae_output, pattern="B C T H W", min_value=-1, max_value=1):
+ # Transform a torch.Tensor to list of PIL.Image
+ if pattern != "T H W C":
+ vae_output = reduce(vae_output, f"{pattern} -> T H W C", reduction="mean")
+ video = [self.vae_output_to_image(image, pattern="H W C", min_value=min_value, max_value=max_value) for image in vae_output]
+ return video
+
+ def output_audio_format_check(self, audio_output):
+ # output standard foramt: [C, T], output dtype: float()
+ # remove batch dim
+ if audio_output.ndim == 3:
+ audio_output = audio_output.squeeze(0)
+ return audio_output.float().cpu()
+
+ def load_models_to_device(self, model_names):
+ if self.vram_management_enabled:
+ # offload models
+ for name, model in self.named_children():
+ if name not in model_names:
+ if hasattr(model, "vram_management_enabled") and model.vram_management_enabled:
+ if hasattr(model, "offload"):
+ model.offload()
+ else:
+ for module in model.modules():
+ if hasattr(module, "offload"):
+ module.offload()
+ getattr(torch, self.device_type).empty_cache()
+ # onload models
+ for name, model in self.named_children():
+ if name in model_names:
+ if hasattr(model, "vram_management_enabled") and model.vram_management_enabled:
+ if hasattr(model, "onload"):
+ model.onload()
+ else:
+ for module in model.modules():
+ if hasattr(module, "onload"):
+ module.onload()
+
+
+ def generate_noise(self, shape, seed=None, rand_device="cpu", rand_torch_dtype=torch.float32, device=None, torch_dtype=None):
+ # Initialize Gaussian noise
+ generator = None if seed is None else torch.Generator(rand_device).manual_seed(seed)
+ noise = torch.randn(shape, generator=generator, device=rand_device, dtype=rand_torch_dtype)
+ noise = noise.to(dtype=torch_dtype or self.torch_dtype, device=device or self.device)
+ return noise
+
+
+ def get_vram(self):
+ device = self.device if not IS_NPU_AVAILABLE else get_device_name()
+ return getattr(torch, self.device_type).mem_get_info(device)[1] / (1024 ** 3)
+
+ def get_module(self, model, name):
+ if "." in name:
+ name, suffix = name[:name.index(".")], name[name.index(".") + 1:]
+ if name.isdigit():
+ return self.get_module(model[int(name)], suffix)
+ else:
+ return self.get_module(getattr(model, name), suffix)
+ else:
+ return getattr(model, name)
+
+ def freeze_except(self, model_names):
+ self.eval()
+ self.requires_grad_(False)
+ for name in model_names:
+ module = self.get_module(self, name)
+ if module is None:
+ print(f"No {name} models in the pipeline. We cannot enable training on the model. If this occurs during the data processing stage, it is normal.")
+ continue
+ module.train()
+ module.requires_grad_(True)
+
+
+ def blend_with_mask(self, base, addition, mask):
+ return base * (1 - mask) + addition * mask
+
+
+ def step(self, scheduler, latents, progress_id, noise_pred, input_latents=None, inpaint_mask=None, **kwargs):
+ timestep = scheduler.timesteps[progress_id]
+ if inpaint_mask is not None:
+ noise_pred_expected = scheduler.return_to_timestep(scheduler.timesteps[progress_id], latents, input_latents)
+ noise_pred = self.blend_with_mask(noise_pred_expected, noise_pred, inpaint_mask)
+ latents_next = scheduler.step(noise_pred, timestep, latents)
+ return latents_next
+
+
+ def split_pipeline_units(self, model_names: list[str]):
+ return PipelineUnitGraph().split_pipeline_units(self.units, model_names)
+
+
+ def flush_vram_management_device(self, device):
+ for module in self.modules():
+ if isinstance(module, AutoTorchModule):
+ module.offload_device = device
+ module.onload_device = device
+ module.preparing_device = device
+ module.computation_device = device
+
+
+ def check_quant_hot_load(self, module: torch.nn.Module):
+ # A quantized weight cannot absorb a fused LoRA, so hot-loading is the only option for it.
+ if getattr(module, "quantize_config", None) is None:
+ return module
+ return self.enable_lora_hot_loading(module)
+
+
+ def load_lora(
+ self,
+ module: torch.nn.Module,
+ lora_config: Union[ModelConfig, str] = None,
+ alpha=1,
+ hotload=None,
+ state_dict=None,
+ verbose=1,
+ ):
+ module = self.check_quant_hot_load(module)
+ if state_dict is None:
+ if isinstance(lora_config, str):
+ lora = load_state_dict(lora_config, torch_dtype=self.torch_dtype, device=self.device)
+ else:
+ lora_config.download_if_necessary()
+ lora = load_state_dict(lora_config.path, torch_dtype=self.torch_dtype, device=self.device)
+ else:
+ lora = state_dict
+ lora_loader = self.lora_loader(torch_dtype=self.torch_dtype, device=self.device)
+ lora = lora_loader.convert_state_dict(lora)
+ if hotload is None:
+ hotload = hasattr(module, "vram_management_enabled") and getattr(module, "vram_management_enabled")
+ if hotload:
+ if not (hasattr(module, "vram_management_enabled") and getattr(module, "vram_management_enabled")):
+ raise ValueError("VRAM Management is not enabled. LoRA hotloading is not supported.")
+ updated_num = 0
+ for _, module in module.named_modules():
+ if isinstance(module, LoRAHotLoadMixin):
+ name = module.name
+ lora_a_name = f'{name}.lora_A.weight'
+ lora_b_name = f'{name}.lora_B.weight'
+ if lora_a_name in lora and lora_b_name in lora:
+ updated_num += 1
+ module.lora_A_weights.append(lora[lora_a_name] * alpha)
+ module.lora_B_weights.append(lora[lora_b_name])
+ if verbose >= 1:
+ print(f"{updated_num} tensors are patched by LoRA. You can use `pipe.clear_lora()` to clear all LoRA layers.")
+ else:
+ lora_loader.fuse_lora_to_base_model(module, lora, alpha=alpha)
+
+
+ def clear_lora(self, verbose=1):
+ cleared_num = 0
+ for name, module in self.named_modules():
+ if isinstance(module, LoRAHotLoadMixin):
+ if hasattr(module, "lora_A_weights"):
+ if len(module.lora_A_weights) > 0:
+ cleared_num += 1
+ module.lora_A_weights.clear()
+ if hasattr(module, "lora_B_weights"):
+ module.lora_B_weights.clear()
+ if verbose >= 1:
+ print(f"{cleared_num} LoRA layers are cleared.")
+
+
+ def download_and_load_models(self, model_configs: list[ModelConfig] = [], vram_limit: float = None):
+ model_pool = ModelPool()
+ for model_config in model_configs:
+ model_config.download_if_necessary()
+ vram_config = model_config.vram_config()
+ vram_config["computation_dtype"] = vram_config["computation_dtype"] or self.torch_dtype
+ vram_config["computation_device"] = vram_config["computation_device"] or self.device
+ model_pool.auto_load_model(
+ model_config.path,
+ vram_config=vram_config,
+ vram_limit=vram_limit,
+ clear_parameters=model_config.clear_parameters,
+ state_dict=model_config.state_dict,
+ quantize=model_config.quantize,
+ )
+ return model_pool
+
+
+ def check_vram_management_state(self):
+ vram_management_enabled = False
+ for module in self.children():
+ if hasattr(module, "vram_management_enabled") and getattr(module, "vram_management_enabled"):
+ vram_management_enabled = True
+ return vram_management_enabled
+
+
+ def cfg_guided_model_fn(self, model_fn, cfg_scale, inputs_shared, inputs_posi, inputs_nega, **inputs_others):
+ # Positive side forward
+ if inputs_shared.get("positive_only_lora", None) is not None:
+ self.load_lora(self.dit, state_dict=inputs_shared["positive_only_lora"], verbose=0)
+ noise_pred_posi = model_fn(**inputs_posi, **inputs_shared, **inputs_others)
+ if inputs_shared.get("positive_only_lora", None) is not None:
+ self.clear_lora(verbose=0)
+
+ if cfg_scale != 1.0:
+ # Negative side forward
+ if inputs_shared.get("negative_only_lora", None) is not None:
+ self.load_lora(self.dit, state_dict=inputs_shared["negative_only_lora"], verbose=0)
+ noise_pred_nega = model_fn(**inputs_nega, **inputs_shared, **inputs_others)
+ if inputs_shared.get("negative_only_lora", None) is not None:
+ self.clear_lora(verbose=0)
+
+ if isinstance(noise_pred_posi, tuple):
+ # Separately handling different output types of latents, eg. video and audio latents.
+ noise_pred = tuple(
+ n_nega + cfg_scale * (n_posi - n_nega)
+ for n_posi, n_nega in zip(noise_pred_posi, noise_pred_nega)
+ )
+ else:
+ noise_pred = noise_pred_nega + cfg_scale * (noise_pred_posi - noise_pred_nega)
+ else:
+ noise_pred = noise_pred_posi
+ return noise_pred
+
+
+ def load_training_template_model(self, model_config: ModelConfig = None):
+ if model_config is not None:
+ model_config.download_if_necessary()
+ self.template_model = load_template_model(model_config.path, torch_dtype=self.torch_dtype, device=self.device)
+ self.template_data_processor = load_template_data_processor(model_config.path)()
+
+
+ def enable_lora_hot_loading(self, model: torch.nn.Module):
+ if hasattr(model, "vram_management_enabled") and getattr(model, "vram_management_enabled"):
+ return model
+ module_map = {torch.nn.Linear: AutoWrappedLinear}
+ vram_config = {
+ "offload_dtype": self.torch_dtype,
+ "offload_device": self.device,
+ "onload_dtype": self.torch_dtype,
+ "onload_device": self.device,
+ "preparing_dtype": self.torch_dtype,
+ "preparing_device": self.device,
+ "computation_dtype": self.torch_dtype,
+ "computation_device": self.device,
+ }
+ model = enable_vram_management(model, module_map, vram_config=vram_config, quantize=getattr(model, "quantize_config", None))
+ return model
+
+ def compile_pipeline(self, mode: str = "default", dynamic: bool = True, fullgraph: bool = False, compile_models: list = None, **kwargs):
+ """
+ compile the pipeline with torch.compile. The models that will be compiled are determined by the `compilable_models` attribute of the pipeline.
+ If a model has `_repeated_blocks` attribute, we will compile these blocks with regional compilation. Otherwise, we will compile the whole model.
+ See https://docs.pytorch.org/docs/stable/generated/torch.compile.html#torch.compile for details about compilation arguments.
+ Args:
+ mode: The compilation mode, which will be passed to `torch.compile`, options are "default", "reduce-overhead", "max-autotune" and "max-autotune-no-cudagraphs. Default to "default".
+ dynamic: Whether to enable dynamic graph compilation to support dynamic input shapes, which will be passed to `torch.compile`. Default to True (recommended).
+ fullgraph: Whether to use full graph compilation, which will be passed to `torch.compile`. Default to False (recommended).
+ compile_models: The list of model names to be compiled. If None, we will compile the models in `pipeline.compilable_models`. Default to None.
+ **kwargs: Other arguments for `torch.compile`.
+ """
+ compile_models = compile_models or getattr(self, "compilable_models", [])
+ if len(compile_models) == 0:
+ print("No compilable models in the pipeline. Skip compilation.")
+ return
+ for name in compile_models:
+ model = getattr(self, name, None)
+ if model is None:
+ print(f"Model '{name}' not found in the pipeline.")
+ continue
+ repeated_blocks = getattr(model, "_repeated_blocks", None)
+ # regional compilation for repeated blocks.
+ if repeated_blocks is not None:
+ for submod in model.modules():
+ if submod.__class__.__name__ in repeated_blocks:
+ submod.compile(mode=mode, dynamic=dynamic, fullgraph=fullgraph, **kwargs)
+ # compile the whole model.
+ else:
+ model.compile(mode=mode, dynamic=dynamic, fullgraph=fullgraph, **kwargs)
+ print(f"{name} is compiled with mode={mode}, dynamic={dynamic}, fullgraph={fullgraph}.")
+
+
+class PipelineUnitGraph:
+ def __init__(self):
+ pass
+
+ def build_edges(self, units: list[PipelineUnit]):
+ # Establish dependencies between units
+ # to search for subsequent related computation units.
+ last_compute_unit_id = {}
+ edges = []
+ for unit_id, unit in enumerate(units):
+ for input_param in unit.fetch_input_params():
+ if input_param in last_compute_unit_id:
+ edges.append((last_compute_unit_id[input_param], unit_id))
+ for output_param in unit.fetch_output_params():
+ last_compute_unit_id[output_param] = unit_id
+ return edges
+
+ def build_chains(self, units: list[PipelineUnit]):
+ # Establish updating chains for each variable
+ # to track their computation process.
+ params = sum([unit.fetch_input_params() + unit.fetch_output_params() for unit in units], [])
+ params = sorted(list(set(params)))
+ chains = {param: [] for param in params}
+ for unit_id, unit in enumerate(units):
+ for param in unit.fetch_output_params():
+ chains[param].append(unit_id)
+ return chains
+
+ def search_direct_unit_ids(self, units: list[PipelineUnit], model_names: list[str]):
+ # Search for units that directly participate in the model's computation.
+ related_unit_ids = []
+ for unit_id, unit in enumerate(units):
+ for model_name in model_names:
+ if unit.onload_model_names is not None and model_name in unit.onload_model_names:
+ related_unit_ids.append(unit_id)
+ break
+ return related_unit_ids
+
+ def search_related_unit_ids(self, edges, start_unit_ids, direction="target"):
+ # Search for subsequent related computation units.
+ related_unit_ids = [unit_id for unit_id in start_unit_ids]
+ while True:
+ neighbors = []
+ for source, target in edges:
+ if direction == "target" and source in related_unit_ids and target not in related_unit_ids:
+ neighbors.append(target)
+ elif direction == "source" and source not in related_unit_ids and target in related_unit_ids:
+ neighbors.append(source)
+ neighbors = sorted(list(set(neighbors)))
+ if len(neighbors) == 0:
+ break
+ else:
+ related_unit_ids.extend(neighbors)
+ related_unit_ids = sorted(list(set(related_unit_ids)))
+ return related_unit_ids
+
+ def search_updating_unit_ids(self, units: list[PipelineUnit], chains, related_unit_ids):
+ # If the input parameters of this subgraph are updated outside the subgraph,
+ # search for the units where these updates occur.
+ first_compute_unit_id = {}
+ for unit_id in related_unit_ids:
+ for param in units[unit_id].fetch_input_params():
+ if param not in first_compute_unit_id:
+ first_compute_unit_id[param] = unit_id
+ updating_unit_ids = []
+ for param in first_compute_unit_id:
+ unit_id = first_compute_unit_id[param]
+ chain = chains[param]
+ if unit_id in chain and chain.index(unit_id) != len(chain) - 1:
+ for unit_id_ in chain[chain.index(unit_id) + 1:]:
+ if unit_id_ not in related_unit_ids:
+ updating_unit_ids.append(unit_id_)
+ related_unit_ids.extend(updating_unit_ids)
+ related_unit_ids = sorted(list(set(related_unit_ids)))
+ return related_unit_ids
+
+ def split_pipeline_units(self, units: list[PipelineUnit], model_names: list[str]):
+ # Split the computation graph,
+ # separating all model-related computations.
+ related_unit_ids = self.search_direct_unit_ids(units, model_names)
+ edges = self.build_edges(units)
+ chains = self.build_chains(units)
+ while True:
+ num_related_unit_ids = len(related_unit_ids)
+ related_unit_ids = self.search_related_unit_ids(edges, related_unit_ids, "target")
+ related_unit_ids = self.search_updating_unit_ids(units, chains, related_unit_ids)
+ if len(related_unit_ids) == num_related_unit_ids:
+ break
+ else:
+ num_related_unit_ids = len(related_unit_ids)
+ related_units = [units[i] for i in related_unit_ids]
+ unrelated_units = [units[i] for i in range(len(units)) if i not in related_unit_ids]
+ return related_units, unrelated_units
+
+
+class PipelineUnitRunner:
+ def __init__(self):
+ pass
+
+ def __call__(self, unit: PipelineUnit, pipe: BasePipeline, inputs_shared: dict, inputs_posi: dict, inputs_nega: dict) -> tuple[dict, dict]:
+ if unit.take_over:
+ # Let the pipeline unit take over this function.
+ inputs_shared, inputs_posi, inputs_nega = unit.process(pipe, inputs_shared=inputs_shared, inputs_posi=inputs_posi, inputs_nega=inputs_nega)
+ elif unit.seperate_cfg:
+ # Positive side
+ processor_inputs = {name: inputs_posi.get(name_) for name, name_ in unit.input_params_posi.items()}
+ if unit.input_params is not None:
+ for name in unit.input_params:
+ processor_inputs[name] = inputs_shared.get(name)
+ processor_outputs = unit.process(pipe, **processor_inputs)
+ inputs_posi.update(processor_outputs)
+ # Negative side
+ if inputs_shared["cfg_scale"] != 1:
+ processor_inputs = {name: inputs_nega.get(name_) for name, name_ in unit.input_params_nega.items()}
+ if unit.input_params is not None:
+ for name in unit.input_params:
+ processor_inputs[name] = inputs_shared.get(name)
+ processor_outputs = unit.process(pipe, **processor_inputs)
+ inputs_nega.update(processor_outputs)
+ else:
+ inputs_nega.update(processor_outputs)
+ else:
+ processor_inputs = {name: inputs_shared.get(name) for name in unit.input_params}
+ processor_outputs = unit.process(pipe, **processor_inputs)
+ inputs_shared.update(processor_outputs)
+ return inputs_shared, inputs_posi, inputs_nega
diff --git a/diffsynth/diffusion/ddim_scheduler.py b/diffsynth/diffusion/ddim_scheduler.py
new file mode 100644
index 0000000000000000000000000000000000000000..09a5bf3cd382b7689e2778387ff6b03552819bf3
--- /dev/null
+++ b/diffsynth/diffusion/ddim_scheduler.py
@@ -0,0 +1,109 @@
+import torch, math
+
+
+class DDIMScheduler():
+
+ def __init__(self, num_train_timesteps=1000, beta_start=0.00085, beta_end=0.012, beta_schedule="scaled_linear", prediction_type="epsilon", rescale_zero_terminal_snr=False):
+ self.num_train_timesteps = num_train_timesteps
+ if beta_schedule == "scaled_linear":
+ betas = torch.square(torch.linspace(math.sqrt(beta_start), math.sqrt(beta_end), num_train_timesteps, dtype=torch.float32))
+ elif beta_schedule == "linear":
+ betas = torch.linspace(beta_start, beta_end, num_train_timesteps, dtype=torch.float32)
+ else:
+ raise NotImplementedError(f"{beta_schedule} is not implemented")
+ self.alphas_cumprod = torch.cumprod(1.0 - betas, dim=0)
+ if rescale_zero_terminal_snr:
+ self.alphas_cumprod = self.rescale_zero_terminal_snr(self.alphas_cumprod)
+ self.alphas_cumprod = self.alphas_cumprod.tolist()
+ self.set_timesteps(10)
+ self.prediction_type = prediction_type
+ self.training = False
+
+
+ def rescale_zero_terminal_snr(self, alphas_cumprod):
+ alphas_bar_sqrt = alphas_cumprod.sqrt()
+
+ # Store old values.
+ alphas_bar_sqrt_0 = alphas_bar_sqrt[0].clone()
+ alphas_bar_sqrt_T = alphas_bar_sqrt[-1].clone()
+
+ # Shift so the last timestep is zero.
+ alphas_bar_sqrt -= alphas_bar_sqrt_T
+
+ # Scale so the first timestep is back to the old value.
+ alphas_bar_sqrt *= alphas_bar_sqrt_0 / (alphas_bar_sqrt_0 - alphas_bar_sqrt_T)
+
+ # Convert alphas_bar_sqrt to betas
+ alphas_bar = alphas_bar_sqrt.square() # Revert sqrt
+
+ return alphas_bar
+
+
+ def set_timesteps(self, num_inference_steps, denoising_strength=1.0, training=False, **kwargs):
+ # The timesteps are aligned to 999...0, which is different from other implementations,
+ # but I think this implementation is more reasonable in theory.
+ max_timestep = max(round(self.num_train_timesteps * denoising_strength) - 1, 0)
+ num_inference_steps = min(num_inference_steps, max_timestep + 1)
+ if num_inference_steps == 1:
+ self.timesteps = torch.Tensor([max_timestep])
+ else:
+ step_length = max_timestep / (num_inference_steps - 1)
+ self.timesteps = torch.Tensor([round(max_timestep - i*step_length) for i in range(num_inference_steps)])
+ self.training = training
+
+
+ def denoise(self, model_output, sample, alpha_prod_t, alpha_prod_t_prev):
+ if self.prediction_type == "epsilon":
+ weight_e = math.sqrt(1 - alpha_prod_t_prev) - math.sqrt(alpha_prod_t_prev * (1 - alpha_prod_t) / alpha_prod_t)
+ weight_x = math.sqrt(alpha_prod_t_prev / alpha_prod_t)
+ prev_sample = sample * weight_x + model_output * weight_e
+ elif self.prediction_type == "v_prediction":
+ weight_e = -math.sqrt(alpha_prod_t_prev * (1 - alpha_prod_t)) + math.sqrt(alpha_prod_t * (1 - alpha_prod_t_prev))
+ weight_x = math.sqrt(alpha_prod_t * alpha_prod_t_prev) + math.sqrt((1 - alpha_prod_t) * (1 - alpha_prod_t_prev))
+ prev_sample = sample * weight_x + model_output * weight_e
+ else:
+ raise NotImplementedError(f"{self.prediction_type} is not implemented")
+ return prev_sample
+
+
+ def step(self, model_output, timestep, sample, to_final=False):
+ alpha_prod_t = self.alphas_cumprod[int(timestep.flatten().tolist()[0])]
+ if isinstance(timestep, torch.Tensor):
+ timestep = timestep.cpu()
+ timestep_id = torch.argmin((self.timesteps - timestep).abs())
+ if to_final or timestep_id + 1 >= len(self.timesteps):
+ alpha_prod_t_prev = 1.0
+ else:
+ timestep_prev = int(self.timesteps[timestep_id + 1])
+ alpha_prod_t_prev = self.alphas_cumprod[timestep_prev]
+
+ return self.denoise(model_output, sample, alpha_prod_t, alpha_prod_t_prev)
+
+
+ def return_to_timestep(self, timestep, sample, sample_stablized):
+ alpha_prod_t = self.alphas_cumprod[int(timestep.flatten().tolist()[0])]
+ noise_pred = (sample - math.sqrt(alpha_prod_t) * sample_stablized) / math.sqrt(1 - alpha_prod_t)
+ return noise_pred
+
+
+ def add_noise(self, original_samples, noise, timestep):
+ timestep_id = max(min(int(timestep.flatten().tolist()[0]), len(self.alphas_cumprod)-1), 0)
+ sqrt_alpha_prod = math.sqrt(self.alphas_cumprod[timestep_id])
+ sqrt_one_minus_alpha_prod = math.sqrt(1 - self.alphas_cumprod[timestep_id])
+ noisy_samples = sqrt_alpha_prod * original_samples + sqrt_one_minus_alpha_prod * noise
+ return noisy_samples
+
+
+ def training_target(self, sample, noise, timestep):
+ if self.prediction_type == "epsilon":
+ return noise
+ else:
+ timestep_id = max(min(int(timestep.flatten().tolist()[0]), len(self.alphas_cumprod)-1), 0)
+ sqrt_alpha_prod = math.sqrt(self.alphas_cumprod[timestep_id])
+ sqrt_one_minus_alpha_prod = math.sqrt(1 - self.alphas_cumprod[timestep_id])
+ target = sqrt_alpha_prod * noise - sqrt_one_minus_alpha_prod * sample
+ return target
+
+
+ def training_weight(self, timestep):
+ return 1.0
diff --git a/diffsynth/diffusion/dmd2.py b/diffsynth/diffusion/dmd2.py
new file mode 100644
index 0000000000000000000000000000000000000000..8f407ab64180247d27e533d18efcf66349487ce4
--- /dev/null
+++ b/diffsynth/diffusion/dmd2.py
@@ -0,0 +1,601 @@
+from dataclasses import dataclass
+from typing import Optional
+import torch
+import torch.nn.functional as F
+from accelerate import Accelerator
+from tqdm import tqdm
+from .runner import get_optimizer_class, initialize_deepspeed_gradient_checkpointing
+from diffsynth.core import OffloadTrainingManager
+
+
+@dataclass
+class DMD2Config:
+ student_update_freq: int = 5
+ student_sample_steps: int = 4
+ student_sample_type: str = "sde"
+ student_schedule: str = "uniform"
+ student_t_list: Optional[list[float]] = None
+ matching_t_min: float = 0.001
+ matching_t_max: float = 0.999
+ matching_t_sampling: str = "uniform"
+ matching_t_mean: float = 0.0
+ matching_t_std: float = 1.0
+ gan_loss_weight: float = 0.03
+ gan_r1_reg_weight: float = 0.0
+ gan_r1_reg_alpha: float = 0.1
+ fake_score_learning_rate: Optional[float] = None
+ discriminator_learning_rate: Optional[float] = None
+ feature_indices: Optional[list[int]] = None
+ teacher_cfg_scale: float = 1.0
+ student_grad_clip_norm: Optional[float] = 10.0
+
+
+def _get_dmd2_pipe_model(module, model_name):
+ if model_name is None:
+ return None
+ return getattr(module.pipe, model_name)
+
+
+def get_dmd2_student_model(module):
+ return _get_dmd2_pipe_model(module, module.dmd2_student_model_name)
+
+
+def get_dmd2_teacher_model(module):
+ return _get_dmd2_pipe_model(module, module.dmd2_teacher_model_name)
+
+
+def get_dmd2_fake_score_model(module):
+ return _get_dmd2_pipe_model(module, module.dmd2_fake_score_model_name)
+
+
+def get_dmd2_discriminator(module):
+ return _get_dmd2_pipe_model(module, module.dmd2_discriminator_model_name)
+
+
+def _dmd2_model_state_names(module, model_name, param_names=None):
+ if model_name is None:
+ return set()
+ model = getattr(module.pipe, model_name, None)
+ if model is None:
+ return set()
+ names = model.state_dict().keys() if param_names is None else param_names
+ return {f"pipe.{model_name}.{name}" for name in names}
+
+
+def export_dmd2_trainable_state_dict(module, state_dict, remove_prefix=None):
+ student_names = _dmd2_model_state_names(module, module.dmd2_student_model_name, module._dmd2_student_param_names)
+ state_names = set(student_names)
+ if remove_prefix is None:
+ state_names.update(_dmd2_model_state_names(module, module.dmd2_teacher_model_name))
+ state_names.update(_dmd2_model_state_names(module, module.dmd2_fake_score_model_name))
+ state_names.update(_dmd2_model_state_names(module, module.dmd2_discriminator_model_name))
+ state_dict = {name: param for name, param in state_dict.items() if name in state_names}
+ if remove_prefix is not None:
+ state_dict = {name[len(remove_prefix):] if name.startswith(remove_prefix) else name: param for name, param in state_dict.items()}
+ return state_dict
+
+
+def set_dmd2_train_phase(module, student_phase: bool):
+ student_model = get_dmd2_student_model(module)
+ student_model.train(student_phase)
+ for name, param in student_model.named_parameters():
+ param.requires_grad = student_phase and name in module._dmd2_student_param_names
+
+ fake_score_model = get_dmd2_fake_score_model(module)
+ fake_score_model.train(not student_phase)
+ for name, param in fake_score_model.named_parameters():
+ param.requires_grad = (not student_phase) and name in module._dmd2_fake_score_param_names
+
+ discriminator = get_dmd2_discriminator(module)
+ if discriminator is not None:
+ discriminator.train(not student_phase)
+ discriminator.requires_grad_(not student_phase)
+
+
+def _expand_like(value, target, dtype=None):
+ if dtype is None:
+ dtype = target.dtype
+ if not isinstance(value, torch.Tensor):
+ value = torch.tensor(value, device=target.device, dtype=dtype)
+ value = value.to(device=target.device, dtype=dtype)
+ while value.ndim < target.ndim:
+ value = value.view(*value.shape, 1)
+ return value
+
+
+def flow_to_x0(latents, flow, sigma):
+ original_dtype = latents.dtype
+ latents = latents.to(torch.float64)
+ flow = flow.to(torch.float64)
+ sigma = _expand_like(sigma, latents)
+ return (latents - sigma * flow).to(original_dtype)
+
+
+def _forward_process(x0, eps, sigma):
+ original_dtype = x0.dtype
+ x0 = x0.to(torch.float64)
+ eps = eps.to(torch.float64)
+ sigma = _expand_like(sigma, x0)
+ return ((1 - sigma) * x0 + sigma * eps).to(original_dtype)
+
+
+def _scale_noise(noise, sigma):
+ original_dtype = noise.dtype
+ noise = noise.to(torch.float64)
+ sigma = _expand_like(sigma, noise)
+ return (noise * sigma).to(original_dtype)
+
+
+def make_dmd2_student_schedule(
+ pipe,
+ num_steps,
+ device,
+ student_schedule="uniform",
+ student_t_list=None,
+ matching_t_max=0.999,
+):
+ time_dtype = torch.float64
+ if student_t_list is not None:
+ sigmas = torch.tensor(student_t_list, device=device, dtype=time_dtype)
+ if sigmas[-1].item() != 0:
+ raise ValueError("`dmd2_student_t_list` must include a final 0.")
+ if len(sigmas) != num_steps + 1:
+ raise ValueError("The student sigma schedule length must equal `student_sample_steps + 1`.")
+ timesteps = sigmas * pipe.scheduler.num_train_timesteps
+ return sigmas, timesteps
+ if student_schedule == "uniform":
+ sigma_start = min(float(matching_t_max), 0.999)
+ sigmas = torch.linspace(sigma_start, 0.0, num_steps + 1, device=device, dtype=time_dtype)
+ timesteps = sigmas * pipe.scheduler.num_train_timesteps
+ return sigmas, timesteps
+ raise ValueError(f"Unsupported DMD2 student schedule: {student_schedule}")
+
+
+def _variational_score_distillation_loss(gen_data, teacher_x0, fake_score_x0):
+ dims = tuple(range(1, teacher_x0.ndim))
+ with torch.no_grad():
+ weight = 1 / ((gen_data.float() - teacher_x0.float()).abs().mean(dim=dims, keepdim=True) + 1e-6)
+ weight = weight.to(dtype=gen_data.dtype)
+ pseudo_target = gen_data - (fake_score_x0 - teacher_x0) * weight
+ loss = 0.5 * F.mse_loss(gen_data.float(), pseudo_target.float(), reduction="mean")
+ return loss
+
+
+def _gan_loss_generator(fake_logits):
+ assert fake_logits.ndim == 2, f"fake_logits has shape {fake_logits.shape}"
+ gan_loss = F.softplus(-fake_logits).mean()
+ return gan_loss
+
+
+def _gan_loss_discriminator(real_logits, fake_logits):
+ assert fake_logits.ndim == 2, f"fake_logits has shape {fake_logits.shape}"
+ assert real_logits.ndim == 2, f"real_logits has shape {real_logits.shape}"
+ gan_loss = F.softplus(fake_logits).mean() + F.softplus(-real_logits).mean()
+ return gan_loss
+
+
+class DMD2Loss:
+ def __init__(self, config: DMD2Config):
+ self.config = config
+
+ def _model_forward_x0(
+ self,
+ module,
+ model,
+ model_fn,
+ latents,
+ timestep,
+ sigma,
+ inputs_shared,
+ inputs_posi,
+ ):
+ shared = dict(inputs_shared)
+ posi = dict(inputs_posi)
+ shared["latents"] = latents
+ flow = model_fn(
+ module.pipe,
+ model,
+ timestep=timestep,
+ progress_id=0,
+ num_inference_steps=1,
+ inputs_shared=shared,
+ inputs_posi=posi,
+ )
+ return flow_to_x0(latents, flow, sigma)
+
+ def _model_forward_features(
+ self,
+ module,
+ model,
+ model_fn,
+ latents,
+ timestep,
+ inputs_shared,
+ inputs_posi,
+ ):
+ discriminator = get_dmd2_discriminator(module)
+ if discriminator is None:
+ raise ValueError("DMD2 feature extraction requires a discriminator.")
+ if self.config.feature_indices is None:
+ raise ValueError("DMD2 feature extraction requires `dmd2_feature_indices`.")
+ shared = dict(inputs_shared)
+ posi = dict(inputs_posi)
+ shared["latents"] = latents
+ return model_fn(
+ module.pipe,
+ model,
+ timestep=timestep,
+ progress_id=0,
+ num_inference_steps=1,
+ inputs_shared=shared,
+ inputs_posi=posi,
+ feature_indices=set(self.config.feature_indices),
+ return_features=True,
+ )
+
+ def _teacher_forward_x0(
+ self,
+ module,
+ latents,
+ timestep,
+ sigma,
+ inputs_shared,
+ inputs_posi,
+ inputs_nega,
+ ):
+ teacher_x0_pos = self._model_forward_x0(
+ module,
+ get_dmd2_teacher_model(module),
+ module.dmd2_model_fn_teacher,
+ latents,
+ timestep,
+ sigma,
+ inputs_shared,
+ inputs_posi,
+ )
+ if self.config.teacher_cfg_scale <= 1.0:
+ return teacher_x0_pos
+
+
+ x0_neg = self._model_forward_x0(
+ module,
+ get_dmd2_teacher_model(module),
+ module.dmd2_model_fn_teacher,
+ latents,
+ timestep,
+ sigma,
+ inputs_shared,
+ inputs_nega,
+ )
+ return x0_neg + float(self.config.teacher_cfg_scale) * (teacher_x0_pos - x0_neg)
+
+ def _sample_matching_timestep(self, pipe, device, dtype, batch_size=1, inputs_shared=None, real_data=None):
+ if self.config.matching_t_min > self.config.matching_t_max:
+ raise ValueError("`dmd2_matching_t_min` must be <= `dmd2_matching_t_max`.")
+ time_dtype = torch.float64
+ if self.config.matching_t_sampling == "uniform":
+ sigma = torch.rand(batch_size, device=device, dtype=time_dtype)
+ sigma = sigma * (self.config.matching_t_max - self.config.matching_t_min) + self.config.matching_t_min
+ timestep = sigma * pipe.scheduler.num_train_timesteps
+ return timestep, sigma
+ if self.config.matching_t_sampling == "logitnormal":
+ sigma = torch.sigmoid(
+ torch.randn(batch_size, device=device, dtype=time_dtype) * self.config.matching_t_std
+ + self.config.matching_t_mean
+ )
+ sigma = sigma * (self.config.matching_t_max - self.config.matching_t_min) + self.config.matching_t_min
+ sigma = sigma.clamp(self.config.matching_t_min, self.config.matching_t_max)
+ timestep = sigma * pipe.scheduler.num_train_timesteps
+ return timestep, sigma
+ raise ValueError(f"Unsupported DMD2 matching timestep sampling: {self.config.matching_t_sampling}")
+
+ def _generate_student_data(self, module, real_data, inputs_shared, inputs_posi):
+ pipe = module.pipe
+ device, dtype = real_data.device, real_data.dtype
+ batch_size = real_data.shape[0]
+ student_sigmas, student_timesteps = make_dmd2_student_schedule(
+ pipe,
+ self.config.student_sample_steps,
+ device,
+ student_schedule=self.config.student_schedule,
+ student_t_list=self.config.student_t_list,
+ matching_t_max=self.config.matching_t_max,
+ )
+ if len(student_sigmas) != self.config.student_sample_steps + 1:
+ raise ValueError("The student sigma schedule length must equal `student_sample_steps + 1`.")
+
+ if self.config.student_sample_steps == 1:
+ timestep = student_timesteps[0:1].expand(batch_size)
+ sigma = student_sigmas[0:1].expand(batch_size)
+ input_student = _scale_noise(torch.randn_like(real_data), sigma)
+ else:
+ step_id = torch.randint(0, self.config.student_sample_steps, (batch_size,), device=device)
+ sigma = student_sigmas[step_id]
+ timestep = student_timesteps[step_id]
+ eps_student = torch.randn_like(real_data)
+ input_student = _forward_process(real_data, eps_student, sigma)
+ gen_data = self._model_forward_x0(
+ module,
+ get_dmd2_student_model(module),
+ module.dmd2_model_fn_student,
+ input_student,
+ timestep,
+ sigma,
+ inputs_shared,
+ inputs_posi,
+ )
+ return gen_data, input_student
+
+ def _compute_real_feat(self, module, real_data, timestep, sigma, eps, inputs_shared, inputs_posi):
+ real_timestep, real_sigma, real_eps = timestep, sigma, eps
+ perturbed_real = _forward_process(real_data, real_eps, real_sigma)
+ real_feat = self._model_forward_features(
+ module,
+ get_dmd2_teacher_model(module),
+ module.dmd2_model_fn_teacher,
+ perturbed_real,
+ real_timestep,
+ inputs_shared,
+ inputs_posi,
+ )
+ return real_feat, real_timestep, real_sigma
+
+ def _student_update_step(self, module, real_data, inputs_shared, inputs_posi, inputs_nega):
+ gen_data, input_student = self._generate_student_data(module, real_data, inputs_shared, inputs_posi)
+ timestep, sigma = self._sample_matching_timestep(
+ module.pipe,
+ real_data.device,
+ real_data.dtype,
+ real_data.shape[0],
+ inputs_shared=inputs_shared,
+ real_data=real_data,
+ )
+ eps = torch.randn_like(real_data)
+ perturbed_data = _forward_process(gen_data, eps, sigma)
+
+ with torch.no_grad():
+ fake_score_x0 = self._model_forward_x0(
+ module,
+ get_dmd2_fake_score_model(module),
+ module.dmd2_model_fn_fake_score,
+ perturbed_data,
+ timestep,
+ sigma,
+ inputs_shared,
+ inputs_posi,
+ )
+
+ with torch.no_grad():
+ teacher_x0 = self._teacher_forward_x0(
+ module, perturbed_data, timestep, sigma, inputs_shared, inputs_posi, inputs_nega
+ )
+ if self.config.gan_loss_weight > 0:
+ fake_feat = self._model_forward_features(
+ module,
+ get_dmd2_teacher_model(module),
+ module.dmd2_model_fn_teacher,
+ perturbed_data,
+ timestep,
+ inputs_shared,
+ inputs_posi,
+ )
+ fake_logits_gen = get_dmd2_discriminator(module)(fake_feat)
+ gan_loss_gen = _gan_loss_generator(fake_logits_gen)
+ else:
+ gan_loss_gen = torch.zeros((), device=real_data.device, dtype=torch.float32)
+
+ vsd_loss = _variational_score_distillation_loss(gen_data, teacher_x0.detach(), fake_score_x0)
+ gan_loss_weighted = self.config.gan_loss_weight * gan_loss_gen
+ loss = vsd_loss + gan_loss_weighted
+
+ return {"total_loss": loss}
+
+ def _fake_score_discriminator_update_step(self, module, real_data, inputs_shared, inputs_posi, inputs_nega):
+ with torch.no_grad():
+ gen_data, _ = self._generate_student_data(module, real_data, inputs_shared, inputs_posi)
+ timestep, sigma = self._sample_matching_timestep(
+ module.pipe,
+ real_data.device,
+ real_data.dtype,
+ real_data.shape[0],
+ inputs_shared=inputs_shared,
+ real_data=real_data,
+ )
+ eps = torch.randn_like(real_data)
+ x_t_sg = _forward_process(gen_data, eps, sigma)
+
+ fake_score_x0 = self._model_forward_x0(
+ module,
+ get_dmd2_fake_score_model(module),
+ module.dmd2_model_fn_fake_score,
+ x_t_sg,
+ timestep,
+ sigma,
+ inputs_shared,
+ inputs_posi,
+ )
+ fake_score_loss = F.mse_loss(fake_score_x0.float(), gen_data.float(), reduction="mean")
+
+ gan_loss_disc = torch.zeros_like(fake_score_loss)
+ gan_loss_ar1 = torch.zeros_like(fake_score_loss)
+ gan_loss_logit_reg = torch.zeros_like(fake_score_loss)
+ if self.config.gan_loss_weight > 0:
+ with torch.no_grad():
+ fake_feat = self._model_forward_features(
+ module,
+ get_dmd2_teacher_model(module),
+ module.dmd2_model_fn_teacher,
+ x_t_sg,
+ timestep,
+ inputs_shared,
+ inputs_posi,
+ )
+ real_feat, real_timestep, real_sigma = self._compute_real_feat(
+ module, real_data, timestep, sigma, eps, inputs_shared, inputs_posi
+ )
+ discriminator = get_dmd2_discriminator(module)
+ real_logits = discriminator(real_feat)
+ fake_logits = discriminator(fake_feat)
+ gan_loss_disc = _gan_loss_discriminator(real_logits, fake_logits)
+ if self.config.gan_r1_reg_weight > 0:
+ perturbed_real_alpha = real_data + self.config.gan_r1_reg_alpha * torch.randn_like(real_data)
+ with torch.no_grad():
+ real_feat_alpha = self._model_forward_features(
+ module,
+ get_dmd2_teacher_model(module),
+ module.dmd2_model_fn_teacher,
+ perturbed_real_alpha,
+ real_timestep,
+ inputs_shared,
+ inputs_posi,
+ )
+ real_logits_alpha = discriminator(real_feat_alpha)
+ gan_loss_ar1 = F.mse_loss(real_logits, real_logits_alpha, reduction="mean")
+
+ loss = (
+ fake_score_loss
+ + gan_loss_disc
+ + self.config.gan_r1_reg_weight * gan_loss_ar1
+ )
+ return {"total_loss": loss}
+
+ def __call__(self, module, inputs, iteration):
+ inputs_shared, inputs_posi, inputs_nega = inputs
+ real_data = inputs_shared.get("input_latents")
+ if real_data is None:
+ raise ValueError("DMD2 requires image latents from the dataset. Please provide training images.")
+ student_phase = iteration % self.config.student_update_freq == 0
+ set_dmd2_train_phase(module, student_phase)
+ if student_phase:
+ return self._student_update_step(module, real_data, inputs_shared, inputs_posi, inputs_nega)
+ return self._fake_score_discriminator_update_step(module, real_data, inputs_shared, inputs_posi, inputs_nega)
+
+
+def _trainable_params(module):
+ return [param for param in module.parameters() if param.requires_grad]
+
+
+def _dmd2_current_optimizers(config, optimizers, iteration):
+ if iteration % config.student_update_freq == 0:
+ return [optimizers["student"]], [optimizers["student_scheduler"]]
+ current_optimizers = [optimizers["fake_score"]]
+ current_schedulers = [optimizers["fake_score_scheduler"]]
+ if "discriminator" in optimizers:
+ current_optimizers.append(optimizers["discriminator"])
+ current_schedulers.append(optimizers["discriminator_scheduler"])
+ return current_optimizers, current_schedulers
+
+
+def launch_dmd2_training_task(
+ accelerator: Accelerator,
+ dataset: torch.utils.data.Dataset,
+ model,
+ model_logger,
+ learning_rate: float = 1e-5,
+ weight_decay: float = 1e-2,
+ num_workers: int = 1,
+ save_steps: int = None,
+ num_epochs: int = 1,
+ enable_model_cpu_offload: bool = False,
+ enable_optimizer_cpu_offload: bool = False,
+ cpu_offload_split_threshold: int = None,
+ customized_optimizer: str = None,
+ args=None,
+ **kwargs,
+):
+ if args is not None:
+ learning_rate = args.learning_rate
+ weight_decay = args.weight_decay
+ num_workers = args.dataset_num_workers
+ save_steps = args.save_steps
+ num_epochs = args.num_epochs
+ enable_model_cpu_offload = args.enable_model_cpu_offload
+ enable_optimizer_cpu_offload = args.enable_optimizer_cpu_offload
+ cpu_offload_split_threshold = args.cpu_offload_split_threshold
+ customized_optimizer = args.customized_optimizer
+
+ optimizer_class = get_optimizer_class(customized_optimizer)
+ config = model.dmd2_config
+ fake_score_lr = config.fake_score_learning_rate or learning_rate
+ discriminator_lr = config.discriminator_learning_rate or learning_rate
+
+ student_model = get_dmd2_student_model(model)
+ fake_score_model = get_dmd2_fake_score_model(model)
+ discriminator = get_dmd2_discriminator(model)
+ has_discriminator = discriminator is not None
+
+ student_optimizer = optimizer_class(_trainable_params(student_model), lr=learning_rate, weight_decay=weight_decay)
+ fake_score_optimizer = optimizer_class(fake_score_model.parameters(), lr=fake_score_lr, weight_decay=weight_decay)
+ student_scheduler = torch.optim.lr_scheduler.ConstantLR(student_optimizer, factor=1.0, total_iters=1)
+ fake_score_scheduler = torch.optim.lr_scheduler.ConstantLR(fake_score_optimizer, factor=1.0, total_iters=1)
+
+ optimizers = {
+ "student": student_optimizer,
+ "fake_score": fake_score_optimizer,
+ "student_scheduler": student_scheduler,
+ "fake_score_scheduler": fake_score_scheduler,
+ }
+ prepare_items = [model, student_optimizer, fake_score_optimizer, student_scheduler, fake_score_scheduler]
+ if has_discriminator:
+ discriminator_optimizer = optimizer_class(discriminator.parameters(), lr=discriminator_lr, weight_decay=weight_decay)
+ discriminator_scheduler = torch.optim.lr_scheduler.ConstantLR(discriminator_optimizer, factor=1.0, total_iters=1)
+ optimizers["discriminator"] = discriminator_optimizer
+ optimizers["discriminator_scheduler"] = discriminator_scheduler
+ prepare_items.extend([discriminator_optimizer, discriminator_scheduler])
+
+ dataloader = torch.utils.data.DataLoader(dataset, shuffle=True, collate_fn=lambda x: x[0], num_workers=num_workers)
+ prepare_items.append(dataloader)
+
+ if enable_model_cpu_offload:
+ prepared = accelerator.prepare(*prepare_items[1:])
+ model.pipe.device = accelerator.device
+ offload_manager = OffloadTrainingManager(
+ model,
+ accelerator.device,
+ enable_optimizer_cpu_offload,
+ cpu_offload_split_threshold,
+ )
+ prepared_tail = list(prepared)
+ else:
+ model.to(device=accelerator.device)
+ prepared = accelerator.prepare(*prepare_items)
+ model = prepared[0]
+ prepared_tail = list(prepared[1:])
+
+ optimizers["student"] = prepared_tail.pop(0)
+ optimizers["fake_score"] = prepared_tail.pop(0)
+ optimizers["student_scheduler"] = prepared_tail.pop(0)
+ optimizers["fake_score_scheduler"] = prepared_tail.pop(0)
+ if has_discriminator:
+ optimizers["discriminator"] = prepared_tail.pop(0)
+ optimizers["discriminator_scheduler"] = prepared_tail.pop(0)
+ dataloader = prepared_tail.pop(0)
+
+ initialize_deepspeed_gradient_checkpointing(accelerator)
+ iteration = 0
+ for epoch_id in range(num_epochs):
+ for data in tqdm(dataloader):
+ with accelerator.accumulate(model):
+ if dataset.load_from_cache:
+ loss_map = model({}, inputs=data, iteration=iteration)
+ else:
+ loss_map = model(data, iteration=iteration)
+ loss = loss_map["total_loss"]
+ current_optimizers, current_schedulers = _dmd2_current_optimizers(config, optimizers, iteration)
+ accelerator.backward(loss)
+ if enable_model_cpu_offload:
+ offload_manager.after_backward()
+ if iteration % config.student_update_freq == 0 and config.student_grad_clip_norm is not None and config.student_grad_clip_norm > 0:
+ accelerator.clip_grad_norm_(_trainable_params(get_dmd2_student_model(accelerator.unwrap_model(model))), config.student_grad_clip_norm)
+ for optimizer in current_optimizers:
+ optimizer.step()
+ for scheduler in current_schedulers:
+ scheduler.step()
+ for optimizer in current_optimizers:
+ optimizer.zero_grad(set_to_none=True)
+ model_logger.on_step_end(accelerator, model, save_steps, loss=loss)
+ iteration += 1
+ if save_steps is None:
+ model_logger.on_epoch_end(accelerator, model, epoch_id)
+
+ model_logger.on_training_end(accelerator, model, save_steps)
\ No newline at end of file
diff --git a/diffsynth/diffusion/flow_match.py b/diffsynth/diffusion/flow_match.py
new file mode 100644
index 0000000000000000000000000000000000000000..350fb3d755d15a15de1c1e8a0bd43c6739756d23
--- /dev/null
+++ b/diffsynth/diffusion/flow_match.py
@@ -0,0 +1,442 @@
+import torch, math
+import numpy as np
+from typing_extensions import Literal
+
+
+class FlowMatchScheduler():
+
+ def __init__(self, template: Literal["FLUX.1", "Wan", "Qwen-Image", "FLUX.2", "Z-Image", "LTX-2", "Qwen-Image-Lightning", "ERNIE-Image", "ACE-Step", "Ideogram4", "Krea-2", "Boogu", "MiniMax-H3", "MiniMax-Music3", "LingBot-Video"] = "FLUX.1"):
+ self.set_timesteps_fn = {
+ "FLUX.1": FlowMatchScheduler.set_timesteps_flux,
+ "Wan": FlowMatchScheduler.set_timesteps_wan,
+ "Qwen-Image": FlowMatchScheduler.set_timesteps_qwen_image,
+ "FLUX.2": FlowMatchScheduler.set_timesteps_flux2,
+ "Z-Image": FlowMatchScheduler.set_timesteps_z_image,
+ "LTX-2": FlowMatchScheduler.set_timesteps_ltx2,
+ "Qwen-Image-Lightning": FlowMatchScheduler.set_timesteps_qwen_image_lightning,
+ "ERNIE-Image": FlowMatchScheduler.set_timesteps_ernie_image,
+ "ACE-Step": FlowMatchScheduler.set_timesteps_ace_step,
+ "HiDream-O1-Image": FlowMatchScheduler.set_timesteps_hidream_o1_image,
+ "Ideogram4": FlowMatchScheduler.set_timesteps_ideogram4,
+ "Krea-2": FlowMatchScheduler.set_timesteps_krea2,
+ "Boogu": FlowMatchScheduler.set_timesteps_boogu,
+ "MiniMax-H3": FlowMatchScheduler.set_timesteps_minimax_h3,
+ "MiniMax-Music3": FlowMatchScheduler.set_timesteps_minimax_music3,
+ "LingBot-Video": FlowMatchScheduler.set_timesteps_lingbot_video,
+ }.get(template, FlowMatchScheduler.set_timesteps_flux)
+ self.num_train_timesteps = 1000
+
+ @staticmethod
+ def set_timesteps_flux(num_inference_steps=100, denoising_strength=1.0, shift=None):
+ sigma_min = 0.003/1.002
+ sigma_max = 1.0
+ shift = 3 if shift is None else shift
+ num_train_timesteps = 1000
+ sigma_start = sigma_min + (sigma_max - sigma_min) * denoising_strength
+ sigmas = torch.linspace(sigma_start, sigma_min, num_inference_steps)
+ sigmas = shift * sigmas / (1 + (shift - 1) * sigmas)
+ timesteps = sigmas * num_train_timesteps
+ return sigmas, timesteps
+
+ @staticmethod
+ def set_timesteps_wan(num_inference_steps=100, denoising_strength=1.0, shift=None):
+ sigma_min = 0.0
+ sigma_max = 1.0
+ shift = 5 if shift is None else shift
+ num_train_timesteps = 1000
+ sigma_start = sigma_min + (sigma_max - sigma_min) * denoising_strength
+ sigmas = torch.linspace(sigma_start, sigma_min, num_inference_steps + 1)[:-1]
+ sigmas = shift * sigmas / (1 + (shift - 1) * sigmas)
+ timesteps = sigmas * num_train_timesteps
+ return sigmas, timesteps
+
+ @staticmethod
+ def _calculate_shift_qwen_image(image_seq_len, base_seq_len=256, max_seq_len=8192, base_shift=0.5, max_shift=0.9):
+ m = (max_shift - base_shift) / (max_seq_len - base_seq_len)
+ b = base_shift - m * base_seq_len
+ mu = image_seq_len * m + b
+ return mu
+
+ @staticmethod
+ def set_timesteps_qwen_image(num_inference_steps=100, denoising_strength=1.0, exponential_shift_mu=None, dynamic_shift_len=None):
+ sigma_min = 0.0
+ sigma_max = 1.0
+ num_train_timesteps = 1000
+ shift_terminal = 0.02
+ # Sigmas
+ sigma_start = sigma_min + (sigma_max - sigma_min) * denoising_strength
+ sigmas = torch.linspace(sigma_start, sigma_min, num_inference_steps + 1)[:-1]
+ # Mu
+ if exponential_shift_mu is not None:
+ mu = exponential_shift_mu
+ elif dynamic_shift_len is not None:
+ mu = FlowMatchScheduler._calculate_shift_qwen_image(dynamic_shift_len)
+ else:
+ mu = 0.8
+ sigmas = math.exp(mu) / (math.exp(mu) + (1 / sigmas - 1))
+ # Shift terminal
+ one_minus_z = 1 - sigmas
+ scale_factor = one_minus_z[-1] / (1 - shift_terminal)
+ sigmas = 1 - (one_minus_z / scale_factor)
+ # Timesteps
+ timesteps = sigmas * num_train_timesteps
+ return sigmas, timesteps
+
+ @staticmethod
+ def set_timesteps_lingbot_video(num_inference_steps=100, denoising_strength=1.0, shift=None, t_thresh=None, sigma_tail_steps=0):
+ sigma_min = 0.0
+ sigma_max = 1.0
+ shift = 5 if shift is None else shift
+ num_train_timesteps = 1000
+ sigma_start = sigma_min + (sigma_max - sigma_min) * denoising_strength
+ sigmas = torch.linspace(sigma_start, sigma_min, num_inference_steps + 1)[:-1]
+ sigmas = shift * sigmas / (1 + (shift - 1) * sigmas)
+ if t_thresh is not None:
+ # Refinement schedule: keep the sub-threshold part of the shifted grid, pin the first
+ # sigma exactly at t_thresh, then append extra low-noise steps that end at sigma_min.
+ sigmas = sigmas[sigmas <= t_thresh + 1e-6]
+ if sigmas.numel() == 0 or abs(float(sigmas[0]) - t_thresh) > 1e-6:
+ sigmas = torch.cat([torch.tensor([t_thresh], dtype=sigmas.dtype), sigmas])
+ if sigma_tail_steps > 0:
+ tail_start = float(sigmas[-1])
+ tail = torch.linspace(tail_start, min(sigma_min, tail_start), sigma_tail_steps + 2)[1:-1]
+ sigmas = torch.cat([sigmas, tail.to(dtype=sigmas.dtype)])
+ timesteps = sigmas * num_train_timesteps
+ return sigmas, timesteps
+
+ @staticmethod
+ def set_timesteps_qwen_image_lightning(num_inference_steps=100, denoising_strength=1.0, exponential_shift_mu=None, dynamic_shift_len=None):
+ sigma_min = 0.0
+ sigma_max = 1.0
+ num_train_timesteps = 1000
+ base_shift = math.log(3)
+ max_shift = math.log(3)
+ # Sigmas
+ sigma_start = sigma_min + (sigma_max - sigma_min) * denoising_strength
+ sigmas = torch.linspace(sigma_start, sigma_min, num_inference_steps + 1)[:-1]
+ # Mu
+ if exponential_shift_mu is not None:
+ mu = exponential_shift_mu
+ elif dynamic_shift_len is not None:
+ mu = FlowMatchScheduler._calculate_shift_qwen_image(dynamic_shift_len, base_shift=base_shift, max_shift=max_shift)
+ else:
+ mu = 0.8
+ sigmas = math.exp(mu) / (math.exp(mu) + (1 / sigmas - 1))
+ # Timesteps
+ timesteps = sigmas * num_train_timesteps
+ return sigmas, timesteps
+
+ @staticmethod
+ def compute_empirical_mu(image_seq_len, num_steps):
+ a1, b1 = 8.73809524e-05, 1.89833333
+ a2, b2 = 0.00016927, 0.45666666
+
+ if image_seq_len > 4300:
+ mu = a2 * image_seq_len + b2
+ return float(mu)
+
+ m_200 = a2 * image_seq_len + b2
+ m_10 = a1 * image_seq_len + b1
+
+ a = (m_200 - m_10) / 190.0
+ b = m_200 - 200.0 * a
+ mu = a * num_steps + b
+
+ return float(mu)
+
+ @staticmethod
+ def set_timesteps_flux2(num_inference_steps=100, denoising_strength=1.0, dynamic_shift_len=None):
+ sigma_min = 1 / num_inference_steps
+ sigma_max = 1.0
+ num_train_timesteps = 1000
+ sigma_start = sigma_min + (sigma_max - sigma_min) * denoising_strength
+ sigmas = torch.linspace(sigma_start, sigma_min, num_inference_steps)
+ if dynamic_shift_len is None:
+ # If you ask me why I set mu=0.8,
+ # I can only say that it yields better training results.
+ mu = 0.8
+ else:
+ mu = FlowMatchScheduler.compute_empirical_mu(dynamic_shift_len, num_inference_steps)
+ sigmas = math.exp(mu) / (math.exp(mu) + (1 / sigmas - 1))
+ timesteps = sigmas * num_train_timesteps
+ return sigmas, timesteps
+
+ @staticmethod
+ def set_timesteps_ernie_image(num_inference_steps=50, denoising_strength=1.0, shift=3.0):
+ sigma_min = 0.0
+ sigma_max = 1.0
+ num_train_timesteps = 1000
+ sigma_start = sigma_min + (sigma_max - sigma_min) * denoising_strength
+ sigmas = torch.linspace(sigma_start, sigma_min, num_inference_steps + 1)[:-1]
+ if shift is not None and shift != 1.0:
+ sigmas = shift * sigmas / (1 + (shift - 1) * sigmas)
+ timesteps = sigmas * num_train_timesteps
+ return sigmas, timesteps
+
+ @staticmethod
+ def set_timesteps_ace_step(num_inference_steps=8, denoising_strength=1.0, shift=3.0):
+ num_train_timesteps = 1000
+ sigma_start = denoising_strength
+ sigmas = torch.linspace(sigma_start, 0.0, num_inference_steps + 1)[:-1]
+ if shift is not None and shift != 1.0:
+ sigmas = shift * sigmas / (1 + (shift - 1) * sigmas)
+ timesteps = sigmas * num_train_timesteps
+ return sigmas, timesteps
+
+ @staticmethod
+ def set_timesteps_z_image(num_inference_steps=100, denoising_strength=1.0, shift=None, target_timesteps=None):
+ sigma_min = 0.0
+ sigma_max = 1.0
+ shift = 3 if shift is None else shift
+ num_train_timesteps = 1000
+ sigma_start = sigma_min + (sigma_max - sigma_min) * denoising_strength
+ sigmas = torch.linspace(sigma_start, sigma_min, num_inference_steps + 1)[:-1]
+ sigmas = shift * sigmas / (1 + (shift - 1) * sigmas)
+ timesteps = sigmas * num_train_timesteps
+ if target_timesteps is not None:
+ target_timesteps = target_timesteps.to(dtype=timesteps.dtype, device=timesteps.device)
+ for timestep in target_timesteps:
+ timestep_id = torch.argmin((timesteps - timestep).abs())
+ timesteps[timestep_id] = timestep
+ return sigmas, timesteps
+
+ @staticmethod
+ def set_timesteps_joyai_image(num_inference_steps=100, denoising_strength=1.0, shift=None):
+ sigma_min = 0.0
+ sigma_max = 1.0
+ shift = 4.0 if shift is None else shift
+ num_train_timesteps = 1000
+ sigma_start = sigma_min + (sigma_max - sigma_min) * denoising_strength
+ sigmas = torch.linspace(sigma_start, sigma_min, num_inference_steps + 1)[:-1]
+ sigmas = shift * sigmas / (1 + (shift - 1) * sigmas)
+ timesteps = sigmas * num_train_timesteps
+ return sigmas, timesteps
+
+ @staticmethod
+ def set_timesteps_hidream_o1_image(num_inference_steps=28, denoising_strength=1.0, shift=None, special_case=None, **kwargs):
+ num_train_timesteps = 1000
+ shift = 3.0 if shift is None else shift
+ if special_case == "dev":
+ timesteps_list = [
+ 999, 987, 974, 960, 945, 929, 913, 895, 877, 857, 836, 814, 790, 764, 737,
+ 707, 675, 640, 602, 560, 515, 464, 409, 347, 278, 199, 110, 8,
+ ]
+ sigmas = torch.tensor([t / 1000.0 for t in timesteps_list], dtype=torch.float32)
+ timesteps = torch.tensor(timesteps_list, dtype=torch.float32)
+ return sigmas, timesteps
+ else:
+ sigma_min = 0.0
+ sigma_max = 1.0
+ sigma_start = sigma_min + (sigma_max - sigma_min) * denoising_strength
+ sigmas = torch.linspace(sigma_start, sigma_min, num_inference_steps + 1)[:-1]
+ sigmas = shift * sigmas / (1 + (shift - 1) * sigmas)
+ timesteps = sigmas * num_train_timesteps
+ return sigmas, timesteps
+
+ @staticmethod
+ def set_timesteps_ideogram4(num_inference_steps=50, denoising_strength=1.0, image_resolution=(1024, 1024), mu=0.0, std=1.5):
+ num_pixels = image_resolution[0] * image_resolution[1]
+ known_pixels = 512 * 512
+ mean = mu + 0.5 * math.log(num_pixels / known_pixels)
+ logsnr_min = -15.0
+ logsnr_max = 18.0
+ t_min = 1.0 / (1 + math.exp(0.5 * logsnr_max))
+ t_max = 1.0 / (1 + math.exp(0.5 * logsnr_min))
+ step_intervals = torch.linspace(0.0, denoising_strength, num_inference_steps + 1, dtype=torch.float64)
+ sigmas = []
+ for i in range(num_inference_steps + 1):
+ z = torch.special.ndtri(step_intervals[i])
+ y = mean + std * z
+ t_ = torch.special.expit(y)
+ t_ = 1 - t_
+ t_ = t_.clamp(t_min, t_max)
+ sigmas.append(float(t_.to(torch.float32)))
+ sigmas = torch.tensor(sigmas, dtype=torch.float32)
+ one_minus_t = (1 - sigmas)[:-1].flip(0)
+ sigma_start = one_minus_t[0] * denoising_strength
+ if one_minus_t[0] > 0:
+ one_minus_t = one_minus_t * (sigma_start / one_minus_t[0])
+ sigmas = sigmas.flip(dims=(0,))
+ timesteps = sigmas[:-1]
+ sigmas = (1 - sigmas)[:-1]
+ return sigmas, timesteps
+
+ def set_timesteps_boogu(num_inference_steps=50, denoising_strength=1.0, sigmas=None):
+ if sigmas is not None:
+ sigmas = torch.tensor(sigmas, dtype=torch.float32)
+ timesteps = 1 - sigmas
+ return sigmas, timesteps
+ t_arr = np.linspace(1-denoising_strength, 1, num_inference_steps + 1, dtype=np.float32)
+ mu = 1.15
+ sigma = 1
+ eps = 1e-8
+ t1 = 1.0 - t_arr
+ t1 = np.clip(t1, eps, 1.0 - eps)
+ num = math.exp(mu)
+ denom = num + np.power(1.0 / t1 - 1.0, sigma)
+ y = num / denom
+ t_arr = 1.0 - y
+ timesteps = torch.from_numpy(t_arr).float()[:-1]
+ sigmas = 1 - timesteps
+ return sigmas, timesteps
+
+ @staticmethod
+ def set_timesteps_krea2(num_inference_steps=28, denoising_strength=1.0, dynamic_shift_len=None, y1=0.5, y2=1.15, mu=None):
+ x1 = 256
+ x2 = 6400
+ sigma = 1
+ ts = torch.linspace(denoising_strength, 0, num_inference_steps + 1)
+ if mu is None and dynamic_shift_len is None:
+ # Training
+ mu = 0.8
+ elif mu is None:
+ # Raw
+ slope = (y2 - y1) / (x2 - x1)
+ mu = slope * dynamic_shift_len + (y1 - slope * x1)
+ ts = math.exp(mu) / (math.exp(mu) + (1.0 / ts - 1.0) ** sigma)
+ sigmas, timesteps = ts[:-1], ts[:-1]
+ return sigmas, timesteps
+
+ @staticmethod
+ def set_timesteps_ltx2(num_inference_steps=100, denoising_strength=1.0, dynamic_shift_len=None, terminal=0.1, special_case=None):
+ num_train_timesteps = 1000
+ if special_case == "stage2":
+ sigmas = torch.Tensor([0.909375, 0.725, 0.421875])
+ elif special_case == "distilled_stage1":
+ sigmas = torch.Tensor([1.0, 0.99375, 0.9875, 0.98125, 0.975, 0.909375, 0.725, 0.421875])
+ else:
+ dynamic_shift_len = dynamic_shift_len or 4096
+ sigma_shift = FlowMatchScheduler._calculate_shift_qwen_image(
+ image_seq_len=dynamic_shift_len,
+ base_seq_len=1024,
+ max_seq_len=4096,
+ base_shift=0.95,
+ max_shift=2.05,
+ )
+ sigma_min = 0.0
+ sigma_max = 1.0
+ sigma_start = sigma_min + (sigma_max - sigma_min) * denoising_strength
+ sigmas = torch.linspace(sigma_start, sigma_min, num_inference_steps + 1)[:-1]
+ sigmas = math.exp(sigma_shift) / (math.exp(sigma_shift) + (1 / sigmas - 1))
+ # Shift terminal
+ one_minus_z = 1.0 - sigmas
+ scale_factor = one_minus_z[-1] / (1 - terminal)
+ sigmas = 1.0 - (one_minus_z / scale_factor)
+ timesteps = sigmas * num_train_timesteps
+ return sigmas, timesteps
+
+ @staticmethod
+ def set_timesteps_minimax_h3(num_inference_steps=50, denoising_strength=1.0, shift=2.22):
+ num_train_timesteps = 1000
+ base = torch.linspace(denoising_strength, 0.0, num_inference_steps+1, dtype=torch.float32)[:-1]
+ sigmas = shift * base / (1 + (shift - 1) * base)
+ timesteps = sigmas * num_train_timesteps
+ return sigmas, timesteps
+
+ @staticmethod
+ def set_timesteps_minimax_music3(num_inference_steps=30, denoising_strength=1.0):
+ num_train_timesteps = 1000
+ sigmas = torch.linspace(denoising_strength, denoising_strength / num_inference_steps, num_inference_steps)
+ timesteps = sigmas * num_train_timesteps
+ return sigmas, timesteps
+
+ def set_training_weight(self):
+ steps = 1000
+ x = self.sigmas * self.num_train_timesteps
+ y = torch.exp(-2 * ((x - steps / 2) / steps) ** 2)
+ y_shifted = y - y.min()
+ bsmntw_weighing = y_shifted * (steps / y_shifted.sum())
+ if len(self.timesteps) != 1000:
+ # This is an empirical formula.
+ bsmntw_weighing = bsmntw_weighing * (len(self.timesteps) / steps)
+ bsmntw_weighing = bsmntw_weighing + bsmntw_weighing[1]
+ self.linear_timesteps_weights = bsmntw_weighing
+
+ def set_timesteps(self, num_inference_steps=100, denoising_strength=1.0, training=False, **kwargs):
+ self.sigmas, self.timesteps = self.set_timesteps_fn(
+ num_inference_steps=num_inference_steps,
+ denoising_strength=denoising_strength,
+ **kwargs,
+ )
+ if training:
+ self.set_training_weight()
+ self.training = True
+ else:
+ self.training = False
+
+ def step(self, model_output, timestep, sample, to_final=False, **kwargs):
+ if isinstance(timestep, torch.Tensor):
+ timestep = timestep.cpu()
+ timestep_id = torch.argmin((self.timesteps - timestep).abs())
+ sigma = self.sigmas[timestep_id]
+ if to_final or timestep_id + 1 >= len(self.timesteps):
+ sigma_ = 0
+ else:
+ sigma_ = self.sigmas[timestep_id + 1]
+ prev_sample = sample + model_output * (sigma_ - sigma)
+ return prev_sample
+
+ def return_to_timestep(self, timestep, sample, sample_stablized):
+ if isinstance(timestep, torch.Tensor):
+ timestep = timestep.cpu()
+ timestep_id = torch.argmin((self.timesteps - timestep).abs())
+ sigma = self.sigmas[timestep_id]
+ model_output = (sample - sample_stablized) / sigma
+ return model_output
+
+ def add_noise(self, original_samples, noise, timestep):
+ if isinstance(timestep, torch.Tensor):
+ timestep = timestep.cpu()
+ timestep_id = torch.argmin((self.timesteps - timestep).abs())
+ sigma = self.sigmas[timestep_id]
+ sample = (1 - sigma) * original_samples + sigma * noise
+ return sample
+
+ def training_target(self, sample, noise, timestep):
+ target = noise - sample
+ return target
+
+ def training_weight(self, timestep):
+ timestep_id = torch.argmin((self.timesteps - timestep.to(self.timesteps.device)).abs())
+ weights = self.linear_timesteps_weights[timestep_id]
+ return weights
+
+
+class HiDreamO1FlashScheduler(FlowMatchScheduler):
+
+ def __init__(self, noise_scale_start=7.5, noise_scale_end=7.5, noise_clip_std=2.5):
+ self.set_timesteps_fn = HiDreamO1FlashScheduler.set_timesteps_hidream_o1_image_dev
+ self.num_train_timesteps = 1000
+ self.noise_clip_std = noise_clip_std
+ num_steps = 28
+ self.noise_scale_schedule = [
+ noise_scale_start + (noise_scale_end - noise_scale_start) * i / (num_steps - 1)
+ for i in range(num_steps)
+ ]
+
+ @staticmethod
+ def set_timesteps_hidream_o1_image_dev(**kwargs):
+ timesteps_list = [
+ 999, 987, 974, 960, 945, 929, 913, 895, 877, 857, 836, 814, 790, 764, 737,
+ 707, 675, 640, 602, 560, 515, 464, 409, 347, 278, 199, 110, 8,
+ ]
+ sigmas = torch.tensor([t / 1000.0 for t in timesteps_list], dtype=torch.float32)
+ timesteps = torch.tensor(timesteps_list, dtype=torch.float32)
+ return sigmas, timesteps
+
+ def clip_noise(self, noise):
+ if self.noise_clip_std > 0:
+ noise_std = noise.std().item()
+ clip_val = self.noise_clip_std * noise_std
+ noise = noise.clamp(min=-clip_val, max=clip_val)
+ return noise
+
+ def step(self, model_output, timestep, sample):
+ timestep_id = torch.argmin((self.timesteps - timestep).abs())
+ sigma = self.sigmas[timestep_id]
+ sigma_ = self.sigmas[timestep_id + 1] if timestep_id + 1 < len(self.sigmas) else 0
+ denoised = sample - model_output * sigma
+
+ noise = self.clip_noise(torch.randn(denoised.shape, device=denoised.device, dtype=denoised.dtype))
+ sample = sigma_ * noise * self.noise_scale_schedule[timestep_id] + (1.0 - sigma_) * denoised
+ return sample
diff --git a/diffsynth/diffusion/logger.py b/diffsynth/diffusion/logger.py
new file mode 100644
index 0000000000000000000000000000000000000000..02de8636c711e8d4820030e39495e229fd57430f
--- /dev/null
+++ b/diffsynth/diffusion/logger.py
@@ -0,0 +1,170 @@
+import json, os, torch
+from accelerate import Accelerator
+
+
+class TensorBoardLogger:
+ def __init__(self, log_dir):
+ from torch.utils.tensorboard import SummaryWriter
+ self.writer = SummaryWriter(log_dir=log_dir)
+ print(f"TensorBoard is enabled. Run `tensorboard --logdir={log_dir}` to visualize the training progress.")
+
+ def log(self, key, value, step):
+ self.writer.add_scalar(key, value, step)
+
+ def close(self):
+ if self.writer is not None:
+ self.writer.close()
+
+
+class SwanLabLogger:
+ def __init__(self, project_name="DiffSynth-Studio", log_dir=None):
+ import swanlab
+ project_name = os.environ.get("SWANLAB_PROJECT", project_name)
+ self.swanlab = swanlab
+ self.swanlab.init(project=project_name, logdir=log_dir)
+ print(f"SwanLab is enabled. Project: {project_name}")
+
+ def log(self, key, value, step):
+ self.swanlab.log({key: value}, step=step)
+
+ def close(self):
+ self.swanlab.finish()
+
+
+class WandbLogger:
+ def __init__(self, project_name="DiffSynth-Studio", log_dir=None):
+ import wandb
+ project_name = os.environ.get("WANDB_PROJECT", project_name)
+ if log_dir is not None:
+ os.makedirs(log_dir, exist_ok=True)
+ self.wandb = wandb
+ config = None
+ if log_dir is not None:
+ training_args_path = os.path.join(os.path.dirname(log_dir), "training_args.json")
+ if os.path.isfile(training_args_path):
+ with open(training_args_path, "r", encoding="utf-8") as handle:
+ config = json.load(handle)
+ self.run = self.wandb.init(project=project_name, dir=log_dir, config=config)
+ print(f"Wandb is enabled. Project: {project_name}")
+
+ def log(self, key, value, step):
+ if isinstance(value, torch.Tensor):
+ value = value.detach().float().mean().item()
+ self.wandb.log({key: value}, step=step)
+
+ def log_metrics(self, metrics, step):
+ values = {
+ key: value.detach().float().mean().item() if isinstance(value, torch.Tensor) else value
+ for key, value in metrics.items()
+ }
+ self.wandb.log(values, step=step)
+
+ def close(self):
+ self.wandb.finish()
+
+class CSVLogger:
+ def __init__(self, log_dir):
+ import csv
+ os.makedirs(log_dir, exist_ok=True)
+ self.path = os.path.join(log_dir, "loss.csv")
+ file_exists = os.path.isfile(self.path) and os.path.getsize(self.path) > 0
+ self.file = open(self.path, "a", encoding="utf-8", newline="")
+ self.writer = csv.DictWriter(self.file, fieldnames=["step", "key", "value"], lineterminator="\n")
+ if not file_exists:
+ self.writer.writeheader()
+ self.file.flush()
+
+ def log(self, key, value, step):
+ if isinstance(value, torch.Tensor):
+ value = value.detach().float().mean().item()
+ self.writer.writerow({"step": step, "key": key, "value": value})
+ self.file.flush()
+
+ def close(self):
+ if not self.file.closed:
+ self.file.close()
+
+
+class ModelLogger:
+ def __init__(
+ self, output_path, remove_prefix_in_ckpt=None, state_dict_converter=lambda x: x,
+ enable_tensorboard_log=False,
+ enable_swanlab_log=False, swanlab_project="DiffSynth-Studio",
+ enable_wandb_log=False, wandb_project="DiffSynth-Studio",
+ enable_csv_log=False,
+ ):
+ self.output_path = output_path
+ self.remove_prefix_in_ckpt = remove_prefix_in_ckpt
+ self.state_dict_converter = state_dict_converter
+ self.num_steps = 0
+ # Loggers
+ self.enable_tensorboard_log = enable_tensorboard_log
+ self.enable_swanlab_log = enable_swanlab_log
+ self.swanlab_project = swanlab_project
+ self.enable_wandb_log = enable_wandb_log
+ self.wandb_project = wandb_project
+ self.enable_csv_log = enable_csv_log
+ self.loggers = []
+ self.loggers_initialized = False
+
+ def init_loggers(self):
+ if self.enable_tensorboard_log:
+ self.loggers.append(TensorBoardLogger(os.path.join(self.output_path, "tensorboard_log")))
+ if self.enable_swanlab_log:
+ self.loggers.append(SwanLabLogger(project_name=self.swanlab_project, log_dir=os.path.join(self.output_path, "swanlab_log")))
+ if self.enable_wandb_log:
+ self.loggers.append(WandbLogger(project_name=self.wandb_project, log_dir=os.path.join(self.output_path, "wandb_log")))
+ if self.enable_csv_log:
+ self.loggers.append(CSVLogger(self.output_path))
+ self.loggers_initialized = True
+
+ def on_step_end(self, accelerator: Accelerator, model: torch.nn.Module, save_steps=None, **kwargs):
+ self.num_steps += 1
+ loss = kwargs.get("loss")
+ metrics = dict(kwargs.get("metrics") or {})
+ if loss is not None:
+ metrics["train/loss"] = loss
+ if metrics:
+ self.log_metrics(accelerator, metrics, self.num_steps)
+ if save_steps is not None and self.num_steps % save_steps == 0:
+ self.save_model(accelerator, model, f"step-{self.num_steps}.safetensors")
+
+ def log_metrics(self, accelerator: Accelerator, metrics, step):
+ if not accelerator.is_main_process:
+ return
+ if not self.loggers_initialized:
+ self.init_loggers()
+ for logger in self.loggers:
+ if hasattr(logger, "log_metrics"):
+ logger.log_metrics(metrics, step)
+ else:
+ for key, value in metrics.items():
+ logger.log(key, value, step)
+
+ def on_epoch_end(self, accelerator: Accelerator, model: torch.nn.Module, epoch_id):
+ self.save_model(accelerator, model, f"epoch-{epoch_id}.safetensors")
+
+ def on_training_end(self, accelerator: Accelerator, model: torch.nn.Module, save_steps=None, skip_model_saving=False):
+ if not skip_model_saving and save_steps is not None and self.num_steps % save_steps != 0:
+ self.save_model(accelerator, model, f"step-{self.num_steps}.safetensors")
+ for logger in self.loggers:
+ logger.close()
+
+ def save_model(self, accelerator: Accelerator, model: torch.nn.Module, file_name):
+ accelerator.wait_for_everyone()
+ deepspeed_plugin = getattr(accelerator.state, "deepspeed_plugin", None)
+ deepspeed_config = getattr(deepspeed_plugin, "deepspeed_config", {}) or {}
+ zero_stage = deepspeed_config.get("zero_optimization", {}).get("stage")
+ if zero_stage == 3 and hasattr(model, "_zero3_consolidated_16bit_state_dict"):
+ # LoRA runs freeze the base model. Excluding frozen parameters here
+ # avoids materializing a complete 33B CPU state dict before the normal
+ # export_trainable_state_dict filter is applied.
+ state_dict = model._zero3_consolidated_16bit_state_dict(exclude_frozen_parameters=True)
+ else:
+ state_dict = accelerator.get_state_dict(model)
+ if accelerator.is_main_process:
+ state_dict = accelerator.unwrap_model(model).export_trainable_state_dict(state_dict, remove_prefix=self.remove_prefix_in_ckpt)
+ state_dict = self.state_dict_converter(state_dict)
+ os.makedirs(self.output_path, exist_ok=True)
+ path = os.path.join(self.output_path, file_name)
+ accelerator.save(state_dict, path, safe_serialization=True)
diff --git a/diffsynth/diffusion/loss.py b/diffsynth/diffusion/loss.py
new file mode 100644
index 0000000000000000000000000000000000000000..b9c53f63c26fb38282cec3ba225e706afe086302
--- /dev/null
+++ b/diffsynth/diffusion/loss.py
@@ -0,0 +1,225 @@
+from .base_pipeline import BasePipeline
+import torch
+
+
+def FlowMatchSFTLoss(pipe: BasePipeline, **inputs):
+ if "lora" in inputs:
+ # Image-to-LoRA models need to load lora here.
+ pipe.clear_lora(verbose=0)
+ pipe.load_lora(pipe.dit, state_dict=inputs["lora"], hotload=True, verbose=0)
+
+ max_timestep_boundary = int(inputs.get("max_timestep_boundary", 1) * len(pipe.scheduler.timesteps))
+ min_timestep_boundary = int(inputs.get("min_timestep_boundary", 0) * len(pipe.scheduler.timesteps))
+
+ timestep_id = torch.randint(min_timestep_boundary, max_timestep_boundary, (1,))
+ timestep = pipe.scheduler.timesteps[timestep_id].to(dtype=pipe.torch_dtype, device=pipe.device)
+
+ noise = torch.randn_like(inputs["input_latents"]) * inputs.get("noise_scale", 1.0)
+ inputs["latents"] = pipe.scheduler.add_noise(inputs["input_latents"], noise, timestep)
+ training_target = pipe.scheduler.training_target(inputs["input_latents"], noise, timestep)
+
+ if "first_frame_latents" in inputs:
+ inputs["latents"][:, :, 0:1] = inputs["first_frame_latents"]
+
+ models = {name: getattr(pipe, name) for name in pipe.in_iteration_models}
+ noise_pred = pipe.model_fn(**models, **inputs, timestep=timestep)
+
+ if "first_frame_latents" in inputs:
+ noise_pred = noise_pred[:, :, 1:]
+ training_target = training_target[:, :, 1:]
+
+ loss = torch.nn.functional.mse_loss(noise_pred.float(), training_target.float())
+ loss = loss * pipe.scheduler.training_weight(timestep)
+ return loss
+
+
+def FlowMatchSFTAudioVideoLoss(pipe: BasePipeline, **inputs):
+ max_timestep_boundary = int(inputs.get("max_timestep_boundary", 1) * len(pipe.scheduler.timesteps))
+ min_timestep_boundary = int(inputs.get("min_timestep_boundary", 0) * len(pipe.scheduler.timesteps))
+
+ timestep_id = torch.randint(min_timestep_boundary, max_timestep_boundary, (1,))
+ timestep = pipe.scheduler.timesteps[timestep_id].to(dtype=pipe.torch_dtype, device=pipe.device)
+
+ # video
+ noise = torch.randn_like(inputs["input_latents"])
+ inputs["video_latents"] = pipe.scheduler.add_noise(inputs["input_latents"], noise, timestep)
+ training_target = pipe.scheduler.training_target(inputs["input_latents"], noise, timestep)
+
+ # audio
+ if inputs.get("audio_input_latents") is not None:
+ audio_noise = torch.randn_like(inputs["audio_input_latents"])
+ inputs["audio_latents"] = pipe.scheduler.add_noise(inputs["audio_input_latents"], audio_noise, timestep)
+ training_target_audio = pipe.scheduler.training_target(inputs["audio_input_latents"], audio_noise, timestep)
+
+ models = {name: getattr(pipe, name) for name in pipe.in_iteration_models}
+ noise_pred, noise_pred_audio = pipe.model_fn(**models, **inputs, timestep=timestep)
+
+ loss = torch.nn.functional.mse_loss(noise_pred.float(), training_target.float())
+ loss = loss * pipe.scheduler.training_weight(timestep)
+ if inputs.get("audio_input_latents") is not None:
+ loss_audio = torch.nn.functional.mse_loss(noise_pred_audio.float(), training_target_audio.float())
+ loss_audio = loss_audio * pipe.scheduler.training_weight(timestep)
+ loss = loss + loss_audio
+ return loss
+
+
+def FlowMatchSFTMiniMaxH3AudioVideoLoss(pipe: BasePipeline, training_cfg_scale: float = 1.0, audio_loss_weight: float = 1.0, inputs_nega: dict | None = None, **inputs):
+ max_timestep_boundary = int(inputs.get("max_timestep_boundary", 1) * len(pipe.scheduler.timesteps))
+ min_timestep_boundary = int(inputs.get("min_timestep_boundary", 0) * len(pipe.scheduler.timesteps))
+
+ timestep_id = torch.randint(min_timestep_boundary, max_timestep_boundary, (1,))
+ timestep_video = pipe.scheduler.timesteps[timestep_id].to(dtype=torch.float32, device=pipe.device)
+ timestep_audio = pipe.scheduler_audio.timesteps[timestep_id].to(dtype=torch.float32, device=pipe.device)
+
+ noise = torch.randn_like(inputs["input_latents"])
+ inputs["video_latents"] = pipe.scheduler.add_noise(inputs["input_latents"], noise, timestep_video)
+ training_target = pipe.scheduler.training_target(inputs["input_latents"], noise, timestep_video)
+
+ if "audio_input_latents" in inputs:
+ audio_noise = torch.randn_like(inputs["audio_input_latents"])
+ inputs["audio_latents"] = pipe.scheduler_audio.add_noise(inputs["audio_input_latents"], audio_noise, timestep_audio)
+ training_target_audio = pipe.scheduler_audio.training_target(inputs["audio_input_latents"], audio_noise, timestep_audio)
+
+ models = {name: getattr(pipe, name) for name in pipe.in_iteration_models}
+ if training_cfg_scale > 1.0:
+ if not inputs_nega:
+ raise ValueError(
+ "MiniMax-H3 CFG-aware training requires unconditional inputs. "
+ "When using split training, rebuild the data cache with the same "
+ "--training_cfg_scale value."
+ )
+ inputs_uncond = {**inputs, **inputs_nega}
+ inputs_uncond["use_gradient_checkpointing"] = False
+ inputs_uncond["use_gradient_checkpointing_offload"] = False
+ with torch.no_grad():
+ noise_pred_uncond, noise_pred_audio_uncond = pipe.model_fn(
+ **models, **inputs_uncond,
+ timestep_video=timestep_video, timestep_audio=timestep_audio,
+ )
+
+ noise_pred, noise_pred_audio = pipe.model_fn(
+ **models, **inputs,
+ timestep_video=timestep_video, timestep_audio=timestep_audio,
+ )
+
+ if training_cfg_scale > 1.0:
+ # The checkpoint's conditional prediction has CFG distilled into it.
+ # Rearrange the CFG equation to recover the raw velocity fitted to the
+ # standard flow-matching target, using the current model as the teacher.
+ noise_pred = (noise_pred + (training_cfg_scale - 1.0) * noise_pred_uncond) / training_cfg_scale
+ noise_pred_audio = (noise_pred_audio + (training_cfg_scale - 1.0) * noise_pred_audio_uncond) / training_cfg_scale
+
+ loss_video = torch.nn.functional.mse_loss(noise_pred.float(), training_target.float())
+ loss_video = loss_video * pipe.scheduler.training_weight(timestep_video)
+ loss = loss_video
+ loss_components = {"train/video_loss": loss_video}
+ if "audio_input_latents" in inputs:
+ loss_audio = torch.nn.functional.mse_loss(noise_pred_audio.float(), training_target_audio.float())
+ loss_audio = loss_audio * pipe.scheduler_audio.training_weight(timestep_audio) * audio_loss_weight
+ loss = loss + loss_audio
+ loss_components["train/audio_loss"] = loss_audio
+ # Backpropagation still uses the exact same summed scalar as before. The
+ # The runner detaches these values for ordinary logging. Keeping the live
+ # component scalars here also permits optional diagnostic grad-norm passes;
+ # the real optimizer update still uses only the summed scalar above.
+ return loss, loss_components
+
+
+def DirectDistillLoss(pipe: BasePipeline, **inputs):
+ pipe.scheduler.set_timesteps(inputs["num_inference_steps"])
+ pipe.scheduler.training = True
+ models = {name: getattr(pipe, name) for name in pipe.in_iteration_models}
+ for progress_id, timestep in enumerate(pipe.scheduler.timesteps):
+ timestep = timestep.unsqueeze(0).to(dtype=pipe.torch_dtype, device=pipe.device)
+ noise_pred = pipe.model_fn(**models, **inputs, timestep=timestep, progress_id=progress_id)
+ inputs["latents"] = pipe.step(pipe.scheduler, progress_id=progress_id, noise_pred=noise_pred, **inputs)
+ loss = torch.nn.functional.mse_loss(inputs["latents"].float(), inputs["input_latents"].float())
+ return loss
+
+
+class TrajectoryImitationLoss(torch.nn.Module):
+ def __init__(self):
+ super().__init__()
+ self.initialized = False
+
+ def initialize(self, device):
+ import lpips # TODO: remove it
+ self.loss_fn = lpips.LPIPS(net='alex').to(device)
+ self.initialized = True
+
+ def fetch_trajectory(self, pipe: BasePipeline, timesteps_student, inputs_shared, inputs_posi, inputs_nega, num_inference_steps, cfg_scale):
+ trajectory = [inputs_shared["latents"].clone()]
+
+ pipe.scheduler.set_timesteps(num_inference_steps, target_timesteps=timesteps_student)
+ models = {name: getattr(pipe, name) for name in pipe.in_iteration_models}
+ for progress_id, timestep in enumerate(pipe.scheduler.timesteps):
+ timestep = timestep.unsqueeze(0).to(dtype=pipe.torch_dtype, device=pipe.device)
+ noise_pred = pipe.cfg_guided_model_fn(
+ pipe.model_fn, cfg_scale,
+ inputs_shared, inputs_posi, inputs_nega,
+ **models, timestep=timestep, progress_id=progress_id
+ )
+ inputs_shared["latents"] = pipe.step(pipe.scheduler, progress_id=progress_id, noise_pred=noise_pred.detach(), **inputs_shared)
+
+ trajectory.append(inputs_shared["latents"].clone())
+ return pipe.scheduler.timesteps, trajectory
+
+ def align_trajectory(self, pipe: BasePipeline, timesteps_teacher, trajectory_teacher, inputs_shared, inputs_posi, inputs_nega, num_inference_steps, cfg_scale):
+ loss = 0
+ pipe.scheduler.set_timesteps(num_inference_steps, training=True)
+ models = {name: getattr(pipe, name) for name in pipe.in_iteration_models}
+ for progress_id, timestep in enumerate(pipe.scheduler.timesteps):
+ timestep = timestep.unsqueeze(0).to(dtype=pipe.torch_dtype, device=pipe.device)
+
+ progress_id_teacher = torch.argmin((timesteps_teacher - timestep).abs())
+ inputs_shared["latents"] = trajectory_teacher[progress_id_teacher]
+
+ noise_pred = pipe.cfg_guided_model_fn(
+ pipe.model_fn, cfg_scale,
+ inputs_shared, inputs_posi, inputs_nega,
+ **models, timestep=timestep, progress_id=progress_id
+ )
+
+ sigma = pipe.scheduler.sigmas[progress_id]
+ sigma_ = 0 if progress_id + 1 >= len(pipe.scheduler.timesteps) else pipe.scheduler.sigmas[progress_id + 1]
+ if progress_id + 1 >= len(pipe.scheduler.timesteps):
+ latents_ = trajectory_teacher[-1]
+ else:
+ progress_id_teacher = torch.argmin((timesteps_teacher - pipe.scheduler.timesteps[progress_id + 1]).abs())
+ latents_ = trajectory_teacher[progress_id_teacher]
+
+ denom = sigma_ - sigma
+ denom = torch.sign(denom) * torch.clamp(denom.abs(), min=1e-6)
+ target = (latents_ - inputs_shared["latents"]) / denom
+ loss = loss + torch.nn.functional.mse_loss(noise_pred.float(), target.float()) * pipe.scheduler.training_weight(timestep)
+ return loss
+
+ def compute_regularization(self, pipe: BasePipeline, trajectory_teacher, inputs_shared, inputs_posi, inputs_nega, num_inference_steps, cfg_scale):
+ inputs_shared["latents"] = trajectory_teacher[0]
+ pipe.scheduler.set_timesteps(num_inference_steps)
+ models = {name: getattr(pipe, name) for name in pipe.in_iteration_models}
+ for progress_id, timestep in enumerate(pipe.scheduler.timesteps):
+ timestep = timestep.unsqueeze(0).to(dtype=pipe.torch_dtype, device=pipe.device)
+ noise_pred = pipe.cfg_guided_model_fn(
+ pipe.model_fn, cfg_scale,
+ inputs_shared, inputs_posi, inputs_nega,
+ **models, timestep=timestep, progress_id=progress_id
+ )
+ inputs_shared["latents"] = pipe.step(pipe.scheduler, progress_id=progress_id, noise_pred=noise_pred.detach(), **inputs_shared)
+
+ image_pred = pipe.vae_decoder(inputs_shared["latents"])
+ image_real = pipe.vae_decoder(trajectory_teacher[-1])
+ loss = self.loss_fn(image_pred.float(), image_real.float())
+ return loss
+
+ def forward(self, pipe: BasePipeline, inputs_shared, inputs_posi, inputs_nega):
+ if not self.initialized:
+ self.initialize(pipe.device)
+ with torch.no_grad():
+ pipe.scheduler.set_timesteps(8)
+ timesteps_teacher, trajectory_teacher = self.fetch_trajectory(inputs_shared["teacher"], pipe.scheduler.timesteps, inputs_shared, inputs_posi, inputs_nega, 50, 2)
+ timesteps_teacher = timesteps_teacher.to(dtype=pipe.torch_dtype, device=pipe.device)
+ loss_1 = self.align_trajectory(pipe, timesteps_teacher, trajectory_teacher, inputs_shared, inputs_posi, inputs_nega, 8, 1)
+ loss_2 = self.compute_regularization(pipe, trajectory_teacher, inputs_shared, inputs_posi, inputs_nega, 8, 1)
+ loss = loss_1 + loss_2
+ return loss
diff --git a/diffsynth/diffusion/parsers.py b/diffsynth/diffusion/parsers.py
new file mode 100644
index 0000000000000000000000000000000000000000..6223e86a6b3a503af312cb7b6fb7706a02768d71
--- /dev/null
+++ b/diffsynth/diffusion/parsers.py
@@ -0,0 +1,140 @@
+import argparse
+
+
+def add_dataset_base_config(parser: argparse.ArgumentParser):
+ parser.add_argument("--dataset_base_path", type=str, default="", required=True, help="Base path of the dataset.")
+ parser.add_argument("--dataset_metadata_path", type=str, default=None, help="Path to the metadata file of the dataset.")
+ parser.add_argument("--dataset_repeat", type=int, default=1, help="Number of times to repeat the dataset per epoch.")
+ parser.add_argument("--dataset_num_workers", type=int, default=0, help="Number of workers for data loading.")
+ parser.add_argument("--data_file_keys", type=str, default="image,video", help="Data file keys in the metadata. Comma-separated.")
+ parser.add_argument("--dataset_cache_file_list", type=str, default=None, help="Optional newline-delimited .pth manifest. When set, train from this cache snapshot instead of recursively scanning every .pth under dataset_base_path.")
+ parser.add_argument("--validation_cache_file_list", type=str, default=None, help="Optional newline-delimited validation .pth manifest. Validation samples are never added to the training dataset.")
+ return parser
+
+def add_image_size_config(parser: argparse.ArgumentParser):
+ parser.add_argument("--height", type=int, default=None, help="Height of images. Leave `height` and `width` empty to enable dynamic resolution.")
+ parser.add_argument("--width", type=int, default=None, help="Width of images. Leave `height` and `width` empty to enable dynamic resolution.")
+ parser.add_argument("--max_pixels", type=int, default=1024*1024, help="Maximum number of pixels per frame, used for dynamic resolution.")
+ return parser
+
+def add_video_size_config(parser: argparse.ArgumentParser):
+ parser.add_argument("--height", type=int, default=None, help="Height of images. Leave `height` and `width` empty to enable dynamic resolution.")
+ parser.add_argument("--width", type=int, default=None, help="Width of images. Leave `height` and `width` empty to enable dynamic resolution.")
+ parser.add_argument("--max_pixels", type=int, default=1024*1024, help="Maximum number of pixels per frame, used for dynamic resolution.")
+ parser.add_argument("--num_frames", type=int, default=81, help="Number of frames per video. Frames are sampled from the video prefix.")
+ return parser
+
+def add_model_config(parser: argparse.ArgumentParser):
+ parser.add_argument("--model_paths", type=str, default=None, help="Paths to load models. In JSON format.")
+ parser.add_argument("--model_id_with_origin_paths", type=str, default=None, help="Model ID with origin paths, e.g., Wan-AI/Wan2.1-T2V-1.3B:diffusion_pytorch_model*.safetensors. Comma-separated.")
+ parser.add_argument("--extra_inputs", default=None, type=str, help="Additional model inputs, comma-separated.")
+ parser.add_argument("--fp8_models", default=None, type=str, help="Models with FP8 precision, comma-separated.")
+ parser.add_argument("--offload_models", default=None, type=str, help="Models with offload, comma-separated. Only used in splited training.")
+ parser.add_argument("--resume_from_checkpoint", default=None, type=str, help="Resume training from checkpoint file. Only single model training is supported.")
+ return parser
+
+def add_quant_config(parser: argparse.ArgumentParser):
+ parser.add_argument("--quant_options", type=str, default=None, help="Dynamically quantize loaded models. Semicolon-separated entries, each `:[/]`, where `` matches an entry in `--model_paths`/`--model_id_with_origin_paths`, `method` is a registered method (e.g. bitsandbytes_nf4), and `exclude_modules` optionally lists layers kept in full precision.")
+ return parser
+
+def add_training_config(parser: argparse.ArgumentParser):
+ parser.add_argument("--learning_rate", type=float, default=1e-4, help="Learning rate.")
+ parser.add_argument("--num_epochs", type=int, default=1, help="Number of epochs.")
+ parser.add_argument("--trainable_models", type=str, default=None, help="Models to train, e.g., dit, vae, text_encoder.")
+ parser.add_argument("--find_unused_parameters", default=False, action="store_true", help="Whether to find unused parameters in DDP.")
+ parser.add_argument("--weight_decay", type=float, default=0.01, help="Weight decay.")
+ parser.add_argument("--task", type=str, default="sft", required=False, help="Task type.")
+ parser.add_argument("--customized_optimizer", type=str, default=None, help="Customized optimizer, e.g., `bitsandbytes.optim.Adam8bit` and `torch.optim.Adam`. The default optimizer is `torch.optim.AdamW`.")
+ parser.add_argument("--max_train_steps", type=int, default=None, help="Optional hard limit on optimizer steps, useful for end-to-end smoke tests.")
+ parser.add_argument("--skip_model_saving", default=False, action="store_true", help="Run training without writing model checkpoints. Intended for smoke tests and profiling.")
+ parser.add_argument("--validation_steps", type=int, default=None, help="Run validation every N optimizer steps. Disabled when unset.")
+ parser.add_argument("--validation_at_start", default=False, action="store_true", help="Run validation once at global step 0 before the first optimizer update.")
+ parser.add_argument("--validation_num_inference_steps", type=int, default=50, help="Number of full denoising steps per validation sample.")
+ parser.add_argument("--validation_seed", type=int, default=42, help="Base noise seed for deterministic full-denoising validation.")
+ parser.add_argument("--validation_flow_shift", type=float, default=12.0, help="Video inference flow shift used during validation.")
+ parser.add_argument("--validation_audio_flow_shift", type=float, default=3.0, help="Audio inference flow shift used during validation.")
+ parser.add_argument("--validation_cfg_scale", type=float, default=1.0, help="Classifier-free guidance scale used during full-denoising validation.")
+ parser.add_argument("--validation_save_outputs", default=False, action="store_true", help="Decode and save each generated validation sample under output_path/validation.")
+ parser.add_argument("--log_component_grad_norms", default=False, action="store_true", help="Log total, video-loss, and audio-loss parameter gradient norms. Component norms require two diagnostic forward/backward passes and do not affect optimizer updates.")
+ parser.add_argument("--grad_norm_steps", type=int, default=1, help="Measure component gradient norms every N optimizer steps when --log_component_grad_norms is enabled.")
+ parser.add_argument("--skip_existing_cache", default=False, action="store_true", help="For data processing only: keep existing rank/local-index .pth files and process only missing outputs. Resume with the same world size and preprocessing settings.")
+ parser.add_argument("--skip_data_process_oom", default=False, action="store_true", help="For data processing only: record CUDA OOM samples and continue with the next sample instead of terminating the worker.")
+ parser.add_argument("--retry_data_process_oom", default=False, action="store_true", help="For data processing only: retry samples that already have an OOM marker instead of skipping them during cache resume.")
+ parser.add_argument("--skip_data_process_errors", default=False, action="store_true", help="For data processing only: record per-sample media loading errors and continue instead of terminating the worker.")
+ parser.add_argument("--retry_data_process_errors", default=False, action="store_true", help="For data processing only: retry samples that already have a data-error marker during cache resume.")
+ return parser
+
+def add_output_config(parser: argparse.ArgumentParser):
+ parser.add_argument("--output_path", type=str, default="./models", help="Output save path.")
+ parser.add_argument("--remove_prefix_in_ckpt", type=str, default="pipe.dit.", help="Remove prefix in ckpt.")
+ parser.add_argument("--save_steps", type=int, default=None, help="Number of checkpoint saving invervals. If None, checkpoints will be saved every epoch.")
+ return parser
+
+def add_lora_config(parser: argparse.ArgumentParser):
+ parser.add_argument("--lora_base_model", type=str, default=None, help="Which model LoRA is added to.")
+ parser.add_argument("--lora_target_modules", type=str, default="q,k,v,o,ffn.0,ffn.2", help="Which layers LoRA is added to.")
+ parser.add_argument("--lora_rank", type=int, default=32, help="Rank of LoRA.")
+ parser.add_argument("--lora_checkpoint", type=str, default=None, help="Path to the LoRA checkpoint. If provided, LoRA will be loaded from this checkpoint.")
+ parser.add_argument("--preset_lora_path", type=str, default=None, help="Path to the preset LoRA checkpoint. If provided, this LoRA will be fused to the base model.")
+ parser.add_argument("--preset_lora_model", type=str, default=None, help="Which model the preset LoRA is fused to.")
+ return parser
+
+def add_gradient_config(parser: argparse.ArgumentParser):
+ parser.add_argument("--use_gradient_checkpointing", default=False, action="store_true", help="Whether to use gradient checkpointing.")
+ parser.add_argument("--use_gradient_checkpointing_offload", default=False, action="store_true", help="Whether to offload gradient checkpointing to CPU memory.")
+ parser.add_argument("--gradient_accumulation_steps", type=int, default=1, help="Gradient accumulation steps.")
+ return parser
+
+def add_template_model_config(parser: argparse.ArgumentParser):
+ parser.add_argument("--template_model_id_or_path", type=str, default=None, help="Model ID of path of template models.")
+ parser.add_argument("--enable_lora_hot_loading", default=False, action="store_true", help="Whether to enable LoRA hot-loading. Only available for image-to-lora models.")
+ return parser
+
+def add_offload_training_config(parser: argparse.ArgumentParser):
+ parser.add_argument("--enable_model_cpu_offload", default=False, action="store_true", help="Enable layer offload training. Weights are kept on CPU and loaded to GPU one layer at a time.")
+ parser.add_argument("--enable_optimizer_cpu_offload", default=False, action="store_true", help="When --enable_model_cpu_offload is enabled, run optimizer on CPU. All params are offloaded to CPU. Default is False (trainable params stay on GPU, optimizer on GPU).")
+ parser.add_argument("--cpu_offload_split_threshold", type=int, default=None, help="Experimental! When --enable_model_cpu_offload is enabled, modules with total params above this threshold (in MB) are recursively split into children. None means offload every leaf module directly. Default: None.")
+ return parser
+
+def add_logger_config(parser: argparse.ArgumentParser):
+ parser.add_argument("--enable_tensorboard_log", default=False, action="store_true", help="Enable tensorboard for logging.")
+ parser.add_argument("--enable_swanlab_log", default=False, action="store_true", help="Enable swanlab for logging.")
+ parser.add_argument("--swanlab_project", type=str, default="DiffSynth-Studio", help="SwanLab project name.")
+ parser.add_argument("--enable_wandb_log", default=False, action="store_true", help="Enable wandb for logging.")
+ parser.add_argument("--wandb_project", type=str, default="DiffSynth-Studio", help="Wandb project name.")
+ parser.add_argument("--enable_csv_log", default=False, action="store_true", help="Enable CSV loss logging.")
+ return parser
+
+def add_dmd2_config(parser: argparse.ArgumentParser):
+ parser.add_argument("--dmd2_student_update_freq", type=int, default=5, help="Update student once every N DMD2 iterations.")
+ parser.add_argument("--dmd2_student_sample_steps", type=int, default=4, help="Number of distilled student sampling steps.")
+ parser.add_argument("--dmd2_student_sample_type", type=str, default="sde", choices=["sde", "ode"], help="Student sampling type.")
+ parser.add_argument("--dmd2_student_schedule", type=str, default="uniform", choices=["uniform"], help="Student sigma schedule.")
+ parser.add_argument("--dmd2_student_t_list", type=str, default=None, help="Optional student sigma schedule, including the final 0.")
+ parser.add_argument("--dmd2_matching_t_min", type=float, default=0.001, help="Minimum matching sigma sampled for DMD2.")
+ parser.add_argument("--dmd2_matching_t_max", type=float, default=0.999, help="Maximum matching sigma sampled for DMD2.")
+ parser.add_argument("--dmd2_matching_t_sampling", type=str, default="uniform", choices=["uniform", "logitnormal"], help="Sample matching sigma.")
+ parser.add_argument("--dmd2_matching_t_mean", type=float, default=0.0, help="Mean for logitnormal matching timestep sampling.")
+ parser.add_argument("--dmd2_matching_t_std", type=float, default=1.0, help="Std for logitnormal matching timestep sampling.")
+ parser.add_argument("--dmd2_gan_loss_weight", type=float, default=0.03, help="Generator GAN loss weight.")
+ parser.add_argument("--dmd2_gan_r1_reg_weight", type=float, default=0.0, help="Approximate R1 regularization weight for the discriminator.")
+ parser.add_argument("--dmd2_gan_r1_reg_alpha", type=float, default=0.1, help="Noise scale for approximate R1 regularization.")
+ parser.add_argument("--dmd2_fake_score_learning_rate", type=float, default=None, help="Learning rate for the fake score model.")
+ parser.add_argument("--dmd2_discriminator_learning_rate", type=float, default=None, help="Learning rate for the discriminator.")
+ parser.add_argument("--dmd2_feature_indices", type=str, default=None, help="Model feature indices used by the DMD2 discriminator.")
+ parser.add_argument("--dmd2_teacher_cfg_scale", type=float, default=1.0, help="CFG scale.")
+ parser.add_argument("--dmd2_student_grad_clip_norm", type=float, default=10.0, help="Clip student gradients to this norm.")
+ return parser
+
+def add_general_config(parser: argparse.ArgumentParser):
+ parser = add_dataset_base_config(parser)
+ parser = add_model_config(parser)
+ parser = add_quant_config(parser)
+ parser = add_training_config(parser)
+ parser = add_output_config(parser)
+ parser = add_lora_config(parser)
+ parser = add_gradient_config(parser)
+ parser = add_template_model_config(parser)
+ parser = add_offload_training_config(parser)
+ parser = add_logger_config(parser)
+ return parser
diff --git a/diffsynth/diffusion/runner.py b/diffsynth/diffusion/runner.py
new file mode 100644
index 0000000000000000000000000000000000000000..f06089cb1267a50aeb62e218a8f47fe195798128
--- /dev/null
+++ b/diffsynth/diffusion/runner.py
@@ -0,0 +1,645 @@
+import gc, os, json, time, torch, importlib, traceback
+from tqdm import tqdm
+from accelerate import Accelerator
+from .training_module import DiffusionTrainingModule
+from .logger import ModelLogger
+from diffsynth.core import OffloadTrainingManager
+
+
+def get_optimizer_class(customized_optimizer=None):
+ if customized_optimizer is None:
+ return torch.optim.AdamW
+ else:
+ module_name, class_name = customized_optimizer.rsplit(".", 1)
+ module = importlib.import_module(module_name)
+ print(f"Customized opimizer `{customized_optimizer}` imported.")
+ return getattr(module, class_name)
+
+
+def _current_global_grad_norm(accelerator, model):
+ """Return the L2 norm of the current trainable gradients, including ZeRO-3 shards."""
+ parameters = [parameter for parameter in model.parameters() if parameter.requires_grad]
+ uses_zero3 = any(hasattr(parameter, "ds_id") for parameter in parameters)
+ squared_norm = torch.zeros((), dtype=torch.float64, device=accelerator.device)
+ found_gradient = False
+ if uses_zero3:
+ from deepspeed.utils import safe_get_local_grad
+ for parameter in parameters:
+ if not hasattr(parameter, "ds_id"):
+ continue
+ gradient = safe_get_local_grad(parameter)
+ if gradient is not None:
+ squared_norm += gradient.detach().double().square().sum()
+ found_gradient = True
+ squared_norm = accelerator.reduce(squared_norm, reduction="sum")
+ else:
+ # DDP gradients have already been averaged and are replicated, so do not
+ # sum the same full gradient once more across ranks.
+ for parameter in parameters:
+ if parameter.grad is not None:
+ squared_norm += parameter.grad.detach().double().square().sum()
+ found_gradient = True
+ if not found_gradient:
+ raise RuntimeError("No trainable gradients were available for grad-norm logging")
+ return squared_norm.sqrt().float()
+
+
+def _forward_cached_training_loss(model, data, load_from_cache):
+ model_output = model({}, inputs=data) if load_from_cache else model(data)
+ if (
+ isinstance(model_output, tuple)
+ and len(model_output) == 2
+ and isinstance(model_output[1], dict)
+ ):
+ return model_output
+ return model_output, {}
+
+
+def _diagnostic_component_grad_norms(accelerator, model, optimizer, data, load_from_cache):
+ """Measure component parameter gradients without applying an optimizer update.
+
+ Two fresh forward/backward passes avoid retaining the 33B-model graph. RNG
+ states are restored before each pass and again on exit, so these diagnostics
+ do not perturb subsequent training noise/timestep sampling.
+ """
+ component_keys = (
+ ("train/video_loss", "train/video_grad_norm"),
+ ("train/audio_loss", "train/audio_grad_norm"),
+ )
+ cpu_rng_state = torch.random.get_rng_state()
+ cuda_rng_state = None
+ if accelerator.device.type == "cuda":
+ cuda_rng_state = torch.cuda.get_rng_state(accelerator.device)
+ deepspeed_boundary = None
+ if hasattr(model, "set_gradient_accumulation_boundary"):
+ deepspeed_boundary = getattr(model, "_is_gradient_accumulation_boundary", None)
+ # Diagnostic passes must reduce their standalone gradients immediately,
+ # independent of the real training accumulation counter.
+ model.set_gradient_accumulation_boundary(True)
+ results = {}
+ try:
+ for component_key, metric_key in component_keys:
+ torch.random.set_rng_state(cpu_rng_state)
+ if cuda_rng_state is not None:
+ torch.cuda.set_rng_state(cuda_rng_state, accelerator.device)
+ optimizer.zero_grad()
+ _, components = _forward_cached_training_loss(model, data, load_from_cache)
+ component_loss = components.get(component_key)
+ if component_loss is None:
+ raise RuntimeError(
+ f"Model did not return {component_key}; cannot measure component grad norm"
+ )
+ backward_kwargs = {"scale_wrt_gas": False} if hasattr(
+ model, "set_gradient_accumulation_boundary"
+ ) else {}
+ accelerator.backward(component_loss, **backward_kwargs)
+ results[metric_key] = _current_global_grad_norm(accelerator, model).detach()
+ finally:
+ optimizer.zero_grad()
+ torch.random.set_rng_state(cpu_rng_state)
+ if cuda_rng_state is not None:
+ torch.cuda.set_rng_state(cuda_rng_state, accelerator.device)
+ if hasattr(model, "set_gradient_accumulation_boundary"):
+ model.set_gradient_accumulation_boundary(deepspeed_boundary)
+ return results
+
+
+def save_training_args(args):
+ output_path = getattr(args, "output_path", None) if args is not None else None
+ if output_path is None:
+ return
+ try:
+ os.makedirs(args.output_path, exist_ok=True)
+ save_path = os.path.join(args.output_path, "training_args.json")
+ with open(save_path, "w", encoding="utf-8") as f:
+ json.dump(vars(args), f, indent=4, ensure_ascii=False, default=str)
+ print(f"Training arguments saved to `{save_path}`.")
+ except Exception as e:
+ print(f"Warning: failed to save training arguments: {e}")
+
+
+def launch_training_task(
+ accelerator: Accelerator,
+ dataset: torch.utils.data.Dataset,
+ model: DiffusionTrainingModule,
+ model_logger: ModelLogger,
+ learning_rate: float = 1e-5,
+ weight_decay: float = 1e-2,
+ num_workers: int = 1,
+ save_steps: int = None,
+ num_epochs: int = 1,
+ enable_model_cpu_offload: bool = False,
+ enable_optimizer_cpu_offload: bool = False,
+ cpu_offload_split_threshold: int = None,
+ customized_optimizer: str = None,
+ max_train_steps: int = None,
+ skip_model_saving: bool = False,
+ validation_dataset=None,
+ args = None,
+ **kwargs,
+):
+ if args is not None:
+ learning_rate = args.learning_rate
+ weight_decay = args.weight_decay
+ num_workers = args.dataset_num_workers
+ save_steps = args.save_steps
+ num_epochs = args.num_epochs
+ enable_model_cpu_offload = args.enable_model_cpu_offload
+ enable_optimizer_cpu_offload = args.enable_optimizer_cpu_offload
+ cpu_offload_split_threshold = args.cpu_offload_split_threshold
+ customized_optimizer = args.customized_optimizer
+ max_train_steps = args.max_train_steps
+ skip_model_saving = args.skip_model_saving
+
+ if accelerator.is_main_process:
+ save_training_args(args)
+
+ optimizer_class = get_optimizer_class(customized_optimizer)
+ optimizer = optimizer_class(model.trainable_modules(), lr=learning_rate, weight_decay=weight_decay)
+ # DeepSpeed performs the optimizer update inside accelerator.backward().
+ # Keeping a separate PyTorch ConstantLR here would remain attached to the
+ # pre-DeepSpeed optimizer (and emit a misleading step-order warning), while
+ # the actual DeepSpeed optimizer already keeps a constant learning rate.
+ use_deepspeed = getattr(accelerator.state, "deepspeed_plugin", None) is not None
+ scheduler = None if use_deepspeed else torch.optim.lr_scheduler.ConstantLR(optimizer)
+ # Cached video/reference tensors have variable temporal and sequence shapes,
+ # so this training path intentionally uses one sample per rank.
+ dataloader = torch.utils.data.DataLoader(dataset, batch_size=1, shuffle=True, collate_fn=lambda x: x[0], num_workers=num_workers)
+
+ if enable_model_cpu_offload:
+ if scheduler is None:
+ optimizer, dataloader = accelerator.prepare(optimizer, dataloader)
+ else:
+ optimizer, dataloader, scheduler = accelerator.prepare(optimizer, dataloader, scheduler)
+ model.pipe.device = accelerator.device
+ offload_manager = OffloadTrainingManager(model, accelerator.device, enable_optimizer_cpu_offload, cpu_offload_split_threshold)
+ else:
+ model.to(device=accelerator.device)
+ if scheduler is None:
+ model, optimizer, dataloader = accelerator.prepare(model, optimizer, dataloader)
+ else:
+ model, optimizer, dataloader, scheduler = accelerator.prepare(model, optimizer, dataloader, scheduler)
+
+ initialize_deepspeed_gradient_checkpointing(accelerator)
+ validation_steps = getattr(args, "validation_steps", None) if args is not None else None
+ validation_at_start = bool(getattr(args, "validation_at_start", False)) if args is not None else False
+ log_component_grad_norms = bool(
+ getattr(args, "log_component_grad_norms", False)
+ ) if args is not None else False
+ grad_norm_steps = int(getattr(args, "grad_norm_steps", 1)) if args is not None else 1
+ if validation_steps is not None and validation_steps <= 0:
+ raise ValueError("--validation_steps must be positive")
+ if (validation_at_start or validation_steps is not None) and validation_dataset is None:
+ raise ValueError("Validation was requested but --validation_cache_file_list was not provided")
+ if grad_norm_steps <= 0:
+ raise ValueError("--grad_norm_steps must be positive")
+ if validation_at_start:
+ run_full_denoise_validation(accelerator, model, validation_dataset, model_logger, step=0, args=args)
+
+ reached_max_steps = False
+ accumulated_loss = None
+ accumulated_loss_metrics = {}
+ accumulated_micro_steps = 0
+ for epoch_id in range(num_epochs):
+ progress_bar = tqdm(
+ dataloader,
+ desc=f"epoch {epoch_id + 1} micro-batches",
+ disable=not accelerator.is_local_main_process,
+ )
+ for data in progress_bar:
+ with accelerator.accumulate(model):
+ loss, loss_metrics = _forward_cached_training_loss(
+ model, data, dataset.load_from_cache,
+ )
+ accelerator.backward(loss)
+ if enable_model_cpu_offload:
+ offload_manager.after_backward()
+ measure_grad_norms = (
+ accelerator.sync_gradients
+ and log_component_grad_norms
+ and (model_logger.num_steps + 1) % grad_norm_steps == 0
+ )
+ grad_norm_metrics = {}
+ if measure_grad_norms:
+ # This is the exact accumulated gradient used by the real
+ # optimizer update below.
+ grad_norm_metrics["train/grad_norm"] = _current_global_grad_norm(
+ accelerator, model,
+ ).detach()
+ optimizer.step()
+ if scheduler is not None:
+ scheduler.step()
+ optimizer.zero_grad()
+ detached_loss = loss.detach()
+ accumulated_loss = detached_loss if accumulated_loss is None else accumulated_loss + detached_loss
+ for key, value in loss_metrics.items():
+ detached_value = value.detach()
+ accumulated_loss_metrics[key] = accumulated_loss_metrics.get(key, 0) + detached_value
+ accumulated_micro_steps += 1
+ # Only count/log/save/validate after a real optimizer update.
+ # With gradient accumulation, intermediate optimizer.step() calls
+ # are no-ops managed by Accelerate/DeepSpeed.
+ if accelerator.sync_gradients:
+ if measure_grad_norms:
+ grad_norm_metrics.update(_diagnostic_component_grad_norms(
+ accelerator, model, optimizer, data, dataset.load_from_cache,
+ ))
+ mean_loss = accumulated_loss / accumulated_micro_steps
+ mean_loss = accelerator.reduce(mean_loss, reduction="mean")
+ mean_loss_metrics = {
+ key: accelerator.reduce(value / accumulated_micro_steps, reduction="mean")
+ for key, value in accumulated_loss_metrics.items()
+ }
+ accumulated_loss = None
+ accumulated_loss_metrics = {}
+ accumulated_micro_steps = 0
+ mean_loss_metrics.update(grad_norm_metrics)
+ model_logger.on_step_end(
+ accelerator, model, save_steps,
+ loss=mean_loss, metrics=mean_loss_metrics,
+ )
+ progress_bar.set_postfix(optimizer_step=model_logger.num_steps)
+ if validation_steps is not None and model_logger.num_steps % validation_steps == 0:
+ run_full_denoise_validation(
+ accelerator, model, validation_dataset, model_logger,
+ step=model_logger.num_steps, args=args,
+ )
+ if max_train_steps is not None and model_logger.num_steps >= max_train_steps:
+ reached_max_steps = True
+ break
+ if not skip_model_saving and save_steps is None:
+ model_logger.on_epoch_end(accelerator, model, epoch_id)
+ if reached_max_steps:
+ break
+
+ report_peak_cuda_memory(accelerator, model_logger.output_path)
+ model_logger.on_training_end(accelerator, model, save_steps, skip_model_saving=skip_model_saving)
+ accelerator.end_training()
+
+
+def _write_json_atomic(path, payload):
+ temp_path = f"{path}.tmp"
+ with open(temp_path, "w", encoding="utf-8") as handle:
+ json.dump(payload, handle, ensure_ascii=False, indent=2, default=str)
+ handle.write("\n")
+ os.replace(temp_path, path)
+
+
+class _DataProcessShardDataset(torch.utils.data.Dataset):
+ """Shard data by rank and check cache files before decoding source media."""
+
+ def __init__(
+ self,
+ dataset,
+ output_path,
+ process_index,
+ num_processes,
+ skip_existing_cache=False,
+ retry_data_process_oom=False,
+ skip_data_process_errors=False,
+ retry_data_process_errors=False,
+ ):
+ self.dataset = dataset
+ self.folder = os.path.join(output_path, str(process_index))
+ self.process_index = process_index
+ self.num_processes = num_processes
+ self.skip_existing_cache = skip_existing_cache
+ self.retry_data_process_oom = retry_data_process_oom
+ self.skip_data_process_errors = skip_data_process_errors
+ self.retry_data_process_errors = retry_data_process_errors
+ os.makedirs(self.folder, exist_ok=True)
+
+ def __len__(self):
+ remaining = len(self.dataset) - self.process_index
+ return max(0, (remaining + self.num_processes - 1) // self.num_processes)
+
+ def cache_paths(self, local_data_id):
+ return (
+ os.path.join(self.folder, f"{local_data_id}.pth"),
+ os.path.join(self.folder, f"{local_data_id}.oom.json"),
+ os.path.join(self.folder, f"{local_data_id}.data_error.json"),
+ )
+
+ def __getitem__(self, local_data_id):
+ save_path, oom_marker_path, data_error_path = self.cache_paths(local_data_id)
+ if self.skip_existing_cache:
+ if os.path.isfile(save_path):
+ return local_data_id, "cache", None
+ if not self.retry_data_process_oom and os.path.isfile(oom_marker_path):
+ return local_data_id, "oom_marker", None
+ if not self.retry_data_process_errors and os.path.isfile(data_error_path):
+ return local_data_id, "data_error_marker", None
+
+ global_data_id = self.process_index + local_data_id * self.num_processes
+ try:
+ return local_data_id, None, self.dataset[global_data_id]
+ except Exception as error:
+ if not self.skip_data_process_errors:
+ raise
+ return local_data_id, "load_error", {
+ "rank": self.process_index,
+ "local_data_id": local_data_id,
+ "global_data_id": global_data_id,
+ "error_type": type(error).__name__,
+ "error": str(error),
+ "traceback": traceback.format_exc(),
+ "created_at_utc": time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime()),
+ }
+
+
+def launch_data_process_task(
+ accelerator: Accelerator,
+ dataset: torch.utils.data.Dataset,
+ model: DiffusionTrainingModule,
+ model_logger: ModelLogger,
+ num_workers: int = 8,
+ args = None,
+ **kwargs,
+):
+ if args is not None:
+ num_workers = args.dataset_num_workers
+ enable_model_cpu_offload = args.enable_model_cpu_offload
+ enable_optimizer_cpu_offload = args.enable_optimizer_cpu_offload
+ cpu_offload_split_threshold = args.cpu_offload_split_threshold
+ skip_existing_cache = args.skip_existing_cache
+ skip_data_process_oom = getattr(args, "skip_data_process_oom", False)
+ retry_data_process_oom = getattr(args, "retry_data_process_oom", False)
+ skip_data_process_errors = getattr(args, "skip_data_process_errors", False)
+ retry_data_process_errors = getattr(args, "retry_data_process_errors", False)
+ else:
+ enable_model_cpu_offload = False
+ enable_optimizer_cpu_offload = False
+ cpu_offload_split_threshold = None
+ skip_existing_cache = False
+ skip_data_process_oom = False
+ retry_data_process_oom = False
+ skip_data_process_errors = False
+ retry_data_process_errors = False
+
+ # Data processing has no cross-rank computation. Shard explicitly so cache
+ # existence can be checked before UnifiedDataset decodes video/audio data.
+ # This matches Accelerate's batch_size=1, split_batches=False mapping.
+ rank_dataset = _DataProcessShardDataset(
+ dataset,
+ model_logger.output_path,
+ accelerator.process_index,
+ accelerator.num_processes,
+ skip_existing_cache=skip_existing_cache,
+ retry_data_process_oom=retry_data_process_oom,
+ skip_data_process_errors=skip_data_process_errors,
+ retry_data_process_errors=retry_data_process_errors,
+ )
+ dataloader = torch.utils.data.DataLoader(
+ rank_dataset,
+ shuffle=False,
+ collate_fn=lambda items: items[0],
+ num_workers=num_workers,
+ )
+ if enable_model_cpu_offload:
+ offload_manager = OffloadTrainingManager(
+ model,
+ accelerator.device,
+ enable_optimizer_cpu_offload,
+ cpu_offload_split_threshold,
+ )
+ model.pipe.device = accelerator.device
+ else:
+ # Do not wrap the preprocessing model in DDP: ranks are independent and
+ # never synchronize gradients or buffers in this task.
+ model.to(device=accelerator.device)
+
+ processed_count = 0
+ skipped_cache_count = 0
+ skipped_oom_marker_count = 0
+ skipped_data_error_marker_count = 0
+ new_oom_count = 0
+ new_data_error_count = 0
+ progress_bar = tqdm(
+ dataloader,
+ desc=f"rank {accelerator.process_index}",
+ position=accelerator.local_process_index,
+ )
+ for data_id, status, data in progress_bar:
+ if status == "cache":
+ skipped_cache_count += 1
+ continue
+ if status == "oom_marker":
+ skipped_oom_marker_count += 1
+ continue
+ if status == "data_error_marker":
+ skipped_data_error_marker_count += 1
+ continue
+
+ save_path, oom_marker_path, data_error_path = rank_dataset.cache_paths(data_id)
+ if status == "load_error":
+ _write_json_atomic(data_error_path, data)
+ new_data_error_count += 1
+ tqdm.write(
+ f"Rank {accelerator.process_index}: data loading failed at local sample {data_id} "
+ f"({data['error_type']}: {data['error']}); recorded {data_error_path} and continuing."
+ )
+ continue
+
+ with accelerator.accumulate(model):
+ with torch.no_grad():
+ source = data.get("source") if isinstance(data, dict) else None
+ temp_save_path = f"{save_path}.tmp"
+ try:
+ processed_data = model(data)
+ if enable_model_cpu_offload:
+ offload_manager.after_backward()
+ torch.save(processed_data, temp_save_path)
+ os.replace(temp_save_path, save_path)
+ for stale_marker in (oom_marker_path, data_error_path):
+ if os.path.isfile(stale_marker):
+ os.unlink(stale_marker)
+ processed_count += 1
+ del processed_data
+ except RuntimeError as error:
+ is_cuda_oom = isinstance(error, torch.OutOfMemoryError) or "CUDA out of memory" in str(error)
+ if not skip_data_process_oom or not is_cuda_oom:
+ raise
+ error_message = str(error)
+ error.__traceback__ = None
+ if os.path.isfile(temp_save_path):
+ os.unlink(temp_save_path)
+ del data
+ gc.collect()
+ if torch.cuda.is_available():
+ torch.cuda.empty_cache()
+ marker = {
+ "rank": accelerator.process_index,
+ "local_data_id": data_id,
+ "global_data_id": accelerator.process_index + data_id * accelerator.num_processes,
+ "source": source,
+ "error": error_message,
+ "created_at_utc": time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime()),
+ }
+ _write_json_atomic(oom_marker_path, marker)
+ new_oom_count += 1
+ tqdm.write(
+ f"Rank {accelerator.process_index}: CUDA OOM at local sample {data_id}; "
+ f"recorded {oom_marker_path} and continuing."
+ )
+ continue
+ progress_bar.set_postfix(
+ processed=processed_count,
+ cached=skipped_cache_count,
+ oom=new_oom_count,
+ data_error=new_data_error_count,
+ )
+ print(
+ f"Rank {accelerator.process_index} data processing complete: processed={processed_count}, "
+ f"cached={skipped_cache_count}, new_oom={new_oom_count}, "
+ f"cached_oom={skipped_oom_marker_count}, new_data_error={new_data_error_count}, "
+ f"cached_data_error={skipped_data_error_marker_count}."
+ )
+ accelerator.end_training()
+
+
+def initialize_deepspeed_gradient_checkpointing(accelerator: Accelerator):
+ if getattr(accelerator.state, "deepspeed_plugin", None) is not None:
+ ds_config = accelerator.state.deepspeed_plugin.deepspeed_config
+ if "activation_checkpointing" in ds_config:
+ import deepspeed
+ act_config = ds_config["activation_checkpointing"]
+ deepspeed.checkpointing.configure(
+ mpu_=None,
+ partition_activations=act_config.get("partition_activations", False),
+ checkpoint_in_cpu=act_config.get("cpu_checkpointing", False),
+ contiguous_checkpointing=act_config.get("contiguous_memory_optimization", False)
+ )
+ else:
+ print("Do not find activation_checkpointing config in deepspeed config, skip initializing deepspeed gradient checkpointing.")
+
+
+def report_peak_cuda_memory(accelerator: Accelerator, output_path: str):
+ if accelerator.device.type != "cuda":
+ return
+ torch.cuda.synchronize(accelerator.device)
+ local = torch.tensor(
+ [torch.cuda.max_memory_allocated(accelerator.device), torch.cuda.max_memory_reserved(accelerator.device)],
+ dtype=torch.int64,
+ device=accelerator.device,
+ )
+ gathered = accelerator.gather(local).reshape(-1, 2).cpu().tolist()
+ if accelerator.is_main_process:
+ mib = 1024 ** 2
+ ranks = [
+ {
+ "rank": rank,
+ "max_allocated_bytes": int(values[0]),
+ "max_reserved_bytes": int(values[1]),
+ "max_allocated_mib": round(values[0] / mib, 2),
+ "max_reserved_mib": round(values[1] / mib, 2),
+ }
+ for rank, values in enumerate(gathered)
+ ]
+ report = {
+ "world_size": len(ranks),
+ "per_rank": ranks,
+ "max_allocated_mib": max(item["max_allocated_mib"] for item in ranks),
+ "max_reserved_mib": max(item["max_reserved_mib"] for item in ranks),
+ "note": "PyTorch CUDA allocator peak; excludes non-PyTorch driver allocations.",
+ }
+ os.makedirs(output_path, exist_ok=True)
+ report_path = os.path.join(output_path, "peak_cuda_memory.json")
+ with open(report_path, "w", encoding="utf-8") as handle:
+ json.dump(report, handle, indent=2)
+ handle.write("\n")
+ print(f"Peak CUDA memory report saved to `{report_path}`: {report}")
+
+
+@torch.no_grad()
+def run_full_denoise_validation(accelerator, model, validation_dataset, model_logger, step, args):
+ """Run the same validation samples on every ZeRO rank.
+
+ ZeRO-3 forward passes contain parameter collectives, so every rank must execute
+ the same number and order of denoising calls. Each metric is reduced across
+ ranks to catch any numerical divergence, then averaged across validation items.
+ """
+ if len(validation_dataset) == 0:
+ raise ValueError("Validation cache manifest is empty")
+ accelerator.wait_for_everyone()
+ was_training = model.training
+ model.eval()
+ started_at = time.time()
+ totals = {}
+ sample_records = []
+ iterator = tqdm(
+ range(len(validation_dataset)),
+ desc=f"validation@{step}",
+ disable=not accelerator.is_main_process,
+ )
+ for sample_id in iterator:
+ data = validation_dataset[sample_id]
+ validation_result = model(
+ {},
+ inputs=data,
+ validation_config={
+ "num_inference_steps": args.validation_num_inference_steps,
+ "seed": args.validation_seed + sample_id,
+ "flow_shift": args.validation_flow_shift,
+ "audio_flow_shift": args.validation_audio_flow_shift,
+ "cfg_scale": args.validation_cfg_scale,
+ "save_outputs": args.validation_save_outputs,
+ },
+ )
+ if (
+ isinstance(validation_result, tuple)
+ and len(validation_result) == 2
+ and isinstance(validation_result[1], dict)
+ ):
+ metrics, outputs = validation_result
+ else:
+ metrics, outputs = validation_result, None
+ reduced_metrics = {}
+ for key, value in metrics.items():
+ reduced = accelerator.reduce(value.detach(), reduction="mean")
+ scalar = float(reduced.float().item())
+ totals[key] = totals.get(key, 0.0) + scalar
+ reduced_metrics[key] = scalar
+ sample_record = {"sample_id": sample_id, **reduced_metrics}
+ if accelerator.is_main_process and outputs is not None:
+ from diffsynth.utils.data.audio_video import write_video_audio
+ validation_dir = os.path.join(
+ model_logger.output_path, "validation", f"step-{step:09d}",
+ )
+ os.makedirs(validation_dir, exist_ok=True)
+ output_path = os.path.join(validation_dir, f"sample-{sample_id:02d}.mp4")
+ write_video_audio(
+ video=outputs["video"], audio=outputs.get("audio"), output_path=output_path,
+ fps=24, audio_sample_rate=outputs.get("audio_sample_rate"),
+ )
+ sample_record["output_path"] = output_path
+ print(f"Saved validation sample {sample_id} to `{output_path}`.")
+ if accelerator.is_main_process:
+ sample_records.append(sample_record)
+ del data, metrics, outputs, validation_result
+
+ averaged = {f"validation/{key}": value / len(validation_dataset) for key, value in totals.items()}
+ if "validation/latent_mse" in averaged:
+ # A canonical validation loss key for W&B dashboards and sweep tooling.
+ # The modality-specific MSE values remain available alongside it.
+ averaged["validation/loss"] = averaged["validation/latent_mse"]
+ elapsed = time.time() - started_at
+ averaged["validation/seconds"] = elapsed
+ model_logger.log_metrics(accelerator, averaged, step)
+ if accelerator.is_main_process:
+ os.makedirs(model_logger.output_path, exist_ok=True)
+ history_path = os.path.join(model_logger.output_path, "validation_history.jsonl")
+ record = {
+ "step": step,
+ "num_samples": len(validation_dataset),
+ "num_inference_steps": args.validation_num_inference_steps,
+ "cfg_scale": args.validation_cfg_scale,
+ "elapsed_seconds": elapsed,
+ "metrics": averaged,
+ "samples": sample_records,
+ }
+ with open(history_path, "a", encoding="utf-8") as handle:
+ handle.write(json.dumps(record, ensure_ascii=False) + "\n")
+ print(f"Validation step {step}: {averaged}. Details appended to `{history_path}`.")
+ if was_training:
+ model.train()
+ accelerator.wait_for_everyone()
diff --git a/diffsynth/diffusion/template.py b/diffsynth/diffusion/template.py
new file mode 100644
index 0000000000000000000000000000000000000000..3c03b1b17f76241178eeb20bf7978f464582778a
--- /dev/null
+++ b/diffsynth/diffusion/template.py
@@ -0,0 +1,208 @@
+import torch, os, importlib, warnings, json, inspect
+from typing import Dict, List, Tuple, Union
+from ..core import ModelConfig, load_model
+from ..core.device.npu_compatible_device import get_device_type
+from ..utils.lora.merge import merge_lora
+
+
+KVCache = Dict[str, Tuple[torch.Tensor, torch.Tensor]]
+
+
+class TemplateModel(torch.nn.Module):
+ def __init__(self):
+ super().__init__()
+
+ @torch.no_grad()
+ def process_inputs(self, **kwargs):
+ return {}
+
+ def forward(self, **kwargs):
+ raise NotImplementedError()
+
+
+def check_template_model_format(model):
+ if not hasattr(model, "process_inputs"):
+ raise NotImplementedError("`process_inputs` is not implemented in the Template model.")
+ if "kwargs" not in inspect.signature(model.process_inputs).parameters:
+ raise NotImplementedError("`**kwargs` is not included in `process_inputs`.")
+ if not hasattr(model, "forward"):
+ raise NotImplementedError("`forward` is not implemented in the Template model.")
+ if "kwargs" not in inspect.signature(model.forward).parameters:
+ raise NotImplementedError("`**kwargs` is not included in `forward`.")
+
+
+def load_template_model(path, torch_dtype=torch.bfloat16, device="cuda", verbose=1):
+ spec = importlib.util.spec_from_file_location("template_model", os.path.join(path, "model.py"))
+ module = importlib.util.module_from_spec(spec)
+ spec.loader.exec_module(module)
+ template_model_path = getattr(module, 'TEMPLATE_MODEL_PATH') if hasattr(module, 'TEMPLATE_MODEL_PATH') else None
+ if template_model_path is not None:
+ # With `TEMPLATE_MODEL_PATH`, a pretrained model will be loaded.
+ model = load_model(
+ model_class=getattr(module, 'TEMPLATE_MODEL'),
+ config=getattr(module, 'TEMPLATE_MODEL_CONFIG') if hasattr(module, 'TEMPLATE_MODEL_CONFIG') else None,
+ path=os.path.join(path, getattr(module, 'TEMPLATE_MODEL_PATH')),
+ torch_dtype=torch_dtype,
+ device=device,
+ )
+ else:
+ # Without `TEMPLATE_MODEL_PATH`, a randomly initialized model or a non-model module will be loaded.
+ model = module.TEMPLATE_MODEL()
+ if hasattr(model, "to"):
+ model = model.to(dtype=torch_dtype, device=device)
+ if hasattr(model, "eval"):
+ model = model.eval()
+ check_template_model_format(model)
+ if verbose > 0:
+ metadata = {
+ "model_architecture": getattr(module, 'TEMPLATE_MODEL').__name__,
+ "code_path": os.path.join(path, "model.py"),
+ "weight_path": template_model_path,
+ }
+ print(f"Template model loaded: {json.dumps(metadata, indent=4)}")
+ return model
+
+
+def load_template_data_processor(path):
+ spec = importlib.util.spec_from_file_location("template_model", os.path.join(path, "model.py"))
+ module = importlib.util.module_from_spec(spec)
+ spec.loader.exec_module(module)
+ if hasattr(module, 'TEMPLATE_DATA_PROCESSOR'):
+ processor = getattr(module, 'TEMPLATE_DATA_PROCESSOR')
+ return processor
+ else:
+ return None
+
+
+class TemplatePipeline(torch.nn.Module):
+ def __init__(
+ self,
+ torch_dtype: torch.dtype = torch.bfloat16,
+ device: Union[str, torch.device] = get_device_type(),
+ model_configs: list[ModelConfig] = [],
+ lazy_loading: bool = False,
+ ):
+ super().__init__()
+ self.torch_dtype = torch_dtype
+ self.device = device
+ self.model_configs = model_configs
+ self.lazy_loading = lazy_loading
+ if lazy_loading:
+ for model_config in model_configs:
+ TemplatePipeline.check_vram_config(model_config)
+ model_config.download_if_necessary()
+ self.models = None
+ else:
+ models = []
+ for model_config in model_configs:
+ TemplatePipeline.check_vram_config(model_config)
+ model_config.download_if_necessary()
+ model = load_template_model(model_config.path, torch_dtype=torch_dtype, device=device)
+ models.append(model)
+ self.models = torch.nn.ModuleList(models)
+
+ def merge_kv_cache(self, kv_cache_list: List[KVCache]) -> KVCache:
+ names = {}
+ for kv_cache in kv_cache_list:
+ for name in kv_cache:
+ names[name] = None
+ kv_cache_merged = {}
+ for name in names:
+ kv_list = [kv_cache.get(name) for kv_cache in kv_cache_list]
+ kv_list = [kv for kv in kv_list if kv is not None]
+ if len(kv_list) > 0:
+ k = torch.concat([kv[0] for kv in kv_list], dim=1)
+ v = torch.concat([kv[1] for kv in kv_list], dim=1)
+ kv_cache_merged[name] = (k, v)
+ return kv_cache_merged
+
+ def merge_sequential_embeddings(self, embeddings: List[torch.Tensor]):
+ return torch.concat(embeddings, dim=-2)
+
+ def merge_template_cache(self, template_cache_list):
+ params = sorted(list(set(sum([list(template_cache.keys()) for template_cache in template_cache_list], []))))
+ template_cache_merged = {}
+ for param in params:
+ data = [template_cache[param] for template_cache in template_cache_list if param in template_cache]
+ if param == "kv_cache":
+ data = self.merge_kv_cache(data)
+ elif param == "lora":
+ data = merge_lora(data)
+ elif param == "text_embedding":
+ data = self.merge_sequential_embeddings(data)
+ elif len(data) == 1:
+ data = data[0]
+ else:
+ print(f"Conflict detected: `{param}` appears in the outputs of multiple Template models. Only the first one will be retained.")
+ data = data[0]
+ template_cache_merged[param] = data
+ return template_cache_merged
+
+ @staticmethod
+ def check_vram_config(model_config: ModelConfig):
+ params = [
+ model_config.offload_device, model_config.offload_dtype,
+ model_config.onload_device, model_config.onload_dtype,
+ model_config.preparing_device, model_config.preparing_dtype,
+ model_config.computation_device, model_config.computation_dtype,
+ ]
+ for param in params:
+ if param is not None:
+ warnings.warn("TemplatePipeline doesn't support VRAM management. VRAM config will be ignored.")
+
+ @staticmethod
+ def from_pretrained(
+ torch_dtype: torch.dtype = torch.bfloat16,
+ device: Union[str, torch.device] = get_device_type(),
+ model_configs: list[ModelConfig] = [],
+ lazy_loading: bool = False,
+ ):
+ pipe = TemplatePipeline(torch_dtype, device, model_configs, lazy_loading)
+ return pipe
+
+ def fetch_model(self, model_id):
+ if self.lazy_loading:
+ model_config = self.model_configs[model_id]
+ model_config.download_if_necessary()
+ model = load_template_model(model_config.path, torch_dtype=self.torch_dtype, device=self.device)
+ else:
+ model = self.models[model_id]
+ return model
+
+ def call_single_side(self, pipe=None, inputs: List[Dict] = None):
+ model = None
+ onload_model_id = -1
+ template_cache = []
+ for i in inputs:
+ model_id = i.get("model_id", 0)
+ if model_id != onload_model_id:
+ model = self.fetch_model(model_id)
+ onload_model_id = model_id
+ cache = model.process_inputs(pipe=pipe, **i)
+ cache = model.forward(pipe=pipe, **cache)
+ template_cache.append(cache)
+ template_cache = self.merge_template_cache(template_cache)
+ return template_cache
+
+ @torch.no_grad()
+ def __call__(
+ self,
+ pipe=None,
+ template_inputs: List[Dict] = None,
+ negative_template_inputs: List[Dict] = None,
+ **kwargs,
+ ):
+ template_cache = self.call_single_side(pipe=pipe, inputs=template_inputs or [])
+ negative_template_cache = self.call_single_side(pipe=pipe, inputs=negative_template_inputs or [])
+ required_params = list(inspect.signature(pipe.__call__).parameters.keys())
+ for param in template_cache:
+ if param in required_params:
+ kwargs[param] = template_cache[param]
+ else:
+ print(f"`{param}` is not included in the inputs of `{pipe.__class__.__name__}`. This parameter will be ignored.")
+ for param in negative_template_cache:
+ if "negative_" + param in required_params:
+ kwargs["negative_" + param] = negative_template_cache[param]
+ else:
+ print(f"`{'negative_' + param}` is not included in the inputs of `{pipe.__class__.__name__}`. This parameter will be ignored.")
+ return pipe(**kwargs)
diff --git a/diffsynth/diffusion/training_module.py b/diffsynth/diffusion/training_module.py
new file mode 100644
index 0000000000000000000000000000000000000000..7fe8df88e07751735ed364f3f3d1cc839c8a7a50
--- /dev/null
+++ b/diffsynth/diffusion/training_module.py
@@ -0,0 +1,412 @@
+import torch, json, os, inspect
+from ..core import ModelConfig, load_state_dict, QuantizeConfig
+from ..utils.controlnet import ControlNetInput
+from .base_pipeline import PipelineUnit
+
+
+class GeneralUnit_RemoveCache(PipelineUnit):
+ # Only used for training
+ def __init__(self, required_params=tuple(), force_remove_params_shared=tuple(), force_remove_params_posi=tuple(), force_remove_params_nega=tuple()):
+ super().__init__(take_over=True)
+ self.required_params = required_params
+ self.force_remove_params_shared = force_remove_params_shared
+ self.force_remove_params_posi = force_remove_params_posi
+ self.force_remove_params_nega = force_remove_params_nega
+
+ def process_params(self, inputs, required_params, force_remove_params):
+ inputs_ = {}
+ for name, param in inputs.items():
+ if name in required_params and name not in force_remove_params:
+ inputs_[name] = param
+ return inputs_
+
+ def process(self, pipe, inputs_shared, inputs_posi, inputs_nega):
+ inputs_shared = self.process_params(inputs_shared, self.required_params, self.force_remove_params_shared)
+ inputs_posi = self.process_params(inputs_posi, self.required_params, self.force_remove_params_posi)
+ inputs_nega = self.process_params(inputs_nega, self.required_params, self.force_remove_params_nega)
+ return inputs_shared, inputs_posi, inputs_nega
+
+
+class GeneralUnit_TemplateProcessInputs(PipelineUnit):
+ # Only used for training
+ def __init__(self, data_processor):
+ super().__init__(
+ input_params=("template_inputs",),
+ output_params=("template_inputs",),
+ )
+ self.data_processor = data_processor
+
+ def process(self, pipe, template_inputs):
+ if not hasattr(pipe, "template_model") or template_inputs is None:
+ return {}
+ if self.data_processor is not None:
+ template_inputs = self.data_processor(**template_inputs)
+ template_inputs = pipe.template_model.process_inputs(pipe=pipe, **template_inputs)
+ return {"template_inputs": template_inputs}
+
+
+class GeneralUnit_TemplateForward(PipelineUnit):
+ # Only used for training
+ def __init__(self, use_gradient_checkpointing=False, use_gradient_checkpointing_offload=False):
+ super().__init__(
+ input_params=("template_inputs",),
+ output_params=("kv_cache",),
+ onload_model_names=("template_model",)
+ )
+ self.use_gradient_checkpointing = use_gradient_checkpointing
+ self.use_gradient_checkpointing_offload = use_gradient_checkpointing_offload
+
+ def process(self, pipe, template_inputs):
+ if not hasattr(pipe, "template_model") or template_inputs is None:
+ return {}
+ template_cache = pipe.template_model.forward(
+ **template_inputs,
+ pipe=pipe,
+ use_gradient_checkpointing=self.use_gradient_checkpointing,
+ use_gradient_checkpointing_offload=self.use_gradient_checkpointing_offload,
+ )
+ return template_cache
+
+
+class DiffusionTrainingModule(torch.nn.Module):
+ def __init__(self):
+ super().__init__()
+
+
+ def to(self, *args, **kwargs):
+ for name, model in self.named_children():
+ model.to(*args, **kwargs)
+ return self
+
+
+ def trainable_modules(self):
+ trainable_modules = filter(lambda p: p.requires_grad, self.parameters())
+ return trainable_modules
+
+
+ def trainable_param_names(self):
+ trainable_param_names = list(filter(lambda named_param: named_param[1].requires_grad, self.named_parameters()))
+ trainable_param_names = set([named_param[0] for named_param in trainable_param_names])
+ return trainable_param_names
+
+
+ def add_lora_to_model(self, model, target_modules, lora_rank, lora_alpha=None, upcast_dtype=None):
+ from peft import LoraConfig, inject_adapter_in_model
+ if lora_alpha is None:
+ lora_alpha = lora_rank
+ if isinstance(target_modules, list) and len(target_modules) == 1:
+ target_modules = target_modules[0]
+ lora_config = LoraConfig(r=lora_rank, lora_alpha=lora_alpha, target_modules=target_modules)
+ model = inject_adapter_in_model(lora_config, model)
+ if upcast_dtype is not None:
+ for param in model.parameters():
+ if param.requires_grad:
+ param.data = param.to(upcast_dtype)
+ return model
+
+
+ def mapping_lora_state_dict(self, state_dict):
+ new_state_dict = {}
+ for key, value in state_dict.items():
+ if "lora_A.weight" in key or "lora_B.weight" in key:
+ new_key = key.replace("lora_A.weight", "lora_A.default.weight").replace("lora_B.weight", "lora_B.default.weight")
+ new_state_dict[new_key] = value
+ elif "lora_A.default.weight" in key or "lora_B.default.weight" in key:
+ new_state_dict[key] = value
+ return new_state_dict
+
+
+ def export_trainable_state_dict(self, state_dict, remove_prefix=None):
+ trainable_param_names = self.trainable_param_names()
+ state_dict = {name: param for name, param in state_dict.items() if name in trainable_param_names}
+ if remove_prefix is not None:
+ state_dict_ = {}
+ for name, param in state_dict.items():
+ if name.startswith(remove_prefix):
+ name = name[len(remove_prefix):]
+ state_dict_[name] = param
+ state_dict = state_dict_
+ return state_dict
+
+
+ def transfer_data_to_device(self, data, device, torch_float_dtype=None):
+ if data is None:
+ return data
+ elif isinstance(data, torch.Tensor):
+ data = data.to(device)
+ if torch_float_dtype is not None and data.dtype in [torch.float, torch.float16, torch.bfloat16]:
+ data = data.to(torch_float_dtype)
+ return data
+ elif isinstance(data, tuple):
+ data = tuple(self.transfer_data_to_device(x, device, torch_float_dtype) for x in data)
+ return data
+ elif isinstance(data, list):
+ data = list(self.transfer_data_to_device(x, device, torch_float_dtype) for x in data)
+ return data
+ elif isinstance(data, dict):
+ data = {i: self.transfer_data_to_device(data[i], device, torch_float_dtype) for i in data}
+ return data
+ else:
+ return data
+
+ def parse_vram_config(self, fp8=False, offload=False, device="cpu"):
+ if fp8:
+ return {
+ "offload_dtype": torch.float8_e4m3fn,
+ "offload_device": device,
+ "onload_dtype": torch.float8_e4m3fn,
+ "onload_device": device,
+ "preparing_dtype": torch.float8_e4m3fn,
+ "preparing_device": device,
+ "computation_dtype": torch.bfloat16,
+ "computation_device": device,
+ }
+ elif offload:
+ return {
+ "offload_dtype": "disk",
+ "offload_device": "disk",
+ "onload_dtype": "disk",
+ "onload_device": "disk",
+ "preparing_dtype": torch.bfloat16,
+ "preparing_device": device,
+ "computation_dtype": torch.bfloat16,
+ "computation_device": device,
+ "clear_parameters": True,
+ }
+ else:
+ return {}
+
+ def parse_quant_options(self, quant_options):
+ quant_map = {}
+ if quant_options is None:
+ return quant_map
+ for entry in quant_options.split(";"):
+ if entry == "":
+ continue
+ model_string, sep, spec = entry.rpartition(":")
+ if sep == "":
+ raise ValueError(f"Failed to parse quant option: `{entry}`. Expected `:[/]`.")
+ if spec == "":
+ continue
+ method, _, excludes = spec.partition("/")
+ exclude_modules = excludes.split(",") if excludes != "" else None
+ quant_config = QuantizeConfig(method=method, exclude_modules=exclude_modules)
+ if not quant_config.backend.capabilities().get("is_differentiable", True):
+ raise ValueError(f"Quantization method `{method}` is not differentiable, so it cannot be used for training (frozen quantized layers must pass gradients through to LoRA branches). Choose a method whose backend declares `is_differentiable=True`.")
+ quant_map[self.normalize_quant_key(model_string)] = quant_config
+ return quant_map
+
+ def normalize_quant_key(self, model_string):
+ if isinstance(model_string, str) and model_string.startswith("[") and model_string.endswith("]"):
+ try:
+ model_string = json.loads(model_string)
+ except json.JSONDecodeError:
+ raise ValueError(f"Failed to parse quant option model string: `{model_string}`. A multi-file model must be written as a JSON list of files, matching its `--model_paths` entry.")
+ if isinstance(model_string, list):
+ return tuple(model_string)
+ return model_string
+
+ def get_quant_config(self, quant_map, model_string):
+ quant_config = quant_map.get(self.normalize_quant_key(model_string))
+ if quant_config is None and len(quant_map) > 0:
+ print(f"No quant option matches `{model_string}`. This model is loaded without quantization. Parsed quant options: {({key: config.method for key, config in quant_map.items()})}.")
+ return quant_config
+
+ def parse_model_configs(self, model_paths, model_id_with_origin_paths, fp8_models=None, offload_models=None, quant_options=None, device="cpu"):
+ fp8_models = [] if fp8_models is None else fp8_models.split(",")
+ offload_models = [] if offload_models is None else offload_models.split(",")
+ quant_map = self.parse_quant_options(quant_options)
+ model_configs = []
+ if model_paths is not None:
+ model_paths = json.loads(model_paths)
+ for path in model_paths:
+ vram_config = self.parse_vram_config(
+ fp8=path in fp8_models,
+ offload=path in offload_models,
+ device=device
+ )
+ model_configs.append(ModelConfig(path=path, quantize=self.get_quant_config(quant_map, path), **vram_config))
+ if model_id_with_origin_paths is not None:
+ model_id_with_origin_paths = model_id_with_origin_paths.split(",")
+ for model_id_with_origin_path in model_id_with_origin_paths:
+ vram_config = self.parse_vram_config(
+ fp8=model_id_with_origin_path in fp8_models,
+ offload=model_id_with_origin_path in offload_models,
+ device=device
+ )
+ config = self.parse_path_or_model_id(model_id_with_origin_path)
+ model_configs.append(ModelConfig(model_id=config.model_id, origin_file_pattern=config.origin_file_pattern, quantize=self.get_quant_config(quant_map, model_id_with_origin_path), **vram_config))
+ return model_configs
+
+
+ def parse_path_or_model_id(self, model_id_with_origin_path, default_value=None):
+ if model_id_with_origin_path is None:
+ return default_value
+ elif os.path.exists(model_id_with_origin_path):
+ return ModelConfig(path=model_id_with_origin_path)
+ else:
+ if ":" not in model_id_with_origin_path:
+ raise ValueError(f"Failed to parse model config: {model_id_with_origin_path}. This is neither a valid path nor in the format of `model_id/origin_file_pattern`.")
+ split_id = model_id_with_origin_path.rfind(":")
+ model_id = model_id_with_origin_path[:split_id]
+ origin_file_pattern = model_id_with_origin_path[split_id + 1:]
+ return ModelConfig(model_id=model_id, origin_file_pattern=origin_file_pattern)
+
+
+ def auto_detect_lora_target_modules(
+ self,
+ model: torch.nn.Module,
+ search_for_linear=False,
+ linear_detector=lambda x: min(x.in_features, x.out_features) >= 512,
+ block_list_detector=lambda x: isinstance(x, torch.nn.ModuleList) and len(x) > 1,
+ name_prefix="",
+ ):
+ lora_target_modules = []
+ if search_for_linear:
+ for name, module in model.named_modules():
+ module_name = name_prefix + ["", "."][name_prefix != ""] + name
+ if isinstance(module, torch.nn.Linear) and linear_detector(module):
+ lora_target_modules.append(module_name)
+ else:
+ for name, module in model.named_children():
+ module_name = name_prefix + ["", "."][name_prefix != ""] + name
+ lora_target_modules += self.auto_detect_lora_target_modules(
+ module,
+ search_for_linear=block_list_detector(module),
+ linear_detector=linear_detector,
+ block_list_detector=block_list_detector,
+ name_prefix=module_name,
+ )
+ return lora_target_modules
+
+
+ def parse_lora_target_modules(self, model, lora_target_modules):
+ if lora_target_modules == "":
+ print("No LoRA target modules specified. The framework will automatically search for them.")
+ lora_target_modules = self.auto_detect_lora_target_modules(model)
+ print(f"LoRA will be patched at {lora_target_modules}.")
+ else:
+ lora_target_modules = lora_target_modules.split(",")
+ return lora_target_modules
+
+
+ def load_training_template_model(self, pipe, path_or_model_id, use_gradient_checkpointing=False, use_gradient_checkpointing_offload=False):
+ if path_or_model_id is None:
+ return pipe
+ model_config = self.parse_path_or_model_id(path_or_model_id)
+ pipe.load_training_template_model(model_config)
+ template_units = [
+ GeneralUnit_TemplateProcessInputs(pipe.template_data_processor),
+ GeneralUnit_TemplateForward(use_gradient_checkpointing, use_gradient_checkpointing_offload),
+ ]
+ pipe.units = template_units + pipe.units
+ return pipe
+
+
+ def switch_pipe_to_training_mode(
+ self,
+ pipe,
+ trainable_models=None,
+ lora_base_model=None, lora_target_modules="", lora_rank=32, lora_checkpoint=None,
+ preset_lora_path=None, preset_lora_model=None,
+ task="sft",
+ ):
+ # Scheduler
+ pipe.scheduler.set_timesteps(1000, training=True)
+
+ # Freeze untrainable models
+ pipe.freeze_except([] if trainable_models is None else trainable_models.split(","))
+
+ # Preset LoRA
+ if preset_lora_path is not None:
+ pipe.load_lora(getattr(pipe, preset_lora_model), preset_lora_path)
+
+ # FP8
+ # FP8 relies on a model-specific memory management scheme.
+ # It is delegated to the subclass.
+
+ # Add LoRA to the base models
+ if lora_base_model is not None and not task.endswith(":data_process"):
+ if lora_base_model != "" and ((not hasattr(pipe, lora_base_model)) or getattr(pipe, lora_base_model) is None):
+ print(f"No {lora_base_model} models in the pipeline. We cannot patch LoRA on the model. If this occurs during the data processing stage, it is normal.")
+ return
+ model = pipe if lora_base_model == "" else getattr(pipe, lora_base_model)
+ model = self.add_lora_to_model(
+ model,
+ target_modules=self.parse_lora_target_modules(model, lora_target_modules),
+ lora_rank=lora_rank,
+ upcast_dtype=pipe.torch_dtype,
+ )
+ if lora_checkpoint is not None:
+ if isinstance(lora_checkpoint, str):
+ lora = load_state_dict(lora_checkpoint)
+ else:
+ lora = lora_checkpoint
+ lora_loader = pipe.lora_loader(torch_dtype=pipe.torch_dtype, device=pipe.device)
+ lora = lora_loader.convert_state_dict(lora)
+ lora = self.mapping_lora_state_dict(lora)
+ load_result = model.load_state_dict(lora, strict=False)
+ print(f"LoRA checkpoint loaded. Total {len(lora)} keys")
+ if len(load_result[1]) > 0:
+ print(f"Warning, LoRA key mismatch! Unexpected keys in LoRA checkpoint: {load_result[1]}")
+ if lora_base_model != "":
+ setattr(pipe, lora_base_model, model)
+
+
+ def split_pipeline_units(
+ self, task, pipe,
+ trainable_models=None, lora_base_model=None,
+ # TODO: set `remove_unnecessary_params` to `True` by default
+ remove_unnecessary_params=False,
+ # TODO: move `loss_required_params` to `loss.py`
+ loss_required_params=("input_latents", "max_timestep_boundary", "min_timestep_boundary", "first_frame_latents", "video_latents", "audio_input_latents", "num_inference_steps", "cfg_scale"),
+ force_remove_params_shared=tuple(),
+ force_remove_params_posi=tuple(),
+ force_remove_params_nega=tuple(),
+ ):
+ models_require_backward = []
+ if trainable_models is not None:
+ models_require_backward += trainable_models.split(",")
+ if lora_base_model is not None:
+ models_require_backward += [lora_base_model]
+ if task.endswith(":data_process"):
+ other_units, pipe.units = pipe.split_pipeline_units(models_require_backward)
+ if remove_unnecessary_params:
+ required_params = list(loss_required_params) + [i for i in inspect.signature(self.pipe.model_fn).parameters]
+ for unit in other_units:
+ required_params.extend(unit.fetch_input_params())
+ required_params = sorted(list(set(required_params)))
+ pipe.units.append(GeneralUnit_RemoveCache(required_params, force_remove_params_shared, force_remove_params_posi, force_remove_params_nega))
+ elif task.endswith(":train"):
+ pipe.units, _ = pipe.split_pipeline_units(models_require_backward)
+ return pipe
+
+ def parse_extra_inputs(self, data, extra_inputs, inputs_shared):
+ controlnet_keys_map = (
+ ("blockwise_controlnet_", "blockwise_controlnet_inputs",),
+ ("controlnet_", "controlnet_inputs"),
+ )
+ controlnet_inputs = {}
+ for extra_input in extra_inputs:
+ for prefix, name in controlnet_keys_map:
+ if extra_input.startswith(prefix):
+ if name not in controlnet_inputs:
+ controlnet_inputs[name] = {}
+ controlnet_inputs[name][extra_input.replace(prefix, "")] = data[extra_input]
+ break
+ else:
+ inputs_shared[extra_input] = data[extra_input]
+ for name, params in controlnet_inputs.items():
+ inputs_shared[name] = [ControlNetInput(**params)]
+ return inputs_shared
+
+ def resume_from_checkpoint(self, path, remove_prefix_in_ckpt):
+ if path is None:
+ return
+ state_dict = load_state_dict(path)
+ if remove_prefix_in_ckpt is not None:
+ state_dict = {remove_prefix_in_ckpt + i: state_dict[i] for i in state_dict}
+ missing_keys, unexpected_keys = self.load_state_dict(state_dict, strict=False)
+ if len(unexpected_keys) != 0:
+ raise ValueError(f"Cannot load checkpoint: {path}. {len(unexpected_keys)} keys are unexpected.")
+ print(f"Loaded checkpoint from {path}. {len(state_dict)} keys are available.")
diff --git a/diffsynth/metrics/__init__.py b/diffsynth/metrics/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..1eb52ccdead5a31854e8bb5f3b97810c1b6a28ed
--- /dev/null
+++ b/diffsynth/metrics/__init__.py
@@ -0,0 +1,32 @@
+from ..core import ModelConfig
+from .aesthetic import AestheticMetric
+from .base import Metric
+from .bioclip import BioCLIPMetric
+from .clip import CLIPMetric
+from .fid import FIDMetric
+from .hpsv2 import HPSv2Metric
+from .hpsv3 import HPSv3Metric
+from .image_reward import ImageRewardMetric
+from .lpips import LPIPSMetric
+from .pickscore import PickScoreMetric
+from .qwen_image_bench import QwenImageBenchMetric
+from .unified_reward_2 import UnifiedReward2Metric
+from .unified_reward_edit import UnifiedRewardEditMetric
+
+
+__all__ = [
+ "Metric",
+ "ModelConfig",
+ "PickScoreMetric",
+ "ImageRewardMetric",
+ "HPSv2Metric",
+ "HPSv3Metric",
+ "CLIPMetric",
+ "BioCLIPMetric",
+ "AestheticMetric",
+ "FIDMetric",
+ "LPIPSMetric",
+ "QwenImageBenchMetric",
+ "UnifiedReward2Metric",
+ "UnifiedRewardEditMetric",
+]
diff --git a/diffsynth/metrics/aesthetic.py b/diffsynth/metrics/aesthetic.py
new file mode 100644
index 0000000000000000000000000000000000000000..23d9b19571dda591bf009b097a721f0ddbbf5190
--- /dev/null
+++ b/diffsynth/metrics/aesthetic.py
@@ -0,0 +1,42 @@
+import torch
+from ..core import ModelConfig
+from ..core.device.npu_compatible_device import get_device_type
+from ..models.aesthetic import AestheticModel
+from .base import Metric
+from transformers import CLIPImageProcessor
+
+class AestheticMetric(Metric):
+ def __init__(self, model: AestheticModel):
+ super().__init__()
+ self.model = model
+
+ @classmethod
+ def from_pretrained(
+ cls,
+ model_config: ModelConfig = ModelConfig(model_id="DiffSynth-Studio/ImageMetrics", origin_file_pattern="Aesthetic/model.safetensors"),
+ processor_config: ModelConfig = ModelConfig(model_id="DiffSynth-Studio/ImageMetrics", origin_file_pattern="Aesthetic/"),
+ torch_dtype: torch.dtype = None,
+ device: torch.device = get_device_type(),
+ processor_kwargs: dict = None,
+ vram_limit: float = None,
+ ):
+
+ processor_kwargs = processor_kwargs or {}
+ model_pool = cls.download_and_load_models([model_config], torch_dtype=torch_dtype, device=device, vram_limit=vram_limit)
+ model = model_pool.fetch_model("image_metrics_aesthetic")
+ processor_config.download_if_necessary()
+ model.processor = CLIPImageProcessor.from_pretrained(processor_config.path, **processor_kwargs)
+ model.layers = model.layers.float()
+ model = model.eval()
+ return cls(model)
+
+ @torch.no_grad()
+ def score(self, images):
+ scores = self.model(images)
+ return self.tensor_to_list(scores)
+
+ def compute(self, images):
+ return self.score(images)
+
+ def forward(self, images):
+ return self.score(images)
diff --git a/diffsynth/metrics/base.py b/diffsynth/metrics/base.py
new file mode 100644
index 0000000000000000000000000000000000000000..784bf8eb1b2ef5281342c8bb493e333096f697c4
--- /dev/null
+++ b/diffsynth/metrics/base.py
@@ -0,0 +1,28 @@
+import torch
+from ..core import ModelConfig
+from ..models.model_loader import ModelPool
+
+class Metric(torch.nn.Module):
+
+ @staticmethod
+ def tensor_to_list(value):
+ if torch.is_tensor(value):
+ value = value.detach().cpu().tolist()
+ return value if isinstance(value, list) else [value]
+
+ @staticmethod
+ def download_and_load_models(model_configs: list[ModelConfig], torch_dtype: torch.dtype = torch.float32, device="cuda", vram_limit: float = None):
+ model_pool = ModelPool()
+ for model_config in model_configs:
+ model_config.download_if_necessary()
+ vram_config = model_config.vram_config()
+ vram_config["computation_dtype"] = vram_config["computation_dtype"] or torch_dtype or torch.float32
+ vram_config["computation_device"] = vram_config["computation_device"] or device
+ model_pool.auto_load_model(
+ model_config.path,
+ vram_config=vram_config,
+ vram_limit=vram_limit,
+ clear_parameters=model_config.clear_parameters,
+ state_dict=model_config.state_dict,
+ )
+ return model_pool
diff --git a/diffsynth/metrics/bioclip.py b/diffsynth/metrics/bioclip.py
new file mode 100644
index 0000000000000000000000000000000000000000..35bc0c9fa7c40f2a76b17505d48b64e6c756926f
--- /dev/null
+++ b/diffsynth/metrics/bioclip.py
@@ -0,0 +1,45 @@
+import torch
+from transformers import CLIPTokenizer
+from ..core import ModelConfig
+from ..core.device.npu_compatible_device import get_device_type
+from ..models.bioclip import BioCLIPv2Compute
+from .base import Metric
+
+
+class BioCLIPMetric(Metric):
+ def __init__(self, model: BioCLIPv2Compute):
+ super().__init__()
+ self.model = model
+
+ @classmethod
+ def from_pretrained(
+ cls,
+ model_config: ModelConfig = ModelConfig(model_id="DiffSynth-Studio/ImageMetrics", origin_file_pattern="BioCLIPv2/open_clip_model.safetensors"),
+ tokenizer_config: ModelConfig = ModelConfig(model_id="DiffSynth-Studio/ImageMetrics", origin_file_pattern="BioCLIPv2/"),
+ torch_dtype: torch.dtype = None,
+ device: torch.device = get_device_type(),
+ max_length: int = 77,
+ vram_limit: float = None,
+ ):
+ model_pool = cls.download_and_load_models([model_config], torch_dtype=torch_dtype, device=device, vram_limit=vram_limit)
+ model = model_pool.fetch_model("image_metrics_bioclip_v2")
+ tokenizer_config.download_if_necessary()
+ tokenizer = CLIPTokenizer.from_pretrained(tokenizer_config.path)
+ model = BioCLIPv2Compute(model=model, tokenizer=tokenizer, max_length=max_length).eval()
+ return cls(model)
+
+ @torch.no_grad()
+ def score(self, prompt, images):
+ scores = self.model(prompt, images)
+ return self.tensor_to_list(scores)
+
+ @torch.no_grad()
+ def similarity_matrix(self, prompt, images):
+ scores = self.model.similarity_matrix(prompt, images)
+ return self.tensor_to_list(scores)
+
+ def compute(self, prompt, images):
+ return self.score(prompt, images)
+
+ def forward(self, prompt, images):
+ return self.score(prompt, images)
diff --git a/diffsynth/metrics/clip.py b/diffsynth/metrics/clip.py
new file mode 100644
index 0000000000000000000000000000000000000000..f15ad95f0fccfa154562fe7552e34feb7f56ea42
--- /dev/null
+++ b/diffsynth/metrics/clip.py
@@ -0,0 +1,55 @@
+import torch
+from transformers import AutoProcessor
+from ..core import ModelConfig
+from ..core.device.npu_compatible_device import get_device_type
+from ..models.clip import CLIPModel
+from .base import Metric
+
+class CLIPMetric(Metric):
+ def __init__(self, model: CLIPModel):
+ super().__init__()
+ self.model = model
+
+ @classmethod
+ def from_pretrained(
+ cls,
+ model_config: ModelConfig = ModelConfig(model_id="DiffSynth-Studio/ImageMetrics", origin_file_pattern="CLIP-ViT-H-14-laion2B-s32B-b79K/model.safetensors"),
+ processor_config: ModelConfig = ModelConfig(model_id="DiffSynth-Studio/ImageMetrics", origin_file_pattern="CLIP-ViT-H-14-laion2B-s32B-b79K/"),
+ torch_dtype: torch.dtype = None,
+ device: torch.device = get_device_type(),
+ max_length: int = 77,
+ processor_kwargs: dict = None,
+ vram_limit: float = None,
+ ):
+
+ processor_kwargs = processor_kwargs or {}
+ model_pool = cls.download_and_load_models([model_config], torch_dtype=torch_dtype, device=device, vram_limit=vram_limit)
+ model = model_pool.fetch_model("image_metrics_clip_hf")
+ processor_config.download_if_necessary()
+ processor = AutoProcessor.from_pretrained(processor_config.path, **processor_kwargs)
+ model = CLIPModel(model=model, processor=processor, max_length=max_length).eval()
+ return cls(model)
+
+ @torch.no_grad()
+ def score(
+ self,
+ prompt: str | list[str],
+ images,
+ ):
+ scores = self.model(prompt, images)
+ return self.tensor_to_list(scores)
+
+ @torch.no_grad()
+ def similarity_matrix(
+ self,
+ prompt: str | list[str],
+ images,
+ ):
+ scores = self.model.similarity_matrix(prompt, images)
+ return self.tensor_to_list(scores)
+
+ def compute(self, prompt: str | list[str], images):
+ return self.score(prompt, images)
+
+ def forward(self, prompt: str | list[str], images):
+ return self.score(prompt, images)
diff --git a/diffsynth/metrics/fid.py b/diffsynth/metrics/fid.py
new file mode 100644
index 0000000000000000000000000000000000000000..1cc4429c1fbef55445870db710327e83ac143a71
--- /dev/null
+++ b/diffsynth/metrics/fid.py
@@ -0,0 +1,37 @@
+import torch
+
+from ..core import ModelConfig
+from ..core.device.npu_compatible_device import get_device_type
+from ..models.fid import FIDModel
+from .base import Metric
+
+
+class FIDMetric(Metric):
+ def __init__(self, model: FIDModel):
+ super().__init__()
+ self.model = model
+
+ @classmethod
+ def from_pretrained(
+ cls,
+ model_config: ModelConfig = ModelConfig(model_id="DiffSynth-Studio/ImageMetrics", origin_file_pattern="FID/model.safetensors"),
+ device: torch.device = get_device_type(),
+ batch_size: int = 16,
+ num_workers: int = 0,
+ vram_limit: float = None,
+ ):
+ model_pool = cls.download_and_load_models([model_config], torch_dtype=torch.float32, device=device, vram_limit=vram_limit)
+ model = model_pool.fetch_model("image_metrics_fid_inception")
+ model = FIDModel(model=model, device=device, batch_size=batch_size, num_workers=num_workers)
+ return cls(model)
+
+ @torch.no_grad()
+ def compute(self, reference_images, generated_images, batch_size: int = None, num_workers: int = None):
+ score = self.model.compute(reference_images, generated_images, batch_size=batch_size, num_workers=num_workers)
+ return score.detach().cpu().item() if torch.is_tensor(score) else float(score)
+
+ def statistics(self, images, batch_size: int = None, num_workers: int = None):
+ return self.model.statistics(images, batch_size=batch_size, num_workers=num_workers)
+
+ def forward(self, reference_images, generated_images):
+ return self.compute(reference_images, generated_images)
diff --git a/diffsynth/metrics/hpsv2.py b/diffsynth/metrics/hpsv2.py
new file mode 100644
index 0000000000000000000000000000000000000000..309bb8d3a49c5b5a8fa4d8f719496b849d80b694
--- /dev/null
+++ b/diffsynth/metrics/hpsv2.py
@@ -0,0 +1,41 @@
+import torch
+from transformers import AutoProcessor
+from ..core import ModelConfig
+from ..core.device.npu_compatible_device import get_device_type
+from ..models.hpsv2 import HPSv2Model
+from .base import Metric
+
+class HPSv2Metric(Metric):
+ def __init__(self, model: HPSv2Model):
+ super().__init__()
+ self.model = model
+
+ @classmethod
+ def from_pretrained(
+ cls,
+ model_config: ModelConfig = ModelConfig(model_id="DiffSynth-Studio/ImageMetrics", origin_file_pattern="HPSv2/model.safetensors"),
+ processor_config: ModelConfig = ModelConfig(model_id="DiffSynth-Studio/ImageMetrics", origin_file_pattern="HPSv2/"),
+ torch_dtype: torch.dtype = None,
+ device: torch.device = get_device_type(),
+ processor_kwargs: dict = None,
+ vram_limit: float = None,
+ ):
+
+ processor_kwargs = processor_kwargs or {}
+ model_pool = cls.download_and_load_models([model_config], torch_dtype=torch_dtype, device=device, vram_limit=vram_limit)
+ model = model_pool.fetch_model("image_metrics_hpsv2")
+ processor_config.download_if_necessary()
+ processor = AutoProcessor.from_pretrained(processor_config.path, **processor_kwargs)
+ model = HPSv2Model(model=model, processor=processor).eval()
+ return cls(model)
+
+ @torch.no_grad()
+ def score(self, prompt: str | list[str], images):
+ scores = self.model(prompt, images)
+ return self.tensor_to_list(scores)
+
+ def compute(self, prompt: str | list[str], images):
+ return self.score(prompt, images)
+
+ def forward(self, prompt: str | list[str], images):
+ return self.score(prompt, images)
diff --git a/diffsynth/metrics/hpsv3.py b/diffsynth/metrics/hpsv3.py
new file mode 100644
index 0000000000000000000000000000000000000000..83b93f45b4e84f39410901051563f225922938b2
--- /dev/null
+++ b/diffsynth/metrics/hpsv3.py
@@ -0,0 +1,63 @@
+from transformers import AutoProcessor
+import torch
+from ..core import ModelConfig
+from ..core.device.npu_compatible_device import get_device_type
+from ..models.hpsv3 import HPSv3Model
+from .base import Metric
+
+
+class HPSv3Metric(Metric):
+ def __init__(self, model: HPSv3Model):
+ super().__init__()
+ self.model = model
+
+ @classmethod
+ def from_pretrained(
+ cls,
+ model_config: ModelConfig = ModelConfig(model_id="DiffSynth-Studio/ImageMetrics", origin_file_pattern="HPSv3/model.safetensors"),
+ processor_config: ModelConfig = ModelConfig(model_id="DiffSynth-Studio/ImageMetrics", origin_file_pattern="HPSv3/"),
+ torch_dtype: torch.dtype = torch.bfloat16,
+ device: torch.device = get_device_type(),
+ score_index: int = 0,
+ use_special_tokens: bool = True,
+ max_pixels: int = 256 * 28 * 28,
+ min_pixels: int = 256 * 28 * 28,
+ processor_kwargs: dict = None,
+ vram_limit: float = None,
+ ):
+
+ processor_kwargs = processor_kwargs or {}
+ model_pool = cls.download_and_load_models([model_config], torch_dtype=torch_dtype, device=device, vram_limit=vram_limit)
+ model = model_pool.fetch_model("image_metrics_hpsv3")
+ processor_config.download_if_necessary()
+ processor = AutoProcessor.from_pretrained(processor_config.path, padding_side="right", **processor_kwargs)
+ if use_special_tokens:
+ special_tokens = ["<|Reward|>"]
+ processor.tokenizer.add_special_tokens({"additional_special_tokens": special_tokens})
+ model.special_token_ids = processor.tokenizer.convert_tokens_to_ids(special_tokens)
+ model.reward_token = "special"
+ model.config.tokenizer_padding_side = processor.tokenizer.padding_side
+ model.config.pad_token_id = processor.tokenizer.pad_token_id
+ if hasattr(model.config, "text_config"):
+ model.config.text_config.pad_token_id = processor.tokenizer.pad_token_id
+ model.rm_head.to(torch.float32)
+ model = HPSv3Model(
+ model=model,
+ processor=processor,
+ use_special_tokens=use_special_tokens,
+ max_pixels=max_pixels,
+ min_pixels=min_pixels,
+ score_index=score_index,
+ ).eval()
+ return cls(model)
+
+ @torch.no_grad()
+ def score(self, prompt: str | list[str], images):
+ scores = self.model(prompt, images)
+ return self.tensor_to_list(scores)
+
+ def compute(self, prompt: str | list[str], images):
+ return self.score(prompt, images)
+
+ def forward(self, prompt: str | list[str], images):
+ return self.score(prompt, images)
diff --git a/diffsynth/metrics/image_reward.py b/diffsynth/metrics/image_reward.py
new file mode 100644
index 0000000000000000000000000000000000000000..85a042e8b7789d436985cd25b0d9ed3ad40c8564
--- /dev/null
+++ b/diffsynth/metrics/image_reward.py
@@ -0,0 +1,48 @@
+import torch
+from transformers import BertTokenizer
+from ..core import ModelConfig
+from ..core.device.npu_compatible_device import get_device_type
+from ..models.image_reward import ImageRewardModel
+from .base import Metric
+
+class ImageRewardMetric(Metric):
+ def __init__(self, model: ImageRewardModel):
+ super().__init__()
+ self.model = model
+
+ @classmethod
+ def from_pretrained(
+ cls,
+ model_config: ModelConfig = ModelConfig(model_id="DiffSynth-Studio/ImageMetrics", origin_file_pattern="ImageReward/model.safetensors"),
+ tokenizer_config: ModelConfig = ModelConfig(model_id="DiffSynth-Studio/ImageMetrics", origin_file_pattern="ImageReward/"),
+ torch_dtype: torch.dtype = None,
+ device: torch.device = get_device_type(),
+ max_length: int = 35,
+ tokenizer_kwargs: dict = None,
+ vram_limit: float = None,
+ ):
+
+ tokenizer_kwargs = tokenizer_kwargs or {}
+ model_pool = cls.download_and_load_models([model_config], torch_dtype=torch_dtype, device=device, vram_limit=vram_limit)
+ model = model_pool.fetch_model("image_metrics_image_reward")
+ tokenizer_config.download_if_necessary()
+ tokenizer = BertTokenizer.from_pretrained(tokenizer_config.path, **tokenizer_kwargs)
+ tokenizer.add_special_tokens({"bos_token": "[DEC]"})
+ tokenizer.add_special_tokens({"additional_special_tokens": ["[ENC]"]})
+ tokenizer.enc_token_id = tokenizer.convert_tokens_to_ids("[ENC]")
+ model.tokenizer = tokenizer
+ model.max_length = max_length
+ model.mlp = model.mlp.float()
+ model = model.eval()
+ return cls(model)
+
+ @torch.no_grad()
+ def score(self, prompt: str | list[str], images):
+ scores = self.model(prompt, images)
+ return self.tensor_to_list(scores)
+
+ def compute(self, prompt: str | list[str], images):
+ return self.score(prompt, images)
+
+ def forward(self, prompt: str | list[str], images):
+ return self.score(prompt, images)
diff --git a/diffsynth/metrics/lpips.py b/diffsynth/metrics/lpips.py
new file mode 100644
index 0000000000000000000000000000000000000000..dc5ae363463440706e3c311cf7756aa4841ddc79
--- /dev/null
+++ b/diffsynth/metrics/lpips.py
@@ -0,0 +1,63 @@
+import torch
+
+from ..core import ModelConfig
+from ..core.device.npu_compatible_device import get_device_type
+from ..models.lpips import LPIPSModel, LPIPS_NET_CHOICES, LPIPSCompute
+from .base import Metric
+
+
+_LPIPS_DEFAULT_FILES = {
+ "alex": "LPIPS/alexnet.safetensors",
+ "vgg": "LPIPS/vgg.safetensors",
+ "squeeze": "LPIPS/squeezenet.safetensors",
+}
+
+_LPIPS_MODEL_NAMES = {
+ "alex": "image_metrics_lpips_alex",
+ "vgg": "image_metrics_lpips_vgg",
+ "squeeze": "image_metrics_lpips_squeeze",
+}
+
+
+class LPIPSMetric(Metric):
+ def __init__(self, model: LPIPSCompute):
+ super().__init__()
+ self.model = model
+
+ @classmethod
+ def from_pretrained(
+ cls,
+ net: str = "alex",
+ model_config: ModelConfig = None,
+ device: torch.device = get_device_type(),
+ batch_size: int = 16,
+ target_size: int = 512,
+ vram_limit: float = None,
+ ):
+ if net not in LPIPS_NET_CHOICES:
+ raise ValueError(f"net must be one of {LPIPS_NET_CHOICES}, got {net!r}")
+ if model_config is None:
+ model_config = ModelConfig(
+ model_id="DiffSynth-Studio/ImageMetrics",
+ origin_file_pattern=_LPIPS_DEFAULT_FILES[net],
+ )
+ model_pool = cls.download_and_load_models([model_config], torch_dtype=torch.float32, device=device, vram_limit=vram_limit)
+ backbone = model_pool.fetch_model(_LPIPS_MODEL_NAMES[net])
+ if backbone is None:
+ raise RuntimeError(
+ f"Failed to load LPIPS model for net={net!r}. The provided weights do not match the registered hash for {_LPIPS_MODEL_NAMES[net]}."
+ )
+ compute_model = LPIPSCompute(
+ model=backbone,
+ device=device,
+ batch_size=batch_size,
+ target_size=target_size,
+ )
+ return cls(compute_model)
+
+ @torch.no_grad()
+ def compute(self, image_a, image_b) -> float:
+ return self.model.compute(image_a, image_b)
+
+ def forward(self, image_a, image_b):
+ return self.compute(image_a, image_b)
diff --git a/diffsynth/metrics/pickscore.py b/diffsynth/metrics/pickscore.py
new file mode 100644
index 0000000000000000000000000000000000000000..608e544848e47ac323b955f98e92c1b458829043
--- /dev/null
+++ b/diffsynth/metrics/pickscore.py
@@ -0,0 +1,59 @@
+from transformers import AutoProcessor
+import torch
+from ..core import ModelConfig
+from ..core.device.npu_compatible_device import get_device_type
+from ..models.pickscore import PickScoreModel
+from .base import Metric
+
+class PickScoreMetric(Metric):
+ def __init__(self, model: PickScoreModel):
+ super().__init__()
+ self.model = model
+
+ @classmethod
+ def from_pretrained(
+ cls,
+ model_config: ModelConfig = ModelConfig(model_id="DiffSynth-Studio/ImageMetrics", origin_file_pattern="PickScore/model.safetensors"),
+ processor_config: ModelConfig = ModelConfig(model_id="DiffSynth-Studio/ImageMetrics", origin_file_pattern="PickScore/"),
+ torch_dtype: torch.dtype = None,
+ device: torch.device = get_device_type(),
+ max_length: int = 77,
+ processor_kwargs: dict = None,
+ vram_limit: float = None,
+ ):
+
+ processor_kwargs = processor_kwargs or {}
+ model_pool = cls.download_and_load_models([model_config], torch_dtype=torch_dtype, device=device, vram_limit=vram_limit)
+ model = model_pool.fetch_model("image_metrics_clip_hf")
+ processor_config.download_if_necessary()
+ processor = AutoProcessor.from_pretrained(processor_config.path, **processor_kwargs)
+ model = PickScoreModel(model=model, processor=processor, max_length=max_length).eval()
+ return cls(model)
+
+ @torch.no_grad()
+ def score(
+ self,
+ prompt: str | list[str],
+ images,
+ ):
+ scores = self.model(prompt, images)
+ return self.tensor_to_list(scores)
+
+ @torch.no_grad()
+ def probabilities(
+ self,
+ prompt: str | list[str],
+ images,
+ ):
+ scores = self.model(prompt, images)
+ probabilities = torch.softmax(scores, dim=-1)
+ return self.tensor_to_list(probabilities)
+
+ def calc_probs(self, prompt: str | list[str], images):
+ return self.probabilities(prompt, images)
+
+ def compute(self, prompt: str | list[str], images):
+ return self.score(prompt, images)
+
+ def forward(self, prompt: str | list[str], images):
+ return self.score(prompt, images)
diff --git a/diffsynth/metrics/qwen_image_bench.py b/diffsynth/metrics/qwen_image_bench.py
new file mode 100644
index 0000000000000000000000000000000000000000..a7a3f65580c42f9c492688c38cf883eea40a7d16
--- /dev/null
+++ b/diffsynth/metrics/qwen_image_bench.py
@@ -0,0 +1,70 @@
+import torch
+from transformers import AutoProcessor
+from ..core import ModelConfig
+from ..core.device.npu_compatible_device import get_device_type
+from ..models.qwen_image_bench import QwenImageBenchModel
+from .base import Metric
+from transformers.utils import logging
+logging.set_verbosity_error()
+
+
+class QwenImageBenchMetric(Metric):
+ def __init__(self, model: QwenImageBenchModel):
+ super().__init__()
+ self.model = model
+
+ @classmethod
+ def from_pretrained(
+ cls,
+ model_config: ModelConfig = ModelConfig(
+ model_id="Qwen/Qwen-Image-Bench",
+ origin_file_pattern="model-*.safetensors",
+ ),
+ processor_config: ModelConfig = ModelConfig(
+ model_id="Qwen/Qwen-Image-Bench",
+ origin_file_pattern="",
+ ),
+ torch_dtype: torch.dtype = None,
+ device: torch.device = get_device_type(),
+ max_new_tokens: int = 4096,
+ resize_long_edge: int = 1024,
+ processor_kwargs: dict = None,
+ vram_limit: float = None,
+ ):
+ processor_kwargs = processor_kwargs or {}
+ model_pool = cls.download_and_load_models(
+ [model_config],
+ torch_dtype=torch_dtype or torch.bfloat16,
+ device=device,
+ vram_limit=vram_limit,
+ )
+ model = model_pool.fetch_model("image_metrics_qwen_image_bench")
+ if model is None:
+ raise ValueError("Cannot find model: image_metrics_qwen_image_bench")
+ if hasattr(model, "model"):
+ model = model.model
+
+ processor_config.download_if_necessary()
+ processor = AutoProcessor.from_pretrained(processor_config.path, **processor_kwargs)
+ model = QwenImageBenchModel(
+ model=model,
+ processor=processor,
+ max_new_tokens=max_new_tokens,
+ resize_long_edge=resize_long_edge,
+ ).eval()
+ return cls(model)
+
+ @torch.no_grad()
+ def evaluate(self, prompt: str | list[str] | None, images, dimensions=None):
+ return self.model(prompt, images, dimensions=dimensions)
+
+ @torch.no_grad()
+ def score(self, prompt: str | list[str] | None, images, dimensions=None):
+ outputs = self.evaluate(prompt, images, dimensions=dimensions)
+ return [self.model._primary_score(output) for output in outputs]
+
+ def compute(self, prompt: str | list[str] | None, images, dimensions=None):
+ return self.score(prompt, images, dimensions=dimensions)
+
+ def forward(self, prompt: str | list[str] | None, images, dimensions=None):
+ return self.score(prompt, images, dimensions=dimensions)
diff --git a/diffsynth/metrics/unified_reward_2.py b/diffsynth/metrics/unified_reward_2.py
new file mode 100644
index 0000000000000000000000000000000000000000..abad17a38e3cd2da430b148f29624b2321238da8
--- /dev/null
+++ b/diffsynth/metrics/unified_reward_2.py
@@ -0,0 +1,69 @@
+import torch
+from transformers import AutoProcessor
+
+from ..core import ModelConfig
+from ..core.device.npu_compatible_device import get_device_type
+from ..models.unified_reward_2 import UnifiedReward2Model
+from .base import Metric
+from transformers.utils import logging
+logging.set_verbosity_error()
+
+
+class UnifiedReward2Metric(Metric):
+ def __init__(self, model: UnifiedReward2Model):
+ super().__init__()
+ self.model = model
+
+ @classmethod
+ def from_pretrained(
+ cls,
+ model_config: ModelConfig = ModelConfig(
+ model_id="DiffSynth-Studio/ImageMetrics",
+ origin_file_pattern="UnifiedReward-2.0-qwen35-9b/model-*.safetensors",
+ ),
+ processor_config: ModelConfig = ModelConfig(
+ model_id="DiffSynth-Studio/ImageMetrics",
+ origin_file_pattern="UnifiedReward-2.0-qwen35-9b/",
+ ),
+ torch_dtype: torch.dtype = None,
+ device: torch.device = get_device_type(),
+ max_new_tokens: int = 1024,
+ processor_kwargs: dict = None,
+ vram_limit: float = None,
+ ):
+ processor_kwargs = processor_kwargs or {}
+ model_pool = cls.download_and_load_models(
+ [model_config],
+ torch_dtype=torch_dtype or torch.bfloat16,
+ device=device,
+ vram_limit=vram_limit,
+ )
+ model = model_pool.fetch_model("image_metrics_unified_reward_2")
+ if model is None:
+ raise ValueError("Cannot find model: image_metrics_unified_reward_2")
+ if hasattr(model, "model"):
+ model = model.model
+
+ processor_config.download_if_necessary()
+ processor = AutoProcessor.from_pretrained(processor_config.path, **processor_kwargs)
+ model = UnifiedReward2Model(
+ model=model,
+ processor=processor,
+ max_new_tokens=max_new_tokens,
+ ).eval()
+ return cls(model)
+
+ @torch.no_grad()
+ def evaluate(self, prompt: str | list[str] | None, images):
+ return self.model(prompt, images)
+
+ @torch.no_grad()
+ def score(self, prompt: str | list[str] | None, images):
+ outputs = self.evaluate(prompt, images)
+ return [self.model._primary_score(output) for output in outputs]
+
+ def compute(self, prompt: str | list[str] | None, images):
+ return self.score(prompt, images)
+
+ def forward(self, prompt: str | list[str] | None, images):
+ return self.score(prompt, images)
diff --git a/diffsynth/metrics/unified_reward_edit.py b/diffsynth/metrics/unified_reward_edit.py
new file mode 100644
index 0000000000000000000000000000000000000000..fefdc724c23abc2b00a146be711b1e940c32ddb8
--- /dev/null
+++ b/diffsynth/metrics/unified_reward_edit.py
@@ -0,0 +1,97 @@
+import torch
+from transformers import AutoProcessor
+from ..core import ModelConfig
+from ..core.device.npu_compatible_device import get_device_type
+from ..models.unified_reward_edit import UnifiedRewardEditModel
+from .base import Metric
+from transformers.utils import logging
+logging.set_verbosity_error()
+
+
+DEFAULT_UNIFIED_REWARD_TASK = "edit_pointwise_score"
+
+
+class UnifiedRewardEditMetric(Metric):
+ def __init__(self, model: UnifiedRewardEditModel):
+ super().__init__()
+ self.model = model
+
+ @classmethod
+ def from_pretrained(
+ cls,
+ model_config: ModelConfig = ModelConfig(
+ model_id="DiffSynth-Studio/ImageMetrics",
+ origin_file_pattern="UnifiedReward-Edit-qwen3vl-8b/model-*.safetensors",
+ ),
+ processor_config: ModelConfig = ModelConfig(
+ model_id="DiffSynth-Studio/ImageMetrics",
+ origin_file_pattern="UnifiedReward-Edit-qwen3vl-8b/",
+ ),
+ torch_dtype: torch.dtype = None,
+ device: torch.device = get_device_type(),
+ task: str = DEFAULT_UNIFIED_REWARD_TASK,
+ max_new_tokens: int = 256,
+ processor_kwargs: dict = None,
+ vram_limit: float = None,
+ ):
+ processor_kwargs = processor_kwargs or {}
+ model_pool = cls.download_and_load_models(
+ [model_config],
+ torch_dtype=torch_dtype or torch.bfloat16,
+ device=device,
+ vram_limit=vram_limit,
+ )
+ model = model_pool.fetch_model("image_metrics_unified_reward_edit")
+ if model is None:
+ model = model_pool.fetch_model("joyai_image_text_encoder")
+ if model is None:
+ raise ValueError("Cannot find model: image_metrics_unified_reward_edit")
+ if hasattr(model, "model"):
+ model = model.model
+
+ processor_config.download_if_necessary()
+ processor = AutoProcessor.from_pretrained(processor_config.path, **processor_kwargs)
+ model = UnifiedRewardEditModel(
+ model=model,
+ processor=processor,
+ task=task,
+ max_new_tokens=max_new_tokens,
+ ).eval()
+ return cls(model)
+
+
+ @staticmethod
+ def _primary_score(parsed: dict, task: str):
+ if task == "edit_pairwise_rank":
+ winner = parsed.get("winner")
+ if isinstance(winner, (int, float)):
+ return int(winner)
+ if isinstance(winner, str):
+ winner = winner.lower()
+ if "equally" in winner or "tie" in winner or ("image 1" in winner and "image 2" in winner):
+ return 0
+ if "image 1" in winner or "first image" in winner:
+ return 1
+ if "image 2" in winner or "second image" in winner:
+ return 2
+ return 0
+ if task == "edit_pairwise_score":
+ return [parsed.get("image_1_score"), parsed.get("image_2_score")]
+ return parsed.get("score")
+
+
+ @torch.no_grad()
+ def evaluate(self, prompt: str | list[str] | None, images, task: str = DEFAULT_UNIFIED_REWARD_TASK):
+ outputs = self.model(prompt, images, task=task)
+ return [{**output, "score": self._primary_score(output, task)} for output in outputs]
+
+ @torch.no_grad()
+ def score(self, prompt: str | list[str] | None, images, task: str = DEFAULT_UNIFIED_REWARD_TASK):
+ outputs = self.evaluate(prompt, images, task=task)
+ return [output["score"] for output in outputs]
+
+ def compute(self, prompt: str | list[str] | None, images, task: str = DEFAULT_UNIFIED_REWARD_TASK):
+ return self.score(prompt, images, task=task)
+
+ def forward(self, prompt: str | list[str] | None, images, task: str = DEFAULT_UNIFIED_REWARD_TASK):
+ return self.score(prompt, images, task=task)
diff --git a/diffsynth/models/ace_step_conditioner.py b/diffsynth/models/ace_step_conditioner.py
new file mode 100644
index 0000000000000000000000000000000000000000..8f3942d6bf015c5db05c8d53a062b2a2e3b4e2be
--- /dev/null
+++ b/diffsynth/models/ace_step_conditioner.py
@@ -0,0 +1,674 @@
+# Copyright 2025 The ACESTEO Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+import math
+from typing import Optional
+
+import torch
+import torch.nn.functional as F
+from torch import nn
+
+from ..core.attention import attention_forward
+from ..core.gradient import gradient_checkpoint_forward
+
+from transformers.cache_utils import Cache
+from transformers.modeling_flash_attention_utils import FlashAttentionKwargs
+from transformers.modeling_outputs import BaseModelOutput
+from transformers.processing_utils import Unpack
+from transformers.utils import can_return_tuple, logging
+from transformers.models.qwen3.modeling_qwen3 import (
+ Qwen3MLP,
+ Qwen3RMSNorm,
+ Qwen3RotaryEmbedding,
+ apply_rotary_pos_emb,
+)
+
+logger = logging.get_logger(__name__)
+
+
+def create_4d_mask(
+ seq_len: int,
+ dtype: torch.dtype,
+ device: torch.device,
+ attention_mask: Optional[torch.Tensor] = None,
+ sliding_window: Optional[int] = None,
+ is_sliding_window: bool = False,
+ is_causal: bool = True,
+) -> torch.Tensor:
+ indices = torch.arange(seq_len, device=device)
+ diff = indices.unsqueeze(1) - indices.unsqueeze(0)
+ valid_mask = torch.ones((seq_len, seq_len), device=device, dtype=torch.bool)
+ if is_causal:
+ valid_mask = valid_mask & (diff >= 0)
+ if is_sliding_window and sliding_window is not None:
+ if is_causal:
+ valid_mask = valid_mask & (diff <= sliding_window)
+ else:
+ valid_mask = valid_mask & (torch.abs(diff) <= sliding_window)
+ valid_mask = valid_mask.unsqueeze(0).unsqueeze(0)
+ if attention_mask is not None:
+ padding_mask_4d = attention_mask.view(attention_mask.shape[0], 1, 1, seq_len).to(torch.bool)
+ valid_mask = valid_mask & padding_mask_4d
+ min_dtype = torch.finfo(dtype).min
+ mask_tensor = torch.full(valid_mask.shape, min_dtype, dtype=dtype, device=device)
+ mask_tensor.masked_fill_(valid_mask, 0.0)
+ return mask_tensor
+
+
+def pack_sequences(hidden1: torch.Tensor, hidden2: torch.Tensor, mask1: torch.Tensor, mask2: torch.Tensor):
+ hidden_cat = torch.cat([hidden1, hidden2], dim=1)
+ mask_cat = torch.cat([mask1, mask2], dim=1)
+ B, L, D = hidden_cat.shape
+ sort_idx = mask_cat.argsort(dim=1, descending=True, stable=True)
+ hidden_left = torch.gather(hidden_cat, 1, sort_idx.unsqueeze(-1).expand(B, L, D))
+ lengths = mask_cat.sum(dim=1)
+ new_mask = (torch.arange(L, dtype=torch.long, device=hidden_cat.device).unsqueeze(0) < lengths.unsqueeze(1))
+ return hidden_left, new_mask
+
+
+class Lambda(nn.Module):
+ def __init__(self, func):
+ super().__init__()
+ self.func = func
+
+ def forward(self, x):
+ return self.func(x)
+
+
+class AceStepAttention(nn.Module):
+ def __init__(
+ self,
+ hidden_size: int,
+ num_attention_heads: int,
+ num_key_value_heads: int,
+ rms_norm_eps: float,
+ attention_bias: bool,
+ attention_dropout: float,
+ layer_types: list,
+ head_dim: Optional[int] = None,
+ sliding_window: Optional[int] = None,
+ layer_idx: int = 0,
+ is_cross_attention: bool = False,
+ is_causal: bool = False,
+ ):
+ super().__init__()
+ self.layer_idx = layer_idx
+ self.head_dim = head_dim or hidden_size // num_attention_heads
+ self.num_key_value_groups = num_attention_heads // num_key_value_heads
+ self.scaling = self.head_dim ** -0.5
+ self.attention_dropout = attention_dropout
+ if is_cross_attention:
+ is_causal = False
+ self.is_causal = is_causal
+ self.is_cross_attention = is_cross_attention
+
+ self.q_proj = nn.Linear(hidden_size, num_attention_heads * self.head_dim, bias=attention_bias)
+ self.k_proj = nn.Linear(hidden_size, num_key_value_heads * self.head_dim, bias=attention_bias)
+ self.v_proj = nn.Linear(hidden_size, num_key_value_heads * self.head_dim, bias=attention_bias)
+ self.o_proj = nn.Linear(num_attention_heads * self.head_dim, hidden_size, bias=attention_bias)
+ self.q_norm = Qwen3RMSNorm(self.head_dim, eps=rms_norm_eps)
+ self.k_norm = Qwen3RMSNorm(self.head_dim, eps=rms_norm_eps)
+ self.attention_type = layer_types[layer_idx]
+ self.sliding_window = sliding_window if layer_types[layer_idx] == "sliding_attention" else None
+
+ def forward(
+ self,
+ hidden_states: torch.Tensor,
+ attention_mask: Optional[torch.Tensor],
+ past_key_value: Optional[Cache] = None,
+ cache_position: Optional[torch.LongTensor] = None,
+ encoder_hidden_states: Optional[torch.Tensor] = None,
+ position_embeddings: tuple[torch.Tensor, torch.Tensor] = None,
+ output_attentions: Optional[bool] = False,
+ **kwargs: Unpack[FlashAttentionKwargs],
+ ) -> tuple[torch.Tensor, Optional[torch.Tensor], Optional[tuple[torch.Tensor]]]:
+ input_shape = hidden_states.shape[:-1]
+ hidden_shape = (*input_shape, -1, self.head_dim)
+
+ query_states = self.q_norm(self.q_proj(hidden_states).view(hidden_shape)).transpose(1, 2)
+
+ is_cross_attention = self.is_cross_attention and encoder_hidden_states is not None
+
+ if is_cross_attention:
+ encoder_hidden_shape = (*encoder_hidden_states.shape[:-1], -1, self.head_dim)
+ if past_key_value is not None:
+ is_updated = past_key_value.is_updated.get(self.layer_idx)
+ curr_past_key_value = past_key_value.cross_attention_cache
+ if not is_updated:
+ key_states = self.k_norm(self.k_proj(encoder_hidden_states).view(encoder_hidden_shape)).transpose(1, 2)
+ value_states = self.v_proj(encoder_hidden_states).view(encoder_hidden_shape).transpose(1, 2)
+ key_states, value_states = curr_past_key_value.update(key_states, value_states, self.layer_idx)
+ past_key_value.is_updated[self.layer_idx] = True
+ else:
+ key_states = curr_past_key_value.layers[self.layer_idx].keys
+ value_states = curr_past_key_value.layers[self.layer_idx].values
+ else:
+ key_states = self.k_norm(self.k_proj(encoder_hidden_states).view(encoder_hidden_shape)).transpose(1, 2)
+ value_states = self.v_proj(encoder_hidden_states).view(encoder_hidden_shape).transpose(1, 2)
+
+ else:
+ key_states = self.k_norm(self.k_proj(hidden_states).view(hidden_shape)).transpose(1, 2)
+ value_states = self.v_proj(hidden_states).view(hidden_shape).transpose(1, 2)
+ if position_embeddings is not None:
+ cos, sin = position_embeddings
+ query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin)
+
+ if past_key_value is not None:
+ cache_kwargs = {"sin": sin, "cos": cos, "cache_position": cache_position}
+ key_states, value_states = past_key_value.update(key_states, value_states, self.layer_idx, cache_kwargs)
+
+ attn_output = attention_forward(
+ query_states, key_states, value_states,
+ q_pattern="b n s d", k_pattern="b n s d", v_pattern="b n s d", out_pattern="b n s d",
+ window_size=None if attention_mask is None else attention_mask["window_size"],
+ )
+ attn_weights = None
+
+ attn_output = attn_output.transpose(1, 2).flatten(2, 3).contiguous()
+ attn_output = self.o_proj(attn_output)
+ return attn_output, attn_weights
+
+
+class AceStepEncoderLayer(nn.Module):
+ def __init__(
+ self,
+ hidden_size: int,
+ intermediate_size: int,
+ num_attention_heads: int,
+ num_key_value_heads: int,
+ rms_norm_eps: float,
+ attention_bias: bool,
+ attention_dropout: float,
+ layer_types: list,
+ head_dim: Optional[int] = None,
+ sliding_window: Optional[int] = None,
+ layer_idx: int = 0,
+ ):
+ super().__init__()
+ self.hidden_size = hidden_size
+ self.layer_idx = layer_idx
+
+ self.self_attn = AceStepAttention(
+ hidden_size=hidden_size,
+ num_attention_heads=num_attention_heads,
+ num_key_value_heads=num_key_value_heads,
+ rms_norm_eps=rms_norm_eps,
+ attention_bias=attention_bias,
+ attention_dropout=attention_dropout,
+ layer_types=layer_types,
+ head_dim=head_dim,
+ sliding_window=sliding_window,
+ layer_idx=layer_idx,
+ is_cross_attention=False,
+ is_causal=False,
+ )
+ self.input_layernorm = Qwen3RMSNorm(hidden_size, eps=rms_norm_eps)
+ self.post_attention_layernorm = Qwen3RMSNorm(hidden_size, eps=rms_norm_eps)
+
+ mlp_config = type('Config', (), {
+ 'hidden_size': hidden_size,
+ 'intermediate_size': intermediate_size,
+ 'hidden_act': 'silu',
+ })()
+ self.mlp = Qwen3MLP(mlp_config)
+ self.attention_type = layer_types[layer_idx]
+
+ def forward(
+ self,
+ hidden_states: torch.Tensor,
+ position_embeddings: tuple[torch.Tensor, torch.Tensor],
+ attention_mask: Optional[torch.Tensor] = None,
+ position_ids: Optional[torch.LongTensor] = None,
+ output_attentions: Optional[bool] = False,
+ **kwargs,
+ ) -> tuple[torch.FloatTensor, Optional[tuple[torch.FloatTensor, torch.FloatTensor]]]:
+ residual = hidden_states
+ hidden_states = self.input_layernorm(hidden_states)
+ hidden_states, self_attn_weights = self.self_attn(
+ hidden_states=hidden_states,
+ position_embeddings=position_embeddings,
+ attention_mask=attention_mask,
+ position_ids=position_ids,
+ output_attentions=output_attentions,
+ use_cache=False,
+ past_key_value=None,
+ **kwargs,
+ )
+ hidden_states = residual + hidden_states
+
+ residual = hidden_states
+ hidden_states = self.post_attention_layernorm(hidden_states)
+ hidden_states = self.mlp(hidden_states)
+ hidden_states = residual + hidden_states
+
+ outputs = (hidden_states,)
+ if output_attentions:
+ outputs += (self_attn_weights,)
+ return outputs
+
+
+class AceStepLyricEncoder(nn.Module):
+ def __init__(
+ self,
+ hidden_size: int = 2048,
+ intermediate_size: int = 6144,
+ num_hidden_layers: int = 24,
+ num_attention_heads: int = 16,
+ num_key_value_heads: int = 8,
+ rms_norm_eps: float = 1e-6,
+ attention_bias: bool = False,
+ attention_dropout: float = 0.0,
+ layer_types: Optional[list] = None,
+ head_dim: Optional[int] = None,
+ sliding_window: Optional[int] = 128,
+ use_sliding_window: bool = True,
+ use_cache: bool = True,
+ rope_theta: float = 1000000,
+ max_position_embeddings: int = 32768,
+ initializer_range: float = 0.02,
+ text_hidden_dim: int = 1024,
+ num_lyric_encoder_hidden_layers: int = 8,
+ **kwargs,
+ ):
+ super().__init__()
+ self.num_lyric_encoder_hidden_layers = num_lyric_encoder_hidden_layers
+ self.text_hidden_dim = text_hidden_dim
+ self.hidden_size = hidden_size
+ self.intermediate_size = intermediate_size
+ self.num_hidden_layers = num_hidden_layers
+ self.num_attention_heads = num_attention_heads
+ self.num_key_value_heads = num_key_value_heads
+ self.rms_norm_eps = rms_norm_eps
+ self.attention_bias = attention_bias
+ self.attention_dropout = attention_dropout
+ self.layer_types = layer_types or (["sliding_attention", "full_attention"] * (num_hidden_layers // 2))
+ self.head_dim = head_dim or hidden_size // num_attention_heads
+ self.sliding_window = sliding_window
+ self.use_sliding_window = use_sliding_window
+ self.use_cache = use_cache
+ self.rope_theta = rope_theta
+ self.max_position_embeddings = max_position_embeddings
+ self.initializer_range = initializer_range
+ self._attn_implementation = kwargs.get("_attn_implementation", "sdpa")
+
+ self.embed_tokens = nn.Linear(text_hidden_dim, hidden_size)
+ self.norm = Qwen3RMSNorm(hidden_size, eps=rms_norm_eps)
+ rope_config = type('RopeConfig', (), {
+ 'hidden_size': hidden_size,
+ 'num_attention_heads': num_attention_heads,
+ 'num_key_value_heads': num_key_value_heads,
+ 'head_dim': head_dim,
+ 'max_position_embeddings': max_position_embeddings,
+ 'rope_theta': rope_theta,
+ 'rope_parameters': {'rope_type': 'default', 'rope_theta': rope_theta},
+ 'rms_norm_eps': rms_norm_eps,
+ 'attention_bias': attention_bias,
+ 'attention_dropout': attention_dropout,
+ 'hidden_act': 'silu',
+ 'intermediate_size': intermediate_size,
+ 'layer_types': self.layer_types,
+ 'sliding_window': sliding_window,
+ '_attn_implementation': self._attn_implementation,
+ })()
+ self.rotary_emb = Qwen3RotaryEmbedding(rope_config)
+ self.gradient_checkpointing = False
+
+ self.layers = nn.ModuleList([
+ AceStepEncoderLayer(
+ hidden_size=hidden_size,
+ intermediate_size=intermediate_size,
+ num_attention_heads=num_attention_heads,
+ num_key_value_heads=num_key_value_heads,
+ rms_norm_eps=rms_norm_eps,
+ attention_bias=attention_bias,
+ attention_dropout=attention_dropout,
+ layer_types=self.layer_types,
+ head_dim=head_dim,
+ sliding_window=sliding_window,
+ layer_idx=layer_idx,
+ )
+ for layer_idx in range(num_lyric_encoder_hidden_layers)
+ ])
+
+
+ @can_return_tuple
+ def forward(
+ self,
+ input_ids: Optional[torch.LongTensor] = None,
+ attention_mask: Optional[torch.Tensor] = None,
+ position_ids: Optional[torch.LongTensor] = None,
+ inputs_embeds: Optional[torch.FloatTensor] = None,
+ output_attentions: Optional[bool] = None,
+ output_hidden_states: Optional[bool] = None,
+ **flash_attn_kwargs: Unpack[FlashAttentionKwargs],
+ ) -> BaseModelOutput:
+ output_attentions = output_attentions if output_attentions is not None else False
+ output_hidden_states = output_hidden_states if output_hidden_states is not None else False
+
+ assert input_ids is None, "Only `inputs_embeds` is supported for the lyric encoder."
+ assert attention_mask is not None, "Attention mask must be provided for the lyric encoder."
+ assert inputs_embeds is not None, "Inputs embeddings must be provided for the lyric encoder."
+
+ inputs_embeds = self.embed_tokens(inputs_embeds)
+ cache_position = torch.arange(0, inputs_embeds.shape[1], device=inputs_embeds.device)
+
+ if position_ids is None:
+ position_ids = cache_position.unsqueeze(0)
+
+ seq_len = inputs_embeds.shape[1]
+ dtype = inputs_embeds.dtype
+ device = inputs_embeds.device
+
+ full_attn_mask = None
+ sliding_attn_mask = None
+ if self.use_sliding_window:
+ sliding_attn_mask = {"window_size": self.sliding_window}
+
+ self_attn_mask_mapping = {
+ "full_attention": full_attn_mask,
+ "sliding_attention": sliding_attn_mask,
+ }
+
+ hidden_states = inputs_embeds
+ position_embeddings = self.rotary_emb(hidden_states, position_ids)
+
+ all_hidden_states = () if output_hidden_states else None
+ all_self_attns = () if output_attentions else None
+
+ for layer_module in self.layers[: self.num_lyric_encoder_hidden_layers]:
+ if output_hidden_states:
+ all_hidden_states += (hidden_states,)
+
+ layer_outputs = layer_module(
+ hidden_states, position_embeddings,
+ self_attn_mask_mapping[layer_module.attention_type],
+ position_ids, output_attentions,
+ **flash_attn_kwargs,
+ )
+ hidden_states = layer_outputs[0]
+ if output_attentions:
+ all_self_attns += (layer_outputs[1],)
+
+ hidden_states = self.norm(hidden_states)
+ if output_hidden_states:
+ all_hidden_states += (hidden_states,)
+
+ return BaseModelOutput(
+ last_hidden_state=hidden_states,
+ hidden_states=all_hidden_states,
+ attentions=all_self_attns,
+ )
+
+
+class AceStepTimbreEncoder(nn.Module):
+ def __init__(
+ self,
+ hidden_size: int = 2048,
+ intermediate_size: int = 6144,
+ num_hidden_layers: int = 24,
+ num_attention_heads: int = 16,
+ num_key_value_heads: int = 8,
+ rms_norm_eps: float = 1e-6,
+ attention_bias: bool = False,
+ attention_dropout: float = 0.0,
+ layer_types: Optional[list] = None,
+ head_dim: Optional[int] = None,
+ sliding_window: Optional[int] = 128,
+ use_sliding_window: bool = True,
+ use_cache: bool = True,
+ rope_theta: float = 1000000,
+ max_position_embeddings: int = 32768,
+ initializer_range: float = 0.02,
+ timbre_hidden_dim: int = 64,
+ num_timbre_encoder_hidden_layers: int = 4,
+ **kwargs,
+ ):
+ super().__init__()
+ self.hidden_size = hidden_size
+ self.intermediate_size = intermediate_size
+ self.num_hidden_layers = num_hidden_layers
+ self.num_attention_heads = num_attention_heads
+ self.num_key_value_heads = num_key_value_heads
+ self.rms_norm_eps = rms_norm_eps
+ self.attention_bias = attention_bias
+ self.attention_dropout = attention_dropout
+ self.layer_types = layer_types or (["sliding_attention", "full_attention"] * (num_hidden_layers // 2))
+ self.head_dim = head_dim or hidden_size // num_attention_heads
+ self.sliding_window = sliding_window
+ self.use_sliding_window = use_sliding_window
+ self.use_cache = use_cache
+ self.rope_theta = rope_theta
+ self.max_position_embeddings = max_position_embeddings
+ self.initializer_range = initializer_range
+ self.timbre_hidden_dim = timbre_hidden_dim
+ self.num_timbre_encoder_hidden_layers = num_timbre_encoder_hidden_layers
+ self._attn_implementation = kwargs.get("_attn_implementation", "sdpa")
+
+ self.embed_tokens = nn.Linear(timbre_hidden_dim, hidden_size)
+ self.norm = Qwen3RMSNorm(hidden_size, eps=rms_norm_eps)
+ rope_config = type('RopeConfig', (), {
+ 'hidden_size': hidden_size,
+ 'num_attention_heads': num_attention_heads,
+ 'num_key_value_heads': num_key_value_heads,
+ 'head_dim': head_dim,
+ 'max_position_embeddings': max_position_embeddings,
+ 'rope_theta': rope_theta,
+ 'rope_parameters': {'rope_type': 'default', 'rope_theta': rope_theta},
+ 'rms_norm_eps': rms_norm_eps,
+ 'attention_bias': attention_bias,
+ 'attention_dropout': attention_dropout,
+ 'hidden_act': 'silu',
+ 'intermediate_size': intermediate_size,
+ 'layer_types': self.layer_types,
+ 'sliding_window': sliding_window,
+ '_attn_implementation': self._attn_implementation,
+ })()
+ self.rotary_emb = Qwen3RotaryEmbedding(rope_config)
+ self.gradient_checkpointing = False
+ self.special_token = nn.Parameter(torch.randn(1, 1, hidden_size))
+ self.layers = nn.ModuleList([
+ AceStepEncoderLayer(
+ hidden_size=hidden_size,
+ intermediate_size=intermediate_size,
+ num_attention_heads=num_attention_heads,
+ num_key_value_heads=num_key_value_heads,
+ rms_norm_eps=rms_norm_eps,
+ attention_bias=attention_bias,
+ attention_dropout=attention_dropout,
+ layer_types=self.layer_types,
+ head_dim=head_dim,
+ sliding_window=sliding_window,
+ layer_idx=layer_idx,
+ )
+ for layer_idx in range(num_timbre_encoder_hidden_layers)
+ ])
+
+ def unpack_timbre_embeddings(self, timbre_embs_packed, refer_audio_order_mask):
+ N, d = timbre_embs_packed.shape
+ device = timbre_embs_packed.device
+ dtype = timbre_embs_packed.dtype
+ B = int(refer_audio_order_mask.max().item() + 1)
+ counts = torch.bincount(refer_audio_order_mask, minlength=B)
+ max_count = counts.max().item()
+ sorted_indices = torch.argsort(refer_audio_order_mask * N + torch.arange(N, device=device), stable=True)
+ sorted_batch_ids = refer_audio_order_mask[sorted_indices]
+ positions = torch.arange(N, device=device)
+ batch_starts = torch.cat([torch.tensor([0], device=device), torch.cumsum(counts, dim=0)[:-1]])
+ positions_in_sorted = positions - batch_starts[sorted_batch_ids]
+ inverse_indices = torch.empty_like(sorted_indices)
+ inverse_indices[sorted_indices] = torch.arange(N, device=device)
+ positions_in_batch = positions_in_sorted[inverse_indices]
+ indices_2d = refer_audio_order_mask * max_count + positions_in_batch
+ one_hot = F.one_hot(indices_2d, num_classes=B * max_count).to(dtype)
+ timbre_embs_flat = one_hot.t() @ timbre_embs_packed
+ timbre_embs_unpack = timbre_embs_flat.reshape(B, max_count, d)
+ mask_flat = (one_hot.sum(dim=0) > 0).long()
+ new_mask = mask_flat.reshape(B, max_count)
+ return timbre_embs_unpack, new_mask
+
+ @can_return_tuple
+ def forward(
+ self,
+ refer_audio_acoustic_hidden_states_packed: Optional[torch.FloatTensor] = None,
+ refer_audio_order_mask: Optional[torch.LongTensor] = None,
+ attention_mask: Optional[torch.Tensor] = None,
+ **flash_attn_kwargs: Unpack[FlashAttentionKwargs],
+ ) -> BaseModelOutput:
+ inputs_embeds = refer_audio_acoustic_hidden_states_packed
+ inputs_embeds = self.embed_tokens(inputs_embeds)
+ seq_len = inputs_embeds.shape[1]
+ cache_position = torch.arange(0, seq_len, device=inputs_embeds.device)
+ position_ids = cache_position.unsqueeze(0)
+
+ dtype = inputs_embeds.dtype
+ device = inputs_embeds.device
+
+ full_attn_mask = None
+ sliding_attn_mask = None
+ if self.use_sliding_window:
+ sliding_attn_mask = {"window_size": self.sliding_window}
+
+ self_attn_mask_mapping = {
+ "full_attention": full_attn_mask,
+ "sliding_attention": sliding_attn_mask,
+ }
+
+ hidden_states = inputs_embeds
+ position_embeddings = self.rotary_emb(hidden_states, position_ids)
+
+ for layer_module in self.layers[: self.num_timbre_encoder_hidden_layers]:
+ layer_outputs = layer_module(
+ hidden_states, position_embeddings,
+ self_attn_mask_mapping[layer_module.attention_type],
+ position_ids,
+ **flash_attn_kwargs,
+ )
+ hidden_states = layer_outputs[0]
+
+ hidden_states = self.norm(hidden_states)
+ hidden_states = hidden_states[:, 0, :]
+ # For packed input: reshape [1, T, D] -> [T, D] for unpacking
+ timbre_embs_unpack, timbre_embs_mask = self.unpack_timbre_embeddings(hidden_states, refer_audio_order_mask)
+ return timbre_embs_unpack, timbre_embs_mask
+
+
+class AceStepConditionEncoder(nn.Module):
+ def __init__(
+ self,
+ hidden_size: int = 2048,
+ intermediate_size: int = 6144,
+ num_hidden_layers: int = 24,
+ num_attention_heads: int = 16,
+ num_key_value_heads: int = 8,
+ rms_norm_eps: float = 1e-6,
+ attention_bias: bool = False,
+ attention_dropout: float = 0.0,
+ layer_types: Optional[list] = None,
+ head_dim: Optional[int] = None,
+ sliding_window: Optional[int] = 128,
+ use_sliding_window: bool = True,
+ use_cache: bool = True,
+ rope_theta: float = 1000000,
+ max_position_embeddings: int = 32768,
+ initializer_range: float = 0.02,
+ text_hidden_dim: int = 1024,
+ timbre_hidden_dim: int = 64,
+ num_lyric_encoder_hidden_layers: int = 8,
+ num_timbre_encoder_hidden_layers: int = 4,
+ **kwargs,
+ ):
+ super().__init__()
+ self.hidden_size = hidden_size
+ self.intermediate_size = intermediate_size
+ self.num_hidden_layers = num_hidden_layers
+ self.num_attention_heads = num_attention_heads
+ self.num_key_value_heads = num_key_value_heads
+ self.rms_norm_eps = rms_norm_eps
+ self.attention_bias = attention_bias
+ self.attention_dropout = attention_dropout
+ self.layer_types = layer_types or (["sliding_attention", "full_attention"] * (num_hidden_layers // 2))
+ self.head_dim = head_dim or hidden_size // num_attention_heads
+ self.sliding_window = sliding_window
+ self.use_sliding_window = use_sliding_window
+ self.use_cache = use_cache
+ self.rope_theta = rope_theta
+ self.max_position_embeddings = max_position_embeddings
+ self.initializer_range = initializer_range
+ self.text_hidden_dim = text_hidden_dim
+ self.timbre_hidden_dim = timbre_hidden_dim
+ self.num_lyric_encoder_hidden_layers = num_lyric_encoder_hidden_layers
+ self.num_timbre_encoder_hidden_layers = num_timbre_encoder_hidden_layers
+ self._attn_implementation = kwargs.get("_attn_implementation", "sdpa")
+
+ self.text_projector = nn.Linear(text_hidden_dim, hidden_size, bias=False)
+ self.null_condition_emb = nn.Parameter(torch.randn(1, 1, hidden_size))
+ self.lyric_encoder = AceStepLyricEncoder(
+ hidden_size=hidden_size,
+ intermediate_size=intermediate_size,
+ num_attention_heads=num_attention_heads,
+ num_key_value_heads=num_key_value_heads,
+ rms_norm_eps=rms_norm_eps,
+ attention_bias=attention_bias,
+ attention_dropout=attention_dropout,
+ layer_types=layer_types,
+ head_dim=head_dim,
+ sliding_window=sliding_window,
+ use_sliding_window=use_sliding_window,
+ rope_theta=rope_theta,
+ max_position_embeddings=max_position_embeddings,
+ initializer_range=initializer_range,
+ text_hidden_dim=text_hidden_dim,
+ num_lyric_encoder_hidden_layers=num_lyric_encoder_hidden_layers,
+ )
+ self.timbre_encoder = AceStepTimbreEncoder(
+ hidden_size=hidden_size,
+ intermediate_size=intermediate_size,
+ num_attention_heads=num_attention_heads,
+ num_key_value_heads=num_key_value_heads,
+ rms_norm_eps=rms_norm_eps,
+ attention_bias=attention_bias,
+ attention_dropout=attention_dropout,
+ layer_types=layer_types,
+ head_dim=head_dim,
+ sliding_window=sliding_window,
+ use_sliding_window=use_sliding_window,
+ rope_theta=rope_theta,
+ max_position_embeddings=max_position_embeddings,
+ initializer_range=initializer_range,
+ timbre_hidden_dim=timbre_hidden_dim,
+ num_timbre_encoder_hidden_layers=num_timbre_encoder_hidden_layers,
+ )
+
+ def forward(
+ self,
+ text_hidden_states: Optional[torch.FloatTensor] = None,
+ text_attention_mask: Optional[torch.Tensor] = None,
+ lyric_hidden_states: Optional[torch.LongTensor] = None,
+ lyric_attention_mask: Optional[torch.Tensor] = None,
+ reference_latents: Optional[torch.Tensor] = None,
+ refer_audio_order_mask: Optional[torch.LongTensor] = None,
+ ):
+ text_hidden_states = self.text_projector(text_hidden_states)
+ lyric_encoder_outputs = self.lyric_encoder(
+ inputs_embeds=lyric_hidden_states,
+ attention_mask=lyric_attention_mask,
+ )
+ lyric_hidden_states = lyric_encoder_outputs.last_hidden_state
+ timbre_embs_unpack, timbre_embs_mask = self.timbre_encoder(reference_latents, refer_audio_order_mask)
+ encoder_hidden_states, encoder_attention_mask = pack_sequences(
+ lyric_hidden_states, timbre_embs_unpack, lyric_attention_mask, timbre_embs_mask
+ )
+ encoder_hidden_states, encoder_attention_mask = pack_sequences(
+ encoder_hidden_states, text_hidden_states, encoder_attention_mask, text_attention_mask
+ )
+ return encoder_hidden_states, encoder_attention_mask
diff --git a/diffsynth/models/ace_step_dit.py b/diffsynth/models/ace_step_dit.py
new file mode 100644
index 0000000000000000000000000000000000000000..0642e4fc9d356a19a0e1eff80167b4f4b600c4be
--- /dev/null
+++ b/diffsynth/models/ace_step_dit.py
@@ -0,0 +1,480 @@
+# Copyright 2025 The ACESTEO Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+import math
+from typing import Optional
+
+import torch
+import torch.nn.functional as F
+from torch import nn
+
+from ..core.attention.attention import attention_forward
+from ..core import gradient_checkpoint_forward
+
+from transformers.cache_utils import Cache, DynamicCache, EncoderDecoderCache
+from transformers.modeling_flash_attention_utils import FlashAttentionKwargs
+from transformers.modeling_outputs import BaseModelOutput
+from transformers.processing_utils import Unpack
+from transformers.utils import logging
+
+from transformers.models.qwen3.modeling_qwen3 import (
+ Qwen3MLP,
+ Qwen3RMSNorm,
+ Qwen3RotaryEmbedding,
+ apply_rotary_pos_emb,
+)
+
+
+class TimestepEmbedding(nn.Module):
+ def __init__(self, in_channels, time_embed_dim, scale=1):
+ super().__init__()
+ self.linear_1 = nn.Linear(in_channels, time_embed_dim, bias=True)
+ self.act1 = nn.SiLU()
+ self.linear_2 = nn.Linear(time_embed_dim, time_embed_dim, bias=True)
+ self.in_channels = in_channels
+
+ self.act2 = nn.SiLU()
+ self.time_proj = nn.Linear(time_embed_dim, time_embed_dim * 6)
+ self.scale = scale
+
+ def timestep_embedding(self, t, dim, max_period=10000):
+ t = t * self.scale
+ half = dim // 2
+ freqs = torch.exp(
+ -math.log(max_period) * torch.arange(start=0, end=half, dtype=torch.float32) / half
+ ).to(device=t.device)
+ args = t[:, None].float() * freqs[None]
+ embedding = torch.cat([torch.cos(args), torch.sin(args)], dim=-1)
+ if dim % 2:
+ embedding = torch.cat([embedding, torch.zeros_like(embedding[:, :1])], dim=-1)
+ return embedding
+
+ def forward(self, t):
+ t_freq = self.timestep_embedding(t, self.in_channels)
+ temb = self.linear_1(t_freq.to(t.dtype))
+ temb = self.act1(temb)
+ temb = self.linear_2(temb)
+ timestep_proj = self.time_proj(self.act2(temb)).unflatten(1, (6, -1))
+ return temb, timestep_proj
+
+
+class AceStepAttention(nn.Module):
+ def __init__(
+ self,
+ hidden_size: int,
+ num_attention_heads: int,
+ num_key_value_heads: int,
+ rms_norm_eps: float,
+ attention_bias: bool,
+ attention_dropout: float,
+ layer_types: list,
+ head_dim: Optional[int] = None,
+ sliding_window: Optional[int] = None,
+ layer_idx: int = 0,
+ is_cross_attention: bool = False,
+ is_causal: bool = False,
+ ):
+ super().__init__()
+ self.layer_idx = layer_idx
+ self.head_dim = head_dim or hidden_size // num_attention_heads
+ self.num_key_value_groups = num_attention_heads // num_key_value_heads
+ self.scaling = self.head_dim ** -0.5
+ self.attention_dropout = attention_dropout
+ if is_cross_attention:
+ is_causal = False
+ self.is_causal = is_causal
+ self.is_cross_attention = is_cross_attention
+
+ self.q_proj = nn.Linear(hidden_size, num_attention_heads * self.head_dim, bias=attention_bias)
+ self.k_proj = nn.Linear(hidden_size, num_key_value_heads * self.head_dim, bias=attention_bias)
+ self.v_proj = nn.Linear(hidden_size, num_key_value_heads * self.head_dim, bias=attention_bias)
+ self.o_proj = nn.Linear(num_attention_heads * self.head_dim, hidden_size, bias=attention_bias)
+ self.q_norm = Qwen3RMSNorm(self.head_dim, eps=rms_norm_eps)
+ self.k_norm = Qwen3RMSNorm(self.head_dim, eps=rms_norm_eps)
+ self.attention_type = layer_types[layer_idx]
+ self.sliding_window = sliding_window if layer_types[layer_idx] == "sliding_attention" else None
+
+ def forward(
+ self,
+ hidden_states: torch.Tensor,
+ attention_mask: Optional[torch.Tensor],
+ encoder_hidden_states: Optional[torch.Tensor] = None,
+ position_embeddings: tuple[torch.Tensor, torch.Tensor] = None,
+ **kwargs: Unpack[FlashAttentionKwargs],
+ ) -> tuple[torch.Tensor, Optional[torch.Tensor], Optional[tuple[torch.Tensor]]]:
+ # Project and normalize query states
+ input_shape = hidden_states.shape[:-1]
+ hidden_shape = (*input_shape, -1, self.head_dim)
+ query_states = self.q_norm(self.q_proj(hidden_states).view(hidden_shape)).transpose(1, 2)
+
+ # Process KV
+ if self.is_cross_attention:
+ encoder_hidden_shape = (*encoder_hidden_states.shape[:-1], -1, self.head_dim)
+ key_states = self.k_norm(self.k_proj(encoder_hidden_states).view(encoder_hidden_shape)).transpose(1, 2)
+ value_states = self.v_proj(encoder_hidden_states).view(encoder_hidden_shape).transpose(1, 2)
+ else:
+ key_states = self.k_norm(self.k_proj(hidden_states).view(hidden_shape)).transpose(1, 2)
+ value_states = self.v_proj(hidden_states).view(hidden_shape).transpose(1, 2)
+ if position_embeddings is not None:
+ cos, sin = position_embeddings
+ query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin)
+
+ # Use DiffSynth unified attention
+ attn_output = attention_forward(
+ query_states, key_states, value_states,
+ q_pattern="b n s d", k_pattern="b n s d", v_pattern="b n s d", out_pattern="b s (n d)",
+ window_size=None if attention_mask is None else attention_mask["window_size"],
+ )
+ attn_output = self.o_proj(attn_output)
+ return attn_output
+
+
+class AceStepDiTLayer(nn.Module):
+ def __init__(
+ self,
+ hidden_size: int,
+ num_attention_heads: int,
+ num_key_value_heads: int,
+ intermediate_size: int,
+ rms_norm_eps: float,
+ attention_bias: bool,
+ attention_dropout: float,
+ layer_types: list,
+ head_dim: Optional[int] = None,
+ sliding_window: Optional[int] = None,
+ layer_idx: int = 0,
+ use_cross_attention: bool = True,
+ ):
+ super().__init__()
+
+ self.self_attn_norm = Qwen3RMSNorm(hidden_size, eps=rms_norm_eps)
+ self.self_attn = AceStepAttention(
+ hidden_size=hidden_size,
+ num_attention_heads=num_attention_heads,
+ num_key_value_heads=num_key_value_heads,
+ rms_norm_eps=rms_norm_eps,
+ attention_bias=attention_bias,
+ attention_dropout=attention_dropout,
+ layer_types=layer_types,
+ head_dim=head_dim,
+ sliding_window=sliding_window,
+ layer_idx=layer_idx,
+ )
+
+ self.use_cross_attention = use_cross_attention
+ if self.use_cross_attention:
+ self.cross_attn_norm = Qwen3RMSNorm(hidden_size, eps=rms_norm_eps)
+ self.cross_attn = AceStepAttention(
+ hidden_size=hidden_size,
+ num_attention_heads=num_attention_heads,
+ num_key_value_heads=num_key_value_heads,
+ rms_norm_eps=rms_norm_eps,
+ attention_bias=attention_bias,
+ attention_dropout=attention_dropout,
+ layer_types=layer_types,
+ head_dim=head_dim,
+ sliding_window=sliding_window,
+ layer_idx=layer_idx,
+ is_cross_attention=True,
+ )
+
+ self.mlp_norm = Qwen3RMSNorm(hidden_size, eps=rms_norm_eps)
+ self.mlp = Qwen3MLP(
+ config=type('Config', (), {
+ 'hidden_size': hidden_size,
+ 'intermediate_size': intermediate_size,
+ 'hidden_act': 'silu',
+ })()
+ )
+
+ self.scale_shift_table = nn.Parameter(torch.randn(1, 6, hidden_size) / hidden_size**0.5)
+ self.attention_type = layer_types[layer_idx]
+
+ def forward(
+ self,
+ hidden_states: torch.Tensor,
+ position_embeddings: tuple[torch.Tensor, torch.Tensor],
+ temb: torch.Tensor,
+ attention_mask: Optional[torch.Tensor] = None,
+ encoder_hidden_states: Optional[torch.Tensor] = None,
+ encoder_attention_mask: Optional[torch.Tensor] = None,
+ ) -> torch.Tensor:
+
+ # Extract scale-shift parameters for adaptive layer norm from timestep embeddings
+ # 6 values: (shift_msa, scale_msa, gate_msa, c_shift_msa, c_scale_msa, c_gate_msa)
+ shift_msa, scale_msa, gate_msa, c_shift_msa, c_scale_msa, c_gate_msa = (
+ self.scale_shift_table.to(temb.device) + temb
+ ).chunk(6, dim=1)
+
+ # Step 1: Self-attention with adaptive layer norm (AdaLN)
+ # Apply adaptive normalization: norm(x) * (1 + scale) + shift
+ norm_hidden_states = (self.self_attn_norm(hidden_states) * (1 + scale_msa) + shift_msa).type_as(hidden_states)
+ attn_output = self.self_attn(
+ hidden_states=norm_hidden_states,
+ position_embeddings=position_embeddings,
+ attention_mask=attention_mask,
+ )
+ # Apply gated residual connection: x = x + attn_output * gate
+ hidden_states = (hidden_states + attn_output * gate_msa).type_as(hidden_states)
+
+ # Step 2: Cross-attention (if enabled) for conditioning on encoder outputs
+ if self.use_cross_attention:
+ norm_hidden_states = self.cross_attn_norm(hidden_states).type_as(hidden_states)
+ attn_output = self.cross_attn(
+ hidden_states=norm_hidden_states,
+ encoder_hidden_states=encoder_hidden_states,
+ attention_mask=encoder_attention_mask,
+ )
+ # Standard residual connection for cross-attention
+ hidden_states = hidden_states + attn_output
+
+ # Step 3: Feed-forward (MLP) with adaptive layer norm
+ # Apply adaptive normalization for MLP: norm(x) * (1 + scale) + shift
+ norm_hidden_states = (self.mlp_norm(hidden_states) * (1 + c_scale_msa) + c_shift_msa).type_as(hidden_states)
+ ff_output = self.mlp(norm_hidden_states)
+ # Apply gated residual connection: x = x + mlp_output * gate
+ hidden_states = (hidden_states + ff_output * c_gate_msa).type_as(hidden_states)
+
+ outputs = (hidden_states,)
+
+ return outputs
+
+
+class Lambda(nn.Module):
+ def __init__(self, func):
+ super().__init__()
+ self.func = func
+
+ def forward(self, x):
+ return self.func(x)
+
+
+class AceStepDiTModel(nn.Module):
+ def __init__(
+ self,
+ hidden_size: int = 2048,
+ intermediate_size: int = 6144,
+ num_hidden_layers: int = 24,
+ num_attention_heads: int = 16,
+ num_key_value_heads: int = 8,
+ rms_norm_eps: float = 1e-6,
+ attention_bias: bool = False,
+ attention_dropout: float = 0.0,
+ layer_types: Optional[list] = None,
+ head_dim: Optional[int] = None,
+ sliding_window: Optional[int] = 128,
+ use_sliding_window: bool = True,
+ use_cache: bool = True,
+ rope_theta: float = 1000000,
+ max_position_embeddings: int = 32768,
+ initializer_range: float = 0.02,
+ patch_size: int = 2,
+ in_channels: int = 192,
+ audio_acoustic_hidden_dim: int = 64,
+ encoder_hidden_size: Optional[int] = None,
+ **kwargs,
+ ):
+ super().__init__()
+
+ self.layer_types = layer_types or (["sliding_attention", "full_attention"] * (num_hidden_layers // 2))
+ self.use_sliding_window = use_sliding_window
+ self.sliding_window = sliding_window
+ self.use_cache = use_cache
+ encoder_hidden_size = encoder_hidden_size or hidden_size
+
+ # Rotary position embeddings for transformer layers
+ rope_config = type('RopeConfig', (), {
+ 'hidden_size': hidden_size,
+ 'num_attention_heads': num_attention_heads,
+ 'num_key_value_heads': num_key_value_heads,
+ 'head_dim': head_dim,
+ 'max_position_embeddings': max_position_embeddings,
+ 'rope_theta': rope_theta,
+ 'rope_parameters': {'rope_type': 'default', 'rope_theta': rope_theta},
+ 'rms_norm_eps': rms_norm_eps,
+ 'attention_bias': attention_bias,
+ 'attention_dropout': attention_dropout,
+ 'hidden_act': 'silu',
+ 'intermediate_size': intermediate_size,
+ 'layer_types': self.layer_types,
+ 'sliding_window': sliding_window,
+ })()
+ self.rotary_emb = Qwen3RotaryEmbedding(rope_config)
+
+ # Stack of DiT transformer layers
+ self.layers = nn.ModuleList([
+ AceStepDiTLayer(
+ hidden_size=hidden_size,
+ num_attention_heads=num_attention_heads,
+ num_key_value_heads=num_key_value_heads,
+ intermediate_size=intermediate_size,
+ rms_norm_eps=rms_norm_eps,
+ attention_bias=attention_bias,
+ attention_dropout=attention_dropout,
+ layer_types=self.layer_types,
+ head_dim=head_dim,
+ sliding_window=sliding_window,
+ layer_idx=layer_idx,
+ )
+ for layer_idx in range(num_hidden_layers)
+ ])
+
+ self.patch_size = patch_size
+
+ # Input projection: patch embedding using 1D convolution
+ self.proj_in = nn.Sequential(
+ Lambda(lambda x: x.transpose(1, 2)),
+ nn.Conv1d(
+ in_channels=in_channels,
+ out_channels=hidden_size,
+ kernel_size=patch_size,
+ stride=patch_size,
+ padding=0,
+ ),
+ Lambda(lambda x: x.transpose(1, 2)),
+ )
+
+ # Timestep embeddings for diffusion conditioning
+ self.time_embed = TimestepEmbedding(in_channels=256, time_embed_dim=hidden_size)
+ self.time_embed_r = TimestepEmbedding(in_channels=256, time_embed_dim=hidden_size)
+
+ # Project encoder hidden states to model dimension
+ self.condition_embedder = nn.Linear(encoder_hidden_size, hidden_size, bias=True)
+
+ # Output normalization and projection
+ self.norm_out = Qwen3RMSNorm(hidden_size, eps=rms_norm_eps)
+ self.proj_out = nn.Sequential(
+ Lambda(lambda x: x.transpose(1, 2)),
+ nn.ConvTranspose1d(
+ in_channels=hidden_size,
+ out_channels=audio_acoustic_hidden_dim,
+ kernel_size=patch_size,
+ stride=patch_size,
+ padding=0,
+ ),
+ Lambda(lambda x: x.transpose(1, 2)),
+ )
+ self.scale_shift_table = nn.Parameter(torch.randn(1, 2, hidden_size) / hidden_size**0.5)
+
+ def forward(
+ self,
+ hidden_states: torch.Tensor,
+ timestep: torch.Tensor,
+ timestep_r: torch.Tensor,
+ attention_mask: torch.Tensor,
+ encoder_hidden_states: torch.Tensor,
+ encoder_attention_mask: torch.Tensor,
+ context_latents: torch.Tensor,
+ use_cache: Optional[bool] = False,
+ past_key_values: Optional[EncoderDecoderCache] = None,
+ cache_position: Optional[torch.LongTensor] = None,
+ position_ids: Optional[torch.LongTensor] = None,
+ output_attentions: Optional[bool] = False,
+ return_hidden_states: int = None,
+ custom_layers_config: Optional[dict] = None,
+ enable_early_exit: bool = False,
+ use_gradient_checkpointing: bool = False,
+ use_gradient_checkpointing_offload: bool = False,
+ residual = None,
+ output_residual = False,
+ **flash_attn_kwargs: Unpack[FlashAttentionKwargs],
+ ):
+ # Compute timestep embeddings for diffusion conditioning
+ # Two embeddings: one for timestep t, one for timestep difference (t - r)
+ temb_t, timestep_proj_t = self.time_embed(timestep)
+ temb_r, timestep_proj_r = self.time_embed_r(timestep - timestep_r)
+ # Combine embeddings
+ temb = temb_t + temb_r
+ timestep_proj = timestep_proj_t + timestep_proj_r
+
+ # Concatenate context latents (source latents + chunk masks) with hidden states
+ hidden_states = torch.cat([context_latents, hidden_states], dim=-1)
+ # Record original sequence length for later restoration after padding
+ original_seq_len = hidden_states.shape[1]
+ # Apply padding if sequence length is not divisible by patch_size
+ # This ensures proper patch extraction
+ if hidden_states.shape[1] % self.patch_size != 0:
+ pad_length = self.patch_size - (hidden_states.shape[1] % self.patch_size)
+ hidden_states = F.pad(hidden_states, (0, 0, 0, pad_length), mode='constant', value=0)
+
+ # Project input to patches and project encoder states
+ hidden_states = self.proj_in(hidden_states)
+ encoder_hidden_states = self.condition_embedder(encoder_hidden_states)
+
+ # Cache positions and Position IDs
+ cache_position = torch.arange(0, hidden_states.shape[1], device=hidden_states.device)
+ position_ids = cache_position.unsqueeze(0)
+
+ # Build mask mapping
+ self_attn_mask_mapping = {
+ "full_attention": None,
+ "sliding_attention": {"window_size": self.sliding_window},
+ "encoder_attention_mask": None,
+ }
+
+ # Create position embeddings to be shared across all decoder layers
+ position_embeddings = self.rotary_emb(hidden_states, position_ids)
+
+ # Process through transformer layers
+ generated_residual = []
+ for index_block, layer_module in enumerate(self.layers):
+ # Prepare layer arguments
+ layer_args = (
+ hidden_states,
+ position_embeddings,
+ timestep_proj,
+ self_attn_mask_mapping[layer_module.attention_type],
+ encoder_hidden_states,
+ self_attn_mask_mapping["encoder_attention_mask"],
+ )
+
+ # Use gradient checkpointing if enabled
+ layer_outputs = gradient_checkpoint_forward(
+ layer_module,
+ use_gradient_checkpointing,
+ use_gradient_checkpointing_offload,
+ *layer_args,
+ )
+ hidden_states = layer_outputs[0]
+
+ # Residual control
+ if residual is not None:
+ block_residual = residual[index_block]
+ if block_residual.shape[1] > hidden_states.shape[1]:
+ block_residual = block_residual[:, :hidden_states.shape[1]]
+ elif block_residual.shape[1] < hidden_states.shape[1]:
+ block_residual = torch.concat([block_residual, torch.zeros_like(hidden_states)[:, :hidden_states.shape[1] - block_residual.shape[1]]], dim=1)
+ hidden_states = hidden_states + block_residual
+ if output_residual:
+ generated_residual.append(hidden_states)
+
+ if return_hidden_states:
+ return hidden_states
+ if output_residual:
+ return generated_residual
+
+ # Extract scale-shift parameters for adaptive output normalization
+ shift, scale = (self.scale_shift_table.to(temb.device) + temb.unsqueeze(1)).chunk(2, dim=1)
+ shift = shift.to(hidden_states.device)
+ scale = scale.to(hidden_states.device)
+
+ # Apply adaptive layer norm: norm(x) * (1 + scale) + shift
+ hidden_states = (self.norm_out(hidden_states) * (1 + scale) + shift).type_as(hidden_states)
+ # Project output: de-patchify back to original sequence format
+ hidden_states = self.proj_out(hidden_states)
+
+ # Crop back to original sequence length to ensure exact length match (remove padding)
+ hidden_states = hidden_states[:, :original_seq_len, :]
+
+ outputs = (hidden_states, past_key_values)
+ return outputs
diff --git a/diffsynth/models/ace_step_residual_fsq.py b/diffsynth/models/ace_step_residual_fsq.py
new file mode 100644
index 0000000000000000000000000000000000000000..211371cc438f92a289d183eb11118418cf4be326
--- /dev/null
+++ b/diffsynth/models/ace_step_residual_fsq.py
@@ -0,0 +1,569 @@
+"""
+Code adapted from https://github.com/lucidrains/vector-quantize-pytorch/blob/master/vector_quantize_pytorch/residual_fsq.py
+"""
+from functools import wraps, partial
+from contextlib import nullcontext
+
+import random
+from math import ceil
+
+import torch
+from torch import nn, tensor, Tensor, int32, tanh, atanh, clamp
+from torch.nn import Module
+import torch.nn.functional as F
+from torch.amp import autocast
+import torch.distributed as dist
+from einops import rearrange, reduce, pack, unpack
+
+
+# helper functions
+
+def exists(v):
+ return v is not None
+
+def default(*args):
+ for arg in args:
+ if exists(arg):
+ return arg
+ return None
+
+def identity(t):
+ return t
+
+def maybe(fn):
+ @wraps(fn)
+ def inner(x, *args, **kwargs):
+ if not exists(x):
+ return x
+ return fn(x, *args, **kwargs)
+ return inner
+
+def pack_one(t, pattern):
+ return pack([t], pattern)
+
+def unpack_one(t, ps, pattern):
+ return unpack(t, ps, pattern)[0]
+
+# tensor helpers
+
+def round_ste(z):
+ """ round with straight through gradients. """
+ zhat = z.round()
+ return z + (zhat - z).detach()
+
+def floor_ste(z):
+ """ floor with straight through gradients. """
+ zhat = z.floor()
+ return z + (zhat - z).detach()
+
+# main class
+
+class FSQ(Module):
+ def __init__(
+ self,
+ levels: list[int] | tuple[int, ...],
+ dim: int | None = None,
+ num_codebooks = 1,
+ keep_num_codebooks_dim: bool | None = None,
+ scale: float | None = None,
+ allowed_dtypes: tuple[torch.dtype, ...] = (torch.float32, torch.float64),
+ channel_first = False,
+ projection_has_bias = True,
+ return_indices = True,
+ force_quantization_f32 = True,
+ preserve_symmetry = False,
+ noise_dropout = 0.,
+ bound_hard_clamp = False, # for residual fsq, if input is pre-softclamped to the right range
+ orthogonal_rotation = False # increase codebook utilization. ensure levels are symmetric! https://arxiv.org/abs/2307.13304v2
+ ):
+ super().__init__()
+
+ assert not (any([l == 2 for l in levels]) and not preserve_symmetry), 'turn on `preserve_symmetry` for using any levels == 2, or use a greater level'
+
+ if isinstance(levels, tuple):
+ levels = list(levels)
+
+ _levels = tensor(levels, dtype = int32)
+ self.register_buffer('_levels', _levels, persistent = False)
+
+ _basis = torch.cumprod(tensor([1] + levels[:-1]), dim = 0, dtype = int32)
+ self.register_buffer('_basis', _basis, persistent = False)
+
+ self.scale = scale
+
+ assert not (noise_dropout > 0 and not preserve_symmetry)
+ self.preserve_symmetry = preserve_symmetry
+ self.noise_dropout = noise_dropout
+
+ codebook_dim = len(levels)
+ self.codebook_dim = codebook_dim
+
+ effective_codebook_dim = codebook_dim * num_codebooks
+ self.num_codebooks = num_codebooks
+ self.effective_codebook_dim = effective_codebook_dim
+
+ keep_num_codebooks_dim = default(keep_num_codebooks_dim, num_codebooks > 1)
+ assert not (num_codebooks > 1 and not keep_num_codebooks_dim)
+ self.keep_num_codebooks_dim = keep_num_codebooks_dim
+
+ self.dim = default(dim, len(_levels) * num_codebooks)
+
+ self.channel_first = channel_first
+
+ has_projections = self.dim != effective_codebook_dim
+ self.project_in = nn.Linear(self.dim, effective_codebook_dim, bias = projection_has_bias) if has_projections else nn.Identity()
+ self.project_out = nn.Linear(effective_codebook_dim, self.dim, bias = projection_has_bias) if has_projections else nn.Identity()
+
+ self.has_projections = has_projections
+
+ self.return_indices = return_indices
+
+ if return_indices:
+ self.codebook_size = self._levels.prod().item()
+ implicit_codebook = self._indices_to_codes(torch.arange(self.codebook_size))
+ self.register_buffer('implicit_codebook', implicit_codebook, persistent = False)
+
+ self.allowed_dtypes = allowed_dtypes
+ self.force_quantization_f32 = force_quantization_f32
+
+ # allow for a hard clamp
+
+ self.bound_hard_clamp = bound_hard_clamp
+
+ self.orthogonal_rotation = orthogonal_rotation
+
+ if orthogonal_rotation:
+ is_symmetric = len(set(levels)) == 1
+ if not is_symmetric:
+ print('orthogonal_rotation is not recommended for FSQ with asymmetric levels (i.e. where the number of bins differ across dimensions)')
+
+ orthogonal_rot = torch.empty(codebook_dim, codebook_dim)
+ nn.init.orthogonal_(orthogonal_rot)
+ self.register_buffer('orthogonal_rot', orthogonal_rot)
+
+ def bound(self, z, eps = 1e-3, hard_clamp = False):
+ """ Bound `z`, an array of shape (..., d). """
+ maybe_tanh = tanh if not hard_clamp else partial(clamp, min = -1., max = 1.)
+ maybe_atanh = atanh if not hard_clamp else identity
+
+ half_l = (self._levels - 1) * (1 + eps) / 2
+ offset = torch.where(self._levels % 2 == 0, 0.5, 0.0)
+ shift = maybe_atanh(offset / half_l)
+ bounded_z = maybe_tanh(z + shift) * half_l - offset
+ half_width = self._levels // 2
+ return round_ste(bounded_z) / half_width
+
+ # symmetry-preserving and noise-approximated quantization, section 3.2 in https://arxiv.org/abs/2411.19842
+
+ def symmetry_preserving_bound(self, z, hard_clamp = False):
+ """ QL(x) = 2 / (L - 1) * [(L - 1) * (tanh(x) + 1) / 2 + 0.5] - 1 """
+ maybe_tanh = tanh if not hard_clamp else partial(clamp, min = -1., max = 1.)
+
+ levels_minus_1 = (self._levels - 1)
+ scale = 2. / levels_minus_1
+ bracket = (levels_minus_1 * (maybe_tanh(z) + 1) / 2.) + 0.5
+ bracket = floor_ste(bracket)
+ return scale * bracket - 1.
+
+ def quantize(self, z):
+ """ Quantizes z, returns quantized zhat, same shape as z. """
+
+ shape, device, preserve_symmetry = z.shape[0], z.device, self.preserve_symmetry
+ bound_fn = self.symmetry_preserving_bound if preserve_symmetry else self.bound
+
+ return bound_fn(z, hard_clamp = self.bound_hard_clamp)
+
+ def maybe_apply_noise(self, bounded_z):
+ noise_dropout = self.noise_dropout
+
+ if not self.training or noise_dropout == 0.:
+ return bounded_z
+
+ # determine where to add a random offset elementwise
+ # if using noise dropout
+
+ offset_mask = torch.full_like(bounded_z, noise_dropout).bernoulli_().bool()
+ offset = torch.rand_like(bounded_z) - 0.5
+
+ bounded_z = torch.where(offset_mask, bounded_z + offset, bounded_z)
+
+ return bounded_z.clamp(-1., 1.)
+
+ def _scale_and_shift(self, zhat_normalized):
+ if self.preserve_symmetry:
+ return (zhat_normalized + 1.) / (2. / (self._levels - 1))
+
+ half_width = self._levels // 2
+ return (zhat_normalized * half_width) + half_width
+
+ def _scale_and_shift_inverse(self, zhat):
+ if self.preserve_symmetry:
+ return zhat * (2. / (self._levels - 1)) - 1.
+
+ half_width = self._levels // 2
+ return (zhat - half_width) / half_width
+
+ def _indices_to_codes(self, indices):
+ level_indices = self.indices_to_level_indices(indices)
+ codes = self._scale_and_shift_inverse(level_indices)
+ return codes
+
+ def indices_to_level_indices(self, indices):
+ """ Converts indices to indices at each level, perhaps needed for a transformer with factorized embeddings """
+ indices = rearrange(indices, '... -> ... 1')
+ codes_non_centered = (indices // self._basis) % self._levels
+ return codes_non_centered
+
+ def codes_to_indices(self, zhat):
+ """ Converts a `code` to an index in the codebook. """
+ assert zhat.shape[-1] == self.codebook_dim
+ zhat = self._scale_and_shift(zhat)
+ return (zhat * self._basis).sum(dim = -1).round().to(int32)
+
+ def indices_to_codes(self, indices):
+ """ Inverse of `codes_to_indices`. """
+ assert exists(indices)
+
+ is_img_or_video = indices.ndim >= (3 + int(self.keep_num_codebooks_dim))
+
+ codes = self._indices_to_codes(indices)
+
+ if self.orthogonal_rotation:
+ codes = codes @ self.orthogonal_rot.t()
+
+ if self.keep_num_codebooks_dim:
+ codes = rearrange(codes, '... c d -> ... (c d)')
+
+ codes = self.project_out(codes)
+
+ if is_img_or_video or self.channel_first:
+ codes = rearrange(codes, 'b ... d -> b d ...')
+
+ return codes
+
+ def forward(self, z):
+ """
+ einstein notation
+ b - batch
+ n - sequence (or flattened spatial dimensions)
+ d - feature dimension
+ c - number of codebook dim
+ """
+
+ is_img_or_video = z.ndim >= 4
+ need_move_channel_last = is_img_or_video or self.channel_first
+
+ # standardize image or video into (batch, seq, dimension)
+
+ if need_move_channel_last:
+ z = rearrange(z, 'b d ... -> b ... d')
+ z, ps = pack_one(z, 'b * d')
+
+ assert z.shape[-1] == self.dim, f'expected dimension of {self.dim} but found dimension of {z.shape[-1]}'
+
+ z = self.project_in(z)
+
+ z = rearrange(z, 'b n (c d) -> b n c d', c = self.num_codebooks)
+
+ if self.orthogonal_rotation:
+ z = z @ self.orthogonal_rot
+
+ # whether to force quantization step to be full precision or not
+
+ force_f32 = self.force_quantization_f32
+ quantization_context = partial(autocast, 'cuda', enabled = False) if force_f32 else nullcontext
+
+ with quantization_context():
+ orig_dtype = z.dtype
+
+ if force_f32 and orig_dtype not in self.allowed_dtypes:
+ z = z.float()
+
+ codes = self.quantize(z)
+
+ # returning indices could be optional
+
+ indices = None
+
+ if self.return_indices:
+ indices = self.codes_to_indices(codes)
+
+ codes = self.maybe_apply_noise(codes)
+
+ if self.orthogonal_rotation:
+ codes = codes @ self.orthogonal_rot.t()
+
+ codes = rearrange(codes, 'b n c d -> b n (c d)')
+
+ codes = codes.to(orig_dtype)
+
+ # project out
+
+ out = self.project_out(codes)
+
+ # reconstitute image or video dimensions
+
+ if need_move_channel_last:
+ out = unpack_one(out, ps, 'b * d')
+ out = rearrange(out, 'b ... d -> b d ...')
+
+ indices = maybe(unpack_one)(indices, ps, 'b * c')
+
+ if not self.keep_num_codebooks_dim and self.return_indices:
+ indices = maybe(rearrange)(indices, '... 1 -> ...')
+
+ # return quantized output and indices
+
+ return out, indices
+
+# helper functions
+
+def first(l):
+ return l[0]
+
+def default_residual_fsq(val, d):
+ return val if exists(val) else d
+
+def round_up_multiple(num, mult):
+ return ceil(num / mult) * mult
+
+# distributed helpers
+
+def is_distributed():
+ return dist.is_initialized() and dist.get_world_size() > 1
+
+def get_maybe_sync_seed(device, max_size = 10_000):
+ rand_int = torch.randint(0, max_size, (), device = device)
+
+ if is_distributed():
+ dist.all_reduce(rand_int)
+
+ return rand_int.item()
+
+# main class
+
+class ResidualFSQ(Module):
+ """ Follows Algorithm 1. in https://arxiv.org/pdf/2107.03312.pdf """
+
+ def __init__(
+ self,
+ *,
+ levels: list[int],
+ num_quantizers,
+ dim = None,
+ is_channel_first = False,
+ quantize_dropout = False,
+ quantize_dropout_cutoff_index = 0,
+ quantize_dropout_multiple_of = 1,
+ soft_clamp_input_value: float | list[float] | Tensor | None = None,
+ bound_hard_clamp = True,
+ **kwargs
+ ):
+ super().__init__()
+ codebook_dim = len(levels)
+ dim = default_residual_fsq(dim, codebook_dim)
+
+ requires_projection = codebook_dim != dim
+ self.project_in = nn.Linear(dim, codebook_dim) if requires_projection else nn.Identity()
+ self.project_out = nn.Linear(codebook_dim, dim) if requires_projection else nn.Identity()
+ self.has_projections = requires_projection
+
+ self.is_channel_first = is_channel_first
+ self.num_quantizers = num_quantizers
+
+ # layers
+
+ self.levels = levels
+ self.layers = nn.ModuleList([])
+
+ levels_tensor = tensor(levels)
+ assert (levels_tensor > 1).all()
+
+ scales = []
+
+ for ind in range(num_quantizers):
+ scales.append(levels_tensor.float() ** -ind)
+
+ fsq = FSQ(
+ levels = levels,
+ dim = codebook_dim,
+ preserve_symmetry = True,
+ bound_hard_clamp = bound_hard_clamp,
+ **kwargs
+ )
+
+ self.layers.append(fsq)
+
+ assert all([not fsq.has_projections for fsq in self.layers])
+
+ self.codebook_size = self.layers[0].codebook_size
+
+ self.register_buffer('scales', torch.stack(scales), persistent = False)
+
+ self.quantize_dropout = quantize_dropout and num_quantizers > 1
+
+ assert quantize_dropout_cutoff_index >= 0
+
+ self.quantize_dropout_cutoff_index = quantize_dropout_cutoff_index
+ self.quantize_dropout_multiple_of = quantize_dropout_multiple_of # encodec paper proposes structured dropout, believe this was set to 4
+
+ # soft clamping the input value
+
+ if bound_hard_clamp:
+ assert not exists(soft_clamp_input_value)
+ soft_clamp_input_value = 1 + (1 / (levels_tensor - 1))
+
+ if isinstance(soft_clamp_input_value, (list, float)):
+ soft_clamp_input_value = tensor(soft_clamp_input_value)
+
+ self.register_buffer('soft_clamp_input_value', soft_clamp_input_value, persistent = False)
+
+ @property
+ def codebooks(self):
+ codebooks = [layer.implicit_codebook for layer in self.layers]
+ codebooks = torch.stack(codebooks, dim = 0)
+ return codebooks
+
+ def get_codes_from_indices(self, indices):
+
+ batch, quantize_dim = indices.shape[0], indices.shape[-1]
+
+ # may also receive indices in the shape of 'b h w q' (accept_image_fmap)
+
+ indices, ps = pack([indices], 'b * q')
+
+ # because of quantize dropout, one can pass in indices that are coarse
+ # and the network should be able to reconstruct
+
+ if quantize_dim < self.num_quantizers:
+ assert self.quantize_dropout > 0., 'quantize dropout must be greater than 0 if you wish to reconstruct from a signal with less fine quantizations'
+ indices = F.pad(indices, (0, self.num_quantizers - quantize_dim), value = -1)
+
+ # take care of quantizer dropout
+
+ mask = indices == -1
+ indices = indices.masked_fill(mask, 0) # have it fetch a dummy code to be masked out later
+
+ all_codes = torch.stack([F.embedding(indices[:, :, i], self.codebooks[i]) for i in range(self.num_quantizers)], dim=0) # (q, b, n, d)
+
+ # mask out any codes that were dropout-ed
+
+ all_codes = all_codes.masked_fill(rearrange(mask, 'b n q -> q b n 1'), 0.)
+
+ # scale the codes
+
+ scales = rearrange(self.scales, 'q d -> q 1 1 d')
+ all_codes = all_codes * scales
+
+ # if (accept_image_fmap = True) then return shape (quantize, batch, height, width, dimension)
+
+ all_codes, = unpack(all_codes, ps, 'q b * d')
+
+ return all_codes
+
+ def get_output_from_indices(self, indices):
+ codes = self.get_codes_from_indices(indices)
+ codes_summed = reduce(codes, 'q ... -> ...', 'sum')
+ return self.project_out(codes_summed)
+
+ def forward(
+ self,
+ x,
+ return_all_codes = False,
+ rand_quantize_dropout_fixed_seed = None
+ ):
+ num_quant, quant_dropout_multiple_of, device = self.num_quantizers, self.quantize_dropout_multiple_of, x.device
+
+ # handle channel first
+
+ if self.is_channel_first:
+ x = rearrange(x, 'b d ... -> b ... d')
+ x, ps = pack([x], 'b * d')
+
+ # maybe project in
+
+ x = self.project_in(x)
+
+ # maybe softclamp input before residual layers
+
+ if exists(self.soft_clamp_input_value):
+ clamp_value = self.soft_clamp_input_value
+ x = (x / clamp_value).tanh() * clamp_value
+
+ # ready some variables to be accumulated
+
+ quantized_out = 0.
+ residual = x
+
+ all_indices = []
+
+ should_quantize_dropout = self.training and self.quantize_dropout and torch.is_grad_enabled()
+
+ # sample a layer index at which to dropout further residual quantization
+ # also prepare null indices
+
+ if should_quantize_dropout:
+
+ # check if seed is manually passed in
+
+ if not exists(rand_quantize_dropout_fixed_seed):
+ rand_quantize_dropout_fixed_seed = get_maybe_sync_seed(device)
+
+ rand = random.Random(rand_quantize_dropout_fixed_seed)
+
+ rand_quantize_dropout_index = rand.randrange(self.quantize_dropout_cutoff_index, num_quant)
+
+ if quant_dropout_multiple_of != 1:
+ rand_quantize_dropout_index = round_up_multiple(rand_quantize_dropout_index + 1, quant_dropout_multiple_of) - 1
+
+ null_indices = torch.full(x.shape[:2], -1., device = device, dtype = torch.long)
+
+ # go through the layers
+
+ with autocast('cuda', enabled = False):
+ for quantizer_index, (layer, scale) in enumerate(zip(self.layers, self.scales)):
+
+ if should_quantize_dropout and quantizer_index > rand_quantize_dropout_index:
+ all_indices.append(null_indices)
+ continue
+
+ quantized, indices = layer(residual / scale)
+
+ quantized = quantized * scale
+
+ residual = residual - quantized.detach()
+ quantized_out = quantized_out + quantized
+
+ all_indices.append(indices)
+
+ # project out, if needed
+
+ quantized_out = self.project_out(quantized_out)
+
+ # stack all indices
+
+ all_indices = torch.stack(all_indices, dim = -1)
+
+ # channel first out
+
+ if self.is_channel_first:
+ quantized_out, = unpack(quantized_out, ps, 'b * d')
+ all_indices, = unpack(all_indices, ps, 'b * d')
+
+ quantized_out = rearrange(quantized_out, 'b ... d -> b d ...')
+ all_indices = rearrange(all_indices, 'b ... d -> b d ...')
+
+ # return
+
+ ret = (quantized_out, all_indices)
+
+ if not return_all_codes:
+ return ret
+
+ # whether to return all codes from all codebooks across layers
+
+ all_codes = self.get_codes_from_indices(all_indices)
+
+ # will return all codes in shape (quantizer, batch, sequence length, codebook dimension)
+
+ return (*ret, all_codes)
diff --git a/diffsynth/models/ace_step_text_encoder.py b/diffsynth/models/ace_step_text_encoder.py
new file mode 100644
index 0000000000000000000000000000000000000000..1d0ac57f6bf76bcd3920f5b9ceffb2f2f694fc0e
--- /dev/null
+++ b/diffsynth/models/ace_step_text_encoder.py
@@ -0,0 +1,53 @@
+import torch
+
+
+class AceStepTextEncoder(torch.nn.Module):
+ def __init__(
+ self,
+ ):
+ super().__init__()
+ from transformers import Qwen3Config, Qwen3Model
+
+ config = Qwen3Config(
+ attention_bias=False,
+ attention_dropout=0.0,
+ bos_token_id=151643,
+ dtype="bfloat16",
+ eos_token_id=151643,
+ head_dim=128,
+ hidden_act="silu",
+ hidden_size=1024,
+ initializer_range=0.02,
+ intermediate_size=3072,
+ layer_types=["full_attention"] * 28,
+ max_position_embeddings=32768,
+ max_window_layers=28,
+ model_type="qwen3",
+ num_attention_heads=16,
+ num_hidden_layers=28,
+ num_key_value_heads=8,
+ pad_token_id=151643,
+ rms_norm_eps=1e-06,
+ rope_scaling=None,
+ rope_theta=1000000,
+ sliding_window=None,
+ tie_word_embeddings=True,
+ use_cache=True,
+ use_sliding_window=False,
+ vocab_size=151669,
+ )
+
+ self.model = Qwen3Model(config)
+
+ @torch.no_grad()
+ def forward(
+ self,
+ input_ids: torch.LongTensor,
+ attention_mask: torch.Tensor,
+ ):
+ outputs = self.model(
+ input_ids=input_ids,
+ attention_mask=attention_mask,
+ return_dict=True,
+ )
+ return outputs.last_hidden_state
diff --git a/diffsynth/models/ace_step_tokenizer.py b/diffsynth/models/ace_step_tokenizer.py
new file mode 100644
index 0000000000000000000000000000000000000000..9a9eb0aaa4bf7a3a0aae6adef65cff84750bae8f
--- /dev/null
+++ b/diffsynth/models/ace_step_tokenizer.py
@@ -0,0 +1,720 @@
+# Copyright 2025 The ACESTEO Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+from typing import Optional
+
+import torch
+import torch.nn as nn
+from einops import rearrange
+
+from ..core.attention import attention_forward
+
+from transformers.cache_utils import Cache
+from transformers.modeling_flash_attention_utils import FlashAttentionKwargs
+from transformers.processing_utils import Unpack
+from transformers.utils import can_return_tuple, logging
+from transformers.models.qwen3.modeling_qwen3 import (
+ Qwen3MLP,
+ Qwen3RMSNorm,
+ Qwen3RotaryEmbedding,
+ apply_rotary_pos_emb,
+)
+from .ace_step_residual_fsq import ResidualFSQ
+
+logger = logging.get_logger(__name__)
+
+
+def create_4d_mask(
+ seq_len: int,
+ dtype: torch.dtype,
+ device: torch.device,
+ attention_mask: Optional[torch.Tensor] = None,
+ sliding_window: Optional[int] = None,
+ is_sliding_window: bool = False,
+ is_causal: bool = True,
+) -> torch.Tensor:
+ indices = torch.arange(seq_len, device=device)
+ diff = indices.unsqueeze(1) - indices.unsqueeze(0)
+ valid_mask = torch.ones((seq_len, seq_len), device=device, dtype=torch.bool)
+ if is_causal:
+ valid_mask = valid_mask & (diff >= 0)
+ if is_sliding_window and sliding_window is not None:
+ if is_causal:
+ valid_mask = valid_mask & (diff <= sliding_window)
+ else:
+ valid_mask = valid_mask & (torch.abs(diff) <= sliding_window)
+ valid_mask = valid_mask.unsqueeze(0).unsqueeze(0)
+ if attention_mask is not None:
+ padding_mask_4d = attention_mask.view(attention_mask.shape[0], 1, 1, seq_len).to(torch.bool)
+ valid_mask = valid_mask & padding_mask_4d
+ min_dtype = torch.finfo(dtype).min
+ mask_tensor = torch.full(valid_mask.shape, min_dtype, dtype=dtype, device=device)
+ mask_tensor.masked_fill_(valid_mask, 0.0)
+ return mask_tensor
+
+
+class Lambda(nn.Module):
+ def __init__(self, func):
+ super().__init__()
+ self.func = func
+
+ def forward(self, x):
+ return self.func(x)
+
+
+class AceStepAttention(nn.Module):
+ def __init__(
+ self,
+ hidden_size: int,
+ num_attention_heads: int,
+ num_key_value_heads: int,
+ rms_norm_eps: float,
+ attention_bias: bool,
+ attention_dropout: float,
+ layer_types: list,
+ head_dim: Optional[int] = None,
+ sliding_window: Optional[int] = None,
+ layer_idx: int = 0,
+ is_cross_attention: bool = False,
+ is_causal: bool = False,
+ ):
+ super().__init__()
+ self.layer_idx = layer_idx
+ self.head_dim = head_dim or hidden_size // num_attention_heads
+ self.num_key_value_groups = num_attention_heads // num_key_value_heads
+ self.scaling = self.head_dim ** -0.5
+ self.attention_dropout = attention_dropout
+ if is_cross_attention:
+ is_causal = False
+ self.is_causal = is_causal
+ self.is_cross_attention = is_cross_attention
+
+ self.q_proj = nn.Linear(hidden_size, num_attention_heads * self.head_dim, bias=attention_bias)
+ self.k_proj = nn.Linear(hidden_size, num_key_value_heads * self.head_dim, bias=attention_bias)
+ self.v_proj = nn.Linear(hidden_size, num_key_value_heads * self.head_dim, bias=attention_bias)
+ self.o_proj = nn.Linear(num_attention_heads * self.head_dim, hidden_size, bias=attention_bias)
+ self.q_norm = Qwen3RMSNorm(self.head_dim, eps=rms_norm_eps)
+ self.k_norm = Qwen3RMSNorm(self.head_dim, eps=rms_norm_eps)
+ self.attention_type = layer_types[layer_idx]
+ self.sliding_window = sliding_window if layer_types[layer_idx] == "sliding_attention" else None
+
+ def forward(
+ self,
+ hidden_states: torch.Tensor,
+ attention_mask: Optional[torch.Tensor],
+ past_key_value: Optional[Cache] = None,
+ cache_position: Optional[torch.LongTensor] = None,
+ encoder_hidden_states: Optional[torch.Tensor] = None,
+ position_embeddings: tuple[torch.Tensor, torch.Tensor] = None,
+ output_attentions: Optional[bool] = False,
+ **kwargs: Unpack[FlashAttentionKwargs],
+ ) -> tuple[torch.Tensor, Optional[torch.Tensor]]:
+ input_shape = hidden_states.shape[:-1]
+ hidden_shape = (*input_shape, -1, self.head_dim)
+
+ query_states = self.q_norm(self.q_proj(hidden_states).view(hidden_shape)).transpose(1, 2)
+
+ is_cross_attention = self.is_cross_attention and encoder_hidden_states is not None
+
+ if is_cross_attention:
+ encoder_hidden_shape = (*encoder_hidden_states.shape[:-1], -1, self.head_dim)
+ if past_key_value is not None:
+ is_updated = past_key_value.is_updated.get(self.layer_idx)
+ curr_past_key_value = past_key_value.cross_attention_cache
+ if not is_updated:
+ key_states = self.k_norm(self.k_proj(encoder_hidden_states).view(encoder_hidden_shape)).transpose(1, 2)
+ value_states = self.v_proj(encoder_hidden_states).view(encoder_hidden_shape).transpose(1, 2)
+ key_states, value_states = curr_past_key_value.update(key_states, value_states, self.layer_idx)
+ past_key_value.is_updated[self.layer_idx] = True
+ else:
+ key_states = curr_past_key_value.layers[self.layer_idx].keys
+ value_states = curr_past_key_value.layers[self.layer_idx].values
+ else:
+ key_states = self.k_norm(self.k_proj(encoder_hidden_states).view(encoder_hidden_shape)).transpose(1, 2)
+ value_states = self.v_proj(encoder_hidden_states).view(encoder_hidden_shape).transpose(1, 2)
+
+ else:
+ key_states = self.k_norm(self.k_proj(hidden_states).view(hidden_shape)).transpose(1, 2)
+ value_states = self.v_proj(hidden_states).view(hidden_shape).transpose(1, 2)
+ if position_embeddings is not None:
+ cos, sin = position_embeddings
+ query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin)
+
+ if past_key_value is not None:
+ cache_kwargs = {"sin": sin, "cos": cos, "cache_position": cache_position}
+ key_states, value_states = past_key_value.update(key_states, value_states, self.layer_idx, cache_kwargs)
+
+ if self.num_key_value_groups > 1:
+ key_states = key_states.unsqueeze(2).expand(-1, -1, self.num_key_value_groups, -1, -1).flatten(1, 2)
+ value_states = value_states.unsqueeze(2).expand(-1, -1, self.num_key_value_groups, -1, -1).flatten(1, 2)
+
+ attn_output = attention_forward(
+ query_states, key_states, value_states,
+ q_pattern="b n s d", k_pattern="b n s d", v_pattern="b n s d", out_pattern="b n s d",
+ attn_mask=attention_mask,
+ )
+ attn_weights = None
+
+ attn_output = attn_output.transpose(1, 2).flatten(2, 3).contiguous()
+ attn_output = self.o_proj(attn_output)
+ return attn_output, attn_weights
+
+
+class AceStepEncoderLayer(nn.Module):
+ def __init__(
+ self,
+ hidden_size: int,
+ intermediate_size: int,
+ num_attention_heads: int,
+ num_key_value_heads: int,
+ rms_norm_eps: float,
+ attention_bias: bool,
+ attention_dropout: float,
+ layer_types: list,
+ head_dim: Optional[int] = None,
+ sliding_window: Optional[int] = None,
+ layer_idx: int = 0,
+ ):
+ super().__init__()
+ self.hidden_size = hidden_size
+ self.layer_idx = layer_idx
+
+ self.self_attn = AceStepAttention(
+ hidden_size=hidden_size,
+ num_attention_heads=num_attention_heads,
+ num_key_value_heads=num_key_value_heads,
+ rms_norm_eps=rms_norm_eps,
+ attention_bias=attention_bias,
+ attention_dropout=attention_dropout,
+ layer_types=layer_types,
+ head_dim=head_dim,
+ sliding_window=sliding_window,
+ layer_idx=layer_idx,
+ is_cross_attention=False,
+ is_causal=False,
+ )
+ self.input_layernorm = Qwen3RMSNorm(hidden_size, eps=rms_norm_eps)
+ self.post_attention_layernorm = Qwen3RMSNorm(hidden_size, eps=rms_norm_eps)
+
+ mlp_config = type('Config', (), {
+ 'hidden_size': hidden_size,
+ 'intermediate_size': intermediate_size,
+ 'hidden_act': 'silu',
+ })()
+ self.mlp = Qwen3MLP(mlp_config)
+ self.attention_type = layer_types[layer_idx]
+
+ def forward(
+ self,
+ hidden_states: torch.Tensor,
+ position_embeddings: tuple[torch.Tensor, torch.Tensor],
+ attention_mask: Optional[torch.Tensor] = None,
+ position_ids: Optional[torch.LongTensor] = None,
+ output_attentions: Optional[bool] = False,
+ **kwargs,
+ ) -> tuple[torch.FloatTensor, Optional[tuple[torch.FloatTensor, torch.FloatTensor]]]:
+ residual = hidden_states
+ hidden_states = self.input_layernorm(hidden_states)
+ hidden_states, self_attn_weights = self.self_attn(
+ hidden_states=hidden_states,
+ position_embeddings=position_embeddings,
+ attention_mask=attention_mask,
+ position_ids=position_ids,
+ output_attentions=output_attentions,
+ use_cache=False,
+ past_key_value=None,
+ **kwargs,
+ )
+ hidden_states = residual + hidden_states
+
+ residual = hidden_states
+ hidden_states = self.post_attention_layernorm(hidden_states)
+ hidden_states = self.mlp(hidden_states)
+ hidden_states = residual + hidden_states
+
+ outputs = (hidden_states,)
+ if output_attentions:
+ outputs += (self_attn_weights,)
+ return outputs
+
+
+class AttentionPooler(nn.Module):
+ """Pools every pool_window_size frames into 1 representation via transformer + CLS token."""
+
+ def __init__(
+ self,
+ hidden_size: int = 2048,
+ intermediate_size: int = 6144,
+ num_attention_heads: int = 16,
+ num_key_value_heads: int = 8,
+ rms_norm_eps: float = 1e-6,
+ attention_bias: bool = False,
+ attention_dropout: float = 0.0,
+ layer_types: Optional[list] = None,
+ head_dim: Optional[int] = None,
+ sliding_window: Optional[int] = 128,
+ use_sliding_window: bool = True,
+ rope_theta: float = 1000000,
+ max_position_embeddings: int = 32768,
+ initializer_range: float = 0.02,
+ num_attention_pooler_hidden_layers: int = 2,
+ **kwargs,
+ ):
+ super().__init__()
+ self.hidden_size = hidden_size
+ self.intermediate_size = intermediate_size
+ self.num_attention_heads = num_attention_heads
+ self.num_key_value_heads = num_key_value_heads
+ self.rms_norm_eps = rms_norm_eps
+ self.attention_bias = attention_bias
+ self.attention_dropout = attention_dropout
+ # Default matches target library config (24 alternating entries).
+ self.layer_types = layer_types or (["sliding_attention", "full_attention"] * 12)
+ self.head_dim = head_dim or hidden_size // num_attention_heads
+ self.sliding_window = sliding_window
+ self.use_sliding_window = use_sliding_window
+ self.rope_theta = rope_theta
+ self.max_position_embeddings = max_position_embeddings
+ self.initializer_range = initializer_range
+ self.num_attention_pooler_hidden_layers = num_attention_pooler_hidden_layers
+ self._attn_implementation = kwargs.get("_attn_implementation", "sdpa")
+
+ self.embed_tokens = nn.Linear(hidden_size, hidden_size)
+ self.norm = Qwen3RMSNorm(hidden_size, eps=rms_norm_eps)
+ # Slice layer_types to our own layer count
+ pooler_layer_types = self.layer_types[:num_attention_pooler_hidden_layers]
+ rope_config = type('RopeConfig', (), {
+ 'hidden_size': hidden_size,
+ 'num_attention_heads': num_attention_heads,
+ 'num_key_value_heads': num_key_value_heads,
+ 'head_dim': head_dim,
+ 'max_position_embeddings': max_position_embeddings,
+ 'rope_theta': rope_theta,
+ 'rope_parameters': {'rope_type': 'default', 'rope_theta': rope_theta},
+ 'rms_norm_eps': rms_norm_eps,
+ 'attention_bias': attention_bias,
+ 'attention_dropout': attention_dropout,
+ 'hidden_act': 'silu',
+ 'intermediate_size': intermediate_size,
+ 'layer_types': pooler_layer_types,
+ 'sliding_window': sliding_window,
+ '_attn_implementation': self._attn_implementation,
+ })()
+ self.rotary_emb = Qwen3RotaryEmbedding(rope_config)
+ self.gradient_checkpointing = False
+ self.special_token = nn.Parameter(torch.randn(1, 1, hidden_size) * 0.02)
+ self.layers = nn.ModuleList([
+ AceStepEncoderLayer(
+ hidden_size=hidden_size,
+ intermediate_size=intermediate_size,
+ num_attention_heads=num_attention_heads,
+ num_key_value_heads=num_key_value_heads,
+ rms_norm_eps=rms_norm_eps,
+ attention_bias=attention_bias,
+ attention_dropout=attention_dropout,
+ layer_types=pooler_layer_types,
+ head_dim=head_dim,
+ sliding_window=sliding_window,
+ layer_idx=layer_idx,
+ )
+ for layer_idx in range(num_attention_pooler_hidden_layers)
+ ])
+
+ @can_return_tuple
+ def forward(
+ self,
+ x,
+ attention_mask: Optional[torch.Tensor] = None,
+ **flash_attn_kwargs: Unpack[FlashAttentionKwargs],
+ ) -> torch.Tensor:
+ B, T, P, D = x.shape
+ x = self.embed_tokens(x)
+ special_tokens = self.special_token.expand(B, T, 1, -1).to(x.device)
+ x = torch.cat([special_tokens, x], dim=2)
+ x = rearrange(x, "b t p c -> (b t) p c")
+
+ cache_position = torch.arange(0, x.shape[1], device=x.device)
+ position_ids = cache_position.unsqueeze(0)
+ hidden_states = x
+ position_embeddings = self.rotary_emb(hidden_states, position_ids)
+
+ seq_len = x.shape[1]
+ dtype = x.dtype
+ device = x.device
+
+ full_attn_mask = create_4d_mask(
+ seq_len=seq_len, dtype=dtype, device=device,
+ attention_mask=attention_mask, sliding_window=None,
+ is_sliding_window=False, is_causal=False
+ )
+ sliding_attn_mask = None
+ if self.use_sliding_window:
+ sliding_attn_mask = create_4d_mask(
+ seq_len=seq_len, dtype=dtype, device=device,
+ attention_mask=attention_mask, sliding_window=self.sliding_window,
+ is_sliding_window=True, is_causal=False
+ )
+
+ self_attn_mask_mapping = {
+ "full_attention": full_attn_mask,
+ "sliding_attention": sliding_attn_mask,
+ }
+
+ for layer_module in self.layers:
+ layer_outputs = layer_module(
+ hidden_states, position_embeddings,
+ attention_mask=self_attn_mask_mapping[layer_module.attention_type],
+ **flash_attn_kwargs,
+ )
+ hidden_states = layer_outputs[0]
+
+ hidden_states = self.norm(hidden_states)
+ cls_output = hidden_states[:, 0, :]
+ return rearrange(cls_output, "(b t) c -> b t c", b=B)
+
+
+class AceStepAudioTokenizer(nn.Module):
+ """Converts continuous acoustic features (VAE latents) into discrete quantized tokens.
+
+ Input: [B, T, 64] (VAE latent dim)
+ Output: quantized [B, T/5, 2048], indices [B, T/5, 1]
+ """
+
+ def __init__(
+ self,
+ hidden_size: int = 2048,
+ intermediate_size: int = 6144,
+ num_attention_heads: int = 16,
+ num_key_value_heads: int = 8,
+ rms_norm_eps: float = 1e-6,
+ attention_bias: bool = False,
+ attention_dropout: float = 0.0,
+ layer_types: Optional[list] = None,
+ head_dim: Optional[int] = None,
+ sliding_window: Optional[int] = 128,
+ use_sliding_window: bool = True,
+ rope_theta: float = 1000000,
+ max_position_embeddings: int = 32768,
+ initializer_range: float = 0.02,
+ audio_acoustic_hidden_dim: int = 64,
+ pool_window_size: int = 5,
+ fsq_dim: int = 2048,
+ fsq_input_levels: list = None,
+ fsq_input_num_quantizers: int = 1,
+ num_attention_pooler_hidden_layers: int = 2,
+ **kwargs,
+ ):
+ super().__init__()
+ self.hidden_size = hidden_size
+ self.intermediate_size = intermediate_size
+ self.num_attention_heads = num_attention_heads
+ self.num_key_value_heads = num_key_value_heads
+ self.rms_norm_eps = rms_norm_eps
+ self.attention_bias = attention_bias
+ self.attention_dropout = attention_dropout
+ # Default matches target library config (24 alternating entries).
+ self.layer_types = layer_types or (["sliding_attention", "full_attention"] * 12)
+ self.head_dim = head_dim or hidden_size // num_attention_heads
+ self.sliding_window = sliding_window
+ self.use_sliding_window = use_sliding_window
+ self.rope_theta = rope_theta
+ self.max_position_embeddings = max_position_embeddings
+ self.initializer_range = initializer_range
+ self.audio_acoustic_hidden_dim = audio_acoustic_hidden_dim
+ self.pool_window_size = pool_window_size
+ self.fsq_dim = fsq_dim
+ self.fsq_input_levels = fsq_input_levels or [8, 8, 8, 5, 5, 5]
+ self.fsq_input_num_quantizers = fsq_input_num_quantizers
+ self.num_attention_pooler_hidden_layers = num_attention_pooler_hidden_layers
+ self._attn_implementation = kwargs.get("_attn_implementation", "sdpa")
+
+ self.audio_acoustic_proj = nn.Linear(audio_acoustic_hidden_dim, hidden_size)
+ # Slice layer_types for the attention pooler
+ pooler_layer_types = self.layer_types[:num_attention_pooler_hidden_layers]
+ self.attention_pooler = AttentionPooler(
+ hidden_size=hidden_size,
+ intermediate_size=intermediate_size,
+ num_attention_heads=num_attention_heads,
+ num_key_value_heads=num_key_value_heads,
+ rms_norm_eps=rms_norm_eps,
+ attention_bias=attention_bias,
+ attention_dropout=attention_dropout,
+ layer_types=pooler_layer_types,
+ head_dim=head_dim,
+ sliding_window=sliding_window,
+ use_sliding_window=use_sliding_window,
+ rope_theta=rope_theta,
+ max_position_embeddings=max_position_embeddings,
+ initializer_range=initializer_range,
+ num_attention_pooler_hidden_layers=num_attention_pooler_hidden_layers,
+ )
+ self.quantizer = ResidualFSQ(
+ dim=self.fsq_dim,
+ levels=self.fsq_input_levels,
+ num_quantizers=self.fsq_input_num_quantizers,
+ force_quantization_f32=False,
+ )
+
+ @can_return_tuple
+ def forward(
+ self,
+ hidden_states: Optional[torch.FloatTensor] = None,
+ **flash_attn_kwargs: Unpack[FlashAttentionKwargs],
+ ) -> tuple[torch.Tensor, torch.Tensor]:
+ hidden_states = self.audio_acoustic_proj(hidden_states)
+ hidden_states = self.attention_pooler(hidden_states)
+ quantized, indices = self.quantizer(hidden_states)
+ return quantized, indices
+
+ def tokenize(self, x):
+ """Convenience: takes [B, T, 64], rearranges to patches, runs forward."""
+ x = rearrange(x, 'n (t_patch p) d -> n t_patch p d', p=self.pool_window_size)
+ return self.forward(x)
+
+
+class AudioTokenDetokenizer(nn.Module):
+ """Converts quantized audio tokens back to continuous acoustic representations.
+
+ Input: [B, T/5, hidden_size] (quantized vectors)
+ Output: [B, T, 64] (VAE-latent-shaped continuous features)
+ """
+
+ def __init__(
+ self,
+ hidden_size: int = 2048,
+ intermediate_size: int = 6144,
+ num_attention_heads: int = 16,
+ num_key_value_heads: int = 8,
+ rms_norm_eps: float = 1e-6,
+ attention_bias: bool = False,
+ attention_dropout: float = 0.0,
+ layer_types: Optional[list] = None,
+ head_dim: Optional[int] = None,
+ sliding_window: Optional[int] = 128,
+ use_sliding_window: bool = True,
+ rope_theta: float = 1000000,
+ max_position_embeddings: int = 32768,
+ initializer_range: float = 0.02,
+ pool_window_size: int = 5,
+ audio_acoustic_hidden_dim: int = 64,
+ num_attention_pooler_hidden_layers: int = 2,
+ **kwargs,
+ ):
+ super().__init__()
+ self.hidden_size = hidden_size
+ self.intermediate_size = intermediate_size
+ self.num_attention_heads = num_attention_heads
+ self.num_key_value_heads = num_key_value_heads
+ self.rms_norm_eps = rms_norm_eps
+ self.attention_bias = attention_bias
+ self.attention_dropout = attention_dropout
+ # Default matches target library config (24 alternating entries).
+ self.layer_types = layer_types or (["sliding_attention", "full_attention"] * 12)
+ self.head_dim = head_dim or hidden_size // num_attention_heads
+ self.sliding_window = sliding_window
+ self.use_sliding_window = use_sliding_window
+ self.rope_theta = rope_theta
+ self.max_position_embeddings = max_position_embeddings
+ self.initializer_range = initializer_range
+ self.pool_window_size = pool_window_size
+ self.audio_acoustic_hidden_dim = audio_acoustic_hidden_dim
+ self.num_attention_pooler_hidden_layers = num_attention_pooler_hidden_layers
+ self._attn_implementation = kwargs.get("_attn_implementation", "sdpa")
+
+ self.embed_tokens = nn.Linear(hidden_size, hidden_size)
+ self.norm = Qwen3RMSNorm(hidden_size, eps=rms_norm_eps)
+ # Slice layer_types to our own layer count (use num_audio_decoder_hidden_layers)
+ detok_layer_types = self.layer_types[:num_attention_pooler_hidden_layers]
+ rope_config = type('RopeConfig', (), {
+ 'hidden_size': hidden_size,
+ 'num_attention_heads': num_attention_heads,
+ 'num_key_value_heads': num_key_value_heads,
+ 'head_dim': head_dim,
+ 'max_position_embeddings': max_position_embeddings,
+ 'rope_theta': rope_theta,
+ 'rope_parameters': {'rope_type': 'default', 'rope_theta': rope_theta},
+ 'rms_norm_eps': rms_norm_eps,
+ 'attention_bias': attention_bias,
+ 'attention_dropout': attention_dropout,
+ 'hidden_act': 'silu',
+ 'intermediate_size': intermediate_size,
+ 'layer_types': detok_layer_types,
+ 'sliding_window': sliding_window,
+ '_attn_implementation': self._attn_implementation,
+ })()
+ self.rotary_emb = Qwen3RotaryEmbedding(rope_config)
+ self.gradient_checkpointing = False
+ self.special_tokens = nn.Parameter(torch.randn(1, pool_window_size, hidden_size) * 0.02)
+ self.layers = nn.ModuleList([
+ AceStepEncoderLayer(
+ hidden_size=hidden_size,
+ intermediate_size=intermediate_size,
+ num_attention_heads=num_attention_heads,
+ num_key_value_heads=num_key_value_heads,
+ rms_norm_eps=rms_norm_eps,
+ attention_bias=attention_bias,
+ attention_dropout=attention_dropout,
+ layer_types=detok_layer_types,
+ head_dim=head_dim,
+ sliding_window=sliding_window,
+ layer_idx=layer_idx,
+ )
+ for layer_idx in range(num_attention_pooler_hidden_layers)
+ ])
+ self.proj_out = nn.Linear(hidden_size, audio_acoustic_hidden_dim)
+
+ @can_return_tuple
+ def forward(
+ self,
+ x,
+ attention_mask: Optional[torch.Tensor] = None,
+ **flash_attn_kwargs: Unpack[FlashAttentionKwargs],
+ ) -> torch.Tensor:
+ B, T, D = x.shape
+ x = self.embed_tokens(x)
+ x = x.unsqueeze(2).repeat(1, 1, self.pool_window_size, 1)
+ special_tokens = self.special_tokens.expand(B, T, -1, -1)
+ x = x + special_tokens.to(x.device)
+ x = rearrange(x, "b t p c -> (b t) p c")
+
+ cache_position = torch.arange(0, x.shape[1], device=x.device)
+ position_ids = cache_position.unsqueeze(0)
+ hidden_states = x
+ position_embeddings = self.rotary_emb(hidden_states, position_ids)
+
+ seq_len = x.shape[1]
+ dtype = x.dtype
+ device = x.device
+
+ full_attn_mask = create_4d_mask(
+ seq_len=seq_len, dtype=dtype, device=device,
+ attention_mask=attention_mask, sliding_window=None,
+ is_sliding_window=False, is_causal=False
+ )
+ sliding_attn_mask = None
+ if self.use_sliding_window:
+ sliding_attn_mask = create_4d_mask(
+ seq_len=seq_len, dtype=dtype, device=device,
+ attention_mask=attention_mask, sliding_window=self.sliding_window,
+ is_sliding_window=True, is_causal=False
+ )
+
+ self_attn_mask_mapping = {
+ "full_attention": full_attn_mask,
+ "sliding_attention": sliding_attn_mask,
+ }
+
+ for layer_module in self.layers:
+ layer_outputs = layer_module(
+ hidden_states, position_embeddings,
+ attention_mask=self_attn_mask_mapping[layer_module.attention_type],
+ **flash_attn_kwargs,
+ )
+ hidden_states = layer_outputs[0]
+
+ hidden_states = self.norm(hidden_states)
+ hidden_states = self.proj_out(hidden_states)
+ return rearrange(hidden_states, "(b t) p c -> b (t p) c", b=B, p=self.pool_window_size)
+
+
+class AceStepTokenizer(nn.Module):
+ """Container for AceStepAudioTokenizer + AudioTokenDetokenizer.
+
+ Provides encode/decode convenience methods for VAE latent discretization.
+ Used in cover song mode to convert source audio latents to discrete tokens
+ and back to continuous conditioning hints.
+ """
+
+ def __init__(
+ self,
+ hidden_size: int = 2048,
+ intermediate_size: int = 6144,
+ num_attention_heads: int = 16,
+ num_key_value_heads: int = 8,
+ rms_norm_eps: float = 1e-6,
+ attention_bias: bool = False,
+ attention_dropout: float = 0.0,
+ layer_types: Optional[list] = None,
+ head_dim: Optional[int] = None,
+ sliding_window: Optional[int] = 128,
+ use_sliding_window: bool = True,
+ rope_theta: float = 1000000,
+ max_position_embeddings: int = 32768,
+ initializer_range: float = 0.02,
+ audio_acoustic_hidden_dim: int = 64,
+ pool_window_size: int = 5,
+ fsq_dim: int = 2048,
+ fsq_input_levels: list = None,
+ fsq_input_num_quantizers: int = 1,
+ num_attention_pooler_hidden_layers: int = 2,
+ num_audio_decoder_hidden_layers: int = 24,
+ **kwargs,
+ ):
+ super().__init__()
+ if layer_types is None:
+ layer_types = ["sliding_attention", "full_attention"] * 12
+ self.tokenizer = AceStepAudioTokenizer(
+ hidden_size=hidden_size,
+ intermediate_size=intermediate_size,
+ num_attention_heads=num_attention_heads,
+ num_key_value_heads=num_key_value_heads,
+ rms_norm_eps=rms_norm_eps,
+ attention_bias=attention_bias,
+ attention_dropout=attention_dropout,
+ layer_types=layer_types,
+ head_dim=head_dim,
+ sliding_window=sliding_window,
+ use_sliding_window=use_sliding_window,
+ rope_theta=rope_theta,
+ max_position_embeddings=max_position_embeddings,
+ initializer_range=initializer_range,
+ audio_acoustic_hidden_dim=audio_acoustic_hidden_dim,
+ pool_window_size=pool_window_size,
+ fsq_dim=fsq_dim,
+ fsq_input_levels=fsq_input_levels,
+ fsq_input_num_quantizers=fsq_input_num_quantizers,
+ num_attention_pooler_hidden_layers=num_attention_pooler_hidden_layers,
+ **kwargs,
+ )
+ self.detokenizer = AudioTokenDetokenizer(
+ hidden_size=hidden_size,
+ intermediate_size=intermediate_size,
+ num_attention_heads=num_attention_heads,
+ num_key_value_heads=num_key_value_heads,
+ rms_norm_eps=rms_norm_eps,
+ attention_bias=attention_bias,
+ attention_dropout=attention_dropout,
+ layer_types=layer_types,
+ head_dim=head_dim,
+ sliding_window=sliding_window,
+ use_sliding_window=use_sliding_window,
+ rope_theta=rope_theta,
+ max_position_embeddings=max_position_embeddings,
+ initializer_range=initializer_range,
+ pool_window_size=pool_window_size,
+ audio_acoustic_hidden_dim=audio_acoustic_hidden_dim,
+ num_attention_pooler_hidden_layers=num_attention_pooler_hidden_layers,
+ **kwargs,
+ )
+
+ def encode(self, hidden_states: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
+ """VAE latent [B, T, 64] → discrete tokens."""
+ return self.tokenizer(hidden_states)
+
+ def decode(self, quantized: torch.Tensor) -> torch.Tensor:
+ """Discrete tokens [B, T/5, hidden_size] → continuous [B, T, 64]."""
+ return self.detokenizer(quantized)
+
+ def tokenize(self, x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
+ """Convenience: [B, T, 64] → quantized + indices via patch rearrangement."""
+ return self.tokenizer.tokenize(x)
diff --git a/diffsynth/models/ace_step_vae.py b/diffsynth/models/ace_step_vae.py
new file mode 100644
index 0000000000000000000000000000000000000000..7a67da343a697c8c70374880cbc39c8a9a5b8cf9
--- /dev/null
+++ b/diffsynth/models/ace_step_vae.py
@@ -0,0 +1,383 @@
+# Copyright 2025 The ACESTEO Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+import math
+from typing import Optional
+
+import torch
+import torch.nn as nn
+from torch.nn.utils import weight_norm, remove_weight_norm
+
+
+class Snake1d(nn.Module):
+ """Snake activation: x + 1/(beta+eps) * sin(alpha*x)^2."""
+
+ def __init__(self, hidden_dim: int, logscale: bool = True):
+ super().__init__()
+ self.alpha = nn.Parameter(torch.zeros(1, hidden_dim, 1))
+ self.beta = nn.Parameter(torch.zeros(1, hidden_dim, 1))
+ self.logscale = logscale
+
+ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+ shape = hidden_states.shape
+ alpha = torch.exp(self.alpha) if self.logscale else self.alpha
+ beta = torch.exp(self.beta) if self.logscale else self.beta
+ hidden_states = hidden_states.reshape(shape[0], shape[1], -1)
+ hidden_states = hidden_states + (beta + 1e-9).reciprocal() * torch.sin(alpha * hidden_states).pow(2)
+ return hidden_states.reshape(shape)
+
+
+class OobleckResidualUnit(nn.Module):
+ """Residual unit: Snake1d → Conv1d(dilated) → Snake1d → Conv1d(1×1) + skip."""
+
+ def __init__(self, dimension: int = 16, dilation: int = 1):
+ super().__init__()
+ pad = ((7 - 1) * dilation) // 2
+ self.snake1 = Snake1d(dimension)
+ self.conv1 = weight_norm(nn.Conv1d(dimension, dimension, kernel_size=7, dilation=dilation, padding=pad))
+ self.snake2 = Snake1d(dimension)
+ self.conv2 = weight_norm(nn.Conv1d(dimension, dimension, kernel_size=1))
+
+ def forward(self, hidden_state: torch.Tensor) -> torch.Tensor:
+ output = self.conv1(self.snake1(hidden_state))
+ output = self.conv2(self.snake2(output))
+ padding = (hidden_state.shape[-1] - output.shape[-1]) // 2
+ if padding > 0:
+ hidden_state = hidden_state[..., padding:-padding]
+ return hidden_state + output
+
+
+class OobleckEncoderBlock(nn.Module):
+ """Encoder block: 3 residual units + downsampling conv."""
+
+ def __init__(self, input_dim: int, output_dim: int, stride: int = 1):
+ super().__init__()
+ self.res_unit1 = OobleckResidualUnit(input_dim, dilation=1)
+ self.res_unit2 = OobleckResidualUnit(input_dim, dilation=3)
+ self.res_unit3 = OobleckResidualUnit(input_dim, dilation=9)
+ self.snake1 = Snake1d(input_dim)
+ self.conv1 = weight_norm(
+ nn.Conv1d(input_dim, output_dim, kernel_size=2 * stride, stride=stride, padding=math.ceil(stride / 2))
+ )
+
+ def forward(self, hidden_state: torch.Tensor) -> torch.Tensor:
+ hidden_state = self.res_unit1(hidden_state)
+ hidden_state = self.res_unit2(hidden_state)
+ hidden_state = self.snake1(self.res_unit3(hidden_state))
+ return self.conv1(hidden_state)
+
+
+class OobleckDecoderBlock(nn.Module):
+ """Decoder block: upsampling conv + 3 residual units."""
+
+ def __init__(self, input_dim: int, output_dim: int, stride: int = 1):
+ super().__init__()
+ self.snake1 = Snake1d(input_dim)
+ self.conv_t1 = weight_norm(
+ nn.ConvTranspose1d(
+ input_dim, output_dim, kernel_size=2 * stride, stride=stride, padding=math.ceil(stride / 2),
+ )
+ )
+ self.res_unit1 = OobleckResidualUnit(output_dim, dilation=1)
+ self.res_unit2 = OobleckResidualUnit(output_dim, dilation=3)
+ self.res_unit3 = OobleckResidualUnit(output_dim, dilation=9)
+
+ def forward(self, hidden_state: torch.Tensor) -> torch.Tensor:
+ hidden_state = self.snake1(hidden_state)
+ hidden_state = self.conv_t1(hidden_state)
+ hidden_state = self.res_unit1(hidden_state)
+ hidden_state = self.res_unit2(hidden_state)
+ return self.res_unit3(hidden_state)
+
+
+class OobleckEncoder(nn.Module):
+ """Full encoder: audio → latent representation [B, encoder_hidden_size, T'].
+
+ conv1 → [blocks] → snake1 → conv2
+ """
+
+ def __init__(
+ self,
+ encoder_hidden_size: int = 128,
+ audio_channels: int = 2,
+ downsampling_ratios: list = None,
+ channel_multiples: list = None,
+ ):
+ super().__init__()
+ downsampling_ratios = downsampling_ratios or [2, 4, 4, 6, 10]
+ channel_multiples = channel_multiples or [1, 2, 4, 8, 16]
+ channel_multiples = [1] + channel_multiples
+
+ self.conv1 = weight_norm(nn.Conv1d(audio_channels, encoder_hidden_size, kernel_size=7, padding=3))
+
+ self.block = nn.ModuleList()
+ for stride_index, stride in enumerate(downsampling_ratios):
+ self.block.append(
+ OobleckEncoderBlock(
+ input_dim=encoder_hidden_size * channel_multiples[stride_index],
+ output_dim=encoder_hidden_size * channel_multiples[stride_index + 1],
+ stride=stride,
+ )
+ )
+
+ d_model = encoder_hidden_size * channel_multiples[-1]
+ self.snake1 = Snake1d(d_model)
+ self.conv2 = weight_norm(nn.Conv1d(d_model, encoder_hidden_size, kernel_size=3, padding=1))
+
+ def forward(self, hidden_state: torch.Tensor) -> torch.Tensor:
+ hidden_state = self.conv1(hidden_state)
+ for block in self.block:
+ hidden_state = block(hidden_state)
+ hidden_state = self.snake1(hidden_state)
+ return self.conv2(hidden_state)
+
+
+class OobleckDecoder(nn.Module):
+ """Full decoder: latent → audio waveform [B, audio_channels, T].
+
+ conv1 → [blocks] → snake1 → conv2(no bias)
+ """
+
+ def __init__(
+ self,
+ channels: int = 128,
+ input_channels: int = 64,
+ audio_channels: int = 2,
+ upsampling_ratios: list = None,
+ channel_multiples: list = None,
+ ):
+ super().__init__()
+ upsampling_ratios = upsampling_ratios or [10, 6, 4, 4, 2]
+ channel_multiples = channel_multiples or [1, 2, 4, 8, 16]
+ channel_multiples = [1] + channel_multiples
+
+ self.conv1 = weight_norm(nn.Conv1d(input_channels, channels * channel_multiples[-1], kernel_size=7, padding=3))
+
+ self.block = nn.ModuleList()
+ for stride_index, stride in enumerate(upsampling_ratios):
+ self.block.append(
+ OobleckDecoderBlock(
+ input_dim=channels * channel_multiples[len(upsampling_ratios) - stride_index],
+ output_dim=channels * channel_multiples[len(upsampling_ratios) - stride_index - 1],
+ stride=stride,
+ )
+ )
+
+ self.snake1 = Snake1d(channels)
+ # conv2 has no bias (matches checkpoint: only weight_g/weight_v, no bias key)
+ self.conv2 = weight_norm(nn.Conv1d(channels, audio_channels, kernel_size=7, padding=3, bias=False))
+
+ def forward(self, hidden_state: torch.Tensor) -> torch.Tensor:
+ hidden_state = self.conv1(hidden_state)
+ for block in self.block:
+ hidden_state = block(hidden_state)
+ hidden_state = self.snake1(hidden_state)
+ return self.conv2(hidden_state)
+
+
+class OobleckDiagonalGaussianDistribution(object):
+ def __init__(self, parameters: torch.Tensor, deterministic: bool = False):
+ self.parameters = parameters
+ self.mean, self.scale = parameters.chunk(2, dim=1)
+ self.std = nn.functional.softplus(self.scale) + 1e-4
+ self.var = self.std * self.std
+ self.logvar = torch.log(self.var)
+ self.deterministic = deterministic
+
+ def sample(self, generator: torch.Generator | None = None) -> torch.Tensor:
+ # make sure sample is on the same device as the parameters and has same dtype
+ sample = torch.randn(
+ self.mean.shape,
+ generator=generator,
+ device=self.parameters.device,
+ dtype=self.parameters.dtype,
+ )
+ x = self.mean + self.std * sample
+ return x
+
+ def kl(self, other: "OobleckDiagonalGaussianDistribution" = None) -> torch.Tensor:
+ if self.deterministic:
+ return torch.Tensor([0.0])
+ else:
+ if other is None:
+ return (self.mean * self.mean + self.var - self.logvar - 1.0).sum(1).mean()
+ else:
+ normalized_diff = torch.pow(self.mean - other.mean, 2) / other.var
+ var_ratio = self.var / other.var
+ logvar_diff = self.logvar - other.logvar
+
+ kl = normalized_diff + var_ratio + logvar_diff - 1
+
+ kl = kl.sum(1).mean()
+ return kl
+
+
+class AceStepVAE(nn.Module):
+ """Audio VAE for ACE-Step (AutoencoderOobleck architecture).
+
+ Encodes audio waveform → latent, decodes latent → audio waveform.
+ Uses Snake1d activations and weight-normalized convolutions.
+ """
+
+ def __init__(
+ self,
+ encoder_hidden_size: int = 128,
+ downsampling_ratios: list = None,
+ channel_multiples: list = None,
+ decoder_channels: int = 128,
+ decoder_input_channels: int = 64,
+ audio_channels: int = 2,
+ sampling_rate: int = 48000,
+ ):
+ super().__init__()
+ downsampling_ratios = downsampling_ratios or [2, 4, 4, 6, 10]
+ channel_multiples = channel_multiples or [1, 2, 4, 8, 16]
+ upsampling_ratios = downsampling_ratios[::-1]
+
+ self.encoder = OobleckEncoder(
+ encoder_hidden_size=encoder_hidden_size,
+ audio_channels=audio_channels,
+ downsampling_ratios=downsampling_ratios,
+ channel_multiples=channel_multiples,
+ )
+ self.decoder = OobleckDecoder(
+ channels=decoder_channels,
+ input_channels=decoder_input_channels,
+ audio_channels=audio_channels,
+ upsampling_ratios=upsampling_ratios,
+ channel_multiples=channel_multiples,
+ )
+ self.sampling_rate = sampling_rate
+
+ def tiled_encode(self, x: torch.Tensor, tile_size: int = 10240, tile_stride: int = 5120) -> torch.Tensor:
+ batch_size, audio_channels, T_audio = x.shape
+ up = self.upsampling_factor
+
+ lat_T = round(T_audio / up)
+
+ tiles = []
+ for t in range(0, T_audio, tile_stride):
+ if t - tile_stride >= 0 and t - tile_stride + tile_size >= T_audio:
+ continue
+ tiles.append((t, min(t + tile_size, T_audio)))
+
+ encoder_out_channels = self.encoder.conv2.out_channels
+ values = torch.zeros((batch_size, encoder_out_channels, lat_T), dtype=x.dtype, device="cpu")
+ weight = torch.zeros((1, 1, lat_T), dtype=x.dtype, device="cpu")
+
+ for t_start, t_end in tiles:
+ tile_audio = x[:, :, t_start:t_end]
+ tile_latent = self.encoder(tile_audio).to("cpu")
+ lat_len = tile_latent.shape[-1]
+
+ lat_start = round(t_start / up)
+ lat_end = lat_start + lat_len
+
+ border = (tile_size - tile_stride) // up
+ mask = self._build_blend_mask_1d(
+ lat_len, border,
+ is_left_bound=(t_start == 0),
+ is_right_bound=(t_end >= T_audio),
+ dtype=x.dtype, device="cpu",
+ )
+
+ values[:, :, lat_start:lat_end] += tile_latent * mask
+ weight[:, :, lat_start:lat_end] += mask
+
+ weight = weight.clamp(min=1e-8)
+ latent = values / weight
+ output = OobleckDiagonalGaussianDistribution(latent.to(x.device)).sample()
+ return output
+
+ def encode(self, x: torch.Tensor, tiled: bool = False, tile_size: int = 10240, tile_stride: int = 5120) -> torch.Tensor:
+ """Audio waveform [B, audio_channels, T] → latent [B, decoder_input_channels, T"]."""
+ if tiled:
+ return self.tiled_encode(x, tile_size=tile_size, tile_stride=tile_stride)
+ h = self.encoder(x)
+ output = OobleckDiagonalGaussianDistribution(h).sample()
+ return output
+
+ @property
+ def upsampling_factor(self) -> int:
+ ratios = [10, 6, 4, 4, 2]
+ result = 1
+ for r in ratios:
+ result *= r
+ return result
+
+ def _build_blend_mask_1d(self, length: int, border_width: int, is_left_bound: bool, is_right_bound: bool,
+ dtype: torch.dtype, device: torch.device) -> torch.Tensor:
+ mask = torch.ones(length, dtype=dtype, device=device)
+ if border_width <= 0:
+ return mask
+ if not is_left_bound:
+ ramp = torch.linspace(0.0, 1.0, border_width, dtype=dtype, device=device)
+ mask[:border_width] = ramp
+ if not is_right_bound:
+ ramp = torch.linspace(1.0, 0.0, border_width, dtype=dtype, device=device)
+ mask[-border_width:] = ramp
+ return mask.unsqueeze(0).unsqueeze(0)
+
+ def tiled_decode(self, z: torch.Tensor, tile_size: int = 512, tile_stride: int = 256) -> torch.Tensor:
+ batch_size, channels, T = z.shape
+ up = self.upsampling_factor
+ out_length = up * T
+
+ tiles = []
+ for t in range(0, T, tile_stride):
+ if t - tile_stride >= 0 and t - tile_stride + tile_size >= T:
+ continue
+ tiles.append((t, min(t + tile_size, T)))
+
+ audio_channels = self.decoder.conv2.out_channels
+ values = torch.zeros((batch_size, audio_channels, out_length), dtype=z.dtype, device="cpu")
+ weight = torch.zeros((1, 1, out_length), dtype=z.dtype, device="cpu")
+
+ for t_start, t_end in tiles:
+ tile_latent = z[:, :, t_start:t_end]
+ tile_output = self.decoder(tile_latent).to("cpu")
+ tile_len = tile_output.shape[-1]
+
+ out_start = t_start * up
+ out_end = out_start + tile_len
+
+ border = (tile_size - tile_stride) * up
+ mask = self._build_blend_mask_1d(
+ tile_len, border,
+ is_left_bound=(t_start == 0),
+ is_right_bound=(t_end >= T),
+ dtype=z.dtype, device="cpu",
+ )
+
+ values[:, :, out_start:out_end] += tile_output * mask
+ weight[:, :, out_start:out_end] += mask
+
+ weight = weight.clamp(min=1e-8)
+ return values / weight
+
+ def decode(self, z: torch.Tensor, tiled: bool = False, tile_size: int = 512, tile_stride: int = 256) -> torch.Tensor:
+ """Latent [B, decoder_input_channels, T] → audio waveform [B, audio_channels, T]."""
+ if tiled:
+ return self.tiled_decode(z, tile_size=tile_size, tile_stride=tile_stride)
+ return self.decoder(z)
+
+ def forward(self, sample: torch.Tensor) -> torch.Tensor:
+ """Full round-trip: encode → decode."""
+ z = self.encode(sample)
+ return self.decode(z)
+
+ def remove_weight_norm(self):
+ """Remove weight normalization from all conv layers (for export/inference)."""
+ for module in self.modules():
+ if isinstance(module, nn.Conv1d) or isinstance(module, nn.ConvTranspose1d):
+ remove_weight_norm(module)
diff --git a/diffsynth/models/aesthetic.py b/diffsynth/models/aesthetic.py
new file mode 100644
index 0000000000000000000000000000000000000000..b1ed41462cf106cffefdabf9f28247c540c0e783
--- /dev/null
+++ b/diffsynth/models/aesthetic.py
@@ -0,0 +1,90 @@
+from typing import Union
+import torch
+from PIL import Image
+
+ImageInput = Union[Image.Image, list[Image.Image], tuple[Image.Image, ...]]
+
+class AestheticMLP(torch.nn.Module):
+ def __init__(self, input_size: int):
+ super().__init__()
+ self.input_size = input_size
+ self.layers = torch.nn.Sequential(
+ torch.nn.Linear(input_size, 1024),
+ torch.nn.Dropout(0.2),
+ torch.nn.Linear(1024, 128),
+ torch.nn.Dropout(0.2),
+ torch.nn.Linear(128, 64),
+ torch.nn.Dropout(0.1),
+ torch.nn.Linear(64, 16),
+ torch.nn.Linear(16, 1),
+ )
+
+ def forward(self, x):
+ return self.layers(x)
+
+
+def _as_image_list(images: ImageInput):
+ if isinstance(images, Image.Image):
+ images = [images]
+ return [image.convert("RGB") for image in images]
+
+
+class AestheticModel(torch.nn.Module):
+ def __init__(
+ self,
+ mlp: AestheticMLP = None,
+ vision_model: torch.nn.Module = None,
+ visual_projection: torch.nn.Module = None,
+ processor=None
+ ):
+ super().__init__()
+ if vision_model is None:
+ vision_model, visual_projection = self.default_vision_model()
+ if mlp is None:
+ mlp = AestheticMLP(768)
+
+ self.vision_model = vision_model
+ self.visual_projection = visual_projection
+ self.processor = processor
+ self.layers = mlp.layers
+
+ @staticmethod
+ def default_vision_model():
+ from transformers import CLIPVisionConfig, CLIPVisionModel
+
+ config = CLIPVisionConfig(
+ hidden_size=1024,
+ intermediate_size=4096,
+ num_attention_heads=16,
+ num_hidden_layers=24,
+ image_size=224,
+ patch_size=14,
+ hidden_act="quick_gelu",
+ layer_norm_eps=1e-5,
+ projection_dim=768,
+ )
+ return CLIPVisionModel(config), torch.nn.Linear(config.hidden_size, config.projection_dim, bias=False)
+
+ @property
+ def device(self):
+ return next(self.parameters(), torch.tensor([])).device
+
+ @property
+ def dtype(self):
+ return next(self.parameters(), torch.tensor(0.0)).dtype
+
+ @torch.no_grad()
+ def get_image_features(self, images):
+ images = _as_image_list(images)
+ inputs = self.processor(images=images, return_tensors="pt")
+ pixel_values = inputs["pixel_values"].to(device=self.device, dtype=self.dtype)
+
+ image_features = self.vision_model(pixel_values=pixel_values, return_dict=True).pooler_output
+ image_features = self.visual_projection(image_features)
+
+ return torch.nn.functional.normalize(image_features, dim=-1)
+
+ @torch.no_grad()
+ def forward(self, images):
+ image_features = self.get_image_features(images)
+ return self.layers(image_features).squeeze(-1)
\ No newline at end of file
diff --git a/diffsynth/models/anima_dit.py b/diffsynth/models/anima_dit.py
new file mode 100644
index 0000000000000000000000000000000000000000..d7519802beeabeff358f873fd3a6bd83cfc54ede
--- /dev/null
+++ b/diffsynth/models/anima_dit.py
@@ -0,0 +1,1307 @@
+# original code from: comfy/ldm/cosmos/predict2.py
+
+import torch
+from torch import nn
+from einops import rearrange, repeat
+from einops.layers.torch import Rearrange
+import logging
+from typing import Callable, Optional, Tuple, List
+import math
+from torchvision import transforms
+from ..core.attention import attention_forward
+from ..core.gradient import gradient_checkpoint_forward
+
+
+class VideoPositionEmb(nn.Module):
+ def forward(self, x_B_T_H_W_C: torch.Tensor, fps=Optional[torch.Tensor], device=None, dtype=None) -> torch.Tensor:
+ """
+ It delegates the embedding generation to generate_embeddings function.
+ """
+ B_T_H_W_C = x_B_T_H_W_C.shape
+ embeddings = self.generate_embeddings(B_T_H_W_C, fps=fps, device=device, dtype=dtype)
+
+ return embeddings
+
+ def generate_embeddings(self, B_T_H_W_C: torch.Size, fps=Optional[torch.Tensor], device=None):
+ raise NotImplementedError
+
+
+def normalize(x: torch.Tensor, dim: Optional[List[int]] = None, eps: float = 0) -> torch.Tensor:
+ """
+ Normalizes the input tensor along specified dimensions such that the average square norm of elements is adjusted.
+
+ Args:
+ x (torch.Tensor): The input tensor to normalize.
+ dim (list, optional): The dimensions over which to normalize. If None, normalizes over all dimensions except the first.
+ eps (float, optional): A small constant to ensure numerical stability during division.
+
+ Returns:
+ torch.Tensor: The normalized tensor.
+ """
+ if dim is None:
+ dim = list(range(1, x.ndim))
+ norm = torch.linalg.vector_norm(x, dim=dim, keepdim=True, dtype=torch.float32)
+ norm = torch.add(eps, norm, alpha=math.sqrt(norm.numel() / x.numel()))
+ return x / norm.to(x.dtype)
+
+
+class LearnablePosEmbAxis(VideoPositionEmb):
+ def __init__(
+ self,
+ *, # enforce keyword arguments
+ interpolation: str,
+ model_channels: int,
+ len_h: int,
+ len_w: int,
+ len_t: int,
+ device=None,
+ dtype=None,
+ **kwargs,
+ ):
+ """
+ Args:
+ interpolation (str): we curretly only support "crop", ideally when we need extrapolation capacity, we should adjust frequency or other more advanced methods. they are not implemented yet.
+ """
+ del kwargs # unused
+ super().__init__()
+ self.interpolation = interpolation
+ assert self.interpolation in ["crop"], f"Unknown interpolation method {self.interpolation}"
+
+ self.pos_emb_h = nn.Parameter(torch.empty(len_h, model_channels, device=device, dtype=dtype))
+ self.pos_emb_w = nn.Parameter(torch.empty(len_w, model_channels, device=device, dtype=dtype))
+ self.pos_emb_t = nn.Parameter(torch.empty(len_t, model_channels, device=device, dtype=dtype))
+
+ def generate_embeddings(self, B_T_H_W_C: torch.Size, fps=Optional[torch.Tensor], device=None, dtype=None) -> torch.Tensor:
+ B, T, H, W, _ = B_T_H_W_C
+ if self.interpolation == "crop":
+ emb_h_H = self.pos_emb_h[:H].to(device=device, dtype=dtype)
+ emb_w_W = self.pos_emb_w[:W].to(device=device, dtype=dtype)
+ emb_t_T = self.pos_emb_t[:T].to(device=device, dtype=dtype)
+ emb = (
+ repeat(emb_t_T, "t d-> b t h w d", b=B, h=H, w=W)
+ + repeat(emb_h_H, "h d-> b t h w d", b=B, t=T, w=W)
+ + repeat(emb_w_W, "w d-> b t h w d", b=B, t=T, h=H)
+ )
+ assert list(emb.shape)[:4] == [B, T, H, W], f"bad shape: {list(emb.shape)[:4]} != {B, T, H, W}"
+ else:
+ raise ValueError(f"Unknown interpolation method {self.interpolation}")
+
+ return normalize(emb, dim=-1, eps=1e-6)
+
+
+class VideoRopePosition3DEmb(VideoPositionEmb):
+ def __init__(
+ self,
+ *, # enforce keyword arguments
+ head_dim: int,
+ len_h: int,
+ len_w: int,
+ len_t: int,
+ base_fps: int = 24,
+ h_extrapolation_ratio: float = 1.0,
+ w_extrapolation_ratio: float = 1.0,
+ t_extrapolation_ratio: float = 1.0,
+ enable_fps_modulation: bool = True,
+ device=None,
+ **kwargs, # used for compatibility with other positional embeddings; unused in this class
+ ):
+ del kwargs
+ super().__init__()
+ self.base_fps = base_fps
+ self.max_h = len_h
+ self.max_w = len_w
+ self.enable_fps_modulation = enable_fps_modulation
+
+ dim = head_dim
+ dim_h = dim // 6 * 2
+ dim_w = dim_h
+ dim_t = dim - 2 * dim_h
+ assert dim == dim_h + dim_w + dim_t, f"bad dim: {dim} != {dim_h} + {dim_w} + {dim_t}"
+ self.register_buffer(
+ "dim_spatial_range",
+ torch.arange(0, dim_h, 2, device=device)[: (dim_h // 2)].float() / dim_h,
+ persistent=False,
+ )
+ self.register_buffer(
+ "dim_temporal_range",
+ torch.arange(0, dim_t, 2, device=device)[: (dim_t // 2)].float() / dim_t,
+ persistent=False,
+ )
+
+ self.h_ntk_factor = h_extrapolation_ratio ** (dim_h / (dim_h - 2))
+ self.w_ntk_factor = w_extrapolation_ratio ** (dim_w / (dim_w - 2))
+ self.t_ntk_factor = t_extrapolation_ratio ** (dim_t / (dim_t - 2))
+
+ def generate_embeddings(
+ self,
+ B_T_H_W_C: torch.Size,
+ fps: Optional[torch.Tensor] = None,
+ h_ntk_factor: Optional[float] = None,
+ w_ntk_factor: Optional[float] = None,
+ t_ntk_factor: Optional[float] = None,
+ device=None,
+ dtype=None,
+ ):
+ """
+ Generate embeddings for the given input size.
+
+ Args:
+ B_T_H_W_C (torch.Size): Input tensor size (Batch, Time, Height, Width, Channels).
+ fps (Optional[torch.Tensor], optional): Frames per second. Defaults to None.
+ h_ntk_factor (Optional[float], optional): Height NTK factor. If None, uses self.h_ntk_factor.
+ w_ntk_factor (Optional[float], optional): Width NTK factor. If None, uses self.w_ntk_factor.
+ t_ntk_factor (Optional[float], optional): Time NTK factor. If None, uses self.t_ntk_factor.
+
+ Returns:
+ Not specified in the original code snippet.
+ """
+ h_ntk_factor = h_ntk_factor if h_ntk_factor is not None else self.h_ntk_factor
+ w_ntk_factor = w_ntk_factor if w_ntk_factor is not None else self.w_ntk_factor
+ t_ntk_factor = t_ntk_factor if t_ntk_factor is not None else self.t_ntk_factor
+
+ h_theta = 10000.0 * h_ntk_factor
+ w_theta = 10000.0 * w_ntk_factor
+ t_theta = 10000.0 * t_ntk_factor
+
+ h_spatial_freqs = 1.0 / (h_theta**self.dim_spatial_range.to(device=device))
+ w_spatial_freqs = 1.0 / (w_theta**self.dim_spatial_range.to(device=device))
+ temporal_freqs = 1.0 / (t_theta**self.dim_temporal_range.to(device=device))
+
+ B, T, H, W, _ = B_T_H_W_C
+ seq = torch.arange(max(H, W, T), dtype=torch.float, device=device)
+ uniform_fps = (fps is None) or isinstance(fps, (int, float)) or (fps.min() == fps.max())
+ assert (
+ uniform_fps or B == 1 or T == 1
+ ), "For video batch, batch size should be 1 for non-uniform fps. For image batch, T should be 1"
+ half_emb_h = torch.outer(seq[:H].to(device=device), h_spatial_freqs)
+ half_emb_w = torch.outer(seq[:W].to(device=device), w_spatial_freqs)
+
+ # apply sequence scaling in temporal dimension
+ if fps is None or self.enable_fps_modulation is False: # image case
+ half_emb_t = torch.outer(seq[:T].to(device=device), temporal_freqs)
+ else:
+ half_emb_t = torch.outer(seq[:T].to(device=device) / fps * self.base_fps, temporal_freqs)
+
+ half_emb_h = torch.stack([torch.cos(half_emb_h), -torch.sin(half_emb_h), torch.sin(half_emb_h), torch.cos(half_emb_h)], dim=-1)
+ half_emb_w = torch.stack([torch.cos(half_emb_w), -torch.sin(half_emb_w), torch.sin(half_emb_w), torch.cos(half_emb_w)], dim=-1)
+ half_emb_t = torch.stack([torch.cos(half_emb_t), -torch.sin(half_emb_t), torch.sin(half_emb_t), torch.cos(half_emb_t)], dim=-1)
+
+ em_T_H_W_D = torch.cat(
+ [
+ repeat(half_emb_t, "t d x -> t h w d x", h=H, w=W),
+ repeat(half_emb_h, "h d x -> t h w d x", t=T, w=W),
+ repeat(half_emb_w, "w d x -> t h w d x", t=T, h=H),
+ ]
+ , dim=-2,
+ )
+
+ return rearrange(em_T_H_W_D, "t h w d (i j) -> (t h w) d i j", i=2, j=2).float()
+
+
+def apply_rotary_pos_emb(
+ t: torch.Tensor,
+ freqs: torch.Tensor,
+) -> torch.Tensor:
+ t_ = t.reshape(*t.shape[:-1], 2, -1).movedim(-2, -1).unsqueeze(-2).float()
+ t_out = freqs[..., 0] * t_[..., 0] + freqs[..., 1] * t_[..., 1]
+ t_out = t_out.movedim(-1, -2).reshape(*t.shape).type_as(t)
+ return t_out
+
+
+# ---------------------- Feed Forward Network -----------------------
+class GPT2FeedForward(nn.Module):
+ def __init__(self, d_model: int, d_ff: int, device=None, dtype=None, operations=None) -> None:
+ super().__init__()
+ self.activation = nn.GELU()
+ self.layer1 = operations.Linear(d_model, d_ff, bias=False, device=device, dtype=dtype)
+ self.layer2 = operations.Linear(d_ff, d_model, bias=False, device=device, dtype=dtype)
+
+ self._layer_id = None
+ self._dim = d_model
+ self._hidden_dim = d_ff
+
+ def forward(self, x: torch.Tensor) -> torch.Tensor:
+ x = self.layer1(x)
+
+ x = self.activation(x)
+ x = self.layer2(x)
+ return x
+
+
+def torch_attention_op(q_B_S_H_D: torch.Tensor, k_B_S_H_D: torch.Tensor, v_B_S_H_D: torch.Tensor, transformer_options: Optional[dict] = {}) -> torch.Tensor:
+ """Computes multi-head attention using PyTorch's native implementation.
+
+ This function provides a PyTorch backend alternative to Transformer Engine's attention operation.
+ It rearranges the input tensors to match PyTorch's expected format, computes scaled dot-product
+ attention, and rearranges the output back to the original format.
+
+ The input tensor names use the following dimension conventions:
+
+ - B: batch size
+ - S: sequence length
+ - H: number of attention heads
+ - D: head dimension
+
+ Args:
+ q_B_S_H_D: Query tensor with shape (batch, seq_len, n_heads, head_dim)
+ k_B_S_H_D: Key tensor with shape (batch, seq_len, n_heads, head_dim)
+ v_B_S_H_D: Value tensor with shape (batch, seq_len, n_heads, head_dim)
+
+ Returns:
+ Attention output tensor with shape (batch, seq_len, n_heads * head_dim)
+ """
+ in_q_shape = q_B_S_H_D.shape
+ in_k_shape = k_B_S_H_D.shape
+ q_B_H_S_D = rearrange(q_B_S_H_D, "b ... h k -> b h ... k").view(in_q_shape[0], in_q_shape[-2], -1, in_q_shape[-1])
+ k_B_H_S_D = rearrange(k_B_S_H_D, "b ... h v -> b h ... v").view(in_k_shape[0], in_k_shape[-2], -1, in_k_shape[-1])
+ v_B_H_S_D = rearrange(v_B_S_H_D, "b ... h v -> b h ... v").view(in_k_shape[0], in_k_shape[-2], -1, in_k_shape[-1])
+ return attention_forward(q_B_H_S_D, k_B_H_S_D, v_B_H_S_D, out_pattern="b s (n d)")
+
+
+class Attention(nn.Module):
+ """
+ A flexible attention module supporting both self-attention and cross-attention mechanisms.
+
+ This module implements a multi-head attention layer that can operate in either self-attention
+ or cross-attention mode. The mode is determined by whether a context dimension is provided.
+ The implementation uses scaled dot-product attention and supports optional bias terms and
+ dropout regularization.
+
+ Args:
+ query_dim (int): The dimensionality of the query vectors.
+ context_dim (int, optional): The dimensionality of the context (key/value) vectors.
+ If None, the module operates in self-attention mode using query_dim. Default: None
+ n_heads (int, optional): Number of attention heads for multi-head attention. Default: 8
+ head_dim (int, optional): The dimension of each attention head. Default: 64
+ dropout (float, optional): Dropout probability applied to the output. Default: 0.0
+ qkv_format (str, optional): Format specification for QKV tensors. Default: "bshd"
+ backend (str, optional): Backend to use for the attention operation. Default: "transformer_engine"
+
+ Examples:
+ >>> # Self-attention with 512 dimensions and 8 heads
+ >>> self_attn = Attention(query_dim=512)
+ >>> x = torch.randn(32, 16, 512) # (batch_size, seq_len, dim)
+ >>> out = self_attn(x) # (32, 16, 512)
+
+ >>> # Cross-attention
+ >>> cross_attn = Attention(query_dim=512, context_dim=256)
+ >>> query = torch.randn(32, 16, 512)
+ >>> context = torch.randn(32, 8, 256)
+ >>> out = cross_attn(query, context) # (32, 16, 512)
+ """
+
+ def __init__(
+ self,
+ query_dim: int,
+ context_dim: Optional[int] = None,
+ n_heads: int = 8,
+ head_dim: int = 64,
+ dropout: float = 0.0,
+ device=None,
+ dtype=None,
+ operations=None,
+ ) -> None:
+ super().__init__()
+ logging.debug(
+ f"Setting up {self.__class__.__name__}. Query dim is {query_dim}, context_dim is {context_dim} and using "
+ f"{n_heads} heads with a dimension of {head_dim}."
+ )
+ self.is_selfattn = context_dim is None # self attention
+
+ context_dim = query_dim if context_dim is None else context_dim
+ inner_dim = head_dim * n_heads
+
+ self.n_heads = n_heads
+ self.head_dim = head_dim
+ self.query_dim = query_dim
+ self.context_dim = context_dim
+
+ self.q_proj = operations.Linear(query_dim, inner_dim, bias=False, device=device, dtype=dtype)
+ self.q_norm = operations.RMSNorm(self.head_dim, eps=1e-6, device=device, dtype=dtype)
+
+ self.k_proj = operations.Linear(context_dim, inner_dim, bias=False, device=device, dtype=dtype)
+ self.k_norm = operations.RMSNorm(self.head_dim, eps=1e-6, device=device, dtype=dtype)
+
+ self.v_proj = operations.Linear(context_dim, inner_dim, bias=False, device=device, dtype=dtype)
+ self.v_norm = nn.Identity()
+
+ self.output_proj = operations.Linear(inner_dim, query_dim, bias=False, device=device, dtype=dtype)
+ self.output_dropout = nn.Dropout(dropout) if dropout > 1e-4 else nn.Identity()
+
+ self.attn_op = torch_attention_op
+
+ self._query_dim = query_dim
+ self._context_dim = context_dim
+ self._inner_dim = inner_dim
+
+ def compute_qkv(
+ self,
+ x: torch.Tensor,
+ context: Optional[torch.Tensor] = None,
+ rope_emb: Optional[torch.Tensor] = None,
+ ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+ q = self.q_proj(x)
+ context = x if context is None else context
+ k = self.k_proj(context)
+ v = self.v_proj(context)
+ q, k, v = map(
+ lambda t: rearrange(t, "b ... (h d) -> b ... h d", h=self.n_heads, d=self.head_dim),
+ (q, k, v),
+ )
+
+ def apply_norm_and_rotary_pos_emb(
+ q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, rope_emb: Optional[torch.Tensor]
+ ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+ q = self.q_norm(q)
+ k = self.k_norm(k)
+ v = self.v_norm(v)
+ if self.is_selfattn and rope_emb is not None: # only apply to self-attention!
+ q = apply_rotary_pos_emb(q, rope_emb)
+ k = apply_rotary_pos_emb(k, rope_emb)
+ return q, k, v
+
+ q, k, v = apply_norm_and_rotary_pos_emb(q, k, v, rope_emb)
+
+ return q, k, v
+
+ def compute_attention(self, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, transformer_options: Optional[dict] = {}) -> torch.Tensor:
+ result = self.attn_op(q, k, v, transformer_options=transformer_options) # [B, S, H, D]
+ return self.output_dropout(self.output_proj(result))
+
+ def forward(
+ self,
+ x: torch.Tensor,
+ context: Optional[torch.Tensor] = None,
+ rope_emb: Optional[torch.Tensor] = None,
+ transformer_options: Optional[dict] = {},
+ ) -> torch.Tensor:
+ """
+ Args:
+ x (Tensor): The query tensor of shape [B, Mq, K]
+ context (Optional[Tensor]): The key tensor of shape [B, Mk, K] or use x as context [self attention] if None
+ """
+ q, k, v = self.compute_qkv(x, context, rope_emb=rope_emb)
+ return self.compute_attention(q, k, v, transformer_options=transformer_options)
+
+
+class Timesteps(nn.Module):
+ def __init__(self, num_channels: int):
+ super().__init__()
+ self.num_channels = num_channels
+
+ def forward(self, timesteps_B_T: torch.Tensor) -> torch.Tensor:
+ assert timesteps_B_T.ndim == 2, f"Expected 2D input, got {timesteps_B_T.ndim}"
+ timesteps = timesteps_B_T.flatten().float()
+ half_dim = self.num_channels // 2
+ exponent = -math.log(10000) * torch.arange(half_dim, dtype=torch.float32, device=timesteps.device)
+ exponent = exponent / (half_dim - 0.0)
+
+ emb = torch.exp(exponent)
+ emb = timesteps[:, None].float() * emb[None, :]
+
+ sin_emb = torch.sin(emb)
+ cos_emb = torch.cos(emb)
+ emb = torch.cat([cos_emb, sin_emb], dim=-1)
+
+ return rearrange(emb, "(b t) d -> b t d", b=timesteps_B_T.shape[0], t=timesteps_B_T.shape[1])
+
+
+class TimestepEmbedding(nn.Module):
+ def __init__(self, in_features: int, out_features: int, use_adaln_lora: bool = False, device=None, dtype=None, operations=None):
+ super().__init__()
+ logging.debug(
+ f"Using AdaLN LoRA Flag: {use_adaln_lora}. We enable bias if no AdaLN LoRA for backward compatibility."
+ )
+ self.in_dim = in_features
+ self.out_dim = out_features
+ self.linear_1 = operations.Linear(in_features, out_features, bias=not use_adaln_lora, device=device, dtype=dtype)
+ self.activation = nn.SiLU()
+ self.use_adaln_lora = use_adaln_lora
+ if use_adaln_lora:
+ self.linear_2 = operations.Linear(out_features, 3 * out_features, bias=False, device=device, dtype=dtype)
+ else:
+ self.linear_2 = operations.Linear(out_features, out_features, bias=False, device=device, dtype=dtype)
+
+ def forward(self, sample: torch.Tensor) -> Tuple[torch.Tensor, Optional[torch.Tensor]]:
+ emb = self.linear_1(sample)
+ emb = self.activation(emb)
+ emb = self.linear_2(emb)
+
+ if self.use_adaln_lora:
+ adaln_lora_B_T_3D = emb
+ emb_B_T_D = sample
+ else:
+ adaln_lora_B_T_3D = None
+ emb_B_T_D = emb
+
+ return emb_B_T_D, adaln_lora_B_T_3D
+
+
+class PatchEmbed(nn.Module):
+ """
+ PatchEmbed is a module for embedding patches from an input tensor by applying either 3D or 2D convolutional layers,
+ depending on the . This module can process inputs with temporal (video) and spatial (image) dimensions,
+ making it suitable for video and image processing tasks. It supports dividing the input into patches
+ and embedding each patch into a vector of size `out_channels`.
+
+ Parameters:
+ - spatial_patch_size (int): The size of each spatial patch.
+ - temporal_patch_size (int): The size of each temporal patch.
+ - in_channels (int): Number of input channels. Default: 3.
+ - out_channels (int): The dimension of the embedding vector for each patch. Default: 768.
+ - bias (bool): If True, adds a learnable bias to the output of the convolutional layers. Default: True.
+ """
+
+ def __init__(
+ self,
+ spatial_patch_size: int,
+ temporal_patch_size: int,
+ in_channels: int = 3,
+ out_channels: int = 768,
+ device=None, dtype=None, operations=None
+ ):
+ super().__init__()
+ self.spatial_patch_size = spatial_patch_size
+ self.temporal_patch_size = temporal_patch_size
+
+ self.proj = nn.Sequential(
+ Rearrange(
+ "b c (t r) (h m) (w n) -> b t h w (c r m n)",
+ r=temporal_patch_size,
+ m=spatial_patch_size,
+ n=spatial_patch_size,
+ ),
+ operations.Linear(
+ in_channels * spatial_patch_size * spatial_patch_size * temporal_patch_size, out_channels, bias=False, device=device, dtype=dtype
+ ),
+ )
+ self.dim = in_channels * spatial_patch_size * spatial_patch_size * temporal_patch_size
+
+ def forward(self, x: torch.Tensor) -> torch.Tensor:
+ """
+ Forward pass of the PatchEmbed module.
+
+ Parameters:
+ - x (torch.Tensor): The input tensor of shape (B, C, T, H, W) where
+ B is the batch size,
+ C is the number of channels,
+ T is the temporal dimension,
+ H is the height, and
+ W is the width of the input.
+
+ Returns:
+ - torch.Tensor: The embedded patches as a tensor, with shape b t h w c.
+ """
+ assert x.dim() == 5
+ _, _, T, H, W = x.shape
+ assert (
+ H % self.spatial_patch_size == 0 and W % self.spatial_patch_size == 0
+ ), f"H,W {(H, W)} should be divisible by spatial_patch_size {self.spatial_patch_size}"
+ assert T % self.temporal_patch_size == 0
+ x = self.proj(x)
+ return x
+
+
+class FinalLayer(nn.Module):
+ """
+ The final layer of video DiT.
+ """
+
+ def __init__(
+ self,
+ hidden_size: int,
+ spatial_patch_size: int,
+ temporal_patch_size: int,
+ out_channels: int,
+ use_adaln_lora: bool = False,
+ adaln_lora_dim: int = 256,
+ device=None, dtype=None, operations=None
+ ):
+ super().__init__()
+ self.layer_norm = operations.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6)
+ self.linear = operations.Linear(
+ hidden_size, spatial_patch_size * spatial_patch_size * temporal_patch_size * out_channels, bias=False, device=device, dtype=dtype
+ )
+ self.hidden_size = hidden_size
+ self.n_adaln_chunks = 2
+ self.use_adaln_lora = use_adaln_lora
+ self.adaln_lora_dim = adaln_lora_dim
+ if use_adaln_lora:
+ self.adaln_modulation = nn.Sequential(
+ nn.SiLU(),
+ operations.Linear(hidden_size, adaln_lora_dim, bias=False, device=device, dtype=dtype),
+ operations.Linear(adaln_lora_dim, self.n_adaln_chunks * hidden_size, bias=False, device=device, dtype=dtype),
+ )
+ else:
+ self.adaln_modulation = nn.Sequential(
+ nn.SiLU(), operations.Linear(hidden_size, self.n_adaln_chunks * hidden_size, bias=False, device=device, dtype=dtype)
+ )
+
+ def forward(
+ self,
+ x_B_T_H_W_D: torch.Tensor,
+ emb_B_T_D: torch.Tensor,
+ adaln_lora_B_T_3D: Optional[torch.Tensor] = None,
+ ):
+ if self.use_adaln_lora:
+ assert adaln_lora_B_T_3D is not None
+ shift_B_T_D, scale_B_T_D = (
+ self.adaln_modulation(emb_B_T_D) + adaln_lora_B_T_3D[:, :, : 2 * self.hidden_size]
+ ).chunk(2, dim=-1)
+ else:
+ shift_B_T_D, scale_B_T_D = self.adaln_modulation(emb_B_T_D).chunk(2, dim=-1)
+
+ shift_B_T_1_1_D, scale_B_T_1_1_D = rearrange(shift_B_T_D, "b t d -> b t 1 1 d"), rearrange(
+ scale_B_T_D, "b t d -> b t 1 1 d"
+ )
+
+ def _fn(
+ _x_B_T_H_W_D: torch.Tensor,
+ _norm_layer: nn.Module,
+ _scale_B_T_1_1_D: torch.Tensor,
+ _shift_B_T_1_1_D: torch.Tensor,
+ ) -> torch.Tensor:
+ return _norm_layer(_x_B_T_H_W_D) * (1 + _scale_B_T_1_1_D) + _shift_B_T_1_1_D
+
+ x_B_T_H_W_D = _fn(x_B_T_H_W_D, self.layer_norm, scale_B_T_1_1_D, shift_B_T_1_1_D)
+ x_B_T_H_W_O = self.linear(x_B_T_H_W_D)
+ return x_B_T_H_W_O
+
+
+class Block(nn.Module):
+ """
+ A transformer block that combines self-attention, cross-attention and MLP layers with AdaLN modulation.
+ Each component (self-attention, cross-attention, MLP) has its own layer normalization and AdaLN modulation.
+
+ Parameters:
+ x_dim (int): Dimension of input features
+ context_dim (int): Dimension of context features for cross-attention
+ num_heads (int): Number of attention heads
+ mlp_ratio (float): Multiplier for MLP hidden dimension. Default: 4.0
+ use_adaln_lora (bool): Whether to use AdaLN-LoRA modulation. Default: False
+ adaln_lora_dim (int): Hidden dimension for AdaLN-LoRA layers. Default: 256
+
+ The block applies the following sequence:
+ 1. Self-attention with AdaLN modulation
+ 2. Cross-attention with AdaLN modulation
+ 3. MLP with AdaLN modulation
+
+ Each component uses skip connections and layer normalization.
+ """
+
+ def __init__(
+ self,
+ x_dim: int,
+ context_dim: int,
+ num_heads: int,
+ mlp_ratio: float = 4.0,
+ use_adaln_lora: bool = False,
+ adaln_lora_dim: int = 256,
+ device=None,
+ dtype=None,
+ operations=None,
+ ):
+ super().__init__()
+ self.x_dim = x_dim
+ self.layer_norm_self_attn = operations.LayerNorm(x_dim, elementwise_affine=False, eps=1e-6, device=device, dtype=dtype)
+ self.self_attn = Attention(x_dim, None, num_heads, x_dim // num_heads, device=device, dtype=dtype, operations=operations)
+
+ self.layer_norm_cross_attn = operations.LayerNorm(x_dim, elementwise_affine=False, eps=1e-6, device=device, dtype=dtype)
+ self.cross_attn = Attention(
+ x_dim, context_dim, num_heads, x_dim // num_heads, device=device, dtype=dtype, operations=operations
+ )
+
+ self.layer_norm_mlp = operations.LayerNorm(x_dim, elementwise_affine=False, eps=1e-6, device=device, dtype=dtype)
+ self.mlp = GPT2FeedForward(x_dim, int(x_dim * mlp_ratio), device=device, dtype=dtype, operations=operations)
+
+ self.use_adaln_lora = use_adaln_lora
+ if self.use_adaln_lora:
+ self.adaln_modulation_self_attn = nn.Sequential(
+ nn.SiLU(),
+ operations.Linear(x_dim, adaln_lora_dim, bias=False, device=device, dtype=dtype),
+ operations.Linear(adaln_lora_dim, 3 * x_dim, bias=False, device=device, dtype=dtype),
+ )
+ self.adaln_modulation_cross_attn = nn.Sequential(
+ nn.SiLU(),
+ operations.Linear(x_dim, adaln_lora_dim, bias=False, device=device, dtype=dtype),
+ operations.Linear(adaln_lora_dim, 3 * x_dim, bias=False, device=device, dtype=dtype),
+ )
+ self.adaln_modulation_mlp = nn.Sequential(
+ nn.SiLU(),
+ operations.Linear(x_dim, adaln_lora_dim, bias=False, device=device, dtype=dtype),
+ operations.Linear(adaln_lora_dim, 3 * x_dim, bias=False, device=device, dtype=dtype),
+ )
+ else:
+ self.adaln_modulation_self_attn = nn.Sequential(nn.SiLU(), operations.Linear(x_dim, 3 * x_dim, bias=False, device=device, dtype=dtype))
+ self.adaln_modulation_cross_attn = nn.Sequential(nn.SiLU(), operations.Linear(x_dim, 3 * x_dim, bias=False, device=device, dtype=dtype))
+ self.adaln_modulation_mlp = nn.Sequential(nn.SiLU(), operations.Linear(x_dim, 3 * x_dim, bias=False, device=device, dtype=dtype))
+
+ def forward(
+ self,
+ x_B_T_H_W_D: torch.Tensor,
+ emb_B_T_D: torch.Tensor,
+ crossattn_emb: torch.Tensor,
+ rope_emb_L_1_1_D: Optional[torch.Tensor] = None,
+ adaln_lora_B_T_3D: Optional[torch.Tensor] = None,
+ extra_per_block_pos_emb: Optional[torch.Tensor] = None,
+ transformer_options: Optional[dict] = {},
+ ) -> torch.Tensor:
+ residual_dtype = x_B_T_H_W_D.dtype
+ compute_dtype = emb_B_T_D.dtype
+ if extra_per_block_pos_emb is not None:
+ x_B_T_H_W_D = x_B_T_H_W_D + extra_per_block_pos_emb
+
+ if self.use_adaln_lora:
+ shift_self_attn_B_T_D, scale_self_attn_B_T_D, gate_self_attn_B_T_D = (
+ self.adaln_modulation_self_attn(emb_B_T_D) + adaln_lora_B_T_3D
+ ).chunk(3, dim=-1)
+ shift_cross_attn_B_T_D, scale_cross_attn_B_T_D, gate_cross_attn_B_T_D = (
+ self.adaln_modulation_cross_attn(emb_B_T_D) + adaln_lora_B_T_3D
+ ).chunk(3, dim=-1)
+ shift_mlp_B_T_D, scale_mlp_B_T_D, gate_mlp_B_T_D = (
+ self.adaln_modulation_mlp(emb_B_T_D) + adaln_lora_B_T_3D
+ ).chunk(3, dim=-1)
+ else:
+ shift_self_attn_B_T_D, scale_self_attn_B_T_D, gate_self_attn_B_T_D = self.adaln_modulation_self_attn(
+ emb_B_T_D
+ ).chunk(3, dim=-1)
+ shift_cross_attn_B_T_D, scale_cross_attn_B_T_D, gate_cross_attn_B_T_D = self.adaln_modulation_cross_attn(
+ emb_B_T_D
+ ).chunk(3, dim=-1)
+ shift_mlp_B_T_D, scale_mlp_B_T_D, gate_mlp_B_T_D = self.adaln_modulation_mlp(emb_B_T_D).chunk(3, dim=-1)
+
+ # Reshape tensors from (B, T, D) to (B, T, 1, 1, D) for broadcasting
+ shift_self_attn_B_T_1_1_D = rearrange(shift_self_attn_B_T_D, "b t d -> b t 1 1 d")
+ scale_self_attn_B_T_1_1_D = rearrange(scale_self_attn_B_T_D, "b t d -> b t 1 1 d")
+ gate_self_attn_B_T_1_1_D = rearrange(gate_self_attn_B_T_D, "b t d -> b t 1 1 d")
+
+ shift_cross_attn_B_T_1_1_D = rearrange(shift_cross_attn_B_T_D, "b t d -> b t 1 1 d")
+ scale_cross_attn_B_T_1_1_D = rearrange(scale_cross_attn_B_T_D, "b t d -> b t 1 1 d")
+ gate_cross_attn_B_T_1_1_D = rearrange(gate_cross_attn_B_T_D, "b t d -> b t 1 1 d")
+
+ shift_mlp_B_T_1_1_D = rearrange(shift_mlp_B_T_D, "b t d -> b t 1 1 d")
+ scale_mlp_B_T_1_1_D = rearrange(scale_mlp_B_T_D, "b t d -> b t 1 1 d")
+ gate_mlp_B_T_1_1_D = rearrange(gate_mlp_B_T_D, "b t d -> b t 1 1 d")
+
+ B, T, H, W, D = x_B_T_H_W_D.shape
+
+ def _fn(_x_B_T_H_W_D, _norm_layer, _scale_B_T_1_1_D, _shift_B_T_1_1_D):
+ return _norm_layer(_x_B_T_H_W_D) * (1 + _scale_B_T_1_1_D) + _shift_B_T_1_1_D
+
+ normalized_x_B_T_H_W_D = _fn(
+ x_B_T_H_W_D,
+ self.layer_norm_self_attn,
+ scale_self_attn_B_T_1_1_D,
+ shift_self_attn_B_T_1_1_D,
+ )
+ result_B_T_H_W_D = rearrange(
+ self.self_attn(
+ # normalized_x_B_T_HW_D,
+ rearrange(normalized_x_B_T_H_W_D.to(compute_dtype), "b t h w d -> b (t h w) d"),
+ None,
+ rope_emb=rope_emb_L_1_1_D,
+ transformer_options=transformer_options,
+ ),
+ "b (t h w) d -> b t h w d",
+ t=T,
+ h=H,
+ w=W,
+ )
+ x_B_T_H_W_D = x_B_T_H_W_D + gate_self_attn_B_T_1_1_D.to(residual_dtype) * result_B_T_H_W_D.to(residual_dtype)
+
+ def _x_fn(
+ _x_B_T_H_W_D: torch.Tensor,
+ layer_norm_cross_attn: Callable,
+ _scale_cross_attn_B_T_1_1_D: torch.Tensor,
+ _shift_cross_attn_B_T_1_1_D: torch.Tensor,
+ transformer_options: Optional[dict] = {},
+ ) -> torch.Tensor:
+ _normalized_x_B_T_H_W_D = _fn(
+ _x_B_T_H_W_D, layer_norm_cross_attn, _scale_cross_attn_B_T_1_1_D, _shift_cross_attn_B_T_1_1_D
+ )
+ _result_B_T_H_W_D = rearrange(
+ self.cross_attn(
+ rearrange(_normalized_x_B_T_H_W_D.to(compute_dtype), "b t h w d -> b (t h w) d"),
+ crossattn_emb,
+ rope_emb=rope_emb_L_1_1_D,
+ transformer_options=transformer_options,
+ ),
+ "b (t h w) d -> b t h w d",
+ t=T,
+ h=H,
+ w=W,
+ )
+ return _result_B_T_H_W_D
+
+ result_B_T_H_W_D = _x_fn(
+ x_B_T_H_W_D,
+ self.layer_norm_cross_attn,
+ scale_cross_attn_B_T_1_1_D,
+ shift_cross_attn_B_T_1_1_D,
+ transformer_options=transformer_options,
+ )
+ x_B_T_H_W_D = result_B_T_H_W_D.to(residual_dtype) * gate_cross_attn_B_T_1_1_D.to(residual_dtype) + x_B_T_H_W_D
+
+ normalized_x_B_T_H_W_D = _fn(
+ x_B_T_H_W_D,
+ self.layer_norm_mlp,
+ scale_mlp_B_T_1_1_D,
+ shift_mlp_B_T_1_1_D,
+ )
+ result_B_T_H_W_D = self.mlp(normalized_x_B_T_H_W_D.to(compute_dtype))
+ x_B_T_H_W_D = x_B_T_H_W_D + gate_mlp_B_T_1_1_D.to(residual_dtype) * result_B_T_H_W_D.to(residual_dtype)
+ return x_B_T_H_W_D
+
+
+class MiniTrainDIT(nn.Module):
+ """
+ A clean impl of DIT that can load and reproduce the training results of the original DIT model in~(cosmos 1)
+ A general implementation of adaln-modulated VIT-like~(DiT) transformer for video processing.
+
+ Args:
+ max_img_h (int): Maximum height of the input images.
+ max_img_w (int): Maximum width of the input images.
+ max_frames (int): Maximum number of frames in the video sequence.
+ in_channels (int): Number of input channels (e.g., RGB channels for color images).
+ out_channels (int): Number of output channels.
+ patch_spatial (tuple): Spatial resolution of patches for input processing.
+ patch_temporal (int): Temporal resolution of patches for input processing.
+ concat_padding_mask (bool): If True, includes a mask channel in the input to handle padding.
+ model_channels (int): Base number of channels used throughout the model.
+ num_blocks (int): Number of transformer blocks.
+ num_heads (int): Number of heads in the multi-head attention layers.
+ mlp_ratio (float): Expansion ratio for MLP blocks.
+ crossattn_emb_channels (int): Number of embedding channels for cross-attention.
+ pos_emb_cls (str): Type of positional embeddings.
+ pos_emb_learnable (bool): Whether positional embeddings are learnable.
+ pos_emb_interpolation (str): Method for interpolating positional embeddings.
+ min_fps (int): Minimum frames per second.
+ max_fps (int): Maximum frames per second.
+ use_adaln_lora (bool): Whether to use AdaLN-LoRA.
+ adaln_lora_dim (int): Dimension for AdaLN-LoRA.
+ rope_h_extrapolation_ratio (float): Height extrapolation ratio for RoPE.
+ rope_w_extrapolation_ratio (float): Width extrapolation ratio for RoPE.
+ rope_t_extrapolation_ratio (float): Temporal extrapolation ratio for RoPE.
+ extra_per_block_abs_pos_emb (bool): Whether to use extra per-block absolute positional embeddings.
+ extra_h_extrapolation_ratio (float): Height extrapolation ratio for extra embeddings.
+ extra_w_extrapolation_ratio (float): Width extrapolation ratio for extra embeddings.
+ extra_t_extrapolation_ratio (float): Temporal extrapolation ratio for extra embeddings.
+ """
+
+ def __init__(
+ self,
+ max_img_h: int,
+ max_img_w: int,
+ max_frames: int,
+ in_channels: int,
+ out_channels: int,
+ patch_spatial: int, # tuple,
+ patch_temporal: int,
+ concat_padding_mask: bool = True,
+ # attention settings
+ model_channels: int = 768,
+ num_blocks: int = 10,
+ num_heads: int = 16,
+ mlp_ratio: float = 4.0,
+ # cross attention settings
+ crossattn_emb_channels: int = 1024,
+ # positional embedding settings
+ pos_emb_cls: str = "sincos",
+ pos_emb_learnable: bool = False,
+ pos_emb_interpolation: str = "crop",
+ min_fps: int = 1,
+ max_fps: int = 30,
+ use_adaln_lora: bool = False,
+ adaln_lora_dim: int = 256,
+ rope_h_extrapolation_ratio: float = 1.0,
+ rope_w_extrapolation_ratio: float = 1.0,
+ rope_t_extrapolation_ratio: float = 1.0,
+ extra_per_block_abs_pos_emb: bool = False,
+ extra_h_extrapolation_ratio: float = 1.0,
+ extra_w_extrapolation_ratio: float = 1.0,
+ extra_t_extrapolation_ratio: float = 1.0,
+ rope_enable_fps_modulation: bool = True,
+ image_model=None,
+ device=None,
+ dtype=None,
+ operations=None,
+ ) -> None:
+ super().__init__()
+ self.dtype = dtype
+ self.max_img_h = max_img_h
+ self.max_img_w = max_img_w
+ self.max_frames = max_frames
+ self.in_channels = in_channels
+ self.out_channels = out_channels
+ self.patch_spatial = patch_spatial
+ self.patch_temporal = patch_temporal
+ self.num_heads = num_heads
+ self.num_blocks = num_blocks
+ self.model_channels = model_channels
+ self.concat_padding_mask = concat_padding_mask
+ # positional embedding settings
+ self.pos_emb_cls = pos_emb_cls
+ self.pos_emb_learnable = pos_emb_learnable
+ self.pos_emb_interpolation = pos_emb_interpolation
+ self.min_fps = min_fps
+ self.max_fps = max_fps
+ self.rope_h_extrapolation_ratio = rope_h_extrapolation_ratio
+ self.rope_w_extrapolation_ratio = rope_w_extrapolation_ratio
+ self.rope_t_extrapolation_ratio = rope_t_extrapolation_ratio
+ self.extra_per_block_abs_pos_emb = extra_per_block_abs_pos_emb
+ self.extra_h_extrapolation_ratio = extra_h_extrapolation_ratio
+ self.extra_w_extrapolation_ratio = extra_w_extrapolation_ratio
+ self.extra_t_extrapolation_ratio = extra_t_extrapolation_ratio
+ self.rope_enable_fps_modulation = rope_enable_fps_modulation
+
+ self.build_pos_embed(device=device, dtype=dtype)
+ self.use_adaln_lora = use_adaln_lora
+ self.adaln_lora_dim = adaln_lora_dim
+ self.t_embedder = nn.Sequential(
+ Timesteps(model_channels),
+ TimestepEmbedding(model_channels, model_channels, use_adaln_lora=use_adaln_lora, device=device, dtype=dtype, operations=operations,),
+ )
+
+ in_channels = in_channels + 1 if concat_padding_mask else in_channels
+ self.x_embedder = PatchEmbed(
+ spatial_patch_size=patch_spatial,
+ temporal_patch_size=patch_temporal,
+ in_channels=in_channels,
+ out_channels=model_channels,
+ device=device, dtype=dtype, operations=operations,
+ )
+
+ self.blocks = nn.ModuleList(
+ [
+ Block(
+ x_dim=model_channels,
+ context_dim=crossattn_emb_channels,
+ num_heads=num_heads,
+ mlp_ratio=mlp_ratio,
+ use_adaln_lora=use_adaln_lora,
+ adaln_lora_dim=adaln_lora_dim,
+ device=device, dtype=dtype, operations=operations,
+ )
+ for _ in range(num_blocks)
+ ]
+ )
+
+ self.final_layer = FinalLayer(
+ hidden_size=self.model_channels,
+ spatial_patch_size=self.patch_spatial,
+ temporal_patch_size=self.patch_temporal,
+ out_channels=self.out_channels,
+ use_adaln_lora=self.use_adaln_lora,
+ adaln_lora_dim=self.adaln_lora_dim,
+ device=device, dtype=dtype, operations=operations,
+ )
+
+ self.t_embedding_norm = operations.RMSNorm(model_channels, eps=1e-6, device=device, dtype=dtype)
+
+ def build_pos_embed(self, device=None, dtype=None) -> None:
+ if self.pos_emb_cls == "rope3d":
+ cls_type = VideoRopePosition3DEmb
+ else:
+ raise ValueError(f"Unknown pos_emb_cls {self.pos_emb_cls}")
+
+ logging.debug(f"Building positional embedding with {self.pos_emb_cls} class, impl {cls_type}")
+ kwargs = dict(
+ model_channels=self.model_channels,
+ len_h=self.max_img_h // self.patch_spatial,
+ len_w=self.max_img_w // self.patch_spatial,
+ len_t=self.max_frames // self.patch_temporal,
+ max_fps=self.max_fps,
+ min_fps=self.min_fps,
+ is_learnable=self.pos_emb_learnable,
+ interpolation=self.pos_emb_interpolation,
+ head_dim=self.model_channels // self.num_heads,
+ h_extrapolation_ratio=self.rope_h_extrapolation_ratio,
+ w_extrapolation_ratio=self.rope_w_extrapolation_ratio,
+ t_extrapolation_ratio=self.rope_t_extrapolation_ratio,
+ enable_fps_modulation=self.rope_enable_fps_modulation,
+ device=device,
+ )
+ self.pos_embedder = cls_type(
+ **kwargs, # type: ignore
+ )
+
+ if self.extra_per_block_abs_pos_emb:
+ kwargs["h_extrapolation_ratio"] = self.extra_h_extrapolation_ratio
+ kwargs["w_extrapolation_ratio"] = self.extra_w_extrapolation_ratio
+ kwargs["t_extrapolation_ratio"] = self.extra_t_extrapolation_ratio
+ kwargs["device"] = device
+ kwargs["dtype"] = dtype
+ self.extra_pos_embedder = LearnablePosEmbAxis(
+ **kwargs, # type: ignore
+ )
+
+ def prepare_embedded_sequence(
+ self,
+ x_B_C_T_H_W: torch.Tensor,
+ fps: Optional[torch.Tensor] = None,
+ padding_mask: Optional[torch.Tensor] = None,
+ ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[torch.Tensor]]:
+ """
+ Prepares an embedded sequence tensor by applying positional embeddings and handling padding masks.
+
+ Args:
+ x_B_C_T_H_W (torch.Tensor): video
+ fps (Optional[torch.Tensor]): Frames per second tensor to be used for positional embedding when required.
+ If None, a default value (`self.base_fps`) will be used.
+ padding_mask (Optional[torch.Tensor]): current it is not used
+
+ Returns:
+ Tuple[torch.Tensor, Optional[torch.Tensor]]:
+ - A tensor of shape (B, T, H, W, D) with the embedded sequence.
+ - An optional positional embedding tensor, returned only if the positional embedding class
+ (`self.pos_emb_cls`) includes 'rope'. Otherwise, None.
+
+ Notes:
+ - If `self.concat_padding_mask` is True, a padding mask channel is concatenated to the input tensor.
+ - The method of applying positional embeddings depends on the value of `self.pos_emb_cls`.
+ - If 'rope' is in `self.pos_emb_cls` (case insensitive), the positional embeddings are generated using
+ the `self.pos_embedder` with the shape [T, H, W].
+ - If "fps_aware" is in `self.pos_emb_cls`, the positional embeddings are generated using the
+ `self.pos_embedder` with the fps tensor.
+ - Otherwise, the positional embeddings are generated without considering fps.
+ """
+ if self.concat_padding_mask:
+ if padding_mask is None:
+ padding_mask = torch.zeros(x_B_C_T_H_W.shape[0], 1, x_B_C_T_H_W.shape[3], x_B_C_T_H_W.shape[4], dtype=x_B_C_T_H_W.dtype, device=x_B_C_T_H_W.device)
+ else:
+ padding_mask = transforms.functional.resize(
+ padding_mask, list(x_B_C_T_H_W.shape[-2:]), interpolation=transforms.InterpolationMode.NEAREST
+ )
+ x_B_C_T_H_W = torch.cat(
+ [x_B_C_T_H_W, padding_mask.unsqueeze(1).repeat(1, 1, x_B_C_T_H_W.shape[2], 1, 1)], dim=1
+ )
+ x_B_T_H_W_D = self.x_embedder(x_B_C_T_H_W)
+
+ if self.extra_per_block_abs_pos_emb:
+ extra_pos_emb = self.extra_pos_embedder(x_B_T_H_W_D, fps=fps, device=x_B_C_T_H_W.device, dtype=x_B_C_T_H_W.dtype)
+ else:
+ extra_pos_emb = None
+
+ if "rope" in self.pos_emb_cls.lower():
+ return x_B_T_H_W_D, self.pos_embedder(x_B_T_H_W_D, fps=fps, device=x_B_C_T_H_W.device), extra_pos_emb
+ x_B_T_H_W_D = x_B_T_H_W_D + self.pos_embedder(x_B_T_H_W_D, device=x_B_C_T_H_W.device) # [B, T, H, W, D]
+
+ return x_B_T_H_W_D, None, extra_pos_emb
+
+ def unpatchify(self, x_B_T_H_W_M: torch.Tensor) -> torch.Tensor:
+ x_B_C_Tt_Hp_Wp = rearrange(
+ x_B_T_H_W_M,
+ "B T H W (p1 p2 t C) -> B C (T t) (H p1) (W p2)",
+ p1=self.patch_spatial,
+ p2=self.patch_spatial,
+ t=self.patch_temporal,
+ )
+ return x_B_C_Tt_Hp_Wp
+
+ def pad_to_patch_size(self, img, patch_size=(2, 2), padding_mode="circular"):
+ if padding_mode == "circular" and (torch.jit.is_tracing() or torch.jit.is_scripting()):
+ padding_mode = "reflect"
+
+ pad = ()
+ for i in range(img.ndim - 2):
+ pad = (0, (patch_size[i] - img.shape[i + 2] % patch_size[i]) % patch_size[i]) + pad
+
+ return torch.nn.functional.pad(img, pad, mode=padding_mode)
+
+ def forward(
+ self,
+ x: torch.Tensor,
+ timesteps: torch.Tensor,
+ context: torch.Tensor,
+ fps: Optional[torch.Tensor] = None,
+ padding_mask: Optional[torch.Tensor] = None,
+ use_gradient_checkpointing=False,
+ use_gradient_checkpointing_offload=False,
+ **kwargs,
+ ):
+ orig_shape = list(x.shape)
+ x = self.pad_to_patch_size(x, (self.patch_temporal, self.patch_spatial, self.patch_spatial))
+ x_B_C_T_H_W = x
+ timesteps_B_T = timesteps
+ crossattn_emb = context
+ """
+ Args:
+ x: (B, C, T, H, W) tensor of spatial-temp inputs
+ timesteps: (B, ) tensor of timesteps
+ crossattn_emb: (B, N, D) tensor of cross-attention embeddings
+ """
+ x_B_T_H_W_D, rope_emb_L_1_1_D, extra_pos_emb_B_T_H_W_D_or_T_H_W_B_D = self.prepare_embedded_sequence(
+ x_B_C_T_H_W,
+ fps=fps,
+ padding_mask=padding_mask,
+ )
+
+ if timesteps_B_T.ndim == 1:
+ timesteps_B_T = timesteps_B_T.unsqueeze(1)
+ t_embedding_B_T_D, adaln_lora_B_T_3D = self.t_embedder[1](self.t_embedder[0](timesteps_B_T).to(x_B_T_H_W_D.dtype))
+ t_embedding_B_T_D = self.t_embedding_norm(t_embedding_B_T_D)
+
+ # for logging purpose
+ affline_scale_log_info = {}
+ affline_scale_log_info["t_embedding_B_T_D"] = t_embedding_B_T_D.detach()
+ self.affline_scale_log_info = affline_scale_log_info
+ self.affline_emb = t_embedding_B_T_D
+ self.crossattn_emb = crossattn_emb
+
+ if extra_pos_emb_B_T_H_W_D_or_T_H_W_B_D is not None:
+ assert (
+ x_B_T_H_W_D.shape == extra_pos_emb_B_T_H_W_D_or_T_H_W_B_D.shape
+ ), f"{x_B_T_H_W_D.shape} != {extra_pos_emb_B_T_H_W_D_or_T_H_W_B_D.shape}"
+
+ block_kwargs = {
+ "rope_emb_L_1_1_D": rope_emb_L_1_1_D.unsqueeze(1).unsqueeze(0),
+ "adaln_lora_B_T_3D": adaln_lora_B_T_3D,
+ "extra_per_block_pos_emb": extra_pos_emb_B_T_H_W_D_or_T_H_W_B_D,
+ "transformer_options": kwargs.get("transformer_options", {}),
+ }
+
+ # The residual stream for this model has large values. To make fp16 compute_dtype work, we keep the residual stream
+ # in fp32, but run attention and MLP modules in fp16.
+ # An alternate method that clamps fp16 values "works" in the sense that it makes coherent images, but there is noticeable
+ # quality degradation and visual artifacts.
+ if x_B_T_H_W_D.dtype == torch.float16:
+ x_B_T_H_W_D = x_B_T_H_W_D.float()
+
+ for block in self.blocks:
+ x_B_T_H_W_D = gradient_checkpoint_forward(
+ block,
+ use_gradient_checkpointing=use_gradient_checkpointing,
+ use_gradient_checkpointing_offload=use_gradient_checkpointing_offload,
+ x_B_T_H_W_D=x_B_T_H_W_D,
+ emb_B_T_D=t_embedding_B_T_D,
+ crossattn_emb=crossattn_emb,
+ **block_kwargs,
+ )
+
+ x_B_T_H_W_O = self.final_layer(x_B_T_H_W_D.to(crossattn_emb.dtype), t_embedding_B_T_D, adaln_lora_B_T_3D=adaln_lora_B_T_3D)
+ x_B_C_Tt_Hp_Wp = self.unpatchify(x_B_T_H_W_O)[:, :, :orig_shape[-3], :orig_shape[-2], :orig_shape[-1]]
+ return x_B_C_Tt_Hp_Wp
+
+
+def rotate_half(x):
+ x1 = x[..., : x.shape[-1] // 2]
+ x2 = x[..., x.shape[-1] // 2 :]
+ return torch.cat((-x2, x1), dim=-1)
+
+
+def apply_rotary_pos_emb2(x, cos, sin, unsqueeze_dim=1):
+ cos = cos.unsqueeze(unsqueeze_dim)
+ sin = sin.unsqueeze(unsqueeze_dim)
+ x_embed = (x * cos) + (rotate_half(x) * sin)
+ return x_embed
+
+
+class RotaryEmbedding(nn.Module):
+ def __init__(self, head_dim):
+ super().__init__()
+ self.rope_theta = 10000
+ inv_freq = 1.0 / (self.rope_theta ** (torch.arange(0, head_dim, 2, dtype=torch.int64).to(dtype=torch.float) / head_dim))
+ self.register_buffer("inv_freq", inv_freq, persistent=False)
+
+ @torch.no_grad()
+ def forward(self, x, position_ids):
+ inv_freq_expanded = self.inv_freq[None, :, None].float().expand(position_ids.shape[0], -1, 1).to(x.device)
+ position_ids_expanded = position_ids[:, None, :].float()
+
+ device_type = x.device.type if isinstance(x.device.type, str) and x.device.type != "mps" else "cpu"
+ with torch.autocast(device_type=device_type, enabled=False): # Force float32
+ freqs = (inv_freq_expanded.float() @ position_ids_expanded.float()).transpose(1, 2)
+ emb = torch.cat((freqs, freqs), dim=-1)
+ cos = emb.cos()
+ sin = emb.sin()
+
+ return cos.to(dtype=x.dtype), sin.to(dtype=x.dtype)
+
+
+class LLMAdapterAttention(nn.Module):
+ def __init__(self, query_dim, context_dim, n_heads, head_dim, device=None, dtype=None, operations=None):
+ super().__init__()
+
+ inner_dim = head_dim * n_heads
+ self.n_heads = n_heads
+ self.head_dim = head_dim
+ self.query_dim = query_dim
+ self.context_dim = context_dim
+
+ self.q_proj = operations.Linear(query_dim, inner_dim, bias=False, device=device, dtype=dtype)
+ self.q_norm = operations.RMSNorm(self.head_dim, eps=1e-6, device=device, dtype=dtype)
+
+ self.k_proj = operations.Linear(context_dim, inner_dim, bias=False, device=device, dtype=dtype)
+ self.k_norm = operations.RMSNorm(self.head_dim, eps=1e-6, device=device, dtype=dtype)
+
+ self.v_proj = operations.Linear(context_dim, inner_dim, bias=False, device=device, dtype=dtype)
+
+ self.o_proj = operations.Linear(inner_dim, query_dim, bias=False, device=device, dtype=dtype)
+
+ def forward(self, x, mask=None, context=None, position_embeddings=None, position_embeddings_context=None):
+ context = x if context is None else context
+ input_shape = x.shape[:-1]
+ q_shape = (*input_shape, self.n_heads, self.head_dim)
+ context_shape = context.shape[:-1]
+ kv_shape = (*context_shape, self.n_heads, self.head_dim)
+
+ query_states = self.q_norm(self.q_proj(x).view(q_shape)).transpose(1, 2)
+ key_states = self.k_norm(self.k_proj(context).view(kv_shape)).transpose(1, 2)
+ value_states = self.v_proj(context).view(kv_shape).transpose(1, 2)
+
+ if position_embeddings is not None:
+ assert position_embeddings_context is not None
+ cos, sin = position_embeddings
+ query_states = apply_rotary_pos_emb2(query_states, cos, sin)
+ cos, sin = position_embeddings_context
+ key_states = apply_rotary_pos_emb2(key_states, cos, sin)
+
+ attn_output = torch.nn.functional.scaled_dot_product_attention(query_states, key_states, value_states, attn_mask=mask)
+
+ attn_output = attn_output.transpose(1, 2).reshape(*input_shape, -1).contiguous()
+ attn_output = self.o_proj(attn_output)
+ return attn_output
+
+ def init_weights(self):
+ torch.nn.init.zeros_(self.o_proj.weight)
+
+
+class LLMAdapterTransformerBlock(nn.Module):
+ def __init__(self, source_dim, model_dim, num_heads=16, mlp_ratio=4.0, use_self_attn=False, layer_norm=False, device=None, dtype=None, operations=None):
+ super().__init__()
+ self.use_self_attn = use_self_attn
+
+ if self.use_self_attn:
+ self.norm_self_attn = operations.LayerNorm(model_dim, device=device, dtype=dtype) if layer_norm else operations.RMSNorm(model_dim, eps=1e-6, device=device, dtype=dtype)
+ self.self_attn = LLMAdapterAttention(
+ query_dim=model_dim,
+ context_dim=model_dim,
+ n_heads=num_heads,
+ head_dim=model_dim//num_heads,
+ device=device,
+ dtype=dtype,
+ operations=operations,
+ )
+
+ self.norm_cross_attn = operations.LayerNorm(model_dim, device=device, dtype=dtype) if layer_norm else operations.RMSNorm(model_dim, eps=1e-6, device=device, dtype=dtype)
+ self.cross_attn = LLMAdapterAttention(
+ query_dim=model_dim,
+ context_dim=source_dim,
+ n_heads=num_heads,
+ head_dim=model_dim//num_heads,
+ device=device,
+ dtype=dtype,
+ operations=operations,
+ )
+
+ self.norm_mlp = operations.LayerNorm(model_dim, device=device, dtype=dtype) if layer_norm else operations.RMSNorm(model_dim, eps=1e-6, device=device, dtype=dtype)
+ self.mlp = nn.Sequential(
+ operations.Linear(model_dim, int(model_dim * mlp_ratio), device=device, dtype=dtype),
+ nn.GELU(),
+ operations.Linear(int(model_dim * mlp_ratio), model_dim, device=device, dtype=dtype)
+ )
+
+ def forward(self, x, context, target_attention_mask=None, source_attention_mask=None, position_embeddings=None, position_embeddings_context=None):
+ if self.use_self_attn:
+ normed = self.norm_self_attn(x)
+ attn_out = self.self_attn(normed, mask=target_attention_mask, position_embeddings=position_embeddings, position_embeddings_context=position_embeddings)
+ x = x + attn_out
+
+ normed = self.norm_cross_attn(x)
+ attn_out = self.cross_attn(normed, mask=source_attention_mask, context=context, position_embeddings=position_embeddings, position_embeddings_context=position_embeddings_context)
+ x = x + attn_out
+
+ x = x + self.mlp(self.norm_mlp(x))
+ return x
+
+ def init_weights(self):
+ torch.nn.init.zeros_(self.mlp[2].weight)
+ self.cross_attn.init_weights()
+
+
+class LLMAdapter(nn.Module):
+ def __init__(
+ self,
+ source_dim=1024,
+ target_dim=1024,
+ model_dim=1024,
+ num_layers=6,
+ num_heads=16,
+ use_self_attn=True,
+ layer_norm=False,
+ device=None,
+ dtype=None,
+ operations=None,
+ ):
+ super().__init__()
+
+ self.embed = operations.Embedding(32128, target_dim, device=device, dtype=dtype)
+ if model_dim != target_dim:
+ self.in_proj = operations.Linear(target_dim, model_dim, device=device, dtype=dtype)
+ else:
+ self.in_proj = nn.Identity()
+ self.rotary_emb = RotaryEmbedding(model_dim//num_heads)
+ self.blocks = nn.ModuleList([
+ LLMAdapterTransformerBlock(source_dim, model_dim, num_heads=num_heads, use_self_attn=use_self_attn, layer_norm=layer_norm, device=device, dtype=dtype, operations=operations) for _ in range(num_layers)
+ ])
+ self.out_proj = operations.Linear(model_dim, target_dim, device=device, dtype=dtype)
+ self.norm = operations.RMSNorm(target_dim, eps=1e-6, device=device, dtype=dtype)
+
+ def forward(self, source_hidden_states, target_input_ids, target_attention_mask=None, source_attention_mask=None):
+ if target_attention_mask is not None:
+ target_attention_mask = target_attention_mask.to(torch.bool)
+ if target_attention_mask.ndim == 2:
+ target_attention_mask = target_attention_mask.unsqueeze(1).unsqueeze(1)
+
+ if source_attention_mask is not None:
+ source_attention_mask = source_attention_mask.to(torch.bool)
+ if source_attention_mask.ndim == 2:
+ source_attention_mask = source_attention_mask.unsqueeze(1).unsqueeze(1)
+
+ context = source_hidden_states
+ x = self.in_proj(self.embed(target_input_ids).to(context.dtype))
+ position_ids = torch.arange(x.shape[1], device=x.device).unsqueeze(0)
+ position_ids_context = torch.arange(context.shape[1], device=x.device).unsqueeze(0)
+ position_embeddings = self.rotary_emb(x, position_ids)
+ position_embeddings_context = self.rotary_emb(x, position_ids_context)
+ for block in self.blocks:
+ x = block(x, context, target_attention_mask=target_attention_mask, source_attention_mask=source_attention_mask, position_embeddings=position_embeddings, position_embeddings_context=position_embeddings_context)
+ return self.norm(self.out_proj(x))
+
+
+class AnimaDiT(MiniTrainDIT):
+
+ _repeated_blocks = ["Block"]
+
+ def __init__(self):
+ kwargs = {'image_model': 'anima', 'max_img_h': 240, 'max_img_w': 240, 'max_frames': 128, 'in_channels': 16, 'out_channels': 16, 'patch_spatial': 2, 'patch_temporal': 1, 'model_channels': 2048, 'concat_padding_mask': True, 'crossattn_emb_channels': 1024, 'pos_emb_cls': 'rope3d', 'pos_emb_learnable': True, 'pos_emb_interpolation': 'crop', 'min_fps': 1, 'max_fps': 30, 'use_adaln_lora': True, 'adaln_lora_dim': 256, 'num_blocks': 28, 'num_heads': 16, 'extra_per_block_abs_pos_emb': False, 'rope_h_extrapolation_ratio': 4.0, 'rope_w_extrapolation_ratio': 4.0, 'rope_t_extrapolation_ratio': 1.0, 'extra_h_extrapolation_ratio': 1.0, 'extra_w_extrapolation_ratio': 1.0, 'extra_t_extrapolation_ratio': 1.0, 'rope_enable_fps_modulation': False, 'dtype': torch.bfloat16, 'device': None, 'operations': torch.nn}
+ super().__init__(**kwargs)
+ self.llm_adapter = LLMAdapter(device=kwargs.get("device"), dtype=kwargs.get("dtype"), operations=kwargs.get("operations"))
+
+ def preprocess_text_embeds(self, text_embeds, text_ids, t5xxl_weights=None):
+ if text_ids is not None:
+ out = self.llm_adapter(text_embeds, text_ids)
+ if t5xxl_weights is not None:
+ out = out * t5xxl_weights
+
+ if out.shape[1] < 512:
+ out = torch.nn.functional.pad(out, (0, 0, 0, 512 - out.shape[1]))
+ return out
+ else:
+ return text_embeds
+
+ def forward(
+ self,
+ x, timesteps, context,
+ use_gradient_checkpointing=False,
+ use_gradient_checkpointing_offload=False,
+ **kwargs
+ ):
+ t5xxl_ids = kwargs.pop("t5xxl_ids", None)
+ if t5xxl_ids is not None:
+ context = self.preprocess_text_embeds(context, t5xxl_ids, t5xxl_weights=kwargs.pop("t5xxl_weights", None))
+ return super().forward(
+ x, timesteps, context,
+ use_gradient_checkpointing=use_gradient_checkpointing, use_gradient_checkpointing_offload=use_gradient_checkpointing_offload,
+ **kwargs
+ )
diff --git a/diffsynth/models/bioclip.py b/diffsynth/models/bioclip.py
new file mode 100644
index 0000000000000000000000000000000000000000..abcd383fba5d0733ed292af8e050cce524739517
--- /dev/null
+++ b/diffsynth/models/bioclip.py
@@ -0,0 +1,118 @@
+import torch
+from PIL import Image
+from torchvision import transforms
+from transformers import CLIPModel as HFCLIPModel, CLIPConfig
+
+
+class BioCLIPv2Model(HFCLIPModel):
+ def __init__(self):
+ super().__init__(self._build_config())
+
+ @staticmethod
+ def _build_config():
+ return CLIPConfig(
+ projection_dim=768,
+ logit_scale_init_value=2.6592,
+ text_config={
+ "hidden_size": 768,
+ "intermediate_size": 3072,
+ "num_attention_heads": 12,
+ "num_hidden_layers": 12,
+ "max_position_embeddings": 77,
+ "vocab_size": 49408,
+ "hidden_act": "gelu",
+ "layer_norm_eps": 1e-5,
+ "projection_dim": 768,
+ "bos_token_id": 0,
+ "eos_token_id": 2,
+ "pad_token_id": 1,
+ },
+ vision_config={
+ "hidden_size": 1024,
+ "intermediate_size": 4096,
+ "num_attention_heads": 16,
+ "num_hidden_layers": 24,
+ "image_size": 224,
+ "patch_size": 14,
+ "hidden_act": "gelu",
+ "layer_norm_eps": 1e-5,
+ "projection_dim": 768,
+ },
+ )
+
+
+class BioCLIPv2Compute(torch.nn.Module):
+ MEAN = (0.48145466, 0.4578275, 0.40821073)
+ STD = (0.26862954, 0.26130258, 0.27577711)
+
+ def __init__(self, model: BioCLIPv2Model, tokenizer, max_length: int = 77):
+ super().__init__()
+ self.model = model
+ self.tokenizer = tokenizer
+ self.max_length = max_length
+ self.image_transform = transforms.Compose([
+ transforms.Resize(224, interpolation=transforms.InterpolationMode.BICUBIC),
+ transforms.CenterCrop(224),
+ transforms.ToTensor(),
+ transforms.Normalize(self.MEAN, self.STD),
+ ])
+
+ @property
+ def device(self):
+ return next(self.model.parameters()).device
+
+ @property
+ def dtype(self):
+ return next(self.model.parameters()).dtype
+
+ def _preprocess_images(self, images):
+ if isinstance(images, Image.Image):
+ images = [images]
+ images = [img.convert("RGB") for img in images]
+ pixel_values = torch.stack([self.image_transform(img) for img in images])
+ return pixel_values.to(device=self.device, dtype=self.dtype)
+
+ def _tokenize(self, text):
+ if isinstance(text, str):
+ text = [text]
+ tokens = self.tokenizer(
+ text, padding=True, truncation=True,
+ max_length=self.max_length, return_tensors="pt",
+ )
+ return {k: v.to(self.device) for k, v in tokens.items()}
+
+ @torch.no_grad()
+ def get_image_features(self, images):
+ pixel_values = self._preprocess_images(images)
+ features = self.model.get_image_features(pixel_values=pixel_values)
+ return torch.nn.functional.normalize(features, dim=-1)
+
+ @torch.no_grad()
+ def get_text_features(self, text):
+ tokens = self._tokenize(text)
+ features = self.model.get_text_features(**tokens)
+ return torch.nn.functional.normalize(features, dim=-1)
+
+ @torch.no_grad()
+ def forward(self, text: str | list[str], images):
+ if isinstance(text, str):
+ text = [text]
+ if isinstance(images, Image.Image):
+ images = [images]
+ if len(text) == 1 and len(images) > 1:
+ text = text * len(images)
+ if len(images) == 1 and len(text) > 1:
+ images = images * len(text)
+ image_features = self.get_image_features(images)
+ text_features = self.get_text_features(text)
+ scores = (text_features * image_features).sum(dim=-1)
+ scores = self.model.logit_scale.exp() * scores
+ return scores
+
+ @torch.no_grad()
+ def similarity_matrix(self, text: str | list[str], images):
+ image_features = self.get_image_features(images)
+ text_features = self.get_text_features(text)
+ scores = text_features @ image_features.T
+ scores = self.model.logit_scale.exp() * scores
+ return scores
diff --git a/diffsynth/models/boogu_image_dit.py b/diffsynth/models/boogu_image_dit.py
new file mode 100644
index 0000000000000000000000000000000000000000..f1c01dea64d37720c3d4c079f7ccb14e34fd98bc
--- /dev/null
+++ b/diffsynth/models/boogu_image_dit.py
@@ -0,0 +1,3282 @@
+import inspect
+import itertools
+from typing import Any, Dict, List, Optional, Tuple, Union
+
+import numpy as np
+import torch
+import torch.nn as nn
+from einops import rearrange
+from torch.nn import RMSNorm
+import math
+import torch.nn.functional as F
+from einops import repeat
+from dataclasses import dataclass
+
+from ..core.gradient import gradient_checkpoint_forward
+
+
+class Attention(nn.Module):
+ def __init__(
+ self,
+ query_dim: int,
+ cross_attention_dim: Optional[int] = None,
+ heads: int = 8,
+ kv_heads: Optional[int] = None,
+ dim_head: int = 64,
+ dropout: float = 0.0,
+ bias: bool = False,
+ upcast_attention: bool = False,
+ upcast_softmax: bool = False,
+ cross_attention_norm: Optional[str] = None,
+ cross_attention_norm_num_groups: int = 32,
+ qk_norm: Optional[str] = None,
+ added_kv_proj_dim: Optional[int] = None,
+ added_proj_bias: Optional[bool] = True,
+ norm_num_groups: Optional[int] = None,
+ spatial_norm_dim: Optional[int] = None,
+ out_bias: bool = True,
+ scale_qk: bool = True,
+ only_cross_attention: bool = False,
+ eps: float = 1e-5,
+ rescale_output_factor: float = 1.0,
+ residual_connection: bool = False,
+ _from_deprecated_attn_block: bool = False,
+ processor=None,
+ out_dim: int = None,
+ out_context_dim: int = None,
+ context_pre_only=None,
+ pre_only=False,
+ elementwise_affine: bool = True,
+ is_causal: bool = False,
+ ):
+ super().__init__()
+
+ self.inner_dim = out_dim if out_dim is not None else dim_head * heads
+ self.inner_kv_dim = self.inner_dim if kv_heads is None else dim_head * kv_heads
+ self.query_dim = query_dim
+ self.use_bias = bias
+ self.is_cross_attention = cross_attention_dim is not None
+ self.cross_attention_dim = cross_attention_dim if cross_attention_dim is not None else query_dim
+ self.upcast_attention = upcast_attention
+ self.upcast_softmax = upcast_softmax
+ self.rescale_output_factor = rescale_output_factor
+ self.residual_connection = residual_connection
+ self.dropout = dropout
+ self.fused_projections = False
+ self.out_dim = out_dim if out_dim is not None else query_dim
+ self.out_context_dim = out_context_dim if out_context_dim is not None else query_dim
+ self.context_pre_only = context_pre_only
+ self.pre_only = pre_only
+ self.is_causal = is_causal
+
+ self._from_deprecated_attn_block = _from_deprecated_attn_block
+
+ self.scale_qk = scale_qk
+ self.scale = dim_head ** -0.5 if self.scale_qk else 1.0
+
+ self.heads = out_dim // dim_head if out_dim is not None else heads
+ self.sliceable_head_dim = heads
+
+ self.added_kv_proj_dim = added_kv_proj_dim
+ self.only_cross_attention = only_cross_attention
+
+ if self.added_kv_proj_dim is None and self.only_cross_attention:
+ raise ValueError(
+ "`only_cross_attention` can only be set to True if `added_kv_proj_dim` is not None."
+ )
+
+ if norm_num_groups is not None:
+ self.group_norm = nn.GroupNorm(num_channels=query_dim, num_groups=norm_num_groups, eps=eps, affine=True)
+ else:
+ self.group_norm = None
+
+ if spatial_norm_dim is not None:
+ self.spatial_norm = None
+ else:
+ self.spatial_norm = None
+
+ if qk_norm is None:
+ self.norm_q = None
+ self.norm_k = None
+ elif qk_norm == "layer_norm":
+ self.norm_q = nn.LayerNorm(dim_head, eps=eps, elementwise_affine=elementwise_affine)
+ self.norm_k = nn.LayerNorm(dim_head, eps=eps, elementwise_affine=elementwise_affine)
+ elif qk_norm == "rms_norm":
+ self.norm_q = RMSNorm(dim_head, eps=eps, elementwise_affine=elementwise_affine)
+ self.norm_k = RMSNorm(dim_head, eps=eps, elementwise_affine=elementwise_affine)
+ elif qk_norm == "rms_norm_across_heads":
+ self.norm_q = RMSNorm(dim_head * heads, eps=eps)
+ self.norm_k = RMSNorm(dim_head * kv_heads, eps=eps)
+ else:
+ raise ValueError(f"unknown qk_norm: {qk_norm}.")
+
+ if cross_attention_norm is None:
+ self.norm_cross = None
+ elif cross_attention_norm == "layer_norm":
+ self.norm_cross = nn.LayerNorm(self.cross_attention_dim)
+ elif cross_attention_norm == "group_norm":
+ if self.added_kv_proj_dim is not None:
+ norm_cross_num_channels = added_kv_proj_dim
+ else:
+ norm_cross_num_channels = self.cross_attention_dim
+ self.norm_cross = nn.GroupNorm(
+ num_channels=norm_cross_num_channels, num_groups=cross_attention_norm_num_groups, eps=1e-5, affine=True
+ )
+ else:
+ raise ValueError(f"unknown cross_attention_norm: {cross_attention_norm}.")
+
+ self.to_q = nn.Linear(query_dim, self.inner_dim, bias=bias)
+
+ if not self.only_cross_attention:
+ self.to_k = nn.Linear(self.cross_attention_dim, self.inner_kv_dim, bias=bias)
+ self.to_v = nn.Linear(self.cross_attention_dim, self.inner_kv_dim, bias=bias)
+ else:
+ self.to_k = None
+ self.to_v = None
+
+ self.added_proj_bias = added_proj_bias
+ if self.added_kv_proj_dim is not None:
+ self.add_k_proj = nn.Linear(added_kv_proj_dim, self.inner_kv_dim, bias=added_proj_bias)
+ self.add_v_proj = nn.Linear(added_kv_proj_dim, self.inner_kv_dim, bias=added_proj_bias)
+ if self.context_pre_only is not None:
+ self.add_q_proj = nn.Linear(added_kv_proj_dim, self.inner_dim, bias=added_proj_bias)
+ else:
+ self.add_q_proj = None
+ self.add_k_proj = None
+ self.add_v_proj = None
+
+ if not self.pre_only:
+ self.to_out = nn.ModuleList([])
+ self.to_out.append(nn.Linear(self.inner_dim, self.out_dim, bias=out_bias))
+ self.to_out.append(nn.Dropout(dropout))
+ else:
+ self.to_out = None
+
+ if self.context_pre_only is not None and not self.context_pre_only:
+ self.to_add_out = nn.Linear(self.inner_dim, self.out_context_dim, bias=out_bias)
+ else:
+ self.to_add_out = None
+
+ if qk_norm is not None and added_kv_proj_dim is not None:
+ if qk_norm == "layer_norm":
+ self.norm_added_q = nn.LayerNorm(dim_head, eps=eps, elementwise_affine=elementwise_affine)
+ self.norm_added_k = nn.LayerNorm(dim_head, eps=eps, elementwise_affine=elementwise_affine)
+ elif qk_norm == "rms_norm":
+ self.norm_added_q = RMSNorm(dim_head, eps=eps)
+ self.norm_added_k = RMSNorm(dim_head, eps=eps)
+ elif qk_norm == "rms_norm_across_heads":
+ self.norm_added_q = None
+ self.norm_added_k = RMSNorm(dim_head * kv_heads, eps=eps)
+ else:
+ raise ValueError(f"unknown qk_norm: {qk_norm}.")
+ else:
+ self.norm_added_q = None
+ self.norm_added_k = None
+
+ if processor is None:
+ processor = AttnProcessor2_0() if hasattr(F, "scaled_dot_product_attention") and self.scale_qk else AttnProcessor()
+ self.set_processor(processor)
+
+ def set_processor(self, processor) -> None:
+ if (
+ hasattr(self, "processor")
+ and isinstance(self.processor, torch.nn.Module)
+ and not isinstance(processor, torch.nn.Module)
+ ):
+ pass
+ self.processor = processor
+
+ def get_processor(self, return_deprecated_lora: bool = False):
+ if not return_deprecated_lora:
+ return self.processor
+
+ def forward(
+ self,
+ hidden_states: torch.Tensor,
+ encoder_hidden_states: Optional[torch.Tensor] = None,
+ attention_mask: Optional[torch.Tensor] = None,
+ **cross_attention_kwargs,
+ ) -> torch.Tensor:
+ attn_parameters = set(inspect.signature(self.processor.__call__).parameters.keys())
+ quiet_attn_parameters = {"ip_adapter_masks", "ip_hidden_states"}
+ unused_kwargs = [
+ k for k, _ in cross_attention_kwargs.items() if k not in attn_parameters and k not in quiet_attn_parameters
+ ]
+ cross_attention_kwargs = {k: w for k, w in cross_attention_kwargs.items() if k in attn_parameters}
+
+ return self.processor(
+ self,
+ hidden_states,
+ encoder_hidden_states=encoder_hidden_states,
+ attention_mask=attention_mask,
+ **cross_attention_kwargs,
+ )
+
+ def norm_encoder_hidden_states(self, encoder_hidden_states: torch.Tensor) -> torch.Tensor:
+ assert self.norm_cross is not None, "self.norm_cross must be defined to call self.norm_encoder_hidden_states"
+
+ if isinstance(self.norm_cross, nn.LayerNorm):
+ encoder_hidden_states = self.norm_cross(encoder_hidden_states)
+ elif isinstance(self.norm_cross, nn.GroupNorm):
+ encoder_hidden_states = encoder_hidden_states.transpose(1, 2)
+ encoder_hidden_states = self.norm_cross(encoder_hidden_states)
+ encoder_hidden_states = encoder_hidden_states.transpose(1, 2)
+ else:
+ assert False
+
+ return encoder_hidden_states
+
+
+class AttnProcessor:
+ def __call__(
+ self,
+ attn: Attention,
+ hidden_states: torch.Tensor,
+ encoder_hidden_states: Optional[torch.Tensor] = None,
+ attention_mask: Optional[torch.Tensor] = None,
+ ) -> torch.Tensor:
+ batch_size, sequence_length, _ = (
+ hidden_states.shape if encoder_hidden_states is None else encoder_hidden_states.shape
+ )
+
+ if attn.group_norm is not None:
+ hidden_states = attn.group_norm(hidden_states.transpose(1, 2)).transpose(1, 2)
+
+ query = attn.to_q(hidden_states)
+
+ if encoder_hidden_states is None:
+ encoder_hidden_states = hidden_states
+
+ if attn.norm_cross:
+ encoder_hidden_states = attn.norm_encoder_hidden_states(encoder_hidden_states)
+
+ key = attn.to_k(encoder_hidden_states)
+ value = attn.to_v(encoder_hidden_states)
+
+ query = attn.head_to_batch_dim(query)
+ key = attn.head_to_batch_dim(key)
+ value = attn.head_to_batch_dim(value)
+
+ attention_probs = attn.get_attention_scores(query, key, attention_mask)
+ hidden_states = torch.bmm(attention_probs, value)
+ hidden_states = attn.batch_to_head_dim(hidden_states)
+
+ if attn.to_out is not None:
+ hidden_states = attn.to_out[0](hidden_states)
+ hidden_states = attn.to_out[1](hidden_states)
+
+ return hidden_states
+
+
+class AttnProcessor2_0:
+ def __call__(
+ self,
+ attn: Attention,
+ hidden_states: torch.Tensor,
+ encoder_hidden_states: Optional[torch.Tensor] = None,
+ attention_mask: Optional[torch.Tensor] = None,
+ ) -> torch.Tensor:
+ batch_size, sequence_length, _ = (
+ hidden_states.shape if encoder_hidden_states is None else encoder_hidden_states.shape
+ )
+
+ if attn.group_norm is not None:
+ hidden_states = attn.group_norm(hidden_states.transpose(1, 2)).transpose(1, 2)
+
+ query = attn.to_q(hidden_states)
+
+ if encoder_hidden_states is None:
+ encoder_hidden_states = hidden_states
+
+ if attn.norm_cross:
+ encoder_hidden_states = attn.norm_encoder_hidden_states(encoder_hidden_states)
+
+ key = attn.to_k(encoder_hidden_states)
+ value = attn.to_v(encoder_hidden_states)
+
+ inner_dim = key.shape[-1]
+ head_dim = inner_dim // attn.heads
+
+ query = query.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
+ key = key.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
+ value = value.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
+
+ if attn.norm_q is not None:
+ query = attn.norm_q(query)
+ if attn.norm_k is not None:
+ key = attn.norm_k(key)
+
+ hidden_states = F.scaled_dot_product_attention(
+ query, key, value, attn_mask=attention_mask, dropout_p=0.0, is_causal=attn.is_causal, scale=attn.scale
+ )
+ hidden_states = hidden_states.transpose(1, 2).reshape(batch_size, -1, attn.heads * head_dim)
+ hidden_states = hidden_states.to(query.dtype)
+
+ if attn.to_out is not None:
+ hidden_states = attn.to_out[0](hidden_states)
+ hidden_states = attn.to_out[1](hidden_states)
+
+ return hidden_states
+
+
+def _get_taylor_cache_entry(
+ cache_dic: Dict, current: Dict, create: bool = False
+) -> Dict:
+ cache_root = cache_dic["cache"][-1]
+ stream = current["stream"]
+ layer = current["layer"]
+ module = current["module"]
+
+ if create:
+ return (
+ cache_root.setdefault(stream, {})
+ .setdefault(layer, {})
+ .setdefault(module, {})
+ )
+ return cache_root[stream][layer][module]
+
+
+def _tree_sub(lhs, rhs):
+ if isinstance(lhs, tuple):
+ return tuple(_tree_sub(x, y) for x, y in zip(lhs, rhs))
+ return lhs - rhs
+
+
+def _tree_div(value, divisor):
+ if isinstance(value, tuple):
+ return tuple(_tree_div(x, divisor) for x in value)
+ return value / divisor
+
+
+def _tree_add(lhs, rhs):
+ if lhs is None:
+ return rhs
+ if isinstance(lhs, tuple):
+ return tuple(_tree_add(x, y) for x, y in zip(lhs, rhs))
+ return lhs + rhs
+
+
+def _tree_mul(value, scalar):
+ if isinstance(value, tuple):
+ return tuple(_tree_mul(x, scalar) for x in value)
+ return value * scalar
+
+
+def derivative_approximation(cache_dic: Dict, current: Dict, feature: torch.Tensor):
+ """
+ Build/update Taylor coefficients from the latest feature tensor.
+
+ Args:
+ cache_dic: Global cache dict storing per-stream/layer/module states.
+ current: Current execution state with keys like `stream`, `layer`,
+ `module`, and `step`.
+ feature: Current feature tensor to use as 0-th order term.
+ """
+ difference_distance = (
+ current["activated_steps"][-1] - current["activated_steps"][-2]
+ )
+
+ cache_entry = _get_taylor_cache_entry(cache_dic, current, create=True)
+ updated_taylor_factors = {}
+ updated_taylor_factors[0] = feature
+
+ for i in range(cache_dic["max_order"]):
+ if (cache_entry.get(i, None) is not None) and (
+ current["step"] > cache_dic["first_enhance"] - 2
+ ):
+ updated_taylor_factors[i + 1] = (
+ updated_taylor_factors[i] - cache_entry[i]
+ ) / difference_distance
+ else:
+ break
+
+ cache_dic["cache"][-1][current["stream"]][current["layer"]][current["module"]] = (
+ updated_taylor_factors
+ )
+
+
+def derivative_approximation_4_double_stream(
+ cache_dic: Dict, current: Dict, feature: tuple
+):
+ """
+ Build/update Taylor coefficients for double-stream outputs.
+ """
+ difference_distance = (
+ current["activated_steps"][-1] - current["activated_steps"][-2]
+ )
+
+ cache_entry = _get_taylor_cache_entry(cache_dic, current, create=True)
+ updated_taylor_factors = {}
+ updated_taylor_factors[0] = feature
+
+ for i in range(cache_dic["max_order"]):
+ if (cache_entry.get(i, None) is not None) and (
+ current["step"] > cache_dic["first_enhance"] - 2
+ ):
+ updated_taylor_factors[i + 1] = _tree_div(
+ _tree_sub(updated_taylor_factors[i], cache_entry[i]),
+ difference_distance,
+ )
+ else:
+ break
+
+ cache_dic["cache"][-1][current["stream"]][current["layer"]][current["module"]] = (
+ updated_taylor_factors
+ )
+
+
+def taylor_formula(cache_dic: Dict, current: Dict) -> torch.Tensor:
+ """
+ Reconstruct feature estimate using cached Taylor coefficients.
+
+ Returns:
+ A tensor with the same shape as cached feature tensors for the
+ current stream/layer/module.
+ """
+ x = current["step"] - current["activated_steps"][-1]
+ output = 0
+ cache_entry = _get_taylor_cache_entry(cache_dic, current)
+
+ for i in range(len(cache_entry)):
+ output += (1 / math.factorial(i)) * cache_entry[i] * (x**i)
+
+ return output
+
+
+def taylor_formula_4_double_stream(cache_dic: Dict, current: Dict) -> tuple:
+ """
+ Reconstruct double-stream outputs using cached Taylor coefficients.
+ """
+ x = current["step"] - current["activated_steps"][-1]
+ output = None
+ cache_entry = _get_taylor_cache_entry(cache_dic, current)
+
+ for i in range(len(cache_entry)):
+ output = _tree_add(
+ output,
+ _tree_mul(cache_entry[i], (1 / math.factorial(i)) * (x**i)),
+ )
+
+ return output
+
+
+def taylor_cache_init(cache_dic: Dict, current: Dict):
+ """
+ Initialize Taylor storage for the first step/module access.
+
+ The target location is
+ `cache_dic['cache'][-1][stream][layer][module]`.
+ """
+ if (current["step"] == 0) and (cache_dic["taylor_cache"]):
+ cache_root = cache_dic["cache"][-1]
+ cache_root.setdefault(current["stream"], {}).setdefault(current["layer"], {})[
+ current["module"]
+ ] = {}
+
+
+class BooguImagePromptTuningRotaryPosEmbed(nn.Module):
+ """
+ Rotary Position Embedding for Prompt Tuning tokens.
+
+ This class generates rotary position embeddings specifically for prompt tuning tokens.
+ Since prompt tokens are treated as text tokens, we use text-style position encoding
+ with a fixed sequence length equal to num_trainable_prompt_tokens.
+
+ Args:
+ theta: Base frequency for rotary embeddings
+ axes_dim: Dimensions for each axis (tuple like (32, 32, 32))
+ num_trainable_prompt_tokens: Number of trainable prompt tokens
+ """
+
+ def __init__(self, theta: int, dim: int, num_trainable_prompt_tokens: int):
+ super().__init__()
+ self.theta = theta
+ self.num_trainable_prompt_tokens = num_trainable_prompt_tokens
+ # For text tokens, only use the first dimension (text/temporal dimension)
+ self.dim = dim # Extract text dimension from tuple
+
+ def forward(
+ self, batch_size: int, device: torch.device, use_causal_mask: bool = False
+ ) -> Tuple[torch.Tensor, torch.Tensor]:
+ """
+ Generate rotary position embeddings and attention mask for prompt tuning.
+
+ Args:
+ batch_size: Batch size
+ device: Target device for tensors
+ use_causal_mask: Whether to use causal attention mask
+
+ Returns:
+ Tuple of (rotary_embeddings, attention_mask)
+ - rotary_embeddings: [B, num_tokens, instruction_dim//2] - RoPE embeddings for prompt tokens (complex form)
+ - attention_mask: [B, num_tokens] or [B, num_tokens, num_tokens] - Attention mask
+ """
+ # Generate 1D rotary embeddings for text-style tokens
+ freqs_dtype = (
+ torch.float32 if torch.backends.mps.is_available() else torch.float64
+ )
+
+ # get_1d_rotary_pos_embed(dim, seq_len) returns [seq_len, dim//2]
+ # Because RoPE uses complex representation, each dimension is split into sin/cos pairs
+ text_freqs_cis = get_1d_rotary_pos_embed(
+ self.dim, # This should be 32 (text dimension)
+ self.num_trainable_prompt_tokens, # Sequence length
+ theta=self.theta,
+ freqs_dtype=freqs_dtype,
+ )
+
+ # For prompt tuning, we create simple sequential position embeddings
+ # Each prompt token gets a unique position ID: 0, 1, 2, ..., num_tokens-1
+ position_indices = torch.arange(
+ self.num_trainable_prompt_tokens,
+ dtype=torch.int64,
+ device=text_freqs_cis.device,
+ )
+
+ # Select the appropriate rotary embeddings for each position
+ # text_freqs_cis is [num_tokens, instruction_dim//2], we want [num_tokens, instruction_dim//2]
+ rotary_emb = text_freqs_cis[
+ position_indices
+ ] # [num_tokens, instruction_dim//2]
+
+ # Expand to batch size and move to target device
+ rotary_emb = (
+ rotary_emb.unsqueeze(0).expand(batch_size, -1, -1).to(device)
+ ) # [B, num_tokens, instruction_dim//2]
+
+ # Create attention mask based on use_causal_mask parameter
+ if use_causal_mask:
+ # Create causal mask: only future tokens can attend to past tokens
+ # Lower triangular matrix where mask[i, j] = True if i >= j
+ causal_mask = torch.tril(
+ torch.ones(
+ self.num_trainable_prompt_tokens,
+ self.num_trainable_prompt_tokens,
+ dtype=torch.bool,
+ device=device,
+ )
+ ) # [num_tokens, num_tokens]
+
+ # Expand to batch size [B, num_tokens, num_tokens]
+ attention_mask = causal_mask.unsqueeze(0).expand(batch_size, -1, -1)
+ else:
+ # Non-causal mask: all tokens can attend to each other (all True)
+ attention_mask = torch.ones(
+ batch_size,
+ self.num_trainable_prompt_tokens,
+ dtype=torch.bool,
+ device=device,
+ ) # [B, num_tokens]
+
+ return rotary_emb, attention_mask
+
+
+def get_1d_rotary_pos_embed(
+ dim: int,
+ pos: np.ndarray | int,
+ theta: float = 10000.0,
+ use_real=False,
+ linear_factor=1.0,
+ ntk_factor=1.0,
+ repeat_interleave_real=True,
+ freqs_dtype=torch.float32, # torch.float32, torch.float64 (flux)
+):
+ """
+ Precompute the frequency tensor for complex exponentials (cis) with given dimensions.
+
+ This function calculates a frequency tensor with complex exponentials using the given dimension 'dim' and the end
+ index 'end'. The 'theta' parameter scales the frequencies. The returned tensor contains complex values in complex64
+ data type.
+
+ Args:
+ dim (`int`): Dimension of the frequency tensor.
+ pos (`np.ndarray` or `int`): Position indices for the frequency tensor. [S] or scalar
+ theta (`float`, *optional*, defaults to 10000.0):
+ Scaling factor for frequency computation. Defaults to 10000.0.
+ use_real (`bool`, *optional*):
+ If True, return real part and imaginary part separately. Otherwise, return complex numbers.
+ linear_factor (`float`, *optional*, defaults to 1.0):
+ Scaling factor for the context extrapolation. Defaults to 1.0.
+ ntk_factor (`float`, *optional*, defaults to 1.0):
+ Scaling factor for the NTK-Aware RoPE. Defaults to 1.0.
+ repeat_interleave_real (`bool`, *optional*, defaults to `True`):
+ If `True` and `use_real`, real part and imaginary part are each interleaved with themselves to reach `dim`.
+ Otherwise, they are concateanted with themselves.
+ freqs_dtype (`torch.float32` or `torch.float64`, *optional*, defaults to `torch.float32`):
+ the dtype of the frequency tensor.
+ Returns:
+ `torch.Tensor`: Precomputed frequency tensor with complex exponentials. [S, D/2]
+ """
+ assert dim % 2 == 0
+
+ if isinstance(pos, int):
+ pos = torch.arange(pos)
+ if isinstance(pos, np.ndarray):
+ pos = torch.from_numpy(pos) # type: ignore # [S]
+
+ theta = theta * ntk_factor
+ freqs = (
+ 1.0 / (theta ** (torch.arange(0, dim, 2, dtype=freqs_dtype, device=pos.device) / dim)) / linear_factor
+ ) # [D/2]
+ freqs = torch.outer(pos, freqs) # type: ignore # [S, D/2]
+ is_npu = freqs.device.type == "npu"
+ if is_npu:
+ freqs = freqs.float()
+ if use_real and repeat_interleave_real:
+ # flux, hunyuan-dit, cogvideox
+ freqs_cos = freqs.cos().repeat_interleave(2, dim=1, output_size=freqs.shape[1] * 2).float() # [S, D]
+ freqs_sin = freqs.sin().repeat_interleave(2, dim=1, output_size=freqs.shape[1] * 2).float() # [S, D]
+ return freqs_cos, freqs_sin
+ elif use_real:
+ # stable audio, allegro
+ freqs_cos = torch.cat([freqs.cos(), freqs.cos()], dim=-1).float() # [S, D]
+ freqs_sin = torch.cat([freqs.sin(), freqs.sin()], dim=-1).float() # [S, D]
+ return freqs_cos, freqs_sin
+ else:
+ # lumina
+ freqs_cis = torch.polar(torch.ones_like(freqs), freqs) # complex64 # [S, D/2]
+ return freqs_cis
+
+
+class BooguImageDoubleStreamRotaryPosEmbed(nn.Module):
+ def __init__(
+ self,
+ theta: int,
+ axes_dim: Tuple[int, int, int],
+ axes_lens: Tuple[int, int, int] = (300, 512, 512),
+ patch_size: int = 2,
+ ):
+ super().__init__()
+ self.theta = theta
+ self.axes_dim = axes_dim
+ self.axes_lens = axes_lens
+ self.patch_size = patch_size
+
+ @staticmethod
+ def get_freqs_cis(
+ axes_dim: Tuple[int, int, int], axes_lens: Tuple[int, int, int], theta: int
+ ) -> List[torch.Tensor]:
+ freqs_cis = []
+ freqs_dtype = (
+ torch.float32 if torch.backends.mps.is_available() else torch.float64
+ )
+ for i, (d, e) in enumerate(zip(axes_dim, axes_lens)):
+ emb = get_1d_rotary_pos_embed(d, e, theta=theta, freqs_dtype=freqs_dtype)
+ freqs_cis.append(emb)
+ return freqs_cis
+
+ def _get_freqs_cis(self, freqs_cis, ids: torch.Tensor) -> torch.Tensor:
+ device = ids.device
+ if ids.device.type == "mps":
+ ids = ids.to("cpu")
+
+ result = []
+ for i in range(len(self.axes_dim)):
+ freqs = freqs_cis[i].to(ids.device)
+ index = ids[:, :, i : i + 1].repeat(1, 1, freqs.shape[-1]).to(torch.int64)
+ result.append(
+ torch.gather(
+ freqs.unsqueeze(0).repeat(index.shape[0], 1, 1), dim=1, index=index
+ )
+ )
+ return torch.cat(result, dim=-1).to(device)
+
+ def forward(
+ self,
+ freqs_cis,
+ attention_mask,
+ l_effective_ref_img_len,
+ l_effective_img_len,
+ ref_img_sizes,
+ img_sizes,
+ device,
+ ):
+ batch_size = len(attention_mask)
+ p = self.patch_size
+
+ encoder_seq_len = attention_mask.shape[1]
+ l_effective_cap_len = attention_mask.sum(dim=1).tolist()
+
+ seq_lengths = [
+ cap_len + sum(ref_img_len) + img_len
+ for cap_len, ref_img_len, img_len in zip(
+ l_effective_cap_len, l_effective_ref_img_len, l_effective_img_len
+ )
+ ]
+
+ max_seq_len = max(seq_lengths)
+ max_ref_img_len = max(
+ [sum(ref_img_len) for ref_img_len in l_effective_ref_img_len]
+ )
+ max_img_len = max(l_effective_img_len)
+
+ # Create position IDs
+ position_ids = torch.zeros(
+ batch_size, max_seq_len, 3, dtype=torch.int32, device=device
+ )
+
+ for i, (cap_seq_len, seq_len) in enumerate(
+ zip(l_effective_cap_len, seq_lengths)
+ ):
+ # add text position ids
+ position_ids[i, :cap_seq_len] = repeat(
+ torch.arange(cap_seq_len, dtype=torch.int32, device=device), "l -> l 3"
+ )
+
+ pe_shift = cap_seq_len
+ pe_shift_len = cap_seq_len
+
+ if ref_img_sizes[i] is not None:
+ for ref_img_size, ref_img_len in zip(
+ ref_img_sizes[i], l_effective_ref_img_len[i]
+ ):
+ H, W = ref_img_size
+ ref_H_tokens, ref_W_tokens = H // p, W // p
+ assert ref_H_tokens * ref_W_tokens == ref_img_len
+ # add image position ids
+
+ row_ids = repeat(
+ torch.arange(ref_H_tokens, dtype=torch.int32, device=device),
+ "h -> h w",
+ w=ref_W_tokens,
+ ).flatten()
+ col_ids = repeat(
+ torch.arange(ref_W_tokens, dtype=torch.int32, device=device),
+ "w -> h w",
+ h=ref_H_tokens,
+ ).flatten()
+ position_ids[i, pe_shift_len : pe_shift_len + ref_img_len, 0] = (
+ pe_shift
+ )
+ position_ids[i, pe_shift_len : pe_shift_len + ref_img_len, 1] = (
+ row_ids
+ )
+ position_ids[i, pe_shift_len : pe_shift_len + ref_img_len, 2] = (
+ col_ids
+ )
+
+ pe_shift += max(ref_H_tokens, ref_W_tokens)
+ pe_shift_len += ref_img_len
+
+ H, W = img_sizes[i]
+ H_tokens, W_tokens = H // p, W // p
+ assert H_tokens * W_tokens == l_effective_img_len[i]
+
+ row_ids = repeat(
+ torch.arange(H_tokens, dtype=torch.int32, device=device),
+ "h -> h w",
+ w=W_tokens,
+ ).flatten()
+ col_ids = repeat(
+ torch.arange(W_tokens, dtype=torch.int32, device=device),
+ "w -> h w",
+ h=H_tokens,
+ ).flatten()
+
+ assert pe_shift_len + l_effective_img_len[i] == seq_len
+ position_ids[i, pe_shift_len:seq_len, 0] = pe_shift
+ position_ids[i, pe_shift_len:seq_len, 1] = row_ids
+ position_ids[i, pe_shift_len:seq_len, 2] = col_ids
+
+ # Get combined rotary embeddings
+ freqs_cis = self._get_freqs_cis(freqs_cis, position_ids)
+
+ # create separate rotary embeddings for captions and images
+ cap_freqs_cis = torch.zeros(
+ batch_size,
+ encoder_seq_len,
+ freqs_cis.shape[-1],
+ device=device,
+ dtype=freqs_cis.dtype,
+ )
+ ref_img_freqs_cis = torch.zeros(
+ batch_size,
+ max_ref_img_len,
+ freqs_cis.shape[-1],
+ device=device,
+ dtype=freqs_cis.dtype,
+ )
+ img_freqs_cis = torch.zeros(
+ batch_size,
+ max_img_len,
+ freqs_cis.shape[-1],
+ device=device,
+ dtype=freqs_cis.dtype,
+ )
+
+ # Calculate combined image sequence lengths (ref_img + img) for each sample
+ combined_img_seq_lengths = [
+ sum(ref_img_len) + img_len
+ for ref_img_len, img_len in zip(
+ l_effective_ref_img_len, l_effective_img_len
+ )
+ ]
+ max_combined_img_len = max(combined_img_seq_lengths)
+
+ # Create combined image rotary embeddings
+ combined_img_freqs_cis = torch.zeros(
+ batch_size,
+ max_combined_img_len,
+ freqs_cis.shape[-1],
+ device=device,
+ dtype=freqs_cis.dtype,
+ )
+
+ for i, (cap_seq_len, ref_img_len, img_len, seq_len) in enumerate(
+ zip(
+ l_effective_cap_len,
+ l_effective_ref_img_len,
+ l_effective_img_len,
+ seq_lengths,
+ )
+ ):
+ cap_freqs_cis[i, :cap_seq_len] = freqs_cis[i, :cap_seq_len]
+ ref_img_freqs_cis[i, : sum(ref_img_len)] = freqs_cis[
+ i, cap_seq_len : cap_seq_len + sum(ref_img_len)
+ ]
+ img_freqs_cis[i, :img_len] = freqs_cis[
+ i,
+ cap_seq_len + sum(ref_img_len) : cap_seq_len
+ + sum(ref_img_len)
+ + img_len,
+ ]
+
+ # Combined image rotary embeddings: ref_img + img (same order as img_patch_embed_and_refine)
+ combined_img_freqs_cis[i, : sum(ref_img_len)] = freqs_cis[
+ i, cap_seq_len : cap_seq_len + sum(ref_img_len)
+ ]
+ combined_img_freqs_cis[i, sum(ref_img_len) : sum(ref_img_len) + img_len] = (
+ freqs_cis[
+ i,
+ cap_seq_len + sum(ref_img_len) : cap_seq_len
+ + sum(ref_img_len)
+ + img_len,
+ ]
+ )
+
+ return (
+ cap_freqs_cis,
+ ref_img_freqs_cis,
+ img_freqs_cis,
+ freqs_cis,
+ l_effective_cap_len,
+ seq_lengths,
+ combined_img_freqs_cis,
+ combined_img_seq_lengths,
+ )
+
+
+class LuminaRMSNormZero(nn.Module):
+ """
+ Norm layer adaptive RMS normalization zero.
+
+ Parameters:
+ embedding_dim (`int`): The size of each embedding vector.
+ """
+
+ def __init__(
+ self,
+ embedding_dim: int,
+ norm_eps: float,
+ norm_elementwise_affine: bool,
+ ):
+ super().__init__()
+ self.silu = nn.SiLU()
+ self.linear = nn.Linear(
+ min(embedding_dim, 1024),
+ 4 * embedding_dim,
+ bias=True,
+ )
+
+ self.norm = RMSNorm(embedding_dim, eps=norm_eps)
+
+ def forward(
+ self,
+ x: torch.Tensor,
+ emb: Optional[torch.Tensor] = None,
+ ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
+ emb = self.linear(self.silu(emb))
+ scale_msa, gate_msa, scale_mlp, gate_mlp = emb.chunk(4, dim=1)
+ x = self.norm(x) * (1 + scale_msa[:, None])
+ return x, gate_msa, scale_mlp, gate_mlp
+
+
+class LuminaLayerNormContinuous(nn.Module):
+ def __init__(
+ self,
+ embedding_dim: int,
+ conditioning_embedding_dim: int,
+ # NOTE: It is a bit weird that the norm layer can be configured to have scale and shift parameters
+ # because the output is immediately scaled and shifted by the projected conditioning embeddings.
+ # Note that AdaLayerNorm does not let the norm layer have scale and shift parameters.
+ # However, this is how it was implemented in the original code, and it's rather likely you should
+ # set `elementwise_affine` to False.
+ elementwise_affine=True,
+ eps=1e-5,
+ bias=True,
+ norm_type="layer_norm",
+ out_dim: Optional[int] = None,
+ ):
+ super().__init__()
+
+ # AdaLN
+ self.silu = nn.SiLU()
+ self.linear_1 = nn.Linear(conditioning_embedding_dim, embedding_dim, bias=bias)
+
+ if norm_type == "layer_norm":
+ self.norm = nn.LayerNorm(embedding_dim, eps, elementwise_affine, bias)
+ elif norm_type == "rms_norm":
+ self.norm = RMSNorm(
+ embedding_dim, eps=eps, elementwise_affine=elementwise_affine
+ )
+ else:
+ raise ValueError(f"unknown norm_type {norm_type}")
+
+ self.linear_2 = None
+ if out_dim is not None:
+ self.linear_2 = nn.Linear(embedding_dim, out_dim, bias=bias)
+
+ def forward(
+ self,
+ x: torch.Tensor,
+ conditioning_embedding: torch.Tensor,
+ ) -> torch.Tensor:
+ # convert back to the original dtype in case `conditioning_embedding`` is upcasted to float32 (needed for hunyuanDiT)
+ emb = self.linear_1(self.silu(conditioning_embedding).to(x.dtype))
+ scale = emb
+ x = self.norm(x) * (1 + scale)[:, None, :]
+
+ if self.linear_2 is not None:
+ x = self.linear_2(x)
+
+ return x
+
+
+def swiglu(x, y):
+ return F.silu(x.float(), inplace=False).to(x.dtype) * y
+
+
+class LuminaFeedForward(nn.Module):
+ r"""
+ A feed-forward layer.
+
+ Parameters:
+ hidden_size (`int`):
+ The dimensionality of the hidden layers in the model. This parameter determines the width of the model's
+ hidden representations.
+ intermediate_size (`int`): The intermediate dimension of the feedforward layer.
+ multiple_of (`int`, *optional*): Value to ensure hidden dimension is a multiple
+ of this value.
+ ffn_dim_multiplier (float, *optional*): Custom multiplier for hidden
+ dimension. Defaults to None.
+ """
+
+ def __init__(
+ self,
+ dim: int,
+ inner_dim: int,
+ multiple_of: Optional[int] = 256,
+ ffn_dim_multiplier: Optional[float] = None,
+ ):
+ super().__init__()
+ self.swiglu = swiglu
+
+ # custom hidden_size factor multiplier
+ if ffn_dim_multiplier is not None:
+ inner_dim = int(ffn_dim_multiplier * inner_dim)
+ inner_dim = multiple_of * ((inner_dim + multiple_of - 1) // multiple_of)
+
+ self.linear_1 = nn.Linear(
+ dim,
+ inner_dim,
+ bias=False,
+ )
+ self.linear_2 = nn.Linear(
+ inner_dim,
+ dim,
+ bias=False,
+ )
+ self.linear_3 = nn.Linear(
+ dim,
+ inner_dim,
+ bias=False,
+ )
+
+ def forward(self, x):
+ h1, h2 = self.linear_1(x), self.linear_3(x)
+ swiglu_fn = self.swiglu
+ return self.linear_2(swiglu_fn(h1, h2))
+
+
+def get_timestep_embedding(
+ timesteps: torch.Tensor,
+ embedding_dim: int,
+ flip_sin_to_cos: bool = False,
+ downscale_freq_shift: float = 1,
+ scale: float = 1,
+ max_period: int = 10000,
+) -> torch.Tensor:
+ """
+ This matches the implementation in Denoising Diffusion Probabilistic Models: Create sinusoidal timestep embeddings.
+
+ Args
+ timesteps (torch.Tensor):
+ a 1-D Tensor of N indices, one per batch element. These may be fractional.
+ embedding_dim (int):
+ the dimension of the output.
+ flip_sin_to_cos (bool):
+ Whether the embedding order should be `cos, sin` (if True) or `sin, cos` (if False)
+ downscale_freq_shift (float):
+ Controls the delta between frequencies between dimensions
+ scale (float):
+ Scaling factor applied to the embeddings.
+ max_period (int):
+ Controls the maximum frequency of the embeddings
+ Returns
+ torch.Tensor: an [N x dim] Tensor of positional embeddings.
+ """
+ assert len(timesteps.shape) == 1, "Timesteps should be a 1d-array"
+
+ half_dim = embedding_dim // 2
+ exponent = -math.log(max_period) * torch.arange(
+ start=0, end=half_dim, dtype=torch.float32, device=timesteps.device
+ )
+ exponent = exponent / (half_dim - downscale_freq_shift)
+
+ emb = torch.exp(exponent)
+ emb = timesteps[:, None].float() * emb[None, :]
+
+ # scale embeddings
+ emb = scale * emb
+
+ # concat sine and cosine embeddings
+ emb = torch.cat([torch.sin(emb), torch.cos(emb)], dim=-1)
+
+ # flip sine and cosine embeddings
+ if flip_sin_to_cos:
+ emb = torch.cat([emb[:, half_dim:], emb[:, :half_dim]], dim=-1)
+
+ # zero pad
+ if embedding_dim % 2 == 1:
+ emb = torch.nn.functional.pad(emb, (0, 1, 0, 0))
+ return emb
+
+
+class Timesteps(nn.Module):
+ def __init__(self, num_channels: int, flip_sin_to_cos: bool, downscale_freq_shift: float, scale: int = 1):
+ super().__init__()
+ self.num_channels = num_channels
+ self.flip_sin_to_cos = flip_sin_to_cos
+ self.downscale_freq_shift = downscale_freq_shift
+ self.scale = scale
+
+ def forward(self, timesteps: torch.Tensor) -> torch.Tensor:
+ t_emb = get_timestep_embedding(
+ timesteps,
+ self.num_channels,
+ flip_sin_to_cos=self.flip_sin_to_cos,
+ downscale_freq_shift=self.downscale_freq_shift,
+ scale=self.scale,
+ )
+ return t_emb
+
+
+class TimestepEmbedding(nn.Module):
+ def __init__(
+ self,
+ in_channels: int,
+ time_embed_dim: int,
+ act_fn: str = "silu",
+ out_dim: int = None,
+ post_act_fn: Optional[str] = None,
+ cond_proj_dim=None,
+ sample_proj_bias=True,
+ ):
+ super().__init__()
+
+ self.linear_1 = nn.Linear(in_channels, time_embed_dim, sample_proj_bias)
+
+ if cond_proj_dim is not None:
+ self.cond_proj = nn.Linear(cond_proj_dim, in_channels, bias=False)
+ else:
+ self.cond_proj = None
+
+ self.act = torch.nn.SiLU()
+
+ if out_dim is not None:
+ time_embed_dim_out = out_dim
+ else:
+ time_embed_dim_out = time_embed_dim
+ self.linear_2 = nn.Linear(time_embed_dim, time_embed_dim_out, sample_proj_bias)
+
+ if post_act_fn is None:
+ self.post_act = None
+ else:
+ self.post_act = torch.nn.SiLU()
+
+ self.initialize_weights()
+
+ def initialize_weights(self):
+ nn.init.normal_(self.linear_1.weight, std=0.02)
+ nn.init.zeros_(self.linear_1.bias)
+ nn.init.normal_(self.linear_2.weight, std=0.02)
+ nn.init.zeros_(self.linear_2.bias)
+
+ def forward(self, sample, condition=None):
+ if condition is not None:
+ sample = sample + self.cond_proj(condition)
+ sample = self.linear_1(sample)
+
+ if self.act is not None:
+ sample = self.act(sample)
+
+ sample = self.linear_2(sample)
+
+ if self.post_act is not None:
+ sample = self.post_act(sample)
+ return sample
+
+
+class Lumina2CombinedTimestepCaptionEmbedding(nn.Module):
+ def __init__(
+ self,
+ hidden_size: int = 4096,
+ instruction_feat_dim: int = 2048,
+ frequency_embedding_size: int = 256,
+ norm_eps: float = 1e-5,
+ timestep_scale: float = 1.0,
+ ) -> None:
+ super().__init__()
+
+ self.time_proj = Timesteps(
+ num_channels=frequency_embedding_size,
+ flip_sin_to_cos=True,
+ downscale_freq_shift=0.0,
+ scale=timestep_scale,
+ )
+
+ self.timestep_embedder = TimestepEmbedding(
+ in_channels=frequency_embedding_size, time_embed_dim=min(hidden_size, 1024)
+ )
+
+ self.caption_embedder = nn.Sequential(
+ RMSNorm(instruction_feat_dim, eps=norm_eps),
+ nn.Linear(instruction_feat_dim, hidden_size, bias=True),
+ )
+
+ self._initialize_weights()
+
+ def _initialize_weights(self):
+ nn.init.trunc_normal_(self.caption_embedder[1].weight, std=0.02)
+ nn.init.zeros_(self.caption_embedder[1].bias)
+
+ def forward(
+ self,
+ timestep: torch.Tensor,
+ instruction_hidden_states: torch.Tensor,
+ dtype: torch.dtype,
+ ) -> Tuple[torch.Tensor, torch.Tensor]:
+ timestep_proj = self.time_proj(timestep).to(dtype=dtype)
+ time_embed = self.timestep_embedder(timestep_proj)
+ caption_embed = self.caption_embedder(instruction_hidden_states)
+ return time_embed, caption_embed
+
+
+def apply_rotary_emb(
+ x: torch.Tensor,
+ freqs_cis: Union[torch.Tensor, Tuple[torch.Tensor]],
+ use_real: bool = True,
+ use_real_unbind_dim: int = -1,
+) -> Tuple[torch.Tensor, torch.Tensor]:
+ """
+ Apply rotary embeddings to input tensors using the given frequency tensor. This function applies rotary embeddings
+ to the given query or key 'x' tensors using the provided frequency tensor 'freqs_cis'. The input tensors are
+ reshaped as complex numbers, and the frequency tensor is reshaped for broadcasting compatibility. The resulting
+ tensors contain rotary embeddings and are returned as real tensors.
+
+ Args:
+ x (`torch.Tensor`):
+ Query or key tensor to apply rotary embeddings. [B, H, S, D] xk (torch.Tensor): Key tensor to apply
+ freqs_cis (`Tuple[torch.Tensor]`): Precomputed frequency tensor for complex exponentials. ([S, D], [S, D],)
+
+ Returns:
+ Tuple[torch.Tensor, torch.Tensor]: Tuple of modified query tensor and key tensor with rotary embeddings.
+ """
+ if use_real:
+ cos, sin = freqs_cis # [S, D]
+ cos = cos[None, None]
+ sin = sin[None, None]
+ cos, sin = cos.to(x.device), sin.to(x.device)
+
+ if use_real_unbind_dim == -1:
+ # Used for flux, cogvideox, hunyuan-dit
+ x_real, x_imag = x.reshape(*x.shape[:-1], -1, 2).unbind(
+ -1
+ ) # [B, S, H, D//2]
+ x_rotated = torch.stack([-x_imag, x_real], dim=-1).flatten(3)
+ elif use_real_unbind_dim == -2:
+ # Used for Stable Audio, Boogu and CogView4
+ x_real, x_imag = x.reshape(*x.shape[:-1], 2, -1).unbind(
+ -2
+ ) # [B, S, H, D//2]
+ x_rotated = torch.cat([-x_imag, x_real], dim=-1)
+ else:
+ raise ValueError(
+ f"`use_real_unbind_dim={use_real_unbind_dim}` but should be -1 or -2."
+ )
+
+ out = (x.float() * cos + x_rotated.float() * sin).to(x.dtype)
+
+ return out
+ else:
+ # used for lumina
+ x_rotated = torch.view_as_complex(
+ x.float().reshape(*x.shape[:-1], x.shape[-1] // 2, 2)
+ )
+ freqs_cis = freqs_cis.unsqueeze(2)
+ x_out = torch.view_as_real(x_rotated * freqs_cis).flatten(3)
+
+ return x_out.type_as(x)
+
+
+class BooguImageAttnProcessor:
+ """
+ Processor for implementing scaled dot-product attention with flash attention and variable length sequences.
+
+ This processor is optimized for PyTorch 2.0 and implements:
+ - Flash attention with variable length sequences
+ - Rotary position embeddings (RoPE)
+ - Query-Key normalization
+ - Proportional attention scaling
+
+ Args:
+ None
+
+ Raises:
+ ImportError: If PyTorch version is less than 2.0
+ """
+
+ def __init__(self) -> None:
+ """Initialize the attention processor."""
+ if not hasattr(F, "scaled_dot_product_attention"):
+ raise ImportError(
+ "BooguImageAttnProcessorFlash2Varlen requires PyTorch 2.0. "
+ "Please upgrade PyTorch to version 2.0 or later."
+ )
+
+ def __call__(
+ self,
+ attn: Attention,
+ hidden_states: torch.Tensor,
+ encoder_hidden_states: torch.Tensor,
+ attention_mask: Optional[torch.Tensor] = None,
+ image_rotary_emb: Optional[torch.Tensor] = None,
+ base_sequence_length: Optional[int] = None,
+ ) -> torch.Tensor:
+ """
+ Process attention computation with flash attention.
+
+ Args:
+ attn: Attention module
+ hidden_states: Hidden states tensor of shape (batch_size, seq_len, hidden_dim)
+ encoder_hidden_states: Encoder hidden states tensor
+ attention_mask: Optional attention mask tensor
+ image_rotary_emb: Optional rotary embeddings for image tokens
+ base_sequence_length: Optional base sequence length for proportional attention
+
+ Returns:
+ torch.Tensor: Processed hidden states after attention computation
+ """
+ batch_size, sequence_length, _ = hidden_states.shape
+
+ # Get Query-Key-Value Pair
+ query = attn.to_q(hidden_states)
+ key = attn.to_k(encoder_hidden_states)
+ value = attn.to_v(encoder_hidden_states)
+
+ query_dim = query.shape[-1]
+ inner_dim = key.shape[-1]
+ head_dim = query_dim // attn.heads
+ dtype = query.dtype
+
+ # Get key-value heads
+ kv_heads = inner_dim // head_dim
+
+ # Reshape tensors for attention computation
+ query = query.view(batch_size, -1, attn.heads, head_dim)
+ key = key.view(batch_size, -1, kv_heads, head_dim)
+ value = value.view(batch_size, -1, kv_heads, head_dim)
+
+ # Apply Query-Key normalization
+ if attn.norm_q is not None:
+ query = attn.norm_q(query)
+ if attn.norm_k is not None:
+ key = attn.norm_k(key)
+
+ # Apply Rotary Position Embeddings
+ if image_rotary_emb is not None:
+ query = apply_rotary_emb(query, image_rotary_emb, use_real=False)
+ key = apply_rotary_emb(key, image_rotary_emb, use_real=False)
+
+ query, key = query.to(dtype), key.to(dtype)
+
+ # Calculate attention scale
+ if base_sequence_length is not None:
+ softmax_scale = (
+ math.sqrt(math.log(sequence_length, base_sequence_length)) * attn.scale
+ )
+ else:
+ softmax_scale = attn.scale
+
+ # sdpa expects attn_mask with shape (B, H, Q, K) as boolean (True keeps, False masks)
+ if attention_mask is not None:
+ attention_mask = attention_mask.bool()
+ if attention_mask.dim() == 2:
+ # Standard padding mask [B, L] -> [B, 1, 1, L]
+ attention_mask = attention_mask.view(batch_size, 1, 1, -1)
+ elif attention_mask.dim() == 3:
+ # Robust causal + padding mask construction
+ # Infer valid lengths from diagonal, then build lower-triangular mask within valid lengths
+ B, L, _ = attention_mask.shape
+ diag_valid = torch.diagonal(attention_mask, dim1=-2, dim2=-1)
+ lengths = diag_valid.sum(dim=-1) # [B]
+ arange_L = torch.arange(L, device=attention_mask.device)
+ # Padding masks for queries and keys: shape [B, L]
+ q_valid = arange_L.unsqueeze(0) < lengths.unsqueeze(1)
+ k_valid = q_valid # same lengths assumed
+ # Lower-triangular causal mask [L, L]
+ causal = torch.tril(
+ torch.ones(L, L, dtype=torch.bool, device=attention_mask.device)
+ )
+ # Combine: [B, L, L]
+ combined = causal & q_valid.unsqueeze(-1) & k_valid.unsqueeze(-2)
+ attention_mask = combined.unsqueeze(1) # [B, 1, L, L]
+ else:
+ raise ValueError(
+ f"Unsupported attention_mask shape: {attention_mask.shape}"
+ )
+
+ query = query.transpose(1, 2)
+ key = key.transpose(1, 2)
+ value = value.transpose(1, 2)
+
+ # explicitly repeat key and value to match query length, otherwise using enable_gqa=True results in MATH backend of sdpa in our test of pytorch2.6
+ key = key.repeat_interleave(query.size(-3) // key.size(-3), -3)
+ value = value.repeat_interleave(query.size(-3) // value.size(-3), -3)
+
+ hidden_states = F.scaled_dot_product_attention(
+ query, key, value, attn_mask=attention_mask, scale=softmax_scale
+ )
+ hidden_states = hidden_states.transpose(1, 2).reshape(
+ batch_size, -1, attn.heads * head_dim
+ )
+ hidden_states = hidden_states.type_as(query)
+
+ # Apply output projection
+ hidden_states = attn.to_out[0](hidden_states)
+ hidden_states = attn.to_out[1](hidden_states)
+
+ return hidden_states
+
+
+class BooguImageDoubleStreamSelfAttnProcessor(nn.Module):
+ """
+ Double-stream self-attention processor without flash attention.
+
+ This processor implements double-stream attention where:
+ - Instruction and image features are processed separately to generate QKV
+ - QKV are concatenated and processed together for cross-modal attention
+ - Uses PyTorch's scaled_dot_product_attention for computation
+ - Supports both standard and causal attention masks
+
+ Args:
+ head_dim: Dimension of each attention head
+ num_attention_heads: Number of attention heads for queries
+ num_kv_heads: Number of key-value heads
+ qkv_bias: Whether to use bias in QKV linear layers
+ """
+
+ def __init__(
+ self,
+ head_dim: int,
+ num_attention_heads: int,
+ num_kv_heads: int,
+ qkv_bias: bool = False,
+ ) -> None:
+ """Initialize the double-stream attention processor."""
+ super().__init__()
+ if not hasattr(F, "scaled_dot_product_attention"):
+ raise ImportError(
+ "BooguImageDoubleStreamSelfAttnProcessor requires PyTorch 2.0. "
+ "Please upgrade PyTorch to version 2.0 or later."
+ )
+
+ # Calculate dimensions
+ self.head_dim = head_dim
+ self.num_attention_heads = num_attention_heads
+ self.num_kv_heads = num_kv_heads
+
+ query_dim = head_dim * num_attention_heads
+ kv_dim = head_dim * num_kv_heads
+
+ # Initialize separate Q, K, V linear layers for instruction and image
+ # Query uses num_attention_heads, Key/Value use num_kv_heads
+ self.img_to_q = nn.Linear(query_dim, query_dim, bias=qkv_bias)
+ self.img_to_k = nn.Linear(query_dim, kv_dim, bias=qkv_bias)
+ self.img_to_v = nn.Linear(query_dim, kv_dim, bias=qkv_bias)
+
+ self.instruct_to_q = nn.Linear(query_dim, query_dim, bias=qkv_bias)
+ self.instruct_to_k = nn.Linear(query_dim, kv_dim, bias=qkv_bias)
+ self.instruct_to_v = nn.Linear(query_dim, kv_dim, bias=qkv_bias)
+
+ # Additional output projection layers for instruction and image streams
+ self.instruct_out = nn.Linear(query_dim, query_dim, bias=qkv_bias)
+ self.img_out = nn.Linear(query_dim, query_dim, bias=qkv_bias)
+
+ # Initialize weights
+ self.initialize_weights()
+
+ def initialize_weights(self) -> None:
+ """
+ Initialize the weights of the double-stream attention processor.
+
+ Uses Xavier uniform initialization for linear layers and zero initialization for biases.
+ """
+ # Initialize image stream QKV projection layers
+ nn.init.xavier_uniform_(self.img_to_q.weight)
+ nn.init.xavier_uniform_(self.img_to_k.weight)
+ nn.init.xavier_uniform_(self.img_to_v.weight)
+
+ # Initialize instruction stream QKV projection layers
+ nn.init.xavier_uniform_(self.instruct_to_q.weight)
+ nn.init.xavier_uniform_(self.instruct_to_k.weight)
+ nn.init.xavier_uniform_(self.instruct_to_v.weight)
+
+ # Initialize separate output projection layers
+ nn.init.xavier_uniform_(self.instruct_out.weight)
+ nn.init.xavier_uniform_(self.img_out.weight)
+
+ # Initialize biases if they exist
+ if self.img_to_q.bias is not None:
+ nn.init.zeros_(self.img_to_q.bias)
+ nn.init.zeros_(self.img_to_k.bias)
+ nn.init.zeros_(self.img_to_v.bias)
+ nn.init.zeros_(self.instruct_to_q.bias)
+ nn.init.zeros_(self.instruct_to_k.bias)
+ nn.init.zeros_(self.instruct_to_v.bias)
+ nn.init.zeros_(self.instruct_out.bias)
+ nn.init.zeros_(self.img_out.bias)
+
+ def _concat_instruction_image_features(
+ self,
+ img_hidden_states_list: List[torch.Tensor],
+ instruct_hidden_states_list: List[torch.Tensor],
+ encoder_seq_lengths: List[int],
+ seq_lengths: List[int],
+ ) -> List[torch.Tensor]:
+ """
+ Concatenate instruction (text & image) and reference image features (instruction first, then image).
+
+ Args:
+ img_hidden_states_list: List of image tensors [img_query, img_key, img_value]
+ instruct_hidden_states_list: List of instruction tensors [instruct_query, instruct_key, instruct_value]
+ encoder_seq_lengths: Instruction sequence lengths for each sample [B]
+ seq_lengths: Total sequence lengths for each sample [B]
+
+ Returns:
+ List of concatenated tensors [query, key, value]
+ """
+ assert len(img_hidden_states_list) == len(instruct_hidden_states_list), (
+ f"Length mismatch: img_list={len(img_hidden_states_list)}, instruct_list={len(instruct_hidden_states_list)}"
+ )
+
+ batch_size = img_hidden_states_list[0].shape[0]
+ max_seq_len = max(seq_lengths)
+
+ concatenated_list = []
+
+ for img_tensor, instruct_tensor in zip(
+ img_hidden_states_list, instruct_hidden_states_list
+ ):
+ # Ensure tensors are on the same device
+ device = img_tensor.device
+ if instruct_tensor.device != device:
+ instruct_tensor = instruct_tensor.to(device)
+
+ # Create output tensor with proper shape [B, max_seq_len, feature_dim]
+ feature_dim = img_tensor.shape[-1]
+ concatenated = img_tensor.new_zeros(batch_size, max_seq_len, feature_dim)
+
+ # Concatenate instruction first, then image for each sample
+ for i, (encoder_seq_len, seq_len) in enumerate(
+ zip(encoder_seq_lengths, seq_lengths)
+ ):
+ # Place instruction tokens first
+ concatenated[i, :encoder_seq_len] = instruct_tensor[i, :encoder_seq_len]
+ # Place image tokens after instruction
+ concatenated[i, encoder_seq_len:seq_len] = img_tensor[
+ i, : seq_len - encoder_seq_len
+ ]
+
+ concatenated_list.append(concatenated)
+
+ return concatenated_list
+
+ def _split_instruction_image_features(
+ self,
+ hidden_states_list: List[torch.Tensor],
+ encoder_seq_lengths: List[int],
+ seq_lengths: List[int],
+ ) -> List[Tuple[torch.Tensor, torch.Tensor]]:
+ """
+ Split concatenated features back to instruction and image features.
+ Inverse operation of _concat_instruction_image_features.
+
+ Args:
+ hidden_states_list: List of concatenated tensors (usually just one element)
+ encoder_seq_lengths: Instruction sequence lengths for each sample [B]
+ seq_lengths: Total sequence lengths for each sample [B]
+
+ Returns:
+ List of tuples, each containing (instruct_hidden_states, img_hidden_states)
+ """
+ result_list = []
+
+ for hidden_states in hidden_states_list:
+ batch_size = hidden_states.shape[0]
+ feature_dim = hidden_states.shape[-1]
+
+ # Get maximum lengths for instruction and image
+ max_instruct_len = max(encoder_seq_lengths)
+ max_img_len = max(
+ seq_len - encoder_seq_len
+ for seq_len, encoder_seq_len in zip(seq_lengths, encoder_seq_lengths)
+ )
+
+ # Create output tensors [B, max_len, feature_dim]
+ instruct_hidden_states = hidden_states.new_zeros(
+ batch_size, max_instruct_len, feature_dim
+ )
+ img_hidden_states = hidden_states.new_zeros(
+ batch_size, max_img_len, feature_dim
+ )
+
+ # Split each sample back to instruction and image
+ for i, (encoder_seq_len, seq_len) in enumerate(
+ zip(encoder_seq_lengths, seq_lengths)
+ ):
+ img_len = seq_len - encoder_seq_len
+
+ # Extract instruction portion
+ instruct_hidden_states[i, :encoder_seq_len] = hidden_states[
+ i, :encoder_seq_len
+ ]
+ # Extract image portion
+ img_hidden_states[i, :img_len] = hidden_states[
+ i, encoder_seq_len:seq_len
+ ]
+
+ result_list.append((instruct_hidden_states, img_hidden_states))
+
+ return result_list
+
+ def __call__(
+ self,
+ attn: Attention,
+ img_hidden_states: torch.Tensor,
+ instruct_hidden_states: torch.Tensor,
+ joint_attention_mask: Optional[torch.Tensor] = None,
+ rotary_emb: Optional[torch.Tensor] = None,
+ encoder_seq_lengths: List[
+ int
+ ] = None, # [B] - Instruction sequence lengths for each sample
+ seq_lengths: List[int] = None, # [B] - Total sequence lengths for each sample
+ base_sequence_length: Optional[int] = None,
+ ) -> torch.Tensor:
+ """
+ Process double-stream self-attention computation with PyTorch's scaled_dot_product_attention.
+
+ Args:
+ attn: Attention module
+ img_hidden_states: Image hidden states tensor [B, L_img, D]
+ instruct_hidden_states: Instruction hidden states tensor [B, L_instruct, D]
+ joint_attention_mask: Combined attention mask [B, L_total]
+ rotary_emb: Rotary embeddings for the joint sequence
+ encoder_seq_lengths: Instruction sequence lengths for each sample [B]
+ seq_lengths: Total sequence lengths for each sample [B]
+ base_sequence_length: Optional base sequence length for proportional attention
+
+ Returns:
+ torch.Tensor: Processed hidden states after attention computation
+ """
+ batch_size = img_hidden_states.shape[0]
+ L_instruct = instruct_hidden_states.shape[1]
+ L_img = img_hidden_states.shape[1]
+
+ # Ensure Q, K, V linear layers are on the same device as input tensors
+ device = img_hidden_states.device
+ for layer in [
+ self.img_to_q,
+ self.img_to_k,
+ self.img_to_v,
+ self.instruct_to_q,
+ self.instruct_to_k,
+ self.instruct_to_v,
+ self.instruct_out,
+ self.img_out,
+ ]:
+ if (
+ (layer.weight.device != device)
+ and (str(layer.weight.device).lower() != "meta")
+ and (str(device).lower() not in {"meta", "auto"})
+ ):
+ layer = layer.to(device)
+
+ # Generate Q, K, V for image and instruction streams (NO head reshaping yet)
+ img_query = self.img_to_q(img_hidden_states) # [B, L_img, query_dim]
+ img_key = self.img_to_k(img_hidden_states) # [B, L_img, kv_dim]
+ img_value = self.img_to_v(img_hidden_states) # [B, L_img, kv_dim]
+
+ instruct_query = self.instruct_to_q(
+ instruct_hidden_states
+ ) # [B, L_instruct, query_dim]
+ instruct_key = self.instruct_to_k(
+ instruct_hidden_states
+ ) # [B, L_instruct, kv_dim]
+ instruct_value = self.instruct_to_v(
+ instruct_hidden_states
+ ) # [B, L_instruct, kv_dim]
+
+ # Use helper function to concatenate QKV (instruction first, then image)
+ img_list = [img_query, img_key, img_value] # [B, L_img, feature_dim] each
+ instruct_list = [
+ instruct_query,
+ instruct_key,
+ instruct_value,
+ ] # [B, L_instruct, feature_dim] each
+ concatenated_list = self._concat_instruction_image_features(
+ img_list, instruct_list, encoder_seq_lengths, seq_lengths
+ )
+ query, key, value = concatenated_list # [B, max_seq_len, feature_dim] each
+
+ # From here, follow exactly the same logic as BooguImageAttnProcessor
+ sequence_length = max(seq_lengths)
+
+ query_dim = query.shape[-1]
+ inner_dim = key.shape[-1]
+ head_dim = query_dim // attn.heads
+ dtype = query.dtype
+
+ # Get key-value heads
+ kv_heads = inner_dim // head_dim
+
+ # Reshape tensors for attention computation
+ query = query.view(batch_size, -1, attn.heads, head_dim)
+ key = key.view(batch_size, -1, kv_heads, head_dim)
+ value = value.view(batch_size, -1, kv_heads, head_dim)
+
+ # Apply Query-Key normalization
+ if attn.norm_q is not None:
+ query = attn.norm_q(query)
+ if attn.norm_k is not None:
+ key = attn.norm_k(key)
+
+ # Apply Rotary Position Embeddings
+ if rotary_emb is not None:
+ query = apply_rotary_emb(query, rotary_emb, use_real=False)
+ key = apply_rotary_emb(key, rotary_emb, use_real=False)
+
+ query, key = query.to(dtype), key.to(dtype)
+
+ # Calculate attention scale
+ if base_sequence_length is not None:
+ softmax_scale = (
+ math.sqrt(math.log(sequence_length, base_sequence_length)) * attn.scale
+ )
+ else:
+ softmax_scale = attn.scale
+
+ # scaled_dot_product_attention expects attention_mask shape to be
+ # (batch, heads, source_length, target_length)
+ if joint_attention_mask is not None:
+ joint_attention_mask = joint_attention_mask.bool()
+ if joint_attention_mask.dim() == 2:
+ # Standard mask [B, seq_len] -> [B, 1, 1, seq_len]
+ joint_attention_mask = joint_attention_mask.view(batch_size, 1, 1, -1)
+ elif joint_attention_mask.dim() == 3:
+ # Causal mask [B, seq_len, seq_len] -> [B, 1, seq_len, seq_len]
+ joint_attention_mask = joint_attention_mask.unsqueeze(1)
+ else:
+ raise ValueError(
+ f"Unsupported joint_attention_mask shape: {joint_attention_mask.shape}"
+ )
+
+ query = query.transpose(1, 2)
+ key = key.transpose(1, 2)
+ value = value.transpose(1, 2)
+
+ # explicitly repeat key and value to match query length, otherwise using enable_gqa=True results in MATH backend of sdpa in our test of pytorch2.6
+ key = key.repeat_interleave(query.size(-3) // key.size(-3), -3)
+ value = value.repeat_interleave(query.size(-3) // value.size(-3), -3)
+
+ hidden_states = F.scaled_dot_product_attention(
+ query, key, value, attn_mask=joint_attention_mask, scale=softmax_scale
+ )
+ hidden_states = hidden_states.transpose(1, 2).reshape(
+ batch_size, -1, attn.heads * head_dim
+ )
+ hidden_states = hidden_states.type_as(query)
+
+ # Split hidden_states back to instruction and image, apply separate output projections, then merge
+ split_results = self._split_instruction_image_features(
+ [hidden_states], encoder_seq_lengths, seq_lengths
+ )
+ instruct_hidden_states, img_hidden_states = split_results[
+ 0
+ ] # [B, max_instruct_len, feature_dim], [B, max_img_len, feature_dim]
+
+ # Apply separate output projections for instruction and image
+ instruct_projected = self.instruct_out(
+ instruct_hidden_states
+ ) # [B, max_instruct_len, feature_dim]
+ img_projected = self.img_out(img_hidden_states) # [B, max_img_len, feature_dim]
+
+ # Merge back to joint representation
+ merged_list = self._concat_instruction_image_features(
+ [img_projected], [instruct_projected], encoder_seq_lengths, seq_lengths
+ )
+ hidden_states = merged_list[0] # [B, max_seq_len, feature_dim]
+
+ # Apply final output projection
+ hidden_states = attn.to_out[0](hidden_states)
+ hidden_states = attn.to_out[1](hidden_states)
+
+ return hidden_states
+
+
+class PromptEmbedding(torch.nn.Module):
+ _supports_gradient_checkpointing = True
+ _no_split_modules = ["BooguImageTransformerBlock"]
+ _skip_layerwise_casting_patterns = ["prompt_token_embedding", "norm"]
+
+ def __init__(self, prompt_tuning_configs):
+ super().__init__()
+
+ num_trainable_prompt_tokens = prompt_tuning_configs.get(
+ "num_trainable_prompt_tokens", 32
+ )
+ hidden_size = prompt_tuning_configs.get("hidden_size", 2048)
+ num_attention_heads = prompt_tuning_configs.get("num_attention_heads", 32)
+ num_kv_heads = prompt_tuning_configs.get("num_kv_heads", 8)
+ multiple_of = prompt_tuning_configs.get("multiple_of", 256)
+ ffn_dim_multiplier = prompt_tuning_configs.get("ffn_dim_multiplier", None)
+ norm_eps = prompt_tuning_configs.get("norm_eps", 1e-5)
+ num_layers = prompt_tuning_configs.get("num_layers", 2)
+ theta = prompt_tuning_configs.get("theta", 10000)
+
+ self.register_to_config(
+ num_trainable_prompt_tokens=num_trainable_prompt_tokens,
+ hidden_size=hidden_size,
+ num_attention_heads=num_attention_heads,
+ num_kv_heads=num_kv_heads,
+ multiple_of=multiple_of,
+ ffn_dim_multiplier=ffn_dim_multiplier,
+ norm_eps=norm_eps,
+ num_layers=num_layers,
+ theta=theta,
+ )
+
+ self.prompt_tuning_configs = prompt_tuning_configs
+
+ prompt_emb_head_dim = self.config.hidden_size // self.config.num_attention_heads
+
+ self.prompt_token_embedding = nn.Embedding(
+ num_embeddings=self.config.num_trainable_prompt_tokens,
+ embedding_dim=self.config.hidden_size,
+ )
+
+ # Rotary embedding for prompt tokens.
+ self.prompt_rope_embedder = BooguImagePromptTuningRotaryPosEmbed(
+ theta=self.config.theta,
+ dim=prompt_emb_head_dim,
+ num_trainable_prompt_tokens=self.config.num_trainable_prompt_tokens,
+ )
+
+ self.prompt_tuning_layers = nn.ModuleList(
+ [
+ BooguImageTransformerBlock(
+ dim=self.config.hidden_size,
+ num_attention_heads=self.config.num_attention_heads,
+ num_kv_heads=self.config.num_kv_heads,
+ multiple_of=self.config.multiple_of,
+ ffn_dim_multiplier=self.config.ffn_dim_multiplier,
+ norm_eps=self.config.norm_eps,
+ modulation=False,
+ )
+ for _ in range(self.config.num_layers)
+ ]
+ )
+
+ self.gradient_checkpointing = False
+
+ self.initialize_weights()
+
+ def initialize_weights(self) -> None:
+ # Small std keeps prompt tuning stable at init.
+ nn.init.normal_(self.prompt_token_embedding.weight, mean=0.0, std=0.02)
+
+ def forward(self, idx=None, batch_size=1, device=None, use_causal_mask=True):
+ if idx is None:
+ prompt_embeddings = self.prompt_token_embedding.weight
+ else:
+ prompt_embeddings = self.prompt_token_embedding(idx)
+
+ # Expand to [B, num_tokens, hidden_dim].
+ hidden_states = prompt_embeddings.unsqueeze(0).expand(batch_size, -1, -1)
+
+ rotary_emb, attention_mask = self.prompt_rope_embedder(
+ batch_size, device, use_causal_mask
+ )
+
+ for i, layer in enumerate(self.prompt_tuning_layers):
+ if torch.is_grad_enabled() and self.gradient_checkpointing:
+ hidden_states = self._gradient_checkpointing_func(
+ layer,
+ hidden_states,
+ attention_mask,
+ rotary_emb,
+ )
+ else:
+ hidden_states = layer(
+ hidden_states,
+ attention_mask,
+ rotary_emb,
+ )
+ return hidden_states
+
+ @classmethod
+ def from_config(cls, config, **kwargs):
+ # `config` is loaded from config.json.
+ instance = cls(prompt_tuning_configs=config)
+
+ weight_dtype = kwargs.get("weight_dtype", None)
+ if weight_dtype is not None:
+ for p in instance.parameters():
+ p.data = p.data.to(dtype=weight_dtype)
+
+ return instance
+
+
+class BooguImageTransformerBlock(nn.Module):
+ """
+ Basic Boogu-Image transformer block: attention + MLP + RMSNorm.
+ """
+
+ def __init__(
+ self,
+ dim: int,
+ num_attention_heads: int,
+ num_kv_heads: int,
+ multiple_of: int,
+ ffn_dim_multiplier: float,
+ norm_eps: float,
+ modulation: bool = True,
+ ) -> None:
+ """Initialize the transformer block."""
+ super().__init__()
+ self.head_dim = dim // num_attention_heads
+ self.modulation = modulation
+
+ processor = BooguImageAttnProcessor()
+
+ # Initialize attention layer
+ self.attn = Attention(
+ query_dim=dim,
+ cross_attention_dim=None,
+ dim_head=dim // num_attention_heads,
+ qk_norm="rms_norm",
+ heads=num_attention_heads,
+ kv_heads=num_kv_heads,
+ eps=1e-5,
+ bias=False,
+ out_bias=False,
+ processor=processor,
+ )
+
+ # Initialize feed-forward network
+ self.feed_forward = LuminaFeedForward(
+ dim=dim,
+ inner_dim=4 * dim,
+ multiple_of=multiple_of,
+ ffn_dim_multiplier=ffn_dim_multiplier,
+ )
+
+ # Initialize normalization layers
+ if modulation:
+ self.norm1 = LuminaRMSNormZero(
+ embedding_dim=dim, norm_eps=norm_eps, norm_elementwise_affine=True
+ )
+ else:
+ self.norm1 = RMSNorm(dim, eps=norm_eps)
+
+ self.ffn_norm1 = RMSNorm(dim, eps=norm_eps)
+ self.norm2 = RMSNorm(dim, eps=norm_eps)
+ self.ffn_norm2 = RMSNorm(dim, eps=norm_eps)
+
+ self.initialize_weights()
+
+ def initialize_weights(self) -> None:
+ """Initialize linear weights and modulation parameters."""
+ nn.init.xavier_uniform_(self.attn.to_q.weight)
+ nn.init.xavier_uniform_(self.attn.to_k.weight)
+ nn.init.xavier_uniform_(self.attn.to_v.weight)
+ nn.init.xavier_uniform_(self.attn.to_out[0].weight)
+
+ nn.init.xavier_uniform_(self.feed_forward.linear_1.weight)
+ nn.init.xavier_uniform_(self.feed_forward.linear_2.weight)
+ nn.init.xavier_uniform_(self.feed_forward.linear_3.weight)
+
+ if self.modulation:
+ nn.init.zeros_(self.norm1.linear.weight)
+ nn.init.zeros_(self.norm1.linear.bias)
+
+ def forward(
+ self,
+ hidden_states: torch.Tensor,
+ attention_mask: torch.Tensor,
+ image_rotary_emb: torch.Tensor,
+ temb: Optional[torch.Tensor] = None,
+ ) -> torch.Tensor:
+ """
+ Forward pass of the transformer block.
+
+ Args:
+ hidden_states: Input hidden states tensor
+ attention_mask: Attention mask tensor
+ image_rotary_emb: Rotary embeddings for image tokens
+ temb: Optional timestep embedding tensor
+
+ Returns:
+ torch.Tensor: Output hidden states after transformer block processing
+ """
+
+ enable_taylorseer = getattr(self, "enable_taylorseer", False)
+
+ if enable_taylorseer:
+ if self.modulation:
+ if temb is None:
+ raise ValueError("temb must be provided when modulation is enabled")
+
+ if self.current["type"] == "full":
+ self.current["module"] = "total"
+ taylor_cache_init(cache_dic=self.cache_dic, current=self.current)
+
+ norm_hidden_states, gate_msa, scale_mlp, gate_mlp = self.norm1(
+ hidden_states, temb
+ )
+ attn_output = self.attn(
+ hidden_states=norm_hidden_states,
+ encoder_hidden_states=norm_hidden_states,
+ attention_mask=attention_mask,
+ image_rotary_emb=image_rotary_emb,
+ )
+ hidden_states = hidden_states + gate_msa.unsqueeze(
+ 1
+ ).tanh() * self.norm2(attn_output)
+ mlp_output = self.feed_forward(
+ self.ffn_norm1(hidden_states) * (1 + scale_mlp.unsqueeze(1))
+ )
+ hidden_states = hidden_states + gate_mlp.unsqueeze(
+ 1
+ ).tanh() * self.ffn_norm2(mlp_output)
+
+ derivative_approximation(
+ cache_dic=self.cache_dic,
+ current=self.current,
+ feature=hidden_states,
+ )
+
+ elif self.current["type"] == "Taylor":
+ self.current["module"] = "total"
+ hidden_states = taylor_formula(
+ cache_dic=self.cache_dic, current=self.current
+ )
+ else:
+ norm_hidden_states = self.norm1(hidden_states)
+ attn_output = self.attn(
+ hidden_states=norm_hidden_states,
+ encoder_hidden_states=norm_hidden_states,
+ attention_mask=attention_mask,
+ image_rotary_emb=image_rotary_emb,
+ )
+ hidden_states = hidden_states + self.norm2(attn_output)
+ mlp_output = self.feed_forward(self.ffn_norm1(hidden_states))
+ hidden_states = hidden_states + self.ffn_norm2(mlp_output)
+ else:
+ if self.modulation:
+ if temb is None:
+ raise ValueError("temb must be provided when modulation is enabled")
+ norm_hidden_states, gate_msa, scale_mlp, gate_mlp = self.norm1(
+ hidden_states, temb
+ )
+
+ attn_output = self.attn(
+ hidden_states=norm_hidden_states,
+ encoder_hidden_states=norm_hidden_states,
+ attention_mask=attention_mask,
+ image_rotary_emb=image_rotary_emb,
+ )
+ hidden_states = hidden_states + gate_msa.unsqueeze(
+ 1
+ ).tanh() * self.norm2(attn_output)
+ mlp_output = self.feed_forward(
+ self.ffn_norm1(hidden_states) * (1 + scale_mlp.unsqueeze(1))
+ )
+ hidden_states = hidden_states + gate_mlp.unsqueeze(
+ 1
+ ).tanh() * self.ffn_norm2(mlp_output)
+ else:
+ norm_hidden_states = self.norm1(hidden_states)
+ attn_output = self.attn(
+ hidden_states=norm_hidden_states,
+ encoder_hidden_states=norm_hidden_states,
+ attention_mask=attention_mask,
+ image_rotary_emb=image_rotary_emb,
+ )
+ hidden_states = hidden_states + self.norm2(attn_output)
+ mlp_output = self.feed_forward(self.ffn_norm1(hidden_states))
+ hidden_states = hidden_states + self.ffn_norm2(mlp_output)
+
+ return hidden_states
+
+
+class BooguImageDoubleStreamTransformerBlock(nn.Module):
+ """
+ Boogu-Image double-stream block.
+ Here "double-stream" is the same idea as a "dual-stream" layer:
+ instruction tokens and image tokens are processed in parallel streams.
+ """
+
+ def __init__(
+ self,
+ dim: int,
+ num_attention_heads: int,
+ num_kv_heads: int,
+ multiple_of: int,
+ ffn_dim_multiplier: float,
+ norm_eps: float,
+ modulation: bool = True,
+ ) -> None:
+ """Initialize the double stream transformer block."""
+ super().__init__()
+ self.head_dim = dim // num_attention_heads
+ self.num_attention_heads = num_attention_heads
+ self.modulation = modulation
+ self.hidden_size = dim
+
+ processor = BooguImageAttnProcessor()
+
+ double_stream_processor = BooguImageDoubleStreamSelfAttnProcessor(
+ head_dim=self.head_dim,
+ num_attention_heads=num_attention_heads,
+ num_kv_heads=num_kv_heads,
+ qkv_bias=False,
+ )
+
+ # Image stream components.
+ self.img_instruct_attn = Attention(
+ query_dim=dim,
+ cross_attention_dim=None,
+ dim_head=dim // num_attention_heads,
+ qk_norm="rms_norm",
+ heads=num_attention_heads,
+ kv_heads=num_kv_heads,
+ eps=1e-5,
+ bias=False,
+ out_bias=False,
+ processor=double_stream_processor,
+ )
+
+ self.img_self_attn = Attention(
+ query_dim=dim,
+ cross_attention_dim=None,
+ dim_head=dim // num_attention_heads,
+ qk_norm="rms_norm",
+ heads=num_attention_heads,
+ kv_heads=num_kv_heads,
+ eps=1e-5,
+ bias=False,
+ out_bias=False,
+ processor=processor,
+ )
+
+ self.img_feed_forward = LuminaFeedForward(
+ dim=dim,
+ inner_dim=4 * dim,
+ multiple_of=multiple_of,
+ ffn_dim_multiplier=ffn_dim_multiplier,
+ )
+
+ if modulation:
+ # Image modulation terms: cross-attn, MLP, self-attn.
+ self.img_norm1 = LuminaRMSNormZero(
+ embedding_dim=dim, norm_eps=norm_eps, norm_elementwise_affine=True
+ )
+ self.img_norm2 = LuminaRMSNormZero(
+ embedding_dim=dim, norm_eps=norm_eps, norm_elementwise_affine=True
+ )
+ self.img_norm3 = LuminaRMSNormZero(
+ embedding_dim=dim, norm_eps=norm_eps, norm_elementwise_affine=True
+ )
+ else:
+ self.img_norm1 = RMSNorm(dim, eps=norm_eps)
+ self.img_norm2 = RMSNorm(dim, eps=norm_eps)
+ self.img_norm3 = RMSNorm(dim, eps=norm_eps)
+
+ self.img_ffn_norm1 = RMSNorm(dim, eps=norm_eps)
+ self.img_attn_norm = RMSNorm(dim, eps=norm_eps)
+ self.img_self_attn_norm = RMSNorm(dim, eps=norm_eps)
+ self.img_ffn_norm2 = RMSNorm(dim, eps=norm_eps)
+
+ # Instruction stream components.
+ self.instruct_feed_forward = LuminaFeedForward(
+ dim=dim,
+ inner_dim=4 * dim,
+ multiple_of=multiple_of,
+ ffn_dim_multiplier=ffn_dim_multiplier,
+ )
+
+ if modulation:
+ # Instruction modulation terms: cross-attn, MLP.
+ self.instruct_norm1 = LuminaRMSNormZero(
+ embedding_dim=dim, norm_eps=norm_eps, norm_elementwise_affine=True
+ )
+ self.instruct_norm2 = LuminaRMSNormZero(
+ embedding_dim=dim, norm_eps=norm_eps, norm_elementwise_affine=True
+ )
+ else:
+ self.instruct_norm1 = RMSNorm(dim, eps=norm_eps)
+ self.instruct_norm2 = RMSNorm(dim, eps=norm_eps)
+
+ self.instruct_ffn_norm1 = RMSNorm(dim, eps=norm_eps)
+ self.instruct_attn_norm = RMSNorm(dim, eps=norm_eps)
+ self.instruct_ffn_norm2 = RMSNorm(dim, eps=norm_eps)
+
+ self.initialize_weights()
+
+ # double_stream_processor owns its own q/k/v projections.
+ for param in self.img_instruct_attn.to_q.parameters():
+ param.requires_grad = False
+ for param in self.img_instruct_attn.to_k.parameters():
+ param.requires_grad = False
+ for param in self.img_instruct_attn.to_v.parameters():
+ param.requires_grad = False
+
+ del self.img_instruct_attn.to_k
+ del self.img_instruct_attn.to_v
+ del self.img_instruct_attn.to_q
+
+ def initialize_weights(self) -> None:
+ """Initialize linear weights and modulation parameters."""
+ nn.init.xavier_uniform_(self.img_instruct_attn.to_out[0].weight)
+
+ # Keep Xavier init consistent across Boogu-Image blocks.
+ nn.init.xavier_uniform_(self.img_self_attn.to_q.weight)
+ nn.init.xavier_uniform_(self.img_self_attn.to_k.weight)
+ nn.init.xavier_uniform_(self.img_self_attn.to_v.weight)
+ nn.init.xavier_uniform_(self.img_self_attn.to_out[0].weight)
+
+ nn.init.xavier_uniform_(self.img_feed_forward.linear_1.weight)
+ nn.init.xavier_uniform_(self.img_feed_forward.linear_2.weight)
+ nn.init.xavier_uniform_(self.img_feed_forward.linear_3.weight)
+
+ nn.init.xavier_uniform_(self.instruct_feed_forward.linear_1.weight)
+ nn.init.xavier_uniform_(self.instruct_feed_forward.linear_2.weight)
+ nn.init.xavier_uniform_(self.instruct_feed_forward.linear_3.weight)
+
+ # Initialize modulation parameters
+ if self.modulation:
+ nn.init.zeros_(self.img_norm1.linear.weight)
+ nn.init.zeros_(self.img_norm1.linear.bias)
+ nn.init.zeros_(self.img_norm2.linear.weight)
+ nn.init.zeros_(self.img_norm2.linear.bias)
+ nn.init.zeros_(self.img_norm3.linear.weight)
+ nn.init.zeros_(self.img_norm3.linear.bias)
+
+ nn.init.zeros_(self.instruct_norm1.linear.weight)
+ nn.init.zeros_(self.instruct_norm1.linear.bias)
+ nn.init.zeros_(self.instruct_norm2.linear.weight)
+ nn.init.zeros_(self.instruct_norm2.linear.bias)
+
+ def forward(
+ self,
+ img_hidden_states: torch.Tensor, # [B, L_img, D] - Image tokens (ref_img + noise_img)
+ instruct_hidden_states: torch.Tensor, # [B, L_instruct, D] - Instruction tokens
+ img_attention_mask: torch.Tensor, # [B, L_img] - Attention mask for [ref_img + noise_img]
+ joint_attention_mask: torch.Tensor, # [B, L_total] - Combined attention mask for [instruct + img]
+ image_rotary_emb: torch.Tensor, # [B, L_img, head_dim] - Rotary embeddings for [ref_img + noise_img]
+ rotary_emb: torch.Tensor, # [B, L_total, head_dim] - Rotary embeddings for [instruct + img]
+ temb: Optional[torch.Tensor] = None, # [B, 1024] - Timestep embeddings
+ encoder_seq_lengths: List[
+ int
+ ] = None, # [B] - Instruction sequence lengths for each sample
+ seq_lengths: List[int] = None, # [B] - Total sequence lengths for each sample
+ ) -> Tuple[torch.Tensor, torch.Tensor]:
+ """
+ Run one dual-stream (double-stream) block step.
+ Returns updated `(img_hidden_states, instruct_hidden_states)`.
+ """
+ if self.modulation and temb is None:
+ raise ValueError("temb must be provided when modulation is enabled")
+
+ enable_taylorseer = getattr(self, "enable_taylorseer", False)
+ if enable_taylorseer:
+ self.current["module"] = "total"
+ if self.current["type"] == "Taylor":
+ return taylor_formula_4_double_stream(
+ cache_dic=self.cache_dic, current=self.current
+ )
+ if self.current["type"] == "full":
+ taylor_cache_init(cache_dic=self.cache_dic, current=self.current)
+
+ # Extract dimensions
+ batch_size = img_hidden_states.shape[0]
+ L_instruct = instruct_hidden_states.shape[1] # Instruction sequence length
+ L_img = img_hidden_states.shape[
+ 1
+ ] # Image sequence length (ref_img + noise_img)
+
+ if self.modulation:
+ # Step 1: modulation for both streams.
+ img_norm1_out, img_gate_msa, img_scale_mlp, img_gate_mlp = self.img_norm1(
+ img_hidden_states, temb
+ )
+ img_norm2_out, img_shift_mlp, _, _ = self.img_norm2(img_hidden_states, temb)
+ img_norm3_out, img_gate_self, _, _ = self.img_norm3(img_hidden_states, temb)
+
+ (
+ instruct_norm1_out,
+ instruct_gate_msa,
+ instruct_scale_mlp,
+ instruct_gate_mlp,
+ ) = self.instruct_norm1(instruct_hidden_states, temb)
+ instruct_norm2_out, instruct_shift_mlp, _, _ = self.instruct_norm2(
+ instruct_hidden_states, temb
+ )
+
+ # Step 2: joint attention on [instruct + img].
+ # Call processor directly because Attention.forward does not expose these dual-stream args.
+ joint_attn_out = self.img_instruct_attn.processor(
+ attn=self.img_instruct_attn,
+ img_hidden_states=img_norm1_out,
+ instruct_hidden_states=instruct_norm1_out,
+ joint_attention_mask=joint_attention_mask,
+ rotary_emb=rotary_emb,
+ encoder_seq_lengths=encoder_seq_lengths,
+ seq_lengths=seq_lengths,
+ )
+
+ # Split back into instruction/image segments.
+ instruct_attn_out = instruct_hidden_states.new_zeros(
+ batch_size, L_instruct, self.hidden_size
+ )
+ img_attn_out = img_hidden_states.new_zeros(
+ batch_size, L_img, self.hidden_size
+ )
+ for i, (encoder_seq_len, seq_len) in enumerate(
+ zip(encoder_seq_lengths, seq_lengths)
+ ):
+ instruct_attn_out[i, :encoder_seq_len] = joint_attn_out[
+ i, :encoder_seq_len
+ ]
+ img_attn_out[i, : seq_len - encoder_seq_len] = joint_attn_out[
+ i, encoder_seq_len:seq_len
+ ]
+
+ # Step 3: image self-attention.
+ img_self_attn_out = self.img_self_attn(
+ hidden_states=img_norm3_out,
+ encoder_hidden_states=img_norm3_out,
+ attention_mask=img_attention_mask,
+ image_rotary_emb=image_rotary_emb,
+ )
+
+ # Step 4: residual updates.
+ img_hidden_states = img_hidden_states + img_gate_msa.unsqueeze(
+ 1
+ ).tanh() * self.img_attn_norm(img_attn_out)
+ img_hidden_states = img_hidden_states + img_gate_self.unsqueeze(
+ 1
+ ).tanh() * self.img_self_attn_norm(img_self_attn_out)
+
+ img_mlp_input = (
+ 1 + img_scale_mlp.unsqueeze(1)
+ ) * img_norm2_out + img_shift_mlp.unsqueeze(1)
+ img_mlp_out = self.img_feed_forward(self.img_ffn_norm1(img_mlp_input))
+ img_hidden_states = img_hidden_states + img_gate_mlp.unsqueeze(
+ 1
+ ).tanh() * self.img_ffn_norm2(img_mlp_out)
+
+ instruct_hidden_states = (
+ instruct_hidden_states
+ + instruct_gate_msa.unsqueeze(1).tanh()
+ * self.instruct_attn_norm(instruct_attn_out)
+ )
+
+ instruct_mlp_input = (
+ 1 + instruct_scale_mlp.unsqueeze(1)
+ ) * instruct_norm2_out + instruct_shift_mlp.unsqueeze(1)
+ instruct_mlp_out = self.instruct_feed_forward(
+ self.instruct_ffn_norm1(instruct_mlp_input)
+ )
+ instruct_hidden_states = (
+ instruct_hidden_states
+ + instruct_gate_mlp.unsqueeze(1).tanh()
+ * self.instruct_ffn_norm2(instruct_mlp_out)
+ )
+
+ else:
+ # Non-modulated branch used by context-style blocks.
+ img_norm1_out = self.img_norm1(img_hidden_states)
+ img_norm3_out = self.img_norm3(img_hidden_states)
+ instruct_norm1_out = self.instruct_norm1(instruct_hidden_states)
+
+ # Same processor path as above.
+ joint_attn_out = self.img_instruct_attn.processor(
+ attn=self.img_instruct_attn,
+ img_hidden_states=img_norm1_out,
+ instruct_hidden_states=instruct_norm1_out,
+ joint_attention_mask=joint_attention_mask,
+ rotary_emb=rotary_emb,
+ encoder_seq_lengths=encoder_seq_lengths,
+ seq_lengths=seq_lengths,
+ )
+
+ instruct_attn_out = instruct_hidden_states.new_zeros(
+ batch_size, L_instruct, self.hidden_size
+ )
+ img_attn_out = img_hidden_states.new_zeros(
+ batch_size, L_img, self.hidden_size
+ )
+ for i, (encoder_seq_len, seq_len) in enumerate(
+ zip(encoder_seq_lengths, seq_lengths)
+ ):
+ instruct_attn_out[i, :encoder_seq_len] = joint_attn_out[
+ i, :encoder_seq_len
+ ]
+ img_attn_out[i, : seq_len - encoder_seq_len] = joint_attn_out[
+ i, encoder_seq_len:seq_len
+ ]
+
+ img_self_attn_out = self.img_self_attn(
+ hidden_states=img_norm3_out,
+ encoder_hidden_states=img_norm3_out,
+ attention_mask=img_attention_mask,
+ image_rotary_emb=image_rotary_emb,
+ )
+
+ img_hidden_states = img_hidden_states + self.img_attn_norm(img_attn_out)
+ img_hidden_states = img_hidden_states + self.img_self_attn_norm(
+ img_self_attn_out
+ )
+ img_norm2_out = self.img_norm2(img_hidden_states)
+ img_mlp_out = self.img_feed_forward(self.img_ffn_norm1(img_norm2_out))
+ img_hidden_states = img_hidden_states + self.img_ffn_norm2(img_mlp_out)
+
+ instruct_hidden_states = instruct_hidden_states + self.instruct_attn_norm(
+ instruct_attn_out
+ )
+ instruct_norm2_out = self.instruct_norm2(instruct_hidden_states)
+ instruct_mlp_out = self.instruct_feed_forward(
+ self.instruct_ffn_norm1(instruct_norm2_out)
+ )
+ instruct_hidden_states = instruct_hidden_states + self.instruct_ffn_norm2(
+ instruct_mlp_out
+ )
+
+ if enable_taylorseer and self.current["type"] == "full":
+ derivative_approximation_4_double_stream(
+ cache_dic=self.cache_dic,
+ current=self.current,
+ feature=(img_hidden_states, instruct_hidden_states),
+ )
+
+ return img_hidden_states, instruct_hidden_states
+
+
+BooguImageSingleStreamTransformerBlock = BooguImageTransformerBlock
+
+
+@dataclass
+class BooguImageDiTConfig:
+ patch_size: int = 2
+ in_channels: int = 16
+ out_channels: Optional[int] = None
+ hidden_size: int = 3360
+ num_layers: int = 40
+ num_double_stream_layers: int = 8
+ num_refiner_layers: int = 2
+ num_attention_heads: int = 28
+ num_kv_heads: int = 7
+ multiple_of: int = 256
+ ffn_dim_multiplier: Optional[float] = None
+ norm_eps: float = 1e-5
+ axes_dim_rope: Tuple[int, int, int] = (40, 40, 40)
+ axes_lens: Tuple[int, int, int] = (2048, 1664, 1664)
+ instruction_feature_configs: Dict[str, Any] = None
+ prompt_tuning_configs: Dict[str, Any] = None
+ timestep_scale: float = 1000.0
+
+
+class BooguImageDiT(torch.nn.Module):
+ """
+ Boogu-Image transformer with mixed stream topology.
+ Early layers use double-stream (aka dual-stream) processing, then switch
+ to single-stream joint processing.
+ """
+
+ _supports_gradient_checkpointing = True
+ _no_split_modules = [
+ "BooguImageTransformerBlock",
+ "BooguImageSingleStreamTransformerBlock",
+ "BooguImageDoubleStreamTransformerBlock",
+ "PromptEmbedding",
+ "nn.Embedding",
+ ]
+ _repeated_blocks = [
+ "BooguImageTransformerBlock",
+ "BooguImageSingleStreamTransformerBlock",
+ "BooguImageDoubleStreamTransformerBlock",
+ ]
+ _skip_layerwise_casting_patterns = ["x_embedder", "norm", "embedding"]
+
+ def __init__(
+ self,
+ patch_size: int = 2,
+ in_channels: int = 16,
+ out_channels: Optional[int] = None,
+ hidden_size: int = 3360,
+ num_layers: int = 40,
+ num_double_stream_layers: int = 8,
+ num_refiner_layers: int = 2,
+ num_attention_heads: int = 28,
+ num_kv_heads: int = 7,
+ multiple_of: int = 256,
+ ffn_dim_multiplier: Optional[float] = None,
+ norm_eps: float = 1e-5,
+ axes_dim_rope: Tuple[int, int, int] = (40, 40, 40),
+ axes_lens: Tuple[int, int, int] = (2048, 1664, 1664),
+ # instruction_feat_dim: int = 1024,
+ instruction_feature_configs: Dict[str, Any] = {'instruction_feat_dim': 4096, 'num_instruction_feature_layers': 1, 'reduce_type': 'mean'},
+ prompt_tuning_configs: Dict[str, Any] = {'ffn_dim_multiplier': None, 'hidden_size': 4096, 'multiple_of': 256, 'norm_eps': 1e-05, 'num_attention_heads': 32, 'num_kv_heads': 8, 'num_layers': 0, 'num_trainable_prompt_tokens': 0, 'use_causal_mask': True, 'use_prompt_tuning': False},
+ timestep_scale: float = 1000.0,
+ ) -> None:
+ """Initialize the Boogu-Image mixed single-double stream transformer model."""
+ super().__init__()
+ self.config = BooguImageDiTConfig(
+ patch_size,
+ in_channels,
+ out_channels,
+ hidden_size,
+ num_layers,
+ num_double_stream_layers,
+ num_refiner_layers,
+ num_attention_heads,
+ num_kv_heads,
+ multiple_of,
+ ffn_dim_multiplier,
+ norm_eps,
+ axes_dim_rope,
+ axes_lens,
+ instruction_feature_configs,
+ prompt_tuning_configs,
+ timestep_scale,
+ )
+
+ # Validate configuration
+ if (hidden_size // num_attention_heads) != sum(axes_dim_rope):
+ raise ValueError(
+ f"hidden_size // num_attention_heads ({hidden_size // num_attention_heads}) "
+ f"must equal sum(axes_dim_rope) ({sum(axes_dim_rope)})"
+ )
+
+ if num_double_stream_layers > num_layers:
+ raise ValueError(
+ f"num_double_stream_layers ({num_double_stream_layers}) cannot be greater than "
+ f"num_layers ({num_layers})"
+ )
+
+ self.out_channels = out_channels or in_channels
+ self.num_double_stream_layers = num_double_stream_layers
+ self.num_single_stream_layers = num_layers - num_double_stream_layers
+ self.instruction_feature_configs = instruction_feature_configs
+ self.prompt_tuning_configs = prompt_tuning_configs
+ self.preprocessed_instruction_feat_dim = (
+ self.cal_preprocessed_instruction_feat_dim(instruction_feature_configs)
+ )
+
+ # Initialize embeddings
+ self.rope_embedder = BooguImageDoubleStreamRotaryPosEmbed(
+ theta=10000,
+ axes_dim=axes_dim_rope,
+ axes_lens=axes_lens,
+ patch_size=patch_size,
+ )
+
+ self.x_embedder = nn.Linear(
+ in_features=patch_size * patch_size * in_channels,
+ out_features=hidden_size,
+ )
+
+ self.ref_image_patch_embedder = nn.Linear(
+ in_features=patch_size * patch_size * in_channels,
+ out_features=hidden_size,
+ )
+
+ self.time_caption_embed = Lumina2CombinedTimestepCaptionEmbedding(
+ hidden_size=hidden_size,
+ instruction_feat_dim=self.preprocessed_instruction_feat_dim,
+ norm_eps=norm_eps,
+ timestep_scale=timestep_scale,
+ )
+
+ # Refiner layers.
+ self.noise_refiner = nn.ModuleList(
+ [
+ BooguImageTransformerBlock(
+ hidden_size,
+ num_attention_heads,
+ num_kv_heads,
+ multiple_of,
+ ffn_dim_multiplier,
+ norm_eps,
+ modulation=True,
+ )
+ for _ in range(num_refiner_layers)
+ ]
+ )
+
+ self.ref_image_refiner = nn.ModuleList(
+ [
+ BooguImageTransformerBlock(
+ hidden_size,
+ num_attention_heads,
+ num_kv_heads,
+ multiple_of,
+ ffn_dim_multiplier,
+ norm_eps,
+ modulation=True,
+ )
+ for _ in range(num_refiner_layers)
+ ]
+ )
+
+ self.context_refiner = nn.ModuleList(
+ [
+ BooguImageTransformerBlock(
+ hidden_size,
+ num_attention_heads,
+ num_kv_heads,
+ multiple_of,
+ ffn_dim_multiplier,
+ norm_eps,
+ modulation=False,
+ )
+ for _ in range(num_refiner_layers)
+ ]
+ )
+
+ # Mixed architecture: dual-stream first, then single-stream.
+ # Here "double-stream" and "dual-stream" mean the same thing.
+ self.double_stream_layers = nn.ModuleList(
+ [
+ BooguImageDoubleStreamTransformerBlock(
+ hidden_size,
+ num_attention_heads,
+ num_kv_heads,
+ multiple_of,
+ ffn_dim_multiplier,
+ norm_eps,
+ modulation=True,
+ )
+ for _ in range(num_double_stream_layers)
+ ]
+ )
+
+ # Single-stream layers process the fused sequence.
+ self.single_stream_layers = nn.ModuleList(
+ [
+ BooguImageSingleStreamTransformerBlock(
+ hidden_size,
+ num_attention_heads,
+ num_kv_heads,
+ multiple_of,
+ ffn_dim_multiplier,
+ norm_eps,
+ modulation=True,
+ )
+ for _ in range(self.num_single_stream_layers)
+ ]
+ )
+
+ # Output norm and projection.
+ self.norm_out = LuminaLayerNormContinuous(
+ embedding_dim=hidden_size,
+ conditioning_embedding_dim=min(hidden_size, 1024),
+ elementwise_affine=False,
+ eps=1e-6,
+ bias=True,
+ out_dim=patch_size * patch_size * self.out_channels,
+ )
+
+ # Distinguish multiple reference images.
+ self.image_index_embedding = nn.Parameter(
+ torch.randn(5, hidden_size)
+ ) # support max 5 ref images
+
+ self.gradient_checkpointing = False
+
+ self.initialize_weights()
+
+ # TeaCache settings
+ self.enable_teacache = False
+ self.enable_taylorseer = False
+ self.enable_teacache_for_all_layers = False
+ self.enable_taylorseer_for_all_layers = False
+ self.teacache_rel_l1_thresh = 0.05
+ self.teacache_params = None
+
+ coefficients = [-5.48259225, 11.48772289, -4.47407401, 2.47730926, -0.03316487]
+ self.rescale_func = np.poly1d(coefficients)
+
+ self.layers = list(self.double_stream_layers) + list(self.single_stream_layers)
+
+ def initialize_weights(self) -> None:
+ """
+ Initialize the weights of the model.
+
+ Uses Xavier uniform initialization for linear layers.
+ """
+ nn.init.xavier_uniform_(self.x_embedder.weight)
+ nn.init.constant_(self.x_embedder.bias, 0.0)
+
+ nn.init.xavier_uniform_(self.ref_image_patch_embedder.weight)
+ nn.init.constant_(self.ref_image_patch_embedder.bias, 0.0)
+
+ nn.init.zeros_(self.norm_out.linear_1.weight)
+ nn.init.zeros_(self.norm_out.linear_1.bias)
+ nn.init.zeros_(self.norm_out.linear_2.weight)
+ nn.init.zeros_(self.norm_out.linear_2.bias)
+
+ nn.init.normal_(self.image_index_embedding, std=0.02)
+
+ def img_patch_embed_and_refine(
+ self,
+ hidden_states,
+ ref_image_hidden_states,
+ padded_img_mask,
+ padded_ref_img_mask,
+ noise_rotary_emb,
+ ref_img_rotary_emb,
+ l_effective_ref_img_len,
+ l_effective_img_len,
+ temb,
+ ):
+ """Embed image patches and run the refiner blocks."""
+ batch_size = len(hidden_states)
+ max_combined_img_len = max(
+ [
+ img_len + sum(ref_img_len)
+ for img_len, ref_img_len in zip(
+ l_effective_img_len, l_effective_ref_img_len
+ )
+ ]
+ )
+
+ hidden_states = self.x_embedder(hidden_states)
+ ref_image_hidden_states = self.ref_image_patch_embedder(ref_image_hidden_states)
+
+ for i in range(batch_size):
+ shift = 0
+ for j, ref_img_len in enumerate(l_effective_ref_img_len[i]):
+ ref_image_hidden_states[i, shift : shift + ref_img_len, :] = (
+ ref_image_hidden_states[i, shift : shift + ref_img_len, :]
+ + self.image_index_embedding[j].to(dtype=ref_image_hidden_states.dtype, device=ref_image_hidden_states.device)
+ )
+ shift += ref_img_len
+
+ for layer in self.noise_refiner:
+ hidden_states = layer(
+ hidden_states, padded_img_mask, noise_rotary_emb, temb
+ )
+
+ flat_l_effective_ref_img_len = list(itertools.chain(*l_effective_ref_img_len))
+ num_ref_images = len(flat_l_effective_ref_img_len)
+ max_ref_img_len = max(flat_l_effective_ref_img_len)
+
+ batch_ref_img_mask = ref_image_hidden_states.new_zeros(
+ num_ref_images, max_ref_img_len, dtype=torch.bool
+ )
+ batch_ref_image_hidden_states = ref_image_hidden_states.new_zeros(
+ num_ref_images, max_ref_img_len, self.config.hidden_size
+ )
+ batch_ref_img_rotary_emb = hidden_states.new_zeros(
+ num_ref_images,
+ max_ref_img_len,
+ ref_img_rotary_emb.shape[-1],
+ dtype=ref_img_rotary_emb.dtype,
+ )
+ batch_temb = temb.new_zeros(num_ref_images, *temb.shape[1:], dtype=temb.dtype)
+
+ # Flatten reference images into a temporary batch.
+ idx = 0
+ for i in range(batch_size):
+ shift = 0
+ for ref_img_len in l_effective_ref_img_len[i]:
+ batch_ref_img_mask[idx, :ref_img_len] = True
+ batch_ref_image_hidden_states[idx, :ref_img_len] = (
+ ref_image_hidden_states[i, shift : shift + ref_img_len]
+ )
+ batch_ref_img_rotary_emb[idx, :ref_img_len] = ref_img_rotary_emb[
+ i, shift : shift + ref_img_len
+ ]
+ batch_temb[idx] = temb[i]
+ shift += ref_img_len
+ idx += 1
+
+ # Refine each reference-image sample.
+ for layer in self.ref_image_refiner:
+ batch_ref_image_hidden_states = layer(
+ batch_ref_image_hidden_states,
+ batch_ref_img_mask,
+ batch_ref_img_rotary_emb,
+ batch_temb,
+ )
+
+ # Restore reference-image sequence layout.
+ idx = 0
+ for i in range(batch_size):
+ shift = 0
+ for ref_img_len in l_effective_ref_img_len[i]:
+ ref_image_hidden_states[i, shift : shift + ref_img_len] = (
+ batch_ref_image_hidden_states[idx, :ref_img_len]
+ )
+ shift += ref_img_len
+ idx += 1
+
+ combined_img_hidden_states = hidden_states.new_zeros(
+ batch_size, max_combined_img_len, self.config.hidden_size
+ )
+ for i, (ref_img_len, img_len) in enumerate(
+ zip(l_effective_ref_img_len, l_effective_img_len)
+ ):
+ combined_img_hidden_states[i, : sum(ref_img_len)] = ref_image_hidden_states[
+ i, : sum(ref_img_len)
+ ]
+ combined_img_hidden_states[
+ i, sum(ref_img_len) : sum(ref_img_len) + img_len
+ ] = hidden_states[i, :img_len]
+
+ return combined_img_hidden_states
+
+ def flat_and_pad_to_seq(self, hidden_states, ref_image_hidden_states):
+ """Flatten patch tokens and pad to batched sequences."""
+ batch_size = len(hidden_states)
+ p = self.config.patch_size
+ device = hidden_states[0].device
+
+ img_sizes = [(img.size(1), img.size(2)) for img in hidden_states]
+ l_effective_img_len = [(H // p) * (W // p) for (H, W) in img_sizes]
+
+ if ref_image_hidden_states is not None:
+ ref_img_sizes = [
+ [(img.size(1), img.size(2)) for img in imgs]
+ if imgs is not None
+ else None
+ for imgs in ref_image_hidden_states
+ ]
+ l_effective_ref_img_len = [
+ [
+ (ref_img_size[0] // p) * (ref_img_size[1] // p)
+ for ref_img_size in _ref_img_sizes
+ ]
+ if _ref_img_sizes is not None
+ else [0]
+ for _ref_img_sizes in ref_img_sizes
+ ]
+ else:
+ ref_img_sizes = [None for _ in range(batch_size)]
+ l_effective_ref_img_len = [[0] for _ in range(batch_size)]
+
+ max_ref_img_len = max(
+ [sum(ref_img_len) for ref_img_len in l_effective_ref_img_len]
+ )
+ max_img_len = max(l_effective_img_len)
+
+ # Reference-image patch embeddings.
+ flat_ref_img_hidden_states = []
+ for i in range(batch_size):
+ if ref_img_sizes[i] is not None:
+ imgs = []
+ for ref_img in ref_image_hidden_states[i]:
+ C, H, W = ref_img.size()
+ ref_img = rearrange(
+ ref_img, "c (h p1) (w p2) -> (h w) (p1 p2 c)", p1=p, p2=p
+ )
+ imgs.append(ref_img)
+
+ img = torch.cat(imgs, dim=0)
+ flat_ref_img_hidden_states.append(img)
+ else:
+ flat_ref_img_hidden_states.append(None)
+
+ # Noise-image patch embeddings.
+ flat_hidden_states = []
+ for i in range(batch_size):
+ img = hidden_states[i]
+ C, H, W = img.size()
+
+ img = rearrange(img, "c (h p1) (w p2) -> (h w) (p1 p2 c)", p1=p, p2=p)
+ flat_hidden_states.append(img)
+
+ padded_ref_img_hidden_states = torch.zeros(
+ batch_size,
+ max_ref_img_len,
+ flat_hidden_states[0].shape[-1],
+ device=device,
+ dtype=flat_hidden_states[0].dtype,
+ )
+ padded_ref_img_mask = torch.zeros(
+ batch_size, max_ref_img_len, dtype=torch.bool, device=device
+ )
+ for i in range(batch_size):
+ if ref_img_sizes[i] is not None:
+ padded_ref_img_hidden_states[i, : sum(l_effective_ref_img_len[i])] = (
+ flat_ref_img_hidden_states[i]
+ )
+ padded_ref_img_mask[i, : sum(l_effective_ref_img_len[i])] = True
+
+ padded_hidden_states = torch.zeros(
+ batch_size,
+ max_img_len,
+ flat_hidden_states[0].shape[-1],
+ device=device,
+ dtype=flat_hidden_states[0].dtype,
+ )
+ padded_img_mask = torch.zeros(
+ batch_size, max_img_len, dtype=torch.bool, device=device
+ )
+ for i in range(batch_size):
+ padded_hidden_states[i, : l_effective_img_len[i]] = flat_hidden_states[i]
+ padded_img_mask[i, : l_effective_img_len[i]] = True
+
+ return (
+ padded_hidden_states,
+ padded_ref_img_hidden_states,
+ padded_img_mask,
+ padded_ref_img_mask,
+ l_effective_ref_img_len,
+ l_effective_img_len,
+ ref_img_sizes,
+ img_sizes,
+ )
+
+ def cal_preprocessed_instruction_feat_dim(
+ self, instruction_feature_configs: Dict[str, Any]
+ ):
+ num_instruction_feat_layers = max(
+ instruction_feature_configs.get("num_instruction_feat_layers", 1), 1
+ )
+ instruction_feat_dim = instruction_feature_configs.get(
+ "instruction_feat_dim", 4096
+ )
+ reduce_type = instruction_feature_configs.get("reduce_type", "concat")
+ if "cat" in reduce_type.lower():
+ return num_instruction_feat_layers * instruction_feat_dim
+ elif "mean" in reduce_type.lower():
+ return instruction_feat_dim
+ else:
+ raise ValueError(f"Invalid reduce_type: {reduce_type}")
+
+ def preprocess_instruction_hidden_states(
+ self, raw_instruction_hidden_states, instruction_feature_configs: Dict[str, Any]
+ ):
+ num_instruction_feat_layers = max(
+ instruction_feature_configs.get("num_instruction_feat_layers", 1), 1
+ )
+ instruction_feat_dim = instruction_feature_configs.get(
+ "instruction_feat_dim", 4096
+ )
+ reduce_type = instruction_feature_configs.get("reduce_type", "concat")
+
+ instruction_hidden_states = None
+ if isinstance(raw_instruction_hidden_states, torch.Tensor):
+ instruction_hidden_states = raw_instruction_hidden_states
+ elif isinstance(raw_instruction_hidden_states, (list, tuple)):
+ assert len(raw_instruction_hidden_states) == num_instruction_feat_layers
+ if "cat" in reduce_type.lower():
+ instruction_hidden_states = torch.cat(
+ raw_instruction_hidden_states, dim=-1
+ )
+ elif "mean" in reduce_type.lower():
+ instruction_hidden_states = torch.mean(
+ torch.stack(raw_instruction_hidden_states), dim=0
+ )
+ else:
+ raise ValueError(f"Invalid reduce_type: {reduce_type}")
+ else:
+ raise ValueError(
+ f"Invalid type of raw_instruction_hidden_states, expected torch.Tensor or list, but got {type(raw_instruction_hidden_states)}"
+ )
+
+ assert (
+ self.preprocessed_instruction_feat_dim
+ == instruction_hidden_states.shape[-1]
+ )
+
+ return instruction_hidden_states
+
+ def forward(
+ self,
+ hidden_states: Union[torch.Tensor, List[torch.Tensor]],
+ timestep: torch.Tensor,
+ instruction_hidden_states: torch.Tensor,
+ freqs_cis: torch.Tensor,
+ instruction_attention_mask: torch.Tensor,
+ ref_image_hidden_states: Optional[List[List[torch.Tensor]]] = None,
+ attention_kwargs: Optional[Dict[str, Any]] = None,
+ return_dict: bool = False,
+ use_gradient_checkpointing: bool = False,
+ use_gradient_checkpointing_offload: bool = False,
+ ) -> Union[torch.Tensor, torch.Tensor]:
+ """
+ Forward pass:
+ context/refiner -> dual-stream (double-stream) -> fusion -> single-stream -> projection.
+ """
+ instruction_hidden_states = self.preprocess_instruction_hidden_states(
+ instruction_hidden_states, self.instruction_feature_configs
+ )
+
+ enable_taylorseer = getattr(self, "enable_taylorseer", False)
+
+ if attention_kwargs is not None:
+ attention_kwargs = attention_kwargs.copy()
+ lora_scale = attention_kwargs.pop("scale", 1.0)
+ else:
+ lora_scale = 1.0
+
+ if False:
+ # weight the lora layers by setting `lora_scale` for each PEFT layer
+ # scale_lora_layers(self, lora_scale)
+ pass
+ else:
+ if (
+ attention_kwargs is not None
+ and attention_kwargs.get("scale", None) is not None
+ ):
+ # logger.warning(
+ # "Passing `scale` via `attention_kwargs` when not using the PEFT backend is ineffective."
+ # )
+ pass
+
+ # === 1. Initial processing (same as original Boogu-Image) ===
+ batch_size = len(hidden_states)
+ is_hidden_states_tensor = isinstance(hidden_states, torch.Tensor)
+
+ if is_hidden_states_tensor:
+ assert hidden_states.ndim == 4
+ hidden_states = [_hidden_states for _hidden_states in hidden_states]
+
+ device = hidden_states[0].device
+
+ # Timestep and instruction embedding.
+ temb, instruction_hidden_states = self.time_caption_embed(
+ timestep, instruction_hidden_states, hidden_states[0].dtype
+ )
+
+ # Flatten and pad token sequences.
+ (
+ hidden_states,
+ ref_image_hidden_states,
+ img_mask,
+ ref_img_mask,
+ l_effective_ref_img_len,
+ l_effective_img_len,
+ ref_img_sizes,
+ img_sizes,
+ ) = self.flat_and_pad_to_seq(hidden_states, ref_image_hidden_states)
+
+ # Build rotary embeddings and sequence lengths.
+ (
+ context_rotary_emb,
+ ref_img_rotary_emb,
+ noise_rotary_emb,
+ rotary_emb,
+ encoder_seq_lengths,
+ seq_lengths,
+ combined_img_rotary_emb,
+ combined_img_seq_lengths,
+ ) = self.rope_embedder(
+ freqs_cis,
+ instruction_attention_mask,
+ l_effective_ref_img_len,
+ l_effective_img_len,
+ ref_img_sizes,
+ img_sizes,
+ device,
+ )
+
+ # Context refinement.
+ for layer in self.context_refiner:
+ instruction_hidden_states = layer(
+ instruction_hidden_states,
+ instruction_attention_mask,
+ context_rotary_emb,
+ )
+
+ # Image patch embedding and refinement.
+ combined_img_hidden_states = self.img_patch_embed_and_refine(
+ hidden_states,
+ ref_image_hidden_states,
+ img_mask,
+ ref_img_mask,
+ noise_rotary_emb,
+ ref_img_rotary_emb,
+ l_effective_ref_img_len,
+ l_effective_img_len,
+ temb,
+ )
+
+ # Dual-stream (double-stream) stage.
+ instruct_hidden_states = instruction_hidden_states
+ img_hidden_states = combined_img_hidden_states
+
+ # Joint mask for [instruct + image].
+ max_seq_len = max(seq_lengths)
+ joint_attention_mask = hidden_states.new_zeros(
+ batch_size, max_seq_len, dtype=torch.bool
+ )
+ for i, seq_len in enumerate(seq_lengths):
+ joint_attention_mask[i, :seq_len] = True
+
+ # Run dual-stream blocks.
+ if self.num_double_stream_layers > 0:
+ # Image-only mask for [ref + noise].
+ max_img_len = max(combined_img_seq_lengths)
+ img_attention_mask = hidden_states.new_zeros(
+ batch_size, max_img_len, dtype=torch.bool
+ )
+ for i, img_seq_len in enumerate(combined_img_seq_lengths):
+ img_attention_mask[i, :img_seq_len] = True
+
+ enable_double_stream_taylorseer = (
+ enable_taylorseer and self.enable_taylorseer_for_all_layers
+ )
+ enable_double_stream_teacache = (
+ self.enable_teacache and self.enable_teacache_for_all_layers
+ )
+
+ if enable_double_stream_teacache:
+ first_double_stream_layer = self.double_stream_layers[0]
+ img_modulated_inp, _, _, _ = first_double_stream_layer.img_norm1(
+ img_hidden_states.clone(), temb
+ )
+ instruct_modulated_inp, _, _, _ = (
+ first_double_stream_layer.instruct_norm1(
+ instruct_hidden_states.clone(), temb
+ )
+ )
+ previous_double_modulated_inp = getattr(
+ self.teacache_params, "previous_double_modulated_inp", None
+ )
+ if (
+ self.teacache_params.is_first_or_last_step
+ or previous_double_modulated_inp is None
+ ):
+ should_calc_double_stream = True
+ self.teacache_params.double_accumulated_rel_l1_distance = 0
+ else:
+ img_rel_l1 = (
+ img_modulated_inp - previous_double_modulated_inp[0]
+ ).abs().mean() / previous_double_modulated_inp[0].abs().mean()
+ instruct_rel_l1 = (
+ instruct_modulated_inp - previous_double_modulated_inp[1]
+ ).abs().mean() / previous_double_modulated_inp[1].abs().mean()
+ rel_l1 = (img_rel_l1 + instruct_rel_l1) * 0.5
+ self.teacache_params.double_accumulated_rel_l1_distance += (
+ self.rescale_func(rel_l1.cpu().item())
+ )
+ if (
+ self.teacache_params.double_accumulated_rel_l1_distance
+ < self.teacache_rel_l1_thresh
+ ):
+ should_calc_double_stream = False
+ else:
+ should_calc_double_stream = True
+ self.teacache_params.double_accumulated_rel_l1_distance = 0
+ self.teacache_params.previous_double_modulated_inp = (
+ img_modulated_inp,
+ instruct_modulated_inp,
+ )
+ else:
+ should_calc_double_stream = True
+
+ if enable_double_stream_teacache and not should_calc_double_stream:
+ img_residual, instruct_residual = (
+ self.teacache_params.previous_double_residual
+ )
+ img_hidden_states = img_hidden_states + img_residual
+ instruct_hidden_states = instruct_hidden_states + instruct_residual
+ else:
+ if enable_double_stream_taylorseer:
+ self.current["stream"] = "double_stream_layers"
+
+ if enable_double_stream_teacache:
+ ori_img_hidden_states = img_hidden_states.clone()
+ ori_instruct_hidden_states = instruct_hidden_states.clone()
+
+ for layer_idx, layer in enumerate(self.double_stream_layers):
+ if enable_double_stream_taylorseer:
+ layer.current = self.current
+ layer.cache_dic = self.cache_dic
+ layer.enable_taylorseer = True
+ self.current["layer"] = layer_idx
+ else:
+ layer.enable_taylorseer = False
+
+ img_hidden_states, instruct_hidden_states = gradient_checkpoint_forward(
+ layer,
+ use_gradient_checkpointing,
+ use_gradient_checkpointing_offload,
+ img_hidden_states,
+ instruct_hidden_states,
+ img_attention_mask,
+ joint_attention_mask,
+ combined_img_rotary_emb,
+ rotary_emb,
+ temb,
+ encoder_seq_lengths,
+ seq_lengths,
+ )
+
+ if enable_double_stream_teacache:
+ self.teacache_params.previous_double_residual = (
+ img_hidden_states - ori_img_hidden_states,
+ instruct_hidden_states - ori_instruct_hidden_states,
+ )
+
+ # Fuse streams to joint sequence.
+ joint_hidden_states = hidden_states.new_zeros(
+ batch_size, max(seq_lengths), self.config.hidden_size
+ )
+ for i, (encoder_seq_len, seq_len) in enumerate(
+ zip(encoder_seq_lengths, seq_lengths)
+ ):
+ joint_hidden_states[i, :encoder_seq_len] = instruct_hidden_states[
+ i, :encoder_seq_len
+ ]
+ joint_hidden_states[i, encoder_seq_len:seq_len] = img_hidden_states[
+ i, : seq_len - encoder_seq_len
+ ]
+
+ # Single-stream stage.
+ hidden_states = joint_hidden_states
+
+ # TeaCache optimization.
+ if self.enable_teacache and len(self.single_stream_layers) > 0:
+ teacache_hidden_states = hidden_states.clone()
+ teacache_temb = temb.clone()
+ modulated_inp, _, _, _ = self.single_stream_layers[0].norm1(
+ teacache_hidden_states, teacache_temb
+ )
+ if self.teacache_params.is_first_or_last_step:
+ should_calc = True
+ self.teacache_params.accumulated_rel_l1_distance = 0
+ else:
+ self.teacache_params.accumulated_rel_l1_distance += self.rescale_func(
+ (
+ (modulated_inp - self.teacache_params.previous_modulated_inp)
+ .abs()
+ .mean()
+ / self.teacache_params.previous_modulated_inp.abs().mean()
+ )
+ .cpu()
+ .item()
+ )
+ if (
+ self.teacache_params.accumulated_rel_l1_distance
+ < self.teacache_rel_l1_thresh
+ ):
+ should_calc = False
+ else:
+ should_calc = True
+ self.teacache_params.accumulated_rel_l1_distance = 0
+ self.teacache_params.previous_modulated_inp = modulated_inp
+ else:
+ should_calc = True
+
+ if self.enable_teacache and not should_calc:
+ hidden_states += self.teacache_params.previous_residual
+ else:
+ if enable_taylorseer:
+ self.current["stream"] = "single_stream_layers"
+
+ if self.enable_teacache:
+ ori_hidden_states = hidden_states.clone()
+
+ for layer_idx, layer in enumerate(self.single_stream_layers):
+ if enable_taylorseer:
+ layer.current = self.current
+ layer.cache_dic = self.cache_dic
+ layer.enable_taylorseer = True
+ self.current["layer"] = self.num_double_stream_layers + layer_idx
+
+ if torch.is_grad_enabled() and use_gradient_checkpointing:
+ hidden_states = gradient_checkpoint_forward(
+ layer,
+ use_gradient_checkpointing,
+ use_gradient_checkpointing_offload,
+ hidden_states,
+ joint_attention_mask,
+ rotary_emb,
+ temb,
+ )
+ else:
+ hidden_states = layer(
+ hidden_states, joint_attention_mask, rotary_emb, temb
+ )
+
+ if self.enable_teacache:
+ self.teacache_params.previous_residual = (
+ hidden_states - ori_hidden_states
+ )
+
+ # Output projection.
+ hidden_states = self.norm_out(hidden_states, temb)
+
+ # Reshape back to image format.
+ p = self.config.patch_size
+ output = []
+ for i, (img_size, img_len, seq_len) in enumerate(
+ zip(img_sizes, l_effective_img_len, seq_lengths)
+ ):
+ height, width = img_size
+ img_tokens = hidden_states[i][seq_len - img_len : seq_len]
+ img_output = rearrange(
+ img_tokens,
+ "(h w) (p1 p2 c) -> c (h p1) (w p2)",
+ h=height // p,
+ w=width // p,
+ p1=p,
+ p2=p,
+ )
+ output.append(img_output)
+
+ if is_hidden_states_tensor:
+ output = torch.stack(output, dim=0)
+
+ # TaylorSeer step counter.
+ if enable_taylorseer:
+ self.current["step"] += 1
+
+ return output
diff --git a/diffsynth/models/clip.py b/diffsynth/models/clip.py
new file mode 100644
index 0000000000000000000000000000000000000000..723107d8c5e2d683537078eee86d0e5db7f46d79
--- /dev/null
+++ b/diffsynth/models/clip.py
@@ -0,0 +1,153 @@
+from typing import Union
+import torch
+from PIL import Image
+from transformers import CLIPModel as HFCLIPModel
+
+ImageInput = Union[Image.Image, list[Image.Image], tuple[Image.Image, ...]]
+
+def _feature_tensor(output, feature_name: str):
+ if torch.is_tensor(output):
+ return output
+ for name in ("image_embeds", "text_embeds", "pooler_output"):
+ value = getattr(output, name, None)
+ if torch.is_tensor(value):
+ return value
+ if isinstance(output, (list, tuple)):
+ for value in output:
+ if torch.is_tensor(value):
+ return value
+ raise TypeError(f"{feature_name} must be a tensor or a model output with projected features.")
+
+
+class ImageMetricsCLIPModel(HFCLIPModel):
+ def __init__(self, variant: str = "h14"):
+ super().__init__(self.config(variant))
+
+ @staticmethod
+ def config(variant: str):
+ from transformers import CLIPConfig
+ return CLIPConfig(
+ projection_dim=1024,
+ logit_scale_init_value=2.6592,
+ text_config={
+ "hidden_size": 1024,
+ "intermediate_size": 4096,
+ "num_attention_heads": 16,
+ "num_hidden_layers": 24,
+ "max_position_embeddings": 77,
+ "vocab_size": 49408,
+ "hidden_act": "quick_gelu",
+ "layer_norm_eps": 1e-5,
+ "projection_dim": 1024,
+ "bos_token_id": 0,
+ "eos_token_id": 2,
+ "pad_token_id": 1,
+ },
+ vision_config={
+ "hidden_size": 1280,
+ "intermediate_size": 5120,
+ "num_attention_heads": 16,
+ "num_hidden_layers": 32,
+ "image_size": 224,
+ "patch_size": 14,
+ "hidden_act": "quick_gelu",
+ "layer_norm_eps": 1e-5,
+ "projection_dim": 1024,
+ },
+ )
+ raise ValueError(f"Unsupported ImageMetrics CLIP variant: {variant}")
+
+
+class CLIPModel(torch.nn.Module):
+ def __init__(self, model: torch.nn.Module, processor, max_length: int = 77):
+ super().__init__()
+ self.model = model
+ self.processor = processor
+ self.max_length = max_length
+
+ @property
+ def device(self):
+ return next(self.parameters(), torch.tensor([])).device
+
+ @property
+ def dtype(self):
+ return next(self.parameters(), torch.tensor(0.0)).dtype
+
+ def _normalize_pairs(self, text, images):
+ if isinstance(text, str):
+ text = [text]
+ else:
+ text = list(text)
+
+ if isinstance(images, Image.Image):
+ images = [images]
+ images = [image.convert("RGB") for image in images]
+
+ if len(text) == 1 and len(images) > 1:
+ text = text * len(images)
+ if len(images) == 1 and len(text) > 1:
+ images = images * len(text)
+
+ if len(text) != len(images):
+ raise ValueError(f"Expected the same number of prompts and images, got {len(text)} and {len(images)}.")
+ return text, images
+
+ def _processor_call(self, **kwargs):
+ inputs = self.processor(
+ padding=True,
+ truncation=True,
+ max_length=self.max_length,
+ return_tensors="pt",
+ **kwargs,
+ ).to(self.device)
+
+ if self.dtype != torch.float32:
+ inputs = {
+ name: (
+ value.to(dtype=self.dtype)
+ if torch.is_tensor(value) and torch.is_floating_point(value)
+ else value
+ )
+ for name, value in inputs.items()
+ }
+ return inputs
+
+ @torch.no_grad()
+ def get_image_features(self, images: ImageInput):
+ if isinstance(images, Image.Image):
+ images = [images]
+ images = [image.convert("RGB") for image in images]
+
+ image_inputs = self._processor_call(images=images)
+ image_features = _feature_tensor(self.model.get_image_features(**image_inputs), "image_features")
+
+ return torch.nn.functional.normalize(image_features, dim=-1)
+
+ @torch.no_grad()
+ def get_text_features(self, text: Union[str, list[str]]):
+ text_inputs = self._processor_call(text=text)
+ text_features = _feature_tensor(self.model.get_text_features(**text_inputs), "text_features")
+
+ return torch.nn.functional.normalize(text_features, dim=-1)
+
+ @torch.no_grad()
+ def similarity_matrix(self, text: Union[str, list[str]], images: ImageInput):
+ image_features = self.get_image_features(images)
+ text_features = self.get_text_features(text)
+
+ scores = text_features @ image_features.T
+ if hasattr(self.model, "logit_scale"):
+ scores = self.model.logit_scale.exp() * scores
+ return scores
+
+ @torch.no_grad()
+ def forward(self, text: Union[str, list[str]], images: ImageInput):
+ text, images = self._normalize_pairs(text, images)
+
+ image_features = self.get_image_features(images)
+ text_features = self.get_text_features(text)
+
+ scores = (text_features * image_features).sum(dim=-1)
+ if hasattr(self.model, "logit_scale"):
+ scores = self.model.logit_scale.exp() * scores
+ return scores
\ No newline at end of file
diff --git a/diffsynth/models/demucs.py b/diffsynth/models/demucs.py
new file mode 100644
index 0000000000000000000000000000000000000000..8e40ce37e237abb27d6c759b1811e0f404decbc5
--- /dev/null
+++ b/diffsynth/models/demucs.py
@@ -0,0 +1,483 @@
+import math
+import torch
+import torch as th
+from einops import rearrange
+from torch import nn
+from torch.nn import functional as F
+import fractions
+import torchaudio
+
+
+def spectro(x, n_fft, hop_length):
+ *other, length = x.shape
+ x = x.reshape(-1, length)
+ z = th.stft(x, n_fft, hop_length, window=th.hann_window(n_fft).to(x),
+ win_length=n_fft, normalized=True, center=True,
+ return_complex=True, pad_mode="reflect")
+ _, freqs, frame = z.shape
+ return z.view(*other, freqs, frame)
+
+
+def ispectro(z, hop_length, length):
+ *other, freqs, frames = z.shape
+ n_fft = 2 * freqs - 2
+ z = z.view(-1, freqs, frames)
+ x = th.istft(z, n_fft, hop_length, window=th.hann_window(n_fft).to(z.real),
+ win_length=n_fft, normalized=True, length=length, center=True)
+ return x.view(*other, x.shape[-1])
+
+
+class LayerScale(nn.Module):
+ def __init__(self, channels, init=0, channel_last=False):
+ super().__init__()
+ self.channel_last = channel_last
+ self.scale = nn.Parameter(torch.full((channels,), float(init)))
+
+ def forward(self, x):
+ return self.scale * x if self.channel_last else self.scale[:, None] * x
+
+
+class DConv(nn.Module):
+ def __init__(self, channels, compress=8, depth=2, init=1e-3, kernel=3):
+ super().__init__()
+ hidden = int(channels / compress)
+ self.layers = nn.ModuleList()
+ for d in range(depth):
+ dilation = 2 ** d
+ padding = dilation * (kernel // 2)
+ self.layers.append(nn.Sequential(
+ nn.Conv1d(channels, hidden, kernel, dilation=dilation, padding=padding),
+ nn.GroupNorm(1, hidden), nn.GELU(),
+ nn.Conv1d(hidden, 2 * channels, 1),
+ nn.GroupNorm(1, 2 * channels), nn.GLU(1),
+ LayerScale(channels, init),
+ ))
+
+ def forward(self, x):
+ for layer in self.layers:
+ x = x + layer(x)
+ return x
+
+
+class ScaledEmbedding(nn.Module):
+ def __init__(self, num_embeddings, embedding_dim, scale=10.0):
+ super().__init__()
+ self.embedding = nn.Embedding(num_embeddings, embedding_dim)
+ weight = torch.cumsum(self.embedding.weight.data, dim=0)
+ weight = weight / torch.arange(1, num_embeddings + 1).to(weight).sqrt()[:, None]
+ self.embedding.weight.data[:] = weight / scale
+ self.scale = scale
+
+ def forward(self, x):
+ return self.embedding(x) * self.scale
+
+
+class HEncLayer(nn.Module):
+ def __init__(self, chin, chout, kernel_size=8, stride=4, freq=True, context=0,
+ dconv_kw={}):
+ super().__init__()
+ self.freq = freq
+ self.stride = stride
+ pad = kernel_size // 4
+ klass = nn.Conv2d if freq else nn.Conv1d
+ if freq:
+ kernel_size, stride, pad = [kernel_size, 1], [stride, 1], [pad, 0]
+ self.conv = klass(chin, chout, kernel_size, stride, pad)
+ self.rewrite = klass(chout, 2 * chout, 1 + 2 * context, 1, context)
+ self.dconv = DConv(chout, **dconv_kw)
+
+ def forward(self, x):
+ if not self.freq:
+ if x.dim() == 4:
+ B, C, Fr, T = x.shape
+ x = x.view(B, -1, T)
+ le = x.shape[-1]
+ if le % self.stride:
+ x = F.pad(x, (0, self.stride - (le % self.stride)))
+ y = F.gelu(self.conv(x))
+ if self.freq:
+ B, C, Fr, T = y.shape
+ y = y.permute(0, 2, 1, 3).reshape(-1, C, T)
+ y = self.dconv(y).view(B, Fr, C, T).permute(0, 2, 1, 3)
+ else:
+ y = self.dconv(y)
+ return F.glu(self.rewrite(y), dim=1)
+
+
+class HDecLayer(nn.Module):
+ def __init__(self, chin, chout, last=False, kernel_size=8, stride=4, freq=True,
+ context=1, dconv_kw={}):
+ super().__init__()
+ pad = kernel_size // 4
+ self.pad = pad
+ self.last = last
+ self.freq = freq
+ self.chin = chin
+ klass = nn.Conv2d if freq else nn.Conv1d
+ klass_tr = nn.ConvTranspose2d if freq else nn.ConvTranspose1d
+ ks, st = ([kernel_size, 1], [stride, 1]) if freq else (kernel_size, stride)
+ self.conv_tr = klass_tr(chin, chout, ks, st)
+ self.rewrite = klass(chin, 2 * chin, 1 + 2 * context, 1, context)
+ self.dconv = DConv(chin, **dconv_kw)
+
+ def forward(self, x, skip, length):
+ if self.freq and x.dim() == 3:
+ B, C, T = x.shape
+ x = x.view(B, self.chin, -1, T)
+ x = x + skip
+ y = F.glu(self.rewrite(x), dim=1)
+ if self.freq:
+ B, C, Fr, T = y.shape
+ y = y.permute(0, 2, 1, 3).reshape(-1, C, T)
+ y = self.dconv(y).view(B, Fr, C, T).permute(0, 2, 1, 3)
+ else:
+ y = self.dconv(y)
+ z = self.conv_tr(y)
+ if self.freq:
+ z = z[..., self.pad: -self.pad, :]
+ else:
+ z = z[..., self.pad: self.pad + length]
+ if not self.last:
+ z = F.gelu(z)
+ return z, y
+
+
+def create_sin_embedding(length, dim, device, max_period=10000):
+ pos = torch.arange(length, device=device).view(-1, 1, 1)
+ half_dim = dim // 2
+ adim = torch.arange(half_dim, device=device).view(1, 1, -1)
+ phase = pos / (max_period ** (adim / (half_dim - 1)))
+ return torch.cat([torch.cos(phase), torch.sin(phase)], dim=-1)
+
+
+def create_2d_sin_embedding(d_model, height, width, device, max_period=10000):
+ pe = torch.zeros(d_model, height, width)
+ d_model = d_model // 2
+ div_term = torch.exp(torch.arange(0.0, d_model, 2) * -(math.log(max_period) / d_model))
+ pos_w = torch.arange(0.0, width).unsqueeze(1)
+ pos_h = torch.arange(0.0, height).unsqueeze(1)
+ pe[0:d_model:2, :, :] = torch.sin(pos_w * div_term).transpose(0, 1).unsqueeze(1).repeat(1, height, 1)
+ pe[1:d_model:2, :, :] = torch.cos(pos_w * div_term).transpose(0, 1).unsqueeze(1).repeat(1, height, 1)
+ pe[d_model::2, :, :] = torch.sin(pos_h * div_term).transpose(0, 1).unsqueeze(2).repeat(1, 1, width)
+ pe[d_model + 1:: 2, :, :] = torch.cos(pos_h * div_term).transpose(0, 1).unsqueeze(2).repeat(1, 1, width)
+ return pe[None, :].to(device)
+
+
+class MyGroupNorm(nn.GroupNorm):
+ def forward(self, x):
+ return super().forward(x.transpose(1, 2)).transpose(1, 2)
+
+
+class MyTransformerEncoderLayer(nn.TransformerEncoderLayer):
+ def __init__(self, d_model, nhead, dim_feedforward, dropout, activation):
+ super().__init__(d_model=d_model, nhead=nhead, dim_feedforward=dim_feedforward,
+ dropout=dropout, activation=activation, batch_first=True,
+ norm_first=True)
+ self.norm_out = MyGroupNorm(1, d_model)
+ self.gamma_1 = LayerScale(d_model, 1e-4, True)
+ self.gamma_2 = LayerScale(d_model, 1e-4, True)
+
+ def forward(self, x, **unused):
+ x = x + self.gamma_1(self._sa_block(self.norm1(x), None, None))
+ x = x + self.gamma_2(self._ff_block(self.norm2(x)))
+ return self.norm_out(x)
+
+
+class CrossTransformerEncoderLayer(nn.Module):
+ def __init__(self, d_model, nhead, dim_feedforward, dropout, activation):
+ super().__init__()
+ self.cross_attn = nn.MultiheadAttention(d_model, nhead, dropout=dropout, batch_first=True)
+ self.linear1 = nn.Linear(d_model, dim_feedforward)
+ self.dropout = nn.Dropout(dropout)
+ self.linear2 = nn.Linear(dim_feedforward, d_model)
+ self.norm1 = nn.LayerNorm(d_model)
+ self.norm2 = nn.LayerNorm(d_model)
+ self.norm3 = nn.LayerNorm(d_model)
+ self.norm_out = MyGroupNorm(1, d_model)
+ self.gamma_1 = LayerScale(d_model, 1e-4, True)
+ self.gamma_2 = LayerScale(d_model, 1e-4, True)
+ self.dropout1 = nn.Dropout(dropout)
+ self.dropout2 = nn.Dropout(dropout)
+ self.activation = activation
+
+ def forward(self, q, k):
+ x = q + self.gamma_1(self.dropout1(
+ self.cross_attn(self.norm1(q), self.norm2(k), self.norm2(k), need_weights=False)[0]))
+ x = x + self.gamma_2(self._ff_block(self.norm3(x)))
+ return self.norm_out(x)
+
+ def _ff_block(self, x):
+ return self.dropout2(self.linear2(self.dropout(self.activation(self.linear1(x)))))
+
+
+class CrossTransformerEncoder(nn.Module):
+ def __init__(self, dim, hidden_scale=4.0, num_heads=8, num_layers=5, dropout=0.0,
+ max_period=10000.0, gelu=True):
+ super().__init__()
+ hidden_dim = int(dim * hidden_scale)
+ self.num_layers = num_layers
+ self.max_period = max_period
+ self.norm_in = nn.LayerNorm(dim)
+ self.norm_in_t = nn.LayerNorm(dim)
+ activation = F.gelu if gelu else F.relu
+ self.layers = nn.ModuleList()
+ self.layers_t = nn.ModuleList()
+ common = dict(d_model=dim, nhead=num_heads, dim_feedforward=hidden_dim,
+ dropout=dropout, activation=activation)
+ for idx in range(num_layers):
+ klass = MyTransformerEncoderLayer if idx % 2 == 0 else CrossTransformerEncoderLayer
+ self.layers.append(klass(**common))
+ self.layers_t.append(klass(**common))
+
+ def forward(self, x, xt):
+ B, C, Fr, T1 = x.shape
+ pos_2d = rearrange(create_2d_sin_embedding(C, Fr, T1, x.device, self.max_period),
+ "b c fr t1 -> b (t1 fr) c")
+ x = (self.norm_in(rearrange(x, "b c fr t1 -> b (t1 fr) c")) + pos_2d).to(x.dtype)
+
+ B, C, T2 = xt.shape
+ pos = rearrange(create_sin_embedding(T2, C, x.device, self.max_period), "t2 b c -> b t2 c")
+ xt = (self.norm_in_t(rearrange(xt, "b c t2 -> b t2 c")) + pos).to(xt.dtype)
+
+ for idx in range(self.num_layers):
+ if idx % 2 == 0:
+ x = self.layers[idx](x)
+ xt = self.layers_t[idx](xt)
+ else:
+ old_x = x
+ x = self.layers[idx](x, xt)
+ xt = self.layers_t[idx](xt, old_x)
+
+ x = rearrange(x, "b (t1 fr) c -> b c fr t1", t1=T1)
+ xt = rearrange(xt, "b t2 c -> b c t2")
+ return x, xt
+
+
+class HTDemucs(nn.Module):
+ def __init__(self, sources=['drums', 'bass', 'other', 'vocals'], audio_channels=2, channels=48, growth=2, nfft=4096, cac=True,
+ depth=4, freq_emb=0.2, emb_scale=10, kernel_size=8, time_stride=2, stride=4,
+ context=1, context_enc=0, dconv_depth=2, dconv_comp=8, dconv_init=1e-3,
+ bottom_channels=512, t_layers=5, t_hidden_scale=4.0, t_heads=8, t_dropout=0.0,
+ t_max_period=10000.0, t_gelu=True, samplerate=44100, segment=fractions.Fraction(39, 5), **unused):
+ super().__init__()
+ self.audio_channels = audio_channels
+ self.sources = sources
+ self.depth = depth
+ self.bottom_channels = bottom_channels
+ self.samplerate = samplerate
+ self.segment = segment
+ self.nfft = nfft
+ self.hop_length = nfft // 4
+
+ self.encoder = nn.ModuleList()
+ self.decoder = nn.ModuleList()
+ self.tencoder = nn.ModuleList()
+ self.tdecoder = nn.ModuleList()
+
+ chin = audio_channels
+ chin_z = chin * 2
+ chout = chout_z = channels
+ freqs = nfft // 2
+
+ for index in range(depth):
+ dconv_kw = {"depth": dconv_depth, "compress": dconv_comp, "init": dconv_init}
+ kw = {"kernel_size": kernel_size, "stride": stride, "dconv_kw": dconv_kw}
+ kwt = dict(kw, freq=False)
+
+ self.encoder.append(HEncLayer(chin_z, chout_z, context=context_enc, **kw))
+ self.tencoder.append(HEncLayer(chin, chout, freq=False, context=context_enc, **kw))
+
+ if index == 0:
+ chin = self.audio_channels * len(self.sources)
+ chin_z = chin * 2
+ self.decoder.insert(0, HDecLayer(chout_z, chin_z, last=index == 0, context=context, **kw))
+ self.tdecoder.insert(0, HDecLayer(chout, chin, last=index == 0, context=context, **kwt))
+
+ chin, chin_z = chout, chout_z
+ chout = int(growth * chout)
+ chout_z = int(growth * chout_z)
+ freqs //= stride
+ if index == 0:
+ self.freq_emb = ScaledEmbedding(freqs, chin_z, scale=emb_scale)
+ self.freq_emb_scale = freq_emb
+
+ transformer_channels = channels * growth ** (depth - 1)
+ self.channel_upsampler = nn.Conv1d(transformer_channels, bottom_channels, 1)
+ self.channel_downsampler = nn.Conv1d(bottom_channels, transformer_channels, 1)
+ self.channel_upsampler_t = nn.Conv1d(transformer_channels, bottom_channels, 1)
+ self.channel_downsampler_t = nn.Conv1d(bottom_channels, transformer_channels, 1)
+
+ self.crosstransformer = CrossTransformerEncoder(
+ dim=bottom_channels, hidden_scale=t_hidden_scale, num_heads=t_heads,
+ num_layers=t_layers, dropout=t_dropout, max_period=t_max_period, gelu=t_gelu)
+
+ def _spec(self, x):
+ hl, nfft = self.hop_length, self.nfft
+ le = int(math.ceil(x.shape[-1] / hl))
+ pad = hl // 2 * 3
+ x = F.pad(x, (pad, pad + le * hl - x.shape[-1]), mode="reflect")
+ z = spectro(x.float(), nfft, hl)[..., :-1, :]
+ return z[..., 2: 2 + le]
+
+ def _ispec(self, z, length):
+ hl = self.hop_length
+ z = F.pad(z, (0, 0, 0, 1))
+ z = F.pad(z, (2, 2))
+ pad = hl // 2 * 3
+ le = hl * int(math.ceil(length / hl)) + 2 * pad
+ return ispectro(z, hl, le)[..., pad: pad + length]
+
+ def _magnitude(self, z):
+ # cac: 复数实虚部当作额外通道
+ B, C, Fr, T = z.shape
+ return torch.view_as_real(z).permute(0, 1, 4, 2, 3).reshape(B, C * 2, Fr, T)
+
+ def _mask(self, m):
+ # cac: m 即完整复数谱(通道形式), 还原成复数
+ B, S, C, Fr, T = m.shape
+ out = m.view(B, S, -1, 2, Fr, T).permute(0, 1, 2, 4, 5, 3)
+ return torch.view_as_complex(out.contiguous())
+
+ def forward(self, mix):
+ length = mix.shape[-1]
+ z = self._spec(mix)
+ x = self._magnitude(z).to(device=mix.device, dtype=mix.dtype)
+
+ B, C, Fq, T = x.shape
+ mean = x.mean(dim=(1, 2, 3), keepdim=True)
+ std = x.std(dim=(1, 2, 3), keepdim=True)
+ x = (x - mean) / (1e-5 + std)
+
+ xt = mix
+ meant = xt.mean(dim=(1, 2), keepdim=True)
+ stdt = xt.std(dim=(1, 2), keepdim=True)
+ xt = (xt - meant) / (1e-5 + stdt)
+
+ saved, saved_t, lengths, lengths_t = [], [], [], []
+ for idx, encode in enumerate(self.encoder):
+ lengths.append(x.shape[-1])
+ lengths_t.append(xt.shape[-1])
+ xt = self.tencoder[idx](xt)
+ saved_t.append(xt)
+ x = encode(x)
+ if idx == 0:
+ frs = torch.arange(x.shape[-2], device=x.device)
+ emb = self.freq_emb(frs).t()[None, :, :, None].expand_as(x)
+ x = x + self.freq_emb_scale * emb
+ saved.append(x)
+
+ b, c, f, t = x.shape
+ x = rearrange(x, "b c f t-> b c (f t)")
+ x = self.channel_upsampler(x)
+ x = rearrange(x, "b c (f t)-> b c f t", f=f)
+ xt = self.channel_upsampler_t(xt)
+
+ x, xt = self.crosstransformer(x, xt)
+
+ x = rearrange(x, "b c f t-> b c (f t)")
+ x = self.channel_downsampler(x)
+ x = rearrange(x, "b c (f t)-> b c f t", f=f)
+ xt = self.channel_downsampler_t(xt)
+
+ for idx, decode in enumerate(self.decoder):
+ x, pre = decode(x, saved.pop(-1), lengths.pop(-1))
+ xt, _ = self.tdecoder[idx](xt, saved_t.pop(-1), lengths_t.pop(-1))
+
+ S = len(self.sources)
+ x = x.view(B, S, -1, Fq, T) * std[:, None] + mean[:, None]
+ x = self._ispec(self._mask(x.float()), length).to(mix.dtype)
+
+ xt = xt.view(B, S, -1, length) * stdt[:, None] + meant[:, None]
+ return xt + x
+
+ def extract_track(self, audio, sample_rate, track="vocals"):
+ if sample_rate != 44100:
+ audio = torchaudio.functional.resample(audio, sample_rate, 44100)
+ ref = audio.mean(0)
+ audio = (audio - ref.mean()) / (ref.std() + 1e-8)
+ out = apply_model(self, audio[None], shifts=1, split=True, overlap=0.25,
+ device=audio.device, progress=True)[0]
+ out = out * (ref.std() + 1e-8) + ref.mean()
+ out = out / max(1.01 * out.abs().max(), 1)
+ out = out.clamp_(-1, 1).cpu()
+ out = out[self.sources.index(track)]
+ return out
+
+def center_trim(tensor, length):
+ delta = tensor.size(-1) - length
+ if delta:
+ tensor = tensor[..., delta // 2: -(delta - delta // 2)]
+ return tensor
+
+
+class TensorChunk:
+ def __init__(self, tensor, offset=0, length=None):
+ total = tensor.shape[-1]
+ if isinstance(tensor, TensorChunk):
+ self.tensor, self.offset = tensor.tensor, offset + tensor.offset
+ else:
+ self.tensor, self.offset = tensor, offset
+ self.length = total - offset if length is None else min(total - offset, length)
+ self.device = tensor.device
+
+ @property
+ def shape(self):
+ shape = list(self.tensor.shape)
+ shape[-1] = self.length
+ return shape
+
+ def padded(self, target_length):
+ delta = target_length - self.length
+ total = self.tensor.shape[-1]
+ start = self.offset - delta // 2
+ end = start + target_length
+ correct_start, correct_end = max(0, start), min(total, end)
+ out = F.pad(self.tensor[..., correct_start:correct_end],
+ (correct_start - start, end - correct_end))
+ return out
+
+
+def apply_model(model, mix, shifts=1, split=True, overlap=0.25, device=None, segment=None,
+ progress=False):
+ device = th.device(device) if device is not None else mix.device
+ batch, channels, length = mix.shape
+
+ if shifts:
+ max_shift = int(0.5 * model.samplerate)
+ mix = mix if isinstance(mix, TensorChunk) else TensorChunk(mix)
+ padded_mix = mix.padded(length + 2 * max_shift)
+ out = 0.0
+ for _ in range(shifts):
+ offset = int(th.randint(0, max_shift, (1,)).item())
+ shifted = TensorChunk(padded_mix, offset, length + max_shift - offset)
+ shifted_out = apply_model(model, shifted, shifts=0, split=split, overlap=overlap,
+ device=device, segment=segment, progress=progress)
+ out += shifted_out[..., max_shift - offset:]
+ return out / shifts
+
+ if split:
+ out = th.zeros(batch, len(model.sources), channels, length, device=mix.device)
+ sum_weight = th.zeros(length, device=mix.device)
+ if segment is None:
+ segment = model.segment
+ segment_length = int(model.samplerate * segment)
+ stride = int((1 - overlap) * segment_length)
+ offsets = range(0, length, stride)
+ scale = float(format(stride / model.samplerate, ".2f"))
+ weight = th.cat([th.arange(1, segment_length // 2 + 1, device=device),
+ th.arange(segment_length - segment_length // 2, 0, -1, device=device)])
+ weight = weight / weight.max()
+ iterator = offsets
+ for offset in iterator:
+ chunk = TensorChunk(mix, offset, segment_length)
+ chunk_out = apply_model(model, chunk, shifts=0, split=False, device=device, segment=segment)
+ cl = chunk_out.shape[-1]
+ out[..., offset:offset + segment_length] += (weight[:cl] * chunk_out).to(mix.device)
+ sum_weight[offset:offset + segment_length] += weight[:cl].to(mix.device)
+ return out / sum_weight
+
+ valid_length = int(segment * model.samplerate)
+ mix = mix if isinstance(mix, TensorChunk) else TensorChunk(mix)
+ padded_mix = mix.padded(valid_length).to(device)
+ out = model(padded_mix)
+ return center_trim(out, length)
diff --git a/diffsynth/models/dinov3_image_encoder.py b/diffsynth/models/dinov3_image_encoder.py
new file mode 100644
index 0000000000000000000000000000000000000000..66ee4168a52139ba0e88e1cd64f7bb0fdd807545
--- /dev/null
+++ b/diffsynth/models/dinov3_image_encoder.py
@@ -0,0 +1,100 @@
+import torch, warnings
+try:
+ from transformers.models.dinov3_vit.modeling_dinov3_vit import DINOv3ViTModel
+except:
+ warnings.warn(f"Cannot import `DINOv3ViTModel`. `DINOv3ImageEncoder` is not available. Please update `transformers` by `pip install -U transformers`.")
+ DINOv3ViTModel = torch.nn.Module
+from ..core.device.npu_compatible_device import get_device_type
+
+
+class DINOv3ImageEncoder(DINOv3ViTModel):
+ def __init__(self):
+ from transformers.models.dinov3_vit.modeling_dinov3_vit import DINOv3ViTConfig
+ from transformers import DINOv3ViTImageProcessor
+ config = DINOv3ViTConfig(
+ architectures = [
+ "DINOv3ViTModel"
+ ],
+ attention_dropout = 0.0,
+ drop_path_rate = 0.0,
+ dtype = "float32",
+ hidden_act = "silu",
+ hidden_size = 4096,
+ image_size = 224,
+ initializer_range = 0.02,
+ intermediate_size = 8192,
+ key_bias = False,
+ layer_norm_eps = 1e-05,
+ layerscale_value = 1.0,
+ mlp_bias = True,
+ model_type = "dinov3_vit",
+ num_attention_heads = 32,
+ num_channels = 3,
+ num_hidden_layers = 40,
+ num_register_tokens = 4,
+ patch_size = 16,
+ pos_embed_jitter = None,
+ pos_embed_rescale = 2.0,
+ pos_embed_shift = None,
+ proj_bias = True,
+ query_bias = False,
+ rope_theta = 100.0,
+ transformers_version = "4.56.1",
+ use_gated_mlp = True,
+ value_bias = False
+ )
+ super().__init__(config)
+ self.processor = DINOv3ViTImageProcessor(
+ crop_size = None,
+ data_format = "channels_first",
+ default_to_square = True,
+ device = None,
+ disable_grouping = None,
+ do_center_crop = None,
+ do_convert_rgb = None,
+ do_normalize = True,
+ do_rescale = True,
+ do_resize = True,
+ image_mean = [
+ 0.485,
+ 0.456,
+ 0.406
+ ],
+ image_processor_type = "DINOv3ViTImageProcessor",
+ image_std = [
+ 0.229,
+ 0.224,
+ 0.225
+ ],
+ input_data_format = None,
+ resample = 2,
+ rescale_factor = 0.00392156862745098,
+ return_tensors = None,
+ size = {
+ "height": 224,
+ "width": 224
+ }
+ )
+
+ def forward(self, image, torch_dtype=torch.bfloat16, device=get_device_type()):
+ inputs = self.processor(images=image, return_tensors="pt")
+ pixel_values = inputs["pixel_values"].to(dtype=torch_dtype, device=device)
+ bool_masked_pos = None
+ head_mask = None
+
+ pixel_values = pixel_values.to(torch_dtype)
+ hidden_states = self.embeddings(pixel_values, bool_masked_pos=bool_masked_pos)
+ position_embeddings = self.rope_embeddings(pixel_values)
+
+ for i, layer_module in enumerate(self.model.layer):
+ layer_head_mask = head_mask[i] if head_mask is not None else None
+ hidden_states = layer_module(
+ hidden_states,
+ attention_mask=layer_head_mask,
+ position_embeddings=position_embeddings,
+ )
+
+ sequence_output = self.norm(hidden_states)
+ pooled_output = sequence_output[:, 0, :]
+
+ return pooled_output
diff --git a/diffsynth/models/ernie_image_dit.py b/diffsynth/models/ernie_image_dit.py
new file mode 100644
index 0000000000000000000000000000000000000000..1becf78923961be7483afe644874a2545cd740fd
--- /dev/null
+++ b/diffsynth/models/ernie_image_dit.py
@@ -0,0 +1,362 @@
+"""
+Ernie-Image DiT for DiffSynth-Studio.
+
+Refactored from diffusers ErnieImageTransformer2DModel to use DiffSynth core modules.
+Default parameters from actual checkpoint config.json (PaddlePaddle/ERNIE-Image transformer).
+"""
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from typing import Optional, Tuple
+
+from ..core.attention import attention_forward
+from ..core.gradient import gradient_checkpoint_forward
+from .flux2_dit import Timesteps, TimestepEmbedding
+
+
+def rope(pos: torch.Tensor, dim: int, theta: int) -> torch.Tensor:
+ assert dim % 2 == 0
+ scale = torch.arange(0, dim, 2, dtype=torch.float64, device=pos.device) / dim
+ omega = 1.0 / (theta ** scale)
+ out = torch.einsum("...n,d->...nd", pos, omega)
+ return out.float()
+
+
+class ErnieImageEmbedND3(nn.Module):
+ def __init__(self, dim: int, theta: int, axes_dim: Tuple[int, int, int]):
+ super().__init__()
+ self.dim = dim
+ self.theta = theta
+ self.axes_dim = list(axes_dim)
+
+ def forward(self, ids: torch.Tensor) -> torch.Tensor:
+ emb = torch.cat([rope(ids[..., i], self.axes_dim[i], self.theta) for i in range(3)], dim=-1)
+ emb = emb.unsqueeze(2)
+ return torch.stack([emb, emb], dim=-1).reshape(*emb.shape[:-1], -1)
+
+
+class ErnieImagePatchEmbedDynamic(nn.Module):
+ def __init__(self, in_channels: int, embed_dim: int, patch_size: int):
+ super().__init__()
+ self.patch_size = patch_size
+ self.proj = nn.Conv2d(in_channels, embed_dim, kernel_size=patch_size, stride=patch_size, bias=True)
+
+ def forward(self, x: torch.Tensor) -> torch.Tensor:
+ x = self.proj(x)
+ batch_size, dim, height, width = x.shape
+ return x.reshape(batch_size, dim, height * width).transpose(1, 2).contiguous()
+
+
+class ErnieImageSingleStreamAttnProcessor:
+ def __call__(
+ self,
+ attn: "ErnieImageAttention",
+ hidden_states: torch.Tensor,
+ attention_mask: Optional[torch.Tensor] = None,
+ freqs_cis: Optional[torch.Tensor] = None,
+ ) -> torch.Tensor:
+ query = attn.to_q(hidden_states)
+ key = attn.to_k(hidden_states)
+ value = attn.to_v(hidden_states)
+
+ query = query.unflatten(-1, (attn.heads, -1))
+ key = key.unflatten(-1, (attn.heads, -1))
+ value = value.unflatten(-1, (attn.heads, -1))
+
+ if attn.norm_q is not None:
+ query = attn.norm_q(query)
+ if attn.norm_k is not None:
+ key = attn.norm_k(key)
+
+ def apply_rotary_emb(x_in: torch.Tensor, freqs_cis: torch.Tensor) -> torch.Tensor:
+ rot_dim = freqs_cis.shape[-1]
+ x, x_pass = x_in[..., :rot_dim], x_in[..., rot_dim:]
+ cos_ = torch.cos(freqs_cis).to(x.dtype)
+ sin_ = torch.sin(freqs_cis).to(x.dtype)
+ x1, x2 = x.chunk(2, dim=-1)
+ x_rotated = torch.cat((-x2, x1), dim=-1)
+ return torch.cat((x * cos_ + x_rotated * sin_, x_pass), dim=-1)
+
+ if freqs_cis is not None:
+ query = apply_rotary_emb(query, freqs_cis)
+ key = apply_rotary_emb(key, freqs_cis)
+
+ if attention_mask is not None and attention_mask.ndim == 2:
+ attention_mask = attention_mask[:, None, None, :]
+
+ hidden_states = attention_forward(
+ query, key, value,
+ q_pattern="b s n d",
+ k_pattern="b s n d",
+ v_pattern="b s n d",
+ out_pattern="b s n d",
+ attn_mask=attention_mask,
+ )
+
+ hidden_states = hidden_states.flatten(2, 3)
+ hidden_states = hidden_states.to(query.dtype)
+ output = attn.to_out[0](hidden_states)
+
+ return output
+
+
+class ErnieImageAttention(nn.Module):
+ def __init__(
+ self,
+ query_dim: int,
+ heads: int = 8,
+ dim_head: int = 64,
+ dropout: float = 0.0,
+ bias: bool = False,
+ qk_norm: str = "rms_norm",
+ out_bias: bool = True,
+ eps: float = 1e-5,
+ out_dim: int = None,
+ elementwise_affine: bool = True,
+ ):
+ super().__init__()
+
+ self.head_dim = dim_head
+ self.inner_dim = out_dim if out_dim is not None else dim_head * heads
+ self.query_dim = query_dim
+ self.out_dim = out_dim if out_dim is not None else query_dim
+ self.heads = out_dim // dim_head if out_dim is not None else heads
+
+ self.use_bias = bias
+ self.dropout = dropout
+
+ self.to_q = nn.Linear(query_dim, self.inner_dim, bias=bias)
+ self.to_k = nn.Linear(query_dim, self.inner_dim, bias=bias)
+ self.to_v = nn.Linear(query_dim, self.inner_dim, bias=bias)
+
+ if qk_norm == "layer_norm":
+ self.norm_q = nn.LayerNorm(dim_head, eps=eps, elementwise_affine=elementwise_affine)
+ self.norm_k = nn.LayerNorm(dim_head, eps=eps, elementwise_affine=elementwise_affine)
+ elif qk_norm == "rms_norm":
+ self.norm_q = nn.RMSNorm(dim_head, eps=eps, elementwise_affine=elementwise_affine)
+ self.norm_k = nn.RMSNorm(dim_head, eps=eps, elementwise_affine=elementwise_affine)
+ else:
+ raise ValueError(
+ f"unknown qk_norm: {qk_norm}. Should be one of None, 'layer_norm', 'rms_norm'."
+ )
+
+ self.to_out = nn.ModuleList([])
+ self.to_out.append(nn.Linear(self.inner_dim, self.out_dim, bias=out_bias))
+
+ self.processor = ErnieImageSingleStreamAttnProcessor()
+
+ def forward(
+ self,
+ hidden_states: torch.Tensor,
+ attention_mask: Optional[torch.Tensor] = None,
+ image_rotary_emb: Optional[torch.Tensor] = None,
+ ) -> torch.Tensor:
+ return self.processor(self, hidden_states, attention_mask, image_rotary_emb)
+
+
+class ErnieImageFeedForward(nn.Module):
+ def __init__(self, hidden_size: int, ffn_hidden_size: int):
+ super().__init__()
+ self.gate_proj = nn.Linear(hidden_size, ffn_hidden_size, bias=False)
+ self.up_proj = nn.Linear(hidden_size, ffn_hidden_size, bias=False)
+ self.linear_fc2 = nn.Linear(ffn_hidden_size, hidden_size, bias=False)
+
+ def forward(self, x: torch.Tensor) -> torch.Tensor:
+ return self.linear_fc2(self.up_proj(x) * F.gelu(self.gate_proj(x)))
+
+
+class ErnieImageRMSNorm(nn.Module):
+ def __init__(self, dim: int, eps: float = 1e-6):
+ super().__init__()
+ self.eps = eps
+ self.weight = nn.Parameter(torch.ones(dim))
+
+ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+ input_dtype = hidden_states.dtype
+ variance = hidden_states.to(torch.float32).pow(2).mean(-1, keepdim=True)
+ hidden_states = hidden_states * torch.rsqrt(variance + self.eps)
+ hidden_states = hidden_states * self.weight
+ return hidden_states.to(input_dtype)
+
+
+class ErnieImageSharedAdaLNBlock(nn.Module):
+ def __init__(
+ self,
+ hidden_size: int,
+ num_heads: int,
+ ffn_hidden_size: int,
+ eps: float = 1e-6,
+ qk_layernorm: bool = True,
+ ):
+ super().__init__()
+ self.adaLN_sa_ln = ErnieImageRMSNorm(hidden_size, eps=eps)
+ self.self_attention = ErnieImageAttention(
+ query_dim=hidden_size,
+ dim_head=hidden_size // num_heads,
+ heads=num_heads,
+ qk_norm="rms_norm" if qk_layernorm else None,
+ eps=eps,
+ bias=False,
+ out_bias=False,
+ )
+ self.adaLN_mlp_ln = ErnieImageRMSNorm(hidden_size, eps=eps)
+ self.mlp = ErnieImageFeedForward(hidden_size, ffn_hidden_size)
+
+ def forward(
+ self,
+ x: torch.Tensor,
+ rotary_pos_emb: torch.Tensor,
+ temb: Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor],
+ attention_mask: Optional[torch.Tensor] = None,
+ ) -> torch.Tensor:
+ shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = temb
+ residual = x
+ x = self.adaLN_sa_ln(x)
+ x = (x.float() * (1 + scale_msa.float()) + shift_msa.float()).to(x.dtype)
+ x_bsh = x.permute(1, 0, 2)
+ attn_out = self.self_attention(x_bsh, attention_mask=attention_mask, image_rotary_emb=rotary_pos_emb)
+ attn_out = attn_out.permute(1, 0, 2)
+ x = residual + (gate_msa.float() * attn_out.float()).to(x.dtype)
+ residual = x
+ x = self.adaLN_mlp_ln(x)
+ x = (x.float() * (1 + scale_mlp.float()) + shift_mlp.float()).to(x.dtype)
+ return residual + (gate_mlp.float() * self.mlp(x).float()).to(x.dtype)
+
+
+class ErnieImageAdaLNContinuous(nn.Module):
+ def __init__(self, hidden_size: int, eps: float = 1e-6):
+ super().__init__()
+ self.norm = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=eps)
+ self.linear = nn.Linear(hidden_size, hidden_size * 2)
+
+ def forward(self, x: torch.Tensor, conditioning: torch.Tensor) -> torch.Tensor:
+ scale, shift = self.linear(conditioning).chunk(2, dim=-1)
+ x = self.norm(x)
+ x = x * (1 + scale.unsqueeze(0)) + shift.unsqueeze(0)
+ return x
+
+
+class ErnieImageDiT(nn.Module):
+ """
+ Ernie-Image DiT model for DiffSynth-Studio.
+
+ Architecture: SharedAdaLN + RoPE 3D + Joint Image-Text Attention.
+ Internal format: [S, B, H] for transformer blocks, [B, S, H] for attention.
+ """
+
+ def __init__(
+ self,
+ hidden_size: int = 4096,
+ num_attention_heads: int = 32,
+ num_layers: int = 36,
+ ffn_hidden_size: int = 12288,
+ in_channels: int = 128,
+ out_channels: int = 128,
+ patch_size: int = 1,
+ text_in_dim: int = 3072,
+ rope_theta: int = 256,
+ rope_axes_dim: Tuple[int, int, int] = (32, 48, 48),
+ eps: float = 1e-6,
+ qk_layernorm: bool = True,
+ ):
+ super().__init__()
+ self.hidden_size = hidden_size
+ self.num_heads = num_attention_heads
+ self.head_dim = hidden_size // num_attention_heads
+ self.num_layers = num_layers
+ self.patch_size = patch_size
+ self.in_channels = in_channels
+ self.out_channels = out_channels
+ self.text_in_dim = text_in_dim
+
+ self.x_embedder = ErnieImagePatchEmbedDynamic(in_channels, hidden_size, patch_size)
+ self.text_proj = nn.Linear(text_in_dim, hidden_size, bias=False) if text_in_dim != hidden_size else None
+ self.time_proj = Timesteps(hidden_size, flip_sin_to_cos=False, downscale_freq_shift=0)
+ self.time_embedding = TimestepEmbedding(hidden_size, hidden_size)
+ self.pos_embed = ErnieImageEmbedND3(dim=self.head_dim, theta=rope_theta, axes_dim=rope_axes_dim)
+ self.adaLN_modulation = nn.Sequential(nn.SiLU(), nn.Linear(hidden_size, 6 * hidden_size))
+ nn.init.zeros_(self.adaLN_modulation[-1].weight)
+ nn.init.zeros_(self.adaLN_modulation[-1].bias)
+ self.layers = nn.ModuleList([
+ ErnieImageSharedAdaLNBlock(hidden_size, num_attention_heads, ffn_hidden_size, eps, qk_layernorm=qk_layernorm)
+ for _ in range(num_layers)
+ ])
+ self.final_norm = ErnieImageAdaLNContinuous(hidden_size, eps)
+ self.final_linear = nn.Linear(hidden_size, patch_size * patch_size * out_channels)
+ nn.init.zeros_(self.final_linear.weight)
+ nn.init.zeros_(self.final_linear.bias)
+
+ def forward(
+ self,
+ hidden_states: torch.Tensor,
+ timestep: torch.Tensor,
+ text_bth: torch.Tensor,
+ text_lens: torch.Tensor,
+ use_gradient_checkpointing: bool = False,
+ use_gradient_checkpointing_offload: bool = False,
+ ) -> torch.Tensor:
+ device, dtype = hidden_states.device, hidden_states.dtype
+ B, C, H, W = hidden_states.shape
+ p, Hp, Wp = self.patch_size, H // self.patch_size, W // self.patch_size
+ N_img = Hp * Wp
+
+ img_sbh = self.x_embedder(hidden_states).transpose(0, 1).contiguous()
+
+ if self.text_proj is not None and text_bth.numel() > 0:
+ text_bth = self.text_proj(text_bth)
+ Tmax = text_bth.shape[1]
+ text_sbh = text_bth.transpose(0, 1).contiguous()
+
+ x = torch.cat([img_sbh, text_sbh], dim=0)
+ S = x.shape[0]
+
+ text_ids = torch.cat([
+ torch.arange(Tmax, device=device, dtype=torch.float32).view(1, Tmax, 1).expand(B, -1, -1),
+ torch.zeros((B, Tmax, 2), device=device)
+ ], dim=-1) if Tmax > 0 else torch.zeros((B, 0, 3), device=device)
+ grid_yx = torch.stack(
+ torch.meshgrid(torch.arange(Hp, device=device, dtype=torch.float32),
+ torch.arange(Wp, device=device, dtype=torch.float32), indexing="ij"),
+ dim=-1
+ ).reshape(-1, 2)
+ image_ids = torch.cat([
+ text_lens.float().view(B, 1, 1).expand(-1, N_img, -1),
+ grid_yx.view(1, N_img, 2).expand(B, -1, -1)
+ ], dim=-1)
+ rotary_pos_emb = self.pos_embed(torch.cat([image_ids, text_ids], dim=1))
+
+ valid_text = torch.arange(Tmax, device=device).view(1, Tmax) < text_lens.view(B, 1) if Tmax > 0 else torch.zeros((B, 0), device=device, dtype=torch.bool)
+ attention_mask = torch.cat([
+ torch.ones((B, N_img), device=device, dtype=torch.bool),
+ valid_text
+ ], dim=1)[:, None, None, :]
+
+ sample = self.time_proj(timestep.to(dtype))
+ sample = sample.to(self.time_embedding.linear_1.weight.dtype)
+ c = self.time_embedding(sample)
+ shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = [
+ t.unsqueeze(0).expand(S, -1, -1).contiguous()
+ for t in self.adaLN_modulation(c).chunk(6, dim=-1)
+ ]
+
+ for layer in self.layers:
+ temb = [shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp]
+ if torch.is_grad_enabled() and use_gradient_checkpointing:
+ x = gradient_checkpoint_forward(
+ layer,
+ use_gradient_checkpointing,
+ use_gradient_checkpointing_offload,
+ x,
+ rotary_pos_emb,
+ temb,
+ attention_mask,
+ )
+ else:
+ x = layer(x, rotary_pos_emb, temb, attention_mask)
+
+ x = self.final_norm(x, c).type_as(x)
+ patches = self.final_linear(x)[:N_img].transpose(0, 1).contiguous()
+ output = patches.view(B, Hp, Wp, p, p, self.out_channels).permute(0, 5, 1, 3, 2, 4).contiguous().view(B, self.out_channels, H, W)
+
+ return output
diff --git a/diffsynth/models/ernie_image_text_encoder.py b/diffsynth/models/ernie_image_text_encoder.py
new file mode 100644
index 0000000000000000000000000000000000000000..17460b262fba6dd07e87c1f07618abc417899aa7
--- /dev/null
+++ b/diffsynth/models/ernie_image_text_encoder.py
@@ -0,0 +1,76 @@
+"""
+Ernie-Image TextEncoder for DiffSynth-Studio.
+
+Wraps transformers Ministral3Model to output text embeddings.
+Pattern: lazy import + manual config dict + torch.nn.Module wrapper.
+Only loads the text (language) model, ignoring vision components.
+"""
+
+import torch
+
+
+class ErnieImageTextEncoder(torch.nn.Module):
+ """
+ Text encoder using Ministral3Model (transformers).
+ Only the text_config portion of the full Mistral3Model checkpoint.
+ Uses the base model (no lm_head) since the checkpoint only has embeddings.
+ """
+
+ def __init__(self):
+ super().__init__()
+ from transformers import Ministral3Config, Ministral3Model
+
+ text_config = {
+ "attention_dropout": 0.0,
+ "bos_token_id": 1,
+ "dtype": "bfloat16",
+ "eos_token_id": 2,
+ "head_dim": 128,
+ "hidden_act": "silu",
+ "hidden_size": 3072,
+ "initializer_range": 0.02,
+ "intermediate_size": 9216,
+ "max_position_embeddings": 262144,
+ "model_type": "ministral3",
+ "num_attention_heads": 32,
+ "num_hidden_layers": 26,
+ "num_key_value_heads": 8,
+ "pad_token_id": 11,
+ "rms_norm_eps": 1e-05,
+ "rope_parameters": {
+ "beta_fast": 32.0,
+ "beta_slow": 1.0,
+ "factor": 16.0,
+ "llama_4_scaling_beta": 0.1,
+ "mscale": 1.0,
+ "mscale_all_dim": 1.0,
+ "original_max_position_embeddings": 16384,
+ "rope_theta": 1000000.0,
+ "rope_type": "yarn",
+ "type": "yarn",
+ },
+ "sliding_window": None,
+ "tie_word_embeddings": True,
+ "use_cache": True,
+ "vocab_size": 131072,
+ }
+ config = Ministral3Config(**text_config)
+ self.model = Ministral3Model(config)
+ self.config = config
+
+ def forward(
+ self,
+ input_ids=None,
+ attention_mask=None,
+ position_ids=None,
+ **kwargs,
+ ):
+ outputs = self.model(
+ input_ids=input_ids,
+ attention_mask=attention_mask,
+ position_ids=position_ids,
+ output_hidden_states=True,
+ return_dict=True,
+ **kwargs,
+ )
+ return (outputs.hidden_states,)
diff --git a/diffsynth/models/fid.py b/diffsynth/models/fid.py
new file mode 100644
index 0000000000000000000000000000000000000000..af36719751c3a9b4b243f74a8c84b286f4c10f9f
--- /dev/null
+++ b/diffsynth/models/fid.py
@@ -0,0 +1,238 @@
+import os
+from typing import Iterable, Union
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from PIL import Image
+from torchvision import transforms
+from torchvision.models import inception_v3
+from torchvision.models.inception import InceptionA, InceptionC, InceptionE
+
+ImageInput = Union[str, os.PathLike, Image.Image]
+
+IMAGE_EXTENSIONS = {".bmp", ".jpg", ".jpeg", ".pgm", ".png", ".ppm", ".tif", ".tiff", ".webp"}
+
+
+def _image_files(path: Union[str, os.PathLike]):
+ path = os.fspath(path)
+ if os.path.isfile(path):
+ if os.path.splitext(path)[1].lower() not in IMAGE_EXTENSIONS:
+ raise ValueError(f"Unsupported image extension for FID: {path}")
+ return [path]
+ if not os.path.exists(path):
+ raise FileNotFoundError(f"FID path does not exist: {path}")
+ files = []
+ for root, dirs, names in os.walk(path):
+ dirs.sort()
+ for name in sorted(names):
+ if os.path.splitext(name)[1].lower() in IMAGE_EXTENSIONS:
+ files.append(os.path.join(root, name))
+ if not files:
+ raise ValueError(f"No images found under {path}.")
+ return files
+
+
+class _ImageDataset(torch.utils.data.Dataset):
+ def __init__(self, images: Iterable[ImageInput], transform):
+ self.images = list(images)
+ self.transform = transform
+
+ def __len__(self):
+ return len(self.images)
+
+ def __getitem__(self, index):
+ image = self.images[index]
+ if isinstance(image, (str, os.PathLike)):
+ image = Image.open(image)
+ if not isinstance(image, Image.Image):
+ raise TypeError(f"FID expects PIL images or image paths, but received {type(image)}.")
+ return self.transform(image.convert("RGB"))
+
+
+class FIDInceptionModel(nn.Module):
+ def __init__(self):
+ super().__init__()
+ self.model = _fid_inception_v3()
+
+ def forward(self, images):
+ images = 2 * images - 1
+ return self.model(images)
+
+
+def _fid_inception_v3(weights_path: str = None):
+ model = inception_v3(weights=None, aux_logits=False, num_classes=1008, init_weights=False)
+ model.Mixed_5b = _FIDInceptionA(192, pool_features=32)
+ model.Mixed_5c = _FIDInceptionA(256, pool_features=64)
+ model.Mixed_5d = _FIDInceptionA(288, pool_features=64)
+ model.Mixed_6b = _FIDInceptionC(768, channels_7x7=128)
+ model.Mixed_6c = _FIDInceptionC(768, channels_7x7=160)
+ model.Mixed_6d = _FIDInceptionC(768, channels_7x7=160)
+ model.Mixed_6e = _FIDInceptionC(768, channels_7x7=192)
+ model.Mixed_7b = _FIDInceptionE1(1280)
+ model.Mixed_7c = _FIDInceptionE2(2048)
+ if weights_path is not None:
+ model.load_state_dict(torch.load(weights_path, map_location="cpu"))
+ model.fc = nn.Identity()
+ return model
+
+
+class _FIDInceptionA(InceptionA):
+ def forward(self, x):
+ branch1x1 = self.branch1x1(x)
+
+ branch5x5 = self.branch5x5_1(x)
+ branch5x5 = self.branch5x5_2(branch5x5)
+
+ branch3x3dbl = self.branch3x3dbl_1(x)
+ branch3x3dbl = self.branch3x3dbl_2(branch3x3dbl)
+ branch3x3dbl = self.branch3x3dbl_3(branch3x3dbl)
+
+ branch_pool = F.avg_pool2d(x, kernel_size=3, stride=1, padding=1, count_include_pad=False)
+ branch_pool = self.branch_pool(branch_pool)
+
+ return torch.cat([branch1x1, branch5x5, branch3x3dbl, branch_pool], 1)
+
+
+class _FIDInceptionC(InceptionC):
+ def forward(self, x):
+ branch1x1 = self.branch1x1(x)
+
+ branch7x7 = self.branch7x7_1(x)
+ branch7x7 = self.branch7x7_2(branch7x7)
+ branch7x7 = self.branch7x7_3(branch7x7)
+
+ branch7x7dbl = self.branch7x7dbl_1(x)
+ branch7x7dbl = self.branch7x7dbl_2(branch7x7dbl)
+ branch7x7dbl = self.branch7x7dbl_3(branch7x7dbl)
+ branch7x7dbl = self.branch7x7dbl_4(branch7x7dbl)
+ branch7x7dbl = self.branch7x7dbl_5(branch7x7dbl)
+
+ branch_pool = F.avg_pool2d(x, kernel_size=3, stride=1, padding=1, count_include_pad=False)
+ branch_pool = self.branch_pool(branch_pool)
+
+ return torch.cat([branch1x1, branch7x7, branch7x7dbl, branch_pool], 1)
+
+
+class _FIDInceptionE1(InceptionE):
+ def forward(self, x):
+ branch1x1 = self.branch1x1(x)
+
+ branch3x3 = self.branch3x3_1(x)
+ branch3x3 = torch.cat([self.branch3x3_2a(branch3x3), self.branch3x3_2b(branch3x3)], 1)
+
+ branch3x3dbl = self.branch3x3dbl_1(x)
+ branch3x3dbl = self.branch3x3dbl_2(branch3x3dbl)
+ branch3x3dbl = torch.cat([self.branch3x3dbl_3a(branch3x3dbl), self.branch3x3dbl_3b(branch3x3dbl)], 1)
+
+ branch_pool = F.avg_pool2d(x, kernel_size=3, stride=1, padding=1, count_include_pad=False)
+ branch_pool = self.branch_pool(branch_pool)
+
+ return torch.cat([branch1x1, branch3x3, branch3x3dbl, branch_pool], 1)
+
+
+class _FIDInceptionE2(InceptionE):
+ def forward(self, x):
+ branch1x1 = self.branch1x1(x)
+
+ branch3x3 = self.branch3x3_1(x)
+ branch3x3 = torch.cat([self.branch3x3_2a(branch3x3), self.branch3x3_2b(branch3x3)], 1)
+
+ branch3x3dbl = self.branch3x3dbl_1(x)
+ branch3x3dbl = self.branch3x3dbl_2(branch3x3dbl)
+ branch3x3dbl = torch.cat([self.branch3x3dbl_3a(branch3x3dbl), self.branch3x3dbl_3b(branch3x3dbl)], 1)
+
+ branch_pool = F.max_pool2d(x, kernel_size=3, stride=1, padding=1)
+ branch_pool = self.branch_pool(branch_pool)
+
+ return torch.cat([branch1x1, branch3x3, branch3x3dbl, branch_pool], 1)
+
+
+class FIDModel(torch.nn.Module):
+ def __init__(self, model: torch.nn.Module, device: Union[str, torch.device] = "cpu", batch_size: int = 50, num_workers: int = 0):
+ super().__init__()
+ self.model = model
+ self.batch_size = batch_size
+ self.num_workers = num_workers
+ self.transform = transforms.Compose(
+ [
+ transforms.Resize((299, 299), interpolation=transforms.InterpolationMode.BICUBIC),
+ transforms.ToTensor(),
+ ]
+ )
+ self.to(device)
+
+ @property
+ def device(self):
+ try:
+ return next(self.model.parameters()).device
+ except StopIteration:
+ return torch.device("cpu")
+
+ def _as_images(self, images):
+ if isinstance(images, (str, os.PathLike)):
+ return _image_files(images)
+ if isinstance(images, Image.Image):
+ return [images]
+ return list(images)
+
+ @torch.no_grad()
+ def get_activations(self, images, batch_size: int = None, num_workers: int = None):
+ images = self._as_images(images)
+ batch_size = self.batch_size if batch_size is None else batch_size
+ num_workers = self.num_workers if num_workers is None else num_workers
+ dataset = _ImageDataset(images, transform=self.transform)
+ dataloader = torch.utils.data.DataLoader(dataset, batch_size=min(batch_size, len(dataset)), shuffle=False, num_workers=num_workers)
+ activations = []
+ self.model.eval()
+ for batch in dataloader:
+ batch = batch.to(self.device)
+ features = self.model(batch)
+ if isinstance(features, tuple):
+ features = features[0]
+ if features.ndim == 4:
+ features = F.adaptive_avg_pool2d(features, output_size=(1, 1)).flatten(1)
+ activations.append(features.detach().cpu().to(torch.float64))
+ return torch.cat(activations, dim=0)
+
+ def statistics(self, images, batch_size: int = None, num_workers: int = None):
+ activations = self.get_activations(images, batch_size=batch_size, num_workers=num_workers)
+ return self.activation_statistics(activations)
+
+ @staticmethod
+ def activation_statistics(activations):
+ activations = activations.to(torch.float64)
+ mean = activations.mean(dim=0)
+ centered = activations - mean
+ if activations.shape[0] <= 1:
+ covariance = torch.zeros((activations.shape[1], activations.shape[1]), dtype=torch.float64)
+ else:
+ covariance = centered.T @ centered / (activations.shape[0] - 1)
+ return mean, covariance
+
+ @staticmethod
+ def _sqrtm_psd(matrix, eps: float = 1e-10):
+ matrix = (matrix + matrix.T) * 0.5
+ eigenvalues, eigenvectors = torch.linalg.eigh(matrix)
+ eigenvalues = eigenvalues.clamp_min(eps).sqrt()
+ return (eigenvectors * eigenvalues.unsqueeze(0)) @ eigenvectors.T
+
+ @classmethod
+ def frechet_distance(cls, mean1, covariance1, mean2, covariance2, eps: float = 1e-6):
+ mean1 = mean1.to(torch.float64)
+ covariance1 = covariance1.to(torch.float64)
+ mean2 = mean2.to(torch.float64)
+ covariance2 = covariance2.to(torch.float64)
+ diff = mean1 - mean2
+ offset = torch.eye(covariance1.shape[0], dtype=torch.float64) * eps
+ sqrt_cov1 = cls._sqrtm_psd(covariance1 + offset)
+ covmean = cls._sqrtm_psd(sqrt_cov1 @ (covariance2 + offset) @ sqrt_cov1)
+ distance = diff.dot(diff) + torch.trace(covariance1) + torch.trace(covariance2) - 2 * torch.trace(covmean)
+ return distance.clamp_min(0)
+
+ def compute(self, reference_images, generated_images, batch_size: int = None, num_workers: int = None):
+ mean1, covariance1 = self.statistics(reference_images, batch_size=batch_size, num_workers=num_workers)
+ mean2, covariance2 = self.statistics(generated_images, batch_size=batch_size, num_workers=num_workers)
+ return self.frechet_distance(mean1, covariance1, mean2, covariance2)
+
+ def forward(self, reference_images, generated_images):
+ return self.compute(reference_images, generated_images)
diff --git a/diffsynth/models/flux2_dit.py b/diffsynth/models/flux2_dit.py
new file mode 100644
index 0000000000000000000000000000000000000000..69d66a9538132aceae8ebca50769c942200722c0
--- /dev/null
+++ b/diffsynth/models/flux2_dit.py
@@ -0,0 +1,998 @@
+import inspect
+from typing import Any, Dict, List, Optional, Tuple, Union
+
+import torch, math
+import torch.nn as nn
+import torch.nn.functional as F
+import numpy as np
+from ..core.attention import attention_forward
+from ..core.gradient import gradient_checkpoint_forward
+
+
+def get_timestep_embedding(
+ timesteps: torch.Tensor,
+ embedding_dim: int,
+ flip_sin_to_cos: bool = False,
+ downscale_freq_shift: float = 1,
+ scale: float = 1,
+ max_period: int = 10000,
+) -> torch.Tensor:
+ """
+ This matches the implementation in Denoising Diffusion Probabilistic Models: Create sinusoidal timestep embeddings.
+
+ Args
+ timesteps (torch.Tensor):
+ a 1-D Tensor of N indices, one per batch element. These may be fractional.
+ embedding_dim (int):
+ the dimension of the output.
+ flip_sin_to_cos (bool):
+ Whether the embedding order should be `cos, sin` (if True) or `sin, cos` (if False)
+ downscale_freq_shift (float):
+ Controls the delta between frequencies between dimensions
+ scale (float):
+ Scaling factor applied to the embeddings.
+ max_period (int):
+ Controls the maximum frequency of the embeddings
+ Returns
+ torch.Tensor: an [N x dim] Tensor of positional embeddings.
+ """
+ assert len(timesteps.shape) == 1, "Timesteps should be a 1d-array"
+
+ half_dim = embedding_dim // 2
+ exponent = -math.log(max_period) * torch.arange(
+ start=0, end=half_dim, dtype=torch.float32, device=timesteps.device
+ )
+ exponent = exponent / (half_dim - downscale_freq_shift)
+
+ emb = torch.exp(exponent)
+ emb = timesteps[:, None].float() * emb[None, :]
+
+ # scale embeddings
+ emb = scale * emb
+
+ # concat sine and cosine embeddings
+ emb = torch.cat([torch.sin(emb), torch.cos(emb)], dim=-1)
+
+ # flip sine and cosine embeddings
+ if flip_sin_to_cos:
+ emb = torch.cat([emb[:, half_dim:], emb[:, :half_dim]], dim=-1)
+
+ # zero pad
+ if embedding_dim % 2 == 1:
+ emb = torch.nn.functional.pad(emb, (0, 1, 0, 0))
+ return emb
+
+
+class TimestepEmbedding(nn.Module):
+ def __init__(
+ self,
+ in_channels: int,
+ time_embed_dim: int,
+ act_fn: str = "silu",
+ out_dim: int = None,
+ post_act_fn: Optional[str] = None,
+ cond_proj_dim=None,
+ sample_proj_bias=True,
+ ):
+ super().__init__()
+
+ self.linear_1 = nn.Linear(in_channels, time_embed_dim, sample_proj_bias)
+
+ if cond_proj_dim is not None:
+ self.cond_proj = nn.Linear(cond_proj_dim, in_channels, bias=False)
+ else:
+ self.cond_proj = None
+
+ self.act = torch.nn.SiLU()
+
+ if out_dim is not None:
+ time_embed_dim_out = out_dim
+ else:
+ time_embed_dim_out = time_embed_dim
+ self.linear_2 = nn.Linear(time_embed_dim, time_embed_dim_out, sample_proj_bias)
+
+ if post_act_fn is None:
+ self.post_act = None
+
+ def forward(self, sample, condition=None):
+ if condition is not None:
+ sample = sample + self.cond_proj(condition)
+ sample = self.linear_1(sample)
+
+ if self.act is not None:
+ sample = self.act(sample)
+
+ sample = self.linear_2(sample)
+
+ if self.post_act is not None:
+ sample = self.post_act(sample)
+ return sample
+
+
+class Timesteps(nn.Module):
+ def __init__(self, num_channels: int, flip_sin_to_cos: bool, downscale_freq_shift: float, scale: int = 1):
+ super().__init__()
+ self.num_channels = num_channels
+ self.flip_sin_to_cos = flip_sin_to_cos
+ self.downscale_freq_shift = downscale_freq_shift
+ self.scale = scale
+
+ def forward(self, timesteps: torch.Tensor) -> torch.Tensor:
+ t_emb = get_timestep_embedding(
+ timesteps,
+ self.num_channels,
+ flip_sin_to_cos=self.flip_sin_to_cos,
+ downscale_freq_shift=self.downscale_freq_shift,
+ scale=self.scale,
+ )
+ return t_emb
+
+
+class AdaLayerNormContinuous(nn.Module):
+ r"""
+ Adaptive normalization layer with a norm layer (layer_norm or rms_norm).
+
+ Args:
+ embedding_dim (`int`): Embedding dimension to use during projection.
+ conditioning_embedding_dim (`int`): Dimension of the input condition.
+ elementwise_affine (`bool`, defaults to `True`):
+ Boolean flag to denote if affine transformation should be applied.
+ eps (`float`, defaults to 1e-5): Epsilon factor.
+ bias (`bias`, defaults to `True`): Boolean flag to denote if bias should be use.
+ norm_type (`str`, defaults to `"layer_norm"`):
+ Normalization layer to use. Values supported: "layer_norm", "rms_norm".
+ """
+
+ def __init__(
+ self,
+ embedding_dim: int,
+ conditioning_embedding_dim: int,
+ # NOTE: It is a bit weird that the norm layer can be configured to have scale and shift parameters
+ # because the output is immediately scaled and shifted by the projected conditioning embeddings.
+ # Note that AdaLayerNorm does not let the norm layer have scale and shift parameters.
+ # However, this is how it was implemented in the original code, and it's rather likely you should
+ # set `elementwise_affine` to False.
+ elementwise_affine=True,
+ eps=1e-5,
+ bias=True,
+ norm_type="layer_norm",
+ ):
+ super().__init__()
+ self.silu = nn.SiLU()
+ self.linear = nn.Linear(conditioning_embedding_dim, embedding_dim * 2, bias=bias)
+ if norm_type == "layer_norm":
+ self.norm = nn.LayerNorm(embedding_dim, eps, elementwise_affine, bias)
+
+ def forward(self, x: torch.Tensor, conditioning_embedding: torch.Tensor) -> torch.Tensor:
+ # convert back to the original dtype in case `conditioning_embedding`` is upcasted to float32 (needed for hunyuanDiT)
+ emb = self.linear(self.silu(conditioning_embedding).to(x.dtype))
+ scale, shift = torch.chunk(emb, 2, dim=1)
+ x = self.norm(x) * (1 + scale)[:, None, :] + shift[:, None, :]
+ return x
+
+
+def get_1d_rotary_pos_embed(
+ dim: int,
+ pos: Union[np.ndarray, int],
+ theta: float = 10000.0,
+ use_real=False,
+ linear_factor=1.0,
+ ntk_factor=1.0,
+ repeat_interleave_real=True,
+ freqs_dtype=torch.float32, # torch.float32, torch.float64 (flux)
+):
+ """
+ Precompute the frequency tensor for complex exponentials (cis) with given dimensions.
+
+ This function calculates a frequency tensor with complex exponentials using the given dimension 'dim' and the end
+ index 'end'. The 'theta' parameter scales the frequencies. The returned tensor contains complex values in complex64
+ data type.
+
+ Args:
+ dim (`int`): Dimension of the frequency tensor.
+ pos (`np.ndarray` or `int`): Position indices for the frequency tensor. [S] or scalar
+ theta (`float`, *optional*, defaults to 10000.0):
+ Scaling factor for frequency computation. Defaults to 10000.0.
+ use_real (`bool`, *optional*):
+ If True, return real part and imaginary part separately. Otherwise, return complex numbers.
+ linear_factor (`float`, *optional*, defaults to 1.0):
+ Scaling factor for the context extrapolation. Defaults to 1.0.
+ ntk_factor (`float`, *optional*, defaults to 1.0):
+ Scaling factor for the NTK-Aware RoPE. Defaults to 1.0.
+ repeat_interleave_real (`bool`, *optional*, defaults to `True`):
+ If `True` and `use_real`, real part and imaginary part are each interleaved with themselves to reach `dim`.
+ Otherwise, they are concateanted with themselves.
+ freqs_dtype (`torch.float32` or `torch.float64`, *optional*, defaults to `torch.float32`):
+ the dtype of the frequency tensor.
+ Returns:
+ `torch.Tensor`: Precomputed frequency tensor with complex exponentials. [S, D/2]
+ """
+ assert dim % 2 == 0
+
+ if isinstance(pos, int):
+ pos = torch.arange(pos)
+ if isinstance(pos, np.ndarray):
+ pos = torch.from_numpy(pos) # type: ignore # [S]
+
+ theta = theta * ntk_factor
+ freqs = (
+ 1.0 / (theta ** (torch.arange(0, dim, 2, dtype=freqs_dtype, device=pos.device) / dim)) / linear_factor
+ ) # [D/2]
+ freqs = torch.outer(pos, freqs) # type: ignore # [S, D/2]
+ is_npu = freqs.device.type == "npu"
+ if is_npu:
+ freqs = freqs.float()
+ if use_real and repeat_interleave_real:
+ # flux, hunyuan-dit, cogvideox
+ freqs_cos = freqs.cos().repeat_interleave(2, dim=1, output_size=freqs.shape[1] * 2).float() # [S, D]
+ freqs_sin = freqs.sin().repeat_interleave(2, dim=1, output_size=freqs.shape[1] * 2).float() # [S, D]
+ return freqs_cos, freqs_sin
+ elif use_real:
+ # stable audio, allegro
+ freqs_cos = torch.cat([freqs.cos(), freqs.cos()], dim=-1).float() # [S, D]
+ freqs_sin = torch.cat([freqs.sin(), freqs.sin()], dim=-1).float() # [S, D]
+ return freqs_cos, freqs_sin
+ else:
+ # lumina
+ freqs_cis = torch.polar(torch.ones_like(freqs), freqs) # complex64 # [S, D/2]
+ return freqs_cis
+
+
+def apply_rotary_emb(
+ x: torch.Tensor,
+ freqs_cis: Union[torch.Tensor, Tuple[torch.Tensor]],
+ use_real: bool = True,
+ use_real_unbind_dim: int = -1,
+ sequence_dim: int = 2,
+) -> Tuple[torch.Tensor, torch.Tensor]:
+ """
+ Apply rotary embeddings to input tensors using the given frequency tensor. This function applies rotary embeddings
+ to the given query or key 'x' tensors using the provided frequency tensor 'freqs_cis'. The input tensors are
+ reshaped as complex numbers, and the frequency tensor is reshaped for broadcasting compatibility. The resulting
+ tensors contain rotary embeddings and are returned as real tensors.
+
+ Args:
+ x (`torch.Tensor`):
+ Query or key tensor to apply rotary embeddings. [B, H, S, D] xk (torch.Tensor): Key tensor to apply
+ freqs_cis (`Tuple[torch.Tensor]`): Precomputed frequency tensor for complex exponentials. ([S, D], [S, D],)
+
+ Returns:
+ Tuple[torch.Tensor, torch.Tensor]: Tuple of modified query tensor and key tensor with rotary embeddings.
+ """
+ if use_real:
+ cos, sin = freqs_cis # [S, D]
+ if sequence_dim == 2:
+ cos = cos[None, None, :, :]
+ sin = sin[None, None, :, :]
+ elif sequence_dim == 1:
+ cos = cos[None, :, None, :]
+ sin = sin[None, :, None, :]
+ else:
+ raise ValueError(f"`sequence_dim={sequence_dim}` but should be 1 or 2.")
+
+ cos, sin = cos.to(x.device), sin.to(x.device)
+
+ if use_real_unbind_dim == -1:
+ # Used for flux, cogvideox, hunyuan-dit
+ x_real, x_imag = x.reshape(*x.shape[:-1], -1, 2).unbind(-1) # [B, H, S, D//2]
+ x_rotated = torch.stack([-x_imag, x_real], dim=-1).flatten(3)
+ elif use_real_unbind_dim == -2:
+ # Used for Stable Audio, OmniGen, CogView4 and Cosmos
+ x_real, x_imag = x.reshape(*x.shape[:-1], 2, -1).unbind(-2) # [B, H, S, D//2]
+ x_rotated = torch.cat([-x_imag, x_real], dim=-1)
+ else:
+ raise ValueError(f"`use_real_unbind_dim={use_real_unbind_dim}` but should be -1 or -2.")
+
+ out = (x.float() * cos + x_rotated.float() * sin).to(x.dtype)
+
+ return out
+ else:
+ # used for lumina
+ x_rotated = torch.view_as_complex(x.float().reshape(*x.shape[:-1], -1, 2))
+ freqs_cis = freqs_cis.unsqueeze(2)
+ x_out = torch.view_as_real(x_rotated * freqs_cis).flatten(3)
+
+ return x_out.type_as(x)
+
+def _get_projections(attn: "Flux2Attention", hidden_states, encoder_hidden_states=None):
+ query = attn.to_q(hidden_states)
+ key = attn.to_k(hidden_states)
+ value = attn.to_v(hidden_states)
+
+ encoder_query = encoder_key = encoder_value = None
+ if encoder_hidden_states is not None and attn.added_kv_proj_dim is not None:
+ encoder_query = attn.add_q_proj(encoder_hidden_states)
+ encoder_key = attn.add_k_proj(encoder_hidden_states)
+ encoder_value = attn.add_v_proj(encoder_hidden_states)
+
+ return query, key, value, encoder_query, encoder_key, encoder_value
+
+
+def _get_fused_projections(attn: "Flux2Attention", hidden_states, encoder_hidden_states=None):
+ query, key, value = attn.to_qkv(hidden_states).chunk(3, dim=-1)
+
+ encoder_query = encoder_key = encoder_value = (None,)
+ if encoder_hidden_states is not None and hasattr(attn, "to_added_qkv"):
+ encoder_query, encoder_key, encoder_value = attn.to_added_qkv(encoder_hidden_states).chunk(3, dim=-1)
+
+ return query, key, value, encoder_query, encoder_key, encoder_value
+
+
+def _get_qkv_projections(attn: "Flux2Attention", hidden_states, encoder_hidden_states=None):
+ return _get_projections(attn, hidden_states, encoder_hidden_states)
+
+
+class Flux2SwiGLU(nn.Module):
+ """
+ Flux 2 uses a SwiGLU-style activation in the transformer feedforward sub-blocks, but with the linear projection
+ layer fused into the first linear layer of the FF sub-block. Thus, this module has no trainable parameters.
+ """
+
+ def __init__(self):
+ super().__init__()
+ self.gate_fn = nn.SiLU()
+
+ def forward(self, x: torch.Tensor) -> torch.Tensor:
+ x1, x2 = x.chunk(2, dim=-1)
+ x = self.gate_fn(x1) * x2
+ return x
+
+
+class Flux2FeedForward(nn.Module):
+ def __init__(
+ self,
+ dim: int,
+ dim_out: Optional[int] = None,
+ mult: float = 3.0,
+ inner_dim: Optional[int] = None,
+ bias: bool = False,
+ ):
+ super().__init__()
+ if inner_dim is None:
+ inner_dim = int(dim * mult)
+ dim_out = dim_out or dim
+
+ # Flux2SwiGLU will reduce the dimension by half
+ self.linear_in = nn.Linear(dim, inner_dim * 2, bias=bias)
+ self.act_fn = Flux2SwiGLU()
+ self.linear_out = nn.Linear(inner_dim, dim_out, bias=bias)
+
+ def forward(self, x: torch.Tensor) -> torch.Tensor:
+ x = self.linear_in(x)
+ x = self.act_fn(x)
+ x = self.linear_out(x)
+ return x
+
+
+class Flux2Attention(torch.nn.Module):
+ def __init__(
+ self,
+ query_dim: int,
+ heads: int = 8,
+ dim_head: int = 64,
+ dropout: float = 0.0,
+ bias: bool = False,
+ added_kv_proj_dim: Optional[int] = None,
+ added_proj_bias: Optional[bool] = True,
+ out_bias: bool = True,
+ eps: float = 1e-5,
+ out_dim: int = None,
+ elementwise_affine: bool = True,
+ ):
+ super().__init__()
+
+ self.head_dim = dim_head
+ self.inner_dim = out_dim if out_dim is not None else dim_head * heads
+ self.query_dim = query_dim
+ self.out_dim = out_dim if out_dim is not None else query_dim
+ self.heads = out_dim // dim_head if out_dim is not None else heads
+
+ self.use_bias = bias
+ self.dropout = dropout
+
+ self.added_kv_proj_dim = added_kv_proj_dim
+ self.added_proj_bias = added_proj_bias
+
+ self.to_q = torch.nn.Linear(query_dim, self.inner_dim, bias=bias)
+ self.to_k = torch.nn.Linear(query_dim, self.inner_dim, bias=bias)
+ self.to_v = torch.nn.Linear(query_dim, self.inner_dim, bias=bias)
+
+ # QK Norm
+ self.norm_q = torch.nn.RMSNorm(dim_head, eps=eps, elementwise_affine=elementwise_affine)
+ self.norm_k = torch.nn.RMSNorm(dim_head, eps=eps, elementwise_affine=elementwise_affine)
+
+ self.to_out = torch.nn.ModuleList([])
+ self.to_out.append(torch.nn.Linear(self.inner_dim, self.out_dim, bias=out_bias))
+ self.to_out.append(torch.nn.Dropout(dropout))
+
+ if added_kv_proj_dim is not None:
+ self.norm_added_q = torch.nn.RMSNorm(dim_head, eps=eps)
+ self.norm_added_k = torch.nn.RMSNorm(dim_head, eps=eps)
+ self.add_q_proj = torch.nn.Linear(added_kv_proj_dim, self.inner_dim, bias=added_proj_bias)
+ self.add_k_proj = torch.nn.Linear(added_kv_proj_dim, self.inner_dim, bias=added_proj_bias)
+ self.add_v_proj = torch.nn.Linear(added_kv_proj_dim, self.inner_dim, bias=added_proj_bias)
+ self.to_add_out = torch.nn.Linear(self.inner_dim, query_dim, bias=out_bias)
+
+ def forward(
+ self,
+ hidden_states: torch.Tensor,
+ encoder_hidden_states: Optional[torch.Tensor] = None,
+ attention_mask: Optional[torch.Tensor] = None,
+ image_rotary_emb: Optional[torch.Tensor] = None,
+ kv_cache = None,
+ **kwargs,
+ ) -> torch.Tensor:
+ query, key, value, encoder_query, encoder_key, encoder_value = _get_qkv_projections(
+ self, hidden_states, encoder_hidden_states
+ )
+
+ query = query.unflatten(-1, (self.heads, -1))
+ key = key.unflatten(-1, (self.heads, -1))
+ value = value.unflatten(-1, (self.heads, -1))
+
+ query = self.norm_q(query)
+ key = self.norm_k(key)
+
+ if self.added_kv_proj_dim is not None:
+ encoder_query = encoder_query.unflatten(-1, (self.heads, -1))
+ encoder_key = encoder_key.unflatten(-1, (self.heads, -1))
+ encoder_value = encoder_value.unflatten(-1, (self.heads, -1))
+
+ encoder_query = self.norm_added_q(encoder_query)
+ encoder_key = self.norm_added_k(encoder_key)
+
+ query = torch.cat([encoder_query, query], dim=1)
+ key = torch.cat([encoder_key, key], dim=1)
+ value = torch.cat([encoder_value, value], dim=1)
+
+ if image_rotary_emb is not None:
+ query = apply_rotary_emb(query, image_rotary_emb, sequence_dim=1)
+ key = apply_rotary_emb(key, image_rotary_emb, sequence_dim=1)
+
+ if kv_cache is not None:
+ key = torch.concat([key, kv_cache[0]], dim=1)
+ value = torch.concat([value, kv_cache[1]], dim=1)
+ hidden_states = attention_forward(
+ query,
+ key,
+ value,
+ q_pattern="b s n d", k_pattern="b s n d", v_pattern="b s n d", out_pattern="b s n d",
+ )
+ hidden_states = hidden_states.flatten(2, 3)
+ hidden_states = hidden_states.to(query.dtype)
+
+ if encoder_hidden_states is not None:
+ encoder_hidden_states, hidden_states = hidden_states.split_with_sizes(
+ [encoder_hidden_states.shape[1], hidden_states.shape[1] - encoder_hidden_states.shape[1]], dim=1
+ )
+ encoder_hidden_states = self.to_add_out(encoder_hidden_states)
+
+ hidden_states = self.to_out[0](hidden_states)
+ hidden_states = self.to_out[1](hidden_states)
+
+ if encoder_hidden_states is not None:
+ return hidden_states, encoder_hidden_states
+ else:
+ return hidden_states
+
+
+class Flux2ParallelSelfAttention(torch.nn.Module):
+ def __init__(
+ self,
+ query_dim: int,
+ heads: int = 8,
+ dim_head: int = 64,
+ dropout: float = 0.0,
+ bias: bool = False,
+ out_bias: bool = True,
+ eps: float = 1e-5,
+ out_dim: int = None,
+ elementwise_affine: bool = True,
+ mlp_ratio: float = 4.0,
+ mlp_mult_factor: int = 2,
+ processor=None,
+ ):
+ super().__init__()
+
+ self.head_dim = dim_head
+ self.inner_dim = out_dim if out_dim is not None else dim_head * heads
+ self.query_dim = query_dim
+ self.out_dim = out_dim if out_dim is not None else query_dim
+ self.heads = out_dim // dim_head if out_dim is not None else heads
+
+ self.use_bias = bias
+ self.dropout = dropout
+
+ self.mlp_ratio = mlp_ratio
+ self.mlp_hidden_dim = int(query_dim * self.mlp_ratio)
+ self.mlp_mult_factor = mlp_mult_factor
+
+ # Fused QKV projections + MLP input projection
+ self.to_qkv_mlp_proj = torch.nn.Linear(
+ self.query_dim, self.inner_dim * 3 + self.mlp_hidden_dim * self.mlp_mult_factor, bias=bias
+ )
+ self.mlp_act_fn = Flux2SwiGLU()
+
+ # QK Norm
+ self.norm_q = torch.nn.RMSNorm(dim_head, eps=eps, elementwise_affine=elementwise_affine)
+ self.norm_k = torch.nn.RMSNorm(dim_head, eps=eps, elementwise_affine=elementwise_affine)
+
+ # Fused attention output projection + MLP output projection
+ self.to_out = torch.nn.Linear(self.inner_dim + self.mlp_hidden_dim, self.out_dim, bias=out_bias)
+
+ def forward(
+ self,
+ hidden_states: torch.Tensor,
+ attention_mask: Optional[torch.Tensor] = None,
+ image_rotary_emb: Optional[torch.Tensor] = None,
+ kv_cache = None,
+ **kwargs,
+ ) -> torch.Tensor:
+ # Parallel in (QKV + MLP in) projection
+ hidden_states = self.to_qkv_mlp_proj(hidden_states)
+ qkv, mlp_hidden_states = torch.split(
+ hidden_states, [3 * self.inner_dim, self.mlp_hidden_dim * self.mlp_mult_factor], dim=-1
+ )
+
+ # Handle the attention logic
+ query, key, value = qkv.chunk(3, dim=-1)
+
+ query = query.unflatten(-1, (self.heads, -1))
+ key = key.unflatten(-1, (self.heads, -1))
+ value = value.unflatten(-1, (self.heads, -1))
+
+ query = self.norm_q(query)
+ key = self.norm_k(key)
+
+ if image_rotary_emb is not None:
+ query = apply_rotary_emb(query, image_rotary_emb, sequence_dim=1)
+ key = apply_rotary_emb(key, image_rotary_emb, sequence_dim=1)
+
+ if kv_cache is not None:
+ key = torch.concat([key, kv_cache[0]], dim=1)
+ value = torch.concat([value, kv_cache[1]], dim=1)
+ hidden_states = attention_forward(
+ query,
+ key,
+ value,
+ q_pattern="b s n d", k_pattern="b s n d", v_pattern="b s n d", out_pattern="b s n d",
+ )
+ hidden_states = hidden_states.flatten(2, 3)
+ hidden_states = hidden_states.to(query.dtype)
+
+ # Handle the feedforward (FF) logic
+ mlp_hidden_states = self.mlp_act_fn(mlp_hidden_states)
+
+ # Concatenate and parallel output projection
+ hidden_states = torch.cat([hidden_states, mlp_hidden_states], dim=-1)
+ hidden_states = self.to_out(hidden_states)
+
+ return hidden_states
+
+
+class Flux2SingleTransformerBlock(nn.Module):
+ def __init__(
+ self,
+ dim: int,
+ num_attention_heads: int,
+ attention_head_dim: int,
+ mlp_ratio: float = 3.0,
+ eps: float = 1e-6,
+ bias: bool = False,
+ ):
+ super().__init__()
+
+ self.norm = nn.LayerNorm(dim, elementwise_affine=False, eps=eps)
+
+ # Note that the MLP in/out linear layers are fused with the attention QKV/out projections, respectively; this
+ # is often called a "parallel" transformer block. See the [ViT-22B paper](https://arxiv.org/abs/2302.05442)
+ # for a visual depiction of this type of transformer block.
+ self.attn = Flux2ParallelSelfAttention(
+ query_dim=dim,
+ dim_head=attention_head_dim,
+ heads=num_attention_heads,
+ out_dim=dim,
+ bias=bias,
+ out_bias=bias,
+ eps=eps,
+ mlp_ratio=mlp_ratio,
+ mlp_mult_factor=2,
+ )
+
+ def forward(
+ self,
+ hidden_states: torch.Tensor,
+ encoder_hidden_states: Optional[torch.Tensor],
+ temb_mod_params: Tuple[torch.Tensor, torch.Tensor, torch.Tensor],
+ image_rotary_emb: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
+ joint_attention_kwargs: Optional[Dict[str, Any]] = None,
+ split_hidden_states: bool = False,
+ text_seq_len: Optional[int] = None,
+ kv_cache = None,
+ ) -> Tuple[torch.Tensor, torch.Tensor]:
+ # If encoder_hidden_states is None, hidden_states is assumed to have encoder_hidden_states already
+ # concatenated
+ if encoder_hidden_states is not None:
+ text_seq_len = encoder_hidden_states.shape[1]
+ hidden_states = torch.cat([encoder_hidden_states, hidden_states], dim=1)
+
+ mod_shift, mod_scale, mod_gate = temb_mod_params
+
+ norm_hidden_states = self.norm(hidden_states)
+ norm_hidden_states = (1 + mod_scale) * norm_hidden_states + mod_shift
+
+ joint_attention_kwargs = joint_attention_kwargs or {}
+ attn_output = self.attn(
+ hidden_states=norm_hidden_states,
+ image_rotary_emb=image_rotary_emb,
+ kv_cache=kv_cache,
+ **joint_attention_kwargs,
+ )
+
+ hidden_states = hidden_states + mod_gate * attn_output
+ if hidden_states.dtype == torch.float16:
+ hidden_states = hidden_states.clip(-65504, 65504)
+
+ if split_hidden_states:
+ encoder_hidden_states, hidden_states = hidden_states[:, :text_seq_len], hidden_states[:, text_seq_len:]
+ return encoder_hidden_states, hidden_states
+ else:
+ return hidden_states
+
+
+class Flux2TransformerBlock(nn.Module):
+ def __init__(
+ self,
+ dim: int,
+ num_attention_heads: int,
+ attention_head_dim: int,
+ mlp_ratio: float = 3.0,
+ eps: float = 1e-6,
+ bias: bool = False,
+ ):
+ super().__init__()
+ self.mlp_hidden_dim = int(dim * mlp_ratio)
+
+ self.norm1 = nn.LayerNorm(dim, elementwise_affine=False, eps=eps)
+ self.norm1_context = nn.LayerNorm(dim, elementwise_affine=False, eps=eps)
+
+ self.attn = Flux2Attention(
+ query_dim=dim,
+ added_kv_proj_dim=dim,
+ dim_head=attention_head_dim,
+ heads=num_attention_heads,
+ out_dim=dim,
+ bias=bias,
+ added_proj_bias=bias,
+ out_bias=bias,
+ eps=eps,
+ )
+
+ self.norm2 = nn.LayerNorm(dim, elementwise_affine=False, eps=eps)
+ self.ff = Flux2FeedForward(dim=dim, dim_out=dim, mult=mlp_ratio, bias=bias)
+
+ self.norm2_context = nn.LayerNorm(dim, elementwise_affine=False, eps=eps)
+ self.ff_context = Flux2FeedForward(dim=dim, dim_out=dim, mult=mlp_ratio, bias=bias)
+
+ def forward(
+ self,
+ hidden_states: torch.Tensor,
+ encoder_hidden_states: torch.Tensor,
+ temb_mod_params_img: Tuple[Tuple[torch.Tensor, torch.Tensor, torch.Tensor], ...],
+ temb_mod_params_txt: Tuple[Tuple[torch.Tensor, torch.Tensor, torch.Tensor], ...],
+ image_rotary_emb: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
+ joint_attention_kwargs: Optional[Dict[str, Any]] = None,
+ kv_cache = None,
+ ) -> Tuple[torch.Tensor, torch.Tensor]:
+ joint_attention_kwargs = joint_attention_kwargs or {}
+
+ # Modulation parameters shape: [1, 1, self.dim]
+ (shift_msa, scale_msa, gate_msa), (shift_mlp, scale_mlp, gate_mlp) = temb_mod_params_img
+ (c_shift_msa, c_scale_msa, c_gate_msa), (c_shift_mlp, c_scale_mlp, c_gate_mlp) = temb_mod_params_txt
+
+ # Img stream
+ norm_hidden_states = self.norm1(hidden_states)
+ norm_hidden_states = (1 + scale_msa) * norm_hidden_states + shift_msa
+
+ # Conditioning txt stream
+ norm_encoder_hidden_states = self.norm1_context(encoder_hidden_states)
+ norm_encoder_hidden_states = (1 + c_scale_msa) * norm_encoder_hidden_states + c_shift_msa
+
+ # Attention on concatenated img + txt stream
+ attention_outputs = self.attn(
+ hidden_states=norm_hidden_states,
+ encoder_hidden_states=norm_encoder_hidden_states,
+ image_rotary_emb=image_rotary_emb,
+ kv_cache=kv_cache,
+ **joint_attention_kwargs,
+ )
+
+ attn_output, context_attn_output = attention_outputs
+
+ # Process attention outputs for the image stream (`hidden_states`).
+ attn_output = gate_msa * attn_output
+ hidden_states = hidden_states + attn_output
+
+ norm_hidden_states = self.norm2(hidden_states)
+ norm_hidden_states = norm_hidden_states * (1 + scale_mlp) + shift_mlp
+
+ ff_output = self.ff(norm_hidden_states)
+ hidden_states = hidden_states + gate_mlp * ff_output
+
+ # Process attention outputs for the text stream (`encoder_hidden_states`).
+ context_attn_output = c_gate_msa * context_attn_output
+ encoder_hidden_states = encoder_hidden_states + context_attn_output
+
+ norm_encoder_hidden_states = self.norm2_context(encoder_hidden_states)
+ norm_encoder_hidden_states = norm_encoder_hidden_states * (1 + c_scale_mlp) + c_shift_mlp
+
+ context_ff_output = self.ff_context(norm_encoder_hidden_states)
+ encoder_hidden_states = encoder_hidden_states + c_gate_mlp * context_ff_output
+ if encoder_hidden_states.dtype == torch.float16:
+ encoder_hidden_states = encoder_hidden_states.clip(-65504, 65504)
+
+ return encoder_hidden_states, hidden_states
+
+
+class Flux2PosEmbed(nn.Module):
+ # modified from https://github.com/black-forest-labs/flux/blob/c00d7c60b085fce8058b9df845e036090873f2ce/src/flux/modules/layers.py#L11
+ def __init__(self, theta: int, axes_dim: List[int]):
+ super().__init__()
+ self.theta = theta
+ self.axes_dim = axes_dim
+
+ def forward(self, ids: torch.Tensor) -> torch.Tensor:
+ # Expected ids shape: [S, len(self.axes_dim)]
+ cos_out = []
+ sin_out = []
+ pos = ids.float()
+ is_mps = ids.device.type == "mps"
+ is_npu = ids.device.type == "npu"
+ freqs_dtype = torch.float32 if (is_mps or is_npu) else torch.float64
+ # Unlike Flux 1, loop over len(self.axes_dim) rather than ids.shape[-1]
+ for i in range(len(self.axes_dim)):
+ cos, sin = get_1d_rotary_pos_embed(
+ self.axes_dim[i],
+ pos[..., i],
+ theta=self.theta,
+ repeat_interleave_real=True,
+ use_real=True,
+ freqs_dtype=freqs_dtype,
+ )
+ cos_out.append(cos)
+ sin_out.append(sin)
+ freqs_cos = torch.cat(cos_out, dim=-1).to(ids.device)
+ freqs_sin = torch.cat(sin_out, dim=-1).to(ids.device)
+ return freqs_cos, freqs_sin
+
+
+class Flux2TimestepGuidanceEmbeddings(nn.Module):
+ def __init__(
+ self,
+ in_channels: int = 256,
+ embedding_dim: int = 6144,
+ bias: bool = False,
+ guidance_embeds: bool = True,
+ ):
+ super().__init__()
+
+ self.time_proj = Timesteps(num_channels=in_channels, flip_sin_to_cos=True, downscale_freq_shift=0)
+ self.timestep_embedder = TimestepEmbedding(
+ in_channels=in_channels, time_embed_dim=embedding_dim, sample_proj_bias=bias
+ )
+
+ if guidance_embeds:
+ self.guidance_embedder = TimestepEmbedding(
+ in_channels=in_channels, time_embed_dim=embedding_dim, sample_proj_bias=bias
+ )
+ else:
+ self.guidance_embedder = None
+
+ def forward(self, timestep: torch.Tensor, guidance: torch.Tensor) -> torch.Tensor:
+ timesteps_proj = self.time_proj(timestep)
+ timesteps_emb = self.timestep_embedder(timesteps_proj.to(timestep.dtype)) # (N, D)
+
+ if guidance is not None and self.guidance_embedder is not None:
+ guidance_proj = self.time_proj(guidance)
+ guidance_emb = self.guidance_embedder(guidance_proj.to(guidance.dtype)) # (N, D)
+ time_guidance_emb = timesteps_emb + guidance_emb
+ return time_guidance_emb
+ else:
+ return timesteps_emb
+
+
+class Flux2Modulation(nn.Module):
+ def __init__(self, dim: int, mod_param_sets: int = 2, bias: bool = False):
+ super().__init__()
+ self.mod_param_sets = mod_param_sets
+
+ self.linear = nn.Linear(dim, dim * 3 * self.mod_param_sets, bias=bias)
+ self.act_fn = nn.SiLU()
+
+ def forward(self, temb: torch.Tensor) -> Tuple[Tuple[torch.Tensor, torch.Tensor, torch.Tensor], ...]:
+ mod = self.act_fn(temb)
+ mod = self.linear(mod)
+
+ if mod.ndim == 2:
+ mod = mod.unsqueeze(1)
+ mod_params = torch.chunk(mod, 3 * self.mod_param_sets, dim=-1)
+ # Return tuple of 3-tuples of modulation params shift/scale/gate
+ return tuple(mod_params[3 * i : 3 * (i + 1)] for i in range(self.mod_param_sets))
+
+
+class Flux2DiT(torch.nn.Module):
+
+ _repeated_blocks = ["Flux2TransformerBlock", "Flux2SingleTransformerBlock"]
+
+ def __init__(
+ self,
+ patch_size: int = 1,
+ in_channels: int = 128,
+ out_channels: Optional[int] = None,
+ num_layers: int = 8,
+ num_single_layers: int = 48,
+ attention_head_dim: int = 128,
+ num_attention_heads: int = 48,
+ joint_attention_dim: int = 15360,
+ timestep_guidance_channels: int = 256,
+ mlp_ratio: float = 3.0,
+ axes_dims_rope: Tuple[int, ...] = (32, 32, 32, 32),
+ rope_theta: int = 2000,
+ eps: float = 1e-6,
+ guidance_embeds: bool = True,
+ ):
+ super().__init__()
+ self.out_channels = out_channels or in_channels
+ self.inner_dim = num_attention_heads * attention_head_dim
+
+ # 1. Sinusoidal positional embedding for RoPE on image and text tokens
+ self.pos_embed = Flux2PosEmbed(theta=rope_theta, axes_dim=axes_dims_rope)
+
+ # 2. Combined timestep + guidance embedding
+ self.time_guidance_embed = Flux2TimestepGuidanceEmbeddings(
+ in_channels=timestep_guidance_channels,
+ embedding_dim=self.inner_dim,
+ bias=False,
+ guidance_embeds=guidance_embeds,
+ )
+
+ # 3. Modulation (double stream and single stream blocks share modulation parameters, resp.)
+ # Two sets of shift/scale/gate modulation parameters for the double stream attn and FF sub-blocks
+ self.double_stream_modulation_img = Flux2Modulation(self.inner_dim, mod_param_sets=2, bias=False)
+ self.double_stream_modulation_txt = Flux2Modulation(self.inner_dim, mod_param_sets=2, bias=False)
+ # Only one set of modulation parameters as the attn and FF sub-blocks are run in parallel for single stream
+ self.single_stream_modulation = Flux2Modulation(self.inner_dim, mod_param_sets=1, bias=False)
+
+ # 4. Input projections
+ self.x_embedder = nn.Linear(in_channels, self.inner_dim, bias=False)
+ self.context_embedder = nn.Linear(joint_attention_dim, self.inner_dim, bias=False)
+
+ # 5. Double Stream Transformer Blocks
+ self.transformer_blocks = nn.ModuleList(
+ [
+ Flux2TransformerBlock(
+ dim=self.inner_dim,
+ num_attention_heads=num_attention_heads,
+ attention_head_dim=attention_head_dim,
+ mlp_ratio=mlp_ratio,
+ eps=eps,
+ bias=False,
+ )
+ for _ in range(num_layers)
+ ]
+ )
+
+ # 6. Single Stream Transformer Blocks
+ self.single_transformer_blocks = nn.ModuleList(
+ [
+ Flux2SingleTransformerBlock(
+ dim=self.inner_dim,
+ num_attention_heads=num_attention_heads,
+ attention_head_dim=attention_head_dim,
+ mlp_ratio=mlp_ratio,
+ eps=eps,
+ bias=False,
+ )
+ for _ in range(num_single_layers)
+ ]
+ )
+
+ # 7. Output layers
+ self.norm_out = AdaLayerNormContinuous(
+ self.inner_dim, self.inner_dim, elementwise_affine=False, eps=eps, bias=False
+ )
+ self.proj_out = nn.Linear(self.inner_dim, patch_size * patch_size * self.out_channels, bias=False)
+
+ self.gradient_checkpointing = False
+
+ def forward(
+ self,
+ hidden_states: torch.Tensor,
+ encoder_hidden_states: torch.Tensor = None,
+ timestep: torch.LongTensor = None,
+ img_ids: torch.Tensor = None,
+ txt_ids: torch.Tensor = None,
+ guidance: torch.Tensor = None,
+ joint_attention_kwargs: Optional[Dict[str, Any]] = None,
+ kv_cache = None,
+ use_gradient_checkpointing=False,
+ use_gradient_checkpointing_offload=False,
+ ):
+ # 0. Handle input arguments
+ if joint_attention_kwargs is not None:
+ joint_attention_kwargs = joint_attention_kwargs.copy()
+ lora_scale = joint_attention_kwargs.pop("scale", 1.0)
+ else:
+ lora_scale = 1.0
+
+ num_txt_tokens = encoder_hidden_states.shape[1]
+
+ # 1. Calculate timestep embedding and modulation parameters
+ timestep = timestep.to(hidden_states.dtype) * 1000
+
+ if guidance is not None:
+ guidance = guidance.to(hidden_states.dtype) * 1000
+
+ temb = self.time_guidance_embed(timestep, guidance)
+
+ double_stream_mod_img = self.double_stream_modulation_img(temb)
+ double_stream_mod_txt = self.double_stream_modulation_txt(temb)
+ single_stream_mod = self.single_stream_modulation(temb)[0]
+
+ # 2. Input projection for image (hidden_states) and conditioning text (encoder_hidden_states)
+ hidden_states = self.x_embedder(hidden_states)
+ encoder_hidden_states = self.context_embedder(encoder_hidden_states)
+
+ # 3. Calculate RoPE embeddings from image and text tokens
+ # NOTE: the below logic means that we can't support batched inference with images of different resolutions or
+ # text prompts of differents lengths. Is this a use case we want to support?
+ if img_ids.ndim == 3:
+ img_ids = img_ids[0]
+ if txt_ids.ndim == 3:
+ txt_ids = txt_ids[0]
+
+ image_rotary_emb = self.pos_embed(img_ids)
+ text_rotary_emb = self.pos_embed(txt_ids)
+ concat_rotary_emb = (
+ torch.cat([text_rotary_emb[0], image_rotary_emb[0]], dim=0),
+ torch.cat([text_rotary_emb[1], image_rotary_emb[1]], dim=0),
+ )
+
+ # 4. Double Stream Transformer Blocks
+ for block_id, block in enumerate(self.transformer_blocks):
+ encoder_hidden_states, hidden_states = gradient_checkpoint_forward(
+ block,
+ use_gradient_checkpointing=use_gradient_checkpointing,
+ use_gradient_checkpointing_offload=use_gradient_checkpointing_offload,
+ hidden_states=hidden_states,
+ encoder_hidden_states=encoder_hidden_states,
+ temb_mod_params_img=double_stream_mod_img,
+ temb_mod_params_txt=double_stream_mod_txt,
+ image_rotary_emb=concat_rotary_emb,
+ joint_attention_kwargs=joint_attention_kwargs,
+ kv_cache=None if kv_cache is None else kv_cache.get(f"double_{block_id}"),
+ )
+ # Concatenate text and image streams for single-block inference
+ hidden_states = torch.cat([encoder_hidden_states, hidden_states], dim=1)
+
+ # 5. Single Stream Transformer Blocks
+ for block_id, block in enumerate(self.single_transformer_blocks):
+ hidden_states = gradient_checkpoint_forward(
+ block,
+ use_gradient_checkpointing=use_gradient_checkpointing,
+ use_gradient_checkpointing_offload=use_gradient_checkpointing_offload,
+ hidden_states=hidden_states,
+ encoder_hidden_states=None,
+ temb_mod_params=single_stream_mod,
+ image_rotary_emb=concat_rotary_emb,
+ joint_attention_kwargs=joint_attention_kwargs,
+ kv_cache=None if kv_cache is None else kv_cache.get(f"single_{block_id}"),
+ )
+ # Remove text tokens from concatenated stream
+ hidden_states = hidden_states[:, num_txt_tokens:, ...]
+
+ # 6. Output layers
+ hidden_states = self.norm_out(hidden_states, temb)
+ output = self.proj_out(hidden_states)
+
+ return output
diff --git a/diffsynth/models/flux2_text_encoder.py b/diffsynth/models/flux2_text_encoder.py
new file mode 100644
index 0000000000000000000000000000000000000000..f3c68411f3160655ebefd49dfb6424b19373a301
--- /dev/null
+++ b/diffsynth/models/flux2_text_encoder.py
@@ -0,0 +1,58 @@
+from transformers import Mistral3ForConditionalGeneration, Mistral3Config
+
+
+class Flux2TextEncoder(Mistral3ForConditionalGeneration):
+ def __init__(self):
+ config = Mistral3Config(**{
+ "architectures": [
+ "Mistral3ForConditionalGeneration"
+ ],
+ "dtype": "bfloat16",
+ "image_token_index": 10,
+ "model_type": "mistral3",
+ "multimodal_projector_bias": False,
+ "projector_hidden_act": "gelu",
+ "spatial_merge_size": 2,
+ "text_config": {
+ "attention_dropout": 0.0,
+ "dtype": "bfloat16",
+ "head_dim": 128,
+ "hidden_act": "silu",
+ "hidden_size": 5120,
+ "initializer_range": 0.02,
+ "intermediate_size": 32768,
+ "max_position_embeddings": 131072,
+ "model_type": "mistral",
+ "num_attention_heads": 32,
+ "num_hidden_layers": 40,
+ "num_key_value_heads": 8,
+ "rms_norm_eps": 1e-05,
+ "rope_theta": 1000000000.0,
+ "sliding_window": None,
+ "use_cache": True,
+ "vocab_size": 131072
+ },
+ "transformers_version": "4.57.1",
+ "vision_config": {
+ "attention_dropout": 0.0,
+ "dtype": "bfloat16",
+ "head_dim": 64,
+ "hidden_act": "silu",
+ "hidden_size": 1024,
+ "image_size": 1540,
+ "initializer_range": 0.02,
+ "intermediate_size": 4096,
+ "model_type": "pixtral",
+ "num_attention_heads": 16,
+ "num_channels": 3,
+ "num_hidden_layers": 24,
+ "patch_size": 14,
+ "rope_theta": 10000.0
+ },
+ "vision_feature_layer": -1
+ })
+ super().__init__(config)
+
+ def forward(self, input_ids = None, pixel_values = None, attention_mask = None, position_ids = None, past_key_values = None, inputs_embeds = None, labels = None, use_cache = None, output_attentions = None, output_hidden_states = None, return_dict = None, cache_position = None, logits_to_keep = 0, image_sizes = None, **kwargs):
+ return super().forward(input_ids, pixel_values, attention_mask, position_ids, past_key_values, inputs_embeds, labels, use_cache, output_attentions, output_hidden_states, return_dict, cache_position, logits_to_keep, image_sizes, **kwargs)
+
diff --git a/diffsynth/models/flux2_vae.py b/diffsynth/models/flux2_vae.py
new file mode 100644
index 0000000000000000000000000000000000000000..c7904b17618cc3c0811e42fde0f80ecd8f15f7ee
--- /dev/null
+++ b/diffsynth/models/flux2_vae.py
@@ -0,0 +1,2322 @@
+# Copyright 2025 The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+import math
+from typing import Dict, Optional, Tuple, Union, Callable
+
+import torch
+import torch.nn as nn
+from einops import rearrange
+import torch.nn.functional as F
+import inspect
+
+ACT2CLS = {
+ "swish": nn.SiLU,
+ "silu": nn.SiLU,
+ "mish": nn.Mish,
+ "gelu": nn.GELU,
+ "relu": nn.ReLU,
+}
+
+def get_activation(act_fn: str) -> nn.Module:
+ """Helper function to get activation function from string.
+
+ Args:
+ act_fn (str): Name of activation function.
+
+ Returns:
+ nn.Module: Activation function.
+ """
+
+ act_fn = act_fn.lower()
+ if act_fn in ACT2CLS:
+ return ACT2CLS[act_fn]()
+ else:
+ raise ValueError(f"activation function {act_fn} not found in ACT2FN mapping {list(ACT2CLS.keys())}")
+
+class ResnetBlock2D(nn.Module):
+ r"""
+ A Resnet block.
+
+ Parameters:
+ in_channels (`int`): The number of channels in the input.
+ out_channels (`int`, *optional*, default to be `None`):
+ The number of output channels for the first conv2d layer. If None, same as `in_channels`.
+ dropout (`float`, *optional*, defaults to `0.0`): The dropout probability to use.
+ temb_channels (`int`, *optional*, default to `512`): the number of channels in timestep embedding.
+ groups (`int`, *optional*, default to `32`): The number of groups to use for the first normalization layer.
+ groups_out (`int`, *optional*, default to None):
+ The number of groups to use for the second normalization layer. if set to None, same as `groups`.
+ eps (`float`, *optional*, defaults to `1e-6`): The epsilon to use for the normalization.
+ non_linearity (`str`, *optional*, default to `"swish"`): the activation function to use.
+ time_embedding_norm (`str`, *optional*, default to `"default"` ): Time scale shift config.
+ By default, apply timestep embedding conditioning with a simple shift mechanism. Choose "scale_shift" for a
+ stronger conditioning with scale and shift.
+ kernel (`torch.Tensor`, optional, default to None): FIR filter, see
+ [`~models.resnet.FirUpsample2D`] and [`~models.resnet.FirDownsample2D`].
+ output_scale_factor (`float`, *optional*, default to be `1.0`): the scale factor to use for the output.
+ use_in_shortcut (`bool`, *optional*, default to `True`):
+ If `True`, add a 1x1 nn.conv2d layer for skip-connection.
+ up (`bool`, *optional*, default to `False`): If `True`, add an upsample layer.
+ down (`bool`, *optional*, default to `False`): If `True`, add a downsample layer.
+ conv_shortcut_bias (`bool`, *optional*, default to `True`): If `True`, adds a learnable bias to the
+ `conv_shortcut` output.
+ conv_2d_out_channels (`int`, *optional*, default to `None`): the number of channels in the output.
+ If None, same as `out_channels`.
+ """
+
+ def __init__(
+ self,
+ *,
+ in_channels: int,
+ out_channels: Optional[int] = None,
+ conv_shortcut: bool = False,
+ dropout: float = 0.0,
+ temb_channels: int = 512,
+ groups: int = 32,
+ groups_out: Optional[int] = None,
+ pre_norm: bool = True,
+ eps: float = 1e-6,
+ non_linearity: str = "swish",
+ skip_time_act: bool = False,
+ time_embedding_norm: str = "default", # default, scale_shift,
+ kernel: Optional[torch.Tensor] = None,
+ output_scale_factor: float = 1.0,
+ use_in_shortcut: Optional[bool] = None,
+ up: bool = False,
+ down: bool = False,
+ conv_shortcut_bias: bool = True,
+ conv_2d_out_channels: Optional[int] = None,
+ ):
+ super().__init__()
+ if time_embedding_norm == "ada_group":
+ raise ValueError(
+ "This class cannot be used with `time_embedding_norm==ada_group`, please use `ResnetBlockCondNorm2D` instead",
+ )
+ if time_embedding_norm == "spatial":
+ raise ValueError(
+ "This class cannot be used with `time_embedding_norm==spatial`, please use `ResnetBlockCondNorm2D` instead",
+ )
+
+ self.pre_norm = True
+ self.in_channels = in_channels
+ out_channels = in_channels if out_channels is None else out_channels
+ self.out_channels = out_channels
+ self.use_conv_shortcut = conv_shortcut
+ self.up = up
+ self.down = down
+ self.output_scale_factor = output_scale_factor
+ self.time_embedding_norm = time_embedding_norm
+ self.skip_time_act = skip_time_act
+
+ if groups_out is None:
+ groups_out = groups
+
+ self.norm1 = torch.nn.GroupNorm(num_groups=groups, num_channels=in_channels, eps=eps, affine=True)
+
+ self.conv1 = nn.Conv2d(in_channels, out_channels, kernel_size=3, stride=1, padding=1)
+
+ if temb_channels is not None:
+ if self.time_embedding_norm == "default":
+ self.time_emb_proj = nn.Linear(temb_channels, out_channels)
+ elif self.time_embedding_norm == "scale_shift":
+ self.time_emb_proj = nn.Linear(temb_channels, 2 * out_channels)
+ else:
+ raise ValueError(f"unknown time_embedding_norm : {self.time_embedding_norm} ")
+ else:
+ self.time_emb_proj = None
+
+ self.norm2 = torch.nn.GroupNorm(num_groups=groups_out, num_channels=out_channels, eps=eps, affine=True)
+
+ self.dropout = torch.nn.Dropout(dropout)
+ conv_2d_out_channels = conv_2d_out_channels or out_channels
+ self.conv2 = nn.Conv2d(out_channels, conv_2d_out_channels, kernel_size=3, stride=1, padding=1)
+
+ self.nonlinearity = get_activation(non_linearity)
+
+ self.upsample = self.downsample = None
+ if self.up:
+ if kernel == "fir":
+ fir_kernel = (1, 3, 3, 1)
+ self.upsample = lambda x: upsample_2d(x, kernel=fir_kernel)
+ elif kernel == "sde_vp":
+ self.upsample = partial(F.interpolate, scale_factor=2.0, mode="nearest")
+ else:
+ self.upsample = Upsample2D(in_channels, use_conv=False)
+ elif self.down:
+ if kernel == "fir":
+ fir_kernel = (1, 3, 3, 1)
+ self.downsample = lambda x: downsample_2d(x, kernel=fir_kernel)
+ elif kernel == "sde_vp":
+ self.downsample = partial(F.avg_pool2d, kernel_size=2, stride=2)
+ else:
+ self.downsample = Downsample2D(in_channels, use_conv=False, padding=1, name="op")
+
+ self.use_in_shortcut = self.in_channels != conv_2d_out_channels if use_in_shortcut is None else use_in_shortcut
+
+ self.conv_shortcut = None
+ if self.use_in_shortcut:
+ self.conv_shortcut = nn.Conv2d(
+ in_channels,
+ conv_2d_out_channels,
+ kernel_size=1,
+ stride=1,
+ padding=0,
+ bias=conv_shortcut_bias,
+ )
+
+ def forward(self, input_tensor: torch.Tensor, temb: torch.Tensor, *args, **kwargs) -> torch.Tensor:
+ if len(args) > 0 or kwargs.get("scale", None) is not None:
+ deprecation_message = "The `scale` argument is deprecated and will be ignored. Please remove it, as passing it will raise an error in the future. `scale` should directly be passed while calling the underlying pipeline component i.e., via `cross_attention_kwargs`."
+ deprecate("scale", "1.0.0", deprecation_message)
+
+ hidden_states = input_tensor
+
+ hidden_states = self.norm1(hidden_states)
+ hidden_states = self.nonlinearity(hidden_states)
+
+ if self.upsample is not None:
+ # upsample_nearest_nhwc fails with large batch sizes. see https://github.com/huggingface/diffusers/issues/984
+ if hidden_states.shape[0] >= 64:
+ input_tensor = input_tensor.contiguous()
+ hidden_states = hidden_states.contiguous()
+ input_tensor = self.upsample(input_tensor)
+ hidden_states = self.upsample(hidden_states)
+ elif self.downsample is not None:
+ input_tensor = self.downsample(input_tensor)
+ hidden_states = self.downsample(hidden_states)
+
+ hidden_states = self.conv1(hidden_states)
+
+ if self.time_emb_proj is not None:
+ if not self.skip_time_act:
+ temb = self.nonlinearity(temb)
+ temb = self.time_emb_proj(temb)[:, :, None, None]
+
+ if self.time_embedding_norm == "default":
+ if temb is not None:
+ hidden_states = hidden_states + temb
+ hidden_states = self.norm2(hidden_states)
+ elif self.time_embedding_norm == "scale_shift":
+ if temb is None:
+ raise ValueError(
+ f" `temb` should not be None when `time_embedding_norm` is {self.time_embedding_norm}"
+ )
+ time_scale, time_shift = torch.chunk(temb, 2, dim=1)
+ hidden_states = self.norm2(hidden_states)
+ hidden_states = hidden_states * (1 + time_scale) + time_shift
+ else:
+ hidden_states = self.norm2(hidden_states)
+
+ hidden_states = self.nonlinearity(hidden_states)
+
+ hidden_states = self.dropout(hidden_states)
+ hidden_states = self.conv2(hidden_states)
+
+ if self.conv_shortcut is not None:
+ input_tensor = self.conv_shortcut(input_tensor.contiguous())
+
+ output_tensor = (input_tensor + hidden_states) / self.output_scale_factor
+
+ return output_tensor
+
+class Downsample2D(nn.Module):
+ """A 2D downsampling layer with an optional convolution.
+
+ Parameters:
+ channels (`int`):
+ number of channels in the inputs and outputs.
+ use_conv (`bool`, default `False`):
+ option to use a convolution.
+ out_channels (`int`, optional):
+ number of output channels. Defaults to `channels`.
+ padding (`int`, default `1`):
+ padding for the convolution.
+ name (`str`, default `conv`):
+ name of the downsampling 2D layer.
+ """
+
+ def __init__(
+ self,
+ channels: int,
+ use_conv: bool = False,
+ out_channels: Optional[int] = None,
+ padding: int = 1,
+ name: str = "conv",
+ kernel_size=3,
+ norm_type=None,
+ eps=None,
+ elementwise_affine=None,
+ bias=True,
+ ):
+ super().__init__()
+ self.channels = channels
+ self.out_channels = out_channels or channels
+ self.use_conv = use_conv
+ self.padding = padding
+ stride = 2
+ self.name = name
+
+ if norm_type == "ln_norm":
+ self.norm = nn.LayerNorm(channels, eps, elementwise_affine)
+ elif norm_type == "rms_norm":
+ self.norm = RMSNorm(channels, eps, elementwise_affine)
+ elif norm_type is None:
+ self.norm = None
+ else:
+ raise ValueError(f"unknown norm_type: {norm_type}")
+
+ if use_conv:
+ conv = nn.Conv2d(
+ self.channels, self.out_channels, kernel_size=kernel_size, stride=stride, padding=padding, bias=bias
+ )
+ else:
+ assert self.channels == self.out_channels
+ conv = nn.AvgPool2d(kernel_size=stride, stride=stride)
+
+ # TODO(Suraj, Patrick) - clean up after weight dicts are correctly renamed
+ if name == "conv":
+ self.Conv2d_0 = conv
+ self.conv = conv
+ elif name == "Conv2d_0":
+ self.conv = conv
+ else:
+ self.conv = conv
+
+ def forward(self, hidden_states: torch.Tensor, *args, **kwargs) -> torch.Tensor:
+ if len(args) > 0 or kwargs.get("scale", None) is not None:
+ deprecation_message = "The `scale` argument is deprecated and will be ignored. Please remove it, as passing it will raise an error in the future. `scale` should directly be passed while calling the underlying pipeline component i.e., via `cross_attention_kwargs`."
+ deprecate("scale", "1.0.0", deprecation_message)
+ assert hidden_states.shape[1] == self.channels
+
+ if self.norm is not None:
+ hidden_states = self.norm(hidden_states.permute(0, 2, 3, 1)).permute(0, 3, 1, 2)
+
+ if self.use_conv and self.padding == 0:
+ pad = (0, 1, 0, 1)
+ hidden_states = F.pad(hidden_states, pad, mode="constant", value=0)
+
+ assert hidden_states.shape[1] == self.channels
+
+ hidden_states = self.conv(hidden_states)
+
+ return hidden_states
+
+class Upsample2D(nn.Module):
+ """A 2D upsampling layer with an optional convolution.
+
+ Parameters:
+ channels (`int`):
+ number of channels in the inputs and outputs.
+ use_conv (`bool`, default `False`):
+ option to use a convolution.
+ use_conv_transpose (`bool`, default `False`):
+ option to use a convolution transpose.
+ out_channels (`int`, optional):
+ number of output channels. Defaults to `channels`.
+ name (`str`, default `conv`):
+ name of the upsampling 2D layer.
+ """
+
+ def __init__(
+ self,
+ channels: int,
+ use_conv: bool = False,
+ use_conv_transpose: bool = False,
+ out_channels: Optional[int] = None,
+ name: str = "conv",
+ kernel_size: Optional[int] = None,
+ padding=1,
+ norm_type=None,
+ eps=None,
+ elementwise_affine=None,
+ bias=True,
+ interpolate=True,
+ ):
+ super().__init__()
+ self.channels = channels
+ self.out_channels = out_channels or channels
+ self.use_conv = use_conv
+ self.use_conv_transpose = use_conv_transpose
+ self.name = name
+ self.interpolate = interpolate
+
+ if norm_type == "ln_norm":
+ self.norm = nn.LayerNorm(channels, eps, elementwise_affine)
+ elif norm_type == "rms_norm":
+ self.norm = RMSNorm(channels, eps, elementwise_affine)
+ elif norm_type is None:
+ self.norm = None
+ else:
+ raise ValueError(f"unknown norm_type: {norm_type}")
+
+ conv = None
+ if use_conv_transpose:
+ if kernel_size is None:
+ kernel_size = 4
+ conv = nn.ConvTranspose2d(
+ channels, self.out_channels, kernel_size=kernel_size, stride=2, padding=padding, bias=bias
+ )
+ elif use_conv:
+ if kernel_size is None:
+ kernel_size = 3
+ conv = nn.Conv2d(self.channels, self.out_channels, kernel_size=kernel_size, padding=padding, bias=bias)
+
+ # TODO(Suraj, Patrick) - clean up after weight dicts are correctly renamed
+ if name == "conv":
+ self.conv = conv
+ else:
+ self.Conv2d_0 = conv
+
+ def forward(self, hidden_states: torch.Tensor, output_size: Optional[int] = None, *args, **kwargs) -> torch.Tensor:
+ if len(args) > 0 or kwargs.get("scale", None) is not None:
+ deprecation_message = "The `scale` argument is deprecated and will be ignored. Please remove it, as passing it will raise an error in the future. `scale` should directly be passed while calling the underlying pipeline component i.e., via `cross_attention_kwargs`."
+ deprecate("scale", "1.0.0", deprecation_message)
+
+ assert hidden_states.shape[1] == self.channels
+
+ if self.norm is not None:
+ hidden_states = self.norm(hidden_states.permute(0, 2, 3, 1)).permute(0, 3, 1, 2)
+
+ if self.use_conv_transpose:
+ return self.conv(hidden_states)
+
+ # Cast to float32 to as 'upsample_nearest2d_out_frame' op does not support bfloat16 until PyTorch 2.1
+ # https://github.com/pytorch/pytorch/issues/86679#issuecomment-1783978767
+ dtype = hidden_states.dtype
+ if dtype == torch.bfloat16:
+ hidden_states = hidden_states.to(torch.float32)
+
+ # upsample_nearest_nhwc fails with large batch sizes. see https://github.com/huggingface/diffusers/issues/984
+ if hidden_states.shape[0] >= 64:
+ hidden_states = hidden_states.contiguous()
+
+ # if `output_size` is passed we force the interpolation output
+ # size and do not make use of `scale_factor=2`
+ if self.interpolate:
+ # upsample_nearest_nhwc also fails when the number of output elements is large
+ # https://github.com/pytorch/pytorch/issues/141831
+ scale_factor = (
+ 2 if output_size is None else max([f / s for f, s in zip(output_size, hidden_states.shape[-2:])])
+ )
+ if hidden_states.numel() * scale_factor > pow(2, 31):
+ hidden_states = hidden_states.contiguous()
+
+ if output_size is None:
+ hidden_states = F.interpolate(hidden_states, scale_factor=2.0, mode="nearest")
+ else:
+ hidden_states = F.interpolate(hidden_states, size=output_size, mode="nearest")
+
+ # Cast back to original dtype
+ if dtype == torch.bfloat16:
+ hidden_states = hidden_states.to(dtype)
+
+ # TODO(Suraj, Patrick) - clean up after weight dicts are correctly renamed
+ if self.use_conv:
+ if self.name == "conv":
+ hidden_states = self.conv(hidden_states)
+ else:
+ hidden_states = self.Conv2d_0(hidden_states)
+
+ return hidden_states
+
+
+class Attention(nn.Module):
+ r"""
+ A cross attention layer.
+
+ Parameters:
+ query_dim (`int`):
+ The number of channels in the query.
+ cross_attention_dim (`int`, *optional*):
+ The number of channels in the encoder_hidden_states. If not given, defaults to `query_dim`.
+ heads (`int`, *optional*, defaults to 8):
+ The number of heads to use for multi-head attention.
+ kv_heads (`int`, *optional*, defaults to `None`):
+ The number of key and value heads to use for multi-head attention. Defaults to `heads`. If
+ `kv_heads=heads`, the model will use Multi Head Attention (MHA), if `kv_heads=1` the model will use Multi
+ Query Attention (MQA) otherwise GQA is used.
+ dim_head (`int`, *optional*, defaults to 64):
+ The number of channels in each head.
+ dropout (`float`, *optional*, defaults to 0.0):
+ The dropout probability to use.
+ bias (`bool`, *optional*, defaults to False):
+ Set to `True` for the query, key, and value linear layers to contain a bias parameter.
+ upcast_attention (`bool`, *optional*, defaults to False):
+ Set to `True` to upcast the attention computation to `float32`.
+ upcast_softmax (`bool`, *optional*, defaults to False):
+ Set to `True` to upcast the softmax computation to `float32`.
+ cross_attention_norm (`str`, *optional*, defaults to `None`):
+ The type of normalization to use for the cross attention. Can be `None`, `layer_norm`, or `group_norm`.
+ cross_attention_norm_num_groups (`int`, *optional*, defaults to 32):
+ The number of groups to use for the group norm in the cross attention.
+ added_kv_proj_dim (`int`, *optional*, defaults to `None`):
+ The number of channels to use for the added key and value projections. If `None`, no projection is used.
+ norm_num_groups (`int`, *optional*, defaults to `None`):
+ The number of groups to use for the group norm in the attention.
+ spatial_norm_dim (`int`, *optional*, defaults to `None`):
+ The number of channels to use for the spatial normalization.
+ out_bias (`bool`, *optional*, defaults to `True`):
+ Set to `True` to use a bias in the output linear layer.
+ scale_qk (`bool`, *optional*, defaults to `True`):
+ Set to `True` to scale the query and key by `1 / sqrt(dim_head)`.
+ only_cross_attention (`bool`, *optional*, defaults to `False`):
+ Set to `True` to only use cross attention and not added_kv_proj_dim. Can only be set to `True` if
+ `added_kv_proj_dim` is not `None`.
+ eps (`float`, *optional*, defaults to 1e-5):
+ An additional value added to the denominator in group normalization that is used for numerical stability.
+ rescale_output_factor (`float`, *optional*, defaults to 1.0):
+ A factor to rescale the output by dividing it with this value.
+ residual_connection (`bool`, *optional*, defaults to `False`):
+ Set to `True` to add the residual connection to the output.
+ _from_deprecated_attn_block (`bool`, *optional*, defaults to `False`):
+ Set to `True` if the attention block is loaded from a deprecated state dict.
+ processor (`AttnProcessor`, *optional*, defaults to `None`):
+ The attention processor to use. If `None`, defaults to `AttnProcessor2_0` if `torch 2.x` is used and
+ `AttnProcessor` otherwise.
+ """
+
+ def __init__(
+ self,
+ query_dim: int,
+ cross_attention_dim: Optional[int] = None,
+ heads: int = 8,
+ kv_heads: Optional[int] = None,
+ dim_head: int = 64,
+ dropout: float = 0.0,
+ bias: bool = False,
+ upcast_attention: bool = False,
+ upcast_softmax: bool = False,
+ cross_attention_norm: Optional[str] = None,
+ cross_attention_norm_num_groups: int = 32,
+ qk_norm: Optional[str] = None,
+ added_kv_proj_dim: Optional[int] = None,
+ added_proj_bias: Optional[bool] = True,
+ norm_num_groups: Optional[int] = None,
+ spatial_norm_dim: Optional[int] = None,
+ out_bias: bool = True,
+ scale_qk: bool = True,
+ only_cross_attention: bool = False,
+ eps: float = 1e-5,
+ rescale_output_factor: float = 1.0,
+ residual_connection: bool = False,
+ _from_deprecated_attn_block: bool = False,
+ processor: Optional["AttnProcessor"] = None,
+ out_dim: int = None,
+ out_context_dim: int = None,
+ context_pre_only=None,
+ pre_only=False,
+ elementwise_affine: bool = True,
+ is_causal: bool = False,
+ ):
+ super().__init__()
+
+ # To prevent circular import.
+ # from .normalization import FP32LayerNorm, LpNorm, RMSNorm
+
+ self.inner_dim = out_dim if out_dim is not None else dim_head * heads
+ self.inner_kv_dim = self.inner_dim if kv_heads is None else dim_head * kv_heads
+ self.query_dim = query_dim
+ self.use_bias = bias
+ self.is_cross_attention = cross_attention_dim is not None
+ self.cross_attention_dim = cross_attention_dim if cross_attention_dim is not None else query_dim
+ self.upcast_attention = upcast_attention
+ self.upcast_softmax = upcast_softmax
+ self.rescale_output_factor = rescale_output_factor
+ self.residual_connection = residual_connection
+ self.dropout = dropout
+ self.fused_projections = False
+ self.out_dim = out_dim if out_dim is not None else query_dim
+ self.out_context_dim = out_context_dim if out_context_dim is not None else query_dim
+ self.context_pre_only = context_pre_only
+ self.pre_only = pre_only
+ self.is_causal = is_causal
+
+ # we make use of this private variable to know whether this class is loaded
+ # with an deprecated state dict so that we can convert it on the fly
+ self._from_deprecated_attn_block = _from_deprecated_attn_block
+
+ self.scale_qk = scale_qk
+ self.scale = dim_head**-0.5 if self.scale_qk else 1.0
+
+ self.heads = out_dim // dim_head if out_dim is not None else heads
+ # for slice_size > 0 the attention score computation
+ # is split across the batch axis to save memory
+ # You can set slice_size with `set_attention_slice`
+ self.sliceable_head_dim = heads
+
+ self.added_kv_proj_dim = added_kv_proj_dim
+ self.only_cross_attention = only_cross_attention
+
+ if self.added_kv_proj_dim is None and self.only_cross_attention:
+ raise ValueError(
+ "`only_cross_attention` can only be set to True if `added_kv_proj_dim` is not None. Make sure to set either `only_cross_attention=False` or define `added_kv_proj_dim`."
+ )
+
+ if norm_num_groups is not None:
+ self.group_norm = nn.GroupNorm(num_channels=query_dim, num_groups=norm_num_groups, eps=eps, affine=True)
+ else:
+ self.group_norm = None
+
+ if spatial_norm_dim is not None:
+ self.spatial_norm = SpatialNorm(f_channels=query_dim, zq_channels=spatial_norm_dim)
+ else:
+ self.spatial_norm = None
+
+ if qk_norm is None:
+ self.norm_q = None
+ self.norm_k = None
+ elif qk_norm == "layer_norm":
+ self.norm_q = nn.LayerNorm(dim_head, eps=eps, elementwise_affine=elementwise_affine)
+ self.norm_k = nn.LayerNorm(dim_head, eps=eps, elementwise_affine=elementwise_affine)
+ elif qk_norm == "fp32_layer_norm":
+ self.norm_q = FP32LayerNorm(dim_head, elementwise_affine=False, bias=False, eps=eps)
+ self.norm_k = FP32LayerNorm(dim_head, elementwise_affine=False, bias=False, eps=eps)
+ elif qk_norm == "layer_norm_across_heads":
+ # Lumina applies qk norm across all heads
+ self.norm_q = nn.LayerNorm(dim_head * heads, eps=eps)
+ self.norm_k = nn.LayerNorm(dim_head * kv_heads, eps=eps)
+ elif qk_norm == "rms_norm":
+ self.norm_q = RMSNorm(dim_head, eps=eps, elementwise_affine=elementwise_affine)
+ self.norm_k = RMSNorm(dim_head, eps=eps, elementwise_affine=elementwise_affine)
+ elif qk_norm == "rms_norm_across_heads":
+ # LTX applies qk norm across all heads
+ self.norm_q = RMSNorm(dim_head * heads, eps=eps)
+ self.norm_k = RMSNorm(dim_head * kv_heads, eps=eps)
+ elif qk_norm == "l2":
+ self.norm_q = LpNorm(p=2, dim=-1, eps=eps)
+ self.norm_k = LpNorm(p=2, dim=-1, eps=eps)
+ else:
+ raise ValueError(
+ f"unknown qk_norm: {qk_norm}. Should be one of None, 'layer_norm', 'fp32_layer_norm', 'layer_norm_across_heads', 'rms_norm', 'rms_norm_across_heads', 'l2'."
+ )
+
+ if cross_attention_norm is None:
+ self.norm_cross = None
+ elif cross_attention_norm == "layer_norm":
+ self.norm_cross = nn.LayerNorm(self.cross_attention_dim)
+ elif cross_attention_norm == "group_norm":
+ if self.added_kv_proj_dim is not None:
+ # The given `encoder_hidden_states` are initially of shape
+ # (batch_size, seq_len, added_kv_proj_dim) before being projected
+ # to (batch_size, seq_len, cross_attention_dim). The norm is applied
+ # before the projection, so we need to use `added_kv_proj_dim` as
+ # the number of channels for the group norm.
+ norm_cross_num_channels = added_kv_proj_dim
+ else:
+ norm_cross_num_channels = self.cross_attention_dim
+
+ self.norm_cross = nn.GroupNorm(
+ num_channels=norm_cross_num_channels, num_groups=cross_attention_norm_num_groups, eps=1e-5, affine=True
+ )
+ else:
+ raise ValueError(
+ f"unknown cross_attention_norm: {cross_attention_norm}. Should be None, 'layer_norm' or 'group_norm'"
+ )
+
+ self.to_q = nn.Linear(query_dim, self.inner_dim, bias=bias)
+
+ if not self.only_cross_attention:
+ # only relevant for the `AddedKVProcessor` classes
+ self.to_k = nn.Linear(self.cross_attention_dim, self.inner_kv_dim, bias=bias)
+ self.to_v = nn.Linear(self.cross_attention_dim, self.inner_kv_dim, bias=bias)
+ else:
+ self.to_k = None
+ self.to_v = None
+
+ self.added_proj_bias = added_proj_bias
+ if self.added_kv_proj_dim is not None:
+ self.add_k_proj = nn.Linear(added_kv_proj_dim, self.inner_kv_dim, bias=added_proj_bias)
+ self.add_v_proj = nn.Linear(added_kv_proj_dim, self.inner_kv_dim, bias=added_proj_bias)
+ if self.context_pre_only is not None:
+ self.add_q_proj = nn.Linear(added_kv_proj_dim, self.inner_dim, bias=added_proj_bias)
+ else:
+ self.add_q_proj = None
+ self.add_k_proj = None
+ self.add_v_proj = None
+
+ if not self.pre_only:
+ self.to_out = nn.ModuleList([])
+ self.to_out.append(nn.Linear(self.inner_dim, self.out_dim, bias=out_bias))
+ self.to_out.append(nn.Dropout(dropout))
+ else:
+ self.to_out = None
+
+ if self.context_pre_only is not None and not self.context_pre_only:
+ self.to_add_out = nn.Linear(self.inner_dim, self.out_context_dim, bias=out_bias)
+ else:
+ self.to_add_out = None
+
+ if qk_norm is not None and added_kv_proj_dim is not None:
+ if qk_norm == "layer_norm":
+ self.norm_added_q = nn.LayerNorm(dim_head, eps=eps, elementwise_affine=elementwise_affine)
+ self.norm_added_k = nn.LayerNorm(dim_head, eps=eps, elementwise_affine=elementwise_affine)
+ elif qk_norm == "fp32_layer_norm":
+ self.norm_added_q = FP32LayerNorm(dim_head, elementwise_affine=False, bias=False, eps=eps)
+ self.norm_added_k = FP32LayerNorm(dim_head, elementwise_affine=False, bias=False, eps=eps)
+ elif qk_norm == "rms_norm":
+ self.norm_added_q = RMSNorm(dim_head, eps=eps)
+ self.norm_added_k = RMSNorm(dim_head, eps=eps)
+ elif qk_norm == "rms_norm_across_heads":
+ # Wan applies qk norm across all heads
+ # Wan also doesn't apply a q norm
+ self.norm_added_q = None
+ self.norm_added_k = RMSNorm(dim_head * kv_heads, eps=eps)
+ else:
+ raise ValueError(
+ f"unknown qk_norm: {qk_norm}. Should be one of `None,'layer_norm','fp32_layer_norm','rms_norm'`"
+ )
+ else:
+ self.norm_added_q = None
+ self.norm_added_k = None
+
+ # set attention processor
+ # We use the AttnProcessor2_0 by default when torch 2.x is used which uses
+ # torch.nn.functional.scaled_dot_product_attention for native Flash/memory_efficient_attention
+ # but only if it has the default `scale` argument. TODO remove scale_qk check when we move to torch 2.1
+ if processor is None:
+ processor = (
+ AttnProcessor2_0() if hasattr(F, "scaled_dot_product_attention") and self.scale_qk else AttnProcessor()
+ )
+ self.set_processor(processor)
+
+ def set_use_xla_flash_attention(
+ self,
+ use_xla_flash_attention: bool,
+ partition_spec: Optional[Tuple[Optional[str], ...]] = None,
+ is_flux=False,
+ ) -> None:
+ r"""
+ Set whether to use xla flash attention from `torch_xla` or not.
+
+ Args:
+ use_xla_flash_attention (`bool`):
+ Whether to use pallas flash attention kernel from `torch_xla` or not.
+ partition_spec (`Tuple[]`, *optional*):
+ Specify the partition specification if using SPMD. Otherwise None.
+ """
+ if use_xla_flash_attention:
+ if not is_torch_xla_available:
+ raise "torch_xla is not available"
+ elif is_torch_xla_version("<", "2.3"):
+ raise "flash attention pallas kernel is supported from torch_xla version 2.3"
+ elif is_spmd() and is_torch_xla_version("<", "2.4"):
+ raise "flash attention pallas kernel using SPMD is supported from torch_xla version 2.4"
+ else:
+ if is_flux:
+ processor = XLAFluxFlashAttnProcessor2_0(partition_spec)
+ else:
+ processor = XLAFlashAttnProcessor2_0(partition_spec)
+ else:
+ processor = (
+ AttnProcessor2_0() if hasattr(F, "scaled_dot_product_attention") and self.scale_qk else AttnProcessor()
+ )
+ self.set_processor(processor)
+
+ def set_use_npu_flash_attention(self, use_npu_flash_attention: bool) -> None:
+ r"""
+ Set whether to use npu flash attention from `torch_npu` or not.
+
+ """
+ if use_npu_flash_attention:
+ processor = AttnProcessorNPU()
+ else:
+ # set attention processor
+ # We use the AttnProcessor2_0 by default when torch 2.x is used which uses
+ # torch.nn.functional.scaled_dot_product_attention for native Flash/memory_efficient_attention
+ # but only if it has the default `scale` argument. TODO remove scale_qk check when we move to torch 2.1
+ processor = (
+ AttnProcessor2_0() if hasattr(F, "scaled_dot_product_attention") and self.scale_qk else AttnProcessor()
+ )
+ self.set_processor(processor)
+
+ def set_use_memory_efficient_attention_xformers(
+ self, use_memory_efficient_attention_xformers: bool, attention_op: Optional[Callable] = None
+ ) -> None:
+ r"""
+ Set whether to use memory efficient attention from `xformers` or not.
+
+ Args:
+ use_memory_efficient_attention_xformers (`bool`):
+ Whether to use memory efficient attention from `xformers` or not.
+ attention_op (`Callable`, *optional*):
+ The attention operation to use. Defaults to `None` which uses the default attention operation from
+ `xformers`.
+ """
+ is_custom_diffusion = hasattr(self, "processor") and isinstance(
+ self.processor,
+ (CustomDiffusionAttnProcessor, CustomDiffusionXFormersAttnProcessor, CustomDiffusionAttnProcessor2_0),
+ )
+ is_added_kv_processor = hasattr(self, "processor") and isinstance(
+ self.processor,
+ (
+ AttnAddedKVProcessor,
+ AttnAddedKVProcessor2_0,
+ SlicedAttnAddedKVProcessor,
+ XFormersAttnAddedKVProcessor,
+ ),
+ )
+ is_ip_adapter = hasattr(self, "processor") and isinstance(
+ self.processor,
+ (IPAdapterAttnProcessor, IPAdapterAttnProcessor2_0, IPAdapterXFormersAttnProcessor),
+ )
+ is_joint_processor = hasattr(self, "processor") and isinstance(
+ self.processor,
+ (
+ JointAttnProcessor2_0,
+ XFormersJointAttnProcessor,
+ ),
+ )
+
+ if use_memory_efficient_attention_xformers:
+ if is_added_kv_processor and is_custom_diffusion:
+ raise NotImplementedError(
+ f"Memory efficient attention is currently not supported for custom diffusion for attention processor type {self.processor}"
+ )
+ if not is_xformers_available():
+ raise ModuleNotFoundError(
+ (
+ "Refer to https://github.com/facebookresearch/xformers for more information on how to install"
+ " xformers"
+ ),
+ name="xformers",
+ )
+ elif not torch.cuda.is_available():
+ raise ValueError(
+ "torch.cuda.is_available() should be True but is False. xformers' memory efficient attention is"
+ " only available for GPU "
+ )
+ else:
+ try:
+ # Make sure we can run the memory efficient attention
+ dtype = None
+ if attention_op is not None:
+ op_fw, op_bw = attention_op
+ dtype, *_ = op_fw.SUPPORTED_DTYPES
+ q = torch.randn((1, 2, 40), device="cuda", dtype=dtype)
+ _ = xformers.ops.memory_efficient_attention(q, q, q)
+ except Exception as e:
+ raise e
+
+ if is_custom_diffusion:
+ processor = CustomDiffusionXFormersAttnProcessor(
+ train_kv=self.processor.train_kv,
+ train_q_out=self.processor.train_q_out,
+ hidden_size=self.processor.hidden_size,
+ cross_attention_dim=self.processor.cross_attention_dim,
+ attention_op=attention_op,
+ )
+ processor.load_state_dict(self.processor.state_dict())
+ if hasattr(self.processor, "to_k_custom_diffusion"):
+ processor.to(self.processor.to_k_custom_diffusion.weight.device)
+ elif is_added_kv_processor:
+ # TODO(Patrick, Suraj, William) - currently xformers doesn't work for UnCLIP
+ # which uses this type of cross attention ONLY because the attention mask of format
+ # [0, ..., -10.000, ..., 0, ...,] is not supported
+ # throw warning
+ logger.info(
+ "Memory efficient attention with `xformers` might currently not work correctly if an attention mask is required for the attention operation."
+ )
+ processor = XFormersAttnAddedKVProcessor(attention_op=attention_op)
+ elif is_ip_adapter:
+ processor = IPAdapterXFormersAttnProcessor(
+ hidden_size=self.processor.hidden_size,
+ cross_attention_dim=self.processor.cross_attention_dim,
+ num_tokens=self.processor.num_tokens,
+ scale=self.processor.scale,
+ attention_op=attention_op,
+ )
+ processor.load_state_dict(self.processor.state_dict())
+ if hasattr(self.processor, "to_k_ip"):
+ processor.to(
+ device=self.processor.to_k_ip[0].weight.device, dtype=self.processor.to_k_ip[0].weight.dtype
+ )
+ elif is_joint_processor:
+ processor = XFormersJointAttnProcessor(attention_op=attention_op)
+ else:
+ processor = XFormersAttnProcessor(attention_op=attention_op)
+ else:
+ if is_custom_diffusion:
+ attn_processor_class = (
+ CustomDiffusionAttnProcessor2_0
+ if hasattr(F, "scaled_dot_product_attention")
+ else CustomDiffusionAttnProcessor
+ )
+ processor = attn_processor_class(
+ train_kv=self.processor.train_kv,
+ train_q_out=self.processor.train_q_out,
+ hidden_size=self.processor.hidden_size,
+ cross_attention_dim=self.processor.cross_attention_dim,
+ )
+ processor.load_state_dict(self.processor.state_dict())
+ if hasattr(self.processor, "to_k_custom_diffusion"):
+ processor.to(self.processor.to_k_custom_diffusion.weight.device)
+ elif is_ip_adapter:
+ processor = IPAdapterAttnProcessor2_0(
+ hidden_size=self.processor.hidden_size,
+ cross_attention_dim=self.processor.cross_attention_dim,
+ num_tokens=self.processor.num_tokens,
+ scale=self.processor.scale,
+ )
+ processor.load_state_dict(self.processor.state_dict())
+ if hasattr(self.processor, "to_k_ip"):
+ processor.to(
+ device=self.processor.to_k_ip[0].weight.device, dtype=self.processor.to_k_ip[0].weight.dtype
+ )
+ else:
+ # set attention processor
+ # We use the AttnProcessor2_0 by default when torch 2.x is used which uses
+ # torch.nn.functional.scaled_dot_product_attention for native Flash/memory_efficient_attention
+ # but only if it has the default `scale` argument. TODO remove scale_qk check when we move to torch 2.1
+ processor = (
+ AttnProcessor2_0()
+ if hasattr(F, "scaled_dot_product_attention") and self.scale_qk
+ else AttnProcessor()
+ )
+
+ self.set_processor(processor)
+
+ def set_attention_slice(self, slice_size: int) -> None:
+ r"""
+ Set the slice size for attention computation.
+
+ Args:
+ slice_size (`int`):
+ The slice size for attention computation.
+ """
+ if slice_size is not None and slice_size > self.sliceable_head_dim:
+ raise ValueError(f"slice_size {slice_size} has to be smaller or equal to {self.sliceable_head_dim}.")
+
+ if slice_size is not None and self.added_kv_proj_dim is not None:
+ processor = SlicedAttnAddedKVProcessor(slice_size)
+ elif slice_size is not None:
+ processor = SlicedAttnProcessor(slice_size)
+ elif self.added_kv_proj_dim is not None:
+ processor = AttnAddedKVProcessor()
+ else:
+ # set attention processor
+ # We use the AttnProcessor2_0 by default when torch 2.x is used which uses
+ # torch.nn.functional.scaled_dot_product_attention for native Flash/memory_efficient_attention
+ # but only if it has the default `scale` argument. TODO remove scale_qk check when we move to torch 2.1
+ processor = (
+ AttnProcessor2_0() if hasattr(F, "scaled_dot_product_attention") and self.scale_qk else AttnProcessor()
+ )
+
+ self.set_processor(processor)
+
+ def set_processor(self, processor: "AttnProcessor") -> None:
+ r"""
+ Set the attention processor to use.
+
+ Args:
+ processor (`AttnProcessor`):
+ The attention processor to use.
+ """
+ # if current processor is in `self._modules` and if passed `processor` is not, we need to
+ # pop `processor` from `self._modules`
+ if (
+ hasattr(self, "processor")
+ and isinstance(self.processor, torch.nn.Module)
+ and not isinstance(processor, torch.nn.Module)
+ ):
+ logger.info(f"You are removing possibly trained weights of {self.processor} with {processor}")
+ self._modules.pop("processor")
+
+ self.processor = processor
+
+ def get_processor(self, return_deprecated_lora: bool = False) -> "AttentionProcessor":
+ r"""
+ Get the attention processor in use.
+
+ Args:
+ return_deprecated_lora (`bool`, *optional*, defaults to `False`):
+ Set to `True` to return the deprecated LoRA attention processor.
+
+ Returns:
+ "AttentionProcessor": The attention processor in use.
+ """
+ if not return_deprecated_lora:
+ return self.processor
+
+ def forward(
+ self,
+ hidden_states: torch.Tensor,
+ encoder_hidden_states: Optional[torch.Tensor] = None,
+ attention_mask: Optional[torch.Tensor] = None,
+ **cross_attention_kwargs,
+ ) -> torch.Tensor:
+ r"""
+ The forward method of the `Attention` class.
+
+ Args:
+ hidden_states (`torch.Tensor`):
+ The hidden states of the query.
+ encoder_hidden_states (`torch.Tensor`, *optional*):
+ The hidden states of the encoder.
+ attention_mask (`torch.Tensor`, *optional*):
+ The attention mask to use. If `None`, no mask is applied.
+ **cross_attention_kwargs:
+ Additional keyword arguments to pass along to the cross attention.
+
+ Returns:
+ `torch.Tensor`: The output of the attention layer.
+ """
+ # The `Attention` class can call different attention processors / attention functions
+ # here we simply pass along all tensors to the selected processor class
+ # For standard processors that are defined here, `**cross_attention_kwargs` is empty
+
+ attn_parameters = set(inspect.signature(self.processor.__call__).parameters.keys())
+ quiet_attn_parameters = {"ip_adapter_masks", "ip_hidden_states"}
+ unused_kwargs = [
+ k for k, _ in cross_attention_kwargs.items() if k not in attn_parameters and k not in quiet_attn_parameters
+ ]
+ if len(unused_kwargs) > 0:
+ logger.warning(
+ f"cross_attention_kwargs {unused_kwargs} are not expected by {self.processor.__class__.__name__} and will be ignored."
+ )
+ cross_attention_kwargs = {k: w for k, w in cross_attention_kwargs.items() if k in attn_parameters}
+
+ return self.processor(
+ self,
+ hidden_states,
+ encoder_hidden_states=encoder_hidden_states,
+ attention_mask=attention_mask,
+ **cross_attention_kwargs,
+ )
+
+ def batch_to_head_dim(self, tensor: torch.Tensor) -> torch.Tensor:
+ r"""
+ Reshape the tensor from `[batch_size, seq_len, dim]` to `[batch_size // heads, seq_len, dim * heads]`. `heads`
+ is the number of heads initialized while constructing the `Attention` class.
+
+ Args:
+ tensor (`torch.Tensor`): The tensor to reshape.
+
+ Returns:
+ `torch.Tensor`: The reshaped tensor.
+ """
+ head_size = self.heads
+ batch_size, seq_len, dim = tensor.shape
+ tensor = tensor.reshape(batch_size // head_size, head_size, seq_len, dim)
+ tensor = tensor.permute(0, 2, 1, 3).reshape(batch_size // head_size, seq_len, dim * head_size)
+ return tensor
+
+ def head_to_batch_dim(self, tensor: torch.Tensor, out_dim: int = 3) -> torch.Tensor:
+ r"""
+ Reshape the tensor from `[batch_size, seq_len, dim]` to `[batch_size, seq_len, heads, dim // heads]` `heads` is
+ the number of heads initialized while constructing the `Attention` class.
+
+ Args:
+ tensor (`torch.Tensor`): The tensor to reshape.
+ out_dim (`int`, *optional*, defaults to `3`): The output dimension of the tensor. If `3`, the tensor is
+ reshaped to `[batch_size * heads, seq_len, dim // heads]`.
+
+ Returns:
+ `torch.Tensor`: The reshaped tensor.
+ """
+ head_size = self.heads
+ if tensor.ndim == 3:
+ batch_size, seq_len, dim = tensor.shape
+ extra_dim = 1
+ else:
+ batch_size, extra_dim, seq_len, dim = tensor.shape
+ tensor = tensor.reshape(batch_size, seq_len * extra_dim, head_size, dim // head_size)
+ tensor = tensor.permute(0, 2, 1, 3)
+
+ if out_dim == 3:
+ tensor = tensor.reshape(batch_size * head_size, seq_len * extra_dim, dim // head_size)
+
+ return tensor
+
+ def get_attention_scores(
+ self, query: torch.Tensor, key: torch.Tensor, attention_mask: Optional[torch.Tensor] = None
+ ) -> torch.Tensor:
+ r"""
+ Compute the attention scores.
+
+ Args:
+ query (`torch.Tensor`): The query tensor.
+ key (`torch.Tensor`): The key tensor.
+ attention_mask (`torch.Tensor`, *optional*): The attention mask to use. If `None`, no mask is applied.
+
+ Returns:
+ `torch.Tensor`: The attention probabilities/scores.
+ """
+ dtype = query.dtype
+ if self.upcast_attention:
+ query = query.float()
+ key = key.float()
+
+ if attention_mask is None:
+ baddbmm_input = torch.empty(
+ query.shape[0], query.shape[1], key.shape[1], dtype=query.dtype, device=query.device
+ )
+ beta = 0
+ else:
+ baddbmm_input = attention_mask
+ beta = 1
+
+ attention_scores = torch.baddbmm(
+ baddbmm_input,
+ query,
+ key.transpose(-1, -2),
+ beta=beta,
+ alpha=self.scale,
+ )
+ del baddbmm_input
+
+ if self.upcast_softmax:
+ attention_scores = attention_scores.float()
+
+ attention_probs = attention_scores.softmax(dim=-1)
+ del attention_scores
+
+ attention_probs = attention_probs.to(dtype)
+
+ return attention_probs
+
+ def prepare_attention_mask(
+ self, attention_mask: torch.Tensor, target_length: int, batch_size: int, out_dim: int = 3
+ ) -> torch.Tensor:
+ r"""
+ Prepare the attention mask for the attention computation.
+
+ Args:
+ attention_mask (`torch.Tensor`):
+ The attention mask to prepare.
+ target_length (`int`):
+ The target length of the attention mask. This is the length of the attention mask after padding.
+ batch_size (`int`):
+ The batch size, which is used to repeat the attention mask.
+ out_dim (`int`, *optional*, defaults to `3`):
+ The output dimension of the attention mask. Can be either `3` or `4`.
+
+ Returns:
+ `torch.Tensor`: The prepared attention mask.
+ """
+ head_size = self.heads
+ if attention_mask is None:
+ return attention_mask
+
+ current_length: int = attention_mask.shape[-1]
+ if current_length != target_length:
+ if attention_mask.device.type == "mps":
+ # HACK: MPS: Does not support padding by greater than dimension of input tensor.
+ # Instead, we can manually construct the padding tensor.
+ padding_shape = (attention_mask.shape[0], attention_mask.shape[1], target_length)
+ padding = torch.zeros(padding_shape, dtype=attention_mask.dtype, device=attention_mask.device)
+ attention_mask = torch.cat([attention_mask, padding], dim=2)
+ else:
+ # TODO: for pipelines such as stable-diffusion, padding cross-attn mask:
+ # we want to instead pad by (0, remaining_length), where remaining_length is:
+ # remaining_length: int = target_length - current_length
+ # TODO: re-enable tests/models/test_models_unet_2d_condition.py#test_model_xattn_padding
+ attention_mask = F.pad(attention_mask, (0, target_length), value=0.0)
+
+ if out_dim == 3:
+ if attention_mask.shape[0] < batch_size * head_size:
+ attention_mask = attention_mask.repeat_interleave(
+ head_size, dim=0, output_size=attention_mask.shape[0] * head_size
+ )
+ elif out_dim == 4:
+ attention_mask = attention_mask.unsqueeze(1)
+ attention_mask = attention_mask.repeat_interleave(
+ head_size, dim=1, output_size=attention_mask.shape[1] * head_size
+ )
+
+ return attention_mask
+
+ def norm_encoder_hidden_states(self, encoder_hidden_states: torch.Tensor) -> torch.Tensor:
+ r"""
+ Normalize the encoder hidden states. Requires `self.norm_cross` to be specified when constructing the
+ `Attention` class.
+
+ Args:
+ encoder_hidden_states (`torch.Tensor`): Hidden states of the encoder.
+
+ Returns:
+ `torch.Tensor`: The normalized encoder hidden states.
+ """
+ assert self.norm_cross is not None, "self.norm_cross must be defined to call self.norm_encoder_hidden_states"
+
+ if isinstance(self.norm_cross, nn.LayerNorm):
+ encoder_hidden_states = self.norm_cross(encoder_hidden_states)
+ elif isinstance(self.norm_cross, nn.GroupNorm):
+ # Group norm norms along the channels dimension and expects
+ # input to be in the shape of (N, C, *). In this case, we want
+ # to norm along the hidden dimension, so we need to move
+ # (batch_size, sequence_length, hidden_size) ->
+ # (batch_size, hidden_size, sequence_length)
+ encoder_hidden_states = encoder_hidden_states.transpose(1, 2)
+ encoder_hidden_states = self.norm_cross(encoder_hidden_states)
+ encoder_hidden_states = encoder_hidden_states.transpose(1, 2)
+ else:
+ assert False
+
+ return encoder_hidden_states
+
+ @torch.no_grad()
+ def fuse_projections(self, fuse=True):
+ device = self.to_q.weight.data.device
+ dtype = self.to_q.weight.data.dtype
+
+ if not self.is_cross_attention:
+ # fetch weight matrices.
+ concatenated_weights = torch.cat([self.to_q.weight.data, self.to_k.weight.data, self.to_v.weight.data])
+ in_features = concatenated_weights.shape[1]
+ out_features = concatenated_weights.shape[0]
+
+ # create a new single projection layer and copy over the weights.
+ self.to_qkv = nn.Linear(in_features, out_features, bias=self.use_bias, device=device, dtype=dtype)
+ self.to_qkv.weight.copy_(concatenated_weights)
+ if self.use_bias:
+ concatenated_bias = torch.cat([self.to_q.bias.data, self.to_k.bias.data, self.to_v.bias.data])
+ self.to_qkv.bias.copy_(concatenated_bias)
+
+ else:
+ concatenated_weights = torch.cat([self.to_k.weight.data, self.to_v.weight.data])
+ in_features = concatenated_weights.shape[1]
+ out_features = concatenated_weights.shape[0]
+
+ self.to_kv = nn.Linear(in_features, out_features, bias=self.use_bias, device=device, dtype=dtype)
+ self.to_kv.weight.copy_(concatenated_weights)
+ if self.use_bias:
+ concatenated_bias = torch.cat([self.to_k.bias.data, self.to_v.bias.data])
+ self.to_kv.bias.copy_(concatenated_bias)
+
+ # handle added projections for SD3 and others.
+ if (
+ getattr(self, "add_q_proj", None) is not None
+ and getattr(self, "add_k_proj", None) is not None
+ and getattr(self, "add_v_proj", None) is not None
+ ):
+ concatenated_weights = torch.cat(
+ [self.add_q_proj.weight.data, self.add_k_proj.weight.data, self.add_v_proj.weight.data]
+ )
+ in_features = concatenated_weights.shape[1]
+ out_features = concatenated_weights.shape[0]
+
+ self.to_added_qkv = nn.Linear(
+ in_features, out_features, bias=self.added_proj_bias, device=device, dtype=dtype
+ )
+ self.to_added_qkv.weight.copy_(concatenated_weights)
+ if self.added_proj_bias:
+ concatenated_bias = torch.cat(
+ [self.add_q_proj.bias.data, self.add_k_proj.bias.data, self.add_v_proj.bias.data]
+ )
+ self.to_added_qkv.bias.copy_(concatenated_bias)
+
+ self.fused_projections = fuse
+
+class AttnProcessor2_0:
+ r"""
+ Processor for implementing scaled dot-product attention (enabled by default if you're using PyTorch 2.0).
+ """
+
+ def __init__(self):
+ if not hasattr(F, "scaled_dot_product_attention"):
+ raise ImportError("AttnProcessor2_0 requires PyTorch 2.0, to use it, please upgrade PyTorch to 2.0.")
+
+ def __call__(
+ self,
+ attn: Attention,
+ hidden_states: torch.Tensor,
+ encoder_hidden_states: Optional[torch.Tensor] = None,
+ attention_mask: Optional[torch.Tensor] = None,
+ temb: Optional[torch.Tensor] = None,
+ *args,
+ **kwargs,
+ ) -> torch.Tensor:
+ if len(args) > 0 or kwargs.get("scale", None) is not None:
+ deprecation_message = "The `scale` argument is deprecated and will be ignored. Please remove it, as passing it will raise an error in the future. `scale` should directly be passed while calling the underlying pipeline component i.e., via `cross_attention_kwargs`."
+ deprecate("scale", "1.0.0", deprecation_message)
+
+ residual = hidden_states
+ if attn.spatial_norm is not None:
+ hidden_states = attn.spatial_norm(hidden_states, temb)
+
+ input_ndim = hidden_states.ndim
+
+ if input_ndim == 4:
+ batch_size, channel, height, width = hidden_states.shape
+ hidden_states = hidden_states.view(batch_size, channel, height * width).transpose(1, 2)
+
+ batch_size, sequence_length, _ = (
+ hidden_states.shape if encoder_hidden_states is None else encoder_hidden_states.shape
+ )
+
+ if attention_mask is not None:
+ attention_mask = attn.prepare_attention_mask(attention_mask, sequence_length, batch_size)
+ # scaled_dot_product_attention expects attention_mask shape to be
+ # (batch, heads, source_length, target_length)
+ attention_mask = attention_mask.view(batch_size, attn.heads, -1, attention_mask.shape[-1])
+
+ if attn.group_norm is not None:
+ hidden_states = attn.group_norm(hidden_states.transpose(1, 2)).transpose(1, 2)
+
+ query = attn.to_q(hidden_states)
+
+ if encoder_hidden_states is None:
+ encoder_hidden_states = hidden_states
+ elif attn.norm_cross:
+ encoder_hidden_states = attn.norm_encoder_hidden_states(encoder_hidden_states)
+
+ key = attn.to_k(encoder_hidden_states)
+ value = attn.to_v(encoder_hidden_states)
+
+ inner_dim = key.shape[-1]
+ head_dim = inner_dim // attn.heads
+
+ query = query.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
+
+ key = key.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
+ value = value.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
+
+ if attn.norm_q is not None:
+ query = attn.norm_q(query)
+ if attn.norm_k is not None:
+ key = attn.norm_k(key)
+
+ # the output of sdp = (batch, num_heads, seq_len, head_dim)
+ # TODO: add support for attn.scale when we move to Torch 2.1
+ hidden_states = F.scaled_dot_product_attention(
+ query, key, value, attn_mask=attention_mask, dropout_p=0.0, is_causal=False
+ )
+
+ hidden_states = hidden_states.transpose(1, 2).reshape(batch_size, -1, attn.heads * head_dim)
+ hidden_states = hidden_states.to(query.dtype)
+
+ # linear proj
+ hidden_states = attn.to_out[0](hidden_states)
+ # dropout
+ hidden_states = attn.to_out[1](hidden_states)
+
+ if input_ndim == 4:
+ hidden_states = hidden_states.transpose(-1, -2).reshape(batch_size, channel, height, width)
+
+ if attn.residual_connection:
+ hidden_states = hidden_states + residual
+
+ hidden_states = hidden_states / attn.rescale_output_factor
+
+ return hidden_states
+
+class UNetMidBlock2D(nn.Module):
+ """
+ A 2D UNet mid-block [`UNetMidBlock2D`] with multiple residual blocks and optional attention blocks.
+
+ Args:
+ in_channels (`int`): The number of input channels.
+ temb_channels (`int`): The number of temporal embedding channels.
+ dropout (`float`, *optional*, defaults to 0.0): The dropout rate.
+ num_layers (`int`, *optional*, defaults to 1): The number of residual blocks.
+ resnet_eps (`float`, *optional*, 1e-6 ): The epsilon value for the resnet blocks.
+ resnet_time_scale_shift (`str`, *optional*, defaults to `default`):
+ The type of normalization to apply to the time embeddings. This can help to improve the performance of the
+ model on tasks with long-range temporal dependencies.
+ resnet_act_fn (`str`, *optional*, defaults to `swish`): The activation function for the resnet blocks.
+ resnet_groups (`int`, *optional*, defaults to 32):
+ The number of groups to use in the group normalization layers of the resnet blocks.
+ attn_groups (`Optional[int]`, *optional*, defaults to None): The number of groups for the attention blocks.
+ resnet_pre_norm (`bool`, *optional*, defaults to `True`):
+ Whether to use pre-normalization for the resnet blocks.
+ add_attention (`bool`, *optional*, defaults to `True`): Whether to add attention blocks.
+ attention_head_dim (`int`, *optional*, defaults to 1):
+ Dimension of a single attention head. The number of attention heads is determined based on this value and
+ the number of input channels.
+ output_scale_factor (`float`, *optional*, defaults to 1.0): The output scale factor.
+
+ Returns:
+ `torch.Tensor`: The output of the last residual block, which is a tensor of shape `(batch_size, in_channels,
+ height, width)`.
+
+ """
+
+ def __init__(
+ self,
+ in_channels: int,
+ temb_channels: int,
+ dropout: float = 0.0,
+ num_layers: int = 1,
+ resnet_eps: float = 1e-6,
+ resnet_time_scale_shift: str = "default", # default, spatial
+ resnet_act_fn: str = "swish",
+ resnet_groups: int = 32,
+ attn_groups: Optional[int] = None,
+ resnet_pre_norm: bool = True,
+ add_attention: bool = True,
+ attention_head_dim: int = 1,
+ output_scale_factor: float = 1.0,
+ ):
+ super().__init__()
+ resnet_groups = resnet_groups if resnet_groups is not None else min(in_channels // 4, 32)
+ self.add_attention = add_attention
+
+ if attn_groups is None:
+ attn_groups = resnet_groups if resnet_time_scale_shift == "default" else None
+
+ # there is always at least one resnet
+ if resnet_time_scale_shift == "spatial":
+ resnets = [
+ ResnetBlockCondNorm2D(
+ in_channels=in_channels,
+ out_channels=in_channels,
+ temb_channels=temb_channels,
+ eps=resnet_eps,
+ groups=resnet_groups,
+ dropout=dropout,
+ time_embedding_norm="spatial",
+ non_linearity=resnet_act_fn,
+ output_scale_factor=output_scale_factor,
+ )
+ ]
+ else:
+ resnets = [
+ ResnetBlock2D(
+ in_channels=in_channels,
+ out_channels=in_channels,
+ temb_channels=temb_channels,
+ eps=resnet_eps,
+ groups=resnet_groups,
+ dropout=dropout,
+ time_embedding_norm=resnet_time_scale_shift,
+ non_linearity=resnet_act_fn,
+ output_scale_factor=output_scale_factor,
+ pre_norm=resnet_pre_norm,
+ )
+ ]
+ attentions = []
+
+ if attention_head_dim is None:
+ logger.warning(
+ f"It is not recommend to pass `attention_head_dim=None`. Defaulting `attention_head_dim` to `in_channels`: {in_channels}."
+ )
+ attention_head_dim = in_channels
+
+ for _ in range(num_layers):
+ if self.add_attention:
+ attentions.append(
+ Attention(
+ in_channels,
+ heads=in_channels // attention_head_dim,
+ dim_head=attention_head_dim,
+ rescale_output_factor=output_scale_factor,
+ eps=resnet_eps,
+ norm_num_groups=attn_groups,
+ spatial_norm_dim=temb_channels if resnet_time_scale_shift == "spatial" else None,
+ residual_connection=True,
+ bias=True,
+ upcast_softmax=True,
+ _from_deprecated_attn_block=True,
+ )
+ )
+ else:
+ attentions.append(None)
+
+ if resnet_time_scale_shift == "spatial":
+ resnets.append(
+ ResnetBlockCondNorm2D(
+ in_channels=in_channels,
+ out_channels=in_channels,
+ temb_channels=temb_channels,
+ eps=resnet_eps,
+ groups=resnet_groups,
+ dropout=dropout,
+ time_embedding_norm="spatial",
+ non_linearity=resnet_act_fn,
+ output_scale_factor=output_scale_factor,
+ )
+ )
+ else:
+ resnets.append(
+ ResnetBlock2D(
+ in_channels=in_channels,
+ out_channels=in_channels,
+ temb_channels=temb_channels,
+ eps=resnet_eps,
+ groups=resnet_groups,
+ dropout=dropout,
+ time_embedding_norm=resnet_time_scale_shift,
+ non_linearity=resnet_act_fn,
+ output_scale_factor=output_scale_factor,
+ pre_norm=resnet_pre_norm,
+ )
+ )
+
+ self.attentions = nn.ModuleList(attentions)
+ self.resnets = nn.ModuleList(resnets)
+
+ self.gradient_checkpointing = False
+
+ def forward(self, hidden_states: torch.Tensor, temb: Optional[torch.Tensor] = None) -> torch.Tensor:
+ hidden_states = self.resnets[0](hidden_states, temb)
+ for attn, resnet in zip(self.attentions, self.resnets[1:]):
+ if torch.is_grad_enabled() and self.gradient_checkpointing:
+ if attn is not None:
+ hidden_states = attn(hidden_states, temb=temb)
+ hidden_states = self._gradient_checkpointing_func(resnet, hidden_states, temb)
+ else:
+ if attn is not None:
+ hidden_states = attn(hidden_states, temb=temb)
+ hidden_states = resnet(hidden_states, temb)
+
+ return hidden_states
+
+class DownEncoderBlock2D(nn.Module):
+ def __init__(
+ self,
+ in_channels: int,
+ out_channels: int,
+ dropout: float = 0.0,
+ num_layers: int = 1,
+ resnet_eps: float = 1e-6,
+ resnet_time_scale_shift: str = "default",
+ resnet_act_fn: str = "swish",
+ resnet_groups: int = 32,
+ resnet_pre_norm: bool = True,
+ output_scale_factor: float = 1.0,
+ add_downsample: bool = True,
+ downsample_padding: int = 1,
+ ):
+ super().__init__()
+ resnets = []
+
+ for i in range(num_layers):
+ in_channels = in_channels if i == 0 else out_channels
+ if resnet_time_scale_shift == "spatial":
+ resnets.append(
+ ResnetBlockCondNorm2D(
+ in_channels=in_channels,
+ out_channels=out_channels,
+ temb_channels=None,
+ eps=resnet_eps,
+ groups=resnet_groups,
+ dropout=dropout,
+ time_embedding_norm="spatial",
+ non_linearity=resnet_act_fn,
+ output_scale_factor=output_scale_factor,
+ )
+ )
+ else:
+ resnets.append(
+ ResnetBlock2D(
+ in_channels=in_channels,
+ out_channels=out_channels,
+ temb_channels=None,
+ eps=resnet_eps,
+ groups=resnet_groups,
+ dropout=dropout,
+ time_embedding_norm=resnet_time_scale_shift,
+ non_linearity=resnet_act_fn,
+ output_scale_factor=output_scale_factor,
+ pre_norm=resnet_pre_norm,
+ )
+ )
+
+ self.resnets = nn.ModuleList(resnets)
+
+ if add_downsample:
+ self.downsamplers = nn.ModuleList(
+ [
+ Downsample2D(
+ out_channels, use_conv=True, out_channels=out_channels, padding=downsample_padding, name="op"
+ )
+ ]
+ )
+ else:
+ self.downsamplers = None
+
+ def forward(self, hidden_states: torch.Tensor, *args, **kwargs) -> torch.Tensor:
+ if len(args) > 0 or kwargs.get("scale", None) is not None:
+ deprecation_message = "The `scale` argument is deprecated and will be ignored. Please remove it, as passing it will raise an error in the future. `scale` should directly be passed while calling the underlying pipeline component i.e., via `cross_attention_kwargs`."
+ deprecate("scale", "1.0.0", deprecation_message)
+
+ for resnet in self.resnets:
+ hidden_states = resnet(hidden_states, temb=None)
+
+ if self.downsamplers is not None:
+ for downsampler in self.downsamplers:
+ hidden_states = downsampler(hidden_states)
+
+ return hidden_states
+
+
+class UpDecoderBlock2D(nn.Module):
+ def __init__(
+ self,
+ in_channels: int,
+ out_channels: int,
+ resolution_idx: Optional[int] = None,
+ dropout: float = 0.0,
+ num_layers: int = 1,
+ resnet_eps: float = 1e-6,
+ resnet_time_scale_shift: str = "default", # default, spatial
+ resnet_act_fn: str = "swish",
+ resnet_groups: int = 32,
+ resnet_pre_norm: bool = True,
+ output_scale_factor: float = 1.0,
+ add_upsample: bool = True,
+ temb_channels: Optional[int] = None,
+ ):
+ super().__init__()
+ resnets = []
+
+ for i in range(num_layers):
+ input_channels = in_channels if i == 0 else out_channels
+
+ if resnet_time_scale_shift == "spatial":
+ resnets.append(
+ ResnetBlockCondNorm2D(
+ in_channels=input_channels,
+ out_channels=out_channels,
+ temb_channels=temb_channels,
+ eps=resnet_eps,
+ groups=resnet_groups,
+ dropout=dropout,
+ time_embedding_norm="spatial",
+ non_linearity=resnet_act_fn,
+ output_scale_factor=output_scale_factor,
+ )
+ )
+ else:
+ resnets.append(
+ ResnetBlock2D(
+ in_channels=input_channels,
+ out_channels=out_channels,
+ temb_channels=temb_channels,
+ eps=resnet_eps,
+ groups=resnet_groups,
+ dropout=dropout,
+ time_embedding_norm=resnet_time_scale_shift,
+ non_linearity=resnet_act_fn,
+ output_scale_factor=output_scale_factor,
+ pre_norm=resnet_pre_norm,
+ )
+ )
+
+ self.resnets = nn.ModuleList(resnets)
+
+ if add_upsample:
+ self.upsamplers = nn.ModuleList([Upsample2D(out_channels, use_conv=True, out_channels=out_channels)])
+ else:
+ self.upsamplers = None
+
+ self.resolution_idx = resolution_idx
+
+ def forward(self, hidden_states: torch.Tensor, temb: Optional[torch.Tensor] = None) -> torch.Tensor:
+ for resnet in self.resnets:
+ hidden_states = resnet(hidden_states, temb=temb)
+
+ if self.upsamplers is not None:
+ for upsampler in self.upsamplers:
+ hidden_states = upsampler(hidden_states)
+
+ return hidden_states
+
+class Encoder(nn.Module):
+ r"""
+ The `Encoder` layer of a variational autoencoder that encodes its input into a latent representation.
+
+ Args:
+ in_channels (`int`, *optional*, defaults to 3):
+ The number of input channels.
+ out_channels (`int`, *optional*, defaults to 3):
+ The number of output channels.
+ down_block_types (`Tuple[str, ...]`, *optional*, defaults to `("DownEncoderBlock2D",)`):
+ The types of down blocks to use. See `~diffusers.models.unet_2d_blocks.get_down_block` for available
+ options.
+ block_out_channels (`Tuple[int, ...]`, *optional*, defaults to `(64,)`):
+ The number of output channels for each block.
+ layers_per_block (`int`, *optional*, defaults to 2):
+ The number of layers per block.
+ norm_num_groups (`int`, *optional*, defaults to 32):
+ The number of groups for normalization.
+ act_fn (`str`, *optional*, defaults to `"silu"`):
+ The activation function to use. See `~diffusers.models.activations.get_activation` for available options.
+ double_z (`bool`, *optional*, defaults to `True`):
+ Whether to double the number of output channels for the last block.
+ """
+
+ def __init__(
+ self,
+ in_channels: int = 3,
+ out_channels: int = 3,
+ down_block_types: Tuple[str, ...] = ("DownEncoderBlock2D",),
+ block_out_channels: Tuple[int, ...] = (64,),
+ layers_per_block: int = 2,
+ norm_num_groups: int = 32,
+ act_fn: str = "silu",
+ double_z: bool = True,
+ mid_block_add_attention=True,
+ ):
+ super().__init__()
+ self.layers_per_block = layers_per_block
+
+ self.conv_in = nn.Conv2d(
+ in_channels,
+ block_out_channels[0],
+ kernel_size=3,
+ stride=1,
+ padding=1,
+ )
+
+ self.down_blocks = nn.ModuleList([])
+
+ # down
+ output_channel = block_out_channels[0]
+ for i, down_block_type in enumerate(down_block_types):
+ input_channel = output_channel
+ output_channel = block_out_channels[i]
+ is_final_block = i == len(block_out_channels) - 1
+
+ down_block = DownEncoderBlock2D(
+ num_layers=self.layers_per_block,
+ in_channels=input_channel,
+ out_channels=output_channel,
+ add_downsample=not is_final_block,
+ resnet_eps=1e-6,
+ downsample_padding=0,
+ resnet_act_fn=act_fn,
+ resnet_groups=norm_num_groups,
+ # attention_head_dim=output_channel,
+ # temb_channels=None,
+ )
+ self.down_blocks.append(down_block)
+
+ # mid
+ self.mid_block = UNetMidBlock2D(
+ in_channels=block_out_channels[-1],
+ resnet_eps=1e-6,
+ resnet_act_fn=act_fn,
+ output_scale_factor=1,
+ resnet_time_scale_shift="default",
+ attention_head_dim=block_out_channels[-1],
+ resnet_groups=norm_num_groups,
+ temb_channels=None,
+ add_attention=mid_block_add_attention,
+ )
+
+ # out
+ self.conv_norm_out = nn.GroupNorm(num_channels=block_out_channels[-1], num_groups=norm_num_groups, eps=1e-6)
+ self.conv_act = nn.SiLU()
+
+ conv_out_channels = 2 * out_channels if double_z else out_channels
+ self.conv_out = nn.Conv2d(block_out_channels[-1], conv_out_channels, 3, padding=1)
+
+ self.gradient_checkpointing = False
+
+ def forward(self, sample: torch.Tensor) -> torch.Tensor:
+ r"""The forward method of the `Encoder` class."""
+
+ sample = self.conv_in(sample)
+
+ if torch.is_grad_enabled() and self.gradient_checkpointing:
+ # down
+ for down_block in self.down_blocks:
+ sample = self._gradient_checkpointing_func(down_block, sample)
+ # middle
+ sample = self._gradient_checkpointing_func(self.mid_block, sample)
+
+ else:
+ # down
+ for down_block in self.down_blocks:
+ sample = down_block(sample)
+
+ # middle
+ sample = self.mid_block(sample)
+
+ # post-process
+ sample = self.conv_norm_out(sample)
+ sample = self.conv_act(sample)
+ sample = self.conv_out(sample)
+
+ return sample
+
+class Decoder(nn.Module):
+ r"""
+ The `Decoder` layer of a variational autoencoder that decodes its latent representation into an output sample.
+
+ Args:
+ in_channels (`int`, *optional*, defaults to 3):
+ The number of input channels.
+ out_channels (`int`, *optional*, defaults to 3):
+ The number of output channels.
+ up_block_types (`Tuple[str, ...]`, *optional*, defaults to `("UpDecoderBlock2D",)`):
+ The types of up blocks to use. See `~diffusers.models.unet_2d_blocks.get_up_block` for available options.
+ block_out_channels (`Tuple[int, ...]`, *optional*, defaults to `(64,)`):
+ The number of output channels for each block.
+ layers_per_block (`int`, *optional*, defaults to 2):
+ The number of layers per block.
+ norm_num_groups (`int`, *optional*, defaults to 32):
+ The number of groups for normalization.
+ act_fn (`str`, *optional*, defaults to `"silu"`):
+ The activation function to use. See `~diffusers.models.activations.get_activation` for available options.
+ norm_type (`str`, *optional*, defaults to `"group"`):
+ The normalization type to use. Can be either `"group"` or `"spatial"`.
+ """
+
+ def __init__(
+ self,
+ in_channels: int = 3,
+ out_channels: int = 3,
+ up_block_types: Tuple[str, ...] = ("UpDecoderBlock2D",),
+ block_out_channels: Tuple[int, ...] = (64,),
+ layers_per_block: int = 2,
+ norm_num_groups: int = 32,
+ act_fn: str = "silu",
+ norm_type: str = "group", # group, spatial
+ mid_block_add_attention=True,
+ ):
+ super().__init__()
+ self.layers_per_block = layers_per_block
+
+ self.conv_in = nn.Conv2d(
+ in_channels,
+ block_out_channels[-1],
+ kernel_size=3,
+ stride=1,
+ padding=1,
+ )
+
+ self.up_blocks = nn.ModuleList([])
+
+ temb_channels = in_channels if norm_type == "spatial" else None
+
+ # mid
+ self.mid_block = UNetMidBlock2D(
+ in_channels=block_out_channels[-1],
+ resnet_eps=1e-6,
+ resnet_act_fn=act_fn,
+ output_scale_factor=1,
+ resnet_time_scale_shift="default" if norm_type == "group" else norm_type,
+ attention_head_dim=block_out_channels[-1],
+ resnet_groups=norm_num_groups,
+ temb_channels=temb_channels,
+ add_attention=mid_block_add_attention,
+ )
+
+ # up
+ reversed_block_out_channels = list(reversed(block_out_channels))
+ output_channel = reversed_block_out_channels[0]
+ for i, up_block_type in enumerate(up_block_types):
+ prev_output_channel = output_channel
+ output_channel = reversed_block_out_channels[i]
+
+ is_final_block = i == len(block_out_channels) - 1
+
+ up_block = UpDecoderBlock2D(
+ num_layers=self.layers_per_block + 1,
+ in_channels=prev_output_channel,
+ out_channels=output_channel,
+ # prev_output_channel=prev_output_channel,
+ add_upsample=not is_final_block,
+ resnet_eps=1e-6,
+ resnet_act_fn=act_fn,
+ resnet_groups=norm_num_groups,
+ # attention_head_dim=output_channel,
+ temb_channels=temb_channels,
+ resnet_time_scale_shift=norm_type,
+ )
+ self.up_blocks.append(up_block)
+ prev_output_channel = output_channel
+
+ # out
+ if norm_type == "spatial":
+ self.conv_norm_out = SpatialNorm(block_out_channels[0], temb_channels)
+ else:
+ self.conv_norm_out = nn.GroupNorm(num_channels=block_out_channels[0], num_groups=norm_num_groups, eps=1e-6)
+ self.conv_act = nn.SiLU()
+ self.conv_out = nn.Conv2d(block_out_channels[0], out_channels, 3, padding=1)
+
+ self.gradient_checkpointing = False
+
+ def forward(
+ self,
+ sample: torch.Tensor,
+ latent_embeds: Optional[torch.Tensor] = None,
+ ) -> torch.Tensor:
+ r"""The forward method of the `Decoder` class."""
+
+ sample = self.conv_in(sample)
+
+ if torch.is_grad_enabled() and self.gradient_checkpointing:
+ # middle
+ sample = self._gradient_checkpointing_func(self.mid_block, sample, latent_embeds)
+
+ # up
+ for up_block in self.up_blocks:
+ sample = self._gradient_checkpointing_func(up_block, sample, latent_embeds)
+ else:
+ # middle
+ sample = self.mid_block(sample, latent_embeds)
+
+ # up
+ for up_block in self.up_blocks:
+ sample = up_block(sample, latent_embeds)
+
+ # post-process
+ if latent_embeds is None:
+ sample = self.conv_norm_out(sample)
+ else:
+ sample = self.conv_norm_out(sample, latent_embeds)
+ sample = self.conv_act(sample)
+ sample = self.conv_out(sample)
+
+ return sample
+
+
+class Flux2VAE(torch.nn.Module):
+ r"""
+ A VAE model with KL loss for encoding images into latents and decoding latent representations into images.
+
+ This model inherits from [`ModelMixin`]. Check the superclass documentation for it's generic methods implemented
+ for all models (such as downloading or saving).
+
+ Parameters:
+ in_channels (int, *optional*, defaults to 3): Number of channels in the input image.
+ out_channels (int, *optional*, defaults to 3): Number of channels in the output.
+ down_block_types (`Tuple[str]`, *optional*, defaults to `("DownEncoderBlock2D",)`):
+ Tuple of downsample block types.
+ up_block_types (`Tuple[str]`, *optional*, defaults to `("UpDecoderBlock2D",)`):
+ Tuple of upsample block types.
+ block_out_channels (`Tuple[int]`, *optional*, defaults to `(64,)`):
+ Tuple of block output channels.
+ act_fn (`str`, *optional*, defaults to `"silu"`): The activation function to use.
+ latent_channels (`int`, *optional*, defaults to 4): Number of channels in the latent space.
+ sample_size (`int`, *optional*, defaults to `32`): Sample input size.
+ force_upcast (`bool`, *optional*, default to `True`):
+ If enabled it will force the VAE to run in float32 for high image resolution pipelines, such as SD-XL. VAE
+ can be fine-tuned / trained to a lower range without losing too much precision in which case `force_upcast`
+ can be set to `False` - see: https://huggingface.co/madebyollin/sdxl-vae-fp16-fix
+ mid_block_add_attention (`bool`, *optional*, default to `True`):
+ If enabled, the mid_block of the Encoder and Decoder will have attention blocks. If set to false, the
+ mid_block will only have resnet blocks
+ """
+
+ _supports_gradient_checkpointing = True
+ _no_split_modules = ["BasicTransformerBlock", "ResnetBlock2D"]
+
+ def __init__(
+ self,
+ in_channels: int = 3,
+ out_channels: int = 3,
+ down_block_types: Tuple[str, ...] = (
+ "DownEncoderBlock2D",
+ "DownEncoderBlock2D",
+ "DownEncoderBlock2D",
+ "DownEncoderBlock2D",
+ ),
+ up_block_types: Tuple[str, ...] = (
+ "UpDecoderBlock2D",
+ "UpDecoderBlock2D",
+ "UpDecoderBlock2D",
+ "UpDecoderBlock2D",
+ ),
+ block_out_channels: Tuple[int, ...] = (
+ 128,
+ 256,
+ 512,
+ 512,
+ ),
+ layers_per_block: int = 2,
+ act_fn: str = "silu",
+ latent_channels: int = 32,
+ norm_num_groups: int = 32,
+ sample_size: int = 1024, # YiYi notes: not sure
+ force_upcast: bool = True,
+ use_quant_conv: bool = True,
+ use_post_quant_conv: bool = True,
+ mid_block_add_attention: bool = True,
+ batch_norm_eps: float = 1e-4,
+ batch_norm_momentum: float = 0.1,
+ patch_size: Tuple[int, int] = (2, 2),
+ ):
+ super().__init__()
+
+ # pass init params to Encoder
+ self.encoder = Encoder(
+ in_channels=in_channels,
+ out_channels=latent_channels,
+ down_block_types=down_block_types,
+ block_out_channels=block_out_channels,
+ layers_per_block=layers_per_block,
+ act_fn=act_fn,
+ norm_num_groups=norm_num_groups,
+ double_z=True,
+ mid_block_add_attention=mid_block_add_attention,
+ )
+
+ # pass init params to Decoder
+ self.decoder = Decoder(
+ in_channels=latent_channels,
+ out_channels=out_channels,
+ up_block_types=up_block_types,
+ block_out_channels=block_out_channels,
+ layers_per_block=layers_per_block,
+ norm_num_groups=norm_num_groups,
+ act_fn=act_fn,
+ mid_block_add_attention=mid_block_add_attention,
+ )
+
+ self.quant_conv = nn.Conv2d(2 * latent_channels, 2 * latent_channels, 1) if use_quant_conv else None
+ self.post_quant_conv = nn.Conv2d(latent_channels, latent_channels, 1) if use_post_quant_conv else None
+
+ self.bn = nn.BatchNorm2d(
+ math.prod(patch_size) * latent_channels,
+ eps=batch_norm_eps,
+ momentum=batch_norm_momentum,
+ affine=False,
+ track_running_stats=True,
+ )
+
+ self.use_slicing = False
+ self.use_tiling = False
+
+ @property
+ # Copied from diffusers.models.unets.unet_2d_condition.UNet2DConditionModel.attn_processors
+ def attn_processors(self):
+ r"""
+ Returns:
+ `dict` of attention processors: A dictionary containing all attention processors used in the model with
+ indexed by its weight name.
+ """
+ # set recursively
+ processors = {}
+
+ def fn_recursive_add_processors(name: str, module: torch.nn.Module, processors):
+ if hasattr(module, "get_processor"):
+ processors[f"{name}.processor"] = module.get_processor()
+
+ for sub_name, child in module.named_children():
+ fn_recursive_add_processors(f"{name}.{sub_name}", child, processors)
+
+ return processors
+
+ for name, module in self.named_children():
+ fn_recursive_add_processors(name, module, processors)
+
+ return processors
+
+ # Copied from diffusers.models.unets.unet_2d_condition.UNet2DConditionModel.set_attn_processor
+ def set_attn_processor(self, processor):
+ r"""
+ Sets the attention processor to use to compute attention.
+
+ Parameters:
+ processor (`dict` of `AttentionProcessor` or only `AttentionProcessor`):
+ The instantiated processor class or a dictionary of processor classes that will be set as the processor
+ for **all** `Attention` layers.
+
+ If `processor` is a dict, the key needs to define the path to the corresponding cross attention
+ processor. This is strongly recommended when setting trainable attention processors.
+
+ """
+ count = len(self.attn_processors.keys())
+
+ if isinstance(processor, dict) and len(processor) != count:
+ raise ValueError(
+ f"A dict of processors was passed, but the number of processors {len(processor)} does not match the"
+ f" number of attention layers: {count}. Please make sure to pass {count} processor classes."
+ )
+
+ def fn_recursive_attn_processor(name: str, module: torch.nn.Module, processor):
+ if hasattr(module, "set_processor"):
+ if not isinstance(processor, dict):
+ module.set_processor(processor)
+ else:
+ module.set_processor(processor.pop(f"{name}.processor"))
+
+ for sub_name, child in module.named_children():
+ fn_recursive_attn_processor(f"{name}.{sub_name}", child, processor)
+
+ for name, module in self.named_children():
+ fn_recursive_attn_processor(name, module, processor)
+
+ def _encode(self, x: torch.Tensor) -> torch.Tensor:
+ batch_size, num_channels, height, width = x.shape
+
+ if self.use_tiling and (width > self.tile_sample_min_size or height > self.tile_sample_min_size):
+ return self._tiled_encode(x)
+
+ enc = self.encoder(x)
+ if self.quant_conv is not None:
+ enc = self.quant_conv(enc)
+
+ return enc
+
+ def encode(
+ self, x: torch.Tensor, return_dict: bool = True
+ ):
+ """
+ Encode a batch of images into latents.
+
+ Args:
+ x (`torch.Tensor`): Input batch of images.
+ return_dict (`bool`, *optional*, defaults to `True`):
+ Whether to return a [`~models.autoencoder_kl.AutoencoderKLOutput`] instead of a plain tuple.
+
+ Returns:
+ The latent representations of the encoded images. If `return_dict` is True, a
+ [`~models.autoencoder_kl.AutoencoderKLOutput`] is returned, otherwise a plain `tuple` is returned.
+ """
+ if self.use_slicing and x.shape[0] > 1:
+ encoded_slices = [self._encode(x_slice) for x_slice in x.split(1)]
+ h = torch.cat(encoded_slices)
+ else:
+ h = self._encode(x)
+
+
+ h = rearrange(h, "B C (H P) (W Q) -> B (C P Q) H W", P=2, Q=2)
+ h = h[:, :128]
+ latents_bn_mean = self.bn.running_mean.view(1, -1, 1, 1).to(h.device, h.dtype)
+ latents_bn_std = torch.sqrt(self.bn.running_var.view(1, -1, 1, 1) + 0.0001).to(
+ h.device, h.dtype
+ )
+ h = (h - latents_bn_mean) / latents_bn_std
+ return h
+
+ def _decode(self, z: torch.Tensor, return_dict: bool = True):
+ if self.use_tiling and (z.shape[-1] > self.tile_latent_min_size or z.shape[-2] > self.tile_latent_min_size):
+ return self.tiled_decode(z, return_dict=return_dict)
+
+ if self.post_quant_conv is not None:
+ z = self.post_quant_conv(z)
+
+ dec = self.decoder(z)
+
+ if not return_dict:
+ return (dec,)
+
+ return dec
+
+ def decode(
+ self, z: torch.FloatTensor, return_dict: bool = True, generator=None
+ ):
+ latents_bn_mean = self.bn.running_mean.view(1, -1, 1, 1).to(z.device, z.dtype)
+ latents_bn_std = torch.sqrt(self.bn.running_var.view(1, -1, 1, 1) + 0.0001).to(
+ z.device, z.dtype
+ )
+ z = z * latents_bn_std + latents_bn_mean
+ z = rearrange(z, "B (C P Q) H W -> B C (H P) (W Q)", P=2, Q=2)
+ """
+ Decode a batch of images.
+
+ Args:
+ z (`torch.Tensor`): Input batch of latent vectors.
+ return_dict (`bool`, *optional*, defaults to `True`):
+ Whether to return a [`~models.vae.DecoderOutput`] instead of a plain tuple.
+
+ Returns:
+ [`~models.vae.DecoderOutput`] or `tuple`:
+ If return_dict is True, a [`~models.vae.DecoderOutput`] is returned, otherwise a plain `tuple` is
+ returned.
+
+ """
+ if self.use_slicing and z.shape[0] > 1:
+ decoded_slices = [self._decode(z_slice) for z_slice in z.split(1)]
+ decoded = torch.cat(decoded_slices)
+ else:
+ decoded = self._decode(z)
+
+ if not return_dict:
+ return (decoded,)
+
+ return decoded
+
+ def blend_v(self, a: torch.Tensor, b: torch.Tensor, blend_extent: int) -> torch.Tensor:
+ blend_extent = min(a.shape[2], b.shape[2], blend_extent)
+ for y in range(blend_extent):
+ b[:, :, y, :] = a[:, :, -blend_extent + y, :] * (1 - y / blend_extent) + b[:, :, y, :] * (y / blend_extent)
+ return b
+
+ def blend_h(self, a: torch.Tensor, b: torch.Tensor, blend_extent: int) -> torch.Tensor:
+ blend_extent = min(a.shape[3], b.shape[3], blend_extent)
+ for x in range(blend_extent):
+ b[:, :, :, x] = a[:, :, :, -blend_extent + x] * (1 - x / blend_extent) + b[:, :, :, x] * (x / blend_extent)
+ return b
+
+ def _tiled_encode(self, x: torch.Tensor) -> torch.Tensor:
+ r"""Encode a batch of images using a tiled encoder.
+
+ When this option is enabled, the VAE will split the input tensor into tiles to compute encoding in several
+ steps. This is useful to keep memory use constant regardless of image size. The end result of tiled encoding is
+ different from non-tiled encoding because each tile uses a different encoder. To avoid tiling artifacts, the
+ tiles overlap and are blended together to form a smooth output. You may still see tile-sized changes in the
+ output, but they should be much less noticeable.
+
+ Args:
+ x (`torch.Tensor`): Input batch of images.
+
+ Returns:
+ `torch.Tensor`:
+ The latent representation of the encoded videos.
+ """
+
+ overlap_size = int(self.tile_sample_min_size * (1 - self.tile_overlap_factor))
+ blend_extent = int(self.tile_latent_min_size * self.tile_overlap_factor)
+ row_limit = self.tile_latent_min_size - blend_extent
+
+ # Split the image into 512x512 tiles and encode them separately.
+ rows = []
+ for i in range(0, x.shape[2], overlap_size):
+ row = []
+ for j in range(0, x.shape[3], overlap_size):
+ tile = x[:, :, i : i + self.tile_sample_min_size, j : j + self.tile_sample_min_size]
+ tile = self.encoder(tile)
+ if self.config.use_quant_conv:
+ tile = self.quant_conv(tile)
+ row.append(tile)
+ rows.append(row)
+ result_rows = []
+ for i, row in enumerate(rows):
+ result_row = []
+ for j, tile in enumerate(row):
+ # blend the above tile and the left tile
+ # to the current tile and add the current tile to the result row
+ if i > 0:
+ tile = self.blend_v(rows[i - 1][j], tile, blend_extent)
+ if j > 0:
+ tile = self.blend_h(row[j - 1], tile, blend_extent)
+ result_row.append(tile[:, :, :row_limit, :row_limit])
+ result_rows.append(torch.cat(result_row, dim=3))
+
+ enc = torch.cat(result_rows, dim=2)
+ return enc
+
+ def tiled_encode(self, x: torch.Tensor, return_dict: bool = True):
+ r"""Encode a batch of images using a tiled encoder.
+
+ When this option is enabled, the VAE will split the input tensor into tiles to compute encoding in several
+ steps. This is useful to keep memory use constant regardless of image size. The end result of tiled encoding is
+ different from non-tiled encoding because each tile uses a different encoder. To avoid tiling artifacts, the
+ tiles overlap and are blended together to form a smooth output. You may still see tile-sized changes in the
+ output, but they should be much less noticeable.
+
+ Args:
+ x (`torch.Tensor`): Input batch of images.
+ return_dict (`bool`, *optional*, defaults to `True`):
+ Whether or not to return a [`~models.autoencoder_kl.AutoencoderKLOutput`] instead of a plain tuple.
+
+ Returns:
+ [`~models.autoencoder_kl.AutoencoderKLOutput`] or `tuple`:
+ If return_dict is True, a [`~models.autoencoder_kl.AutoencoderKLOutput`] is returned, otherwise a plain
+ `tuple` is returned.
+ """
+
+ overlap_size = int(self.tile_sample_min_size * (1 - self.tile_overlap_factor))
+ blend_extent = int(self.tile_latent_min_size * self.tile_overlap_factor)
+ row_limit = self.tile_latent_min_size - blend_extent
+
+ # Split the image into 512x512 tiles and encode them separately.
+ rows = []
+ for i in range(0, x.shape[2], overlap_size):
+ row = []
+ for j in range(0, x.shape[3], overlap_size):
+ tile = x[:, :, i : i + self.tile_sample_min_size, j : j + self.tile_sample_min_size]
+ tile = self.encoder(tile)
+ if self.config.use_quant_conv:
+ tile = self.quant_conv(tile)
+ row.append(tile)
+ rows.append(row)
+ result_rows = []
+ for i, row in enumerate(rows):
+ result_row = []
+ for j, tile in enumerate(row):
+ # blend the above tile and the left tile
+ # to the current tile and add the current tile to the result row
+ if i > 0:
+ tile = self.blend_v(rows[i - 1][j], tile, blend_extent)
+ if j > 0:
+ tile = self.blend_h(row[j - 1], tile, blend_extent)
+ result_row.append(tile[:, :, :row_limit, :row_limit])
+ result_rows.append(torch.cat(result_row, dim=3))
+
+ moments = torch.cat(result_rows, dim=2)
+ return moments
+
+ def tiled_decode(self, z: torch.Tensor, return_dict: bool = True):
+ r"""
+ Decode a batch of images using a tiled decoder.
+
+ Args:
+ z (`torch.Tensor`): Input batch of latent vectors.
+ return_dict (`bool`, *optional*, defaults to `True`):
+ Whether or not to return a [`~models.vae.DecoderOutput`] instead of a plain tuple.
+
+ Returns:
+ [`~models.vae.DecoderOutput`] or `tuple`:
+ If return_dict is True, a [`~models.vae.DecoderOutput`] is returned, otherwise a plain `tuple` is
+ returned.
+ """
+ overlap_size = int(self.tile_latent_min_size * (1 - self.tile_overlap_factor))
+ blend_extent = int(self.tile_sample_min_size * self.tile_overlap_factor)
+ row_limit = self.tile_sample_min_size - blend_extent
+
+ # Split z into overlapping 64x64 tiles and decode them separately.
+ # The tiles have an overlap to avoid seams between tiles.
+ rows = []
+ for i in range(0, z.shape[2], overlap_size):
+ row = []
+ for j in range(0, z.shape[3], overlap_size):
+ tile = z[:, :, i : i + self.tile_latent_min_size, j : j + self.tile_latent_min_size]
+ if self.config.use_post_quant_conv:
+ tile = self.post_quant_conv(tile)
+ decoded = self.decoder(tile)
+ row.append(decoded)
+ rows.append(row)
+ result_rows = []
+ for i, row in enumerate(rows):
+ result_row = []
+ for j, tile in enumerate(row):
+ # blend the above tile and the left tile
+ # to the current tile and add the current tile to the result row
+ if i > 0:
+ tile = self.blend_v(rows[i - 1][j], tile, blend_extent)
+ if j > 0:
+ tile = self.blend_h(row[j - 1], tile, blend_extent)
+ result_row.append(tile[:, :, :row_limit, :row_limit])
+ result_rows.append(torch.cat(result_row, dim=3))
+
+ dec = torch.cat(result_rows, dim=2)
+ if not return_dict:
+ return (dec,)
+
+ return dec
+
+ def forward(
+ self,
+ sample: torch.Tensor,
+ sample_posterior: bool = False,
+ return_dict: bool = True,
+ generator: Optional[torch.Generator] = None,
+ ):
+ r"""
+ Args:
+ sample (`torch.Tensor`): Input sample.
+ sample_posterior (`bool`, *optional*, defaults to `False`):
+ Whether to sample from the posterior.
+ return_dict (`bool`, *optional*, defaults to `True`):
+ Whether or not to return a [`DecoderOutput`] instead of a plain tuple.
+ """
+ x = sample
+ posterior = self.encode(x).latent_dist
+ if sample_posterior:
+ z = posterior.sample(generator=generator)
+ else:
+ z = posterior.mode()
+ dec = self.decode(z).sample
+
+ if not return_dict:
+ return (dec,)
+
+ return dec
diff --git a/diffsynth/models/flux_controlnet.py b/diffsynth/models/flux_controlnet.py
new file mode 100644
index 0000000000000000000000000000000000000000..7fb1138bb74f7b55c6e92ac312098e4168829f0a
--- /dev/null
+++ b/diffsynth/models/flux_controlnet.py
@@ -0,0 +1,384 @@
+import torch
+from einops import rearrange, repeat
+from .flux_dit import RoPEEmbedding, TimestepEmbeddings, FluxJointTransformerBlock, FluxSingleTransformerBlock, RMSNorm
+# from .utils import hash_state_dict_keys, init_weights_on_device
+from contextlib import contextmanager
+
+def hash_state_dict_keys(state_dict, with_shape=True):
+ keys_str = convert_state_dict_keys_to_single_str(state_dict, with_shape=with_shape)
+ keys_str = keys_str.encode(encoding="UTF-8")
+ return hashlib.md5(keys_str).hexdigest()
+
+@contextmanager
+def init_weights_on_device(device = torch.device("meta"), include_buffers :bool = False):
+
+ old_register_parameter = torch.nn.Module.register_parameter
+ if include_buffers:
+ old_register_buffer = torch.nn.Module.register_buffer
+
+ def register_empty_parameter(module, name, param):
+ old_register_parameter(module, name, param)
+ if param is not None:
+ param_cls = type(module._parameters[name])
+ kwargs = module._parameters[name].__dict__
+ kwargs["requires_grad"] = param.requires_grad
+ module._parameters[name] = param_cls(module._parameters[name].to(device), **kwargs)
+
+ def register_empty_buffer(module, name, buffer, persistent=True):
+ old_register_buffer(module, name, buffer, persistent=persistent)
+ if buffer is not None:
+ module._buffers[name] = module._buffers[name].to(device)
+
+ def patch_tensor_constructor(fn):
+ def wrapper(*args, **kwargs):
+ kwargs["device"] = device
+ return fn(*args, **kwargs)
+
+ return wrapper
+
+ if include_buffers:
+ tensor_constructors_to_patch = {
+ torch_function_name: getattr(torch, torch_function_name)
+ for torch_function_name in ["empty", "zeros", "ones", "full"]
+ }
+ else:
+ tensor_constructors_to_patch = {}
+
+ try:
+ torch.nn.Module.register_parameter = register_empty_parameter
+ if include_buffers:
+ torch.nn.Module.register_buffer = register_empty_buffer
+ for torch_function_name in tensor_constructors_to_patch.keys():
+ setattr(torch, torch_function_name, patch_tensor_constructor(getattr(torch, torch_function_name)))
+ yield
+ finally:
+ torch.nn.Module.register_parameter = old_register_parameter
+ if include_buffers:
+ torch.nn.Module.register_buffer = old_register_buffer
+ for torch_function_name, old_torch_function in tensor_constructors_to_patch.items():
+ setattr(torch, torch_function_name, old_torch_function)
+
+class FluxControlNet(torch.nn.Module):
+ def __init__(self, disable_guidance_embedder=False, num_joint_blocks=5, num_single_blocks=10, num_mode=0, mode_dict={}, additional_input_dim=0):
+ super().__init__()
+ self.pos_embedder = RoPEEmbedding(3072, 10000, [16, 56, 56])
+ self.time_embedder = TimestepEmbeddings(256, 3072)
+ self.guidance_embedder = None if disable_guidance_embedder else TimestepEmbeddings(256, 3072)
+ self.pooled_text_embedder = torch.nn.Sequential(torch.nn.Linear(768, 3072), torch.nn.SiLU(), torch.nn.Linear(3072, 3072))
+ self.context_embedder = torch.nn.Linear(4096, 3072)
+ self.x_embedder = torch.nn.Linear(64, 3072)
+
+ self.blocks = torch.nn.ModuleList([FluxJointTransformerBlock(3072, 24) for _ in range(num_joint_blocks)])
+ self.single_blocks = torch.nn.ModuleList([FluxSingleTransformerBlock(3072, 24) for _ in range(num_single_blocks)])
+
+ self.controlnet_blocks = torch.nn.ModuleList([torch.nn.Linear(3072, 3072) for _ in range(num_joint_blocks)])
+ self.controlnet_single_blocks = torch.nn.ModuleList([torch.nn.Linear(3072, 3072) for _ in range(num_single_blocks)])
+
+ self.mode_dict = mode_dict
+ self.controlnet_mode_embedder = torch.nn.Embedding(num_mode, 3072) if len(mode_dict) > 0 else None
+ self.controlnet_x_embedder = torch.nn.Linear(64 + additional_input_dim, 3072)
+
+
+ def prepare_image_ids(self, latents):
+ batch_size, _, height, width = latents.shape
+ latent_image_ids = torch.zeros(height // 2, width // 2, 3)
+ latent_image_ids[..., 1] = latent_image_ids[..., 1] + torch.arange(height // 2)[:, None]
+ latent_image_ids[..., 2] = latent_image_ids[..., 2] + torch.arange(width // 2)[None, :]
+
+ latent_image_id_height, latent_image_id_width, latent_image_id_channels = latent_image_ids.shape
+
+ latent_image_ids = latent_image_ids[None, :].repeat(batch_size, 1, 1, 1)
+ latent_image_ids = latent_image_ids.reshape(
+ batch_size, latent_image_id_height * latent_image_id_width, latent_image_id_channels
+ )
+ latent_image_ids = latent_image_ids.to(device=latents.device, dtype=latents.dtype)
+
+ return latent_image_ids
+
+
+ def patchify(self, hidden_states):
+ hidden_states = rearrange(hidden_states, "B C (H P) (W Q) -> B (H W) (C P Q)", P=2, Q=2)
+ return hidden_states
+
+
+ def align_res_stack_to_original_blocks(self, res_stack, num_blocks, hidden_states):
+ if len(res_stack) == 0:
+ return [torch.zeros_like(hidden_states)] * num_blocks
+ interval = (num_blocks + len(res_stack) - 1) // len(res_stack)
+ aligned_res_stack = [res_stack[block_id // interval] for block_id in range(num_blocks)]
+ return aligned_res_stack
+
+
+ def forward(
+ self,
+ hidden_states,
+ controlnet_conditioning,
+ timestep, prompt_emb, pooled_prompt_emb, guidance, text_ids, image_ids=None,
+ processor_id=None,
+ tiled=False, tile_size=128, tile_stride=64,
+ **kwargs
+ ):
+ if image_ids is None:
+ image_ids = self.prepare_image_ids(hidden_states)
+
+ conditioning = self.time_embedder(timestep, hidden_states.dtype) + self.pooled_text_embedder(pooled_prompt_emb)
+ if self.guidance_embedder is not None:
+ guidance = guidance * 1000
+ conditioning = conditioning + self.guidance_embedder(guidance, hidden_states.dtype)
+ prompt_emb = self.context_embedder(prompt_emb)
+ if self.controlnet_mode_embedder is not None: # Different from FluxDiT
+ processor_id = torch.tensor([self.mode_dict[processor_id]], dtype=torch.int)
+ processor_id = repeat(processor_id, "D -> B D", B=1).to(text_ids.device)
+ prompt_emb = torch.concat([self.controlnet_mode_embedder(processor_id), prompt_emb], dim=1)
+ text_ids = torch.cat([text_ids[:, :1], text_ids], dim=1)
+ image_rotary_emb = self.pos_embedder(torch.cat((text_ids, image_ids), dim=1))
+
+ hidden_states = self.patchify(hidden_states)
+ hidden_states = self.x_embedder(hidden_states)
+ controlnet_conditioning = self.patchify(controlnet_conditioning) # Different from FluxDiT
+ hidden_states = hidden_states + self.controlnet_x_embedder(controlnet_conditioning) # Different from FluxDiT
+
+ controlnet_res_stack = []
+ for block, controlnet_block in zip(self.blocks, self.controlnet_blocks):
+ hidden_states, prompt_emb = block(hidden_states, prompt_emb, conditioning, image_rotary_emb)
+ controlnet_res_stack.append(controlnet_block(hidden_states))
+
+ controlnet_single_res_stack = []
+ hidden_states = torch.cat([prompt_emb, hidden_states], dim=1)
+ for block, controlnet_block in zip(self.single_blocks, self.controlnet_single_blocks):
+ hidden_states, prompt_emb = block(hidden_states, prompt_emb, conditioning, image_rotary_emb)
+ controlnet_single_res_stack.append(controlnet_block(hidden_states[:, prompt_emb.shape[1]:]))
+
+ controlnet_res_stack = self.align_res_stack_to_original_blocks(controlnet_res_stack, 19, hidden_states[:, prompt_emb.shape[1]:])
+ controlnet_single_res_stack = self.align_res_stack_to_original_blocks(controlnet_single_res_stack, 38, hidden_states[:, prompt_emb.shape[1]:])
+
+ return controlnet_res_stack, controlnet_single_res_stack
+
+
+ # @staticmethod
+ # def state_dict_converter():
+ # return FluxControlNetStateDictConverter()
+
+ def quantize(self):
+ def cast_to(weight, dtype=None, device=None, copy=False):
+ if device is None or weight.device == device:
+ if not copy:
+ if dtype is None or weight.dtype == dtype:
+ return weight
+ return weight.to(dtype=dtype, copy=copy)
+
+ r = torch.empty_like(weight, dtype=dtype, device=device)
+ r.copy_(weight)
+ return r
+
+ def cast_weight(s, input=None, dtype=None, device=None):
+ if input is not None:
+ if dtype is None:
+ dtype = input.dtype
+ if device is None:
+ device = input.device
+ weight = cast_to(s.weight, dtype, device)
+ return weight
+
+ def cast_bias_weight(s, input=None, dtype=None, device=None, bias_dtype=None):
+ if input is not None:
+ if dtype is None:
+ dtype = input.dtype
+ if bias_dtype is None:
+ bias_dtype = dtype
+ if device is None:
+ device = input.device
+ bias = None
+ weight = cast_to(s.weight, dtype, device)
+ bias = cast_to(s.bias, bias_dtype, device)
+ return weight, bias
+
+ class quantized_layer:
+ class QLinear(torch.nn.Linear):
+ def __init__(self, *args, **kwargs):
+ super().__init__(*args, **kwargs)
+
+ def forward(self,input,**kwargs):
+ weight,bias= cast_bias_weight(self,input)
+ return torch.nn.functional.linear(input,weight,bias)
+
+ class QRMSNorm(torch.nn.Module):
+ def __init__(self, module):
+ super().__init__()
+ self.module = module
+
+ def forward(self,hidden_states,**kwargs):
+ weight= cast_weight(self.module,hidden_states)
+ input_dtype = hidden_states.dtype
+ variance = hidden_states.to(torch.float32).square().mean(-1, keepdim=True)
+ hidden_states = hidden_states * torch.rsqrt(variance + self.module.eps)
+ hidden_states = hidden_states.to(input_dtype) * weight
+ return hidden_states
+
+ class QEmbedding(torch.nn.Embedding):
+ def __init__(self, *args, **kwargs):
+ super().__init__(*args, **kwargs)
+
+ def forward(self,input,**kwargs):
+ weight= cast_weight(self,input)
+ return torch.nn.functional.embedding(
+ input, weight, self.padding_idx, self.max_norm,
+ self.norm_type, self.scale_grad_by_freq, self.sparse)
+
+ def replace_layer(model):
+ for name, module in model.named_children():
+ if isinstance(module,quantized_layer.QRMSNorm):
+ continue
+ if isinstance(module, torch.nn.Linear):
+ with init_weights_on_device():
+ new_layer = quantized_layer.QLinear(module.in_features,module.out_features)
+ new_layer.weight = module.weight
+ if module.bias is not None:
+ new_layer.bias = module.bias
+ setattr(model, name, new_layer)
+ elif isinstance(module, RMSNorm):
+ if hasattr(module,"quantized"):
+ continue
+ module.quantized= True
+ new_layer = quantized_layer.QRMSNorm(module)
+ setattr(model, name, new_layer)
+ elif isinstance(module,torch.nn.Embedding):
+ rows, cols = module.weight.shape
+ new_layer = quantized_layer.QEmbedding(
+ num_embeddings=rows,
+ embedding_dim=cols,
+ _weight=module.weight,
+ # _freeze=module.freeze,
+ padding_idx=module.padding_idx,
+ max_norm=module.max_norm,
+ norm_type=module.norm_type,
+ scale_grad_by_freq=module.scale_grad_by_freq,
+ sparse=module.sparse)
+ setattr(model, name, new_layer)
+ else:
+ replace_layer(module)
+
+ replace_layer(self)
+
+
+
+class FluxControlNetStateDictConverter:
+ def __init__(self):
+ pass
+
+ def from_diffusers(self, state_dict):
+ hash_value = hash_state_dict_keys(state_dict)
+ global_rename_dict = {
+ "context_embedder": "context_embedder",
+ "x_embedder": "x_embedder",
+ "time_text_embed.timestep_embedder.linear_1": "time_embedder.timestep_embedder.0",
+ "time_text_embed.timestep_embedder.linear_2": "time_embedder.timestep_embedder.2",
+ "time_text_embed.guidance_embedder.linear_1": "guidance_embedder.timestep_embedder.0",
+ "time_text_embed.guidance_embedder.linear_2": "guidance_embedder.timestep_embedder.2",
+ "time_text_embed.text_embedder.linear_1": "pooled_text_embedder.0",
+ "time_text_embed.text_embedder.linear_2": "pooled_text_embedder.2",
+ "norm_out.linear": "final_norm_out.linear",
+ "proj_out": "final_proj_out",
+ }
+ rename_dict = {
+ "proj_out": "proj_out",
+ "norm1.linear": "norm1_a.linear",
+ "norm1_context.linear": "norm1_b.linear",
+ "attn.to_q": "attn.a_to_q",
+ "attn.to_k": "attn.a_to_k",
+ "attn.to_v": "attn.a_to_v",
+ "attn.to_out.0": "attn.a_to_out",
+ "attn.add_q_proj": "attn.b_to_q",
+ "attn.add_k_proj": "attn.b_to_k",
+ "attn.add_v_proj": "attn.b_to_v",
+ "attn.to_add_out": "attn.b_to_out",
+ "ff.net.0.proj": "ff_a.0",
+ "ff.net.2": "ff_a.2",
+ "ff_context.net.0.proj": "ff_b.0",
+ "ff_context.net.2": "ff_b.2",
+ "attn.norm_q": "attn.norm_q_a",
+ "attn.norm_k": "attn.norm_k_a",
+ "attn.norm_added_q": "attn.norm_q_b",
+ "attn.norm_added_k": "attn.norm_k_b",
+ }
+ rename_dict_single = {
+ "attn.to_q": "a_to_q",
+ "attn.to_k": "a_to_k",
+ "attn.to_v": "a_to_v",
+ "attn.norm_q": "norm_q_a",
+ "attn.norm_k": "norm_k_a",
+ "norm.linear": "norm.linear",
+ "proj_mlp": "proj_in_besides_attn",
+ "proj_out": "proj_out",
+ }
+ state_dict_ = {}
+ for name, param in state_dict.items():
+ if name.endswith(".weight") or name.endswith(".bias"):
+ suffix = ".weight" if name.endswith(".weight") else ".bias"
+ prefix = name[:-len(suffix)]
+ if prefix in global_rename_dict:
+ state_dict_[global_rename_dict[prefix] + suffix] = param
+ elif prefix.startswith("transformer_blocks."):
+ names = prefix.split(".")
+ names[0] = "blocks"
+ middle = ".".join(names[2:])
+ if middle in rename_dict:
+ name_ = ".".join(names[:2] + [rename_dict[middle]] + [suffix[1:]])
+ state_dict_[name_] = param
+ elif prefix.startswith("single_transformer_blocks."):
+ names = prefix.split(".")
+ names[0] = "single_blocks"
+ middle = ".".join(names[2:])
+ if middle in rename_dict_single:
+ name_ = ".".join(names[:2] + [rename_dict_single[middle]] + [suffix[1:]])
+ state_dict_[name_] = param
+ else:
+ state_dict_[name] = param
+ else:
+ state_dict_[name] = param
+ for name in list(state_dict_.keys()):
+ if ".proj_in_besides_attn." in name:
+ name_ = name.replace(".proj_in_besides_attn.", ".to_qkv_mlp.")
+ param = torch.concat([
+ state_dict_[name.replace(".proj_in_besides_attn.", f".a_to_q.")],
+ state_dict_[name.replace(".proj_in_besides_attn.", f".a_to_k.")],
+ state_dict_[name.replace(".proj_in_besides_attn.", f".a_to_v.")],
+ state_dict_[name],
+ ], dim=0)
+ state_dict_[name_] = param
+ state_dict_.pop(name.replace(".proj_in_besides_attn.", f".a_to_q."))
+ state_dict_.pop(name.replace(".proj_in_besides_attn.", f".a_to_k."))
+ state_dict_.pop(name.replace(".proj_in_besides_attn.", f".a_to_v."))
+ state_dict_.pop(name)
+ for name in list(state_dict_.keys()):
+ for component in ["a", "b"]:
+ if f".{component}_to_q." in name:
+ name_ = name.replace(f".{component}_to_q.", f".{component}_to_qkv.")
+ param = torch.concat([
+ state_dict_[name.replace(f".{component}_to_q.", f".{component}_to_q.")],
+ state_dict_[name.replace(f".{component}_to_q.", f".{component}_to_k.")],
+ state_dict_[name.replace(f".{component}_to_q.", f".{component}_to_v.")],
+ ], dim=0)
+ state_dict_[name_] = param
+ state_dict_.pop(name.replace(f".{component}_to_q.", f".{component}_to_q."))
+ state_dict_.pop(name.replace(f".{component}_to_q.", f".{component}_to_k."))
+ state_dict_.pop(name.replace(f".{component}_to_q.", f".{component}_to_v."))
+ if hash_value == "78d18b9101345ff695f312e7e62538c0":
+ extra_kwargs = {"num_mode": 10, "mode_dict": {"canny": 0, "tile": 1, "depth": 2, "blur": 3, "pose": 4, "gray": 5, "lq": 6}}
+ elif hash_value == "b001c89139b5f053c715fe772362dd2a":
+ extra_kwargs = {"num_single_blocks": 0}
+ elif hash_value == "52357cb26250681367488a8954c271e8":
+ extra_kwargs = {"num_joint_blocks": 6, "num_single_blocks": 0, "additional_input_dim": 4}
+ elif hash_value == "0cfd1740758423a2a854d67c136d1e8c":
+ extra_kwargs = {"num_joint_blocks": 4, "num_single_blocks": 1}
+ elif hash_value == "7f9583eb8ba86642abb9a21a4b2c9e16":
+ extra_kwargs = {"num_joint_blocks": 4, "num_single_blocks": 10}
+ elif hash_value == "43ad5aaa27dd4ee01b832ed16773fa52":
+ extra_kwargs = {"num_joint_blocks": 6, "num_single_blocks": 0}
+ else:
+ extra_kwargs = {}
+ return state_dict_, extra_kwargs
+
+
+ def from_civitai(self, state_dict):
+ return self.from_diffusers(state_dict)
diff --git a/diffsynth/models/flux_dit.py b/diffsynth/models/flux_dit.py
new file mode 100644
index 0000000000000000000000000000000000000000..46fa861dedeaee37c7f58f477f544f802f6deffa
--- /dev/null
+++ b/diffsynth/models/flux_dit.py
@@ -0,0 +1,398 @@
+import torch
+from .general_modules import TimestepEmbeddings, AdaLayerNorm, RMSNorm
+from einops import rearrange
+
+
+def interact_with_ipadapter(hidden_states, q, ip_k, ip_v, scale=1.0):
+ batch_size, num_tokens = hidden_states.shape[0:2]
+ ip_hidden_states = torch.nn.functional.scaled_dot_product_attention(q, ip_k, ip_v)
+ ip_hidden_states = ip_hidden_states.transpose(1, 2).reshape(batch_size, num_tokens, -1)
+ hidden_states = hidden_states + scale * ip_hidden_states
+ return hidden_states
+
+
+class RoPEEmbedding(torch.nn.Module):
+ def __init__(self, dim, theta, axes_dim):
+ super().__init__()
+ self.dim = dim
+ self.theta = theta
+ self.axes_dim = axes_dim
+
+
+ def rope(self, pos: torch.Tensor, dim: int, theta: int) -> torch.Tensor:
+ assert dim % 2 == 0, "The dimension must be even."
+
+ scale = torch.arange(0, dim, 2, dtype=torch.float64, device=pos.device) / dim
+ omega = 1.0 / (theta**scale)
+
+ batch_size, seq_length = pos.shape
+ out = torch.einsum("...n,d->...nd", pos, omega)
+ cos_out = torch.cos(out)
+ sin_out = torch.sin(out)
+
+ stacked_out = torch.stack([cos_out, -sin_out, sin_out, cos_out], dim=-1)
+ out = stacked_out.view(batch_size, -1, dim // 2, 2, 2)
+ return out.float()
+
+
+ def forward(self, ids):
+ n_axes = ids.shape[-1]
+ emb = torch.cat([self.rope(ids[..., i], self.axes_dim[i], self.theta) for i in range(n_axes)], dim=-3)
+ return emb.unsqueeze(1)
+
+
+
+class FluxJointAttention(torch.nn.Module):
+ def __init__(self, dim_a, dim_b, num_heads, head_dim, only_out_a=False):
+ super().__init__()
+ self.num_heads = num_heads
+ self.head_dim = head_dim
+ self.only_out_a = only_out_a
+
+ self.a_to_qkv = torch.nn.Linear(dim_a, dim_a * 3)
+ self.b_to_qkv = torch.nn.Linear(dim_b, dim_b * 3)
+
+ self.norm_q_a = RMSNorm(head_dim, eps=1e-6)
+ self.norm_k_a = RMSNorm(head_dim, eps=1e-6)
+ self.norm_q_b = RMSNorm(head_dim, eps=1e-6)
+ self.norm_k_b = RMSNorm(head_dim, eps=1e-6)
+
+ self.a_to_out = torch.nn.Linear(dim_a, dim_a)
+ if not only_out_a:
+ self.b_to_out = torch.nn.Linear(dim_b, dim_b)
+
+
+ def apply_rope(self, xq, xk, freqs_cis):
+ xq_ = xq.float().reshape(*xq.shape[:-1], -1, 1, 2)
+ xk_ = xk.float().reshape(*xk.shape[:-1], -1, 1, 2)
+ xq_out = freqs_cis[..., 0] * xq_[..., 0] + freqs_cis[..., 1] * xq_[..., 1]
+ xk_out = freqs_cis[..., 0] * xk_[..., 0] + freqs_cis[..., 1] * xk_[..., 1]
+ return xq_out.reshape(*xq.shape).type_as(xq), xk_out.reshape(*xk.shape).type_as(xk)
+
+ def forward(self, hidden_states_a, hidden_states_b, image_rotary_emb, attn_mask=None, ipadapter_kwargs_list=None):
+ batch_size = hidden_states_a.shape[0]
+
+ # Part A
+ qkv_a = self.a_to_qkv(hidden_states_a)
+ qkv_a = qkv_a.view(batch_size, -1, 3 * self.num_heads, self.head_dim).transpose(1, 2)
+ q_a, k_a, v_a = qkv_a.chunk(3, dim=1)
+ q_a, k_a = self.norm_q_a(q_a), self.norm_k_a(k_a)
+
+ # Part B
+ qkv_b = self.b_to_qkv(hidden_states_b)
+ qkv_b = qkv_b.view(batch_size, -1, 3 * self.num_heads, self.head_dim).transpose(1, 2)
+ q_b, k_b, v_b = qkv_b.chunk(3, dim=1)
+ q_b, k_b = self.norm_q_b(q_b), self.norm_k_b(k_b)
+
+ q = torch.concat([q_b, q_a], dim=2)
+ k = torch.concat([k_b, k_a], dim=2)
+ v = torch.concat([v_b, v_a], dim=2)
+
+ q, k = self.apply_rope(q, k, image_rotary_emb)
+
+ hidden_states = torch.nn.functional.scaled_dot_product_attention(q, k, v, attn_mask=attn_mask)
+ hidden_states = hidden_states.transpose(1, 2).reshape(batch_size, -1, self.num_heads * self.head_dim)
+ hidden_states = hidden_states.to(q.dtype)
+ hidden_states_b, hidden_states_a = hidden_states[:, :hidden_states_b.shape[1]], hidden_states[:, hidden_states_b.shape[1]:]
+ if ipadapter_kwargs_list is not None:
+ hidden_states_a = interact_with_ipadapter(hidden_states_a, q_a, **ipadapter_kwargs_list)
+ hidden_states_a = self.a_to_out(hidden_states_a)
+ if self.only_out_a:
+ return hidden_states_a
+ else:
+ hidden_states_b = self.b_to_out(hidden_states_b)
+ return hidden_states_a, hidden_states_b
+
+
+
+class FluxJointTransformerBlock(torch.nn.Module):
+ def __init__(self, dim, num_attention_heads):
+ super().__init__()
+ self.norm1_a = AdaLayerNorm(dim)
+ self.norm1_b = AdaLayerNorm(dim)
+
+ self.attn = FluxJointAttention(dim, dim, num_attention_heads, dim // num_attention_heads)
+
+ self.norm2_a = torch.nn.LayerNorm(dim, elementwise_affine=False, eps=1e-6)
+ self.ff_a = torch.nn.Sequential(
+ torch.nn.Linear(dim, dim*4),
+ torch.nn.GELU(approximate="tanh"),
+ torch.nn.Linear(dim*4, dim)
+ )
+
+ self.norm2_b = torch.nn.LayerNorm(dim, elementwise_affine=False, eps=1e-6)
+ self.ff_b = torch.nn.Sequential(
+ torch.nn.Linear(dim, dim*4),
+ torch.nn.GELU(approximate="tanh"),
+ torch.nn.Linear(dim*4, dim)
+ )
+
+
+ def forward(self, hidden_states_a, hidden_states_b, temb, image_rotary_emb, attn_mask=None, ipadapter_kwargs_list=None):
+ norm_hidden_states_a, gate_msa_a, shift_mlp_a, scale_mlp_a, gate_mlp_a = self.norm1_a(hidden_states_a, emb=temb)
+ norm_hidden_states_b, gate_msa_b, shift_mlp_b, scale_mlp_b, gate_mlp_b = self.norm1_b(hidden_states_b, emb=temb)
+
+ # Attention
+ attn_output_a, attn_output_b = self.attn(norm_hidden_states_a, norm_hidden_states_b, image_rotary_emb, attn_mask, ipadapter_kwargs_list)
+
+ # Part A
+ hidden_states_a = hidden_states_a + gate_msa_a * attn_output_a
+ norm_hidden_states_a = self.norm2_a(hidden_states_a) * (1 + scale_mlp_a) + shift_mlp_a
+ hidden_states_a = hidden_states_a + gate_mlp_a * self.ff_a(norm_hidden_states_a)
+
+ # Part B
+ hidden_states_b = hidden_states_b + gate_msa_b * attn_output_b
+ norm_hidden_states_b = self.norm2_b(hidden_states_b) * (1 + scale_mlp_b) + shift_mlp_b
+ hidden_states_b = hidden_states_b + gate_mlp_b * self.ff_b(norm_hidden_states_b)
+
+ return hidden_states_a, hidden_states_b
+
+
+
+class FluxSingleAttention(torch.nn.Module):
+ def __init__(self, dim_a, dim_b, num_heads, head_dim):
+ super().__init__()
+ self.num_heads = num_heads
+ self.head_dim = head_dim
+
+ self.a_to_qkv = torch.nn.Linear(dim_a, dim_a * 3)
+
+ self.norm_q_a = RMSNorm(head_dim, eps=1e-6)
+ self.norm_k_a = RMSNorm(head_dim, eps=1e-6)
+
+
+ def apply_rope(self, xq, xk, freqs_cis):
+ xq_ = xq.float().reshape(*xq.shape[:-1], -1, 1, 2)
+ xk_ = xk.float().reshape(*xk.shape[:-1], -1, 1, 2)
+ xq_out = freqs_cis[..., 0] * xq_[..., 0] + freqs_cis[..., 1] * xq_[..., 1]
+ xk_out = freqs_cis[..., 0] * xk_[..., 0] + freqs_cis[..., 1] * xk_[..., 1]
+ return xq_out.reshape(*xq.shape).type_as(xq), xk_out.reshape(*xk.shape).type_as(xk)
+
+
+ def forward(self, hidden_states, image_rotary_emb):
+ batch_size = hidden_states.shape[0]
+
+ qkv_a = self.a_to_qkv(hidden_states)
+ qkv_a = qkv_a.view(batch_size, -1, 3 * self.num_heads, self.head_dim).transpose(1, 2)
+ q_a, k_a, v = qkv_a.chunk(3, dim=1)
+ q_a, k_a = self.norm_q_a(q_a), self.norm_k_a(k_a)
+
+ q, k = self.apply_rope(q_a, k_a, image_rotary_emb)
+
+ hidden_states = torch.nn.functional.scaled_dot_product_attention(q, k, v)
+ hidden_states = hidden_states.transpose(1, 2).reshape(batch_size, -1, self.num_heads * self.head_dim)
+ hidden_states = hidden_states.to(q.dtype)
+ return hidden_states
+
+
+
+class AdaLayerNormSingle(torch.nn.Module):
+ def __init__(self, dim):
+ super().__init__()
+ self.silu = torch.nn.SiLU()
+ self.linear = torch.nn.Linear(dim, 3 * dim, bias=True)
+ self.norm = torch.nn.LayerNorm(dim, elementwise_affine=False, eps=1e-6)
+
+
+ def forward(self, x, emb):
+ emb = self.linear(self.silu(emb))
+ shift_msa, scale_msa, gate_msa = emb.chunk(3, dim=1)
+ x = self.norm(x) * (1 + scale_msa[:, None]) + shift_msa[:, None]
+ return x, gate_msa
+
+
+
+class FluxSingleTransformerBlock(torch.nn.Module):
+ def __init__(self, dim, num_attention_heads):
+ super().__init__()
+ self.num_heads = num_attention_heads
+ self.head_dim = dim // num_attention_heads
+ self.dim = dim
+
+ self.norm = AdaLayerNormSingle(dim)
+ self.to_qkv_mlp = torch.nn.Linear(dim, dim * (3 + 4))
+ self.norm_q_a = RMSNorm(self.head_dim, eps=1e-6)
+ self.norm_k_a = RMSNorm(self.head_dim, eps=1e-6)
+
+ self.proj_out = torch.nn.Linear(dim * 5, dim)
+
+
+ def apply_rope(self, xq, xk, freqs_cis):
+ xq_ = xq.float().reshape(*xq.shape[:-1], -1, 1, 2)
+ xk_ = xk.float().reshape(*xk.shape[:-1], -1, 1, 2)
+ xq_out = freqs_cis[..., 0] * xq_[..., 0] + freqs_cis[..., 1] * xq_[..., 1]
+ xk_out = freqs_cis[..., 0] * xk_[..., 0] + freqs_cis[..., 1] * xk_[..., 1]
+ return xq_out.reshape(*xq.shape).type_as(xq), xk_out.reshape(*xk.shape).type_as(xk)
+
+
+ def process_attention(self, hidden_states, image_rotary_emb, attn_mask=None, ipadapter_kwargs_list=None):
+ batch_size = hidden_states.shape[0]
+
+ qkv = hidden_states.view(batch_size, -1, 3 * self.num_heads, self.head_dim).transpose(1, 2)
+ q, k, v = qkv.chunk(3, dim=1)
+ q, k = self.norm_q_a(q), self.norm_k_a(k)
+
+ q, k = self.apply_rope(q, k, image_rotary_emb)
+
+ hidden_states = torch.nn.functional.scaled_dot_product_attention(q, k, v, attn_mask=attn_mask)
+ hidden_states = hidden_states.transpose(1, 2).reshape(batch_size, -1, self.num_heads * self.head_dim)
+ hidden_states = hidden_states.to(q.dtype)
+ if ipadapter_kwargs_list is not None:
+ hidden_states = interact_with_ipadapter(hidden_states, q, **ipadapter_kwargs_list)
+ return hidden_states
+
+
+ def forward(self, hidden_states_a, hidden_states_b, temb, image_rotary_emb, attn_mask=None, ipadapter_kwargs_list=None):
+ residual = hidden_states_a
+ norm_hidden_states, gate = self.norm(hidden_states_a, emb=temb)
+ hidden_states_a = self.to_qkv_mlp(norm_hidden_states)
+ attn_output, mlp_hidden_states = hidden_states_a[:, :, :self.dim * 3], hidden_states_a[:, :, self.dim * 3:]
+
+ attn_output = self.process_attention(attn_output, image_rotary_emb, attn_mask, ipadapter_kwargs_list)
+ mlp_hidden_states = torch.nn.functional.gelu(mlp_hidden_states, approximate="tanh")
+
+ hidden_states_a = torch.cat([attn_output, mlp_hidden_states], dim=2)
+ hidden_states_a = gate.unsqueeze(1) * self.proj_out(hidden_states_a)
+ hidden_states_a = residual + hidden_states_a
+
+ return hidden_states_a, hidden_states_b
+
+
+
+class AdaLayerNormContinuous(torch.nn.Module):
+ def __init__(self, dim):
+ super().__init__()
+ self.silu = torch.nn.SiLU()
+ self.linear = torch.nn.Linear(dim, dim * 2, bias=True)
+ self.norm = torch.nn.LayerNorm(dim, eps=1e-6, elementwise_affine=False)
+
+ def forward(self, x, conditioning):
+ emb = self.linear(self.silu(conditioning))
+ shift, scale = torch.chunk(emb, 2, dim=1)
+ x = self.norm(x) * (1 + scale)[:, None] + shift[:, None]
+ return x
+
+
+
+class FluxDiT(torch.nn.Module):
+
+ _repeated_blocks = ["FluxJointTransformerBlock", "FluxSingleTransformerBlock"]
+
+ def __init__(self, disable_guidance_embedder=False, input_dim=64, num_blocks=19):
+ super().__init__()
+ self.pos_embedder = RoPEEmbedding(3072, 10000, [16, 56, 56])
+ self.time_embedder = TimestepEmbeddings(256, 3072)
+ self.guidance_embedder = None if disable_guidance_embedder else TimestepEmbeddings(256, 3072)
+ self.pooled_text_embedder = torch.nn.Sequential(torch.nn.Linear(768, 3072), torch.nn.SiLU(), torch.nn.Linear(3072, 3072))
+ self.context_embedder = torch.nn.Linear(4096, 3072)
+ self.x_embedder = torch.nn.Linear(input_dim, 3072)
+
+ self.blocks = torch.nn.ModuleList([FluxJointTransformerBlock(3072, 24) for _ in range(num_blocks)])
+ self.single_blocks = torch.nn.ModuleList([FluxSingleTransformerBlock(3072, 24) for _ in range(38)])
+
+ self.final_norm_out = AdaLayerNormContinuous(3072)
+ self.final_proj_out = torch.nn.Linear(3072, 64)
+
+ self.input_dim = input_dim
+
+
+ def patchify(self, hidden_states):
+ hidden_states = rearrange(hidden_states, "B C (H P) (W Q) -> B (H W) (C P Q)", P=2, Q=2)
+ return hidden_states
+
+
+ def unpatchify(self, hidden_states, height, width):
+ hidden_states = rearrange(hidden_states, "B (H W) (C P Q) -> B C (H P) (W Q)", P=2, Q=2, H=height//2, W=width//2)
+ return hidden_states
+
+
+ def prepare_image_ids(self, latents):
+ batch_size, _, height, width = latents.shape
+ latent_image_ids = torch.zeros(height // 2, width // 2, 3)
+ latent_image_ids[..., 1] = latent_image_ids[..., 1] + torch.arange(height // 2)[:, None]
+ latent_image_ids[..., 2] = latent_image_ids[..., 2] + torch.arange(width // 2)[None, :]
+
+ latent_image_id_height, latent_image_id_width, latent_image_id_channels = latent_image_ids.shape
+
+ latent_image_ids = latent_image_ids[None, :].repeat(batch_size, 1, 1, 1)
+ latent_image_ids = latent_image_ids.reshape(
+ batch_size, latent_image_id_height * latent_image_id_width, latent_image_id_channels
+ )
+ latent_image_ids = latent_image_ids.to(device=latents.device, dtype=latents.dtype)
+
+ return latent_image_ids
+
+
+ def construct_mask(self, entity_masks, prompt_seq_len, image_seq_len):
+ N = len(entity_masks)
+ batch_size = entity_masks[0].shape[0]
+ total_seq_len = N * prompt_seq_len + image_seq_len
+ patched_masks = [self.patchify(entity_masks[i]) for i in range(N)]
+ attention_mask = torch.ones((batch_size, total_seq_len, total_seq_len), dtype=torch.bool).to(device=entity_masks[0].device)
+
+ image_start = N * prompt_seq_len
+ image_end = N * prompt_seq_len + image_seq_len
+ # prompt-image mask
+ for i in range(N):
+ prompt_start = i * prompt_seq_len
+ prompt_end = (i + 1) * prompt_seq_len
+ image_mask = torch.sum(patched_masks[i], dim=-1) > 0
+ image_mask = image_mask.unsqueeze(1).repeat(1, prompt_seq_len, 1)
+ # prompt update with image
+ attention_mask[:, prompt_start:prompt_end, image_start:image_end] = image_mask
+ # image update with prompt
+ attention_mask[:, image_start:image_end, prompt_start:prompt_end] = image_mask.transpose(1, 2)
+ # prompt-prompt mask
+ for i in range(N):
+ for j in range(N):
+ if i != j:
+ prompt_start_i = i * prompt_seq_len
+ prompt_end_i = (i + 1) * prompt_seq_len
+ prompt_start_j = j * prompt_seq_len
+ prompt_end_j = (j + 1) * prompt_seq_len
+ attention_mask[:, prompt_start_i:prompt_end_i, prompt_start_j:prompt_end_j] = False
+
+ attention_mask = attention_mask.float()
+ attention_mask[attention_mask == 0] = float('-inf')
+ attention_mask[attention_mask == 1] = 0
+ return attention_mask
+
+
+ def process_entity_masks(self, hidden_states, prompt_emb, entity_prompt_emb, entity_masks, text_ids, image_ids, repeat_dim):
+ max_masks = 0
+ attention_mask = None
+ prompt_embs = [prompt_emb]
+ if entity_masks is not None:
+ # entity_masks
+ batch_size, max_masks = entity_masks.shape[0], entity_masks.shape[1]
+ entity_masks = entity_masks.repeat(1, 1, repeat_dim, 1, 1)
+ entity_masks = [entity_masks[:, i, None].squeeze(1) for i in range(max_masks)]
+ # global mask
+ global_mask = torch.ones_like(entity_masks[0]).to(device=hidden_states.device, dtype=hidden_states.dtype)
+ entity_masks = entity_masks + [global_mask] # append global to last
+ # attention mask
+ attention_mask = self.construct_mask(entity_masks, prompt_emb.shape[1], hidden_states.shape[1])
+ attention_mask = attention_mask.to(device=hidden_states.device, dtype=hidden_states.dtype)
+ attention_mask = attention_mask.unsqueeze(1)
+ # embds: n_masks * b * seq * d
+ local_embs = [entity_prompt_emb[:, i, None].squeeze(1) for i in range(max_masks)]
+ prompt_embs = local_embs + prompt_embs # append global to last
+ prompt_embs = [self.context_embedder(prompt_emb) for prompt_emb in prompt_embs]
+ prompt_emb = torch.cat(prompt_embs, dim=1)
+
+ # positional embedding
+ text_ids = torch.cat([text_ids] * (max_masks + 1), dim=1)
+ image_rotary_emb = self.pos_embedder(torch.cat((text_ids, image_ids), dim=1))
+ return prompt_emb, image_rotary_emb, attention_mask
+
+
+ def forward(
+ self,
+ hidden_states,
+ timestep, prompt_emb, pooled_prompt_emb, guidance, text_ids, image_ids=None,
+ tiled=False, tile_size=128, tile_stride=64, entity_prompt_emb=None, entity_masks=None,
+ use_gradient_checkpointing=False,
+ **kwargs
+ ):
+ # (Deprecated) The real forward is in `pipelines.flux_image`.
+ return None
diff --git a/diffsynth/models/flux_infiniteyou.py b/diffsynth/models/flux_infiniteyou.py
new file mode 100644
index 0000000000000000000000000000000000000000..861538a4b02fb6a52edee662b6efcd60f78f6916
--- /dev/null
+++ b/diffsynth/models/flux_infiniteyou.py
@@ -0,0 +1,129 @@
+import math
+import torch
+import torch.nn as nn
+
+
+# FFN
+def FeedForward(dim, mult=4):
+ inner_dim = int(dim * mult)
+ return nn.Sequential(
+ nn.LayerNorm(dim),
+ nn.Linear(dim, inner_dim, bias=False),
+ nn.GELU(),
+ nn.Linear(inner_dim, dim, bias=False),
+ )
+
+
+def reshape_tensor(x, heads):
+ bs, length, width = x.shape
+ #(bs, length, width) --> (bs, length, n_heads, dim_per_head)
+ x = x.view(bs, length, heads, -1)
+ # (bs, length, n_heads, dim_per_head) --> (bs, n_heads, length, dim_per_head)
+ x = x.transpose(1, 2)
+ # (bs, n_heads, length, dim_per_head) --> (bs*n_heads, length, dim_per_head)
+ x = x.reshape(bs, heads, length, -1)
+ return x
+
+
+class PerceiverAttention(nn.Module):
+
+ def __init__(self, *, dim, dim_head=64, heads=8):
+ super().__init__()
+ self.scale = dim_head**-0.5
+ self.dim_head = dim_head
+ self.heads = heads
+ inner_dim = dim_head * heads
+
+ self.norm1 = nn.LayerNorm(dim)
+ self.norm2 = nn.LayerNorm(dim)
+
+ self.to_q = nn.Linear(dim, inner_dim, bias=False)
+ self.to_kv = nn.Linear(dim, inner_dim * 2, bias=False)
+ self.to_out = nn.Linear(inner_dim, dim, bias=False)
+
+ def forward(self, x, latents):
+ """
+ Args:
+ x (torch.Tensor): image features
+ shape (b, n1, D)
+ latent (torch.Tensor): latent features
+ shape (b, n2, D)
+ """
+ x = self.norm1(x)
+ latents = self.norm2(latents)
+
+ b, l, _ = latents.shape
+
+ q = self.to_q(latents)
+ kv_input = torch.cat((x, latents), dim=-2)
+ k, v = self.to_kv(kv_input).chunk(2, dim=-1)
+
+ q = reshape_tensor(q, self.heads)
+ k = reshape_tensor(k, self.heads)
+ v = reshape_tensor(v, self.heads)
+
+ # attention
+ scale = 1 / math.sqrt(math.sqrt(self.dim_head))
+ weight = (q * scale) @ (k * scale).transpose(-2, -1) # More stable with f16 than dividing afterwards
+ weight = torch.softmax(weight.float(), dim=-1).type(weight.dtype)
+ out = weight @ v
+
+ out = out.permute(0, 2, 1, 3).reshape(b, l, -1)
+
+ return self.to_out(out)
+
+
+class InfiniteYouImageProjector(nn.Module):
+
+ def __init__(
+ self,
+ dim=1280,
+ depth=4,
+ dim_head=64,
+ heads=20,
+ num_queries=8,
+ embedding_dim=512,
+ output_dim=4096,
+ ff_mult=4,
+ ):
+ super().__init__()
+ self.latents = nn.Parameter(torch.randn(1, num_queries, dim) / dim**0.5)
+ self.proj_in = nn.Linear(embedding_dim, dim)
+
+ self.proj_out = nn.Linear(dim, output_dim)
+ self.norm_out = nn.LayerNorm(output_dim)
+
+ self.layers = nn.ModuleList([])
+ for _ in range(depth):
+ self.layers.append(
+ nn.ModuleList([
+ PerceiverAttention(dim=dim, dim_head=dim_head, heads=heads),
+ FeedForward(dim=dim, mult=ff_mult),
+ ]))
+
+ def forward(self, x):
+
+ latents = self.latents.repeat(x.size(0), 1, 1)
+ latents = latents.to(dtype=x.dtype, device=x.device)
+
+ x = self.proj_in(x)
+
+ for attn, ff in self.layers:
+ latents = attn(x, latents) + latents
+ latents = ff(latents) + latents
+
+ latents = self.proj_out(latents)
+ return self.norm_out(latents)
+
+ @staticmethod
+ def state_dict_converter():
+ return FluxInfiniteYouImageProjectorStateDictConverter()
+
+
+class FluxInfiniteYouImageProjectorStateDictConverter:
+
+ def __init__(self):
+ pass
+
+ def from_diffusers(self, state_dict):
+ return state_dict['image_proj']
diff --git a/diffsynth/models/flux_ipadapter.py b/diffsynth/models/flux_ipadapter.py
new file mode 100644
index 0000000000000000000000000000000000000000..31176fc2c2a508388502b45dc27e4d2218f16eec
--- /dev/null
+++ b/diffsynth/models/flux_ipadapter.py
@@ -0,0 +1,110 @@
+from .general_modules import RMSNorm
+from transformers import SiglipVisionModel, SiglipVisionConfig
+import torch
+
+
+class SiglipVisionModelSO400M(SiglipVisionModel):
+ def __init__(self):
+ config = SiglipVisionConfig(
+ hidden_size=1152,
+ image_size=384,
+ intermediate_size=4304,
+ model_type="siglip_vision_model",
+ num_attention_heads=16,
+ num_hidden_layers=27,
+ patch_size=14,
+ architectures=["SiglipModel"],
+ initializer_factor=1.0,
+ torch_dtype="float32",
+ transformers_version="4.37.0.dev0"
+ )
+ super().__init__(config)
+
+class MLPProjModel(torch.nn.Module):
+ def __init__(self, cross_attention_dim=768, id_embeddings_dim=512, num_tokens=4):
+ super().__init__()
+
+ self.cross_attention_dim = cross_attention_dim
+ self.num_tokens = num_tokens
+
+ self.proj = torch.nn.Sequential(
+ torch.nn.Linear(id_embeddings_dim, id_embeddings_dim*2),
+ torch.nn.GELU(),
+ torch.nn.Linear(id_embeddings_dim*2, cross_attention_dim*num_tokens),
+ )
+ self.norm = torch.nn.LayerNorm(cross_attention_dim)
+
+ def forward(self, id_embeds):
+ x = self.proj(id_embeds)
+ x = x.reshape(-1, self.num_tokens, self.cross_attention_dim)
+ x = self.norm(x)
+ return x
+
+class IpAdapterModule(torch.nn.Module):
+ def __init__(self, num_attention_heads, attention_head_dim, input_dim):
+ super().__init__()
+ self.num_heads = num_attention_heads
+ self.head_dim = attention_head_dim
+ output_dim = num_attention_heads * attention_head_dim
+ self.to_k_ip = torch.nn.Linear(input_dim, output_dim, bias=False)
+ self.to_v_ip = torch.nn.Linear(input_dim, output_dim, bias=False)
+ self.norm_added_k = RMSNorm(attention_head_dim, eps=1e-5, elementwise_affine=False)
+
+
+ def forward(self, hidden_states):
+ batch_size = hidden_states.shape[0]
+ # ip_k
+ ip_k = self.to_k_ip(hidden_states)
+ ip_k = ip_k.view(batch_size, -1, self.num_heads, self.head_dim).transpose(1, 2)
+ ip_k = self.norm_added_k(ip_k)
+ # ip_v
+ ip_v = self.to_v_ip(hidden_states)
+ ip_v = ip_v.view(batch_size, -1, self.num_heads, self.head_dim).transpose(1, 2)
+ return ip_k, ip_v
+
+
+class FluxIpAdapter(torch.nn.Module):
+ def __init__(self, num_attention_heads=24, attention_head_dim=128, cross_attention_dim=4096, num_tokens=128, num_blocks=57):
+ super().__init__()
+ self.ipadapter_modules = torch.nn.ModuleList([IpAdapterModule(num_attention_heads, attention_head_dim, cross_attention_dim) for _ in range(num_blocks)])
+ self.image_proj = MLPProjModel(cross_attention_dim=cross_attention_dim, id_embeddings_dim=1152, num_tokens=num_tokens)
+ self.set_adapter()
+
+ def set_adapter(self):
+ self.call_block_id = {i:i for i in range(len(self.ipadapter_modules))}
+
+ def forward(self, hidden_states, scale=1.0):
+ hidden_states = self.image_proj(hidden_states)
+ hidden_states = hidden_states.view(1, -1, hidden_states.shape[-1])
+ ip_kv_dict = {}
+ for block_id in self.call_block_id:
+ ipadapter_id = self.call_block_id[block_id]
+ ip_k, ip_v = self.ipadapter_modules[ipadapter_id](hidden_states)
+ ip_kv_dict[block_id] = {
+ "ip_k": ip_k,
+ "ip_v": ip_v,
+ "scale": scale
+ }
+ return ip_kv_dict
+
+ @staticmethod
+ def state_dict_converter():
+ return FluxIpAdapterStateDictConverter()
+
+
+class FluxIpAdapterStateDictConverter:
+ def __init__(self):
+ pass
+
+ def from_diffusers(self, state_dict):
+ state_dict_ = {}
+ for name in state_dict["ip_adapter"]:
+ name_ = 'ipadapter_modules.' + name
+ state_dict_[name_] = state_dict["ip_adapter"][name]
+ for name in state_dict["image_proj"]:
+ name_ = "image_proj." + name
+ state_dict_[name_] = state_dict["image_proj"][name]
+ return state_dict_
+
+ def from_civitai(self, state_dict):
+ return self.from_diffusers(state_dict)
diff --git a/diffsynth/models/flux_lora_encoder.py b/diffsynth/models/flux_lora_encoder.py
new file mode 100644
index 0000000000000000000000000000000000000000..13589b0611f3140479ef4faa3b7a29371caa447b
--- /dev/null
+++ b/diffsynth/models/flux_lora_encoder.py
@@ -0,0 +1,521 @@
+import torch
+from einops import rearrange
+
+
+def low_version_attention(query, key, value, attn_bias=None):
+ scale = 1 / query.shape[-1] ** 0.5
+ query = query * scale
+ attn = torch.matmul(query, key.transpose(-2, -1))
+ if attn_bias is not None:
+ attn = attn + attn_bias
+ attn = attn.softmax(-1)
+ return attn @ value
+
+
+class Attention(torch.nn.Module):
+
+ def __init__(self, q_dim, num_heads, head_dim, kv_dim=None, bias_q=False, bias_kv=False, bias_out=False):
+ super().__init__()
+ dim_inner = head_dim * num_heads
+ kv_dim = kv_dim if kv_dim is not None else q_dim
+ self.num_heads = num_heads
+ self.head_dim = head_dim
+
+ self.to_q = torch.nn.Linear(q_dim, dim_inner, bias=bias_q)
+ self.to_k = torch.nn.Linear(kv_dim, dim_inner, bias=bias_kv)
+ self.to_v = torch.nn.Linear(kv_dim, dim_inner, bias=bias_kv)
+ self.to_out = torch.nn.Linear(dim_inner, q_dim, bias=bias_out)
+
+ def interact_with_ipadapter(self, hidden_states, q, ip_k, ip_v, scale=1.0):
+ batch_size = q.shape[0]
+ ip_k = ip_k.view(batch_size, -1, self.num_heads, self.head_dim).transpose(1, 2)
+ ip_v = ip_v.view(batch_size, -1, self.num_heads, self.head_dim).transpose(1, 2)
+ ip_hidden_states = torch.nn.functional.scaled_dot_product_attention(q, ip_k, ip_v)
+ hidden_states = hidden_states + scale * ip_hidden_states
+ return hidden_states
+
+ def torch_forward(self, hidden_states, encoder_hidden_states=None, attn_mask=None, ipadapter_kwargs=None, qkv_preprocessor=None):
+ if encoder_hidden_states is None:
+ encoder_hidden_states = hidden_states
+
+ batch_size = encoder_hidden_states.shape[0]
+
+ q = self.to_q(hidden_states)
+ k = self.to_k(encoder_hidden_states)
+ v = self.to_v(encoder_hidden_states)
+
+ q = q.view(batch_size, -1, self.num_heads, self.head_dim).transpose(1, 2)
+ k = k.view(batch_size, -1, self.num_heads, self.head_dim).transpose(1, 2)
+ v = v.view(batch_size, -1, self.num_heads, self.head_dim).transpose(1, 2)
+
+ if qkv_preprocessor is not None:
+ q, k, v = qkv_preprocessor(q, k, v)
+
+ hidden_states = torch.nn.functional.scaled_dot_product_attention(q, k, v, attn_mask=attn_mask)
+ if ipadapter_kwargs is not None:
+ hidden_states = self.interact_with_ipadapter(hidden_states, q, **ipadapter_kwargs)
+ hidden_states = hidden_states.transpose(1, 2).reshape(batch_size, -1, self.num_heads * self.head_dim)
+ hidden_states = hidden_states.to(q.dtype)
+
+ hidden_states = self.to_out(hidden_states)
+
+ return hidden_states
+
+ def xformers_forward(self, hidden_states, encoder_hidden_states=None, attn_mask=None):
+ if encoder_hidden_states is None:
+ encoder_hidden_states = hidden_states
+
+ q = self.to_q(hidden_states)
+ k = self.to_k(encoder_hidden_states)
+ v = self.to_v(encoder_hidden_states)
+
+ q = rearrange(q, "b f (n d) -> (b n) f d", n=self.num_heads)
+ k = rearrange(k, "b f (n d) -> (b n) f d", n=self.num_heads)
+ v = rearrange(v, "b f (n d) -> (b n) f d", n=self.num_heads)
+
+ if attn_mask is not None:
+ hidden_states = low_version_attention(q, k, v, attn_bias=attn_mask)
+ else:
+ import xformers.ops as xops
+ hidden_states = xops.memory_efficient_attention(q, k, v)
+ hidden_states = rearrange(hidden_states, "(b n) f d -> b f (n d)", n=self.num_heads)
+
+ hidden_states = hidden_states.to(q.dtype)
+ hidden_states = self.to_out(hidden_states)
+
+ return hidden_states
+
+ def forward(self, hidden_states, encoder_hidden_states=None, attn_mask=None, ipadapter_kwargs=None, qkv_preprocessor=None):
+ return self.torch_forward(hidden_states, encoder_hidden_states=encoder_hidden_states, attn_mask=attn_mask, ipadapter_kwargs=ipadapter_kwargs, qkv_preprocessor=qkv_preprocessor)
+
+
+
+
+
+class CLIPEncoderLayer(torch.nn.Module):
+ def __init__(self, embed_dim, intermediate_size, num_heads=12, head_dim=64, use_quick_gelu=True):
+ super().__init__()
+ self.attn = Attention(q_dim=embed_dim, num_heads=num_heads, head_dim=head_dim, bias_q=True, bias_kv=True, bias_out=True)
+ self.layer_norm1 = torch.nn.LayerNorm(embed_dim)
+ self.layer_norm2 = torch.nn.LayerNorm(embed_dim)
+ self.fc1 = torch.nn.Linear(embed_dim, intermediate_size)
+ self.fc2 = torch.nn.Linear(intermediate_size, embed_dim)
+
+ self.use_quick_gelu = use_quick_gelu
+
+ def quickGELU(self, x):
+ return x * torch.sigmoid(1.702 * x)
+
+ def forward(self, hidden_states, attn_mask=None):
+ residual = hidden_states
+
+ hidden_states = self.layer_norm1(hidden_states)
+ hidden_states = self.attn(hidden_states, attn_mask=attn_mask)
+ hidden_states = residual + hidden_states
+
+ residual = hidden_states
+ hidden_states = self.layer_norm2(hidden_states)
+ hidden_states = self.fc1(hidden_states)
+ if self.use_quick_gelu:
+ hidden_states = self.quickGELU(hidden_states)
+ else:
+ hidden_states = torch.nn.functional.gelu(hidden_states)
+ hidden_states = self.fc2(hidden_states)
+ hidden_states = residual + hidden_states
+
+ return hidden_states
+
+
+class SDTextEncoder(torch.nn.Module):
+ def __init__(self, embed_dim=768, vocab_size=49408, max_position_embeddings=77, num_encoder_layers=12, encoder_intermediate_size=3072):
+ super().__init__()
+
+ # token_embedding
+ self.token_embedding = torch.nn.Embedding(vocab_size, embed_dim)
+
+ # position_embeds (This is a fixed tensor)
+ self.position_embeds = torch.nn.Parameter(torch.zeros(1, max_position_embeddings, embed_dim))
+
+ # encoders
+ self.encoders = torch.nn.ModuleList([CLIPEncoderLayer(embed_dim, encoder_intermediate_size) for _ in range(num_encoder_layers)])
+
+ # attn_mask
+ self.attn_mask = self.attention_mask(max_position_embeddings)
+
+ # final_layer_norm
+ self.final_layer_norm = torch.nn.LayerNorm(embed_dim)
+
+ def attention_mask(self, length):
+ mask = torch.empty(length, length)
+ mask.fill_(float("-inf"))
+ mask.triu_(1)
+ return mask
+
+ def forward(self, input_ids, clip_skip=1):
+ embeds = self.token_embedding(input_ids) + self.position_embeds
+ attn_mask = self.attn_mask.to(device=embeds.device, dtype=embeds.dtype)
+ for encoder_id, encoder in enumerate(self.encoders):
+ embeds = encoder(embeds, attn_mask=attn_mask)
+ if encoder_id + clip_skip == len(self.encoders):
+ break
+ embeds = self.final_layer_norm(embeds)
+ return embeds
+
+ @staticmethod
+ def state_dict_converter():
+ return SDTextEncoderStateDictConverter()
+
+
+class SDTextEncoderStateDictConverter:
+ def __init__(self):
+ pass
+
+ def from_diffusers(self, state_dict):
+ rename_dict = {
+ "text_model.embeddings.token_embedding.weight": "token_embedding.weight",
+ "text_model.embeddings.position_embedding.weight": "position_embeds",
+ "text_model.final_layer_norm.weight": "final_layer_norm.weight",
+ "text_model.final_layer_norm.bias": "final_layer_norm.bias"
+ }
+ attn_rename_dict = {
+ "self_attn.q_proj": "attn.to_q",
+ "self_attn.k_proj": "attn.to_k",
+ "self_attn.v_proj": "attn.to_v",
+ "self_attn.out_proj": "attn.to_out",
+ "layer_norm1": "layer_norm1",
+ "layer_norm2": "layer_norm2",
+ "mlp.fc1": "fc1",
+ "mlp.fc2": "fc2",
+ }
+ state_dict_ = {}
+ for name in state_dict:
+ if name in rename_dict:
+ param = state_dict[name]
+ if name == "text_model.embeddings.position_embedding.weight":
+ param = param.reshape((1, param.shape[0], param.shape[1]))
+ state_dict_[rename_dict[name]] = param
+ elif name.startswith("text_model.encoder.layers."):
+ param = state_dict[name]
+ names = name.split(".")
+ layer_id, layer_type, tail = names[3], ".".join(names[4:-1]), names[-1]
+ name_ = ".".join(["encoders", layer_id, attn_rename_dict[layer_type], tail])
+ state_dict_[name_] = param
+ return state_dict_
+
+ def from_civitai(self, state_dict):
+ rename_dict = {
+ "cond_stage_model.transformer.text_model.embeddings.token_embedding.weight": "token_embedding.weight",
+ "cond_stage_model.transformer.text_model.encoder.layers.0.layer_norm1.bias": "encoders.0.layer_norm1.bias",
+ "cond_stage_model.transformer.text_model.encoder.layers.0.layer_norm1.weight": "encoders.0.layer_norm1.weight",
+ "cond_stage_model.transformer.text_model.encoder.layers.0.layer_norm2.bias": "encoders.0.layer_norm2.bias",
+ "cond_stage_model.transformer.text_model.encoder.layers.0.layer_norm2.weight": "encoders.0.layer_norm2.weight",
+ "cond_stage_model.transformer.text_model.encoder.layers.0.mlp.fc1.bias": "encoders.0.fc1.bias",
+ "cond_stage_model.transformer.text_model.encoder.layers.0.mlp.fc1.weight": "encoders.0.fc1.weight",
+ "cond_stage_model.transformer.text_model.encoder.layers.0.mlp.fc2.bias": "encoders.0.fc2.bias",
+ "cond_stage_model.transformer.text_model.encoder.layers.0.mlp.fc2.weight": "encoders.0.fc2.weight",
+ "cond_stage_model.transformer.text_model.encoder.layers.0.self_attn.k_proj.bias": "encoders.0.attn.to_k.bias",
+ "cond_stage_model.transformer.text_model.encoder.layers.0.self_attn.k_proj.weight": "encoders.0.attn.to_k.weight",
+ "cond_stage_model.transformer.text_model.encoder.layers.0.self_attn.out_proj.bias": "encoders.0.attn.to_out.bias",
+ "cond_stage_model.transformer.text_model.encoder.layers.0.self_attn.out_proj.weight": "encoders.0.attn.to_out.weight",
+ "cond_stage_model.transformer.text_model.encoder.layers.0.self_attn.q_proj.bias": "encoders.0.attn.to_q.bias",
+ "cond_stage_model.transformer.text_model.encoder.layers.0.self_attn.q_proj.weight": "encoders.0.attn.to_q.weight",
+ "cond_stage_model.transformer.text_model.encoder.layers.0.self_attn.v_proj.bias": "encoders.0.attn.to_v.bias",
+ "cond_stage_model.transformer.text_model.encoder.layers.0.self_attn.v_proj.weight": "encoders.0.attn.to_v.weight",
+ "cond_stage_model.transformer.text_model.encoder.layers.1.layer_norm1.bias": "encoders.1.layer_norm1.bias",
+ "cond_stage_model.transformer.text_model.encoder.layers.1.layer_norm1.weight": "encoders.1.layer_norm1.weight",
+ "cond_stage_model.transformer.text_model.encoder.layers.1.layer_norm2.bias": "encoders.1.layer_norm2.bias",
+ "cond_stage_model.transformer.text_model.encoder.layers.1.layer_norm2.weight": "encoders.1.layer_norm2.weight",
+ "cond_stage_model.transformer.text_model.encoder.layers.1.mlp.fc1.bias": "encoders.1.fc1.bias",
+ "cond_stage_model.transformer.text_model.encoder.layers.1.mlp.fc1.weight": "encoders.1.fc1.weight",
+ "cond_stage_model.transformer.text_model.encoder.layers.1.mlp.fc2.bias": "encoders.1.fc2.bias",
+ "cond_stage_model.transformer.text_model.encoder.layers.1.mlp.fc2.weight": "encoders.1.fc2.weight",
+ "cond_stage_model.transformer.text_model.encoder.layers.1.self_attn.k_proj.bias": "encoders.1.attn.to_k.bias",
+ "cond_stage_model.transformer.text_model.encoder.layers.1.self_attn.k_proj.weight": "encoders.1.attn.to_k.weight",
+ "cond_stage_model.transformer.text_model.encoder.layers.1.self_attn.out_proj.bias": "encoders.1.attn.to_out.bias",
+ "cond_stage_model.transformer.text_model.encoder.layers.1.self_attn.out_proj.weight": "encoders.1.attn.to_out.weight",
+ "cond_stage_model.transformer.text_model.encoder.layers.1.self_attn.q_proj.bias": "encoders.1.attn.to_q.bias",
+ "cond_stage_model.transformer.text_model.encoder.layers.1.self_attn.q_proj.weight": "encoders.1.attn.to_q.weight",
+ "cond_stage_model.transformer.text_model.encoder.layers.1.self_attn.v_proj.bias": "encoders.1.attn.to_v.bias",
+ "cond_stage_model.transformer.text_model.encoder.layers.1.self_attn.v_proj.weight": "encoders.1.attn.to_v.weight",
+ "cond_stage_model.transformer.text_model.encoder.layers.10.layer_norm1.bias": "encoders.10.layer_norm1.bias",
+ "cond_stage_model.transformer.text_model.encoder.layers.10.layer_norm1.weight": "encoders.10.layer_norm1.weight",
+ "cond_stage_model.transformer.text_model.encoder.layers.10.layer_norm2.bias": "encoders.10.layer_norm2.bias",
+ "cond_stage_model.transformer.text_model.encoder.layers.10.layer_norm2.weight": "encoders.10.layer_norm2.weight",
+ "cond_stage_model.transformer.text_model.encoder.layers.10.mlp.fc1.bias": "encoders.10.fc1.bias",
+ "cond_stage_model.transformer.text_model.encoder.layers.10.mlp.fc1.weight": "encoders.10.fc1.weight",
+ "cond_stage_model.transformer.text_model.encoder.layers.10.mlp.fc2.bias": "encoders.10.fc2.bias",
+ "cond_stage_model.transformer.text_model.encoder.layers.10.mlp.fc2.weight": "encoders.10.fc2.weight",
+ "cond_stage_model.transformer.text_model.encoder.layers.10.self_attn.k_proj.bias": "encoders.10.attn.to_k.bias",
+ "cond_stage_model.transformer.text_model.encoder.layers.10.self_attn.k_proj.weight": "encoders.10.attn.to_k.weight",
+ "cond_stage_model.transformer.text_model.encoder.layers.10.self_attn.out_proj.bias": "encoders.10.attn.to_out.bias",
+ "cond_stage_model.transformer.text_model.encoder.layers.10.self_attn.out_proj.weight": "encoders.10.attn.to_out.weight",
+ "cond_stage_model.transformer.text_model.encoder.layers.10.self_attn.q_proj.bias": "encoders.10.attn.to_q.bias",
+ "cond_stage_model.transformer.text_model.encoder.layers.10.self_attn.q_proj.weight": "encoders.10.attn.to_q.weight",
+ "cond_stage_model.transformer.text_model.encoder.layers.10.self_attn.v_proj.bias": "encoders.10.attn.to_v.bias",
+ "cond_stage_model.transformer.text_model.encoder.layers.10.self_attn.v_proj.weight": "encoders.10.attn.to_v.weight",
+ "cond_stage_model.transformer.text_model.encoder.layers.11.layer_norm1.bias": "encoders.11.layer_norm1.bias",
+ "cond_stage_model.transformer.text_model.encoder.layers.11.layer_norm1.weight": "encoders.11.layer_norm1.weight",
+ "cond_stage_model.transformer.text_model.encoder.layers.11.layer_norm2.bias": "encoders.11.layer_norm2.bias",
+ "cond_stage_model.transformer.text_model.encoder.layers.11.layer_norm2.weight": "encoders.11.layer_norm2.weight",
+ "cond_stage_model.transformer.text_model.encoder.layers.11.mlp.fc1.bias": "encoders.11.fc1.bias",
+ "cond_stage_model.transformer.text_model.encoder.layers.11.mlp.fc1.weight": "encoders.11.fc1.weight",
+ "cond_stage_model.transformer.text_model.encoder.layers.11.mlp.fc2.bias": "encoders.11.fc2.bias",
+ "cond_stage_model.transformer.text_model.encoder.layers.11.mlp.fc2.weight": "encoders.11.fc2.weight",
+ "cond_stage_model.transformer.text_model.encoder.layers.11.self_attn.k_proj.bias": "encoders.11.attn.to_k.bias",
+ "cond_stage_model.transformer.text_model.encoder.layers.11.self_attn.k_proj.weight": "encoders.11.attn.to_k.weight",
+ "cond_stage_model.transformer.text_model.encoder.layers.11.self_attn.out_proj.bias": "encoders.11.attn.to_out.bias",
+ "cond_stage_model.transformer.text_model.encoder.layers.11.self_attn.out_proj.weight": "encoders.11.attn.to_out.weight",
+ "cond_stage_model.transformer.text_model.encoder.layers.11.self_attn.q_proj.bias": "encoders.11.attn.to_q.bias",
+ "cond_stage_model.transformer.text_model.encoder.layers.11.self_attn.q_proj.weight": "encoders.11.attn.to_q.weight",
+ "cond_stage_model.transformer.text_model.encoder.layers.11.self_attn.v_proj.bias": "encoders.11.attn.to_v.bias",
+ "cond_stage_model.transformer.text_model.encoder.layers.11.self_attn.v_proj.weight": "encoders.11.attn.to_v.weight",
+ "cond_stage_model.transformer.text_model.encoder.layers.2.layer_norm1.bias": "encoders.2.layer_norm1.bias",
+ "cond_stage_model.transformer.text_model.encoder.layers.2.layer_norm1.weight": "encoders.2.layer_norm1.weight",
+ "cond_stage_model.transformer.text_model.encoder.layers.2.layer_norm2.bias": "encoders.2.layer_norm2.bias",
+ "cond_stage_model.transformer.text_model.encoder.layers.2.layer_norm2.weight": "encoders.2.layer_norm2.weight",
+ "cond_stage_model.transformer.text_model.encoder.layers.2.mlp.fc1.bias": "encoders.2.fc1.bias",
+ "cond_stage_model.transformer.text_model.encoder.layers.2.mlp.fc1.weight": "encoders.2.fc1.weight",
+ "cond_stage_model.transformer.text_model.encoder.layers.2.mlp.fc2.bias": "encoders.2.fc2.bias",
+ "cond_stage_model.transformer.text_model.encoder.layers.2.mlp.fc2.weight": "encoders.2.fc2.weight",
+ "cond_stage_model.transformer.text_model.encoder.layers.2.self_attn.k_proj.bias": "encoders.2.attn.to_k.bias",
+ "cond_stage_model.transformer.text_model.encoder.layers.2.self_attn.k_proj.weight": "encoders.2.attn.to_k.weight",
+ "cond_stage_model.transformer.text_model.encoder.layers.2.self_attn.out_proj.bias": "encoders.2.attn.to_out.bias",
+ "cond_stage_model.transformer.text_model.encoder.layers.2.self_attn.out_proj.weight": "encoders.2.attn.to_out.weight",
+ "cond_stage_model.transformer.text_model.encoder.layers.2.self_attn.q_proj.bias": "encoders.2.attn.to_q.bias",
+ "cond_stage_model.transformer.text_model.encoder.layers.2.self_attn.q_proj.weight": "encoders.2.attn.to_q.weight",
+ "cond_stage_model.transformer.text_model.encoder.layers.2.self_attn.v_proj.bias": "encoders.2.attn.to_v.bias",
+ "cond_stage_model.transformer.text_model.encoder.layers.2.self_attn.v_proj.weight": "encoders.2.attn.to_v.weight",
+ "cond_stage_model.transformer.text_model.encoder.layers.3.layer_norm1.bias": "encoders.3.layer_norm1.bias",
+ "cond_stage_model.transformer.text_model.encoder.layers.3.layer_norm1.weight": "encoders.3.layer_norm1.weight",
+ "cond_stage_model.transformer.text_model.encoder.layers.3.layer_norm2.bias": "encoders.3.layer_norm2.bias",
+ "cond_stage_model.transformer.text_model.encoder.layers.3.layer_norm2.weight": "encoders.3.layer_norm2.weight",
+ "cond_stage_model.transformer.text_model.encoder.layers.3.mlp.fc1.bias": "encoders.3.fc1.bias",
+ "cond_stage_model.transformer.text_model.encoder.layers.3.mlp.fc1.weight": "encoders.3.fc1.weight",
+ "cond_stage_model.transformer.text_model.encoder.layers.3.mlp.fc2.bias": "encoders.3.fc2.bias",
+ "cond_stage_model.transformer.text_model.encoder.layers.3.mlp.fc2.weight": "encoders.3.fc2.weight",
+ "cond_stage_model.transformer.text_model.encoder.layers.3.self_attn.k_proj.bias": "encoders.3.attn.to_k.bias",
+ "cond_stage_model.transformer.text_model.encoder.layers.3.self_attn.k_proj.weight": "encoders.3.attn.to_k.weight",
+ "cond_stage_model.transformer.text_model.encoder.layers.3.self_attn.out_proj.bias": "encoders.3.attn.to_out.bias",
+ "cond_stage_model.transformer.text_model.encoder.layers.3.self_attn.out_proj.weight": "encoders.3.attn.to_out.weight",
+ "cond_stage_model.transformer.text_model.encoder.layers.3.self_attn.q_proj.bias": "encoders.3.attn.to_q.bias",
+ "cond_stage_model.transformer.text_model.encoder.layers.3.self_attn.q_proj.weight": "encoders.3.attn.to_q.weight",
+ "cond_stage_model.transformer.text_model.encoder.layers.3.self_attn.v_proj.bias": "encoders.3.attn.to_v.bias",
+ "cond_stage_model.transformer.text_model.encoder.layers.3.self_attn.v_proj.weight": "encoders.3.attn.to_v.weight",
+ "cond_stage_model.transformer.text_model.encoder.layers.4.layer_norm1.bias": "encoders.4.layer_norm1.bias",
+ "cond_stage_model.transformer.text_model.encoder.layers.4.layer_norm1.weight": "encoders.4.layer_norm1.weight",
+ "cond_stage_model.transformer.text_model.encoder.layers.4.layer_norm2.bias": "encoders.4.layer_norm2.bias",
+ "cond_stage_model.transformer.text_model.encoder.layers.4.layer_norm2.weight": "encoders.4.layer_norm2.weight",
+ "cond_stage_model.transformer.text_model.encoder.layers.4.mlp.fc1.bias": "encoders.4.fc1.bias",
+ "cond_stage_model.transformer.text_model.encoder.layers.4.mlp.fc1.weight": "encoders.4.fc1.weight",
+ "cond_stage_model.transformer.text_model.encoder.layers.4.mlp.fc2.bias": "encoders.4.fc2.bias",
+ "cond_stage_model.transformer.text_model.encoder.layers.4.mlp.fc2.weight": "encoders.4.fc2.weight",
+ "cond_stage_model.transformer.text_model.encoder.layers.4.self_attn.k_proj.bias": "encoders.4.attn.to_k.bias",
+ "cond_stage_model.transformer.text_model.encoder.layers.4.self_attn.k_proj.weight": "encoders.4.attn.to_k.weight",
+ "cond_stage_model.transformer.text_model.encoder.layers.4.self_attn.out_proj.bias": "encoders.4.attn.to_out.bias",
+ "cond_stage_model.transformer.text_model.encoder.layers.4.self_attn.out_proj.weight": "encoders.4.attn.to_out.weight",
+ "cond_stage_model.transformer.text_model.encoder.layers.4.self_attn.q_proj.bias": "encoders.4.attn.to_q.bias",
+ "cond_stage_model.transformer.text_model.encoder.layers.4.self_attn.q_proj.weight": "encoders.4.attn.to_q.weight",
+ "cond_stage_model.transformer.text_model.encoder.layers.4.self_attn.v_proj.bias": "encoders.4.attn.to_v.bias",
+ "cond_stage_model.transformer.text_model.encoder.layers.4.self_attn.v_proj.weight": "encoders.4.attn.to_v.weight",
+ "cond_stage_model.transformer.text_model.encoder.layers.5.layer_norm1.bias": "encoders.5.layer_norm1.bias",
+ "cond_stage_model.transformer.text_model.encoder.layers.5.layer_norm1.weight": "encoders.5.layer_norm1.weight",
+ "cond_stage_model.transformer.text_model.encoder.layers.5.layer_norm2.bias": "encoders.5.layer_norm2.bias",
+ "cond_stage_model.transformer.text_model.encoder.layers.5.layer_norm2.weight": "encoders.5.layer_norm2.weight",
+ "cond_stage_model.transformer.text_model.encoder.layers.5.mlp.fc1.bias": "encoders.5.fc1.bias",
+ "cond_stage_model.transformer.text_model.encoder.layers.5.mlp.fc1.weight": "encoders.5.fc1.weight",
+ "cond_stage_model.transformer.text_model.encoder.layers.5.mlp.fc2.bias": "encoders.5.fc2.bias",
+ "cond_stage_model.transformer.text_model.encoder.layers.5.mlp.fc2.weight": "encoders.5.fc2.weight",
+ "cond_stage_model.transformer.text_model.encoder.layers.5.self_attn.k_proj.bias": "encoders.5.attn.to_k.bias",
+ "cond_stage_model.transformer.text_model.encoder.layers.5.self_attn.k_proj.weight": "encoders.5.attn.to_k.weight",
+ "cond_stage_model.transformer.text_model.encoder.layers.5.self_attn.out_proj.bias": "encoders.5.attn.to_out.bias",
+ "cond_stage_model.transformer.text_model.encoder.layers.5.self_attn.out_proj.weight": "encoders.5.attn.to_out.weight",
+ "cond_stage_model.transformer.text_model.encoder.layers.5.self_attn.q_proj.bias": "encoders.5.attn.to_q.bias",
+ "cond_stage_model.transformer.text_model.encoder.layers.5.self_attn.q_proj.weight": "encoders.5.attn.to_q.weight",
+ "cond_stage_model.transformer.text_model.encoder.layers.5.self_attn.v_proj.bias": "encoders.5.attn.to_v.bias",
+ "cond_stage_model.transformer.text_model.encoder.layers.5.self_attn.v_proj.weight": "encoders.5.attn.to_v.weight",
+ "cond_stage_model.transformer.text_model.encoder.layers.6.layer_norm1.bias": "encoders.6.layer_norm1.bias",
+ "cond_stage_model.transformer.text_model.encoder.layers.6.layer_norm1.weight": "encoders.6.layer_norm1.weight",
+ "cond_stage_model.transformer.text_model.encoder.layers.6.layer_norm2.bias": "encoders.6.layer_norm2.bias",
+ "cond_stage_model.transformer.text_model.encoder.layers.6.layer_norm2.weight": "encoders.6.layer_norm2.weight",
+ "cond_stage_model.transformer.text_model.encoder.layers.6.mlp.fc1.bias": "encoders.6.fc1.bias",
+ "cond_stage_model.transformer.text_model.encoder.layers.6.mlp.fc1.weight": "encoders.6.fc1.weight",
+ "cond_stage_model.transformer.text_model.encoder.layers.6.mlp.fc2.bias": "encoders.6.fc2.bias",
+ "cond_stage_model.transformer.text_model.encoder.layers.6.mlp.fc2.weight": "encoders.6.fc2.weight",
+ "cond_stage_model.transformer.text_model.encoder.layers.6.self_attn.k_proj.bias": "encoders.6.attn.to_k.bias",
+ "cond_stage_model.transformer.text_model.encoder.layers.6.self_attn.k_proj.weight": "encoders.6.attn.to_k.weight",
+ "cond_stage_model.transformer.text_model.encoder.layers.6.self_attn.out_proj.bias": "encoders.6.attn.to_out.bias",
+ "cond_stage_model.transformer.text_model.encoder.layers.6.self_attn.out_proj.weight": "encoders.6.attn.to_out.weight",
+ "cond_stage_model.transformer.text_model.encoder.layers.6.self_attn.q_proj.bias": "encoders.6.attn.to_q.bias",
+ "cond_stage_model.transformer.text_model.encoder.layers.6.self_attn.q_proj.weight": "encoders.6.attn.to_q.weight",
+ "cond_stage_model.transformer.text_model.encoder.layers.6.self_attn.v_proj.bias": "encoders.6.attn.to_v.bias",
+ "cond_stage_model.transformer.text_model.encoder.layers.6.self_attn.v_proj.weight": "encoders.6.attn.to_v.weight",
+ "cond_stage_model.transformer.text_model.encoder.layers.7.layer_norm1.bias": "encoders.7.layer_norm1.bias",
+ "cond_stage_model.transformer.text_model.encoder.layers.7.layer_norm1.weight": "encoders.7.layer_norm1.weight",
+ "cond_stage_model.transformer.text_model.encoder.layers.7.layer_norm2.bias": "encoders.7.layer_norm2.bias",
+ "cond_stage_model.transformer.text_model.encoder.layers.7.layer_norm2.weight": "encoders.7.layer_norm2.weight",
+ "cond_stage_model.transformer.text_model.encoder.layers.7.mlp.fc1.bias": "encoders.7.fc1.bias",
+ "cond_stage_model.transformer.text_model.encoder.layers.7.mlp.fc1.weight": "encoders.7.fc1.weight",
+ "cond_stage_model.transformer.text_model.encoder.layers.7.mlp.fc2.bias": "encoders.7.fc2.bias",
+ "cond_stage_model.transformer.text_model.encoder.layers.7.mlp.fc2.weight": "encoders.7.fc2.weight",
+ "cond_stage_model.transformer.text_model.encoder.layers.7.self_attn.k_proj.bias": "encoders.7.attn.to_k.bias",
+ "cond_stage_model.transformer.text_model.encoder.layers.7.self_attn.k_proj.weight": "encoders.7.attn.to_k.weight",
+ "cond_stage_model.transformer.text_model.encoder.layers.7.self_attn.out_proj.bias": "encoders.7.attn.to_out.bias",
+ "cond_stage_model.transformer.text_model.encoder.layers.7.self_attn.out_proj.weight": "encoders.7.attn.to_out.weight",
+ "cond_stage_model.transformer.text_model.encoder.layers.7.self_attn.q_proj.bias": "encoders.7.attn.to_q.bias",
+ "cond_stage_model.transformer.text_model.encoder.layers.7.self_attn.q_proj.weight": "encoders.7.attn.to_q.weight",
+ "cond_stage_model.transformer.text_model.encoder.layers.7.self_attn.v_proj.bias": "encoders.7.attn.to_v.bias",
+ "cond_stage_model.transformer.text_model.encoder.layers.7.self_attn.v_proj.weight": "encoders.7.attn.to_v.weight",
+ "cond_stage_model.transformer.text_model.encoder.layers.8.layer_norm1.bias": "encoders.8.layer_norm1.bias",
+ "cond_stage_model.transformer.text_model.encoder.layers.8.layer_norm1.weight": "encoders.8.layer_norm1.weight",
+ "cond_stage_model.transformer.text_model.encoder.layers.8.layer_norm2.bias": "encoders.8.layer_norm2.bias",
+ "cond_stage_model.transformer.text_model.encoder.layers.8.layer_norm2.weight": "encoders.8.layer_norm2.weight",
+ "cond_stage_model.transformer.text_model.encoder.layers.8.mlp.fc1.bias": "encoders.8.fc1.bias",
+ "cond_stage_model.transformer.text_model.encoder.layers.8.mlp.fc1.weight": "encoders.8.fc1.weight",
+ "cond_stage_model.transformer.text_model.encoder.layers.8.mlp.fc2.bias": "encoders.8.fc2.bias",
+ "cond_stage_model.transformer.text_model.encoder.layers.8.mlp.fc2.weight": "encoders.8.fc2.weight",
+ "cond_stage_model.transformer.text_model.encoder.layers.8.self_attn.k_proj.bias": "encoders.8.attn.to_k.bias",
+ "cond_stage_model.transformer.text_model.encoder.layers.8.self_attn.k_proj.weight": "encoders.8.attn.to_k.weight",
+ "cond_stage_model.transformer.text_model.encoder.layers.8.self_attn.out_proj.bias": "encoders.8.attn.to_out.bias",
+ "cond_stage_model.transformer.text_model.encoder.layers.8.self_attn.out_proj.weight": "encoders.8.attn.to_out.weight",
+ "cond_stage_model.transformer.text_model.encoder.layers.8.self_attn.q_proj.bias": "encoders.8.attn.to_q.bias",
+ "cond_stage_model.transformer.text_model.encoder.layers.8.self_attn.q_proj.weight": "encoders.8.attn.to_q.weight",
+ "cond_stage_model.transformer.text_model.encoder.layers.8.self_attn.v_proj.bias": "encoders.8.attn.to_v.bias",
+ "cond_stage_model.transformer.text_model.encoder.layers.8.self_attn.v_proj.weight": "encoders.8.attn.to_v.weight",
+ "cond_stage_model.transformer.text_model.encoder.layers.9.layer_norm1.bias": "encoders.9.layer_norm1.bias",
+ "cond_stage_model.transformer.text_model.encoder.layers.9.layer_norm1.weight": "encoders.9.layer_norm1.weight",
+ "cond_stage_model.transformer.text_model.encoder.layers.9.layer_norm2.bias": "encoders.9.layer_norm2.bias",
+ "cond_stage_model.transformer.text_model.encoder.layers.9.layer_norm2.weight": "encoders.9.layer_norm2.weight",
+ "cond_stage_model.transformer.text_model.encoder.layers.9.mlp.fc1.bias": "encoders.9.fc1.bias",
+ "cond_stage_model.transformer.text_model.encoder.layers.9.mlp.fc1.weight": "encoders.9.fc1.weight",
+ "cond_stage_model.transformer.text_model.encoder.layers.9.mlp.fc2.bias": "encoders.9.fc2.bias",
+ "cond_stage_model.transformer.text_model.encoder.layers.9.mlp.fc2.weight": "encoders.9.fc2.weight",
+ "cond_stage_model.transformer.text_model.encoder.layers.9.self_attn.k_proj.bias": "encoders.9.attn.to_k.bias",
+ "cond_stage_model.transformer.text_model.encoder.layers.9.self_attn.k_proj.weight": "encoders.9.attn.to_k.weight",
+ "cond_stage_model.transformer.text_model.encoder.layers.9.self_attn.out_proj.bias": "encoders.9.attn.to_out.bias",
+ "cond_stage_model.transformer.text_model.encoder.layers.9.self_attn.out_proj.weight": "encoders.9.attn.to_out.weight",
+ "cond_stage_model.transformer.text_model.encoder.layers.9.self_attn.q_proj.bias": "encoders.9.attn.to_q.bias",
+ "cond_stage_model.transformer.text_model.encoder.layers.9.self_attn.q_proj.weight": "encoders.9.attn.to_q.weight",
+ "cond_stage_model.transformer.text_model.encoder.layers.9.self_attn.v_proj.bias": "encoders.9.attn.to_v.bias",
+ "cond_stage_model.transformer.text_model.encoder.layers.9.self_attn.v_proj.weight": "encoders.9.attn.to_v.weight",
+ "cond_stage_model.transformer.text_model.final_layer_norm.bias": "final_layer_norm.bias",
+ "cond_stage_model.transformer.text_model.final_layer_norm.weight": "final_layer_norm.weight",
+ "cond_stage_model.transformer.text_model.embeddings.position_embedding.weight": "position_embeds"
+ }
+ state_dict_ = {}
+ for name in state_dict:
+ if name in rename_dict:
+ param = state_dict[name]
+ if name == "cond_stage_model.transformer.text_model.embeddings.position_embedding.weight":
+ param = param.reshape((1, param.shape[0], param.shape[1]))
+ state_dict_[rename_dict[name]] = param
+ return state_dict_
+
+
+
+class LoRALayerBlock(torch.nn.Module):
+ def __init__(self, L, dim_in, dim_out):
+ super().__init__()
+ self.x = torch.nn.Parameter(torch.randn(1, L, dim_in))
+ self.layer_norm = torch.nn.LayerNorm(dim_out)
+
+ def forward(self, lora_A, lora_B):
+ x = self.x @ lora_A.T @ lora_B.T
+ x = self.layer_norm(x)
+ return x
+
+
+class LoRAEmbedder(torch.nn.Module):
+ def __init__(self, lora_patterns=None, L=1, out_dim=2048):
+ super().__init__()
+ if lora_patterns is None:
+ lora_patterns = self.default_lora_patterns()
+
+ model_dict = {}
+ for lora_pattern in lora_patterns:
+ name, dim = lora_pattern["name"], lora_pattern["dim"]
+ model_dict[name.replace(".", "___")] = LoRALayerBlock(L, dim[0], dim[1])
+ self.model_dict = torch.nn.ModuleDict(model_dict)
+
+ proj_dict = {}
+ for lora_pattern in lora_patterns:
+ layer_type, dim = lora_pattern["type"], lora_pattern["dim"]
+ if layer_type not in proj_dict:
+ proj_dict[layer_type.replace(".", "___")] = torch.nn.Linear(dim[1], out_dim)
+ self.proj_dict = torch.nn.ModuleDict(proj_dict)
+
+ self.lora_patterns = lora_patterns
+
+
+ def default_lora_patterns(self):
+ lora_patterns = []
+ lora_dict = {
+ "attn.a_to_qkv": (3072, 9216), "attn.a_to_out": (3072, 3072), "ff_a.0": (3072, 12288), "ff_a.2": (12288, 3072), "norm1_a.linear": (3072, 18432),
+ "attn.b_to_qkv": (3072, 9216), "attn.b_to_out": (3072, 3072), "ff_b.0": (3072, 12288), "ff_b.2": (12288, 3072), "norm1_b.linear": (3072, 18432),
+ }
+ for i in range(19):
+ for suffix in lora_dict:
+ lora_patterns.append({
+ "name": f"blocks.{i}.{suffix}",
+ "dim": lora_dict[suffix],
+ "type": suffix,
+ })
+ lora_dict = {"to_qkv_mlp": (3072, 21504), "proj_out": (15360, 3072), "norm.linear": (3072, 9216)}
+ for i in range(38):
+ for suffix in lora_dict:
+ lora_patterns.append({
+ "name": f"single_blocks.{i}.{suffix}",
+ "dim": lora_dict[suffix],
+ "type": suffix,
+ })
+ return lora_patterns
+
+ def forward(self, lora):
+ lora_emb = []
+ for lora_pattern in self.lora_patterns:
+ name, layer_type = lora_pattern["name"], lora_pattern["type"]
+ lora_A = lora[name + ".lora_A.weight"]
+ lora_B = lora[name + ".lora_B.weight"]
+ lora_out = self.model_dict[name.replace(".", "___")](lora_A, lora_B)
+ lora_out = self.proj_dict[layer_type.replace(".", "___")](lora_out)
+ lora_emb.append(lora_out)
+ lora_emb = torch.concat(lora_emb, dim=1)
+ return lora_emb
+
+
+class FluxLoRAEncoder(torch.nn.Module):
+ def __init__(self, embed_dim=4096, encoder_intermediate_size=8192, num_encoder_layers=1, num_embeds_per_lora=16, num_special_embeds=1):
+ super().__init__()
+ self.num_embeds_per_lora = num_embeds_per_lora
+ # embedder
+ self.embedder = LoRAEmbedder(L=num_embeds_per_lora, out_dim=embed_dim)
+
+ # encoders
+ self.encoders = torch.nn.ModuleList([CLIPEncoderLayer(embed_dim, encoder_intermediate_size, num_heads=32, head_dim=128) for _ in range(num_encoder_layers)])
+
+ # special embedding
+ self.special_embeds = torch.nn.Parameter(torch.randn(1, num_special_embeds, embed_dim))
+ self.num_special_embeds = num_special_embeds
+
+ # final layer
+ self.final_layer_norm = torch.nn.LayerNorm(embed_dim)
+ self.final_linear = torch.nn.Linear(embed_dim, embed_dim)
+
+ def forward(self, lora):
+ lora_embeds = self.embedder(lora)
+ special_embeds = self.special_embeds.to(dtype=lora_embeds.dtype, device=lora_embeds.device)
+ embeds = torch.concat([special_embeds, lora_embeds], dim=1)
+ for encoder_id, encoder in enumerate(self.encoders):
+ embeds = encoder(embeds)
+ embeds = embeds[:, :self.num_special_embeds]
+ embeds = self.final_layer_norm(embeds)
+ embeds = self.final_linear(embeds)
+ return embeds
+
+ @staticmethod
+ def state_dict_converter():
+ return FluxLoRAEncoderStateDictConverter()
+
+
+class FluxLoRAEncoderStateDictConverter:
+ def from_civitai(self, state_dict):
+ return state_dict
diff --git a/diffsynth/models/flux_lora_patcher.py b/diffsynth/models/flux_lora_patcher.py
new file mode 100644
index 0000000000000000000000000000000000000000..5d8fc8cea03bbbc658d8e1869432d903b6ae7ce9
--- /dev/null
+++ b/diffsynth/models/flux_lora_patcher.py
@@ -0,0 +1,306 @@
+import torch, math
+from ..core.loader import load_state_dict
+from typing import Union
+
+class GeneralLoRALoader:
+ def __init__(self, device="cpu", torch_dtype=torch.float32):
+ self.device = device
+ self.torch_dtype = torch_dtype
+
+
+ def get_name_dict(self, lora_state_dict):
+ lora_name_dict = {}
+ for key in lora_state_dict:
+ if ".lora_B." not in key:
+ continue
+ keys = key.split(".")
+ if len(keys) > keys.index("lora_B") + 2:
+ keys.pop(keys.index("lora_B") + 1)
+ keys.pop(keys.index("lora_B"))
+ if keys[0] == "diffusion_model":
+ keys.pop(0)
+ keys.pop(-1)
+ target_name = ".".join(keys)
+ lora_name_dict[target_name] = (key, key.replace(".lora_B.", ".lora_A."))
+ return lora_name_dict
+
+
+ def load(self, model: torch.nn.Module, state_dict_lora, alpha=1.0):
+ updated_num = 0
+ lora_name_dict = self.get_name_dict(state_dict_lora)
+ for name, module in model.named_modules():
+ if name in lora_name_dict:
+ weight_up = state_dict_lora[lora_name_dict[name][0]].to(device=self.device, dtype=self.torch_dtype)
+ weight_down = state_dict_lora[lora_name_dict[name][1]].to(device=self.device, dtype=self.torch_dtype)
+ if len(weight_up.shape) == 4:
+ weight_up = weight_up.squeeze(3).squeeze(2)
+ weight_down = weight_down.squeeze(3).squeeze(2)
+ weight_lora = alpha * torch.mm(weight_up, weight_down).unsqueeze(2).unsqueeze(3)
+ else:
+ weight_lora = alpha * torch.mm(weight_up, weight_down)
+ state_dict = module.state_dict()
+ state_dict["weight"] = state_dict["weight"].to(device=self.device, dtype=self.torch_dtype) + weight_lora
+ module.load_state_dict(state_dict)
+ updated_num += 1
+ print(f"{updated_num} tensors are updated by LoRA.")
+
+class FluxLoRALoader(GeneralLoRALoader):
+ def __init__(self, device="cpu", torch_dtype=torch.float32):
+ super().__init__(device=device, torch_dtype=torch_dtype)
+
+ self.diffusers_rename_dict = {
+ "transformer.single_transformer_blocks.blockid.attn.to_k.lora_A.weight":"single_blocks.blockid.a_to_k.lora_A.default.weight",
+ "transformer.single_transformer_blocks.blockid.attn.to_k.lora_B.weight":"single_blocks.blockid.a_to_k.lora_B.default.weight",
+ "transformer.single_transformer_blocks.blockid.attn.to_q.lora_A.weight":"single_blocks.blockid.a_to_q.lora_A.default.weight",
+ "transformer.single_transformer_blocks.blockid.attn.to_q.lora_B.weight":"single_blocks.blockid.a_to_q.lora_B.default.weight",
+ "transformer.single_transformer_blocks.blockid.attn.to_v.lora_A.weight":"single_blocks.blockid.a_to_v.lora_A.default.weight",
+ "transformer.single_transformer_blocks.blockid.attn.to_v.lora_B.weight":"single_blocks.blockid.a_to_v.lora_B.default.weight",
+ "transformer.single_transformer_blocks.blockid.norm.linear.lora_A.weight":"single_blocks.blockid.norm.linear.lora_A.default.weight",
+ "transformer.single_transformer_blocks.blockid.norm.linear.lora_B.weight":"single_blocks.blockid.norm.linear.lora_B.default.weight",
+ "transformer.single_transformer_blocks.blockid.proj_mlp.lora_A.weight":"single_blocks.blockid.proj_in_besides_attn.lora_A.default.weight",
+ "transformer.single_transformer_blocks.blockid.proj_mlp.lora_B.weight":"single_blocks.blockid.proj_in_besides_attn.lora_B.default.weight",
+ "transformer.single_transformer_blocks.blockid.proj_out.lora_A.weight":"single_blocks.blockid.proj_out.lora_A.default.weight",
+ "transformer.single_transformer_blocks.blockid.proj_out.lora_B.weight":"single_blocks.blockid.proj_out.lora_B.default.weight",
+ "transformer.transformer_blocks.blockid.attn.add_k_proj.lora_A.weight":"blocks.blockid.attn.b_to_k.lora_A.default.weight",
+ "transformer.transformer_blocks.blockid.attn.add_k_proj.lora_B.weight":"blocks.blockid.attn.b_to_k.lora_B.default.weight",
+ "transformer.transformer_blocks.blockid.attn.add_q_proj.lora_A.weight":"blocks.blockid.attn.b_to_q.lora_A.default.weight",
+ "transformer.transformer_blocks.blockid.attn.add_q_proj.lora_B.weight":"blocks.blockid.attn.b_to_q.lora_B.default.weight",
+ "transformer.transformer_blocks.blockid.attn.add_v_proj.lora_A.weight":"blocks.blockid.attn.b_to_v.lora_A.default.weight",
+ "transformer.transformer_blocks.blockid.attn.add_v_proj.lora_B.weight":"blocks.blockid.attn.b_to_v.lora_B.default.weight",
+ "transformer.transformer_blocks.blockid.attn.to_add_out.lora_A.weight":"blocks.blockid.attn.b_to_out.lora_A.default.weight",
+ "transformer.transformer_blocks.blockid.attn.to_add_out.lora_B.weight":"blocks.blockid.attn.b_to_out.lora_B.default.weight",
+ "transformer.transformer_blocks.blockid.attn.to_k.lora_A.weight":"blocks.blockid.attn.a_to_k.lora_A.default.weight",
+ "transformer.transformer_blocks.blockid.attn.to_k.lora_B.weight":"blocks.blockid.attn.a_to_k.lora_B.default.weight",
+ "transformer.transformer_blocks.blockid.attn.to_out.0.lora_A.weight":"blocks.blockid.attn.a_to_out.lora_A.default.weight",
+ "transformer.transformer_blocks.blockid.attn.to_out.0.lora_B.weight":"blocks.blockid.attn.a_to_out.lora_B.default.weight",
+ "transformer.transformer_blocks.blockid.attn.to_q.lora_A.weight":"blocks.blockid.attn.a_to_q.lora_A.default.weight",
+ "transformer.transformer_blocks.blockid.attn.to_q.lora_B.weight":"blocks.blockid.attn.a_to_q.lora_B.default.weight",
+ "transformer.transformer_blocks.blockid.attn.to_v.lora_A.weight":"blocks.blockid.attn.a_to_v.lora_A.default.weight",
+ "transformer.transformer_blocks.blockid.attn.to_v.lora_B.weight":"blocks.blockid.attn.a_to_v.lora_B.default.weight",
+ "transformer.transformer_blocks.blockid.ff.net.0.proj.lora_A.weight":"blocks.blockid.ff_a.0.lora_A.default.weight",
+ "transformer.transformer_blocks.blockid.ff.net.0.proj.lora_B.weight":"blocks.blockid.ff_a.0.lora_B.default.weight",
+ "transformer.transformer_blocks.blockid.ff.net.2.lora_A.weight":"blocks.blockid.ff_a.2.lora_A.default.weight",
+ "transformer.transformer_blocks.blockid.ff.net.2.lora_B.weight":"blocks.blockid.ff_a.2.lora_B.default.weight",
+ "transformer.transformer_blocks.blockid.ff_context.net.0.proj.lora_A.weight":"blocks.blockid.ff_b.0.lora_A.default.weight",
+ "transformer.transformer_blocks.blockid.ff_context.net.0.proj.lora_B.weight":"blocks.blockid.ff_b.0.lora_B.default.weight",
+ "transformer.transformer_blocks.blockid.ff_context.net.2.lora_A.weight":"blocks.blockid.ff_b.2.lora_A.default.weight",
+ "transformer.transformer_blocks.blockid.ff_context.net.2.lora_B.weight":"blocks.blockid.ff_b.2.lora_B.default.weight",
+ "transformer.transformer_blocks.blockid.norm1.linear.lora_A.weight":"blocks.blockid.norm1_a.linear.lora_A.default.weight",
+ "transformer.transformer_blocks.blockid.norm1.linear.lora_B.weight":"blocks.blockid.norm1_a.linear.lora_B.default.weight",
+ "transformer.transformer_blocks.blockid.norm1_context.linear.lora_A.weight":"blocks.blockid.norm1_b.linear.lora_A.default.weight",
+ "transformer.transformer_blocks.blockid.norm1_context.linear.lora_B.weight":"blocks.blockid.norm1_b.linear.lora_B.default.weight",
+ }
+
+ self.civitai_rename_dict = {
+ "lora_unet_double_blocks_blockid_img_mod_lin.lora_down.weight": "blocks.blockid.norm1_a.linear.lora_A.default.weight",
+ "lora_unet_double_blocks_blockid_img_mod_lin.lora_up.weight": "blocks.blockid.norm1_a.linear.lora_B.default.weight",
+ "lora_unet_double_blocks_blockid_txt_mod_lin.lora_down.weight": "blocks.blockid.norm1_b.linear.lora_A.default.weight",
+ "lora_unet_double_blocks_blockid_txt_mod_lin.lora_up.weight": "blocks.blockid.norm1_b.linear.lora_B.default.weight",
+ "lora_unet_double_blocks_blockid_img_attn_qkv.lora_down.weight": "blocks.blockid.attn.a_to_qkv.lora_A.default.weight",
+ "lora_unet_double_blocks_blockid_img_attn_qkv.lora_up.weight": "blocks.blockid.attn.a_to_qkv.lora_B.default.weight",
+ "lora_unet_double_blocks_blockid_txt_attn_qkv.lora_down.weight": "blocks.blockid.attn.b_to_qkv.lora_A.default.weight",
+ "lora_unet_double_blocks_blockid_txt_attn_qkv.lora_up.weight": "blocks.blockid.attn.b_to_qkv.lora_B.default.weight",
+ "lora_unet_double_blocks_blockid_img_attn_proj.lora_down.weight": "blocks.blockid.attn.a_to_out.lora_A.default.weight",
+ "lora_unet_double_blocks_blockid_img_attn_proj.lora_up.weight": "blocks.blockid.attn.a_to_out.lora_B.default.weight",
+ "lora_unet_double_blocks_blockid_txt_attn_proj.lora_down.weight": "blocks.blockid.attn.b_to_out.lora_A.default.weight",
+ "lora_unet_double_blocks_blockid_txt_attn_proj.lora_up.weight": "blocks.blockid.attn.b_to_out.lora_B.default.weight",
+ "lora_unet_double_blocks_blockid_img_mlp_0.lora_down.weight": "blocks.blockid.ff_a.0.lora_A.default.weight",
+ "lora_unet_double_blocks_blockid_img_mlp_0.lora_up.weight": "blocks.blockid.ff_a.0.lora_B.default.weight",
+ "lora_unet_double_blocks_blockid_img_mlp_2.lora_down.weight": "blocks.blockid.ff_a.2.lora_A.default.weight",
+ "lora_unet_double_blocks_blockid_img_mlp_2.lora_up.weight": "blocks.blockid.ff_a.2.lora_B.default.weight",
+ "lora_unet_double_blocks_blockid_txt_mlp_0.lora_down.weight": "blocks.blockid.ff_b.0.lora_A.default.weight",
+ "lora_unet_double_blocks_blockid_txt_mlp_0.lora_up.weight": "blocks.blockid.ff_b.0.lora_B.default.weight",
+ "lora_unet_double_blocks_blockid_txt_mlp_2.lora_down.weight": "blocks.blockid.ff_b.2.lora_A.default.weight",
+ "lora_unet_double_blocks_blockid_txt_mlp_2.lora_up.weight": "blocks.blockid.ff_b.2.lora_B.default.weight",
+ "lora_unet_single_blocks_blockid_modulation_lin.lora_down.weight": "single_blocks.blockid.norm.linear.lora_A.default.weight",
+ "lora_unet_single_blocks_blockid_modulation_lin.lora_up.weight": "single_blocks.blockid.norm.linear.lora_B.default.weight",
+ "lora_unet_single_blocks_blockid_linear1.lora_down.weight": "single_blocks.blockid.to_qkv_mlp.lora_A.default.weight",
+ "lora_unet_single_blocks_blockid_linear1.lora_up.weight": "single_blocks.blockid.to_qkv_mlp.lora_B.default.weight",
+ "lora_unet_single_blocks_blockid_linear2.lora_down.weight": "single_blocks.blockid.proj_out.lora_A.default.weight",
+ "lora_unet_single_blocks_blockid_linear2.lora_up.weight": "single_blocks.blockid.proj_out.lora_B.default.weight",
+ }
+
+ def load(self, model: torch.nn.Module, state_dict_lora, alpha=1.0):
+ super().load(model, state_dict_lora, alpha)
+
+
+ def convert_state_dict(self,state_dict):
+
+ def guess_block_id(name,model_resource):
+ if model_resource == 'civitai':
+ names = name.split("_")
+ for i in names:
+ if i.isdigit():
+ return i, name.replace(f"_{i}_", "_blockid_")
+ if model_resource == 'diffusers':
+ names = name.split(".")
+ for i in names:
+ if i.isdigit():
+ return i, name.replace(f"transformer_blocks.{i}.", "transformer_blocks.blockid.")
+ return None, None
+
+ def guess_resource(state_dict):
+ for k in state_dict:
+ if "lora_unet_" in k:
+ return 'civitai'
+ elif k.startswith("transformer."):
+ return 'diffusers'
+ else:
+ None
+
+ model_resource = guess_resource(state_dict)
+ if model_resource is None:
+ return state_dict
+
+ rename_dict = self.diffusers_rename_dict if model_resource == 'diffusers' else self.civitai_rename_dict
+ def guess_alpha(state_dict):
+ for name, param in state_dict.items():
+ if ".alpha" in name:
+ for suffix in [".lora_down.weight", ".lora_A.weight"]:
+ name_ = name.replace(".alpha", suffix)
+ if name_ in state_dict:
+ lora_alpha = param.item() / state_dict[name_].shape[0]
+ lora_alpha = math.sqrt(lora_alpha)
+ return lora_alpha
+
+ return 1
+
+ alpha = guess_alpha(state_dict)
+
+ state_dict_ = {}
+ for name, param in state_dict.items():
+ block_id, source_name = guess_block_id(name,model_resource)
+ if alpha != 1:
+ param *= alpha
+ if source_name in rename_dict:
+ target_name = rename_dict[source_name]
+ target_name = target_name.replace(".blockid.", f".{block_id}.")
+ state_dict_[target_name] = param
+ else:
+ state_dict_[name] = param
+
+ if model_resource == 'diffusers':
+ for name in list(state_dict_.keys()):
+ if "single_blocks." in name and ".a_to_q." in name:
+ mlp = state_dict_.get(name.replace(".a_to_q.", ".proj_in_besides_attn."), None)
+ if mlp is None:
+ dim = 4
+ if 'lora_A' in name:
+ dim = 1
+ mlp = torch.zeros(dim * state_dict_[name].shape[0],
+ *state_dict_[name].shape[1:],
+ dtype=state_dict_[name].dtype)
+ else:
+ state_dict_.pop(name.replace(".a_to_q.", ".proj_in_besides_attn."))
+ if 'lora_A' in name:
+ param = torch.concat([
+ state_dict_.pop(name),
+ state_dict_.pop(name.replace(".a_to_q.", ".a_to_k.")),
+ state_dict_.pop(name.replace(".a_to_q.", ".a_to_v.")),
+ mlp,
+ ], dim=0)
+ elif 'lora_B' in name:
+ d, r = state_dict_[name].shape
+ param = torch.zeros((3*d+mlp.shape[0], 3*r+mlp.shape[1]), dtype=state_dict_[name].dtype, device=state_dict_[name].device)
+ param[:d, :r] = state_dict_.pop(name)
+ param[d:2*d, r:2*r] = state_dict_.pop(name.replace(".a_to_q.", ".a_to_k."))
+ param[2*d:3*d, 2*r:3*r] = state_dict_.pop(name.replace(".a_to_q.", ".a_to_v."))
+ param[3*d:, 3*r:] = mlp
+ else:
+ param = torch.concat([
+ state_dict_.pop(name),
+ state_dict_.pop(name.replace(".a_to_q.", ".a_to_k.")),
+ state_dict_.pop(name.replace(".a_to_q.", ".a_to_v.")),
+ mlp,
+ ], dim=0)
+ name_ = name.replace(".a_to_q.", ".to_qkv_mlp.")
+ state_dict_[name_] = param
+ for name in list(state_dict_.keys()):
+ for component in ["a", "b"]:
+ if f".{component}_to_q." in name:
+ name_ = name.replace(f".{component}_to_q.", f".{component}_to_qkv.")
+ concat_dim = 0
+ if 'lora_A' in name:
+ param = torch.concat([
+ state_dict_[name.replace(f".{component}_to_q.", f".{component}_to_q.")],
+ state_dict_[name.replace(f".{component}_to_q.", f".{component}_to_k.")],
+ state_dict_[name.replace(f".{component}_to_q.", f".{component}_to_v.")],
+ ], dim=0)
+ elif 'lora_B' in name:
+ origin = state_dict_[name.replace(f".{component}_to_q.", f".{component}_to_q.")]
+ d, r = origin.shape
+ # print(d, r)
+ param = torch.zeros((3*d, 3*r), dtype=origin.dtype, device=origin.device)
+ param[:d, :r] = state_dict_[name.replace(f".{component}_to_q.", f".{component}_to_q.")]
+ param[d:2*d, r:2*r] = state_dict_[name.replace(f".{component}_to_q.", f".{component}_to_k.")]
+ param[2*d:3*d, 2*r:3*r] = state_dict_[name.replace(f".{component}_to_q.", f".{component}_to_v.")]
+ else:
+ param = torch.concat([
+ state_dict_[name.replace(f".{component}_to_q.", f".{component}_to_q.")],
+ state_dict_[name.replace(f".{component}_to_q.", f".{component}_to_k.")],
+ state_dict_[name.replace(f".{component}_to_q.", f".{component}_to_v.")],
+ ], dim=0)
+ state_dict_[name_] = param
+ state_dict_.pop(name.replace(f".{component}_to_q.", f".{component}_to_q."))
+ state_dict_.pop(name.replace(f".{component}_to_q.", f".{component}_to_k."))
+ state_dict_.pop(name.replace(f".{component}_to_q.", f".{component}_to_v."))
+ return state_dict_
+
+
+class LoraMerger(torch.nn.Module):
+ def __init__(self, dim):
+ super().__init__()
+ self.weight_base = torch.nn.Parameter(torch.randn((dim,)))
+ self.weight_lora = torch.nn.Parameter(torch.randn((dim,)))
+ self.weight_cross = torch.nn.Parameter(torch.randn((dim,)))
+ self.weight_out = torch.nn.Parameter(torch.ones((dim,)))
+ self.bias = torch.nn.Parameter(torch.randn((dim,)))
+ self.activation = torch.nn.Sigmoid()
+ self.norm_base = torch.nn.LayerNorm(dim, eps=1e-5)
+ self.norm_lora = torch.nn.LayerNorm(dim, eps=1e-5)
+
+ def forward(self, base_output, lora_outputs):
+ norm_base_output = self.norm_base(base_output)
+ norm_lora_outputs = self.norm_lora(lora_outputs)
+ gate = self.activation(
+ norm_base_output * self.weight_base \
+ + norm_lora_outputs * self.weight_lora \
+ + norm_base_output * norm_lora_outputs * self.weight_cross + self.bias
+ )
+ output = base_output + (self.weight_out * gate * lora_outputs).sum(dim=0)
+ return output
+
+class FluxLoraPatcher(torch.nn.Module):
+ def __init__(self, lora_patterns=None):
+ super().__init__()
+ if lora_patterns is None:
+ lora_patterns = self.default_lora_patterns()
+ model_dict = {}
+ for lora_pattern in lora_patterns:
+ name, dim = lora_pattern["name"], lora_pattern["dim"]
+ model_dict[name.replace(".", "___")] = LoraMerger(dim)
+ self.model_dict = torch.nn.ModuleDict(model_dict)
+
+ def default_lora_patterns(self):
+ lora_patterns = []
+ lora_dict = {
+ "attn.a_to_qkv": 9216, "attn.a_to_out": 3072, "ff_a.0": 12288, "ff_a.2": 3072, "norm1_a.linear": 18432,
+ "attn.b_to_qkv": 9216, "attn.b_to_out": 3072, "ff_b.0": 12288, "ff_b.2": 3072, "norm1_b.linear": 18432,
+ }
+ for i in range(19):
+ for suffix in lora_dict:
+ lora_patterns.append({
+ "name": f"blocks.{i}.{suffix}",
+ "dim": lora_dict[suffix]
+ })
+ lora_dict = {"to_qkv_mlp": 21504, "proj_out": 3072, "norm.linear": 9216}
+ for i in range(38):
+ for suffix in lora_dict:
+ lora_patterns.append({
+ "name": f"single_blocks.{i}.{suffix}",
+ "dim": lora_dict[suffix]
+ })
+ return lora_patterns
+
+ def forward(self, base_output, lora_outputs, name):
+ return self.model_dict[name.replace(".", "___")](base_output, lora_outputs)
diff --git a/diffsynth/models/flux_redux.py b/diffsynth/models/flux_redux.py
new file mode 100644
index 0000000000000000000000000000000000000000..99cdef69c982745347fe2d3a32ef4696b268e046
--- /dev/null
+++ b/diffsynth/models/flux_redux.py
@@ -0,0 +1,13 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+class FluxReduxImageEncoder(nn.Module):
+ def __init__(self, redux_dim: int = 1152, txt_in_features: int = 4096):
+ super().__init__()
+ self.redux_up = nn.Linear(redux_dim, txt_in_features * 3)
+ self.redux_down = nn.Linear(txt_in_features * 3, txt_in_features)
+
+ def forward(self, x: torch.Tensor) -> torch.Tensor:
+ return self.redux_down(F.silu(self.redux_up(x)))
diff --git a/diffsynth/models/flux_text_encoder_clip.py b/diffsynth/models/flux_text_encoder_clip.py
new file mode 100644
index 0000000000000000000000000000000000000000..1425423ce6d1df946198a16a7e96078ab8fed807
--- /dev/null
+++ b/diffsynth/models/flux_text_encoder_clip.py
@@ -0,0 +1,112 @@
+import torch
+
+
+class Attention(torch.nn.Module):
+
+ def __init__(self, q_dim, num_heads, head_dim, kv_dim=None, bias_q=False, bias_kv=False, bias_out=False):
+ super().__init__()
+ dim_inner = head_dim * num_heads
+ kv_dim = kv_dim if kv_dim is not None else q_dim
+ self.num_heads = num_heads
+ self.head_dim = head_dim
+
+ self.to_q = torch.nn.Linear(q_dim, dim_inner, bias=bias_q)
+ self.to_k = torch.nn.Linear(kv_dim, dim_inner, bias=bias_kv)
+ self.to_v = torch.nn.Linear(kv_dim, dim_inner, bias=bias_kv)
+ self.to_out = torch.nn.Linear(dim_inner, q_dim, bias=bias_out)
+
+ def forward(self, hidden_states, encoder_hidden_states=None, attn_mask=None):
+ if encoder_hidden_states is None:
+ encoder_hidden_states = hidden_states
+
+ batch_size = encoder_hidden_states.shape[0]
+
+ q = self.to_q(hidden_states)
+ k = self.to_k(encoder_hidden_states)
+ v = self.to_v(encoder_hidden_states)
+
+ q = q.view(batch_size, -1, self.num_heads, self.head_dim).transpose(1, 2)
+ k = k.view(batch_size, -1, self.num_heads, self.head_dim).transpose(1, 2)
+ v = v.view(batch_size, -1, self.num_heads, self.head_dim).transpose(1, 2)
+
+ hidden_states = torch.nn.functional.scaled_dot_product_attention(q, k, v, attn_mask=attn_mask)
+ hidden_states = hidden_states.transpose(1, 2).reshape(batch_size, -1, self.num_heads * self.head_dim)
+ hidden_states = hidden_states.to(q.dtype)
+
+ hidden_states = self.to_out(hidden_states)
+
+ return hidden_states
+
+
+class CLIPEncoderLayer(torch.nn.Module):
+ def __init__(self, embed_dim, intermediate_size, num_heads=12, head_dim=64, use_quick_gelu=True):
+ super().__init__()
+ self.attn = Attention(q_dim=embed_dim, num_heads=num_heads, head_dim=head_dim, bias_q=True, bias_kv=True, bias_out=True)
+ self.layer_norm1 = torch.nn.LayerNorm(embed_dim)
+ self.layer_norm2 = torch.nn.LayerNorm(embed_dim)
+ self.fc1 = torch.nn.Linear(embed_dim, intermediate_size)
+ self.fc2 = torch.nn.Linear(intermediate_size, embed_dim)
+
+ self.use_quick_gelu = use_quick_gelu
+
+ def quickGELU(self, x):
+ return x * torch.sigmoid(1.702 * x)
+
+ def forward(self, hidden_states, attn_mask=None):
+ residual = hidden_states
+
+ hidden_states = self.layer_norm1(hidden_states)
+ hidden_states = self.attn(hidden_states, attn_mask=attn_mask)
+ hidden_states = residual + hidden_states
+
+ residual = hidden_states
+ hidden_states = self.layer_norm2(hidden_states)
+ hidden_states = self.fc1(hidden_states)
+ if self.use_quick_gelu:
+ hidden_states = self.quickGELU(hidden_states)
+ else:
+ hidden_states = torch.nn.functional.gelu(hidden_states)
+ hidden_states = self.fc2(hidden_states)
+ hidden_states = residual + hidden_states
+
+ return hidden_states
+
+
+class FluxTextEncoderClip(torch.nn.Module):
+ def __init__(self, embed_dim=768, vocab_size=49408, max_position_embeddings=77, num_encoder_layers=12, encoder_intermediate_size=3072):
+ super().__init__()
+
+ # token_embedding
+ self.token_embedding = torch.nn.Embedding(vocab_size, embed_dim)
+
+ # position_embeds (This is a fixed tensor)
+ self.position_embeds = torch.nn.Parameter(torch.zeros(1, max_position_embeddings, embed_dim))
+
+ # encoders
+ self.encoders = torch.nn.ModuleList([CLIPEncoderLayer(embed_dim, encoder_intermediate_size) for _ in range(num_encoder_layers)])
+
+ # attn_mask
+ self.attn_mask = self.attention_mask(max_position_embeddings)
+
+ # final_layer_norm
+ self.final_layer_norm = torch.nn.LayerNorm(embed_dim)
+
+ def attention_mask(self, length):
+ mask = torch.empty(length, length)
+ mask.fill_(float("-inf"))
+ mask.triu_(1)
+ return mask
+
+ def forward(self, input_ids, clip_skip=2, extra_mask=None):
+ embeds = self.token_embedding(input_ids)
+ embeds = embeds + self.position_embeds.to(dtype=embeds.dtype, device=input_ids.device)
+ attn_mask = self.attn_mask.to(device=embeds.device, dtype=embeds.dtype)
+ if extra_mask is not None:
+ attn_mask[:, extra_mask[0]==0] = float("-inf")
+ for encoder_id, encoder in enumerate(self.encoders):
+ embeds = encoder(embeds, attn_mask=attn_mask)
+ if encoder_id + clip_skip == len(self.encoders):
+ hidden_states = embeds
+ embeds = self.final_layer_norm(embeds)
+ pooled_embeds = embeds[torch.arange(embeds.shape[0]), input_ids.to(dtype=torch.int).argmax(dim=-1)]
+ return pooled_embeds, hidden_states
diff --git a/diffsynth/models/flux_text_encoder_t5.py b/diffsynth/models/flux_text_encoder_t5.py
new file mode 100644
index 0000000000000000000000000000000000000000..ee72e4a89b2089b62c6ea86c4aef91755ee9ee9a
--- /dev/null
+++ b/diffsynth/models/flux_text_encoder_t5.py
@@ -0,0 +1,43 @@
+import torch
+from transformers import T5EncoderModel, T5Config
+
+
+class FluxTextEncoderT5(T5EncoderModel):
+ def __init__(self):
+ config = T5Config(**{
+ "architectures": [
+ "T5EncoderModel"
+ ],
+ "classifier_dropout": 0.0,
+ "d_ff": 10240,
+ "d_kv": 64,
+ "d_model": 4096,
+ "decoder_start_token_id": 0,
+ "dense_act_fn": "gelu_new",
+ "dropout_rate": 0.1,
+ "dtype": "bfloat16",
+ "eos_token_id": 1,
+ "feed_forward_proj": "gated-gelu",
+ "initializer_factor": 1.0,
+ "is_encoder_decoder": True,
+ "is_gated_act": True,
+ "layer_norm_epsilon": 1e-06,
+ "model_type": "t5",
+ "num_decoder_layers": 24,
+ "num_heads": 64,
+ "num_layers": 24,
+ "output_past": True,
+ "pad_token_id": 0,
+ "relative_attention_max_distance": 128,
+ "relative_attention_num_buckets": 32,
+ "tie_word_embeddings": False,
+ "transformers_version": "4.57.1",
+ "use_cache": True,
+ "vocab_size": 32128
+ })
+ super().__init__(config)
+
+ def forward(self, input_ids):
+ outputs = super().forward(input_ids=input_ids)
+ prompt_emb = outputs.last_hidden_state
+ return prompt_emb
diff --git a/diffsynth/models/flux_vae.py b/diffsynth/models/flux_vae.py
new file mode 100644
index 0000000000000000000000000000000000000000..5eabeaee6ad54f0b1c02f1b24cd2ccd84d0238bf
--- /dev/null
+++ b/diffsynth/models/flux_vae.py
@@ -0,0 +1,451 @@
+import torch
+from einops import rearrange, repeat
+
+
+class TileWorker:
+ def __init__(self):
+ pass
+
+
+ def mask(self, height, width, border_width):
+ # Create a mask with shape (height, width).
+ # The centre area is filled with 1, and the border line is filled with values in range (0, 1].
+ x = torch.arange(height).repeat(width, 1).T
+ y = torch.arange(width).repeat(height, 1)
+ mask = torch.stack([x + 1, height - x, y + 1, width - y]).min(dim=0).values
+ mask = (mask / border_width).clip(0, 1)
+ return mask
+
+
+ def tile(self, model_input, tile_size, tile_stride, tile_device, tile_dtype):
+ # Convert a tensor (b, c, h, w) to (b, c, tile_size, tile_size, tile_num)
+ batch_size, channel, _, _ = model_input.shape
+ model_input = model_input.to(device=tile_device, dtype=tile_dtype)
+ unfold_operator = torch.nn.Unfold(
+ kernel_size=(tile_size, tile_size),
+ stride=(tile_stride, tile_stride)
+ )
+ model_input = unfold_operator(model_input)
+ model_input = model_input.view((batch_size, channel, tile_size, tile_size, -1))
+
+ return model_input
+
+
+ def tiled_inference(self, forward_fn, model_input, tile_batch_size, inference_device, inference_dtype, tile_device, tile_dtype):
+ # Call y=forward_fn(x) for each tile
+ tile_num = model_input.shape[-1]
+ model_output_stack = []
+
+ for tile_id in range(0, tile_num, tile_batch_size):
+
+ # process input
+ tile_id_ = min(tile_id + tile_batch_size, tile_num)
+ x = model_input[:, :, :, :, tile_id: tile_id_]
+ x = x.to(device=inference_device, dtype=inference_dtype)
+ x = rearrange(x, "b c h w n -> (n b) c h w")
+
+ # process output
+ y = forward_fn(x)
+ y = rearrange(y, "(n b) c h w -> b c h w n", n=tile_id_-tile_id)
+ y = y.to(device=tile_device, dtype=tile_dtype)
+ model_output_stack.append(y)
+
+ model_output = torch.concat(model_output_stack, dim=-1)
+ return model_output
+
+
+ def io_scale(self, model_output, tile_size):
+ # Determine the size modification happened in forward_fn
+ # We only consider the same scale on height and width.
+ io_scale = model_output.shape[2] / tile_size
+ return io_scale
+
+
+ def untile(self, model_output, height, width, tile_size, tile_stride, border_width, tile_device, tile_dtype):
+ # The reversed function of tile
+ mask = self.mask(tile_size, tile_size, border_width)
+ mask = mask.to(device=tile_device, dtype=tile_dtype)
+ mask = rearrange(mask, "h w -> 1 1 h w 1")
+ model_output = model_output * mask
+
+ fold_operator = torch.nn.Fold(
+ output_size=(height, width),
+ kernel_size=(tile_size, tile_size),
+ stride=(tile_stride, tile_stride)
+ )
+ mask = repeat(mask[0, 0, :, :, 0], "h w -> 1 (h w) n", n=model_output.shape[-1])
+ model_output = rearrange(model_output, "b c h w n -> b (c h w) n")
+ model_output = fold_operator(model_output) / fold_operator(mask)
+
+ return model_output
+
+
+ def tiled_forward(self, forward_fn, model_input, tile_size, tile_stride, tile_batch_size=1, tile_device="cpu", tile_dtype=torch.float32, border_width=None):
+ # Prepare
+ inference_device, inference_dtype = model_input.device, model_input.dtype
+ height, width = model_input.shape[2], model_input.shape[3]
+ border_width = int(tile_stride*0.5) if border_width is None else border_width
+
+ # tile
+ model_input = self.tile(model_input, tile_size, tile_stride, tile_device, tile_dtype)
+
+ # inference
+ model_output = self.tiled_inference(forward_fn, model_input, tile_batch_size, inference_device, inference_dtype, tile_device, tile_dtype)
+
+ # resize
+ io_scale = self.io_scale(model_output, tile_size)
+ height, width = int(height*io_scale), int(width*io_scale)
+ tile_size, tile_stride = int(tile_size*io_scale), int(tile_stride*io_scale)
+ border_width = int(border_width*io_scale)
+
+ # untile
+ model_output = self.untile(model_output, height, width, tile_size, tile_stride, border_width, tile_device, tile_dtype)
+
+ # Done!
+ model_output = model_output.to(device=inference_device, dtype=inference_dtype)
+ return model_output
+
+
+class ConvAttention(torch.nn.Module):
+
+ def __init__(self, q_dim, num_heads, head_dim, kv_dim=None, bias_q=False, bias_kv=False, bias_out=False):
+ super().__init__()
+ dim_inner = head_dim * num_heads
+ kv_dim = kv_dim if kv_dim is not None else q_dim
+ self.num_heads = num_heads
+ self.head_dim = head_dim
+
+ self.to_q = torch.nn.Conv2d(q_dim, dim_inner, kernel_size=(1, 1), bias=bias_q)
+ self.to_k = torch.nn.Conv2d(kv_dim, dim_inner, kernel_size=(1, 1), bias=bias_kv)
+ self.to_v = torch.nn.Conv2d(kv_dim, dim_inner, kernel_size=(1, 1), bias=bias_kv)
+ self.to_out = torch.nn.Conv2d(dim_inner, q_dim, kernel_size=(1, 1), bias=bias_out)
+
+ def forward(self, hidden_states, encoder_hidden_states=None, attn_mask=None):
+ if encoder_hidden_states is None:
+ encoder_hidden_states = hidden_states
+
+ batch_size = encoder_hidden_states.shape[0]
+
+ conv_input = rearrange(hidden_states, "B L C -> B C L 1")
+ q = self.to_q(conv_input)
+ q = rearrange(q[:, :, :, 0], "B C L -> B L C")
+ conv_input = rearrange(encoder_hidden_states, "B L C -> B C L 1")
+ k = self.to_k(conv_input)
+ v = self.to_v(conv_input)
+ k = rearrange(k[:, :, :, 0], "B C L -> B L C")
+ v = rearrange(v[:, :, :, 0], "B C L -> B L C")
+
+ q = q.view(batch_size, -1, self.num_heads, self.head_dim).transpose(1, 2)
+ k = k.view(batch_size, -1, self.num_heads, self.head_dim).transpose(1, 2)
+ v = v.view(batch_size, -1, self.num_heads, self.head_dim).transpose(1, 2)
+
+ hidden_states = torch.nn.functional.scaled_dot_product_attention(q, k, v, attn_mask=attn_mask)
+ hidden_states = hidden_states.transpose(1, 2).reshape(batch_size, -1, self.num_heads * self.head_dim)
+ hidden_states = hidden_states.to(q.dtype)
+
+ conv_input = rearrange(hidden_states, "B L C -> B C L 1")
+ hidden_states = self.to_out(conv_input)
+ hidden_states = rearrange(hidden_states[:, :, :, 0], "B C L -> B L C")
+
+ return hidden_states
+
+
+class Attention(torch.nn.Module):
+
+ def __init__(self, q_dim, num_heads, head_dim, kv_dim=None, bias_q=False, bias_kv=False, bias_out=False):
+ super().__init__()
+ dim_inner = head_dim * num_heads
+ kv_dim = kv_dim if kv_dim is not None else q_dim
+ self.num_heads = num_heads
+ self.head_dim = head_dim
+
+ self.to_q = torch.nn.Linear(q_dim, dim_inner, bias=bias_q)
+ self.to_k = torch.nn.Linear(kv_dim, dim_inner, bias=bias_kv)
+ self.to_v = torch.nn.Linear(kv_dim, dim_inner, bias=bias_kv)
+ self.to_out = torch.nn.Linear(dim_inner, q_dim, bias=bias_out)
+
+ def forward(self, hidden_states, encoder_hidden_states=None, attn_mask=None):
+ if encoder_hidden_states is None:
+ encoder_hidden_states = hidden_states
+
+ batch_size = encoder_hidden_states.shape[0]
+
+ q = self.to_q(hidden_states)
+ k = self.to_k(encoder_hidden_states)
+ v = self.to_v(encoder_hidden_states)
+
+ q = q.view(batch_size, -1, self.num_heads, self.head_dim).transpose(1, 2)
+ k = k.view(batch_size, -1, self.num_heads, self.head_dim).transpose(1, 2)
+ v = v.view(batch_size, -1, self.num_heads, self.head_dim).transpose(1, 2)
+
+ hidden_states = torch.nn.functional.scaled_dot_product_attention(q, k, v, attn_mask=attn_mask)
+ hidden_states = hidden_states.transpose(1, 2).reshape(batch_size, -1, self.num_heads * self.head_dim)
+ hidden_states = hidden_states.to(q.dtype)
+
+ hidden_states = self.to_out(hidden_states)
+
+ return hidden_states
+
+
+class VAEAttentionBlock(torch.nn.Module):
+
+ def __init__(self, num_attention_heads, attention_head_dim, in_channels, num_layers=1, norm_num_groups=32, eps=1e-5, use_conv_attention=True):
+ super().__init__()
+ inner_dim = num_attention_heads * attention_head_dim
+
+ self.norm = torch.nn.GroupNorm(num_groups=norm_num_groups, num_channels=in_channels, eps=eps, affine=True)
+
+ if use_conv_attention:
+ self.transformer_blocks = torch.nn.ModuleList([
+ ConvAttention(
+ inner_dim,
+ num_attention_heads,
+ attention_head_dim,
+ bias_q=True,
+ bias_kv=True,
+ bias_out=True
+ )
+ for d in range(num_layers)
+ ])
+ else:
+ self.transformer_blocks = torch.nn.ModuleList([
+ Attention(
+ inner_dim,
+ num_attention_heads,
+ attention_head_dim,
+ bias_q=True,
+ bias_kv=True,
+ bias_out=True
+ )
+ for d in range(num_layers)
+ ])
+
+ def forward(self, hidden_states, time_emb, text_emb, res_stack):
+ batch, _, height, width = hidden_states.shape
+ residual = hidden_states
+
+ hidden_states = self.norm(hidden_states)
+ inner_dim = hidden_states.shape[1]
+ hidden_states = hidden_states.permute(0, 2, 3, 1).reshape(batch, height * width, inner_dim)
+
+ for block in self.transformer_blocks:
+ hidden_states = block(hidden_states)
+
+ hidden_states = hidden_states.reshape(batch, height, width, inner_dim).permute(0, 3, 1, 2).contiguous()
+ hidden_states = hidden_states + residual
+
+ return hidden_states, time_emb, text_emb, res_stack
+
+
+class ResnetBlock(torch.nn.Module):
+ def __init__(self, in_channels, out_channels, temb_channels=None, groups=32, eps=1e-5):
+ super().__init__()
+ self.norm1 = torch.nn.GroupNorm(num_groups=groups, num_channels=in_channels, eps=eps, affine=True)
+ self.conv1 = torch.nn.Conv2d(in_channels, out_channels, kernel_size=3, stride=1, padding=1)
+ if temb_channels is not None:
+ self.time_emb_proj = torch.nn.Linear(temb_channels, out_channels)
+ self.norm2 = torch.nn.GroupNorm(num_groups=groups, num_channels=out_channels, eps=eps, affine=True)
+ self.conv2 = torch.nn.Conv2d(out_channels, out_channels, kernel_size=3, stride=1, padding=1)
+ self.nonlinearity = torch.nn.SiLU()
+ self.conv_shortcut = None
+ if in_channels != out_channels:
+ self.conv_shortcut = torch.nn.Conv2d(in_channels, out_channels, kernel_size=1, stride=1, padding=0, bias=True)
+
+ def forward(self, hidden_states, time_emb, text_emb, res_stack, **kwargs):
+ x = hidden_states
+ x = self.norm1(x)
+ x = self.nonlinearity(x)
+ x = self.conv1(x)
+ if time_emb is not None:
+ emb = self.nonlinearity(time_emb)
+ emb = self.time_emb_proj(emb)[:, :, None, None]
+ x = x + emb
+ x = self.norm2(x)
+ x = self.nonlinearity(x)
+ x = self.conv2(x)
+ if self.conv_shortcut is not None:
+ hidden_states = self.conv_shortcut(hidden_states)
+ hidden_states = hidden_states + x
+ return hidden_states, time_emb, text_emb, res_stack
+
+
+class UpSampler(torch.nn.Module):
+ def __init__(self, channels):
+ super().__init__()
+ self.conv = torch.nn.Conv2d(channels, channels, 3, padding=1)
+
+ def forward(self, hidden_states, time_emb, text_emb, res_stack, **kwargs):
+ hidden_states = torch.nn.functional.interpolate(hidden_states, scale_factor=2.0, mode="nearest")
+ hidden_states = self.conv(hidden_states)
+ return hidden_states, time_emb, text_emb, res_stack
+
+
+class DownSampler(torch.nn.Module):
+ def __init__(self, channels, padding=1, extra_padding=False):
+ super().__init__()
+ self.conv = torch.nn.Conv2d(channels, channels, 3, stride=2, padding=padding)
+ self.extra_padding = extra_padding
+
+ def forward(self, hidden_states, time_emb, text_emb, res_stack, **kwargs):
+ if self.extra_padding:
+ hidden_states = torch.nn.functional.pad(hidden_states, (0, 1, 0, 1), mode="constant", value=0)
+ hidden_states = self.conv(hidden_states)
+ return hidden_states, time_emb, text_emb, res_stack
+
+
+class FluxVAEDecoder(torch.nn.Module):
+ def __init__(self, use_conv_attention=True):
+ super().__init__()
+ self.scaling_factor = 0.3611
+ self.shift_factor = 0.1159
+ self.conv_in = torch.nn.Conv2d(16, 512, kernel_size=3, padding=1) # Different from SD 1.x
+
+ self.blocks = torch.nn.ModuleList([
+ # UNetMidBlock2D
+ ResnetBlock(512, 512, eps=1e-6),
+ VAEAttentionBlock(1, 512, 512, 1, eps=1e-6, use_conv_attention=use_conv_attention),
+ ResnetBlock(512, 512, eps=1e-6),
+ # UpDecoderBlock2D
+ ResnetBlock(512, 512, eps=1e-6),
+ ResnetBlock(512, 512, eps=1e-6),
+ ResnetBlock(512, 512, eps=1e-6),
+ UpSampler(512),
+ # UpDecoderBlock2D
+ ResnetBlock(512, 512, eps=1e-6),
+ ResnetBlock(512, 512, eps=1e-6),
+ ResnetBlock(512, 512, eps=1e-6),
+ UpSampler(512),
+ # UpDecoderBlock2D
+ ResnetBlock(512, 256, eps=1e-6),
+ ResnetBlock(256, 256, eps=1e-6),
+ ResnetBlock(256, 256, eps=1e-6),
+ UpSampler(256),
+ # UpDecoderBlock2D
+ ResnetBlock(256, 128, eps=1e-6),
+ ResnetBlock(128, 128, eps=1e-6),
+ ResnetBlock(128, 128, eps=1e-6),
+ ])
+
+ self.conv_norm_out = torch.nn.GroupNorm(num_channels=128, num_groups=32, eps=1e-6)
+ self.conv_act = torch.nn.SiLU()
+ self.conv_out = torch.nn.Conv2d(128, 3, kernel_size=3, padding=1)
+
+ def tiled_forward(self, sample, tile_size=64, tile_stride=32):
+ hidden_states = TileWorker().tiled_forward(
+ lambda x: self.forward(x),
+ sample,
+ tile_size,
+ tile_stride,
+ tile_device=sample.device,
+ tile_dtype=sample.dtype
+ )
+ return hidden_states
+
+ def forward(self, sample, tiled=False, tile_size=64, tile_stride=32, **kwargs):
+ # For VAE Decoder, we do not need to apply the tiler on each layer.
+ if tiled:
+ return self.tiled_forward(sample, tile_size=tile_size, tile_stride=tile_stride)
+
+ # 1. pre-process
+ hidden_states = sample / self.scaling_factor + self.shift_factor
+ hidden_states = self.conv_in(hidden_states)
+ time_emb = None
+ text_emb = None
+ res_stack = None
+
+ # 2. blocks
+ for i, block in enumerate(self.blocks):
+ hidden_states, time_emb, text_emb, res_stack = block(hidden_states, time_emb, text_emb, res_stack)
+
+ # 3. output
+ hidden_states = self.conv_norm_out(hidden_states)
+ hidden_states = self.conv_act(hidden_states)
+ hidden_states = self.conv_out(hidden_states)
+
+ return hidden_states
+
+
+class FluxVAEEncoder(torch.nn.Module):
+ def __init__(self, use_conv_attention=True):
+ super().__init__()
+ self.scaling_factor = 0.3611
+ self.shift_factor = 0.1159
+ self.conv_in = torch.nn.Conv2d(3, 128, kernel_size=3, padding=1)
+
+ self.blocks = torch.nn.ModuleList([
+ # DownEncoderBlock2D
+ ResnetBlock(128, 128, eps=1e-6),
+ ResnetBlock(128, 128, eps=1e-6),
+ DownSampler(128, padding=0, extra_padding=True),
+ # DownEncoderBlock2D
+ ResnetBlock(128, 256, eps=1e-6),
+ ResnetBlock(256, 256, eps=1e-6),
+ DownSampler(256, padding=0, extra_padding=True),
+ # DownEncoderBlock2D
+ ResnetBlock(256, 512, eps=1e-6),
+ ResnetBlock(512, 512, eps=1e-6),
+ DownSampler(512, padding=0, extra_padding=True),
+ # DownEncoderBlock2D
+ ResnetBlock(512, 512, eps=1e-6),
+ ResnetBlock(512, 512, eps=1e-6),
+ # UNetMidBlock2D
+ ResnetBlock(512, 512, eps=1e-6),
+ VAEAttentionBlock(1, 512, 512, 1, eps=1e-6, use_conv_attention=use_conv_attention),
+ ResnetBlock(512, 512, eps=1e-6),
+ ])
+
+ self.conv_norm_out = torch.nn.GroupNorm(num_channels=512, num_groups=32, eps=1e-6)
+ self.conv_act = torch.nn.SiLU()
+ self.conv_out = torch.nn.Conv2d(512, 32, kernel_size=3, padding=1)
+
+ def tiled_forward(self, sample, tile_size=64, tile_stride=32):
+ hidden_states = TileWorker().tiled_forward(
+ lambda x: self.forward(x),
+ sample,
+ tile_size,
+ tile_stride,
+ tile_device=sample.device,
+ tile_dtype=sample.dtype
+ )
+ return hidden_states
+
+ def forward(self, sample, tiled=False, tile_size=64, tile_stride=32, **kwargs):
+ # For VAE Decoder, we do not need to apply the tiler on each layer.
+ if tiled:
+ return self.tiled_forward(sample, tile_size=tile_size, tile_stride=tile_stride)
+
+ # 1. pre-process
+ hidden_states = self.conv_in(sample)
+ time_emb = None
+ text_emb = None
+ res_stack = None
+
+ # 2. blocks
+ for i, block in enumerate(self.blocks):
+ hidden_states, time_emb, text_emb, res_stack = block(hidden_states, time_emb, text_emb, res_stack)
+
+ # 3. output
+ hidden_states = self.conv_norm_out(hidden_states)
+ hidden_states = self.conv_act(hidden_states)
+ hidden_states = self.conv_out(hidden_states)
+ hidden_states = hidden_states[:, :16]
+ hidden_states = (hidden_states - self.shift_factor) * self.scaling_factor
+
+ return hidden_states
+
+ def encode_video(self, sample, batch_size=8):
+ B = sample.shape[0]
+ hidden_states = []
+
+ for i in range(0, sample.shape[2], batch_size):
+
+ j = min(i + batch_size, sample.shape[2])
+ sample_batch = rearrange(sample[:,:,i:j], "B C T H W -> (B T) C H W")
+
+ hidden_states_batch = self(sample_batch)
+ hidden_states_batch = rearrange(hidden_states_batch, "(B T) C H W -> B C T H W", B=B)
+
+ hidden_states.append(hidden_states_batch)
+
+ hidden_states = torch.concat(hidden_states, dim=2)
+ return hidden_states
diff --git a/diffsynth/models/flux_value_control.py b/diffsynth/models/flux_value_control.py
new file mode 100644
index 0000000000000000000000000000000000000000..549dbc93b41343a42266af11584e2e7d39a17cd6
--- /dev/null
+++ b/diffsynth/models/flux_value_control.py
@@ -0,0 +1,56 @@
+import torch
+from .general_modules import TemporalTimesteps
+
+
+class MultiValueEncoder(torch.nn.Module):
+ def __init__(self, encoders=()):
+ super().__init__()
+ if not isinstance(encoders, list):
+ encoders = [encoders]
+ self.encoders = torch.nn.ModuleList(encoders)
+
+ def __call__(self, values, dtype):
+ emb = []
+ for encoder, value in zip(self.encoders, values):
+ if value is not None:
+ value = value.unsqueeze(0)
+ emb.append(encoder(value, dtype))
+ emb = torch.concat(emb, dim=0)
+ return emb
+
+
+class SingleValueEncoder(torch.nn.Module):
+ def __init__(self, dim_in=256, dim_out=4096, prefer_len=32, computation_device=None):
+ super().__init__()
+ self.prefer_len = prefer_len
+ self.prefer_proj = TemporalTimesteps(num_channels=dim_in, flip_sin_to_cos=True, downscale_freq_shift=0, computation_device=computation_device)
+ self.prefer_value_embedder = torch.nn.Sequential(
+ torch.nn.Linear(dim_in, dim_out), torch.nn.SiLU(), torch.nn.Linear(dim_out, dim_out)
+ )
+ self.positional_embedding = torch.nn.Parameter(
+ torch.randn(self.prefer_len, dim_out)
+ )
+
+ def forward(self, value, dtype):
+ value = value * 1000
+ emb = self.prefer_proj(value).to(dtype)
+ emb = self.prefer_value_embedder(emb).squeeze(0)
+ base_embeddings = emb.expand(self.prefer_len, -1)
+ positional_embedding = self.positional_embedding.to(dtype=base_embeddings.dtype, device=base_embeddings.device)
+ learned_embeddings = base_embeddings + positional_embedding
+ return learned_embeddings
+
+ @staticmethod
+ def state_dict_converter():
+ return SingleValueEncoderStateDictConverter()
+
+
+class SingleValueEncoderStateDictConverter:
+ def __init__(self):
+ pass
+
+ def from_diffusers(self, state_dict):
+ return state_dict
+
+ def from_civitai(self, state_dict):
+ return state_dict
diff --git a/diffsynth/models/general_modules.py b/diffsynth/models/general_modules.py
new file mode 100644
index 0000000000000000000000000000000000000000..1e97ba6e80447ab9f10eed9424edbd6e3d147cb4
--- /dev/null
+++ b/diffsynth/models/general_modules.py
@@ -0,0 +1,146 @@
+import torch, math
+
+
+def get_timestep_embedding(
+ timesteps: torch.Tensor,
+ embedding_dim: int,
+ flip_sin_to_cos: bool = False,
+ downscale_freq_shift: float = 1,
+ scale: float = 1,
+ max_period: int = 10000,
+ computation_device = None,
+ align_dtype_to_timestep = False,
+):
+ assert len(timesteps.shape) == 1, "Timesteps should be a 1d-array"
+
+ half_dim = embedding_dim // 2
+ exponent = -math.log(max_period) * torch.arange(
+ start=0, end=half_dim, dtype=torch.float32, device=timesteps.device if computation_device is None else computation_device
+ )
+ exponent = exponent / (half_dim - downscale_freq_shift)
+
+ emb = torch.exp(exponent)
+ if align_dtype_to_timestep:
+ emb = emb.to(timesteps.dtype)
+ emb = timesteps[:, None].float() * emb[None, :]
+
+ # scale embeddings
+ emb = scale * emb
+
+ # concat sine and cosine embeddings
+ emb = torch.cat([torch.sin(emb), torch.cos(emb)], dim=-1)
+
+ # flip sine and cosine embeddings
+ if flip_sin_to_cos:
+ emb = torch.cat([emb[:, half_dim:], emb[:, :half_dim]], dim=-1)
+
+ # zero pad
+ if embedding_dim % 2 == 1:
+ emb = torch.nn.functional.pad(emb, (0, 1, 0, 0))
+ return emb
+
+
+class TemporalTimesteps(torch.nn.Module):
+ def __init__(self, num_channels: int, flip_sin_to_cos: bool, downscale_freq_shift: float, computation_device = None, scale=1, align_dtype_to_timestep=False):
+ super().__init__()
+ self.num_channels = num_channels
+ self.flip_sin_to_cos = flip_sin_to_cos
+ self.downscale_freq_shift = downscale_freq_shift
+ self.computation_device = computation_device
+ self.scale = scale
+ self.align_dtype_to_timestep = align_dtype_to_timestep
+
+ def forward(self, timesteps):
+ t_emb = get_timestep_embedding(
+ timesteps,
+ self.num_channels,
+ flip_sin_to_cos=self.flip_sin_to_cos,
+ downscale_freq_shift=self.downscale_freq_shift,
+ computation_device=self.computation_device,
+ scale=self.scale,
+ align_dtype_to_timestep=self.align_dtype_to_timestep,
+ )
+ return t_emb
+
+
+class DiffusersCompatibleTimestepProj(torch.nn.Module):
+ def __init__(self, dim_in, dim_out):
+ super().__init__()
+ self.linear_1 = torch.nn.Linear(dim_in, dim_out)
+ self.act = torch.nn.SiLU()
+ self.linear_2 = torch.nn.Linear(dim_out, dim_out)
+
+ def forward(self, x):
+ x = self.linear_1(x)
+ x = self.act(x)
+ x = self.linear_2(x)
+ return x
+
+
+class TimestepEmbeddings(torch.nn.Module):
+ def __init__(self, dim_in, dim_out, computation_device=None, diffusers_compatible_format=False, scale=1, align_dtype_to_timestep=False, use_additional_t_cond=False):
+ super().__init__()
+ self.time_proj = TemporalTimesteps(num_channels=dim_in, flip_sin_to_cos=True, downscale_freq_shift=0, computation_device=computation_device, scale=scale, align_dtype_to_timestep=align_dtype_to_timestep)
+ if diffusers_compatible_format:
+ self.timestep_embedder = DiffusersCompatibleTimestepProj(dim_in, dim_out)
+ else:
+ self.timestep_embedder = torch.nn.Sequential(
+ torch.nn.Linear(dim_in, dim_out), torch.nn.SiLU(), torch.nn.Linear(dim_out, dim_out)
+ )
+ self.use_additional_t_cond = use_additional_t_cond
+ if use_additional_t_cond:
+ self.addition_t_embedding = torch.nn.Embedding(2, dim_out)
+
+ def forward(self, timestep, dtype, addition_t_cond=None):
+ time_emb = self.time_proj(timestep).to(dtype)
+ time_emb = self.timestep_embedder(time_emb)
+ if addition_t_cond is not None:
+ addition_t_emb = self.addition_t_embedding(addition_t_cond)
+ addition_t_emb = addition_t_emb.to(dtype=dtype)
+ time_emb = time_emb + addition_t_emb
+ return time_emb
+
+
+class RMSNorm(torch.nn.Module):
+ def __init__(self, dim, eps, elementwise_affine=True):
+ super().__init__()
+ self.eps = eps
+ if elementwise_affine:
+ self.weight = torch.nn.Parameter(torch.ones((dim,)))
+ else:
+ self.weight = None
+
+ def forward(self, hidden_states):
+ input_dtype = hidden_states.dtype
+ variance = hidden_states.to(torch.float32).square().mean(-1, keepdim=True)
+ hidden_states = hidden_states * torch.rsqrt(variance + self.eps)
+ hidden_states = hidden_states.to(input_dtype)
+ if self.weight is not None:
+ hidden_states = hidden_states * self.weight
+ return hidden_states
+
+
+class AdaLayerNorm(torch.nn.Module):
+ def __init__(self, dim, single=False, dual=False):
+ super().__init__()
+ self.single = single
+ self.dual = dual
+ self.linear = torch.nn.Linear(dim, dim * [[6, 2][single], 9][dual])
+ self.norm = torch.nn.LayerNorm(dim, elementwise_affine=False, eps=1e-6)
+
+ def forward(self, x, emb):
+ emb = self.linear(torch.nn.functional.silu(emb))
+ if self.single:
+ scale, shift = emb.unsqueeze(1).chunk(2, dim=2)
+ x = self.norm(x) * (1 + scale) + shift
+ return x
+ elif self.dual:
+ shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp, shift_msa2, scale_msa2, gate_msa2 = emb.unsqueeze(1).chunk(9, dim=2)
+ norm_x = self.norm(x)
+ x = norm_x * (1 + scale_msa) + shift_msa
+ norm_x2 = norm_x * (1 + scale_msa2) + shift_msa2
+ return x, gate_msa, shift_mlp, scale_mlp, gate_mlp, norm_x2, gate_msa2
+ else:
+ shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = emb.unsqueeze(1).chunk(6, dim=2)
+ x = self.norm(x) * (1 + scale_msa) + shift_msa
+ return x, gate_msa, shift_mlp, scale_mlp, gate_mlp
diff --git a/diffsynth/models/hidream_common.py b/diffsynth/models/hidream_common.py
new file mode 100644
index 0000000000000000000000000000000000000000..e4d1945df3090110b3d8447fd99cf6b32d54a1a0
--- /dev/null
+++ b/diffsynth/models/hidream_common.py
@@ -0,0 +1,373 @@
+import torch, math
+from typing import Optional, Any, Tuple, List, Sequence, Dict, Iterable
+from PIL import Image, ImageDraw
+import einops
+
+
+PATCH_SIZE = 32
+
+
+DEFAULT_COLORS = [
+ (255, 0, 0),
+ (0, 180, 0),
+ (0, 0, 255),
+ (204, 180, 0),
+ (255, 0, 255),
+ (0, 255, 255),
+ (128, 0, 0),
+ (0, 128, 0),
+ (0, 0, 128),
+ (128, 128, 0),
+]
+
+
+def add_special_tokens(tokenizer):
+ """Attach the special-token shortcuts that the pipeline relies on."""
+ tokenizer.boi_token = "<|boi_token|>"
+ tokenizer.bor_token = "<|bor_token|>"
+ tokenizer.eor_token = "<|eor_token|>"
+ tokenizer.bot_token = "<|bot_token|>"
+ tokenizer.tms_token = "<|tms_token|>"
+
+
+def unpatchify(x, height, width):
+ h_patches = height // PATCH_SIZE
+ w_patches = width // PATCH_SIZE
+ x = einops.rearrange(
+ x, 'B (H W) (C p1 p2) -> B C (H p1) (W p2)',
+ H=h_patches, W=w_patches, p1=PATCH_SIZE, p2=PATCH_SIZE,
+ )
+ return x
+
+
+def patchify(x):
+ x = einops.rearrange(
+ x, 'B C (H p1) (W p2) -> B (H W) (C p1 p2)',
+ p1=PATCH_SIZE, p2=PATCH_SIZE,
+ )
+ return x
+
+
+def resize_pilimage(pil_image, image_size, resampler=Image.BICUBIC):
+ patch_size = PATCH_SIZE
+ while min(*pil_image.size) >= 2 * image_size:
+ pil_image = pil_image.resize(
+ tuple(x // 2 for x in pil_image.size), resample=Image.BOX
+ )
+
+ m = patch_size
+ width, height = pil_image.width, pil_image.height
+ S_max = image_size * image_size
+ scale = S_max / (width * height)
+ scale = math.sqrt(scale)
+
+ new_sizes = [
+ (round(width * scale) // m * m, round(height * scale) // m * m),
+ (round(width * scale) // m * m, math.floor(height * scale) // m * m),
+ (math.floor(width * scale) // m * m, round(height * scale) // m * m),
+ (math.floor(width * scale) // m * m, math.floor(height * scale) // m * m),
+ ]
+ new_sizes = sorted(new_sizes, key=lambda x: x[0] * x[1], reverse=True)
+
+ for new_size in new_sizes:
+ if new_size[0] * new_size[1] <= S_max:
+ break
+
+ s1 = width / new_size[0]
+ s2 = height / new_size[1]
+ if s1 < s2:
+ pil_image = pil_image.resize([new_size[0], round(height / s1)], resample=resampler)
+ top = (round(height / s1) - new_size[1]) // 2
+ pil_image = pil_image.crop((0, top, new_size[0], top + new_size[1]))
+ else:
+ pil_image = pil_image.resize([round(width / s2), new_size[1]], resample=resampler)
+ left = (round(width / s2) - new_size[0]) // 2
+ pil_image = pil_image.crop((left, 0, left + new_size[0], new_size[1]))
+
+ return pil_image
+
+
+def calculate_dimensions(max_size, ratio):
+ width = math.sqrt(max_size * max_size * ratio)
+ height = width / ratio
+ width = int(width / 32) * 32
+ height = int(height / 32) * 32
+ return width, height
+
+
+def get_rope_index_fix_point(
+ spatial_merge_size,
+ image_token_id,
+ video_token_id,
+ vision_start_token_id,
+ input_ids: Optional[torch.LongTensor] = None,
+ image_grid_thw: Optional[torch.LongTensor] = None,
+ video_grid_thw: Optional[torch.LongTensor] = None,
+ attention_mask: Optional[torch.Tensor] = None,
+ skip_vision_start_token=None,
+ fix_point=4096,
+) -> tuple[torch.Tensor, torch.Tensor]:
+ if video_grid_thw is not None:
+ video_grid_thw = torch.repeat_interleave(video_grid_thw, video_grid_thw[:, 0], dim=0)
+ video_grid_thw[:, 0] = 1
+
+ mrope_position_deltas = []
+ if input_ids is not None and (image_grid_thw is not None or video_grid_thw is not None):
+ total_input_ids = input_ids
+ if attention_mask is None:
+ attention_mask = torch.ones_like(total_input_ids)
+ position_ids = torch.ones(
+ 3,
+ input_ids.shape[0],
+ input_ids.shape[1],
+ dtype=input_ids.dtype,
+ device=input_ids.device,
+ )
+ image_index, video_index = 0, 0
+ attention_mask = attention_mask.to(total_input_ids.device)
+ for i, input_ids in enumerate(total_input_ids):
+ input_ids = input_ids[attention_mask[i] == 1]
+ image_nums, video_nums = 0, 0
+ vision_start_indices = torch.argwhere(input_ids == vision_start_token_id).squeeze(1)
+ vision_tokens = input_ids[vision_start_indices + 1]
+ image_nums = (vision_tokens == image_token_id).sum()
+ video_nums = (vision_tokens == video_token_id).sum()
+ input_tokens = input_ids.tolist()
+ llm_pos_ids_list: list = []
+ st = 0
+ remain_images, remain_videos = image_nums, video_nums
+ for _ in range(image_nums + video_nums):
+ if image_token_id in input_tokens and remain_images > 0:
+ ed_image = input_tokens.index(image_token_id, st)
+ else:
+ ed_image = len(input_tokens) + 1
+ if video_token_id in input_tokens and remain_videos > 0:
+ ed_video = input_tokens.index(video_token_id, st)
+ else:
+ ed_video = len(input_tokens) + 1
+ if ed_image < ed_video:
+ t, h, w = (
+ image_grid_thw[image_index][0],
+ image_grid_thw[image_index][1],
+ image_grid_thw[image_index][2],
+ )
+ image_index += 1
+ remain_images -= 1
+ ed = ed_image
+ else:
+ t, h, w = (
+ video_grid_thw[video_index][0],
+ video_grid_thw[video_index][1],
+ video_grid_thw[video_index][2],
+ )
+ video_index += 1
+ remain_videos -= 1
+ ed = ed_video
+ llm_grid_t, llm_grid_h, llm_grid_w = (
+ t.item(),
+ h.item() // spatial_merge_size,
+ w.item() // spatial_merge_size,
+ )
+ text_len = ed - st
+
+ text_len -= skip_vision_start_token[image_index - 1]
+ text_len = max(0, text_len)
+
+ st_idx = llm_pos_ids_list[-1].max() + 1 if len(llm_pos_ids_list) > 0 else 0
+ llm_pos_ids_list.append(torch.arange(text_len).view(1, -1).expand(3, -1) + st_idx)
+
+ t_index = torch.arange(llm_grid_t).view(-1, 1).expand(-1, llm_grid_h * llm_grid_w).flatten()
+ h_index = torch.arange(llm_grid_h).view(1, -1, 1).expand(llm_grid_t, -1, llm_grid_w).flatten()
+ w_index = torch.arange(llm_grid_w).view(1, 1, -1).expand(llm_grid_t, llm_grid_h, -1).flatten()
+
+ if skip_vision_start_token[image_index - 1]:
+ if fix_point > 0:
+ fix_point = fix_point - st_idx
+ llm_pos_ids_list.append(torch.stack([t_index, h_index, w_index]) + fix_point + st_idx)
+ fix_point = 0
+ else:
+ llm_pos_ids_list.append(torch.stack([t_index, h_index, w_index]) + text_len + st_idx)
+ st = ed + llm_grid_t * llm_grid_h * llm_grid_w
+
+ if st < len(input_tokens):
+ st_idx = llm_pos_ids_list[-1].max() + 1 if len(llm_pos_ids_list) > 0 else 0
+ text_len = len(input_tokens) - st
+ llm_pos_ids_list.append(torch.arange(text_len).view(1, -1).expand(3, -1) + st_idx)
+
+ llm_positions = torch.cat(llm_pos_ids_list, dim=1).reshape(3, -1)
+ position_ids[..., i, attention_mask[i] == 1] = llm_positions.to(position_ids.device)
+ mrope_position_deltas.append(llm_positions.max() + 1 - len(total_input_ids[i]))
+ mrope_position_deltas = torch.tensor(mrope_position_deltas, device=input_ids.device).unsqueeze(1)
+ return position_ids, mrope_position_deltas
+ else:
+ if attention_mask is not None:
+ position_ids = attention_mask.long().cumsum(-1) - 1
+ position_ids.masked_fill_(attention_mask == 0, 1)
+ position_ids = position_ids.unsqueeze(0).expand(3, -1, -1).to(attention_mask.device)
+ max_position_ids = position_ids.max(0, keepdim=False)[0].max(-1, keepdim=True)[0]
+ mrope_position_deltas = max_position_ids + 1 - attention_mask.shape[-1]
+ else:
+ position_ids = (
+ torch.arange(input_ids.shape[1], device=input_ids.device)
+ .view(1, 1, -1)
+ .expand(3, input_ids.shape[0], -1)
+ )
+ mrope_position_deltas = torch.zeros(
+ [input_ids.shape[0], 1],
+ device=input_ids.device,
+ dtype=input_ids.dtype,
+ )
+ return position_ids, mrope_position_deltas
+
+
+def _unwrap_boxes(data: Any) -> Any:
+ if isinstance(data, dict):
+ for key in ("layout_bboxes", "bboxes", "boxes", "bbox_list"):
+ if key in data:
+ return data[key]
+ return data
+
+
+def _as_bbox_and_text(item: Any) -> Tuple[Sequence[float], str]:
+ if isinstance(item, dict):
+ bbox = item.get("bbox") or item.get("box")
+ text = str(item.get("text") or item.get("label") or "")
+ if bbox is None:
+ raise ValueError(f"Missing bbox in layout item: {item!r}")
+ return bbox, text
+ if isinstance(item, (list, tuple)) and len(item) == 4:
+ return item, ""
+ raise ValueError(f"Unsupported layout bbox item: {item!r}")
+
+
+def _xxyy_relative_to_absolute_bbox(bbox: Sequence[float], width: int, height: int) -> List[int]:
+ if len(bbox) != 4:
+ raise ValueError(f"Expected bbox with 4 values, got: {bbox!r}")
+ x1, x2, y1, y2 = [float(v) for v in bbox]
+
+ # Inference layout input is xxyy relative coordinates: [x1, x2, y1, y2].
+ # Values in [0, 1] are the intended format. Keep 0-100 support for convenience.
+ max_abs = max(abs(x1), abs(y1), abs(x2), abs(y2))
+ if max_abs <= 1.0:
+ x1, x2 = x1 * width, x2 * width
+ y1, y2 = y1 * height, y2 * height
+ elif max_abs <= 100.0:
+ x1, x2 = x1 / 100.0 * width, x2 / 100.0 * width
+ y1, y2 = y1 / 100.0 * height, y2 / 100.0 * height
+
+ x1, x2 = sorted((x1, x2))
+ y1, y2 = sorted((y1, y2))
+ x1 = max(0, min(width - 1, int(round(x1))))
+ y1 = max(0, min(height - 1, int(round(y1))))
+ x2 = max(0, min(width - 1, int(round(x2))))
+ y2 = max(0, min(height - 1, int(round(y2))))
+ if x2 <= x1 or y2 <= y1:
+ raise ValueError(f"Invalid bbox after scaling/clamping: {[x1, y1, x2, y2]!r}")
+ return [x1, y1, x2, y2]
+
+
+def parse_layout_bboxes(layout_bboxes: Any, width: int, height: int) -> List[Dict[str, Any]]:
+ """Convert xxyy relative layout boxes into the training-side bbox layout format."""
+ raw_boxes = _unwrap_boxes(layout_bboxes)
+ if not isinstance(raw_boxes, list):
+ raise ValueError("layout_bboxes must be a list, or a dict containing one of: layout_bboxes/bboxes/boxes")
+
+ parsed = []
+ for idx, item in enumerate(raw_boxes):
+ bbox, text = _as_bbox_and_text(item)
+ parsed.append({
+ "bbox": _xxyy_relative_to_absolute_bbox(bbox, width, height),
+ "color": "",
+ "text": text,
+ "image": None,
+ "_orig_idx": idx,
+ })
+ return parsed
+
+
+def _bbox_area(item: Dict[str, Any]) -> int:
+ x1, y1, x2, y2 = item["bbox"]
+ return max(0, x2 - x1) * max(0, y2 - y1)
+
+
+def get_render_params(image_width: int, image_height: int) -> Tuple[int, int]:
+ edge = math.sqrt(image_width * image_height)
+ max_font_size = int(edge * 0.07)
+ max_bbox_line_width = int(edge * 0.05)
+ return max_font_size, max_bbox_line_width
+
+
+def draw_bbox_layout(
+ bbox_list: List[Dict[str, Any]],
+ image_width: int,
+ image_height: int,
+ max_bbox: int = 5,
+ max_bbox_line_width: int | None = None,
+ bbox_line_gap: int | None = None,
+ return_color: bool = False,
+):
+ """Draw a black layout image with colored boxes, matching the training-side layout style."""
+ if max_bbox_line_width is None:
+ _, max_bbox_line_width = get_render_params(image_width, image_height)
+ if bbox_line_gap is None:
+ bbox_line_gap = max(1, max_bbox_line_width // max_bbox)
+
+ image = Image.new("RGB", (image_width, image_height), (0, 0, 0))
+ draw = ImageDraw.Draw(image)
+ color_list = [None] * len(bbox_list)
+ sorted_bboxes = sorted(bbox_list, key=_bbox_area, reverse=True)[:max_bbox]
+
+ for sorted_idx, item in enumerate(sorted_bboxes):
+ color = DEFAULT_COLORS[sorted_idx % len(DEFAULT_COLORS)]
+ orig_idx = int(item.get("_orig_idx", sorted_idx))
+ if 0 <= orig_idx < len(color_list):
+ color_list[orig_idx] = color
+ line_width = max(max_bbox_line_width - sorted_idx * bbox_line_gap, 5)
+ draw.rectangle([int(v) for v in item["bbox"]], outline=color, width=line_width)
+
+ if return_color:
+ return image, color_list
+ return image
+
+
+def add_outer_border_keep_size(pil: Image.Image, color: Iterable[int], width: int) -> Image.Image:
+ """Draw a border inside the image without changing its size."""
+ img = pil.convert("RGB").copy()
+ color_tuple = tuple(int(c) for c in color)
+ width = max(0, int(width))
+ if width == 0:
+ return img
+
+ draw = ImageDraw.Draw(img)
+ w, h = img.size
+ for t in range(width):
+ draw.rectangle([t, t, w - 1 - t, h - 1 - t], outline=color_tuple)
+ return img
+
+
+def create_layout_reference_images(
+ ref_pils: List[Image.Image],
+ layout_bboxes: Any,
+ image_width: int,
+ image_height: int,
+ ref_max_size: int | None = None,
+) -> Tuple[List[str], str]:
+ """Create bordered ref images plus one layout image; returns paths to pass as ref_images."""
+ parsed_boxes = parse_layout_bboxes(layout_bboxes, image_width, image_height)
+ layout_image, color_list = draw_bbox_layout(
+ parsed_boxes,
+ image_width=image_width,
+ image_height=image_height,
+ return_color=True,
+ )
+
+ output_refs = []
+ for idx, ref in enumerate(ref_pils):
+ if ref_max_size is not None:
+ ref = resize_pilimage(ref, ref_max_size)
+ color = color_list[idx] if idx < len(color_list) and color_list[idx] is not None else DEFAULT_COLORS[idx % len(DEFAULT_COLORS)]
+ line_width = int(math.sqrt(ref.width * ref.height) * 0.04)
+ bordered = add_outer_border_keep_size(ref, color, line_width)
+ output_refs.append(bordered)
+ output_refs.append(layout_image)
+ return output_refs
diff --git a/diffsynth/models/hidream_o1_image_dit.py b/diffsynth/models/hidream_o1_image_dit.py
new file mode 100644
index 0000000000000000000000000000000000000000..b42e391a80838b6844a86a508979b6ef9b17de98
--- /dev/null
+++ b/diffsynth/models/hidream_o1_image_dit.py
@@ -0,0 +1,1910 @@
+# Copyright 2025 The Qwen Team and The HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+import functools
+import math
+from collections import OrderedDict
+from abc import ABC, abstractmethod
+import dataclasses
+from dataclasses import dataclass
+from typing import Any, Optional, Union
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch import Tensor
+
+from ..core.attention.attention import attention_forward
+from ..core import gradient_checkpoint_forward
+
+from transformers.modeling_utils import PreTrainedModel
+from transformers.configuration_utils import PretrainedConfig
+
+
+class ClassInstantier(OrderedDict):
+ def __getitem__(self, key):
+ content = super().__getitem__(key)
+ cls, kwargs = content if isinstance(content, tuple) else (content, {})
+ return cls(**kwargs)
+
+
+class SiLUActivation(nn.Module):
+
+ def forward(self, input: Tensor) -> Tensor:
+ return nn.functional.silu(input)
+
+
+class GELUTanh(nn.Module):
+ def __init__(self, use_gelu_tanh_python: bool = False):
+ super().__init__()
+ if use_gelu_tanh_python:
+ self.act = self._gelu_tanh_python
+ else:
+ self.act = functools.partial(nn.functional.gelu, approximate="tanh")
+
+ def _gelu_tanh_python(self, input: Tensor) -> Tensor:
+ return input * 0.5 * (1.0 + torch.tanh(math.sqrt(2.0 / math.pi) * (input + 0.044715 * torch.pow(input, 3.0))))
+
+ def forward(self, input: Tensor) -> Tensor:
+ return self.act(input)
+
+
+ACT2CLS = {
+ "gelu_pytorch_tanh": GELUTanh,
+ "silu": SiLUActivation,
+}
+ACT2FN = ClassInstantier(ACT2CLS)
+
+
+class CacheLayerMixin(ABC):
+ """Base, abstract class for a single layer's cache."""
+ is_compileable = False
+
+ def __init__(self):
+ self.keys: Optional[torch.Tensor] = None
+ self.values: Optional[torch.Tensor] = None
+ self.is_initialized = False
+
+ def __repr__(self):
+ return f"{self.__class__.__name__}"
+
+ @abstractmethod
+ def lazy_initialization(self, key_states: torch.Tensor): ...
+
+ @abstractmethod
+ def update(self, key_states: torch.Tensor, value_states: torch.Tensor, cache_kwargs: Optional[dict[str, Any]] = None) -> tuple[torch.Tensor, torch.Tensor]: ...
+
+ @abstractmethod
+ def get_mask_sizes(self, cache_position: torch.Tensor) -> tuple[int, int]: ...
+
+ @abstractmethod
+ def get_seq_length(self) -> int: ...
+
+ @abstractmethod
+ def get_max_cache_shape(self) -> int: ...
+
+ def offload(self):
+ if self.is_initialized:
+ self.keys = self.keys.to("cpu", non_blocking=True)
+ self.values = self.values.to("cpu", non_blocking=True)
+
+ def prefetch(self):
+ if self.is_initialized and self.keys.device != self.device:
+ self.keys = self.keys.to(self.device, non_blocking=True)
+ self.values = self.values.to(self.device, non_blocking=True)
+
+ def reset(self) -> None:
+ if self.is_initialized:
+ self.keys.zero_()
+ self.values.zero_()
+ if hasattr(self, "cumulative_length"):
+ self.cumulative_length = 0
+
+ def reorder_cache(self, beam_idx: torch.LongTensor) -> None:
+ if self.get_seq_length() > 0:
+ self.keys = self.keys.index_select(0, beam_idx.to(self.keys.device))
+ self.values = self.values.index_select(0, beam_idx.to(self.values.device))
+
+
+class DynamicLayer(CacheLayerMixin):
+ is_sliding = False
+
+ def lazy_initialization(self, key_states: torch.Tensor):
+ self.dtype, self.device = key_states.dtype, key_states.device
+ self.keys = torch.tensor([], dtype=self.dtype, device=self.device)
+ self.values = torch.tensor([], dtype=self.dtype, device=self.device)
+ self.is_initialized = True
+
+ def update(self, key_states: torch.Tensor, value_states: torch.Tensor, cache_kwargs: Optional[dict[str, Any]] = None) -> tuple[torch.Tensor, torch.Tensor]:
+ if not self.is_initialized:
+ self.lazy_initialization(key_states)
+ self.keys = torch.cat([self.keys, key_states], dim=-2)
+ self.values = torch.cat([self.values, value_states], dim=-2)
+ return self.keys, self.values
+
+ def get_mask_sizes(self, cache_position: torch.Tensor) -> tuple[int, int]:
+ kv_offset = 0
+ query_length = cache_position.shape[0]
+ kv_length = self.get_seq_length() + query_length
+ return kv_length, kv_offset
+
+ def get_seq_length(self) -> int:
+ if not self.is_initialized or self.keys.numel() == 0:
+ return 0
+ return self.keys.shape[-2]
+
+ def get_max_cache_shape(self) -> int:
+ return -1
+
+ def crop(self, max_length: int) -> None:
+ if max_length < 0:
+ max_length = self.get_seq_length() - abs(max_length)
+ if self.get_seq_length() <= max_length:
+ return
+ self.keys = self.keys[..., :max_length, :]
+ self.values = self.values[..., :max_length, :]
+
+ def batch_repeat_interleave(self, repeats: int) -> None:
+ if self.get_seq_length() > 0:
+ self.keys = self.keys.repeat_interleave(repeats, dim=0)
+ self.values = self.values.repeat_interleave(repeats, dim=0)
+
+ def batch_select_indices(self, indices: torch.Tensor) -> None:
+ if self.get_seq_length() > 0:
+ self.keys = self.keys[indices, ...]
+ self.values = self.values[indices, ...]
+
+
+class Cache:
+ def __init__(self, layers: Optional[list[CacheLayerMixin]] = None, layer_class_to_replicate: Optional[type[CacheLayerMixin]] = None, offloading: bool = False, offload_only_non_sliding: bool = True):
+ if layers is not None and layer_class_to_replicate is not None:
+ raise ValueError("Provide exactly one of `layers` or `layer_class_to_replicate`.")
+ if layers is None and layer_class_to_replicate is None:
+ raise ValueError("Provide exactly one of `layers` or `layer_class_to_replicate`.")
+ self.layers = layers if layers is not None else []
+ self.layer_class_to_replicate = layer_class_to_replicate
+ self.offloading = offloading
+ if self.offloading:
+ self.only_non_sliding = offload_only_non_sliding
+ self.prefetch_stream = torch.cuda.Stream()
+
+ def __repr__(self):
+ return f"{self.__class__.__name__}(layers={self.layers})"
+
+ def update(self, key_states: torch.Tensor, value_states: torch.Tensor, layer_idx: int, cache_kwargs: Optional[dict[str, Any]] = None) -> tuple[torch.Tensor, torch.Tensor]:
+ if self.layer_class_to_replicate is not None:
+ while len(self.layers) <= layer_idx:
+ self.layers.append(self.layer_class_to_replicate())
+ if self.offloading:
+ torch.cuda.default_stream(key_states.device).wait_stream(self.prefetch_stream)
+ self.prefetch(layer_idx + 1, self.only_non_sliding)
+ keys, values = self.layers[layer_idx].update(key_states, value_states, cache_kwargs)
+ if self.offloading:
+ self.offload(layer_idx, self.only_non_sliding)
+ return keys, values
+
+ def get_seq_length(self, layer_idx: int = 0) -> int:
+ if layer_idx >= len(self.layers):
+ return 0
+ return self.layers[layer_idx].get_seq_length()
+
+ def get_mask_sizes(self, cache_position: torch.Tensor, layer_idx: int) -> tuple[int, int]:
+ if layer_idx >= len(self.layers):
+ return cache_position.shape[0], 0
+ return self.layers[layer_idx].get_mask_sizes(cache_position)
+
+ def get_max_cache_shape(self, layer_idx: int = 0) -> int:
+ if layer_idx >= len(self.layers):
+ return -1
+ return self.layers[layer_idx].get_max_cache_shape()
+
+ def reset(self):
+ for layer_idx in range(len(self.layers)):
+ self.layers[layer_idx].reset()
+
+ def reorder_cache(self, beam_idx: torch.LongTensor):
+ for layer_idx in range(len(self.layers)):
+ self.layers[layer_idx].reorder_cache(beam_idx)
+
+ def crop(self, max_length: int):
+ for layer_idx in range(len(self.layers)):
+ self.layers[layer_idx].crop(max_length)
+
+ def batch_repeat_interleave(self, repeats: int):
+ for layer_idx in range(len(self.layers)):
+ self.layers[layer_idx].batch_repeat_interleave(repeats)
+
+ def batch_select_indices(self, indices: torch.Tensor):
+ for layer_idx in range(len(self.layers)):
+ self.layers[layer_idx].batch_select_indices(indices)
+
+ def prefetch(self, layer_idx: int, only_non_sliding: bool = True):
+ if only_non_sliding:
+ try:
+ layer_idx = layer_idx + self.is_sliding[layer_idx:].index(False)
+ except ValueError:
+ layer_idx = self.is_sliding.index(False)
+ else:
+ layer_idx = layer_idx if layer_idx < len(self.layers) else 0
+ with torch.cuda.stream(self.prefetch_stream):
+ self.layers[layer_idx].prefetch()
+
+ def offload(self, layer_idx: int, only_non_sliding: bool = True):
+ if not (only_non_sliding and self.is_sliding[layer_idx]):
+ self.layers[layer_idx].offload()
+
+ @property
+ def is_sliding(self) -> list[bool]:
+ return [getattr(layer, "is_sliding", False) for layer in self.layers]
+
+ def __getitem__(self, layer_idx: int) -> tuple[torch.Tensor, torch.Tensor]:
+ if layer_idx < len(self.layers):
+ return self.layers[layer_idx].keys, self.layers[layer_idx].values
+ else:
+ raise KeyError(f"Cache only has {len(self.layers)} layers, attempted to access layer with index {layer_idx}")
+
+ def __iter__(self):
+ for layer_idx in range(len(self)):
+ yield (self.layers[layer_idx].keys, self.layers[layer_idx].values)
+
+ def __len__(self):
+ return len(self.layers)
+
+ @property
+ def is_compileable(self) -> bool:
+ if len(self.layers) == 0:
+ return False
+ return all(layer.is_compileable for layer in self.layers)
+
+ @property
+ def is_initialized(self) -> bool:
+ return len(self.layers) > 0 and all(layer.is_initialized for layer in self.layers)
+
+
+class ModelOutput(OrderedDict):
+ """Base class for model outputs that allows additional fields."""
+ def __post_init__(self):
+ if dataclasses.is_dataclass(self):
+ self.__dict__.update({f.name: getattr(self, f.name) for f in dataclasses.fields(self)})
+
+ def __getitem__(self, key):
+ return getattr(self, key)
+
+ def __setitem__(self, key, value):
+ setattr(self, key, value)
+
+ def __iter__(self):
+ if dataclasses.is_dataclass(self):
+ for f in dataclasses.fields(self):
+ val = getattr(self, f.name)
+ if val is not None:
+ yield val
+ else:
+ for key in self.keys():
+ yield self[key]
+
+ def keys(self):
+ if dataclasses.is_dataclass(self):
+ return [f.name for f in dataclasses.fields(self) if getattr(self, f.name) is not None]
+ return list(self.__dict__.keys())
+
+ def values(self):
+ if dataclasses.is_dataclass(self):
+ return [getattr(self, f.name) for f in dataclasses.fields(self) if getattr(self, f.name) is not None]
+ return list(self.__dict__.values())
+
+ def items(self):
+ if dataclasses.is_dataclass(self):
+ return [(f.name, getattr(self, f.name)) for f in dataclasses.fields(self) if getattr(self, f.name) is not None]
+ return list(self.__dict__.items())
+
+ def __contains__(self, key):
+ return hasattr(self, key) and getattr(self, key) is not None
+
+
+@dataclass
+class BaseModelOutputWithPast(ModelOutput):
+ last_hidden_state: Optional[torch.FloatTensor] = None
+ past_key_values: Optional[Cache] = None
+ hidden_states: Optional[tuple[torch.FloatTensor, ...]] = None
+ attentions: Optional[tuple[torch.FloatTensor, ...]] = None
+
+
+def _compute_default_rope_parameters(config, device=None, seq_len=None):
+ base = config.rope_theta
+ partial_rotary_factor = getattr(config, "partial_rotary_factor", 1.0)
+ head_dim = getattr(config, "head_dim", None) or config.hidden_size // config.num_attention_heads
+ dim = int(head_dim * partial_rotary_factor)
+ attention_factor = 1.0
+ inv_freq = 1.0 / (base ** (torch.arange(0, dim, 2, dtype=torch.int64).to(device=device, dtype=torch.float) / dim))
+ return inv_freq, attention_factor
+
+
+ROPE_INIT_FUNCTIONS = {
+ "default": _compute_default_rope_parameters,
+}
+
+
+class Qwen3VLVisionConfig(PretrainedConfig):
+
+ model_type = "qwen3_vl"
+ base_config_key = "vision_config"
+
+ def __init__(
+ self,
+ depth=27,
+ hidden_size=1152,
+ hidden_act="gelu_pytorch_tanh",
+ intermediate_size=4304,
+ num_heads=16,
+ in_channels=3,
+ patch_size=16,
+ spatial_merge_size=2,
+ temporal_patch_size=2,
+ out_hidden_size=3584,
+ num_position_embeddings=2304,
+ deepstack_visual_indexes=[8, 16, 24],
+ initializer_range=0.02,
+ **kwargs,
+ ):
+ super().__init__(**kwargs)
+
+ self.depth = depth
+ self.hidden_size = hidden_size
+ self.hidden_act = hidden_act
+ self.intermediate_size = intermediate_size
+ self.num_heads = num_heads
+ self.in_channels = in_channels
+ self.patch_size = patch_size
+ self.spatial_merge_size = spatial_merge_size
+ self.temporal_patch_size = temporal_patch_size
+ self.out_hidden_size = out_hidden_size
+ self.num_position_embeddings = num_position_embeddings
+ self.initializer_range = initializer_range
+ self.deepstack_visual_indexes = deepstack_visual_indexes
+
+
+class Qwen3VLTextConfig(PretrainedConfig):
+
+ model_type = "qwen3_vl_text"
+ base_config_key = "text_config"
+
+ def __init__(
+ self,
+ vocab_size=151936,
+ hidden_size=4096,
+ intermediate_size=22016,
+ num_hidden_layers=32,
+ num_attention_heads=32,
+ num_key_value_heads=32,
+ head_dim=128,
+ hidden_act="silu",
+ max_position_embeddings=128000,
+ initializer_range=0.02,
+ rms_norm_eps=1e-6,
+ use_cache=True,
+ tie_word_embeddings=False,
+ rope_theta=5000000.0,
+ rope_scaling=None,
+ attention_bias=False,
+ attention_dropout=0.0,
+ **kwargs,
+ ):
+ self.vocab_size = vocab_size
+ self.max_position_embeddings = max_position_embeddings
+ self.hidden_size = hidden_size
+ self.intermediate_size = intermediate_size
+ self.num_hidden_layers = num_hidden_layers
+ self.num_attention_heads = num_attention_heads
+ if num_key_value_heads is None:
+ num_key_value_heads = num_attention_heads
+ self.num_key_value_heads = num_key_value_heads
+ self.head_dim = head_dim
+ self.hidden_act = hidden_act
+ self.initializer_range = initializer_range
+ self.rms_norm_eps = rms_norm_eps
+ self.use_cache = use_cache
+ self.rope_theta = rope_theta
+ self.rope_scaling = rope_scaling
+ self.attention_bias = attention_bias
+ self.attention_dropout = attention_dropout
+ super().__init__(tie_word_embeddings=tie_word_embeddings, **kwargs)
+
+
+class Qwen3VLConfig(PretrainedConfig):
+
+ model_type = "qwen3_vl"
+ sub_configs = {"vision_config": Qwen3VLVisionConfig, "text_config": Qwen3VLTextConfig}
+ keys_to_ignore_at_inference = ["past_key_values"]
+
+ def __init__(
+ self,
+ text_config=None,
+ vision_config=None,
+ image_token_id=151655,
+ video_token_id=151656,
+ vision_start_token_id=151652,
+ vision_end_token_id=151653,
+ tie_word_embeddings=False,
+ **kwargs,
+ ):
+ if isinstance(vision_config, dict):
+ self.vision_config = self.sub_configs["vision_config"](**vision_config)
+ elif vision_config is None:
+ self.vision_config = self.sub_configs["vision_config"]()
+
+ if isinstance(text_config, dict):
+ self.text_config = self.sub_configs["text_config"](**text_config)
+ elif text_config is None:
+ self.text_config = self.sub_configs["text_config"]()
+
+ self.image_token_id = image_token_id
+ self.video_token_id = video_token_id
+ self.vision_start_token_id = vision_start_token_id
+ self.vision_end_token_id = vision_end_token_id
+ super().__init__(**kwargs, tie_word_embeddings=tie_word_embeddings)
+
+
+class Qwen3VLVisionMLP(nn.Module):
+ def __init__(self, config):
+ super().__init__()
+ self.hidden_size = config.hidden_size
+ self.intermediate_size = config.intermediate_size
+ self.linear_fc1 = nn.Linear(self.hidden_size, self.intermediate_size, bias=True)
+ self.linear_fc2 = nn.Linear(self.intermediate_size, self.hidden_size, bias=True)
+ self.act_fn = ACT2FN[config.hidden_act]
+
+ def forward(self, hidden_state):
+ return self.linear_fc2(self.act_fn(self.linear_fc1(hidden_state)))
+
+
+class Qwen3VLVisionPatchEmbed(nn.Module):
+ def __init__(self, config) -> None:
+ super().__init__()
+ self.patch_size = config.patch_size
+ self.temporal_patch_size = config.temporal_patch_size
+ self.in_channels = config.in_channels
+ self.embed_dim = config.hidden_size
+
+ kernel_size = [self.temporal_patch_size, self.patch_size, self.patch_size]
+ self.proj = nn.Conv3d(self.in_channels, self.embed_dim, kernel_size=kernel_size, stride=kernel_size, bias=True)
+
+ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+ target_dtype = self.proj.weight.dtype
+ hidden_states = hidden_states.view(
+ -1, self.in_channels, self.temporal_patch_size, self.patch_size, self.patch_size
+ )
+ hidden_states = self.proj(hidden_states.to(dtype=target_dtype)).view(-1, self.embed_dim)
+ return hidden_states
+
+
+class Qwen3VLVisionRotaryEmbedding(nn.Module):
+ inv_freq: torch.Tensor
+
+ def __init__(self, dim: int, theta: float = 10000.0) -> None:
+ super().__init__()
+ inv_freq = 1.0 / (theta ** (torch.arange(0, dim, 2, dtype=torch.float) / dim))
+ self.register_buffer("inv_freq", inv_freq, persistent=False)
+
+ def forward(self, seqlen: int) -> torch.Tensor:
+ seq = torch.arange(seqlen, device=self.inv_freq.device, dtype=self.inv_freq.dtype)
+ freqs = torch.outer(seq, self.inv_freq)
+ return freqs
+
+class Qwen3VLVisionPatchMerger(nn.Module):
+ def __init__(self, config, use_postshuffle_norm=False) -> None:
+ super().__init__()
+ self.hidden_size = config.hidden_size * (config.spatial_merge_size**2)
+ self.use_postshuffle_norm = use_postshuffle_norm
+ self.norm = nn.LayerNorm(self.hidden_size if use_postshuffle_norm else config.hidden_size, eps=1e-6)
+ self.linear_fc1 = nn.Linear(self.hidden_size, self.hidden_size)
+ self.act_fn = nn.GELU()
+ self.linear_fc2 = nn.Linear(self.hidden_size, config.out_hidden_size)
+
+ def forward(self, x: torch.Tensor) -> torch.Tensor:
+ x = self.norm(x.view(-1, self.hidden_size) if self.use_postshuffle_norm else x).view(-1, self.hidden_size)
+ x = self.linear_fc2(self.act_fn(self.linear_fc1(x)))
+ return x
+
+
+def rotate_half(x):
+ """Rotates half the hidden dims of the input."""
+ x1 = x[..., : x.shape[-1] // 2]
+ x2 = x[..., x.shape[-1] // 2 :]
+ return torch.cat((-x2, x1), dim=-1)
+
+
+def apply_rotary_pos_emb_vision(
+ q: torch.Tensor, k: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor
+) -> tuple[torch.Tensor, torch.Tensor]:
+ orig_q_dtype = q.dtype
+ orig_k_dtype = k.dtype
+ q, k = q.float(), k.float()
+ cos, sin = cos.unsqueeze(-2).float(), sin.unsqueeze(-2).float()
+ q_embed = (q * cos) + (rotate_half(q) * sin)
+ k_embed = (k * cos) + (rotate_half(k) * sin)
+ q_embed = q_embed.to(orig_q_dtype)
+ k_embed = k_embed.to(orig_k_dtype)
+ return q_embed, k_embed
+
+
+def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor:
+ """
+ Equivalent of torch.repeat_interleave(x, dim=1, repeats=n_rep).
+ The hidden states go from (batch, num_key_value_heads, seqlen, head_dim)
+ to (batch, num_attention_heads, seqlen, head_dim)
+ """
+ batch, num_key_value_heads, slen, head_dim = hidden_states.shape
+ if n_rep == 1:
+ return hidden_states
+ hidden_states = hidden_states[:, :, None, :, :].expand(batch, num_key_value_heads, n_rep, slen, head_dim)
+ return hidden_states.reshape(batch, num_key_value_heads * n_rep, slen, head_dim)
+
+
+def apply_rotary_pos_emb(q, k, cos, sin, unsqueeze_dim=1):
+ """Applies rotary position embeddings to query and key states."""
+ cos = cos.unsqueeze(unsqueeze_dim)
+ sin = sin.unsqueeze(unsqueeze_dim)
+ q_embed = (q * cos) + (rotate_half(q) * sin)
+ k_embed = (k * cos) + (rotate_half(k) * sin)
+ return q_embed, k_embed
+
+
+class Qwen3VLVisionAttention(nn.Module):
+ def __init__(self, config) -> None:
+ super().__init__()
+ self.dim = config.hidden_size
+ self.num_heads = config.num_heads
+ self.head_dim = self.dim // self.num_heads
+ self.num_key_value_groups = 1
+ self.qkv = nn.Linear(self.dim, self.dim * 3, bias=True)
+ self.proj = nn.Linear(self.dim, self.dim)
+ self.scaling = self.head_dim**-0.5
+ self.config = config
+ self.attention_dropout = 0.0
+ self.is_causal = False
+
+ def forward(
+ self,
+ hidden_states: torch.Tensor,
+ cu_seqlens: torch.Tensor,
+ rotary_pos_emb: Optional[torch.Tensor] = None,
+ position_embeddings: Optional[tuple[torch.Tensor, torch.Tensor]] = None,
+ **kwargs,
+ ) -> torch.Tensor:
+ seq_length = hidden_states.shape[0]
+ query_states, key_states, value_states = (
+ self.qkv(hidden_states).reshape(seq_length, 3, self.num_heads, -1).permute(1, 0, 2, 3).unbind(0)
+ )
+ cos, sin = position_embeddings
+ query_states, key_states = apply_rotary_pos_emb_vision(query_states, key_states, cos, sin)
+
+ query_states = query_states.transpose(0, 1).unsqueeze(0)
+ key_states = key_states.transpose(0, 1).unsqueeze(0)
+ value_states = value_states.transpose(0, 1).unsqueeze(0)
+
+ # Process each chunk separately using DiffSynth attention_forward
+ lengths = cu_seqlens[1:] - cu_seqlens[:-1]
+ splits = [
+ torch.split(tensor, lengths.tolist(), dim=2)
+ for tensor in (query_states, key_states, value_states)
+ ]
+
+ attn_outputs = [
+ attention_forward(
+ q, k, v,
+ q_pattern="b n s d", k_pattern="b n s d", v_pattern="b n s d",
+ out_pattern="b n s d",
+ scale=self.scaling,
+ )
+ for q, k, v in zip(*splits)
+ ]
+ attn_output = torch.cat(attn_outputs, dim=2) # [B, N, total_S, D]
+ attn_output = attn_output.transpose(1, 2).contiguous() # [B, total_S, N, D]
+ attn_output = attn_output.reshape(seq_length, -1).contiguous()
+ attn_output = self.proj(attn_output)
+ return attn_output
+
+
+class Qwen3VLVisionBlock(nn.Module):
+ def __init__(self, config, attn_implementation: str = "sdpa") -> None:
+ super().__init__()
+ self.norm1 = nn.LayerNorm(config.hidden_size, eps=1e-6)
+ self.norm2 = nn.LayerNorm(config.hidden_size, eps=1e-6)
+ self.attn = Qwen3VLVisionAttention(config=config)
+ self.mlp = Qwen3VLVisionMLP(config=config)
+
+ def forward(
+ self,
+ hidden_states: torch.Tensor,
+ cu_seqlens: torch.Tensor,
+ rotary_pos_emb: Optional[torch.Tensor] = None,
+ position_embeddings: Optional[tuple[torch.Tensor, torch.Tensor]] = None,
+ **kwargs,
+ ) -> torch.Tensor:
+ hidden_states = hidden_states + self.attn(
+ self.norm1(hidden_states),
+ cu_seqlens=cu_seqlens,
+ rotary_pos_emb=rotary_pos_emb,
+ position_embeddings=position_embeddings,
+ **kwargs,
+ )
+ hidden_states = hidden_states + self.mlp(self.norm2(hidden_states))
+ return hidden_states
+
+
+class Qwen3VLTextRMSNorm(nn.Module):
+ def __init__(self, hidden_size, eps=1e-6):
+ super().__init__()
+ self.weight = nn.Parameter(torch.ones(hidden_size))
+ self.variance_epsilon = eps
+
+ def forward(self, hidden_states):
+ input_dtype = hidden_states.dtype
+ hidden_states = hidden_states.float()
+ variance = hidden_states.pow(2).mean(-1, keepdim=True)
+ hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon)
+ return self.weight * hidden_states.to(input_dtype)
+
+
+def eager_attention_forward(
+ module,
+ query: torch.Tensor,
+ key: torch.Tensor,
+ value: torch.Tensor,
+ attention_mask: Optional[torch.Tensor],
+ scaling: float,
+ dropout: float = 0.0,
+ **kwargs,
+):
+ key_states = repeat_kv(key, module.num_key_value_groups)
+ value_states = repeat_kv(value, module.num_key_value_groups)
+
+ attn_weights = torch.matmul(query, key_states.transpose(2, 3)) * scaling
+ if attention_mask is not None:
+ causal_mask = attention_mask[:, :, :, : key_states.shape[-2]]
+ attn_weights = attn_weights + causal_mask
+
+ attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to(query.dtype)
+ attn_weights = nn.functional.dropout(attn_weights, p=dropout, training=module.training)
+ attn_output = torch.matmul(attn_weights, value_states)
+ attn_output = attn_output.transpose(1, 2).contiguous()
+
+ return attn_output, attn_weights
+
+
+class Qwen3VLTextMLP(nn.Module):
+ def __init__(self, config):
+ super().__init__()
+ self.config = config
+ self.hidden_size = config.hidden_size
+ self.intermediate_size = config.intermediate_size
+ self.gate_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False)
+ self.up_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False)
+ self.down_proj = nn.Linear(self.intermediate_size, self.hidden_size, bias=False)
+ self.act_fn = ACT2FN[config.hidden_act]
+
+ def forward(self, x):
+ down_proj = self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x))
+ return down_proj
+
+
+class Qwen3VLTextDecoderLayer(nn.Module):
+ def __init__(self, config, layer_idx: int):
+ super().__init__()
+ self.hidden_size = config.hidden_size
+
+ self.self_attn = Qwen3VLTextAttention(config=config, layer_idx=layer_idx)
+
+ self.mlp = Qwen3VLTextMLP(config)
+ self.input_layernorm = Qwen3VLTextRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+ self.post_attention_layernorm = Qwen3VLTextRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+
+ def forward(
+ self,
+ hidden_states: torch.Tensor,
+ position_embeddings: tuple[torch.Tensor, torch.Tensor],
+ attention_mask: Optional[torch.Tensor] = None,
+ position_ids: Optional[torch.LongTensor] = None,
+ past_key_values: Optional[Cache] = None,
+ use_cache: Optional[bool] = False,
+ cache_position: Optional[torch.LongTensor] = None,
+ **kwargs,
+ ) -> torch.Tensor:
+ residual = hidden_states
+ hidden_states = self.input_layernorm(hidden_states)
+ hidden_states, _ = self.self_attn(
+ hidden_states=hidden_states,
+ attention_mask=attention_mask,
+ position_ids=position_ids,
+ past_key_values=past_key_values,
+ use_cache=use_cache,
+ cache_position=cache_position,
+ position_embeddings=position_embeddings,
+ **kwargs,
+ )
+ hidden_states = residual + hidden_states
+
+ residual = hidden_states
+ hidden_states = self.post_attention_layernorm(hidden_states)
+ hidden_states = self.mlp(hidden_states)
+ hidden_states = residual + hidden_states
+ return hidden_states
+
+
+@dataclass
+class Qwen3VLModelOutputWithPast(ModelOutput):
+ last_hidden_state: Optional[torch.FloatTensor] = None
+ past_key_values: Optional[Cache] = None
+ hidden_states: Optional[tuple[torch.FloatTensor]] = None
+ attentions: Optional[tuple[torch.FloatTensor]] = None
+ rope_deltas: Optional[torch.LongTensor] = None
+ x_pred: Optional[torch.FloatTensor] = None
+ mid_results: Optional[list] = None
+
+
+class Qwen3VLTextAttention(nn.Module):
+ """Multi-headed attention from 'Attention Is All You Need' paper"""
+
+ def __init__(self, config, layer_idx: int):
+ super().__init__()
+ self.config = config
+ self.layer_idx = layer_idx
+ self.head_dim = getattr(config, "head_dim", config.hidden_size // config.num_attention_heads)
+ self.num_key_value_groups = config.num_attention_heads // config.num_key_value_heads
+ self.scaling = self.head_dim**-0.5
+ self.attention_dropout = config.attention_dropout
+ self.is_causal = True
+
+ self.q_proj = nn.Linear(
+ config.hidden_size, config.num_attention_heads * self.head_dim, bias=config.attention_bias
+ )
+ self.k_proj = nn.Linear(
+ config.hidden_size, config.num_key_value_heads * self.head_dim, bias=config.attention_bias
+ )
+ self.v_proj = nn.Linear(
+ config.hidden_size, config.num_key_value_heads * self.head_dim, bias=config.attention_bias
+ )
+ self.o_proj = nn.Linear(
+ config.num_attention_heads * self.head_dim, config.hidden_size, bias=config.attention_bias
+ )
+ self.q_norm = Qwen3VLTextRMSNorm(self.head_dim, eps=config.rms_norm_eps)
+ self.k_norm = Qwen3VLTextRMSNorm(self.head_dim, eps=config.rms_norm_eps)
+
+ def forward(
+ self,
+ hidden_states: torch.Tensor,
+ position_embeddings: tuple[torch.Tensor, torch.Tensor],
+ attention_mask: Optional[torch.Tensor],
+ past_key_values: Optional[Cache] = None,
+ cache_position: Optional[torch.LongTensor] = None,
+ **kwargs,
+ ) -> tuple[torch.Tensor, Optional[torch.Tensor]]:
+ input_shape = hidden_states.shape[:-1]
+ hidden_shape = (*input_shape, -1, self.head_dim)
+
+ query_states = self.q_norm(self.q_proj(hidden_states).view(hidden_shape)).transpose(1, 2)
+ key_states = self.k_norm(self.k_proj(hidden_states).view(hidden_shape)).transpose(1, 2)
+ value_states = self.v_proj(hidden_states).view(hidden_shape).transpose(1, 2)
+
+ cos, sin = position_embeddings
+ query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin)
+
+ if past_key_values is not None:
+ cache_kwargs = {"sin": sin, "cos": cos, "cache_position": cache_position}
+ key_states, value_states = past_key_values.update(key_states, value_states, self.layer_idx, cache_kwargs)
+
+ # Repeat KV for GQA before attention_forward
+ key_states = repeat_kv(key_states, self.num_key_value_groups)
+ value_states = repeat_kv(value_states, self.num_key_value_groups)
+
+ # Use DiffSynth attention_forward
+ attn_output = attention_forward(
+ query_states, key_states, value_states,
+ q_pattern="b n s d", k_pattern="b n s d", v_pattern="b n s d",
+ out_pattern="b n s d",
+ attn_mask=attention_mask,
+ scale=self.scaling,
+ )
+
+ attn_weights = None
+
+ # Flatten and project
+ attn_output = attn_output.transpose(1, 2).flatten(2, 3).contiguous()
+ attn_output = self.o_proj(attn_output)
+ return attn_output, attn_weights
+
+
+class BottleneckPatchEmbed(nn.Module):
+ def __init__(self, config, patch_size=32, in_chans=3, pca_dim=1024, embed_dim=4096, bias=True):
+ super().__init__()
+ self.config = config
+ self.pca_dim = pca_dim
+ self.embed_dim = embed_dim
+ self.patch_size = patch_size
+ self.in_chans = in_chans
+
+ self.proj1 = nn.Linear(patch_size * patch_size * in_chans, pca_dim, bias=False)
+ self.proj2 = nn.Linear(pca_dim, embed_dim, bias=bias)
+
+ def forward(self, x: torch.Tensor) -> torch.Tensor:
+ x = self.proj2(self.proj1(x))
+ return x
+
+
+class FinalLayer(nn.Module):
+ def __init__(self, config, hidden_size, patch_size, out_channels):
+ super().__init__()
+ self.linear = nn.Linear(hidden_size, patch_size * patch_size * out_channels, bias=True)
+
+ def forward(self, x, adaln_input=None):
+ x = self.linear(x)
+ return x
+
+
+class TimestepEmbedder(nn.Module):
+ def __init__(self, config, hidden_size, frequency_embedding_size=256):
+ super().__init__()
+ self.mlp = nn.Sequential(
+ nn.Linear(frequency_embedding_size, hidden_size, bias=True),
+ nn.SiLU(),
+ nn.Linear(hidden_size, hidden_size, bias=True),
+ )
+ self.frequency_embedding_size = frequency_embedding_size
+
+ @staticmethod
+ def timestep_embedding(t, dim, max_period=10000):
+ half = dim // 2
+ freqs = torch.exp(
+ -math.log(max_period) * torch.arange(start=0, end=half, dtype=torch.float32, device=t.device) / half
+ )
+ args = t.float()[:, None] * freqs[None, :]
+ embedding = torch.cat([torch.cos(args), torch.sin(args)], dim=-1)
+ if dim % 2:
+ embedding = torch.cat([embedding, torch.zeros_like(embedding[:, :1])], dim=-1)
+ return embedding
+
+ def forward(self, t):
+ t_freq = self.timestep_embedding(t * 1000, self.frequency_embedding_size)
+ t_emb = self.mlp(t_freq.to(self.mlp[0].weight.dtype))
+ return t_emb
+
+
+class Qwen3VLPreTrainedModel(PreTrainedModel):
+ config: Qwen3VLConfig
+ base_model_prefix = "model"
+ supports_gradient_checkpointing = True
+ _no_split_modules = ["Qwen3VLTextDecoderLayer", "Qwen3VLVisionBlock"]
+ _skip_keys_device_placement = "past_key_values"
+ _supports_flash_attn = True
+ _supports_sdpa = True
+ _can_compile_fullgraph = True
+ _supports_attention_backend = True
+ _can_record_outputs = {
+ "hidden_states": Qwen3VLTextDecoderLayer,
+ "attentions": Qwen3VLTextAttention,
+ }
+
+
+class Qwen3VLModel(Qwen3VLPreTrainedModel):
+ config: Qwen3VLConfig
+ base_model_prefix = "model"
+ supports_gradient_checkpointing = True
+ _no_split_modules = ["Qwen3VLTextDecoderLayer", "Qwen3VLVisionBlock"]
+ _skip_keys_device_placement = "past_key_values"
+ _supports_flash_attn = True
+ _supports_sdpa = True
+ _can_compile_fullgraph = True
+ _supports_attention_backend = True
+
+ def __init__(self, config):
+ super().__init__(config)
+ self.language_model = Qwen3VLTextModel(config.text_config)
+ self.visual = Qwen3VLVisionModel(config.vision_config)
+
+ self.patch_size = 32
+ self.in_channels = 3
+ hidden_size = config.text_config.hidden_size
+ bottleneck_dim = hidden_size // 4
+
+ self.t_embedder1 = TimestepEmbedder(config, hidden_size)
+ self.x_embedder = BottleneckPatchEmbed(
+ config,
+ patch_size=self.patch_size,
+ in_chans=self.in_channels,
+ pca_dim=bottleneck_dim,
+ embed_dim=hidden_size,
+ bias=True,
+ )
+ self.final_layer2 = FinalLayer(
+ config,
+ hidden_size=hidden_size,
+ patch_size=self.patch_size,
+ out_channels=self.in_channels,
+ )
+
+ self.tms_token_id = 151673
+ self.rope_deltas = None
+
+ def get_input_embeddings(self):
+ return self.language_model.get_input_embeddings()
+
+ def set_input_embeddings(self, value):
+ self.language_model.set_input_embeddings(value)
+
+ def set_decoder(self, decoder):
+ self.language_model = decoder
+
+ def get_decoder(self):
+ return self.language_model
+
+ @property
+ def language_model(self):
+ return self._language_model
+
+ @language_model.setter
+ def language_model(self, value):
+ self._language_model = value
+
+ @property
+ def visual(self):
+ return self._visual
+
+ @visual.setter
+ def visual(self, value):
+ self._visual = value
+
+ def get_video_features(
+ self, pixel_values_videos: torch.FloatTensor, video_grid_thw: Optional[torch.LongTensor] = None
+ ):
+ return self.get_image_features(pixel_values_videos, video_grid_thw)
+
+ def get_image_features(self, pixel_values: torch.FloatTensor, image_grid_thw: Optional[torch.LongTensor] = None):
+ pixel_values = pixel_values.type(self.visual.dtype)
+ image_embeds, deepstack_image_embeds = self.visual(pixel_values, grid_thw=image_grid_thw)
+ split_sizes = (image_grid_thw.prod(-1) // self.visual.spatial_merge_size**2).tolist()
+ image_embeds = torch.split(image_embeds, split_sizes)
+ return image_embeds, deepstack_image_embeds
+
+ def get_placeholder_mask(
+ self,
+ input_ids: torch.LongTensor,
+ inputs_embeds: torch.FloatTensor,
+ image_features: Optional[torch.FloatTensor] = None,
+ video_features: Optional[torch.FloatTensor] = None,
+ ):
+ if input_ids is None:
+ special_image_mask = inputs_embeds == self.get_input_embeddings()(
+ torch.tensor(self.config.image_token_id, dtype=torch.long, device=inputs_embeds.device)
+ )
+ special_image_mask = special_image_mask.all(-1)
+ special_video_mask = inputs_embeds == self.get_input_embeddings()(
+ torch.tensor(self.config.video_token_id, dtype=torch.long, device=inputs_embeds.device)
+ )
+ special_video_mask = special_video_mask.all(-1)
+ else:
+ special_image_mask = input_ids == self.config.image_token_id
+ special_video_mask = input_ids == self.config.video_token_id
+
+ n_image_tokens = special_image_mask.sum()
+ special_image_mask = special_image_mask.unsqueeze(-1).expand_as(inputs_embeds).to(inputs_embeds.device)
+ if image_features is not None and inputs_embeds[special_image_mask].numel() != image_features.numel():
+ raise ValueError(
+ f"Image features and image tokens do not match: tokens: {n_image_tokens}, features {image_features.shape[0]}"
+ )
+
+ n_video_tokens = special_video_mask.sum()
+ special_video_mask = special_video_mask.unsqueeze(-1).expand_as(inputs_embeds).to(inputs_embeds.device)
+ if video_features is not None and inputs_embeds[special_video_mask].numel() != video_features.numel():
+ raise ValueError(
+ f"Videos features and video tokens do not match: tokens: {n_video_tokens}, features {video_features.shape[0]}"
+ )
+
+ return special_image_mask, special_video_mask
+
+ def _run_decoder_flash(self, inputs_embeds, position_ids, token_types, return_mid_results_layers=None,
+ use_gradient_checkpointing=False, use_gradient_checkpointing_offload=False):
+ """Run decoder layers with flash attention two-pass approach.
+
+ Replicates the Megatron attention pattern:
+ 1. Causal attention on AR tokens only (text)
+ 2. Full (bidirectional) attention on ALL tokens
+ 3. Replace AR positions with causal result (index_copy)
+ """
+ text_model = self.language_model
+
+ if position_ids.ndim == 2:
+ position_ids = position_ids[None, ...].expand(3, position_ids.shape[0], -1)
+ elif position_ids.ndim == 3 and position_ids.shape[0] == 4:
+ position_ids = position_ids[1:]
+ position_embeddings = text_model.rotary_emb(inputs_embeds, position_ids)
+ cos, sin = position_embeddings
+
+ is_gen = token_types[0].bool()
+ idx_ar = torch.nonzero(~is_gen, as_tuple=False).squeeze(-1)
+
+ hidden_states = inputs_embeds
+ mid_results = [] if return_mid_results_layers else None
+
+ def _flash_layer_forward(hidden_states, decoder_layer, cos, sin, idx_ar):
+ """Flash attention layer with two-pass approach using DiffSynth attention_forward."""
+ original_attn_forward = decoder_layer.self_attn.forward
+
+ def _custom_flash_attn(hidden_states, position_embeddings, attention_mask=None, **kwargs):
+ attn = decoder_layer.self_attn
+ input_shape = hidden_states.shape[:-1]
+ head_dim = attn.head_dim
+ hidden_shape = (*input_shape, -1, head_dim)
+
+ q = attn.q_norm(attn.q_proj(hidden_states).view(hidden_shape))
+ k = attn.k_norm(attn.k_proj(hidden_states).view(hidden_shape))
+ v = attn.v_proj(hidden_states).view(hidden_shape)
+
+ cos_pe, sin_pe = position_embeddings
+ q_r = q.transpose(1, 2)
+ k_r = k.transpose(1, 2)
+ q_r, k_r = apply_rotary_pos_emb(q_r, k_r, cos_pe, sin_pe)
+ q = q_r.transpose(1, 2).contiguous()
+ k = k_r.transpose(1, 2).contiguous()
+ v = v.contiguous()
+
+ softmax_scale = head_dim ** -0.5
+
+ # Rearrange to [B, H, S, D] for attention_forward
+ q_bn = q.transpose(1, 2).contiguous()
+ k_bn = k.transpose(1, 2).contiguous()
+ v_bn = v.transpose(1, 2).contiguous()
+
+ # Handle GQA: repeat K/V heads to match Q heads for attention_forward
+ n_rep = attn.num_key_value_groups
+ if n_rep > 1:
+ k_bn = k_bn.repeat_interleave(n_rep, dim=1)
+ v_bn = v_bn.repeat_interleave(n_rep, dim=1)
+
+ # Two-pass attention using attention_forward
+ # Pass 1: causal attention on AR tokens only
+ q_ar = q_bn[:, :, idx_ar].contiguous()
+ k_ar = k_bn[:, :, idx_ar].contiguous()
+ v_ar = v_bn[:, :, idx_ar].contiguous()
+ out_ar = attention_forward(
+ q_ar, k_ar, v_ar,
+ q_pattern="b n s d", k_pattern="b n s d", v_pattern="b n s d",
+ out_pattern="b n s d",
+ is_causal=True,
+ scale=softmax_scale,
+ )
+
+ # Pass 2: full (bidirectional) attention on all tokens
+ out_full = attention_forward(
+ q_bn, k_bn, v_bn,
+ q_pattern="b n s d", k_pattern="b n s d", v_pattern="b n s d",
+ out_pattern="b n s d",
+ is_causal=False,
+ scale=softmax_scale,
+ )
+
+ # Replace AR positions with causal result, rearrange back to [B, S, H, D]
+ out_full = out_full.clone()
+ out_full[:, :, idx_ar] = out_ar
+ out_full = out_full.transpose(1, 2).contiguous()
+
+ attn_output = out_full.reshape(*input_shape, -1).contiguous()
+ attn_output = attn.o_proj(attn_output)
+ return attn_output, None
+
+ decoder_layer.self_attn.forward = _custom_flash_attn
+ try:
+ hidden_states = decoder_layer(
+ hidden_states,
+ position_embeddings=(cos, sin),
+ )
+ finally:
+ decoder_layer.self_attn.forward = original_attn_forward
+
+ return hidden_states
+
+ for layer_idx, decoder_layer in enumerate(text_model.layers):
+ hidden_states = gradient_checkpoint_forward(
+ _flash_layer_forward,
+ use_gradient_checkpointing=use_gradient_checkpointing,
+ use_gradient_checkpointing_offload=use_gradient_checkpointing_offload,
+ hidden_states=hidden_states,
+ decoder_layer=decoder_layer,
+ cos=cos,
+ sin=sin,
+ idx_ar=idx_ar,
+ )
+
+ if return_mid_results_layers is not None and layer_idx in return_mid_results_layers:
+ mid_results.append(hidden_states)
+
+ hidden_states = text_model.norm(hidden_states)
+ return hidden_states, mid_results
+
+ def _forward_generation(self, input_ids, position_ids, vinputs, timestep, token_types,
+ attention_mask=None, pixel_values=None, pixel_values_videos=None,
+ image_grid_thw=None, video_grid_thw=None,
+ return_mid_results_layers=None,
+ use_gradient_checkpointing=False,
+ use_gradient_checkpointing_offload=False,
+ **kwargs):
+ """Forward pass for image generation (denoising step)."""
+ inputs_embeds = self.get_input_embeddings()(input_ids)
+
+ if pixel_values is not None:
+ image_embeds, _ = self.get_image_features(pixel_values, image_grid_thw)
+ image_embeds = torch.cat(image_embeds, dim=0).to(inputs_embeds.device, inputs_embeds.dtype)
+ image_mask, _ = self.get_placeholder_mask(
+ input_ids, inputs_embeds=inputs_embeds, image_features=image_embeds)
+ inputs_embeds = inputs_embeds.masked_scatter(image_mask, image_embeds)
+ elif torch.is_grad_enabled():
+ # t2i task: no pixel_values, but we must run the vision encoder with a
+ # tiny dummy input so that EVERY rank has non-None (zero) gradients for
+ # vision-encoder parameters. This keeps the FSDP reduce-scatter and the
+ # replicate-group all-reduce symmetric across t2i and ref-task ranks,
+ # preventing collective hangs at backward / clip_grad_norm_.
+ # The dummy output is zeroed out before being added to inputs_embeds, so
+ # the forward result is numerically identical to the no-pixel_values path.
+ pe = self.visual.patch_embed
+ t_sz = pe.temporal_patch_size
+ m_sz = self.visual.spatial_merge_size
+ n_patches = t_sz * m_sz * m_sz
+ patch_dim = pe.in_channels * t_sz * pe.patch_size * pe.patch_size
+ fake_pv = torch.zeros(n_patches, patch_dim,
+ device=inputs_embeds.device,
+ dtype=pe.proj.weight.dtype)
+ fake_grid = torch.tensor([[t_sz, m_sz, m_sz]],
+ dtype=torch.long, device=inputs_embeds.device)
+ fake_embs, _ = self.get_image_features(fake_pv, fake_grid)
+ fake_embs = torch.cat(fake_embs, dim=0).to(inputs_embeds.dtype)
+ inputs_embeds = inputs_embeds + fake_embs.sum() * inputs_embeds.new_zeros([])
+
+ if pixel_values_videos is not None:
+ video_embeds, _ = self.get_video_features(pixel_values_videos, video_grid_thw)
+ video_embeds = torch.cat(video_embeds, dim=0).to(inputs_embeds.device, inputs_embeds.dtype)
+ _, video_mask = self.get_placeholder_mask(
+ input_ids, inputs_embeds=inputs_embeds, video_features=video_embeds)
+ inputs_embeds = inputs_embeds.masked_scatter(video_mask, video_embeds)
+
+ if isinstance(timestep, list):
+ timestep = torch.cat(timestep, dim=0)
+ timestep = timestep.to(inputs_embeds.device)
+ t_emb = self.t_embedder1(timestep)
+
+ tms_mask = (input_ids == self.tms_token_id)
+ tms_mask_3d = tms_mask.unsqueeze(-1).expand_as(inputs_embeds)
+ t_emb_expanded = t_emb.unsqueeze(1).expand_as(inputs_embeds)
+ inputs_embeds = torch.where(tms_mask_3d, t_emb_expanded, inputs_embeds)
+
+ if isinstance(vinputs, list):
+ vinputs = torch.cat(vinputs, dim=0)
+ vinputs = vinputs.to(inputs_embeds.device)
+ vinputs_embedded = self.x_embedder(vinputs).to(inputs_embeds.dtype)
+ inputs_embeds = torch.cat([inputs_embeds, vinputs_embedded], dim=1)
+
+ batch_size, total_seq_len, _ = inputs_embeds.shape
+
+ if isinstance(token_types, list):
+ token_types = torch.cat(token_types, dim=0)
+ token_types = token_types.to(inputs_embeds.device)
+ if token_types.dim() == 1:
+ token_types = token_types.unsqueeze(0)
+ elif token_types.dim() == 2 and token_types.shape[-1] == 1 and token_types.shape[0] == total_seq_len:
+ token_types = token_types.squeeze(-1).unsqueeze(0)
+ if token_types.shape[0] == 1 and batch_size > 1:
+ token_types = token_types.expand(batch_size, -1)
+
+ mid_results = None
+
+ hidden_states, mid_results = self._run_decoder_flash(
+ inputs_embeds, position_ids, token_types,
+ return_mid_results_layers=return_mid_results_layers,
+ use_gradient_checkpointing=use_gradient_checkpointing,
+ use_gradient_checkpointing_offload=use_gradient_checkpointing_offload,
+ )
+
+ x_pred = self.final_layer2(hidden_states)
+
+ return Qwen3VLModelOutputWithPast(
+ last_hidden_state=hidden_states,
+ x_pred=x_pred,
+ mid_results=mid_results,
+ )
+
+ def forward(
+ self,
+ input_ids: torch.LongTensor = None,
+ attention_mask: Optional[torch.Tensor] = None,
+ position_ids: Optional[torch.LongTensor] = None,
+ past_key_values: Optional[Cache] = None,
+ inputs_embeds: Optional[torch.FloatTensor] = None,
+ pixel_values: Optional[torch.Tensor] = None,
+ pixel_values_videos: Optional[torch.FloatTensor] = None,
+ image_grid_thw: Optional[torch.LongTensor] = None,
+ video_grid_thw: Optional[torch.LongTensor] = None,
+ cache_position: Optional[torch.LongTensor] = None,
+ vinputs: Optional[torch.Tensor] = None,
+ timestep: Optional[torch.Tensor] = None,
+ token_types: Optional[torch.Tensor] = None,
+ return_mid_results_layers: Optional[list] = None,
+ use_gradient_checkpointing: bool = False,
+ use_gradient_checkpointing_offload: bool = False,
+ **kwargs,
+ ) -> Union[tuple, Qwen3VLModelOutputWithPast]:
+ if vinputs is not None:
+ return self._forward_generation(
+ input_ids=input_ids, position_ids=position_ids,
+ vinputs=vinputs, timestep=timestep, token_types=token_types,
+ attention_mask=attention_mask,
+ pixel_values=pixel_values, pixel_values_videos=pixel_values_videos,
+ image_grid_thw=image_grid_thw, video_grid_thw=video_grid_thw,
+ return_mid_results_layers=return_mid_results_layers,
+ use_gradient_checkpointing=use_gradient_checkpointing,
+ use_gradient_checkpointing_offload=use_gradient_checkpointing_offload,
+ **kwargs)
+
+ if (input_ids is None) ^ (inputs_embeds is not None):
+ raise ValueError("You must specify exactly one of input_ids or inputs_embeds")
+
+ if inputs_embeds is None:
+ inputs_embeds = self.get_input_embeddings()(input_ids)
+
+ image_mask = None
+ video_mask = None
+
+ if pixel_values is not None:
+ image_embeds, deepstack_image_embeds = self.get_image_features(pixel_values, image_grid_thw)
+ image_embeds = torch.cat(image_embeds, dim=0).to(inputs_embeds.device, inputs_embeds.dtype)
+ image_mask, _ = self.get_placeholder_mask(
+ input_ids, inputs_embeds=inputs_embeds, image_features=image_embeds
+ )
+ inputs_embeds = inputs_embeds.masked_scatter(image_mask, image_embeds)
+
+ if pixel_values_videos is not None:
+ video_embeds, deepstack_video_embeds = self.get_video_features(pixel_values_videos, video_grid_thw)
+ video_embeds = torch.cat(video_embeds, dim=0).to(inputs_embeds.device, inputs_embeds.dtype)
+ _, video_mask = self.get_placeholder_mask(
+ input_ids, inputs_embeds=inputs_embeds, video_features=video_embeds
+ )
+ inputs_embeds = inputs_embeds.masked_scatter(video_mask, video_embeds)
+
+ visual_pos_masks = None
+ deepstack_visual_embeds = None
+ if image_mask is not None and video_mask is not None:
+ image_mask = image_mask[..., 0]
+ video_mask = video_mask[..., 0]
+ visual_pos_masks = image_mask | video_mask
+ deepstack_visual_embeds = []
+ image_mask_joint = image_mask[visual_pos_masks]
+ video_mask_joint = video_mask[visual_pos_masks]
+ for img_embed, vid_embed in zip(deepstack_image_embeds, deepstack_video_embeds):
+ embed_joint = img_embed.new_zeros(visual_pos_masks.sum(), img_embed.shape[-1]).to(img_embed.device)
+ embed_joint[image_mask_joint, :] = img_embed
+ embed_joint[video_mask_joint, :] = vid_embed
+ deepstack_visual_embeds.append(embed_joint)
+ elif image_mask is not None:
+ image_mask = image_mask[..., 0]
+ visual_pos_masks = image_mask
+ deepstack_visual_embeds = deepstack_image_embeds
+ elif video_mask is not None:
+ video_mask = video_mask[..., 0]
+ visual_pos_masks = video_mask
+ deepstack_visual_embeds = deepstack_video_embeds
+
+ if position_ids is None:
+ attention_mask_tensor = (
+ attention_mask if not isinstance(attention_mask, dict) else attention_mask["full_attention"]
+ )
+ if attention_mask_tensor is not None and attention_mask_tensor.ndim == 4:
+ attention_mask_tensor = torch.diagonal(attention_mask_tensor[:, 0], dim1=1, dim2=2)
+ if attention_mask_tensor.dtype.is_floating_point:
+ attention_mask_tensor = attention_mask_tensor / torch.finfo(attention_mask_tensor.dtype).min
+ attention_mask_tensor = (1.0 - attention_mask_tensor).int()
+
+ if (cache_position is not None and cache_position[0] == 0) or (past_key_values is None or past_key_values.get_seq_length() == 0):
+ position_ids, rope_deltas = self.get_rope_index(
+ input_ids, image_grid_thw, video_grid_thw, attention_mask
+ )
+ self.rope_deltas = rope_deltas
+ else:
+ q_len = inputs_embeds.shape[1]
+ position_ids = torch.arange(q_len, device=inputs_embeds.device)
+ position_ids = position_ids.view(1, -1).expand(inputs_embeds.shape[0], -1)
+ position_ids = position_ids.unsqueeze(0).expand(3, -1, -1)
+
+ outputs = self.language_model(
+ input_ids=None,
+ position_ids=position_ids,
+ attention_mask=attention_mask,
+ past_key_values=past_key_values,
+ inputs_embeds=inputs_embeds,
+ cache_position=cache_position,
+ visual_pos_masks=visual_pos_masks,
+ deepstack_visual_embeds=deepstack_visual_embeds,
+ use_gradient_checkpointing=use_gradient_checkpointing,
+ use_gradient_checkpointing_offload=use_gradient_checkpointing_offload,
+ **kwargs,
+ )
+
+ return Qwen3VLModelOutputWithPast(
+ last_hidden_state=outputs.last_hidden_state,
+ past_key_values=outputs.past_key_values,
+ rope_deltas=self.rope_deltas,
+ )
+
+ def get_rope_index(
+ self,
+ input_ids: Optional[torch.LongTensor] = None,
+ image_grid_thw: Optional[torch.LongTensor] = None,
+ video_grid_thw: Optional[torch.LongTensor] = None,
+ attention_mask: Optional[torch.Tensor] = None,
+ ) -> tuple[torch.Tensor, torch.Tensor]:
+ if video_grid_thw is not None:
+ video_grid_thw = torch.repeat_interleave(video_grid_thw, video_grid_thw[:, 0], dim=0)
+ video_grid_thw[:, 0] = 1
+
+ spatial_merge_size = self.config.vision_config.spatial_merge_size
+ image_token_id = self.config.image_token_id
+ video_token_id = self.config.video_token_id
+ vision_start_token_id = self.config.vision_start_token_id
+ mrope_position_deltas = []
+ if input_ids is not None and (image_grid_thw is not None or video_grid_thw is not None):
+ total_input_ids = input_ids
+ if attention_mask is None:
+ attention_mask = torch.ones_like(total_input_ids)
+ position_ids = torch.ones(
+ 3, input_ids.shape[0], input_ids.shape[1],
+ dtype=input_ids.dtype, device=input_ids.device,
+ )
+ image_index, video_index = 0, 0
+ attention_mask = attention_mask.to(total_input_ids.device)
+ for i, input_ids_i in enumerate(total_input_ids):
+ input_ids_i = input_ids_i[attention_mask[i] == 1]
+ image_nums, video_nums = 0, 0
+ vision_start_indices = torch.argwhere(input_ids_i == vision_start_token_id).squeeze(1)
+ vision_tokens = input_ids_i[vision_start_indices + 1]
+ image_nums = (vision_tokens == image_token_id).sum()
+ video_nums = (vision_tokens == video_token_id).sum()
+ input_tokens = input_ids_i.tolist()
+ llm_pos_ids_list: list = []
+ st = 0
+ remain_images, remain_videos = image_nums, video_nums
+ for _ in range(image_nums + video_nums):
+ if image_token_id in input_tokens and remain_images > 0:
+ ed_image = input_tokens.index(image_token_id, st)
+ else:
+ ed_image = len(input_tokens) + 1
+ if video_token_id in input_tokens and remain_videos > 0:
+ ed_video = input_tokens.index(video_token_id, st)
+ else:
+ ed_video = len(input_tokens) + 1
+ if ed_image < ed_video:
+ t, h, w = (
+ image_grid_thw[image_index][0],
+ image_grid_thw[image_index][1],
+ image_grid_thw[image_index][2],
+ )
+ image_index += 1
+ remain_images -= 1
+ ed = ed_image
+ else:
+ t, h, w = (
+ video_grid_thw[video_index][0],
+ video_grid_thw[video_index][1],
+ video_grid_thw[video_index][2],
+ )
+ video_index += 1
+ remain_videos -= 1
+ ed = ed_video
+ llm_grid_t, llm_grid_h, llm_grid_w = (
+ t.item(),
+ h.item() // spatial_merge_size,
+ w.item() // spatial_merge_size,
+ )
+ text_len = ed - st
+ st_idx = llm_pos_ids_list[-1].max() + 1 if len(llm_pos_ids_list) > 0 else 0
+ llm_pos_ids_list.append(torch.arange(text_len).view(1, -1).expand(3, -1) + st_idx)
+ t_index = torch.arange(llm_grid_t).view(-1, 1).expand(-1, llm_grid_h * llm_grid_w).flatten()
+ h_index = torch.arange(llm_grid_h).view(1, -1, 1).expand(llm_grid_t, -1, llm_grid_w).flatten()
+ w_index = torch.arange(llm_grid_w).view(1, 1, -1).expand(llm_grid_t, llm_grid_h, -1).flatten()
+ llm_pos_ids_list.append(torch.stack([t_index, h_index, w_index]) + text_len + st_idx)
+ st = ed + llm_grid_t * llm_grid_h * llm_grid_w
+
+ if st < len(input_tokens):
+ st_idx = llm_pos_ids_list[-1].max() + 1 if len(llm_pos_ids_list) > 0 else 0
+ text_len = len(input_tokens) - st
+ llm_pos_ids_list.append(torch.arange(text_len).view(1, -1).expand(3, -1) + st_idx)
+
+ llm_positions = torch.cat(llm_pos_ids_list, dim=1).reshape(3, -1)
+ position_ids[..., i, attention_mask[i] == 1] = llm_positions.to(position_ids.device)
+ mrope_position_deltas.append(llm_positions.max() + 1 - len(total_input_ids[i]))
+ mrope_position_deltas = torch.tensor(mrope_position_deltas, device=input_ids.device).unsqueeze(1)
+ return position_ids, mrope_position_deltas
+ else:
+ if attention_mask is not None:
+ position_ids = attention_mask.long().cumsum(-1) - 1
+ position_ids.masked_fill_(attention_mask == 0, 1)
+ position_ids = position_ids.unsqueeze(0).expand(3, -1, -1).to(attention_mask.device)
+ max_position_ids = position_ids.max(0, keepdim=False)[0].max(-1, keepdim=True)[0]
+ mrope_position_deltas = max_position_ids + 1 - attention_mask.shape[-1]
+ else:
+ position_ids = (
+ torch.arange(input_ids.shape[1], device=input_ids.device)
+ .view(1, 1, -1)
+ .expand(3, input_ids.shape[0], -1)
+ )
+ mrope_position_deltas = torch.zeros(
+ [input_ids.shape[0], 1], device=input_ids.device, dtype=input_ids.dtype,
+ )
+ return position_ids, mrope_position_deltas
+
+
+class Qwen3VLTextModel(Qwen3VLPreTrainedModel):
+ config: Qwen3VLTextConfig
+
+ def __init__(self, config):
+ super().__init__(config)
+ self.vocab_size = config.vocab_size
+
+ self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size)
+ self.layers = nn.ModuleList(
+ [Qwen3VLTextDecoderLayer(config, layer_idx) for layer_idx in range(config.num_hidden_layers)]
+ )
+ self.norm = Qwen3VLTextRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+ self.rotary_emb = Qwen3VLRotaryEmbedding(config=config)
+
+ self.post_init()
+
+ def get_input_embeddings(self):
+ return self.embed_tokens
+
+ def set_input_embeddings(self, value):
+ self.embed_tokens = value
+
+ def forward(
+ self,
+ input_ids: torch.LongTensor = None,
+ position_ids: Optional[torch.LongTensor] = None,
+ attention_mask: Optional[torch.Tensor] = None,
+ past_key_values: Optional[Cache] = None,
+ inputs_embeds: Optional[torch.FloatTensor] = None,
+ cache_position: Optional[torch.LongTensor] = None,
+ visual_pos_masks: Optional[torch.Tensor] = None,
+ deepstack_visual_embeds: Optional[torch.FloatTensor] = None,
+ use_cache: Optional[bool] = None,
+ return_mid_results_layers: Optional[list] = None,
+ use_gradient_checkpointing: bool = False,
+ use_gradient_checkpointing_offload: bool = False,
+ **kwargs,
+ ) -> BaseModelOutputWithPast:
+ if (input_ids is None) ^ (inputs_embeds is not None):
+ raise ValueError("You must specify exactly one of input_ids or inputs_embeds")
+
+ if inputs_embeds is None:
+ inputs_embeds = self.embed_tokens(input_ids)
+
+ if use_cache is None:
+ use_cache = self.config.use_cache
+
+ if position_ids is None:
+ position_ids = torch.arange(
+ inputs_embeds.shape[1], device=inputs_embeds.device
+ ).unsqueeze(0).expand(inputs_embeds.shape[0], -1)
+
+ position_embeddings = self.rotary_emb(inputs_embeds, position_ids)
+ hidden_states = inputs_embeds
+ mid_results = [] if return_mid_results_layers else None
+
+ for layer_idx, decoder_layer in enumerate(self.layers):
+ hidden_states = gradient_checkpoint_forward(
+ decoder_layer,
+ use_gradient_checkpointing=use_gradient_checkpointing,
+ use_gradient_checkpointing_offload=False,
+ hidden_states=hidden_states,
+ position_embeddings=position_embeddings,
+ attention_mask=attention_mask,
+ past_key_values=past_key_values,
+ use_cache=use_cache,
+ cache_position=cache_position,
+ )
+
+ if return_mid_results_layers is not None and layer_idx in return_mid_results_layers:
+ mid_results.append(hidden_states)
+
+ hidden_states = self.norm(hidden_states)
+
+ output = BaseModelOutputWithPast(
+ last_hidden_state=hidden_states,
+ past_key_values=past_key_values,
+ )
+ if return_mid_results_layers is not None:
+ output.mid_results = mid_results
+ return output
+
+
+class Qwen3VLVisionModel(Qwen3VLPreTrainedModel):
+ config: Qwen3VLVisionConfig
+
+ def __init__(self, config):
+ super().__init__(config)
+ self.spatial_merge_size = config.spatial_merge_size
+ self.patch_size = config.patch_size
+ self.spatial_merge_unit = self.spatial_merge_size * self.spatial_merge_size
+
+ self.patch_embed = Qwen3VLVisionPatchEmbed(config)
+
+ self.pos_embed = nn.Embedding(config.num_position_embeddings, config.hidden_size)
+ self.num_grid_per_side = int(config.num_position_embeddings**0.5)
+
+ head_dim = config.hidden_size // config.num_heads
+ self.rotary_pos_emb = Qwen3VLVisionRotaryEmbedding(head_dim // 2)
+
+ self.blocks = nn.ModuleList(
+ [Qwen3VLVisionBlock(config) for _ in range(config.depth)]
+ )
+
+ self.merger = Qwen3VLVisionPatchMerger(config, use_postshuffle_norm=False)
+ self.deepstack_visual_indexes = config.deepstack_visual_indexes
+ self.deepstack_merger_list = nn.ModuleList(
+ [
+ Qwen3VLVisionPatchMerger(
+ config=config,
+ use_postshuffle_norm=True,
+ )
+ for _ in range(len(config.deepstack_visual_indexes))
+ ]
+ )
+ self.post_init()
+
+ def get_dtype(self) -> torch.dtype:
+ return self.blocks[0].mlp.linear_fc1.weight.dtype
+
+ def fast_pos_embed_interpolate(self, grid_thw):
+ grid_ts, grid_hs, grid_ws = grid_thw[:, 0], grid_thw[:, 1], grid_thw[:, 2]
+
+ idx_list = [[] for _ in range(4)]
+ weight_list = [[] for _ in range(4)]
+
+ for t, h, w in zip(grid_ts, grid_hs, grid_ws):
+ h_idxs = torch.linspace(0, self.num_grid_per_side - 1, h, device=self.pos_embed.weight.device)
+ w_idxs = torch.linspace(0, self.num_grid_per_side - 1, w, device=self.pos_embed.weight.device)
+
+ h_idxs_floor = h_idxs.int()
+ w_idxs_floor = w_idxs.int()
+ h_idxs_ceil = (h_idxs.int() + 1).clip(max=self.num_grid_per_side - 1)
+ w_idxs_ceil = (w_idxs.int() + 1).clip(max=self.num_grid_per_side - 1)
+
+ dh = h_idxs - h_idxs_floor
+ dw = w_idxs - w_idxs_floor
+
+ base_h = h_idxs_floor * self.num_grid_per_side
+ base_h_ceil = h_idxs_ceil * self.num_grid_per_side
+
+ indices = [
+ (base_h[None].T + w_idxs_floor[None]).flatten(),
+ (base_h[None].T + w_idxs_ceil[None]).flatten(),
+ (base_h_ceil[None].T + w_idxs_floor[None]).flatten(),
+ (base_h_ceil[None].T + w_idxs_ceil[None]).flatten(),
+ ]
+
+ weights = [
+ ((1 - dh)[None].T * (1 - dw)[None]).flatten(),
+ ((1 - dh)[None].T * dw[None]).flatten(),
+ (dh[None].T * (1 - dw)[None]).flatten(),
+ (dh[None].T * dw[None]).flatten(),
+ ]
+
+ for i in range(4):
+ idx_list[i].extend(indices[i].tolist())
+ weight_list[i].extend(weights[i].tolist())
+
+ idx_tensor = torch.tensor(idx_list, dtype=torch.long, device=self.pos_embed.weight.device)
+ weight_tensor = torch.tensor(
+ weight_list, dtype=self.pos_embed.weight.dtype, device=self.pos_embed.weight.device
+ )
+ pos_embeds = self.pos_embed(idx_tensor) * weight_tensor[:, :, None]
+ patch_pos_embeds = pos_embeds[0] + pos_embeds[1] + pos_embeds[2] + pos_embeds[3]
+
+ patch_pos_embeds = patch_pos_embeds.split([h * w for h, w in zip(grid_hs, grid_ws)])
+
+ patch_pos_embeds_permute = []
+ merge_size = self.spatial_merge_size
+ for pos_embed, t, h, w in zip(patch_pos_embeds, grid_ts, grid_hs, grid_ws):
+ pos_embed = pos_embed.repeat(t, 1)
+ pos_embed = (
+ pos_embed.view(t, h // merge_size, merge_size, w // merge_size, merge_size, -1)
+ .permute(0, 1, 3, 2, 4, 5)
+ .flatten(0, 4)
+ )
+ patch_pos_embeds_permute.append(pos_embed)
+ patch_pos_embeds = torch.cat(patch_pos_embeds_permute)
+ return patch_pos_embeds
+
+ def forward(
+ self,
+ pixel_values: torch.FloatTensor,
+ grid_thw: Optional[torch.LongTensor] = None,
+ ) -> tuple[torch.Tensor, list]:
+ hidden_states = self.patch_embed(pixel_values)
+
+ pos_embeds = self.fast_pos_embed_interpolate(grid_thw)
+ hidden_states = hidden_states + pos_embeds
+
+ rotary_pos_emb = self.rot_pos_emb(grid_thw)
+
+ seq_len, _ = hidden_states.size()
+ hidden_states = hidden_states.reshape(seq_len, -1)
+ rotary_pos_emb = rotary_pos_emb.reshape(seq_len, -1)
+ emb = torch.cat((rotary_pos_emb, rotary_pos_emb), dim=-1)
+ position_embeddings = (emb.cos(), emb.sin())
+
+ cu_seqlens = torch.repeat_interleave(grid_thw[:, 1] * grid_thw[:, 2], grid_thw[:, 0]).cumsum(
+ dim=0,
+ dtype=grid_thw.dtype if torch.jit.is_tracing() else torch.int32,
+ )
+ cu_seqlens = F.pad(cu_seqlens, (1, 0), value=0)
+
+ deepstack_feature_lists = []
+ for layer_num, blk in enumerate(self.blocks):
+ hidden_states = blk(
+ hidden_states,
+ cu_seqlens=cu_seqlens,
+ position_embeddings=position_embeddings,
+ )
+ if layer_num in self.deepstack_visual_indexes:
+ deepstack_feature = self.deepstack_merger_list[self.deepstack_visual_indexes.index(layer_num)](
+ hidden_states
+ )
+ deepstack_feature_lists.append(deepstack_feature)
+
+ hidden_states = self.merger(hidden_states)
+
+ return hidden_states, deepstack_feature_lists
+
+ def rot_pos_emb(self, grid_thw: torch.Tensor) -> torch.Tensor:
+ merge_size = self.spatial_merge_size
+
+ max_hw = int(grid_thw[:, 1:].max().item())
+ freq_table = self.rotary_pos_emb(max_hw) # (max_hw, dim // 2)
+ device = freq_table.device
+
+ total_tokens = int(torch.prod(grid_thw, dim=1).sum().item())
+ pos_ids = torch.empty((total_tokens, 2), dtype=torch.long, device=device)
+
+ offset = 0
+ for num_frames, height, width in grid_thw:
+ merged_h, merged_w = height // merge_size, width // merge_size
+
+ block_rows = torch.arange(merged_h, device=device)
+ block_cols = torch.arange(merged_w, device=device)
+ intra_row = torch.arange(merge_size, device=device)
+ intra_col = torch.arange(merge_size, device=device)
+
+ # Compute full-resolution positions
+ row_idx = block_rows[:, None, None, None] * merge_size + intra_row[None, None, :, None]
+ col_idx = block_cols[None, :, None, None] * merge_size + intra_col[None, None, None, :]
+
+ row_idx = row_idx.expand(merged_h, merged_w, merge_size, merge_size).reshape(-1)
+ col_idx = col_idx.expand(merged_h, merged_w, merge_size, merge_size).reshape(-1)
+
+ coords = torch.stack((row_idx, col_idx), dim=-1)
+
+ if num_frames > 1:
+ coords = coords.repeat(num_frames, 1)
+
+ num_tokens = coords.shape[0]
+ pos_ids[offset : offset + num_tokens] = coords
+ offset += num_tokens
+
+ embeddings = freq_table[pos_ids] # lookup rotary embeddings
+ embeddings = embeddings.flatten(1)
+ return embeddings
+
+
+class Qwen3VLRotaryEmbedding(nn.Module):
+ def __init__(self, config, device=None):
+ super().__init__()
+ if hasattr(config, "rope_scaling") and config.rope_scaling is not None:
+ self.rope_type = config.rope_scaling.get("rope_type", "default")
+ else:
+ self.rope_type = "default"
+ self.max_seq_len_cached = config.max_position_embeddings
+ self.original_max_seq_len = config.max_position_embeddings
+ self.config = config
+ self.rope_init_fn = ROPE_INIT_FUNCTIONS[self.rope_type]
+ inv_freq, self.attention_scaling = self.rope_init_fn(self.config, device)
+ self.register_buffer("inv_freq", inv_freq.detach().clone(), persistent=False)
+ self.original_inv_freq = self.inv_freq
+ self.mrope_section = config.rope_scaling.get("mrope_section", [24, 20, 20])
+
+ @staticmethod
+ def compute_default_rope_parameters(config, device=None, seq_len=None):
+ return _compute_default_rope_parameters(config, device=None, seq_len=None)
+
+ @staticmethod
+ def apply_interleaved_mrope(freqs, mrope_section):
+ freqs_t = freqs[0]
+ for dim, offset in enumerate((1, 2), start=1):
+ length = mrope_section[dim] * 3
+ idx = slice(offset, length, 3)
+ freqs_t[..., idx] = freqs[dim, ..., idx]
+ return freqs_t
+
+ def forward(self, x, position_ids):
+ if position_ids.ndim == 2:
+ position_ids = position_ids[None, ...].expand(3, position_ids.shape[0], -1)
+ inv_freq_expanded = self.original_inv_freq[None, None, :, None].float().to(device=x.device).expand(3, position_ids.shape[1], -1, 1)
+ position_ids_expanded = position_ids[:, :, None, :].float()
+ device_type = x.device.type if isinstance(x.device.type, str) and x.device.type != "mps" else "cpu"
+ with torch.autocast(device_type=device_type, enabled=False):
+ freqs = (inv_freq_expanded.float() @ position_ids_expanded.float()).transpose(2, 3)
+ freqs = self.apply_interleaved_mrope(freqs, self.mrope_section)
+ emb = torch.cat((freqs, freqs), dim=-1)
+ cos = emb.cos() * self.attention_scaling
+ sin = emb.sin() * self.attention_scaling
+ return cos.to(x.dtype), sin.to(x.dtype)
+
+
+@dataclass
+class Qwen3VLCausalLMOutputWithPast(ModelOutput):
+ loss: Optional[torch.FloatTensor] = None
+ logits: Optional[torch.FloatTensor] = None
+ past_key_values: Optional[Cache] = None
+ hidden_states: Optional[tuple[torch.FloatTensor]] = None
+ attentions: Optional[tuple[torch.FloatTensor]] = None
+ rope_deltas: Optional[torch.LongTensor] = None
+ x_pred: Optional[torch.FloatTensor] = None
+ mid_results: Optional[list] = None
+
+
+def _build_hidream_config():
+ text_config = Qwen3VLTextConfig(
+ hidden_size=4096,
+ num_hidden_layers=36,
+ num_attention_heads=32,
+ num_key_value_heads=8,
+ intermediate_size=12288,
+ vocab_size=151936,
+ max_position_embeddings=262144,
+ head_dim=128,
+ attention_bias=False,
+ hidden_act="silu",
+ rms_norm_eps=1e-6,
+ use_cache=True,
+ bos_token_id=151643,
+ eos_token_id=151645,
+ rope_theta=5000000,
+ attention_dropout=0.0,
+ initializer_range=0.02,
+ rope_scaling={"rope_type": "default", "mrope_section": [24, 20, 20], "mrope_interleaved": True},
+ )
+ vision_config = Qwen3VLVisionConfig(
+ hidden_size=1152,
+ depth=27,
+ num_heads=16,
+ intermediate_size=4304,
+ patch_size=16,
+ spatial_merge_size=2,
+ in_channels=3,
+ out_hidden_size=4096,
+ deepstack_visual_indexes=[8, 16, 24],
+ temporal_patch_size=2,
+ num_position_embeddings=2304,
+ hidden_act="gelu_pytorch_tanh",
+ initializer_range=0.02,
+ )
+ config = Qwen3VLConfig()
+ config.text_config = text_config
+ config.vision_config = vision_config
+ config.image_token_id = 151655
+ config.video_token_id = 151656
+ config.vision_start_token_id = 151652
+ config.vision_end_token_id = 151653
+ return config
+
+
+class HiDreamO1ImageModel(Qwen3VLPreTrainedModel):
+ _checkpoint_conversion_mapping = {}
+ _tied_weights_keys = ["lm_head.weight"]
+ accepts_loss_kwargs = False
+ config: Qwen3VLConfig
+
+ def __init__(self):
+ config = _build_hidream_config()
+ super().__init__(config)
+ self.model = Qwen3VLModel(config)
+ self.lm_head = nn.Linear(config.text_config.hidden_size, config.text_config.vocab_size, bias=False)
+ self.post_init()
+
+ def get_input_embeddings(self):
+ return self.model.get_input_embeddings()
+
+ def set_input_embeddings(self, value):
+ self.model.set_input_embeddings(value)
+
+ def set_decoder(self, decoder):
+ self.model.set_decoder(decoder)
+
+ def get_decoder(self):
+ return self.model.get_decoder()
+
+ @property
+ def language_model(self):
+ return self.model.language_model
+
+ @property
+ def visual(self):
+ return self.model.visual
+
+ def forward(
+ self,
+ input_ids: torch.LongTensor = None,
+ attention_mask: Optional[torch.Tensor] = None,
+ position_ids: Optional[torch.LongTensor] = None,
+ past_key_values: Optional[Cache] = None,
+ inputs_embeds: Optional[torch.FloatTensor] = None,
+ labels: Optional[torch.LongTensor] = None,
+ pixel_values: Optional[torch.Tensor] = None,
+ pixel_values_videos: Optional[torch.FloatTensor] = None,
+ image_grid_thw: Optional[torch.LongTensor] = None,
+ video_grid_thw: Optional[torch.LongTensor] = None,
+ cache_position: Optional[torch.LongTensor] = None,
+ logits_to_keep: Union[int, torch.Tensor] = 0,
+ vinputs: Optional[torch.Tensor] = None,
+ timestep: Optional[torch.Tensor] = None,
+ token_types: Optional[torch.Tensor] = None,
+ return_mid_results_layers: Optional[list] = None,
+ use_gradient_checkpointing: bool = False,
+ use_gradient_checkpointing_offload: bool = False,
+ **kwargs,
+ ) -> Union[tuple, Qwen3VLCausalLMOutputWithPast]:
+ outputs = self.model(
+ input_ids=input_ids,
+ pixel_values=pixel_values,
+ pixel_values_videos=pixel_values_videos,
+ image_grid_thw=image_grid_thw,
+ video_grid_thw=video_grid_thw,
+ position_ids=position_ids,
+ attention_mask=attention_mask,
+ past_key_values=past_key_values,
+ inputs_embeds=inputs_embeds,
+ cache_position=cache_position,
+ vinputs=vinputs,
+ timestep=timestep,
+ token_types=token_types,
+ return_mid_results_layers=return_mid_results_layers,
+ use_gradient_checkpointing=use_gradient_checkpointing,
+ use_gradient_checkpointing_offload=use_gradient_checkpointing_offload,
+ **kwargs,
+ )
+
+ if vinputs is not None:
+ return Qwen3VLCausalLMOutputWithPast(
+ x_pred=outputs.x_pred,
+ mid_results=outputs.mid_results if hasattr(outputs, 'mid_results') else None,
+ )
+
+ hidden_states = outputs[0]
+ slice_indices = slice(-logits_to_keep, None) if isinstance(logits_to_keep, int) else logits_to_keep
+ logits = self.lm_head(hidden_states[:, slice_indices, :])
+
+ loss = None
+ if labels is not None:
+ loss = self.loss_function(logits=logits, labels=labels, vocab_size=self.config.text_config.vocab_size)
+
+ return Qwen3VLCausalLMOutputWithPast(
+ loss=loss,
+ logits=logits,
+ past_key_values=outputs.past_key_values,
+ rope_deltas=outputs.rope_deltas,
+ )
diff --git a/diffsynth/models/hpsv2.py b/diffsynth/models/hpsv2.py
new file mode 100644
index 0000000000000000000000000000000000000000..c8cc4c924983e370a843761422fe4b403b112012
--- /dev/null
+++ b/diffsynth/models/hpsv2.py
@@ -0,0 +1,92 @@
+from typing import Union
+import torch
+from PIL import Image
+
+ImageInput = Union[Image.Image, list[Image.Image], tuple[Image.Image, ...]]
+
+def _as_list(value):
+ if isinstance(value, (list, tuple)):
+ return list(value)
+ return [value]
+
+def _feature_tensor(output, feature_name: str):
+ if torch.is_tensor(output):
+ return output
+ for name in ("image_embeds", "text_embeds", "pooler_output"):
+ value = getattr(output, name, None)
+ if torch.is_tensor(value):
+ return value
+ if isinstance(output, (list, tuple)):
+ for value in output:
+ if torch.is_tensor(value):
+ return value
+ raise TypeError(f"{feature_name} must be a tensor or a model output with projected features.")
+
+
+class HPSv2Model(torch.nn.Module):
+ def __init__(self, model: torch.nn.Module, processor):
+ super().__init__()
+ self.model = model
+ self.processor = processor
+
+ @property
+ def device(self):
+ return next(self.parameters(), torch.tensor([])).device
+
+ @property
+ def dtype(self):
+ return next(self.parameters(), torch.tensor(0.0)).dtype
+
+ def _normalize_inputs(self, prompts, images):
+ images = _as_list(images)
+ prompts = _as_list(prompts)
+
+ if len(prompts) == 1 and len(images) > 1:
+ prompts = prompts * len(images)
+ if len(images) == 1 and len(prompts) > 1:
+ images = images * len(prompts)
+
+ if len(prompts) != len(images):
+ raise ValueError(f"Expected the same number of prompts and images, got {len(prompts)} and {len(images)}.")
+ return prompts, images
+
+ @torch.no_grad()
+ def forward(self, prompts: Union[str, list[str]], images: ImageInput):
+ prompts, images = self._normalize_inputs(prompts, images)
+ images = [image.convert("RGB") for image in images]
+
+ inputs = self.processor(
+ text=prompts,
+ images=images,
+ padding=True,
+ truncation=True,
+ return_tensors="pt"
+ ).to(self.device)
+
+ if self.dtype != torch.float32:
+ inputs = {
+ name: (
+ value.to(dtype=self.dtype)
+ if torch.is_tensor(value) and torch.is_floating_point(value)
+ else value
+ )
+ for name, value in inputs.items()
+ }
+
+ image_features = _feature_tensor(
+ self.model.get_image_features(pixel_values=inputs["pixel_values"]),
+ "image_features",
+ )
+ text_features = _feature_tensor(
+ self.model.get_text_features(input_ids=inputs["input_ids"], attention_mask=inputs.get("attention_mask")),
+ "text_features",
+ )
+
+ image_features = torch.nn.functional.normalize(image_features, dim=-1)
+ text_features = torch.nn.functional.normalize(text_features, dim=-1)
+
+ scores = (image_features * text_features).sum(dim=-1)
+ if hasattr(self.model, "logit_scale"):
+ scores = self.model.logit_scale.exp() * scores
+
+ return scores
\ No newline at end of file
diff --git a/diffsynth/models/hpsv3.py b/diffsynth/models/hpsv3.py
new file mode 100644
index 0000000000000000000000000000000000000000..b22966b6dadf62eee27f4c6cda0a8183d3c5c34e
--- /dev/null
+++ b/diffsynth/models/hpsv3.py
@@ -0,0 +1,353 @@
+import math
+from typing import Optional, Union
+import torch
+from PIL import Image
+from transformers import Qwen2VLForConditionalGeneration
+
+ImageInput = Union[Image.Image, list[Image.Image], tuple[Image.Image, ...]]
+
+HPSV3_INSTRUCTION = """
+You are tasked with evaluating a generated image based on Visual Quality and Text Alignment and give a overall score to estimate the human preference. Please provide a rating from 0 to 10, with 0 being the worst and 10 being the best.
+
+**Visual Quality:**
+Evaluate the overall visual quality of the image. The following sub-dimensions should be considered:
+- **Reasonableness:** The image should not contain any significant biological or logical errors, such as abnormal body structures or nonsensical environmental setups.
+- **Clarity:** Evaluate the sharpness and visibility of the image. The image should be clear and easy to interpret, with no blurring or indistinct areas.
+- **Detail Richness:** Consider the level of detail in textures, materials, lighting, and other visual elements (e.g., hair, clothing, shadows).
+- **Aesthetic and Creativity:** Assess the artistic aspects of the image, including the color scheme, composition, atmosphere, depth of field, and the overall creative appeal. The scene should convey a sense of harmony and balance.
+- **Safety:** The image should not contain harmful or inappropriate content, such as political, violent, or adult material. If such content is present, the image quality and satisfaction score should be the lowest possible.
+
+**Text Alignment:**
+Assess how well the image matches the textual prompt across the following sub-dimensions:
+- **Subject Relevance** Evaluate how accurately the subject(s) in the image (e.g., person, animal, object) align with the textual description. The subject should match the description in terms of number, appearance, and behavior.
+- **Style Relevance:** If the prompt specifies a particular artistic or stylistic style, evaluate how well the image adheres to this style.
+- **Contextual Consistency**: Assess whether the background, setting, and surrounding elements in the image logically fit the scenario described in the prompt. The environment should support and enhance the subject without contradictions.
+- **Attribute Fidelity**: Check if specific attributes mentioned in the prompt (e.g., colors, clothing, accessories, expressions, actions) are faithfully represented in the image. Minor deviations may be acceptable, but critical attributes should be preserved.
+- **Semantic Coherence**: Evaluate whether the overall meaning and intent of the prompt are captured in the image. The generated content should not introduce elements that conflict with or distort the original description.
+Textual prompt - {text_prompt}
+
+
+"""
+
+HPSV3_PROMPT_WITH_SPECIAL_TOKEN = """
+Please provide the overall ratings of this image: <|Reward|>
+
+END
+"""
+
+HPSV3_PROMPT_WITHOUT_SPECIAL_TOKEN = """
+Please provide the overall ratings of this image:
+"""
+
+def _as_list(value):
+ if isinstance(value, (list, tuple)):
+ return list(value)
+ return [value]
+
+def _round_by_factor(number, factor):
+ return round(number / factor) * factor
+
+def _ceil_by_factor(number, factor):
+ return math.ceil(number / factor) * factor
+
+def _floor_by_factor(number, factor):
+ return math.floor(number / factor) * factor
+
+def _smart_resize(height, width, factor=28, min_pixels=256 * 28 * 28, max_pixels=256 * 28 * 28):
+ if max(height, width) / min(height, width) > 200:
+ raise ValueError(f"Image aspect ratio must be smaller than 200, got {max(height, width) / min(height, width)}.")
+
+ h_bar = max(factor, _round_by_factor(height, factor))
+ w_bar = max(factor, _round_by_factor(width, factor))
+
+ if h_bar * w_bar > max_pixels:
+ beta = math.sqrt((height * width) / max_pixels)
+ h_bar = _floor_by_factor(height / beta, factor)
+ w_bar = _floor_by_factor(width / beta, factor)
+ elif h_bar * w_bar < min_pixels:
+ beta = math.sqrt(min_pixels / (height * width))
+ h_bar = _ceil_by_factor(height * beta, factor)
+ w_bar = _ceil_by_factor(width * beta, factor)
+
+ return h_bar, w_bar
+
+
+class HPSv3RewardModelMixin:
+ def init_reward_head(
+ self,
+ output_dim=2,
+ reward_token="special",
+ special_token_ids=None,
+ rm_head_type="ranknet",
+ rm_head_kwargs=None,
+ ):
+ self.output_dim = output_dim
+ self.reward_token = "special" if special_token_ids is not None else reward_token
+ self.special_token_ids = special_token_ids
+
+ hidden_size = getattr(self.config, "hidden_size", None)
+ if hidden_size is None and hasattr(self.config, "text_config"):
+ hidden_size = self.config.text_config.hidden_size
+
+ if rm_head_type == "ranknet":
+ rm_head_kwargs = rm_head_kwargs or {}
+ hidden = rm_head_kwargs.get("hidden_size", 1024)
+ dropout = rm_head_kwargs.get("dropout", 0.05)
+ self.rm_head = torch.nn.Sequential(
+ torch.nn.Linear(hidden_size, hidden),
+ torch.nn.ReLU(),
+ torch.nn.Dropout(dropout),
+ torch.nn.Linear(hidden, 16),
+ torch.nn.ReLU(),
+ torch.nn.Linear(16, output_dim),
+ )
+ else:
+ self.rm_head = torch.nn.Linear(hidden_size, output_dim, bias=False)
+
+ self.rm_head.to(torch.float32)
+
+ def forward(
+ self,
+ input_ids: torch.LongTensor = None,
+ attention_mask: Optional[torch.Tensor] = None,
+ position_ids: Optional[torch.LongTensor] = None,
+ past_key_values: Optional[list[torch.FloatTensor]] = None,
+ inputs_embeds: Optional[torch.FloatTensor] = None,
+ labels: Optional[torch.LongTensor] = None,
+ use_cache: Optional[bool] = None,
+ output_attentions: Optional[bool] = None,
+ output_hidden_states: Optional[bool] = None,
+ return_dict: Optional[bool] = None,
+ pixel_values: Optional[torch.Tensor] = None,
+ pixel_values_videos: Optional[torch.FloatTensor] = None,
+ image_grid_thw: Optional[torch.LongTensor] = None,
+ video_grid_thw: Optional[torch.LongTensor] = None,
+ rope_deltas: Optional[torch.LongTensor] = None,
+ **kwargs,
+ ):
+ output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
+ output_hidden_states = output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
+ return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+
+ kwargs.pop("logits_to_keep", None)
+ mm_token_type_ids = kwargs.pop("mm_token_type_ids", None)
+
+ outputs = self.model(
+ input_ids=input_ids,
+ attention_mask=attention_mask,
+ position_ids=position_ids,
+ past_key_values=past_key_values,
+ inputs_embeds=inputs_embeds,
+ use_cache=use_cache,
+ pixel_values=pixel_values,
+ pixel_values_videos=pixel_values_videos,
+ image_grid_thw=image_grid_thw,
+ video_grid_thw=video_grid_thw,
+ rope_deltas=rope_deltas,
+ mm_token_type_ids=mm_token_type_ids,
+ output_attentions=output_attentions,
+ output_hidden_states=output_hidden_states,
+ return_dict=return_dict,
+ **kwargs,
+ )
+
+ hidden_states = outputs.last_hidden_state if hasattr(outputs, "last_hidden_state") else outputs[0]
+ logits = self.rm_head(hidden_states.to(next(self.rm_head.parameters()).dtype))
+
+ batch_size = input_ids.shape[0] if input_ids is not None else inputs_embeds.shape[0]
+ pad_token_id = getattr(self.config, "pad_token_id", None)
+ if pad_token_id is None and hasattr(self.config, "text_config"):
+ pad_token_id = getattr(self.config.text_config, "pad_token_id", None)
+
+ if pad_token_id is None and batch_size != 1:
+ raise ValueError("Cannot handle batch sizes > 1 if no padding token is defined.")
+
+ if pad_token_id is None:
+ sequence_lengths = -1
+ elif input_ids is not None:
+ sequence_lengths = torch.eq(input_ids, pad_token_id).int().argmax(-1) - 1
+ sequence_lengths = sequence_lengths % input_ids.shape[-1]
+ sequence_lengths = sequence_lengths.to(logits.device)
+ else:
+ sequence_lengths = -1
+
+ if self.reward_token == "last":
+ pooled_logits = logits[torch.arange(batch_size, device=logits.device), sequence_lengths]
+ elif self.reward_token == "mean":
+ valid_lengths = torch.clamp(sequence_lengths, min=0, max=logits.size(1) - 1)
+ pooled_logits = torch.stack([logits[i, : valid_lengths[i]].mean(dim=0) for i in range(batch_size)])
+ elif self.reward_token == "special":
+ if self.special_token_ids is None:
+ raise ValueError("HPSv3 reward_token='special' requires special_token_ids.")
+ special_token_mask = torch.zeros_like(input_ids, dtype=torch.bool)
+ for special_token_id in self.special_token_ids:
+ special_token_mask = special_token_mask | (input_ids == special_token_id)
+ pooled_logits = logits[special_token_mask, ...].view(batch_size, -1)
+ else:
+ raise ValueError(f"Invalid HPSv3 reward token mode: {self.reward_token}")
+
+ return {"logits": pooled_logits}
+
+
+class HPSv3Qwen2VLRewardModel(HPSv3RewardModelMixin, Qwen2VLForConditionalGeneration):
+ def __init__(
+ self,
+ config=None,
+ vocab_size=None,
+ output_dim=2,
+ reward_token="special",
+ special_token_ids=None,
+ rm_head_type="ranknet",
+ rm_head_kwargs=None,
+ ):
+ if config is None:
+ config = self.default_config(vocab_size or 151658)
+ elif vocab_size is not None and hasattr(config, "text_config"):
+ config.text_config.vocab_size = vocab_size
+
+ super().__init__(config)
+ self.init_reward_head(
+ output_dim=output_dim,
+ reward_token=reward_token,
+ special_token_ids=special_token_ids,
+ rm_head_type=rm_head_type,
+ rm_head_kwargs=rm_head_kwargs,
+ )
+
+ @staticmethod
+ def default_config(vocab_size=151658):
+ from transformers import Qwen2VLConfig
+
+ return Qwen2VLConfig(
+ text_config={
+ "vocab_size": vocab_size,
+ "hidden_size": 3584,
+ "intermediate_size": 18944,
+ "num_hidden_layers": 28,
+ "num_attention_heads": 28,
+ "num_key_value_heads": 4,
+ "hidden_act": "silu",
+ "max_position_embeddings": 32768,
+ "initializer_range": 0.02,
+ "rms_norm_eps": 1e-6,
+ "use_cache": True,
+ "use_sliding_window": False,
+ "sliding_window": 32768,
+ "max_window_layers": 28,
+ "attention_dropout": 0.0,
+ "rope_parameters": {
+ "rope_type": "default",
+ "type": "mrope",
+ "mrope_section": [16, 24, 24],
+ "rope_theta": 1000000.0,
+ },
+ "bos_token_id": 151643,
+ "eos_token_id": 151645,
+ },
+ vision_config={
+ "depth": 32,
+ "embed_dim": 1280,
+ "hidden_size": 3584,
+ "mlp_ratio": 4,
+ "num_heads": 16,
+ "in_channels": 3,
+ "patch_size": 14,
+ "spatial_merge_size": 2,
+ "temporal_patch_size": 2,
+ },
+ image_token_id=151655,
+ video_token_id=151656,
+ vision_start_token_id=151652,
+ vision_end_token_id=151653,
+ tie_word_embeddings=False,
+ )
+
+
+class HPSv3Model(torch.nn.Module):
+ def __init__(
+ self,
+ model,
+ processor,
+ use_special_tokens=True,
+ max_pixels=256 * 28 * 28,
+ min_pixels=256 * 28 * 28,
+ score_index=0,
+ ):
+ super().__init__()
+ self.model = model
+ self.processor = processor
+ self.use_special_tokens = use_special_tokens
+ self.max_pixels = max_pixels
+ self.min_pixels = min_pixels
+ self.score_index = score_index
+
+ @property
+ def device(self):
+ return next(self.parameters(), torch.tensor([])).device
+
+ def _normalize_inputs(self, prompts, images):
+ images = _as_list(images)
+ prompts = _as_list(prompts)
+
+ if len(prompts) == 1 and len(images) > 1:
+ prompts = prompts * len(images)
+ if len(images) == 1 and len(prompts) > 1:
+ images = images * len(prompts)
+
+ if len(prompts) != len(images):
+ raise ValueError(f"Expected the same number of prompts and images, got {len(prompts)} and {len(images)}.")
+ return prompts, images
+
+ def _prepare_images(self, images):
+ prepared = []
+ for image in images:
+ image = image.convert("RGB")
+ height, width = image.height, image.width
+
+ resized_height, resized_width = _smart_resize(
+ height,
+ width,
+ min_pixels=self.min_pixels,
+ max_pixels=self.max_pixels,
+ )
+ prepared.append(image.resize((resized_width, resized_height), Image.BICUBIC))
+
+ return prepared
+
+ def _messages(self, prompts, images):
+ suffix = HPSV3_PROMPT_WITH_SPECIAL_TOKEN if self.use_special_tokens else HPSV3_PROMPT_WITHOUT_SPECIAL_TOKEN
+ messages = []
+
+ for prompt, image in zip(prompts, images):
+ messages.append(
+ [
+ {
+ "role": "user",
+ "content": [
+ {"type": "image", "image": image},
+ {"type": "text", "text": HPSV3_INSTRUCTION.format(text_prompt=prompt) + suffix},
+ ],
+ }
+ ]
+ )
+ return messages
+
+ def _prepare_batch(self, prompts, images):
+ images = self._prepare_images(images)
+ messages = self._messages(prompts, images)
+
+ text = self.processor.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
+ batch = self.processor(text=text, images=images, padding=True, return_tensors="pt")
+
+ return batch.to(self.device)
+
+ @torch.no_grad()
+ def forward(self, prompts: Union[str, list[str]], images):
+ prompts, images = self._normalize_inputs(prompts, images)
+ batch = self._prepare_batch(prompts, images)
+
+ rewards = self.model(return_dict=True, **batch)["logits"]
+ if rewards.ndim == 2:
+ return rewards[:, self.score_index]
+
+ return rewards
\ No newline at end of file
diff --git a/diffsynth/models/ideogram4_dit.py b/diffsynth/models/ideogram4_dit.py
new file mode 100644
index 0000000000000000000000000000000000000000..0f6c3e4a9d64f44fe64d97b33bee11b6abd209e3
--- /dev/null
+++ b/diffsynth/models/ideogram4_dit.py
@@ -0,0 +1,470 @@
+import math
+from dataclasses import dataclass
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from ..core.gradient import gradient_checkpoint_forward
+from ..core.quant import QuantBackend, BackendConfig, register_quant_backend, register_quant_method
+
+LLM_TOKEN_INDICATOR = 3
+OUTPUT_IMAGE_INDICATOR = 2
+IMAGE_POSITION_OFFSET = 65536
+QWEN3_VL_ACTIVATION_LAYERS = (0, 3, 6, 9, 12, 15, 18, 21, 24, 27, 30, 33, 35)
+
+FP8_E4M3_MAX = 448.0
+FP8_WEIGHT_DTYPE = torch.float8_e4m3fn
+
+
+class Fp8Linear(torch.nn.Linear):
+ """Linear layer holding an e4m3 float8 weight + per-row float32 scale."""
+
+ # Names of the packed tensors whose dtype must survive `to(dtype)` / `half()` / etc.
+ dtype_guarded_tensor_names: tuple = ("weight", "weight_scale")
+
+ weight: torch.Tensor
+ weight_scale: torch.Tensor
+ bias: torch.Tensor | None
+
+ def __init__(
+ self,
+ in_features: int,
+ out_features: int,
+ bias: bool,
+ compute_dtype: torch.dtype,
+ ) -> None:
+ with torch.device("meta"):
+ super().__init__(in_features, out_features, bias=bias, dtype=compute_dtype)
+ del self.weight
+ self.compute_dtype = compute_dtype
+ self.register_buffer(
+ "weight",
+ torch.empty(out_features, in_features, dtype=FP8_WEIGHT_DTYPE, device="meta"),
+ )
+ self.register_buffer("weight_scale", torch.empty(out_features, dtype=torch.float32, device="meta"))
+ if self.bias is not None:
+ self.bias.requires_grad_(False)
+
+ def _apply(self, fn, recurse=True):
+ # All dtype/device conversions (`to` / `half` / `float` / ...) funnel through
+ # `_apply`; wrap `fn` so that a conversion which would change the dtype of a
+ # tensor listed in `dtype_guarded_tensor_names` is redone as a device-only move,
+ # keeping contract guarantee (b). The bias and everything else convert as usual.
+ protected = {id(tensor) for name in self.dtype_guarded_tensor_names
+ if (tensor := getattr(self, name, None)) is not None}
+
+ def guard(tensor):
+ converted = fn(tensor)
+ if id(tensor) in protected and converted.dtype != tensor.dtype:
+ return tensor.to(device=converted.device)
+ return converted
+
+ return super()._apply(guard, recurse)
+
+ def forward(self, x: torch.Tensor) -> torch.Tensor:
+ w = self.weight.to(x.dtype) * self.weight_scale.to(x.dtype).unsqueeze(1)
+ bias = self.bias.to(x.dtype) if self.bias is not None else None
+ return F.linear(x, w, bias)
+
+
+# Backend for the released fp8 checkpoints, whose quantized Linears follow the
+# compressed-tensors layout: `X.weight` in float8_e4m3 plus one float32 scale per output
+# channel in `X.weight_scale`. `Fp8Linear` names its buffers accordingly, so loading such a
+# checkpoint is a pure structure replacement.
+@register_quant_backend("ideogram4_fp8")
+class Ideogram4Fp8QuantBackend(QuantBackend):
+ def capabilities(self):
+ return {**super().capabilities(), "is_serializable": True, "is_differentiable": True}
+
+ def create_quantized_linear_shell(self, linear, compute_dtype):
+ with torch.device("meta"):
+ return Fp8Linear(linear.in_features, linear.out_features, bias=linear.bias is not None, compute_dtype=compute_dtype)
+
+ def quantized_linear_classes(self) -> tuple:
+ return (Fp8Linear,)
+
+ def dequantize_to_linear(self, module, compute_dtype, compute_device=None, model_device=None):
+ if compute_device is not None:
+ module = module.to(device=compute_device)
+ fp_weight = (module.weight.to(compute_dtype) * module.weight_scale.to(compute_dtype).unsqueeze(1))
+ linear = nn.Linear(module.in_features, module.out_features, bias=module.bias is not None, device="meta")
+ linear.weight = nn.Parameter(fp_weight, requires_grad=False)
+ if module.bias is not None:
+ linear.bias = nn.Parameter(module.bias.to(dtype=compute_dtype, device=fp_weight.device), requires_grad=False)
+ if model_device is not None:
+ linear = linear.to(device=model_device)
+ return linear
+
+
+@dataclass
+class Ideogram4Fp8Config(BackendConfig):
+ """No knobs: the checkpoint's compressed-tensors layout fully determines this path."""
+
+
+register_quant_method(
+ "ideogram4_fp8",
+ "ideogram4_fp8",
+ Ideogram4Fp8Config.from_kwargs,
+ label="8bit, fp8 (compressed-tensors layout, per-channel scale), weight-only",
+)
+
+
+@dataclass
+class Ideogram4Config:
+ emb_dim: int = 4608
+ num_layers: int = 34
+ num_heads: int = 18
+ intermediate_size: int = 12288
+ adanln_dim: int = 512
+ in_channels: int = 128
+ llm_features_dim: int = 4096 * len(QWEN3_VL_ACTIVATION_LAYERS)
+ rope_theta: int = 5_000_000
+ mrope_section: tuple[int, ...] = (24, 20, 20)
+ norm_eps: float = 1e-5
+
+
+def _rotate_half(x: torch.Tensor) -> torch.Tensor:
+ half = x.shape[-1] // 2
+ x1 = x[..., :half]
+ x2 = x[..., half:]
+ return torch.cat((-x2, x1), dim=-1)
+
+
+def _apply_rotary_pos_emb(
+ q: torch.Tensor,
+ k: torch.Tensor,
+ cos: torch.Tensor,
+ sin: torch.Tensor,
+) -> tuple[torch.Tensor, torch.Tensor]:
+ cos = cos.unsqueeze(1)
+ sin = sin.unsqueeze(1)
+ q_embed = (q * cos) + (_rotate_half(q) * sin)
+ k_embed = (k * cos) + (_rotate_half(k) * sin)
+ return q_embed, k_embed
+
+
+class Ideogram4MRoPE(nn.Module):
+ inv_freq: torch.Tensor
+
+ def __init__(
+ self,
+ head_dim: int,
+ base: int,
+ mrope_section: tuple[int, ...],
+ ) -> None:
+ super().__init__()
+ inv_freq = 1.0 / (
+ base ** (torch.arange(0, head_dim, 2, dtype=torch.float32) / head_dim)
+ )
+ self.register_buffer("inv_freq", inv_freq, persistent=False)
+ self.base = base
+ self.mrope_section = tuple(mrope_section)
+ self.head_dim = head_dim
+
+ @torch.no_grad()
+ def forward(self, position_ids: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
+ assert position_ids.ndim == 3 and position_ids.shape[-1] == 3
+ batch_size, seq_len, _ = position_ids.shape
+
+ pos = position_ids.permute(2, 0, 1).to(dtype=torch.float32)
+ inv_freq = 1.0 / (
+ self.base ** (torch.arange(0, self.head_dim, 2, dtype=torch.float32, device=pos.device) / self.head_dim)
+ )
+ inv_freq = inv_freq[None, None, :, None].expand(3, batch_size, -1, 1)
+ freqs = inv_freq @ pos.unsqueeze(2)
+ freqs = freqs.transpose(2, 3)
+
+ freqs_t = freqs[0].clone()
+ for axis, offset in ((1, 1), (2, 2)):
+ length = self.mrope_section[axis] * 3
+ idx = torch.arange(offset, length, 3, device=freqs_t.device)
+ freqs_t[..., idx] = freqs[axis][..., idx]
+
+ emb = torch.cat((freqs_t, freqs_t), dim=-1)
+ return emb.cos(), emb.sin()
+
+
+class Ideogram4RMSNorm(nn.Module):
+ def __init__(self, dim: int, eps: float = 1e-6) -> None:
+ super().__init__()
+ self.weight = nn.Parameter(torch.ones(dim))
+ self.eps = eps
+
+ def forward(self, x: torch.Tensor) -> torch.Tensor:
+ return F.rms_norm(x, self.weight.shape, self.weight, self.eps)
+
+
+class Ideogram4Attention(nn.Module):
+ def __init__(self, hidden_size: int, num_heads: int, eps: float = 1e-5) -> None:
+ super().__init__()
+ assert hidden_size % num_heads == 0
+ self.hidden_size = hidden_size
+ self.num_heads = num_heads
+ self.head_dim = hidden_size // num_heads
+
+ self.qkv = nn.Linear(hidden_size, hidden_size * 3, bias=False)
+ self.norm_q = Ideogram4RMSNorm(self.head_dim, eps=eps)
+ self.norm_k = Ideogram4RMSNorm(self.head_dim, eps=eps)
+ self.o = nn.Linear(hidden_size, hidden_size, bias=False)
+
+ def forward(
+ self,
+ x: torch.Tensor,
+ segment_ids: torch.Tensor,
+ cos: torch.Tensor,
+ sin: torch.Tensor,
+ ) -> torch.Tensor:
+ batch_size, seq_len, _ = x.shape
+
+ qkv = self.qkv(x)
+ qkv = qkv.view(batch_size, seq_len, 3, self.num_heads, self.head_dim)
+ q, k, v = qkv.unbind(dim=2)
+
+ q = self.norm_q(q)
+ k = self.norm_k(k)
+
+ q = q.transpose(1, 2)
+ k = k.transpose(1, 2)
+ v = v.transpose(1, 2)
+
+ q, k = _apply_rotary_pos_emb(q, k, cos, sin)
+
+ attn_mask = (segment_ids.unsqueeze(2) == segment_ids.unsqueeze(1)).unsqueeze(1)
+
+ out = F.scaled_dot_product_attention(q, k, v, attn_mask=attn_mask)
+ out = out.transpose(1, 2).reshape(batch_size, seq_len, self.hidden_size)
+ return self.o(out)
+
+
+class Ideogram4MLP(nn.Module):
+ def __init__(self, dim: int, hidden_dim: int) -> None:
+ super().__init__()
+ self.w1 = nn.Linear(dim, hidden_dim, bias=False)
+ self.w2 = nn.Linear(hidden_dim, dim, bias=False)
+ self.w3 = nn.Linear(dim, hidden_dim, bias=False)
+
+ def forward(self, x: torch.Tensor) -> torch.Tensor:
+ return self.w2(F.silu(self.w1(x)) * self.w3(x))
+
+
+class Ideogram4TransformerBlock(nn.Module):
+ def __init__(
+ self,
+ hidden_size: int,
+ intermediate_size: int,
+ num_heads: int,
+ norm_eps: float,
+ adanln_dim: int,
+ ) -> None:
+ super().__init__()
+ self.attention = Ideogram4Attention(hidden_size, num_heads, eps=1e-5)
+ self.feed_forward = Ideogram4MLP(hidden_size, intermediate_size)
+
+ self.attention_norm1 = Ideogram4RMSNorm(hidden_size, eps=norm_eps)
+ self.ffn_norm1 = Ideogram4RMSNorm(hidden_size, eps=norm_eps)
+ self.attention_norm2 = Ideogram4RMSNorm(hidden_size, eps=norm_eps)
+ self.ffn_norm2 = Ideogram4RMSNorm(hidden_size, eps=norm_eps)
+
+ self.adaln_modulation = nn.Linear(adanln_dim, 4 * hidden_size, bias=True)
+
+ def forward(
+ self,
+ x: torch.Tensor,
+ segment_ids: torch.Tensor,
+ cos: torch.Tensor,
+ sin: torch.Tensor,
+ adaln_input: torch.Tensor,
+ ) -> torch.Tensor:
+ mod = self.adaln_modulation(adaln_input)
+ scale_msa, gate_msa, scale_mlp, gate_mlp = mod.chunk(4, dim=-1)
+ gate_msa = torch.tanh(gate_msa)
+ gate_mlp = torch.tanh(gate_mlp)
+ scale_msa = 1.0 + scale_msa
+ scale_mlp = 1.0 + scale_mlp
+
+ attn_out = self.attention(
+ self.attention_norm1(x) * scale_msa,
+ segment_ids=segment_ids,
+ cos=cos,
+ sin=sin,
+ )
+ x = x + gate_msa * self.attention_norm2(attn_out)
+ x = x + gate_mlp * self.ffn_norm2(self.feed_forward(self.ffn_norm1(x) * scale_mlp))
+ return x
+
+
+def _sinusoidal_embedding(
+ t: torch.Tensor, dim: int, scale: float = 1e4
+) -> torch.Tensor:
+ t = t.to(torch.float32)
+ half = dim // 2
+ freq = math.log(scale) / (half - 1)
+ freq = torch.exp(torch.arange(half, dtype=torch.float32, device=t.device) * -freq)
+ emb = t.unsqueeze(-1) * freq
+ emb = torch.cat([torch.sin(emb), torch.cos(emb)], dim=-1)
+ if dim % 2 == 1:
+ emb = F.pad(emb, (0, 1))
+ return emb
+
+
+class Ideogram4EmbedScalar(nn.Module):
+ def __init__(self, dim: int, input_range: tuple[float, float]) -> None:
+ super().__init__()
+ self.dim = dim
+ self.range_min, self.range_max = input_range
+ assert self.range_max > self.range_min
+ self.mlp_in = nn.Linear(dim, dim, bias=True)
+ self.mlp_out = nn.Linear(dim, dim, bias=True)
+
+ def forward(self, x: torch.Tensor) -> torch.Tensor:
+ x = x.to(torch.float32)
+ scaled = 1e4 * (x - self.range_min) / (self.range_max - self.range_min)
+ emb = _sinusoidal_embedding(scaled, self.dim)
+ emb = emb.to(
+ getattr(self.mlp_in, "compute_dtype", None) or getattr(self.mlp_in, "computation_dtype", None) or self.mlp_in.weight.dtype
+ )
+ emb = F.silu(self.mlp_in(emb))
+ return self.mlp_out(emb)
+
+
+class Ideogram4FinalLayer(nn.Module):
+ def __init__(self, hidden_size: int, out_channels: int, adanln_dim: int) -> None:
+ super().__init__()
+ self.norm_final = nn.LayerNorm(hidden_size, eps=1e-6, elementwise_affine=False)
+ self.linear = nn.Linear(hidden_size, out_channels, bias=True)
+ self.adaln_modulation = nn.Linear(adanln_dim, hidden_size, bias=True)
+
+ def forward(self, x: torch.Tensor, c: torch.Tensor) -> torch.Tensor:
+ scale = 1.0 + self.adaln_modulation(F.silu(c))
+ return self.linear(self.norm_final(x) * scale)
+
+
+class Ideogram4DiT(nn.Module):
+ """Ideogram 4 flow-matching transformer."""
+
+ def __init__(self, config: Ideogram4Config | dict = None, **kwargs) -> None:
+ super().__init__()
+ if config is None:
+ config = Ideogram4Config()
+ elif isinstance(config, dict):
+ config = Ideogram4Config(**config)
+ self.config = config
+ self.patch_size = 2
+
+ head_dim = config.emb_dim // config.num_heads
+
+ self.input_proj = nn.Linear(config.in_channels, config.emb_dim, bias=True)
+ self.llm_cond_norm = Ideogram4RMSNorm(config.llm_features_dim, eps=1e-6)
+ self.llm_cond_proj = nn.Linear(config.llm_features_dim, config.emb_dim, bias=True)
+ self.t_embedding = Ideogram4EmbedScalar(config.emb_dim, input_range=(0.0, 1.0))
+ self.adaln_proj = nn.Linear(config.emb_dim, config.adanln_dim, bias=True)
+
+ self.embed_image_indicator = nn.Embedding(2, config.emb_dim)
+
+ self.rotary_emb = Ideogram4MRoPE(
+ head_dim=head_dim,
+ base=config.rope_theta,
+ mrope_section=config.mrope_section,
+ )
+
+ self.layers = nn.ModuleList(
+ [
+ Ideogram4TransformerBlock(
+ hidden_size=config.emb_dim,
+ intermediate_size=config.intermediate_size,
+ num_heads=config.num_heads,
+ norm_eps=config.norm_eps,
+ adanln_dim=config.adanln_dim,
+ )
+ for _ in range(config.num_layers)
+ ]
+ )
+
+ self.final_layer = Ideogram4FinalLayer(
+ hidden_size=config.emb_dim,
+ out_channels=config.in_channels,
+ adanln_dim=config.adanln_dim,
+ )
+
+ @property
+ def device(self) -> torch.device:
+ return next(self.parameters()).device
+
+ def forward(
+ self,
+ *,
+ llm_features: torch.Tensor,
+ x: torch.Tensor,
+ t: torch.Tensor,
+ position_ids: torch.Tensor,
+ segment_ids: torch.Tensor,
+ indicator: torch.Tensor,
+ use_gradient_checkpointing: bool = False,
+ use_gradient_checkpointing_offload: bool = False,
+ ) -> torch.Tensor:
+ """Velocity prediction.
+
+ Args:
+ llm_features: (B, L, llm_features_dim) Qwen3-VL conditioning features.
+ x: (B, L, in_channels) noise tokens.
+ t: (B,) or (B, L) flow-matching time in [0, 1].
+ position_ids: (B, L, 3) (t, h, w) positions for MRoPE.
+ segment_ids: (B, L) sample id within a packed batch.
+ indicator: (B, L) per-token role: LLM_TOKEN_INDICATOR or OUTPUT_IMAGE_INDICATOR.
+
+ Returns:
+ (B, L, in_channels) velocity prediction in float32.
+ """
+ batch_size, seq_len, in_channels = x.shape
+ assert in_channels == self.config.in_channels
+
+ param_dtype = (
+ getattr(self.input_proj, "compute_dtype", None) or getattr(self.input_proj, "computation_dtype", None) or self.input_proj.weight.dtype
+ )
+ x = x.to(param_dtype)
+ t = t.to(param_dtype)
+ llm_features = llm_features.to(param_dtype)
+
+ indicator = indicator.to(torch.long)
+ llm_token_mask = (indicator == LLM_TOKEN_INDICATOR).to(x.dtype).unsqueeze(-1)
+ output_image_mask = (indicator == OUTPUT_IMAGE_INDICATOR).to(x.dtype).unsqueeze(-1)
+
+ llm_features = llm_features * llm_token_mask
+ x = x * output_image_mask
+
+ x = self.input_proj(x) * output_image_mask
+
+ t_cond = self.t_embedding(t)
+ if t.dim() == 1:
+ t_cond = t_cond.unsqueeze(1)
+ adaln_input = F.silu(self.adaln_proj(t_cond))
+
+ llm_features = self.llm_cond_norm(llm_features)
+ llm_features = self.llm_cond_proj(llm_features) * llm_token_mask
+
+ h = x + llm_features
+
+ image_indicator_embedding = self.embed_image_indicator(
+ (indicator == OUTPUT_IMAGE_INDICATOR).to(torch.long)
+ )
+ h = h + image_indicator_embedding
+
+ cos, sin = self.rotary_emb(position_ids)
+ cos = cos.to(h.dtype)
+ sin = sin.to(h.dtype)
+
+ for layer in self.layers:
+ h = gradient_checkpoint_forward(
+ layer,
+ use_gradient_checkpointing=use_gradient_checkpointing,
+ use_gradient_checkpointing_offload=use_gradient_checkpointing_offload,
+ x=h,
+ segment_ids=segment_ids,
+ cos=cos,
+ sin=sin,
+ adaln_input=adaln_input,
+ )
+
+ out = self.final_layer(h, c=adaln_input)
+ return out.to(torch.float32)
diff --git a/diffsynth/models/ideogram4_text_encoder.py b/diffsynth/models/ideogram4_text_encoder.py
new file mode 100644
index 0000000000000000000000000000000000000000..628527cf06bfa50ae23a856cd38b0348e4fd58f2
--- /dev/null
+++ b/diffsynth/models/ideogram4_text_encoder.py
@@ -0,0 +1,137 @@
+import torch
+import torch.nn as nn
+from .ideogram4_dit import Ideogram4Fp8QuantBackend # noqa: F401
+
+LLM_TOKEN_INDICATOR = 3
+QWEN3_VL_ACTIVATION_LAYERS = (0, 3, 6, 9, 12, 15, 18, 21, 24, 27, 30, 33, 35)
+
+_DEFAULT_TEXT_ENCODER_CONFIG = {
+ "architectures": ["Qwen3VLModel"],
+ "dtype": "bfloat16",
+ "image_token_id": 151655,
+ "model_type": "qwen3_vl",
+ "text_config": {
+ "attention_bias": False,
+ "attention_dropout": 0.0,
+ "bos_token_id": 151643,
+ "dtype": "bfloat16",
+ "eos_token_id": 151645,
+ "head_dim": 128,
+ "hidden_act": "silu",
+ "hidden_size": 4096,
+ "initializer_range": 0.02,
+ "intermediate_size": 12288,
+ "max_position_embeddings": 262144,
+ "model_type": "qwen3_vl_text",
+ "num_attention_heads": 32,
+ "num_hidden_layers": 36,
+ "num_key_value_heads": 8,
+ "pad_token_id": None,
+ "rms_norm_eps": 1e-06,
+ "rope_parameters": {
+ "mrope_interleaved": True,
+ "mrope_section": [24, 20, 20],
+ "rope_theta": 5000000,
+ "rope_type": "default",
+ },
+ "use_cache": True,
+ "vocab_size": 151936,
+ },
+ "tie_word_embeddings": False,
+ "transformers_version": "5.8.0",
+ "video_token_id": 151656,
+ "vision_config": {
+ "deepstack_visual_indexes": [8, 16, 24],
+ "depth": 27,
+ "dtype": "bfloat16",
+ "hidden_act": "gelu_pytorch_tanh",
+ "hidden_size": 1152,
+ "in_channels": 3,
+ "initializer_range": 0.02,
+ "intermediate_size": 4304,
+ "model_type": "qwen3_vl_vision",
+ "num_heads": 16,
+ "num_position_embeddings": 2304,
+ "out_hidden_size": 4096,
+ "patch_size": 16,
+ "spatial_merge_size": 2,
+ "temporal_patch_size": 2,
+ },
+ "vision_end_token_id": 151653,
+ "vision_start_token_id": 151652,
+ "ideogram_fp8_weight_only": True,
+}
+
+
+class Ideogram4TextEncoder(nn.Module):
+ """Qwen3-VL-8B-Instruct wrapper that extracts hidden states from specific layers."""
+
+ def __init__(self, config_path: str = None, **kwargs) -> None:
+ super().__init__()
+ from transformers import AutoConfig, AutoModel
+ if config_path is None:
+ config_kwargs = {k: v for k, v in _DEFAULT_TEXT_ENCODER_CONFIG.items() if k != "model_type"}
+ config = AutoConfig.for_model("qwen3_vl", **config_kwargs)
+ else:
+ config = AutoConfig.from_pretrained(config_path, trust_remote_code=True)
+ self.model = AutoModel.from_config(config, trust_remote_code=True)
+ self.config = config
+
+ def forward(
+ self,
+ token_ids: torch.Tensor,
+ attention_mask: torch.Tensor,
+ text_position_ids: torch.Tensor,
+ ) -> torch.Tensor:
+ """Extract hidden states from specific layers of Qwen3-VL.
+
+ Args:
+ token_ids: (B, L) token ids
+ attention_mask: (B, L) attention mask
+ text_position_ids: (B, L) position ids for text tokens
+
+ Returns:
+ (B, L, hidden_size * num_activation_layers) concatenated hidden states
+ """
+ from transformers.masking_utils import create_causal_mask
+
+ language_model = self.model.language_model
+
+ inputs_embeds = language_model.embed_tokens(token_ids)
+
+ position_ids_4d = text_position_ids[None, ...].expand(4, text_position_ids.shape[0], -1)
+ text_position_ids_4d = position_ids_4d[0]
+ mrope_position_ids = position_ids_4d[1:]
+
+ causal_mask = create_causal_mask(
+ config=language_model.config,
+ inputs_embeds=inputs_embeds,
+ attention_mask=attention_mask,
+ past_key_values=None,
+ position_ids=text_position_ids_4d,
+ )
+ position_embeddings = language_model.rotary_emb(inputs_embeds, mrope_position_ids)
+
+ tap_set = set(QWEN3_VL_ACTIVATION_LAYERS)
+ captured: dict[int, torch.Tensor] = {}
+ hidden_states = inputs_embeds
+ for layer_idx, decoder_layer in enumerate(language_model.layers):
+ hidden_states = decoder_layer(
+ hidden_states,
+ attention_mask=causal_mask,
+ position_ids=text_position_ids_4d,
+ past_key_values=None,
+ position_embeddings=position_embeddings,
+ )
+ if layer_idx in tap_set:
+ captured[layer_idx] = hidden_states
+
+ selected = [captured[i] for i in QWEN3_VL_ACTIVATION_LAYERS]
+ stacked = torch.stack(selected, dim=0)
+ stacked = torch.permute(stacked, (1, 2, 3, 0))
+ batch_size, seq_len, hidden_size = stacked.shape[:3]
+ stacked = stacked.reshape(batch_size, seq_len, -1)
+
+ text_mask = attention_mask.to(stacked.dtype).unsqueeze(-1)
+ stacked = stacked * text_mask
+ return stacked.to(torch.float32)
diff --git a/diffsynth/models/ideogram4_vae.py b/diffsynth/models/ideogram4_vae.py
new file mode 100644
index 0000000000000000000000000000000000000000..5ad38cd420ba1b2d49c2218bed4f6720eb006aaf
--- /dev/null
+++ b/diffsynth/models/ideogram4_vae.py
@@ -0,0 +1,74 @@
+import torch
+from einops import rearrange
+
+
+LATENT_SHIFT = (
+ 0.01984364, 0.10149707, 0.29689495, 0.27188619, -0.21445648, -0.15979549,
+ 0.05021099, -0.15083604, -0.15360136, -0.20131799, 0.01922352, 0.0622626,
+ 0.10140969, -0.06739428, 0.3758261, -0.233712, 0.35164491, -0.02590912,
+ -0.0271935, -0.10833897, -0.1476848, -0.01130957, -0.2298372, 0.23526423,
+ -0.10893522, 0.11957631, 0.04047799, 0.3134589, -0.17225064, -0.18646109,
+ -0.34691978, -0.03571246, 0.02583857, 0.10190072, 0.28402294, 0.26952152,
+ -0.21634675, -0.17938656, 0.04358909, -0.15007621, -0.1548502, -0.18971131,
+ 0.02710861, 0.05609494, 0.10697846, -0.06854968, 0.38167698, -0.24269937,
+ 0.35705471, -0.03063305, -0.02946109, -0.11244286, -0.14336038, -0.01362137,
+ -0.21863696, 0.23228983, -0.11739769, 0.11693044, 0.02563311, 0.31356594,
+ -0.17420591, -0.19006285, -0.34905377, -0.04025005, 0.01924137, 0.07652984,
+ 0.2995608, 0.2628057, -0.22011674, -0.12715361, 0.04879879, -0.14075719,
+ -0.15935895, -0.2123584, 0.01974813, 0.05523547, 0.10011992, -0.06428964,
+ 0.37781868, -0.21491644, 0.34254215, -0.03153528, -0.0310082, -0.10761415,
+ -0.14730405, -0.02475182, -0.2285588, 0.2515081, -0.10445128, 0.12446,
+ 0.07062869, 0.30880162, -0.18016875, -0.18869164, -0.34533499, -0.0129177,
+ 0.02578168, 0.07993659, 0.28642181, 0.26038408, -0.22459419, -0.14820155,
+ 0.04059549, -0.14043529, -0.16111187, -0.2020305, 0.02602069, 0.04852717,
+ 0.10432153, -0.06309942, 0.38402443, -0.22397003, 0.34814481, -0.03774432,
+ -0.03381438, -0.11245691, -0.14128767, -0.02853208, -0.21752016, 0.24872463,
+ -0.11399775, 0.1222687, 0.05620835, 0.309178, -0.18065738, -0.19401479,
+ -0.34495114, -0.01760592,
+)
+
+LATENT_SCALE = (
+ 1.63933691, 1.70204478, 1.73642566, 1.90004803, 1.6675316, 1.69059584,
+ 1.56853198, 1.62314944, 1.89106626, 1.58086668, 1.60822129, 1.60962993,
+ 1.63322129, 1.56074359, 1.73419528, 1.7919265, 1.64040632, 1.66802808,
+ 1.60390303, 1.75480492, 1.63187587, 1.64334594, 1.61722884, 1.60146046,
+ 1.63459219, 1.55291476, 1.68771497, 1.68415657, 1.78966054, 1.66631641,
+ 1.65626686, 1.65976433, 1.63487607, 1.69513249, 1.72933756, 1.91310663,
+ 1.67035057, 1.72286863, 1.56719251, 1.61934825, 1.88628859, 1.56911539,
+ 1.59455129, 1.60829869, 1.62470611, 1.56052853, 1.73677003, 1.77563606,
+ 1.63732541, 1.66370527, 1.59508952, 1.75153949, 1.63029275, 1.64517667,
+ 1.61659342, 1.59722044, 1.64103121, 1.5408531, 1.68610394, 1.67772755,
+ 1.78998563, 1.66621713, 1.65458955, 1.66041308, 1.64710857, 1.68163503,
+ 1.74000294, 1.92784786, 1.67411194, 1.67395548, 1.57406532, 1.62199356,
+ 1.87618195, 1.5584375, 1.57438785, 1.61711053, 1.63094305, 1.55644029,
+ 1.73124302, 1.80666627, 1.6463621, 1.65932006, 1.60816188, 1.75682671,
+ 1.64695873, 1.63121722, 1.61380832, 1.60478651, 1.63396035, 1.53505068,
+ 1.65534289, 1.67132281, 1.80317197, 1.6767314, 1.65700938, 1.68426259,
+ 1.65339716, 1.67540638, 1.73298504, 1.94067348, 1.67893609, 1.70635117,
+ 1.5730906, 1.61928553, 1.87148809, 1.56244866, 1.56697152, 1.61584394,
+ 1.62759496, 1.55480378, 1.73484107, 1.79055143, 1.64688773, 1.66121492,
+ 1.60135887, 1.75254572, 1.64798332, 1.62989921, 1.61381592, 1.60792883,
+ 1.63939668, 1.53075757, 1.65371318, 1.66801185, 1.80029087, 1.67591476,
+ 1.65655173, 1.68533454,
+)
+
+def get_latent_norm(device: torch.device) -> tuple[torch.Tensor, torch.Tensor]:
+ shift = torch.tensor(LATENT_SHIFT, dtype=torch.float32, device=device)
+ scale = torch.tensor(LATENT_SCALE, dtype=torch.float32, device=device)
+ return shift, scale
+
+def decode(vae, latents, height, width, torch_dtype):
+ latent_shift, latent_scale = get_latent_norm(latents.device)
+ latents = latents.float() * latent_scale + latent_shift
+ latents = rearrange(latents, "B (H W) (P Q C) -> B C (H P) (W Q)", P=2, Q=2, H=height//16, W=width//16).to(torch.bfloat16)
+ latents = latents.to(torch_dtype)
+ image = vae._decode(latents)
+ return image
+
+def encode(vae, image, height, width, torch_dtype):
+ latents = vae._encode(image)[:, :32]
+ latent_shift, latent_scale = get_latent_norm(latents.device)
+ latents = rearrange(latents, "B C (H P) (W Q) -> B (H W) (P Q C)", P=2, Q=2, H=height//16, W=width//16).to(torch.bfloat16)
+ latents = (latents.float() - latent_shift) / latent_scale
+ latents = latents.to(torch_dtype)
+ return latents
diff --git a/diffsynth/models/image_reward.py b/diffsynth/models/image_reward.py
new file mode 100644
index 0000000000000000000000000000000000000000..da3f39fc206805edbde99ee286f0e8f2ccff07f7
--- /dev/null
+++ b/diffsynth/models/image_reward.py
@@ -0,0 +1,206 @@
+from typing import Union
+import torch
+from PIL import Image
+from torchvision.transforms import CenterCrop, Compose, Normalize, Resize, ToTensor
+from torchvision.transforms import InterpolationMode
+
+BICUBIC = InterpolationMode.BICUBIC
+ImageInput = Union[Image.Image, list[Image.Image], tuple[Image.Image, ...]]
+
+def _convert_image_to_rgb(image):
+ return image.convert("RGB")
+
+def _image_transform(image_size):
+ return Compose(
+ [
+ Resize(image_size, interpolation=BICUBIC),
+ CenterCrop(image_size),
+ _convert_image_to_rgb,
+ ToTensor(),
+ Normalize((0.48145466, 0.4578275, 0.40821073), (0.26862954, 0.26130258, 0.27577711)),
+ ]
+ )
+
+def _as_list(value):
+ if isinstance(value, (list, tuple)):
+ return list(value)
+ return [value]
+
+
+class ImageRewardMLP(torch.nn.Module):
+ def __init__(self, input_size):
+ super().__init__()
+ self.layers = torch.nn.Sequential(
+ torch.nn.Linear(input_size, 1024),
+ torch.nn.Dropout(0.2),
+ torch.nn.Linear(1024, 128),
+ torch.nn.Dropout(0.2),
+ torch.nn.Linear(128, 64),
+ torch.nn.Dropout(0.1),
+ torch.nn.Linear(64, 16),
+ torch.nn.Linear(16, 1),
+ )
+
+ for name, param in self.layers.named_parameters():
+ if "weight" in name:
+ torch.nn.init.normal_(param, mean=0.0, std=1.0 / (input_size + 1))
+ if "bias" in name:
+ torch.nn.init.constant_(param, val=0)
+
+ def forward(self, x):
+ return self.layers(x)
+
+
+class ImageRewardModel(torch.nn.Module):
+ def __init__(self, blip=None, tokenizer=None, image_size=224, max_length=35, mean=0.16717362830052426, std=1.0333394966054072):
+ super().__init__()
+ if blip is None:
+ blip = self.default_blip_model()
+
+ self.blip = blip
+ self.tokenizer = tokenizer
+ self.preprocess = _image_transform(image_size)
+ self.max_length = max_length
+ self.mlp = ImageRewardMLP(blip.config.text_config.hidden_size)
+
+ self.register_buffer("score_mean", torch.tensor(float(mean)), persistent=False)
+ self.register_buffer("score_std", torch.tensor(float(std)), persistent=False)
+
+ @staticmethod
+ def default_blip_model():
+ from transformers import BlipConfig, BlipForImageTextRetrieval
+
+ vision_hidden_size = 1024
+ text_config = ImageRewardModel._load_text_config(None)
+ config = BlipConfig(
+ vision_config={
+ "hidden_size": vision_hidden_size,
+ "intermediate_size": vision_hidden_size * 4,
+ "num_hidden_layers": 24,
+ "num_attention_heads": 16,
+ "image_size": 224,
+ "patch_size": 16,
+ "hidden_act": "gelu",
+ "layer_norm_eps": 1e-6,
+ },
+ text_config={
+ **text_config,
+ "vocab_size": 30524,
+ "encoder_hidden_size": vision_hidden_size,
+ "add_cross_attention": True,
+ "is_decoder": True,
+ },
+ projection_dim=256,
+ )
+ return BlipForImageTextRetrieval(config)
+
+ @staticmethod
+ def _load_text_config(med_config_path):
+ return {
+ "hidden_size": 768,
+ "intermediate_size": 3072,
+ "num_hidden_layers": 12,
+ "num_attention_heads": 12,
+ "max_position_embeddings": 512,
+ "vocab_size": 30524,
+ "hidden_act": "gelu",
+ "layer_norm_eps": 1e-12,
+ "attention_probs_dropout_prob": 0.1,
+ "hidden_dropout_prob": 0.1,
+ "pad_token_id": 0,
+ "type_vocab_size": 2,
+ }
+
+ @staticmethod
+ def convert_key_value(key, value):
+ if key.startswith("blip.visual_encoder."):
+ suffix = key[len("blip.visual_encoder.") :]
+
+ if suffix == "cls_token":
+ return "blip.vision_model.embeddings.class_embedding", value
+ if suffix == "pos_embed":
+ return "blip.vision_model.embeddings.position_embedding", value
+ if suffix.startswith("patch_embed.proj."):
+ return "blip.vision_model.embeddings.patch_embedding." + suffix[len("patch_embed.proj.") :], value
+
+ if suffix.startswith("blocks."):
+ parts = suffix.split(".")
+ layer = parts[1]
+ rest = ".".join(parts[2:])
+ prefix = f"blip.vision_model.encoder.layers.{layer}."
+ mapping = {
+ "norm1.": "layer_norm1.",
+ "attn.qkv.": "self_attn.qkv.",
+ "attn.proj.": "self_attn.projection.",
+ "norm2.": "layer_norm2.",
+ "mlp.fc1.": "mlp.fc1.",
+ "mlp.fc2.": "mlp.fc2.",
+ }
+ for source, target in mapping.items():
+ if rest.startswith(source):
+ return prefix + target + rest[len(source) :], value
+
+ if suffix.startswith("norm."):
+ return "blip.vision_model.post_layernorm." + suffix[len("norm.") :], value
+
+ return None, value
+ return key, value
+
+ @property
+ def device(self):
+ return next(self.parameters(), torch.tensor([])).device
+
+ @property
+ def dtype(self):
+ return next(self.parameters(), torch.tensor(0.0)).dtype
+
+ def _tokenize(self, prompts):
+ return self.tokenizer(
+ prompts,
+ padding="max_length",
+ truncation=True,
+ max_length=self.max_length,
+ return_tensors="pt",
+ ).to(self.device)
+
+ def _preprocess_images(self, images):
+ tensors = [self.preprocess(image.convert("RGB")) for image in images]
+ return torch.stack(tensors, dim=0).to(device=self.device, dtype=self.dtype)
+
+ def _normalize_inputs(self, prompts, images):
+ images = _as_list(images)
+ prompts = _as_list(prompts)
+
+ if len(prompts) == 1 and len(images) > 1:
+ prompts = prompts * len(images)
+ if len(images) == 1 and len(prompts) > 1:
+ images = images * len(prompts)
+
+ if len(prompts) != len(images):
+ raise ValueError(f"Expected the same number of prompts and images, got {len(prompts)} and {len(images)}.")
+
+ return prompts, images
+
+ @torch.no_grad()
+ def forward(self, prompts: Union[str, list[str]], images):
+ prompts, images = self._normalize_inputs(prompts, images)
+ text_input = self._tokenize(prompts)
+ image_tensor = self._preprocess_images(images)
+
+ image_output = self.blip.vision_model(pixel_values=image_tensor, return_dict=True)
+ image_embeds = image_output.last_hidden_state
+ image_atts = torch.ones(image_embeds.size()[:-1], dtype=torch.long, device=self.device)
+
+ text_output = self.blip.text_encoder(
+ input_ids=text_input.input_ids,
+ attention_mask=text_input.attention_mask,
+ encoder_hidden_states=image_embeds,
+ encoder_attention_mask=image_atts,
+ return_dict=True,
+ )
+
+ text_features = text_output.last_hidden_state[:, 0, :]
+ rewards = self.mlp(text_features).squeeze(-1)
+ rewards = (rewards - self.score_mean) / self.score_std
+
+ return rewards
\ No newline at end of file
diff --git a/diffsynth/models/joyai_image_dit.py b/diffsynth/models/joyai_image_dit.py
new file mode 100644
index 0000000000000000000000000000000000000000..ba88f606019553ad00bb651855098d796c39115f
--- /dev/null
+++ b/diffsynth/models/joyai_image_dit.py
@@ -0,0 +1,636 @@
+import math
+from typing import Dict, List, Optional, Tuple, Union
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from einops import rearrange
+
+from ..core.attention import attention_forward
+from ..core.gradient import gradient_checkpoint_forward
+
+
+def get_timestep_embedding(
+ timesteps: torch.Tensor,
+ embedding_dim: int,
+ flip_sin_to_cos: bool = False,
+ downscale_freq_shift: float = 1,
+ scale: float = 1,
+ max_period: int = 10000,
+) -> torch.Tensor:
+ assert len(timesteps.shape) == 1, "Timesteps should be a 1d-array"
+ half_dim = embedding_dim // 2
+ exponent = -math.log(max_period) * torch.arange(
+ start=0, end=half_dim, dtype=torch.float32, device=timesteps.device
+ )
+ exponent = exponent / (half_dim - downscale_freq_shift)
+ emb = torch.exp(exponent)
+ emb = timesteps[:, None].float() * emb[None, :]
+ emb = scale * emb
+ emb = torch.cat([torch.sin(emb), torch.cos(emb)], dim=-1)
+ if flip_sin_to_cos:
+ emb = torch.cat([emb[:, half_dim:], emb[:, :half_dim]], dim=-1)
+ if embedding_dim % 2 == 1:
+ emb = torch.nn.functional.pad(emb, (0, 1, 0, 0))
+ return emb
+
+
+class Timesteps(nn.Module):
+ def __init__(self, num_channels: int, flip_sin_to_cos: bool, downscale_freq_shift: float, scale: int = 1):
+ super().__init__()
+ self.num_channels = num_channels
+ self.flip_sin_to_cos = flip_sin_to_cos
+ self.downscale_freq_shift = downscale_freq_shift
+ self.scale = scale
+
+ def forward(self, timesteps: torch.Tensor) -> torch.Tensor:
+ return get_timestep_embedding(
+ timesteps,
+ self.num_channels,
+ flip_sin_to_cos=self.flip_sin_to_cos,
+ downscale_freq_shift=self.downscale_freq_shift,
+ scale=self.scale,
+ )
+
+
+class TimestepEmbedding(nn.Module):
+ def __init__(
+ self,
+ in_channels: int,
+ time_embed_dim: int,
+ act_fn: str = "silu",
+ out_dim: int = None,
+ post_act_fn: Optional[str] = None,
+ cond_proj_dim=None,
+ sample_proj_bias=True,
+ ):
+ super().__init__()
+ self.linear_1 = nn.Linear(in_channels, time_embed_dim, sample_proj_bias)
+ if cond_proj_dim is not None:
+ self.cond_proj = nn.Linear(cond_proj_dim, in_channels, bias=False)
+ else:
+ self.cond_proj = None
+ self.act = nn.SiLU()
+ time_embed_dim_out = out_dim if out_dim is not None else time_embed_dim
+ self.linear_2 = nn.Linear(time_embed_dim, time_embed_dim_out, sample_proj_bias)
+ self.post_act = nn.SiLU() if post_act_fn == "silu" else None
+
+ def forward(self, sample, condition=None):
+ if condition is not None:
+ sample = sample + self.cond_proj(condition)
+ sample = self.linear_1(sample)
+ if self.act is not None:
+ sample = self.act(sample)
+ sample = self.linear_2(sample)
+ if self.post_act is not None:
+ sample = self.post_act(sample)
+ return sample
+
+
+class PixArtAlphaTextProjection(nn.Module):
+ def __init__(self, in_features, hidden_size, out_features=None, act_fn="gelu_tanh"):
+ super().__init__()
+ if out_features is None:
+ out_features = hidden_size
+ self.linear_1 = nn.Linear(in_features=in_features, out_features=hidden_size, bias=True)
+ if act_fn == "gelu_tanh":
+ self.act_1 = nn.GELU(approximate="tanh")
+ elif act_fn == "silu":
+ self.act_1 = nn.SiLU()
+ else:
+ self.act_1 = nn.GELU(approximate="tanh")
+ self.linear_2 = nn.Linear(in_features=hidden_size, out_features=out_features, bias=True)
+
+ def forward(self, caption):
+ hidden_states = self.linear_1(caption)
+ hidden_states = self.act_1(hidden_states)
+ hidden_states = self.linear_2(hidden_states)
+ return hidden_states
+
+
+class GELU(nn.Module):
+ def __init__(self, dim_in: int, dim_out: int, approximate: str = "none", bias: bool = True):
+ super().__init__()
+ self.proj = nn.Linear(dim_in, dim_out, bias=bias)
+ self.approximate = approximate
+
+ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+ hidden_states = self.proj(hidden_states)
+ hidden_states = F.gelu(hidden_states, approximate=self.approximate)
+ return hidden_states
+
+
+class FeedForward(nn.Module):
+ def __init__(
+ self,
+ dim: int,
+ dim_out: Optional[int] = None,
+ mult: int = 4,
+ dropout: float = 0.0,
+ activation_fn: str = "geglu",
+ final_dropout: bool = False,
+ inner_dim=None,
+ bias: bool = True,
+ ):
+ super().__init__()
+ if inner_dim is None:
+ inner_dim = int(dim * mult)
+ dim_out = dim_out if dim_out is not None else dim
+
+ # Build activation + projection matching diffusers pattern
+ if activation_fn == "gelu":
+ act_fn = GELU(dim, inner_dim, bias=bias)
+ elif activation_fn == "gelu-approximate":
+ act_fn = GELU(dim, inner_dim, approximate="tanh", bias=bias)
+ else:
+ act_fn = GELU(dim, inner_dim, bias=bias)
+
+ self.net = nn.ModuleList([])
+ self.net.append(act_fn)
+ self.net.append(nn.Dropout(dropout))
+ self.net.append(nn.Linear(inner_dim, dim_out, bias=bias))
+ if final_dropout:
+ self.net.append(nn.Dropout(dropout))
+
+ def forward(self, hidden_states: torch.Tensor, *args, **kwargs) -> torch.Tensor:
+ for module in self.net:
+ hidden_states = module(hidden_states)
+ return hidden_states
+
+
+def _to_tuple(x, dim=2):
+ if isinstance(x, int):
+ return (x,) * dim
+ elif len(x) == dim:
+ return x
+ else:
+ raise ValueError(f"Expected length {dim} or int, but got {x}")
+
+
+def get_meshgrid_nd(start, *args, dim=2):
+ if len(args) == 0:
+ num = _to_tuple(start, dim=dim)
+ start = (0,) * dim
+ stop = num
+ elif len(args) == 1:
+ start = _to_tuple(start, dim=dim)
+ stop = _to_tuple(args[0], dim=dim)
+ num = [stop[i] - start[i] for i in range(dim)]
+ elif len(args) == 2:
+ start = _to_tuple(start, dim=dim)
+ stop = _to_tuple(args[0], dim=dim)
+ num = _to_tuple(args[1], dim=dim)
+ else:
+ raise ValueError(f"len(args) should be 0, 1 or 2, but got {len(args)}")
+ axis_grid = []
+ for i in range(dim):
+ a, b, n = start[i], stop[i], num[i]
+ g = torch.linspace(a, b, n + 1, dtype=torch.float32)[:n]
+ axis_grid.append(g)
+ grid = torch.meshgrid(*axis_grid, indexing="ij")
+ grid = torch.stack(grid, dim=0)
+ return grid
+
+
+def reshape_for_broadcast(freqs_cis, x, head_first=False):
+ ndim = x.ndim
+ assert 0 <= 1 < ndim
+ if isinstance(freqs_cis, tuple):
+ if head_first:
+ assert freqs_cis[0].shape == (x.shape[-2], x.shape[-1])
+ shape = [d if i == ndim - 2 or i == ndim - 1 else 1 for i, d in enumerate(x.shape)]
+ else:
+ assert freqs_cis[0].shape == (x.shape[1], x.shape[-1])
+ shape = [d if i == 1 or i == ndim - 1 else 1 for i, d in enumerate(x.shape)]
+ return freqs_cis[0].view(*shape), freqs_cis[1].view(*shape)
+ else:
+ if head_first:
+ assert freqs_cis.shape == (x.shape[-2], x.shape[-1])
+ shape = [d if i == ndim - 2 or i == ndim - 1 else 1 for i, d in enumerate(x.shape)]
+ else:
+ assert freqs_cis.shape == (x.shape[1], x.shape[-1])
+ shape = [d if i == 1 or i == ndim - 1 else 1 for i, d in enumerate(x.shape)]
+ return freqs_cis.view(*shape)
+
+
+def rotate_half(x):
+ x_real, x_imag = x.float().reshape(*x.shape[:-1], -1, 2).unbind(-1)
+ return torch.stack([-x_imag, x_real], dim=-1).flatten(3)
+
+
+def apply_rotary_emb(xq, xk, freqs_cis, head_first=False):
+ cos, sin = reshape_for_broadcast(freqs_cis, xq, head_first)
+ cos, sin = cos.to(xq.device), sin.to(xq.device)
+ xq_out = (xq.float() * cos + rotate_half(xq.float()) * sin).type_as(xq)
+ xk_out = (xk.float() * cos + rotate_half(xk.float()) * sin).type_as(xk)
+ return xq_out, xk_out
+
+
+def get_1d_rotary_pos_embed(dim, pos, theta=10000.0, use_real=False, theta_rescale_factor=1.0, interpolation_factor=1.0):
+ if isinstance(pos, int):
+ pos = torch.arange(pos).float()
+ if theta_rescale_factor != 1.0:
+ theta *= theta_rescale_factor ** (dim / (dim - 2))
+ freqs = 1.0 / (theta ** (torch.arange(0, dim, 2)[: (dim // 2)].float() / dim))
+ freqs = torch.outer(pos * interpolation_factor, freqs)
+ if use_real:
+ freqs_cos = freqs.cos().repeat_interleave(2, dim=1)
+ freqs_sin = freqs.sin().repeat_interleave(2, dim=1)
+ return freqs_cos, freqs_sin
+ else:
+ return torch.polar(torch.ones_like(freqs), freqs)
+
+
+def get_nd_rotary_pos_embed(rope_dim_list, start, *args, theta=10000.0, use_real=False,
+ txt_rope_size=None, theta_rescale_factor=1.0, interpolation_factor=1.0):
+ grid = get_meshgrid_nd(start, *args, dim=len(rope_dim_list))
+ if isinstance(theta_rescale_factor, (int, float)):
+ theta_rescale_factor = [theta_rescale_factor] * len(rope_dim_list)
+ elif isinstance(theta_rescale_factor, list) and len(theta_rescale_factor) == 1:
+ theta_rescale_factor = [theta_rescale_factor[0]] * len(rope_dim_list)
+ if isinstance(interpolation_factor, (int, float)):
+ interpolation_factor = [interpolation_factor] * len(rope_dim_list)
+ elif isinstance(interpolation_factor, list) and len(interpolation_factor) == 1:
+ interpolation_factor = [interpolation_factor[0]] * len(rope_dim_list)
+ embs = []
+ for i in range(len(rope_dim_list)):
+ emb = get_1d_rotary_pos_embed(
+ rope_dim_list[i], grid[i].reshape(-1), theta,
+ use_real=use_real, theta_rescale_factor=theta_rescale_factor[i],
+ interpolation_factor=interpolation_factor[i],
+ )
+ embs.append(emb)
+ if use_real:
+ vis_emb = (torch.cat([emb[0] for emb in embs], dim=1), torch.cat([emb[1] for emb in embs], dim=1))
+ else:
+ vis_emb = torch.cat(embs, dim=1)
+ if txt_rope_size is not None:
+ embs_txt = []
+ vis_max_ids = grid.view(-1).max().item()
+ grid_txt = torch.arange(txt_rope_size) + vis_max_ids + 1
+ for i in range(len(rope_dim_list)):
+ emb = get_1d_rotary_pos_embed(
+ rope_dim_list[i], grid_txt, theta,
+ use_real=use_real, theta_rescale_factor=theta_rescale_factor[i],
+ interpolation_factor=interpolation_factor[i],
+ )
+ embs_txt.append(emb)
+ if use_real:
+ txt_emb = (torch.cat([emb[0] for emb in embs_txt], dim=1), torch.cat([emb[1] for emb in embs_txt], dim=1))
+ else:
+ txt_emb = torch.cat(embs_txt, dim=1)
+ else:
+ txt_emb = None
+ return vis_emb, txt_emb
+
+
+class ModulateWan(nn.Module):
+ def __init__(self, hidden_size: int, factor: int, dtype=None, device=None):
+ super().__init__()
+ self.factor = factor
+ factory_kwargs = {"dtype": dtype, "device": device}
+ self.modulate_table = nn.Parameter(
+ torch.zeros(1, factor, hidden_size, **factory_kwargs) / hidden_size**0.5,
+ requires_grad=True
+ )
+
+ def forward(self, x: torch.Tensor) -> torch.Tensor:
+ if len(x.shape) != 3:
+ x = x.unsqueeze(1)
+ return [o.squeeze(1) for o in (self.modulate_table + x).chunk(self.factor, dim=1)]
+
+
+def modulate(x, shift=None, scale=None):
+ if scale is None and shift is None:
+ return x
+ elif shift is None:
+ return x * (1 + scale.unsqueeze(1))
+ elif scale is None:
+ return x + shift.unsqueeze(1)
+ else:
+ return x * (1 + scale.unsqueeze(1)) + shift.unsqueeze(1)
+
+
+def apply_gate(x, gate=None, tanh=False):
+ if gate is None:
+ return x
+ if tanh:
+ return x * gate.unsqueeze(1).tanh()
+ else:
+ return x * gate.unsqueeze(1)
+
+
+def load_modulation(modulate_type: str, hidden_size: int, factor: int, dtype=None, device=None):
+ factory_kwargs = {"dtype": dtype, "device": device}
+ if modulate_type == 'wanx':
+ return ModulateWan(hidden_size, factor, **factory_kwargs)
+ raise ValueError(f"Unknown modulation type: {modulate_type}. Only 'wanx' is supported.")
+
+
+class RMSNorm(nn.Module):
+ def __init__(self, dim: int, elementwise_affine=True, eps: float = 1e-6, device=None, dtype=None):
+ factory_kwargs = {"device": device, "dtype": dtype}
+ super().__init__()
+ self.eps = eps
+ if elementwise_affine:
+ self.weight = nn.Parameter(torch.ones(dim, **factory_kwargs))
+
+ def _norm(self, x):
+ return x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps)
+
+ def forward(self, x):
+ output = self._norm(x.float()).type_as(x)
+ if hasattr(self, "weight"):
+ output = output * self.weight
+ return output
+
+
+class MMDoubleStreamBlock(nn.Module):
+ """
+ A multimodal dit block with separate modulation for
+ text and image/video, see more details (SD3): https://arxiv.org/abs/2403.03206
+ (Flux.1): https://github.com/black-forest-labs/flux
+ """
+
+ def __init__(
+ self,
+ hidden_size: int,
+ heads_num: int,
+ mlp_width_ratio: float,
+ mlp_act_type: str = "gelu_tanh",
+ dtype: Optional[torch.dtype] = None,
+ device: Optional[torch.device] = None,
+ dit_modulation_type: Optional[str] = "wanx",
+ ):
+ factory_kwargs = {"device": device, "dtype": dtype}
+ super().__init__()
+ self.dit_modulation_type = dit_modulation_type
+ self.heads_num = heads_num
+ head_dim = hidden_size // heads_num
+ mlp_hidden_dim = int(hidden_size * mlp_width_ratio)
+
+ self.img_mod = load_modulation(
+ modulate_type=self.dit_modulation_type,
+ hidden_size=hidden_size, factor=6, **factory_kwargs,
+ )
+ self.img_norm1 = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6, **factory_kwargs)
+ self.img_attn_qkv = nn.Linear(hidden_size, hidden_size * 3, bias=True, **factory_kwargs)
+ self.img_attn_q_norm = RMSNorm(head_dim, elementwise_affine=True, eps=1e-6, **factory_kwargs)
+ self.img_attn_k_norm = RMSNorm(head_dim, elementwise_affine=True, eps=1e-6, **factory_kwargs)
+ self.img_attn_proj = nn.Linear(hidden_size, hidden_size, bias=True, **factory_kwargs)
+ self.img_norm2 = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6, **factory_kwargs)
+ self.img_mlp = FeedForward(hidden_size, inner_dim=mlp_hidden_dim, activation_fn="gelu-approximate")
+
+ self.txt_mod = load_modulation(
+ modulate_type=self.dit_modulation_type,
+ hidden_size=hidden_size, factor=6, **factory_kwargs,
+ )
+ self.txt_norm1 = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6, **factory_kwargs)
+ self.txt_attn_qkv = nn.Linear(hidden_size, hidden_size * 3, bias=True, **factory_kwargs)
+ self.txt_attn_q_norm = RMSNorm(head_dim, elementwise_affine=True, eps=1e-6, **factory_kwargs)
+ self.txt_attn_k_norm = RMSNorm(head_dim, elementwise_affine=True, eps=1e-6, **factory_kwargs)
+ self.txt_attn_proj = nn.Linear(hidden_size, hidden_size, bias=True, **factory_kwargs)
+ self.txt_norm2 = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6, **factory_kwargs)
+ self.txt_mlp = FeedForward(hidden_size, inner_dim=mlp_hidden_dim, activation_fn="gelu-approximate")
+
+ def forward(
+ self,
+ img: torch.Tensor,
+ txt: torch.Tensor,
+ vec: torch.Tensor,
+ vis_freqs_cis: tuple = None,
+ txt_freqs_cis: tuple = None,
+ attn_kwargs: Optional[dict] = {},
+ ) -> Tuple[torch.Tensor, torch.Tensor]:
+ (
+ img_mod1_shift, img_mod1_scale, img_mod1_gate,
+ img_mod2_shift, img_mod2_scale, img_mod2_gate,
+ ) = self.img_mod(vec)
+ (
+ txt_mod1_shift, txt_mod1_scale, txt_mod1_gate,
+ txt_mod2_shift, txt_mod2_scale, txt_mod2_gate,
+ ) = self.txt_mod(vec)
+
+ img_modulated = self.img_norm1(img)
+ img_modulated = modulate(img_modulated, shift=img_mod1_shift, scale=img_mod1_scale)
+ img_qkv = self.img_attn_qkv(img_modulated)
+ img_q, img_k, img_v = rearrange(img_qkv, "B L (K H D) -> K B L H D", K=3, H=self.heads_num)
+ img_q = self.img_attn_q_norm(img_q).to(img_v)
+ img_k = self.img_attn_k_norm(img_k).to(img_v)
+
+ if vis_freqs_cis is not None:
+ img_qq, img_kk = apply_rotary_emb(img_q, img_k, vis_freqs_cis, head_first=False)
+ img_q, img_k = img_qq, img_kk
+
+ txt_modulated = self.txt_norm1(txt)
+ txt_modulated = modulate(txt_modulated, shift=txt_mod1_shift, scale=txt_mod1_scale)
+ txt_qkv = self.txt_attn_qkv(txt_modulated)
+ txt_q, txt_k, txt_v = rearrange(txt_qkv, "B L (K H D) -> K B L H D", K=3, H=self.heads_num)
+ txt_q = self.txt_attn_q_norm(txt_q).to(txt_v)
+ txt_k = self.txt_attn_k_norm(txt_k).to(txt_v)
+
+ if txt_freqs_cis is not None:
+ raise NotImplementedError("RoPE text is not supported for inference")
+
+ q = torch.cat((img_q, txt_q), dim=1)
+ k = torch.cat((img_k, txt_k), dim=1)
+ v = torch.cat((img_v, txt_v), dim=1)
+
+ # Use DiffSynth unified attention
+ attn_out = attention_forward(
+ q, k, v,
+ q_pattern="b s n d", k_pattern="b s n d", v_pattern="b s n d", out_pattern="b s n d",
+ )
+
+ attn_out = attn_out.flatten(2, 3)
+ img_attn, txt_attn = attn_out[:, : img.shape[1]], attn_out[:, img.shape[1]:]
+
+ img = img + apply_gate(self.img_attn_proj(img_attn), gate=img_mod1_gate)
+ img = img + apply_gate(
+ self.img_mlp(modulate(self.img_norm2(img), shift=img_mod2_shift, scale=img_mod2_scale)),
+ gate=img_mod2_gate,
+ )
+
+ txt = txt + apply_gate(self.txt_attn_proj(txt_attn), gate=txt_mod1_gate)
+ txt = txt + apply_gate(
+ self.txt_mlp(modulate(self.txt_norm2(txt), shift=txt_mod2_shift, scale=txt_mod2_scale)),
+ gate=txt_mod2_gate,
+ )
+
+ return img, txt
+
+
+class WanTimeTextImageEmbedding(nn.Module):
+ def __init__(
+ self,
+ dim: int,
+ time_freq_dim: int,
+ time_proj_dim: int,
+ text_embed_dim: int,
+ image_embed_dim: Optional[int] = None,
+ pos_embed_seq_len: Optional[int] = None,
+ ):
+ super().__init__()
+ self.timesteps_proj = Timesteps(num_channels=time_freq_dim, flip_sin_to_cos=True, downscale_freq_shift=0)
+ self.time_embedder = TimestepEmbedding(in_channels=time_freq_dim, time_embed_dim=dim)
+ self.act_fn = nn.SiLU()
+ self.time_proj = nn.Linear(dim, time_proj_dim)
+ self.text_embedder = PixArtAlphaTextProjection(text_embed_dim, dim, act_fn="gelu_tanh")
+
+ def forward(self, timestep: torch.Tensor, encoder_hidden_states: torch.Tensor):
+ timestep = self.timesteps_proj(timestep)
+ time_embedder_dtype = next(iter(self.time_embedder.parameters())).dtype
+ if timestep.dtype != time_embedder_dtype and time_embedder_dtype != torch.int8:
+ timestep = timestep.to(time_embedder_dtype)
+ temb = self.time_embedder(timestep).type_as(encoder_hidden_states)
+ timestep_proj = self.time_proj(self.act_fn(temb))
+ encoder_hidden_states = self.text_embedder(encoder_hidden_states)
+ return temb, timestep_proj, encoder_hidden_states
+
+
+class JoyAIImageDiT(nn.Module):
+ _supports_gradient_checkpointing = True
+
+ def __init__(
+ self,
+ patch_size: list = [1, 2, 2],
+ in_channels: int = 16,
+ out_channels: int = 16,
+ hidden_size: int = 4096,
+ heads_num: int = 32,
+ text_states_dim: int = 4096,
+ mlp_width_ratio: float = 4.0,
+ mm_double_blocks_depth: int = 40,
+ rope_dim_list: List[int] = [16, 56, 56],
+ rope_type: str = 'rope',
+ dtype: Optional[torch.dtype] = None,
+ device: Optional[torch.device] = None,
+ dit_modulation_type: str = "wanx",
+ theta: int = 10000,
+ ):
+ super().__init__()
+ self.out_channels = out_channels or in_channels
+ self.patch_size = patch_size
+ self.hidden_size = hidden_size
+ self.heads_num = heads_num
+ self.rope_dim_list = rope_dim_list
+ self.dit_modulation_type = dit_modulation_type
+ self.mm_double_blocks_depth = mm_double_blocks_depth
+ self.rope_type = rope_type
+ self.theta = theta
+
+ factory_kwargs = {"device": device, "dtype": dtype}
+
+ if hidden_size % heads_num != 0:
+ raise ValueError(f"Hidden size {hidden_size} must be divisible by heads_num {heads_num}")
+
+ self.img_in = nn.Conv3d(in_channels, hidden_size, kernel_size=patch_size, stride=patch_size)
+
+ self.condition_embedder = WanTimeTextImageEmbedding(
+ dim=hidden_size,
+ time_freq_dim=256,
+ time_proj_dim=hidden_size * 6,
+ text_embed_dim=text_states_dim,
+ )
+
+ self.double_blocks = nn.ModuleList([
+ MMDoubleStreamBlock(
+ self.hidden_size, self.heads_num,
+ mlp_width_ratio=mlp_width_ratio,
+ dit_modulation_type=self.dit_modulation_type,
+ **factory_kwargs,
+ )
+ for _ in range(mm_double_blocks_depth)
+ ])
+
+ self.norm_out = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6)
+ self.proj_out = nn.Linear(hidden_size, self.out_channels * math.prod(patch_size), **factory_kwargs)
+
+ def get_rotary_pos_embed(self, vis_rope_size, txt_rope_size=None):
+ target_ndim = 3
+ if len(vis_rope_size) != target_ndim:
+ vis_rope_size = [1] * (target_ndim - len(vis_rope_size)) + vis_rope_size
+ head_dim = self.hidden_size // self.heads_num
+ rope_dim_list = self.rope_dim_list
+ if rope_dim_list is None:
+ rope_dim_list = [head_dim // target_ndim for _ in range(target_ndim)]
+ assert sum(rope_dim_list) == head_dim
+ vis_freqs, txt_freqs = get_nd_rotary_pos_embed(
+ rope_dim_list, vis_rope_size,
+ txt_rope_size=txt_rope_size if self.rope_type == 'mrope' else None,
+ theta=self.theta, use_real=True, theta_rescale_factor=1,
+ )
+ return vis_freqs, txt_freqs
+
+ def forward(
+ self,
+ hidden_states: torch.Tensor,
+ timestep: torch.Tensor,
+ encoder_hidden_states: torch.Tensor = None,
+ encoder_hidden_states_mask: torch.Tensor = None,
+ return_dict: bool = True,
+ use_gradient_checkpointing: bool = False,
+ use_gradient_checkpointing_offload: bool = False,
+ ) -> Union[torch.Tensor, Dict[str, torch.Tensor]]:
+ is_multi_item = (len(hidden_states.shape) == 6)
+ num_items = 0
+ if is_multi_item:
+ num_items = hidden_states.shape[1]
+ if num_items > 1:
+ assert self.patch_size[0] == 1, "For multi-item input, patch_size[0] must be 1"
+ hidden_states = torch.cat([hidden_states[:, -1:], hidden_states[:, :-1]], dim=1)
+ hidden_states = rearrange(hidden_states, 'b n c t h w -> b c (n t) h w')
+
+ batch_size, _, ot, oh, ow = hidden_states.shape
+ tt, th, tw = ot // self.patch_size[0], oh // self.patch_size[1], ow // self.patch_size[2]
+
+ if encoder_hidden_states_mask is None:
+ encoder_hidden_states_mask = torch.ones(
+ (encoder_hidden_states.shape[0], encoder_hidden_states.shape[1]),
+ dtype=torch.bool,
+ ).to(encoder_hidden_states.device)
+
+ img = self.img_in(hidden_states).flatten(2).transpose(1, 2)
+ temb, vec, txt = self.condition_embedder(timestep, encoder_hidden_states)
+ if vec.shape[-1] > self.hidden_size:
+ vec = vec.unflatten(1, (6, -1))
+
+ txt_seq_len = txt.shape[1]
+ img_seq_len = img.shape[1]
+
+ vis_freqs_cis, txt_freqs_cis = self.get_rotary_pos_embed(
+ vis_rope_size=(tt, th, tw),
+ txt_rope_size=txt_seq_len if self.rope_type == 'mrope' else None,
+ )
+
+ for block in self.double_blocks:
+ img, txt = gradient_checkpoint_forward(
+ block,
+ use_gradient_checkpointing=use_gradient_checkpointing,
+ use_gradient_checkpointing_offload=use_gradient_checkpointing_offload,
+ img=img, txt=txt, vec=vec,
+ vis_freqs_cis=vis_freqs_cis, txt_freqs_cis=txt_freqs_cis,
+ attn_kwargs={},
+ )
+
+ img_len = img.shape[1]
+ x = torch.cat((img, txt), 1)
+ img = x[:, :img_len, ...]
+
+ img = self.proj_out(self.norm_out(img))
+ img = self.unpatchify(img, tt, th, tw)
+
+ if is_multi_item:
+ img = rearrange(img, 'b c (n t) h w -> b n c t h w', n=num_items)
+ if num_items > 1:
+ img = torch.cat([img[:, 1:], img[:, :1]], dim=1)
+
+ return img
+
+ def unpatchify(self, x, t, h, w):
+ c = self.out_channels
+ pt, ph, pw = self.patch_size
+ assert t * h * w == x.shape[1]
+ x = x.reshape(shape=(x.shape[0], t, h, w, pt, ph, pw, c))
+ x = torch.einsum("nthwopqc->nctohpwq", x)
+ return x.reshape(shape=(x.shape[0], c, t * pt, h * ph, w * pw))
diff --git a/diffsynth/models/joyai_image_text_encoder.py b/diffsynth/models/joyai_image_text_encoder.py
new file mode 100644
index 0000000000000000000000000000000000000000..292c5f6bd472a9ac909b6c7f0a60d0f1a780e096
--- /dev/null
+++ b/diffsynth/models/joyai_image_text_encoder.py
@@ -0,0 +1,82 @@
+import torch
+from typing import Optional
+
+
+class JoyAIImageTextEncoder(torch.nn.Module):
+ def __init__(self):
+ super().__init__()
+ from transformers import Qwen3VLConfig, Qwen3VLForConditionalGeneration
+
+ config = Qwen3VLConfig(
+ text_config={
+ "attention_bias": False,
+ "attention_dropout": 0.0,
+ "bos_token_id": 151643,
+ "eos_token_id": 151645,
+ "head_dim": 128,
+ "hidden_act": "silu",
+ "hidden_size": 4096,
+ "initializer_range": 0.02,
+ "intermediate_size": 12288,
+ "max_position_embeddings": 262144,
+ "model_type": "qwen3_vl_text",
+ "num_attention_heads": 32,
+ "num_hidden_layers": 36,
+ "num_key_value_heads": 8,
+ "rms_norm_eps": 1e-6,
+ "rope_scaling": {
+ "mrope_interleaved": True,
+ "mrope_section": [24, 20, 20],
+ "rope_type": "default",
+ },
+ "rope_theta": 5000000,
+ "use_cache": True,
+ "vocab_size": 151936,
+ },
+ vision_config={
+ "deepstack_visual_indexes": [8, 16, 24],
+ "depth": 27,
+ "hidden_act": "gelu_pytorch_tanh",
+ "hidden_size": 1152,
+ "in_channels": 3,
+ "initializer_range": 0.02,
+ "intermediate_size": 4304,
+ "model_type": "qwen3_vl",
+ "num_heads": 16,
+ "num_position_embeddings": 2304,
+ "out_hidden_size": 4096,
+ "patch_size": 16,
+ "spatial_merge_size": 2,
+ "temporal_patch_size": 2,
+ },
+ image_token_id=151655,
+ video_token_id=151656,
+ vision_start_token_id=151652,
+ vision_end_token_id=151653,
+ tie_word_embeddings=False,
+ )
+
+ self.model = Qwen3VLForConditionalGeneration(config)
+ self.config = config
+
+ def forward(
+ self,
+ input_ids: torch.LongTensor = None,
+ attention_mask: Optional[torch.Tensor] = None,
+ pixel_values: Optional[torch.Tensor] = None,
+ image_grid_thw: Optional[torch.LongTensor] = None,
+ **kwargs,
+ ):
+ pre_norm_output = [None]
+ def hook_fn(module, args, kwargs_output=None):
+ pre_norm_output[0] = args[0]
+ self.model.model.language_model.norm.register_forward_hook(hook_fn)
+ _ = self.model(
+ input_ids=input_ids,
+ pixel_values=pixel_values,
+ image_grid_thw=image_grid_thw,
+ attention_mask=attention_mask,
+ output_hidden_states=True,
+ **kwargs,
+ )
+ return pre_norm_output[0]
\ No newline at end of file
diff --git a/diffsynth/models/krea2_dit.py b/diffsynth/models/krea2_dit.py
new file mode 100644
index 0000000000000000000000000000000000000000..a0e2c0bbfa4bc9dafe6b68ff4bec200c2ba7f0fe
--- /dev/null
+++ b/diffsynth/models/krea2_dit.py
@@ -0,0 +1,419 @@
+import math
+from dataclasses import dataclass
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from einops import rearrange
+from torch import Tensor
+from torch.nn.attention import SDPBackend, sdpa_kernel
+from ..core.gradient import gradient_checkpoint_forward
+
+
+def rope(pos: Tensor, dim: int, theta: float = 1e4, ntk: float = 1.0) -> Tensor:
+ scale = torch.arange(0, dim, 2, dtype=torch.float64, device=pos.device) / dim
+ omega = 1.0 / ((theta * ntk) ** scale)
+ out = torch.einsum("...n,d->...nd", pos, omega)
+ out = torch.stack(
+ [torch.cos(out), -torch.sin(out), torch.sin(out), torch.cos(out)], dim=-1
+ )
+ out = rearrange(out, "b n d (i j) -> b n d i j", i=2, j=2)
+ return out.float()
+
+
+def ropeapply(xq: Tensor, xk: Tensor, freqs: Tensor) -> tuple[Tensor, Tensor]:
+ xq_ = xq.float().reshape(*xq.shape[:-1], -1, 1, 2)
+ xk_ = xk.float().reshape(*xk.shape[:-1], -1, 1, 2)
+ freqs = freqs[:, None, :, :, :]
+ xq_ = freqs[..., 0] * xq_[..., 0] + freqs[..., 1] * xq_[..., 1]
+ xk_ = freqs[..., 0] * xk_[..., 0] + freqs[..., 1] * xk_[..., 1]
+ return xq_.reshape(*xq.shape).to(xq.dtype), xk_.reshape(*xk.shape).to(xk.dtype)
+
+
+def attention(
+ q: Tensor,
+ k: Tensor,
+ v: Tensor,
+ mask: Tensor | None = None,
+ scale: float | None = None,
+ gqa: bool = False,
+) -> Tensor:
+ with sdpa_kernel(SDPBackend.CUDNN_ATTENTION):
+ x = F.scaled_dot_product_attention(
+ q, k, v, attn_mask=mask, scale=scale, enable_gqa=gqa
+ )
+ return rearrange(x, "B H L D -> B L (H D)")
+
+
+def _mask(mask: Tensor) -> Tensor:
+ return mask.unsqueeze(1).unsqueeze(2) * mask.unsqueeze(1).unsqueeze(3)
+
+
+def temb(
+ t: Tensor,
+ dim: int,
+ period: float = 1e4,
+ tfactor: float = 1e3,
+ device: torch.device = None,
+ dtype: torch.dtype = None,
+) -> Tensor:
+ half = dim // 2
+ freqs = torch.exp(
+ -math.log(period)
+ * torch.arange(half, dtype=torch.float32, device=device)
+ / half
+ )
+ args = (t.float() * tfactor)[:, None, None] * freqs
+ sin, cos = torch.sin(args), torch.cos(args)
+ return torch.cat((cos, sin), dim=-1).to(dtype=dtype)
+
+
+@dataclass
+class SingleMMDiTConfig:
+ features: int
+ tdim: int
+ txtdim: int
+ heads: int
+ multiplier: int
+ layers: int
+ patch: int
+ channels: int
+ bias: bool = False
+ theta: float = 1e3
+ kvheads: int | None = None
+ txtlayers: int = 1
+ txtheads: int = 20
+ txtkvheads: int = 20
+
+ def get(self, *args, **kwargs):
+ # For compatibility with low-version peft
+ return None
+
+
+class SimpleModulation(torch.nn.Module):
+ def __init__(self, dim: int):
+ super().__init__()
+ self.lin = torch.nn.Parameter(torch.zeros(2, dim))
+ self.multiplier = 2
+
+ def forward(self, vec: Tensor):
+ out = vec + rearrange(self.lin, "two d -> 1 two d")
+ scale, shift = out.chunk(self.multiplier, dim=1)
+ return scale, shift
+
+
+class DoubleSharedModulation(torch.nn.Module):
+ def __init__(self, dim: int):
+ super().__init__()
+ self.lin = torch.nn.Parameter(torch.zeros(6 * dim))
+
+ def forward(self, vec: Tensor):
+ out = vec + self.lin
+ prescale, preshift, pregate, postscale, postshift, postgate = out.chunk(
+ 6, dim=-1
+ )
+ return prescale, preshift, pregate, postscale, postshift, postgate
+
+
+class PositionalEncoding(torch.nn.Module):
+ def __init__(self, dim, axdims: list[int], theta: float = 1e2, ntk: float = 1.0):
+ super().__init__()
+ self.axdims = axdims
+ self.theta = theta
+ self.ntk = ntk
+
+ def forward(self, pos: Tensor) -> Tensor:
+ return torch.cat(
+ [
+ rope(pos[..., i], d, self.theta, self.ntk)
+ for i, d in enumerate(self.axdims)
+ ],
+ dim=-3,
+ )
+
+
+class QKNorm(torch.nn.Module):
+ def __init__(self, dim: int):
+ super().__init__()
+ self.qnorm = RMSNorm(dim)
+ self.knorm = RMSNorm(dim)
+
+ def forward(self, q: Tensor, k: Tensor, v: Tensor) -> tuple[Tensor, Tensor, Tensor]:
+ return self.qnorm(q), self.knorm(k), v
+
+
+class RMSNorm(torch.nn.Module):
+ def __init__(self, features: int, eps: float = 1e-05, device: torch.device = None):
+ super().__init__()
+ self.features = features
+ self.eps = eps
+ self.scale = torch.nn.Parameter(
+ torch.zeros(features, device=device, dtype=torch.float32)
+ )
+
+ def forward(self, x: Tensor) -> Tensor:
+ t, dtype = x.float(), x.dtype
+ t = F.rms_norm(
+ t, (self.features,), eps=self.eps, weight=(self.scale.float() + 1.0)
+ )
+ return t.to(dtype)
+
+
+class SwiGLU(torch.nn.Module):
+ def __init__(
+ self, features: int, multiplier: int, bias: bool = False, multiple: int = 128
+ ):
+ super().__init__()
+ mlpdim = int(2 * features / 3) * multiplier
+ mlpdim = multiple * ((mlpdim + multiple - 1) // multiple)
+ self.gate = torch.nn.Linear(features, mlpdim, bias=bias)
+ self.up = torch.nn.Linear(features, mlpdim, bias=bias)
+ self.down = torch.nn.Linear(mlpdim, features, bias=bias)
+
+ def forward(self, x: Tensor) -> Tensor:
+ return self.down(F.silu(self.gate(x)) * self.up(x))
+
+
+class Attention(torch.nn.Module):
+ def __init__(self, dim: int, heads: int, kvheads: int = None, bias: bool = False):
+ super().__init__()
+ self.heads = heads
+ self.kvheads = kvheads if kvheads is not None else heads
+ self.headdim = dim // self.heads
+ self.wq = torch.nn.Linear(dim, self.headdim * self.heads, bias=bias)
+ self.wk = torch.nn.Linear(dim, self.headdim * self.kvheads, bias=bias)
+ self.wv = torch.nn.Linear(dim, self.headdim * self.kvheads, bias=bias)
+ self.gate = torch.nn.Linear(dim, dim, bias=bias)
+ self.qknorm = QKNorm(self.headdim)
+ self.gqa = self.heads != self.kvheads
+ self.wo = torch.nn.Linear(dim, dim, bias=bias)
+
+ def forward(
+ self, qkv: Tensor, freqs: Tensor | None = None, mask: Tensor | None = None
+ ) -> Tensor:
+ q, k, v, gate = self.wq(qkv), self.wk(qkv), self.wv(qkv), self.gate(qkv)
+ q, k, v = (
+ rearrange(q, "B L (H D) -> B H L D", H=self.heads),
+ rearrange(k, "B L (H D) -> B H L D", H=self.kvheads),
+ rearrange(v, "B L (H D) -> B H L D", H=self.kvheads),
+ )
+ q, k, v = self.qknorm(q, k, v)
+ if freqs is not None:
+ q, k = ropeapply(q, k, freqs)
+ out = self.wo(attention(q, k, v, mask=mask, gqa=self.gqa) * F.sigmoid(gate))
+ return out
+
+
+class LastLayer(torch.nn.Module):
+ def __init__(self, features: int, patch: int, channels: int):
+ super().__init__()
+ self.norm = RMSNorm(features)
+ self.linear = torch.nn.Linear(features, patch * patch * channels, bias=True)
+ self.modulation = SimpleModulation(features)
+
+ def forward(self, x: Tensor, tvec: Tensor) -> Tensor:
+ scale, shift = self.modulation(tvec)
+ x = (1 + scale) * self.norm(x) + shift
+ x = self.linear(x)
+ return x
+
+
+class TextFusionBlock(torch.nn.Module):
+ def __init__(
+ self,
+ features: int,
+ heads: int,
+ multiplier: int,
+ bias: bool = False,
+ kvheads: int = None,
+ ):
+ super().__init__()
+ self.prenorm = RMSNorm(features)
+ self.postnorm = RMSNorm(features)
+ self.attn = Attention(dim=features, heads=heads, bias=bias, kvheads=kvheads)
+ self.mlp = SwiGLU(features, multiplier, bias)
+
+ def forward(self, x: Tensor, mask: Tensor | None = None) -> Tensor:
+ x = x + self.attn(self.prenorm(x), mask=mask)
+ x = x + self.mlp(self.postnorm(x))
+ return x
+
+
+class TextFusionTransformer(torch.nn.Module):
+ def __init__(
+ self,
+ num_txt_layers: int,
+ txt_dim: int,
+ heads: int,
+ multiplier: int,
+ bias: bool = False,
+ kvheads: int = None,
+ ):
+ super().__init__()
+ self.layerwise_blocks = torch.nn.ModuleList(
+ [
+ TextFusionBlock(txt_dim, heads, multiplier, bias, kvheads)
+ for _ in range(2)
+ ]
+ )
+ self.projector = torch.nn.Linear(num_txt_layers, 1, bias=False)
+ self.refiner_blocks = torch.nn.ModuleList(
+ [
+ TextFusionBlock(txt_dim, heads, multiplier, bias, kvheads)
+ for _ in range(2)
+ ]
+ )
+
+ def forward(self, x: Tensor, mask: Tensor | None = None, use_gradient_checkpointing=False, use_gradient_checkpointing_offload=False) -> Tensor:
+ b, l, n, d = x.shape
+ x = x.reshape(b * l, n, d)
+ for block in self.layerwise_blocks:
+ x = gradient_checkpoint_forward(
+ block,
+ use_gradient_checkpointing=use_gradient_checkpointing,
+ use_gradient_checkpointing_offload=use_gradient_checkpointing_offload,
+ x=x.contiguous(),
+ mask=None
+ )
+ x = rearrange(x, "(b l) n d -> b l d n", b=b, l=l)
+ x = self.projector(x)
+ x = x.squeeze(-1)
+ for block in self.refiner_blocks:
+ x = block(x, mask=mask)
+ return x
+
+
+class SingleStreamBlock(nn.Module):
+ def __init__(
+ self,
+ features: int,
+ heads: int,
+ multiplier: int,
+ bias: bool = False,
+ kvheads: int = None,
+ ):
+ super().__init__()
+ self.mod = DoubleSharedModulation(features)
+ self.prenorm = RMSNorm(features)
+ self.postnorm = RMSNorm(features)
+ self.attn = Attention(dim=features, heads=heads, bias=bias, kvheads=kvheads)
+ self.mlp = SwiGLU(features, multiplier, bias)
+
+ def forward(
+ self, x: Tensor, vec: Tensor, freqs: Tensor, mask: Tensor | None = None
+ ) -> Tensor:
+ prescale, preshift, pregate, postscale, postshift, postgate = self.mod(vec)
+ x = x + pregate * self.attn(
+ (1 + prescale) * self.prenorm(x) + preshift, freqs, mask
+ )
+ x = x + postgate * self.mlp((1 + postscale) * self.postnorm(x) + postshift)
+ return x
+
+
+class SingleStreamDiT(nn.Module):
+ def __init__(
+ self,
+ config={"features": 6144, "tdim": 256, "txtdim": 2560, "heads": 48, "kvheads": 12, "multiplier": 4, "layers": 28, "patch": 2, "channels": 16, "txtheads": 20, "txtkvheads": 20, "txtlayers": 12},
+ **kwargs
+ ):
+ super().__init__()
+ if config is None:
+ config = SingleMMDiTConfig(**kwargs)
+ elif isinstance(config, dict):
+ config = SingleMMDiTConfig(**config)
+ self.config = config
+ headdim = config.features // config.heads
+ axes = [
+ headdim - 12 * (headdim // 16),
+ 6 * (headdim // 16),
+ 6 * (headdim // 16),
+ ]
+ assert sum(axes) == headdim, f"sum(axes) = {sum(axes)}, headdim = {headdim}"
+ assert all(a % 2 == 0 for a in axes), f"axes = {axes}"
+ self.posemb = PositionalEncoding(
+ config.features, axes, theta=config.theta, ntk=1.0
+ )
+ self.first = nn.Linear(
+ config.channels * config.patch**2, config.features, bias=True
+ )
+ self.blocks = nn.ModuleList(
+ [
+ SingleStreamBlock(
+ config.features,
+ config.heads,
+ config.multiplier,
+ config.bias,
+ config.kvheads,
+ )
+ for _ in range(config.layers)
+ ]
+ )
+ self.tmlp = nn.Sequential(
+ nn.Linear(config.tdim, config.features),
+ nn.GELU(approximate="tanh"),
+ nn.Linear(config.features, config.features),
+ )
+ self.txtfusion = TextFusionTransformer(
+ config.txtlayers,
+ config.txtdim,
+ config.txtheads,
+ config.multiplier,
+ config.bias,
+ config.txtkvheads,
+ )
+ self.txtmlp = nn.Sequential(
+ RMSNorm(config.txtdim),
+ nn.Linear(config.txtdim, config.features),
+ nn.GELU(approximate="tanh"),
+ nn.Linear(config.features, config.features),
+ )
+ self.last = LastLayer(config.features, config.patch, config.channels)
+ self.tproj = nn.Sequential(
+ nn.GELU(approximate="tanh"), nn.Linear(config.features, config.features * 6)
+ )
+
+ def forward(
+ self,
+ img: Tensor,
+ context: Tensor,
+ t: Tensor,
+ pos: Tensor,
+ mask: Tensor | None = None,
+ context_pre_compute: bool = False,
+ use_gradient_checkpointing: bool = False,
+ use_gradient_checkpointing_offload: bool = False,
+ ) -> Tensor:
+ img = self.first(img)
+ t = self.tmlp(temb(t, self.config.tdim, device=img.device, dtype=img.dtype))
+ tvec = self.tproj(t)
+ if not context_pre_compute:
+ txtmask = _mask(mask[:, : context.shape[1]])
+ context = self.txtfusion(
+ context, mask=txtmask,
+ use_gradient_checkpointing=use_gradient_checkpointing,
+ use_gradient_checkpointing_offload=use_gradient_checkpointing_offload,
+ )
+ context = self.txtmlp(context)
+ txtlen, imglen = context.shape[1], img.shape[1]
+ combined = torch.cat((context, img), dim=1)
+ fulllen = combined.shape[1]
+ _padlen = (-fulllen) % 256
+ if _padlen > 0:
+ combined = F.pad(combined, (0, 0, 0, _padlen))
+ mask = F.pad(mask, (0, _padlen), value=False)
+ pos = F.pad(pos, (0, 0, 0, _padlen))
+ mask = _mask(mask)
+ freqs = self.posemb(pos)
+ for block in self.blocks:
+ combined = gradient_checkpoint_forward(
+ block,
+ use_gradient_checkpointing=use_gradient_checkpointing,
+ use_gradient_checkpointing_offload=use_gradient_checkpointing_offload,
+ x=combined,
+ vec=tvec,
+ freqs=freqs,
+ mask=mask,
+ )
+ final = self.last(combined, t)
+ output = final[:, txtlen : txtlen + imglen, :]
+ return output
diff --git a/diffsynth/models/krea2_text_encoder.py b/diffsynth/models/krea2_text_encoder.py
new file mode 100644
index 0000000000000000000000000000000000000000..0671746312163036687d590eb5913c731bbe5f9d
--- /dev/null
+++ b/diffsynth/models/krea2_text_encoder.py
@@ -0,0 +1,79 @@
+import torch
+from torch import Tensor
+
+
+class Krea2TextEncoder(torch.nn.Module):
+ def __init__(self):
+ super().__init__()
+ from transformers import Qwen3VLConfig, Qwen3VLForConditionalGeneration
+ config = Qwen3VLConfig(**{
+ "architectures": ["Qwen3VLForConditionalGeneration"],
+ "image_token_id": 151655,
+ "model_type": "qwen3_vl",
+ "text_config": {
+ "attention_bias": False,
+ "attention_dropout": 0.0,
+ "bos_token_id": 151643,
+ "dtype": "bfloat16",
+ "eos_token_id": 151645,
+ "head_dim": 128,
+ "hidden_act": "silu",
+ "hidden_size": 2560,
+ "initializer_range": 0.02,
+ "intermediate_size": 9728,
+ "max_position_embeddings": 262144,
+ "model_type": "qwen3_vl_text",
+ "num_attention_heads": 32,
+ "num_hidden_layers": 36,
+ "num_key_value_heads": 8,
+ "rms_norm_eps": 1e-06,
+ "rope_scaling": {
+ "mrope_interleaved": True,
+ "mrope_section": [24, 20, 20],
+ "rope_type": "default",
+ },
+ "rope_theta": 5000000,
+ "tie_word_embeddings": True,
+ "use_cache": True,
+ "vocab_size": 151936,
+ },
+ "tie_word_embeddings": True,
+ "transformers_version": "4.57.0.dev0",
+ "video_token_id": 151656,
+ "vision_config": {
+ "deepstack_visual_indexes": [5, 11, 17],
+ "depth": 24,
+ "hidden_act": "gelu_pytorch_tanh",
+ "hidden_size": 1024,
+ "in_channels": 3,
+ "initializer_range": 0.02,
+ "intermediate_size": 4096,
+ "model_type": "qwen3_vl",
+ "num_heads": 16,
+ "num_position_embeddings": 2304,
+ "out_hidden_size": 2560,
+ "patch_size": 16,
+ "spatial_merge_size": 2,
+ "temporal_patch_size": 2,
+ },
+ "vision_end_token_id": 151653,
+ "vision_start_token_id": 151652,
+ })
+ self.model = Qwen3VLForConditionalGeneration(config)
+ self.config = config
+
+ def forward(
+ self,
+ input_ids=None,
+ attention_mask=None,
+ output_hidden_states=True,
+ **kwargs,
+ ):
+ outputs = self.model(
+ input_ids=input_ids,
+ attention_mask=attention_mask,
+ output_hidden_states=True,
+ return_dict=True,
+ **kwargs,
+ )
+ return outputs.hidden_states
diff --git a/diffsynth/models/lingbot_video_dit.py b/diffsynth/models/lingbot_video_dit.py
new file mode 100644
index 0000000000000000000000000000000000000000..9fbe44df14cddf69aad9c99e0ff80f9359824588
--- /dev/null
+++ b/diffsynth/models/lingbot_video_dit.py
@@ -0,0 +1,623 @@
+import math
+from typing import Optional, Tuple
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from ..core.attention import attention_forward
+from ..core.gradient import gradient_checkpoint_forward
+from ..core.device.npu_compatible_device import get_device_type
+
+
+def get_timestep_embedding(
+ timesteps: torch.Tensor,
+ embedding_dim: int,
+ flip_sin_to_cos: bool = True,
+ downscale_freq_shift: float = 0.0,
+ scale: float = 1.0,
+ max_period: int = 10000,
+) -> torch.Tensor:
+ """Sinusoidal timestep embedding."""
+ assert timesteps.ndim == 1
+ half_dim = embedding_dim // 2
+ exponent = -math.log(max_period) * torch.arange(
+ start=0, end=half_dim, dtype=torch.float32, device=timesteps.device
+ )
+ exponent = exponent / (half_dim - downscale_freq_shift)
+ emb = torch.exp(exponent)
+ emb = timesteps[:, None].float() * emb[None, :]
+ emb = scale * emb
+ emb = torch.cat([torch.sin(emb), torch.cos(emb)], dim=-1)
+ if flip_sin_to_cos:
+ emb = torch.cat([emb[:, half_dim:], emb[:, :half_dim]], dim=-1)
+ if embedding_dim % 2 == 1:
+ emb = F.pad(emb, (0, 1, 0, 0))
+ return emb
+
+
+class Timesteps(nn.Module):
+ """Parameter-free sinusoidal timestep projection."""
+
+ def __init__(self, num_channels: int, flip_sin_to_cos: bool = True, downscale_freq_shift: float = 0.0, scale: int = 1):
+ super().__init__()
+ self.num_channels = num_channels
+ self.flip_sin_to_cos = flip_sin_to_cos
+ self.downscale_freq_shift = downscale_freq_shift
+ self.scale = scale
+
+ def forward(self, timesteps: torch.Tensor) -> torch.Tensor:
+ return get_timestep_embedding(
+ timesteps,
+ self.num_channels,
+ flip_sin_to_cos=self.flip_sin_to_cos,
+ downscale_freq_shift=self.downscale_freq_shift,
+ scale=self.scale,
+ )
+
+
+class TimestepEmbedding(nn.Module):
+ """Two-layer timestep MLP (act_fn='silu')."""
+
+ def __init__(self, in_channels: int, time_embed_dim: int, sample_proj_bias: bool = True):
+ super().__init__()
+ self.linear_1 = nn.Linear(in_channels, time_embed_dim, bias=sample_proj_bias)
+ self.act = nn.SiLU()
+ self.linear_2 = nn.Linear(time_embed_dim, time_embed_dim, bias=sample_proj_bias)
+
+ def forward(self, sample: torch.Tensor) -> torch.Tensor:
+ return self.linear_2(self.act(self.linear_1(sample)))
+
+
+class LingBotVideoRMSNorm(nn.Module):
+ """RMSNorm with fp32 accumulation."""
+
+ def __init__(self, dim: int, eps: float = 1e-6):
+ super().__init__()
+ self.weight = nn.Parameter(torch.ones(dim))
+ self.variance_epsilon = eps
+
+ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+ input_dtype = hidden_states.dtype
+ hidden_states = hidden_states.to(torch.float32)
+ variance = hidden_states.pow(2).mean(-1, keepdim=True)
+ hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon)
+ return (self.weight * hidden_states).to(input_dtype)
+
+
+def apply_rotary_emb(x: torch.Tensor, freqs_cis: torch.Tensor) -> torch.Tensor:
+ """Apply complex RoPE to `(B, S, H, D)` attention tensors."""
+ with torch.amp.autocast(get_device_type(), enabled=False):
+ x_c = torch.view_as_complex(x.float().reshape(*x.shape[:-1], -1, 2))
+ out = torch.view_as_real(x_c * freqs_cis.unsqueeze(2)).flatten(3)
+ return out.type_as(x)
+
+
+class LingBotVideoRotaryEmbedding(nn.Module):
+ """Complex64 RoPE table indexed by 3D position ids. Holds no persistent state."""
+
+ def __init__(self, axes_dims: Tuple[int, ...], axes_lens: Tuple[int, ...], theta: float):
+ super().__init__()
+ self.axes_dims = tuple(axes_dims)
+ self.axes_lens = list(axes_lens)
+ self.theta = theta
+ self.freqs_cis = None
+
+ @staticmethod
+ def precompute_freqs_cis(dim: Tuple[int, ...], end: Tuple[int, ...], theta: float):
+ freqs_cis = []
+ for d, e in zip(dim, end):
+ freqs = 1.0 / (theta ** (torch.arange(0, d, 2, dtype=torch.float64, device="cpu") / d))
+ timestep = torch.arange(e, device=freqs.device, dtype=torch.float64)
+ freqs = torch.outer(timestep, freqs).float()
+ freqs_cis.append(torch.polar(torch.ones_like(freqs), freqs).to(torch.complex64))
+ return freqs_cis
+
+ def forward(self, position_ids: torch.Tensor) -> torch.Tensor:
+ # position_ids: (S, 3) int -> (S, head_dim/2) complex64
+ device = position_ids.device
+ max_vals = position_ids.max(dim=0).values.tolist()
+ needs_rebuild = self.freqs_cis is None or any(m >= l for m, l in zip(max_vals, self.axes_lens))
+ if needs_rebuild:
+ for i in range(len(self.axes_lens)):
+ if max_vals[i] >= self.axes_lens[i]:
+ self.axes_lens[i] = int(max_vals[i] * 1.5) + 1
+ self.freqs_cis = self.precompute_freqs_cis(self.axes_dims, tuple(self.axes_lens), theta=self.theta)
+ self.freqs_cis = [freqs_cis.to(device) for freqs_cis in self.freqs_cis]
+ elif self.freqs_cis[0].device != device:
+ self.freqs_cis = [freqs_cis.to(device) for freqs_cis in self.freqs_cis]
+ return torch.cat([self.freqs_cis[i][position_ids[:, i]] for i in range(len(self.axes_dims))], dim=-1)
+
+
+def make_joint_position_ids(text_len: int, grid_t: int, grid_h: int, grid_w: int, device: torch.device) -> torch.Tensor:
+ """3D positions in [video; text] order. Text t-axis is 1..text_len; video t-axis starts at text_len+1."""
+ tt = torch.arange(grid_t, device=device, dtype=torch.int32) + (text_len + 1)
+ hh = torch.arange(grid_h, device=device, dtype=torch.int32)
+ ww = torch.arange(grid_w, device=device, dtype=torch.int32)
+ grid = torch.stack(torch.meshgrid(tt, hh, ww, indexing="ij"), dim=-1).flatten(0, 2)
+ text_t = torch.arange(text_len, device=device, dtype=torch.int32) + 1
+ text_pos = torch.stack([text_t, torch.zeros_like(text_t), torch.zeros_like(text_t)], dim=-1)
+ return torch.cat([grid, text_pos], dim=0) # (Nx + L, 3)
+
+
+def _cat_interleave(a: torch.Tensor, len_a: list, b: torch.Tensor, len_b: list) -> torch.Tensor:
+ a_split = torch.split(a, len_a, dim=1)
+ b_split = torch.split(b, len_b, dim=1)
+ blocks = []
+ for x_part, text_part in zip(a_split, b_split):
+ blocks.extend([x_part, text_part])
+ return torch.cat(blocks, dim=1)
+
+
+class LingBotVideoTextEmbedder(nn.Module):
+ """RMSNorm(text_dim) -> Linear -> SiLU -> Linear."""
+
+ def __init__(self, text_dim: int, hidden_size: int):
+ super().__init__()
+ self.norm = LingBotVideoRMSNorm(text_dim, eps=1e-6)
+ self.linear_1 = nn.Linear(text_dim, hidden_size, bias=True)
+ self.linear_2 = nn.Linear(hidden_size, hidden_size, bias=True)
+
+ def forward(self, x: torch.Tensor) -> torch.Tensor:
+ x = self.norm(x)
+ return self.linear_2(F.silu(self.linear_1(x)))
+
+
+class LingBotVideoAttention(nn.Module):
+ def __init__(self, hidden_size, num_heads, norm_eps, qkv_bias, out_bias):
+ super().__init__()
+ self.num_heads = num_heads
+ self.head_dim = hidden_size // num_heads
+ self.to_q = nn.Linear(hidden_size, hidden_size, bias=qkv_bias)
+ self.to_k = nn.Linear(hidden_size, hidden_size, bias=qkv_bias)
+ self.to_v = nn.Linear(hidden_size, hidden_size, bias=qkv_bias)
+ self.norm_q = LingBotVideoRMSNorm(self.head_dim, norm_eps)
+ self.norm_k = LingBotVideoRMSNorm(self.head_dim, norm_eps)
+ self.to_out = nn.Linear(hidden_size, hidden_size, bias=out_bias)
+
+ def forward(self, x, rotary_emb, attention_mask=None):
+ q = self.to_q(x).unflatten(2, (self.num_heads, self.head_dim))
+ k = self.to_k(x).unflatten(2, (self.num_heads, self.head_dim))
+ v = self.to_v(x).unflatten(2, (self.num_heads, self.head_dim))
+ q = apply_rotary_emb(self.norm_q(q), rotary_emb)
+ k = apply_rotary_emb(self.norm_k(k), rotary_emb)
+ # q/k/v are (B, S, H, D).
+ out = attention_forward(
+ q, k, v,
+ q_pattern="b s n d", k_pattern="b s n d", v_pattern="b s n d", out_pattern="b s n d",
+ attn_mask=attention_mask,
+ )
+ return self.to_out(out.flatten(2, 3).type_as(x))
+
+
+class LingBotVideoMLP(nn.Module):
+ def __init__(self, hidden_size, intermediate_size):
+ super().__init__()
+ self.gate_proj = nn.Linear(hidden_size, intermediate_size, bias=False)
+ self.up_proj = nn.Linear(hidden_size, intermediate_size, bias=False)
+ self.down_proj = nn.Linear(intermediate_size, hidden_size, bias=False)
+
+ def forward(self, x):
+ return self.down_proj(F.silu(self.gate_proj(x)) * self.up_proj(x))
+
+
+class LingBotVideoRouter(nn.Module):
+ def __init__(self, hidden_size, num_experts, top_k, score_func, norm_topk_prob, n_group, topk_group, route_scale):
+ super().__init__()
+ self.num_experts = num_experts
+ self.top_k = top_k
+ self.score_func = score_func
+ self.norm_topk_prob = norm_topk_prob
+ self.n_group = n_group
+ self.topk_group = topk_group
+ self.route_scale = route_scale
+ self.weight = nn.Parameter(torch.empty(num_experts, hidden_size))
+ self.e_score_correction_bias = nn.Parameter(torch.zeros(num_experts), requires_grad=False)
+
+ def _group_limited_topk(self, scores_for_choice):
+ seq_len = scores_for_choice.shape[0]
+ experts_per_group = self.num_experts // self.n_group
+ grouped = scores_for_choice.view(seq_len, self.n_group, experts_per_group)
+ group_scores = grouped.topk(2, dim=-1)[0].sum(dim=-1)
+ group_idx = torch.topk(group_scores, k=self.topk_group, dim=-1, sorted=False)[1]
+ group_mask = torch.zeros_like(group_scores)
+ group_mask.scatter_(1, group_idx, 1)
+ score_mask = group_mask.unsqueeze(-1).expand(seq_len, self.n_group, experts_per_group).reshape(seq_len, -1)
+ masked = scores_for_choice.masked_fill(~score_mask.bool(), float("-inf"))
+ return torch.topk(masked, k=self.top_k, dim=-1, sorted=False)[1]
+
+ def forward(self, tokens: torch.Tensor):
+ with torch.amp.autocast(tokens.device.type, enabled=False):
+ logits = F.linear(tokens.float(), self.weight.to(device=tokens.device, dtype=torch.float32))
+ if self.score_func == "softmax":
+ scores = F.softmax(logits, dim=-1)
+ else:
+ scores = logits.sigmoid()
+ scores_for_choice = scores + self.e_score_correction_bias.to(device=scores.device, dtype=scores.dtype).unsqueeze(0)
+ if self.n_group is not None and self.n_group > 1:
+ top_indices = self._group_limited_topk(scores_for_choice)
+ else:
+ top_indices = torch.topk(scores_for_choice, k=self.top_k, dim=-1, sorted=False)[1]
+ top_scores = scores.gather(1, top_indices)
+ if self.top_k > 1 and self.norm_topk_prob:
+ top_scores = top_scores / (top_scores.sum(dim=-1, keepdim=True) + 1e-20)
+ top_scores = top_scores * self.route_scale
+ return top_indices, top_scores.to(tokens.dtype)
+
+
+class LingBotVideoGroupedExperts(nn.Module):
+ """Weight layout: w1 [E,I,H], w2 [E,H,I], w3 [E,I,H]."""
+
+ def __init__(self, num_experts, hidden_size, intermediate_size):
+ super().__init__()
+ self.num_experts = num_experts
+ self.w1 = nn.Parameter(torch.empty(num_experts, intermediate_size, hidden_size))
+ self.w2 = nn.Parameter(torch.empty(num_experts, hidden_size, intermediate_size))
+ self.w3 = nn.Parameter(torch.empty(num_experts, intermediate_size, hidden_size))
+
+
+def _round_up_to_multiple(value: int, multiple: int) -> int:
+ return ((value + multiple - 1) // multiple) * multiple
+
+
+class LingBotVideoSparseMoeBlock(nn.Module):
+ """MoE FFN with a grouped_mm expert path and a per-expert for-loop fallback."""
+
+ def __init__(self, hidden_size, intermediate_size, num_experts, top_k, moe_intermediate_size,
+ score_func, norm_topk_prob, n_group, topk_group, routed_scaling_factor, n_shared_experts):
+ super().__init__()
+ self.hidden_size = hidden_size
+ self.num_experts = num_experts
+ self.router = LingBotVideoRouter(
+ hidden_size, num_experts, top_k, score_func, norm_topk_prob, n_group, topk_group, routed_scaling_factor,
+ )
+ self.experts = LingBotVideoGroupedExperts(num_experts, hidden_size, moe_intermediate_size)
+ self.shared_experts = None
+ if n_shared_experts is not None and n_shared_experts > 0:
+ self.shared_experts = LingBotVideoMLP(hidden_size, moe_intermediate_size * n_shared_experts)
+
+ @staticmethod
+ def _reorder_tokens(tokens, top_scores, top_indices, num_experts):
+ num_tokens = tokens.shape[0]
+ top_k = top_indices.shape[1]
+ flat_scores = top_scores.reshape(-1)
+ flat_indices = top_indices.reshape(-1)
+ active_positions = torch.where(flat_scores != 0)[0]
+ active_experts = flat_indices[active_positions]
+
+ counts = torch.zeros(num_experts, device=tokens.device, dtype=torch.int64)
+ counts.scatter_add_(0, active_experts, torch.ones_like(active_experts, dtype=torch.int64))
+
+ sort_order = torch.argsort(active_experts, stable=True)
+ sorted_positions = active_positions[sort_order]
+ sorted_scores = flat_scores[sorted_positions]
+ original_token_idx = sorted_positions // top_k
+ permuted_tokens = tokens[original_token_idx]
+ return permuted_tokens, counts, sorted_positions, sorted_scores, num_tokens, top_k
+
+ @staticmethod
+ def _pad_grouped_tokens(tokens, counts, align: int = 8):
+ num_tokens = tokens.shape[0]
+ num_experts = int(counts.shape[0])
+ max_len = _round_up_to_multiple(num_tokens + num_experts * align, align)
+ counts_i64 = counts.to(torch.int64)
+ total_per_expert = torch.clamp_min(counts_i64, align)
+ aligned_counts = ((total_per_expert + align - 1) // align * align).to(torch.int32)
+ write_offsets = torch.cumsum(aligned_counts, dim=0) - aligned_counts
+ start_indices = torch.cumsum(counts_i64, dim=0) - counts_i64
+
+ fill_value = num_tokens
+ permuted_indices = torch.full((max_len,), fill_value, dtype=torch.int64, device=tokens.device)
+ for expert_idx in range(num_experts):
+ length = int(counts_i64[expert_idx].item())
+ if length == 0:
+ continue
+ write_start = int(write_offsets[expert_idx].item())
+ start = int(start_indices[expert_idx].item())
+ permuted_indices[write_start:write_start + length] = torch.arange(
+ start, start + length, device=tokens.device, dtype=torch.int64
+ )
+
+ tokens_with_pad = torch.vstack((tokens, tokens.new_zeros((tokens.shape[-1],))))
+ input_shape = tokens_with_pad.shape
+ return input_shape, tokens_with_pad[permuted_indices], permuted_indices, aligned_counts
+
+ @staticmethod
+ def _unpad_grouped_tokens(output, input_shape, permuted_indices):
+ unpermuted = output.new_empty(input_shape)
+ unpermuted[permuted_indices, :] = output
+ return unpermuted[:-1]
+
+ def _run_grouped_experts(self, tokens, counts):
+ if not hasattr(torch, "_grouped_mm") or tokens.device.type != "cuda":
+ return self._run_experts_for_loop(tokens, counts)
+ input_shape, padded_tokens, permuted_indices, aligned_counts = self._pad_grouped_tokens(tokens, counts)
+ offsets = torch.cumsum(aligned_counts, dim=0, dtype=torch.int32)
+ w1, w2, w3 = (w.to(device=tokens.device, dtype=torch.bfloat16) for w in (self.experts.w1, self.experts.w2, self.experts.w3))
+ h = F.silu(torch._grouped_mm(padded_tokens.bfloat16(), w1.transpose(-2, -1), offs=offsets))
+ h = h * torch._grouped_mm(padded_tokens.bfloat16(), w3.transpose(-2, -1), offs=offsets)
+ out = torch._grouped_mm(h, w2.transpose(-2, -1), offs=offsets).type_as(padded_tokens)
+ return self._unpad_grouped_tokens(out, input_shape, permuted_indices)
+
+ def _run_experts_for_loop(self, tokens, counts):
+ count_list = counts.tolist()
+ splits = torch.split(tokens, count_list, dim=0)
+ outputs = []
+ for expert_idx, expert_tokens in enumerate(splits):
+ if expert_tokens.numel() == 0:
+ continue
+ h = F.silu(expert_tokens @ self.experts.w1[expert_idx].to(device=expert_tokens.device, dtype=expert_tokens.dtype).transpose(-2, -1))
+ h = h * (expert_tokens @ self.experts.w3[expert_idx].to(device=expert_tokens.device, dtype=expert_tokens.dtype).transpose(-2, -1))
+ h = h @ self.experts.w2[expert_idx].to(device=expert_tokens.device, dtype=expert_tokens.dtype).transpose(-2, -1)
+ outputs.append(h)
+ if not outputs:
+ return tokens.new_zeros(tokens.shape)
+ return torch.cat(outputs, dim=0)
+
+ @staticmethod
+ def _restore_tokens(expert_output, sorted_positions, sorted_scores, num_tokens, top_k):
+ dim = expert_output.shape[-1]
+ unsorted = torch.zeros((num_tokens * top_k, dim), dtype=expert_output.dtype, device=expert_output.device)
+ unsorted[sorted_positions] = expert_output
+ unsorted = unsorted.reshape(num_tokens, top_k, dim)
+
+ scores_unsorted = torch.zeros(num_tokens * top_k, dtype=sorted_scores.dtype, device=sorted_scores.device)
+ scores_unsorted[sorted_positions] = sorted_scores
+ scores_unsorted = scores_unsorted.reshape(num_tokens, top_k, 1)
+ return (unsorted.float() * scores_unsorted).sum(dim=1).to(expert_output.dtype)
+
+ def _run_selected_experts(self, tokens, top_scores, top_indices):
+ permuted_tokens, counts, sorted_positions, sorted_scores, num_tokens, top_k = self._reorder_tokens(
+ tokens, top_scores, top_indices, self.router.num_experts
+ )
+ expert_output = self._run_grouped_experts(permuted_tokens, counts)
+ return self._restore_tokens(expert_output, sorted_positions, sorted_scores, num_tokens, top_k)
+
+ def forward(self, hidden_states: torch.Tensor, padding_mask: Optional[torch.Tensor] = None):
+ # hidden_states: (B, S, H); padding_mask: (B*S,) with 1=valid (only needed when B>1)
+ B = hidden_states.shape[0]
+ tokens = hidden_states.view(-1, self.hidden_size)
+ top_indices, top_scores = self.router(tokens)
+ if padding_mask is not None:
+ pm = padding_mask.unsqueeze(-1).to(top_scores.dtype)
+ top_scores = top_scores * pm
+ top_scores = top_scores / (top_scores.sum(dim=-1, keepdim=True) + 1e-9)
+ top_scores = top_scores * self.router.route_scale
+
+ out = self._run_selected_experts(tokens, top_scores, top_indices)
+
+ out = out.view(B, -1, self.hidden_size)
+ if self.shared_experts is not None:
+ out = out + self.shared_experts(hidden_states)
+ return out
+
+
+class LingBotVideoBlock(nn.Module):
+ def __init__(self, hidden_size, num_attention_heads, intermediate_size, norm_eps, qkv_bias, out_bias,
+ num_experts, num_experts_per_tok, moe_intermediate_size, decoder_sparse_step, mlp_only_layers,
+ n_shared_experts, score_func, norm_topk_prob, n_group, topk_group, routed_scaling_factor, layer_idx):
+ super().__init__()
+ self.layer_idx = layer_idx
+ h = hidden_size
+ self.scale_shift_table = nn.Parameter(torch.zeros(1, 6 * h))
+ self.norm1 = LingBotVideoRMSNorm(h, norm_eps)
+ self.attn = LingBotVideoAttention(h, num_attention_heads, norm_eps, qkv_bias, out_bias)
+ self.norm_post_attn = LingBotVideoRMSNorm(h, norm_eps)
+ self.norm2 = LingBotVideoRMSNorm(h, norm_eps)
+ # Sparsity decision: mlp_only_layers + decoder_sparse_step + num_experts.
+ if layer_idx not in mlp_only_layers and (num_experts > 0 and (layer_idx + 1) % decoder_sparse_step == 0):
+ self.ffn = LingBotVideoSparseMoeBlock(
+ h, intermediate_size, num_experts, num_experts_per_tok, moe_intermediate_size,
+ score_func, norm_topk_prob, n_group, topk_group, routed_scaling_factor, n_shared_experts,
+ )
+ else:
+ self.ffn = LingBotVideoMLP(h, intermediate_size)
+ self.norm_post_ffn = LingBotVideoRMSNorm(h, norm_eps)
+
+ def forward(self, x, temb6, rotary_emb, attention_mask=None, moe_padding_mask=None):
+ expected_tokens = x.shape[0] * x.shape[1]
+ if temb6.ndim != 2 or temb6.shape[0] != expected_tokens:
+ raise ValueError(
+ "LingBotVideoBlock expects token-level temb6 with shape (B*S, 6D); "
+ f"got {tuple(temb6.shape)} for hidden states {tuple(x.shape)}."
+ )
+ mod = temb6.view(x.shape[0], x.shape[1], -1) + self.scale_shift_table.to(dtype=temb6.dtype, device=temb6.device).unsqueeze(0)
+ shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = mod.chunk(6, dim=-1)
+ gate_msa, gate_mlp = gate_msa.tanh(), gate_mlp.tanh()
+ scale_msa, scale_mlp = 1.0 + scale_msa, 1.0 + scale_mlp
+
+ bulk_dtype = x.dtype
+ attn_in = (self.norm1(x) * scale_msa + shift_msa).to(bulk_dtype)
+ attn_out = self.attn(attn_in, rotary_emb, attention_mask)
+ x = x + (gate_msa * self.norm_post_attn(attn_out)).to(x.dtype)
+
+ ffn_in = (self.norm2(x) * scale_mlp + shift_mlp).to(bulk_dtype)
+ if isinstance(self.ffn, LingBotVideoSparseMoeBlock):
+ ffn_out = self.ffn(ffn_in, padding_mask=moe_padding_mask)
+ else:
+ ffn_out = self.ffn(ffn_in)
+ x = x + (gate_mlp * self.norm_post_ffn(ffn_out)).to(x.dtype)
+ return x
+
+
+class LingBotVideoDiT(nn.Module):
+ """LingBot-Video MoE DiT.
+
+ Supports both the Dense (`num_experts=0`, FFN = MLP) and MoE
+ (`num_experts>0`, FFN = sparse MoE) variants from a single class.
+ """
+
+ _repeated_blocks = ["LingBotVideoBlock"]
+
+ def __init__(
+ self,
+ patch_size: Tuple[int, int, int] = (1, 2, 2),
+ in_channels: int = 16,
+ out_channels: int = 16,
+ hidden_size: int = 2048,
+ num_attention_heads: int = 16,
+ depth: int = 24,
+ intermediate_size: int = 6144,
+ text_dim: int = 2560,
+ freq_dim: int = 256,
+ norm_eps: float = 1e-6,
+ rope_theta: float = 256.0,
+ axes_dims: Tuple[int, int, int] = (32, 48, 48),
+ axes_lens: Tuple[int, int, int] = (8192, 1024, 1024),
+ qkv_bias: bool = False,
+ out_bias: bool = True,
+ patch_embed_bias: bool = True,
+ timestep_mlp_bias: bool = True,
+ num_experts: int = 0,
+ num_experts_per_tok: int = 8,
+ moe_intermediate_size: int = 512,
+ decoder_sparse_step: int = 1,
+ mlp_only_layers: Tuple[int, ...] = (),
+ n_shared_experts: Optional[int] = None,
+ score_func: str = "sigmoid",
+ norm_topk_prob: bool = True,
+ n_group: Optional[int] = None,
+ topk_group: Optional[int] = None,
+ routed_scaling_factor: float = 1.0,
+ ):
+ super().__init__()
+ head_dim = hidden_size // num_attention_heads
+ assert head_dim == sum(axes_dims), f"head_dim {head_dim} != sum(axes_dims) {sum(axes_dims)}"
+ mlp_only_layers = tuple(mlp_only_layers)
+
+ # Config attributes used by forward / pipeline.
+ self.patch_size = tuple(patch_size)
+ self.in_channels = in_channels
+ self.out_channels = out_channels
+ self.hidden_size = hidden_size
+ self.num_attention_heads = num_attention_heads
+ self.depth = depth
+ self.gradient_checkpointing = False
+
+ self.patch_embedder = nn.Linear(in_channels * math.prod(patch_size), hidden_size, bias=patch_embed_bias)
+ self.time_proj = Timesteps(freq_dim, flip_sin_to_cos=True, downscale_freq_shift=0)
+ self.time_embedder = TimestepEmbedding(freq_dim, hidden_size, sample_proj_bias=timestep_mlp_bias)
+ self.time_modulation = nn.Sequential(nn.SiLU(), nn.Linear(hidden_size, 6 * hidden_size))
+ self.text_embedder = LingBotVideoTextEmbedder(text_dim, hidden_size)
+ self.rope = LingBotVideoRotaryEmbedding(axes_dims, axes_lens, rope_theta)
+ self.blocks = nn.ModuleList([
+ LingBotVideoBlock(
+ hidden_size=hidden_size, num_attention_heads=num_attention_heads, intermediate_size=intermediate_size,
+ norm_eps=norm_eps, qkv_bias=qkv_bias, out_bias=out_bias, num_experts=num_experts,
+ num_experts_per_tok=num_experts_per_tok, moe_intermediate_size=moe_intermediate_size,
+ decoder_sparse_step=decoder_sparse_step, mlp_only_layers=mlp_only_layers,
+ n_shared_experts=n_shared_experts, score_func=score_func, norm_topk_prob=norm_topk_prob,
+ n_group=n_group, topk_group=topk_group, routed_scaling_factor=routed_scaling_factor, layer_idx=i,
+ )
+ for i in range(depth)
+ ])
+ self.norm_out = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=norm_eps)
+ self.norm_out_modulation = nn.Sequential(nn.SiLU(), nn.Linear(hidden_size, 2 * hidden_size))
+ self.proj_out = nn.Linear(hidden_size, math.prod(patch_size) * out_channels)
+
+ def forward(
+ self,
+ hidden_states: torch.Tensor, # (B, C, T, H, W)
+ timestep: torch.Tensor, # (B,) in [0, 1000] (= sigma * 1000)
+ encoder_hidden_states: torch.Tensor, # (B, L, text_dim)
+ encoder_attention_mask: Optional[torch.Tensor] = None, # (B, L) 1=valid
+ use_gradient_checkpointing: bool = False,
+ use_gradient_checkpointing_offload: bool = False,
+ **kwargs,
+ ):
+ B, C, T, H, W = hidden_states.shape
+ pF, pH, pW = self.patch_size
+ gt, gh, gw = T // pF, H // pH, W // pW
+ n_video = gt * gh * gw
+ L = encoder_hidden_states.shape[1]
+ device = hidden_states.device
+ if encoder_attention_mask is not None:
+ text_lens = encoder_attention_mask.sum(dim=-1).long()
+ else:
+ text_lens = torch.full((B,), L, dtype=torch.long, device=device)
+ text_lens_list = [int(v) for v in text_lens.detach().cpu().tolist()]
+ packed_batch = B > 1
+
+ # patchify: token order (f h w), feature order (pf ph pw c)
+ patch_tokens = hidden_states.reshape(B, C, gt, pF, gh, pH, gw, pW)
+ patch_tokens = patch_tokens.permute(0, 2, 4, 6, 3, 5, 7, 1).reshape(B, n_video, pF * pH * pW * C)
+
+ if packed_batch:
+ x = torch.cat([self.patch_embedder(patch_tokens[i:i + 1]) for i in range(B)], dim=1)
+ text_parts = [self.text_embedder(encoder_hidden_states[i:i + 1, :text_lens_list[i], :]) for i in range(B)]
+ text = torch.cat(text_parts, dim=1)
+ joint = _cat_interleave(x, [n_video] * B, text, text_lens_list)
+ else:
+ x = self.patch_embedder(patch_tokens)
+ text = self.text_embedder(encoder_hidden_states)
+ joint = torch.cat([x, text], dim=1) # [video; text]
+ joint_seq_len = joint.shape[1]
+
+ # Per-sample RoPE: video t-axis start = real text length of this sample + 1.
+ rotary_parts = [self.rope(make_joint_position_ids(text_lens_list[i], gt, gh, gw, device)) for i in range(B)]
+ if packed_batch:
+ rotary = torch.cat(rotary_parts, dim=0).unsqueeze(0)
+ else:
+ rotary = torch.stack(rotary_parts, dim=0) # (B, S, head_dim/2) complex64
+
+ attention_mask = None
+ moe_padding_mask = None
+ if packed_batch:
+ # Block-diagonal mask so packed samples only attend within their own block.
+ sample_seq_lens = [n_video + tl for tl in text_lens_list]
+ total = sum(sample_seq_lens)
+ block_mask = torch.zeros((total, total), dtype=torch.bool, device=device)
+ start = 0
+ for slen in sample_seq_lens:
+ block_mask[start:start + slen, start:start + slen] = True
+ start += slen
+ attention_mask = block_mask[None, None, :, :] # (1,1,total,total)
+ else:
+ has_padding = encoder_attention_mask is not None and bool((text_lens < L).any())
+ if has_padding:
+ key_mask = torch.cat(
+ [torch.ones(B, n_video, dtype=torch.bool, device=device), encoder_attention_mask.bool()], dim=1
+ )
+ attention_mask = key_mask[:, None, None, :] # (B,1,1,S) -> SDPA broadcast
+ moe_padding_mask = key_mask.reshape(-1).float() # (B*S,)
+
+ # Timestep -> per-token modulation.
+ timestep_proj = self.time_proj(timestep.float())
+ t_emb = self.time_embedder(timestep_proj.to(joint.dtype)) # (B, D)
+ if packed_batch:
+ temb_input = torch.cat(
+ [t_emb[i:i + 1].unsqueeze(1).expand(1, sample_seq_lens[i], -1) for i in range(B)], dim=1
+ ) # (1, total, D)
+ else:
+ temb_input = t_emb.unsqueeze(1).expand(B, joint_seq_len, -1) # (B, S, D)
+ b_eff, s_eff = temb_input.shape[0], temb_input.shape[1]
+ temb6 = self.time_modulation(temb_input.reshape(b_eff * s_eff, -1)) # (B*S, 6D)
+
+ for block in self.blocks:
+ joint = gradient_checkpoint_forward(
+ block,
+ use_gradient_checkpointing=use_gradient_checkpointing,
+ use_gradient_checkpointing_offload=use_gradient_checkpointing_offload,
+ x=joint, temb6=temb6, rotary_emb=rotary,
+ attention_mask=attention_mask, moe_padding_mask=moe_padding_mask,
+ )
+
+ final_mod = self.norm_out_modulation(temb_input.reshape(joint.shape[0] * joint.shape[1], -1))
+ shift, scale = final_mod.reshape(joint.shape[0], joint.shape[1], -1).chunk(2, dim=-1)
+ final_hidden = self.norm_out(joint) * (1.0 + scale) + shift
+ projected = self.proj_out(final_hidden.to(joint.dtype))
+
+ if packed_batch:
+ split_lengths = []
+ for tl in text_lens_list:
+ split_lengths.extend([n_video, tl])
+ parts = torch.split(projected, split_lengths, dim=1)
+ x = torch.cat(parts[::2], dim=1).reshape(B, n_video, -1)
+ else:
+ x = projected[:, :n_video]
+
+ # unpatchify
+ Cout = self.out_channels
+ x = x.reshape(B, gt, gh, gw, pF, pH, pW, Cout)
+ x = x.permute(0, 7, 1, 4, 2, 5, 3, 6).reshape(B, Cout, T, H, W)
+ return x
diff --git a/diffsynth/models/longcat_video_dit.py b/diffsynth/models/longcat_video_dit.py
new file mode 100644
index 0000000000000000000000000000000000000000..dbe1c21b9304c484c72cb17e66ff97e7c85bc16b
--- /dev/null
+++ b/diffsynth/models/longcat_video_dit.py
@@ -0,0 +1,902 @@
+from typing import List, Optional, Tuple
+
+import math
+import torch
+import torch.nn as nn
+import torch.amp as amp
+
+import numpy as np
+import torch.nn.functional as F
+from einops import rearrange, repeat
+from .wan_video_dit import flash_attention
+from ..core.device.npu_compatible_device import get_device_type
+from ..core.gradient import gradient_checkpoint_forward
+
+
+class RMSNorm_FP32(torch.nn.Module):
+ def __init__(self, dim: int, eps: float):
+ super().__init__()
+ self.eps = eps
+ self.weight = nn.Parameter(torch.ones(dim))
+
+ def _norm(self, x):
+ return x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps)
+
+ def forward(self, x):
+ output = self._norm(x.float()).type_as(x)
+ return output * self.weight
+
+
+def broadcat(tensors, dim=-1):
+ num_tensors = len(tensors)
+ shape_lens = set(list(map(lambda t: len(t.shape), tensors)))
+ assert len(shape_lens) == 1, "tensors must all have the same number of dimensions"
+ shape_len = list(shape_lens)[0]
+ dim = (dim + shape_len) if dim < 0 else dim
+ dims = list(zip(*map(lambda t: list(t.shape), tensors)))
+ expandable_dims = [(i, val) for i, val in enumerate(dims) if i != dim]
+ assert all(
+ [*map(lambda t: len(set(t[1])) <= 2, expandable_dims)]
+ ), "invalid dimensions for broadcastable concatentation"
+ max_dims = list(map(lambda t: (t[0], max(t[1])), expandable_dims))
+ expanded_dims = list(map(lambda t: (t[0], (t[1],) * num_tensors), max_dims))
+ expanded_dims.insert(dim, (dim, dims[dim]))
+ expandable_shapes = list(zip(*map(lambda t: t[1], expanded_dims)))
+ tensors = list(map(lambda t: t[0].expand(*t[1]), zip(tensors, expandable_shapes)))
+ return torch.cat(tensors, dim=dim)
+
+
+def rotate_half(x):
+ x = rearrange(x, "... (d r) -> ... d r", r=2)
+ x1, x2 = x.unbind(dim=-1)
+ x = torch.stack((-x2, x1), dim=-1)
+ return rearrange(x, "... d r -> ... (d r)")
+
+
+class RotaryPositionalEmbedding(nn.Module):
+
+ def __init__(self,
+ head_dim,
+ cp_split_hw=None
+ ):
+ """Rotary positional embedding for 3D
+ Reference : https://blog.eleuther.ai/rotary-embeddings/
+ Paper: https://arxiv.org/pdf/2104.09864.pdf
+ Args:
+ dim: Dimension of embedding
+ base: Base value for exponential
+ """
+ super().__init__()
+ self.head_dim = head_dim
+ assert self.head_dim % 8 == 0, 'Dim must be a multiply of 8 for 3D RoPE.'
+ self.cp_split_hw = cp_split_hw
+ # We take the assumption that the longest side of grid will not larger than 512, i.e, 512 * 8 = 4098 input pixels
+ self.base = 10000
+ self.freqs_dict = {}
+
+ def register_grid_size(self, grid_size):
+ if grid_size not in self.freqs_dict:
+ self.freqs_dict.update({
+ grid_size: self.precompute_freqs_cis_3d(grid_size)
+ })
+
+ def precompute_freqs_cis_3d(self, grid_size):
+ num_frames, height, width = grid_size
+ dim_t = self.head_dim - 4 * (self.head_dim // 6)
+ dim_h = 2 * (self.head_dim // 6)
+ dim_w = 2 * (self.head_dim // 6)
+ freqs_t = 1.0 / (self.base ** (torch.arange(0, dim_t, 2)[: (dim_t // 2)].float() / dim_t))
+ freqs_h = 1.0 / (self.base ** (torch.arange(0, dim_h, 2)[: (dim_h // 2)].float() / dim_h))
+ freqs_w = 1.0 / (self.base ** (torch.arange(0, dim_w, 2)[: (dim_w // 2)].float() / dim_w))
+ grid_t = np.linspace(0, num_frames, num_frames, endpoint=False, dtype=np.float32)
+ grid_h = np.linspace(0, height, height, endpoint=False, dtype=np.float32)
+ grid_w = np.linspace(0, width, width, endpoint=False, dtype=np.float32)
+ grid_t = torch.from_numpy(grid_t).float()
+ grid_h = torch.from_numpy(grid_h).float()
+ grid_w = torch.from_numpy(grid_w).float()
+ freqs_t = torch.einsum("..., f -> ... f", grid_t, freqs_t)
+ freqs_h = torch.einsum("..., f -> ... f", grid_h, freqs_h)
+ freqs_w = torch.einsum("..., f -> ... f", grid_w, freqs_w)
+ freqs_t = repeat(freqs_t, "... n -> ... (n r)", r=2)
+ freqs_h = repeat(freqs_h, "... n -> ... (n r)", r=2)
+ freqs_w = repeat(freqs_w, "... n -> ... (n r)", r=2)
+ freqs = broadcat((freqs_t[:, None, None, :], freqs_h[None, :, None, :], freqs_w[None, None, :, :]), dim=-1)
+ # (T H W D)
+ freqs = rearrange(freqs, "T H W D -> (T H W) D")
+ # if self.cp_split_hw[0] * self.cp_split_hw[1] > 1:
+ # with torch.no_grad():
+ # freqs = rearrange(freqs, "(T H W) D -> T H W D", T=num_frames, H=height, W=width)
+ # freqs = context_parallel_util.split_cp_2d(freqs, seq_dim_hw=(1, 2), split_hw=self.cp_split_hw)
+ # freqs = rearrange(freqs, "T H W D -> (T H W) D")
+
+ return freqs
+
+ def forward(self, q, k, grid_size):
+ """3D RoPE.
+
+ Args:
+ query: [B, head, seq, head_dim]
+ key: [B, head, seq, head_dim]
+ Returns:
+ query and key with the same shape as input.
+ """
+
+ if grid_size not in self.freqs_dict:
+ self.register_grid_size(grid_size)
+
+ freqs_cis = self.freqs_dict[grid_size].to(q.device)
+ q_, k_ = q.float(), k.float()
+ freqs_cis = freqs_cis.float().to(q.device)
+ cos, sin = freqs_cis.cos(), freqs_cis.sin()
+ cos, sin = rearrange(cos, 'n d -> 1 1 n d'), rearrange(sin, 'n d -> 1 1 n d')
+ q_ = (q_ * cos) + (rotate_half(q_) * sin)
+ k_ = (k_ * cos) + (rotate_half(k_) * sin)
+
+ return q_.type_as(q), k_.type_as(k)
+
+
+class Attention(nn.Module):
+ def __init__(
+ self,
+ dim: int,
+ num_heads: int,
+ enable_flashattn3: bool = False,
+ enable_flashattn2: bool = False,
+ enable_xformers: bool = False,
+ enable_bsa: bool = False,
+ bsa_params: dict = None,
+ cp_split_hw: Optional[List[int]] = None
+ ) -> None:
+ super().__init__()
+ assert dim % num_heads == 0, "dim should be divisible by num_heads"
+ self.dim = dim
+ self.num_heads = num_heads
+ self.head_dim = dim // num_heads
+ self.scale = self.head_dim**-0.5
+ self.enable_flashattn3 = enable_flashattn3
+ self.enable_flashattn2 = enable_flashattn2
+ self.enable_xformers = enable_xformers
+ self.enable_bsa = enable_bsa
+ self.bsa_params = bsa_params
+ self.cp_split_hw = cp_split_hw
+
+ self.qkv = nn.Linear(dim, dim * 3, bias=True)
+ self.q_norm = RMSNorm_FP32(self.head_dim, eps=1e-6)
+ self.k_norm = RMSNorm_FP32(self.head_dim, eps=1e-6)
+ self.proj = nn.Linear(dim, dim)
+
+ self.rope_3d = RotaryPositionalEmbedding(
+ self.head_dim,
+ cp_split_hw=cp_split_hw
+ )
+
+ def _process_attn(self, q, k, v, shape):
+ q = rearrange(q, "B H S D -> B S (H D)")
+ k = rearrange(k, "B H S D -> B S (H D)")
+ v = rearrange(v, "B H S D -> B S (H D)")
+ x = flash_attention(q, k, v, num_heads=self.num_heads)
+ x = rearrange(x, "B S (H D) -> B H S D", H=self.num_heads)
+ return x
+
+ def forward(self, x: torch.Tensor, shape=None, num_cond_latents=None, return_kv=False) -> torch.Tensor:
+ """
+ """
+ B, N, C = x.shape
+ qkv = self.qkv(x)
+
+ qkv_shape = (B, N, 3, self.num_heads, self.head_dim)
+ qkv = qkv.view(qkv_shape).permute((2, 0, 3, 1, 4)) # [3, B, H, N, D]
+ q, k, v = qkv.unbind(0)
+ q, k = self.q_norm(q), self.k_norm(k)
+
+ if return_kv:
+ k_cache, v_cache = k.clone(), v.clone()
+
+ q, k = self.rope_3d(q, k, shape)
+
+ # cond mode
+ if num_cond_latents is not None and num_cond_latents > 0:
+ num_cond_latents_thw = num_cond_latents * (N // shape[0])
+ # process the condition tokens
+ q_cond = q[:, :, :num_cond_latents_thw].contiguous()
+ k_cond = k[:, :, :num_cond_latents_thw].contiguous()
+ v_cond = v[:, :, :num_cond_latents_thw].contiguous()
+ x_cond = self._process_attn(q_cond, k_cond, v_cond, shape)
+ # process the noise tokens
+ q_noise = q[:, :, num_cond_latents_thw:].contiguous()
+ x_noise = self._process_attn(q_noise, k, v, shape)
+ # merge x_cond and x_noise
+ x = torch.cat([x_cond, x_noise], dim=2).contiguous()
+ else:
+ x = self._process_attn(q, k, v, shape)
+
+ x_output_shape = (B, N, C)
+ x = x.transpose(1, 2) # [B, H, N, D] --> [B, N, H, D]
+ x = x.reshape(x_output_shape) # [B, N, H, D] --> [B, N, C]
+ x = self.proj(x)
+
+ if return_kv:
+ return x, (k_cache, v_cache)
+ else:
+ return x
+
+ def forward_with_kv_cache(self, x: torch.Tensor, shape=None, num_cond_latents=None, kv_cache=None) -> torch.Tensor:
+ """
+ """
+ B, N, C = x.shape
+ qkv = self.qkv(x)
+
+ qkv_shape = (B, N, 3, self.num_heads, self.head_dim)
+ qkv = qkv.view(qkv_shape).permute((2, 0, 3, 1, 4)) # [3, B, H, N, D]
+ q, k, v = qkv.unbind(0)
+ q, k = self.q_norm(q), self.k_norm(k)
+
+ T, H, W = shape
+ k_cache, v_cache = kv_cache
+ assert k_cache.shape[0] == v_cache.shape[0] and k_cache.shape[0] in [1, B]
+ if k_cache.shape[0] == 1:
+ k_cache = k_cache.repeat(B, 1, 1, 1)
+ v_cache = v_cache.repeat(B, 1, 1, 1)
+
+ if num_cond_latents is not None and num_cond_latents > 0:
+ k_full = torch.cat([k_cache, k], dim=2).contiguous()
+ v_full = torch.cat([v_cache, v], dim=2).contiguous()
+ q_padding = torch.cat([torch.empty_like(k_cache), q], dim=2).contiguous()
+ q_padding, k_full = self.rope_3d(q_padding, k_full, (T + num_cond_latents, H, W))
+ q = q_padding[:, :, -N:].contiguous()
+
+ x = self._process_attn(q, k_full, v_full, shape)
+
+ x_output_shape = (B, N, C)
+ x = x.transpose(1, 2) # [B, H, N, D] --> [B, N, H, D]
+ x = x.reshape(x_output_shape) # [B, N, H, D] --> [B, N, C]
+ x = self.proj(x)
+
+ return x
+
+
+class MultiHeadCrossAttention(nn.Module):
+ def __init__(
+ self,
+ dim,
+ num_heads,
+ enable_flashattn3=False,
+ enable_flashattn2=False,
+ enable_xformers=False,
+ ):
+ super(MultiHeadCrossAttention, self).__init__()
+ assert dim % num_heads == 0, "d_model must be divisible by num_heads"
+
+ self.dim = dim
+ self.num_heads = num_heads
+ self.head_dim = dim // num_heads
+
+ self.q_linear = nn.Linear(dim, dim)
+ self.kv_linear = nn.Linear(dim, dim * 2)
+ self.proj = nn.Linear(dim, dim)
+
+ self.q_norm = RMSNorm_FP32(self.head_dim, eps=1e-6)
+ self.k_norm = RMSNorm_FP32(self.head_dim, eps=1e-6)
+
+ self.enable_flashattn3 = enable_flashattn3
+ self.enable_flashattn2 = enable_flashattn2
+ self.enable_xformers = enable_xformers
+
+ def _process_cross_attn(self, x, cond, kv_seqlen):
+ B, N, C = x.shape
+ assert C == self.dim and cond.shape[2] == self.dim
+
+ q = self.q_linear(x).view(1, -1, self.num_heads, self.head_dim)
+ kv = self.kv_linear(cond).view(1, -1, 2, self.num_heads, self.head_dim)
+ k, v = kv.unbind(2)
+
+ q, k = self.q_norm(q), self.k_norm(k)
+
+ q = rearrange(q, "B S H D -> B S (H D)")
+ k = rearrange(k, "B S H D -> B S (H D)")
+ v = rearrange(v, "B S H D -> B S (H D)")
+ x = flash_attention(q, k, v, num_heads=self.num_heads)
+
+ x = x.view(B, -1, C)
+ x = self.proj(x)
+ return x
+
+ def forward(self, x, cond, kv_seqlen, num_cond_latents=None, shape=None):
+ """
+ x: [B, N, C]
+ cond: [B, M, C]
+ """
+ if num_cond_latents is None or num_cond_latents == 0:
+ return self._process_cross_attn(x, cond, kv_seqlen)
+ else:
+ B, N, C = x.shape
+ if num_cond_latents is not None and num_cond_latents > 0:
+ assert shape is not None, "SHOULD pass in the shape"
+ num_cond_latents_thw = num_cond_latents * (N // shape[0])
+ x_noise = x[:, num_cond_latents_thw:] # [B, N_noise, C]
+ output_noise = self._process_cross_attn(x_noise, cond, kv_seqlen) # [B, N_noise, C]
+ output = torch.cat([
+ torch.zeros((B, num_cond_latents_thw, C), dtype=output_noise.dtype, device=output_noise.device),
+ output_noise
+ ], dim=1).contiguous()
+ else:
+ raise NotImplementedError
+
+ return output
+
+
+class LayerNorm_FP32(nn.LayerNorm):
+ def __init__(self, dim, eps, elementwise_affine):
+ super().__init__(dim, eps=eps, elementwise_affine=elementwise_affine)
+
+ def forward(self, inputs: torch.Tensor) -> torch.Tensor:
+ origin_dtype = inputs.dtype
+ out = F.layer_norm(
+ inputs.float(),
+ self.normalized_shape,
+ None if self.weight is None else self.weight.float(),
+ None if self.bias is None else self.bias.float() ,
+ self.eps
+ ).to(origin_dtype)
+ return out
+
+
+def modulate_fp32(norm_func, x, shift, scale):
+ # Suppose x is (B, N, D), shift is (B, -1, D), scale is (B, -1, D)
+ # ensure the modulation params be fp32
+ assert shift.dtype == torch.float32, scale.dtype == torch.float32
+ dtype = x.dtype
+ x = norm_func(x.to(torch.float32))
+ x = x * (scale + 1) + shift
+ x = x.to(dtype)
+ return x
+
+
+class FinalLayer_FP32(nn.Module):
+ """
+ The final layer of DiT.
+ """
+
+ def __init__(self, hidden_size, num_patch, out_channels, adaln_tembed_dim):
+ super().__init__()
+ self.hidden_size = hidden_size
+ self.num_patch = num_patch
+ self.out_channels = out_channels
+ self.adaln_tembed_dim = adaln_tembed_dim
+
+ self.norm_final = LayerNorm_FP32(hidden_size, elementwise_affine=False, eps=1e-6)
+ self.linear = nn.Linear(hidden_size, num_patch * out_channels, bias=True)
+ self.adaLN_modulation = nn.Sequential(nn.SiLU(), nn.Linear(adaln_tembed_dim, 2 * hidden_size, bias=True))
+
+ def forward(self, x, t, latent_shape):
+ # timestep shape: [B, T, C]
+ assert t.dtype == torch.float32
+ B, N, C = x.shape
+ T, _, _ = latent_shape
+
+ with amp.autocast(get_device_type(), dtype=torch.float32):
+ shift, scale = self.adaLN_modulation(t).unsqueeze(2).chunk(2, dim=-1) # [B, T, 1, C]
+ x = modulate_fp32(self.norm_final, x.view(B, T, -1, C), shift, scale).view(B, N, C)
+ x = self.linear(x)
+ return x
+
+
+class FeedForwardSwiGLU(nn.Module):
+ def __init__(
+ self,
+ dim: int,
+ hidden_dim: int,
+ multiple_of: int = 256,
+ ffn_dim_multiplier: Optional[float] = None,
+ ):
+ super().__init__()
+ hidden_dim = int(2 * hidden_dim / 3)
+ # custom dim factor multiplier
+ if ffn_dim_multiplier is not None:
+ hidden_dim = int(ffn_dim_multiplier * hidden_dim)
+ hidden_dim = multiple_of * ((hidden_dim + multiple_of - 1) // multiple_of)
+
+ self.dim = dim
+ self.hidden_dim = hidden_dim
+ self.w1 = nn.Linear(dim, hidden_dim, bias=False)
+ self.w2 = nn.Linear(hidden_dim, dim, bias=False)
+ self.w3 = nn.Linear(dim, hidden_dim, bias=False)
+
+ def forward(self, x):
+ return self.w2(F.silu(self.w1(x)) * self.w3(x))
+
+
+class TimestepEmbedder(nn.Module):
+ """
+ Embeds scalar timesteps into vector representations.
+ """
+
+ def __init__(self, t_embed_dim, frequency_embedding_size=256):
+ super().__init__()
+ self.t_embed_dim = t_embed_dim
+ self.frequency_embedding_size = frequency_embedding_size
+ self.mlp = nn.Sequential(
+ nn.Linear(frequency_embedding_size, t_embed_dim, bias=True),
+ nn.SiLU(),
+ nn.Linear(t_embed_dim, t_embed_dim, bias=True),
+ )
+
+ @staticmethod
+ def timestep_embedding(t, dim, max_period=10000):
+ """
+ Create sinusoidal timestep embeddings.
+ :param t: a 1-D Tensor of N indices, one per batch element.
+ These may be fractional.
+ :param dim: the dimension of the output.
+ :param max_period: controls the minimum frequency of the embeddings.
+ :return: an (N, D) Tensor of positional embeddings.
+ """
+ half = dim // 2
+ freqs = torch.exp(-math.log(max_period) * torch.arange(start=0, end=half, dtype=torch.float32) / half)
+ freqs = freqs.to(device=t.device)
+ args = t[:, None].float() * freqs[None]
+ embedding = torch.cat([torch.cos(args), torch.sin(args)], dim=-1)
+ if dim % 2:
+ embedding = torch.cat([embedding, torch.zeros_like(embedding[:, :1])], dim=-1)
+ return embedding
+
+ def forward(self, t, dtype):
+ t_freq = self.timestep_embedding(t, self.frequency_embedding_size)
+ if t_freq.dtype != dtype:
+ t_freq = t_freq.to(dtype)
+ t_emb = self.mlp(t_freq)
+ return t_emb
+
+
+class CaptionEmbedder(nn.Module):
+ """
+ Embeds class labels into vector representations.
+ """
+
+ def __init__(self, in_channels, hidden_size):
+ super().__init__()
+ self.in_channels = in_channels
+ self.hidden_size = hidden_size
+ self.y_proj = nn.Sequential(
+ nn.Linear(in_channels, hidden_size, bias=True),
+ nn.GELU(approximate="tanh"),
+ nn.Linear(hidden_size, hidden_size, bias=True),
+ )
+
+ def forward(self, caption):
+ B, _, N, C = caption.shape
+ caption = self.y_proj(caption)
+ return caption
+
+
+class PatchEmbed3D(nn.Module):
+ """Video to Patch Embedding.
+
+ Args:
+ patch_size (int): Patch token size. Default: (2,4,4).
+ in_chans (int): Number of input video channels. Default: 3.
+ embed_dim (int): Number of linear projection output channels. Default: 96.
+ norm_layer (nn.Module, optional): Normalization layer. Default: None
+ """
+
+ def __init__(
+ self,
+ patch_size=(2, 4, 4),
+ in_chans=3,
+ embed_dim=96,
+ norm_layer=None,
+ flatten=True,
+ ):
+ super().__init__()
+ self.patch_size = patch_size
+ self.flatten = flatten
+
+ self.in_chans = in_chans
+ self.embed_dim = embed_dim
+
+ self.proj = nn.Conv3d(in_chans, embed_dim, kernel_size=patch_size, stride=patch_size)
+ if norm_layer is not None:
+ self.norm = norm_layer(embed_dim)
+ else:
+ self.norm = None
+
+ def forward(self, x):
+ """Forward function."""
+ # padding
+ _, _, D, H, W = x.size()
+ if W % self.patch_size[2] != 0:
+ x = F.pad(x, (0, self.patch_size[2] - W % self.patch_size[2]))
+ if H % self.patch_size[1] != 0:
+ x = F.pad(x, (0, 0, 0, self.patch_size[1] - H % self.patch_size[1]))
+ if D % self.patch_size[0] != 0:
+ x = F.pad(x, (0, 0, 0, 0, 0, self.patch_size[0] - D % self.patch_size[0]))
+
+ B, C, T, H, W = x.shape
+ x = self.proj(x) # (B C T H W)
+ if self.norm is not None:
+ D, Wh, Ww = x.size(2), x.size(3), x.size(4)
+ x = x.flatten(2).transpose(1, 2)
+ x = self.norm(x)
+ x = x.transpose(1, 2).view(-1, self.embed_dim, D, Wh, Ww)
+ if self.flatten:
+ x = x.flatten(2).transpose(1, 2) # BCTHW -> BNC
+ return x
+
+
+class LongCatSingleStreamBlock(nn.Module):
+ def __init__(
+ self,
+ hidden_size: int,
+ num_heads: int,
+ mlp_ratio: int,
+ adaln_tembed_dim: int,
+ enable_flashattn3: bool = False,
+ enable_flashattn2: bool = False,
+ enable_xformers: bool = False,
+ enable_bsa: bool = False,
+ bsa_params=None,
+ cp_split_hw=None
+ ):
+ super().__init__()
+
+ self.hidden_size = hidden_size
+
+ # scale and gate modulation
+ self.adaLN_modulation = nn.Sequential(
+ nn.SiLU(),
+ nn.Linear(adaln_tembed_dim, 6 * hidden_size, bias=True)
+ )
+
+ self.mod_norm_attn = LayerNorm_FP32(hidden_size, eps=1e-6, elementwise_affine=False)
+ self.mod_norm_ffn = LayerNorm_FP32(hidden_size, eps=1e-6, elementwise_affine=False)
+ self.pre_crs_attn_norm = LayerNorm_FP32(hidden_size, eps=1e-6, elementwise_affine=True)
+
+ self.attn = Attention(
+ dim=hidden_size,
+ num_heads=num_heads,
+ enable_flashattn3=enable_flashattn3,
+ enable_flashattn2=enable_flashattn2,
+ enable_xformers=enable_xformers,
+ enable_bsa=enable_bsa,
+ bsa_params=bsa_params,
+ cp_split_hw=cp_split_hw
+ )
+ self.cross_attn = MultiHeadCrossAttention(
+ dim=hidden_size,
+ num_heads=num_heads,
+ enable_flashattn3=enable_flashattn3,
+ enable_flashattn2=enable_flashattn2,
+ enable_xformers=enable_xformers,
+ )
+ self.ffn = FeedForwardSwiGLU(dim=hidden_size, hidden_dim=int(hidden_size * mlp_ratio))
+
+ def forward(self, x, y, t, y_seqlen, latent_shape, num_cond_latents=None, return_kv=False, kv_cache=None, skip_crs_attn=False):
+ """
+ x: [B, N, C]
+ y: [1, N_valid_tokens, C]
+ t: [B, T, C_t]
+ y_seqlen: [B]; type of a list
+ latent_shape: latent shape of a single item
+ """
+ x_dtype = x.dtype
+
+ B, N, C = x.shape
+ T, _, _ = latent_shape # S != T*H*W in case of CP split on H*W.
+
+ # compute modulation params in fp32
+ with amp.autocast(device_type=get_device_type(), dtype=torch.float32):
+ shift_msa, scale_msa, gate_msa, \
+ shift_mlp, scale_mlp, gate_mlp = \
+ self.adaLN_modulation(t).unsqueeze(2).chunk(6, dim=-1) # [B, T, 1, C]
+
+ # self attn with modulation
+ x_m = modulate_fp32(self.mod_norm_attn, x.view(B, T, -1, C), shift_msa, scale_msa).view(B, N, C)
+
+ if kv_cache is not None:
+ kv_cache = (kv_cache[0].to(x.device), kv_cache[1].to(x.device))
+ attn_outputs = self.attn.forward_with_kv_cache(x_m, shape=latent_shape, num_cond_latents=num_cond_latents, kv_cache=kv_cache)
+ else:
+ attn_outputs = self.attn(x_m, shape=latent_shape, num_cond_latents=num_cond_latents, return_kv=return_kv)
+
+ if return_kv:
+ x_s, kv_cache = attn_outputs
+ else:
+ x_s = attn_outputs
+
+ with amp.autocast(device_type=get_device_type(), dtype=torch.float32):
+ x = x + (gate_msa * x_s.view(B, -1, N//T, C)).view(B, -1, C) # [B, N, C]
+ x = x.to(x_dtype)
+
+ # cross attn
+ if not skip_crs_attn:
+ if kv_cache is not None:
+ num_cond_latents = None
+ x = x + self.cross_attn(self.pre_crs_attn_norm(x), y, y_seqlen, num_cond_latents=num_cond_latents, shape=latent_shape)
+
+ # ffn with modulation
+ x_m = modulate_fp32(self.mod_norm_ffn, x.view(B, -1, N//T, C), shift_mlp, scale_mlp).view(B, -1, C)
+ x_s = self.ffn(x_m)
+ with amp.autocast(device_type=get_device_type(), dtype=torch.float32):
+ x = x + (gate_mlp * x_s.view(B, -1, N//T, C)).view(B, -1, C) # [B, N, C]
+ x = x.to(x_dtype)
+
+ if return_kv:
+ return x, kv_cache
+ else:
+ return x
+
+
+class LongCatVideoTransformer3DModel(torch.nn.Module):
+ def __init__(
+ self,
+ in_channels: int = 16,
+ out_channels: int = 16,
+ hidden_size: int = 4096,
+ depth: int = 48,
+ num_heads: int = 32,
+ caption_channels: int = 4096,
+ mlp_ratio: int = 4,
+ adaln_tembed_dim: int = 512,
+ frequency_embedding_size: int = 256,
+ # default params
+ patch_size: Tuple[int] = (1, 2, 2),
+ # attention config
+ enable_flashattn3: bool = False,
+ enable_flashattn2: bool = True,
+ enable_xformers: bool = False,
+ enable_bsa: bool = False,
+ bsa_params: dict = {'sparsity': 0.9375, 'chunk_3d_shape_q': [4, 4, 4], 'chunk_3d_shape_k': [4, 4, 4]},
+ cp_split_hw: Optional[List[int]] = [1, 1],
+ text_tokens_zero_pad: bool = True,
+ ) -> None:
+ super().__init__()
+
+ self.patch_size = patch_size
+ self.in_channels = in_channels
+ self.out_channels = out_channels
+ self.cp_split_hw = cp_split_hw
+
+ self.x_embedder = PatchEmbed3D(patch_size, in_channels, hidden_size)
+ self.t_embedder = TimestepEmbedder(t_embed_dim=adaln_tembed_dim, frequency_embedding_size=frequency_embedding_size)
+ self.y_embedder = CaptionEmbedder(
+ in_channels=caption_channels,
+ hidden_size=hidden_size,
+ )
+
+ self.blocks = nn.ModuleList(
+ [
+ LongCatSingleStreamBlock(
+ hidden_size=hidden_size,
+ num_heads=num_heads,
+ mlp_ratio=mlp_ratio,
+ adaln_tembed_dim=adaln_tembed_dim,
+ enable_flashattn3=enable_flashattn3,
+ enable_flashattn2=enable_flashattn2,
+ enable_xformers=enable_xformers,
+ enable_bsa=enable_bsa,
+ bsa_params=bsa_params,
+ cp_split_hw=cp_split_hw
+ )
+ for i in range(depth)
+ ]
+ )
+
+ self.final_layer = FinalLayer_FP32(
+ hidden_size,
+ np.prod(self.patch_size),
+ out_channels,
+ adaln_tembed_dim,
+ )
+
+ self.gradient_checkpointing = False
+ self.text_tokens_zero_pad = text_tokens_zero_pad
+
+ self.lora_dict = {}
+ self.active_loras = []
+
+ def enable_loras(self, lora_key_list=[]):
+ self.disable_all_loras()
+
+ module_loras = {} # {module_name: [lora1, lora2, ...]}
+ model_device = next(self.parameters()).device
+ model_dtype = next(self.parameters()).dtype
+
+ for lora_key in lora_key_list:
+ if lora_key in self.lora_dict:
+ for lora in self.lora_dict[lora_key].loras:
+ lora.to(model_device, dtype=model_dtype, non_blocking=True)
+ module_name = lora.lora_name.replace("lora___lorahyphen___", "").replace("___lorahyphen___", ".")
+ if module_name not in module_loras:
+ module_loras[module_name] = []
+ module_loras[module_name].append(lora)
+ self.active_loras.append(lora_key)
+
+ for module_name, loras in module_loras.items():
+ module = self._get_module_by_name(module_name)
+ if not hasattr(module, 'org_forward'):
+ module.org_forward = module.forward
+ module.forward = self._create_multi_lora_forward(module, loras)
+
+ def _create_multi_lora_forward(self, module, loras):
+ def multi_lora_forward(x, *args, **kwargs):
+ weight_dtype = x.dtype
+ org_output = module.org_forward(x, *args, **kwargs)
+
+ total_lora_output = 0
+ for lora in loras:
+ if lora.use_lora:
+ lx = lora.lora_down(x.to(lora.lora_down.weight.dtype))
+ lx = lora.lora_up(lx)
+ lora_output = lx.to(weight_dtype) * lora.multiplier * lora.alpha_scale
+ total_lora_output += lora_output
+
+ return org_output + total_lora_output
+
+ return multi_lora_forward
+
+ def _get_module_by_name(self, module_name):
+ try:
+ module = self
+ for part in module_name.split('.'):
+ module = getattr(module, part)
+ return module
+ except AttributeError as e:
+ raise ValueError(f"Cannot find module: {module_name}, error: {e}")
+
+ def disable_all_loras(self):
+ for name, module in self.named_modules():
+ if hasattr(module, 'org_forward'):
+ module.forward = module.org_forward
+ delattr(module, 'org_forward')
+
+ for lora_key, lora_network in self.lora_dict.items():
+ for lora in lora_network.loras:
+ lora.to("cpu")
+
+ self.active_loras.clear()
+
+ def enable_bsa(self,):
+ for block in self.blocks:
+ block.attn.enable_bsa = True
+
+ def disable_bsa(self,):
+ for block in self.blocks:
+ block.attn.enable_bsa = False
+
+ def forward(
+ self,
+ hidden_states,
+ timestep,
+ encoder_hidden_states,
+ encoder_attention_mask=None,
+ num_cond_latents=0,
+ return_kv=False,
+ kv_cache_dict={},
+ skip_crs_attn=False,
+ offload_kv_cache=False,
+ use_gradient_checkpointing=False,
+ use_gradient_checkpointing_offload=False,
+ ):
+
+ B, _, T, H, W = hidden_states.shape
+
+ N_t = T // self.patch_size[0]
+ N_h = H // self.patch_size[1]
+ N_w = W // self.patch_size[2]
+
+ assert self.patch_size[0]==1, "Currently, 3D x_embedder should not compress the temporal dimension."
+
+ # expand the shape of timestep from [B] to [B, T]
+ if len(timestep.shape) == 1:
+ timestep = timestep.unsqueeze(1).expand(-1, N_t).clone() # [B, T]
+ timestep[:, :num_cond_latents] = 0
+
+ dtype = hidden_states.dtype
+ hidden_states = hidden_states.to(dtype)
+ timestep = timestep.to(dtype)
+ encoder_hidden_states = encoder_hidden_states.to(dtype)
+
+ hidden_states = self.x_embedder(hidden_states) # [B, N, C]
+
+ with amp.autocast(device_type=get_device_type(), dtype=torch.float32):
+ t = self.t_embedder(timestep.float().flatten(), dtype=torch.float32).reshape(B, N_t, -1) # [B, T, C_t]
+
+ encoder_hidden_states = self.y_embedder(encoder_hidden_states) # [B, 1, N_token, C]
+
+ if self.text_tokens_zero_pad and encoder_attention_mask is not None:
+ encoder_hidden_states = encoder_hidden_states * encoder_attention_mask[:, None, :, None]
+ encoder_attention_mask = (encoder_attention_mask * 0 + 1).to(encoder_attention_mask.dtype)
+
+ if encoder_attention_mask is not None:
+ encoder_attention_mask = encoder_attention_mask.squeeze(1).squeeze(1)
+ encoder_hidden_states = encoder_hidden_states.squeeze(1).masked_select(encoder_attention_mask.unsqueeze(-1) != 0).view(1, -1, hidden_states.shape[-1]) # [1, N_valid_tokens, C]
+ y_seqlens = encoder_attention_mask.sum(dim=1).tolist() # [B]
+ else:
+ y_seqlens = [encoder_hidden_states.shape[2]] * encoder_hidden_states.shape[0]
+ encoder_hidden_states = encoder_hidden_states.squeeze(1).view(1, -1, hidden_states.shape[-1])
+
+ # if self.cp_split_hw[0] * self.cp_split_hw[1] > 1:
+ # hidden_states = rearrange(hidden_states, "B (T H W) C -> B T H W C", T=N_t, H=N_h, W=N_w)
+ # hidden_states = context_parallel_util.split_cp_2d(hidden_states, seq_dim_hw=(2, 3), split_hw=self.cp_split_hw)
+ # hidden_states = rearrange(hidden_states, "B T H W C -> B (T H W) C")
+
+ # blocks
+ kv_cache_dict_ret = {}
+ for i, block in enumerate(self.blocks):
+ block_outputs = gradient_checkpoint_forward(
+ block,
+ use_gradient_checkpointing=use_gradient_checkpointing,
+ use_gradient_checkpointing_offload=use_gradient_checkpointing_offload,
+ x=hidden_states,
+ y=encoder_hidden_states,
+ t=t,
+ y_seqlen=y_seqlens,
+ latent_shape=(N_t, N_h, N_w),
+ num_cond_latents=num_cond_latents,
+ return_kv=return_kv,
+ kv_cache=kv_cache_dict.get(i, None),
+ skip_crs_attn=skip_crs_attn,
+ )
+
+ if return_kv:
+ hidden_states, kv_cache = block_outputs
+ if offload_kv_cache:
+ kv_cache_dict_ret[i] = (kv_cache[0].cpu(), kv_cache[1].cpu())
+ else:
+ kv_cache_dict_ret[i] = (kv_cache[0].contiguous(), kv_cache[1].contiguous())
+ else:
+ hidden_states = block_outputs
+
+ hidden_states = self.final_layer(hidden_states, t, (N_t, N_h, N_w)) # [B, N, C=T_p*H_p*W_p*C_out]
+
+ # if self.cp_split_hw[0] * self.cp_split_hw[1] > 1:
+ # hidden_states = context_parallel_util.gather_cp_2d(hidden_states, shape=(N_t, N_h, N_w), split_hw=self.cp_split_hw)
+
+ hidden_states = self.unpatchify(hidden_states, N_t, N_h, N_w) # [B, C_out, H, W]
+
+ # cast to float32 for better accuracy
+ hidden_states = hidden_states.to(torch.float32)
+
+ if return_kv:
+ return hidden_states, kv_cache_dict_ret
+ else:
+ return hidden_states
+
+
+ def unpatchify(self, x, N_t, N_h, N_w):
+ """
+ Args:
+ x (torch.Tensor): of shape [B, N, C]
+
+ Return:
+ x (torch.Tensor): of shape [B, C_out, T, H, W]
+ """
+ T_p, H_p, W_p = self.patch_size
+ x = rearrange(
+ x,
+ "B (N_t N_h N_w) (T_p H_p W_p C_out) -> B C_out (N_t T_p) (N_h H_p) (N_w W_p)",
+ N_t=N_t,
+ N_h=N_h,
+ N_w=N_w,
+ T_p=T_p,
+ H_p=H_p,
+ W_p=W_p,
+ C_out=self.out_channels,
+ )
+ return x
+
+ @staticmethod
+ def state_dict_converter():
+ return LongCatVideoTransformer3DModelDictConverter()
+
+
+class LongCatVideoTransformer3DModelDictConverter:
+ def __init__(self):
+ pass
+
+ def from_diffusers(self, state_dict):
+ return state_dict
+
+ def from_civitai(self, state_dict):
+ return state_dict
+
diff --git a/diffsynth/models/lpips.py b/diffsynth/models/lpips.py
new file mode 100644
index 0000000000000000000000000000000000000000..4c7847c89e1c78a3b6c450c04fdb5464c24be904
--- /dev/null
+++ b/diffsynth/models/lpips.py
@@ -0,0 +1,351 @@
+import os
+from collections import defaultdict
+from pathlib import Path
+from typing import Union
+
+import torch
+import torch.nn as nn
+from PIL import Image
+from torchvision import transforms
+
+ImageInput = Union[str, os.PathLike, Image.Image]
+
+IMAGE_EXTENSIONS = {".bmp", ".jpg", ".jpeg", ".pgm", ".png", ".ppm", ".tif", ".tiff", ".webp"}
+
+LPIPS_NET_CHOICES = ("alex", "vgg", "squeeze")
+
+
+def _list_image_files(path: Union[str, os.PathLike]):
+ path = os.fspath(path)
+ if not os.path.isdir(path):
+ raise ValueError(f"Expected a directory for LPIPS, got: {path}")
+ files = []
+ for entry in sorted(os.listdir(path)):
+ full = os.path.join(path, entry)
+ if os.path.isfile(full) and os.path.splitext(entry)[1].lower() in IMAGE_EXTENSIONS:
+ files.append(full)
+ if not files:
+ raise ValueError(f"No images found under {path}.")
+ return files
+
+
+def _pair_directories_by_stem(dir_a, dir_b):
+ files_a = _list_image_files(dir_a)
+ files_b = _list_image_files(dir_b)
+ by_stem_a = defaultdict(list)
+ for f in files_a:
+ by_stem_a[Path(f).stem].append(f)
+ by_stem_b = defaultdict(list)
+ for f in files_b:
+ by_stem_b[Path(f).stem].append(f)
+ common = sorted(set(by_stem_a.keys()) & set(by_stem_b.keys()))
+ if not common:
+ raise ValueError(f"No matching filename stems between {dir_a} and {dir_b}.")
+ pairs = []
+ for stem in common:
+ pairs.append((sorted(by_stem_a[stem])[0], sorted(by_stem_b[stem])[0]))
+ return pairs
+
+
+def _open_rgb(image: ImageInput) -> Image.Image:
+ if isinstance(image, (str, os.PathLike)):
+ image = Image.open(image)
+ if not isinstance(image, Image.Image):
+ raise TypeError(f"LPIPS expects PIL images or image paths, got {type(image)}.")
+ return image.convert("RGB")
+
+
+class _AlexFeatures(nn.Module):
+ def __init__(self):
+ super().__init__()
+ self.slice1 = nn.Sequential()
+ self.slice2 = nn.Sequential()
+ self.slice3 = nn.Sequential()
+ self.slice4 = nn.Sequential()
+ self.slice5 = nn.Sequential()
+ self.slice1.add_module("0", nn.Conv2d(3, 64, kernel_size=11, stride=4, padding=2))
+ self.slice1.add_module("1", nn.ReLU(inplace=True))
+ self.slice2.add_module("2", nn.MaxPool2d(kernel_size=3, stride=2))
+ self.slice2.add_module("3", nn.Conv2d(64, 192, kernel_size=5, padding=2))
+ self.slice2.add_module("4", nn.ReLU(inplace=True))
+ self.slice3.add_module("5", nn.MaxPool2d(kernel_size=3, stride=2))
+ self.slice3.add_module("6", nn.Conv2d(192, 384, kernel_size=3, padding=1))
+ self.slice3.add_module("7", nn.ReLU(inplace=True))
+ self.slice4.add_module("8", nn.Conv2d(384, 256, kernel_size=3, padding=1))
+ self.slice4.add_module("9", nn.ReLU(inplace=True))
+ self.slice5.add_module("10", nn.Conv2d(256, 256, kernel_size=3, padding=1))
+ self.slice5.add_module("11", nn.ReLU(inplace=True))
+
+ def forward(self, x):
+ h1 = self.slice1(x)
+ h2 = self.slice2(h1)
+ h3 = self.slice3(h2)
+ h4 = self.slice4(h3)
+ h5 = self.slice5(h4)
+ return [h1, h2, h3, h4, h5]
+
+
+class _VGG16Features(nn.Module):
+ def __init__(self):
+ super().__init__()
+ self.slice1 = nn.Sequential()
+ self.slice2 = nn.Sequential()
+ self.slice3 = nn.Sequential()
+ self.slice4 = nn.Sequential()
+ self.slice5 = nn.Sequential()
+ cfg = [
+ (1, 0, nn.Conv2d(3, 64, 3, padding=1)),
+ (1, 1, nn.ReLU(inplace=True)),
+ (1, 2, nn.Conv2d(64, 64, 3, padding=1)),
+ (1, 3, nn.ReLU(inplace=True)),
+ (2, 4, nn.MaxPool2d(2, 2)),
+ (2, 5, nn.Conv2d(64, 128, 3, padding=1)),
+ (2, 6, nn.ReLU(inplace=True)),
+ (2, 7, nn.Conv2d(128, 128, 3, padding=1)),
+ (2, 8, nn.ReLU(inplace=True)),
+ (3, 9, nn.MaxPool2d(2, 2)),
+ (3, 10, nn.Conv2d(128, 256, 3, padding=1)),
+ (3, 11, nn.ReLU(inplace=True)),
+ (3, 12, nn.Conv2d(256, 256, 3, padding=1)),
+ (3, 13, nn.ReLU(inplace=True)),
+ (3, 14, nn.Conv2d(256, 256, 3, padding=1)),
+ (3, 15, nn.ReLU(inplace=True)),
+ (4, 16, nn.MaxPool2d(2, 2)),
+ (4, 17, nn.Conv2d(256, 512, 3, padding=1)),
+ (4, 18, nn.ReLU(inplace=True)),
+ (4, 19, nn.Conv2d(512, 512, 3, padding=1)),
+ (4, 20, nn.ReLU(inplace=True)),
+ (4, 21, nn.Conv2d(512, 512, 3, padding=1)),
+ (4, 22, nn.ReLU(inplace=True)),
+ (5, 23, nn.MaxPool2d(2, 2)),
+ (5, 24, nn.Conv2d(512, 512, 3, padding=1)),
+ (5, 25, nn.ReLU(inplace=True)),
+ (5, 26, nn.Conv2d(512, 512, 3, padding=1)),
+ (5, 27, nn.ReLU(inplace=True)),
+ (5, 28, nn.Conv2d(512, 512, 3, padding=1)),
+ (5, 29, nn.ReLU(inplace=True)),
+ ]
+ for slice_idx, orig_idx, module in cfg:
+ getattr(self, f"slice{slice_idx}").add_module(str(orig_idx), module)
+
+ def forward(self, x):
+ h1 = self.slice1(x)
+ h2 = self.slice2(h1)
+ h3 = self.slice3(h2)
+ h4 = self.slice4(h3)
+ h5 = self.slice5(h4)
+ return [h1, h2, h3, h4, h5]
+
+
+class _Fire(nn.Module):
+ def __init__(self, in_channels, squeeze_channels, expand1x1_channels, expand3x3_channels):
+ super().__init__()
+ self.squeeze = nn.Conv2d(in_channels, squeeze_channels, kernel_size=1)
+ self.squeeze_activation = nn.ReLU(inplace=True)
+ self.expand1x1 = nn.Conv2d(squeeze_channels, expand1x1_channels, kernel_size=1)
+ self.expand1x1_activation = nn.ReLU(inplace=True)
+ self.expand3x3 = nn.Conv2d(squeeze_channels, expand3x3_channels, kernel_size=3, padding=1)
+ self.expand3x3_activation = nn.ReLU(inplace=True)
+
+ def forward(self, x):
+ x = self.squeeze_activation(self.squeeze(x))
+ return torch.cat(
+ [
+ self.expand1x1_activation(self.expand1x1(x)),
+ self.expand3x3_activation(self.expand3x3(x)),
+ ],
+ dim=1,
+ )
+
+
+class _SqueezeNet11Features(nn.Module):
+ def __init__(self):
+ super().__init__()
+ self.slice1 = nn.Sequential()
+ self.slice2 = nn.Sequential()
+ self.slice3 = nn.Sequential()
+ self.slice4 = nn.Sequential()
+ self.slice5 = nn.Sequential()
+ self.slice6 = nn.Sequential()
+ self.slice7 = nn.Sequential()
+ self.slice1.add_module("0", nn.Conv2d(3, 64, kernel_size=3, stride=2))
+ self.slice1.add_module("1", nn.ReLU(inplace=True))
+ self.slice2.add_module("2", nn.MaxPool2d(kernel_size=3, stride=2, ceil_mode=True))
+ self.slice2.add_module("3", _Fire(64, 16, 64, 64))
+ self.slice2.add_module("4", _Fire(128, 16, 64, 64))
+ self.slice3.add_module("5", nn.MaxPool2d(kernel_size=3, stride=2, ceil_mode=True))
+ self.slice3.add_module("6", _Fire(128, 32, 128, 128))
+ self.slice3.add_module("7", _Fire(256, 32, 128, 128))
+ self.slice4.add_module("8", nn.MaxPool2d(kernel_size=3, stride=2, ceil_mode=True))
+ self.slice4.add_module("9", _Fire(256, 48, 192, 192))
+ self.slice5.add_module("10", _Fire(384, 48, 192, 192))
+ self.slice6.add_module("11", _Fire(384, 64, 256, 256))
+ self.slice7.add_module("12", _Fire(512, 64, 256, 256))
+
+ def forward(self, x):
+ h1 = self.slice1(x)
+ h2 = self.slice2(h1)
+ h3 = self.slice3(h2)
+ h4 = self.slice4(h3)
+ h5 = self.slice5(h4)
+ h6 = self.slice6(h5)
+ h7 = self.slice7(h6)
+ return [h1, h2, h3, h4, h5, h6, h7]
+
+
+_NET_CONFIG = {
+ "alex": {"factory": _AlexFeatures, "channels": (64, 192, 384, 256, 256)},
+ "vgg": {"factory": _VGG16Features, "channels": (64, 128, 256, 512, 512)},
+ "squeeze": {"factory": _SqueezeNet11Features, "channels": (64, 128, 256, 384, 384, 512, 512)},
+}
+
+
+class _ScalingLayer(nn.Module):
+ def __init__(self):
+ super().__init__()
+ self.register_buffer("shift", torch.tensor([-0.030, -0.088, -0.188]).view(1, 3, 1, 1))
+ self.register_buffer("scale", torch.tensor([0.458, 0.448, 0.450]).view(1, 3, 1, 1))
+
+ def forward(self, x):
+ return (x - self.shift) / self.scale
+
+
+class _NetLinLayer(nn.Module):
+ def __init__(self, chn_in, use_dropout=True):
+ super().__init__()
+ layers = []
+ if use_dropout:
+ layers.append(nn.Dropout())
+ layers.append(nn.Conv2d(chn_in, 1, kernel_size=1, stride=1, padding=0, bias=False))
+ self.model = nn.Sequential(*layers)
+
+ def forward(self, x):
+ return self.model(x)
+
+
+def _normalize_tensor(x, eps=1e-10):
+ norm = torch.sqrt(torch.sum(x**2, dim=1, keepdim=True))
+ return x / (norm + eps)
+
+
+def _spatial_average(x):
+ return x.mean(dim=(2, 3), keepdim=True)
+
+
+class LPIPSModel(nn.Module):
+ def __init__(self, net: str = "alex", use_dropout: bool = True):
+ super().__init__()
+ if net not in _NET_CONFIG:
+ raise ValueError(f"net must be one of {LPIPS_NET_CHOICES}, got {net!r}")
+ self.net_name = net
+ self.scaling_layer = _ScalingLayer()
+ self.net = _NET_CONFIG[net]["factory"]()
+ chns = _NET_CONFIG[net]["channels"]
+ for i, chn in enumerate(chns):
+ setattr(self, f"lin{i}", _NetLinLayer(chn, use_dropout=use_dropout))
+ self.num_layers = len(chns)
+ for p in self.parameters():
+ p.requires_grad = False
+
+ def forward(self, in0, in1):
+ in0 = self.scaling_layer(in0)
+ in1 = self.scaling_layer(in1)
+ feats0 = self.net(in0)
+ feats1 = self.net(in1)
+ val = 0
+ for i in range(self.num_layers):
+ diff = (_normalize_tensor(feats0[i]) - _normalize_tensor(feats1[i])) ** 2
+ lin = getattr(self, f"lin{i}")
+ val = val + _spatial_average(lin(diff))
+ return val.view(-1)
+
+
+class LPIPSCompute(nn.Module):
+ def __init__(
+ self,
+ model: LPIPSModel,
+ device: Union[str, torch.device] = "cpu",
+ batch_size: int = 16,
+ target_size: int = 512,
+ ):
+ super().__init__()
+ self.model = model
+ self.batch_size = batch_size
+ self.target_size = target_size
+ self._resize_transform = transforms.Compose(
+ [
+ transforms.Resize(target_size, interpolation=transforms.InterpolationMode.BICUBIC),
+ transforms.CenterCrop(target_size),
+ transforms.ToTensor(),
+ ]
+ )
+ self._raw_transform = transforms.ToTensor()
+ self.to(device)
+
+ @property
+ def device(self):
+ try:
+ return next(self.model.parameters()).device
+ except StopIteration:
+ return torch.device("cpu")
+
+ def _to_tensor(self, image: Image.Image, do_resize: bool) -> torch.Tensor:
+ transform = self._resize_transform if do_resize else self._raw_transform
+ x = transform(image).clamp(0.0, 1.0) * 2.0 - 1.0
+ return x
+
+ @torch.no_grad()
+ def _compute_pair(self, img_a: Image.Image, img_b: Image.Image, do_resize: bool) -> float:
+ x0 = self._to_tensor(img_a, do_resize).unsqueeze(0).to(self.device)
+ x1 = self._to_tensor(img_b, do_resize).unsqueeze(0).to(self.device)
+ return float(self.model(x0, x1).item())
+
+ @torch.no_grad()
+ def _compute_pairs(self, pairs, do_resize: bool) -> float:
+ scores = []
+ batch_size = max(1, self.batch_size)
+ for start in range(0, len(pairs), batch_size):
+ chunk = pairs[start : start + batch_size]
+ xs0 = torch.stack([self._to_tensor(_open_rgb(a), do_resize) for a, _ in chunk]).to(self.device)
+ xs1 = torch.stack([self._to_tensor(_open_rgb(b), do_resize) for _, b in chunk]).to(self.device)
+ scores.append(self.model(xs0, xs1).detach().cpu())
+ merged = torch.cat(scores, dim=0)
+ return float(merged.mean().item())
+
+ @staticmethod
+ def _is_dir(value) -> bool:
+ return isinstance(value, (str, os.PathLike)) and os.path.isdir(os.fspath(value))
+
+ @staticmethod
+ def _is_image_input(value) -> bool:
+ if isinstance(value, Image.Image):
+ return True
+ if isinstance(value, (str, os.PathLike)):
+ return os.path.isfile(os.fspath(value))
+ return False
+
+ def compute(self, image_a, image_b) -> float:
+ a_is_dir = self._is_dir(image_a)
+ b_is_dir = self._is_dir(image_b)
+ if a_is_dir != b_is_dir:
+ raise ValueError("LPIPS.compute requires both inputs to be directories or both to be single images.")
+
+ if a_is_dir:
+ pairs = _pair_directories_by_stem(image_a, image_b)
+ sizes = set()
+ for path_a, path_b in pairs:
+ with Image.open(path_a) as ia, Image.open(path_b) as ib:
+ sizes.add(ia.size)
+ sizes.add(ib.size)
+ do_resize = len(sizes) > 1
+ return self._compute_pairs(pairs, do_resize=do_resize)
+
+ if not (self._is_image_input(image_a) and self._is_image_input(image_b)):
+ raise ValueError("LPIPS.compute inputs must be image paths, PIL images, or directories.")
+ img_a = _open_rgb(image_a)
+ img_b = _open_rgb(image_b)
+ do_resize = img_a.size != img_b.size
+ return self._compute_pair(img_a, img_b, do_resize=do_resize)
+
+ def forward(self, image_a, image_b):
+ return self.compute(image_a, image_b)
diff --git a/diffsynth/models/ltx2_audio_vae.py b/diffsynth/models/ltx2_audio_vae.py
new file mode 100644
index 0000000000000000000000000000000000000000..8a58f97245d3d14715740c7c04e8e53b3162dfa1
--- /dev/null
+++ b/diffsynth/models/ltx2_audio_vae.py
@@ -0,0 +1,1874 @@
+from typing import Set, Tuple, Optional, List
+from enum import Enum
+import math
+import einops
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import torchaudio
+from .ltx2_common import VideoLatentShape, AudioLatentShape, Patchifier, NormType, build_normalization_layer
+
+
+class AudioProcessor(nn.Module):
+ """Converts audio waveforms to log-mel spectrograms with optional resampling."""
+
+ def __init__(
+ self,
+ sample_rate: int = 16000,
+ mel_bins: int = 64,
+ mel_hop_length: int = 160,
+ n_fft: int = 1024,
+ ) -> None:
+ super().__init__()
+ self.sample_rate = sample_rate
+ self.mel_transform = torchaudio.transforms.MelSpectrogram(
+ sample_rate=sample_rate,
+ n_fft=n_fft,
+ win_length=n_fft,
+ hop_length=mel_hop_length,
+ f_min=0.0,
+ f_max=sample_rate / 2.0,
+ n_mels=mel_bins,
+ window_fn=torch.hann_window,
+ center=True,
+ pad_mode="reflect",
+ power=1.0,
+ mel_scale="slaney",
+ norm="slaney",
+ )
+
+ def resample_waveform(
+ self,
+ waveform: torch.Tensor,
+ source_rate: int,
+ target_rate: int,
+ ) -> torch.Tensor:
+ """Resample waveform to target sample rate if needed."""
+ if source_rate == target_rate:
+ return waveform
+ resampled = torchaudio.functional.resample(waveform, source_rate, target_rate)
+ return resampled.to(device=waveform.device, dtype=waveform.dtype)
+
+ def waveform_to_mel(
+ self,
+ waveform: torch.Tensor,
+ waveform_sample_rate: int,
+ ) -> torch.Tensor:
+ """Convert waveform to log-mel spectrogram [batch, channels, time, n_mels]."""
+ original_dtype = waveform.dtype
+ waveform = self.resample_waveform(waveform, waveform_sample_rate, self.sample_rate)
+
+ self.mel_transform = self.mel_transform.to(dtype=torch.float32)
+ mel = self.mel_transform(waveform.to(dtype=torch.float32))
+ mel = torch.log(torch.clamp(mel, min=1e-5))
+
+ mel = mel.to(device=waveform.device, dtype=original_dtype)
+ return mel.permute(0, 1, 3, 2).contiguous()
+
+
+class AudioPatchifier(Patchifier):
+ def __init__(
+ self,
+ patch_size: int,
+ sample_rate: int = 16000,
+ hop_length: int = 160,
+ audio_latent_downsample_factor: int = 4,
+ is_causal: bool = True,
+ shift: int = 0,
+ ):
+ """
+ Patchifier tailored for spectrogram/audio latents.
+ Args:
+ patch_size: Number of mel bins combined into a single patch. This
+ controls the resolution along the frequency axis.
+ sample_rate: Original waveform sampling rate. Used to map latent
+ indices back to seconds so downstream consumers can align audio
+ and video cues.
+ hop_length: Window hop length used for the spectrogram. Determines
+ how many real-time samples separate two consecutive latent frames.
+ audio_latent_downsample_factor: Ratio between spectrogram frames and
+ latent frames; compensates for additional downsampling inside the
+ VAE encoder.
+ is_causal: When True, timing is shifted to account for causal
+ receptive fields so timestamps do not peek into the future.
+ shift: Integer offset applied to the latent indices. Enables
+ constructing overlapping windows from the same latent sequence.
+ """
+ self.hop_length = hop_length
+ self.sample_rate = sample_rate
+ self.audio_latent_downsample_factor = audio_latent_downsample_factor
+ self.is_causal = is_causal
+ self.shift = shift
+ self._patch_size = (1, patch_size, patch_size)
+
+ @property
+ def patch_size(self) -> Tuple[int, int, int]:
+ return self._patch_size
+
+ def get_token_count(self, tgt_shape: AudioLatentShape) -> int:
+ return tgt_shape.frames
+
+ def _get_audio_latent_time_in_sec(
+ self,
+ start_latent: int,
+ end_latent: int,
+ dtype: torch.dtype,
+ device: Optional[torch.device] = None,
+ ) -> torch.Tensor:
+ """
+ Converts latent indices into real-time seconds while honoring causal
+ offsets and the configured hop length.
+ Args:
+ start_latent: Inclusive start index inside the latent sequence. This
+ sets the first timestamp returned.
+ end_latent: Exclusive end index. Determines how many timestamps get
+ generated.
+ dtype: Floating-point dtype used for the returned tensor, allowing
+ callers to control precision.
+ device: Target device for the timestamp tensor. When omitted the
+ computation occurs on CPU to avoid surprising GPU allocations.
+ """
+ if device is None:
+ device = torch.device("cpu")
+
+ audio_latent_frame = torch.arange(start_latent, end_latent, dtype=dtype, device=device)
+
+ audio_mel_frame = audio_latent_frame * self.audio_latent_downsample_factor
+
+ if self.is_causal:
+ # Frame offset for causal alignment.
+ # The "+1" ensures the timestamp corresponds to the first sample that is fully available.
+ causal_offset = 1
+ audio_mel_frame = (audio_mel_frame + causal_offset - self.audio_latent_downsample_factor).clip(min=0)
+
+ return audio_mel_frame * self.hop_length / self.sample_rate
+
+ def _compute_audio_timings(
+ self,
+ batch_size: int,
+ num_steps: int,
+ device: Optional[torch.device] = None,
+ ) -> torch.Tensor:
+ """
+ Builds a `(B, 1, T, 2)` tensor containing timestamps for each latent frame.
+ This helper method underpins `get_patch_grid_bounds` for the audio patchifier.
+ Args:
+ batch_size: Number of sequences to broadcast the timings over.
+ num_steps: Number of latent frames (time steps) to convert into timestamps.
+ device: Device on which the resulting tensor should reside.
+ """
+ resolved_device = device
+ if resolved_device is None:
+ resolved_device = torch.device("cpu")
+
+ start_timings = self._get_audio_latent_time_in_sec(
+ self.shift,
+ num_steps + self.shift,
+ torch.float32,
+ resolved_device,
+ )
+ start_timings = start_timings.unsqueeze(0).expand(batch_size, -1).unsqueeze(1)
+
+ end_timings = self._get_audio_latent_time_in_sec(
+ self.shift + 1,
+ num_steps + self.shift + 1,
+ torch.float32,
+ resolved_device,
+ )
+ end_timings = end_timings.unsqueeze(0).expand(batch_size, -1).unsqueeze(1)
+
+ return torch.stack([start_timings, end_timings], dim=-1)
+
+ def patchify(
+ self,
+ audio_latents: torch.Tensor,
+ ) -> torch.Tensor:
+ """
+ Flattens the audio latent tensor along time. Use `get_patch_grid_bounds`
+ to derive timestamps for each latent frame based on the configured hop
+ length and downsampling.
+ Args:
+ audio_latents: Latent tensor to patchify.
+ Returns:
+ Flattened patch tokens tensor. Use `get_patch_grid_bounds` to compute the
+ corresponding timing metadata when needed.
+ """
+ audio_latents = einops.rearrange(
+ audio_latents,
+ "b c t f -> b t (c f)",
+ )
+
+ return audio_latents
+
+ def unpatchify(
+ self,
+ audio_latents: torch.Tensor,
+ output_shape: AudioLatentShape,
+ ) -> torch.Tensor:
+ """
+ Restores the `(B, C, T, F)` spectrogram tensor from flattened patches.
+ Use `get_patch_grid_bounds` to recompute the timestamps that describe each
+ frame's position in real time.
+ Args:
+ audio_latents: Latent tensor to unpatchify.
+ output_shape: Shape of the unpatched output tensor.
+ Returns:
+ Unpatched latent tensor. Use `get_patch_grid_bounds` to compute the timing
+ metadata associated with the restored latents.
+ """
+ # audio_latents shape: (batch, time, freq * channels)
+ audio_latents = einops.rearrange(
+ audio_latents,
+ "b t (c f) -> b c t f",
+ c=output_shape.channels,
+ f=output_shape.mel_bins,
+ )
+
+ return audio_latents
+
+ def unpatchify_audio(
+ self,
+ audio_latents: torch.Tensor,
+ channels: int,
+ mel_bins: int
+ ) -> torch.Tensor:
+ audio_latents = einops.rearrange(
+ audio_latents,
+ "b t (c f) -> b c t f",
+ c=channels,
+ f=mel_bins,
+ )
+ return audio_latents
+
+ def get_patch_grid_bounds(
+ self,
+ output_shape: AudioLatentShape | VideoLatentShape,
+ device: Optional[torch.device] = None,
+ ) -> torch.Tensor:
+ """
+ Return the temporal bounds `[inclusive start, exclusive end)` for every
+ patch emitted by `patchify`. For audio this corresponds to timestamps in
+ seconds aligned with the original spectrogram grid.
+ The returned tensor has shape `[batch_size, 1, time_steps, 2]`, where:
+ - axis 1 (size 1) represents the temporal dimension
+ - axis 3 (size 2) stores the `[start, end)` timestamps per patch
+ Args:
+ output_shape: Audio grid specification describing the number of time steps.
+ device: Target device for the returned tensor.
+ """
+ if not isinstance(output_shape, AudioLatentShape):
+ raise ValueError("AudioPatchifier expects AudioLatentShape when computing coordinates")
+
+ return self._compute_audio_timings(output_shape.batch, output_shape.frames, device)
+
+
+class AttentionType(Enum):
+ """Enum for specifying the attention mechanism type."""
+
+ VANILLA = "vanilla"
+ LINEAR = "linear"
+ NONE = "none"
+
+
+class AttnBlock(torch.nn.Module):
+ def __init__(
+ self,
+ in_channels: int,
+ norm_type: NormType = NormType.GROUP,
+ ) -> None:
+ super().__init__()
+ self.in_channels = in_channels
+
+ self.norm = build_normalization_layer(in_channels, normtype=norm_type)
+ self.q = torch.nn.Conv2d(in_channels, in_channels, kernel_size=1, stride=1, padding=0)
+ self.k = torch.nn.Conv2d(in_channels, in_channels, kernel_size=1, stride=1, padding=0)
+ self.v = torch.nn.Conv2d(in_channels, in_channels, kernel_size=1, stride=1, padding=0)
+ self.proj_out = torch.nn.Conv2d(in_channels, in_channels, kernel_size=1, stride=1, padding=0)
+
+ def forward(self, x: torch.Tensor) -> torch.Tensor:
+ h_ = x
+ h_ = self.norm(h_)
+ q = self.q(h_)
+ k = self.k(h_)
+ v = self.v(h_)
+
+ # compute attention
+ b, c, h, w = q.shape
+ q = q.reshape(b, c, h * w).contiguous()
+ q = q.permute(0, 2, 1).contiguous() # b,hw,c
+ k = k.reshape(b, c, h * w).contiguous() # b,c,hw
+ w_ = torch.bmm(q, k).contiguous() # b,hw,hw w[b,i,j]=sum_c q[b,i,c]k[b,c,j]
+ w_ = w_ * (int(c) ** (-0.5))
+ w_ = torch.nn.functional.softmax(w_, dim=2)
+
+ # attend to values
+ v = v.reshape(b, c, h * w).contiguous()
+ w_ = w_.permute(0, 2, 1).contiguous() # b,hw,hw (first hw of k, second of q)
+ h_ = torch.bmm(v, w_).contiguous() # b, c,hw (hw of q) h_[b,c,j] = sum_i v[b,c,i] w_[b,i,j]
+ h_ = h_.reshape(b, c, h, w).contiguous()
+
+ h_ = self.proj_out(h_)
+
+ return x + h_
+
+
+def make_attn(
+ in_channels: int,
+ attn_type: AttentionType = AttentionType.VANILLA,
+ norm_type: NormType = NormType.GROUP,
+) -> torch.nn.Module:
+ match attn_type:
+ case AttentionType.VANILLA:
+ return AttnBlock(in_channels, norm_type=norm_type)
+ case AttentionType.NONE:
+ return torch.nn.Identity()
+ case AttentionType.LINEAR:
+ raise NotImplementedError(f"Attention type {attn_type.value} is not supported yet.")
+ case _:
+ raise ValueError(f"Unknown attention type: {attn_type}")
+
+
+class CausalityAxis(Enum):
+ """Enum for specifying the causality axis in causal convolutions."""
+
+ NONE = None
+ WIDTH = "width"
+ HEIGHT = "height"
+ WIDTH_COMPATIBILITY = "width-compatibility"
+
+
+class CausalConv2d(torch.nn.Module):
+ """
+ A causal 2D convolution.
+ This layer ensures that the output at time `t` only depends on inputs
+ at time `t` and earlier. It achieves this by applying asymmetric padding
+ to the time dimension (width) before the convolution.
+ """
+
+ def __init__(
+ self,
+ in_channels: int,
+ out_channels: int,
+ kernel_size: int | tuple[int, int],
+ stride: int = 1,
+ dilation: int | tuple[int, int] = 1,
+ groups: int = 1,
+ bias: bool = True,
+ causality_axis: CausalityAxis = CausalityAxis.HEIGHT,
+ ) -> None:
+ super().__init__()
+
+ self.causality_axis = causality_axis
+
+ # Ensure kernel_size and dilation are tuples
+ kernel_size = torch.nn.modules.utils._pair(kernel_size)
+ dilation = torch.nn.modules.utils._pair(dilation)
+
+ # Calculate padding dimensions
+ pad_h = (kernel_size[0] - 1) * dilation[0]
+ pad_w = (kernel_size[1] - 1) * dilation[1]
+
+ # The padding tuple for F.pad is (pad_left, pad_right, pad_top, pad_bottom)
+ match self.causality_axis:
+ case CausalityAxis.NONE:
+ self.padding = (pad_w // 2, pad_w - pad_w // 2, pad_h // 2, pad_h - pad_h // 2)
+ case CausalityAxis.WIDTH | CausalityAxis.WIDTH_COMPATIBILITY:
+ self.padding = (pad_w, 0, pad_h // 2, pad_h - pad_h // 2)
+ case CausalityAxis.HEIGHT:
+ self.padding = (pad_w // 2, pad_w - pad_w // 2, pad_h, 0)
+ case _:
+ raise ValueError(f"Invalid causality_axis: {causality_axis}")
+
+ # The internal convolution layer uses no padding, as we handle it manually
+ self.conv = torch.nn.Conv2d(
+ in_channels,
+ out_channels,
+ kernel_size,
+ stride=stride,
+ padding=0,
+ dilation=dilation,
+ groups=groups,
+ bias=bias,
+ )
+
+ def forward(self, x: torch.Tensor) -> torch.Tensor:
+ # Apply causal padding before convolution
+ x = F.pad(x, self.padding)
+ return self.conv(x)
+
+
+def make_conv2d(
+ in_channels: int,
+ out_channels: int,
+ kernel_size: int | tuple[int, int],
+ stride: int = 1,
+ padding: tuple[int, int, int, int] | None = None,
+ dilation: int = 1,
+ groups: int = 1,
+ bias: bool = True,
+ causality_axis: CausalityAxis | None = None,
+) -> torch.nn.Module:
+ """
+ Create a 2D convolution layer that can be either causal or non-causal.
+ Args:
+ in_channels: Number of input channels
+ out_channels: Number of output channels
+ kernel_size: Size of the convolution kernel
+ stride: Convolution stride
+ padding: Padding (if None, will be calculated based on causal flag)
+ dilation: Dilation rate
+ groups: Number of groups for grouped convolution
+ bias: Whether to use bias
+ causality_axis: Dimension along which to apply causality.
+ Returns:
+ Either a regular Conv2d or CausalConv2d layer
+ """
+ if causality_axis is not None:
+ # For causal convolution, padding is handled internally by CausalConv2d
+ return CausalConv2d(in_channels, out_channels, kernel_size, stride, dilation, groups, bias, causality_axis)
+ else:
+ # For non-causal convolution, use symmetric padding if not specified
+ if padding is None:
+ padding = kernel_size // 2 if isinstance(kernel_size, int) else tuple(k // 2 for k in kernel_size)
+
+ return torch.nn.Conv2d(
+ in_channels,
+ out_channels,
+ kernel_size,
+ stride,
+ padding,
+ dilation,
+ groups,
+ bias,
+ )
+
+
+
+LRELU_SLOPE = 0.1
+
+
+class ResBlock1(torch.nn.Module):
+ def __init__(self, channels: int, kernel_size: int = 3, dilation: Tuple[int, int, int] = (1, 3, 5)):
+ super(ResBlock1, self).__init__()
+ self.convs1 = torch.nn.ModuleList(
+ [
+ torch.nn.Conv1d(
+ channels,
+ channels,
+ kernel_size,
+ 1,
+ dilation=dilation[0],
+ padding="same",
+ ),
+ torch.nn.Conv1d(
+ channels,
+ channels,
+ kernel_size,
+ 1,
+ dilation=dilation[1],
+ padding="same",
+ ),
+ torch.nn.Conv1d(
+ channels,
+ channels,
+ kernel_size,
+ 1,
+ dilation=dilation[2],
+ padding="same",
+ ),
+ ]
+ )
+
+ self.convs2 = torch.nn.ModuleList(
+ [
+ torch.nn.Conv1d(
+ channels,
+ channels,
+ kernel_size,
+ 1,
+ dilation=1,
+ padding="same",
+ ),
+ torch.nn.Conv1d(
+ channels,
+ channels,
+ kernel_size,
+ 1,
+ dilation=1,
+ padding="same",
+ ),
+ torch.nn.Conv1d(
+ channels,
+ channels,
+ kernel_size,
+ 1,
+ dilation=1,
+ padding="same",
+ ),
+ ]
+ )
+
+ def forward(self, x: torch.Tensor) -> torch.Tensor:
+ for conv1, conv2 in zip(self.convs1, self.convs2, strict=True):
+ xt = torch.nn.functional.leaky_relu(x, LRELU_SLOPE)
+ xt = conv1(xt)
+ xt = torch.nn.functional.leaky_relu(xt, LRELU_SLOPE)
+ xt = conv2(xt)
+ x = xt + x
+ return x
+
+
+class ResBlock2(torch.nn.Module):
+ def __init__(self, channels: int, kernel_size: int = 3, dilation: Tuple[int, int] = (1, 3)):
+ super(ResBlock2, self).__init__()
+ self.convs = torch.nn.ModuleList(
+ [
+ torch.nn.Conv1d(
+ channels,
+ channels,
+ kernel_size,
+ 1,
+ dilation=dilation[0],
+ padding="same",
+ ),
+ torch.nn.Conv1d(
+ channels,
+ channels,
+ kernel_size,
+ 1,
+ dilation=dilation[1],
+ padding="same",
+ ),
+ ]
+ )
+
+ def forward(self, x: torch.Tensor) -> torch.Tensor:
+ for conv in self.convs:
+ xt = torch.nn.functional.leaky_relu(x, LRELU_SLOPE)
+ xt = conv(xt)
+ x = xt + x
+ return x
+
+
+class ResnetBlock(torch.nn.Module):
+ def __init__(
+ self,
+ *,
+ in_channels: int,
+ out_channels: int | None = None,
+ conv_shortcut: bool = False,
+ dropout: float = 0.0,
+ temb_channels: int = 512,
+ norm_type: NormType = NormType.GROUP,
+ causality_axis: CausalityAxis = CausalityAxis.HEIGHT,
+ ) -> None:
+ super().__init__()
+ self.causality_axis = causality_axis
+
+ if self.causality_axis != CausalityAxis.NONE and norm_type == NormType.GROUP:
+ raise ValueError("Causal ResnetBlock with GroupNorm is not supported.")
+ self.in_channels = in_channels
+ out_channels = in_channels if out_channels is None else out_channels
+ self.out_channels = out_channels
+ self.use_conv_shortcut = conv_shortcut
+
+ self.norm1 = build_normalization_layer(in_channels, normtype=norm_type)
+ self.non_linearity = torch.nn.SiLU()
+ self.conv1 = make_conv2d(in_channels, out_channels, kernel_size=3, stride=1, causality_axis=causality_axis)
+ if temb_channels > 0:
+ self.temb_proj = torch.nn.Linear(temb_channels, out_channels)
+ self.norm2 = build_normalization_layer(out_channels, normtype=norm_type)
+ self.dropout = torch.nn.Dropout(dropout)
+ self.conv2 = make_conv2d(out_channels, out_channels, kernel_size=3, stride=1, causality_axis=causality_axis)
+ if self.in_channels != self.out_channels:
+ if self.use_conv_shortcut:
+ self.conv_shortcut = make_conv2d(
+ in_channels, out_channels, kernel_size=3, stride=1, causality_axis=causality_axis
+ )
+ else:
+ self.nin_shortcut = make_conv2d(
+ in_channels, out_channels, kernel_size=1, stride=1, causality_axis=causality_axis
+ )
+
+ def forward(
+ self,
+ x: torch.Tensor,
+ temb: torch.Tensor | None = None,
+ ) -> torch.Tensor:
+ h = x
+ h = self.norm1(h)
+ h = self.non_linearity(h)
+ h = self.conv1(h)
+
+ if temb is not None:
+ h = h + self.temb_proj(self.non_linearity(temb))[:, :, None, None]
+
+ h = self.norm2(h)
+ h = self.non_linearity(h)
+ h = self.dropout(h)
+ h = self.conv2(h)
+
+ if self.in_channels != self.out_channels:
+ x = self.conv_shortcut(x) if self.use_conv_shortcut else self.nin_shortcut(x)
+
+ return x + h
+
+
+class Downsample(torch.nn.Module):
+ """
+ A downsampling layer that can use either a strided convolution
+ or average pooling. Supports standard and causal padding for the
+ convolutional mode.
+ """
+
+ def __init__(
+ self,
+ in_channels: int,
+ with_conv: bool,
+ causality_axis: CausalityAxis = CausalityAxis.WIDTH,
+ ) -> None:
+ super().__init__()
+ self.with_conv = with_conv
+ self.causality_axis = causality_axis
+
+ if self.causality_axis != CausalityAxis.NONE and not self.with_conv:
+ raise ValueError("causality is only supported when `with_conv=True`.")
+
+ if self.with_conv:
+ # Do time downsampling here
+ # no asymmetric padding in torch conv, must do it ourselves
+ self.conv = torch.nn.Conv2d(in_channels, in_channels, kernel_size=3, stride=2, padding=0)
+
+ def forward(self, x: torch.Tensor) -> torch.Tensor:
+ if self.with_conv:
+ # Padding tuple is in the order: (left, right, top, bottom).
+ match self.causality_axis:
+ case CausalityAxis.NONE:
+ pad = (0, 1, 0, 1)
+ case CausalityAxis.WIDTH:
+ pad = (2, 0, 0, 1)
+ case CausalityAxis.HEIGHT:
+ pad = (0, 1, 2, 0)
+ case CausalityAxis.WIDTH_COMPATIBILITY:
+ pad = (1, 0, 0, 1)
+ case _:
+ raise ValueError(f"Invalid causality_axis: {self.causality_axis}")
+
+ x = torch.nn.functional.pad(x, pad, mode="constant", value=0)
+ x = self.conv(x)
+ else:
+ # This branch is only taken if with_conv=False, which implies causality_axis is NONE.
+ x = torch.nn.functional.avg_pool2d(x, kernel_size=2, stride=2)
+
+ return x
+
+
+def build_downsampling_path( # noqa: PLR0913
+ *,
+ ch: int,
+ ch_mult: Tuple[int, ...],
+ num_resolutions: int,
+ num_res_blocks: int,
+ resolution: int,
+ temb_channels: int,
+ dropout: float,
+ norm_type: NormType,
+ causality_axis: CausalityAxis,
+ attn_type: AttentionType,
+ attn_resolutions: Set[int],
+ resamp_with_conv: bool,
+) -> tuple[torch.nn.ModuleList, int]:
+ """Build the downsampling path with residual blocks, attention, and downsampling layers."""
+ down_modules = torch.nn.ModuleList()
+ curr_res = resolution
+ in_ch_mult = (1, *tuple(ch_mult))
+ block_in = ch
+
+ for i_level in range(num_resolutions):
+ block = torch.nn.ModuleList()
+ attn = torch.nn.ModuleList()
+ block_in = ch * in_ch_mult[i_level]
+ block_out = ch * ch_mult[i_level]
+
+ for _ in range(num_res_blocks):
+ block.append(
+ ResnetBlock(
+ in_channels=block_in,
+ out_channels=block_out,
+ temb_channels=temb_channels,
+ dropout=dropout,
+ norm_type=norm_type,
+ causality_axis=causality_axis,
+ )
+ )
+ block_in = block_out
+ if curr_res in attn_resolutions:
+ attn.append(make_attn(block_in, attn_type=attn_type, norm_type=norm_type))
+
+ down = torch.nn.Module()
+ down.block = block
+ down.attn = attn
+ if i_level != num_resolutions - 1:
+ down.downsample = Downsample(block_in, resamp_with_conv, causality_axis=causality_axis)
+ curr_res = curr_res // 2
+ down_modules.append(down)
+
+ return down_modules, block_in
+
+
+class Upsample(torch.nn.Module):
+ def __init__(
+ self,
+ in_channels: int,
+ with_conv: bool,
+ causality_axis: CausalityAxis = CausalityAxis.HEIGHT,
+ ) -> None:
+ super().__init__()
+ self.with_conv = with_conv
+ self.causality_axis = causality_axis
+ if self.with_conv:
+ self.conv = make_conv2d(in_channels, in_channels, kernel_size=3, stride=1, causality_axis=causality_axis)
+
+ def forward(self, x: torch.Tensor) -> torch.Tensor:
+ x = torch.nn.functional.interpolate(x, scale_factor=2.0, mode="nearest")
+ if self.with_conv:
+ x = self.conv(x)
+ # Drop FIRST element in the causal axis to undo encoder's padding, while keeping the length 1 + 2 * n.
+ # For example, if the input is [0, 1, 2], after interpolation, the output is [0, 0, 1, 1, 2, 2].
+ # The causal convolution will pad the first element as [-, -, 0, 0, 1, 1, 2, 2],
+ # So the output elements rely on the following windows:
+ # 0: [-,-,0]
+ # 1: [-,0,0]
+ # 2: [0,0,1]
+ # 3: [0,1,1]
+ # 4: [1,1,2]
+ # 5: [1,2,2]
+ # Notice that the first and second elements in the output rely only on the first element in the input,
+ # while all other elements rely on two elements in the input.
+ # So we can drop the first element to undo the padding (rather than the last element).
+ # This is a no-op for non-causal convolutions.
+ match self.causality_axis:
+ case CausalityAxis.NONE:
+ pass # x remains unchanged
+ case CausalityAxis.HEIGHT:
+ x = x[:, :, 1:, :]
+ case CausalityAxis.WIDTH:
+ x = x[:, :, :, 1:]
+ case CausalityAxis.WIDTH_COMPATIBILITY:
+ pass # x remains unchanged
+ case _:
+ raise ValueError(f"Invalid causality_axis: {self.causality_axis}")
+
+ return x
+
+
+def build_upsampling_path( # noqa: PLR0913
+ *,
+ ch: int,
+ ch_mult: Tuple[int, ...],
+ num_resolutions: int,
+ num_res_blocks: int,
+ resolution: int,
+ temb_channels: int,
+ dropout: float,
+ norm_type: NormType,
+ causality_axis: CausalityAxis,
+ attn_type: AttentionType,
+ attn_resolutions: Set[int],
+ resamp_with_conv: bool,
+ initial_block_channels: int,
+) -> tuple[torch.nn.ModuleList, int]:
+ """Build the upsampling path with residual blocks, attention, and upsampling layers."""
+ up_modules = torch.nn.ModuleList()
+ block_in = initial_block_channels
+ curr_res = resolution // (2 ** (num_resolutions - 1))
+
+ for level in reversed(range(num_resolutions)):
+ stage = torch.nn.Module()
+ stage.block = torch.nn.ModuleList()
+ stage.attn = torch.nn.ModuleList()
+ block_out = ch * ch_mult[level]
+
+ for _ in range(num_res_blocks + 1):
+ stage.block.append(
+ ResnetBlock(
+ in_channels=block_in,
+ out_channels=block_out,
+ temb_channels=temb_channels,
+ dropout=dropout,
+ norm_type=norm_type,
+ causality_axis=causality_axis,
+ )
+ )
+ block_in = block_out
+ if curr_res in attn_resolutions:
+ stage.attn.append(make_attn(block_in, attn_type=attn_type, norm_type=norm_type))
+
+ if level != 0:
+ stage.upsample = Upsample(block_in, resamp_with_conv, causality_axis=causality_axis)
+ curr_res *= 2
+
+ up_modules.insert(0, stage)
+
+ return up_modules, block_in
+
+
+class PerChannelStatistics(nn.Module):
+ """
+ Per-channel statistics for normalizing and denormalizing the latent representation.
+ This statics is computed over the entire dataset and stored in model's checkpoint under AudioVAE state_dict.
+ """
+
+ def __init__(self, latent_channels: int = 128) -> None:
+ super().__init__()
+ self.register_buffer("std-of-means", torch.empty(latent_channels))
+ self.register_buffer("mean-of-means", torch.empty(latent_channels))
+
+ def un_normalize(self, x: torch.Tensor) -> torch.Tensor:
+ return (x * self.get_buffer("std-of-means").to(x)) + self.get_buffer("mean-of-means").to(x)
+
+ def normalize(self, x: torch.Tensor) -> torch.Tensor:
+ return (x - self.get_buffer("mean-of-means").to(x)) / self.get_buffer("std-of-means").to(x)
+
+
+LATENT_DOWNSAMPLE_FACTOR = 4
+
+
+def build_mid_block(
+ channels: int,
+ temb_channels: int,
+ dropout: float,
+ norm_type: NormType,
+ causality_axis: CausalityAxis,
+ attn_type: AttentionType,
+ add_attention: bool,
+) -> torch.nn.Module:
+ """Build the middle block with two ResNet blocks and optional attention."""
+ mid = torch.nn.Module()
+ mid.block_1 = ResnetBlock(
+ in_channels=channels,
+ out_channels=channels,
+ temb_channels=temb_channels,
+ dropout=dropout,
+ norm_type=norm_type,
+ causality_axis=causality_axis,
+ )
+ mid.attn_1 = make_attn(channels, attn_type=attn_type, norm_type=norm_type) if add_attention else torch.nn.Identity()
+ mid.block_2 = ResnetBlock(
+ in_channels=channels,
+ out_channels=channels,
+ temb_channels=temb_channels,
+ dropout=dropout,
+ norm_type=norm_type,
+ causality_axis=causality_axis,
+ )
+ return mid
+
+
+def run_mid_block(mid: torch.nn.Module, features: torch.Tensor) -> torch.Tensor:
+ """Run features through the middle block."""
+ features = mid.block_1(features, temb=None)
+ features = mid.attn_1(features)
+ return mid.block_2(features, temb=None)
+
+
+class LTX2AudioEncoder(torch.nn.Module):
+ """
+ Encoder that compresses audio spectrograms into latent representations.
+ The encoder uses a series of downsampling blocks with residual connections,
+ attention mechanisms, and configurable causal convolutions.
+ """
+
+ def __init__( # noqa: PLR0913
+ self,
+ *,
+ ch: int = 128,
+ ch_mult: Tuple[int, ...] = (1, 2, 4),
+ num_res_blocks: int = 2,
+ attn_resolutions: Set[int] = set(),
+ dropout: float = 0.0,
+ resamp_with_conv: bool = True,
+ in_channels: int = 2,
+ resolution: int = 256,
+ z_channels: int = 8,
+ double_z: bool = True,
+ attn_type: AttentionType = AttentionType.VANILLA,
+ mid_block_add_attention: bool = False,
+ norm_type: NormType = NormType.PIXEL,
+ causality_axis: CausalityAxis = CausalityAxis.HEIGHT,
+ sample_rate: int = 16000,
+ mel_hop_length: int = 160,
+ n_fft: int = 1024,
+ is_causal: bool = True,
+ mel_bins: int = 64,
+ **_ignore_kwargs,
+ ) -> None:
+ """
+ Initialize the Encoder.
+ Args:
+ Arguments are configuration parameters, loaded from the audio VAE checkpoint config
+ (audio_vae.model.params.ddconfig):
+ ch: Base number of feature channels used in the first convolution layer.
+ ch_mult: Multiplicative factors for the number of channels at each resolution level.
+ num_res_blocks: Number of residual blocks to use at each resolution level.
+ attn_resolutions: Spatial resolutions (e.g., in time/frequency) at which to apply attention.
+ resolution: Input spatial resolution of the spectrogram (height, width).
+ z_channels: Number of channels in the latent representation.
+ norm_type: Normalization layer type to use within the network (e.g., group, batch).
+ causality_axis: Axis along which convolutions should be causal (e.g., time axis).
+ sample_rate: Audio sample rate in Hz for the input signals.
+ mel_hop_length: Hop length used when computing the mel spectrogram.
+ n_fft: FFT size used to compute the spectrogram.
+ mel_bins: Number of mel-frequency bins in the input spectrogram.
+ in_channels: Number of channels in the input spectrogram tensor.
+ double_z: If True, predict both mean and log-variance (doubling latent channels).
+ is_causal: If True, use causal convolutions suitable for streaming setups.
+ dropout: Dropout probability used in residual and mid blocks.
+ attn_type: Type of attention mechanism to use in attention blocks.
+ resamp_with_conv: If True, perform resolution changes using strided convolutions.
+ mid_block_add_attention: If True, add an attention block in the mid-level of the encoder.
+ """
+ super().__init__()
+
+ self.per_channel_statistics = PerChannelStatistics(latent_channels=ch)
+ self.sample_rate = sample_rate
+ self.mel_hop_length = mel_hop_length
+ self.n_fft = n_fft
+ self.is_causal = is_causal
+ self.mel_bins = mel_bins
+
+ self.patchifier = AudioPatchifier(
+ patch_size=1,
+ audio_latent_downsample_factor=LATENT_DOWNSAMPLE_FACTOR,
+ sample_rate=sample_rate,
+ hop_length=mel_hop_length,
+ is_causal=is_causal,
+ )
+
+ self.ch = ch
+ self.temb_ch = 0
+ self.num_resolutions = len(ch_mult)
+ self.num_res_blocks = num_res_blocks
+ self.resolution = resolution
+ self.in_channels = in_channels
+ self.z_channels = z_channels
+ self.double_z = double_z
+ self.norm_type = norm_type
+ self.causality_axis = causality_axis
+ self.attn_type = attn_type
+
+ # downsampling
+ self.conv_in = make_conv2d(
+ in_channels,
+ self.ch,
+ kernel_size=3,
+ stride=1,
+ causality_axis=self.causality_axis,
+ )
+
+ self.non_linearity = torch.nn.SiLU()
+
+ self.down, block_in = build_downsampling_path(
+ ch=ch,
+ ch_mult=ch_mult,
+ num_resolutions=self.num_resolutions,
+ num_res_blocks=num_res_blocks,
+ resolution=resolution,
+ temb_channels=self.temb_ch,
+ dropout=dropout,
+ norm_type=self.norm_type,
+ causality_axis=self.causality_axis,
+ attn_type=self.attn_type,
+ attn_resolutions=attn_resolutions,
+ resamp_with_conv=resamp_with_conv,
+ )
+
+ self.mid = build_mid_block(
+ channels=block_in,
+ temb_channels=self.temb_ch,
+ dropout=dropout,
+ norm_type=self.norm_type,
+ causality_axis=self.causality_axis,
+ attn_type=self.attn_type,
+ add_attention=mid_block_add_attention,
+ )
+
+ self.norm_out = build_normalization_layer(block_in, normtype=self.norm_type)
+ self.conv_out = make_conv2d(
+ block_in,
+ 2 * z_channels if double_z else z_channels,
+ kernel_size=3,
+ stride=1,
+ causality_axis=self.causality_axis,
+ )
+
+ def forward(self, spectrogram: torch.Tensor) -> torch.Tensor:
+ """
+ Encode audio spectrogram into latent representations.
+ Args:
+ spectrogram: Input spectrogram of shape (batch, channels, time, frequency)
+ Returns:
+ Encoded latent representation of shape (batch, channels, frames, mel_bins)
+ """
+ h = self.conv_in(spectrogram)
+ h = self._run_downsampling_path(h)
+ h = run_mid_block(self.mid, h)
+ h = self._finalize_output(h)
+
+ return self._normalize_latents(h)
+
+ def _run_downsampling_path(self, h: torch.Tensor) -> torch.Tensor:
+ for level in range(self.num_resolutions):
+ stage = self.down[level]
+ for block_idx in range(self.num_res_blocks):
+ h = stage.block[block_idx](h, temb=None)
+ if stage.attn:
+ h = stage.attn[block_idx](h)
+
+ if level != self.num_resolutions - 1:
+ h = stage.downsample(h)
+
+ return h
+
+ def _finalize_output(self, h: torch.Tensor) -> torch.Tensor:
+ h = self.norm_out(h)
+ h = self.non_linearity(h)
+ return self.conv_out(h)
+
+ def _normalize_latents(self, latent_output: torch.Tensor) -> torch.Tensor:
+ """
+ Normalize encoder latents using per-channel statistics.
+ When the encoder is configured with ``double_z=True``, the final
+ convolution produces twice the number of latent channels, typically
+ interpreted as two concatenated tensors along the channel dimension
+ (e.g., mean and variance or other auxiliary parameters).
+ This method intentionally uses only the first half of the channels
+ (the "mean" component) as input to the patchifier and normalization
+ logic. The remaining channels are left unchanged by this method and
+ are expected to be consumed elsewhere in the VAE pipeline.
+ If ``double_z=False``, the encoder output already contains only the
+ mean latents and the chunking operation simply returns that tensor.
+ """
+ means = torch.chunk(latent_output, 2, dim=1)[0]
+ latent_shape = AudioLatentShape(
+ batch=means.shape[0],
+ channels=means.shape[1],
+ frames=means.shape[2],
+ mel_bins=means.shape[3],
+ )
+ latent_patched = self.patchifier.patchify(means)
+ latent_normalized = self.per_channel_statistics.normalize(latent_patched)
+ return self.patchifier.unpatchify(latent_normalized, latent_shape)
+
+
+class LTX2AudioDecoder(torch.nn.Module):
+ """
+ Symmetric decoder that reconstructs audio spectrograms from latent features.
+ The decoder mirrors the encoder structure with configurable channel multipliers,
+ attention resolutions, and causal convolutions.
+ """
+
+ def __init__( # noqa: PLR0913
+ self,
+ *,
+ ch: int = 128,
+ out_ch: int = 2,
+ ch_mult: Tuple[int, ...] = (1, 2, 4),
+ num_res_blocks: int = 2,
+ attn_resolutions: Set[int] = set(),
+ resolution: int=256,
+ z_channels: int=8,
+ norm_type: NormType = NormType.PIXEL,
+ causality_axis: CausalityAxis = CausalityAxis.HEIGHT,
+ dropout: float = 0.0,
+ mid_block_add_attention: bool = False,
+ sample_rate: int = 16000,
+ mel_hop_length: int = 160,
+ is_causal: bool = True,
+ mel_bins: int | None = 64,
+ ) -> None:
+ """
+ Initialize the Decoder.
+ Args:
+ Arguments are configuration parameters, loaded from the audio VAE checkpoint config
+ (audio_vae.model.params.ddconfig):
+ - ch, out_ch, ch_mult, num_res_blocks, attn_resolutions
+ - resolution, z_channels
+ - norm_type, causality_axis
+ """
+ super().__init__()
+
+ # Internal behavioural defaults that are not driven by the checkpoint.
+ resamp_with_conv = True
+ attn_type = AttentionType.VANILLA
+
+ # Per-channel statistics for denormalizing latents
+ self.per_channel_statistics = PerChannelStatistics(latent_channels=ch)
+ self.sample_rate = sample_rate
+ self.mel_hop_length = mel_hop_length
+ self.is_causal = is_causal
+ self.mel_bins = mel_bins
+ self.patchifier = AudioPatchifier(
+ patch_size=1,
+ audio_latent_downsample_factor=LATENT_DOWNSAMPLE_FACTOR,
+ sample_rate=sample_rate,
+ hop_length=mel_hop_length,
+ is_causal=is_causal,
+ )
+
+ self.ch = ch
+ self.temb_ch = 0
+ self.num_resolutions = len(ch_mult)
+ self.num_res_blocks = num_res_blocks
+ self.resolution = resolution
+ self.out_ch = out_ch
+ self.give_pre_end = False
+ self.tanh_out = False
+ self.norm_type = norm_type
+ self.z_channels = z_channels
+ self.channel_multipliers = ch_mult
+ self.attn_resolutions = attn_resolutions
+ self.causality_axis = causality_axis
+ self.attn_type = attn_type
+
+ base_block_channels = ch * self.channel_multipliers[-1]
+ base_resolution = resolution // (2 ** (self.num_resolutions - 1))
+ self.z_shape = (1, z_channels, base_resolution, base_resolution)
+
+ self.conv_in = make_conv2d(
+ z_channels, base_block_channels, kernel_size=3, stride=1, causality_axis=self.causality_axis
+ )
+ self.non_linearity = torch.nn.SiLU()
+ self.mid = build_mid_block(
+ channels=base_block_channels,
+ temb_channels=self.temb_ch,
+ dropout=dropout,
+ norm_type=self.norm_type,
+ causality_axis=self.causality_axis,
+ attn_type=self.attn_type,
+ add_attention=mid_block_add_attention,
+ )
+ self.up, final_block_channels = build_upsampling_path(
+ ch=ch,
+ ch_mult=ch_mult,
+ num_resolutions=self.num_resolutions,
+ num_res_blocks=num_res_blocks,
+ resolution=resolution,
+ temb_channels=self.temb_ch,
+ dropout=dropout,
+ norm_type=self.norm_type,
+ causality_axis=self.causality_axis,
+ attn_type=self.attn_type,
+ attn_resolutions=attn_resolutions,
+ resamp_with_conv=resamp_with_conv,
+ initial_block_channels=base_block_channels,
+ )
+
+ self.norm_out = build_normalization_layer(final_block_channels, normtype=self.norm_type)
+ self.conv_out = make_conv2d(
+ final_block_channels, out_ch, kernel_size=3, stride=1, causality_axis=self.causality_axis
+ )
+
+ def forward(self, sample: torch.Tensor) -> torch.Tensor:
+ """
+ Decode latent features back to audio spectrograms.
+ Args:
+ sample: Encoded latent representation of shape (batch, channels, frames, mel_bins)
+ Returns:
+ Reconstructed audio spectrogram of shape (batch, channels, time, frequency)
+ """
+ sample, target_shape = self._denormalize_latents(sample)
+
+ h = self.conv_in(sample)
+ h = run_mid_block(self.mid, h)
+ h = self._run_upsampling_path(h)
+ h = self._finalize_output(h)
+
+ return self._adjust_output_shape(h, target_shape)
+
+ def _denormalize_latents(self, sample: torch.Tensor) -> tuple[torch.Tensor, AudioLatentShape]:
+ latent_shape = AudioLatentShape(
+ batch=sample.shape[0],
+ channels=sample.shape[1],
+ frames=sample.shape[2],
+ mel_bins=sample.shape[3],
+ )
+
+ sample_patched = self.patchifier.patchify(sample)
+ sample_denormalized = self.per_channel_statistics.un_normalize(sample_patched)
+ sample = self.patchifier.unpatchify(sample_denormalized, latent_shape)
+
+ target_frames = latent_shape.frames * LATENT_DOWNSAMPLE_FACTOR
+ if self.causality_axis != CausalityAxis.NONE:
+ target_frames = max(target_frames - (LATENT_DOWNSAMPLE_FACTOR - 1), 1)
+
+ target_shape = AudioLatentShape(
+ batch=latent_shape.batch,
+ channels=self.out_ch,
+ frames=target_frames,
+ mel_bins=self.mel_bins if self.mel_bins is not None else latent_shape.mel_bins,
+ )
+
+ return sample, target_shape
+
+ def _adjust_output_shape(
+ self,
+ decoded_output: torch.Tensor,
+ target_shape: AudioLatentShape,
+ ) -> torch.Tensor:
+ """
+ Adjust output shape to match target dimensions for variable-length audio.
+ This function handles the common case where decoded audio spectrograms need to be
+ resized to match a specific target shape.
+ Args:
+ decoded_output: Tensor of shape (batch, channels, time, frequency)
+ target_shape: AudioLatentShape describing (batch, channels, time, mel bins)
+ Returns:
+ Tensor adjusted to match target_shape exactly
+ """
+ # Current output shape: (batch, channels, time, frequency)
+ _, _, current_time, current_freq = decoded_output.shape
+ target_channels = target_shape.channels
+ target_time = target_shape.frames
+ target_freq = target_shape.mel_bins
+
+ # Step 1: Crop first to avoid exceeding target dimensions
+ decoded_output = decoded_output[
+ :, :target_channels, : min(current_time, target_time), : min(current_freq, target_freq)
+ ]
+
+ # Step 2: Calculate padding needed for time and frequency dimensions
+ time_padding_needed = target_time - decoded_output.shape[2]
+ freq_padding_needed = target_freq - decoded_output.shape[3]
+
+ # Step 3: Apply padding if needed
+ if time_padding_needed > 0 or freq_padding_needed > 0:
+ # PyTorch padding format: (pad_left, pad_right, pad_top, pad_bottom)
+ # For audio: pad_left/right = frequency, pad_top/bottom = time
+ padding = (
+ 0,
+ max(freq_padding_needed, 0), # frequency padding (left, right)
+ 0,
+ max(time_padding_needed, 0), # time padding (top, bottom)
+ )
+ decoded_output = F.pad(decoded_output, padding)
+
+ # Step 4: Final safety crop to ensure exact target shape
+ decoded_output = decoded_output[:, :target_channels, :target_time, :target_freq]
+
+ return decoded_output
+
+ def _run_upsampling_path(self, h: torch.Tensor) -> torch.Tensor:
+ for level in reversed(range(self.num_resolutions)):
+ stage = self.up[level]
+ for block_idx, block in enumerate(stage.block):
+ h = block(h, temb=None)
+ if stage.attn:
+ h = stage.attn[block_idx](h)
+
+ if level != 0 and hasattr(stage, "upsample"):
+ h = stage.upsample(h)
+
+ return h
+
+ def _finalize_output(self, h: torch.Tensor) -> torch.Tensor:
+ if self.give_pre_end:
+ return h
+
+ h = self.norm_out(h)
+ h = self.non_linearity(h)
+ h = self.conv_out(h)
+ return torch.tanh(h) if self.tanh_out else h
+
+
+def get_padding(kernel_size: int, dilation: int = 1) -> int:
+ return int((kernel_size * dilation - dilation) / 2)
+
+
+# ---------------------------------------------------------------------------
+# Anti-aliased resampling helpers (kaiser-sinc filters) for BigVGAN v2
+# Adopted from https://github.com/NVIDIA/BigVGAN
+# ---------------------------------------------------------------------------
+
+
+def _sinc(x: torch.Tensor) -> torch.Tensor:
+ return torch.where(
+ x == 0,
+ torch.tensor(1.0, device=x.device, dtype=x.dtype),
+ torch.sin(math.pi * x) / math.pi / x,
+ )
+
+
+def kaiser_sinc_filter1d(cutoff: float, half_width: float, kernel_size: int) -> torch.Tensor:
+ even = kernel_size % 2 == 0
+ half_size = kernel_size // 2
+ delta_f = 4 * half_width
+ amplitude = 2.285 * (half_size - 1) * math.pi * delta_f + 7.95
+ if amplitude > 50.0:
+ beta = 0.1102 * (amplitude - 8.7)
+ elif amplitude >= 21.0:
+ beta = 0.5842 * (amplitude - 21) ** 0.4 + 0.07886 * (amplitude - 21.0)
+ else:
+ beta = 0.0
+ window = torch.kaiser_window(kernel_size, beta=beta, periodic=False)
+ time = torch.arange(-half_size, half_size) + 0.5 if even else torch.arange(kernel_size) - half_size
+ if cutoff == 0:
+ filter_ = torch.zeros_like(time)
+ else:
+ filter_ = 2 * cutoff * window * _sinc(2 * cutoff * time)
+ filter_ /= filter_.sum()
+ return filter_.view(1, 1, kernel_size)
+
+
+class LowPassFilter1d(nn.Module):
+ def __init__(
+ self,
+ cutoff: float = 0.5,
+ half_width: float = 0.6,
+ stride: int = 1,
+ padding: bool = True,
+ padding_mode: str = "replicate",
+ kernel_size: int = 12,
+ ) -> None:
+ super().__init__()
+ if cutoff < -0.0:
+ raise ValueError("Minimum cutoff must be larger than zero.")
+ if cutoff > 0.5:
+ raise ValueError("A cutoff above 0.5 does not make sense.")
+ self.kernel_size = kernel_size
+ self.even = kernel_size % 2 == 0
+ self.pad_left = kernel_size // 2 - int(self.even)
+ self.pad_right = kernel_size // 2
+ self.stride = stride
+ self.padding = padding
+ self.padding_mode = padding_mode
+ self.register_buffer("filter", kaiser_sinc_filter1d(cutoff, half_width, kernel_size))
+
+ def forward(self, x: torch.Tensor) -> torch.Tensor:
+ _, n_channels, _ = x.shape
+ if self.padding:
+ x = F.pad(x, (self.pad_left, self.pad_right), mode=self.padding_mode)
+ return F.conv1d(x, self.filter.expand(n_channels, -1, -1), stride=self.stride, groups=n_channels)
+
+
+class UpSample1d(nn.Module):
+ def __init__(
+ self,
+ ratio: int = 2,
+ kernel_size: int | None = None,
+ persistent: bool = True,
+ window_type: str = "kaiser",
+ ) -> None:
+ super().__init__()
+ self.ratio = ratio
+ self.stride = ratio
+
+ if window_type == "hann":
+ # Hann-windowed sinc filter equivalent to torchaudio.functional.resample
+ rolloff = 0.99
+ lowpass_filter_width = 6
+ width = math.ceil(lowpass_filter_width / rolloff)
+ self.kernel_size = 2 * width * ratio + 1
+ self.pad = width
+ self.pad_left = 2 * width * ratio
+ self.pad_right = self.kernel_size - ratio
+ time_axis = (torch.arange(self.kernel_size) / ratio - width) * rolloff
+ time_clamped = time_axis.clamp(-lowpass_filter_width, lowpass_filter_width)
+ window = torch.cos(time_clamped * math.pi / lowpass_filter_width / 2) ** 2
+ sinc_filter = (torch.sinc(time_axis) * window * rolloff / ratio).view(1, 1, -1)
+ else:
+ # Kaiser-windowed sinc filter (BigVGAN default).
+ self.kernel_size = int(6 * ratio // 2) * 2 if kernel_size is None else kernel_size
+ self.pad = self.kernel_size // ratio - 1
+ self.pad_left = self.pad * self.stride + (self.kernel_size - self.stride) // 2
+ self.pad_right = self.pad * self.stride + (self.kernel_size - self.stride + 1) // 2
+ sinc_filter = kaiser_sinc_filter1d(
+ cutoff=0.5 / ratio,
+ half_width=0.6 / ratio,
+ kernel_size=self.kernel_size,
+ )
+
+ self.register_buffer("filter", sinc_filter, persistent=persistent)
+
+ def forward(self, x: torch.Tensor) -> torch.Tensor:
+ _, n_channels, _ = x.shape
+ x = F.pad(x, (self.pad, self.pad), mode="replicate")
+ filt = self.filter.to(dtype=x.dtype, device=x.device).expand(n_channels, -1, -1)
+ x = self.ratio * F.conv_transpose1d(x, filt, stride=self.stride, groups=n_channels)
+ return x[..., self.pad_left : -self.pad_right]
+
+
+class DownSample1d(nn.Module):
+ def __init__(self, ratio: int = 2, kernel_size: int | None = None) -> None:
+ super().__init__()
+ self.ratio = ratio
+ self.kernel_size = int(6 * ratio // 2) * 2 if kernel_size is None else kernel_size
+ self.lowpass = LowPassFilter1d(
+ cutoff=0.5 / ratio,
+ half_width=0.6 / ratio,
+ stride=ratio,
+ kernel_size=self.kernel_size,
+ )
+
+ def forward(self, x: torch.Tensor) -> torch.Tensor:
+ return self.lowpass(x)
+
+
+class Activation1d(nn.Module):
+ def __init__(
+ self,
+ activation: nn.Module,
+ up_ratio: int = 2,
+ down_ratio: int = 2,
+ up_kernel_size: int = 12,
+ down_kernel_size: int = 12,
+ ) -> None:
+ super().__init__()
+ self.act = activation
+ self.upsample = UpSample1d(up_ratio, up_kernel_size)
+ self.downsample = DownSample1d(down_ratio, down_kernel_size)
+
+ def forward(self, x: torch.Tensor) -> torch.Tensor:
+ x = self.upsample(x)
+ x = self.act(x)
+ return self.downsample(x)
+
+
+class Snake(nn.Module):
+ def __init__(
+ self,
+ in_features: int,
+ alpha: float = 1.0,
+ alpha_trainable: bool = True,
+ alpha_logscale: bool = True,
+ ) -> None:
+ super().__init__()
+ self.alpha_logscale = alpha_logscale
+ self.alpha = nn.Parameter(torch.zeros(in_features) if alpha_logscale else torch.ones(in_features) * alpha)
+ self.alpha.requires_grad = alpha_trainable
+ self.eps = 1e-9
+
+ def forward(self, x: torch.Tensor) -> torch.Tensor:
+ alpha = self.alpha.unsqueeze(0).unsqueeze(-1)
+ if self.alpha_logscale:
+ alpha = torch.exp(alpha)
+ return x + (1.0 / (alpha + self.eps)) * torch.sin(x * alpha).pow(2)
+
+
+class SnakeBeta(nn.Module):
+ def __init__(
+ self,
+ in_features: int,
+ alpha: float = 1.0,
+ alpha_trainable: bool = True,
+ alpha_logscale: bool = True,
+ ) -> None:
+ super().__init__()
+ self.alpha_logscale = alpha_logscale
+ self.alpha = nn.Parameter(torch.zeros(in_features) if alpha_logscale else torch.ones(in_features) * alpha)
+ self.alpha.requires_grad = alpha_trainable
+ self.beta = nn.Parameter(torch.zeros(in_features) if alpha_logscale else torch.ones(in_features) * alpha)
+ self.beta.requires_grad = alpha_trainable
+ self.eps = 1e-9
+
+ def forward(self, x: torch.Tensor) -> torch.Tensor:
+ alpha = self.alpha.unsqueeze(0).unsqueeze(-1)
+ beta = self.beta.unsqueeze(0).unsqueeze(-1)
+ if self.alpha_logscale:
+ alpha = torch.exp(alpha)
+ beta = torch.exp(beta)
+ return x + (1.0 / (beta + self.eps)) * torch.sin(x * alpha).pow(2)
+
+
+class AMPBlock1(nn.Module):
+ def __init__(
+ self,
+ channels: int,
+ kernel_size: int = 3,
+ dilation: tuple[int, int, int] = (1, 3, 5),
+ activation: str = "snake",
+ ) -> None:
+ super().__init__()
+ act_cls = SnakeBeta if activation == "snakebeta" else Snake
+ self.convs1 = nn.ModuleList(
+ [
+ nn.Conv1d(
+ channels,
+ channels,
+ kernel_size,
+ 1,
+ dilation=dilation[0],
+ padding=get_padding(kernel_size, dilation[0]),
+ ),
+ nn.Conv1d(
+ channels,
+ channels,
+ kernel_size,
+ 1,
+ dilation=dilation[1],
+ padding=get_padding(kernel_size, dilation[1]),
+ ),
+ nn.Conv1d(
+ channels,
+ channels,
+ kernel_size,
+ 1,
+ dilation=dilation[2],
+ padding=get_padding(kernel_size, dilation[2]),
+ ),
+ ]
+ )
+
+ self.convs2 = nn.ModuleList(
+ [
+ nn.Conv1d(channels, channels, kernel_size, 1, dilation=1, padding=get_padding(kernel_size, 1)),
+ nn.Conv1d(channels, channels, kernel_size, 1, dilation=1, padding=get_padding(kernel_size, 1)),
+ nn.Conv1d(channels, channels, kernel_size, 1, dilation=1, padding=get_padding(kernel_size, 1)),
+ ]
+ )
+
+ self.acts1 = nn.ModuleList([Activation1d(act_cls(channels)) for _ in range(len(self.convs1))])
+ self.acts2 = nn.ModuleList([Activation1d(act_cls(channels)) for _ in range(len(self.convs2))])
+
+ def forward(self, x: torch.Tensor) -> torch.Tensor:
+ for c1, c2, a1, a2 in zip(self.convs1, self.convs2, self.acts1, self.acts2, strict=True):
+ xt = a1(x)
+ xt = c1(xt)
+ xt = a2(xt)
+ xt = c2(xt)
+ x = x + xt
+ return x
+
+
+class LTX2Vocoder(torch.nn.Module):
+ """
+ LTX2Vocoder model for synthesizing audio from Mel spectrograms.
+ Args:
+ resblock_kernel_sizes: List of kernel sizes for the residual blocks.
+ This value is read from the checkpoint at `config.vocoder.resblock_kernel_sizes`.
+ upsample_rates: List of upsampling rates.
+ This value is read from the checkpoint at `config.vocoder.upsample_rates`.
+ upsample_kernel_sizes: List of kernel sizes for the upsampling layers.
+ This value is read from the checkpoint at `config.vocoder.upsample_kernel_sizes`.
+ resblock_dilation_sizes: List of dilation sizes for the residual blocks.
+ This value is read from the checkpoint at `config.vocoder.resblock_dilation_sizes`.
+ upsample_initial_channel: Initial number of channels for the upsampling layers.
+ This value is read from the checkpoint at `config.vocoder.upsample_initial_channel`.
+ resblock: Type of residual block to use ("1", "2", or "AMP1").
+ This value is read from the checkpoint at `config.vocoder.resblock`.
+ output_sampling_rate: Waveform sample rate.
+ This value is read from the checkpoint at `config.vocoder.output_sampling_rate`.
+ activation: Activation type for BigVGAN v2 ("snake" or "snakebeta"). Only used when resblock="AMP1".
+ use_tanh_at_final: Apply tanh at the output (when apply_final_activation=True).
+ apply_final_activation: Whether to apply the final tanh/clamp activation.
+ use_bias_at_final: Whether to use bias in the final conv layer.
+ """
+
+ def __init__( # noqa: PLR0913
+ self,
+ resblock_kernel_sizes: List[int] | None = [3, 7, 11],
+ upsample_rates: List[int] | None = [6, 5, 2, 2, 2],
+ upsample_kernel_sizes: List[int] | None = [16, 15, 8, 4, 4],
+ resblock_dilation_sizes: List[List[int]] | None = [[1, 3, 5], [1, 3, 5], [1, 3, 5]],
+ upsample_initial_channel: int = 1024,
+ resblock: str = "1",
+ output_sampling_rate: int = 24000,
+ activation: str = "snake",
+ use_tanh_at_final: bool = True,
+ apply_final_activation: bool = True,
+ use_bias_at_final: bool = True,
+ ) -> None:
+ super().__init__()
+
+ # Mutable default values are not supported as default arguments.
+ if resblock_kernel_sizes is None:
+ resblock_kernel_sizes = [3, 7, 11]
+ if upsample_rates is None:
+ upsample_rates = [6, 5, 2, 2, 2]
+ if upsample_kernel_sizes is None:
+ upsample_kernel_sizes = [16, 15, 8, 4, 4]
+ if resblock_dilation_sizes is None:
+ resblock_dilation_sizes = [[1, 3, 5], [1, 3, 5], [1, 3, 5]]
+
+ self.output_sampling_rate = output_sampling_rate
+ self.num_kernels = len(resblock_kernel_sizes)
+ self.num_upsamples = len(upsample_rates)
+ self.use_tanh_at_final = use_tanh_at_final
+ self.apply_final_activation = apply_final_activation
+ self.is_amp = resblock == "AMP1"
+
+ # All production checkpoints are stereo: 128 input channels (2 stereo channels x 64 mel
+ # bins each), 2 output channels.
+ self.conv_pre = nn.Conv1d(
+ in_channels=128,
+ out_channels=upsample_initial_channel,
+ kernel_size=7,
+ stride=1,
+ padding=3,
+ )
+ resblock_cls = ResBlock1 if resblock == "1" else AMPBlock1
+
+ self.ups = nn.ModuleList(
+ nn.ConvTranspose1d(
+ upsample_initial_channel // (2**i),
+ upsample_initial_channel // (2 ** (i + 1)),
+ kernel_size,
+ stride,
+ padding=(kernel_size - stride) // 2,
+ )
+ for i, (stride, kernel_size) in enumerate(zip(upsample_rates, upsample_kernel_sizes, strict=True))
+ )
+
+ final_channels = upsample_initial_channel // (2 ** len(upsample_rates))
+ self.resblocks = nn.ModuleList()
+
+ for i in range(len(upsample_rates)):
+ ch = upsample_initial_channel // (2 ** (i + 1))
+ for kernel_size, dilations in zip(resblock_kernel_sizes, resblock_dilation_sizes, strict=True):
+ if self.is_amp:
+ self.resblocks.append(resblock_cls(ch, kernel_size, dilations, activation=activation))
+ else:
+ self.resblocks.append(resblock_cls(ch, kernel_size, dilations))
+
+ if self.is_amp:
+ self.act_post: nn.Module = Activation1d(SnakeBeta(final_channels))
+ else:
+ self.act_post = nn.LeakyReLU()
+
+ # All production checkpoints are stereo: this final conv maps `final_channels` to 2 output channels (stereo).
+ self.conv_post = nn.Conv1d(
+ in_channels=final_channels,
+ out_channels=2,
+ kernel_size=7,
+ stride=1,
+ padding=3,
+ bias=use_bias_at_final,
+ )
+
+ def forward(self, x: torch.Tensor) -> torch.Tensor:
+ """
+ Forward pass of the vocoder.
+ Args:
+ x: Input Mel spectrogram tensor. Can be either:
+ - 3D: (batch_size, time, mel_bins) for mono
+ - 4D: (batch_size, 2, time, mel_bins) for stereo
+ Returns:
+ Audio waveform tensor of shape (batch_size, out_channels, audio_length)
+ """
+ x = x.transpose(2, 3) # (batch, channels, time, mel_bins) -> (batch, channels, mel_bins, time)
+
+ if x.dim() == 4: # stereo
+ assert x.shape[1] == 2, "Input must have 2 channels for stereo"
+ x = einops.rearrange(x, "b s c t -> b (s c) t")
+
+ x = self.conv_pre(x)
+
+ for i in range(self.num_upsamples):
+ if not self.is_amp:
+ x = F.leaky_relu(x, LRELU_SLOPE)
+ x = self.ups[i](x)
+ start = i * self.num_kernels
+ end = start + self.num_kernels
+
+ # Evaluate all resblocks with the same input tensor so they can run
+ # independently (and thus in parallel on accelerator hardware) before
+ # aggregating their outputs via mean.
+ block_outputs = torch.stack(
+ [self.resblocks[idx](x) for idx in range(start, end)],
+ dim=0,
+ )
+ x = block_outputs.mean(dim=0)
+
+ x = self.act_post(x)
+ x = self.conv_post(x)
+
+ if self.apply_final_activation:
+ x = torch.tanh(x) if self.use_tanh_at_final else torch.clamp(x, -1, 1)
+
+ return x
+
+
+class _STFTFn(nn.Module):
+ """Implements STFT as a convolution with precomputed DFT x Hann-window bases.
+ The DFT basis rows (real and imaginary parts interleaved) multiplied by the causal
+ Hann window are stored as buffers and loaded from the checkpoint. Using the exact
+ bfloat16 bases from training ensures the mel values fed to the BWE generator are
+ bit-identical to what it was trained on.
+ """
+
+ def __init__(self, filter_length: int, hop_length: int, win_length: int) -> None:
+ super().__init__()
+ self.hop_length = hop_length
+ self.win_length = win_length
+ n_freqs = filter_length // 2 + 1
+ self.register_buffer("forward_basis", torch.zeros(n_freqs * 2, 1, filter_length))
+ self.register_buffer("inverse_basis", torch.zeros(n_freqs * 2, 1, filter_length))
+
+ def forward(self, y: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
+ """Compute magnitude and phase spectrogram from a batch of waveforms.
+ Applies causal (left-only) padding of win_length - hop_length samples so that
+ each output frame depends only on past and present input — no lookahead.
+ Args:
+ y: Waveform tensor of shape (B, T).
+ Returns:
+ magnitude: Linear amplitude spectrogram, shape (B, n_freqs, T_frames).
+ phase: Phase spectrogram in radians, shape (B, n_freqs, T_frames).
+ """
+ if y.dim() == 2:
+ y = y.unsqueeze(1) # (B, 1, T)
+ left_pad = max(0, self.win_length - self.hop_length) # causal: left-only
+ y = F.pad(y, (left_pad, 0))
+ spec = F.conv1d(y, self.forward_basis, stride=self.hop_length, padding=0)
+ n_freqs = spec.shape[1] // 2
+ real, imag = spec[:, :n_freqs], spec[:, n_freqs:]
+ magnitude = torch.sqrt(real**2 + imag**2)
+ phase = torch.atan2(imag.float(), real.float()).to(real.dtype)
+ return magnitude, phase
+
+
+class MelSTFT(nn.Module):
+ """Causal log-mel spectrogram module whose buffers are loaded from the checkpoint.
+ Computes a log-mel spectrogram by running the causal STFT (_STFTFn) on the input
+ waveform and projecting the linear magnitude spectrum onto the mel filterbank.
+ The module's state dict layout matches the 'mel_stft.*' keys stored in the checkpoint
+ (mel_basis, stft_fn.forward_basis, stft_fn.inverse_basis).
+ """
+
+ def __init__(
+ self,
+ filter_length: int,
+ hop_length: int,
+ win_length: int,
+ n_mel_channels: int,
+ ) -> None:
+ super().__init__()
+ self.stft_fn = _STFTFn(filter_length, hop_length, win_length)
+
+ # Initialized to zeros; load_state_dict overwrites with the checkpoint's
+ # exact bfloat16 filterbank (vocoder.mel_stft.mel_basis, shape [n_mels, n_freqs]).
+ n_freqs = filter_length // 2 + 1
+ self.register_buffer("mel_basis", torch.zeros(n_mel_channels, n_freqs))
+
+ def mel_spectrogram(self, y: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
+ """Compute log-mel spectrogram and auxiliary spectral quantities.
+ Args:
+ y: Waveform tensor of shape (B, T).
+ Returns:
+ log_mel: Log-compressed mel spectrogram, shape (B, n_mel_channels, T_frames).
+ magnitude: Linear amplitude spectrogram, shape (B, n_freqs, T_frames).
+ phase: Phase spectrogram in radians, shape (B, n_freqs, T_frames).
+ energy: Per-frame energy (L2 norm over frequency), shape (B, T_frames).
+ """
+ magnitude, phase = self.stft_fn(y)
+ energy = torch.norm(magnitude, dim=1)
+ mel = torch.matmul(self.mel_basis.to(magnitude.dtype), magnitude)
+ log_mel = torch.log(torch.clamp(mel, min=1e-5))
+ return log_mel, magnitude, phase, energy
+
+
+class LTX2VocoderWithBWE(nn.Module):
+ """LTX2Vocoder with bandwidth extension (BWE) upsampling.
+ Chains a mel-to-wav vocoder with a BWE module that upsamples the output
+ to a higher sample rate. The BWE computes a mel spectrogram from the
+ vocoder output, runs it through a second generator to predict a residual,
+ and adds it to a sinc-resampled skip connection.
+ """
+
+ def __init__(
+ self,
+ input_sampling_rate: int = 16000,
+ output_sampling_rate: int = 48000,
+ hop_length: int = 80,
+ ) -> None:
+ super().__init__()
+ self.vocoder = LTX2Vocoder(
+ resblock_kernel_sizes=[3, 7, 11],
+ upsample_rates=[5, 2, 2, 2, 2, 2],
+ upsample_kernel_sizes=[11, 4, 4, 4, 4, 4],
+ resblock_dilation_sizes=[[1, 3, 5], [1, 3, 5], [1, 3, 5]],
+ upsample_initial_channel=1536,
+ resblock="AMP1",
+ activation="snakebeta",
+ use_tanh_at_final=False,
+ apply_final_activation=True,
+ use_bias_at_final=False,
+ output_sampling_rate=input_sampling_rate,
+ )
+ self.bwe_generator = LTX2Vocoder(
+ resblock_kernel_sizes=[3, 7, 11],
+ upsample_rates=[6, 5, 2, 2, 2],
+ upsample_kernel_sizes=[12, 11, 4, 4, 4],
+ resblock_dilation_sizes=[[1, 3, 5], [1, 3, 5], [1, 3, 5]],
+ upsample_initial_channel=512,
+ resblock="AMP1",
+ activation="snakebeta",
+ use_tanh_at_final=False,
+ apply_final_activation=False,
+ use_bias_at_final=False,
+ output_sampling_rate=output_sampling_rate,
+ )
+
+ self.mel_stft = MelSTFT(
+ filter_length=512,
+ hop_length=hop_length,
+ win_length=512,
+ n_mel_channels=64,
+ )
+ self.input_sampling_rate = input_sampling_rate
+ self.output_sampling_rate = output_sampling_rate
+ self.hop_length = hop_length
+ # Compute the resampler on CPU so the sinc filter is materialized even when
+ # the model is constructed on meta device (SingleGPUModelBuilder pattern).
+ # The filter is not stored in the checkpoint (persistent=False).
+ with torch.device("cpu"):
+ self.resampler = UpSample1d(
+ ratio=output_sampling_rate // input_sampling_rate, persistent=False, window_type="hann"
+ )
+
+ @property
+ def conv_pre(self) -> nn.Conv1d:
+ return self.vocoder.conv_pre
+
+ @property
+ def conv_post(self) -> nn.Conv1d:
+ return self.vocoder.conv_post
+
+ def _compute_mel(self, audio: torch.Tensor) -> torch.Tensor:
+ """Compute log-mel spectrogram from waveform using causal STFT bases.
+ Args:
+ audio: Waveform tensor of shape (B, C, T).
+ Returns:
+ mel: Log-mel spectrogram of shape (B, C, n_mels, T_frames).
+ """
+ batch, n_channels, _ = audio.shape
+ flat = audio.reshape(batch * n_channels, -1) # (B*C, T)
+ mel, _, _, _ = self.mel_stft.mel_spectrogram(flat) # (B*C, n_mels, T_frames)
+ return mel.reshape(batch, n_channels, mel.shape[1], mel.shape[2]) # (B, C, n_mels, T_frames)
+
+ def forward(self, mel_spec: torch.Tensor) -> torch.Tensor:
+ """Run the full vocoder + BWE forward pass.
+ Args:
+ mel_spec: Mel spectrogram of shape (B, 2, T, mel_bins) for stereo
+ or (B, T, mel_bins) for mono. Same format as LTX2Vocoder.forward.
+ Returns:
+ Waveform tensor of shape (B, out_channels, T_out) clipped to [-1, 1].
+ """
+ x = self.vocoder(mel_spec)
+ _, _, length_low_rate = x.shape
+ output_length = length_low_rate * self.output_sampling_rate // self.input_sampling_rate
+
+ # Pad to multiple of hop_length for exact mel frame count
+ remainder = length_low_rate % self.hop_length
+ if remainder != 0:
+ x = F.pad(x, (0, self.hop_length - remainder))
+
+ # Compute mel spectrogram from vocoder output: (B, C, n_mels, T_frames)
+ mel = self._compute_mel(x)
+
+ # LTX2Vocoder.forward expects (B, C, T, mel_bins) — transpose before calling bwe_generator
+ mel_for_bwe = mel.transpose(2, 3) # (B, C, T_frames, mel_bins)
+ residual = self.bwe_generator(mel_for_bwe)
+ skip = self.resampler(x)
+ assert residual.shape == skip.shape, f"residual {residual.shape} != skip {skip.shape}"
+
+ return torch.clamp(residual + skip, -1, 1)[..., :output_length]
diff --git a/diffsynth/models/ltx2_common.py b/diffsynth/models/ltx2_common.py
new file mode 100644
index 0000000000000000000000000000000000000000..a658ec6ecb76bbf7372a0e0893ec52b2a8307ddd
--- /dev/null
+++ b/diffsynth/models/ltx2_common.py
@@ -0,0 +1,388 @@
+from dataclasses import dataclass
+from typing import NamedTuple, Protocol, Tuple
+import torch
+from torch import nn
+from enum import Enum
+
+
+class VideoPixelShape(NamedTuple):
+ """
+ Shape of the tensor representing the video pixel array. Assumes BGR channel format.
+ """
+
+ batch: int
+ frames: int
+ height: int
+ width: int
+ fps: float
+
+
+class SpatioTemporalScaleFactors(NamedTuple):
+ """
+ Describes the spatiotemporal downscaling between decoded video space and
+ the corresponding VAE latent grid.
+ """
+
+ time: int
+ width: int
+ height: int
+
+ @classmethod
+ def default(cls) -> "SpatioTemporalScaleFactors":
+ return cls(time=8, width=32, height=32)
+
+
+VIDEO_SCALE_FACTORS = SpatioTemporalScaleFactors.default()
+
+
+class VideoLatentShape(NamedTuple):
+ """
+ Shape of the tensor representing video in VAE latent space.
+ The latent representation is a 5D tensor with dimensions ordered as
+ (batch, channels, frames, height, width). Spatial and temporal dimensions
+ are downscaled relative to pixel space according to the VAE's scale factors.
+ """
+
+ batch: int
+ channels: int
+ frames: int
+ height: int
+ width: int
+
+ def to_torch_shape(self) -> torch.Size:
+ return torch.Size([self.batch, self.channels, self.frames, self.height, self.width])
+
+ @staticmethod
+ def from_torch_shape(shape: torch.Size) -> "VideoLatentShape":
+ return VideoLatentShape(
+ batch=shape[0],
+ channels=shape[1],
+ frames=shape[2],
+ height=shape[3],
+ width=shape[4],
+ )
+
+ def mask_shape(self) -> "VideoLatentShape":
+ return self._replace(channels=1)
+
+ @staticmethod
+ def from_pixel_shape(
+ shape: VideoPixelShape,
+ latent_channels: int = 128,
+ scale_factors: SpatioTemporalScaleFactors = VIDEO_SCALE_FACTORS,
+ ) -> "VideoLatentShape":
+ frames = (shape.frames - 1) // scale_factors[0] + 1
+ height = shape.height // scale_factors[1]
+ width = shape.width // scale_factors[2]
+
+ return VideoLatentShape(
+ batch=shape.batch,
+ channels=latent_channels,
+ frames=frames,
+ height=height,
+ width=width,
+ )
+
+ def upscale(self, scale_factors: SpatioTemporalScaleFactors = VIDEO_SCALE_FACTORS) -> "VideoLatentShape":
+ return self._replace(
+ channels=3,
+ frames=(self.frames - 1) * scale_factors.time + 1,
+ height=self.height * scale_factors.height,
+ width=self.width * scale_factors.width,
+ )
+
+
+class AudioLatentShape(NamedTuple):
+ """
+ Shape of audio in VAE latent space: (batch, channels, frames, mel_bins).
+ mel_bins is the number of frequency bins from the mel-spectrogram encoding.
+ """
+
+ batch: int
+ channels: int
+ frames: int
+ mel_bins: int
+
+ def to_torch_shape(self) -> torch.Size:
+ return torch.Size([self.batch, self.channels, self.frames, self.mel_bins])
+
+ def mask_shape(self) -> "AudioLatentShape":
+ return self._replace(channels=1, mel_bins=1)
+
+ @staticmethod
+ def from_torch_shape(shape: torch.Size) -> "AudioLatentShape":
+ return AudioLatentShape(
+ batch=shape[0],
+ channels=shape[1],
+ frames=shape[2],
+ mel_bins=shape[3],
+ )
+
+ @staticmethod
+ def from_duration(
+ batch: int,
+ duration: float,
+ channels: int = 8,
+ mel_bins: int = 16,
+ sample_rate: int = 16000,
+ hop_length: int = 160,
+ audio_latent_downsample_factor: int = 4,
+ ) -> "AudioLatentShape":
+ latents_per_second = float(sample_rate) / float(hop_length) / float(audio_latent_downsample_factor)
+
+ return AudioLatentShape(
+ batch=batch,
+ channels=channels,
+ frames=round(duration * latents_per_second),
+ mel_bins=mel_bins,
+ )
+
+ @staticmethod
+ def from_video_pixel_shape(
+ shape: VideoPixelShape,
+ channels: int = 8,
+ mel_bins: int = 16,
+ sample_rate: int = 16000,
+ hop_length: int = 160,
+ audio_latent_downsample_factor: int = 4,
+ ) -> "AudioLatentShape":
+ return AudioLatentShape.from_duration(
+ batch=shape.batch,
+ duration=float(shape.frames) / float(shape.fps),
+ channels=channels,
+ mel_bins=mel_bins,
+ sample_rate=sample_rate,
+ hop_length=hop_length,
+ audio_latent_downsample_factor=audio_latent_downsample_factor,
+ )
+
+
+@dataclass(frozen=True)
+class LatentState:
+ """
+ State of latents during the diffusion denoising process.
+ Attributes:
+ latent: The current noisy latent tensor being denoised.
+ denoise_mask: Mask encoding the denoising strength for each token (1 = full denoising, 0 = no denoising).
+ positions: Positional indices for each latent element, used for positional embeddings.
+ clean_latent: Initial state of the latent before denoising, may include conditioning latents.
+ """
+
+ latent: torch.Tensor
+ denoise_mask: torch.Tensor
+ positions: torch.Tensor
+ clean_latent: torch.Tensor
+
+ def clone(self) -> "LatentState":
+ return LatentState(
+ latent=self.latent.clone(),
+ denoise_mask=self.denoise_mask.clone(),
+ positions=self.positions.clone(),
+ clean_latent=self.clean_latent.clone(),
+ )
+
+
+class NormType(Enum):
+ """Normalization layer types: GROUP (GroupNorm) or PIXEL (per-location RMS norm)."""
+
+ GROUP = "group"
+ PIXEL = "pixel"
+
+
+class PixelNorm(nn.Module):
+ """
+ Per-pixel (per-location) RMS normalization layer.
+ For each element along the chosen dimension, this layer normalizes the tensor
+ by the root-mean-square of its values across that dimension:
+ y = x / sqrt(mean(x^2, dim=dim, keepdim=True) + eps)
+ """
+
+ def __init__(self, dim: int = 1, eps: float = 1e-8) -> None:
+ """
+ Args:
+ dim: Dimension along which to compute the RMS (typically channels).
+ eps: Small constant added for numerical stability.
+ """
+ super().__init__()
+ self.dim = dim
+ self.eps = eps
+
+ def forward(self, x: torch.Tensor) -> torch.Tensor:
+ """
+ Apply RMS normalization along the configured dimension.
+ """
+ # Compute mean of squared values along `dim`, keep dimensions for broadcasting.
+ mean_sq = torch.mean(x**2, dim=self.dim, keepdim=True)
+ # Normalize by the root-mean-square (RMS).
+ rms = torch.sqrt(mean_sq + self.eps)
+ return x / rms
+
+
+def build_normalization_layer(
+ in_channels: int, *, num_groups: int = 32, normtype: NormType = NormType.GROUP
+) -> nn.Module:
+ """
+ Create a normalization layer based on the normalization type.
+ Args:
+ in_channels: Number of input channels
+ num_groups: Number of groups for group normalization
+ normtype: Type of normalization: "group" or "pixel"
+ Returns:
+ A normalization layer
+ """
+ if normtype == NormType.GROUP:
+ return torch.nn.GroupNorm(num_groups=num_groups, num_channels=in_channels, eps=1e-6, affine=True)
+ if normtype == NormType.PIXEL:
+ return PixelNorm(dim=1, eps=1e-6)
+ raise ValueError(f"Invalid normalization type: {normtype}")
+
+
+def rms_norm(x: torch.Tensor, weight: torch.Tensor | None = None, eps: float = 1e-6) -> torch.Tensor:
+ """Root-mean-square (RMS) normalize `x` over its last dimension.
+ Thin wrapper around `torch.nn.functional.rms_norm` that infers the normalized
+ shape and forwards `weight` and `eps`.
+ """
+ return torch.nn.functional.rms_norm(x, (x.shape[-1],), weight=weight, eps=eps)
+
+
+@dataclass(frozen=True)
+class Modality:
+ """
+ Input data for a single modality (video or audio) in the transformer.
+ Bundles the latent tokens, timestep embeddings, positional information,
+ and text conditioning context for processing by the diffusion transformer.
+ Attributes:
+ latent: Patchified latent tokens, shape ``(B, T, D)`` where *B* is
+ the batch size, *T* is the total number of tokens (noisy +
+ conditioning), and *D* is the input dimension.
+ timesteps: Per-token timestep embeddings, shape ``(B, T)``.
+ positions: Positional coordinates, shape ``(B, 3, T)`` for video
+ (time, height, width) or ``(B, 1, T)`` for audio.
+ context: Text conditioning embeddings from the prompt encoder.
+ enabled: Whether this modality is active in the current forward pass.
+ context_mask: Optional mask for the text context tokens.
+ attention_mask: Optional 2-D self-attention mask, shape ``(B, T, T)``.
+ Values in ``[0, 1]`` where ``1`` = full attention and ``0`` = no
+ attention. ``None`` means unrestricted (full) attention between
+ all tokens. Built incrementally by conditioning items; see
+ :class:`~ltx_core.conditioning.types.attention_strength_wrapper.ConditioningItemAttentionStrengthWrapper`.
+ """
+
+ latent: (
+ torch.Tensor
+ ) # Shape: (B, T, D) where B is the batch size, T is the number of tokens, and D is input dimension
+ sigma: torch.Tensor # Shape: (B,). Current sigma value, used for cross-attention timestep calculation.
+ timesteps: torch.Tensor # Shape: (B, T) where T is the number of timesteps
+ positions: (
+ torch.Tensor
+ ) # Shape: (B, 3, T) for video, where 3 is the number of dimensions and T is the number of tokens
+ context: torch.Tensor
+ enabled: bool = True
+ context_mask: torch.Tensor | None = None
+ attention_mask: torch.Tensor | None = None
+
+
+def to_denoised(
+ sample: torch.Tensor,
+ velocity: torch.Tensor,
+ sigma: float | torch.Tensor,
+ calc_dtype: torch.dtype = torch.float32,
+) -> torch.Tensor:
+ """
+ Convert the sample and its denoising velocity to denoised sample.
+ Returns:
+ Denoised sample
+ """
+ if isinstance(sigma, torch.Tensor):
+ sigma = sigma.to(calc_dtype)
+ return (sample.to(calc_dtype) - velocity.to(calc_dtype) * sigma).to(sample.dtype)
+
+
+
+class Patchifier(Protocol):
+ """
+ Protocol for patchifiers that convert latent tensors into patches and assemble them back.
+ """
+
+ def patchify(
+ self,
+ latents: torch.Tensor,
+ ) -> torch.Tensor:
+ ...
+ """
+ Convert latent tensors into flattened patch tokens.
+ Args:
+ latents: Latent tensor to patchify.
+ Returns:
+ Flattened patch tokens tensor.
+ """
+
+ def unpatchify(
+ self,
+ latents: torch.Tensor,
+ output_shape: AudioLatentShape | VideoLatentShape,
+ ) -> torch.Tensor:
+ """
+ Converts latent tensors between spatio-temporal formats and flattened sequence representations.
+ Args:
+ latents: Patch tokens that must be rearranged back into the latent grid constructed by `patchify`.
+ output_shape: Shape of the output tensor. Note that output_shape is either AudioLatentShape or
+ VideoLatentShape.
+ Returns:
+ Dense latent tensor restored from the flattened representation.
+ """
+
+ @property
+ def patch_size(self) -> Tuple[int, int, int]:
+ ...
+ """
+ Returns the patch size as a tuple of (temporal, height, width) dimensions
+ """
+
+ def get_patch_grid_bounds(
+ self,
+ output_shape: AudioLatentShape | VideoLatentShape,
+ device: torch.device | None = None,
+ ) -> torch.Tensor:
+ ...
+ """
+ Compute metadata describing where each latent patch resides within the
+ grid specified by `output_shape`.
+ Args:
+ output_shape: Target grid layout for the patches.
+ device: Target device for the returned tensor.
+ Returns:
+ Tensor containing patch coordinate metadata such as spatial or temporal intervals.
+ """
+
+
+def get_pixel_coords(
+ latent_coords: torch.Tensor,
+ scale_factors: SpatioTemporalScaleFactors,
+ causal_fix: bool = False,
+) -> torch.Tensor:
+ """
+ Map latent-space `[start, end)` coordinates to their pixel-space equivalents by scaling
+ each axis (frame/time, height, width) with the corresponding VAE downsampling factors.
+ Optionally compensate for causal encoding that keeps the first frame at unit temporal scale.
+ Args:
+ latent_coords: Tensor of latent bounds shaped `(batch, 3, num_patches, 2)`.
+ scale_factors: SpatioTemporalScaleFactors tuple `(temporal, height, width)` with integer scale factors applied
+ per axis.
+ causal_fix: When True, rewrites the temporal axis of the first frame so causal VAEs
+ that treat frame zero differently still yield non-negative timestamps.
+ """
+ # Broadcast the VAE scale factors so they align with the `(batch, axis, patch, bound)` layout.
+ broadcast_shape = [1] * latent_coords.ndim
+ broadcast_shape[1] = -1 # axis dimension corresponds to (frame/time, height, width)
+ scale_tensor = torch.tensor(scale_factors, device=latent_coords.device).view(*broadcast_shape)
+
+ # Apply per-axis scaling to convert latent bounds into pixel-space coordinates.
+ pixel_coords = latent_coords * scale_tensor
+
+ if causal_fix:
+ # VAE temporal stride for the very first frame is 1 instead of `scale_factors[0]`.
+ # Shift and clamp to keep the first-frame timestamps causal and non-negative.
+ pixel_coords[:, 0, ...] = (pixel_coords[:, 0, ...] + 1 - scale_factors[0]).clamp(min=0)
+
+ return pixel_coords
diff --git a/diffsynth/models/ltx2_dit.py b/diffsynth/models/ltx2_dit.py
new file mode 100644
index 0000000000000000000000000000000000000000..9df0ed3a78b5ee074386c67d2f2bc931fd6ed19b
--- /dev/null
+++ b/diffsynth/models/ltx2_dit.py
@@ -0,0 +1,1683 @@
+import math
+import functools
+from dataclasses import dataclass, replace
+from enum import Enum
+from typing import Optional, Tuple, Callable
+import numpy as np
+import torch
+from einops import rearrange
+from .ltx2_common import rms_norm, Modality
+from ..core.attention.attention import attention_forward
+from ..core import gradient_checkpoint_forward
+
+
+def get_timestep_embedding(
+ timesteps: torch.Tensor,
+ embedding_dim: int,
+ flip_sin_to_cos: bool = False,
+ downscale_freq_shift: float = 1,
+ scale: float = 1,
+ max_period: int = 10000,
+) -> torch.Tensor:
+ """
+ This matches the implementation in Denoising Diffusion Probabilistic Models: Create sinusoidal timestep embeddings.
+ Args
+ timesteps (torch.Tensor):
+ a 1-D Tensor of N indices, one per batch element. These may be fractional.
+ embedding_dim (int):
+ the dimension of the output.
+ flip_sin_to_cos (bool):
+ Whether the embedding order should be `cos, sin` (if True) or `sin, cos` (if False)
+ downscale_freq_shift (float):
+ Controls the delta between frequencies between dimensions
+ scale (float):
+ Scaling factor applied to the embeddings.
+ max_period (int):
+ Controls the maximum frequency of the embeddings
+ Returns
+ torch.Tensor: an [N x dim] Tensor of positional embeddings.
+ """
+ assert len(timesteps.shape) == 1, "Timesteps should be a 1d-array"
+
+ half_dim = embedding_dim // 2
+ exponent = -math.log(max_period) * torch.arange(start=0, end=half_dim, dtype=torch.float32, device=timesteps.device)
+ exponent = exponent / (half_dim - downscale_freq_shift)
+
+ emb = torch.exp(exponent)
+ emb = timesteps[:, None].float() * emb[None, :]
+
+ # scale embeddings
+ emb = scale * emb
+
+ # concat sine and cosine embeddings
+ emb = torch.cat([torch.sin(emb), torch.cos(emb)], dim=-1)
+
+ # flip sine and cosine embeddings
+ if flip_sin_to_cos:
+ emb = torch.cat([emb[:, half_dim:], emb[:, :half_dim]], dim=-1)
+
+ # zero pad
+ if embedding_dim % 2 == 1:
+ emb = torch.nn.functional.pad(emb, (0, 1, 0, 0))
+ return emb
+
+
+class TimestepEmbedding(torch.nn.Module):
+ def __init__(
+ self,
+ in_channels: int,
+ time_embed_dim: int,
+ out_dim: int | None = None,
+ post_act_fn: str | None = None,
+ cond_proj_dim: int | None = None,
+ sample_proj_bias: bool = True,
+ ):
+ super().__init__()
+
+ self.linear_1 = torch.nn.Linear(in_channels, time_embed_dim, sample_proj_bias)
+
+ if cond_proj_dim is not None:
+ self.cond_proj = torch.nn.Linear(cond_proj_dim, in_channels, bias=False)
+ else:
+ self.cond_proj = None
+
+ self.act = torch.nn.SiLU()
+ time_embed_dim_out = out_dim if out_dim is not None else time_embed_dim
+
+ self.linear_2 = torch.nn.Linear(time_embed_dim, time_embed_dim_out, sample_proj_bias)
+
+ if post_act_fn is None:
+ self.post_act = None
+
+ def forward(self, sample: torch.Tensor, condition: torch.Tensor | None = None) -> torch.Tensor:
+ if condition is not None:
+ sample = sample + self.cond_proj(condition)
+ sample = self.linear_1(sample)
+
+ if self.act is not None:
+ sample = self.act(sample)
+
+ sample = self.linear_2(sample)
+
+ if self.post_act is not None:
+ sample = self.post_act(sample)
+ return sample
+
+
+class Timesteps(torch.nn.Module):
+ def __init__(self, num_channels: int, flip_sin_to_cos: bool, downscale_freq_shift: float, scale: int = 1):
+ super().__init__()
+ self.num_channels = num_channels
+ self.flip_sin_to_cos = flip_sin_to_cos
+ self.downscale_freq_shift = downscale_freq_shift
+ self.scale = scale
+
+ def forward(self, timesteps: torch.Tensor) -> torch.Tensor:
+ t_emb = get_timestep_embedding(
+ timesteps,
+ self.num_channels,
+ flip_sin_to_cos=self.flip_sin_to_cos,
+ downscale_freq_shift=self.downscale_freq_shift,
+ scale=self.scale,
+ )
+ return t_emb
+
+
+class PixArtAlphaCombinedTimestepSizeEmbeddings(torch.nn.Module):
+ """
+ For PixArt-Alpha.
+ Reference:
+ https://github.com/PixArt-alpha/PixArt-alpha/blob/0f55e922376d8b797edd44d25d0e7464b260dcab/diffusion/model/nets/PixArtMS.py#L164C9-L168C29
+ """
+
+ def __init__(
+ self,
+ embedding_dim: int,
+ size_emb_dim: int,
+ ):
+ super().__init__()
+
+ self.outdim = size_emb_dim
+ self.time_proj = Timesteps(num_channels=256, flip_sin_to_cos=True, downscale_freq_shift=0)
+ self.timestep_embedder = TimestepEmbedding(in_channels=256, time_embed_dim=embedding_dim)
+
+ def forward(
+ self,
+ timestep: torch.Tensor,
+ hidden_dtype: torch.dtype,
+ ) -> torch.Tensor:
+ timesteps_proj = self.time_proj(timestep)
+ timesteps_emb = self.timestep_embedder(timesteps_proj.to(dtype=hidden_dtype)) # (N, D)
+ return timesteps_emb
+
+
+class PerturbationType(Enum):
+ """Types of attention perturbations for STG (Spatio-Temporal Guidance)."""
+
+ SKIP_A2V_CROSS_ATTN = "skip_a2v_cross_attn"
+ SKIP_V2A_CROSS_ATTN = "skip_v2a_cross_attn"
+ SKIP_VIDEO_SELF_ATTN = "skip_video_self_attn"
+ SKIP_AUDIO_SELF_ATTN = "skip_audio_self_attn"
+
+
+@dataclass(frozen=True)
+class Perturbation:
+ """A single perturbation specifying which attention type to skip and in which blocks."""
+
+ type: PerturbationType
+ blocks: list[int] | None # None means all blocks
+
+ def is_perturbed(self, perturbation_type: PerturbationType, block: int) -> bool:
+ if self.type != perturbation_type:
+ return False
+
+ if self.blocks is None:
+ return True
+
+ return block in self.blocks
+
+
+@dataclass(frozen=True)
+class PerturbationConfig:
+ """Configuration holding a list of perturbations for a single sample."""
+
+ perturbations: list[Perturbation] | None
+
+ def is_perturbed(self, perturbation_type: PerturbationType, block: int) -> bool:
+ if self.perturbations is None:
+ return False
+
+ return any(perturbation.is_perturbed(perturbation_type, block) for perturbation in self.perturbations)
+
+ @staticmethod
+ def empty() -> "PerturbationConfig":
+ return PerturbationConfig([])
+
+
+@dataclass(frozen=True)
+class BatchedPerturbationConfig:
+ """Perturbation configurations for a batch, with utilities for generating attention masks."""
+
+ perturbations: list[PerturbationConfig]
+
+ def mask(
+ self, perturbation_type: PerturbationType, block: int, device, dtype: torch.dtype
+ ) -> torch.Tensor:
+ mask = torch.ones((len(self.perturbations),), device=device, dtype=dtype)
+ for batch_idx, perturbation in enumerate(self.perturbations):
+ if perturbation.is_perturbed(perturbation_type, block):
+ mask[batch_idx] = 0
+
+ return mask
+
+ def mask_like(self, perturbation_type: PerturbationType, block: int, values: torch.Tensor) -> torch.Tensor:
+ mask = self.mask(perturbation_type, block, values.device, values.dtype)
+ return mask.view(mask.numel(), *([1] * len(values.shape[1:])))
+
+ def any_in_batch(self, perturbation_type: PerturbationType, block: int) -> bool:
+ return any(perturbation.is_perturbed(perturbation_type, block) for perturbation in self.perturbations)
+
+ def all_in_batch(self, perturbation_type: PerturbationType, block: int) -> bool:
+ return all(perturbation.is_perturbed(perturbation_type, block) for perturbation in self.perturbations)
+
+ @staticmethod
+ def empty(batch_size: int) -> "BatchedPerturbationConfig":
+ return BatchedPerturbationConfig([PerturbationConfig.empty() for _ in range(batch_size)])
+
+
+
+ADALN_NUM_BASE_PARAMS = 6
+# Cross-attention AdaLN adds 3 more (scale, shift, gate) for the CA norm.
+ADALN_NUM_CROSS_ATTN_PARAMS = 3
+
+
+def adaln_embedding_coefficient(cross_attention_adaln: bool) -> int:
+ """Total number of AdaLN parameters per block."""
+ return ADALN_NUM_BASE_PARAMS + (ADALN_NUM_CROSS_ATTN_PARAMS if cross_attention_adaln else 0)
+
+
+class AdaLayerNormSingle(torch.nn.Module):
+ r"""
+ Norm layer adaptive layer norm single (adaLN-single).
+ As proposed in PixArt-Alpha (see: https://arxiv.org/abs/2310.00426; Section 2.3).
+ Parameters:
+ embedding_dim (`int`): The size of each embedding vector.
+ use_additional_conditions (`bool`): To use additional conditions for normalization or not.
+ """
+
+ def __init__(self, embedding_dim: int, embedding_coefficient: int = 6):
+ super().__init__()
+
+ self.emb = PixArtAlphaCombinedTimestepSizeEmbeddings(
+ embedding_dim,
+ size_emb_dim=embedding_dim // 3,
+ )
+
+ self.silu = torch.nn.SiLU()
+ self.linear = torch.nn.Linear(embedding_dim, embedding_coefficient * embedding_dim, bias=True)
+
+ def forward(
+ self,
+ timestep: torch.Tensor,
+ hidden_dtype: Optional[torch.dtype] = None,
+ ) -> Tuple[torch.Tensor, torch.Tensor]:
+ embedded_timestep = self.emb(timestep, hidden_dtype=hidden_dtype)
+ return self.linear(self.silu(embedded_timestep)), embedded_timestep
+
+
+class LTXRopeType(Enum):
+ INTERLEAVED = "interleaved"
+ SPLIT = "split"
+
+
+def apply_rotary_emb(
+ input_tensor: torch.Tensor,
+ freqs_cis: Tuple[torch.Tensor, torch.Tensor],
+ rope_type: LTXRopeType = LTXRopeType.INTERLEAVED,
+) -> torch.Tensor:
+ if rope_type == LTXRopeType.INTERLEAVED:
+ return apply_interleaved_rotary_emb(input_tensor, *freqs_cis)
+ elif rope_type == LTXRopeType.SPLIT:
+ return apply_split_rotary_emb(input_tensor, *freqs_cis)
+ else:
+ raise ValueError(f"Invalid rope type: {rope_type}")
+
+
+
+def apply_interleaved_rotary_emb(
+ input_tensor: torch.Tensor, cos_freqs: torch.Tensor, sin_freqs: torch.Tensor
+) -> torch.Tensor:
+ t_dup = rearrange(input_tensor, "... (d r) -> ... d r", r=2)
+ t1, t2 = t_dup.unbind(dim=-1)
+ t_dup = torch.stack((-t2, t1), dim=-1)
+ input_tensor_rot = rearrange(t_dup, "... d r -> ... (d r)")
+
+ out = input_tensor * cos_freqs + input_tensor_rot * sin_freqs
+
+ return out
+
+
+def apply_split_rotary_emb(
+ input_tensor: torch.Tensor, cos_freqs: torch.Tensor, sin_freqs: torch.Tensor
+) -> torch.Tensor:
+ needs_reshape = False
+ if input_tensor.ndim != 4 and cos_freqs.ndim == 4:
+ b, h, t, _ = cos_freqs.shape
+ input_tensor = input_tensor.reshape(b, t, h, -1).swapaxes(1, 2)
+ needs_reshape = True
+
+ split_input = rearrange(input_tensor, "... (d r) -> ... d r", d=2)
+ first_half_input = split_input[..., :1, :]
+ second_half_input = split_input[..., 1:, :]
+
+ output = split_input * cos_freqs.unsqueeze(-2)
+ first_half_output = output[..., :1, :]
+ second_half_output = output[..., 1:, :]
+
+ first_half_output.addcmul_(-sin_freqs.unsqueeze(-2), second_half_input)
+ second_half_output.addcmul_(sin_freqs.unsqueeze(-2), first_half_input)
+
+ output = rearrange(output, "... d r -> ... (d r)")
+ if needs_reshape:
+ output = output.swapaxes(1, 2).reshape(b, t, -1)
+
+ return output
+
+
+@functools.lru_cache(maxsize=5)
+def generate_freq_grid_np(
+ positional_embedding_theta: float, positional_embedding_max_pos_count: int, inner_dim: int
+) -> torch.Tensor:
+ theta = positional_embedding_theta
+ start = 1
+ end = theta
+
+ n_elem = 2 * positional_embedding_max_pos_count
+ pow_indices = np.power(
+ theta,
+ np.linspace(
+ np.log(start) / np.log(theta),
+ np.log(end) / np.log(theta),
+ inner_dim // n_elem,
+ dtype=np.float64,
+ ),
+ )
+ return torch.tensor(pow_indices * math.pi / 2, dtype=torch.float32)
+
+
+@functools.lru_cache(maxsize=5)
+def generate_freq_grid_pytorch(
+ positional_embedding_theta: float, positional_embedding_max_pos_count: int, inner_dim: int
+) -> torch.Tensor:
+ theta = positional_embedding_theta
+ start = 1
+ end = theta
+ n_elem = 2 * positional_embedding_max_pos_count
+
+ indices = theta ** (
+ torch.linspace(
+ math.log(start, theta),
+ math.log(end, theta),
+ inner_dim // n_elem,
+ dtype=torch.float32,
+ )
+ )
+ indices = indices.to(dtype=torch.float32)
+
+ indices = indices * math.pi / 2
+
+ return indices
+
+
+def get_fractional_positions(indices_grid: torch.Tensor, max_pos: list[int]) -> torch.Tensor:
+ n_pos_dims = indices_grid.shape[1]
+ assert n_pos_dims == len(max_pos), (
+ f"Number of position dimensions ({n_pos_dims}) must match max_pos length ({len(max_pos)})"
+ )
+ fractional_positions = torch.stack(
+ [indices_grid[:, i] / max_pos[i] for i in range(n_pos_dims)],
+ dim=-1,
+ )
+ return fractional_positions
+
+
+def generate_freqs(
+ indices: torch.Tensor, indices_grid: torch.Tensor, max_pos: list[int], use_middle_indices_grid: bool
+) -> torch.Tensor:
+ if use_middle_indices_grid:
+ assert len(indices_grid.shape) == 4
+ assert indices_grid.shape[-1] == 2
+ indices_grid_start, indices_grid_end = indices_grid[..., 0], indices_grid[..., 1]
+ indices_grid = (indices_grid_start + indices_grid_end) / 2.0
+ elif len(indices_grid.shape) == 4:
+ indices_grid = indices_grid[..., 0]
+
+ fractional_positions = get_fractional_positions(indices_grid, max_pos)
+ indices = indices.to(device=fractional_positions.device)
+
+ freqs = (indices * (fractional_positions.unsqueeze(-1) * 2 - 1)).transpose(-1, -2).flatten(2)
+ return freqs
+
+
+def split_freqs_cis(freqs: torch.Tensor, pad_size: int, num_attention_heads: int) -> tuple[torch.Tensor, torch.Tensor]:
+ cos_freq = freqs.cos()
+ sin_freq = freqs.sin()
+
+ if pad_size != 0:
+ cos_padding = torch.ones_like(cos_freq[:, :, :pad_size])
+ sin_padding = torch.zeros_like(sin_freq[:, :, :pad_size])
+
+ cos_freq = torch.concatenate([cos_padding, cos_freq], axis=-1)
+ sin_freq = torch.concatenate([sin_padding, sin_freq], axis=-1)
+
+ # Reshape freqs to be compatible with multi-head attention
+ b = cos_freq.shape[0]
+ t = cos_freq.shape[1]
+
+ cos_freq = cos_freq.reshape(b, t, num_attention_heads, -1)
+ sin_freq = sin_freq.reshape(b, t, num_attention_heads, -1)
+
+ cos_freq = torch.swapaxes(cos_freq, 1, 2) # (B,H,T,D//2)
+ sin_freq = torch.swapaxes(sin_freq, 1, 2) # (B,H,T,D//2)
+ return cos_freq, sin_freq
+
+
+def interleaved_freqs_cis(freqs: torch.Tensor, pad_size: int) -> tuple[torch.Tensor, torch.Tensor]:
+ cos_freq = freqs.cos().repeat_interleave(2, dim=-1)
+ sin_freq = freqs.sin().repeat_interleave(2, dim=-1)
+ if pad_size != 0:
+ cos_padding = torch.ones_like(cos_freq[:, :, :pad_size])
+ sin_padding = torch.zeros_like(cos_freq[:, :, :pad_size])
+ cos_freq = torch.cat([cos_padding, cos_freq], dim=-1)
+ sin_freq = torch.cat([sin_padding, sin_freq], dim=-1)
+ return cos_freq, sin_freq
+
+
+def precompute_freqs_cis(
+ indices_grid: torch.Tensor,
+ dim: int,
+ out_dtype: torch.dtype,
+ theta: float = 10000.0,
+ max_pos: list[int] | None = None,
+ use_middle_indices_grid: bool = False,
+ num_attention_heads: int = 32,
+ rope_type: LTXRopeType = LTXRopeType.INTERLEAVED,
+ freq_grid_generator: Callable[[float, int, int, torch.device], torch.Tensor] = generate_freq_grid_pytorch,
+) -> tuple[torch.Tensor, torch.Tensor]:
+ if max_pos is None:
+ max_pos = [20, 2048, 2048]
+
+ indices = freq_grid_generator(theta, indices_grid.shape[1], dim)
+ freqs = generate_freqs(indices, indices_grid, max_pos, use_middle_indices_grid)
+
+ if rope_type == LTXRopeType.SPLIT:
+ expected_freqs = dim // 2
+ current_freqs = freqs.shape[-1]
+ pad_size = expected_freqs - current_freqs
+ cos_freq, sin_freq = split_freqs_cis(freqs, pad_size, num_attention_heads)
+ else:
+ # 2 because of cos and sin by 3 for (t, x, y), 1 for temporal only
+ n_elem = 2 * indices_grid.shape[1]
+ cos_freq, sin_freq = interleaved_freqs_cis(freqs, dim % n_elem)
+ return cos_freq.to(out_dtype), sin_freq.to(out_dtype)
+
+
+class Attention(torch.nn.Module):
+ def __init__(
+ self,
+ query_dim: int,
+ context_dim: int | None = None,
+ heads: int = 8,
+ dim_head: int = 64,
+ norm_eps: float = 1e-6,
+ rope_type: LTXRopeType = LTXRopeType.INTERLEAVED,
+ apply_gated_attention: bool = False,
+ ) -> None:
+ super().__init__()
+ self.rope_type = rope_type
+
+ inner_dim = dim_head * heads
+ context_dim = query_dim if context_dim is None else context_dim
+
+ self.heads = heads
+ self.dim_head = dim_head
+
+ self.q_norm = torch.nn.RMSNorm(inner_dim, eps=norm_eps)
+ self.k_norm = torch.nn.RMSNorm(inner_dim, eps=norm_eps)
+
+ self.to_q = torch.nn.Linear(query_dim, inner_dim, bias=True)
+ self.to_k = torch.nn.Linear(context_dim, inner_dim, bias=True)
+ self.to_v = torch.nn.Linear(context_dim, inner_dim, bias=True)
+
+ # Optional per-head gating
+ if apply_gated_attention:
+ self.to_gate_logits = torch.nn.Linear(query_dim, heads, bias=True)
+ else:
+ self.to_gate_logits = None
+
+ self.to_out = torch.nn.Sequential(torch.nn.Linear(inner_dim, query_dim, bias=True), torch.nn.Identity())
+
+ def forward(
+ self,
+ x: torch.Tensor,
+ context: torch.Tensor | None = None,
+ mask: torch.Tensor | None = None,
+ pe: torch.Tensor | None = None,
+ k_pe: torch.Tensor | None = None,
+ perturbation_mask: torch.Tensor | None = None,
+ all_perturbed: bool = False,
+ ) -> torch.Tensor:
+ q = self.to_q(x)
+ context = x if context is None else context
+ k = self.to_k(context)
+ v = self.to_v(context)
+
+ q = self.q_norm(q)
+ k = self.k_norm(k)
+
+ if pe is not None:
+ q = apply_rotary_emb(q, pe, self.rope_type)
+ k = apply_rotary_emb(k, pe if k_pe is None else k_pe, self.rope_type)
+
+ # Reshape for attention_forward using unflatten
+ q = q.unflatten(-1, (self.heads, self.dim_head))
+ k = k.unflatten(-1, (self.heads, self.dim_head))
+ v = v.unflatten(-1, (self.heads, self.dim_head))
+
+ out = attention_forward(
+ q=q,
+ k=k,
+ v=v,
+ q_pattern="b s n d",
+ k_pattern="b s n d",
+ v_pattern="b s n d",
+ out_pattern="b s n d",
+ attn_mask=mask
+ )
+
+ # Reshape back to original format
+ out = out.flatten(2, 3)
+
+ # Apply per-head gating if enabled
+ if self.to_gate_logits is not None:
+ gate_logits = self.to_gate_logits(x) # (B, T, H)
+ b, t, _ = out.shape
+ # Reshape to (B, T, H, D) for per-head gating
+ out = out.view(b, t, self.heads, self.dim_head)
+ # Apply gating: 2 * sigmoid(x) so that zero-init gives identity (2 * 0.5 = 1.0)
+ gates = 2.0 * torch.sigmoid(gate_logits) # (B, T, H)
+ out = out * gates.unsqueeze(-1) # (B, T, H, D) * (B, T, H, 1)
+ # Reshape back to (B, T, H*D)
+ out = out.view(b, t, self.heads * self.dim_head)
+
+ return self.to_out(out)
+
+
+class PixArtAlphaTextProjection(torch.nn.Module):
+ """
+ Projects caption embeddings. Also handles dropout for classifier-free guidance.
+ Adapted from https://github.com/PixArt-alpha/PixArt-alpha/blob/master/diffusion/model/nets/PixArt_blocks.py
+ """
+
+ def __init__(self, in_features: int, hidden_size: int, out_features: int | None = None, act_fn: str = "gelu_tanh"):
+ super().__init__()
+ if out_features is None:
+ out_features = hidden_size
+ self.linear_1 = torch.nn.Linear(in_features=in_features, out_features=hidden_size, bias=True)
+ if act_fn == "gelu_tanh":
+ self.act_1 = torch.nn.GELU(approximate="tanh")
+ elif act_fn == "silu":
+ self.act_1 = torch.nn.SiLU()
+ else:
+ raise ValueError(f"Unknown activation function: {act_fn}")
+ self.linear_2 = torch.nn.Linear(in_features=hidden_size, out_features=out_features, bias=True)
+
+ def forward(self, caption: torch.Tensor) -> torch.Tensor:
+ hidden_states = self.linear_1(caption)
+ hidden_states = self.act_1(hidden_states)
+ hidden_states = self.linear_2(hidden_states)
+ return hidden_states
+
+@dataclass(frozen=True)
+class TransformerArgs:
+ x: torch.Tensor
+ context: torch.Tensor
+ context_mask: torch.Tensor
+ timesteps: torch.Tensor
+ embedded_timestep: torch.Tensor
+ positional_embeddings: torch.Tensor
+ cross_positional_embeddings: torch.Tensor | None
+ cross_scale_shift_timestep: torch.Tensor | None
+ cross_gate_timestep: torch.Tensor | None
+ enabled: bool
+ prompt_timestep: torch.Tensor | None = None
+ self_attention_mask: torch.Tensor | None = (
+ None # Additive log-space self-attention bias (B, 1, T, T), None = full attention
+ )
+
+
+class TransformerArgsPreprocessor:
+ def __init__( # noqa: PLR0913
+ self,
+ patchify_proj: torch.nn.Linear,
+ adaln: AdaLayerNormSingle,
+ inner_dim: int,
+ max_pos: list[int],
+ num_attention_heads: int,
+ use_middle_indices_grid: bool,
+ timestep_scale_multiplier: int,
+ double_precision_rope: bool,
+ positional_embedding_theta: float,
+ rope_type: LTXRopeType,
+ caption_projection: torch.nn.Module | None = None,
+ prompt_adaln: AdaLayerNormSingle | None = None,
+ ) -> None:
+ self.patchify_proj = patchify_proj
+ self.adaln = adaln
+ self.inner_dim = inner_dim
+ self.max_pos = max_pos
+ self.num_attention_heads = num_attention_heads
+ self.use_middle_indices_grid = use_middle_indices_grid
+ self.timestep_scale_multiplier = timestep_scale_multiplier
+ self.double_precision_rope = double_precision_rope
+ self.positional_embedding_theta = positional_embedding_theta
+ self.rope_type = rope_type
+ self.caption_projection = caption_projection
+ self.prompt_adaln = prompt_adaln
+
+ def _prepare_timestep(
+ self, timestep: torch.Tensor, adaln: AdaLayerNormSingle, batch_size: int, hidden_dtype: torch.dtype
+ ) -> tuple[torch.Tensor, torch.Tensor]:
+ """Prepare timestep embeddings."""
+ timestep_scaled = timestep * self.timestep_scale_multiplier
+ timestep, embedded_timestep = adaln(
+ timestep_scaled.flatten(),
+ hidden_dtype=hidden_dtype,
+ )
+ # Second dimension is 1 or number of tokens (if timestep_per_token)
+ timestep = timestep.view(batch_size, -1, timestep.shape[-1])
+ embedded_timestep = embedded_timestep.view(batch_size, -1, embedded_timestep.shape[-1])
+ return timestep, embedded_timestep
+
+ def _prepare_context(
+ self,
+ context: torch.Tensor,
+ x: torch.Tensor,
+ ) -> torch.Tensor:
+ """Prepare context for transformer blocks."""
+ if self.caption_projection is not None:
+ context = self.caption_projection(context)
+ batch_size = x.shape[0]
+ return context.view(batch_size, -1, x.shape[-1])
+
+ def _prepare_attention_mask(self, attention_mask: torch.Tensor | None, x_dtype: torch.dtype) -> torch.Tensor | None:
+ """Prepare attention mask."""
+ if attention_mask is None or torch.is_floating_point(attention_mask):
+ return attention_mask
+
+ return (attention_mask - 1).to(x_dtype).reshape(
+ (attention_mask.shape[0], 1, -1, attention_mask.shape[-1])
+ ) * torch.finfo(x_dtype).max
+
+ def _prepare_self_attention_mask(
+ self, attention_mask: torch.Tensor | None, x_dtype: torch.dtype
+ ) -> torch.Tensor | None:
+ """Prepare self-attention mask by converting [0,1] values to additive log-space bias.
+ Input shape: (B, T, T) with values in [0, 1].
+ Output shape: (B, 1, T, T) with 0.0 for full attention and a large negative value
+ for masked positions.
+ Positions with attention_mask <= 0 are fully masked (mapped to the dtype's minimum
+ representable value). Strictly positive entries are converted via log-space for
+ smooth attenuation, with small values clamped for numerical stability.
+ Returns None if input is None (no masking).
+ """
+ if attention_mask is None:
+ return None
+
+ # Convert [0, 1] attention mask to additive log-space bias:
+ # 1.0 -> log(1.0) = 0.0 (no bias, full attention)
+ # 0.0 -> finfo.min (fully masked)
+ finfo = torch.finfo(x_dtype)
+ eps = finfo.tiny
+
+ bias = torch.full_like(attention_mask, finfo.min, dtype=x_dtype)
+ positive = attention_mask > 0
+ if positive.any():
+ bias[positive] = torch.log(attention_mask[positive].clamp(min=eps)).to(x_dtype)
+
+ return bias.unsqueeze(1) # (B, 1, T, T) for head broadcast
+
+ def _prepare_positional_embeddings(
+ self,
+ positions: torch.Tensor,
+ inner_dim: int,
+ max_pos: list[int],
+ use_middle_indices_grid: bool,
+ num_attention_heads: int,
+ x_dtype: torch.dtype,
+ ) -> torch.Tensor:
+ """Prepare positional embeddings."""
+ freq_grid_generator = generate_freq_grid_np if self.double_precision_rope else generate_freq_grid_pytorch
+ pe = precompute_freqs_cis(
+ positions,
+ dim=inner_dim,
+ out_dtype=x_dtype,
+ theta=self.positional_embedding_theta,
+ max_pos=max_pos,
+ use_middle_indices_grid=use_middle_indices_grid,
+ num_attention_heads=num_attention_heads,
+ rope_type=self.rope_type,
+ freq_grid_generator=freq_grid_generator,
+ )
+ return pe
+
+ def prepare(
+ self,
+ modality: Modality,
+ cross_modality: Modality | None = None, # noqa: ARG002
+ ) -> TransformerArgs:
+ x = self.patchify_proj(modality.latent)
+ batch_size = x.shape[0]
+ timestep, embedded_timestep = self._prepare_timestep(
+ modality.timesteps, self.adaln, batch_size, modality.latent.dtype
+ )
+ prompt_timestep = None
+ if self.prompt_adaln is not None:
+ prompt_timestep, _ = self._prepare_timestep(
+ modality.sigma, self.prompt_adaln, batch_size, modality.latent.dtype
+ )
+ context = self._prepare_context(modality.context, x)
+ attention_mask = self._prepare_attention_mask(modality.context_mask, modality.latent.dtype)
+ pe = self._prepare_positional_embeddings(
+ positions=modality.positions,
+ inner_dim=self.inner_dim,
+ max_pos=self.max_pos,
+ use_middle_indices_grid=self.use_middle_indices_grid,
+ num_attention_heads=self.num_attention_heads,
+ x_dtype=modality.latent.dtype,
+ )
+ self_attention_mask = self._prepare_self_attention_mask(modality.attention_mask, modality.latent.dtype)
+ return TransformerArgs(
+ x=x,
+ context=context,
+ context_mask=attention_mask,
+ timesteps=timestep,
+ embedded_timestep=embedded_timestep,
+ positional_embeddings=pe,
+ cross_positional_embeddings=None,
+ cross_scale_shift_timestep=None,
+ cross_gate_timestep=None,
+ enabled=modality.enabled,
+ prompt_timestep=prompt_timestep,
+ self_attention_mask=self_attention_mask,
+ )
+
+
+class MultiModalTransformerArgsPreprocessor:
+ def __init__( # noqa: PLR0913
+ self,
+ patchify_proj: torch.nn.Linear,
+ adaln: AdaLayerNormSingle,
+ cross_scale_shift_adaln: AdaLayerNormSingle,
+ cross_gate_adaln: AdaLayerNormSingle,
+ inner_dim: int,
+ max_pos: list[int],
+ num_attention_heads: int,
+ cross_pe_max_pos: int,
+ use_middle_indices_grid: bool,
+ audio_cross_attention_dim: int,
+ timestep_scale_multiplier: int,
+ double_precision_rope: bool,
+ positional_embedding_theta: float,
+ rope_type: LTXRopeType,
+ av_ca_timestep_scale_multiplier: int,
+ caption_projection: torch.nn.Module | None = None,
+ prompt_adaln: AdaLayerNormSingle | None = None,
+ ) -> None:
+ self.simple_preprocessor = TransformerArgsPreprocessor(
+ patchify_proj=patchify_proj,
+ adaln=adaln,
+ inner_dim=inner_dim,
+ max_pos=max_pos,
+ num_attention_heads=num_attention_heads,
+ use_middle_indices_grid=use_middle_indices_grid,
+ timestep_scale_multiplier=timestep_scale_multiplier,
+ double_precision_rope=double_precision_rope,
+ positional_embedding_theta=positional_embedding_theta,
+ rope_type=rope_type,
+ caption_projection=caption_projection,
+ prompt_adaln=prompt_adaln,
+ )
+ self.cross_scale_shift_adaln = cross_scale_shift_adaln
+ self.cross_gate_adaln = cross_gate_adaln
+ self.cross_pe_max_pos = cross_pe_max_pos
+ self.audio_cross_attention_dim = audio_cross_attention_dim
+ self.av_ca_timestep_scale_multiplier = av_ca_timestep_scale_multiplier
+
+ def prepare(
+ self,
+ modality: Modality,
+ cross_modality: Modality | None = None,
+ ) -> TransformerArgs:
+ transformer_args = self.simple_preprocessor.prepare(modality)
+ if cross_modality is None:
+ return transformer_args
+
+ if cross_modality.sigma.numel() > 1:
+ if cross_modality.sigma.shape[0] != modality.timesteps.shape[0]:
+ raise ValueError("Cross modality sigma must have the same batch size as the modality")
+ if cross_modality.sigma.ndim != 1:
+ raise ValueError("Cross modality sigma must be a 1D tensor")
+
+ cross_timestep = cross_modality.sigma.view(
+ modality.timesteps.shape[0], 1, *[1] * len(modality.timesteps.shape[2:])
+ )
+
+ cross_pe = self.simple_preprocessor._prepare_positional_embeddings(
+ positions=modality.positions[:, 0:1, :],
+ inner_dim=self.audio_cross_attention_dim,
+ max_pos=[self.cross_pe_max_pos],
+ use_middle_indices_grid=True,
+ num_attention_heads=self.simple_preprocessor.num_attention_heads,
+ x_dtype=modality.latent.dtype,
+ )
+
+ cross_scale_shift_timestep, cross_gate_timestep = self._prepare_cross_attention_timestep(
+ timestep=cross_timestep,
+ timestep_scale_multiplier=self.simple_preprocessor.timestep_scale_multiplier,
+ batch_size=transformer_args.x.shape[0],
+ hidden_dtype=modality.latent.dtype,
+ )
+
+ return replace(
+ transformer_args,
+ cross_positional_embeddings=cross_pe,
+ cross_scale_shift_timestep=cross_scale_shift_timestep,
+ cross_gate_timestep=cross_gate_timestep,
+ )
+
+ def _prepare_cross_attention_timestep(
+ self,
+ timestep: torch.Tensor | None,
+ timestep_scale_multiplier: int,
+ batch_size: int,
+ hidden_dtype: torch.dtype,
+ ) -> tuple[torch.Tensor, torch.Tensor]:
+ """Prepare cross attention timestep embeddings."""
+ timestep = timestep * timestep_scale_multiplier
+
+ av_ca_factor = self.av_ca_timestep_scale_multiplier / timestep_scale_multiplier
+
+ scale_shift_timestep, _ = self.cross_scale_shift_adaln(
+ timestep.flatten(),
+ hidden_dtype=hidden_dtype,
+ )
+ scale_shift_timestep = scale_shift_timestep.view(batch_size, -1, scale_shift_timestep.shape[-1])
+ gate_noise_timestep, _ = self.cross_gate_adaln(
+ timestep.flatten() * av_ca_factor,
+ hidden_dtype=hidden_dtype,
+ )
+ gate_noise_timestep = gate_noise_timestep.view(batch_size, -1, gate_noise_timestep.shape[-1])
+
+ return scale_shift_timestep, gate_noise_timestep
+
+
+@dataclass
+class TransformerConfig:
+ dim: int
+ heads: int
+ d_head: int
+ context_dim: int
+ apply_gated_attention: bool = False
+ cross_attention_adaln: bool = False
+
+
+class BasicAVTransformerBlock(torch.nn.Module):
+ def __init__(
+ self,
+ idx: int,
+ video: TransformerConfig | None = None,
+ audio: TransformerConfig | None = None,
+ rope_type: LTXRopeType = LTXRopeType.INTERLEAVED,
+ norm_eps: float = 1e-6,
+ ):
+ super().__init__()
+
+ self.idx = idx
+ if video is not None:
+ self.attn1 = Attention(
+ query_dim=video.dim,
+ heads=video.heads,
+ dim_head=video.d_head,
+ context_dim=None,
+ rope_type=rope_type,
+ norm_eps=norm_eps,
+ apply_gated_attention=video.apply_gated_attention,
+ )
+ self.attn2 = Attention(
+ query_dim=video.dim,
+ context_dim=video.context_dim,
+ heads=video.heads,
+ dim_head=video.d_head,
+ rope_type=rope_type,
+ norm_eps=norm_eps,
+ apply_gated_attention=video.apply_gated_attention,
+ )
+ self.ff = FeedForward(video.dim, dim_out=video.dim)
+ video_sst_size = adaln_embedding_coefficient(video.cross_attention_adaln)
+ self.scale_shift_table = torch.nn.Parameter(torch.empty(video_sst_size, video.dim))
+
+ if audio is not None:
+ self.audio_attn1 = Attention(
+ query_dim=audio.dim,
+ heads=audio.heads,
+ dim_head=audio.d_head,
+ context_dim=None,
+ rope_type=rope_type,
+ norm_eps=norm_eps,
+ apply_gated_attention=audio.apply_gated_attention,
+ )
+ self.audio_attn2 = Attention(
+ query_dim=audio.dim,
+ context_dim=audio.context_dim,
+ heads=audio.heads,
+ dim_head=audio.d_head,
+ rope_type=rope_type,
+ norm_eps=norm_eps,
+ apply_gated_attention=audio.apply_gated_attention,
+ )
+ self.audio_ff = FeedForward(audio.dim, dim_out=audio.dim)
+ audio_sst_size = adaln_embedding_coefficient(audio.cross_attention_adaln)
+ self.audio_scale_shift_table = torch.nn.Parameter(torch.empty(audio_sst_size, audio.dim))
+
+ if audio is not None and video is not None:
+ # Q: Video, K,V: Audio
+ self.audio_to_video_attn = Attention(
+ query_dim=video.dim,
+ context_dim=audio.dim,
+ heads=audio.heads,
+ dim_head=audio.d_head,
+ rope_type=rope_type,
+ norm_eps=norm_eps,
+ apply_gated_attention=video.apply_gated_attention,
+ )
+
+ # Q: Audio, K,V: Video
+ self.video_to_audio_attn = Attention(
+ query_dim=audio.dim,
+ context_dim=video.dim,
+ heads=audio.heads,
+ dim_head=audio.d_head,
+ rope_type=rope_type,
+ norm_eps=norm_eps,
+ apply_gated_attention=audio.apply_gated_attention,
+ )
+
+ self.scale_shift_table_a2v_ca_audio = torch.nn.Parameter(torch.empty(5, audio.dim))
+ self.scale_shift_table_a2v_ca_video = torch.nn.Parameter(torch.empty(5, video.dim))
+
+ self.cross_attention_adaln = (video is not None and video.cross_attention_adaln) or (
+ audio is not None and audio.cross_attention_adaln
+ )
+
+ if self.cross_attention_adaln and video is not None:
+ self.prompt_scale_shift_table = torch.nn.Parameter(torch.empty(2, video.dim))
+ if self.cross_attention_adaln and audio is not None:
+ self.audio_prompt_scale_shift_table = torch.nn.Parameter(torch.empty(2, audio.dim))
+
+ self.norm_eps = norm_eps
+
+ def get_ada_values(
+ self, scale_shift_table: torch.Tensor, batch_size: int, timestep: torch.Tensor, indices: slice
+ ) -> tuple[torch.Tensor, ...]:
+ num_ada_params = scale_shift_table.shape[0]
+
+ ada_values = (
+ scale_shift_table[indices].unsqueeze(0).unsqueeze(0).to(device=timestep.device, dtype=timestep.dtype)
+ + timestep.reshape(batch_size, timestep.shape[1], num_ada_params, -1)[:, :, indices, :]
+ ).unbind(dim=2)
+ return ada_values
+
+ def get_av_ca_ada_values(
+ self,
+ scale_shift_table: torch.Tensor,
+ batch_size: int,
+ scale_shift_timestep: torch.Tensor,
+ gate_timestep: torch.Tensor,
+ scale_shift_indices: slice,
+ num_scale_shift_values: int = 4,
+ ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+ scale_shift_ada_values = self.get_ada_values(
+ scale_shift_table[:num_scale_shift_values, :], batch_size, scale_shift_timestep, scale_shift_indices
+ )
+ gate_ada_values = self.get_ada_values(
+ scale_shift_table[num_scale_shift_values:, :], batch_size, gate_timestep, slice(None, None)
+ )
+
+ scale, shift = (t.squeeze(2) for t in scale_shift_ada_values)
+ (gate,) = (t.squeeze(2) for t in gate_ada_values)
+
+ return scale, shift, gate
+
+ def _apply_text_cross_attention(
+ self,
+ x: torch.Tensor,
+ context: torch.Tensor,
+ attn: Attention,
+ scale_shift_table: torch.Tensor,
+ prompt_scale_shift_table: torch.Tensor | None,
+ timestep: torch.Tensor,
+ prompt_timestep: torch.Tensor | None,
+ context_mask: torch.Tensor | None,
+ cross_attention_adaln: bool = False,
+ ) -> torch.Tensor:
+ """Apply text cross-attention, with optional AdaLN modulation."""
+ if cross_attention_adaln:
+ shift_q, scale_q, gate = self.get_ada_values(scale_shift_table, x.shape[0], timestep, slice(6, 9))
+ return apply_cross_attention_adaln(
+ x,
+ context,
+ attn,
+ shift_q,
+ scale_q,
+ gate,
+ prompt_scale_shift_table,
+ prompt_timestep,
+ context_mask,
+ self.norm_eps,
+ )
+ return attn(rms_norm(x, eps=self.norm_eps), context=context, mask=context_mask)
+
+ def forward( # noqa: PLR0915
+ self,
+ video: TransformerArgs | None,
+ audio: TransformerArgs | None,
+ perturbations: BatchedPerturbationConfig | None = None,
+ ) -> tuple[TransformerArgs | None, TransformerArgs | None]:
+ if video is None and audio is None:
+ raise ValueError("At least one of video or audio must be provided")
+
+ batch_size = (video or audio).x.shape[0]
+
+ if perturbations is None:
+ perturbations = BatchedPerturbationConfig.empty(batch_size)
+
+ vx = video.x if video is not None else None
+ ax = audio.x if audio is not None else None
+
+ run_vx = video is not None and video.enabled and vx.numel() > 0
+ run_ax = audio is not None and audio.enabled and ax.numel() > 0
+
+ run_a2v = run_vx and (audio is not None and ax.numel() > 0)
+ run_v2a = run_ax and (video is not None and vx.numel() > 0)
+
+ if run_vx:
+ vshift_msa, vscale_msa, vgate_msa = self.get_ada_values(
+ self.scale_shift_table, vx.shape[0], video.timesteps, slice(0, 3)
+ )
+ norm_vx = rms_norm(vx, eps=self.norm_eps) * (1 + vscale_msa) + vshift_msa
+ del vshift_msa, vscale_msa
+
+ all_perturbed = perturbations.all_in_batch(PerturbationType.SKIP_VIDEO_SELF_ATTN, self.idx)
+ none_perturbed = not perturbations.any_in_batch(PerturbationType.SKIP_VIDEO_SELF_ATTN, self.idx)
+ v_mask = (
+ perturbations.mask_like(PerturbationType.SKIP_VIDEO_SELF_ATTN, self.idx, vx)
+ if not all_perturbed and not none_perturbed
+ else None
+ )
+ vx = (
+ vx
+ + self.attn1(
+ norm_vx,
+ pe=video.positional_embeddings,
+ mask=video.self_attention_mask,
+ perturbation_mask=v_mask,
+ all_perturbed=all_perturbed,
+ )
+ * vgate_msa
+ )
+ del vgate_msa, norm_vx, v_mask
+ vx = vx + self._apply_text_cross_attention(
+ vx,
+ video.context,
+ self.attn2,
+ self.scale_shift_table,
+ getattr(self, "prompt_scale_shift_table", None),
+ video.timesteps,
+ video.prompt_timestep,
+ video.context_mask,
+ cross_attention_adaln=self.cross_attention_adaln,
+ )
+
+ if run_ax:
+ ashift_msa, ascale_msa, agate_msa = self.get_ada_values(
+ self.audio_scale_shift_table, ax.shape[0], audio.timesteps, slice(0, 3)
+ )
+
+ norm_ax = rms_norm(ax, eps=self.norm_eps) * (1 + ascale_msa) + ashift_msa
+ del ashift_msa, ascale_msa
+ all_perturbed = perturbations.all_in_batch(PerturbationType.SKIP_AUDIO_SELF_ATTN, self.idx)
+ none_perturbed = not perturbations.any_in_batch(PerturbationType.SKIP_AUDIO_SELF_ATTN, self.idx)
+ a_mask = (
+ perturbations.mask_like(PerturbationType.SKIP_AUDIO_SELF_ATTN, self.idx, ax)
+ if not all_perturbed and not none_perturbed
+ else None
+ )
+ ax = (
+ ax
+ + self.audio_attn1(
+ norm_ax,
+ pe=audio.positional_embeddings,
+ mask=audio.self_attention_mask,
+ perturbation_mask=a_mask,
+ all_perturbed=all_perturbed,
+ )
+ * agate_msa
+ )
+ del agate_msa, norm_ax, a_mask
+ ax = ax + self._apply_text_cross_attention(
+ ax,
+ audio.context,
+ self.audio_attn2,
+ self.audio_scale_shift_table,
+ getattr(self, "audio_prompt_scale_shift_table", None),
+ audio.timesteps,
+ audio.prompt_timestep,
+ audio.context_mask,
+ cross_attention_adaln=self.cross_attention_adaln,
+ )
+
+ # Audio - Video cross attention.
+ if run_a2v or run_v2a:
+ vx_norm3 = rms_norm(vx, eps=self.norm_eps)
+ ax_norm3 = rms_norm(ax, eps=self.norm_eps)
+
+ if run_a2v and not perturbations.all_in_batch(PerturbationType.SKIP_A2V_CROSS_ATTN, self.idx):
+ scale_ca_video_a2v, shift_ca_video_a2v, gate_out_a2v = self.get_av_ca_ada_values(
+ self.scale_shift_table_a2v_ca_video,
+ vx.shape[0],
+ video.cross_scale_shift_timestep,
+ video.cross_gate_timestep,
+ slice(0, 2),
+ )
+ vx_scaled = vx_norm3 * (1 + scale_ca_video_a2v) + shift_ca_video_a2v
+ del scale_ca_video_a2v, shift_ca_video_a2v
+
+ scale_ca_audio_a2v, shift_ca_audio_a2v, _ = self.get_av_ca_ada_values(
+ self.scale_shift_table_a2v_ca_audio,
+ ax.shape[0],
+ audio.cross_scale_shift_timestep,
+ audio.cross_gate_timestep,
+ slice(0, 2),
+ )
+ ax_scaled = ax_norm3 * (1 + scale_ca_audio_a2v) + shift_ca_audio_a2v
+ del scale_ca_audio_a2v, shift_ca_audio_a2v
+ a2v_mask = perturbations.mask_like(PerturbationType.SKIP_A2V_CROSS_ATTN, self.idx, vx)
+ vx = vx + (
+ self.audio_to_video_attn(
+ vx_scaled,
+ context=ax_scaled,
+ pe=video.cross_positional_embeddings,
+ k_pe=audio.cross_positional_embeddings,
+ )
+ * gate_out_a2v
+ * a2v_mask
+ )
+ del gate_out_a2v, a2v_mask, vx_scaled, ax_scaled
+
+ if run_v2a and not perturbations.all_in_batch(PerturbationType.SKIP_V2A_CROSS_ATTN, self.idx):
+ scale_ca_audio_v2a, shift_ca_audio_v2a, gate_out_v2a = self.get_av_ca_ada_values(
+ self.scale_shift_table_a2v_ca_audio,
+ ax.shape[0],
+ audio.cross_scale_shift_timestep,
+ audio.cross_gate_timestep,
+ slice(2, 4),
+ )
+ ax_scaled = ax_norm3 * (1 + scale_ca_audio_v2a) + shift_ca_audio_v2a
+ del scale_ca_audio_v2a, shift_ca_audio_v2a
+ scale_ca_video_v2a, shift_ca_video_v2a, _ = self.get_av_ca_ada_values(
+ self.scale_shift_table_a2v_ca_video,
+ vx.shape[0],
+ video.cross_scale_shift_timestep,
+ video.cross_gate_timestep,
+ slice(2, 4),
+ )
+ vx_scaled = vx_norm3 * (1 + scale_ca_video_v2a) + shift_ca_video_v2a
+ del scale_ca_video_v2a, shift_ca_video_v2a
+ v2a_mask = perturbations.mask_like(PerturbationType.SKIP_V2A_CROSS_ATTN, self.idx, ax)
+ ax = ax + (
+ self.video_to_audio_attn(
+ ax_scaled,
+ context=vx_scaled,
+ pe=audio.cross_positional_embeddings,
+ k_pe=video.cross_positional_embeddings,
+ )
+ * gate_out_v2a
+ * v2a_mask
+ )
+ del gate_out_v2a, v2a_mask, ax_scaled, vx_scaled
+
+ del vx_norm3, ax_norm3
+
+ if run_vx:
+ vshift_mlp, vscale_mlp, vgate_mlp = self.get_ada_values(
+ self.scale_shift_table, vx.shape[0], video.timesteps, slice(3, 6)
+ )
+ vx_scaled = rms_norm(vx, eps=self.norm_eps) * (1 + vscale_mlp) + vshift_mlp
+ vx = vx + self.ff(vx_scaled) * vgate_mlp
+
+ del vshift_mlp, vscale_mlp, vgate_mlp, vx_scaled
+
+ if run_ax:
+ ashift_mlp, ascale_mlp, agate_mlp = self.get_ada_values(
+ self.audio_scale_shift_table, ax.shape[0], audio.timesteps, slice(3, 6)
+ )
+ ax_scaled = rms_norm(ax, eps=self.norm_eps) * (1 + ascale_mlp) + ashift_mlp
+ ax = ax + self.audio_ff(ax_scaled) * agate_mlp
+
+ del ashift_mlp, ascale_mlp, agate_mlp, ax_scaled
+
+ return replace(video, x=vx) if video is not None else None, replace(audio, x=ax) if audio is not None else None
+
+
+def apply_cross_attention_adaln(
+ x: torch.Tensor,
+ context: torch.Tensor,
+ attn: Attention,
+ q_shift: torch.Tensor,
+ q_scale: torch.Tensor,
+ q_gate: torch.Tensor,
+ prompt_scale_shift_table: torch.Tensor,
+ prompt_timestep: torch.Tensor,
+ context_mask: torch.Tensor | None = None,
+ norm_eps: float = 1e-6,
+) -> torch.Tensor:
+ batch_size = x.shape[0]
+ shift_kv, scale_kv = (
+ prompt_scale_shift_table[None, None].to(device=x.device, dtype=x.dtype)
+ + prompt_timestep.reshape(batch_size, prompt_timestep.shape[1], 2, -1)
+ ).unbind(dim=2)
+ attn_input = rms_norm(x, eps=norm_eps) * (1 + q_scale) + q_shift
+ encoder_hidden_states = context * (1 + scale_kv) + shift_kv
+ return attn(attn_input, context=encoder_hidden_states, mask=context_mask) * q_gate
+
+
+class GELUApprox(torch.nn.Module):
+ def __init__(self, dim_in: int, dim_out: int) -> None:
+ super().__init__()
+ self.proj = torch.nn.Linear(dim_in, dim_out)
+
+ def forward(self, x: torch.Tensor) -> torch.Tensor:
+ return torch.nn.functional.gelu(self.proj(x), approximate="tanh")
+
+
+class FeedForward(torch.nn.Module):
+ def __init__(self, dim: int, dim_out: int, mult: int = 4) -> None:
+ super().__init__()
+ inner_dim = int(dim * mult)
+ project_in = GELUApprox(dim, inner_dim)
+
+ self.net = torch.nn.Sequential(project_in, torch.nn.Identity(), torch.nn.Linear(inner_dim, dim_out))
+
+ def forward(self, x: torch.Tensor) -> torch.Tensor:
+ return self.net(x)
+
+
+class LTXModelType(Enum):
+ AudioVideo = "ltx av model"
+ VideoOnly = "ltx video only model"
+ AudioOnly = "ltx audio only model"
+
+ def is_video_enabled(self) -> bool:
+ return self in (LTXModelType.AudioVideo, LTXModelType.VideoOnly)
+
+ def is_audio_enabled(self) -> bool:
+ return self in (LTXModelType.AudioVideo, LTXModelType.AudioOnly)
+
+
+class LTXModel(torch.nn.Module):
+ """
+ LTX model transformer implementation.
+ This class implements the transformer blocks for the LTX model.
+ """
+ _repeated_blocks = ["BasicAVTransformerBlock"]
+
+ def __init__( # noqa: PLR0913
+ self,
+ *,
+ model_type: LTXModelType = LTXModelType.AudioVideo,
+ num_attention_heads: int = 32,
+ attention_head_dim: int = 128,
+ in_channels: int = 128,
+ out_channels: int = 128,
+ num_layers: int = 48,
+ cross_attention_dim: int = 4096,
+ norm_eps: float = 1e-06,
+ caption_channels: int = 3840,
+ positional_embedding_theta: float = 10000.0,
+ positional_embedding_max_pos: list[int] | None = [20, 2048, 2048],
+ timestep_scale_multiplier: int = 1000,
+ use_middle_indices_grid: bool = True,
+ audio_num_attention_heads: int = 32,
+ audio_attention_head_dim: int = 64,
+ audio_in_channels: int = 128,
+ audio_out_channels: int = 128,
+ audio_cross_attention_dim: int = 2048,
+ audio_positional_embedding_max_pos: list[int] | None = [20],
+ av_ca_timestep_scale_multiplier: int = 1000,
+ rope_type: LTXRopeType = LTXRopeType.SPLIT,
+ double_precision_rope: bool = True,
+ apply_gated_attention: bool = False,
+ cross_attention_adaln: bool = False,
+ ):
+ super().__init__()
+ self._enable_gradient_checkpointing = False
+ self.use_middle_indices_grid = use_middle_indices_grid
+ self.rope_type = rope_type
+ self.double_precision_rope = double_precision_rope
+ self.timestep_scale_multiplier = timestep_scale_multiplier
+ self.positional_embedding_theta = positional_embedding_theta
+ self.model_type = model_type
+ self.cross_attention_adaln = cross_attention_adaln
+ cross_pe_max_pos = None
+ if model_type.is_video_enabled():
+ if positional_embedding_max_pos is None:
+ positional_embedding_max_pos = [20, 2048, 2048]
+ self.positional_embedding_max_pos = positional_embedding_max_pos
+ self.num_attention_heads = num_attention_heads
+ self.inner_dim = num_attention_heads * attention_head_dim
+ self._init_video(
+ in_channels=in_channels,
+ out_channels=out_channels,
+ caption_channels=caption_channels,
+ norm_eps=norm_eps,
+ )
+
+ if model_type.is_audio_enabled():
+ if audio_positional_embedding_max_pos is None:
+ audio_positional_embedding_max_pos = [20]
+ self.audio_positional_embedding_max_pos = audio_positional_embedding_max_pos
+ self.audio_num_attention_heads = audio_num_attention_heads
+ self.audio_inner_dim = self.audio_num_attention_heads * audio_attention_head_dim
+ self._init_audio(
+ in_channels=audio_in_channels,
+ out_channels=audio_out_channels,
+ caption_channels=caption_channels,
+ norm_eps=norm_eps,
+ )
+
+ if model_type.is_video_enabled() and model_type.is_audio_enabled():
+ cross_pe_max_pos = max(self.positional_embedding_max_pos[0], self.audio_positional_embedding_max_pos[0])
+ self.av_ca_timestep_scale_multiplier = av_ca_timestep_scale_multiplier
+ self.audio_cross_attention_dim = audio_cross_attention_dim
+ self._init_audio_video(num_scale_shift_values=4)
+
+ self._init_preprocessors(cross_pe_max_pos)
+ # Initialize transformer blocks
+ self._init_transformer_blocks(
+ num_layers=num_layers,
+ attention_head_dim=attention_head_dim if model_type.is_video_enabled() else 0,
+ cross_attention_dim=cross_attention_dim,
+ audio_attention_head_dim=audio_attention_head_dim if model_type.is_audio_enabled() else 0,
+ audio_cross_attention_dim=audio_cross_attention_dim,
+ norm_eps=norm_eps,
+ apply_gated_attention=apply_gated_attention,
+ )
+
+ @property
+ def _adaln_embedding_coefficient(self) -> int:
+ return adaln_embedding_coefficient(self.cross_attention_adaln)
+
+ def _init_video(
+ self,
+ in_channels: int,
+ out_channels: int,
+ caption_channels: int,
+ norm_eps: float,
+ ) -> None:
+ """Initialize video-specific components."""
+ # Video input components
+ self.patchify_proj = torch.nn.Linear(in_channels, self.inner_dim, bias=True)
+ self.adaln_single = AdaLayerNormSingle(self.inner_dim, embedding_coefficient=self._adaln_embedding_coefficient)
+ self.prompt_adaln_single = AdaLayerNormSingle(self.inner_dim, embedding_coefficient=2) if self.cross_attention_adaln else None
+
+ # Video caption projection
+ if caption_channels is not None:
+ self.caption_projection = PixArtAlphaTextProjection(
+ in_features=caption_channels,
+ hidden_size=self.inner_dim,
+ )
+
+ # Video output components
+ self.scale_shift_table = torch.nn.Parameter(torch.empty(2, self.inner_dim))
+ self.norm_out = torch.nn.LayerNorm(self.inner_dim, elementwise_affine=False, eps=norm_eps)
+ self.proj_out = torch.nn.Linear(self.inner_dim, out_channels)
+
+ def _init_audio(
+ self,
+ in_channels: int,
+ out_channels: int,
+ caption_channels: int,
+ norm_eps: float,
+ ) -> None:
+ """Initialize audio-specific components."""
+
+ # Audio input components
+ self.audio_patchify_proj = torch.nn.Linear(in_channels, self.audio_inner_dim, bias=True)
+
+ self.audio_adaln_single = AdaLayerNormSingle(self.audio_inner_dim, embedding_coefficient=self._adaln_embedding_coefficient)
+ self.audio_prompt_adaln_single = AdaLayerNormSingle(self.audio_inner_dim, embedding_coefficient=2) if self.cross_attention_adaln else None
+
+ # Audio caption projection
+ if caption_channels is not None:
+ self.audio_caption_projection = PixArtAlphaTextProjection(
+ in_features=caption_channels,
+ hidden_size=self.audio_inner_dim,
+ )
+
+ # Audio output components
+ self.audio_scale_shift_table = torch.nn.Parameter(torch.empty(2, self.audio_inner_dim))
+ self.audio_norm_out = torch.nn.LayerNorm(self.audio_inner_dim, elementwise_affine=False, eps=norm_eps)
+ self.audio_proj_out = torch.nn.Linear(self.audio_inner_dim, out_channels)
+
+ def _init_audio_video(
+ self,
+ num_scale_shift_values: int,
+ ) -> None:
+ """Initialize audio-video cross-attention components."""
+ self.av_ca_video_scale_shift_adaln_single = AdaLayerNormSingle(
+ self.inner_dim,
+ embedding_coefficient=num_scale_shift_values,
+ )
+
+ self.av_ca_audio_scale_shift_adaln_single = AdaLayerNormSingle(
+ self.audio_inner_dim,
+ embedding_coefficient=num_scale_shift_values,
+ )
+
+ self.av_ca_a2v_gate_adaln_single = AdaLayerNormSingle(
+ self.inner_dim,
+ embedding_coefficient=1,
+ )
+
+ self.av_ca_v2a_gate_adaln_single = AdaLayerNormSingle(
+ self.audio_inner_dim,
+ embedding_coefficient=1,
+ )
+
+ def _init_preprocessors(
+ self,
+ cross_pe_max_pos: int | None = None,
+ ) -> None:
+ """Initialize preprocessors for LTX."""
+
+ if self.model_type.is_video_enabled() and self.model_type.is_audio_enabled():
+ self.video_args_preprocessor = MultiModalTransformerArgsPreprocessor(
+ patchify_proj=self.patchify_proj,
+ adaln=self.adaln_single,
+ cross_scale_shift_adaln=self.av_ca_video_scale_shift_adaln_single,
+ cross_gate_adaln=self.av_ca_a2v_gate_adaln_single,
+ inner_dim=self.inner_dim,
+ max_pos=self.positional_embedding_max_pos,
+ num_attention_heads=self.num_attention_heads,
+ cross_pe_max_pos=cross_pe_max_pos,
+ use_middle_indices_grid=self.use_middle_indices_grid,
+ audio_cross_attention_dim=self.audio_cross_attention_dim,
+ timestep_scale_multiplier=self.timestep_scale_multiplier,
+ double_precision_rope=self.double_precision_rope,
+ positional_embedding_theta=self.positional_embedding_theta,
+ rope_type=self.rope_type,
+ av_ca_timestep_scale_multiplier=self.av_ca_timestep_scale_multiplier,
+ caption_projection=getattr(self, "caption_projection", None),
+ prompt_adaln=getattr(self, "prompt_adaln_single", None),
+ )
+ self.audio_args_preprocessor = MultiModalTransformerArgsPreprocessor(
+ patchify_proj=self.audio_patchify_proj,
+ adaln=self.audio_adaln_single,
+ cross_scale_shift_adaln=self.av_ca_audio_scale_shift_adaln_single,
+ cross_gate_adaln=self.av_ca_v2a_gate_adaln_single,
+ inner_dim=self.audio_inner_dim,
+ max_pos=self.audio_positional_embedding_max_pos,
+ num_attention_heads=self.audio_num_attention_heads,
+ cross_pe_max_pos=cross_pe_max_pos,
+ use_middle_indices_grid=self.use_middle_indices_grid,
+ audio_cross_attention_dim=self.audio_cross_attention_dim,
+ timestep_scale_multiplier=self.timestep_scale_multiplier,
+ double_precision_rope=self.double_precision_rope,
+ positional_embedding_theta=self.positional_embedding_theta,
+ rope_type=self.rope_type,
+ av_ca_timestep_scale_multiplier=self.av_ca_timestep_scale_multiplier,
+ caption_projection=getattr(self, "audio_caption_projection", None),
+ prompt_adaln=getattr(self, "audio_prompt_adaln_single", None),
+ )
+ elif self.model_type.is_video_enabled():
+ self.video_args_preprocessor = TransformerArgsPreprocessor(
+ patchify_proj=self.patchify_proj,
+ adaln=self.adaln_single,
+ inner_dim=self.inner_dim,
+ max_pos=self.positional_embedding_max_pos,
+ num_attention_heads=self.num_attention_heads,
+ use_middle_indices_grid=self.use_middle_indices_grid,
+ timestep_scale_multiplier=self.timestep_scale_multiplier,
+ double_precision_rope=self.double_precision_rope,
+ positional_embedding_theta=self.positional_embedding_theta,
+ rope_type=self.rope_type,
+ caption_projection=getattr(self, "caption_projection", None),
+ prompt_adaln=getattr(self, "prompt_adaln_single", None),
+ )
+ elif self.model_type.is_audio_enabled():
+ self.audio_args_preprocessor = TransformerArgsPreprocessor(
+ patchify_proj=self.audio_patchify_proj,
+ adaln=self.audio_adaln_single,
+ inner_dim=self.audio_inner_dim,
+ max_pos=self.audio_positional_embedding_max_pos,
+ num_attention_heads=self.audio_num_attention_heads,
+ use_middle_indices_grid=self.use_middle_indices_grid,
+ timestep_scale_multiplier=self.timestep_scale_multiplier,
+ double_precision_rope=self.double_precision_rope,
+ positional_embedding_theta=self.positional_embedding_theta,
+ rope_type=self.rope_type,
+ caption_projection=getattr(self, "audio_caption_projection", None),
+ prompt_adaln=getattr(self, "audio_prompt_adaln_single", None),
+ )
+
+ def _init_transformer_blocks(
+ self,
+ num_layers: int,
+ attention_head_dim: int,
+ cross_attention_dim: int,
+ audio_attention_head_dim: int,
+ audio_cross_attention_dim: int,
+ norm_eps: float,
+ apply_gated_attention: bool,
+ ) -> None:
+ """Initialize transformer blocks for LTX."""
+ video_config = (
+ TransformerConfig(
+ dim=self.inner_dim,
+ heads=self.num_attention_heads,
+ d_head=attention_head_dim,
+ context_dim=cross_attention_dim,
+ apply_gated_attention=apply_gated_attention,
+ cross_attention_adaln=self.cross_attention_adaln,
+ )
+ if self.model_type.is_video_enabled()
+ else None
+ )
+ audio_config = (
+ TransformerConfig(
+ dim=self.audio_inner_dim,
+ heads=self.audio_num_attention_heads,
+ d_head=audio_attention_head_dim,
+ context_dim=audio_cross_attention_dim,
+ apply_gated_attention=apply_gated_attention,
+ cross_attention_adaln=self.cross_attention_adaln,
+ )
+ if self.model_type.is_audio_enabled()
+ else None
+ )
+ self.transformer_blocks = torch.nn.ModuleList(
+ [
+ BasicAVTransformerBlock(
+ idx=idx,
+ video=video_config,
+ audio=audio_config,
+ rope_type=self.rope_type,
+ norm_eps=norm_eps,
+ )
+ for idx in range(num_layers)
+ ]
+ )
+
+ def set_gradient_checkpointing(self, enable: bool) -> None:
+ """Enable or disable gradient checkpointing for transformer blocks.
+ Gradient checkpointing trades compute for memory by recomputing activations
+ during the backward pass instead of storing them. This can significantly
+ reduce memory usage at the cost of ~20-30% slower training.
+ Args:
+ enable: Whether to enable gradient checkpointing
+ """
+ self._enable_gradient_checkpointing = enable
+
+ def _process_transformer_blocks(
+ self,
+ video: TransformerArgs | None,
+ audio: TransformerArgs | None,
+ perturbations: BatchedPerturbationConfig,
+ use_gradient_checkpointing: bool = False,
+ use_gradient_checkpointing_offload: bool = False,
+ ) -> tuple[TransformerArgs, TransformerArgs]:
+ """Process transformer blocks for LTXAV."""
+
+ # Process transformer blocks
+ for block in self.transformer_blocks:
+ video, audio = gradient_checkpoint_forward(
+ block,
+ use_gradient_checkpointing,
+ use_gradient_checkpointing_offload,
+ video=video,
+ audio=audio,
+ perturbations=perturbations,
+ )
+
+ return video, audio
+
+ def _process_output(
+ self,
+ scale_shift_table: torch.Tensor,
+ norm_out: torch.nn.LayerNorm,
+ proj_out: torch.nn.Linear,
+ x: torch.Tensor,
+ embedded_timestep: torch.Tensor,
+ ) -> torch.Tensor:
+ """Process output for LTXV."""
+ # Apply scale-shift modulation
+ scale_shift_values = (
+ scale_shift_table[None, None].to(device=x.device, dtype=x.dtype) + embedded_timestep[:, :, None]
+ )
+ shift, scale = scale_shift_values[:, :, 0], scale_shift_values[:, :, 1]
+
+ x = norm_out(x)
+ x = x * (1 + scale) + shift
+ x = proj_out(x)
+ return x
+
+ def _forward(
+ self,
+ video: Modality | None,
+ audio: Modality | None,
+ perturbations: BatchedPerturbationConfig,
+ use_gradient_checkpointing: bool = False,
+ use_gradient_checkpointing_offload: bool = False,
+ ) -> tuple[torch.Tensor, torch.Tensor]:
+ """
+ Forward pass for LTX models.
+ Returns:
+ Processed output tensors
+ """
+ if not self.model_type.is_video_enabled() and video is not None:
+ raise ValueError("Video is not enabled for this model")
+ if not self.model_type.is_audio_enabled() and audio is not None:
+ raise ValueError("Audio is not enabled for this model")
+
+ video_args = self.video_args_preprocessor.prepare(video, audio) if video is not None else None
+ audio_args = self.audio_args_preprocessor.prepare(audio, video) if audio is not None else None
+ # Process transformer blocks
+ video_out, audio_out = self._process_transformer_blocks(
+ video=video_args,
+ audio=audio_args,
+ perturbations=perturbations,
+ use_gradient_checkpointing=use_gradient_checkpointing,
+ use_gradient_checkpointing_offload=use_gradient_checkpointing_offload,
+ )
+
+ # Process output
+ vx = (
+ self._process_output(
+ self.scale_shift_table, self.norm_out, self.proj_out, video_out.x, video_out.embedded_timestep
+ )
+ if video_out is not None
+ else None
+ )
+ ax = (
+ self._process_output(
+ self.audio_scale_shift_table,
+ self.audio_norm_out,
+ self.audio_proj_out,
+ audio_out.x,
+ audio_out.embedded_timestep,
+ )
+ if audio_out is not None
+ else None
+ )
+ return vx, ax
+
+ def forward(self, video_latents, video_positions, video_context, video_timesteps, audio_latents, audio_positions, audio_context, audio_timesteps, sigma, use_gradient_checkpointing=False, use_gradient_checkpointing_offload=False):
+ cross_pe_max_pos = None
+ if self.model_type.is_video_enabled() and self.model_type.is_audio_enabled():
+ cross_pe_max_pos = max(self.positional_embedding_max_pos[0], self.audio_positional_embedding_max_pos[0])
+ self._init_preprocessors(cross_pe_max_pos)
+ video = Modality(video_latents, sigma, video_timesteps, video_positions, video_context)
+ audio = Modality(audio_latents, sigma, audio_timesteps, audio_positions, audio_context) if audio_latents is not None else None
+ vx, ax = self._forward(video=video, audio=audio, perturbations=None, use_gradient_checkpointing=use_gradient_checkpointing, use_gradient_checkpointing_offload=use_gradient_checkpointing_offload)
+ return vx, ax
diff --git a/diffsynth/models/ltx2_text_encoder.py b/diffsynth/models/ltx2_text_encoder.py
new file mode 100644
index 0000000000000000000000000000000000000000..e4f3b1a3e1995b1adefadf31348a9fe19b9f8f02
--- /dev/null
+++ b/diffsynth/models/ltx2_text_encoder.py
@@ -0,0 +1,549 @@
+import math
+import torch
+import torch.nn as nn
+from einops import rearrange
+from transformers import Gemma3ForConditionalGeneration, Gemma3Config, AutoTokenizer
+from .ltx2_dit import (LTXRopeType, generate_freq_grid_np, generate_freq_grid_pytorch, precompute_freqs_cis, Attention,
+ FeedForward)
+from .ltx2_common import rms_norm
+
+
+class LTX2TextEncoder(Gemma3ForConditionalGeneration):
+ def __init__(self):
+ config = Gemma3Config(
+ **{
+ "architectures": ["Gemma3ForConditionalGeneration"],
+ "boi_token_index": 255999,
+ "dtype": "bfloat16",
+ "eoi_token_index": 256000,
+ "eos_token_id": [1, 106],
+ "image_token_index": 262144,
+ "initializer_range": 0.02,
+ "mm_tokens_per_image": 256,
+ "model_type": "gemma3",
+ "text_config": {
+ "_sliding_window_pattern": 6,
+ "attention_bias": False,
+ "attention_dropout": 0.0,
+ "attn_logit_softcapping": None,
+ "cache_implementation": "hybrid",
+ "dtype": "bfloat16",
+ "final_logit_softcapping": None,
+ "head_dim": 256,
+ "hidden_activation": "gelu_pytorch_tanh",
+ "hidden_size": 3840,
+ "initializer_range": 0.02,
+ "intermediate_size": 15360,
+ "layer_types": [
+ "sliding_attention", "sliding_attention", "sliding_attention", "sliding_attention",
+ "sliding_attention", "full_attention", "sliding_attention", "sliding_attention",
+ "sliding_attention", "sliding_attention", "sliding_attention", "full_attention",
+ "sliding_attention", "sliding_attention", "sliding_attention", "sliding_attention",
+ "sliding_attention", "full_attention", "sliding_attention", "sliding_attention",
+ "sliding_attention", "sliding_attention", "sliding_attention", "full_attention",
+ "sliding_attention", "sliding_attention", "sliding_attention", "sliding_attention",
+ "sliding_attention", "full_attention", "sliding_attention", "sliding_attention",
+ "sliding_attention", "sliding_attention", "sliding_attention", "full_attention",
+ "sliding_attention", "sliding_attention", "sliding_attention", "sliding_attention",
+ "sliding_attention", "full_attention", "sliding_attention", "sliding_attention",
+ "sliding_attention", "sliding_attention", "sliding_attention", "full_attention"
+ ],
+ "max_position_embeddings": 131072,
+ "model_type": "gemma3_text",
+ "num_attention_heads": 16,
+ "num_hidden_layers": 48,
+ "num_key_value_heads": 8,
+ "query_pre_attn_scalar": 256,
+ "rms_norm_eps": 1e-06,
+ "rope_local_base_freq": 10000,
+ "rope_scaling": {
+ "factor": 8.0,
+ "rope_type": "linear"
+ },
+ "rope_theta": 1000000,
+ "sliding_window": 1024,
+ "sliding_window_pattern": 6,
+ "use_bidirectional_attention": False,
+ "use_cache": True,
+ "vocab_size": 262208
+ },
+ "transformers_version": "4.57.3",
+ "vision_config": {
+ "attention_dropout": 0.0,
+ "dtype": "bfloat16",
+ "hidden_act": "gelu_pytorch_tanh",
+ "hidden_size": 1152,
+ "image_size": 896,
+ "intermediate_size": 4304,
+ "layer_norm_eps": 1e-06,
+ "model_type": "siglip_vision_model",
+ "num_attention_heads": 16,
+ "num_channels": 3,
+ "num_hidden_layers": 27,
+ "patch_size": 14,
+ "vision_use_head": False
+ }
+ })
+ super().__init__(config)
+
+
+class LTXVGemmaTokenizer:
+ """
+ Tokenizer wrapper for Gemma models compatible with LTXV processes.
+ This class wraps HuggingFace's `AutoTokenizer` for use with Gemma text encoders,
+ ensuring correct settings and output formatting for downstream consumption.
+ """
+
+ def __init__(self, tokenizer_path: str, max_length: int = 1024):
+ """
+ Initialize the tokenizer.
+ Args:
+ tokenizer_path (str): Path to the pretrained tokenizer files or model directory.
+ max_length (int, optional): Max sequence length for encoding. Defaults to 256.
+ """
+ self.tokenizer = AutoTokenizer.from_pretrained(
+ tokenizer_path, local_files_only=True, model_max_length=max_length
+ )
+ # Gemma expects left padding for chat-style prompts; for plain text it doesn't matter much.
+ self.tokenizer.padding_side = "left"
+ if self.tokenizer.pad_token is None:
+ self.tokenizer.pad_token = self.tokenizer.eos_token
+
+ self.max_length = max_length
+
+ def tokenize_with_weights(self, text: str, return_word_ids: bool = False) -> dict[str, list[tuple[int, int]]]:
+ """
+ Tokenize the given text and return token IDs and attention weights.
+ Args:
+ text (str): The input string to tokenize.
+ return_word_ids (bool, optional): If True, includes the token's position (index) in the output tuples.
+ If False (default), omits the indices.
+ Returns:
+ dict[str, list[tuple[int, int]]] OR dict[str, list[tuple[int, int, int]]]:
+ A dictionary with a "gemma" key mapping to:
+ - a list of (token_id, attention_mask) tuples if return_word_ids is False;
+ - a list of (token_id, attention_mask, index) tuples if return_word_ids is True.
+ Example:
+ >>> tokenizer = LTXVGemmaTokenizer("path/to/tokenizer", max_length=8)
+ >>> tokenizer.tokenize_with_weights("hello world")
+ {'gemma': [(1234, 1), (5678, 1), (2, 0), ...]}
+ """
+ text = text.strip()
+ encoded = self.tokenizer(
+ text,
+ padding="max_length",
+ max_length=self.max_length,
+ truncation=True,
+ return_tensors="pt",
+ )
+ input_ids = encoded.input_ids
+ attention_mask = encoded.attention_mask
+ tuples = [
+ (token_id, attn, i) for i, (token_id, attn) in enumerate(zip(input_ids[0], attention_mask[0], strict=True))
+ ]
+ out = {"gemma": tuples}
+
+ if not return_word_ids:
+ # Return only (token_id, attention_mask) pairs, omitting token position
+ out = {k: [(t, w) for t, w, _ in v] for k, v in out.items()}
+
+ return out
+
+
+class GemmaFeaturesExtractorProjLinear(nn.Module):
+ """
+ Feature extractor module for Gemma models.
+ This module applies a single linear projection to the input tensor.
+ It expects a flattened feature tensor of shape (batch_size, 3840*49).
+ The linear layer maps this to a (batch_size, 3840) embedding.
+ Attributes:
+ aggregate_embed (nn.Linear): Linear projection layer.
+ """
+
+ def __init__(self) -> None:
+ """
+ Initialize the GemmaFeaturesExtractorProjLinear module.
+ The input dimension is expected to be 3840 * 49, and the output is 3840.
+ """
+ super().__init__()
+ self.aggregate_embed = nn.Linear(3840 * 49, 3840, bias=False)
+
+ def forward(
+ self,
+ hidden_states: torch.Tensor,
+ attention_mask: torch.Tensor,
+ padding_side: str = "left",
+ ) -> tuple[torch.Tensor, torch.Tensor | None]:
+ encoded = torch.stack(hidden_states, dim=-1) if isinstance(hidden_states, (list, tuple)) else hidden_states
+ dtype = encoded.dtype
+ sequence_lengths = attention_mask.sum(dim=-1)
+ normed = _norm_and_concat_padded_batch(encoded, sequence_lengths, padding_side)
+ features = self.aggregate_embed(normed.to(dtype))
+ return features, features
+
+
+class GemmaSeperatedFeaturesExtractorProjLinear(nn.Module):
+ """22B: per-token RMS norm → rescale → dual aggregate embeds"""
+
+ def __init__(
+ self,
+ num_layers: int,
+ embedding_dim: int,
+ video_inner_dim: int,
+ audio_inner_dim: int,
+ ):
+ super().__init__()
+ in_dim = embedding_dim * num_layers
+ self.video_aggregate_embed = torch.nn.Linear(in_dim, video_inner_dim, bias=True)
+ self.audio_aggregate_embed = torch.nn.Linear(in_dim, audio_inner_dim, bias=True)
+ self.embedding_dim = embedding_dim
+
+ def forward(
+ self,
+ hidden_states: torch.Tensor,
+ attention_mask: torch.Tensor,
+ padding_side: str = "left", # noqa: ARG002
+ ) -> tuple[torch.Tensor, torch.Tensor | None]:
+ encoded = torch.stack(hidden_states, dim=-1) if isinstance(hidden_states, (list, tuple)) else hidden_states
+ normed = norm_and_concat_per_token_rms(encoded, attention_mask)
+ normed = normed.to(encoded.dtype)
+ v_dim = self.video_aggregate_embed.out_features
+ video = self.video_aggregate_embed(_rescale_norm(normed, v_dim, self.embedding_dim))
+ audio = None
+ if self.audio_aggregate_embed is not None:
+ a_dim = self.audio_aggregate_embed.out_features
+ audio = self.audio_aggregate_embed(_rescale_norm(normed, a_dim, self.embedding_dim))
+ return video, audio
+
+
+
+class _BasicTransformerBlock1D(nn.Module):
+ def __init__(
+ self,
+ dim: int,
+ heads: int,
+ dim_head: int,
+ rope_type: LTXRopeType = LTXRopeType.INTERLEAVED,
+ apply_gated_attention: bool = False,
+ ):
+ super().__init__()
+
+ self.attn1 = Attention(
+ query_dim=dim,
+ heads=heads,
+ dim_head=dim_head,
+ rope_type=rope_type,
+ apply_gated_attention=apply_gated_attention,
+ )
+
+ self.ff = FeedForward(
+ dim,
+ dim_out=dim,
+ )
+
+ def forward(
+ self,
+ hidden_states: torch.Tensor,
+ attention_mask: torch.Tensor | None = None,
+ pe: torch.Tensor | None = None,
+ ) -> torch.Tensor:
+ # Notice that normalization is always applied before the real computation in the following blocks.
+
+ # 1. Normalization Before Self-Attention
+ norm_hidden_states = rms_norm(hidden_states)
+
+ norm_hidden_states = norm_hidden_states.squeeze(1)
+
+ # 2. Self-Attention
+ attn_output = self.attn1(norm_hidden_states, mask=attention_mask, pe=pe)
+
+ hidden_states = attn_output + hidden_states
+ if hidden_states.ndim == 4:
+ hidden_states = hidden_states.squeeze(1)
+
+ # 3. Normalization before Feed-Forward
+ norm_hidden_states = rms_norm(hidden_states)
+
+ # 4. Feed-forward
+ ff_output = self.ff(norm_hidden_states)
+
+ hidden_states = ff_output + hidden_states
+ if hidden_states.ndim == 4:
+ hidden_states = hidden_states.squeeze(1)
+
+ return hidden_states
+
+
+class Embeddings1DConnector(nn.Module):
+ """
+ Embeddings1DConnector applies a 1D transformer-based processing to sequential embeddings (e.g., for video, audio, or
+ other modalities). It supports rotary positional encoding (rope), optional causal temporal positioning, and can
+ substitute padded positions with learnable registers. The module is highly configurable for head size, number of
+ layers, and register usage.
+ Args:
+ attention_head_dim (int): Dimension of each attention head (default=128).
+ num_attention_heads (int): Number of attention heads (default=30).
+ num_layers (int): Number of transformer layers (default=2).
+ positional_embedding_theta (float): Scaling factor for position embedding (default=10000.0).
+ positional_embedding_max_pos (list[int] | None): Max positions for positional embeddings (default=[1]).
+ causal_temporal_positioning (bool): If True, uses causal attention (default=False).
+ num_learnable_registers (int | None): Number of learnable registers to replace padded tokens. If None, disables
+ register replacement. (default=128)
+ rope_type (LTXRopeType): The RoPE variant to use (default=DEFAULT_ROPE_TYPE).
+ double_precision_rope (bool): Use double precision rope calculation (default=False).
+ """
+
+ _supports_gradient_checkpointing = True
+
+ def __init__(
+ self,
+ attention_head_dim: int = 128,
+ num_attention_heads: int = 30,
+ num_layers: int = 2,
+ positional_embedding_theta: float = 10000.0,
+ positional_embedding_max_pos: list[int] | None = [4096],
+ causal_temporal_positioning: bool = False,
+ num_learnable_registers: int | None = 128,
+ rope_type: LTXRopeType = LTXRopeType.SPLIT,
+ double_precision_rope: bool = True,
+ apply_gated_attention: bool = False,
+ ):
+ super().__init__()
+ self.num_attention_heads = num_attention_heads
+ self.inner_dim = num_attention_heads * attention_head_dim
+ self.causal_temporal_positioning = causal_temporal_positioning
+ self.positional_embedding_theta = positional_embedding_theta
+ self.positional_embedding_max_pos = (
+ positional_embedding_max_pos if positional_embedding_max_pos is not None else [1]
+ )
+ self.rope_type = rope_type
+ self.double_precision_rope = double_precision_rope
+ self.transformer_1d_blocks = nn.ModuleList(
+ [
+ _BasicTransformerBlock1D(
+ dim=self.inner_dim,
+ heads=num_attention_heads,
+ dim_head=attention_head_dim,
+ rope_type=rope_type,
+ apply_gated_attention=apply_gated_attention,
+ )
+ for _ in range(num_layers)
+ ]
+ )
+
+ self.num_learnable_registers = num_learnable_registers
+ if self.num_learnable_registers:
+ self.learnable_registers = nn.Parameter(
+ torch.rand(self.num_learnable_registers, self.inner_dim, dtype=torch.bfloat16) * 2.0 - 1.0
+ )
+
+ def _replace_padded_with_learnable_registers(
+ self, hidden_states: torch.Tensor, attention_mask: torch.Tensor
+ ) -> tuple[torch.Tensor, torch.Tensor]:
+ assert hidden_states.shape[1] % self.num_learnable_registers == 0, (
+ f"Hidden states sequence length {hidden_states.shape[1]} must be divisible by num_learnable_registers "
+ f"{self.num_learnable_registers}."
+ )
+
+ num_registers_duplications = hidden_states.shape[1] // self.num_learnable_registers
+ learnable_registers = torch.tile(self.learnable_registers, (num_registers_duplications, 1))
+ attention_mask_binary = (attention_mask.squeeze(1).squeeze(1).unsqueeze(-1) >= -9000.0).int()
+
+ non_zero_hidden_states = hidden_states[:, attention_mask_binary.squeeze().bool(), :]
+ non_zero_nums = non_zero_hidden_states.shape[1]
+ pad_length = hidden_states.shape[1] - non_zero_nums
+ adjusted_hidden_states = nn.functional.pad(non_zero_hidden_states, pad=(0, 0, 0, pad_length), value=0)
+ flipped_mask = torch.flip(attention_mask_binary, dims=[1])
+ hidden_states = flipped_mask * adjusted_hidden_states + (1 - flipped_mask) * learnable_registers
+
+ attention_mask = torch.full_like(
+ attention_mask,
+ 0.0,
+ dtype=attention_mask.dtype,
+ device=attention_mask.device,
+ )
+
+ return hidden_states, attention_mask
+
+ def forward(
+ self,
+ hidden_states: torch.Tensor,
+ attention_mask: torch.Tensor | None = None,
+ ) -> tuple[torch.Tensor, torch.Tensor]:
+ """
+ Forward pass of Embeddings1DConnector.
+ Args:
+ hidden_states (torch.Tensor): Input tensor of embeddings (shape [batch, seq_len, feature_dim]).
+ attention_mask (torch.Tensor|None): Optional mask for valid tokens (shape compatible with hidden_states).
+ Returns:
+ tuple[torch.Tensor, torch.Tensor]: Processed features and the corresponding (possibly modified) mask.
+ """
+ if self.num_learnable_registers:
+ hidden_states, attention_mask = self._replace_padded_with_learnable_registers(hidden_states, attention_mask)
+
+ indices_grid = torch.arange(hidden_states.shape[1], dtype=torch.float32, device=hidden_states.device)
+ indices_grid = indices_grid[None, None, :]
+ freq_grid_generator = generate_freq_grid_np if self.double_precision_rope else generate_freq_grid_pytorch
+ freqs_cis = precompute_freqs_cis(
+ indices_grid=indices_grid,
+ dim=self.inner_dim,
+ out_dtype=hidden_states.dtype,
+ theta=self.positional_embedding_theta,
+ max_pos=self.positional_embedding_max_pos,
+ num_attention_heads=self.num_attention_heads,
+ rope_type=self.rope_type,
+ freq_grid_generator=freq_grid_generator,
+ )
+
+ for block in self.transformer_1d_blocks:
+ hidden_states = block(hidden_states, attention_mask=attention_mask, pe=freqs_cis)
+
+ hidden_states = rms_norm(hidden_states)
+
+ return hidden_states, attention_mask
+
+
+class LTX2TextEncoderPostModules(nn.Module):
+ def __init__(
+ self,
+ separated_audio_video: bool = False,
+ embedding_dim_gemma: int = 3840,
+ num_layers_gemma: int = 49,
+ video_attention_heads: int = 32,
+ video_attention_head_dim: int = 128,
+ audio_attention_heads: int = 32,
+ audio_attention_head_dim: int = 64,
+ num_connector_layers: int = 2,
+ apply_gated_attention: bool = False,
+ ):
+ super().__init__()
+ if not separated_audio_video:
+ self.feature_extractor_linear = GemmaFeaturesExtractorProjLinear()
+ self.embeddings_connector = Embeddings1DConnector()
+ self.audio_embeddings_connector = Embeddings1DConnector()
+ else:
+ # LTX-2.3
+ self.feature_extractor_linear = GemmaSeperatedFeaturesExtractorProjLinear(
+ num_layers_gemma, embedding_dim_gemma, video_attention_heads * video_attention_head_dim,
+ audio_attention_heads * audio_attention_head_dim)
+ self.embeddings_connector = Embeddings1DConnector(
+ attention_head_dim=video_attention_head_dim,
+ num_attention_heads=video_attention_heads,
+ num_layers=num_connector_layers,
+ apply_gated_attention=apply_gated_attention,
+ )
+ self.audio_embeddings_connector = Embeddings1DConnector(
+ attention_head_dim=audio_attention_head_dim,
+ num_attention_heads=audio_attention_heads,
+ num_layers=num_connector_layers,
+ apply_gated_attention=apply_gated_attention,
+ )
+
+ def create_embeddings(
+ self,
+ video_features: torch.Tensor,
+ audio_features: torch.Tensor | None,
+ additive_attention_mask: torch.Tensor,
+ ) -> tuple[torch.Tensor, torch.Tensor | None, torch.Tensor]:
+ video_encoded, video_mask = self.embeddings_connector(video_features, additive_attention_mask)
+ video_encoded, binary_mask = _to_binary_mask(video_encoded, video_mask)
+ audio_encoded, _ = self.audio_embeddings_connector(audio_features, additive_attention_mask)
+
+ return video_encoded, audio_encoded, binary_mask
+
+ def process_hidden_states(
+ self,
+ hidden_states: tuple[torch.Tensor, ...],
+ attention_mask: torch.Tensor,
+ padding_side: str = "left",
+ ):
+ video_feats, audio_feats = self.feature_extractor_linear(hidden_states, attention_mask, padding_side)
+ additive_mask = _convert_to_additive_mask(attention_mask, video_feats.dtype)
+ video_enc, audio_enc, binary_mask = self.create_embeddings(video_feats, audio_feats, additive_mask)
+ return video_enc, audio_enc, binary_mask
+
+
+def _norm_and_concat_padded_batch(
+ encoded_text: torch.Tensor,
+ sequence_lengths: torch.Tensor,
+ padding_side: str = "right",
+) -> torch.Tensor:
+ """Normalize and flatten multi-layer hidden states, respecting padding.
+ Performs per-batch, per-layer normalization using masked mean and range,
+ then concatenates across the layer dimension.
+ Args:
+ encoded_text: Hidden states of shape [batch, seq_len, hidden_dim, num_layers].
+ sequence_lengths: Number of valid (non-padded) tokens per batch item.
+ padding_side: Whether padding is on "left" or "right".
+ Returns:
+ Normalized tensor of shape [batch, seq_len, hidden_dim * num_layers],
+ with padded positions zeroed out.
+ """
+ b, t, d, l = encoded_text.shape # noqa: E741
+ device = encoded_text.device
+ # Build mask: [B, T, 1, 1]
+ token_indices = torch.arange(t, device=device)[None, :] # [1, T]
+ if padding_side == "right":
+ # For right padding, valid tokens are from 0 to sequence_length-1
+ mask = token_indices < sequence_lengths[:, None] # [B, T]
+ elif padding_side == "left":
+ # For left padding, valid tokens are from (T - sequence_length) to T-1
+ start_indices = t - sequence_lengths[:, None] # [B, 1]
+ mask = token_indices >= start_indices # [B, T]
+ else:
+ raise ValueError(f"padding_side must be 'left' or 'right', got {padding_side}")
+ mask = rearrange(mask, "b t -> b t 1 1")
+ eps = 1e-6
+ # Compute masked mean: [B, 1, 1, L]
+ masked = encoded_text.masked_fill(~mask, 0.0)
+ denom = (sequence_lengths * d).view(b, 1, 1, 1)
+ mean = masked.sum(dim=(1, 2), keepdim=True) / (denom + eps)
+ # Compute masked min/max: [B, 1, 1, L]
+ x_min = encoded_text.masked_fill(~mask, float("inf")).amin(dim=(1, 2), keepdim=True)
+ x_max = encoded_text.masked_fill(~mask, float("-inf")).amax(dim=(1, 2), keepdim=True)
+ range_ = x_max - x_min
+ # Normalize only the valid tokens
+ normed = 8 * (encoded_text - mean) / (range_ + eps)
+ # concat to be [Batch, T, D * L] - this preserves the original structure
+ normed = normed.reshape(b, t, -1) # [B, T, D * L]
+ # Apply mask to preserve original padding (set padded positions to 0)
+ mask_flattened = rearrange(mask, "b t 1 1 -> b t 1").expand(-1, -1, d * l)
+ normed = normed.masked_fill(~mask_flattened, 0.0)
+
+ return normed
+
+
+def _convert_to_additive_mask(attention_mask: torch.Tensor, dtype: torch.dtype) -> torch.Tensor:
+ return (attention_mask - 1).to(dtype).reshape(
+ (attention_mask.shape[0], 1, -1, attention_mask.shape[-1])) * torch.finfo(dtype).max
+
+def _to_binary_mask(encoded: torch.Tensor, encoded_mask: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
+ """Convert connector output mask to binary mask and apply to encoded tensor."""
+ binary_mask = (encoded_mask < 0.000001).to(torch.int64)
+ binary_mask = binary_mask.reshape([encoded.shape[0], encoded.shape[1], 1])
+ encoded = encoded * binary_mask
+ return encoded, binary_mask
+
+
+def norm_and_concat_per_token_rms(
+ encoded_text: torch.Tensor,
+ attention_mask: torch.Tensor,
+) -> torch.Tensor:
+ """Per-token RMSNorm normalization for V2 models.
+ Args:
+ encoded_text: [B, T, D, L]
+ attention_mask: [B, T] binary mask
+ Returns:
+ [B, T, D*L] normalized tensor with padding zeroed out.
+ """
+ B, T, D, L = encoded_text.shape # noqa: N806
+ variance = torch.mean(encoded_text**2, dim=2, keepdim=True) # [B,T,1,L]
+ normed = encoded_text * torch.rsqrt(variance + 1e-6)
+ normed = normed.reshape(B, T, D * L)
+ mask_3d = attention_mask.bool().unsqueeze(-1) # [B, T, 1]
+ return torch.where(mask_3d, normed, torch.zeros_like(normed))
+
+
+def _rescale_norm(x: torch.Tensor, target_dim: int, source_dim: int) -> torch.Tensor:
+ """Rescale normalization: x * sqrt(target_dim / source_dim)."""
+ return x * math.sqrt(target_dim / source_dim)
diff --git a/diffsynth/models/ltx2_upsampler.py b/diffsynth/models/ltx2_upsampler.py
new file mode 100644
index 0000000000000000000000000000000000000000..862ca14bc320880f2b10c933fadc575c8ebad681
--- /dev/null
+++ b/diffsynth/models/ltx2_upsampler.py
@@ -0,0 +1,313 @@
+import math
+from typing import Optional, Tuple
+import torch
+from einops import rearrange
+import torch.nn.functional as F
+from .ltx2_video_vae import LTX2VideoEncoder
+
+class PixelShuffleND(torch.nn.Module):
+ """
+ N-dimensional pixel shuffle operation for upsampling tensors.
+ Args:
+ dims (int): Number of dimensions to apply pixel shuffle to.
+ - 1: Temporal (e.g., frames)
+ - 2: Spatial (e.g., height and width)
+ - 3: Spatiotemporal (e.g., depth, height, width)
+ upscale_factors (tuple[int, int, int], optional): Upscaling factors for each dimension.
+ For dims=1, only the first value is used.
+ For dims=2, the first two values are used.
+ For dims=3, all three values are used.
+ The input tensor is rearranged so that the channel dimension is split into
+ smaller channels and upscaling factors, and the upscaling factors are moved
+ into the corresponding spatial/temporal dimensions.
+ Note:
+ This operation is equivalent to the patchifier operation in for the models. Consider
+ using this class instead.
+ """
+
+ def __init__(self, dims: int, upscale_factors: tuple[int, int, int] = (2, 2, 2)):
+ super().__init__()
+ assert dims in [1, 2, 3], "dims must be 1, 2, or 3"
+ self.dims = dims
+ self.upscale_factors = upscale_factors
+
+ def forward(self, x: torch.Tensor) -> torch.Tensor:
+ if self.dims == 3:
+ return rearrange(
+ x,
+ "b (c p1 p2 p3) d h w -> b c (d p1) (h p2) (w p3)",
+ p1=self.upscale_factors[0],
+ p2=self.upscale_factors[1],
+ p3=self.upscale_factors[2],
+ )
+ elif self.dims == 2:
+ return rearrange(
+ x,
+ "b (c p1 p2) h w -> b c (h p1) (w p2)",
+ p1=self.upscale_factors[0],
+ p2=self.upscale_factors[1],
+ )
+ elif self.dims == 1:
+ return rearrange(
+ x,
+ "b (c p1) f h w -> b c (f p1) h w",
+ p1=self.upscale_factors[0],
+ )
+ else:
+ raise ValueError(f"Unsupported dims: {self.dims}")
+
+
+class ResBlock(torch.nn.Module):
+ """
+ Residual block with two convolutional layers, group normalization, and SiLU activation.
+ Args:
+ channels (int): Number of input and output channels.
+ mid_channels (Optional[int]): Number of channels in the intermediate convolution layer. Defaults to `channels`
+ if not specified.
+ dims (int): Dimensionality of the convolution (2 for Conv2d, 3 for Conv3d). Defaults to 3.
+ """
+
+ def __init__(self, channels: int, mid_channels: Optional[int] = None, dims: int = 3):
+ super().__init__()
+ if mid_channels is None:
+ mid_channels = channels
+
+ conv = torch.nn.Conv2d if dims == 2 else torch.nn.Conv3d
+
+ self.conv1 = conv(channels, mid_channels, kernel_size=3, padding=1)
+ self.norm1 = torch.nn.GroupNorm(32, mid_channels)
+ self.conv2 = conv(mid_channels, channels, kernel_size=3, padding=1)
+ self.norm2 = torch.nn.GroupNorm(32, channels)
+ self.activation = torch.nn.SiLU()
+
+ def forward(self, x: torch.Tensor) -> torch.Tensor:
+ residual = x
+ x = self.conv1(x)
+ x = self.norm1(x)
+ x = self.activation(x)
+ x = self.conv2(x)
+ x = self.norm2(x)
+ x = self.activation(x + residual)
+ return x
+
+
+class BlurDownsample(torch.nn.Module):
+ """
+ Anti-aliased spatial downsampling by integer stride using a fixed separable binomial kernel.
+ Applies only on H,W. Works for dims=2 or dims=3 (per-frame).
+ """
+
+ def __init__(self, dims: int, stride: int, kernel_size: int = 5) -> None:
+ super().__init__()
+ assert dims in (2, 3)
+ assert isinstance(stride, int)
+ assert stride >= 1
+ assert kernel_size >= 3
+ assert kernel_size % 2 == 1
+ self.dims = dims
+ self.stride = stride
+ self.kernel_size = kernel_size
+
+ # 5x5 separable binomial kernel using binomial coefficients [1, 4, 6, 4, 1] from
+ # the 4th row of Pascal's triangle. This kernel is used for anti-aliasing and
+ # provides a smooth approximation of a Gaussian filter (often called a "binomial filter").
+ # The 2D kernel is constructed as the outer product and normalized.
+ k = torch.tensor([math.comb(kernel_size - 1, k) for k in range(kernel_size)])
+ k2d = k[:, None] @ k[None, :]
+ k2d = (k2d / k2d.sum()).float() # shape (kernel_size, kernel_size)
+ self.register_buffer("kernel", k2d[None, None, :, :]) # (1, 1, kernel_size, kernel_size)
+
+ def forward(self, x: torch.Tensor) -> torch.Tensor:
+ if self.stride == 1:
+ return x
+
+ if self.dims == 2:
+ return self._apply_2d(x)
+ else:
+ # dims == 3: apply per-frame on H,W
+ b, _, f, _, _ = x.shape
+ x = rearrange(x, "b c f h w -> (b f) c h w")
+ x = self._apply_2d(x)
+ h2, w2 = x.shape[-2:]
+ x = rearrange(x, "(b f) c h w -> b c f h w", b=b, f=f, h=h2, w=w2)
+ return x
+
+ def _apply_2d(self, x2d: torch.Tensor) -> torch.Tensor:
+ c = x2d.shape[1]
+ weight = self.kernel.expand(c, 1, self.kernel_size, self.kernel_size) # depthwise
+ x2d = F.conv2d(x2d, weight=weight, bias=None, stride=self.stride, padding=self.kernel_size // 2, groups=c)
+ return x2d
+
+
+def _rational_for_scale(scale: float) -> Tuple[int, int]:
+ mapping = {0.75: (3, 4), 1.5: (3, 2), 2.0: (2, 1), 4.0: (4, 1)}
+ if float(scale) not in mapping:
+ raise ValueError(f"Unsupported scale {scale}. Choose from {list(mapping.keys())}")
+ return mapping[float(scale)]
+
+
+class SpatialRationalResampler(torch.nn.Module):
+ """
+ Fully-learned rational spatial scaling: up by 'num' via PixelShuffle, then anti-aliased
+ downsample by 'den' using fixed blur + stride. Operates on H,W only.
+ For dims==3, work per-frame for spatial scaling (temporal axis untouched).
+ Args:
+ mid_channels (`int`): Number of intermediate channels for the convolution layer
+ scale (`float`): Spatial scaling factor. Supported values are:
+ - 0.75: Downsample by 3/4 (reduce spatial size)
+ - 1.5: Upsample by 3/2 (increase spatial size)
+ - 2.0: Upsample by 2x (double spatial size)
+ - 4.0: Upsample by 4x (quadruple spatial size)
+ Any other value will raise a ValueError.
+ """
+
+ def __init__(self, mid_channels: int, scale: float):
+ super().__init__()
+ self.scale = float(scale)
+ self.num, self.den = _rational_for_scale(self.scale)
+ self.conv = torch.nn.Conv2d(mid_channels, (self.num**2) * mid_channels, kernel_size=3, padding=1)
+ self.pixel_shuffle = PixelShuffleND(2, upscale_factors=(self.num, self.num))
+ self.blur_down = BlurDownsample(dims=2, stride=self.den)
+
+ def forward(self, x: torch.Tensor) -> torch.Tensor:
+ b, _, f, _, _ = x.shape
+ x = rearrange(x, "b c f h w -> (b f) c h w")
+ x = self.conv(x)
+ x = self.pixel_shuffle(x)
+ x = self.blur_down(x)
+ x = rearrange(x, "(b f) c h w -> b c f h w", b=b, f=f)
+ return x
+
+
+class LTX2LatentUpsampler(torch.nn.Module):
+ """
+ Model to upsample VAE latents spatially and/or temporally.
+ Args:
+ in_channels (`int`): Number of channels in the input latent
+ mid_channels (`int`): Number of channels in the middle layers
+ num_blocks_per_stage (`int`): Number of ResBlocks to use in each stage (pre/post upsampling)
+ dims (`int`): Number of dimensions for convolutions (2 or 3)
+ spatial_upsample (`bool`): Whether to spatially upsample the latent
+ temporal_upsample (`bool`): Whether to temporally upsample the latent
+ spatial_scale (`float`): Scale factor for spatial upsampling
+ rational_resampler (`bool`): Whether to use a rational resampler for spatial upsampling
+ """
+ def __init__(
+ self,
+ in_channels: int = 128,
+ mid_channels: int = 1024,
+ num_blocks_per_stage: int = 4,
+ dims: int = 3,
+ spatial_upsample: bool = True,
+ temporal_upsample: bool = False,
+ spatial_scale: float = 2.0,
+ rational_resampler: bool = True,
+ ):
+ super().__init__()
+
+ self.in_channels = in_channels
+ self.mid_channels = mid_channels
+ self.num_blocks_per_stage = num_blocks_per_stage
+ self.dims = dims
+ self.spatial_upsample = spatial_upsample
+ self.temporal_upsample = temporal_upsample
+ self.spatial_scale = float(spatial_scale)
+ self.rational_resampler = rational_resampler
+
+ conv = torch.nn.Conv2d if dims == 2 else torch.nn.Conv3d
+
+ self.initial_conv = conv(in_channels, mid_channels, kernel_size=3, padding=1)
+ self.initial_norm = torch.nn.GroupNorm(32, mid_channels)
+ self.initial_activation = torch.nn.SiLU()
+
+ self.res_blocks = torch.nn.ModuleList([ResBlock(mid_channels, dims=dims) for _ in range(num_blocks_per_stage)])
+
+ if spatial_upsample and temporal_upsample:
+ self.upsampler = torch.nn.Sequential(
+ torch.nn.Conv3d(mid_channels, 8 * mid_channels, kernel_size=3, padding=1),
+ PixelShuffleND(3),
+ )
+ elif spatial_upsample:
+ if rational_resampler:
+ self.upsampler = SpatialRationalResampler(mid_channels=mid_channels, scale=self.spatial_scale)
+ else:
+ self.upsampler = torch.nn.Sequential(
+ torch.nn.Conv2d(mid_channels, 4 * mid_channels, kernel_size=3, padding=1),
+ PixelShuffleND(2),
+ )
+ elif temporal_upsample:
+ self.upsampler = torch.nn.Sequential(
+ torch.nn.Conv3d(mid_channels, 2 * mid_channels, kernel_size=3, padding=1),
+ PixelShuffleND(1),
+ )
+ else:
+ raise ValueError("Either spatial_upsample or temporal_upsample must be True")
+
+ self.post_upsample_res_blocks = torch.nn.ModuleList(
+ [ResBlock(mid_channels, dims=dims) for _ in range(num_blocks_per_stage)]
+ )
+
+ self.final_conv = conv(mid_channels, in_channels, kernel_size=3, padding=1)
+
+ def forward(self, latent: torch.Tensor) -> torch.Tensor:
+ b, _, f, _, _ = latent.shape
+
+ if self.dims == 2:
+ x = rearrange(latent, "b c f h w -> (b f) c h w")
+ x = self.initial_conv(x)
+ x = self.initial_norm(x)
+ x = self.initial_activation(x)
+
+ for block in self.res_blocks:
+ x = block(x)
+
+ x = self.upsampler(x)
+
+ for block in self.post_upsample_res_blocks:
+ x = block(x)
+
+ x = self.final_conv(x)
+ x = rearrange(x, "(b f) c h w -> b c f h w", b=b, f=f)
+ else:
+ x = self.initial_conv(latent)
+ x = self.initial_norm(x)
+ x = self.initial_activation(x)
+
+ for block in self.res_blocks:
+ x = block(x)
+
+ if self.temporal_upsample:
+ x = self.upsampler(x)
+ # remove the first frame after upsampling.
+ # This is done because the first frame encodes one pixel frame.
+ x = x[:, :, 1:, :, :]
+ elif isinstance(self.upsampler, SpatialRationalResampler):
+ x = self.upsampler(x)
+ else:
+ x = rearrange(x, "b c f h w -> (b f) c h w")
+ x = self.upsampler(x)
+ x = rearrange(x, "(b f) c h w -> b c f h w", b=b, f=f)
+
+ for block in self.post_upsample_res_blocks:
+ x = block(x)
+
+ x = self.final_conv(x)
+
+ return x
+
+
+def upsample_video(latent: torch.Tensor, video_encoder: LTX2VideoEncoder, upsampler: "LTX2LatentUpsampler") -> torch.Tensor:
+ """
+ Apply upsampling to the latent representation using the provided upsampler,
+ with normalization and un-normalization based on the video encoder's per-channel statistics.
+ Args:
+ latent: Input latent tensor of shape [B, C, F, H, W].
+ video_encoder: VideoEncoder with per_channel_statistics for normalization.
+ upsampler: LTX2LatentUpsampler module to perform upsampling.
+ Returns:
+ torch.Tensor: Upsampled and re-normalized latent tensor.
+ """
+ latent = video_encoder.per_channel_statistics.un_normalize(latent)
+ latent = upsampler(latent)
+ latent = video_encoder.per_channel_statistics.normalize(latent)
+ return latent
diff --git a/diffsynth/models/ltx2_video_vae.py b/diffsynth/models/ltx2_video_vae.py
new file mode 100644
index 0000000000000000000000000000000000000000..a70dc80e830e6910976451247b11334e824f34d4
--- /dev/null
+++ b/diffsynth/models/ltx2_video_vae.py
@@ -0,0 +1,2322 @@
+import itertools
+import math
+import einops
+from dataclasses import replace, dataclass
+from typing import Any, Callable, Iterator, List, NamedTuple, Tuple, Union, Optional
+import torch
+from einops import rearrange
+from torch import nn
+from torch.nn import functional as F
+from enum import Enum
+from .ltx2_common import PixelNorm, SpatioTemporalScaleFactors, VideoLatentShape, Patchifier, AudioLatentShape
+from .ltx2_dit import PixArtAlphaCombinedTimestepSizeEmbeddings
+
+VAE_SPATIAL_FACTOR = 32
+VAE_TEMPORAL_FACTOR = 8
+
+
+class VideoLatentPatchifier(Patchifier):
+ def __init__(self, patch_size: int):
+ # Patch sizes for video latents.
+ self._patch_size = (
+ 1, # temporal dimension
+ patch_size, # height dimension
+ patch_size, # width dimension
+ )
+
+ @property
+ def patch_size(self) -> Tuple[int, int, int]:
+ return self._patch_size
+
+ def get_token_count(self, tgt_shape: VideoLatentShape) -> int:
+ return math.prod(tgt_shape.to_torch_shape()[2:]) // math.prod(self._patch_size)
+
+ def patchify(
+ self,
+ latents: torch.Tensor,
+ ) -> torch.Tensor:
+ latents = einops.rearrange(
+ latents,
+ "b c (f p1) (h p2) (w p3) -> b (f h w) (c p1 p2 p3)",
+ p1=self._patch_size[0],
+ p2=self._patch_size[1],
+ p3=self._patch_size[2],
+ )
+
+ return latents
+
+ def unpatchify(
+ self,
+ latents: torch.Tensor,
+ output_shape: VideoLatentShape,
+ ) -> torch.Tensor:
+ assert self._patch_size[0] == 1, "Temporal patch size must be 1 for symmetric patchifier"
+
+ patch_grid_frames = output_shape.frames // self._patch_size[0]
+ patch_grid_height = output_shape.height // self._patch_size[1]
+ patch_grid_width = output_shape.width // self._patch_size[2]
+
+ latents = einops.rearrange(
+ latents,
+ "b (f h w) (c p q) -> b c f (h p) (w q)",
+ f=patch_grid_frames,
+ h=patch_grid_height,
+ w=patch_grid_width,
+ p=self._patch_size[1],
+ q=self._patch_size[2],
+ )
+
+ return latents
+
+ def unpatchify_video(
+ self,
+ latents: torch.Tensor,
+ frames: int,
+ height: int,
+ width: int,
+ ) -> torch.Tensor:
+ latents = einops.rearrange(
+ latents,
+ "b (f h w) (c p q) -> b c f (h p) (w q)",
+ f=frames,
+ h=height // self._patch_size[1],
+ w=width // self._patch_size[2],
+ p=self._patch_size[1],
+ q=self._patch_size[2],
+ )
+ return latents
+
+ def get_patch_grid_bounds(
+ self,
+ output_shape: AudioLatentShape | VideoLatentShape,
+ device: Optional[torch.device] = None,
+ ) -> torch.Tensor:
+ """
+ Return the per-dimension bounds [inclusive start, exclusive end) for every
+ patch produced by `patchify`. The bounds are expressed in the original
+ video grid coordinates: frame/time, height, and width.
+ The resulting tensor is shaped `[batch_size, 3, num_patches, 2]`, where:
+ - axis 1 (size 3) enumerates (frame/time, height, width) dimensions
+ - axis 3 (size 2) stores `[start, end)` indices within each dimension
+ Args:
+ output_shape: Video grid description containing frames, height, and width.
+ device: Device of the latent tensor.
+ """
+ if not isinstance(output_shape, VideoLatentShape):
+ raise ValueError("VideoLatentPatchifier expects VideoLatentShape when computing coordinates")
+
+ frames = output_shape.frames
+ height = output_shape.height
+ width = output_shape.width
+ batch_size = output_shape.batch
+
+ # Validate inputs to ensure positive dimensions
+ assert frames > 0, f"frames must be positive, got {frames}"
+ assert height > 0, f"height must be positive, got {height}"
+ assert width > 0, f"width must be positive, got {width}"
+ assert batch_size > 0, f"batch_size must be positive, got {batch_size}"
+
+ # Generate grid coordinates for each dimension (frame, height, width)
+ # We use torch.arange to create the starting coordinates for each patch.
+ # indexing='ij' ensures the dimensions are in the order (frame, height, width).
+ grid_coords = torch.meshgrid(
+ torch.arange(start=0, end=frames, step=self._patch_size[0], device=device),
+ torch.arange(start=0, end=height, step=self._patch_size[1], device=device),
+ torch.arange(start=0, end=width, step=self._patch_size[2], device=device),
+ indexing="ij",
+ )
+
+ # Stack the grid coordinates to create the start coordinates tensor.
+ # Shape becomes (3, grid_f, grid_h, grid_w)
+ patch_starts = torch.stack(grid_coords, dim=0)
+
+ # Create a tensor containing the size of a single patch:
+ # (frame_patch_size, height_patch_size, width_patch_size).
+ # Reshape to (3, 1, 1, 1) to enable broadcasting when adding to the start coordinates.
+ patch_size_delta = torch.tensor(
+ self._patch_size,
+ device=patch_starts.device,
+ dtype=patch_starts.dtype,
+ ).view(3, 1, 1, 1)
+
+ # Calculate end coordinates: start + patch_size
+ # Shape becomes (3, grid_f, grid_h, grid_w)
+ patch_ends = patch_starts + patch_size_delta
+
+ # Stack start and end coordinates together along the last dimension
+ # Shape becomes (3, grid_f, grid_h, grid_w, 2), where the last dimension is [start, end]
+ latent_coords = torch.stack((patch_starts, patch_ends), dim=-1)
+
+ # Broadcast to batch size and flatten all spatial/temporal dimensions into one sequence.
+ # Final Shape: (batch_size, 3, num_patches, 2)
+ latent_coords = einops.repeat(
+ latent_coords,
+ "c f h w bounds -> b c (f h w) bounds",
+ b=batch_size,
+ bounds=2,
+ )
+
+ return latent_coords
+
+
+class NormLayerType(Enum):
+ GROUP_NORM = "group_norm"
+ PIXEL_NORM = "pixel_norm"
+
+
+class LogVarianceType(Enum):
+ PER_CHANNEL = "per_channel"
+ UNIFORM = "uniform"
+ CONSTANT = "constant"
+ NONE = "none"
+
+
+class PaddingModeType(Enum):
+ ZEROS = "zeros"
+ REFLECT = "reflect"
+ REPLICATE = "replicate"
+ CIRCULAR = "circular"
+
+
+class DualConv3d(nn.Module):
+
+ def __init__(
+ self,
+ in_channels: int,
+ out_channels: int,
+ kernel_size: int,
+ stride: Union[int, Tuple[int, int, int]] = 1,
+ padding: Union[int, Tuple[int, int, int]] = 0,
+ dilation: Union[int, Tuple[int, int, int]] = 1,
+ groups: int = 1,
+ bias: bool = True,
+ padding_mode: str = "zeros",
+ ) -> None:
+ super(DualConv3d, self).__init__()
+
+ self.in_channels = in_channels
+ self.out_channels = out_channels
+ self.padding_mode = padding_mode
+ # Ensure kernel_size, stride, padding, and dilation are tuples of length 3
+ if isinstance(kernel_size, int):
+ kernel_size = (kernel_size, kernel_size, kernel_size)
+ if kernel_size == (1, 1, 1):
+ raise ValueError("kernel_size must be greater than 1. Use make_linear_nd instead.")
+ if isinstance(stride, int):
+ stride = (stride, stride, stride)
+ if isinstance(padding, int):
+ padding = (padding, padding, padding)
+ if isinstance(dilation, int):
+ dilation = (dilation, dilation, dilation)
+
+ # Set parameters for convolutions
+ self.groups = groups
+ self.bias = bias
+
+ # Define the size of the channels after the first convolution
+ intermediate_channels = out_channels if in_channels < out_channels else in_channels
+
+ # Define parameters for the first convolution
+ self.weight1 = nn.Parameter(
+ torch.Tensor(
+ intermediate_channels,
+ in_channels // groups,
+ 1,
+ kernel_size[1],
+ kernel_size[2],
+ ))
+ self.stride1 = (1, stride[1], stride[2])
+ self.padding1 = (0, padding[1], padding[2])
+ self.dilation1 = (1, dilation[1], dilation[2])
+ if bias:
+ self.bias1 = nn.Parameter(torch.Tensor(intermediate_channels))
+ else:
+ self.register_parameter("bias1", None)
+
+ # Define parameters for the second convolution
+ self.weight2 = nn.Parameter(torch.Tensor(out_channels, intermediate_channels // groups, kernel_size[0], 1, 1))
+ self.stride2 = (stride[0], 1, 1)
+ self.padding2 = (padding[0], 0, 0)
+ self.dilation2 = (dilation[0], 1, 1)
+ if bias:
+ self.bias2 = nn.Parameter(torch.Tensor(out_channels))
+ else:
+ self.register_parameter("bias2", None)
+
+ # Initialize weights and biases
+ self.reset_parameters()
+
+ def reset_parameters(self) -> None:
+ nn.init.kaiming_uniform_(self.weight1, a=torch.sqrt(5))
+ nn.init.kaiming_uniform_(self.weight2, a=torch.sqrt(5))
+ if self.bias:
+ fan_in1, _ = nn.init._calculate_fan_in_and_fan_out(self.weight1)
+ bound1 = 1 / torch.sqrt(fan_in1)
+ nn.init.uniform_(self.bias1, -bound1, bound1)
+ fan_in2, _ = nn.init._calculate_fan_in_and_fan_out(self.weight2)
+ bound2 = 1 / torch.sqrt(fan_in2)
+ nn.init.uniform_(self.bias2, -bound2, bound2)
+
+ def forward(
+ self,
+ x: torch.Tensor,
+ use_conv3d: bool = False,
+ skip_time_conv: bool = False,
+ ) -> torch.Tensor:
+ if use_conv3d:
+ return self.forward_with_3d(x=x, skip_time_conv=skip_time_conv)
+ else:
+ return self.forward_with_2d(x=x, skip_time_conv=skip_time_conv)
+
+ def forward_with_3d(self, x: torch.Tensor, skip_time_conv: bool = False) -> torch.Tensor:
+ # First convolution
+ x = F.conv3d(
+ x,
+ self.weight1,
+ self.bias1,
+ self.stride1,
+ self.padding1,
+ self.dilation1,
+ self.groups,
+ padding_mode=self.padding_mode,
+ )
+
+ if skip_time_conv:
+ return x
+
+ # Second convolution
+ x = F.conv3d(
+ x,
+ self.weight2,
+ self.bias2,
+ self.stride2,
+ self.padding2,
+ self.dilation2,
+ self.groups,
+ padding_mode=self.padding_mode,
+ )
+
+ return x
+
+ def forward_with_2d(self, x: torch.Tensor, skip_time_conv: bool = False) -> torch.Tensor:
+ b, _, _, h, w = x.shape
+
+ # First 2D convolution
+ x = rearrange(x, "b c d h w -> (b d) c h w")
+ # Squeeze the depth dimension out of weight1 since it's 1
+ weight1 = self.weight1.squeeze(2)
+ # Select stride, padding, and dilation for the 2D convolution
+ stride1 = (self.stride1[1], self.stride1[2])
+ padding1 = (self.padding1[1], self.padding1[2])
+ dilation1 = (self.dilation1[1], self.dilation1[2])
+ x = F.conv2d(
+ x,
+ weight1,
+ self.bias1,
+ stride1,
+ padding1,
+ dilation1,
+ self.groups,
+ padding_mode=self.padding_mode,
+ )
+
+ _, _, h, w = x.shape
+
+ if skip_time_conv:
+ x = rearrange(x, "(b d) c h w -> b c d h w", b=b)
+ return x
+
+ # Second convolution which is essentially treated as a 1D convolution across the 'd' dimension
+ x = rearrange(x, "(b d) c h w -> (b h w) c d", b=b)
+
+ # Reshape weight2 to match the expected dimensions for conv1d
+ weight2 = self.weight2.squeeze(-1).squeeze(-1)
+ # Use only the relevant dimension for stride, padding, and dilation for the 1D convolution
+ stride2 = self.stride2[0]
+ padding2 = self.padding2[0]
+ dilation2 = self.dilation2[0]
+ x = F.conv1d(
+ x,
+ weight2,
+ self.bias2,
+ stride2,
+ padding2,
+ dilation2,
+ self.groups,
+ padding_mode=self.padding_mode,
+ )
+ x = rearrange(x, "(b h w) c d -> b c d h w", b=b, h=h, w=w)
+
+ return x
+
+ @property
+ def weight(self) -> torch.Tensor:
+ return self.weight2
+
+
+class CausalConv3d(nn.Module):
+
+ def __init__(
+ self,
+ in_channels: int,
+ out_channels: int,
+ kernel_size: int = 3,
+ stride: Union[int, Tuple[int]] = 1,
+ dilation: int = 1,
+ groups: int = 1,
+ bias: bool = True,
+ spatial_padding_mode: PaddingModeType = PaddingModeType.ZEROS,
+ ) -> None:
+ super().__init__()
+
+ self.in_channels = in_channels
+ self.out_channels = out_channels
+
+ kernel_size = (kernel_size, kernel_size, kernel_size)
+ self.time_kernel_size = kernel_size[0]
+
+ dilation = (dilation, 1, 1)
+
+ height_pad = kernel_size[1] // 2
+ width_pad = kernel_size[2] // 2
+ padding = (0, height_pad, width_pad)
+
+ self.conv = nn.Conv3d(
+ in_channels,
+ out_channels,
+ kernel_size,
+ stride=stride,
+ dilation=dilation,
+ padding=padding,
+ padding_mode=spatial_padding_mode.value,
+ groups=groups,
+ bias=bias,
+ )
+
+ def forward(self, x: torch.Tensor, causal: bool = True) -> torch.Tensor:
+ if causal:
+ first_frame_pad = x[:, :, :1, :, :].repeat((1, 1, self.time_kernel_size - 1, 1, 1))
+ x = torch.concatenate((first_frame_pad, x), dim=2)
+ else:
+ first_frame_pad = x[:, :, :1, :, :].repeat((1, 1, (self.time_kernel_size - 1) // 2, 1, 1))
+ last_frame_pad = x[:, :, -1:, :, :].repeat((1, 1, (self.time_kernel_size - 1) // 2, 1, 1))
+ x = torch.concatenate((first_frame_pad, x, last_frame_pad), dim=2)
+ x = self.conv(x)
+ return x
+
+ @property
+ def weight(self) -> torch.Tensor:
+ return self.conv.weight
+
+
+def make_conv_nd( # noqa: PLR0913
+ dims: Union[int, Tuple[int, int]],
+ in_channels: int,
+ out_channels: int,
+ kernel_size: int,
+ stride: int = 1,
+ padding: int = 0,
+ dilation: int = 1,
+ groups: int = 1,
+ bias: bool = True,
+ causal: bool = False,
+ spatial_padding_mode: PaddingModeType = PaddingModeType.ZEROS,
+ temporal_padding_mode: PaddingModeType = PaddingModeType.ZEROS,
+) -> nn.Module:
+ if not (spatial_padding_mode == temporal_padding_mode or causal):
+ raise NotImplementedError("spatial and temporal padding modes must be equal")
+ if dims == 2:
+ return nn.Conv2d(
+ in_channels=in_channels,
+ out_channels=out_channels,
+ kernel_size=kernel_size,
+ stride=stride,
+ padding=padding,
+ dilation=dilation,
+ groups=groups,
+ bias=bias,
+ padding_mode=spatial_padding_mode.value,
+ )
+ elif dims == 3:
+ if causal:
+ return CausalConv3d(
+ in_channels=in_channels,
+ out_channels=out_channels,
+ kernel_size=kernel_size,
+ stride=stride,
+ dilation=dilation,
+ groups=groups,
+ bias=bias,
+ spatial_padding_mode=spatial_padding_mode,
+ )
+ return nn.Conv3d(
+ in_channels=in_channels,
+ out_channels=out_channels,
+ kernel_size=kernel_size,
+ stride=stride,
+ padding=padding,
+ dilation=dilation,
+ groups=groups,
+ bias=bias,
+ padding_mode=spatial_padding_mode.value,
+ )
+ elif dims == (2, 1):
+ return DualConv3d(
+ in_channels=in_channels,
+ out_channels=out_channels,
+ kernel_size=kernel_size,
+ stride=stride,
+ padding=padding,
+ bias=bias,
+ padding_mode=spatial_padding_mode.value,
+ )
+ else:
+ raise ValueError(f"unsupported dimensions: {dims}")
+
+
+def make_linear_nd(
+ dims: int,
+ in_channels: int,
+ out_channels: int,
+ bias: bool = True,
+) -> nn.Module:
+ if dims == 2:
+ return nn.Conv2d(in_channels=in_channels, out_channels=out_channels, kernel_size=1, bias=bias)
+ elif dims in (3, (2, 1)):
+ return nn.Conv3d(in_channels=in_channels, out_channels=out_channels, kernel_size=1, bias=bias)
+ else:
+ raise ValueError(f"unsupported dimensions: {dims}")
+
+
+def patchify(x: torch.Tensor, patch_size_hw: int, patch_size_t: int = 1) -> torch.Tensor:
+ """
+ Rearrange spatial dimensions into channels. Divides image into patch_size x patch_size blocks
+ and moves pixels from each block into separate channels (space-to-depth).
+ Args:
+ x: Input tensor (4D or 5D)
+ patch_size_hw: Spatial patch size for height and width. With patch_size_hw=4, divides HxW into 4x4 blocks.
+ patch_size_t: Temporal patch size for frames. Default=1 (no temporal patching).
+ For 5D: (B, C, F, H, W) -> (B, Cx(patch_size_hw^2)x(patch_size_t), F/patch_size_t, H/patch_size_hw, W/patch_size_hw)
+ Example: (B, 3, 33, 512, 512) with patch_size_hw=4, patch_size_t=1 -> (B, 48, 33, 128, 128)
+ """
+ if patch_size_hw == 1 and patch_size_t == 1:
+ return x
+ if x.dim() == 4:
+ x = rearrange(x, "b c (h q) (w r) -> b (c r q) h w", q=patch_size_hw, r=patch_size_hw)
+ elif x.dim() == 5:
+ x = rearrange(
+ x,
+ "b c (f p) (h q) (w r) -> b (c p r q) f h w",
+ p=patch_size_t,
+ q=patch_size_hw,
+ r=patch_size_hw,
+ )
+ else:
+ raise ValueError(f"Invalid input shape: {x.shape}")
+
+ return x
+
+
+def unpatchify(x: torch.Tensor, patch_size_hw: int, patch_size_t: int = 1) -> torch.Tensor:
+ """
+ Rearrange channels back into spatial dimensions. Inverse of patchify - moves pixels from
+ channels back into patch_size x patch_size blocks (depth-to-space).
+ Args:
+ x: Input tensor (4D or 5D)
+ patch_size_hw: Spatial patch size for height and width. With patch_size_hw=4, expands HxW by 4x.
+ patch_size_t: Temporal patch size for frames. Default=1 (no temporal expansion).
+ For 5D: (B, Cx(patch_size_hw^2)x(patch_size_t), F, H, W) -> (B, C, Fxpatch_size_t, Hxpatch_size_hw, Wxpatch_size_hw)
+ Example: (B, 48, 33, 128, 128) with patch_size_hw=4, patch_size_t=1 -> (B, 3, 33, 512, 512)
+ """
+ if patch_size_hw == 1 and patch_size_t == 1:
+ return x
+
+ if x.dim() == 4:
+ x = rearrange(x, "b (c r q) h w -> b c (h q) (w r)", q=patch_size_hw, r=patch_size_hw)
+ elif x.dim() == 5:
+ x = rearrange(
+ x,
+ "b (c p r q) f h w -> b c (f p) (h q) (w r)",
+ p=patch_size_t,
+ q=patch_size_hw,
+ r=patch_size_hw,
+ )
+
+ return x
+
+
+class PerChannelStatistics(nn.Module):
+ """
+ Per-channel statistics for normalizing and denormalizing the latent representation.
+ This statics is computed over the entire dataset and stored in model's checkpoint under VAE state_dict.
+ """
+
+ def __init__(self, latent_channels: int = 128):
+ super().__init__()
+ self.register_buffer("std-of-means", torch.empty(latent_channels))
+ self.register_buffer("mean-of-means", torch.empty(latent_channels))
+
+ def un_normalize(self, x: torch.Tensor) -> torch.Tensor:
+ return (x * self.get_buffer("std-of-means").view(1, -1, 1, 1, 1).to(x)) + self.get_buffer("mean-of-means").view(
+ 1, -1, 1, 1, 1).to(x)
+
+ def normalize(self, x: torch.Tensor) -> torch.Tensor:
+ return (x - self.get_buffer("mean-of-means").view(1, -1, 1, 1, 1).to(x)) / self.get_buffer("std-of-means").view(
+ 1, -1, 1, 1, 1).to(x)
+
+
+class ResnetBlock3D(nn.Module):
+ r"""
+ A Resnet block.
+ Parameters:
+ in_channels (`int`): The number of channels in the input.
+ out_channels (`int`, *optional*, default to be `None`):
+ The number of output channels for the first conv layer. If None, same as `in_channels`.
+ dropout (`float`, *optional*, defaults to `0.0`): The dropout probability to use.
+ groups (`int`, *optional*, default to `32`): The number of groups to use for the first normalization layer.
+ eps (`float`, *optional*, defaults to `1e-6`): The epsilon to use for the normalization.
+ """
+
+ def __init__(
+ self,
+ dims: Union[int, Tuple[int, int]],
+ in_channels: int,
+ out_channels: Optional[int] = None,
+ dropout: float = 0.0,
+ groups: int = 32,
+ eps: float = 1e-6,
+ norm_layer: NormLayerType = NormLayerType.PIXEL_NORM,
+ inject_noise: bool = False,
+ timestep_conditioning: bool = False,
+ spatial_padding_mode: PaddingModeType = PaddingModeType.ZEROS,
+ ):
+ super().__init__()
+ self.in_channels = in_channels
+ out_channels = in_channels if out_channels is None else out_channels
+ self.out_channels = out_channels
+ self.inject_noise = inject_noise
+
+ if norm_layer == NormLayerType.GROUP_NORM:
+ self.norm1 = nn.GroupNorm(num_groups=groups, num_channels=in_channels, eps=eps, affine=True)
+ elif norm_layer == NormLayerType.PIXEL_NORM:
+ self.norm1 = PixelNorm()
+
+ self.non_linearity = nn.SiLU()
+
+ self.conv1 = make_conv_nd(
+ dims,
+ in_channels,
+ out_channels,
+ kernel_size=3,
+ stride=1,
+ padding=1,
+ causal=True,
+ spatial_padding_mode=spatial_padding_mode,
+ )
+
+ if inject_noise:
+ self.per_channel_scale1 = nn.Parameter(torch.zeros((in_channels, 1, 1)))
+
+ if norm_layer == NormLayerType.GROUP_NORM:
+ self.norm2 = nn.GroupNorm(num_groups=groups, num_channels=out_channels, eps=eps, affine=True)
+ elif norm_layer == NormLayerType.PIXEL_NORM:
+ self.norm2 = PixelNorm()
+
+ self.dropout = torch.nn.Dropout(dropout)
+
+ self.conv2 = make_conv_nd(
+ dims,
+ out_channels,
+ out_channels,
+ kernel_size=3,
+ stride=1,
+ padding=1,
+ causal=True,
+ spatial_padding_mode=spatial_padding_mode,
+ )
+
+ if inject_noise:
+ self.per_channel_scale2 = nn.Parameter(torch.zeros((in_channels, 1, 1)))
+
+ self.conv_shortcut = (make_linear_nd(dims=dims, in_channels=in_channels, out_channels=out_channels)
+ if in_channels != out_channels else nn.Identity())
+
+ # Using GroupNorm with 1 group is equivalent to LayerNorm but works with (B, C, ...) layout
+ # avoiding the need for dimension rearrangement used in standard nn.LayerNorm
+ self.norm3 = (nn.GroupNorm(num_groups=1, num_channels=in_channels, eps=eps, affine=True)
+ if in_channels != out_channels else nn.Identity())
+
+ self.timestep_conditioning = timestep_conditioning
+
+ if timestep_conditioning:
+ self.scale_shift_table = nn.Parameter(torch.zeros(4, in_channels))
+
+ def _feed_spatial_noise(
+ self,
+ hidden_states: torch.Tensor,
+ per_channel_scale: torch.Tensor,
+ generator: Optional[torch.Generator] = None,
+ ) -> torch.Tensor:
+ spatial_shape = hidden_states.shape[-2:]
+ device = hidden_states.device
+ dtype = hidden_states.dtype
+
+ # similar to the "explicit noise inputs" method in style-gan
+ spatial_noise = torch.randn(spatial_shape, device=device, dtype=dtype, generator=generator)[None]
+ scaled_noise = (spatial_noise * per_channel_scale)[None, :, None, ...]
+ hidden_states = hidden_states + scaled_noise
+
+ return hidden_states
+
+ def forward(
+ self,
+ input_tensor: torch.Tensor,
+ causal: bool = True,
+ timestep: Optional[torch.Tensor] = None,
+ generator: Optional[torch.Generator] = None,
+ ) -> torch.Tensor:
+ hidden_states = input_tensor
+ batch_size = hidden_states.shape[0]
+
+ hidden_states = self.norm1(hidden_states)
+ if self.timestep_conditioning:
+ if timestep is None:
+ raise ValueError("'timestep' parameter must be provided when 'timestep_conditioning' is True")
+ ada_values = self.scale_shift_table[None, ..., None, None, None].to(
+ device=hidden_states.device, dtype=hidden_states.dtype) + timestep.reshape(
+ batch_size,
+ 4,
+ -1,
+ timestep.shape[-3],
+ timestep.shape[-2],
+ timestep.shape[-1],
+ )
+ shift1, scale1, shift2, scale2 = ada_values.unbind(dim=1)
+
+ hidden_states = hidden_states * (1 + scale1) + shift1
+
+ hidden_states = self.non_linearity(hidden_states)
+
+ hidden_states = self.conv1(hidden_states, causal=causal)
+
+ if self.inject_noise:
+ hidden_states = self._feed_spatial_noise(
+ hidden_states,
+ self.per_channel_scale1.to(device=hidden_states.device, dtype=hidden_states.dtype),
+ generator=generator,
+ )
+
+ hidden_states = self.norm2(hidden_states)
+
+ if self.timestep_conditioning:
+ hidden_states = hidden_states * (1 + scale2) + shift2
+
+ hidden_states = self.non_linearity(hidden_states)
+
+ hidden_states = self.dropout(hidden_states)
+
+ hidden_states = self.conv2(hidden_states, causal=causal)
+
+ if self.inject_noise:
+ hidden_states = self._feed_spatial_noise(
+ hidden_states,
+ self.per_channel_scale2.to(device=hidden_states.device, dtype=hidden_states.dtype),
+ generator=generator,
+ )
+
+ input_tensor = self.norm3(input_tensor)
+
+ batch_size = input_tensor.shape[0]
+
+ input_tensor = self.conv_shortcut(input_tensor)
+
+ output_tensor = input_tensor + hidden_states
+
+ return output_tensor
+
+
+class UNetMidBlock3D(nn.Module):
+ """
+ A 3D UNet mid-block [`UNetMidBlock3D`] with multiple residual blocks.
+ Args:
+ in_channels (`int`): The number of input channels.
+ dropout (`float`, *optional*, defaults to 0.0): The dropout rate.
+ num_layers (`int`, *optional*, defaults to 1): The number of residual blocks.
+ resnet_eps (`float`, *optional*, 1e-6 ): The epsilon value for the resnet blocks.
+ resnet_groups (`int`, *optional*, defaults to 32):
+ The number of groups to use in the group normalization layers of the resnet blocks.
+ norm_layer (`str`, *optional*, defaults to `group_norm`):
+ The normalization layer to use. Can be either `group_norm` or `pixel_norm`.
+ inject_noise (`bool`, *optional*, defaults to `False`):
+ Whether to inject noise into the hidden states.
+ timestep_conditioning (`bool`, *optional*, defaults to `False`):
+ Whether to condition the hidden states on the timestep.
+ Returns:
+ `torch.Tensor`: The output of the last residual block, which is a tensor of shape `(batch_size,
+ in_channels, height, width)`.
+ """
+
+ def __init__(
+ self,
+ dims: Union[int, Tuple[int, int]],
+ in_channels: int,
+ dropout: float = 0.0,
+ num_layers: int = 1,
+ resnet_eps: float = 1e-6,
+ resnet_groups: int = 32,
+ norm_layer: NormLayerType = NormLayerType.GROUP_NORM,
+ inject_noise: bool = False,
+ timestep_conditioning: bool = False,
+ spatial_padding_mode: PaddingModeType = PaddingModeType.ZEROS,
+ ):
+ super().__init__()
+ resnet_groups = resnet_groups if resnet_groups is not None else min(in_channels // 4, 32)
+
+ self.timestep_conditioning = timestep_conditioning
+
+ if timestep_conditioning:
+ self.time_embedder = PixArtAlphaCombinedTimestepSizeEmbeddings(embedding_dim=in_channels * 4,
+ size_emb_dim=0)
+
+ self.res_blocks = nn.ModuleList([
+ ResnetBlock3D(
+ dims=dims,
+ in_channels=in_channels,
+ out_channels=in_channels,
+ eps=resnet_eps,
+ groups=resnet_groups,
+ dropout=dropout,
+ norm_layer=norm_layer,
+ inject_noise=inject_noise,
+ timestep_conditioning=timestep_conditioning,
+ spatial_padding_mode=spatial_padding_mode,
+ ) for _ in range(num_layers)
+ ])
+
+ def forward(
+ self,
+ hidden_states: torch.Tensor,
+ causal: bool = True,
+ timestep: Optional[torch.Tensor] = None,
+ generator: Optional[torch.Generator] = None,
+ ) -> torch.Tensor:
+ timestep_embed = None
+ if self.timestep_conditioning:
+ if timestep is None:
+ raise ValueError("'timestep' parameter must be provided when 'timestep_conditioning' is True")
+ batch_size = hidden_states.shape[0]
+ timestep_embed = self.time_embedder(
+ timestep=timestep.flatten(),
+ hidden_dtype=hidden_states.dtype,
+ )
+ timestep_embed = timestep_embed.view(batch_size, timestep_embed.shape[-1], 1, 1, 1)
+
+ for resnet in self.res_blocks:
+ hidden_states = resnet(
+ hidden_states,
+ causal=causal,
+ timestep=timestep_embed,
+ generator=generator,
+ )
+
+ return hidden_states
+
+
+class SpaceToDepthDownsample(nn.Module):
+
+ def __init__(
+ self,
+ dims: Union[int, Tuple[int, int]],
+ in_channels: int,
+ out_channels: int,
+ stride: Tuple[int, int, int],
+ spatial_padding_mode: PaddingModeType = PaddingModeType.ZEROS,
+ ):
+ super().__init__()
+ self.stride = stride
+ self.group_size = in_channels * math.prod(stride) // out_channels
+ self.conv = make_conv_nd(
+ dims=dims,
+ in_channels=in_channels,
+ out_channels=out_channels // math.prod(stride),
+ kernel_size=3,
+ stride=1,
+ causal=True,
+ spatial_padding_mode=spatial_padding_mode,
+ )
+
+ def forward(
+ self,
+ x: torch.Tensor,
+ causal: bool = True,
+ ) -> torch.Tensor:
+ if self.stride[0] == 2:
+ x = torch.cat([x[:, :, :1, :, :], x], dim=2) # duplicate first frames for padding
+
+ # skip connection
+ x_in = rearrange(
+ x,
+ "b c (d p1) (h p2) (w p3) -> b (c p1 p2 p3) d h w",
+ p1=self.stride[0],
+ p2=self.stride[1],
+ p3=self.stride[2],
+ )
+ x_in = rearrange(x_in, "b (c g) d h w -> b c g d h w", g=self.group_size)
+ x_in = x_in.mean(dim=2)
+
+ # conv
+ x = self.conv(x, causal=causal)
+ x = rearrange(
+ x,
+ "b c (d p1) (h p2) (w p3) -> b (c p1 p2 p3) d h w",
+ p1=self.stride[0],
+ p2=self.stride[1],
+ p3=self.stride[2],
+ )
+
+ x = x + x_in
+
+ return x
+
+
+class DepthToSpaceUpsample(nn.Module):
+
+ def __init__(
+ self,
+ dims: int | Tuple[int, int],
+ in_channels: int,
+ stride: Tuple[int, int, int],
+ residual: bool = False,
+ out_channels_reduction_factor: int = 1,
+ spatial_padding_mode: PaddingModeType = PaddingModeType.ZEROS,
+ ):
+ super().__init__()
+ self.stride = stride
+ self.out_channels = math.prod(stride) * in_channels // out_channels_reduction_factor
+ self.conv = make_conv_nd(
+ dims=dims,
+ in_channels=in_channels,
+ out_channels=self.out_channels,
+ kernel_size=3,
+ stride=1,
+ causal=True,
+ spatial_padding_mode=spatial_padding_mode,
+ )
+ self.residual = residual
+ self.out_channels_reduction_factor = out_channels_reduction_factor
+
+ def forward(
+ self,
+ x: torch.Tensor,
+ causal: bool = True,
+ ) -> torch.Tensor:
+ if self.residual:
+ # Reshape and duplicate the input to match the output shape
+ x_in = rearrange(
+ x,
+ "b (c p1 p2 p3) d h w -> b c (d p1) (h p2) (w p3)",
+ p1=self.stride[0],
+ p2=self.stride[1],
+ p3=self.stride[2],
+ )
+ num_repeat = math.prod(self.stride) // self.out_channels_reduction_factor
+ x_in = x_in.repeat(1, num_repeat, 1, 1, 1)
+ if self.stride[0] == 2:
+ x_in = x_in[:, :, 1:, :, :]
+ x = self.conv(x, causal=causal)
+ x = rearrange(
+ x,
+ "b (c p1 p2 p3) d h w -> b c (d p1) (h p2) (w p3)",
+ p1=self.stride[0],
+ p2=self.stride[1],
+ p3=self.stride[2],
+ )
+ if self.stride[0] == 2:
+ x = x[:, :, 1:, :, :]
+ if self.residual:
+ x = x + x_in
+ return x
+
+
+def compute_trapezoidal_mask_1d(
+ length: int,
+ ramp_left: int,
+ ramp_right: int,
+ left_starts_from_0: bool = False,
+) -> torch.Tensor:
+ """
+ Generate a 1D trapezoidal blending mask with linear ramps.
+ Args:
+ length: Output length of the mask.
+ ramp_left: Fade-in length on the left.
+ ramp_right: Fade-out length on the right.
+ left_starts_from_0: Whether the ramp starts from 0 or first non-zero value.
+ Useful for temporal tiles where the first tile is causal.
+ Returns:
+ A 1D tensor of shape `(length,)` with values in [0, 1].
+ """
+ if length <= 0:
+ raise ValueError("Mask length must be positive.")
+
+ ramp_left = max(0, min(ramp_left, length))
+ ramp_right = max(0, min(ramp_right, length))
+
+ mask = torch.ones(length)
+
+ if ramp_left > 0:
+ interval_length = ramp_left + 1 if left_starts_from_0 else ramp_left + 2
+ fade_in = torch.linspace(0.0, 1.0, interval_length)[:-1]
+ if not left_starts_from_0:
+ fade_in = fade_in[1:]
+ mask[:ramp_left] *= fade_in
+
+ if ramp_right > 0:
+ fade_out = torch.linspace(1.0, 0.0, steps=ramp_right + 2)[1:-1]
+ mask[-ramp_right:] *= fade_out
+
+ return mask.clamp_(0, 1)
+
+
+@dataclass(frozen=True)
+class SpatialTilingConfig:
+ """Configuration for dividing each frame into spatial tiles with optional overlap.
+ Args:
+ tile_size_in_pixels (int): Size of each tile in pixels. Must be at least 64 and divisible by 32.
+ tile_overlap_in_pixels (int, optional): Overlap between tiles in pixels. Must be divisible by 32. Defaults to 0.
+ """
+
+ tile_size_in_pixels: int
+ tile_overlap_in_pixels: int = 0
+
+ def __post_init__(self) -> None:
+ if self.tile_size_in_pixels < 64:
+ raise ValueError(f"tile_size_in_pixels must be at least 64, got {self.tile_size_in_pixels}")
+ if self.tile_size_in_pixels % 32 != 0:
+ raise ValueError(f"tile_size_in_pixels must be divisible by 32, got {self.tile_size_in_pixels}")
+ if self.tile_overlap_in_pixels % 32 != 0:
+ raise ValueError(f"tile_overlap_in_pixels must be divisible by 32, got {self.tile_overlap_in_pixels}")
+ if self.tile_overlap_in_pixels >= self.tile_size_in_pixels:
+ raise ValueError(
+ f"Overlap must be less than tile size, got {self.tile_overlap_in_pixels} and {self.tile_size_in_pixels}"
+ )
+
+
+@dataclass(frozen=True)
+class TemporalTilingConfig:
+ """Configuration for dividing a video into temporal tiles (chunks of frames) with optional overlap.
+ Args:
+ tile_size_in_frames (int): Number of frames in each tile. Must be at least 16 and divisible by 8.
+ tile_overlap_in_frames (int, optional): Number of overlapping frames between consecutive tiles.
+ Must be divisible by 8. Defaults to 0.
+ """
+
+ tile_size_in_frames: int
+ tile_overlap_in_frames: int = 0
+
+ def __post_init__(self) -> None:
+ if self.tile_size_in_frames < 16:
+ raise ValueError(f"tile_size_in_frames must be at least 16, got {self.tile_size_in_frames}")
+ if self.tile_size_in_frames % 8 != 0:
+ raise ValueError(f"tile_size_in_frames must be divisible by 8, got {self.tile_size_in_frames}")
+ if self.tile_overlap_in_frames % 8 != 0:
+ raise ValueError(f"tile_overlap_in_frames must be divisible by 8, got {self.tile_overlap_in_frames}")
+ if self.tile_overlap_in_frames >= self.tile_size_in_frames:
+ raise ValueError(
+ f"Overlap must be less than tile size, got {self.tile_overlap_in_frames} and {self.tile_size_in_frames}"
+ )
+
+
+@dataclass(frozen=True)
+class TilingConfig:
+ """Configuration for splitting video into tiles with optional overlap.
+ Attributes:
+ spatial_config: Configuration for splitting spatial dimensions into tiles.
+ temporal_config: Configuration for splitting temporal dimension into tiles.
+ """
+
+ spatial_config: SpatialTilingConfig | None = None
+ temporal_config: TemporalTilingConfig | None = None
+
+ @classmethod
+ def default(cls) -> "TilingConfig":
+ return cls(
+ spatial_config=SpatialTilingConfig(tile_size_in_pixels=512, tile_overlap_in_pixels=64),
+ temporal_config=TemporalTilingConfig(tile_size_in_frames=64, tile_overlap_in_frames=24),
+ )
+
+
+@dataclass(frozen=True)
+class DimensionIntervals:
+ """Intervals which a single dimension of the latent space is split into.
+ Each interval is defined by its start, end, left ramp, and right ramp.
+ The start and end are the indices of the first and last element (exclusive) in the interval.
+ Ramps are regions of the interval where the value of the mask tensor is
+ interpolated between 0 and 1 for blending with neighboring intervals.
+ The left ramp and right ramp values are the lengths of the left and right ramps.
+ """
+
+ starts: List[int]
+ ends: List[int]
+ left_ramps: List[int]
+ right_ramps: List[int]
+
+
+@dataclass(frozen=True)
+class LatentIntervals:
+ """Intervals which the latent tensor of given shape is split into.
+ Each dimension of the latent space is split into intervals based on the length along said dimension.
+ """
+
+ original_shape: torch.Size
+ dimension_intervals: Tuple[DimensionIntervals, ...]
+
+
+# Operation to split a single dimension of the tensor into intervals based on the length along the dimension.
+SplitOperation = Callable[[int], DimensionIntervals]
+# Operation to map the intervals in input dimension to slices and masks along a corresponding output dimension.
+MappingOperation = Callable[[DimensionIntervals], tuple[list[slice], list[torch.Tensor | None]]]
+
+
+def default_split_operation(length: int) -> DimensionIntervals:
+ return DimensionIntervals(starts=[0], ends=[length], left_ramps=[0], right_ramps=[0])
+
+
+DEFAULT_SPLIT_OPERATION: SplitOperation = default_split_operation
+
+
+def default_mapping_operation(_intervals: DimensionIntervals,) -> tuple[list[slice], list[torch.Tensor | None]]:
+ return [slice(0, None)], [None]
+
+
+DEFAULT_MAPPING_OPERATION: MappingOperation = default_mapping_operation
+
+
+class Tile(NamedTuple):
+ """
+ Represents a single tile.
+ Attributes:
+ in_coords:
+ Tuple of slices specifying where to cut the tile from the INPUT tensor.
+ out_coords:
+ Tuple of slices specifying where this tile's OUTPUT should be placed in the reconstructed OUTPUT tensor.
+ masks_1d:
+ Per-dimension masks in OUTPUT units.
+ These are used to create all-dimensional blending mask.
+ Methods:
+ blend_mask:
+ Create a single N-D mask from the per-dimension masks.
+ """
+
+ in_coords: Tuple[slice, ...]
+ out_coords: Tuple[slice, ...]
+ masks_1d: Tuple[Tuple[torch.Tensor, ...]]
+
+ @property
+ def blend_mask(self) -> torch.Tensor:
+ num_dims = len(self.out_coords)
+ per_dimension_masks: List[torch.Tensor] = []
+
+ for dim_idx in range(num_dims):
+ mask_1d = self.masks_1d[dim_idx]
+ view_shape = [1] * num_dims
+ if mask_1d is None:
+ # Broadcast mask along this dimension (length 1).
+ one = torch.ones(1)
+
+ view_shape[dim_idx] = 1
+ per_dimension_masks.append(one.view(*view_shape))
+ continue
+
+ # Reshape (L,) -> (1, ..., L, ..., 1) so masks across dimensions broadcast-multiply.
+ view_shape[dim_idx] = mask_1d.shape[0]
+ per_dimension_masks.append(mask_1d.view(*view_shape))
+
+ # Multiply per-dimension masks to form the full N-D mask (separable blending window).
+ combined_mask = per_dimension_masks[0]
+ for mask in per_dimension_masks[1:]:
+ combined_mask = combined_mask * mask
+
+ return combined_mask
+
+
+def create_tiles_from_intervals_and_mappers(
+ intervals: LatentIntervals,
+ mappers: List[MappingOperation],
+) -> List[Tile]:
+ full_dim_input_slices = []
+ full_dim_output_slices = []
+ full_dim_masks_1d = []
+ for axis_index in range(len(intervals.original_shape)):
+ dimension_intervals = intervals.dimension_intervals[axis_index]
+ starts = dimension_intervals.starts
+ ends = dimension_intervals.ends
+ input_slices = [slice(s, e) for s, e in zip(starts, ends, strict=True)]
+ output_slices, masks_1d = mappers[axis_index](dimension_intervals)
+ full_dim_input_slices.append(input_slices)
+ full_dim_output_slices.append(output_slices)
+ full_dim_masks_1d.append(masks_1d)
+
+ tiles = []
+ tile_in_coords = list(itertools.product(*full_dim_input_slices))
+ tile_out_coords = list(itertools.product(*full_dim_output_slices))
+ tile_mask_1ds = list(itertools.product(*full_dim_masks_1d))
+ for in_coord, out_coord, mask_1d in zip(tile_in_coords, tile_out_coords, tile_mask_1ds, strict=True):
+ tiles.append(Tile(
+ in_coords=in_coord,
+ out_coords=out_coord,
+ masks_1d=mask_1d,
+ ))
+ return tiles
+
+
+def create_tiles(
+ latent_shape: torch.Size,
+ splitters: List[SplitOperation],
+ mappers: List[MappingOperation],
+) -> List[Tile]:
+ if len(splitters) != len(latent_shape):
+ raise ValueError(f"Number of splitters must be equal to number of dimensions in latent shape, "
+ f"got {len(splitters)} and {len(latent_shape)}")
+ if len(mappers) != len(latent_shape):
+ raise ValueError(f"Number of mappers must be equal to number of dimensions in latent shape, "
+ f"got {len(mappers)} and {len(latent_shape)}")
+ intervals = [splitter(length) for splitter, length in zip(splitters, latent_shape, strict=True)]
+ latent_intervals = LatentIntervals(original_shape=latent_shape, dimension_intervals=tuple(intervals))
+ return create_tiles_from_intervals_and_mappers(latent_intervals, mappers)
+
+
+def _make_encoder_block(
+ block_name: str,
+ block_config: dict[str, Any],
+ in_channels: int,
+ convolution_dimensions: int,
+ norm_layer: NormLayerType,
+ norm_num_groups: int,
+ spatial_padding_mode: PaddingModeType,
+) -> Tuple[nn.Module, int]:
+ out_channels = in_channels
+
+ if block_name == "res_x":
+ block = UNetMidBlock3D(
+ dims=convolution_dimensions,
+ in_channels=in_channels,
+ num_layers=block_config["num_layers"],
+ resnet_eps=1e-6,
+ resnet_groups=norm_num_groups,
+ norm_layer=norm_layer,
+ spatial_padding_mode=spatial_padding_mode,
+ )
+ elif block_name == "res_x_y":
+ out_channels = in_channels * block_config.get("multiplier", 2)
+ block = ResnetBlock3D(
+ dims=convolution_dimensions,
+ in_channels=in_channels,
+ out_channels=out_channels,
+ eps=1e-6,
+ groups=norm_num_groups,
+ norm_layer=norm_layer,
+ spatial_padding_mode=spatial_padding_mode,
+ )
+ elif block_name == "compress_time":
+ block = make_conv_nd(
+ dims=convolution_dimensions,
+ in_channels=in_channels,
+ out_channels=out_channels,
+ kernel_size=3,
+ stride=(2, 1, 1),
+ causal=True,
+ spatial_padding_mode=spatial_padding_mode,
+ )
+ elif block_name == "compress_space":
+ block = make_conv_nd(
+ dims=convolution_dimensions,
+ in_channels=in_channels,
+ out_channels=out_channels,
+ kernel_size=3,
+ stride=(1, 2, 2),
+ causal=True,
+ spatial_padding_mode=spatial_padding_mode,
+ )
+ elif block_name == "compress_all":
+ block = make_conv_nd(
+ dims=convolution_dimensions,
+ in_channels=in_channels,
+ out_channels=out_channels,
+ kernel_size=3,
+ stride=(2, 2, 2),
+ causal=True,
+ spatial_padding_mode=spatial_padding_mode,
+ )
+ elif block_name == "compress_all_x_y":
+ out_channels = in_channels * block_config.get("multiplier", 2)
+ block = make_conv_nd(
+ dims=convolution_dimensions,
+ in_channels=in_channels,
+ out_channels=out_channels,
+ kernel_size=3,
+ stride=(2, 2, 2),
+ causal=True,
+ spatial_padding_mode=spatial_padding_mode,
+ )
+ elif block_name == "compress_all_res":
+ out_channels = in_channels * block_config.get("multiplier", 2)
+ block = SpaceToDepthDownsample(
+ dims=convolution_dimensions,
+ in_channels=in_channels,
+ out_channels=out_channels,
+ stride=(2, 2, 2),
+ spatial_padding_mode=spatial_padding_mode,
+ )
+ elif block_name == "compress_space_res":
+ out_channels = in_channels * block_config.get("multiplier", 2)
+ block = SpaceToDepthDownsample(
+ dims=convolution_dimensions,
+ in_channels=in_channels,
+ out_channels=out_channels,
+ stride=(1, 2, 2),
+ spatial_padding_mode=spatial_padding_mode,
+ )
+ elif block_name == "compress_time_res":
+ out_channels = in_channels * block_config.get("multiplier", 2)
+ block = SpaceToDepthDownsample(
+ dims=convolution_dimensions,
+ in_channels=in_channels,
+ out_channels=out_channels,
+ stride=(2, 1, 1),
+ spatial_padding_mode=spatial_padding_mode,
+ )
+ else:
+ raise ValueError(f"unknown block: {block_name}")
+
+ return block, out_channels
+
+
+class LTX2VideoEncoder(nn.Module):
+ _DEFAULT_NORM_NUM_GROUPS = 32
+ """
+ Variational Autoencoder Encoder. Encodes video frames into a latent representation.
+ The encoder compresses the input video through a series of downsampling operations controlled by
+ patch_size and encoder_blocks. The output is a normalized latent tensor with shape (B, 128, F', H', W').
+ Compression Behavior:
+ The total compression is determined by:
+ 1. Initial spatial compression via patchify: H -> H/4, W -> W/4 (patch_size=4)
+ 2. Sequential compression through encoder_blocks based on their stride patterns
+ Compression blocks apply 2x compression in specified dimensions:
+ - "compress_time" / "compress_time_res": temporal only
+ - "compress_space" / "compress_space_res": spatial only (H and W)
+ - "compress_all" / "compress_all_res": all dimensions (F, H, W)
+ - "res_x" / "res_x_y": no compression
+ Standard LTX Video configuration:
+ - patch_size=4
+ - encoder_blocks: 1x compress_space_res, 1x compress_time_res, 2x compress_all_res
+ - Final dimensions: F' = 1 + (F-1)/8, H' = H/32, W' = W/32
+ - Example: (B, 3, 33, 512, 512) -> (B, 128, 5, 16, 16)
+ - Note: Input must have 1 + 8*k frames (e.g., 1, 9, 17, 25, 33...)
+ Args:
+ convolution_dimensions: The number of dimensions to use in convolutions (2D or 3D).
+ in_channels: The number of input channels. For RGB images, this is 3.
+ out_channels: The number of output channels (latent channels). For latent channels, this is 128.
+ encoder_blocks: The list of blocks to construct the encoder. Each block is a tuple of (block_name, params)
+ where params is either an int (num_layers) or a dict with configuration.
+ patch_size: The patch size for initial spatial compression. Should be a power of 2.
+ norm_layer: The normalization layer to use. Can be either `group_norm` or `pixel_norm`.
+ latent_log_var: The log variance mode. Can be either `per_channel`, `uniform`, `constant` or `none`.
+ """
+
+ def __init__(
+ self,
+ convolution_dimensions: int = 3,
+ in_channels: int = 3,
+ out_channels: int = 128,
+ patch_size: int = 4,
+ norm_layer: NormLayerType = NormLayerType.PIXEL_NORM,
+ latent_log_var: LogVarianceType = LogVarianceType.UNIFORM,
+ encoder_spatial_padding_mode: PaddingModeType = PaddingModeType.ZEROS,
+ encoder_version: str = "ltx-2",
+ ):
+ super().__init__()
+ if encoder_version == "ltx-2":
+ encoder_blocks = [
+ ['res_x', {'num_layers': 4}],
+ ['compress_space_res', {'multiplier': 2}],
+ ['res_x', {'num_layers': 6}],
+ ['compress_time_res', {'multiplier': 2}],
+ ['res_x', {'num_layers': 6}],
+ ['compress_all_res', {'multiplier': 2}],
+ ['res_x', {'num_layers': 2}],
+ ['compress_all_res', {'multiplier': 2}],
+ ['res_x', {'num_layers': 2}]
+ ]
+ else:
+ # LTX-2.3
+ encoder_blocks = [
+ ["res_x", {"num_layers": 4}],
+ ["compress_space_res", {"multiplier": 2}],
+ ["res_x", {"num_layers": 6}],
+ ["compress_time_res", {"multiplier": 2}],
+ ["res_x", {"num_layers": 4}],
+ ["compress_all_res", {"multiplier": 2}],
+ ["res_x", {"num_layers": 2}],
+ ["compress_all_res", {"multiplier": 1}],
+ ["res_x", {"num_layers": 2}]
+ ]
+ self.patch_size = patch_size
+ self.norm_layer = norm_layer
+ self.latent_channels = out_channels
+ self.latent_log_var = latent_log_var
+ self._norm_num_groups = self._DEFAULT_NORM_NUM_GROUPS
+
+ # Per-channel statistics for normalizing latents
+ self.per_channel_statistics = PerChannelStatistics(latent_channels=out_channels)
+
+ in_channels = in_channels * patch_size**2
+ feature_channels = out_channels
+
+ self.conv_in = make_conv_nd(
+ dims=convolution_dimensions,
+ in_channels=in_channels,
+ out_channels=feature_channels,
+ kernel_size=3,
+ stride=1,
+ padding=1,
+ causal=True,
+ spatial_padding_mode=encoder_spatial_padding_mode,
+ )
+
+ self.down_blocks = nn.ModuleList([])
+
+ for block_name, block_params in encoder_blocks:
+ # Convert int to dict format for uniform handling
+ block_config = {"num_layers": block_params} if isinstance(block_params, int) else block_params
+
+ block, feature_channels = _make_encoder_block(
+ block_name=block_name,
+ block_config=block_config,
+ in_channels=feature_channels,
+ convolution_dimensions=convolution_dimensions,
+ norm_layer=norm_layer,
+ norm_num_groups=self._norm_num_groups,
+ spatial_padding_mode=encoder_spatial_padding_mode,
+ )
+
+ self.down_blocks.append(block)
+
+ # out
+ if norm_layer == NormLayerType.GROUP_NORM:
+ self.conv_norm_out = nn.GroupNorm(num_channels=feature_channels, num_groups=self._norm_num_groups, eps=1e-6)
+ elif norm_layer == NormLayerType.PIXEL_NORM:
+ self.conv_norm_out = PixelNorm()
+
+ self.conv_act = nn.SiLU()
+
+ conv_out_channels = out_channels
+ if latent_log_var == LogVarianceType.PER_CHANNEL:
+ conv_out_channels *= 2
+ elif latent_log_var in {LogVarianceType.UNIFORM, LogVarianceType.CONSTANT}:
+ conv_out_channels += 1
+ elif latent_log_var != LogVarianceType.NONE:
+ raise ValueError(f"Invalid latent_log_var: {latent_log_var}")
+
+ self.conv_out = make_conv_nd(
+ dims=convolution_dimensions,
+ in_channels=feature_channels,
+ out_channels=conv_out_channels,
+ kernel_size=3,
+ padding=1,
+ causal=True,
+ spatial_padding_mode=encoder_spatial_padding_mode,
+ )
+
+ def forward(self, sample: torch.Tensor) -> torch.Tensor:
+ r"""
+ Encode video frames into normalized latent representation.
+ Args:
+ sample: Input video (B, C, F, H, W). F must be 1 + 8*k (e.g., 1, 9, 17, 25, 33...).
+ Returns:
+ Normalized latent means (B, 128, F', H', W') where F' = 1+(F-1)/8, H' = H/32, W' = W/32.
+ Example: (B, 3, 33, 512, 512) -> (B, 128, 5, 16, 16).
+ """
+ # Validate frame count
+ frames_count = sample.shape[2]
+ if ((frames_count - 1) % 8) != 0:
+ frames_to_crop = (frames_count - 1) % 8
+ sample = sample[:, :, :-frames_to_crop, ...]
+
+ # Initial spatial compression: trade spatial resolution for channel depth
+ # This reduces H,W by patch_size and increases channels, making convolutions more efficient
+ # Example: (B, 3, F, 512, 512) -> (B, 48, F, 128, 128) with patch_size=4
+ sample = patchify(sample, patch_size_hw=self.patch_size, patch_size_t=1)
+ sample = self.conv_in(sample)
+
+ for down_block in self.down_blocks:
+ sample = down_block(sample)
+
+ sample = self.conv_norm_out(sample)
+ sample = self.conv_act(sample)
+ sample = self.conv_out(sample)
+
+ if self.latent_log_var == LogVarianceType.UNIFORM:
+ # Uniform Variance: model outputs N means and 1 shared log-variance channel.
+ # We need to expand the single logvar to match the number of means channels
+ # to create a format compatible with PER_CHANNEL (means + logvar, each with N channels).
+ # Sample shape: (B, N+1, ...) where N = latent_channels (e.g., 128 means + 1 logvar = 129)
+ # Target shape: (B, 2*N, ...) where first N are means, last N are logvar
+
+ if sample.shape[1] < 2:
+ raise ValueError(f"Invalid channel count for UNIFORM mode: expected at least 2 channels "
+ f"(N means + 1 logvar), got {sample.shape[1]}")
+
+ # Extract means (first N channels) and logvar (last 1 channel)
+ means = sample[:, :-1, ...] # (B, N, ...)
+ logvar = sample[:, -1:, ...] # (B, 1, ...)
+
+ # Repeat logvar N times to match means channels
+ # Use expand/repeat pattern that works for both 4D and 5D tensors
+ num_channels = means.shape[1]
+ repeat_shape = [1, num_channels] + [1] * (sample.ndim - 2)
+ repeated_logvar = logvar.repeat(*repeat_shape) # (B, N, ...)
+
+ # Concatenate to create (B, 2*N, ...) format: [means, repeated_logvar]
+ sample = torch.cat([means, repeated_logvar], dim=1)
+ elif self.latent_log_var == LogVarianceType.CONSTANT:
+ sample = sample[:, :-1, ...]
+ approx_ln_0 = -30 # this is the minimal clamp value in DiagonalGaussianDistribution objects
+ sample = torch.cat(
+ [sample, torch.ones_like(sample, device=sample.device) * approx_ln_0],
+ dim=1,
+ )
+
+ # Split into means and logvar, then normalize means
+ means, _ = torch.chunk(sample, 2, dim=1)
+ return self.per_channel_statistics.normalize(means)
+
+
+ def tiled_encode_video(
+ self,
+ video: torch.Tensor,
+ tile_size: int = 512,
+ tile_overlap: int = 128,
+ ) -> torch.Tensor:
+ """Encode video using spatial tiling for memory efficiency.
+ Splits the video into overlapping spatial tiles, encodes each tile separately,
+ and blends the results using linear feathering in the overlap regions.
+ Args:
+ video: Input tensor of shape [B, C, F, H, W]
+ tile_size: Tile size in pixels (must be divisible by 32)
+ tile_overlap: Overlap between tiles in pixels (must be divisible by 32)
+ Returns:
+ Encoded latent tensor [B, C_latent, F_latent, H_latent, W_latent]
+ """
+ batch, _channels, frames, height, width = video.shape
+ device = video.device
+ dtype = video.dtype
+
+ # Validate tile parameters
+ if tile_size % VAE_SPATIAL_FACTOR != 0:
+ raise ValueError(f"tile_size must be divisible by {VAE_SPATIAL_FACTOR}, got {tile_size}")
+ if tile_overlap % VAE_SPATIAL_FACTOR != 0:
+ raise ValueError(f"tile_overlap must be divisible by {VAE_SPATIAL_FACTOR}, got {tile_overlap}")
+ if tile_overlap >= tile_size:
+ raise ValueError(f"tile_overlap ({tile_overlap}) must be less than tile_size ({tile_size})")
+
+ # If video fits in a single tile, use regular encoding
+ if height <= tile_size and width <= tile_size:
+ return self.forward(video)
+
+ # Calculate output dimensions
+ # VAE compresses: H -> H/32, W -> W/32, F -> 1 + (F-1)/8
+ output_height = height // VAE_SPATIAL_FACTOR
+ output_width = width // VAE_SPATIAL_FACTOR
+ output_frames = 1 + (frames - 1) // VAE_TEMPORAL_FACTOR
+
+ # Latent channels (128 for LTX-2)
+ # Get from a small test encode or assume 128
+ latent_channels = 128
+
+ # Initialize output and weight tensors
+ output = torch.zeros(
+ (batch, latent_channels, output_frames, output_height, output_width),
+ device=device,
+ dtype=dtype,
+ )
+ weights = torch.zeros(
+ (batch, 1, output_frames, output_height, output_width),
+ device=device,
+ dtype=dtype,
+ )
+
+ # Calculate tile positions with overlap
+ # Step size is tile_size - tile_overlap
+ step_h = tile_size - tile_overlap
+ step_w = tile_size - tile_overlap
+
+ h_positions = list(range(0, max(1, height - tile_overlap), step_h))
+ w_positions = list(range(0, max(1, width - tile_overlap), step_w))
+
+ # Ensure last tile covers the edge
+ if h_positions[-1] + tile_size < height:
+ h_positions.append(height - tile_size)
+ if w_positions[-1] + tile_size < width:
+ w_positions.append(width - tile_size)
+
+ # Remove duplicates and sort
+ h_positions = sorted(set(h_positions))
+ w_positions = sorted(set(w_positions))
+
+ # Overlap in latent space
+ overlap_out_h = tile_overlap // VAE_SPATIAL_FACTOR
+ overlap_out_w = tile_overlap // VAE_SPATIAL_FACTOR
+
+ # Process each tile
+ for h_pos in h_positions:
+ for w_pos in w_positions:
+ # Calculate tile boundaries in input space
+ h_start = max(0, h_pos)
+ w_start = max(0, w_pos)
+ h_end = min(h_start + tile_size, height)
+ w_end = min(w_start + tile_size, width)
+
+ # Ensure tile dimensions are divisible by VAE_SPATIAL_FACTOR
+ tile_h = ((h_end - h_start) // VAE_SPATIAL_FACTOR) * VAE_SPATIAL_FACTOR
+ tile_w = ((w_end - w_start) // VAE_SPATIAL_FACTOR) * VAE_SPATIAL_FACTOR
+
+ if tile_h < VAE_SPATIAL_FACTOR or tile_w < VAE_SPATIAL_FACTOR:
+ continue
+
+ # Adjust end positions
+ h_end = h_start + tile_h
+ w_end = w_start + tile_w
+
+ # Extract tile
+ tile = video[:, :, :, h_start:h_end, w_start:w_end]
+
+ # Encode tile
+ encoded_tile = self.forward(tile)
+
+ # Get actual encoded dimensions
+ _, _, tile_out_frames, tile_out_height, tile_out_width = encoded_tile.shape
+
+ # Calculate output positions
+ out_h_start = h_start // VAE_SPATIAL_FACTOR
+ out_w_start = w_start // VAE_SPATIAL_FACTOR
+ out_h_end = min(out_h_start + tile_out_height, output_height)
+ out_w_end = min(out_w_start + tile_out_width, output_width)
+
+ # Trim encoded tile if necessary
+ actual_tile_h = out_h_end - out_h_start
+ actual_tile_w = out_w_end - out_w_start
+ encoded_tile = encoded_tile[:, :, :, :actual_tile_h, :actual_tile_w]
+
+ # Create blending mask with linear feathering at edges
+ mask = torch.ones(
+ (1, 1, tile_out_frames, actual_tile_h, actual_tile_w),
+ device=device,
+ dtype=dtype,
+ )
+
+ # Apply feathering at edges (linear blend in overlap regions)
+ # Left edge
+ if h_pos > 0 and overlap_out_h > 0 and overlap_out_h < actual_tile_h:
+ fade_in = torch.linspace(0.0, 1.0, overlap_out_h + 2, device=device, dtype=dtype)[1:-1]
+ mask[:, :, :, :overlap_out_h, :] *= fade_in.view(1, 1, 1, -1, 1)
+
+ # Right edge (bottom in height dimension)
+ if h_end < height and overlap_out_h > 0 and overlap_out_h < actual_tile_h:
+ fade_out = torch.linspace(1.0, 0.0, overlap_out_h + 2, device=device, dtype=dtype)[1:-1]
+ mask[:, :, :, -overlap_out_h:, :] *= fade_out.view(1, 1, 1, -1, 1)
+
+ # Top edge (left in width dimension)
+ if w_pos > 0 and overlap_out_w > 0 and overlap_out_w < actual_tile_w:
+ fade_in = torch.linspace(0.0, 1.0, overlap_out_w + 2, device=device, dtype=dtype)[1:-1]
+ mask[:, :, :, :, :overlap_out_w] *= fade_in.view(1, 1, 1, 1, -1)
+
+ # Bottom edge (right in width dimension)
+ if w_end < width and overlap_out_w > 0 and overlap_out_w < actual_tile_w:
+ fade_out = torch.linspace(1.0, 0.0, overlap_out_w + 2, device=device, dtype=dtype)[1:-1]
+ mask[:, :, :, :, -overlap_out_w:] *= fade_out.view(1, 1, 1, 1, -1)
+
+ # Accumulate weighted results
+ output[:, :, :, out_h_start:out_h_end, out_w_start:out_w_end] += encoded_tile * mask
+ weights[:, :, :, out_h_start:out_h_end, out_w_start:out_w_end] += mask
+
+ # Normalize by weights (avoid division by zero)
+ output = output / (weights + 1e-8)
+
+ return output
+
+ def encode(
+ self,
+ video: torch.Tensor,
+ tiled=False,
+ tile_size_in_pixels: Optional[int] = 512,
+ tile_overlap_in_pixels: Optional[int] = 128,
+ **kwargs,
+ ) -> torch.Tensor:
+ if video.ndim == 4:
+ video = video.unsqueeze(0) # [C, F, H, W] -> [B, C, F, H, W]
+ # Choose encoding method based on tiling flag
+ if tiled:
+ latents = self.tiled_encode_video(
+ video=video,
+ tile_size=tile_size_in_pixels,
+ tile_overlap=tile_overlap_in_pixels,
+ )
+ else:
+ # Encode video - VAE expects [B, C, F, H, W], returns [B, C, F', H', W']
+ latents = self.forward(video)
+ return latents
+
+
+def _make_decoder_block(
+ block_name: str,
+ block_config: dict[str, Any],
+ in_channels: int,
+ convolution_dimensions: int,
+ norm_layer: NormLayerType,
+ timestep_conditioning: bool,
+ norm_num_groups: int,
+ spatial_padding_mode: PaddingModeType,
+) -> Tuple[nn.Module, int]:
+ out_channels = in_channels
+ if block_name == "res_x":
+ block = UNetMidBlock3D(
+ dims=convolution_dimensions,
+ in_channels=in_channels,
+ num_layers=block_config["num_layers"],
+ resnet_eps=1e-6,
+ resnet_groups=norm_num_groups,
+ norm_layer=norm_layer,
+ inject_noise=block_config.get("inject_noise", False),
+ timestep_conditioning=timestep_conditioning,
+ spatial_padding_mode=spatial_padding_mode,
+ )
+ elif block_name == "attn_res_x":
+ block = UNetMidBlock3D(
+ dims=convolution_dimensions,
+ in_channels=in_channels,
+ num_layers=block_config["num_layers"],
+ resnet_groups=norm_num_groups,
+ norm_layer=norm_layer,
+ inject_noise=block_config.get("inject_noise", False),
+ timestep_conditioning=timestep_conditioning,
+ attention_head_dim=block_config["attention_head_dim"],
+ spatial_padding_mode=spatial_padding_mode,
+ )
+ elif block_name == "res_x_y":
+ out_channels = in_channels // block_config.get("multiplier", 2)
+ block = ResnetBlock3D(
+ dims=convolution_dimensions,
+ in_channels=in_channels,
+ out_channels=out_channels,
+ eps=1e-6,
+ groups=norm_num_groups,
+ norm_layer=norm_layer,
+ inject_noise=block_config.get("inject_noise", False),
+ timestep_conditioning=False,
+ spatial_padding_mode=spatial_padding_mode,
+ )
+ elif block_name == "compress_time":
+ out_channels = in_channels // block_config.get("multiplier", 1)
+ block = DepthToSpaceUpsample(
+ dims=convolution_dimensions,
+ in_channels=in_channels,
+ stride=(2, 1, 1),
+ out_channels_reduction_factor=block_config.get("multiplier", 1),
+ spatial_padding_mode=spatial_padding_mode,
+ )
+ elif block_name == "compress_space":
+ out_channels = in_channels // block_config.get("multiplier", 1)
+ block = DepthToSpaceUpsample(
+ dims=convolution_dimensions,
+ in_channels=in_channels,
+ stride=(1, 2, 2),
+ out_channels_reduction_factor=block_config.get("multiplier", 1),
+ spatial_padding_mode=spatial_padding_mode,
+ )
+ elif block_name == "compress_all":
+ out_channels = in_channels // block_config.get("multiplier", 1)
+ block = DepthToSpaceUpsample(
+ dims=convolution_dimensions,
+ in_channels=in_channels,
+ stride=(2, 2, 2),
+ residual=block_config.get("residual", False),
+ out_channels_reduction_factor=block_config.get("multiplier", 1),
+ spatial_padding_mode=spatial_padding_mode,
+ )
+ else:
+ raise ValueError(f"unknown layer: {block_name}")
+
+ return block, out_channels
+
+
+class LTX2VideoDecoder(nn.Module):
+ _DEFAULT_NORM_NUM_GROUPS = 32
+ """
+ Variational Autoencoder Decoder. Decodes latent representation into video frames.
+ The decoder upsamples latents through a series of upsampling operations (inverse of encoder).
+ Output dimensions: F = 8x(F'-1) + 1, H = 32xH', W = 32xW' for standard LTX Video configuration.
+ Upsampling blocks expand dimensions by 2x in specified dimensions:
+ - "compress_time": temporal only
+ - "compress_space": spatial only (H and W)
+ - "compress_all": all dimensions (F, H, W)
+ - "res_x" / "res_x_y" / "attn_res_x": no upsampling
+ Causal Mode:
+ causal=False (standard): Symmetric padding, allows future frame dependencies.
+ causal=True: Causal padding, each frame depends only on past/current frames.
+ First frame removed after temporal upsampling in both modes. Output shape unchanged.
+ Example: (B, 128, 5, 16, 16) -> (B, 3, 33, 512, 512) for both modes.
+ Args:
+ convolution_dimensions: The number of dimensions to use in convolutions (2D or 3D).
+ in_channels: The number of input channels (latent channels). Default is 128.
+ out_channels: The number of output channels. For RGB images, this is 3.
+ decoder_blocks: The list of blocks to construct the decoder. Each block is a tuple of (block_name, params)
+ where params is either an int (num_layers) or a dict with configuration.
+ patch_size: Final spatial expansion factor. For standard LTX Video, use 4 for 4x spatial expansion:
+ H -> Hx4, W -> Wx4. Should be a power of 2.
+ norm_layer: The normalization layer to use. Can be either `group_norm` or `pixel_norm`.
+ causal: Whether to use causal convolutions. For standard LTX Video, use False for symmetric padding.
+ When True, uses causal padding (past/current frames only).
+ timestep_conditioning: Whether to condition the decoder on timestep for denoising.
+ """
+
+ def __init__(
+ self,
+ convolution_dimensions: int = 3,
+ in_channels: int = 128,
+ out_channels: int = 3,
+ decoder_blocks: List[Tuple[str, int | dict]] = [], # noqa: B006
+ patch_size: int = 4,
+ norm_layer: NormLayerType = NormLayerType.PIXEL_NORM,
+ causal: bool = False,
+ timestep_conditioning: bool = False,
+ decoder_spatial_padding_mode: PaddingModeType = PaddingModeType.REFLECT,
+ decoder_version: str = "ltx-2",
+ base_channels: int = 128,
+ ):
+ super().__init__()
+
+ # Spatiotemporal downscaling between decoded video space and VAE latents.
+ # According to the LTXV paper, the standard configuration downsamples
+ # video inputs by a factor of 8 in the temporal dimension and 32 in
+ # each spatial dimension (height and width). This parameter determines how
+ # many video frames and pixels correspond to a single latent cell.
+ if decoder_version == "ltx-2":
+ decoder_blocks = [
+ ['res_x', {'num_layers': 5, 'inject_noise': False}],
+ ['compress_all', {'residual': True, 'multiplier': 2}],
+ ['res_x', {'num_layers': 5, 'inject_noise': False}],
+ ['compress_all', {'residual': True, 'multiplier': 2}],
+ ['res_x', {'num_layers': 5, 'inject_noise': False}],
+ ['compress_all', {'residual': True, 'multiplier': 2}],
+ ['res_x', {'num_layers': 5, 'inject_noise': False}]
+ ]
+ else:
+ # LTX-2.3
+ decoder_blocks = [
+ ["res_x", {"num_layers": 4}],
+ ["compress_space", {"multiplier": 2}],
+ ["res_x", {"num_layers": 6}],
+ ["compress_time", {"multiplier": 2}],
+ ["res_x", {"num_layers": 4}],
+ ["compress_all", {"multiplier": 1}],
+ ["res_x", {"num_layers": 2}],
+ ["compress_all", {"multiplier": 2}],
+ ["res_x", {"num_layers": 2}]
+ ]
+ self.video_downscale_factors = SpatioTemporalScaleFactors(
+ time=8,
+ width=32,
+ height=32,
+ )
+
+ self.patch_size = patch_size
+ out_channels = out_channels * patch_size**2
+ self.causal = causal
+ self.timestep_conditioning = timestep_conditioning
+ self._norm_num_groups = self._DEFAULT_NORM_NUM_GROUPS
+
+ # Per-channel statistics for denormalizing latents
+ self.per_channel_statistics = PerChannelStatistics(latent_channels=in_channels)
+
+ # Noise and timestep parameters for decoder conditioning
+ self.decode_noise_scale = 0.025
+ self.decode_timestep = 0.05
+
+ # LTX VAE decoder architecture uses 3 upsampler blocks with multiplier equals to 2.
+ # Hence the total feature_channels is multiplied by 8 (2^3).
+ feature_channels = base_channels * 8
+
+ self.conv_in = make_conv_nd(
+ dims=convolution_dimensions,
+ in_channels=in_channels,
+ out_channels=feature_channels,
+ kernel_size=3,
+ stride=1,
+ padding=1,
+ causal=True,
+ spatial_padding_mode=decoder_spatial_padding_mode,
+ )
+
+ self.up_blocks = nn.ModuleList([])
+
+ for block_name, block_params in list(reversed(decoder_blocks)):
+ # Convert int to dict format for uniform handling
+ block_config = {"num_layers": block_params} if isinstance(block_params, int) else block_params
+
+ block, feature_channels = _make_decoder_block(
+ block_name=block_name,
+ block_config=block_config,
+ in_channels=feature_channels,
+ convolution_dimensions=convolution_dimensions,
+ norm_layer=norm_layer,
+ timestep_conditioning=timestep_conditioning,
+ norm_num_groups=self._norm_num_groups,
+ spatial_padding_mode=decoder_spatial_padding_mode,
+ )
+
+ self.up_blocks.append(block)
+
+ if norm_layer == NormLayerType.GROUP_NORM:
+ self.conv_norm_out = nn.GroupNorm(num_channels=feature_channels, num_groups=self._norm_num_groups, eps=1e-6)
+ elif norm_layer == NormLayerType.PIXEL_NORM:
+ self.conv_norm_out = PixelNorm()
+
+ self.conv_act = nn.SiLU()
+ self.conv_out = make_conv_nd(
+ dims=convolution_dimensions,
+ in_channels=feature_channels,
+ out_channels=out_channels,
+ kernel_size=3,
+ padding=1,
+ causal=True,
+ spatial_padding_mode=decoder_spatial_padding_mode,
+ )
+
+ if timestep_conditioning:
+ self.timestep_scale_multiplier = nn.Parameter(torch.tensor(1000.0))
+ self.last_time_embedder = PixArtAlphaCombinedTimestepSizeEmbeddings(embedding_dim=feature_channels * 2,
+ size_emb_dim=0)
+ self.last_scale_shift_table = nn.Parameter(torch.empty(2, feature_channels))
+
+ def forward(
+ self,
+ sample: torch.Tensor,
+ timestep: torch.Tensor | None = None,
+ generator: torch.Generator | None = None,
+ ) -> torch.Tensor:
+ r"""
+ Decode latent representation into video frames.
+ Args:
+ sample: Latent tensor (B, 128, F', H', W').
+ timestep: Timestep for conditioning (if timestep_conditioning=True). Uses default 0.05 if None.
+ generator: Random generator for deterministic noise injection (if inject_noise=True in blocks).
+ Returns:
+ Decoded video (B, 3, F, H, W) where F = 8x(F'-1) + 1, H = 32xH', W = 32xW'.
+ Example: (B, 128, 5, 16, 16) -> (B, 3, 33, 512, 512).
+ Note: First frame is removed after temporal upsampling regardless of causal mode.
+ When causal=False, allows future frame dependencies in convolutions but maintains same output shape.
+ """
+ batch_size = sample.shape[0]
+
+ # Add noise if timestep conditioning is enabled
+ if self.timestep_conditioning:
+ noise = (torch.randn(
+ sample.size(),
+ generator=generator,
+ dtype=sample.dtype,
+ device=sample.device,
+ ) * self.decode_noise_scale)
+
+ sample = noise + (1.0 - self.decode_noise_scale) * sample
+
+ # Denormalize latents
+ sample = self.per_channel_statistics.un_normalize(sample)
+
+ # Use default decode_timestep if timestep not provided
+ if timestep is None and self.timestep_conditioning:
+ timestep = torch.full((batch_size,), self.decode_timestep, device=sample.device, dtype=sample.dtype)
+
+ sample = self.conv_in(sample, causal=self.causal)
+
+ scaled_timestep = None
+ if self.timestep_conditioning:
+ if timestep is None:
+ raise ValueError("'timestep' parameter must be provided when 'timestep_conditioning' is True")
+ scaled_timestep = timestep * self.timestep_scale_multiplier.to(sample)
+
+ for up_block in self.up_blocks:
+ if isinstance(up_block, UNetMidBlock3D):
+ block_kwargs = {
+ "causal": self.causal,
+ "timestep": scaled_timestep if self.timestep_conditioning else None,
+ "generator": generator,
+ }
+ sample = up_block(sample, **block_kwargs)
+ elif isinstance(up_block, ResnetBlock3D):
+ sample = up_block(sample, causal=self.causal, generator=generator)
+ else:
+ sample = up_block(sample, causal=self.causal)
+
+ sample = self.conv_norm_out(sample)
+
+ if self.timestep_conditioning:
+ embedded_timestep = self.last_time_embedder(
+ timestep=scaled_timestep.flatten(),
+ hidden_dtype=sample.dtype,
+ )
+ embedded_timestep = embedded_timestep.view(batch_size, embedded_timestep.shape[-1], 1, 1, 1)
+ ada_values = self.last_scale_shift_table[None, ..., None, None, None].to(
+ device=sample.device, dtype=sample.dtype) + embedded_timestep.reshape(
+ batch_size,
+ 2,
+ -1,
+ embedded_timestep.shape[-3],
+ embedded_timestep.shape[-2],
+ embedded_timestep.shape[-1],
+ )
+ shift, scale = ada_values.unbind(dim=1)
+ sample = sample * (1 + scale) + shift
+
+ sample = self.conv_act(sample)
+ sample = self.conv_out(sample, causal=self.causal)
+
+ # Final spatial expansion: reverse the initial patchify from encoder
+ # Moves pixels from channels back to spatial dimensions
+ # Example: (B, 48, F, 128, 128) -> (B, 3, F, 512, 512) with patch_size=4
+ sample = unpatchify(sample, patch_size_hw=self.patch_size, patch_size_t=1)
+
+ return sample
+
+ def _prepare_tiles(
+ self,
+ latent: torch.Tensor,
+ tiling_config: TilingConfig | None = None,
+ ) -> List[Tile]:
+ splitters = [DEFAULT_SPLIT_OPERATION] * len(latent.shape)
+ mappers = [DEFAULT_MAPPING_OPERATION] * len(latent.shape)
+ if tiling_config is not None and tiling_config.spatial_config is not None:
+ cfg = tiling_config.spatial_config
+ long_side = max(latent.shape[3], latent.shape[4])
+
+ def enable_on_axis(axis_idx: int, factor: int) -> None:
+ size = cfg.tile_size_in_pixels // factor
+ overlap = cfg.tile_overlap_in_pixels // factor
+ axis_length = latent.shape[axis_idx]
+ lower_threshold = max(2, overlap + 1)
+ tile_size = max(lower_threshold, round(size * axis_length / long_side))
+ splitters[axis_idx] = split_in_spatial(tile_size, overlap)
+ mappers[axis_idx] = to_mapping_operation(map_spatial_slice, factor)
+
+ enable_on_axis(3, self.video_downscale_factors.height)
+ enable_on_axis(4, self.video_downscale_factors.width)
+
+ if tiling_config is not None and tiling_config.temporal_config is not None:
+ cfg = tiling_config.temporal_config
+ tile_size = cfg.tile_size_in_frames // self.video_downscale_factors.time
+ overlap = cfg.tile_overlap_in_frames // self.video_downscale_factors.time
+ splitters[2] = split_in_temporal(tile_size, overlap)
+ mappers[2] = to_mapping_operation(map_temporal_slice, self.video_downscale_factors.time)
+
+ return create_tiles(latent.shape, splitters, mappers)
+
+ def tiled_decode(
+ self,
+ latent: torch.Tensor,
+ tiling_config: TilingConfig | None = None,
+ timestep: torch.Tensor | None = None,
+ generator: torch.Generator | None = None,
+ ) -> Iterator[torch.Tensor]:
+ """
+ Decode a latent tensor into video frames using tiled processing.
+ Splits the latent tensor into tiles, decodes each tile individually,
+ and yields video chunks as they become available.
+ Args:
+ latent: Input latent tensor (B, C, F', H', W').
+ tiling_config: Tiling configuration for the latent tensor.
+ timestep: Optional timestep for decoder conditioning.
+ generator: Optional random generator for deterministic decoding.
+ Yields:
+ Video chunks (B, C, T, H, W) by temporal slices;
+ """
+
+ # Calculate full video shape from latent shape to get spatial dimensions
+ full_video_shape = VideoLatentShape.from_torch_shape(latent.shape).upscale(self.video_downscale_factors)
+ tiles = self._prepare_tiles(latent, tiling_config)
+
+ temporal_groups = self._group_tiles_by_temporal_slice(tiles)
+
+ # State for temporal overlap handling
+ previous_chunk = None
+ previous_weights = None
+ previous_temporal_slice = None
+
+ for temporal_group_tiles in temporal_groups:
+ curr_temporal_slice = temporal_group_tiles[0].out_coords[2]
+
+ # Calculate the shape of the temporal buffer for this group of tiles.
+ # The temporal length depends on whether this is the first tile (starts at 0) or not.
+ # - First tile: (frames - 1) * scale + 1
+ # - Subsequent tiles: frames * scale
+ # This logic is handled by TemporalAxisMapping and reflected in out_coords.
+ temporal_tile_buffer_shape = full_video_shape._replace(frames=curr_temporal_slice.stop -
+ curr_temporal_slice.start,)
+
+ buffer = torch.zeros(
+ temporal_tile_buffer_shape.to_torch_shape(),
+ device=latent.device,
+ dtype=latent.dtype,
+ )
+
+ curr_weights = self._accumulate_temporal_group_into_buffer(
+ group_tiles=temporal_group_tiles,
+ buffer=buffer,
+ latent=latent,
+ timestep=timestep,
+ generator=generator,
+ )
+
+ # Blend with previous temporal chunk if it exists
+ if previous_chunk is not None:
+ # Check if current temporal slice overlaps with previous temporal slice
+ if previous_temporal_slice.stop > curr_temporal_slice.start:
+ overlap_len = previous_temporal_slice.stop - curr_temporal_slice.start
+ temporal_overlap_slice = slice(curr_temporal_slice.start - previous_temporal_slice.start, None)
+
+ # The overlap is already masked before it reaches this step. Each tile is accumulated into buffer
+ # with its trapezoidal mask, and curr_weights accumulates the same mask. In the overlap blend we add
+ # the masked values (buffer[...]) and the corresponding weights (curr_weights[...]) into the
+ # previous buffers, then later normalize by weights.
+ previous_chunk[:, :, temporal_overlap_slice, :, :] += buffer[:, :, slice(0, overlap_len), :, :]
+ previous_weights[:, :, temporal_overlap_slice, :, :] += curr_weights[:, :,
+ slice(0, overlap_len), :, :]
+
+ buffer[:, :, slice(0, overlap_len), :, :] = previous_chunk[:, :, temporal_overlap_slice, :, :]
+ curr_weights[:, :, slice(0, overlap_len), :, :] = previous_weights[:, :,
+ temporal_overlap_slice, :, :]
+
+ # Yield the non-overlapping part of the previous chunk
+ previous_weights = previous_weights.clamp(min=1e-8)
+ yield_len = curr_temporal_slice.start - previous_temporal_slice.start
+ yield (previous_chunk / previous_weights)[:, :, :yield_len, :, :]
+
+ # Update state for next iteration
+ previous_chunk = buffer
+ previous_weights = curr_weights
+ previous_temporal_slice = curr_temporal_slice
+
+ # Yield any remaining chunk
+ if previous_chunk is not None:
+ previous_weights = previous_weights.clamp(min=1e-8)
+ yield previous_chunk / previous_weights
+
+ def _group_tiles_by_temporal_slice(self, tiles: List[Tile]) -> List[List[Tile]]:
+ """Group tiles by their temporal output slice."""
+ if not tiles:
+ return []
+
+ groups = []
+ current_slice = tiles[0].out_coords[2]
+ current_group = []
+
+ for tile in tiles:
+ tile_slice = tile.out_coords[2]
+ if tile_slice == current_slice:
+ current_group.append(tile)
+ else:
+ groups.append(current_group)
+ current_slice = tile_slice
+ current_group = [tile]
+
+ # Add the final group
+ if current_group:
+ groups.append(current_group)
+
+ return groups
+
+ def _accumulate_temporal_group_into_buffer(
+ self,
+ group_tiles: List[Tile],
+ buffer: torch.Tensor,
+ latent: torch.Tensor,
+ timestep: torch.Tensor | None,
+ generator: torch.Generator | None,
+ ) -> torch.Tensor:
+ """
+ Decode and accumulate all tiles of a temporal group into a local buffer.
+ The buffer is local to the group and always starts at time 0; temporal coordinates
+ are rebased by subtracting temporal_slice.start.
+ """
+ temporal_slice = group_tiles[0].out_coords[2]
+
+ weights = torch.zeros_like(buffer)
+
+ for tile in group_tiles:
+ decoded_tile = self.forward(latent[tile.in_coords], timestep, generator)
+ mask = tile.blend_mask.to(device=buffer.device, dtype=buffer.dtype)
+ temporal_offset = tile.out_coords[2].start - temporal_slice.start
+ # Use the tile's output coordinate length, not the decoded tile's length,
+ # as the decoder may produce a different number of frames than expected
+ expected_temporal_len = tile.out_coords[2].stop - tile.out_coords[2].start
+ decoded_temporal_len = decoded_tile.shape[2]
+
+ # Ensure we don't exceed the buffer or decoded tile bounds
+ actual_temporal_len = min(expected_temporal_len, decoded_temporal_len, buffer.shape[2] - temporal_offset)
+
+ chunk_coords = (
+ slice(None), # batch
+ slice(None), # channels
+ slice(temporal_offset, temporal_offset + actual_temporal_len),
+ tile.out_coords[3], # height
+ tile.out_coords[4], # width
+ )
+
+ # Slice decoded_tile and mask to match the actual length we're writing
+ decoded_slice = decoded_tile[:, :, :actual_temporal_len, :, :]
+ mask_slice = mask[:, :, :actual_temporal_len, :, :] if mask.shape[2] > 1 else mask
+
+ buffer[chunk_coords] += decoded_slice * mask_slice
+ weights[chunk_coords] += mask_slice
+
+ return weights
+
+ def decode(
+ self,
+ latent: torch.Tensor,
+ tiled=False,
+ tile_size_in_pixels: Optional[int] = 512,
+ tile_overlap_in_pixels: Optional[int] = 128,
+ tile_size_in_frames: Optional[int] = 128,
+ tile_overlap_in_frames: Optional[int] = 24,
+ ) -> torch.Tensor:
+ if tiled:
+ tiling_config = TilingConfig(
+ spatial_config=SpatialTilingConfig(
+ tile_size_in_pixels=tile_size_in_pixels,
+ tile_overlap_in_pixels=tile_overlap_in_pixels,
+ ),
+ temporal_config=TemporalTilingConfig(
+ tile_size_in_frames=tile_size_in_frames,
+ tile_overlap_in_frames=tile_overlap_in_frames,
+ ),
+ )
+ tiles = self.tiled_decode(latent, tiling_config)
+ return torch.cat(list(tiles), dim=2)
+ else:
+ return self.forward(latent)
+
+def decode_video(
+ latent: torch.Tensor,
+ video_decoder: LTX2VideoDecoder,
+ tiling_config: TilingConfig | None = None,
+ generator: torch.Generator | None = None,
+) -> Iterator[torch.Tensor]:
+ """
+ Decode a video latent tensor with the given decoder.
+ Args:
+ latent: Tensor [c, f, h, w]
+ video_decoder: Decoder module.
+ tiling_config: Optional tiling settings.
+ generator: Optional random generator for deterministic decoding.
+ Yields:
+ Decoded chunk [f, h, w, c], uint8 in [0, 255].
+ """
+
+ def convert_to_uint8(frames: torch.Tensor) -> torch.Tensor:
+ frames = (((frames + 1.0) / 2.0).clamp(0.0, 1.0) * 255.0).to(torch.uint8)
+ frames = rearrange(frames[0], "c f h w -> f h w c")
+ return frames
+
+ if tiling_config is not None:
+ for frames in video_decoder.tiled_decode(latent, tiling_config, generator=generator):
+ return convert_to_uint8(frames)
+ else:
+ decoded_video = video_decoder(latent, generator=generator)
+ return convert_to_uint8(decoded_video)
+
+
+def get_video_chunks_number(num_frames: int, tiling_config: TilingConfig | None = None) -> int:
+ """
+ Get the number of video chunks for a given number of frames and tiling configuration.
+ Args:
+ num_frames: Number of frames in the video.
+ tiling_config: Tiling configuration.
+ Returns:
+ Number of video chunks.
+ """
+ if not tiling_config or not tiling_config.temporal_config:
+ return 1
+ cfg = tiling_config.temporal_config
+ frame_stride = cfg.tile_size_in_frames - cfg.tile_overlap_in_frames
+ return (num_frames - 1 + frame_stride - 1) // frame_stride
+
+
+def split_in_spatial(size: int, overlap: int) -> SplitOperation:
+
+ def split(dimension_size: int) -> DimensionIntervals:
+ if dimension_size <= size:
+ return DEFAULT_SPLIT_OPERATION(dimension_size)
+ amount = (dimension_size + size - 2 * overlap - 1) // (size - overlap)
+ starts = [i * (size - overlap) for i in range(amount)]
+ ends = [start + size for start in starts]
+ ends[-1] = dimension_size
+ left_ramps = [0] + [overlap] * (amount - 1)
+ right_ramps = [overlap] * (amount - 1) + [0]
+ return DimensionIntervals(starts=starts, ends=ends, left_ramps=left_ramps, right_ramps=right_ramps)
+
+ return split
+
+
+def split_in_temporal(size: int, overlap: int) -> SplitOperation:
+ non_causal_split = split_in_spatial(size, overlap)
+
+ def split(dimension_size: int) -> DimensionIntervals:
+ if dimension_size <= size:
+ return DEFAULT_SPLIT_OPERATION(dimension_size)
+ intervals = non_causal_split(dimension_size)
+ starts = intervals.starts
+ starts[1:] = [s - 1 for s in starts[1:]]
+ left_ramps = intervals.left_ramps
+ left_ramps[1:] = [r + 1 for r in left_ramps[1:]]
+ return replace(intervals, starts=starts, left_ramps=left_ramps)
+
+ return split
+
+
+def to_mapping_operation(
+ map_func: Callable[[int, int, int, int, int], Tuple[slice, torch.Tensor]],
+ scale: int,
+) -> MappingOperation:
+
+ def map_op(intervals: DimensionIntervals) -> tuple[list[slice], list[torch.Tensor | None]]:
+ output_slices: list[slice] = []
+ masks_1d: list[torch.Tensor | None] = []
+ number_of_slices = len(intervals.starts)
+ for i in range(number_of_slices):
+ start = intervals.starts[i]
+ end = intervals.ends[i]
+ left_ramp = intervals.left_ramps[i]
+ right_ramp = intervals.right_ramps[i]
+ output_slice, mask_1d = map_func(start, end, left_ramp, right_ramp, scale)
+ output_slices.append(output_slice)
+ masks_1d.append(mask_1d)
+ return output_slices, masks_1d
+
+ return map_op
+
+
+def map_temporal_slice(begin: int, end: int, left_ramp: int, right_ramp: int, scale: int) -> Tuple[slice, torch.Tensor]:
+ start = begin * scale
+ stop = 1 + (end - 1) * scale
+ left_ramp = 1 + (left_ramp - 1) * scale
+ right_ramp = right_ramp * scale
+
+ return slice(start, stop), compute_trapezoidal_mask_1d(stop - start, left_ramp, right_ramp, True)
+
+
+def map_spatial_slice(begin: int, end: int, left_ramp: int, right_ramp: int, scale: int) -> Tuple[slice, torch.Tensor]:
+ start = begin * scale
+ stop = end * scale
+ left_ramp = left_ramp * scale
+ right_ramp = right_ramp * scale
+
+ return slice(start, stop), compute_trapezoidal_mask_1d(stop - start, left_ramp, right_ramp, False)
diff --git a/diffsynth/models/minimax_h3_audio_vae.py b/diffsynth/models/minimax_h3_audio_vae.py
new file mode 100644
index 0000000000000000000000000000000000000000..3be65d94653d80a326618f395c9ef10be82822be
--- /dev/null
+++ b/diffsynth/models/minimax_h3_audio_vae.py
@@ -0,0 +1,484 @@
+# SPDX-License-Identifier: Apache-2.0
+import math
+from typing import List
+
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch.nn import Parameter
+from ..core.attention import attention_forward
+
+
+class WarpedTensor(torch.nn.Module):
+ def __init__(self, weight=None, shape=None):
+ super().__init__()
+ if weight is not None:
+ self.weight = torch.nn.Parameter(weight)
+ else:
+ self.weight = torch.nn.Parameter(torch.empty(shape))
+
+ def forward(self):
+ # Offload/ZeRO forward hooks may replace ``weight.data`` immediately
+ # after this module returns. Do not expose storage that those hooks can
+ # turn into a zero-sized placeholder before the caller consumes it.
+ return self.weight.clone()
+
+
+class WeightNormedConv1d(torch.nn.Module):
+ def __init__(self, in_channels, out_channels, kernel_size, stride = 1, padding = 0, dilation = 1, groups = 1, bias = True, padding_mode = "zeros", device=None, dtype=None):
+ super().__init__()
+ self.weight_g = torch.nn.Parameter(torch.empty((out_channels, 1, 1)))
+ self.weight_v = torch.nn.Parameter(torch.empty((out_channels, in_channels, kernel_size)))
+ self.stride = stride
+ self.padding = padding
+ self.dilation = dilation
+ self.groups = groups
+ if bias:
+ self.bias = torch.nn.Parameter(torch.empty((out_channels,)))
+ else:
+ self.bias = None
+
+ def forward(self, input):
+ dtype = self.weight_g.dtype
+ weight = (self.weight_g.float() * self.weight_v.float() / torch.norm(self.weight_v.float(), dim=(1, 2), keepdim=True)).to(dtype)
+ return torch.nn.functional.conv1d(
+ input, weight, self.bias, self.stride, self.padding, self.dilation, self.groups
+ )
+
+
+class WeightNormedConvTranspose1d(torch.nn.Module):
+ def __init__(self, in_channels, out_channels, kernel_size, stride = 1, padding = 0, output_padding = 0, dilation = 1, groups = 1, bias = True, padding_mode = "zeros", device=None, dtype=None):
+ super().__init__()
+ self.weight_g = torch.nn.Parameter(torch.empty((in_channels, 1, 1)))
+ self.weight_v = torch.nn.Parameter(torch.empty((in_channels, out_channels, kernel_size)))
+ self.kernel_size = kernel_size
+ self.stride = stride
+ self.padding = (padding,)
+ self.output_padding = output_padding
+ self.dilation = dilation
+ self.groups = groups
+ if bias:
+ self.bias = torch.nn.Parameter(torch.empty((out_channels,)))
+ else:
+ self.bias = None
+
+ def _output_padding(
+ self,
+ input,
+ output_size: list[int] | None,
+ stride: list[int],
+ padding: list[int],
+ kernel_size: list[int],
+ num_spatial_dims: int,
+ dilation: list[int] | None = None,
+ ) -> list[int]:
+ if output_size is None:
+ ret = self.output_padding
+ else:
+ has_batch_dim = input.dim() == num_spatial_dims + 2
+ num_non_spatial_dims = 2 if has_batch_dim else 1
+ if len(output_size) == num_non_spatial_dims + num_spatial_dims:
+ output_size = output_size[num_non_spatial_dims:]
+ if len(output_size) != num_spatial_dims:
+ raise ValueError(
+ f"ConvTranspose{num_spatial_dims}D: for {input.dim()}D input, output_size must have {num_spatial_dims} "
+ f"or {num_non_spatial_dims + num_spatial_dims} elements (got {len(output_size)})"
+ )
+
+ min_sizes = torch.jit.annotate(list[int], [])
+ max_sizes = torch.jit.annotate(list[int], [])
+ for d in range(num_spatial_dims):
+ dim_size = (
+ (input.size(d + num_non_spatial_dims) - 1) * stride[d]
+ - 2 * padding[d]
+ + (dilation[d] if dilation is not None else 1)
+ * (kernel_size[d] - 1)
+ + 1
+ )
+ min_sizes.append(dim_size)
+ max_sizes.append(min_sizes[d] + stride[d] - 1)
+
+ for i in range(len(output_size)):
+ size = output_size[i]
+ min_size = min_sizes[i]
+ max_size = max_sizes[i]
+ if size < min_size or size > max_size:
+ raise ValueError(
+ f"requested an output size of {output_size}, but valid sizes range "
+ f"from {min_sizes} to {max_sizes} (for an input of {input.size()[2:]})"
+ )
+
+ res = torch.jit.annotate(list[int], [])
+ for d in range(num_spatial_dims):
+ res.append(output_size[d] - min_sizes[d])
+
+ ret = res
+ return ret
+
+ def forward(self, input, output_size=None):
+ dtype = self.weight_g.dtype
+ weight = (self.weight_g.float() * self.weight_v.float() / torch.norm(self.weight_v.float(), dim=(1, 2), keepdim=True)).to(dtype)
+ assert isinstance(self.padding, tuple)
+ num_spatial_dims = 1
+ output_padding = self._output_padding(
+ input,
+ output_size,
+ self.stride, # type: ignore[arg-type]
+ self.padding, # type: ignore[arg-type]
+ self.kernel_size, # type: ignore[arg-type]
+ num_spatial_dims,
+ self.dilation, # type: ignore[arg-type]
+ )
+ return F.conv_transpose1d(
+ input,
+ weight,
+ self.bias,
+ self.stride,
+ self.padding,
+ output_padding,
+ self.groups,
+ self.dilation,
+ )
+
+
+class Snake(nn.Module):
+ def __init__(self, channels, beta=False):
+ super().__init__()
+ if not beta:
+ self.alpha = nn.Parameter(torch.ones(1, channels, 1))
+ self.beta = None
+ else:
+ self.alpha = nn.Parameter(torch.ones(channels))
+ self.beta = Parameter(torch.zeros(channels)) if beta else None
+
+ def forward(self, x):
+ alpha = self.alpha
+ if len(alpha.shape) == 1: alpha = alpha.unsqueeze(0).unsqueeze(-1).exp()
+ beta = self.alpha if self.beta is None else self.beta
+ if len(beta.shape) == 1: beta = beta.unsqueeze(0).unsqueeze(-1).exp()
+ x = x + (beta + 1e-9).reciprocal() * torch.sin(alpha * x).pow(2)
+ return x
+
+
+def kaiser_sinc_filter1d(cutoff, half_width, kernel_size):
+ even = kernel_size % 2 == 0
+ half_size = kernel_size // 2
+ delta_f = 4 * half_width
+ A = 2.285 * (half_size - 1) * math.pi * delta_f + 7.95
+ if A > 50.0:
+ beta = 0.1102 * (A - 8.7)
+ elif A >= 21.0:
+ beta = 0.5842 * (A - 21) ** 0.4 + 0.07886 * (A - 21.0)
+ else:
+ beta = 0.0
+ window = torch.kaiser_window(kernel_size, beta=beta, periodic=False)
+ if even:
+ time = torch.arange(-half_size, half_size, device=window.device) + 0.5
+ else:
+ time = torch.arange(kernel_size, device=window.device) - half_size
+ if cutoff == 0:
+ filter_ = torch.zeros_like(time)
+ else:
+ filter_ = 2 * cutoff * window * torch.sinc(2 * cutoff * time)
+ filter_ /= filter_.sum()
+ return filter_.view(1, 1, kernel_size)
+
+
+class LowPassFilter1d(nn.Module):
+ def __init__(self, cutoff=0.25, half_width=0.3, stride: int = 2, kernel_size: int = 12):
+ super().__init__()
+ self.kernel_size = kernel_size
+ self.stride = stride
+ self.filter = torch.nn.Parameter(kaiser_sinc_filter1d(cutoff, half_width, kernel_size))
+
+ def forward(self, x):
+ x = F.pad(x, (self.kernel_size // 2 - 1, self.kernel_size // 2), mode="replicate")
+ out = F.conv1d(x, self.filter.expand(x.shape[1], -1, -1), stride=self.stride, groups=x.shape[1])
+ return out
+
+
+class UpSample1d(nn.Module):
+ def __init__(self, ratio=2, kernel_size=None):
+ super().__init__()
+ self.ratio = ratio
+ self.kernel_size = int(6 * ratio // 2) * 2 if kernel_size is None else kernel_size
+ self.stride = ratio
+ self.pad = self.kernel_size // ratio - 1
+ self.pad_left = self.pad * self.stride + (self.kernel_size - self.stride) // 2
+ self.pad_right = self.pad * self.stride + (self.kernel_size - self.stride + 1) // 2
+ filter = kaiser_sinc_filter1d(cutoff=0.5 / ratio, half_width=0.6 / ratio, kernel_size=self.kernel_size)
+ self.filter = torch.nn.Parameter(filter)
+
+ def forward(self, x):
+ _, C, _ = x.shape
+ x = F.pad(x, (self.pad, self.pad), mode="replicate")
+ x = self.ratio * F.conv_transpose1d(x, self.filter.expand(C, -1, -1).to(dtype=x.dtype, device=x.device), stride=self.stride, groups=C)
+ x = x[..., self.pad_left : -self.pad_right]
+ return x
+
+
+class DownSample1d(nn.Module):
+ def __init__(self, ratio=2, kernel_size=None):
+ super().__init__()
+ self.ratio = ratio
+ self.kernel_size = int(6 * ratio // 2) * 2 if kernel_size is None else kernel_size
+ self.lowpass = LowPassFilter1d(cutoff=0.5 / ratio, half_width=0.6 / ratio, stride=ratio, kernel_size=self.kernel_size)
+
+ def forward(self, x):
+ return self.lowpass(x)
+
+
+class Activation1d(nn.Module):
+ def __init__(
+ self,
+ activation,
+ up_ratio: int = 2,
+ down_ratio: int = 2,
+ up_kernel_size: int = 12,
+ down_kernel_size: int = 12,
+ ):
+ super().__init__()
+ self.act = activation
+ self.upsample = UpSample1d(up_ratio, up_kernel_size)
+ self.downsample = DownSample1d(down_ratio, down_kernel_size)
+
+ def forward(self, x):
+ x = self.upsample(x)
+ x = self.act(x)
+ x = self.downsample(x)
+ return x
+
+
+class GeGluMlp(nn.Module):
+ def __init__(self, in_features, hidden_features):
+ super().__init__()
+ self.norm = nn.LayerNorm(in_features)
+ self.act = nn.GELU(approximate="tanh")
+ self.w0 = nn.Linear(in_features, hidden_features)
+ self.w1 = nn.Linear(in_features, hidden_features)
+ self.w2 = nn.Linear(hidden_features, in_features)
+
+ def forward(self, x):
+ x = self.norm(x)
+ x = self.act(self.w0(x)) * self.w1(x)
+ x = self.w2(x)
+ return x
+
+
+class CausalAttention(nn.Module):
+ def __init__(self, in_dim, out_dim, num_heads):
+ super().__init__()
+ self.qkv = nn.Linear(in_dim, in_dim * 3, bias=False)
+ self.out_dim = out_dim
+ self.num_heads = num_heads
+ self.proj = nn.Linear(out_dim, out_dim)
+ self.q_bias = WarpedTensor(torch.empty(in_dim))
+ self.zero_k_bias = WarpedTensor(torch.empty(in_dim))
+ self.v_bias = WarpedTensor(torch.empty(in_dim))
+
+ def forward(self, x: torch.Tensor) -> torch.Tensor:
+ q, k, v = self.qkv(x).chunk(3, dim=-1)
+ q, k, v = q + self.q_bias(), k + self.zero_k_bias(), v + self.v_bias()
+ x = attention_forward(q, k, v, q_pattern="b s (n d)", k_pattern="b s (n d)", v_pattern="b s (n d)", out_pattern="b n s d", dims={"n": self.num_heads}, is_causal=True)
+ x = torch.mean(x, dim=1)
+ x = F.adaptive_avg_pool1d(x, self.out_dim)
+ x = self.proj(x)
+ return x
+
+
+class AttnProjection(nn.Module):
+ def __init__(self, in_dim, out_dim, num_heads, mlp_ratio=2):
+ super().__init__()
+ self.norm1 = nn.LayerNorm(in_dim)
+ self.attn = CausalAttention(in_dim, out_dim, num_heads)
+ self.proj = nn.Linear(in_dim, out_dim)
+ self.norm3 = nn.LayerNorm(in_dim)
+ self.norm2 = nn.LayerNorm(out_dim)
+ self.mlp = GeGluMlp(in_features=out_dim, hidden_features=out_dim * mlp_ratio)
+
+ def forward(self, x):
+ x = self.proj(self.norm3(x)) + self.attn(self.norm1(x))
+ x = x + self.mlp(self.norm2(x))
+ return x
+
+
+class ResidualUnit(nn.Module):
+ def __init__(self, dim: int = 16, dilation: int = 1):
+ super().__init__()
+ pad = ((7 - 1) * dilation) // 2
+ self.block = nn.Sequential(
+ Snake(dim),
+ WeightNormedConv1d(dim, dim, kernel_size=7, dilation=dilation, padding=pad),
+ Snake(dim),
+ WeightNormedConv1d(dim, dim, kernel_size=1),
+ )
+
+ def forward(self, x):
+ y = self.block(x)
+ pad = (x.shape[-1] - y.shape[-1]) // 2
+ if pad > 0:
+ x = x[..., pad:-pad]
+ return x + y
+
+
+class EncoderBlock(nn.Module):
+ def __init__(self, dim: int = 16, stride: int = 1):
+ super().__init__()
+ self.block = nn.Sequential(
+ ResidualUnit(dim // 2, dilation=1),
+ ResidualUnit(dim // 2, dilation=3),
+ ResidualUnit(dim // 2, dilation=9),
+ Snake(dim // 2),
+ WeightNormedConv1d(
+ dim // 2,
+ dim,
+ kernel_size=2 * stride,
+ stride=stride,
+ padding=math.ceil(stride / 2),
+ ),
+ )
+
+ def forward(self, x):
+ return self.block(x)
+
+
+class Encoder(nn.Module):
+ def __init__(self, d_model: int = 64, strides: list = [2, 4, 8, 8], d_latent: int = 64):
+ super().__init__()
+ self.block = [WeightNormedConv1d(1, d_model, kernel_size=7, padding=3)]
+ for stride in strides:
+ d_model *= 2
+ self.block += [EncoderBlock(d_model, stride=stride)]
+ self.block += [
+ Snake(d_model),
+ WeightNormedConv1d(d_model, d_latent, kernel_size=3, padding=1),
+ ]
+ self.block = nn.Sequential(*self.block)
+
+ def forward(self, x):
+ return self.block(x)
+
+
+class AMPBlock1(torch.nn.Module):
+ def __init__(self, channels, kernel_size=3, dilation=(1, 3, 5)):
+ super().__init__()
+ self.convs1 = nn.ModuleList([WeightNormedConv1d(channels, channels, kernel_size, stride=1, dilation=d, padding=int((kernel_size * d - d) / 2)) for d in dilation])
+ self.convs2 = nn.ModuleList([WeightNormedConv1d(channels, channels, kernel_size, stride=1, dilation=1, padding=int((kernel_size - 1) / 2)) for _ in range(len(dilation))])
+ self.activations = nn.ModuleList([Activation1d(activation=Snake(channels, beta=True)) for _ in range(2 * len(dilation))])
+
+ def forward(self, x):
+ acts1, acts2 = self.activations[::2], self.activations[1::2]
+ for c1, c2, a1, a2 in zip(self.convs1, self.convs2, acts1, acts2):
+ xt = a1(x)
+ xt = c1(xt)
+ xt = a2(xt)
+ xt = c2(xt)
+ x = xt + x
+ return x
+
+
+class BigVGAN(torch.nn.Module):
+ def __init__(self):
+ super().__init__()
+ self.conv_pre = WeightNormedConv1d(2048, 1024, 7, 1, padding=3)
+ self.ups = nn.ModuleList([
+ nn.ModuleList([WeightNormedConvTranspose1d(1024 // (2**i), 1024 // (2**(i+1)), k, u, padding=(k-u)//2)])
+ for i, (u, k) in enumerate(zip([5, 5, 2, 2, 2, 2, 2], [9, 9, 4, 4, 4, 4, 4]))
+ ])
+ self.resblocks = nn.ModuleList()
+ for i in range(7):
+ ch = 1024 // (2 ** (i + 1))
+ for k in [3, 7, 11]:
+ self.resblocks.append(AMPBlock1(ch, k, (1, 3, 5)))
+ self.activation_post = Activation1d(activation=Snake(ch, beta=True))
+ self.conv_post = WeightNormedConv1d(ch, 1, 7, 1, padding=3, bias=False)
+
+ def forward(self, x):
+ x = self.conv_pre(x)
+ for i, up in enumerate(self.ups):
+ x = up[0](x)
+ xs = self.resblocks[i * 3](x)
+ for j in range(1, 3):
+ xs += self.resblocks[i * 3 + j](x)
+ x = xs / 3
+ x = self.activation_post(x)
+ x = self.conv_post(x)
+ x = torch.clamp(x, min=-1.0, max=1.0)
+ return x
+
+
+_AUDIO_LATENTS_MEAN = [-0.020211687488382354, 0.3876466479950502, -0.04398279799186767, -0.28591514936373, 0.08179686214561671, -0.35782641352446604, 0.040623809960919084, -0.01552534501956604, -0.223362481667332, 0.1821006842509091, 0.2941778783780663, -0.07901167601970885, -0.056815072777201, -0.3699028221860095, -0.31616315591624855, 0.5905951377425391, -0.052139568068853864, 0.013673160263486295, -0.03691647864630577, 0.09732660653298163, -0.3394662328788498, -0.30685677538541667, -0.24504598907458763, -0.034698524462007344, 0.02868032184767538, -0.21217779266454084, -0.1678263169941987, 0.3221287889040614, -0.1223055851554907, 0.4356604928128464, -0.0502599202236253, 0.3979258376211797]
+_AUDIO_LATENTS_STD = [1.6895524230479284, 2.76263727217653, 1.7945344281264435, 1.6801681847309828, 1.6390226546605453, 2.7788298348882177, 1.7659090095747236, 1.6199757612137327, 2.6336525640336896, 1.8539356672817833, 2.5056497896915633, 1.811019237886178, 1.9579657790720237, 1.6685498243529284, 1.4922469314453364, 3.298670198067373, 1.9491804496832168, 1.8720003270431442, 1.8334080103291832, 1.6488070416529093, 1.6176957696319716, 1.9131449234774398, 1.5695245398428617, 1.6943659940415912, 1.8318420762504692, 1.5540637421583379, 1.9344930328968526, 1.599198216109855, 1.718045989838149, 1.6307219190837705, 1.8661226051202384, 1.5613768203168363]
+
+
+class MiniMaxH3AudioVAE(nn.Module):
+
+ def __init__(
+ self,
+ encoder_dim: int = 64,
+ encoder_rates: List[int] = [2, 4, 4, 5, 5],
+ decoder_dim: int = 1024,
+ decoder_rates: List[int] = [5, 5, 2, 2, 2, 2, 2],
+ sample_rate: int = 32000,
+ vae_latent_channels: int = 32,
+ attn_proj: bool = True,
+ decoder_type: str = "bigvgan",
+ ):
+ super().__init__()
+ self.sample_rate = sample_rate
+ self.attn_proj = attn_proj
+ latent_dim = encoder_dim * (2 ** len(encoder_rates))
+ self.hop_length = int(np.prod(encoder_rates))
+ self.encoder = Encoder(encoder_dim, encoder_rates, latent_dim)
+ self.mean_proj = nn.Conv1d(vae_latent_channels, vae_latent_channels, 1)
+ self.logs_proj = nn.Conv1d(vae_latent_channels, vae_latent_channels, 1)
+ self.dec_in_proj = nn.Conv1d(vae_latent_channels, latent_dim, 1)
+ self.decoder = BigVGAN()
+ self.pre_block = AttnProjection(latent_dim, vae_latent_channels, num_heads=8)
+
+ def _latent_stats(self, device):
+ mean = torch.tensor(_AUDIO_LATENTS_MEAN, device=device, dtype=torch.float32).view(1, -1, 1)
+ std = torch.tensor(_AUDIO_LATENTS_STD, device=device, dtype=torch.float32).view(1, -1, 1)
+ return mean, std
+
+ def preprocess(self, audio_data):
+ length = audio_data.shape[-1]
+ right_pad = math.ceil(length / self.hop_length) * self.hop_length - length
+ return nn.functional.pad(audio_data, (0, right_pad))
+
+ def encode(self, audio_data: torch.Tensor) -> torch.Tensor:
+ z = self.encoder(audio_data)
+ if self.attn_proj:
+ z = self.pre_block(z.transpose(1, 2)).transpose(1, 2)
+ mean = self.mean_proj(z)
+ return mean
+
+ def decode(self, z: torch.Tensor) -> torch.Tensor:
+ z = self.dec_in_proj(z)
+ return self.decoder(z)
+
+ @torch.no_grad()
+ def encode_audio(self, waveform: torch.Tensor, dtype=None) -> torch.Tensor:
+ if waveform.dim() != 2:
+ raise ValueError(f"expected waveform [C, L], got {list(waveform.shape)}")
+ out_dtype = waveform.dtype
+ audio_data = waveform.unsqueeze(1).to(dtype if dtype is not None else out_dtype)
+ latent = self.encode(self.preprocess(audio_data))
+ latent_channels = len(_AUDIO_LATENTS_MEAN)
+ if latent.dim() != 3 or latent.shape[1] != latent_channels:
+ raise ValueError(
+ f"expected audio latent [C, {latent_channels}, T], got {list(latent.shape)}"
+ )
+ mean, std = self._latent_stats(latent.device)
+ return ((latent.to(torch.float32) - mean) / std).to(out_dtype)
+
+ @torch.no_grad()
+ def decode_audio(self, latents: torch.Tensor, dtype=None) -> torch.Tensor:
+ mean, std = self._latent_stats(latents.device)
+ z = (latents.to(latents.device, torch.float32) * std + mean).to(dtype if dtype is not None else latents.dtype)
+ waveform = self.decode(z)
+ return waveform.transpose(0, 1).to(latents.dtype)
+
+ def forward(self, z: torch.Tensor) -> torch.Tensor:
+ return self.decode(z)
diff --git a/diffsynth/models/minimax_h3_dit.py b/diffsynth/models/minimax_h3_dit.py
new file mode 100644
index 0000000000000000000000000000000000000000..2a4684e7a2519e63567511dc384c6e7fd7d447ad
--- /dev/null
+++ b/diffsynth/models/minimax_h3_dit.py
@@ -0,0 +1,398 @@
+from __future__ import annotations
+
+import math
+
+import torch
+import torch.nn as nn
+
+from ..core.attention import attention_forward
+from ..core.gradient import gradient_checkpoint_forward
+
+MINIMAX_H3_ADALN_MODALITY_NUM = 3
+_PATCH_T, _PATCH_H, _PATCH_W = 1, 2, 2
+
+
+def patchify_video(latent: torch.Tensor) -> torch.Tensor:
+ # [1,24,T,H,W] -> [T*(H/2)*(W/2), 96]
+ b, c, ft, fh, fw = (int(x) for x in latent.shape)
+ t, h, w = ft // _PATCH_T, fh // _PATCH_H, fw // _PATCH_W
+ packed = latent.reshape(b, c, t, _PATCH_T, h, _PATCH_H, w, _PATCH_W)
+ packed = torch.einsum("nctrhpwq->nthwcrpq", packed)
+ return packed.reshape(b * t * h * w, c * _PATCH_T * _PATCH_H * _PATCH_W).contiguous()
+
+
+def unpatchify_video(rows: torch.Tensor, ft: int, fh: int, fw: int, channel: int = 24) -> torch.Tensor:
+ # [T*(H/2)*(W/2), 96] -> [1,24,T,H,W] (inverse of patchify_video; ft/fh/fw match its input)
+ t, h, w = ft // _PATCH_T, fh // _PATCH_H, fw // _PATCH_W
+ packed = rows.reshape(-1, t, h, w, channel, _PATCH_T, _PATCH_H, _PATCH_W)
+ latent = torch.einsum("nthwcrpq->nctrhpwq", packed)
+ return latent.reshape(-1, channel, ft, fh, fw).contiguous()
+
+
+def pack_audio(latent: torch.Tensor) -> torch.Tensor:
+ # [audio_channel, 32, T] -> [audio_channel*T, 32] (channel-major)
+ ac, ld, steps = (int(x) for x in latent.shape)
+ return latent.permute(0, 2, 1).reshape(ac * steps, ld).contiguous()
+
+
+def unpack_audio(rows: torch.Tensor, audio_channel: int, steps: int, latent_dim: int = 32) -> torch.Tensor:
+ # [audio_channel*T, 32] -> [audio_channel, 32, T]
+ return rows.reshape(audio_channel, steps, latent_dim).permute(0, 2, 1).contiguous()
+
+
+def _norm(size: int, *, eps: float) -> nn.RMSNorm:
+ return nn.RMSNorm(size, eps=eps)
+
+
+def _rotate_half(x: torch.Tensor) -> torch.Tensor:
+ x1, x2 = torch.chunk(x, 2, dim=-1)
+ return torch.cat((-x2, x1), dim=-1)
+
+
+def _apply_rope(x: torch.Tensor, freqs: torch.Tensor) -> torch.Tensor:
+ rot_dim = freqs.shape[-1]
+ x_rot, x_pass = x[..., :rot_dim], x[..., rot_dim:]
+ cos = torch.cos(freqs).to(x.dtype).unsqueeze(1)
+ sin = torch.sin(freqs).to(x.dtype).unsqueeze(1)
+ x_rot = (x_rot * cos) + (_rotate_half(x_rot) * sin)
+ return torch.cat((x_rot, x_pass), dim=-1)
+
+
+def _modulate_scale_shift(x, shift, scale, indices):
+ # Cast back to x: index_select on the AdaLN params can promote the expression.
+ return (x * (1.0 + scale.index_select(0, indices)) + shift.index_select(0, indices)).to(x.dtype)
+
+
+def _modulate_gate(x, gate, other, indices):
+ return (x + gate.index_select(0, indices) * other).to(x.dtype)
+
+
+def _sdpa_varlen_attention(q, k, v, cu_seqlens, softmax_scale):
+ out = torch.empty_like(q)
+ bounds = cu_seqlens.tolist()
+ for start, stop in zip(bounds[:-1], bounds[1:]):
+ if stop == start:
+ continue
+ seg_q = q[start:stop].transpose(0, 1).unsqueeze(0)
+ seg_k = k[start:stop].transpose(0, 1).unsqueeze(0)
+ seg_v = v[start:stop].transpose(0, 1).unsqueeze(0)
+ seg_out = attention_forward(seg_q, seg_k, seg_v, scale=softmax_scale)
+ out[start:stop] = seg_out.squeeze(0).transpose(0, 1)
+ return out
+
+
+class MiniMaxH3Rope(nn.Module):
+ def __init__(self, inv_freq_len: int) -> None:
+ super().__init__()
+ self.inv_freq_len = inv_freq_len
+ self.inv_freq = nn.Parameter(self._build_inv_freq())
+
+ def _build_inv_freq(self, device=None) -> torch.Tensor:
+ steps = torch.arange(0, self.inv_freq_len, dtype=torch.float32, device=device)
+ return 1.0 / (10000.0 ** (steps / self.inv_freq_len))
+
+ def forward(self, img_position_ids: torch.Tensor) -> torch.Tensor:
+ if img_position_ids.dim() != 3 or img_position_ids.shape[0] != 1:
+ raise ValueError(f"img_position_ids must be [1, S, 3], got {list(img_position_ids.shape)}")
+ pos = img_position_ids[0].to(torch.float32)
+ inv_freq = self._build_inv_freq(pos.device)
+ per_axis = pos.unsqueeze(-1) * inv_freq.view(1, 1, -1)
+ t_f, h_f, w_f = per_axis.unbind(dim=1)
+ half = torch.cat((t_f, h_f, w_f), dim=-1)
+ return torch.cat((half, half), dim=-1)
+
+
+class MiniMaxH3TimeEmbedder(nn.Module):
+ def __init__(self, timestep_input_dim, time_embed_hidden_size, time_embed_dim):
+ super().__init__()
+ self.frequency_embedding_size = timestep_input_dim
+ self.proj_in = nn.Linear(timestep_input_dim, time_embed_hidden_size, bias=True)
+ self.proj_out = nn.Linear(time_embed_hidden_size, time_embed_dim, bias=True)
+
+ def forward(self, t: torch.Tensor, *, dtype: torch.dtype) -> torch.Tensor:
+ half = self.frequency_embedding_size // 2
+ freqs = torch.exp(
+ -math.log(10000.0) * torch.arange(half, dtype=torch.float32, device=t.device) / half
+ )
+ args = t.to(torch.float32)[:, None] * freqs[None]
+ t_freq = torch.cat([torch.cos(args), torch.sin(args)], dim=-1)
+ hidden = self.proj_in(t_freq.to(dtype))
+ hidden = nn.functional.silu(hidden)
+ return self.proj_out(hidden)
+
+
+class MiniMaxH3Attention(nn.Module):
+ def __init__(self, hidden_size, num_attention_heads, attention_head_dim, qk_norm_eps):
+ super().__init__()
+ self.num_heads = num_attention_heads
+ self.head_dim = attention_head_dim
+ inner_dim = self.num_heads * self.head_dim
+ self.softmax_scale = self.head_dim**-0.5
+ self.qkv_proj = nn.Linear(hidden_size, inner_dim * 3, bias=False)
+ self.q_norm = _norm(attention_head_dim, eps=qk_norm_eps)
+ self.k_norm = _norm(attention_head_dim, eps=qk_norm_eps)
+ self.out_proj = nn.Linear(inner_dim, hidden_size, bias=False)
+
+ def forward(self, x, *, rope_freqs, cu_seqlens, max_seqlen=None):
+ total = x.shape[0]
+ qkv = self.qkv_proj(x)
+ qkv = qkv.view(total, self.num_heads, 3, self.head_dim)
+ q = qkv[:, :, 0, :]
+ k = qkv[:, :, 1, :]
+ v = qkv[:, :, 2, :]
+ q = self.q_norm(q)
+ k = self.k_norm(k)
+ if rope_freqs is not None:
+ q = _apply_rope(q, rope_freqs)
+ k = _apply_rope(k, rope_freqs)
+ out = _sdpa_varlen_attention(q, k, v, cu_seqlens=cu_seqlens, softmax_scale=self.softmax_scale)
+ out = out.reshape(total, self.num_heads * self.head_dim)
+ return self.out_proj(out)
+
+
+class MiniMaxH3MLP(nn.Module):
+ def __init__(self, hidden_size, ffn_hidden_size):
+ super().__init__()
+ self.fc1 = nn.Linear(hidden_size, ffn_hidden_size * 2, bias=False)
+ self.fc2 = nn.Linear(ffn_hidden_size, hidden_size, bias=False)
+
+ def forward(self, x):
+ hidden = self.fc1(x)
+ gate, up = hidden.chunk(2, dim=-1)
+ hidden = nn.functional.silu(gate) * up
+ return self.fc2(hidden)
+
+
+class MiniMaxH3AdalnProj(nn.Module):
+ def __init__(self, hidden_size, time_embed_dim, out_features, *, expand_ratio, modality_num):
+ super().__init__()
+ if out_features != expand_ratio * hidden_size * modality_num:
+ raise ValueError(
+ f"adaln out_features mismatch: {out_features} != {expand_ratio}*{hidden_size}*{modality_num}"
+ )
+ self.expand_ratio = expand_ratio
+ self.modality_num = modality_num
+ self.hidden_size = hidden_size
+ self.linear = nn.Linear(time_embed_dim, out_features, bias=True)
+
+ def forward(self, t_emb):
+ x = nn.functional.silu(t_emb)
+ x = self.linear(x)
+ m = x.shape[0]
+ x = x.view(m * self.modality_num, self.expand_ratio * self.hidden_size)
+ return tuple(x.chunk(self.expand_ratio, dim=-1))
+
+
+class MiniMaxH3TokenRefinerBlock(nn.Module):
+ def __init__(self, hidden_size, num_attention_heads, attention_head_dim, ffn_hidden_size, norm_eps, qk_norm_eps):
+ super().__init__()
+ self.norm1 = _norm(hidden_size, eps=norm_eps)
+ self.norm2 = _norm(hidden_size, eps=norm_eps)
+ self.attn = MiniMaxH3Attention(hidden_size, num_attention_heads, attention_head_dim, qk_norm_eps)
+ self.mlp = MiniMaxH3MLP(hidden_size, ffn_hidden_size)
+
+ def forward(self, x, *, cu_seqlens, max_seqlen):
+ x = x + self.attn(self.norm1(x), rope_freqs=None, cu_seqlens=cu_seqlens, max_seqlen=max_seqlen)
+ x = x + self.mlp(self.norm2(x))
+ return x
+
+
+class MiniMaxH3TokenRefiner(nn.Module):
+ def __init__(self, num_layers, hidden_size, num_attention_heads, attention_head_dim, ffn_hidden_size, norm_eps, qk_norm_eps, final_norm_eps):
+ super().__init__()
+ self.blocks = nn.ModuleList(
+ [MiniMaxH3TokenRefinerBlock(hidden_size, num_attention_heads, attention_head_dim, ffn_hidden_size, norm_eps, qk_norm_eps) for _ in range(num_layers)]
+ )
+ self.final_norm = _norm(hidden_size, eps=final_norm_eps)
+
+ def forward(self, x, *, cu_seqlens, max_seqlen):
+ for block in self.blocks:
+ x = block(x, cu_seqlens=cu_seqlens, max_seqlen=max_seqlen)
+ return self.final_norm(x)
+
+
+class MiniMaxH3DiTBlock(nn.Module):
+ def __init__(self, hidden_size, num_attention_heads, attention_head_dim, ffn_hidden_size, time_embed_dim, adaln_out_features, norm_eps, qk_norm_eps):
+ super().__init__()
+ self.norm1 = _norm(hidden_size, eps=norm_eps)
+ self.norm2 = _norm(hidden_size, eps=norm_eps)
+ self.attn = MiniMaxH3Attention(hidden_size, num_attention_heads, attention_head_dim, qk_norm_eps)
+ self.mlp = MiniMaxH3MLP(hidden_size, ffn_hidden_size)
+ self.adaln_proj = MiniMaxH3AdalnProj(hidden_size, time_embed_dim, adaln_out_features, expand_ratio=6, modality_num=MINIMAX_H3_ADALN_MODALITY_NUM)
+
+ def forward(self, x, *, t_emb, combined_indices, rope_freqs, cu_seqlens, max_seqlen):
+ shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = self.adaln_proj(t_emb)
+ residual = x
+ h = self.norm1(x)
+ h = _modulate_scale_shift(h, shift_msa, scale_msa, combined_indices)
+ h = self.attn(h, rope_freqs=rope_freqs, cu_seqlens=cu_seqlens, max_seqlen=max_seqlen)
+ x = _modulate_gate(residual, gate_msa, h, combined_indices)
+ residual = x
+ h = self.norm2(x)
+ h = _modulate_scale_shift(h, shift_mlp, scale_mlp, combined_indices)
+ h = self.mlp(h)
+ return _modulate_gate(residual, gate_mlp, h, combined_indices)
+
+
+class MiniMaxH3FinalLayer(nn.Module):
+ def __init__(self, hidden_size, time_embed_dim, final_adaln_out_features, latents_dim, audio_latents_dim, patch_size, final_norm_eps):
+ super().__init__()
+ video_patch_dim = latents_dim * patch_size[0] * patch_size[1] * patch_size[2]
+ self.norm = _norm(hidden_size, eps=final_norm_eps)
+ self.adaln_proj = MiniMaxH3AdalnProj(hidden_size, time_embed_dim, final_adaln_out_features, expand_ratio=2, modality_num=1)
+ self.video_out = nn.Linear(hidden_size, video_patch_dim, bias=True)
+ self.audio_out = nn.Linear(hidden_size, audio_latents_dim, bias=True)
+
+ def forward(self, x, *, t_emb, inverse_indices):
+ shift, scale = self.adaln_proj(t_emb)
+ h = self.norm(x)
+ h = _modulate_scale_shift(h, shift, scale, inverse_indices)
+ video = self.video_out(h)
+ audio = self.audio_out(h)
+ return video, audio
+
+
+class MiniMaxH3DiT(nn.Module):
+ _repeated_blocks = ["MiniMaxH3DiTBlock"]
+
+ def __init__(
+ self,
+ num_layers: int = 50,
+ token_refiner_num_layers: int = 2,
+ hidden_size: int = 5376,
+ num_attention_heads: int = 56,
+ attention_head_dim: int = 128,
+ ffn_hidden_size: int = 14336,
+ latents_dim: int = 24,
+ audio_latents_dim: int = 32,
+ patch_size: tuple = (1, 2, 2),
+ text_dim: int = 5120,
+ timestep_input_dim: int = 256,
+ time_embed_hidden_size: int = 5376,
+ time_embed_dim: int = 2688,
+ adaln_out_features: int = 96768,
+ final_adaln_out_features: int = 10752,
+ rope_inv_freq_len: int = 16,
+ norm_eps: float = 1e-5,
+ qk_norm_eps: float = 1e-5,
+ final_norm_eps: float = 1e-5,
+ **kwargs,
+ ):
+ super().__init__()
+ self.hidden_size = hidden_size
+ self.num_attention_heads = num_attention_heads
+ self.num_channels_latents = latents_dim
+ video_patch_dim = latents_dim * patch_size[0] * patch_size[1] * patch_size[2]
+
+ self.video_patch_proj = nn.Linear(video_patch_dim, hidden_size, bias=True)
+ self.audio_patch_proj = nn.Linear(audio_latents_dim, hidden_size, bias=True)
+ self.condition_proj = nn.Linear(text_dim, hidden_size, bias=True)
+ self.time_embedder = MiniMaxH3TimeEmbedder(timestep_input_dim, time_embed_hidden_size, time_embed_dim)
+ self.rope = MiniMaxH3Rope(rope_inv_freq_len)
+ self.token_refiner = MiniMaxH3TokenRefiner(
+ token_refiner_num_layers, hidden_size, num_attention_heads, attention_head_dim,
+ ffn_hidden_size, norm_eps, qk_norm_eps, final_norm_eps,
+ )
+ self.blocks = nn.ModuleList(
+ [MiniMaxH3DiTBlock(hidden_size, num_attention_heads, attention_head_dim, ffn_hidden_size, time_embed_dim, adaln_out_features, norm_eps, qk_norm_eps) for _ in range(num_layers)]
+ )
+ self.final_layer = MiniMaxH3FinalLayer(hidden_size, time_embed_dim, final_adaln_out_features, latents_dim, audio_latents_dim, patch_size, final_norm_eps)
+
+ def _embed(self, *, x, audio_x, text_embeddings_selected, unique_timesteps, img_pos, audio_pos, text_pos, refiner_cu_seqlens, refiner_max_seqlen, seq_len, device):
+ dtype = text_embeddings_selected.dtype
+ x_rows = x.view(-1, x.shape[-1]).index_select(0, img_pos).to(dtype)
+ video_embed = self.video_patch_proj(x_rows)
+ audio_rows = audio_x.view(-1, audio_x.shape[-1]).index_select(0, audio_pos).to(dtype)
+ audio_embed = self.audio_patch_proj(audio_rows)
+ text_rows = text_embeddings_selected.to(device=device)
+ text_embed = self.condition_proj(text_rows)
+ text_embed = self.token_refiner(text_embed, cu_seqlens=refiner_cu_seqlens, max_seqlen=refiner_max_seqlen)
+
+ embeddings = torch.zeros((seq_len, self.hidden_size), device=device, dtype=dtype)
+ embeddings.index_add_(0, text_pos, text_embed.to(dtype)[: text_pos.shape[0]])
+ embeddings.index_add_(0, img_pos, video_embed.to(dtype)[: img_pos.shape[0]])
+ embeddings.index_add_(0, audio_pos, audio_embed.to(dtype)[: audio_pos.shape[0]])
+
+ t_emb = self.time_embedder(unique_timesteps, dtype=dtype)
+ return embeddings, t_emb
+
+ def forward(
+ self,
+ x,
+ audio_x,
+ img_position_ids,
+ unique_timesteps,
+ inverse_indices,
+ update_mask,
+ token_tags,
+ prompt_embeds,
+ img_pos_info,
+ audio_pos_info,
+ text_pos_info,
+ img_pos_for_infer_output_info,
+ packed_seq_params,
+ refiner_packed_seq_params,
+ use_gradient_checkpointing=False,
+ use_gradient_checkpointing_offload=False,
+ update_audio_mask=None,
+ skip_mask_out_condition=False,
+ ) -> tuple[torch.Tensor, torch.Tensor]:
+ inverse_indices = inverse_indices.view(-1).to(torch.long)
+ token_tags = token_tags.view(-1).to(torch.long)
+ text_selected = prompt_embeds
+
+ img_pos = img_pos_info["position_ids"].view(-1).to(torch.long)
+ audio_pos = audio_pos_info["position_ids"].view(-1).to(torch.long)
+ text_pos = text_pos_info["position_ids"].view(-1).to(torch.long)
+ infer_out_pos = img_pos_for_infer_output_info["position_ids"].view(-1).to(torch.long)
+
+ cu_seqlens = packed_seq_params["cu_seqlens_q"].to(torch.int32)
+ max_seqlen = int(packed_seq_params["max_seqlen_q"])
+ refiner_cu = refiner_packed_seq_params["cu_seqlens_q"].to(torch.int32)
+ refiner_max = int(refiner_packed_seq_params["max_seqlen_q"])
+
+ if x.dim() != 3 or x.shape[0] != 1:
+ raise ValueError(f"x must be [1, S, C], got {list(x.shape)}")
+ seq_len = int(x.shape[1])
+ device = x.device
+
+ rope_freqs = self.rope(img_position_ids).to(device)
+
+ decoder_input, t_emb = self._embed(
+ x=x, audio_x=audio_x, text_embeddings_selected=text_selected,
+ unique_timesteps=unique_timesteps.view(-1).to(device),
+ img_pos=img_pos.to(device), audio_pos=audio_pos.to(device), text_pos=text_pos.to(device),
+ refiner_cu_seqlens=refiner_cu.to(device), refiner_max_seqlen=refiner_max,
+ seq_len=seq_len, device=device,
+ )
+
+ combined_indices = (inverse_indices * MINIMAX_H3_ADALN_MODALITY_NUM + token_tags.clamp(min=0)).to(device)
+ inverse_indices = inverse_indices.to(device)
+
+ hidden = decoder_input
+ cu_seqlens = cu_seqlens.to(device)
+ for block in self.blocks:
+ hidden = gradient_checkpoint_forward(
+ block,
+ use_gradient_checkpointing,
+ use_gradient_checkpointing_offload,
+ hidden,
+ t_emb=t_emb,
+ combined_indices=combined_indices,
+ rope_freqs=rope_freqs,
+ cu_seqlens=cu_seqlens,
+ max_seqlen=max_seqlen,
+ )
+
+ video_logits, audio_logits = self.final_layer(hidden, t_emb=t_emb, inverse_indices=inverse_indices)
+
+ video_logits = video_logits.index_select(0, infer_out_pos.to(device))
+ audio_logits = audio_logits.index_select(0, audio_pos.to(device))
+ if not skip_mask_out_condition:
+ update_mask = update_mask.view(-1).to(device)
+ if update_mask.shape[0] != video_logits.shape[0]:
+ raise ValueError(f"update_mask length mismatch: {update_mask.shape[0]} != {video_logits.shape[0]}")
+ video_logits = video_logits * update_mask.unsqueeze(-1)
+ if update_audio_mask is not None:
+ audio_logits = audio_logits * update_audio_mask.view(-1).unsqueeze(-1)
+ return video_logits, audio_logits
diff --git a/diffsynth/models/minimax_h3_dit_comfy.py b/diffsynth/models/minimax_h3_dit_comfy.py
new file mode 100644
index 0000000000000000000000000000000000000000..50a6d1b7a056ecf51ad6d1f406e45da8ed4086b0
--- /dev/null
+++ b/diffsynth/models/minimax_h3_dit_comfy.py
@@ -0,0 +1,78 @@
+import types
+
+import torch
+from torch import nn
+
+from .minimax_h3_dit import MiniMaxH3AdalnProj, MiniMaxH3DiT, _apply_rope, _sdpa_varlen_attention
+
+
+class MiniMaxH3ComfyPrunedAdalnProj(MiniMaxH3AdalnProj):
+ def forward(self, t_emb):
+ x = self.linear(t_emb)
+ m = x.shape[0]
+ x = x.view(m * self.modality_num, self.expand_ratio * self.hidden_size)
+ return tuple(x.chunk(self.expand_ratio, dim=-1))
+
+
+def _comfy_attention_forward(self, x, *, rope_freqs, cu_seqlens, max_seqlen=None):
+ total = x.shape[0]
+ qkv = self.qkv_proj(x)
+ qkv = qkv.view(total, 3, self.num_heads, self.head_dim)
+ q = qkv[:, 0, :, :]
+ k = qkv[:, 1, :, :]
+ v = qkv[:, 2, :, :]
+ q = self.q_norm(q)
+ k = self.k_norm(k)
+ if rope_freqs is not None:
+ q = _apply_rope(q, rope_freqs)
+ k = _apply_rope(k, rope_freqs)
+ out = _sdpa_varlen_attention(q, k, v, cu_seqlens=cu_seqlens, softmax_scale=self.softmax_scale)
+ out = out.reshape(total, self.num_heads * self.head_dim)
+ return self.out_proj(out)
+
+
+class MiniMaxH3ComfyPrunedTimeEmbedder(nn.Module):
+ def __init__(self, table_getter):
+ super().__init__()
+ self._table_getter = table_getter
+
+ def forward(self, t: torch.Tensor, *, dtype: torch.dtype) -> torch.Tensor:
+ table = self._table_getter().to(device=t.device, dtype=torch.float32)
+ pos = t.to(torch.float32).clamp(0.0, 1.0) * (table.shape[0] - 1)
+ i0 = pos.floor().long().clamp(max=table.shape[0] - 2)
+ weight = (pos - i0).unsqueeze(1)
+ t_emb = torch.lerp(table[i0], table[i0 + 1], weight)
+ return t_emb.to(dtype)
+
+
+def _to_comfy_pruned_adaln_proj(proj: MiniMaxH3AdalnProj) -> MiniMaxH3ComfyPrunedAdalnProj:
+ curve_proj = MiniMaxH3ComfyPrunedAdalnProj(
+ proj.hidden_size,
+ proj.linear.in_features,
+ proj.linear.out_features,
+ expand_ratio=proj.expand_ratio,
+ modality_num=proj.modality_num,
+ )
+ curve_proj.linear = proj.linear
+ return curve_proj
+
+
+class MiniMaxH3DiTComfy(MiniMaxH3DiT):
+ def __init__(self, **kwargs):
+ super().__init__(**kwargs)
+ for block in self.blocks:
+ block.attn.forward = types.MethodType(_comfy_attention_forward, block.attn)
+ for block in self.token_refiner.blocks:
+ block.attn.forward = types.MethodType(_comfy_attention_forward, block.attn)
+
+
+class MiniMaxH3DiTComfyPruned(MiniMaxH3DiTComfy):
+ def __init__(self, adaln_curve_grid: int = 1025, time_embed_dim: int = 8, **kwargs):
+ super().__init__(time_embed_dim=time_embed_dim, **kwargs)
+ self.register_buffer(
+ "adaln_t_table", torch.empty(adaln_curve_grid, time_embed_dim, dtype=torch.float32)
+ )
+ self.time_embedder = MiniMaxH3ComfyPrunedTimeEmbedder(lambda: self.adaln_t_table)
+ for block in self.blocks:
+ block.adaln_proj = _to_comfy_pruned_adaln_proj(block.adaln_proj)
+ self.final_layer.adaln_proj = _to_comfy_pruned_adaln_proj(self.final_layer.adaln_proj)
diff --git a/diffsynth/models/minimax_h3_text_encoder.py b/diffsynth/models/minimax_h3_text_encoder.py
new file mode 100644
index 0000000000000000000000000000000000000000..2b3a894524a2ec99a206fc993b3f007c7b8f76e3
--- /dev/null
+++ b/diffsynth/models/minimax_h3_text_encoder.py
@@ -0,0 +1,231 @@
+import torch
+from transformers import Qwen3VLConfig, Qwen3VLModel
+
+VISION_START = "<|vision_start|>"
+VISION_END = "<|vision_end|>"
+IMAGE_PAD = "<|image_pad|>"
+VIDEO_PAD = "<|video_pad|>"
+
+PRESENTATION_TEXT_TAG = 1
+PRESENTATION_VIDEO_TAG = 0
+
+QWEN_VIDEO_SAMPLE_FPS = 2.0
+QWEN_TEMPORAL_PATCH = 2
+MINIMAX_SUPPORTED_FPS = 24
+
+class MiniMaxH3TextEncoder(torch.nn.Module):
+ def __init__(self, num_retained_layers: int = 50):
+ super().__init__()
+ self.num_retained_layers = num_retained_layers
+ config = Qwen3VLConfig(**{
+ "architectures": ["Qwen3VLForConditionalGeneration"],
+ "image_token_id": 151655,
+ "model_type": "qwen3_vl",
+ "text_config": {
+ "attention_bias": False,
+ "attention_dropout": 0.0,
+ "bos_token_id": 151643,
+ "dtype": "bfloat16",
+ "eos_token_id": 151645,
+ "head_dim": 128,
+ "hidden_act": "silu",
+ "hidden_size": 5120,
+ "initializer_range": 0.02,
+ "intermediate_size": 25600,
+ "max_position_embeddings": 262144,
+ "model_type": "qwen3_vl_text",
+ "num_attention_heads": 64,
+ "num_hidden_layers": num_retained_layers,
+ "num_key_value_heads": 8,
+ "rms_norm_eps": 1e-06,
+ "rope_scaling": {
+ "mrope_interleaved": True,
+ "mrope_section": [24, 20, 20],
+ "rope_type": "default",
+ },
+ "rope_theta": 5000000,
+ "use_cache": True,
+ "vocab_size": 151936,
+ },
+ "tie_word_embeddings": False,
+ "transformers_version": "4.57.0.dev0",
+ "video_token_id": 151656,
+ "vision_config": {
+ "deepstack_visual_indexes": [8, 16, 24],
+ "depth": 27,
+ "hidden_act": "gelu_pytorch_tanh",
+ "hidden_size": 1152,
+ "in_channels": 3,
+ "initializer_range": 0.02,
+ "intermediate_size": 4304,
+ "model_type": "qwen3_vl",
+ "num_heads": 16,
+ "num_position_embeddings": 2304,
+ "out_hidden_size": 5120,
+ "patch_size": 16,
+ "spatial_merge_size": 2,
+ "temporal_patch_size": 2,
+ },
+ "vision_end_token_id": 151653,
+ "vision_start_token_id": 151652,
+ })
+ self.model = Qwen3VLModel(config)
+ self.model.language_model.norm = torch.nn.Identity()
+ self.config = config
+ self.image_token_id = config.image_token_id
+ self.video_token_id = config.video_token_id
+
+ def forward(
+ self,
+ input_ids=None,
+ attention_mask=None,
+ pixel_values=None,
+ image_grid_thw=None,
+ pixel_values_videos=None,
+ video_grid_thw=None,
+ **kwargs,
+ ):
+ mm_token_type_ids = None
+ if pixel_values is not None or pixel_values_videos is not None:
+ mm_token_type_ids = torch.zeros_like(input_ids, dtype=torch.int32)
+ mm_token_type_ids[input_ids == self.image_token_id] = 1
+ mm_token_type_ids[input_ids == self.video_token_id] = 2
+ outputs = self.model(
+ input_ids=input_ids,
+ attention_mask=attention_mask,
+ pixel_values=pixel_values,
+ image_grid_thw=image_grid_thw,
+ pixel_values_videos=pixel_values_videos,
+ video_grid_thw=video_grid_thw,
+ mm_token_type_ids=mm_token_type_ids,
+ output_hidden_states=False,
+ use_cache=False,
+ return_dict=True,
+ **kwargs,
+ )
+ return outputs.last_hidden_state[0]
+
+
+def image_token_counts(processor, images):
+ vision = processor.image_processor(images=images, return_tensors="pt")
+ merge = int(processor.image_processor.merge_size) ** 2
+ counts = [int(vision["image_grid_thw"][i].prod().item()) // merge for i in range(len(images))]
+ return vision["pixel_values"], vision["image_grid_thw"], counts
+
+
+def video_token_counts(processor, videos, timestamps_per_video):
+ vout = processor.video_processor(videos=videos, do_sample_frames=False, return_tensors="pt")
+ grid = vout["video_grid_thw"]
+ merge = int(processor.image_processor.merge_size) ** 2
+ block_counts, block_timestamps = [], []
+ for index, timestamps in enumerate(timestamps_per_video):
+ n_blocks = int(grid[index, 0])
+ per_block = int(grid[index, 1]) * int(grid[index, 2]) // merge
+ block_counts.append([per_block] * n_blocks)
+ block_timestamps.append([float(t) for t in timestamps])
+ return vout["pixel_values_videos"], grid, block_counts, block_timestamps
+
+
+def _text_ids(tokenizer, text: str) -> list[int]:
+ return list(tokenizer(text, add_special_tokens=False)["input_ids"])
+
+
+def _vision_block_ids(tokenizer, pad_token: str, count: int) -> list[int]:
+ return (
+ [tokenizer.convert_tokens_to_ids(VISION_START)]
+ + [tokenizer.convert_tokens_to_ids(pad_token)] * int(count)
+ + [tokenizer.convert_tokens_to_ids(VISION_END)]
+ )
+
+class _Presentation:
+ def __init__(self):
+ self.ids: list[int] = []
+ self.tags: list[int] = []
+
+ def text(self, token_ids: list[int]):
+ self.ids += token_ids
+ self.tags += [PRESENTATION_TEXT_TAG] * len(token_ids)
+
+ def vision(self, token_ids: list[int]):
+ self.ids += token_ids
+ self.tags += [PRESENTATION_VIDEO_TAG] * len(token_ids)
+
+ def build(self):
+ return (torch.tensor(self.ids, dtype=torch.long), torch.tensor(self.tags, dtype=torch.long))
+
+
+def presentation_t2va(tokenizer, prompt: str):
+ presentation = _Presentation()
+ presentation.text(_text_ids(tokenizer, prompt))
+ return presentation.build()
+
+
+def presentation_fl2va(tokenizer, prompt: str, image_token_counts):
+ presentation = _Presentation()
+ for index, count in enumerate(image_token_counts, start=1):
+ presentation.text(_text_ids(tokenizer, f": "))
+ presentation.vision(_vision_block_ids(tokenizer, IMAGE_PAD, count))
+ presentation.text(_text_ids(tokenizer, prompt))
+ return presentation.build()
+
+
+def presentation_ref2va(tokenizer, prompt: str, condition_labels, image_token_counts, video_block_token_counts, video_block_timestamps):
+ if not prompt:
+ raise ValueError("prompt must be non-empty")
+ presentation = _Presentation()
+ image_seen = 0
+ video_seen = 0
+ for cond_type, ordinal in condition_labels:
+ if cond_type == "image":
+ image_seen += 1
+ if image_seen > len(image_token_counts):
+ raise ValueError("image_token_count required for an image reference")
+ count = int(image_token_counts[image_seen - 1])
+ if count <= 0:
+ raise ValueError("image_token_count required for an image reference")
+ presentation.text(_text_ids(tokenizer, f": "))
+ presentation.vision(_vision_block_ids(tokenizer, IMAGE_PAD, count))
+ elif cond_type == "audio":
+ presentation.text(_text_ids(tokenizer, f"