update

.gitignore

@@ -0,0 +1,2 @@
+__pycache__
+.DS_Store

app.py

@@ -1,10 +1,70 @@
 import gradio as gr
+import torch
 import os
-def greet(name):
-    return "Hello " + name + "!!"
+from miditoolkit import MidiFile
+from src.model.generate import batch_performance_render, map_midi
+from src.model.pianoformer import PianoT5Gemma
 
-from transformers import AutoModel, AutoTokenizer
+# ------------------------------
+# Load model
+# ------------------------------
+def load_model():
+    print("Loading model...")
+    model = PianoT5Gemma.from_pretrained(
+        "yhj137/pianist-transformer-rendering",
+        token=os.environ.get("hf_token"),
+        torch_dtype=torch.bfloat16
+    )
+    model.eval()
+    return model
 
-model = AutoModel.from_pretrained("yhj137/pianist-transformer-rendering", token=os.environ["hf_token"])
-demo = gr.Interface(fn=greet, inputs="text", outputs="text")
-demo.launch()
+model = load_model()
+
+# ------------------------------
+# Define inference function
+# ------------------------------
+def render_midi(midi_file, temperature, top_p, top_k):
+    # Save uploaded file temporarily
+    input_path = midi_file.name
+    midi = MidiFile(input_path)
+
+    # Run inference
+    res = batch_performance_render(
+        model,
+        [midi],
+        temperature=temperature,
+        top_p=top_p,
+        top_k=top_k,
+        device="cpu"  # change to "cuda" if you use GPU space
+    )
+
+    # Map result and save
+    mapped = map_midi(midi, res[0])
+    out_path = "output.mid"
+    mapped.dump(out_path)
+
+    return out_path
+
+# ------------------------------
+# Build Gradio interface
+# ------------------------------
+demo = gr.Interface(
+    fn=render_midi,
+    inputs=[
+        gr.File(label="Upload a Score MIDI File (.mid)", file_types=[".mid"]),
+        gr.Slider(0.1, 2.0, value=1.0, step=0.01, label="Temperature"),
+        gr.Slider(0.1, 1.0, value=0.95, step=0.01, label="Top-p"),
+        gr.Slider(1, 100, value=50, step=1, label="Top-k"),
+    ],
+    outputs=gr.File(label="Rendered Performance MIDI"),
+    title="🎹 Pianist Transformer Rendering",
+    description=(
+        "Upload a symbolic (score) MIDI file and let the Pianist Transformer render it into "
+        "a more expressive performance MIDI. Adjust decoding parameters below to control "
+        "the expressiveness and randomness of the output."
+    ),
+    examples=None,
+)
+
+if __name__ == "__main__":
+    demo.launch()

src/inference/inference.py

@@ -0,0 +1,33 @@
+from src.model.generate import batch_performance_render, map_midi
+from src.model.pianoformer import PianoT5Gemma, PianoT5GemmaConfig
+import torch
+from datasets import load_dataset
+import os
+from miditoolkit import MidiFile
+from src.utils.midi import midi_to_ids, ids_to_midi
+import random
+
+if __name__ == "__main__":    
+    model = PianoT5Gemma.from_pretrained(
+        "models/sft/",
+        torch_dtype=torch.bfloat16
+    )#.cuda()
+
+    midis = []
+    for i in range(1):
+        midis.append(MidiFile(f"data/midis/testset/score/{i}.mid"))
+
+    res = batch_performance_render(
+        model, 
+        midis, 
+        temperature=1.0,
+        top_p=0.95,
+        device="cpu"
+    )
+
+    if not os.path.exists("data/midis/testset/inference"):
+        os.makedirs("data/midis/testset/inference")
+
+    for i, mid in enumerate(res):
+        mid = map_midi(midis[i], mid)
+        mid.dump(f"data/midis/testset/inference/{i}.mid")

src/model/generate.py

