Merge branch 'main' of github.com:meta-pytorch/torchcodec into worktree-stream-sync-events

Author: NicolasHug

examples/decoding/blocks.py

@@ -22,7 +22,7 @@
 .. code-block::
 
    Demuxer  ->  PacketDecoder  ->  ColorConverter
-    Packet       DecodedFrame        RGB Frame
+    Packet         RawFrame          RGB Frame
 
 The blocks are passive: they never create threads, and they release the GIL.
 You decide how they are composed, on which threads, and where to stop. Below
@@ -65,7 +65,8 @@
 #
 # ``PacketDecoder`` and ``ColorConverter`` both accept ``device="cuda"``:
 # decoding then runs on NVDEC and the color conversion on the GPU, and the
-# frames never leave the device. Demuxing always happens on the CPU.
+# frames never leave the device. Demuxing always happens on the CPU. Left
+# unspecified, ``device`` is the current default device.
 from torchcodec.decoders._blocks import ColorConverter, Demuxer, PacketDecoder
 
 demuxer = Demuxer(video_path)
@@ -74,10 +75,10 @@
 
 frames = []
 for packet in demuxer:
-    for decoded_frame in packet_decoder.decode(packet):
-        frames.append(color_converter.convert(decoded_frame))
-for decoded_frame in packet_decoder.drain():
-    frames.append(color_converter.convert(decoded_frame))
+    for raw_frame in packet_decoder.decode(packet):
+        frames.append(color_converter.convert(raw_frame))
+for raw_frame in packet_decoder.drain():
+    frames.append(color_converter.convert(raw_frame))
 
 print(f"{len(frames)} frames, {frames[0].data.shape = }, "
       f"{frames[0].pts_seconds = }, {frames[0].data.device = }")
@@ -104,9 +105,9 @@ def decode(packet_decoder, packets):
     yield from packet_decoder.drain()
 
 
-def color_convert(color_converter, decoded_frames):
-    for decoded_frame in decoded_frames:
-        yield color_converter.convert(decoded_frame)
+def color_convert(color_converter, raw_frames):
+    for raw_frame in raw_frames:
+        yield color_converter.convert(raw_frame)
 
 
 def prefetch(upstream, buffer_size=8):
@@ -143,8 +144,8 @@ def convert_on_own_thread():
     demuxer = Demuxer(video_path)
     packet_decoder = PacketDecoder(demuxer, device=device)
     color_converter = ColorConverter(device=device)
-    decoded_frames = prefetch(decode(packet_decoder, demux(demuxer)))
-    return color_convert(color_converter, decoded_frames)
+    raw_frames = prefetch(decode(packet_decoder, demux(demuxer)))
+    return color_convert(color_converter, raw_frames)
 
 
 def demux_on_own_thread():
@@ -197,13 +198,12 @@ def demux_on_own_thread():
 # Raw frames
 # ----------
 #
-# Color conversion is optional. A ``DecodedFrame`` can hand out the decoder's
-# own planes as tensor views, with no copy and no conversion.
+# Color conversion is optional. A ``RawFrame`` can hand out the decoder's own
+# planes as tensor views, with no copy and no conversion.
 demuxer = Demuxer(video_path)
 packet_decoder = PacketDecoder(demuxer, device=device)
-decoded_frame = next(decode(packet_decoder, demux(demuxer)))
+raw_frame = next(decode(packet_decoder, demux(demuxer)))
 
-raw_frame = decoded_frame.materialize()
 Y, U, V = raw_frame.planes
 print(f"{raw_frame.pix_fmt = }, {raw_frame.bit_depth = }, "
       f"{raw_frame.colorspace = }, {raw_frame.color_range = }")
@@ -215,8 +215,9 @@ def demux_on_own_thread():
 # views over a single plane. Writing through them is visible downstream.
 #
 # Being the decoder's own planes, they are also never rotated - a video whose
-# container asks for a rotation gives you the samples as they were encoded.
-# ``ColorConverter`` applies the rotation for you.
+# container asks for a rotation gives you the samples as they were encoded, and
+# ``raw_frame.rotation_degrees`` tells you what to apply. ``ColorConverter``
+# applies it for you.
 #
 # So we can do the color conversion ourselves. Here it's plain PyTorch ops -
 # it could just as well be a Triton or CUDA kernel, fused with whatever your
@@ -246,7 +247,7 @@ def upsample(plane):
 
 
 ours = yuv420_to_rgb(Y, U, V)
-reference = ColorConverter(device=device).convert(decoded_frame).data
+reference = ColorConverter(device=device).convert(raw_frame).data
 print(f"{ours.shape = }, mean abs diff vs ColorConverter: "
       f"{(ours.float() - reference.float()).abs().mean():.2f}")
 
