ImageUpscaleWithModel reserves 4.83 GB for a 17.9 MB model and evicts other loaded models

[Dhevenddra]

Expected Behavior

Upscale Image (using Model) should reserve roughly what the upscale actually needs. A checkpoint that is already resident should stay resident when there is room for both.

Actual Behavior

The node asks load_models_gpu for a fixed amount that does not depend on the size of the upscale model. For a float32 image and a scale-4 model that works out to 4.83 GB, which is more than a 4 GB card has in total, so free_memory() unloads everything else before the upscale runs. A checkpoint that was loaded before the upscale is gone afterwards and has to be reloaded for the next generation.

From comfy_extras/nodes_upscale_model.py:

memory_required = (512 * 512 * 3) * image.element_size() * max(upscale_model.scale, 1.0) * 384.0
memory_required += image.nelement() * image.element_size()
model_management.load_models_gpu([upscale_model.patcher], memory_required=memory_required, force_full_load=True)

786432 * 4 bytes * 4 * 384.0 = 4.83 GB. The actual weights of RealESRGAN_x4plus_anime_6B are 17.9 MB, so the request is about 270x the model. The first term is also independent of the input image, since it is fixed at one 512x512 tile.

memory_required becomes extra_mem inside load_models_gpu, which is passed to free_memory(total_memory_required[device] * 1.1 + extra_mem, device). That is the call that does the unloading.

The line already carries #The 384.0 is an estimate of how much some of these models take, TODO: make it more accurate, so this is a known approximation. I am filing it because the consequence on small cards is bigger than it looks.

To be clear about scope: this is not the crash in #15433. That one is fixed by force_full_load=True from #15437, which I verified by bisection. Because the load is now forced, the estimate no longer blocks the upscale. It only evicts everything else.

Steps to Reproduce

On a 4 GB card:

1. Load any SD1.5 checkpoint so it becomes resident.
2. Run Upscale Image (using Model) with any small ESRGAN model.
3. The checkpoint is no longer loaded afterwards.

Reproduced directly against the node rather than through a workflow:

model = comfy.sd.load_checkpoint_guess_config(CKPT, ...)[0]
comfy.model_management.load_models_gpu([model])
print([type(m.model.model).__name__ for m in mm.current_loaded_models])

num.ImageUpscaleWithModel().execute(upscale_model, torch.rand(1, 512, 512, 3))
print([type(m.model.model).__name__ for m in mm.current_loaded_models])

Debug Logs

ComfyUI master b323a34
Windows, NVIDIA GeForce RTX 3050 Laptop 4 GB, torch 2.10.0+cu130
upscale model: RealESRGAN_x4plus_anime_6B.pth (17.9 MB of weights, scale 4)

free VRAM at start: 3.46 GB
after loading checkpoint  -> loaded models: ['BaseModel']
   free VRAM: 1.69 GB

upscale node will request memory_required = 4.83 GB
   total VRAM on this card = 4.29 GB
   request exceeds the whole card by 1.1x

after the upscale         -> loaded models: ['RRDBNet']

BaseModel is present before the upscale and gone after it.

Other

Changing only the 384.0 factor to 24.0 and rerunning the same script keeps both models resident, and the upscale still completes:

after the upscale         -> loaded models: ['RRDBNet', 'BaseModel']

I am not proposing 24.0 as a value. It is only there to show that the factor is what drives the eviction. A heavier upscale model would need more headroom than that, and picking a new constant would have the same problem as the current one. Deriving the estimate from the tile size and the loaded model looks like the direction, since the tile is already fixed at 512 with an OOM backoff loop underneath that halves it.

Happy to open a PR if you have a preference for how the estimate should be computed.

[YiGeSama]

Adding a data point from my setup, since the numbers in this report originally came from it.

Environment: RTX 3050 Ti Laptop 4 GB (~3.2-3.3 GB free), ComfyUI v0.33.0, RealESRGAN_x4plus_anime_6B.pth (17.9 MB on disk).

In a real workflow the upscale node typically runs while a checkpoint is already resident on the GPU. The node computes memory_required from the input image size and the fixed x384 fudge factor, then passes it to load_models_gpu — even though the actual weights are only ~18 MB. The resulting 4.83 GB figure (for a 512x512 float32 input at scale 4) trips the VRAM manager's eviction logic, so the resident checkpoint gets offloaded and has to be reloaded afterwards. That's the user-visible cost today: not a crash anymore (thanks to #15437), but load/unload churn around every upscale step — plus extra latency and disk reads on low-end cards.

Formula for the record: 512*512*3 * 4 bytes * scale(4) * 384.0 ≈ 4.83 GB, while the real model is three orders of magnitude smaller.

An estimate based on the actual model size (parameter count x element size, with a saner headroom) instead of the fixed x384 factor would fix both the misleading number and the eviction behavior.

[a-yeyang]

I traced the current reservation path and would like to work on this. The fixed 384.0 multiplier is passed to load_models_gpu() as transient inference memory, while model weights are accounted for separately, so the estimate can evict resident checkpoints well beyond the buffers used by tiled upscaling.

I plan to replace it with a shape-, dtype-, and device-aware bounded reservation covering the moved full input, the largest tile’s float32 conversion/output, and—only when the intermediate output shares the execution device—the full output plus the one-channel blending divisor used by the multi-tile path. Model-specific activations will remain protected by the existing minimum inference reservation and OOM tile-size backoff.

I’ll add CPU-only regression tests covering sub-tile and 512×512 inputs, batches, float16 conversion, non-integer scales/rounding, and both CPU and execution-device intermediate outputs. If you are already preparing a fix, please let me know so I do not duplicate it.