_extension.py 3.5 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103
  1. import os
  2. import warnings
  3. from pathlib import Path
  4. import torch
  5. from torchaudio._internal import module_utils as _mod_utils # noqa: F401
  6. _LIB_DIR = Path(__file__).parent / "lib"
  7. def _get_lib_path(lib: str):
  8. suffix = "pyd" if os.name == "nt" else "so"
  9. path = _LIB_DIR / f"{lib}.{suffix}"
  10. return path
  11. def _load_lib(lib: str) -> bool:
  12. """Load extension module
  13. Note:
  14. In case `torchaudio` is deployed with `pex` format, the library file
  15. is not in a standard location.
  16. In this case, we expect that `libtorchaudio` is available somewhere
  17. in the search path of dynamic loading mechanism, so that importing
  18. `_torchaudio` will have library loader find and load `libtorchaudio`.
  19. This is the reason why the function should not raising an error when the library
  20. file is not found.
  21. Returns:
  22. bool:
  23. True if the library file is found AND the library loaded without failure.
  24. False if the library file is not found (like in the case where torchaudio
  25. is deployed with pex format, thus the shared library file is
  26. in a non-standard location.).
  27. If the library file is found but there is an issue loading the library,
  28. (such as missing dependency) then this function raises the exception as-is.
  29. Raises:
  30. Exception:
  31. If the library file is found, but there is an issue loading the library file,
  32. (when underlying `ctype.DLL` throws an exception), this function will pass
  33. the exception as-is, instead of catching it and returning bool.
  34. The expected case is `OSError` thrown by `ctype.DLL` when a dynamic dependency
  35. is not found.
  36. This behavior was chosen because the expected failure case is not recoverable.
  37. If a dependency is missing, then users have to install it.
  38. """
  39. path = _get_lib_path(lib)
  40. if not path.exists():
  41. return False
  42. torch.ops.load_library(path)
  43. torch.classes.load_library(path)
  44. return True
  45. _FFMPEG_INITIALIZED = False
  46. def _init_ffmpeg():
  47. global _FFMPEG_INITIALIZED
  48. if _FFMPEG_INITIALIZED:
  49. return
  50. if not torch.ops.torchaudio.is_ffmpeg_available():
  51. raise RuntimeError(
  52. "torchaudio is not compiled with FFmpeg integration. Please set USE_FFMPEG=1 when compiling torchaudio."
  53. )
  54. try:
  55. _load_lib("libtorchaudio_ffmpeg")
  56. except OSError as err:
  57. raise ImportError("FFmpeg libraries are not found. Please install FFmpeg.") from err
  58. import torchaudio._torchaudio_ffmpeg # noqa
  59. torch.ops.torchaudio.ffmpeg_init()
  60. if torch.ops.torchaudio.ffmpeg_get_log_level() > 8:
  61. torch.ops.torchaudio.ffmpeg_set_log_level(8)
  62. _FFMPEG_INITIALIZED = True
  63. def _init_extension():
  64. if not _mod_utils.is_module_available("torchaudio._torchaudio"):
  65. warnings.warn("torchaudio C++ extension is not available.")
  66. return
  67. _load_lib("libtorchaudio")
  68. # This import is for initializing the methods registered via PyBind11
  69. # This has to happen after the base library is loaded
  70. from torchaudio import _torchaudio # noqa
  71. # Because this part is executed as part of `import torchaudio`, we ignore the
  72. # initialization failure.
  73. # If the FFmpeg integration is not properly initialized, then detailed error
  74. # will be raised when client code attempts to import the dedicated feature.
  75. try:
  76. _init_ffmpeg()
  77. except Exception:
  78. pass
  79. _init_extension()