main/inference/separator_music.py

import os
import sys
import time
import logging
import argparse
import logging.handlers

import soundfile as sf
import noisereduce as nr

from pydub import AudioSegment
from distutils.util import strtobool

now_dir = os.getcwd()
sys.path.append(now_dir)

from main.configs.config import Config
from main.library.algorithm.separator import Separator

translations = Config().translations

log_file = os.path.join("assets", "logs", "separator.log")
logger = logging.getLogger(__name__)

if logger.hasHandlers(): logger.handlers.clear()
else:
    console_handler = logging.StreamHandler()
    console_formatter = logging.Formatter(fmt="\n%(asctime)s.%(msecs)03d | %(levelname)s | %(module)s | %(message)s", datefmt="%Y-%m-%d %H:%M:%S")

    console_handler.setFormatter(console_formatter)
    console_handler.setLevel(logging.INFO)

    file_handler = logging.handlers.RotatingFileHandler(log_file, maxBytes=5*1024*1024, backupCount=3, encoding='utf-8')
    file_formatter = logging.Formatter(fmt="\n%(asctime)s.%(msecs)03d | %(levelname)s | %(module)s | %(message)s", datefmt="%Y-%m-%d %H:%M:%S")

    file_handler.setFormatter(file_formatter)
    file_handler.setLevel(logging.DEBUG)

    logger.addHandler(console_handler)
    logger.addHandler(file_handler)
    logger.setLevel(logging.DEBUG)

demucs_models = {
    "HT-Tuned": "htdemucs_ft.yaml",
    "HT-Normal": "htdemucs.yaml",
    "HD_MMI": "hdemucs_mmi.yaml",
    "HT_6S": "htdemucs_6s.yaml"
}

mdx_models = {
    "Main_340": "UVR-MDX-NET_Main_340.onnx",
    "Main_390": "UVR-MDX-NET_Main_390.onnx",
    "Main_406": "UVR-MDX-NET_Main_406.onnx",
    "Main_427": "UVR-MDX-NET_Main_427.onnx",
    "Main_438": "UVR-MDX-NET_Main_438.onnx",
    "Inst_full_292": "UVR-MDX-NET-Inst_full_292.onnx",
    "Inst_HQ_1": "UVR-MDX-NET_Inst_HQ_1.onnx",
    "Inst_HQ_2": "UVR-MDX-NET_Inst_HQ_2.onnx",
    "Inst_HQ_3": "UVR-MDX-NET_Inst_HQ_3.onnx",
    "Inst_HQ_4": "UVR-MDX-NET-Inst_HQ_4.onnx",
    "Kim_Vocal_1": "Kim_Vocal_1.onnx",
    "Kim_Vocal_2": "Kim_Vocal_2.onnx",
    "Kim_Inst": "Kim_Inst.onnx",
    "Inst_187_beta": "UVR-MDX-NET_Inst_187_beta.onnx",
    "Inst_82_beta": "UVR-MDX-NET_Inst_82_beta.onnx",
    "Inst_90_beta": "UVR-MDX-NET_Inst_90_beta.onnx",
    "Voc_FT": "UVR-MDX-NET-Voc_FT.onnx",
    "Crowd_HQ": "UVR-MDX-NET_Crowd_HQ_1.onnx",
    "MDXNET_9482": "UVR_MDXNET_9482.onnx",
    "Inst_1": "UVR-MDX-NET-Inst_1.onnx",
    "Inst_2": "UVR-MDX-NET-Inst_2.onnx",
    "Inst_3": "UVR-MDX-NET-Inst_3.onnx",
    "MDXNET_1_9703": "UVR_MDXNET_1_9703.onnx",
    "MDXNET_2_9682": "UVR_MDXNET_2_9682.onnx",
    "MDXNET_3_9662": "UVR_MDXNET_3_9662.onnx",
    "Inst_Main": "UVR-MDX-NET-Inst_Main.onnx",
    "MDXNET_Main": "UVR_MDXNET_Main.onnx"
}

