lfeats 0.1.4__tar.gz → 0.2.0__tar.gz

{lfeats-0.1.4 → lfeats-0.2.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: lfeats
-Version: 0.1.4
+Version: 0.2.0
 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
@@ -18,6 +18,7 @@ Classifier: Programming Language :: Python :: 3.12
 Classifier: Programming Language :: Python :: 3.13
 Classifier: Programming Language :: Python :: 3.14
 Requires-Python: >=3.10
+Requires-Dist: filelock>=3.10.0
 Requires-Dist: huggingface-hub>=0.23.0
 Requires-Dist: hydra-core>=1.3.0
 Requires-Dist: hyperpyyaml>=0.0.1

{lfeats-0.1.4 → lfeats-0.2.0}/lfeats/interfaces/extractor.py

@@ -90,6 +90,18 @@ class Extractor:
         """
         self.model_manager.get_model().load(self.cache_dir, quiet)
 
+    def to(self, device: str) -> None:
+        """Move the model to the specified device.
+
+        Parameters
+        ----------
+        device : str
+            The device to move the model to (e.g., 'cpu' or 'cuda').
+
+        """
+        self.model_manager.to(device)
+        self.resampler_manager.to(device)
+
     def __call__(
         self,
         source: np.ndarray | torch.Tensor | Audio,

{lfeats-0.1.4 → lfeats-0.2.0}/lfeats/interfaces/resampler.py

@@ -45,6 +45,17 @@ class Resampler:
                 f"Supported resamplers are: {[k for k in RESAMPLER_MAP.keys()]}"
             ) from e
 
+    def to(self, device: str) -> None:
+        """Move the resampler to the specified device.
+
+        Parameters
+        ----------
+        device : str
+            The device to move the resampler to (e.g., 'cpu' or 'cuda').
+
+        """
+        self.resampler_manager.to(device)
+
     def __call__(
         self,
         source: np.ndarray | torch.Tensor | Audio,

{lfeats-0.1.4 → lfeats-0.2.0}/lfeats/models/base.py

@@ -25,6 +25,7 @@ class BaseModel(ABC):
         """
         self.device = device
 
+        self.model = None
         self._model_id = None  # To be defined in subclasses
 
     @abstractmethod
@@ -42,6 +43,24 @@ class BaseModel(ABC):
         """
         raise NotImplementedError
 
+    def to(self, device: str) -> None:
+        """Move the model to the specified device.
+
+        Parameters
+        ----------
+        device : str
+            The device to move the model to (e.g., 'cpu' or 'cuda').
+
+        """
+        self.device = device
+        if self.model is not None and hasattr(self.model, "to"):
+            self.model.to(device)
+            if hasattr(self.model, "device"):
+                try:
+                    self.model.device = device
+                except AttributeError:
+                    pass
+
     def extract_features(self, audio: Audio, layers: list[int]) -> Features:
         """Extract features from the input audio data.
 

{lfeats-0.1.4 → lfeats-0.2.0}/lfeats/models/contentvec.py

@@ -68,8 +68,6 @@ class ContentVecModel(FrameLevelFeatureModel):
         )
         self._model_id = f"contentvec-{self.variant.value}"
 
-        self.model = None
-
     def load(self, model_dir: str, quiet: bool = False) -> None:
         """Load the model from the specified directory.
 

{lfeats-0.1.4 → lfeats-0.2.0}/lfeats/models/data2vec.py

@@ -58,9 +58,6 @@ class Data2VecModel(FrameLevelFeatureModel):
         self.variant = validate_enum(variant, Data2VecVariant, Data2VecVariant.BASE)
         self._model_id = f"data2vec-{self.variant.value}"
 
-        self.processor = None
-        self.model = None
-
     def load(self, model_dir: str, quiet: bool = False) -> None:
         """Load the model from the specified directory.
 

{lfeats-0.1.4 → lfeats-0.2.0}/lfeats/models/data2vec2.py

@@ -58,9 +58,6 @@ class Data2Vec2Model(FrameLevelFeatureModel):
         self.variant = validate_enum(variant, Data2Vec2Variant, Data2Vec2Variant.BASE)
         self._model_id = f"data2vec2-{self.variant.value}"
 
-        self.processor = None
-        self.model = None
-
     def load(self, model_dir: str, quiet: bool = False) -> None:
         """Load the model from the specified directory.
 

{lfeats-0.1.4 → lfeats-0.2.0}/lfeats/models/ecapa_tdnn.py

@@ -40,8 +40,6 @@ class EcapaTDNNModel(UtteranceLevelFeatureModel):
         self.variant = validate_enum(variant, EcapaTDNNVariant, EcapaTDNNVariant.BASE)
         self._model_id = f"ecapa-tdnn-{self.variant.value}"
 
-        self.model = None
-
     def load(self, model_dir: str, quiet: bool = False) -> None:
         """Load the model from the specified directory.
 
@@ -78,11 +76,11 @@ class EcapaTDNNModel(UtteranceLevelFeatureModel):
             self.model = EncoderClassifier.from_hparams(
                 source="speechbrain/spkrec-ecapa-voxceleb",
                 fetch_config=fetch_config,
+                run_opts={"device": self.device},
             )
             if self.model is None:
                 raise RuntimeError("Failed to load the model.")
             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.

{lfeats-0.1.4 → lfeats-0.2.0}/lfeats/models/emotion2vec.py

@@ -55,8 +55,6 @@ class Emotion2VecModel(FrameLevelFeatureModel):
         )
         self._model_id = f"emotion2vec-{self.variant.value}"
 
-        self.model = None
-
     def load(self, model_dir: str, quiet: bool = False) -> None:
         """Load the model from the specified directory.
 
@@ -78,7 +76,7 @@ class Emotion2VecModel(FrameLevelFeatureModel):
                 repo_id=repo_id,
                 filename=filename,
                 repo_type="model",
-                local_dir=model_dir,
+                cache_dir=model_dir,
             )
 
         setup_third_party_path()

{lfeats-0.1.4 → lfeats-0.2.0}/lfeats/models/emotion2vec_plus.py

@@ -58,8 +58,6 @@ class Emotion2VecPlusModel(FrameLevelFeatureModel):
         )
         self._model_id = f"emotion2vec+-{self.variant.value}"
 
-        self.model = None
-
     def load(self, model_dir: str, quiet: bool = False) -> None:
         """Load the model from the specified directory.
 
@@ -81,13 +79,13 @@ class Emotion2VecPlusModel(FrameLevelFeatureModel):
                 repo_id=repo_id,
                 filename="model.pt",
                 repo_type="model",
-                local_dir=os.path.join(model_dir, sanitize(self.model_id)),
+                cache_dir=os.path.join(model_dir, sanitize(self.model_id)),
             )
             config = hf_hub_download(
                 repo_id=repo_id,
                 filename="config.yaml",
                 repo_type="model",
-                local_dir=os.path.join(model_dir, sanitize(self.model_id)),
+                cache_dir=os.path.join(model_dir, sanitize(self.model_id)),
             )
 
         from hyperpyyaml import load_hyperpyyaml

{lfeats-0.1.4 → lfeats-0.2.0}/lfeats/models/hubert.py

@@ -60,8 +60,6 @@ class HuBERTModel(FrameLevelFeatureModel):
         self.variant = validate_enum(variant, HuBERTVariant, HuBERTVariant.BASE)
         self._model_id = f"hubert-{self.variant.value}"
 
-        self.model = None
-
     def load(self, model_dir: str, quiet: bool = False) -> None:
         """Load the model from the specified directory.
 

{lfeats-0.1.4 → lfeats-0.2.0}/lfeats/models/manager.py

@@ -32,6 +32,19 @@ class ModelManager:
 
         self._cache: dict[str, BaseModel] = {}
 
+    def to(self, device: str) -> None:
+        """Move all models to the specified device.
+
+        Parameters
+        ----------
+        device : str
+            The device to move the models to (e.g., 'cpu' or 'cuda').
+
+        """
+        self.device = device
+        for model in self._cache.values():
+            model.to(device)
+
     def get_model(self) -> BaseModel:
         """Get the model instance.
 

{lfeats-0.1.4 → lfeats-0.2.0}/lfeats/models/next_tdnn.py

