mirrors/ComfyUI
1
0
Fork 0
mirror of https://github.com/comfyanonymous/ComfyUI.git synced 2025-09-12 12:37:01 +00:00
Code Issues Packages Projects Releases Wiki Activity

Leave text_encoder on the CPU when it can handle it.

Browse Source
This commit is contained in:
comfyanonymous
2023-07-01 14:38:51 -04:00
parent b6a60fa696
commit 3b6fe51c1d
2 changed files with 10 additions and 4 deletions
Download Patch File Download Diff File Expand all files Collapse all files

9
comfy/model_management.py
View File

@@ -333,14 +333,19 @@ def unet_offload_device():
return torch.device("cpu")
def text_encoder_offload_device():
if args.gpu_only:
if args.gpu_only or vram_state == VRAMState.SHARED:
return get_torch_device()
else:
return torch.device("cpu")
def text_encoder_device():
if vram_state == VRAMState.HIGH_VRAM or vram_state == VRAMState.SHARED or vram_state == VRAMState.NORMAL_VRAM:
if args.gpu_only or vram_state == VRAMState.SHARED:
return get_torch_device()
elif vram_state == VRAMState.HIGH_VRAM or vram_state == VRAMState.NORMAL_VRAM:
if torch.get_num_threads() < 8: #leaving the text encoder on the CPU is faster than shifting it if the CPU is fast enough.
return get_torch_device()
else:
return torch.device("cpu")
else:
return torch.device("cpu")