From c92129ac632ae8f245726b1c8dbb743424cfc411 Mon Sep 17 00:00:00 2001 From: Marc-Andre Ferland Date: Mon, 17 Oct 2022 02:27:30 -0400 Subject: [PATCH] Improved detection of missing cuda:0 and added warning to console about how to fix. --- ui/sd_internal/runtime.py | 2 +- ui/sd_internal/task_manager.py | 9 ++++++--- ui/server.py | 29 ++++++++++++++++++++++++----- 3 files changed, 31 insertions(+), 9 deletions(-) diff --git a/ui/sd_internal/runtime.py b/ui/sd_internal/runtime.py index 5756aaf2..0a667417 100644 --- a/ui/sd_internal/runtime.py +++ b/ui/sd_internal/runtime.py @@ -130,7 +130,7 @@ def device_init(device_selection=None): device_selection = device_selection.lower() if device_selection.startswith('gpu:'): device_selection = int(device_selection[4:]) - if device_selection != 'cuda' and device_selection != 'current' and device_selection != 'gpu': + if device_selection != 'current' and device_selection != 'gpu': if device_select(device_selection): if isinstance(device_selection, int): print(f'Setting GPU:{device_selection} as active') diff --git a/ui/sd_internal/task_manager.py b/ui/sd_internal/task_manager.py index a47e4263..22146c4d 100644 --- a/ui/sd_internal/task_manager.py +++ b/ui/sd_internal/task_manager.py @@ -299,16 +299,19 @@ def thread_render(device): print(f'Session {task.request.session_id} task {id(task)} completed.') current_state = ServerStates.Online +def is_first_cuda_device(device): + from . import runtime # When calling runtime from outside thread_render DO NOT USE thread specific attributes or functions. + return runtime.is_first_cuda_device(device) + def is_alive(name=None): - from . import runtime # When calling runtime from here DO NOT USE thread specific attributes or functions. if not manager_lock.acquire(blocking=True, timeout=LOCK_TIMEOUT): raise Exception('is_alive' + ERR_LOCK_FAILED) nbr_alive = 0 try: for rthread in render_threads: thread_name = rthread.name[len(THREAD_NAME_PREFIX):].lower() if name is not None: - if runtime.is_first_cuda_device(name): - if not runtime.is_first_cuda_device(thread_name): + if is_first_cuda_device(name): + if not is_first_cuda_device(thread_name): continue elif thread_name != name: continue diff --git a/ui/server.py b/ui/server.py index 55786906..8c64c59b 100644 --- a/ui/server.py +++ b/ui/server.py @@ -78,16 +78,27 @@ def setConfig(config): print(traceback.format_exc()) if 'render_devices' in config: - gpu_devices = filter(lambda dev: dev.startswith('GPU:'), config['render_devices']) + gpu_devices = filter(lambda dev: dev.lower().startswith('gpu') or dev.lower().startswith('cuda'), config['render_devices']) else: gpu_devices = [] + has_first_cuda_device = False + for device in gpu_devices: + if not task_manager.is_first_cuda_device(device): continue + has_first_cuda_device = True + break + if len(gpu_devices) > 0 and not has_first_cuda_device: + print('WARNING: GFPGANer only works on CPU or GPU:0, use CUDA_VISIBLE_DEVICES if GFPGANer is needed on a specific GPU.') + print('Using CUDA_VISIBLE_DEVICES will remap the selected devices starting at GPU:0 fixing GFPGANer') + try: # config.bat config_bat = [ f"@set update_branch={config['update_branch']}" ] - if len(gpu_devices) > 0: - config_sh.append(f"@set CUDA_VISIBLE_DEVICES={','.join(gpu_devices)}") + if len(gpu_devices) > 0 and not has_first_cuda_device: + config_sh.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') config_bat_path = os.path.join(CONFIG_DIR, 'config.bat') with open(config_bat_path, 'w') as f: f.write(f.write('\r\n'.join(config_bat))) @@ -99,8 +110,10 @@ def setConfig(config): '#!/bin/bash' f"export update_branch={config['update_branch']}" ] - if len(gpu_devices) > 0: - config_sh.append(f"CUDA_VISIBLE_DEVICES={','.join(gpu_devices)}") + if len(gpu_devices) > 0 and not has_first_cuda_device: + 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') config_sh_path = os.path.join(CONFIG_DIR, 'config.sh') with open(config_sh_path, 'w') as f: f.write('\n'.join(config_sh)) @@ -345,5 +358,11 @@ if task_manager.is_alive() <= 0: # No running devices, apply defaults. if task_manager.is_alive('cpu') <= 0 and allow_cpu: task_manager.start_render_thread('cpu') +if task_manager.is_alive(0) <= 0: # Missing cuda:0, warn the user. + print('WARNING: GFPGANer only works on CPU or GPU:0, use CUDA_VISIBLE_DEVICES if GFPGANer is needed on a specific GPU.') + print('Using CUDA_VISIBLE_DEVICES will remap the selected devices starting at GPU:0 fixing GFPGANer') + print('Add the line "@set CUDA_VISIBLE_DEVICES=N" where N is the GPUs to use to config.bat') + print('Add the line "CUDA_VISIBLE_DEVICES=N" where N is the GPUs to use to config.sh') + # start the browser ui import webbrowser; webbrowser.open('http://localhost:9000') \ No newline at end of file