refactor: update audio loading and saving methods to use torchcodec

music_dcae/music_dcae_pipeline.py

@@ -45,7 +45,7 @@ class MusicDCAE(ModelMixin, ConfigMixin, FromOriginalModelMixin):
         self.shift_factor = -1.9091
 
     def load_audio(self, audio_path):
-        audio, sr = torchaudio.load(audio_path)
+        audio, sr = torchaudio.load_with_torchcodec(audio_path)
         return audio, sr
 
     def forward_mel(self, audios):
@@ -121,7 +121,7 @@ class MusicDCAE(ModelMixin, ConfigMixin, FromOriginalModelMixin):
 
 if __name__ == "__main__":
 
-    audio, sr = torchaudio.load("test.wav")
+    audio, sr = torchaudio.load_with_torchcodec("test.wav")
     audio_lengths = torch.tensor([audio.shape[1]])
     audios = audio.unsqueeze(0)
     
@@ -137,5 +137,5 @@ if __name__ == "__main__":
     print("latents shape: ", latents.shape)
     print("latent_lengths: ", latent_lengths)
     print("sr: ", sr)
-    torchaudio.save("test_reconstructed.flac", pred_wavs[0], sr)
+    torchaudio.save_with_torchcodec("test_reconstructed.flac", pred_wavs[0], sr)
     print("test_reconstructed.flac")

pipeline_ace_step.py

@@ -36,7 +36,6 @@ from apg_guidance import (
     cfg_double_condition_forward,
 )
 import torchaudio
-import torio
 
 
 torch.backends.cudnn.benchmark = False
@@ -1428,12 +1427,11 @@ class ACEStepPipeline:
             f"{base_path}/output_{time.strftime('%Y%m%d%H%M%S')}_{idx}.{format}"
         )
         target_wav = target_wav.float()
-        torchaudio.save(
+        torchaudio.save_with_torchcodec(
             output_path_flac,
             target_wav,
             sample_rate=sample_rate,
-            format=format,
-            compression=torio.io.CodecConfig(bit_rate=320000),
+            compression=320000,
         )
         return output_path_flac
 

requirements.txt

@@ -10,6 +10,7 @@ pytorch_lightning==2.5.1
 soundfile==0.13.1
 torch==2.8.0
 torchaudio==2.8.0
+torchcodec>=0.2
 torchvision==0.23.0
 tqdm
 transformers==4.50.0