Update app.py

app.py

@@ -1,17 +1,19 @@
 from fastapi import FastAPI, HTTPException, Response
 from pydantic import BaseModel
-from transformers import AutoTokenizer, AutoModelForTextToWaveform
+from transformers import AutoTokenizer, AutoModel
 import torch
-from scipy.io.wavfile import write
 import numpy as np
 import io
+from scipy.io.wavfile import write
+from PIL import Image
+import riffusion
 
 app = FastAPI()
 
-# Load model and tokenizer
-model_name = "facebook/musicgen-medium"
+# Load Riffusion model
+model_name = "riffusion/riffusion-model-v1"
+model = AutoModel.from_pretrained(model_name)
 tokenizer = AutoTokenizer.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)
@@ -24,15 +26,18 @@ def generate_music(request: MusicRequest):
     try:
         inputs = tokenizer(request.prompt, return_tensors="pt").to(device)
         with torch.no_grad():
-            audio_values = model.generate(**inputs)
+            spectrogram = model.generate(**inputs)
+        
+        # Convert spectrogram to an image (since Riffusion outputs spectrograms)
+        spectrogram_image = Image.fromarray((spectrogram.cpu().numpy().squeeze() * 255).astype(np.uint8))
         
-        sampling_rate = model.config.sampling_rate
-        audio_values = audio_values.cpu().numpy().squeeze()
+        # Convert spectrogram to audio
+        audio_values, sampling_rate = riffusion.audio_processing.spectrogram_to_audio(spectrogram_image)
         
-        # Normalize audio values to fit int16 range
+        # Normalize and convert to int16
         audio_values = np.clip(audio_values * 32767, -32768, 32767).astype(np.int16)
         
-        # Convert audio to bytes
+        # Convert to WAV format
         audio_bytes = io.BytesIO()
         write(audio_bytes, sampling_rate, audio_values)
         audio_bytes.seek(0)
@@ -43,4 +48,4 @@ def generate_music(request: MusicRequest):
 
 @app.get("/")
 def root():
-    return {"message": "Welcome to the Music Generation API"}
+    return {"message": "Welcome to the Riffusion Music Generation API"}