""" Contrastive Language-Audio Pretraining Model from LAION -------------------------------------------------------- Paper: https://arxiv.org/abs/2211.06687 Authors (equal contributions): Ke Chen, Yusong Wu, Tianyu Zhang, Yuchen Hui Support: LAION """ import os import logging import torch import librosa from clap_module import create_model from .training.data import get_audio_features from .training.data import int16_to_float32, float32_to_int16 from transformers import RobertaTokenizer import wget from clap_module.factory import load_state_dict class CLAP_Module(torch.nn.Module): def __init__(self, enable_fusion=False, device=None, amodel= 'HTSAT-tiny', tmodel='roberta') -> None: """Initialize CLAP Model Parameters ---------- enable_fusion: bool if true, it will create the fusion clap model, otherwise non-fusion clap model (default: false) device: str if None, it will automatically detect the device (gpu or cpu) amodel: str audio encoder architecture, default: HTSAT-tiny tmodel: str text encoder architecture, default: roberta """ super(CLAP_Module, self).__init__() if device is None: device = 'cuda:0' if torch.cuda.is_available() else 'cpu' precision = 'fp32' if enable_fusion: fusion_type = 'aff_2d' model, model_cfg = create_model( amodel, tmodel, precision=precision, device=device, enable_fusion=enable_fusion, fusion_type=fusion_type ) else: model, model_cfg = create_model( amodel, tmodel, precision=precision, device=device, enable_fusion=enable_fusion ) self.enable_fusion = enable_fusion self.model = model self.model_cfg = model_cfg self.tokenize = RobertaTokenizer.from_pretrained('roberta-base') def tokenizer(self, text): result = self.tokenize( text, padding="max_length", truncation=True, max_length=77, return_tensors="pt", ) return result def load_ckpt(self, ckpt = None, model_id = -1, verbose = True): """Load the pretrained checkpoint of CLAP model Parameters ---------- ckpt: str if ckpt is specified, the model will load this ckpt, otherwise the model will download the ckpt from zenodo. \n For fusion model, it will download the 630k+audioset fusion model (id=3). For non-fusion model, it will download the 630k+audioset model (id=1). model_id: if model_id is specified, you can download our best ckpt, as: id = 0 --> 630k non-fusion ckpt \n id = 1 --> 630k+audioset non-fusion ckpt \n id = 2 --> 630k fusion ckpt \n id = 3 --> 630k+audioset fusion ckpt \n Note that if your model is specied as non-fusion model but you download a fusion model ckpt, you will face an error. """ download_link = 'https://huggingface.co/lukewys/laion_clap/resolve/main/' download_names = [ '630k-best.pt', '630k-audioset-best.pt', '630k-fusion-best.pt', '630k-audioset-fusion-best.pt' ] if ckpt is not None: logging.info(f'Load the specified checkpoint {ckpt} from users.') else: logging.info(f'Load our best checkpoint in the paper.') if model_id == -1: model_id = 3 if self.enable_fusion else 1 package_dir = os.path.dirname(os.path.realpath(__file__)) weight_file_name = download_names[model_id] ckpt = os.path.join(package_dir, weight_file_name) if os.path.exists(ckpt): logging.info(f'The checkpoint is already downloaded') else: logging.info('Downloading laion_clap weight files...') ckpt = wget.download(download_link + weight_file_name, os.path.dirname(ckpt)) logging.info('Download completed!') logging.info('Load Checkpoint...') ckpt = load_state_dict(ckpt, skip_params=True) self.model.load_state_dict(ckpt) if verbose: param_names = [n for n, p in self.model.named_parameters()] for n in param_names: logging.info(n, "\t", "Loaded" if n in ckpt else "Unloaded") def get_audio_embedding_from_filelist(self, x, use_tensor=False): """get audio embeddings from the audio file list Parameters ---------- x: List[str] (N,): an audio file list to extract features, audio files can have different lengths (as we have the feature fusion machanism) use_tensor: boolean: if True, it will return the torch tensor, preserving the gradient (default: False). Returns ---------- audio_embed : numpy.darray | torch.Tensor (N,D): audio embeddings that extracted from audio files """ self.model.eval() audio_input = [] for f in x: # load the waveform of the shape (T,), should resample to 48000 audio_waveform, _ = librosa.load(f, sr=48000) # quantize audio_waveform = int16_to_float32(float32_to_int16(audio_waveform)) audio_waveform = torch.from_numpy(audio_waveform).float() temp_dict = {} temp_dict = get_audio_features( temp_dict, audio_waveform, 480000, data_truncating='fusion' if self.enable_fusion else 'rand_trunc', data_filling='repeatpad', audio_cfg=self.model_cfg['audio_cfg'], require_grad=audio_waveform.requires_grad ) audio_input.append(temp_dict) audio_embed = self.model.get_audio_embedding(audio_input) if not use_tensor: audio_embed = audio_embed.detach().cpu().numpy() return audio_embed def get_audio_embedding_from_data(self, x, use_tensor=False): """get audio embeddings from the audio data Parameters ---------- x: np.darray | torch.Tensor (N,T): audio data, must be mono audio tracks. use_tensor: boolean: if True, x should be the tensor input and the output will be the tesnor, preserving the gradient (default: False). Note that if 'use tensor' is set to True, it will not do the quantize of the audio waveform (otherwise the gradient will not be preserved). Returns ---------- audio embed: numpy.darray | torch.Tensor (N,D): audio embeddings that extracted from audio files """ self.model.eval() audio_input = [] for audio_waveform in x: # quantize if not use_tensor: audio_waveform = int16_to_float32(float32_to_int16(audio_waveform)) audio_waveform = torch.from_numpy(audio_waveform).float() temp_dict = {} temp_dict = get_audio_features( temp_dict, audio_waveform, 480000, data_truncating='fusion' if self.enable_fusion else 'rand_trunc', data_filling='repeatpad', audio_cfg=self.model_cfg['audio_cfg'], require_grad=audio_waveform.requires_grad ) audio_input.append(temp_dict) audio_embed = self.model.get_audio_embedding(audio_input) if not use_tensor: audio_embed = audio_embed.detach().cpu().numpy() return audio_embed def get_text_embedding(self, x, tokenizer = None, use_tensor = False): """get text embeddings from texts Parameters ---------- x: List[str] (N,): text list tokenizer: func: the tokenizer function, if not provided (None), will use the default Roberta tokenizer. use_tensor: boolean: if True, the output will be the tesnor, preserving the gradient (default: False). Returns ---------- text_embed : numpy.darray | torch.Tensor (N,D): text embeddings that extracted from texts """ self.model.eval() if tokenizer is not None: text_input = tokenizer(x) else: text_input = self.tokenizer(x) text_embed = self.model.get_text_embedding(text_input) if not use_tensor: text_embed = text_embed.detach().cpu().numpy() return text_embed