Update tabs/voice_blender/voice_blender.py

tabs/voice_blender/voice_blender.py

@@ -1,6 +1,9 @@
 import os, sys
 import gradio as gr
 import shutil
+import requests
+from zipfile import ZipFile
+from tempfile import TemporaryDirectory
 
 now_dir = os.getcwd()
 sys.path.append(now_dir)
@@ -10,10 +13,28 @@ from core import run_model_blender_script
 
 i18n = I18nAuto()
 
+def download_and_extract_model(url):
+    temp_dir = TemporaryDirectory()
+    zip_path = os.path.join(temp_dir.name, "model.zip")
 
-def update_model_fusion(dropbox):
-    return dropbox, None
+    with requests.get(url, stream=True) as r:
+        with open(zip_path, 'wb') as f:
+            shutil.copyfileobj(r.raw, f)
 
+    with ZipFile(zip_path, 'r') as zip_ref:
+        zip_ref.extractall(temp_dir.name)
+    
+    pth_files = [os.path.join(temp_dir.name, f) for f in os.listdir(temp_dir.name) if f.endswith('.pth')]
+
+    if pth_files:
+        return pth_files[0]
+    else:
+        raise FileNotFoundError("No se encontró un archivo .pth en el zip descargado.")
+
+def download_and_set_paths(link_a, link_b):
+    path_a = download_and_extract_model(link_a)
+    path_b = download_and_extract_model(link_b)
+    return path_a, path_b
 
 def voice_blender_tab():
     gr.Markdown(i18n("## Voice Blender"))
@@ -33,26 +54,30 @@ def voice_blender_tab():
         )
         with gr.Row():
             with gr.Column():
-                model_fusion_a_dropbox = gr.File(
-                    label=i18n("Drag and drop your model here"), type="filepath"
-                )
-                model_fusion_a = gr.Textbox(
-                    label=i18n("Path to Model"),
+                model_fusion_a_link = gr.Textbox(
+                    label=i18n("Hugging Face Link for Model A"),
                     value="",
                     interactive=True,
-                    placeholder=i18n("Enter path to model"),
-                    info=i18n("You can also use a custom path."),
+                    placeholder=i18n("Enter Hugging Face link for model A"),
                 )
-            with gr.Column():
-                model_fusion_b_dropbox = gr.File(
-                    label=i18n("Drag and drop your model here"), type="filepath"
+                model_fusion_a_path = gr.Textbox(
+                    label=i18n("Path to Model A"),
+                    value="",
+                    interactive=False,
+                    placeholder=i18n("Path will be set automatically after download"),
                 )
-                model_fusion_b = gr.Textbox(
-                    label=i18n("Path to Model"),
+            with gr.Column():
+                model_fusion_b_link = gr.Textbox(
+                    label=i18n("Hugging Face Link for Model B"),
                     value="",
                     interactive=True,
-                    placeholder=i18n("Enter path to model"),
-                    info=i18n("You can also use a custom path."),
+                    placeholder=i18n("Enter Hugging Face link for model B"),
+                )
+                model_fusion_b_path = gr.Textbox(
+                    label=i18n("Path to Model B"),
+                    value="",
+                    interactive=False,
+                    placeholder=i18n("Path will be set automatically after download"),
                 )
         alpha_a = gr.Slider(
             minimum=0,
@@ -75,25 +100,19 @@ def voice_blender_tab():
                 label=i18n("Download Model"), type="filepath", interactive=False
             )
 
+    model_fusion_button.click(
+        fn=download_and_set_paths,
+        inputs=[model_fusion_a_link, model_fusion_b_link],
+        outputs=[model_fusion_a_path, model_fusion_b_path],
+    )
+
     model_fusion_button.click(
         fn=run_model_blender_script,
         inputs=[
             model_fusion_name,
-            model_fusion_a,
-            model_fusion_b,
+            model_fusion_a_path,
+            model_fusion_b_path,
             alpha_a,
         ],
         outputs=[model_fusion_output_info, model_fusion_pth_output],
-    )
-
-    model_fusion_a_dropbox.upload(
-        fn=update_model_fusion,
-        inputs=model_fusion_a_dropbox,
-        outputs=[model_fusion_a, model_fusion_a_dropbox],
-    )
-
-    model_fusion_b_dropbox.upload(
-        fn=update_model_fusion,
-        inputs=model_fusion_b_dropbox,
-        outputs=[model_fusion_b, model_fusion_b_dropbox],
-    )
+    )