From 12cf688cbe3ece45b22e95b979d6e0010597f5a9 Mon Sep 17 00:00:00 2001 From: Vladimir Mandic Date: Thu, 16 Nov 2023 15:26:34 -0500 Subject: [PATCH] authentication and locking improvements --- CHANGELOG.md | 3 ++ modules/api/api.py | 9 +++-- modules/hashes.py | 2 +- modules/middleware.py | 9 +++-- modules/sd_hijack.py | 7 ++-- modules/sd_models.py | 2 +- modules/sd_models_compile.py | 1 + modules/shared.py | 72 ++++++++++++++++++++++-------------- modules/upscaler.py | 1 - webui.py | 2 + 10 files changed, 68 insertions(+), 40 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 633d94f58..c4d97ffee 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -39,6 +39,9 @@ - Add option to create only subimages in XYZ grid, thanks @midcoastal - Support custom upscalers in subfolders - Add additional image info when loading image in process tab + - Better file locking when sharing config and/or models between multiple instances + - Handle custom API endpoints when using auth + - Show logged in user in log when accessing via UI and/or API - Support `--ckpt none` to skip loading a model - **Fixes** - Fix `params.txt` saved before actual image diff --git a/modules/api/api.py b/modules/api/api.py index 09ee1e116..ef0825926 100644 --- a/modules/api/api.py +++ b/modules/api/api.py @@ -145,7 +145,7 @@ class Api: self.add_api_route("/sdapi/v1/reload-checkpoint", self.reloadapi, methods=["POST"]) self.add_api_route("/sdapi/v1/scripts", self.get_scripts_list, methods=["GET"], response_model=models.ScriptsList) self.add_api_route("/sdapi/v1/script-info", self.get_script_info, methods=["GET"], response_model=List[models.ScriptInfo]) - self.add_api_route("/sdapi/v1/log", self.get_log_buffer, methods=["GET"], response_model=List) # bypass auth + self.add_api_route("/sdapi/v1/log", self.get_log_buffer, methods=["GET"], response_model=List) self.add_api_route("/sdapi/v1/start", self.session_start, methods=["GET"]) self.add_api_route("/sdapi/v1/motd", self.get_motd, methods=["GET"], response_model=str) self.add_api_route("/sdapi/v1/extra-networks", self.get_extra_networks, methods=["GET"], response_model=List[models.ExtraNetworkItem]) @@ -153,11 +153,12 @@ class Api: self.default_script_arg_img2img = [] def add_api_route(self, path: str, endpoint, **kwargs): - if shared.cmd_opts.auth or shared.cmd_opts.auth_file: + if (shared.cmd_opts.auth or shared.cmd_opts.auth_file) and shared.cmd_opts.api_only: return self.app.add_api_route(path, endpoint, dependencies=[Depends(self.auth)], **kwargs) return self.app.add_api_route(path, endpoint, **kwargs) def auth(self, credentials: HTTPBasicCredentials = Depends(HTTPBasic())): + # this is only needed for api-only since otherwise auth is handled in gradio/routes.py if credentials.username in self.credentials: if compare_digest(credentials.password, self.credentials[credentials.username]): return True @@ -170,7 +171,9 @@ class Api: return lines def session_start(self, req: Request, agent: Optional[str] = None): - shared.log.info(f'Browser session: client={req.client.host} agent={agent}') + token = req.cookies.get("access-token") or req.cookies.get("access-token-unsecure") + user = self.app.tokens.get(token) + shared.log.info(f'Browser session: user={user} client={req.client.host} agent={agent}') return {} def get_motd(self): diff --git a/modules/hashes.py b/modules/hashes.py index e4d13f2e3..892368d3c 100644 --- a/modules/hashes.py +++ b/modules/hashes.py @@ -16,7 +16,7 @@ def dump_cache(): def cache(subsection): global cache_data # pylint: disable=global-statement if cache_data is None: - cache_data = {} if not os.path.isfile(cache_filename) else shared.readfile(cache_filename) + cache_data = {} if not os.path.isfile(cache_filename) else shared.readfile(cache_filename, lock=True) s = cache_data.get(subsection, {}) cache_data[subsection] = s return s diff --git a/modules/middleware.py b/modules/middleware.py index e188b4a3d..b119a9c4f 100644 --- a/modules/middleware.py +++ b/modules/middleware.py @@ -1,6 +1,5 @@ import ssl import time -import datetime import logging from asyncio.exceptions import CancelledError import anyio @@ -14,6 +13,7 @@ from fastapi.encoders import jsonable_encoder from installer import log import modules.errors as errors + errors.install() @@ -42,11 +42,12 @@ def setup_middleware(app: FastAPI, cmd_opts): duration = str(round(time.time() - ts, 4)) res.headers["X-Process-Time"] = duration endpoint = req.scope.get('path', 'err') + token = req.cookies.get("access-token") or req.cookies.get("access-token-unsecure") if (cmd_opts.api_log or cmd_opts.api_only) and endpoint.startswith('/sdapi'): - if endpoint.endswith('/sdapi/v1/log'): + if '/sdapi/v1/log' in endpoint: return res - log.info('API {t} {code} {prot}/{ver} {method} {endpoint} {cli} {duration}'.format( # pylint: disable=consider-using-f-string, logging-format-interpolation - t = datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S.%f"), + log.info('API {user} {code} {prot}/{ver} {method} {endpoint} {cli} {duration}'.format( # pylint: disable=consider-using-f-string, logging-format-interpolation + user = app.tokens.get(token), code = res.status_code, ver = req.scope.get('http_version', '0.0'), cli = req.scope.get('client', ('0:0.0.0', 0))[0], diff --git a/modules/sd_hijack.py b/modules/sd_hijack.py index 36c8846e1..dd5c68f7c 100644 --- a/modules/sd_hijack.py +++ b/modules/sd_hijack.py @@ -185,7 +185,7 @@ class StableDiffusionModelHijack: import torch._dynamo # pylint: disable=unused-import,redefined-outer-name if shared.opts.cuda_compile_backend == "openvino_fx": torch._dynamo.reset() # pylint: disable=protected-access - from modules.intel.openvino import openvino_fx, openvino_clear_caches # pylint: disable=unused-import + from modules.intel.openvino import openvino_fx, openvino_clear_caches # pylint: disable=unused-import, no-name-in-module openvino_clear_caches() torch._dynamo.eval_frame.check_if_dynamo_supported = lambda: True # pylint: disable=protected-access log_level = logging.WARNING if opts.cuda_compile_verbose else logging.CRITICAL # pylint: disable=protected-access @@ -199,11 +199,12 @@ class StableDiffusionModelHijack: hidet.torch.dynamo_config.use_tensor_core(True) hidet.torch.dynamo_config.search_space(2) m.model = torch.compile(m.model, mode=opts.cuda_compile_mode, backend=opts.cuda_compile_backend, fullgraph=opts.cuda_compile_fullgraph, dynamic=False) - from installer import setup_logging - setup_logging() shared.log.info("Model complilation done.") except Exception as err: shared.log.warning(f"Model compile not supported: {err}") + finally: + from installer import setup_logging + setup_logging() self.optimization_method = apply_optimizations() self.clip = m.cond_stage_model diff --git a/modules/sd_models.py b/modules/sd_models.py index 145c76766..fc96b00c5 100644 --- a/modules/sd_models.py +++ b/modules/sd_models.py @@ -327,7 +327,7 @@ def read_metadata_from_safetensors(filename): if not os.path.isfile(sd_metadata_file): sd_metadata = {} else: - sd_metadata = shared.readfile(sd_metadata_file) + sd_metadata = shared.readfile(sd_metadata_file, lock=True) res = sd_metadata.get(filename, None) if res is not None: return res diff --git a/modules/sd_models_compile.py b/modules/sd_models_compile.py index 711245754..fe9f355d1 100644 --- a/modules/sd_models_compile.py +++ b/modules/sd_models_compile.py @@ -93,6 +93,7 @@ def compile_torch(sd_model): try: import torch._dynamo # pylint: disable=unused-import,redefined-outer-name torch._dynamo.reset() # pylint: disable=protected-access + shared.log.debug(f"Model compile available backends: {torch._dynamo.list_backends()}") # pylint: disable=protected-access if shared.opts.ipex_optimize: optimize_ipex(sd_model) if shared.opts.cuda_compile_backend == "openvino_fx": diff --git a/modules/shared.py b/modules/shared.py index 17ce2bd54..2ab35e858 100644 --- a/modules/shared.py +++ b/modules/shared.py @@ -181,51 +181,69 @@ def temp_disable_extensions(): return disabled -def readfile(filename, silent=False): +def readfile(filename, silent=False, lock=False): data = {} + lock_file = None + locked = False try: if not os.path.exists(filename): return {} - with fasteners.InterProcessLock(f"{filename}.lock"): - with open(filename, "r", encoding="utf8") as file: - data = json.load(file) - if type(data) is str: - data = json.loads(data) - if not silent: - log.debug(f'Read: file="{filename}" len={len(data)}') + if lock: + lock_file = fasteners.InterProcessReaderWriterLock(f"{filename}.lock", logger=log) + locked = lock_file.acquire_read_lock(blocking=True, timeout=3) + with open(filename, "r", encoding="utf8") as file: + data = json.load(file) + if type(data) is str: + data = json.loads(data) + if not silent: + log.debug(f'Read: file="{filename}" json={len(data)} bytes={os.path.getsize(filename)}') except Exception as e: if not silent: log.error(f'Reading failed: {filename} {e}') return {} + finally: + if lock_file is not None: + lock_file.release_read_lock() + if locked and os.path.exists(f"{filename}.lock"): + os.remove(f"{filename}.lock") return data def writefile(data, filename, mode='w', silent=False): + lock = None + locked = False + def default(obj): log.error(f"Saving: {filename} not a valid object: {obj}") return str(obj) try: - with fasteners.InterProcessLock(f"{filename}.lock"): - # skipkeys=True, ensure_ascii=True, check_circular=True, allow_nan=True - if type(data) == dict: - output = json.dumps(data, indent=2, default=default) - elif type(data) == list: - output = json.dumps(data, indent=2, default=default) - elif isinstance(data, object): - simple = {} - for k in data.__dict__: - if data.__dict__[k] is not None: - simple[k] = data.__dict__[k] - output = json.dumps(simple, indent=2, default=default) - else: - raise ValueError('not a valid object') - if not silent: - log.debug(f'Save: file="{filename}" len={len(output)}') - with open(filename, mode, encoding="utf8") as file: - file.write(output) + # skipkeys=True, ensure_ascii=True, check_circular=True, allow_nan=True + if type(data) == dict: + output = json.dumps(data, indent=2, default=default) + elif type(data) == list: + output = json.dumps(data, indent=2, default=default) + elif isinstance(data, object): + simple = {} + for k in data.__dict__: + if data.__dict__[k] is not None: + simple[k] = data.__dict__[k] + output = json.dumps(simple, indent=2, default=default) + else: + raise ValueError('not a valid object') + lock = fasteners.InterProcessReaderWriterLock(f"{filename}.lock", logger=log) + locked = lock.acquire_write_lock(blocking=True, timeout=3) + with open(filename, mode, encoding="utf8") as file: + file.write(output) + if not silent: + log.debug(f'Save: file="{filename}" json={len(data)} bytes={len(output)}') except Exception as e: log.error(f'Saving failed: {filename} {e}') + finally: + if lock is not None: + lock.release_read_lock() + if locked and os.path.exists(f"{filename}.lock"): + os.remove(f"{filename}.lock") if devices.backend == "cpu": @@ -699,7 +717,7 @@ class Options: log.debug(f'Created default config: {filename}') self.save(filename) return - self.data = readfile(filename) + self.data = readfile(filename, lock=True) if self.data.get('quicksettings') is not None and self.data.get('quicksettings_list') is None: self.data['quicksettings_list'] = [i.strip() for i in self.data.get('quicksettings').split(',')] unknown_settings = [] diff --git a/modules/upscaler.py b/modules/upscaler.py index e9ec14a54..9228bc658 100644 --- a/modules/upscaler.py +++ b/modules/upscaler.py @@ -66,7 +66,6 @@ class Upscaler: scalers.append(scaler) loaded.append(file_name) modules.shared.log.debug(f'Upscaler type={self.name} folder="{folder}" model="{model_name}" path="{file_name}"') - print(f'Upscaler type={self.name} folder="{folder}" model="{model_name}" path="{file_name}"') def find_scalers(self): scalers = [] diff --git a/webui.py b/webui.py index 5b0692ed2..103d78ec7 100644 --- a/webui.py +++ b/webui.py @@ -247,6 +247,8 @@ def start_ui(): with open(cmd_opts.auth_file, 'r', encoding="utf8") as file: for line in file.readlines(): gradio_auth_creds += [x.strip() for x in line.split(',') if x.strip()] + if len(gradio_auth_creds) > 0: + log.info(f'Authentication enabled: {gradio_auth_creds}') global local_url # pylint: disable=global-statement stdout = io.StringIO()