ComfyUI Node

NNT Define LSTM Layer

NNT Define LSTM Layer — sequence modeling's favorite cell, defined but not yet compiled

By inventorado·Created 2 years ago·Updated 2 years ago· 69
NNT Define LSTM Layer
  • LAYER_STACK
  • LAYER_STACK
input_size10
hidden_size20
num_layers1
biasTrue
batch_firstTrue
dropout0.0
bidirectionalFalse

LSTMs are the long-short-term-memory workhorses of sequence modeling - the thing you reach for when your data has order and your network needs to remember what happened twenty steps ago. This node is the Neural Network Toolkit's way of putting one in a layer stack.

The inputs

Straightforward, and they map one-to-one onto PyTorch's nn.LSTM:

  • input_size (default 10) - features per time step coming in.
  • hidden_size (default 20) - width of the hidden state; the "memory capacity."
  • num_layers (default 1) - stack multiple LSTM layers on top of each other.
  • bias (True), batch_first (True), dropout (0), bidirectional (False) - the standard options. With batch_first=True (the default), your input is (batch, seq_len, features), which is the layout you'll actually want. bidirectional=True doubles the hidden state (both directions concatenated).

Output: LAYER_STACK, chained onward. There's also an NntDefineRNNLayer for plain RNNs and NntDefineGRULayer for the cheaper cousin; they're siblings with identical input fields.

How it works - and the honest caveat

Like every define node, it appends a dict - {'type': 'LSTM', 'input_size': 10, 'hidden_size': 20, ...} - to the stack list. The stack travels through the graph as a plain Python list, and it's NntCompileModel's job to turn it into real modules.

Here's where to set expectations. As of the current source, the compile path instantiates the classic layer types (Conv, Linear, Flatten, Reshape, pooling, norms) - and LSTM isn't in that dispatch list yet. The _create_* machinery exists for several attention types and the code even knows how to report nn.LSTM layers, but building one from a stack definition is still being wired up. So this node is best treated as: (a) a way to declare your intent and learn the stack format, and (b) a node whose runtime support is a work in progress. The pack's README is upfront that it's a learning tool in active development, not a production substitute for PyTorch.

If you genuinely need a trainable LSTM right now, the honest recommendation is to write it in PyTorch - the toolkit itself says it's not meant to replace that.

Common issues

  • Compiled model doesn't contain the LSTM - see above; the layer definition travels, the instantiation is pending. Watch for this before you assume your model trains.
  • Wrong tensor layout - when you do get to the point of feeding sequences in, keep batch_first=True or your (seq_len, batch, features) defaults will bite you.

Install

Ships with the pack, no per-node steps:

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

or ComfyUI Manager → "ComfyUI Neural Network Toolkit NNT", restart, and the LSTM node shows up under NNT Neural Network Toolkit/Layers. The pack's RNN-flavored example workflow leans more on Dense/Conv/Reshape nodes than on this one - worth knowing before you go hunting for a full RNN pipeline.

CategoryNNT Neural Network Toolkit/Layers

Inputs (8)

NameTypeDefaultDescription
input_sizeINT101–2048
hidden_sizeINT201–2048
num_layersINT11–10
biasCOMBOTrue2 options: True, False
batch_firstCOMBOTrue2 options: True, False
dropoutFLOAT0.00–1
bidirectionalCOMBOFalse2 options: True, False
LAYER_STACKoptLIST

Outputs (1)

NameTypeDescription
LAYER_STACKLIST