124 lines
4.6 KiB
Python
124 lines
4.6 KiB
Python
"""
|
|
Regression tests for the AudioNormalize and AudioApplyBatchedGain nodes in ComfyUI-AudioBatch.
|
|
"""
|
|
|
|
import torch
|
|
import pytest
|
|
import bootstrap # noqa: F401
|
|
from nodes.nodes_audio import AudioNormalize, AudioApplyBatchedGain
|
|
|
|
|
|
# Helper function
|
|
def create_dummy_audio(batch_size, channels, samples, sr, peak_value=0.5, device='cpu'):
|
|
# Create audio that is not normalized to 1.0 to make testing meaningful
|
|
waveform = torch.randn(batch_size, channels, samples, device=device)
|
|
# Normalize to a known peak value
|
|
current_peak, _ = torch.max(torch.abs(waveform), dim=2, keepdim=True)
|
|
current_peak, _ = torch.max(current_peak, dim=1, keepdim=True)
|
|
waveform = waveform / (current_peak + 1e-9) * peak_value
|
|
return {"waveform": waveform, "sample_rate": sr}
|
|
|
|
|
|
@pytest.fixture
|
|
def normalize_node():
|
|
return AudioNormalize()
|
|
|
|
|
|
@pytest.fixture
|
|
def apply_gain_node():
|
|
return AudioApplyBatchedGain()
|
|
|
|
|
|
# --- Tests for AudioNormalize ---
|
|
|
|
def test_normalize_simple(normalize_node):
|
|
"""Tests if a simple audio clip is normalized to the target peak level."""
|
|
sr = 44100
|
|
original_peak = 0.25
|
|
target_peak = 1.0
|
|
audio = create_dummy_audio(1, 1, sr, sr, peak_value=original_peak)
|
|
|
|
(normalized_audio, original_peak_level) = normalize_node.normalize_audio(audio, target_peak)
|
|
|
|
# Assertions
|
|
new_peak = torch.max(torch.abs(normalized_audio['waveform']))
|
|
assert new_peak.item() == pytest.approx(target_peak, abs=1e-6)
|
|
assert original_peak_level.item() == pytest.approx(original_peak, abs=1e-6)
|
|
|
|
|
|
def test_normalize_with_batch(normalize_node):
|
|
"""Tests that each item in a batch is normalized independently."""
|
|
sr = 44100
|
|
# Create a batch with different peak levels
|
|
audio1 = create_dummy_audio(1, 2, sr, sr, peak_value=0.5)
|
|
audio2 = create_dummy_audio(1, 2, sr, sr, peak_value=0.1)
|
|
|
|
batched_waveform = torch.cat((audio1['waveform'], audio2['waveform']), dim=0)
|
|
batched_audio = {"waveform": batched_waveform, "sample_rate": sr}
|
|
|
|
target_peak = 0.9
|
|
(normalized_audio, original_peak_levels) = normalize_node.normalize_audio(batched_audio, target_peak)
|
|
|
|
# Check item 1
|
|
peak1 = torch.max(torch.abs(normalized_audio['waveform'][0]))
|
|
assert peak1.item() == pytest.approx(target_peak, abs=1e-6)
|
|
assert original_peak_levels[0].item() == pytest.approx(0.5, abs=1e-6)
|
|
|
|
# Check item 2
|
|
peak2 = torch.max(torch.abs(normalized_audio['waveform'][1]))
|
|
assert peak2.item() == pytest.approx(target_peak, abs=1e-6)
|
|
assert original_peak_levels[1].item() == pytest.approx(0.1, abs=1e-6)
|
|
|
|
assert original_peak_levels.shape == (2,)
|
|
|
|
|
|
def test_normalize_silent_audio(normalize_node):
|
|
"""Tests that silent audio remains silent and doesn't cause errors."""
|
|
sr = 44100
|
|
silent_audio = {"waveform": torch.zeros(1, 1, sr), "sample_rate": sr}
|
|
|
|
(normalized_audio, original_peak_level) = normalize_node.normalize_audio(silent_audio, 1.0)
|
|
|
|
# Assert that the output is still silent
|
|
assert torch.all(normalized_audio['waveform'] == 0)
|
|
# The original peak level should be 0
|
|
assert original_peak_level.item() == 0.0
|
|
|
|
|
|
# --- Tests for AudioApplyBatchedGain ---
|
|
|
|
def test_apply_gain_and_revert_normalization(normalize_node, apply_gain_node):
|
|
"""
|
|
An integration test to verify that applying the `original_peak_level`
|
|
from the normalize node correctly reverts the audio to its original state.
|
|
"""
|
|
sr = 44100
|
|
audio_orig = create_dummy_audio(3, 2, sr, sr, peak_value=0.3)
|
|
|
|
# 1. Normalize the audio
|
|
(normalized_audio, original_peak_levels) = normalize_node.normalize_audio(audio_orig, 1.0)
|
|
|
|
# 2. Apply the original peak levels as gain to the normalized audio
|
|
(reverted_audio,) = apply_gain_node.apply_gain(normalized_audio, original_peak_levels)
|
|
|
|
# Assertions
|
|
# The reverted audio should be almost identical to the original audio
|
|
assert reverted_audio['sample_rate'] == audio_orig['sample_rate']
|
|
assert reverted_audio['waveform'].shape == audio_orig['waveform'].shape
|
|
# Check that the content is the same (with a small tolerance for float math)
|
|
assert torch.allclose(reverted_audio['waveform'], audio_orig['waveform'], atol=1e-6)
|
|
|
|
|
|
def test_apply_gain_mismatched_batch_size_error(apply_gain_node):
|
|
"""
|
|
Tests that an error is raised if the audio batch size and gain values
|
|
batch size do not match.
|
|
"""
|
|
sr = 44100
|
|
audio = create_dummy_audio(3, 1, sr, sr)
|
|
# Gain values for only 2 items, but audio has 3
|
|
gain_values = torch.tensor([0.5, 0.5])
|
|
|
|
with pytest.raises(ValueError, match="Batch size of audio .* and gain_values .* must match"):
|
|
apply_gain_node.apply_gain(audio, gain_values)
|