load models cuda

main.py

@@ -71,12 +71,24 @@ def chords_string_to_list(chords: str):
     chrd_times = [x.split(',') for x in chords[1:-1].split('),(')]
     return [(x[0], float(x[1])) for x in chrd_times]
 
+# Create necessary directories
+os.makedirs("models", exist_ok=True)
+
+@spaces.GPU
 def load_model(version='facebook/jasco-chords-drums-400M'):
     global MODEL
     print("Loading model", version)
     if MODEL is None or MODEL.name != version:
         MODEL = None
-        MODEL = JASCO.get_pretrained(version)
+        try:
+            MODEL = JASCO.get_pretrained(version, device='cuda')
+            MODEL.name = version
+        except Exception as e:
+            raise gr.Error(f"Error loading model: {str(e)}")
+        
+        if MODEL is None:
+            raise gr.Error("Failed to load model")
+    return MODEL
 
 @spaces.GPU
 def _do_predictions(texts, chords, melody_matrix, drum_prompt, progress=False, gradio_progress=None, **gen_kwargs):