[Partner Nodes] fix(ByteDance): encode stereo reference audio without doubling its duration (#15177)

* [Partner Nodes] fix(ByteDance): encode stereo reference audio without doubling its duration

Signed-off-by: bigcat88 <bigcat88@icloud.com>

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))

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)