initial deploy

.gitignore

@@ -0,0 +1,5 @@
+.venv/
+__pycache__/
+*.pyc
+.DS_Store
+PiCoGenv2-repo/

README.md

@@ -1,11 +1,11 @@
 ---
 title: PiCoGen
-emoji: 🏆
+emoji: 🎹
 colorFrom: pink
 colorTo: red
 sdk: gradio
-sdk_version: 6.20.0
-python_version: '3.12'
+sdk_version: 5.28.0
+python_version: '3.10'
 app_file: app.py
 pinned: false
 license: mit

app.py

@@ -0,0 +1,167 @@
+import sys
+
+sys.stdout.reconfigure(line_buffering=True)
+
+try:
+    import spaces
+except ImportError:
+    # keep @spaces.GPU usable as a no-op; ZeroGPU requires this exact name.
+    class spaces:
+        class GPU:
+            def __init__(self, func=None, duration=60):
+                self.func = func
+
+            def __call__(self, *args, **kwargs):
+                if self.func is not None:
+                    return self.func(*args, **kwargs)
+                func = args[0]
+                return func
+
+import tempfile
+import threading
+
+import gradio as gr
+import soundfile as sf
+import torch
+from pyharp import ModelCard, build_endpoint
+
+import picogen2
+from picogen2.mirtoolkit.beat_this import BeatThis
+from picogen2.mirtoolkit.sheetsage import SheetSage
+
+# SheetSage (this model's audio feature extractor) was trained on ~24s segments and is
+# most accurate on short clips; longer songs also risk exceeding the GPU time budget below.
+MAX_INPUT_SECONDS = 30.0
+
+DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
+
+decoder = None
+decoder_ready = False  # has the decoder been moved onto the GPU yet?
+tokenizer = None
+beat_detector = None
+sheetsage_model = None
+
+model_loading = True
+model_error = None
+
+
+def load_assets():
+    """Downloads PiCoGen2's checkpoint, pre-fetches SheetSage's (large) asset cache, and
+    builds the beat tracker, all on CPU so the app can start serving while this runs."""
+    global decoder, tokenizer, beat_detector, model_loading, model_error
+    try:
+        tokenizer = picogen2.Tokenizer()
+        decoder = picogen2.PiCoGenDecoder.from_pretrained(device="cpu")
+
+        try:
+            import sheetsage.assets as sheetsage_assets
+
+            for tag in sheetsage_assets.get_asset_tags():
+                sheetsage_assets.retrieve_asset(tag)
+        except Exception as e:
+            # SheetSage lazily re-downloads whatever's missing on first use, so this is
+            # a warm-up, not a hard requirement.
+            print(f"SheetSage asset pre-fetch incomplete, will retry on first request: {e}")
+
+        beat_detector = BeatThis(cuda=False)
+        print("Models loaded (CPU).")
+    except Exception as e:
+        model_error = str(e)
+        print(f"Load error: {e}")
+    finally:
+        model_loading = False
+
+
+threading.Thread(target=load_assets, daemon=True).start()
+
+
+model_card = ModelCard(
+    name="PiCoGen2",
+    description=(
+        "Generates a piano cover from a short pop song clip: SheetSage extracts melody/"
+        "harmony audio features, and a GPT-NeoX decoder turns them into piano notes."
+    ),
+    author="Chih-Pin Tan, Hsin Ai, Yi-Hsin Chang, Shuen-Huei Guan, Yi-Hsuan Yang",
+    tags=["music generation", "piano cover", "midi"],
+)
+
+
+@spaces.GPU(duration=300)
+@torch.inference_mode()
+def process_fn(input_audio_path: str, temperature: float) -> str:
+    """Detects beats, extracts SheetSage audio features, and generates a piano cover."""
+    global decoder, sheetsage_model, decoder_ready
+
+    if model_loading:
+        raise gr.Error("Model is still loading, please wait a moment and try again.")
+    if decoder is None:
+        raise gr.Error(f"Model failed to load: {model_error}")
+    if DEVICE != "cuda":
+        raise gr.Error("This model requires a GPU; none is available.")
+
+    duration = sf.info(input_audio_path).duration
+    if duration > MAX_INPUT_SECONDS:
+        raise gr.Error(
+            f"Input is {duration:.1f}s long; please trim it to {MAX_INPUT_SECONDS:.0f}s "
+            "or shorter."
+        )
+
+    if not decoder_ready:
+        decoder = decoder.to(DEVICE)  # only safe here, inside @spaces.GPU
+        decoder_ready = True
+    if sheetsage_model is None:
+        sheetsage_model = SheetSage()  # constructs Jukebox on the GPU internally
+
+    beats, downbeats = beat_detector(input_audio_path)
+    beat_information = {"beats": beats.tolist(), "downbeats": downbeats.tolist()}
+
+    sheetsage_output = sheetsage_model(
+        audio_path=input_audio_path, beat_information=beat_information
+    )
+
+    out_events = picogen2.decode(
+        model=decoder,
+        tokenizer=tokenizer,
+        beat_information=beat_information,
+        melody_last_embs=sheetsage_output["melody_last_hidden_state"],
+        harmony_last_embs=sheetsage_output["harmony_last_hidden_state"],
+        temperature=temperature,
+        device=DEVICE,
+    )
+
+    with tempfile.NamedTemporaryFile(suffix=".mid", delete=False) as f:
+        output_midi_path = f.name
+    tokenizer.events_to_midi(out_events).dump(output_midi_path)
+    return output_midi_path
+
+
+with gr.Blocks() as demo:
+    input_components = [
+        gr.Audio(
+            type="filepath",
+            label="Input Audio",
+            info=f"Short song clip, up to {MAX_INPUT_SECONDS:.0f}s",
+        ).harp_required(True),
+        gr.Slider(
+            minimum=0.1,
+            maximum=2.0,
+            step=0.1,
+            value=1.0,
+            label="Temperature",
+            info="Sampling temperature for the piano decoder (default: 1.0, per repo config)",
+        ),
+    ]
+    output_components = [
+        gr.File(type="filepath", label="Piano Cover", file_types=[".mid", ".midi"]).set_info(
+            "Generated piano cover, as a MIDI file."
+        ),
+    ]
+
+    build_endpoint(
+        model_card=model_card,
+        input_components=input_components,
+        output_components=output_components,
+        process_fn=process_fn,
+    )
+
+demo.queue().launch(pwa=True)

packages.txt

@@ -0,0 +1,3 @@
+ffmpeg
+libopenmpi-dev
+wget

picogen2/__init__.py

@@ -0,0 +1,6 @@
+from .infer import decode
+from .model import PiCoGenDecoder
+from .repr import Tokenizer
+from .version import VERSION, VERSION_SHORT
+
+__all__ = ["decode", "PiCoGenDecoder", "Tokenizer", "VERSION", "VERSION_SHORT"]

picogen2/assets.py

