@@ -156,13 +156,20 @@ BetaCudaDeviceInterface::~BetaCudaDeviceInterface() {
156156 }
157157}
158158
159- void BetaCudaDeviceInterface::initializeInterface ( AVStream* avStream) {
159+ void BetaCudaDeviceInterface::initialize ( const AVStream* avStream) {
160160 torch::Tensor dummyTensorForCudaInitialization = torch::empty (
161161 {1 }, torch::TensorOptions ().dtype (torch::kUInt8 ).device (device_));
162162
163163 TORCH_CHECK (avStream != nullptr , " AVStream cannot be null" );
164164 timeBase_ = avStream->time_base ;
165165
166+ auto cudaDevice = torch::Device (torch::kCUDA );
167+ defaultCudaInterface_ =
168+ std::unique_ptr<DeviceInterface>(createDeviceInterface (cudaDevice));
169+ AVCodecContext dummyCodecContext = {};
170+ defaultCudaInterface_->initialize (avStream);
171+ defaultCudaInterface_->registerHardwareDeviceWithCodec (&dummyCodecContext);
172+
166173 const AVCodecParameters* codecpar = avStream->codecpar ;
167174 TORCH_CHECK (codecpar != nullptr , " CodecParameters cannot be null" );
168175
@@ -523,8 +530,6 @@ void BetaCudaDeviceInterface::flush() {
523530}
524531
525532void BetaCudaDeviceInterface::convertAVFrameToFrameOutput (
526- const VideoStreamOptions& videoStreamOptions,
527- const AVRational& timeBase,
528533 UniqueAVFrame& avFrame,
529534 FrameOutput& frameOutput,
530535 std::optional<torch::Tensor> preAllocatedOutputTensor) {
@@ -535,20 +540,8 @@ void BetaCudaDeviceInterface::convertAVFrameToFrameOutput(
535540 // TODONVDEC P1: we use the 'default' cuda device interface for color
536541 // conversion. That's a temporary hack to make things work. we should abstract
537542 // the color conversion stuff separately.
538- if (!defaultCudaInterface_) {
539- auto cudaDevice = torch::Device (torch::kCUDA );
540- defaultCudaInterface_ =
541- std::unique_ptr<DeviceInterface>(createDeviceInterface (cudaDevice));
542- AVCodecContext dummyCodecContext = {};
543- defaultCudaInterface_->initializeContext (&dummyCodecContext);
544- }
545-
546543 defaultCudaInterface_->convertAVFrameToFrameOutput (
547- videoStreamOptions,
548- timeBase,
549- avFrame,
550- frameOutput,
551- preAllocatedOutputTensor);
544+ avFrame, frameOutput, preAllocatedOutputTensor);
552545}
553546
554547BetaCudaDeviceInterface::FrameBuffer::Slot*
0 commit comments