diffsynth 2.2.0__tar.gz → 2.2.1__tar.gz

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: diffsynth
-Version: 2.2.0
+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.2.0 → 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.2.0 → 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.2.0 → 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.0 → diffsynth-2.2.1}/diffsynth/core/attention/attention.py

@@ -15,7 +15,7 @@ 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", "sage_attention", "xformers"]:
+        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"

{diffsynth-2.2.0 → diffsynth-2.2.1}/diffsynth/core/quant/backends/entropack.py

@@ -45,13 +45,13 @@ if ENTROPACK_AVAILABLE:
     @dataclass
     class EntroPackINT8CodedConfig(BackendConfig, ep.LatticeRANSConfig):
         linear_cls = ep.CompressedINT8Linear
-        use_convrot: bool = False
+        use_convrot: bool = True
         convrot_groupsize: int = 256
 
     @dataclass
     class EntroPackINT8RawConfig(BackendConfig, ep.RawConfig):
         linear_cls = ep.CompressedINT8Linear
-        use_convrot: bool = False
+        use_convrot: bool = True
         convrot_groupsize: int = 256
 
 

diffsynth-2.2.1/diffsynth/diffusion/dmd2.py

@@ -0,0 +1,355 @@
+import math, torch
+import torch.nn.functional as F
+from tqdm import tqdm
+from accelerate.utils import broadcast, DistributedType
+from diffsynth.core import load_state_dict
+from .runner import exclude_quantized_params_from_ddp_sync, get_optimizer_class, save_training_args
+from .training_module import DiffusionTrainingModule
+
+
+class DMD2ClassifierHead(torch.nn.Module):
+    def __init__(self, input_channels, channels=1280):
+        super().__init__()
+        groups = math.gcd(32, channels)
+        if channels // groups < 2:
+            raise ValueError("The classifier needs at least two channels per GroupNorm group.")
+        layers = [torch.nn.AdaptiveAvgPool2d(32), torch.nn.Conv2d(input_channels, channels, 1)]
+        for stride, padding in ((2, 1), (2, 1), (2, 1), (4, 0)):
+            layers.extend([torch.nn.Conv2d(channels, channels, 4, stride, padding),
+                           torch.nn.GroupNorm(groups, channels), torch.nn.SiLU()])
+        self.layers = torch.nn.Sequential(*layers, torch.nn.Conv2d(channels, 1, 1))
+
+    def forward(self, features):
+        return self.layers(features.float()).flatten(1)
+
+
+class DMD2FeatureExtractor:
+    def __init__(self, model, module_path):
+        self.module = model.get_submodule(module_path)
+
+    def run(self, forward, *args, **kwargs):
+        features = []
+        handle = self.module.register_forward_hook(lambda module, inputs, output: features.append(output))
+        try:
+            forward(*args, **kwargs)
+        finally:
+            handle.remove()
+        if len(features) != 1:
+            raise ValueError("The GAN feature module must run exactly once per forward.")
+        return features[0]
+
+
+class DMD2TrainingModule(DiffusionTrainingModule):
+    def switch_pipe_to_dmd2_training_mode(
+        self, pipe,
+        student_sigmas, student_timesteps, teacher_cfg=3.,
+        min_timestep_dmd=20, max_timestep_dmd=981,
+        lora_base_model=None, lora_target_modules="", lora_rank=128, lora_dtype=None, lora_checkpoint=None,
+        preset_lora_path=None, preset_lora_model=None,
+        use_gan=True, generator_gan_loss_weight=0.005, discriminator_gan_loss_weight=0.01,
+        classifier_channels=1280, feature_module_of_dit_for_gan=None,
+        trainable_models=None, noise_sampling="uniform",
+    ):
+        # Training config
+        self.gan = use_gan
+        self.teacher_cfg = teacher_cfg
+        self.lora_base_model = lora_base_model
+        self.trainable_models = trainable_models
+        self.generator_gan_loss_weight = generator_gan_loss_weight
+        self.discriminator_gan_loss_weight = discriminator_gan_loss_weight
+        self.min_timestep_dmd = min_timestep_dmd
+        self.max_timestep_dmd = max_timestep_dmd
+        self.noise_sampling = noise_sampling
+        if lora_base_model is None and trainable_models is None:
+            raise ValueError("DMD2 requires trainable_models or lora_base_model.")
+
+        # Freeze models
+        pipe.freeze_except([] if trainable_models is None else trainable_models.split(","))
+        if preset_lora_path is not None:
+            pipe.load_lora(getattr(pipe, preset_lora_model), preset_lora_path)
+        model = getattr(pipe, lora_base_model or trainable_models)
+        self.adapter_layers = []
+
+        if self.trainable_models is not None and self.lora_base_model is None:
+            from copy import deepcopy
+            self.fake_model = deepcopy(model)
+            self.teacher_model = deepcopy(model).eval().requires_grad_(False)
+
+        # Add LoRA
+        if lora_base_model is not None:
+            from peft.tuners.tuners_utils import BaseTunerLayer
+            lora_dtype = pipe.torch_dtype if lora_dtype is None else {"bf16": torch.bfloat16, "fp32": torch.float32}[lora_dtype]
+            model = self.add_lora_to_model(model, lora_target_modules, lora_rank, upcast_dtype=lora_dtype)
+            model = self.add_lora_to_model(model, lora_target_modules, lora_rank, upcast_dtype=lora_dtype, adapter_name="fake")
+            setattr(pipe, lora_base_model, model)
+
+            # Load student LoRA
+            if lora_checkpoint is not None:
+                if isinstance(lora_checkpoint, str):
+                    lora = load_state_dict(lora_checkpoint)
+                else:
+                    lora = lora_checkpoint
+                lora_loader = pipe.lora_loader(torch_dtype=pipe.torch_dtype, device=pipe.device)
+                lora = lora_loader.convert_state_dict(lora)
+                lora = self.mapping_lora_state_dict(lora)
+                load_result = model.load_state_dict(lora, strict=False)
+                print(f"LoRA checkpoint loaded. Total {len(lora)} keys")
+                if len(load_result[1]) > 0:
+                    print(f"Warning, LoRA key mismatch! Unexpected keys in LoRA checkpoint: {load_result[1]}")
+
+            self.adapter_layers = [layer for layer in model.modules() if isinstance(layer, BaseTunerLayer)]
+            for layer in self.adapter_layers:
+                layer.set_adapter(["default", "fake"])
+
+        # Initialize GAN
+        self.classifier = self.feature_extractor = None
+        if self.gan:
+            if feature_module_of_dit_for_gan is None:
+                raise ValueError("GAN training requires feature_module_of_dit_for_gan.")
+            self.classifier = DMD2ClassifierHead(self.gan_feature_channels, classifier_channels).to(pipe.device)
+            self.feature_extractor = DMD2FeatureExtractor(self.fake_model if self.trainable_models is not None and self.lora_base_model is None else model, feature_module_of_dit_for_gan)
+
+        # Set student schedule
+        pipe.scheduler.sigmas = self.parse_float_list(student_sigmas)
+        pipe.scheduler.timesteps = self.parse_float_list(student_timesteps)
+        pipe.scheduler.training = True
+
+    def add_lora_to_model(self, model, target_modules, lora_rank, lora_alpha=None, upcast_dtype=None, adapter_name="default"):
+        from peft import LoraConfig, inject_adapter_in_model
+        if lora_alpha is None:
+            lora_alpha = lora_rank
+        target_modules = self.parse_lora_target_modules(model, target_modules)
+        if isinstance(target_modules, list) and len(target_modules) == 1:
+            target_modules = target_modules[0]
+        lora_config = LoraConfig(r=lora_rank, lora_alpha=lora_alpha, target_modules=target_modules)
+        model = inject_adapter_in_model(lora_config, model, adapter_name=adapter_name)
+        if upcast_dtype is not None:
+            for param in model.parameters():
+                if param.requires_grad:
+                    param.data = param.to(upcast_dtype)
+        return model
+
+    def parse_float_list(self, x):
+        return torch.tensor(list(map(float, x.split(",")))) if isinstance(x, str) else x
+
+    def select_model(self, model_role):
+        if self.trainable_models is not None and self.lora_base_model is None:
+            return {"generator": getattr(self.pipe, self.trainable_models), "fake": self.fake_model, "teacher": self.teacher_model}[model_role]
+        adapter = {"generator": "default", "fake": "fake", "teacher": []}[model_role]
+        for layer in self.adapter_layers:
+            layer.set_adapter(adapter)
+        return getattr(self.pipe, self.lora_base_model)
+
+    def generator_parameters(self):
+        if self.trainable_models is not None and self.lora_base_model is None:
+            return list(getattr(self.pipe, self.trainable_models).parameters())
+        model = getattr(self.pipe, self.lora_base_model)
+        return [p for name, p in model.named_parameters() if ".lora_A.default." in name or ".lora_B.default." in name]
+
+    def critic_parameters(self):
+        if self.trainable_models is not None and self.lora_base_model is None:
+            parameters = list(self.fake_model.parameters())
+        else:
+            model = getattr(self.pipe, self.lora_base_model)
+            parameters = [p for name, p in model.named_parameters() if ".lora_A.fake." in name or ".lora_B.fake." in name]
+        return parameters + (list(self.classifier.parameters()) if self.gan else [])
+
+    def prepare_batch(self, data):
+        """Subclasses must return (inputs_shared, inputs_posi, inputs_nega) for their pipeline."""
+        raise NotImplementedError("Implement prepare_batch in the model's training module.")
+
+    def predict(self, x, timestep, inputs, model_role, features=False):
+        model = self.select_model(model_role)
+        inputs_shared, inputs_posi, inputs_nega = inputs
+        inputs_shared = dict(inputs_shared, latents=x.to(self.pipe.torch_dtype))
+        models = {name: getattr(self.pipe, name) for name in self.pipe.in_iteration_models}
+        models[self.lora_base_model or self.trainable_models] = model
+        if features:
+            output = self.feature_extractor.run(self.pipe.model_fn, **models, **inputs_shared, **inputs_posi, timestep=timestep)
+            output = self.reshape_gan_features(output, inputs_shared)
+        else:
+            output = self.pipe.cfg_guided_model_fn(
+                lambda **kwargs: self.pipe.model_fn(**kwargs).float(),
+                self.teacher_cfg if model_role == "teacher" else 1.,
+                inputs_shared, inputs_posi, inputs_nega, **models, timestep=timestep,
+            )
+        return output.float()
+
+    def rollout(self, inputs, step, track_grad=False):
+        x = inputs[0]["noise"].float()
+        for index in range(step + 1):
+            with torch.set_grad_enabled(track_grad and index == step):
+                x = x.to(self.pipe.torch_dtype).float()
+                timestep = self.pipe.scheduler.timesteps[index].to(x.device).reshape(1)
+                velocity = self.predict(x, timestep, inputs, "generator")
+                x = self.pipe.scheduler.step(velocity, timestep, x, to_final=index == step)
+        return x
+
+    def sample_noise(self, clean, distribution_matching, min_timestep_dmd=20, max_timestep_dmd=981, noise_sampling=None):
+        noise_sampling = self.noise_sampling if noise_sampling is None else noise_sampling
+        if noise_sampling == "beta":
+            beta = torch.distributions.Beta(torch.tensor(2.5, device=clean.device), torch.tensor(1., device=clean.device))
+            level = beta.sample((clean.shape[0],)).clamp((min_timestep_dmd + 1) / 1000, max_timestep_dmd / 1000)
+        else:
+            low, high = (min_timestep_dmd, max_timestep_dmd) if distribution_matching else (0, 1000)
+            level = (torch.randint(low, high, (clean.shape[0],), device=clean.device).float() + 1) / 1000
+        noise = torch.randn_like(clean)
+        sigma = level.reshape(-1, *([1] * (clean.ndim - 1)))
+        return (1 - sigma) * clean + sigma * noise, level, noise
+
+    def teacher_clean(self, noised, inputs):
+        noisy, level, _ = noised
+        velocity = self.predict(noisy, level * 1000, inputs, "teacher")
+        return noisy - level.reshape(-1, *([1] * (noisy.ndim - 1))) * velocity
+
+    def fake_clean(self, noised, inputs):
+        noisy, level, _ = noised
+        velocity = self.predict(noisy, level * 1000, inputs, "fake")
+        return noisy - level.reshape(-1, *([1] * (noisy.ndim - 1))) * velocity
+
+    def logits(self, clean, inputs):
+        noisy, level, _ = self.sample_noise(clean, False, noise_sampling="uniform")
+        features = self.predict(noisy, level * 1000, inputs, "fake", features=True)
+        return self.classifier(features)
+
+    def export_trainable_state_dict(self, state_dict, remove_prefix=None):
+        generator_param_ids = {id(param) for param in self.generator_parameters()}
+        trainable_param_names = {name for name, param in self.named_parameters() if id(param) in generator_param_ids}
+        state_dict = {name: param for name, param in state_dict.items() if name in trainable_param_names}
+        if remove_prefix is not None:
+            state_dict = {name.removeprefix(remove_prefix): param for name, param in state_dict.items()}
+        return state_dict
+
+    def distribution_matching_loss(self, student_sample, teacher, fake):
+        with torch.no_grad():
+            dims = tuple(range(1, student_sample.ndim))
+            scale = (student_sample.float() - teacher.float()).abs().mean(dim=dims, keepdim=True).clamp_min(1e-8)
+            direction = torch.nan_to_num((fake.float() - teacher.float()) / scale)
+        return 0.5 * F.mse_loss(student_sample.float(), (student_sample.float() - direction).detach())
+
+    def forward_generator(self, inputs, rollout_step, update_generator):
+        if update_generator:
+            student_sample = self.rollout(inputs, rollout_step, True)
+            with torch.no_grad():
+                noised = self.sample_noise(student_sample.detach(), True, self.min_timestep_dmd, self.max_timestep_dmd)
+                teacher = self.teacher_clean(noised, inputs)
+                fake = self.fake_clean(noised, inputs)
+            loss = self.distribution_matching_loss(student_sample, teacher, fake)
+            if self.gan:
+                gan_sample = student_sample.detach().requires_grad_(True)
+                logits = self.logits(gan_sample, inputs)
+                gan_loss = F.softplus(-logits).mean()
+                gradient = torch.autograd.grad(gan_loss, gan_sample)[0].detach()
+                loss = loss + self.generator_gan_loss_weight * (gan_loss.detach() + ((student_sample - student_sample.detach()) * gradient).sum())
+            self.select_model("generator")
+        else:
+            with torch.no_grad():
+                student_sample = self.rollout(inputs, rollout_step)
+            loss = None
+        return loss, student_sample.detach()
+
+    def forward_critic(self, inputs, student_sample):
+        inputs_shared, inputs_posi, inputs_nega = inputs
+        noisy, level, noise = self.sample_noise(student_sample.detach(), False, self.min_timestep_dmd, self.max_timestep_dmd)
+        velocity = self.predict(noisy, level * 1000, inputs, "fake")
+        loss = F.mse_loss(velocity.float(), noise.float() - student_sample.detach().float())
+        if self.gan:
+            real_logits = self.logits(inputs_shared['input_latents'].float(), inputs)
+            fake_logits = self.logits(student_sample.detach(), inputs)
+            loss = loss + self.discriminator_gan_loss_weight * (F.softplus(-real_logits).mean() + F.softplus(fake_logits).mean())
+        return loss
+
+    def forward(self, data, phase, rollout_step=None, update_generator=True, inputs=None, student_sample=None):
+        if phase == 'generator':
+            with torch.no_grad():
+                inputs = self.prepare_batch(data)
+            loss, student_sample = self.forward_generator(inputs, rollout_step, update_generator)
+            return loss, inputs, student_sample
+        elif phase == 'critic':
+            return self.forward_critic(inputs, student_sample)
+        else:
+            raise ValueError(f"Unknown DMD2 phase: {phase}")
+
+
+def launch_dmd2_training_task(
+    accelerator, dataset, model, model_logger,
+    learning_rate=1e-4, weight_decay=0.01, num_workers=0, save_steps=None, num_epochs=1,
+    enable_model_cpu_offload=False, customized_optimizer=None, warm_up_steps=5,
+    fake_updates_per_generator=5, fake_warmup_steps=100, fake_learning_rate=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
+        customized_optimizer = args.customized_optimizer
+        warm_up_steps = getattr(args, "warm_up_steps", warm_up_steps)
+        fake_updates_per_generator = args.dmd2_fake_updates_per_generator
+        fake_warmup_steps = args.dmd2_fake_warmup_steps
+        fake_learning_rate = args.dmd2_fake_learning_rate
+    num_steps = len(model.pipe.scheduler.sigmas)
+
+    if accelerator.distributed_type == DistributedType.DEEPSPEED:
+        raise ValueError("DMD2 training does not support DeepSpeed.")
+    if enable_model_cpu_offload:
+        raise ValueError("DMD2 training does not support model CPU offload.")
+    if fake_updates_per_generator < 1:
+        raise ValueError('Update ratio must be positive.')
+    if fake_warmup_steps < 0:
+        raise ValueError('Fake warmup steps must be nonnegative.')
+    if accelerator.is_main_process:
+        save_training_args(args)
+
+    generator_params, critic_params = model.generator_parameters(), model.critic_parameters()
+    dataloader = torch.utils.data.DataLoader(dataset, shuffle=True, collate_fn=lambda batch: batch[0], num_workers=num_workers)
+    optimizer_class = get_optimizer_class(customized_optimizer)
+    fake_learning_rate = learning_rate if fake_learning_rate is None else fake_learning_rate
+    generator_optimizer = optimizer_class(generator_params, lr=learning_rate, weight_decay=weight_decay)
+    critic_optimizer = optimizer_class(critic_params, lr=fake_learning_rate, weight_decay=weight_decay)
+    generator_scheduler = torch.optim.lr_scheduler.LambdaLR(generator_optimizer, lambda s: min(1.0, (s + 1) / max(1, warm_up_steps)))
+    critic_scheduler = torch.optim.lr_scheduler.LambdaLR(critic_optimizer, lambda s: min(1.0, (s + 1) / max(1, warm_up_steps)))
+
+    model.to(device=accelerator.device)
+    exclude_quantized_params_from_ddp_sync(accelerator, model)
+    model, generator_optimizer, critic_optimizer, dataloader = accelerator.prepare(
+        model, generator_optimizer, critic_optimizer, dataloader,
+    )
+
+    for epoch in range(num_epochs):
+        for data in (progress_bar := tqdm(dataloader, disable=not accelerator.is_main_process)):
+            if accelerator.sync_gradients:
+                rollout_step = broadcast(torch.randint(num_steps, (), device=accelerator.device)).item()
+                step = model_logger.num_steps - fake_warmup_steps
+                update_generator = step >= 0 and step % fake_updates_per_generator == 0
+            with accelerator.accumulate(model):
+                with torch.set_grad_enabled(update_generator):
+                    loss, inputs, student_sample = model(data, phase='generator', rollout_step=rollout_step, update_generator=update_generator)
+                if update_generator:
+                    accelerator.backward(loss)
+                loss = model({}, phase='critic', inputs=inputs, student_sample=student_sample)
+                accelerator.backward(loss)
+
+                if accelerator.sync_gradients:
+                    if update_generator:
+                        accelerator.clip_grad_norm_(generator_params, 10.)
+                    accelerator.clip_grad_norm_(critic_params, 10.)
+
+                if update_generator:
+                    generator_optimizer.step()
+                    if accelerator.sync_gradients and not getattr(generator_optimizer, "step_was_skipped", False):
+                        generator_scheduler.step()
+                critic_optimizer.step()
+                if accelerator.sync_gradients and not getattr(critic_optimizer, "step_was_skipped", False):
+                    critic_scheduler.step()
+                generator_optimizer.zero_grad(set_to_none=True)
+                critic_optimizer.zero_grad(set_to_none=True)
+                if accelerator.sync_gradients:
+                    model_logger.on_step_end(accelerator, model, save_steps)
+                    progress_bar.set_postfix(generator_updates=(max(0, model_logger.num_steps - fake_warmup_steps) + fake_updates_per_generator - 1) // fake_updates_per_generator)
+        if save_steps is None:
+            model_logger.on_epoch_end(accelerator, model, epoch)
+
+    model_logger.on_training_end(accelerator, model, save_steps)

{diffsynth-2.2.0 → diffsynth-2.2.1}/diffsynth/diffusion/flow_match.py

@@ -5,7 +5,7 @@ from typing_extensions import Literal
 
 class FlowMatchScheduler():
 
-    def __init__(self, template: Literal["FLUX.1", "Wan", "Qwen-Image", "FLUX.2", "FLUX.3", "Z-Image", "LTX-2", "Qwen-Image-Lightning", "ERNIE-Image", "MOSS-Image", "ACE-Step", "Ideogram4", "Krea-2", "Boogu", "MiniMax-H3", "MiniMax-Music3", "LingBot-Video", "SenseNova-U1", "YuE2", "Fibo"] = "FLUX.1"):
+    def __init__(self, template: Literal["FLUX.1", "Wan", "Qwen-Image", "FLUX.2", "FLUX.3", "Z-Image", "LTX-2", "Qwen-Image-Lightning", "ERNIE-Image", "MOSS-Image", "ACE-Step", "Ideogram4", "Krea-2", "Boogu", "MiniMax-H3", "MiniMax-H3-Image", "MiniMax-Music3", "LingBot-Video", "SenseNova-U1", "YuE2", "Fibo"] = "FLUX.1"):
         self.set_timesteps_fn = {
             "FLUX.1": FlowMatchScheduler.set_timesteps_flux,
             "Wan": FlowMatchScheduler.set_timesteps_wan,
@@ -23,6 +23,7 @@ class FlowMatchScheduler():
             "Krea-2": FlowMatchScheduler.set_timesteps_krea2,
             "Boogu": FlowMatchScheduler.set_timesteps_boogu,
             "MiniMax-H3": FlowMatchScheduler.set_timesteps_minimax_h3,
+            "MiniMax-H3-Image": FlowMatchScheduler.set_timesteps_minimax_h3_image,
             "MiniMax-Music3": FlowMatchScheduler.set_timesteps_minimax_music3,
             "LingBot-Video": FlowMatchScheduler.set_timesteps_lingbot_video,
             "SenseNova-U1": FlowMatchScheduler.set_timesteps_sensenova_u1,
@@ -393,6 +394,20 @@ class FlowMatchScheduler():
         timesteps = sigmas * num_train_timesteps
         return sigmas, timesteps
 
+    @staticmethod
+    def set_timesteps_minimax_h3_image(num_inference_steps=50, denoising_strength=1.0, shift=2.22, distilled=False):
+        if distilled:
+            num_train_timesteps = 1000
+            sample_sigmas = torch.tensor([
+                1.0, 0.9675236940383911, 0.8875485062599182, 0.75, 0.5327739119529724, 0.23200887441635132, 0.0,
+            ], dtype=torch.float32)
+            timesteps = 1.0 - sample_sigmas[:-1]
+            sigmas = 1.0 - timesteps
+            timesteps = sigmas * num_train_timesteps
+            return sigmas, timesteps
+        else:
+            return FlowMatchScheduler.set_timesteps_minimax_h3(num_inference_steps=num_inference_steps, denoising_strength=denoising_strength, shift=shift)
+
     @staticmethod
     def set_timesteps_minimax_music3(num_inference_steps=30, denoising_strength=1.0):
         num_train_timesteps = 1000