Upload folder using huggingface_hub

.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

.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 }}

.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/

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.

diffsynth/__init__.py

@@ -0,0 +1 @@
+from .core import *

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

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,
+}

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 *

diffsynth/core/attention/__init__.py

@@ -0,0 +1 @@
+from .attention import attention_forward

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.")

diffsynth/core/data/__init__.py

@@ -0,0 +1 @@
+from .unified_dataset import UnifiedDataset

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

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

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

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()

diffsynth/core/gradient/__init__.py

@@ -0,0 +1 @@
+from .gradient_checkpoint import gradient_checkpoint_forward

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

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

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,
+        }

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()

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

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)

diffsynth/core/offload_training/__init__.py

@@ -0,0 +1 @@
+from .manager import OffloadTrainingManager

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()

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<<k = x (unchanged).
+    For non-power-of-two: (x-1).bit_length() exceeds floor-log2(x), rounding up to next 2^n.
+    """
+    return 1 if x <= 1 else 1 << (x - 1).bit_length()
+
+def _prev_power_of_two(x: int) -> 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()

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

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

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

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")

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)")

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")

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

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)

diffsynth/core/vram/__init__.py

@@ -0,0 +1,2 @@
+from .initialization import skip_model_initialization
+from .layers import *

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

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

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

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 *

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

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

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)

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

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)

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

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 `<model_string>:<method>[/<exclude_modules>]`, where `<model_string>` 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

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()