add progress bar to modelmerger

1 files changed, 8 insertions, 5 deletions

diff --git a/modules/ui.py b/modules/ui.py
index 09a3c92e..aeee7853 100644
--- a/modules/ui.py
+++ b/modules/ui.py
@@ -1208,8 +1208,9 @@ def create_ui():
                 with gr.Row():

                     modelmerger_merge = gr.Button(elem_id="modelmerger_merge", value="Merge", variant='primary')

 

-            with gr.Column(variant='panel'):

-                submit_result = gr.Textbox(elem_id="modelmerger_result", show_label=False)

+            with gr.Column(variant='compact', elem_id="modelmerger_results_container"):

+                with gr.Group(elem_id="modelmerger_results_panel"):

+                    modelmerger_result = gr.HTML(elem_id="modelmerger_result", show_label=False)

 

     with gr.Blocks(analytics_enabled=False) as train_interface:

         with gr.Row().style(equal_height=False):

@@ -1753,12 +1754,14 @@ def create_ui():
                 print("Error loading/saving model file:", file=sys.stderr)

                 print(traceback.format_exc(), file=sys.stderr)

                 modules.sd_models.list_models()  # to remove the potentially missing models from the list

-                return [f"Error merging checkpoints: {e}"] + [gr.Dropdown.update(choices=modules.sd_models.checkpoint_tiles()) for _ in range(4)]

+                return [*[gr.Dropdown.update(choices=modules.sd_models.checkpoint_tiles()) for _ in range(4)], f"Error merging checkpoints: {e}"]

             return results

 

         modelmerger_merge.click(

-            fn=modelmerger,

+            fn=wrap_gradio_gpu_call(modelmerger, extra_outputs=lambda: [gr.update() for _ in range(4)]),

+            _js='modelmerger',

             inputs=[

+                dummy_component,

                 primary_model_name,

                 secondary_model_name,

                 tertiary_model_name,

@@ -1770,11 +1773,11 @@ def create_ui():
                 config_source,

             ],

             outputs=[

-                submit_result,

                 primary_model_name,

                 secondary_model_name,

                 tertiary_model_name,

                 component_dict['sd_model_checkpoint'],

+                modelmerger_result,

             ]

         )