@@ -0,0 +1,377 @@
+from src.utils.midi import ids_to_midi, midi_to_ids
+from src.model.pianoformer import PianoT5Gemma
+from miditoolkit import MidiFile
+from torch.nn.utils.rnn import pad_sequence
+from transformers import LogitsProcessorList, LogitsProcessor
+from tqdm import tqdm
+import torch
+from src.utils.midi import normalize_midi
+from miditoolkit import MidiFile, Note, TempoChange, Instrument, ControlChange
+import bisect
+
+class BatchSparseForcedTokenProcessor(LogitsProcessor):
+    def __init__(self, input_ids, config, target_len, origin_len, already, weight, progress_callback):
+        self.batch_map = [{j: input_ids[i][j] for j in range(0, len(input_ids[i]), 8)} for i in range(len(input_ids))]
+        self.valid_id_range = config.valid_id_range
+        self.target_len = target_len
+        self.origin_len = origin_len
+        self.already = already
+        self.weight = weight
+        self.progress_callback = progress_callback
+
+    def __call__(self, input_ids: torch.LongTensor, scores: torch.FloatTensor) -> torch.FloatTensor:
+        if self.progress_callback:
+            self.progress_callback(
+                (input_ids.shape[1] - self.origin_len) / (self.target_len - self.origin_len) * self.weight + self.already
+            )
+        step = input_ids.shape[1] - 1
+        batch_size = scores.shape[0]
+        for i in range(batch_size):
+            sample_map = self.batch_map[i]
+            if step in sample_map:
+                forced_token_id = sample_map[step]
+                scores[i] = float('-inf')
+                scores[i, forced_token_id] = 0.0
+            else:
+                step = step % 8
+                scores[i, :self.valid_id_range[step][0]] = float('-inf')
+                scores[i, self.valid_id_range[step][1]:] = float('-inf')
+                #if step % 8 > 3:
+                #    scores = scores / 0.95
+        return scores
+
+@torch.no_grad()
+def batch_performance_render(
+        model, 
+        score_midi_objs, 
+        max_context_length=4096, 
+        overlap_ratio=0.5, 
+        temperature=1.0,
+        top_p=0.95,
+        device="cpu",
+        progress_callback=None
+    ):
+    def slide_window(total_len, window_len):
+        if total_len <= window_len:
+            return [(0, total_len)]
+        window_len = window_len // 8 * 8
+        out = []
+        start = 0
+        while start + window_len <= total_len:
+            out.append((start, start + window_len))
+            start += int(window_len * (1 - overlap_ratio)) // 8 * 8
+        if out[-1][1] != total_len:
+            out.append((start, total_len))
+        return out
+    if max_context_length > 4096:
+        raise ValueError("You should set max_context_length <= 4096!")
+    batch_ids = [torch.tensor(midi_to_ids(model.config, score_midi_obj), dtype=torch.long).to(device) for score_midi_obj in score_midi_objs]
+    len_list = [len(batch_ids[i]) for i in range(len(batch_ids))]
+    
+    input_ids = pad_sequence(batch_ids, batch_first=True, padding_value=model.config.pad_token_id)
+    windows = slide_window(input_ids.shape[1], max_context_length)
+    #print(windows)
+    output_list = []
+    res_tensor = None
+    for i in tqdm(range(len(windows))):
+        start, end = windows[i]
+        logits_processor = LogitsProcessorList([
+            BatchSparseForcedTokenProcessor(
+                input_ids[:,start:end], 
+                model.config, 
+                end,
+                start,
+                i / len(windows),
+                1 / len(windows),
+                progress_callback,
+            )
+        ])
+        if i == 0:
+            output = model.generate(
+                input_ids[:,start:end], 
+                do_sample=True, 
+                max_new_tokens=end-start, 
+                logits_processor=logits_processor,
+                temperature=temperature,
+                top_p=top_p,
+            )
+            res_tensor = output[:,1:]
+        else:
+            last_start, last_end = windows[i-1]
+            length = int(((last_end-last_start) - (start-last_start)) * 0.2)
+            decoder_input_ids = output_list[i-1][:, start-last_start:last_end-last_start - length]
+            start_tensor = torch.tensor([[model.config.bos_token_id] for _ in range(input_ids.shape[0])], dtype=torch.long).to(device)
+            decoder_input_ids = torch.cat([start_tensor, decoder_input_ids], dim=1)
+            #print(decoder_input_ids.shape)
+            output = model.generate(
+                input_ids[:,start:end], 
+                decoder_input_ids=decoder_input_ids,
+                do_sample=True, 
+                max_new_tokens=end-last_end+length, 
+                logits_processor=logits_processor,
+                temperature=temperature,
+                top_p=top_p,
+            )
+            res_tensor = torch.cat([res_tensor[:,:-length], output[:,-(end-last_end+length):]], dim=1)
+        output_list.append(output)
+    res_tensor = res_tensor.cpu().numpy().tolist()
+    #print(res_tensor)
+    res = []
+    for i in range(len(res_tensor)):
+        #print(res_tensor[i][:len_list[i]])
+        res.append(ids_to_midi(model.config, res_tensor[i][:len_list[i]]))
+    return res
+
+
+def map_midi(score_midi_obj, performance_midi_obj):
+    def compute_duration(start_time, target_duration, tempo_list):
+        if target_duration <= 0:
+            return 0
+        if not tempo_list:
+            # 如果没有提供tempo信息，则假定为默认的120 BPM
+            tempo_list = [TempoChange(120, 0)]
+
+        # --- 步骤1: 定位start_time所在的BPM区间 ---
+        # 提取所有tempo变化的时间点
+        tempo_times = [t.time for t in tempo_list]
+        # 使用二分查找找到start_time应该插入的位置
+        # bisect_right返回的是插入点索引，因此当前生效的tempo在索引-1的位置
+        start_tempo_idx = bisect.bisect_right(tempo_times, start_time) - 1
+        # 如果start_time在第一个tempo变化之前，索引会是-1，修正为0
+        if start_tempo_idx < 0:
+            start_tempo_idx = 0
+
+        # --- 步骤2: 初始化循环变量 ---
+        total_ticks_duration = 0.0
+        time_remaining_ms = float(target_duration)
+        current_tick = start_time
+        current_tempo_idx = start_tempo_idx
+
+        # --- 步骤3: 循环处理每个BPM区间，直到消耗完target_duration ---
+        # 使用一个极小值(epsilon)来处理浮点数精度问题
+        while time_remaining_ms > 1e-9:
+            current_tempo_event = tempo_list[current_tempo_idx]
+            current_bpm = current_tempo_event.tempo
+
+            # 计算在当前BPM下，每个tick持续多少毫秒
+            # 1分钟 = 60,000毫秒
+            # 每分钟节拍数 = bpm
+            # 每拍tick数 = TICK_PER_BEAT
+            # ms_per_tick = (毫秒/分钟) / (节拍/分钟) / (tick/节拍) = (60000 / bpm) / TICK_PER_BEAT
+            ms_per_tick = (60 * 1000.0 / current_bpm) / 500
+
+            # 确定当前BPM区间的结束点
+            # 如果是最后一个tempo，则它会一直持续下去
+            end_of_segment_tick = float('inf')
+            if current_tempo_idx + 1 < len(tempo_list):
+                end_of_segment_tick = tempo_list[current_tempo_idx + 1].time
+
+            # 计算从当前位置到本BPM区间结束，有多少tick
+            ticks_in_segment = end_of_segment_tick - current_tick
+            # 这些tick总共持续多少毫秒
+            ms_in_segment = ticks_in_segment * ms_per_tick
+
+            # --- 步骤4: 判断与更新 ---
+            if time_remaining_ms <= ms_in_segment:
+                # 如果剩余需要的时间，在本BPM区间内就能满足
+                # 计算还需要多少tick来凑够剩余的毫秒数
+                ticks_needed = time_remaining_ms / ms_per_tick
+                total_ticks_duration += ticks_needed
+                # 时间已全部消耗完毕，跳出循环
+                time_remaining_ms = 0
+            else:
+                # 如果本BPM区间的时间不够用
+                # 消耗掉整个区间的tick和毫秒数
+                total_ticks_duration += ticks_in_segment
+                time_remaining_ms -= ms_in_segment
+                
+                # 更新“指针”，移动到下一个BPM区间的起点
+                current_tick = end_of_segment_tick
+                current_tempo_idx += 1
+
+        # 返回四舍五入后的总tick数
+        return round(total_ticks_duration)
+
+    def ms_to_tick(target_ms, tempo_list):
+        # --- 边缘情况处理 ---
+        if target_ms <= 0:
+            return 0
+        if not tempo_list:
+            # 如果没有提供tempo信息，则假定为默认的120 BPM
+            tempo_list = [TempoChange(120, 0)]
+
+        # --- 步骤1: 初始化累加器 ---
+        accumulated_ms = 0.0
+
+        # --- 步骤2: 遍历所有“有终点”的BPM区间 ---
+        # 我们遍历到倒数第二个元素，因为每个循环处理的是 tempo[i] 到 tempo[i+1] 的区间
+        for i in range(len(tempo_list) - 1):
+            current_tempo_event = tempo_list[i]
+            next_tempo_event = tempo_list[i+1]
+
+            current_bpm = current_tempo_event.tempo
+            
+            # 计算当前区间的tick数和对应的毫秒数
+            ticks_in_segment = next_tempo_event.time - current_tempo_event.time
+            
+            # 如果区间长度为0，直接跳过，避免除零错误
+            if ticks_in_segment == 0:
+                continue
+                
+            ms_per_tick = (60 * 1000.0 / current_bpm) / 500
+            ms_in_segment = ticks_in_segment * ms_per_tick
+
+            # --- 步骤3: 判断目标是否在本区间内 ---
+            if target_ms <= accumulated_ms + ms_in_segment:
+                # 目标在本区间内！
+                ms_into_segment = target_ms - accumulated_ms
+                ticks_needed = ms_into_segment / ms_per_tick
+                
+                # 最终tick = 本区间起始tick + 在本区间内转换出的tick
+                final_tick = current_tempo_event.time + ticks_needed
+                return round(final_tick)
+            
+            # 如果目标不在本区间，则累加本区间的总毫秒数，继续下一个循环
+            accumulated_ms += ms_in_segment
+
+        # --- 步骤4: 如果循环结束仍未返回，说明目标在最后一个BPM区间内 ---
+        last_tempo_event = tempo_list[-1]
+        last_bpm = last_tempo_event.tempo
+        
+        ms_per_tick = (60 * 1000.0 / last_bpm) / 500
+
+        # 计算进入最后一个区间后，还需要多少毫秒
+        ms_into_segment = target_ms - accumulated_ms
+        ticks_needed = ms_into_segment / ms_per_tick
+
+        # 最终tick = 最后一个区间的起始tick + 剩余毫秒转换的tick
+        final_tick = last_tempo_event.time + ticks_needed
+        return round(final_tick)
+
+    norm_score = normalize_midi(score_midi_obj)
+    norm_performance = normalize_midi(performance_midi_obj)
+
+    score_notes = norm_score.instruments[0].notes
+    performance_notes = norm_performance.instruments[0].notes
+    performance_ccs = norm_performance.instruments[0].control_changes
+
+    start_list = []
+    last = -1
+    score_start = score_notes[0].start
+    performance_start = performance_notes[0].start
+    for i in range(len(score_notes)):
+        performance_notes[i].end -= performance_start
+        performance_notes[i].start -= performance_start
+        score_notes[i].end -= score_start
+        score_notes[i].start -= score_start
+        if score_notes[i].start != last:
+            start_list.append((score_notes[i].start, performance_notes[i].start, i))
+            last = score_notes[i].start
+    
+    for i in range(len(performance_ccs)):
+        performance_ccs[i].time -= performance_start
+
+    score_interval_list = []
+    performance_interval_list = []
+
+    for i in range(len(start_list)-1):
+        score_interval_list.append(start_list[i+1][0] - start_list[i][0])
+        performance_interval_list.append(start_list[i+1][1] - start_list[i][1])
+    #print(score_interval_list)
+    #print(performance_interval_list)
+
+    tempo_list = []
+    start_note_offset = []
+    for i in range(len(score_interval_list)):
+        if performance_interval_list[i] != 0:
+            bpm = 120.0 / performance_interval_list[i] * score_interval_list[i]
+        else:
+            bpm = 300
+
+        if bpm > 300:
+            start_note_offset.append(300 / 120.0 * performance_interval_list[i] - score_interval_list[i])
+        elif bpm < 10:
+            start_note_offset.append(10 / 120.0 * performance_interval_list[i] - score_interval_list[i])
+        else:
+            start_note_offset.append(0)
+        tempo_list.append(max(min(bpm, 300), 10))
+        #tempo_list.append(120.0 / performance_interval_list[i] * score_interval_list[i])
+    #print(tempo_list)
+
+    for i in range(1, len(start_note_offset)):
+        start_note_offset[i] += start_note_offset[i-1]
+    #print(start_note_offset)
+
+    #print(len(tempo_list))
+    #print(len(start_list))
+    note_tempo_list = []
+    note_performance_align = []
+    note_start_offset = [0]
+    cnt = 0
+    for i in range(len(score_notes)):
+        if cnt < len(start_list) - 2 and i >= start_list[cnt + 1][2]:
+            cnt += 1
+        note_tempo_list.append(tempo_list[cnt])
+        note_performance_align.append(start_list[cnt][1])
+        note_start_offset.append(start_note_offset[cnt])
+    #print(note_start_offset)
+
+    for i in range(len(score_notes)):
+        score_notes[i].start += note_start_offset[i]
+    note_interval_list = [0]
+    for i in range(len(score_notes)-1):
+        note_interval_list.append(score_notes[i+1].start - score_notes[i].start)
+
+    #print(note_tempo_list)
+    #print(note_performance_align)
+
+    #for i in range(len(performance_notes)):
+        #print(performance_notes[i].start)
+
+    micro_shift_list = [0]
+    cnt = 1
+    last_time = 0
+    for i in range(1, len(score_notes)):
+        last_time += note_interval_list[i] / note_tempo_list[i-1] * 120
+        micro_shift_list.append((performance_notes[i].start - last_time) / 120 * note_tempo_list[i-1])
+        #last_time = note_performance_align[i]
+        #print(last_time)
+
+    #print(micro_shift_list)
+    #plt.plot(tempo_list)
+
+    res = MidiFile(ticks_per_beat=500)
+    res_notes = []
+    start_time_list = []
+    tempo_list_filter = []
+    cc_list = []
+    last = -1
+    for i in range(len(score_notes)):
+        start_time_list.append(round(score_notes[i].start + micro_shift_list[i]))
+        #res_notes.append(Note(performance_notes[i].velocity, score_notes[i].pitch, round(score_notes[i].start + micro_shift_list[i]), round(score_notes[i].start + micro_shift_list[i]) + 100))
+        #res.tempo_changes.append(TempoChange(round(note_tempo_list[i]), round(score_notes[i].start + micro_shift_list[i])))
+        #print(last , round(note_tempo_list[i]))
+        if last != round(note_tempo_list[i]):
+            tempo_list_filter.append(TempoChange(round(note_tempo_list[i]), round(score_notes[i].start + micro_shift_list[i])))
+            last = round(note_tempo_list[i])
+    for i in range(len(score_notes)):
+        res_notes.append(
+            Note(
+                performance_notes[i].velocity, 
+                score_notes[i].pitch, 
+                start_time_list[i], 
+                start_time_list[i]+compute_duration(start_time_list[i], performance_notes[i].duration, tempo_list_filter)
+            )
+        )
+
+    for cc in performance_ccs:
+        cc_list.append(ControlChange(64, cc.value, ms_to_tick(cc.time, tempo_list_filter)))
+        
+    #print(tempo_list_filter)
+    res.tempo_changes = tempo_list_filter
+    res.instruments.append(Instrument(program=0, is_drum=False, name="Piano", notes=res_notes, control_changes=cc_list))
+    return res
+
+
+if __name__ == "__main__":
+    pass

