updated gpu decorators

app.py

@@ -155,7 +155,10 @@ def generate_audio(audio_model, f0s, invert_audio_fn, singers=[3], num_steps=100
     
 @spaces.GPU(duration=150)
 def generate(pitch, num_samples=1, num_steps=100, singers=[3], outfolder='temp', audio_seq_len=750, pitch_qt=None, type='response', invert_pitch_fn=None):
-    
+    global pitch_model, audio_model
+    # move the models to device
+    pitch_model = pitch_model.to(device)
+    audio_model = audio_model.to(device)
     logging.log(logging.INFO, 'Generate function')
     # load pitch values onto GPU
     pitch = torch.tensor(pitch).float().unsqueeze(0).unsqueeze(0).to(device)
@@ -173,7 +176,6 @@ def generate(pitch, num_samples=1, num_steps=100, singers=[3], outfolder='temp',
         # if there is not pitch quantile transformer, undo the default quantile transformation that occurs
         def undo_qt(x, min_clip=200):
             pitch= pitch_qt.inverse_transform(x.reshape(-1, 1)).reshape(1, -1)
-            print(pitch.device)
             pitch = torch.round(pitch) # round to nearest integer, done in preprocessing of pitch contour fed into model
             pitch[pitch < 200] = np.nan
             return pitch
@@ -197,7 +199,7 @@ audio_model, audio_qt, audio_seq_len, invert_audio_fn = load_audio_fns(
     os.path.join(audio_path, 'last.ckpt'),
     qt_path = os.path.join(audio_path, 'qt.joblib'),
     config_path = os.path.join(audio_path, 'config.gin'),
-    device = device
+    device = 'cpu'
 )
 
 
@@ -209,7 +211,7 @@ def load_pitch_model(model_selection):
         model_type = 'diffusion', \
         config_path = os.path.join(pitch_path, 'config.gin'), \
         qt_path = os.path.join(pitch_path, 'qt.joblib'), \
-        device = device
+        device = 'cpu'
     )
     return pitch_model, pitch_qt, pitch_task_fn, invert_pitch_fn