Update app.py

app.py

@@ -37,6 +37,7 @@ print('Loading aux Giant Music Transformer modules...')
 import matplotlib.pyplot as plt
 
 import gradio as gr
+import spaces
 
 print('=' * 70)
 print('PyTorch version:', torch.__version__)
@@ -47,50 +48,6 @@ print('=' * 70)
 
 #==================================================================================
 
-print('=' * 70)
-print('Instantiating model...')
-
-device_type = 'cuda'
-dtype = 'bfloat16'
-
-ptdtype = {'bfloat16': torch.bfloat16, 'float16': torch.float16}[dtype]
-ctx = torch.amp.autocast(device_type=device_type, dtype=ptdtype)
-
-SEQ_LEN = 8192
-PAD_IDX = 19463
-
-model = TransformerWrapper(
-        num_tokens = PAD_IDX+1,
-        max_seq_len = SEQ_LEN,
-        attn_layers = Decoder(dim = 2048,
-                              depth = 8,
-                              heads = 32,
-                              rotary_pos_emb = True,
-                              attn_flash = True
-                              )
-)
-
-model = AutoregressiveWrapper(model, ignore_index=PAD_IDX, pad_value=PAD_IDX)
-
-print('=' * 70)
-print('Loading model checkpoint...')
-
-model_path = 'Giant-Music-Transformer/Models/Medium/Giant_Music_Transformer_Medium_Trained_Model_10446_steps_0.7202_loss_0.8233_acc.pth'
-
-model.load_state_dict(torch.load(model_path))
-
-print('=' * 70)
-
-model.cuda()
-model.eval()
-
-print('Done!')
-print('=' * 70)
-print('Model will use', dtype, 'precision...')
-print('=' * 70)
-
-#==================================================================================
-
 SOUDFONT_PATH = 'SGM-v2.01-YamahaGrand-Guit-Bass-v2.7'
 
 NUM_OUT_BATCHES = 8
@@ -249,6 +206,7 @@ def save_midi(tokens, batch_number=None):
 
 #==================================================================================
 
+@spaces.GPU
 def generate_music(prime, 
                     num_gen_tokens, 
                     num_gen_batches,
@@ -258,6 +216,55 @@ def generate_music(prime,
                     model_sampling_top_p
                   ):
 
+
+    #==============================================================================
+    
+    print('=' * 70)
+    print('Instantiating model...')
+    
+    device_type = 'cuda'
+    dtype = 'bfloat16'
+    
+    ptdtype = {'bfloat16': torch.bfloat16, 'float16': torch.float16}[dtype]
+    ctx = torch.amp.autocast(device_type=device_type, dtype=ptdtype)
+    
+    SEQ_LEN = 8192
+    PAD_IDX = 19463
+    
+    model = TransformerWrapper(
+            num_tokens = PAD_IDX+1,
+            max_seq_len = SEQ_LEN,
+            attn_layers = Decoder(dim = 2048,
+                                  depth = 8,
+                                  heads = 32,
+                                  rotary_pos_emb = True,
+                                  attn_flash = True
+                                  )
+    )
+    
+    model = AutoregressiveWrapper(model, ignore_index=PAD_IDX, pad_value=PAD_IDX)
+    
+    print('=' * 70)
+    print('Loading model checkpoint...')
+    
+    model_path = 'Giant-Music-Transformer/Models/Medium/Giant_Music_Transformer_Medium_Trained_Model_10446_steps_0.7202_loss_0.8233_acc.pth'
+    
+    model.load_state_dict(torch.load(model_path))
+    
+    print('=' * 70)
+    
+    model.cuda()
+    model.eval()
+    
+    print('Done!')
+    print('=' * 70)
+    print('Model will use', dtype, 'precision...')
+    print('=' * 70)
+
+    #==============================================================================
+
+    print('Generating...')
+
     if not prime:
         inputs = [19461]
 
@@ -290,6 +297,9 @@ def generate_music(prime,
     
     output = out.tolist()
     
+    print('Done!')
+    print('=' * 70)    
+    
     return output
     
 #==================================================================================