add cache

app.py

@@ -13,6 +13,7 @@ import torch
 import torch.nn.functional as F
 import tqdm
 from huggingface_hub import hf_hub_download
+from transformers import DynamicCache
 
 import MIDI
 from midi_model import MIDIModel, MIDIModelConfig
@@ -51,12 +52,14 @@ def generate(model: MIDIModel, prompt=None, batch_size=1, max_len=512, temp=1.0,
  input_tensor = torch.from_numpy(prompt).to(dtype=torch.long, device=model.device)
  cur_len = input_tensor.shape[1]
  bar = tqdm.tqdm(desc="generating", total=max_len - cur_len)
+ cache1 = DynamicCache()
  with bar:
  while cur_len < max_len:
  end = [False] * batch_size
- hidden = model.forward(input_tensor)[:, -1]
+ hidden = model.forward(input_tensor[:, -1:], cache=cache1)[:, -1]
  next_token_seq = None
  event_names = [""] * batch_size
+ cache2 = DynamicCache()
  for i in range(max_token_seq):
  mask = torch.zeros((batch_size, tokenizer.vocab_size), dtype=torch.int64, device=model.device)
  for b in range(batch_size):
@@ -81,7 +84,11 @@ def generate(model: MIDIModel, prompt=None, batch_size=1, max_len=512, temp=1.0,
  mask_ids = [i for i in mask_ids if i not in disable_channels]
  mask[b, mask_ids] = 1
  mask = mask.unsqueeze(1)
- logits = model.forward_token(hidden, next_token_seq)[:, -1:]
+ x = next_token_seq
+ if i != 0:
+ hidden = None
+ x = x[:, -1:]
+ logits = model.forward_token(hidden, x, cache=cache2)[:, -1:]
  scores = torch.softmax(logits / temp, dim=-1) * mask
  samples = model.sample_top_p_k(scores, top_p, top_k, generator=generator)
  if i == 0:
@@ -118,21 +125,10 @@ def send_msgs(msgs):
  return json.dumps(msgs)
 
 
-def calc_time(x):
- return 5.849e-5*x**2 + 0.04781*x + 0.1168
-
 def get_duration(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):
- if tab == 0:
- start_events = 1
- elif tab == 1 and mid is not None:
- start_events = midi_events
- elif tab == 2 and mid_seq is not None:
- start_events = len(mid_seq[0])
- else:
- start_events = 1
- t = calc_time(start_events + gen_events) - calc_time(start_events) + 5
+ t = gen_events // 20 + 5
  if "large" in model_name:
  t *= 2
  return t

midi_model.py

@@ -6,7 +6,7 @@ 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 import LlamaModel, LlamaConfig, DynamicCache
 from transformers.integrations import PeftAdapterMixin
 
 from midi_tokenizer import MIDITokenizerV1, MIDITokenizerV2, MIDITokenizer
@@ -83,30 +83,40 @@ class MIDIModel(nn.Module, PeftAdapterMixin):
  set_peft_model_state_dict(self, adapter_state_dict, "default")
  return model.merge_and_unload()
 
- def forward_token(self, hidden_state, x=None):
+ def forward_token(self, hidden_state=None, x=None, cache=None):
  """
 
  :param hidden_state: (batch_size, n_embd)
  :param x: (batch_size, token_sequence_length)
+ :param cache: Cache
  :return: (batch_size, 1 + token_sequence_length, vocab_size)
  """
- hidden_state = hidden_state.unsqueeze(1) # (batch_size, 1, n_embd)
+ if hidden_state is not None:
+ #if you use cache, you don't need to pass in hidden_state
+ hidden_state = hidden_state.unsqueeze(1) # (batch_size, 1, n_embd)
  if x is not None:
  x = self.net_token.embed_tokens(x)
- hidden_state = torch.cat([hidden_state, x], dim=1)
- hidden_state = self.net_token.forward(inputs_embeds=hidden_state).last_hidden_state
+ if hidden_state is not None:
+ x = torch.cat([hidden_state, x], dim=1)
+ hidden_state = x
+ hidden_state = self.net_token.forward(inputs_embeds=hidden_state,
+ past_key_values=cache,
+ use_cache=cache is not None).last_hidden_state
  return self.lm_head(hidden_state)
 
- def forward(self, x):
+ def forward(self, x, cache = None):
  """
  :param x: (batch_size, midi_sequence_length, token_sequence_length)
+ :param cache: Cache
  :return: hidden (batch_size, midi_sequence_length, n_embd)
  """
 
  # merge token sequence
  x = self.net.embed_tokens(x)
  x = x.sum(dim=-2)
- x = self.net.forward(inputs_embeds=x)
+ x = self.net.forward(inputs_embeds=x,
+ past_key_values=cache,
+ use_cache=cache is not None)
  return x.last_hidden_state
 
  def sample_top_p_k(self, probs, p, k, generator=None):
@@ -149,12 +159,14 @@ class MIDIModel(nn.Module, PeftAdapterMixin):
 
  cur_len = input_tensor.shape[1]
  bar = tqdm.tqdm(desc="generating", total=max_len - cur_len)
+ cache1 = DynamicCache()
  with bar:
  while cur_len < max_len:
  end = [False] * batch_size
- hidden = self.forward(input_tensor)[:, -1]
+ hidden = self.forward(input_tensor[:,-1:], cache=cache1)[:, -1]
  next_token_seq = None
  event_names = [""] * batch_size
+ cache2 = DynamicCache()
  for i in range(max_token_seq):
  mask = torch.zeros((batch_size, tokenizer.vocab_size), dtype=torch.int64, device=self.device)
  for b in range(batch_size):
@@ -170,7 +182,12 @@ class MIDIModel(nn.Module, PeftAdapterMixin):
  continue
  mask[b, tokenizer.parameter_ids[param_names[i - 1]]] = 1
  mask = mask.unsqueeze(1)
- logits = self.forward_token(hidden, next_token_seq)[:, -1:]
+ x = next_token_seq
+ if i != 0:
+ # cached
+ hidden = None
+ x = x[:, -1:]
+ logits = self.forward_token(hidden, x, cache=cache2)[:, -1:]
  scores = torch.softmax(logits / temp, dim=-1) * mask
  samples = self.sample_top_p_k(scores, top_p, top_k, generator=generator)
  if i == 0: