layer diffuse enhancements

CHANGELOG.md

@@ -1,6 +1,6 @@
 # Change Log for SD.Next
 
-## Update for 2024-05-27
+## Update for 2024-05-28
 
 - **Features**:
   - **ModernUI** preview of the new [ModernUI](https://github.com/BinaryQuantumSoul/sdnext-modernui)  

html/locale_en.json

@@ -212,8 +212,8 @@
   {"id":"shutdown_submit","label":"Shutdown server","localized":"","hint":"Shutdown server"},
   {"id":"settings_preview_theme","label":"Preview theme","localized":"","hint":"Show theme preview"},
   {"id":"defaults_submit","label":"Restore defaults","localized":"","hint":"Restore default server settings"},
-  {"id":"sett_unload_sd_model","label":"Unload checkpoint","localized":"","hint":"Unload currently loaded model checkpoint"},
-  {"id":"sett_reload_sd_model","label":"Reload checkpoint","localized":"","hint":"Reload currently selected model checkpoint"}
+  {"id":"sett_unload_sd_model","label":"Unload model","localized":"","hint":"Unload currently loaded model"},
+  {"id":"sett_reload_sd_model","label":"Reload model","localized":"","hint":"Reload currently selected model"}
 ],
 "settings sections": [
   {"id":"","label":"Execution & Models","localized":"","hint":"Settings related to execution backend, models, and prompt attention"},

javascript/sdnext.css

@@ -5,7 +5,7 @@
   --color-debug: #7F7F7F;
   --color-info: #D4D4D4;
   --color-warning: #FF9900;
-  --color-error: #BE0000
+  --color-error: #BE0000;
 }
 a { font-weight: bold; cursor: pointer; }
 h2 { margin-top: 1em !important; font-size: var(--text-xxl) !important; }

modules/sd_models.py

@@ -1486,7 +1486,7 @@ def load_model(checkpoint_info=None, already_loaded_state_dict=None, timer=None,
     shared.log.info(f'Model load finished: {memory_stats()} cached={len(checkpoints_loaded.keys())}')
 
 
-def reload_model_weights(sd_model=None, info=None, reuse_dict=False, op='model'):
+def reload_model_weights(sd_model=None, info=None, reuse_dict=False, op='model', force=False):
     load_dict = shared.opts.sd_model_dict != model_data.sd_dict
     from modules import lowvram, sd_hijack
     checkpoint_info = info or select_checkpoint(op=op) # are we selecting model or dictionary
@@ -1508,7 +1508,7 @@ def reload_model_weights(sd_model=None, info=None, reuse_dict=False, op='model')
         current_checkpoint_info = None
     else:
         current_checkpoint_info = getattr(sd_model, 'sd_checkpoint_info', None)
-        if current_checkpoint_info is not None and checkpoint_info is not None and current_checkpoint_info.filename == checkpoint_info.filename:
+        if current_checkpoint_info is not None and checkpoint_info is not None and current_checkpoint_info.filename == checkpoint_info.filename and not force:
             return None
         if shared.backend == shared.Backend.ORIGINAL and (shared.cmd_opts.lowvram or shared.cmd_opts.medvram):
             lowvram.send_everything_to_cpu()

modules/ui.py

@@ -276,8 +276,8 @@ def create_ui(startup_timer = None):
         with gr.Row(elem_id="system_row"):
             restart_submit = gr.Button(value="Restart server", variant='primary', elem_id="restart_submit")
             shutdown_submit = gr.Button(value="Shutdown server", variant='primary', elem_id="shutdown_submit")
-            unload_sd_model = gr.Button(value='Unload checkpoint', variant='primary', elem_id="sett_unload_sd_model")
-            reload_sd_model = gr.Button(value='Reload checkpoint', variant='primary', elem_id="sett_reload_sd_model")
+            unload_sd_model = gr.Button(value='Unload model', variant='primary', elem_id="sett_unload_sd_model")
+            reload_sd_model = gr.Button(value='Reload model', variant='primary', elem_id="sett_reload_sd_model")
             enable_profiling = gr.Button(value='Start profiling', variant='primary', elem_id="enable_profiling")
 
         with gr.Tabs(elem_id="system") as system_tabs:
@@ -362,7 +362,7 @@ def create_ui(startup_timer = None):
             modules.sd_models.unload_model_weights(op='refiner')
 
         def reload_sd_weights():
-            modules.sd_models.reload_model_weights()
+            modules.sd_models.reload_model_weights(force=True)
 
         def switch_profiling():
             shared.cmd_opts.profile = not shared.cmd_opts.profile

scripts/layerdiffuse.py

@@ -1,5 +1,5 @@
 import gradio as gr
-from modules import scripts, shared
+from modules import shared, scripts, sd_models
 
 
 class Script(scripts.Script):
@@ -14,23 +14,43 @@ class Script(scripts.Script):
         from modules import layerdiffuse
         if not shared.sd_loaded:
             shared.log.error('LayerDiffuse: model not loaded')
-            return
+            return self.is_active()
         if shared.sd_model_type != 'sd' and shared.sd_model_type != 'sdxl':
             shared.log.error(f'LayerDiffuse: incorrect base model: class={shared.sd_model.__class__.__name__} type={shared.sd_model_type}')
-            return
+            return self.is_active()
         if hasattr(shared.sd_model, 'layerdiffusion'):
             shared.log.warning('LayerDiffuse: already applied')
-            return
+            return self.is_active()
         layerdiffuse.apply_layerdiffuse()
+        return self.is_active()
+
+    def reload(self):
+        sd_models.reload_model_weights(force=True)
+        return self.is_active()
+
+    def is_active(self):
+        if not shared.sd_loaded:
+            return '<div style="color: darkred">LayerDiffuse: model not loaded</div><br>'
+        if shared.sd_model_type != 'sd' and shared.sd_model_type != 'sdxl':
+            return '<div style="color: darkred">LayerDiffuse: incorrect base model</div><br>'
+        if hasattr(shared.sd_model, 'layerdiffusion'):
+            return '<div style="color: darkgreen">LayerDiffuse: active</div><br>'
+        return '<div style="color: darkgray">LayerDiffuse: inactive</div><br>'
 
     def ui(self, _is_img2img):
         with gr.Row():
             gr.HTML("""
                     <a href="https://github.com/rootonchair/diffuser_layerdiffuse">&nbsp LayerDiffuse</a><br><br>
-                    <div>Click once to permanently apply to current model</div>
-                    <div>Reload model to unapply</div><br>
+                    <div>- Click Apply to model to apply LayerDiffuse to current model</div>
+                    <div>- Click Reload model to remove LayerDiffuse from current model</div><br>
                     """)
         with gr.Row():
+            active = gr.HTML('')
+        with gr.Row():
+            check_btn = gr.Button('Check status', variant='primary')
             apply_btn = gr.Button('Apply to model', variant='primary')
-            apply_btn.click(fn=self.apply)
+            reload_btn = gr.Button('Reload model', variant='primary')
+            check_btn.click(fn=self.is_active, inputs=[], outputs=[active])
+            apply_btn.click(fn=self.apply, inputs=[], outputs=[active])
+            reload_btn.click(fn=self.reload, inputs=[], outputs=[active])
         return []