From 9e4fa360b29c95d765bab06adb00865141e34cac Mon Sep 17 00:00:00 2001
From: moto <855818+mthrok@users.noreply.github.com>
Date: Wed, 29 Dec 2021 22:11:55 -0500
Subject: [PATCH] add seek

---
 torchaudio/csrc/ffmpeg/prototype.cpp | 5 +++++
 torchaudio/csrc/ffmpeg/streamer.cpp  | 8 ++++++++
 torchaudio/csrc/ffmpeg/streamer.h    | 2 ++
 3 files changed, 15 insertions(+)

diff --git a/torchaudio/csrc/ffmpeg/prototype.cpp b/torchaudio/csrc/ffmpeg/prototype.cpp
index e1c94b32ab0..7ff9415c710 100644
--- a/torchaudio/csrc/ffmpeg/prototype.cpp
+++ b/torchaudio/csrc/ffmpeg/prototype.cpp
@@ -98,6 +98,10 @@ int64_t find_best_video_stream(S s) {
   return s->s.find_best_video_stream();
 }
 
+void seek(S s, int64_t timestamp) {
+  s->s.seek(timestamp);
+}
+
 template <typename... Args>
 std::string string_format(const std::string& format, Args... args) {
   char buffer[512];
@@ -309,6 +313,7 @@ TORCH_LIBRARY_FRAGMENT(torchaudio, m) {
   m.def(
       "torchaudio::ffmpeg_streamer_find_best_video_stream",
       find_best_video_stream);
+  m.def("torchaudio::ffmpeg_streamer_seek", seek);
   m.def(
       "torchaudio::ffmpeg_streamer_add_basic_audio_stream",
       add_basic_audio_stream);
diff --git a/torchaudio/csrc/ffmpeg/streamer.cpp b/torchaudio/csrc/ffmpeg/streamer.cpp
index cd59eca3cb1..836ec2b3f94 100644
--- a/torchaudio/csrc/ffmpeg/streamer.cpp
+++ b/torchaudio/csrc/ffmpeg/streamer.cpp
@@ -133,6 +133,14 @@ bool Streamer::is_buffer_ready() const {
 ////////////////////////////////////////////////////////////////////////////////
 // Configure methods
 ////////////////////////////////////////////////////////////////////////////////
+void Streamer::seek(double timestamp) {
+  int64_t ts = static_cast<int64_t>(timestamp * AV_TIME_BASE);
+  int ret = avformat_seek_file(pFormatContext, -1, INT64_MIN, ts, INT64_MAX, 0);
+  if (ret < 0) {
+    throw std::runtime_error(std::string("Failed to seek: ") + av_err2str(ret));
+  }
+}
+
 void Streamer::add_audio_stream(
     int i,
     int frames_per_chunk,
diff --git a/torchaudio/csrc/ffmpeg/streamer.h b/torchaudio/csrc/ffmpeg/streamer.h
index 612048ef859..bb8eb5d3981 100644
--- a/torchaudio/csrc/ffmpeg/streamer.h
+++ b/torchaudio/csrc/ffmpeg/streamer.h
@@ -60,6 +60,8 @@ class Streamer {
   //////////////////////////////////////////////////////////////////////////////
   // Configure methods
   //////////////////////////////////////////////////////////////////////////////
+  void seek(double timestamp);
+
   void add_audio_stream(
       int i,
       int frames_per_chunk,