Add stream_index to TorchAudioDecoder

Author: Dan-Flores

benchmarks/decoders/benchmark_decoders_library.py

@@ -372,7 +372,7 @@ def __init__(
         num_ffmpeg_threads=None,
         device="cpu",
         seek_mode="exact",
-        stream_index: int | None = None,
+        stream_index=None,
     ):
         self._num_ffmpeg_threads = num_ffmpeg_threads
         self._device = device
@@ -536,19 +536,22 @@ def decode_first_n_frames(self, video_file, n):
 
 
 class TorchAudioDecoder(AbstractDecoder):
-    def __init__(self):
+    def __init__(self, stream_index=None):
         import torchaudio  # noqa: F401
 
         self.torchaudio = torchaudio
 
         from torchvision.transforms import v2 as transforms_v2
 
         self.transforms_v2 = transforms_v2
+        self._stream_index = int(stream_index) if stream_index else None
 
     def decode_frames(self, video_file, pts_list):
         stream_reader = self.torchaudio.io.StreamReader(src=video_file)
         stream_reader.add_basic_video_stream(
-            frames_per_chunk=1, decoder_option={"threads": "0"}
+            frames_per_chunk=1,
+            decoder_option={"threads": "0"},
+            stream_index=self._stream_index,
         )
         frames = []
         for pts in pts_list:
@@ -561,7 +564,9 @@ def decode_frames(self, video_file, pts_list):
     def decode_first_n_frames(self, video_file, n):
         stream_reader = self.torchaudio.io.StreamReader(src=video_file)
         stream_reader.add_basic_video_stream(
-            frames_per_chunk=1, decoder_option={"threads": "0"}
+            frames_per_chunk=1,
+            decoder_option={"threads": "0"},
+            stream_index=self._stream_index,
         )
         frames = []
         frame_cnt = 0
@@ -576,7 +581,9 @@ def decode_first_n_frames(self, video_file, n):
     def decode_and_resize(self, video_file, pts_list, height, width, device):
         stream_reader = self.torchaudio.io.StreamReader(src=video_file)
         stream_reader.add_basic_video_stream(
-            frames_per_chunk=1, decoder_option={"threads": "1"}
+            frames_per_chunk=1,
+            decoder_option={"threads": "1"},
+            stream_index=self._stream_index,
         )
         frames = []
         for pts in pts_list: