Add AudioConverter to the Blocks API

Completes the audio pipeline: AudioDemuxer -> PacketDecoder -> AudioConverter,
turning RawAudioSamples into normalized float32 AudioSamples, optionally
resampled and remixed to another channel count.

Unlike ColorConverter this block is a stream processor rather than a function
of its input, because resampling is an interpolation filter: swresample holds
the tail of each frame back until the next arrives. So convert() can return
fewer samples than it was given, drain() is needed for the last ones, frames
must be fed in order, and reset() is needed after a seek. None of that is true
when sample_rate is left unset - format conversion and channel remixing are
frame-local - but drain()/reset() are in the documented loop from the start so
that adding sample_rate later can't silently truncate anyone's audio.

We deliberately don't implement SingleStreamDecoder's pre-roll or its
resampling alignment grid. The grid works by *dropping* up to
isr/gcd(isr,osr)-1 input samples, which is only safe behind a pre-roll; without
one it would eat the caller's own samples (up to ~1s for 44100 -> 16001).
Consequence: samples decoded after a seek don't line up bit-for-bit with a
whole-file decode. Decoding from the start does, which the tests pin.

create_swr_context() grows an overload taking channel counts rather than an
AVFrame, and the swr_convert() output-size bound moves into a shared helper.

Author: NicolasHug

examples/decoding/blocks.py

@@ -29,6 +29,9 @@
 we illustrate a few things this enables: overlapping stages on multiple
 threads, accessing raw (YUV) frames, and decoding streams of unknown -
 possibly infinite - length.
