Prune Diffusion Model Keys

The same key-editing idea silveroxides' ComfyUI-ModelUtils applies to embeddings and LoRAs, scaled up to a full diffusion model checkpoint. Load a model, remove specific tensor keys, save the result as a new .safetensors file - this time on a file that's measured in gigabytes rather than kilobytes or megabytes.

Why you'd prune a checkpoint at all. Checkpoints occasionally carry more than the loader actually needs - EMA and non-EMA weight sets bundled together being the classic long-standing example in this ecosystem, or components you don't want (a baked-in VAE you'd rather load separately, keys left over from whatever export pipeline produced the file). Some strict loaders complain loudly about "unexpected keys" at load time even when the model works fine; trimming the offending keys ahead of time is a way to get a clean load without touching the weights you actually need. This is a targeted edit - the tensors that survive are byte-identical to the source.

How it works. diffusionmodel_name picks the checkpoint from your models folder. It loads the full file, drops every tensor whose key matches an entry in keys_to_prune, and writes the result as a new file. Nothing about the original is modified.

Inputs that matter. keys_to_prune is a multiline field - the exact key names you want gone, one per line. use_regex switches from literal matching to regex, useful for cutting a whole prefix's worth of keys (an entire EMA block, say) without listing every one by hand. output_filename defaults to pruned_model. The single output, output_path, reports where the new file landed.

Installing it. ComfyUI Manager, search "Model Utility Toolkit" - or cd ComfyUI/custom_nodes && git clone https://github.com/silveroxides/ComfyUI-ModelUtils and restart. No extra dependencies flagged in the README.

This is not a quick node. Full diffusion model checkpoints run 2-12GB depending on architecture, and this loads the whole thing into memory to do its work - expect it to take real time and real RAM, especially on a large model. There's no lazy_load/force_clear_cache toggle here the way there is on this pack's extraction and merge nodes, so budget accordingly before you queue it on your biggest checkpoint.

Know your key names first. Guessing at keys_to_prune wastes a run - a key that doesn't match exactly (or, with use_regex off, isn't an exact literal match) is silently skipped, and you get back a "pruned" file that's identical to the source with no error telling you why. If you're not sure of the exact tensor names, read the safetensors JSON header directly before running this - it's the same 8-byte-length-plus-JSON structure every safetensors file uses, and reading it executes nothing.

One execution quirk. This node isn't flagged as a workflow output, so if output_path doesn't feed into anything downstream, ComfyUI's graph pruning can skip the node entirely and the file never gets written - even though saving happens as a side effect of running the node, not through the graph connection itself. Wire the output to something, even just a Show Text node, before hitting Queue.