add lora

app.py

@@ -142,7 +142,12 @@ def get_duration(model_name, tab, mid_seq, continuation_state, continuation_sele
 def run(model_name, tab, mid_seq, continuation_state, continuation_select, instruments, drum_kit, bpm, time_sig,
         key_sig, mid, midi_events, reduce_cc_st, remap_track_channel, add_default_instr, remove_empty_channels,
         seed, seed_rand, gen_events, temp, top_p, top_k, allow_cc):
-    model = models[model_name]
+    model, lora_name = models[model_name]
+    if lora_name is None and model.peft_loaded():
+        model.disable_adapters()
+    elif lora_name is not None:
+        model.enable_adapters()
+        model.set_adapter(lora_name)
     model.to(device=opt.device)
     tokenizer = model.tokenizer
     bpm = int(bpm)
@@ -253,7 +258,7 @@ def finish_run(model_name, mid_seq):
     if mid_seq is None:
         outputs = [None] * OUTPUT_BATCH_SIZE
         return *outputs, []
-    tokenizer = models[model_name].tokenizer
+    tokenizer = models[model_name][0].tokenizer
     outputs = []
     end_msgs = [create_msg("progress", [0, 0])]
     if not os.path.exists("outputs"):
@@ -277,7 +282,7 @@ def render_audio(model_name, mid_seq, should_render_audio):
     if (not should_render_audio) or mid_seq is None:
         outputs = [None] * OUTPUT_BATCH_SIZE
         return tuple(outputs)
-    tokenizer = models[model_name].tokenizer
+    tokenizer = models[model_name][0].tokenizer
     outputs = []
     if not os.path.exists("outputs"):
         os.mkdir("outputs")
@@ -294,7 +299,7 @@ def render_audio(model_name, mid_seq, should_render_audio):
 def undo_continuation(model_name, mid_seq, continuation_state):
     if mid_seq is None or len(continuation_state) < 2:
         return mid_seq, continuation_state, send_msgs([])
-    tokenizer = models[model_name].tokenizer
+    tokenizer = models[model_name][0].tokenizer
     if isinstance(continuation_state[-1], list):
         mid_seq = continuation_state[-1]
     else:
@@ -364,12 +369,21 @@ if __name__ == "__main__":
     thread_pool = ThreadPoolExecutor(max_workers=OUTPUT_BATCH_SIZE)
     synthesizer = MidiSynthesizer(soundfont_path)
     models_info = {
-        "generic pretrain model (tv2o-medium) by skytnt": ["skytnt/midi-model-tv2o-medium", "", "tv2o-medium"],
-        "generic pretrain model (tv2o-large) by asigalov61": ["asigalov61/Music-Llama", "", "tv2o-large"],
-        "generic pretrain model (tv2o-medium) by asigalov61": ["asigalov61/Music-Llama-Medium", "", "tv2o-medium"],
-        "generic pretrain model (tv1-medium) by skytnt": ["skytnt/midi-model", "", "tv1-medium"],
-        "j-pop finetune model (tv2o-medium) by skytnt": ["skytnt/midi-model-ft", "jpop-tv2o-medium/", "tv2o-medium"],
-        "touhou finetune model (tv2o-medium) by skytnt": ["skytnt/midi-model-ft", "touhou-tv2o-medium/", "tv2o-medium"],
+        "generic pretrain model (tv2o-medium) by skytnt": [
+            "skytnt/midi-model-tv2o-medium", "", "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", {}
+        ],
+        "generic pretrain model (tv2o-medium) by asigalov61": [
+            "asigalov61/Music-Llama-Medium", "", "tv2o-medium", {}
+        ],
+        "generic pretrain model (tv1-medium) by skytnt": [
+            "skytnt/midi-model", "", "tv1-medium", {}
+        ]
     }
     models = {}
     if opt.device == "cuda":
@@ -379,14 +393,20 @@ 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) in models_info.items():
+    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)
+        for lora_name, lora_repo in loras.items():
+            model.load_adapter(lora_repo, lora_name)
+        if loras:
+            model.disable_adapters()
         model.to(device="cpu", dtype=torch.float32).eval()
-        models[name] = model
+        models[name] = model, None
+        for lora_name, lora_repo in loras.items():
+            models[f"{name} with {lora_name} lora"] = model, lora_name
 
     load_javascript()
     app = gr.Blocks()

midi_model.py

@@ -6,6 +6,7 @@ import torch.nn as nn
 import torch.nn.functional as F
 import tqdm
 from transformers import LlamaModel, LlamaConfig
+from transformers.integrations import PeftAdapterMixin
 
 from midi_tokenizer import MIDITokenizerV1, MIDITokenizerV2, MIDITokenizer
 
@@ -55,7 +56,7 @@ class MIDIModelConfig:
             raise ValueError(f"Unknown model size {size}")
 
 
-class MIDIModel(nn.Module):
+class MIDIModel(nn.Module, PeftAdapterMixin):
     def __init__(self, config: MIDIModelConfig, *args, **kwargs):
         super(MIDIModel, self).__init__()
         self.tokenizer = config.tokenizer
@@ -69,6 +70,9 @@ class MIDIModel(nn.Module):
             self.device = kwargs["device"]
         return super(MIDIModel, self).to(*args, **kwargs)
 
+    def peft_loaded(self):
+        return self._hf_peft_config_loaded
+
     def forward_token(self, hidden_state, x=None):
         """