feat(npu): add MindIE-SD compile fusion and MINDIE attention backend

diffsynth_engine/args.py

@@ -18,7 +18,7 @@ def _parse_tuple(value: str) -> Tuple[int, int] | int:
 
 
 def _parse_attention_params(
-    attn_type: str,
+    attn_type: str | None,
     sparge_topk: float | None = None,
 ) -> AttentionParams | None:
     """Parse attention parameters based on attention type"""
@@ -100,8 +100,8 @@ def parse_cli_args() -> Dict[str, Any]:
     attn_group.add_argument(
         "--attn-type",
         type=str,
-        default="sdpa",
-        help="Attention type (default: sdpa)",
+        default=None,
+        help="Attention type (default: auto, SDPA on GPU, mindie on NPU)",
     )
     attn_group.add_argument(
         "--sparge-topk",
@@ -171,7 +171,7 @@ def parse_cli_args() -> Dict[str, Any]:
     args_dict["vae_tile_stride"] = _parse_tuple(args.vae_tile_stride)
 
     # Attention configuration
-    attn_type = args.attn_type.lower()
+    attn_type = args.attn_type.lower() if args.attn_type is not None else None
     args_dict["attn_type"] = attn_type
     args_dict["attn_params"] = _parse_attention_params(attn_type, args.sparge_topk)
 

diffsynth_engine/configs/base.py

@@ -6,6 +6,7 @@
 from diffsynth_engine.layers.attention import AttentionType
 from diffsynth_engine.registry import get_attn_backend
 from diffsynth_engine.utils import logging
+from diffsynth_engine.utils.platform import get_device_type, is_mindie_sd_available
 
 logger = logging.get_logger(__name__)
 
@@ -27,8 +28,7 @@ class PipelineConfig:
     model_dtype: torch.dtype = torch.bfloat16
     text_encoder_dtype: torch.dtype = torch.bfloat16
     vae_dtype: torch.dtype = torch.float32
-    device: str | torch.device = "cuda"
-
+    device: str | torch.device = "auto"
     pipeline_class_name: str | None = None
 
     # vae
@@ -37,7 +37,7 @@ class PipelineConfig:
     vae_tile_stride: int | Tuple[int, int] = (192, 192)
 
     # attention
-    attn_type: AttentionType | str = AttentionType.SDPA
+    attn_type: AttentionType | str = "auto"
     attn_params: Optional[AttentionParams] = None
 
     # parallelism
@@ -56,6 +56,11 @@ def from_dict(cls, args_dict: Dict[str, Any]) -> "PipelineConfig":
         return cls(**filtered_dict)
 
     def __post_init__(self):
+        if self.device == "auto":
+            self.device = get_device_type() 
+        if self.attn_type == "auto":
+            self.attn_type = "mindie" if is_mindie_sd_available() else "sdpa"
+
         self.attn_type = str(self.attn_type)
         init_parallel_config(self)
         validate_attn_config(self)

diffsynth_engine/layers/attention/backends/abstract.py

@@ -26,6 +26,7 @@ class AttentionType(str, enum.Enum):
     SAGE2 = "sage2"
     SAGE3 = "sage3"
     SPARGE = "sparge"
+    MINDIE = "mindie"
 
     def __str__(self) -> str:
         return self.value

diffsynth_engine/layers/attention/backends/mindie_attn.py

@@ -0,0 +1,71 @@
+import torch
+
+from diffsynth_engine.layers.attention.backends.abstract import (
+    AttentionBackend,
+    AttentionImpl,
+    AttentionMetadata,
+    AttentionType,
+)
+from diffsynth_engine.utils import logging
+from diffsynth_engine.utils.platform import is_npu_available
+
+logger = logging.get_logger(__name__)
+
+try:
+    from mindiesd.layers.flash_attn.attention_forward import attention_forward
+
+    MINDIESD_ATTN_AVAILABLE = is_npu_available() 
+except ImportError:
+    MINDIESD_ATTN_AVAILABLE = False
+
+class MindieAttentionBackend(AttentionBackend):
+    @staticmethod
+    def check_availability() -> None:
+        if not MINDIESD_ATTN_AVAILABLE:
+            error_msg = "MindiesdAttention backend is not available. Please visit https://gitcode.com/Ascend/MindIE-SD/blob/dev/docs/zh/features/core_layers.md and follow the installation instructions."
+            logger.error(error_msg)
+            raise RuntimeError(error_msg)
+
+    @staticmethod
+    def get_type() -> str:
+        return str(AttentionType.MINDIE)
+
+    @staticmethod
+    def get_impl_cls() -> type["AttentionImpl"]:
+        return MindieAttentionImpl
+
+    @staticmethod
+    def get_supported_head_sizes() -> list[int]:
+        return []
+
+
+class MindieAttentionImpl(AttentionImpl):
+    def __init__(
+        self,
+        num_heads: int,
+        head_size: int,
+        softmax_scale: float | None = None,
+        causal: bool = False,
+        num_kv_heads: int | None = None,
+        **extra_impl_args,
+    ) -> None:
+        self.scale = softmax_scale or (head_size**-0.5)
+
+    def forward(
+        self,
+        query: torch.Tensor,
+        key: torch.Tensor,
+        value: torch.Tensor,
+        attn_mask: torch.Tensor | None = None,
+        attn_metadata: AttentionMetadata | None = None,
+        **kwargs,
+    ) -> torch.Tensor:
+        return attention_forward(
+            query=query,
+            key=key,
+            value=value,
+            attn_mask=attn_mask,
+            scale=self.scale,
+            fused=True,
+            head_first=False,
+        )

diffsynth_engine/pipelines/base.py

@@ -96,7 +96,12 @@ def init_transformer(
             model.to(device=pipeline_config.device)
 
         del state_dict
-        return model
+
+        # Pipeline loads via from_config + load_state_dict, so compile must be
+        # applied here (DiffusionModel.from_pretrained is no longer on this path).
+        from diffsynth_engine.platform.npu import apply_mindie_sd_compile
+
+        return apply_mindie_sd_compile(model)
 
     @staticmethod
     def init_text_encoder(

diffsynth_engine/platform/npu/__init__.py

@@ -0,0 +1,8 @@
+from diffsynth_engine.utils.platform import is_mindie_sd_available
+
+from .compilation import apply_mindie_sd_compile
+
+__all__ = [
+    "apply_mindie_sd_compile",
+    "is_mindie_sd_available",
+]

diffsynth_engine/platform/npu/compilation.py

@@ -0,0 +1,47 @@
+import os
+import torch
+
+from diffsynth_engine.utils.platform import is_mindie_sd_available
+
+
+def apply_mindie_sd_compile(model: torch.nn.Module) -> torch.nn.Module:
+    """Apply MindIE-SD torch.compile backend fusion patterns.
+
+    Enables RMSNorm, RoPE, AdaLayerNorm, and MulAdd op fusion for NPU execution.
+    FastGELU fusion is disabled (npu_fast_gelu was slower than native GELU in practice).
+    On non-NPU devices or when mindiesd is unavailable, returns the model unchanged.
+
+    Compilation strategy:
+        If the model defines ``_repeated_blocks`` (e.g. ``["QwenImageTransformerBlock"]``),
+        compiles only the repeated submodules individually to reduce compilation time
+        and graph break risk. Otherwise falls back to full-model ``torch.compile``.
+        Set ``DIFFSYNTH_DISABLE_COMPILE=1`` to disable.
+    """
+    if not is_mindie_sd_available() or os.environ.get("DIFFSYNTH_DISABLE_COMPILE", "0") == "1":
+        return model
+
+    from mindiesd.compilation import CompilationConfig, MindieSDBackend
+
+    CompilationConfig.fusion_patterns.enable_fast_gelu = False
+
+    repeated_blocks = getattr(model, "_repeated_blocks", None)
+    if repeated_blocks:
+        paths = [
+            name
+            for name, submodule in model.named_modules()
+            if submodule.__class__.__name__ in repeated_blocks
+        ]
+        # Sort by depth in descending order (submodules before parent modules)
+        paths.sort(key=lambda p: len(p.split(".")), reverse=True)
+
+        backend = MindieSDBackend()
+        for name in paths:
+            parent_name, _, child_name = name.rpartition(".")
+            parent = model.get_submodule(parent_name) if parent_name else model
+            submodule = getattr(parent, child_name)
+            setattr(parent, child_name, torch.compile(submodule, backend=backend))
+
+        return model
+
+    # Fall back to full-model compilation for models without `_repeated_blocks` defined.
+    return torch.compile(model, backend=MindieSDBackend())

diffsynth_engine/registry.py

@@ -7,6 +7,7 @@
 from diffsynth_engine.utils import logging
 from diffsynth_engine.utils.constants import MODEL_INDEX_NAME
 from diffsynth_engine.utils.import_utils import LazyImport
+from diffsynth_engine.utils.platform import is_mindie_sd_available
 
 if TYPE_CHECKING:
     from diffsynth_engine.layers.attention.backends.abstract import AttentionBackend
@@ -37,6 +38,7 @@
     "sage3": "diffsynth_engine.layers.attention.backends.sage_attn_3:SageAttention3Backend",
     "sdpa": "diffsynth_engine.layers.attention.backends.sdpa:SDPABackend",
     "sparge": "diffsynth_engine.layers.attention.backends.sparge_attn:SpargeAttentionBackend",
+    "mindie": "diffsynth_engine.layers.attention.backends.mindie_attn:MindieAttentionBackend",
 }
 
 PIPELINE_REGISTRY: dict[str, LazyImport] = {}
@@ -118,7 +120,7 @@ def get_attn_backend(attn_type: str | None = None) -> type["AttentionBackend"]:
         _attention_backend_registry_initialized = True
 
     if attn_type is None:
-        attn_type = "sdpa"
+        attn_type = "mindie" if is_mindie_sd_available() else "sdpa"
     if attn_type not in ATTENTION_BACKEND_REGISTRY:
         available_backends = sorted(ATTENTION_BACKEND_REGISTRY)
         raise ValueError(f"Attention backend {attn_type!r} not found. Available backends: {available_backends}")

diffsynth_engine/utils/platform.py

@@ -1,4 +1,5 @@
 import torch
+from functools import cache
 
 
 def _is_cuda() -> bool:
@@ -12,8 +13,28 @@ def _is_rocm() -> bool:
 def _is_mps() -> bool:
     return torch.backends.mps.is_available()
 
+@cache
+def is_npu_available() -> bool:
+    try:
+        import torch_npu
+
+        return torch_npu.npu.is_available()
+    except ImportError:
+        return False
+
+@cache
+def is_mindie_sd_available() -> bool:
+    try:
+        import mindiesd  # noqa: F401
+
+        return is_npu_available()
+    except ImportError:
+        return False
+
 
 def get_device(local_rank: int) -> torch.device:
+    if is_npu_available():
+        return torch.device("npu", local_rank)
     if _is_cuda() or _is_rocm():
         return torch.device("cuda", local_rank)
     if _is_mps():
@@ -23,6 +44,8 @@ def get_device(local_rank: int) -> torch.device:
 
 
 def get_device_type() -> str:
+    if is_npu_available():
+        return "npu"
     if _is_cuda() or _is_rocm():
         return "cuda"
     if _is_mps():
@@ -32,6 +55,8 @@ def get_device_type() -> str:
 
 
 def get_torch_distributed_backend() -> str:
+    if is_npu_available():
+        return "hccl"
     if _is_cuda() or _is_rocm():
         return "nccl"
     if _is_mps():