Comfy-Org / Comfy-Org/ComfyUI

Use too many VRAM (and very slow) when running big models after update

Open
#5,737 7 comments 0 reactions 0 assignees View on GitHub
Potential Bug
Dominant language
Python
Stars
133k
Forks
15.7k
Avg merge
1d 7h
Merged PRs (30d)
158

Description

### Expected Behavior

On Windows 11, I have 16GB VRAM NVIDIA card, before I can run full size model of Mochi, i.e. first example in [https://comfyanonymous.github.io/ComfyUI_examples/mochi/](https://comfyanonymous.github.io/ComfyUI_examples/mochi/) in about 3 minutes, because ComfyUI will partially load the model.

```
got prompt
Requested to load MochiTEModel_
Loading 1 new model
loaded completely 0.0 9083.38671875 True
model weight dtype torch.bfloat16, manual cast: None
model_type FLOW
Requested to load GenmoMochi
Loading 1 new model
loaded partially 11656.891287994385 11653.267669677734 0
43%|███████████████████████████████████▌ | 13/30 [01:25<01:49, 6.46s/it]
```

![Good VRAM usage](https://github.com/user-attachments/assets/584e96cd-d4ca-4ab0-967d-ee216d5055a0)

ComfyUI didn't use any *shared* GPU memory, which is good, because shared GPU memory is *quite* slow.

### Actual Behavior

After updating, ComfyUI seems to use more dedicated GPU memory than it can. So it starts to use shared GPU memory, which is very slow.

```
got prompt
Requested to load MochiTEModel_
Loading 1 new model
loaded completely 0.0 18166.773441314697 True
model weight dtype torch.bfloat16, manual cast: None
model_type FLOW
Requested to load GenmoMochi
Loading 1 new model
loaded partially 11656.891287994385 19126.324310302734 0
3%|██▋ | 1/30 [04:15<2:03:26, 255.41s/it]
```

![Bad VRAM usage](https://github.com/user-attachments/assets/21e0dbf9-2bce-4118-9a24-32e9f9398d83)

### Steps to Reproduce

1. Update ComfyUI to newest version.
2. Run large model and large data (video) with not very large VRAM graphics card with KSampler.

### Debug Logs

```powershell
As above.
```

### Other

After using git bisect, I found this commit cause the issue: 41444b5236c613ee5ff7d09aa620f94b702f30eb.
Temporary workaround would be just `git revert 41444b5236c613ee5ff7d09aa620f94b702f30eb`.

This is for the KSampler part. The VAE Decode part seems to always use beyond dedicated GPU memory even in old version for large data. I have already worked around this by making it always use tiled VAE decoding, but this should be investigated also:

comfy/sd.py:
![VAE VRAM usage workaround](https://github.com/user-attachments/assets/b0ffdc13-241e-4bc8-b50a-32a0233cace9)

The result would be with a little difference between original VAE decoding and tiled VAE decoding.

But for LTX-Video, the above trick causes error in VAE Decode: (Issue #5735)
```
RuntimeError: The expanded size of the tensor (192) must match the existing size (768) at non-singleton dimension 4. Target sizes: [1, 3, 104, 128, 192]. Tensor sizes: [3, 97, 512, 768]
```

Contributor guide

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.