kara_models = {
    "Version-1": "UVR_MDXNET_KARA.onnx",
    "Version-2": "UVR_MDXNET_KARA_2.onnx"
}


def parse_arguments() -> tuple:
    parser = argparse.ArgumentParser()
    parser.add_argument("--input_path", type=str, required=True)
    parser.add_argument("--output_path", type=str, default="./audios")
    parser.add_argument("--format", type=str, default="wav")
    parser.add_argument("--shifts", type=int, default=10)
    parser.add_argument("--segments_size", type=int, default=256)
    parser.add_argument("--overlap", type=float, default=0.25)
    parser.add_argument("--mdx_hop_length", type=int, default=1024)
    parser.add_argument("--mdx_batch_size", type=int, default=1)
    parser.add_argument("--clean_audio", type=lambda x: bool(strtobool(x)), default=False)
    parser.add_argument("--clean_strength", type=float, default=0.7)
    parser.add_argument("--backing_denoise", type=lambda x: bool(strtobool(x)), default=False)
    parser.add_argument("--demucs_model", type=str, default="HT-Normal")
    parser.add_argument("--kara_model", type=str, default="Version-1")
    parser.add_argument("--mdx_model", type=str, default="Main_340")
    parser.add_argument("--backing", type=lambda x: bool(strtobool(x)), default=False)
    parser.add_argument("--mdx", type=lambda x: bool(strtobool(x)), default=False)
    parser.add_argument("--mdx_denoise", type=lambda x: bool(strtobool(x)), default=False)
    parser.add_argument("--reverb", type=lambda x: bool(strtobool(x)), default=False)
    parser.add_argument("--reverb_denoise", type=lambda x: bool(strtobool(x)), default=False)
    parser.add_argument("--backing_reverb", type=lambda x: bool(strtobool(x)), default=False)

    args = parser.parse_args()
    return args

