Update app.py

app.py

@@ -11,6 +11,7 @@ import re
 import tqdm
 
 import gradio as gr
+import spaces
 
 from x_transformer_1_23_2 import *
 import random
@@ -116,20 +117,14 @@ print('Loading model checkpoint...')
 
 model_path = 'Popular_Hook_Transformer_Small_Trained_Model_10869_steps_0.2308_loss_0.9252_acc.pth'
 
-model.load_state_dict(torch.load(model_path, map_location=DEVICE))
-
-print('=' * 70)
-
-model.to(DEVICE)
-model.eval()
-
-ctx = torch.amp.autocast(device_type=DEVICE, dtype=torch.bfloat16)
+model.load_state_dict(torch.load(model_path, map_location='cpu'))
 
 print('Done!')
 print('=' * 70)
 
 #====================================================================================
 
+@spaces.GPU
 def Generate_POP_Section(input_comp_section, 
                          input_mode_time, 
                          input_mode_dur,
@@ -167,9 +162,13 @@ def Generate_POP_Section(input_comp_section,
 
         seq += input_seq
 
+    model.to(DEVICE)
+    model.eval()
+
     x = torch.LongTensor(seq).to(DEVICE)
     
-    with ctx:
+    with torch.amp.autocast(device_type=DEVICE, dtype=torch.bfloat16):
+        
         out = model.generate(x,
                              512-len(seq),
                              temperature=0.9,