Update qwenimage/qwen_fa3_processor.py

qwenimage/qwen_fa3_processor.py

@@ -2,6 +2,23 @@ import torch
 from typing import Optional, Tuple
 from diffusers.models.transformers.transformer_qwenimage import apply_rotary_emb_qwen
 
+try:
+    from kernels import get_kernel
+    _k = get_kernel("kernels-community/vllm-flash-attn3")
+    _flash_attn_func = _k.flash_attn_func
+except Exception as e:
+    _flash_attn_func = None
+    _kernels_err = e
+
+
+def _ensure_fa3_available():
+    if _flash_attn_func is None:
+        raise ImportError(
+            "FlashAttention-3 via Hugging Face `kernels` is required. "
+            "Tried `get_kernel('kernels-community/vllm-flash-attn3')` and failed with:\n"
+            f"{_kernels_err}"
+            
+
 @torch.library.custom_op("flash::flash_attn_func", mutates_args=())
 def flash_attn_func(
     q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, causal: bool = False
@@ -35,12 +52,7 @@ class QwenDoubleStreamAttnProcessorFA3:
     _attention_backend = "fa3"  # for parity with your other processors, not used internally
 
     def __init__(self):
-        try:
-            from flash_attn.flash_attn_interface import flash_attn_interface_func
-        except ImportError:
-            raise ImportError(
-                "flash_attention v3 package is required to be installed"
-            )
+        _ensure_fa3_available()
 
     @torch.no_grad()
     def __call__(