fix the mel/prediction in the trainer

e2_tts_pytorch/e2_tts.py

@@ -61,7 +61,7 @@ def __getitem__(self, shapes: str):
 Int   = TorchTyping(jaxtyping.Int)
 Bool  = TorchTyping(jaxtyping.Bool)
 
-E2TTSReturn = namedtuple('E2TTS', ['loss', 'cond', 'pred'])
+E2TTSReturn = namedtuple('E2TTS', ['loss', 'cond', 'pred_flow', 'pred_data'])
 
 # helpers
 
@@ -1168,4 +1168,4 @@ def forward(
 
         loss = loss[rand_span_mask].mean()
 
-        return E2TTSReturn(loss, cond, pred)
+        return E2TTSReturn(loss, cond, pred, x0 + pred)

e2_tts_pytorch/trainer.py

@@ -35,8 +35,12 @@ def exists(v):
 def default(v, d):
     return v if exists(v) else d
 
+def to_numpy(t):
+    return t.detach().cpu().numpy()
+
 # plot spectrogram 
 def plot_spectrogram(spectrogram):
+    spectrogram = to_numpy(spectrogram)
     fig, ax = plt.subplots(figsize=(10, 4))
     im = ax.imshow(spectrogram.T, aspect="auto", origin="lower", interpolation="none")
     plt.colorbar(im, ax=ax)
@@ -238,7 +242,7 @@ def train(self, train_dataset, epochs, batch_size, num_workers=12, save_step=100
                         dur_loss = self.duration_predictor(mel_spec, lens=batch.get('durations'))
                         self.writer.add_scalar('duration loss', dur_loss.item(), global_step)
 
-                    loss, cond, pred = self.model(mel_spec, text=text_inputs, lens=mel_lengths)
+                    loss, cond, pred, pred_data = self.model(mel_spec, text=text_inputs, lens=mel_lengths)
                     self.accelerator.backward(loss)
 
                     if self.max_grad_norm > 0 and self.accelerator.sync_gradients:
@@ -262,9 +266,9 @@ def train(self, train_dataset, epochs, batch_size, num_workers=12, save_step=100
 
                 if global_step % save_step == 0:
                     self.save_checkpoint(global_step)
-                    self.writer.add_figure("mel/target", plot_spectrogram(mel_spec[0,:,:].detach().cpu().numpy()), global_step)
-                    self.writer.add_figure("mel/mask", plot_spectrogram(cond[0,:,:].detach().cpu().numpy()), global_step)
-                    self.writer.add_figure("mel/prediction", plot_spectrogram(pred[0,:,:].detach().cpu().numpy()), global_step)
+                    self.writer.add_figure("mel/target", plot_spectrogram(mel_spec[0,:,:]), global_step)
+                    self.writer.add_figure("mel/mask", plot_spectrogram(cond[0,:,:]), global_step)
+                    self.writer.add_figure("mel/prediction", plot_spectrogram(pred_data[0,:,:]), global_step)
 
             epoch_loss /= len(train_dataloader)
             if self.accelerator.is_local_main_process:

pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "e2-tts-pytorch"
-version = "1.0.5"
+version = "1.0.6"
 description = "E2-TTS in Pytorch"
 authors = [
     { name = "Phil Wang", email = "lucidrains@gmail.com" }