authentication and locking improvements

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

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

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

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],

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

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

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":

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 = []

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 = []

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