From 5e461e9b6b707d80716e27f6d01106ac39dafe4b Mon Sep 17 00:00:00 2001 From: Marc-Andre Ferland Date: Tue, 18 Oct 2022 13:21:15 -0400 Subject: [PATCH] Fixed is_alive with render_threads that can update the device name after starting. --- ui/sd_internal/task_manager.py | 14 ++++++++++-- ui/server.py | 40 +++++++++++++++++++++------------- 2 files changed, 37 insertions(+), 17 deletions(-) diff --git a/ui/sd_internal/task_manager.py b/ui/sd_internal/task_manager.py index 52a6b183..917387fe 100644 --- a/ui/sd_internal/task_manager.py +++ b/ui/sd_internal/task_manager.py @@ -3,7 +3,7 @@ import traceback TASK_TTL = 15 * 60 # Discard last session's task timeout -import queue, threading, time +import queue, threading, time, weakref from typing import Any, Generator, Hashable, Optional, Union from pydantic import BaseModel @@ -165,6 +165,7 @@ current_model_path = None tasks_queue = [] task_cache = TaskCache() default_model_to_load = None +weak_thread_data = weakref.WeakKeyDictionary() def preload_model(file_path=None): global current_state, current_state_error, current_model_path @@ -189,11 +190,17 @@ def preload_model(file_path=None): def thread_render(device): global current_state, current_state_error, current_model_path from . import runtime + weak_thread_data[threading.current_thread()] = { + 'device': device + } try: runtime.device_init(device) except: print(traceback.format_exc()) return + weak_thread_data[threading.current_thread()] = { + 'device': runtime.thread_data.device + } preload_model() current_state = ServerStates.Online while True: @@ -308,8 +315,11 @@ def is_alive(name=None): nbr_alive = 0 try: for rthread in render_threads: - thread_name = rthread.name[len(THREAD_NAME_PREFIX):].lower() if name is not None: + weak_data = weak_thread_data.get(rthread) + if weak_data is None or weak_data['device'] is None: + continue + thread_name = str(weak_data['device']).lower() if is_first_cuda_device(name): if not is_first_cuda_device(thread_name): continue diff --git a/ui/server.py b/ui/server.py index c913364f..acb5719b 100644 --- a/ui/server.py +++ b/ui/server.py @@ -26,6 +26,7 @@ APP_CONFIG_DEFAULT_MODELS = [ 'sd-v1-4', # Default fallback. ] +import asyncio from fastapi import FastAPI, HTTPException from fastapi.staticfiles import StaticFiles from starlette.responses import FileResponse, JSONResponse, StreamingResponse @@ -36,6 +37,7 @@ from typing import Any, Generator, Hashable, List, Optional, Union from sd_internal import Request, Response, task_manager +LOOP = asyncio.get_event_loop() app = FastAPI() modifiers_cache = None @@ -348,21 +350,29 @@ if 'render_devices' in config: # Start a new thread for each device. for device in config['render_devices']: task_manager.start_render_thread(device) -allow_cpu = False -if task_manager.is_alive() <= 0: # No running devices, apply defaults. - # Select best device GPU device using free memory if more than one device. - task_manager.start_render_thread('auto') - allow_cpu = True - -# Allow CPU to be used for renders if not already enabled in current config. -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') +async def check_status(): + device_count = 0 + for i in range(10): # Wait for devices to register and/or change names. + new_count = task_manager.is_alive() + if device_count != new_count: + device_count = new_count + await asyncio.sleep(3) + else: + break; + allow_cpu = False + if task_manager.is_alive() <= 0: # No running devices, apply defaults. + # Select best device GPU device using free memory if more than one device. + task_manager.start_render_thread('auto') + allow_cpu = True + # Allow CPU to be used for renders if not already enabled in current config. + if task_manager.is_alive('cpu') <= 0 and allow_cpu: + task_manager.start_render_thread('cpu') + if not task_manager.is_alive(0) <= 0: + 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') +LOOP.create_task(check_status()) # start the browser ui import webbrowser; webbrowser.open('http://localhost:9000') \ No newline at end of file