initial deploy

.gitignore

@@ -0,0 +1,4 @@
+__pycache__/
+*.pyc
+.DS_Store
+beat-this-repo/

README.md

@@ -1,13 +1,14 @@
 ---
 title: Beat This
-emoji: 👀
+emoji: 🥁
 colorFrom: blue
 colorTo: yellow
 sdk: gradio
-sdk_version: 6.20.0
-python_version: '3.13'
+sdk_version: 5.28.0
+python_version: '3.11'
 app_file: app.py
 pinned: false
+license: mit
 ---
 
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+Beat and downbeat tracking with [Beat This!](https://github.com/CPJKU/beat_this), accessible in HARP.

SOURCES.md

@@ -0,0 +1,4 @@
+# Sources — beat-this
+
+- Source repo: https://github.com/CPJKU/beat_this
+- Paper: https://arxiv.org/abs/2407.21658

app.py

@@ -0,0 +1,82 @@
+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
+
+
+from pyharp import ModelCard, build_endpoint
+from pyharp.labels import LabelList, OutputLabel
+
+import gradio as gr
+import torch
+
+from beat_this.inference import File2Beats
+
+DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
+
+# built on CPU at import time so a broken checkpoint fails fast in the logs;
+# moved to DEVICE lazily on first request, since ZeroGPU only allows CUDA
+# calls made inside an @spaces.GPU-decorated call.
+model = File2Beats(checkpoint_path="final0.ckpt", device="cpu", dbn=False)
+model_ready = DEVICE == "cpu"
+
+model_card = ModelCard(
+    name="Beat This!",
+    description="Detects the beat and downbeat (start-of-bar) positions in a piece of music.",
+    author="Francesco Foscarin, Jan Schlüter",
+    tags=["beat tracking", "rhythm"],
+)
+
+
+@spaces.GPU
+@torch.inference_mode()
+def process_fn(input_audio_path: str) -> LabelList:
+    """Finds beat and downbeat times in the input audio and returns them as labeled markers."""
+    global model, model_ready
+    if not model_ready:
+        model = File2Beats(checkpoint_path="final0.ckpt", device=DEVICE, dbn=False)
+        model_ready = True
+
+    beats, downbeats = model(input_audio_path)
+    downbeat_times = set(downbeats.tolist())
+
+    output_labels = LabelList()
+    for t in beats:
+        label = "downbeat" if float(t) in downbeat_times else "beat"
+        output_labels.append(OutputLabel(t=float(t), label=label))
+    return output_labels
+
+
+with gr.Blocks() as demo:
+    input_components = [
+        gr.Audio(type="filepath", label="Input Audio").harp_required(True),
+    ]
+    output_components = [
+        gr.JSON(label="Beats").set_info(
+            "Detected beat and downbeat times, labeled \"beat\" or \"downbeat\" (start of bar)."
+        ),
+    ]
+
+    build_endpoint(
+        model_card=model_card,
+        input_components=input_components,
+        output_components=output_components,
+        process_fn=process_fn,
+    )
+
+if __name__ == "__main__":
+    demo.queue().launch(pwa=True)

beat_this/inference.py

@@ -0,0 +1,315 @@
+import inspect
+
+import numpy as np
+import soxr
+import torch
+import torch.nn.functional as F
+
+from beat_this.model.beat_tracker import BeatThis
+from beat_this.model.postprocessor import Postprocessor
+from beat_this.preprocessing import LogMelSpect, load_audio
+from beat_this.utils import replace_state_dict_key, save_beat_tsv
+
+CHECKPOINT_URL = "https://cloud.cp.jku.at/public.php/dav/files/7ik4RrBKTS273gp"
+
+
+def load_checkpoint(checkpoint_path: str, device: str | torch.device = "cpu") -> dict:
+    """
+    Load a BeatThis checkpoint as a dictionary.
+
+    Args:
+        checkpoint_path (str, optional): The path to the checkpoint. Can be a local path, a URL, or a shortname.
+        device (torch.device or str): The device to load the model on.
+
+    Returns:
+        dict: The loaded checkpoint dictionary.
+    """
+    try:
+        # try interpreting as local file name
+        weights_only = {"weights_only": True} if torch.__version__ >= "2" else {}
+        return torch.load(checkpoint_path, map_location=device, **weights_only)
+    except FileNotFoundError:
+        try:
+            if not (
+                str(checkpoint_path).startswith("https://")
+                or str(checkpoint_path).startswith("http://")
+            ):
+                # interpret it as a name of one of our checkpoints
+                checkpoint_url = f"{CHECKPOINT_URL}/{checkpoint_path}.ckpt"
+                file_name = f"beat_this-{checkpoint_path}.ckpt"
+            else:
+                # try interpreting as a URL
+                checkpoint_url = checkpoint_path
+                file_name = None
+            return torch.hub.load_state_dict_from_url(
+                checkpoint_url,
+                file_name=file_name,
+                map_location=device,
+            )
+        except Exception:
+            raise ValueError(
+                "Could not load the checkpoint given the provided name",
+                checkpoint_path,
+            )
+
+
+def load_model(
+    checkpoint_path: str | None = "final0", device: str | torch.device = "cpu"
+) -> BeatThis:
+    """
+    Load a BeatThis model from a checkpoint.
+
+    Args:
+        checkpoint_path (str, optional): The path to the checkpoint. Can be a local path, a URL, or a shortname.
+        device (torch.device or str): The device to load the model on.
+
+    Returns:
+        BeatThis: The loaded model.
+    """
+    if checkpoint_path is not None:
+        checkpoint = load_checkpoint(checkpoint_path, device)
+        # Retrieve the model hyperparameters as it could be the small model
+        hparams = checkpoint["hyper_parameters"]
+        # Filter only those hyperparameters that apply to the model itself
+        hparams = {
+            k: v
+            for k, v in hparams.items()
+            if k in set(inspect.signature(BeatThis).parameters)
+        }
+        # Create the uninitialized model
+        model = BeatThis(**hparams)
+        # The PLBeatThis (LightningModule) state_dict contains the BeatThis
+        # state_dict under the "model." prefix; remove the prefix to load it
+        state_dict = replace_state_dict_key(checkpoint["state_dict"], "model.", "")
+        model.load_state_dict(state_dict)
+    else:
+        model = BeatThis()
+    return model.to(device).eval()
+
+
+def zeropad(spect: torch.Tensor, left: int = 0, right: int = 0):
+    """
+    Pads a tensor spectrogram matrix of shape (time x bins) with `left` frames in the beginning and `right` frames in the end.
+    """
+    if left == 0 and right == 0:
+        return spect
+    else:
+        return F.pad(spect, (0, 0, left, right), "constant", 0)
+
+
+def split_piece(
+    spect: torch.Tensor,
+    chunk_size: int,
+    border_size: int = 6,
+    avoid_short_end: bool = True,
+):
+    """
+    Split a tensor spectrogram matrix of shape (time x bins) into time chunks of `chunk_size` and return the chunks and starting positions.
+    The `border_size` is the number of frames assumed to be discarded in the predictions on either side (since the model was not trained on the input edges due to the max-pool in the loss).
+    To cater for this, the first and last chunk are padded by `border_size` on the beginning and end, respectively, and consecutive chunks overlap by `border_size`.
+    If `avoid_short_end` is true, the last chunk start is shifted left to ends at the end of the piece, therefore the last chunk can potentially overlap with previous chunks more than border_size, otherwise it will be a shorter segment.
+    If the piece is shorter than `chunk_size`, avoid_short_end is ignored and the piece is returned as a single shorter chunk.
+
+    Args:
+        spect (torch.Tensor): The input spectrogram tensor of shape (time x bins).
+        chunk_size (int): The size of the chunks to produce.
+        border_size (int, optional): The size of the border to overlap between chunks. Defaults to 6.
+        avoid_short_end (bool, optional): If True, the last chunk is shifted left to end at the end of the piece. Defaults to True.
+    """
+    # generate the start and end indices
+    starts = np.arange(
+        -border_size, len(spect) - border_size, chunk_size - 2 * border_size
+    )
+    if avoid_short_end and len(spect) > chunk_size - 2 * border_size:
+        # if we avoid short ends, move the last index to the end of the piece - (chunk_size - border_size)
+        starts[-1] = len(spect) - (chunk_size - border_size)
+    # generate the chunks
+    chunks = [
+        zeropad(
+            spect[max(start, 0) : min(start + chunk_size, len(spect))],
+            left=max(0, -start),
+            right=max(0, min(border_size, start + chunk_size - len(spect))),
+        )
+        for start in starts
+    ]
+    return chunks, starts
+
+
+def aggregate_prediction(
+    pred_chunks: list,
+    starts: list,
+    full_size: int,
+    chunk_size: int,
+    border_size: int,
+    overlap_mode: str,
+    device: str | torch.device,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """
+    Aggregates the predictions for the whole piece based on the given prediction chunks.
+
+    Args:
+        pred_chunks (list): List of prediction chunks, where each chunk is a dictionary containing 'beat' and 'downbeat' predictions.
+        starts (list): List of start positions for each prediction chunk.
+        full_size (int): Size of the full piece.
+        chunk_size (int): Size of each prediction chunk.
+        border_size (int): Size of the border to be discarded from each prediction chunk.
+        overlap_mode (str): Mode for handling overlapping predictions. Can be 'keep_first' or 'keep_last'.
+        device (torch.device): Device to be used for the predictions.
+
+    Returns:
+        tuple: A tuple containing the aggregated beat predictions and downbeat predictions as torch tensors for the whole piece.
+    """
+    if border_size > 0:
+        # cut the predictions to discard the border
+        pred_chunks = [
+            {
+                "beat": pchunk["beat"][border_size:-border_size],
+                "downbeat": pchunk["downbeat"][border_size:-border_size],
+            }
+            for pchunk in pred_chunks
+        ]
+    # aggregate the predictions for the whole piece
+    piece_prediction_beat = torch.full((full_size,), -1000.0, device=device)
+    piece_prediction_downbeat = torch.full((full_size,), -1000.0, device=device)
+    if overlap_mode == "keep_first":
+        # process in reverse order, so predictions of earlier excerpts overwrite later ones
+        pred_chunks = reversed(list(pred_chunks))
+        starts = reversed(list(starts))
+    for start, pchunk in zip(starts, pred_chunks):
+        piece_prediction_beat[
+            start + border_size : start + chunk_size - border_size
+        ] = pchunk["beat"]
+        piece_prediction_downbeat[
+            start + border_size : start + chunk_size - border_size
+        ] = pchunk["downbeat"]
+    return piece_prediction_beat, piece_prediction_downbeat
+
+
+def split_predict_aggregate(
+    spect: torch.Tensor,
+    chunk_size: int,
+    border_size: int,
+    overlap_mode: str,
+    model: torch.nn.Module,
+) -> dict:
+    """
+    Function for pieces that are longer than the training length of the model.
+    Split the input piece into chunks, run the model on them, and aggregate the predictions.
+    The spect is supposed to be a torch tensor of shape (time x bins), i.e., unbatched, and the output is also unbatched.
+
+    Args:
+        spect (torch.Tensor): the input piece
+        chunk_size (int): the length of the chunks
+        border_size (int): the size of the border that is discarded from the predictions
+        overlap_mode (str): how to handle overlaps between chunks
+        model (torch.nn.Module): the model to run
+
+    Returns:
+        dict: the model framewise predictions for the hole piece as a dictionary containing 'beat' and 'downbeat' predictions.
+    """
+    # split the piece into chunks
+    chunks, starts = split_piece(
+        spect, chunk_size, border_size=border_size, avoid_short_end=True
+    )
+    # run the model
+    pred_chunks = [model(chunk.unsqueeze(0)) for chunk in chunks]
+    # remove the extra dimension in beat and downbeat prediction due to batch size 1
+    pred_chunks = [
+        {"beat": p["beat"][0], "downbeat": p["downbeat"][0]} for p in pred_chunks
+    ]
+    piece_prediction_beat, piece_prediction_downbeat = aggregate_prediction(
+        pred_chunks,
+        starts,
+        spect.shape[0],
+        chunk_size,
+        border_size,
+        overlap_mode,
+        spect.device,
+    )
+    # save it to model_prediction
+    return {"beat": piece_prediction_beat, "downbeat": piece_prediction_downbeat}
+
+
+class Spect2Frames:
+    """
+    Class for extracting framewise beat and downbeat predictions (logits) from a spectrogram.
+    """
+
+    def __init__(self, checkpoint_path="final0", device="cpu", float16=False):
+        super().__init__()
+        self.device = torch.device(device)
+        self.float16 = float16
+        self.model = load_model(checkpoint_path, self.device)
+
+    def spect2frames(self, spect):
+        with torch.inference_mode():
+            with torch.autocast(enabled=self.float16, device_type=self.device.type):
+                model_prediction = split_predict_aggregate(
+                    spect=spect,
+                    chunk_size=1500,
+                    overlap_mode="keep_first",
+                    border_size=6,
+                    model=self.model,
+                )
+        return model_prediction["beat"].float(), model_prediction["downbeat"].float()
+
+    def __call__(self, spect):
+        return self.spect2frames(spect)
+
+
+class Audio2Frames(Spect2Frames):
+    """
+    Class for extracting framewise beat and downbeat predictions (logits) from an audio tensor.
+    """
+
+    def __init__(self, checkpoint_path="final0", device="cpu", float16=False):
+        super().__init__(checkpoint_path, device, float16)
+        self.spect = LogMelSpect(device=self.device)
+
+    def signal2spect(self, signal, sr):
+        if signal.ndim == 2:
+            signal = signal.mean(1)
+        elif signal.ndim != 1:
+            raise ValueError(f"Expected 1D or 2D signal, got shape {signal.shape}")
+        if sr != 22050:
+            signal = soxr.resample(signal, in_rate=sr, out_rate=22050)
+        signal = torch.tensor(signal, dtype=torch.float32, device=self.device)
+        return self.spect(signal)
+
+    def __call__(self, signal, sr):
+        spect = self.signal2spect(signal, sr)
+        return self.spect2frames(spect)
+
+
+class Audio2Beats(Audio2Frames):
+    """
+    Class for extracting beat and downbeat positions (in seconds) from an audio tensor.
+
+    Args:
+        checkpoint_path (str): Path to the model checkpoint file. It can be a local path, a URL, or a key from the CHECKPOINT_URL dictionary. Default is "final0", which will load the model trained on all data except GTZAN with seed 0.
+        device (str): Device to use for inference. Default is "cpu".
+        float16 (bool): Whether to use half precision floating point arithmetic. Default is False.
+        dbn (bool): Whether to use the madmom DBN for post-processing. Default is False.
+    """
+
+    def __init__(
+        self, checkpoint_path="final0", device="cpu", float16=False, dbn=False
+    ):
+        super().__init__(checkpoint_path, device, float16)
+        self.frames2beats = Postprocessor(type="dbn" if dbn else "minimal")
+
+    def __call__(self, signal, sr):
+        beat_logits, downbeat_logits = super().__call__(signal, sr)
+        return self.frames2beats(beat_logits, downbeat_logits)
+
+
+class File2Beats(Audio2Beats):
+    def __call__(self, audio_path):
+        signal, sr = load_audio(audio_path)
+        return super().__call__(signal, sr)
+
+
+class File2File(File2Beats):
+    def __call__(self, audio_path, output_path):
+        downbeats, beats = super().__call__(audio_path)
+        save_beat_tsv(downbeats, beats, output_path)

beat_this/model/beat_tracker.py

@@ -0,0 +1,346 @@
+"""
+Model definitions for the Beat This! beat tracker.
+"""
+
+import contextlib
+from collections import OrderedDict
+
+import torch
+from einops import rearrange
+from einops.layers.torch import Rearrange
+from rotary_embedding_torch import RotaryEmbedding
+from torch import nn
+
+from beat_this.model import roformer
+from beat_this.utils import replace_state_dict_key
+
+
+class BeatThis(nn.Module):
+    """
+    A neural network model for beat tracking. It is composed of three main components:
+    - a frontend that processes the input spectrogram,
+    - a series of transformer blocks that process the output of the frontend,
+    - a head that produces the final beat and downbeat predictions.
+
+    Args:
+        spect_dim (int): The dimension of the input spectrogram (default: 128).
+        transformer_dim (int): The dimension of the main transformer blocks (default: 512).
+        ff_mult (int): The multiplier for the feed-forward dimension in the transformer blocks (default: 4).
+        n_layers (int): The number of transformer blocks (default: 6).
+        head_dim (int): The dimension of each attention head for the partial transformers in the frontend and the transformer blocks (default: 32).
+        stem_dim (int): The out dimension of the stem convolutional layer (default: 32).
+        dropout (dict): A dictionary specifying the dropout rates for different parts of the model
+            (default: {"frontend": 0.1, "transformer": 0.2}).
+        sum_head (bool): Whether to use a SumHead for the final predictions (default: True) or plain independent projections.
+        partial_transformers (bool): Whether to include partial frequency- and time-transformers in the frontend (default: True)
+    """
+
+    def __init__(
+        self,
+        spect_dim: int = 128,
+        transformer_dim: int = 512,
+        ff_mult: int = 4,
+        n_layers: int = 6,
+        head_dim: int = 32,
+        stem_dim: int = 32,
+        dropout: dict = {"frontend": 0.1, "transformer": 0.2},
+        sum_head: bool = True,
+        partial_transformers: bool = True,
+    ):
+        super().__init__()
+        # shared rotary embedding for frontend blocks and transformer blocks
+        rotary_embed = RotaryEmbedding(head_dim)
+
+        # create the frontend
+        # - stem
+        stem = self.make_stem(spect_dim, stem_dim)
+        spect_dim //= 4  # frequencies were convolved with stride 4
+        # - three frontend blocks
+        frontend_blocks = []
+        dim = stem_dim
+        for _ in range(3):
+            frontend_blocks.append(
+                self.make_frontend_block(
+                    dim,
+                    dim * 2,
+                    partial_transformers,
+                    head_dim,
+                    rotary_embed,
+                    dropout["frontend"],
+                )
+            )
+            dim *= 2
+            spect_dim //= 2  # frequencies were convolved with stride 2
+        frontend_blocks = nn.Sequential(*frontend_blocks)
+        # - linear projection to transformer dimensionality
+        concat = Rearrange("b c f t -> b t (c f)")
+        linear = nn.Linear(dim * spect_dim, transformer_dim)
+        self.frontend = nn.Sequential(
+            OrderedDict(stem=stem, blocks=frontend_blocks, concat=concat, linear=linear)
+        )
+
+        # create the transformer blocks
+        assert (
+            transformer_dim % head_dim == 0
+        ), "transformer_dim must be divisible by head_dim"
+        n_heads = transformer_dim // head_dim
+        self.transformer_blocks = roformer.Transformer(
+            dim=transformer_dim,
+            depth=n_layers,
+            heads=n_heads,
+            attn_dropout=dropout["transformer"],
+            ff_dropout=dropout["transformer"],
+            rotary_embed=rotary_embed,
+            ff_mult=ff_mult,
+            dim_head=head_dim,
+            norm_output=True,
+        )
+
+        # create the output heads
+        if sum_head:
+            self.task_heads = SumHead(transformer_dim)
+        else:
+            self.task_heads = Head(transformer_dim)
+
+        # init all weights
+        self.apply(self._init_weights)
+
+    @staticmethod
+    def make_stem(spect_dim: int, stem_dim: int) -> nn.Module:
+        return nn.Sequential(
+            OrderedDict(
+                rearrange_tf=Rearrange("b t f -> b f t"),
+                bn1d=nn.BatchNorm1d(spect_dim),
+                add_channel=Rearrange("b f t -> b 1 f t"),
+                conv2d=nn.Conv2d(
+                    in_channels=1,
+                    out_channels=stem_dim,
+                    kernel_size=(4, 3),
+                    stride=(4, 1),
+                    padding=(0, 1),
+                    bias=False,
+                ),
+                bn2d=nn.BatchNorm2d(stem_dim),
+                activation=nn.GELU(),
+            )
+        )
+
+    @staticmethod
+    def make_frontend_block(
+        in_dim: int,
+        out_dim: int,
+        partial_transformers: bool = True,
+        head_dim: int | None = 32,
+        rotary_embed: RotaryEmbedding | None = None,
+        dropout: float = 0.1,
+    ) -> nn.Module:
+        if partial_transformers and (head_dim is None or rotary_embed is None):
+            raise ValueError(
+                "Must specify head_dim and rotary_embed for using partial_transformers"
+            )
+        return nn.Sequential(
+            OrderedDict(
+                partial=(
+                    PartialFTTransformer(
+                        dim=in_dim,
+                        dim_head=head_dim,
+                        n_head=in_dim // head_dim,
+                        rotary_embed=rotary_embed,
+                        dropout=dropout,
+                    )
+                    if partial_transformers
+                    else nn.Identity()
+                ),
+                # conv block
+                conv2d=nn.Conv2d(
+                    in_channels=in_dim,
+                    out_channels=out_dim,
+                    kernel_size=(2, 3),
+                    stride=(2, 1),
+                    padding=(0, 1),
+                    bias=False,
+                ),
+                # out_channels : 64, 128, 256
+                # freqs : 16, 8, 4 (due to the stride=2)
+                norm=nn.BatchNorm2d(out_dim),
+                activation=nn.GELU(),
+            )
+        )
+
+    @staticmethod
+    def _init_weights(module: nn.Module):
+        if isinstance(module, (nn.Linear, nn.Conv1d)):
+            torch.nn.init.normal_(module.weight, mean=0.0, std=0.02)
+            if module.bias is not None:
+                torch.nn.init.zeros_(module.bias)
+        elif isinstance(module, nn.Conv2d):
+            torch.nn.init.kaiming_normal_(
+                module.weight, mode="fan_out", nonlinearity="relu"
+            )
+            if module.bias is not None:
+                torch.nn.init.zeros_(module.bias)
+        elif isinstance(module, nn.Embedding):
+            torch.nn.init.normal_(module.weight, mean=0.0, std=0.02)
+            if module.padding_idx is not None:
+                with torch.no_grad():
+                    module.weight[module.padding_idx].fill_(0)
+
+    def forward(self, x):
+        x = self.frontend(x)
+        x = self.transformer_blocks(x)
+        x = self.task_heads(x)
+        return x
+
+    def _load_from_state_dict(self, state_dict, prefix, *args, **kwargs):
+        # remove _orig_mod prefixes for compiled models
+        state_dict = replace_state_dict_key(state_dict, "_orig_mod.", "")
+        super()._load_from_state_dict(state_dict, prefix, *args, **kwargs)
+
+    def state_dict(self, *args, **kwargs):
+        state_dict = super().state_dict(*args, **kwargs)
+        # remove _orig_mod prefixes for compiled models
+        state_dict = replace_state_dict_key(state_dict, "_orig_mod.", "")
+        return state_dict
+
+
+class PartialRoformer(nn.Module):
+    """
+    Takes a (batch, channels, freqs, time) input, applies self-attention and
+    a feed-forward block either only across frequencies or only across time.
+    Returns a tensor of the same shape as the input.
+    """
+
+    def __init__(
+        self,
+        dim: int,
+        dim_head: int,
+        n_head: int,
+        direction: str,
+        rotary_embed: RotaryEmbedding,
+        dropout: float,
+    ):
+        super().__init__()
+
+        assert dim % dim_head == 0, "dim must be divisible by dim_head"
+        assert dim // dim_head == n_head, "n_head must be equal to dim // dim_head"
+        self.direction = direction[0].lower()
+        if self.direction not in "ft":
+            raise ValueError(f"direction must be F or T, got {direction}")
+        self.attn = roformer.Attention(
+            dim,
+            heads=n_head,
+            dim_head=dim_head,
+            dropout=dropout,
+            rotary_embed=rotary_embed,
+        )
+        self.ff = roformer.FeedForward(dim, dropout=dropout)
+
+    def forward(self, x):
+        b = len(x)
+        if self.direction == "f":
+            pattern = "(b t) f c"
+        elif self.direction == "t":
+            pattern = "(b f) t c"
+        x = rearrange(x, f"b c f t -> {pattern}")
+        x = x + self.attn(x)
+        x = x + self.ff(x)
+        x = rearrange(x, f"{pattern} -> b c f t", b=b)
+        return x
+
+
+class PartialFTTransformer(nn.Module):
+    """
+    Takes a (batch, channels, freqs, time) input, applies self-attention and
+    a feed-forward block once across frequencies and once across time. Same
+    as applying two PartialRoformer() in sequence, but encapsulated in a single
+    module. Returns a tensor of the same shape as the input.
+    """
+
+    def __init__(
+        self,
+        dim: int,
+        dim_head: int,
+        n_head: int,
+        rotary_embed: RotaryEmbedding,
+        dropout: float,
+    ):
+        super().__init__()
+
+        assert dim % dim_head == 0, "dim must be divisible by dim_head"
+        assert dim // dim_head == n_head, "n_head must be equal to dim // dim_head"
+        # frequency directed partial transformer
+        self.attnF = roformer.Attention(
+            dim,
+            heads=n_head,
+            dim_head=dim_head,
+            dropout=dropout,
+            rotary_embed=rotary_embed,
+        )
+        self.ffF = roformer.FeedForward(dim, dropout=dropout)
+        # time directed partial transformer
+        self.attnT = roformer.Attention(
+            dim,
+            heads=n_head,
+            dim_head=dim_head,
+            dropout=dropout,
+            rotary_embed=rotary_embed,
+        )
+        self.ffT = roformer.FeedForward(dim, dropout=dropout)
+
+    def forward(self, x):
+        b = len(x)
+        # frequency directed partial transformer
+        x = rearrange(x, "b c f t -> (b t) f c")
+        x = x + self.attnF(x)
+        x = x + self.ffF(x)
+        # time directed partial transformer
+        x = rearrange(x, "(b t) f c ->(b f) t c", b=b)
+        x = x + self.attnT(x)
+        x = x + self.ffT(x)
+        x = rearrange(x, "(b f) t c -> b c f t", b=b)
+        return x
+
+
+class SumHead(nn.Module):
+    """
+    A PyTorch module that produces the final beat and downbeat prediction logits.
+    The beats are a sum of all beats and all downbeats predictions, to reduce the prediction
+    of downbeats which are not beats.
+    """
+
+    def __init__(self, input_dim):
+        super().__init__()
+        self.beat_downbeat_lin = nn.Linear(input_dim, 2)
+
+    def forward(self, x):
+        beat_downbeat = self.beat_downbeat_lin(x)
+        # separate beat from downbeat
+        beat, downbeat = rearrange(beat_downbeat, "b t c -> c b t", c=2)
+        # aggregate beats and downbeats prediction
+        # autocast to float16 disabled to avoid numerical issues causing NaNs
+        if hasattr(
+            torch.amp, "is_autocast_available"
+        ) and not torch.amp.is_autocast_available(beat.device.type):
+            # but do not try disabling if the device does not support autocast
+            disable_autocast = contextlib.nullcontext()
+        else:
+            disable_autocast = torch.autocast(beat.device.type, enabled=False)
+        with disable_autocast:
+            beat = beat.float() + downbeat.float()
+        return {"beat": beat, "downbeat": downbeat}
+
+
+class Head(nn.Module):
+    """
+    A PyToch module that produces the final beat and downbeat prediction logits with independent linear layers outputs.
+    """
+
+    def __init__(self, input_dim):
+        super().__init__()
+        self.beat_downbeat_lin = nn.Linear(input_dim, 2)
+
+    def forward(self, x):
+        beat_downbeat = self.beat_downbeat_lin(x)
+        # separate beat from downbeat
+        beat, downbeat = rearrange(beat_downbeat, "b t c -> c b t", c=2)
+        return {"beat": beat, "downbeat": downbeat}

beat_this/model/postprocessor.py

@@ -0,0 +1,197 @@
+from concurrent.futures import ThreadPoolExecutor
+
+import numpy as np
+import torch
+import torch.nn.functional as F
+from einops import rearrange
+
+
+class Postprocessor:
+    """Postprocessor for the beat and downbeat predictions of the model.
+    The postprocessor takes the (framewise) model predictions (beat and downbeats) and the padding mask,
+    and returns the postprocessed beat and downbeat as list of times in seconds.
+    The beats and downbeats can be 1D arrays (for only 1 piece) or 2D arrays, if a batch of pieces is considered.
+    The output dimensionality is the same as the input dimensionality.
+    Two types of postprocessing are implemented:
+        - minimal: a simple postprocessing that takes the maximum of the framewise predictions,
+        and removes adjacent peaks.
+        - dbn: a postprocessing based on the Dynamic Bayesian Network proposed by Böck et al.
+    Args:
+        type (str): the type of postprocessing to apply. Either "minimal" or "dbn". Default is "minimal".
+        fps (int): the frames per second of the model framewise predictions. Default is 50.
+    """
+
+    def __init__(self, type: str = "minimal", fps: int = 50):
+        assert type in ["minimal", "dbn"]
+        self.type = type
+        self.fps = fps
+        if type == "dbn":
+            from madmom.features.downbeats import DBNDownBeatTrackingProcessor
+
+            self.dbn = DBNDownBeatTrackingProcessor(
+                beats_per_bar=[3, 4],
+                min_bpm=55.0,
+                max_bpm=215.0,
+                fps=self.fps,
+                transition_lambda=100,
+            )
+
+    def __call__(
+        self,
+        beat: torch.Tensor,
+        downbeat: torch.Tensor,
+        padding_mask: torch.Tensor | None = None,
+    ) -> tuple[np.ndarray, np.ndarray]:
+        """
+        Apply postprocessing to the input beat and downbeat tensors. Works with batched and unbatched inputs.
+        The output is a list of times in seconds, or a list of lists of times in seconds, if the input is batched.
+
+        Args:
+            beat (torch.Tensor): The input beat tensor.
+            downbeat (torch.Tensor): The input downbeat tensor.
+            padding_mask (torch.Tensor, optional): The padding mask tensor. Defaults to None.
+
+        Returns:
+            torch.Tensor: The postprocessed beat tensor.
+            torch.Tensor: The postprocessed downbeat tensor.
+        """
+        batched = False if beat.ndim == 1 else True
+        if padding_mask is None:
+            padding_mask = torch.ones_like(beat, dtype=torch.bool)
+
+        # if beat and downbeat are 1D tensors, add a batch dimension
+        if not batched:
+            beat = beat.unsqueeze(0)
+            downbeat = downbeat.unsqueeze(0)
+            padding_mask = padding_mask.unsqueeze(0)
+
+        if self.type == "minimal":
+            postp_beat, postp_downbeat = self.postp_minimal(
+                beat, downbeat, padding_mask
+            )
+        elif self.type == "dbn":
+            postp_beat, postp_downbeat = self.postp_dbn(beat, downbeat, padding_mask)
+        else:
+            raise ValueError("Invalid postprocessing type")
+
+        # remove the batch dimension if it was added
+        if not batched:
+            postp_beat = postp_beat[0]
+            postp_downbeat = postp_downbeat[0]
+
+        # update the model prediction dict
+        return postp_beat, postp_downbeat
+
+    def postp_minimal(self, beat, downbeat, padding_mask):
+        # concatenate beat and downbeat in the same tensor of shape (B, T, 2)
+        packed_pred = rearrange(
+            [beat, downbeat], "c b t -> b t c", b=beat.shape[0], t=beat.shape[1], c=2
+        )
+        # set padded elements to -1000 (= probability zero even in float64) so they don't influence the maxpool
+        pred_logits = packed_pred.masked_fill(~padding_mask.unsqueeze(-1), -1000)
+        # reshape to (2*B, T) to apply max pooling
+        pred_logits = rearrange(pred_logits, "b t c -> (c b) t")
+        # pick maxima within +/- 70ms
+        pred_peaks = pred_logits.masked_fill(
+            pred_logits != F.max_pool1d(pred_logits, 7, 1, 3), -1000
+        )
+        # keep maxima with over 0.5 probability (logit > 0)
+        pred_peaks = pred_peaks > 0
+        #  rearrange back to two tensors of shape (B, T)
+        beat_peaks, downbeat_peaks = rearrange(
+            pred_peaks, "(c b) t -> c b t", b=beat.shape[0], t=beat.shape[1], c=2
+        )
+        # run the piecewise operations
+        with ThreadPoolExecutor() as executor:
+            postp_beat, postp_downbeat = zip(
+                *executor.map(
+                    self._postp_minimal_item, beat_peaks, downbeat_peaks, padding_mask
+                )
+            )
+        return postp_beat, postp_downbeat
+
+    def _postp_minimal_item(self, padded_beat_peaks, padded_downbeat_peaks, mask):
+        """Function to compute the operations that must be computed piece by piece, and cannot be done in batch."""
+        # unpad the predictions by truncating the padding positions
+        beat_peaks = padded_beat_peaks[mask]
+        downbeat_peaks = padded_downbeat_peaks[mask]
+        # pass from a boolean array to a list of times in frames.
+        beat_frame = torch.nonzero(beat_peaks).cpu().numpy()[:, 0]
+        downbeat_frame = torch.nonzero(downbeat_peaks).cpu().numpy()[:, 0]
+        # remove adjacent peaks
+        beat_frame = deduplicate_peaks(beat_frame, width=1)
+        downbeat_frame = deduplicate_peaks(downbeat_frame, width=1)
+        # convert from frame to seconds
+        beat_time = beat_frame / self.fps
+        downbeat_time = downbeat_frame / self.fps
+        # move the downbeat to the nearest beat
+        if (
+            len(beat_time) > 0
+        ):  # skip if there are no beats, like in the first training steps
+            for i, d_time in enumerate(downbeat_time):
+                beat_idx = np.argmin(np.abs(beat_time - d_time))
+                downbeat_time[i] = beat_time[beat_idx]
+        # remove duplicate downbeat times (if some db were moved to the same position)
+        downbeat_time = np.unique(downbeat_time)
+        return beat_time, downbeat_time
+
+    def postp_dbn(self, beat, downbeat, padding_mask):
+        beat_prob = beat.double().sigmoid()
+        downbeat_prob = downbeat.double().sigmoid()
+        # limit lower and upper bound, since 0 and 1 create problems in the DBN
+        epsilon = 1e-5
+        beat_prob = beat_prob * (1 - epsilon) + epsilon / 2
+        downbeat_prob = downbeat_prob * (1 - epsilon) + epsilon / 2
+        with ThreadPoolExecutor() as executor:
+            postp_beat, postp_downbeat = zip(
+                *executor.map(
+                    self._postp_dbn_item, beat_prob, downbeat_prob, padding_mask
+                )
+            )
+        return postp_beat, postp_downbeat
+
+    def _postp_dbn_item(self, padded_beat_prob, padded_downbeat_prob, mask):
+        """Function to compute the operations that must be computed piece by piece, and cannot be done in batch."""
+        # unpad the predictions by truncating the padding positions
+        beat_prob = padded_beat_prob[mask]
+        downbeat_prob = padded_downbeat_prob[mask]
+        # build an artificial multiclass prediction, as suggested by Böck et al.
+        # again we limit the lower bound to avoid problems with the DBN
+        epsilon = 1e-5
+        combined_act = np.vstack(
+            (
+                np.maximum(
+                    beat_prob.cpu().numpy() - downbeat_prob.cpu().numpy(), epsilon / 2
+                ),
+                downbeat_prob.cpu().numpy(),
+            )
+        ).T
+        # run the DBN
+        dbn_out = self.dbn(combined_act)
+        postp_beat = dbn_out[:, 0]
+        postp_downbeat = dbn_out[dbn_out[:, 1] == 1][:, 0]
+        return postp_beat, postp_downbeat
+
+
+def deduplicate_peaks(peaks, width=1) -> np.ndarray:
+    """
+    Replaces groups of adjacent peak frame indices that are each not more
+    than `width` frames apart by the average of the frame indices.
+    """
+    result = []
+    peaks = map(int, peaks)  # ensure we get ordinary Python int objects
+    try:
+        p = next(peaks)
+    except StopIteration:
+        return np.array(result)
+    c = 1
+    for p2 in peaks:
+        if p2 - p <= width:
+            c += 1
+            p += (p2 - p) / c  # update mean
+        else:
+            result.append(p)
+            p = p2
+            c = 1
+    result.append(p)
+    return np.array(result)

beat_this/model/roformer.py

@@ -0,0 +1,181 @@
+"""
+Transformer with rotary position embedding, adapted from Phil Wang's repository
+at https://github.com/lucidrains/BS-RoFormer (under MIT License).
+"""
+
+import torch
+import torch.nn.functional as F
+from einops import rearrange
+from torch import nn
+from torch.nn import Module, ModuleList
+
+# helper functions
+
+
+def exists(val):
+    return val is not None
+
+
+# norm
+
+
+class RMSNorm(Module):
+    def __init__(self, size, dim=-1):
+        super().__init__()
+        self.scale = size**0.5
+        if dim >= 0:
+            raise ValueError(f"dim must be negative, got {dim}")
+        self.gamma = nn.Parameter(torch.ones((size,) + (1,) * (abs(dim) - 1)))
+        self.dim = dim
+
+    def forward(self, x):
+        return F.normalize(x, dim=self.dim) * self.scale * self.gamma
+
+
+# feedforward
+
+
+class FeedForward(Module):
+    def __init__(
+        self,
+        dim,
+        mult=4,
+        dropout=0.0,
+        dim_out=None,
+    ):
+        super().__init__()
+        if dim_out is None:
+            dim_out = dim
+        dim_inner = int(dim * mult)
+        self.activation = nn.GELU()
+        self.net = nn.Sequential(
+            RMSNorm(dim),
+            nn.Linear(dim, dim_inner),
+            self.activation,
+            nn.Dropout(dropout),
+            nn.Linear(dim_inner, dim_out),
+            nn.Dropout(dropout),
+        )
+
+    def forward(self, x):
+        return self.net(x)
+
+
+# attention
+
+
+class Attend(nn.Module):
+    def __init__(self, dropout=0.0, scale=None):
+        super().__init__()
+        self.dropout = dropout
+        self.scale = scale
+
+    def forward(self, q, k, v):
+        if exists(self.scale):
+            default_scale = q.shape[-1] ** -0.5
+            q = q * (self.scale / default_scale)
+
+        return F.scaled_dot_product_attention(
+            q, k, v, dropout_p=self.dropout if self.training else 0.0
+        )
+
+
+class Attention(Module):
+    def __init__(
+        self,
+        dim,
+        heads=8,
+        dim_head=64,
+        dropout=0.0,
+        rotary_embed=None,
+        gating=True,
+    ):
+        super().__init__()
+        self.heads = heads
+        self.scale = dim_head**-0.5
+        dim_inner = heads * dim_head
+
+        self.rotary_embed = rotary_embed
+
+        self.attend = Attend(dropout=dropout)
+
+        self.norm = RMSNorm(dim)
+        self.to_qkv = nn.Linear(dim, dim_inner * 3, bias=False)
+
+        if gating:
+            self.to_gates = nn.Linear(dim, heads)
+        else:
+            self.to_gates = None
+
+        self.to_out = nn.Sequential(
+            nn.Linear(dim_inner, dim, bias=False), nn.Dropout(dropout)
+        )
+
+    def forward(self, x):
+        x = self.norm(x)
+
+        q, k, v = rearrange(
+            self.to_qkv(x), "b n (qkv h d) -> qkv b h n d", qkv=3, h=self.heads
+        )
+
+        if exists(self.rotary_embed):
+            q = self.rotary_embed.rotate_queries_or_keys(q)
+            k = self.rotary_embed.rotate_queries_or_keys(k)
+
+        out = self.attend(q, k, v)
+
+        if exists(self.to_gates):
+            gates = self.to_gates(x)
+            out = out * rearrange(gates, "b n h -> b h n 1").sigmoid()
+
+        out = rearrange(out, "b h n d -> b n (h d)")
+        return self.to_out(out)
+
+
+# Roformer
+
+
+class Transformer(Module):
+    def __init__(
+        self,
+        *,
+        dim,
+        depth,
+        dim_head=32,
+        heads=16,
+        attn_dropout=0.1,
+        ff_dropout=0.1,
+        ff_mult=4,
+        norm_output=True,
+        rotary_embed=None,
+        gating=True,
+    ):
+        super().__init__()
+        self.layers = ModuleList([])
+
+        for _ in range(depth):
+            ff = FeedForward(dim=dim, mult=ff_mult, dropout=ff_dropout)
+            self.layers.append(
+                ModuleList(
+                    [
+                        Attention(
+                            dim=dim,
+                            dim_head=dim_head,
+                            heads=heads,
+                            dropout=attn_dropout,
+                            rotary_embed=rotary_embed,
+                            gating=gating,
+                        ),
+                        ff,
+                    ]
+                )
+            )
+
+        self.norm = RMSNorm(dim) if norm_output else nn.Identity()
+
+    def forward(self, x):
+        for attn, ff in self.layers:
+            x = attn(x) + x
+            x = ff(x) + x
+        x = self.norm(x)
+        return x

beat_this/preprocessing.py

@@ -0,0 +1,59 @@
+import numpy as np
+import torch
+import torchaudio
+
+
+def load_audio(path, dtype="float64"):
+    try:
+        waveform, samplerate = torchaudio.load(path, channels_first=False)
+        waveform = np.asanyarray(waveform.squeeze().numpy(), dtype=dtype)
+        return waveform, samplerate
+    except Exception:
+        # in case torchaudio fails, try soundfile
+        try:
+            import soundfile as sf
+
+            return sf.read(path, dtype=dtype)
+        except Exception:
+            # some files are not readable by soundfile, try madmom
+            try:
+                import madmom
+
+                return madmom.io.load_audio_file(str(path), dtype=dtype)
+            except Exception:
+                raise RuntimeError(f'Could not load audio from "{path}".')
+
+
+class LogMelSpect(torch.nn.Module):
+    def __init__(
+        self,
+        sample_rate=22050,
+        n_fft=1024,
+        hop_length=441,
+        f_min=30,
+        f_max=11000,
+        n_mels=128,
+        mel_scale="slaney",
+        normalized="frame_length",
+        power=1,
+        log_multiplier=1000,
+        device="cpu",
+    ):
+        super().__init__()
+        self.spect_class = torchaudio.transforms.MelSpectrogram(
+            sample_rate=sample_rate,
+            n_fft=n_fft,
+            hop_length=hop_length,
+            f_min=f_min,
+            f_max=f_max,
+            n_mels=n_mels,
+            mel_scale=mel_scale,
+            normalized=normalized,
+            power=power,
+        ).to(device)
+        self.log_multiplier = log_multiplier
+
+    def forward(self, x):
+        """Input is a waveform as a monodimensional array of shape T,
+        output is a 2D log mel spectrogram of shape (F,128)."""
+        return torch.log1p(self.log_multiplier * self.spect_class(x).T)

beat_this/utils.py

@@ -0,0 +1,111 @@
+from itertools import chain
+from pathlib import Path
+
+import numpy as np
+
+
+def index_to_framewise(index, length):
+    """Convert an index to a framewise sequence"""
+    sequence = np.zeros(length, dtype=bool)
+    sequence[index] = True
+    return sequence
+
+
+def filename_to_augmentation(filename):
+    """Convert a filename to an augmentation factor."""
+    parts = Path(filename).stem.split("_")
+    augmentations = {}
+    for part in parts[1:]:
+        if part.startswith("ps"):
+            augmentations["shift"] = int(part[2:])
+        elif part.startswith("ts"):
+            augmentations["stretch"] = int(part[2:])
+    return augmentations
+
+
+def infer_beat_numbers(beats: np.ndarray, downbeats: np.ndarray) -> np.ndarray:
+    """
+    From beat and downbeat times, infer a number for each beat such that each downbeat
+    is associated with a 1 and beats in between are counted upwards.
+    The function requires that all downbeats are also listed as beats.
+
+    Args:
+        beats (numpy.ndarray): Array of beat positions in seconds (including downbeats).
+        downbeats (numpy.ndarray): Array of downbeat positions in seconds.
+
+    Returns:
+        numbers (numpy.ndarray): Array of integer beat numbers.
+    """
+    # check if all downbeats are beats
+    if not np.all(np.isin(downbeats, beats)):
+        raise ValueError("Not all downbeats are beats.")
+
+    # handle pickup measure, by considering the beat count of the first full measure
+    if len(downbeats) >= 2:
+        # find the number of beats between the first two downbeats
+        first_downbeat, second_downbeat = np.searchsorted(beats, downbeats[:2])
+        beats_in_first_measure = second_downbeat - first_downbeat
+        # find the number of beats before the first downbeat
+        pickup_beats = first_downbeat
+        # derive where to start counting
+        if pickup_beats < beats_in_first_measure:
+            start_counter = beats_in_first_measure - pickup_beats
+        else:
+            print(
+                "WARNING: There are more beats in the pickup measure than in the first measure. The beat count will start from 2 without trying to estimate the length of the pickup measure."
+            )
+            start_counter = 1
+    else:
+        print(
+            "WARNING: There are less than two downbeats in the predictions. Something may be wrong. The beat count will start from 2 without trying to estimate the length of the pickup measure."
+        )
+        start_counter = 1
+
+    # assemble the beat numbers
+    numbers = []
+    counter = start_counter
+    downbeats = chain(downbeats, [-1])
+    next_downbeat = next(downbeats)
+    for beat in beats:
+        if beat == next_downbeat:
+            counter = 1
+            next_downbeat = next(downbeats)
+        else:
+            counter += 1
+        numbers.append(counter)
+    return np.asarray(numbers)
+
+
+def save_beat_tsv(beats: np.ndarray, downbeats: np.ndarray, outpath: str) -> None:
+    """
+    Save beat information to a tab-separated file in the standard .beats format:
+    each line has a time in seconds, a tab, and a beat number (1 = downbeat).
+    The function requires that all downbeats are also listed as beats.
+
+    Args:
+        beats (numpy.ndarray): Array of beat positions in seconds (including downbeats).
+        downbeats (numpy.ndarray): Array of downbeat positions in seconds.
+        outpath (str): Path to the output TSV file.
+
+    Returns:
+        None
+    """
+    # infer beat numbers
+    numbers = infer_beat_numbers(beats, downbeats)
+
+    # write the beat file
+    Path(outpath).parent.mkdir(parents=True, exist_ok=True)
+    try:
+        with open(outpath, "w") as f:
+            f.writelines(f"{beat}\t{number}\n" for beat, number in zip(beats, numbers))
+    except KeyboardInterrupt:
+        outpath.unlink()  # avoid half-written files
+
+
+def replace_state_dict_key(state_dict: dict, old: str, new: str):
+    """Replaces `old` in all keys of `state_dict` with `new`."""
+    keys = list(state_dict.keys())  # take snapshot of the keys
+    for key in keys:
+        if old in key:
+            state_dict[key.replace(old, new)] = state_dict.pop(key)
+    return state_dict

final0.ckpt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:8c328b45f59d8dd3dff219253ff6a8d6482be57d0133a29140e2febbf8eb8331
+size 81058141

model.json

@@ -0,0 +1,6 @@
+{
+  "name": "beat-this",
+  "package_dir": "beat_this",
+  "entry_point": "beat_this.inference.File2Beats",
+  "checkpoint": {"repo": "teamup-tech/beat-this", "filename": "final0.ckpt", "size_mb": 81}
+}

requirements.txt

@@ -0,0 +1,9 @@
+git+https://github.com/TEAMuP-dev/pyharp.git@develop
+# model-specific deps below:
+numpy>=1.20
+torch>=2
+torchaudio
+einops
+rotary-embedding-torch
+soxr
+soundfile