def main():
    start_time = time.time()

    try:
        args = parse_arguments()

        input_path = args.input_path
        output_path = args.output_path
        export_format = args.format
        shifts = args.shifts
        segments_size = args.segments_size
        overlap = args.overlap
        hop_length = args.mdx_hop_length
        batch_size = args.mdx_batch_size
        clean_audio = args.clean_audio
        clean_strength = args.clean_strength
        backing_denoise = args.backing_denoise
        demucs_model = args.demucs_model
        kara_model = args.kara_model
        backing = args.backing
        mdx = args.mdx
        mdx_model = args.mdx_model
        mdx_denoise = args.mdx_denoise
        reverb = args.reverb
        reverb_denoise = args.reverb_denoise
        backing_reverb = args.backing_reverb

        if backing_reverb and not reverb: 
            logger.warning(translations["turn_on_dereverb"])
            sys.exit(1)

        if backing_reverb and not backing: 
            logger.warning(translations["turn_on_separator_backing"])
            sys.exit(1)

        logger.debug(f"{translations['audio_path']}: {input_path}")
        logger.debug(f"{translations['output_path']}: {output_path}")
        logger.debug(f"{translations['export_format']}: {export_format}")
        if not mdx: logger.debug(f"{translations['shift']}: {shifts}") 
        logger.debug(f"{translations['segments_size']}: {segments_size}")
        logger.debug(f"{translations['overlap']}: {overlap}")
        if clean_audio: logger.debug(f"{translations['clear_audio']}: {clean_audio}")
        if clean_audio: logger.debug(f"{translations['clean_strength']}: {clean_strength}")
        if not mdx: logger.debug(f"{translations['demucs_model']}: {demucs_model}")
        if backing: logger.debug(f"{translations['denoise_backing']}: {backing_denoise}")
        if backing: logger.debug(f"{translations['backing_model_ver']}: {kara_model}")
        if backing: logger.debug(f"{translations['separator_backing']}: {backing}")
        if mdx: logger.debug(f"{translations['use_mdx']}: {mdx}")
        if mdx: logger.debug(f"{translations['mdx_model']}: {mdx_model}")
        if mdx: logger.debug(f"{translations['denoise_mdx']}: {mdx_denoise}")
        if mdx or backing or reverb: logger.debug(f"Hop length: {hop_length}")
        if mdx or backing or reverb: logger.debug(f"{translations['batch_size']}: {batch_size}")
        if reverb: logger.debug(f"{translations['dereveb_audio']}: {reverb}")
        if reverb: logger.debug(f"{translations['denoise_dereveb']}: {reverb_denoise}")
        if reverb: logger.debug(f"{translations['dereveb_backing']}: {backing_reverb}")

        if not mdx: vocals, instruments = separator_music_demucs(input_path, output_path, export_format, shifts, overlap, segments_size, demucs_model)
        else: vocals, instruments = separator_music_mdx(input_path, output_path, export_format, segments_size, overlap, mdx_denoise, mdx_model, hop_length, batch_size)

        if backing: main_vocals, backing_vocals = separator_backing(vocals, output_path, export_format, segments_size, overlap, backing_denoise, kara_model, hop_length, batch_size)
        if reverb: vocals_no_reverb, main_vocals_no_reverb, backing_vocals_no_reverb = separator_reverb(output_path, export_format, segments_size, overlap, reverb_denoise, reverb, backing, backing_reverb, hop_length, batch_size)

        original_output = os.path.join(output_path, f"Original_Vocals_No_Reverb.{export_format}") if reverb else os.path.join(output_path, f"Original_Vocals.{export_format}")
        main_output = os.path.join(output_path, f"Main_Vocals_No_Reverb.{export_format}") if reverb and backing else os.path.join(output_path, f"Main_Vocals.{export_format}")
        backing_output = os.path.join(output_path, f"Backing_Vocals_No_Reverb.{export_format}") if reverb and backing_reverb else os.path.join(output_path, f"Backing_Vocals.{export_format}")
        
        if clean_audio:
            logger.info(f"{translations['clear_audio']}...")
            vocal_data, vocal_sr = sf.read(vocals_no_reverb if reverb else vocals)
            main_data, main_sr = sf.read(main_vocals_no_reverb if reverb and backing else main_vocals)
            backing_data, backing_sr = sf.read(backing_vocals_no_reverb if reverb and backing_reverb else backing_vocals)

            vocals_clean = nr.reduce_noise(y=vocal_data, prop_decrease=clean_strength)
            
            sf.write(original_output, vocals_clean, vocal_sr, format=export_format)

            if backing:
                mains_clean = nr.reduce_noise(y=main_data, sr=main_sr, prop_decrease=clean_strength)
                backing_clean = nr.reduce_noise(y=backing_data, sr=backing_sr, prop_decrease=clean_strength)
                sf.write(main_output, mains_clean, main_sr, format=export_format)
                sf.write(backing_output, backing_clean, backing_sr, format=export_format)          

            logger.info(translations["clean_audio_success"])
            return original_output, instruments, main_output, backing_output
    except Exception as e:
        logger.error(f"{translations['separator_error']}: {e}")
        return None, None, None, None
    
    elapsed_time = time.time() - start_time
    logger.info(translations["separator_success"].format(elapsed_time=f"{elapsed_time:.2f}"))
    

def separator_music_demucs(input, output, format, shifts, overlap, segments_size, demucs_model):
    if not os.path.exists(input): 
        logger.warning(translations["input_not_valid"])
        sys.exit(1)
    
    if not os.path.exists(output): 
        logger.warning(translations["output_not_valid"])
        sys.exit(1)
    
    for i in [f"Original_Vocals.{format}", f"Instruments.{format}"]:
        if os.path.exists(os.path.join(output, i)): os.remove(os.path.join(output, i))

    model = demucs_models.get(demucs_model)

    segment_size = segments_size / 2

    logger.info(f"{translations['separator_process_2']}...")

    demucs_output = separator_main(audio_file=input, model_filename=model, output_format=format, output_dir=output, demucs_segment_size=segment_size, demucs_shifts=shifts, demucs_overlap=overlap)
    
    for f in demucs_output:
        path = os.path.join(output, f)

        if not os.path.exists(path): logger.error(translations["not_found"].format(name=path))

        if '_(Drums)_' in f: drums = path
        elif '_(Bass)_' in f: bass = path
        elif '_(Other)_' in f: other = path
        elif '_(Vocals)_' in f: os.rename(path, os.path.join(output, f"Original_Vocals.{format}"))

    AudioSegment.from_file(drums).overlay(AudioSegment.from_file(bass)).overlay(AudioSegment.from_file(other)).export(os.path.join(output, f"Instruments.{format}"), format=format)

    for f in [drums, bass, other]:
        if os.path.exists(f): os.remove(f)
    
    logger.info(translations["separator_success_2"])
    return os.path.join(output, f"Original_Vocals.{format}"), os.path.join(output, f"Instruments.{format}")

