diff --git a/README.md b/README.md index faa656d..5dd6990 100644 --- a/README.md +++ b/README.md @@ -16,6 +16,7 @@ workflows, especially when dealing with multiple audio inputs or outputs. - [7. Audio Information](#7-audio-information) - [8. Audio Cut](#8-audio-cut) - [9. Audio Blend](#9-audio-blend) + - [10. Audio Test Signal Generator](#10-audio-test-signal-generator) - [🚀 Installation](#-installation) - [📦 Dependencies](#-dependencies) - [🖼️ Examples](#️-examples) @@ -158,6 +159,30 @@ workflows, especially when dealing with multiple audio inputs or outputs. - **Output:** - `audio_out` (AUDIO): The blended audio. +### 10. Audio Test Signal Generator + - **Display Name:** `Audio Test Signal Generator` + - **Internal Name:** `SET_AudioTestSignalGenerator` + - **Category:** `audio/generation` + - **Description:** Generates various standard audio test signals. This node is perfect for creating predictable, clean audio to test or debug other audio processing nodes in your workflow without needing external files. + - **Inputs:** + - `waveform_type` (COMBO): The type of signal to generate. + - **Periodic:** `sine`, `square`, `sawtooth`, `triangle`. + - **Sweep:** `sweep` (a linear frequency chirp from `frequency` to `frequency_end`). + - **Noise:** `white_noise`, `pink_noise`, `brownian_noise`. + - **Special:** `impulse` (a single sample at full amplitude), `silence` (a signal of all zeros, to which a DC offset can be added). + - `frequency` (FLOAT): The fundamental frequency in Hz for periodic waveforms, or the starting frequency for a sweep. + - `frequency_end` (FLOAT): The ending frequency in Hz, used only for the `sweep` waveform type. + - `amplitude` (FLOAT): The peak amplitude of the signal, from 0.0 to 1.0. This is ignored for `silence`. + - `dc_offset` (FLOAT): A constant value to add to the entire signal, shifting it up or down. For the `silence` type, this directly sets the signal's value. + - `phase` (FLOAT): A phase offset in degrees (0-360) for periodic waveforms. + - `duration` (STRING): The total duration of the audio. Supports seconds (`"5.5"`) or time format (`"HH:MM:SS.ss"`). + - `sample_rate` (INT): The sample rate for the output audio (e.g., 44100). + - `batch_size` (INT): The number of identical audio clips to generate in the output batch. + - `channels` (INT): The number of audio channels (1 for mono, 2 for stereo, etc.). + - `seed` (INT, optional): A seed for random number generation, ensuring that noise waveforms are reproducible. + - **Output:** + - `audio_out` (AUDIO): The generated test signal. + ## 🚀 Installation You can install the nodes from the ComfyUI nodes manager, the name is *Audio Batch*, or just do it manually: diff --git a/source/nodes/nodes_audio.py b/source/nodes/nodes_audio.py index 1ec67ec..6fddbec 100644 --- a/source/nodes/nodes_audio.py +++ b/source/nodes/nodes_audio.py @@ -5,6 +5,7 @@ # Project: ComfyUI-AudioBatch # From code generated by Gemini 2.5 Pro import torch +import torchaudio import torchaudio.transforms as T from .utils.aligner import AudioBatchAligner from .utils.logger import main_logger @@ -15,6 +16,7 @@ BASE_CATEGORY = "audio" BATCH_CATEGORY = "batch" CONV_CATEGORY = "conversion" MANIPULATION_CATEGORY = "manipulation" +GEN_CATEGORY = "generation" class AudioBatch: @@ -457,3 +459,163 @@ class AudioBlend: output_audio = {"waveform": blended_waveform, "sample_rate": target_sr} return (output_audio,) + + +class AudioTestSignalGenerator: + WAVEFORM_TYPES = ["sine", "square", "sawtooth", "triangle", "sweep", + "white_noise", "pink_noise", "brownian_noise", + "impulse", "silence"] + + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "waveform_type": (cls.WAVEFORM_TYPES, {"default": "sine"}), + "frequency": ("FLOAT", {"default": 440.0, "min": 0.0, "max": 22050.0, "step": 1.0, + "tooltip": "Frequency in Hz for periodic waveforms (or start frequency for sweep)."}), + "frequency_end": ("FLOAT", {"default": 880.0, "min": 0.0, "max": 22050.0, "step": 1.0, + "tooltip": "End frequency in Hz for 'sweep' waveform type."}), + "amplitude": ("FLOAT", {"default": 0.5, "min": 0.0, "max": 1.0, "step": 0.01, + "tooltip": "Amplitude of the signal (0.0 to 1.0)."}), + "dc_offset": ("FLOAT", {"default": 0.0, "min": -1.0, "max": 1.0, "step": 0.01, + "tooltip": "DC offset to add to the signal."}), + "phase": ("FLOAT", {"default": 0.0, "min": 0.0, "max": 360.0, "step": 1.0, + "tooltip": "Phase offset in degrees (0-360)."}), + # "duration_seconds": ("FLOAT", {"default": 3.0, "min": 0.01, "max": 3600.0, "step": 0.1}), + "duration": ("STRING", { + "default": "3.0", + "tooltip": "Duration time in HH:MM:SS.ss format, or just a float for seconds." + }), + "sample_rate": ("INT", {"default": 44100, "min": 1, "max": 192000, "step": 100}), + "batch_size": ("INT", {"default": 1, "min": 1, "max": 64, "step": 1}), + "channels": ("INT", {"default": 1, "min": 1, "max": 8, "step": 1, + "tooltip": "Number of audio channels (1=mono, 2=stereo, etc.)."}), + }, + "optional": { # To allow for deterministic noise generation + "seed": ("INT", {"default": 0, "min": 0, "max": 0xffffffffffffffff}), + } + } + + RETURN_TYPES = ("AUDIO",) + RETURN_NAMES = ("audio_out",) + FUNCTION = "generate_signal" + CATEGORY = BASE_CATEGORY + "/" + GEN_CATEGORY + DESCRIPTION = "Generates various audio test signals like sine waves, noise, and sweeps." + UNIQUE_NAME = "SET_AudioTestSignalGenerator" + DISPLAY_NAME = "Audio Test Signal Generator" + + def generate_signal(self, waveform_type: str, frequency: float, frequency_end: float, + amplitude: float, dc_offset: float, phase: float, + duration: str, sample_rate: int, + batch_size: int, channels: int, seed: int = 0): + + duration_seconds = parse_time_to_seconds(duration) + num_samples = int(duration_seconds * sample_rate) + if num_samples <= 0: + logger.warning("Duration and sample rate result in 0 or negative samples. Outputting empty audio.") + waveform = torch.empty((batch_size, channels, 0), dtype=torch.float32) + return ({"waveform": waveform, "sample_rate": sample_rate},) + + # Create the time vector `t` for one channel + t = torch.linspace(0., duration_seconds, num_samples) + + # Convert phase from degrees to radians + phase_rad = torch.deg2rad(torch.tensor(phase)) + + waveform_1d = None + waveform_multichannel = None + + if waveform_type == "sine": + waveform_1d = torch.sin(2 * torch.pi * frequency * t + phase_rad) + + elif waveform_type == "square": + # A square wave is the sign of a sine wave + waveform_1d = torch.sign(torch.sin(2 * torch.pi * frequency * t + phase_rad)) + # Avoid pure 0s for a cleaner square wave if sine hits exactly 0 + waveform_1d[waveform_1d == 0] = 1.0 + + elif waveform_type == "sawtooth": + # A sawtooth wave is related to the modulo of the time component + waveform_1d = 2 * (t * frequency - torch.floor(0.5 + t * frequency)) + + elif waveform_type == "triangle": + # A triangle wave can be derived from the absolute value of a sawtooth + sawtooth = 2 * (t * frequency - torch.floor(0.5 + t * frequency)) + waveform_1d = 2 * torch.abs(sawtooth) - 1 + + elif waveform_type == "sweep": + # Linear frequency sweep (chirp) from `frequency` to `frequency_end` + logger.debug(f"Generating linear sweep from {frequency}Hz to {frequency_end}Hz.") + # Instantaneous frequency: f(t) = f_start + (f_end - f_start) * (t / duration) + # Phase is the integral of instantaneous frequency + f_start = frequency + f_end = frequency_end + k = (f_end - f_start) / duration_seconds # Sweep rate + # Phase = 2 * pi * (f_start * t + (k/2) * t^2) + phase_sweep = 2 * torch.pi * (f_start * t + (k / 2.0) * (t ** 2)) + waveform_1d = torch.sin(phase_sweep + phase_rad) + + elif "noise" in waveform_type: + generator = torch.Generator().manual_seed(seed) + noise_shape = (batch_size, channels, num_samples) + + # Use torchaudio's colored_noise if available for pink/brownian + if waveform_type != "white_noise" and hasattr(torchaudio.functional, 'colored_noise'): + try: + exponent = -1 if waveform_type == "pink_noise" else -2 + # Reshape for torchaudio function and then back + flat_noise = torchaudio.functional.colored_noise( + size=(batch_size * channels, num_samples), + exponent=exponent, + generator=generator + ) + waveform_multichannel = flat_noise.reshape(noise_shape) + except Exception as e: + logger.warning(f"torchaudio.functional.colored_noise failed ({e}). Falling back to simpler methods.") + waveform_multichannel = None # Signal fallback + + if waveform_multichannel is None: # Fallback or white_noise + white_noise = torch.rand(noise_shape, generator=generator) * 2 - 1 + if waveform_type == "brownian_noise": + brownian_unscaled = torch.cumsum(white_noise, dim=2) + max_abs, _ = torch.max(torch.abs(brownian_unscaled), dim=2, keepdim=True) + waveform_multichannel = brownian_unscaled / (max_abs + 1e-9) + else: # For white_noise or as a fallback for pink_noise + if waveform_type == "pink_noise": + logger.warning("Using white noise as a fallback for pink noise generation.") + waveform_multichannel = white_noise + + elif waveform_type == "impulse": + waveform_1d = torch.zeros(num_samples) + if num_samples > 0: + waveform_1d[0] = 1.0 # Single sample impulse at the beginning + + elif waveform_type == "silence": + # The amplitude and dc_offset will be applied correctly later. + waveform_1d = torch.zeros(num_samples) + + else: + logger.error(f"Unknown waveform type: {waveform_type}. Defaulting to silence.") + waveform_1d = torch.zeros(num_samples) + + # Apply amplitude and DC offset to periodic/single-channel waveforms + if waveform_1d is not None: + waveform_1d = waveform_1d * amplitude + dc_offset + # Expand to the correct number of channels and batch size + # (1, num_samples) -> (1, 1, num_samples) -> (batch_size, channels, num_samples) + waveform_multichannel = waveform_1d.unsqueeze(0).unsqueeze(0).repeat(batch_size, channels, 1) + elif 'noise' not in waveform_type: # if waveform_multichannel was not created in noise block + logger.error("Waveform generation failed. Outputting silence.") + waveform_multichannel = torch.zeros((batch_size, channels, num_samples)) + + # Final clip to ensure [-1, 1] range, as some operations might exceed it slightly + final_waveform = torch.clamp(waveform_multichannel, -1.0, 1.0) + + logger.info(f"Generated '{waveform_type}' signal. Shape: {final_waveform.shape}, SR: {sample_rate}Hz") + + output_audio = { + "waveform": final_waveform, + "sample_rate": sample_rate + } + + return (output_audio,) diff --git a/source/tests/test_audio_test_signal_generator.py b/source/tests/test_audio_test_signal_generator.py new file mode 100644 index 0000000..2be96b1 --- /dev/null +++ b/source/tests/test_audio_test_signal_generator.py @@ -0,0 +1,315 @@ +""" +Functional and regression tests for the AudioTestSignalGenerator node. +""" + +import bootstrap # noqa: F401 +import logging +import pytest +import torch +import torchaudio.transforms as T + +# Import the class and utility function to be tested +from nodes.nodes_audio import AudioTestSignalGenerator, logger as node_logger +# from utils.misc import parse_time_to_seconds + +# Configure logging for tests to see output from the node +node_logger.setLevel(logging.DEBUG) + + +# --- Pytest Fixtures (Reusable Setup) --- + +@pytest.fixture +def signal_gen_node(): + """Provides an instance of the AudioTestSignalGenerator node for tests.""" + return AudioTestSignalGenerator() + + +# --- Helper Functions for Strict Testing --- + +def get_peak_frequency(waveform: torch.Tensor, sample_rate: int) -> float: + """Helper to find the dominant frequency in a waveform using FFT.""" + if waveform.ndim > 1: + # Use the first channel and first batch item for analysis + waveform = waveform.squeeze() + if waveform.ndim > 1: + waveform = waveform[0] + + # Perform Real Fast Fourier Transform + fft_result = torch.fft.rfft(waveform) + # Get frequency bins for the FFT result + freq_bins = torch.fft.rfftfreq(n=waveform.size(-1), d=1.0 / sample_rate) + # Find the index of the maximum magnitude in the FFT + peak_index = torch.argmax(torch.abs(fft_result)) + # Get the frequency corresponding to that peak index + peak_freq = freq_bins[peak_index].item() + return peak_freq + + +# --- Test Cases --- + +def test_generator_basic_sine_shape_and_type(signal_gen_node): + """Tests basic generation of a sine wave, checking output structure.""" + (result_audio,) = signal_gen_node.generate_signal( + waveform_type="sine", frequency=440.0, frequency_end=880.0, + amplitude=0.8, dc_offset=0.0, phase=0.0, duration="2.0", + sample_rate=44100, batch_size=2, channels=2, seed=0 + ) + + assert isinstance(result_audio, dict) + assert "waveform" in result_audio + assert "sample_rate" in result_audio + + waveform = result_audio['waveform'] + sr = result_audio['sample_rate'] + + assert isinstance(waveform, torch.Tensor) + assert sr == 44100 + expected_samples = int(2.0 * sr) + assert waveform.shape == (2, 2, expected_samples) # B, C, N + assert waveform.dtype == torch.float32 + + +@pytest.mark.parametrize("duration_str, expected_seconds", [ + ("5", 5.0), + ("2.5", 2.5), + ("0:10", 10.0), + ("0:01.5", 1.5), + ("1:02:03", 3723.0) +]) +def test_generator_duration_parsing(signal_gen_node, duration_str, expected_seconds): + """Tests if the duration string is parsed correctly to produce the right length.""" + sr = 16000 + (result_audio,) = signal_gen_node.generate_signal( + waveform_type="sine", frequency=100.0, frequency_end=200.0, + amplitude=1.0, dc_offset=0.0, phase=0.0, duration=duration_str, + sample_rate=sr, batch_size=1, channels=1, seed=0 + ) + + expected_samples = int(expected_seconds * sr) + actual_samples = result_audio['waveform'].shape[2] + assert actual_samples == expected_samples + + +def test_generator_fft_verification_sine(signal_gen_node): + """ + Strict Test: Verifies that a generated sine wave has the correct peak frequency. + """ + sr = 44100 + target_freq = 1000.0 + + (result_audio,) = signal_gen_node.generate_signal( + waveform_type="sine", frequency=target_freq, frequency_end=0.0, + amplitude=0.9, dc_offset=0.0, phase=0.0, duration="1.0", + sample_rate=sr, batch_size=1, channels=1, seed=0 + ) + + peak_freq = get_peak_frequency(result_audio['waveform'], result_audio['sample_rate']) + + print(f"Target Frequency: {target_freq} Hz, Detected Peak Frequency: {peak_freq:.2f} Hz") + # Assert that the detected frequency is very close to the target. + # The tolerance accounts for the resolution of the FFT (bin width). + assert peak_freq == pytest.approx(target_freq, abs=1.0) + + +def test_generator_amplitude_and_offset(signal_gen_node): + """ + Strict Test: Verifies that amplitude and DC offset are correctly applied. + """ + amplitude = 0.7 + offset = -0.1 + + (result_audio,) = signal_gen_node.generate_signal( + waveform_type="sine", frequency=100.0, frequency_end=0.0, + amplitude=amplitude, dc_offset=offset, phase=90.0, # Phase 90 ensures we hit the max value at t=0 + duration="1.0", sample_rate=44100, batch_size=1, channels=1, seed=0 + ) + + waveform = result_audio['waveform'] + + expected_max = amplitude + offset + expected_min = -amplitude + offset + + # Check that the max and min values of the waveform are close to expected + assert torch.max(waveform).item() == pytest.approx(expected_max, abs=1e-5) + assert torch.min(waveform).item() == pytest.approx(expected_min, abs=1e-5) + + +def test_generator_noise_reproducibility_with_seed(signal_gen_node): + """ + Strict Test: Verifies that using the same seed produces identical noise, + and different seeds produce different noise. + """ + params = { + "waveform_type": "white_noise", "frequency": 0.0, "frequency_end": 0.0, + "amplitude": 1.0, "dc_offset": 0.0, "phase": 0.0, "duration": "0.5", + "sample_rate": 16000, "batch_size": 1, "channels": 1 + } + + # Generate twice with the same seed + (result1,) = signal_gen_node.generate_signal(**params, seed=123) + (result2,) = signal_gen_node.generate_signal(**params, seed=123) + + # Generate once with a different seed + (result3,) = signal_gen_node.generate_signal(**params, seed=456) + + # Assertions + assert torch.equal(result1['waveform'], result2['waveform']) + assert not torch.equal(result1['waveform'], result3['waveform']) + + +def test_generator_impulse_content(signal_gen_node): + """ + Strict Test: Verifies that an impulse waveform is correct. + """ + amplitude = 0.9 + offset = 0.05 + + (result_audio,) = signal_gen_node.generate_signal( + waveform_type="impulse", frequency=0.0, frequency_end=0.0, + amplitude=amplitude, dc_offset=offset, phase=0.0, duration="0.1", + sample_rate=100, batch_size=1, channels=1, seed=0 + ) + + waveform = result_audio['waveform'].squeeze() # Get 1D tensor + + # The first sample should be the impulse + expected_impulse_value = amplitude + offset + assert waveform[0].item() == pytest.approx(expected_impulse_value) + + # All other samples should be just the DC offset + assert torch.all(torch.isclose(waveform[1:], torch.tensor(offset))) + + +def linear_regression(x, y): + """ + Performs a simple linear regression and returns the slope and intercept. + Uses torch for calculations. + """ + x_mean = torch.mean(x) + y_mean = torch.mean(y) + numerator = torch.sum((x - x_mean) * (y - y_mean)) + denominator = torch.sum((x - x_mean) ** 2) + if denominator == 0: + return 0.0, y_mean.item() # No slope if x is constant + slope = numerator / denominator + intercept = y_mean - slope * x_mean + return slope.item(), intercept.item() + + +def test_generator_sweep_frequency_trend(signal_gen_node): + """ + Strict Test (Corrected Again): Verifies a linear sweep has a constant rate of + frequency change using STFT and linear regression. + """ + sr = 44100 + f_start = 500.0 + f_end = 8000.0 + duration_seconds = 2.0 + + (result_audio,) = signal_gen_node.generate_signal( + waveform_type="sweep", frequency=f_start, frequency_end=f_end, + amplitude=1.0, dc_offset=0.0, phase=0.0, duration=str(duration_seconds), + sample_rate=sr, batch_size=1, channels=1, seed=0 + ) + + waveform = result_audio['waveform'].squeeze() + + # --- Use STFT to analyze the frequency over time --- + n_fft = 2048 + hop_length = 512 + + spectrogram_transform = T.Spectrogram(n_fft=n_fft, hop_length=hop_length, power=2.0) + spectrogram = spectrogram_transform(waveform) # (freq_bins, time_frames) + + # Get the frequency and time axes for the spectrogram + fft_freqs = torch.fft.rfftfreq(n=n_fft, d=1.0 / sr) + time_bins = torch.linspace(0, duration_seconds, steps=spectrogram.shape[1]) + + # For each time slice, find the index of the peak frequency + peak_freq_indices = torch.argmax(torch.abs(spectrogram), dim=0) + + # Convert these indices to actual frequencies + detected_peak_freqs = fft_freqs[peak_freq_indices] + + # We might get some noise at the very beginning/end, so let's analyze the stable middle part + start_idx_for_fit = 10 + end_idx_for_fit = -10 + + if len(time_bins) <= start_idx_for_fit * 2: # Handle very short signals + start_idx_for_fit = 0 + end_idx_for_fit = len(time_bins) + + time_for_fit = time_bins[start_idx_for_fit:end_idx_for_fit] + freqs_for_fit = detected_peak_freqs[start_idx_for_fit:end_idx_for_fit] + + # --- Perform Linear Regression --- + # We are fitting a line: frequency = slope * time + intercept + detected_slope, detected_intercept = linear_regression(time_for_fit, freqs_for_fit) + + # --- Calculate Theoretical Values --- + theoretical_slope = (f_end - f_start) / duration_seconds + theoretical_intercept = f_start + + logging.info("\n--- Sweep Linear Regression Test ---") + logging.info(f"Theoretical Slope: {theoretical_slope:.2f} Hz/s") + logging.info(f"Detected Slope: {detected_slope:.2f} Hz/s") + logging.info(f"Theoretical Intercept (Start Freq): {theoretical_intercept:.2f} Hz") + logging.info(f"Detected Intercept: {detected_intercept:.2f} Hz") + + # Assertions: + # 1. The detected slope should be very close to the theoretical rate of change. + assert detected_slope == pytest.approx(theoretical_slope, rel=0.05) # 5% tolerance is reasonable + + # 2. The detected intercept should be close to the starting frequency. + # This is still subject to some error from the first few windows, so give it some tolerance. + assert detected_intercept == pytest.approx(theoretical_intercept, rel=0.1) # 10% tolerance + + +def test_generator_silence_and_dc_offset(signal_gen_node): + """ + Strict Test: Verifies the 'silence' waveform type and that it correctly + applies the dc_offset. + """ + sample_rate = 1000 # Use a low SR for a small tensor + duration = "0.5" # 0.5 seconds + num_samples = int(0.5 * sample_rate) + target_offset = 0.25 + + # Generate a "silence" waveform with a non-zero DC offset. + # Amplitude should have no effect. + (result_audio,) = signal_gen_node.generate_signal( + waveform_type="silence", + frequency=100.0, # Should be ignored + frequency_end=800.0, # Should be ignored + amplitude=0.8, # Should be ignored + dc_offset=target_offset, + phase=0.0, + duration=duration, + sample_rate=sample_rate, + batch_size=2, + channels=2, + seed=0 + ) + + waveform = result_audio['waveform'] + + # 1. Check shape and SR + assert waveform.shape == (2, 2, num_samples) + assert result_audio['sample_rate'] == sample_rate + + # 2. Check content + # The entire tensor should be filled with the dc_offset value. + expected_waveform = torch.full_like(waveform, fill_value=target_offset) + + assert torch.allclose(waveform, expected_waveform) + + # 3. Test pure silence (dc_offset = 0) + (result_pure_silence,) = signal_gen_node.generate_signal( + waveform_type="silence", + dc_offset=0.0, + # ... other params ... + frequency=100.0, frequency_end=800.0, amplitude=0.8, phase=0.0, duration=duration, + sample_rate=sample_rate, batch_size=1, channels=1, seed=0 + ) + + assert torch.all(result_pure_silence['waveform'] == 0.0)