Merge branch 'master' into master

1 files changed, 65 insertions, 36 deletions

diff --git a/webui.py b/webui.py
index b8cccd54..fe0ce321 100644
--- a/webui.py
+++ b/webui.py
@@ -1,38 +1,37 @@
 import os

 import threading

-

-from modules import devices

-from modules.paths import script_path

+import time

+import importlib

 import signal

 import threading

-import modules.paths

+

+from fastapi.middleware.gzip import GZipMiddleware

+

+from modules.paths import script_path

+

+from modules import devices, sd_samplers

 import modules.codeformer_model as codeformer

-import modules.esrgan_model as esrgan

-import modules.bsrgan_model as bsrgan

 import modules.extras

 import modules.face_restoration

 import modules.gfpgan_model as gfpgan

 import modules.img2img

-import modules.ldsr_model as ldsr

+

 import modules.lowvram

-import modules.realesrgan_model as realesrgan

+import modules.paths

 import modules.scripts

 import modules.sd_hijack

 import modules.sd_models

 import modules.shared as shared

-import modules.swinir_model as swinir

 import modules.txt2img

+

 import modules.ui

+from modules import devices

 from modules import modelloader

 from modules.paths import script_path

 from modules.shared import cmd_opts

+import modules.hypernetworks.hypernetwork

+

 

-modelloader.cleanup_models()

-modules.sd_models.setup_model(cmd_opts.ckpt_dir)

-codeformer.setup_model(cmd_opts.codeformer_models_path)

-gfpgan.setup_model(cmd_opts.gfpgan_models_path)

-shared.face_restorers.append(modules.face_restoration.FaceRestoration())

-modelloader.load_upscalers()

 queue_lock = threading.Lock()

 

 

@@ -46,7 +45,7 @@ def wrap_queued_call(func):
     return f

 

 

-def wrap_gradio_gpu_call(func):

+def wrap_gradio_gpu_call(func, extra_outputs=None):

     def f(*args, **kwargs):

         devices.torch_gc()

 

@@ -57,7 +56,9 @@ def wrap_gradio_gpu_call(func):
         shared.state.current_latent = None

         shared.state.current_image = None

         shared.state.current_image_sampling_step = 0

+        shared.state.skipped = False

         shared.state.interrupted = False

+        shared.state.textinfo = None

 

         with queue_lock:

             res = func(*args, **kwargs)

@@ -69,16 +70,27 @@ def wrap_gradio_gpu_call(func):
 

         return res

 

-    return modules.ui.wrap_gradio_call(f)

+    return modules.ui.wrap_gradio_call(f, extra_outputs=extra_outputs)

 

+def initialize():

+    modelloader.cleanup_models()

+    modules.sd_models.setup_model()

+    codeformer.setup_model(cmd_opts.codeformer_models_path)

+    gfpgan.setup_model(cmd_opts.gfpgan_models_path)

+    shared.face_restorers.append(modules.face_restoration.FaceRestoration())

+    modelloader.load_upscalers()

 

-modules.scripts.load_scripts(os.path.join(script_path, "scripts"))

+    modules.scripts.load_scripts(os.path.join(script_path, "scripts"))

 

-shared.sd_model = modules.sd_models.load_model()

-shared.opts.onchange("sd_model_checkpoint", wrap_queued_call(lambda: modules.sd_models.reload_model_weights(shared.sd_model)))

+    shared.sd_model = modules.sd_models.load_model()

+    shared.opts.onchange("sd_model_checkpoint", wrap_queued_call(lambda: modules.sd_models.reload_model_weights(shared.sd_model)))

+    shared.opts.onchange("sd_hypernetwork", wrap_queued_call(lambda: modules.hypernetworks.hypernetwork.load_hypernetwork(shared.opts.sd_hypernetwork)))

+    shared.opts.onchange("sd_hypernetwork_strength", modules.hypernetworks.hypernetwork.apply_strength)

 

 

 def webui():

+    initialize()

+    

     # make the program just exit at ctrl+c without waiting for anything

     def sigint_handler(sig, frame):

         print(f'Interrupted with signal {sig} in {frame}')

@@ -86,22 +98,39 @@ def webui():
 

     signal.signal(signal.SIGINT, sigint_handler)

 

-    demo = modules.ui.create_ui(

-        txt2img=wrap_gradio_gpu_call(modules.txt2img.txt2img),

-        img2img=wrap_gradio_gpu_call(modules.img2img.img2img),

-        run_extras=wrap_gradio_gpu_call(modules.extras.run_extras),

-        run_pnginfo=modules.extras.run_pnginfo,

-        run_modelmerger=modules.extras.run_modelmerger

-    )

-

-    demo.launch(

-        share=cmd_opts.share,

-        server_name="0.0.0.0" if cmd_opts.listen else None,

-        server_port=cmd_opts.port,

-        debug=cmd_opts.gradio_debug,

-        auth=[tuple(cred.split(':')) for cred in cmd_opts.gradio_auth.strip('"').split(',')] if cmd_opts.gradio_auth else None,

-        inbrowser=cmd_opts.autolaunch,

-    )

+    while 1:

+

+        demo = modules.ui.create_ui(wrap_gradio_gpu_call=wrap_gradio_gpu_call)

+        

+        app, local_url, share_url = demo.launch(

+            share=cmd_opts.share,

+            server_name="0.0.0.0" if cmd_opts.listen else None,

+            server_port=cmd_opts.port,

+            debug=cmd_opts.gradio_debug,

+            auth=[tuple(cred.split(':')) for cred in cmd_opts.gradio_auth.strip('"').split(',')] if cmd_opts.gradio_auth else None,

+            inbrowser=cmd_opts.autolaunch,

+            prevent_thread_lock=True

+        )

+        

+        app.add_middleware(GZipMiddleware, minimum_size=1000)

+

+        while 1:

+            time.sleep(0.5)

+            if getattr(demo, 'do_restart', False):

+                time.sleep(0.5)

+                demo.close()

+                time.sleep(0.5)

+                break

+

+        sd_samplers.set_samplers()

+

+        print('Reloading Custom Scripts')

+        modules.scripts.reload_scripts(os.path.join(script_path, "scripts"))

+        print('Reloading modules: modules.ui')

+        importlib.reload(modules.ui)

+        print('Refreshing Model List')

+        modules.sd_models.list_models()

+        print('Restarting Gradio')

 

 

 if __name__ == "__main__":