ComfyUI Node

AddTransformerLayer

Bolt an actual attention layer onto your ComfyUI model

By TashaSkyUp·Created about a year ago·Updated about a year ago· 1
AddTransformerLayer
  • model
  • TORCH_MODEL
input_features
num_heads
feedforward_dim
dropout0.10

The whole modern AI stack - image gen, LLMs, video models - runs on attention, and AddTransformerLayer is how you get a slice of that into the EternalKernel PyTorch Nodes pack: it appends a torch.nn.TransformerEncoderLayer to your nn.Sequential. That's the full multi-head self-attention block plus feedforward, the thing transformers are made of.

Reach for it when your toy sequence problem (a bit of time-series, a small text task, anything where position matters) needs more than a dense layer can give. One layer won't make you GPT, but it's the real thing - the same building block under every LLM you've used. If you just want probabilities out of a classifier, don't bother; attention is for relationships between positions, and it's heavy for what a single linear layer does cheaper.

How it works

It builds nn.TransformerEncoderLayer(d_model=input_features, nhead=num_heads, dim_feedforward=feedforward_dim, dropout=dropout) and inserts it at the end of your Sequential, mutating in place and returning it. Straightforward - the subtleties are in the input shapes.

The big one: PyTorch's TransformerEncoderLayer expects input shaped (sequence, batch, features) (batch_first=False by default). If you're used to ComfyUI's (batch, height, width, channels) or even (batch, sequence, features), you have to transpose into (seq, batch, features) before this layer, and back after. There's no permute node in this pack, so budget for that.

Inputs

  • model (TORCH_MODEL) - the Sequential. Start from SequentialModelProvider.
  • input_features - the model dimension (d_model), the feature size of each token. Must match the last dimension of your input tensor.
  • num_heads - the number of attention heads. input_features should be divisible by num_heads (8 features with 2 heads → 4 per head). The node won't validate this; PyTorch will error at runtime if it doesn't divide evenly.
  • feedforward_dim - the hidden size of the internal MLP, typically 2–4× input_features.
  • dropout - default 0.1, applied inside the block. Fine as-is.

Output: one TORCH_MODEL.

Install

ComfyUI Manager, search "EternalKernel PyTorch Nodes", or:

cd ComfyUI/custom_nodes
git clone https://github.com/TashaSkyUp/EternalKernelPytorchNodes
cd EternalKernelPytorchNodes
pip install -r requirements.txt

Restart ComfyUI; node under ETK/pytorch. No model downloads - this is a from-scratch layer. Requirements are the standard ComfyUI stack plus scipy, scikit-learn, transformers, einops.

Common issues

  • Shape ordering. Feeding (batch, seq, features) straight in is the #1 failure. Transpose to (seq, batch, features) first. And note the return will also be in that layout - flip it back before the next layer.
  • input_features not divisible by num_heads. Runtime error, not caught at the node. Keep the math clean.
  • Self-attention needs enough data. A transformer layer on a handful of samples overfits or never learns anything meaningful. This pack's MNIST-style playground data will work; a tiny bespoke dataset may just memorize.
  • "Not an nn.Sequential." Same rule as every builder node here.

It's a quiet pack - no tutorials, nobody to copy. But the layer is stock torch, so standard transformer guidance applies, and the source is one readable file. Pack quirk worth remembering: it patches ComfyUI's validator to ignore return_type_mismatch, so a bad wire might not produce the error message you'd expect.

CategoryETK/pytorch

Inputs (5)

NameTypeDefaultDescription
modelTORCH_MODEL
input_featuresINT1–16777216
num_headsINT
feedforward_dimINT
dropoutFLOAT0.100–1

Outputs (1)

NameTypeDescription
TORCH_MODELTORCH_MODEL