lfeats 0.2.2__tar.gz → 0.2.3__tar.gz

{lfeats-0.2.2 → lfeats-0.2.3}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.5
 Name: lfeats
-Version: 0.2.2
+Version: 0.2.3
 Summary: A unified interface to extract hidden representations from speech foundation models
 Project-URL: Documentation, https://takenori-y.github.io/lfeats/stable/
 Project-URL: Source, https://github.com/takenori-y/lfeats
@@ -130,6 +130,7 @@ pip install git+https://github.com/takenori-y/lfeats.git@master
 | `unispeech-sat` | `base` | 12 | 768 | [arXiv](https://arxiv.org/abs/2110.05752) | [GitHub](https://github.com/microsoft/UniSpeech) | [🤗](https://huggingface.co/microsoft/unispeech-sat-base) |
 | | `base+` | 12 | 768 | | | [🤗](https://huggingface.co/microsoft/unispeech-sat-base-plus) |
 | | `large` | 24 | 1024 | | | [🤗](https://huggingface.co/microsoft/unispeech-sat-large) |
+| `w2v-bert2` | `base` | 24 | 1024 | [arXiv](https://arxiv.org/abs/2312.05187) | [GitHub](https://github.com/facebookresearch/seamless_communication) | [🤗](https://huggingface.co/facebook/w2v-bert-2.0) |
 | `wav2vec2` | `base` | 12 | 768 | [arXiv](https://arxiv.org/abs/2006.11477) | [GitHub](https://github.com/facebookresearch/fairseq/tree/main/examples/wav2vec) | |
 | | `large` | 24 | 1024 | | | |
 | | `xlsr` | 24 | 1024 | [arXiv](https://arxiv.org/abs/2006.13979) | | |
@@ -149,7 +150,9 @@ pip install git+https://github.com/takenori-y/lfeats.git@master
 
 | Model Name | Model Variant | Hop Size [ms] | Dimension | Paper | Source | Model Hub |
 | :--- | :--- | ---: | ---: | :---: | :---: | :---: |
+| `dacvae` | `base` | 40 | 128 | [arXiv](https://arxiv.org/abs/2410.13720) | [GitHub](https://github.com/facebookresearch/dacvae) | [🤗](https://huggingface.co/facebook/dacvae-watermarked) |
 | `higgs-audio` | `v2` | 40 | 1024 | [Blog](https://github.com/Render-AI-Code/cog-higgs-audio/blob/main/tech_blogs/TOKENIZER_BLOG.md) | [GitHub](https://github.com/boson-ai/higgs-audio) | [🤗](https://huggingface.co/eustlb/higgs-audio-v2-tokenizer) |
+| `mimi` | `base` | 80 | 512 | [arXiv](https://arxiv.org/abs/2410.00037) | [GitHub](https://github.com/kyutai-labs/moshi) | [🤗](https://huggingface.co/kyutai/mimi) |
 | `x-codec` | `hubert` | 20 | 1024 | [arXiv](https://arxiv.org/abs/2408.17175) | [GitHub](https://github.com/zhenye234/xcodec) | [🤗](https://huggingface.co/hf-audio/xcodec-hubert-librispeech) |
 | | `wavlm` | 20 | 1024 | | | [🤗](https://huggingface.co/hf-audio/xcodec-wavlm-more-data) |
 
@@ -169,6 +172,13 @@ pip install git+https://github.com/takenori-y/lfeats.git@master
 | | `b4` | 192 | | | |
 | | `b5` | 192 | | | |
 | | `b6` | 192 | | | |
+| `redimnet2` | `b0` | 192 | [arXiv](https://arxiv.org/abs/2603.11841) | [GitHub](https://github.com/PalabraAI/redimnet2) | |
+| | `b1` | 192 | | | |
+| | `b2` | 192 | | | |
+| | `b3` | 192 | | | |
+| | `b4` | 192 | | | |
+| | `b5` | 192 | | | |
+| | `b6` | 192 | | | |
 | `wavlm-sv` | `base` | 512 | [arXiv](https://arxiv.org/abs/2110.13900) | [GitHub](https://github.com/microsoft/unilm/tree/master/wavlm) | [🤗](https://huggingface.co/microsoft/wavlm-base-sv) |
 | | `base+` | 512 | | | [🤗](https://huggingface.co/microsoft/wavlm-base-plus-sv) |
 | `x-vector` | `base` | 512 | [IEEE](https://ieeexplore.ieee.org/document/8461375) | [GitHub](https://github.com/speechbrain/speechbrain) | [🤗](https://huggingface.co/speechbrain/spkrec-xvect-voxceleb) |
@@ -352,10 +362,11 @@ This project is released under the MIT License.
 
 | Repository | License |
 | :--- | :--- |
+| [DAC-VAE](https://github.com/facebookresearch/dacvae) | Apache-2.0 |
 | [fairseq](https://github.com/facebookresearch/fairseq) | MIT |
-| [NeXt_TDNN_ASV](https://github.com/dmlguq456/NeXt_TDNN_ASV) | Apache-2.0 |
+| [NeXt-TDNN](https://github.com/dmlguq456/NeXt_TDNN_ASV) | Apache-2.0 |
 | [R-Spin](https://github.com/vectominist/rspin) | MIT |
 | [S3PRL](https://github.com/s3prl/s3prl) | Apache-2.0 |
 | [SpeechBrain](https://github.com/speechbrain/speechbrain) | Apache-2.0 |
 | [Spin](https://github.com/vectominist/spin) | MIT |
-| [timm](https://github.com/huggingface/pytorch-image-models) | Apache-2.0 |
+| [TIMM](https://github.com/huggingface/pytorch-image-models) | Apache-2.0 |

{lfeats-0.2.2 → lfeats-0.2.3}/README.md

@@ -78,6 +78,7 @@ pip install git+https://github.com/takenori-y/lfeats.git@master
 | `unispeech-sat` | `base` | 12 | 768 | [arXiv](https://arxiv.org/abs/2110.05752) | [GitHub](https://github.com/microsoft/UniSpeech) | [🤗](https://huggingface.co/microsoft/unispeech-sat-base) |
 | | `base+` | 12 | 768 | | | [🤗](https://huggingface.co/microsoft/unispeech-sat-base-plus) |
 | | `large` | 24 | 1024 | | | [🤗](https://huggingface.co/microsoft/unispeech-sat-large) |
+| `w2v-bert2` | `base` | 24 | 1024 | [arXiv](https://arxiv.org/abs/2312.05187) | [GitHub](https://github.com/facebookresearch/seamless_communication) | [🤗](https://huggingface.co/facebook/w2v-bert-2.0) |
 | `wav2vec2` | `base` | 12 | 768 | [arXiv](https://arxiv.org/abs/2006.11477) | [GitHub](https://github.com/facebookresearch/fairseq/tree/main/examples/wav2vec) | |
 | | `large` | 24 | 1024 | | | |
 | | `xlsr` | 24 | 1024 | [arXiv](https://arxiv.org/abs/2006.13979) | | |
@@ -97,7 +98,9 @@ pip install git+https://github.com/takenori-y/lfeats.git@master
 
 | Model Name | Model Variant | Hop Size [ms] | Dimension | Paper | Source | Model Hub |
 | :--- | :--- | ---: | ---: | :---: | :---: | :---: |
+| `dacvae` | `base` | 40 | 128 | [arXiv](https://arxiv.org/abs/2410.13720) | [GitHub](https://github.com/facebookresearch/dacvae) | [🤗](https://huggingface.co/facebook/dacvae-watermarked) |
 | `higgs-audio` | `v2` | 40 | 1024 | [Blog](https://github.com/Render-AI-Code/cog-higgs-audio/blob/main/tech_blogs/TOKENIZER_BLOG.md) | [GitHub](https://github.com/boson-ai/higgs-audio) | [🤗](https://huggingface.co/eustlb/higgs-audio-v2-tokenizer) |
+| `mimi` | `base` | 80 | 512 | [arXiv](https://arxiv.org/abs/2410.00037) | [GitHub](https://github.com/kyutai-labs/moshi) | [🤗](https://huggingface.co/kyutai/mimi) |
 | `x-codec` | `hubert` | 20 | 1024 | [arXiv](https://arxiv.org/abs/2408.17175) | [GitHub](https://github.com/zhenye234/xcodec) | [🤗](https://huggingface.co/hf-audio/xcodec-hubert-librispeech) |
 | | `wavlm` | 20 | 1024 | | | [🤗](https://huggingface.co/hf-audio/xcodec-wavlm-more-data) |
 
@@ -117,6 +120,13 @@ pip install git+https://github.com/takenori-y/lfeats.git@master
 | | `b4` | 192 | | | |
 | | `b5` | 192 | | | |
 | | `b6` | 192 | | | |
+| `redimnet2` | `b0` | 192 | [arXiv](https://arxiv.org/abs/2603.11841) | [GitHub](https://github.com/PalabraAI/redimnet2) | |
+| | `b1` | 192 | | | |
+| | `b2` | 192 | | | |
+| | `b3` | 192 | | | |
+| | `b4` | 192 | | | |
+| | `b5` | 192 | | | |
+| | `b6` | 192 | | | |
 | `wavlm-sv` | `base` | 512 | [arXiv](https://arxiv.org/abs/2110.13900) | [GitHub](https://github.com/microsoft/unilm/tree/master/wavlm) | [🤗](https://huggingface.co/microsoft/wavlm-base-sv) |
 | | `base+` | 512 | | | [🤗](https://huggingface.co/microsoft/wavlm-base-plus-sv) |
 | `x-vector` | `base` | 512 | [IEEE](https://ieeexplore.ieee.org/document/8461375) | [GitHub](https://github.com/speechbrain/speechbrain) | [🤗](https://huggingface.co/speechbrain/spkrec-xvect-voxceleb) |
@@ -300,10 +310,11 @@ This project is released under the MIT License.
 
 | Repository | License |
 | :--- | :--- |
+| [DAC-VAE](https://github.com/facebookresearch/dacvae) | Apache-2.0 |
 | [fairseq](https://github.com/facebookresearch/fairseq) | MIT |
-| [NeXt_TDNN_ASV](https://github.com/dmlguq456/NeXt_TDNN_ASV) | Apache-2.0 |
+| [NeXt-TDNN](https://github.com/dmlguq456/NeXt_TDNN_ASV) | Apache-2.0 |
 | [R-Spin](https://github.com/vectominist/rspin) | MIT |
 | [S3PRL](https://github.com/s3prl/s3prl) | Apache-2.0 |
 | [SpeechBrain](https://github.com/speechbrain/speechbrain) | Apache-2.0 |
 | [Spin](https://github.com/vectominist/spin) | MIT |
-| [timm](https://github.com/huggingface/pytorch-image-models) | Apache-2.0 |
+| [TIMM](https://github.com/huggingface/pytorch-image-models) | Apache-2.0 |

{lfeats-0.2.2 → lfeats-0.2.3}/lfeats/cli.py

@@ -122,6 +122,7 @@ def get_arguments() -> argparse.Namespace:
         "--reduction",
         type=str,
         default="auto",
+        choices=["none", "mean", "auto"],
         help=(
             "The reduction method to apply to the extracted features. Can be 'none', "
             "'mean', or 'auto'."
@@ -183,12 +184,14 @@ def main() -> None:
     # Parse the layers argument.
     if args.layers in ("all", "last"):
         layers = args.layers
-    elif "," in args.layers:
-        layers = [int(layer.strip()) for layer in args.layers.split(",")]
-    elif args.layers.isdigit():
-        layers = int(args.layers)
     else:
-        raise ValueError(f"Invalid layers argument: {args.layers}")
+        try:
+            if "," in args.layers:
+                layers = [int(layer.strip()) for layer in args.layers.split(",")]
+            else:
+                layers = int(args.layers)
+        except ValueError as e:
+            raise ValueError(f"Invalid layers argument: {args.layers}") from e
 
     output_ext = {
         "npz": "npz",

{lfeats-0.2.2 → lfeats-0.2.3}/lfeats/interfaces/extractor.py

@@ -210,8 +210,10 @@ class Extractor:
         if model.granularity != Granularity.FRAME or reduction == "mean":
             raise ValueError("Upsampling is only supported for frame-level features.")
 
-        # Prepare the audio data and validate the upsample factor.
+        # Prepare the audio data at the model's sample rate so that the shifts below
+        # are measured in the same unit as the frame shift.
         audio = create_audio_object(source, sample_rate)
+        audio = self._resample(audio, model.sample_rate)
         B, T = audio.data.shape
         frame_shift = model.frame_shift
         if frame_shift % upsample_factor != 0:
@@ -225,8 +227,10 @@ class Extractor:
         shifted_waveforms = audio.zeros((B * upsample_factor, T))
         for i in range(upsample_factor):
             offset = i * step
-            end = T - offset
-            shifted_waveforms[i::upsample_factor, :end] = audio.data[:, offset:]  # type: ignore
+            if offset < T:
+                shifted_waveforms[i::upsample_factor, : T - offset] = audio.data[  # type: ignore
+                    :, offset:
+                ]
 
         # Extract features from the shifted waveforms.
         features = self._extract(
@@ -236,6 +240,7 @@ class Extractor:
             center=center,
             chunk_length_sec=chunk_length_sec,
             overlap_length_sec=overlap_length_sec,
+            normalize=normalize,
         )
 
         # Interleave the features from the shifted waveforms.
@@ -321,17 +326,16 @@ class Extractor:
         model.load(self.cache_dir, quiet=False)
         if model.chunk_length_sec is not None:
             chunk_length_sec = model.chunk_length_sec
+            if chunk_length_sec <= overlap_length_sec:
+                raise ValueError(
+                    f"overlap_length_sec ({overlap_length_sec}) must be less than the "
+                    f"model's fixed chunk length ({chunk_length_sec})."
+                )
         normalized_layers = self._normalize_layers(layers, model.num_layers + 1)
 
         # Prepare the audio data.
         audio = create_audio_object(source, sample_rate)
-
-        # Resample the audio if needed.
-        if audio.sample_rate != model.sample_rate:
-            resampler = self.resampler_manager.get_resampler(
-                audio.sample_rate, model.sample_rate
-            )
-            audio = resampler.resample(audio)
+        audio = self._resample(audio, model.sample_rate)
         expected_num_frames = self._get_num_frames(audio.length, model.frame_shift)
 
         # Pad the audio if needed.
@@ -403,6 +407,28 @@ class Extractor:
 
         return features
 
+    def _resample(self, audio: Audio, sample_rate: int) -> Audio:
+        """Resample the audio to the specified sample rate if needed.
+
+        Parameters
+        ----------
+        audio : Audio
+            The input audio data.
+
+        sample_rate : int
+            The target sample rate in Hz.
+
+        Returns
+        -------
+        out : Audio
+            The resampled audio data.
+
+        """
+        if audio.sample_rate == sample_rate:
+            return audio
+        resampler = self.resampler_manager.get_resampler(audio.sample_rate, sample_rate)
+        return resampler.resample(audio)
+
     @staticmethod
     def _get_num_frames(length: int, frame_shift: int) -> int:
         """Calculate the number of frames for a given sample length.

{lfeats-0.2.2 → lfeats-0.2.3}/lfeats/models/__init__.py

@@ -4,6 +4,7 @@
 """The module for audio feature extraction models."""
 
 from .contentvec import ContentVecModel
+from .dacvae import DACVAEModel
 from .data2vec import Data2VecModel
 from .data2vec2 import Data2Vec2Model
 from .ecapa_tdnn import EcapaTDNNModel
@@ -12,14 +13,17 @@ from .emotion2vec_plus import Emotion2VecPlusModel
 from .higgs_audio import HiggsAudioTokenizerModel
 from .hubert import HuBERTModel
 from .manager import ModelManager
+from .mimi import MimiModel
 from .next_tdnn import NeXtTDNNModel
 from .r_spin import RSpinModel
 from .r_vector import RVectorModel
 from .redimnet import ReDimNetModel
+from .redimnet2 import ReDimNet2Model
 from .spidr import SpidRModel
 from .spin import SpinModel
 from .sslzip import SSLZipModel
 from .unispeech_sat import UniSpeechSATModel
+from .w2v_bert2 import W2VBert2Model
 from .wav2vec2 import Wav2Vec2Model
 from .wavlm import WavLMModel
 from .wavlm_sv import WavLMSVModel
@@ -29,6 +33,7 @@ from .x_vector import XVectorModel
 
 MODEL_MAP = {
     "contentvec": ContentVecModel,
+    "dacvae": DACVAEModel,
     "data2vec": Data2VecModel,
     "data2vec2": Data2Vec2Model,
     "ecapa-tdnn": EcapaTDNNModel,
@@ -36,14 +41,17 @@ MODEL_MAP = {
     "emotion2vec+": Emotion2VecPlusModel,
     "higgs-audio": HiggsAudioTokenizerModel,
     "hubert": HuBERTModel,
+    "mimi": MimiModel,
     "next-tdnn": NeXtTDNNModel,
     "r-spin": RSpinModel,
     "r-vector": RVectorModel,
     "redimnet": ReDimNetModel,
+    "redimnet2": ReDimNet2Model,
     "spidr": SpidRModel,
     "spin": SpinModel,
     "sslzip": SSLZipModel,
     "unispeech-sat": UniSpeechSATModel,
+    "w2v-bert2": W2VBert2Model,
     "wav2vec2": Wav2Vec2Model,
     "wavlm": WavLMModel,
     "wavlm-sv": WavLMSVModel,

lfeats-0.2.3/lfeats/models/dacvae.py

@@ -0,0 +1,176 @@
+# Copyright (c) 2026 Takenori Yoshimura
+# Released under the MIT License.
+
+"""A module for the DAC-VAE model."""
+
+from enum import Enum
+
+import torch
+from huggingface_hub import hf_hub_download
+
+from ..interfaces.types import Audio, Features
+from ..utils.io import silence_hf_hub
+from ..utils.validation import validate_enum, validate_length
+from .base import TokenLevelFeatureModel
+
+
+class DACVAEVariant(str, Enum):
+    """Enumeration of supported DAC-VAE model variants."""
+
+    BASE = "base"
+
+    @property
+    def model_name(self) -> str:
+        """Return the model name corresponding to the variant.
+
+        Returns
+        -------
+        out : str
+            The model name corresponding to the variant.
+
+        """
+        return "facebook/dacvae-watermarked"
+
+
+class DACVAEModel(TokenLevelFeatureModel):
+    """A class for the DAC-VAE model.
+
+    The decoder of the original model is not constructed because only the encoder is
+    required to extract features. This also means that the watermarking module, which
+    is a part of the decoder, is never used. In addition, the loudness normalization
+    performed in the original inference script is not applied.
+
+    """
+
+    def __init__(self, variant: str | None = None, device: str = "cpu") -> None:
+        """Initialize the DAC-VAE model.
+
+        Parameters
+        ----------
+        variant : str | None, optional
+            The variant of the model to use.
+
+        device : str, optional
+            The device to run the model on (e.g., 'cpu' or 'cuda').
+
+        """
+        super().__init__(variant, device)
+
+        self.variant = validate_enum(variant, DACVAEVariant, DACVAEVariant.BASE)
+        self._model_id = f"dacvae-{self.variant.value}"
+
+    def load(self, model_dir: str, quiet: bool = False) -> None:
+        """Load the model from the specified directory.
+
+        Parameters
+        ----------
+        model_dir : str
+            The directory where the model checkpoint will be stored.
+
+        quiet : bool, optional
+            Whether to suppress output during the loading process.
+
+        Raises
+        ------
+        RuntimeError
+            If the model checkpoint does not contain the expected parameters.
+
+        """
+        if self.model is not None:
+            return
+
+        with silence_hf_hub(quiet):
+            model_path = hf_hub_download(
+                repo_id=self.variant.model_name,
+                filename="weights.pth",
+                repo_type="model",
+                cache_dir=model_dir,
+            )
+
+        checkpoint = torch.load(
+            model_path, map_location=torch.device("cpu"), weights_only=False
+        )
+
+        from lfeats.third_party.dacvae import DACVAE
+
+        self.model = DACVAE(**checkpoint["metadata"]["kwargs"])
+
+        # The checkpoint contains the parameters of the decoder, which are not used.
+        missing_keys, _ = self.model.load_state_dict(
+            checkpoint["state_dict"], strict=False
+        )
+        if missing_keys:
+            raise RuntimeError(f"Missing parameters in the checkpoint: {missing_keys}")
+
+        self.model.eval()
+        self.model.to(self.device)
+
+    def extract_features_impl(self, audio: Audio, layers: list[int]) -> Features:
+        """Extract features from the input audio using the model.
+
+        Parameters
+        ----------
+        audio : Audio
+            The input audio data with shape (B, T).
+
+        layers : list[int]
+            The layer(s) from which to extract features.
+
+        Returns
+        -------
+        out : Features
+            The extracted features.
+
+        Raises
+        ------
+        RuntimeError
+            If the model is not loaded.
+
+        """
+        if self.model is None:
+            raise RuntimeError("Model not loaded. Call 'load' method first.")
+
+        with torch.inference_mode():
+            inputs = audio.tensor.to(self.device)
+            inputs = validate_length(inputs, self.frame_shift // 2 + 1)
+            vectors = self.model.encode(inputs.unsqueeze(1))  # (B, D, N)
+            vectors = vectors.transpose(1, 2)
+            vectors = vectors[:, : audio.length // self.frame_shift]  # due to padding
+
+        return Features(data=vectors, source=self.model_id)
+
+    @property
+    def center_offset(self) -> int:
+        """Get the center offset of the model.
+
+        Returns
+        -------
+        out : int
+            The center offset in samples.
+
+        """
+        return self.frame_shift // 2
+
+    @property
+    def frame_shift(self) -> int:
+        """Get the frame shift of the model.
+
+        Returns
+        -------
+        out : int
+            The frame shift in samples.
+
+        """
+        return int(40.0 * self.sample_rate / 1000)
+
+    @property
+    def sample_rate(self) -> int:
+        """Get the sample rate required by the model.
+
+        Returns
+        -------
+        out : int
+            The sample rate in Hz.
+
+        """
+        return 48000

{lfeats-0.2.2 → lfeats-0.2.3}/lfeats/models/higgs_audio.py

@@ -123,6 +123,22 @@ class HiggsAudioTokenizerModel(TokenLevelFeatureModel):
 
         return Features(data=vectors, source=self.model_id)
 
+    @property
+    def center_offset(self) -> int:
+        """Get the center offset of the model.
+
+        The model concatenates the outputs of the semantic branch (HuBERT) and the
+        acoustic branch (DAC encoder), which are delayed by different amounts. Since
+        no single offset can align both branches, no compensation is applied.
+
+        Returns
+        -------
+        out : int
+            The center offset in samples.
+
+        """
+        return 0
+
     @property
     def frame_shift(self) -> int:
         """Get the frame shift of the model.

lfeats-0.2.3/lfeats/models/mimi.py

@@ -0,0 +1,161 @@
+# Copyright (c) 2026 Takenori Yoshimura
+# Released under the MIT License.
+
+"""A module for the Mimi model."""
+
+from enum import Enum
+from typing import Any
+
+import torch
+
+from ..interfaces.types import Audio, Features
+from ..utils.io import setup_transformers
+from ..utils.validation import validate_enum
+from .base import TokenLevelFeatureModel
+
+
+class MimiVariant(str, Enum):
+    """Enumeration of supported Mimi variants."""
+
+    BASE = "base"
+
+    @property
+    def model_name(self) -> str:
+        """Return the model name corresponding to the variant.
+
+        Returns
+        -------
+        out : str
+            The model name corresponding to the variant.
+
+        """
+        return "kyutai/mimi"
+
+
+class MimiModel(TokenLevelFeatureModel):
+    """A class for the Mimi model."""
+
+    def __init__(self, variant: str | None = None, device: str = "cpu") -> None:
+        """Initialize the Mimi model.
+
+        Parameters
+        ----------
+        variant : str | None, optional
+            The variant of the model to use.
+
+        device : str, optional
+            The device to run the model on (e.g., 'cpu' or 'cuda').
+
+        """
+        super().__init__(variant, device)
+
+        self.variant = validate_enum(variant, MimiVariant, MimiVariant.BASE)
+        self._model_id = f"mimi-{self.variant.value}"
+
+        self.feature_extractor = None
+
+    def load(self, model_dir: str, quiet: bool = False) -> None:
+        """Load the model from the specified directory.
+
+        Parameters
+        ----------
+        model_dir : str
+            The directory where the model checkpoint will be stored.
+
+        quiet : bool, optional
+            Whether to suppress output during the loading process.
+
+        """
+        if self.model is not None:
+            return
+
+        from transformers import AutoFeatureExtractor
+        from transformers import MimiModel as _MimiModel
+
+        with setup_transformers(quiet):
+            self.feature_extractor = AutoFeatureExtractor.from_pretrained(
+                self.variant.model_name, cache_dir=model_dir
+            )
+            self.model = _MimiModel.from_pretrained(
+                self.variant.model_name, cache_dir=model_dir
+            )
+            self.model.eval()
+            self.model.to(self.device)  # type: ignore
+
+    def extract_features_impl(self, audio: Audio, layers: list[int]) -> Features:
+        """Extract features from the input audio using the model.
+
+        Parameters
+        ----------
+        audio : Audio
+            The input audio data with shape (B, T).
+
+        layers : list[int]
+            The layer(s) from which to extract features.
+
+        Returns
+        -------
+        out : Features
+            The extracted features.
+
+        Raises
+        ------
+        RuntimeError
+            If the model is not loaded.
+
+        """
+        if self.feature_extractor is None or self.model is None:
+            raise RuntimeError("Model not loaded. Call 'load' method first.")
+
+        with torch.inference_mode():
+            inputs = self.feature_extractor(
+                raw_audio=[x for x in audio.array],
+                sampling_rate=self.feature_extractor.sampling_rate,
+                return_tensors="pt",
+            ).to(self.device)
+
+            encoder_outputs: Any = self.model.encode(
+                inputs["input_values"], inputs["padding_mask"]
+            )
+            indices = encoder_outputs.audio_codes  # (B, Q, N)
+            vectors = self.model.quantizer.decode(indices)  # (B, D, N)
+            vectors = vectors.transpose(1, 2)
+            vectors = vectors[:, : audio.length // self.frame_shift]  # due to padding
+
+        return Features(data=vectors, source=self.model_id)
+
+    @property
+    def center_offset(self) -> int:
+        """Get the center offset of the model.
+
+        Returns
+        -------
+        out : int
+            The center offset in samples.
+
+        """
+        return self.frame_shift // 2
+
+    @property
+    def frame_shift(self) -> int:
+        """Get the frame shift of the model.
+
+        Returns
+        -------
+        out : int
+            The frame shift in samples.
+
+        """
+        return int(80.0 * self.sample_rate / 1000)
+
+    @property
+    def sample_rate(self) -> int:
+        """Get the sample rate required by the model.
+
+        Returns
+        -------
+        out : int
+            The sample rate in Hz.
+
+        """
+        return 24000