| 177 | |
| 178 | ################ Condition Encoders ################ |
| 179 | class AudioEncoder(BaseModel): |
| 180 | def __init__(self, opt): |
| 181 | super().__init__() |
| 182 | self.opt = opt |
| 183 | self.only_last_features = opt.only_last_features |
| 184 | |
| 185 | self.num_frames_for_clip = int(opt.wav2vec_sec * self.opt.fps) |
| 186 | self.num_prev_frames = int(opt.num_prev_frames) |
| 187 | |
| 188 | self.wav2vec2 = Wav2VecModel.from_pretrained(opt.wav2vec_model_path, local_files_only = True) |
| 189 | self.wav2vec2.feature_extractor._freeze_parameters() |
| 190 | |
| 191 | for name, param in self.wav2vec2.named_parameters(): |
| 192 | param.requires_grad = False |
| 193 | |
| 194 | audio_input_dim = 768 if opt.only_last_features else 12 * 768 |
| 195 | |
| 196 | self.audio_projection = nn.Sequential( |
| 197 | nn.Linear(audio_input_dim, opt.dim_w), |
| 198 | nn.LayerNorm(opt.dim_w), |
| 199 | nn.SiLU() |
| 200 | ) |
| 201 | |
| 202 | def get_wav2vec2_feature(self, a: torch.Tensor, seq_len:int) -> torch.Tensor: |
| 203 | a = self.wav2vec2(a, seq_len=seq_len, output_hidden_states = not self.only_last_features) |
| 204 | if self.only_last_features: |
| 205 | a = a.last_hidden_state |
| 206 | else: |
| 207 | a = torch.stack(a.hidden_states[1:], dim=1).permute(0, 2, 1, 3) |
| 208 | a = a.reshape(a.shape[0], a.shape[1], -1) |
| 209 | return a |
| 210 | |
| 211 | def forward(self, a:torch.Tensor, prev_a:torch.Tensor = None) -> torch.Tensor: |
| 212 | if prev_a is not None: |
| 213 | a = torch.cat([prev_a, a], dim = 1) |
| 214 | if a.shape[1] % int( (self.num_frames_for_clip + self.num_prev_frames) * self.opt.sampling_rate / self.opt.fps) != 0: |
| 215 | a = F.pad(a, (0, int((self.num_frames_for_clip + self.num_prev_frames) * self.opt.sampling_rate / self.opt.fps) - a.shape[1]), mode='replicate') |
| 216 | a = self.get_wav2vec2_feature(a, seq_len = self.num_frames_for_clip + self.num_prev_frames) |
| 217 | else: |
| 218 | if a.shape[1] % int( self.num_frames_for_clip * self.opt.sampling_rate / self.opt.fps) != 0: |
| 219 | a = F.pad(a, (0, int(self.num_frames_for_clip * self.opt.sampling_rate / self.opt.fps) - a.shape[1]), mode = 'replicate') |
| 220 | a = self.get_wav2vec2_feature(a, seq_len = self.num_frames_for_clip) |
| 221 | |
| 222 | return self.audio_projection(a) # frame by frame |
| 223 | |
| 224 | @torch.no_grad() |
| 225 | def inference(self, a: torch.Tensor, seq_len:int) -> torch.Tensor: |
| 226 | if a.shape[1] % int(seq_len * self.opt.sampling_rate / self.opt.fps) != 0: |
| 227 | a = F.pad(a, (0, int(seq_len * self.opt.sampling_rate / self.opt.fps) - a.shape[1]), mode = 'replicate') |
| 228 | a = self.get_wav2vec2_feature(a, seq_len=seq_len) |
| 229 | return self.audio_projection(a) |
| 230 | |
| 231 | |
| 232 | |