Make sox selective (#1338)

torchaudio/csrc/CMakeLists.txt

@@ -5,13 +5,9 @@ get_property(TORCHAUDIO_THIRD_PARTIES GLOBAL PROPERTY TORCHAUDIO_THIRD_PARTIES)
 ################################################################################
 set(
   LIBTORCHAUDIO_SOURCES
-  sox/io.cpp
-  sox/utils.cpp
-  sox/effects.cpp
-  sox/effects_chain.cpp
-  sox/types.cpp
   lfilter.cpp
   overdrive.cpp
+  utils.cpp
   )
 
 if(BUILD_TRANSDUCER)
@@ -22,6 +18,18 @@ if(BUILD_KALDI)
   list(APPEND LIBTORCHAUDIO_SOURCES kaldi.cpp)
 endif()
 
+if(BUILD_SOX)
+  set(
+    SOX_SOURCES
+    sox/io.cpp
+    sox/utils.cpp
+    sox/effects.cpp
+    sox/effects_chain.cpp
+    sox/types.cpp
+  )
+  list(APPEND LIBTORCHAUDIO_SOURCES ${SOX_SOURCES})
+endif()
+
 ################################################################################
 # libtorchaudio.so
 ################################################################################
@@ -76,6 +84,10 @@ if (BUILD_TORCHAUDIO_PYTHON_EXTENSION)
     _torchaudio PRIVATE TORCH_API_INCLUDE_EXTENSION_H
   )
 
