prabaerode/zero-shot-tts

1import os2 3import torch4import torch.nn.functional as F5import torchaudio6 7from f5_tts.infer.utils_infer import load_checkpoint, load_vocoder, save_spectrogram8from f5_tts.model import CFM, DiT, UNetT9from f5_tts.model.utils import convert_char_to_pinyin, get_tokenizer10 11device = "cuda" if torch.cuda.is_available() else "mps" if torch.backends.mps.is_available() else "cpu"12 13 14# --------------------- Dataset Settings -------------------- #15 16target_sample_rate = 2400017n_mel_channels = 10018hop_length = 25619win_length = 102420n_fft = 102421mel_spec_type = "vocos"  # 'vocos' or 'bigvgan'22target_rms = 0.123 24tokenizer = "pinyin"25dataset_name = "Emilia_ZH_EN"26 27 28# ---------------------- infer setting ---------------------- #29 30seed = None  # int | None31 32exp_name = "F5TTS_Base"  # F5TTS_Base | E2TTS_Base33ckpt_step = 120000034 35nfe_step = 32  # 16, 3236cfg_strength = 2.037ode_method = "euler"  # euler | midpoint38sway_sampling_coef = -1.039speed = 1.040 41if exp_name == "F5TTS_Base":42    model_cls = DiT43    model_cfg = dict(dim=1024, depth=22, heads=16, ff_mult=2, text_dim=512, conv_layers=4)44 45elif exp_name == "E2TTS_Base":46    model_cls = UNetT47    model_cfg = dict(dim=1024, depth=24, heads=16, ff_mult=4)48 49ckpt_path = f"ckpts/{exp_name}/model_{ckpt_step}.safetensors"50output_dir = "tests"51 52# [leverage https://github.com/MahmoudAshraf97/ctc-forced-aligner to get char level alignment]53# pip install git+https://github.com/MahmoudAshraf97/ctc-forced-aligner.git54# [write the origin_text into a file, e.g. tests/test_edit.txt]55# ctc-forced-aligner --audio_path "src/f5_tts/infer/examples/basic/basic_ref_en.wav" --text_path "tests/test_edit.txt" --language "zho" --romanize --split_size "char"56# [result will be saved at same path of audio file]57# [--language "zho" for Chinese, "eng" for English]58# [if local ckpt, set --alignment_model "../checkpoints/mms-300m-1130-forced-aligner"]59 60audio_to_edit = "src/f5_tts/infer/examples/basic/basic_ref_en.wav"61origin_text = "Some call me nature, others call me mother nature."62target_text = "Some call me optimist, others call me realist."63parts_to_edit = [64    [1.42, 2.44],65    [4.04, 4.9],66]  # stard_ends of "nature" & "mother nature", in seconds67fix_duration = [68    1.2,69    1,70]  # fix duration for "optimist" & "realist", in seconds71 72# audio_to_edit = "src/f5_tts/infer/examples/basic/basic_ref_zh.wav"73# origin_text = "对，这就是我，万人敬仰的太乙真人。"74# target_text = "对，那就是你，万人敬仰的太白金星。"75# parts_to_edit = [[0.84, 1.4], [1.92, 2.4], [4.26, 6.26], ]76# fix_duration = None  # use origin text duration77 78 79# -------------------------------------------------#80 81use_ema = True82 83if not os.path.exists(output_dir):84    os.makedirs(output_dir)85 86# Vocoder model87local = False88if mel_spec_type == "vocos":89    vocoder_local_path = "../checkpoints/charactr/vocos-mel-24khz"90elif mel_spec_type == "bigvgan":91    vocoder_local_path = "../checkpoints/bigvgan_v2_24khz_100band_256x"92vocoder = load_vocoder(vocoder_name=mel_spec_type, is_local=local, local_path=vocoder_local_path)93 94# Tokenizer95vocab_char_map, vocab_size = get_tokenizer(dataset_name, tokenizer)96 97# Model98model = CFM(99    transformer=model_cls(**model_cfg, text_num_embeds=vocab_size, mel_dim=n_mel_channels),100    mel_spec_kwargs=dict(101        n_fft=n_fft,102        hop_length=hop_length,103        win_length=win_length,104        n_mel_channels=n_mel_channels,105        target_sample_rate=target_sample_rate,106        mel_spec_type=mel_spec_type,107    ),108    odeint_kwargs=dict(109        method=ode_method,110    ),111    vocab_char_map=vocab_char_map,112).to(device)113 114dtype = torch.float32 if mel_spec_type == "bigvgan" else None115model = load_checkpoint(model, ckpt_path, device, dtype=dtype, use_ema=use_ema)116 117# Audio118audio, sr = torchaudio.load(audio_to_edit)119if audio.shape[0] > 1:120    audio = torch.mean(audio, dim=0, keepdim=True)121rms = torch.sqrt(torch.mean(torch.square(audio)))122if rms < target_rms:123    audio = audio * target_rms / rms124if sr != target_sample_rate:125    resampler = torchaudio.transforms.Resample(sr, target_sample_rate)126    audio = resampler(audio)127offset = 0128audio_ = torch.zeros(1, 0)129edit_mask = torch.zeros(1, 0, dtype=torch.bool)130for part in parts_to_edit:131    start, end = part132    part_dur = end - start if fix_duration is None else fix_duration.pop(0)133    part_dur = part_dur * target_sample_rate134    start = start * target_sample_rate135    audio_ = torch.cat((audio_, audio[:, round(offset) : round(start)], torch.zeros(1, round(part_dur))), dim=-1)136    edit_mask = torch.cat(137        (138            edit_mask,139            torch.ones(1, round((start - offset) / hop_length), dtype=torch.bool),140            torch.zeros(1, round(part_dur / hop_length), dtype=torch.bool),141        ),142        dim=-1,143    )144    offset = end * target_sample_rate145# audio = torch.cat((audio_, audio[:, round(offset):]), dim = -1)146edit_mask = F.pad(edit_mask, (0, audio.shape[-1] // hop_length - edit_mask.shape[-1] + 1), value=True)147audio = audio.to(device)148edit_mask = edit_mask.to(device)149 150# Text151text_list = [target_text]152if tokenizer == "pinyin":153    final_text_list = convert_char_to_pinyin(text_list)154else:155    final_text_list = [text_list]156print(f"text  : {text_list}")157print(f"pinyin: {final_text_list}")158 159# Duration160ref_audio_len = 0161duration = audio.shape[-1] // hop_length162 163# Inference164with torch.inference_mode():165    generated, trajectory = model.sample(166        cond=audio,167        text=final_text_list,168        duration=duration,169        steps=nfe_step,170        cfg_strength=cfg_strength,171        sway_sampling_coef=sway_sampling_coef,172        seed=seed,173        edit_mask=edit_mask,174    )175    print(f"Generated mel: {generated.shape}")176 177    # Final result178    generated = generated.to(torch.float32)179    generated = generated[:, ref_audio_len:, :]180    gen_mel_spec = generated.permute(0, 2, 1)181    if mel_spec_type == "vocos":182        generated_wave = vocoder.decode(gen_mel_spec)183    elif mel_spec_type == "bigvgan":184        generated_wave = vocoder(gen_mel_spec)185 186    if rms < target_rms:187        generated_wave = generated_wave * rms / target_rms188 189    save_spectrogram(gen_mel_spec[0].cpu().numpy(), f"{output_dir}/speech_edit_out.png")190    torchaudio.save(f"{output_dir}/speech_edit_out.wav", generated_wave.squeeze(0).cpu(), target_sample_rate)191    print(f"Generated wav: {generated_wave.shape}")192