diff --git a/comfy/ldm/minimax/model.py b/comfy/ldm/minimax/model.py index b6feb8860..e83033bc0 100644 --- a/comfy/ldm/minimax/model.py +++ b/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): diff --git a/tests-unit/comfy_test/test_minimax_attention.py b/tests-unit/comfy_test/test_minimax_attention.py new file mode 100644 index 000000000..97e9a6116 --- /dev/null +++ b/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))