upl base

.gitignore

@@ -0,0 +1,10 @@
+*__pycache__*
+output/*
+rename.sh
+test.py
+gpt2-abcmusic/*
+*.pth
+flagged/*
+mscore3/*
+key.txt
+feedbacks/*

app.py

@@ -0,0 +1,552 @@
+import re
+import os
+import json
+import time
+import torch
+import random
+import shutil
+import argparse
+import warnings
+import gradio as gr
+import soundfile as sf
+from transformers import GPT2Config
+from model import Patchilizer, TunesFormer
+from convert import abc2xml, xml2img, xml2, transpose_octaves_abc
+from utils import (
+    PATCH_NUM_LAYERS,
+    PATCH_LENGTH,
+    CHAR_NUM_LAYERS,
+    PATCH_SIZE,
+    SHARE_WEIGHTS,
+    TEMP_DIR,
+    WEIGHTS_DIR,
+    DEVICE,
+)
+
+
+def get_args(parser: argparse.ArgumentParser):
+    parser.add_argument(
+        "-num_tunes",
+        type=int,
+        default=1,
+        help="the number of independently computed returned tunes",
+    )
+    parser.add_argument(
+        "-max_patch",
+        type=int,
+        default=128,
+        help="integer to define the maximum length in tokens of each tune",
+    )
+    parser.add_argument(
+        "-top_p",
+        type=float,
+        default=0.8,
+        help="float to define the tokens that are within the sample operation of text generation",
+    )
+    parser.add_argument(
+        "-top_k",
+        type=int,
+        default=8,
+        help="integer to define the tokens that are within the sample operation of text generation",
+    )
+    parser.add_argument(
+        "-temperature",
+        type=float,
+        default=1.2,
+        help="the temperature of the sampling operation",
+    )
+    parser.add_argument("-seed", type=int, default=None, help="seed for randomstate")
+    parser.add_argument(
+        "-show_control_code",
+        type=bool,
+        default=False,
+        help="whether to show control code",
+    )
+    return parser.parse_args()
+
+
+def get_abc_key_val(text: str, key="K"):
+    pattern = re.escape(key) + r":(.*?)\n"
+    match = re.search(pattern, text)
+    if match:
+        return match.group(1).strip()
+    else:
+        return None
+
+
+def adjust_volume(in_audio: str, dB_change: int):
+    y, sr = sf.read(in_audio)
+    sf.write(in_audio, y * 10 ** (dB_change / 20), sr)
+
+
+def generate_music(
+    args,
+    emo: str,
+    weights: str,
+    outdir=TEMP_DIR,
+    fix_tempo=None,
+    fix_pitch=None,
+    fix_volume=None,
+):
+    patchilizer = Patchilizer()
+    patch_config = GPT2Config(
+        num_hidden_layers=PATCH_NUM_LAYERS,
+        max_length=PATCH_LENGTH,
+        max_position_embeddings=PATCH_LENGTH,
+        vocab_size=1,
+    )
+    char_config = GPT2Config(
+        num_hidden_layers=CHAR_NUM_LAYERS,
+        max_length=PATCH_SIZE,
+        max_position_embeddings=PATCH_SIZE,
+        vocab_size=128,
+    )
+    model = TunesFormer(patch_config, char_config, share_weights=SHARE_WEIGHTS)
+    checkpoint = torch.load(weights)
+    model.load_state_dict(checkpoint["model"])
+    model = model.to(DEVICE)
+    model.eval()
+    prompt = f"A:{emo}\n"
+    tunes = ""
+    num_tunes = args.num_tunes
+    max_patch = args.max_patch
+    top_p = args.top_p
+    top_k = args.top_k
+    temperature = args.temperature
+    seed = args.seed
+    show_control_code = args.show_control_code
+    print(" Hyper parms ".center(60, "#"), "\n")
+    args_dict: dict = vars(args)
+    for arg in args_dict.keys():
+        print(f"{arg}: {str(args_dict[arg])}")
+
+    print("\n", " Output tunes ".center(60, "#"))
+    start_time = time.time()
+    for i in range(num_tunes):
+        title = f"T:{emo} Fragment\n"
+        artist = f"C:Generated by AI\n"
+        tune = f"X:{str(i + 1)}\n{title}{artist}{prompt}"
+        lines = re.split(r"(\n)", tune)
+        tune = ""
+        skip = False
+        for line in lines:
+            if show_control_code or line[:2] not in ["S:", "B:", "E:"]:
+                if not skip:
+                    print(line, end="")
+                    tune += line
+
+                skip = False
+
+            else:
+                skip = True
+
+        input_patches = torch.tensor(
+            [patchilizer.encode(prompt, add_special_patches=True)[:-1]],
+            device=DEVICE,
+        )
+        if tune == "":
+            tokens = None
+
+        else:
+            prefix = patchilizer.decode(input_patches[0])
+            remaining_tokens = prompt[len(prefix) :]
+            tokens = torch.tensor(
+                [patchilizer.bos_token_id] + [ord(c) for c in remaining_tokens],
+                device=DEVICE,
+            )
+
+        while input_patches.shape[1] < max_patch:
+            predicted_patch, seed = model.generate(
+                input_patches,
+                tokens,
+                top_p=top_p,
+                top_k=top_k,
+                temperature=temperature,
+                seed=seed,
+            )
+            tokens = None
+            if predicted_patch[0] != patchilizer.eos_token_id:
+                next_bar = patchilizer.decode([predicted_patch])
+                if show_control_code or next_bar[:2] not in ["S:", "B:", "E:"]:
+                    print(next_bar, end="")
+                    tune += next_bar
+
+                if next_bar == "":
+                    break
+
+                next_bar = remaining_tokens + next_bar
+                remaining_tokens = ""
+                predicted_patch = torch.tensor(
+                    patchilizer.bar2patch(next_bar),
+                    device=DEVICE,
+                ).unsqueeze(0)
+                input_patches = torch.cat(
+                    [input_patches, predicted_patch.unsqueeze(0)],
+                    dim=1,
+                )
+
+            else:
+                break
+
+        tunes += f"{tune}\n\n"
+        print("\n")
+
+    # fix tempo
+    if fix_tempo != None:
+        tempo = f"Q:{fix_tempo}\n"
+
+    else:
+        tempo = f"Q:{random.randint(88, 132)}\n"
+        if emo == "Q1":
+            tempo = f"Q:{random.randint(160, 184)}\n"
+        elif emo == "Q2":
+            tempo = f"Q:{random.randint(184, 228)}\n"
+        elif emo == "Q3":
+            tempo = f"Q:{random.randint(40, 69)}\n"
+        elif emo == "Q4":
+            tempo = f"Q:{random.randint(40, 69)}\n"
+
+    Q_val = get_abc_key_val(tunes, "Q")
+    if Q_val:
+        tunes = tunes.replace(f"Q:{Q_val}\n", "")
+
+    tunes = tunes.replace(f"A:{emo}\n", tempo)
+    # fix mode:major/minor
+    mode = "major" if emo == "Q1" or emo == "Q4" else "minor"
+    K_val = get_abc_key_val(tunes)
+    if mode == "major" and K_val and "m" in K_val:
+        tunes = tunes.replace(f"\nK:{K_val}\n", f"\nK:{K_val.split('m')[0]}\n")
+
+    elif mode == "minor" and K_val and not "m" in K_val:
+        tunes = tunes.replace(f"\nK:{K_val}\n", f"\nK:{K_val.lower()}min\n")
+
+    print("Generation time: {:.2f} seconds".format(time.time() - start_time))
+    timestamp = time.strftime("%a_%d_%b_%Y_%H_%M_%S", time.localtime())
+    try:
+        # fix avg_pitch (octave)
+        if fix_pitch != None:
+            if fix_pitch:
+                tunes, xml = transpose_octaves_abc(
+                    tunes,
+                    f"{outdir}/{timestamp}.musicxml",
+                    fix_pitch,
+                )
+                tunes = tunes.replace(title + title, title)
+                os.rename(xml, f"{outdir}/[{emo}]{timestamp}.musicxml")
+                xml = f"{outdir}/[{emo}]{timestamp}.musicxml"
+
+        else:
+            if mode == "minor":
+                offset = -12
+                if emo == "Q2":
+                    offset -= 12
+
+                tunes, xml = transpose_octaves_abc(
+                    tunes,
+                    f"{outdir}/{timestamp}.musicxml",
+                    offset,
+                )
+                tunes = tunes.replace(title + title, title)
+                os.rename(xml, f"{outdir}/[{emo}]{timestamp}.musicxml")
+                xml = f"{outdir}/[{emo}]{timestamp}.musicxml"
+
+            else:
+                xml = abc2xml(tunes, f"{outdir}/[{emo}]{timestamp}.musicxml")
+
+        audio = xml2(xml, "wav")
+        if fix_volume != None:
+            if fix_volume:
+                adjust_volume(audio, fix_volume)
+
+        elif os.path.exists(audio):
+            if emo == "Q1":
+                adjust_volume(audio, 5)
+
+            elif emo == "Q2":
+                adjust_volume(audio, 10)
+
+        mxl = xml2(xml, "mxl")
+        midi = xml2(xml, "mid")
+        pdf, jpg = xml2img(xml)
+        return audio, midi, pdf, xml, mxl, tunes, jpg
+
+    except Exception as e:
+        print(f"{e}")
+        return generate_music(args, emo, weights)
+
+
+def inference(dataset: str, v: str, a: str, add_chord: bool):
+    if os.path.exists(TEMP_DIR):
+        shutil.rmtree(TEMP_DIR)
+
+    os.makedirs(TEMP_DIR, exist_ok=True)
+    emotion = "Q1"
+    if v == "Low" and a == "High":
+        emotion = "Q2"
+
+    elif v == "Low" and a == "Low":
+        emotion = "Q3"
+
+    elif v == "High" and a == "Low":
+        emotion = "Q4"
+
+    parser = argparse.ArgumentParser()
+    args = get_args(parser)
+    return generate_music(
+        args,
+        emo=emotion,
+        weights=f"{WEIGHTS_DIR}/{dataset.lower()}/weights.pth",
+    )
+
+
+def infer(
+    dataset: str,
+    pitch_std: str,
+    mode: str,
+    tempo: int,
+    octave: int,
+    rms: int,
+    add_chord: bool,
+):
+    if os.path.exists(TEMP_DIR):
+        shutil.rmtree(TEMP_DIR)
+
+    os.makedirs(TEMP_DIR, exist_ok=True)
+    emotion = "Q1"
+    if mode == "Minor" and pitch_std == "High":
+        emotion = "Q2"
+
+    elif mode == "Minor" and pitch_std == "Low":
+        emotion = "Q3"
+
+    elif mode == "Major" and pitch_std == "Low":
+        emotion = "Q4"
+
+    parser = argparse.ArgumentParser()
+    args = get_args(parser)
+    return generate_music(
+        args,
+        emo=emotion,
+        weights=f"{WEIGHTS_DIR}/{dataset.lower()}/weights.pth",
+        fix_tempo=tempo,
+        fix_pitch=octave,
+        fix_volume=rms,
+    )
+
+
+def feedback(fixed_emo: str, source_dir="./flagged", target_dir="./feedbacks"):
+    if not fixed_emo:
+        return "Please select feedback before submitting! "
+
+    os.makedirs(target_dir, exist_ok=True)
+    for root, _, files in os.walk(source_dir):
+        for file in files:
+            if file.endswith(".mxl"):
+                prompt_emo = file.split("]")[0][1:]
+                if prompt_emo != fixed_emo:
+                    file_path = os.path.join(root, file)
+                    target_path = os.path.join(
+                        target_dir, file.replace(".mxl", f"_{fixed_emo}.mxl")
+                    )
+                    shutil.copy(file_path, target_path)
+                    return f"Copied {file_path} to {target_path}"
+
+                else:
+                    return "Thanks for your feedback!"
+
+    return "No .mxl files found in the source directory."
+
+
+def save_template(
+    label: str,
+    pitch_std: str,
+    mode: str,
+    tempo: int,
+    octave: int,
+    rms: int,
+):
+    if (
+        label
+        and pitch_std
+        and mode
+        and tempo != None
+        and octave != None
+        and rms != None
+    ):
+        json_str = json.dumps(
+            {
+                "label": label,
+                "pitch_std": pitch_std == "High",
+                "mode": mode == "Major",
+                "tempo": tempo,
+                "octave": octave,
+                "volume": rms,
+            }
+        )
+
+        with open("./feedbacks/templates.jsonl", "a", encoding="utf-8") as file:
+            file.write(json_str + "\n")
+
+
+if __name__ == "__main__":
+    warnings.filterwarnings("ignore")
+    if os.path.exists("./flagged"):
+        shutil.rmtree("./flagged")
+
+    with gr.Blocks() as demo:
+        with gr.Row():
+            with gr.Column():
+                dataset_option = gr.Dropdown(
+                    ["VGMIDI", "EMOPIA", "Rough4Q"],
+                    label="Dataset",
+                    value="Rough4Q",
+                )
+                gr.Markdown(
+                    "# Generate by emotion condition<br><img width='100%' src='https://www.modelscope.cn/studio/monetjoe/EMusicGen/resolve/master/4q.jpg'>"
+                )
+                valence_radio = gr.Radio(
+                    ["Low", "High"],
+                    label="Valence (reflects negative-positive levels of emotion)",
+                    value="High",
+                )
+                arousal_radio = gr.Radio(
+                    ["Low", "High"],
+                    label="Arousal (reflects the calmness-intensity of the emotion)",
+                    value="High",
+                )
+                chord_check = gr.Checkbox(
+                    label="Generate chords",
+                    value=False,
+                )
+                gen_btn = gr.Button("Generate")
+                gr.Markdown("# Generate by feature control")
+                std_option = gr.Radio(
+                    ["Low", "High"],
+                    label="Pitch SD",
+                    value="High",
+                )
+                mode_option = gr.Radio(
+                    ["Minor", "Major"],
+                    label="Mode",
+                    value="Major",
+                )
+                tempo_option = gr.Slider(
+                    minimum=40,
+                    maximum=228,
+                    step=1,
+                    value=120,
+                    label="Tempo (BPM)",
+                )
+                octave_option = gr.Slider(
+                    minimum=-24,
+                    maximum=24,
+                    step=12,
+                    value=0,
+                    label="Octave (±12)",
+                )
+                volume_option = gr.Slider(
+                    minimum=-5,
+                    maximum=10,
+                    step=5,
+                    value=0,
+                    label="Volume (dB)",
+                )
+                chord_check_2 = gr.Checkbox(
+                    label="Generate chords",
+                    value=False,
+                )
+                gen_btn_2 = gr.Button("Generate")
+                template_radio = gr.Radio(
+                    ["Q1", "Q2", "Q3", "Q4"],
+                    label="The emotion to which the current template belongs",
+                )
+                save_btn = gr.Button("Save template")
+                gr.Markdown(
+                    """
+## Cite
+```bibtex
+@article{Zhou2024EMusicGen,
+title     = {EMusicGen: Emotion-Conditioned Melody Generation in ABC Notation},
+author    = {Monan Zhou, Xiaobing Li, Feng Yu and Wei Li},
+month     = {Sep},
+year      = {2024},
+publisher = {GitHub},
+version   = {0.1},
+url       = {https://github.com/monetjoe/EMusicGen}
+}
+```
+"""
+                )
+
+            with gr.Column():
+                wav_audio = gr.Audio(label="Audio", type="filepath")
+                midi_file = gr.File(label="Download MIDI")
+                pdf_file = gr.File(label="Download PDF score")
+                xml_file = gr.File(label="Download MusicXML")
+                mxl_file = gr.File(label="Download MXL")
+                abc_textbox = gr.Textbox(
+                    label="ABC notation",
+                    show_copy_button=True,
+                )
+                staff_img = gr.Image(label="Staff", type="filepath")
+
+        with gr.Row():
+            gr.Interface(
+                fn=feedback,
+                inputs=gr.Radio(
+                    ["Q1", "Q2", "Q3", "Q4"],
+                    label="Feedback: the emotion you believe the generated result should belong to",
+                ),
+                outputs=gr.Textbox(show_copy_button=False, show_label=False),
+                allow_flagging="never",
+            )
+
+        gen_btn.click(
+            fn=inference,
+            inputs=[dataset_option, valence_radio, arousal_radio, chord_check],
+            outputs=[
+                wav_audio,
+                midi_file,
+                pdf_file,
+                xml_file,
+                mxl_file,
+                abc_textbox,
+                staff_img,
+            ],
+        )
+
+        gen_btn_2.click(
+            fn=infer,
+            inputs=[
+                dataset_option,
+                std_option,
+                mode_option,
+                tempo_option,
+                octave_option,
+                volume_option,
+                chord_check,
+            ],
+            outputs=[
+                wav_audio,
+                midi_file,
+                pdf_file,
+                xml_file,
+                mxl_file,
+                abc_textbox,
+                staff_img,
+            ],
+        )
+
+        save_btn.click(
+            fn=save_template,
+            inputs=[
+                template_radio,
+                std_option,
+                mode_option,
+                tempo_option,
+                octave_option,
+                volume_option,
+            ],
+        )
+
+    demo.launch()

conda.txt

@@ -0,0 +1,5 @@
+python=3.10
+pytorch=1.12.1
+torchvision=0.13.1
+torchaudio=0.12.1
+cudatoolkit=11.3.1

convert.py

@@ -0,0 +1,91 @@
+import os
+import fitz
+import subprocess
+from PIL import Image
+from music21 import converter, interval, clef, stream
+from utils import MSCORE
+
+
+def abc2xml(abc_content, output_xml_path):
+    score = converter.parse(abc_content, format="abc")
+    score.write("musicxml", fp=output_xml_path, encoding="utf-8")
+    return output_xml_path
+
+
+def xml2(xml_path: str, target_fmt: str):
+    src_fmt = os.path.basename(xml_path).split(".")[-1]
+    if not "." in target_fmt:
+        target_fmt = "." + target_fmt
+
+    target_file = xml_path.replace(f".{src_fmt}", target_fmt)
+    print(subprocess.run([MSCORE, "-o", target_file, xml_path]))
+    return target_file
+
+
+def pdf2img(pdf_path: str):
+    output_path = pdf_path.replace(".pdf", ".jpg")
+    doc = fitz.open(pdf_path)
+    # 创建一个图像列表
+    images = []
+    for page_number in range(doc.page_count):
+        page = doc[page_number]
+        # 将页面渲染为图像
+        image = page.get_pixmap()
+        # 将图像添加到列表
+        images.append(
+            Image.frombytes("RGB", [image.width, image.height], image.samples)
+        )
+    # 竖向合并图像
+    merged_image = Image.new(
+        "RGB", (images[0].width, sum(image.height for image in images))
+    )
+    y_offset = 0
+    for image in images:
+        merged_image.paste(image, (0, y_offset))
+        y_offset += image.height
+    # 保存合并后的图像为JPG
+    merged_image.save(output_path, "JPEG")
+    # 关闭PDF文档
+    doc.close()
+    return output_path
+
+
+def xml2img(xml_file: str):
+    ext = os.path.basename(xml_file).split(".")[-1]
+    pdf_score = xml_file.replace(f".{ext}", ".pdf")
+    command = [MSCORE, "-o", pdf_score, xml_file]
+    result = subprocess.run(command)
+    print(result)
+    return pdf_score, pdf2img(pdf_score)
+
+
+# xml to abc
+def xml2abc(input_xml_file: str):
+    result = subprocess.run(
+        ["python", "-Xfrozen_modules=off", "./xml2abc.py", input_xml_file],
+        stdout=subprocess.PIPE,
+        text=True,
+    )
+
+    if result.returncode == 0:
+        return str(result.stdout).strip()
+
+    return ""
+
+
+def transpose_octaves_abc(abc_notation: str, out_xml_file: str, offset=-12):
+    score = converter.parse(abc_notation)
+    if offset < 0:
+        for part in score.parts:
+            for measure in part.getElementsByClass(stream.Measure):
+                # 检查当前小节的谱号
+                if measure.clef:
+                    measure.clef = clef.BassClef()
+
+    octaves_interval = interval.Interval(offset)
+    # 遍历每个音符，将其上下移八度
+    for note in score.recurse().notes:
+        note.transpose(octaves_interval, inPlace=True)
+
+    score.write("musicxml", fp=out_xml_file)
+    return xml2abc(out_xml_file), out_xml_file

model.py

@@ -0,0 +1,350 @@
+import re
+import torch
+import random
+from tqdm import tqdm
+from unidecode import unidecode
+from torch.utils.data import Dataset
+from transformers import GPT2Model, GPT2LMHeadModel, PreTrainedModel
+from samplings import top_p_sampling, top_k_sampling, temperature_sampling
+from utils import PATCH_SIZE, PATCH_LENGTH, PATCH_SAMPLING_BATCH_SIZE
+
+
+class Patchilizer:
+    """
+    A class for converting music bars to patches and vice versa.
+    """
+
+    def __init__(self):
+        self.delimiters = ["|:", "::", ":|", "[|", "||", "|]", "|"]
+        self.regexPattern = f"({'|'.join(map(re.escape, self.delimiters))})"
+        self.pad_token_id = 0
+        self.bos_token_id = 1
+        self.eos_token_id = 2
+
+    def split_bars(self, body):
+        """
+        Split a body of music into individual bars.
+        """
+        bars = re.split(self.regexPattern, "".join(body))
+        bars = list(filter(None, bars))
+        # remove empty strings
+        if bars[0] in self.delimiters:
+            bars[1] = bars[0] + bars[1]
+            bars = bars[1:]
+
+        bars = [bars[i * 2] + bars[i * 2 + 1] for i in range(len(bars) // 2)]
+        return bars
+
+    def bar2patch(self, bar, patch_size=PATCH_SIZE):
+        """
+        Convert a bar into a patch of specified length.
+        """
+        patch = [self.bos_token_id] + [ord(c) for c in bar] + [self.eos_token_id]
+        patch = patch[:patch_size]
+        patch += [self.pad_token_id] * (patch_size - len(patch))
+        return patch
+
+    def patch2bar(self, patch):
+        """
+        Convert a patch into a bar.
+        """
+        return "".join(
+            chr(idx) if idx > self.eos_token_id else ""
+            for idx in patch
+            if idx != self.eos_token_id
+        )
+
+    def encode(
+        self,
+        abc_code,
+        patch_length=PATCH_LENGTH,
+        patch_size=PATCH_SIZE,
+        add_special_patches=False,
+    ):
+        """
+        Encode music into patches of specified length.
+        """
+        lines = unidecode(abc_code).split("\n")
+        lines = list(filter(None, lines))  # remove empty lines
+
+        body = ""
+        patches = []
+
+        for line in lines:
+            if len(line) > 1 and (
+                (line[0].isalpha() and line[1] == ":") or line.startswith("%%score")
+            ):
+                if body:
+                    bars = self.split_bars(body)
+                    patches.extend(
+                        self.bar2patch(
+                            bar + "\n" if idx == len(bars) - 1 else bar, patch_size
+                        )
+                        for idx, bar in enumerate(bars)
+                    )
+                    body = ""
+
+                patches.append(self.bar2patch(line + "\n", patch_size))
+
+            else:
+                body += line + "\n"
+
+        if body:
+            patches.extend(
+                self.bar2patch(bar, patch_size) for bar in self.split_bars(body)
+            )
+
+        if add_special_patches:
+            bos_patch = [self.bos_token_id] * (patch_size - 1) + [self.eos_token_id]
+            eos_patch = [self.bos_token_id] + [self.eos_token_id] * (patch_size - 1)
+            patches = [bos_patch] + patches + [eos_patch]
+
+        return patches[:patch_length]
+
+    def decode(self, patches):
+        """
+        Decode patches into music.
+        """
+        return "".join(self.patch2bar(patch) for patch in patches)
+
+
+class PatchLevelDecoder(PreTrainedModel):
+    """
+    An Patch-level Decoder model for generating patch features in an auto-regressive manner.
+    It inherits PreTrainedModel from transformers.
+    """
+
+    def __init__(self, config):
+        super().__init__(config)
+        self.patch_embedding = torch.nn.Linear(PATCH_SIZE * 128, config.n_embd)
+        torch.nn.init.normal_(self.patch_embedding.weight, std=0.02)
+        self.base = GPT2Model(config)
+
+    def forward(self, patches: torch.Tensor) -> torch.Tensor:
+        """
+        The forward pass of the patch-level decoder model.
+        :param patches: the patches to be encoded
+        :return: the encoded patches
+        """
+        patches = torch.nn.functional.one_hot(patches, num_classes=128).float()
+        patches = patches.reshape(len(patches), -1, PATCH_SIZE * 128)
+        patches = self.patch_embedding(patches.to(self.device))
+
+        return self.base(inputs_embeds=patches)
+
+
+class CharLevelDecoder(PreTrainedModel):
+    """
+    A Char-level Decoder model for generating the characters within each bar patch sequentially.
+    It inherits PreTrainedModel from transformers.
+    """
+
+    def __init__(self, config):
+        super().__init__(config)
+        self.pad_token_id = 0
+        self.bos_token_id = 1
+        self.eos_token_id = 2
+        self.base = GPT2LMHeadModel(config)
+
+    def forward(
+        self,
+        encoded_patches: torch.Tensor,
+        target_patches: torch.Tensor,
+        patch_sampling_batch_size: int,
+    ):
+        """
+        The forward pass of the char-level decoder model.
+        :param encoded_patches: the encoded patches
+        :param target_patches: the target patches
+        :return: the decoded patches
+        """
+        # preparing the labels for model training
+        target_masks = target_patches == self.pad_token_id
+        labels = target_patches.clone().masked_fill_(target_masks, -100)
+
+        # masking the labels for model training
+        target_masks = torch.ones_like(labels)
+        target_masks = target_masks.masked_fill_(labels == -100, 0)
+
+        # select patches
+        if (
+            patch_sampling_batch_size != 0
+            and patch_sampling_batch_size < target_patches.shape[0]
+        ):
+            indices = list(range(len(target_patches)))
+            random.shuffle(indices)
+            selected_indices = sorted(indices[:patch_sampling_batch_size])
+
+            target_patches = target_patches[selected_indices, :]
+            target_masks = target_masks[selected_indices, :]
+            encoded_patches = encoded_patches[selected_indices, :]
+            labels = labels[selected_indices, :]
+
+        # get input embeddings
+        inputs_embeds = torch.nn.functional.embedding(
+            target_patches, self.base.transformer.wte.weight
+        )
+
+        # concatenate the encoded patches with the input embeddings
+        inputs_embeds = torch.cat(
+            (encoded_patches.unsqueeze(1), inputs_embeds[:, 1:, :]), dim=1
+        )
+
+        return self.base(
+            inputs_embeds=inputs_embeds, attention_mask=target_masks, labels=labels
+        )
+
+    def generate(self, encoded_patch: torch.Tensor, tokens: torch.Tensor):
+        """
+        The generate function for generating a patch based on the encoded patch and already generated tokens.
+        :param encoded_patch: the encoded patch
+        :param tokens: already generated tokens in the patch
+        :return: the probability distribution of next token
+        """
+        encoded_patch = encoded_patch.reshape(1, 1, -1)
+        tokens = tokens.reshape(1, -1)
+
+        # Get input embeddings
+        tokens = torch.nn.functional.embedding(tokens, self.base.transformer.wte.weight)
+
+        # Concatenate the encoded patch with the input embeddings
+        tokens = torch.cat((encoded_patch, tokens[:, 1:, :]), dim=1)
+
+        # Get output from model
+        outputs = self.base(inputs_embeds=tokens)
+
+        # Get probabilities of next token
+        probs = torch.nn.functional.softmax(outputs.logits.squeeze(0)[-1], dim=-1)
+
+        return probs
+
+
+class TunesFormer(PreTrainedModel):
+    """
+    TunesFormer is a hierarchical music generation model based on bar patching.
+    It includes a patch-level decoder and a character-level decoder.
+    It inherits PreTrainedModel from transformers.
+    """
+
+    def __init__(self, encoder_config, decoder_config, share_weights=False):
+        super().__init__(encoder_config)
+        self.pad_token_id = 0
+        self.bos_token_id = 1
+        self.eos_token_id = 2
+        if share_weights:
+            max_layers = max(
+                encoder_config.num_hidden_layers, decoder_config.num_hidden_layers
+            )
+
+            max_context_size = max(encoder_config.max_length, decoder_config.max_length)
+
+            max_position_embeddings = max(
+                encoder_config.max_position_embeddings,
+                decoder_config.max_position_embeddings,
+            )
+
+            encoder_config.num_hidden_layers = max_layers
+            encoder_config.max_length = max_context_size
+            encoder_config.max_position_embeddings = max_position_embeddings
+            decoder_config.num_hidden_layers = max_layers
+            decoder_config.max_length = max_context_size
+            decoder_config.max_position_embeddings = max_position_embeddings
+
+        self.patch_level_decoder = PatchLevelDecoder(encoder_config)
+        self.char_level_decoder = CharLevelDecoder(decoder_config)
+
+        if share_weights:
+            self.patch_level_decoder.base = self.char_level_decoder.base.transformer
+
+    def forward(
+        self,
+        patches: torch.Tensor,
+        patch_sampling_batch_size: int = PATCH_SAMPLING_BATCH_SIZE,
+    ):
+        """
+        The forward pass of the TunesFormer model.
+        :param patches: the patches to be both encoded and decoded
+        :return: the decoded patches
+        """
+        patches = patches.reshape(len(patches), -1, PATCH_SIZE)
+        encoded_patches = self.patch_level_decoder(patches)["last_hidden_state"]
+
+        return self.char_level_decoder(
+            encoded_patches.squeeze(0)[:-1, :],
+            patches.squeeze(0)[1:, :],
+            patch_sampling_batch_size,
+        )
+
+    def generate(
+        self,
+        patches: torch.Tensor,
+        tokens: torch.Tensor,
+        top_p: float = 1,
+        top_k: int = 0,
+        temperature: float = 1,
+        seed: int = None,
+    ):
+        """
+        The generate function for generating patches based on patches.
+        :param patches: the patches to be encoded
+        :return: the generated patches
+        """
+        patches = patches.reshape(len(patches), -1, PATCH_SIZE)
+        encoded_patches = self.patch_level_decoder(patches)["last_hidden_state"]
+
+        if tokens == None:
+            tokens = torch.tensor([self.bos_token_id], device=self.device)
+
+        generated_patch = []
+        random.seed(seed)
+
+        while True:
+            if seed != None:
+                n_seed = random.randint(0, 1000000)
+                random.seed(n_seed)
+
+            else:
+                n_seed = None
+
+            prob = (
+                self.char_level_decoder.generate(encoded_patches[0][-1], tokens)
+                .cpu()
+                .detach()
+                .numpy()
+            )
+
+            prob = top_p_sampling(prob, top_p=top_p, return_probs=True)
+            prob = top_k_sampling(prob, top_k=top_k, return_probs=True)
+
+            token = temperature_sampling(prob, temperature=temperature, seed=n_seed)
+
+            generated_patch.append(token)
+            if token == self.eos_token_id or len(tokens) >= PATCH_SIZE - 1:
+                break
+
+            else:
+                tokens = torch.cat(
+                    (tokens, torch.tensor([token], device=self.device)), dim=0
+                )
+
+        return generated_patch, n_seed
+
+
+class PatchilizedData(Dataset):
+    def __init__(self, items, patchilizer):
+        self.texts = []
+
+        for item in tqdm(items):
+            text = item["control code"] + "\n".join(
+                item["abc notation"].split("\n")[1:]
+            )
+            input_patch = patchilizer.encode(text, add_special_patches=True)
+            input_patch = torch.tensor(input_patch)
+            if torch.sum(input_patch) != 0:
+                self.texts.append(input_patch)
+
+    def __len__(self):
+        return len(self.texts)
+
+    def __getitem__(self, idx):
+        return self.texts[idx]

requirements.txt

@@ -0,0 +1,8 @@
+music21
+pymupdf
+autopep8
+unidecode
+pillow==9.4.0
+samplings==0.1.7
+modelscope==1.15
+transformers==4.18.0

utils.py

@@ -0,0 +1,62 @@
+import os
+import sys
+import time
+import torch
+import requests
+import subprocess
+from tqdm import tqdm
+from modelscope.hub.api import HubApi
+from modelscope import snapshot_download
+
+HubApi().login(os.getenv("ms_app_key"))
+TEMP_DIR = "./flagged"
+WEIGHTS_DIR = snapshot_download("monetjoe/EMusicGen", cache_dir="./__pycache__")
+DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+PATCH_LENGTH = 128  # Patch Length
+PATCH_SIZE = 32  # Patch Size
+PATCH_NUM_LAYERS = 9  # Number of layers in the encoder
+CHAR_NUM_LAYERS = 3  # Number of layers in the decoder
+PATCH_SAMPLING_BATCH_SIZE = 0  # Batch size for training patch, 0 for full context
+LOAD_FROM_CHECKPOINT = True  # Whether to load weights from a checkpoint
+SHARE_WEIGHTS = False  # Whether to share weights between the encoder and decoder
+
+
+def download(filename: str, url: str):
+    try:
+        response = requests.get(url, stream=True)
+        total_size = int(response.headers.get("content-length", 0))
+        chunk_size = 1024
+
+        with open(filename, "wb") as file, tqdm(
+            desc=f"Downloading {filename} from {url}...",
+            total=total_size,
+            unit="B",
+            unit_scale=True,
+            unit_divisor=1024,
+        ) as bar:
+            for data in response.iter_content(chunk_size=chunk_size):
+                size = file.write(data)
+                bar.update(size)
+
+    except Exception as e:
+        print(f"Error: {e}")
+        time.sleep(10)
+        download(filename, url)
+
+
+if sys.platform.startswith("linux"):
+    apkname = "MuseScore.AppImage"
+    extra_dir = "squashfs-root"
+    download(
+        filename=apkname,
+        url="https://www.modelscope.cn/studio/MuGeminorum/piano_transcription/resolve/master/MuseScore.AppImage",
+    )
+    if not os.path.exists(extra_dir):
+        subprocess.run(["chmod", "+x", f"./{apkname}"])
+        subprocess.run([f"./{apkname}", "--appimage-extract"])
+
+    MSCORE = f"./{extra_dir}/AppRun"
+    os.environ["QT_QPA_PLATFORM"] = "offscreen"
+
+else:
+    MSCORE = os.getenv("mscore")