Support multiple FFmpeg versions (#3464)

Summary: Pull Request resolved: #3464

Differential Revision: D47300223

Pulled By: mthrok

fbshipit-source-id: b9562a0597f79d7feb6d3bf9172754610b01b39f

third_party/ffmpeg/CMakeLists.txt

@@ -2,7 +2,8 @@
 # This file defines the following FFmpeg libraries using pre-built binaries.
 
 add_library(ffmpeg4 INTERFACE)
-add_library(ffmpeg ALIAS ffmpeg4)
+add_library(ffmpeg5 INTERFACE)
+add_library(ffmpeg6 INTERFACE)
 
 ################################################################################
 
@@ -17,12 +18,32 @@ if (APPLE)
       URL ${base_url}/2023-07-06/macos_arm64/4.1.8.tar.gz
       URL_HASH SHA256=a44b8152b7f204ce5050fc7f6fd2bbbafe7ae4e45f03e135f3b45dd9a08f404e
       )
+    FetchContent_Declare(
+      f5
+      URL ${base_url}/2023-07-06/macos_arm64/5.0.3.tar.gz
+      URL_HASH SHA256=316fe8378afadcf63089acf3ad53a626fd3c26cc558b96ce1dc94d2a78f4deb4
+      )
+    FetchContent_Declare(
+      f6
+      URL ${base_url}/2023-07-06/macos_arm64/6.0.tar.gz
+      URL_HASH SHA256=5d1da9626f8cb817d6c558a2c61085a3d39a8d9f725a6f69f4658bea8efa9389
+      )
   elseif ("${CMAKE_SYSTEM_PROCESSOR}" STREQUAL "x86_64")
     FetchContent_Declare(
       f4
       URL ${base_url}/2023-07-06/macos_x86_64/4.1.8.tar.gz
       URL_HASH SHA256=392d5af0b24535bfc69d6244e7595e5f07117b93d94505d0a4b34c82ae479f48
       )
+    FetchContent_Declare(
+      f5
+      URL ${base_url}/2023-07-06/macos_x86_64/5.0.3.tar.gz
+      URL_HASH SHA256=d0b49575d3b174cfcca53b3049641855e48028cf22dd32f3334bbec4ca94f43e
+      )
+    FetchContent_Declare(
+      f6
+      URL ${base_url}/2023-07-06/macos_x86_64/6.0.tar.gz
+      URL_HASH SHA256=eabc01eb7d9e714e484d5e1b27bf7d921e87c1f3c00334abd1729e158d6db862
+      )
   else ()
     message(
       FATAL_ERROR
@@ -34,13 +55,33 @@ elseif (UNIX)
       f4
       URL ${base_url}/2023-07-06/linux_aarch64/4.1.8.tar.gz
       URL_HASH SHA256=aae0b713040e30ceebe0d0bc82353d3d9054055c7af8a4f4abc1766015ab7681
-    )
+      )
+    FetchContent_Declare(
+      f5
+      URL ${base_url}/2023-07-06/linux_aarch64/5.0.3.tar.gz
+      URL_HASH SHA256=65c663206982ee3f0ff88436d8869d191b46061e01e753518c77ecc13ea0236d
+      )
+    FetchContent_Declare(
+      f6
+      URL ${base_url}/2023-07-06/linux_aarch64/6.0.tar.gz
+      URL_HASH SHA256=ec762fd41ea7b8d9ad4f810f53fd78a565f2bc6f680afe56d555c80f3d35adef
+      )
   elseif ("${CMAKE_SYSTEM_PROCESSOR}" STREQUAL "x86_64")
     FetchContent_Declare(
       f4
       URL ${base_url}/2023-07-06/linux_x86_64/4.1.8.tar.gz
       URL_HASH SHA256=52e53b8857739bdd54f9d8541e22569b57f6c6f16504ee83963c2ed3e7061a23
       )
+    FetchContent_Declare(
+      f5
+      URL ${base_url}/2023-07-06/linux_x86_64/5.0.3.tar.gz
+      URL_HASH SHA256=de3c75c99b9ce33de7efdc144566804ae5880457ce71e185b3f592dc452edce7
+      )
+    FetchContent_Declare(
+      f6
+      URL ${base_url}/2023-07-06/linux_x86_64/6.0.tar.gz
+      URL_HASH SHA256=04d3916404bab5efadd29f68361b7d13ea71e6242c6473edcb747a41a9fb97a6
+      )
   else ()
     # Possible case ppc64le (though it's not officially supported.)
     message(
@@ -53,10 +94,22 @@ elseif(MSVC)
     URL ${base_url}/2023-07-06/windows/4.1.8.tar.gz
     URL_HASH SHA256=c45cd36e0575490f97ace07365bb67c5e1cbe9f3e6a4272d035c19348df96790
     )
