NNT Define Flatten Layer

This is the boring, essential node that makes CNN + classifier stacks work. A convolutional layer outputs a 4D tensor (batch, channels, height, width). A dense layer expects a 2D tensor (batch, features). NntDefineFlattenLayer is the bridge: it squashes everything after the batch dimension into one long vector so the dense block can chew on it.

The node has exactly one required input - LAYER_STACK - and one output, LIST. You wire the stack from your conv/pooling layers in, and the flattened stack comes out. No numbers to set, nothing to get wrong. start_dim is fixed at 1 and end_dim at -1: flatten everything except the batch.

How it works

It appends {'type': 'Flatten', 'start_dim': 1, 'end_dim': -1} to the layer stack. At compile time, NntCompileModel turns that into nn.Flatten(start_dim=1, end_dim=-1). Handy detail: the compile loop actually auto-inserts a flatten when a Linear layer follows a multi-dimensional shape - so your stack will often work without an explicit flatten node. But being explicit is better for learning (and for reading your own architecture back), which is exactly why the pack ships a dedicated node for it.

When you'd actually use it

The classic CNN shape:

• NntInputLayer [3, 32, 32]
• NntDefineConvLayer (32 filters) → Pooling
• NntDefineConvLayer (64 filters) → Pooling
• NntDefineFlattenLayer
• NntDefineDenseLayer (classifier head)
• NntCompileModel

That's the MNIST/CIFAR pattern in miniature. Without the flatten, the dense layer gets handed a 4D tensor and either errors or silently flattens it the way it wants, and shape bookkeeping gets confusing fast.

Common issues

Honestly: there aren't many, because it's one line. The one real mistake is putting it in the wrong place - flatten before the convs and you've destroyed the spatial structure for no reason. Flatten goes after your feature-extraction layers and before the classifier.

Also worth knowing: this is one of the layer types the pack's compile path handles reliably. Some of the pack's fancier transformer layer nodes are still being wired into compile; Flatten is not one of those. It's a safe, boring, works-every-time node.

Install

Ships with the pack:

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

or via ComfyUI Manager under "ComfyUI Neural Network Toolkit NNT", then restart ComfyUI. It's a thin node - a couple of minutes of reading here is more effort than the node itself takes to set up.