From 7f151cbebad81b6e28bcbf21af6008ad21997980 Mon Sep 17 00:00:00 2001 From: JeLuF Date: Mon, 31 Oct 2022 00:48:18 +0100 Subject: [PATCH] Copy CUDA_VISIBLE_DEVICES to config.*, it it has been set Don't delete CUDA_VISIBLE_DEVICES settings when generating a new config file --- ui/server.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/ui/server.py b/ui/server.py index c722894f..0df08785 100644 --- a/ui/server.py +++ b/ui/server.py @@ -111,6 +111,8 @@ def setConfig(config): config_bat.append('::Set the devices visible inside SD-UI here') config_bat.append(f"::@set CUDA_VISIBLE_DEVICES={','.join(gpu_devices)}") # Needs better detection for edge cases, add as a comment for now. print('Add the line "@set CUDA_VISIBLE_DEVICES=N" where N is the GPUs to use to config.bat') + if os.getenv('CUDA_VISIBLE_DEVICES') is not None: + config_bat.append(f"@set CUDA_VISIBLE_DEVICES={os.getenv('CUDA_VISIBLE_DEVICES')}") config_bat_path = os.path.join(CONFIG_DIR, 'config.bat') with open(config_bat_path, 'w', encoding='utf-8') as f: f.write('\r\n'.join(config_bat)) @@ -126,6 +128,8 @@ def setConfig(config): config_sh.append('#Set the devices visible inside SD-UI here') config_sh.append(f"#CUDA_VISIBLE_DEVICES={','.join(gpu_devices)}") # Needs better detection for edge cases, add as a comment for now. print('Add the line "CUDA_VISIBLE_DEVICES=N" where N is the GPUs to use to config.sh') + if os.getenv('CUDA_VISIBLE_DEVICES') is not None: + config_sh.append(f"export CUDA_VISIBLE_DEVICES=\"{os.getenv('CUDA_VISIBLE_DEVICES')}\"") config_sh_path = os.path.join(CONFIG_DIR, 'config.sh') with open(config_sh_path, 'w', encoding='utf-8') as f: f.write('\n'.join(config_sh)) @@ -430,4 +434,4 @@ if display_warning or task_manager.is_alive(0) <= 0: del display_warning # start the browser ui -import webbrowser; webbrowser.open('http://localhost:9000') \ No newline at end of file +import webbrowser; webbrowser.open('http://localhost:9000')