Update app.py

app.py

@@ -4,7 +4,6 @@ from transformers import AutoTokenizer, AutoModelForTextToWaveform
 import torch
 from scipy.io.wavfile import write
 import numpy as np
-import uuid
 import io
 
 app = FastAPI()
@@ -12,7 +11,7 @@ app = FastAPI()
 # Load model and tokenizer
 model_name = "facebook/musicgen-medium"
 tokenizer = AutoTokenizer.from_pretrained(model_name)
-model = AutoModelForTextToWaveform.from_pretrained(model_name)
+model = AutoModelForTextToWaveform.from_pretrained(model_name, attn_implementation="eager")
 
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 model.to(device)
@@ -30,9 +29,12 @@ def generate_music(request: MusicRequest):
         sampling_rate = model.config.sampling_rate
         audio_values = audio_values.cpu().numpy().squeeze()
         
+        # Normalize audio values to fit int16 range
+        audio_values = np.clip(audio_values * 32767, -32768, 32767).astype(np.int16)
+        
         # Convert audio to bytes
         audio_bytes = io.BytesIO()
-        write(audio_bytes, sampling_rate, np.int16(audio_values * 32767))
+        write(audio_bytes, sampling_rate, audio_values)
         audio_bytes.seek(0)
         
         return Response(content=audio_bytes.read(), media_type="audio/wav", headers={"Content-Disposition": "attachment; filename=generated_music.wav"})