torchcodec 0.7.0__cp310-cp310-win_amd64.whl → 0.8.0__cp310-cp310-win_amd64.whl

torchcodec/_core/Encoder.h

@@ -57,7 +57,6 @@ class AudioEncoder {
   bool encodeWasCalled_ = false;
   int64_t lastEncodedAVFramePts_ = 0;
 };
-} // namespace facebook::torchcodec
 
 /* clang-format off */
 //
@@ -121,3 +120,54 @@ class AudioEncoder {
 //
 //
 /* clang-format on */
+
+class VideoEncoder {
+ public:
+  ~VideoEncoder();
+
+  // Rule of Five requires that we define copy and move
+  // constructors and assignment operators.
+  // Both are deleted because we have unique_ptr members
+  VideoEncoder(const VideoEncoder&) = delete;
+  VideoEncoder& operator=(const VideoEncoder&) = delete;
+
+  // Move assignment operator deleted since we have a const member
+  VideoEncoder(VideoEncoder&&) = default;
+  VideoEncoder& operator=(VideoEncoder&&) = delete;
+
+  VideoEncoder(
+      const torch::Tensor& frames,
+      int frameRate,
+      std::string_view fileName,
+      const VideoStreamOptions& videoStreamOptions);
+
+  void encode();
+
+ private:
+  void initializeEncoder(const VideoStreamOptions& videoStreamOptions);
+  UniqueAVFrame convertTensorToAVFrame(
+      const torch::Tensor& frame,
+      int frameIndex);
+  void encodeFrame(AutoAVPacket& autoAVPacket, const UniqueAVFrame& avFrame);
+  void flushBuffers();
+
+  UniqueEncodingAVFormatContext avFormatContext_;
+  UniqueAVCodecContext avCodecContext_;
+  AVStream* avStream_;
+  UniqueSwsContext swsContext_;
+
+  const torch::Tensor frames_;
+  int inFrameRate_;
+
+  int inWidth_ = -1;
+  int inHeight_ = -1;
+  AVPixelFormat inPixelFormat_ = AV_PIX_FMT_NONE;
+
+  int outWidth_ = -1;
+  int outHeight_ = -1;
+  AVPixelFormat outPixelFormat_ = AV_PIX_FMT_NONE;
+
+  bool encodeWasCalled_ = false;
+};
+
+} // namespace facebook::torchcodec

torchcodec/_core/FFMPEGCommon.cpp

@@ -8,6 +8,11 @@
 
 #include <c10/util/Exception.h>
 
