add vae and slicer modules

create evaluate module as template
update app file to slice given songs

app.py

@@ -1,7 +1,37 @@
+import os
+import sys
+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
 
-def greet(name):
-    return "Hello " + name + "!!"
 
-demo = gr.Interface(fn=greet, inputs="text", outputs="text")
-demo.launch()
+vae = VAE()
+vae.load_state_dict(torch.load('vae_model_state_dict.pth', map_location=torch.device('cpu')))
+vae.to(config.device)
+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):
+        get_slices(audio_path)
+    
+    return
+
+    
+# Define the Gradio interface
+interface = gr.Interface(
+    fn=generate_new_track,
+    inputs=gr.Files(file_count="multiple", label="Upload Your Audio Files"),
+    outputs=gr.Audio(type="filepath", label="Generated Track"),
+    title="AMUSE: Music Generation",
+    description="Upload audio files and generate new tracks based on them using AMUSE."
+)
+
+# Launch the interface
+interface.launch()

config.py

@@ -0,0 +1,14 @@
+import torch
+class Config:
+    def __init__(self):
+        self.image_size: int = 512  # the generated image resolution
+        self.sample_rate: int = 44100  # the sample rate of the audio
+        self.guidance_scale: float = 1  # the guidance scale for the diffusion process
+        self.mixed_precision: str = "fp16"  # `no` for float32, `fp16` for automatic mixed precision
+        self.hub_model_id: str = "alppo/amuse"  # the name of the repository to create on the HF Hub
+        self.hub_dataset_id: str = "alppo/music"  # the name of the dataset to create on the HF Hub
+        self.seed: int = 0
+        self.device = "cuda" if torch.cuda.is_available() else "cpu"
+
+
+config = Config()

evaluate_module.py

@@ -0,0 +1 @@
+# to do : evaluate function

requirements.txt

@@ -0,0 +1,4 @@
+diffusers
+torch
+librosa
+soundfile

slicer_module.py

@@ -0,0 +1,30 @@
+import os
+import librosa
+import soundfile as sf
+
+def get_slices(file_path, sample_rate=44100, slice_duration=10, output_dir='slices'):
+    os.makedirs(output_dir, exist_ok=True)
+
+    audio, sr = librosa.load(file_path, sr=sample_rate)
+    
+    slice_samples = slice_duration * sample_rate
+
+    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')
+        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')
+        sf.write(output_file, audio_slice, sample_rate)
+
+if __name__ == "__main__":
+    get_slices('rock_song_009.wav')

vae_model.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:a8e271f76eea5c196e3d2a3b69b3a605610ae53fab76f1b2d451d96d13dfcfcd
+size 277897466

vae_model_state_dict.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:131c1e7d7707a044b41f794bd18297549553f29ccdf693f0a4eeaed177e28006
+size 277883050

vae_module.py

@@ -0,0 +1,131 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+class ResidualBlock(nn.Module):
+    def __init__(self, channels):
+        super(ResidualBlock, self).__init__()
+        self.conv1 = nn.Conv2d(channels, channels, kernel_size=3, padding=1)
+        self.bn1 = nn.BatchNorm2d(channels)
+        self.relu = nn.ReLU(inplace=True)
+        self.conv2 = nn.Conv2d(channels, channels, kernel_size=3, padding=1)
+        self.bn2 = nn.BatchNorm2d(channels)
+
+    def forward(self, x):
+        residual = x
+        out = self.conv1(x)
+        out = self.bn1(out)
+        out = self.relu(out)
+        out = self.conv2(out)
+        out = self.bn2(out)
+        out += residual
+        out = self.relu(out)
+        return out
+
+class Encoder(nn.Module):
+    def __init__(self, input_channels=1, hidden_dims=[64, 128, 256, 512, 1024], latent_dim=32):
+        super(Encoder, self).__init__()
+        self.hidden_dims = hidden_dims
+
+        # Build Encoder with Residual Blocks
+        modules = []
+        for h_dim in hidden_dims:
+            modules.append(
+                nn.Sequential(
+                    nn.Conv2d(input_channels, h_dim, kernel_size=3, stride=2, padding=1),
+                    nn.BatchNorm2d(h_dim),
+                    nn.LeakyReLU(),
+                    ResidualBlock(h_dim)  # Adding a residual block
+                )
+            )
+            input_channels = h_dim
+
+        self.encoder = nn.Sequential(*modules)
+        self.fc_mu = nn.Linear(hidden_dims[-1]*hidden_dims[-3], latent_dim)
+        self.fc_var = nn.Linear(hidden_dims[-1]*hidden_dims[-3], latent_dim)
+
+    def forward(self, x):
+        for layer in self.encoder:
+            x = layer(x)
+        x = torch.flatten(x, start_dim=1)
+        mu = self.fc_mu(x)
+        log_var = self.fc_var(x)
+        return mu, log_var
+
+class Decoder(nn.Module):
+    def __init__(self, latent_dim=32, output_channels=1, hidden_dims=[64, 128, 256, 512, 1024]):
+        super(Decoder, self).__init__()
+        self.hidden_dims = hidden_dims
+        # Reversing the order for the decoder
+        hidden_dims = hidden_dims[::-1]
+        self.decoder_input = nn.Linear(latent_dim, hidden_dims[0]*hidden_dims[2])
+
+        # Build Decoder with Residual Blocks
+        modules = []
+        for i in range(len(hidden_dims) - 1):
+            modules.append(
+                nn.Sequential(
+                    nn.ConvTranspose2d(hidden_dims[i], hidden_dims[i+1], kernel_size=3, stride=2, padding=1, output_padding=1),
+                    nn.BatchNorm2d(hidden_dims[i+1]),
+                    nn.LeakyReLU(),
+                    ResidualBlock(hidden_dims[i+1])  # Adding a residual block
+                )
+            )
+
+        self.decoder = nn.Sequential(*modules)
+        self.final_layer = nn.Sequential(
+            nn.ConvTranspose2d(hidden_dims[-1], hidden_dims[-1], kernel_size=3, stride=2, padding=1, output_padding=1),
+            nn.BatchNorm2d(hidden_dims[-1]),
+            nn.LeakyReLU(),
+            nn.Conv2d(hidden_dims[-1], output_channels, kernel_size=3, padding=1),
+            nn.Sigmoid()
+        )
+
+    def forward(self, z):
+        z = self.decoder_input(z)
+        z = z.view(-1, 1024, 16, 16)
+        for layer in self.decoder:
+            z = layer(z)
+        result = self.final_layer(z)
+        return result
+
+class VAE(nn.Module):
+    def __init__(self,
+                 input_channels=1,
+                 latent_dim=32,
+                 hidden_dims=None):
+        super(VAE, self).__init__()
+
+        if hidden_dims is None:
+            hidden_dims = [64, 128, 256, 512, 1024]
+
+        self.encoder = Encoder(input_channels=input_channels,
+                               hidden_dims=hidden_dims,
+                               latent_dim=latent_dim)
+
+        self.decoder = Decoder(latent_dim=latent_dim,
+                               output_channels=input_channels,
+                               hidden_dims=hidden_dims)
+
+    def encode(self, input):
+        mu, log_var = self.encoder(input)
+        return mu, log_var
+
+    def reparameterize(self, mu, log_var):
+        std = torch.exp(0.5 * log_var)
+        eps = torch.randn_like(std)
+        return mu + eps * std
+
+    def decode(self, z):
+        return self.decoder(z)
+
+    def forward(self, input):
+        mu, log_var = self.encode(input)
+        z = self.reparameterize(mu, log_var)
+        return self.decode(z), mu, log_var
+
+# Loss function for VAE
+def loss_function(recon_x, x, mu, log_var):
+    BCE = F.binary_cross_entropy(recon_x, x, reduction='sum')
+    KLD = -0.5 * torch.sum(1 + log_var - mu.pow(2) - log_var.exp())
+    return BCE + KLD