Add qscale to CodecConfig option (#3224)

Summary:
This commit adds the equivalent of `qscale` option in FFmpeg to StreamWriter.CodecConfig.
`qscale` enables variable bit rate.

The following figure illustrates the difference between currently available configs.
From top to bottom; original, `compression_level=1`, `compression_level=9`, `bit_rate=192k`, `bit_rate=8k`, `qscale=9`, `qscale=1`.
![Figure_1](https://user-images.githubusercontent.com/855818/228990681-368bf84f-00a7-4248-80ac-6ee728da8f1a.png)

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

Reviewed By: hwangjeff

Differential Revision: D44563633

Pulled By: mthrok

fbshipit-source-id: ff74cd803b5abf1222f087e3e46ba7d81a35f672

torchaudio/csrc/ffmpeg/pybind/pybind.cpp

@@ -36,7 +36,7 @@ PYBIND11_MODULE(_torchaudio_ffmpeg, m) {
       .def_readwrite("frames", &Chunk::frames)
       .def_readwrite("pts", &Chunk::pts);
   py::class_<CodecConfig>(m, "CodecConfig", py::module_local())
-      .def(py::init<int, int, int, int>());
+      .def(py::init<int, int, const c10::optional<int>&, int, int>());
   py::class_<StreamWriter>(m, "StreamWriter", py::module_local())
       .def(py::init<const std::string&, const c10::optional<std::string>&>())
       .def("set_metadata", &StreamWriter::set_metadata)

torchaudio/csrc/ffmpeg/stream_writer/encode_process.cpp

@@ -370,6 +370,10 @@ void configure_audio_codec_ctx(
     if (cfg.compression_level != -1) {
       codec_ctx->compression_level = cfg.compression_level;
     }
+    if (cfg.qscale) {
+      codec_ctx->flags |= AV_CODEC_FLAG_QSCALE;
+      codec_ctx->global_quality = FF_QP2LAMBDA * cfg.qscale.value();
+    }
   }
 }
 
@@ -500,6 +504,10 @@ void configure_video_codec_ctx(
     if (cfg.max_b_frames != -1) {
       ctx->max_b_frames = cfg.max_b_frames;
     }
+    if (cfg.qscale) {
+      ctx->flags |= AV_CODEC_FLAG_QSCALE;
+      ctx->global_quality = FF_QP2LAMBDA * cfg.qscale.value();
+    }
   }
 }
 

torchaudio/csrc/ffmpeg/stream_writer/types.h

@@ -5,6 +5,13 @@ struct CodecConfig {
   int bit_rate = -1;
   int compression_level = -1;
 
+  // qscale corresponds to ffmpeg CLI's qscale.
+  // Example: MP3
+  // https://trac.ffmpeg.org/wiki/Encode/MP3
+  // This should be set like
+  // https://github.com/FFmpeg/FFmpeg/blob/n4.3.2/fftools/ffmpeg_opt.c#L1550
+  const c10::optional<int> qscale = -1;
+
   // video
   int gop_size = -1;
   int max_b_frames = -1;

torchaudio/io/_stream_writer.py

@@ -141,6 +141,13 @@ class StreamWriter:
         compression_level: int = -1
         """Compression level"""
 
+        qscale: Optional[int] = None
+        """Global quality factor. Enables variable bit rate. Valid values depend on encoder.
+
+        For example: MP3 takes ``0`` - ``9`` (https://trac.ffmpeg.org/wiki/Encode/MP3) while
+        libvorbis takes ``-1`` - ``10``.
+        """
+
         gop_size: int = -1
         """The number of pictures in a group of pictures, or 0 for intra_only"""
 
@@ -148,7 +155,7 @@ class StreamWriter:
         """maximum number of B-frames between non-B-frames."""
 
         def __post_init__(self):
-            super().__init__(self.bit_rate, self.compression_level, self.gop_size, self.max_b_frames)
+            super().__init__(self.bit_rate, self.compression_level, self.qscale, self.gop_size, self.max_b_frames)
 
     def __init__(
         self,