ImageUpscaleWithModel fails on low-VRAM GPUs (4GB): "Input type (torch.cuda.FloatTensor) and weight type (torch.FloatTensor) should be the same" (v0.29+ regression, broken in 0.30/0.31)
- Dominant language
- Python
- Stars
- 133k
- Forks
- 15.7k
- Avg merge
- 1d 6h
- Merged PRs (30d)
- 155
Description
## Bug Description
`ImageUpscaleWithModel` (core **Upscale Image (using Model)** node) crashes on low-VRAM GPUs in v0.31.0 with:
```
RuntimeError: Input type (torch.cuda.FloatTensor) and weight type (torch.FloatTensor) should be the same
```
This is a **regression**: the same workflow ran fine through v0.28 (older code called `upscale_model.to(device)` unconditionally, so the weights were always moved to GPU). The breakage was introduced in **v0.29.0** (commit f8a3fd9d, PR #15063) when the node was switched to a ModelPatcher + `load_models_gpu(memory_required=...)`; v0.30.0 and v0.31.0 have identical code.
## Environment
- ComfyUI v0.31.0 (also present on current master `dd79c643a`; not fixed in v0.31.1)
- Windows 10, NVIDIA GeForce RTX 3050 Ti Laptop, **4 GB VRAM** (~3.2–3.4 GB free)
- PyTorch 2.7.0+cu128, Python 3.11.6
- Upscale model: `RealESRGAN_x4plus_anime_6B.pth` (~67 MB)
## Root Cause
v0.31 rewrote the node to load the upscale model through a `ModelPatcher` with an upfront VRAM estimate (`comfy_extras/nodes_upscale_model.py`):
```python
memory_required = (512 * 512 * 3) * image.element_size() * max(upscale_model.scale, 1.0) * 384.0 #The 384.0 is an estimate ...
memory_required += image.nelement() * image.element_size()
model_management.load_models_gpu([upscale_model.patcher], memory_required=memory_required)
```
With a float32 image and scale 4 the estimate is **~4.5 GB** (the `384.0` fudge factor dominates). On a 4 GB card with ~3.4 GB free, `load_models_gpu` concludes the model cannot fit and leaves the weights on CPU — see log:
```
Requested to load RRDBNet
0 models unloaded.
loaded completely; 0.00 MB usable, 0.00 MB loaded, full load: False
```
But the node still moves the input tensor to CUDA:
```python
in_img = image.movedim(-1,-3).to(device)
```
→ `F.conv2d` receives CUDA input × CPU weights → the error above. The actual model is only ~67 MB; it's the *estimate* that doesn't fit, so the load is skipped entirely.
## Reproduction (verified locally)
```python
# ups = UpscaleModelLoader.execute(...) via the real node code path
mm.load_models_gpu([ups.patcher], memory_required=memory_required) # estimate ~4.5 GB, free ~3.2 GB
sorted({str(p.device) for p in ups.model.parameters()}) # -> ['cpu'] (nothing loaded)
ups(torch.rand(1, 768, 1344, 3).movedim(-1, -3).to('cuda').float()) # RuntimeError: Input type (torch.cuda.FloatTensor)
# and weight type (torch.FloatTensor) should be the same
```
## Suggested Fix
Force the (tiny) weights onto the load device after the load call — restores pre-0.31 behavior:
```python
device = upscale_model.patcher.load_device
...
model_management.load_models_gpu([upscale_model.patcher], memory_required=memory_required)
upscale_model.patcher.model.to(device)
```
Verified locally: weights land on `cuda:0` and the node completes a 4× upscale of a 1344×768 image (5376×3072 output) with no error.
The `384.0` factor in the `memory_required` formula is also worth revisiting — for typical ESRGAN-style models it overestimates the footprint by orders of magnitude, and that inflated estimate is exactly what makes `load_models_gpu` refuse to load on small GPUs.
Contributor guide
Research direction
Start in comfy_extras/nodes_upscale_model.py at the ImageUpscaleWithModel execution path and inspect the load_models_gpu call and subsequent tensor device movement. Reproduce on a 4GB CUDA setup, then verify that model and input devices match and that the 4× upscale completes; also check how the memory_required estimate affects loading.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- python, pytorch
- Domain
- backend, machine-learning, performance
- Issue type
- Bug
- Difficulty
- 3/5
- Estimated time
- 1-2 days
- Activity status
- Quiet
- Clarity
- Clearly specified
- Newbie friendliness
- 72/100