ThreadAbort/E2-F5-TTS

1import os2import re3 4import torch5import torchaudio6from einops import rearrange7from ema_pytorch import EMA8from vocos import Vocos9 10from model import CFM, UNetT, DiT, MMDiT11from model.utils import (12    get_tokenizer, 13    convert_char_to_pinyin, 14    save_spectrogram,15)16 17device = "cuda" if torch.cuda.is_available() else "cpu"18 19 20# --------------------- Dataset Settings -------------------- #21 22target_sample_rate = 2400023n_mel_channels = 10024hop_length = 25625target_rms = 0.126 27tokenizer = "pinyin"28dataset_name = "Emilia_ZH_EN"29 30 31# ---------------------- infer setting ---------------------- #32 33seed = None  # int | None34 35exp_name = "F5TTS_Base"  # F5TTS_Base | E2TTS_Base36ckpt_step = 120000037 38nfe_step = 32  # 16, 3239cfg_strength = 2.40ode_method = 'euler'  # euler | midpoint41sway_sampling_coef = -1.42speed = 1.43fix_duration = 27  # None (will linear estimate. if code-switched, consider fix) | float (total in seconds, include ref audio) 44 45if exp_name == "F5TTS_Base":46    model_cls = DiT47    model_cfg = dict(dim = 1024, depth = 22, heads = 16, ff_mult = 2, text_dim = 512, conv_layers = 4)48 49elif exp_name == "E2TTS_Base":50    model_cls = UNetT51    model_cfg = dict(dim = 1024, depth = 24, heads = 16, ff_mult = 4)52 53checkpoint = torch.load(f"ckpts/{exp_name}/model_{ckpt_step}.pt", map_location=device)54output_dir = "tests"55 56ref_audio = "tests/ref_audio/test_en_1_ref_short.wav"57ref_text = "Some call me nature, others call me mother nature."58gen_text = "I don't really care what you call me. I've been a silent spectator, watching species evolve, empires rise and fall. But always remember, I am mighty and enduring. Respect me and I'll nurture you; ignore me and you shall face the consequences."59 60# ref_audio = "tests/ref_audio/test_zh_1_ref_short.wav"61# ref_text = "对，这就是我，万人敬仰的太乙真人。"62# gen_text = "突然，身边一阵笑声。我看着他们，意气风发地挺直了胸膛，甩了甩那稍显肉感的双臂，轻笑道：\"我身上的肉，是为了掩饰我爆棚的魅力，否则，岂不吓坏了你们呢？\""63 64 65# -------------------------------------------------#66 67use_ema = True68 69if not os.path.exists(output_dir):70    os.makedirs(output_dir)71 72# Vocoder model73local = False74if local:75    vocos_local_path = "../checkpoints/charactr/vocos-mel-24khz"76    vocos = Vocos.from_hparams(f"{vocos_local_path}/config.yaml")77    state_dict = torch.load(f"{vocos_local_path}/pytorch_model.bin", map_location=device)78    vocos.load_state_dict(state_dict)79    vocos.eval()80else:81    vocos = Vocos.from_pretrained("charactr/vocos-mel-24khz")82 83# Tokenizer84vocab_char_map, vocab_size = get_tokenizer(dataset_name, tokenizer)85 86# Model87model = CFM(88    transformer = model_cls(89        **model_cfg,90        text_num_embeds = vocab_size, 91        mel_dim = n_mel_channels92    ),93    mel_spec_kwargs = dict(94        target_sample_rate = target_sample_rate, 95        n_mel_channels = n_mel_channels,96        hop_length = hop_length,97    ),98    odeint_kwargs = dict(99        method = ode_method,100    ),101    vocab_char_map = vocab_char_map,102).to(device)103 104if use_ema == True:105    ema_model = EMA(model, include_online_model = False).to(device)106    ema_model.load_state_dict(checkpoint['ema_model_state_dict'])107    ema_model.copy_params_from_ema_to_model()108else:109    model.load_state_dict(checkpoint['model_state_dict'])110 111# Audio112audio, sr = torchaudio.load(ref_audio)113rms = torch.sqrt(torch.mean(torch.square(audio)))114if rms < target_rms:115    audio = audio * target_rms / rms116if sr != target_sample_rate:117    resampler = torchaudio.transforms.Resample(sr, target_sample_rate)118    audio = resampler(audio)119audio = audio.to(device)120 121# Text122text_list = [ref_text + gen_text]123if tokenizer == "pinyin":124    final_text_list = convert_char_to_pinyin(text_list)125else:126    final_text_list = [text_list]127print(f"text  : {text_list}")128print(f"pinyin: {final_text_list}")129 130# Duration131ref_audio_len = audio.shape[-1] // hop_length132if fix_duration is not None:133    duration = int(fix_duration * target_sample_rate / hop_length)134else:  # simple linear scale calcul135    zh_pause_punc = r"。，、；：？！"136    ref_text_len = len(ref_text) + len(re.findall(zh_pause_punc, ref_text))137    gen_text_len = len(gen_text) + len(re.findall(zh_pause_punc, gen_text))138    duration = ref_audio_len + int(ref_audio_len / ref_text_len * gen_text_len / speed)139 140# Inference141with torch.inference_mode():142    generated, trajectory = model.sample(143        cond = audio,144        text = final_text_list,145        duration = duration,146        steps = nfe_step,147        cfg_strength = cfg_strength,148        sway_sampling_coef = sway_sampling_coef,149        seed = seed,150    )151print(f"Generated mel: {generated.shape}")152 153# Final result154generated = generated[:, ref_audio_len:, :]155generated_mel_spec = rearrange(generated, '1 n d -> 1 d n')156generated_wave = vocos.decode(generated_mel_spec.cpu())157if rms < target_rms:158    generated_wave = generated_wave * rms / target_rms159 160save_spectrogram(generated_mel_spec[0].cpu().numpy(), f"{output_dir}/test_single.png")161torchaudio.save(f"{output_dir}/test_single.wav", generated_wave, target_sample_rate)162print(f"Generated wav: {generated_wave.shape}")163