Update app.py

app.py

@@ -1,27 +1,9 @@
 import gradio as gr
-import os
 import torch
 import torchaudio
+from google.colab import drive
+drive.mount('/content/drive')
 from torch import nn
-from torch.utils.data import DataLoader, Dataset
-from torchvision import transforms
-import numpy as np
-import IPython.display as ipd
-
-# Generate new kick drum samples
-generator.eval()
-with torch.no_grad():
-    for i in range(num_generated_samples):
-        noise = torch.randn(1, latent_dim).to(device)
-        generated_sample = generator(noise).squeeze().cpu()
-
-        # Save the generated sample
-        output_filename = f"generated_kick_{i+1}.wav"
-        torchaudio.save(output_filename, generated_sample.unsqueeze(0), 16000)
-
-        # Play the generated sample
-        print(f"Generated Sample {i+1}:")
-        display(ipd.Audio(output_filename))
 
 # Load the saved generator model
 class Generator(nn.Module):
@@ -39,11 +21,13 @@ class Generator(nn.Module):
     def forward(self, x):
         return self.generator(x)
 
+
 latent_dim = 100
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 generator = Generator(latent_dim).to(device)
-generator_model_path = "generator_model.pkl"
-generator.load_state_dict(torch.load(generator_model_path))
+generator_model_path = '/content/drive/MyDrive/generator_model.pkl'
+generator.load_state_dict(torch.load(generator_model_path, map_location=device))
+
 
 def generate_kick_drums():
     # Define the number of samples you want to generate
@@ -65,10 +49,11 @@ def generate_kick_drums():
 
     return tuple(output_files)
 
+# Define Gradio interface
 def gradio_interface():
     generate_button = gr.Interface(fn=generate_kick_drums,
                                    inputs=None, 
-                                   outputs=[gr.Audio(type='filepath', label=f"generated_kick_{i+1}") for i in range(3)], 
+                                   outputs=[gr.Audio(type='filepath', label=f"generated_kick_{i}") for i in range(3)], 
                                    live=True)
     generate_button.launch(debug=True)