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()