wangbo/ComfyUI
1
0
Fork 0
mirror of https://github.com/comfyanonymous/ComfyUI.git synced 2026-03-08 02:37:42 +08:00
Code Issues Actions 11 Packages Projects Releases Wiki Activity

Initial load device should be cpu when using dynamic vram. (#12766)

Browse Source
This commit is contained in:
comfyanonymous 2026-03-04 13:33:14 -08:00 committed by GitHub
parent 8811db52db
commit ac4a943ff3
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
1 changed files with 8 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

10
comfy/model_management.py
View File

@ -830,11 +830,14 @@ def unet_offload_device():
return torch.device("cpu") return torch.device("cpu")
def unet_inital_load_device(parameters, dtype): def unet_inital_load_device(parameters, dtype):
cpu_dev = torch.device("cpu")
if comfy.memory_management.aimdo_enabled:
return cpu_dev
torch_dev = get_torch_device() torch_dev = get_torch_device()
if vram_state == VRAMState.HIGH_VRAM or vram_state == VRAMState.SHARED: if vram_state == VRAMState.HIGH_VRAM or vram_state == VRAMState.SHARED:
return torch_dev return torch_dev
cpu_dev = torch.device("cpu")
if DISABLE_SMART_MEMORY or vram_state == VRAMState.NO_VRAM: if DISABLE_SMART_MEMORY or vram_state == VRAMState.NO_VRAM:
return cpu_dev return cpu_dev
@ -842,7 +845,7 @@ def unet_inital_load_device(parameters, dtype):
mem_dev = get_free_memory(torch_dev) mem_dev = get_free_memory(torch_dev)
mem_cpu = get_free_memory(cpu_dev) mem_cpu = get_free_memory(cpu_dev)
if mem_dev > mem_cpu and model_size < mem_dev and comfy.memory_management.aimdo_enabled: if mem_dev > mem_cpu and model_size < mem_dev:
return torch_dev return torch_dev
else: else:
return cpu_dev return cpu_dev
@ -945,6 +948,9 @@ def text_encoder_device():
return torch.device("cpu") return torch.device("cpu")
def text_encoder_initial_device(load_device, offload_device, model_size=0): def text_encoder_initial_device(load_device, offload_device, model_size=0):
if comfy.memory_management.aimdo_enabled:
return offload_device
if load_device == offload_device or model_size <= 1024 * 1024 * 1024: if load_device == offload_device or model_size <= 1024 * 1024 * 1024:
return offload_device return offload_device