From 8bdaffe3b8542ce2a54d2012f938121f0d9182de Mon Sep 17 00:00:00 2001
From: John Slegers <slegersjohn@gmail.com>
Date: Mon, 12 Aug 2024 03:01:55 +0200
Subject: [PATCH] Recognize CPU state, low VRAM & high VRAM when no flags are
 set

When flags like args.cpu, args.highvram or args.gpu_only aren't set, model_management should detect the correct state automatically
---
 comfy/model_management.py | 36 +++++++++++++++++++++++++-----------
 1 file changed, 25 insertions(+), 11 deletions(-)

diff --git a/comfy/model_management.py b/comfy/model_management.py
index a0105131d..3c8781690 100644
--- a/comfy/model_management.py
+++ b/comfy/model_management.py
@@ -81,6 +81,10 @@ except:
 if args.cpu:
     cpu_state = CPUState.CPU
 
+def is_nvidia():
+    global cpu_state
+    return cpu_state == CPUState.GPU and torch.cuda.is_available()
+
 def is_intel_xpu():
     global cpu_state
     global xpu_available
@@ -89,6 +93,10 @@ def is_intel_xpu():
             return True
     return False
 
+if args.cpu or (not is_nvidia() and not is_intel_xpu()):
+    cpu_state = CPUState.CPU
+
+
 def get_torch_device():
     global directml_enabled
     global cpu_state
@@ -107,11 +115,12 @@ def get_torch_device():
 
 def get_total_memory(dev=None, torch_total_too=False):
     global directml_enabled
+    global total_ram
     if dev is None:
         dev = get_torch_device()
 
     if hasattr(dev, 'type') and (dev.type == 'cpu' or dev.type == 'mps'):
-        mem_total = psutil.virtual_memory().total
+        mem_total = total_ram
         mem_total_torch = mem_total
     else:
         if directml_enabled:
@@ -134,9 +143,21 @@ def get_total_memory(dev=None, torch_total_too=False):
     else:
         return mem_total
 
-total_vram = get_total_memory(get_torch_device()) / (1024 * 1024)
-total_ram = psutil.virtual_memory().total / (1024 * 1024)
-logging.info("Total VRAM {:0.0f} MB, total RAM {:0.0f} MB".format(total_vram, total_ram))
+total_ram = psutil.virtual_memory().total
+total_vram = get_total_memory(get_torch_device())
+logging.info(
+    "Total VRAM {:0.0f} MB, total RAM {:0.0f} MB".format(
+        total_vram / (1024 * 1024),
+        total_ram / (1024 * 1024)
+    )
+)
+
+if cpu_state != CPUState.CPU:
+    if total_ram < 6 * (1024 * 1024 * 1024) :
+        set_vram_to = VRAMState.LOW_VRAM
+        lowvram_available = True
+    if total_ram > 12 * (1024 * 1024 * 1024) :
+        vram_state = VRAMState.HIGH_VRAM
 
 try:
     logging.info("pytorch version: {}".format(torch.version.__version__))
@@ -173,13 +194,6 @@ else:
     except:
         XFORMERS_IS_AVAILABLE = False
 
-def is_nvidia():
-    global cpu_state
-    if cpu_state == CPUState.GPU:
-        if torch.version.cuda:
-            return True
-    return False
-
 ENABLE_PYTORCH_ATTENTION = False
 if args.use_pytorch_cross_attention:
     ENABLE_PYTORCH_ATTENTION = True