Fix issue with regular torch version.

2be27427 · comfyanonymous · 89a0767a · 2be27427
Commit 2be27427 authored Aug 17, 2023 by comfyanonymous
Show whitespace changes
Inline Side-by-side

Showing with 5 additions and 1 deletion

comfy/model_management.py comfy/model_management.py +5 -1

No files found.
--- a/comfy/model_management.py
+++ b/comfy/model_management.py
@@ -390,7 +390,11 @@ def unet_inital_load_device(parameters, dtype):
        return torch_dev

    cpu_dev = torch.device("cpu")
-    model_size = dtype.itemsize * parameters
+    dtype_size = 4
+    if dtype == torch.float16 or dtype == torch.bfloat16:
+        dtype_size = 2
+
+    model_size = dtype_size * parameters

    mem_dev = get_free_memory(torch_dev)
    mem_cpu = get_free_memory(cpu_dev)