uvr5/multiprocess_cuda_infer.py

import logging
import multiprocessing
from concurrent.futures import ProcessPoolExecutor
from random import shuffle
from tqdm import tqdm
import sys, wave
import torch, torchaudio
import hashlib
import time, os, psutil

# Make sure we resolve imports relative to this bundled copy of uvr5
THIS_FILE = os.path.abspath(__file__)
UVR5_ROOT = os.path.dirname(THIS_FILE)
if UVR5_ROOT not in sys.path:
    sys.path.append(UVR5_ROOT)

from gui_data.constants import *
from lib_v5.vr_network.model_param_init import ModelParameters
import argparse, json
import onnx
import onnxruntime as ort
import traceback
from datetime import datetime

logging.getLogger("numba").setLevel(logging.WARNING)
logging.getLogger("matplotlib").setLevel(logging.WARNING)

class ModelData():
    def __init__(self, 
                    model_path,
                    audio_path,
                    result_path,
                    process_method,
                    device, 
                    save_background=True,
                    is_pre_proc_model=False,
                    base_dir=UVR5_ROOT, 
                    **kwargs):
        self.__dict__.update(kwargs)

        BASE_PATH = result_path
        VR_MODELS_DIR = os.path.join(base_dir, 'models', 'VR_Models')
        VR_HASH_JSON = os.path.join(VR_MODELS_DIR, 'model_data', 'model_data.json')
        VR_PARAM_DIR = os.path.join(base_dir, 'lib_v5', 'vr_network', 'modelparams')
        SAMPLE_CLIP_PATH = os.path.join(BASE_PATH, 'temp_sample_clips')

        MDX_MIXER_PATH = os.path.join(base_dir, 'lib_v5', 'mixer.ckpt')
        # MDX_MODELS_DIR = os.path.join(base_dir, 'models', 'MDX_Net_Models')
        # MDX_HASH_DIR = (base_dir, 'models', 'MDX_Net_Models', 'model_data')
        MDX_HASH_JSON = os.path.join(base_dir, 'model_data.json')
        MDX_MODEL_NAME_SELECT = os.path.join(base_dir, 'model_name_mapper.json')

        self.model_name = self.model_name
        self.aggression_setting = float(int(self.aggression_setting)/100) # 1 - 20
        self.window_size = int(self.window_size)
        self.batch_size = int(self.batch_size) if self.batch_size.isdigit() else 1
        self.mdx_batch_size = 1 if self.mdx_batch_size == DEF_OPT else int(self.mdx_batch_size)
        self.is_mdx_ckpt = False
        self.crop_size = int(self.crop_size) 
        self.is_high_end_process = 'mirroring' if self.is_high_end_process else 'None'
        self.post_process_threshold = float(self.post_process_threshold)
        self.model_capacity = 32, 128
        self.model_path = model_path
        self.result_path = result_path
        self.model_basename = os.path.splitext(os.path.basename(self.model_path))[0]
        self.mixer_path = MDX_MIXER_PATH
        self.process_method = process_method
        self.is_pre_proc_model = is_pre_proc_model
        self.vr_is_secondary_model = self.vr_is_secondary_model_activate
        self.mdx_is_secondary_model = self.mdx_is_secondary_model_activate
        self.is_ensemble_mode = False
        self.secondary_model = None
        self.primary_model_primary_stem = None
        self.primary_stem = None
        self.secondary_stem = None
        self.secondary_model_scale = None
        self.is_demucs_pre_proc_model_inst_mix = False
        self.device = device
        self.save_background = save_background

        if type(audio_path)==str and os.path.isdir(audio_path):
            self.inputPaths = os.listdir(audio_path)
            self.inputPaths = [os.path.join(audio_path, x) for x in self.inputPaths if x[-4:]=='.wav']
        elif type(audio_path)==str and audio_path[-4:] == '.wav':
            self.inputPaths = [audio_path]
        elif type(audio_path) == list and audio_path[0][-4:] == '.wav':
            self.inputPaths = audio_path
        else:
            print(f"Invalid audio_path {audio_path}")

        self.get_model_hash()
        
        if self.process_method == VR_ARCH_TYPE:
            self.model_data = json.loads(open(VR_HASH_JSON, 'r', encoding='utf-8').read())[self.model_hash]
            if self.model_data:
                vr_model_param = os.path.join(VR_PARAM_DIR, "{}.json".format(self.model_data["vr_model_param"]))
                self.primary_stem = self.model_data["primary_stem"]
                self.secondary_stem = STEM_PAIR_MAPPER[self.primary_stem]
                self.vr_model_param = ModelParameters(vr_model_param)
                self.model_samplerate = self.vr_model_param.param['sr']
                if "nout" in self.model_data.keys() and "nout_lstm" in self.model_data.keys():
                    self.model_capacity = self.model_data["nout"], self.model_data["nout_lstm"]
                    self.is_vr_51_model = True
            else:
                self.model_status = False
        
        
        if self.process_method == MDX_ARCH_TYPE:
            self.is_vr_51_model = False
            self.margin = int(self.margin)
            self.model_samplerate = self.margin
            self.chunks = self.determine_auto_chunks(self.chunks) if self.is_chunk_mdxnet else 0
            self.model_data = json.loads(open(MDX_HASH_JSON, 'r', encoding='utf-8').read())[self.model_hash]
            if self.model_data:
                self.is_secondary_model = self.mdx_is_secondary_model
                self.compensate = self.model_data["compensate"]
                self.mdx_dim_f_set = self.model_data["mdx_dim_f_set"]
                self.mdx_dim_t_set = self.model_data["mdx_dim_t_set"]
                self.mdx_n_fft_scale_set = self.model_data["mdx_n_fft_scale_set"]
                self.primary_stem = self.model_data["primary_stem"]
                self.secondary_stem = STEM_PAIR_MAPPER[self.primary_stem]
            else:
                self.model_status = False


    def determine_auto_chunks(self, chunks):
        """Determines appropriate chunk size based on user computer specs"""
        gpu = 0 if torch.cuda.device_count() > 0 else -1
        if OPERATING_SYSTEM == 'Darwin':
            gpu = -1

        if chunks == BATCH_MODE:
            chunks = 0
            #self.chunks_var.set(AUTO_SELECT)

        if chunks == 'Full':
            chunk_set = 0
        elif chunks == 'Auto':
            if gpu == 0:
                gpu_mem = round(torch.cuda.get_device_properties(0).total_memory/1.074e+9)
                if gpu_mem <= int(6):
                    chunk_set = int(5)
                if gpu_mem in [7, 8, 9, 10, 11, 12, 13, 14, 15]:
                    chunk_set = int(10)
                if gpu_mem >= int(16):
                    chunk_set = int(40)
            if gpu == -1:
                sys_mem = psutil.virtual_memory().total >> 30
                if sys_mem <= int(4):
                    chunk_set = int(1)
                if sys_mem in [5, 6, 7, 8]:
                    chunk_set = int(10)
                if sys_mem in [9, 10, 11, 12, 13, 14, 15, 16]:
                    chunk_set = int(25)
                if sys_mem >= int(17):
                    chunk_set = int(60) 
        elif chunks == '0':
            chunk_set = 0
        else:
            chunk_set = int(chunks)
        print("chunks: ", gpu_mem, chunk_set)
        return chunk_set


    def get_model_hash(self):
        self.model_hash = None
        
        if not os.path.isfile(self.model_path):
            self.model_status = False
            self.model_hash is None
        else:
            if not self.model_hash:
                try:
                    with open(self.model_path, 'rb') as f:
                        f.seek(- 10000 * 1024, 2)
                        self.model_hash = hashlib.md5(f.read()).hexdigest()
                except:
                    self.model_hash = hashlib.md5(open(self.model_path,'rb').read()).hexdigest()


