diff --git a/README.md b/README.md index 01248787b..9d89c75b7 100644 --- a/README.md +++ b/README.md @@ -211,6 +211,20 @@ format you want. Refer to Nvidia's GPU support matrix for more details conda install -c conda-forge "torchcodec=*=*cuda*" ``` +### Specify FFmpeg Path Manually +If torchcodec cannot detect the FFmpeg installation correctly, you can set the `TORCHCODEC_FFMPEG_DIR` environment variable to the directory containing the FFmpeg shared libraries. + +For example, in conda environments this is typically: + +```bash +# Conda on Linux/macOS +export TORCHCODEC_FFMPEG_DIR="$CONDA_PREFIX/lib" +``` +```powershell +# Conda on Windows +$env:TORCHCODEC_FFMPEG_DIR = "$env:CONDA_PREFIX\Library\bin" +``` + ## Benchmark Results The following was generated by running [our benchmark script](./benchmarks/decoders/generate_readme_data.py) on a lightly loaded 22-core machine with an Nvidia A100 with diff --git a/src/torchcodec/_core/ops.py b/src/torchcodec/_core/ops.py index 7a5c9dc33..a7040b0bd 100644 --- a/src/torchcodec/_core/ops.py +++ b/src/torchcodec/_core/ops.py @@ -5,6 +5,8 @@ # LICENSE file in the root directory of this source tree. +import ctypes +import glob import io import json import os @@ -22,13 +24,30 @@ expose_ffmpeg_dlls = nullcontext -if sys.platform == "win32" and hasattr(os, "add_dll_directory"): +if ffmpeg_dir := os.getenv("TORCHCODEC_FFMPEG_DIR"): + if hasattr(os, "add_dll_directory"): + + def expose_ffmpeg_dlls(): # type: ignore[no-redef] # noqa: F811 + return os.add_dll_directory(str(ffmpeg_dir)) + + else: + ffmpeg_library_glob = "*.dylib" if sys.platform == "darwin" else "*.so*" + ffmpeg_library_paths = glob.glob(os.path.join(ffmpeg_dir, ffmpeg_library_glob)) + if not ffmpeg_library_paths: + raise RuntimeError( + "TORCHCODEC_FFMPEG_DIR is set, but no FFmpeg shared libraries " + f"were found in {repr(ffmpeg_dir)}." + ) + for ffmpeg_library_path in ffmpeg_library_paths: + ctypes.CDLL(ffmpeg_library_path) + +elif sys.platform == "win32" and hasattr(os, "add_dll_directory"): # On windows we try to locate the FFmpeg DLLs and temporarily add them to # the DLL search path. This seems to be needed on some users machine, but # not on our CI. We don't know why. if ffmpeg_path := shutil.which("ffmpeg"): - def expose_ffmpeg_dlls(): # noqa: F811 + def expose_ffmpeg_dlls(): # type: ignore[no-redef] # noqa: F811 ffmpeg_dir = Path(ffmpeg_path).parent.absolute() return os.add_dll_directory(str(ffmpeg_dir)) # that's the actual CM