Update app.py

app.py

@@ -1,10 +1,7 @@
 from audiocraft.models import MusicGen
 import streamlit as st
-import os
 import torch
 import torchaudio
-import numpy as np
-import base64
 from io import BytesIO
 
 @st.cache_resource
@@ -17,12 +14,11 @@ def generate_music_tensors(description, duration: int):
     print("Duration:", duration)
     model = load_model()
 
-    # Experiment with different generation parameters for improved quality
     model.set_generation_params(
         use_sampling=True,
-        top_k=300,          # Increase top_k for more diversity
-        top_p=0.85,         # Probability threshold for token sampling
-        temperature=0.8,    # Control randomness; lower values = more focused output
+        top_k=300,
+        top_p=0.85,
+        temperature=0.8,
         duration=duration
     )
 
@@ -35,13 +31,11 @@ def generate_music_tensors(description, duration: int):
 
 def save_audio_to_bytes(samples: torch.Tensor):
     sample_rate = 32000
-    assert samples.dim() == 2 or samples.dim() == 3
+    assert samples.dim() == 3  # Expecting (batch, channels, samples)
+    samples = samples[0]  # Take the first batch item
     samples = samples.detach().cpu()
 
-    if samples.dim() == 2:
-        samples = samples[None, ...]
-    
-    # Save audio to a byte buffer instead of file for easier download
+    # Save audio to a byte buffer instead of a file for easier download
     byte_io = BytesIO()
     torchaudio.save(byte_io, samples, sample_rate=sample_rate, format="wav")
     byte_io.seek(0)  # Reset buffer position to the beginning for reading
@@ -66,7 +60,7 @@ def main():
             "Description": text_area,
             "Selected duration": time_slider
         })
-        
+
         st.write("We will back with your music....please enjoy doing the rest of your tasks while we come back in some time :)")
         st.subheader("Generated Music")
         music_tensors = generate_music_tensors(text_area, time_slider)