DiffSpeech / mfa_usr /run_mfa_align.py
RayeRen's picture
init
d1b91e7
import argparse
import glob
import importlib
import os
import subprocess
from utils.commons.hparams import set_hparams, hparams
from utils.commons.multiprocess_utils import multiprocess_run_tqdm
from utils.os_utils import remove_file
from utils.text.encoding import get_encoding
def process_item(idx, txt_fn):
base_fn = os.path.splitext(txt_fn)[0]
basename = os.path.basename(base_fn)
if os.path.exists(base_fn + '.wav'):
wav_fn = base_fn + '.wav'
elif os.path.exists(base_fn + '.mp3'):
wav_fn = base_fn + '.mp3'
else:
return
# process text
encoding = get_encoding(txt_fn)
with open(txt_fn, encoding=encoding) as f:
txt_raw = " ".join(f.readlines()).strip()
phs, _, phs_for_align, _ = preprocesser.process_text(txt_processor, txt_raw, hparams['preprocess_args'])
os.makedirs(f'{mfa_process_dir}/{basename}', exist_ok=True)
with open(f'{mfa_process_dir}/{basename}/{basename}.lab', 'w') as f:
f.write(phs_for_align)
# process wav
new_wav_fn = preprocesser.process_wav(basename, wav_fn, mfa_process_dir, preprocess_args)
subprocess.check_call(f'cp "{new_wav_fn}" "{mfa_process_dir}/{basename}/{basename}.wav"', shell=True)
if __name__ == "__main__":
set_hparams()
parser = argparse.ArgumentParser(description='')
parser.add_argument('--input_dir', type=str, default='', help='input dir')
args, unknown = parser.parse_known_args()
input_dir = args.input_dir
processed_data_dir = hparams['processed_data_dir']
preprocess_args = hparams['preprocess_args']
preprocess_args['sox_to_wav'] = True
preprocess_args['trim_all_sil'] = True
# preprocess_args['trim_sil'] = True
# preprocess_args['denoise'] = True
pkg = ".".join(hparams["preprocess_cls"].split(".")[:-1])
cls_name = hparams["preprocess_cls"].split(".")[-1]
process_cls = getattr(importlib.import_module(pkg), cls_name)
preprocesser = process_cls()
txt_processor = preprocesser.txt_processor
num_workers = int(os.getenv('N_PROC', os.cpu_count()))
mfa_process_dir = f'{input_dir}/mfa_inputs'
remove_file(mfa_process_dir, f'{input_dir}/mfa_tmp')
os.makedirs(mfa_process_dir, exist_ok=True)
os.makedirs(f'{mfa_process_dir}/processed_tmp', exist_ok=True)
for res in multiprocess_run_tqdm(
process_item, list(enumerate(glob.glob(f'{input_dir}/*.txt')))):
pass
remove_file(f'{mfa_process_dir}/processed_tmp')
subprocess.check_call(
f'mfa align {mfa_process_dir} ' # process dir
f'{hparams["processed_data_dir"]}/mfa_dict.txt ' # dict
f'{input_dir}/mfa_model.zip ' # model
f'{input_dir}/mfa_outputs -t {input_dir}/mfa_tmp -j {num_workers} '
f' && cp -rf {input_dir}/mfa_outputs/*/* {input_dir}/'
f' && cp -rf {mfa_process_dir}/*/* {input_dir}/'
f' && rm -rf {input_dir}/mfa_tmp {input_dir}/mfa_outputs {mfa_process_dir}', # remove tmp dir
shell=True)