Add setup method to extensions

extensions/api/script.py

@@ -8,8 +8,6 @@ params = {
     'port': 5000,
 }
 
-server = None
-
 class Handler(BaseHTTPRequestHandler):
     def do_GET(self):
         if self.path == '/api/v1/model':
@@ -75,7 +73,6 @@ class Handler(BaseHTTPRequestHandler):
 
 
 def run_server():
-    global server
     server_addr = ('0.0.0.0' if shared.args.listen else '127.0.0.1', params['port'])
     server = ThreadingHTTPServer(server_addr, Handler)
     if shared.args.share: 
@@ -89,6 +86,5 @@ def run_server():
         print(f'Starting KoboldAI compatible api at http://{server_addr[0]}:{server_addr[1]}/api')
     server.serve_forever()
 
-def ui():
-    if server is None:
+def setup():
     Thread(target=run_server, daemon=True).start()

modules/extensions.py

@@ -5,6 +5,7 @@ import modules.shared as shared
 
 state = {}
 available_extensions = []
+setup_called = False
 
 def load_extensions():
     global state
@@ -44,6 +45,13 @@ def create_extensions_block():
                 if _id in shared.settings:
                     extension.params[param] = shared.settings[_id]
 
+    # Running setup function
+    if not setup_called:
+        for extension, name in iterator():
+            if hasattr(extension, "setup"):
+                extension.setup()
+        setup_called = True
+
     # Creating the extension ui elements
     if len(state) > 0:
         with gr.Box(elem_id="extensions"):