"vscode:/vscode.git/clone" did not exist on "80be0744a63d47f62f5f41d0507d7b68161e08da"
Commit 3b6fe51c authored by comfyanonymous's avatar comfyanonymous
Browse files

Leave text_encoder on the CPU when it can handle it.

parent b6a60fa6
...@@ -333,14 +333,19 @@ def unet_offload_device(): ...@@ -333,14 +333,19 @@ def unet_offload_device():
return torch.device("cpu") return torch.device("cpu")
def text_encoder_offload_device(): def text_encoder_offload_device():
if args.gpu_only: if args.gpu_only or vram_state == VRAMState.SHARED:
return get_torch_device() return get_torch_device()
else: else:
return torch.device("cpu") return torch.device("cpu")
def text_encoder_device(): 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() 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: else:
return torch.device("cpu") return torch.device("cpu")
......
...@@ -533,8 +533,9 @@ class CLIP: ...@@ -533,8 +533,9 @@ class CLIP:
load_device = model_management.text_encoder_device() load_device = model_management.text_encoder_device()
offload_device = model_management.text_encoder_offload_device() offload_device = model_management.text_encoder_offload_device()
self.cond_stage_model = clip(**(params)) self.cond_stage_model = clip(**(params))
if model_management.should_use_fp16(load_device): #TODO: make sure this doesn't have a quality loss before enabling.
self.cond_stage_model.half() # if model_management.should_use_fp16(load_device):
# self.cond_stage_model.half()
self.cond_stage_model = self.cond_stage_model.to() self.cond_stage_model = self.cond_stage_model.to()
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment