Merge 8fd62980e7 into aaabf34258

comfy/ldm/minimax/model.py

@@ -156,6 +156,7 @@ class Attention(nn.Module):
         self.out_proj = operations.Linear(inner, hidden, bias=False, dtype=dtype, device=device)
 
     def forward(self, x, rope_freqs=None, transformer_options={}):
+        patches = transformer_options.get("patches", {})
         s = x.shape[0]
         q, k, v = self.qkv_proj(x).split(self.heads * self.head_dim, dim=-1)
         v = v.view(s, self.heads, self.head_dim)
@@ -177,11 +178,43 @@ class Attention(nn.Module):
         else:
             q = self.q_norm(q.view(s, self.heads, self.head_dim))
             k = self.k_norm(k.view(s, self.heads, self.head_dim))
+        extra_options = None
+        if "attn1_patch" in patches or "attn1_output_patch" in patches:
+            extra_options = {
+                key: value
+                for key, value in transformer_options.items()
+                if key not in ("patches", "patches_replace")
+            }
+            extra_options["n_heads"] = self.heads
+            extra_options["dim_head"] = self.head_dim
+
         v = v.clone()
-        q = AttentionTensorContainer(q.transpose(0, 1).unsqueeze(0))
-        k = AttentionTensorContainer(k.transpose(0, 1).unsqueeze(0))
-        v = AttentionTensorContainer(v.transpose(0, 1).unsqueeze(0))
+        if "attn1_patch" in patches:
+            q = q.reshape(1, s, -1)
+            k = k.reshape(1, s, -1)
+            v = v.reshape(1, s, -1)
+            for p in patches["attn1_patch"]:
+                out = p(q, k, v, extra_options=extra_options)
+                if isinstance(out, dict):
+                    q, k, v = out.get("q", q), out.get("k", k), out.get("v", v)
+                else:
+                    q, k, v = out
+            q = q.view(q.shape[0], q.shape[1], self.heads, self.head_dim).transpose(1, 2)
+            k = k.view(k.shape[0], k.shape[1], self.heads, self.head_dim).transpose(1, 2)
+            v = v.view(v.shape[0], v.shape[1], self.heads, self.head_dim).transpose(1, 2)
+        else:
+            q = q.transpose(0, 1).unsqueeze(0)
+            k = k.transpose(0, 1).unsqueeze(0)
+            v = v.transpose(0, 1).unsqueeze(0)
+        q = AttentionTensorContainer(q)
+        k = AttentionTensorContainer(k)
+        v = AttentionTensorContainer(v)
         out = optimized_attention(q, k, v, self.heads, mask=None, skip_reshape=True, transformer_options=transformer_options)
+
+        if "attn1_output_patch" in patches:
+            for p in patches["attn1_output_patch"]:
+                out = p(out, extra_options)
+
         return self.out_proj(out.squeeze(0))
 
 
@@ -643,7 +676,10 @@ class MiniMaxH3Model(nn.Module):
         patches_replace = transformer_options.get("patches_replace", {})
         blocks_replace = patches_replace.get("dit", {})
         prefetch_queue = comfy.model_prefetch.make_prefetch_queue(list(self.blocks), device, transformer_options)
+        transformer_options["total_blocks"] = len(self.blocks)
+        transformer_options["block_type"] = "double"
         for i, block in enumerate(self.blocks):
+            transformer_options["block_index"] = i
             comfy.model_prefetch.prefetch_queue_pop(prefetch_queue, device, block)
             if ("double_block", i) in blocks_replace:
                 def block_wrap(args):

tests-unit/comfy_test/test_minimax_attention.py

@@ -0,0 +1,53 @@
+import torch
+
+import comfy.ldm.minimax.model as minimax
+
+
+class QKV(torch.nn.Module):
+    def forward(self, x):
+        return torch.cat((x, x + 1, x + 2), dim=-1)
+
+
+def test_attention_patch_accepts_tuple_and_mapping_callbacks(monkeypatch):
+    attention = minimax.Attention(4, 2, 2, 1e-6, operations=torch.nn)
+    attention.qkv_proj = QKV()
+    attention.q_norm = torch.nn.Identity()
+    attention.k_norm = torch.nn.Identity()
+    attention.out_proj = torch.nn.Identity()
+    seen = {}
+
+    def tuple_patch(q, k, v, extra_options):
+        assert extra_options["block_index"] == 3
+        return q + 1, k + 1, v + 1
+
+    def mapping_patch(q, k, v, pe=None, attn_mask=None, extra_options=None):
+        assert pe is None
+        assert attn_mask is None
+        assert extra_options["n_heads"] == 2
+        return {"q": q * 2, "v": v * 3}
+
+    def output_patch(out, extra_options):
+        assert extra_options["block_index"] == 3
+        return out + 4
+
+    def fake_attention(q, k, v, *args, **kwargs):
+        seen.update(q=q, k=k, v=v)
+        return v.transpose(1, 2).reshape(1, v.shape[2], -1)
+
+    monkeypatch.setattr(minimax, "optimized_attention", fake_attention)
+    x = torch.zeros(2, 4)
+    output = attention(
+        x,
+        transformer_options={
+            "block_index": 3,
+            "patches": {
+                "attn1_patch": [tuple_patch, mapping_patch],
+                "attn1_output_patch": [output_patch],
+            },
+        },
+    )
+
+    assert torch.equal(seen["q"], torch.full((1, 2, 2, 2), 2.0))
+    assert torch.equal(seen["k"], torch.full((1, 2, 2, 2), 2.0))
+    assert torch.equal(seen["v"], torch.full((1, 2, 2, 2), 9.0))
+    assert torch.equal(output, torch.full((2, 4), 13.0))