Update app.py

app.py

@@ -1,9 +1,9 @@
-import scipy.io.wavfile as wav
 import io
 from transformers import AutoProcessor, MusicgenForConditionalGeneration
 from IPython.display import Audio
 import torch
 import streamlit as st
+import wave
 
 def mu_gen(prompt):
     processor = AutoProcessor.from_pretrained("facebook/musicgen-small")
@@ -13,7 +13,7 @@ def mu_gen(prompt):
     model.to(device)
 
     inputs = processor(
-        text=[str(prompt)],  # This line is correct
+        text=[str(prompt)],
         padding=True,
         return_tensors="pt",
     )
@@ -25,11 +25,15 @@ def mu_gen(prompt):
     sampling_rate = model.config.audio_encoder.sampling_rate
 
     # Convert audio data to WAV format
-    wav_data = audio_values[0].numpy().tobytes()
+    wav_data = audio_values[0].numpy()
 
     # Create an in-memory WAV file
     with io.BytesIO() as wav_file:
-        wav.write(wav_file, sampling_rate, wav_data)
+        with wave.open(wav_file, 'wb') as wf:
+            wf.setnchannels(1)  # Mono audio
+            wf.setsampwidth(2)  # 16-bit audio
+            wf.setframerate(sampling_rate)
+            wf.writeframes(wav_data.tobytes())
         wav_bytes = wav_file.getvalue()
 
     return wav_bytes  # Return the WAV audio data as bytes