+  FetchContent_Declare(
+    f5
+    URL ${base_url}/2023-07-06/windows/5.0.3.tar.gz
+    URL_HASH SHA256=e2daa10799909e366cb1b4b91a217d35f6749290dcfeea40ecae3d5b05a46cb3
+    )
+  FetchContent_Declare(
+    f6
+    URL ${base_url}/2023-07-06/windows/6.0.tar.gz
+    URL_HASH SHA256=098347eca8cddb5aaa61e9ecc1a00548c645fc59b4f7346b3d91414aa00a9cf6
+    )
 endif()
 
-FetchContent_MakeAvailable(f4)
+FetchContent_MakeAvailable(f4 f5 f6)
 target_include_directories(ffmpeg4 INTERFACE ${f4_SOURCE_DIR}/include)
+target_include_directories(ffmpeg5 INTERFACE ${f5_SOURCE_DIR}/include)
+target_include_directories(ffmpeg6 INTERFACE ${f6_SOURCE_DIR}/include)
 
 if(APPLE)
   target_link_libraries(
@@ -68,6 +121,24 @@ if(APPLE)
     ${f4_SOURCE_DIR}/lib/libavdevice.58.dylib
     ${f4_SOURCE_DIR}/lib/libavfilter.7.dylib
     )
+  target_link_libraries(
+    ffmpeg5
+    INTERFACE
+    ${f5_SOURCE_DIR}/lib/libavutil.57.dylib
+    ${f5_SOURCE_DIR}/lib/libavcodec.59.dylib
+    ${f5_SOURCE_DIR}/lib/libavformat.59.dylib
+    ${f5_SOURCE_DIR}/lib/libavdevice.59.dylib
+    ${f5_SOURCE_DIR}/lib/libavfilter.8.dylib
+    )
+  target_link_libraries(
+    ffmpeg6
+    INTERFACE
+    ${f6_SOURCE_DIR}/lib/libavutil.58.dylib
+    ${f6_SOURCE_DIR}/lib/libavcodec.60.dylib
+    ${f6_SOURCE_DIR}/lib/libavformat.60.dylib
+    ${f6_SOURCE_DIR}/lib/libavdevice.60.dylib
+    ${f6_SOURCE_DIR}/lib/libavfilter.9.dylib
+    )
 elseif (UNIX)
   target_link_libraries(
     ffmpeg4
@@ -78,6 +149,24 @@ elseif (UNIX)
     ${f4_SOURCE_DIR}/lib/libavdevice.so.58
     ${f4_SOURCE_DIR}/lib/libavfilter.so.7
     )
+  target_link_libraries(
+    ffmpeg5
+    INTERFACE
+    ${f5_SOURCE_DIR}/lib/libavutil.so.57
+    ${f5_SOURCE_DIR}/lib/libavcodec.so.59
+    ${f5_SOURCE_DIR}/lib/libavformat.so.59
+    ${f5_SOURCE_DIR}/lib/libavdevice.so.59
+    ${f5_SOURCE_DIR}/lib/libavfilter.so.8
+    )
+  target_link_libraries(
+    ffmpeg6
+    INTERFACE
+    ${f6_SOURCE_DIR}/lib/libavutil.so.58
+    ${f6_SOURCE_DIR}/lib/libavcodec.so.60
+    ${f6_SOURCE_DIR}/lib/libavformat.so.60
+    ${f6_SOURCE_DIR}/lib/libavdevice.so.60
+    ${f6_SOURCE_DIR}/lib/libavfilter.so.9
+    )
 elseif(MSVC)
   target_link_libraries(
     ffmpeg4
@@ -88,4 +177,22 @@ elseif(MSVC)
     ${f4_SOURCE_DIR}/bin/avdevice.lib
     ${f4_SOURCE_DIR}/bin/avfilter.lib
     )
+  target_link_libraries(
+    ffmpeg5
+    INTERFACE
+    ${f5_SOURCE_DIR}/bin/avutil.lib
+    ${f5_SOURCE_DIR}/bin/avcodec.lib
+    ${f5_SOURCE_DIR}/bin/avformat.lib
+    ${f5_SOURCE_DIR}/bin/avdevice.lib
+    ${f5_SOURCE_DIR}/bin/avfilter.lib
+    )
+  target_link_libraries(
+    ffmpeg6
+    INTERFACE
+    ${f6_SOURCE_DIR}/bin/avutil.lib
+    ${f6_SOURCE_DIR}/bin/avcodec.lib
+    ${f6_SOURCE_DIR}/bin/avformat.lib
+    ${f6_SOURCE_DIR}/bin/avdevice.lib
+    ${f6_SOURCE_DIR}/bin/avfilter.lib
+    )
 endif()