def separator_backing(input, output, format, segments_size, overlap, denoise, kara_model, hop_length, batch_size):
    if not os.path.exists(input): 
        logger.warning(translations["input_not_valid"])
        sys.exit(1)
    
    if not os.path.exists(output): 
        logger.warning(translations["output_not_valid"])
        sys.exit(1)
    
    for f in [f"Main_Vocals.{format}", f"Backing_Vocals.{format}"]:
        if os.path.exists(os.path.join(output, f)): os.remove(os.path.join(output, f))

    model_2 = kara_models.get(kara_model)

    logger.info(f"{translations['separator_process_backing']}...")
    backing_outputs = separator_main(audio_file=input, model_filename=model_2, output_format=format, output_dir=output, mdx_segment_size=segments_size, mdx_overlap=overlap, mdx_batch_size=batch_size, mdx_hop_length=hop_length, mdx_enable_denoise=denoise)

    main_output = os.path.join(output, f"Main_Vocals.{format}")
    backing_output = os.path.join(output, f"Backing_Vocals.{format}")

    for f in backing_outputs:
        path = os.path.join(output, f)

        if not os.path.exists(path): logger.error(translations["not_found"].format(name=path))

        if '_(Instrumental)_' in f: os.rename(path, backing_output)
        elif '_(Vocals)_' in f: os.rename(path, main_output)

    logger.info(translations["separator_process_backing_success"])
    return main_output, backing_output

def separator_music_mdx(input, output, format, segments_size, overlap, denoise, mdx_model, hop_length, batch_size):
    if not os.path.exists(input): 
        logger.warning(translations["input_not_valid"])
        sys.exit(1)
    
    if not os.path.exists(output): 
        logger.warning(translations["output_not_valid"])
        sys.exit(1)

    for i in [f"Original_Vocals.{format}", f"Instruments.{format}"]:
        if os.path.exists(os.path.join(output, i)): os.remove(os.path.join(output, i))
    
    model_3 = mdx_models.get(mdx_model)

    logger.info(f"{translations['separator_process_2']}...")
    output_music = separator_main(audio_file=input, model_filename=model_3, output_format=format, output_dir=output, mdx_segment_size=segments_size, mdx_overlap=overlap, mdx_batch_size=batch_size, mdx_hop_length=hop_length, mdx_enable_denoise=denoise)
    
    original_output = os.path.join(output, f"Original_Vocals.{format}")
    instruments_output = os.path.join(output, f"Instruments.{format}")

    for f in output_music:
        path = os.path.join(output, f)

        if not os.path.exists(path): logger.error(translations["not_found"].format(name=path))

        if '_(Instrumental)_' in f: os.rename(path, instruments_output)
        elif '_(Vocals)_' in f: os.rename(path, original_output)

    logger.info(translations["separator_process_backing_success"])
    return original_output, instruments_output

