diff --git a/diffsynth_engine/args.py b/diffsynth_engine/args.py index d73a030..dda2291 100644 --- a/diffsynth_engine/args.py +++ b/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) diff --git a/diffsynth_engine/configs/base.py b/diffsynth_engine/configs/base.py index b44c472..a3e032c 100644 --- a/diffsynth_engine/configs/base.py +++ b/diffsynth_engine/configs/base.py @@ -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 | None = None # None = auto-detect attn_params: Optional[AttentionParams] = None # parallelism @@ -56,7 +56,8 @@ def from_dict(cls, args_dict: Dict[str, Any]) -> "PipelineConfig": return cls(**filtered_dict) def __post_init__(self): - self.attn_type = str(self.attn_type) + if self.attn_type is not None: + self.attn_type = str(self.attn_type) init_parallel_config(self) validate_attn_config(self) diff --git a/diffsynth_engine/layers/attention/backends/abstract.py b/diffsynth_engine/layers/attention/backends/abstract.py index e9b5fdd..35be745 100644 --- a/diffsynth_engine/layers/attention/backends/abstract.py +++ b/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 diff --git a/diffsynth_engine/layers/attention/backends/mindie_attn.py b/diffsynth_engine/layers/attention/backends/mindie_attn.py new file mode 100644 index 0000000..e0c847c --- /dev/null +++ b/diffsynth_engine/layers/attention/backends/mindie_attn.py @@ -0,0 +1,62 @@ +import torch + +from diffsynth_engine.layers.attention.backends.abstract import ( + AttentionBackend, + AttentionImpl, + AttentionMetadata, + AttentionType, +) +from diffsynth_engine.utils.platform import is_npu_available + + +class MindieAttentionBackend(AttentionBackend): + @staticmethod + def check_availability() -> None: + if not is_npu_available(): + raise RuntimeError("NPU is not available, cannot use MINDIE attention backend") + + @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: + from mindiesd.layers.flash_attn.attention_forward import attention_forward + + return attention_forward( + query=query, + key=key, + value=value, + attn_mask=attn_mask, + scale=self.scale, + fused=True, + head_first=False, + ) diff --git a/diffsynth_engine/pipelines/base.py b/diffsynth_engine/pipelines/base.py index 601e895..5848ddd 100644 --- a/diffsynth_engine/pipelines/base.py +++ b/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( diff --git a/diffsynth_engine/platform/__init__.py b/diffsynth_engine/platform/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/diffsynth_engine/platform/npu/__init__.py b/diffsynth_engine/platform/npu/__init__.py new file mode 100644 index 0000000..c33ceef --- /dev/null +++ b/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", +] diff --git a/diffsynth_engine/platform/npu/compilation.py b/diffsynth_engine/platform/npu/compilation.py new file mode 100644 index 0000000..4cb4a66 --- /dev/null +++ b/diffsynth_engine/platform/npu/compilation.py @@ -0,0 +1,19 @@ +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. + """ + if not is_mindie_sd_available(): + return model + + from mindiesd.compilation import CompilationConfig, MindieSDBackend + + CompilationConfig.fusion_patterns.enable_fast_gelu = False + return torch.compile(model, backend=MindieSDBackend()) diff --git a/diffsynth_engine/registry.py b/diffsynth_engine/registry.py index b6b9f4e..8b735a7 100644 --- a/diffsynth_engine/registry.py +++ b/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_npu_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_npu_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}") diff --git a/diffsynth_engine/utils/platform.py b/diffsynth_engine/utils/platform.py index 575f573..4e11c39 100644 --- a/diffsynth_engine/utils/platform.py +++ b/diffsynth_engine/utils/platform.py @@ -13,7 +13,27 @@ def _is_mps() -> bool: return torch.backends.mps.is_available() +def is_npu_available() -> bool: + try: + import torch_npu + + return torch_npu.npu.is_available() + except ImportError: + return False + + +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 +43,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 +54,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():