Update app.py

app.py

@@ -1,47 +1,51 @@
+import scipy.io.wavfile as wav
+import io
 from transformers import AutoProcessor, MusicgenForConditionalGeneration
 from IPython.display import Audio
-import scipy
 import torch
 import streamlit as st
 
-
 def mu_gen(prompt):
-  processor = AutoProcessor.from_pretrained("facebook/musicgen-small")
-  model = MusicgenForConditionalGeneration.from_pretrained("facebook/musicgen-small")
+    processor = AutoProcessor.from_pretrained("facebook/musicgen-small")
+    model = MusicgenForConditionalGeneration.from_pretrained("facebook/musicgen-small")
+
+    device = torch.device("cpu")
+    model.to(device)
 
-  device = torch.device("cpu")
-  model.to(device)
+    inputs = processor(
+        text=[str(prompt)],  # This line is correct
+        padding=True,
+        return_tensors="pt",
+    )
 
-  inputs = processor(
-      text = [str(prompt)],     # This line is correct
-      padding=True,
-      return_tensors="pt",
-  )
-  
-  inputs = {key: value.to(device) for key, value in inputs.items()}
+    inputs = {key: value.to(device) for key, value in inputs.items()}
 
     # Generate audio on CPU
-  audio_values = model.generate(**inputs, max_new_tokens=256)
-  sampling_rate = model.config.audio_encoder.sampling_rate
+    audio_values = model.generate(**inputs, max_new_tokens=256)
+    sampling_rate = model.config.audio_encoder.sampling_rate
 
-    # Create an Audio object from the generated audio
-  result = Audio(audio_values[0].numpy(), rate=sampling_rate)
+    # Convert audio data to WAV format
+    wav_data = audio_values[0].numpy().tobytes()
 
-  return result
+    # Create an in-memory WAV file
+    with io.BytesIO() as wav_file:
+        wav.write(wav_file, sampling_rate, wav_data)
+        wav_bytes = wav_file.getvalue()
 
+    return wav_bytes  # Return the WAV audio data as bytes
 
 def main():
-    st.title("Text to music")
+    st.title("Text to Music")
 
     # Input text prompt
     title = st.text_input('Write a prompt (จะใช้เวลาค่อนข้างมากในการสร้างเนื่องจากใช้ CPU ในการรันโมเดล)', "")
 
-    if st.button('Generate Image'):
+    if st.button('Generate Music'):
         # Call the mu_gen function to generate music using the 'title' prompt
-        generated_music = mu_gen(title)  # Replace 'prompt' with 'title'
+        generated_music = mu_gen(title)
 
-        # Display the generated audio
-        st.audio(generated_music, format='audio/mpeg', start_time=0)
+        # Display the generated audio in WAV format
+        st.audio(generated_music, format='audio/wav', start_time=0)
 
 if __name__ == '__main__':
     main()