Swap decoder/encoder implementation

Differential Revision: D50677606

Pull Request resolved: #3681

cmake/TorchAudioHelper.cmake

@@ -43,7 +43,7 @@ endfunction()
 
 function(torio_library name source include_dirs link_libraries compile_defs)
   _library(
-    torchaudio/lib
+    torio/lib
     "${name}"
     "${source}"
     "${include_dirs}"
@@ -104,7 +104,7 @@ if (BUILD_TORCHAUDIO_PYTHON_EXTENSION)
   endfunction()
   function(torio_extension name sources include_dirs libraries definitions)
     _extension(
-      torchaudio/lib
+      torio/lib
       "${name}"
       "${sources}"
       "${include_dirs}"

docs/source/_templates/autosummary/io_class.rst

@@ -73,7 +73,7 @@ Support Structures
 
 {{ item | underline("~") }}
 
-.. autoclass:: torchaudio.io._stream_reader.{{item}}()
+.. autoclass:: torio.io._streaming_media_decoder.{{item}}()
    :members:
 
 {%- endfor %}

docs/source/conf.py

@@ -385,6 +385,13 @@ def inject_minigalleries(app, what, name, obj, options, lines):
 def setup(app):
     app.connect("autodoc-process-docstring", inject_minigalleries)
 
+    from torchaudio.io import StreamReader, StreamWriter
+
+    # need to assign the names here, otherwise autodoc won't document these classes,
+    # and will instead just say 'alias of ...'
+    StreamReader.__name__ = "StreamReader"
+    StreamWriter.__name__ = "StreamWriter"
+
 
 from custom_directives import CustomCardEnd, CustomCardItem, CustomCardStart, SupportedDevices, SupportedProperties
 

src/torchaudio/_backend/utils.py

@@ -4,8 +4,9 @@
 
 import torch
 
-from torchaudio._extension import lazy_import_ffmpeg_ext, lazy_import_sox_ext
+from torchaudio._extension import lazy_import_sox_ext
 from torchaudio.io import CodecConfig
+from torio._extension import lazy_import_ffmpeg_ext
 
 from . import soundfile_backend
 

src/torchaudio/_extension/__init__.py

@@ -4,7 +4,7 @@
 
 from torchaudio._internal.module_utils import fail_with_message, is_module_available, no_op
 
-from .utils import _check_cuda_version, _init_dll_path, _init_ffmpeg, _init_sox, _LazyImporter, _load_lib
+from .utils import _check_cuda_version, _init_dll_path, _init_sox, _LazyImporter, _load_lib
 
 _LG = logging.getLogger(__name__)
 
@@ -18,7 +18,6 @@
     "_IS_TORCHAUDIO_EXT_AVAILABLE",
     "_IS_RIR_AVAILABLE",
     "lazy_import_sox_ext",
-    "lazy_import_ffmpeg_ext",
 ]
 
 
@@ -57,18 +56,6 @@ def lazy_import_sox_ext():
     return _SOX_EXT
 
 
-_FFMPEG_EXT = None
-
-
-def lazy_import_ffmpeg_ext():
-    """Load FFmpeg integration based on availability in lazy manner"""
-
-    global _FFMPEG_EXT
-    if _FFMPEG_EXT is None:
-        _FFMPEG_EXT = _LazyImporter("_torchaudio_ffmpeg", _init_ffmpeg)
-    return _FFMPEG_EXT
-
-
 fail_if_no_rir = (
     no_op
     if _IS_RIR_AVAILABLE

src/torchaudio/_extension/utils.py

@@ -4,7 +4,6 @@
 They should not depend on external state.
 Anything that depends on external state should happen in __init__.py
 """
-
 import importlib
 import logging
 import os
@@ -107,57 +106,6 @@ def _init_sox():
     return ext
 
 
-_FFMPEG_VERS = ["6", "5", "4", ""]
-
-
-def _find_versionsed_ffmpeg_extension(version: str):
-    _LG.debug("Attempting to load FFmpeg%s", version)
-
-    ext = f"torchaudio.lib._torio_ffmpeg{version}"
-    lib = f"libtorio_ffmpeg{version}"
-
-    if not importlib.util.find_spec(ext):
-        raise RuntimeError(f"FFmpeg{version} extension is not available.")
-
-    _load_lib(lib)
-    return importlib.import_module(ext)
-
-
-def _find_ffmpeg_extension(ffmpeg_vers):
-    for ffmpeg_ver in ffmpeg_vers:
-        try:
-            return _find_versionsed_ffmpeg_extension(ffmpeg_ver)
-        except Exception:
-            _LG.debug("Failed to load FFmpeg%s extension.", ffmpeg_ver, exc_info=True)
-            continue
-    raise ImportError(
-        f"Failed to intialize FFmpeg extension. Tried versions: {ffmpeg_vers}. "
-        "Enable DEBUG logging to see more details about the error."
-    )
-
-
-def _get_ffmpeg_versions():
-    ffmpeg_vers = _FFMPEG_VERS
-    # User override
-    if (ffmpeg_ver := os.environ.get("TORCHAUDIO_USE_FFMPEG_VERSION")) is not None:
-        if ffmpeg_ver not in ffmpeg_vers:
-            raise ValueError(
-                f"The FFmpeg version '{ffmpeg_ver}' (read from TORCHAUDIO_USE_FFMPEG_VERSION) "
-                f"is not one of supported values. Possible values are {ffmpeg_vers}"
-            )
-        ffmpeg_vers = [ffmpeg_ver]
-    return ffmpeg_vers
-
-
-def _init_ffmpeg():
-    ffmpeg_vers = _get_ffmpeg_versions()
-    ext = _find_ffmpeg_extension(ffmpeg_vers)
-    ext.init()
-    if ext.get_log_level() > 8:
-        ext.set_log_level(8)
-    return ext
-
-
 class _LazyImporter(types.ModuleType):
     """Lazily import module/extension."""
 

src/torchaudio/io/__init__.py

@@ -1,7 +1,7 @@
+from torio.io import CodecConfig, StreamingMediaDecoder as StreamReader, StreamingMediaEncoder as StreamWriter
+
 from ._effector import AudioEffector
 from ._playback import play_audio
-from ._stream_reader import StreamReader
-from ._stream_writer import CodecConfig, StreamWriter
 
 
 __all__ = [

src/torchaudio/io/_effector.py

@@ -4,8 +4,8 @@
 import torch
 from torch import Tensor
 
-from ._stream_reader import _get_afilter_desc, StreamReader
-from ._stream_writer import CodecConfig, StreamWriter
+from torio.io._streaming_media_decoder import _get_afilter_desc, StreamingMediaDecoder as StreamReader
+from torio.io._streaming_media_encoder import CodecConfig, StreamingMediaEncoder as StreamWriter
 
 
 class _StreamingIOBuffer:

src/torchaudio/utils/__init__.py

@@ -1,4 +1,6 @@
-from . import ffmpeg_utils, sox_utils
+from torio.utils import ffmpeg_utils
+
+from . import sox_utils
 from .download import download_asset