Create AniTalker/extract_audio_features.py

AniTalker/extract_audio_features.py

@@ -0,0 +1,51 @@
+import argparse
+import os
+import numpy as np
+import librosa
+import torch
+from tqdm import tqdm
+from transformers import Wav2Vec2FeatureExtractor, HubertModel
+
+def main(args):
+    if not torch.cuda.is_available() and args.computed_device == 'cuda':
+        print('CUDA is not available on this device. Switching to CPU.')
+        args.computed_device = 'cpu'
+    
+    device = torch.device(args.computed_device)
+    model = HubertModel.from_pretrained(args.model_path).to(device)
+    feature_extractor = Wav2Vec2FeatureExtractor.from_pretrained(args.model_path)
+    model.feature_extractor._freeze_parameters()
+    model.eval()
+
+    os.makedirs(args.audio_feature_saved_path, exist_ok=True)
+
+    for wavfile in tqdm(os.listdir(args.audio_dir_path)):
+        npy_save_path = os.path.join(args.audio_feature_saved_path, os.path.splitext(os.path.basename(wavfile))[0] + '.npy')
+        
+        if os.path.exists(npy_save_path):
+            continue
+
+        audio, sr = librosa.load(os.path.join(args.audio_dir_path, wavfile), sr=16000)
+        input_values = feature_extractor(audio, sampling_rate=16000, padding=True, do_normalize=True, return_tensors="pt").input_values
+        input_values = input_values.to(device)
+        ws_feats = []
+        with torch.no_grad():
+            outputs = model(input_values, output_hidden_states=True)
+            for i in range(len(outputs.hidden_states)):
+                ws_feats.append(outputs.hidden_states[i].detach().cpu().numpy())
+            ws_feat_obj = np.array(ws_feats)
+            ws_feat_obj = np.squeeze(ws_feat_obj, 1)
+
+            if args.padding_to_align_audio:
+                ws_feat_obj = np.pad(ws_feat_obj, ((0, 0), (0, 1), (0, 0)), 'edge')
+            np.save(npy_save_path, ws_feat_obj)
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(description="Extract audio features using a pre-trained HuBERT model.")
+    parser.add_argument("--model_path", type=str, default='weights/chinese-hubert-large', help="Path to the pre-trained model weights.")
+    parser.add_argument("--audio_dir_path", type=str, default='./audio_samples/raw_audios/', help="Directory containing raw audio files.")
+    parser.add_argument("--audio_feature_saved_path", type=str, default='./audio_samples/audio_features/', help="Directory where extracted audio features will be saved.")
+    parser.add_argument("--computed_device", type=str, default='cuda', choices=['cuda', 'cpu'], help="Device to compute the audio features on. Use 'cuda' for GPU or 'cpu' for CPU.")
+    parser.add_argument("--padding_to_align_audio", type=bool, default=True, help="Whether to pad the audio to align features.")
+    args = parser.parse_args()
+    main(args)