79 lines
2.7 KiB
Python
79 lines
2.7 KiB
Python
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)
|