test_infer_batch.py

import os
import time
import random
from tqdm import tqdm
import argparse

import torch
import torchaudio
from accelerate import Accelerator
from einops import rearrange
from ema_pytorch import EMA
from vocos import Vocos

from model import CFM, UNetT, DiT
from model.utils import (
    get_tokenizer, 
    get_seedtts_testset_metainfo, 
    get_librispeech_test_clean_metainfo, 
    get_inference_prompt,
)

accelerator = Accelerator()
device = f"cuda:{accelerator.process_index}"


# --------------------- Dataset Settings -------------------- #

target_sample_rate = 24000
n_mel_channels = 100
hop_length = 256
target_rms = 0.1

tokenizer = "pinyin"


# ---------------------- infer setting ---------------------- #

parser = argparse.ArgumentParser(description="batch inference")

parser.add_argument('-s', '--seed', default=None, type=int)
parser.add_argument('-d', '--dataset', default="Emilia_ZH_EN")
parser.add_argument('-n', '--expname', required=True)
parser.add_argument('-c', '--ckptstep', default=1200000, type=int)

parser.add_argument('-nfe', '--nfestep', default=32, type=int)
parser.add_argument('-o', '--odemethod', default="euler")
parser.add_argument('-ss', '--swaysampling', default=-1, type=float)

parser.add_argument('-t', '--testset', required=True)

args = parser.parse_args()


seed = args.seed
dataset_name = args.dataset
exp_name = args.expname
ckpt_step = args.ckptstep
checkpoint = torch.load(f"ckpts/{exp_name}/model_{ckpt_step}.pt", map_location=device)

nfe_step = args.nfestep
ode_method = args.odemethod
sway_sampling_coef = args.swaysampling

testset = args.testset


infer_batch_size = 1  # max frames. 1 for ddp single inference (recommended)
cfg_strength = 2.
speed = 1.
use_truth_duration = False
no_ref_audio = False


if exp_name == "F5TTS_Base":
    model_cls = DiT
    model_cfg = dict(dim = 1024, depth = 22, heads = 16, ff_mult = 2, text_dim = 512, conv_layers = 4)

elif exp_name == "E2TTS_Base":
    model_cls = UNetT
    model_cfg = dict(dim = 1024, depth = 24, heads = 16, ff_mult = 4)


if testset == "ls_pc_test_clean":
    metalst = "data/librispeech_pc_test_clean_cross_sentence.lst"
    librispeech_test_clean_path = "<SOME_PATH>/LibriSpeech/test-clean"  # test-clean path
    metainfo = get_librispeech_test_clean_metainfo(metalst, librispeech_test_clean_path)
    
elif testset == "seedtts_test_zh":
    metalst = "data/seedtts_testset/zh/meta.lst"
    metainfo = get_seedtts_testset_metainfo(metalst)

elif testset == "seedtts_test_en":
    metalst = "data/seedtts_testset/en/meta.lst"
    metainfo = get_seedtts_testset_metainfo(metalst)


# path to save genereted wavs
if seed is None: seed = random.randint(-10000, 10000)
output_dir = f"results/{exp_name}_{ckpt_step}/{testset}/" \
    f"seed{seed}_{ode_method}_nfe{nfe_step}" \
    f"{f'_ss{sway_sampling_coef}' if sway_sampling_coef else ''}" \
    f"_cfg{cfg_strength}_speed{speed}" \
    f"{'_gt-dur' if use_truth_duration else ''}" \
    f"{'_no-ref-audio' if no_ref_audio else ''}"


# -------------------------------------------------#

use_ema = True

prompts_all = get_inference_prompt(
    metainfo, 
    speed = speed, 
    tokenizer = tokenizer, 
    target_sample_rate = target_sample_rate, 
    n_mel_channels = n_mel_channels,
    hop_length = hop_length,
    target_rms = target_rms,
    use_truth_duration = use_truth_duration,
    infer_batch_size = infer_batch_size,
)

# Vocoder model
local = False
if local:
    vocos_local_path = "../checkpoints/charactr/vocos-mel-24khz"
    vocos = Vocos.from_hparams(f"{vocos_local_path}/config.yaml")
    state_dict = torch.load(f"{vocos_local_path}/pytorch_model.bin", map_location=device)
    vocos.load_state_dict(state_dict)
    vocos.eval()
else:
    vocos = Vocos.from_pretrained("charactr/vocos-mel-24khz")

# Tokenizer
vocab_char_map, vocab_size = get_tokenizer(dataset_name, tokenizer)

# Model
model = CFM(
    transformer = model_cls(
        **model_cfg,
        text_num_embeds = vocab_size, 
        mel_dim = n_mel_channels
    ),
    mel_spec_kwargs = dict(
        target_sample_rate = target_sample_rate, 
        n_mel_channels = n_mel_channels,
        hop_length = hop_length,
    ),
    odeint_kwargs = dict(
        method = ode_method,
    ),
    vocab_char_map = vocab_char_map,
).to(device)

if use_ema == True:
    ema_model = EMA(model, include_online_model = False).to(device)
    ema_model.load_state_dict(checkpoint['ema_model_state_dict'])
    ema_model.copy_params_from_ema_to_model()
else:
    model.load_state_dict(checkpoint['model_state_dict'])

if not os.path.exists(output_dir) and accelerator.is_main_process:
    os.makedirs(output_dir)

# start batch inference
accelerator.wait_for_everyone()
start = time.time()

with accelerator.split_between_processes(prompts_all) as prompts:

    for prompt in tqdm(prompts, disable=not accelerator.is_local_main_process):
        utts, ref_rms_list, ref_mels, ref_mel_lens, total_mel_lens, final_text_list = prompt
        ref_mels = ref_mels.to(device)
        ref_mel_lens = torch.tensor(ref_mel_lens, dtype = torch.long).to(device)
        total_mel_lens = torch.tensor(total_mel_lens, dtype = torch.long).to(device)
        
        # Inference
        with torch.inference_mode():
            generated, _ = model.sample(
                cond = ref_mels,
                text = final_text_list,
                duration = total_mel_lens,
                lens = ref_mel_lens,
                steps = nfe_step,
                cfg_strength = cfg_strength,
                sway_sampling_coef = sway_sampling_coef,
                no_ref_audio = no_ref_audio,
                seed = seed,
            )
        # Final result
        for i, gen in enumerate(generated):
            gen = gen[ref_mel_lens[i]:total_mel_lens[i], :].unsqueeze(0)
            gen_mel_spec = rearrange(gen, '1 n d -> 1 d n')
            generated_wave = vocos.decode(gen_mel_spec.cpu())
            if ref_rms_list[i] < target_rms:
                generated_wave = generated_wave * ref_rms_list[i] / target_rms
            torchaudio.save(f"{output_dir}/{utts[i]}.wav", generated_wave, target_sample_rate)

accelerator.wait_for_everyone()
if accelerator.is_main_process:
    timediff = time.time() - start
    print(f"Done batch inference in {timediff / 60 :.2f} minutes.")