Ptn Chained Model With Attention Mask

Ptn Chained Model With Attention Mask is the mask-aware version of the pack's chaining node. It does the same sequential model_a → model_b composition as Ptn Chained Model, but its forward pass also carries an attention mask, and it lets you say which of the two models actually needs that mask. If you're building a Transformer-style network by hand - the pack's "build a Transformer encoder from scratch" workflow - this is the node that wires masked multi-head attention into the chain.

Why you'd reach for it

Attention layers take an extra argument: a mask that tells them which positions to attend to and which to ignore (padding tokens, or positions you want hidden in a causal setup). A plain chained model can't express that, because its forward is just model_a(x); model_b(x) - no room for a second argument. This node adds the mask plumbing: if model_a_mask_req is on, model_a receives (inputs, mask); if model_b_mask_req is on, model_b does too. In practice that means you can chain something like Embedding → Multihead Attention (masked) → Linear and have the mask flow through exactly where it's needed without every layer having to accept it.

How it works

Under the hood it's a small nn.Module whose forward is:

x = model_a(inputs, mask) if model_a_mask_req else model_a(inputs)
x = model_b(x, mask) if model_b_mask_req else model_b(x)
if closure: x = closure(x)

So a model that doesn't take a mask just gets the tensor; a model that does gets tensor plus mask. The mask itself is produced upstream (typically by a tokenizer or dataset node) and fed to this node's forward along with the inputs when the chain runs. Output is a PTMODEL, ready for another chain or a training node.

The inputs

• model_a, model_b - the two models to chain.
• model_a_mask_req - set true if model_a's forward takes (inputs, mask), e.g. a Ptn Multihead Attention node.
• model_b_mask_req - same, for model_b.
• closure (optional) - a PTCALLABLE applied at the end.

Output: PTMODEL.

Installing the pack

Part of the "Training" category. Install via ComfyUI Manager (search "ComfyUI-Pt-Wrapper") or:

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

Restart ComfyUI. Install drags in transformers, datasets, peft, scikit-learn, gensim and more; no model downloads at install.

Common issues

• Flag mismatch - the classic failure. Set model_a_mask_req true on a model that doesn't accept a mask, and you get a takes 1 positional argument but 2 were given error; set it false on a masked model and it'll complain about the missing mask. The flags have to describe the model, not your intent.
• Mask shape - the mask must match what the attention layer expects (usually (batch, 1, seq, seq) for PyTorch's scaled-dot-product style attention). Shape errors surface at first forward pass.
• Attention without masking - if your sequence has padding but you chained attention without the mask, the model will attend to pad tokens and quietly degrade. That's not an error, just a subtle accuracy hit - the whole reason this node exists.
• Thin support net - the pack is a single-author educational project, basically invisible in r/comfyui discussion. The "building_transformer_from_scratch" doc on the repo is the authoritative walkthrough for this node.