codemo/fish-speech-1

1import random2from multiprocessing import Pool3from pathlib import Path4 5import click6import librosa7import torch.nn.functional as F8import torchaudio9from tqdm import tqdm10 11from tools.file import AUDIO_EXTENSIONS, list_files12 13threshold = 10 ** (-50 / 20.0)14 15 16def process(file):17    waveform, sample_rate = torchaudio.load(str(file), backend="sox")18    loudness = librosa.feature.rms(19        y=waveform.numpy().squeeze(), frame_length=2048, hop_length=512, center=True20    )[0]21    for i in range(len(loudness) - 1, 0, -1):22        if loudness[i] > threshold:23            break24 25    silent_time = (len(loudness) - i) * 512 / sample_rate26 27    if silent_time <= 0.3:28        random_time = random.uniform(0.3, 0.7)29        waveform = F.pad(30            waveform, (0, int(random_time * sample_rate)), mode="constant", value=031        )32 33    torchaudio.save(uri=str(file), src=waveform, sample_rate=sample_rate)34 35 36@click.command()37@click.argument("source", type=Path)38@click.option("--num-workers", type=int, default=12)39def main(source, num_workers):40    files = list(list_files(source, AUDIO_EXTENSIONS, recursive=True))41 42    with Pool(num_workers) as p:43        list(tqdm(p.imap_unordered(process, files), total=len(files)))44 45 46if __name__ == "__main__":47    main()48