Spaces:
Running on Zero
Running on Zero
| import pdb | |
| import numpy as np | |
| import torch | |
| import torch.nn as nn | |
| import torch.nn.functional as F | |
| from models.layers.layer import BasicBlock | |
| from einops import rearrange | |
| import pickle | |
| import math | |
| from models.wavlm.WavLM import WavLM, WavLMConfig | |
| class ExactLengthAdjuster(nn.Module): | |
| """ | |
| Layer that ensures the output has exactly the target length along the time dimension. | |
| It either adds or removes frames as needed. | |
| """ | |
| def __init__(self, target_length=196): | |
| super(ExactLengthAdjuster, self).__init__() | |
| self.target_length = target_length | |
| def forward(self, x): | |
| # x is expected to be [batch, channels, time] | |
| current_length = x.shape[2] | |
| if current_length == self.target_length: | |
| return x | |
| elif current_length < self.target_length: | |
| # Need to add frames | |
| frames_to_add = self.target_length - current_length | |
| # Duplicate the last frame as many times as needed | |
| last_frame = x[:, :, -1:] | |
| extra_frames = last_frame.repeat(1, 1, frames_to_add) | |
| return torch.cat([x, extra_frames], dim=2) | |
| else: | |
| # Need to remove frames | |
| # Just truncate to the target length | |
| return x[:, :, :self.target_length] | |
| class WavEncoder(nn.Module): | |
| def __init__(self, out_dim, audio_in=2, target_length=256): | |
| super().__init__() | |
| self.out_dim = out_dim | |
| self.feat_extractor = nn.Sequential( | |
| BasicBlock(audio_in, out_dim//4, 15, 5, first_dilation=1700, downsample=True), | |
| BasicBlock(out_dim//4, out_dim//4, 15, 6, first_dilation=0, downsample=True), | |
| BasicBlock(out_dim//4, out_dim//4, 15, 1, first_dilation=7, ), | |
| BasicBlock(out_dim//4, out_dim//2, 15, 6, first_dilation=0, downsample=True), | |
| BasicBlock(out_dim//2, out_dim//2, 15, 1, first_dilation=7), | |
| BasicBlock(out_dim//2, out_dim, 15, 3, first_dilation=0,downsample=True), | |
| ) | |
| self.length_adjuster = ExactLengthAdjuster(target_length=target_length) | |
| def forward(self, wav_data): | |
| if wav_data.dim() == 2: | |
| wav_data = wav_data.unsqueeze(1) | |
| else: | |
| wav_data = wav_data.transpose(1, 2) | |
| out = self.feat_extractor(wav_data) | |
| out = self.length_adjuster(out) | |
| return out.transpose(1, 2) | |
| class ModalityEncoder(nn.Module): | |
| def __init__(self, | |
| data_path, | |
| t_fix_pre, | |
| audio_dim, | |
| audio_in=2, | |
| raw_audio=False, | |
| latent_dim=256, | |
| audio_fps=30, | |
| use_exp=False, | |
| target_length=256, | |
| spatial_temporal=False | |
| ): | |
| super().__init__() | |
| self.raw_audio = raw_audio | |
| self.latent_dim = latent_dim | |
| self.audio_fps = audio_fps | |
| self.WavEncoder = WavEncoder(audio_dim, audio_in=audio_in, target_length=target_length) | |
| self.text_encoder_body = nn.Linear(300, audio_dim) | |
| with open(f"{data_path}weights/vocab.pkl", 'rb') as f: | |
| self.lang_model = pickle.load(f) | |
| pre_trained_embedding = self.lang_model.word_embedding_weights | |
| self.text_pre_encoder_body = nn.Embedding.from_pretrained(torch.FloatTensor(pre_trained_embedding),freeze=t_fix_pre) | |
| word_dim = pre_trained_embedding.shape[1] | |
| if self.raw_audio: | |
| # load the pre-trained wavlm model | |
| # self.load_and_freeze_wavlm() | |
| self.audio_projection = nn.Linear(1024, audio_dim) | |
| if self.raw_audio: | |
| if use_exp: | |
| self.mix_audio_text = nn.Linear(audio_dim*3, self.latent_dim * (4 if spatial_temporal else 1)) | |
| else: | |
| self.mix_audio_text = nn.Linear(audio_dim*3, self.latent_dim * (3 if spatial_temporal else 1)) | |
| else: | |
| if use_exp: | |
| self.mix_audio_text = nn.Linear(audio_dim*2, self.latent_dim * (4 if spatial_temporal else 1)) | |
| else: | |
| self.mix_audio_text = nn.Linear(audio_dim*2, self.latent_dim * (3 if spatial_temporal else 1)) | |
| def forward(self, audio, word, raw_audio=None, squeeze_scale=4): | |
| # Initial features extraction - single transpose each | |
| # [B, T, D] -> [T, B, D] | |
| audio_feat = self.WavEncoder(audio) | |
| text_feat = self.text_encoder_body(self.text_pre_encoder_body(word)) | |
| if raw_audio is not None and self.raw_audio: | |
| # Keep the same transpose pattern for consistency | |
| # raw_feat = self.extract_wavlm_feats(raw_audio) | |
| raw_feat = self.audio_projection(raw_audio) | |
| at_feat = torch.cat([audio_feat, raw_feat, text_feat], dim=2) | |
| else: | |
| at_feat = torch.cat([audio_feat, text_feat], dim=2) # [B, T, D] | |
| at_feat = self.mix_audio_text(at_feat) # [B, T, D'] | |
| at_feat = F.avg_pool1d(at_feat.transpose(1, 2), squeeze_scale) | |
| at_feat = at_feat.transpose(1, 2) # [B, T/scale, D'] | |
| return at_feat | |
| def load_and_freeze_wavlm(self, wavlm_path='./dataloaders/wavlm/WavLM-Base+.pt'): | |
| checkpoint = torch.load(wavlm_path) | |
| self.wavlm_cfg = WavLMConfig(checkpoint['cfg']) | |
| self.audio_encoder = WavLM(self.wavlm_cfg) | |
| self.audio_encoder.load_state_dict(checkpoint['model']) | |
| self.audio_encoder.eval() | |
| for param in self.audio_encoder.parameters(): | |
| param.requires_grad = False | |
| def extract_wavlm_feats(self, wav_input_16khz): | |
| assert self.audio_encoder is not None, "Please load the wavlm model first" | |
| # check the input type | |
| if isinstance(wav_input_16khz, np.ndarray): | |
| wav_input_16khz = torch.from_numpy(wav_input_16khz) | |
| if wav_input_16khz.dim() == 1: | |
| wav_input_16khz = wav_input_16khz.unsqueeze(0) | |
| wav_input_16khz = wav_input_16khz.cuda() | |
| if self.wavlm_cfg.normalize: | |
| wav_input_16khz = F.layer_norm(wav_input_16khz, wav_input_16khz.shape) | |
| wavlm_feats = self.audio_encoder.extract_features(wav_input_16khz)[0] | |
| wavlm_feats = wavlm_feats.detach() # (bs, seq_len, dim) | |
| target_size = math.ceil(wavlm_feats.shape[1] / 50 * self.audio_fps) | |
| wavlm_feats = F.interpolate( | |
| wavlm_feats.transpose(1, 2), | |
| size=target_size, | |
| align_corners=True, | |
| mode='linear' | |
| ).transpose(1, 2) | |
| return wavlm_feats | |