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

Author: NicolasHug

pyproject.toml

@@ -54,4 +54,8 @@ norecursedirs = ["third-party-interface"]
 # the CI, where we definitely want the 'slow' tests to run.
 addopts = "-v -m 'not slow'"
 
+# Tells pytest to keep around temporary paths (via tmp_path fixture) for
+# *failed tests* only. Otherwise, the /tmp folder can grow too large.
+tmp_path_retention_policy = "failed"
+
 testpaths = ["test"]

src/torchcodec/_core/Encoder.cpp

@@ -1057,6 +1057,7 @@ MultiStreamEncoder::MultiStreamEncoder() {
 }
 
 void MultiStreamEncoder::open(std::string_view fileName) {
+  STD_TORCH_CHECK(!closed_, "Cannot open after close() was called.");
   STD_TORCH_CHECK(!headerWritten_, "open() was already called.");
 
   AVFormatContext* avFormatContext = nullptr;
@@ -1086,6 +1087,7 @@ void MultiStreamEncoder::open(std::string_view fileName) {
 void MultiStreamEncoder::open(
     std::string_view formatName,
     std::unique_ptr<AVIOContextHolder> avioContextHolder) {
+  STD_TORCH_CHECK(!closed_, "Cannot open after close() was called.");
   STD_TORCH_CHECK(!headerWritten_, "open() was already called.");
 
   avioContextHolder_ = std::move(avioContextHolder);
@@ -1148,6 +1150,13 @@ int MultiStreamEncoder::addAudioStream(
   STD_TORCH_CHECK(sampleRate > 0, "sample_rate must be > 0, got ", sampleRate);
   STD_TORCH_CHECK(
       numChannels > 0, "num_channels must be > 0, got ", numChannels);
+  STD_TORCH_CHECK(
+      numChannels <= AV_NUM_DATA_POINTERS,
+      "Trying to encode ",
+      numChannels,
+      " channels, but FFmpeg only supports ",
+      AV_NUM_DATA_POINTERS,
+      " channels per frame.");
 
   AudioStream audioStream;
   audioStream.inSampleRate = sampleRate;
@@ -1412,6 +1421,7 @@ void MultiStreamEncoder::openStreamsAndWriteHeader() {
 void MultiStreamEncoder::addFrames(
     const torch::stable::Tensor& frames,
     int streamIndex) {
+  STD_TORCH_CHECK(!closed_, "Cannot add frames after close() was called.");
   STD_TORCH_CHECK(headerWritten_, "Call open() before addFrames().");
   STD_TORCH_CHECK(
       streamIndex >= 0 && streamIndex < static_cast<int>(videoStreams_.size()),
@@ -1501,6 +1511,7 @@ void MultiStreamEncoder::encodeVideoFrame(
 void MultiStreamEncoder::addSamples(
     const torch::stable::Tensor& samples,
     int streamIndex) {
+  STD_TORCH_CHECK(!closed_, "Cannot add samples after close() was called.");
   STD_TORCH_CHECK(headerWritten_, "Call open() before addSamples().");
   STD_TORCH_CHECK(
       streamIndex >= 0 && streamIndex < static_cast<int>(audioStreams_.size()),

src/torchcodec/encoders/_multi_stream_encoder.py

@@ -89,11 +89,18 @@ def add_audio(
 
     # TODO MultiStreamEncoder: Maybe there should 2 separate methods, one for
     # file, one for file-like.
-    def open(self, dest, *, format: str | None = None) -> None:
+    def open(self, dest, *, format: str | None = None) -> "StreamingEncoder":
         if format is not None:
             _core.streaming_encoder_open_file_like(self._encoder_tensor, format, dest)
         else:
             _core.streaming_encoder_open_file(self._encoder_tensor, str(dest))
+        return self
 
     def close(self) -> None:
         _core.streaming_encoder_close(self._encoder_tensor)
+
+    def __enter__(self) -> "StreamingEncoder":
+        return self
+
+    def __exit__(self, exc_type, exc_val, exc_tb) -> None:
+        self.close()