@@ -66,8 +66,6 @@ class NeXtTDNNModel(UtteranceLevelFeatureModel):
         self.variant = validate_enum(variant, NeXtTDNNVariant, NeXtTDNNVariant.BASE_V2)
         self._model_id = f"next-tdnn-{self.variant.value}"
 
-        self.model = None
-
     def load(self, model_dir: str, quiet: bool = False) -> None:
         """Load the model from the specified directory.
 

{lfeats-0.1.4 → lfeats-0.2.0}/lfeats/models/r_spin.py

@@ -63,8 +63,6 @@ class RSpinModel(FrameLevelFeatureModel):
         self.variant = validate_enum(variant, RSpinVariant, RSpinVariant.WAVLM_256)
         self._model_id = f"r-spin-{self.variant.value}"
 
-        self.model = None
-
     def load(self, model_dir: str, quiet: bool = False) -> None:
         """Load the model from the specified directory.
 

{lfeats-0.1.4 → lfeats-0.2.0}/lfeats/models/r_vector.py

@@ -40,8 +40,6 @@ class RVectorModel(UtteranceLevelFeatureModel):
         self.variant = validate_enum(variant, RVectorVariant, RVectorVariant.BASE)
         self._model_id = f"r-vector-{self.variant.value}"
 
-        self.model = None
-
     def load(self, model_dir: str, quiet: bool = False) -> None:
         """Load the model from the specified directory.
 
@@ -78,11 +76,11 @@ class RVectorModel(UtteranceLevelFeatureModel):
             self.model = EncoderClassifier.from_hparams(
                 source="speechbrain/spkrec-resnet-voxceleb",
                 fetch_config=fetch_config,
+                run_opts={"device": self.device},
             )
             if self.model is None:
                 raise RuntimeError("Failed to load the model.")
             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.

{lfeats-0.1.4 → lfeats-0.2.0}/lfeats/models/spidr.py

@@ -4,12 +4,11 @@
 """A module for the SpidR model."""
 
 from enum import Enum
-from typing import Any
 
 import torch
 
 from ..interfaces.types import Audio, Features
-from ..utils.io import set_torch_hub_dir
+from ..utils.io import safe_torch_hub_load
 from ..utils.validation import validate_enum
 from .base import FrameLevelFeatureModel
 
@@ -40,8 +39,6 @@ class SpidRModel(FrameLevelFeatureModel):
         self.variant = validate_enum(variant, SpidRVariant, SpidRVariant.BASE)
         self._model_id = f"spidr-{self.variant.value}"
 
-        self.model = None
-
     def load(self, model_dir: str, quiet: bool = False) -> None:
         """Load the model from the specified directory.
 
@@ -57,10 +54,9 @@ class SpidRModel(FrameLevelFeatureModel):
         if self.model is not None:
             return
 
-        with set_torch_hub_dir(model_dir):
-            self.model: Any = torch.hub.load(
-                "facebookresearch/spidr", "spidr_base", verbose=not quiet
-            )
+        self.model = safe_torch_hub_load(
+            "facebookresearch/spidr", "spidr_base", model_dir, quiet=quiet
+        )
         self.model.eval()
         self.model.to(self.device)
 

{lfeats-0.1.4 → lfeats-0.2.0}/lfeats/models/spin.py

@@ -117,8 +117,6 @@ class SpinModel(FrameLevelFeatureModel):
         self.variant = validate_enum(variant, SpinVariant, SpinVariant.HUBERT_256)
         self._model_id = f"spin-{self.variant.value}"
 
