to_tensor, AVIOTensorContext fix

Author: Daniel Flores

src/torchcodec/_core/AVIOTensorContext.cpp

@@ -75,6 +75,11 @@ int write(void* opaque, const uint8_t* buf, int buf_size) {
   uint8_t* outputTensorData = tensorContext->data.data_ptr<uint8_t>();
   std::memcpy(outputTensorData + tensorContext->current, buf, bufSize);
   tensorContext->current += bufSize;
+  // Track the maximum position written so getOutputTensor's narrow() does not
+  // truncate the file if final seek was backwards
+  if (tensorContext->current > tensorContext->max) {
+    tensorContext->max = tensorContext->current;
+  }
   return buf_size;
 }
 
@@ -101,7 +106,7 @@ int64_t seek(void* opaque, int64_t offset, int whence) {
 } // namespace
 
 AVIOFromTensorContext::AVIOFromTensorContext(torch::Tensor data)
-    : tensorContext_{data, 0} {
+    : tensorContext_{data, 0, 0} {
   TORCH_CHECK(data.numel() > 0, "data must not be empty");
   TORCH_CHECK(data.is_contiguous(), "data must be contiguous");
   TORCH_CHECK(data.scalar_type() == torch::kUInt8, "data must be kUInt8");
@@ -110,14 +115,17 @@ AVIOFromTensorContext::AVIOFromTensorContext(torch::Tensor data)
 }
 
 AVIOToTensorContext::AVIOToTensorContext()
-    : tensorContext_{torch::empty({INITIAL_TENSOR_SIZE}, {torch::kUInt8}), 0} {
+    : tensorContext_{
+          torch::empty({INITIAL_TENSOR_SIZE}, {torch::kUInt8}),
+          0,
+          0} {
   createAVIOContext(
       nullptr, &write, &seek, &tensorContext_, /*isForWriting=*/true);
 }
 
 torch::Tensor AVIOToTensorContext::getOutputTensor() {
   return tensorContext_.data.narrow(
-      /*dim=*/0, /*start=*/0, /*length=*/tensorContext_.current);
+      /*dim=*/0, /*start=*/0, /*length=*/tensorContext_.max);
 }
 
 } // namespace facebook::torchcodec

src/torchcodec/_core/AVIOTensorContext.h

@@ -16,6 +16,7 @@ namespace detail {
 struct TensorContext {
   torch::Tensor data;
   int64_t current;
+  int64_t max;
 };
 
 } // namespace detail

src/torchcodec/_core/Encoder.cpp

@@ -4,10 +4,6 @@
 #include "src/torchcodec/_core/Encoder.h"
 #include "torch/types.h"
 
