diff --git a/torchaudio/_extension.py b/torchaudio/_extension.py index e1f5eaed22e..6a6a2a8b32c 100644 --- a/torchaudio/_extension.py +++ b/torchaudio/_extension.py @@ -53,6 +53,28 @@ def _load_lib(lib: str) -> bool: return True +_FFMPEG_INITIALIZED = False + + +def _init_ffmpeg(): + global _FFMPEG_INITIALIZED + if _FFMPEG_INITIALIZED: + return + + if not torch.ops.torchaudio.is_ffmpeg_available(): + raise RuntimeError("torchaudio is not compiled with FFmpeg integration.") + + try: + _load_lib("libtorchaudio_ffmpeg") + except OSError as err: + raise ImportError("FFmpeg libraries are not found. Please install FFmpeg.") from err + + import torchaudio._torchaudio_ffmpeg # noqa + torch.ops.torchaudio.ffmpeg_init() + + _FFMPEG_INITIALIZED = True + + def _init_extension(): if not _mod_utils.is_module_available("torchaudio._torchaudio"): warnings.warn("torchaudio C++ extension is not available.") @@ -63,5 +85,14 @@ def _init_extension(): # This has to happen after the base library is loaded from torchaudio import _torchaudio # noqa + # Because this part is executed as part of `import torchaudio`, we ignore the + # initialization failure. + # If the FFmpeg integration is not properly initialized, then detailed error + # will be raised when client code attempts to import the dedicated feature. + try: + _init_ffmpeg() + except Exception: + pass + _init_extension() diff --git a/torchaudio/csrc/CMakeLists.txt b/torchaudio/csrc/CMakeLists.txt index 5e8b2e5f22a..e73a82f39ce 100644 --- a/torchaudio/csrc/CMakeLists.txt +++ b/torchaudio/csrc/CMakeLists.txt @@ -108,6 +108,14 @@ if(OpenMP_CXX_FOUND) ) endif() +if(USE_FFMPEG) + list( + APPEND + LIBTORCHAUDIO_COMPILE_DEFINITIONS + USE_FFMPEG + ) +endif() + #------------------------------------------------------------------------------# # END OF CUSTOMIZATION LOGICS #------------------------------------------------------------------------------# diff --git a/torchaudio/csrc/utils.cpp b/torchaudio/csrc/utils.cpp index e8050bb69a3..9fa807fd7df 100644 --- a/torchaudio/csrc/utils.cpp +++ b/torchaudio/csrc/utils.cpp @@ -20,11 +20,22 @@ bool is_kaldi_available() { #endif } +// It tells whether torchaudio was compiled with ffmpeg +// not the runtime availability. +bool is_ffmpeg_available() { +#ifdef USE_FFMPEG + return true; +#else + return false; +#endif +} + } // namespace TORCH_LIBRARY_FRAGMENT(torchaudio, m) { m.def("torchaudio::is_sox_available", &is_sox_available); m.def("torchaudio::is_kaldi_available", &is_kaldi_available); + m.def("torchaudio::is_ffmpeg_available", &is_ffmpeg_available); } } // namespace torchaudio diff --git a/torchaudio/io/__init__.py b/torchaudio/io/__init__.py index 65f01e1a957..fbee8c7daee 100644 --- a/torchaudio/io/__init__.py +++ b/torchaudio/io/__init__.py @@ -1,4 +1,5 @@ -_INITIALIZED = False +import torchaudio + _LAZILY_IMPORTED = [ "StreamReader", "StreamReaderSourceStream", @@ -8,32 +9,10 @@ ] -def _init_extension(): - import torch - import torchaudio - - try: - torchaudio._extension._load_lib("libtorchaudio_ffmpeg") - import torchaudio._torchaudio_ffmpeg - except OSError as err: - raise ImportError( - "Stream API requires FFmpeg libraries (libavformat and such). Please install FFmpeg 4." - ) from err - try: - torch.ops.torchaudio.ffmpeg_init() - except RuntimeError as err: - raise RuntimeError( - "Stream API requires FFmpeg binding. Please set USE_FFMPEG=1 when building from source." - ) from err - - global _INITIALIZED - _INITIALIZED = True - - def __getattr__(name: str): if name in _LAZILY_IMPORTED: - if not _INITIALIZED: - _init_extension() + + torchaudio._extension._init_ffmpeg() from . import _stream_reader