@@ -0,0 +1,85 @@
+import shutil
+import subprocess
+from pathlib import Path
+
+from .utils import logger
+
+CACHE_DIR = Path.home() / ".cache" / "picogen2"
+
+URL_MODEL = "https://zenodo.org/records/13380452/files/model_ft_00070000?download=1"
+URL_VOCAB = "https://raw.githubusercontent.com/tanchihpin0517/PiCoGen/v2/assets/vocab.json"
+URL_CONFIG = "https://raw.githubusercontent.com/tanchihpin0517/PiCoGen/v2/assets/config.json"
+URL_TEST_SONG = "https://www.dropbox.com/scl/fi/zj68yghtn0cwtwnqj7vrx/pop.00000.wav?rlkey=bejuh89wehbc8psl9ujmqa73u&st=kb265uvz&dl=0"
+
+
+def default_cache_dir_decorator(func):
+    def wrapper(*args, **kwargs):
+        CACHE_DIR.mkdir(parents=True, exist_ok=True)
+        return func(*args, **kwargs)
+
+    return wrapper
+
+
+@default_cache_dir_decorator
+def checkpoint_file():
+    default_ckpt_file = CACHE_DIR / "model_ft_00070000"
+    if not default_ckpt_file.exists():
+        logger.warning("Download default model from {}".format(URL_MODEL))
+        logger.warning("Save to {}".format(default_ckpt_file))
+
+        _download(URL_MODEL, default_ckpt_file)
+
+    return default_ckpt_file
+
+
+@default_cache_dir_decorator
+def vocab_file():
+    default_vocab_file = CACHE_DIR / "vocab.json"
+    if not default_vocab_file.exists():
+        logger.warning("Download default vocab from {}".format(URL_VOCAB))
+        logger.warning("Save to {}".format(default_vocab_file))
+
+        _download(URL_VOCAB, default_vocab_file)
+
+    return default_vocab_file
+
+
+@default_cache_dir_decorator
+def config_file():
+    default_config_file = CACHE_DIR / "config.json"
+    if not default_config_file.exists():
+        logger.warning("Download default config from {}".format(URL_CONFIG))
+        logger.warning("Save to {}".format(default_config_file))
+
+        _download(URL_CONFIG, default_config_file)
+
+    return default_config_file
+
+
+@default_cache_dir_decorator
+def test_song():
+    default_test_song = CACHE_DIR / "pop.00000.wav"
+    if not default_test_song.exists():
+        logger.warning("Download default test song from {}".format(URL_TEST_SONG))
+        logger.warning("Save to {}".format(default_test_song))
+
+        _download(URL_TEST_SONG, default_test_song)
+
+    return default_test_song
+
+
+def _download(url, output_file_path, verbose=True):
+    if verbose:
+        logger.info(f"Downloading {url} to {output_file_path}")
+
+    if shutil.which("wget") is None:
+        logger.error("wget is not installed. Please install wget to download the model.")
+        raise FileNotFoundError("`wget` is not installed")
+
+    try:
+        subprocess.run(["wget", url, "-O", str(output_file_path)], check=True)
+    except subprocess.CalledProcessError as e:
+        logger.error(f"Failed to download file from {url}: {e}")
+        if output_file_path.exists():
+            output_file_path.unlink()
+        raise e

picogen2/infer.py