src/model/pianoformer.py

@@ -0,0 +1,459 @@
+from typing import Callable, Optional, Union
+from transformers import T5GemmaModel, T5GemmaConfig, T5GemmaModuleConfig, T5GemmaPreTrainedModel, T5GemmaForConditionalGeneration, AutoTokenizer
+import torch
+from transformers.models.t5gemma.modeling_t5gemma import (
+    T5GemmaLMHead, 
+    GenerationMixin, 
+    logger,
+    T5GemmaSelfAttention,
+    T5GemmaEncoderLayer,
+    T5GemmaRMSNorm,
+    T5GemmaRotaryEmbedding,
+    make_default_2d_attention_mask,
+    create_causal_mask,
+    bidirectional_mask_function,
+    create_sliding_window_causal_mask,
+    sliding_window_bidirectional_mask_function,
+    T5GemmaDecoder
+)
+
+from transformers.modeling_outputs import (
+    BaseModelOutput,
+    BaseModelOutputWithPastAndCrossAttentions,
+    Seq2SeqLMOutput,
+    Seq2SeqModelOutput,
+    SequenceClassifierOutput,
+    TokenClassifierOutput,
+)
+from transformers.cache_utils import Cache, DynamicCache, EncoderDecoderCache
+from transformers.processing_utils import Unpack
+from transformers.utils import TransformersKwargs, auto_docstring, can_return_tuple, is_torchdynamo_compiling, logging
+import torch.nn as nn
+
+class PianoT5GemmaConfig(T5GemmaConfig):
+
+    def __init__(
+            self, 
+            hidden_size=768, 
+            intermediate_size=3072,
+            num_attention_heads=8,
+            num_key_value_heads=4,
+            head_dim=128,
+            encoder_layers_num=8,
+            decoder_layers_num=4,
+            **kwargs
+        ):
+        total_vocab_size = 5389
+
+        self.mask_token_id = 1
+        self.bos_token_id = 2
+        self.play_token_id = 4
+        self.pitch_start = 5
+        self.velocity_start = 5 + 128
+        self.timing_start = 5 + 128 + 128
+        self.pedal_start = 5 + 128 + 128 + 5000
+        self.hidden_size = hidden_size
+
+        self.valid_id_range = [
+            (5, 133),
+            (261, 5252),   
+            (133, 261),
+            (261, 5261),
+            (5261, 5389),
+            (5261, 5389),
+            (5261, 5389),
+            (5261, 5389),
+        ]
+
+        encoder_config = T5GemmaModuleConfig(
+            vocab_size=total_vocab_size,
+            hidden_size=hidden_size,
+            intermediate_size=intermediate_size,
+            num_hidden_layers=encoder_layers_num,
+            num_attention_heads=num_attention_heads,
+            num_key_value_heads=num_key_value_heads,
+            head_dim=head_dim,
+            pad_token_id=0,
+            bos_token_id=2,
+            eos_token_id=3,
+        )
+        decoder_config = T5GemmaModuleConfig(
+            vocab_size=total_vocab_size,
+            hidden_size=hidden_size,
+            intermediate_size=intermediate_size,
+            num_hidden_layers=decoder_layers_num,
+            num_attention_heads=num_attention_heads,
+            num_key_value_heads=num_key_value_heads,
+            head_dim=head_dim,
+            pad_token_id=0,
+            bos_token_id=2,
+            eos_token_id=3,
+        )
+
+        super().__init__(
+            encoder=encoder_config,
+            decoder=decoder_config,
+            vocab_size=total_vocab_size,
+            **kwargs,
+        )
+
+class PianoEncoderEmbeddings(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        self.word_embeddings = nn.Embedding(config.vocab_size, config.hidden_size, padding_idx=config.pad_token_id)
+        if config.hidden_size % 8 != 0:
+            raise ValueError("Invalid hidden size: must be a multiple of 8.")
+        self.projection_layers = nn.ModuleList([nn.Linear(config.hidden_size, config.hidden_size // 8) for i in range(8)])
+        self.hidden_size = config.hidden_size
+
+    def forward(
+        self,
+        input_ids: Optional[torch.LongTensor] = None,
+    ) -> torch.Tensor:
+        input_shape = input_ids.size()
+
+        batch_size = input_shape[0]
+        seq_length = input_shape[1]
+
+        inputs_embeds = self.word_embeddings(input_ids)
+        grouped_embeds = inputs_embeds.view(batch_size, seq_length // 8, 8, -1)
+        projection_list = []
+        for i in range(8):
+            projection_list.append(self.projection_layers[i](grouped_embeds[:,:,i,:]))
+        projection_cat = torch.cat(projection_list, dim=-1)
+        inputs_embeds = projection_cat.view(batch_size, -1, self.hidden_size)
+        embeddings = inputs_embeds
+        return embeddings
+
+
+class PianoT5GemmaEncoder(T5GemmaPreTrainedModel):
+    _can_record_outputs = {
+        "attentions": T5GemmaSelfAttention,
+        "hidden_states": T5GemmaEncoderLayer,
+    }
+
+    def __init__(self, config):
+        super().__init__(config)
+        self.padding_idx = config.pad_token_id
+        self.vocab_size = config.vocab_size
+
+        self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size, self.padding_idx)
+        self.embeddings = PianoEncoderEmbeddings(config)
+        self.norm = T5GemmaRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+        self.rotary_emb = T5GemmaRotaryEmbedding(config=config)
+        self.gradient_checkpointing = False
+
+        self.layers = nn.ModuleList(
+            [T5GemmaEncoderLayer(config, layer_idx) for layer_idx in range(config.num_hidden_layers)]
+        )
+        self.dropout = nn.Dropout(config.dropout_rate)
+
+        # Initialize weights and apply final processing
+        self.post_init()
+
+    def forward(
+        self,
+        input_ids: Optional[torch.LongTensor] = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        inputs_embeds: Optional[torch.FloatTensor] = None,
+        **kwargs: Unpack[TransformersKwargs],
+    ) -> BaseModelOutput:
+        if (input_ids is None) ^ (inputs_embeds is not None):
+            raise ValueError("You must specify exactly one of input_ids or inputs_embeds")
+
+        if inputs_embeds is None:
+            inputs_embeds = self.embeddings(input_ids)
+            input_ids = None
+
+        cache_position = torch.arange(0, inputs_embeds.shape[1], device=inputs_embeds.device)
+
+        if position_ids is None:
+            position_ids = cache_position.unsqueeze(0)
+
+        if attention_mask is not None:
+            B, L = attention_mask.shape
+            block_mask = attention_mask.view(B, L // 8, 8)
+            mask2 = block_mask.any(dim=-1).long()
+            attention_mask = mask2.view(B, -1)
+
+        if attention_mask is None:
+            attention_mask = make_default_2d_attention_mask(input_ids, inputs_embeds, self.config.pad_token_id)
+        
+        if not isinstance(self_attn_mask_mapping := attention_mask, dict):
+            mask_kwargs = {
+                "config": self.config,
+                "input_embeds": inputs_embeds,
+                "attention_mask": attention_mask,
+                "cache_position": cache_position,
+                "past_key_values": None,
+                "position_ids": position_ids,
+            }
+            self_attn_mask_mapping = {
+                "full_attention": create_causal_mask(
+                    **mask_kwargs,
+                    or_mask_function=bidirectional_mask_function(attention_mask),
+                ),
+                "sliding_attention": create_sliding_window_causal_mask(
+                    **mask_kwargs,
+                    or_mask_function=sliding_window_bidirectional_mask_function(self.config.sliding_window),
+                    and_mask_function=bidirectional_mask_function(attention_mask),
+                ),
+            }
+
+        hidden_states = inputs_embeds
+        position_embeddings = self.rotary_emb(hidden_states, position_ids)
+
+        normalizer = torch.tensor(self.config.hidden_size**0.5, dtype=hidden_states.dtype)
+        hidden_states = hidden_states * normalizer
+        hidden_states = self.dropout(hidden_states)
+
+        for layer_module in self.layers[: self.config.num_hidden_layers]:
+            hidden_states = layer_module(
+                hidden_states,
+                position_embeddings,
+                self_attn_mask_mapping[layer_module.attention_type],
+                position_ids,
+                **kwargs,
+            )
+        hidden_states = self.norm(hidden_states)
+        hidden_states = self.dropout(hidden_states)
+        return BaseModelOutput(
+            last_hidden_state=hidden_states,
+        )
+
+class PianoT5GemmaModel(T5GemmaPreTrainedModel):
+    def __init__(self, config: T5GemmaConfig):
+        super().__init__(config)
+
+        if not config.is_encoder_decoder:
+            raise ValueError("T5GemmaModel only support encoder-decoder modeling. Use `T5GemmaEncoderModel` instead.")
+
+        self.encoder = PianoT5GemmaEncoder(config.encoder)
+        self.decoder = T5GemmaDecoder(config.decoder)
+
+        self.post_init()
+
+    def get_encoder(self):
+        return self.encoder
+
+    def get_decoder(self):
+        return self.decoder
+
+    def get_input_embeddings(self):
+        return self.encoder.get_input_embeddings()
+
+    def set_input_embeddings(self, new_embeddings):
+        return self.encoder.set_input_embeddings(new_embeddings)
+
+    def forward(
+        self,
+        input_ids: Optional[torch.LongTensor] = None,
+        attention_mask: Optional[torch.FloatTensor] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        decoder_input_ids: Optional[torch.LongTensor] = None,
+        decoder_attention_mask: Optional[torch.BoolTensor] = None,
+        decoder_position_ids: Optional[torch.LongTensor] = None,
+        encoder_outputs: Optional[BaseModelOutput] = None,
+        past_key_values: Optional[EncoderDecoderCache] = None,
+        inputs_embeds: Optional[torch.Tensor] = None,
+        decoder_inputs_embeds: Optional[torch.Tensor] = None,
+        use_cache: Optional[bool] = None,
+        cache_position: Optional[torch.LongTensor] = None,
+        **kwargs: Unpack[TransformersKwargs],
+    ) -> Seq2SeqModelOutput:
+        r"""
+        decoder_position_ids (`torch.LongTensor` of shape `(batch_size, decoder_sequence_length)`, *optional*):
+            Indices of positions of each decoder input sequence tokens in the position embeddings. Selected in the range `[0,
+            config.decoder.n_positions - 1]`. [What are position IDs?](../glossary#position-ids)
+        """
+        if encoder_outputs is None:
+            encoder_outputs = self.encoder(
+                input_ids=input_ids,
+                attention_mask=attention_mask,
+                position_ids=position_ids,
+                inputs_embeds=inputs_embeds,
+                **kwargs,
+            )
+
+        encoder_hidden_states = encoder_outputs.last_hidden_state
+
+        decoder_outputs = self.decoder(
+            input_ids=decoder_input_ids,
+            attention_mask=decoder_attention_mask,
+            position_ids=decoder_position_ids,
+            inputs_embeds=decoder_inputs_embeds,
+            past_key_values=past_key_values,
+            encoder_hidden_states=encoder_hidden_states,
+            encoder_attention_mask=attention_mask,
+            use_cache=use_cache,
+            cache_position=cache_position,
+            **kwargs,
+        )
+
+        return Seq2SeqModelOutput(
+            last_hidden_state=decoder_outputs.last_hidden_state,
+            past_key_values=decoder_outputs.past_key_values,
+            decoder_hidden_states=decoder_outputs.hidden_states
+            if kwargs.get("output_hidden_states", False)
+            else (decoder_outputs.last_hidden_state,),
+            decoder_attentions=decoder_outputs.attentions,
+            cross_attentions=decoder_outputs.cross_attentions,
+            encoder_last_hidden_state=encoder_outputs.last_hidden_state,
+            encoder_hidden_states=encoder_outputs.hidden_states,
+            encoder_attentions=encoder_outputs.attentions,
+        )
+
+
+class PianoT5Gemma(T5GemmaPreTrainedModel, GenerationMixin):
+    _tied_weights_keys = ["model.decoder.embed_tokens.weight", "lm_head.out_proj.weight"]
+    _tp_plan = {"lm_head.out_proj": "colwise_rep"}
+    _pp_plan = {"lm_head.out_proj": (["hidden_states"], ["logits"])}
+
+    def __init__(self, config: PianoT5GemmaConfig):
+        config.is_encoder_decoder = True
+        super().__init__(config)
+        self.embeddings = PianoEncoderEmbeddings(config)
+        self.model = PianoT5GemmaModel(config)
+        self.vocab_size = config.decoder.vocab_size
+        self.lm_head = T5GemmaLMHead(config.decoder.hidden_size, self.vocab_size)
+        self.loss_type = "ForMaskedLM"
+
+        self.post_init()
+
+    def set_output_embeddings(self, new_embeddings):
+        self.lm_head.out_proj = new_embeddings
+
+    def get_output_embeddings(self):
+        return self.lm_head.out_proj
+
+    def _tie_weights(self):
+        # Decoder input and output embeddings are tied.
+        if self.config.tie_word_embeddings:
+            self._tie_or_clone_weights(self.lm_head.out_proj, self.get_decoder().get_input_embeddings())
+
+    def get_encoder(self):
+        return self.model.encoder
+
+    def get_decoder(self):
+        return self.model.decoder
+
+    def forward(
+        self,
+        input_ids: Optional[torch.LongTensor] = None,
+        attention_mask: Optional[torch.FloatTensor] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        decoder_input_ids: Optional[torch.LongTensor] = None,
+        decoder_attention_mask: Optional[torch.BoolTensor] = None,
+        decoder_position_ids: Optional[torch.LongTensor] = None,
+        encoder_outputs: Optional[BaseModelOutput] = None,
+        past_key_values: Optional[EncoderDecoderCache] = None,
+        inputs_embeds: Optional[torch.FloatTensor] = None,
+        decoder_inputs_embeds: Optional[torch.FloatTensor] = None,
+        labels: Optional[torch.LongTensor] = None,
+        use_cache: Optional[bool] = None,
+        cache_position: Optional[torch.LongTensor] = None,
+        logits_to_keep: Union[int, torch.Tensor] = 0,
+        **kwargs: Unpack[TransformersKwargs],
+    ) -> Union[tuple[torch.FloatTensor], Seq2SeqLMOutput]:
+        r"""
+        decoder_position_ids (`torch.LongTensor` of shape `(batch_size, decoder_sequence_length)`, *optional*):
+            Indices of positions of each decoder input sequence tokens in the position embeddings. Selected in the range `[0,
+            config.decoder.n_positions - 1]`. [What are position IDs?](../glossary#position-ids)
+        labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
+            Labels for computing the masked language modeling loss. Indices should either be in `[0, ...,
+            config.vocab_size]` or -100 (see `input_ids` docstring). Tokens with indices set to `-100` are ignored
+            (masked), the loss is only computed for the tokens with labels in `[0, ..., config.vocab_size]`.
+        """
+        if self.training and self.config._attn_implementation != "eager":
+            msg = (
+                "It is strongly recommended to train T5Gemma models with the `eager` attention implementation "
+                f"instead of `{self.config._attn_implementation}`. Use `eager` with `AutoModelForCausalLM.from_pretrained('<path-to-checkpoint>', attn_implementation='eager')`."
+            )
+            if is_torchdynamo_compiling():
+                raise ValueError(msg)
+            else:
+                logger.warning_once(msg)
+
+        if labels is not None and decoder_input_ids is None and decoder_inputs_embeds is None:
+            # get decoder inputs from shifting lm labels to the right
+            decoder_input_ids = self._shift_right(labels)
+
+        #if input_ids is not None:
+        #    inputs_embeds = self.embeddings(input_ids)
+        
+        #if attention_mask is not None:
+        #    B, L = attention_mask.shape
+        #    block_mask = attention_mask.view(B, L // 8, 8)
+        #    mask2 = block_mask.any(dim=-1).long()
+        #    attention_mask = mask2.view(B, -1)
+
+        #print(attention_mask)
+
+        decoder_outputs: Seq2SeqModelOutput = self.model(
+            input_ids=input_ids,
+            attention_mask=attention_mask,
+            position_ids=position_ids,
+            decoder_input_ids=decoder_input_ids,
+            decoder_attention_mask=decoder_attention_mask,
+            decoder_position_ids=decoder_position_ids,
+            encoder_outputs=encoder_outputs,
+            past_key_values=past_key_values,
+            inputs_embeds=inputs_embeds,
+            decoder_inputs_embeds=decoder_inputs_embeds,
+            use_cache=use_cache,
+            cache_position=cache_position,
+            **kwargs,
+        )
+
+        hidden_states = decoder_outputs.last_hidden_state
+        # Only compute necessary logits, and do not upcast them to float if we are not computing the loss
+        slice_indices = slice(-logits_to_keep, None) if isinstance(logits_to_keep, int) else logits_to_keep
+        logits = self.lm_head(hidden_states[:, slice_indices, :])
+        decoder_config = self.get_decoder().config
+        if decoder_config.final_logit_softcapping is not None:
+            logits = logits / decoder_config.final_logit_softcapping
+            logits = torch.tanh(logits)
+            logits = logits * decoder_config.final_logit_softcapping
+
+        loss = None
+        if labels is not None:
+            # Input has right-shifted so we directly perform masked lm loss
+            loss = self.loss_function(logits, labels, self.vocab_size, **kwargs)
+
+        return Seq2SeqLMOutput(
+            loss=loss,
+            logits=logits,
+            past_key_values=decoder_outputs.past_key_values,
+            decoder_hidden_states=decoder_outputs.decoder_hidden_states,
+            decoder_attentions=decoder_outputs.decoder_attentions,
+            cross_attentions=decoder_outputs.cross_attentions,
+            encoder_last_hidden_state=decoder_outputs.encoder_last_hidden_state,
+            encoder_hidden_states=decoder_outputs.encoder_hidden_states,
+            encoder_attentions=decoder_outputs.encoder_attentions,
+        )
+
+    def prepare_decoder_input_ids_from_labels(self, labels: torch.Tensor):
+        return self._shift_right(labels)
+
+
+
+if __name__ == "__main__":
+    config = PianoT5GemmaConfig()
+    test = PianoEncoderEmbeddings(config)
+    model = PianoT5Gemma(config).cuda()
+    #encoder_config = T5GemmaModuleConfig(num_hidden_layers=1)
+    #decoder_config = T5GemmaModuleConfig(num_hidden_layers=1)
+    #config = T5GemmaConfig(encoder_config, decoder_config, attn_implementation='eager')
+
+    #model = T5GemmaForConditionalGeneration(config).cuda()
+
+    toy_ids = torch.tensor([[1,2,3,4,1,2,3,4,1,2,3,4,1,2,3,4]], dtype=torch.long).cuda()
+    #tokenizer = AutoTokenizer.from_pretrained("google/t5gemma-2b-2b-ul2")
+    #input_text = "Write me a poem about Machine Learning. Answer:"
+    #input_ids = tokenizer(input_text, return_tensors="pt").to("cuda")
+    print(model.generate(toy_ids, decoder_input_ids=toy_ids, max_new_tokens=32))
+
+   
+    #print(model(input_ids=toy_ids, decoder_input_ids=toy_ids).logits.shape)
+

src/utils/func.py

@@ -0,0 +1,5 @@
+import inspect
+
+def filter_valid_args(arg_dict, class_type):
+    valid_keys = inspect.signature(class_type).parameters.keys()
+    return {k: v for k, v in arg_dict.items() if k in valid_keys}

src/utils/midi.py

@@ -0,0 +1,602 @@
+from miditoolkit import MidiFile, Note, Instrument, TempoChange, ControlChange
+import bisect
+import numpy as np
+import os
+from copy import copy
+import random
+from collections import defaultdict
+
+"""
+def normalize_midi(midi_obj, target_ticks_per_beat = 500, target_tempo = 120):
+    ticks_per_beat = midi_obj.ticks_per_beat
+    merged_events = []
+    for i in range(len(midi_obj.instruments)):
+        filter_control_changes = []
+        for cc in midi_obj.instruments[i].control_changes:
+            if cc.number == 64:
+                filter_control_changes.append(cc)
+        merged_events.extend(midi_obj.instruments[i].notes + filter_control_changes)
+    merged_events.sort(key=lambda x: (x.start, x.pitch) if isinstance(x, Note) else (x.time, x.number))
+    
+    time_interval = []
+    last_time = 0
+    for note in merged_events:
+        if isinstance(note, Note):
+            time_interval.append(note.start - last_time)
+            last_time = note.start
+        else:
+            time_interval.append(note.time - last_time)
+            last_time = note.time
+
+    output_notes = []
+    output_cc = []
+    ind = -1
+    now_tempo = 120
+    now_time = 0
+    for i, note in enumerate(merged_events):
+        if isinstance(note, Note):
+            time = note.start
+        else:
+            time = note.time
+        while ind + 1 < len(midi_obj.tempo_changes) and time >= midi_obj.tempo_changes[ind+1].time:
+            now_tempo = midi_obj.tempo_changes[ind+1].tempo
+            ind += 1
+        ratio = target_ticks_per_beat * target_tempo / now_tempo / ticks_per_beat
+        start_time = time_interval[i] * ratio + now_time
+        if isinstance(note, Note):
+            end_time = (note.end - note.start) * ratio + start_time
+            output_notes.append(Note(note.velocity, note.pitch, round(start_time), round(end_time)))
+        else:
+            output_cc.append(ControlChange(64, note.value, round(start_time)))
+        now_time = round(start_time)
+    
+    output_midi_obj = MidiFile(ticks_per_beat=target_ticks_per_beat)
+    output_midi_obj.instruments.append(Instrument(program=0, is_drum=False, name="Piano", notes=output_notes, control_changes=output_cc))
+    output_midi_obj.tempo_changes.append(TempoChange(target_tempo, 0))
+    for note in output_notes:
+        output_midi_obj.max_tick = max(output_midi_obj.max_tick, note.end)
+    for cc in output_cc:
+        output_midi_obj.max_tick = max(output_midi_obj.max_tick, cc.time)
+    return output_midi_obj
+"""
+
+"""
+def normalize_midi(midi_obj, target_ticks_per_beat=500, target_tempo=120):    
+    # 创建一个新的、干净的MidiFile对象用于输出
+    output_midi_obj = MidiFile(ticks_per_beat=target_ticks_per_beat)
+    output_midi_obj.tempo_changes.append(TempoChange(target_tempo, 0))
+    
+    # 获取原始MIDI的tick到秒的精确映射
+    # 这是最关键的一步，partitura和miditoolkit都有类似功能
+    # miditoolkit的get_tick_to_time_mapping()可以处理所有tempo变化
+    tick_to_time_map = midi_obj.get_tick_to_time_mapping()
+    
+    # 计算从秒转换回目标tick的比例因子
+    # 目标MIDI中，每秒对应的tick数 = target_ticks_per_beat * (target_tempo / 60)
+    seconds_to_target_ticks_factor = target_ticks_per_beat * (target_tempo / 60.0)
+
+    merged_notes = []
+    merged_cc = []
+
+    # 遍历所有乐器轨道
+    for instrument in midi_obj.instruments:
+        # 只处理非鼓组的乐器
+        if not instrument.is_drum:
+            # --- 处理音符 (Notes) ---
+            for note in instrument.notes:
+                # 1. 将原始tick转换为绝对秒数
+                start_time_sec = tick_to_time_map[note.start]
+                end_time_sec = tick_to_time_map[note.end]
+                
+                # 2. 将绝对秒数转换为目标tick
+                new_start_tick = round(start_time_sec * seconds_to_target_ticks_factor)
+                new_end_tick = round(end_time_sec * seconds_to_target_ticks_factor)
+                
+                # 避免duration为0的音符
+                if new_start_tick == new_end_tick:
+                    new_end_tick += 1
+
+                merged_notes.append(Note(velocity=note.velocity, 
+                                         pitch=note.pitch, 
+                                         start=new_start_tick, 
+                                         end=new_end_tick))
+            
+            # --- 处理延音踏板 (CC #64) ---
+            for cc in instrument.control_changes:
+                if cc.number == 64:
+                    # 1. 将原始tick转换为绝对秒数
+                    time_sec = tick_to_time_map[cc.time]
+                    
+                    # 2. 将绝对秒数转换为目标tick
+                    new_time_tick = round(time_sec * seconds_to_target_ticks_factor)
+                    
+                    merged_cc.append(ControlChange(number=64, 
+                                                   value=cc.value, 
+                                                   time=new_time_tick))
+
+    # --- 排序并创建新乐器 ---
+    # 按开始时间排序，对于同时开始的事件，CC优先于Note
+    merged_notes.sort(key=lambda x: (x.start, x.pitch))
+    merged_cc.sort(key=lambda x: (x.time, x.number))
+    
+    output_instrument = Instrument(program=0, is_drum=False, name="Piano")
+    output_instrument.notes = merged_notes
+    output_instrument.control_changes = merged_cc
+    output_midi_obj.instruments.append(output_instrument)
+    
+    # --- 正确计算 max_tick ---
+    max_tick = 0
+    if output_instrument.notes:
+        max_tick = max(max_tick, max(n.end for n in output_instrument.notes))
+    if output_instrument.control_changes:
+        max_tick = max(max_tick, max(c.time for c in output_instrument.control_changes))
+    
+    output_midi_obj.max_tick = max_tick
+
+    return output_midi_obj
+"""
+
+def normalize_midi(midi_obj, target_ticks_per_beat=500, target_tempo=120):
+    """
+    将一个MidiFile对象标准化：
+    1. 合并所有轨道的钢琴音符和延音踏板事件。
+    2. 将所有时间信息（包括tempo变化）统一转换为一个固定的ticks_per_beat和tempo。
+    3. 清理重叠音符以避免解析错误。
+    4. 正确计算并设置max_tick。
+
+    Args:
+        midi_obj (MidiFile): 原始的MidiFile对象。
+        target_ticks_per_beat (int): 目标ticks_per_beat.
+        target_tempo (float): 目标tempo (BPM).
+
+    Returns:
+        MidiFile: 标准化后的新MidiFile对象。
+    """
+    
+    # 创建一个新的、干净的MidiFile对象用于输出
+    output_midi_obj = MidiFile(ticks_per_beat=target_ticks_per_beat)
+    output_midi_obj.tempo_changes.append(TempoChange(target_tempo, 0))
+    
+    tick_to_time_map = midi_obj.get_tick_to_time_mapping()
+    seconds_to_target_ticks_factor = target_ticks_per_beat * (target_tempo / 60.0)
+
+    # --- 1. 收集并转换所有音符 ---
+    all_converted_notes = []
+    for instrument in midi_obj.instruments:
+        if not instrument.is_drum:
+            for note in instrument.notes:
+                start_time_sec = tick_to_time_map[note.start]
+                end_time_sec = tick_to_time_map[note.end]
+                
+                new_start_tick = round(start_time_sec * seconds_to_target_ticks_factor)
+                new_end_tick = round(end_time_sec * seconds_to_target_ticks_factor)
+                
+                if new_start_tick >= new_end_tick:
+                    # 确保音符至少有1 tick的长度
+                    new_end_tick = new_start_tick + 1
+
+                all_converted_notes.append(Note(velocity=note.velocity, 
+                                                pitch=note.pitch, 
+                                                start=new_start_tick, 
+                                                end=new_end_tick))
+
+    # --- 2. 清理重叠音符 (关键新增部分) ---
+    # 首先按音高分组，然后按开始时间排序
+    notes_by_pitch = defaultdict(list)
+    for note in all_converted_notes:
+        notes_by_pitch[note.pitch].append(note)
+
+    merged_notes = []
+    for pitch in sorted(notes_by_pitch.keys()):
+        # 对每个音高的音符列表按开始时间排序
+        sorted_notes = sorted(notes_by_pitch[pitch], key=lambda n: n.start)
+        
+        # 迭代并修复重叠
+        if len(sorted_notes) > 1:
+            for i in range(len(sorted_notes) - 1):
+                current_note = sorted_notes[i]
+                next_note = sorted_notes[i+1]
+                
+                # 如果当前音符的结束时间晚于或等于下一个音符的开始时间
+                if current_note.end >= next_note.start:
+                    # 修正当前音符的结束时间，让它在下一个音符开始前结束
+                    # 我们可以让它在下一个音符开始时就结束
+                    current_note.end = next_note.start
+                    # 如果修复后导致时长为0，则丢弃该音符（或者设置为1 tick，这里选择前者更干净）
+                    if current_note.start >= current_note.end:
+                         # 标记为待删除，而不是直接删除，以避免迭代问题
+                         current_note.pitch = -1 # 用一个无效音高作为标记
+
+        # 将处理过的（且未被标记删除的）音符添加到最终列表
+        merged_notes.extend([n for n in sorted_notes if n.pitch != -1])
+
+    # --- 3. 收集并转换CC事件 ---
+    merged_cc = []
+    for instrument in midi_obj.instruments:
+        if not instrument.is_drum:
+            for cc in instrument.control_changes:
+                if cc.number == 64:
+                    time_sec = tick_to_time_map[cc.time]
+                    new_time_tick = round(time_sec * seconds_to_target_ticks_factor)
+                    merged_cc.append(ControlChange(number=64, 
+                                                   value=cc.value, 
+                                                   time=new_time_tick))
+
+    # --- 4. 排序并创建新乐器 ---
+    merged_notes.sort(key=lambda x: (x.start, x.pitch))
+    merged_cc.sort(key=lambda x: (x.time, x.number))
+    
+    output_instrument = Instrument(program=0, is_drum=False, name="Piano")
+    output_instrument.notes = merged_notes
+    output_instrument.control_changes = merged_cc
+    output_midi_obj.instruments.append(output_instrument)
+    
+    # --- 5. 正确计算 max_tick ---
+    max_tick = 0
+    if output_instrument.notes:
+        max_tick = max(max_tick, max(n.end for n in output_instrument.notes if n.end is not None))
+    if output_instrument.control_changes:
+        max_tick = max(max_tick, max(c.time for c in output_instrument.control_changes if c.time is not None))
+    
+    # 添加一个小的buffer，确保最后一个事件不会被截断
+    output_midi_obj.max_tick = max_tick + target_ticks_per_beat 
+
+    return output_midi_obj
+
+def midi_to_ids(config, midi_obj, normalize=True):
+    def get_pedal(time_list, ccs, time):
+        i = bisect.bisect_right(time_list, time)
+        if i == 0:
+            return 0
+        else:
+            return ccs[i-1].value
+    if normalize:
+        norm_midi_obj = normalize_midi(midi_obj)
+    else:
+        norm_midi_obj = midi_obj
+    time_list = [cc.time for cc in norm_midi_obj.instruments[0].control_changes]
+    #print(time_list)
+    intervals = []
+    last_time = 0
+    for note in norm_midi_obj.instruments[0].notes:
+        intervals.append(note.start - last_time)
+        last_time = note.start
+    intervals.append(4990)
+
+    ids = []
+    last_time = 0
+    for i, note in enumerate(norm_midi_obj.instruments[0].notes):
+        interval = config.timing_start + intervals[i]
+        #print(interval - interval_start)
+
+        pitch = config.pitch_start + note.pitch
+        velocity = config.velocity_start + note.velocity
+        duration = config.timing_start + note.duration
+        last_time = last_time + intervals[i]
+
+        pedal1 = config.pedal_start + get_pedal(time_list, norm_midi_obj.instruments[0].control_changes, last_time)
+        pedal2 = config.pedal_start + get_pedal(time_list, norm_midi_obj.instruments[0].control_changes, last_time + intervals[i+1] * 1 / 4)
+        pedal3 = config.pedal_start + get_pedal(time_list, norm_midi_obj.instruments[0].control_changes, last_time + intervals[i+1] * 2 / 4)
+        pedal4 = config.pedal_start + get_pedal(time_list, norm_midi_obj.instruments[0].control_changes, last_time + intervals[i+1] * 3 / 4)
+        
+        pitch = min(config.valid_id_range[0][1] - 1, max(config.valid_id_range[0][0], pitch))
+        interval = min(config.valid_id_range[1][1] - 1, max(config.valid_id_range[1][0], interval))
+        velocity = min(config.valid_id_range[2][1] - 1, max(config.valid_id_range[2][0], velocity))
+        duration = min(config.valid_id_range[3][1] - 1, max(config.valid_id_range[3][0], duration))
+        pedal1 = min(config.valid_id_range[4][1] - 1, max(config.valid_id_range[4][0], pedal1))
+        pedal2 = min(config.valid_id_range[5][1] - 1, max(config.valid_id_range[5][0], pedal2))
+        pedal3 = min(config.valid_id_range[6][1] - 1, max(config.valid_id_range[6][0], pedal3))
+        pedal4 = min(config.valid_id_range[7][1] - 1, max(config.valid_id_range[7][0], pedal4))
+
+        ids.extend([pitch, interval, velocity, duration, pedal1, pedal2, pedal3, pedal4])
+    return ids
+
+def ids_to_midi(config, ids, target_ticks_per_beat = 500, target_tempo = 120):
+    note_list = []
+    cc_list = []
+    intervals = []
+    for i in range(0, len(ids), 8):
+        intervals.append(ids[i+1] - config.timing_start)
+    intervals.append(4990)
+    
+    last_time = 0
+    for i in range(0, len(ids), 8):
+        interval = intervals[i // 8]
+        pitch = ids[i] - config.pitch_start
+        velocity = ids[i+2] - config.velocity_start
+        duration = ids[i+3] - config.timing_start
+        pedal1 = ids[i+4] - config.pedal_start
+        pedal2 = ids[i+5] - config.pedal_start
+        pedal3 = ids[i+6] - config.pedal_start
+        pedal4 = ids[i+7] - config.pedal_start
+        note_list.append(Note(velocity, pitch, last_time + interval, last_time + interval + duration))
+        last_time += interval
+        #cc_list.append(ControlChange(64, pedal1, last_time))
+        #cc_list.append(ControlChange(64, pedal2, round(last_time + min(intervals[i // 8 + 1] * 1 / 10, 5))))
+        #cc_list.append(ControlChange(64, pedal3, round(last_time + max(intervals[i // 8 + 1] * 8 / 10, intervals[i // 8 + 1] * 8 / 10 - 10))))
+        #cc_list.append(ControlChange(64, pedal4, round(last_time + max(intervals[i // 8 + 1] * 9 / 10, intervals[i // 8 + 1] * 9 / 10 - 5))))
+        cc_list.append(ControlChange(64, pedal1, last_time))
+        cc_list.append(ControlChange(64, pedal2, round(last_time + intervals[i // 8 + 1] * 1 / 4)))
+        cc_list.append(ControlChange(64, pedal3, round(last_time + intervals[i // 8 + 1] * 2 / 4)))
+        cc_list.append(ControlChange(64, pedal4, round(last_time + intervals[i // 8 + 1] * 3 / 4)))
+
+    max_tick = 0
+    for note in note_list:
+        max_tick = max(max_tick, note.end)
+    for cc in cc_list:
+        max_tick = max(max_tick, cc.time)
+    max_tick = max_tick + 1
+
+    output = MidiFile(ticks_per_beat=target_ticks_per_beat)
+    output.instruments.append(Instrument(program=0, is_drum=False, name="Piano", notes=note_list, control_changes=cc_list))
+    output.tempo_changes.append(TempoChange(target_tempo, 0))
+    output.max_tick = max_tick
+    
+    return output
+
+def read_corresp(corresp_path):
+    out = []
+    performacne_id_list = []
+    with open(corresp_path, "r") as f:
+        align_txt = f.readlines()
+
+    score_ids_map = {}
+    performance_ids_map = {}
+    score_temp_list = []
+    performance_temp_list = set()
+    for line in align_txt[1:]:
+        informs = line.split("\t")
+        if informs[0] != '*':
+            score_temp_list.append((float(informs[1]), int(informs[3]), int(informs[0])))
+        if informs[5] != '*':
+            performance_temp_list.add((float(informs[6]), int(informs[8]), int(informs[5])))
+    performance_temp_list = list(performance_temp_list)
+    score_temp_list.sort()
+    performance_temp_list.sort()
+    for i, inform in enumerate(score_temp_list):
+        score_ids_map[inform[2]] = i
+    for i, inform in enumerate(performance_temp_list):
+        performance_ids_map[inform[2]] = i
+
+    for line in align_txt[1:]:
+        informs = line.split("\t")
+        if informs[0] == '*':
+            break
+        if informs[5] != '*':
+            out.append((score_ids_map[int(informs[0])], performance_ids_map[int(informs[5])]))
+        else:
+            out.append((score_ids_map[int(informs[0])], -1))
+    
+    for line in align_txt[1:]:
+        informs = line.split("\t")
+        if informs[5] != '*':
+            performacne_id_list.append(performance_ids_map[int(informs[5])])
+    if out[0][1] == -1:
+        out[0] = (out[0][0], min(performacne_id_list))
+    if out[-1][1] == -1:
+        out[-1] = (out[-1][0], max(performacne_id_list)) 
+    out.sort()
+    return out
+
+def interpolate(a, b):
+    a = np.array(a) + np.linspace(0, 1e-5, len(a))
+    b = np.array(b)
+    known_inds = np.where(~np.isnan(b))[0]
+    x_known = a[known_inds]
+    y_known = b[known_inds]
+    res = np.interp(a, x_known, y_known)
+    res[known_inds] = b[known_inds]
+    return [round(i) for i in res.tolist()]
+
+def segment_sequences(x, label, unknown_ids, total_notes, max_consecutive_missing, min_segment_notes):
+
+    if not unknown_ids:
+        if total_notes >= min_segment_notes:
+            return [x], [label]
+        else:
+            return [], []
+
+    x_segments = []
+    label_segments = []
+    
+    unknown_set = set(unknown_ids)
+    
+    last_cut_note_idx = 0
+    consecutive_missing_count = 0
+
+    for i in range(total_notes):
+        if i in unknown_set:
+            consecutive_missing_count += 1
+        else:
+            consecutive_missing_count = 0
+
+        if consecutive_missing_count >= max_consecutive_missing:
+            segment_end_note_idx = i - consecutive_missing_count + 1
+            
+            if segment_end_note_idx - last_cut_note_idx >= min_segment_notes:
+                start_token = last_cut_note_idx * 8
+                end_token = segment_end_note_idx * 8
+                
+                x_segments.append(x[start_token:end_token])
+                label_segments.append(label[start_token:end_token])
+            
+            last_cut_note_idx = i + 1
+            consecutive_missing_count = 0
+
+    if total_notes - last_cut_note_idx >= min_segment_notes:
+        start_token = last_cut_note_idx * 8
+        x_segments.append(x[start_token:])
+        label_segments.append(label[start_token:])
+        
+    return x_segments, label_segments
+
+def align_score_and_performance(config, score_midi_obj, performance_midi_obj):
+    norm_score_midi_obj = normalize_midi(score_midi_obj)
+    norm_performance_midi_obj = normalize_midi(performance_midi_obj)
+    
+    norm_score_midi_obj.dump("temp/score.mid")
+    norm_performance_midi_obj.dump("temp/performance.mid")
+
+    os.chdir("./tools/AlignmentTool")
+    os.system(f"timeout 120s ./MIDIToMIDIAlign.sh ../../temp/performance ../../temp/score")
+    os.chdir("./../../") 
+
+    corresp_list = read_corresp("temp/score_corresp.txt")
+    aligned_midi_obj = MidiFile(ticks_per_beat=500)
+    score_notes = norm_score_midi_obj.instruments[0].notes
+    performance_notes = norm_performance_midi_obj.instruments[0].notes
+    score_start_list = []
+    output_notes = []
+    output_ccs = []
+    vel_list = []
+    start_list = []
+    duration_list = []
+    unknown_ids = []
+    for i, ids in enumerate(corresp_list):
+        if ids[1] != -1:
+            vel_list.append(performance_notes[ids[1]].velocity)
+            start_list.append(performance_notes[ids[1]].start)
+            duration_list.append(performance_notes[ids[1]].end - performance_notes[ids[1]].start)
+        else:
+            vel_list.append(np.nan)
+            duration_list.append(np.nan)
+            unknown_ids.append(i)
+        score_start_list.append(score_notes[ids[0]].start)
+    start_list.sort()
+    temp = []
+    cnt = 0
+    for i in range(len(corresp_list)):
+        if i not in unknown_ids:
+            temp.append(start_list[cnt])
+            cnt += 1
+        else:
+            temp.append(np.nan)
+    start_list = interpolate(score_start_list, temp)
+    vel_list = interpolate(start_list, vel_list)
+    duration_list = interpolate(start_list, duration_list)
+
+    end_list = []
+    for i, ids in enumerate(corresp_list):
+        end = start_list[i]+duration_list[i]
+        end_list.append(end)
+        output_notes.append(Note(vel_list[i], score_notes[ids[0]].pitch, start_list[i], end))
+    max_tick = max(end_list) + 4999
+    for cc in norm_performance_midi_obj.instruments[0].control_changes:
+        if cc.time <= max_tick:
+            output_ccs.append(cc)
+        else:
+            break
+
+    aligned_midi_obj.instruments.append(Instrument(program=0, is_drum=False, name="Piano", notes=output_notes, control_changes=output_ccs))
+    x = midi_to_ids(config, norm_score_midi_obj)
+    label = midi_to_ids(config, aligned_midi_obj, normalize=False)
+    assert(len(x) == len(label))
+    for i in range(len(x)):
+        if i % 8 == 0:
+            assert(x[i] == label[i])
+
+    total_notes = len(score_notes)
+    xs, labels = segment_sequences(
+        x,
+        label,
+        unknown_ids,
+        total_notes,
+        5,
+        64,
+    )
+    return xs, labels
+
+def enhanced_ids(config, ids):
+    res = copy(ids)
+    retry = 10
+    for i in range(len(res)):
+        j = i % 8
+        if j == 3:
+            value = res[i] - config.valid_id_range[j][0]
+            if value == 10:
+                noise = 0
+                for _ in range(retry):
+                    n = round(np.random.randn() * 5)
+                    if n >= -9 and n <= 5:
+                        noise = n
+                        break
+            else:
+                noise = 0
+                for _ in range(retry):
+                    n = round(np.random.randn() * 5)
+                    if n >= -4 and n <= 5:
+                        noise = n
+                        break
+            value = min(max(value + noise, 0), 4999)
+            res[i] = config.valid_id_range[j][0] + value
+        elif j == 2:
+            value = res[i] - config.valid_id_range[j][0]
+            if value == 5:
+                noise = 0
+                for _ in range(retry):
+                    n = round(np.random.randn() * 2.5)
+                    if n >= -4 and n <= 2:
+                        noise = n
+                        break
+            elif value == 120:
+                noise = 0
+                for _ in range(retry):
+                    n = round(np.random.randn() * 2.5)
+                    if n >= -2 and n <= 7:
+                        noise = n
+                        break
+            else:
+                noise = 0
+                for _ in range(retry):
+                    n = round(np.random.randn() * 2.5)
+                    if n >= -2 and n <= 2:
+                        noise = n
+                        break
+            value = min(max(value + noise, 0), 127)
+            res[i] = config.valid_id_range[j][0] + value
+        elif j == 1:
+            value = res[i] - config.valid_id_range[j][0]
+            noise = 0
+            for _ in range(retry):
+                n = round(np.random.randn() * 5)
+                if n >= -4 and n <= 5:
+                    noise = n
+                    break
+            value = min(max(value + noise, 0), 4990)
+            res[i] = config.valid_id_range[j][0] + value
+    return res
+
+def enhanced_ids_uniform(config, ids):
+    res = copy(ids)
+    for i in range(len(res)):
+        j = i % 8
+        if j == 3:
+            value = res[i] - config.valid_id_range[j][0]
+            if value == 10:
+                noise = random.randint(-9, 5)
+            else:
+                noise = random.randint(-4, 5)
+            value = min(max(value + noise, 0), 4999)
+            res[i] = config.valid_id_range[j][0] + value
+        elif j == 2:
+            value = res[i] - config.valid_id_range[j][0]
+            if value == 5:
+                noise = random.randint(-4, 2)
+            elif value == 120:
+                noise = random.randint(-2, 7)
+            else:
+                noise = random.randint(-2, 2)
+            value = min(max(value + noise, 0), 127)
+            res[i] = config.valid_id_range[j][0] + value
+        elif j == 1:
+            value = res[i] - config.valid_id_range[j][0]
+            noise = random.randint(-4, 5)
+            value = min(max(value + noise, 0), 4990)
+            res[i] = config.valid_id_range[j][0] + value
+    return res
+
+#if __name__ == "__main__":
+#    midi_obj = MidiFile("data/midi/test/2.mid")
+#    ids = midi_to_ids(midi_obj)
+#    midi = ids_to_midi(ids)
+#    midi.dump("data/rebuild/2.mid")