class Inference():
    def __init__(self, model_data: ModelData, device):
        self.device = device
        self.n_fft = model_data.mdx_n_fft_scale_set
        self.is_normalization = model_data.is_normalization
        self.compensate = model_data.compensate
        self.dim_f, self.dim_t = model_data.mdx_dim_f_set, 2**model_data.mdx_dim_t_set
        self.mdx_batch_size = model_data.mdx_batch_size
        self.is_denoise = model_data.is_denoise
        self.hop = 1024
        self.dim_c = 4
        self.chunks = model_data.chunks
        self.margin = model_data.margin
        self.adjust = 1
        self.progress_value = 0

        self.n_bins = self.n_fft//2+1
        self.trim = self.n_fft//2
        self.chunk_size = self.hop * (self.dim_t-1)
        self.window = torch.hann_window(window_length=self.n_fft, periodic=False).to(self.device)
        self.freq_pad = torch.zeros([1, self.dim_c, self.n_bins-self.dim_f, self.dim_t]).to(self.device)
        self.gen_size = self.chunk_size-2*self.trim
        self.save_background = model_data.save_background


    def stft(self, x):
        x = x.reshape([-1, self.chunk_size])
        x = torch.stft(x, n_fft=self.n_fft, hop_length=self.hop, window=self.window, center=True,return_complex=True)
        x=torch.view_as_real(x)
        x = x.permute([0,3,1,2])
        x = x.reshape([-1,2,2,self.n_bins,self.dim_t]).reshape([-1,self.dim_c,self.n_bins,self.dim_t])
        return x[:,:,:self.dim_f]

    def istft(self, x, freq_pad=None):
        freq_pad = self.freq_pad.repeat([x.shape[0],1,1,1]) if freq_pad is None else freq_pad
        x = torch.cat([x, freq_pad], -2)
        x = x.reshape([-1,2,2,self.n_bins,self.dim_t]).reshape([-1,2,self.n_bins,self.dim_t])
        x = x.permute([0,2,3,1])
        x=x.contiguous()
        x=torch.view_as_complex(x)
        x = torch.istft(x, n_fft=self.n_fft, hop_length=self.hop, window=self.window, center=True)
        return x.reshape([-1,2,self.chunk_size])


    def load_model(self, model_path, threads, device='cpu'):
        model = onnx.load_model(model_path)
        if torch.cuda.is_available() and device != 'cpu':
            providers = [("CUDAExecutionProvider", {"device_id": torch.cuda.current_device(),
                                                    "user_compute_stream": str(torch.cuda.current_stream().cuda_stream)})]
        else:
            providers = ["CPUExecutionProvider"]

        sess_options = ort.SessionOptions()
        sess_options.intra_op_num_threads = threads
        # sess_options.enable_profiling = True # debug 时开启
        self.ort_ = ort.InferenceSession(model.SerializeToString(), sess_options=sess_options, providers=providers)

        self.model_run = lambda spek:self.ort_.run(None, {'input': spek.cpu().numpy()})[0]


    def initialize_mix(self, mix):
        mix_waves = []
        n_sample = mix.shape[1]
        pad = self.gen_size - n_sample%self.gen_size
        zero_pad = torch.zeros((2,self.trim), device=mix.device)
        # print("mix:", mix.shape, mix.device, "zero_pad:", zero_pad.shape, zero_pad.device)
        mix_p = torch.cat((zero_pad, mix, torch.zeros((2,pad), device=mix.device), zero_pad), 1)
        i = 0
        while i < n_sample + pad:
            waves = mix_p[:, i:i+self.chunk_size]
            mix_waves.append(waves.unsqueeze(0))
            i += self.gen_size
            # print("debug 7:", i, waves, waves.shape, self.gen_size)
        mix_waves = torch.cat(mix_waves, 0).to(self.device)
        # print("debug 8:", mix_waves, mix_waves.shape, self.device, pad)
        return mix_waves, pad


    def run_model(self, mix, is_match_mix=False):
        
        spek = self.stft(mix.to(self.device))*self.adjust
        spek[:, :, :3, :] *= 0 
        # print("spek input:", spek.device, spek.shape)
        if is_match_mix:
            spec_pred = spek.to(self.device)
        else:
            spec_pred = -self.model_run(-spek)*0.5+self.model_run(spek)*0.5 if self.is_denoise else self.model_run(spek)
            spec_pred = torch.from_numpy(spec_pred).to(self.device)

        # print("is_denoise:", self.is_denoise, "spec_pred:", spec_pred.dtype, type(spec_pred))
        return self.istft(spec_pred).to(self.device)[:,:,self.trim:-self.trim].transpose(0,1).reshape(2, -1)


    def demix_base(self, mix, is_match_mix=False, device='cpu'):
        chunked_sources = []
        
        for slice in mix:
            # print("debug 6:", mix, slice, is_match_mix)
            sources = []
            tar_waves_ = []
            mix_p = mix[slice]
            # print("demix_base: ", mix_p.shape, mix_p.device)
            mix_waves, pad = self.initialize_mix(mix_p.to(device))
            mix_waves = mix_waves.split(self.mdx_batch_size)
            with torch.no_grad():
                for mix_wave in mix_waves:
                    # self.running_inference_progress_bar(len(mix)*len(mix_waves), is_match_mix=is_match_mix)
                    # print("debug10:", mix_wave, mix_wave.shape, is_match_mix)
                    tar_waves = self.run_model(mix_wave, is_match_mix=is_match_mix)
                    tar_waves_.append(tar_waves)

                tar_waves = torch.cat(tar_waves_, axis=-1)[:, :-pad]
                start = 0 if slice == 0 else self.margin
                end = None if slice == list(mix.keys())[::-1][0] or self.margin == 0 else -self.margin
                sources.append(tar_waves[:,start:end]*(1/self.adjust))
            chunked_sources = torch.cat(sources, axis=-1)
        # print("debug 11:",chunked_sources, len(chunked_sources), chunked_sources.shape)
        # sources = torch.cat(chunked_sources, axis=-1)
        sources = chunked_sources
        # print("debug 4:", sources, sources.shape)
        return sources

    def onnx_inference(self, wav_path, save_dir, device):
        start_time = time.time()
        input_audio, sr = torchaudio.load(wav_path, channels_first=True)
        input_audio = input_audio.to(device)
        # input_audio = input_audio.mean(dim=0).unsqueeze(0)  # stereo to mono
        if input_audio.shape[0] == 1:
            input_audio = torch.cat((input_audio, input_audio), 0) # mono to stereo
        if sr != 44100:
            input_audio = torchaudio.functional.resample(input_audio.squeeze(), sr, 44100)

        output_audio = self.demix_base({0:input_audio.squeeze()}, is_match_mix=False, device=device)
        torchaudio.save(
            os.path.join(save_dir, os.path.basename(wav_path).replace(".wav", "_vocal.wav")),
            output_audio.cpu(),
            44100,
        )

        if self.save_background:
            raw_mix = self.demix_base({0:input_audio.squeeze()}, is_match_mix=True)
            secondary_source, raw_mix = normalize_two_stem(output_audio*self.compensate, raw_mix, self.is_normalization)
            secondary_source = (-secondary_source+raw_mix)
            torchaudio.save(
                os.path.join(save_dir, os.path.basename(wav_path)).replace(".wav", "_background.wav"),
                secondary_source.cpu(),
                44100,
            )
        process_time = time.time() - start_time
        print(f"{datetime.now()} {wav_path} denoised time: {process_time:.3f}s audio len: {output_audio.shape[-1]/44100:.3f}s RTF: {output_audio.shape[-1]/44100/process_time:.3f}")
        
        vocal_path = os.path.join(save_dir, os.path.basename(wav_path).replace(".wav", "_vocal.wav"))
        bg_path = os.path.join(save_dir, os.path.basename(wav_path).replace(".wav", "_background.wav")) if self.save_background else ""
        
        return vocal_path, bg_path

