flash

app.py

@@ -223,13 +223,16 @@ if __name__ == "__main__":
                    "touhou finetune model": ["skytnt/midi-model-ft", "touhou/"],
                    }
     device = "cuda" if torch.cuda.is_available() else "cpu"
+    if device=="cuda": # flash attn
+        torch.backends.cuda.enable_mem_efficient_sdp(True)
+        torch.backends.cuda.enable_flash_sdp(True)
     models = {}
     tokenizer = MIDITokenizer()
     for name, (repo_id, path) in models_info.items():
 
         model_path = hf_hub_download_retry(repo_id=repo_id, filename=f"{path}model.ckpt")
         model = MIDIModel(tokenizer).to(device=device)
-        ckpt = torch.load(model_path)
+        ckpt = torch.load(model_path, weights_only=True)
         state_dict = ckpt.get("state_dict", ckpt)
         model.load_state_dict(state_dict, strict=False)
         model.eval()

midi_model.py

@@ -9,7 +9,7 @@ from midi_tokenizer import MIDITokenizer
 
 
 class MIDIModel(nn.Module):
-    def __init__(self, tokenizer: MIDITokenizer, n_layer=12, n_head=16, n_embd=1024, n_inner=4096, flash=False,
+    def __init__(self, tokenizer: MIDITokenizer, n_layer=12, n_head=16, n_embd=1024, n_inner=4096,
                  *args, **kwargs):
         super(MIDIModel, self).__init__()
         self.tokenizer = tokenizer
@@ -21,9 +21,6 @@ class MIDIModel(nn.Module):
                                                 hidden_size=n_embd, num_attention_heads=n_head // 4,
                                                 num_hidden_layers=n_layer // 4, intermediate_size=n_inner // 4,
                                                 pad_token_id=tokenizer.pad_id, max_position_embeddings=4096))
-        if flash:
-            torch.backends.cuda.enable_mem_efficient_sdp(True)
-            torch.backends.cuda.enable_flash_sdp(True)
         self.lm_head = nn.Linear(n_embd, tokenizer.vocab_size, bias=False)
         self.device = "cpu"