@@ -376,7 +376,7 @@ void BetaCudaDeviceInterface::initialize_video(
376376
377377void BetaCudaDeviceInterface::send_seqhdr_packet () {
378378 // This must be called at parser initialization, and after each flush.
379- // See TODO for details.
379+ // See https://github.com/meta-pytorch/torchcodec/pull/1503 for details.
380380 // FFmpeg's nvcuviddec.c does the same thing (not the nvdec.c one, because it
381381 // doesn't rely on the nvcuvid parser):
382382 // -
@@ -752,7 +752,8 @@ UniqueAVFrame BetaCudaDeviceInterface::convert_cuda_frame_to_av_frame(
752752 // Note that we used to rely on videoFormat_.frame_rate for this, but that
753753 // proved less accurate than FFmpeg.
754754 set_duration (
755- av_frame, compute_safe_duration (frame_rate_avg_from_ffmpeg_, time_base_));
755+ *av_frame,
756+ compute_safe_duration (frame_rate_avg_from_ffmpeg_, time_base_));
756757
757758 // We need to assign the frame colorspace. This is crucial for proper color
758759 // conversion. NVCUVID stores that in the matrix_coefficients field, but
@@ -824,7 +825,7 @@ void BetaCudaDeviceInterface::flush() {
824825}
825826
826827UniqueAVFrame BetaCudaDeviceInterface::transfer_cpu_frame_to_gpu (
827- UniqueAVFrame & cpu_frame,
828+ const AVFrame & cpu_frame,
828829 AVPixelFormat target_pix_fmt) {
829830 // This is called in the context of the CPU fallback: the frame was decoded on
830831 // 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(
838839 // (rounded up) width and height, even if the original CPU frame had odd
839840 // dimensions.
840841
841- STD_TORCH_CHECK (cpu_frame != nullptr , " CPU frame cannot be null" );
842842 // NV12 = 1 byte per sample, P016 = 2 bytes per sample
843843 STD_TORCH_CHECK (
844844 target_pix_fmt == AV_PIX_FMT_NV12 || target_pix_fmt == AV_PIX_FMT_P016LE ,
845845 " targetPixFmt must be NV12 or P016LE" );
846846 int bytes_per_sample = (target_pix_fmt == AV_PIX_FMT_P016LE ) ? 2 : 1 ;
847847
848- int width = cpu_frame-> width ;
849- int height = cpu_frame-> height ;
848+ int width = cpu_frame. width ;
849+ int height = cpu_frame. height ;
850850 int even_width = round_up_to_even (width);
851851 int even_height = round_up_to_even (height);
852852
@@ -868,8 +868,8 @@ UniqueAVFrame BetaCudaDeviceInterface::transfer_cpu_frame_to_gpu(
868868 SwsConfig sws_config (
869869 width,
870870 height,
871- static_cast <AVPixelFormat>(cpu_frame-> format ),
872- cpu_frame-> colorspace ,
871+ static_cast <AVPixelFormat>(cpu_frame. format ),
872+ cpu_frame. colorspace ,
873873 even_width,
874874 even_height,
875875 target_pix_fmt);
@@ -881,8 +881,8 @@ UniqueAVFrame BetaCudaDeviceInterface::transfer_cpu_frame_to_gpu(
881881
882882 int converted_height = sws_scale (
883883 sws_context_.get (),
884- cpu_frame-> data ,
885- cpu_frame-> linesize ,
884+ cpu_frame. data ,
885+ cpu_frame. linesize ,
886886 0 ,
887887 height,
888888 intermediate_cpu_frame->data ,
@@ -944,7 +944,7 @@ UniqueAVFrame BetaCudaDeviceInterface::transfer_cpu_frame_to_gpu(
944944 " Failed to copy UV plane to GPU: " ,
945945 cudaGetErrorString (err));
946946
947- ret = av_frame_copy_props (gpu_frame.get (), cpu_frame. get () );
947+ ret = av_frame_copy_props (gpu_frame.get (), & cpu_frame);
948948 STD_TORCH_CHECK (
949949 ret >= 0 ,
950950 " Failed to copy frame properties: " ,
@@ -967,7 +967,7 @@ UniqueAVFrame BetaCudaDeviceInterface::transfer_cpu_frame_to_gpu(
967967}
968968
969969void BetaCudaDeviceInterface::convert_av_frame_to_frame_output (
970- UniqueAVFrame & av_frame,
970+ const AVFrame & av_frame,
971971 FrameOutput& frame_output,
972972 std::optional<torch::stable::Tensor> pre_allocated_output_tensor) {
973973 if (cpu_fallback_) {
@@ -979,7 +979,7 @@ void BetaCudaDeviceInterface::convert_av_frame_to_frame_output(
979979 // do the color conversion on the CPU and then send the full RGB frame to
980980 // the GPU.
981981 const AVPixFmtDescriptor* desc =
982- av_pix_fmt_desc_get (static_cast <AVPixelFormat>(av_frame-> format ));
982+ av_pix_fmt_desc_get (static_cast <AVPixelFormat>(av_frame. format ));
983983 bool is444 = desc && desc->log2_chroma_w == 0 && desc->log2_chroma_h == 0 ;
984984 if (is444) {
985985 FrameOutput cpu_frame_output;
@@ -1001,28 +1001,29 @@ void BetaCudaDeviceInterface::convert_av_frame_to_frame_output(
10011001
10021002 // Capture original dimensions before transferCpuFrameToGpu()
10031003 // may round them up to even.
1004- FrameDims original_dims (av_frame-> height , av_frame-> width );
1004+ FrameDims original_dims (av_frame. height , av_frame. width );
10051005
1006- UniqueAVFrame gpu_frame;
1006+ // On the CPU fallback we own the GPU frame we just created; otherwise the
1007+ // input frame is already what we need, and we only observe it.
1008+ UniqueAVFrame transferred_frame;
10071009 if (cpu_fallback_) {
10081010 AVPixelFormat target_pix_fmt = (output_dtype_ == OutputDtype::FLOAT32 )
10091011 ? AV_PIX_FMT_P016LE
10101012 : AV_PIX_FMT_NV12 ;
1011- gpu_frame = transfer_cpu_frame_to_gpu (av_frame, target_pix_fmt);
1012- } else {
1013- gpu_frame = std::move (av_frame);
1013+ transferred_frame = transfer_cpu_frame_to_gpu (av_frame, target_pix_fmt);
10141014 }
1015+ const AVFrame& gpu_frame = cpu_fallback_ ? *transferred_frame : av_frame;
10151016
10161017 STD_TORCH_CHECK (
1017- gpu_frame-> format == AV_PIX_FMT_NV12 ||
1018- gpu_frame-> format == AV_PIX_FMT_P016LE ,
1018+ gpu_frame. format == AV_PIX_FMT_NV12 ||
1019+ gpu_frame. format == AV_PIX_FMT_P016LE ,
10191020 " Expected NV12 or P016LE format frame" );
10201021
10211022 cudaStream_t nvdec_stream = get_current_cuda_stream (device_.index ());
10221023
10231024 auto convert_frame = [&](std::optional<torch::stable::Tensor> pre_alloc)
10241025 -> torch::stable::Tensor {
1025- bool is_p016 = (gpu_frame-> format == AV_PIX_FMT_P016LE );
1026+ bool is_p016 = (gpu_frame. format == AV_PIX_FMT_P016LE );
10261027 int bit_depth = 8 ;
10271028 if (is_p016) {
10281029 bit_depth = cpu_fallback_
0 commit comments