-        self.model = None
-
     def load(self, model_dir: str, quiet: bool = False) -> None:
         """Load the model from the specified directory.
 
@@ -139,7 +137,7 @@ class SpinModel(FrameLevelFeatureModel):
                 repo_id="vectominist/spin_ckpt",
                 filename=self.variant.checkpoint_filename,
                 repo_type="dataset",
-                local_dir=model_dir,
+                cache_dir=model_dir,
             )
 
         with lightning_mock_context():
@@ -147,11 +145,12 @@ class SpinModel(FrameLevelFeatureModel):
                 model_path, map_location=torch.device("cpu"), weights_only=False
             )
 
-        from lfeats.third_party.s3prl.util.download import set_dir
+        from lfeats.third_party.s3prl.util.download import set_dir, set_progress
         from lfeats.third_party.spin.model import SpinModel as _SpinModel
         from lfeats.third_party.spin.util import len_to_padding
 
         set_dir(model_dir)
+        set_progress(not quiet)
         self.model = _SpinModel(checkpoint["hyper_parameters"])
         self.model.load_state_dict(checkpoint["state_dict"])
         self.model.eval()

{lfeats-0.1.4 → lfeats-0.2.0}/lfeats/models/sslzip.py

@@ -63,6 +63,19 @@ class SSLZipModel(FrameLevelFeatureModel):
         self.upstream = HuBERTModel(variant=HuBERTVariant.BASE.value, device=device)
         self.model = None
 
+    def to(self, device: str) -> None:
+        """Move the model to the specified device.
+
+        Parameters
+        ----------
+        device : str
+            The device to move the model to (e.g., 'cpu' or 'cuda'). Note that the ONNX
+            model will be loaded on the specified device when calling the 'load' method.
+
+        """
+        self.device = device
+        self.upstream.to(device)
+
     def load(self, model_dir: str, quiet: bool = False) -> None:
         """Load the model from the specified directory.
 
@@ -86,7 +99,7 @@ class SSLZipModel(FrameLevelFeatureModel):
                 repo_id=repo_id,
                 filename=filename,
                 repo_type="model",
-                local_dir=model_dir,
+                cache_dir=model_dir,
             )
 
         import onnxruntime as ort

{lfeats-0.1.4 → lfeats-0.2.0}/lfeats/models/unispeech_sat.py

@@ -56,8 +56,6 @@ class UniSpeechSATModel(FrameLevelFeatureModel):
         )
         self._model_id = f"unispeech-sat-{self.variant.value}"
 
-        self.model = None
-
     def load(self, model_dir: str, quiet: bool = False) -> None:
         """Load the model from the specified directory.
 

{lfeats-0.1.4 → lfeats-0.2.0}/lfeats/models/wav2vec2.py

@@ -64,9 +64,6 @@ class Wav2Vec2Model(FrameLevelFeatureModel):
         self.variant = validate_enum(variant, Wav2Vec2Variant, Wav2Vec2Variant.BASE)
         self._model_id = f"wav2vec2-{self.variant.value}"
 
-        self.processor = None
-        self.model = None
-
     def load(self, model_dir: str, quiet: bool = False) -> None:
         """Load the model from the specified directory.
 

{lfeats-0.1.4 → lfeats-0.2.0}/lfeats/models/wavlm.py

@@ -54,8 +54,6 @@ class WavLMModel(FrameLevelFeatureModel):
         self.variant = validate_enum(variant, WavLMVariant, WavLMVariant.BASE_PLUS)
         self._model_id = f"wavlm-{self.variant.value}"
 
-        self.model = None
-
     def load(self, model_dir: str, quiet: bool = False) -> None:
         """Load the model from the specified directory.
 

{lfeats-0.1.4 → lfeats-0.2.0}/lfeats/models/whisper.py

@@ -58,7 +58,6 @@ class WhisperModel(FrameLevelFeatureModel):
         self._model_id = f"whisper-{self.variant.value}"
 
         self.processor = None
-        self.model = None
 
     def load(self, model_dir: str, quiet: bool = False) -> None:
         """Load the model from the specified directory.

{lfeats-0.1.4 → lfeats-0.2.0}/lfeats/models/x_vector.py

@@ -40,8 +40,6 @@ class XVectorModel(UtteranceLevelFeatureModel):
         self.variant = validate_enum(variant, XVectorVariant, XVectorVariant.BASE)
         self._model_id = f"x-vector-{self.variant.value}"
 
