feat(npu): add Ascend platform abstraction and Qwen-Image NPU support

diffsynth_engine/configs/pipeline.py

@@ -35,6 +35,7 @@ class AttnImpl(Enum):
     SAGE = "sage"  # Sage Attention
     SPARGE = "sparge"  # Sparge Attention
     VSA = "vsa"  # Video Sparse Attention
+    MINDIE = "mindie"  # Mindie Attention
 
 
 @dataclass

diffsynth_engine/models/basic/attention.py

@@ -15,6 +15,7 @@
     SPARGE_ATTN_AVAILABLE,
     VIDEO_SPARSE_ATTN_AVAILABLE,
     AITER_AVAILABLE,
+    MINDIE_AVAILABLE,
 )
 from diffsynth_engine.utils.platform import DTYPE_FP8
 
@@ -107,6 +108,16 @@ def sparge_attn(
         distributed_video_sparse_attn,
     )
 
+if MINDIE_AVAILABLE:
+    from mindiesd.layers.flash_attn.attention_forward import attention_forward
+
+    def mindie_attn(q, k, v, attn_mask=None, scale=None):
+        return attention_forward(
+            query=q, key=k, value=v,
+            attn_mask=attn_mask, scale=scale,
+            fused=True, head_first=False,
+        )
+
 
 def eager_attn(q, k, v, attn_mask=None, scale=None):
     q = q.transpose(1, 2)
@@ -152,6 +163,7 @@ def attention(
         "sage",
         "sparge",
         "vsa",
+        "mindie",
     ]
     flash_attn3_compatible = q.shape[-1] <= FA3_MAX_HEADDIM
     if attn_impl is None or attn_impl == "auto":
@@ -192,6 +204,8 @@ def attention(
                 )
         if XFORMERS_AVAILABLE:
             return xformers_attn(q, k, v, attn_mask=attn_mask, scale=scale)
+        if MINDIE_AVAILABLE:
+            return mindie_attn(q, k, v, attn_mask=attn_mask, scale=scale)
         if SDPA_AVAILABLE:
             return sdpa_attn(q, k, v, attn_mask=attn_mask, scale=scale)
         if FLASH_ATTN_2_AVAILABLE:
@@ -263,6 +277,8 @@ def attention(
                 cdfthreshd=kwargs.get("cdfthreshd", 0.98),
                 pvthreshd=kwargs.get("pvthreshd", 50),
             )
