Merge branch 'main' of https://github.com/meta-pytorch/torchcodec into encode_gpu

Author: Dan-Flores

src/torchcodec/_core/Encoder.cpp

@@ -674,7 +674,7 @@ VideoEncoder::~VideoEncoder() {
 
 VideoEncoder::VideoEncoder(
     const torch::Tensor& frames,
-    int frameRate,
+    double frameRate,
     std::string_view fileName,
     const VideoStreamOptions& videoStreamOptions)
     : frames_(validateFrames(frames, videoStreamOptions.device)),
@@ -707,7 +707,7 @@ VideoEncoder::VideoEncoder(
 
 VideoEncoder::VideoEncoder(
     const torch::Tensor& frames,
-    int frameRate,
+    double frameRate,
     std::string_view formatName,
     std::unique_ptr<AVIOContextHolder> avioContextHolder,
     const VideoStreamOptions& videoStreamOptions)
@@ -812,9 +812,9 @@ void VideoEncoder::initializeEncoder(
   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};
+  // TODO-VideoEncoder: Add and utilize output frame_rate option
+  avCodecContext_->framerate = av_d2q(inFrameRate_, INT_MAX);
+  avCodecContext_->time_base = av_inv_q(avCodecContext_->framerate);
 
   // Set flag for containers that require extradata to be in the codec context
   if (avFormatContext_->oformat->flags & AVFMT_GLOBALHEADER) {
@@ -866,6 +866,10 @@ void VideoEncoder::initializeEncoder(
 
   // Set the stream time base to encode correct frame timestamps
   avStream_->time_base = avCodecContext_->time_base;
+  // Set the stream frame rate to store correct frame durations for some
+  // containers (webm, mkv)
+  avStream_->r_frame_rate = avCodecContext_->framerate;
+
   status = avcodec_parameters_from_context(
       avStream_->codecpar, avCodecContext_.get());
   TORCH_CHECK(

src/torchcodec/_core/Encoder.h

@@ -144,13 +144,13 @@ class VideoEncoder {
 
   VideoEncoder(
       const torch::Tensor& frames,
-      int frameRate,
+      double frameRate,
       std::string_view fileName,
       const VideoStreamOptions& videoStreamOptions);
 
   VideoEncoder(
       const torch::Tensor& frames,
-      int frameRate,
+      double frameRate,
       std::string_view formatName,
       std::unique_ptr<AVIOContextHolder> avioContextHolder,
       const VideoStreamOptions& videoStreamOptions);
@@ -170,7 +170,7 @@ class VideoEncoder {
   UniqueSwsContext swsContext_;
 
   const torch::Tensor frames_;
-  int inFrameRate_;
+  double inFrameRate_;
 
   int inWidth_ = -1;
   int inHeight_ = -1;

src/torchcodec/_core/custom_ops.cpp

@@ -37,11 +37,11 @@ TORCH_LIBRARY(torchcodec_ns, m) {
   m.def(
       "_encode_audio_to_file_like(Tensor samples, int sample_rate, str format, int file_like_context, int? bit_rate=None, int? num_channels=None, int? desired_sample_rate=None) -> ()");
   m.def(
-      "encode_video_to_file(Tensor frames, int frame_rate, str filename, str device=\"cpu\", str? codec=None, str? pixel_format=None, float? crf=None, str? preset=None, str[]? extra_options=None) -> ()");
+      "encode_video_to_file(Tensor frames, float frame_rate, str filename, str device=\"cpu\", str? codec=None, str? pixel_format=None, float? crf=None, str? preset=None, str[]? extra_options=None) -> ()");
   m.def(
-      "encode_video_to_tensor(Tensor frames, int frame_rate, str format, str device=\"cpu\", str? codec=None, str? pixel_format=None, float? crf=None, str? preset=None, str[]? extra_options=None) -> Tensor");
+      "encode_video_to_tensor(Tensor frames, float frame_rate, str format, str device=\"cpu\", str? codec=None, str? pixel_format=None, float? crf=None, str? preset=None, str[]? extra_options=None) -> Tensor");
   m.def(
-      "_encode_video_to_file_like(Tensor frames, int frame_rate, str format, int file_like_context, str device=\"cpu\",str? codec=None, str? pixel_format=None, float? crf=None, str? preset=None, str[]? extra_options=None) -> ()");
+      "_encode_video_to_file_like(Tensor frames, float frame_rate, str format, int file_like_context, str device=\"cpu\",str? codec=None, str? pixel_format=None, float? crf=None, str? preset=None, str[]? extra_options=None) -> ()");
   m.def(
       "create_from_tensor(Tensor video_tensor, str? seek_mode=None) -> Tensor");
   m.def(
@@ -638,7 +638,7 @@ void _encode_audio_to_file_like(
 
 void encode_video_to_file(
     const at::Tensor& frames,
-    int64_t frame_rate,
+    double frame_rate,
     std::string_view file_name,
     std::string_view device = "cpu",
     std::optional<std::string_view> codec = std::nullopt,
@@ -658,17 +658,12 @@ void encode_video_to_file(
         unflattenExtraOptions(extra_options.value());
   }
 
-  VideoEncoder(
-      frames,
-      validateInt64ToInt(frame_rate, "frame_rate"),
-      file_name,
-      videoStreamOptions)
-      .encode();
+  VideoEncoder(frames, frame_rate, file_name, videoStreamOptions).encode();
 }
 
 at::Tensor encode_video_to_tensor(
     const at::Tensor& frames,
-    int64_t frame_rate,
+    double frame_rate,
     std::string_view format,
     std::string_view device = "cpu",
     std::optional<std::string_view> codec = std::nullopt,
@@ -691,7 +686,7 @@ at::Tensor encode_video_to_tensor(
 
   return VideoEncoder(
              frames,
-             validateInt64ToInt(frame_rate, "frame_rate"),
+             frame_rate,
              format,
              std::move(avioContextHolder),
              videoStreamOptions)
@@ -700,7 +695,7 @@ at::Tensor encode_video_to_tensor(
 
 void _encode_video_to_file_like(
     const at::Tensor& frames,
-    int64_t frame_rate,
+    double frame_rate,
     std::string_view format,
     int64_t file_like_context,
     std::string_view device = "cpu",
@@ -729,7 +724,7 @@ void _encode_video_to_file_like(
 
   VideoEncoder encoder(
       frames,
-      validateInt64ToInt(frame_rate, "frame_rate"),
+      frame_rate,
       format,
       std::move(avioContextHolder),
       videoStreamOptions);

src/torchcodec/_core/ops.py

@@ -210,7 +210,7 @@ def encode_audio_to_file_like(
 
 def encode_video_to_file_like(
     frames: torch.Tensor,
-    frame_rate: int,
+    frame_rate: float,
     format: str,
     file_like: Union[io.RawIOBase, io.BufferedIOBase],
     device: Optional[str] = "cpu",
@@ -332,7 +332,7 @@ def _encode_audio_to_file_like_abstract(
 @register_fake("torchcodec_ns::encode_video_to_file")
 def encode_video_to_file_abstract(
     frames: torch.Tensor,
-    frame_rate: int,
+    frame_rate: float,
     filename: str,
     device: str = "cpu",
     codec: Optional[str] = None,
@@ -347,7 +347,7 @@ def encode_video_to_file_abstract(
 @register_fake("torchcodec_ns::encode_video_to_tensor")
 def encode_video_to_tensor_abstract(
     frames: torch.Tensor,
-    frame_rate: int,
+    frame_rate: float,
     format: str,
     device: str = "cpu",
     codec: Optional[str] = None,
@@ -362,7 +362,7 @@ def encode_video_to_tensor_abstract(
 @register_fake("torchcodec_ns::_encode_video_to_file_like")
 def _encode_video_to_file_like_abstract(
     frames: torch.Tensor,
-    frame_rate: int,
+    frame_rate: float,
     format: str,
     file_like_context: int,
     device: str = "cpu",

src/torchcodec/encoders/_video_encoder.py

@@ -15,7 +15,7 @@ class VideoEncoder:
             tensor of shape ``(N, C, H, W)`` where N is the number of frames,
             C is 3 channels (RGB), H is height, and W is width.
             Values must be uint8 in the range ``[0, 255]``.
-        frame_rate (int): The frame rate of the **input** ``frames``. Also defines the encoded **output** frame rate.
+        frame_rate (float): The frame rate of the **input** ``frames``. Also defines the encoded **output** frame rate.
         device (str or torch.device, optional): The device to use for encoding. Default: "cpu".
             If you pass a CUDA device, frames will be encoded on GPU.
             Note: The "beta" CUDA backend is not supported for encoding.
@@ -25,7 +25,7 @@ def __init__(
         self,
         frames: Tensor,
         *,
-        frame_rate: int,
+        frame_rate: float,
         device: Optional[Union[str, torch_device]] = "cpu",
     ):
         torch._C._log_api_usage_once("torchcodec.encoders.VideoEncoder")