app.py

# Imports
import gradio as gr
import spaces
import re
import torch
import torchaudio
import numpy as np
import tempfile
import click
import soundfile as sf

from einops import rearrange
from vocos import Vocos
from pydub import AudioSegment, silence
from model import CFM, UNetT, DiT, MMDiT
from cached_path import cached_path
from model.utils import (load_checkpoint, get_tokenizer, convert_char_to_pinyin, save_spectrogram)

# Pre-Initialize
DEVICE = "auto"
if DEVICE == "auto":
    DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
print(f"[SYSTEM] | Using {DEVICE} type compute device.")

target_sample_rate = 24000
n_mel_channels = 100
hop_length = 256
ode_method = "euler"

def load_model(repo_name, exp_name, model_cls, model_cfg, ckpt_step):
    ckpt_path = str(cached_path(f"hf://SWivid/{repo_name}/{exp_name}/model_{ckpt_step}.safetensors"))
    vocab_char_map, vocab_size = get_tokenizer("Emilia_ZH_EN", "pinyin")
    
    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)
    
    model = load_checkpoint(model, ckpt_path, DEVICE, use_ema = True)
    return model
    
# Variables
DEFAULT_MODEL = "F5"
DEFAULT_REMOVE_SILENCES = True

DEFAULT_STEPS = 32
DEFAULT_SPEED = 1
DEFAULT_CROSS_FADE = 0.15

target_rms = 0.1
cfg_strength = 2.0
sway_sampling_coef = -1.0

silence_offset = 25
silence_min_len = 500

vocos = Vocos.from_pretrained("charactr/vocos-mel-24khz")

F5TTS_model_cfg = dict(dim=1024, depth=22, heads=16, ff_mult=2, text_dim=512, conv_layers=4)
F5TTS_ema_model = load_model("F5-TTS", "F5TTS_Base", DiT, F5TTS_model_cfg, 1200000)

E2TTS_model_cfg = dict(dim=1024, depth=24, heads=16, ff_mult=4)
E2TTS_ema_model = load_model("E2-TTS", "E2TTS_Base", UNetT, E2TTS_model_cfg, 1200000)

css = '''
.gradio-container{max-width: 560px !important}
h1{text-align:center}
footer {
    visibility: hidden
}
'''

# Functions
@spaces.GPU(duration=30)
def infer_batch(input_batches, reference_audio, reference_input, model_choice=DEFAULT_MODEL, remove_silences=DEFAULT_REMOVE_SILENCES, steps=DEFAULT_STEPS, speed=DEFAULT_SPEED, cross_fade=DEFAULT_CROSS_FADE):
    if model_choice == "F5":
        ema_model = F5TTS_ema_model
    elif model_choice == "E2":
        ema_model = E2TTS_ema_model

    print("Waiting for inference...")

    audio, sr = reference_audio
    if audio.shape[0] > 1:
        audio = torch.mean(audio, dim=0, keepdim=True)

    rms = torch.sqrt(torch.mean(torch.square(audio)))
    if rms < target_rms:
        audio = audio * target_rms / rms
    if sr != target_sample_rate:
        resampler = torchaudio.transforms.Resample(sr, target_sample_rate)
        audio = resampler(audio)
        
    audio = audio.to(DEVICE)
    generated_waves = []

    if len(reference_input[-1].encode('utf-8')) == 1:
        reference_input = reference_input + " "

    print("Inferencing each batch...")
        
    for i, input in enumerate(input_batches):
        text_list = [reference_input + input]
        final_text_list = convert_char_to_pinyin(text_list)

        reference_audio_len = audio.shape[-1] // hop_length
        zh_pause_punc = r"。，、；：？！"
        reference_input_len = len(reference_input.encode('utf-8')) + 3 * len(re.findall(zh_pause_punc, reference_input))
        input_len = len(input.encode('utf-8')) + 3 * len(re.findall(zh_pause_punc, input))
        duration = reference_audio_len + int(reference_audio_len / reference_input_len * input_len / speed)

        # Inference
        with torch.inference_mode():
            generated, _ = ema_model.sample(cond=audio, text=final_text_list, duration=duration, steps=steps, cfg_strength=cfg_strength, sway_sampling_coef=sway_sampling_coef)

        generated = generated[:, reference_audio_len:, :]
        generated_mel_spec = rearrange(generated, "1 n d -> 1 d n")
        generated_wave = vocos.decode(generated_mel_spec.cpu())
        
        if rms < target_rms:
            generated_wave = generated_wave * rms / target_rms

        generated_wave = generated_wave.squeeze().cpu().numpy()
        generated_waves.append(generated_wave)

    # Handle combining generated waves with cross-fading
    print("Handling combining and cross-fading...")
    
    if cross_fade <= 0:
        final_wave = np.concatenate(generated_waves)
    else:
        final_wave = generated_waves[0]
        for i in range(1, len(generated_waves)):
            prev_wave = final_wave
            next_wave = generated_waves[i]

            cross_fade_samples = int(cross_fade * target_sample_rate)
            cross_fade_samples = min(cross_fade_samples, len(prev_wave), len(next_wave))

            if cross_fade_samples <= 0:
                final_wave = np.concatenate([prev_wave, next_wave])
                continue

            prev_overlap = prev_wave[-cross_fade_samples:]
            next_overlap = next_wave[:cross_fade_samples]

            fade_out = np.linspace(1, 0, cross_fade_samples)
            fade_in = np.linspace(0, 1, cross_fade_samples)

            cross_faded_overlap = prev_overlap * fade_out + next_overlap * fade_in

            new_wave = np.concatenate([prev_wave[:-cross_fade_samples], cross_faded_overlap, next_wave[cross_fade_samples:]])

            final_wave = new_wave

    # Handle removing silences
    print("Handling removing silences...")
    
    if remove_silences:
        with tempfile.NamedTemporaryFile(delete=False, suffix=".wav") as f:
            sf.write(f.name, final_wave, target_sample_rate)
            aseg = AudioSegment.from_file(f.name)
                
            non_silent_segs = silence.split_on_silence(aseg, min_silence_len=silence_min_len, silence_thresh=aseg.dBFS - silence_offset, keep_silence=250)
            
            non_silent_wave = AudioSegment.empty()
            for seg in non_silent_segs:
                non_silent_wave += seg
    
            aseg = non_silent_wave
            aseg.export(f.name, format="wav")
            final_wave, _ = torchaudio.load(f.name)
        final_wave = final_wave.squeeze().cpu().numpy()

    print("Done!")
    return (target_sample_rate, final_wave)

