Allow multiple audio streams in MultiStreamEncoder

src/torchcodec/_core/Encoder.cpp

@@ -1142,21 +1142,20 @@ int MultiStreamEncoder::addVideoStream(
   return static_cast<int>(videoStreams_.size() - 1);
 }
 
-void MultiStreamEncoder::addAudioStream(
+int MultiStreamEncoder::addAudioStream(
     int sampleRate,
     int numChannels,
     std::optional<int> bitRate) {
-  STD_TORCH_CHECK(
-      !audioStream_.has_value(),
-      "An audio stream has already been added. Cannot add another.");
   STD_TORCH_CHECK(sampleRate > 0, "sample_rate must be > 0, got ", sampleRate);
   STD_TORCH_CHECK(
       numChannels > 0, "num_channels must be > 0, got ", numChannels);
 
-  audioStream_ = AudioStream{};
-  audioStream_->inSampleRate = sampleRate;
-  audioStream_->inNumChannels = numChannels;
-  audioStream_->options.bitRate = bitRate;
+  AudioStream audioStream;
+  audioStream.inSampleRate = sampleRate;
+  audioStream.inNumChannels = numChannels;
+  audioStream.options.bitRate = bitRate;
+  audioStreams_.push_back(std::move(audioStream));
+  return static_cast<int>(audioStreams_.size() - 1);
 }
 
 void MultiStreamEncoder::initializeVideoStream(VideoStream& videoStream) {
@@ -1313,8 +1312,7 @@ void MultiStreamEncoder::initializeVideoStream(VideoStream& videoStream) {
       getFFMPEGErrorStringFromErrorCode(status));
 }
 
-void MultiStreamEncoder::initializeAudioStream() {
-  auto& audioStream = *audioStream_;
+void MultiStreamEncoder::initializeAudioStream(AudioStream& audioStream) {
   // We use the AVFormatContext's default codec for that
   // specific format/container.
   const AVCodec* avCodec =
@@ -1390,14 +1388,14 @@ void MultiStreamEncoder::initializeAudioStream() {
 
 void MultiStreamEncoder::openStreamsAndWriteHeader() {
   STD_TORCH_CHECK(
-      !videoStreams_.empty() || audioStream_.has_value(),
+      !videoStreams_.empty() || !audioStreams_.empty(),
       "Call addVideoStream() or addAudioStream() before open().");
 
   for (auto& videoStream : videoStreams_) {
     initializeVideoStream(videoStream);
   }
-  if (audioStream_.has_value()) {
-    initializeAudioStream();
+  for (auto& audioStream : audioStreams_) {
+    initializeAudioStream(audioStream);
   }
 
   int status = avformat_write_header(
@@ -1499,26 +1497,31 @@ void MultiStreamEncoder::encodeVideoFrame(
   }
 }
 
-void MultiStreamEncoder::addSamples(const torch::stable::Tensor& samples) {
+void MultiStreamEncoder::addSamples(
+    const torch::stable::Tensor& samples,
+    int streamIndex) {
   STD_TORCH_CHECK(headerWritten_, "Call open() before addSamples().");
   STD_TORCH_CHECK(
-      audioStream_.has_value(),
-      "No audio stream has been added. Call addAudioStream() first.");
+      streamIndex >= 0 && streamIndex < static_cast<int>(audioStreams_.size()),
+      "Invalid stream index ",
+      streamIndex,
+      ". Number of audio streams: ",
+      audioStreams_.size());
+  auto& audioStream = audioStreams_[streamIndex];
   auto validatedSamples = validateSamples(samples);
   STD_TORCH_CHECK(
       static_cast<int>(validatedSamples.sizes()[0]) ==
-          audioStream_->inNumChannels,
+          audioStream.inNumChannels,
       "Expected ",
-      audioStream_->inNumChannels,
+      audioStream.inNumChannels,
       " channels, got ",
       validatedSamples.sizes()[0]);
-  encodeAudioSamples(validatedSamples);
+  encodeAudioSamples(validatedSamples, audioStream);
 }
 
 void MultiStreamEncoder::encodeAudioSamples(
-    const torch::stable::Tensor& samples) {
-  auto& audioStream = *audioStream_;
-
+    const torch::stable::Tensor& samples,
+    AudioStream& audioStream) {
   UniqueAVFrame avFrame = allocateAVFrame(
       audioStream.frameSize,
       audioStream.inSampleRate,
@@ -1754,11 +1757,12 @@ void MultiStreamEncoder::maybeFlushSwrAndFifo(
 }
 
 void MultiStreamEncoder::flushBuffers() {
-  if (audioStream_.has_value() && audioStream_->avStream != nullptr) {
-    AutoAVPacket audioAVPacket;
-    auto& audioStream = *audioStream_;
-    maybeFlushSwrAndFifo(audioAVPacket, audioStream);
-    encodeAudioFrame(audioAVPacket, UniqueAVFrame(nullptr), audioStream);
+  for (auto& audioStream : audioStreams_) {
+    if (audioStream.avStream != nullptr) {
+      AutoAVPacket audioAVPacket;
+      maybeFlushSwrAndFifo(audioAVPacket, audioStream);
+      encodeAudioFrame(audioAVPacket, UniqueAVFrame(nullptr), audioStream);
+    }
   }
   for (auto& videoStream : videoStreams_) {
     if (videoStream.avStream != nullptr) {

src/torchcodec/_core/Encoder.h

@@ -201,7 +201,7 @@ class FORCE_PUBLIC_VISIBILITY MultiStreamEncoder {
       std::optional<std::string> preset = std::nullopt,
       std::optional<std::map<std::string, std::string>> extraOptions =
           std::nullopt);
-  void addAudioStream(
+  int addAudioStream(
       int sampleRate,
       int numChannels,
       std::optional<int> bitRate = std::nullopt);
@@ -210,7 +210,7 @@ class FORCE_PUBLIC_VISIBILITY MultiStreamEncoder {
       std::string_view formatName,
       std::unique_ptr<AVIOContextHolder> avioContextHolder);
   void addFrames(const torch::stable::Tensor& frames, int streamIndex);
-  void addSamples(const torch::stable::Tensor& samples);
+  void addSamples(const torch::stable::Tensor& samples, int streamIndex);
   void close();
 
  private:
@@ -243,8 +243,10 @@ class FORCE_PUBLIC_VISIBILITY MultiStreamEncoder {
       AutoAVPacket& autoAVPacket,
       const UniqueAVFrame& avFrame,
       VideoStream& videoStream);
-  void initializeAudioStream();
-  void encodeAudioSamples(const torch::stable::Tensor& samples);
+  void initializeAudioStream(AudioStream& audioStream);
+  void encodeAudioSamples(
+      const torch::stable::Tensor& samples,
+      AudioStream& audioStream);
   UniqueAVFrame maybeConvertAudioAVFrame(
       const UniqueAVFrame& avFrame,
       AudioStream& audioStream);
@@ -264,7 +266,7 @@ class FORCE_PUBLIC_VISIBILITY MultiStreamEncoder {
 
   UniqueEncodingAVFormatContext avFormatContext_;
   std::vector<VideoStream> videoStreams_;
-  std::optional<AudioStream> audioStream_;
+  std::vector<AudioStream> audioStreams_;
   bool headerWritten_ = false;
   UniqueAVDictionary avFormatOptions_;
 

src/torchcodec/_core/custom_ops.cpp

@@ -94,14 +94,14 @@ STABLE_TORCH_LIBRARY(torchcodec_ns, m) {
   m.def(
       "streaming_encoder_add_video_stream(Tensor(a!) encoder, int height, int width, float frame_rate, str device=\"cpu\", str? codec=None, str? pixel_format=None, float? crf=None, str? preset=None, str[]? extra_options=None) -> int");
   m.def(
-      "streaming_encoder_add_audio_stream(Tensor(a!) encoder, int sample_rate, int num_channels, int? bit_rate=None) -> ()");
+      "streaming_encoder_add_audio_stream(Tensor(a!) encoder, int sample_rate, int num_channels, int? bit_rate=None) -> int");
   m.def("streaming_encoder_open_file(Tensor(a!) encoder, str filename) -> ()");
   m.def(
       "streaming_encoder_open_file_like(Tensor(a!) encoder, str format, int file_like_context) -> ()");
   m.def(
       "streaming_encoder_add_frames(Tensor(a!) encoder, Tensor frames, int stream_index) -> ()");
   m.def(
-      "streaming_encoder_add_samples(Tensor(a!) encoder, Tensor samples) -> ()");
+      "streaming_encoder_add_samples(Tensor(a!) encoder, Tensor samples, int stream_index) -> ()");
   m.def("set_nvdec_cache_capacity(int capacity) -> ()");
   m.def("get_nvdec_cache_capacity() -> int");
   m.def("_get_nvdec_cache_size(int device_index) -> int");
@@ -1258,15 +1258,16 @@ void streaming_encoder_open_file_like(
       format, std::move(avioContextHolder));
 }
 
-void streaming_encoder_add_audio_stream(
+int64_t streaming_encoder_add_audio_stream(
     torch::stable::Tensor& encoder,
     int64_t sample_rate,
     int64_t num_channels,
     std::optional<int64_t> bit_rate = std::nullopt) {
-  unwrapTensorToGetMultiStreamEncoder(encoder)->addAudioStream(
-      validateInt64ToInt(sample_rate, "sample_rate"),
-      validateInt64ToInt(num_channels, "num_channels"),
-      validateOptionalInt64ToInt(bit_rate, "bit_rate"));
+  return static_cast<int64_t>(
+      unwrapTensorToGetMultiStreamEncoder(encoder)->addAudioStream(
+          validateInt64ToInt(sample_rate, "sample_rate"),
+          validateInt64ToInt(num_channels, "num_channels"),
+          validateOptionalInt64ToInt(bit_rate, "bit_rate")));
 }
 
 void streaming_encoder_add_frames(
@@ -1279,8 +1280,10 @@ void streaming_encoder_add_frames(
 
 void streaming_encoder_add_samples(
     torch::stable::Tensor& encoder,
-    const torch::stable::Tensor& samples) {
-  unwrapTensorToGetMultiStreamEncoder(encoder)->addSamples(samples);
+    const torch::stable::Tensor& samples,
+    int64_t stream_index) {
+  unwrapTensorToGetMultiStreamEncoder(encoder)->addSamples(
+      samples, static_cast<int>(stream_index));
 }
 
 torch::stable::Tensor create_wav_decoder_from_file(

src/torchcodec/_core/ops.py

@@ -657,8 +657,8 @@ def streaming_encoder_add_audio_stream_abstract(
     sample_rate: int,
     num_channels: int,
     bit_rate: int | None = None,
-) -> None:
-    return
+) -> int:
+    return 0
 
 
 @register_fake("torchcodec_ns::streaming_encoder_open_file")
@@ -682,7 +682,7 @@ def streaming_encoder_add_frames_abstract(
 
 @register_fake("torchcodec_ns::streaming_encoder_add_samples")
 def streaming_encoder_add_samples_abstract(
-    encoder: torch.Tensor, samples: torch.Tensor
+    encoder: torch.Tensor, samples: torch.Tensor, stream_index: int
 ) -> None:
     return
 

src/torchcodec/encoders/_multi_stream_encoder.py

@@ -5,6 +5,12 @@
 from torchcodec import _core
 
 
+# TODO MultiStreamEncoder: the stream_index values here are per media-type,
+# while everywhere else in the code base (and particularly in the public decoder
+# APIs) they are absolute across all media types. That'll quickly becomes
+# confusing, and we should definitely not expose this one as-is. We should either:
+# - keep it private but rename it to something that's not stream_index
+# - make it absolute per container, if we ever want to expose it.
 class _VideoStream:
     def __init__(self, encoder_tensor: Tensor, stream_index: int):
         self._encoder_tensor = encoder_tensor
@@ -17,11 +23,14 @@ def write(self, frames: Tensor) -> None:
 
 
 class _AudioStream:
-    def __init__(self, encoder_tensor: Tensor):
+    def __init__(self, encoder_tensor: Tensor, stream_index: int):
         self._encoder_tensor = encoder_tensor
+        self._stream_index = stream_index
 
     def write(self, samples: Tensor) -> None:
-        _core.streaming_encoder_add_samples(self._encoder_tensor, samples)
+        _core.streaming_encoder_add_samples(
+            self._encoder_tensor, samples, self._stream_index
+        )
 
 
 class StreamingEncoder:
@@ -65,13 +74,13 @@ def add_audio(
         num_channels: int,
         bit_rate: int | None = None,
     ) -> _AudioStream:
-        _core.streaming_encoder_add_audio_stream(
+        stream_index = _core.streaming_encoder_add_audio_stream(
             self._encoder_tensor,
             sample_rate=sample_rate,
             num_channels=num_channels,
             bit_rate=bit_rate,
         )
-        return _AudioStream(self._encoder_tensor)
+        return _AudioStream(self._encoder_tensor, stream_index)
 
     # TODO MultiStreamEncoder: Maybe there should 2 separate methods, one for
     # file, one for file-like.

test/test_encoders.py

@@ -1729,13 +1729,6 @@ def test_fragmented_mp4(self, format, tmp_path, method, device):
             atol=atol,
         )
 
-    @pytest.mark.parametrize("method", ("to_file", "to_file_like"))
-    def test_add_audio_twice_errors(self, tmp_path, method):
-        enc, _, open_kwargs = self._create_encoder(method, tmp_path, "mp4")
-        enc.add_audio(sample_rate=44100, num_channels=2)
-        with pytest.raises(RuntimeError, match="already been added"):
-            enc.add_audio(sample_rate=16000, num_channels=1)
-
     @pytest.mark.parametrize("method", ("to_file", "to_file_like"))
     @pytest.mark.parametrize("device", cpu_and_oss_cuda)
     def test_write_frames_mismatched_dimensions_errors(self, tmp_path, method, device):
@@ -1944,8 +1937,9 @@ def test_multiple_video_streams_and_audio(self, tmp_path, method):
         source_frames_small = torch.randint(0, 256, (8, 3, 128, 128), dtype=torch.uint8)
 
         source_audio = AudioDecoder(str(NASA_AUDIO_MP3_44100.path)).get_all_samples()
-        source_samples = source_audio.data
+        source_samples_stereo = source_audio.data
         sample_rate = source_audio.sample_rate
+        source_samples_mono = source_samples_stereo[:1]
 
         enc, encoder_output, open_kwargs = self._create_encoder(method, tmp_path, "mp4")
         video_big = enc.add_video(
@@ -1954,17 +1948,27 @@ def test_multiple_video_streams_and_audio(self, tmp_path, method):
         video_small = enc.add_video(
             height=128, width=128, frame_rate=25.0, pixel_format="yuv444p", crf=0
         )
-        audio = enc.add_audio(
+        audio_stereo = enc.add_audio(
             sample_rate=sample_rate,
-            num_channels=source_samples.shape[0],
+            num_channels=2,
+        )
+        audio_mono = enc.add_audio(
+            sample_rate=sample_rate,
+            num_channels=1,
         )
         enc.open(**open_kwargs)
         video_big.write(source_frames_big[:3])
         video_small.write(source_frames_small[:4])
-        audio.write(source_samples[:, : source_samples.shape[1] // 2])
+        audio_stereo.write(
+            source_samples_stereo[:, : source_samples_stereo.shape[1] // 2]
+        )
+        audio_mono.write(source_samples_mono[:, : source_samples_mono.shape[1] // 2])
         video_big.write(source_frames_big[3:])
         video_small.write(source_frames_small[4:])
-        audio.write(source_samples[:, source_samples.shape[1] // 2 :])
+        audio_stereo.write(
+            source_samples_stereo[:, source_samples_stereo.shape[1] // 2 :]
+        )
+        audio_mono.write(source_samples_mono[:, source_samples_mono.shape[1] // 2 :])
         enc.close()
 
         source = self._get_decoder_source(encoder_output)
@@ -1989,15 +1993,29 @@ def test_multiple_video_streams_and_audio(self, tmp_path, method):
             decoded_small_frames, source_frames_small, percentage=99, atol=2
         )
 
-        audio_decoder = AudioDecoder(source)
-        decoded_audio = audio_decoder.get_all_samples()
-        assert decoded_audio.sample_rate == sample_rate
+        # stream_index is absolute: 0, 1 are video; 2, 3 are audio
+        decoded_stereo = AudioDecoder(source, stream_index=2).get_all_samples()
+        assert decoded_stereo.sample_rate == sample_rate
+        assert decoded_stereo.data.shape[0] == 2
         num_samples_to_compare = min(
-            decoded_audio.data.shape[1], source_samples.shape[1]
+            decoded_stereo.data.shape[1], source_samples_stereo.shape[1]
         )
         assert_tensor_close_on_at_least(
-            decoded_audio.data[:, :num_samples_to_compare],
-            source_samples[:, :num_samples_to_compare],
-            percentage=99,
+            decoded_stereo.data[:, :num_samples_to_compare],
+            source_samples_stereo[:, :num_samples_to_compare],
+            percentage=98,
+            atol=0.01,
+        )
+
+        decoded_mono = AudioDecoder(source, stream_index=3).get_all_samples()
+        assert decoded_mono.sample_rate == sample_rate
+        assert decoded_mono.data.shape[0] == 1
+        num_samples_to_compare = min(
+            decoded_mono.data.shape[1], source_samples_mono.shape[1]
+        )
+        assert_tensor_close_on_at_least(
+            decoded_mono.data[:, :num_samples_to_compare],
+            source_samples_mono[:, :num_samples_to_compare],
+            percentage=98,
             atol=0.01,
         )