Stragglers

Author: scotts

src/torchcodec/_core/BetaCudaDeviceInterface.cpp

@@ -156,13 +156,20 @@ BetaCudaDeviceInterface::~BetaCudaDeviceInterface() {
   }
 }
 
-void BetaCudaDeviceInterface::initializeInterface(AVStream* avStream) {
+void BetaCudaDeviceInterface::initialize(const AVStream* avStream) {
   torch::Tensor dummyTensorForCudaInitialization = torch::empty(
       {1}, torch::TensorOptions().dtype(torch::kUInt8).device(device_));
 
   TORCH_CHECK(avStream != nullptr, "AVStream cannot be null");
   timeBase_ = avStream->time_base;
 
+  auto cudaDevice = torch::Device(torch::kCUDA);
+  defaultCudaInterface_ =
+      std::unique_ptr<DeviceInterface>(createDeviceInterface(cudaDevice));
+  AVCodecContext dummyCodecContext = {};
+  defaultCudaInterface_->initialize(avStream);
+  defaultCudaInterface_->registerHardwareDeviceWithCodec(&dummyCodecContext);
+
   const AVCodecParameters* codecpar = avStream->codecpar;
   TORCH_CHECK(codecpar != nullptr, "CodecParameters cannot be null");
 
@@ -523,8 +530,6 @@ void BetaCudaDeviceInterface::flush() {
 }
 
 void BetaCudaDeviceInterface::convertAVFrameToFrameOutput(
-    const VideoStreamOptions& videoStreamOptions,
-    const AVRational& timeBase,
     UniqueAVFrame& avFrame,
     FrameOutput& frameOutput,
     std::optional<torch::Tensor> preAllocatedOutputTensor) {
@@ -535,20 +540,8 @@ void BetaCudaDeviceInterface::convertAVFrameToFrameOutput(
   // TODONVDEC P1: we use the 'default' cuda device interface for color
   // conversion. That's a temporary hack to make things work. we should abstract
   // the color conversion stuff separately.
-  if (!defaultCudaInterface_) {
-    auto cudaDevice = torch::Device(torch::kCUDA);
-    defaultCudaInterface_ =
-        std::unique_ptr<DeviceInterface>(createDeviceInterface(cudaDevice));
-    AVCodecContext dummyCodecContext = {};
-    defaultCudaInterface_->initializeContext(&dummyCodecContext);
-  }
-
   defaultCudaInterface_->convertAVFrameToFrameOutput(
-      videoStreamOptions,
-      timeBase,
-      avFrame,
-      frameOutput,
-      preAllocatedOutputTensor);
+      avFrame, frameOutput, preAllocatedOutputTensor);
 }
 
 BetaCudaDeviceInterface::FrameBuffer::Slot*

src/torchcodec/_core/BetaCudaDeviceInterface.h

@@ -37,11 +37,9 @@ class BetaCudaDeviceInterface : public DeviceInterface {
   explicit BetaCudaDeviceInterface(const torch::Device& device);
   virtual ~BetaCudaDeviceInterface();
 
-  void initializeInterface(AVStream* stream) override;
+  void initialize(const AVStream* avStream) override;
 
   void convertAVFrameToFrameOutput(
-      const VideoStreamOptions& videoStreamOptions,
-      const AVRational& timeBase,
       UniqueAVFrame& avFrame,
       FrameOutput& frameOutput,
       std::optional<torch::Tensor> preAllocatedOutputTensor =

src/torchcodec/_core/CpuDeviceInterface.cpp

@@ -46,10 +46,9 @@ CpuDeviceInterface::CpuDeviceInterface(const torch::Device& device)
       device_.type() == torch::kCPU, "Unsupported device: ", device_.str());
 }
 
-void CpuDeviceInterface::initialize(
-    [[maybe_unused]] AVCodecContext* codecContext,
-    const AVRational& timeBase) {
-  timeBase_ = timeBase;
+void CpuDeviceInterface::initialize(const AVStream* avStream) {
+  TORCH_CHECK(avStream != nullptr, "avStream is null");
+  timeBase_ = avStream->time_base;
 }
 
 void CpuDeviceInterface::initializeVideo(

src/torchcodec/_core/CpuDeviceInterface.h

@@ -23,9 +23,7 @@ class CpuDeviceInterface : public DeviceInterface {
     return std::nullopt;
   }
 
-  virtual void initialize(
-      [[maybe_unused]] AVCodecContext* codecContext,
-      const AVRational& timeBase) override;
+  virtual void initialize(const AVStream* avStream) override;
 
   virtual void initializeVideo(
       const VideoStreamOptions& videoStreamOptions,

src/torchcodec/_core/CudaDeviceInterface.h

@@ -20,8 +20,7 @@ class CudaDeviceInterface : public DeviceInterface {
 
   std::optional<const AVCodec*> findCodec(const AVCodecID& codecId) override;
 
-  void initialize(AVCodecContext* codecContext, const AVRational& timeBase)
-      override;
+  void initialize(const AVStream* avStream) override;
 
   void initializeVideo(
       const VideoStreamOptions& videoStreamOptions,
@@ -30,6 +29,8 @@ class CudaDeviceInterface : public DeviceInterface {
       [[maybe_unused]] const std::optional<FrameDims>& resizedOutputDims)
       override;
 
+  void registerHardwareDeviceWithCodec(AVCodecContext* codecContext) override;
+
   void convertAVFrameToFrameOutput(
       UniqueAVFrame& avFrame,
       FrameOutput& frameOutput,
@@ -42,6 +43,10 @@ class CudaDeviceInterface : public DeviceInterface {
   // 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_;