tools/setup_helpers/extension.py

@@ -67,8 +67,12 @@ def get_ext_modules():
     if _USE_FFMPEG:
         modules.extend(
             [
-                Extension(name="torchaudio.lib.libtorchaudio_ffmpeg", sources=[]),
-                Extension(name="torchaudio.lib._torchaudio_ffmpeg", sources=[]),
+                Extension(name="torchaudio.lib.libtorchaudio_ffmpeg4", sources=[]),
+                Extension(name="torchaudio.lib._torchaudio_ffmpeg4", sources=[]),
+                Extension(name="torchaudio.lib.libtorchaudio_ffmpeg5", sources=[]),
+                Extension(name="torchaudio.lib._torchaudio_ffmpeg5", sources=[]),
+                Extension(name="torchaudio.lib.libtorchaudio_ffmpeg6", sources=[]),
+                Extension(name="torchaudio.lib._torchaudio_ffmpeg6", sources=[]),
             ]
         )
     return modules

torchaudio/_extension/__init__.py

@@ -21,6 +21,7 @@
     "_IS_RIR_AVAILABLE",
     "_SOX_INITIALIZED",
     "_FFMPEG_INITIALIZED",
+    "_FFMPEG_EXT",
 ]
 
 
@@ -60,9 +61,10 @@
 
 # Initialize FFmpeg-related features
 _FFMPEG_INITIALIZED = False
-if is_module_available("torchaudio.lib._torchaudio_ffmpeg"):
+_FFMPEG_EXT = None
+if _IS_TORCHAUDIO_EXT_AVAILABLE:
     try:
-        _init_ffmpeg()
+        _FFMPEG_EXT = _init_ffmpeg()
         _FFMPEG_INITIALIZED = True
     except Exception:
         # The initialization of FFmpeg extension will fail if supported FFmpeg

torchaudio/_extension/utils.py

