address comments

Author: Molly Xu

src/torchcodec/_core/BetaCudaDeviceInterface.cpp

@@ -814,7 +814,6 @@ UniqueAVFrame BetaCudaDeviceInterface::transferCpuFrameToGpuNV12(
 void BetaCudaDeviceInterface::convertAVFrameToFrameOutput(
     UniqueAVFrame& avFrame,
     FrameOutput& frameOutput,
-    [[maybe_unused]] AVMediaType mediaType,
     std::optional<torch::Tensor> preAllocatedOutputTensor) {
   UniqueAVFrame gpuFrame =
       cpuFallback_ ? transferCpuFrameToGpuNV12(avFrame) : std::move(avFrame);

src/torchcodec/_core/BetaCudaDeviceInterface.h

@@ -46,9 +46,7 @@ class BetaCudaDeviceInterface : public DeviceInterface {
   void convertAVFrameToFrameOutput(
       UniqueAVFrame& avFrame,
       FrameOutput& frameOutput,
-      AVMediaType mediaType,
-      std::optional<torch::Tensor> preAllocatedOutputTensor =
-          std::nullopt) override;
+      std::optional<torch::Tensor> preAllocatedOutputTensor) override;
 
   int sendPacket(ReferenceAVPacket& packet) override;
   int sendEOFPacket() override;

src/torchcodec/_core/CpuDeviceInterface.cpp

@@ -35,6 +35,7 @@ void CpuDeviceInterface::initializeVideo(
     const VideoStreamOptions& videoStreamOptions,
     const std::vector<std::unique_ptr<Transform>>& transforms,
     const std::optional<FrameDims>& resizedOutputDims) {
+  avMediaType_ = AVMEDIA_TYPE_VIDEO;
   videoStreamOptions_ = videoStreamOptions;
   resizedOutputDims_ = resizedOutputDims;
 
@@ -88,6 +89,7 @@ void CpuDeviceInterface::initializeVideo(
 
 void CpuDeviceInterface::initializeAudio(
     const AudioStreamOptions& audioStreamOptions) {
+  avMediaType_ = AVMEDIA_TYPE_AUDIO;
   audioStreamOptions_ = audioStreamOptions;
   initialized_ = true;
 }
@@ -123,11 +125,10 @@ ColorConversionLibrary CpuDeviceInterface::getColorConversionLibrary(
 void CpuDeviceInterface::convertAVFrameToFrameOutput(
     UniqueAVFrame& avFrame,
     FrameOutput& frameOutput,
-    AVMediaType mediaType,
     std::optional<torch::Tensor> preAllocatedOutputTensor) {
   TORCH_CHECK(initialized_, "CpuDeviceInterface was not initialized.");
 
-  if (mediaType == AVMEDIA_TYPE_AUDIO) {
+  if (avMediaType_ == AVMEDIA_TYPE_AUDIO) {
     convertAudioAVFrameToFrameOutput(avFrame, frameOutput);
   } else {
     convertVideoAVFrameToFrameOutput(
@@ -390,7 +391,8 @@ std::optional<torch::Tensor> CpuDeviceInterface::maybeFlushAudioBuffers() {
   if (!swrContext_) {
     return std::nullopt;
   }
-  auto numRemainingSamples = swr_get_out_samples(swrContext_.get(), 0);
+  auto numRemainingSamples = // this is an upper bound
+      swr_get_out_samples(swrContext_.get(), 0);
 
   if (numRemainingSamples == 0) {
     return std::nullopt;

src/torchcodec/_core/CpuDeviceInterface.h

@@ -41,9 +41,7 @@ class CpuDeviceInterface : public DeviceInterface {
   void convertAVFrameToFrameOutput(
       UniqueAVFrame& avFrame,
       FrameOutput& frameOutput,
-      AVMediaType mediaType,
-      std::optional<torch::Tensor> preAllocatedOutputTensor =
-          std::nullopt) override;
+      std::optional<torch::Tensor> preAllocatedOutputTensor) override;
 
   std::string getDetails() override;
 

src/torchcodec/_core/CudaDeviceInterface.cpp

@@ -238,7 +238,6 @@ UniqueAVFrame CudaDeviceInterface::maybeConvertAVFrameToNV12OrRGB24(
 void CudaDeviceInterface::convertAVFrameToFrameOutput(
     UniqueAVFrame& avFrame,
     FrameOutput& frameOutput,
-    [[maybe_unused]] AVMediaType mediaType,
     std::optional<torch::Tensor> preAllocatedOutputTensor) {
   validatePreAllocatedTensorShape(preAllocatedOutputTensor, avFrame);
 
@@ -272,8 +271,7 @@ void CudaDeviceInterface::convertAVFrameToFrameOutput(
     } else {
       // Reason 2 above. We need to do a full conversion which requires an
       // actual CPU device.
-      cpuInterface_->convertAVFrameToFrameOutput(
-          avFrame, cpuFrameOutput, AVMEDIA_TYPE_VIDEO);
+      cpuInterface_->convertAVFrameToFrameOutput(avFrame, cpuFrameOutput);
     }
 
     // Finally, we need to send the frame back to the GPU. Note that the

src/torchcodec/_core/CudaDeviceInterface.h

@@ -37,9 +37,7 @@ class CudaDeviceInterface : public DeviceInterface {
   void convertAVFrameToFrameOutput(
       UniqueAVFrame& avFrame,
       FrameOutput& frameOutput,
-      AVMediaType mediaType,
-      std::optional<torch::Tensor> preAllocatedOutputTensor =
-          std::nullopt) override;
+      std::optional<torch::Tensor> preAllocatedOutputTensor) override;
 
   std::string getDetails() override;
 

src/torchcodec/_core/DeviceInterface.h

@@ -90,7 +90,6 @@ class DeviceInterface {
   virtual void convertAVFrameToFrameOutput(
       UniqueAVFrame& avFrame,
       FrameOutput& frameOutput,
-      AVMediaType mediaType,
       std::optional<torch::Tensor> preAllocatedOutputTensor = std::nullopt) = 0;
 
   // ------------------------------------------
@@ -142,6 +141,7 @@ class DeviceInterface {
  protected:
   torch::Device device_;
   SharedAVCodecContext codecContext_;
+  AVMediaType avMediaType_;
 };
 
 using CreateDeviceInterfaceFn =

src/torchcodec/_core/SingleStreamDecoder.cpp

@@ -1289,15 +1289,14 @@ FrameOutput SingleStreamDecoder::convertAVFrameToFrameOutput(
     std::optional<torch::Tensor> preAllocatedOutputTensor) {
   // Convert the frame to tensor.
   FrameOutput frameOutput;
-  auto& streamInfo = streamInfos_[activeStreamIndex_];
   frameOutput.ptsSeconds = ptsToSeconds(
       getPtsOrDts(avFrame),
       formatContext_->streams[activeStreamIndex_]->time_base);
   frameOutput.durationSeconds = ptsToSeconds(
       getDuration(avFrame),
       formatContext_->streams[activeStreamIndex_]->time_base);
   deviceInterface_->convertAVFrameToFrameOutput(
-      avFrame, frameOutput, streamInfo.avMediaType, preAllocatedOutputTensor);
+      avFrame, frameOutput, preAllocatedOutputTensor);
   return frameOutput;
 }