140 lines
4.5 KiB
Python
140 lines
4.5 KiB
Python
# Copyright (c) 2023 Amphion.
|
|
#
|
|
# This source code is licensed under the MIT license found in the
|
|
# LICENSE file in the root directory of this source tree.
|
|
|
|
# This code is modified from https://huggingface.co/m-a-p/MERT-v1-330M
|
|
|
|
import torch
|
|
from tqdm import tqdm
|
|
import numpy as np
|
|
|
|
from transformers import Wav2Vec2FeatureExtractor
|
|
from transformers import AutoModel
|
|
import torchaudio
|
|
import torchaudio.transforms as T
|
|
from sklearn.preprocessing import StandardScaler
|
|
|
|
|
|
def mert_encoder(model, processor, audio_path, hps):
|
|
"""
|
|
# mert default sr: 24000
|
|
"""
|
|
with torch.no_grad():
|
|
resample_rate = processor.sampling_rate
|
|
device = next(model.parameters()).device
|
|
|
|
input_audio, sampling_rate = torchaudio.load(audio_path)
|
|
input_audio = input_audio.squeeze()
|
|
|
|
if sampling_rate != resample_rate:
|
|
resampler = T.Resample(sampling_rate, resample_rate)
|
|
input_audio = resampler(input_audio)
|
|
|
|
inputs = processor(
|
|
input_audio, sampling_rate=resample_rate, return_tensors="pt"
|
|
).to(
|
|
device
|
|
) # {input_values: tensor, attention_mask: tensor}
|
|
|
|
outputs = model(**inputs, output_hidden_states=True) # list: len is 25
|
|
|
|
# [25 layer, Time steps, 1024 feature_dim]
|
|
# all_layer_hidden_states = torch.stack(outputs.hidden_states).squeeze()
|
|
# mert_features.append(all_layer_hidden_states)
|
|
|
|
feature = outputs.hidden_states[
|
|
hps.mert_feature_layer
|
|
].squeeze() # [1, frame len, 1024] -> [frame len, 1024]
|
|
|
|
return feature.cpu().detach().numpy()
|
|
|
|
|
|
def mert_features_normalization(raw_mert_features):
|
|
normalized_mert_features = list()
|
|
|
|
mert_features = np.array(raw_mert_features)
|
|
scaler = StandardScaler().fit(mert_features)
|
|
for raw_mert_feature in raw_mert_feature:
|
|
normalized_mert_feature = scaler.transform(raw_mert_feature)
|
|
normalized_mert_features.append(normalized_mert_feature)
|
|
return normalized_mert_features
|
|
|
|
|
|
def get_mapped_mert_features(raw_mert_features, mapping_features, fast_mapping=True):
|
|
source_hop = 320
|
|
target_hop = 256
|
|
|
|
factor = np.gcd(source_hop, target_hop)
|
|
source_hop //= factor
|
|
target_hop //= factor
|
|
print(
|
|
"Mapping source's {} frames => target's {} frames".format(
|
|
target_hop, source_hop
|
|
)
|
|
)
|
|
|
|
mert_features = []
|
|
for index, mapping_feat in enumerate(tqdm(mapping_features)):
|
|
# mapping_feat: (mels_frame_len, n_mels)
|
|
target_len = mapping_feat.shape[0]
|
|
|
|
# (frame_len, 1024)
|
|
raw_feats = raw_mert_features[index].cpu().numpy()
|
|
source_len, width = raw_feats.shape
|
|
|
|
# const ~= target_len * target_hop
|
|
const = source_len * source_hop // target_hop * target_hop
|
|
|
|
# (source_len * source_hop, dim)
|
|
up_sampling_feats = np.repeat(raw_feats, source_hop, axis=0)
|
|
# (const, dim) -> (const/target_hop, target_hop, dim) -> (const/target_hop, dim)
|
|
down_sampling_feats = np.average(
|
|
up_sampling_feats[:const].reshape(-1, target_hop, width), axis=1
|
|
)
|
|
|
|
err = abs(target_len - len(down_sampling_feats))
|
|
if err > 3:
|
|
print("index:", index)
|
|
print("mels:", mapping_feat.shape)
|
|
print("raw mert vector:", raw_feats.shape)
|
|
print("up_sampling:", up_sampling_feats.shape)
|
|
print("const:", const)
|
|
print("down_sampling_feats:", down_sampling_feats.shape)
|
|
exit()
|
|
if len(down_sampling_feats) < target_len:
|
|
# (1, dim) -> (err, dim)
|
|
end = down_sampling_feats[-1][None, :].repeat(err, axis=0)
|
|
down_sampling_feats = np.concatenate([down_sampling_feats, end], axis=0)
|
|
|
|
# (target_len, dim)
|
|
feats = down_sampling_feats[:target_len]
|
|
mert_features.append(feats)
|
|
|
|
return mert_features
|
|
|
|
|
|
def load_mert_model(hps):
|
|
print("Loading MERT Model: ", hps.mert_model)
|
|
|
|
# Load model
|
|
model_name = hps.mert_model
|
|
model = AutoModel.from_pretrained(model_name, trust_remote_code=True)
|
|
|
|
if torch.cuda.is_available():
|
|
model = model.cuda()
|
|
|
|
# model = model.eval()
|
|
|
|
preprocessor = Wav2Vec2FeatureExtractor.from_pretrained(
|
|
model_name, trust_remote_code=True
|
|
)
|
|
return model, preprocessor
|
|
|
|
|
|
# loading the corresponding preprocessor config
|
|
# def load_preprocessor (model_name="m-a-p/MERT-v1-330M"):
|
|
# print('load_preprocessor...')
|
|
# preprocessor = Wav2Vec2FeatureExtractor.from_pretrained(model_name,trust_remote_code=True)
|
|
# return preprocessor
|