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: