Update app.py

app.py

@@ -25,12 +25,13 @@ in_space = os.getenv("SYSTEM") == "spaces"
 
 # =================================================================================================
 
+@spaces.GPU
 def generate_drums(notes_times,
                    max_drums_limit = 8,
                    num_memory_tokens = 4096,
                    temperature=0.9):
 
-    x = torch.tensor([notes_times] * 1, dtype=torch.long, device=DEVICE)
+    x = torch.tensor([notes_times] * 1, dtype=torch.long, device='cuda')
 
     o = 128
 
@@ -62,6 +63,44 @@ def GenerateDrums(input_midi, input_num_tokens):
     print('Req start time: {:%Y-%m-%d %H:%M:%S}'.format(datetime.datetime.now(PDT)))
     start_time = reqtime.time()
 
+    print('Loading model...')
+
+    SEQ_LEN = 8192 # Models seq len
+    PAD_IDX = 393 # Models pad index
+    DEVICE = 'cuda' # 'cuda'
+
+    # instantiate the model
+
+    model = TransformerWrapper(
+        num_tokens = PAD_IDX+1,
+        max_seq_len = SEQ_LEN,
+        attn_layers = Decoder(dim = 1024, depth = 4, heads = 8, attn_flash = True)
+        )
+    
+    model = AutoregressiveWrapper(model, ignore_index = PAD_IDX)
+
+    model.to(DEVICE)
+    print('=' * 70)
+
+    print('Loading model checkpoint...')
+
+    model.load_state_dict(
+        torch.load('Ultimate_Drums_Transformer_Small_Trained_Model_VER3_VEL_11222_steps_0.5749_loss_0.8085_acc.pth',
+                   map_location=DEVICE))
+    print('=' * 70)
+
+    model.eval()
+
+    if DEVICE == 'cpu':
+        dtype = torch.bfloat16
+    else:
+        dtype = torch.float16
+
+    ctx = torch.amp.autocast(device_type=DEVICE, dtype=dtype)
+
+    print('Done!')
+    print('=' * 70)
+
     fn = os.path.basename(input_midi.name)
     fn1 = fn.split('.')[0]
 
@@ -239,45 +278,7 @@ if __name__ == "__main__":
     opt = parser.parse_args()
 
     soundfont = "SGM-v2.01-YamahaGrand-Guit-Bass-v2.7.sf2"
-
-    print('Loading model...')
-
-    SEQ_LEN = 8192 # Models seq len
-    PAD_IDX = 393 # Models pad index
-    DEVICE = 'cuda' # 'cuda'
-
-    # instantiate the model
-
-    model = TransformerWrapper(
-        num_tokens = PAD_IDX+1,
-        max_seq_len = SEQ_LEN,
-        attn_layers = Decoder(dim = 1024, depth = 4, heads = 8, attn_flash = True)
-        )
-    
-    model = AutoregressiveWrapper(model, ignore_index = PAD_IDX)
-
-    model.to(DEVICE)
-    print('=' * 70)
-
-    print('Loading model checkpoint...')
-
-    model.load_state_dict(
-        torch.load('Ultimate_Drums_Transformer_Small_Trained_Model_VER3_VEL_11222_steps_0.5749_loss_0.8085_acc.pth',
-                   map_location=DEVICE))
-    print('=' * 70)
-
-    model.eval()
-
-    if DEVICE == 'cpu':
-        dtype = torch.bfloat16
-    else:
-        dtype = torch.float16
-
-    ctx = torch.amp.autocast(device_type=DEVICE, dtype=dtype)
-
-    print('Done!')
-    print('=' * 70)
-    
+   
     app = gr.Blocks()
     with app:
         gr.Markdown("<h1 style='text-align: center; margin-bottom: 1rem'>Ultimate Drums Transformer</h1>")