Added custom multihead attention code.

__init__.py

@@ -182,6 +182,7 @@ from .modules.pytorch_wrapper.ptn_max_pool2d import PtnMaxPool2d
 from .modules.pytorch_wrapper.ptn_model_with_closure import PtnModelWithClosure
 from .modules.pytorch_wrapper.ptn_mse_loss import PtnMSELoss
 from .modules.pytorch_wrapper.ptn_multihead_attention import PtnMultiheadAttention
+from .modules.pytorch_wrapper.ptn_multihead_attention_custom import PtnMultiheadAttentionCustom
 from .modules.pytorch_wrapper.ptn_nll_loss import PtnNLLLoss
 from .modules.pytorch_wrapper.ptn_pre_add_channel_axis import PtnPreAddChannelAxis
 from .modules.pytorch_wrapper.ptn_pre_flatten import PtnPreFlatten
@@ -402,6 +403,7 @@ NODE_CLASS_MAPPINGS: Dict[str, Type[T]] = {
     "PtnMaxPool2d": PtnMaxPool2d,
     "PtnModelWithClosure": PtnModelWithClosure,
     "PtnMultiheadAttention": PtnMultiheadAttention,
+    "PtnMultiheadAttentionCustom": PtnMultiheadAttentionCustom,
     "PtnNLLLoss": PtnNLLLoss,
     "PtnPreAddChannelAxis": PtnPreAddChannelAxis,
     "PtnPreFlatten": PtnPreFlatten,
@@ -622,6 +624,7 @@ NODE_DISPLAY_NAME_MAPPINGS: Dict[str, str] = {
     "PtnMaxPool2d": "Ptn Max Pool 2d",
     "PtnModelWithClosure": "Ptn Model With Closure",
     "PtnMultiheadAttention": "Ptn Multihead Attention",
+    "PtnMultiheadAttentionCustom": "Ptn Multihead Attention Custom",
     "PtnNLLLoss": "Ptn NLL Loss",
     "PtnPreAddChannelAxis": "Ptn Pre Add Channel Axis",
     "PtnPreFlatten": "Ptn Pre Flatten",

docs/building_transformer_from_scratch.md

@@ -167,3 +167,11 @@ With this workflow, you should reach around **85% validation accuracy** on the H
 This is a solid performance for a Transformer built entirely from scratch.
 
 ---
+
+### If you want to understand how Multi-head Attention is implemented
+
+Refer to `modules/pytorch_wrapper/ptn_multihead_attention_custom.py`, which contains the implementation of Multi-head Attention and its wrapper node.  
+You can also check out the example workflow in `examples/workflows/mha_custom.json`, which demonstrates how to use this wrapper node.  
+This workflow achieves around 85–86% val accuracy on the IMDB dataset.
+
+---

docs/reference/node_reference.md

@@ -176,6 +176,7 @@
 | [Ptn Masked Mean Pooling](ptn_masked_mean_pooling.md) | Constructs a masked mean pooling layer. |
 | [Ptn Max Pool 2d](ptn_max_pool2d.md) | A maxpool layer.   |
 | [Ptn Multihead Attention](ptn_multihead_attention.md) | A Multihead attention model. |
+| [Ptn Multihead Attention Custom](ptn_multihead_attention_custom.md) | A Multihead attention model. |
 | [Ptn Pre Add Channel Axis](ptn_pre_add_channel_axis.md) | Adds a channel axis after the batch axis if the input is rank 3 (bs, h, w) |
 | [Ptn Pre Flatten](ptn_pre_flatten.md) | Flattens the input tensor before processing the tensor in the specified model. |
 | [Ptn RNN](ptn_rnn.md) | A recurrent neural network (RNN) model consisting of one or more of a recurrent layer.   |

docs/reference/ptn_multihead_attention_custom.md

@@ -0,0 +1,32 @@
+# Ptn Multihead Attention Custom
+A Multihead attention model.
+
+    Args:
+        embed_dim (int): Hidden dimension (d_model) for each token of the Transformer.
+        num_heads  (int): Number of attention heads.
+        dropout (float): Dropout rate for output.
+        bias (bool): If True, bias will be added to the linear projections of query, key, value, and the final output.
+        kdim (int): Dimension of k.
+        vdim (int): Dimension of v.
+        batch_first (bool): Input shape format: True for (B, Seq, Feature), False for (Seq, B, Feature).
+
+## Input
+| Name | Data type |
+|---|---|
+| embed_dim | Int |
+| num_heads | Int |
+| dropout | Float |
+| bias | Boolean |
+| kdim | Int |
+| vdim | Int |
+| batch_first | Boolean |
+
+## Output
+| Data type |
+|---|
+| Ptmodel |
+
+<HR>
+Category: PyTorch wrapper - Model
+
+ComfyUI Pt Wrapper Node Reference. © 2025 Hide Inada (HowToSD.com). All rights reserved.

modules/pytorch_wrapper/ptn_multihead_attention_custom.py

@@ -0,0 +1,209 @@
+from typing import Any, Dict, Tuple, Optional
+import math
+import torch
+import torch.nn as nn
+from einops import einops
+
+class PtnMultiheadAttentionCustomDef(nn.Module):
+    """
+    Implements MultiheadAttention wrapper.
+
+    Attributes:
+        embed_dim (int): Hidden dimension (d_model) for each token of the Transformer.
+        num_heads  (int): Number of attention heads.
+        dropout (float): Dropout rate for output.
+        bias (bool): If True, bias will be added to the linear projections of query, key, value, and the final output.
+        add_bias_kv (bool): If True, bias will be added to k and v sequences.
+        add_zero_attn (bool): If True, a zero batch will be added  to k and v sequences.
+        kdim (int): Dimension of k.
+        vdim (int): Dimension of v.
+        batch_first (bool): Input shape format: True for (B, Seq, Feature), False for (Seq, B, Feature).
+
+    pragma: skip_doc
+    """
+    def __init__(self,
+        embed_dim: int,
+        num_heads: int,
+        dropout: float,
+        bias: bool,
+        kdim: int,
+        vdim: int,
+        batch_first: bool
+    ) -> None:
+        super().__init__()
+
+        # Compute dimensions
+        self.embed_dim = embed_dim
+        self.num_heads = num_heads
+        self.head_dim = embed_dim // num_heads
+        assert embed_dim % num_heads == 0, "embed_dim must be divisible by num_heads"
+
+        # Projection matrix
+        self.q_proj = nn.Linear(embed_dim, embed_dim, bias=bias)
+        self.k_proj = nn.Linear(kdim, embed_dim, bias=bias)
+        self.v_proj = nn.Linear(vdim, embed_dim, bias=bias)
+        self.out_proj = nn.Linear(embed_dim, embed_dim, bias=bias)
+
+        if dropout:
+            self.dropout = nn.Dropout(dropout)
+        else:
+            self.dropout = None
+
+        self.batch_first = batch_first
+
+
+    def forward(self,
+                inputs: torch.Tensor,
+                mask:Optional[torch.Tensor]=None) -> torch.Tensor:
+        """
+        Args:
+            inputs (torch.Tensor): Input tensor of shape (batch_size, seq_len, embed_dim) if batch_first is True.
+            mask:Optional[torch.Tensor]: An optional attention mask where 1 indicates a valid token and 0 indicates padding.
+                This is the inverse of PyTorch's expected key_padding_mask, which masks out positions where the value is True.
+
+        Returns:
+            torch.Tensor: Output tensor after multi-head self-attention.
+        """
+        if inputs.dim() != 3:
+            raise ValueError("Inputs is not a rank 3 tensor.")
+
+        x = inputs
+
+        if not self.batch_first:  # (seq, batch, features) to (batch, seq, feature)
+            x = einops.rearrange(x, 's b f -> b s f')
+
+        q = self.q_proj(x)
+        k = self.k_proj(x)
+        v = self.v_proj(x)
+
+        # Split feature axis into heads × (feature per head).
+        # The size of this per-head feature is denoted by "d" (aka "head_dim").
+        # Rearrange so that dimensions are (batch, heads, seq, d).
+        # (h d) means h * d.
+        q = einops.rearrange(q, 'b s (h d) -> b h s d', h=self.num_heads)
+        k = einops.rearrange(k, 'b s (h d) -> b h s d', h=self.num_heads)
+        v = einops.rearrange(v, 'b s (h d) -> b h s d', h=self.num_heads)
+
+        # dot product q @ k.T / sqrt(head_dim)
+        kt = torch.transpose(k, -1, -2)
+        dot = torch.matmul(q, kt) / math.sqrt(self.head_dim)  # torch.sqrt requires a tensor
+
+        # q @ k = (b h s d) @ (b h d s) = (b h s s)
+        if dot.size() != (q.size(0), q.size(1), q.size(2), q.size(2)):
+            raise ValueError(f"Unexpected size for q@k.T: {dot.size()}")
+
+        # Apply mask (1 = valid token, 0 = padding)
+        # If token contains 0, the attention score for the cell
+        # will be replaced with -inf, which will be converted to 0
+        # in softmax.
+        # e.g. one row of attention score
+        # pos  0  1  2 ... 1023
+        #      5  2  -5     -20
+        # mask
+        #      1  1  0        0
+        # result
+        #      5  2   -inf   -inf
+        if mask is not None:
+            mask = mask.to(torch.bool)
+            mask = mask.unsqueeze(1).unsqueeze(2)  # (B, 1, 1, S)
+            dot = dot.masked_fill(mask == 0, float('-inf'))
+
+        # Normalize
+        dot = torch.softmax(dot, -1)
+
+        # Apply attention weights to value vectors
+        v = torch.matmul(dot, v)
+
+        # Rearrange
+        v = einops.rearrange(v, 'b h s d -> b s (h d)')
+
+        x = self.out_proj(v)
+
+        if self.dropout:
+            x = self.dropout(x)
+
+        if not self.batch_first:
+            x = einops.rearrange(x, 'b s f -> s b f')
+
+        return x
+
+
+class PtnMultiheadAttentionCustom:
+    """
+    A Multihead attention model.
+
+        Args:
+            embed_dim (int): Hidden dimension (d_model) for each token of the Transformer.
+            num_heads  (int): Number of attention heads.
+            dropout (float): Dropout rate for output.
+            bias (bool): If True, bias will be added to the linear projections of query, key, value, and the final output.
+            kdim (int): Dimension of k.
+            vdim (int): Dimension of v.
+            batch_first (bool): Input shape format: True for (B, Seq, Feature), False for (Seq, B, Feature).
+            
+    category: PyTorch wrapper - Model
+    """
+
+    @classmethod
+    def INPUT_TYPES(cls) -> Dict[str, Any]:
+        """
+        Defines the input types for the function.
+
+        Returns:
+            Dict[str, Any]: A dictionary specifying required input types.
+        """
+        return {
+            "required": {
+                "embed_dim": ("INT", {"default": 512, "min": 1, "max": 1e6}),
+                "num_heads": ("INT", {"default": 8, "min": 1, "max": 256}),
+                "dropout": ("FLOAT", {"default": 0, "min": 0, "max": 1}),
+                "bias": ("BOOLEAN", {"default": True}),
+                "kdim": ("INT", {"default": 512, "min": 1, "max": 1e6}),
+                "vdim": ("INT", {"default": 512, "min": 1, "max": 1e6}),
+                "batch_first": ("BOOLEAN", {"default": True})
+            }
+        }
+
+    RETURN_TYPES: Tuple[str] = ("PTMODEL",)
+    FUNCTION: str = "f"
+    CATEGORY: str = "Training"
+
+    @classmethod
+    def IS_CHANGED(cls, **kw):
+        return float("NaN")
+
+    def f(self,
+        embed_dim: int,
+        num_heads: int,
+        dropout: float,
+        bias: bool,
+        kdim: int,
+        vdim: int,
+        batch_first: bool
+        ) -> Tuple[nn.Module]:
+        """
+        Constructs a Multihead attention model.
+
+        Args:
+            embed_dim (int): Hidden dimension (d_model) for each token of the Transformer.
+            num_heads  (int): Number of attention heads.
+            dropout (float): Dropout rate for output.
+            bias (bool): If True, bias will be added to the linear projections of query, key, value, and the final output.
+            kdim (int): Dimension of k.
+            vdim (int): Dimension of v.
+            batch_first (bool): Input shape format: True for (B, Seq, Feature), False for (Seq, B, Feature).
+
+        Returns:
+            Tuple[nn.Module]: A tuple containing the instantiated PyTorch model.
+        """
+        with torch.inference_mode(False):
+            model = PtnMultiheadAttentionCustomDef(
+                embed_dim=embed_dim,
+                num_heads=num_heads,
+                dropout=dropout,
+                bias=bias,
+                kdim=kdim,
+                vdim=vdim,
+                batch_first=batch_first
+            )
+            return (model,)

pyproject.toml

@@ -1,7 +1,7 @@
 [project]
 name = "pt-wrapper"
 description = "PyTorch extension for ComfyUI featuring extensive PyTorch wrapper nodes for seamless tensor operations and PyTorch model training."
-version = "4.0.1"
+version = "4.1.0"
 license = {file = "LICENSE"}
 
 [project.urls]

tests/pytorch_wrapper/ptn_multihead_attention_custom_test.py

@@ -0,0 +1,62 @@
+import os
+import sys
+import unittest
+import torch
+import torch.nn as nn
+
+PROJECT_ROOT = os.path.realpath(os.path.join(os.path.dirname(__file__), "..", ".."))
+MODULE_ROOT = os.path.join(PROJECT_ROOT, "modules")
+sys.path.append(MODULE_ROOT)
+
+from pytorch_wrapper.ptn_multihead_attention_custom import PtnMultiheadAttentionCustom
+from pytorch_wrapper.ptn_multihead_attention_custom import PtnMultiheadAttentionCustomDef
+
+class TestPtnMultiheadAttentionCustom(unittest.TestCase):
+    def setUp(self) -> None:
+        """Set up test instance."""
+        self.node = PtnMultiheadAttentionCustom()
+
+    def test_fbatch_first(self) -> None:
+        """Test forward pass with batch first."""
+        model, = self.node.f(
+            embed_dim=16,
+            num_heads=4,
+            dropout=0,
+            bias=True,
+            kdim=16,
+            vdim=16,
+            batch_first=True
+        )
+
+        self.assertIsInstance(model, PtnMultiheadAttentionCustomDef)
+
+        # Create dummy input: batch_size=2, seq_len=5, embed_dim=16
+        x = torch.rand(2, 5, 16)
+        attention_custom_mask = torch.ones(2, 5)
+        attn_output = model(x, attention_custom_mask)
+        self.assertEqual(attn_output.shape, (2, 5, 16))
+
+    def test_forward_seq_first(self) -> None:
+        """Test forward pass with seq first."""
+        model, = self.node.f(
+            embed_dim=16,
+            num_heads=4,
+            dropout=0.1,
+            bias=True,
+            kdim=16,
+            vdim=16,
+            batch_first=False
+        )
+
+        self.assertIsInstance(model, PtnMultiheadAttentionCustomDef)
+
+        # Create dummy input: batch_size=2, seq_len=5, embed_dim=16
+        # Use seq first.
+        x = torch.rand(5, 2, 16)
+        attention_custom_mask = torch.ones(2, 5)
+        attn_output = model(x, attention_custom_mask)
+        self.assertEqual(attn_output.shape, (5, 2, 16))
+
+
+if __name__ == "__main__":
+    unittest.main()