|
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 |
|
from generator_module import Generator |
|
import shutil |
|
|
|
slices_folder = 'slices' |
|
|
|
if os.path.exists(slices_folder): |
|
shutil.rmtree(slices_folder) |
|
|
|
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, progress=gr.Progress(track_tqdm=True)): |
|
for i, audio_path in enumerate(audio_paths): |
|
print(audio_paths, audio_path) |
|
get_slices(audio_path) |
|
|
|
embedding = get_embedding() |
|
print("sample latent", embedding.shape) |
|
|
|
generator = Generator(config, model, noise_scheduler, vae, embedding, progress_callback=progress) |
|
generator.generate() |
|
|
|
return config.generated_track_path |
|
|
|
def get_embedding(): |
|
latents = [] |
|
slices_dir = 'slices' |
|
|
|
for slice_file in os.listdir(slices_dir): |
|
if slice_file.endswith('.wav'): |
|
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) |
|
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 |
|
|
|
|
|
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 = ( |
|
"<h3>Welcome to the AMUSE music generation app</h3>" |
|
"<p>Here's how it works:</p>" |
|
"<ol>" |
|
"<li><strong>Upload Your Audio Files:</strong> Provide audio files from which the taste will be extracted, " |
|
"and a new track will be generated accordingly. The audio files should be in .wav format!</li>" |
|
"<li><strong>Process:</strong> The app slices the audio, extracts features, and generates a new track using a VAE and a diffusion model.</li>" |
|
"<li><strong>Progress:</strong> The progress bar will show the generation process in real-time. Note that this takes a significant amount of time, " |
|
"so you may leave the site in the free version and come back later to see the result.</li>" |
|
"<li><strong>Download:</strong> Once the track is generated, you can download it directly.</li>" |
|
"</ol>" |
|
"<h4>Notes:</h4>" |
|
"<ul>" |
|
"<li>As mentioned earlier, it takes a significant amount of time to generate a new track in the free version of HF Spaces. " |
|
"So, submit your tracks and forget about it for a little while :) Then come back to see the new track.</li>" |
|
"<li>Ensure your audio files are clean and of good quality for the best results (sample rate: 44100 and .wav format).</li>" |
|
"</ul>" |
|
) |
|
) |
|
|
|
|
|
|
|
|
|
interface.launch() |
|
|
