Pass AVFrames by reference, not by smart pointer

src/torchcodec/_core/BetaCudaDeviceInterface.cpp

@@ -752,7 +752,8 @@ UniqueAVFrame BetaCudaDeviceInterface::convert_cuda_frame_to_av_frame(
   // Note that we used to rely on videoFormat_.frame_rate for this, but that
   // proved less accurate than FFmpeg.
   set_duration(
-      av_frame, compute_safe_duration(frame_rate_avg_from_ffmpeg_, time_base_));
+      *av_frame,
+      compute_safe_duration(frame_rate_avg_from_ffmpeg_, time_base_));
 
   // We need to assign the frame colorspace. This is crucial for proper color
   // conversion. NVCUVID stores that in the matrix_coefficients field, but
@@ -824,7 +825,7 @@ void BetaCudaDeviceInterface::flush() {
 }
 
 UniqueAVFrame BetaCudaDeviceInterface::transfer_cpu_frame_to_gpu(
-    UniqueAVFrame& cpu_frame,
+    const AVFrame& cpu_frame,
     AVPixelFormat target_pix_fmt) {
   // This is called in the context of the CPU fallback: the frame was decoded on
   // the CPU, and in this function we convert that frame into NV12 or P016
@@ -838,15 +839,14 @@ UniqueAVFrame BetaCudaDeviceInterface::transfer_cpu_frame_to_gpu(
   // (rounded up) width and height, even if the original CPU frame had odd
   // dimensions.
 
-  STD_TORCH_CHECK(cpu_frame != nullptr, "CPU frame cannot be null");
   // NV12 = 1 byte per sample, P016 = 2 bytes per sample
   STD_TORCH_CHECK(
       target_pix_fmt == AV_PIX_FMT_NV12 || target_pix_fmt == AV_PIX_FMT_P016LE,
       "targetPixFmt must be NV12 or P016LE");
   int bytes_per_sample = (target_pix_fmt == AV_PIX_FMT_P016LE) ? 2 : 1;
 
-  int width = cpu_frame->width;
-  int height = cpu_frame->height;
+  int width = cpu_frame.width;
+  int height = cpu_frame.height;
   int even_width = round_up_to_even(width);
   int even_height = round_up_to_even(height);
 
@@ -868,8 +868,8 @@ UniqueAVFrame BetaCudaDeviceInterface::transfer_cpu_frame_to_gpu(
   SwsConfig sws_config(
       width,
       height,
-      static_cast<AVPixelFormat>(cpu_frame->format),
-      cpu_frame->colorspace,
+      static_cast<AVPixelFormat>(cpu_frame.format),
+      cpu_frame.colorspace,
       even_width,
       even_height,
       target_pix_fmt);
@@ -881,8 +881,8 @@ UniqueAVFrame BetaCudaDeviceInterface::transfer_cpu_frame_to_gpu(
 
   int converted_height = sws_scale(
       sws_context_.get(),
-      cpu_frame->data,
-      cpu_frame->linesize,
+      cpu_frame.data,
+      cpu_frame.linesize,
       0,
       height,
       intermediate_cpu_frame->data,
@@ -944,7 +944,7 @@ UniqueAVFrame BetaCudaDeviceInterface::transfer_cpu_frame_to_gpu(
       "Failed to copy UV plane to GPU: ",
       cudaGetErrorString(err));
 
-  ret = av_frame_copy_props(gpu_frame.get(), cpu_frame.get());
+  ret = av_frame_copy_props(gpu_frame.get(), &cpu_frame);
   STD_TORCH_CHECK(
       ret >= 0,
       "Failed to copy frame properties: ",
@@ -967,7 +967,7 @@ UniqueAVFrame BetaCudaDeviceInterface::transfer_cpu_frame_to_gpu(
 }
 
 void BetaCudaDeviceInterface::convert_av_frame_to_frame_output(
-    UniqueAVFrame& av_frame,
+    const AVFrame& av_frame,
     FrameOutput& frame_output,
     std::optional<torch::stable::Tensor> pre_allocated_output_tensor) {
   if (cpu_fallback_) {
@@ -979,7 +979,7 @@ void BetaCudaDeviceInterface::convert_av_frame_to_frame_output(
     // 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>(av_frame->format));
+        av_pix_fmt_desc_get(static_cast<AVPixelFormat>(av_frame.format));
     bool is444 = desc && desc->log2_chroma_w == 0 && desc->log2_chroma_h == 0;
     if (is444) {
       FrameOutput cpu_frame_output;
@@ -1001,28 +1001,29 @@ void BetaCudaDeviceInterface::convert_av_frame_to_frame_output(
 
   // Capture original dimensions before transferCpuFrameToGpu()
   // may round them up to even.
-  FrameDims original_dims(av_frame->height, av_frame->width);
+  FrameDims original_dims(av_frame.height, av_frame.width);
 
-  UniqueAVFrame gpu_frame;
+  // On the CPU fallback we own the GPU frame we just created; otherwise the
+  // input frame is already what we need, and we only observe it.
+  UniqueAVFrame transferred_frame;
   if (cpu_fallback_) {
     AVPixelFormat target_pix_fmt = (output_dtype_ == OutputDtype::FLOAT32)
         ? AV_PIX_FMT_P016LE
         : AV_PIX_FMT_NV12;
-    gpu_frame = transfer_cpu_frame_to_gpu(av_frame, target_pix_fmt);
-  } else {
-    gpu_frame = std::move(av_frame);
+    transferred_frame = transfer_cpu_frame_to_gpu(av_frame, target_pix_fmt);
   }
+  const AVFrame& gpu_frame = cpu_fallback_ ? *transferred_frame : av_frame;
 
   STD_TORCH_CHECK(
-      gpu_frame->format == AV_PIX_FMT_NV12 ||
-          gpu_frame->format == AV_PIX_FMT_P016LE,
+      gpu_frame.format == AV_PIX_FMT_NV12 ||
+          gpu_frame.format == AV_PIX_FMT_P016LE,
       "Expected NV12 or P016LE format frame");
 
   cudaStream_t nvdec_stream = get_current_cuda_stream(device_.index());
 
   auto convert_frame = [&](std::optional<torch::stable::Tensor> pre_alloc)
       -> torch::stable::Tensor {
-    bool is_p016 = (gpu_frame->format == AV_PIX_FMT_P016LE);
+    bool is_p016 = (gpu_frame.format == AV_PIX_FMT_P016LE);
     int bit_depth = 8;
     if (is_p016) {
       bit_depth = cpu_fallback_

src/torchcodec/_core/BetaCudaDeviceInterface.h

@@ -52,7 +52,7 @@ class BetaCudaDeviceInterface : public DeviceInterface {
       OutputDtype requested_dtype) const override;
 
   void convert_av_frame_to_frame_output(
-      UniqueAVFrame& av_frame,
+      const AVFrame& av_frame,
       FrameOutput& frame_output,
       std::optional<torch::stable::Tensor> pre_allocated_output_tensor)
       override;
@@ -92,7 +92,7 @@ class BetaCudaDeviceInterface : public DeviceInterface {
       const CUVIDPARSERDISPINFO& disp_info);
 
   UniqueAVFrame transfer_cpu_frame_to_gpu(
-      UniqueAVFrame& cpu_frame,
+      const AVFrame& cpu_frame,
       AVPixelFormat target_pix_fmt);
 
   void apply_rotation(

src/torchcodec/_core/ColorConverter.cpp

@@ -45,7 +45,7 @@ ColorConverter::ColorConverter(
       /*resized_output_dims=*/std::nullopt);
 }
 
-torch::stable::Tensor ColorConverter::convert(UniqueAVFrame& av_frame) {
+torch::stable::Tensor ColorConverter::convert(const AVFrame& av_frame) {
   FrameOutput frame_output;
   device_interface_->convert_av_frame_to_frame_output(
       av_frame, frame_output, std::nullopt);

src/torchcodec/_core/ColorConverter.h

@@ -21,7 +21,7 @@ class FORCE_PUBLIC_VISIBILITY ColorConverter {
       const StableDevice& device = StableDevice(kStableCPU),
       std::string_view device_variant = "default");
 
-  torch::stable::Tensor convert(UniqueAVFrame& av_frame);
+  torch::stable::Tensor convert(const AVFrame& av_frame);
 
  private:
   std::unique_ptr<DeviceInterface> device_interface_;

src/torchcodec/_core/CpuDeviceInterface.cpp

@@ -200,7 +200,7 @@ ColorConversionLibrary CpuDeviceInterface::get_color_conversion_library(
 }
 
 void CpuDeviceInterface::convert_av_frame_to_frame_output(
-    UniqueAVFrame& av_frame,
+    const AVFrame& av_frame,
     FrameOutput& frame_output,
     std::optional<torch::stable::Tensor> pre_allocated_output_tensor) {
   STD_TORCH_CHECK(initialized_, "CpuDeviceInterface was not initialized.");
@@ -223,7 +223,7 @@ void CpuDeviceInterface::convert_av_frame_to_frame_output(
 // Dimension order of the preAllocatedOutputTensor must be HWC, regardless of
 // `dimension_order` parameter. It's up to callers to re-shape it if needed.
 void CpuDeviceInterface::convert_video_av_frame_to_frame_output(
-    UniqueAVFrame& av_frame,
+    const AVFrame& av_frame,
     FrameOutput& frame_output,
     std::optional<torch::stable::Tensor> pre_allocated_output_tensor) {
   // Note that we ignore the dimensions from the metadata; we don't even bother
@@ -239,7 +239,7 @@ void CpuDeviceInterface::convert_video_av_frame_to_frame_output(
   // Both cases cause problems for our batch APIs, as we allocate
   // FrameBatchOutputs based on the the stream metadata. But single-frame APIs
   // can still work in such situations, so they should.
-  auto input_dims = FrameDims(av_frame->height, av_frame->width);
+  auto input_dims = FrameDims(av_frame.height, av_frame.width);
   auto output_dims = resized_output_dims_.value_or(input_dims);
 
   if (pre_allocated_output_tensor.has_value()) {
@@ -264,12 +264,12 @@ void CpuDeviceInterface::convert_video_av_frame_to_frame_output(
         pre_allocated_output_tensor.value_or(allocate_empty_hwc_tensor(
             output_dims, kStableCPU, video_stream_options_.output_dtype));
 
-    auto av_frame_format = static_cast<AVPixelFormat>(av_frame->format);
+    auto av_frame_format = static_cast<AVPixelFormat>(av_frame.format);
     SwsConfig sws_config(
-        av_frame->width,
-        av_frame->height,
+        av_frame.width,
+        av_frame.height,
         av_frame_format,
-        av_frame->colorspace,
+        av_frame.colorspace,
         output_dims.width,
         output_dims.height,
         output_pixel_format_);
@@ -326,15 +326,15 @@ void CpuDeviceInterface::convert_video_av_frame_to_frame_output(
 
 torch::stable::Tensor
 CpuDeviceInterface::convert_av_frame_to_tensor_using_filter_graph(
-    const UniqueAVFrame& av_frame,
+    const AVFrame& av_frame,
     const FrameDims& output_dims) {
-  auto av_frame_format = static_cast<AVPixelFormat>(av_frame->format);
+  auto av_frame_format = static_cast<AVPixelFormat>(av_frame.format);
 
   FiltersConfig filters_config(
-      av_frame->width,
-      av_frame->height,
+      av_frame.width,
+      av_frame.height,
       av_frame_format,
-      av_frame->sample_aspect_ratio,
+      av_frame.sample_aspect_ratio,
       output_dims.width,
       output_dims.height,
       output_pixel_format_,
@@ -346,17 +346,17 @@ CpuDeviceInterface::convert_av_frame_to_tensor_using_filter_graph(
         std::make_unique<FilterGraph>(filters_config, video_stream_options_);
     prev_filters_config_ = std::move(filters_config);
   }
-  return rgb_av_frame_to_tensor(filter_graph_->convert(av_frame));
+  return rgb_av_frame_to_tensor(*filter_graph_->convert(av_frame));
 }
 
 void CpuDeviceInterface::convert_audio_av_frame_to_frame_output(
-    UniqueAVFrame& src_av_frame,
+    const AVFrame& src_av_frame,
     FrameOutput& frame_output) {
   AVSampleFormat src_sample_format =
-      static_cast<AVSampleFormat>(src_av_frame->format);
+      static_cast<AVSampleFormat>(src_av_frame.format);
   AVSampleFormat out_sample_format = AV_SAMPLE_FMT_FLTP;
 
-  int src_sample_rate = src_av_frame->sample_rate;
+  int src_sample_rate = src_av_frame.sample_rate;
   int out_sample_rate =
       audio_stream_options_.sample_rate.value_or(src_sample_rate);
 
@@ -397,10 +397,9 @@ void CpuDeviceInterface::convert_audio_av_frame_to_frame_output(
         out_sample_rate,
         out_num_channels);
   }
-  const UniqueAVFrame& av_frame =
-      must_convert ? converted_av_frame : src_av_frame;
+  const AVFrame& av_frame = must_convert ? *converted_av_frame : src_av_frame;
 
-  AVSampleFormat format = static_cast<AVSampleFormat>(av_frame->format);
+  AVSampleFormat format = static_cast<AVSampleFormat>(av_frame.format);
   STD_TORCH_CHECK(
       format == out_sample_format,
       "Something went wrong, the frame didn't get converted to the desired format. ",
@@ -419,7 +418,7 @@ void CpuDeviceInterface::convert_audio_av_frame_to_frame_output(
       num_channels,
       " instead.");
 
-  auto num_samples = av_frame->nb_samples;
+  auto num_samples = av_frame.nb_samples;
 
   frame_output.data = torch::stable::empty({num_channels, num_samples});
 
@@ -431,7 +430,7 @@ void CpuDeviceInterface::convert_audio_av_frame_to_frame_output(
          ++channel, output_channel_data += num_bytes_per_channel) {
       std::memcpy(
           output_channel_data,
-          av_frame->extended_data[channel],
+          av_frame.extended_data[channel],
           num_bytes_per_channel);
     }
   }

src/torchcodec/_core/CpuDeviceInterface.h

@@ -41,7 +41,7 @@ class CpuDeviceInterface : public DeviceInterface {
       override;
 
   void convert_av_frame_to_frame_output(
-      UniqueAVFrame& av_frame,
+      const AVFrame& av_frame,
       FrameOutput& frame_output,
       std::optional<torch::stable::Tensor> pre_allocated_output_tensor)
       override;
@@ -59,16 +59,16 @@ class CpuDeviceInterface : public DeviceInterface {
 
  private:
   void convert_audio_av_frame_to_frame_output(
-      UniqueAVFrame& src_av_frame,
+      const AVFrame& src_av_frame,
       FrameOutput& frame_output);
 
   void convert_video_av_frame_to_frame_output(
-      UniqueAVFrame& av_frame,
+      const AVFrame& av_frame,
       FrameOutput& frame_output,
       std::optional<torch::stable::Tensor> pre_allocated_output_tensor);
 
   torch::stable::Tensor convert_av_frame_to_tensor_using_filter_graph(
-      const UniqueAVFrame& av_frame,
+      const AVFrame& av_frame,
       const FrameDims& output_dims);
 
   ColorConversionLibrary get_color_conversion_library(