121 lines
4.2 KiB
Python
121 lines
4.2 KiB
Python
import hydra
|
|
import librosa
|
|
import torch
|
|
import yaml
|
|
from prodict import Prodict
|
|
import torchaudio
|
|
|
|
from musiclm_pytorch import AudioSpectrogramTransformerPretrained, TextTransformerPretrained, MuLaN, MuLaNEmbedder
|
|
from omegaconf import DictConfig
|
|
import os
|
|
|
|
def get_pretrained_config(root, name):
|
|
if root is None:
|
|
return name
|
|
path = os.path.join(root, name)
|
|
#获取snapshots目录下的目录
|
|
config_dir = os.path.join(path, 'snapshots')
|
|
config_files = os.listdir(config_dir)
|
|
assert len(config_files) == 1
|
|
config_path = os.path.join(config_dir, config_files[0])
|
|
return config_path
|
|
|
|
def create_MuLaN_from_config(config: DictConfig):
|
|
"""
|
|
Create a MuLaN model from a configuration file.
|
|
"""
|
|
pretraind_root = config.model.pretraind_root
|
|
|
|
audio_model_name = get_pretrained_config(pretraind_root, config.model.audio_model.name)
|
|
audio_transformer = AudioSpectrogramTransformerPretrained(
|
|
model_name = audio_model_name,
|
|
model_dim = config.model.audio_model.model_dim,
|
|
use_layer_idx = config.model.audio_model.use_layer_idx,
|
|
**config.model.audio_transformer
|
|
)
|
|
text_model_name = get_pretrained_config(pretraind_root, config.model.text_model.name)
|
|
text_transformer = TextTransformerPretrained(
|
|
model_name = text_model_name,
|
|
**config.model.text_transformer
|
|
)
|
|
|
|
mulan = MuLaN(
|
|
audio_transformer = audio_transformer,
|
|
text_transformer = text_transformer,
|
|
**config.model.mulan
|
|
)
|
|
|
|
return mulan
|
|
|
|
|
|
def create_CLAP_model( model_kwargs = {}, ckpt_path = None ):
|
|
from musiclm_pytorch import SoftmaxContrastiveLearning
|
|
import laion_clap
|
|
|
|
from torch import nn
|
|
import torch
|
|
from torchaudio.functional import resample
|
|
|
|
import numpy as np
|
|
|
|
from functools import partial
|
|
|
|
# quantization
|
|
def int16_to_float32(x):
|
|
return (x / 32767.0).float()
|
|
|
|
def float32_to_int16(x):
|
|
x = torch.clip(x, min=-1., max=1.)
|
|
return (x * 32767.).int()
|
|
|
|
model = laion_clap.CLAP_Module(enable_fusion=False, **model_kwargs)
|
|
if ckpt_path is not None:
|
|
model.load_ckpt(ckpt_path)
|
|
else:
|
|
model.load_ckpt()
|
|
|
|
class CLAP_Model(nn.Module):
|
|
def __init__(self, model, sr = 24000, decoupled_contrastive_learning = True):
|
|
super().__init__()
|
|
self.model = model
|
|
self.model.eval()
|
|
self.orig_sr = sr
|
|
|
|
klass = partial(SoftmaxContrastiveLearning, decoupled_contrastive_learning = decoupled_contrastive_learning)
|
|
self.contrast = klass()
|
|
|
|
|
|
def forward(self, wavs, raw_texts):
|
|
with torch.no_grad():
|
|
wavs = int16_to_float32(float32_to_int16(resample(wavs, self.orig_sr, 48000)))
|
|
audio_latents = self.model.get_audio_embedding_from_data(x = wavs, use_tensor=True).float()
|
|
text_latents = model.get_text_embedding(raw_texts, use_tensor=True)
|
|
cl_loss = self.contrast(audio_latents, text_latents)
|
|
return cl_loss
|
|
|
|
clap = CLAP_Model(model)
|
|
return clap
|
|
|
|
def get_mulan(config):
|
|
with open(config, "r") as stream:
|
|
mulan_config = yaml.safe_load(stream)
|
|
mulan_config = Prodict.from_dict(mulan_config)
|
|
ckpt_path = mulan_config.checkpoint_path
|
|
mulan = create_MuLaN_from_config(mulan_config)
|
|
mulan_embedder = MuLaNEmbedder(mulan, checkpoint_path = ckpt_path)
|
|
mulan_embedder.eval()
|
|
|
|
return mulan_embedder
|
|
|
|
def extract_mert_embeds(mulan_embd_extractor, layer_num, filename):
|
|
input_audios, fs = torchaudio.load(filename)
|
|
mulan_sr = 24000
|
|
if(fs!=mulan_sr):
|
|
input_audios = torchaudio.functional.resample(input_audios, fs, mulan_sr)
|
|
fs = mulan_sr
|
|
# print(input_audios.shape)
|
|
inputs = mulan_embd_extractor.mulan.audio.processor(input_audios, sampling_rate=mulan_embd_extractor.mulan.audio.sr, return_tensors="pt")
|
|
input_values = inputs['input_values'].squeeze(0).to(input_audios.device, dtype = input_audios.dtype)
|
|
prompt_embeds = mulan_embd_extractor.mulan.audio.model(input_values, output_hidden_states=True).hidden_states[layer_num] # batch_size, Time steps, 1024 feature_dim
|
|
return prompt_embeds
|