2026-05-15 nightly release (da5b1cb)

Author: pytorchbot

.github/workflows/cpp_tests.yaml

@@ -19,7 +19,7 @@ jobs:
     strategy:
       fail-fast: false
       matrix:
-        ffmpeg-version-for-tests: ['4.4.2', '5.1.2', '6.1.1', '7.0.1']
+        ffmpeg-version-for-tests: ['7.0.1']
     steps:
       - name: Check out repo
         uses: actions/checkout@v6

.github/workflows/windows_wheel.yaml

@@ -36,7 +36,7 @@ jobs:
       with-rocm: disable
       with-cuda: disable
       build-python-only: "disable"
-      python-versions: '["3.10", "3.11", "3.12", "3.13"]'
+      python-versions: '["3.10", "3.11", "3.12", "3.13", "3.14"]'
 
 
   build:

src/torchcodec/_core/Encoder.cpp

@@ -1060,6 +1060,7 @@ MultiStreamEncoder::MultiStreamEncoder() {
 }
 
 void MultiStreamEncoder::open(std::string_view fileName) {
+  STD_TORCH_CHECK(!closed_, "Cannot open after close() was called.");
   STD_TORCH_CHECK(!headerWritten_, "open() was already called.");
 
   AVFormatContext* avFormatContext = nullptr;
@@ -1089,6 +1090,7 @@ void MultiStreamEncoder::open(std::string_view fileName) {
 void MultiStreamEncoder::open(
     std::string_view formatName,
     std::unique_ptr<AVIOContextHolder> avioContextHolder) {
+  STD_TORCH_CHECK(!closed_, "Cannot open after close() was called.");
   STD_TORCH_CHECK(!headerWritten_, "open() was already called.");
 
   avioContextHolder_ = std::move(avioContextHolder);
@@ -1113,7 +1115,7 @@ void MultiStreamEncoder::open(
   openStreamsAndWriteHeader();
 }
 
-void MultiStreamEncoder::addVideoStream(
+int MultiStreamEncoder::addVideoStream(
     int height,
     int width,
     double frameRate,
@@ -1123,48 +1125,46 @@ void MultiStreamEncoder::addVideoStream(
     std::optional<double> crf,
     std::optional<std::string> preset,
     std::optional<std::map<std::string, std::string>> extraOptions) {
-  STD_TORCH_CHECK(
-      !videoStream_.has_value(),
-      "A video stream has already been added. Cannot add another.");
   STD_TORCH_CHECK(height > 0, "height must be > 0, got ", height);
   STD_TORCH_CHECK(width > 0, "width must be > 0, got ", width);
   STD_TORCH_CHECK(frameRate > 0, "frame_rate must be > 0, got ", frameRate);
-  videoStream_ = VideoStream{};
+  VideoStream videoStream;
   StableDevice stableDevice(std::move(device));
-  // The NVDEC CUDA interface is decode-only; encoders need the FFmpeg-based
-  // one.
-  videoStream_->deviceInterface = createDeviceInterface(
+  videoStream.deviceInterface = createDeviceInterface(
       stableDevice, stableDevice.type() == kStableCUDA ? "ffmpeg" : "default");
-  videoStream_->inHeight = height;
-  videoStream_->inWidth = width;
-  videoStream_->inFrameRate = frameRate;
-  videoStream_->options.codec = std::move(codec);
-  videoStream_->options.pixelFormat = std::move(pixelFormat);
-  videoStream_->options.crf = crf;
-  videoStream_->options.preset = std::move(preset);
-  videoStream_->options.extraOptions = std::move(extraOptions);
+  videoStream.inHeight = height;
+  videoStream.inWidth = width;
+  videoStream.inFrameRate = frameRate;
+  videoStream.options.codec = std::move(codec);
+  videoStream.options.pixelFormat = std::move(pixelFormat);
+  videoStream.options.crf = crf;
+  videoStream.options.preset = std::move(preset);
+  videoStream.options.extraOptions = std::move(extraOptions);
+  videoStreams_.push_back(std::move(videoStream));
+  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::optional<int> bitRate,
+    std::optional<int> outNumChannels,
+    std::optional<int> outSampleRate) {
   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;
+  audioStream.options.numChannels = outNumChannels;
+  audioStream.options.sampleRate = outSampleRate;
+  audioStreams_.push_back(std::move(audioStream));
+  return static_cast<int>(audioStreams_.size() - 1);
 }
 
-void MultiStreamEncoder::initializeVideoStream() {
-  // TODO MultiStreamEncoder: Iterate over all video streams
-  auto& videoStream = *videoStream_;
+void MultiStreamEncoder::initializeVideoStream(VideoStream& videoStream) {
   auto deviceType = videoStream.deviceInterface->device().type();
 
   const AVCodec* avCodec = nullptr;
@@ -1318,8 +1318,7 @@ void MultiStreamEncoder::initializeVideoStream() {
       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 =
@@ -1340,15 +1339,18 @@ void MultiStreamEncoder::initializeAudioStream() {
   // well when "-b:a" isn't specified.
   audioStream.avCodecContext->bit_rate = desiredBitRate.value_or(0);
 
-  // TODO MultiStreamEncoder: support output numChannels and sampleRate
-  validateNumChannels(*avCodec, audioStream.inNumChannels);
-  setDefaultChannelLayout(
-      audioStream.avCodecContext, audioStream.inNumChannels);
+  int outNumChannels =
+      audioStream.options.numChannels.value_or(audioStream.inNumChannels);
+  audioStream.outNumChannels = outNumChannels;
+  validateNumChannels(*avCodec, outNumChannels);
+  setDefaultChannelLayout(audioStream.avCodecContext, outNumChannels);
 
-  validateSampleRate(*avCodec, audioStream.inSampleRate);
-  audioStream.avCodecContext->sample_rate = audioStream.inSampleRate;
-  audioStream.avCodecContext->time_base =
-      AVRational{1, audioStream.inSampleRate};
+  int outSampleRate =
+      audioStream.options.sampleRate.value_or(audioStream.inSampleRate);
+  audioStream.outSampleRate = outSampleRate;
+  validateSampleRate(*avCodec, outSampleRate);
+  audioStream.avCodecContext->sample_rate = outSampleRate;
+  audioStream.avCodecContext->time_base = AVRational{1, outSampleRate};
 
   // Input samples are expected to be FLTP. Not all encoders support FLTP, so we
   // may need to convert the samples into a supported output sample format,
@@ -1387,22 +1389,22 @@ void MultiStreamEncoder::initializeAudioStream() {
   // sized batches.
   auto avAudioFifo = av_audio_fifo_alloc(
       audioStream.avCodecContext->sample_fmt,
-      audioStream.inNumChannels,
+      outNumChannels,
       audioStream.frameSize * 2);
   STD_TORCH_CHECK(avAudioFifo != nullptr, "Couldn't create AVAudioFifo.");
   audioStream.avAudioFifo.reset(avAudioFifo);
 }
 
 void MultiStreamEncoder::openStreamsAndWriteHeader() {
   STD_TORCH_CHECK(
-      videoStream_.has_value() || audioStream_.has_value(),
+      !videoStreams_.empty() || !audioStreams_.empty(),
       "Call addVideoStream() or addAudioStream() before open().");
 
-  if (videoStream_.has_value()) {
-    initializeVideoStream();
+  for (auto& videoStream : videoStreams_) {
+    initializeVideoStream(videoStream);
   }
-  if (audioStream_.has_value()) {
-    initializeAudioStream();
+  for (auto& audioStream : audioStreams_) {
+    initializeAudioStream(audioStream);
   }
 
   int status = avformat_write_header(
@@ -1414,49 +1416,60 @@ void MultiStreamEncoder::openStreamsAndWriteHeader() {
   headerWritten_ = true;
 }
 
-void MultiStreamEncoder::addFrames(const torch::stable::Tensor& frames) {
+void MultiStreamEncoder::addFrames(
+    const torch::stable::Tensor& frames,
+    int streamIndex) {
   // TODO MultiStreamEncoder: Specify which video stream to add frames to
+  STD_TORCH_CHECK(!closed_, "Cannot add frames after close() was called.");
   STD_TORCH_CHECK(headerWritten_, "Call open() before addFrames().");
+  STD_TORCH_CHECK(
+      streamIndex >= 0 && streamIndex < static_cast<int>(videoStreams_.size()),
+      "Invalid stream index ",
+      streamIndex,
+      ". Number of video streams: ",
+      videoStreams_.size());
+  auto& videoStream = videoStreams_[streamIndex];
   auto validatedFrames = validateFrames(
       frames,
-      videoStream_->avCodecContext.get(),
-      videoStream_->deviceInterface.get());
+      videoStream.avCodecContext.get(),
+      videoStream.deviceInterface.get());
 
   AutoAVPacket autoAVPacket;
   // TODO MultiStreamEncoder: Consider using accessor for potential performance
   // improvement
   int numFrames = static_cast<int>(validatedFrames.sizes()[0]);
   for (int i = 0; i < numFrames; ++i) {
     torch::stable::Tensor currFrame = selectRow(validatedFrames, i);
-    int frameIndex = videoStream_->numEncodedFrames + i;
+    int frameIndex = videoStream.numEncodedFrames + i;
     UniqueAVFrame avFrame =
-        videoStream_->deviceInterface->convertTensorToAVFrameForEncoding(
-            currFrame, frameIndex, videoStream_->avCodecContext.get());
+        videoStream.deviceInterface->convertTensorToAVFrameForEncoding(
+            currFrame, frameIndex, videoStream.avCodecContext.get());
     STD_TORCH_CHECK(
         avFrame != nullptr,
         "convertTensorToAVFrameForEncoding failed for frame ",
         frameIndex,
         " on device: ",
         deviceTypeName(validatedFrames.device().type()));
-    encodeVideoFrame(autoAVPacket, avFrame);
+    encodeVideoFrame(autoAVPacket, avFrame, videoStream);
   }
-  videoStream_->numEncodedFrames += numFrames;
+  videoStream.numEncodedFrames += numFrames;
 }
 
 void MultiStreamEncoder::encodeVideoFrame(
     AutoAVPacket& autoAVPacket,
-    const UniqueAVFrame& avFrame) {
+    const UniqueAVFrame& avFrame,
+    VideoStream& videoStream) {
   auto status =
-      avcodec_send_frame(videoStream_->avCodecContext.get(), avFrame.get());
+      avcodec_send_frame(videoStream.avCodecContext.get(), avFrame.get());
   STD_TORCH_CHECK(
       status == AVSUCCESS,
       "Error while sending frame: ",
       getFFMPEGErrorStringFromErrorCode(status));
 
   while (status >= 0) {
     ReferenceAVPacket packet(autoAVPacket);
-    status = avcodec_receive_packet(
-        videoStream_->avCodecContext.get(), packet.get());
+    status =
+        avcodec_receive_packet(videoStream.avCodecContext.get(), packet.get());
     if (status == AVERROR(EAGAIN) || status == AVERROR_EOF) {
       if (status == AVERROR_EOF) {
         // Flush remaining buffered packets
@@ -1482,9 +1495,9 @@ void MultiStreamEncoder::encodeVideoFrame(
     }
     av_packet_rescale_ts(
         packet.get(),
-        videoStream_->avCodecContext->time_base,
-        videoStream_->avStream->time_base);
-    packet->stream_index = videoStream_->avStream->index;
+        videoStream.avCodecContext->time_base,
+        videoStream.avStream->time_base);
+    packet->stream_index = videoStream.avStream->index;
 
     status = av_interleaved_write_frame(avFormatContext_.get(), packet.get());
     STD_TORCH_CHECK(
@@ -1494,26 +1507,32 @@ void MultiStreamEncoder::encodeVideoFrame(
   }
 }
 
-void MultiStreamEncoder::addSamples(const torch::stable::Tensor& samples) {
+void MultiStreamEncoder::addSamples(
+    const torch::stable::Tensor& samples,
+    int streamIndex) {
+  STD_TORCH_CHECK(!closed_, "Cannot add samples after close() was called.");
   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,
@@ -1564,8 +1583,8 @@ UniqueAVFrame MultiStreamEncoder::maybeConvertAudioAVFrame(
     AudioStream& audioStream) {
   if (static_cast<AVSampleFormat>(avFrame->format) ==
           audioStream.avCodecContext->sample_fmt &&
-      getNumChannels(avFrame) == audioStream.inNumChannels &&
-      avFrame->sample_rate == audioStream.inSampleRate) {
+      getNumChannels(avFrame) == audioStream.outNumChannels &&
+      avFrame->sample_rate == audioStream.outSampleRate) {
     // Note: the clone references the same underlying data, it's a cheap copy.
     return UniqueAVFrame(av_frame_clone(avFrame.get()));
   }
@@ -1575,9 +1594,9 @@ UniqueAVFrame MultiStreamEncoder::maybeConvertAudioAVFrame(
         static_cast<AVSampleFormat>(avFrame->format),
         audioStream.avCodecContext->sample_fmt,
         avFrame->sample_rate,
-        audioStream.inSampleRate,
+        audioStream.outSampleRate,
         avFrame,
-        audioStream.inNumChannels));
+        audioStream.outNumChannels));
   }
   // convertAudioAVFrameSamples uses avFrame's extended_data field, so we ensure
   // it's the same as data. This should always be the case since we validated
@@ -1589,10 +1608,10 @@ UniqueAVFrame MultiStreamEncoder::maybeConvertAudioAVFrame(
       audioStream.swrContext,
       avFrame,
       audioStream.avCodecContext->sample_fmt,
-      audioStream.inSampleRate,
-      audioStream.inNumChannels);
+      audioStream.outSampleRate,
+      audioStream.outNumChannels);
 
-  if (avFrame->sample_rate == audioStream.inSampleRate) {
+  if (avFrame->sample_rate == audioStream.outSampleRate) {
     STD_TORCH_CHECK(
         convertedAVFrame->nb_samples == avFrame->nb_samples,
         "convertedAVFrame->nb_samples=",
@@ -1629,7 +1648,7 @@ void MultiStreamEncoder::encodeAudioFrameThroughFifo(
   UniqueAVFrame newavFrame = allocateAVFrame(
       audioStream.frameSize,
       audioStream.avCodecContext->sample_rate,
-      audioStream.inNumChannels,
+      audioStream.outNumChannels,
       audioStream.avCodecContext->sample_fmt);
 
   // Explaining the while bound:
@@ -1730,8 +1749,8 @@ void MultiStreamEncoder::maybeFlushSwrAndFifo(
     if (numRemainingSamples > 0) {
       swrFrame = allocateAVFrame(
           numRemainingSamples,
-          audioStream.inSampleRate,
-          audioStream.inNumChannels,
+          audioStream.outSampleRate,
+          audioStream.outNumChannels,
           audioStream.avCodecContext->sample_fmt);
       int actualNumRemainingSamples = swr_convert(
           audioStream.swrContext.get(),
@@ -1749,16 +1768,18 @@ 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);
-  }
-  if (videoStream_.has_value() && videoStream_->avStream != nullptr) {
-    AutoAVPacket videoAVPacket;
-    // Send null frame to signal end of input
-    encodeVideoFrame(videoAVPacket, UniqueAVFrame(nullptr));
+  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) {
+      AutoAVPacket videoAVPacket;
+      encodeVideoFrame(videoAVPacket, UniqueAVFrame(nullptr), videoStream);
+    }
   }
 }