adding flowvae

flowae/models/diffusion/fm.py

@@ -12,6 +12,7 @@ class FM:
         self.timescale = timescale
         self.use_immiscible = use_immiscible
         self.k_candidates = k_candidates
+        print('use_immiscible: ', use_immiscible, 'k_candidates: ', k_candidates)   
         
     def alpha(self, t):
         return 1.0 - t

flowae/models/networks/consistency_audio_decoder_unet.py

@@ -274,9 +274,9 @@ class AudioDiffusionUNet(nn.Module):
                 size=x.shape[-1], 
                 mode='linear'  # or 'linear' for smoother interpolation
             )
-        
+        print('shape of z_proj: ', z_proj.shape)
         # Add latent conditioning to audio features
-        return x + z_proj
+        return torch.cat([x, z_proj], dim=1)
 
     def forward(self, x, t=None, z_dec=None) -> torch.Tensor:
         """
@@ -288,11 +288,13 @@ class AudioDiffusionUNet(nn.Module):
             z_dec: [batch, 64, n_frames] - latent conditioning (any length)
         """
         # Embed audio input
+        print('shape of x: ', x.shape, 'shape of z_dec: ', z_dec.shape)
         x = self.embed_audio(x)  # [batch, c0, samples]
-        
+        print('shape of x: ', x.shape)
         # Add latent conditioning
         if z_dec is not None:
             x = self.condition_with_latents(x, z_dec)
+        print('shape of x: ', x.shape)
 
         # Embed timestep
         if t is None:

flowae/trainers/audio_ldm_trainer.py

@@ -412,15 +412,30 @@ class AudioLDMTrainer(BaseTrainer):
             for i in range(min(signal.batch_size, 5)):  # Save up to 5 per batch
                 idx = int(os.environ['RANK']) + cnt * int(os.environ['WORLD_SIZE'])
                 if max_samples is None or idx < max_samples:
+                    tmp_recon = recons[i].audio_data.cpu().numpy()
+                    if tmp_recon.dim() == 3:
+                        tmp_recon = tmp_recon.squeeze(0)
+                    elif tmp_recon.dim() == 1:
+                        tmp_recon = tmp_recon.unsqueeze(0)
+                    tmp_recon = tmp_recon.T
+
+                    tmp_signal = signal[i].audio_data.cpu().numpy()
+                    if tmp_signal.dim() == 3:
+                        tmp_signal = tmp_signal.squeeze(0)
+                    elif tmp_signal.dim() == 1:
+                        tmp_signal = tmp_signal.unsqueeze(0)
+                    tmp_signal = tmp_signal.T
+
+                    
                     # Save as wav files
                     sf.write(
                         os.path.join(cache_gen_dir, f'{idx}.wav'),
-                        recons[i].audio_data.cpu().numpy().T,
+                        tmp_recon,
                         int(recons[i].sample_rate)
                     )
                     sf.write(
                         os.path.join(cache_gt_dir, f'{idx}.wav'),
-                        signal[i].audio_data.cpu().numpy().T,
+                        tmp_signal,
                         int(signal[i].sample_rate)
                     )
                 cnt += 1
@@ -493,14 +508,29 @@ class AudioLDMTrainer(BaseTrainer):
             for i in range(min(gt_signal.batch_size, 5)):
                 idx = int(os.environ['RANK']) + cnt * int(os.environ['WORLD_SIZE'])
                 if max_samples is None or idx < max_samples:
+                    tmp_recon = pred_signal[i].audio_data.cpu().numpy()
+                    if tmp_recon.dim() == 3:
+                        tmp_recon = tmp_recon.squeeze(0)
+                    elif tmp_recon.dim() == 1:
+                        tmp_recon = tmp_recon.unsqueeze(0)
+                    tmp_recon = tmp_recon.T
+
+                    tmp_signal = gt_signal[i].audio_data.cpu().numpy()
+                    if tmp_signal.dim() == 3:
+                        tmp_signal = tmp_signal.squeeze(0)
+                    elif tmp_signal.dim() == 1:
+                        tmp_signal = tmp_signal.unsqueeze(0)
+                    tmp_signal = tmp_signal.T
+
+
                     sf.write(
                         os.path.join(cache_gen_dir, f'{idx}.wav'),
-                        pred_signal[i].audio_data.cpu().numpy().T,
+                        tmp_recon,
                         int(pred_signal[i].sample_rate)
                     )
                     sf.write(
                         os.path.join(cache_gt_dir, f'{idx}.wav'),
-                        gt_signal[i].audio_data.cpu().numpy().T,
+                        tmp_signal,
                         int(gt_signal[i].sample_rate)
                     )
                 cnt += 1