@spaces.GPU(duration=30)
def infer(input, reference_audio, reference_input, model_choice=DEFAULT_MODEL, remove_silences=DEFAULT_REMOVE_SILENCES, steps=DEFAULT_STEPS, speed=DEFAULT_SPEED, cross_fade=DEFAULT_CROSS_FADE):

    print("Modifying reference audio...")
    with tempfile.NamedTemporaryFile(delete=False, suffix=".wav") as f:
        aseg = AudioSegment.from_file(reference_audio)
        audio_duration = len(aseg)
    
        if audio_duration > 15000:
            gr.Warning("Audio is over 15s, clipping to only first 15s.")
            aseg = aseg[:15000]
    
        aseg.export(f.name, format="wav")
        ref_audio = f.name

    # Ensure it ends with period.
    print("Modifying reference input...")
    
    if not reference_input.endswith(". "):
        if reference_input.endswith("."):
            reference_input += " "
        else:
            reference_input += ". "

    print("Loading reference audio...")
    audio, sr = torchaudio.load(ref_audio)

    # Split input into chunks
    print("--------------------------------------------- INPUT")
    print(f"Input: {input}")
    print(f"Reference Input: {reference_input}")
    print(f"Parameters: {model_choice}, {remove_silences}, {steps}, {speed}, {cross_fade}")
    print("---------------------------------------------------")
    
    max_chars = int(len(reference_input.encode('utf-8')) / (audio.shape[-1] / sr) * (25 - audio.shape[-1] / sr))
    input_batches = chunk_text(input, max_chars=max_chars)

    print(f"------------------------------------------ BATCHES")
    for i, batch_text in enumerate(input_batches):
        print(f" {i}: ", batch_text)
        
    print("---------------------------------------------------")
    
    return infer_batch(input_batches, (audio, sr), reference_input, model_choice, remove_silences, steps, speed, cross_fade)

def chunk_text(text, max_chars=135):
    chunks = []
    current_chunk = ""
    
    # Split input into sentences with punctuations
    sentences = re.split(r'(?<=[;:,.!?])\s+|(?<=[；：，。！？])', text)

    for sentence in sentences:
        if len(current_chunk.encode('utf-8')) + len(sentence.encode('utf-8')) <= max_chars:
            current_chunk += sentence + " " if sentence and len(sentence[-1].encode('utf-8')) == 1 else sentence
        else:
            if current_chunk:
                chunks.append(current_chunk.strip())
            current_chunk = sentence + " " if sentence and len(sentence[-1].encode('utf-8')) == 1 else sentence

    if current_chunk:
        chunks.append(current_chunk.strip())

    print("-------------------------------------------- CHUNKS")
    print(chunks)
    print("---------------------------------------------------")
    
    return chunks

def cloud():
    print("[CLOUD] | Space maintained.")

# Initialize
with gr.Blocks(css=css) as main:
    with gr.Column():
        gr.Markdown("🪄 Speak text to audio.")

    with gr.Column():
        input = gr.Textbox(lines=1, value="", label="Input")
        reference_audio = gr.Audio(sources="upload", type="filepath", label="Reference Audio")
        reference_input = gr.Textbox(lines=1, value="", label="Reference Text")
        model_choice = gr.Radio(["F5", "E2"], label="TTS Model", value=DEFAULT_MODEL)
    
        remove_silences = gr.Checkbox(value=DEFAULT_REMOVE_SILENCES, label="Remove Silences")
        
        steps = gr.Slider(minimum=1, maximum=64, value=DEFAULT_STEPS, step=1, label="Steps")
        speed = gr.Slider(minimum=0.3, maximum=2.0, value=DEFAULT_SPEED, step=0.1, label="Speed")
        cross_fade = gr.Slider(minimum=0.0, maximum=1.0, value=DEFAULT_CROSS_FADE, step=0.01, label="Audio Cross-Fade Duration Between Sentences")
                
        submit = gr.Button("▶")
        maintain = gr.Button("☁️")

    with gr.Column():
        output = gr.Audio(label="Output")

    submit.click(infer, inputs=[input, reference_audio, reference_input, model_choice, remove_silences, steps, speed, cross_fade], outputs=output, queue=False)
    maintain.click(cloud, inputs=[], outputs=[], queue=False)

main.launch(show_api=True)