Memory safe wrap (middlek)

You've seen it: a workflow that runs fine once, then generation number two dies with an out-of-memory error. The usual culprit isn't your checkpoint - it's some custom node that loads its own model and never tells ComfyUI it exists. This node is the fix for that, and the whole mem-safe-wrapper pack exists to make that fix.

What it actually does

ComfyUI's "smart memory management" works because models are wrapped in a ModelPatcher. That wrapper is what gets tracked in model_management.current_loaded_models, which is how ComfyUI decides what to keep in VRAM, what to park in system RAM, and what to free when you're low. The system is genuinely good at this - on a fresh install you can stack SDXL and a couple of controlnets and it'll shuffle them around without you thinking about it.

The problem: a lot of third-party nodes just return a raw torch.nn.Module from their OUTPUT_NODE, not a ModelPatcher. ComfyUI has no idea that model exists, so it loads into VRAM and stays there forever. After a few generations you're OOM even though ComfyUI "should" have freed stuff.

Memory safe wrap (middlek) takes that orphaned model and wraps it in a proper ModelPatcher subclass, so ComfyUI's memory manager can finally see and manage it. It's a compatibility shim, and it's aimed squarely at custom node developers - though if you're running a node that leaks VRAM, it works from the workflow side too.

How it works

The wrapper class (MemSafeWrapper) subclasses ComfyUI's ModelPatcher. On creation it moves the model to the offload device (by default ComfyUI's intermediate_device(), i.e. system RAM) at the current UNet dtype. When anything calls into the model, a __getattr__ hook fires model_management.load_model_gpu(self), pulling it into VRAM just in time for the forward pass - then ComfyUI can push it right back out when it's done.

That "track it, load it on demand, offload it again" cycle is exactly what normal checkpoints do. The wrapper is what makes your custom model behave like a first-class citizen.

The inputs that matter

The node takes a MODEL in and hands a MODEL back out, plus three settings - all default to auto, and for most people auto is the right answer:

• device (auto / cpu / gpu) - where the model executes. auto uses ComfyUI's current torch device.
• offload (auto / cpu / gpu) - where the weights sit while idle. auto is RAM, which is what you want.
• dtype (auto / fp32 / fp16 / bf16) - what precision to cast to. auto follows ComfyUI's UNet dtype. Drop to bf16 on newer cards or fp16 elsewhere if VRAM is tight; fp32 is only for debugging weird numerical bugs.

The single output is a MODEL, which you wire into anything that takes one - a KSampler, a patcher, another wrapper. Nothing else changes in your workflow.

Install

Standard custom node install - this pack has no extra dependencies (it only imports torch and ComfyUI's own model_management / model_patcher) and downloads no model files:

cd {your_comfyui_path}/custom_nodes
git clone https://github.com/MiddleKD/ComfyUI-mem-safe-wrapper

Then restart ComfyUI. Or just search "ComfyUI-mem-safe-wrapper" in ComfyUI Manager and hit Install.

Where people get burned

• Attribute warnings on load. The wrapper logs a warning if the model's attributes overlap with ModelPatcher's - the model's versions get ignored in that conflict. It's a heads-up, not a crash, but if your node suddenly misbehaves after wrapping, that's the first thing to check.
• It needs a real nn.Module. Wrapping something that's already a ModelPatcher is pointless - you'd be double-wrapping. This node is for the orphaned models.
• It won't shrink VRAM usage. It makes memory manageable - ComfyUI can offload and free it - it doesn't make the model smaller. For that you want a GGUF quantized version.

Honest caveat: this is a small, single-author pack (MiddleKD, Korean README) with essentially no community footprint. It's a clever, well-scoped tool rather than a widely-tested one. If it bites you, the GitHub issues page is your only real support channel.