@@ -156,13 +156,20 @@ BetaCudaDeviceInterface::~BetaCudaDeviceInterface() {
156
156
}
157
157
}
158
158
159
- void BetaCudaDeviceInterface::initializeInterface ( AVStream* avStream) {
159
+ void BetaCudaDeviceInterface::initialize ( const AVStream* avStream) {
160
160
torch::Tensor dummyTensorForCudaInitialization = torch::empty (
161
161
{1 }, torch::TensorOptions ().dtype (torch::kUInt8 ).device (device_));
162
162
163
163
TORCH_CHECK (avStream != nullptr , " AVStream cannot be null" );
164
164
timeBase_ = avStream->time_base ;
165
165
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
+
166
173
const AVCodecParameters* codecpar = avStream->codecpar ;
167
174
TORCH_CHECK (codecpar != nullptr , " CodecParameters cannot be null" );
168
175
@@ -523,8 +530,6 @@ void BetaCudaDeviceInterface::flush() {
523
530
}
524
531
525
532
void BetaCudaDeviceInterface::convertAVFrameToFrameOutput (
526
- const VideoStreamOptions& videoStreamOptions,
527
- const AVRational& timeBase,
528
533
UniqueAVFrame& avFrame,
529
534
FrameOutput& frameOutput,
530
535
std::optional<torch::Tensor> preAllocatedOutputTensor) {
@@ -535,20 +540,8 @@ void BetaCudaDeviceInterface::convertAVFrameToFrameOutput(
535
540
// TODONVDEC P1: we use the 'default' cuda device interface for color
536
541
// conversion. That's a temporary hack to make things work. we should abstract
537
542
// 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
-
546
543
defaultCudaInterface_->convertAVFrameToFrameOutput (
547
- videoStreamOptions,
548
- timeBase,
549
- avFrame,
550
- frameOutput,
551
- preAllocatedOutputTensor);
544
+ avFrame, frameOutput, preAllocatedOutputTensor);
552
545
}
553
546
554
547
BetaCudaDeviceInterface::FrameBuffer::Slot*
0 commit comments