ImageUpscaleWithModel reserves 4.83 GB for a 17.9 MB model and evicts other loaded models
- Dominant language
- Python
- Stars
- 133k
- Forks
- 15.7k
- Avg merge
- 1d 7h
- Merged PRs (30d)
- 158
Description
### 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`:
```python
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:
```python
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.
Contributor guide
Research direction
Start in comfy_extras/nodes_upscale_model.py and trace the memory_required argument through model_management.load_models_gpu to free_memory. Reproduce the issue with the direct ImageUpscaleWithModel script on a constrained GPU, then compare the requested memory with the upscale model and tile behavior. Done means the upscale completes without an oversized reservation and an already resident checkpoint remains loaded when both models fit.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- python, pytorch
- Domain
- backend, machine-learning, performance
- Issue type
- Bug
- Difficulty
- 4/5
- Estimated time
- 3-5 days
- Activity status
- Active
- Clarity
- Mostly clear
- Newbie friendliness
- 58/100