torchcodec 0.7.0__cp312-cp312-win_amd64.whl → 0.8.1__cp312-cp312-win_amd64.whl

torchcodec/_core/Encoder.cpp

@@ -33,21 +33,22 @@ torch::Tensor validateSamples(const torch::Tensor& samples) {
 }
 
 void validateSampleRate(const AVCodec& avCodec, int sampleRate) {
-  if (avCodec.supported_samplerates == nullptr) {
+  const int* supportedSampleRates = getSupportedSampleRates(avCodec);
+  if (supportedSampleRates == nullptr) {
     return;
   }
 
-  for (auto i = 0; avCodec.supported_samplerates[i] != 0; ++i) {
-    if (sampleRate == avCodec.supported_samplerates[i]) {
+  for (auto i = 0; supportedSampleRates[i] != 0; ++i) {
+    if (sampleRate == supportedSampleRates[i]) {
       return;
     }
   }
   std::stringstream supportedRates;
-  for (auto i = 0; avCodec.supported_samplerates[i] != 0; ++i) {
+  for (auto i = 0; supportedSampleRates[i] != 0; ++i) {
     if (i > 0) {
       supportedRates << ", ";
     }
-    supportedRates << avCodec.supported_samplerates[i];
+    supportedRates << supportedSampleRates[i];
   }
 
   TORCH_CHECK(
@@ -73,19 +74,22 @@ static const std::vector<AVSampleFormat> preferredFormatsOrder = {
     AV_SAMPLE_FMT_U8};
 
 AVSampleFormat findBestOutputSampleFormat(const AVCodec& avCodec) {
+  const AVSampleFormat* supportedSampleFormats =
+      getSupportedOutputSampleFormats(avCodec);
+
   // Find a sample format that the encoder supports. We prefer using FLT[P],
   // since this is the format of the input samples. If FLTP isn't supported
   // then we'll need to convert the AVFrame's format. Our heuristic is to encode
   // into the format with the highest resolution.
-  if (avCodec.sample_fmts == nullptr) {
+  if (supportedSampleFormats == nullptr) {
     // Can't really validate anything in this case, best we can do is hope that
     // FLTP is supported by the encoder. If not, FFmpeg will raise.
     return AV_SAMPLE_FMT_FLTP;
   }
 
   for (AVSampleFormat preferredFormat : preferredFormatsOrder) {
-    for (int i = 0; avCodec.sample_fmts[i] != -1; ++i) {
-      if (avCodec.sample_fmts[i] == preferredFormat) {
+    for (int i = 0; supportedSampleFormats[i] != -1; ++i) {
+      if (supportedSampleFormats[i] == preferredFormat) {
         return preferredFormat;
       }
     }
@@ -93,7 +97,7 @@ AVSampleFormat findBestOutputSampleFormat(const AVCodec& avCodec) {
   // We should always find a match in preferredFormatsOrder, so we should always
   // return earlier. But in the event that a future FFmpeg version defines an
   // additional sample format that isn't in preferredFormatsOrder, we fallback:
-  return avCodec.sample_fmts[0];
+  return supportedSampleFormats[0];
 }
 
 } // namespace
@@ -511,4 +515,337 @@ void AudioEncoder::flushBuffers() {
 
   encodeFrame(autoAVPacket, UniqueAVFrame(nullptr));
 }
+
+namespace {
+
+torch::Tensor validateFrames(const torch::Tensor& frames) {
+  TORCH_CHECK(
+      frames.dtype() == torch::kUInt8,
+      "frames must have uint8 dtype, got ",
+      frames.dtype());
+  TORCH_CHECK(
+      frames.dim() == 4,
+      "frames must have 4 dimensions (N, C, H, W), got ",
+      frames.dim());
+  TORCH_CHECK(
+      frames.sizes()[1] == 3,
+      "frame must have 3 channels (R, G, B), got ",
+      frames.sizes()[1]);
+  // TODO-VideoEncoder: Investigate if non-contiguous frames can be accepted
+  return frames.contiguous();
+}
+
+} // namespace
+
+VideoEncoder::~VideoEncoder() {
+  // TODO-VideoEncoder: Unify destructor with ~AudioEncoder()
+  if (avFormatContext_ && avFormatContext_->pb) {
+    if (avFormatContext_->pb->error == 0) {
+      avio_flush(avFormatContext_->pb);
+    }
+    if (!avioContextHolder_) {
+      if (avFormatContext_->pb->error == 0) {
+        avio_close(avFormatContext_->pb);
+      }
+      avFormatContext_->pb = nullptr;
+    }
+  }
+}
+
+VideoEncoder::VideoEncoder(
+    const torch::Tensor& frames,
+    int frameRate,
+    std::string_view fileName,
+    const VideoStreamOptions& videoStreamOptions)
+    : frames_(validateFrames(frames)), inFrameRate_(frameRate) {
+  setFFmpegLogLevel();
+
+  // Allocate output format context
+  AVFormatContext* avFormatContext = nullptr;
+  int status = avformat_alloc_output_context2(
+      &avFormatContext, nullptr, nullptr, fileName.data());
+
+  TORCH_CHECK(
+      avFormatContext != nullptr,
+      "Couldn't allocate AVFormatContext. ",
+      "The destination file is ",
+      fileName,
+      ", check the desired extension? ",
+      getFFMPEGErrorStringFromErrorCode(status));
+  avFormatContext_.reset(avFormatContext);
+
+  status = avio_open(&avFormatContext_->pb, fileName.data(), AVIO_FLAG_WRITE);
+  TORCH_CHECK(
+      status >= 0,
+      "avio_open failed. The destination file is ",
+      fileName,
+      ", make sure it's a valid path? ",
+      getFFMPEGErrorStringFromErrorCode(status));
+  initializeEncoder(videoStreamOptions);
+}
+
+VideoEncoder::VideoEncoder(
+    const torch::Tensor& frames,
+    int frameRate,
+    std::string_view formatName,
+    std::unique_ptr<AVIOContextHolder> avioContextHolder,
+    const VideoStreamOptions& videoStreamOptions)
+    : frames_(validateFrames(frames)),
+      inFrameRate_(frameRate),
+      avioContextHolder_(std::move(avioContextHolder)) {
+  setFFmpegLogLevel();
+  // Map mkv -> matroska when used as format name
+  formatName = (formatName == "mkv") ? "matroska" : formatName;
+  AVFormatContext* avFormatContext = nullptr;
+  int status = avformat_alloc_output_context2(
+      &avFormatContext, nullptr, formatName.data(), nullptr);
+
+  TORCH_CHECK(
+      avFormatContext != nullptr,
+      "Couldn't allocate AVFormatContext. ",
+      "Check the desired format? Got format=",
+      formatName,
+      ". ",
+      getFFMPEGErrorStringFromErrorCode(status));
+  avFormatContext_.reset(avFormatContext);
+
+  avFormatContext_->pb = avioContextHolder_->getAVIOContext();
+
+  initializeEncoder(videoStreamOptions);
+}
+
+void VideoEncoder::initializeEncoder(
+    const VideoStreamOptions& videoStreamOptions) {
+  const AVCodec* avCodec =
+      avcodec_find_encoder(avFormatContext_->oformat->video_codec);
+  TORCH_CHECK(avCodec != nullptr, "Video codec not found");
+
+  AVCodecContext* avCodecContext = avcodec_alloc_context3(avCodec);
+  TORCH_CHECK(avCodecContext != nullptr, "Couldn't allocate codec context.");
+  avCodecContext_.reset(avCodecContext);
+
+  // Store dimension order and input pixel format
+  // TODO-VideoEncoder: Remove assumption that tensor in NCHW format
+  auto sizes = frames_.sizes();
+  inPixelFormat_ = AV_PIX_FMT_GBRP;
+  inHeight_ = static_cast<int>(sizes[2]);
+  inWidth_ = static_cast<int>(sizes[3]);
+
+  // Use specified dimensions or input dimensions
+  // TODO-VideoEncoder: Allow height and width to be set
+  outWidth_ = inWidth_;
+  outHeight_ = inHeight_;
+
+  // TODO-VideoEncoder: Enable other pixel formats
+  // Let FFmpeg choose best pixel format to minimize loss
+  outPixelFormat_ = avcodec_find_best_pix_fmt_of_list(
+      getSupportedPixelFormats(*avCodec), // List of supported formats
+      AV_PIX_FMT_GBRP, // We reorder input to GBRP currently
+      0, // No alpha channel
+      nullptr // Discard conversion loss information
+  );
+  TORCH_CHECK(outPixelFormat_ != -1, "Failed to find best pix fmt")
+
+  // Configure codec parameters
+  avCodecContext_->codec_id = avCodec->id;
+  avCodecContext_->width = outWidth_;
+  avCodecContext_->height = outHeight_;
+  avCodecContext_->pix_fmt = outPixelFormat_;
+  // TODO-VideoEncoder: Verify that frame_rate and time_base are correct
+  avCodecContext_->time_base = {1, inFrameRate_};
+  avCodecContext_->framerate = {inFrameRate_, 1};
+
+  // Set flag for containers that require extradata to be in the codec context
+  if (avFormatContext_->oformat->flags & AVFMT_GLOBALHEADER) {
+    avCodecContext_->flags |= AV_CODEC_FLAG_GLOBAL_HEADER;
+  }
+
+  // Apply videoStreamOptions
+  AVDictionary* options = nullptr;
+  if (videoStreamOptions.crf.has_value()) {
+    av_dict_set(
+        &options,
+        "crf",
+        std::to_string(videoStreamOptions.crf.value()).c_str(),
+        0);
+  }
+  int status = avcodec_open2(avCodecContext_.get(), avCodec, &options);
+  av_dict_free(&options);
+
+  TORCH_CHECK(
+      status == AVSUCCESS,
+      "avcodec_open2 failed: ",
+      getFFMPEGErrorStringFromErrorCode(status));
+
+  avStream_ = avformat_new_stream(avFormatContext_.get(), nullptr);
+  TORCH_CHECK(avStream_ != nullptr, "Couldn't create new stream.");
+
+  // Set the stream time base to encode correct frame timestamps
+  avStream_->time_base = avCodecContext_->time_base;
+  status = avcodec_parameters_from_context(
+      avStream_->codecpar, avCodecContext_.get());
+  TORCH_CHECK(
+      status == AVSUCCESS,
+      "avcodec_parameters_from_context failed: ",
+      getFFMPEGErrorStringFromErrorCode(status));
+}
+
+void VideoEncoder::encode() {
+  // To be on the safe side we enforce that encode() can only be called once
+  TORCH_CHECK(!encodeWasCalled_, "Cannot call encode() twice.");
+  encodeWasCalled_ = true;
+
+  int status = avformat_write_header(avFormatContext_.get(), nullptr);
+  TORCH_CHECK(
+      status == AVSUCCESS,
+      "Error in avformat_write_header: ",
+      getFFMPEGErrorStringFromErrorCode(status));
+
+  AutoAVPacket autoAVPacket;
+  int numFrames = static_cast<int>(frames_.sizes()[0]);
+  for (int i = 0; i < numFrames; ++i) {
+    torch::Tensor currFrame = frames_[i];
+    UniqueAVFrame avFrame = convertTensorToAVFrame(currFrame, i);
+    encodeFrame(autoAVPacket, avFrame);
+  }
+
+  flushBuffers();
+
+  status = av_write_trailer(avFormatContext_.get());
+  TORCH_CHECK(
+      status == AVSUCCESS,
+      "Error in av_write_trailer: ",
+      getFFMPEGErrorStringFromErrorCode(status));
+}
+
+UniqueAVFrame VideoEncoder::convertTensorToAVFrame(
+    const torch::Tensor& frame,
+    int frameIndex) {
+  // Initialize and cache scaling context if it does not exist
+  if (!swsContext_) {
+    swsContext_.reset(sws_getContext(
+        inWidth_,
+        inHeight_,
+        inPixelFormat_,
+        outWidth_,
+        outHeight_,
+        outPixelFormat_,
+        SWS_BICUBIC, // Used by FFmpeg CLI
+        nullptr,
+        nullptr,
+        nullptr));
+    TORCH_CHECK(swsContext_ != nullptr, "Failed to create scaling context");
+  }
+
+  UniqueAVFrame avFrame(av_frame_alloc());
+  TORCH_CHECK(avFrame != nullptr, "Failed to allocate AVFrame");
+
+  // Set output frame properties
+  avFrame->format = outPixelFormat_;
+  avFrame->width = outWidth_;
+  avFrame->height = outHeight_;
+  avFrame->pts = frameIndex;
+
+  int status = av_frame_get_buffer(avFrame.get(), 0);
+  TORCH_CHECK(status >= 0, "Failed to allocate frame buffer");
+
+  // Need to convert/scale the frame
+  // Create temporary frame with input format
+  UniqueAVFrame inputFrame(av_frame_alloc());
+  TORCH_CHECK(inputFrame != nullptr, "Failed to allocate input AVFrame");
+
+  inputFrame->format = inPixelFormat_;
+  inputFrame->width = inWidth_;
+  inputFrame->height = inHeight_;
+
+  uint8_t* tensorData = static_cast<uint8_t*>(frame.data_ptr());
+
+  // TODO-VideoEncoder: Reorder tensor if in NHWC format
+  int channelSize = inHeight_ * inWidth_;
+  // Reorder RGB -> GBR for AV_PIX_FMT_GBRP format
+  // TODO-VideoEncoder: Determine if FFmpeg supports planar RGB input format
+  inputFrame->data[0] = tensorData + channelSize;
+  inputFrame->data[1] = tensorData + (2 * channelSize);
+  inputFrame->data[2] = tensorData;
+
+  inputFrame->linesize[0] = inWidth_;
+  inputFrame->linesize[1] = inWidth_;
+  inputFrame->linesize[2] = inWidth_;
+
+  status = sws_scale(
+      swsContext_.get(),
+      inputFrame->data,
+      inputFrame->linesize,
+      0,
+      inputFrame->height,
+      avFrame->data,
+      avFrame->linesize);
+  TORCH_CHECK(status == outHeight_, "sws_scale failed");
+  return avFrame;
+}
+
+torch::Tensor VideoEncoder::encodeToTensor() {
+  TORCH_CHECK(
+      avioContextHolder_ != nullptr,
+      "Cannot encode to tensor, avio tensor context doesn't exist.");
+  encode();
+  auto avioToTensorContext =
+      dynamic_cast<AVIOToTensorContext*>(avioContextHolder_.get());
+  TORCH_CHECK(avioToTensorContext != nullptr, "Invalid AVIO context holder.");
+  return avioToTensorContext->getOutputTensor();
+}
+
+void VideoEncoder::encodeFrame(
+    AutoAVPacket& autoAVPacket,
+    const UniqueAVFrame& avFrame) {
+  auto status = avcodec_send_frame(avCodecContext_.get(), avFrame.get());
+  TORCH_CHECK(
+      status == AVSUCCESS,
+      "Error while sending frame: ",
+      getFFMPEGErrorStringFromErrorCode(status));
+
+  while (status >= 0) {
+    ReferenceAVPacket packet(autoAVPacket);
+    status = avcodec_receive_packet(avCodecContext_.get(), packet.get());
+    if (status == AVERROR(EAGAIN) || status == AVERROR_EOF) {
+      if (status == AVERROR_EOF) {
+        // Flush remaining buffered packets
+        status = av_interleaved_write_frame(avFormatContext_.get(), nullptr);
+        TORCH_CHECK(
+            status == AVSUCCESS,
+            "Failed to flush packet: ",
+            getFFMPEGErrorStringFromErrorCode(status));
+      }
+      return;
+    }
+    TORCH_CHECK(
+        status >= 0,
+        "Error receiving packet: ",
+        getFFMPEGErrorStringFromErrorCode(status));
+
+    // The code below is borrowed from torchaudio:
+    // https://github.com/pytorch/audio/blob/b6a3368a45aaafe05f1a6a9f10c68adc5e944d9e/src/libtorio/ffmpeg/stream_writer/encoder.cpp#L46
+    // Setting packet->duration to 1 allows the last frame to be properly
+    // encoded, and needs to be set before calling av_packet_rescale_ts.
+    if (packet->duration == 0) {
+      packet->duration = 1;
+    }
+    av_packet_rescale_ts(
+        packet.get(), avCodecContext_->time_base, avStream_->time_base);
+    packet->stream_index = avStream_->index;
+
+    status = av_interleaved_write_frame(avFormatContext_.get(), packet.get());
+    TORCH_CHECK(
+        status == AVSUCCESS,
+        "Error in av_interleaved_write_frame: ",
+        getFFMPEGErrorStringFromErrorCode(status));
+  }
+}
+
+void VideoEncoder::flushBuffers() {
+  AutoAVPacket autoAVPacket;
+  // Send null frame to signal end of input
+  encodeFrame(autoAVPacket, UniqueAVFrame(nullptr));
+}
+
 } // namespace facebook::torchcodec

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,65 @@ 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);
+
+  VideoEncoder(
+      const torch::Tensor& frames,
+      int frameRate,
+      std::string_view formatName,
+      std::unique_ptr<AVIOContextHolder> avioContextHolder,
+      const VideoStreamOptions& videoStreamOptions);
+
+  void encode();
+
+  torch::Tensor encodeToTensor();
+
+ 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_ = nullptr;
+  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;
+
+  std::unique_ptr<AVIOContextHolder> avioContextHolder_;
+
+  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)
@@ -73,7 +149,7 @@ int getNumChannels(const UniqueAVFrame& avFrame) {
 #endif
 }
 
-int getNumChannels(const UniqueAVCodecContext& avCodecContext) {
+int getNumChannels(const SharedAVCodecContext& avCodecContext) {
 #if LIBAVFILTER_VERSION_MAJOR > 8 || \
     (LIBAVFILTER_VERSION_MAJOR == 8 && LIBAVFILTER_VERSION_MINOR >= 44)
   return avCodecContext->ch_layout.nb_channels;
@@ -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