implement global and api exception handler

automatic.sh

@@ -140,7 +140,7 @@ if [ "$MODE" == install ]; then
 fi
 
 if [ "$MODE" == clean ]; then
-  CMD="--disable-opt-split-attention --disable-console-progressbars --api"
+  CMD="--disable-opt-split-attention"
   "$PYTHON" launch.py $CMD
   exit 0
 fi
@@ -153,8 +153,8 @@ if [ $MODE == optimized ]; then
   CMD="$CMD"
 fi
 
-exec accelerate launch --no_python --quiet --num_cpu_threads_per_process=6 "$PYTHON" $CMD
-# exec "$PYTHON" $CMD
+# exec accelerate launch --no_python --quiet --num_cpu_threads_per_process=6 "$PYTHON" $CMD
+exec "$PYTHON" $CMD
 
 # export LD_PRELOAD=libtcmalloc.so
 # TORCH_CUDA_ARCH_LIST="8.6"

launch.py

@@ -372,7 +372,16 @@ def tests(test_dir):
 
 
 def start():
-    print(f"Launching {'API server' if '--nowebui' in sys.argv else 'Web UI'} with arguments: {' '.join(sys.argv[1:])}")
+    print(f"Launching server with arguments: {' '.join(sys.argv[1:])}")
+
+    try:
+        from rich.traceback import install
+        from rich.console import Console
+        console = Console()
+        install(show_locals=True, max_frames=2, extra_lines=1, word_wrap=False, width=min([console.width, 200]))
+    except:
+        pass # if rich is not installed do nothing
+
     import webui
     if '--nowebui' in sys.argv:
         webui.api_only()

modules/api/api.py

@@ -6,8 +6,11 @@ import uvicorn
 from threading import Lock
 from io import BytesIO
 from gradio.processing_utils import decode_base64_to_file
-from fastapi import APIRouter, Depends, FastAPI, HTTPException, Request, Response
+from fastapi import APIRouter, Depends, FastAPI, Request, Response
 from fastapi.security import HTTPBasic, HTTPBasicCredentials
+from fastapi.exceptions import HTTPException
+from fastapi.responses import JSONResponse
+from fastapi.encoders import jsonable_encoder
 from secrets import compare_digest
 
 import modules.shared as shared
@@ -90,6 +93,16 @@ def encode_pil_to_base64(image):
     return base64.b64encode(bytes_data)
 
 def api_middleware(app: FastAPI):
+    rich_available = True
+    try:
+        import anyio # importing just so it can be placed on silent list
+        import starlette # importing just so it can be placed on silent list
+        from rich.console import Console
+        console = Console()
+    except:
+        import traceback
+        rich_available = False
+
     @app.middleware("http")
     async def log_and_time(req: Request, call_next):
         ts = time.time()
@@ -110,6 +123,36 @@ def api_middleware(app: FastAPI):
             ))
         return res
 
+    def handle_exception(request: Request, e: Exception):
+        err = {
+            "error": type(e).__name__,
+            "detail": vars(e).get('detail', ''),
+            "body": vars(e).get('body', ''),
+            "errors": str(e),
+        }
+        print(f"API error: {request.method}: {request.url} {err}")
+        if not isinstance(e, HTTPException): # do not print backtrace on known httpexceptions
+            if rich_available:
+                console.print_exception(show_locals=True, max_frames=2, extra_lines=1, suppress=[anyio, starlette], word_wrap=False, width=min([console.width, 200]))
+            else:
+                traceback.print_exc()
+        return JSONResponse(status_code=vars(e).get('status_code', 500), content=jsonable_encoder(err))
+
+    @app.middleware("http")
+    async def exception_handling(request: Request, call_next):
+        try:
+            return await call_next(request)
+        except Exception as e:
+            return handle_exception(request, e)
+
+    @app.exception_handler(Exception)
+    async def fastapi_exception_handler(request: Request, e: Exception):
+        return handle_exception(request, e)
+
+    @app.exception_handler(HTTPException)
+    async def http_exception_handler(request: Request, e: HTTPException):
+        return handle_exception(request, e)
+
 
 class Api:
     def __init__(self, app: FastAPI, queue_lock: Lock):

requirements.txt

@@ -34,6 +34,7 @@ pytorch_lightning
 realesrgan
 requests
 resize-right
+rich
 safetensors
 scikit-image
 timm

ui-config.json

@@ -1431,5 +1431,15 @@
     "img2img/Threshold B/maximum": 1024,
     "img2img/Threshold B/step": 1,
     "img2img/Resize Mode/visible": true,
-    "img2img/Resize Mode/value": "Scale to Fit (Inner Fit)"
+    "img2img/Resize Mode/value": "Scale to Fit (Inner Fit)",
+    "customscript/seed_travel.py/txt2img/Desired min SSIM threshold (% of threshold)/visible": true,
+    "customscript/seed_travel.py/txt2img/Desired min SSIM threshold (% of threshold)/value": 75,
+    "customscript/seed_travel.py/txt2img/Desired min SSIM threshold (% of threshold)/minimum": 0,
+    "customscript/seed_travel.py/txt2img/Desired min SSIM threshold (% of threshold)/maximum": 100,
+    "customscript/seed_travel.py/txt2img/Desired min SSIM threshold (% of threshold)/step": 1,
+    "customscript/seed_travel.py/img2img/Desired min SSIM threshold (% of threshold)/visible": true,
+    "customscript/seed_travel.py/img2img/Desired min SSIM threshold (% of threshold)/value": 75,
+    "customscript/seed_travel.py/img2img/Desired min SSIM threshold (% of threshold)/minimum": 0,
+    "customscript/seed_travel.py/img2img/Desired min SSIM threshold (% of threshold)/maximum": 100,
+    "customscript/seed_travel.py/img2img/Desired min SSIM threshold (% of threshold)/step": 1
 }

webui.py

@@ -179,7 +179,7 @@ def initialize():
 
     # make the program just exit at ctrl+c without waiting for anything
     def sigint_handler(sig, frame):
-        print(f'Interrupted with signal {sig} in {frame}')
+        print('Exiting')
         os._exit(0)
 
     signal.signal(signal.SIGINT, sigint_handler)
@@ -262,8 +262,11 @@ def webui():
             debug=cmd_opts.gradio_debug,
             auth=[tuple(cred.split(':')) for cred in gradio_auth_creds] if gradio_auth_creds else None,
             inbrowser=cmd_opts.autolaunch,
-            prevent_thread_lock=True
+            prevent_thread_lock=True,
         )
+        for dep in shared.demo.dependencies:
+            dep['show_progress'] = False  # disable gradio css animation on component update
+
         # app is instance of FastAPI server
         # shared.demo.server is instance of gradio class which inherits from uvicorn.Server
         # shared.demo.config is instance of uvicorn.Config
@@ -293,10 +296,7 @@ def webui():
 
         print(f"Startup time: {startup_timer.summary()}.")
 
-        try:
-            wait_on_server(shared.demo)
-        except KeyboardInterrupt as e:
-            pass
+        wait_on_server(shared.demo)
         print('Restarting UI...')
 
         startup_timer.reset()