Update audioldm_train/modules/latent_diffusion/ddpm.py

audioldm_train/modules/latent_diffusion/ddpm.py

@@ -1335,7 +1335,7 @@ class LatentDiffusion(DDPM):
         waveform = self.first_stage_model.vocoder(mel)
         waveform = waveform.cpu().detach().numpy()
         if save:
-            self.save_waveform(waveform, savepath, name, n_gen)
+            self.save_waveform(waveform, savepath="./")
         return waveform
 
     def encode_first_stage(self, x):
@@ -1818,44 +1818,31 @@ class LatentDiffusion(DDPM):
             **kwargs,
         )
 
-    def save_waveform(self, waveform, savepath, name="outwav", n_gen=1):
-        print(f'debug_name : {name}')
-        if type(name) != str and len(name[0][1]) > 1:
-            name = list(name[0][1])
-            name = [_.decode() if type(_) is bytes else _ for _ in name]
-        n_gen = int(waveform.shape[0] / len(name))
-        assert len(name) * n_gen == waveform.shape[0]
-        lenn = len(name)
-        for i in range(n_gen - 1):
-            for x in range(lenn):
-                name.append(name[x]) 
-        assert len(name) == waveform.shape[0]
-        for i in range(waveform.shape[0]):
-            if type(name) is str:
-                path = os.path.join(savepath, "%s_%s_%s.wav" % (self.global_step, i, name))
-            elif type(name) is list:
-                path = os.path.join(
-                    savepath,
-                    "%s.wav"
-                    % (
-                        os.path.basename(name[i])
+    def save_waveform(self, waveform, savepath="./", name="awesome.wav", n_gen=1):
+      print(f'debug_name : {name}')
+      
+      # If `name` is a list, join the elements into a string or select the first element
+      if isinstance(name, list):
+          name = "_".join(name)  # Joins the list elements with an underscore
+          name += ".wav"  # Ensures the file has a `.wav` extension
+      elif not isinstance(name, str):
+          raise TypeError("Name must be a string or list")
+      
+      # Normalize the energy of the waveform
+      todo_waveform = waveform[0, 0]  # Assuming you are only saving the first waveform
+      todo_waveform = (todo_waveform / np.max(np.abs(todo_waveform))) * 0.8
+      
+      # Define the path where to save the file
+      path = os.path.join(savepath, name)
+      
+      try:
+          # Save the waveform to the specified path
+          sf.write(path, todo_waveform, samplerate=self.sampling_rate)
+          print(f'Waveform saved at -> {path}')
+      except Exception as e:
+          print(f'Error saving waveform: {e}')
+
 
-                        if (not ".wav" in name[i])
-                        else os.path.basename(name[i]).split(".")[0]
-                    ),
-                )
-            else:
-                # import pdb 
-                # pdb.set_trace()
-                raise NotImplementedError
-            todo_waveform = waveform[i, 0]
-            todo_waveform = (
-                todo_waveform / np.max(np.abs(todo_waveform))
-            ) * 0.8  # Normalize the energy of the generation output
-            try:
-                sf.write(path, todo_waveform, samplerate=self.sampling_rate)
-            except:
-                print('waveform name ERROR!!!!!!!!!!!!')
 
     @torch.no_grad()
     def sample_log(
@@ -2054,7 +2041,7 @@ class LatentDiffusion(DDPM):
                         print("Choose the following indexes:", best_index)
                     except Exception as e:
                         print("Warning: while calculating CLAP score (not fatal), ", e)
-                self.save_waveform(waveform, waveform_save_path, name=fnames, n_gen=n_gen)
+                self.save_waveform(waveform, savepath="./")
         return waveform_save_path