diff --git a/comfy_api_nodes/util/conversions.py b/comfy_api_nodes/util/conversions.py index 9cd644fc0..f46cac3f8 100644 --- a/comfy_api_nodes/util/conversions.py +++ b/comfy_api_nodes/util/conversions.py @@ -266,16 +266,14 @@ def audio_tensor_to_contiguous_ndarray(waveform: torch.Tensor) -> np.ndarray: waveform: a tensor of shape (1, channels, samples) derived from a Comfy `AUDIO` type. Returns: - Contiguous numpy array of the audio waveform. If the audio was batched, - the first item is taken. + Contiguous numpy array of the audio waveform. + + Raises: + ValueError: If the waveform is not shaped (1, channels, samples). """ if waveform.ndim != 3 or waveform.shape[0] != 1: raise ValueError("Expected waveform tensor shape (1, channels, samples)") - # If batch is > 1, take first item - if waveform.shape[0] > 1: - waveform = waveform[0] - # Prepare for av: remove batch dim, move to CPU, make contiguous, convert to numpy array audio_data_np = waveform.squeeze(0).cpu().contiguous().numpy() if audio_data_np.dtype != np.float32: @@ -285,20 +283,21 @@ def audio_tensor_to_contiguous_ndarray(waveform: torch.Tensor) -> np.ndarray: def audio_input_to_mp3(audio: Input.Audio) -> BytesIO: - waveform = audio["waveform"].cpu() + audio_data_np = audio_tensor_to_contiguous_ndarray(audio["waveform"]) + sample_rate = int(audio["sample_rate"]) output_buffer = BytesIO() output_container = av.open(output_buffer, mode="w", format="mp3") - out_stream = output_container.add_stream("libmp3lame", rate=audio["sample_rate"]) + out_stream = output_container.add_stream("libmp3lame", rate=sample_rate) out_stream.bit_rate = 320000 frame = av.AudioFrame.from_ndarray( - waveform.movedim(0, 1).reshape(1, -1).float().numpy(), - format="flt", - layout="mono" if waveform.shape[0] == 1 else "stereo", + audio_data_np, + format="fltp", + layout="stereo" if audio_data_np.shape[0] > 1 else "mono", ) - frame.sample_rate = audio["sample_rate"] + frame.sample_rate = sample_rate frame.pts = 0 output_container.mux(out_stream.encode(frame)) output_container.mux(out_stream.encode(None)) diff --git a/tests-unit/comfy_api_nodes_test/audio_conversions_test.py b/tests-unit/comfy_api_nodes_test/audio_conversions_test.py new file mode 100644 index 000000000..d7c9b3899 --- /dev/null +++ b/tests-unit/comfy_api_nodes_test/audio_conversions_test.py @@ -0,0 +1,78 @@ +import math + +import av +import numpy as np +import pytest +import torch + +from comfy.cli_args import args + +if not torch.cuda.is_available(): + args.cpu = True + +from comfy_api_nodes.util.conversions import audio_input_to_mp3 # noqa: E402 + +SAMPLE_RATE = 48000 +DURATION = 2.0 +LEFT_HZ = 440.0 +RIGHT_HZ = 880.0 + + +def tone(freq, duration=DURATION, sample_rate=SAMPLE_RATE): + t = torch.arange(int(sample_rate * duration), dtype=torch.float32) / sample_rate + return 0.5 * torch.sin(2 * math.pi * freq * t) + + +@pytest.fixture +def stereo_audio(): + """Comfy AUDIO with two tones that stay distinguishable through mp3.""" + waveform = torch.stack([tone(LEFT_HZ), tone(RIGHT_HZ)]).unsqueeze(0) + return {"waveform": waveform, "sample_rate": SAMPLE_RATE} + + +@pytest.fixture +def mono_audio(): + return {"waveform": tone(LEFT_HZ).unsqueeze(0).unsqueeze(0), "sample_rate": SAMPLE_RATE} + + +def decode(buffer): + """(planes[C, N], sample_rate, channels) of an encoded mp3 buffer""" + buffer.seek(0) + with av.open(buffer, mode="r") as container: + stream = container.streams.audio[0] + planes = [] + for frame in container.decode(audio=0): + array = frame.to_ndarray() + if frame.format.is_planar: + planes.append(array) + else: + planes.append(array.reshape(-1, len(frame.layout.channels)).T) + return np.concatenate(planes, axis=1), stream.codec_context.sample_rate, len(stream.layout.channels) + + +def dominant_hz(signal, sample_rate): + """Peak frequency, ignoring the encoder's padding at either edge""" + edge = int(0.2 * sample_rate) + window = signal[edge:-edge] + spectrum = np.abs(np.fft.rfft(window * np.hanning(window.size))) + return np.fft.rfftfreq(window.size, 1.0 / sample_rate)[int(np.argmax(spectrum))] + + +def test_stereo_duration_is_preserved(stereo_audio): + planes, sample_rate, channels = decode(audio_input_to_mp3(stereo_audio)) + assert channels == 2 + assert sample_rate == SAMPLE_RATE + assert planes.shape[1] / sample_rate == pytest.approx(DURATION, abs=0.15) + + +def test_stereo_channels_are_not_concatenated(stereo_audio): + """The channels must be interleaved; concatenating them plays the clip twice.""" + planes, sample_rate, _ = decode(audio_input_to_mp3(stereo_audio)) + assert dominant_hz(planes[0], sample_rate) == pytest.approx(LEFT_HZ, abs=15) + assert dominant_hz(planes[1], sample_rate) == pytest.approx(RIGHT_HZ, abs=15) + + +def test_mono_duration_is_preserved(mono_audio): + planes, sample_rate, _ = decode(audio_input_to_mp3(mono_audio)) + assert planes.shape[1] / sample_rate == pytest.approx(DURATION, abs=0.15) + assert dominant_hz(planes[0], sample_rate) == pytest.approx(LEFT_HZ, abs=15)