Use torchcodec for loading

Author: samanklesaria

requirements.txt

@@ -1,5 +1,6 @@
 # Minimum runtime dependencies
 torch
+torchcodec
 
 # Optional runtime dependencies
 kaldi_io

src/torchaudio/datasets/cmuarctic.py

@@ -4,6 +4,7 @@
 from typing import Tuple, Union
 
 import torchaudio
+from torchaudio.utils import load_torchcodec
 from torch import Tensor
 from torch.utils.data import Dataset
 from torchaudio._internal import download_url_to_file
@@ -43,7 +44,7 @@ def load_cmuarctic_item(line: str, path: str, folder_audio: str, ext_audio: str)
     file_audio = os.path.join(path, folder_audio, utterance_id + ext_audio)
 
     # Load audio
-    waveform, sample_rate = torchaudio.load(file_audio)
+    waveform, sample_rate = load_torchcodec(file_audio)
 
     return (waveform, sample_rate, transcript, utterance_id.split("_")[1])
 

src/torchaudio/datasets/commonvoice.py

@@ -6,6 +6,7 @@
 import torchaudio
 from torch import Tensor
 from torch.utils.data import Dataset
+from torchaudio.utils import load_torchcodec
 
 
 def load_commonvoice_item(
@@ -20,7 +21,7 @@ def load_commonvoice_item(
     filename = os.path.join(path, folder_audio, fileid)
     if not filename.endswith(ext_audio):
         filename += ext_audio
-    waveform, sample_rate = torchaudio.load(filename)
+    waveform, sample_rate = load_torchcodec(filename)
 
     dic = dict(zip(header, line))
 

src/torchaudio/datasets/dr_vctk.py

@@ -6,6 +6,7 @@
 from torch.utils.data import Dataset
 from torchaudio._internal import download_url_to_file
 from torchaudio.datasets.utils import _extract_zip
+from torchaudio.utils import load_torchcodec
 
 
 _URL = "https://datashare.ed.ac.uk/bitstream/handle/10283/3038/DR-VCTK.zip"
@@ -75,8 +76,8 @@ def _load_dr_vctk_item(self, filename: str) -> Tuple[Tensor, int, Tensor, int, s
         source, channel_id = self._config[filename]
         file_clean_audio = self._clean_audio_dir / filename
         file_noisy_audio = self._noisy_audio_dir / filename
-        waveform_clean, sample_rate_clean = torchaudio.load(file_clean_audio)
-        waveform_noisy, sample_rate_noisy = torchaudio.load(file_noisy_audio)
+        waveform_clean, sample_rate_clean = load_torchcodec(file_clean_audio)
+        waveform_noisy, sample_rate_noisy = load_torchcodec(file_noisy_audio)
         return (
             waveform_clean,
             sample_rate_clean,

src/torchaudio/datasets/gtzan.py

@@ -7,6 +7,7 @@
 from torch.utils.data import Dataset
 from torchaudio._internal import download_url_to_file
 from torchaudio.datasets.utils import _extract_tar
+from torchaudio.utils import load_torchcodec
 
 # The following lists prefixed with `filtered_` provide a filtered split
 # that:
@@ -990,7 +991,7 @@ def load_gtzan_item(fileid: str, path: str, ext_audio: str) -> Tuple[Tensor, str
 
     # Read wav
     file_audio = os.path.join(path, label, fileid + ext_audio)
-    waveform, sample_rate = torchaudio.load(file_audio)
+    waveform, sample_rate = load_torchcodec(file_audio)
 
     return waveform, sample_rate, label
 

src/torchaudio/datasets/librilight_limited.py

@@ -8,6 +8,7 @@
 from torchaudio._internal import download_url_to_file
 from torchaudio.datasets.librispeech import _get_librispeech_metadata
 from torchaudio.datasets.utils import _extract_tar
+from torchaudio.utils import load_torchcodec
 
 
 _ARCHIVE_NAME = "librispeech_finetuning"
@@ -104,7 +105,7 @@ def __getitem__(self, n: int) -> Tuple[Tensor, int, str, int, int, int]:
         """
         file_path, fileid = self._fileids_paths[n]
         metadata = _get_librispeech_metadata(fileid, self._path, file_path, self._ext_audio, self._ext_txt)
-        waveform, _ = torchaudio.load(os.path.join(self._path, metadata[0]))
+        waveform, _ = load_torchcodec(os.path.join(self._path, metadata[0]))
         return (waveform,) + metadata[1:]
 
     def __len__(self) -> int:

src/torchaudio/datasets/libritts.py

@@ -7,6 +7,7 @@
 from torch.utils.data import Dataset
 from torchaudio._internal import download_url_to_file
 from torchaudio.datasets.utils import _extract_tar
+from torchaudio.utils import load_torchcodec
 
 URL = "train-clean-100"
 FOLDER_IN_ARCHIVE = "LibriTTS"
@@ -41,7 +42,7 @@ def load_libritts_item(
     file_audio = os.path.join(path, speaker_id, chapter_id, file_audio)
 
     # Load audio
-    waveform, sample_rate = torchaudio.load(file_audio)
+    waveform, sample_rate = load_torchcodec(file_audio)
 
     # Load original text
     with open(original_text) as ft:

src/torchaudio/datasets/ljspeech.py

@@ -8,7 +8,7 @@
 from torch.utils.data import Dataset
 from torchaudio._internal import download_url_to_file
 from torchaudio.datasets.utils import _extract_tar
-
+from torchaudio.utils import load_torchcodec
 
 _RELEASE_CONFIGS = {
     "release1": {
@@ -94,7 +94,7 @@ def __getitem__(self, n: int) -> Tuple[Tensor, int, str, str]:
         fileid_audio = self._path / (fileid + ".wav")
 
         # Load audio
-        waveform, sample_rate = torchaudio.load(fileid_audio)
+        waveform, sample_rate = load_torchcodec(fileid_audio)
 
         return (
             waveform,

src/torchaudio/datasets/musdb_hq.py

@@ -7,6 +7,7 @@
 from torch.utils.data import Dataset
 from torchaudio._internal import download_url_to_file
 from torchaudio.datasets.utils import _extract_zip
+from torchaudio.utils import load_torchcodec
 
 _URL = "https://zenodo.org/record/3338373/files/musdb18hq.zip"
 _CHECKSUM = "baac80d0483c61d74b2e5f3be75fa557eec52898339e6aa45c1fa48833c5d21d"
@@ -87,7 +88,7 @@ def _load_sample(self, n: int) -> Tuple[torch.Tensor, int, int, str]:
         num_frames = None
         for source in self.sources:
             track = self._get_track(name, source)
-            wav, sr = torchaudio.load(str(track))
+            wav, sr = load_torchcodec(str(track))
             if sr != _SAMPLE_RATE:
                 raise ValueError(f"expected sample rate {_SAMPLE_RATE}, but got {sr}")
             if num_frames is None:

src/torchaudio/datasets/tedlium.py

@@ -7,6 +7,7 @@
 from torch.utils.data import Dataset
 from torchaudio._internal import download_url_to_file
 from torchaudio.datasets.utils import _extract_tar
+from torchaudio.utils import load_torchcodec
 
 
 _RELEASE_CONFIGS = {
@@ -163,12 +164,7 @@ def _load_audio(self, path: str, start_time: float, end_time: float, sample_rate
         Returns:
             [Tensor, int]: Audio tensor representation and sample rate
         """
-        start_time = int(float(start_time) * sample_rate)
-        end_time = int(float(end_time) * sample_rate)
-
-        kwargs = {"frame_offset": start_time, "num_frames": end_time - start_time}
-
-        return torchaudio.load(path, **kwargs)
+        return load_torchcodec(path, start_seconds=float(start_time), stop_seconds=float(end_time))
 
     def __getitem__(self, n: int) -> Tuple[Tensor, int, str, int, int, int]:
         """Load the n-th sample from the dataset.