# Copyright (c) Meta Platforms, Inc. and affiliates. # All rights reserved. # # This source code is licensed under the BSD-style license found in the # LICENSE file in the root directory of this source tree. import functools import importlib import importlib.util import sys import traceback from pathlib import Path from types import ModuleType import torch # Note that this value must match the value used as PYBIND_OPS_MODULE_NAME when we compile _core/pybind_ops.cpp. # If the values do not match, we will not be able to import the C++ shared library as a Python module at runtime. _PYBIND_OPS_MODULE_NAME = "core_pybind_ops" # Copy pasted from torchvision # https://github.com/pytorch/vision/blob/947ae1dc71867f28021d5bc0ff3a19c249236e2a/torchvision/_internally_replaced_utils.py#L25 def _get_extension_path(lib_name: str) -> str: extension_suffixes = [] if sys.platform == "linux" or sys.platform.startswith("freebsd"): extension_suffixes = importlib.machinery.EXTENSION_SUFFIXES elif sys.platform == "darwin": extension_suffixes = importlib.machinery.EXTENSION_SUFFIXES + [".dylib"] elif sys.platform in ("win32", "cygwin"): extension_suffixes = importlib.machinery.EXTENSION_SUFFIXES + [".dll", ".pyd"] else: raise NotImplementedError(f"{sys.platform = } is not not supported") loader_details = ( importlib.machinery.ExtensionFileLoader, extension_suffixes, ) extfinder = importlib.machinery.FileFinder( str(Path(__file__).parent), loader_details ) ext_specs = extfinder.find_spec(lib_name) if ext_specs is None: raise ImportError(f"No spec found for {lib_name}") if ext_specs.origin is None: raise ImportError(f"Existing spec found for {lib_name} does not have an origin") return ext_specs.origin def _load_pybind11_module(module_name: str, library_path: str) -> ModuleType: spec = importlib.util.spec_from_file_location( module_name, library_path, ) if spec is None or spec.loader is None: raise ImportError( f"Unable to load spec or spec.loader for module {module_name} from path {library_path}" ) mod = importlib.util.module_from_spec(spec) spec.loader.exec_module(mod) return mod def load_image_library() -> None: """Load the FFmpeg-free image library""" image_library_path = _get_extension_path("libtorchcodec_image") torch.ops.load_library(image_library_path) @functools.cache def load_pybind_ops() -> ModuleType: """Load the pybind module providing file-like support (create_file_like_context). This module links no FFmpeg, so it's loaded eagerly (regardless of whether FFmpeg is available) and shared by both the image encoders and the FFmpeg encoders/decoders. """ library_path = _get_extension_path("libtorchcodec_pybind_ops") return _load_pybind11_module(_PYBIND_OPS_MODULE_NAME, library_path) @functools.cache def load_heic_library() -> None: """Load the HEIC library (libtorchcodec_heic). It's loaded lazily, only when HEIC decoding is requested, never at import: because it links libheif, an optional runtime dependency we don't bundle. importing torchcodec must not require libheif to be installed. """ torch.ops.load_library(_get_extension_path("libtorchcodec_heic")) @functools.cache def load_core_libraries() -> tuple[int, str]: """Load the FFmpeg-dependent shared libraries, memoizing the result. This raises if the libraries cannot be loaded, typically because FFmpeg couldn't be found. This is called: - at import time (import torchcodec) within a try/except, to determine whether FFmpeg is available, and then load the corresponding ops if it is. If FFmpeg isn't available, we don't propagate the exception and `import torchcodec` still works: we still want to support the image decoders, which don't need FFmpeg. - at runtime in all FFmpeg-dependent public entry-points like VideoDecoder, so that if FFmpeg isn't available the user gets a clear error message. Since this is @functools.cache'd, the libraries are only ever loaded once on success. Failures are *not* cached (functools.cache only caches return values), so when FFmpeg is missing the load is retried and re-raises on each call. That's fine, it only happens on the error path. We successively try to load the shared libraries for each version of FFmpeg that we support, starting with the highest version and working our way down. Once we can load the core and custom-ops libraries for a version of FFmpeg, we have succeeded and we stop. We load them via torch.ops.load_library(), which registers the custom ops with PyTorch's runtime so they're accessible through torch.ops. """ exceptions = [] for ffmpeg_major_version in (9, 8, 7, 6, 5, 4): core_library_name = f"libtorchcodec_core{ffmpeg_major_version}" custom_ops_library_name = f"libtorchcodec_custom_ops{ffmpeg_major_version}" try: core_library_path = _get_extension_path(core_library_name) torch.ops.load_library(core_library_path) torch.ops.load_library(_get_extension_path(custom_ops_library_name)) return (ffmpeg_major_version, core_library_path) except Exception: # Capture the full traceback for this exception exc_traceback = traceback.format_exc() exceptions.append((ffmpeg_major_version, exc_traceback)) traceback_info = ( "\n[start of libtorchcodec loading traceback]\n" + "\n".join(f"FFmpeg version {v}:\n{tb}" for v, tb in exceptions) + "[end of libtorchcodec loading traceback]." ) raise RuntimeError( f"""Could not load libtorchcodec. Likely causes: 1. FFmpeg is not properly installed in your environment. We support versions 4, 5, 6, 7, 8, and 9, and we attempt to load libtorchcodec for each of those versions. Errors for versions not installed on your system are expected; only the error for your installed FFmpeg version is relevant. On Windows, ensure you've installed the "full-shared" version which ships DLLs. 2. The PyTorch version ({torch.__version__}) is not compatible with this version of TorchCodec. Refer to the version compatibility table: https://github.com/pytorch/torchcodec?tab=readme-ov-file#installing-torchcodec. 3. Another runtime dependency; see exceptions below. The following exceptions were raised as we tried to load libtorchcodec: """ f"{traceback_info}" )