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

Summary: Pull Request resolved: #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, int]] = 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, int]] = 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, int]] = 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, int]] = None,
     ) -> None:
+        if not isinstance(compression, (int, None)):
+            raise ValueError(
+                "SoX backend expects non-`None` value for argument `compression` to be of ",
+                f"type `int`, 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,