webUI.py

import os
import json
import subprocess
from typing import List, Dict, Callable, Union

import gradio as gr
import librosa
import numpy as np
import soundfile

from UVR_interface import root, UVRInterface, VR_MODELS_DIR, MDX_MODELS_DIR, DEMUCS_MODELS_DIR
from gui_data.constants import *

# Preload only a very small default set so the Space can start quickly.
# All other VR/MDX models are downloaded on demand when selected.
DEFAULT_STARTUP_MODELS = {
    VR_ARCH_TYPE: {"1_HP-UVR"},
    MDX_ARCH_TYPE: {"UVR-MDX-NET-Inst_Main"},
    DEMUCS_ARCH_TYPE: set(),
}


class UVRWebUI:
    def __init__(self, uvr: UVRInterface, online_data_path: str) -> None:
        self.uvr = uvr
        self.input_temp_dir = "__temp"
        self.export_path = "out"

        os.makedirs(self.input_temp_dir, exist_ok=True)
        os.makedirs(self.export_path, exist_ok=True)
        os.makedirs(VR_MODELS_DIR, exist_ok=True)
        os.makedirs(MDX_MODELS_DIR, exist_ok=True)
        os.makedirs(DEMUCS_MODELS_DIR, exist_ok=True)

        self.models_url = self.get_models_url(online_data_path)
        self.download_url_by_local_name = self.build_download_url_by_local_name()
        self.preload_default_models()
        self.define_layout()

    def get_models_url(self, models_info_path: str) -> Dict[str, Dict]:
        with open(models_info_path, "r", encoding="utf-8") as f:
            online_data = json.loads(f.read())

        models_url: Dict[str, Dict] = {}
        for arch, download_list_key in zip(
            [VR_ARCH_TYPE, MDX_ARCH_TYPE],
            ["vr_download_list", "mdx_download_list"],
        ):
            models_url[arch] = {
                model: NORMAL_REPO + model_path
                for model, model_path in online_data[download_list_key].items()
            }
        models_url[DEMUCS_ARCH_TYPE] = online_data["demucs_download_list"]
        return models_url

    def build_download_url_by_local_name(self) -> Dict[str, Dict[str, str]]:
        download_map: Dict[str, Dict[str, str]] = {
            VR_ARCH_TYPE: {},
            MDX_ARCH_TYPE: {},
            DEMUCS_ARCH_TYPE: {},
        }

        for arch in [VR_ARCH_TYPE, MDX_ARCH_TYPE]:
            for _, model_url in self.models_url[arch].items():
                local_name = os.path.splitext(os.path.basename(model_url))[0]
                download_map[arch][local_name] = model_url

        return download_map

    def get_model_dir_and_suffix(self, arch: str):
        model_config = {
            VR_ARCH_TYPE: (VR_MODELS_DIR, ".pth"),
            MDX_ARCH_TYPE: (MDX_MODELS_DIR, ".onnx"),
            DEMUCS_ARCH_TYPE: (DEMUCS_MODELS_DIR, ".yaml"),
        }
        try:
            return model_config[arch]
        except KeyError as exc:
            raise ValueError(f"Unknown arch type: {arch}") from exc

    def get_local_models(self, arch: str) -> List[str]:
        model_dir, suffix = self.get_model_dir_and_suffix(arch)
        os.makedirs(model_dir, exist_ok=True)
        return sorted(
            [os.path.splitext(f)[0] for f in os.listdir(model_dir) if f.endswith(suffix)]
        )

    def get_available_models(self, arch: str) -> List[str]:
        available = set(self.get_local_models(arch))
        available.update(self.download_url_by_local_name.get(arch, {}).keys())
        return sorted(available)

    def should_preload(self, arch: str, model_name: str) -> bool:
        return model_name in DEFAULT_STARTUP_MODELS.get(arch, set())

    def download_file(self, model_url: str, model_dir: str) -> None:
        os.makedirs(model_dir, exist_ok=True)
        cmd = [
            "aria2c",
            "--optimize-concurrent-downloads",
            "--console-log-level=error",
            "--summary-interval=10",
            "-j5",
            "-x16",
            "-s16",
            "-k1M",
            "-c",
            "-d",
            model_dir,
            "-Z",
            model_url,
        ]
        try:
            subprocess.run(cmd, check=True)
        except FileNotFoundError as exc:
            raise RuntimeError("aria2c is not installed in this environment.") from exc
        except subprocess.CalledProcessError as exc:
            raise RuntimeError(f"Failed to download model: {model_url}") from exc

    def ensure_model_exists(self, arch: str, model_name: str) -> str:
        model_dir, suffix = self.get_model_dir_and_suffix(arch)
        local_path = os.path.join(model_dir, f"{model_name}{suffix}")
        if os.path.exists(local_path):
            return local_path

        model_url = self.download_url_by_local_name.get(arch, {}).get(model_name)
        if not model_url:
            raise gr.Error(f"No download URL found for model: {model_name}")

        print(f"Downloading on demand: {model_name}")
        self.download_file(model_url, model_dir)

        if not os.path.exists(local_path):
            raise gr.Error(f"Model download finished but file was not found: {local_path}")

        return local_path

    def preload_default_models(self) -> None:
        for arch in [VR_ARCH_TYPE, MDX_ARCH_TYPE]:
            model_dir, suffix = self.get_model_dir_and_suffix(arch)
            downloaded_any = False

            for model_name, model_url in self.download_url_by_local_name[arch].items():
                if not self.should_preload(arch, model_name):
                    continue

                local_path = os.path.join(model_dir, f"{model_name}{suffix}")
                if os.path.exists(local_path):
                    continue

                print(f"Preloading default {arch} model: {model_name}")
                self.download_file(model_url, model_dir)
                downloaded_any = True

            if downloaded_any:
                print(f"Default {arch} models downloaded successfully.")
            else:
                print(f"No default {arch} model download needed.")

        print(f"Ignoring category: {DEMUCS_ARCH_TYPE}")

    def set_arch_setting_value(self, arch: str, setting1, setting2):
        if arch == VR_ARCH_TYPE:
            root.window_size_var.set(setting1)
            root.aggression_setting_var.set(setting2)
        elif arch == MDX_ARCH_TYPE:
            root.mdx_batch_size_var.set(setting1)
            root.compensate_var.set(setting2)
        elif arch == DEMUCS_ARCH_TYPE:
            pass

    def arch_select_update(self, arch: str) -> List[Dict]:
        choices = self.get_available_models(arch)
        if arch == VR_ARCH_TYPE:
            model_update = gr.update(
                choices=choices,
                value=CHOOSE_MODEL,
                label=SELECT_VR_MODEL_MAIN_LABEL,
            )
            setting1_update = gr.update(
                choices=VR_WINDOW,
                label=WINDOW_SIZE_MAIN_LABEL,
                value=root.window_size_var.get(),
            )
            setting2_update = gr.update(
                choices=VR_AGGRESSION,
                label=AGGRESSION_SETTING_MAIN_LABEL,
                value=root.aggression_setting_var.get(),
            )
        elif arch == MDX_ARCH_TYPE:
            model_update = gr.update(
                choices=choices,
                value=CHOOSE_MODEL,
                label=CHOOSE_MDX_MODEL_MAIN_LABEL,
            )
            setting1_update = gr.update(
                choices=BATCH_SIZE,
                label=BATCHES_MDX_MAIN_LABEL,
                value=root.mdx_batch_size_var.get(),
            )
            setting2_update = gr.update(
                choices=VOL_COMPENSATION,
                label=VOL_COMP_MDX_MAIN_LABEL,
                value=root.compensate_var.get(),
            )
        elif arch == DEMUCS_ARCH_TYPE:
            model_update = gr.update(
                choices=choices,
                value=CHOOSE_MODEL,
                label=CHOOSE_DEMUCS_MODEL_MAIN_LABEL,
            )
            raise gr.Error(f"{DEMUCS_ARCH_TYPE} not implemented")
        else:
            raise gr.Error(f"Unknown arch type: {arch}")
        return [model_update, setting1_update, setting2_update]

    def model_select_update(self, arch: str, model_name: str) -> List[Union[str, Dict, None]]:
        if model_name == CHOOSE_MODEL:
            return [None for _ in range(4)]

        if arch in [VR_ARCH_TYPE, MDX_ARCH_TYPE]:
            self.ensure_model_exists(arch, model_name)

        model, = self.uvr.assemble_model_data(model_name, arch)
        if not model.model_status:
            raise gr.Error(f"Cannot get model data, model hash = {model.model_hash}")

        stem1_check_update = gr.update(label=f"{model.primary_stem} Only")
        stem2_check_update = gr.update(label=f"{model.secondary_stem} Only")
        stem1_out_update = gr.update(label=f"Output {model.primary_stem}")
        stem2_out_update = gr.update(label=f"Output {model.secondary_stem}")

        return [
            stem1_check_update,
            stem2_check_update,
            stem1_out_update,
            stem2_out_update,
        ]

    def checkbox_set_root_value(self, checkbox: gr.Checkbox, root_attr: str):
        checkbox.change(lambda value: root.__getattribute__(root_attr).set(value), inputs=checkbox)

    def set_checkboxes_exclusive(
        self,
        checkboxes: List[gr.Checkbox],
        pure_callbacks: List[Callable],
        exclusive_value=True,
    ):
        def exclusive_onchange(i, callback_i):
            def new_onchange(*check_values):
                if check_values[i] == exclusive_value:
                    return_values = []
                    for j, value_j in enumerate(check_values):
                        if j != i and value_j == exclusive_value:
                            return_values.append(not exclusive_value)
                        else:
                            return_values.append(value_j)
                else:
                    return_values = check_values
                callback_i(check_values[i])
                return return_values

            return new_onchange

        for i, (checkbox, callback) in enumerate(zip(checkboxes, pure_callbacks)):
            checkbox.change(exclusive_onchange(i, callback), inputs=checkboxes, outputs=checkboxes)

    def process(
        self,
        input_audio,
        input_filename,
        model_name,
        arch,
        setting1,
        setting2,
        progress=gr.Progress(),
    ):
        if input_audio is None:
            raise gr.Error("Please upload an audio file before processing.")
        if not model_name or model_name == CHOOSE_MODEL:
            raise gr.Error("Please select a model before processing.")

        def set_progress_func(step, inference_iterations=0):
            progress_curr = step + inference_iterations
            progress(progress_curr)

        if arch in [VR_ARCH_TYPE, MDX_ARCH_TYPE]:
            self.ensure_model_exists(arch, model_name)

        sampling_rate, audio = input_audio
        audio = np.asarray(audio)
        if np.issubdtype(audio.dtype, np.integer):
            audio = audio.astype(np.float32) / np.iinfo(audio.dtype).max
        else:
            audio = audio.astype(np.float32)

        if len(audio.shape) > 1:
            audio = librosa.to_mono(audio.transpose(1, 0))

        safe_filename = os.path.basename(input_filename.strip()) if input_filename else "temp.wav"
        if not safe_filename:
            safe_filename = "temp.wav"
        if not safe_filename.lower().endswith(".wav"):
            safe_filename = f"{safe_filename}.wav"

        input_path = os.path.join(self.input_temp_dir, safe_filename)
        soundfile.write(input_path, audio, sampling_rate, format="wav")

        self.set_arch_setting_value(arch, setting1, setting2)

        separator = self.uvr.process(
            model_name=model_name,
            arch_type=arch,
            audio_file=input_path,
            export_path=self.export_path,
            is_model_sample_mode=root.model_sample_mode_var.get(),
            set_progress_func=set_progress_func,
        )

        primary_audio = None
        secondary_audio = None
        msg = ""
        if not separator.is_secondary_stem_only:
            primary_stem_path = os.path.join(
                separator.export_path,
                f"{separator.audio_file_base}_({separator.primary_stem}).wav",
            )
            audio_out, rate = soundfile.read(primary_stem_path)
            primary_audio = (rate, audio_out)
            msg += f"{separator.primary_stem} saved at {primary_stem_path}\n"
        if not separator.is_primary_stem_only:
            secondary_stem_path = os.path.join(
                separator.export_path,
                f"{separator.audio_file_base}_({separator.secondary_stem}).wav",
            )
            audio_out, rate = soundfile.read(secondary_stem_path)
            secondary_audio = (rate, audio_out)
            msg += f"{separator.secondary_stem} saved at {secondary_stem_path}\n"

        if os.path.exists(input_path):
            os.remove(input_path)

        return primary_audio, secondary_audio, msg.strip()

    def define_layout(self):
        with gr.Blocks() as app:
            self.app = app
            gr.HTML("<h1> 🎵 Ultimate Vocal Remover WebUI 🎵 </h1>")
            gr.Markdown("This is an experimental demo with CPU. Duplicate the space for use in private")
            gr.Markdown(
                "[![Duplicate this Space](https://huggingface.co/datasets/huggingface/badges/raw/main/duplicate-this-space-sm-dark.svg)](https://huggingface.co/spaces/r3gm/Ultimate-Vocal-Remover-WebUI?duplicate=true)\n\n"
            )
            with gr.Tabs():
                with gr.TabItem("process"):
                    with gr.Row():
                        self.arch_choice = gr.Dropdown(
                            choices=[VR_ARCH_TYPE, MDX_ARCH_TYPE],
                            value=VR_ARCH_TYPE,
                            label=CHOOSE_PROC_METHOD_MAIN_LABEL,
                            interactive=True,
                        )
                        self.model_choice = gr.Dropdown(
                            choices=self.get_available_models(VR_ARCH_TYPE),
                            value=CHOOSE_MODEL,
                            label=SELECT_VR_MODEL_MAIN_LABEL + " 👋Select a model",
                            interactive=True,
                        )
                    with gr.Row():
                        self.arch_setting1 = gr.Dropdown(
                            choices=VR_WINDOW,
                            value=root.window_size_var.get(),
                            label=WINDOW_SIZE_MAIN_LABEL + " 👋Select one",
                            interactive=True,
                        )
                        self.arch_setting2 = gr.Dropdown(
                            choices=VR_AGGRESSION,
                            value=root.aggression_setting_var.get(),
                            label=AGGRESSION_SETTING_MAIN_LABEL,
                            interactive=True,
                        )
                    with gr.Row():
                        self.use_gpu = gr.Checkbox(
                            label="Rhythmic Transmutation Device",
                            value=True,
                            interactive=True,
                        )
                        self.primary_stem_only = gr.Checkbox(
                            label=f"{PRIMARY_STEM} only",
                            value=root.is_primary_stem_only_var.get(),
                            interactive=True,
                        )
                        self.secondary_stem_only = gr.Checkbox(
                            label=f"{SECONDARY_STEM} only",
                            value=root.is_secondary_stem_only_var.get(),
                            interactive=True,
                        )
                        self.sample_mode = gr.Checkbox(
                            label=SAMPLE_MODE_CHECKBOX(root.model_sample_mode_duration_var.get()),
                            value=root.model_sample_mode_var.get(),
                            interactive=True,
                        )

                    with gr.Row():
                        self.input_filename = gr.Textbox(
                            label="Input filename",
                            value="temp.wav",
                            interactive=True,
                        )
                    with gr.Row():
                        self.audio_in = gr.Audio(label="Input audio", interactive=True)
                    with gr.Row():
                        self.process_submit = gr.Button(START_PROCESSING, variant="primary")
                    with gr.Row():
                        self.primary_stem_out = gr.Audio(
                            label=f"Output {PRIMARY_STEM}",
                            interactive=False,
                        )
                        self.secondary_stem_out = gr.Audio(
                            label=f"Output {SECONDARY_STEM}",
                            interactive=False,
                        )
                    with gr.Row():
                        self.out_message = gr.Textbox(
                            label="Output Message",
                            interactive=False,
                        )

                with gr.TabItem("settings"):
                    with gr.Tabs():
                        with gr.TabItem("Settings Guide"):
                            pass
                        with gr.TabItem("Additional Settigns"):
                            self.wav_type = gr.Dropdown(
                                choices=WAV_TYPE,
                                label="Wav Type",
                                value="PCM_16",
                                interactive=True,
                            )
                            self.mp3_rate = gr.Dropdown(
                                choices=MP3_BIT_RATES,
                                label="MP3 Bitrate",
                                value="320k",
                                interactive=True,
                            )
                        with gr.TabItem("Download models"):

                            def md_url(url, text=None):
                                if text is None:
                                    text = url
                                return f"[{text}]({url})"

                            with gr.Row():
                                vr_models = self.models_url[VR_ARCH_TYPE]
                                self.vr_download_choice = gr.Dropdown(
                                    choices=list(vr_models.keys()),
                                    label=f"Select {VR_ARCH_TYPE} Model",
                                    interactive=True,
                                )
                                self.vr_download_url = gr.Markdown()
                                self.vr_download_choice.change(
                                    lambda model: md_url(vr_models[model]) if model else "",
                                    inputs=self.vr_download_choice,
                                    outputs=self.vr_download_url,
                                )
                            with gr.Row(variant="panel"):
                                mdx_models = self.models_url[MDX_ARCH_TYPE]
                                self.mdx_download_choice = gr.Dropdown(
                                    choices=list(mdx_models.keys()),
                                    label=f"Select {MDX_ARCH_TYPE} Model",
                                    interactive=True,
                                )
                                self.mdx_download_url = gr.Markdown()
                                self.mdx_download_choice.change(
                                    lambda model: md_url(mdx_models[model]) if model else "",
                                    inputs=self.mdx_download_choice,
                                    outputs=self.mdx_download_url,
                                )
                            with gr.Row(variant="panel"):
                                demucs_models: Dict[str, Dict] = self.models_url[DEMUCS_ARCH_TYPE]
                                self.demucs_download_choice = gr.Dropdown(
                                    choices=list(demucs_models.keys()),
                                    label=f"Select {DEMUCS_ARCH_TYPE} Model",
                                    interactive=True,
                                )
                                self.demucs_download_url = gr.Markdown()
                                self.demucs_download_choice.change(
                                    lambda model: "\n".join(
                                        [
                                            "- " + md_url(url, text=filename)
                                            for filename, url in demucs_models[model].items()
                                        ]
                                    ) if model else "",
                                    inputs=self.demucs_download_choice,
                                    outputs=self.demucs_download_url,
                                )

            self.arch_choice.change(
                self.arch_select_update,
                inputs=self.arch_choice,
                outputs=[self.model_choice, self.arch_setting1, self.arch_setting2],
            )
            self.model_choice.change(
                self.model_select_update,
                inputs=[self.arch_choice, self.model_choice],
                outputs=[
                    self.primary_stem_only,
                    self.secondary_stem_only,
                    self.primary_stem_out,
                    self.secondary_stem_out,
                ],
            )

            self.checkbox_set_root_value(self.use_gpu, "is_gpu_conversion_var")
            self.checkbox_set_root_value(self.sample_mode, "model_sample_mode_var")
            self.set_checkboxes_exclusive(
                [self.primary_stem_only, self.secondary_stem_only],
                [
                    lambda value: root.is_primary_stem_only_var.set(value),
                    lambda value: root.is_secondary_stem_only_var.set(value),
                ],
            )

            self.process_submit.click(
                self.process,
                inputs=[
                    self.audio_in,
                    self.input_filename,
                    self.model_choice,
                    self.arch_choice,
                    self.arch_setting1,
                    self.arch_setting2,
                ],
                outputs=[
                    self.primary_stem_out,
                    self.secondary_stem_out,
                    self.out_message,
                ],
            )

    def launch(self, **kwargs):
        kwargs.setdefault("ssr_mode", False)
        self.app.queue().launch(**kwargs)


uvr = UVRInterface()
uvr.cached_sources_clear()

webui = UVRWebUI(uvr, online_data_path="models/download_checks.json")
webui.launch()