NPU Adaption for Sanna

Author: J石页

src/diffusers/models/attention_processor.py

@@ -294,6 +294,10 @@ def __init__(
             processor = (
                 AttnProcessor2_0() if hasattr(F, "scaled_dot_product_attention") and self.scale_qk else AttnProcessor()
             )
+
+        if is_torch_npu_available():
+            if isinstance(processor, AttnProcessor2_0):
+                processor = AttnProcessorNPU()
         self.set_processor(processor)
 
     def set_use_xla_flash_attention(

src/diffusers/models/transformers/sana_transformer.py

@@ -120,12 +120,6 @@ def __init__(
         # 2. Cross Attention
         if cross_attention_dim is not None:
             self.norm2 = nn.LayerNorm(dim, elementwise_affine=norm_elementwise_affine, eps=norm_eps)
-
-            if is_torch_npu_available():
-                attn_processor = AttnProcessorNPU()
-            else:
-                attn_processor = AttnProcessor2_0()
-
             self.attn2 = Attention(
                 query_dim=dim,
                 cross_attention_dim=cross_attention_dim,
@@ -134,7 +128,7 @@ def __init__(
                 dropout=dropout,
                 bias=True,
                 out_bias=attention_out_bias,
-                processor=attn_processor,
+                processor=AttnProcessor2_0(),
             )
 
         # 3. Feed-forward