kkvc-hf/Style-Bert-VITS2-AS2

1import argparse2from concurrent.futures import ThreadPoolExecutor, as_completed3from multiprocessing import cpu_count4from pathlib import Path5from typing import Any6 7import librosa8import pyloudnorm as pyln9import soundfile10from numpy.typing import NDArray11from tqdm import tqdm12 13from config import get_config14from style_bert_vits2.logging import logger15from style_bert_vits2.utils.stdout_wrapper import SAFE_STDOUT16 17 18DEFAULT_BLOCK_SIZE: float = 0.400  # seconds19 20 21class BlockSizeException(Exception):22    pass23 24 25def normalize_audio(data: NDArray[Any], sr: int):26    meter = pyln.Meter(sr, block_size=DEFAULT_BLOCK_SIZE)  # create BS.1770 meter27    try:28        loudness = meter.integrated_loudness(data)29    except ValueError as e:30        raise BlockSizeException(e)31 32    data = pyln.normalize.loudness(data, loudness, -23.0)33    return data34 35 36def resample(37    file: Path,38    input_dir: Path,39    output_dir: Path,40    target_sr: int,41    normalize: bool,42    trim: bool,43):44    """45    fileを読み込んで、target_srなwavファイルに変換して、46    output_dirの中に、input_dirからの相対パスを保つように保存する47    """48    try:49        # librosaが読めるファイルかチェック50        # wav以外にもmp3やoggやflacなども読める51        wav: NDArray[Any]52        sr: int53        wav, sr = librosa.load(file, sr=target_sr)54        if normalize:55            try:56                wav = normalize_audio(wav, sr)57            except BlockSizeException:58                print("")59                logger.info(60                    f"Skip normalize due to less than {DEFAULT_BLOCK_SIZE} second audio: {file}"61                )62        if trim:63            wav, _ = librosa.effects.trim(wav, top_db=30)64        relative_path = file.relative_to(input_dir)65        # ここで拡張子が.wav以外でも.wavに置き換えられる66        output_path = output_dir / relative_path.with_suffix(".wav")67        output_path.parent.mkdir(parents=True, exist_ok=True)68        soundfile.write(output_path, wav, sr)69    except Exception as e:70        logger.warning(f"Cannot load file, so skipping: {file}, {e}")71 72 73if __name__ == "__main__":74    config = get_config()75    parser = argparse.ArgumentParser()76    parser.add_argument(77        "--sr",78        type=int,79        default=config.resample_config.sampling_rate,80        help="sampling rate",81    )82    parser.add_argument(83        "--input_dir",84        "-i",85        type=str,86        default=config.resample_config.in_dir,87        help="path to source dir",88    )89    parser.add_argument(90        "--output_dir",91        "-o",92        type=str,93        default=config.resample_config.out_dir,94        help="path to target dir",95    )96    parser.add_argument(97        "--num_processes",98        type=int,99        default=4,100        help="cpu_processes",101    )102    parser.add_argument(103        "--normalize",104        action="store_true",105        default=False,106        help="loudness normalize audio",107    )108    parser.add_argument(109        "--trim",110        action="store_true",111        default=False,112        help="trim silence (start and end only)",113    )114    args = parser.parse_args()115 116    if args.num_processes == 0:117        processes = cpu_count() - 2 if cpu_count() > 4 else 1118    else:119        processes: int = args.num_processes120 121    input_dir = Path(args.input_dir)122    output_dir = Path(args.output_dir)123    logger.info(f"Resampling {input_dir} to {output_dir}")124    sr = int(args.sr)125    normalize: bool = args.normalize126    trim: bool = args.trim127 128    # 後でlibrosaに読ませて有効な音声ファイルかチェックするので、全てのファイルを取得129    original_files = [f for f in input_dir.rglob("*") if f.is_file()]130 131    if len(original_files) == 0:132        logger.error(f"No files found in {input_dir}")133        raise ValueError(f"No files found in {input_dir}")134 135    output_dir.mkdir(parents=True, exist_ok=True)136 137    with ThreadPoolExecutor(max_workers=processes) as executor:138        futures = [139            executor.submit(resample, file, input_dir, output_dir, sr, normalize, trim)140            for file in original_files141        ]142        for future in tqdm(143            as_completed(futures), total=len(original_files), file=SAFE_STDOUT144        ):145            pass146 147    logger.info("Resampling Done!")148