Add a launch argument for non_blocking=True
- 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
Assessment
This issue has not been assessed yet.