From 10d07d44b1dab7786a36cbec61a03a42ff0cd345 Mon Sep 17 00:00:00 2001
From: Michele Rossi <mikymod@gmail.com>
Date: Mon, 6 Jul 2026 20:19:01 +0200
Subject: [PATCH] Patch Qwen transformer attention with matmul instead of SDPA

Monkey-patches dispatch_attention_fn to use simple Q@K^T@V matmul,
bypassing SDPA dispatch entirely. Enables editing at full 1024x1024
(was limited to 512x512 due to SDPA crashes) and is ~3x faster.
---
 vnassets/attention.py | 42 ++++++++++++++++++++++++++++++------------
 vnassets/edit.py      | 14 +++-----------
 2 files changed, 33 insertions(+), 23 deletions(-)

diff --git a/vnassets/attention.py b/vnassets/attention.py
index 4fa0928..fc46e8b 100644
--- a/vnassets/attention.py
+++ b/vnassets/attention.py
@@ -1,7 +1,7 @@
-"""Custom attention processor that avoids SDPA on ROCm."""
+"""Custom attention processors that avoid SDPA on ROCm."""
 import torch
 import torch.nn.functional as F
-from diffusers.models.attention_processor import Attention
+from diffusers.models.attention_processor import Attention, AttnProcessor
 
 
 def simple_attention_forward(
@@ -15,37 +15,45 @@ def simple_attention_forward(
     batch_size, seq_len, _ = hidden_states.shape
     inner_dim = attn.inner_dim if hasattr(attn, 'inner_dim') else hidden_states.shape[-1]
 
-    # Project to Q, K, V
     query = attn.to_q(hidden_states)
     key = attn.to_k(hidden_states if encoder_hidden_states is None else encoder_hidden_states)
     value = attn.to_v(hidden_states if encoder_hidden_states is None else encoder_hidden_states)
 
-    # Reshape to (batch, heads, seq, head_dim)
     head_dim = inner_dim // attn.heads
     query = query.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
     key = key.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
     value = value.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
 
-    # Attention: softmax(Q @ K^T / sqrt(d)) @ V
     scale = head_dim ** -0.5
     attn_weights = query @ key.transpose(-2, -1) * scale
     attn_weights = F.softmax(attn_weights, dim=-1).to(query.dtype)
     hidden_states = attn_weights @ value
 
-    # Reshape back
     hidden_states = hidden_states.transpose(1, 2).reshape(batch_size, -1, attn.heads * head_dim)
 
-    # Output projection
     hidden_states = attn.to_out[0](hidden_states)
     if len(attn.to_out) > 1:
-        hidden_states = attn.to_out[1](hidden_states)  # dropout
+        hidden_states = attn.to_out[1](hidden_states)
     return hidden_states
 
 
+def _matmul_attention(query, key, value, attn_mask, dropout_p, is_causal, backend, parallel_config):
+    """Simple matmul attention replacing dispatch_attention_fn. Returns 4D tensor."""
+    batch_size, seq_len, num_heads, head_dim = query.shape
+    scale = head_dim ** -0.5
+    q = query.transpose(1, 2)
+    k = key.transpose(1, 2)
+    v = value.transpose(1, 2)
+    w = q @ k.transpose(-2, -1) * scale
+    if attn_mask is not None:
+        w = w + attn_mask
+    w = F.softmax(w, dim=-1).to(q.dtype)
+    out = w @ v
+    return out.transpose(1, 2)
+
+
 def patch_unet_attention(unet):
     """Replace all attention processors with simple matmul-based attention."""
-    from diffusers.models.attention_processor import AttnProcessor
-
     class SimpleAttnProcessor(AttnProcessor):
         def __call__(
             self,
@@ -57,5 +65,15 @@ def patch_unet_attention(unet):
         ) -> torch.Tensor:
             return simple_attention_forward(attn, hidden_states, encoder_hidden_states, attention_mask, **kwargs)
 
-    processor = SimpleAttnProcessor()
-    unet.set_attn_processor(processor)
+    unet.set_attn_processor(SimpleAttnProcessor())
+
+
+def patch_qwen_transformer(transformer):
+    """Patch Qwen transformer to use matmul attention instead of SDPA."""
+    from diffusers.models.attention_dispatch import dispatch_attention_fn
+    # Monkey-patch the dispatch function
+    import diffusers.models.attention_dispatch as ad
+    ad.dispatch_attention_fn = _matmul_attention
+    # Also patch the module that imports it
+    import diffusers.models.transformers.transformer_qwenimage as tq
+    tq.dispatch_attention_fn = _matmul_attention
diff --git a/vnassets/edit.py b/vnassets/edit.py
index b11f290..bd67ce9 100644
--- a/vnassets/edit.py
+++ b/vnassets/edit.py
@@ -14,6 +14,8 @@ from diffusers.models.autoencoders import AutoencoderKLQwenImage
 from diffusers.models.transformers import QwenImageTransformer2DModel
 from transformers import Qwen2_5_VLForConditionalGeneration, Qwen2Tokenizer, Qwen2VLProcessor
 
+from .attention import patch_qwen_transformer
+
 
 TEXT_ENCODER_ID = "Qwen/Qwen2.5-VL-7B-Instruct"
 VAE_ID = "Qwen/Qwen-Image"
@@ -89,21 +91,11 @@ def edit(
         transformer=transformer,
     )
     pipe.to(device)
+    patch_qwen_transformer(transformer)
 
     t_load = time.perf_counter() - t0
 
-    # Resize large inputs to avoid GPU crashes. The Qwen transformer does
-    # O(N^2) attention on image patches. 512px is safe on this GPU.
     input_image = Image.open(input_path).convert("RGB")
-    w, h = input_image.size
-    max_dim = 512
-    if w > max_dim or h > max_dim:
-        scale = max_dim / max(w, h)
-        new_w, new_h = int(w * scale), int(h * scale)
-        new_w = (new_w // 16) * 16
-        new_h = (new_h // 16) * 16
-        input_image = input_image.resize((new_w, new_h), Image.LANCZOS)
-        print(f"Resized input {w}x{h} -> {new_w}x{new_h}")
     generator = torch.Generator(device=device).manual_seed(seed)
 
     t1 = time.perf_counter()