Update

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 生成完毕!")

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)

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