fix on slice overwrite on slicer module

app.py

@@ -1,12 +1,13 @@
 import os
 import sys
+import numpy as np
 import torch
 import gradio as gr
 from vae_module import VAE, Encoder, Decoder, loss_function
 from config import config
 from slicer_module import get_slices
 from diffusers import UNet2DConditionModel, DDPMScheduler
-
+from mel_module import Mel
 
 vae = VAE()
 vae.load_state_dict(torch.load('vae_model_state_dict.pth', map_location=torch.device('cpu')))
@@ -16,14 +17,45 @@ vae.eval()
 model = UNet2DConditionModel.from_pretrained(config.hub_model_id, subfolder="unet")
 noise_scheduler = DDPMScheduler.from_pretrained(config.hub_model_id, subfolder="scheduler")
 
+
+
 def generate_new_track(audio_paths):
 
     for i, audio_path in enumerate(audio_paths):
+        print(audio_paths,audio_path)
         get_slices(audio_path)
-    
-    return
+
+    embedding = get_embedding()
+    print(embedding)
+
 
     
+def get_embedding(): # returns middle point of given audio files latent representations
+    latents = []
+    slices_dir = 'slices'
+    
+    for slice_file in os.listdir(slices_dir):
+        if slice_file.endswith('.wav'): # make sure the file is audio 
+            mel = Mel(os.path.join(slices_dir, slice_file))
+            spectrogram = mel.get_spectrogram()
+            tensor = torch.tensor(spectrogram).float().unsqueeze(0).unsqueeze(0)
+            mu, log_var = vae.encode(tensor)
+            latent = torch.cat((mu, log_var), dim=1)
+            min_val = latent.min()
+            max_val = latent.max()
+            normalized_tensor = 2 * ((latent - min_val) / (max_val - min_val)) - 1
+            latent = normalized_tensor.unsqueeze(0)
+            print(latent.shape)
+            latents.append(latent)
+    
+    if not latents:
+        return None
+    
+    latents_tensor = torch.cat(latents, dim=0)
+    mean_latent = latents_tensor.mean(dim=0, keepdim=True)
+    return mean_latent
+    
+
 # Define the Gradio interface
 interface = gr.Interface(
     fn=generate_new_track,

mel_module.py

@@ -0,0 +1,101 @@
+from typing import Optional
+from config import config
+import numpy as np
+import librosa
+from PIL import Image 
+
+import warnings
+warnings.filterwarnings("ignore", category=UserWarning, module='librosa')
+
+class Mel:
+    def __init__(
+            self,
+            file_path: str = None,
+            spectrogram: Optional[np.ndarray] = None,
+            image: Image.Image = None,
+            x_res: int = config.image_size,
+            y_res: int = config.image_size,
+            sample_rate: int = config.sample_rate,
+            n_fft: int = 2048,
+            hop_length: int = 882,
+            top_db: int = 80,
+            n_iter: int = 32,
+            ):
+        self.hop_length = hop_length
+        self.sr = sample_rate
+        self.n_fft = n_fft
+        self.top_db = top_db
+        self.n_iter = n_iter
+        self.x_res = x_res
+        self.y_res = y_res
+        self.n_mels = self.y_res
+        self.slice_size = self.x_res * self.hop_length - 1
+        self.file_path = file_path
+        self.spectrogram = spectrogram
+        self.image = image
+
+        if file_path is not None and not isinstance(file_path, str):
+            raise ValueError("file_path must be a string")
+        if spectrogram is not None and not isinstance(spectrogram, np.ndarray):
+            raise ValueError("spectrogram must be an ndarray")
+        if image is not None and not isinstance(image, Image.Image):
+            raise ValueError("image must be a PIL Image")
+
+        if file_path is not None:
+            self.load_file()
+        elif image is not None:
+            self.load_spectrogram()
+        elif spectrogram is not None:
+            self.load_image()
+        else:
+            print("Both file path and image are None!")
+
+    def load_file(self):
+        try:
+            # Load audio
+            if ".wav" in self.file_path:
+                audio, _ = librosa.load(self.file_path, mono=True, sr=self.sr)
+            # Pad audio if necessary
+            if len(audio) < self.x_res * self.hop_length:
+                audio = np.concatenate([audio, np.zeros((self.x_res * self.hop_length - len(audio),))])
+            # Compute mel spectrogram
+            S = librosa.feature.melspectrogram(
+                y=audio, sr=self.sr, n_fft=self.n_fft, hop_length=self.hop_length, n_mels=self.n_mels, fmax=self.sr//2
+            )
+            log_S = librosa.power_to_db(S, ref=np.max, top_db=self.top_db)
+            log_S = log_S[:self.y_res, :self.x_res]  # Ensure the spectrogram is of the desired size
+            self.spectrogram = (((log_S + self.top_db) * 255 / self.top_db).clip(0, 255) + 0.5).astype(np.uint8)
+            self.image = Image.fromarray(self.spectrogram)
+        except Exception as e:
+            print(f"Error loading {self.file_path}: {e}")
+
+    def load_spectrogram(self):
+        self.spectrogram = np.array(self.image)
+
+    def load_image(self):
+        self.spectrogram = self.spectrogram.astype("uint8")
+        self.image = Image.fromarray(self.spectrogram)
+
+    def get_spectrogram(self):
+        return self.spectrogram
+
+    def get_image(self):
+        return self.image
+
+    def get_audio(self):
+        log_S = self.spectrogram.astype("float") * self.top_db / 255 - self.top_db
+        S = librosa.db_to_power(log_S)
+        audio = librosa.feature.inverse.mel_to_audio(
+            S, sr=self.sr, n_fft=self.n_fft, hop_length=self.hop_length, n_iter=self.n_iter
+        )
+        display(Audio(audio, rate=self.sr))
+
+    def plot_spectrogram(self):
+        plt.figure(figsize=(10, 4))
+        plt.imshow(self.spectrogram, aspect='auto', origin='lower', cmap='viridis')
+        plt.colorbar(label='Magnitude')
+        plt.title('Mel Spectrogram')
+        plt.xlabel('Time (frames)')
+        plt.ylabel('Frequency (Mel bins)')
+        plt.tight_layout()
+        plt.show()

requirements.txt

@@ -1,4 +1,8 @@
+accelerate
 diffusers
 torch
 librosa
-soundfile
+soundfile
+Optional
+pillow
+numpy

slicer_module.py

@@ -6,24 +6,28 @@ def get_slices(file_path, sample_rate=44100, slice_duration=10, output_dir='slic
     os.makedirs(output_dir, exist_ok=True)
 
     audio, sr = librosa.load(file_path, sr=sample_rate)
-    
     slice_samples = slice_duration * sample_rate
 
+    existing_slices = [f for f in os.listdir(output_dir) if f.endswith('.wav')]
+    if existing_slices:
+        max_index = max(int(f.split('_')[1].split('.')[0]) for f in existing_slices)
+        start_index = max_index + 1
+    else:
+        start_index = 0
+
     num_slices = len(audio) // slice_samples
     for i in range(num_slices):
         start_sample = i * slice_samples
         end_sample = start_sample + slice_samples
         audio_slice = audio[start_sample:end_sample]
 
-        # save it into /slices
-        output_file = os.path.join(output_dir, f'slice_{i:04d}.wav')
+        output_file = os.path.join(output_dir, f'slice_{start_index + i:04d}.wav')
         sf.write(output_file, audio_slice, sample_rate)
 
-    # handle last slice
     if len(audio) % slice_samples != 0:
         start_sample = num_slices * slice_samples
         audio_slice = audio[start_sample:]
-        output_file = os.path.join(output_dir, f'slice_{num_slices:04d}.wav')
+        output_file = os.path.join(output_dir, f'slice_{start_index + num_slices:04d}.wav')
         sf.write(output_file, audio_slice, sample_rate)
 
 if __name__ == "__main__":