From bd0b43f7e6dd002cf1c2558a76a3e44a1bb7241d Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E6=BA=90=E6=96=87=E9=9B=A8?=
 <41315874+fumiama@users.noreply.github.com>
Date: Tue, 4 Jun 2024 15:45:17 +0900
Subject: [PATCH] optimize(web): program cleanup on exit

---
 infer-web.py | 34 +++++++++++++++++++++-------------
 1 file changed, 21 insertions(+), 13 deletions(-)

diff --git a/infer-web.py b/infer-web.py
index def0515..21b0224 100644
--- a/infer-web.py
+++ b/infer-web.py
@@ -1638,16 +1638,24 @@ with gr.Blocks(title="RVC WebUI") as app:
             except:
                 gr.Markdown(traceback.format_exc())
 
-    try:
-        if config.iscolab:
-            app.queue(max_size=1022).launch(share=True, max_threads=511)
-        else:
-            app.queue(max_size=1022).launch(
-                max_threads=511,
-                server_name="0.0.0.0",
-                inbrowser=not config.noautoopen,
-                server_port=config.listen_port,
-                quiet=True,
-            )
-    except Exception as e:
-        logger.error(str(e))
+try:
+    import signal
+    def cleanup(signum, frame):
+        signame = signal.Signals(signum).name
+        print(f'Got signal {signame} ({signum})')
+        app.close()
+        sys.exit(0)
+    signal.signal(signal.SIGINT, cleanup)
+    signal.signal(signal.SIGTERM, cleanup)
+    if config.iscolab:
+        app.queue(max_size=1022).launch(share=True, max_threads=511)
+    else:
+        app.queue(max_size=1022).launch(
+            max_threads=511,
+            server_name="0.0.0.0",
+            inbrowser=not config.noautoopen,
+            server_port=config.listen_port,
+            quiet=True,
+        )
+except Exception as e:
+    logger.error(str(e))