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):
  """