# MIT License # # Copyright (c) 2026 audio-embeddings contributors # # Permission is hereby granted, free of charge, to any person obtaining a copy # of this software and associated documentation files (the "Software"), to deal # in the Software without restriction, including without limitation the rights # to use, copy, modify, merge, publish, distribute, sublicense, and/or sell # copies of the Software, and to permit persons to whom the Software is # furnished to do so, subject to the following conditions: # # The above copyright notice and this permission notice shall be included in all # copies or substantial portions of the Software. # # THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR # IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, # FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE # AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER # LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, # OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE # SOFTWARE. from typing import Optional import torch import torch.nn as nn import torchaudio class Spectrogram(nn.Module): """ Mel-frequency audio representation with optional temporal derivatives. Args: sample_rate (int): Sample rate of the audio. n_fft (int): Size of FFT. win_length (Optional[int]): Window length. Defaults to n_fft. win_length_ms (Optional[float]): Window length in milliseconds. Overrides win_length if provided. hop_length (Optional[int]): Hop length. Defaults to win_length // 2. hop_length_ms (Optional[float]): Hop length in milliseconds. Overrides hop_length if provided. n_mels (int): Number of mel filterbanks. f_min (float): Minimum frequency. f_max (Optional[float]): Maximum frequency. power (float): Power of the magnitude. representation (str): ``log_mel`` or phase-aware ``complex_mel``. complex_log_dynamic_range_db (float): Retained complex-mel log range. add_delta (bool): Append one temporal derivative per base channel. add_delta_delta (bool): Append a second derivative per base channel. delta_win_length (int): Odd regression window used for derivatives. """ def __init__( self, sample_rate: int = 32000, n_fft: int = 4096, win_length: Optional[int] = None, win_length_ms: Optional[float] = None, hop_length: Optional[int] = None, hop_length_ms: Optional[float] = None, n_mels: int = 128, f_min: float = 0.0, f_max: Optional[float] = None, power: float = 2.0, representation: str = "log_mel", complex_log_dynamic_range_db: float = 80.0, add_delta: bool = False, add_delta_delta: bool = False, delta_win_length: int = 5, ): super().__init__() if win_length is None: if win_length_ms is None: win_length = n_fft else: win_length = int(sample_rate * win_length_ms / 1000) if hop_length is None: if hop_length_ms is None: hop_length = win_length // 2 else: hop_length = int(sample_rate * hop_length_ms / 1000) representation = representation.strip().lower().replace("-", "_") if representation not in {"log_mel", "complex_mel"}: raise ValueError( "representation must be 'log_mel' or 'complex_mel', " f"got {representation!r}" ) if complex_log_dynamic_range_db <= 0.0: raise ValueError( "complex_log_dynamic_range_db must be positive, " f"got {complex_log_dynamic_range_db}" ) if delta_win_length < 3 or delta_win_length % 2 == 0: raise ValueError( f"delta_win_length must be an odd integer >= 3, got {delta_win_length}" ) self.representation = representation self.complex_log_dynamic_range_db = float(complex_log_dynamic_range_db) self.add_delta = bool(add_delta) self.add_delta_delta = bool(add_delta_delta) self.delta_win_length = int(delta_win_length) self.base_output_channels = 1 if representation == "log_mel" else 2 derivative_orders = int(self.add_delta) + int(self.add_delta_delta) self.output_channels = self.base_output_channels * (1 + derivative_orders) if representation == "log_mel": self.mel_spec = torchaudio.transforms.MelSpectrogram( sample_rate=sample_rate, n_fft=n_fft, win_length=win_length, hop_length=hop_length, n_mels=n_mels, f_min=f_min, f_max=f_max, power=power, normalized=True, ) self.amplitude_to_db = torchaudio.transforms.AmplitudeToDB() else: # Keep the historical ``mel_spec`` attribute for sample-rate and # hop-length discovery in callbacks and HEAR adapters. self.mel_spec = torchaudio.transforms.Spectrogram( n_fft=n_fft, win_length=win_length, hop_length=hop_length, power=None, normalized=True, ) self.mel_spec.sample_rate = sample_rate self.amplitude_to_db = None mel_fb = torchaudio.functional.melscale_fbanks( n_freqs=n_fft // 2 + 1, f_min=f_min, f_max=float(sample_rate / 2 if f_max is None else f_max), n_mels=n_mels, sample_rate=sample_rate, norm=None, mel_scale="htk", ) self.register_buffer("mel_fb", mel_fb) def _complex_mel(self, x: torch.Tensor) -> torch.Tensor: if x.shape[1] != 1: raise ValueError( "complex_mel expects mono waveform input [B, 1, T], " f"got {tuple(x.shape)}" ) complex_spec = self.mel_spec(x[:, 0]) mel_fb = self.mel_fb.to(dtype=complex_spec.dtype) complex_mel = torch.matmul( complex_spec.transpose(-1, -2), mel_fb, ).transpose(-1, -2) magnitude = complex_mel.abs() eps = torch.finfo(magnitude.dtype).eps reference = magnitude.amax(dim=(-2, -1), keepdim=True) relative_magnitude = magnitude / reference.clamp_min(eps) floor_ratio = 10.0 ** (-self.complex_log_dynamic_range_db / 20.0) log_magnitude_db = 20.0 * torch.log10(relative_magnitude.clamp_min(floor_ratio)) log_radius = ( log_magnitude_db + self.complex_log_dynamic_range_db ) / self.complex_log_dynamic_range_db log_radius = log_radius.clamp_(0.0, 1.0) log_radius = torch.where( reference > eps, log_radius, torch.zeros_like(log_radius) ) unit_phase = complex_mel / magnitude.clamp_min(eps) compressed = unit_phase * log_radius return torch.stack([compressed.real, compressed.imag], dim=1) def _append_deltas(self, spec: torch.Tensor) -> torch.Tensor: if not self.add_delta and not self.add_delta_delta: return spec delta = torchaudio.functional.compute_deltas( spec, win_length=self.delta_win_length, mode="replicate", ) channels = [spec] if self.add_delta: channels.append(delta) if self.add_delta_delta: channels.append( torchaudio.functional.compute_deltas( delta, win_length=self.delta_win_length, mode="replicate", ) ) return torch.cat(channels, dim=1) def forward(self, x: torch.Tensor) -> torch.Tensor: """ Forward pass. Args: x (torch.Tensor): Input waveform [B, C, T] or [B, T]. Returns: torch.Tensor: Mel-frequency features [B, C, F, T]. """ if x.ndim == 2: x = x.unsqueeze(1) if x.ndim != 3: raise ValueError(f"Expected waveform [B, C, T], got {tuple(x.shape)}") if self.representation == "complex_mel": spec = self._complex_mel(x) else: spec = self.mel_spec(x) spec = self.amplitude_to_db(spec) return self._append_deltas(spec)