merge lora into model

app.py

@@ -142,12 +142,7 @@ 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, 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 = models[model_name]
     model.to(device=opt.device)
     tokenizer = model.tokenizer
     bpm = int(bpm)
@@ -258,7 +253,7 @@ def finish_run(model_name, mid_seq):
     if mid_seq is None:
         outputs = [None] * OUTPUT_BATCH_SIZE
         return *outputs, []
-    tokenizer = models[model_name][0].tokenizer
+    tokenizer = models[model_name].tokenizer
     outputs = []
     end_msgs = [create_msg("progress", [0, 0])]
     if not os.path.exists("outputs"):
@@ -282,7 +277,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][0].tokenizer
+    tokenizer = models[model_name].tokenizer
     outputs = []
     if not os.path.exists("outputs"):
         os.mkdir("outputs")
@@ -293,13 +288,15 @@ def render_audio(model_name, mid_seq, should_render_audio):
         audio_futures.append(audio_future)
     for future in audio_futures:
         outputs.append((44100, future.result()))
+    if OUTPUT_BATCH_SIZE == 1:
+        return outputs[0]
     return tuple(outputs)
 
 
 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][0].tokenizer
+    tokenizer = models[model_name].tokenizer
     if isinstance(continuation_state[-1], list):
         mid_seq = continuation_state[-1]
     else:
@@ -399,14 +396,15 @@ if __name__ == "__main__":
         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, None
+        models[name] = model
         for lora_name, lora_repo in loras.items():
-            models[f"{name} with {lora_name} lora"] = model, lora_name
+            model = MIDIModel(config=MIDIModelConfig.from_name(config))
+            model.load_state_dict(state_dict, strict=False)
+            print(f"loading lora {lora_repo} for {name}")
+            model = model.load_merge_lora(lora_repo)
+            model.to(device="cpu", dtype=torch.float32).eval()
+            models[f"{name} with {lora_name} lora"] = model
 
     load_javascript()
     app = gr.Blocks()

midi_model.py

@@ -5,6 +5,7 @@ import torch
 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
 from transformers.integrations import PeftAdapterMixin
 
@@ -75,6 +76,13 @@ class MIDIModel(nn.Module, PeftAdapterMixin):
     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)
+        set_peft_model_state_dict(self, adapter_state_dict, "default")
+        return model.merge_and_unload()
+
     def forward_token(self, hidden_state, x=None):
         """