+        if attn_impl == "mindie":
+            return mindie_attn(q, k, v, attn_mask=attn_mask, scale=scale)
         if attn_impl == "vsa":
             return video_sparse_attn(
                 q,
@@ -354,7 +370,10 @@ def long_context_attention(
         "sage",
         "sparge",
         "vsa",
+        "mindie",
     ]
+    if attn_impl == "mindie":
+        raise RuntimeError("mindie long-context attention is not supported yet")
     assert attn_mask is None, "long context attention does not support attention mask"
     flash_attn3_compatible = q.shape[-1] <= FA3_MAX_HEADDIM
     if attn_impl is None or attn_impl == "auto":

diffsynth_engine/pipelines/qwen_image.py

@@ -376,7 +376,9 @@ def update_weights(self, state_dicts: QwenImageStateDicts) -> None:
         self.update_component(self.vae, state_dicts.vae, self.config.device, self.config.vae_dtype)
 
     def compile(self):
-        self.dit.compile_repeated_blocks()
+        from diffsynth_engine.platforms import resolve_platform
+        platform_cls = resolve_platform(self.config.device)
+        self.dit.compile_repeated_blocks(**platform_cls.compile_kwargs())
 
     def load_loras(self, lora_list: List[Tuple[str, float]], fused: bool = True, save_original_weight: bool = False):
         assert self.config.tp_degree is None or self.config.tp_degree == 1, (

diffsynth_engine/platforms/__init__.py

@@ -0,0 +1,148 @@
+from __future__ import annotations
+
+import platform as host_platform
+from functools import lru_cache
+from typing import Type
+
+import torch
+
+from .ascend import (
+    AscendPlatform,
+    probe_ascend_capabilities,
+    probe_ascend_feature,
+    reset_ascend_capability_cache,
+)
+from .base import PlatformBackend, PlatformCapabilities
+
+
+class CPUPlatform(PlatformBackend):
+    name = "cpu"
+    device_type = "cpu"
+
+
+class CUDAPlatform(PlatformBackend):
+    name = "cuda"
+    device_type = "cuda"
+
+    @classmethod
+    def is_available(cls) -> bool:
+        return torch.cuda.is_available()
+
+    @classmethod
+    def set_device(cls, index: int | str | torch.device) -> None:
+        torch.cuda.set_device(index)
+
+    @classmethod
+    def synchronize(cls) -> None:
+        torch.cuda.synchronize()
+
+    @classmethod
+    def empty_cache(cls) -> None:
+        torch.cuda.empty_cache()
+
+    @classmethod
+    def distributed_backend(cls) -> str:
+        return "nccl"
+
+
+class ROCmPlatform(CUDAPlatform):
+    name = "rocm"
+
+
+class MPSPlatform(PlatformBackend):
+    name = "mps"
+    device_type = "mps"
+
+    @classmethod
+    def is_available(cls) -> bool:
+        return torch.backends.mps.is_available()
+
+    @classmethod
+    def synchronize(cls) -> None:
+        torch.mps.synchronize()
+
+    @classmethod
+    def empty_cache(cls) -> None:
+        torch.mps.empty_cache()
+
+
+_PLATFORM_REGISTRY: dict[str, Type[PlatformBackend]] = {
+    "cpu": CPUPlatform,
+    "cuda": ROCmPlatform if torch.version.hip else CUDAPlatform,
+    "mps": MPSPlatform,
+    "npu": AscendPlatform,
+}
+
+
+def register_platform(device_type: str, platform_cls: Type[PlatformBackend], *, overwrite: bool = False) -> None:
+    if device_type in _PLATFORM_REGISTRY and not overwrite:
+        raise ValueError(f"Platform for device type {device_type!r} is already registered")
+    _PLATFORM_REGISTRY[device_type] = platform_cls
+
+
+@lru_cache(maxsize=None)
+def auto_detect_device() -> str:
+    """auto detect device type in order of cuda(gpu/rocm), npu, mps, cpu"""
+    for device_type in ("cuda", "npu", "mps", "cpu"):
+        try:
+            if resolve_platform(device_type).is_available():
+                return device_type
+        except Exception:
+            continue
+    return "cpu"
+
+
+def get_device_type(device: str | torch.device | None = None) -> str:
+    if device is None or (isinstance(device, str) and device.lower() in ("auto", "")):
+        return auto_detect_device()
+    if isinstance(device, torch.device):
+        return device.type
+    return str(device).split(":", 1)[0].lower()
+
+
+def resolve_platform(device: str | torch.device) -> Type[PlatformBackend]:
+    device_type = get_device_type(device)
+    try:
+        return _PLATFORM_REGISTRY[device_type]
+    except KeyError as exc:
+        available = ", ".join(sorted(_PLATFORM_REGISTRY))
+        raise ValueError(f"Unsupported device type {device_type!r}. Registered device types: {available}") from exc
+
+
+def get_preferred_fp8_dtype(device: str | torch.device = "cuda") -> torch.dtype:
+    platform_cls = resolve_platform(device)
+    if platform_cls is ROCmPlatform and platform_cls.is_available():
+        properties = torch.cuda.get_device_properties(0)
+        if "gfx94" in properties.gcnArchName:
+            return torch.float8_e4m3fnuz
+    return torch.float8_e4m3fn
+
+
+def pin_memory(
+    tensor: torch.Tensor,
+    device: str | torch.device | None = None,
+) -> torch.Tensor:
+    if host_platform.system() != "Linux":
+        return tensor
+    platform_cls = resolve_platform(get_device_type(device))
+    return platform_cls.pin_memory(tensor)
+
+
+__all__ = [
+    "AscendPlatform",
+    "CPUPlatform",
+    "CUDAPlatform",
+    "MPSPlatform",
+    "PlatformBackend",
+    "PlatformCapabilities",
+    "ROCmPlatform",
+    "auto_detect_device",
+    "get_device_type",
+    "get_preferred_fp8_dtype",
+    "pin_memory",
+    "probe_ascend_capabilities",
+    "probe_ascend_feature",
+    "register_platform",
+    "reset_ascend_capability_cache",
+    "resolve_platform",
+]