Use Encoder as backend of AudioEncoder and VideoEncoder (meta-pytorch#1478)

Author: NicolasHug

src/torchcodec/_core/Encoder.cpp

@@ -14,69 +14,18 @@ extern "C" {
 }
 
 namespace facebook::torchcodec {
-class FORCE_PUBLIC_VISIBILITY AudioEncoder {
- public:
-  ~AudioEncoder();
-
-  AudioEncoder(
-      const torch::stable::Tensor& samples,
-      int sampleRate,
-      std::string_view fileName,
-      const AudioStreamOptions& audioStreamOptions);
-
-  AudioEncoder(
-      const torch::stable::Tensor& samples,
-      int sampleRate,
-      std::string_view formatName,
-      std::unique_ptr<AVIOContextHolder> avioContextHolder,
-      const AudioStreamOptions& audioStreamOptions);
-
-  void encode();
-
-  torch::stable::Tensor encodeToTensor();
-
- private:
-  void initializeEncoder(const AudioStreamOptions& audioStreamOptions);
-  UniqueAVFrame maybeConvertAVFrame(const UniqueAVFrame& avFrame);
-  void encodeFrameThroughFifo(
-      AutoAVPacket& autoAVPacket,
-      const UniqueAVFrame& avFrame,
-      bool flushFifo = false);
-  void encodeFrame(AutoAVPacket& autoAVPacket, const UniqueAVFrame& avFrame);
-  void maybeFlushSwrBuffers(AutoAVPacket& autoAVPacket);
-  void flushBuffers();
-
-  UniqueEncodingAVFormatContext avFormatContext_;
-  UniqueAVCodecContext avCodecContext_;
-  int streamIndex_;
-  UniqueSwrContext swrContext_;
-  AudioStreamOptions audioStreamOptions;
-
-  const torch::stable::Tensor samples_;
-
-  int outNumChannels_ = -1;
-  int outSampleRate_ = -1;
-  int inSampleRate_ = -1;
-
-  UniqueAVAudioFifo avAudioFifo_;
-
-  std::unique_ptr<AVIOContextHolder> avioContextHolder_;
-
-  bool encodeWasCalled_ = false;
-  int64_t lastEncodedAVFramePts_ = 0;
-};
 
 /* clang-format off */
 //
 // Note: [Encoding loop, sample rate conversion and FIFO]
 //
 // The input samples are in a given format, sample rate, and number of channels.
 // We may want to change these properties before encoding. The conversion is
-// done in maybeConvertAVFrame() and we rely on libswresample. When sample rate
-// conversion is needed, this means two things:
+// done in maybeConvertAudioAVFrame() and we rely on libswresample. When sample
+// rate conversion is needed, this means two things:
 // - swr will be storing samples in its internal buffers, which we'll need to
 //   flush at the very end of the encoding process.
-// - the converted AVFrame we get back from maybeConvertAVFrame() typically
+// - the converted AVFrame we get back from maybeConvertAudioAVFrame() typically
 //   won't have the same number of samples as the original AVFrame. And that's
 //   a problem, because some encoders expect AVFrames with a specific and
 //   constant number of samples. If we were to send it as-is, we'd get an error
@@ -85,100 +34,14 @@ class FORCE_PUBLIC_VISIBILITY AudioEncoder {
 //   from which we can pull the exact number of samples that we need. Note that
 //   this involves at least 2 additional copies.
 //
-// To be clear, the FIFO is only used if BOTH the following conditions are met:
-//  - sample rate conversion is needed (inSampleRate_ != outSampleRate_)
-//  - the encoder expects a specific number of samples per AVFrame (fixed frame size)
-//    This is not the case for all encoders, e.g. WAV doesn't care about frame size.
-//
-// Also, the FIFO is either:
-// - used for every single frame during the encoding process, or
-// - not used at all.
-// There is no scenario where a given Encoder() instance would sometimes use a
-// FIFO, sometimes not.
-//
-// Drawing made with https://asciiflow.com/, can be copy/pasted there if it
-// needs editing:
-//
-// ┌─One─iteration─of─main─encoding─loop─(encode())───────────────────────────────────────────┐
-// │                                                                                          │
-// │                        Converts:                                                         │
-// │                         - num channels                                                   │
-// │                         - format                                                         │
-// │                         - sample rate                                                    │
-// │                        If sample rate is converted, stores data in swr buffers,          │
-// │                        which will need to be flushed by maybeFlushSwrBuffers()           │
-// │                                                                                          │
-// │                               ▲                                                          │
-// │                               │                 ┌─EncodeFrameThroughFifo()──────────────┐│
-// │                               │                 │                                       ││
-// │    AVFrame  ──────►  MaybeConvertAVFrame()───▲──│─┬──► NO FIFO ─►┬──▲────►encodeFrame() ││
-// │    with                                      │  │ │              │  │                   ││
-// │    input                                     │  │ │              │  │                   ││
-// │    samples                                   │  │ │              │  │                   ││
-// │                                              │  │ │              │  │                   ││
-// │                                              │  │ └────► FIFO ─►─┘  │                   ││
-// │                                              │  └───────────────────┼───────────────────┘│
-// └──────────────────────────────────────────────┼──────────────────────┼────────────────────┘
-//                                                │                      │
-//  AVFrame from  maybeFlushSwrBuffers()       ───┘                      │
-//  Only if sample rate conversion was needed                       nullptr, to flush
-//  The call to maybeFlushSwrBuffers() will                         FFmpeg buffers
-//  also instruct to flush the FIFO, if it exists.
-//
-//
+// Unlike the old single-stream AudioEncoder where the FIFO was only used when
+// sample rate conversion was needed with a fixed-frame-size encoder, the
+// MultiStreamEncoder always creates a FIFO. This is because addSamples() can
+// be called multiple times with arbitrary chunk sizes, so we always need to
+// buffer and re-chunk samples into frame_size-sized batches.
 //
 /* clang-format on */
 
-class FORCE_PUBLIC_VISIBILITY VideoEncoder {
- public:
-  ~VideoEncoder();
-
-  // Rule of Five requires that we define copy and move
-  // constructors and assignment operators.
-  // Both are deleted because we have unique_ptr members
-  VideoEncoder(const VideoEncoder&) = delete;
-  VideoEncoder& operator=(const VideoEncoder&) = delete;
-
-  // Move operators deleted since UniqueAVDictionary member is not movable
-  VideoEncoder(VideoEncoder&&) = delete;
-  VideoEncoder& operator=(VideoEncoder&&) = delete;
-
-  VideoEncoder(
-      const torch::stable::Tensor& frames,
-      double frameRate,
-      std::string_view fileName,
-      const VideoStreamOptions& videoStreamOptions);
-
-  VideoEncoder(
-      const torch::stable::Tensor& frames,
-      double frameRate,
-      std::string_view formatName,
-      std::unique_ptr<AVIOContextHolder> avioContextHolder,
-      const VideoStreamOptions& videoStreamOptions);
-
-  void encode();
-
-  torch::stable::Tensor encodeToTensor();
-
- private:
-  void initializeEncoder(const VideoStreamOptions& videoStreamOptions);
-  void encodeFrame(AutoAVPacket& autoAVPacket, const UniqueAVFrame& avFrame);
-  void flushBuffers();
-
-  UniqueEncodingAVFormatContext avFormatContext_;
-  UniqueAVCodecContext avCodecContext_;
-  AVStream* avStream_ = nullptr;
-
-  const torch::stable::Tensor frames_;
-  double inFrameRate_;
-
-  std::unique_ptr<AVIOContextHolder> avioContextHolder_;
-  std::unique_ptr<DeviceInterface> deviceInterface_;
-
-  bool encodeWasCalled_ = false;
-  UniqueAVDictionary avFormatOptions_;
-};
-
 class FORCE_PUBLIC_VISIBILITY MultiStreamEncoder {
  public:
   ~MultiStreamEncoder();

src/torchcodec/_core/Encoder.h

@@ -20,12 +20,6 @@
     core_library_path,
     create_streaming_encoder,
     create_wav_decoder_from_file,
-    encode_audio_to_file,
-    encode_audio_to_file_like,
-    encode_audio_to_tensor,
-    encode_video_to_file,
-    encode_video_to_file_like,
-    encode_video_to_tensor,
     ffmpeg_major_version,
     get_ffmpeg_library_versions,
     get_frame_at_index,