-        self.model = None
-
     def load(self, model_dir: str, quiet: bool = False) -> None:
         """Load the model from the specified directory.
 
@@ -78,11 +76,11 @@ class XVectorModel(UtteranceLevelFeatureModel):
             self.model = EncoderClassifier.from_hparams(
                 source="speechbrain/spkrec-xvect-voxceleb",
                 fetch_config=fetch_config,
+                run_opts={"device": self.device},
             )
             if self.model is None:
                 raise RuntimeError("Failed to load the model.")
             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.

{lfeats-0.1.4 → lfeats-0.2.0}/lfeats/resamplers/base.py

@@ -35,6 +35,17 @@ class BaseResampler(ABC):
         self.dst_rate = dst_rate
         self.device = device
 
+    def to(self, device: str) -> None:
+        """Move the resampler to the specified device.
+
+        Parameters
+        ----------
+        device : str
+            The device to move the resampler to (e.g., 'cpu' or 'cuda').
+
+        """
+        self.device = device
+
     def resample(self, audio: Audio) -> Audio:
         """Resample the given audio to the target sample rate.
 

{lfeats-0.1.4 → lfeats-0.2.0}/lfeats/resamplers/lilfilter.py

@@ -56,7 +56,19 @@ class LilFilterResampler(BaseResampler):
             dst_rate,
             dtype=torch.float32,
         )
-        self.resampler.weights = self.resampler.weights.to(self.device)
+        self.to(device)
+
+    def to(self, device: str) -> None:
+        """Move the resampler to the specified device.
+
+        Parameters
+        ----------
+        device : str
+            The device to move the resampler to (e.g., 'cpu' or 'cuda').
+
+        """
+        super().to(device)
+        self.resampler.weights = self.resampler.weights.to(device)
 
     def resample_impl(self, audio: Audio) -> Audio:
         """Resample the given audio to the target sample rate.

{lfeats-0.1.4 → lfeats-0.2.0}/lfeats/resamplers/manager.py

@@ -32,6 +32,19 @@ class ResamplerManager:
 
         self._cache: dict[tuple[int, int], BaseResampler] = {}
 
+    def to(self, device: str) -> None:
+        """Move all resamplers to the specified device.
+
+        Parameters
+        ----------
+        device : str
+            The device to move the resamplers to (e.g., 'cpu' or 'cuda').
+
+        """
+        self.device = device
+        for resampler in self._cache.values():
+            resampler.to(device)
+
     def get_resampler(self, src_rate: int, dst_rate: int) -> BaseResampler:
         """Get the resampler instance.
 

{lfeats-0.1.4 → lfeats-0.2.0}/lfeats/resamplers/torchaudio.py

@@ -73,7 +73,20 @@ class TorchAudioResampler(BaseResampler):
 
         self.resampler = torchaudio.transforms.Resample(
             orig_freq=src_rate, new_freq=dst_rate, dtype=torch.float32, **params
-        ).to(device)
+        )
+        self.to(device)
+
+    def to(self, device: str) -> None:
+        """Move the resampler to the specified device.
+
+        Parameters
+        ----------
+        device : str
+            The device to move the resampler to (e.g., 'cpu' or 'cuda').
+
+        """
+        super().to(device)
+        self.resampler.to(device)
 
     def resample_impl(self, audio: Audio) -> Audio:
         """Resample the given audio to the target sample rate.

{lfeats-0.1.4 → lfeats-0.2.0}/lfeats/third_party/s3prl/util/download.py

@@ -24,10 +24,13 @@ logger = logging.getLogger(__name__)
 
 
 _download_dir = Path.home() / ".cache" / "s3prl" / "download"
+_progress = True
 
 __all__ = [
     "get_dir",
     "set_dir",
+    "get_progress",
+    "set_progress",
     "download",
     "urls_to_filepaths",
 ]
@@ -43,6 +46,15 @@ def set_dir(d):
     _download_dir = Path(d)
 
 
+def get_progress():
+    return _progress
+
+
+def set_progress(progress):
+    global _progress
+    _progress = progress
+
+
 def _download_url_to_file(url, dst, hash_prefix=None, progress=True):
     """
     This function is not thread-safe. Please ensure only a single
@@ -69,8 +81,9 @@ def _download_url_to_file(url, dst, hash_prefix=None, progress=True):
         if hash_prefix is not None:
             sha256 = hashlib.sha256()
 
