(self, batch, batch_idx)
| 486 | return timesteps |
| 487 | |
| 488 | def run_step(self, batch, batch_idx): |
| 489 | self.plot_step(batch, batch_idx) |
| 490 | ( |
| 491 | keys, |
| 492 | target_latents, |
| 493 | attention_mask, |
| 494 | encoder_text_hidden_states, |
| 495 | text_attention_mask, |
| 496 | speaker_embds, |
| 497 | lyric_token_ids, |
| 498 | lyric_mask, |
| 499 | mert_ssl_hidden_states, |
| 500 | mhubert_ssl_hidden_states, |
| 501 | ) = self.preprocess(batch) |
| 502 | |
| 503 | target_image = target_latents |
| 504 | device = target_image.device |
| 505 | dtype = target_image.dtype |
| 506 | # Step 1: Generate random noise, initialize settings |
| 507 | noise = torch.randn_like(target_image, device=device) |
| 508 | bsz = target_image.shape[0] |
| 509 | timesteps = self.get_timestep(bsz, device) |
| 510 | |
| 511 | # Add noise according to flow matching. |
| 512 | sigmas = self.get_sd3_sigmas( |
| 513 | timesteps=timesteps, device=device, n_dim=target_image.ndim, dtype=dtype |
| 514 | ) |
| 515 | noisy_image = sigmas * noise + (1.0 - sigmas) * target_image |
| 516 | |
| 517 | # This is the flow-matching target for vanilla SD3. |
| 518 | target = target_image |
| 519 | |
| 520 | # SSL constraints for CLAP and vocal_latent_channel2 |
| 521 | all_ssl_hiden_states = [] |
| 522 | if mert_ssl_hidden_states is not None: |
| 523 | all_ssl_hiden_states.append(mert_ssl_hidden_states) |
| 524 | if mhubert_ssl_hidden_states is not None: |
| 525 | all_ssl_hiden_states.append(mhubert_ssl_hidden_states) |
| 526 | |
| 527 | # N x H -> N x c x W x H |
| 528 | x = noisy_image |
| 529 | # Step 5: Predict noise |
| 530 | transformer_output = self.transformers( |
| 531 | hidden_states=x, |
| 532 | attention_mask=attention_mask, |
| 533 | encoder_text_hidden_states=encoder_text_hidden_states, |
| 534 | text_attention_mask=text_attention_mask, |
| 535 | speaker_embeds=speaker_embds, |
| 536 | lyric_token_idx=lyric_token_ids, |
| 537 | lyric_mask=lyric_mask, |
| 538 | timestep=timesteps.to(device).to(dtype), |
| 539 | ssl_hidden_states=all_ssl_hiden_states, |
| 540 | ) |
| 541 | model_pred = transformer_output.sample |
| 542 | proj_losses = transformer_output.proj_losses |
| 543 | |
| 544 | # Follow: Section 5 of https://arxiv.org/abs/2206.00364. |
| 545 | # Preconditioning of the model outputs. |
no test coverage detected