(self, batch, train=True)
| 314 | return last_hidden_states, attention_mask |
| 315 | |
| 316 | def preprocess(self, batch, train=True): |
| 317 | target_wavs = batch["target_wavs"] |
| 318 | wav_lengths = batch["wav_lengths"] |
| 319 | |
| 320 | dtype = target_wavs.dtype |
| 321 | bs = target_wavs.shape[0] |
| 322 | device = target_wavs.device |
| 323 | |
| 324 | # SSL constraints |
| 325 | mert_ssl_hidden_states = None |
| 326 | mhubert_ssl_hidden_states = None |
| 327 | if train: |
| 328 | with torch.amp.autocast(device_type="cuda", dtype=dtype): |
| 329 | mert_ssl_hidden_states = self.infer_mert_ssl(target_wavs, wav_lengths) |
| 330 | mhubert_ssl_hidden_states = self.infer_mhubert_ssl( |
| 331 | target_wavs, wav_lengths |
| 332 | ) |
| 333 | |
| 334 | # 1: text embedding |
| 335 | texts = batch["prompts"] |
| 336 | encoder_text_hidden_states, text_attention_mask = self.get_text_embeddings( |
| 337 | texts, device |
| 338 | ) |
| 339 | encoder_text_hidden_states = encoder_text_hidden_states.to(dtype) |
| 340 | |
| 341 | target_latents, _ = self.dcae.encode(target_wavs, wav_lengths) |
| 342 | attention_mask = torch.ones( |
| 343 | bs, target_latents.shape[-1], device=device, dtype=dtype |
| 344 | ) |
| 345 | |
| 346 | speaker_embds = batch["speaker_embs"].to(dtype) |
| 347 | keys = batch["keys"] |
| 348 | lyric_token_ids = batch["lyric_token_ids"] |
| 349 | lyric_mask = batch["lyric_masks"] |
| 350 | |
| 351 | # cfg |
| 352 | if train: |
| 353 | full_cfg_condition_mask = torch.where( |
| 354 | (torch.rand(size=(bs,), device=device) < 0.15), |
| 355 | torch.zeros(size=(bs,), device=device), |
| 356 | torch.ones(size=(bs,), device=device), |
| 357 | ).long() |
| 358 | # N x T x 768 |
| 359 | encoder_text_hidden_states = torch.where( |
| 360 | full_cfg_condition_mask.unsqueeze(1).unsqueeze(1).bool(), |
| 361 | encoder_text_hidden_states, |
| 362 | torch.zeros_like(encoder_text_hidden_states), |
| 363 | ) |
| 364 | |
| 365 | full_cfg_condition_mask = torch.where( |
| 366 | (torch.rand(size=(bs,), device=device) < 0.50), |
| 367 | torch.zeros(size=(bs,), device=device), |
| 368 | torch.ones(size=(bs,), device=device), |
| 369 | ).long() |
| 370 | # N x 512 |
| 371 | speaker_embds = torch.where( |
| 372 | full_cfg_condition_mask.unsqueeze(1).bool(), |
| 373 | speaker_embds, |
no test coverage detected