Upload 5 files

app.py

@@ -1,237 +1,267 @@
-import re
-import os
-import time
-import torch
-import shutil
-import argparse
-import gradio as gr
-from utils import *
-from config import *
-from convert import *
-from transformers import GPT2Config
-import warnings
-
-warnings.filterwarnings("ignore")
-
-
-def get_args(parser):
-    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=True,
-        help="whether to show control code",
-    )
-    args = parser.parse_args()
-
-    return args
-
-
-def generate_abc(args, region):
-    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)
-
-    filename = WEIGHT_PATH
-
-    if os.path.exists(filename):
-        print(f"Weights already exist at '{filename}'. Loading...")
-
-    else:
-        download()
-
-    checkpoint = torch.load(filename, map_location=torch.device("cpu"))
-    model.load_state_dict(checkpoint["model"])
-    model = model.to(device)
-    model.eval()
-
-    prompt = template(region)
-
-    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(" HYPERPARAMETERS ".center(60, "#"), "\n")
-    args = vars(args)
-
-    for key in args.keys():
-        print(f"{key}: {str(args[key])}")
-
-    print("\n", " OUTPUT TUNES ".center(60, "#"))
-
-    start_time = time.time()
-
-    for i in range(num_tunes):
-        title_artist = f"T:{region} Fragment\nC: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")
-
-    print("Generation time: {:.2f} seconds".format(time.time() - start_time))
-    os.makedirs("./tmp", exist_ok=True)
-    timestamp = time.strftime("%a_%d_%b_%Y_%H_%M_%S", time.localtime())
-    out_midi = abc_to_midi(tunes, f"./tmp/[{region}]{timestamp}.mid")
-    out_xml = abc_to_musicxml(tunes, f"./tmp/[{region}]{timestamp}.musicxml")
-    out_mxl = musicxml_to_mxl(f"./tmp/[{region}]{timestamp}.musicxml")
-    pdf_file, jpg_file = mxl2jpg(out_mxl)
-    wav_file = midi2wav(out_midi)
-
-    return tunes, out_midi, pdf_file, out_xml, out_mxl, jpg_file, wav_file
-
-
-def inference(region):
-    if os.path.exists("./tmp"):
-        shutil.rmtree("./tmp")
-
-    parser = argparse.ArgumentParser()
-    args = get_args(parser)
-    return generate_abc(args, region)
-
-
-with gr.Blocks() as demo:
-    with gr.Row():
-        with gr.Column():
-            region_opt = gr.Dropdown(
-                choices=["Mondstadt", "Liyue", "Inazuma", "Sumeru", "Fontaine"],
-                value="Mondstadt",
-                label="Region genre",
-            )
-            gen_btn = gr.Button("Generate")
-
-        with gr.Column():
-            wav_output = gr.Audio(label="Audio", type="filepath")
-            dld_midi = gr.components.File(label="Download MIDI")
-            pdf_score = gr.components.File(label="Download PDF score")
-            dld_xml = gr.components.File(label="Download MusicXML")
-            dld_mxl = gr.components.File(label="Download MXL")
-            abc_output = gr.Textbox(label="abc score", show_copy_button=True)
-            img_score = gr.Image(label="Staff", type="filepath")
-
-    gen_btn.click(
-        inference,
-        inputs=region_opt,
-        outputs=[
-            abc_output,
-            dld_midi,
-            pdf_score,
-            dld_xml,
-            dld_mxl,
-            img_score,
-            wav_output,
-        ],
-    )
-
-demo.launch(share=True)
+import re
+import os
+import time
+import torch
+import shutil
+import argparse
+import gradio as gr
+from config import *
+from utils import Patchilizer, TunesFormer, DEVICE
+from convert import abc_to_midi, abc_to_musicxml, musicxml_to_mxl, mxl2jpg, midi2wav
+from modelscope import snapshot_download
+from transformers import GPT2Config
+import warnings
+
+warnings.filterwarnings("ignore")
+
+# 模型下载
+MODEL_DIR = snapshot_download("MuGeminorum/hoyoGPT")
+
+
+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",
+    )
+    args = parser.parse_args()
+
+    return args
+
+
+def generate_abc(args, epochs: str, region: str):
+    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)
+    filename = f"{MODEL_DIR}/{epochs}/weights.pth"
+    checkpoint = torch.load(filename, map_location=torch.device("cpu"))
+    model.load_state_dict(checkpoint["model"])
+    model = model.to(DEVICE)
+    model.eval()
+
+    prompt = f"A:{region}\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(" HYPERPARAMETERS ".center(60, "#"), "\n")
+    arg_dict: dict = vars(args)
+
+    for key in arg_dict.keys():
+        print(f"{key}: {str(arg_dict[key])}")
+
+    print("\n", " OUTPUT TUNES ".center(60, "#"))
+
+    start_time = time.time()
+    for i in range(num_tunes):
+        title_artist = f"T:{region} Fragment\nC: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")
+
+    print("Generation time: {:.2f} seconds".format(time.time() - start_time))
+    os.makedirs(TEMP_DIR, exist_ok=True)
+    timestamp = time.strftime("%a_%d_%b_%Y_%H_%M_%S", time.localtime())
+    try:
+        out_midi = abc_to_midi(tunes, f"{TEMP_DIR}/[{region}]{timestamp}.mid")
+        out_xml = abc_to_musicxml(tunes, f"{TEMP_DIR}/[{region}]{timestamp}.musicxml")
+        out_mxl = musicxml_to_mxl(f"{TEMP_DIR}/[{region}]{timestamp}.musicxml")
+        pdf_file, jpg_file = mxl2jpg(out_mxl)
+        wav_file = midi2wav(out_midi)
+
+        return tunes, out_midi, pdf_file, out_xml, out_mxl, jpg_file, wav_file
+
+    except Exception as e:
+        print(f"Invalid abc generated: {e}, retrying...")
+        return generate_abc(args, epochs, region)
+
+
+def inference(epochs: str, region: str):
+    Teyvat = {
+        "蒙德": "Mondstadt",
+        "璃月": "Liyue",
+        "稻妻": "Inazuma",
+        "须弥": "Sumeru",
+        "枫丹": "Fontaine",
+    }
+
+    if os.path.exists(TEMP_DIR):
+        shutil.rmtree(TEMP_DIR)
+
+    parser = argparse.ArgumentParser()
+    args = get_args(parser)
+    return generate_abc(args, epochs, Teyvat[region])
+
+
+if __name__ == "__main__":
+    with gr.Blocks() as demo:
+        gr.Markdown(
+            """<center>欢迎使用此创空间，此创空间由bilibili <a href="https://space.bilibili.com/30620472">@亦真亦幻Studio</a> 基于 Tunesformer 开源项目制作，完全免费。</center>"""
+        )
+        with gr.Row():
+            with gr.Column():
+                weight_opt = gr.Dropdown(
+                    choices=["5", "15"],
+                    value="15",
+                    label="模型选择(epochs)",
+                )
+                region_opt = gr.Dropdown(
+                    choices=["蒙德", "璃月", "稻妻", "须弥", "枫丹"],
+                    value="蒙德",
+                    label="地区风格",
+                )
+                gen_btn = gr.Button("生成")
+                gr.Markdown(
+                    """
+    <center>
+    当前模型还在调试中，由于训练数据由 MIDI 转换而来，存在大量谱面不规范问题，导致很多生成结果有谱面不规范等问题。<br>
+
+    计划在原神主线杀青后，所有国家地区角色全部开放后，二创音乐会齐全且样本均衡，届时重新微调模型并添加现实风格筛选辅助游戏各国家输出把关，以提升输出区分度与质量。
+
+    数据来源：<a href="https://musescore.org">MuseScore</a><br>
+    Tag 嵌入数据来源：<a href="https://genshin-impact.fandom.com/wiki/Genshin_Impact_Wiki">Genshin Impact Wiki | Fandom</a><br>
+    模型基础：<a href="https://github.com/sander-wood/tunesformer">Tunesformer</a>
+
+    注：崩铁方面数据工程正在运作中，未来也希望随主线杀青而基线化。</center>"""
+                )
+
+            with gr.Column():
+                wav_output = gr.Audio(label="音频", type="filepath")
+                dld_midi = gr.components.File(label="下载 MIDI")
+                pdf_score = gr.components.File(label="下载 PDF 乐谱")
+                dld_xml = gr.components.File(label="下载 MusicXML")
+                dld_mxl = gr.components.File(label="下载 MXL")
+                abc_output = gr.Textbox(label="abc notation", show_copy_button=True)
+                img_score = gr.Image(label="五线谱", type="filepath")
+
+        gen_btn.click(
+            inference,
+            inputs=[weight_opt, region_opt],
+            outputs=[
+                abc_output,
+                dld_midi,
+                pdf_score,
+                dld_xml,
+                dld_mxl,
+                img_score,
+                wav_output,
+            ],
+        )
+
+    demo.launch()

