API breakdown: CUDA

src/torchcodec/_core/ColorConverter.cpp

@@ -28,14 +28,11 @@ ColorConverter::ColorConverter(
   options.device = device;
 
   // No user transforms and no stream: the converter is stream-agnostic and
-  // derives everything it needs from each frame.
-  //
-  // TODO_API_BREAKDOWN Need to refac/rethink all this. It seems unnatural that
-  // the color-converter needs its own device_interface_, but at the same time
-  // the color-conversion *must* be third-party aware, and the only way to
-  // achieve that for now is via the interface.
-  // This will become very relevant when we tackle CUDA, so we can defer until
-  // then. For now this is an OK hack.
+  // derives everything it needs (dimensions, pixel format, and on CUDA the
+  // hardware context + stream) from each frame. This is what keeps the block
+  // self-contained -- it needs no reference to the decoder that produced the
+  // frame, on either CPU or CUDA. The DeviceInterface is the vendor extension
+  // point, so color conversion still goes through it to stay third-party aware.
   std::vector<std::unique_ptr<Transform>> no_transforms;
   device_interface_->initialize_video(
       /*av_stream=*/nullptr,

src/torchcodec/_core/ColorConverter.h

@@ -15,6 +15,9 @@
 
 namespace facebook::torchcodec {
 
+// Self-contained color-conversion block: owns its own DeviceInterface and is
+// not bound to any decoder or stream. Everything it needs is derived from each
+// frame, so one converter can process frames from any video, on CPU or CUDA.
 class FORCE_PUBLIC_VISIBILITY ColorConverter {
  public:
   explicit ColorConverter(

src/torchcodec/_core/CudaDeviceInterface.cpp

@@ -120,11 +120,20 @@ void CudaDeviceInterface::initialize_video(
     const VideoStreamOptions& video_stream_options,
     [[maybe_unused]] const std::vector<std::unique_ptr<Transform>>& transforms,
     [[maybe_unused]] const std::optional<FrameDims>& resized_output_dims) {
-  STD_TORCH_CHECK(av_stream != nullptr, "avStream is null");
-  time_base_ = av_stream->time_base;
+  // av_stream may be null: that's how a standalone ColorConverter configures
+  // conversion without being bound to a decoder/stream. Everything the convert
+  // path needs is derived from each frame (dimensions, pixel format, hardware
+  // context and CUDA stream), so a null stream is fine here. The time base only
+  // feeds the (rarely used) nv12 filtergraph's pts and doesn't affect pixels,
+  // so we fall back to a default. Mirrors CpuDeviceInterface::initialize_video.
+  time_base_ = (av_stream != nullptr) ? av_stream->time_base
+                                      : AVRational{1, AV_TIME_BASE};
   video_stream_options_ = video_stream_options;
 
   // TODO: Ideally, we should keep all interface implementations independent.
+  // The CPU interface is only used as a fallback when NVDEC can't decode a
+  // codec and hands us back an already-decoded CPU frame; color-converting that
+  // frame is itself stream-agnostic, so it too tolerates a null stream.
   cpu_interface_ = create_device_interface(kStableCPU);
   STD_TORCH_CHECK(
       cpu_interface_ != nullptr, "Failed to create CPU device interface");

src/torchcodec/_core/custom_ops.cpp

@@ -231,15 +231,15 @@ AVPacket* unwrap_tensor_to_packet(torch::stable::Tensor& tensor) {
 }
 
 torch::stable::Tensor wrap_frame_pointer_to_tensor(AVFrame* frame) {
-  // Owning handle: the frame is freed when the handle tensor's refcount drops.
-  // ColorConverter borrows the frame during conversion (on CPU it does not free
-  // it), so the handle stays the sole owner and there is no leak even if a
-  // frame is never converted. (GPU conversion would consume the frame; GPU is
-  // not exposed through these ops yet.)
-  auto deleter = [frame](void*) {
-    AVFrame* f = frame;
-    av_frame_free(&f);
-  };
+  // NON-owning handle (no-op deleter): ownership is transferred to
+  // _blocks_convert_frame, which consumes the frame. This is the only model
+  // that's correct across devices: CPU convert reads the frame (the op's local
+  // UniqueAVFrame frees it), the CUDA/NVDEC path std::move's it, and the CUDA
+  // ffmpeg path may replace it with a filtered frame -- in every case convert
+  // ends up owning (and freeing) whatever frame is live, so the handle must not
+  // also free. Trade-off: a frame that is never converted leaks; in practice
+  // pipelines convert every frame.
+  auto deleter = [](void*) {};
   int64_t sizes[] = {1};
   int64_t strides[] = {1};
   return torch::stable::from_blob(
@@ -933,12 +933,12 @@ torch::stable::Tensor _blocks_convert_frame(
   ColorConverter* converter_ptr =
       unwrap_tensor_to_pointer<ColorConverter>(converter);
   AVFrame* raw_frame = unwrap_tensor_to_frame(frame);
-  // Borrow the frame for conversion, then release() so the handle keeps
-  // ownership and frees it when its tensor is dropped (CPU path).
-  UniqueAVFrame borrowed(raw_frame);
-  torch::stable::Tensor data = converter_ptr->convert(borrowed);
-  borrowed.release();
-  return data;
+  // Take ownership of the frame (the handle is non-owning) and let `owned` free
+  // whatever frame is live after convert: CPU leaves `owned` holding raw_frame;
+  // the CUDA/NVDEC path moves it out (owned becomes null); the CUDA ffmpeg path
+  // may replace it with a filtered frame. Single free in every case.
+  UniqueAVFrame owned(raw_frame);
+  return converter_ptr->convert(owned);
 }
 
 // For testing only. We need to implement this operation as a core library

src/torchcodec/decoders/_blocks/_color_converter.py

@@ -22,16 +22,21 @@ class ColorConverter:
     :class:`DecodedFrame` into an RGB :class:`~torchcodec._frame.Frame`
     (CHW, uint8 -- matching ``VideoDecoder``'s default output).
 
-    Not bound to anything: everything it needs (dims, pixel format, colorspace)
-    comes from the frame itself, so one converter can process frames from any
-    video. Passive and *not* thread-safe: use one ``ColorConverter`` per thread.
+    Self-contained and not bound to anything: everything it needs (dims, pixel
+    format, colorspace, and on CUDA the hardware context) comes from the frame
+    itself, so one converter can process frames from any video -- it needs no
+    reference to the decoder that produced them. ``device`` selects where the
+    conversion runs; on CUDA use ``device_variant="ffmpeg"`` to match a
+    ``PacketDecoder`` configured the same way.
+
+    Passive and *not* thread-safe: use one ``ColorConverter`` per thread.
 
     Note: automatic rotation (from stream side data) is not applied, since this
     block is intentionally stream-agnostic.
     """
 
-    def __init__(self):
-        self._handle = _blocks_create_color_converter()
+    def __init__(self, device: str = "cpu", device_variant: str = "default"):
+        self._handle = _blocks_create_color_converter(device, device_variant)
 
     def convert(self, decoded_frame: DecodedFrame) -> Frame:
         data = _blocks_convert_frame(self._handle, decoded_frame._handle)

src/torchcodec/decoders/_blocks/_packet_decoder.py

@@ -29,10 +29,25 @@ class PacketDecoder:
     use one ``PacketDecoder`` per thread. FFmpeg's internal codec thread count
     is kept at 1 for now (not exposed); parallelism comes from composing blocks
     on your own threads.
+
+    On CUDA, use ``device_variant="ffmpeg"`` (FFmpeg NVDEC hwaccel): it produces
+    self-contained, thread-movable frames. The default NVDEC variant hands out
+    transient GPU surfaces that the blocks don't support yet.
     """
 
-    def __init__(self, demuxer: Demuxer):
-        self._handle = _blocks_create_packet_decoder(demuxer._handle, num_threads=1)
+    def __init__(
+        self,
+        demuxer: Demuxer,
+        *,
+        device: str = "cpu",
+        device_variant: str = "default",
+    ):
+        self._handle = _blocks_create_packet_decoder(
+            demuxer._handle,
+            num_threads=1,
+            device=device,
+            device_variant=device_variant,
+        )
 
     def _drain(self) -> list[DecodedFrame]:
         frames = []

test/test_decoders.py

@@ -3351,42 +3351,44 @@ def drain():
 
         return drain()
 
-    def _decoded_frames(self, path):
-        # demux + decode, as a single generator of DecodedFrames (pts order).
+    @staticmethod
+    def _make_blocks(path, device="cpu"):
+        # Build the three blocks for `device`. All blocks are self-contained: the
+        # ColorConverter is not bound to the decoder even on CUDA. On CUDA we use
+        # the "ffmpeg" variant for both decoder and converter.
+        variant = "default" if device == "cpu" else "ffmpeg"
         demuxer = Demuxer(path)
-        decoder = PacketDecoder(demuxer)
+        decoder = PacketDecoder(demuxer, device=device, device_variant=variant)
+        converter = ColorConverter(device=device, device_variant=variant)
+        return demuxer, decoder, converter
+
+    def _decoded_frames(self, path, device="cpu"):
+        # demux + decode, as a single generator of DecodedFrames (pts order).
+        demuxer, decoder, _ = self._make_blocks(path, device)
         return self._decode(decoder, self._demux(demuxer))
 
-    def _decode_sequential(self, path):
+    def _decode_sequential(self, path, device="cpu"):
         # demux -> decode -> color-convert, all on the calling thread.
-        demuxer = Demuxer(path)
-        decoder = PacketDecoder(demuxer)
-        converter = ColorConverter()
+        demuxer, decoder, converter = self._make_blocks(path, device)
         return list(
             self._convert(converter, self._decode(decoder, self._demux(demuxer)))
         )
 
-    def _decode_prefetch_frames(self, path):
+    def _decode_prefetch_frames(self, path, device="cpu"):
         # [demux + decode] on one thread || [color-convert] on another.
-        demuxer = Demuxer(path)
-        decoder = PacketDecoder(demuxer)
-        converter = ColorConverter()
+        demuxer, decoder, converter = self._make_blocks(path, device)
         frames = self.prefetch(self._decode(decoder, self._demux(demuxer)))
         return list(self._convert(converter, frames))
 
-    def _decode_prefetch_packets(self, path):
+    def _decode_prefetch_packets(self, path, device="cpu"):
         # [demux] on one thread || [decode + color-convert] on another.
-        demuxer = Demuxer(path)
-        decoder = PacketDecoder(demuxer)
-        converter = ColorConverter()
+        demuxer, decoder, converter = self._make_blocks(path, device)
         packets = self.prefetch(self._demux(demuxer))
         return list(self._convert(converter, self._decode(decoder, packets)))
 
-    def _decode_prefetch_packets_and_frames(self, path):
+    def _decode_prefetch_packets_and_frames(self, path, device="cpu"):
         # [demux] || [decode] || [color-convert], each on its own thread.
-        demuxer = Demuxer(path)
-        decoder = PacketDecoder(demuxer)
-        converter = ColorConverter()
+        demuxer, decoder, converter = self._make_blocks(path, device)
         packets = self.prefetch(self._demux(demuxer))
         frames = self.prefetch(self._decode(decoder, packets))
         return list(self._convert(converter, frames))
@@ -3450,6 +3452,35 @@ def test_color_converter_reused_across_videos(self):
             assert got.data.shape == ref.data.shape
             torch.testing.assert_close(got.data, ref.data, atol=0, rtol=0)
 
+    @needs_cuda
+    @pytest.mark.parametrize(
+        "decode_method",
+        (
+            _decode_sequential,
+            _decode_prefetch_frames,
+            _decode_prefetch_packets,
+            _decode_prefetch_packets_and_frames,
+        ),
+        ids=lambda f: f.__name__.removeprefix("_decode_"),
+    )
+    def test_matches_video_decoder_cuda(self, decode_method):
+        # On CUDA, the ColorConverter must share the decoder's DeviceInterface,
+        # so we bind it to the decoder. Blocks target the "ffmpeg" CUDA variant
+        # (self-contained, thread-movable frames); the reference VideoDecoder
+        # uses the same backend so the color conversion matches.
+        video = NASA_VIDEO
+        with set_cuda_backend("ffmpeg"):
+            got = self._to_frame_batch(decode_method(self, video.path, device="cuda"))
+            ref = VideoDecoder(video.path, device="cuda").get_all_frames()
+
+        assert got.data.device.type == "cuda"
+        assert got.data.shape == ref.data.shape
+        assert_frames_equal(got.data, ref.data)
+        torch.testing.assert_close(got.pts_seconds, ref.pts_seconds, atol=0, rtol=0)
+        torch.testing.assert_close(
+            got.duration_seconds, ref.duration_seconds, atol=0, rtol=0
+        )
+
 
 # Small helpers to avoid having to always specify the same skip marks and decode_fn
 def _jpeg_param(*values):