Comfy-Org / Comfy-Org/ComfyUI

Add a launch argument for non_blocking=True

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

Description

### Feature Idea

In `model_management.py`
```
def device_should_use_non_blocking(device):
if not device_supports_non_blocking(device):
return False
return False
# return True #TODO: figure out why this causes memory issues on Nvidia and possibly others
```

https://github.com/comfyanonymous/ComfyUI/blob/7390ff3b1ec2e15017ba4a52d6eaabc4aa4636e3/comfy/model_management.py#L834

Changing this function back to its pre-TODO state results in a large speedup in model patching. (19s -> 6s for Flux LoRAs on my computer.) It probably also speeds up loading in other areas.

This is because the largest bottleneck is the one-by-one blocking transfer of each layer of the unet to the GPU, which is massively accelerated if non_blocking=True.

Are there still memory issues? Changes like this (https://github.com/comfyanonymous/ComfyUI/commit/39f114c44bb99d4a221e8da451d4f2a20119c674) since this TODO was written could mean the same problems that used to cause memory issues may be less relevant than before or non-existent.

Please consider re-adding support for non_blocking=True as a launch argument so users can start trying it out again.

### Existing Solutions

_No response_

### Other

_No response_

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.