Fix 4:4:4 video decoding on CUDA interface (#1415)

Author: NicolasHug

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

test/test_decoders.py

@@ -23,6 +23,7 @@
 )
 from torchcodec.decoders._decoder_utils import _get_cuda_backend
 from torchcodec.decoders._wav_decoder import WavDecoder
+from torchcodec.encoders import VideoEncoder
 from torchcodec.transforms import CenterCrop, RandomCrop, Resize
 
 from .utils import (
@@ -2032,6 +2033,25 @@ def test_nvdec_cuda_interface_cpu_fallback(self):
 
         assert psnr(ref_frame.data, nvdec_frame.data) > 25
 
+    @needs_cuda
+    def test_nvdec_cpu_fallback_yuv444(self, tmp_path):
+        # Non-regression test for https://github.com/meta-pytorch/torchcodec/issues/1414
+        num_frames = 5
+        frames = torch.randint(0, 256, size=(num_frames, 3, 64, 64), dtype=torch.uint8)
+        path = str(tmp_path / "yuv444.mp4")
+        VideoEncoder(frames=frames, frame_rate=30).to_file(
+            path, pixel_format="yuv444p", crf=0
+        )
+
+        cpu_decoder = VideoDecoder(path, device="cpu")
+        cuda_decoder = VideoDecoder(path, device="cuda")
+        assert cuda_decoder.cpu_fallback
+
+        cpu_frames = cpu_decoder.get_frames_in_range(start=0, stop=num_frames).data
+        cuda_frames = cuda_decoder.get_frames_in_range(start=0, stop=num_frames).data
+
+        torch.testing.assert_close(cpu_frames, cuda_frames.cpu(), rtol=0, atol=0)
+
     @needs_cuda
     def test_nvdec_cuda_interface_error(self):
         with pytest.raises(RuntimeError, match="torch_parse_device_string"):