--disable-smart-memory now disables loading model directly to vram.

This commit is contained in:
comfyanonymous 2023-08-20 04:00:53 -04:00
parent c9b562aed1
commit e9469e732d
1 changed files with 3 additions and 0 deletions

View File

@ -397,6 +397,9 @@ def unet_inital_load_device(parameters, dtype):
return torch_dev return torch_dev
cpu_dev = torch.device("cpu") cpu_dev = torch.device("cpu")
if DISABLE_SMART_MEMORY:
return cpu_dev
dtype_size = 4 dtype_size = 4
if dtype == torch.float16 or dtype == torch.bfloat16: if dtype == torch.float16 or dtype == torch.bfloat16:
dtype_size = 2 dtype_size = 2