+
+Audio works the same way, through ``AudioDemuxer`` and ``AudioConverter``; we
+come back to it at the end.
 """
 
 # %%
@@ -427,3 +430,71 @@ def start_live_stream():
 
 ffmpeg.kill()
 ffmpeg.wait()
+
+# %%
+# Audio
+# -----
+#
+# Audio has the same three stages, and ``PacketDecoder`` is the same block:
+#
+# .. code-block::
+#
+#    AudioDemuxer  ->  PacketDecoder   ->  AudioConverter
+#      Packet         RawAudioSamples       AudioSamples
+#
+# What ``PacketDecoder`` hands out follows the demuxer it was built from, so
+# audio comes out as ``RawAudioSamples``: the codec's own samples, in the codec's
+# own sample type, as a ``[num_channels, num_samples]`` tensor.
+from torchcodec.decoders._blocks import AudioConverter, AudioDemuxer
+
+audio_path = temp_dir / "audio.wav"
+subprocess.run(
+    [
+        "ffmpeg", "-y", "-hide_banner", "-loglevel", "error",
+        "-f", "lavfi", "-i", "sine=frequency=440:sample_rate=44100:duration=5",
+        "-c:a", "pcm_s16le", str(audio_path),
+    ],
+    check=True,
+)
+
+demuxer = AudioDemuxer(audio_path)
+packet_decoder = PacketDecoder(demuxer)
+raw = next(iter(packet_decoder.decode(next(iter(demuxer)))))
+print(f"{raw.sample_format = }, {raw.data.dtype = }, {raw.data.shape = }, "
+      f"{raw.sample_rate = }")
+
+# %%
+# Those are the true source samples - 16-bit integers here, not floats in
+# ``[-1, 1]``. ``AudioConverter`` is what normalizes them, and it can resample
+# and change the channel count on the way.
+#
+# It differs from ``ColorConverter`` in one important way: resampling is an
+# interpolation filter, so the sample it emits at a given instant depends on
+# input samples on *both* sides of it. The converter therefore holds the tail
+# of each frame back until the next one arrives - which is why ``convert()``
+# can return fewer samples than it was given, and why the pipeline ends with
+# ``drain()``. Leave that call out and you lose the end of the stream.
+demuxer = AudioDemuxer(audio_path)
+packet_decoder = PacketDecoder(demuxer)
+audio_converter = AudioConverter(sample_rate=16_000, num_channels=1)
+
+chunks = []
+for packet in demuxer:
+    chunks += [audio_converter.convert(raw) for raw in packet_decoder.decode(packet)]
+chunks += [audio_converter.convert(raw) for raw in packet_decoder.drain()]
+chunks.append(audio_converter.drain())  # don't forget me
+
+samples = torch.cat([chunk.data for chunk in chunks], dim=1)
+print(f"{samples.shape = }, {samples.dtype = }, "
+      f"{chunks[-1].data.shape[1]} samples came out of drain()")
+
+# %%
+# .. warning::
+#
+#    These blocks do no pre-roll. A lossy codec's first frames after a seek are
+#    subtly wrong until it re-primes, and a resampler started mid-stream emits
+#    samples on a grid of its own, so samples decoded after a
+#    ``demuxer.seek()`` do not line up bit-for-bit with the same region of a
+#    whole-file decode. Decoding a margin before your target and discarding it
+#    is up to you. :class:`~torchcodec.decoders.AudioDecoder` does all of this
+#    for you.

src/torchcodec/_core/AudioCommon.cpp

@@ -6,6 +6,8 @@
 
 #include "AudioCommon.h"
 
+#include <vector>
+
 namespace facebook::torchcodec {
 
 torch::headeronly::ScalarType sample_format_dtype(
@@ -35,4 +37,65 @@ torch::headeronly::ScalarType sample_format_dtype(
   return kStableUInt8;
 }
 
+AVSampleFormat planar_sample_format(torch::headeronly::ScalarType dtype) {
+  switch (dtype) {
+    case kStableUInt8:
+      return AV_SAMPLE_FMT_U8P;
+    case kStableInt16:
+      return AV_SAMPLE_FMT_S16P;
+    case kStableInt32:
+      return AV_SAMPLE_FMT_S32P;
+    case kStableInt64:
+      return AV_SAMPLE_FMT_S64P;
+    case kStableFloat32:
+      return AV_SAMPLE_FMT_FLTP;
+    case kStableFloat64:
+      return AV_SAMPLE_FMT_DBLP;
+    default:
+      break;
+  }
+  STD_TORCH_CHECK(
+      false,
+      "Unsupported dtype for audio samples. Expected one of uint8, int16, "
+      "int32, int64, float32 or float64.");
+  return AV_SAMPLE_FMT_NONE;
+}
+
+torch::stable::Tensor swr_convert_to_tensor(
+    const UniqueSwrContext& swr_context,
+    const uint8_t* const* src_planes,
+    int num_src_samples,
+    int num_out_channels,
+    int64_t num_out_samples_bound) {
+  torch::stable::Tensor out = torch::stable::empty(
+      {num_out_channels, num_out_samples_bound}, kStableFloat32);
+  if (num_out_samples_bound == 0) {
+    return out;
+  }
+
+  // A contiguous [num_out_channels, N] float32 tensor is exactly FLTP: one
+  // plane per row.
+  int64_t bytes_per_channel =
+      num_out_samples_bound * av_get_bytes_per_sample(kAudioOutSampleFormat);
+  auto* base = static_cast<uint8_t*>(out.mutable_data_ptr());
+  std::vector<uint8_t*> out_planes(num_out_channels);
+  for (int channel = 0; channel < num_out_channels; ++channel) {
+    out_planes[channel] = base + channel * bytes_per_channel;
+  }
+
+  int num_out_samples = swr_convert(
+      swr_context.get(),
+      out_planes.data(),
+      static_cast<int>(num_out_samples_bound),
+      src_planes,
+      num_src_samples);
+  STD_TORCH_CHECK(
+      num_out_samples >= 0,
+      "Error in swr_convert: ",
+      get_ffmpeg_error_string_from_error_code(num_out_samples));
+
+  return torch::stable::narrow(
+      out, /*dim=*/1, /*start=*/0, /*length=*/num_out_samples);
+}
+
 } // namespace facebook::torchcodec

src/torchcodec/_core/AudioCommon.h

@@ -15,9 +15,29 @@
 
 namespace facebook::torchcodec {
 
+// Decoded audio always leaves torchcodec as normalized float32 samples,
+// one plane per channel.
+constexpr AVSampleFormat kAudioOutSampleFormat = AV_SAMPLE_FMT_FLTP;
+
 // The dtype that holds `sample_format`'s samples exactly. Planar and packed
 // variants of a format share a sample type, which is why this doesn't care
 // which one it is given.
 torch::headeronly::ScalarType sample_format_dtype(AVSampleFormat sample_format);
 
+// The inverse: the sample format that a contiguous [num_channels, num_samples]
+// tensor of `dtype` already is. One row per channel is precisely what planar
+// means, which is what lets swresample read such a tensor's rows directly.
+AVSampleFormat planar_sample_format(torch::headeronly::ScalarType dtype);
+
+// Runs swr_convert() straight into a fresh float32 [num_channels, N] tensor,
+// narrowed to the number of samples it actually produced - which is at most
+// `num_out_samples_bound` and is often fewer, since swresample holds back the
+// samples that need future input. Pass a null `src_planes` to flush those out.
+torch::stable::Tensor swr_convert_to_tensor(
+    const UniqueSwrContext& swr_context,
+    const uint8_t* const* src_planes,
+    int num_src_samples,
+    int num_out_channels,
+    int64_t num_out_samples_bound);
+
 } // namespace facebook::torchcodec

src/torchcodec/_core/AudioConverter.cpp

@@ -0,0 +1,127 @@
+// Copyright (c) Meta Platforms, Inc. and affiliates.
+// All rights reserved.
+//
+// This source code is licensed under the BSD-style license found in the
+// LICENSE file in the root directory of this source tree.
+
+#include "AudioConverter.h"
+
+#include <vector>
+
+#include "AudioCommon.h"
+
+namespace facebook::torchcodec {
+
+AudioConverter::AudioConverter(
+    std::optional<int> sample_rate,
+    std::optional<int> num_channels)
+    : requested_sample_rate_(sample_rate),
+      requested_num_channels_(num_channels) {
+  STD_TORCH_CHECK(
+      !sample_rate.has_value() || *sample_rate > 0,
+      "sample_rate must be > 0. Got: ",
+      sample_rate.value_or(0));
+  STD_TORCH_CHECK(
+      !num_channels.has_value() || *num_channels > 0,
+      "num_channels must be > 0. Got: ",
+      num_channels.value_or(0));
+}
+
+void AudioConverter::reset() {
+  swr_context_.reset();
+  src_sample_format_ = AV_SAMPLE_FMT_NONE;
+  src_sample_rate_ = 0;
+  src_num_channels_ = 0;
+  out_sample_rate_ = 0;
+  out_num_channels_ = 0;
+}
+
+torch::stable::Tensor AudioConverter::convert(
+    const torch::stable::Tensor& samples,
+    int sample_rate) {
+  STD_TORCH_CHECK(
+      samples.dim() == 2,
+      "Expected a 2D [num_channels, num_samples] tensor, got a ",
+      samples.dim(),
+      "D one.");
+  STD_TORCH_CHECK(samples.is_contiguous(), "The samples must be contiguous.");
+  STD_TORCH_CHECK(sample_rate > 0, "sample_rate must be > 0.");
+
+  AVSampleFormat src_sample_format =
+      planar_sample_format(samples.scalar_type());
+  int num_channels = static_cast<int>(samples.sizes()[0]);
+  int num_samples = static_cast<int>(samples.sizes()[1]);
+  STD_TORCH_CHECK(
+      num_channels > 0, "The samples must have at least 1 channel.");
+
+  if (swr_context_ == nullptr) {
+    src_sample_format_ = src_sample_format;
+    src_sample_rate_ = sample_rate;
+    src_num_channels_ = num_channels;
+    out_sample_rate_ = requested_sample_rate_.value_or(sample_rate);
+    out_num_channels_ = requested_num_channels_.value_or(num_channels);
+    swr_context_.reset(create_swr_context(
+        src_sample_format_,
+        kAudioOutSampleFormat,
+        src_sample_rate_,
+        out_sample_rate_,
+        src_num_channels_,
+        out_num_channels_));
+  } else {
+    // swresample is configured once, from the first samples we see, and its
+    // buffered state is tied to that configuration. Rather than silently
+    // reconfiguring - which would discard whatever it still holds - we make the
+    // caller decide, by reset()ing.
+    STD_TORCH_CHECK(
+        src_sample_format == src_sample_format_ &&
+            sample_rate == src_sample_rate_ &&
+            num_channels == src_num_channels_,
+        "This AudioConverter was set up for ",
+        src_num_channels_,
+        " channels of ",
+        av_get_sample_fmt_name(src_sample_format_),
+        " at ",
+        src_sample_rate_,
+        " Hz, but got ",
+        num_channels,
+        " channels of ",
+        av_get_sample_fmt_name(src_sample_format),
+        " at ",
+        sample_rate,
+        " Hz. Call reset() to convert a different stream.");
+  }
+
+  const auto* base = static_cast<const uint8_t*>(samples.const_data_ptr());
+  int64_t bytes_per_channel =
+      num_samples * av_get_bytes_per_sample(src_sample_format);
+  std::vector<const uint8_t*> src_planes(num_channels);
+  for (int channel = 0; channel < num_channels; ++channel) {
+    src_planes[channel] = base + channel * bytes_per_channel;
+  }
+
+  return swr_convert_to_tensor(
+      swr_context_,
+      src_planes.data(),
+      num_samples,
+      out_num_channels_,
+      get_swr_output_num_samples_bound(
+          swr_context_, num_samples, src_sample_rate_, out_sample_rate_));
+}
+
+torch::stable::Tensor AudioConverter::drain() {
+  STD_TORCH_CHECK(
+      swr_context_ != nullptr,
+      "This AudioConverter hasn't converted any samples, so there is nothing "
+      "to drain and no way to know what shape the result should have.");
+  // A null input is what tells swr_convert() to flush. Unlike the convert()
+  // path we ask swresample how much it is holding rather than deriving a bound
+  // from an input size, since here there is no input.
+  return swr_convert_to_tensor(
+      swr_context_,
+      /*src_planes=*/nullptr,
+      /*num_src_samples=*/0,
+      out_num_channels_,
+      swr_get_out_samples(swr_context_.get(), 0));
+}
+
+} // namespace facebook::torchcodec

src/torchcodec/_core/AudioConverter.h

@@ -0,0 +1,77 @@
+// Copyright (c) Meta Platforms, Inc. and affiliates.
+// All rights reserved.
+//
+// This source code is licensed under the BSD-style license found in the
+// LICENSE file in the root directory of this source tree.
+
+#pragma once
+
+#include <optional>
+
+#include "FFMPEGCommon.h"
+#include "StableABICompat.h"
+
+namespace facebook::torchcodec {
+
+// Audio conversion building block: turns a decoded frame's samples, in the
+// codec's own sample type, into normalized float32 ones - optionally at another
+// sample rate, and with another channel count.
+//
+// Unlike ColorConverter this is a stream processor, not a function of its
+// input, because resampling is an interpolation filter: the sample it emits at
+// a given instant is a weighted sum of input samples on both sides of it. So
+// swresample holds the tail of each frame back until the next one arrives,
+// which means convert() emits fewer samples than it was given (sometimes none),
+// drain() is needed to get the last ones out, and frames must be fed in order.
+// reset() drops that state, and is what a caller must do after seeking.
+//
+// Not thread-safe.
+class FORCE_PUBLIC_VISIBILITY AudioConverter {
+ public:
+  // Both default to the source's own value, i.e. to no conversion. Note that
+  // the sample *type* is always converted, to float32.
+  explicit AudioConverter(
+      std::optional<int> sample_rate = std::nullopt,
+      std::optional<int> num_channels = std::nullopt);
+
+  // `samples` is a contiguous [num_channels, num_samples] tensor in the
+  // source's own sample type, i.e. exactly what audio_samples() produces.
+  // Returns the samples that are now computable, as float32
+  // [out_num_channels, N] - and N may well be 0.
+  torch::stable::Tensor convert(
+      const torch::stable::Tensor& samples,
+      int sample_rate);
+
+  // The samples swresample was still holding on to. Callers who skip this lose
+  // the tail of the stream.
+  torch::stable::Tensor drain();
+
+  // Drop the resampler's buffered state and start over, so that the next
+  // convert() call reconfigures from the samples it is given.
+  void reset();
+
+  bool has_converted_samples() const {
+    return swr_context_ != nullptr;
+  }
+
+  // Only meaningful once convert() has been called at least once.
+  int out_sample_rate() const {
+    return out_sample_rate_;
+  }
+
+ private:
+  std::optional<int> requested_sample_rate_;
+  std::optional<int> requested_num_channels_;
+
+  UniqueSwrContext swr_context_;
+  // All of these are set when swr_context_ is created, and describe what it was
+  // configured for. The src_ ones are what a later convert() is checked
+  // against: swresample is built once for one input shape.
+  AVSampleFormat src_sample_format_ = AV_SAMPLE_FMT_NONE;
+  int src_sample_rate_ = 0;
+  int src_num_channels_ = 0;
+  int out_sample_rate_ = 0;
+  int out_num_channels_ = 0;
+};
+
+} // namespace facebook::torchcodec