@@ -273,9 +274,8 @@ def upsample(plane):
 
 hdr_demuxer = Demuxer(hdr_video_path)
 hdr_packet_decoder = PacketDecoder(hdr_demuxer, device=device)
-hdr_frame = next(decode(hdr_packet_decoder, demux(hdr_demuxer)))
+hdr_raw = next(decode(hdr_packet_decoder, demux(hdr_demuxer)))
 
-hdr_raw = hdr_frame.materialize()
 hdr_Y = hdr_raw.planes[0]
 print(f"{hdr_raw.pix_fmt = }, {hdr_raw.bit_depth = }, "
       f"{hdr_raw.colorspace = }, {hdr_Y.dtype = }")

src/torchcodec/_core/BetaCudaDeviceInterface.cpp

@@ -384,6 +384,8 @@ BetaCudaDeviceInterface::BetaCudaDeviceInterface(const StableDevice& device)
   STD_TORCH_CHECK(
       device_.type() == kStableCUDA, "Unsupported device: must be CUDA");
 
+  device_ = StableDevice(kStableCUDA, get_device_index(device_));
+
   // Note: now that we have the CudaContextGuard, we might not need to do that
   // anymore. The comment says we need pytorch to create the context - maybe
   // that's true, but that's a very old comment now.

src/torchcodec/_core/ColorConverter.cpp

