From 0c4f76c2e88ba13018b297f9fc4d8669c567123a Mon Sep 17 00:00:00 2001
From: Jeff Hwang <jeffhwang@fb.com>
Date: Tue, 17 Oct 2023 09:42:35 -0700
Subject: [PATCH] Support compression level in i/o dispatcher backend (#3662)

Summary: Pull Request resolved: https://github.com/pytorch/audio/pull/3662

Differential Revision: D50367721
---
 src/torchaudio/_backend/backend.py   | 2 ++
 src/torchaudio/_backend/ffmpeg.py    | 9 +++++++++
 src/torchaudio/_backend/soundfile.py | 5 +++++
 src/torchaudio/_backend/sox.py       | 8 +++++++-
 4 files changed, 23 insertions(+), 1 deletion(-)

diff --git a/src/torchaudio/_backend/backend.py b/src/torchaudio/_backend/backend.py
index da8d1bb7bc4..5749343da9d 100644
--- a/src/torchaudio/_backend/backend.py
+++ b/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
 
diff --git a/src/torchaudio/_backend/ffmpeg.py b/src/torchaudio/_backend/ffmpeg.py
index 0bcf6ee4bd4..97c1e070541 100644
--- a/src/torchaudio/_backend/ffmpeg.py
+++ b/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
diff --git a/src/torchaudio/_backend/soundfile.py b/src/torchaudio/_backend/soundfile.py
index 701e471c98d..63d66f686d6 100644
--- a/src/torchaudio/_backend/soundfile.py
+++ b/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
         )
diff --git a/src/torchaudio/_backend/sox.py b/src/torchaudio/_backend/sox.py
index 592890c95c9..d1a4ac172a2 100644
--- a/src/torchaudio/_backend/sox.py
+++ b/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,