uvr5 / inst.py
lorneluo's picture
add inst.py
60c7461
raw history blame
No virus
4.53 kB
import hashlib
import os
import shutil
import sys
from datetime import datetime
from pathlib import Path
import click
import yaml
from ml_collections import ConfigDict
from UVR import ModelData, AudioTools, Ensembler, DENOISER_MODEL_PATH, DEVERBER_MODEL_PATH, MDX_MODELS_DIR, \
MDX_MIXER_PATH, load_model_hash_data, MDX_MODEL_NAME_SELECT, model_hash_table, MDX_HASH_DIR, MDX_C_CONFIG_PATH
from args import mdx23c_8kfft_instvoc_hq_process_data, htdemucs_ft_process_data, uvr_mdx_net_voc_ft_process_data
from download import download_model, get_model_file
from gui_data.constants import VR_ARCH_TYPE, MDX_ARCH_TYPE, DEMUCS_ARCH_TYPE, ENSEMBLE_MODE, TIME_STRETCH, \
MANUAL_ENSEMBLE, MATCH_INPUTS, ALIGN_INPUTS, ALL_STEMS, DEFAULT, VOCAL_STEM, MP3_BIT_RATES, WAV, DEMUCS_2_SOURCE, \
DEMUCS_2_SOURCE_MAPPER, INST_STEM, CKPT, ONNX, MDX_POP_NFFT, secondary_stem, PRIMARY_STEM, SECONDARY_STEM
from lib_v5 import spec_utils
from separate import (
SeperateDemucs, SeperateMDX, SeperateMDXC, SeperateVR, # Model-related
save_format, clear_gpu_cache, # Utility functions
cuda_available, mps_available, # directml_available,
)
def run_ensemble_models(audio_path, export_path, format=WAV, clean=True):
start = datetime.now()
process_datas = [mdx23c_8kfft_instvoc_hq_process_data, uvr_mdx_net_voc_ft_process_data,
htdemucs_ft_process_data]
# download models
for process_data in process_datas:
download_model(process_data['model_name'])
# create folder
os.makedirs(export_path, exist_ok=True)
temp_export_path = os.path.join(export_path, 'uvr5_' + datetime.now().strftime("%Y-%m-%d %H:%M:%S"))
os.makedirs(temp_export_path, exist_ok=True)
print(f'temp_export_path', temp_export_path)
instrumental_export_paths = []
vocals_export_paths = []
for process_data in process_datas:
current_model = process_data['model_data']
audio_file_base = Path(audio_path).stem + '_' + current_model.model_basename
process_data['export_path'] = temp_export_path
process_data['audio_file_base'] = audio_file_base
process_data['audio_file'] = audio_path
if current_model.process_method == VR_ARCH_TYPE:
seperator = SeperateVR(current_model, process_data)
elif current_model.process_method == MDX_ARCH_TYPE:
seperator = SeperateMDXC(current_model, process_data) if current_model.is_mdx_c else SeperateMDX(
current_model, process_data)
elif current_model.process_method == DEMUCS_ARCH_TYPE:
seperator = SeperateDemucs(current_model, process_data, vocal_stem_path=(audio_path, audio_file_base))
else:
raise Exception(f'model not found')
seperator.seperate()
instrumental_path = Path(temp_export_path) / f"{audio_file_base}_(Instrumental).{format.lower()}"
vocals_path = Path(temp_export_path) / f"{audio_file_base}_(Vocals).{format.lower()}"
instrumental_export_paths.append(str(instrumental_path))
vocals_export_paths.append(str(vocals_path))
# merge each model outputs
vocals_final_path = Path(export_path) / f"{Path(audio_path).stem}.vocal.{format.lower()}"
instrumental_final_path = Path(export_path) / f"{Path(audio_path).stem}.instrumental.{format.lower()}"
ensemble(vocals_export_paths, vocals_final_path)
ensemble(instrumental_export_paths, instrumental_final_path)
if clean:
shutil.rmtree(temp_export_path)
print(f'instrumental_final_path', instrumental_final_path)
print(f'vocals_final_path', vocals_final_path)
print(f'Finished in {datetime.now() - start}')
return instrumental_final_path, vocals_final_path
def ensemble(stem_outputs, stem_save_path, format=WAV):
algorithm = 'Average'
is_normalization = True
spec_utils.ensemble_inputs(stem_outputs, algorithm, is_normalization, 'PCM_16', stem_save_path, is_wave=True)
save_format(stem_save_path, format, '320k')
# /Users/taoluo/Downloads/test/kimk_audio_MDX23C-8KFFT-InstVoc_HQ_(Instrumental).WAV
#
if __name__ == '__main__':
audio_file = '/Users/taoluo/Downloads/assets/audio/kimk_audio.mp3'
audio_file = sys.argv[1]
if not os.path.isfile(audio_file):
raise FileNotFoundError(f"File {audio_file} not exist.")
output_dir = os.path.dirname(audio_file)
print(output_dir)
run_ensemble_models(audio_file, output_dir)