diffsynth 2.1.12__tar.gz → 2.2.1__tar.gz

{diffsynth-2.1.12 → diffsynth-2.2.1}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: diffsynth
-Version: 2.1.12
+Version: 2.2.1
 Summary: Enjoy the magic of Diffusion models!
 Author: ModelScope Team
 License-Expression: Apache-2.0
@@ -20,6 +20,7 @@ Requires-Dist: ftfy
 Requires-Dist: pandas
 Requires-Dist: accelerate
 Requires-Dist: peft
+Requires-Dist: entropack>=0.1.7
 Provides-Extra: audio
 Requires-Dist: av; extra == "audio"
 Requires-Dist: torchaudio; extra == "audio"
@@ -28,7 +29,6 @@ Requires-Dist: librosa; extra == "audio"
 Provides-Extra: quant
 Requires-Dist: bitsandbytes; extra == "quant"
 Requires-Dist: comfy-kitchen; extra == "quant"
-Requires-Dist: entropack[cuda13]; extra == "quant"
 Requires-Dist: torchao>=0.16; extra == "quant"
 Provides-Extra: training
 Requires-Dist: deepspeed; extra == "training"
@@ -63,7 +63,6 @@ Requires-Dist: torchcodec; extra == "all"
 Requires-Dist: librosa; extra == "all"
 Requires-Dist: bitsandbytes; extra == "all"
 Requires-Dist: comfy-kitchen; extra == "all"
-Requires-Dist: entropack[cuda13]; extra == "all"
 Requires-Dist: torchao>=0.16; extra == "all"
 Requires-Dist: deepspeed; extra == "all"
 Requires-Dist: tensorboard; extra == "all"
