1
0
mirror of synced 2024-12-02 19:17:18 +01:00
Retrieval-based-Voice-Conve.../extract_feature_print.py

124 lines
3.8 KiB
Python
Raw Normal View History

import os, sys, traceback
os.environ["PYTORCH_ENABLE_MPS_FALLBACK"] = "1"
os.environ["PYTORCH_MPS_HIGH_WATERMARK_RATIO"] = "0.0"
2023-06-24 10:21:31 +02:00
2023-04-13 17:57:27 +02:00
# device=sys.argv[1]
n_part = int(sys.argv[2])
i_part = int(sys.argv[3])
if len(sys.argv) == 6:
exp_dir = sys.argv[4]
version = sys.argv[5]
else:
i_gpu = sys.argv[4]
exp_dir = sys.argv[5]
os.environ["CUDA_VISIBLE_DEVICES"] = str(i_gpu)
version = sys.argv[6]
2023-03-31 11:54:38 +02:00
import torch
import torch.nn.functional as F
import soundfile as sf
import numpy as np
from fairseq import checkpoint_utils
2023-06-24 10:36:15 +02:00
device = "cpu"
if torch.cuda.is_available():
device = "cuda"
elif torch.backends.mps.is_available():
device = "mps"
2023-05-12 21:29:30 +02:00
f = open("%s/extract_f0_feature.log" % exp_dir, "a+")
2023-03-31 11:54:38 +02:00
def printt(strr):
print(strr)
f.write("%s\n" % strr)
f.flush()
2023-03-31 11:54:38 +02:00
printt(sys.argv)
model_path = "hubert_base.pt"
printt(exp_dir)
wavPath = "%s/1_16k_wavs" % exp_dir
outPath = (
"%s/3_feature256" % exp_dir if version == "v1" else "%s/3_feature768" % exp_dir
)
os.makedirs(outPath, exist_ok=True)
2023-03-31 11:54:38 +02:00
# wave must be 16k, hop_size=320
def readwave(wav_path, normalize=False):
wav, sr = sf.read(wav_path)
assert sr == 16000
feats = torch.from_numpy(wav).float()
if feats.dim() == 2: # double channels
feats = feats.mean(-1)
assert feats.dim() == 1, feats.dim()
if normalize:
with torch.no_grad():
feats = F.layer_norm(feats, feats.shape)
feats = feats.view(1, -1)
return feats
2023-03-31 11:54:38 +02:00
# HuBERT model
printt("load model(s) from {}".format(model_path))
# if hubert model is exist
if os.access(model_path, os.F_OK) == False:
printt(
"Error: Extracting is shut down because %s does not exist, you may download it from https://huggingface.co/lj1995/VoiceConversionWebUI/tree/main"
% model_path
)
exit(0)
2023-03-31 11:54:38 +02:00
models, saved_cfg, task = checkpoint_utils.load_model_ensemble_and_task(
[model_path],
suffix="",
)
model = models[0]
model = model.to(device)
printt("move model to %s" % device)
if device not in ["mps", "cpu"]:
model = model.half()
2023-03-31 11:54:38 +02:00
model.eval()
todo = sorted(list(os.listdir(wavPath)))[i_part::n_part]
n = max(1, len(todo) // 10) # 最多打印十条
if len(todo) == 0:
printt("no-feature-todo")
2023-03-31 11:54:38 +02:00
else:
printt("all-feature-%s" % len(todo))
for idx, file in enumerate(todo):
2023-03-31 11:54:38 +02:00
try:
if file.endswith(".wav"):
wav_path = "%s/%s" % (wavPath, file)
out_path = "%s/%s" % (outPath, file.replace("wav", "npy"))
2023-03-31 11:54:38 +02:00
if os.path.exists(out_path):
continue
2023-03-31 11:54:38 +02:00
feats = readwave(wav_path, normalize=saved_cfg.task.normalize)
padding_mask = torch.BoolTensor(feats.shape).fill_(False)
inputs = {
"source": feats.half().to(device)
2023-05-12 21:29:30 +02:00
if device not in ["mps", "cpu"]
else feats.to(device),
2023-03-31 11:54:38 +02:00
"padding_mask": padding_mask.to(device),
"output_layer": 9 if version == "v1" else 12, # layer 9
2023-03-31 11:54:38 +02:00
}
with torch.no_grad():
logits = model.extract_features(**inputs)
feats = (
model.final_proj(logits[0]) if version == "v1" else logits[0]
)
2023-03-31 11:54:38 +02:00
feats = feats.squeeze(0).float().cpu().numpy()
if np.isnan(feats).sum() == 0:
2023-04-09 17:31:38 +02:00
np.save(out_path, feats, allow_pickle=False)
else:
printt("%s-contains nan" % file)
if idx % n == 0:
printt("now-%s,all-%s,%s,%s" % (len(todo), idx, file, feats.shape))
2023-03-31 11:54:38 +02:00
except:
printt(traceback.format_exc())
2023-04-09 17:31:38 +02:00
printt("all-feature-done")