From 0c8825b2bec3a68836eacf37718306c9c78554a0 Mon Sep 17 00:00:00 2001
From: rucadi <ruben.cano96@gmail.com>
Date: Fri, 16 Dec 2022 18:31:20 +0100
Subject: [PATCH] Add a callback called before reloading the server

---
 modules/script_callbacks.py | 13 ++++++++++++-
 webui.py                    |  1 +
 2 files changed, 13 insertions(+), 1 deletion(-)

diff --git a/modules/script_callbacks.py b/modules/script_callbacks.py
index 715e1830..b646b0f9 100644
--- a/modules/script_callbacks.py
+++ b/modules/script_callbacks.py
@@ -64,6 +64,7 @@ callback_map = dict(
     callbacks_before_component=[],
     callbacks_after_component=[],
     callbacks_on_polling=[],
+    callbacks_on_reload=[],
 )
 
 
@@ -71,7 +72,6 @@ def clear_callbacks():
     for callback_list in callback_map.values():
         callback_list.clear()
 
-
 def app_started_callback(demo: Optional[Blocks], app: FastAPI):
     for c in callback_map['callbacks_app_started']:
         try:
@@ -86,6 +86,14 @@ def app_polling_callback(demo: Optional[Blocks], app: FastAPI):
         except Exception:
             report_exception(c, 'callbacks_on_polling')
 
+def app_reload_callback(demo: Optional[Blocks], app: FastAPI):
+    for c in callback_map['callbacks_on_reload']:
+        try:
+            c.callback()
+        except Exception:
+            report_exception(c, 'callbacks_on_reload')
+
+
 def model_loaded_callback(sd_model):
     for c in callback_map['callbacks_model_loaded']:
         try:
@@ -195,6 +203,9 @@ def on_polling(callback):
     """register a function to be called on each polling of the server."""
     add_callback(callback_map['callbacks_on_polling'], callback)
 
+def on_before_reload(callback):
+    """register a function to be called just before the server reloads."""
+    add_callback(callback_map['callbacks_on_reload'], callback)
 
 def on_model_loaded(callback):
     """register a function to be called when the stable diffusion model is created; the model is
diff --git a/webui.py b/webui.py
index 5f7a53f9..dda34249 100644
--- a/webui.py
+++ b/webui.py
@@ -108,6 +108,7 @@ def wait_on_server(demo=None):
         time.sleep(0.5)
         modules.script_callbacks.app_polling_callback(None, demo)
         if shared.state.need_restart:
+            modules.script_callbacks.app_reload_callback(None, demo)
             shared.state.need_restart = False
             time.sleep(0.5)
             demo.close()