diff --git a/one_click.py b/one_click.py index 6febcbac..1f91f6ad 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 ae4eceb9..d28ca804 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'