diff --git a/emo_gen.py b/emo_gen.py index 71325c8..fda55c6 100644 --- a/emo_gen.py +++ b/emo_gen.py @@ -105,19 +105,6 @@ def process_func( return y -def get_emo(path): - wav, sr = librosa.load(path, 16000) - device = config.bert_gen_config.device - return process_func( - np.expand_dims(wav, 0).astype(np.float), - sr, - model, - processor, - device, - embeddings=True, - ).squeeze(0) - - if __name__ == "__main__": parser = argparse.ArgumentParser() parser.add_argument( @@ -160,4 +147,4 @@ if __name__ == "__main__": emb = model(data.to(device))[0].detach().cpu().numpy() np.save(emo_path, emb) - print("Emo vec 生成完毕!") + print("Emo vec 生成完毕!") diff --git a/get_emo.py b/get_emo.py new file mode 100644 index 0000000..b0fd448 --- /dev/null +++ b/get_emo.py @@ -0,0 +1,26 @@ +from emo_gen import RegressionHead, EmotionModel, AudioDataset, process_func + +import librosa +import numpy as np +import torch +from transformers import Wav2Vec2Processor + +from config import config + +model_name = "./emotional/wav2vec2-large-robust-12-ft-emotion-msp-dim" +device = "cuda" if torch.cuda.is_available() else "cpu" +processor = Wav2Vec2Processor.from_pretrained(model_name) +model = EmotionModel.from_pretrained(model_name).to(device) + + +def get_emo(path): + wav, sr = librosa.load(path, 16000) + device = config.bert_gen_config.device + return process_func( + np.expand_dims(wav, 0).astype(np.float), + sr, + model, + processor, + device, + embeddings=True, + ).squeeze(0) diff --git a/infer.py b/infer.py index 88509c0..51722da 100644 --- a/infer.py +++ b/infer.py @@ -11,7 +11,7 @@ import torch import commons from text import cleaned_text_to_sequence, get_bert -from emo_gen import get_emo +from get_emo import get_emo from text.cleaner import clean_text import utils