From f55755f0d24255a1207a803dc30efe9e4466bfe1 Mon Sep 17 00:00:00 2001
From: Francesco Yoshi Gobbo <yoshi@fgobbo.com>
Date: Mon, 27 Mar 2023 06:48:09 +0200
Subject: [PATCH] code cleanup

---
 comfy/model_management.py | 5 +++--
 1 file changed, 3 insertions(+), 2 deletions(-)

diff --git a/comfy/model_management.py b/comfy/model_management.py
index d9498e29..4aa47ff1 100644
--- a/comfy/model_management.py
+++ b/comfy/model_management.py
@@ -15,6 +15,8 @@ total_vram_available_mb = -1
 import sys
 import psutil
 
+forced_cpu = "--cpu" in sys.argv
+
 set_vram_to = NORMAL_VRAM
 
 try:
@@ -22,7 +24,6 @@ try:
     total_vram = torch.cuda.mem_get_info(torch.cuda.current_device())[1] / (1024 * 1024)
     total_ram = psutil.virtual_memory().total / (1024 * 1024)
     forced_normal_vram = "--normalvram" in sys.argv
-    forced_cpu = "--cpu" in sys.argv
     if not forced_normal_vram and not forced_cpu:
         if total_vram <= 4096:
             print("Trying to enable lowvram mode because your GPU seems to have 4GB or less. If you don't want this use: --normalvram")
@@ -84,7 +85,7 @@ try:
 except:
     pass
 
-if "--cpu" in sys.argv:
+if forced_cpu:
     vram_state = CPU
 
 print("Set vram state to:", ["CPU", "NO VRAM", "LOW VRAM", "NORMAL VRAM", "HIGH VRAM", "MPS"][vram_state])