Update app.py

app.py

@@ -62,54 +62,45 @@ SOUDFONT_PATH = 'SGM-v2.01-YamahaGrand-Guit-Bass-v2.7.sf2'
 
 #==================================================================================
 
-def load_model():
+print('=' * 70)
+print('Instantiating model...')
 
-    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 = 2048
+device_type = 'cuda'
+dtype = 'bfloat16'
 
-    if model_selector == 'with velocity - 3 epochs':
-        PAD_IDX = 512
+ptdtype = {'bfloat16': torch.bfloat16, 'float16': torch.float16}[dtype]
+ctx = torch.amp.autocast(device_type=device_type, dtype=ptdtype)
 
-    else:
-        PAD_IDX = 384
-    
-    model = TransformerWrapper(
-            num_tokens = PAD_IDX+1,
-            max_seq_len = SEQ_LEN,
-            attn_layers = Decoder(dim = 2048,
-                                  depth = 4,
-                                  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_checkpoint = hf_hub_download(repo_id='asigalov61/Guided-Accompaniment-Transformer', filename=MODEL_CHECKPOINTS[model_selector])
-    
-    model.load_state_dict(torch.load(model_checkpoint, map_location='cpu', weights_only=True))
-    
-    model = torch.compile(model, mode='max-autotune')
-    
-    print('=' * 70)
-    print('Done!')
-    print('=' * 70)
-    print('Model will use', dtype, 'precision...')
-    print('=' * 70)
+SEQ_LEN = 4096
+PAD_IDX = 1794
+
+model = TransformerWrapper(
+        num_tokens = PAD_IDX+1,
+        max_seq_len = SEQ_LEN,
+        attn_layers = Decoder(dim = 2048,
+                              depth = 4,
+                              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...')      
 
-    return [model, ctx]
+model_checkpoint = hf_hub_download(repo_id='asigalov61/Guided-Accompaniment-Transformer', filename=MODEL_CHECKPOINTS[model_selector])
+
+model.load_state_dict(torch.load(model_checkpoint, map_location='cpu', weights_only=True))
+
+model = torch.compile(model, mode='max-autotune')
+
+print('=' * 70)
+print('Done!')
+print('=' * 70)
+print('Model will use', dtype, 'precision...')
+print('=' * 70)
 
 #==================================================================================