From 57fed943377bf17506863fafc17e8662e770b53b Mon Sep 17 00:00:00 2001
From: Sean Sube <seansube@gmail.com>
Date: Sun, 5 Mar 2023 21:07:06 -0600
Subject: [PATCH] fix(api): exit worker on memory allocation errors

---
 api/onnx_web/worker/worker.py   | 14 +++++++++-----
 api/scripts/test-memory-leak.sh |  2 +-
 2 files changed, 10 insertions(+), 6 deletions(-)

diff --git a/api/onnx_web/worker/worker.py b/api/onnx_web/worker/worker.py
index dc5d36c0..e53a060c 100644
--- a/api/onnx_web/worker/worker.py
+++ b/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__),
+                )
diff --git a/api/scripts/test-memory-leak.sh b/api/scripts/test-memory-leak.sh
index 6c576519..b8a0d4bb 100755
--- a/api/scripts/test-memory-leak.sh
+++ b/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' \