@@ -114,6 +113,8 @@ See also:
 
 > Currently, the development personnel of this project are limited, with most of the work handled by [Artiprocher](https://github.com/Artiprocher) and [mi804](https://github.com/mi804). Therefore, the progress of new feature development will be relatively slow, and the speed of responding to and resolving issues is limited. We apologize for this and ask developers to understand.
 
+- **November 10, 2026** We now support MiniMax-H3-Image, the image model of the MiniMax-H3 family, including unified text-to-image and reference-image editing inference, low VRAM inference, and full/LoRA training. The checkpoint is distilled to six sampling steps with one transformer output head per step, and it generates about 16-megapixel images without an upscaler. See the [documentation](/docs/en/Model_Details/MiniMax-H3-Image.md) and [example code](/examples/minimax_h3_image/).
+
 - **October 21, 2026** We now support Fibo 2.0, including unified text-to-image and image editing inference, masked and multi-reference editing, text rendering, low VRAM inference, and full/LoRA training for both text-to-image and image editing on [Fibo-2.0-dev](https://www.modelscope.cn/models/briaai/Fibo-2.0-dev). See the [documentation](/docs/en/Model_Details/Fibo.md) and [example code](/examples/fibo/).
 
 - **October 20, 2026** We now support MOSS-Image, including inference and training for [MOSS-Image-Base](https://modelscope.cn/models/openmoss/MOSS-Image-Base) and [MOSS-Image-SFT](https://modelscope.cn/models/openmoss/MOSS-Image-SFT). In addition, we provide an [EntroPack low-bit quantized model](https://www.modelscope.cn/models/DiffSynth-Studio/MOSS-Image-EntroPack), which can run the 3.5-bit quantized model with as little as 4 GB of VRAM; and an [accelerated model](https://modelscope.cn/models/DiffSynth-Studio/MOSS-Image-SFT-TurboLoRA) distillation-accelerated via [OPD training](https://diffsynth-studio-doc.readthedocs.io/en/latest/Training/Distillation.html), supporting 8-step ultra-fast inference.
@@ -404,6 +405,7 @@ Quick start: experience popular and the latest models:
 
 | Architecture | Model ID | Inference | Low VRAM Inference | Full Training | Validation After Full Training | LoRA Training | Validation After LoRA Training |
 |-|-|-|-|-|-|-|-|
+| MiniMax-H3-Image | [MiniMax/MiniMax-H3-Image](https://www.modelscope.cn/models/MiniMax/MiniMax-H3-Image) | [code](/examples/minimax_h3_image/model_inference/MiniMax-H3-Image.py) | [code](/examples/minimax_h3_image/model_inference_low_vram/MiniMax-H3-Image.py) | [code](/examples/minimax_h3_image/model_training/full/MiniMax-H3-Image.sh) | [code](/examples/minimax_h3_image/model_training/validate_full/MiniMax-H3-Image.py) | [code](/examples/minimax_h3_image/model_training/lora/MiniMax-H3-Image.sh) | [code](/examples/minimax_h3_image/model_training/validate_lora/MiniMax-H3-Image.py) |
 | Fibo 2.0 | [briaai/Fibo-2.0-dev](https://www.modelscope.cn/models/briaai/Fibo-2.0-dev) | [code](/examples/fibo/model_inference/Fibo-2.0-dev.py) | [code](/examples/fibo/model_inference_low_vram/Fibo-2.0-dev.py) | [code](/examples/fibo/model_training/full/Fibo-2.0-dev-T2I.sh) | [code](/examples/fibo/model_training/validate_full/Fibo-2.0-dev-T2I.py) | [code](/examples/fibo/model_training/lora/Fibo-2.0-dev-T2I.sh) | [code](/examples/fibo/model_training/validate_lora/Fibo-2.0-dev-T2I.py) |
 | Fibo 2.0 | [briaai/Fibo-2.0-edit-turbo](https://www.modelscope.cn/models/briaai/Fibo-2.0-edit-turbo) | [code](/examples/fibo/model_inference/Fibo-2.0-edit-turbo.py) | [code](/examples/fibo/model_inference_low_vram/Fibo-2.0-edit-turbo.py) | - | - | - | - |
 | Qwen-Image-2.1 | [Qwen/Qwen-Image-2.1](https://www.modelscope.cn/models/Qwen/Qwen-Image-2.1) | [code](/examples/qwen_image_21/model_inference/Qwen-Image-2.1.py) | [code](/examples/qwen_image_21/model_inference_low_vram/Qwen-Image-2.1.py) | [code](/examples/qwen_image_21/model_training/full/Qwen-Image-2.1.sh) | [code](/examples/qwen_image_21/model_training/validate_full/Qwen-Image-2.1.py) | [code](/examples/qwen_image_21/model_training/lora/Qwen-Image-2.1.sh) | [code](/examples/qwen_image_21/model_training/validate_lora/Qwen-Image-2.1.py) |
@@ -418,6 +420,7 @@ Quick start: experience popular and the latest models:
 Model overview:
 
 - Image generation
+    - MiniMax-H3-Image: [Documentation](https://diffsynth-studio-doc.readthedocs.io/en/latest/Model_Details/MiniMax-H3-Image.html), [Example code](/examples/minimax_h3_image/)
     - Fibo 2.0: [Documentation](https://diffsynth-studio-doc.readthedocs.io/en/latest/Model_Details/Fibo.html), [Example code](/examples/fibo/)
     - MOSS-Image: [Documentation](https://diffsynth-studio-doc.readthedocs.io/en/latest/Model_Details/MOSS-Image.html), [Example code](/examples/moss_image/)
     - FLUX.3: [Documentation](https://diffsynth-studio-doc.readthedocs.io/en/latest/Model_Details/FLUX.3.html), [Example code](/examples/flux3/)
@@ -707,6 +710,7 @@ https://github.com/Artiprocher/DiffSynth-Studio/assets/35051019/59fb2f7b-8de0-44
 
 | Architecture | Model ID | Inference | Low VRAM Inference | Full Training | Validation After Full Training | LoRA Training | Validation After LoRA Training |
 |-|-|-|-|-|-|-|-|
+| MiniMax-H3-Image | [MiniMax/MiniMax-H3-Image](https://www.modelscope.cn/models/MiniMax/MiniMax-H3-Image) | [code](/examples/minimax_h3_image/model_inference/MiniMax-H3-Image.py) | [code](/examples/minimax_h3_image/model_inference_low_vram/MiniMax-H3-Image.py) | [code](/examples/minimax_h3_image/model_training/full/MiniMax-H3-Image.sh) | [code](/examples/minimax_h3_image/model_training/validate_full/MiniMax-H3-Image.py) | [code](/examples/minimax_h3_image/model_training/lora/MiniMax-H3-Image.sh) | [code](/examples/minimax_h3_image/model_training/validate_lora/MiniMax-H3-Image.py) |
 | Fibo 2.0 | [briaai/Fibo-2.0-dev](https://www.modelscope.cn/models/briaai/Fibo-2.0-dev) | [code](/examples/fibo/model_inference/Fibo-2.0-dev.py) | [code](/examples/fibo/model_inference_low_vram/Fibo-2.0-dev.py) | [code](/examples/fibo/model_training/full/Fibo-2.0-dev-T2I.sh) | [code](/examples/fibo/model_training/validate_full/Fibo-2.0-dev-T2I.py) | [code](/examples/fibo/model_training/lora/Fibo-2.0-dev-T2I.sh) | [code](/examples/fibo/model_training/validate_lora/Fibo-2.0-dev-T2I.py) |
 | Fibo 2.0 | [briaai/Fibo-2.0-edit-turbo](https://www.modelscope.cn/models/briaai/Fibo-2.0-edit-turbo) | [code](/examples/fibo/model_inference/Fibo-2.0-edit-turbo.py) | [code](/examples/fibo/model_inference_low_vram/Fibo-2.0-edit-turbo.py) | - | - | - | - |
 | MOSS-Image | [openmoss/MOSS-Image-Base](https://www.modelscope.cn/models/openmoss/MOSS-Image-Base) | [code](/examples/moss_image/model_inference/MOSS-Image-Base.py) | [code](/examples/moss_image/model_inference_low_vram/MOSS-Image-Base.py) | [code](/examples/moss_image/model_training/full/MOSS-Image-Base.sh) | [code](/examples/moss_image/model_training/validate_full/MOSS-Image-Base.py) | [code](/examples/moss_image/model_training/lora/MOSS-Image-Base.sh) | [code](/examples/moss_image/model_training/validate_lora/MOSS-Image-Base.py) |

{diffsynth-2.1.12 → diffsynth-2.2.1}/README.md

@@ -41,6 +41,8 @@ See also:
 
 > Currently, the development personnel of this project are limited, with most of the work handled by [Artiprocher](https://github.com/Artiprocher) and [mi804](https://github.com/mi804). Therefore, the progress of new feature development will be relatively slow, and the speed of responding to and resolving issues is limited. We apologize for this and ask developers to understand.
 
+- **November 10, 2026** We now support MiniMax-H3-Image, the image model of the MiniMax-H3 family, including unified text-to-image and reference-image editing inference, low VRAM inference, and full/LoRA training. The checkpoint is distilled to six sampling steps with one transformer output head per step, and it generates about 16-megapixel images without an upscaler. See the [documentation](/docs/en/Model_Details/MiniMax-H3-Image.md) and [example code](/examples/minimax_h3_image/).
+
 - **October 21, 2026** We now support Fibo 2.0, including unified text-to-image and image editing inference, masked and multi-reference editing, text rendering, low VRAM inference, and full/LoRA training for both text-to-image and image editing on [Fibo-2.0-dev](https://www.modelscope.cn/models/briaai/Fibo-2.0-dev). See the [documentation](/docs/en/Model_Details/Fibo.md) and [example code](/examples/fibo/).
 
 - **October 20, 2026** We now support MOSS-Image, including inference and training for [MOSS-Image-Base](https://modelscope.cn/models/openmoss/MOSS-Image-Base) and [MOSS-Image-SFT](https://modelscope.cn/models/openmoss/MOSS-Image-SFT). In addition, we provide an [EntroPack low-bit quantized model](https://www.modelscope.cn/models/DiffSynth-Studio/MOSS-Image-EntroPack), which can run the 3.5-bit quantized model with as little as 4 GB of VRAM; and an [accelerated model](https://modelscope.cn/models/DiffSynth-Studio/MOSS-Image-SFT-TurboLoRA) distillation-accelerated via [OPD training](https://diffsynth-studio-doc.readthedocs.io/en/latest/Training/Distillation.html), supporting 8-step ultra-fast inference.
@@ -331,6 +333,7 @@ Quick start: experience popular and the latest models:
 
 | Architecture | Model ID | Inference | Low VRAM Inference | Full Training | Validation After Full Training | LoRA Training | Validation After LoRA Training |
 |-|-|-|-|-|-|-|-|
+| MiniMax-H3-Image | [MiniMax/MiniMax-H3-Image](https://www.modelscope.cn/models/MiniMax/MiniMax-H3-Image) | [code](/examples/minimax_h3_image/model_inference/MiniMax-H3-Image.py) | [code](/examples/minimax_h3_image/model_inference_low_vram/MiniMax-H3-Image.py) | [code](/examples/minimax_h3_image/model_training/full/MiniMax-H3-Image.sh) | [code](/examples/minimax_h3_image/model_training/validate_full/MiniMax-H3-Image.py) | [code](/examples/minimax_h3_image/model_training/lora/MiniMax-H3-Image.sh) | [code](/examples/minimax_h3_image/model_training/validate_lora/MiniMax-H3-Image.py) |
 | Fibo 2.0 | [briaai/Fibo-2.0-dev](https://www.modelscope.cn/models/briaai/Fibo-2.0-dev) | [code](/examples/fibo/model_inference/Fibo-2.0-dev.py) | [code](/examples/fibo/model_inference_low_vram/Fibo-2.0-dev.py) | [code](/examples/fibo/model_training/full/Fibo-2.0-dev-T2I.sh) | [code](/examples/fibo/model_training/validate_full/Fibo-2.0-dev-T2I.py) | [code](/examples/fibo/model_training/lora/Fibo-2.0-dev-T2I.sh) | [code](/examples/fibo/model_training/validate_lora/Fibo-2.0-dev-T2I.py) |
 | Fibo 2.0 | [briaai/Fibo-2.0-edit-turbo](https://www.modelscope.cn/models/briaai/Fibo-2.0-edit-turbo) | [code](/examples/fibo/model_inference/Fibo-2.0-edit-turbo.py) | [code](/examples/fibo/model_inference_low_vram/Fibo-2.0-edit-turbo.py) | - | - | - | - |
 | Qwen-Image-2.1 | [Qwen/Qwen-Image-2.1](https://www.modelscope.cn/models/Qwen/Qwen-Image-2.1) | [code](/examples/qwen_image_21/model_inference/Qwen-Image-2.1.py) | [code](/examples/qwen_image_21/model_inference_low_vram/Qwen-Image-2.1.py) | [code](/examples/qwen_image_21/model_training/full/Qwen-Image-2.1.sh) | [code](/examples/qwen_image_21/model_training/validate_full/Qwen-Image-2.1.py) | [code](/examples/qwen_image_21/model_training/lora/Qwen-Image-2.1.sh) | [code](/examples/qwen_image_21/model_training/validate_lora/Qwen-Image-2.1.py) |
@@ -345,6 +348,7 @@ Quick start: experience popular and the latest models:
 Model overview:
 
 - Image generation
+    - MiniMax-H3-Image: [Documentation](https://diffsynth-studio-doc.readthedocs.io/en/latest/Model_Details/MiniMax-H3-Image.html), [Example code](/examples/minimax_h3_image/)
     - Fibo 2.0: [Documentation](https://diffsynth-studio-doc.readthedocs.io/en/latest/Model_Details/Fibo.html), [Example code](/examples/fibo/)
     - MOSS-Image: [Documentation](https://diffsynth-studio-doc.readthedocs.io/en/latest/Model_Details/MOSS-Image.html), [Example code](/examples/moss_image/)
     - FLUX.3: [Documentation](https://diffsynth-studio-doc.readthedocs.io/en/latest/Model_Details/FLUX.3.html), [Example code](/examples/flux3/)
@@ -634,6 +638,7 @@ https://github.com/Artiprocher/DiffSynth-Studio/assets/35051019/59fb2f7b-8de0-44
 
 | Architecture | Model ID | Inference | Low VRAM Inference | Full Training | Validation After Full Training | LoRA Training | Validation After LoRA Training |
 |-|-|-|-|-|-|-|-|
+| MiniMax-H3-Image | [MiniMax/MiniMax-H3-Image](https://www.modelscope.cn/models/MiniMax/MiniMax-H3-Image) | [code](/examples/minimax_h3_image/model_inference/MiniMax-H3-Image.py) | [code](/examples/minimax_h3_image/model_inference_low_vram/MiniMax-H3-Image.py) | [code](/examples/minimax_h3_image/model_training/full/MiniMax-H3-Image.sh) | [code](/examples/minimax_h3_image/model_training/validate_full/MiniMax-H3-Image.py) | [code](/examples/minimax_h3_image/model_training/lora/MiniMax-H3-Image.sh) | [code](/examples/minimax_h3_image/model_training/validate_lora/MiniMax-H3-Image.py) |
 | Fibo 2.0 | [briaai/Fibo-2.0-dev](https://www.modelscope.cn/models/briaai/Fibo-2.0-dev) | [code](/examples/fibo/model_inference/Fibo-2.0-dev.py) | [code](/examples/fibo/model_inference_low_vram/Fibo-2.0-dev.py) | [code](/examples/fibo/model_training/full/Fibo-2.0-dev-T2I.sh) | [code](/examples/fibo/model_training/validate_full/Fibo-2.0-dev-T2I.py) | [code](/examples/fibo/model_training/lora/Fibo-2.0-dev-T2I.sh) | [code](/examples/fibo/model_training/validate_lora/Fibo-2.0-dev-T2I.py) |
 | Fibo 2.0 | [briaai/Fibo-2.0-edit-turbo](https://www.modelscope.cn/models/briaai/Fibo-2.0-edit-turbo) | [code](/examples/fibo/model_inference/Fibo-2.0-edit-turbo.py) | [code](/examples/fibo/model_inference_low_vram/Fibo-2.0-edit-turbo.py) | - | - | - | - |
 | MOSS-Image | [openmoss/MOSS-Image-Base](https://www.modelscope.cn/models/openmoss/MOSS-Image-Base) | [code](/examples/moss_image/model_inference/MOSS-Image-Base.py) | [code](/examples/moss_image/model_inference_low_vram/MOSS-Image-Base.py) | [code](/examples/moss_image/model_training/full/MOSS-Image-Base.sh) | [code](/examples/moss_image/model_training/validate_full/MOSS-Image-Base.py) | [code](/examples/moss_image/model_training/lora/MOSS-Image-Base.sh) | [code](/examples/moss_image/model_training/validate_lora/MOSS-Image-Base.py) |

{diffsynth-2.1.12 → diffsynth-2.2.1}/diffsynth/configs/model_configs.py

@@ -1743,6 +1743,29 @@ minimax_music3_series = [
     },
 ]
 
+minimax_h3_image_series = [
+    {
+        # Example: ModelConfig(model_id="MiniMax/MiniMax-H3-Image", origin_file_pattern="transformer/model*.safetensors")
+        # Example: ModelConfig(path="./models/train/MiniMax-H3-Image-full-split/epoch-4.safetensors")
+        "model_hash": "cbc654a2ff91b53b57e2556904bdbc5e",
+        "model_name": "minimax_h3_image_dit",
+        "model_class": "diffsynth.models.minimax_h3_image_dit.MiniMaxH3ImageDiT",
+    },
+    {
+        # Example: ModelConfig(model_id="MiniMax/MiniMax-H3-Image", origin_file_pattern="pdd_fused_heads.safetensors")
+        "model_hash": "af836167abebf4a4701ab436f291a972",
+        "model_name": "minimax_h3_image_pdd_heads",
+        "model_class": "diffsynth.models.minimax_h3_image_dit.MiniMaxH3ImagePDDHeads",
+    },
+    {
+        # Example: ModelConfig(model_id="MiniMax/MiniMax-H3-Image", origin_file_pattern="vae/diffusion_pytorch_model.safetensors")
+        "model_hash": "e16b57e3180a14235ffcaad014ce2772",
+        "model_name": "minimax_h3_image_vae",
+        "model_class": "diffsynth.models.minimax_h3_image_vae.MiniMaxH3ImageVAE",
+        "state_dict_converter": "diffsynth.utils.state_dict_converters.minimax_h3_image_vae.MiniMaxH3ImageVAEStateDictConverter",
+    },
+]
+
 sensenova_u1_series = [
     {
         # Example: ModelConfig(model_id="SenseNova/SenseNova-U1.5-8B-MoT", origin_file_pattern="model*.safetensors")
@@ -1857,5 +1880,5 @@ MODEL_CONFIGS = (
     stable_diffusion_xl_series + stable_diffusion_series + qwen_image_series + qwen_video_edit_series + wan_series + flux_series + flux2_series + ernie_image_series
     + z_image_series + ltx2_series + anima_series + mova_series + joyai_image_series + boogu_image_series + ace_step_series + hidream_o1_image_series
     + image_metrics_series + ideogram4_series + krea2_series + lingbot_video_series + minimax_h3_series + minimax_music3_series + sensenova_u1_series
-    + yue2_series + qwen_image_21_series + flux3_series + moss_image_series + fibo_series
+    + yue2_series + qwen_image_21_series + flux3_series + moss_image_series + fibo_series + minimax_h3_image_series
 )

{diffsynth-2.1.12 → diffsynth-2.2.1}/diffsynth/configs/vram_management_module_maps.py

@@ -10,6 +10,24 @@ flux_general_vram_config = {
 }
 
 VRAM_MANAGEMENT_MODULE_MAPS = {
+    "diffsynth.models.minimax_h3_image_dit.MiniMaxH3ImageDiT": {
+        "torch.nn.Linear": "diffsynth.core.vram.layers.AutoWrappedLinear",
+        "torch.nn.RMSNorm": "diffsynth.core.vram.layers.AutoWrappedModule",
+    },
+    "diffsynth.models.minimax_h3_image_dit.MiniMaxH3ImagePDDHeads": {
+        "diffsynth.models.minimax_h3_image_dit._MiniMaxH3ImagePDDHeadParameters": "diffsynth.core.vram.layers.AutoWrappedModule",
+    },
+    "diffsynth.models.minimax_h3_image_vae.MiniMaxH3ImageVAE": {
+        "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.Conv2d": "diffsynth.core.vram.layers.AutoWrappedModule",
+        "torch.nn.Conv3d": "diffsynth.core.vram.layers.AutoWrappedModule",
+        "torch.nn.GroupNorm": "diffsynth.core.vram.layers.AutoWrappedModule",
+        "diffsynth.models.minimax_h3_video_vae.BaseConv3d": "diffsynth.core.vram.layers.AutoWrappedModule",
+        "diffsynth.models.minimax_h3_video_vae.WarpedTensor": "diffsynth.core.vram.layers.AutoWrappedModule",
+        "diffsynth.models.minimax_h3_video_vae.TemporalIsolatedSpatialGroupNorm": "diffsynth.core.vram.layers.AutoWrappedModule",
+    },
     "diffsynth.models.minimax_h3_dit.MiniMaxH3DiT": {
         "torch.nn.Linear": "diffsynth.core.vram.layers.AutoWrappedLinear",
         "torch.nn.RMSNorm": "diffsynth.core.vram.layers.AutoWrappedModule",
@@ -115,6 +133,9 @@ VRAM_MANAGEMENT_MODULE_MAPS = {
         "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",
+        "diffsynth.models.wan_video_animate_adapter.Blur": "diffsynth.core.vram.layers.AutoWrappedModule",
+        "diffsynth.models.wan_video_animate_adapter.EqualConv2d": "diffsynth.core.vram.layers.AutoWrappedModule",
+        "diffsynth.models.wan_video_animate_adapter.Direction": "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",
@@ -150,6 +171,7 @@ VRAM_MANAGEMENT_MODULE_MAPS = {
         "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",
+        "torch.nn.MultiheadAttention": "diffsynth.core.vram.layers.AutoWrappedModule",
     },
     "diffsynth.models.wan_video_image_encoder.WanImageEncoder": {
         "diffsynth.models.wan_video_image_encoder.VisionTransformer": "diffsynth.core.vram.layers.AutoWrappedModule",

diffsynth-2.2.1/diffsynth/core/attention/__init__.py

@@ -0,0 +1,2 @@
+from .attention import attention_forward, list_attention_implementations
+from .flex_attention import FLEX_ATTN_AVAILABLE

diffsynth-2.2.1/diffsynth/core/attention/attention.py

@@ -0,0 +1,51 @@
+import os
+from .registry import (
+    get_attention_implementation,
+    list_attention_implementations,
+    register_attention_implementation,
+)
+from .flash_attention import CUSTOMIZED_FA_KERNEL_AVAILABLE, FLASH_ATTN_4_AVAILABLE, FLASH_ATTN_3_AVAILABLE, FLASH_ATTN_2_AVAILABLE
+from .sage_attention import SAGE_ATTN_AVAILABLE
+from .xformers_attention import XFORMERS_AVAILABLE
+from .torch_attention import torch_sdpa, torch_sdpa_sliding_window
+from .flex_attention import flex_attention
+
+
+def initialize_attention_priority():
+    if os.environ.get('DIFFSYNTH_ATTENTION_IMPLEMENTATION') is not None:
+        return os.environ.get('DIFFSYNTH_ATTENTION_IMPLEMENTATION').strip().lower()
+    else:
+        for name in ["customized_fa_kernel", "flash_attention_4", "flash_attention_3", "flash_attention_2", "xformers"]:
+            if name in list_attention_implementations():
+                return name
+        return "torch"
+
+
+ATTENTION_IMPLEMENTATION = initialize_attention_priority()
+
+
+def attention_forward(q, k, v, q_pattern="b n s d", k_pattern="b n s d", v_pattern="b n s d", out_pattern="b n s d", dims=None, attn_mask=None, scale=None, is_causal=False, compatibility_mode=False, window_size=None, use_flex=False, score_mod=None):
+    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)
+
+    if compatibility_mode:
+        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, window_size=window_size)
+
+    backend = get_attention_implementation(ATTENTION_IMPLEMENTATION)
+
+    if attn_mask is not None and not backend.supports_mask:
+        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)
+    if window_size is not None and not backend.supports_window:
+        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)
+    if is_causal and not backend.supports_causal:
+        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)
+
+    kwargs = {"scale": scale}
+    if backend.supports_causal:
+        kwargs["is_causal"] = is_causal
+    if backend.supports_mask:
+        kwargs["attn_mask"] = attn_mask
+    if backend.supports_window:
+        kwargs["window_size"] = window_size
+    return backend.func(q, k, v, q_pattern, k_pattern, v_pattern, out_pattern, dims, **kwargs)

diffsynth-2.2.1/diffsynth/core/attention/common.py

@@ -0,0 +1,27 @@
+import inspect
+
+import torch
+from einops import rearrange
+
+try:
+    TORCH_SUPPORT_GQA = "enable_gqa" in inspect.signature(torch.nn.functional.scaled_dot_product_attention).parameters
+except Exception:
+    TORCH_SUPPORT_GQA = False
+
+
+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

diffsynth-2.2.1/diffsynth/core/attention/flash_attention.py

@@ -0,0 +1,118 @@
+import os
+
+from .common import rearrange_qkv, rearrange_out
+from .registry import register_attention_implementation
+
+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:
+    from flash_attn_3 import flash_attn_interface
+    FLASH_ATTN_3_AVAILABLE = True
+except Exception:
+    FLASH_ATTN_3_AVAILABLE = False
+
+try:
+    import flash_attn
+    FLASH_ATTN_2_AVAILABLE = hasattr(flash_attn, "flash_attn_func")
+except Exception:
+    FLASH_ATTN_2_AVAILABLE = False
+
+
+def _require(flag, name):
+    if not flag:
+        raise RuntimeError(f"{name} is not available.")
+
+
+def flash_attention_4(q, k, v, q_pattern="b n s d", k_pattern="b n s d", v_pattern="b n s d", out_pattern="b n s d", dims=None, scale=None, is_causal=False, window_size=None):
+    _require(FLASH_ATTN_4_AVAILABLE, "FlashAttention-4")
+    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, k, v, q_pattern="b n s d", k_pattern="b n s d", v_pattern="b n s d", out_pattern="b n s d", dims=None, scale=None, is_causal=False, window_size=None):
+    _require(FLASH_ATTN_3_AVAILABLE, "FlashAttention-3")
+    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, k, v, q_pattern="b n s d", k_pattern="b n s d", v_pattern="b n s d", out_pattern="b n s d", dims=None, scale=None, is_causal=False, window_size=None):
+    _require(FLASH_ATTN_2_AVAILABLE, "FlashAttention-2")
+    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, k, v, q_pattern="b n s d", k_pattern="b n s d", v_pattern="b n s d", out_pattern="b n s d", dims=None, scale=None, is_causal=False, window_size=None):
+    _require(CUSTOMIZED_FA_KERNEL_AVAILABLE, "customized FlashAttention kernel")
+    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
+
+
+if CUSTOMIZED_FA_KERNEL_AVAILABLE:
+    register_attention_implementation(
+        "customized_fa_kernel",
+        flash_attention_customized_kernel,
+        supports_causal=True,
+        supports_window=True,
+        supports_gqa=True,
+    )
+if FLASH_ATTN_4_AVAILABLE:
+    register_attention_implementation(
+        "flash_attention_4",
+        flash_attention_4,
+        supports_causal=True,
+        supports_window=True,
+        supports_gqa=True,
+    )
+if FLASH_ATTN_3_AVAILABLE:
+    register_attention_implementation(
+        "flash_attention_3",
+        flash_attention_3,
+        supports_causal=True,
+        supports_window=True,
+        supports_gqa=True,
+    )
+if FLASH_ATTN_2_AVAILABLE:
+    register_attention_implementation(
+        "flash_attention_2",
+        flash_attention_2,
+        supports_causal=True,
+        supports_window=True,
+        supports_gqa=True,
+    )

diffsynth-2.2.1/diffsynth/core/attention/flex_attention.py

@@ -0,0 +1,28 @@
+import torch
+
+from .common import rearrange_qkv, rearrange_out
+
+try:
+    from torch.nn.attention.flex_attention import flex_attention as _flex_attention_func
+    FLEX_ATTN_AVAILABLE = True
+except Exception:
+    FLEX_ATTN_AVAILABLE = False
+
+# Compile flex attention lazily
+_compiled_flex_attention_func = None
+
+
+def get_flex_attention_func():
+    global _compiled_flex_attention_func
+    if _compiled_flex_attention_func is None:
+        _compiled_flex_attention_func = torch.compile(_flex_attention_func)
+    return _compiled_flex_attention_func
+
+
+def flex_attention(q, k, v, q_pattern="b n s d", k_pattern="b n s d", v_pattern="b n s d", out_pattern="b n s d", 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 = get_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

diffsynth-2.2.1/diffsynth/core/attention/registry.py

@@ -0,0 +1,48 @@
+from dataclasses import dataclass
+from typing import Callable, Optional
+
+
+@dataclass(frozen=True)
+class AttentionBackend:
+    func: Callable
+    supports_causal: bool = False
+    supports_window: bool = False
+    supports_mask: bool = False
+    supports_gqa: bool = False
+
+
+_ATTENTION_BACKENDS = {}
+
+
+def register_attention_implementation(
+    name: str,
+    func: Optional[Callable] = None,
+    supports_causal: bool = False,
+    supports_window: bool = False,
+    supports_mask: bool = False,
+    supports_gqa: bool = False,
+):
+    key = name.strip().lower()
+    _ATTENTION_BACKENDS[key] = AttentionBackend(
+        func=func,
+        supports_causal=supports_causal,
+        supports_window=supports_window,
+        supports_mask=supports_mask,
+        supports_gqa=supports_gqa,
+    )
+
+
+def get_attention_implementation(name: str) -> AttentionBackend:
+    key = name.strip().lower()
+    if key not in _ATTENTION_BACKENDS:
+        available = ", ".join(list_attention_implementations()) or "none"
+        raise ValueError(
+            f"Unknown attention implementation: {name!r}. "
+            f"Available implementations: {available}. "
+            f"Please check the environment variable `DIFFSYNTH_ATTENTION_IMPLEMENTATION`."
+        )
+    return _ATTENTION_BACKENDS[key]
+
+
+def list_attention_implementations():
+    return sorted(_ATTENTION_BACKENDS)

diffsynth-2.2.1/diffsynth/core/attention/sage_attention.py

@@ -0,0 +1,20 @@
+from .common import rearrange_qkv, rearrange_out
+from .registry import register_attention_implementation
+
+try:
+    from sageattention import sageattn
+    SAGE_ATTN_AVAILABLE = True
+except Exception:
+    SAGE_ATTN_AVAILABLE = False
+
+def sage_attention(q, k, v, q_pattern="b n s d", k_pattern="b n s d", v_pattern="b n s d", out_pattern="b n s d", dims=None, scale=None):
+    if not SAGE_ATTN_AVAILABLE:
+        raise RuntimeError("SageAttention is not available.")
+    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
+
+if SAGE_ATTN_AVAILABLE:
+    register_attention_implementation("sage_attention", sage_attention)

diffsynth-2.2.1/diffsynth/core/attention/torch_attention.py

@@ -0,0 +1,82 @@
+import torch
+from einops import repeat
+
+from .common import TORCH_SUPPORT_GQA, rearrange_qkv, rearrange_out
+from .registry import register_attention_implementation
+
+
+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, window_size=None):
+    if window_size is not None:
+        if is_causal or attn_mask is not None:
+            raise ValueError("Sliding window attention does not support `is_causal=True` or `attn_mask`.")
+        return torch_sdpa_sliding_window(q, k, v, window_size, q_pattern, k_pattern, v_pattern, out_pattern, dims, scale=scale)
+
+    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
+
+
+register_attention_implementation(
+    "torch",
+    torch_sdpa,
+    supports_window=True,
+    supports_causal=True,
+    supports_mask=True,
+    supports_gqa=True,
+)

diffsynth-2.2.1/diffsynth/core/attention/xformers_attention.py

@@ -0,0 +1,22 @@
+from .common import rearrange_qkv, rearrange_out
+from .registry import register_attention_implementation
+
+try:
+    import xformers.ops as xops
+    XFORMERS_AVAILABLE = True
+except Exception:
+    XFORMERS_AVAILABLE = False
+
+
+def xformers_attention(q, k, v, q_pattern="b n s d", k_pattern="b n s d", v_pattern="b n s d", out_pattern="b n s d", dims=None, scale=None):
+    if not XFORMERS_AVAILABLE:
+        raise RuntimeError("xformers is not available.")
+    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
+
+
+if XFORMERS_AVAILABLE:
+    register_attention_implementation("xformers", xformers_attention)