@@ -6,15 +6,18 @@
 """
 
 
+import importlib
+import logging
 import os
+import platform
 from functools import wraps
 from pathlib import Path
 
 import torch
-
 import torchaudio
 from torchaudio._internal.module_utils import is_module_available
 
+_LG = logging.getLogger(__name__)
 _LIB_DIR = Path(__file__).parent.parent / "lib"
 
 
@@ -76,21 +79,63 @@ def _init_sox():
 
 
 def _init_ffmpeg():
-    if not is_module_available("torchaudio.lib._torchaudio_ffmpeg"):
-        raise RuntimeError(
-            "torchaudio is not compiled with FFmpeg integration. Please set USE_FFMPEG=1 when compiling torchaudio."
-        )
+    import logging
+
+    logging.basicConfig(level=logging.DEBUG, format="%(message)s")
+
+    # Library path for non-OSS
+    if _get_lib_path("_torchaudio_ffmpeg").exists():
+        from torchaudio.lib import _torchaudio_ffmpeg
+
+        _torchaudio_ffmpeg.init()
+        if _torchaudio_ffmpeg.get_log_level() > 8:
+            _torchaudio_ffmpeg.set_log_level(8)
+        return _torchaudio_ffmpeg
+
+    # For environment where we do not know which FFmpegs are available.
+    libname_template = {
+        "Linux": "libavutil.so.{ver}",
+        "Darwin": "libavutil.{ver}.dylib",
+        "Windows": "avutil-{ver}.dll",
+    }[platform.system()]
 
-    try:
-        _load_lib("libtorchaudio_ffmpeg")
-    except OSError as err:
-        raise ImportError("FFmpeg libraries are not found. Please install FFmpeg.") from err
+    for ffmpeg_ver, avutil_ver in [(6, 58), (5, 57), (4, 56)]:
+        library = f"libtorchaudio_ffmpeg{ffmpeg_ver}"
+        extension = f"_torchaudio_ffmpeg{ffmpeg_ver}"
 
-    import torchaudio.lib._torchaudio_ffmpeg  # noqa
+        if not _get_lib_path(library).exists():
+            _LG.debug("The library (%s) is not found. Please build torchaudio with USE_FFMPEG=1.", library)
+            continue
 
-    torchaudio.lib._torchaudio_ffmpeg.init()
-    if torchaudio.lib._torchaudio_ffmpeg.get_log_level() > 8:
-        torchaudio.lib._torchaudio_ffmpeg.set_log_level(8)
+        if not _get_lib_path(extension).exists():
+            _LG.debug("The extension module (%s) is not found. Please build torchaudio with USE_FFMPEG=1.", extension)
+            continue
+
+        libavutil = libname_template.format(ver=avutil_ver)
+
+        try:
+            torchaudio.lib._torchaudio.find_avutil(libavutil)
+        except Exception as e:
+            _LG.debug("%s was not found. (%s)", libavutil, e)
+            continue
+
+        try:
+            _load_lib(library)
+        except Exception as e:
+            _LG.debug("Failed to load library, %s. (%s)", library, e)
+            continue
+
+        mod = importlib.import_module(f"torchaudio.lib.{extension}")
+
+        mod.init()
+        if mod.get_log_level() > 8:
+            mod.set_log_level(8)
+        return mod
+    raise ImportError(
+        "Failed to initialize FFmpeg. "
+        "To see the detail of the initialization attempt, "
+        "please set the log level of torchaudio logger to DEBUG."
+    )
 
 
 def _init_dll_path():

torchaudio/csrc/ffmpeg/CMakeLists.txt

@@ -25,10 +25,26 @@ if (USE_CUDA)
 endif()
 
 torchaudio_library(
-  libtorchaudio_ffmpeg
+  libtorchaudio_ffmpeg4
   "${sources}"
   ""
-  "torch;ffmpeg;${additional_lib}"
+  "torch;ffmpeg4;${additional_lib}"
+  ""
+  )
+
+torchaudio_library(
+  libtorchaudio_ffmpeg5
+  "${sources}"
+  ""
+  "torch;ffmpeg5;${additional_lib}"
+  ""
+  )
+
+torchaudio_library(
+  libtorchaudio_ffmpeg6
+  "${sources}"
+  ""
+  "torch;ffmpeg6;${additional_lib}"
   ""
   )
 
@@ -38,10 +54,24 @@ if (BUILD_TORCHAUDIO_PYTHON_EXTENSION)
     pybind/pybind.cpp
     )
   torchaudio_extension(
-    _torchaudio_ffmpeg
+    _torchaudio_ffmpeg4
     "${ext_sources}"
     ""
-    "libtorchaudio_ffmpeg"
+    "libtorchaudio_ffmpeg4"
+    "EXT_NAME=_torchaudio_ffmpeg4"
+    )
+  torchaudio_extension(
+    _torchaudio_ffmpeg5
+    "${ext_sources}"
+    ""
+    "libtorchaudio_ffmpeg5"
+    "EXT_NAME=_torchaudio_ffmpeg5"
+    )
+  torchaudio_extension(
+    _torchaudio_ffmpeg6
+    "${ext_sources}"
     ""
+    "libtorchaudio_ffmpeg6"
+    "EXT_NAME=_torchaudio_ffmpeg6"
     )
 endif ()

torchaudio/csrc/ffmpeg/pybind/pybind.cpp

@@ -186,7 +186,7 @@ struct StreamWriterFileObj : private FileObj, public StreamWriterCustomIO {
             py::hasattr(fileobj, "seek") ? &seek_func : nullptr) {}
 };
 
-PYBIND11_MODULE(_torchaudio_ffmpeg, m) {
+PYBIND11_MODULE(EXT_NAME, m) {
   m.def("init", []() { avdevice_register_all(); });
   m.def("get_log_level", []() { return av_log_get_level(); });
   m.def("set_log_level", [](int level) { av_log_set_level(level); });

torchaudio/csrc/pybind/pybind.cpp

@@ -8,6 +8,7 @@ PYBIND11_MODULE(_torchaudio, m) {
   m.def("is_rir_available", &is_rir_available, "");
   m.def("is_align_available", &is_align_available, "");
   m.def("cuda_version", &cuda_version, "");
+  m.def("find_avutil", &find_avutil, "");
 }
 
 } // namespace

torchaudio/csrc/utils.cpp

@@ -1,4 +1,4 @@
-#include <torch/script.h>
+#include <ATen/DynamicLibrary.h>
 #include <torchaudio/csrc/utils.h>
 
 #ifdef USE_CUDA
@@ -31,4 +31,10 @@ c10::optional<int64_t> cuda_version() {
 #endif
 }
 
+int find_avutil(const char* name) {
+  auto lib = at::DynamicLibrary{name};
+  auto avutil_version = (unsigned (*)(void))(lib.sym("avutil_version"));
+  return avutil_version() >> 16;
+}
+
 } // namespace torchaudio