From 42dd15b4236db922ca063fcdee14974d42158bfb Mon Sep 17 00:00:00 2001
From: Pablo Alonso <pablo.alonso@upf.edu>
Date: Wed, 18 Jun 2025 19:21:40 +0200
Subject: [PATCH 01/15] Support batched audio input and torchaudio melspec

MAEST's forward method implements an `melspectrogram_input` parater.
When the input is 2D it will be treated as a batch of audios or a single
melspectrogram according to it.

torchaudio melspec replaces essentia for parallel, GPU-based
mel-spectrogram extraction when required.
---
 models/helpers/melspectrogram.py           | 60 ++++++++++++++++++++
 models/helpers/melspectrogram_extractor.py | 56 -------------------
 models/maest.py                            | 64 ++++++++++++++--------
 pyproject.toml                             |  2 +-
 4 files changed, 101 insertions(+), 81 deletions(-)
 create mode 100644 models/helpers/melspectrogram.py
 delete mode 100644 models/helpers/melspectrogram_extractor.py

diff --git a/models/helpers/melspectrogram.py b/models/helpers/melspectrogram.py
new file mode 100644
index 0000000..c6e3719
--- /dev/null
+++ b/models/helpers/melspectrogram.py
@@ -0,0 +1,60 @@
+import torch
+from torch.nn import Module
+from torchaudio.transforms import Spectrogram, MelScale
+
+# According to Pytoch, mel-spectrogram should be implemented as a module:
+# https://pytorch.org/audio/stable/transforms.html
+
+# WARNING: The torchaudio implementation is similar but not identical to Essentia's.
+# Relative tolerance is 1e-3 and absolute tolerance is 1e-3. We assume that this has minimal
+# impact in the resulting embeddings.
+
+
+class MelSpectrogram(Module):
+    """Extract mel-spectgrams as a torchaudio module"""
+
+    sr = 16000
+    win_len = 512
+    hop_len = 256
+    power = 2
+    n_mel = 96
+    norm = "slaney"
+    mel_scale_type = "slaney"
+    norm_mean = 2.06755686098554
+    norm_std = 1.268292820667291
+
+    def __init__(self):
+        super().__init__()
+
+        self.spec = Spectrogram(
+            n_fft=self.win_len,
+            win_length=self.win_len,
+            hop_length=self.hop_len,
+            power=self.power,
+        )
+
+        self.mel_scale = MelScale(
+            n_mels=self.n_mel,
+            sample_rate=self.sr,
+            n_stft=self.win_len // 2 + 1,
+            norm=self.norm,
+            mel_scale=self.mel_scale_type,
+        )
+
+    def znorm(self, input_values: torch.Tensor) -> torch.Tensor:
+        return (input_values - (self.norm_mean)) / (self.norm_std * 2)
+
+    def forward(self, waveform: torch.Tensor) -> torch.Tensor:
+        # convert to power spectrogram
+        spec = self.spec(waveform)
+
+        # convert to mel-scale
+        mel = self.mel_scale(spec)
+
+        # apply logC compression
+        logmel = torch.log10(1 + mel * 10000)
+
+        # normalize
+        logmel = self.znorm(logmel)
+
+        return logmel
diff --git a/models/helpers/melspectrogram_extractor.py b/models/helpers/melspectrogram_extractor.py
deleted file mode 100644
index 0f9d380..0000000
--- a/models/helpers/melspectrogram_extractor.py
+++ /dev/null
@@ -1,56 +0,0 @@
-import numpy as np
-from essentia.streaming import (
-    VectorInput,
-    MonoLoader,
-    FrameCutter,
-    TensorflowInputMusiCNN,
-)
-from essentia import Pool, run, reset
-
-
-class MelSpectrogramExtractor:
-    # mel-spec signature params
-    sample_rate = 16000
-    resample_quality = 4
-    frame_size = 512
-    hop_size = 256
-
-    def __init__(self):
-        self.pool = Pool()
-        self.frame_cutter = FrameCutter(
-            frameSize=self.frame_size,
-            hopSize=self.hop_size,
-            silentFrames="keep",
-        )
-        self.melspec = TensorflowInputMusiCNN()
-
-        self.frame_cutter.frame >> self.melspec.frame
-        self.melspec.bands >> (self.pool, "mel_bands")
-
-    def __call__(self, audio):
-        # make sure audio is float32 for Essentia
-        audio = audio.astype(np.float32)
-
-        # set vector input and connect the network
-        vector_input = VectorInput(audio)
-        vector_input.data >> self.frame_cutter.signal
-
-        # run the network and get the mel bands
-        run(vector_input)
-        mel_bands = np.array(self.pool["mel_bands"])
-
-        # clear the pool and reset the vector input
-        reset(vector_input)
-        vector_input.data.disconnect(self.frame_cutter.signal)
-        self.pool.clear()
-        del vector_input
-
-        # return freq, time
-        return mel_bands.T
-
-    def load_audio(self, audio_file):
-        return MonoLoader(
-            filename=audio_file,
-            sampleRate=self.sample_rate,
-            resampleQuality=self.resample_quality,
-        )()
diff --git a/models/maest.py b/models/maest.py
index c2b876d..cd9d337 100644
--- a/models/maest.py
+++ b/models/maest.py
@@ -24,7 +24,7 @@
     trunc_normal_,
     build_model_with_cfg,
 )
