diffsynth-engine 0.6.1.dev24__tar.gz → 0.6.1.dev26__tar.gz

{diffsynth_engine-0.6.1.dev24 → diffsynth_engine-0.6.1.dev26}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: diffsynth_engine
-Version: 0.6.1.dev24
+Version: 0.6.1.dev26
 Author: MuseAI x ModelScope
 Classifier: Programming Language :: Python :: 3
 Classifier: Operating System :: OS Independent

{diffsynth_engine-0.6.1.dev24 → diffsynth_engine-0.6.1.dev26}/diffsynth_engine/__init__.py

@@ -12,11 +12,13 @@ from .configs import (
     WanStateDicts,
     QwenImageStateDicts,
     AttnImpl,
+    SpargeAttentionParams,
+    VideoSparseAttentionParams,
+    LoraConfig,
     ControlNetParams,
     ControlType,
     QwenImageControlNetParams,
     QwenImageControlType,
-    LoraConfig,
 )
 from .pipelines import (
     SDImagePipeline,
@@ -59,6 +61,9 @@ __all__ = [
     "WanStateDicts",
     "QwenImageStateDicts",
     "AttnImpl",
+    "SpargeAttentionParams",
+    "VideoSparseAttentionParams",
+    "LoraConfig",
     "ControlNetParams",
     "ControlType",
     "QwenImageControlNetParams",
@@ -79,7 +84,6 @@ __all__ = [
     "FluxIPAdapterRefTool",
     "FluxReplaceByControlTool",
     "FluxReduxRefTool",
-    "LoraConfig",
     "fetch_model",
     "fetch_modelscope_model",
     "register_fetch_modelscope_model",

{diffsynth_engine-0.6.1.dev24 → diffsynth_engine-0.6.1.dev26}/diffsynth_engine/configs/__init__.py

@@ -17,14 +17,16 @@ from .pipeline import (
     WanStateDicts,
     WanS2VStateDicts,
     QwenImageStateDicts,
-    LoraConfig,
     AttnImpl,
+    SpargeAttentionParams,
+    VideoSparseAttentionParams,
+    LoraConfig,
 )
 from .controlnet import (
     ControlType,
     ControlNetParams,
-    QwenImageControlNetParams,
     QwenImageControlType,
+    QwenImageControlNetParams,
 )
 
 __all__ = [
@@ -46,10 +48,12 @@ __all__ = [
     "WanStateDicts",
     "WanS2VStateDicts",
     "QwenImageStateDicts",
-    "QwenImageControlType",
-    "QwenImageControlNetParams",
+    "AttnImpl",
+    "SpargeAttentionParams",
+    "VideoSparseAttentionParams",
+    "LoraConfig",
     "ControlType",
     "ControlNetParams",
-    "LoraConfig",
-    "AttnImpl",
+    "QwenImageControlType",
+    "QwenImageControlNetParams",
 ]

{diffsynth_engine-0.6.1.dev24 → diffsynth_engine-0.6.1.dev26}/diffsynth_engine/configs/pipeline.py

@@ -5,7 +5,6 @@ from dataclasses import dataclass, field
 from typing import List, Dict, Tuple, Optional
 
 from diffsynth_engine.configs.controlnet import ControlType
-from diffsynth_engine.models.basic.video_sparse_attention import get_vsa_kwargs
 
 
 @dataclass
@@ -27,6 +26,8 @@ class AttnImpl(Enum):
     FA2 = "fa2"  # Flash Attention 2
     FA3 = "fa3"  # Flash Attention 3
     FA3_FP8 = "fa3_fp8"  # Flash Attention 3 with FP8
+    AITER = "aiter"  # Aiter Flash Attention
+    AITER_FP8 = "aiter_fp8"  # Aiter Flash Attention with FP8
     XFORMERS = "xformers"  # XFormers
     SDPA = "sdpa"  # Scaled Dot Product Attention
     SAGE = "sage"  # Sage Attention
@@ -52,23 +53,6 @@ class AttentionConfig:
     dit_attn_impl: AttnImpl = AttnImpl.AUTO
     attn_params: Optional[SpargeAttentionParams | VideoSparseAttentionParams] = None
 
-    def get_attn_kwargs(self, latents: torch.Tensor, device: str) -> Dict:
-        attn_kwargs = {"attn_impl": self.dit_attn_impl.value}
-        if isinstance(self.attn_params, SpargeAttentionParams):
-            assert self.dit_attn_impl == AttnImpl.SPARGE
-            attn_kwargs.update(
-                {
-                    "smooth_k": self.attn_params.smooth_k,
-                    "simthreshd1": self.attn_params.simthreshd1,
-                    "cdfthreshd": self.attn_params.cdfthreshd,
-                    "pvthreshd": self.attn_params.pvthreshd,
-                }
-            )
-        elif isinstance(self.attn_params, VideoSparseAttentionParams):
-            assert self.dit_attn_impl == AttnImpl.VSA
-            attn_kwargs.update(get_vsa_kwargs(latents.shape[2:], (1, 2, 2), self.attn_params.sparsity, device=device))
-        return attn_kwargs
-
 
 @dataclass
 class OptimizationConfig:

{diffsynth_engine-0.6.1.dev24 → diffsynth_engine-0.6.1.dev26}/diffsynth_engine/models/basic/attention.py

@@ -13,6 +13,7 @@ from diffsynth_engine.utils.flag import (
     SAGE_ATTN_AVAILABLE,
     SPARGE_ATTN_AVAILABLE,
     VIDEO_SPARSE_ATTN_AVAILABLE,
+    AITER_AVAILABLE,
 )
 from diffsynth_engine.utils.platform import DTYPE_FP8
 
@@ -93,6 +94,9 @@ if SPARGE_ATTN_AVAILABLE:
         )
         return out.transpose(1, 2)
 
+if AITER_AVAILABLE:
+    from aiter import flash_attn_func as aiter_flash_attn
+    from aiter import flash_attn_fp8_pertensor_func as aiter_flash_attn_fp8
 
 if VIDEO_SPARSE_ATTN_AVAILABLE:
     from diffsynth_engine.models.basic.video_sparse_attention import (
@@ -137,6 +141,8 @@ def attention(
         "fa2",
         "fa3",
         "fa3_fp8",
+        "aiter",
+        "aiter_fp8",
         "xformers",
         "sdpa",
         "sage",
@@ -157,6 +163,13 @@ def attention(
                     logger.debug(
                         "flash_attn_3 does not support attention mask, will use fallback attention implementation"
                     )
+        if AITER_AVAILABLE:
+            if flash_attn3_compatible:
+                return aiter_flash_attn(q, k, v, softmax_scale=scale)
+            else:
+                logger.warning(
+                    f"head_dim={q.shape[-1]}, but aiter_flash_attn only supports head dimension at most {FA3_MAX_HEADDIM}, will use fallback attention implementation"
+                )
         if XFORMERS_AVAILABLE:
             return xformers_attn(q, k, v, attn_mask=attn_mask, scale=scale)
         if SDPA_AVAILABLE:
@@ -183,6 +196,22 @@ def attention(
                 v = v.to(dtype=DTYPE_FP8)
                 out = flash_attn3(q, k, v, softmax_scale=scale)
                 return out.to(dtype=origin_dtype)
+        if attn_impl == "aiter" or attn_impl == "aiter_fp8":
+            if not flash_attn3_compatible:
+                raise RuntimeError(
+                    f"head_dim={q.shape[-1]}, but aiter_flash_attn only supports head dimension at most {FA3_MAX_HEADDIM}"
+                )
+            if attn_mask is not None:
+                raise RuntimeError("aiter_flash_attn does not support attention mask")
+            if attn_impl == "aiter" :
+                return aiter_flash_attn(q, k, v, softmax_scale=scale)
+            else:
+                origin_dtype = q.dtype
+                q = q.to(dtype=DTYPE_FP8)
+                k = k.to(dtype=DTYPE_FP8)
+                v = v.to(dtype=DTYPE_FP8)
+                out = aiter_flash_attn_fp8(q, k, v, softmax_scale=scale)
+                return out.to(dtype=origin_dtype) 
         if attn_impl == "fa2":
             return flash_attn2(q, k, v, softmax_scale=scale)
         if attn_impl == "xformers":
@@ -288,6 +317,8 @@ def long_context_attention(
         "fa2",
         "fa3",
         "fa3_fp8",
+        "aiter",
+        "aiter_fp8",
         "sdpa",
         "sage",
         "sparge",
@@ -303,6 +334,13 @@ def long_context_attention(
                 logger.warning(
                     f"head_dim={q.shape[-1]}, but flash_attn_3 only supports head dimension at most {FA3_MAX_HEADDIM}, will use fallback attention implementation"
                 )
+        if AITER_AVAILABLE:
+            if flash_attn3_compatible:
+                return LongContextAttention(attn_type=AttnType.AITER)(q, k, v, softmax_scale=scale)
+            else:
+                logger.warning(
+                    f"head_dim={q.shape[-1]}, but aiter_flash_attn only supports head dimension at most {FA3_MAX_HEADDIM}, will use fallback attention implementation"
+                )
         if SDPA_AVAILABLE:
             return LongContextAttention(attn_type=AttnType.TORCH)(q, k, v, softmax_scale=scale)
         if FLASH_ATTN_2_AVAILABLE:
@@ -323,6 +361,20 @@ def long_context_attention(
             v = v.to(dtype=DTYPE_FP8)
             out = LongContextAttention(attn_type=AttnType.FA3)(q, k, v, softmax_scale=scale)
             return out.to(dtype=origin_dtype)
+        if attn_impl == "aiter" or attn_impl == "aiter_fp8":
+            if not flash_attn3_compatible:
+                raise RuntimeError(
+                    f"head_dim={q.shape[-1]}, but aiter_flash_attn only supports head dimension at most {FA3_MAX_HEADDIM}"
+                )
+            if attn_impl == "aiter":
+                return LongContextAttention(attn_type=AttnType.AITER)(q, k, v, softmax_scale=scale)
+
+            origin_dtype = q.dtype
+            q = q.to(dtype=DTYPE_FP8)
+            k = k.to(dtype=DTYPE_FP8)
+            v = v.to(dtype=DTYPE_FP8)
+            out = LongContextAttention(attn_type=AttnType.AITER)(q, k, v, softmax_scale=scale)
+            return out.to(dtype=origin_dtype)
         if attn_impl == "fa2":
             return LongContextAttention(attn_type=AttnType.FA)(q, k, v, softmax_scale=scale)
         if attn_impl == "sdpa":

{diffsynth_engine-0.6.1.dev24 → diffsynth_engine-0.6.1.dev26}/diffsynth_engine/models/basic/video_sparse_attention.py

@@ -2,9 +2,12 @@ import torch
 import math
 import functools
 
-from vsa import video_sparse_attn as vsa_core
+from diffsynth_engine.utils.flag import VIDEO_SPARSE_ATTN_AVAILABLE
 from diffsynth_engine.utils.parallel import get_sp_ulysses_group, get_sp_ring_world_size
 
+if VIDEO_SPARSE_ATTN_AVAILABLE:
+    from vsa import video_sparse_attn as vsa_core
+
 VSA_TILE_SIZE = (4, 4, 4)
 
 

{diffsynth_engine-0.6.1.dev24 → diffsynth_engine-0.6.1.dev26}/diffsynth_engine/pipelines/base.py

@@ -5,7 +5,15 @@ from einops import rearrange
 from typing import Dict, List, Tuple, Union, Optional
 from PIL import Image
 
-from diffsynth_engine.configs import BaseConfig, BaseStateDicts, LoraConfig
+from diffsynth_engine.configs import (
+    BaseConfig,
+    BaseStateDicts,
+    LoraConfig,
+    AttnImpl,
+    SpargeAttentionParams,
+    VideoSparseAttentionParams,
+)
+from diffsynth_engine.models.basic.video_sparse_attention import get_vsa_kwargs
 from diffsynth_engine.utils.offload import enable_sequential_cpu_offload, offload_model_to_dict, restore_model_from_dict
 from diffsynth_engine.utils.fp8_linear import enable_fp8_autocast
 from diffsynth_engine.utils.gguf import load_gguf_checkpoint
@@ -33,6 +41,7 @@ class BasePipeline:
         dtype=torch.float16,
     ):
         super().__init__()
+        self.config = None
         self.vae_tiled = vae_tiled
         self.vae_tile_size = vae_tile_size
         self.vae_tile_stride = vae_tile_stride
@@ -48,7 +57,7 @@ class BasePipeline:
         raise NotImplementedError()
 
     @classmethod
-    def from_state_dict(cls, state_dicts: BaseStateDicts, pipeline_config: BaseConfig) -> "BasePipeline":
+    def from_state_dict(cls, state_dicts: BaseStateDicts, config: BaseConfig) -> "BasePipeline":
         raise NotImplementedError()
 
     def update_weights(self, state_dicts: BaseStateDicts) -> None:
@@ -260,6 +269,25 @@ class BasePipeline:
         )
         return init_latents, latents, sigmas, timesteps
 
+    def get_attn_kwargs(self, latents: torch.Tensor) -> Dict:
+        attn_kwargs = {"attn_impl": self.config.dit_attn_impl.value}
+        if isinstance(self.config.attn_params, SpargeAttentionParams):
+            assert self.config.dit_attn_impl == AttnImpl.SPARGE
+            attn_kwargs.update(
+                {
+                    "smooth_k": self.config.attn_params.smooth_k,
+                    "simthreshd1": self.config.attn_params.simthreshd1,
+                    "cdfthreshd": self.config.attn_params.cdfthreshd,
+                    "pvthreshd": self.config.attn_params.pvthreshd,
+                }
+            )
+        elif isinstance(self.config.attn_params, VideoSparseAttentionParams):
+            assert self.config.dit_attn_impl == AttnImpl.VSA
+            attn_kwargs.update(
+                get_vsa_kwargs(latents.shape[2:], (1, 2, 2), self.config.attn_params.sparsity, device=self.device)
+            )
+        return attn_kwargs
+
     def eval(self):
         for model_name in self.model_names:
             model = getattr(self, model_name)

{diffsynth_engine-0.6.1.dev24 → diffsynth_engine-0.6.1.dev26}/diffsynth_engine/pipelines/flux_image.py

@@ -751,7 +751,7 @@ class FluxImagePipeline(BasePipeline):
         latents = latents.to(self.dtype)
         self.load_models_to_device(["dit"])
 
-        attn_kwargs = self.config.get_attn_kwargs(latents, self.device)
+        attn_kwargs = self.get_attn_kwargs(latents)
         noise_pred = self.dit(
             hidden_states=latents,
             timestep=timestep,
@@ -886,7 +886,7 @@ class FluxImagePipeline(BasePipeline):
                 empty_cache()
                 param.model.to(self.device)
 
-            attn_kwargs = self.config.get_attn_kwargs(latents, self.device)
+            attn_kwargs = self.get_attn_kwargs(latents)
             double_block_output, single_block_output = param.model(
                 hidden_states=latents,
                 control_condition=control_condition,

{diffsynth_engine-0.6.1.dev24 → diffsynth_engine-0.6.1.dev26}/diffsynth_engine/pipelines/qwen_image.py

@@ -208,7 +208,9 @@ class QwenImagePipeline(BasePipeline):
             )
         if config.load_encoder:
             logger.info(f"loading state dict from {config.encoder_path} ...")
-            encoder_state_dict = cls.load_model_checkpoint(config.encoder_path, device="cpu", dtype=config.encoder_dtype)
+            encoder_state_dict = cls.load_model_checkpoint(
+                config.encoder_path, device="cpu", dtype=config.encoder_dtype
+            )
 
         state_dicts = QwenImageStateDicts(
             model=model_state_dict,
@@ -547,7 +549,7 @@ class QwenImagePipeline(BasePipeline):
         entity_masks: Optional[List[torch.Tensor]] = None,
     ):
         self.load_models_to_device(["dit"])
-        attn_kwargs = self.config.get_attn_kwargs(latents, self.device)
+        attn_kwargs = self.get_attn_kwargs(latents)
         noise_pred = self.dit(
             image=latents,
             edit=image_latents,

{diffsynth_engine-0.6.1.dev24 → diffsynth_engine-0.6.1.dev26}/diffsynth_engine/pipelines/wan_s2v.py

@@ -394,7 +394,7 @@ class WanSpeech2VideoPipeline(WanVideoPipeline):
         void_audio_input: torch.Tensor | None = None,
     ):
         latents = latents.to(dtype=self.config.model_dtype, device=self.device)
-        attn_kwargs = self.config.get_attn_kwargs(latents, self.device)
+        attn_kwargs = self.get_attn_kwargs(latents)
 
         noise_pred = model(
             x=latents,

{diffsynth_engine-0.6.1.dev24 → diffsynth_engine-0.6.1.dev26}/diffsynth_engine/pipelines/wan_video.py

@@ -144,7 +144,7 @@ class WanVideoPipeline(BasePipeline):
         lora_list: List[Tuple[str, float]],
         fused: bool = True,
         save_original_weight: bool = False,
-        lora_converter: Optional[WanLoRAConverter] = None
+        lora_converter: Optional[WanLoRAConverter] = None,
     ):
         assert self.config.tp_degree is None or self.config.tp_degree == 1, (
             "load LoRA is not allowed when tensor parallel is enabled; "
@@ -156,11 +156,15 @@ class WanVideoPipeline(BasePipeline):
         )
         super().load_loras(lora_list, fused, save_original_weight, lora_converter)
 
-    def load_loras_low_noise(self, lora_list: List[Tuple[str, float]], fused: bool = True, save_original_weight: bool = False):
+    def load_loras_low_noise(
+        self, lora_list: List[Tuple[str, float]], fused: bool = True, save_original_weight: bool = False
+    ):
         assert self.dit2 is not None, "low noise LoRA can only be applied to Wan2.2"
         self.load_loras(lora_list, fused, save_original_weight, self.low_noise_lora_converter)
 
-    def load_loras_high_noise(self, lora_list: List[Tuple[str, float]], fused: bool = True, save_original_weight: bool = False):
+    def load_loras_high_noise(
+        self, lora_list: List[Tuple[str, float]], fused: bool = True, save_original_weight: bool = False
+    ):
         assert self.dit2 is not None, "high noise LoRA can only be applied to Wan2.2"
         self.load_loras(lora_list, fused, save_original_weight)
 
@@ -323,7 +327,7 @@ class WanVideoPipeline(BasePipeline):
 
     def predict_noise(self, model, latents, image_clip_feature, image_y, timestep, context):
         latents = latents.to(dtype=self.config.model_dtype, device=self.device)
-        attn_kwargs = self.config.get_attn_kwargs(latents, self.device)
+        attn_kwargs = self.get_attn_kwargs(latents)
 
         noise_pred = model(
             x=latents,

{diffsynth_engine-0.6.1.dev24 → diffsynth_engine-0.6.1.dev26}/diffsynth_engine/utils/flag.py

@@ -31,6 +31,11 @@ if SDPA_AVAILABLE:
 else:
     logger.info("Torch SDPA is not available")
 
+AITER_AVAILABLE = importlib.util.find_spec("aiter") is not None
+if AITER_AVAILABLE:
+    logger.info("Aiter is available")
+else:
+    logger.info("Aiter is not available")
 
 # 有损
 SAGE_ATTN_AVAILABLE = importlib.util.find_spec("sageattention") is not None

{diffsynth_engine-0.6.1.dev24 → diffsynth_engine-0.6.1.dev26}/diffsynth_engine/utils/parallel.py

@@ -19,8 +19,6 @@ from typing import Dict, List, Set, Type, Union, Optional
 from queue import Empty
 
 import diffsynth_engine.models.basic.attention as attention_ops
-from diffsynth_engine.models import PreTrainedModel
-from diffsynth_engine.pipelines import BasePipeline
 from diffsynth_engine.utils.platform import empty_cache
 from diffsynth_engine.utils import logging
 
@@ -300,14 +298,15 @@ def _worker_loop(
             world_size=world_size,
         )
 
-        def wrap_for_parallel(module: Union[PreTrainedModel, BasePipeline]):
-            if isinstance(module, BasePipeline):
-                for model_name in module.model_names:
-                    if isinstance(submodule := getattr(module, model_name), PreTrainedModel):
+        def wrap_for_parallel(module):
+            if hasattr(module, "model_names"):
+                for model_name in getattr(module, "model_names"):
+                    submodule = getattr(module, model_name)
+                    if getattr(submodule, "_supports_parallelization", False):
                         setattr(module, model_name, wrap_for_parallel(submodule))
                 return module
 
-            if not module._supports_parallelization:
+            if not getattr(module, "_supports_parallelization", False):
                 return module
 
             if tp_degree > 1:

{diffsynth_engine-0.6.1.dev24 → diffsynth_engine-0.6.1.dev26}/diffsynth_engine.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: diffsynth_engine
-Version: 0.6.1.dev24
+Version: 0.6.1.dev26
 Author: MuseAI x ModelScope
 Classifier: Programming Language :: Python :: 3
 Classifier: Operating System :: OS Independent