+  if (BUILD_SOX)
+    target_compile_definitions(_torchaudio PRIVATE INCLUDE_SOX)
+  endif()
+
   target_include_directories(
     _torchaudio
     PRIVATE

torchaudio/csrc/pybind.cpp

 #include <torch/extension.h>
+
+#ifdef INCLUDE_SOX
 #include <torchaudio/csrc/sox/effects.h>
 #include <torchaudio/csrc/sox/io.h>
+#endif
 
 PYBIND11_MODULE(_torchaudio, m) {
+#ifdef INCLUDE_SOX
   m.def(
       "get_info_fileobj",
       &torchaudio::sox_io::get_info_fileobj,
@@ -19,4 +23,5 @@ PYBIND11_MODULE(_torchaudio, m) {
       "apply_effects_fileobj",
       &torchaudio::sox_effects::apply_effects_fileobj,
       "Decode audio data from file-like obj and apply effects.");
+#endif
 }

torchaudio/csrc/utils.cpp

+#include <torch/script.h>
+
+namespace torchaudio {
+
+namespace {
+
+bool is_sox_available() {
+#ifdef INCLUDE_SOX
+  return true;
+#else
+  return false;
+#endif
+}
+
+} // namespace
+
+TORCH_LIBRARY_FRAGMENT(torchaudio, m) {
+  m.def("torchaudio::is_sox_available", &is_sox_available);
+}
+
+} // namespace torchaudio

torchaudio/functional/functional.py

@@ -7,9 +7,7 @@ from typing import Optional, Tuple
 
 import torch
 from torch import Tensor
-from torchaudio._internal import (
-    module_utils as _mod_utils,
-)
+from torchaudio._internal import module_utils as _mod_utils
 import torchaudio
 
 __all__ = [
@@ -1000,7 +998,7 @@ def spectral_centroid(
     return (freqs * specgram).sum(dim=freq_dim) / specgram.sum(dim=freq_dim)
 
 
-@_mod_utils.requires_module('torchaudio._torchaudio')
+@_mod_utils.requires_sox()
 def apply_codec(
     waveform: Tensor,
     sample_rate: int,

torchaudio/sox_effects/__init__.py

@@ -8,7 +8,7 @@ from .sox_effects import (
 )
 
 
-if _mod_utils.is_module_available('torchaudio._torchaudio'):
+if _mod_utils.is_sox_available():
     import atexit
     init_sox_effects()
     atexit.register(shutdown_sox_effects)

torchaudio/sox_effects/sox_effects.py

@@ -8,7 +8,7 @@ from torchaudio._internal import module_utils as _mod_utils
 from torchaudio.utils.sox_utils import list_effects
 
 
-@_mod_utils.requires_module('torchaudio._torchaudio')
+@_mod_utils.requires_sox()
 def init_sox_effects():
     """Initialize resources required to use sox effects.
 
@@ -23,7 +23,7 @@ def init_sox_effects():
     torch.ops.torchaudio.sox_effects_initialize_sox_effects()
 
 
-@_mod_utils.requires_module("torchaudio._torchaudio")
+@_mod_utils.requires_sox()
 def shutdown_sox_effects():
     """Clean up resources required to use sox effects.
 
@@ -37,7 +37,7 @@ def shutdown_sox_effects():
     torch.ops.torchaudio.sox_effects_shutdown_sox_effects()
 
 
-@_mod_utils.requires_module('torchaudio._torchaudio')
+@_mod_utils.requires_sox()
 def effect_names() -> List[str]:
     """Gets list of valid sox effect names
 
@@ -51,7 +51,7 @@ def effect_names() -> List[str]:
     return list(list_effects().keys())
 
 
-@_mod_utils.requires_module('torchaudio._torchaudio')
+@_mod_utils.requires_sox()
 def apply_effects_tensor(
         tensor: torch.Tensor,
         sample_rate: int,
@@ -152,7 +152,7 @@ def apply_effects_tensor(
         tensor, sample_rate, effects, channels_first)
 
 
-@_mod_utils.requires_module('torchaudio._torchaudio')
+@_mod_utils.requires_sox()
 def apply_effects_file(
         path: str,
         effects: List[List[str]],

torchaudio/utils/__init__.py

 from . import (
     sox_utils,
 )
-
 from torchaudio._internal import module_utils as _mod_utils
 
 
-if _mod_utils.is_module_available('torchaudio._torchaudio'):
+if _mod_utils.is_sox_available():
     sox_utils.set_verbosity(1)

torchaudio/utils/sox_utils.py

 from typing import List, Dict
 
 import torch
+from torchaudio._internal import module_utils as _mod_utils
 
-from torchaudio._internal import (
-    module_utils as _mod_utils,
-)
 
-
-@_mod_utils.requires_module('torchaudio._torchaudio')
+@_mod_utils.requires_sox()
 def set_seed(seed: int):
     """Set libsox's PRNG
 
@@ -20,7 +17,7 @@ def set_seed(seed: int):
     torch.ops.torchaudio.sox_utils_set_seed(seed)
 
 
-@_mod_utils.requires_module('torchaudio._torchaudio')
+@_mod_utils.requires_sox()
 def set_verbosity(verbosity: int):
     """Set libsox's verbosity
 
@@ -38,7 +35,7 @@ def set_verbosity(verbosity: int):
     torch.ops.torchaudio.sox_utils_set_verbosity(verbosity)
 
 
-@_mod_utils.requires_module('torchaudio._torchaudio')
+@_mod_utils.requires_sox()
 def set_buffer_size(buffer_size: int):
     """Set buffer size for sox effect chain
 
@@ -51,7 +48,7 @@ def set_buffer_size(buffer_size: int):
     torch.ops.torchaudio.sox_utils_set_buffer_size(buffer_size)
 
 
-@_mod_utils.requires_module('torchaudio._torchaudio')
+@_mod_utils.requires_sox()
 def set_use_threads(use_threads: bool):
     """Set multithread option for sox effect chain
 
@@ -65,7 +62,7 @@ def set_use_threads(use_threads: bool):
     torch.ops.torchaudio.sox_utils_set_use_threads(use_threads)
 
 
-@_mod_utils.requires_module('torchaudio._torchaudio')
+@_mod_utils.requires_sox()
 def list_effects() -> Dict[str, str]:
     """List the available sox effect names
 
@@ -75,7 +72,7 @@ def list_effects() -> Dict[str, str]:
     return dict(torch.ops.torchaudio.sox_utils_list_effects())
 
 
-@_mod_utils.requires_module('torchaudio._torchaudio')
+@_mod_utils.requires_sox()
 def list_read_formats() -> List[str]:
     """List the supported audio formats for read
 
@@ -85,7 +82,7 @@ def list_read_formats() -> List[str]:
     return torch.ops.torchaudio.sox_utils_list_read_formats()
 
 
-@_mod_utils.requires_module('torchaudio._torchaudio')
+@_mod_utils.requires_sox()
 def list_write_formats() -> List[str]:
     """List the supported audio formats for write