torch-l1-snr 0.0.1__py3-none-any.whl

torch_l1_snr-0.0.1.dist-info/METADATA

@@ -0,0 +1,201 @@
+Metadata-Version: 2.4
+Name: torch-l1-snr
+Version: 0.0.1
+Summary: L1-SNR loss functions for audio source separation in PyTorch
+Home-page: https://github.com/crlandsc/torch-l1-snr
+Author: Christopher Landscaping
+Author-email: crlandschoot@gmail.com
+License: MIT
+Classifier: Intended Audience :: Developers
+Classifier: Intended Audience :: Science/Research
+Classifier: License :: OSI Approved :: MIT License
+Classifier: Programming Language :: Python :: 3
+Classifier: Programming Language :: Python :: 3.8
+Classifier: Programming Language :: Python :: 3.9
+Classifier: Programming Language :: Python :: 3.10
+Classifier: Programming Language :: Python :: 3.11
+Classifier: Programming Language :: Python :: 3.12
+Classifier: Operating System :: OS Independent
+Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
+Classifier: Topic :: Multimedia :: Sound/Audio :: Analysis
+Requires-Python: >=3.8
+Description-Content-Type: text/markdown
+License-File: LICENSE
+Requires-Dist: torch
+Requires-Dist: torchaudio
+Requires-Dist: numpy>=1.21.0
+Dynamic: license-file
+
+![torch-l1-snr-logo](https://raw.githubusercontent.com/crlandsc/torch-l1-snr/main/images/logo.png) -->
+
+# NOTE: Repo is currently a work-in-progress and not ready for installation & use.
+
+[![LICENSE](https://img.shields.io/github/license/crlandsc/torch-l1snr)](https://github.com/crlandsc/torch-l1snr/blob/main/LICENSE) [![GitHub Repo stars](https://img.shields.io/github/stars/crlandsc/torch-l1snr)](https://github.com/crlandsc/torch-l1snr/stargazers)
+
+# torch-l1-snr
+
+A PyTorch implementation of L1-based Signal-to-Noise Ratio (SNR) loss functions for audio source separation. This package provides implementations and novel extensions based on concepts from recent academic papers, offering flexible and robust loss functions that can be easily integrated into any PyTorch-based audio separation pipeline.
+
+The core `L1SNRLoss` is based on the loss function described in [1], while `L1SNRDBLoss` and `STFTL1SNRDBLoss` are extensions of the adaptive level-matching regularization technique proposed in [2].
+
+## Features
+
+- **Time-Domain L1SNR Loss**: A basic, time-domain L1-SNR loss, based on [1].
+- **Regularized Time-Domain L1SNRDBLoss**: An extension of the L1SNR loss with adaptive level-matching regularization from [2], plus an optional L1 loss component.
+- **Multi-Resolution STFT L1SNRDBLoss**: A spectrogram-domain version of the loss from [2], calculated over multiple STFT resolutions.
+- **Modular Stem-based Loss**: A wrapper that combines time and spectrogram domain losses and can be configured to run on specific stems.
+- **Efficient & Robust**: Includes optimizations for pure L1 loss calculation and robust handling of `NaN`/`inf` values and short audio segments.
+
+## Installation
+
+<!-- Add PyPI badges once the package is published -->
+<!-- [![PyPI - Python Version](https://img.shields.io/pypi/pyversions/torch-l1snr)](https://pypi.org/project/torch-l1snr/) -->
+<!-- [![PyPI - Version](https://img.shields.io/pypi/v/torch-l1snr)](https://pypi.org/project/torch-l1snr/) -->
+<!-- [![Number of downloads from PyPI per month](https://img.shields.io/pypi/dm/torch-l1snr)](https://pypi.org/project/torch-l1snr/) -->
+
+You can install the package directly from GitHub:
+
+```bash
+pip install git+https://github.com/crlandsc/torch-l1snr.git
+```
+
+Or, you can clone the repository and install it in editable mode for development:
+
+```bash
+git clone https://github.com/crlandsc/torch-l1snr.git
+cd torch-l1snr
+pip install -e .
+```
+
+## Dependencies
+
+- [PyTorch](https://pytorch.org/)
+- [torchaudio](https://pytorch.org/audio/stable/index.html)
+
+## Supported Tensor Shapes
+
+All loss functions in this package (`L1SNRLoss`, `L1SNRDBLoss`, `STFTL1SNRDBLoss`, and `MultiL1SNRDBLoss`) accept standard audio tensors of shape `(batch, samples)` or `(batch, channels, samples)`. For 3D tensors, the channel and sample dimensions are flattened before the time-domain losses are calculated. For the spectrogram-domain loss, a separate STFT is computed for each channel.
+
+## Usage
+
+The loss functions can be imported directly from the `torch_l1snr` package.
+
+### Example: `L1SNRDBLoss` (Time Domain)
+
+```python
+import torch
+from torch_l1snr import L1SNRDBLoss
+
+# Create dummy audio signals
+estimates = torch.randn(4, 32000)  # Batch of 4, 32000 samples
+actuals = torch.randn(4, 32000)
+
+# Initialize the loss function
+# l1_weight=0.1 blends L1SNR with 10% L1 loss
+loss_fn = L1SNRDBLoss(l1_weight=0.1)
+
+# Calculate loss
+loss = loss_fn(estimates, actuals)
+loss.backward()
+
+print(f"L1SNRDBLoss: {loss.item()}")
+```
+
+### Example: `STFTL1SNRDBLoss` (Spectrogram Domain)
+
+```python
+import torch
+from torch_l1snr import STFTL1SNRDBLoss
+
+# Create dummy audio signals
+estimates = torch.randn(4, 32000)
+actuals = torch.randn(4, 32000)
+
+# Initialize the loss function
+# Uses multiple STFT resolutions by default
+loss_fn = STFTL1SNRDBLoss(l1_weight=0.0) # Pure L1SNR + Regularization
+
+# Calculate loss
+loss = loss_fn(estimates, actuals)
+loss.backward()
+
+print(f"STFTL1SNRDBLoss: {loss.item()}")
+```
+
+### Example: `MultiL1SNRDBLoss` for a Combined Time+Spectrogram Loss
+
+This loss combines the time-domain and spectrogram-domain losses into a single, weighted objective function.
+
+```python
+import torch
+from torch_l1snr import MultiL1SNRDBLoss
+
+# Create dummy audio signals
+# Shape: (batch, channels, samples)
+estimates = torch.randn(2, 2, 44100) # Batch of 2, stereo
+actuals = torch.randn(2, 2, 44100)
+
+# --- Configuration ---
+loss_fn = MultiL1SNRDBLoss(
+    weight=1.0,               # Overall weight for this loss
+    spec_weight=0.7,          # 70% spectrogram loss, 30% time-domain loss
+    l1_weight=0.1,           # Use 10% L1, 90% L1SNR+Reg
+)
+loss = loss_fn(estimates, actuals)
+print(f"Multi-domain Loss: {loss.item()}")
+```
+
+## Motivation
+
+The goal of these loss functions is to provide a perceptually-informed and robust alternative to common audio losses like L1, L2 (MSE), and SI-SDR for training audio source separation models.
+
+- **Robustness**: The L1 norm is less sensitive to large outliers than the L2 norm, making it more suitable for audio signals which can have sharp transients.
+- **Perceptual Relevance**: The loss is scaled to decibels (dB), which more closely aligns with human perception of loudness.
+- **Adaptive Regularization**: Prevents the model from collapsing to silent outputs by penalizing mismatches in the overall loudness (dBRMS) between the estimate and the target.
+
+#### Level-Matching Regularization
+
+A key feature of `L1SNRDBLoss` is the adaptive regularization term, as described in [2]. This component calculates the difference in decibel-scaled root-mean-square (dBRMS) levels between the estimated and actual signals. An adaptive weight (`lambda`) is applied to this difference, which increases when the model incorrectly silences a non-silent target. This encourages the model to learn the correct output level and specifically avoids the model collapsing to a trivial silent solution when uncertain.
+
+#### Multi-Resolution Spectrogram Analysis
+
+The `STFTL1SNRDBLoss` module applies the L1SNRDB loss across multiple time-frequency resolutions. By analyzing the signal with different STFT window sizes and hop lengths, the loss function can capture a wider range of artifacts—from short, transient errors to longer, tonal discrepancies. This provides a more comprehensive error signal to the model during training.
+
+#### "All-or-Nothing" Behavior and `l1_weight`
+
+A characteristic of SNR-style losses is that they encourage the model to make definitive, "all-or-nothing" separation decisions. This can be highly effective for well-defined sources, as it pushes the model to be confident in its estimations. However, this can also lead to "confident errors," where the model completely removes a signal component it should have kept.
+
+While the Level-Matching Regularization prevents a *total collapse to silence*, it does not by itself solve this issue of overly confident, hard-boundary separation. To provide a tunable solution, this implementation introduces a novel `l1_weight` hyperparameter. This allows you to create a hybrid loss, blending the decisive L1SNR objective with a standard L1 loss to soften its "all-or-nothing"-style behavior and allow for more nuanced separation.
+
+-   `l1_weight=0.0` (Default): Pure L1SNR (+ regularization).
+-   `l1_weight=1.0`: Pure L1 loss.
+-   `0.0 < l1_weight < 1.0`: A weighted combination of the two.
+
+The implementation is optimized for efficiency: if `l1_weight` is `0.0` or `1.0`, the unused loss component is not computed, saving computational resources.
+
+**Note on Gradient Balancing:** When blending losses (`0.0 < l1_weight < 1.0`), you may need to tune `l1_scale_time` and `l1_scale_spec`. This is to ensure the gradients of the L1 and L1SNR components are balanced, which is crucial for stable training. The default values provide a reasonable starting point, but monitoring the loss components is recommended to ensure they are scaled appropriately.
+
+## Limitations
+
+- The L1SNR loss is not scale-invariant. Unlike SI-SNR, it requires the model's output to be correctly scaled relative to the target.
+- While the dB scaling and regularization are psychoacoustically motivated, the loss does not model more complex perceptual phenomena like auditory masking.
+
+## Contributing
+
+Contributions are welcome! Please open an issue or submit a pull request if you have any improvements or new features to suggest.
+
+## License
+
+This project is licensed under the MIT License - see the [LICENSE](LICENSE) file for details.
+
+## Acknowledgments
+
+The loss functions implemented here are based on the work of the authors of the referenced papers.
+
+## References
+
+[1] K. N. Watcharasupat, C.-W. Wu, Y. Ding, I. Orife, A. J. Hipple, P. A. Williams, S. Kramer, A. Lerch, and W. Wolcott, "A Generalized Bandsplit Neural Network for Cinematic Audio Source Separation," IEEE Open Journal of Signal Processing, 2023. (arXiv:2309.02539)
+
+[2] K. N. Watcharasupat and A. Lerch, "Separate This, and All of these Things Around It: Music Source Separation via Hyperellipsoidal Queries," arXiv:2501.16171.
+
+[3] K. N. Watcharasupat and A. Lerch, "A Stem-Agnostic Single-Decoder System for Music Source Separation Beyond Four Stems," Proceedings of the 25th International Society for Music Information Retrieval Conference, 2024. (arXiv:2406.18747) 

torch_l1_snr-0.0.1.dist-info/RECORD

@@ -0,0 +1,7 @@
+torch_l1_snr-0.0.1.dist-info/licenses/LICENSE,sha256=gBRgAD6TvJXVLS9LVkO0V1GzNHKX2poZLFGtyc6hwq0,1079
+torch_l1snr/__init__.py,sha256=pR9jg3fjTKt_suZoVDC67tqB7EWRkbfaXaPP7pYQrlQ,220
+torch_l1snr/l1snr.py,sha256=aqmtNfT_8A0IRI9jiVGwNse3igBvelQGKnjfe23Xh7w,35304
+torch_l1_snr-0.0.1.dist-info/METADATA,sha256=0EuezNJH0APVDEWIsmqW_6Wc9B1J-Zk3ksg8VknU7kY,10374
+torch_l1_snr-0.0.1.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+torch_l1_snr-0.0.1.dist-info/top_level.txt,sha256=NfaRND6pcjZ7-035d4XAg8xJuz31EEU210Y9xWeFOxc,12
+torch_l1_snr-0.0.1.dist-info/RECORD,,

torch_l1_snr-0.0.1.dist-info/WHEEL

@@ -0,0 +1,5 @@
+Wheel-Version: 1.0
+Generator: setuptools (80.9.0)
+Root-Is-Purelib: true
+Tag: py3-none-any
+

torch_l1_snr-0.0.1.dist-info/licenses/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2025 Christopher Landschoot
+
+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.

torch_l1_snr-0.0.1.dist-info/top_level.txt

@@ -0,0 +1 @@
+torch_l1snr

torch_l1snr/__init__.py

@@ -0,0 +1,15 @@
+from .l1snr import (
+    dbrms,
+    L1SNRLoss,
+    L1SNRDBLoss,
+    STFTL1SNRDBLoss,
+    MultiL1SNRDBLoss,
+)
+
+__all__ = [
+    "dbrms",
+    "L1SNRLoss",
+    "L1SNRDBLoss",
+    "STFTL1SNRDBLoss",
+    "MultiL1SNRDBLoss",
+] 

torch_l1snr/l1snr.py

@@ -0,0 +1,786 @@
+# PyTorch implementation of L1SNR loss functions for audio source separation
+# https://github.com/crlandsc/torch-l1-snr
+# Copyright (c) 2026 crlandsc
+# MIT License
+#
+# This implementation is based on and extends the loss functions described in:
+# [1] "Separate This, and All of these Things Around It: Music Source Separation via Hyperellipsoidal Queries"
+#     Karn N. Watcharasupat, Alexander Lerch
+#     arXiv:2501.16171
+# [2] "A Generalized Bandsplit Neural Network for Cinematic Audio Source Separation"
+#     Karn N. Watcharasupat, Chih-Wei Wu, Yiwei Ding, Iroro Orife, Aaron J. Hipple, Phillip A. Williams, Scott Kramer, Alexander Lerch, William Wolcott
+#     IEEE Open Journal of Signal Processing, 2023
+#     arXiv:2309.02539
+# [3] "A Stem-Agnostic Single-Decoder System for Music Source Separation Beyond Four Stems"
+#     Karn N. Watcharasupat, Alexander Lerch
+#     Proceedings of the 25th International Society for Music Information Retrieval Conference, 2024
+#     arXiv:2406.18747
+
+import torch
+import torch.nn as nn
+from torchaudio.transforms import Spectrogram
+import math
+from typing import Union, Dict, Tuple, Optional, List
+    
+def dbrms(x, eps=1e-8):
+    """
+    Compute RMS level in decibels for a batch of signals.
+    Args:
+        x: (batch, time) or (batch, ...) tensor
+        eps: stability constant
+    Returns:
+        (batch,) tensor of dB RMS
+    """
+    x = x.reshape(x.shape[0], -1)
+    rms = torch.sqrt(torch.mean(x ** 2, dim=-1) + eps)
+    return 20.0 * torch.log10(rms + eps)
+
+
+class L1SNRLoss(torch.nn.Module):
+    """
+    Implements the L1 Signal-to-Noise Ratio (SNR) loss with optional weighted L1 loss 
+    component to balance "all-or-nothing" behavior.
+
+    Paper-aligned D1(ŷ; y) form:
+      D1 = 10 * log10( (||ŷ - y||_1 + eps) / (||y||_1 + eps) )
+    L1SNR_loss = mean(D1)
+
+    When l1_weight > 0, the loss combines L1SNR with scaled L1:
+      loss = (1 - l1_weight) * L1SNR_loss + l1_weight * L1_auto_scaled
+
+    Input Shape:
+        Accepts waveform tensors (time-domain audio) of any shape as long as they are batch-first.
+        Recommended shapes:
+        - [batch, time] for single-source audio
+        - [batch, num_sources, time] for multi-source audio
+        - [batch, num_sources, channels, time] for multi-channel multi-source audio
+
+    Attributes:
+        name (str): Name identifier for the loss.
+        weight (float): Global weight multiplier for the loss.
+        eps (float): Small epsilon for numerical stability in D1 (default 1e-3 per the papers).
+        l1_weight (float): Weight for the L1 term mixed into L1SNR.
+    """
+    def __init__(
+        self,
+        name,
+        weight: float = 1.0,
+        eps: float = 1e-3,
+        l1_weight: float = 0.0,
+    ):
+        super().__init__()
+        self.name = name
+        self.weight = weight
+        self.eps = eps
+        self.l1_weight = l1_weight
+
+    def forward(self, estimates, actuals, *args, **kwargs):
+        batch_size = estimates.shape[0]
+        
+        est_source = estimates.reshape(batch_size, -1)
+        act_source = actuals.reshape(batch_size, -1)
+
+        # L1 errors and reference
+        l1_error = torch.mean(torch.abs(est_source - act_source), dim=-1)
+        l1_true = torch.mean(torch.abs(act_source), dim=-1)
+        
+        # Auto-balanced L1/SNR mixing
+        w = float(self.l1_weight)
+
+        # Pure-L1 shortcut: avoid D1 computation
+        if w >= 1.0:
+            return torch.mean(l1_error) * self.weight
+
+        # If pure SNR (w == 0) we can skip L1 scaling math
+        if w <= 0.0:
+            d1 = 10.0 * torch.log10((l1_error + self.eps) / (l1_true + self.eps))
+            l1snr_loss = torch.mean(d1)
+            return l1snr_loss * self.weight
+
+        # Mixed path
+        d1 = 10.0 * torch.log10((l1_error + self.eps) / (l1_true + self.eps))
+        l1snr_loss = torch.mean(d1)
+
+        c = 10.0 / math.log(10.0)
+        inv_mean = torch.mean(1.0 / (l1_error.detach() + self.eps))
+        # w-independent scaling to match typical gradient magnitudes
+        scale_time = c * inv_mean
+        l1_term = torch.mean(l1_error) * scale_time
+
+        if getattr(self, "balance_per_sample", False):
+            # per-sample w-independent scaling
+            bal = c / (l1_error.detach() + self.eps)
+            l1_term = torch.mean(l1_error * bal)
+
+        if getattr(self, "debug_balance", False):
+            g_d1 = (1.0 - w) * c * inv_mean
+            if getattr(self, "balance_per_sample", False):
+                g_l1 = w * torch.mean(c / (l1_error.detach() + self.eps))
+            else:
+                g_l1 = w * c * inv_mean
+            ratio = (g_l1 / (g_d1 + 1e-12)).item()
+            setattr(self, "last_balance_ratio", ratio)
+
+        loss = (1.0 - w) * l1snr_loss + w * l1_term
+        return loss * self.weight
+
+
+class L1SNRDBLoss(torch.nn.Module):
+    """
+    Implements L1SNR plus adaptive level-matching regularization in the time domain
+    as described in arXiv:2501.16171, with optional L1 loss component to balance
+    "all-or-nothing" behavior.
+    
+    The loss combines three components:
+    1. L1SNR loss: mean(10*log10((l1_error + eps) / (l1_true + eps)))
+    2. Level-matching regularization: λ*|L_pred - L_true|
+       Where λ is adaptively computed based on the signal levels
+    3. Optional L1 loss: mean(l1_error)
+    
+    The complete loss is structured as:
+    When l1_weight < 1.0: total_loss = l1snr_loss + (1-l1_weight) * mean(reg_loss)
+    When l1_weight = 1.0: total_loss = l1_loss (pure L1, bypassing all other computations)
+    
+    The adaptive weighting λ for regularization increases when loud parts of a stem aren't
+    reconstructed properly, helping balance between quality and level preservation.
+    
+    When l1_weight=1.0, this loss efficiently switches to a pure L1 loss calculation,
+    bypassing all SNR and regularization computations for standard L1 behavior.
+    This is useful when you want to avoid the "all-or-nothing" behavior of the SNR-style loss.
+    
+    Input Shape:
+        Accepts waveform tensors (time-domain audio) of any shape as long as they are batch-first.
+        Recommended shapes:
+        - [batch, time] for single-source audio
+        - [batch, num_sources, time] for multi-source audio
+        - [batch, num_sources, channels, time] for multi-channel multi-source audio
+    
+    Attributes:
+        name (str): The name identifier for the loss.
+        weight (float): The overall weight multiplier for the loss.
+        lambda0 (float): Minimum regularization weight (λ_min).
+        delta_lambda (float): Range of extra weight for regularization (Δλ).
+        l1snr_eps (float): Epsilon value for the L1SNR component to avoid log(0).
+        dbrms_eps (float): Epsilon value for dBRMS calculation to avoid log(0).
+        lmin (float): Minimum dBRMS considered non-silent for adaptive weighting.
+        use_regularization (bool): Whether to use level-matching regularization.
+            If False, only the L1SNR (and optional L1) components are used.
+        l1_weight (float): Weight for the L1 loss component. Default 0 (disabled).
+            As this increases, the regularization term is also scaled down proportionally.
+            When set to 1.0, efficiently computes only L1 loss.
+    """
+    def __init__(
+        self, 
+        name, 
+        weight: float = 1.0,
+        lambda0: float = 0.1,
+        delta_lambda: float = 0.9,
+        l1snr_eps: float = 1e-3,
+        dbrms_eps: float = 1e-8,
+        lmin: float = -60.0,
+        use_regularization: bool = True,
+        l1_weight: float = 0.0,
+    ):
+        super().__init__()
+        self.name = name
+        self.weight = weight
+        self.lambda0 = lambda0          # minimum regularization weight
+        self.delta_lambda = delta_lambda # range of extra weight
+        self.l1snr_eps = l1snr_eps
+        self.dbrms_eps = dbrms_eps
+        self.lmin = lmin
+        self.use_regularization = use_regularization
+        
+        # Validate l1_weight is between 0.0 and 1.0 inclusive
+        assert 0.0 <= l1_weight <= 1.0, "l1_weight must be between 0.0 and 1.0 inclusive"
+        self.l1_weight = l1_weight
+        
+        # Initialize component losses based on l1_weight
+        if self.l1_weight == 1.0:
+            # Pure L1 mode - only need L1 loss
+            self.l1snr_loss = None
+            self.l1_loss = torch.nn.L1Loss()
+        else:
+            # Standard mode with L1SNR (and optional weighted L1 if l1_weight > 0)
+            self.l1snr_loss = L1SNRLoss(
+                name="l1_snr",
+                weight=1.0,  # We'll apply the weight at the end
+                eps=l1snr_eps,
+                l1_weight=l1_weight,
+            )
+            self.l1_loss = None
+    
+    @staticmethod
+    def compute_adaptive_weight(L_pred, L_true, L_min, lambda0, delta_lambda, R):
+        """
+        Implements the adaptive weighting of the level-matching regularization term, per arXiv:2501.16171.
+        Args:
+            L_pred: predicted dBRMS, shape (batch,)
+            L_true: reference dBRMS, shape (batch,)
+            L_min: minimum dBRMS considered non-silent (float)
+            lambda0: minimum weight for regularization
+            delta_lambda: range of extra weight for regularization
+            R: |L_pred - L_true|, shape (batch,)
+        Returns:
+            lambda_weight: shape (batch,)
+        """
+        # Compute eta: 1 if L_true > max(L_pred, L_min), else 0
+        max_val = torch.max(L_pred, torch.full_like(L_true, L_min))
+        eta = (L_true > max_val).float()
+        denom = (L_true - L_min).clamp(min=1e-6)
+        clamp_arg = (R / denom).clamp(0.0, 1.0)
+        lam = lambda0 + eta * delta_lambda * clamp_arg
+        return lam.detach()  # Stop-gradient
+    
+    def forward(self, estimates, actuals, *args, **kwargs):
+        batch_size = estimates.shape[0]
+        
+        est_source = estimates.reshape(batch_size, -1)
+        act_source = actuals.reshape(batch_size, -1)
+            
+        # Pure L1 mode - efficient path that bypasses SNR and regularization
+        if self.l1_loss is not None:
+            l1_loss = self.l1_loss(est_source, act_source)
+            return l1_loss * self.weight
+            
+        # Standard mode with L1SNR, regularization, and optional weighted L1
+        # 1. L1SNR reconstruction loss (with L1 component if l1_weight > 0)
+        l1snr_loss = self.l1snr_loss(estimates, actuals, *args, **kwargs)
+        
+        # Only compute and apply regularization if enabled
+        if self.use_regularization:
+            # 2. Level-matching regularization
+            L_true = dbrms(act_source, self.dbrms_eps)   # (batch,)
+            L_pred = dbrms(est_source, self.dbrms_eps)   # (batch,)
+            R = torch.abs(L_pred - L_true)               # (batch,)
+            
+            lambda_weight = self.compute_adaptive_weight(L_pred, L_true, self.lmin, self.lambda0, self.delta_lambda, R)  # (batch,)
+            
+            reg_loss = lambda_weight * R
+            
+            # Scale regularization by the same factor as L1SNR loss
+            l1snr_weight = 1.0 - self.l1_weight
+            total_loss = l1snr_loss + (l1snr_weight * torch.mean(reg_loss))
+        else:
+            # Skip regularization calculation entirely
+            total_loss = l1snr_loss
+        
+        return total_loss * self.weight
+
+
+class STFTL1SNRDBLoss(torch.nn.Module):
+    """
+    Implements L1SNR plus adaptive level-matching regularization in the spectrogram domain
+    as described in arXiv:2501.16171, with multi-resolution STFT and optional L1 loss component
+    to balance "all-or-nothing" behavior.
+    
+    This loss applies the same principles as L1SNRDBLoss but operates in the complex
+    spectrogram domain across multiple time-frequency resolutions. For each resolution:
+    
+    1. L1SNR loss: Computed on the complex STFT representation (real/imaginary parts)
+    2. Level-matching regularization: Applied to STFT magnitudes with adaptive weighting
+    3. Optional L1 loss: Direct L1 penalty on STFT differences
+    
+    Multi-resolution processing helps capture both fine temporal details and frequency
+    characteristics. The loss averages results across all valid STFT resolutions.
+    
+    The complete loss structure is similar to L1SNRDBLoss:
+    When l1_weight < 1.0: total_loss = l1snr_loss + (1-l1_weight) * spec_reg_coef * mean(reg_loss)
+    When l1_weight = 1.0: total_loss = l1_loss (pure L1 in spectrogram domain, bypassing all other computations)
+    
+    When l1_weight=1.0, this loss efficiently switches to a pure L1 loss calculation in the
+    spectrogram domain, bypassing all SNR and regularization computations for standard L1 behavior.
+    This is useful when you want to avoid the "all-or-nothing" behavior of the SNR-style loss.
+    
+    Input Shape:
+        Accepts waveform tensors (time-domain audio) of any shape as long as they are batch-first
+        and time-last. Recommended shapes:
+        - [batch, time] for single-source audio
+        - [batch, num_sources, time] for multi-source audio
+        - [batch, num_sources, channels, time] for multi-channel multi-source audio
+    
+    Attributes:
+        name (str): The name identifier for the loss.
+        weight (float): The overall weight multiplier for the loss.
+        lambda0 (float): Minimum regularization weight (λ_min).
+        delta_lambda (float): Range of extra weight for regularization (Δλ).
+        l1snr_eps (float): Epsilon value for the L1SNR component to avoid log(0).
+        dbrms_eps (float): Epsilon value for dBRMS calculation to avoid log(0).
+        lmin (float): Minimum dBRMS considered non-silent for adaptive weighting.
+        n_ffts (List[int]): List of FFT sizes for multi-resolution STFT analysis.
+        hop_lengths (List[int]): List of hop lengths (STFT time steps) for each resolution.
+        win_lengths (List[int]): List of window lengths for each resolution.
+        window_fn (str): Window function for the STFT ('hann', 'hamming', etc.)
+        min_audio_length (int): Minimum audio length required for processing.
+            If audio is shorter, returns zero loss to avoid errors.
+        use_regularization (bool): Whether to use level-matching regularization.
+            If False, only the L1SNR (and optional L1) components are used.
+        l1_weight (float): Weight for the L1 loss component. Default 0 (disabled).
+            As this increases, the regularization term is also scaled down proportionally.
+            When set to 1.0, efficiently computes only L1 loss.
+    """
+    def __init__(
+        self, 
+        name, 
+        weight: float = 1.0,
+        lambda0: float = 0.1,
+        delta_lambda: float = 0.9,
+        l1snr_eps: float = 1e-3,
+        dbrms_eps: float = 1e-8,
+        lmin: float = -60.0,
+        n_ffts: List[int] = [512, 1024, 2048],
+        hop_lengths: List[int] = [128, 256, 512],
+        win_lengths: List[int] = [512, 1024, 2048],
+        window_fn: str = 'hann',
+        min_audio_length: int = 512,
+        use_regularization: bool = False,
+        spec_reg_coef: float = 0.1,
+        l1_weight: float = 0.0,
+    ):
+        super().__init__()
+        self.name = name
+        self.weight = weight
+        self.min_audio_length = min_audio_length
+        
+        # Validate STFT parameters
+        assert len(n_ffts) == len(hop_lengths) == len(win_lengths), "All STFT parameter lists must have the same length"
+        
+        # Store STFT parameters for validation during forward pass
+        self.n_ffts = n_ffts
+        self.hop_lengths = hop_lengths
+        self.win_lengths = win_lengths
+        self.window_fn_name = window_fn
+        
+        # Validate window sizes
+        for n_fft, win_length in zip(n_ffts, win_lengths):
+            assert n_fft >= win_length, f"FFT size ({n_fft}) must be greater than or equal to window length ({win_length})"
+        
+        # Pre-initialize Spectrogram transforms for maximum efficiency
+        self.spectrogram_transforms = nn.ModuleList()
+        window_fn_callable = getattr(torch, f"{window_fn}_window")
+        
+        for n_fft, hop_length, win_length in zip(n_ffts, hop_lengths, win_lengths):
+            # Create a spectrogram transform for each resolution
+            transform = Spectrogram(
+                n_fft=n_fft,
+                win_length=win_length,
+                hop_length=hop_length,
+                pad_mode="reflect",
+                center=True,
+                window_fn=window_fn_callable,
+                normalized=True,
+                power=None,  # This ensures the output is complex
+            )
+            self.spectrogram_transforms.append(transform)
+        
+        # Parameters for spectrogram domain level-matching
+        self.lambda0 = lambda0
+        self.delta_lambda = delta_lambda
+        self.lmin = lmin
+        self.dbrms_eps = dbrms_eps
+        self.l1snr_eps = l1snr_eps
+        
+        # Add L1 loss parameters and validate
+        assert 0.0 <= l1_weight <= 1.0, "l1_weight must be between 0.0 and 1.0 inclusive"
+        self.l1_weight = l1_weight
+        
+        # Flag for pure L1 mode
+        self.pure_l1_mode = (self.l1_weight == 1.0)
+        # Create L1 loss function for pure L1 mode
+        if self.pure_l1_mode:
+            self.l1_loss = torch.nn.L1Loss()
+        else:
+            self.l1_loss = None
+
+        
+        # Add this parameter to control regularization
+        self.use_regularization = use_regularization
+        # Coefficient to scale spectral regularization (disabled by default)
+        self.spec_reg_coef = spec_reg_coef
+        
+        # Fallback time-domain loss (used when audio is too short for TF processing)
+        self.fallback_time_loss = L1SNRDBLoss(
+            name=f"{name}_fallback_time",
+            weight=1.0,
+            lambda0=self.lambda0,
+            delta_lambda=self.delta_lambda,
+            l1snr_eps=self.l1snr_eps,
+            dbrms_eps=self.dbrms_eps,
+            lmin=self.lmin,
+            use_regularization=False,  # regularizer belongs to TF for this class
+            l1_weight=self.l1_weight,
+        )
+
+        # Simplified tracking
+        self.nan_inf_counts = {"inputs": 0, "spec_loss": 0}
+
+    def _compute_complex_spec_l1snr_loss(self, est_spec, act_spec):
+        """
+        Compute TF-domain loss as per the papers:
+        - D1 on real part + D1 on imaginary part, summed.
+        - Optional L1 mixing applied symmetrically to Re/Im.
+        est_spec, act_spec: complex tensors with shape (B, C, F, T)
+        """
+        # Ensure same shape (assert to avoid silent mismatches)
+        assert est_spec.shape == act_spec.shape, f"Spec shapes must match: {est_spec.shape} vs {act_spec.shape}"
+
+        # Split real/imag
+        est_re, est_im = est_spec.real, est_spec.imag
+        act_re, act_im = act_spec.real, act_spec.imag
+
+        B = est_spec.shape[0]
+
+        # Flatten to (B, -1)
+        est_re = est_re.reshape(B, -1)
+        act_re = act_re.reshape(B, -1)
+        est_im = est_im.reshape(B, -1)
+        act_im = act_im.reshape(B, -1)
+
+        # L1 errors and refs
+        err_re = torch.mean(torch.abs(est_re - act_re), dim=1)
+        ref_re = torch.mean(torch.abs(act_re), dim=1)
+        err_im = torch.mean(torch.abs(est_im - act_im), dim=1)
+        ref_im = torch.mean(torch.abs(act_im), dim=1)
+
+        # Paper-aligned D1 = 10*log10((||e||_1 + eps)/(||y||_1 + eps))
+        d1_re = 10.0 * torch.log10((err_re + self.l1snr_eps) / (ref_re + self.l1snr_eps))
+        d1_im = 10.0 * torch.log10((err_im + self.l1snr_eps) / (ref_im + self.l1snr_eps))
+        d1_sum = torch.mean(d1_re + d1_im)  # mean over batch
+
+        # Pure L1 mode
+        if self.pure_l1_mode:
+            l1_re = torch.mean(err_re)
+            l1_im = torch.mean(err_im)
+            l1_term = 0.5 * (l1_re + l1_im)
+            return l1_term
+
+        # Mixed mode (auto-balanced L1/SNR) with per-batch scaling
+        w = float(self.l1_weight)
+        if 0.0 < w < 1.0:
+            c = 10.0 / math.log(10.0)
+            inv_mean_comp = torch.mean(0.5 * (1.0 / (err_re.detach() + self.l1snr_eps) +
+                                              1.0 / (err_im.detach() + self.l1snr_eps)))
+            # w-independent scaling to match typical gradient magnitudes (factor 2.0 for Re/Im symmetry)
+            scale_spec = 2.0 * c * inv_mean_comp
+            l1_term = 0.5 * (torch.mean(err_re) + torch.mean(err_im)) * scale_spec
+
+            if getattr(self, "balance_per_sample", False):
+                bal_re = c / (err_re.detach() + self.l1snr_eps)
+                bal_im = c / (err_im.detach() + self.l1snr_eps)
+                l1_term = 0.5 * (torch.mean(err_re * bal_re) + torch.mean(err_im * bal_im))
+
+            loss = (1.0 - w) * d1_sum + w * l1_term
+            return loss
+        elif w >= 1.0:
+            # Pure L1
+            l1_term = 0.5 * (torch.mean(err_re) + torch.mean(err_im))
+            return l1_term
+        else:
+            # Pure SNR (D1)
+            return d1_sum
+    
+    def _compute_spec_level_matching(self, est_spec, act_spec):
+        """
+        Compute the level matching regularization term for a spectrogram.
+        """
+        batch_size = est_spec.shape[0]
+        
+        # Make sure dimensions match before operations
+        if est_spec.shape != act_spec.shape:
+            # Resize to match the smaller of the two
+            min_freq = min(est_spec.shape[-2], act_spec.shape[-2])
+            min_time = min(est_spec.shape[-1], act_spec.shape[-1])
+            est_spec = est_spec[..., :min_freq, :min_time]
+            act_spec = act_spec[..., :min_freq, :min_time]
+        
+        # For level-matching regularization, we use magnitude information
+        est_mag = torch.abs(est_spec)
+        act_mag = torch.abs(act_spec)
+        
+        # Reshape once for efficiency
+        est_mag_flat = est_mag.reshape(batch_size, -1)
+        act_mag_flat = act_mag.reshape(batch_size, -1)
+        
+        # Calculate dB levels
+        L_true = dbrms(act_mag_flat, self.dbrms_eps)
+        L_pred = dbrms(est_mag_flat, self.dbrms_eps)
+            
+        R = torch.abs(L_pred - L_true)
+        
+        # Use the adaptive weighting function
+        lambda_weight = L1SNRDBLoss.compute_adaptive_weight(
+            L_pred, L_true, self.lmin, self.lambda0, self.delta_lambda, R
+        )
+        
+        return torch.mean(lambda_weight * R)
+
+    def _validate_audio_length(self, audio_length):
+        """
+        Validates that the audio is long enough for the STFT parameters.
+        """
+        if audio_length < self.min_audio_length:
+            return False
+            
+        for n_fft, hop_length in zip(self.n_ffts, self.hop_lengths):
+            n_frames = (audio_length // hop_length) + 1
+            if n_frames < 2:
+                return False
+                
+        return True
+
+    def forward(self, estimates, actuals, *args, **kwargs):
+        device = estimates.device
+        batch_size = estimates.shape[0]
+        
+        # Basic NaN/Inf handling (simplified)
+        if torch.isnan(estimates).any() or torch.isinf(estimates).any() or torch.isnan(actuals).any() or torch.isinf(actuals).any():
+            self.nan_inf_counts["inputs"] += 1
+            estimates = torch.nan_to_num(estimates, nan=0.0, posinf=1.0, neginf=-1.0)
+            actuals = torch.nan_to_num(actuals, nan=0.0, posinf=1.0, neginf=-1.0)
+        
+        est_source = estimates.reshape(batch_size, -1, estimates.shape[-1])
+        act_source = actuals.reshape(batch_size, -1, actuals.shape[-1])
+        
+        # Validate audio length
+        audio_length = est_source.shape[-1]
+        if not self._validate_audio_length(audio_length):
+            # Fallback to time-domain L1SNR-style loss instead of zero
+            return self.fallback_time_loss(estimates, actuals, *args, **kwargs) * self.weight
+        
+        # Track losses (initialize as tensors on the correct device for stability)
+        total_spec_loss = torch.tensor(0.0, device=device)
+        total_spec_reg_loss = torch.tensor(0.0, device=device)
+        valid_transforms = 0
+        
+        # Ensure transforms are on the correct device
+        self.spectrogram_transforms.to(device)
+        
+        # Process each resolution
+        for i, transform in enumerate(self.spectrogram_transforms):
+            try:
+                # Compute spectrograms using pre-initialized transforms
+                try:
+                    est_spec = transform(est_source)
+                    act_spec = transform(act_source)
+                except RuntimeError as e:
+                    print(f"Error computing spectrogram for resolution {i}: {e}")
+                    print(f"Parameters: n_fft={self.n_ffts[i]}, hop_length={self.hop_lengths[i]}, win_length={self.win_lengths[i]}")
+                    continue
+                
+                # Ensure same (B, C, F, T); crop only (F, T) if needed
+                if est_spec.shape != act_spec.shape:
+                    min_f = min(est_spec.shape[-2], act_spec.shape[-2])
+                    min_t = min(est_spec.shape[-1], act_spec.shape[-1])
+                    est_spec = est_spec[..., :min_f, :min_t]
+                    act_spec = act_spec[..., :min_f, :min_t]
+                
+                # Compute complex spectral loss (either L1 or L1SNR based on self.pure_l1_mode)
+                try:
+                    spec_loss = self._compute_complex_spec_l1snr_loss(est_spec, act_spec)
+                except RuntimeError as e:
+                    print(f"Error computing complex spectral loss for resolution {i}: {e}")
+                    continue
+                
+                # Check for numerical issues
+                if torch.isnan(spec_loss) or torch.isinf(spec_loss):
+                    self.nan_inf_counts["spec_loss"] += 1
+                    continue
+                
+                # Only compute regularization if not in pure L1 mode and regularization is enabled
+                if not self.pure_l1_mode and self.use_regularization:
+                    try:
+                        spec_reg_loss = self._compute_spec_level_matching(est_spec, act_spec)
+                        
+                        # Check for numerical issues
+                        if torch.isnan(spec_reg_loss) or torch.isinf(spec_reg_loss):
+                            self.nan_inf_counts["spec_loss"] += 1
+                            spec_reg_loss = 0.0  # Use zero reg_loss if there are issues
+                        
+                        # Accumulate regularization loss
+                        total_spec_reg_loss += spec_reg_loss
+                    except RuntimeError as e:
+                        print(f"Error computing spectral level-matching for resolution {i}: {e}")
+                
+                # Accumulate loss
+                total_spec_loss += spec_loss
+                valid_transforms += 1
+                                
+            except RuntimeError as e:
+                print(f"Runtime error in spectrogram transform {i}: {e}")
+                continue
+        
+        # If all transforms failed, return zero loss
+        if valid_transforms == 0:
+            print("Warning: All spectrogram transforms failed. Returning zero loss.")
+            return torch.tensor(0.0, device=device)
+        
+        # Average losses across valid transforms
+        avg_spec_loss = total_spec_loss / valid_transforms
+        
+        # For standard mode, apply regularization if enabled
+        if not self.pure_l1_mode and self.use_regularization:
+            avg_spec_reg_loss = total_spec_reg_loss / valid_transforms
+            # Scale spectral regularization by both (1 - l1_weight) and spec_reg_coef
+            l1snr_weight = 1.0 - self.l1_weight
+            final_loss = avg_spec_loss + l1snr_weight * (self.spec_reg_coef * avg_spec_reg_loss)
+        else:
+            final_loss = avg_spec_loss
+
+        return final_loss * self.weight
+
+
+class MultiL1SNRDBLoss(torch.nn.Module):
+    """
+    A modular loss function that combines time-domain and spectrogram-domain L1SNR and
+    adaptive level-matching losses, as described in arXiv:2501.16171, with optional
+    L1 loss component to balance "all-or-nothing" behavior.
+    
+    This implementation uses separate specialized components:
+    - L1SNRDBLoss for time domain processing
+    - STFTL1SNRDBLoss for spectrogram domain processing
+    
+    The loss combines time-domain and spectrogram-domain losses:
+    Loss = weight * [(1-spec_weight) * time_loss + spec_weight * spec_loss]
+    
+    Where time_loss and spec_loss are computed by L1SNRDBLoss and STFTL1SNRDBLoss respectively,
+    each handling their own L1SNR, regularization, and optional L1 components as described
+    in their individual docstrings.
+    
+    When l1_weight=1.0, this loss efficiently switches to a pure L1 loss calculation in both
+    domains, bypassing all SNR and regularization computations for standard L1 behavior.
+    This is useful when you want to avoid the "all-or-nothing" behavior of the SNR-style loss.
+    
+    The regularization components use adaptive weighting based on level differences
+    between estimated and target signals, with weighting controlled by lambda0 and delta_lambda.
+    
+    Input Shape:
+        Accepts waveform tensors (time-domain audio) of any shape as long as they are batch-first
+        and time-last. Recommended shapes:
+        - [batch, time] for single-source audio
+        - [batch, num_sources, time] for multi-source audio
+        - [batch, num_sources, channels, time] for multi-channel multi-source audio
+    
+    Attributes:
+        name (str): The name identifier for the loss.
+        weight (float): The overall weight multiplier for the loss.
+        spec_weight (float): The weight for spectrogram domain loss relative to time domain.
+            Default 0.5 (equal weighting). Set higher to emphasize spectral accuracy.
+        use_time_regularization (bool): Whether to use level-matching regularization in time domain.
+        use_spec_regularization (bool): Whether to use level-matching regularization in spectogram domain.
+        l1_weight (float): Weight for the L1 loss component vs the L1SNR+reg components.
+            Default 0 (disabled). As this increases, the regularization term is also scaled down.
+            When set to 1.0, efficiently computes only L1 loss in both domains.
+        lambda0 (float): Minimum regularization weight for both domains.
+        delta_lambda (float): Range of extra weight for regularization in both domains.
+        time_loss_params (dict): Optional additional parameters to pass to time domain loss.
+        spec_loss_params (dict): Optional additional parameters to pass to spectrogram domain loss.
+    """
+    def __init__(
+        self, 
+        name, 
+        weight: float = 1.0,
+        spec_weight: float = 0.5,  # Balance between time and frequency domain
+        # L1 component parameters
+        l1_weight: float = 0.0, # Weight for the L1 loss component vs (L1SNR + Regularization).
+                               # Note: Regularization term is also scaled by (1.0 - l1_weight).
+                               # When set to 1.0, efficiently computes only L1 loss in both domains.
+        # auto-balanced mixing used
+        # Regularization on/off flags
+        use_time_regularization: bool = True,
+        use_spec_regularization: bool = False, # likely redundant if already using in time domain
+        # Default parameters for both loss components
+        lambda0: float = 0.1,
+        delta_lambda: float = 0.9,
+        l1snr_eps: float = 1e-3,
+        dbrms_eps: float = 1e-8,
+        lmin: float = -60.0,
+        # STFT parameters
+        n_ffts: List[int] = [512, 1024, 2048],
+        hop_lengths: List[int] = [128, 256, 512],
+        win_lengths: List[int] = [512, 1024, 2048],
+        window_fn: str = 'hann',
+        min_audio_length: int = 512,
+        # Allow for separate parameter overrides (e.g. different delta_lambda for time and spec)
+        time_loss_params: dict = None,
+        spec_loss_params: dict = None,
+    ):
+        super().__init__()
+        self.name = name
+        self.weight = weight
+        self.spec_weight = spec_weight
+        
+        # Validate l1_weight is in valid range
+        assert 0.0 <= l1_weight <= 1.0, "l1_weight must be between 0.0 and 1.0 inclusive"
+        self.l1_weight = l1_weight
+        self.use_time_regularization = use_time_regularization
+        self.use_spec_regularization = use_spec_regularization
+        
+        # Set up default parameters
+        default_time_params = {
+            "name": f"{name}_time",
+            "weight": 1.0,  # Will be scaled by the combined loss
+            "lambda0": lambda0,
+            "delta_lambda": delta_lambda,
+            "l1snr_eps": l1snr_eps,
+            "dbrms_eps": dbrms_eps,
+            "lmin": lmin,
+            "l1_weight": l1_weight,
+            "use_regularization": use_time_regularization  # Apply time domain regularization flag
+        }
+        
+        default_spec_params = {
+            "name": f"{name}_spec",
+            "weight": 1.0,  # Will be scaled by the combined loss
+            "lambda0": lambda0,
+            "delta_lambda": delta_lambda,
+            "l1snr_eps": l1snr_eps,
+            "dbrms_eps": dbrms_eps,
+            "lmin": lmin,
+            "n_ffts": n_ffts,
+            "hop_lengths": hop_lengths,
+            "win_lengths": win_lengths,
+            "window_fn": window_fn,
+            "min_audio_length": min_audio_length,
+            "l1_weight": l1_weight,
+
+            "use_regularization": use_spec_regularization  # Apply spectrogram domain regularization flag
+        }
+        
+        # Override with any custom parameters
+        if time_loss_params:
+            default_time_params.update(time_loss_params)
+        if spec_loss_params:
+            default_spec_params.update(spec_loss_params)
+        
+        # Create the specialized loss components
+        # Note: Component losses handle all optimizations internally based on l1_weight
+        # When l1_weight=1.0, they will efficiently bypass SNR and regularization calculations
+        self.time_loss = L1SNRDBLoss(**default_time_params)
+        self.spec_loss = STFTL1SNRDBLoss(**default_spec_params)
+        
+        # For reference only, indicate if we're in pure L1 mode
+        self.pure_l1_mode = (self.l1_weight == 1.0)
+
+    def forward(self, estimates, actuals, *args, **kwargs):
+        """
+        Forward pass to compute the combined multi-domain loss.
+        
+        Args:
+            estimates: Model output predictions, shape [batch, ...] (batch-first, ..., time-last)
+            actuals: Ground truth targets, shape [batch, ...] (batch-first, ..., time-last)
+            *args, **kwargs: Additional arguments passed to sub-losses
+            
+        Returns:
+            Combined weighted loss from time and spectrogram domains
+        """
+        # Compute time domain loss
+        time_loss = self.time_loss(estimates, actuals, *args, **kwargs)
+        
+        # Compute spectrogram domain loss
+        spec_loss = self.spec_loss(estimates, actuals, *args, **kwargs)
+        
+        # Combine with weighting
+        combined_loss = (1 - self.spec_weight) * time_loss + self.spec_weight * spec_loss
+        
+        # Apply overall weight
+        return combined_loss * self.weight