-        tqdm.write(f"Downloading: {url}", file=sys.stderr)
-        tqdm.write(f"Destination: {dst}", file=sys.stderr)
+        if progress:
+            tqdm.write(f"Downloading: {url}", file=sys.stderr)
+            tqdm.write(f"Destination: {dst}", file=sys.stderr)
         with tqdm(
             total=file_size,
             disable=not progress,
@@ -120,12 +133,13 @@ def _download_url_to_file_requests(url, dst, hash_prefix=None, progress=True):
         if hash_prefix is not None:
             sha256 = hashlib.sha256()
 
-        tqdm.write(
-            f"urllib.Request method failed. Trying using another method...",
-            file=sys.stderr,
-        )
-        tqdm.write(f"Downloading: {url}", file=sys.stderr)
-        tqdm.write(f"Destination: {dst}", file=sys.stderr)
+        if progress:
+            tqdm.write(
+                f"urllib.Request method failed. Trying using another method...",
+                file=sys.stderr,
+            )
+            tqdm.write(f"Downloading: {url}", file=sys.stderr)
+            tqdm.write(f"Destination: {dst}", file=sys.stderr)
         with tqdm(
             total=file_size,
             disable=not progress,
@@ -176,9 +190,9 @@ def _download(filepath: Path, url, refresh: bool, new_enough_secs: float = 2.0):
             refresh and (time.time() - os.path.getmtime(filepath)) > new_enough_secs
         ):
             try:
-                _download_url_to_file(url, filepath)
+                _download_url_to_file(url, filepath, progress=_progress)
             except:
-                _download_url_to_file_requests(url, filepath)
+                _download_url_to_file_requests(url, filepath, progress=_progress)
 
     logger.info(f"Using URL's local file: {filepath}")
 

{lfeats-0.1.4 → lfeats-0.2.0}/lfeats/third_party/speechbrain/utils/distributed.py

@@ -8,7 +8,7 @@ Authors:
 """
 
 # import datetime
-# import os
+import os
 from functools import wraps
 from typing import Optional
 

{lfeats-0.1.4 → lfeats-0.2.0}/lfeats/utils/io.py

@@ -4,13 +4,16 @@
 """I/O utilities."""
 
 import logging
+import os
 from collections.abc import Generator
 from contextlib import contextmanager
 from importlib.metadata import PackageNotFoundError, version
+from typing import Any
 
 import soundfile as sf
 import torch
 import torchaudio
+from filelock import FileLock
 
 HF_HTTP_LOGGER = "huggingface_hub.utils._http"
 
@@ -108,14 +111,50 @@ def download_file(url: str, download_dir: str, quiet: bool = False) -> str:
     """
     from parfive import Downloader
 
-    dl = Downloader(progress=not quiet)
+    basename = os.path.basename(url)
+    lock_path = os.path.join(download_dir, f"{basename}.lock")
 
-    dl.enqueue_file(url, path=download_dir)
+    with FileLock(lock_path):
+        dl = Downloader(progress=not quiet)
 
-    results = dl.download()
-    if len(results) == 0:
-        raise RuntimeError(f"Failed to download file: {results.errors}")
-    return results[0]
+        dl.enqueue_file(url, path=download_dir)
+
+        results = dl.download()
+        if len(results) == 0:
+            raise RuntimeError(f"Failed to download file: {results.errors}")
+        return results[0]
+
+
+def safe_torch_hub_load(
+    repo_or_dir: str, model: str, download_dir: str, quiet: bool = False
+) -> Any:
+    """Safely load a model from PyTorch Hub.
+
+    Parameters
+    ----------
+    repo_or_dir : str
+        The repository or directory to load the model from.
+
+    model : str
+        The name of the model defined in the repository.
+
+    download_dir : str
+        The directory to use for caching the downloaded model files.
+
+    quiet : bool, optional
+        Whether to suppress output during the loading process.
+
+    Returns
+    -------
+    out : Any
+        The loaded object.
+
+    """
+    lock_path = os.path.join(download_dir, f"{model}.lock")
+
+    with FileLock(lock_path):
+        with set_torch_hub_dir(download_dir):
+            return torch.hub.load(repo_or_dir, model, verbose=not quiet)
 
 
 def load_audio(path: str) -> tuple[torch.Tensor, int]: