127 lines
4.4 KiB
Python
127 lines
4.4 KiB
Python
"""
|
|
Regression tests for the AudioJoin2Channels and AudioSplit2Channels nodes in ComfyUI-AudioBatch.
|
|
"""
|
|
|
|
import bootstrap # noqa: F401
|
|
import torch
|
|
import pytest
|
|
from nodes.nodes_audio import AudioJoin2Channels, AudioSplit2Channels
|
|
|
|
|
|
# Helper function
|
|
def create_dummy_audio(batch_size, channels, samples, sr, device='cpu'):
|
|
waveform = torch.randn(batch_size, channels, samples, device=device, dtype=torch.float32)
|
|
return {"waveform": waveform, "sample_rate": sr}
|
|
|
|
|
|
@pytest.fixture
|
|
def join_node():
|
|
return AudioJoin2Channels()
|
|
|
|
|
|
@pytest.fixture
|
|
def split_node():
|
|
return AudioSplit2Channels()
|
|
|
|
|
|
# --- Tests for AudioJoin2Channels ---
|
|
|
|
def test_join_simple(join_node):
|
|
"""Tests joining two perfectly matched mono signals."""
|
|
sr = 44100
|
|
samples = 1000
|
|
left_audio = create_dummy_audio(1, 1, samples, sr)
|
|
right_audio = create_dummy_audio(1, 1, samples, sr)
|
|
|
|
(stereo_audio,) = join_node.join_channels(left_audio, right_audio)
|
|
|
|
assert stereo_audio['waveform'].shape == (1, 2, samples)
|
|
# Check if L/R channels match the original mono inputs
|
|
assert torch.allclose(stereo_audio['waveform'][:, 0:1, :], left_audio['waveform'])
|
|
assert torch.allclose(stereo_audio['waveform'][:, 1:2, :], right_audio['waveform'])
|
|
|
|
|
|
def test_join_converts_inputs_to_mono(join_node):
|
|
"""Tests that stereo inputs are correctly converted to mono before joining."""
|
|
sr = 44100
|
|
samples = 1000
|
|
# Left input is stereo, Right is mono
|
|
left_stereo_in = create_dummy_audio(1, 2, samples, sr)
|
|
right_mono_in = create_dummy_audio(1, 1, samples, sr)
|
|
|
|
(result_audio,) = join_node.join_channels(left_stereo_in, right_mono_in)
|
|
|
|
# Left channel of the output should be the average of the stereo input
|
|
expected_left_mono = torch.mean(left_stereo_in['waveform'], dim=1, keepdim=True)
|
|
|
|
assert result_audio['waveform'].shape == (1, 2, samples)
|
|
assert torch.allclose(result_audio['waveform'][:, 0:1, :], expected_left_mono)
|
|
assert torch.allclose(result_audio['waveform'][:, 1:2, :], right_mono_in['waveform'])
|
|
|
|
|
|
def test_join_aligns_and_batches(join_node):
|
|
"""Tests that mismatched SR, length, and batch sizes are handled."""
|
|
# Left: B=2, 1ch, 44.1k SR, 1s long
|
|
audio_left = create_dummy_audio(2, 1, 44100, 44100)
|
|
# Right: B=3, 2ch, 22.05k SR, 0.5s long
|
|
audio_right = create_dummy_audio(3, 2, 11025, 22050)
|
|
|
|
(result_audio,) = join_node.join_channels(audio_left, audio_right)
|
|
|
|
# Expected output shape:
|
|
# Batch size = max(2, 3) = 3
|
|
# Channels = 2 (stereo)
|
|
# SR = 44100 (from left)
|
|
# Samples = 44100 (from left, since right is 11025*2=22050 after resampling)
|
|
assert result_audio['waveform'].shape == (3, 2, 44100)
|
|
assert result_audio['sample_rate'] == 44100
|
|
|
|
# Check the repeated last item for the left channel
|
|
assert torch.allclose(result_audio['waveform'][1, 0, :], result_audio['waveform'][2, 0, :])
|
|
|
|
|
|
# --- Tests for AudioSplit2Channels ---
|
|
|
|
def test_split_simple(split_node):
|
|
"""Tests splitting a standard stereo signal."""
|
|
sr = 44100
|
|
samples = 1000
|
|
stereo_audio = create_dummy_audio(1, 2, samples, sr)
|
|
|
|
(left_out, right_out) = split_node.split_channels(stereo_audio)
|
|
|
|
# Check left channel output
|
|
assert left_out['sample_rate'] == sr
|
|
assert left_out['waveform'].shape == (1, 1, samples)
|
|
assert torch.allclose(left_out['waveform'], stereo_audio['waveform'][:, 0:1, :])
|
|
|
|
# Check right channel output
|
|
assert right_out['sample_rate'] == sr
|
|
assert right_out['waveform'].shape == (1, 1, samples)
|
|
assert torch.allclose(right_out['waveform'], stereo_audio['waveform'][:, 1:2, :])
|
|
|
|
|
|
def test_split_with_batch(split_node):
|
|
"""Tests that splitting preserves the batch dimension."""
|
|
sr = 44100
|
|
samples = 1000
|
|
stereo_batch = create_dummy_audio(5, 2, samples, sr)
|
|
|
|
(left_out, right_out) = split_node.split_channels(stereo_batch)
|
|
|
|
assert left_out['waveform'].shape == (5, 1, samples)
|
|
assert right_out['waveform'].shape == (5, 1, samples)
|
|
|
|
|
|
def test_split_raises_error_on_non_stereo(split_node):
|
|
"""Tests that an error is raised if the input is not stereo."""
|
|
# Test with mono
|
|
mono_audio = create_dummy_audio(1, 1, 1000, 44100)
|
|
with pytest.raises(ValueError, match="Input audio must be stereo"):
|
|
split_node.split_channels(mono_audio)
|
|
|
|
# Test with 3 channels
|
|
three_ch_audio = create_dummy_audio(1, 3, 1000, 44100)
|
|
with pytest.raises(ValueError, match="Input audio must be stereo"):
|
|
split_node.split_channels(three_ch_audio)
|