config.py

@@ -1,19 +1,9 @@
-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
-
-# Number of epochs to train for (if early stopping doesn't intervene)
-NUM_EPOCHS = 5  # 32
-LEARNING_RATE = 5e-5            # Learning rate for the optimizer
-# Batch size for patch during training, 0 for full context
-PATCH_SAMPLING_BATCH_SIZE = 0
-LOAD_FROM_CHECKPOINT = True     # Whether to load weights from a checkpoint
-# Whether to share weights between the encoder and decoder
-SHARE_WEIGHTS = False
-WEIGHT_URL = 'https://huggingface.co/MuGeminorum/hoyoGPT/resolve/main/weights.pth'
-ZH_WEIGHT_URL = 'https://www.modelscope.cn/api/v1/models/MuGeminorum/hoyoGPT/repo?Revision=master&FilePath=weights.pth'
-WEIGHT_PATH = 'weights.pth'
-LOG_PATH = 'logs.txt'
-PROMPT_PATH = 'prompt.txt'
+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
+# Batch size for patch during training, 0 for full context
+PATCH_SAMPLING_BATCH_SIZE = 0
+# Whether to share weights between the encoder and decoder
+SHARE_WEIGHTS = False
+TEMP_DIR = "./tmp"

convert.py

@@ -1,93 +1,88 @@
-import os
-import sys
-import fitz
-import subprocess
-from PIL import Image
-from music21 import converter
-from utils import download
-
-if sys.platform.startswith('linux'):
-    apkname = 'MuseScore.AppImage'
-    extra_dir = 'squashfs-root'
-    download(
-        filename=apkname,
-        url='https://cdn.jsdelivr.net/musescore/v4.2.0/MuseScore-4.2.0.233521125-x86_64.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 = "D:/Program Files/MuseScore 3/bin/MuseScore3.exe"
-
-
-def abc_to_midi(abc_content, output_midi_path):
-    score = converter.parse(abc_content, format='abc')
-    score.write('midi', fp=output_midi_path)
-    return output_midi_path
-
-
-def abc_to_musicxml(abc_content, output_xml_path):
-    score = converter.parse(abc_content, format='abc')
-    score.write('musicxml', fp=output_xml_path)
-    return output_xml_path
-
-
-def musicxml_to_mxl(xml_path):
-    mxl_file = xml_path.replace('.musicxml', '.mxl')
-    command = [mscore, "-o", mxl_file, xml_path]
-    result = subprocess.run(command)
-    print(result)
-    return mxl_file
-
-
-def midi2wav(mid_file: str):
-    wav_file = mid_file.replace('.mid', '.wav')
-    command = [mscore, "-o", wav_file, mid_file]
-    result = subprocess.run(command)
-    print(result)
-    return wav_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 mxl2jpg(mxl_file: str):
-    pdf_score = mxl_file.replace('.mxl', '.pdf')
-    command = [mscore, "-o", pdf_score, mxl_file]
-    result = subprocess.run(command)
-    print(result)
-    return pdf_score, pdf2img(pdf_score)
+import os
+import sys
+import fitz
+import subprocess
+from PIL import Image
+from music21 import converter
+from utils import download
+
+if sys.platform.startswith("linux"):
+    apkname = "MuseScore.AppImage"
+    extra_dir = "squashfs-root"
+    download(
+        filename=apkname,
+        url="https://master.dl.sourceforge.net/project/musescore.mirror/v4.2.0/MuseScore-4.2.0.233521125-x86_64.AppImage?viasf=1",
+    )
+    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 = "D:/Program Files/MuseScore 3/bin/MuseScore3.exe"
+
+
+def abc_to_midi(abc_content, output_midi_path):
+    score = converter.parse(abc_content, format="abc")
+    score.write("midi", fp=output_midi_path)
+    return output_midi_path
+
+
+def abc_to_musicxml(abc_content, output_xml_path):
+    score = converter.parse(abc_content, format="abc")
+    score.write("musicxml", fp=output_xml_path)
+    return output_xml_path
+
+
+def musicxml_to_mxl(xml_path):
+    mxl_file = xml_path.replace(".musicxml", ".mxl")
+    command = [MSCORE, "-o", mxl_file, xml_path]
+    result = subprocess.run(command)
+    print(result)
+    return mxl_file
+
+
+def midi2wav(mid_file: str):
+    wav_file = mid_file.replace(".mid", ".wav")
+    command = [MSCORE, "-o", wav_file, mid_file]
+    result = subprocess.run(command)
+    print(result)
+    return wav_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 mxl2jpg(mxl_file: str):
+    pdf_score = mxl_file.replace(".mxl", ".pdf")
+    command = [MSCORE, "-o", pdf_score, mxl_file]
+    result = subprocess.run(command)
+    print(result)
+    return pdf_score, pdf2img(pdf_score)

requirements.txt

@@ -1,9 +1,9 @@
-transformers==4.18.0
-samplings==0.1.7
-unidecode
-music21
-autopep8
-pillow==9.4.0
-gradio
-pymupdf
+transformers==4.18.0
+samplings==0.1.7
+unidecode
+music21
+autopep8
+pillow==9.4.0
+gradio==4.8.0
+pymupdf
 torch

utils.py

@@ -1,394 +1,377 @@
-import os
-import re
-import torch
-import random
-from config import *
-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
-
-if torch.cuda.is_available():
-    device = torch.device("cuda")
-else:
-    device = torch.device("cpu")
-
-
-def template(region):
-    return f'''A:{region}
-S:2
-B:9
-E:4
-B:9
-L:1/8
-M:3/4
-K:D
-de |"D"'''
-
-
-def download(filename=WEIGHT_PATH, url=WEIGHT_URL):
-    import time
-    import requests
-
-    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 weights to '{filename}'...",
-            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(3)
-        download(filename, ZH_WEIGHT_URL)
-
-
-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]
+import re
+import time
+import torch
+import random
+import requests
+from config import *
+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
+
+DEVICE = torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu")
+
+
+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}, retrying...")
+        time.sleep(10)
+        download(filename, url)
+
+
+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]