ChazzyG/Retrieval-based-Voice-Conversion-WebUI

1import os, sys, traceback2 3# device=sys.argv[1]4n_part = int(sys.argv[2])5i_part = int(sys.argv[3])6if len(sys.argv) == 5:7    exp_dir = sys.argv[4]8else:9    i_gpu = sys.argv[4]10    exp_dir = sys.argv[5]11    os.environ["CUDA_VISIBLE_DEVICES"] = str(i_gpu)12 13import torch14import torch.nn.functional as F15import soundfile as sf16import numpy as np17from fairseq import checkpoint_utils18 19device = torch.device("cuda" if torch.cuda.is_available() else "cpu")20 21f = open("%s/extract_f0_feature.log" % exp_dir, "a+")22 23 24def printt(strr):25    print(strr)26    f.write("%s\n" % strr)27    f.flush()28 29 30printt(sys.argv)31model_path = "hubert_base.pt"32 33printt(exp_dir)34wavPath = "%s/1_16k_wavs" % exp_dir35outPath = "%s/3_feature256" % exp_dir36os.makedirs(outPath, exist_ok=True)37 38 39# wave must be 16k, hop_size=32040def readwave(wav_path, normalize=False):41    wav, sr = sf.read(wav_path)42    assert sr == 1600043    feats = torch.from_numpy(wav).float()44    if feats.dim() == 2:  # double channels45        feats = feats.mean(-1)46    assert feats.dim() == 1, feats.dim()47    if normalize:48        with torch.no_grad():49            feats = F.layer_norm(feats, feats.shape)50    feats = feats.view(1, -1)51    return feats52 53 54# HuBERT model55printt("load model(s) from {}".format(model_path))56models, saved_cfg, task = checkpoint_utils.load_model_ensemble_and_task(57    [model_path],58    suffix="",59)60model = models[0]61model = model.to(device)62printt("move model to %s" % device)63if device != "cpu":64    model = model.half()65model.eval()66 67todo = sorted(list(os.listdir(wavPath)))[i_part::n_part]68n = max(1, len(todo) // 10)  # 最多打印十条69if len(todo) == 0:70    printt("no-feature-todo")71else:72    printt("all-feature-%s" % len(todo))73    for idx, file in enumerate(todo):74        try:75            if file.endswith(".wav"):76                wav_path = "%s/%s" % (wavPath, file)77                out_path = "%s/%s" % (outPath, file.replace("wav", "npy"))78 79                if os.path.exists(out_path):80                    continue81 82                feats = readwave(wav_path, normalize=saved_cfg.task.normalize)83                padding_mask = torch.BoolTensor(feats.shape).fill_(False)84                inputs = {85                    "source": feats.half().to(device)86                    if device != "cpu"87                    else feats.to(device),88                    "padding_mask": padding_mask.to(device),89                    "output_layer": 9,  # layer 990                }91                with torch.no_grad():92                    logits = model.extract_features(**inputs)93                    feats = model.final_proj(logits[0])94 95                feats = feats.squeeze(0).float().cpu().numpy()96                if np.isnan(feats).sum() == 0:97                    np.save(out_path, feats, allow_pickle=False)98                else:99                    printt("%s-contains nan" % file)100                if idx % n == 0:101                    printt("now-%s,all-%s,%s,%s" % (len(todo), idx, file, feats.shape))102        except:103            printt(traceback.format_exc())104    printt("all-feature-done")105