ComfyUI Node

NNT Compile Model

NntCompileModel — the glue that turns a pile of layer nodes into a real model

By inventorado·Created 2 years ago·Updated 2 years ago· 69
NNT Compile Model
  • LAYER_STACK
  • hyperparameters
  • model
  • report
  • script
modeCompile
activation_functionReLU
normalizationNone
padding_modezeros
weight_initkaiming_normal
activation_params{}

Every Neural Network Toolkit workflow funnels through this node. The layer-definition nodes (Dense, Conv, Flatten, and the rest) just build up a Python list - the LAYER_STACK. NntCompileModel is where that list stops being a shopping cart and becomes an actual trainable nn.Module. Without it you have layers and nowhere to run them.

First, the name. "Compile" here means assemble, not torch.compile. There's no JIT, no graph fusion, no speed magic. The node walks your layer stack, instantiates each layer as real PyTorch, tracks the tensor shape through the network, and hands you back a model wrapped in an nn.Sequential. It's "compile the architecture," the way you'd compile a list of parts into a machine.

The inputs that matter

  • LAYER_STACK - the list output of any define-layer node. Chain Dense → Conv → Flatten and wire the last LIST into here. If you want layers, this is non-negotiable.
  • mode - Compile builds the model; Only create script skips the build and just emits the Python source. That second mode is a genuinely nice teaching feature: you get a .py version of your network you can study or run outside ComfyUI.
  • activation_function (default ReLU, 28 choices) and normalization - these are the global defaults applied to layers that don't specify their own. Most define-nodes carry their own activation/norm settings, so these act as fallbacks.
  • weight_init - 12 methods from default to kaiming_normal, xavier_uniform, orthogonal, even dirac. Applies to layers as they're built. Kaiming is the sensible default for ReLU networks; leave it alone until you're experimenting.
  • activation_params - a JSON-ish string, e.g. {"negative_slope": 0.01}, for activations that take arguments.

Outputs: model (MODEL - the thing you feed to NntTrainModel, NntFineTuneModel, or NntAnalyzeModel), report (STRING - build status), and script (STRING - the generated Python, empty in compile mode).

How it works, honestly

The compile loop handles the classic layer types - Input, Conv1d/2d/3d, Linear, Flatten, Reshape, MaxPool2d/AvgPool2d, and the norm layers - and it's real about shapes: it computes the output size after each conv and auto-inserts a flatten before a dense layer when needed. It also wires in weight init, normalization, and dropout per layer.

The honest caveat: the transformer and recurrent layer definitions (LinearAttention, MultiheadAttention, GRU, LSTM and friends) append to the stack, but as of the current source they aren't dispatched by the compile loop - they fall through silently. The pack is a work in progress, and the author says so on the tin. If your compiled model misbehaves or ignores a layer, check whether you've mixed in one of those not-yet-wired layer types. Classic CNNs and MLPs - the MNIST/CIFAR workflows the pack ships - compile and train fine.

Common issues

  • No model output - you set mode to Only create script. The model output is None in that mode by design.
  • Shape mismatch errors - usually a conv output that doesn't line up with the next dense layer. Watch the shapes the report prints, or add a Flatten/Pooling layer.
  • Nothing happens / layer ignored - you used a layer type the compile path doesn't materialize yet (see above).

Install

This is the pack's hub node, so it installs with the pack - no extra steps beyond getting NNT itself:

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

Or search "ComfyUI Neural Network Toolkit NNT" in ComfyUI Manager and restart. The requirements are heavy (transformers, onnx, shap, torchview, graphviz...), so give the first install time to settle.

CategoryNNT Neural Network Toolkit/Models

Inputs (8)

NameTypeDefaultDescription
modeCOMBOCompile2 options: Compile, Only create script
LAYER_STACKLIST
activation_functionCOMBOReLU28 options: None, ELU, GELU, GLU, Hardshrink, Hardsigmoid, +22
normalizationCOMBONone8 options: None, BatchNorm, LayerNorm, InstanceNorm1d, InstanceNorm2d, InstanceNorm3d, +2
padding_modeCOMBOzeros4 options: zeros, reflect, replicate, circular
weight_initCOMBOkaiming_normal12 options: default, normal, uniform, xavier_normal, xavier_uniform, kaiming_normal, +6
activation_paramsSTRING{}
hyperparametersoptDICT

Outputs (3)

NameTypeDescription
modelMODEL
reportSTRING
scriptSTRING