From cb87b2d4832684790d453b9ff5e6cf762f1ad430 Mon Sep 17 00:00:00 2001
From: Dalton <murraydalton1@comcast.net>
Date: Mon, 25 Mar 2024 17:01:38 -0400
Subject: [PATCH] Prepare further startup check & move away from nvidia-smi

---
 modules/startup_checks.py | 66 ++++++++++++++++++++++++++++-----------
 1 file changed, 47 insertions(+), 19 deletions(-)

diff --git a/modules/startup_checks.py b/modules/startup_checks.py
index 9538061ee..746c511e9 100644
--- a/modules/startup_checks.py
+++ b/modules/startup_checks.py
@@ -1,9 +1,10 @@
 import sys
 import os
 import subprocess
-import json
 
-## Functions for checking if user wants to skip startup
+import torch
+
+## Functions for checking if user wants to skip startup checks
 SKIP_STARTUP_CHECKS = "./skip_startup_checks.txt"
 
 def startup_skip_check():
@@ -14,15 +15,29 @@ def startup_skip_check():
         pass
 
 ## Functions for checks
+# Performs gpu brand check
+def get_gpu_brand():
+    try:
+        cuda_is_available = torch.cuda.is_available()
+        if cuda_is_available == True:
+            gpu_brand = torch.cuda.get_device_name()
+            gpu_brand = gpu_brand.split()[0]
+            print("\nGPU brand:", gpu_brand)
+            return gpu_brand
+        else:
+            print("Cuda is not currently available")
+    except:
+        print("Torch or Cuda is not currently available")
+
+# Performs vram check
 def vram_check():
     try:
-        output = subprocess.check_output(['nvidia-smi', '--query-gpu=memory.total', '--format=csv,noheader,nounits'])
-        total_vram = sum(map(int, output.decode('utf-8').strip().split('\n')))
+        total_vram = torch.cuda.get_device_properties(0).total_memory // 1024 ** 2
     except (subprocess.CalledProcessError, FileNotFoundError):
         print("\nUnable to get VRAM amount")
 
     if total_vram is not None:
-        print(f"\nTotal VRAM available: {total_vram} bytes")
+        print(f"Total VRAM: {total_vram} MegaBytes")
         if total_vram > 8192:
             return 0
         elif total_vram <= 8192 and total_vram > 4096:
@@ -33,10 +48,8 @@ def vram_check():
         return 3
 
 # Function to write new args to webui-user.bat file
-def write_commandline_args():
+def write_commandline_args(new_commandline_args):
     bat_file = "./webui-user.bat"
-    ## TODO: Switch this to use actual dynamic commandline args from user options and system info
-    new_commandline_args = ""
 
     # Read webui-user.bat file and save lines
     with open(bat_file, 'r') as file:
@@ -54,6 +67,8 @@ def write_commandline_args():
 
 ## Main starting check sequence
 def start_check_sequence():
+    new_commandline_args = ""
+
     print("Running system checks")
 
     # Initial check to see if they wish to perform startup checks at each launch or not - It's really not necessary to
@@ -65,29 +80,42 @@ def start_check_sequence():
         with open(SKIP_STARTUP_CHECKS, "w") as file:
             file.write("1")
 
-    # Skip webui install at each start?
-    ## TODO: Implement check if already answered
+    ## Skip webui install at each start?
     skip_webui_install = str(input("\nWould you like to skip stable-diffusion webui installation at each startup? (y)es or (n)o: ")).lower()
     while (skip_webui_install != "yes") and (skip_webui_install != "y") and (skip_webui_install != "no") and (skip_webui_install != "n"):
         print("\nPlease choose a valid response (y)es or (n)o")
         skip_webui_install = str(input("Would you like to skip stable-diffusion webui installation at each startup? (y)es or (n)o: ")).lower()
     if (skip_webui_install == "yes") or (skip_webui_install == "y"):
-        ## TODO
-        pass
+        new_commandline_args += "--skip-install "
 
-    # Perform checks and set variables
+    ## Perform checks and set variables
+    # GPU brand
+    gpu_brand = get_gpu_brand()
+    if (gpu_brand == "NVIDIA"):
+        new_commandline_args += "--xformers "
+
+    # VRAM check
     vram_low_med = vram_check()
 
-    if vram_low_med == 0:
+    if (vram_low_med == 0):
         pass
-    elif vram_low_med == 1:
-        # TODO
+    elif (vram_low_med == 1):
         print("Medium vram")
-    elif vram_low_med == 2:
-        # TODO
+        new_commandline_args += "--medvram "
+    elif (vram_low_med == 2):
         print("Low vram")
+        new_commandline_args += "--lowvram "
     else:
-        print("NVIDIA GPU or driver not found.")
+        pass
+
+    # Ask user to confirm command line arguments, perform command line argument write
+    print("\"", new_commandline_args, "\"")
+    confirm_write = str(input("\nWould you like to set these arguments?\nPlease note that this will get rid of any existing arguments\n(y)es or (n)o: ")).lower()
+    while (confirm_write != "yes") and (confirm_write != "y") and (confirm_write != "no") and (confirm_write != "n"):
+        print("\nPlease choose a valid response (y)es or (n)o")
+        confirm_write = str(input("Would you like to skip stable-diffusion webui installation at each startup? (y)es or (n)o: ")).lower()
+    if (confirm_write == "yes") or (confirm_write == "y"):
+        write_commandline_args(new_commandline_args)
 
 # Checks if we should skip loading the startup checks
 skip_startup_checks = startup_skip_check()