add model refresh button to stable audio UI

rsxdalv · Oct 7, 2024 · 94ffe52 · 94ffe52
1 parent 2a70117
commit 94ffe52
Showing 1 changed file with 8 additions and 0 deletions.
diff --git a/tts_webui/stable_audio/stable_audio.py b/tts_webui/stable_audio/stable_audio.py
@@ -14,6 +14,7 @@
 
 from tts_webui.utils.torch_clear_memory import torch_clear_memory
 from tts_webui.utils.prompt_to_title import prompt_to_title
+from tts_webui.tortoise.gr_reload_button import gr_open_button_simple, gr_reload_button
 
 LOCAL_DIR_BASE = os.path.join("data", "models", "stable-audio")
 LOCAL_DIR_BASE_ABSOLUTE = get_path_from_root(*LOCAL_DIR_BASE.split(os.path.sep))
@@ -157,6 +158,13 @@ def model_select_ui():
                 value=pretrained_name,
             )
 
+            gr_open_button_simple(LOCAL_DIR_BASE, api_name="stable_audio_open_models")
+            gr_reload_button().click(
+                fn=lambda: gr.Dropdown(choices=get_model_list()),
+                outputs=[model_select],
+                api_name="stable_audio_refresh_models",
+            )
+
             with gr.Column():
                 gr.Markdown(
                     """