tabs/voice_blender/voice_blender.py

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
import zipfile

i18n = I18nAuto()

# Global variables to store the downloaded model paths
model_paths = {"model_a": None, "model_b": None}


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)

        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)}")


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_b"], model_paths["model_a"], alpha_a)
        return message, model_blended
    except Exception as e:
        return f"An error occurred blending the models: {str(e)}", None


def voice_blender_tab():
    gr.Markdown(i18n("## Voice Blender"))
    gr.Markdown(
        i18n(
            "Select two voice models, set your desired blend percentage, and blend them into an entirely new voice."
        )
    )
    with gr.Column():
        model_fusion_name = gr.Textbox(
            label=i18n("Model Name"),
            info=i18n("Name of the new model."),
            value="",
            max_lines=1,
            interactive=True,
            placeholder=i18n("Enter model name"),
        )
        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,
            label=i18n("Blend Ratio"),
            value=0.5,
            interactive=True,
            info=i18n(
                "Adjusting the position more towards one side or the other will make the model more similar to the first or second."
            ),
        )
        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],
    )

    blend_button.click(
        fn=blend_models,
        inputs=[model_fusion_name, alpha_a],
        outputs=[model_fusion_output_info, model_fusion_pth_output],
    )