+extern "C" {
+#include <libavfilter/avfilter.h>
+#include <libavfilter/buffersink.h>
+}
+
 namespace facebook::torchcodec {
 
 AutoAVPacket::AutoAVPacket() : avPacket_(av_packet_alloc()) {
@@ -56,6 +61,77 @@ int64_t getDuration(const UniqueAVFrame& avFrame) {
 #endif
 }
 
+void setDuration(const UniqueAVFrame& avFrame, int64_t duration) {
+#if LIBAVUTIL_VERSION_MAJOR < 58
+  avFrame->pkt_duration = duration;
+#else
+  avFrame->duration = duration;
+#endif
+}
+
+const int* getSupportedSampleRates(const AVCodec& avCodec) {
+  const int* supportedSampleRates = nullptr;
+#if LIBAVCODEC_VERSION_INT >= AV_VERSION_INT(61, 13, 100) // FFmpeg >= 7.1
+  int numSampleRates = 0;
+  int ret = avcodec_get_supported_config(
+      nullptr,
+      &avCodec,
+      AV_CODEC_CONFIG_SAMPLE_RATE,
+      0,
+      reinterpret_cast<const void**>(&supportedSampleRates),
+      &numSampleRates);
+  if (ret < 0 || supportedSampleRates == nullptr) {
+    // Return nullptr to skip validation in validateSampleRate.
+    return nullptr;
+  }
+#else
+  supportedSampleRates = avCodec.supported_samplerates;
+#endif
+  return supportedSampleRates;
+}
+
+const AVPixelFormat* getSupportedPixelFormats(const AVCodec& avCodec) {
+  const AVPixelFormat* supportedPixelFormats = nullptr;
+#if LIBAVCODEC_VERSION_INT >= AV_VERSION_INT(61, 13, 100) // FFmpeg >= 7.1
+  int numPixelFormats = 0;
+  int ret = avcodec_get_supported_config(
+      nullptr,
+      &avCodec,
+      AV_CODEC_CONFIG_PIX_FORMAT,
+      0,
+      reinterpret_cast<const void**>(&supportedPixelFormats),
+      &numPixelFormats);
+  if (ret < 0 || supportedPixelFormats == nullptr) {
+    TORCH_CHECK(false, "Couldn't get supported pixel formats from encoder.");
+  }
+#else
+  supportedPixelFormats = avCodec.pix_fmts;
+#endif
+  return supportedPixelFormats;
+}
+
+const AVSampleFormat* getSupportedOutputSampleFormats(const AVCodec& avCodec) {
+  const AVSampleFormat* supportedSampleFormats = nullptr;
+#if LIBAVCODEC_VERSION_INT >= AV_VERSION_INT(61, 13, 100) // FFmpeg >= 7.1
+  int numSampleFormats = 0;
+  int ret = avcodec_get_supported_config(
+      nullptr,
+      &avCodec,
+      AV_CODEC_CONFIG_SAMPLE_FORMAT,
+      0,
+      reinterpret_cast<const void**>(&supportedSampleFormats),
+      &numSampleFormats);
+  if (ret < 0 || supportedSampleFormats == nullptr) {
+    // Return nullptr to use default output format in
+    // findBestOutputSampleFormat.
+    return nullptr;
+  }
+#else
+  supportedSampleFormats = avCodec.sample_fmts;
+#endif
+  return supportedSampleFormats;
+}
+
 int getNumChannels(const UniqueAVFrame& avFrame) {
 #if LIBAVFILTER_VERSION_MAJOR > 8 || \
     (LIBAVFILTER_VERSION_MAJOR == 8 && LIBAVFILTER_VERSION_MINOR >= 44)
@@ -109,7 +185,32 @@ void setDefaultChannelLayout(UniqueAVFrame& avFrame, int numChannels) {
 }
 
 void validateNumChannels(const AVCodec& avCodec, int numChannels) {
-#if LIBAVFILTER_VERSION_MAJOR > 7 // FFmpeg > 4
+#if LIBAVCODEC_VERSION_INT >= AV_VERSION_INT(61, 13, 100) // FFmpeg >= 7.1
+  std::stringstream supportedNumChannels;
+  const AVChannelLayout* supportedLayouts = nullptr;
+  int numLayouts = 0;
+  int ret = avcodec_get_supported_config(
+      nullptr,
+      &avCodec,
+      AV_CODEC_CONFIG_CHANNEL_LAYOUT,
+      0,
+      reinterpret_cast<const void**>(&supportedLayouts),
+      &numLayouts);
+  if (ret < 0 || supportedLayouts == nullptr) {
+    // If we can't validate, we must assume it'll be fine. If not, FFmpeg will
+    // eventually raise.
+    return;
+  }
+  for (int i = 0; i < numLayouts; ++i) {
+    if (i > 0) {
+      supportedNumChannels << ", ";
+    }
+    supportedNumChannels << supportedLayouts[i].nb_channels;
+    if (numChannels == supportedLayouts[i].nb_channels) {
+      return;
+    }
+  }
+#elif LIBAVFILTER_VERSION_MAJOR > 7 // FFmpeg > 4
   if (avCodec.ch_layouts == nullptr) {
     // If we can't validate, we must assume it'll be fine. If not, FFmpeg will
     // eventually raise.
@@ -131,7 +232,7 @@ void validateNumChannels(const AVCodec& avCodec, int numChannels) {
     }
     supportedNumChannels << avCodec.ch_layouts[i].nb_channels;
   }
-#else
+#else // FFmpeg <= 4
   if (avCodec.channel_layouts == nullptr) {
     // can't validate, same as above.
     return;
@@ -298,6 +399,70 @@ SwrContext* createSwrContext(
   return swrContext;
 }
 
+AVFilterContext* createBuffersinkFilter(
+    AVFilterGraph* filterGraph,
+    enum AVPixelFormat outputFormat) {
+  const AVFilter* buffersink = avfilter_get_by_name("buffersink");
+  TORCH_CHECK(buffersink != nullptr, "Failed to get buffersink filter.");
+
+  AVFilterContext* sinkContext = nullptr;
+  int status;
+  const char* filterName = "out";
+
+  enum AVPixelFormat pix_fmts[] = {outputFormat, AV_PIX_FMT_NONE};
+
+// av_opt_set_int_list was replaced by av_opt_set_array() in FFmpeg 8.
+#if LIBAVUTIL_VERSION_MAJOR >= 60 // FFmpeg >= 8
+  // Output options like pixel_formats must be set before filter init
+  sinkContext =
+      avfilter_graph_alloc_filter(filterGraph, buffersink, filterName);
+  TORCH_CHECK(
+      sinkContext != nullptr, "Failed to allocate buffersink filter context.");
+
+  // When setting pix_fmts, only the first element is used, so nb_elems = 1
+  // AV_PIX_FMT_NONE acts as a terminator for the array in av_opt_set_int_list
+  status = av_opt_set_array(
+      sinkContext,
+      "pixel_formats",
+      AV_OPT_SEARCH_CHILDREN,
+      0, // start_elem
+      1, // nb_elems
+      AV_OPT_TYPE_PIXEL_FMT,
+      pix_fmts);
+  TORCH_CHECK(
+      status >= 0,
+      "Failed to set pixel format for buffersink filter: ",
+      getFFMPEGErrorStringFromErrorCode(status));
+
+  status = avfilter_init_str(sinkContext, nullptr);
+  TORCH_CHECK(
+      status >= 0,
+      "Failed to initialize buffersink filter: ",
+      getFFMPEGErrorStringFromErrorCode(status));
+#else // FFmpeg <= 7
+  // For older FFmpeg versions, create filter and then set options
+  status = avfilter_graph_create_filter(
+      &sinkContext, buffersink, filterName, nullptr, nullptr, filterGraph);
+  TORCH_CHECK(
+      status >= 0,
+      "Failed to create buffersink filter: ",
+      getFFMPEGErrorStringFromErrorCode(status));
+
+  status = av_opt_set_int_list(
+      sinkContext,
+      "pix_fmts",
+      pix_fmts,
+      AV_PIX_FMT_NONE,
+      AV_OPT_SEARCH_CHILDREN);
+  TORCH_CHECK(
+      status >= 0,
+      "Failed to set pixel formats for buffersink filter: ",
+      getFFMPEGErrorStringFromErrorCode(status));
+#endif
+
+  return sinkContext;
+}
+
 UniqueAVFrame convertAudioAVFrameSamples(
     const UniqueSwrContext& swrContext,
     const UniqueAVFrame& srcAVFrame,
@@ -418,4 +583,26 @@ AVIOContext* avioAllocContext(
       seek);
 }
 
+double ptsToSeconds(int64_t pts, const AVRational& timeBase) {
+  // To perform the multiplication before the division, av_q2d is not used
+  return static_cast<double>(pts) * timeBase.num / timeBase.den;
+}
+
+int64_t secondsToClosestPts(double seconds, const AVRational& timeBase) {
+  return static_cast<int64_t>(
+      std::round(seconds * timeBase.den / timeBase.num));
+}
+
+int64_t computeSafeDuration(
+    const AVRational& frameRate,
+    const AVRational& timeBase) {
+  if (frameRate.num <= 0 || frameRate.den <= 0 || timeBase.num <= 0 ||
+      timeBase.den <= 0) {
+    return 0;
+  } else {
+    return (static_cast<int64_t>(frameRate.den) * timeBase.den) /
+        (static_cast<int64_t>(timeBase.num) * frameRate.num);
+  }
+}
+
 } // namespace facebook::torchcodec

torchcodec/_core/FFMPEGCommon.h

@@ -12,6 +12,7 @@
 
 extern "C" {
 #include <libavcodec/avcodec.h>
+#include <libavcodec/bsf.h>
 #include <libavfilter/avfilter.h>
 #include <libavfilter/buffersrc.h>
 #include <libavformat/avformat.h>
@@ -86,6 +87,8 @@ using UniqueSwrContext =
     std::unique_ptr<SwrContext, Deleterp<SwrContext, void, swr_free>>;
 using UniqueAVAudioFifo = std::
     unique_ptr<AVAudioFifo, Deleter<AVAudioFifo, void, av_audio_fifo_free>>;
+using UniqueAVBSFContext =
+    std::unique_ptr<AVBSFContext, Deleterp<AVBSFContext, void, av_bsf_free>>;
 using UniqueAVBufferRef =
     std::unique_ptr<AVBufferRef, Deleterp<AVBufferRef, void, av_buffer_unref>>;
 using UniqueAVBufferSrcParameters = std::unique_ptr<
@@ -161,6 +164,11 @@ std::string getFFMPEGErrorStringFromErrorCode(int errorCode);
 // struct member representing duration has changed across the versions we
 // support.
 int64_t getDuration(const UniqueAVFrame& frame);
+void setDuration(const UniqueAVFrame& frame, int64_t duration);
+
+const int* getSupportedSampleRates(const AVCodec& avCodec);
+const AVSampleFormat* getSupportedOutputSampleFormats(const AVCodec& avCodec);
+const AVPixelFormat* getSupportedPixelFormats(const AVCodec& avCodec);
 
 int getNumChannels(const UniqueAVFrame& avFrame);
 int getNumChannels(const UniqueAVCodecContext& avCodecContext);
@@ -224,4 +232,14 @@ AVIOContext* avioAllocContext(
     AVIOWriteFunction write_packet,
     AVIOSeekFunction seek);
 
+double ptsToSeconds(int64_t pts, const AVRational& timeBase);
+int64_t secondsToClosestPts(double seconds, const AVRational& timeBase);
+int64_t computeSafeDuration(
+    const AVRational& frameRate,
+    const AVRational& timeBase);
+
+AVFilterContext* createBuffersinkFilter(
+    AVFilterGraph* filterGraph,
+    enum AVPixelFormat outputFormat);
+
 } // namespace facebook::torchcodec

torchcodec/_core/FilterGraph.cpp

@@ -5,6 +5,7 @@
 // LICENSE file in the root directory of this source tree.
 
 #include "src/torchcodec/_core/FilterGraph.h"
+#include "src/torchcodec/_core/FFMPEGCommon.h"
 
 extern "C" {
 #include <libavfilter/buffersink.h>
@@ -13,6 +14,28 @@ extern "C" {
 
 namespace facebook::torchcodec {
 
+FiltersContext::FiltersContext(
+    int inputWidth,
+    int inputHeight,
+    AVPixelFormat inputFormat,
+    AVRational inputAspectRatio,
+    int outputWidth,
+    int outputHeight,
+    AVPixelFormat outputFormat,
+    const std::string& filtergraphStr,
+    AVRational timeBase,
+    AVBufferRef* hwFramesCtx)
+    : inputWidth(inputWidth),
+      inputHeight(inputHeight),
+      inputFormat(inputFormat),
+      inputAspectRatio(inputAspectRatio),
+      outputWidth(outputWidth),
+      outputHeight(outputHeight),
+      outputFormat(outputFormat),
+      filtergraphStr(filtergraphStr),
+      timeBase(timeBase),
+      hwFramesCtx(hwFramesCtx) {}
+
 bool operator==(const AVRational& lhs, const AVRational& rhs) {
   return lhs.num == rhs.num && lhs.den == rhs.den;
 }
@@ -41,7 +64,6 @@ FilterGraph::FilterGraph(
   }
 
   const AVFilter* buffersrc = avfilter_get_by_name("buffer");
-  const AVFilter* buffersink = avfilter_get_by_name("buffersink");
 
   UniqueAVBufferSrcParameters srcParams(av_buffersrc_parameters_alloc());
   TORCH_CHECK(srcParams, "Failed to allocate buffersrc params");
@@ -71,26 +93,10 @@ FilterGraph::FilterGraph(
       "Failed to create filter graph : ",
       getFFMPEGErrorStringFromErrorCode(status));
 
-  status = avfilter_graph_create_filter(
-      &sinkContext_, buffersink, "out", nullptr, nullptr, filterGraph_.get());
-  TORCH_CHECK(
-      status >= 0,
-      "Failed to create filter graph: ",
-      getFFMPEGErrorStringFromErrorCode(status));
-
-  enum AVPixelFormat pix_fmts[] = {
-      filtersContext.outputFormat, AV_PIX_FMT_NONE};
-
-  status = av_opt_set_int_list(
-      sinkContext_,
-      "pix_fmts",
-      pix_fmts,
-      AV_PIX_FMT_NONE,
-      AV_OPT_SEARCH_CHILDREN);
+  sinkContext_ =
+      createBuffersinkFilter(filterGraph_.get(), filtersContext.outputFormat);
   TORCH_CHECK(
-      status >= 0,
-      "Failed to set output pixel formats: ",
-      getFFMPEGErrorStringFromErrorCode(status));
+      sinkContext_ != nullptr, "Failed to create and configure buffersink");
 
   UniqueAVFilterInOut outputs(avfilter_inout_alloc());
   UniqueAVFilterInOut inputs(avfilter_inout_alloc());
@@ -117,7 +123,8 @@ FilterGraph::FilterGraph(
   TORCH_CHECK(
       status >= 0,
       "Failed to parse filter description: ",
-      getFFMPEGErrorStringFromErrorCode(status));
+      getFFMPEGErrorStringFromErrorCode(status),
+      ", provided filters: " + filtersContext.filtergraphStr);
 
   status = avfilter_graph_config(filterGraph_.get(), nullptr);
   TORCH_CHECK(

torchcodec/_core/FilterGraph.h

@@ -19,11 +19,25 @@ struct FiltersContext {
   int outputWidth = 0;
   int outputHeight = 0;
   AVPixelFormat outputFormat = AV_PIX_FMT_NONE;
-
   std::string filtergraphStr;
   AVRational timeBase = {0, 0};
   UniqueAVBufferRef hwFramesCtx;
 
+  FiltersContext() = default;
+  FiltersContext(FiltersContext&&) = default;
+  FiltersContext& operator=(FiltersContext&&) = default;
+  FiltersContext(
+      int inputWidth,
+      int inputHeight,
+      AVPixelFormat inputFormat,
+      AVRational inputAspectRatio,
+      int outputWidth,
+      int outputHeight,
+      AVPixelFormat outputFormat,
+      const std::string& filtergraphStr,
+      AVRational timeBase,
+      AVBufferRef* hwFramesCtx = nullptr);
+
   bool operator==(const FiltersContext&) const;
   bool operator!=(const FiltersContext&) const;
 };

torchcodec/_core/Frame.cpp

@@ -8,24 +8,34 @@
 
 namespace facebook::torchcodec {
 
+FrameBatchOutput::FrameBatchOutput(
+    int64_t numFrames,
+    const FrameDims& outputDims,
+    const torch::Device& device)
+    : ptsSeconds(torch::empty({numFrames}, {torch::kFloat64})),
+      durationSeconds(torch::empty({numFrames}, {torch::kFloat64})) {
+  data = allocateEmptyHWCTensor(outputDims, device, numFrames);
+}
+
 torch::Tensor allocateEmptyHWCTensor(
-    int height,
-    int width,
-    torch::Device device,
+    const FrameDims& frameDims,
+    const torch::Device& device,
     std::optional<int> numFrames) {
   auto tensorOptions = torch::TensorOptions()
                            .dtype(torch::kUInt8)
                            .layout(torch::kStrided)
                            .device(device);
-  TORCH_CHECK(height > 0, "height must be > 0, got: ", height);
-  TORCH_CHECK(width > 0, "width must be > 0, got: ", width);
+  TORCH_CHECK(
+      frameDims.height > 0, "height must be > 0, got: ", frameDims.height);
+  TORCH_CHECK(frameDims.width > 0, "width must be > 0, got: ", frameDims.width);
   if (numFrames.has_value()) {
     auto numFramesValue = numFrames.value();
     TORCH_CHECK(
         numFramesValue >= 0, "numFrames must be >= 0, got: ", numFramesValue);
-    return torch::empty({numFramesValue, height, width, 3}, tensorOptions);
+    return torch::empty(
+        {numFramesValue, frameDims.height, frameDims.width, 3}, tensorOptions);
   } else {
-    return torch::empty({height, width, 3}, tensorOptions);
+    return torch::empty({frameDims.height, frameDims.width, 3}, tensorOptions);
   }
 }
 

torchcodec/_core/Frame.h

@@ -13,6 +13,15 @@
 
 namespace facebook::torchcodec {
 
+struct FrameDims {
+  int height = 0;
+  int width = 0;
+
+  FrameDims() = default;
+
+  FrameDims(int h, int w) : height(h), width(w) {}
+};
+
 // All public video decoding entry points return either a FrameOutput or a
 // FrameBatchOutput.
 // They are the equivalent of the user-facing Frame and FrameBatch classes in
@@ -34,10 +43,10 @@ struct FrameBatchOutput {
   torch::Tensor ptsSeconds; // 1D of shape (N,)
   torch::Tensor durationSeconds; // 1D of shape (N,)
 
-  explicit FrameBatchOutput(
+  FrameBatchOutput(
       int64_t numFrames,
-      const VideoStreamOptions& videoStreamOptions,
-      const StreamMetadata& streamMetadata);
+      const FrameDims& outputDims,
+      const torch::Device& device);
 };
 
 struct AudioFramesOutput {
@@ -49,70 +58,15 @@ struct AudioFramesOutput {
 // FRAME TENSOR ALLOCATION APIs
 // --------------------------------------------------------------------------
 
-// Note [Frame Tensor allocation and height and width]
+// Note [Frame Tensor allocation]
 //
 // We always allocate [N]HWC tensors. The low-level decoding functions all
 // assume HWC tensors, since this is what FFmpeg natively handles. It's up to
 // the high-level decoding entry-points to permute that back to CHW, by calling
 // maybePermuteHWC2CHW().
-//
-// Also, importantly, the way we figure out the the height and width of the
-// output frame tensor varies, and depends on the decoding entry-point. In
-// *decreasing order of accuracy*, we use the following sources for determining
-// height and width:
-// - getHeightAndWidthFromResizedAVFrame(). This is the height and width of the
-//   AVframe, *post*-resizing. This is only used for single-frame decoding APIs,
-//   on CPU, with filtergraph.
-// - getHeightAndWidthFromOptionsOrAVFrame(). This is the height and width from
-//   the user-specified options if they exist, or the height and width of the
-//   AVFrame *before* it is resized. In theory, i.e. if there are no bugs within
-//   our code or within FFmpeg code, this should be exactly the same as
-//   getHeightAndWidthFromResizedAVFrame(). This is used by single-frame
-//   decoding APIs, on CPU with swscale, and on GPU.
-// - getHeightAndWidthFromOptionsOrMetadata(). This is the height and width from
-//   the user-specified options if they exist, or the height and width form the
-//   stream metadata, which itself got its value from the CodecContext, when the
-//   stream was added. This is used by batch decoding APIs, for both GPU and
-//   CPU.
-//
-// The source of truth for height and width really is the (resized) AVFrame: it
-// comes from the decoded ouptut of FFmpeg. The info from the metadata (i.e.
-// from the CodecContext) may not be as accurate. However, the AVFrame is only
-// available late in the call stack, when the frame is decoded, while the
-// CodecContext is available early when a stream is added. This is why we use
-// the CodecContext for pre-allocating batched output tensors (we could
-// pre-allocate those only once we decode the first frame to get the info frame
-// the AVFrame, but that's a more complex logic).
-//
-// Because the sources for height and width may disagree, we may end up with
-// conflicts: e.g. if we pre-allocate a batch output tensor based on the
-// metadata info, but the decoded AVFrame has a different height and width.
-// it is very important to check the height and width assumptions where the
-// tensors memory is used/filled in order to avoid segfaults.
-
-struct FrameDims {
-  int height;
-  int width;
-
-  FrameDims(int h, int w) : height(h), width(w) {}
-};
-
-// There's nothing preventing you from calling this on a non-resized frame, but
-// please don't.
-FrameDims getHeightAndWidthFromResizedAVFrame(const AVFrame& resizedAVFrame);
-
-FrameDims getHeightAndWidthFromOptionsOrMetadata(
-    const VideoStreamOptions& videoStreamOptions,
-    const StreamMetadata& streamMetadata);
-
-FrameDims getHeightAndWidthFromOptionsOrAVFrame(
-    const VideoStreamOptions& videoStreamOptions,
-    const UniqueAVFrame& avFrame);
-
 torch::Tensor allocateEmptyHWCTensor(
-    int height,
-    int width,
-    torch::Device device,
+    const FrameDims& frameDims,
+    const torch::Device& device,
     std::optional<int> numFrames = std::nullopt);
 
 } // namespace facebook::torchcodec

torchcodec/_core/Metadata.h

@@ -44,8 +44,8 @@ struct StreamMetadata {
   std::optional<int64_t> numFramesFromContent;
 
   // Video-only fields derived from the AVCodecContext.
-  std::optional<int64_t> width;
-  std::optional<int64_t> height;
+  std::optional<int> width;
+  std::optional<int> height;
   std::optional<AVRational> sampleAspectRatio;
 
   // Audio-only fields

torchcodec/_core/NVDECCache.cpp

@@ -0,0 +1,70 @@
+// Copyright (c) Meta Platforms, Inc. and affiliates.
+// All rights reserved.
+//
+// This source code is licensed under the BSD-style license found in the
+// LICENSE file in the root directory of this source tree.
+
+#include <torch/types.h>
+#include <mutex>
+
+#include "src/torchcodec/_core/FFMPEGCommon.h"
+#include "src/torchcodec/_core/NVDECCache.h"
+
+#include <cuda_runtime.h> // For cudaGetDevice
+
+extern "C" {
+#include <libavutil/hwcontext_cuda.h>
+#include <libavutil/pixdesc.h>
+}
+
+namespace facebook::torchcodec {
+
+NVDECCache& NVDECCache::getCache(int deviceIndex) {
+  const int MAX_CUDA_GPUS = 128;
+  TORCH_CHECK(
+      deviceIndex >= -1 && deviceIndex < MAX_CUDA_GPUS,
+      "Invalid device index = ",
+      deviceIndex);
+  static NVDECCache cacheInstances[MAX_CUDA_GPUS];
+  if (deviceIndex == -1) {
+    // TODO NVDEC P3: Unify with existing getNonNegativeDeviceIndex()
+    TORCH_CHECK(
+        cudaGetDevice(&deviceIndex) == cudaSuccess,
+        "Failed to get current CUDA device.");
+  }
+  return cacheInstances[deviceIndex];
+}
+
+UniqueCUvideodecoder NVDECCache::getDecoder(CUVIDEOFORMAT* videoFormat) {
+  CacheKey key(videoFormat);
+  std::lock_guard<std::mutex> lock(cacheLock_);
+
+  auto it = cache_.find(key);
+  if (it != cache_.end()) {
+    auto decoder = std::move(it->second);
+    cache_.erase(it);
+    return decoder;
+  }
+
+  return nullptr;
+}
+
+bool NVDECCache::returnDecoder(
+    CUVIDEOFORMAT* videoFormat,
+    UniqueCUvideodecoder decoder) {
+  if (!decoder) {
+    return false;
+  }
+
+  CacheKey key(videoFormat);
+  std::lock_guard<std::mutex> lock(cacheLock_);
+
+  if (cache_.size() >= MAX_CACHE_SIZE) {
+    return false;
+  }
+
+  cache_[key] = std::move(decoder);
+  return true;
+}
+
+} // namespace facebook::torchcodec