ThreadAbort/E2-F5-TTS

1print("WARNING: You are running this unofficial E2/F5 TTS demo locally, it may not be as up-to-date as the hosted version (https://huggingface.co/spaces/mrfakename/E2-F5-TTS)")2 3import os4import re5import torch6import torchaudio7import gradio as gr8import numpy as np9import tempfile10from einops import rearrange11from ema_pytorch import EMA12from vocos import Vocos13from pydub import AudioSegment14from model import CFM, UNetT, DiT, MMDiT15from cached_path import cached_path16from model.utils import (17    get_tokenizer, 18    convert_char_to_pinyin, 19    save_spectrogram,20)21from transformers import pipeline22import librosa23from txtsplit import txtsplit24 25device = "cuda" if torch.cuda.is_available() else "mps" if torch.backends.mps.is_available() else "cpu"26 27pipe = pipeline(28    "automatic-speech-recognition",29    model="openai/whisper-large-v3-turbo",30    torch_dtype=torch.float16,31    device=device,32)33 34# --------------------- Settings -------------------- #35 36target_sample_rate = 2400037n_mel_channels = 10038hop_length = 25639target_rms = 0.140nfe_step = 32  # 16, 3241cfg_strength = 2.042ode_method = 'euler'43sway_sampling_coef = -1.044speed = 1.045# fix_duration = 27  # None or float (duration in seconds)46fix_duration = None47 48def load_model(exp_name, model_cls, model_cfg, ckpt_step):49    checkpoint = torch.load(str(cached_path(f"hf://SWivid/F5-TTS/{exp_name}/model_{ckpt_step}.pt")), map_location=device)50    vocab_char_map, vocab_size = get_tokenizer("Emilia_ZH_EN", "pinyin")51    model = CFM(52        transformer=model_cls(53            **model_cfg,54            text_num_embeds=vocab_size,55            mel_dim=n_mel_channels56        ),57        mel_spec_kwargs=dict(58            target_sample_rate=target_sample_rate,59            n_mel_channels=n_mel_channels,60            hop_length=hop_length,61        ),62        odeint_kwargs=dict(63            method=ode_method,64        ),65        vocab_char_map=vocab_char_map,66    ).to(device)67 68    ema_model = EMA(model, include_online_model=False).to(device)69    ema_model.load_state_dict(checkpoint['ema_model_state_dict'])70    ema_model.copy_params_from_ema_to_model()71 72    return ema_model, model73 74# load models75F5TTS_model_cfg = dict(dim=1024, depth=22, heads=16, ff_mult=2, text_dim=512, conv_layers=4)76E2TTS_model_cfg = dict(dim=1024, depth=24, heads=16, ff_mult=4)77 78F5TTS_ema_model, F5TTS_base_model = load_model("F5TTS_Base", DiT, F5TTS_model_cfg, 1200000)79E2TTS_ema_model, E2TTS_base_model = load_model("E2TTS_Base", UNetT, E2TTS_model_cfg, 1200000)80 81def infer(ref_audio_orig, ref_text, gen_text, exp_name, remove_silence, progress = gr.Progress()):82    print(gen_text)83    gr.Info("Converting audio...")84    with tempfile.NamedTemporaryFile(delete=False, suffix=".wav") as f:85        aseg = AudioSegment.from_file(ref_audio_orig)86        # Convert to mono87        aseg = aseg.set_channels(1)88        audio_duration = len(aseg)89        if audio_duration > 15000:90            gr.Warning("Audio is over 15s, clipping to only first 15s.")91            aseg = aseg[:15000]92        aseg.export(f.name, format="wav")93        ref_audio = f.name94    if exp_name == "F5-TTS":95        ema_model = F5TTS_ema_model96        base_model = F5TTS_base_model97    elif exp_name == "E2-TTS":98        ema_model = E2TTS_ema_model99        base_model = E2TTS_base_model100    101    if not ref_text.strip():102        gr.Info("No reference text provided, transcribing reference audio...")103        ref_text = outputs = pipe(104            ref_audio,105            chunk_length_s=30,106            batch_size=128,107            generate_kwargs={"task": "transcribe"},108            return_timestamps=False,109        )['text'].strip()110        gr.Info("Finished transcription")111    else:112        gr.Info("Using custom reference text...")113    audio, sr = torchaudio.load(ref_audio)114    # Audio115    if audio.shape[0] > 1:116        audio = torch.mean(audio, dim=0, keepdim=True)117    rms = torch.sqrt(torch.mean(torch.square(audio)))118    if rms < target_rms:119        audio = audio * target_rms / rms120    if sr != target_sample_rate:121        resampler = torchaudio.transforms.Resample(sr, target_sample_rate)122        audio = resampler(audio)123    audio = audio.to(device)124    # Chunk125    chunks = txtsplit(gen_text, 100, 150) # 100 chars preferred, 150 max126    results = []127    generated_mel_specs = []128    for chunk in progress.tqdm(chunks):129        # Prepare the text130        text_list = [ref_text + chunk]131        final_text_list = convert_char_to_pinyin(text_list)132    133        # Calculate duration134        ref_audio_len = audio.shape[-1] // hop_length135        # if fix_duration is not None:136        #     duration = int(fix_duration * target_sample_rate / hop_length)137        # else:138        zh_pause_punc = r"。，、；：？！"139        ref_text_len = len(ref_text) + len(re.findall(zh_pause_punc, ref_text))140        gen_text_len = len(gen_text) + len(re.findall(zh_pause_punc, gen_text))141        duration = ref_audio_len + int(ref_audio_len / ref_text_len * gen_text_len / speed)142    143        # inference144        gr.Info(f"Generating audio using {exp_name}")145        with torch.inference_mode():146            generated, _ = base_model.sample(147                cond=audio,148                text=final_text_list,149                duration=duration,150                steps=nfe_step,151                cfg_strength=cfg_strength,152                sway_sampling_coef=sway_sampling_coef,153            )154    155        generated = generated[:, ref_audio_len:, :]156        generated_mel_spec = rearrange(generated, '1 n d -> 1 d n')157        gr.Info("Running vocoder")158        vocos = Vocos.from_pretrained("charactr/vocos-mel-24khz")159        generated_wave = vocos.decode(generated_mel_spec.cpu())160        if rms < target_rms:161            generated_wave = generated_wave * rms / target_rms162    163        # wav -> numpy164        generated_wave = generated_wave.squeeze().cpu().numpy()165        results.append(generated_wave)166    generated_wave = np.concatenate(results)167    if remove_silence:168        gr.Info("Removing audio silences... This may take a moment")169        non_silent_intervals = librosa.effects.split(generated_wave, top_db=30)170        non_silent_wave = np.array([])171        for interval in non_silent_intervals:172            start, end = interval173            non_silent_wave = np.concatenate([non_silent_wave, generated_wave[start:end]])174        generated_wave = non_silent_wave175 176 177    # spectogram178    # with tempfile.NamedTemporaryFile(suffix=".png", delete=False) as tmp_spectrogram:179    #     spectrogram_path = tmp_spectrogram.name180    #     save_spectrogram(generated_mel_spec[0].cpu().numpy(), spectrogram_path)181 182    return (target_sample_rate, generated_wave)183 184with gr.Blocks() as app:185    gr.Markdown("""186# E2/F5 TTS187 188This is an unofficial E2/F5 TTS demo. This demo supports the following TTS models:189 190* [E2-TTS](https://arxiv.org/abs/2406.18009) (Embarrassingly Easy Fully Non-Autoregressive Zero-Shot TTS)191* [F5-TTS](https://arxiv.org/abs/2410.06885) (A Fairytaler that Fakes Fluent and Faithful Speech with Flow Matching)192 193This demo is based on the [F5-TTS](https://github.com/SWivid/F5-TTS) codebase, which is based on an [unofficial E2-TTS implementation](https://github.com/lucidrains/e2-tts-pytorch).194 195The checkpoints support English and Chinese.196 197If you're having issues, try converting your reference audio to WAV or MP3, clipping it to 15s, and shortening your prompt. If you're still running into issues, please open a [community Discussion](https://huggingface.co/spaces/mrfakename/E2-F5-TTS/discussions).198 199Long-form/batched inference + speech editing is coming soon!200 201**NOTE: Reference text will be automatically transcribed with Whisper if not provided. For best results, keep your reference clips short (<15s). Ensure the audio is fully uploaded before generating.**202""")203 204    ref_audio_input = gr.Audio(label="Reference Audio", type="filepath")205    gen_text_input = gr.Textbox(label="Text to Generate (longer text will use chunking)", lines=4)206    model_choice = gr.Radio(choices=["F5-TTS", "E2-TTS"], label="Choose TTS Model", value="F5-TTS")207    generate_btn = gr.Button("Synthesize", variant="primary")208    with gr.Accordion("Advanced Settings", open=False):209        ref_text_input = gr.Textbox(label="Reference Text", info="Leave blank to automatically transcribe the reference audio. If you enter text it will override automatic transcription.", lines=2)210        remove_silence = gr.Checkbox(label="Remove Silences", info="The model tends to produce silences, especially on longer audio. We can manually remove silences if needed. Note that this is an experimental feature and may produce strange results. This will also increase generation time.", value=True)211    212    audio_output = gr.Audio(label="Synthesized Audio")213    # spectrogram_output = gr.Image(label="Spectrogram")214 215    generate_btn.click(infer, inputs=[ref_audio_input, ref_text_input, gen_text_input, model_choice, remove_silence], outputs=[audio_output])216    gr.Markdown("Unofficial demo by [mrfakename](https://x.com/realmrfakename)")217 218 219app.queue().launch()