Ptn Pre Flatten

Dense layers want flat vectors; images arrive as grids. PtnPreFlatten is the adapter that closes that gap: it wraps a model and flattens every input to (batch, features) before the wrapped model runs. For the pack's Fashion-MNIST linear-model workflow - 28×28 pixels into a 784-input MLP - this node is the essential first step, and it's the one to reach for any time your input is rank-3 or rank-4 but your model expects rank-2.

How it works

The wrapper's forward calls torch.flatten(inputs, start_dim=1) - meaning it flattens everything after the batch axis and keeps the batch dimension intact. A (8, 1, 28, 28) tensor becomes (8, 784); a (8, 28, 28) becomes (8, 784) too. Then it passes the flattened tensor to the wrapped model. The flattened feature count just needs to match the model's expected in_features, which is why PtnLinearModel with dim_list="[784,10]" is the natural partner.

The inputs

• model - any PTMODEL, wrapped in the flatten. No other parameters exist; the only output is the wrapped model itself.

It's the counterpart to PtnPreAddChannelAxis (which adds a channel axis for conv nets rather than flattening). Where the channel-axis node is the front door for conv stacks, this one is the front door for MLPs.

How you'd use it

Chain image → PtnPreFlatten → PtnLinearModel → trainer. The pack's Fashion-MNIST linear workflow is exactly this: images come off the data loader, get flattened to 784-vectors here, and train through a linear model with an optimizer and loss. If you're swapping a conv head for an MLP head, put this node in front of the swap and the rest of the graph barely changes.

Installing

Same as every node in the pack. ComfyUI Manager → search "Pt-Wrapper", or:

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

Restart ComfyUI; the pack's requirements install on first launch.

Where people get burned

The silent mismatch is the risk: flattening is happy to produce (batch, 784) for any input shape that multiplies to 784, so if your input isn't actually 28×28 the node still runs and the next node fails with a shape error. When that happens, count your input's feature product before blaming the MLP. And remember it flattens everything after batch - if you need to preserve a channel axis, that's the job of PtnPreAddChannelAxis, not this node.