Fix: Add qwen_fa3_processor.py and update __init__.py

qwenimage/__init__.py

@@ -1,3 +1,2 @@
 from .pipeline_qwenimage_edit_plus import QwenImageEditPlusPipeline
 from .transformer_qwenimage import QwenImageTransformer2DModel
-from .qwen_fa3_processor import QwenDoubleStreamAttnProcessorFA3

qwenimage/qwen_fa3_processor.py

@@ -0,0 +1,90 @@
+"""
+FA3 attention processor for Qwen-Image-Edit with SDPA fallback for Blackwell GPUs.
+"""
+import torch
+import torch.nn.functional as F
+from typing import Optional, Tuple
+from diffusers.models.transformers.transformer_qwenimage import apply_rotary_emb_qwen
+
+def _is_blackwell() -> bool:
+    if not torch.cuda.is_available():
+        return False
+    cap = torch.cuda.get_device_capability()
+    return cap[0] >= 10
+
+_fa3_available = False
+_fa3_unavailable_reason = ""
+_flash_attn_func = None
+
+if _is_blackwell():
+    _fa3_unavailable_reason = "FlashAttention-3 unsupported on Blackwell. Falling back to SDPA."
+else:
+    try:
+        from kernels import get_kernel
+        _k = get_kernel("kernels-community/vllm-flash-attn3")
+        _flash_attn_func = _k.flash_attn_func
+        _fa3_available = True
+    except Exception as e:
+        _fa3_unavailable_reason = f"kernels unavailable: {e}. Falling back to SDPA."
+
+if _fa3_available:
+    @torch.library.custom_op("flash::flash_attn_func", mutates_args=())
+    def flash_attn_func(q, k, v, causal=False):
+        output, _ = _flash_attn_func(q, k, v, causal=causal)
+        return output
+    @flash_attn_func.register_fake
+    def _fa_fake(q, k, v, causal=False):
+        return torch.empty_like(q).contiguous()
+else:
+    def flash_attn_func(q, k, v, causal=False):
+        raise RuntimeError(_fa3_unavailable_reason)
+
+def _sdpa_attention(q, k, v, causal=False):
+    return F.scaled_dot_product_attention(q.transpose(1, 2), k.transpose(1, 2), v.transpose(1, 2), is_causal=causal).transpose(1, 2)
+
+class QwenDoubleStreamAttnProcessorFA3:
+    def __init__(self):
+        if _fa3_available:
+            self._attention_backend = "fa3"
+        else:
+            import warnings
+            warnings.warn(f"QwenDoubleStreamAttnProcessorFA3: {_fa3_unavailable_reason}")
+            self._attention_backend = "sdpa"
+
+    def _attend(self, q, k, v, causal=False):
+        if self._attention_backend == "fa3":
+            return flash_attn_func(q, k, v, causal=causal)
+        return _sdpa_attention(q, k, v, causal=causal)
+
+    @torch.no_grad()
+    def __call__(self, attn, hidden_states, encoder_hidden_states=None, encoder_hidden_states_mask=None, attention_mask=None, image_rotary_emb=None):
+        if encoder_hidden_states is None:
+            raise ValueError("QwenDoubleStreamAttnProcessorFA3 requires encoder_hidden_states.")
+        B, S_img, _ = hidden_states.shape
+        S_txt = encoder_hidden_states.shape[1]
+        img_q = attn.to_q(hidden_states).unflatten(-1, (attn.heads, -1))
+        img_k = attn.to_k(hidden_states).unflatten(-1, (attn.heads, -1))
+        img_v = attn.to_v(hidden_states).unflatten(-1, (attn.heads, -1))
+        txt_q = attn.add_q_proj(encoder_hidden_states).unflatten(-1, (attn.heads, -1))
+        txt_k = attn.add_k_proj(encoder_hidden_states).unflatten(-1, (attn.heads, -1))
+        txt_v = attn.add_v_proj(encoder_hidden_states).unflatten(-1, (attn.heads, -1))
+        if getattr(attn, "norm_q", None) is not None: img_q = attn.norm_q(img_q)
+        if getattr(attn, "norm_k", None) is not None: img_k = attn.norm_k(img_k)
+        if getattr(attn, "norm_added_q", None) is not None: txt_q = attn.norm_added_q(txt_q)
+        if getattr(attn, "norm_added_k", None) is not None: txt_k = attn.norm_added_k(txt_k)
+        if image_rotary_emb is not None:
+            img_freqs, txt_freqs = image_rotary_emb
+            img_q = apply_rotary_emb_qwen(img_q, img_freqs, use_real=False)
+            img_k = apply_rotary_emb_qwen(img_k, img_freqs, use_real=False)
+            txt_q = apply_rotary_emb_qwen(txt_q, txt_freqs, use_real=False)
+            txt_k = apply_rotary_emb_qwen(txt_k, txt_freqs, use_real=False)
+        q = torch.cat([txt_q, img_q], dim=1)
+        k = torch.cat([txt_k, img_k], dim=1)
+        v = torch.cat([txt_v, img_v], dim=1)
+        out = self._attend(q, k, v, causal=False).flatten(2, 3).to(q.dtype)
+        txt_attn_out = out[:, :S_txt, :]
+        img_attn_out = out[:, S_txt:, :]
+        img_attn_out = attn.to_out[0](img_attn_out)
+        if len(attn.to_out) > 1: img_attn_out = attn.to_out[1](img_attn_out)
+        txt_attn_out = attn.to_add_out(txt_attn_out)
+        return img_attn_out, txt_attn_out