import math import os from tqdm import tqdm import librosa import numpy as np import torch from torch.utils.data import Dataset, DataLoader, DistributedSampler import torch.multiprocessing as mp from torch.distributed import init_process_group import torch.distributed as dist from .inference import RMVPE from ....utils.commons.dataset_utils import batch_by_size, build_dataloader # import utils from ....utils.audio import get_wav_num_frames """ A convenient API for batch inference update: add ddp """ class RMVPEInferDataset(Dataset): def __init__(self, wav_fns: list, id_and_sizes=None, sr=24000, hop_size=128, num_workers=0): if id_and_sizes is None: id_and_sizes = [] if type(wav_fns[0]) == str: # wav_paths for idx, wav_path in enumerate(wav_fns): total_frames = get_wav_num_frames(wav_path, sr) id_and_sizes.append((idx, round(total_frames / hop_size))) else: # numpy arrays, mono wavs for idx, wav in enumerate(wav_fns): id_and_sizes.append((idx, round(wav.shape[-1] / hop_size))) self.wav_fns = wav_fns self.id_and_sizes = id_and_sizes self.sr = sr self.num_workers = num_workers def __getitem__(self, idx): if type(self.wav_fns[idx]) == str: wav_fn = self.wav_fns[idx] wav, _ = librosa.core.load(wav_fn, sr=self.sr) else: wav = self.wav_fns[idx] return idx, wav def collater(self, samples: list): return samples def __len__(self): return len(self.wav_fns) def ordered_indices(self): """Return an ordered list of indices. Batches will be constructed based on this order.""" return np.arange(len(self)) def num_tokens(self, index): return self.id_and_sizes[index][1] @torch.no_grad() def extract(wav_fns: list, id_and_sizes=None, ckpt=None, sr=24000, hop_size=128, bsz=128, max_tokens=100000, fmax=900, fmin=50, ds_workers=0): all_gpu_ids = [int(x) for x in os.environ.get("CUDA_VISIBLE_DEVICES", "").split(",") if x != ''] num_gpus = len(all_gpu_ids) dist_config = { "dist_backend": "nccl", "dist_url": "tcp://localhost:54189", "world_size": 1 } # https://discuss.pytorch.org/t/how-to-fix-a-sigsegv-in-pytorch-when-using-distributed-training-e-g-ddp/113518/10#:~:text=Using%20start%20and%20join%20avoids # https://github.com/pytorch/pytorch/issues/40403#issuecomment-648515174 # mp.set_start_method('spawn') if num_gpus > 1: result_queue = mp.Queue() for rank in range(num_gpus): mp.Process(target=extract_worker, args=(rank, wav_fns, id_and_sizes, ckpt, sr, hop_size, bsz, max_tokens, fmax, fmin, dist_config, num_gpus, ds_workers, result_queue,)).start() f0_res = [None] * len(wav_fns) for _ in range(num_gpus): f0_res_dict = result_queue.get() for idx in f0_res_dict: f0_res[idx] = f0_res_dict[idx] del f0_res_dict else: # f0_res = extract_one_process(wav_fns, id_and_sizes, ckpt, sr, hop_size, bsz, max_tokens, fmax, fmin) f0_res_dict = extract_worker(0, wav_fns, id_and_sizes, ckpt, sr, hop_size, bsz, max_tokens, fmax, fmin, dist_config, num_gpus, ds_workers, None) f0_res = [None] * len(wav_fns) for idx in f0_res_dict: f0_res[idx] = f0_res_dict[idx] return f0_res @torch.no_grad() def extract_worker(rank, wav_fns: list, id_and_sizes=None, ckpt=None, sr=24000, hop_size=128, bsz=128, max_tokens=100000, fmax=900, fmin=50, dist_config=None, num_gpus=1, ds_workers=0, q=None): # print(f"rank: {rank}") if num_gpus > 1: init_process_group(backend=dist_config['dist_backend'], init_method=dist_config['dist_url'], world_size=dist_config['world_size'] * num_gpus, rank=rank) dataset = RMVPEInferDataset(wav_fns, id_and_sizes, sr, hop_size, num_workers=ds_workers) # ds_sampler = DistributedSampler(dataset, shuffle=False) if num_gpus > 1 else None # loader = DataLoader(dataset, sampler=ds_sampler, collate_fn=dataset.collator, batch_size=1, num_workers=40, drop_last=False) loader = build_dataloader(dataset, shuffle=False, max_tokens=max_tokens, max_sentences=bsz, use_ddp=num_gpus > 1) loader = tqdm(loader, desc=f'| Processing f0 in [n_ranks={num_gpus}; max_tokens={max_tokens}; max_sentences={bsz}]') if rank == 0 else loader device = torch.device(f"cuda:{int(rank)}") model = RMVPE(ckpt, device=device) f0_res_dict = {} for batch in loader: if batch is None or len(batch) == 0: continue idxs = [item[0] for item in batch] wavs = [item[1] for item in batch] lengths = [(wav.shape[0] + hop_size - 1) // hop_size for wav in wavs] with torch.no_grad(): f0s, uvs = model.get_pitch_batch( wavs, sample_rate=sr, hop_size=hop_size, lengths=lengths, fmax=fmax, fmin=fmin ) for i, idx in enumerate(idxs): f0_res_dict[idx] = f0s[i] if q is not None: q.put(f0_res_dict) else: return f0_res_dict # old version def extract_one_process(wav_fns: list, id_and_sizes=None, ckpt=None, sr=24000, hop_size=128, bsz=128, max_tokens=100000, fmax=900, fmin=50, device='cuda'): assert ckpt is not None rmvpe = RMVPE(ckpt, device=device) if id_and_sizes is None: id_and_sizes = [] if type(wav_fns[0]) == str: # wav_paths for idx, wav_path in enumerate(wav_fns): total_frames = get_wav_num_frames(wav_path, sr) id_and_sizes.append((idx, round(total_frames / hop_size))) else: # numpy arrays, mono wavs for idx, wav in enumerate(wav_fns): id_and_sizes.append((idx, round(wav.shape[-1] / hop_size))) get_size = lambda x: x[1] bs = batch_by_size(id_and_sizes, get_size, max_tokens=max_tokens, max_sentences=bsz) for i in range(len(bs)): bs[i] = [bs[i][j][0] for j in range(len(bs[i]))] f0_res = [None] * len(wav_fns) for batch in tqdm(bs, total=len(bs), desc=f'| Processing f0 in [max_tokens={max_tokens}; max_sentences={bsz}]'): wavs, mel_lengths, lengths = [], [], [] for idx in batch: if type(wav_fns[idx]) == str: wav_fn = wav_fns[idx] wav, _ = librosa.core.load(wav_fn, sr=sr) else: wav = wav_fns[idx] wavs.append(wav) mel_lengths.append(math.ceil((wav.shape[0] + 1) / hop_size)) lengths.append((wav.shape[0] + hop_size - 1) // hop_size) with torch.no_grad(): f0s, uvs = rmvpe.get_pitch_batch( wavs, sample_rate=sr, hop_size=hop_size, lengths=lengths, fmax=fmax, fmin=fmin ) for i, idx in enumerate(batch): f0_res[idx] = f0s[i] if rmvpe is not None: rmvpe.release_cuda() torch.cuda.empty_cache() rmvpe = None return f0_res