Added GPU decorators

app.py

@@ -10,6 +10,11 @@ import pretty_midi
 import librosa
 import soundfile as sf
 from midi2audio import FluidSynth
+import spaces
+
+# Remove CPU forcing since we'll use ZeroGPU
+# os.environ["CUDA_VISIBLE_DEVICES"] = ""
+# torch.set_num_threads(4)
 
 from aria.image_encoder import ImageEncoder
 from aria.aria import ARIA
@@ -163,6 +168,7 @@ def convert_midi_to_wav(midi_path):
         print(f"Error converting MIDI to WAV: {str(e)}")
         return None
 
+@spaces.GPU # Set duration to 60 seconds for music generation
 def generate_music(image, conditioning_type, gen_len, temperature, top_p, min_instruments):
     """Generate music from input image"""
     model = get_model(conditioning_type)

aria/aria.py

@@ -3,6 +3,7 @@ import os
 from PIL import Image
 import numpy as np
 import datetime
+import spaces
 
 from .image_encoder import ImageEncoder
 
@@ -14,6 +15,7 @@ sys.path.append(MIDI_EMOTION_PATH)
 class ARIA:
     """ARIA model that generates music from images based on emotional content."""
     
+    @spaces.GPU  # Model loading should be quick
     def __init__(
         self,
         image_model_checkpoint: str,
@@ -29,21 +31,8 @@ class ARIA:
             conditioning: Type of conditioning to use (continuous_concat, continuous_token, discrete_token)
             device: Device to run on (default: auto-detect)
         """
-        # Initialize CUDA if available
-        if device is None:
-            if not torch.cuda.is_available():
-                self.device = torch.device("cpu")
-            else:
-                try:
-                    # Test CUDA initialization
-                    torch.zeros(1).cuda()
-                    self.device = torch.device("cuda")
-                except RuntimeError:
-                    print("CUDA initialization failed, falling back to CPU")
-                    self.device = torch.device("cpu")
-        else:
-            self.device = torch.device(device)
-            
+        # Initialize device
+        self.device = torch.device("cuda")  # Always use CUDA with ZeroGPU
         print(f"Using device: {self.device}")
         self.conditioning = conditioning
         
@@ -71,6 +60,7 @@ class ARIA:
         self.midi_model.load_state_dict(torch.load(model_fp, map_location=self.device, weights_only=True))
         self.midi_model.eval()
     
+    @spaces.GPU
     @torch.inference_mode()  # More efficient than no_grad for inference
     def generate(
         self,

requirements.txt

@@ -1,11 +1,12 @@
+torch>=2.1.0
+torchvision>=0.16.0
+numpy>=1.21.0
+Pillow>=10.0.0
 gradio>=4.0.0
-torch>=2.0.0
-numpy>=1.24.0
 matplotlib>=3.7.0
-Pillow>=10.0.0
-huggingface-hub>=0.19.0
-pretty-midi>=0.2.10
+huggingface_hub>=0.19.0
+pretty-midi>=0.2.9
 librosa>=0.10.0
-soundfile>=0.12.0 
+soundfile>=0.12.0
 midi2audio>=0.1.1
-transformers>=4.30.0
+transformers>=4.35.0