Update tabs/voice_blender/voice_blender.py

tabs/voice_blender/voice_blender.py

@@ -1,40 +1,62 @@
-import os, sys
+import os
 import gradio as gr
 import shutil
+import tempfile
+from core import run_model_blender_script
+from assets.i18n.i18n import I18nAuto
 import requests
-from zipfile import ZipFile
-from tempfile import TemporaryDirectory
+import zipfile
 
-now_dir = os.getcwd()
-sys.path.append(now_dir)
+i18n = I18nAuto()
 
-from assets.i18n.i18n import I18nAuto
-from core import run_model_blender_script
+# Global variables to store the downloaded model paths
+model_paths = {"model_a": None, "model_b": None}
 
-i18n = I18nAuto()
 
-def download_and_extract_model(url):
-    temp_dir = TemporaryDirectory()
-    zip_path = os.path.join(temp_dir.name, "model.zip")
+def download_and_extract(url, extract_to):
+    try:
+        with requests.get(url, stream=True) as r:
+            r.raise_for_status()
+            with tempfile.NamedTemporaryFile(delete=False) as tmp_file:
+                tmp_file.write(r.content)
+        
+        with zipfile.ZipFile(tmp_file.name, 'r') as zip_ref:
+            zip_ref.extractall(extract_to)
 
-    with requests.get(url, stream=True) as r:
-        with open(zip_path, 'wb') as f:
-            shutil.copyfileobj(r.raw, f)
+        pth_files = [os.path.join(extract_to, f) for f in os.listdir(extract_to) if f.endswith('.pth')]
+        if len(pth_files) == 0:
+            raise FileNotFoundError(f"No .pth files found in {extract_to}")
+        
+        return pth_files[0]  # Assuming there's one .pth file per model
+    except Exception as e:
+        raise RuntimeError(f"Failed to download or extract model from {url}. Error: {str(e)}")
 
-    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_models(model_url_1, model_url_2):
+    try:
+        # Create temporary directories for model extraction
+        temp_dir_1 = tempfile.mkdtemp()
+        temp_dir_2 = tempfile.mkdtemp()
+
+        # Download and extract models
+        model_paths["model_a"] = download_and_extract(model_url_1, temp_dir_1)
+        model_paths["model_b"] = download_and_extract(model_url_2, temp_dir_2)
+
+        return "Models downloaded successfully!", model_paths["model_a"], model_paths["model_b"]
+    except Exception as e:
+        return f"An error occurred during download: {str(e)}", None, None
+
+
+def blend_models(model_name, alpha_a):
+    try:
+        if not model_paths["model_a"] or not model_paths["model_b"]:
+            raise FileNotFoundError("Models not found. Please download the models first.")
+        
+        message, model_blended = run_model_blender_script(model_name, model_paths["model_a"], model_paths["model_b"], alpha_a)
+        return message, model_blended
+    except Exception as e:
+        return f"An error occurred blending the models: {str(e)}", None
 
-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"))
@@ -52,33 +74,15 @@ def voice_blender_tab():
             interactive=True,
             placeholder=i18n("Enter model name"),
         )
-        with gr.Row():
-            with gr.Column():
-                model_fusion_a_link = gr.Textbox(
-                    label=i18n("Hugging Face Link for Model A"),
-                    value="",
-                    interactive=True,
-                    placeholder=i18n("Enter Hugging Face link for model A"),
-                )
-                model_fusion_a_path = gr.Textbox(
-                    label=i18n("Path to Model A"),
-                    value="",
-                    interactive=False,
-                    placeholder=i18n("Path will be set automatically after download"),
-                )
-            with gr.Column():
-                model_fusion_b_link = gr.Textbox(
-                    label=i18n("Hugging Face Link for Model B"),
-                    value="",
-                    interactive=True,
-                    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"),
-                )
+        model_fusion_a_url = gr.Textbox(
+            label=i18n("Model A URL"),
+            placeholder=i18n("Enter Hugging Face URL for Model A"),
+        )
+        model_fusion_b_url = gr.Textbox(
+            label=i18n("Model B URL"),
+            placeholder=i18n("Enter Hugging Face URL for Model B"),
+        )
+        download_button = gr.Button(i18n("Download Models"), variant="primary")
         alpha_a = gr.Slider(
             minimum=0,
             maximum=1,
@@ -89,30 +93,43 @@ def voice_blender_tab():
                 "Adjusting the position more towards one side or the other will make the model more similar to the first or second."
             ),
         )
-        model_fusion_button = gr.Button(i18n("Fusion"), variant="primary")
-        with gr.Row():
-            model_fusion_output_info = gr.Textbox(
-                label=i18n("Output Information"),
-                info=i18n("The output information will be displayed here."),
-                value="",
-            )
-            model_fusion_pth_output = gr.File(
-                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],
+        blend_button = gr.Button(i18n("Blend Models"), variant="primary")
+
+        # Outputs
+        download_info = gr.Textbox(
+            label=i18n("Download Information"),
+            value="",
+            interactive=False,
+        )
+        model_fusion_a_path = gr.Textbox(
+            label=i18n("Model A Path"),
+            value="",
+            interactive=False,
+        )
+        model_fusion_b_path = gr.Textbox(
+            label=i18n("Model B Path"),
+            value="",
+            interactive=False,
+        )
+        model_fusion_output_info = gr.Textbox(
+            label=i18n("Output Information"),
+            value="",
+            interactive=False,
+        )
+        model_fusion_pth_output = gr.File(
+            label=i18n("Download Blended Model"),
+            type="filepath",
+            interactive=False,
+        )
+
+    download_button.click(
+        fn=download_models,
+        inputs=[model_fusion_a_url, model_fusion_b_url],
+        outputs=[download_info, model_fusion_a_path, model_fusion_b_path],
     )
 
-    model_fusion_button.click(
-        fn=run_model_blender_script,
-        inputs=[
-            model_fusion_name,
-            model_fusion_a_path,
-            model_fusion_b_path,
-            alpha_a,
-        ],
+    blend_button.click(
+        fn=blend_models,
+        inputs=[model_fusion_name, alpha_a],
         outputs=[model_fusion_output_info, model_fusion_pth_output],
     )