Comfy-Org / Comfy-Org/ComfyUI

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)

Open
#15,433 7 comments 1 reaction 0 assignees View on GitHub
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

Open the contributing 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

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.