Ptn Residual Connection Model With Attention Mask

PtnResidualConnectionModelWithAttentionMask is the mask-aware sibling of the pack's plain residual connection node. Same idea - wrap a model so its output is added back to its input - except the wrapped model gets two arguments on every forward pass: the input tensor and an attention mask. That one extra argument is the whole reason this node exists, and it's what lets you build a Transformer encoder from scratch in ComfyUI without writing code.

Why it exists

The plain Ptn Residual Connection Model wraps sub-models that take a single tensor. But multi-head attention - the engine of every Transformer - needs to know which tokens are real and which are padding, so it consumes an attention mask as a second input. In the pack's "Building Transformer From Scratch" guide, this node is what wraps each multi-head attention block: the mask flows through the wrapped model, and the residual add happens afterward, exactly where a textbook encoder puts it. If you're replicating that workflow, you'll use this node on the attention side of each encoder layer and the plain residual node on the feed-forward side.

How it works

The source is barely longer than the docstring:

res = inputs
x = self.model(inputs, mask)
x = res + x
if self.closure:
    x = self.closure(x)
return x

Input goes through model(inputs, mask), the original input is added back, and an optional closure function runs last. Same gradient-highway trick as any residual connection; the only difference is that the wrapped model must be built to accept (tensor, mask) - which means it has to be composed from the pack's mask-aware building blocks, like Ptn Multihead Attention or a Ptn Chained Model With Attention Mask.

Inputs and output

• model (PTMODEL) - a sub-model whose forward signature is (inputs, mask). This is the constraint that trips people up: a normal chain node won't work here because it doesn't know how to receive a mask.
• closure (PTCALLABLE, optional) - a function applied after the residual add.

The single PTMODEL output feeds into a Pt Chained Model or a training node. The attention mask itself isn't a direct input to this node - it arrives at runtime through the wrapped model's forward call, which is a subtle but important point: the mask is passed by the training/evaluation machinery when it calls the model, not wired into this node's inputs.

The gotcha to watch

Your wrapped model must accept (x, mask). If you accidentally wrap a single-input model, you'll get a TypeError about too many arguments the first time forward runs - not at graph-build time. That's actually helpful: the pack deliberately splits the two residual nodes precisely so you can't silently feed a mask to a model that ignores it, the way the "you must explicitly specify which layers consume the attention mask" section of the guide warns. Match the node to the model and it just works.

Installing it

Part of ComfyUI-Pt-Wrapper, HowToSD's no-code PyTorch pack (a spin-off of ComfyUI-Data-Analysis). ComfyUI Manager → search "ComfyUI-Pt-Wrapper", or:

cd ComfyUI/custom_nodes
git clone https://github.com/HowToSD/ComfyUI-Pt-Wrapper

Restart after. First boot is slow while requirements.txt installs pandas, scikit-learn, transformers, sentencepiece, peft and the rest. No model downloads for the node itself - the Transformer example workflow fetches the IMDB dataset when you run it.