Update app.py

app.py

@@ -9,6 +9,7 @@ import subprocess
 import socket
 import time
 import atexit
+import torch
 from transformers import AutoModel, AutoProcessor
 
 PROXY_URL = None
@@ -287,7 +288,7 @@ EXAMPLE_YOUTUBE_PROMPTS = [
     ],
 ]
 processor = AutoProcessor.from_pretrained(MODEL_ID)
-model = AutoModel.from_pretrained(MODEL_ID, device_map="auto").eval()
+model = AutoModel.from_pretrained(MODEL_ID, device_map="cpu", torch_dtype=torch.float32).eval()
 
 _youtube_cache = {}
 
@@ -412,7 +413,10 @@ def download_youtube_audio(url, force_reload=False):
 
 @spaces.GPU
 def infer(audio_path, youtube_url, prompt_text):
-    try:        
+    try:
+        device = "cuda" if torch.cuda.is_available() else "cpu"
+        model.to(device)
+        
         final_audio_path = None
         status_message = ""