tests, limiter

.gitignore

@@ -8,6 +8,7 @@ __pycache__/
 
 # development directories
 audio/
+.nox/
 
 # Distribution / packaging
 .Python

README.md

@@ -17,11 +17,9 @@ This repo contains signal processing nodes for [ComfyUI](https://github.com/comf
 
 ### Latests Updates
 
-- **`Clipping removal`** - removed all the clipping functions to ensure audio is not clipped throught processing chain
-- **`Double-precision floating-point`** - changed internal processing to use double precision floating point. 
-- **`Keeping loudness`** - all processing nodes use lufs loudness normalization to keep pre-processing gain.
-- **`GPU support`** - ported some of the nodes to gpu to gain speed
-- **`Removed some nodes`** - if you want them back contact me through this repository. I am planning to add new nodes with better quality
+- **`Gpu limiter - experimental`** - Limiter optimized for parallel GPU execution
+- **`Saturatation - experimental`** - Added basic saturation node with couple algorithms
+- **`Tests suite`** - Added basic tests suite for most of the nodes
 
 ---
 
@@ -48,7 +46,7 @@ Implementation is using the standard shelf filter equations from the Audio EQ Co
 
 ### Enhance Harmonics
 
-Harmonic enhancer boosts selected harmonics to enrich the sound
+Harmonic enhancer boosts selected harmonics to enrich the sound. 
 
 #### Parameters:
 
@@ -229,4 +227,26 @@ This node lets you stretch audio to about 100x it's original speed whil mainting
 #### Parameters:
 
 - **`audio_file`**: input audio file
-- **`gain`**: when to start audio from
+- **`gain`**: when to start audio from
+
+
+### Development and Testing/Profiling
+
+### Dependnecies:
+
+The development dependencies are specified in the pyproject.toml file and are not included in the requirements.txt. 
+To run tests and contribute to development, you'll need the following tools and modules installed. 
+My primary development and testing environment is Ubuntu Linux with an RTX 3090 GPU and CUDA 11.8.
+
+## Required Tools and Libraries:
+
+-**`scalene`**: High-performance CPU and GPU profiler.
+-**`pytest`**: Comprehensive testing framework.
+-**`nox`**: Automation tool for managing tasks.
+-**`ruff/flake8`**: Linter for ensuring clean and consistent code.
+
+### Contributing
+
+Contributions are always welcome! Feel free to fork the repository, make changes, and create a pull request. 
+I'm open to collaboration and willing to make adjustments to improve the project. 
+Let's build something great together! 🚀

core/compression.py

@@ -0,0 +1,104 @@
+#!/usr/bin/env python3
+# -*- coding: utf-8 -*-
+"""
+Author: C0ffymachyne
+License: GPLv3
+Version: 1.0.0
+
+Description:
+    various compression methods
+"""
+
+
+import cupy as cp
+import torch
+import numpy as np
+
+from typing import Tuple, Any
+from ..core.utilitiescuda import read_kernel_by_name
+
+compressor_kernel = read_kernel_by_name(
+    "compressor", kernel_class="compressor", kernel_identifier="compexp_kernel"
+)
+
+
+def compressor(
+    audio_in: torch.Tensor,
+    sample_rate: int,
+    comp: float = -0.3,  # Compression/expansion factor
+    attack: float = 0.1,  # Attack time in ms
+    release: float = 60.0,  # Release time in ms
+    a: float = 0.3,  # Filter parameter < 1
+    device: str = "cuda",
+) -> Tuple[torch.Tensor, Any]:
+    """
+    Compresses or expands stereo audio using an optimized CUDA kernel.
+
+    Parameters:
+        audio_in (torch.Tensor or np.ndarray): Input stereo audio signal with shape (n_samples, 2).
+        sample_rate (int): Sampling rate in Hz.
+        comp (float): Compression/expansion factor.
+        attack (float): Attack time in milliseconds.
+        release (float): Release time in milliseconds.
+        a (float): Filter parameter (< 1) for envelope smoothing.
+        device (str): Device to place the output tensor ('cuda' or 'cpu').
+
+    Returns:
+        torch.Tensor: Compressed stereo audio with shape (n_samples, 2).
+    """
+    # Convert input to NumPy array if necessary
+    audio_in = audio_in.T
+    if isinstance(audio_in, torch.Tensor):
+        audio_in = audio_in.detach().cpu().numpy()
+    else:
+        audio_in = np.asarray(audio_in, dtype=np.float64)
+
+    # 2. Ensure the audio is in shape (n_samples, 2)
+    if audio_in.ndim != 2 or audio_in.shape[1] != 2:
+        raise ValueError(
+            f"Input audio must have shape (n_samples, 2), but got {audio_in.shape}"
+        )
+
+    n_samples, n_channels = audio_in.shape
+
+    # Flatten the audio for kernel processing
+    wav_in_flat = audio_in.flatten()
+
+    # Move data to GPU
+    wav_in_gpu = cp.asarray(wav_in_flat, dtype=cp.float64)
+    wav_out_gpu = cp.zeros_like(wav_in_gpu)
+
+    # Define grid and block dimensions
+    block_size = 256  # Number of threads per block
+    grid_size = n_channels  # One block per channel
+
+    # Launch the CUDA kernel
+    compressor_kernel(
+        (grid_size,),  # Grid dimensions
+        (block_size,),  # Block dimensions
+        (
+            wav_in_gpu,
+            wav_out_gpu,
+            np.int32(n_channels),
+            np.int32(n_samples),
+            np.float64(comp),
+            np.float64(release),
+            np.float64(attack),
+            np.float64(a),
+            np.float64(sample_rate),
+        ),
+    )
+
+    # Retrieve results from GPU
+    wav_out_host = wav_out_gpu.get().astype(np.float64)
+
+    # Reshape the output
+    wav_out_stereo = wav_out_host.reshape((n_samples, n_channels))
+
+    # Convert back to Torch tensor if desired
+    if device == "cuda":
+        out_tensor = torch.from_numpy(wav_out_stereo).to("cuda")
+    else:
+        out_tensor = torch.from_numpy(wav_out_stereo).to("cpu")
+
+    return out_tensor.T, None

core/conversions.py

@@ -12,14 +12,14 @@ Description:
 import torch
 
 
-def db_to_lin(value):
+def db_to_lin(value: float) -> float:
     return 10 ** (value / 20)
 
 
-def lin_to_tb(value):
+def lin_to_tb(value: float) -> torch.Tensor:
     return 20 * torch.log10(torch.abs(value) + 1.0e-24)
 
 
-def get_sign(value):
+def get_sign(value: float) -> torch.Tensor:
     sign = torch.sign(value)
     return sign

core/filters.py

@@ -0,0 +1,161 @@
+#!/usr/bin/env python3
+# -*- coding: utf-8 -*-
+"""
+Author: C0ffymachyne
+License: GPLv3
+Version: 1.0.0
+
+Description:
+    utility functions for various filtering tasks
+"""
+
+import torch
+from scipy.signal import butter
+import torchaudio.functional as AF
+
+
+def anti_aliasing_iir_filter(
+    audio: torch.Tensor, sample_rate: int, cutoff: float = 0.0, order: int = 4
+) -> torch.Tensor:
+    """
+    Apply an anti-aliasing IIR low-pass filter to the audio.
+
+    Parameters:
+        audio (Tensor): [channels, samples] audio signal.
+        sr (int): Sample rate.
+        cutoff (float): Cutoff frequency for the filter. Defaults to Nyquist limit (sr / 2).
+        order (int): Order of the IIR filter.
+
+    Returns:
+        Tensor: Filtered audio signal.
+    """
+    if cutoff == 0:
+        cutoff = sample_rate / 2  # Default to Nyquist frequency
+    nyquist = sample_rate / 2
+    normalized_cutoff = cutoff / nyquist - 0.01
+
+    # Design the Butterworth filter
+    b, a = butter(order, normalized_cutoff, btype="low", output="ba")
+    b = torch.tensor(b, dtype=audio.dtype, device=audio.device)
+    a = torch.tensor(a, dtype=audio.dtype, device=audio.device)
+
+    # Apply the filter
+    filtered_audio = AF.lfilter(audio, b_coeffs=b, a_coeffs=a)
+
+    return filtered_audio
+
+
+def band_stop_filter(
+    audio: torch.Tensor,
+    sample_rate: int,
+    low_cut: float,
+    high_cut: float,
+    filter_order: int = 2,
+) -> torch.Tensor:
+    """
+    Apply a band-stop filter to attenuate lower mid frequencies.
+
+    Parameters:
+        audio (Tensor): [channels, samples] input audio signal.
+        low_cut (float): Lower cutoff frequency of the band in Hz.
+        high_cut (float): Upper cutoff frequency of the band in Hz.
+        sr (int): Sample rate in Hz.
+        filter_order (int): Order of the Butterworth filter.
+
+    Returns:
+        Tensor: Audio signal after band-stop filtering.
+    """
+    nyquist = sample_rate / 2
+    normalized_band = [low_cut / nyquist, high_cut / nyquist]
+
+    # Design band-stop Butterworth filter
+    b, a = butter(filter_order, normalized_band, btype="bandstop", analog=False)
+
+    # Convert coefficients to Torch tensors
+    b = torch.tensor(b, dtype=audio.dtype, device=audio.device)
+    a = torch.tensor(a, dtype=audio.dtype, device=audio.device)
+
+    # Ensure [channels, samples] format
+    if audio.dim() == 1:
+        audio = audio.unsqueeze(0)
+
+    # Apply the filter using lfilter
+    filtered_audio = AF.lfilter(audio, a_coeffs=a, b_coeffs=b, clamp=False)
+
+    return filtered_audio.squeeze(0) if filtered_audio.size(0) == 1 else filtered_audio
+
+
+def low_pass_filter(
+    audio: torch.Tensor, sampler_rate: int, cutoff_freq: float, filter_order: int = 4
+) -> torch.Tensor:
+    """
+    Apply a low-pass filter using a Butterworth filter.
+    """
+    # Design Butterworth filter using SciPy
+    nyquist = sampler_rate / 2
+    normalized_cutoff = cutoff_freq / nyquist
+    b, a = butter(filter_order, normalized_cutoff, btype="low", analog=False)
+
+    # Convert coefficients to Torch tensors
+    b = torch.tensor(b, dtype=audio.dtype, device=audio.device)
+    a = torch.tensor(a, dtype=audio.dtype, device=audio.device)
+
+    # Ensure [channels, samples] format
+    if audio.dim() == 1:
+        audio = audio.unsqueeze(0)
+
+    # Apply the filter using torchaudio
+    filtered_audio = AF.lfilter(audio, a_coeffs=a, b_coeffs=b, clamp=False)
+
+    return filtered_audio.squeeze(0) if filtered_audio.size(0) == 1 else filtered_audio
+
+
+def butter_filter(
+    audio: torch.Tensor,
+    sample_rate: int,
+    cutoff_freq: float,
+    filter_type: str = "low",
+    order: int = 4,
+) -> torch.Tensor:
+    """
+    Create and apply a Butterworth filter (low-pass or high-pass).
+
+    Parameters:
+        audio (Tensor): [channels, samples] input audio signal.
+        cutoff_freq (float): Cutoff frequency in Hz.
+        sr (int): Sample rate in Hz.
+        filter_type (str): "low" for low-pass, "high" for high-pass.
+        order (int): Filter order.
+
+    Returns:
+        Tensor: Filtered audio signal.
+    """
+    nyquist = sample_rate / 2
+    normalized_cutoff = cutoff_freq / nyquist
+    b, a = butter(order, normalized_cutoff, btype=filter_type, analog=False)
+
+    b = torch.tensor(b, dtype=audio.dtype, device=audio.device)
+    a = torch.tensor(a, dtype=audio.dtype, device=audio.device)
+
+    # Ensure [channels, samples] format
+    if audio.dim() == 1:
+        audio = audio.unsqueeze(0)
+
+    filtered_audio = AF.lfilter(audio, a_coeffs=a, b_coeffs=b, clamp=False)
+    return filtered_audio.squeeze(0) if filtered_audio.size(0) == 1 else filtered_audio
+
+
+def butter_low_pass(
+    audio: torch.Tensor, sample_rate: int, cutoff_freq: float, order: int = 4
+) -> torch.Tensor:
+    return butter_filter(
+        audio, sample_rate, cutoff_freq, filter_type="low", order=order
+    )
+
+
+def butter_high_pass(
+    audio: torch.Tensor, sample_rate: int, cutoff_freq: int, order: int = 4
+) -> torch.Tensor:
+    return butter_filter(
+        audio, sample_rate, cutoff_freq, filter_type="high", order=order
+    )

core/harmonics.py

@@ -0,0 +1,180 @@
+#!/usr/bin/env python3
+# -*- coding: utf-8 -*-
+"""
+Author: C0ffymachyne
+License: GPLv3
+Version: 1.0.0
+
+Description:
+    various plotting methods for debugging and visualization
+"""
+
+
+import torch
+import torchaudio
+import torchaudio.functional as F
+from typing import List
+
+
+def enhance_harmonics(
+    audio: torch.Tensor,
+    sample_rate: int,
+    harmonics: List[int] = [1, 3, 5, 7, 9, 11],
+    gain_db: float = 5,
+    base_frequency: float = 0,
+    Q: float = 0.707,
+) -> torch.Tensor:
+
+    pitch = F.detect_pitch_frequency(audio, sample_rate)
+    base_frequency = pitch.mean().item()
+    if base_frequency <= 0:  # Fallback if pitch detection fails
+        base_frequency = 440  # Use a default base frequency
+
+    # Apply EQ boosts to specific harmonic frequencies
+    for harmonic in harmonics:
+        freq = base_frequency * harmonic
+        if freq < sample_rate / 2:  # Ensure it's within the Nyquist frequency
+            audio = F.equalizer_biquad(
+                audio, sample_rate, center_freq=freq, gain=gain_db, Q=Q
+            )
+
+    return audio
+
+
+def enhance_harmonics2(
+    audio: torch.Tensor,
+    sample_rate: int,
+    harmonics: list[int] = [1, 2, 3, 4, 5],
+    gain_db: float = 5,
+    base_frequency: float = 0,
+    Q: float = 0.707,
+) -> torch.Tensor:
+    """
+    Enhance specified harmonics in an audio signal, emulating Distressor-like harmonic enhancement.
+    Parameters:
+        audio (Tensor): Input audio signal (1D or 2D [channels, samples]).
+        sample_rate (int): Sampling rate of the audio.
+        harmonics (list): List of harmonic multipliers to enhance.
+        gain_db (float): Gain to apply to each harmonic.
+        base_frequency (float, optional): Fundamental frequency. If None, it will be estimated.
+        Q (float): Quality factor for the EQ bands.
+    Returns:
+        Tensor: Audio signal with enhanced harmonics.
+    """
+
+    if base_frequency == 0:
+        # Detect the pitch frequency using torchaudio's pitch detection
+        pitch = F.detect_pitch_frequency(audio, sample_rate)
+        base_frequency = pitch.mean().item()
+        if base_frequency <= 0:  # Fallback if pitch detection fails
+            base_frequency = 440  # Default base frequency (A4)
+
+    # Create a copy of the input signal for processing
+    processed_audio: torch.Tensor = audio.clone()
+
+    # Enhance harmonics using biquad EQ for precision
+    for harmonic in harmonics:
+        freq = base_frequency * harmonic
+        if freq < sample_rate / 2:  # Ensure frequency is within Nyquist limit
+            processed_audio = F.equalizer_biquad(
+                processed_audio, sample_rate, center_freq=freq, gain=gain_db, Q=Q
+            )
+
+    # Apply a non-linear saturation for warmth and further harmonic enhancement
+    def non_linear_saturation(audio: torch.Tensor, drive: float = 1.0) -> torch.Tensor:
+        k = torch.tensor(1.0 + drive, dtype=audio.dtype, device=audio.device)
+        return torch.tanh(k * audio) / torch.tanh(k)
+
+    processed_audio = non_linear_saturation(processed_audio, drive=gain_db / 10.0)
+
+    # Blend processed harmonics with the original signal
+    output_audio = audio + processed_audio * (gain_db / 20.0)  # Scale blend by gain
+    return output_audio / torch.max(torch.abs(output_audio))  # Normalize output
+
+
+def batch_equalizer_biquad(
+    audio: torch.Tensor, sample_rate: int, freqs: torch.Tensor, gain_db: float, Q: float
+) -> torch.Tensor:
+    """
+    Apply biquad filters to enhance multiple harmonics in a batch.
+    Parameters:
+        audio (Tensor): Input audio signal (1D or 2D [channels, samples]).
+        sample_rate (int): Sampling rate of the audio.
+        freqs (Tensor): Frequencies for biquad filters.
+        gain_db (float): Gain to apply to each harmonic.
+        Q (float): Quality factor for all filters.
+    Returns:
+        Tensor: Audio signal with harmonics enhanced.
+    """
+    audio = audio.unsqueeze(0) if audio.dim() == 1 else audio
+
+    # Precompute filter coefficients for all frequencies
+    coeffs = [
+        torchaudio.functional.equalizer_biquad(
+            audio, sample_rate, center_freq=f, gain=gain_db, Q=Q
+        )
+        for f in freqs
+    ]
+
+    # Sum the filtered outputs for all harmonics
+    filtered_audio = sum(coeffs)
+
+    return filtered_audio
+
+
+def enhance_harmonics3(
+    audio: torch.Tensor,
+    sample_rate: int,
+    harmonics: List[int] = [1, 2, 3, 4, 5],
+    gain_db: float = 5.0,
+    base_frequency: float = 0,
+    Q: float = 0.707,
+) -> torch.Tensor:
+    """
+    Enhance specified harmonics in an audio signal, efficiently processing harmonics in a batch.
+    Parameters:
+        audio (Tensor): Input audio signal (1D or 2D [channels, samples]).
+        sample_rate (int): Sampling rate of the audio.
+        harmonics (list): List of harmonic multipliers to enhance.
+        gain_db (float): Gain to apply to each harmonic.
+        base_frequency (float, optional): Fundamental frequency. If None, it will be estimated.
+        Q (float): Quality factor for the EQ bands.
+    Returns:
+        Tensor: Audio signal with enhanced harmonics.
+    """
+    if base_frequency is None:
+        # Detect the pitch frequency using torchaudio's pitch detection
+        pitch = torchaudio.functional.detect_pitch_frequency(audio, sample_rate)
+        base_frequency = pitch.mean().item()
+        if base_frequency <= 0:  # Fallback if pitch detection fails
+            base_frequency = 440.0  # Default base frequency (A4)
+
+    # Calculate harmonic frequencies
+    harmonic_freqs = torch.tensor(
+        [base_frequency * h for h in harmonics], device=audio.device
+    )
+
+    # Ensure frequencies are within Nyquist limit
+    harmonic_freqs = harmonic_freqs[harmonic_freqs < sample_rate / 2]
+
+    # Apply batched harmonic enhancement
+    processed_audio = batch_equalizer_biquad(
+        audio, sample_rate, harmonic_freqs, gain_db, Q
+    )
+
+    # Normalize the output
+    return processed_audio / torch.max(torch.abs(processed_audio))
+
+
+def enahnce_harmonics_23(
+    audio: torch.Tensor, sample_rate: int, gain_db_base: int = 0, Q: float = 0.707
+) -> torch.Tensor:
+
+    audio = enhance_harmonics3(
+        audio, sample_rate, harmonics=[2], gain_db=gain_db_base + 1, Q=0.303
+    )
+    audio = enhance_harmonics3(
+        audio, sample_rate, harmonics=[3], gain_db=gain_db_base + 3, Q=0.303
+    )
+
+    return audio

core/io.py

@@ -9,26 +9,27 @@ Description:
     io for audio files and comfy
 """
 
-import torch, torchaudio
+import torch
+import torchaudio
 from pathlib import Path
-from typing import Dict, Tuple, Any
+from typing import Dict, Tuple, Any, Union
 import numpy as np
 
-# this mess needs some cleanup !!
+debug_print = False
 
 
-# when we return audio to comfy, turn it to cpu
 def audio_to_comfy_3d(
-    waveform: torch.Tensor, sample_rate: int, cpu=True, clip: bool = False
-) -> Tuple[Dict]:
+    waveform: torch.Tensor, sample_rate: int, cpu: bool = True, clip: bool = False
+) -> Tuple[Dict[str, Union[torch.Tensor, int]]]:
 
-    # try to fix wrong shape
     if waveform.ndim == 2:
         waveform = waveform.unsqueeze(0)
 
     if waveform.ndim != 3:
         raise RuntimeError(
-            f"returning to_comfy failed {waveform.shape}@{waveform.device}[{waveform.ndim}] has incorrect dimentions, it should be of dimension [batch,channels,audio]"
+            f"returning to_comfy failed \
+                {waveform.shape}@{waveform.device}[{waveform.ndim}]  \
+            has incorrect dimentions, it should be of dimension [batch,channels,audio]"
         )
 
     if cpu:
@@ -43,25 +44,32 @@ def audio_to_comfy_3d(
 
     return_value = ({"waveform": waveform, "sample_rate": sample_rate},)
 
-    print(
-        f"audio_to_comfy_3d '{return_value}' with shape {waveform.shape}@{waveform.device}dim[{waveform.ndim}] and sample rate {sample_rate} Hz."
-    )
+    if debug_print:
+        print(
+            f"audio_to_comfy_3d '{return_value}' with shape \
+                {waveform.shape}@{waveform.device}dim[{waveform.ndim}] \
+                    and sample rate {sample_rate} Hz."
+        )
 
     return return_value
 
 
-# when we return audio to comfy, turn it to cpu
 def audio_to_comfy_3dp1(
-    waveform: torch.Tensor, sample_rate: int, value0: Any, cpu=True, clip: bool = False
-) -> Tuple[Dict]:
+    waveform: torch.Tensor,
+    sample_rate: int,
+    value0: Any,
+    cpu: bool = True,
+    clip: bool = False,
+) -> Tuple[Dict[str, torch.Tensor], Any]:
 
-    # try to fix wrong shape
     if waveform.ndim == 2:
         waveform = waveform.unsqueeze(0)
 
     if waveform.ndim != 3:
         raise RuntimeError(
-            f"returning to_comfy failed {waveform.shape}@{waveform.device}dim[{waveform.ndim}] has incorrect dimentions, it should be of dimension [batch,channels,audio]"
+            f"returning to_comfy failed \
+                 {waveform.shape}@{waveform.device}dim[{waveform.ndim}]\
+                     has incorrect dimentions, it should be of dimension [batch,channels,audio]"
         )
 
     if cpu:
@@ -79,30 +87,33 @@ def audio_to_comfy_3dp1(
         value0,
     )
 
-    print(
-        f"audio_to_comfy_3d '{return_value}' with shape {waveform.shape}@{waveform.device}dim[{waveform.ndim}] and sample rate {sample_rate} Hz."
-    )
+    if debug_print:
+        print(
+            f"audio_to_comfy_3d '{return_value}' with shape \
+            {waveform.shape}@{waveform.device}dim[{waveform.ndim}] \
+                and sample rate {sample_rate} Hz."
+        )
 
     return return_value
 
 
-# when we return audio to comfy, turn it to cpu
 def audio_to_comfy_3dp2(
     waveform: torch.Tensor,
     sample_rate: int,
     value0: Any,
     value1: Any,
-    cpu=True,
+    cpu: bool = True,
     clip: bool = False,
-) -> Tuple[Dict]:
+) -> Tuple[Dict[str, torch.Tensor], Any, Any]:
 
-    # try to fix wrong shape
     if waveform.ndim == 2:
         waveform = waveform.unsqueeze(0)
 
     if waveform.ndim != 3:
         raise RuntimeError(
-            f"returning to_comfy failed {waveform.shape}@{waveform.device}dim[{waveform.ndim}] has incorrect dimentions, it should be of dimension [batch,channels,audio]"
+            f"returning to_comfy failed \
+                {waveform.shape}@{waveform.device}dim[{waveform.ndim}] \
+                    has incorrect dimentions, it should be of dimension [batch,channels,audio]"
         )
 
     if cpu:
@@ -121,23 +132,27 @@ def audio_to_comfy_3dp2(
         value1,
     )
 
-    print(
-        f"audio_to_comfy_3d '{return_value}' with shape {waveform.shape}@{waveform.device}dim[{waveform.ndim}] and sample rate {sample_rate} Hz."
-    )
+    if debug_print:
+        print(
+            f"audio_to_comfy_3d '{return_value}' with shape \
+            {waveform.shape}@{waveform.device}dim[{waveform.ndim}] \
+                and sample rate {sample_rate} Hz."
+        )
 
     return return_value
 
 
-# when we get audio from comfy, load it to gpu
 def audio_from_comfy_2d(
-    audio: Tuple[Dict], repeat=True, try_gpu=True
+    audio: Dict[str, Union[torch.Tensor, int]],
+    repeat: bool = True,
+    try_gpu: bool = True,
 ) -> Tuple[torch.Tensor, int]:
 
     if audio is None:
         raise ValueError("Input audio must be provided.")
 
-    waveform = audio.get("waveform")  # [batch, channels, samples]
-    sample_rate = audio.get("sample_rate")
+    waveform: torch.Tensor = audio.get("waveform")  # [batch, channels, samples]
+    sample_rate: int | Any | None = audio.get("sample_rate")
 
     if waveform.dtype != torch.float64:
         waveform = waveform.to(dtype=torch.float64)
@@ -148,43 +163,66 @@ def audio_from_comfy_2d(
     if not isinstance(waveform, torch.Tensor):
         raise TypeError("Waveform must be a torch.Tensor.")
 
-    print(
-        f"audio_from_comfy_2d with shape {waveform.shape}@{waveform.device}dim[{waveform.ndim}] and sample rate {sample_rate} Hz."
-    )
+    if debug_print:
+        print(
+            f"audio_from_comfy_2d with shape \
+            {waveform.shape}@{waveform.device}dim[{waveform.ndim}] \
+                and sample rate {sample_rate} Hz."
+        )
 
     if waveform.ndim != 3:
         raise ValueError(
-            f"Waveform must be a 3D tensor with shape (batch, channels, samples) not {waveform.shape}@{waveform.device}"
+            f"Waveform must be a 3D tensor with shape (batch, channels, samples) not \
+                {waveform.shape}@{waveform.device}"
         )
 
     if repeat:
         if waveform.ndim == 1:  # add extra channel in case of mono
-            print(f"audio_from_comfy_2d repeat")
+            print("audio_from_comfy_2d repeat")
             waveform = waveform.unsqueeze(0)  # Add channel dimension
             waveform = waveform.repeat(2, 1)  # copy mono to the new channel
 
-    waveform = waveform.squeeze(0).contiguous()
+    waveform = waveform.squeeze(0)
+    # waveform = waveform.contiguous()
 
     if try_gpu:
         device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
         waveform = waveform.to(device)
 
-    print(
-        f"audio_from_comfy_2d with shape {waveform.shape}@{waveform.device}dim[{waveform.ndim}] and sample rate {sample_rate} Hz."
-    )
+    if debug_print:
+        print(
+            f"audio_from_comfy_2d with shape \
+            {waveform.shape}@{waveform.device}dim[{waveform.ndim}] \
+                and sample rate {sample_rate} Hz."
+        )
 
     return waveform, sample_rate
 
 
+def audio_from_comfy_3d_to_disk(
+    audio: Dict[str, Union[torch.Tensor, int]],
+    filepath: str,
+    repeat: bool = True,
+    try_gpu: bool = True,
+) -> None:
+    waveform: torch.Tensor = audio.get("waveform")  # [batch, channels, samples]
+    sample_rate: int | Any | None = audio.get("sample_rate")
+    waveform = waveform.squeeze(0)
+    waveform = waveform.cpu()
+    torchaudio.save(filepath, waveform, sample_rate, buffer_size=4096 * 16)
+
+
 def audio_from_comfy_3d(
-    audio: Tuple[Dict], repeat=True, try_gpu=True
+    audio: Dict[str, Union[torch.Tensor, int]],
+    repeat: bool = True,
+    try_gpu: bool = True,
 ) -> Tuple[torch.Tensor, int]:
 
     if audio is None:
         raise ValueError("Input audio must be provided.")
 
-    waveform = audio.get("waveform")  # [batch, channels, samples]
-    sample_rate = audio.get("sample_rate")
+    waveform: torch.Tensor = audio.get("waveform")  # [batch, channels, samples]
+    sample_rate: Union[torch.Tensor, int] = audio.get("sample_rate")
 
     if waveform.dtype != torch.float64:
         waveform = waveform.to(dtype=torch.float64)
@@ -195,18 +233,23 @@ def audio_from_comfy_3d(
     if not isinstance(waveform, torch.Tensor):
         raise TypeError("Waveform must be a torch.Tensor.")
 
-    print(
-        f"audio_from_comfy_2d with shape {waveform.shape}@{waveform.device}dim[{waveform.ndim}] and sample rate {sample_rate} Hz."
-    )
+    if debug_print:
+        print(
+            f"audio_from_comfy_2d with shape \
+            {waveform.shape}@{waveform.device}dim[{waveform.ndim}] \
+                and sample rate {sample_rate} Hz."
+        )
 
     if waveform.ndim != 3:
         raise ValueError(
-            f"Waveform must be a 3D tensor with shape (batch, channels, samples) not {waveform.shape}@{waveform.device}"
+            f"Waveform must be a 3D tensor with shape \
+                (batch, channels, samples) not \
+                    {waveform.shape}@{waveform.device}"
         )
 
     if repeat:
         if waveform.ndim == 1:  # add extra channel in case of mono
-            print(f"audio_from_comfy_2d repeat")
+            print("audio_from_comfy_2d repeat")
             waveform = waveform.unsqueeze(0)  # Add channel dimension
             waveform = waveform.repeat(2, 1)  # copy mono to the new channel
 
@@ -216,9 +259,12 @@ def audio_from_comfy_3d(
         device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
         waveform = waveform.to(device)
 
-    print(
-        f"audio_from_comfy_3d with shape {waveform.shape}@{waveform.device}dim[{waveform.ndim}] and sample rate {sample_rate} Hz."
-    )
+    if debug_print:
+        print(
+            f"audio_from_comfy_3d with shape \
+            {waveform.shape}@{waveform.device}dim[{waveform.ndim}] \
+                and sample rate {sample_rate} Hz."
+        )
 
     return waveform, sample_rate
 
@@ -226,12 +272,10 @@ def audio_from_comfy_3d(
 def from_disk_as_dict_3d(
     audio_file: str,
     gain: float = 1.0,
-    repeat=True,
+    repeat: bool = True,
     start_time: float = 0.0,
     end_time: float = 0.0,
-) -> Dict:
-    # all loaded audio comes in 3 dimensions
-    # [batch,channels,audio] 3 dimensions
+) -> Tuple[Dict[str, Union[torch.Tensor, int]]]:
 
     if not Path(audio_file).exists():
         raise ValueError(f"File {audio_file} Does Not Exist")
@@ -260,44 +304,62 @@ def from_disk_as_dict_3d(
     # add batch dimension at the front
     waveform = waveform.unsqueeze(0).contiguous()
 
-    # waveform = waveform * gain
-
-    print(
-        f"from_disk_as_dict_3d loaded audio file: '{audio_file}' with shape: {waveform.shape}@{waveform.device}dim[{waveform.ndim}]  and sample rate: {sample_rate} Hz."
-    )
+    if debug_print:
+        print(
+            f"from_disk_as_dict_3d loaded audio file: '{audio_file}' with shape: \
+            {waveform.shape}@{waveform.device}dim[{waveform.ndim}] \
+                and sample rate: {sample_rate} Hz."
+        )
 
     return audio_to_comfy_3d(waveform, sample_rate)
 
 
 def from_disk_as_raw_3d(
-    audio_file: str, repeat: bool = True, try_gpu: bool = False
+    audio_file: str,
+    repeat: bool = True,
+    try_gpu: bool = False,
+    start_seconds: float = 0.0,
+    duration_seconds: float = 0.0,
 ) -> Tuple[torch.Tensor, int]:
-    # all loaded audio comes in 3 dimensions
-    # [batch,channels,audio] 3 dimensions
 
-    waveform, sample_rate = torchaudio.load(audio_file)
-    info = torchaudio.info(audio_file)
-    print(
-        f"from_disk_as_raw_3d loaded audio file: '{audio_file}' with shape {waveform.shape}@{waveform.device}dim[{waveform.ndim}]  and sample rate {sample_rate} Hz."
-    )
+    waveform, sample_rate = torchaudio.load(audio_file, buffer_size=4096 * 16)
+    # info = torchaudio.info(audio_file)
+    if debug_print:
+        print(
+            f"from_disk_as_raw_3d loaded audio file: '{audio_file}' with shape \
+            {waveform.shape}@{waveform.device}dim[{waveform.ndim}] \
+                and sample rate {sample_rate} Hz."
+        )
+
+    if duration_seconds != 0.0:
+        start_sample = int(start_seconds * sample_rate)
+        end_sample = None
+        if duration_seconds is not None:
+            end_sample = start_sample + int(duration_seconds * sample_rate)
+        print(f"slicing from_disk_as_raw_3d: [{start_seconds}:{duration_seconds}]")
+        # Slice the waveform
+        waveform = waveform[:, start_sample:end_sample]
 
     if waveform.dtype != torch.float64:
         waveform = waveform.to(dtype=torch.float64)
 
     if repeat:
-        if info.num_channels == 1:  # add extra channel in case of mono
+        # if info.num_channels == 1:  # add extra channel in case of mono
+        if waveform.ndim == 1:  # add extra channel in case of mono
             waveform = waveform.repeat(2, 1)  # copy mono to the new channel
 
-    # add batch dimension at the front
-    waveform = waveform.unsqueeze(0).contiguous()
+    waveform = waveform.unsqueeze(0)
 
     if try_gpu:
         device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
         waveform = waveform.to(device)
 
-    print(
-        f"from_disk_as_raw_3d loaded audio file: '{audio_file}' with shape {waveform.shape}@{waveform.device}dim[{waveform.ndim}]  and sample rate {sample_rate} Hz."
-    )
+    if debug_print:
+        print(
+            f"from_disk_as_raw_3d loaded audio file: '{audio_file}' with shape \
+            {waveform.shape}@{waveform.device}dim[{waveform.ndim}] \
+                and sample rate {sample_rate} Hz."
+        )
 
     return waveform, sample_rate
 
@@ -311,9 +373,12 @@ def from_disk_as_raw_2d(
     waveform, sample_rate = torchaudio.load(audio_file)
     info = torchaudio.info(audio_file)
 
-    print(
-        f"from_disk_as_raw_2d loaded audio file: '{audio_file}' with shape {waveform.shape}@{waveform.device}dim[{waveform.ndim}]  and sample rate {sample_rate} Hz."
-    )
+    if debug_print:
+        print(
+            f"from_disk_as_raw_2d loaded audio file: '{audio_file}' with shape \
+            {waveform.shape}@{waveform.device}dim[{waveform.ndim}] \
+                and sample rate {sample_rate} Hz."
+        )
 
     if waveform.dtype != torch.float64:
         waveform = waveform.to(dtype=torch.float64)
@@ -327,8 +392,11 @@ def from_disk_as_raw_2d(
         device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
         waveform = waveform.to(device)
 
-    print(
-        f"from_disk_as_raw_2d loaded audio file: '{audio_file}' with shape {waveform.shape}@{waveform.device}dim[{waveform.ndim}] and sample rate {sample_rate} Hz."
-    )
+    if debug_print:
+        print(
+            f"from_disk_as_raw_2d loaded audio file: '{audio_file}' with shape \
+            {waveform.shape}@{waveform.device}dim[{waveform.ndim}] \
+                and sample rate {sample_rate} Hz."
+        )
 
     return waveform, sample_rate

core/kernels/compressor/compressor.cu

@@ -0,0 +1,52 @@
+#ifndef M_PI
+#define M_PI 3.14159265358979323846
+#endif
+
+extern "C" __global__
+void compexp_kernel(
+    const double* wav_in,       // Input audio signal [L0, R0, L1, R1, ..., LN-1, RN-1]
+    double* wav_out,            // Output audio signal
+    const int n_channels,       // Number of channels (e.g., 2 for stereo)
+    const int n_samples,        // Number of samples per channel
+    const double comp,          // Compression/expansion factor
+    const double release,       // Release time in ms
+    const double attack,        // Attack time in ms
+    const double a,             // Filter parameter < 1
+    const double Fs             // Sampling rate in Hz
+) {
+    int ch = blockIdx.x;  // Each block processes one channel
+    int thread_id = threadIdx.x; // Thread within the block
+    int stride = blockDim.x;     // Number of threads in the block
+
+    if (ch >= n_channels) return;
+
+    double attack_coeff = exp(-1.0 / (Fs * (attack * 1e-3)));
+    double release_coeff = exp(-1.0 / (Fs * (release * 1e-3)));
+
+    double h = 0.0;  // Initialize filter state for envelope detection
+
+    // Divide samples across threads in parallel
+    for (int i = thread_id; i < n_samples; i += stride) {
+        int sample_idx = i * n_channels + ch;
+        double sample = wav_in[sample_idx];
+
+        // Envelope detection using attack/release dynamics
+        double abs_sample = fabs(sample);
+        if (abs_sample > h) {
+            h = attack_coeff * (h - abs_sample) + abs_sample;
+        } else {
+            h = release_coeff * (h - abs_sample) + abs_sample;
+        }
+
+        // Apply compression/expansion
+        double gain;
+        if (comp > 0) { // Compression: attenuate higher envelope values
+            gain = pow(h + 1e-8, -comp);
+        } else { // Expansion: boost lower envelope values
+            gain = pow(h + 1e-8, -comp);
+        }
+
+        // Scale output
+        wav_out[sample_idx] = sample * gain;
+    }
+}

core/kernels/limiter/limiter_dev.cu

@@ -0,0 +1,60 @@
+#ifndef M_PI
+#define M_PI 3.14159265358979323846
+#endif
+
+extern "C" __global__
+void limiter_kernel(
+    const double*__restrict__ wav_in,       // Input audio signal [L0, R0, L1, R1, ..., LN-1, RN-1]
+    double* __restrict__ wav_out,            // Output audio signal
+    double* __restrict__ debug_out,          // Debug buffer [envelope, gain]
+    const int n_channels,       // Number of channels (e.g., 2 for stereo)
+    const int n_samples,        // Number of samples per channel
+    double threshold,     // Threshold in percents (0-100)
+    double slope,         // Slope angle in percents (0-100)
+    const double sr,            // Sample rate (samples/sec)
+    double twnd,          // Window time for RMS in ms
+    double tatt,          // Attack time in ms
+    double trel           // Release time in ms
+) {
+    // Only one thread handles the entire stereo pair
+    int ch = blockIdx.x * blockDim.x + threadIdx.x; // Thread processes a single channel
+
+    if (ch >= n_channels) return;
+
+    double attack_coeff = exp(-1.0 / (sr * (tatt * 1e-3)));
+    double release_coeff = exp(-1.0 / (sr * (trel * 1e-3)));
+    double envelope = 0.00;
+    //threshold = .55;
+    //slope = 1.0;
+
+    for (int i = 0; i < n_samples; ++i) {
+
+        double sample = wav_in[i * n_channels + ch];
+
+        // Envelope tracking
+        double abs_sample = fabs(sample);
+        if (abs_sample > envelope) {
+            envelope = attack_coeff * (envelope - abs_sample) + abs_sample;
+        } else {
+            envelope = release_coeff * (envelope - abs_sample) + abs_sample;
+        }
+
+        // Gain calculation
+        double gain = 1.0;
+        if (envelope > threshold) {
+            gain = pow(10.0, -slope * (log10(envelope) - log10(threshold)));
+        }
+        // Upward compression below threshold
+        double upward_compression_gain = 1.0;
+        if (envelope < threshold && envelope > 0.0) {
+            upward_compression_gain = pow(10.0, slope * (log10(threshold) - log10(envelope)));
+        }
+        // Apply gain
+        wav_out[i * n_channels + ch] = sample * gain * upward_compression_gain;
+
+        // Debugging output (envelope and gain)
+        if (debug_out) {
+            debug_out[i * n_channels + ch] = gain;
+        }
+    }
+}

core/kernels/limiter/limiter_down_parallel.cu

@@ -0,0 +1,73 @@
+#ifndef M_PI
+#define M_PI 3.14159265358979323846
+#endif
+
+extern "C" __global__
+void limiter_kernel(
+    const double* __restrict__ wav_in,       // Input audio signal [L0, R0, L1, R1, ..., LN-1, RN-1]
+    double* __restrict__ wav_out,            // Output audio signal
+    const int n_channels,       // Number of channels (e.g., 2 for stereo)
+    const int n_samples,        // Number of samples per channel
+    const double threshold,     // Threshold in linear scale (e.g., 0.5 for 50%)
+    const double slope,         // Slope parameter for gain calculation
+    const double sr,            // Sample rate in Hz
+    const double tatt,          // Attack time in ms
+    const double trel           // Release time in ms
+) {
+    int ch = blockIdx.x;  // Each block processes one channel
+    int thread_id = threadIdx.x; // Thread within the block
+    int stride = blockDim.x;     // Number of threads in the block
+
+    if (ch >= n_channels) return;
+
+    double attack_coeff = exp(-1.0 / (sr * (tatt * 1e-3)));
+    double release_coeff = exp(-1.0 / (sr * (trel * 1e-3)));
+
+    // Use release time to determine precompute samples
+    int precompute_samples = int((trel * 1e-3) * sr); // Convert release time to samples
+
+    // Each thread computes its own range of samples
+    int start_idx = thread_id * (n_samples / stride);
+    int end_idx = (thread_id + 1) * (n_samples / stride);
+
+    // Extend range backward for precomputing
+    int precompute_start_idx = max(0, start_idx - precompute_samples);
+
+    double envelope = 0.0;  // Envelope tracking state
+
+    // Precompute envelope for the extra range
+    for (int i = precompute_start_idx; i < start_idx; ++i) {
+        int sample_idx = i * n_channels + ch;
+        double sample = wav_in[sample_idx];
+        double abs_sample = fabs(sample);
+
+        if (abs_sample > envelope) {
+            envelope = attack_coeff * (envelope - abs_sample) + abs_sample;
+        } else {
+            envelope = release_coeff * (envelope - abs_sample) + abs_sample;
+        }
+    }
+
+    // Process assigned range of samples
+    for (int i = start_idx; i < end_idx; ++i) {
+        int sample_idx = i * n_channels + ch;
+        double sample = wav_in[sample_idx];
+        double abs_sample = fabs(sample);
+
+        // Envelope tracking
+        if (abs_sample > envelope) {
+            envelope = attack_coeff * (envelope - abs_sample) + abs_sample;
+        } else {
+            envelope = release_coeff * (envelope - abs_sample) + abs_sample;
+        }
+
+        // Gain calculation for downward limiting
+        double gain = 1.0;
+        if (envelope > threshold) {
+            gain = pow(10.0, -slope * (log10(envelope) - log10(threshold)));
+        }
+
+        // Apply both gains
+        wav_out[sample_idx] = sample * gain;
+    }
+}

core/kernels/limiter/limiter_hard_clipper.cu

@@ -0,0 +1,45 @@
+#ifndef M_PI
+#define M_PI 3.14159265358979323846
+#endif
+
+extern "C" __global__
+void limiter_kernel(
+    const double* __restrict__ wav_in, // Input audio signal [L0, R0, L1, R1, ..., LN-1, RN-1]
+    double* __restrict__ wav_out,      // Output audio signal
+    const int n_channels,       // Number of channels (e.g., 2 for stereo)
+    const int n_samples,        // Number of samples per channel
+    const double threshold,     // Threshold in linear scale (e.g., 0.5 for 50%)
+    const double slope,         // Slope parameter for gain calculation
+    const double sr,            // Sample rate in Hz
+    const double tatt,          // Attack time in ms
+    const double trel           // Release time in ms
+)
+{
+    int ch = blockIdx.x;  // Each block processes one channel
+    int thread_id = threadIdx.x; // Thread within the block
+    int stride = blockDim.x;     // Number of threads in the block
+
+    if (ch >= n_channels) return;
+
+    int start_idx = thread_id * (n_samples / stride);
+    int end_idx = (thread_id + 1) * (n_samples / stride);
+
+    double clip_limit = threshold;
+    double clip_limit_inv = 1.0 / clip_limit;
+
+    for (int i = start_idx; i < end_idx; ++i)
+    {
+        int sample_idx = i * n_channels + ch;
+        double y = wav_in[sample_idx];
+
+        double abs = fabs(y);
+        if(abs >= clip_limit){
+            if (y < 0){
+                y = -clip_limit;
+            }else{
+                y = clip_limit;
+            }
+        }
+        wav_out[sample_idx] = y;
+    }
+}

core/kernels/limiter/limiter_soft_clipper.cu

@@ -0,0 +1,43 @@
+#ifndef M_PI
+#define M_PI 3.14159265358979323846
+#endif
+
+extern "C" __global__
+void limiter_kernel(
+    const double* __restrict__ wav_in, // Input audio signal [L0, R0, L1, R1, ..., LN-1, RN-1]
+    double* __restrict__ wav_out,      // Output audio signal
+    const int n_channels,       // Number of channels (e.g., 2 for stereo)
+    const int n_samples,        // Number of samples per channel
+    const double threshold,     // Threshold in linear scale (e.g., 0.5 for 50%)
+    const double slope,         // Slope parameter for gain calculation
+    const double sr,            // Sample rate in Hz
+    const double tatt,          // Attack time in ms
+    const double trel           // Release time in ms
+)
+{
+    int ch = blockIdx.x;  // Each block processes one channel
+    int thread_id = threadIdx.x; // Thread within the block
+    int stride = blockDim.x;     // Number of threads in the block
+
+    if (ch >= n_channels) return;
+
+    int start_idx = thread_id * (n_samples / stride);
+    int end_idx = (thread_id + 1) * (n_samples / stride);
+
+    double clip_limit = threshold;
+
+    for (int i = start_idx; i < end_idx; ++i){
+        int sample_idx = i * n_channels + ch;
+        double y = wav_in[sample_idx];
+
+        // cubic soft clipping
+        if (y <= -1.0) {
+            y = -2.0 / 3.0;
+        } else if (y >= 1.0) {
+            y = 2.0 / 3.0;
+        } else {
+            y = y - (1.0 / 3.0) * y * y * y;
+        }
+        wav_out[sample_idx] = y;
+    }
+}

core/kernels/limiter/limiter_updown_parallel.cu

@@ -0,0 +1,79 @@
+#ifndef M_PI
+#define M_PI 3.14159265358979323846
+#endif
+
+extern "C" __global__
+void limiter_kernel(
+    const double* __restrict__ wav_in,       // Input audio signal [L0, R0, L1, R1, ..., LN-1, RN-1]
+    double* __restrict__ wav_out,            // Output audio signal
+    const int n_channels,       // Number of channels (e.g., 2 for stereo)
+    const int n_samples,        // Number of samples per channel
+    const double threshold,     // Threshold in linear scale (e.g., 0.5 for 50%)
+    const double slope,         // Slope parameter for gain calculation
+    const double sr,            // Sample rate in Hz
+    const double tatt,          // Attack time in ms
+    const double trel           // Release time in ms
+) {
+    int ch = blockIdx.x;  // Each block processes one channel
+    int thread_id = threadIdx.x; // Thread within the block
+    int stride = blockDim.x;     // Number of threads in the block
+
+    if (ch >= n_channels) return;
+
+    double attack_coeff = exp(-1.0 / (sr * (tatt * 1e-3)));
+    double release_coeff = exp(-1.0 / (sr * (trel * 1e-3)));
+
+    // Use release time to determine precompute samples
+    int precompute_samples = int((trel * 1e-3) * sr); // Convert release time to samples
+
+    // Each thread computes its own range of samples
+    int start_idx = thread_id * (n_samples / stride);
+    int end_idx = (thread_id + 1) * (n_samples / stride);
+
+    // Extend range backward for precomputing
+    int precompute_start_idx = max(0, start_idx - precompute_samples);
+
+    double envelope = 0.0;  // Envelope tracking state
+
+    // Precompute envelope for the extra range
+    for (int i = precompute_start_idx; i < start_idx; ++i) {
+        int sample_idx = i * n_channels + ch;
+        double sample = wav_in[sample_idx];
+        double abs_sample = fabs(sample);
+
+        if (abs_sample > envelope) {
+            envelope = attack_coeff * (envelope - abs_sample) + abs_sample;
+        } else {
+            envelope = release_coeff * (envelope - abs_sample) + abs_sample;
+        }
+    }
+
+    // Process assigned range of samples
+    for (int i = start_idx; i < end_idx; ++i) {
+        int sample_idx = i * n_channels + ch;
+        double sample = wav_in[sample_idx];
+        double abs_sample = fabs(sample);
+
+        // Envelope tracking
+        if (abs_sample > envelope) {
+            envelope = attack_coeff * (envelope - abs_sample) + abs_sample;
+        } else {
+            envelope = release_coeff * (envelope - abs_sample) + abs_sample;
+        }
+
+        // Gain calculation for downward limiting
+        double gain = 1.0;
+        if (envelope > threshold) {
+            gain = pow(10.0, -slope * (log10(envelope) - log10(threshold)));
+        }
+
+        // Upward compression below threshold
+        double upward_compression_gain = 1.0;
+        if (envelope < threshold && envelope > 0.0) {
+            upward_compression_gain = pow(10.0, slope * (log10(threshold) - log10(envelope)));
+        }
+
+        // Apply both gains
+        wav_out[sample_idx] = sample * gain * upward_compression_gain;
+    }
+}

core/kernels/limiter/limiter_updown_weighted_parallel.cu

@@ -0,0 +1,85 @@
+#ifndef M_PI
+#define M_PI 3.14159265358979323846
+#endif
+
+extern "C" __global__
+void limiter_kernel(
+    const double* __restrict__ wav_in,       // Input audio signal [L0, R0, L1, R1, ..., LN-1, RN-1]
+    double* __restrict__ wav_out,            // Output audio signal
+    const int n_channels,       // Number of channels (e.g., 2 for stereo)
+    const int n_samples,        // Number of samples per channel
+    const double threshold,     // Threshold in linear scale (e.g., 0.5 for 50%)
+    const double slope,         // Slope parameter for gain calculation
+    const double sr,            // Sample rate in Hz
+    const double tatt,          // Attack time in ms
+    const double trel           // Release time in ms
+)
+{
+    int ch = blockIdx.x;  // Each block processes one channel
+    int thread_id = threadIdx.x; // Thread within the block
+    int stride = blockDim.x;     // Number of threads in the block
+
+    if (ch >= n_channels) return;
+
+    double attack_coeff = exp(-1.0 / (sr * (tatt * 1e-3)));
+    double release_coeff = exp(-1.0 / (sr * (trel * 1e-3)));
+
+    // Use release time to determine precompute samples
+    int precompute_samples = int((trel * 1e-3) * sr); // Convert release time to samples
+
+    // Each thread computes its own range of samples
+    int start_idx = thread_id * (n_samples / stride);
+    int end_idx = (thread_id + 1) * (n_samples / stride);
+
+    // Extend range backward for precomputing
+    int precompute_start_idx = max(0, start_idx - precompute_samples);
+
+    double envelope = 0.0;  // Envelope tracking state
+
+    // Precompute envelope for the extra range
+    for (int i = precompute_start_idx; i < start_idx; ++i) {
+        int sample_idx = i * n_channels + ch;
+        double sample = wav_in[sample_idx];
+        double abs_sample = fabs(sample);
+
+        if (abs_sample > envelope) {
+            envelope = attack_coeff * (envelope - abs_sample) + abs_sample;
+        } else {
+            envelope = release_coeff * (envelope - abs_sample) + abs_sample;
+        }
+    }
+
+    // Process assigned range of samples
+    for (int i = start_idx; i < end_idx; ++i) {
+        int sample_idx = i * n_channels + ch;
+        double sample = wav_in[sample_idx];
+        double abs_sample = fabs(sample);
+
+        // Envelope tracking
+        if (abs_sample > envelope) {
+            envelope = attack_coeff * (envelope - abs_sample) + abs_sample;
+        } else {
+            envelope = release_coeff * (envelope - abs_sample) + abs_sample;
+        }
+
+        // Gain calculation for downward limiting
+        double gain = 1.0;
+        if (envelope > threshold) {
+            gain = pow(10.0, -slope * (log10(envelope) - log10(threshold)));
+        }
+
+        // Upward compression below threshold with gradual application
+        double upward_compression_gain = 1.0;
+        if (envelope < threshold && envelope > 0.0) {
+            double t = 1.0 - threshold; // Scaling factor based on threshold
+            // Dynamic factor increases as envelope decreases
+            double dynamic_factor = (threshold - envelope) / threshold;
+            // Clamp to [0, 1]
+            dynamic_factor = fmin(fmax(dynamic_factor, 0.0), 1.0);
+            upward_compression_gain = 1.0 + t * dynamic_factor * pow(10.0, slope * (log10(threshold) - log10(envelope)));
+        }
+
+        // Apply both gains
+        wav_out[sample_idx] = sample * gain * upward_compression_gain;
+    }
+}

core/limiting.py

@@ -0,0 +1,149 @@
+#!/usr/bin/env python3
+# -*- coding: utf-8 -*-
+"""
+Author: C0ffymachyne
+License: GPLv3
+Version: 1.0.0
+
+Description:
+    various conversion methods
+"""
+
+import cupy as cp
+import torch
+import numpy as np
+from typing import List
+
+from ..core.utilitiescuda import read_kernel_by_name
+
+limiter_updown_parallel = read_kernel_by_name(
+    "limiter_updown_parallel",
+    kernel_class="limiter",
+    kernel_identifier="limiter_kernel",
+)
+limiter_updown_weighted_parallel = read_kernel_by_name(
+    "limiter_updown_weighted_parallel",
+    kernel_class="limiter",
+    kernel_identifier="limiter_kernel",
+)
+limiter_down_parallel = read_kernel_by_name(
+    "limiter_down_parallel", kernel_class="limiter", kernel_identifier="limiter_kernel"
+)
+limiter_soft_clipper = read_kernel_by_name(
+    "limiter_soft_clipper", kernel_class="limiter", kernel_identifier="limiter_kernel"
+)
+limiter_hard_clipper = read_kernel_by_name(
+    "limiter_hard_clipper", kernel_class="limiter", kernel_identifier="limiter_kernel"
+)
+limiter_dev = read_kernel_by_name(
+    "limiter_dev", kernel_class="limiter", kernel_identifier="limiter_kernel"
+)
+
+_limiter_kernel_map = {
+    "downward-upward": limiter_updown_weighted_parallel,
+    "downward": limiter_down_parallel,
+    "soft-clipper": limiter_soft_clipper,
+    "hard-clipper": limiter_hard_clipper,
+}
+
+
+def limiter_get_modes() -> List[str]:
+    return list(_limiter_kernel_map.keys())
+
+
+def limiter(
+    audio_in: torch.Tensor | np.ndarray,
+    mode: str = "downward",
+    sample_rate: int = 44100,
+    threshold: float = 0.5,  # Threshold in percents
+    slope: float = 1.0,  # Slope in percents
+    attack_ms: float = 0.008,  # Attack time in ms
+    release_ms: float = 100.0,  # Release time in ms
+) -> torch.Tensor:
+    """
+    Compresses stereo audio using an optimized CUDA kernel with running sum RMS calculation.
+
+    Parameters:
+        audio_in (torch.Tensor or np.ndarray): Input stereo audio signal with shape (n_samples, 2).
+        sample_rate (int): Sampling rate in Hz.
+        threshold (float): Threshold in percents (0-100).
+        slope (float): Slope angle in percents (0-100).
+        rms_window_ms (float): RMS window width in milliseconds.
+        attack_ms (float): Attack time in milliseconds.
+        release_ms (float): Release time in milliseconds.
+        chunk_size (int): Number of samples per chunk.
+        device (str): Device to place the output tensor ('cuda' or 'cpu').
+
+    Returns:
+        torch.Tensor: Compressed stereo audio with shape (n_samples, 2).
+        np.ndarray: Debug information (envelope and gain) with shape (n_samples, 2).
+    """
+    # Convert input to CPU double-precision NumPy array if necessary
+
+    device = audio_in.device
+
+    audio_in = audio_in.T
+    if isinstance(audio_in, torch.Tensor):
+        audio_in = audio_in.detach().cpu()
+        audio_in = audio_in.numpy()
+    else:
+        audio_in = np.asarray(audio_in, dtype=np.float64)
+
+    # Ensure the audio is in shape (n_samples, 2)
+    if audio_in.ndim != 2 or audio_in.shape[1] != 2:
+        raise ValueError(
+            f"Input audio must have shape (n_samples, 2), but got {audio_in.shape}"
+        )
+
+    n_samples = audio_in.shape[0]
+    n_channels = audio_in.shape[1]
+
+    wav_in_flat = audio_in.flatten()
+
+    wav_in_gpu = cp.asarray(wav_in_flat, dtype=cp.float64)
+    wav_out_gpu = cp.zeros_like(wav_in_gpu)
+
+    # Define grid and block dimensions
+    block_size = 64
+    grid_size = n_channels  # (n_channels + block_size - 1) // block_size
+    shared_mem_size = n_channels * cp.float64().nbytes  # Shared memory for envelopes
+
+    # attack_coeff = math.exp(-1.0 / (sample_rate * (attack_ms * 1e-3)))
+    # release_coeff = math.exp(-1.0 / (sample_rate * (release_ms * 1e-3)))
+
+    # print('attack_coeff',attack_coeff)
+    # print('attack_coeff',release_coeff)
+
+    if mode not in limiter_get_modes():
+        raise Exception(f"Limiter Kernel '{mode}' Not Found")
+
+    kernel = _limiter_kernel_map[mode]
+
+    kernel(
+        (grid_size,),
+        (block_size,),
+        (
+            wav_in_gpu,
+            wav_out_gpu,
+            np.int32(n_channels),
+            np.int32(n_samples),
+            np.float64(threshold),
+            np.float64(slope),
+            np.float64(sample_rate),
+            np.float64(attack_ms),
+            np.float64(release_ms),
+        ),
+        shared_mem=shared_mem_size,
+    )
+
+    # Retrieve results
+    wav_out_host = wav_out_gpu.get().astype(np.float64)
+    # Reshape the output
+    wav_out_stereo = wav_out_host.reshape((n_samples, n_channels))
+
+    if device == "cuda":
+        out_tensor = torch.from_numpy(wav_out_stereo).to(device)
+    else:
+        out_tensor = torch.from_numpy(wav_out_stereo).to("cpu")
+
+    return out_tensor.T

core/loudness.py

@@ -11,10 +11,10 @@ Description:
 
 import torch
 import pyloudnorm as pyln  # pip install pyloudnorm
-import torch
+import numpy as np
 
 
-def rms_normalization(audio: torch.Tensor, target_rms: float = 0.1):
+def rms_normalization(audio: torch.Tensor, target_rms: float = 0.1) -> torch.Tensor:
     rms = torch.sqrt(torch.mean(audio**2))
     scaling_factor = target_rms / rms
     normalized_audio = audio * scaling_factor
@@ -22,37 +22,120 @@ def rms_normalization(audio: torch.Tensor, target_rms: float = 0.1):
 
 
 def lufs_normalization(
-    audio: torch.Tensor, sample_rate: float, target_lufs: float = -14.0
-):
+    audio: torch.Tensor, sample_rate: int, target_lufs: float = -14.0
+) -> torch.Tensor:
     meter = pyln.Meter(sample_rate)  # Create a loudness meter
-    audio = audio.T.cpu().numpy()
-    loudness = meter.integrated_loudness(audio)  # Current LUFS
+
+    __audio = audio.T.cpu().numpy()
+    loudness = meter.integrated_loudness(__audio)  # Current LUFS
 
     loudness_offset = target_lufs - loudness
-    normalized_audio = audio * (10 ** (loudness_offset / 20.0))
+    normalized_audio = __audio * (10 ** (loudness_offset / 20.0))
 
     result = torch.from_numpy(normalized_audio)
     result = result.T
 
+    result = result.to(device=audio.device, dtype=result.dtype)
+
     return result
 
 
-def peak_normalization(audio, target_peak=0.9):
+def peak_normalization(audio: torch.Tensor, target_peak: float = 0.9) -> torch.Tensor:
     peak = torch.max(torch.abs(audio))
     scaling_factor = target_peak / peak
     normalized_audio = audio * scaling_factor
     return normalized_audio
 
 
-def get_loudness(audio: torch.Tensor, sample_rate: float) -> float:
+def get_loudness(audio: torch.Tensor, sample_rate: int) -> float:
     meter = pyln.Meter(sample_rate)  # Create a loudness meter
     audio = audio.T
     audio = audio.cpu()
-    loudness = meter.integrated_loudness(audio.numpy())  # Current LUFS
+    loudness: float = float(meter.integrated_loudness(audio.numpy()))  # Current LUFS
     return loudness
 
 
-def automatic_gain_control(audio, target_level=0.7, alpha=0.1):
+def set_loudness2(
+    audio_signal: torch.Tensor, sample_rate: int, target_loudness_db: float = -20.0
+) -> torch.Tensor:
+    """
+    Adjusts the loudness of the audio signal to a target level in decibels.
+
+    Args:
+        audio_signal (torch.Tensor): Input audio signal (channels, samples).
+        sample_rate (int): Sample rate of the audio signal.
+        target_loudness_db (float): Desired loudness in dB (e.g., -20.0 dB).
+
+    Returns:
+        torch.Tensor: Audio signal adjusted to the target loudness.
+    """
+    # Convert PyTorch tensor to NumPy array for loudness calculation
+    audio_np = audio_signal.cpu().numpy().T  # Convert to [samples, channels]
+
+    # Use pyloudnorm Meter to calculate and normalize loudness
+    meter = pyln.Meter(sample_rate)  # Create loudness meter
+    current_loudness = meter.integrated_loudness(audio_np)  # Measure LUFS
+
+    # Compute loudness adjustment gain
+    loudness_offset = target_loudness_db - current_loudness
+    gain_factor = 10 ** (loudness_offset / 20.0)
+
+    # Apply gain to adjust loudness
+    adjusted_audio_np = audio_np * gain_factor
+
+    # Convert back to PyTorch tensor
+    adjusted_audio = torch.from_numpy(adjusted_audio_np.T).to(
+        audio_signal.device, dtype=torch.float32
+    )
+
+    return adjusted_audio
+
+
+def set_loudness(
+    audio_signal: torch.Tensor, sample_rate: int, target_loudness_db: float = -20.0
+) -> torch.Tensor:
+    """
+    Adjusts the loudness of the audio signal to a target level in decibels,
+    ensuring no clipping occurs.
+
+    Args:
+        audio_signal (torch.Tensor): Input audio signal (channels, samples).
+        sample_rate (int): Sample rate of the audio signal.
+        target_loudness_db (float): Desired loudness in dB (e.g., -20.0 dB).
+
+    Returns:
+        torch.Tensor: Audio signal adjusted to the target loudness.
+    """
+    # Convert PyTorch tensor to NumPy array for loudness calculation
+    audio_np = audio_signal.cpu().numpy().T  # Convert to [samples, channels]
+
+    # Use pyloudnorm Meter to calculate and normalize loudness
+    meter = pyln.Meter(sample_rate)  # Create loudness meter
+    current_loudness = meter.integrated_loudness(audio_np)  # Measure LUFS
+
+    # Compute loudness adjustment gain
+    loudness_offset = target_loudness_db - current_loudness
+    gain_factor = 10 ** (loudness_offset / 20.0)
+
+    # Apply gain to adjust loudness
+    adjusted_audio_np = audio_np * gain_factor
+
+    # Prevent clipping by normalizing the peak
+    peak_amplitude = np.max(np.abs(adjusted_audio_np))
+    if peak_amplitude > 1.0:
+        adjusted_audio_np = adjusted_audio_np / peak_amplitude
+
+    # Convert back to PyTorch tensor
+    adjusted_audio = torch.from_numpy(adjusted_audio_np.T).to(
+        audio_signal.device, dtype=torch.float32
+    )
+
+    return adjusted_audio
+
+
+def automatic_gain_control(
+    audio: torch.Tensor, target_level: float = 0.7, alpha: float = 0.1
+) -> torch.Tensor:
     current_level = torch.mean(torch.abs(audio))
     gain = target_level / (current_level + 1e-6)
     smoothed_gain = alpha * gain + (1 - alpha) * 1.0

core/matching.py

@@ -1,34 +0,0 @@
-#!/usr/bin/env python3
-# -*- coding: utf-8 -*-
-"""
-Author: C0ffymachyne
-License: GPLv3
-Version: 1.0.0
-
-Description:
-    loudness matching
-"""
-
-import torch
-import torchaudio.transforms as T
-from torch.optim import Adam
-
-
-def approximate_loudness_matching(audio, target_loudness=-14.0):
-    rms = torch.sqrt(torch.mean(audio**2))
-    current_loudness = 20 * torch.log10(rms + 1e-6)
-    gain = 10 ** ((target_loudness - current_loudness) / 20.0)
-    return audio * gain
-
-
-def waveform_matching(audio, target_audio, lr=0.01, steps=100):
-    audio = audio.clone().requires_grad_(True)
-    optimizer = Adam([audio], lr=lr)
-
-    for _ in range(steps):
-        loss = torch.nn.functional.mse_loss(audio, target_audio)
-        optimizer.zero_grad()
-        loss.backward()
-        optimizer.step()
-
-    return audio.detach()

core/mixing.py

@@ -12,7 +12,12 @@ Description:
 import torch
 
 
-def combine_audio_files(waveform_a, waveform_b, sample_rate, chunk_duration=2.0):
+def combine_audio_files(
+    waveform_a: torch.Tensor,
+    waveform_b: torch.Tensor,
+    sample_rate: int,
+    chunk_duration: float = 2.0,
+) -> torch.Tensor:
     """
     Combine two audio files by alternating 2-second chunks, cropping to the shorter audio.
 

core/plotting.py

@@ -8,33 +8,40 @@ Version: 1.0.0
 Description:
     various plotting methods for debugging and visualization
 """
-
-import torch
-from io import BytesIO
-from PIL import Image
 import numpy as np
+import torch
+import torchaudio
+from PIL import Image
 import matplotlib.pyplot as plt
+from typing import Tuple, List
+from matplotlib.figure import Figure
 
 
-def _figure_to_image(figure, format="png", dpi=96):
-    buf = BytesIO()
-    figure.savefig(buf, format=format, dpi=96)
-    plt.close(figure)
+def _figure_to_image(figure: Figure, dpi: int = 96) -> Image.Image:
+    """Convert a Matplotlib figure to a high-resolution RGB PIL Image."""
+    figure.set_dpi(dpi)
+    figure.canvas.draw()
+    data = np.frombuffer(figure.canvas.tostring_argb(), dtype=np.uint8)
+    width, height = figure.canvas.get_width_height()
+    image = data.reshape((height, width, 4))  # ARGB format
 
-    buf.seek(0)
-    waveform_image = Image.open(buf).convert("RGB")
+    # Convert ARGB to RGB
+    rgb_image = np.zeros((height, width, 3), dtype=np.uint8)
+    rgb_image[:, :, 0] = image[:, :, 1]  # Red
+    rgb_image[:, :, 1] = image[:, :, 2]  # Green
+    rgb_image[:, :, 2] = image[:, :, 3]  # Blue
 
-    return waveform_image
+    return Image.fromarray(rgb_image)
 
 
 def get_wave(
     waveform: torch.Tensor,
     sample_rate: int,
     title: str = "Waveform",
-    xlim=None,
-    ylim=None,
-):
-    waveform = waveform.numpy()
+    xlim: int = 1000,
+    ylim: int = 1000,
+) -> Image:
+    waveform = waveform.cpu().numpy()
 
     num_channels, num_frames = waveform.shape
     time_axis = torch.arange(0, num_frames) / sample_rate
@@ -61,27 +68,237 @@ def get_wave(
 
 
 def get_spectogram(
-    waveform: torch.Tensor, sample_rate: int, title="Spectrogram", xlim=None
-):
-    waveform = waveform.numpy()
-    num_channels, num_frames = waveform.shape
-    time_axis = torch.arange(0, num_frames) / sample_rate
+    waveform: torch.Tensor,
+    sample_rate: int,
+    n_fft: int = 4096,
+    n_mels: int = 512,
+    title: str = "Spectrogram",
+    xlim: int = 8192,
+    dpi: int = 96,  # Set a high DPI for better image resolution
+) -> np.ndarray:
+    """Generate and plot a high-resolution spectrogram from a waveform."""
 
-    figure, axes = plt.subplots(num_channels, 1)
-    if num_channels == 1:
-        axes = [axes]
-    for c in range(num_channels):
-        axes[c].specgram(waveform[c], Fs=sample_rate)
-    if num_channels > 1:
-        axes[c].set_ylabel(f"Channel {c+1}")
-    if xlim:
-        axes[c].set_xlim(xlim)
-    figure.suptitle(title)
+    # Parameters for Mel Spectrogram
+    win_length = n_fft // 2
+    hop_length = n_fft // 4  # Smaller hop for better time resolution
 
-    waveform_image = _figure_to_image(figure)
+    spectrogram_transform = torchaudio.transforms.MelSpectrogram(
+        sample_rate=sample_rate,
+        n_fft=n_fft,
+        win_length=win_length,
+        hop_length=hop_length,
+        center=True,
+        pad_mode="reflect",
+        normalized=True,
+        power=2.0,  # Using power spectrogram
+        norm="slaney",
+        n_mels=n_mels,
+        mel_scale="slaney",
+    ).to(waveform.device, dtype=waveform.dtype)
 
+    # Compute spectrogram
+    mel_spectrogram = spectrogram_transform(waveform).cpu()
+
+    # Convert to decibel scale for better visualization
+    spectrogram = torchaudio.transforms.AmplitudeToDB(top_db=80)(
+        mel_spectrogram
+    ).numpy()
+
+    # Plot the spectrogram
+    num_channels, _ = waveform.shape
+    figure, axes = plt.subplots(
+        num_channels, 1, figsize=(20, 10 * num_channels), squeeze=False, dpi=dpi
+    )
+    figure.suptitle(title, fontsize=16)
+
+    for i, ax in enumerate(axes[:, 0]):  # Unpack axes
+        ax.imshow(
+            spectrogram[i],
+            origin="lower",
+            aspect="auto",
+            extent=[0, xlim, 0, sample_rate / 2],
+            cmap="magma",
+        )
+        ax.set_title(f"Channel {i + 1}", fontsize=14)
+        ax.set_xlabel("Time (frames)", fontsize=12)
+        ax.set_ylabel("Frequency (Hz)", fontsize=12)
+        ax.tick_params(axis="both", which="major", labelsize=10)
+
+    plt.tight_layout(rect=[0, 0.03, 1, 0.95])
+
+    # Convert the figure to a high-resolution image
+    waveform_image = _figure_to_image(figure, dpi=dpi)
+    plt.close(figure)  # Close the figure to free up memory
     return waveform_image
 
 
-def save_image(filepath: str, image: Image):
-    image.save(filepath, format="JPEG", quality=95)
+def save_image(filepath: str, image: Image) -> Image:
+    image.save(filepath, format="PNG", quality=95)
+
+
+def save_harmonic_spectrum(
+    waveform: torch.Tensor,
+    sample_rate: int,
+    output_image: str,
+    figsize: Tuple[int, int] = (12, 6),
+    num_harmonics: int = 10,
+) -> None:
+    """
+    Generate the harmonic spectrum of a waveform and save it as an image.
+
+    Parameters:
+        waveform (torch.Tensor): Audio waveform tensor.
+        sample_rate (int): Sample rate of the audio.
+        output_image (str): Path to save the output image.
+        figsize (tuple): Size of the output figure in inches.
+        num_harmonics (int): Number of harmonics to calculate.
+
+    Returns:
+        None
+    """
+    # Ensure mono audio (combine channels if necessary)
+    if waveform.size(0) > 1:
+        waveform = torch.mean(waveform, dim=0, keepdim=True)
+
+    # Perform FFT to get the frequency domain
+    fft = torch.fft.fft(waveform).to(device=waveform.device)
+    magnitude = torch.abs(fft[0])  # Magnitude of the FFT
+    frequencies = torch.fft.fftfreq(waveform.size(1), d=1 / sample_rate).to(
+        device=waveform.device
+    )
+
+    # Extract the fundamental frequency
+    fundamental_idx = torch.argmax(magnitude[: len(magnitude) // 2])
+    fundamental_freq = frequencies[fundamental_idx]
+
+    # Calculate harmonic frequencies
+    harmonic_frequencies = [fundamental_freq * (i + 1) for i in range(num_harmonics)]
+    harmonic_amplitudes = [
+        magnitude[int(harmonic / sample_rate * len(magnitude))]
+        for harmonic in harmonic_frequencies
+    ]
+    harmonic_frequencies = torch.tensor(harmonic_frequencies).cpu().numpy()
+    harmonic_amplitudes = torch.tensor(harmonic_amplitudes).cpu().numpy()
+
+    # Plot the harmonic spectrum
+    plt.figure(figsize=figsize)
+    plt.plot(harmonic_frequencies, harmonic_amplitudes, color="blue", linewidth=2)
+    plt.title("Harmonic Spectrum")
+    plt.xlabel("Frequency (Hz)")
+    plt.ylabel("Amplitude")
+    plt.grid(True)
+
+    # Save the image
+    plt.tight_layout()
+    plt.savefig(output_image, dpi=300)
+    plt.close()
+
+
+PREDEFINED_COLORS = [
+    "#1f77b4",
+    "#ff7f0e",
+    "#2ca02c",
+    "#d62728",
+    "#9467bd",
+    "#8c564b",
+    "#e377c2",
+    "#7f7f7f",
+    "#bcbd22",
+    "#17becf",
+]
+
+
+def plot_multiple_harmonic_spectra(
+    audio_data: List[Tuple[torch.Tensor, int, str]],
+    output_image: str,
+    figsize: Tuple[int, int] = (12, 6),
+    num_harmonics: int = 16,
+    upper_bound: int = 140000,
+    title: str = "Harmonic Spectrum",
+) -> None:
+    """
+    Optimized: Plot high-resolution harmonic spectra for multiple audio waveforms.
+
+    Parameters:
+        audio_data (list): List of (waveform, sample_rate, label) tuples.
+        output_image (str): Path to save the combined output image.
+        figsize (tuple): Size of the figure.
+        num_harmonics (int): Number of harmonics to calculate.
+        upper_bound (int): Maximum frequency to display.
+        title (str): Title of the plot.
+    """
+    plt.figure(figsize=figsize)
+
+    for idx, (waveform, sample_rate, label) in enumerate(audio_data):
+        # Ensure mono audio
+        if waveform.size(0) > 1:
+            waveform = torch.mean(waveform, dim=0)
+
+        # Move waveform to GPU for efficient computation
+        device = (
+            waveform.device
+            if waveform.is_cuda
+            else "cuda" if torch.cuda.is_available() else "cpu"
+        )
+        waveform = waveform.to(device)
+
+        # High-resolution FFT (Zero-padding for better frequency resolution)
+        n_fft = 1 * waveform.size(0)  # Zero-padding factor of 4
+        fft = torch.fft.fft(waveform, n=n_fft).to(device=waveform.device)
+        magnitude = torch.abs(fft[: n_fft // 2])
+        frequencies = torch.fft.fftfreq(n_fft, d=1 / sample_rate)[: n_fft // 2]
+
+        # Find fundamental frequency
+        fundamental_idx = torch.argmax(
+            magnitude[: len(magnitude) // 4]
+        )  # Search in the first quarter
+        fundamental_freq = frequencies[fundamental_idx]
+
+        # Precompute harmonic frequencies
+        harmonic_freqs = fundamental_freq * torch.arange(
+            1, num_harmonics + 1, device=device
+        )
+        harmonic_indices = (harmonic_freqs / (sample_rate / n_fft)).long()
+        harmonic_amplitudes = magnitude[harmonic_indices].cpu().numpy()
+
+        # Limit harmonic frequencies to upper bound
+        harmonic_freqs = harmonic_freqs[harmonic_freqs <= upper_bound].cpu().numpy()
+        harmonic_amplitudes = harmonic_amplitudes[: len(harmonic_freqs)]
+
+        # Convert data to CPU for plotting
+        frequencies_cpu = frequencies.cpu().numpy()
+        magnitude_cpu = magnitude.cpu().numpy()
+
+        # Plot the full spectrum
+        color = PREDEFINED_COLORS[idx % len(PREDEFINED_COLORS)]
+        plt.plot(
+            frequencies_cpu,
+            magnitude_cpu,
+            color=color,
+            alpha=0.4,
+            label=f"Full Spectrum ({label})",
+        )
+
+        # Overlay harmonic peaks
+        plt.vlines(
+            harmonic_freqs,
+            ymin=0,
+            ymax=harmonic_amplitudes,
+            color=color,
+            linewidth=1.5,
+            linestyle="--",
+            label=f"Harmonics ({label})",
+        )
+
+    plt.title(title)
+    plt.xscale("log")
+    plt.xlim(20, upper_bound)
+    plt.xlabel("Frequency (Hz)")
+    plt.ylabel("Amplitude")
+    plt.legend()
+    plt.grid(True, which="both", linestyle="--", linewidth=0.5)
+
+    # Save and close the figure
+    plt.tight_layout()
+    plt.savefig(output_image, dpi=150)
+    plt.close()

core/sampling.py

@@ -0,0 +1,33 @@
+#!/usr/bin/env python3
+# -*- coding: utf-8 -*-
+"""
+Author: C0ffymachyne
+License: GPLv3
+Version: 1.0.0
+
+Description:
+    utility functions for sampling
+"""
+
+import torch
+import torchaudio
+
+
+def oversample(
+    audio_signal: torch.Tensor, sample_rate: int, factor: int = 8
+) -> torch.Tensor:
+
+    resampler = torchaudio.transforms.Resample(sample_rate, sample_rate * factor).to(
+        device=audio_signal.device, dtype=audio_signal.dtype
+    )
+    return resampler(audio_signal), sample_rate * factor
+
+
+def downsample(
+    audio_signal: torch.Tensor, sample_rate: int, factor: int = 8
+) -> torch.Tensor:
+
+    resampler = torchaudio.transforms.Resample(sample_rate, sample_rate // factor).to(
+        device=audio_signal.device, dtype=audio_signal.dtype
+    )
+    return resampler(audio_signal)

core/saturation.py

@@ -9,96 +9,159 @@ Description:
     various conversion methods
 """
 
-import math
 import torch
-from torch import nn
+import torch.nn as nn
+from torch.nn import Module
+from ..core.sampling import oversample, downsample
+from ..core.harmonics import enahnce_harmonics_23
 
-class TubeValveSimulator:
-    def __init__(self, sample_rate, coefficients=None):
-        """
-        Initialize the Tube/Valve simulator using Volterra Series Expansion.
-
-        Args:
-            sample_rate (int): The sample rate of the audio signal.
-            coefficients (dict): Coefficients for Volterra series terms.
-                                Example: {'linear': 1.0, 'quadratic': 0.5, 'cubic': 0.2}
-        """
-        self.sample_rate = sample_rate
-        self.coefficients = coefficients if coefficients else {
-            'linear': 1.0,
-            'quadratic': 0.5,
-            'cubic': 0.2,
-        }
-
-    def process(self, audio_input):
-        """
-        Process the audio signal to simulate tube/valve characteristics.
-
-        Args:
-            audio_input (torch.Tensor): The input audio signal (1D or 2D tensor).
-
-        Returns:
-            torch.Tensor: The processed audio signal.
-        """
-        # Ensure input is a float tensor
-        #audio_input = audio_input.to(torch.float32)
-
-        # Volterra series expansion
-        linear_term = self.coefficients['linear'] * audio_input
-        quadratic_term = self.coefficients['quadratic'] * audio_input**2
-        cubic_term = self.coefficients['cubic'] * audio_input**3
-
-        # Combine the terms
-        processed_audio = linear_term + quadratic_term + cubic_term
-
-        # Optional soft clipping to simulate saturation
-        #processed_audio = self.soft_clip(processed_audio)
-
-        return processed_audio
-
-    @staticmethod
-    def soft_clip(audio_signal, threshold=0.9):
-        """
-        Apply soft clipping to the audio signal to limit extreme values.
-
-        Args:
-            audio_signal (torch.Tensor): The input audio signal.
-            threshold (float): The clipping threshold.
-
-        Returns:
-            torch.Tensor: The clipped audio signal.
-        """
-        return torch.tanh(audio_signal / threshold) * threshold
-
-########################################
-# Example Usage
-########################################
-if __name__ == "__main__":
-
-    from pathlib import Path
-    import torchaudio
-    from core.io import from_disk_as_raw_2d, audio_to_comfy_3d, audio_from_comfy_2d
-
-    samples_path = Path("ComfyUI_SignalProcessing/audio/inputs/baxandall-normalizer_00003_.flac")
-
-    source_path = samples_path.absolute()
-    source_audio, source_audio_sample_rate = from_disk_as_raw_2d(source_path,try_gpu=True)
-    input = audio_to_comfy_3d(source_audio, source_audio_sample_rate)[0]
+from typing import List
 
 
-    # Initialize the simulator
-    tube_simulator = TubeValveSimulator(source_audio_sample_rate, coefficients={
-        'linear': 1.0,
-        'quadratic': 0.3,
-        'cubic': 0.4,
-    })
+def sigmoid_saturation(audio: torch.Tensor, drive: float = 50.0) -> torch.Tensor:
+    """
+    Apply sigmoid saturation with drive control.
+    Parameters:
+        audio: [channels, samples] input audio signal in the range [-1, 1].
+        drive: Controls the steepness of the sigmoid in range [0, 100.0].
+    Returns:
+        Saturated audio signal.
+    """
+    k = 0.1 + (drive / 100.0) * 10  # Map drive to steepness control
+    normalized = (audio + 1) / 2
+    saturated = 1 / (1 + torch.exp(-k * (normalized - 0.5)))
+    return 2 * saturated - 1
 
-    compressed = tube_simulator.process(source_audio)
-    print("Compressed shape:", compressed.shape)
-    print("Device:", compressed.device)
 
-    torchaudio.save(
-        "audio/outputs/valve.wav",
-        compressed.cpu(),
-        source_audio_sample_rate,
+def tanh_saturation(audio: torch.Tensor, drive: float = 50.0) -> torch.Tensor:
+    """
+    Apply tanh saturation with drive control.
+    Parameters:
+        audio: Input audio signal (Tensor).
+        drive: Controls the strength of tanh effect in range [0, 100.0].
+    Returns:
+        Saturated audio signal.
+    """
+
+    audio = audio / torch.max(torch.abs(audio))
+    k = 0.1 + (drive / 100.0) * 10  # Map drive to scaling factor
+    k_tensor = torch.tensor(
+        k, dtype=audio.dtype, device=audio.device
+    )  # Convert k to Tensor
+    return torch.tanh(
+        k_tensor * audio
+    )  # / torch.tanh(k_tensor)  # Normalize output range
+
+
+def poly_saturation(audio: torch.Tensor, drive: float = 50.0) -> torch.Tensor:
+    """
+    Apply cubic polynomial saturation with drive control.
+    Parameters:
+        audio: Input audio signal.
+        drive: Controls the strength of the cubic term in range [0, 100.0].
+    Returns:
+        Saturated audio signal.
+    """
+    c3 = 0.01 + (drive / 100.0) * 0.3  # Map drive to nonlinearity strength
+    return audio - c3 * audio**3
+
+
+def logarithmic_mapping(audio: torch.Tensor, drive: float = 50.0) -> torch.Tensor:
+    """
+    Apply logarithmic mapping with drive control.
+    Parameters:
+        audio: Input audio signal.
+        drive: Controls the scaling of the logarithmic mapping.
+    Returns:
+        Saturated audio signal.
+    """
+    max_value = 0.1 + (drive / 100.0) * 10  # Map drive to maximum scaling
+    return (
+        torch.sign(audio)
+        * torch.log1p(torch.abs(audio * max_value))
+        / torch.log1p(torch.tensor(max_value))
     )
+
+
+class Saturator(Module):
+    @staticmethod
+    def get_modes() -> List[str]:
+        return ["poly", "soft", "tanh", "sig", "log"]
+
+    def __init__(
+        self,
+        drive: float = 0.5,
+        order: int = 3,
+        sample_rate: int = 48000,
+        mode: str = "poly",
+        oversample_factor: int = 4,
+    ):
+        super(Saturator, self).__init__()
+        self.order: int = order
+        self.sample_rate: int = sample_rate
+        self.mode: str = mode
+        self.drive: float = drive
+        self.oversample_factor: int = oversample_factor
+        self.harmonics_level: int = 0
+
+        order = order  # third order polynomial approximation
+
+        # Input scaling and output gain (adjust as needed)
+        self.input_scale = nn.Parameter(torch.tensor(1.0, dtype=torch.float32))
+        self.output_gain = nn.Parameter(torch.tensor(1.0, dtype=torch.float32))
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+
+        y, sample_rate = oversample(x, self.sample_rate, factor=self.oversample_factor)
+
+        y = enahnce_harmonics_23(y, sample_rate, gain_db_base=self.harmonics_level)
+
+        if self.mode == "poly":
+            y = poly_saturation(y, drive=self.drive)
+        elif self.mode == "tanh":
+            y = tanh_saturation(y, drive=self.drive)
+        elif self.mode == "sig":
+            y = sigmoid_saturation(y, drive=self.drive)
+        elif self.mode == "log":
+            y = logarithmic_mapping(y, drive=self.drive)
+
+        y = downsample(y, sample_rate, factor=self.oversample_factor)
+
+        return y
+
+
+def saturator_get_modes() -> List[str]:
+    return ["poly", "tanh", "sig", "log"]
+
+
+def saturator(
+    audio_in: torch.Tensor,
+    mode: str = "poly",
+    sample_rate: int = 44100,
+    drive: float = 1.5,  # Removed lookahead
+    oversample_factor: int = 4,
+    harmonics_level: float = 1.2,
+) -> torch.Tensor:
+    y = audio_in.clone()
+    # loudness = get_loudness(audio_in, sample_rate)
+    # y = automatic_gain_control(audio_in)
+    # y = y*drive_pre
+    y, _sample_rate = oversample(y, sample_rate, factor=oversample_factor)
+
+    # y = enahnce_harmonics_23(y, _sample_rate, gain_db_base=harmonics_level)
+
+    if mode == "poly":
+        drive = drive * 2
+        y = poly_saturation(y, drive=drive)
+    elif mode == "tanh":
+        drive = drive / 3
+        y = tanh_saturation(y, drive=drive)
+    elif mode == "sig":
+        y = sigmoid_saturation(y, drive=drive)
+    elif mode == "log":
+        y = logarithmic_mapping(y, drive=drive)
+
+    y = downsample(y, sample_rate, factor=oversample_factor)
+    # y = lufs_normalization(y, sample_rate, loudness)
+
+    return y

core/spectral.py

@@ -0,0 +1,394 @@
+#!/usr/bin/env python3
+# -*- coding: utf-8 -*-
+"""
+Author: C0ffymachyne
+License: GPLv3
+Version: 1.0.0
+
+Description:
+    utility functions operating in frequency domain
+"""
+
+import torch
+import torch.nn.functional as F
+import torchaudio.functional as AF
+from scipy.signal import butter
+from typing import Any
+
+
+def tilt(
+    audio: torch.Tensor,
+    sample_rate: int,
+    tilt_amount: float = -1.0,
+    pivot_freq: float = 300,
+    n_fft: int = 4096,
+) -> torch.Tensor:
+    """
+    Apply a spectral tilt to audio.
+
+    Parameters:
+        audio (Tensor): [channels, samples] audio signal.
+        sr (int): Sample rate.
+        tilt_amount (float): Positive for brighter, negative for darker (dB/octave).
+        pivot_freq (float): Pivot frequency in Hz.
+
+    Returns:
+        Tensor: Tilted audio signal.
+    """
+    hop_length = n_fft // 4
+
+    # Perform STFT
+    stft = torch.stft(audio, n_fft=n_fft, hop_length=hop_length, return_complex=True)
+    freqs = torch.fft.rfftfreq(n_fft, 1 / sample_rate).to(audio.device)
+
+    # Calculate tilt filter
+    slope = tilt_amount / (freqs.max() - pivot_freq)
+    tilt_filter = 1 + slope * (freqs - pivot_freq)
+    tilt_filter = tilt_filter.clamp(min=0.1)  # Avoid negative or zero gain
+
+    # Apply filter to magnitude
+    mag = torch.abs(stft)
+    phase = torch.angle(stft)
+    tilted_mag = mag * tilt_filter.unsqueeze(0).unsqueeze(-1)
+
+    # Reconstruct the tilted signal
+    tilted_stft = tilted_mag * torch.exp(1j * phase)
+    tilted_audio = torch.istft(
+        tilted_stft, n_fft=n_fft, hop_length=hop_length, length=audio.size(-1)
+    )
+
+    return tilted_audio
+
+
+def perceptual_weighted_central_frequency(
+    audio: torch.Tensor, sample_rate: int, n_fft: int = 4096, hop_length: int = 1024
+) -> torch.Tensor:
+    """
+    Calculate a perceptually weighted central frequency.
+
+    Parameters:
+        audio (Tensor): [channels, samples] audio signal.
+        sr (int): Sample rate.
+        n_fft (int): FFT size.
+        hop_length (int): Hop length for STFT. Defaults to n_fft // 4.
+
+    Returns:
+        float: Perceptually weighted central frequency in Hz.
+    """
+
+    window = torch.hann_window(n_fft, device=audio.device)  # Hann window for STFT
+    # Perform STFT
+    stft = torch.stft(
+        audio, n_fft=n_fft, hop_length=hop_length, return_complex=True, window=window
+    )
+    mag = torch.abs(stft)  # Magnitude spectrum
+    freqs = torch.fft.rfftfreq(n_fft, 1 / sample_rate).to(
+        audio.device
+    )  # Frequency bins
+
+    # Apply A-weighting (approximation for perceptual loudness)
+    a_weighting = (
+        1.2589
+        * ((12200**2) * (freqs**4))
+        / (
+            ((freqs**2) + (20.6**2))
+            * ((freqs**2) + (12200**2))
+            * torch.sqrt((freqs**2) + (107.7**2))
+        )
+    )
+
+    # Ensure dimensions match for broadcasting
+    a_weighting = a_weighting.unsqueeze(0).unsqueeze(-1)  # [1, freq_bins, 1]
+    weighted_mag = mag * a_weighting  # Apply A-weighting
+
+    # Calculate weighted central frequency
+    spectral_sum = weighted_mag.sum(dim=1)  # Sum across frequency axis
+    weighted_sum = (weighted_mag * freqs.unsqueeze(0).unsqueeze(-1)).sum(
+        dim=1
+    )  # Weighted sum of frequencies
+    central_frequency = (weighted_sum / (spectral_sum + 1e-8)).mean()
+
+    return central_frequency.item()
+
+
+def copy_frequency_amplitudes_with_balance(
+    source: torch.Tensor,
+    target: torch.Tensor,
+    sampler_rate: int,
+    n_fft: int = 4096,
+    hop_length: int = 1024,
+    alpha: float = 0.5,
+) -> torch.Tensor:
+    """
+    Copy frequency amplitudes from the source to the target song with
+    frequency balance adjustment for stereo signals.
+
+    Parameters:
+        source (Tensor): Source audio signal [channels, samples].
+        target (Tensor): Target audio signal [channels, samples].
+        sr (int): Sample rate.
+        n_fft (int): FFT window size.
+        hop_length (int): Hop length for STFT. Defaults to n_fft // 4.
+        alpha (float): Blending factor for magnitude adjustment (0 = target-only, 1 = source-only).
+
+    Returns:
+        Tensor: Modified target audio signal with balanced frequency power.
+    """
+
+    # Ensure the source and target have the same number of channels
+    if source.size(0) != target.size(0):
+        raise ValueError("Source and target must have the same number of channels.")
+
+    # Define frequency bands (in Hz) for low, mid, and high
+    bands = [(0, 200), (200, 2000), (2000, sampler_rate // 2)]
+
+    # Process each channel independently
+    modified_channels = []
+    for ch in range(source.size(0)):
+
+        source_window = torch.hann_window(
+            n_fft, device=source.device
+        )  # Hann window for STFT
+        target_window = torch.hann_window(
+            n_fft, device=target.device
+        )  # Hann window for STFT
+        # Perform STFT on both source and target for the current channel
+        source_stft = torch.stft(
+            source[ch],
+            n_fft=n_fft,
+            hop_length=hop_length,
+            return_complex=True,
+            window=source_window,
+        )
+        target_stft = torch.stft(
+            target[ch],
+            n_fft=n_fft,
+            hop_length=hop_length,
+            return_complex=True,
+            window=target_window,
+        )
+
+        # Extract magnitudes and phases
+        source_mag = torch.abs(source_stft)  # Source magnitudes
+        target_mag = torch.abs(target_stft)  # Target magnitudes
+        target_phase = torch.angle(target_stft)  # Target phases
+
+        # Frequency bins for the FFT
+        freqs = torch.fft.rfftfreq(n_fft, 1 / sampler_rate).to(source.device)
+
+        # Adjust frequency bands
+        for low, high in bands:
+            # Find indices for the frequency range
+            band_indices = (freqs >= low) & (freqs < high)
+
+            # Compute power in the band for source and target
+            source_power = source_mag[band_indices, :].pow(2).mean()
+            target_power = target_mag[band_indices, :].pow(2).mean()
+
+            # Compute scaling factor
+            if target_power > 0:
+                scaling_factor = (source_power / target_power).sqrt()
+            else:
+                scaling_factor = 1.0  # Avoid division by zero
+
+            # Apply scaling to target magnitudes in the band
+            target_mag[band_indices, :] *= scaling_factor
+
+        # Blend magnitudes: alpha * source_mag + (1 - alpha) * target_mag
+        blended_mag = alpha * source_mag + (1 - alpha) * target_mag
+
+        # Reconstruct STFT with blended magnitudes and original target phase
+        modified_stft = blended_mag * torch.exp(1j * target_phase)
+
+        # Perform ISTFT to reconstruct the modified target signal
+        modified_audio = torch.istft(
+            modified_stft,
+            n_fft=n_fft,
+            hop_length=hop_length,
+            length=target.size(-1),
+        )
+
+        # Store the modified channel
+        modified_channels.append(modified_audio)
+
+    # Stack the modified channels back into a single tensor
+    modified_audio = torch.stack(modified_channels, dim=0)
+
+    return modified_audio
+
+
+def get_emphasis_alpha(
+    audio: torch.Tensor, sample_rate: int, n_fft: int = 2048, hop_length: int = 512
+) -> torch.Tensor:
+    """
+    Analyze the audio signal and compute high and low-frequency energy ratios.
+
+    Parameters:
+        audio (Tensor): [channels, samples] input audio signal.
+        sr (int): Sample rate of the audio.
+
+    Returns:
+        float: Ratio of high-frequency energy to total energy.
+    """
+    # Compute the Short-Time Fourier Transform (STFT)
+    # Compute the Short-Time Fourier Transform (STFT)
+    stft = torch.stft(audio, n_fft=n_fft, hop_length=hop_length, return_complex=True)
+    magnitudes = torch.abs(
+        stft
+    )  # Get magnitude spectrum [channels, freq_bins, time_frames]
+
+    # Define frequency ranges
+    freqs = torch.fft.rfftfreq(n=2048, d=1 / sample_rate)  # Frequency bins
+
+    # Create masks for low and high frequencies
+    low_freqs = freqs < 500  # Low frequencies < 500 Hz
+    high_freqs = freqs > 2000  # High frequencies > 2 kHz
+
+    # Apply masks along the frequency dimension
+    low_energy = magnitudes[low_freqs, :].pow(2).sum()  # Energy in low frequencies
+    high_energy = magnitudes[high_freqs, :].pow(2).sum()  # Energy in high frequencies
+
+    # Avoid division by zero
+    if low_energy > 0:
+        ratio = high_energy / low_energy
+    else:
+        ratio = float("inf")  # Dominated by high frequencies
+
+    return ratio.item()
+
+
+def get_emphasis_aplha_clipped(
+    audio: torch.Tensor,
+    sample_rate: int,
+    min_alpha: float = 0.0,
+    max_alpha: float = 0.99,
+) -> Any:
+
+    ratio = get_emphasis_alpha(audio.mean(dim=0), sample_rate)
+
+    # Map the ratio to an alpha value
+    # Higher ratio -> lower alpha (more high-frequency content, less emphasis)
+    alpha = max_alpha - ratio / (1 + ratio) * (max_alpha - min_alpha)
+    alpha = torch.clamp(
+        torch.tensor(alpha), min_alpha, max_alpha
+    ).item()  # Ensure alpha is within range
+
+    return alpha
+
+
+def pre_emphasis(audio: torch.Tensor, alpha: float = 0.7) -> torch.Tensor:
+    """
+    Apply pre-emphasis to boost high frequencies before saturation.
+
+    Parameters:
+        audio (Tensor): [channels, samples] input audio signal.
+        alpha (float): Pre-emphasis factor (0 < alpha < 1).
+
+    Returns:
+        Tensor: Pre-emphasized audio signal.
+    """
+    x_shifted = F.pad(audio.unsqueeze(1), (1, 0))[:, :, :-1].squeeze(1)
+    emphasized = audio - alpha * x_shifted
+    return emphasized
+
+
+def de_emphasis(audio: torch.Tensor, alpha: float = 0.7) -> torch.Tensor:
+    """
+    Apply de-emphasis to restore original frequency balance.
+
+    Parameters:
+        emphasized (Tensor): [channels, samples] input audio signal.
+        alpha (float): De-emphasis factor (0 < alpha < 1).
+
+    Returns:
+        Tensor: De-emphasized audio signal.
+    """
+    # IIR filter coefficients
+    b = torch.tensor([1.0, 0.0], dtype=audio.dtype, device=audio.device)
+    a = torch.tensor([1.0, -alpha], dtype=audio.dtype, device=audio.device)
+
+    # Ensure [channels, samples] format
+    if audio.dim() == 1:
+        audio = audio.unsqueeze(0)
+
+    # Apply IIR filter for de-emphasis
+    x = AF.lfilter(audio, a_coeffs=a, b_coeffs=b, clamp=False)
+
+    # Return to original dimensions
+    if x.size(0) == 1:
+        x = x.squeeze(0)
+    return x
+
+
+def olive_coloring(audio: torch.Tensor, sample_rate: int) -> torch.Tensor:
+    """
+    Apply 'olive' sound coloring to an audio signal.
+
+    Parameters:
+        audio (Tensor): [channels, samples] audio signal.
+        sr (int): Sample rate.
+
+    Returns:
+        Tensor: Colored audio signal.
+    """
+    n_fft = 4096
+    hop_length = n_fft // 4
+    window = torch.hann_window(n_fft, device=audio.device)  # Hann window for STFT
+
+    # Apply spectral tilt
+    stft = torch.stft(
+        audio,
+        n_fft=n_fft,
+        hop_length=hop_length,
+        window=window,
+        return_complex=True,
+    )
+    mag = torch.abs(stft)  # Magnitude spectrum
+    phase = torch.angle(stft)  # Phase spectrum
+    freqs = torch.fft.rfftfreq(n_fft, 1 / sample_rate).to(
+        audio.device
+    )  # Frequency bins
+
+    # Create tilt filter
+    tilt_filter = (freqs + torch.max(freqs)) ** -0.7  # Gentle tilt down the highs
+    tilt_filter = tilt_filter.clamp(
+        0.9, 1.2
+    )  # Prevent excessive attenuation or boosting
+    tilt_filter = tilt_filter.unsqueeze(0).unsqueeze(
+        -1
+    )  # Broadcast across time dimension
+
+    # Blend original magnitude with tilted magnitude
+    tilted_mag = mag * tilt_filter.expand_as(mag)
+    blended_mag = 0.75 * mag + 0.25 * tilted_mag  # Retain some original characteristics
+
+    # Reconstruct audio
+    colored_stft = blended_mag * torch.exp(1j * phase)
+    audio_tilted = torch.istft(
+        colored_stft,
+        n_fft=n_fft,
+        hop_length=hop_length,
+        length=audio.size(-1),
+        window=window,
+    )
+
+    # Apply a gentle high-frequency attenuation (above 15 kHz)
+    nyquist = sample_rate / 2
+    cutoff = 15000 / nyquist
+
+    b, a = butter(1, cutoff, btype="low", output="ba")  # 6 dB/octave slope
+    b = torch.tensor(b, dtype=audio.dtype, device=audio.device)
+    a = torch.tensor(a, dtype=audio.dtype, device=audio.device)
+    audio_lowpassed = AF.lfilter(audio_tilted, a_coeffs=a, b_coeffs=b)
+
+    # Add dynamic saturation (soft clipping)
+    saturated_audio = torch.tanh(1.1 * audio_lowpassed)
+
+    # saturated_audio = torch.atan(0.9 * saturated_audio)
+
+    # Normalize output
+    max_val = torch.max(torch.abs(saturated_audio))
+    if max_val > 0:
+        saturated_audio = saturated_audio / max_val
+
+    return saturated_audio

core/tests.py

@@ -0,0 +1,9 @@
+from pathlib import Path
+
+# Reusable constants
+DATA_ROOT = Path("ComfyUI_SignalProcessing/audio")
+
+
+def get_output_file_path(output_root: Path, test_name: str, mode: str) -> Path:
+    """Generate output file path for a given test."""
+    return output_root / f"{test_name}-{mode}.wav"

core/utilities.py

@@ -19,15 +19,15 @@ pattern = r"^ComfyUI-\d+\.\d+\.\d+$"
 this_file_directory = os.path.join(os.path.dirname(os.path.realpath(__file__)))
 
 
-def find_comfy_root():
+def find_comfy_root() -> str:
     path = os.fspath(this_file_directory)
+
     if isinstance(this_file_directory, bytes):
         sep = b"/"
     else:
         sep = "/"
     tokens = path.split(sep)
     while not re.match(pattern, tokens[-1]):
-        print(tokens[-1])
         tokens.pop(-1)
 
     path = "/".join(tokens)
@@ -37,13 +37,13 @@ def find_comfy_root():
 # add comfy to path for local devepment only
 # find comfy root by going upwards hoping it match regex
 # export coffy_local_dev=1
-def comfy_root_to_syspath():
+def comfy_root_to_syspath() -> None:
     try:
         if os.environ["coffy_local_dev"] == "1":
-            print("coffy_local_dev")
+            pass
         else:
             raise (Exception())
-    except:
+    except Exception:
         return
 
     path = find_comfy_root()

core/utilitiescuda.py

@@ -0,0 +1,32 @@
+#!/usr/bin/env python3
+# -*- coding: utf-8 -*-
+"""
+Author: C0ffymachyne
+License: GPLv3
+Version: 1.0.0
+
+Description:
+    various conversion methods
+"""
+
+import os
+import cupy as cp
+
+
+def read_kernel_by_name(
+    kernel_name: str,
+    kernel_class: str = "limiter",
+    kernel_identifier: str = "limiter_kernel",
+) -> cp.RawKernel:
+    this_directory = os.path.dirname(os.path.abspath(__file__))
+    kernel_relativepath = f"kernels/{kernel_class}/{kernel_name}.cu"
+    kernel_filepath = os.path.join(this_directory, kernel_relativepath)
+    print(f"Loading CUDA kernel... {kernel_relativepath}")
+
+    if not os.path.exists(kernel_filepath):
+        raise FileNotFoundError(f"Kernel file not found: {kernel_filepath}")
+
+    with open(kernel_filepath, "r", encoding="utf-8") as file:  # Open as text
+        code = file.read()  # Read kernel source code as string
+        # Pass code to RawKernel
+        return cp.RawKernel(code=code, name=kernel_identifier, backend="nvrtc")

core/widening.py

@@ -7,13 +7,14 @@ Version: 1.0.0
 
 Description:
     Open Source Stereo Widening Plugin
-    The copyright of the source code is waived by the CCO-1.0 license putting it in public domain with some nuance, read the license please for more details.
+    The copyright of the source code is waived by the CCO-1.0 license putting it
+    in public domain with some nuance, read the license please for more details.
 
     This is directy copy of the part of the code corresponding paper by Orchisama Das
         https://www.dafx.de/paper-archive/2024/papers/DAFx24_paper_92.pdf
 
     License : Creative Commons Zero v1.0 Universal
-    Link To The License : 
+    Link To The License:
         https://github.com/orchidas/StereoWidener/blob/main/LICENSE
 
     Repository : https://github.com/orchidas/StereoWidener
@@ -27,11 +28,13 @@ import numpy.typing as npt
 
 from pathlib import Path
 from enum import Enum
-from typing import Optional, Dict, Union, Tuple
+from typing import Optional, Dict, Union, Tuple, List
 from abc import ABC
 from scipy.signal import sosfilt
 from scipy.signal import fftconvolve
 
+from numpy.typing import NDArray
+import matplotlib.pyplot as plt
 
 this_directory = os.path.dirname(os.path.realpath(__file__))
 data_directory = os.path.join(os.path.split(this_directory)[0], "data", "widener")
@@ -39,9 +42,6 @@ data_directory = os.path.join(os.path.split(this_directory)[0], "data", "widener
 VN_PATH = Path(f"{data_directory}/init_vn_filters.txt")
 OPT_VN_PATH = Path(f"{data_directory}/opt_vn_filters.txt")
 
-print(VN_PATH.absolute())
-print(os.path.dirname(os.path.realpath(__file__)))
-
 
 class FilterbankType(Enum):
     AMP_PRESERVE = "amplitude-preserve"
@@ -54,16 +54,12 @@ class DecorrelationType(Enum):
     OPT_VELVET = "opt_velvet"
 
 
-import numpy as np
-from numpy.typing import ArrayLike, NDArray
-from typing import List
-import matplotlib.pyplot as plt
-
 _HIGH_EPS = 1e9
 
 
 def ms_to_samps(ms: npt.NDArray[float], /, fs: float) -> npt.NDArray[int]:
-    """Calculate the nearest integer number of samples corresponding to the given time duration in milliseconds.
+    """Calculate the nearest integer number of samples
+    corresponding to the given time duration in milliseconds.
 
     Args:
         ms (NDArray): Duration, in milliseconds
@@ -106,7 +102,7 @@ def warp_pole_angle(rho: float, pole_freq: Union[float, np.ndarray]) -> npt.NDAr
 
 def decorrelate_allpass_filters(
     fs: float, nbiquads: int = 250, max_grp_del_ms: float = 30.0
-):
+) -> npt.NDArray:
     """
     Return cascaded allpass SOS sections with randomised phase to perform signal decorrelation
     Args:
@@ -147,7 +143,7 @@ def process_allpass(
     fs: float,
     num_biquads: int = 200,
     max_grp_del_ms: float = 30.0,
-) -> np.ndarray:
+) -> npt.NDArray:
     """
     For an input stereo signal, pass both channels through
     cascade of allpass filters, and return the output
@@ -177,7 +173,7 @@ class LeakyIntegrator:
     """Leaky integrator for signal envelope detection"""
 
     def __init__(
-        self, fs: float, attack_time_ms: float = 5.0, release_time_ms: float = 50.0
+        self, fs: int, attack_time_ms: float = 5.0, release_time_ms: float = 50.0
     ):
         self.fs = fs
         self.attack_time_ms = attack_time_ms
@@ -189,7 +185,8 @@ class LeakyIntegrator:
         Leaky integrator = a first-order IIR low-pass filter.
 
         Args:
-            input_signal (npt.NDArray): The impulse response (should be 1-dimensional array or only the 1st column is taken)
+            input_signal (npt.NDArray): The impulse response (should be 1-dimensional array
+            or only the 1st column is taken)
             fs (float): Sample rate
             attack_time_ms (float): Integrator attack time in milliseconds, by default 5
             release_time_ms (float): Integrator release time in milliseconds, by default 50
@@ -225,7 +222,7 @@ class OnsetDetector:
 
     def __init__(
         self,
-        fs: float,
+        fs: int,
         attack_time_ms: float = 5.0,
         release_time_ms: float = 20.0,
         min_onset_hold_ms: float = 80.0,
@@ -243,11 +240,11 @@ class OnsetDetector:
         """
         self.fs = fs
         self.leaky = LeakyIntegrator(fs, attack_time_ms, release_time_ms)
-        self._onset_flag = []
+        self._onset_flag: Optional[List[bool]] = []
         self.min_onset_hold_samps = int(ms_to_samps(min_onset_hold_ms, self.fs))
         self.min_onset_sep_samps = int(ms_to_samps(min_onset_sep_ms, self.fs))
-        self._threshold = None
-        self._signal_env = None
+        self._threshold: Optional[npt.NDArray] = None
+        self._signal_env: Optional[npt.NDArray] = None
 
     @property
     def signal_env(self) -> NDArray:
@@ -262,11 +259,11 @@ class OnsetDetector:
         return self._running_sum_thres
 
     @property
-    def onset_flag(self) -> List[bool]:
+    def onset_flag(self) -> Optional[List[bool]]:
         return self._onset_flag
 
     @staticmethod
-    def check_local_peak(cur_samp: float, prev_samp: float, next_samp: float):
+    def check_local_peak(cur_samp: float, prev_samp: float, next_samp: float) -> bool:
         """
         Given the current, previous and next samples, check if the current sample
         is a local peak
@@ -289,7 +286,7 @@ class OnsetDetector:
         else:
             return True if cur_samp < prev_samp and cur_samp > next_samp else False
 
-    def process(self, input_signal: NDArray, to_plot: bool = False):
+    def process(self, input_signal: NDArray, to_plot: bool = False) -> None:
         """Given an input signal, find the location of onsets"""
         if input_signal.ndim == 2 and input_signal.shape[1] > 1:
             input_signal = input_signal[:, 0]
@@ -351,10 +348,7 @@ class OnsetDetector:
                     self._onset_flag[k] = False
                     inhibit_counter += 1
 
-        if to_plot:
-            ax = self.plot(input_signal)
-
-    def plot(self, input_signal: NDArray):
+    def plot(self, input_signal: NDArray) -> plt.Axes:
         """Plot the input signal and the detected signal, threshold and onsets"""
         num_samp = len(input_signal)
         time_vector = np.arange(0, num_samp / self.fs, 1.0 / self.fs)
@@ -387,7 +381,7 @@ def process_velvet(
 
     try:
         vn_seq = np.loadtxt(vn_seq_path, dtype="f", delimiter=" ")
-    except:
+    except Exception:
         raise OSError("Error reading file!")
 
     output_signal = np.zeros_like(input_signal)
@@ -409,7 +403,7 @@ class StereoWidener(ABC):
         decorr_type: DecorrelationType,
         beta: float,
         detect_transient: bool = False,
-        onset_detection_params: Optional[Dict] = None,
+        onset_detection_params: Optional[Dict[str, float]] = None,
         xfade_win_len_ms: float = 1.0,
     ):
         """Args:
@@ -439,12 +433,12 @@ class StereoWidener(ABC):
         if self.detect_transient:
             if onset_detection_params is not None:
                 self.onset_detector = OnsetDetector(
-                    self.fs,
+                    int(self.fs),
                     min_onset_hold_ms=onset_detection_params["min_onset_hold_ms"],
                     min_onset_sep_ms=onset_detection_params["min_onset_sep_ms"],
                 )
             else:
-                self.onset_detector = OnsetDetector(self.fs)
+                self.onset_detector = OnsetDetector(int(self.fs))
             self.xfade_win_len_samps = int(ms_to_samps(xfade_win_len_ms, self.fs))
             self.xfade_in_win = half_hann_fade(self.xfade_win_len_samps, fade_out=False)
             self.xfade_out_win = half_hann_fade(self.xfade_win_len_samps, fade_out=True)
@@ -469,7 +463,7 @@ class StereoWidener(ABC):
         self.onset_detector.process(input)
         return self.onset_detector.onset_flag
 
-    def process(self):
+    def process(self) -> None:
         pass
 
 
@@ -477,7 +471,7 @@ def filter_in_subbands(
     input_signal: np.ndarray,
     fs: int,
     bands_per_octave: int = 3,
-    freq_range=(20, 16000),
+    freq_range: Tuple[int, int] = (20, 16000),
     filter_length: int = 4096,
 ) -> Tuple[npt.NDArray, npt.NDArray]:
 
@@ -507,12 +501,11 @@ def calculate_interchannel_cross_correlation_matrix(
     channel_axis: int = 0,
     return_single_coeff: bool = False,
     bands_per_octave: int = 3,
-    freq_range=(20, 16000),
-):
+    freq_range: Tuple[int, int] = (20, 16000),
+) -> npt.NDArray:
     """Returns a matrix of ICC values for each channel axis in signals"""
     if time_axis != -1:
         signals = np.moveaxis(signals, 0, 1)
-        channel_axis = 0
         time_axis = -1
 
     # passthrough filterbank
@@ -553,10 +546,10 @@ class StereoWidenerFrequencyBased(StereoWidener):
     def __init__(
         self,
         input_stereo: np.ndarray,
-        fs: float,
+        fs: int,
         filterbank_type: FilterbankType,
         decorr_type: DecorrelationType,
-        beta: Tuple[float, float],
+        beta: List[float],
         cutoff_freq: float,
     ):
         """Frequency based stereo widener
@@ -574,7 +567,7 @@ class StereoWidenerFrequencyBased(StereoWidener):
         self.cutoff_freq = cutoff_freq
         self.get_filter_coefficients()
 
-    def get_filter_coefficients(self):
+    def get_filter_coefficients(self) -> None:
         if self.filterbank_type == FilterbankType.AMP_PRESERVE:
             # Linkwitz Riley crossover filterbank
             filters = pf.dsp.filter.crossover(
@@ -617,16 +610,18 @@ class StereoWidenerFrequencyBased(StereoWidener):
         highpass_signal = self.highpass_filter_coeffs.process(pf_signal).time
         return np.vstack((lowpass_signal, highpass_signal))
 
-    def update_beta(self, new_beta: Tuple[float, float]):
+    def update_beta(self, new_beta: float) -> None:
         self.beta = new_beta
 
-    def update_cutoff_frequency(self, new_cutoff_freq: float):
+    def update_cutoff_frequency(self, new_cutoff_freq: float) -> None:
         self.cutoff_freq = new_cutoff_freq
         self.get_filter_coefficients()
 
-    def process(self):
-        stereo_output = np.zeros_like(self.input_signal)
-        filtered_input = np.zeros((2, self.input_signal.shape[0], self.num_channels))
+    def process(self) -> npt.NDArray:
+        stereo_output: npt.NDArray = np.zeros_like(self.input_signal)
+        filtered_input: npt.NDArray = np.zeros(
+            (2, self.input_signal.shape[0], self.num_channels)
+        )
         filtered_decorr = np.zeros_like(filtered_input)
 
         for chan in range(self.num_channels):
@@ -644,7 +639,9 @@ class StereoWidenerFrequencyBased(StereoWidener):
 
         return stereo_output
 
-    def calculate_interchannel_coherence(self, output_signal: np.ndarray):
+    def calculate_interchannel_coherence(
+        self, output_signal: np.ndarray
+    ) -> Tuple[npt.NDArray, npt.NDArray]:
         icc_matrix, icc_freqs = calculate_interchannel_cross_correlation_matrix(
             output_signal,
             fs=self.fs,

effects/SignalProcessingConvolutionReverb.py

@@ -10,20 +10,19 @@ Description:
 """
 
 import os
-import torch, torchaudio
+import torch
+import torchaudio
 import torch.nn.functional as F
-from pathlib import Path
+
+from typing import Dict, Tuple, Union
 
 from ..core.utilities import comfy_root_to_syspath
-
-comfy_root_to_syspath()  # add comfy to sys path for dev
-
 from ..core.io import audio_from_comfy_2d, audio_to_comfy_3d, from_disk_as_raw_2d
 from ..core.loudness import lufs_normalization, get_loudness
-
-
 import folder_paths
 
+comfy_root_to_syspath()  # add comfy to sys path for dev
+
 
 class SignalProcessingConvolutionReverb:
     supported_formats = [".wav", ".mp3", ".ogg", ".m4a", ".flac", ".mp4"]
@@ -31,7 +30,7 @@ class SignalProcessingConvolutionReverb:
     ir_directory = os.path.join(os.path.split(this_directory)[0], "audio", "ir")
 
     @classmethod
-    def INPUT_TYPES(s):
+    def INPUT_TYPES(s) -> Dict[str, torch.Tensor]:
 
         files, _ = folder_paths.recursive_search(
             SignalProcessingConvolutionReverb.ir_directory
@@ -43,7 +42,7 @@ class SignalProcessingConvolutionReverb:
                 _, ext = os.path.splitext(file)
                 if ext in SignalProcessingConvolutionReverb.supported_formats:
                     ir_files.append(file)
-            except:
+            except Exception:
                 pass
 
         return {
@@ -67,14 +66,16 @@ class SignalProcessingConvolutionReverb:
     CATEGORY = "Signal Processing"
     FUNCTION = "process"
 
-    def process(self, impulse_response, audio_input, wet_dry):
+    def process(
+        self,
+        impulse_response: str,
+        audio_input: Dict[str, Union[torch.Tensor, int]],
+        wet_dry: float,
+    ) -> Tuple[Dict[str, Union[torch.Tensor, int]]]:
 
         try_gpu: bool = True
         repeat: bool = True
 
-        print(impulse_response)
-        print(audio_input)
-
         waveform, sample_rate = audio_from_comfy_2d(
             audio_input, repeat=repeat, try_gpu=try_gpu
         )
@@ -84,25 +85,23 @@ class SignalProcessingConvolutionReverb:
         it_filepath = os.path.join(
             SignalProcessingConvolutionReverb.ir_directory, impulse_response
         )
+
         ir, ir_sr = from_disk_as_raw_2d(it_filepath, repeat=repeat, try_gpu=try_gpu)
 
         # Resample IR if sampling rates do not match
         if ir_sr != sample_rate:
-            print(f"Resampling impulse response from {ir_sr} Hz to {sample_rate} Hz.")
             resampler = torchaudio.transforms.Resample(
                 orig_freq=ir_sr, new_freq=sample_rate
             ).to(ir.device, dtype=waveform.dtype)
             ir = resampler(ir)
             ir_sr = sample_rate
 
-        print("ensure channels")
         # if wave is mono and ir is not mono
         if waveform.shape[0] == 1 and ir.shape[0] == 2:
             ir = ir.mean(dim=0, keepdim=True)
         if waveform.shape[0] == 2 and ir.shape[0] == 1:
             ir = ir.repeat(2, 1)
 
-        print("apply reverb")
         processed_audio = self.apply_reverb(waveform, sample_rate, ir, wet_dry=wet_dry)
         processed_audio = lufs_normalization(processed_audio, sample_rate, loudness)
 
@@ -162,35 +161,3 @@ class SignalProcessingConvolutionReverb:
         processed_audio = torch.stack(processed_channels)  # Shape: [2, N]
 
         return processed_audio
-
-
-if __name__ == "__main__":
-
-    import torchaudio
-    from pathlib import Path
-    from ..core.io import audio_from_comfy_2d, audio_to_comfy_3d, from_disk_as_raw_2d
-
-    node = SignalProcessingConvolutionReverb()
-    types = node.INPUT_TYPES()
-
-    samples_path = Path("ComfyUI_SignalProcessing/audio/samples")
-    impulses_path = Path("ComfyUI_SignalProcessing/audio/ir")
-
-    samples = list(samples_path.rglob("*.*"))
-    impulses = list(impulses_path.rglob("*.*"))
-
-    print(samples)
-
-    s = samples[1].absolute()
-    i = impulses[0].absolute()
-    s, sr = from_disk_as_raw_2d(s)
-
-    sound = audio_to_comfy_3d(s, sr)[0]
-
-    result = node.process(i, sound, wet_dry=0.5)[0]
-    audio, sr = audio_from_comfy_2d(result)
-    torchaudio.save("ComfyUI_SignalProcessing/audio/tests/reverb.wav", audio.cpu(), sr)
-
-    # set console to comfy ComfyUI-0.2.4/custom_nodes and run below command
-    # export coffy_local_dev=1
-    # python3 -m ComfyUI_SingalProcessing.effects.SignalProcessingConvolutionReverb

effects/SignalProcessingPaulStretch.py

@@ -16,22 +16,21 @@ Description:
 import torch
 
 import math
-from typing import Tuple, List, Dict
+from typing import Tuple, Dict, Any, Union
 
 from ..core.utilities import comfy_root_to_syspath
-
-comfy_root_to_syspath()  # add comfy to sys path for dev
-
 from ..core.io import audio_from_comfy_2d, audio_to_comfy_3d
 from ..core.loudness import lufs_normalization, get_loudness
 
+comfy_root_to_syspath()  # add comfy to sys path for dev
+
 
 class SignalProcessingPaulStretch:
     @classmethod
-    def INPUT_TYPES(cls) -> Dict:
+    def INPUT_TYPES(cls) -> Dict[str, Any]:
         return {
             "required": {
-                "audio": ("AUDIO", {"forceInput": True}),
+                "audio_input": ("AUDIO", {"forceInput": True}),
                 "stretch_factor": (
                     "FLOAT",
                     {"default": 8.0, "min": 1.0, "max": 100.0, "step": 0.1},
@@ -50,17 +49,21 @@ class SignalProcessingPaulStretch:
 
     def process(
         self,
-        audio: Dict[str, torch.Tensor],
+        audio_input: Dict[str, Union[torch.Tensor, int]],
         stretch_factor: float,
         window_size_seconds: float,
-    ) -> Tuple[Dict[str, torch.Tensor]]:
+    ) -> Tuple[Dict[str, Union[torch.Tensor, int]]]:
 
         # Conditional processing: If stretch_factor is 1.0, return original audio
         if stretch_factor == 1.0:
-            return audio_to_comfy_3d(audio["waveform"], audio["sample_rate"])
+            return audio_to_comfy_3d(
+                audio_input["waveform"], audio_input["sample_rate"]
+            )
 
         # Extract waveform and sample_rate
-        waveform, sample_rate = audio_from_comfy_2d(audio, repeat=True, try_gpu=True)
+        waveform, sample_rate = audio_from_comfy_2d(
+            audio_input, repeat=True, try_gpu=True
+        )
         loudness = get_loudness(waveform, sample_rate)
 
         nchannels, nsamples = waveform.shape
@@ -168,58 +171,13 @@ class SignalProcessingPaulStretch:
         while True:
             n = orig_n
             while (n % 2) == 0:
-                n /= 2
+                n //= 2
             while (n % 3) == 0:
-                n /= 3
+                n //= 3
             while (n % 5) == 0:
-                n /= 5
+                n //= 5
 
             if n < 2:
                 break
             orig_n += 1
         return orig_n
-
-
-if __name__ == "__main__":
-
-    import torchaudio
-    from pathlib import Path
-    from ..core.io import audio_from_comfy_2d, audio_to_comfy_3d, from_disk_as_raw_2d
-    from ..core.mixing import combine_audio_files
-
-    node = SignalProcessingPaulStretch()
-    types = node.INPUT_TYPES()
-
-    samples_path = Path("ComfyUI_SignalProcessing/audio/samples")
-
-    samples = list(samples_path.rglob("*.*"))
-
-    print(samples)
-
-    source_path = samples[1].absolute()
-    source_audio, source_audio_sample_rate = from_disk_as_raw_2d(source_path)
-
-    input = audio_to_comfy_3d(source_audio, source_audio_sample_rate)[0]
-
-    window_size_seconds = 0.25
-    stretch_factor = 2.0
-
-    result = node.process(
-        input, stretch_factor=stretch_factor, window_size_seconds=window_size_seconds
-    )[0]
-
-    output_audio, sample_rate_audio = audio_from_comfy_2d(result)
-
-    combined = combine_audio_files(
-        source_audio.cpu(), output_audio.cpu(), sample_rate_audio, chunk_duration=4.0
-    )
-
-    torchaudio.save(
-        "ComfyUI_SignalProcessing/audio/tests/paulstretch.wav",
-        combined.cpu(),
-        sample_rate_audio,
-    )
-
-    # set console to comfy ComfyUI-0.2.4/custom_nodes and run below command
-    # export coffy_local_dev=1
-    # python3 -m ComfyUI_SignalProcessing.effects.SignalProcessingPitchShifter

effects/SignalProcessingPitchShifter.py

@@ -10,24 +10,23 @@ Description:
 """
 
 import torch
-from typing import Tuple, Dict
+from typing import Tuple, Dict, Any, Union
 
 import torchaudio.functional as F
 
 from ..core.utilities import comfy_root_to_syspath
-
-comfy_root_to_syspath()  # add comfy to sys path for dev
-
 from ..core.io import audio_from_comfy_3d, audio_to_comfy_3d
 from ..core.loudness import lufs_normalization, get_loudness
 
+comfy_root_to_syspath()  # add comfy to sys path for dev
+
 
 class SignalProcessingPitchShifter:
     @classmethod
-    def INPUT_TYPES(cls):
+    def INPUT_TYPES(cls) -> Dict[str, Any]:
         return {
             "required": {
-                "input_audio": ("AUDIO",),  # Input audio
+                "audio_input": ("AUDIO",),  # Input audio
                 "pitch_shift_factor": (
                     "INT",
                     {"default": 2, "min": -12 * 4, "max": 12 * 4, "step": 1},
@@ -38,17 +37,17 @@ class SignalProcessingPitchShifter:
 
     RETURN_TYPES = ("AUDIO",)
     RETURN_NAMES = ("output_audio",)
-    CATEGORY = "Audio Processing"
+    CATEGORY = "Signal Processing"
     FUNCTION = "process"
 
     def process(
         self,
-        input_audio: Dict[str, torch.Tensor],
+        audio_input: Dict[str, Union[torch.Tensor, int]],
         pitch_shift_factor: int = 2,
-    ) -> Tuple[Dict[str, torch.Tensor]]:
+    ) -> Tuple[Dict[str, Union[torch.Tensor, int]]]:
 
         try_gpu: bool = True
-        waveform, sample_rate = audio_from_comfy_3d(input_audio, try_gpu=try_gpu)
+        waveform, sample_rate = audio_from_comfy_3d(audio_input, try_gpu=try_gpu)
 
         loudness = get_loudness(waveform, sample_rate)
 
@@ -60,38 +59,3 @@ class SignalProcessingPitchShifter:
         )
 
         return audio_to_comfy_3d(pitch_shifted_waveform, sample_rate)
-
-
-if __name__ == "__main__":
-
-    import torchaudio
-    from pathlib import Path
-    from ..core.io import audio_from_comfy_2d, audio_to_comfy_3d, from_disk_as_raw_2d
-    from ..core.mixing import combine_audio_files
-
-    node = SignalProcessingPitchShifter()
-    types = node.INPUT_TYPES()
-
-    samples_path = Path("ComfyUI_SignalProcessing/audio/samples")
-
-    samples = list(samples_path.rglob("*.*"))
-
-    source_path = samples[1].absolute()
-    source_audio, source_audio_sample_rate = from_disk_as_raw_2d(source_path)
-
-    input = audio_to_comfy_3d(source_audio, source_audio_sample_rate)[0]
-
-    pitch_shift_factor = -4
-
-    result = node.process(input, pitch_shift_factor=pitch_shift_factor)[0]
-
-    output_audio, sample_rate_audio = audio_from_comfy_2d(result)
-
-    combined = combine_audio_files(
-        source_audio.cpu(), output_audio.cpu(), sample_rate_audio, chunk_duration=4.0
-    )
-    torchaudio.save(
-        "ComfyUI_SignalProcessing/audio/tests/pitchshifter.wav",
-        combined.cpu(),
-        sample_rate_audio,
-    )

effects/SignalProcessingStereoWidening.py

@@ -14,9 +14,7 @@ import sys
 import math
 import torch
 
-sys.path.insert(0, os.path.join(os.path.dirname(os.path.realpath(__file__)), "comfy"))
-
-
+from typing import Dict, Any, Tuple, Union
 from ..core.io import audio_from_comfy_2d, audio_to_comfy_3d
 from ..core.loudness import lufs_normalization, get_loudness
 from ..core.widening import (
@@ -25,8 +23,10 @@ from ..core.widening import (
     FilterbankType,
 )
 
+sys.path.insert(0, os.path.join(os.path.dirname(os.path.realpath(__file__)), "comfy"))
 
-def interpolate(t, a, b):
+
+def interpolate(t: float, a: float, b: float) -> float:
     if not 0.0 <= t <= 1.0:
         raise ValueError("t must be in the range [0.0, 1.0]")
     return a + t * (b - a)
@@ -34,7 +34,7 @@ def interpolate(t, a, b):
 
 class SignalProcessingStereoWidening:
     @classmethod
-    def INPUT_TYPES(cls):
+    def INPUT_TYPES(cls) -> Dict[str, Any]:
         return {
             "required": {
                 "mode": (["decorrelation", "simple"],),
@@ -50,12 +50,18 @@ class SignalProcessingStereoWidening:
 
     RETURN_TYPES = ("AUDIO",)
     RETURN_NAMES = ("widened_audio",)
-    CATEGORY = "Audio Processing"
+    CATEGORY = "Signal Processing"
     FUNCTION = "process"
 
-    def process(self, mode, audio_input, width=1.2):
+    def process(
+        self,
+        mode: str,
+        audio_input: Dict[str, Union[torch.Tensor, int]],
+        width: float = 1.2,
+    ) -> Tuple[Dict[str, Union[torch.Tensor, int]]]:
         """
-        Widen stereo audio or convert mono audio to wide stereo using the provided widening algorithm.
+        Widen stereo audio or convert mono audio to wide stereo
+        using the provided widening algorithm.
 
         Parameters:
             audio_input (Dict): Dictionary containing 'waveform' and 'sample_rate'.
@@ -76,7 +82,8 @@ class SignalProcessingStereoWidening:
 
             if channels not in [1, 2]:
                 raise ValueError(
-                    f"Unsupported number of channels: {channels}. Only mono and stereo are supported."
+                    f"Unsupported number of channels: {channels}. \
+                        Only mono and stereo are supported."
                 )
 
             # Calculate coefficients based on the provided width parameter
@@ -146,11 +153,11 @@ class SignalProcessingStereoWidening:
                 sample_rate,
                 filterbank_type,
                 decorellation_type,
-                (beta, beta),
+                [beta, beta],
                 cutoff_frequency_hz,
             )
-
-            widened_waveform = torch.from_numpy(stereoWidener.process())
+            widener_result = stereoWidener.process()
+            widened_waveform = torch.from_numpy(widener_result)
             widened_waveform = widened_waveform.T
 
             widened_waveform = lufs_normalization(
@@ -158,3 +165,5 @@ class SignalProcessingStereoWidening:
             )
 
             return audio_to_comfy_3d(widened_waveform, sample_rate)
+
+        return audio_to_comfy_3d(waveform, sample_rate)

generators/SignalProcessingPadSynth.py

@@ -6,20 +6,21 @@ License: GPLv3
 Version: 1.0.0
 
 Description:
-    Pad Synthesiser port of code from this article https://zynaddsubfx.sourceforge.io/doc/PADsynth/PADsynth.htm#c_implementation
+    Pad Synthesiser port of code from this article :
+    https://zynaddsubfx.sourceforge.io/doc/PADsynth/PADsynth.htm#c_implementation
 """
 
 
 import torch
-import math, json
-from typing import Tuple, List, Dict
+import math
+from typing import Tuple, Dict, Any, Union
 
 from ..core.io import audio_to_comfy_3d
 
 
 class SignalProcessingPadSynth:
     @classmethod
-    def INPUT_TYPES(cls) -> Dict:
+    def INPUT_TYPES(cls) -> Dict[str, Any]:
         return {
             "required": {
                 "sample_rate": (
@@ -52,7 +53,7 @@ class SignalProcessingPadSynth:
         fundamental_freq: float,
         bandwidth_cents: float,
         number_harmonics: int,
-    ) -> Tuple[Dict[str, torch.Tensor]]:
+    ) -> Tuple[Dict[str, Union[torch.Tensor, int]]]:
         """
         Apply PADsynth algorithm to generate audio.
 

generators/SignalProcessingPadSynthChoir.py

@@ -6,17 +6,18 @@ License: GPLv3
 Version: 1.0.0
 
 Description:
-    Pad Synth Choir Synthesizer port of code from this article https://zynaddsubfx.sourceforge.io/doc/PADsynth/PADsynth.htm#c_implementation
+    Pad Synth Choir Synthesizer port of code from this article
+    https://zynaddsubfx.sourceforge.io/doc/PADsynth/PADsynth.htm#c_implementation
 """
 
 import torch
 import math
-from typing import Tuple, List, Dict
+from typing import Tuple, List, Dict, Any, Union
 
 
 class SignalProcessingPadSynthChoir:
     @classmethod
-    def INPUT_TYPES(cls) -> Dict:
+    def INPUT_TYPES(cls) -> Dict[str, Any]:
         return {
             "required": {
                 "samplerate": (
@@ -53,7 +54,7 @@ class SignalProcessingPadSynthChoir:
         num_notes: int,
         bandwidth_cents: float,
         number_harmonics: int,
-    ) -> Tuple[List[Dict[str, torch.Tensor]]]:
+    ) -> Tuple[List[Dict[str, Union[torch.Tensor, int]]], int]:
         """
         Apply PADsynth choir algorithm to generate multiple audio files.
 
@@ -66,7 +67,8 @@ class SignalProcessingPadSynthChoir:
             number_harmonics (int): Number of harmonics to generate.
 
         Returns:
-            Tuple[List[Dict[str, torch.Tensor]]]: List of generated audios with waveform and sample rate.
+            Tuple[List[Dict[str, torch.Tensor]]]:
+                List of generated audios with waveform and sample rate.
         """
 
         # Define FFT size

nodes.py

@@ -29,7 +29,13 @@ from .processors.SignalProcessingBaxandallEQ import (
     SignalProcessingBaxandallEQ,
     SignalProcessingBaxandall3BandEQ,
 )
-from .processors.SignalProcessingHarmonicsEnhancer import SignalProcessingHarmonicsEnhancer
+from .processors.SignalProcessingHarmonicsEnhancer import (
+    SignalProcessingHarmonicsEnhancer,
+)
+from .processors.SignalProcessingSaturation import SignalProcessingSaturation
+from .processors.SignalProcessingLimiter import SignalProcessingLimiter
+
+# from .processors.SignalProcessingCompressor import SignalProcessingCompressor
 
 # visuals
 from .visuals.SignalProcessingSpectrogram import SignalProcessingSpectrogram
@@ -51,7 +57,10 @@ NODE_CLASS_MAPPINGS = {
     "SignalProcessingLoudness": SignalProcessingLoudness,
     "SignalProcessingBaxandallEQ": SignalProcessingBaxandallEQ,
     "SignalProcessingBaxandall3BandEQ": SignalProcessingBaxandall3BandEQ,
-    "SignalProcessingHarmonicsEnhancer" : SignalProcessingHarmonicsEnhancer
+    "SignalProcessingHarmonicsEnhancer": SignalProcessingHarmonicsEnhancer,
+    "SignalProcessingSaturation": SignalProcessingSaturation,
+    "SignalProcessingLimiter": SignalProcessingLimiter,
+    # "SignalProcessingCompressor": SignalProcessingCompressor,
 }
 
 NODE_DISPLAY_NAME_MAPPINGS = {
@@ -70,5 +79,8 @@ NODE_DISPLAY_NAME_MAPPINGS = {
     "SignalProcessingLoudness": "(SP) Loudness",
     "SignalProcessingBaxandallEQ": "(SP) Baxandall EQ",
     "SignalProcessingBaxandall3BandEQ": "(SP) Baxandall 3 Band EQ",
-    "SignalProcessingHarmonicsEnhancer" : "(SP) Enhance Harmonics"
+    "SignalProcessingHarmonicsEnhancer": "(SP) Enhance Harmonics",
+    "SignalProcessingSaturation": "(SP) Saturation",
+    "SignalProcessingLimiter": "(SP) Limiter",
+    # "SignalProcessingCompressor": "(SP) Compressor",
 }

noxfile.py

@@ -0,0 +1,63 @@
+import nox
+from nox import Session
+
+PYTHON_VERSIONS = ["3.10.12"]
+REUSE_ENV = True
+
+
+@nox.session(python=PYTHON_VERSIONS, tags=["style"], reuse_venv=REUSE_ENV)
+def lint(session: Session) -> None:
+
+    session.install("black")
+    session.install("flake8")
+    session.install("mypy")
+
+    session.run("black", ".")
+    session.run("flake8", "--max-line-length=100", "--ignore=E501,E203,W503", ".")
+    session.run(
+        "mypy",
+        ".",
+        "--ignore-missing-imports",
+        "--strict",
+        "--show-error-codes",
+    )
+
+
+@nox.session(python=PYTHON_VERSIONS, tags=["tests"], reuse_venv=REUSE_ENV)
+def tests(session: Session) -> None:
+    """Run pytest tests with Scalene profiling."""
+
+    session.install("scalene", "pytest", "pytest-cov", "pytest-xdist")
+    requirements = nox.project.load_toml("pyproject.toml")["project"]["dependencies"]
+    for _, v in requirements.items():
+        session.install(*v)
+
+    project_name = nox.project.load_toml("pyproject.toml")["project"]["name"]
+
+    pytest_path = session.run("which", "pytest", external=True, silent=True).strip()
+
+    if not session.posargs:
+        with session.cd(".."):
+            session.run(
+                "scalene",
+                # "--profile-all",
+                pytest_path,
+                f"--cov={project_name}",
+                "--cov-report=term",
+                "--cov-report=html",
+                f"{project_name}/tests/",
+                external=True,
+            )
+    else:
+        with session.cd(".."):
+            session.run(
+                "pytest",
+                "--rootdir=.",
+                f"--cov={project_name}",
+                "--cov-report=term",
+                "--cov-report=html",
+                f"{project_name}/tests/",
+                "-k",
+                session.posargs[0],
+                external=True,
+            )

processors/SignalProcessingBaxandallEQ.py

@@ -13,17 +13,15 @@ Description:
 import torch
 import torchaudio
 import math
+from typing import Dict, Any, Tuple, Union
 
 from ..core.io import audio_to_comfy_3d, audio_from_comfy_3d
 from ..core.loudness import lufs_normalization, get_loudness
 
-import torch
-import math
-
 
 class SignalProcessingBaxandallEQ:
     @classmethod
-    def INPUT_TYPES(cls):
+    def INPUT_TYPES(cls) -> Dict[str, Any]:
         return {
             "required": {
                 "audio_input": ("AUDIO",),
@@ -40,10 +38,15 @@ class SignalProcessingBaxandallEQ:
 
     RETURN_TYPES = ("AUDIO",)
     RETURN_NAMES = ("processed_audio",)
-    CATEGORY = "Audio Processing"
+    CATEGORY = "Signal Processing"
     FUNCTION = "process"
 
-    def process(self, audio_input, bass_gain_db=0.0, treble_gain_db=0.0):
+    def process(
+        self,
+        audio_input: Dict[str, Union[torch.Tensor, int]],
+        bass_gain_db: float = 0.0,
+        treble_gain_db: float = 0.0,
+    ) -> Tuple[Dict[str, torch.Tensor]]:
         waveform, sample_rate = audio_from_comfy_3d(audio_input, try_gpu=True)
         loudness = get_loudness(waveform, sample_rate)
 
@@ -72,7 +75,9 @@ class SignalProcessingBaxandallEQ:
         waveform = lufs_normalization(waveform, sample_rate, loudness)
         return audio_to_comfy_3d(waveform, sample_rate)
 
-    def design_rbj_shelf(self, sr, freq, gain_db, shelf_type="low"):
+    def design_rbj_shelf(
+        self, sr: int, freq: float, gain_db: float, shelf_type: str = "low"
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
         # RBJ Audio EQ Cookbook shelf filters
         A = 10.0 ** (gain_db / 40.0)
         w0 = 2 * math.pi * freq / sr
@@ -103,7 +108,7 @@ class SignalProcessingBaxandallEQ:
 
 class SignalProcessingBaxandall3BandEQ:
     @classmethod
-    def INPUT_TYPES(cls):
+    def INPUT_TYPES(cls) -> Dict[str, Any]:
         return {
             "required": {
                 "audio_input": ("AUDIO",),
@@ -140,12 +145,12 @@ class SignalProcessingBaxandall3BandEQ:
 
     RETURN_TYPES = ("AUDIO",)
     RETURN_NAMES = ("processed_audio",)
-    CATEGORY = "Audio Processing"
+    CATEGORY = "Signal Processing"
     FUNCTION = "process"
 
     def process(
         self,
-        audio_input,
+        audio_input: Dict[str, Union[torch.Tensor, int]],
         bass_gain_db: float = 0.0,
         mid_gain_db: float = 0.0,
         treble_gain_db: float = 0.0,
@@ -153,7 +158,7 @@ class SignalProcessingBaxandall3BandEQ:
         mid_freq: float = 1000.0,
         high_freq: float = 10000.0,
         mid_q: float = 0.7,
-    ):
+    ) -> Tuple[Dict[str, Union[torch.Tensor, int]]]:
 
         waveform, sample_rate = audio_from_comfy_3d(audio_input, try_gpu=True)
         device = waveform.device
@@ -187,7 +192,9 @@ class SignalProcessingBaxandall3BandEQ:
 
         return audio_to_comfy_3d(waveform, sample_rate)
 
-    def design_rbj_shelf(self, sr, freq, gain_db, shelf_type="low"):
+    def design_rbj_shelf(
+        self, sr: int, freq: float, gain_db: float, shelf_type: str = "low"
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
         # RBJ audio EQ cookbook formula for shelving filters
         A = 10.0 ** (gain_db / 40.0)
         w0 = 2.0 * math.pi * freq / sr
@@ -216,7 +223,9 @@ class SignalProcessingBaxandall3BandEQ:
         a = torch.tensor([1.0, a1 / a0, a2 / a0], dtype=torch.float64)
         return b, a
 
-    def design_rbj_peak(self, sr, freq, gain_db, Q=0.7):
+    def design_rbj_peak(
+        self, sr: int, freq: float, gain_db: float, Q: float = 0.7
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
         # RBJ audio EQ cookbook peak filter
         A = 10.0 ** (gain_db / 40.0)
         w0 = 2.0 * math.pi * freq / sr
@@ -233,55 +242,3 @@ class SignalProcessingBaxandall3BandEQ:
         b = torch.tensor([b0 / a0, b1 / a0, b2 / a0], dtype=torch.float64)
         a = torch.tensor([1.0, a1 / a0, a2 / a0], dtype=torch.float64)
         return b, a
-
-
-if __name__ == "__main__":
-    import torchaudio
-    from pathlib import Path
-    from ..core.io import from_disk_as_raw_2d, audio_to_comfy_3d, audio_from_comfy_2d
-
-    node = SignalProcessingBaxandallEQ()
-    samples_path = Path("ComfyUI_SignalProcessing/audio/inputs/song.mp4")
-
-    source_path = samples_path.absolute()
-    source_audio, source_audio_sample_rate = from_disk_as_raw_2d(source_path)
-    input = audio_to_comfy_3d(source_audio, source_audio_sample_rate)[0]
-
-    # Test with some gain settings
-    result = node.process(input, bass_gain_db=5.0, treble_gain_db=-3.0)[0]
-
-    output_audio, sample_rate_audio = audio_from_comfy_2d(result)
-
-    # Save output for analysis
-    torchaudio.save(
-        "ComfyUI_SignalProcessing/audio/tests/baxandall_eq.wav",
-        output_audio.cpu(),
-        sample_rate_audio,
-    )
-
-    node = SignalProcessingBaxandall3BandEQ()
-
-    # Example usage:
-    # Provide a test audio file path
-    test_audio_path = Path("ComfyUI_SignalProcessing/audio/samples/test_audio.wav")
-    source_audio, source_audio_sample_rate = from_disk_as_raw_2d(test_audio_path)
-    input_tensors = audio_to_comfy_3d(source_audio, source_audio_sample_rate)[0]
-
-    # Apply EQ with some settings:
-    result = node.process(
-        input_tensors,
-        bass_gain_db=4.0,
-        mid_gain_db=-2.0,
-        treble_gain_db=3.0,
-        low_freq=100,
-        mid_freq=1000,
-        high_freq=10000,
-        mid_q=0.7,
-    )[0]
-
-    output_audio, out_sr = audio_from_comfy_2d(result)
-    torchaudio.save(
-        "ComfyUI_SignalProcessing/audio/tests/3band_baxandall_eq_output.wav",
-        output_audio.cpu(),
-        out_sr,
-    )

processors/SignalProcessingCompressor.py

@@ -0,0 +1,91 @@
+#!/usr/bin/env python3
+# -*- coding: utf-8 -*-
+"""
+Author: C0ffymachyne
+License: GPLv3
+Version: 1.0.0
+
+Description:
+  This file contains a description of a compressor node that utilizes a CUDA-optimized kernel.
+"""
+import torch
+from typing import Dict, Any, Tuple, Union
+
+from ..core.io import audio_to_comfy_3d, audio_from_comfy_2d
+from ..core.compression import compressor
+from ..core.loudness import get_loudness, lufs_normalization
+
+
+class SignalProcessingCompressor:
+    @classmethod
+    def INPUT_TYPES(cls) -> Dict[str, Any]:
+        return {
+            "required": {
+                "audio_input": ("AUDIO",),
+                "comp": (
+                    "FLOAT",
+                    {"default": -0.3, "min": -2.0, "max": 2.0, "step": 0.01},
+                ),
+                "attack": (
+                    "FLOAT",
+                    {"default": 0.1, "min": 0.01, "max": 100.0, "step": 0.01},
+                ),
+                "release": (
+                    "FLOAT",
+                    {"default": 60.0, "min": 0.01, "max": 1000.0, "step": 0.1},
+                ),
+                "filter_param": (
+                    "FLOAT",
+                    {"default": 0.3, "min": 0.0, "max": 1.0, "step": 0.01},
+                ),
+            }
+        }
+
+    RETURN_TYPES = ("AUDIO",)
+    RETURN_NAMES = ("audio",)
+    CATEGORY = "Signal Processing"
+    FUNCTION = "process"
+
+    def process(
+        self,
+        audio_input: Dict[str, Union[torch.Tensor, int]],
+        comp: float = -0.3,  # Compression/expansion factor
+        attack: float = 0.1,  # Attack time in ms
+        release: float = 60.0,  # Release time in ms
+        filter_param: float = 0.3,  # Filter parameter < 1
+    ) -> Tuple[Dict[str, Union[torch.Tensor, int]]]:
+        """
+        Apply compression or expansion to the audio input using CUDA.
+
+        Parameters:
+            audio_input (Dict[str, Union[torch.Tensor, int]]): Input audio waveform and sample rate.
+            comp (float): Compression/expansion factor.
+            attack (float): Attack time in milliseconds.
+            release (float): Release time in milliseconds.
+            filter_param (float): Filter parameter for envelope smoothing.
+
+        Returns:
+            Tuple[Dict[str, Union[torch.Tensor, int]]]: Compressed audio and sample rate.
+        """
+        # Extract waveform and sample rate
+        waveform, sample_rate = audio_from_comfy_2d(audio_input, try_gpu=True)
+
+        loudness = get_loudness(waveform, sample_rate=sample_rate)
+
+        # Apply the compressor kernel
+        filtered_waveform, _ = compressor(
+            waveform,
+            sample_rate,
+            comp=comp,
+            attack=attack,
+            release=release,
+            a=filter_param,
+            device="cuda" if torch.cuda.is_available() else "cpu",
+        )
+
+        filtered_waveform = lufs_normalization(
+            filtered_waveform, sample_rate=sample_rate, target_lufs=loudness
+        )
+
+        # Return the processed audio
+        return audio_to_comfy_3d(filtered_waveform, sample_rate)

processors/SignalProcessingFilter.py

@@ -6,12 +6,12 @@ License: GPLv3
 Version: 1.0.0
 
 Description:
-    Classic Audio filter set 
+     The code defines a classic audio filter set for performing various audio processing tasks such as filtering out unwanted frequencies
 """
 
-import torch, torchaudio
 import torch
-from typing import Dict
+import torchaudio
+from typing import Dict, Any, Tuple, Union
 
 from ..core.io import audio_to_comfy_3d, audio_from_comfy_3d
 from ..core.loudness import lufs_normalization, get_loudness
@@ -19,10 +19,10 @@ from ..core.loudness import lufs_normalization, get_loudness
 
 class SignalProcessingFilter:
     @classmethod
-    def INPUT_TYPES(cls):
+    def INPUT_TYPES(cls) -> Dict[str, Any]:
         return {
             "required": {
-                "audio": ("AUDIO", {"forceInput": True}),
+                "audio_input": ("AUDIO", {"forceInput": True}),
                 "cutoff": (
                     "FLOAT",
                     {"default": 0.5, "min": 0.0, "max": 1.0, "step": 0.01},
@@ -41,15 +41,15 @@ class SignalProcessingFilter:
     RETURN_TYPES = ("AUDIO", "INT")
     RETURN_NAMES = ("audio", "sample_rate")
     CATEGORY = "Signal Processing"
-    FUNCTION = "apply_filter"
+    FUNCTION = "process"
 
-    def apply_filter(
+    def process(
         self,
-        audio: Dict[str, torch.Tensor],
+        audio_input: Dict[str, Union[torch.Tensor, int]],
         cutoff: float,
         filter_type: str,
         q_factor: float,
-    ):
+    ) -> Tuple[Dict[str, Union[torch.Tensor, int]]]:
         """
         Apply a specified filter to the input audio.
 
@@ -63,7 +63,7 @@ class SignalProcessingFilter:
             Tuple[Dict[str, torch.Tensor]]: Filtered audio.
         """
 
-        waveform, sample_rate = audio_from_comfy_3d(audio)
+        waveform, sample_rate = audio_from_comfy_3d(audio_input)
 
         loudness = get_loudness(waveform, sample_rate)
 
@@ -94,13 +94,6 @@ class SignalProcessingFilter:
         elif filter_type in ["bandpass", "bandstop"]:
             center_freq = cutoff_freq
             # Ensure that the bandwidth does not exceed the Nyquist frequency
-            bandwidth = center_freq / q_factor
-            lower_freq = max(
-                center_freq - bandwidth / 2.0, 20.0
-            )  # Prevent dropping below 20 Hz
-            upper_freq = min(
-                center_freq + bandwidth / 2.0, nyquist - 100.0
-            )  # Prevent exceeding Nyquist
 
             if filter_type == "bandpass":
                 filtered_waveform = torchaudio.functional.bandpass_biquad(

processors/SignalProcessingHarmonicsEnhancer.py

@@ -6,31 +6,31 @@ License: GPLv3
 Version: 1.0.0
 
 Description:
-    Classic Audio filter set 
+    Classic Audio filter set
 """
 from ast import literal_eval
-import torch, torchaudio
 import torch
-from typing import Dict
+import torchaudio
+from typing import Dict, Any, List, Tuple, Union
 
 from ..core.utilities import comfy_root_to_syspath
+from ..core.io import audio_to_comfy_3d, audio_from_comfy_2d
+from ..core.loudness import lufs_normalization, get_loudness
 
 comfy_root_to_syspath()  # add comfy to sys path for dev
-from ..core.io import from_disk_as_raw_2d, audio_to_comfy_3d, audio_from_comfy_2d
-from ..core.loudness import lufs_normalization, get_loudness
 
 
 class SignalProcessingHarmonicsEnhancer:
     @classmethod
-    def INPUT_TYPES(cls):
+    def INPUT_TYPES(cls) -> Dict[str, Any]:
         return {
             "required": {
                 "audio_input": ("AUDIO",),
                 "harmonics": ("STRING", {"default": "1,3,5,7,9"}),
                 "mode": (["detect base frequency", "use base frequency"],),
-                "base_frequency" : ("FLOAT", {"default" : 440, "min" : 0, "max" : 20000 }),
-                "gain_db" : ("INT", {"default": 5,"min" : 0, "max" : 50, "step" : 1}),
-                "Q" : ("FLOAT", {"default" : 0.707, "min" : 0, "max" : 1.0, "step" : 0.01 }),
+                "base_frequency": ("FLOAT", {"default": 440, "min": 0, "max": 20000}),
+                "gain_db": ("INT", {"default": 5, "min": 0, "max": 500, "step": 1}),
+                "Q": ("FLOAT", {"default": 0.707, "min": 0, "max": 1.0, "step": 0.01}),
             }
         }
 
@@ -41,86 +41,78 @@ class SignalProcessingHarmonicsEnhancer:
 
     def process(
         self,
-        audio_input: Dict[str, torch.Tensor],
-        harmonics : str = '1,3,5,7,9',
-        mode : str = "detect base frequency",
-        base_frequency : int = 440,
-        gain_db : int = 5,
+        audio_input: Dict[str, Union[torch.Tensor, int]],
+        harmonics: str = "1,3,5,7,9",
+        mode: str = "detect base frequency",
+        base_frequency: int = 440,
+        gain_db: int = 5,
         Q: float = 0.707,
-    ):
-        waveform, sample_rate = audio_from_comfy_2d(audio_input,try_gpu=True)
+    ) -> Tuple[Dict[str, Union[torch.Tensor, int]]]:
+        waveform, sample_rate = audio_from_comfy_2d(audio_input, try_gpu=True)
         loudness = get_loudness(waveform, sample_rate)
+
         try:
-            harmonics = [literal_eval(x) for x in harmonics.split(',')]
-        except:
-            raise RuntimeWarning("Invalid Harmonics Format. Please delimit integers by a comma ',' like this 1,,3,5,7,9 ")
+            harmonics_list: List[int] = [literal_eval(x) for x in harmonics.split(",")]
+        except Exception:
+            raise RuntimeWarning(
+                "Invalid Harmonics Format. Please delimit integers by a comma \
+                    ',' like this 1,,3,5,7,9 "
+            )
         if mode == "detect base frequency":
-            filtered_waveform = self.enhance_harmonics(waveform,sample_rate,harmonics=harmonics,gain_db=gain_db,Q=Q)
+            filtered_waveform = self.enhance_harmonics(
+                waveform, sample_rate, harmonics=harmonics_list, gain_db=gain_db, Q=Q
+            )
         elif mode == "use base frequency":
-            filtered_waveform = self.enhance_harmonics(waveform,sample_rate,harmonics=harmonics,gain_db=gain_db,base_frequency=base_frequency,Q=Q)
+            filtered_waveform = self.enhance_harmonics(
+                waveform,
+                sample_rate,
+                harmonics=harmonics_list,
+                gain_db=gain_db,
+                base_frequency=base_frequency,
+                Q=Q,
+            )
 
         filtered_waveform = lufs_normalization(filtered_waveform, sample_rate, loudness)
         return audio_to_comfy_3d(filtered_waveform, sample_rate)
 
-    def add_harmonics(self, audio : torch.Tensor, gain=1.2):
+    def add_harmonics(self, audio: torch.Tensor, gain: float = 1.2) -> torch.Tensor:
         # Apply saturation using a tanh curve
         harmonic_audio = torch.tanh(audio * gain)
         return harmonic_audio
 
-    def detect_fundamental(self,audio, sample_rate):
+    def detect_fundamental(self, audio: torch.Tensor, sample_rate: int) -> torch.Tensor:
         # Estimate the fundamental frequency using a pitch detection method
         pitch = torchaudio.functional.detect_pitch_frequency(audio, sample_rate)
 
         return pitch
 
-    def detect_fundamental_mean(self,audio, sample_rate):
+    def detect_fundamental_mean(self, audio: torch.Tensor, sample_rate: int) -> int:
         # Estimate the fundamental frequency using a pitch detection method
         pitch = torchaudio.functional.detect_pitch_frequency(audio, sample_rate)
 
-        return pitch.mean().item()
+        return int(pitch.mean().item())
 
-    def enhance_harmonics(self,audio, sample_rate, harmonics=[1,3,5,7,9,11], gain_db=5, base_frequency = None,Q = 0.707):
+    def enhance_harmonics(
+        self,
+        audio: torch.Tensor,
+        sample_rate: int,
+        harmonics: List[int] = [1, 3, 5, 7, 9, 11],
+        gain_db: float = 5,
+        base_frequency: float = 0,
+        Q: float = 0.707,
+    ) -> torch.Tensor:
         # Detect the base frequency
-        if base_frequency == None:
+        if base_frequency == 0:
             base_frequency = self.detect_fundamental_mean(audio, sample_rate)
             if base_frequency <= 0:  # Fallback if pitch detection fails
                 base_frequency = 440  # Use a default base frequency
 
         # Apply EQ boosts to specific harmonic frequencies
-
-        print('base_frequency',base_frequency)
         for harmonic in harmonics:
             freq = base_frequency * harmonic
             if freq < sample_rate / 2:  # Ensure it's within the Nyquist frequency
-                print('center_freq',freq)
-                audio = torchaudio.functional.equalizer_biquad(audio, sample_rate, center_freq=freq, gain=gain_db,Q=Q)
-        
+                audio = torchaudio.functional.equalizer_biquad(
+                    audio, sample_rate, center_freq=freq, gain=gain_db, Q=Q
+                )
+
         return audio
-
-if __name__ == "__main__":
-    import torchaudio
-    from pathlib import Path
-    from ..core.io import from_disk_as_raw_2d, audio_to_comfy_3d, audio_from_comfy_2d
-
-    node = SignalProcessingHarmonicsEnhancer()
-    samples_path = Path("ComfyUI_SignalProcessing/audio/inputs/baxandall-normalizer_00003_.flac")
-
-    source_path = samples_path.absolute()
-    source_audio, source_audio_sample_rate = from_disk_as_raw_2d(source_path)
-    input = audio_to_comfy_3d(source_audio, source_audio_sample_rate)[0]
-
-    # Test with some gain settings
-    result = node.process(input)[0]
-
-    output_audio, sample_rate_audio = audio_from_comfy_2d(result)
-
-    # Save output for analysis
-    torchaudio.save(
-        "ComfyUI_SignalProcessing/audio/outputs/harmonics.wav",
-        output_audio.cpu(),
-        sample_rate_audio,
-    )
-
-    # set console to comfy ComfyUI-0.2.4/custom_nodes and run below command
-    # export coffy_local_dev=1
-    # python3 -m ComfyUI_SignalProcessing.processors.SignalProcessingHarmonicsEnhancer

processors/SignalProcessingLimiter.py

@@ -0,0 +1,75 @@
+#!/usr/bin/env python3
+# -*- coding: utf-8 -*-
+"""
+Author: C0ffymachyne
+License: GPLv3
+Version: 1.0.0
+
+Description:
+    Classic Limiter CUDA Optimized
+
+Reference:
+    https://ccrma.stanford.edu/~jatin/ComplexNonlinearities/Hysteresis.html
+    https://viennatalk.mdw.ac.at/papers/Pap_01_79_Tronchin.pdf
+    https://jatinchowdhury18.medium.com/complex-nonlinearities-episode-3-hysteresis-fdeb2cd3e3f6
+    https://ccrma.stanford.edu/~dtyeh/papers/yeh07_dafx_clipode.pdf
+"""
+import torch
+from typing import Dict, Any, Tuple, Union
+
+from ..core.io import audio_to_comfy_3d, audio_from_comfy_2d
+from ..core.limiting import limiter, limiter_get_modes
+from ..core.loudness import get_loudness, lufs_normalization
+
+
+class SignalProcessingLimiter:
+    @classmethod
+    def INPUT_TYPES(cls) -> Dict[str, Any]:
+        return {
+            "required": {
+                "audio_input": ("AUDIO",),
+                "mode": (limiter_get_modes(),),
+                "threshold": (
+                    "FLOAT",
+                    {"default": 100.0, "min": 0.0, "max": 100.0, "step": 0.1},
+                ),
+                "slope": (
+                    "FLOAT",
+                    {"default": 100.0, "min": 0.0, "max": 100.0, "step": 0.1},
+                ),
+                "release_ms": (
+                    "FLOAT",
+                    {"default": 100.0, "min": 0.0, "max": 1000.0, "step": 0.1},
+                ),
+            }
+        }
+
+    RETURN_TYPES = ("AUDIO",)
+    RETURN_NAMES = ("audio",)
+    CATEGORY = "Signal Processing"
+    FUNCTION = "process"
+
+    def process(
+        self,
+        audio_input: Dict[str, Union[torch.Tensor, int]],
+        mode: str = "downward",
+        threshold: float = 50.0,  # Threshold in percents
+        slope: float = 100.0,  # Slope in percents
+        release_ms: float = 100.0,  # Release time in ms
+    ) -> Tuple[Dict[str, Union[torch.Tensor, int]]]:
+        waveform, sample_rate = audio_from_comfy_2d(audio_input, try_gpu=True)
+
+        loudness = get_loudness(waveform, sample_rate)
+
+        filtered_waveform = limiter(
+            waveform,
+            mode=mode,
+            sample_rate=sample_rate,
+            threshold=threshold / 100.0,
+            slope=slope / 100,
+            release_ms=release_ms,
+        )
+
+        filtered_waveform = lufs_normalization(filtered_waveform, sample_rate, loudness)
+
+        return audio_to_comfy_3d(filtered_waveform, sample_rate)

processors/SignalProcessingLoadAudio.py

@@ -9,21 +9,23 @@ Description:
     Audio loading node
 """
 
-import sys, os
-
-sys.path.insert(0, os.path.join(os.path.dirname(os.path.realpath(__file__)), "comfy"))
+import sys
+import os
+import torch
+from typing import Dict, Tuple, Any, Union
 
 from ..core.io import from_disk_as_dict_3d
-
 import folder_paths
 
+sys.path.insert(0, os.path.join(os.path.dirname(os.path.realpath(__file__)), "comfy"))
+
 
 class SignalProcessingLoadAudio:
     supported_formats = ["wav", "mp3", "ogg", "m4a", "flac", "mp4"]
     input_dir = os.path.join(folder_paths.get_input_directory(), "samples")
 
     @classmethod
-    def INPUT_TYPES(s):
+    def INPUT_TYPES(s) -> Dict[str, Any]:
         supported_extensions = tuple(
             f".{fmt.lower()}" for fmt in SignalProcessingLoadAudio.supported_formats
         )
@@ -49,5 +51,7 @@ class SignalProcessingLoadAudio:
     CATEGORY = "Signal Processing"
     FUNCTION = "process"
 
-    def process(self, audio_file, gain):
+    def process(
+        self, audio_file: str, gain: float
+    ) -> Tuple[Dict[str, Union[torch.Tensor, int]]]:
         return from_disk_as_dict_3d(audio_file=audio_file, gain=gain)

processors/SignalProcessingLoudness.py

@@ -9,13 +9,16 @@ Description:
     Loudness node
 """
 
+import torch
+from typing import Tuple, Dict, Any, Union
+
 from ..core.io import audio_from_comfy_2d
 from ..core.loudness import get_loudness
 
 
 class SignalProcessingLoudness:
     @classmethod
-    def INPUT_TYPES(cls):
+    def INPUT_TYPES(cls) -> Dict[str, Any]:
         return {
             "required": {
                 "audio_input": ("AUDIO",),
@@ -24,10 +27,10 @@ class SignalProcessingLoudness:
 
     RETURN_TYPES = ("FLOAT",)
     RETURN_NAMES = ("loudness",)
-    CATEGORY = "Audio Processing"
+    CATEGORY = "Signal Processing"
     FUNCTION = "process"
 
-    def process(self, audio_input):
+    def process(self, audio_input: Dict[str, Union[torch.Tensor, int]]) -> Tuple[float]:
         waveform, sample_rate = audio_from_comfy_2d(audio_input, try_gpu=True)
 
         loudness: float = get_loudness(waveform, sample_rate)

processors/SignalProcessingMixdown.py

@@ -11,7 +11,7 @@ Description:
 
 import torch
 
-from typing import Tuple, List, Dict, Union
+from typing import Tuple, List, Dict, Union, Any
 import torchaudio
 
 from ..core.io import audio_to_comfy_3d
@@ -20,7 +20,7 @@ from ..core.loudness import lufs_normalization
 
 class SignalProcessingMixdown:
     @classmethod
-    def INPUT_TYPES(cls) -> Dict:
+    def INPUT_TYPES(cls) -> Dict[str, Any]:
         return {
             "required": {
                 "audio_inputs": ("AUDIO_LIST", {"default": []}),
@@ -36,19 +36,21 @@ class SignalProcessingMixdown:
 
     RETURN_TYPES = ("AUDIO",)
     RETURN_NAMES = ("mixed_audio",)
-    CATEGORY = "Audio Processing"
+    CATEGORY = "Signal Processing"
     FUNCTION = "process"
 
     def process(
         self,
         audio_inputs: List[Dict[str, Union[torch.Tensor, int]]],
         gain_factors: List[float] = [],
-    ) -> Tuple[Dict[str, torch.Tensor], int]:
+    ) -> Tuple[Dict[str, Union[torch.Tensor, int]]]:
         """
-        Mix down multiple audio inputs into a single audio output with optional individual volume controls.
+        Mix down multiple audio inputs into a single audio output
+        with optional individual volume controls.
 
         Parameters:
-            audio_inputs (List[Dict]): List of audio inputs, each containing 'waveform' and 'sample_rate'.
+            audio_inputs (List[Dict]): List of audio inputs,
+            each containing 'waveform' and 'sample_rate'.
             output_normalization (float): Normalization factor for the mixed audio (0.0 to 1.0).
             gain_factors (List[float], optional): List of gain factors for each audio input.
 
@@ -66,11 +68,12 @@ class SignalProcessingMixdown:
             gain_factors = [1.0] * num_audios
         elif len(gain_factors) != num_audios:
             raise ValueError(
-                f"Number of gain factors ({len(gain_factors)}) does not match number of audio inputs ({num_audios})."
+                f"Number of gain factors ({len(gain_factors)}) \
+                    does not match number of audio inputs ({num_audios})."
             )
 
         # Extract sample rates and verify consistency
-        sample_rates = [audio["sample_rate"] for audio in audio_inputs]
+        sample_rates: List[int] = [audio["sample_rate"] for audio in audio_inputs]
         target_sample_rate = sample_rates[0]
 
         for idx, sr in enumerate(sample_rates):
@@ -78,20 +81,23 @@ class SignalProcessingMixdown:
                 resampler = torchaudio.transforms.Resample(
                     orig_freq=sr, new_freq=target_sample_rate
                 )
+
+                _waveform: torch.Tensor = audio_inputs[idx]["waveform"]
                 resampler.to(
-                    device=audio_inputs[idx]["waveform"].device,
-                    dtype=audio_inputs[idx]["waveform"].dtype,
+                    device=_waveform.device,
+                    dtype=_waveform.dtype,
                 )
-                audio_inputs[idx]["waveform"] = resampler(audio_inputs[idx]["waveform"])
+                audio_inputs[idx]["waveform"] = resampler(_waveform)
                 audio_inputs[idx]["sample_rate"] = target_sample_rate
 
         # Determine the maximum length among all audio inputs
-        lengths = [audio["waveform"].shape[-1] for audio in audio_inputs]
+        lengthsw: List[torch.Tensor] = [audio["waveform"] for audio in audio_inputs]
+        lengths: List[int] = [wave.shape[-1] for wave in lengthsw]
         max_length = max(lengths)
 
         # Pad or truncate each audio to match the maximum length and apply gain
         for idx, audio in enumerate(audio_inputs):
-            waveform = audio["waveform"]
+            waveform: torch.Tensor = audio["waveform"]
             current_length = waveform.shape[-1]
             gain = gain_factors[idx]
 
@@ -109,7 +115,7 @@ class SignalProcessingMixdown:
             audio["waveform"] = waveform
 
         # Sum all waveforms to create the mix
-        mixed_waveform = torch.zeros_like(audio_inputs[0]["waveform"])
+        mixed_waveform: torch.Tensor = torch.zeros_like(audio_inputs[0]["waveform"])
         for idx, audio in enumerate(audio_inputs):
             mixed_waveform += audio["waveform"]
 

processors/SignalProcessingNormalizer.py

@@ -9,6 +9,8 @@ Description:
     Varios normalizatin techqniues node
 """
 
+import torch
+from typing import Dict, Any, Tuple, Union
 from ..core.io import audio_to_comfy_3d, audio_from_comfy_2d
 from ..core.loudness import (
     rms_normalization,
@@ -20,7 +22,7 @@ from ..core.loudness import (
 
 class SignalProcessingNormalizer:
     @classmethod
-    def INPUT_TYPES(cls):
+    def INPUT_TYPES(cls) -> Dict[str, Any]:
         return {
             "required": {
                 "audio_input": ("AUDIO",),
@@ -50,19 +52,19 @@ class SignalProcessingNormalizer:
 
     RETURN_TYPES = ("AUDIO",)
     RETURN_NAMES = ("processed_audio",)
-    CATEGORY = "Audio Processing"
+    CATEGORY = "Signal Processing"
     FUNCTION = "process"
 
     def process(
         self,
-        audio_input,
+        audio_input: Dict[str, Union[torch.Tensor, int]],
         mode: str,
         target_rms: float,
         target_lufs_db: float,
         target_peak: float,
         target_auto: float,
         target_auto_alpha: float,
-    ):
+    ) -> Tuple[Dict[str, torch.Tensor]]:
 
         waveform, sample_rate = audio_from_comfy_2d(audio_input, try_gpu=True)
 

processors/SignalProcessingSaturation.py

@@ -0,0 +1,61 @@
+#!/usr/bin/env python3
+# -*- coding: utf-8 -*-
+"""
+Author: C0ffymachyne
+License: GPLv3
+Version: 1.0.0
+
+Description:
+    Classic Audio filter set
+
+Reference:
+    https://ccrma.stanford.edu/~jatin/ComplexNonlinearities/Hysteresis.html
+    https://viennatalk.mdw.ac.at/papers/Pap_01_79_Tronchin.pdf
+    https://jatinchowdhury18.medium.com/complex-nonlinearities-episode-3-hysteresis-fdeb2cd3e3f6
+    https://ccrma.stanford.edu/~dtyeh/papers/yeh07_dafx_clipode.pdf
+"""
+
+import torch
+from typing import Dict, Any, Tuple, Union
+
+from ..core.io import audio_to_comfy_3d, audio_from_comfy_2d
+from ..core.saturation import saturator, saturator_get_modes
+from ..core.loudness import get_loudness, lufs_normalization
+
+
+class SignalProcessingSaturation:
+    @classmethod
+    def INPUT_TYPES(cls) -> Dict[str, Any]:
+        return {
+            "required": {
+                "audio_input": ("AUDIO",),
+                "mode": (saturator_get_modes(),),
+                "drive": (
+                    "FLOAT",
+                    {"default": 50.0, "min": 0.0, "max": 200.0, "step": 0.1},
+                ),
+            }
+        }
+
+    RETURN_TYPES = ("AUDIO",)
+    RETURN_NAMES = ("audio",)
+    CATEGORY = "Signal Processing"
+    FUNCTION = "process"
+
+    def process(
+        self,
+        audio_input: Dict[str, Union[torch.Tensor, int]],
+        mode: str = "poly",
+        drive: float = 50.0,
+    ) -> Tuple[Dict[str, Union[torch.Tensor, int]]]:
+        waveform, sample_rate = audio_from_comfy_2d(audio_input, try_gpu=True)
+
+        loudness = get_loudness(waveform, sample_rate)
+
+        filtered_waveform = saturator(
+            waveform, mode=mode, sample_rate=sample_rate, drive=drive
+        )
+
+        filtered_waveform = lufs_normalization(filtered_waveform, sample_rate, loudness)
+
+        return audio_to_comfy_3d(filtered_waveform, sample_rate)

pyproject.toml

@@ -12,3 +12,20 @@ PublisherId = "c0ffymachyne"
 DisplayName = "ComfyUI_SignalProcessing"
 Icon = "images/icon.jpg"
 Models = []
+
+[project.dependencies] 
+torch = [
+    "--index-url",
+    "https://download.pytorch.org/whl/cu118",
+    "torch==2.4.1",
+    "torchaudio==2.4.1",
+    "torchvision==0.19.1"
+]
+other = [
+    "numpy>=1.23.0",
+    "scipy>=1.5.0",
+    "pyfar",
+    "scipy",
+    "pyloudnorm",
+    "cupy"
+]

pytest.ini

@@ -0,0 +1,9 @@
+[pytest]
+markers = 
+  inference: mark as inference test (deselect with '-m "not inference"')
+  execution: mark as execution test (deselect with '-m "not execution"')
+testpaths =
+  tests
+  tests-unit
+addopts = -s
+pythonpath = .

requirements.txt

@@ -4,3 +4,4 @@ numpy
 torch
 torchaudio
 pyloudnorm
+cupy

tests/conftest.py

@@ -0,0 +1,57 @@
+import pytest
+from typing import Dict, Union
+from pathlib import Path
+import torch
+from ..core.io import from_disk_as_raw_3d
+
+# Test data roots
+DATA_ROOT = Path("ComfyUI_SignalProcessing/audio")
+INPUT_FILE = DATA_ROOT / "inputs/002-orig.mp4"
+INPUT_FILE = DATA_ROOT / "inputs/pf-01.mp3"
+INPUT_IR_FILE = DATA_ROOT / "inputs/ir.wav"
+INPUTS_ROOT = DATA_ROOT / "inputs"
+INPUT_FILES = {
+    file.name: file.resolve() for file in INPUTS_ROOT.rglob("*") if file.is_file()
+}
+
+TestData = Dict[str, Union[Dict[str, Union[torch.Tensor, int]], Path]]
+
+
+@pytest.fixture
+def test_data(request) -> TestData:
+
+    test_name = request.node.name  # Automatically get the current test function name
+    test_name = str(request.node.function.__name__)
+
+    OUTPUT_ROOT = DATA_ROOT / f"outputs/{test_name}"
+    INPUTS_ROOT = DATA_ROOT / "inputs"
+
+    param_values = (
+        request.node.callspec.params if hasattr(request.node, "callspec") else {}
+    )
+    pest_param_str = "_".join(f"{key}-{value}" for key, value in param_values.items())
+
+    audio_slice_begin_seconds: float = 60.0
+    audio_slice_duration_seconds: float = 120.0
+
+    # Prepare audio data
+    audio, sample_rate = from_disk_as_raw_3d(
+        str(INPUT_FILE.absolute()),
+        try_gpu=True,
+        start_seconds=audio_slice_begin_seconds,
+        duration_seconds=audio_slice_duration_seconds,
+    )
+    OUTPUT_ROOT.mkdir(parents=True, exist_ok=True)  # Ensure the output directory exists
+
+    audio_to_comfy: Dict[torch.Tensor, int] = {
+        "waveform": audio,
+        "sample_rate": sample_rate,
+    }
+
+    return {
+        "audio": audio_to_comfy,
+        "output_root": OUTPUT_ROOT,
+        "inputs_root": INPUTS_ROOT,
+        "test_name": test_name,
+        "pest_param_str": pest_param_str,
+    }

tests/test_baxandall.py

@@ -0,0 +1,93 @@
+import pytest
+from ..tests.conftest import TestData
+from typing import Tuple
+from ..core.io import audio_from_comfy_3d_to_disk
+from ..processors.SignalProcessingBaxandallEQ import (
+    SignalProcessingBaxandallEQ,
+    SignalProcessingBaxandall3BandEQ,
+)
+
+
+TEST_NAME = "baxandall"
+
+params: list[Tuple[float, float]] = [
+    (9.0, 0.0),
+    (6.0, 0.0),
+    (3.0, 0.0),
+    (1.0, 0.0),
+    (0.0, 9.0),
+    (0.0, 6.0),
+    (0.0, 3.0),
+    (0.0, 1.0),
+    (0.0, 0.0),
+]
+
+
+@pytest.mark.parametrize(
+    "bass_gain_db, treble_gain_db",
+    params,
+)
+def test_baxandalleq_general(
+    test_data: TestData, bass_gain_db: float, treble_gain_db: float
+) -> None:
+
+    node = SignalProcessingBaxandallEQ()
+
+    output = node.process(
+        audio_input=test_data["audio"],
+        bass_gain_db=bass_gain_db,
+        treble_gain_db=treble_gain_db,
+    )[0]
+
+    pest_param_str = test_data["pest_param_str"]
+    output_filepath = test_data["output_root"] / f"{pest_param_str}.wav"
+    audio_from_comfy_3d_to_disk(output, output_filepath)
+
+    assert output_filepath.exists(), f"Output file {output_filepath} was not created."
+    assert output is not None, "Processed audio output is None."
+
+    print(f"test_baxandalleq_general {test_data['output_root']}")
+
+
+@pytest.mark.parametrize(
+    "bass_gain_db,mid_gain_db,treble_gain_db,low_freq,mid_freq,high_freq,mid_q",
+    [
+        (9.0, 9.0, 9.0, 100.0, 1000.0, 10000.0, 0.707),
+        (6.0, 6.0, 6.0, 100.0, 1000.0, 10000.0, 0.707),
+        (3.0, 3.0, 3.0, 100.0, 1000.0, 10000.0, 0.707),
+        (1.0, 1.0, 1.0, 100.0, 1000.0, 10000.0, 0.707),
+        (0.0, 0.0, 0.0, 100.0, 1000.0, 10000.0, 0.707),
+    ],
+)
+def test_baxandalleq3band_general(
+    test_data: TestData,
+    bass_gain_db: float,
+    mid_gain_db: float,
+    treble_gain_db: float,
+    low_freq: float,
+    mid_freq: float,
+    high_freq: float,
+    mid_q: float,
+) -> None:
+
+    node = SignalProcessingBaxandall3BandEQ()
+
+    output = node.process(
+        audio_input=test_data["audio"],
+        bass_gain_db=bass_gain_db,
+        mid_gain_db=mid_gain_db,
+        treble_gain_db=treble_gain_db,
+        low_freq=low_freq,
+        mid_freq=mid_freq,
+        high_freq=high_freq,
+        mid_q=mid_q,
+    )[0]
+
+    pest_param_str = test_data["pest_param_str"]
+    output_filepath = test_data["output_root"] / f"{pest_param_str}.wav"
+    audio_from_comfy_3d_to_disk(output, output_filepath)
+
+    assert output_filepath.exists(), f"Output file {output_filepath} was not created."
+    assert output is not None, "Processed audio output is None."
+
+    print(f"test_baxandalleq3band_general {test_data['output_root']}")

tests/test_compressor.py

@@ -0,0 +1,50 @@
+import pytest
+from ..tests.conftest import TestData
+from typing import Tuple
+from ..core.io import audio_from_comfy_3d_to_disk
+from ..processors.SignalProcessingCompressor import SignalProcessingCompressor
+
+TEST_NAME = "compression"
+
+params: list[Tuple[float, float, float, float]] = [
+    (0.9, 0.1, 60.0, 0.3),
+    (0.3, 0.1, 60.0, 0.3),
+    (0.1, 0.1, 60.0, 0.3),
+    (-0.0, 0.1, 60.0, 0.3),
+    (-0.1, 0.1, 60.0, 0.3),
+    (-0.3, 0.1, 60.0, 0.3),
+    (-0.9, 0.1, 60.0, 0.3),
+]
+
+
+@pytest.mark.parametrize(
+    "comp, attack, release, filter_param",
+    params,
+)
+def test_compressor_general(
+    test_data: TestData, comp: float, attack: float, release: float, filter_param: float
+) -> None:
+    """
+    Test SignalProcessingCompressor with various parameter configurations.
+    """
+    node = SignalProcessingCompressor()
+
+    # Process input audio
+    output = node.process(
+        audio_input=test_data["audio"],
+        comp=comp,
+        attack=attack,
+        release=release,
+        filter_param=filter_param,
+    )[0]
+
+    # Save the output audio
+    pest_param_str = test_data["pest_param_str"]
+    output_filepath = test_data["output_root"] / f"{pest_param_str}.wav"
+    audio_from_comfy_3d_to_disk(output, output_filepath)
+
+    # Assertions
+    assert output_filepath.exists(), f"Output file {output_filepath} was not created."
+    assert output is not None, "Processed audio output is None."
+
+    print(f"test_compressor_general {test_data['output_root']}")

tests/test_convolution_reverb.py

@@ -0,0 +1,36 @@
+import pytest
+from .conftest import TestData
+from ..core.io import audio_from_comfy_3d_to_disk
+from ..effects.SignalProcessingConvolutionReverb import (
+    SignalProcessingConvolutionReverb,
+)
+
+TEST_NAME = "convolution_reverb"
+
+
+@pytest.mark.parametrize(
+    "impulse_response, wet_dry",
+    [("ir.wav", 1.0), ("ir.wav", 0.6), ("ir.wav", 0.3), ("ir.wav", 0.0)],
+)
+def test_convolution_reverb_general(
+    test_data: TestData, impulse_response: str, wet_dry: float
+) -> None:
+    SignalProcessingConvolutionReverb.ir_directory = test_data["inputs_root"] / "ir"
+    node = SignalProcessingConvolutionReverb()
+
+    node.INPUT_TYPES()
+
+    output = node.process(
+        impulse_response=impulse_response,
+        audio_input=test_data["audio"],
+        wet_dry=wet_dry,
+    )[0]
+
+    pest_param_str = test_data["pest_param_str"]
+    output_filepath = test_data["output_root"] / f"{pest_param_str}.wav"
+    audio_from_comfy_3d_to_disk(output, output_filepath)
+
+    assert output_filepath.exists(), f"Output file {output_filepath} was not created."
+    assert output is not None, "Processed audio output is None."
+
+    print(f"test_convolution_reverb_general {test_data['output_root']}")

tests/test_filter.py

@@ -0,0 +1,43 @@
+import pytest
+from ..tests.conftest import TestData
+from ..core.io import audio_from_comfy_3d_to_disk
+from ..processors.SignalProcessingFilter import (
+    SignalProcessingFilter,
+)
+
+TEST_NAME = "filter"
+
+
+@pytest.mark.parametrize(
+    "cutoff, q_factor",
+    [
+        (0.9, 0.707),
+        (0.6, 0.707),
+        (0.3, 0.707),
+        (0.1, 0.707),
+        (0.0, 0.707),
+    ],
+)
+def test_filter_general(test_data: TestData, cutoff: float, q_factor: float) -> None:
+
+    node: SignalProcessingFilter = SignalProcessingFilter()
+    modes: str = node.INPUT_TYPES()["required"]["filter_type"][0]  # Extract modes
+
+    for mode in modes:
+        output = node.process(
+            audio_input=test_data["audio"],
+            cutoff=cutoff,
+            filter_type=mode,
+            q_factor=q_factor,
+        )[0]
+
+        pest_param_str = test_data["pest_param_str"]
+        output_filepath = test_data["output_root"] / f"{mode}-{pest_param_str}.wav"
+        audio_from_comfy_3d_to_disk(output, output_filepath)
+
+        assert (
+            output_filepath.exists()
+        ), f"Output file {output_filepath} was not created."
+        assert output is not None, f"Processed audio output is None for mode {mode}."
+
+    print(f"test_filter_general {test_data['output_root']}")

tests/test_harmonics.py

@@ -0,0 +1,49 @@
+import pytest
+from ..tests.conftest import TestData
+from typing import Tuple
+from ..core.io import audio_from_comfy_3d_to_disk
+from ..processors.SignalProcessingHarmonicsEnhancer import (
+    SignalProcessingHarmonicsEnhancer,
+)
+
+TEST_NAME = "harmonics"
+
+params: list[Tuple[str, int, float, float]] = [
+    ("2,3", 440, 6.0, 0.707),
+    ("2,3", 440, 3.0, 0.707),
+    ("2,3", 440, 1.0, 0.707),
+    ("2,3", 440, 0.0, 0.707),
+]
+
+
+@pytest.mark.parametrize(
+    "harmonics, base_frequency, gain_db, Q",
+    params,
+)
+def test_harmonics_general(
+    test_data: TestData, harmonics: str, base_frequency: int, gain_db: int, Q: float
+) -> None:
+
+    node = SignalProcessingHarmonicsEnhancer()
+    modes: str = node.INPUT_TYPES()["required"]["mode"][0]  # Extract modes
+
+    for mode in modes:
+        output = node.process(
+            audio_input=test_data["audio"],
+            harmonics=harmonics,
+            mode=mode,
+            base_frequency=base_frequency,
+            gain_db=gain_db,
+            Q=Q,
+        )[0]
+
+        pest_param_str = test_data["pest_param_str"]
+        output_filepath = test_data["output_root"] / f"{mode}-{pest_param_str}.wav"
+        audio_from_comfy_3d_to_disk(output, output_filepath)
+
+        assert (
+            output_filepath.exists()
+        ), f"Output file {output_filepath} was not created."
+        assert output is not None, f"Processed audio output is None for mode {mode}."
+
+    print(f"test_harmonics_general {test_data['output_root']}")

tests/test_limiting.py

@@ -0,0 +1,47 @@
+import pytest
+from ..tests.conftest import TestData
+from typing import Tuple
+from ..core.io import audio_from_comfy_3d_to_disk
+from ..processors.SignalProcessingLimiter import SignalProcessingLimiter
+
+TEST_NAME = "limiting"
+
+params: list[Tuple[float, float, float]] = [
+    (0.0, 100.0, 600.0),
+    (10.0, 100.0, 600.0),
+    (30.0, 100.0, 600.0),
+    (90.0, 100.0, 600.0),
+]
+
+
+@pytest.mark.parametrize(
+    "threshold, slope, release_ms",
+    params,
+)
+def test_limiting_general(
+    test_data: TestData, threshold: float, slope: float, release_ms: float
+) -> None:
+
+    node = SignalProcessingLimiter()
+    modes = node.INPUT_TYPES()["required"]["mode"][0]  # Extract modes
+
+    for mode in modes:
+
+        output = node.process(
+            audio_input=test_data["audio"],
+            mode=mode,
+            threshold=threshold,
+            slope=slope,
+            release_ms=release_ms,
+        )[0]
+
+        pest_param_str = test_data["pest_param_str"]
+        output_filepath = test_data["output_root"] / f"{mode}-{pest_param_str}.wav"
+        audio_from_comfy_3d_to_disk(output, output_filepath)
+
+        assert (
+            output_filepath.exists()
+        ), f"Output file {output_filepath} was not created."
+        assert output is not None, f"Processed audio output is None for mode {mode}."
+
+    print(f"test_limiting_general {test_data['output_root']}")

tests/test_normalizer.py

@@ -0,0 +1,51 @@
+import pytest
+from .conftest import TestData
+from ..core.io import audio_from_comfy_3d_to_disk
+from ..processors.SignalProcessingNormalizer import SignalProcessingNormalizer
+
+TEST_NAME = "normalizer"
+
+
+@pytest.mark.parametrize(
+    "target_rms, target_lufs_db, target_peak, target_auto, target_auto_alpha",
+    [
+        (-0.0, -14.0, -0.0, -0.0, -0.0),
+        (0.3, -14.0, 0.3, 0.3, 0.3),
+        (0.6, -14.0, 0.6, 0.6, 0.6),
+        (0.9, -14.0, 0.9, 0.9, 0.9),
+    ],
+)
+def test_normalizer_general(
+    test_data: TestData,
+    target_rms: float,
+    target_lufs_db: float,
+    target_peak: float,
+    target_auto: float,
+    target_auto_alpha: float,
+) -> None:
+
+    node = SignalProcessingNormalizer()
+    modes = node.INPUT_TYPES()["required"]["mode"][0]
+
+    for mode in modes:
+
+        output = node.process(
+            audio_input=test_data["audio"],
+            mode=mode,
+            target_rms=target_rms,
+            target_lufs_db=target_lufs_db,
+            target_peak=target_peak,
+            target_auto=target_auto,
+            target_auto_alpha=target_auto_alpha,
+        )[0]
+
+        pest_param_str = test_data["pest_param_str"]
+        output_filepath = test_data["output_root"] / f"{mode}-{pest_param_str}.wav"
+        audio_from_comfy_3d_to_disk(output, output_filepath)
+
+        assert (
+            output_filepath.exists()
+        ), f"Output file {output_filepath} was not created."
+        assert output is not None, f"Processed audio output is None for mode {mode}."
+
+    print(f"test_normalizer_general {test_data['output_root']}")

tests/test_padsynthchoir.py

@@ -0,0 +1,30 @@
+from pathlib import Path
+from ..tests.conftest import TestData
+from ..generators.SignalProcessingPadSynthChoir import SignalProcessingPadSynthChoir
+
+TEST_NAME = "synth_and_mixdown"
+OUTPUT_ROOT = Path(f"ComfyUI_SignalProcessing/audio/outputs/{TEST_NAME}")
+
+
+def test_synth_and_mixdown(test_data: TestData) -> None:
+
+    OUTPUT_ROOT.mkdir(parents=True, exist_ok=True)
+
+    samplerate = 44100
+    base_freq = 440.0
+    step_size = 4
+    num_notes = 5
+    bandwidth_cents = 60.0
+    number_harmonics = 32
+
+    synth_node = SignalProcessingPadSynthChoir()
+    synth_output, sample_rate = synth_node.process(
+        samplerate=samplerate,
+        base_freq=base_freq,
+        step_size=step_size,
+        num_notes=num_notes,
+        bandwidth_cents=bandwidth_cents,
+        number_harmonics=number_harmonics,
+    )
+
+    print(f"test_synth_and_mixdown {test_data['output_root']}")

tests/test_paulstretch.py

@@ -0,0 +1,42 @@
+import pytest
+from .conftest import TestData
+from ..core.io import audio_from_comfy_3d_to_disk
+from ..effects.SignalProcessingPaulStretch import SignalProcessingPaulStretch
+
+TEST_NAME = "paulstretch"
+
+
+@pytest.mark.parametrize(
+    "stretch_factor, window_size_seconds",
+    [
+        (9.0, 0.25),
+        (6.0, 0.25),
+        (3.0, 0.25),
+        (1.0, 0.25),
+        (2.0, 0.25),
+        (2.0, 0.5),
+        (2.0, 0.75),
+        (2.0, 1.0),
+        (2.0, 3.0),
+        (2.0, 6.0),
+    ],
+)
+def test_paul_stretch_general(
+    test_data: TestData, stretch_factor: float, window_size_seconds: float
+) -> None:
+
+    node = SignalProcessingPaulStretch()
+    output = node.process(
+        audio_input=test_data["audio"],
+        stretch_factor=stretch_factor,
+        window_size_seconds=window_size_seconds,
+    )[0]
+
+    pest_param_str = test_data["pest_param_str"]
+    output_filepath = test_data["output_root"] / f"{pest_param_str}.wav"
+    audio_from_comfy_3d_to_disk(output, output_filepath)
+
+    assert output_filepath.exists(), f"Output file {output_filepath} was not created."
+    assert output is not None, "Processed audio output is None."
+
+    print(f"test_paul_stretch_general {test_data['output_root']}")

tests/test_pitchshift.py

@@ -0,0 +1,28 @@
+import pytest
+from .conftest import TestData
+from ..core.io import audio_from_comfy_3d_to_disk
+from ..effects.SignalProcessingPitchShifter import SignalProcessingPitchShifter
+
+TEST_NAME = "pitchshift"
+
+
+@pytest.mark.parametrize(
+    "pitch_shift_factor",
+    [(3), (1), (0), (-1), (-3)],
+)
+def test_pitch_shift_general(test_data: TestData, pitch_shift_factor: int) -> None:
+
+    node = SignalProcessingPitchShifter()
+
+    output = node.process(
+        audio_input=test_data["audio"], pitch_shift_factor=pitch_shift_factor
+    )[0]
+
+    pest_param_str = test_data["pest_param_str"]
+    output_filepath = test_data["output_root"] / f"{pest_param_str}.wav"
+    audio_from_comfy_3d_to_disk(output, output_filepath)
+
+    assert output_filepath.exists(), f"Output file {output_filepath} was not created."
+    assert output is not None, "Processed audio output is None."
+
+    print(f"test_pitch_shift_general {test_data['output_root']}")

tests/test_plotting.py

@@ -0,0 +1,96 @@
+import pytest
+from torchvision.transforms import ToPILImage
+from .conftest import TestData
+from ..core.plotting import get_spectogram, get_wave, save_image
+from ..visuals.SignalProcessingSpectrogram import SignalProcessingSpectrogram
+
+TEST_NAME = "plotting"
+
+
+@pytest.mark.parametrize(
+    "n_fft, n_mels, xlim",
+    [
+        (4096, 128 * 1, 8192),
+        (4096, 128 * 2, 8192),
+        (4096, 128 * 4, 8192),
+        (4096, 128 * 6, 8192),
+    ],
+)
+def test_plotting_spectogram_general(
+    test_data: TestData, n_fft: int, n_mels: int, xlim: int
+) -> None:
+
+    waveform = test_data["audio"]["waveform"].squeeze(0)
+    sample_rate = test_data["audio"]["sample_rate"]
+
+    print("waveform", waveform.shape)
+
+    spectogram = get_spectogram(
+        waveform, sample_rate=sample_rate, n_fft=n_fft, n_mels=n_mels, xlim=xlim
+    )
+
+    pest_param_str = test_data["pest_param_str"]
+    output_filepath = test_data["output_root"] / f"{pest_param_str}.png"
+
+    save_image(output_filepath, spectogram)
+
+    assert output_filepath.exists(), f"Output file {output_filepath} was not created."
+    assert output_filepath is not None, "Processed audio output is None."
+
+    print(f"test_plotting_spectogram_general {test_data['output_root']}")
+
+
+@pytest.mark.parametrize(
+    "stretch_factor, window_size_seconds",
+    [(8.0, 0.25)],
+)
+def test_plotting_waveform_general(
+    test_data: TestData, stretch_factor: float, window_size_seconds: float
+) -> None:
+
+    waveform = test_data["audio"]["waveform"].squeeze(0)
+    sample_rate = test_data["audio"]["sample_rate"]
+
+    print("waveform", waveform.shape)
+
+    spectogram = get_wave(waveform, sample_rate=sample_rate, xlim=4096)
+
+    pest_param_str = test_data["pest_param_str"]
+    output_filepath = test_data["output_root"] / f"{pest_param_str}.png"
+
+    save_image(output_filepath, spectogram)
+
+    assert output_filepath.exists(), f"Output file {output_filepath} was not created."
+    assert output_filepath is not None, "Processed audio output is None."
+
+    print(f"test_plotting_spectogram_general {test_data['output_root']}")
+
+
+@pytest.mark.parametrize(
+    "stretch_factor, window_size_seconds",
+    [(8.0, 0.25)],
+)
+def test_plotting_waveform_node_general(
+    test_data: TestData, stretch_factor: float, window_size_seconds: float
+) -> None:
+
+    node = SignalProcessingSpectrogram()
+
+    output = node.process(audio_input=test_data["audio"])[0]
+
+    pest_param_str = test_data["pest_param_str"]
+    output_filepath = test_data["output_root"] / f"{pest_param_str}.png"
+
+    to_pil = ToPILImage()
+    print("rgb_image[0] -----------------------------", output[0].shape)
+    rgb_image = output[0][..., :3]
+    rgb_image = rgb_image.permute(2, 0, 1)
+    print("rgb_image -----------------------------", rgb_image.shape)
+    spectogram = to_pil(rgb_image)
+    print("spectogram -----------------------------", spectogram)
+    save_image(output_filepath, spectogram)
+
+    assert output_filepath.exists(), f"Output file {output_filepath} was not created."
+    assert output is not None, "Processed audio output is None."
+
+    print(f"test_pitch_shift_general {test_data['output_root']}")

tests/test_saturation.py

@@ -0,0 +1,37 @@
+import pytest
+from ..tests.conftest import TestData
+from ..core.io import audio_from_comfy_3d_to_disk
+from ..processors.SignalProcessingSaturation import SignalProcessingSaturation
+
+TEST_NAME = "saturation"
+
+
+@pytest.mark.parametrize(
+    "drive",
+    [
+        (90.0),
+        (60.0),
+        (30.0),
+        (10.0),
+        (0.0),
+    ],
+)
+def test_saturation_general(test_data: TestData, drive: float) -> None:
+
+    node = SignalProcessingSaturation()
+    modes = node.INPUT_TYPES()["required"]["mode"][0]  # Extract modes
+
+    for mode in modes:
+
+        output = node.process(audio_input=test_data["audio"], mode=mode, drive=drive)[0]
+
+        pest_param_str = test_data["pest_param_str"]
+        output_filepath = test_data["output_root"] / f"{mode}-{pest_param_str}.wav"
+        audio_from_comfy_3d_to_disk(output, output_filepath)
+
+        assert (
+            output_filepath.exists()
+        ), f"Output file {output_filepath} was not created."
+        assert output is not None, f"Processed audio output is None for mode {mode}."
+
+    print(f"test_saturation_general {test_data['output_root']}")

tests/test_widening.py

@@ -0,0 +1,38 @@
+import pytest
+from ..tests.conftest import TestData
+from ..core.io import audio_from_comfy_3d_to_disk
+from ..effects.SignalProcessingStereoWidening import SignalProcessingStereoWidening
+
+TEST_NAME = "widening"
+
+
+@pytest.mark.parametrize(
+    "width",
+    [
+        (3.0),
+        (2.0),
+        (1.0),
+        (0.5),
+        (0.25),
+        (0.0),
+    ],
+)
+def test_widening_general(test_data: TestData, width: float) -> None:
+
+    node = SignalProcessingStereoWidening()
+    modes = node.INPUT_TYPES()["required"]["mode"][0]
+
+    for mode in modes:
+
+        output = node.process(audio_input=test_data["audio"], mode=mode, width=width)[0]
+
+        pest_param_str = test_data["pest_param_str"]
+        output_filepath = test_data["output_root"] / f"{mode}-{pest_param_str}.wav"
+        audio_from_comfy_3d_to_disk(output, output_filepath)
+
+        assert (
+            output_filepath.exists()
+        ), f"Output file {output_filepath} was not created."
+        assert output is not None, f"Processed audio output is None for mode {mode}."
+
+    print(f"test_widening_general {test_data['output_root']}")

visuals/SignalProcessingSpectrogram.py

@@ -11,22 +11,20 @@ Description:
 
 import torch
 
-from typing import List, Dict
 from PIL import Image
 
+from typing import Dict, Tuple, Any, Type, List
+
 import numpy as np
 import torchaudio
 
 import matplotlib.pyplot as plt
 
-from ..core.io import audio_to_comfy_3d, audio_from_comfy_3d, audio_from_comfy_2d
-from ..core.loudness import lufs_normalization, get_loudness
-
 
 class SignalProcessingSpectrogram:
     @classmethod
-    def INPUT_TYPES(cls):
-        cmaps = ["viridis", "plasma", "inferno", "magma", "cividis"]
+    def INPUT_TYPES(cls: Type["SignalProcessingSpectrogram"]) -> Dict[str, Any]:
+        cmaps: List[str] = ["viridis", "plasma", "inferno", "magma", "cividis"]
         return {
             "required": {
                 "audio_input": ("AUDIO",),
@@ -51,23 +49,37 @@ class SignalProcessingSpectrogram:
 
     RETURN_TYPES = ("IMAGE",)
     RETURN_NAMES = ("spectrogram_image",)
-    CATEGORY = "Audio Processing"
+    CATEGORY = "Signal Processing"
     FUNCTION = "process"
 
     def process(
         self,
-        audio_input,
-        color_map,
-        n_fft=2048,
-        hop_length=512,
-        n_mels=128,
-        top_db=80.0,
-    ):
-        waveform = audio_input.get(
-            "waveform"
-        )  # [channels, samples] or [batch, channels, samples]
+        audio_input: Dict[str, torch.Tensor],
+        color_map: str = "viridis",
+        n_fft: int = 2048,
+        hop_length: int = 512,
+        n_mels: int = 128,
+        top_db: float = 80.0,
+    ) -> Tuple[torch.Tensor]:
+        waveform = audio_input.get("waveform")
         sample_rate = audio_input.get("sample_rate")
 
+        # Validate that waveform and sample_rate are not None
+        if waveform is None:
+            raise ValueError("The 'waveform' key is missing or None in 'audio_input'.")
+        if not isinstance(waveform, torch.Tensor):
+            raise TypeError(
+                f"Expected 'waveform' to be a torch.Tensor, got {type(waveform)}."
+            )
+        if sample_rate is None:
+            raise ValueError(
+                "The 'sample_rate' key is missing or None in 'audio_input'."
+            )
+        if not isinstance(sample_rate, int):
+            raise TypeError(
+                f"Expected 'sample_rate' to be an int, got {type(sample_rate)}."
+            )
+
         # waveform, sample_rate = audio_from_comfy_2d(audio_input)
 
         # Convert to mono by averaging channels

visuals/SignalProcessingWaveform.py

@@ -15,16 +15,16 @@ from PIL import Image
 import numpy as np
 import matplotlib.pyplot as plt
 
+from typing import Any, Dict, Tuple
+
 from ..core.utilities import comfy_root_to_syspath
 
 comfy_root_to_syspath()  # add comfy to sys path for dev
 
-from ..core.io import audio_to_comfy_3d
-
 
 class SignalProcessingWaveform:
     @classmethod
-    def INPUT_TYPES(cls):
+    def INPUT_TYPES(cls) -> Dict[str, Any]:
         return {
             "required": {
                 "audio_input": ("AUDIO",),
@@ -46,22 +46,21 @@ class SignalProcessingWaveform:
 
     RETURN_TYPES = ("IMAGE",)
     RETURN_NAMES = ("waveform_image",)
-    CATEGORY = "Audio Processing"
+    CATEGORY = "Signal Processing"
     FUNCTION = "process"
 
     def process(
         self,
-        audio_input,
-        color="white",
-        background_color="black",
-        width=800,
-        height=200,
-        line_width=1.0,
-    ):
+        audio_input: torch.Tensor,
+        color: str = "white",
+        background_color: str = "black",
+        width: int = 800,
+        height: int = 200,
+        line_width: float = 1.0,
+    ) -> Tuple[torch.Tensor]:
         waveform = audio_input.get(
             "waveform"
         )  # [channels, samples] or [batch, channels, samples]
-        sample_rate = audio_input.get("sample_rate")
 
         # Convert to mono by averaging channels
         if waveform.ndim == 3:
@@ -93,7 +92,7 @@ class SignalProcessingWaveform:
 
         # Plot the waveform
         plt.plot(waveform_np, color=color, linewidth=line_width)
-        plt.ylim(-1.1, 1.1)  # Set y-axis limits to -1 and 1
+        plt.ylim(-1.3, 1.3)  # Set y-axis limits to -1 and 1
         plt.tight_layout(pad=0)
 
         # Save the plot to a buffer
@@ -119,7 +118,7 @@ class SignalProcessingWaveform:
 
 class SignalProcessingWaveform2:
     @classmethod
-    def INPUT_TYPES(cls):
+    def INPUT_TYPES(cls) -> Dict[str, Any]:
         return {
             "required": {
                 "audio_input": ("AUDIO",),
@@ -141,22 +140,21 @@ class SignalProcessingWaveform2:
 
     RETURN_TYPES = ("IMAGE",)
     RETURN_NAMES = ("waveform_image",)
-    CATEGORY = "Audio Processing"
+    CATEGORY = "Signal Processing"
     FUNCTION = "process"
 
     def process(
         self,
-        audio_input,
-        color="black",
-        background_color="white",
-        width=800,
-        height=200,
-        line_width=1.0,
-    ):
+        audio_input: torch.Tensor,
+        color: str = "black",
+        background_color: str = "white",
+        width: int = 800,
+        height: int = 200,
+        line_width: float = 1.0,
+    ) -> Tuple[torch.Tensor]:
         waveform = audio_input.get(
             "waveform"
         )  # [channels, samples] or [batch, channels, samples]
-        sample_rate = audio_input.get("sample_rate")
 
         # Convert to mono by averaging channels
         if waveform.ndim == 3:
@@ -211,21 +209,3 @@ class SignalProcessingWaveform2:
         image = torch.from_numpy(image_np).permute(2, 0, 1).unsqueeze(0)  # [1, 3, H, W]
 
         return (image,)
-
-
-if __name__ == "__main__":
-
-    from pathlib import Path
-    from ..core.io import from_disk_as_raw_2d, audio_to_comfy_3d
-
-    node = SignalProcessingWaveform()
-    samples_path = Path("ComfyUI_SignalProcessing/audio/inputs/song.mp4")
-
-    source_path = samples_path.absolute()
-    source_audio, source_audio_sample_rate = from_disk_as_raw_2d(source_path)
-
-    input = audio_to_comfy_3d(source_audio, source_audio_sample_rate)[0]
-
-    result = node.process(input)[0]
-
-    print("result", result)