diff --git a/comfy/model_management.py b/comfy/model_management.py index a16cf380..cc2ae82a 100644 --- a/comfy/model_management.py +++ b/comfy/model_management.py @@ -581,7 +581,7 @@ def unet_offload_device(): def unet_inital_load_device(parameters, dtype): torch_dev = get_torch_device() - if vram_state == VRAMState.HIGH_VRAM: + if vram_state == VRAMState.HIGH_VRAM or vram_state == VRAMState.SHARED: return torch_dev cpu_dev = torch.device("cpu") @@ -695,7 +695,7 @@ def text_encoder_initial_device(load_device, offload_device, model_size=0): return offload_device if is_device_mps(load_device): - return offload_device + return load_device mem_l = get_free_memory(load_device) mem_o = get_free_memory(offload_device)