Skip to content

Commit 042a35e

Browse files
committed
Merge branch 'main' of github.com:pytorch/torchcodec into fallback-colorconversion
2 parents 89f1547 + 28b0346 commit 042a35e

13 files changed

+391
-23
lines changed

packaging/build_ffmpeg.bat

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,3 +1,9 @@
1+
:: Copyright (c) Meta Platforms, Inc. and affiliates.
2+
:: All rights reserved.
3+
::
4+
:: This source code is licensed under the BSD-style license found in the
5+
:: LICENSE file in the root directory of this source tree.
6+
17
:: Taken from torchaudio
28
@echo off
39

packaging/build_ffmpeg.sh

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,9 @@
11
#!/usr/bin/env bash
2+
# Copyright (c) Meta Platforms, Inc. and affiliates.
3+
# All rights reserved.
4+
#
5+
# This source code is licensed under the BSD-style license found in the
6+
# LICENSE file in the root directory of this source tree.
27

38
# This is taken and adapated from torchaudio, only keeping the parts relevant to
49
# linux.

packaging/check_glibcxx.py

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,3 +1,9 @@
1+
# Copyright (c) Meta Platforms, Inc. and affiliates.
2+
# All rights reserved.
3+
#
4+
# This source code is licensed under the BSD-style license found in the
5+
# LICENSE file in the root directory of this source tree.
6+
17
"""
28
The goal of this script is to ensure that the .so files we ship do not contain
39
symbol versions from libstdc++ that are too recent. This is a very manual way of

packaging/helpers.sh

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,9 @@
11
#!/usr/bin/env bash
2+
# Copyright (c) Meta Platforms, Inc. and affiliates.
3+
# All rights reserved.
4+
#
5+
# This source code is licensed under the BSD-style license found in the
6+
# LICENSE file in the root directory of this source tree.
27

38
_list_wheel_files() {
49
unzip -l "$1" | awk '{print $4}'

packaging/post_build_script.sh

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,9 @@
11
#!/bin/bash
2+
# Copyright (c) Meta Platforms, Inc. and affiliates.
3+
# All rights reserved.
4+
#
5+
# This source code is licensed under the BSD-style license found in the
6+
# LICENSE file in the root directory of this source tree.
27

38
set -ex
49

packaging/pre_build_script.sh

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,9 @@
11
#!/bin/bash
2+
# Copyright (c) Meta Platforms, Inc. and affiliates.
3+
# All rights reserved.
4+
#
5+
# This source code is licensed under the BSD-style license found in the
6+
# LICENSE file in the root directory of this source tree.
27

38
set -ex
49

packaging/vc_env_helper.bat

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,3 +1,9 @@
1+
:: Copyright (c) Meta Platforms, Inc. and affiliates.
2+
:: All rights reserved.
3+
::
4+
:: This source code is licensed under the BSD-style license found in the
5+
:: LICENSE file in the root directory of this source tree.
6+
17
:: Taken from torchaudio
28
@echo on
39

src/torchcodec/_core/BetaCudaDeviceInterface.cpp

Lines changed: 15 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -15,7 +15,7 @@
1515
#include "src/torchcodec/_core/FFMPEGCommon.h"
1616
#include "src/torchcodec/_core/NVDECCache.h"
1717

18-
// #include <cuda_runtime.h> // For cudaStreamSynchronize
18+
#include "src/torchcodec/_core/NVCUVIDRuntimeLoader.h"
1919
#include "src/torchcodec/_core/nvcuvid_include/cuviddec.h"
2020
#include "src/torchcodec/_core/nvcuvid_include/nvcuvid.h"
2121

@@ -155,6 +155,7 @@ std::optional<cudaVideoCodec> validateCodecSupport(AVCodecID codecId) {
155155
bool nativeNVDECSupport(const SharedAVCodecContext& codecContext) {
156156
// Return true iff the input video stream is supported by our NVDEC
157157
// implementation.
158+
158159
auto codecType = validateCodecSupport(codecContext->codec_id);
159160
if (!codecType.has_value()) {
160161
return false;
@@ -228,6 +229,8 @@ BetaCudaDeviceInterface::BetaCudaDeviceInterface(const torch::Device& device)
228229

229230
initializeCudaContextWithPytorch(device_);
230231
nppCtx_ = getNppStreamContext(device_);
232+
233+
nvcuvidAvailable_ = loadNVCUVIDLibrary();
231234
}
232235

233236
BetaCudaDeviceInterface::~BetaCudaDeviceInterface() {
@@ -255,7 +258,7 @@ void BetaCudaDeviceInterface::initialize(
255258
const AVStream* avStream,
256259
const UniqueDecodingAVFormatContext& avFormatCtx,
257260
[[maybe_unused]] const SharedAVCodecContext& codecContext) {
258-
if (!nativeNVDECSupport(codecContext)) {
261+
if (!nvcuvidAvailable_ || !nativeNVDECSupport(codecContext)) {
259262
cpuFallback_ = createDeviceInterface(torch::kCPU);
260263
TORCH_CHECK(
261264
cpuFallback_ != nullptr, "Failed to create CPU device interface");
@@ -825,8 +828,16 @@ void BetaCudaDeviceInterface::convertAVFrameToFrameOutput(
825828
}
826829

827830
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;
830841
}
831842

832843
} // namespace facebook::torchcodec

src/torchcodec/_core/BetaCudaDeviceInterface.h

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -100,6 +100,7 @@ class BetaCudaDeviceInterface : public DeviceInterface {
100100
UniqueNppContext nppCtx_;
101101

102102
std::unique_ptr<DeviceInterface> cpuFallback_;
103+
bool nvcuvidAvailable_ = false;
103104
UniqueSwsContext swsContext_;
104105
SwsFrameContext prevSwsFrameContext_;
105106
};

src/torchcodec/_core/CMakeLists.txt

Lines changed: 1 addition & 19 deletions
Original file line numberDiff line numberDiff line change
@@ -99,7 +99,7 @@ function(make_torchcodec_libraries
9999
)
100100

101101
if(ENABLE_CUDA)
102-
list(APPEND core_sources CudaDeviceInterface.cpp BetaCudaDeviceInterface.cpp NVDECCache.cpp CUDACommon.cpp)
102+
list(APPEND core_sources CudaDeviceInterface.cpp BetaCudaDeviceInterface.cpp NVDECCache.cpp CUDACommon.cpp NVCUVIDRuntimeLoader.cpp)
103103
endif()
104104

105105
set(core_library_dependencies
@@ -108,27 +108,9 @@ function(make_torchcodec_libraries
108108
)
109109

110110
if(ENABLE_CUDA)
111-
# Try to find NVCUVID. Try the normal way first. This should work locally.
112-
find_library(NVCUVID_LIBRARY NAMES nvcuvid)
113-
# If not found, try with version suffix, or hardcoded path. Appears
114-
# to be necessary on the CI.
115-
if(NOT NVCUVID_LIBRARY)
116-
find_library(NVCUVID_LIBRARY NAMES nvcuvid.1 PATHS /usr/lib64 /usr/lib)
117-
endif()
118-
if(NOT NVCUVID_LIBRARY)
119-
set(NVCUVID_LIBRARY "/usr/lib64/libnvcuvid.so.1")
120-
endif()
121-
122-
if(NVCUVID_LIBRARY)
123-
message(STATUS "Found NVCUVID: ${NVCUVID_LIBRARY}")
124-
else()
125-
message(FATAL_ERROR "Could not find NVCUVID library")
126-
endif()
127-
128111
list(APPEND core_library_dependencies
129112
${CUDA_nppi_LIBRARY}
130113
${CUDA_nppicc_LIBRARY}
131-
${NVCUVID_LIBRARY}
132114
)
133115
endif()
134116

0 commit comments

Comments
 (0)