Rifd/Ultimate-Vocal-Remover-WebUI
5
1import os2import json3 4import librosa5import soundfile6import numpy as np7 8import gradio as gr9from UVR_interface import root, UVRInterface, VR_MODELS_DIR, MDX_MODELS_DIR, DEMUCS_MODELS_DIR10from gui_data.constants import *11from typing import List, Dict, Callable, Union12 13 14class UVRWebUI:15 def __init__(self, uvr: UVRInterface, online_data_path: str) -> None:16 self.uvr = uvr17 self.models_url = self.get_models_url(online_data_path)18 self.define_layout()19 20 self.input_temp_dir = "__temp"21 self.export_path = "out"22 if not os.path.exists(self.input_temp_dir):23 os.mkdir(self.input_temp_dir)24 25 def get_models_url(self, models_info_path: str) -> Dict[str, Dict]:26 with open(models_info_path, "r") as f:27 online_data = json.loads(f.read())28 models_url = {}29 for arch, download_list_key in zip([VR_ARCH_TYPE, MDX_ARCH_TYPE], ["vr_download_list", "mdx_download_list"]):30 models_url[arch] = {model: NORMAL_REPO+model_path for model, model_path in online_data[download_list_key].items()}31 models_url[DEMUCS_ARCH_TYPE] = online_data["demucs_download_list"]32 return models_url33 34 def get_local_models(self, arch: str) -> List[str]:35 model_config = {36 VR_ARCH_TYPE: (VR_MODELS_DIR, ".pth"),37 MDX_ARCH_TYPE: (MDX_MODELS_DIR, ".onnx"),38 DEMUCS_ARCH_TYPE: (DEMUCS_MODELS_DIR, ".yaml"),39 }40 try:41 model_dir, suffix = model_config[arch]42 except KeyError:43 raise ValueError(f"Unkown arch type: {arch}")44 return [os.path.splitext(f)[0] for f in os.listdir(model_dir) if f.endswith(suffix)]45 46 def set_arch_setting_value(self, arch: str, setting1, setting2):47 if arch == VR_ARCH_TYPE:48 root.window_size_var.set(setting1)49 root.aggression_setting_var.set(setting2)50 elif arch == MDX_ARCH_TYPE:51 root.mdx_batch_size_var.set(setting1)52 root.compensate_var.set(setting2)53 elif arch == DEMUCS_ARCH_TYPE:54 pass55 56 def arch_select_update(self, arch: str) -> List[Dict]:57 choices = self.get_local_models(arch)58 if arch == VR_ARCH_TYPE:59 model_update = self.model_choice.update(choices=choices, value=CHOOSE_MODEL, label=SELECT_VR_MODEL_MAIN_LABEL)60 setting1_update = self.arch_setting1.update(choices=VR_WINDOW, label=WINDOW_SIZE_MAIN_LABEL, value=root.window_size_var.get())61 setting2_update = self.arch_setting2.update(choices=VR_AGGRESSION, label=AGGRESSION_SETTING_MAIN_LABEL, value=root.aggression_setting_var.get())62 elif arch == MDX_ARCH_TYPE:63 model_update = self.model_choice.update(choices=choices, value=CHOOSE_MODEL, label=CHOOSE_MDX_MODEL_MAIN_LABEL)64 setting1_update = self.arch_setting1.update(choices=BATCH_SIZE, label=BATCHES_MDX_MAIN_LABEL, value=root.mdx_batch_size_var.get())65 setting2_update = self.arch_setting2.update(choices=VOL_COMPENSATION, label=VOL_COMP_MDX_MAIN_LABEL, value=root.compensate_var.get())66 elif arch == DEMUCS_ARCH_TYPE:67 model_update = self.model_choice.update(choices=choices, value=CHOOSE_MODEL, label=CHOOSE_DEMUCS_MODEL_MAIN_LABEL)68 raise gr.Error(f"{DEMUCS_ARCH_TYPE} not implempted")69 else:70 raise gr.Error(f"Unkown arch type: {arch}")71 return [model_update, setting1_update, setting2_update]72 73 def model_select_update(self, arch: str, model_name: str) -> List[Union[str, Dict, None]]:74 if model_name == CHOOSE_MODEL:75 return [None for _ in range(4)]76 model, = self.uvr.assemble_model_data(model_name, arch)77 if not model.model_status:78 raise gr.Error(f"Cannot get model data, model hash = {model.model_hash}")79 80 stem1_check_update = self.primary_stem_only.update(label=f"{model.primary_stem} Only")81 stem2_check_update = self.secondary_stem_only.update(label=f"{model.secondary_stem} Only")82 stem1_out_update = self.primary_stem_out.update(label=f"Output {model.primary_stem}")83 stem2_out_update = self.secondary_stem_out.update(label=f"Output {model.secondary_stem}")84 85 return [stem1_check_update, stem2_check_update, stem1_out_update, stem2_out_update]86 87 def checkbox_set_root_value(self, checkbox: gr.Checkbox, root_attr: str):88 checkbox.change(lambda value: root.__getattribute__(root_attr).set(value), inputs=checkbox)89 90 def set_checkboxes_exclusive(self, checkboxes: List[gr.Checkbox], pure_callbacks: List[Callable], exclusive_value=True):91 def exclusive_onchange(i, callback_i):92 def new_onchange(*check_values):93 if check_values[i] == exclusive_value:94 return_values = []95 for j, value_j in enumerate(check_values):96 if j != i and value_j == exclusive_value:97 return_values.append(not exclusive_value)98 else:99 return_values.append(value_j)100 else:101 return_values = check_values102 callback_i(check_values[i])103 return return_values104 return new_onchange105 106 for i, (checkbox, callback) in enumerate(zip(checkboxes, pure_callbacks)):107 checkbox.change(exclusive_onchange(i, callback), inputs=checkboxes, outputs=checkboxes)108 109 def process(self, input_audio, input_filename, model_name, arch, setting1, setting2, progress=gr.Progress()):110 def set_progress_func(step, inference_iterations=0):111 progress_curr = step + inference_iterations112 progress(progress_curr)113 114 sampling_rate, audio = input_audio115 audio = (audio / np.iinfo(audio.dtype).max).astype(np.float32)116 if len(audio.shape) > 1:117 audio = librosa.to_mono(audio.transpose(1, 0))118 input_path = os.path.join(self.input_temp_dir, input_filename)119 soundfile.write(input_path, audio, sampling_rate, format="wav")120 121 self.set_arch_setting_value(arch, setting1, setting2)122 123 seperator = uvr.process(124 model_name=model_name,125 arch_type=arch,126 audio_file=input_path,127 export_path=self.export_path,128 is_model_sample_mode=root.model_sample_mode_var.get(),129 set_progress_func=set_progress_func,130 )131 132 primary_audio = None133 secondary_audio = None134 msg = ""135 if not seperator.is_secondary_stem_only:136 primary_stem_path = os.path.join(seperator.export_path, f"{seperator.audio_file_base}_({seperator.primary_stem}).wav")137 audio, rate = soundfile.read(primary_stem_path)138 primary_audio = (rate, audio)139 msg += f"{seperator.primary_stem} saved at {primary_stem_path}\n"140 if not seperator.is_primary_stem_only:141 secondary_stem_path = os.path.join(seperator.export_path, f"{seperator.audio_file_base}_({seperator.secondary_stem}).wav")142 audio, rate = soundfile.read(secondary_stem_path)143 secondary_audio = (rate, audio)144 msg += f"{seperator.secondary_stem} saved at {secondary_stem_path}\n"145 146 os.remove(input_path)147 148 return primary_audio, secondary_audio, msg149 150 def define_layout(self):151 with gr.Blocks() as app:152 self.app = app153 gr.HTML("<h1> 🎵 Ultimate Vocal Remover WebUI 🎵 </h1>")154 gr.Markdown("This is an experimental demo with CPU. Duplicate the space for use in private")155 gr.Markdown(156 "[](https://huggingface.co/spaces/r3gm/Ultimate-Vocal-Remover-WebUI?duplicate=true)\n\n"157 ) 158 with gr.Tabs():159 with gr.TabItem("process"):160 with gr.Row():161 self.arch_choice = gr.Dropdown(162 choices=[VR_ARCH_TYPE, MDX_ARCH_TYPE], value=VR_ARCH_TYPE, # choices=[VR_ARCH_TYPE, MDX_ARCH_TYPE, DEMUCS_ARCH_TYPE], value=VR_ARCH_TYPE,163 label=CHOOSE_PROC_METHOD_MAIN_LABEL, interactive=True)164 self.model_choice = gr.Dropdown(165 choices=self.get_local_models(VR_ARCH_TYPE), value=CHOOSE_MODEL,166 label=SELECT_VR_MODEL_MAIN_LABEL+' 👋Select a model', interactive=True)167 with gr.Row():168 self.arch_setting1 = gr.Dropdown(169 choices=VR_WINDOW, value=root.window_size_var.get(),170 label=WINDOW_SIZE_MAIN_LABEL+' 👋Select one', interactive=True)171 self.arch_setting2 = gr.Dropdown(172 choices=VR_AGGRESSION, value=root.aggression_setting_var.get(),173 label=AGGRESSION_SETTING_MAIN_LABEL, interactive=True)174 with gr.Row():175 self.use_gpu = gr.Checkbox(176 label='Rhythmic Transmutation Device', value=True, interactive=True) #label=GPU_CONVERSION_MAIN_LABEL, value=root.is_gpu_conversion_var.get(), interactive=True)177 self.primary_stem_only = gr.Checkbox(178 label=f"{PRIMARY_STEM} only", value=root.is_primary_stem_only_var.get(), interactive=True)179 self.secondary_stem_only = gr.Checkbox(180 label=f"{SECONDARY_STEM} only", value=root.is_secondary_stem_only_var.get(), interactive=True)181 self.sample_mode = gr.Checkbox(182 label=SAMPLE_MODE_CHECKBOX(root.model_sample_mode_duration_var.get()),183 value=root.model_sample_mode_var.get(), interactive=True)184 185 with gr.Row():186 self.input_filename = gr.Textbox(label="Input filename", value="temp.wav", interactive=True)187 with gr.Row():188 self.audio_in = gr.Audio(label="Input audio", interactive=True)189 with gr.Row():190 self.process_submit = gr.Button(START_PROCESSING, variant="primary")191 with gr.Row():192 self.primary_stem_out = gr.Audio(label=f"Output {PRIMARY_STEM}", interactive=False)193 self.secondary_stem_out = gr.Audio(label=f"Output {SECONDARY_STEM}", interactive=False)194 with gr.Row():195 self.out_message = gr.Textbox(label="Output Message", interactive=False, show_progress=False)196 197 with gr.TabItem("settings"):198 with gr.Tabs():199 with gr.TabItem("Settings Guide"):200 pass201 with gr.TabItem("Additional Settigns"):202 self.wav_type = gr.Dropdown(choices=WAV_TYPE, label="Wav Type", value="PCM_16", interactive=True)203 self.mp3_rate = gr.Dropdown(choices=MP3_BIT_RATES, label="MP3 Bitrate", value="320k",interactive=True)204 with gr.TabItem("Download models"):205 206 def md_url(url, text=None):207 if text is None:208 text = url209 return f"[{url}]({url})"210 211 with gr.Row():212 vr_models = self.models_url[VR_ARCH_TYPE]213 self.vr_download_choice = gr.Dropdown(choices=list(vr_models.keys()), label=f"Select {VR_ARCH_TYPE} Model", interactive=True)214 self.vr_download_url = gr.Markdown()215 self.vr_download_choice.change(lambda model: md_url(vr_models[model]), inputs=self.vr_download_choice, outputs=self.vr_download_url)216 with gr.Row(variant="panel"):217 mdx_models = self.models_url[MDX_ARCH_TYPE]218 self.mdx_download_choice = gr.Dropdown(choices=list(mdx_models.keys()), label=f"Select {MDX_ARCH_TYPE} Model", interactive=True)219 self.mdx_download_url = gr.Markdown()220 self.mdx_download_choice.change(lambda model: md_url(mdx_models[model]), inputs=self.mdx_download_choice, outputs=self.mdx_download_url)221 with gr.Row(variant="panel"):222 demucs_models: Dict[str, Dict] = self.models_url[DEMUCS_ARCH_TYPE]223 self.demucs_download_choice = gr.Dropdown(choices=list(demucs_models.keys()), label=f"Select {DEMUCS_ARCH_TYPE} Model", interactive=True)224 self.demucs_download_url = gr.Markdown()225 226 self.demucs_download_choice.change(227 lambda model: "\n".join([228 "- " + md_url(url, text=filename) for filename, url in demucs_models[model].items()]),229 inputs=self.demucs_download_choice,230 outputs=self.demucs_download_url)231 232 self.arch_choice.change(233 self.arch_select_update, inputs=self.arch_choice,234 outputs=[self.model_choice, self.arch_setting1, self.arch_setting2])235 self.model_choice.change(236 self.model_select_update, inputs=[self.arch_choice, self.model_choice],237 outputs=[self.primary_stem_only, self.secondary_stem_only, self.primary_stem_out, self.secondary_stem_out])238 239 self.checkbox_set_root_value(self.use_gpu, 'is_gpu_conversion_var')240 self.checkbox_set_root_value(self.sample_mode, 'model_sample_mode_var')241 self.set_checkboxes_exclusive(242 [self.primary_stem_only, self.secondary_stem_only],243 [lambda value: root.is_primary_stem_only_var.set(value), lambda value: root.is_secondary_stem_only_var.set(value)])244 245 self.process_submit.click(246 self.process,247 inputs=[self.audio_in, self.input_filename, self.model_choice, self.arch_choice, self.arch_setting1, self.arch_setting2],248 outputs=[self.primary_stem_out, self.secondary_stem_out, self.out_message])249 250 def launch(self, **kwargs):251 self.app.queue().launch(**kwargs)252 253 254uvr = UVRInterface()255uvr.cached_sources_clear()256 257webui = UVRWebUI(uvr, online_data_path='models/download_checks.json')258 259 260print(webui.models_url)261model_dict = webui.models_url262 263import os264import wget265 266for category, models in model_dict.items():267 if category in ['VR Arc', 'MDX-Net']:268 if category == 'VR Arc':269 model_path = 'models/VR_Models'270 elif category == 'MDX-Net':271 model_path = 'models/MDX_Net_Models'272 273 for model_name, model_url in models.items():274 cmd = f"aria2c --optimize-concurrent-downloads --console-log-level=error --summary-interval=10 -j5 -x16 -s16 -k1M -c -d {model_path} -Z {model_url}"275 os.system(cmd)276 277 print("Models downloaded successfully.")278 else:279 print(f"Ignoring category: {category}")280 281 282 283 284webui = UVRWebUI(uvr, online_data_path='models/download_checks.json')285webui.launch()286