2026-05-16 nightly release (42c8e86)

Author: pytorchbot

pyproject.toml

@@ -54,4 +54,8 @@ norecursedirs = ["third-party-interface"]
 # the CI, where we definitely want the 'slow' tests to run.
 addopts = "-v -m 'not slow'"
 
+# Tells pytest to keep around temporary paths (via tmp_path fixture) for
+# *failed tests* only. Otherwise, the /tmp folder can grow too large.
+tmp_path_retention_policy = "failed"
+
 testpaths = ["test"]

src/torchcodec/_core/BetaCudaDeviceInterface.cpp

@@ -903,6 +903,33 @@ void BetaCudaDeviceInterface::convertAVFrameToFrameOutput(
     UniqueAVFrame& avFrame,
     FrameOutput& frameOutput,
     std::optional<torch::stable::Tensor> preAllocatedOutputTensor) {
+  if (cpuFallback_) {
+    // When the CPU fallabck happens, we'll try to run the color-conversion on
+    // GPU by sending those CPU frames to the GPU as NV12 (See
+    // transferCpuFrameToGpuNV12() below). However, it's not always possible:
+    // NV12 would downsample 4:4:4 frames and lose chroma resolution, resulting
+    // in poorly decoded frames. So for those, we still do the color conversion
+    // on the CPU and then send the full RGB frame to the GPU.
+    const AVPixFmtDescriptor* desc =
+        av_pix_fmt_desc_get(static_cast<AVPixelFormat>(avFrame->format));
+    if (desc && desc->log2_chroma_w == 0 && desc->log2_chroma_h == 0) {
+      // 4:4:4: converting through NV12 (4:2:0) would lose chroma resolution.
+      FrameOutput cpuFrameOutput;
+      cpuFallback_->convertAVFrameToFrameOutput(avFrame, cpuFrameOutput);
+      if (preAllocatedOutputTensor.has_value()) {
+        torch::stable::copy_(
+            preAllocatedOutputTensor.value(), cpuFrameOutput.data);
+        frameOutput.data = preAllocatedOutputTensor.value();
+      } else {
+        frameOutput.data = torch::stable::to(cpuFrameOutput.data, device_);
+      }
+      if (rotation_ != Rotation::NONE) {
+        applyRotation(frameOutput, preAllocatedOutputTensor);
+      }
+      return;
+    }
+  }
+
   // Capture original dimensions before transferCpuFrameToGpuNV12 may
   // round them up to even for NV12.
   FrameDims originalDims(avFrame->height, avFrame->width);

src/torchcodec/_core/Encoder.cpp

@@ -573,7 +573,6 @@ torch::stable::Tensor validateFrames(
         framesDevice.index());
   }
   if (avCodecContext) {
-    // TODO MultiStreamEncoder: Enable tensors in NHWC shape
     STD_TORCH_CHECK(
         static_cast<int>(frames.sizes()[2]) == avCodecContext->height &&
             static_cast<int>(frames.sizes()[3]) == avCodecContext->width,
@@ -832,13 +831,11 @@ void VideoEncoder::initializeEncoder(
   avCodecContext_.reset(avCodecContext);
 
   // Store dimensions of input frames
-  // TODO-VideoEncoder: (P2) Enable tensors in NHWC shape
   auto sizes = frames_.sizes();
   int inHeight = static_cast<int>(sizes[2]);
   int inWidth = static_cast<int>(sizes[3]);
 
   // Always use input dimensions as output dimensions
-  // TODO-VideoEncoder: (P2) Allow height and width to be set
   int outWidth = inWidth;
   int outHeight = inHeight;
   AVPixelFormat outPixelFormat = AV_PIX_FMT_NONE;
@@ -1153,6 +1150,13 @@ int MultiStreamEncoder::addAudioStream(
   STD_TORCH_CHECK(sampleRate > 0, "sample_rate must be > 0, got ", sampleRate);
   STD_TORCH_CHECK(
       numChannels > 0, "num_channels must be > 0, got ", numChannels);
+  STD_TORCH_CHECK(
+      numChannels <= AV_NUM_DATA_POINTERS,
+      "Trying to encode ",
+      numChannels,
+      " channels, but FFmpeg only supports ",
+      AV_NUM_DATA_POINTERS,
+      " channels per frame.");
 
   AudioStream audioStream;
   audioStream.inSampleRate = sampleRate;
@@ -1210,7 +1214,6 @@ void MultiStreamEncoder::initializeVideoStream(VideoStream& videoStream) {
       avCodecContext != nullptr, "Couldn't allocate codec context.");
   videoStream.avCodecContext.reset(avCodecContext);
 
-  // TODO MultiStreamEncoder: Allow output height and width to be set
   int outHeight = videoStream.inHeight;
   int outWidth = videoStream.inWidth;
   AVPixelFormat outPixelFormat = AV_PIX_FMT_NONE;
@@ -1247,7 +1250,6 @@ void MultiStreamEncoder::initializeVideoStream(VideoStream& videoStream) {
   videoStream.avCodecContext->width = outWidth;
   videoStream.avCodecContext->height = outHeight;
   videoStream.avCodecContext->pix_fmt = outPixelFormat;
-  // TODO MultiStreamEncoder: Add and utilize output frame_rate option
   videoStream.avCodecContext->framerate =
       av_d2q(videoStream.inFrameRate, INT_MAX);
   videoStream.avCodecContext->time_base =
@@ -1419,7 +1421,6 @@ void MultiStreamEncoder::openStreamsAndWriteHeader() {
 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(