def separator_reverb(output, format, segments_size, overlap, denoise, original, main, backing_reverb, hop_length, batch_size):
    if not os.path.exists(output): 
        logger.warning(translations["output_not_valid"])
        sys.exit(1)
    
    for i in [f"Original_Vocals_Reverb.{format}", f"Main_Vocals_Reverb.{format}", f"Original_Vocals_No_Reverb.{format}", f"Main_Vocals_No_Reverb.{format}"]:
        if os.path.exists(os.path.join(output, i)): os.remove(os.path.join(output, i))

    dereveb_path = []

    if original: 
        try:
            dereveb_path.append(os.path.join(output, [f for f in os.listdir(output) if 'Original_Vocals' in f][0]))
        except IndexError:
            logger.warning(translations["not_found_original_vocal"])
            sys.exit(1)
        
    if main:
        try:
            dereveb_path.append(os.path.join(output, [f for f in os.listdir(output) if 'Main_Vocals' in f][0]))
        except IndexError:
            logger.warning(translations["not_found_main_vocal"])
            sys.exit(1)
    
    if backing_reverb:
        try:
            dereveb_path.append(os.path.join(output, [f for f in os.listdir(output) if 'Backing_Vocals' in f][0]))
        except IndexError:
            logger.warning(translations["not_found_backing_vocal"])
            sys.exit(1)
        
    for path in dereveb_path:
        if not os.path.exists(path): 
            logger.warning(translations["not_found"].format(name=path))
            sys.exit(1)
        
        if "Original_Vocals" in path: 
            reverb_path = os.path.join(output, f"Original_Vocals_Reverb.{format}")
            no_reverb_path = os.path.join(output, f"Original_Vocals_No_Reverb.{format}")
            start_title = translations["process_original"]
            end_title = translations["process_original_success"]
        elif "Main_Vocals" in path:
            reverb_path = os.path.join(output, f"Main_Vocals_Reverb.{format}")
            no_reverb_path = os.path.join(output, f"Main_Vocals_No_Reverb.{format}")
            start_title = translations["process_main"]
            end_title = translations["process_main_success"]
        elif "Backing_Vocals" in path:
            reverb_path = os.path.join(output, f"Backing_Vocals_Reverb.{format}")
            no_reverb_path = os.path.join(output, f"Backing_Vocals_No_Reverb.{format}")
            start_title = translations["process_backing"]
            end_title = translations["process_backing_success"]

        logger.info(start_title)
        output_dereveb = separator_main(audio_file=path, model_filename="Reverb_HQ_By_FoxJoy.onnx", output_format=format, output_dir=output, mdx_segment_size=segments_size, mdx_overlap=overlap, mdx_batch_size=batch_size, mdx_hop_length=hop_length, mdx_enable_denoise=denoise)

        for f in output_dereveb:
            path = os.path.join(output, f)

            if not os.path.exists(path): logger.error(translations["not_found"].format(name=path))

            if '_(Reverb)_' in f: os.rename(path, reverb_path)
            elif '_(No Reverb)_' in f: os.rename(path, no_reverb_path)

        logger.info(end_title)

    return (os.path.join(output, f"Original_Vocals_No_Reverb.{format}") if original else None), (os.path.join(output, f"Main_Vocals_No_Reverb.{format}") if main else None), (os.path.join(output, f"Backing_Vocals_No_Reverb.{format}") if backing_reverb else None)

def separator_main(audio_file=None, model_filename="UVR-MDX-NET_Main_340.onnx", output_format="wav", output_dir=".", mdx_segment_size=256, mdx_overlap=0.25, mdx_batch_size=1, mdx_hop_length=1024, mdx_enable_denoise=True, demucs_segment_size=256, demucs_shifts=2, demucs_overlap=0.25):
    separator = Separator(
        log_formatter=file_formatter,
        log_level=logging.INFO,
        output_dir=output_dir,
        output_format=output_format,
        output_bitrate=None,
        normalization_threshold=0.9,
        output_single_stem=None,
        invert_using_spec=False,
        sample_rate=44100,
        mdx_params={
            "hop_length": mdx_hop_length,
            "segment_size": mdx_segment_size,
            "overlap": mdx_overlap,
            "batch_size": mdx_batch_size,
            "enable_denoise": mdx_enable_denoise,
        },
        demucs_params={
            "segment_size": demucs_segment_size,
            "shifts": demucs_shifts,
            "overlap": demucs_overlap,
            "segments_enabled": True,
        }
    )

    separator.load_model(model_filename=model_filename)
    return separator.separate(audio_file)

if __name__ == "__main__": main()