Skip to content

API Reference

The API reference is generated from the package source.

For normal inference, start with load_model(). Use the runtime classes directly when the application manages ONNX files itself.

Models

available_models 

available_models() -> tuple[str, ...]

Return all supported model identifiers.

get_model_spec 

get_model_spec(model_name: str) -> ModelSpec

Return metadata for a supported model.

Parameters:

  • model_name (str) –

    Public model identifier.

Returns:

  • ModelSpec –

    Metadata for the requested model.

Raises:

  • ValueError –

    If the model identifier is unknown.

download_model 

download_model(
    model_name: str,
    cache_dir: str | Path | None = None,
    *,
    force_download: bool = False,
) -> Path

Download and verify an ONNX model from the GitHub release.

Parameters:

  • model_name (str) –

    Public model identifier.

  • cache_dir (str | Path | None, default: None ) –

    Optional model cache directory.

  • force_download (bool, default: False ) –

    Download the manifest and model again when true.

Returns:

  • Path –

    Path to the verified ONNX model.

Raises:

  • RuntimeError –

    If the manifest is invalid or checksum verification fails.

  • ValueError –

    If the model identifier is unknown.

ModelSpec dataclass 

ModelSpec(
    name: str,
    model_name: str,
    dataset: str,
    train_type: str,
)

Describe a released ReDimNet2 model.

Attributes:

  • name (str) –

    Public model identifier used by the package.

  • model_name (str) –

    Upstream ReDimNet2 architecture name.

  • dataset (str) –

    Upstream training dataset identifier.

  • train_type (str) –

    Upstream training type.

filename property 

filename: str

Return the ONNX release filename.

source_filename property 

source_filename: str

Return the upstream PyTorch checkpoint filename.

source_url property 

source_url: str

Return the upstream PyTorch checkpoint URL.

ONNX Runtime

load_model 

load_model(
    model_name: str,
    *,
    model_cache_dir: str | Path | None = None,
    force_download: bool = False,
    providers: Sequence[Provider] | None = None,
    session_options: Any | None = None,
) -> ReDimNet2

Download a model and create a provider-neutral runtime.

Parameters:

  • model_name (str) –

    Public model identifier.

  • model_cache_dir (str | Path | None, default: None ) –

    Optional model download cache directory.

  • force_download (bool, default: False ) –

    Download the model again when true.

  • providers (Sequence[Provider] | None, default: None ) –

    ONNX Runtime providers and optional provider settings.

  • session_options (Any | None, default: None ) –

    Optional ONNX Runtime session options.

Returns:

  • ReDimNet2 –

    A ready-to-use ReDimNet2 runtime.

ReDimNet2 

ReDimNet2(
    model_path: str | Path,
    *,
    providers: Sequence[Provider] | None = None,
    session_options: Any | None = None,
)

Run a ReDimNet2 ONNX model with any ONNX Runtime provider.

Create an inference session for a local ONNX model.

Parameters:

  • model_path (str | Path) –

    Path to a ReDimNet2 ONNX model.

  • providers (Sequence[Provider] | None, default: None ) –

    ONNX Runtime providers and optional provider settings.

  • session_options (Any | None, default: None ) –

    Optional ONNX Runtime session options.

Raises:

  • FileNotFoundError –

    If the model file does not exist.

  • ModuleNotFoundError –

    If ONNX Runtime is not installed.

  • RuntimeError –

    If the model input or output contract is invalid.

Methods:

  • embed_features –

    Generate an embedding from mel features.

  • embed –

    Generate an embedding from a raw 16 kHz waveform.

embed_features 

embed_features(features: ndarray) -> ndarray

Generate an embedding from mel features.

Parameters:

  • features (ndarray) –

    FP32 mel features shaped [1, 1, 72, T]. T must be divisible by four and within the 1-30 second profile.

Returns:

  • ndarray –

    An L2-normalized FP32 embedding shaped [1, 192].

Raises:

  • ValueError –

    If the input shape or time dimension is invalid.

  • RuntimeError –

    If the model returns an invalid embedding.

embed 

embed(waveforms: Any) -> ndarray

Generate an embedding from a raw 16 kHz waveform.

Parameters:

  • waveforms (Any) –

    Mono audio shaped [samples] or [1, samples].

Returns:

  • ndarray –

    An L2-normalized FP32 embedding shaped [1, 192].

Raises:

  • ModuleNotFoundError –

    If the optional waveform extra is not installed.

  • ValueError –

    If the waveform duration or shape is invalid.

TensorRT

load_tensorrt_model 

load_tensorrt_model(
    model_name: str,
    *,
    model_cache_dir: str | Path | None = None,
    engine_cache_dir: str | Path | None = None,
    force_download: bool = False,
    device_id: int = 0,
) -> TensorRTReDimNet2

Download a model and create a TensorRT runtime.

Parameters:

  • model_name (str) –

    Public model identifier.

  • model_cache_dir (str | Path | None, default: None ) –

    Optional model download cache directory.

  • engine_cache_dir (str | Path | None, default: None ) –

    Optional TensorRT engine cache directory.

  • force_download (bool, default: False ) –

    Download the model again when true.

  • device_id (int, default: 0 ) –

    CUDA device index.

Returns:

TensorRTReDimNet2 

TensorRTReDimNet2(
    model_path: str | Path,
    *,
    engine_cache_dir: str | Path | None = None,
    device_id: int = 0,
)

Run ReDimNet2 with FP32 preprocessing/output normalization and a TensorRT FP16 backend.

Create a TensorRT runtime for a local ONNX model.

Parameters:

  • model_path (str | Path) –

    Path to a ReDimNet2 ONNX model.

  • engine_cache_dir (str | Path | None, default: None ) –

    Optional TensorRT engine cache directory.

  • device_id (int, default: 0 ) –

    CUDA device index.

Raises:

  • FileNotFoundError –

    If the model file does not exist.

  • RuntimeError –

    If CUDA or the required ONNX Runtime providers are unavailable.

Methods:

  • preprocess –

    Convert a raw 16 kHz waveform to aligned FP32 mel features.

  • embed_features –

    Generate an embedding from CUDA mel features.

  • embed –

    Generate an embedding from a raw 16 kHz waveform.

preprocess 

preprocess(waveforms: Tensor) -> Tensor

Convert a raw 16 kHz waveform to aligned FP32 mel features.

Parameters:

  • waveforms (Tensor) –

    Mono audio shaped [samples] or [1, samples].

Returns:

  • Tensor –

    FP32 mel features on the configured CUDA device.

embed_features 

embed_features(features: Tensor) -> Tensor

Generate an embedding from CUDA mel features.

Parameters:

  • features (Tensor) –

    FP32 CUDA mel features shaped [1, 1, 72, T].

Returns:

  • Tensor –

    An L2-normalized FP32 CUDA embedding shaped [1, 192].

Raises:

  • ValueError –

    If the input device, dtype, or shape is invalid.

embed 

embed(waveforms: Tensor) -> Tensor

Generate an embedding from a raw 16 kHz waveform.

Parameters:

  • waveforms (Tensor) –

    Mono audio shaped [samples] or [1, samples].

Returns:

  • Tensor –

    An L2-normalized FP32 CUDA embedding shaped [1, 192].