def normalize_two_stem(wave, mix, is_normalize=False):
    """Save output music files"""
    
    maxv = torch.abs(wave).max()
    max_mix = torch.abs(mix).max()
    
    if maxv > 1.0:
        # print(f"\nNormalization Set {is_normalize}: Primary source above threshold for clipping. Max:{maxv}")
        # print(f"\nNormalization Set {is_normalize}: Mixture above threshold for clipping. Max:{max_mix}")
        if is_normalize:
            wave /= maxv
            mix /= maxv
    
    return wave, mix  


def get_wav_duration(file_path):
    with wave.open(file_path, 'rb') as wav_file:
        # 获取音频帧数
        n_frames = wav_file.getnframes()
        # 获取采样率
        framerate = wav_file.getframerate()
        # 计算时长（秒）
        duration = n_frames / float(framerate)
    return duration


def walkFile(data_dir, save_dir):
    res_wavs = []
    res_txts = []
    for root, dirs, files in tqdm(os.walk(data_dir)):
        # 遍历文件
        for f in files:
            if f[-4:] == '.wav':
                wav_path = os.path.join(root, f)
                if not os.path.exists(os.path.join(save_dir, f'{f[:-4]}_Vocals.wav')):
                    res_wavs.append(wav_path)
            # elif f[-4:] == '.csv':
            #     res_txts.append(os.path.join(root, f))
            
    return res_wavs # , res_txts


def process_batch(files, args, device='cpu'):

    configs = json.loads(open(args.config_path, 'r', encoding='utf-8').read())
    model_data = ModelData(
        model_path=args.model_path,
        audio_path = files,
        result_path = args.result_path,
        process_method = args.process_method,
        device = device,
        save_background = args.save_background,
        **configs
    )
    # uvr5_model = Inference_raw(model_data, device)
    # uvr5_model.process_start()

    uvr5_model = Inference(model_data, device)
    uvr5_model.load_model(args.model_path, args.num_processes)
    print(f"Loaded UVR5 model in {device}.")

    for file in files:
        vocal_path, bg_path = uvr5_model.onnx_inference(file, os.path.join(args.result_path, os.path.basename(file)), device)



def parallel_process(filenames, args):
    total_gpu = torch.cuda.device_count()
    print(f'Total GPUs: {total_gpu}')
    with ProcessPoolExecutor(max_workers=args.num_processes*total_gpu) as executor:
        tasks = []
        for i in range(args.num_processes):
            start = int(i * len(filenames) / args.num_processes)
            end = int((i + 1) * len(filenames) / args.num_processes)
            file_chunk = filenames[start:end]
            for n in range(total_gpu):
                chunk = file_chunk[int(n*len(file_chunk)/total_gpu): int((n+1)*len(file_chunk)/total_gpu)]
                device = f"cuda:{n}" if torch.cuda.is_available() else "cpu"
                print("load model in devices: ", args.num_processes, total_gpu, i, n, device)
                tasks.append(executor.submit(process_batch, chunk, args, device))

        for task in tqdm(tasks):
            task.result()


def parallel_process_cpu(filenames, args):
    with ProcessPoolExecutor(max_workers=args.num_processes) as executor:
        tasks = []
        for i in range(args.num_processes):
            start = int(i * len(filenames) / args.num_processes)
            end = int((i + 1) * len(filenames) / args.num_processes)
            chunk = filenames[start:end]
            print("load model in devices: ", args.num_processes, i, "cpu")
            tasks.append(executor.submit(process_batch, chunk, args, "cpu"))
        for task in tqdm(tasks):
            task.result()


if __name__ == "__main__":
    parser = argparse.ArgumentParser()

    parser.add_argument('-m', '--model_path', type=str, default="models/MDX_Net_Models/model_data/Kim_Vocal_1.onnx", help='模型路径')
    parser.add_argument('-c', '--config_path', type=str, default="models/MDX_Net_Models/model_data/MDX-Net-Kim-Vocal1.json", help='配置文件路径') 
    parser.add_argument('-a', '--audio_path', type=str, default="", help='wav文件名列表，放在raw文件夹下')
    parser.add_argument('-r', '--result_path', type=str, default="", help='结果存储路径')
    parser.add_argument('-p', '--process_method', type=str, default="MDX-Net", help='可选方法:["VR Arc", "MDX-Net"]')
    parser.add_argument('-b', '--save_background', type=bool, default=True, help='True:保存人声和背景音，False:只保存人声')
    parser.add_argument('-w', '--num_processes', type=int, default=4, help='You are advised to set the number of processes to the same as the number of CPU cores')

    args = parser.parse_args()

    if not os.path.exists(args.result_path):
        os.makedirs(args.result_path, exist_ok=True)
    if args.save_background:
        os.makedirs(os.path.join(os.path.dirname(args.result_path), "bg_music"), exist_ok=True)

    if os.path.isdir(args.audio_path):
        filenames = walkFile(args.audio_path, args.result_path)
    elif args.audio_path.endswith(".wav"):
        filenames = [args.audio_path]

    # shuffle(filenames)
    print(len(filenames))

    # process_batch(filenames, args, "cpu")

    multiprocessing.set_start_method("spawn", force=True)

    num_processes = args.num_processes
    if num_processes == 0:
        num_processes = os.cpu_count()

    if torch.cuda.is_available():
        parallel_process(filenames, args)
    else:
        parallel_process_cpu(filenames, args)