-extern "C" {
-#include <libavutil/pixdesc.h>
-}
-
 namespace facebook::torchcodec {
 
 namespace {
@@ -543,9 +539,15 @@ torch::Tensor validateFrames(const torch::Tensor& frames) {
 
 VideoEncoder::~VideoEncoder() {
   if (avFormatContext_ && avFormatContext_->pb) {
-    avio_flush(avFormatContext_->pb);
-    avio_close(avFormatContext_->pb);
-    avFormatContext_->pb = nullptr;
+    if (avFormatContext_->pb->error == 0) {
+      avio_flush(avFormatContext_->pb);
+    }
+    if (!avioContextHolder_) {
+      if (avFormatContext_->pb->error == 0) {
+        avio_close(avFormatContext_->pb);
+      }
+      avFormatContext_->pb = nullptr;
+    }
   }
 }
 
@@ -581,6 +583,34 @@ VideoEncoder::VideoEncoder(
   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();
+  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 =
@@ -751,6 +781,17 @@ UniqueAVFrame VideoEncoder::convertTensorToAVFrame(
   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) {

src/torchcodec/_core/Encoder.h

@@ -141,8 +141,17 @@ class VideoEncoder {
       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(
@@ -167,6 +176,8 @@ class VideoEncoder {
   int outHeight_ = -1;
   AVPixelFormat outPixelFormat_ = AV_PIX_FMT_NONE;
 
+  std::unique_ptr<AVIOContextHolder> avioContextHolder_;
+
   bool encodeWasCalled_ = false;
 };
 

src/torchcodec/_core/__init__.py

@@ -26,6 +26,7 @@
     encode_audio_to_file_like,
     encode_audio_to_tensor,
     encode_video_to_file,
+    encode_video_to_tensor,
     get_ffmpeg_library_versions,
     get_frame_at_index,
     get_frame_at_pts,

src/torchcodec/_core/custom_ops.cpp

@@ -32,12 +32,14 @@ TORCH_LIBRARY(torchcodec_ns, m) {
   m.def("create_from_file(str filename, str? seek_mode=None) -> Tensor");
   m.def(
       "encode_audio_to_file(Tensor samples, int sample_rate, str filename, 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, int? crf=None) -> ()");
   m.def(
       "encode_audio_to_tensor(Tensor samples, int sample_rate, str format, int? bit_rate=None, int? num_channels=None, int? desired_sample_rate=None) -> Tensor");
   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, int? crf=None) -> ()");
+  m.def(
+      "encode_video_to_tensor(Tensor frames, int frame_rate, str format, int? crf=None) -> Tensor");
   m.def(
       "create_from_tensor(Tensor video_tensor, str? seek_mode=None) -> Tensor");
   m.def(
@@ -498,21 +500,6 @@ OpsAudioFramesOutput get_frames_by_pts_in_range_audio(
   return makeOpsAudioFramesOutput(result);
 }
 
-void encode_video_to_file(
-    const at::Tensor& frames,
-    int64_t frame_rate,
-    std::string_view file_name,
-    std::optional<int64_t> crf = std::nullopt) {
-  VideoStreamOptions videoStreamOptions;
-  videoStreamOptions.crf = crf;
-  VideoEncoder(
-      frames,
-      validateInt64ToInt(frame_rate, "frame_rate"),
-      file_name,
-      videoStreamOptions)
-      .encode();
-}
-
 void encode_audio_to_file(
     const at::Tensor& samples,
     int64_t sample_rate,
@@ -587,6 +574,38 @@ void _encode_audio_to_file_like(
   encoder.encode();
 }
 
+void encode_video_to_file(
+    const at::Tensor& frames,
+    int64_t frame_rate,
+    std::string_view file_name,
+    std::optional<int64_t> crf = std::nullopt) {
+  VideoStreamOptions videoStreamOptions;
+  videoStreamOptions.crf = crf;
+  VideoEncoder(
+      frames,
+      validateInt64ToInt(frame_rate, "frame_rate"),
+      file_name,
+      videoStreamOptions)
+      .encode();
+}
+
+at::Tensor encode_video_to_tensor(
+    const at::Tensor& frames,
+    int64_t frame_rate,
+    std::string_view format,
+    std::optional<int64_t> crf = std::nullopt) {
+  auto avioContextHolder = std::make_unique<AVIOToTensorContext>();
+  VideoStreamOptions videoStreamOptions;
+  videoStreamOptions.crf = crf;
+  return VideoEncoder(
+             frames,
+             validateInt64ToInt(frame_rate, "frame_rate"),
+             format,
+             std::move(avioContextHolder),
+             videoStreamOptions)
+      .encodeToTensor();
+}
+
 // For testing only. We need to implement this operation as a core library
 // function because what we're testing is round-tripping pts values as
 // double-precision floating point numbers from C++ to Python and back to C++.
@@ -847,9 +866,10 @@ TORCH_LIBRARY_IMPL(torchcodec_ns, BackendSelect, m) {
 
 TORCH_LIBRARY_IMPL(torchcodec_ns, CPU, m) {
   m.impl("encode_audio_to_file", &encode_audio_to_file);
-  m.impl("encode_video_to_file", &encode_video_to_file);
   m.impl("encode_audio_to_tensor", &encode_audio_to_tensor);
   m.impl("_encode_audio_to_file_like", &_encode_audio_to_file_like);
+  m.impl("encode_video_to_file", &encode_video_to_file);
+  m.impl("encode_video_to_tensor", &encode_video_to_tensor);
   m.impl("seek_to_pts", &seek_to_pts);
   m.impl("add_video_stream", &add_video_stream);
   m.impl("_add_video_stream", &_add_video_stream);

src/torchcodec/_core/ops.py

@@ -92,15 +92,18 @@ def load_torchcodec_shared_libraries():
 encode_audio_to_file = torch._dynamo.disallow_in_graph(
     torch.ops.torchcodec_ns.encode_audio_to_file.default
 )
-encode_video_to_file = torch._dynamo.disallow_in_graph(
-    torch.ops.torchcodec_ns.encode_video_to_file.default
-)
 encode_audio_to_tensor = torch._dynamo.disallow_in_graph(
     torch.ops.torchcodec_ns.encode_audio_to_tensor.default
 )
 _encode_audio_to_file_like = torch._dynamo.disallow_in_graph(
     torch.ops.torchcodec_ns._encode_audio_to_file_like.default
 )
+encode_video_to_file = torch._dynamo.disallow_in_graph(
+    torch.ops.torchcodec_ns.encode_video_to_file.default
+)
+encode_video_to_tensor = torch._dynamo.disallow_in_graph(
+    torch.ops.torchcodec_ns.encode_video_to_tensor.default
+)
 create_from_tensor = torch._dynamo.disallow_in_graph(
     torch.ops.torchcodec_ns.create_from_tensor.default
 )
@@ -254,16 +257,6 @@ def encode_audio_to_file_abstract(
     return
 
 
-@register_fake("torchcodec_ns::encode_video_to_file")
-def encode_video_to_file_abstract(
-    frames: torch.Tensor,
-    frame_rate: int,
-    filename: str,
-    crf: Optional[int] = None,
-) -> None:
-    return
-
-
 @register_fake("torchcodec_ns::encode_audio_to_tensor")
 def encode_audio_to_tensor_abstract(
     samples: torch.Tensor,
@@ -289,6 +282,26 @@ def _encode_audio_to_file_like_abstract(
     return
 
 
+@register_fake("torchcodec_ns::encode_video_to_file")
+def encode_video_to_file_abstract(
+    frames: torch.Tensor,
+    frame_rate: int,
+    filename: str,
+    crf: Optional[int],
+) -> None:
+    return
+
+
+@register_fake("torchcodec_ns::encode_video_to_tensor")
+def encode_video_to_tensor_abstract(
+    frames: torch.Tensor,
+    frame_rate: int,
+    format: str,
+    crf: Optional[int],
+) -> torch.Tensor:
+    return torch.empty([], dtype=torch.long)
+
+
 @register_fake("torchcodec_ns::create_from_tensor")
 def create_from_tensor_abstract(
     video_tensor: torch.Tensor, seek_mode: Optional[str]

test/test_ops.py

@@ -29,6 +29,7 @@
     create_from_tensor,
     encode_audio_to_file,
     encode_video_to_file,
+    encode_video_to_tensor,
     get_ffmpeg_library_versions,
     get_frame_at_index,
     get_frame_at_pts,
@@ -1378,14 +1379,18 @@ def test_bad_input(self, tmp_path):
                 filename="./bad/path.mp3",
             )
 
-    def decode(self, file_path) -> torch.Tensor:
-        decoder = create_from_file(str(file_path), seek_mode="approximate")
+    def decode(self, file_path=None, tensor=None) -> torch.Tensor:
+        if file_path is not None:
+            decoder = create_from_file(str(file_path), seek_mode="approximate")
+        elif tensor is not None:
+            decoder = create_from_tensor(tensor, seek_mode="approximate")
         add_video_stream(decoder)
         frames, *_ = get_frames_in_range(decoder, start=0, stop=60)
         return frames
 
     @pytest.mark.parametrize("format", ("mov", "mp4", "mkv", "webm"))
-    def test_video_encoder_round_trip(self, tmp_path, format):
+    @pytest.mark.parametrize("output_method", ("to_file", "to_tensor"))
+    def test_video_encoder_round_trip(self, tmp_path, format, output_method):
         # Test that decode(encode(decode(asset))) == decode(asset)
         ffmpeg_version = get_ffmpeg_major_version()
         # In FFmpeg6, the default codec's best pixel format is lossy for all container formats but webm.
@@ -1399,14 +1404,25 @@ def test_video_encoder_round_trip(self, tmp_path, format):
         ):
             pytest.skip("Codec for webm is not available in this FFmpeg installation.")
         asset = TEST_SRC_2_720P
-        source_frames = self.decode(str(asset.path)).data
+        source_frames = self.decode(file_path=str(asset.path)).data
 
-        encoded_path = str(tmp_path / f"encoder_output.{format}")
         frame_rate = 30  # Frame rate is fixed with num frames decoded
-        encode_video_to_file(
-            frames=source_frames, frame_rate=frame_rate, filename=encoded_path, crf=0
-        )
-        round_trip_frames = self.decode(encoded_path).data
+        if output_method == "to_file":
+            encoded_path = str(tmp_path / f"encoder_output.{format}")
+            encode_video_to_file(
+                frames=source_frames,
+                frame_rate=frame_rate,
+                filename=encoded_path,
+                crf=0,
+            )
+            round_trip_frames = self.decode(file_path=encoded_path).data
+        else:  # to_tensor
+            format = "matroska" if format == "mkv" else format
+            encoded_tensor = encode_video_to_tensor(
+                source_frames, frame_rate, format, crf=0
+            )
+            round_trip_frames = self.decode(tensor=encoded_tensor).data
+
         assert source_frames.shape == round_trip_frames.shape
         assert source_frames.dtype == round_trip_frames.dtype