|
15 | 15 | #include "src/torchcodec/_core/FFMPEGCommon.h" |
16 | 16 | #include "src/torchcodec/_core/NVDECCache.h" |
17 | 17 |
|
18 | | -// #include <cuda_runtime.h> // For cudaStreamSynchronize |
| 18 | +#include "src/torchcodec/_core/NVCUVIDRuntimeLoader.h" |
19 | 19 | #include "src/torchcodec/_core/nvcuvid_include/cuviddec.h" |
20 | 20 | #include "src/torchcodec/_core/nvcuvid_include/nvcuvid.h" |
21 | 21 |
|
@@ -155,6 +155,7 @@ std::optional<cudaVideoCodec> validateCodecSupport(AVCodecID codecId) { |
155 | 155 | bool nativeNVDECSupport(const SharedAVCodecContext& codecContext) { |
156 | 156 | // Return true iff the input video stream is supported by our NVDEC |
157 | 157 | // implementation. |
| 158 | + |
158 | 159 | auto codecType = validateCodecSupport(codecContext->codec_id); |
159 | 160 | if (!codecType.has_value()) { |
160 | 161 | return false; |
@@ -222,6 +223,8 @@ BetaCudaDeviceInterface::BetaCudaDeviceInterface(const torch::Device& device) |
222 | 223 |
|
223 | 224 | initializeCudaContextWithPytorch(device_); |
224 | 225 | nppCtx_ = getNppStreamContext(device_); |
| 226 | + |
| 227 | + nvcuvidAvailable_ = loadNVCUVIDLibrary(); |
225 | 228 | } |
226 | 229 |
|
227 | 230 | BetaCudaDeviceInterface::~BetaCudaDeviceInterface() { |
@@ -249,7 +252,7 @@ void BetaCudaDeviceInterface::initialize( |
249 | 252 | const AVStream* avStream, |
250 | 253 | const UniqueDecodingAVFormatContext& avFormatCtx, |
251 | 254 | [[maybe_unused]] const SharedAVCodecContext& codecContext) { |
252 | | - if (!nativeNVDECSupport(codecContext)) { |
| 255 | + if (!nvcuvidAvailable_ || !nativeNVDECSupport(codecContext)) { |
253 | 256 | cpuFallback_ = createDeviceInterface(torch::kCPU); |
254 | 257 | TORCH_CHECK( |
255 | 258 | cpuFallback_ != nullptr, "Failed to create CPU device interface"); |
@@ -700,8 +703,16 @@ void BetaCudaDeviceInterface::convertAVFrameToFrameOutput( |
700 | 703 | } |
701 | 704 |
|
702 | 705 | std::string BetaCudaDeviceInterface::getDetails() { |
703 | | - return std::string("Beta CUDA Device Interface. Using ") + |
704 | | - (cpuFallback_ ? "CPU fallback." : "NVDEC."); |
| 706 | + std::string details = "Beta CUDA Device Interface."; |
| 707 | + if (cpuFallback_) { |
| 708 | + details += " Using CPU fallback."; |
| 709 | + if (!nvcuvidAvailable_) { |
| 710 | + details += " NVCUVID not available!"; |
| 711 | + } |
| 712 | + } else { |
| 713 | + details += " Using NVDEC."; |
| 714 | + } |
| 715 | + return details; |
705 | 716 | } |
706 | 717 |
|
707 | 718 | } // namespace facebook::torchcodec |
0 commit comments