@@ -0,0 +1,112 @@
+import numpy as np
+import torch
+from tqdm import tqdm
+
+from .model import PiCoGenDecoder
+from .repr import Event
+from .utils import downbeat_time_to_index
+
+
+@torch.no_grad()
+def decode(
+    model,
+    tokenizer,
+    beat_information,
+    melody_last_embs,
+    harmony_last_embs,
+    max_bar_num=None,
+    max_token_num=None,
+    temperature=1.0,
+    device=None,
+):
+    if device is None:
+        device = model.parameters().__next__().device
+
+    starting_bpm = 60 / np.diff(np.array(beat_information["beats"])[:2]).mean()
+
+    TARGET = PiCoGenDecoder.InputClass.TARGET.value
+    CONDITION = PiCoGenDecoder.InputClass.CONDITION.value
+
+    out_events = [
+        Event(etype="spec", value="spec_ss"),
+        tokenizer.get_tempo_event(starting_bpm),
+    ]
+
+    input_seg = [tokenizer.e2i(Event(etype="spec", value="spec_bos"))]
+    input_seg.extend([tokenizer.e2i(e) for e in out_events])  # song start, global tempo
+    need_encode_seg = [0] * len(input_seg)
+    input_cls_seg = [TARGET] * len(input_seg)
+
+    end_event = Event(etype="spec", value="spec_se")
+    bar_start_event = Event(etype="bar", value="bar_start")
+    bar_end_event = Event(etype="bar", value="bar_end")
+
+    total_beats = len(beat_information["beats"])
+    downbeats = downbeat_time_to_index(beat_information["beats"], beat_information["downbeats"])
+    if downbeats[-1] < total_beats:
+        downbeats.append(total_beats - 1)
+    if max_bar_num is not None:
+        downbeats = downbeats[: max_bar_num + 1]
+
+    pbar = tqdm(total=len(downbeats) - 1)
+    for bar_i, b in enumerate(range(len(downbeats) - 1)):
+        last_past_kv = None
+
+        # NOTE: upbeat is handled by SheetSage
+        downbeat_start, downbeat_end = downbeats[b], downbeats[b + 1]
+        # downbeat_start, downbeat_end = downbeats[b]-downbeats[0], downbeats[b+1]-downbeats[0]
+        for j in range(downbeat_start * tokenizer.beat_div, downbeat_end * tokenizer.beat_div):
+            input_seg.append((melody_last_embs[j], harmony_last_embs[j]))
+            input_cls_seg.append(CONDITION)
+            need_encode_seg.append(1)
+        if b == len(downbeats) - 2:  # NOTE: add song_end to the last bar condition
+            input_seg.append(tokenizer.e2i(Event(etype="spec", value="spec_se")))
+            input_cls_seg.append(CONDITION)
+            need_encode_seg.append(0)
+
+        input_seg.append(tokenizer.e2i(bar_start_event))
+        need_encode_seg.append(0)
+        input_cls_seg.append(TARGET)
+        out_events.append(bar_start_event)
+
+        while True:  # generate one bar
+            if len(input_seg) > model.hp.max_seq_len:
+                input_seg = input_seg[-model.hp.max_seq_len // 2 :]
+                input_seg[0] = tokenizer.e2i(Event(etype="spec", value="spec_bos"))
+                need_encode_seg = need_encode_seg[-model.hp.max_seq_len // 2 :]
+                need_encode_seg[0] = 0
+                input_cls_seg = input_cls_seg[-model.hp.max_seq_len // 2 :]
+                input_cls_seg[0] = TARGET
+                last_past_kv = None
+
+            input_cls_ids = torch.LongTensor(input_cls_seg)[None, :].to(device)
+            need_encode = torch.BoolTensor(need_encode_seg)[None, :].to(device)
+
+            output_ids, past_kv = model.generate(
+                input_seg=[input_seg],
+                input_cls_ids=input_cls_ids,
+                need_encode=need_encode,
+                kv_cache=last_past_kv,
+                temperature=temperature,
+            )
+            out_id = output_ids[0][-1].item()
+            out_event = tokenizer.i2e(out_id)
+
+            out_events.append(out_event)
+            input_seg.append(out_id)
+            input_cls_seg.append(TARGET)
+            need_encode_seg.append(0)
+
+            if out_event in (bar_end_event, end_event):
+                break
+            last_past_kv = past_kv
+
+            pbar.set_description(f"length: {len(out_events)}({len(input_seg)})")
+
+        pbar.update(1)
+
+        if max_token_num is not None and len(input_seg) > max_token_num:
+            break
+
+    pbar.close()
+    return out_events

picogen2/mirtoolkit/beat_this.py

@@ -0,0 +1,78 @@
+"""
+CPJKU's beat_this
+- reference: https://github.com/CPJKU/beat_this
+"""
+
+import sys
+from pathlib import Path
+
+import torch
+
+from .utils import load_audio
+
+if not sys.version_info < (3, 10):
+    from beat_this.inference import Audio2Beats
+
+
+class BeatThis:
+    def __init__(self, cuda=None, dbn=True):
+        if cuda is None:
+            cuda = torch.cuda.is_available()
+
+        self.audio2beats = Audio2Beats(
+            device="cuda" if cuda else "cpu",
+            dbn=dbn,
+        )
+        self.use_dbn = dbn
+
+    @torch.no_grad()
+    def __call__(
+        self,
+        file_or_array,
+        sr=None,
+        beats_per_bar=[3, 4],
+        min_bpm=55.0,
+        max_bpm=215.0,
+        fps=50,
+        transition_lambda=100,
+    ):
+        """
+        Function for extracting beat and downbeat positions (in seconds) from a file or a data array.
+
+        Args:
+            file_or_array (str or Path or ndarray): Path to the audio file or numpy array containing the audio data.
+            sr (int, optional): Sample rate of the audio file. Required if `file_or_array` is a numpy array. Defaults to None.
+            if dbn is True:
+                beats_per_bar (list, optional): List of possible beats per bar. Defaults to [3, 4].
+                min_bpm (float, optional): Minimum tempo in BPM. Defaults to 55.0.
+                max_bpm (float, optional): Maximum tempo in BPM. Defaults to 215.0.
+                fps (int, optional): Frames per second. Defaults to 50.
+
+        Returns:
+            beats (ndarray): Array of beat positions in seconds.
+            downbeats (ndarray): Array of downbeat positions in seconds.
+        """
+        if sys.version_info < (3, 10):
+            raise ImportError("Python 3.10 or higher is required to use this function.")
+
+        if isinstance(file_or_array, (str, Path)):
+            audio, sr = load_audio(file_or_array, dtype="float64")
+        else:
+            audio = file_or_array
+
+        if self.use_dbn:
+            from madmom.features.downbeats import DBNDownBeatTrackingProcessor
+
+            # WARN: This is a hacky way to set the DBN parameters
+            dbn = DBNDownBeatTrackingProcessor(
+                beats_per_bar=beats_per_bar,
+                min_bpm=min_bpm,
+                max_bpm=max_bpm,
+                fps=fps,
+                transition_lambda=transition_lambda,
+            )
+            self.audio2beats.frames2beats.dbn = dbn
+
+        beats, downbeats = self.audio2beats(audio, sr)
+
+        return beats, downbeats

picogen2/mirtoolkit/sheetsage.py

@@ -0,0 +1,126 @@
+"""
+SheetSage: lead sheet transcription model
+- reference: https://github.com/chrisdonahue/sheetsage
+"""
+
+import logging
+import shutil
+import subprocess
+import tempfile
+from pathlib import Path
+from typing import Union
+
+from sheetsage.infer import sheetsage as sheetsage_infer
+from tqdm import tqdm as tqdm_fn
+
+
+class SheetSage:
+    def __init__(self):
+        pass
+
+    def __call__(
+        self,
+        audio_path: Union[str, Path] = None,
+        audio_url: str = None,
+        segment_start_hint=None,
+        segment_end_hint=None,
+        use_jukebox=True,
+        measures_per_chunk=8,
+        dynamic_chunking=True,
+        segment_hints_are_downbeats=False,
+        beat_information=None,
+        beats_per_measure_hint=None,
+        beats_per_minute_hint=None,
+        detect_melody=True,
+        detect_harmony=True,
+        melody_threshold=None,
+        harmony_threshold=None,
+        beat_detection_padding=15.0,
+        avoid_chunking_if_possible=True,
+        legacy_behavior=False,
+        status_change_callback=lambda s: logging.info(s.name),
+        return_intermediaries=False,
+        tqdm=tqdm_fn,
+        return_dict=True,
+    ):
+        # assert audio_path or audio_url should be provided but not both
+        assert (audio_path and not audio_url) or (audio_url and not audio_path), (
+            f"One of audio_path or audio_url should be provided but not both: {audio_path}, {audio_url}"
+        )
+
+        assert shutil.which("ffmpeg") is not None, "ffmpeg not found. Please install ffmpeg."
+
+        if audio_path:  # if audio_path is provided
+            ext = "flac"
+            tmp_audio_file = tempfile.NamedTemporaryFile(suffix=f".{ext}")
+            subprocess.run(
+                [
+                    "ffmpeg",
+                    "-i",
+                    str(audio_path),
+                    "-vn",
+                    "-f",
+                    ext,
+                    "-y",
+                    tmp_audio_file.name,
+                ]
+            )
+            audio_path = Path(tmp_audio_file.name)
+
+            assert audio_path.exists(), f"File not found: {audio_path}"
+        else:  # if audio_url is provided
+            ext = "flac"
+            tmp_dir = tempfile.TemporaryDirectory()
+            tmp_audio_file = tmp_dir.name + f"/audio.{ext}"
+            subprocess.run(
+                [
+                    "yt-dlp",
+                    "-x",
+                    "--audio-format",
+                    ext,
+                    "--audio-quality",
+                    "0",
+                    "-o",
+                    tmp_dir.name + "/audio.%(ext)s",
+                    audio_url,
+                ],
+                check=True,
+            )
+            audio_path = Path(tmp_audio_file)
+
+        sheetsage_output = sheetsage_infer(
+            audio_path_bytes_or_url=audio_path,
+            segment_start_hint=segment_start_hint,
+            segment_end_hint=segment_end_hint,
+            use_jukebox=use_jukebox,
+            measures_per_chunk=measures_per_chunk,
+            dynamic_chunking=dynamic_chunking,
+            segment_hints_are_downbeats=segment_hints_are_downbeats,
+            beat_information=beat_information,
+            beats_per_measure_hint=beats_per_measure_hint,
+            beats_per_minute_hint=beats_per_minute_hint,
+            detect_melody=detect_melody,
+            detect_harmony=detect_harmony,
+            melody_threshold=melody_threshold,
+            harmony_threshold=harmony_threshold,
+            beat_detection_padding=beat_detection_padding,
+            avoid_chunking_if_possible=avoid_chunking_if_possible,
+            legacy_behavior=legacy_behavior,
+            status_change_callback=status_change_callback,
+            return_intermediaries=return_intermediaries,
+            tqdm=tqdm,
+        )
+
+        if return_dict:
+            return sheetsage_output
+        else:
+            return (
+                sheetsage_output["lead_sheet"],
+                sheetsage_output["segment_beats"],
+                sheetsage_output["segment_beats_times"],
+                sheetsage_output["chunks_tertiaries"],
+                sheetsage_output["melody_logits"],
+                sheetsage_output["harmony_logits"],
+                sheetsage_output["melody_last_hidden_state"],
+                sheetsage_output["harmony_last_hidden_state"],
+            )

picogen2/mirtoolkit/utils.py

@@ -0,0 +1,80 @@
+import shutil
+import subprocess
+import tempfile
+from pathlib import Path
+
+import torch
+import torchaudio
+
+
+def download(url, file):
+    assert isinstance(url, str)
+    assert isinstance(file, (str, Path))
+    if isinstance(file, str):
+        file = Path(file)
+
+    if shutil.which("wget") is None:
+        raise FileNotFoundError("wget not found. Please install wget.")
+    # Download the file using wget
+    tmp_dir = tempfile.TemporaryDirectory()
+    tmp_file = tmp_dir.name + "/" + file.name
+    subprocess.run(["wget", "-O", tmp_file, url], check=True)
+    shutil.move(tmp_file, file)
+
+    # try:
+    #     # Send a GET request to the URL
+    #     response = requests.get(url, stream=True)
+    #     response.raise_for_status()  # Raise an error for bad status codes
+
+    #     tmp_dir = tempfile.TemporaryDirectory()
+    #     tmp_file = tmp_dir.name + "/" + file.name
+    #     with open(tmp_file, "wb") as f:
+    #         for chunk in response.iter_content(chunk_size=8192):
+    #             f.write(chunk)
+    #     shutil.move(tmp_file, file)
+    # except requests.RequestException as e:
+    #     print(f"An error occurred: {e}")
+
+
+def load_audio(
+    path, dtype="float32", return_tensor=False, channels_first=False, mono=False, sr=None
+):
+    """
+    Load an audio file from the given path.
+    Args:
+        path (str): The path to the audio file.
+        dtype (str, optional): The desired data type of the audio waveform. Defaults to "float64".
+        return_tensor (bool, optional): Whether to return the audio waveform as a tensor. Defaults to False.
+        channels_first (bool, optional): Whether to return the audio waveform with channels as the first dimension. Defaults to False.
+    Returns:
+        tuple or ndarray: If `return_tensor` is False, returns a tuple containing the audio waveform as a numpy ndarray and the sample rate as an integer. If `return_tensor` is True, returns a tuple containing the audio waveform as a PyTorch tensor and the sample rate as an integer.
+    """
+    assert dtype in ["float32", "float64"]
+
+    try:
+        waveform, sample_rate = torchaudio.load(path, channels_first=channels_first)
+        if mono and waveform.shape[0] > 1:
+            if channels_first:
+                waveform = waveform.mean(0)
+            else:
+                waveform = waveform.mean(-1)
+
+        if dtype == "float64":
+            waveform = waveform.to(dtype=torch.float64)
+
+    except Exception:
+        # in case torchaudio fails, try soundfile
+        import soundfile as sf
+
+        waveform, sample_rate = sf.read(path, dtype=dtype)
+        waveform = torch.from_numpy(waveform)
+
+    if sr is not None:
+        # resample the audio to the given sample rate
+        waveform = torchaudio.transforms.Resample(sample_rate, sr)(waveform)
+        sample_rate = sr
+
+    if not return_tensor:
+        waveform = waveform.squeeze().numpy()
+
+    return waveform, sample_rate

picogen2/model.py

@@ -0,0 +1,230 @@
+from enum import Enum
+
+import numpy as np
+import torch
+import torch.nn.functional as F
+from torch import nn
+from transformers import GPTNeoXConfig, GPTNeoXModel
+
+from . import assets
+from .utils import load_checkpoint, load_config, top_p
+
+
+def _get_device(module):
+    return next(module.parameters()).device
+
+
+class ConditionEncoder(nn.Module):
+    def __init__(self, hp):
+        super().__init__()
+        self.l1_encoder = nn.TransformerEncoder(
+            nn.TransformerEncoderLayer(
+                d_model=hp.d_model,
+                nhead=hp.num_heads,
+                dim_feedforward=hp.d_model * 4,
+                dropout=hp.dropout,
+                activation=hp.activation,
+                batch_first=True,
+            ),
+            hp.num_layers_encoder,
+        )
+        self.pos_emb = nn.Embedding(hp.condition_class, hp.d_model)
+        self.bottlenect = nn.Sequential(
+            nn.Linear(hp.d_model, hp.d_bottleneck),
+            nn.ReLU(),
+            nn.Linear(hp.d_bottleneck, hp.d_model),
+        )
+
+    def forward(self, input_embs):
+        B, L, N, D = input_embs.shape
+        pos = torch.arange(N).to(input_embs.device)
+        pos = self.pos_emb(pos)[None, None, :, :].expand(B, L, N, D)
+        input_embs = input_embs + pos
+        out = self.l1_encoder(input_embs.view(B * L, N, D)).view(B, L, N, D)
+        out = out[:, :, 0, :]
+        assert out.shape == (B, L, D)
+        out = self.bottlenect(out)
+        return out
+
+
+class PiCoGenDecoder(nn.Module):
+    class InputClass(Enum):
+        TARGET = 0
+        CONDITION = 1
+
+    def __init__(self, hp):
+        super().__init__()
+        self.hp = hp
+        config = GPTNeoXConfig(
+            vocab_size=hp.vocab_size,
+            hidden_size=hp.d_model,
+            num_hidden_layers=hp.num_layers,
+            num_attention_heads=hp.num_heads,
+            intermediate_size=hp.d_model * 4,
+            hidden_act=hp.activation,
+            hidden_dropout=hp.dropout,
+            max_position_embeddings=hp.max_position_embeddings,
+        )
+        self.model = GPTNeoXModel(config)
+        self.word_emb = nn.Embedding(hp.vocab_size, hp.d_model, padding_idx=0)
+        self.cond_encoder = ConditionEncoder(hp)
+        self.cls_emb = nn.Embedding(
+            hp.token_class, hp.d_model, padding_idx=0
+        )  # 0: target, 1: condition
+        self.lm_head = nn.Linear(hp.d_model, hp.vocab_size)
+
+    @staticmethod
+    def from_pretrained(
+        ckpt_file=None,
+        config_file=None,
+        device="cpu",
+    ):
+        ckpt_file = ckpt_file if ckpt_file is not None else assets.checkpoint_file()
+        config_file = config_file if config_file is not None else assets.config_file()
+        hp = load_config(config_file)
+        model = PiCoGenDecoder(hp)
+        state_dict = load_checkpoint(ckpt_file, device)
+        model.load_state_dict(state_dict["model"])
+        model.to(device)
+        model.eval()
+        return model
+
+    def generate(
+        self, input_seg, input_cls_ids, need_encode, kv_cache=None, temperature=1.0, thres=0.9
+    ):
+        B, L = input_cls_ids.shape
+
+        if kv_cache is None:
+            input_ids = torch.zeros(B, L, device=_get_device(self.word_emb)).long()
+            input_cond_embs = torch.zeros(
+                B,
+                L,
+                self.hp.condition_class,
+                self.hp.d_model,
+                device=_get_device(self.cond_encoder),
+            ).float()
+            for b in range(B):
+                for ll in range(L):
+                    if need_encode[b, ll]:
+                        emb = torch.FloatTensor(np.array(input_seg[b][ll])).to(
+                            _get_device(self.cond_encoder)
+                        )
+                        input_cond_embs[b, ll] = emb
+                    else:
+                        input_ids[b, ll] = input_seg[b][ll]
+        else:  # NOTE: only use the last token as input
+            input_ids = torch.zeros(B, 1, device=_get_device(self.word_emb)).long()
+            input_cond_embs = torch.zeros(
+                B,
+                1,
+                self.hp.condition_class,
+                self.hp.d_model,
+                device=_get_device(self.cond_encoder),
+            ).float()
+            for b in range(B):
+                if need_encode[b, -1]:
+                    emb = torch.FloatTensor(np.array(input_seg[b][-1])).to(
+                        _get_device(self.cond_encoder)
+                    )
+                    input_cond_embs[b, -1] = emb
+                else:
+                    input_ids[b, -1] = input_seg[b][-1]
+            input_cls_ids = input_cls_ids[:, -1:]
+        assert input_ids.shape == input_cls_ids.shape
+
+        input_embs = self.word_emb(input_ids)
+        input_cond_embs = self.cond_encoder(input_cond_embs)
+        input_cls_embs = self.cls_emb(input_cls_ids)
+
+        if kv_cache is None:
+            mask = (input_embs.sum(dim=-1, keepdim=True) != 0).expand(B, L, self.hp.d_model)
+        else:
+            mask = (input_embs.sum(dim=-1, keepdim=True) != 0).expand(B, 1, self.hp.d_model)
+        input_cond_embs[mask] = 0  # NOTE: where input_embs is not zero
+
+        input_embs = input_embs + input_cond_embs + input_cls_embs
+
+        model_out = self.model(
+            inputs_embeds=input_embs,
+            past_key_values=kv_cache,
+        )
+
+        logits = self.lm_head(model_out.last_hidden_state)[:, -1, :]
+        assert logits.shape == (B, self.hp.vocab_size)
+        probs = F.softmax(top_p(logits, thres=thres, temperature=temperature), dim=-1)
+        output_ids = torch.multinomial(probs, num_samples=1)
+        assert output_ids.shape == (B, 1)
+
+        return output_ids, model_out.past_key_values
+
+    def forward(
+        self,
+        input_seqs,
+        input_cls_ids,
+        need_encode,
+        input_ids=None,
+        input_cond_embs=None,
+        labels=None,
+        kv_cache=None,
+    ):
+        B, L = input_cls_ids.shape
+        input_cls_ids = input_cls_ids.to(_get_device(self.cls_emb))
+
+        if input_seqs is not None:
+            assert input_ids is None and input_cond_embs is None
+            input_ids = torch.zeros(B, L, device=_get_device(self.word_emb)).long()
+            input_cond_embs = torch.zeros(
+                B,
+                L,
+                self.hp.condition_class,
+                self.hp.d_model,
+                device=_get_device(self.cond_encoder),
+            ).float()
+            for b in range(B):
+                for ll in range(L):
+                    if need_encode[b, ll]:
+                        emb = torch.FloatTensor(np.array(input_seqs[b][ll])).to(
+                            _get_device(self.cond_encoder)
+                        )
+                        input_cond_embs[b, ll] = emb
+                    else:
+                        input_ids[b, ll] = input_seqs[b][ll]
+        else:
+            assert input_ids is not None and input_cond_embs is not None
+            input_ids = input_ids.to(_get_device(self.word_emb))
+            input_cond_embs = input_cond_embs.to(_get_device(self.cond_encoder))
+
+        input_embs = self.word_emb(input_ids)
+        input_cond_embs = self.cond_encoder(input_cond_embs)
+        input_cls_embs = self.cls_emb(input_cls_ids)
+
+        mask = (input_embs.sum(dim=-1, keepdim=True) != 0).expand(B, L, self.hp.d_model)
+        input_cond_embs[mask] = 0  # NOTE: where input_embs is not zero
+
+        input_embs = input_embs + input_cond_embs + input_cls_embs
+
+        model_out = self.model(
+            inputs_embeds=input_embs,
+            past_key_values=kv_cache,
+        )
+
+        logits = self.lm_head(model_out.last_hidden_state)
+        assert logits.shape == (B, L, self.hp.vocab_size)
+
+        lm_loss = None
+        if labels is not None:
+            assert labels.shape == (B, L)
+            labels = labels.to(logits.device)
+
+            loss_fct = F.cross_entropy
+            lm_loss = loss_fct(logits.view(-1, self.hp.vocab_size), labels.view(-1))
+
+        out = {
+            "loss": lm_loss,
+            "logits": logits,
+            "past_key_values": model_out.past_key_values,
+            "hidden_states": model_out.hidden_states,
+            "attentions": model_out.attentions,
+        }
+
+        return out

picogen2/repr.py

@@ -0,0 +1,530 @@
+import collections
+import json
+from itertools import chain
+from pathlib import Path
+
+import miditoolkit
+import numpy as np
+
+from . import assets
+from .utils import load_config
+
+DEFAULT_SUBBEAT_RANGE = np.arange(0, 64, dtype=int)
+DEFAULT_PIANO_RANGE = np.arange(21, 109, dtype=int)
+DEFAULT_VELOCITY_BINS = np.linspace(0, 124, 31 + 1, dtype=int)  # midi velocity: 0~127
+LS_DEFAULT_VELOCITY = 80
+DEFAULT_BPM_BINS = np.linspace(32, 224, 64 + 1, dtype=int)
+DEFAULT_DURATION_RANGE = np.arange(1, 1 + 32, dtype=int)
+DEFAULT_CHORD_ROOTS = [
+    "A",
+    "A#",
+    "B",
+    "C",
+    "C#",
+    "D",
+    "D#",
+    "E",
+    "F",
+    "F#",
+    "G",
+    "G#",
+]
+DEFAULT_CHORD_QUALITY = [
+    "+",
+    "/o7",
+    "7",
+    "M",
+    "M7",
+    "m",
+    "m7",
+    "o",
+    "o7",
+    "sus2",
+    "sus4",
+]
+VOCAB_SIZE = 500
+
+
+def gen_vocab():
+    spec = [f"spec_{t}" for t in ["pad", "bos", "eos", "unk", "mask", "ss", "se"]]
+    bar = ["bar_start", "bar_end"] + [f"bar_{i}" for i in range(1, 5)] + ["bar_N"]
+    position = [f"position_{i}" for i in DEFAULT_SUBBEAT_RANGE]
+    chord = ["chord_N_N"]
+    for root in DEFAULT_CHORD_ROOTS:
+        for quality in DEFAULT_CHORD_QUALITY:
+            chord.append(f"chord_{root}_{quality}")
+    tempo = [f"tempo_{i}" for i in DEFAULT_BPM_BINS]
+    pitch = [f"pitch_{i}" for i in DEFAULT_PIANO_RANGE]
+    duration = [f"duration_{i}" for i in DEFAULT_DURATION_RANGE]
+    velocity = [f"velocity_{i}" for i in DEFAULT_VELOCITY_BINS]
+    vocab = spec + bar + position + chord + tempo + pitch + duration + velocity
+    vocab = vocab + ["reserved"] * (VOCAB_SIZE - len(vocab))
+    return vocab
+
+
+class Event:
+    def __init__(
+        self,
+        etype,
+        value,
+    ):
+        self.etype = etype
+        self.value = value
+
+    def init_check(self):
+        if self.etype == "spec":
+            assert self.value.split("_")[0] in ["spec"], f"{self.etype}: {self.value}"
+
+        elif self.etype == "bar":
+            assert self.value.split("_")[0] in ["bar"], f"{self.etype}: {self.value}"
+
+        elif self.etype == "metric":
+            assert self.value.split("_")[0] in [
+                "position",
+                "chord",
+                "tempo",
+            ], f"{self.etype}: {self.value}"
+
+        elif self.etype == "note":
+            assert self.value.split("_")[0] in [
+                "pitch",
+                "duration",
+                "velocity",
+            ], f"{self.etype}: {self.value}"
+
+        else:
+            raise ValueError(f"Unknown etype: {self.etype}")
+
+    def unwrap(self, vtype):
+        return vtype(self.value.split("_")[1])
+
+    def __repr__(self):
+        return f"Event({self.etype}: {self.value})"
+
+    def __eq__(self, other):
+        if not isinstance(other, Event):
+            return False
+        return self.etype == other.etype and self.value == other.value
+
+
+class Tokenizer:
+    def __init__(self, vocab_file=None, beat_div=None, ticks_per_beat=None):
+        vocab_file = assets.vocab_file() if vocab_file is None else vocab_file
+
+        if beat_div is None or ticks_per_beat is None:
+            config_file = assets.config_file()
+            hp = load_config(config_file)
+            beat_div = hp.beat_div
+            ticks_per_beat = hp.ticks_per_beat
+
+        self.vocab = Vocab(vocab_file)
+        self.beat_div = beat_div
+        self.ticks_per_beat = ticks_per_beat
+
+    def get_song_from_midi(self, midi):
+        song = midi_to_song(midi, self.beat_div)
+        song["events"] = song_to_events(song)
+
+        song["ls_events"] = extract_leadsheet_from_events(
+            song["events"], song["metadata"]["beat_per_bar"], self.beat_div
+        )
+
+        return song
+
+    def events_to_midi(self, events):
+        midi = events_to_midi(events, self.ticks_per_beat, self.beat_div)
+        return midi
+
+    def e2i(self, e):
+        return self.vocab.t2i[e.value]
+
+    def i2e(self, eid):
+        token = self.vocab.i2t[eid]
+        if token.startswith("spec"):
+            return Event("spec", token)
+        elif token.startswith("bar"):
+            return Event("bar", token)
+        elif token.split("_")[0] in ["position", "chord", "tempo"]:
+            return Event("metric", token)
+        elif token.split("_")[0] in ["pitch", "duration", "velocity"]:
+            return Event("note", token)
+        else:
+            raise ValueError(f"Unknown token: {token}")
+
+    def get_bar_ranges(self, events, from_start=True):
+        return get_bar_ranges(events, from_start)
+
+    @staticmethod
+    def get_tempo_event(bpm):
+        tempo = DEFAULT_BPM_BINS[np.argmin(abs(DEFAULT_BPM_BINS - bpm))]
+        return Event("metric", f"tempo_{tempo}")
+
+    @staticmethod
+    def get_duration_event(duration):
+        duration = DEFAULT_DURATION_RANGE[np.argmin(abs(DEFAULT_DURATION_RANGE - duration))]
+        return Event("note", f"duration_{duration}")
+
+    @staticmethod
+    def get_velocity_event(velocity):
+        velocity = DEFAULT_VELOCITY_BINS[np.argmin(abs(DEFAULT_VELOCITY_BINS - velocity))]
+        return Event("note", f"velocity_{velocity}")
+
+
+class Vocab:
+    def __init__(self, vacab_file):
+        self.i2t = json.loads(vacab_file.read_text())
+        self.t2i = {}
+        for i, t in enumerate(self.i2t):
+            self.t2i[t] = i
+
+    def len(self):
+        return len(self.i2t)
+
+    def __len__(self):
+        return len(self.i2t)
+
+    def __repr__(self):
+        out = []
+        for i, t in enumerate(self.i2t):
+            out.append(f"{i}: {t}")
+        return "\n".join(out)
+
+
+def midi_to_song(midi_obj, beat_div):
+    assert midi_obj.ticks_per_beat % beat_div == 0
+    grid_resol = midi_obj.ticks_per_beat // beat_div
+
+    # load notes
+    instr_notes = collections.defaultdict(list)
+    for instr in midi_obj.instruments:
+        for note in instr.notes:
+            instr_notes[instr.name].append(note)
+        instr_notes[instr.name].sort(key=lambda x: x.start)
+
+    # load chords
+    chords = []
+    for marker in midi_obj.markers:
+        if marker.text.split("_")[0] != "global" and "Boundary" not in marker.text.split("_")[0]:
+            chords.append(marker)
+    chords.sort(key=lambda x: x.time)
+
+    # load tempos
+    tempos = midi_obj.tempo_changes
+    tempos.sort(key=lambda x: x.time)
+
+    # load labels
+    labels = []
+    for marker in midi_obj.markers:
+        if "Boundary" in marker.text.split("_")[0]:
+            labels.append(marker)
+    labels.sort(key=lambda x: x.time)
+
+    # load global bpm
+    global_bpm = None
+    for marker in midi_obj.markers:
+        if marker.text.split("_")[0] == "global" and marker.text.split("_")[1] == "bpm":
+            global_bpm = int(marker.text.split("_")[2])
+
+    # process notes
+    intsr_gird = dict()
+    for key in instr_notes.keys():
+        notes = instr_notes[key]
+        note_grid = collections.defaultdict(list)
+        for note in notes:
+            # quantize start
+            quant_time = round(note.start / grid_resol)
+
+            # duration
+            note_duration = note.end - note.start
+            duration = round(note_duration / grid_resol)
+            duration = max(duration, 1)  # dur >= 1
+
+            # append
+            note_grid[quant_time].append(
+                {
+                    "note": note,
+                    "pitch": note.pitch,
+                    "duration": duration,
+                    "velocity": note.velocity,
+                }
+            )
+
+        # sort
+        for time in note_grid.keys():
+            note_grid[time].sort(key=lambda x: -x["pitch"])
+
+        # set to track
+        intsr_gird[key] = note_grid.copy()
+
+    # process chords
+    chord_grid = collections.defaultdict(list)
+    for chord in chords:
+        quant_time = round(chord.time / grid_resol)
+        # chord_grid[quant_time] = [chord] # NOTE: only one chord per time
+        chord_grid[quant_time].append(chord)
+
+    # process tempo
+    tempo_grid = collections.defaultdict(list)
+    for tempo in tempos:
+        quant_time = round(tempo.time / grid_resol)
+        # tempo.tempo = DEFAULT_BPM_BINS[np.argmin(abs(DEFAULT_BPM_BINS-tempo.tempo))]
+        tempo_grid[quant_time] = [tempo]  # NOTE: only one tempo per time
+
+    all_bpm = [tempo[0].tempo for _, tempo in tempo_grid.items()]
+    assert len(all_bpm) > 0, " No tempo changes in midi file."
+    average_bpm = sum(all_bpm) / len(all_bpm)
+    if global_bpm is None:
+        global_bpm = average_bpm
+
+    # process boundary
+    label_grid = collections.defaultdict(list)
+    for label in labels:
+        quant_time = round(label.time / grid_resol)
+        label_grid[quant_time] = [label]
+
+    # collect
+    song_data = {
+        "notes": intsr_gird,
+        "chords": chord_grid,
+        "tempos": tempo_grid,
+        "labels": label_grid,
+        "metadata": {
+            "global_bpm": global_bpm,
+            "average_bpm": average_bpm,
+            "beat_div": beat_div,
+            "beat_per_bar": midi_obj.time_signature_changes[0].numerator,
+        },
+    }
+    return song_data
+
+
+def song_to_events(song):
+    beat_div = song["metadata"]["beat_div"]
+    beat_per_bar = song["metadata"]["beat_per_bar"]
+    grid_per_bar = beat_div * beat_per_bar
+
+    events = [Event("spec", "spec_ss")]
+    global_tempo = DEFAULT_BPM_BINS[
+        np.argmin(abs(DEFAULT_BPM_BINS - song["metadata"]["global_bpm"]))
+    ]
+    events.append(Event("metric", f"tempo_{global_tempo}"))
+
+    max_grid = list(chain(song["tempos"].keys(), song["chords"].keys()))
+    for _, v in song["notes"].items():
+        max_grid.extend(v.keys())
+    max_grid = max(max_grid)
+
+    for bar_i in range(0, max_grid + 1, grid_per_bar):
+        events.append(Event("bar", "bar_start"))
+        for i in range(bar_i, min(bar_i + grid_per_bar, max_grid + 1)):
+            pos = Event("metric", f"position_{i-bar_i}")
+            tmp = []
+            empty = True
+            if i in song["chords"]:
+                chord_items = song["chords"][i][0].text.split("_")
+                chord = f"{chord_items[0]}_{chord_items[1]}"
+                tmp.append(Event("metric", f"chord_{chord}"))
+                empty = False
+            if i in song["tempos"]:
+                tempo = DEFAULT_BPM_BINS[
+                    np.argmin(abs(DEFAULT_BPM_BINS - song["tempos"][i][0].tempo))
+                ]
+                tmp.append(Event("metric", f"tempo_{tempo}"))
+                empty = False
+            for _, instr in song["notes"].items():
+                if i in instr:
+                    for note in instr[i]:
+                        duration = DEFAULT_DURATION_RANGE[
+                            np.argmin(abs(DEFAULT_DURATION_RANGE - note["duration"]))
+                        ]
+                        velocity = DEFAULT_VELOCITY_BINS[
+                            np.argmin(abs(DEFAULT_VELOCITY_BINS - note["velocity"]))
+                        ]
+                        tmp.append(Event("note", f'pitch_{note["pitch"]}'))
+                        tmp.append(Event("note", f"duration_{duration}"))
+                        tmp.append(Event("note", f"velocity_{velocity}"))
+                        empty = False
+
+            if not empty:
+                events.append(pos)
+            events.extend(tmp)
+        events.append(Event("bar", "bar_end"))
+
+    events.append(Event("spec", "spec_se"))
+
+    return events
+
+
+def events_to_midi(events, ticks_per_beat, grid_div):
+    bar_ranges = get_bar_ranges(events, from_start=False)
+
+    midi = miditoolkit.MidiFile()
+    midi.ticks_per_beat = ticks_per_beat
+    track = miditoolkit.Instrument(program=0, is_drum=False, name="piano")
+    midi.instruments = [track]
+
+    bar_tick = 0
+    subbeat_tick = 0
+
+    pitch, velocity, duration = 0, 0, 0
+    bar_len = 4
+
+    for i, (start, end) in enumerate(bar_ranges):
+        assert ticks_per_beat % grid_div == 0
+        ticks_per_subbeat = ticks_per_beat // grid_div
+
+        for event in events[start:end]:
+            if event.etype == "spec":
+                pass
+            elif event.etype == "bar":
+                if event.value in ["bar_start", "bar_end"]:
+                    pass
+                else:
+                    try:
+                        bar_len = int(event.value.split("_")[1])
+                    except ValueError:
+                        assert event.value == "bar_N"
+
+            elif event.etype == "metric":
+                v = event.value
+                if v.startswith("position"):
+                    pos = int(v.split("_")[1])
+                    subbeat_tick = pos * ticks_per_subbeat
+                elif v.startswith("tempo"):
+                    tempo = int(v.split("_")[1])
+                    m = miditoolkit.TempoChange(time=bar_tick + subbeat_tick, tempo=tempo)
+                    midi.tempo_changes.append(m)
+                elif v.startswith("chord"):
+                    pass
+                else:
+                    raise ValueError(f"Unknown metric: {v}")
+            elif event.etype == "note":
+                v = event.value
+                if v.startswith("pitch"):
+                    pitch = int(v.split("_")[1])
+                elif v.startswith("duration"):
+                    duration = int(v.split("_")[1]) * ticks_per_subbeat
+                elif v.startswith("velocity"):
+                    velocity = int(v.split("_")[1])
+                    n = miditoolkit.Note(
+                        start=bar_tick + subbeat_tick,
+                        end=bar_tick + subbeat_tick + duration,
+                        pitch=pitch,
+                        velocity=velocity,
+                    )
+                    midi.instruments[0].notes.append(n)
+                else:
+                    raise ValueError(f"Unknown note: {v}")
+            else:
+                raise ValueError(f"Unknown event: {type(event)}")
+
+        bar_tick += ticks_per_beat * bar_len
+
+    return midi
+
+
+def extract_leadsheet_from_events(
+    events, beat_per_bar, beat_div, cover_beat=0, min_pitch=60, no_chord=False
+):
+    # algorithm:
+    # - skyline
+    # - filter out notes with pitch < 60
+
+    grids = []
+    for event in events:
+        if event.etype == "bar":
+            for _ in range(beat_per_bar * beat_div):
+                grids.append(list())
+
+    grid_idx = 0
+    bar_count = 0
+    subbeat = 0
+    note_tmp = []
+    first_tempo = True
+    for event in events:
+        if event.etype == "spec":
+            grids[grid_idx].append(event)
+        elif event.etype == "bar":
+            if event.value == "bar_end":
+                bar_count += 1
+                grid_idx = bar_count * (beat_per_bar * beat_div)
+                subbeat = 0
+            grids[grid_idx].append(event)
+        elif event.etype == "metric":
+            if event.value.startswith("position"):
+                subbeat = int(event.value.split("_")[1])
+            if not event.value.startswith("tempo"):
+                grids[grid_idx + subbeat].append(event)
+            else:  # tempo
+                if first_tempo:
+                    grids[grid_idx + subbeat].append(event)
+                    first_tempo = False
+        elif event.etype == "note":
+            note_tmp.append(event)
+            if event.value.startswith("velocity"):
+                pitch = note_tmp[0].value.split("_")[1]
+                if int(pitch) >= min_pitch:  # only keep notes with pitch >= min_pitch
+                    grids[grid_idx + subbeat].append(note_tmp[0])
+                    grids[grid_idx + subbeat].append(note_tmp[1])
+                    grids[grid_idx + subbeat].append(
+                        Event("note", f"velocity_{LS_DEFAULT_VELOCITY}")
+                    )
+                note_tmp = []
+
+    # select the highest note
+    for i, grid in enumerate(grids):
+        notes = [e for e in grid if e.etype == "note"]
+        notes = [(notes[i], notes[i + 1], notes[i + 2]) for i in range(0, len(notes), 3)]
+        notes.sort(key=lambda x: -x[0].unwrap(int))  # sort by pitch
+        grids[i] = [e for e in grid if not e.etype == "note"]
+        if len(notes) > 0:
+            grids[i].extend(notes[0])
+
+    # remove useless metric
+    for i, grid in enumerate(grids):
+        if len(grid) == 0 or not grid[-1].etype == "metric":
+            continue
+        metric = grid[-1]
+        assert metric.etype == "metric"
+        if metric.value.startswith("position"):
+            grids[i].pop()
+            assert len(grid) == 0 or grid[-1].etype == "bar"
+
+    return list(chain(*grids))
+
+
+def get_bar_ranges(events, from_start=True):
+    bar_idx_list = []
+    for i, event in enumerate(events):
+        if event.etype == "bar" and event.value == "bar_start":
+            bar_idx_list.append(i)
+    bar_idx_list = bar_idx_list + [len(events)]
+    if from_start:
+        bar_idx_list[0] = 0  # the first bar starts at 0
+
+    bar_ranges = []
+    for i in range(len(bar_idx_list) - 1):
+        bar_ranges.append((bar_idx_list[i], bar_idx_list[i + 1]))
+    return bar_ranges
+
+
+if __name__ == "__main__":
+    import argparse
+
+    parser = argparse.ArgumentParser()
+    subparsers = parser.add_subparsers(dest="command")
+    cmd_gen_vocab = subparsers.add_parser("gen_vocab")
+    cmd_gen_vocab.add_argument("--output_file", type=Path, required=True)
+    ca = parser.parse_args()
+
+    if ca.command is None:
+        parser.print_help()
+        exit()
+    elif ca.command == "gen_vocab":
+        vocab = gen_vocab()
+        ca.output_file.write_text(json.dumps(vocab, indent=2))
+        vocab = Vocab(ca.output_file)
+        print(vocab)
+        print("vocab size:", vocab.len())
+    else:
+        raise ValueError(f"Unknown command: {ca.command}")

picogen2/utils.py

@@ -0,0 +1,194 @@
+import json
+import logging
+import math
+import pickle
+import shutil
+from dataclasses import dataclass, fields
+from pathlib import Path
+
+import numpy as np
+import questionary
+import torch
+import torch.nn.functional as F
+
+_logger = None
+_level = None
+
+
+@dataclass
+class HyperParam:
+    beat_div: int
+    ticks_per_beat: int
+
+    seed: int
+    learning_rate: float
+    learning_rate_min: float
+    adam_b1: float
+    adam_b2: float
+    sched_T: int
+    warmup_epochs: int
+
+    vocab_size: int
+    token_class: int
+    condition_class: int
+    d_model: int
+    d_bottleneck: int
+    num_layers: int
+    num_layers_encoder: int
+    num_heads: int
+    activation: str
+    dropout: float
+    max_seq_len: int
+    max_position_embeddings: int
+
+    loss_weight: float = 1.0
+
+
+def _get_logger():
+    global _logger
+    if _logger is None:
+        _logger = logging.getLogger("picogen2")
+
+    return _logger
+
+
+class Logger:
+    def setLevel(self, level):
+        global _level, _logger
+        _level = level.upper()
+        _get_logger().setLevel(_level)
+
+    def __getattr__(self, name):
+        return getattr(_get_logger(), name)
+
+    def __repr__(self):
+        return repr(_get_logger())
+
+
+logger = Logger()
+
+
+def check_task_done(task: str, output_dir: Path):
+    done_file = output_dir / f"done_{task}"
+    return done_file.exists()
+
+
+def mark_task_done(task: str, output_dir: Path):
+    done_file = output_dir / f"done_{task}"
+    done_file.touch()
+
+
+def song_dir_name(index: int):
+    return "{:04d}".format(index)
+
+
+def load_config(config_file):
+    config = json.loads(config_file.read_text())
+    hp = HyperParam(**config)
+    logger.info("checkpoint model config:")
+    for v in fields(hp):
+        logger.info(f"\t{v.name}: {getattr(hp, v.name)}")
+    return hp
+
+
+def init_ckpt_dir(ckpt_dir, config_file, config_name="config"):
+    t_path = (ckpt_dir / config_name).with_suffix(config_file.suffix)
+    if not t_path.exists():
+        ckpt_dir.mkdir(exist_ok=True)
+        shutil.copyfile(config_file, t_path)
+    else:
+        # check if config is the same
+        if config_file.read_text() != t_path.read_text():
+            override = questionary.confirm(
+                f'Config file "{config_file}" is not same with checkpoint "{t_path}", override?',
+                default=False,
+            ).ask()
+            if override:
+                shutil.copyfile(config_file, t_path)
+            else:
+                print("Confliction between config file and checkpoint. Exit.")
+                exit()
+                # raise ValueError(f'config file {config_file} and {t_path} are not the same')
+
+
+def save_checkpoint(filepath, obj, verbose=False):
+    print("Saving checkpoint to {} ... ".format(filepath), end="") if verbose else None
+    torch.save(obj, filepath)
+    print("Done.") if verbose else None
+
+
+def scan_checkpoint(cp_dir, prefix):
+    # pattern = os.path.join(cp_dir, prefix + '????????')
+    # cp_list = glob.glob(pattern)
+    cp_list = list(cp_dir.glob(f"{prefix}*"))
+    if len(cp_list) == 0:
+        return None
+    return sorted(cp_list, key=lambda n: int(n.stem.split("_")[-1]))[-1]
+
+
+def load_checkpoint(filepath: Path, device="cpu"):
+    assert filepath.is_file()
+    logger.info("Loading '{}'".format(filepath))
+    checkpoint_dict = torch.load(filepath, map_location=device, weights_only=False)
+    logger.info("Done.")
+    return checkpoint_dict
+
+
+def downbeat_time_to_index(beats, downbeats):
+    downbeat_indices = []
+    beats = np.array(beats)
+    for downbeat in downbeats:
+        idx = np.argmin(np.abs(beats - downbeat))
+        downbeat_indices.append(idx)
+    return downbeat_indices
+
+
+def top_p(logits, thres=0.9, temperature=1.0):
+    assert logits.dim() == 2, logits.shape
+
+    sorted_logits, sorted_indices = torch.sort(logits, descending=True)
+    cum_probs = torch.cumsum(F.softmax(sorted_logits / temperature, dim=-1), dim=-1)
+
+    sorted_indices_to_remove = cum_probs > thres
+    sorted_indices_to_remove[:, 0] = False
+    sorted_logits[sorted_indices_to_remove] = float("-inf")
+
+    return sorted_logits.scatter(1, sorted_indices, sorted_logits)
+
+
+def top_k(logits, thres=0.9):
+    assert logits.dim() == 2
+
+    k = math.ceil((1 - thres) * logits.shape[-1])
+    val, ind = torch.topk(logits, k)
+    probs = torch.full_like(logits, float("-inf"))
+    probs.scatter_(1, ind, val)
+    return probs
+
+
+def normalize(audio, min_y=-1.0, max_y=1.0, eps=1e-6):
+    assert len(audio.shape) == 1
+    max_y -= eps
+    min_y += eps
+    amax = audio.max()
+    amin = audio.min()
+    audio = (max_y - min_y) * (audio - amin) / (amax - amin) + min_y
+    return audio
+
+
+def pickle_load(file):
+    return pickle.load(open(file, "rb"))
+
+
+def pickle_save(data, file):
+    pickle.dump(data, open(file, "wb"))
+
+
+def get_downbeat_indices(beats, downbeats):
+    beats = np.array(beats)
+    downbeats = np.array(downbeats)
+    downbeat_indices = []
+    for downbeat in downbeats:
+        idx = np.argmin(np.abs(beats - downbeat))
+        downbeat_indices.append(idx)
+    return np.array(downbeat_indices)

picogen2/version.py

@@ -0,0 +1,11 @@
+_MAJOR = "0"
+_MINOR = "1"
+# On main and in a nightly release the patch should be one ahead of the last
+# released build.
+_PATCH = "1"
+# This is mainly for nightly builds which have the suffix ".dev$DATE". See
+# https://semver.org/#is-v123-a-semantic-version for the semantics.
+_SUFFIX = ""
+
+VERSION_SHORT = "{0}.{1}".format(_MAJOR, _MINOR)
+VERSION = "{0}.{1}.{2}{3}".format(_MAJOR, _MINOR, _PATCH, _SUFFIX)

requirements.txt

@@ -0,0 +1,12 @@
+git+https://github.com/TEAMuP-dev/pyharp.git@v0.3.0
+# model-specific deps below:
+torch
+torchaudio
+transformers
+miditoolkit
+questionary
+soundfile
+mpi4py
+sheetsage @ git+https://github.com/tanchihpin0517/PiCoGen-sheetsage.git
+beat_this @ https://github.com/CPJKU/beat_this/archive/main.zip
+madmom @ git+https://github.com/CPJKU/madmom.git@0551aa8f48d71a367d92b5d3a347a0cf7cd97cc9