diff --git a/.gitignore b/.gitignore index cf282e3..8da3a88 100644 --- a/.gitignore +++ b/.gitignore @@ -8,6 +8,7 @@ __pycache__/ # development directories audio/ +.nox/ # Distribution / packaging .Python diff --git a/README.md b/README.md index 0eea768..18582dc 100644 --- a/README.md +++ b/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 \ No newline at end of file +- **`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! 🚀 \ No newline at end of file diff --git a/core/compression.py b/core/compression.py new file mode 100644 index 0000000..6deae28 --- /dev/null +++ b/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 diff --git a/core/conversions.py b/core/conversions.py index 8019ae9..d999770 100644 --- a/core/conversions.py +++ b/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 diff --git a/core/filters.py b/core/filters.py new file mode 100644 index 0000000..33930ba --- /dev/null +++ b/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 + ) diff --git a/core/harmonics.py b/core/harmonics.py new file mode 100644 index 0000000..ae7f9ec --- /dev/null +++ b/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 diff --git a/core/io.py b/core/io.py index 1d7224e..7afed12 100644 --- a/core/io.py +++ b/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 diff --git a/core/kernels/compressor/compressor.cu b/core/kernels/compressor/compressor.cu new file mode 100644 index 0000000..e5c7ad2 --- /dev/null +++ b/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; + } +} \ No newline at end of file diff --git a/core/kernels/limiter/limiter_dev.cu b/core/kernels/limiter/limiter_dev.cu new file mode 100644 index 0000000..4a8a353 --- /dev/null +++ b/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; + } + } +} \ No newline at end of file diff --git a/core/kernels/limiter/limiter_down_parallel.cu b/core/kernels/limiter/limiter_down_parallel.cu new file mode 100644 index 0000000..3626121 --- /dev/null +++ b/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; + } +} \ No newline at end of file diff --git a/core/kernels/limiter/limiter_hard_clipper.cu b/core/kernels/limiter/limiter_hard_clipper.cu new file mode 100644 index 0000000..4ad9468 --- /dev/null +++ b/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; + } +} \ No newline at end of file diff --git a/core/kernels/limiter/limiter_soft_clipper.cu b/core/kernels/limiter/limiter_soft_clipper.cu new file mode 100644 index 0000000..5d32f6d --- /dev/null +++ b/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; + } +} \ No newline at end of file diff --git a/core/kernels/limiter/limiter_updown_parallel.cu b/core/kernels/limiter/limiter_updown_parallel.cu new file mode 100644 index 0000000..2fa0f12 --- /dev/null +++ b/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; + } +} \ No newline at end of file diff --git a/core/kernels/limiter/limiter_updown_weighted_parallel.cu b/core/kernels/limiter/limiter_updown_weighted_parallel.cu new file mode 100644 index 0000000..effc30e --- /dev/null +++ b/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; + } +} \ No newline at end of file diff --git a/core/limiting.py b/core/limiting.py new file mode 100644 index 0000000..97c34af --- /dev/null +++ b/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 diff --git a/core/loudness.py b/core/loudness.py index 2caa921..cd30a3e 100644 --- a/core/loudness.py +++ b/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 diff --git a/core/matching.py b/core/matching.py deleted file mode 100644 index 9faae15..0000000 --- a/core/matching.py +++ /dev/null @@ -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() diff --git a/core/mixing.py b/core/mixing.py index f20a034..0f8a05c 100644 --- a/core/mixing.py +++ b/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. diff --git a/core/plotting.py b/core/plotting.py index e498c91..d7c739c 100644 --- a/core/plotting.py +++ b/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() diff --git a/core/sampling.py b/core/sampling.py new file mode 100644 index 0000000..a9f108f --- /dev/null +++ b/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) diff --git a/core/saturation.py b/core/saturation.py index 071bf9e..e62bd4b 100644 --- a/core/saturation.py +++ b/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 diff --git a/core/spectral.py b/core/spectral.py new file mode 100644 index 0000000..2ff4c15 --- /dev/null +++ b/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 diff --git a/core/tests.py b/core/tests.py new file mode 100644 index 0000000..b614763 --- /dev/null +++ b/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" diff --git a/core/utilities.py b/core/utilities.py index 6fdcb50..68d7e93 100644 --- a/core/utilities.py +++ b/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() diff --git a/core/utilitiescuda.py b/core/utilitiescuda.py new file mode 100644 index 0000000..127ec0b --- /dev/null +++ b/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") diff --git a/core/widening.py b/core/widening.py index 789d1c8..9a1afab 100644 --- a/core/widening.py +++ b/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, diff --git a/effects/SignalProcessingConvolutionReverb.py b/effects/SignalProcessingConvolutionReverb.py index 3e299cd..a5a001b 100644 --- a/effects/SignalProcessingConvolutionReverb.py +++ b/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 diff --git a/effects/SignalProcessingPaulStretch.py b/effects/SignalProcessingPaulStretch.py index bbc6325..45cd5bd 100644 --- a/effects/SignalProcessingPaulStretch.py +++ b/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 diff --git a/effects/SignalProcessingPitchShifter.py b/effects/SignalProcessingPitchShifter.py index 6a3d895..cf50174 100644 --- a/effects/SignalProcessingPitchShifter.py +++ b/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, - ) diff --git a/effects/SignalProcessingStereoWidening.py b/effects/SignalProcessingStereoWidening.py index cda25d6..61d0520 100644 --- a/effects/SignalProcessingStereoWidening.py +++ b/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) diff --git a/generators/SignalProcessingPadSynth.py b/generators/SignalProcessingPadSynth.py index 99df412..0dbcd39 100644 --- a/generators/SignalProcessingPadSynth.py +++ b/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. diff --git a/generators/SignalProcessingPadSynthChoir.py b/generators/SignalProcessingPadSynthChoir.py index 41148b4..7af2209 100644 --- a/generators/SignalProcessingPadSynthChoir.py +++ b/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 diff --git a/nodes.py b/nodes.py index f35a4c1..6cd7934 100644 --- a/nodes.py +++ b/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", } diff --git a/noxfile.py b/noxfile.py new file mode 100644 index 0000000..e311d0f --- /dev/null +++ b/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, + ) diff --git a/processors/SignalProcessingBaxandallEQ.py b/processors/SignalProcessingBaxandallEQ.py index 52576ea..4c4d580 100644 --- a/processors/SignalProcessingBaxandallEQ.py +++ b/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, - ) diff --git a/processors/SignalProcessingCompressor.py b/processors/SignalProcessingCompressor.py new file mode 100644 index 0000000..01e7191 --- /dev/null +++ b/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) diff --git a/processors/SignalProcessingFilter.py b/processors/SignalProcessingFilter.py index 148d710..e7f78a4 100644 --- a/processors/SignalProcessingFilter.py +++ b/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( diff --git a/processors/SignalProcessingHarmonicsEnhancer.py b/processors/SignalProcessingHarmonicsEnhancer.py index ce5e398..b524387 100644 --- a/processors/SignalProcessingHarmonicsEnhancer.py +++ b/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 \ No newline at end of file diff --git a/processors/SignalProcessingLimiter.py b/processors/SignalProcessingLimiter.py new file mode 100644 index 0000000..a462ced --- /dev/null +++ b/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) diff --git a/processors/SignalProcessingLoadAudio.py b/processors/SignalProcessingLoadAudio.py index 2e18bd3..c3befe5 100644 --- a/processors/SignalProcessingLoadAudio.py +++ b/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) diff --git a/processors/SignalProcessingLoudness.py b/processors/SignalProcessingLoudness.py index 7beab68..0d4c75c 100644 --- a/processors/SignalProcessingLoudness.py +++ b/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) diff --git a/processors/SignalProcessingMixdown.py b/processors/SignalProcessingMixdown.py index 83f6eb5..e77c4ba 100644 --- a/processors/SignalProcessingMixdown.py +++ b/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"] diff --git a/processors/SignalProcessingNormalizer.py b/processors/SignalProcessingNormalizer.py index ed81efb..abed8e3 100644 --- a/processors/SignalProcessingNormalizer.py +++ b/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) diff --git a/processors/SignalProcessingSaturation.py b/processors/SignalProcessingSaturation.py new file mode 100644 index 0000000..bf2f766 --- /dev/null +++ b/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) diff --git a/pyproject.toml b/pyproject.toml index f5e8cda..49fbf12 100644 --- a/pyproject.toml +++ b/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" +] \ No newline at end of file diff --git a/pytest.ini b/pytest.ini new file mode 100644 index 0000000..a224d8c --- /dev/null +++ b/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 = . diff --git a/requirements.txt b/requirements.txt index 7c2228a..d29df71 100644 --- a/requirements.txt +++ b/requirements.txt @@ -4,3 +4,4 @@ numpy torch torchaudio pyloudnorm +cupy diff --git a/tests/__init__.py b/tests/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/tests/conftest.py b/tests/conftest.py new file mode 100644 index 0000000..d95bba8 --- /dev/null +++ b/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, + } diff --git a/tests/test_baxandall.py b/tests/test_baxandall.py new file mode 100644 index 0000000..c97eabd --- /dev/null +++ b/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']}") diff --git a/tests/test_compressor.py b/tests/test_compressor.py new file mode 100644 index 0000000..b7a3e2b --- /dev/null +++ b/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']}") diff --git a/tests/test_convolution_reverb.py b/tests/test_convolution_reverb.py new file mode 100644 index 0000000..45a8b33 --- /dev/null +++ b/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']}") diff --git a/tests/test_filter.py b/tests/test_filter.py new file mode 100644 index 0000000..76c1b7f --- /dev/null +++ b/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']}") diff --git a/tests/test_harmonics.py b/tests/test_harmonics.py new file mode 100644 index 0000000..c39ecdb --- /dev/null +++ b/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']}") diff --git a/tests/test_limiting.py b/tests/test_limiting.py new file mode 100644 index 0000000..7b66bf9 --- /dev/null +++ b/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']}") diff --git a/tests/test_normalizer.py b/tests/test_normalizer.py new file mode 100644 index 0000000..6a20be3 --- /dev/null +++ b/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']}") diff --git a/tests/test_padsynthchoir.py b/tests/test_padsynthchoir.py new file mode 100644 index 0000000..436bb6e --- /dev/null +++ b/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']}") diff --git a/tests/test_paulstretch.py b/tests/test_paulstretch.py new file mode 100644 index 0000000..402d3fa --- /dev/null +++ b/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']}") diff --git a/tests/test_pitchshift.py b/tests/test_pitchshift.py new file mode 100644 index 0000000..99a1814 --- /dev/null +++ b/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']}") diff --git a/tests/test_plotting.py b/tests/test_plotting.py new file mode 100644 index 0000000..c147833 --- /dev/null +++ b/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']}") diff --git a/tests/test_saturation.py b/tests/test_saturation.py new file mode 100644 index 0000000..fd07a32 --- /dev/null +++ b/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']}") diff --git a/tests/test_widening.py b/tests/test_widening.py new file mode 100644 index 0000000..c1092c9 --- /dev/null +++ b/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']}") diff --git a/visuals/SignalProcessingSpectrogram.py b/visuals/SignalProcessingSpectrogram.py index 265632d..a000552 100644 --- a/visuals/SignalProcessingSpectrogram.py +++ b/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 diff --git a/visuals/SignalProcessingWaveform.py b/visuals/SignalProcessingWaveform.py index 94248dd..f0f90e1 100644 --- a/visuals/SignalProcessingWaveform.py +++ b/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)