code cleanup

This commit is contained in:
Francesco Yoshi Gobbo 2023-03-27 06:48:09 +02:00
parent cf0098d539
commit f55755f0d2
No known key found for this signature in database
GPG Key ID: 1B931AAE47014274
1 changed files with 3 additions and 2 deletions

View File

@ -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])