From f786aa3caa11ac987b5241be65890b185499361a Mon Sep 17 00:00:00 2001
From: erew123 <35898566+erew123@users.noreply.github.com>
Date: Tue, 5 Dec 2023 05:16:16 +0000
Subject: [PATCH] Clean-up Ctrl+C Shutdown (#4802)

---
 one_click.py |  8 ++++++++
 server.py    | 12 ++++++++++++
 2 files changed, 20 insertions(+)

diff --git a/one_click.py b/one_click.py
index 6febcba..1f91f6a 100644
--- a/one_click.py
+++ b/one_click.py
@@ -4,6 +4,7 @@ import hashlib
 import os
 import platform
 import re
+import signal
 import site
 import subprocess
 import sys
@@ -27,6 +28,13 @@ else:
 flags = f"{' '.join([flag for flag in sys.argv[1:] if flag != '--update'])} {CMD_FLAGS}"
 
 
+def signal_handler(sig, frame):
+    sys.exit(0)
+
+
+signal.signal(signal.SIGINT, signal_handler)
+
+
 def is_linux():
     return sys.platform.startswith("linux")
 
diff --git a/server.py b/server.py
index ae4eceb..d28ca80 100644
--- a/server.py
+++ b/server.py
@@ -21,6 +21,7 @@ matplotlib.use('Agg')  # This fixes LaTeX rendering on some systems
 
 import json
 import os
+import signal
 import sys
 import time
 from functools import partial
@@ -55,6 +56,17 @@ from modules.models_settings import (
 from modules.utils import gradio
 
 
+def signal_handler(sig, frame):
+    logger.info(f"Received Ctrl+C. Shutting down Text generation web UI gracefully")
+    if 'interface' in shared.gradio:
+        shared.gradio['interface'].close()
+
+    sys.exit(0)
+
+
+signal.signal(signal.SIGINT, signal_handler)
+
+
 def create_interface():
 
     title = 'Text generation web UI'