fix(api): exit worker on memory allocation errors

api/onnx_web/worker/worker.py

@@ -39,10 +39,14 @@ def worker_main(context: WorkerContext, server: ServerContext):
             logger.info("worker got keyboard interrupt")
             exit(0)
         except ValueError as e:
-            logger.info("value error in worker: %s", e)
+            logger.info("value error in worker, exiting: %s", e)
             exit(1)
         except Exception as e:
-            logger.error(
-                "error while running job: %s",
-                format_exception(type(e), e, e.__traceback__),
-            )
+            if "Failed to allocate memory" in str(e):
+                logger.error("detected out-of-memory error, exiting: %s", e)
+                exit(2)
+            else:
+                logger.error(
+                    "error while running job: %s",
+                    format_exception(type(e), e, e.__traceback__),
+                )

api/scripts/test-memory-leak.sh

@@ -6,7 +6,7 @@ do
   curl "${test_host}/api/txt2img?"\
 'cfg=16.00&steps=30&scheduler=ddim&seed=-1&'\
 'prompt=an+astronaut+eating+a+hamburger&negativePrompt=&'\
-'model=stable-diffusion-onnx-v1-5&platform=any&'\
+'model=stable-diffusion-onnx-v1-5&platform=cuda&'\
 'upscaling=upscaling-real-esrgan-x2-plus&correction=correction-codeformer&'\
 'lpw=false&width=512&height=512&upscaleOrder=correction-both' \
     -X 'POST' \