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"