|
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; |
@@ -228,6 +229,8 @@ BetaCudaDeviceInterface::BetaCudaDeviceInterface(const torch::Device& device) |
228 | 229 |
|
229 | 230 | initializeCudaContextWithPytorch(device_); |
230 | 231 | nppCtx_ = getNppStreamContext(device_); |
| 232 | + |
| 233 | + nvcuvidAvailable_ = loadNVCUVIDLibrary(); |
231 | 234 | } |
232 | 235 |
|
233 | 236 | BetaCudaDeviceInterface::~BetaCudaDeviceInterface() { |
@@ -255,7 +258,7 @@ void BetaCudaDeviceInterface::initialize( |
255 | 258 | const AVStream* avStream, |
256 | 259 | const UniqueDecodingAVFormatContext& avFormatCtx, |
257 | 260 | [[maybe_unused]] const SharedAVCodecContext& codecContext) { |
258 | | - if (!nativeNVDECSupport(codecContext)) { |
| 261 | + if (!nvcuvidAvailable_ || !nativeNVDECSupport(codecContext)) { |
259 | 262 | cpuFallback_ = createDeviceInterface(torch::kCPU); |
260 | 263 | TORCH_CHECK( |
261 | 264 | cpuFallback_ != nullptr, "Failed to create CPU device interface"); |
@@ -825,8 +828,16 @@ void BetaCudaDeviceInterface::convertAVFrameToFrameOutput( |
825 | 828 | } |
826 | 829 |
|
827 | 830 | std::string BetaCudaDeviceInterface::getDetails() { |
828 | | - return std::string("Beta CUDA Device Interface. Using ") + |
829 | | - (cpuFallback_ ? "CPU fallback." : "NVDEC."); |
| 831 | + std::string details = "Beta CUDA Device Interface."; |
| 832 | + if (cpuFallback_) { |
| 833 | + details += " Using CPU fallback."; |
| 834 | + if (!nvcuvidAvailable_) { |
| 835 | + details += " NVCUVID not available!"; |
| 836 | + } |
| 837 | + } else { |
| 838 | + details += " Using NVDEC."; |
| 839 | + } |
| 840 | + return details; |
830 | 841 | } |
831 | 842 |
|
832 | 843 | } // namespace facebook::torchcodec |
0 commit comments