Support compression level in i/o dispatcher backend (pytorch#3662)

Summary: Pull Request resolved: pytorch#3662

Differential Revision: D50367721

src/torchaudio/_backend/backend.py

@@ -3,6 +3,7 @@
 from typing import BinaryIO, Optional, Tuple, Union
 
 from torch import Tensor
+from torchaudio.io import CodecConfig
 
 from .common import AudioMetaData
 
@@ -37,6 +38,7 @@ def save(
         encoding: Optional[str] = None,
         bits_per_sample: Optional[int] = None,
         buffer_size: int = 4096,
+        compression: Optional[Union[CodecConfig, float]] = None,
     ) -> None:
         raise NotImplementedError
 

src/torchaudio/_backend/ffmpeg.py

@@ -228,6 +228,7 @@ def save_audio(
     encoding: Optional[str] = None,
     bits_per_sample: Optional[int] = None,
     buffer_size: int = 4096,
+    compression: Optional[torchaudio.io.CodecConfig] = None,
 ) -> None:
     ext = None
     if hasattr(uri, "write"):
@@ -250,6 +251,7 @@ def save_audio(
         format=_get_sample_format(src.dtype),
         encoder=encoder,
         encoder_format=enc_fmt,
+        codec_config=compression,
     )
     with s.open():
         s.write_audio_chunk(0, src)
@@ -304,7 +306,13 @@ def save(
         encoding: Optional[str] = None,
         bits_per_sample: Optional[int] = None,
         buffer_size: int = 4096,
+        compression: Optional[Union[torchaudio.io.CodecConfig, float]] = None,
     ) -> None:
+        if not isinstance(compression, (torchaudio.io.CodecConfig, None)):
+            raise ValueError(
+                "FFmpeg backend expects non-`None` value for argument `compression` to be of ",
+                f"type `torchaudio.io.CodecConfig`, but received value of type {type(compression)}",
+            )
         save_audio(
             uri,
             src,
@@ -314,6 +322,7 @@ def save(
             encoding,
             bits_per_sample,
             buffer_size,
+            compression,
         )
 
     @staticmethod

src/torchaudio/_backend/soundfile.py

@@ -2,6 +2,7 @@
 from typing import BinaryIO, Optional, Tuple, Union
 
 import torch
+from torchaudio.io import CodecConfig
 
 from . import soundfile_backend
 from .backend import Backend
@@ -35,7 +36,11 @@ def save(
         encoding: Optional[str] = None,
         bits_per_sample: Optional[int] = None,
         buffer_size: int = 4096,
+        compression: Optional[Union[CodecConfig, float]] = None,
     ) -> None:
+        if compression:
+            raise ValueError("soundfile backend does not support argument `compression`.")
+
         soundfile_backend.save(
             uri, src, sample_rate, channels_first, format=format, encoding=encoding, bits_per_sample=bits_per_sample
         )

src/torchaudio/_backend/sox.py

@@ -56,7 +56,13 @@ def save(
         encoding: Optional[str] = None,
         bits_per_sample: Optional[int] = None,
         buffer_size: int = 4096,
+        compression: Optional[Union[torchaudio.io.CodecConfig, float]] = None,
     ) -> None:
+        if not isinstance(compression, (float, None)):
+            raise ValueError(
+                "SoX backend expects non-`None` value for argument `compression` to be of ",
+                f"type `float`, but received value of type {type(compression)}",
+            )
         if hasattr(uri, "write"):
             raise ValueError(
                 "SoX backend does not support writing to file-like objects. ",
@@ -68,7 +74,7 @@ def save(
                 src,
                 sample_rate,
                 channels_first,
-                None,
+                compression,
                 format,
                 encoding,
                 bits_per_sample,

src/torchaudio/_backend/utils.py

@@ -5,6 +5,7 @@
 import torch
 
 from torchaudio._extension import lazy_import_ffmpeg_ext, lazy_import_sox_ext
+from torchaudio.io import CodecConfig
 
 from . import soundfile_backend
 
@@ -229,6 +230,7 @@ def save(
         bits_per_sample: Optional[int] = None,
         buffer_size: int = 4096,
         backend: Optional[str] = None,
+        compression: Optional[Union[CodecConfig, float]] = None,
     ):
         """Save audio data to file.
 
@@ -283,8 +285,13 @@ def save(
 
                 .. seealso::
                    :ref:`backend`
+
+            compression (CodecConfig, float, or None, optional):
+                To fill in.
         """
         backend = dispatcher(uri, format, backend)
-        return backend.save(uri, src, sample_rate, channels_first, format, encoding, bits_per_sample, buffer_size)
+        return backend.save(
+            uri, src, sample_rate, channels_first, format, encoding, bits_per_sample, buffer_size, compression
+        )
 
     return save