torchcodec 0.10.0__cp312-cp312-manylinux_2_28_x86_64.whl

torchcodec/_core/DeviceInterface.cpp

@@ -0,0 +1,117 @@
+// Copyright (c) Meta Platforms, Inc. and affiliates.
+// All rights reserved.
+//
+// This source code is licensed under the BSD-style license found in the
+// LICENSE file in the root directory of this source tree.
+
+#include "DeviceInterface.h"
+#include <map>
+#include <mutex>
+
+namespace facebook::torchcodec {
+
+namespace {
+using DeviceInterfaceMap =
+    std::map<DeviceInterfaceKey, CreateDeviceInterfaceFn>;
+static std::mutex g_interface_mutex;
+
+DeviceInterfaceMap& getDeviceMap() {
+  static DeviceInterfaceMap deviceMap;
+  return deviceMap;
+}
+
+std::string getDeviceType(const std::string& device) {
+  size_t pos = device.find(':');
+  if (pos == std::string::npos) {
+    return device;
+  }
+  return device.substr(0, pos);
+}
+
+} // namespace
+
+bool registerDeviceInterface(
+    const DeviceInterfaceKey& key,
+    CreateDeviceInterfaceFn createInterface) {
+  std::scoped_lock lock(g_interface_mutex);
+  DeviceInterfaceMap& deviceMap = getDeviceMap();
+
+  TORCH_CHECK(
+      deviceMap.find(key) == deviceMap.end(),
+      "Device interface already registered for device type ",
+      key.deviceType,
+      " variant '",
+      key.variant,
+      "'");
+  deviceMap.insert({key, createInterface});
+
+  return true;
+}
+
+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<DeviceInterfaceKey, CreateDeviceInterfaceFn>& arg) {
+        return arg.first.deviceType == deviceTypeEnum &&
+            arg.first.variant == variant;
+      });
+
+  TORCH_CHECK(
+      deviceInterface != deviceMap.end(),
+      "Unsupported device: ",
+      device,
+      " (device type: ",
+      deviceType,
+      ", variant: ",
+      variant,
+      ")");
+}
+
+std::unique_ptr<DeviceInterface> createDeviceInterface(
+    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(
+      false,
+      "No device interface found for device type: ",
+      device.type(),
+      " variant: '",
+      variant,
+      "'");
+}
+
+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

@@ -0,0 +1,191 @@
+// Copyright (c) Meta Platforms, Inc. and affiliates.
+// All rights reserved.
+//
+// This source code is licensed under the BSD-style license found in the
+// LICENSE file in the root directory of this source tree.
+
+#pragma once
+
+#include <torch/types.h>
+#include <functional>
+#include <memory>
+#include <stdexcept>
+#include <string>
+#include "FFMPEGCommon.h"
+#include "Frame.h"
+#include "StreamOptions.h"
+#include "Transform.h"
+
+namespace facebook::torchcodec {
+
+// 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:
+  DeviceInterface(const torch::Device& device) : device_(device) {}
+
+  virtual ~DeviceInterface(){};
+
+  torch::Device& device() {
+    return device_;
+  };
+
+  virtual std::optional<const AVCodec*> findCodec(
+      [[maybe_unused]] const AVCodecID& codecId,
+      [[maybe_unused]] bool isDecoder = true) {
+    return std::nullopt;
+  };
+
+  // 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) {}
+
+  // Initialize the device with parameters specific to audio decoding. There is
+  // a default empty implementation.
+  virtual void initializeAudio(
+      [[maybe_unused]] const AudioStreamOptions& audioStreamOptions) {}
+
+  // Flush any remaining samples from the audio resampler buffer.
+  // When sample rate conversion is involved, some samples may be buffered
+  // between frames for proper interpolation. This function flushes those
+  // buffered samples.
+  // Returns an optional tensor containing the flushed samples, or std::nullopt
+  // if there are no buffered samples or audio is not supported.
+  virtual std::optional<torch::Tensor> maybeFlushAudioBuffers() {
+    return std::nullopt;
+  }
+
+  // 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(
+      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 "";
+  }
+
+  // Pixel format used for encoding on CUDA devices
+  static constexpr AVPixelFormat CUDA_ENCODING_PIXEL_FORMAT = AV_PIX_FMT_NV12;
+
+  // Function used for video encoding, only implemented in CudaDeviceInterface.
+  // It is here to isolate CUDA dependencies from CPU builds
+  // TODO Video-Encoder: Reconsider using video encoding functions in device
+  // interface
+  virtual UniqueAVFrame convertCUDATensorToAVFrameForEncoding(
+      [[maybe_unused]] const torch::Tensor& tensor,
+      [[maybe_unused]] int frameIndex,
+      [[maybe_unused]] AVCodecContext* codecContext) {
+    TORCH_CHECK(false);
+  }
+
+  // Function used for video encoding, only implemented in CudaDeviceInterface.
+  // It is here to isolate CUDA dependencies from CPU builds
+  virtual void setupHardwareFrameContextForEncoding(
+      [[maybe_unused]] AVCodecContext* codecContext) {
+    TORCH_CHECK(false);
+  }
+
+  virtual std::optional<const AVCodec*> findHardwareEncoder(
+      [[maybe_unused]] const AVCodecID& codecId) {
+    TORCH_CHECK(false);
+  }
+
+ protected:
+  torch::Device device_;
+  SharedAVCodecContext codecContext_;
+  AVMediaType avMediaType_;
+};
+
+using CreateDeviceInterfaceFn =
+    std::function<DeviceInterface*(const torch::Device& device)>;
+
+bool registerDeviceInterface(
+    const DeviceInterfaceKey& key,
+    const CreateDeviceInterfaceFn createInterface);
+
+void validateDeviceInterface(
+    const std::string device,
+    const std::string variant);
+
+std::unique_ptr<DeviceInterface> createDeviceInterface(
+    const torch::Device& device,
+    const std::string_view variant = "ffmpeg");
+
+torch::Tensor rgbAVFrameToTensor(const UniqueAVFrame& avFrame);
+
+} // namespace facebook::torchcodec