@@ -15,6 +15,17 @@
 
 namespace facebook::torchcodec {
 
+namespace {
+// Only ever used to name a device in the error message below.
+std::string printable(const StableDevice& device) {
+  std::string name = device_type_name(device.type());
+  if (device.type() != kStableCPU && device.index() >= 0) {
+    name += ":" + std::to_string(device.index());
+  }
+  return name;
+}
+} // namespace
+
 ColorConverter::ColorConverter(
     const StableDevice& device,
     OutputDtypeConfig output_dtype_config)
@@ -23,6 +34,7 @@ ColorConverter::ColorConverter(
   STD_TORCH_CHECK(
       device_interface_ != nullptr,
       "Failed to create device interface. This should never happen, please report.");
+  device_ = device_interface_->device(); // resolved, so we don't have to
 }
 
 void ColorConverter::maybe_initialize_interface(OutputDtype output_dtype) {
@@ -44,7 +56,24 @@ void ColorConverter::maybe_initialize_interface(OutputDtype output_dtype) {
   initialized_output_dtype_ = output_dtype;
 }
 
-torch::stable::Tensor ColorConverter::convert(const AVFrame& av_frame) {
+torch::stable::Tensor ColorConverter::convert(
+    const AVFrame& av_frame,
+    const StableDevice& frame_device) {
+  // TODO_API_BREAKDOWN CC P2: OK, it's not fantastic that we have to pass the
+  // frame's device. Especially given the related design TODO about whether the
+  // RawFrame should carry that device field at all. Maybe it should, maybe it's
+  // overkill. I think the main alternative is to retrieve the device from the
+  // AVFrame, it's possible, but likely requires moving the
+  // StandaloneFrameAttachedData to the public header.
+  STD_TORCH_CHECK(
+      frame_device == device_,
+      "This ColorConverter is on ",
+      printable(device_),
+      " but the frame's samples are on ",
+      printable(frame_device),
+      ". A ColorConverter only converts frames that are already on its own "
+      "device: create one per device, or move the RGB output afterwards.");
+
   OutputDtype output_dtype = resolve_output_dtype(
       output_dtype_config_, static_cast<AVPixelFormat>(av_frame.format));
   maybe_initialize_interface(output_dtype);

src/torchcodec/_core/ColorConverter.h

@@ -8,6 +8,7 @@
 
 #include <memory>
 #include <optional>
+#include <string>
 #include <string_view>
 
 #include "DeviceInterface.h"
@@ -23,7 +24,9 @@ class FORCE_PUBLIC_VISIBILITY ColorConverter {
       const StableDevice& device = StableDevice(kStableCPU),
       OutputDtypeConfig output_dtype_config = OutputDtypeConfig::UINT8);
 
-  torch::stable::Tensor convert(const AVFrame& av_frame);
+  torch::stable::Tensor convert(
+      const AVFrame& av_frame,
+      const StableDevice& frame_device);
 
  private:
   void maybe_initialize_interface(OutputDtype output_dtype);

src/torchcodec/_core/Demuxer.cpp

@@ -60,12 +60,72 @@ Demuxer::Demuxer(
   STD_TORCH_CHECK(raw_context != nullptr, "Failed to allocate AVFormatContext");
   format_context_.reset(raw_context);
 
-  status = avformat_find_stream_info(format_context_.get(), nullptr);
+  select_stream(stream_index);
+}
+
+Demuxer::Demuxer(
+    std::unique_ptr<AVIOContextHolder> avio_context_holder,
+    std::optional<int> stream_index)
+    : avio_context_holder_(std::move(avio_context_holder)) {
+  set_ffmpeg_log_level();
+
+  STD_TORCH_CHECK(avio_context_holder_ != nullptr, "Context holder is null");
+
+  // FFmpeg takes a reference to the pointer in the call to open, so we can't
+  // hand it a unique_ptr. That means we must free the context ourselves if the
+  // open fails.
+  AVFormatContext* raw_context = avformat_alloc_context();
+  STD_TORCH_CHECK(raw_context != nullptr, "Failed to allocate AVFormatContext");
+  raw_context->pb = avio_context_holder_->get_avio_context();
+
+  int status = avformat_open_input(&raw_context, nullptr, nullptr, nullptr);
+  if (status != 0) {
+    avformat_free_context(raw_context);
+    STD_TORCH_CHECK(
+        false,
+        "Could not open input buffer: " +
+            get_ffmpeg_error_string_from_error_code(status));
+  }
+  format_context_.reset(raw_context);
+
+  select_stream(stream_index);
+}
+
+void Demuxer::validate_requested_stream(int stream_index) {
+  int num_streams = static_cast<int>(format_context_->nb_streams);
+  STD_TORCH_CHECK(
+      stream_index >= 0 && stream_index < num_streams,
+      "The stream index ",
+      stream_index,
+      " is not a valid stream. The file has ",
+      num_streams,
+      " streams, so the index must be in [0, ",
+      num_streams - 1,
+      "].");
+
+  AVMediaType media_type =
+      format_context_->streams[stream_index]->codecpar->codec_type;
+  const char* media_type_name = av_get_media_type_string(media_type);
+  STD_TORCH_CHECK(
+      media_type == AVMEDIA_TYPE_VIDEO,
+      "The stream at index ",
+      stream_index,
+      " is not a video stream, it is of type '",
+      media_type_name == nullptr ? "unknown" : media_type_name,
+      "'. Only video streams can be demuxed.");
+}
+
+void Demuxer::select_stream(std::optional<int> stream_index) {
+  int status = avformat_find_stream_info(format_context_.get(), nullptr);
   STD_TORCH_CHECK(
       status >= 0,
       "Failed to find stream info: ",
       get_ffmpeg_error_string_from_error_code(status));
 
+  if (stream_index.has_value()) {
+    validate_requested_stream(*stream_index);
+  }
+
   active_stream_index_ = av_find_best_stream(
       format_context_.get(),
       AVMEDIA_TYPE_VIDEO,
@@ -75,9 +135,8 @@ Demuxer::Demuxer(
       /*flags=*/0);
   STD_TORCH_CHECK(
       active_stream_index_ >= 0,
-      "No valid video stream found in input file (requested index ",
-      stream_index.value_or(-1),
-      ").");
+      "No valid video stream found in input file. Only video streams are "
+      "supported: audio streams cannot be demuxed.");
   stream_ = format_context_->streams[active_stream_index_];
 
   // We only need packets from the active stream, so tell FFmpeg to discard the

src/torchcodec/_core/Demuxer.h

@@ -6,9 +6,11 @@
 
 #pragma once
 
+#include <memory>
 #include <optional>
 #include <string>
 
+#include "AVIOContextHolder.h"
 #include "FFMPEGCommon.h"
 #include "StableABICompat.h"
 
@@ -36,6 +38,10 @@ class FORCE_PUBLIC_VISIBILITY Demuxer {
       const std::string& file_path,
       std::optional<int> stream_index = std::nullopt);
 
+  explicit Demuxer(
+      std::unique_ptr<AVIOContextHolder> avio_context_holder,
+      std::optional<int> stream_index = std::nullopt);
+
   // Returns the next packet for the active stream as a freshly-allocated
   // packet, or a null packet at end of stream.
   UniqueAVPacket next_packet();
@@ -55,6 +61,12 @@ class FORCE_PUBLIC_VISIBILITY Demuxer {
   }
 
  private:
+  void validate_requested_stream(int stream_index);
+  void select_stream(std::optional<int> stream_index);
+
+  // Declared before format_context_ so that it outlives it: the format context
+  // reads through the AVIOContext this holds.
+  std::unique_ptr<AVIOContextHolder> avio_context_holder_;
   UniqueDecodingAVFormatContext format_context_;
   int active_stream_index_ = -1;
   AVStream* stream_ = nullptr;