ComfyUI Node

SetModelTrainable

Freeze or unfreeze an entire model with one dropdown

By TashaSkyUp·Created about a year ago·Updated about a year ago· 1
SetModelTrainable
  • model
  • TORCH_MODEL
requires_gradFalse

Freezing a model - setting every parameter's requires_grad to False - is how you stop training from touching weights you want left alone. SetModelTrainable does exactly that to an entire TORCH_MODEL in one shot, and flips it back on with the same dropdown. It's the transfer-learning workhorse of the pack: keep a feature extractor frozen, let a head train, or vice versa.

It's part of EternalKernel PyTorch Nodes (TashaSkyUp), the raw-PyTorch-in-ComfyUI pack. In the broader ecosystem, "freezing" usually means freezing UNet or text-encoder weights during a LoRA run; here it means the same idea applied to whatever nn.Sequential you built, which is refreshingly plain.

How it works

The implementation is brutally simple: loop over model.parameters() and set param.requires_grad to the flag. False freezes (no gradients, no updates during backprop), True unfreezes. The dropdown defaults to False, and the node returns the same model so you can keep chaining.

One thing to know so you don't build on a false foundation: TrainModel force-unfreezes everything. Its training loop explicitly sets requires_grad = True on all parameters before it starts, so if you freeze a model and then hand it to TrainModel, the freeze is silently undone. Freezing in this pack is useful for: running inference with PyTorchInferenceNode on a model you don't want changing, building a model with ExtractLayersAsModel (which has its own freeze option), or - if you're wiring a manual train/infer split - keeping a pretrained half out of a training pass that doesn't go through TrainModel. Read that as: it works, but check which training path your graph uses.

Inputs and outputs

  • model (required TORCH_MODEL) - the network to freeze or unfreeze.
  • requires_grad (required enum) - True / False, default False.
  • Output: TORCH_MODEL - the same model, with flags set.

Installing it

Shared with the whole pack:

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

Restart ComfyUI and it's under ETK/pytorch - or ComfyUI Manager, searching "EternalKernel PyTorch Nodes".

Troubleshooting

The big one is the TrainModel interaction above - you'll freeze, train, and wonder why your frozen weights drifted. That's the force-unfreeze in the training loop, not a hardware problem. Second, freezing can silently reduce memory use (no gradient buffers for frozen params), but don't expect dramatic VRAM savings on the toy-sized models this pack builds. And since the pack's validate_inputs patch tolerates type mismatches, a None passed in as the model will fail with an unhelpful 'NoneType' object has no attribute 'parameters' - trace the wire, it's usually a node upstream that didn't get fed.

CategoryETK/pytorch

Inputs (2)

NameTypeDefaultDescription
modelTORCH_MODEL
requires_gradCOMBOFalse2 options: True, False

Outputs (1)

NameTypeDescription
TORCH_MODELTORCH_MODEL