torchcodec 0.7.0__cp313-cp313-win_amd64.whl → 0.8.1__cp313-cp313-win_amd64.whl

torchcodec/_core/CudaDeviceInterface.h

@@ -6,8 +6,9 @@
 
 #pragma once
 
-#include <npp.h>
+#include "src/torchcodec/_core/CUDACommon.h"
 #include "src/torchcodec/_core/DeviceInterface.h"
+#include "src/torchcodec/_core/FilterGraph.h"
 
 namespace facebook::torchcodec {
 
@@ -19,19 +20,50 @@ class CudaDeviceInterface : public DeviceInterface {
 
   std::optional<const AVCodec*> findCodec(const AVCodecID& codecId) override;
 
-  void initializeContext(AVCodecContext* codecContext) override;
+  void initialize(
+      const AVStream* avStream,
+      const UniqueDecodingAVFormatContext& avFormatCtx,
+      const SharedAVCodecContext& codecContext) override;
 
-  void convertAVFrameToFrameOutput(
+  void initializeVideo(
       const VideoStreamOptions& videoStreamOptions,
-      const AVRational& timeBase,
+      [[maybe_unused]] const std::vector<std::unique_ptr<Transform>>&
+          transforms,
+      [[maybe_unused]] const std::optional<FrameDims>& resizedOutputDims)
+      override;
+
+  void registerHardwareDeviceWithCodec(AVCodecContext* codecContext) override;
+
+  void convertAVFrameToFrameOutput(
       UniqueAVFrame& avFrame,
       FrameOutput& frameOutput,
       std::optional<torch::Tensor> preAllocatedOutputTensor =
           std::nullopt) override;
 
+  std::string getDetails() override;
+
  private:
-  UniqueAVBufferRef ctx_;
-  std::unique_ptr<NppStreamContext> nppCtx_;
+  // Our CUDA decoding code assumes NV12 format. In order to handle other
+  // kinds of input, we need to convert them to NV12. Our current implementation
+  // does this using filtergraph.
+  UniqueAVFrame maybeConvertAVFrameToNV12OrRGB24(UniqueAVFrame& avFrame);
+
+  // We sometimes encounter frames that cannot be decoded on the CUDA device.
+  // Rather than erroring out, we decode them on the CPU.
+  std::unique_ptr<DeviceInterface> cpuInterface_;
+
+  VideoStreamOptions videoStreamOptions_;
+  AVRational timeBase_;
+
+  UniqueAVBufferRef hardwareDeviceCtx_;
+  UniqueNppContext nppCtx_;
+
+  // This filtergraph instance is only used for NV12 format conversion in
+  // maybeConvertAVFrameToNV12().
+  std::unique_ptr<FiltersContext> nv12ConversionContext_;
+  std::unique_ptr<FilterGraph> nv12Conversion_;
+
+  bool usingCPUFallback_ = false;
 };
 
 } // namespace facebook::torchcodec

torchcodec/_core/DeviceInterface.cpp

@@ -11,7 +11,8 @@
 namespace facebook::torchcodec {
 
 namespace {
-using DeviceInterfaceMap = std::map<torch::DeviceType, CreateDeviceInterfaceFn>;
+using DeviceInterfaceMap =
+    std::map<DeviceInterfaceKey, CreateDeviceInterfaceFn>;
 static std::mutex g_interface_mutex;
 
 DeviceInterfaceMap& getDeviceMap() {
@@ -30,50 +31,87 @@ std::string getDeviceType(const std::string& device) {
 } // namespace
 
 bool registerDeviceInterface(
-    torch::DeviceType deviceType,
+    const DeviceInterfaceKey& key,
     CreateDeviceInterfaceFn createInterface) {
   std::scoped_lock lock(g_interface_mutex);
   DeviceInterfaceMap& deviceMap = getDeviceMap();
 
   TORCH_CHECK(
-      deviceMap.find(deviceType) == deviceMap.end(),
-      "Device interface already registered for ",
-      deviceType);
-  deviceMap.insert({deviceType, createInterface});
+      deviceMap.find(key) == deviceMap.end(),
+      "Device interface already registered for device type ",
+      key.deviceType,
+      " variant '",
+      key.variant,
+      "'");
+  deviceMap.insert({key, createInterface});
 
   return true;
 }
 
-torch::Device createTorchDevice(const std::string device) {
+void validateDeviceInterface(
+    const std::string device,
+    const std::string variant) {
   std::scoped_lock lock(g_interface_mutex);
   std::string deviceType = getDeviceType(device);
+
   DeviceInterfaceMap& deviceMap = getDeviceMap();
 
+  // Find device interface that matches device type and variant
+  torch::DeviceType deviceTypeEnum = torch::Device(deviceType).type();
+
   auto deviceInterface = std::find_if(
       deviceMap.begin(),
       deviceMap.end(),
-      [&](const std::pair<torch::DeviceType, CreateDeviceInterfaceFn>& arg) {
-        return device.rfind(
-                   torch::DeviceTypeName(arg.first, /*lcase*/ true), 0) == 0;
+      [&](const std::pair<DeviceInterfaceKey, CreateDeviceInterfaceFn>& arg) {
+        return arg.first.deviceType == deviceTypeEnum &&
+            arg.first.variant == variant;
       });
-  TORCH_CHECK(
-      deviceInterface != deviceMap.end(), "Unsupported device: ", device);
 
-  return torch::Device(device);
+  TORCH_CHECK(
+      deviceInterface != deviceMap.end(),
+      "Unsupported device: ",
+      device,
+      " (device type: ",
+      deviceType,
+      ", variant: ",
+      variant,
+      ")");
 }
 
 std::unique_ptr<DeviceInterface> createDeviceInterface(
-    const torch::Device& device) {
-  auto deviceType = device.type();
+    const torch::Device& device,
+    const std::string_view variant) {
+  DeviceInterfaceKey key(device.type(), variant);
   std::scoped_lock lock(g_interface_mutex);
   DeviceInterfaceMap& deviceMap = getDeviceMap();
 
+  auto it = deviceMap.find(key);
+  if (it != deviceMap.end()) {
+    return std::unique_ptr<DeviceInterface>(it->second(device));
+  }
+
   TORCH_CHECK(
-      deviceMap.find(deviceType) != deviceMap.end(),
-      "Unsupported device: ",
-      device);
+      false,
+      "No device interface found for device type: ",
+      device.type(),
+      " variant: '",
+      variant,
+      "'");
+}
 
-  return std::unique_ptr<DeviceInterface>(deviceMap[deviceType](device));
+torch::Tensor rgbAVFrameToTensor(const UniqueAVFrame& avFrame) {
+  TORCH_CHECK_EQ(avFrame->format, AV_PIX_FMT_RGB24);
+
+  int height = avFrame->height;
+  int width = avFrame->width;
+  std::vector<int64_t> shape = {height, width, 3};
+  std::vector<int64_t> strides = {avFrame->linesize[0], 3, 1};
+  AVFrame* avFrameClone = av_frame_clone(avFrame.get());
+  auto deleter = [avFrameClone](void*) {
+    UniqueAVFrame avFrameToDelete(avFrameClone);
+  };
+  return torch::from_blob(
+      avFrameClone->data[0], shape, strides, deleter, {torch::kUInt8});
 }
 
 } // namespace facebook::torchcodec

torchcodec/_core/DeviceInterface.h

@@ -14,16 +14,27 @@
 #include "FFMPEGCommon.h"
 #include "src/torchcodec/_core/Frame.h"
 #include "src/torchcodec/_core/StreamOptions.h"
+#include "src/torchcodec/_core/Transform.h"
 
 namespace facebook::torchcodec {
 
-// Note that all these device functions should only be called if the device is
-// not a CPU device. CPU device functions are already implemented in the
-// SingleStreamDecoder implementation.
-// These functions should only be called from within an if block like this:
-// if (device.type() != torch::kCPU) {
-//   deviceFunction(device, ...);
-// }
+// Key for device interface registration with device type + variant support
+struct DeviceInterfaceKey {
+  torch::DeviceType deviceType;
+  std::string_view variant = "ffmpeg"; // e.g., "ffmpeg", "beta", etc.
+
+  bool operator<(const DeviceInterfaceKey& other) const {
+    if (deviceType != other.deviceType) {
+      return deviceType < other.deviceType;
+    }
+    return variant < other.variant;
+  }
+
+  explicit DeviceInterfaceKey(torch::DeviceType type) : deviceType(type) {}
+
+  DeviceInterfaceKey(torch::DeviceType type, const std::string_view& variant)
+      : deviceType(type), variant(variant) {}
+};
 
 class DeviceInterface {
  public:
@@ -35,33 +46,103 @@ class DeviceInterface {
     return device_;
   };
 
-  virtual std::optional<const AVCodec*> findCodec(const AVCodecID& codecId) = 0;
+  virtual std::optional<const AVCodec*> findCodec(
+      [[maybe_unused]] const AVCodecID& codecId) {
+    return std::nullopt;
+  };
 
-  // Initialize the hardware device that is specified in `device`. Some builds
-  // support CUDA and others only support CPU.
-  virtual void initializeContext(AVCodecContext* codecContext) = 0;
+  // Initialize the device with parameters generic to all kinds of decoding.
+  virtual void initialize(
+      const AVStream* avStream,
+      const UniqueDecodingAVFormatContext& avFormatCtx,
+      const SharedAVCodecContext& codecContext) = 0;
+
+  // Initialize the device with parameters specific to video decoding. There is
+  // a default empty implementation.
+  virtual void initializeVideo(
+      [[maybe_unused]] const VideoStreamOptions& videoStreamOptions,
+      [[maybe_unused]] const std::vector<std::unique_ptr<Transform>>&
+          transforms,
+      [[maybe_unused]] const std::optional<FrameDims>& resizedOutputDims) {}
+
+  // In order for decoding to actually happen on an FFmpeg managed hardware
+  // device, we need to register the DeviceInterface managed
+  // AVHardwareDeviceContext with the AVCodecContext. We don't need to do this
+  // on the CPU and if FFmpeg is not managing the hardware device.
+  virtual void registerHardwareDeviceWithCodec(
+      [[maybe_unused]] AVCodecContext* codecContext) {}
 
   virtual void convertAVFrameToFrameOutput(
-      const VideoStreamOptions& videoStreamOptions,
-      const AVRational& timeBase,
       UniqueAVFrame& avFrame,
       FrameOutput& frameOutput,
       std::optional<torch::Tensor> preAllocatedOutputTensor = std::nullopt) = 0;
 
+  // ------------------------------------------
+  // Extension points for custom decoding paths
+  // ------------------------------------------
+
+  // Returns AVSUCCESS on success, AVERROR(EAGAIN) if decoder queue full, or
+  // other AVERROR on failure
+  // Default implementation uses FFmpeg directly
+  virtual int sendPacket(ReferenceAVPacket& avPacket) {
+    TORCH_CHECK(
+        codecContext_ != nullptr,
+        "Codec context not available for default packet sending");
+    return avcodec_send_packet(codecContext_.get(), avPacket.get());
+  }
+
+  // Send an EOF packet to flush the decoder
+  // Returns AVSUCCESS on success, or other AVERROR on failure
+  // Default implementation uses FFmpeg directly
+  virtual int sendEOFPacket() {
+    TORCH_CHECK(
+        codecContext_ != nullptr,
+        "Codec context not available for default EOF packet sending");
+    return avcodec_send_packet(codecContext_.get(), nullptr);
+  }
+
+  // Returns AVSUCCESS on success, AVERROR(EAGAIN) if no frame ready,
+  // AVERROR_EOF if end of stream, or other AVERROR on failure
+  // Default implementation uses FFmpeg directly
+  virtual int receiveFrame(UniqueAVFrame& avFrame) {
+    TORCH_CHECK(
+        codecContext_ != nullptr,
+        "Codec context not available for default frame receiving");
+    return avcodec_receive_frame(codecContext_.get(), avFrame.get());
+  }
+
+  // Flush remaining frames from decoder
+  virtual void flush() {
+    TORCH_CHECK(
+        codecContext_ != nullptr,
+        "Codec context not available for default flushing");
+    avcodec_flush_buffers(codecContext_.get());
+  }
+
+  virtual std::string getDetails() {
+    return "";
+  }
+
  protected:
   torch::Device device_;
+  SharedAVCodecContext codecContext_;
 };
 
 using CreateDeviceInterfaceFn =
     std::function<DeviceInterface*(const torch::Device& device)>;
 
 bool registerDeviceInterface(
-    torch::DeviceType deviceType,
+    const DeviceInterfaceKey& key,
     const CreateDeviceInterfaceFn createInterface);
 
-torch::Device createTorchDevice(const std::string device);
+void validateDeviceInterface(
+    const std::string device,
+    const std::string variant);
 
 std::unique_ptr<DeviceInterface> createDeviceInterface(
-    const torch::Device& device);
+    const torch::Device& device,
+    const std::string_view variant = "ffmpeg");
+
+torch::Tensor rgbAVFrameToTensor(const UniqueAVFrame& avFrame);
 
 } // namespace facebook::torchcodec