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