Spaces:
Running
Running
| # Copyright (c) Meta Platforms, Inc. and affiliates. | |
| # All rights reserved. | |
| # | |
| # This source code is licensed under the license found in the | |
| # LICENSE file in the root directory of this source tree. | |
| from itertools import product | |
| import random | |
| import numpy as np | |
| import torch | |
| import torchaudio | |
| from audiocraft.data.audio import audio_info, audio_read, audio_write, _av_read | |
| from ..common_utils import TempDirMixin, get_white_noise, save_wav | |
| class TestInfo(TempDirMixin): | |
| def test_info_mp3(self): | |
| sample_rates = [8000, 16_000] | |
| channels = [1, 2] | |
| duration = 1. | |
| for sample_rate, ch in product(sample_rates, channels): | |
| wav = get_white_noise(ch, int(sample_rate * duration)) | |
| path = self.get_temp_path('sample_wav.mp3') | |
| save_wav(path, wav, sample_rate) | |
| info = audio_info(path) | |
| assert info.sample_rate == sample_rate | |
| assert info.channels == ch | |
| # we cannot trust torchaudio for num_frames, so we don't check | |
| def _test_info_format(self, ext: str): | |
| sample_rates = [8000, 16_000] | |
| channels = [1, 2] | |
| duration = 1. | |
| for sample_rate, ch in product(sample_rates, channels): | |
| n_frames = int(sample_rate * duration) | |
| wav = get_white_noise(ch, n_frames) | |
| path = self.get_temp_path(f'sample_wav{ext}') | |
| save_wav(path, wav, sample_rate) | |
| info = audio_info(path) | |
| assert info.sample_rate == sample_rate | |
| assert info.channels == ch | |
| assert np.isclose(info.duration, duration, atol=1e-5) | |
| def test_info_wav(self): | |
| self._test_info_format('.wav') | |
| def test_info_flac(self): | |
| self._test_info_format('.flac') | |
| def test_info_ogg(self): | |
| self._test_info_format('.ogg') | |
| def test_info_m4a(self): | |
| # TODO: generate m4a file programmatically | |
| # self._test_info_format('.m4a') | |
| pass | |
| class TestRead(TempDirMixin): | |
| def test_read_full_wav(self): | |
| sample_rates = [8000, 16_000] | |
| channels = [1, 2] | |
| duration = 1. | |
| for sample_rate, ch in product(sample_rates, channels): | |
| n_frames = int(sample_rate * duration) | |
| wav = get_white_noise(ch, n_frames).clamp(-0.99, 0.99) | |
| path = self.get_temp_path('sample_wav.wav') | |
| save_wav(path, wav, sample_rate) | |
| read_wav, read_sr = audio_read(path) | |
| assert read_sr == sample_rate | |
| assert read_wav.shape[0] == wav.shape[0] | |
| assert read_wav.shape[1] == wav.shape[1] | |
| assert torch.allclose(read_wav, wav, rtol=1e-03, atol=1e-04) | |
| def test_read_partial_wav(self): | |
| sample_rates = [8000, 16_000] | |
| channels = [1, 2] | |
| duration = 1. | |
| read_duration = torch.rand(1).item() | |
| for sample_rate, ch in product(sample_rates, channels): | |
| n_frames = int(sample_rate * duration) | |
| read_frames = int(sample_rate * read_duration) | |
| wav = get_white_noise(ch, n_frames).clamp(-0.99, 0.99) | |
| path = self.get_temp_path('sample_wav.wav') | |
| save_wav(path, wav, sample_rate) | |
| read_wav, read_sr = audio_read(path, 0, read_duration) | |
| assert read_sr == sample_rate | |
| assert read_wav.shape[0] == wav.shape[0] | |
| assert read_wav.shape[1] == read_frames | |
| assert torch.allclose(read_wav[..., 0:read_frames], wav[..., 0:read_frames], rtol=1e-03, atol=1e-04) | |
| def test_read_seek_time_wav(self): | |
| sample_rates = [8000, 16_000] | |
| channels = [1, 2] | |
| duration = 1. | |
| read_duration = 1. | |
| for sample_rate, ch in product(sample_rates, channels): | |
| n_frames = int(sample_rate * duration) | |
| wav = get_white_noise(ch, n_frames).clamp(-0.99, 0.99) | |
| path = self.get_temp_path('sample_wav.wav') | |
| save_wav(path, wav, sample_rate) | |
| seek_time = torch.rand(1).item() | |
| read_wav, read_sr = audio_read(path, seek_time, read_duration) | |
| seek_frames = int(sample_rate * seek_time) | |
| expected_frames = n_frames - seek_frames | |
| assert read_sr == sample_rate | |
| assert read_wav.shape[0] == wav.shape[0] | |
| assert read_wav.shape[1] == expected_frames | |
| assert torch.allclose(read_wav, wav[..., seek_frames:], rtol=1e-03, atol=1e-04) | |
| def test_read_seek_time_wav_padded(self): | |
| sample_rates = [8000, 16_000] | |
| channels = [1, 2] | |
| duration = 1. | |
| read_duration = 1. | |
| for sample_rate, ch in product(sample_rates, channels): | |
| n_frames = int(sample_rate * duration) | |
| read_frames = int(sample_rate * read_duration) | |
| wav = get_white_noise(ch, n_frames).clamp(-0.99, 0.99) | |
| path = self.get_temp_path('sample_wav.wav') | |
| save_wav(path, wav, sample_rate) | |
| seek_time = torch.rand(1).item() | |
| seek_frames = int(sample_rate * seek_time) | |
| expected_frames = n_frames - seek_frames | |
| read_wav, read_sr = audio_read(path, seek_time, read_duration, pad=True) | |
| expected_pad_wav = torch.zeros(wav.shape[0], read_frames - expected_frames) | |
| assert read_sr == sample_rate | |
| assert read_wav.shape[0] == wav.shape[0] | |
| assert read_wav.shape[1] == read_frames | |
| assert torch.allclose(read_wav[..., :expected_frames], wav[..., seek_frames:], rtol=1e-03, atol=1e-04) | |
| assert torch.allclose(read_wav[..., expected_frames:], expected_pad_wav) | |
| class TestAvRead(TempDirMixin): | |
| def test_avread_seek_base(self): | |
| sample_rates = [8000, 16_000] | |
| channels = [1, 2] | |
| duration = 2. | |
| for sample_rate, ch in product(sample_rates, channels): | |
| n_frames = int(sample_rate * duration) | |
| wav = get_white_noise(ch, n_frames) | |
| path = self.get_temp_path(f'reference_a_{sample_rate}_{ch}.wav') | |
| save_wav(path, wav, sample_rate) | |
| for _ in range(100): | |
| # seek will always load a full duration segment in the file | |
| seek_time = random.uniform(0.0, 1.0) | |
| seek_duration = random.uniform(0.001, 1.0) | |
| read_wav, read_sr = _av_read(path, seek_time, seek_duration) | |
| assert read_sr == sample_rate | |
| assert read_wav.shape[0] == wav.shape[0] | |
| assert read_wav.shape[-1] == int(seek_duration * sample_rate) | |
| def test_avread_seek_partial(self): | |
| sample_rates = [8000, 16_000] | |
| channels = [1, 2] | |
| duration = 1. | |
| for sample_rate, ch in product(sample_rates, channels): | |
| n_frames = int(sample_rate * duration) | |
| wav = get_white_noise(ch, n_frames) | |
| path = self.get_temp_path(f'reference_b_{sample_rate}_{ch}.wav') | |
| save_wav(path, wav, sample_rate) | |
| for _ in range(100): | |
| # seek will always load a partial segment | |
| seek_time = random.uniform(0.5, 1.) | |
| seek_duration = 1. | |
| expected_num_frames = n_frames - int(seek_time * sample_rate) | |
| read_wav, read_sr = _av_read(path, seek_time, seek_duration) | |
| assert read_sr == sample_rate | |
| assert read_wav.shape[0] == wav.shape[0] | |
| assert read_wav.shape[-1] == expected_num_frames | |
| def test_avread_seek_outofbound(self): | |
| sample_rates = [8000, 16_000] | |
| channels = [1, 2] | |
| duration = 1. | |
| for sample_rate, ch in product(sample_rates, channels): | |
| n_frames = int(sample_rate * duration) | |
| wav = get_white_noise(ch, n_frames) | |
| path = self.get_temp_path(f'reference_c_{sample_rate}_{ch}.wav') | |
| save_wav(path, wav, sample_rate) | |
| seek_time = 1.5 | |
| read_wav, read_sr = _av_read(path, seek_time, 1.) | |
| assert read_sr == sample_rate | |
| assert read_wav.shape[0] == wav.shape[0] | |
| assert read_wav.shape[-1] == 0 | |
| def test_avread_seek_edge(self): | |
| sample_rates = [8000, 16_000] | |
| # some of these values will have | |
| # int(((frames - 1) / sample_rate) * sample_rate) != (frames - 1) | |
| n_frames = [1000, 1001, 1002] | |
| channels = [1, 2] | |
| for sample_rate, ch, frames in product(sample_rates, channels, n_frames): | |
| duration = frames / sample_rate | |
| wav = get_white_noise(ch, frames) | |
| path = self.get_temp_path(f'reference_d_{sample_rate}_{ch}.wav') | |
| save_wav(path, wav, sample_rate) | |
| seek_time = (frames - 1) / sample_rate | |
| seek_frames = int(seek_time * sample_rate) | |
| read_wav, read_sr = _av_read(path, seek_time, duration) | |
| assert read_sr == sample_rate | |
| assert read_wav.shape[0] == wav.shape[0] | |
| assert read_wav.shape[-1] == (frames - seek_frames) | |
| class TestAudioWrite(TempDirMixin): | |
| def test_audio_write_wav(self): | |
| torch.manual_seed(1234) | |
| sample_rates = [8000, 16_000] | |
| n_frames = [1000, 1001, 1002] | |
| channels = [1, 2] | |
| strategies = ["peak", "clip", "rms"] | |
| formats = ["wav", "mp3"] | |
| for sample_rate, ch, frames in product(sample_rates, channels, n_frames): | |
| for format_, strategy in product(formats, strategies): | |
| wav = get_white_noise(ch, frames) | |
| path = self.get_temp_path(f'pred_{sample_rate}_{ch}') | |
| audio_write(path, wav, sample_rate, format_, strategy=strategy) | |
| read_wav, read_sr = torchaudio.load(f'{path}.{format_}') | |
| if format_ == "wav": | |
| assert read_wav.shape == wav.shape | |
| if format_ == "wav" and strategy in ["peak", "rms"]: | |
| rescaled_read_wav = read_wav / read_wav.abs().max() * wav.abs().max() | |
| # for a Gaussian, the typical max scale will be less than ~5x the std. | |
| # The error when writing to disk will ~ 1/2**15, and when rescaling, 5x that. | |
| # For RMS target, rescaling leaves more headroom by default, leading | |
| # to a 20x rescaling typically | |
| atol = (5 if strategy == "peak" else 20) / 2**15 | |
| delta = (rescaled_read_wav - wav).abs().max() | |
| assert torch.allclose(wav, rescaled_read_wav, rtol=0, atol=atol), (delta, atol) | |
| formats = ["wav"] # faster unit tests | |