SetModelTrainable
Freeze or unfreeze an entire model with one dropdown
- model
- TORCH_MODEL
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, defaultFalse. - 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.
Inputs (2)
| Name | Type | Default | Description |
|---|---|---|---|
| model | TORCH_MODEL | — | |
| requires_grad | COMBO | False | 2 options: True, False |
Outputs (1)
| Name | Type | Description |
|---|---|---|
| TORCH_MODEL | TORCH_MODEL | — |