diffsynth 2.1.11__tar.gz → 2.2.0__tar.gz

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: diffsynth
-Version: 2.1.11
+Version: 2.2.0
 Summary: Enjoy the magic of Diffusion models!
 Author: ModelScope Team
 License-Expression: Apache-2.0
@@ -114,6 +114,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.
 
+- **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.
 
 - **October 16, 2026** flux-3-image-klein open-sourced, welcome a new member to the image model family! Support includes text-to-image generation, multi-reference editing, layout control, low VRAM inference, and Base training capabilities. For details, please refer to the [documentation](/docs/en/Model_Details/FLUX.3.md) and [example code](/examples/flux3/). We quantized the model with [EntroPack](https://github.com/modelscope/entropack), providing pre-quantized models in 3-bit, 4-bit, and 6-bit versions at [DiffSynth-Studio/flux-3-image-klein-entropack](https://www.modelscope.cn/models/DiffSynth-Studio/flux-3-image-klein-entropack).
@@ -402,6 +404,8 @@ 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 |
 |-|-|-|-|-|-|-|-|
+| 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) |
 | Qwen-Image-2.1 | [Qwen/Qwen-Image-2.1-Turbo](https://www.modelscope.cn/models/Qwen/Qwen-Image-2.1-Turbo) | [code](/examples/qwen_image_21/model_inference/Qwen-Image-2.1-Turbo.py) | [code](/examples/qwen_image_21/model_inference_low_vram/Qwen-Image-2.1-Turbo.py) | [code](/examples/qwen_image_21/model_training/full/Qwen-Image-2.1-Turbo.sh) | [code](/examples/qwen_image_21/model_training/validate_full/Qwen-Image-2.1-Turbo.py) | [code](/examples/qwen_image_21/model_training/lora/Qwen-Image-2.1-Turbo.sh) | [code](/examples/qwen_image_21/model_training/validate_lora/Qwen-Image-2.1-Turbo.py) |
 | MiniMax-H3 | [MiniMax/MiniMax-H3: FL2VA](https://www.modelscope.cn/models/MiniMax/MiniMax-H3) | [code](/examples/minimax_h3/model_inference/MiniMax-H3-FL2VA.py) | [code](/examples/minimax_h3/model_inference_low_vram/MiniMax-H3-FL2VA.py) | [code](/examples/minimax_h3/model_training/full/MiniMax-H3-FL2VA.sh) | [code](/examples/minimax_h3/model_training/validate_full/MiniMax-H3-FL2VA.py) | [code](/examples/minimax_h3/model_training/lora/MiniMax-H3-FL2VA.sh) | [code](/examples/minimax_h3/model_training/validate_lora/MiniMax-H3-FL2VA.py) |
@@ -414,6 +418,7 @@ Quick start: experience popular and the latest models:
 Model overview:
 
 - Image generation
+    - 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/)
     - Qwen-Image-2.1: [Documentation](https://diffsynth-studio-doc.readthedocs.io/en/latest/Model_Details/Qwen-Image-2.1.html), [Example code](/examples/qwen_image_21/)
@@ -702,6 +707,8 @@ 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 |
 |-|-|-|-|-|-|-|-|
+| 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) |
 | MOSS-Image | [openmoss/MOSS-Image-SFT](https://www.modelscope.cn/models/openmoss/MOSS-Image-SFT) | [code](/examples/moss_image/model_inference/MOSS-Image-SFT.py) | [code](/examples/moss_image/model_inference_low_vram/MOSS-Image-SFT.py) | [code](/examples/moss_image/model_training/full/MOSS-Image-SFT.sh) | [code](/examples/moss_image/model_training/validate_full/MOSS-Image-SFT.py) | [code](/examples/moss_image/model_training/lora/MOSS-Image-SFT.sh) | [code](/examples/moss_image/model_training/validate_lora/MOSS-Image-SFT.py) |
 | MOSS-Image | [DiffSynth-Studio/MOSS-Image-EntroPack](https://www.modelscope.cn/models/DiffSynth-Studio/MOSS-Image-EntroPack) | [code](/examples/moss_image/model_inference/MOSS-Image-EntroPack.py) | [code](/examples/moss_image/model_inference_low_vram/MOSS-Image-EntroPack.py) | - | - | [code](/examples/moss_image/model_training/lora/MOSS-Image-entropack.sh) | [code](/examples/moss_image/model_training/validate_lora/MOSS-Image-entropack.py) |

{diffsynth-2.1.11 → diffsynth-2.2.0}/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.
 
+- **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.
 
 - **October 16, 2026** flux-3-image-klein open-sourced, welcome a new member to the image model family! Support includes text-to-image generation, multi-reference editing, layout control, low VRAM inference, and Base training capabilities. For details, please refer to the [documentation](/docs/en/Model_Details/FLUX.3.md) and [example code](/examples/flux3/). We quantized the model with [EntroPack](https://github.com/modelscope/entropack), providing pre-quantized models in 3-bit, 4-bit, and 6-bit versions at [DiffSynth-Studio/flux-3-image-klein-entropack](https://www.modelscope.cn/models/DiffSynth-Studio/flux-3-image-klein-entropack).
@@ -329,6 +331,8 @@ 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 |
 |-|-|-|-|-|-|-|-|
+| 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) |
 | Qwen-Image-2.1 | [Qwen/Qwen-Image-2.1-Turbo](https://www.modelscope.cn/models/Qwen/Qwen-Image-2.1-Turbo) | [code](/examples/qwen_image_21/model_inference/Qwen-Image-2.1-Turbo.py) | [code](/examples/qwen_image_21/model_inference_low_vram/Qwen-Image-2.1-Turbo.py) | [code](/examples/qwen_image_21/model_training/full/Qwen-Image-2.1-Turbo.sh) | [code](/examples/qwen_image_21/model_training/validate_full/Qwen-Image-2.1-Turbo.py) | [code](/examples/qwen_image_21/model_training/lora/Qwen-Image-2.1-Turbo.sh) | [code](/examples/qwen_image_21/model_training/validate_lora/Qwen-Image-2.1-Turbo.py) |
 | MiniMax-H3 | [MiniMax/MiniMax-H3: FL2VA](https://www.modelscope.cn/models/MiniMax/MiniMax-H3) | [code](/examples/minimax_h3/model_inference/MiniMax-H3-FL2VA.py) | [code](/examples/minimax_h3/model_inference_low_vram/MiniMax-H3-FL2VA.py) | [code](/examples/minimax_h3/model_training/full/MiniMax-H3-FL2VA.sh) | [code](/examples/minimax_h3/model_training/validate_full/MiniMax-H3-FL2VA.py) | [code](/examples/minimax_h3/model_training/lora/MiniMax-H3-FL2VA.sh) | [code](/examples/minimax_h3/model_training/validate_lora/MiniMax-H3-FL2VA.py) |
@@ -341,6 +345,7 @@ Quick start: experience popular and the latest models:
 Model overview:
 
 - Image generation
+    - 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/)
     - Qwen-Image-2.1: [Documentation](https://diffsynth-studio-doc.readthedocs.io/en/latest/Model_Details/Qwen-Image-2.1.html), [Example code](/examples/qwen_image_21/)
@@ -629,6 +634,8 @@ 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 |
 |-|-|-|-|-|-|-|-|
+| 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) |
 | MOSS-Image | [openmoss/MOSS-Image-SFT](https://www.modelscope.cn/models/openmoss/MOSS-Image-SFT) | [code](/examples/moss_image/model_inference/MOSS-Image-SFT.py) | [code](/examples/moss_image/model_inference_low_vram/MOSS-Image-SFT.py) | [code](/examples/moss_image/model_training/full/MOSS-Image-SFT.sh) | [code](/examples/moss_image/model_training/validate_full/MOSS-Image-SFT.py) | [code](/examples/moss_image/model_training/lora/MOSS-Image-SFT.sh) | [code](/examples/moss_image/model_training/validate_lora/MOSS-Image-SFT.py) |
 | MOSS-Image | [DiffSynth-Studio/MOSS-Image-EntroPack](https://www.modelscope.cn/models/DiffSynth-Studio/MOSS-Image-EntroPack) | [code](/examples/moss_image/model_inference/MOSS-Image-EntroPack.py) | [code](/examples/moss_image/model_inference_low_vram/MOSS-Image-EntroPack.py) | - | - | [code](/examples/moss_image/model_training/lora/MOSS-Image-entropack.sh) | [code](/examples/moss_image/model_training/validate_lora/MOSS-Image-entropack.py) |

{diffsynth-2.1.11 → diffsynth-2.2.0}/diffsynth/configs/model_configs.py

@@ -1843,9 +1843,19 @@ flux3_series = [
     },
 ]
 
+fibo_series = [
+    {
+        # Example: ModelConfig(model_id="briaai/Fibo-2.0-dev", origin_file_pattern="transformer/*.safetensors")
+        # Example: ModelConfig(model_id="briaai/Fibo-2.0-edit-turbo", origin_file_pattern="transformer/*.safetensors")
+        "model_hash": "39f09f95705fd6d9b20670f1923d3c12",
+        "model_name": "fibo_dit",
+        "model_class": "diffsynth.models.fibo_dit.FiboDiT",
+    },
+]
+
 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
+    + yue2_series + qwen_image_21_series + flux3_series + moss_image_series + fibo_series
 )

{diffsynth-2.1.11 → diffsynth-2.2.0}/diffsynth/configs/vram_management_module_maps.py

@@ -256,6 +256,14 @@ VRAM_MANAGEMENT_MODULE_MAPS = {
         "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.flux2_vae.Flux2VAELatentNorm": "diffsynth.core.vram.layers.AutoWrappedModule",
+    },
+    "diffsynth.models.fibo_dit.FiboDiT": {
+        "torch.nn.Linear": "diffsynth.core.vram.layers.AutoWrappedLinear",
+        "diffsynth.models.general_modules.RMSNorm": "diffsynth.core.vram.layers.AutoWrappedModule",
+        "torch.nn.LayerNorm": "diffsynth.core.vram.layers.AutoWrappedModule",
+        "diffsynth.models.fibo_dit.FiboPerceiverResampler": "diffsynth.core.vram.layers.AutoWrappedModule",
+        "diffsynth.models.fibo_dit.FiboTransformerBlock": "diffsynth.core.vram.layers.AutoWrappedModule",
     },
     "diffsynth.models.z_image_text_encoder.ZImageTextEncoder": {
         "torch.nn.Linear": "diffsynth.core.vram.layers.AutoWrappedLinear",

diffsynth-2.2.0/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.0/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", "sage_attention", "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.0/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.0/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.0/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.0/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.0/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.0/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.0/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)

{diffsynth-2.1.11 → diffsynth-2.2.0}/diffsynth/core/data/operators.py

@@ -1,6 +1,5 @@
 import math, warnings, time
-import torch, torchvision, imageio, os
-import imageio.v3 as iio
+import torch, os
 import urllib.request
 from io import BytesIO
 from PIL import Image
@@ -100,16 +99,20 @@ class ImageCropAndResize(DataProcessingOperator):
         self.max_pixels = max_pixels
         self.height_division_factor = height_division_factor
         self.width_division_factor = width_division_factor
+        import torchvision
+        self.resize = torchvision.transforms.functional.resize
+        self.center_crop = torchvision.transforms.functional.center_crop
+        self.interpolation_mode = torchvision.transforms.InterpolationMode.BILINEAR
 
     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 = self.resize(
             image,
             (round(height*scale), round(width*scale)),
-            interpolation=torchvision.transforms.InterpolationMode.BILINEAR
+            interpolation=self.interpolation_mode
         )
-        image = torchvision.transforms.functional.center_crop(image, (target_height, target_width))
+        image = self.center_crop(image, (target_height, target_width))
         return image
     
     def get_height_width(self, image):
@@ -143,6 +146,7 @@ class FrameSamplerByRateMixin:
         self.fix_frame_rate = fix_frame_rate
 
     def get_reader(self, data: str):
+        import imageio
         return imageio.get_reader(data)
 
     def get_available_num_frames(self, reader):
@@ -213,6 +217,7 @@ class LoadGIF(DataProcessingOperator):
 
     def get_num_frames(self, path):
         num_frames = self.num_frames
+        import imageio.v3 as iio
         images = iio.imread(path, mode="RGB")
         if len(images) < num_frames:
             num_frames = len(images)
@@ -223,6 +228,7 @@ class LoadGIF(DataProcessingOperator):
     def __call__(self, data: str):
         num_frames = self.get_num_frames(data)
         frames = []
+        import imageio.v3 as iio
         images = iio.imread(data, mode="RGB")
         for img in images:
             frame = Image.fromarray(img)

{diffsynth-2.1.11 → diffsynth-2.2.0}/diffsynth/core/data/unified_dataset.py

@@ -1,5 +1,5 @@
 from .operators import *
-import torch, json, pandas
+import torch, json, os
 
 
 class UnifiedDataset(torch.utils.data.Dataset):
@@ -84,6 +84,7 @@ class UnifiedDataset(torch.utils.data.Dataset):
                     metadata.append(json.loads(line.strip()))
             self.data = metadata
         else:
+            import pandas
             metadata = pandas.read_csv(metadata_path)
             self.data = [metadata.iloc[i].to_dict() for i in range(len(metadata))]
 

{diffsynth-2.1.11 → diffsynth-2.2.0}/diffsynth/core/gradient/gradient_checkpoint.py

@@ -1,11 +1,9 @@
 import torch
+import sys
 
 
-try:
-    import deepspeed
-    _HAS_DEEPSPEED = True
-except ModuleNotFoundError:
-    _HAS_DEEPSPEED = False
+def get_deepspeed():
+    return sys.modules.get("deepspeed")
 
 
 def create_custom_forward(module):
@@ -36,7 +34,8 @@ def gradient_checkpoint_forward(
     *args,
     **kwargs,
 ):
-    if use_gradient_checkpointing and _HAS_DEEPSPEED and deepspeed.checkpointing.is_configured():
+    deepspeed = get_deepspeed() if use_gradient_checkpointing else None
+    if use_gradient_checkpointing and deepspeed is not None and deepspeed.checkpointing.is_configured():
         kwarg_keys = tuple(kwargs)
         all_args = args + tuple(kwargs.values())
         if not judge_args_requires_grad(*all_args):