-from .helpers.melspectrogram_extractor import MelSpectrogramExtractor
+from .helpers.melspectrogram import MelSpectrogram
 from .discogs_labels import discogs_400labels, discogs_519labels
 
 _logger = logging.getLogger("MAEST")
@@ -583,8 +583,8 @@ def __init__(
 
         self.init_weights(weight_init)
 
-    def init_melspectrogram_extractor(self):
-        self.melspectrogram_extractor = MelSpectrogramExtractor()
+    def init_melspectrogram(self):
+        self.melspectrogram = MelSpectrogram()
 
     def init_weights(self, mode=""):
         assert mode in ("jax", "jax_nlhb", "nlhb", "")
@@ -663,8 +663,9 @@ def forward_features(self, x, transformer_block=-1, return_self_attention=False)
                     f"CUT time_new_pos_embed.shape: {time_new_pos_embed.shape}"
                 )
         else:
-            warnings.warn(
-                f"the patches shape:{x.shape} are larger than the expected time encodings {time_new_pos_embed.shape}, x will be cut"
+            raise Exception(
+                f"the patches shape:{x.shape} are larger than the expected time encodings {time_new_pos_embed.shape},"
+                " please reduce the input duration."
             )
             x = x[:, :, :, : time_new_pos_embed.shape[-1]]
         x = x + time_new_pos_embed
@@ -828,36 +829,43 @@ def forward_features(self, x, transformer_block=-1, return_self_attention=False)
             first_RUN = False
             return torch.cat([cls, dist, feats], dim=1)
 
-    def forward(self, x, transformer_block=-1, return_self_attention=False):
+    def forward(
+        self,
+        x,
+        transformer_block: int = -1,
+        return_self_attention: bool = False,
+        melspectrogram_input: bool = False,
+    ):
         global first_RUN
         if first_RUN:
             _logger.debug(f"x size: {len(x)}")
 
         if len(x.shape) == 1:
+            assert melspectrogram_input is False, (
+                "Input is 1D, but melspectrogram_input is True. This is not supported."
+            )
+
             _logger.debug("extracting melspec")
             if not self.melspectrogram_extractor:
                 _logger.debug("initializing melspectrogram extractor")
-                self.init_melspectrogram_extractor()
+                self.init_melspectrogram()
 
             # extract melspec from raw audio
-            x = self.melspectrogram_extractor(x)
-            # normalize
-            x = (x - DISCOGS_MEAN) / (DISCOGS_STD * 2)
-
-            # zero-pad to if if the input is shorter than the expected size
-            pad = self.img_size[1] - x.shape[1]
-
-            if pad > 0.2 * self.img_size[1]:
-                warnings.warn(
-                    "Padding input by more than 20% of the expected size. "
-                    "This may result in poor performance. "
-                    "Consider using a MAEST models trained with a smaller input size."
-                )
-
-            if pad > 0:
-                x = np.pad(x, ((0, 0), (0, pad)), mode="constant")
+            x = self.melspectrogram(x)
+
+            if x.shape[1] >= self.img_size[1]:
+                # cut into equally-sized patches.
+                trim = x.shape[1] % self.img_size[1]
+                if trim:
+                    x = x[:, :-trim]
+                x = x.reshape(self.img_size[0], 1, -1, self.img_size[1])
+                # resort axes: batch, channels, freq, time
+                x = torch.swapaxes(x, 0, 2)
+            else:
+                x.unsqueeze_(0)
+            x.unsqueeze_(1)
 
-        if len(x.shape) == 2:
+        elif len(x.shape) == 2 and melspectrogram_input:
             # need batching
             trim = x.shape[1] % self.img_size[1]
             if trim:
@@ -869,6 +877,14 @@ def forward(self, x, transformer_block=-1, return_self_attention=False):
             x = np.swapaxes(x, 0, 2)
             x = torch.tensor(x)
 
+        elif len(x.shape) == 2 and not melspectrogram_input:
+            if not self.melspectrogram_extractor:
+                _logger.debug("initializing melspectrogram extractor")
+                self.init_melspectrogram()
+
+            x = self.melspectrogram(x)
+            x.unsqueeze_(1)
+
         x = self.forward_features(
             x,
             transformer_block=transformer_block,
diff --git a/pyproject.toml b/pyproject.toml
index cba7484..3d4e9e1 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -14,7 +14,6 @@ readme = "README.md"
 requires-python = ">=3.7"
 
 dependencies = [
-    "essentia",
     "timm~=0.9",
     "torchmetrics~=0.11.4",
     "tqdm~=4.65",
@@ -24,6 +23,7 @@ dependencies = [
     "tensorboard~=2.13.0",
     "numpy~=1.26",
     "torch",
+    "torchaudio",
 ]
 
 [project.urls]

From 157ccd44f4f03e4b78a78e1a622885ddd33d9591 Mon Sep 17 00:00:00 2001
From: Pablo Alonso <pablo.alonso@upf.edu>
Date: Thu, 19 Jun 2025 08:29:16 +0200
Subject: [PATCH 02/15] Ensure that input is a torch.Tensor

---
 models/maest.py | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/models/maest.py b/models/maest.py
index cd9d337..5b9c02d 100644
--- a/models/maest.py
+++ b/models/maest.py
@@ -840,6 +840,8 @@ def forward(
         if first_RUN:
             _logger.debug(f"x size: {len(x)}")
 
+        assert isinstance(x, torch.Tensor), "Input must be a torch.Tensor"
+
         if len(x.shape) == 1:
             assert melspectrogram_input is False, (
                 "Input is 1D, but melspectrogram_input is True. This is not supported."

From 9cb99cc8af61ffe0541db26ad15ff43963877286 Mon Sep 17 00:00:00 2001
From: Pablo Alonso <pablo.alonso@upf.edu>
Date: Thu, 19 Jun 2025 08:32:46 +0200
Subject: [PATCH 03/15] Assert that input audio is not empty

---
 models/maest.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/models/maest.py b/models/maest.py
index 5b9c02d..05bdd97 100644
--- a/models/maest.py
+++ b/models/maest.py
@@ -841,6 +841,7 @@ def forward(
             _logger.debug(f"x size: {len(x)}")
 
         assert isinstance(x, torch.Tensor), "Input must be a torch.Tensor"
+        assert x.nelement() > 0, "Input tensor must not be empty"
 
         if len(x.shape) == 1:
             assert melspectrogram_input is False, (

From f831c81ba6b6576e2f00142788aa507f926b3346 Mon Sep 17 00:00:00 2001
From: Pablo Alonso <pablo.alonso@upf.edu>
Date: Thu, 19 Jun 2025 08:48:50 +0200
Subject: [PATCH 04/15] Fix dimension handling for audio input

---
 models/maest.py | 3 +--
 1 file changed, 1 insertion(+), 2 deletions(-)

diff --git a/models/maest.py b/models/maest.py
index 05bdd97..613de40 100644
--- a/models/maest.py
+++ b/models/maest.py
@@ -865,8 +865,7 @@ def forward(
                 # resort axes: batch, channels, freq, time
                 x = torch.swapaxes(x, 0, 2)
             else:
-                x.unsqueeze_(0)
-            x.unsqueeze_(1)
+                x = x.reshape(1, 1, x.shape[0], x.shape[1])
 
         elif len(x.shape) == 2 and melspectrogram_input:
             # need batching

From 19d4b4104ead85dd0211dc2355dba37eff17ddd5 Mon Sep 17 00:00:00 2001
From: Pablo Alonso <pablo.alonso@upf.edu>
Date: Thu, 19 Jun 2025 09:32:25 +0200
Subject: [PATCH 05/15] Add unit tests

---
 tests/test_maest.py | 57 +++++++++++++++++++++++++++++++++++++++++++++
 1 file changed, 57 insertions(+)
 create mode 100644 tests/test_maest.py

diff --git a/tests/test_maest.py b/tests/test_maest.py
new file mode 100644
index 0000000..8d81d75
--- /dev/null
+++ b/tests/test_maest.py
@@ -0,0 +1,57 @@
+import pytest
+import numpy as np
+import torch
+from maest import get_maest
+
+
+@pytest.fixture
+def model():
+    # Use a simple config for testing; adjust as needed for your environment
+    return get_maest(arch="discogs-maest-30s-pw-129e", pretrained=False)
+
+
+def test_numpy_input(model):
+    input_data = np.random.rand(128, 128)
+    with pytest.raises(Exception):
+        model(input_data)
+
+
+def test_empty_input(model):
+    input_data = torch.empty([])
+    with pytest.raises(Exception):
+        model(input_data)
+
+
+def test_1d_input(model):
+    input_data = torch.rand(10 * 16000).float()
+    activations, _ = model(input_data)
+    assert activations.shape == (1, 400), (
+        "Activations shape should be (batch_size, num_classes)"
+    )
+
+
+def test_2d_audio_activations(model):
+    input_data = torch.rand(2, 10 * 16000).float()
+    # Batch of 2 audio samples
+    activations, _ = model(input_data, melspectrogram_input=False)
+    assert activations.shape == (2, 400), (
+        "Output batch size should match input batch size"
+    )
+
+
+def test_2d_melspec_activations(model):
+    input_data = torch.rand(96, 1875).float()
+    # Batch of 10 mel spectrograms
+    activations, _ = model(input_data, melspectrogram_input=True)
+    assert activations.shape == (1, 400), (
+        "Output batch size should match input batch size"
+    )
+
+
+def test_2d_melspec_embeddings(model):
+    input_data = torch.rand(96, 1875).float()
+    # Batch of 10 mel spectrograms
+    _, embeddings = model(input_data, melspectrogram_input=True, transformer_block=6)
+    assert embeddings.shape == (1, 2304), (
+        "Output batch size should match input batch size"
+    )

From 857afb99fd186d7da19798f45c1256ffbd2338a3 Mon Sep 17 00:00:00 2001
From: Pablo Alonso <pablo.alonso@upf.edu>
Date: Thu, 19 Jun 2025 09:48:53 +0200
Subject: [PATCH 06/15] Replace warning by debug message

Inference with segments shorter than training time is not problem so we
don't need this warning
---
 models/maest.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/models/maest.py b/models/maest.py
index 613de40..755d15d 100644
--- a/models/maest.py
+++ b/models/maest.py
@@ -242,7 +242,7 @@ def __init__(
     def forward(self, x):
         B, C, H, W = x.shape
         if not (H == self.img_size[0] and W == self.img_size[1]):
-            warnings.warn(
+            _logger.debug(
                 f"Input image size ({H}*{W}) doesn't match model ({self.img_size[0]}*{self.img_size[1]})."
             )
         # to do maybe replace weights

From 313078295c16b9b85f6f08d57dfadaa70eba32dc Mon Sep 17 00:00:00 2001
From: Pablo Alonso <pablo.alonso@upf.edu>
Date: Thu, 19 Jun 2025 09:49:41 +0200
Subject: [PATCH 07/15] Update load_pretrain import

---
 models/helpers/vit_helpers.py | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/models/helpers/vit_helpers.py b/models/helpers/vit_helpers.py
index 2f6a45b..6b24e7b 100644
--- a/models/helpers/vit_helpers.py
+++ b/models/helpers/vit_helpers.py
@@ -2,12 +2,13 @@
 Adapted from https://github.com/rwightman/pytorch-image-models/blob/master/timm/models/vision_transformer.py
 
 """
+
 import math
 import warnings
 from copy import deepcopy
 
 import torch
-from timm.models.helpers import load_pretrained
+from timm.models import load_pretrained
 from torch import nn
 
 

From 18105268872b1a5f8154940174080c6644741c87 Mon Sep 17 00:00:00 2001
From: Pablo Alonso <pablo.alonso@upf.edu>
Date: Thu, 19 Jun 2025 09:50:20 +0200
Subject: [PATCH 08/15] We don't need to convert to tensor

---
 models/maest.py | 1 -
 1 file changed, 1 deletion(-)

diff --git a/models/maest.py b/models/maest.py
index 755d15d..9c6ffb9 100644
--- a/models/maest.py
+++ b/models/maest.py
@@ -877,7 +877,6 @@ def forward(
             # Note that the new batch axis needs to be next to the time.
             # resort axes: batch, channels, freq, time
             x = np.swapaxes(x, 0, 2)
-            x = torch.tensor(x)
 
         elif len(x.shape) == 2 and not melspectrogram_input:
             if not self.melspectrogram_extractor:

From 568a4191e0bafb19835d449d583f8422c26839cd Mon Sep 17 00:00:00 2001
From: Pablo Alonso <pablo.alonso@upf.edu>
Date: Thu, 19 Jun 2025 09:56:36 +0200
Subject: [PATCH 09/15] Init melspectrogram extractor by defaul

---
 models/maest.py | 9 ---------
 1 file changed, 9 deletions(-)

diff --git a/models/maest.py b/models/maest.py
index 9c6ffb9..c08322f 100644
--- a/models/maest.py
+++ b/models/maest.py
@@ -495,7 +495,6 @@ def __init__(
         )
         self.num_tokens = 2 if distilled else 1
         self.distilled_type = distilled_type
-        self.melspectrogram_extractor = None
         norm_layer = norm_layer or partial(nn.LayerNorm, eps=1e-6)
         act_layer = act_layer or nn.GELU
         if self.num_classes == 400:
@@ -583,7 +582,6 @@ def __init__(
 
         self.init_weights(weight_init)
 
-    def init_melspectrogram(self):
         self.melspectrogram = MelSpectrogram()
 
     def init_weights(self, mode=""):
@@ -849,9 +847,6 @@ def forward(
             )
 
             _logger.debug("extracting melspec")
-            if not self.melspectrogram_extractor:
-                _logger.debug("initializing melspectrogram extractor")
-                self.init_melspectrogram()
 
             # extract melspec from raw audio
             x = self.melspectrogram(x)
@@ -879,10 +874,6 @@ def forward(
             x = np.swapaxes(x, 0, 2)
 
         elif len(x.shape) == 2 and not melspectrogram_input:
-            if not self.melspectrogram_extractor:
-                _logger.debug("initializing melspectrogram extractor")
-                self.init_melspectrogram()
-
             x = self.melspectrogram(x)
             x.unsqueeze_(1)
 

From 219a932df076448722e74a4a49ceb2a83a5275a6 Mon Sep 17 00:00:00 2001
From: Pablo Alonso <pablo.alonso@upf.edu>
Date: Thu, 19 Jun 2025 10:03:09 +0200
Subject: [PATCH 10/15] Improve forward method doc

---
 models/maest.py | 16 +++++++++++++++-
 1 file changed, 15 insertions(+), 1 deletion(-)

diff --git a/models/maest.py b/models/maest.py
index c08322f..221026e 100644
--- a/models/maest.py
+++ b/models/maest.py
@@ -12,6 +12,7 @@
 import warnings
 from collections import OrderedDict
 from functools import partial
+from typing import Tuple
 
 import numpy as np
 import torch
@@ -833,7 +834,20 @@ def forward(
         transformer_block: int = -1,
         return_self_attention: bool = False,
         melspectrogram_input: bool = False,
-    ):
+    ) -> Tuple[torch.Tensor | None, torch.Tensor]:
+        """
+        Forward pass of the model.
+        Args:
+            x (torch.Tensor): Input tensor, either raw audio or melspectrogram.
+            transformer_block (int): If -1, returns the output of the last block.
+                If >= 0, returns the output of the specified block.
+            return_self_attention (bool): If True, returns self-attention from the specified block.
+            melspectrogram_input (bool): If True, input is expected to be a melspectrogram.
+        Returns:
+            activations (torch.Tensor): Output tensor with the model's predictions. Requires `transformer_block` to be -1.
+            embeddings (torch.Tensor): Output tensor with the model's embeddings of the selected transformer block.
+        """
+
         global first_RUN
         if first_RUN:
             _logger.debug(f"x size: {len(x)}")

From 951f81c55f14cfba0fc62655d2a64098e7a3b70d Mon Sep 17 00:00:00 2001
From: Pablo Alonso <pablo.alonso@upf.edu>
Date: Thu, 19 Jun 2025 10:07:43 +0200
Subject: [PATCH 11/15] Add GH workflow

---
 .github/workflows/python-tests.yml | 25 +++++++++++++++++++++++++
 1 file changed, 25 insertions(+)
 create mode 100644 .github/workflows/python-tests.yml

diff --git a/.github/workflows/python-tests.yml b/.github/workflows/python-tests.yml
new file mode 100644
index 0000000..9b66178
--- /dev/null
+++ b/.github/workflows/python-tests.yml
@@ -0,0 +1,25 @@
+name: Python Tests
+
+on:
+  push:
+    branches: [ main ]
+  pull_request:
+    branches: [ main ]
+
+jobs:
+  test:
+    runs-on: macos-latest
+
+    steps:
+    - uses: actions/checkout@v4
+    - name: Set up Python
+      uses: actions/setup-python@v5
+      with:
+        python-version: '3.10'
+    - name: Install dependencies
+      run: |
+        python -m pip install --upgrade pip
+        pip install pytest
+        pip install -r requirements.txt
+    - name: Run tests
+      run: pytest

From 0f052b17fb5151e52d0db88e234be024f6835566 Mon Sep 17 00:00:00 2001
From: Pablo Alonso <pablo.alonso@upf.edu>
Date: Thu, 19 Jun 2025 10:09:55 +0200
Subject: [PATCH 12/15] Fix workflow

---
 .github/workflows/python-tests.yml | 7 +++----
 1 file changed, 3 insertions(+), 4 deletions(-)

diff --git a/.github/workflows/python-tests.yml b/.github/workflows/python-tests.yml
index 9b66178..0200069 100644
--- a/.github/workflows/python-tests.yml
+++ b/.github/workflows/python-tests.yml
@@ -8,18 +8,17 @@ on:
 
 jobs:
   test:
-    runs-on: macos-latest
+    runs-on: ubuntu-latest
 
     steps:
     - uses: actions/checkout@v4
     - name: Set up Python
       uses: actions/setup-python@v5
       with:
-        python-version: '3.10'
+        python-version: '3.8'
     - name: Install dependencies
       run: |
         python -m pip install --upgrade pip
-        pip install pytest
-        pip install -r requirements.txt
+        pip install .
     - name: Run tests
       run: pytest

From 083a2bd8d4faaf27651ce85b535e467306696afa Mon Sep 17 00:00:00 2001
From: Pablo Alonso <pablo.alonso@upf.edu>
Date: Thu, 19 Jun 2025 10:13:09 +0200
Subject: [PATCH 13/15] Update pyproject

- Bump package version
- Bump min python version
- Add pytest dep
---
 .github/workflows/python-tests.yml | 2 +-
 pyproject.toml                     | 7 ++++---
 2 files changed, 5 insertions(+), 4 deletions(-)

diff --git a/.github/workflows/python-tests.yml b/.github/workflows/python-tests.yml
index 0200069..0eb8f5d 100644
--- a/.github/workflows/python-tests.yml
+++ b/.github/workflows/python-tests.yml
@@ -15,7 +15,7 @@ jobs:
     - name: Set up Python
       uses: actions/setup-python@v5
       with:
-        python-version: '3.8'
+        python-version: '3.10'
     - name: Install dependencies
       run: |
         python -m pip install --upgrade pip
diff --git a/pyproject.toml b/pyproject.toml
index 3d4e9e1..5058421 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -4,14 +4,14 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "maest"
-version = "0.1.0"
+version = "0.1.1"
 description = "MAEST model package"
 authors = [
     { name = "Pablo Alonso", email = "pablo.alonso@upf.edu" }
 ]
 license = { text = "AGPL-3.0 license" }
 readme = "README.md"
-requires-python = ">=3.7"
+requires-python = ">=3.10"
 
 dependencies = [
     "timm~=0.9",
@@ -24,6 +24,7 @@ dependencies = [
     "numpy~=1.26",
     "torch",
     "torchaudio",
+    "pytest",
 ]
 
 [project.urls]
@@ -31,7 +32,7 @@ homepage = "https://github.com/palonso/maest"
 repository = "https://github.com/palonso/maest"
 
 [tool.setuptools]
-packages = ["maest", "maest.helpers"] 
+packages = ["maest", "maest.helpers"]
 
 [tool.setuptools.package-dir]
 "maest" = "models"

From fb62bd09cc16cc29cc19bd8307713b6766373729 Mon Sep 17 00:00:00 2001
From: Pablo Alonso <pablo.alonso@upf.edu>
Date: Thu, 19 Jun 2025 10:14:13 +0200
Subject: [PATCH 14/15] Update doc

---
 README.md | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/README.md b/README.md
index 2b44b23..b97629e 100644
--- a/README.md
+++ b/README.md
@@ -173,7 +173,8 @@ _, embeddings = model(audio, transformer_block=6)
 MAEST is designed to accept `data` in different input formats:
 
 - 1D: 16kHz audio waveform is assumed.
-- 2D: mel-spectrogram is assumed (frequency, time).
+- 2D: (with `melspectrogram_input=False`) audio is assumed (batch, time).
+- 2D: (`melspectrogram_input=True`) mel-spectrogram is assumed (frequency, time).
 - 3D: batched mel-spectrogram (batch, frequency, time).
 - 4D: batched mel-spectrgroam plus singleton channel axis (batch, 1, frequency, time).
 

From 7524a6af6db5b8bff4c7ebadf6495270963f649c Mon Sep 17 00:00:00 2001
From: Pablo Alonso <pablo.alonso@upf.edu>
Date: Thu, 19 Jun 2025 10:14:38 +0200
Subject: [PATCH 15/15] Format README

---
 README.md | 8 +++-----
 1 file changed, 3 insertions(+), 5 deletions(-)

diff --git a/README.md b/README.md
index b97629e..0d3177d 100644
--- a/README.md
+++ b/README.md
@@ -12,7 +12,7 @@ The MAEST models are also available for inference from [Essentia](https://essent
 # Install
 
 Our software has been tested in Ubuntu 22.04 LTS and CentOS 7.5 using Python 3.10.
-If Python 3.10 is unavailable in your environment, we recommend using the [Conda](https://docs.conda.io) package manager to set up the working environment. 
+If Python 3.10 is unavailable in your environment, we recommend using the [Conda](https://docs.conda.io) package manager to set up the working environment.
 
 1. Create a conda environment (optional):
 
@@ -34,7 +34,7 @@ pip install -e .
 
 # Usage
 
-We use [Sacred](https://github.com/IDSIA/sacred) to run, configure and log our experiments. 
+We use [Sacred](https://github.com/IDSIA/sacred) to run, configure and log our experiments.
 The different routines can be run with Sacred commands, and many experiment options can be directly
 configure from the command line.
 
@@ -115,7 +115,7 @@ the paper:
 
 - `extract_embeddings`: returns a [3, 768] vector with the embeddings for each audio file in
   the `predict` dataset. The three dimensions of the first axis correspond to the CLS token, the
-  DIST token and the average of the rest of tokens (see Section 4.1 from the paper). 
+  DIST token and the average of the rest of tokens (see Section 4.1 from the paper).
 - `extract_logits`: returns logits that can be used in the teacher student setup, or transformed into label predictions by applying a `Sigmoid` function.
 
 Each pre-training configuration has its inference version. For example, to extract embeddings
@@ -191,7 +191,6 @@ The following `arch` values are supported:
 - `discogs-maest-30s-pw-73e-ts`
 - `discogs-maest-30s-pw-129e-519l`
 
-
 Additionally, `predict_labels()` is an auxiliary function that applies a sigmoid activation, averages the predictions along the time axes, and returns the label vector for convenience.
 
 ```python
@@ -214,4 +213,3 @@ If you are going to use MAEST as part of your research, please consider citing t
       year={2023},
     }
 ```
-