separate encoding frame ctx init

Author: Dan-Flores

src/torchcodec/_core/CudaDeviceInterface.cpp

@@ -144,6 +144,12 @@ void CudaDeviceInterface::registerHardwareDeviceWithCodec(
       hardwareDeviceCtx_, "Hardware device context has not been initialized");
   TORCH_CHECK(codecContext != nullptr, "codecContext is null");
   codecContext->hw_device_ctx = av_buffer_ref(hardwareDeviceCtx_.get());
+}
+
+void CudaDeviceInterface::setupEncodingContext(AVCodecContext* codecContext) {
+  TORCH_CHECK(
+      hardwareDeviceCtx_, "Hardware device context has not been initialized");
+  TORCH_CHECK(codecContext != nullptr, "codecContext is null");
   // is there any way to preserve actual desired format?
   // codecContext->sw_pix_fmt = codecContext->pix_fmt;
   // Should we always produce AV_PIX_FMT_NV12?

src/torchcodec/_core/CudaDeviceInterface.h

@@ -35,6 +35,8 @@ class CudaDeviceInterface : public DeviceInterface {
 
   void registerHardwareDeviceWithCodec(AVCodecContext* codecContext) override;
 
+  void setupEncodingContext(AVCodecContext* codecContext) override;
+
   void convertAVFrameToFrameOutput(
       UniqueAVFrame& avFrame,
       FrameOutput& frameOutput,

src/torchcodec/_core/DeviceInterface.h

@@ -92,6 +92,12 @@ class DeviceInterface {
   virtual void registerHardwareDeviceWithCodec(
       [[maybe_unused]] AVCodecContext* codecContext) {}
 
+  // Setup device-specific encoding context (e.g., hardware frame contexts).
+  // Called after registerHardwareDeviceWithCodec for encoders.
+  // Default implementation does nothing (suitable for CPU and basic cases).
+  virtual void setupEncodingContext(
+      [[maybe_unused]] AVCodecContext* codecContext) {}
+
   virtual void convertAVFrameToFrameOutput(
       UniqueAVFrame& avFrame,
       FrameOutput& frameOutput,

src/torchcodec/_core/Encoder.cpp

@@ -850,6 +850,9 @@ void VideoEncoder::initializeEncoder(
   // context before calling avcodec_open2().
   deviceInterface_->registerHardwareDeviceWithCodec(avCodecContext_.get());
 
+  // Setup device-specific encoding context (e.g., hardware frame contexts)
+  deviceInterface_->setupEncodingContext(avCodecContext_.get());
+
   int status = avcodec_open2(avCodecContext_.get(), avCodec, &avCodecOptions);
   av_dict_free(&avCodecOptions);
 

test/test_encoders.py

@@ -775,7 +775,7 @@ def test_contiguity(self, method, tmp_path, device):
         num_frames, channels, height, width = 5, 3, 256, 256
         contiguous_frames = torch.randint(
             0, 256, size=(num_frames, channels, height, width), dtype=torch.uint8
-        ).contiguous()
+        ).contiguous().to(device)
         assert contiguous_frames.is_contiguous()
 
         # Permute NCHW to NHWC, then update the memory layout, then permute back