hf

app.py

@@ -365,19 +365,19 @@ if __name__ == "__main__":
  synthesizer = MidiSynthesizer(soundfont_path)
  models_info = {
  "generic pretrain model (tv2o-medium) by skytnt": [
- "skytnt/midi-model-tv2o-medium", "", "tv2o-medium", {
+ "skytnt/midi-model-tv2o-medium", {
  "jpop": "skytnt/midi-model-tv2om-jpop-lora",
  "touhou": "skytnt/midi-model-tv2om-touhou-lora"
  }
  ],
  "generic pretrain model (tv2o-large) by asigalov61": [
- "asigalov61/Music-Llama", "", "tv2o-large", {}
+ "asigalov61/Music-Llama", {}
  ],
  "generic pretrain model (tv2o-medium) by asigalov61": [
- "asigalov61/Music-Llama-Medium", "", "tv2o-medium", {}
+ "asigalov61/Music-Llama-Medium", {}
  ],
  "generic pretrain model (tv1-medium) by skytnt": [
- "skytnt/midi-model", "", "tv1-medium", {}
+ "skytnt/midi-model", {}
  ]
  }
  models = {}
@@ -388,20 +388,15 @@ if __name__ == "__main__":
  torch.backends.cudnn.allow_tf32 = True
  torch.backends.cuda.enable_mem_efficient_sdp(True)
  torch.backends.cuda.enable_flash_sdp(True)
- for name, (repo_id, path, config, loras) in models_info.items():
- model_path = hf_hub_download_retry(repo_id=repo_id, filename=f"{path}model.ckpt")
- model = MIDIModel(config=MIDIModelConfig.from_name(config))
- ckpt = torch.load(model_path, map_location="cpu", weights_only=True)
- state_dict = ckpt.get("state_dict", ckpt)
- model.load_state_dict(state_dict, strict=False)
- model.to(device="cpu", dtype=torch.float32).eval()
+ for name, (repo_id, loras) in models_info.items():
+ model = MIDIModel.from_pretrained(repo_id)
+ model.to(device="cpu", dtype=torch.float32)
  models[name] = model
  for lora_name, lora_repo in loras.items():
- model = MIDIModel(config=MIDIModelConfig.from_name(config))
- model.load_state_dict(state_dict, strict=False)
+ model = MIDIModel.from_pretrained(repo_id)
  print(f"loading lora {lora_repo} for {name}")
  model = model.load_merge_lora(lora_repo)
- model.to(device="cpu", dtype=torch.float32).eval()
+ model.to(device="cpu", dtype=torch.float32)
  models[f"{name} with {lora_name} lora"] = model
 
  load_javascript()

app_onnx.py

@@ -432,18 +432,12 @@ def hf_hub_download_retry(repo_id, filename):
  raise err
 
 
-def get_tokenizer(config_name):
- tv, size = config_name.split("-")
- tv = tv[1:]
- if tv[-1] == "o":
- o = True
- tv = tv[:-1]
- else:
- o = False
- if tv not in ["v1", "v2"]:
- raise ValueError(f"Unknown tokenizer version {tv}")
- tokenizer = MIDITokenizer(tv)
- tokenizer.set_optimise_midi(o)
+def get_tokenizer(repo_id):
+ config_path = hf_hub_download_retry(repo_id=repo_id, filename=f"config.json")
+ with open(config_path, "r") as f:
+ config = json.load(f)
+ tokenizer = MIDITokenizer(config["tokenizer"]["version"])
+ tokenizer.set_optimise_midi(config["tokenizer"]["optimise_midi"])
  return tokenizer
 
 
@@ -468,34 +462,33 @@ if __name__ == "__main__":
  synthesizer = MidiSynthesizer(soundfont_path)
  models_info = {
  "generic pretrain model (tv2o-medium) by skytnt": [
- "skytnt/midi-model-tv2o-medium", "", "tv2o-medium", {
+ "skytnt/midi-model-tv2o-medium", "", {
  "jpop": "skytnt/midi-model-tv2om-jpop-lora",
  "touhou": "skytnt/midi-model-tv2om-touhou-lora"
  }
  ],
  "generic pretrain model (tv2o-large) by asigalov61": [
- "asigalov61/Music-Llama", "", "tv2o-large", {}
+ "asigalov61/Music-Llama", "", {}
  ],
  "generic pretrain model (tv2o-medium) by asigalov61": [
- "asigalov61/Music-Llama-Medium", "", "tv2o-medium", {}
+ "asigalov61/Music-Llama-Medium", "", {}
  ],
  "generic pretrain model (tv1-medium) by skytnt": [
- "skytnt/midi-model", "", "tv1-medium", {}
+ "skytnt/midi-model", "", {}
  ]
  }
  models = {}
  providers = ['CUDAExecutionProvider', 'CPUExecutionProvider']
  device = "cuda"
 
- for name, (repo_id, path, config, loras) in models_info.items():
+ for name, (repo_id, path, loras) in models_info.items():
  model_base_path = hf_hub_download_retry(repo_id=repo_id, filename=f"{path}onnx/model_base.onnx")
  model_token_path = hf_hub_download_retry(repo_id=repo_id, filename=f"{path}onnx/model_token.onnx")
- tokenizer = get_tokenizer(config)
+ tokenizer = get_tokenizer(repo_id)
  models[name] = [model_base_path, model_token_path, tokenizer]
  for lora_name, lora_repo in loras.items():
  model_base_path = hf_hub_download_retry(repo_id=lora_repo, filename=f"onnx/model_base.onnx")
  model_token_path = hf_hub_download_retry(repo_id=lora_repo, filename=f"onnx/model_token.onnx")
- tokenizer = get_tokenizer(config)
  models[f"{name} with {lora_name} lora"] = [model_base_path, model_token_path, tokenizer]
 
  load_javascript()

midi_model.py

@@ -1,4 +1,5 @@
-from typing import Union
+import json
+from typing import Union, Dict, Any
 
 import numpy as np
 import torch
@@ -6,21 +7,57 @@ import torch.nn as nn
 import torch.nn.functional as F
 import tqdm
 from peft import PeftConfig, LoraModel, load_peft_weights, set_peft_model_state_dict
-from transformers import LlamaModel, LlamaConfig, DynamicCache
-from transformers.integrations import PeftAdapterMixin
+from transformers import LlamaModel, LlamaConfig, DynamicCache, PretrainedConfig, PreTrainedModel
 
 from midi_tokenizer import MIDITokenizerV1, MIDITokenizerV2, MIDITokenizer
 
 config_name_list = ["tv1-medium", "tv2-medium", "tv2o-medium", "tv2-large", "tv2o-large"]
 
 
-class MIDIModelConfig:
- def __init__(self, tokenizer: Union[MIDITokenizerV1, MIDITokenizerV2],
- net_config: LlamaConfig, net_token_config: LlamaConfig):
- self.tokenizer = tokenizer
- self.net_config = net_config
- self.net_token_config = net_token_config
- self.n_embd = net_token_config.hidden_size
+class MIDIModelConfig(PretrainedConfig):
+ model_type = "midi_model"
+
+ def __init__(self,
+ tokenizer: Union[MIDITokenizerV1, MIDITokenizerV2, Dict]=None,
+ net_config: Union[LlamaConfig, Dict]=None,
+ net_token_config: Union[LlamaConfig, Dict]=None,
+ **kwargs):
+ super().__init__(**kwargs)
+ if tokenizer:
+ if isinstance(tokenizer, dict):
+ self.tokenizer = MIDITokenizer(tokenizer["version"])
+ self.tokenizer.set_optimise_midi(tokenizer["optimise_midi"])
+ else:
+ self.tokenizer = tokenizer
+ else:
+ self.tokenizer = MIDITokenizer()
+ if net_config:
+ if isinstance(net_config, dict):
+ self.net_config = LlamaConfig(**net_config)
+ else:
+ self.net_config = net_config
+ else:
+ self.net_config = LlamaConfig()
+ if net_token_config:
+ if isinstance(net_token_config, dict):
+ self.net_token_config = LlamaConfig(**net_token_config)
+ else:
+ self.net_token_config = net_token_config
+ else:
+ self.net_token_config = LlamaConfig()
+ self.n_embd = self.net_token_config.hidden_size
+
+ def to_dict(self) -> Dict[str, Any]:
+ d = super().to_dict()
+ d["tokenizer"] = self.tokenizer.to_dict()
+ return d
+
+ def __str__(self):
+ d = {
+ "net": self.net_config.to_json_string(use_diff=False),
+ "net_token": self.net_token_config.to_json_string(use_diff=False)
+ }
+ return json.dumps(d, indent=4)
 
  @staticmethod
  def get_config(tokenizer_ver="v2", optimise_midi=True, n_layer=12, n_head=16, n_embd=1024, n_inner=4096):
@@ -59,27 +96,20 @@ class MIDIModelConfig:
  raise ValueError(f"Unknown model size {size}")
 
 
-class MIDIModel(nn.Module, PeftAdapterMixin):
+class MIDIModel(PreTrainedModel):
+ config_class = MIDIModelConfig
+
  def __init__(self, config: MIDIModelConfig, *args, **kwargs):
- super(MIDIModel, self).__init__()
+ super(MIDIModel, self).__init__(config, *args, **kwargs)
  self.tokenizer = config.tokenizer
  self.net = LlamaModel(config.net_config)
  self.net_token = LlamaModel(config.net_token_config)
  self.lm_head = nn.Linear(config.n_embd, self.tokenizer.vocab_size, bias=False)
- self.device = "cpu"
-
- def to(self, *args, **kwargs):
- if "device" in kwargs:
- self.device = kwargs["device"]
- return super(MIDIModel, self).to(*args, **kwargs)
-
- def peft_loaded(self):
- return self._hf_peft_config_loaded
 
  def load_merge_lora(self, model_id):
  peft_config = PeftConfig.from_pretrained(model_id)
  model = LoraModel(self, peft_config, adapter_name="default")
- adapter_state_dict = load_peft_weights(model_id, device=self.device)
+ adapter_state_dict = load_peft_weights(model_id, device=str(self.device))
  set_peft_model_state_dict(self, adapter_state_dict, "default")
  return model.merge_and_unload()
 
@@ -164,7 +194,7 @@ class MIDIModel(nn.Module, PeftAdapterMixin):
  with bar:
  while cur_len < max_len:
  end = [False] * batch_size
- hidden = self.forward(input_tensor[:,past_len:], cache=cache1)[:, -1]
+ hidden = self.forward(input_tensor[:, past_len:], cache=cache1)[:, -1]
  next_token_seq = None
  event_names = [""] * batch_size
  cache2 = DynamicCache()

midi_tokenizer.py

@@ -1,4 +1,5 @@
 import random
+from typing import Dict, Any
 
 import PIL.Image
 import numpy as np
@@ -33,6 +34,20 @@ class MIDITokenizerV1:
  self.parameter_ids = {p: allocate_ids(s) for p, s in self.event_parameters.items()}
  self.max_token_seq = max([len(ps) for ps in self.events.values()]) + 1
 
+ def to_dict(self) -> Dict[str, Any]:
+ d = {
+ "version":self.version,
+ "optimise_midi":self.optimise_midi,
+ "vocab_size": self.vocab_size,
+ "events": self.events,
+ "event_parameters": self.event_parameters,
+ "max_token_seq": self.max_token_seq,
+ "pad_id": self.pad_id,
+ "bos_id": self.bos_id,
+ "eos_id": self.eos_id,
+ }
+ return d
+
  def set_optimise_midi(self, optimise_midi=True):
  self.optimise_midi = optimise_midi
 
@@ -519,6 +534,20 @@ class MIDITokenizerV2:
  self.parameter_ids = {p: allocate_ids(s) for p, s in self.event_parameters.items()}
  self.max_token_seq = max([len(ps) for ps in self.events.values()]) + 1
 
+ def to_dict(self) -> Dict[str, Any]:
+ d = {
+ "version":self.version,
+ "optimise_midi":self.optimise_midi,
+ "vocab_size": self.vocab_size,
+ "events": self.events,
+ "event_parameters": self.event_parameters,
+ "max_token_seq": self.max_token_seq,
+ "pad_id": self.pad_id,
+ "bos_id": self.bos_id,
+ "eos_id": self.eos_id,
+ }
+ return d
+
  def set_optimise_midi(self, optimise_midi=True):
  self.optimise_midi = optimise_midi