Sync from GitHub repo

This Space is synced from the GitHub repo: https://github.com/SWivid/F5-TTS. Please submit contributions to the Space there

app.py

@@ -1,22 +1,25 @@
 import re
-import torch
-import torchaudio
+import tempfile
+
+import click
 import gradio as gr
 import numpy as np
-import tempfile
-from vocos import Vocos
-from pydub import AudioSegment, silence
-from model import CFM, UNetT, DiT, MMDiT
+import soundfile as sf
+import torchaudio
 from cached_path import cached_path
+from pydub import AudioSegment
+
+from model import DiT, UNetT
 from model.utils import (
-    load_checkpoint,
-    get_tokenizer,
-    convert_char_to_pinyin,
     save_spectrogram,
 )
-from transformers import pipeline
-import click
-import soundfile as sf
+from model.utils_infer import (
+    load_vocoder,
+    load_model,
+    preprocess_ref_audio_text,
+    infer_process,
+    remove_silence_for_generated_wav,
+)
 
 try:
     import spaces
@@ -30,282 +33,47 @@ def gpu_decorator(func):
     else:
         return func
 
-device = (
-    "cuda"
-    if torch.cuda.is_available()
-    else "mps" if torch.backends.mps.is_available() else "cpu"
-)
-
-print(f"Using {device} device")
-
-pipe = pipeline(
-    "automatic-speech-recognition",
-    model="openai/whisper-large-v3-turbo",
-    torch_dtype=torch.float16,
-    device=device,
-)
-vocos = Vocos.from_pretrained("charactr/vocos-mel-24khz")
-
-# --------------------- Settings -------------------- #
-
-target_sample_rate = 24000
-n_mel_channels = 100
-hop_length = 256
-target_rms = 0.1
-nfe_step = 32  # 16, 32
-cfg_strength = 2.0
-ode_method = "euler"
-sway_sampling_coef = -1.0
-speed = 1.0
-fix_duration = None
-
-
-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"))
-    # ckpt_path = f"ckpts/{exp_name}/model_{ckpt_step}.pt"  # .pt | .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
+vocos = load_vocoder()
 
 
 # load models
-F5TTS_model_cfg = dict(
-    dim=1024, depth=22, heads=16, ff_mult=2, text_dim=512, conv_layers=4
-)
-E2TTS_model_cfg = dict(dim=1024, depth=24, heads=16, ff_mult=4)
-
-F5TTS_ema_model = load_model(
-    "F5-TTS", "F5TTS_Base", DiT, F5TTS_model_cfg, 1200000
-)
-E2TTS_ema_model = load_model(
-    "E2-TTS", "E2TTS_Base", UNetT, E2TTS_model_cfg, 1200000
-)
-
-def chunk_text(text, max_chars=135):
-    """
-    Splits the input text into chunks, each with a maximum number of characters.
-
-    Args:
-        text (str): The text to be split.
-        max_chars (int): The maximum number of characters per chunk.
-
-    Returns:
-        List[str]: A list of text chunks.
-    """
-    chunks = []
-    current_chunk = ""
-    # Split the text into sentences based on punctuation followed by whitespace
-    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
+F5TTS_model_cfg = dict(dim=1024, depth=22, heads=16, ff_mult=2, text_dim=512, conv_layers=4)
+F5TTS_ema_model = load_model(DiT, F5TTS_model_cfg, str(cached_path(f"hf://SWivid/F5-TTS/F5TTS_Base/model_1200000.safetensors")))
 
-    if current_chunk:
-        chunks.append(current_chunk.strip())
+E2TTS_model_cfg = dict(dim=1024, depth=24, heads=16, ff_mult=4)
+E2TTS_ema_model = load_model(UNetT, E2TTS_model_cfg, str(cached_path(f"hf://SWivid/E2-TTS/E2TTS_Base/model_1200000.safetensors")))
 
-    return chunks
 
 @gpu_decorator
-def infer_batch(ref_audio, ref_text, gen_text_batches, exp_name, remove_silence, cross_fade_duration=0.15, progress=gr.Progress()):
-    if exp_name == "F5-TTS":
+def infer(ref_audio_orig, ref_text, gen_text, model, remove_silence, cross_fade_duration=0.15, speed=1):
+    
+    ref_audio, ref_text = preprocess_ref_audio_text(ref_audio_orig, ref_text, show_info=gr.Info)
+
+    if model == "F5-TTS":
         ema_model = F5TTS_ema_model
-    elif exp_name == "E2-TTS":
+    elif model == "E2-TTS":
         ema_model = E2TTS_ema_model
 
-    audio, sr = ref_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 = []
-    spectrograms = []
-
-    if len(ref_text[-1].encode('utf-8')) == 1:
-        ref_text = ref_text + " "
-    for i, gen_text in enumerate(progress.tqdm(gen_text_batches)):
-        # Prepare the text
-        text_list = [ref_text + gen_text]
-        final_text_list = convert_char_to_pinyin(text_list)
-
-        # Calculate duration
-        ref_audio_len = audio.shape[-1] // hop_length
-        ref_text_len = len(ref_text.encode('utf-8'))
-        gen_text_len = len(gen_text.encode('utf-8'))
-        duration = ref_audio_len + int(ref_audio_len / ref_text_len * gen_text_len / speed)
-
-        # inference
-        with torch.inference_mode():
-            generated, _ = ema_model.sample(
-                cond=audio,
-                text=final_text_list,
-                duration=duration,
-                steps=nfe_step,
-                cfg_strength=cfg_strength,
-                sway_sampling_coef=sway_sampling_coef,
-            )
-
-        generated = generated.to(torch.float32)
-        generated = generated[:, ref_audio_len:, :]
-        generated_mel_spec = generated.permute(0, 2, 1)
-        generated_wave = vocos.decode(generated_mel_spec.cpu())
-        if rms < target_rms:
-            generated_wave = generated_wave * rms / target_rms
-
-        # wav -> numpy
-        generated_wave = generated_wave.squeeze().cpu().numpy()
-        
-        generated_waves.append(generated_wave)
-        spectrograms.append(generated_mel_spec[0].cpu().numpy())
-
-    # Combine all generated waves with cross-fading
-    if cross_fade_duration <= 0:
-        # Simply concatenate
-        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]
-
-            # Calculate cross-fade samples, ensuring it does not exceed wave lengths
-            cross_fade_samples = int(cross_fade_duration * target_sample_rate)
-            cross_fade_samples = min(cross_fade_samples, len(prev_wave), len(next_wave))
-
-            if cross_fade_samples <= 0:
-                # No overlap possible, concatenate
-                final_wave = np.concatenate([prev_wave, next_wave])
-                continue
-
-            # Overlapping parts
-            prev_overlap = prev_wave[-cross_fade_samples:]
-            next_overlap = next_wave[:cross_fade_samples]
-
-            # Fade out and fade in
-            fade_out = np.linspace(1, 0, cross_fade_samples)
-            fade_in = np.linspace(0, 1, cross_fade_samples)
-
-            # Cross-faded overlap
-            cross_faded_overlap = prev_overlap * fade_out + next_overlap * fade_in
-
-            # Combine
-            new_wave = np.concatenate([
-                prev_wave[:-cross_fade_samples],
-                cross_faded_overlap,
-                next_wave[cross_fade_samples:]
-            ])
-
-            final_wave = new_wave
+    final_wave, final_sample_rate, combined_spectrogram = infer_process(ref_audio, ref_text, gen_text, ema_model, cross_fade_duration=cross_fade_duration, speed=speed, show_info=gr.Info, progress=gr.Progress())
 
     # Remove silence
     if remove_silence:
         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=1000, silence_thresh=-50, keep_silence=500)
-            non_silent_wave = AudioSegment.silent(duration=0)
-            for non_silent_seg in non_silent_segs:
-                non_silent_wave += non_silent_seg
-            aseg = non_silent_wave
-            aseg.export(f.name, format="wav")
+            sf.write(f.name, final_wave, final_sample_rate)
+            remove_silence_for_generated_wav(f.name)
             final_wave, _ = torchaudio.load(f.name)
         final_wave = final_wave.squeeze().cpu().numpy()
 
-    # Create a combined spectrogram
-    combined_spectrogram = np.concatenate(spectrograms, axis=1)
-    
+    # Save the spectrogram
     with tempfile.NamedTemporaryFile(suffix=".png", delete=False) as tmp_spectrogram:
         spectrogram_path = tmp_spectrogram.name
         save_spectrogram(combined_spectrogram, spectrogram_path)
 
-    return (target_sample_rate, final_wave), spectrogram_path
-
-@gpu_decorator
-def infer(ref_audio_orig, ref_text, gen_text, exp_name, remove_silence, cross_fade_duration=0.15):
-
-    print(gen_text)
-
-    gr.Info("Converting audio...")
-    with tempfile.NamedTemporaryFile(delete=False, suffix=".wav") as f:
-        aseg = AudioSegment.from_file(ref_audio_orig)
-
-        non_silent_segs = silence.split_on_silence(
-            aseg, min_silence_len=1000, silence_thresh=-50, keep_silence=1000
-        )
-        non_silent_wave = AudioSegment.silent(duration=0)
-        for non_silent_seg in non_silent_segs:
-            non_silent_wave += non_silent_seg
-        aseg = non_silent_wave
-
-        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
-
-    if not ref_text.strip():
-        gr.Info("No reference text provided, transcribing reference audio...")
-        ref_text = pipe(
-            ref_audio,
-            chunk_length_s=30,
-            batch_size=128,
-            generate_kwargs={"task": "transcribe"},
-            return_timestamps=False,
-        )["text"].strip()
-        gr.Info("Finished transcription")
-    else:
-        gr.Info("Using custom reference text...")
-
-    # Add the functionality to ensure it ends with ". "
-    if not ref_text.endswith(". "):
-        if ref_text.endswith("."):
-            ref_text += " "
-        else:
-            ref_text += ". "
-
-    audio, sr = torchaudio.load(ref_audio)
-
-    # Use the new chunk_text function to split gen_text
-    max_chars = int(len(ref_text.encode('utf-8')) / (audio.shape[-1] / sr) * (25 - audio.shape[-1] / sr))
-    gen_text_batches = chunk_text(gen_text, max_chars=max_chars)
-    print('ref_text', ref_text)
-    for i, batch_text in enumerate(gen_text_batches):
-        print(f'gen_text {i}', batch_text)
-    
-    gr.Info(f"Generating audio using {exp_name} in {len(gen_text_batches)} batches")
-    return infer_batch((audio, sr), ref_text, gen_text_batches, exp_name, remove_silence, cross_fade_duration)
+    return (final_sample_rate, final_wave), spectrogram_path
 
 
 @gpu_decorator
-def generate_podcast(script, speaker1_name, ref_audio1, ref_text1, speaker2_name, ref_audio2, ref_text2, exp_name, remove_silence):
+def generate_podcast(script, speaker1_name, ref_audio1, ref_text1, speaker2_name, ref_audio2, ref_text2, model, remove_silence):
     # Split the script into speaker blocks
     speaker_pattern = re.compile(f"^({re.escape(speaker1_name)}|{re.escape(speaker2_name)}):", re.MULTILINE)
     speaker_blocks = speaker_pattern.split(script)[1:]  # Skip the first empty element
@@ -327,7 +95,7 @@ def generate_podcast(script, speaker1_name, ref_audio1, ref_text1, speaker2_name
             continue  # Skip if the speaker is neither speaker1 nor speaker2
         
         # Generate audio for this block
-        audio, _ = infer(ref_audio, ref_text, text, exp_name, remove_silence)
+        audio, _ = infer(ref_audio, ref_text, text, model, remove_silence)
         
         # Convert the generated audio to a numpy array
         sr, audio_data = audio
@@ -377,10 +145,6 @@ def parse_speechtypes_text(gen_text):
 
     return segments
 
-def update_speed(new_speed):
-    global speed
-    speed = new_speed
-    return f"Speed set to: {speed}"
 
 with gr.Blocks() as app_credits:
     gr.Markdown("""
@@ -413,7 +177,7 @@ with gr.Blocks() as app_tts:
             label="Speed",
             minimum=0.3,
             maximum=2.0,
-            value=speed,
+            value=1.0,
             step=0.1,
             info="Adjust the speed of the audio.",
         )
@@ -425,7 +189,6 @@ with gr.Blocks() as app_tts:
             step=0.01,
             info="Set the duration of the cross-fade between audio clips.",
         )
-    speed_slider.change(update_speed, inputs=speed_slider)
 
     audio_output = gr.Audio(label="Synthesized Audio")
     spectrogram_output = gr.Image(label="Spectrogram")
@@ -439,6 +202,7 @@ with gr.Blocks() as app_tts:
             model_choice,
             remove_silence,
             cross_fade_duration_slider,
+            speed_slider,
         ],
         outputs=[audio_output, spectrogram_output],
     )

inference-cli.py

@@ -1,23 +1,22 @@
 import argparse
 import codecs
 import re
-import tempfile
 from pathlib import Path
 
 import numpy as np
 import soundfile as sf
 import tomli
-import torch
-import torchaudio
-import tqdm
 from cached_path import cached_path
-from pydub import AudioSegment, silence
-from transformers import pipeline
-from vocos import Vocos
 
-from model import CFM, DiT, MMDiT, UNetT
-from model.utils import (convert_char_to_pinyin, get_tokenizer,
-                         load_checkpoint, save_spectrogram)
+from model import DiT, UNetT
+from model.utils_infer import (
+    load_vocoder,
+    load_model,
+    preprocess_ref_audio_text,
+    infer_process,
+    remove_silence_for_generated_wav,
+)
+
 
 parser = argparse.ArgumentParser(
     prog="python3 inference-cli.py",
@@ -104,282 +103,35 @@ wave_path = Path(output_dir)/"out.wav"
 spectrogram_path = Path(output_dir)/"out.png"
 vocos_local_path = "../checkpoints/charactr/vocos-mel-24khz"
 
-device = (
-    "cuda"
-    if torch.cuda.is_available()
-    else "mps" if torch.backends.mps.is_available() else "cpu"
-)
-
-if args.load_vocoder_from_local:
-    print(f"Load vocos from local path {vocos_local_path}")
-    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:
-    print("Download Vocos from huggingface charactr/vocos-mel-24khz")
-    vocos = Vocos.from_pretrained("charactr/vocos-mel-24khz")
-
-print(f"Using {device} device")
-
-# --------------------- Settings -------------------- #
-
-target_sample_rate = 24000
-n_mel_channels = 100
-hop_length = 256
-target_rms = 0.1
-nfe_step = 32  # 16, 32
-cfg_strength = 2.0
-ode_method = "euler"
-sway_sampling_coef = -1.0
-speed = 1.0
-# fix_duration = 27  # None or float (duration in seconds)
-fix_duration = None
+vocos = load_vocoder(is_local=args.load_vocoder_from_local, local_path=vocos_local_path)
 
-def load_model(model_cls, model_cfg, ckpt_path,file_vocab):
-    
-    if file_vocab=="":
-        file_vocab="Emilia_ZH_EN"
-        tokenizer="pinyin"
-    else:
-        tokenizer="custom"
-
-    print("\nvocab : ", vocab_file,tokenizer) 
-    print("tokenizer : ", tokenizer) 
-    print("model : ", ckpt_path,"\n")    
-
-    vocab_char_map, vocab_size = get_tokenizer(file_vocab, tokenizer)
-    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
 
 # load models
-F5TTS_model_cfg = dict(
-    dim=1024, depth=22, heads=16, ff_mult=2, text_dim=512, conv_layers=4
-)
-E2TTS_model_cfg = dict(dim=1024, depth=24, heads=16, ff_mult=4)
-
 if model == "F5-TTS":
-
+    model_cls = DiT
+    model_cfg = dict(dim=1024, depth=22, heads=16, ff_mult=2, text_dim=512, conv_layers=4)
     if ckpt_file == "": 
-       repo_name= "F5-TTS"
-       exp_name = "F5TTS_Base"
-       ckpt_step= 1200000
-       ckpt_file = str(cached_path(f"hf://SWivid/{repo_name}/{exp_name}/model_{ckpt_step}.safetensors"))
-
-    ema_model = load_model(DiT, F5TTS_model_cfg, ckpt_file,vocab_file)
+        repo_name= "F5-TTS"
+        exp_name = "F5TTS_Base"
+        ckpt_step= 1200000
+        ckpt_file = str(cached_path(f"hf://SWivid/{repo_name}/{exp_name}/model_{ckpt_step}.safetensors"))
+        # ckpt_path = f"ckpts/{exp_name}/model_{ckpt_step}.pt"  # .pt | .safetensors; local path
 
 elif model == "E2-TTS":
+    model_cls = UNetT
+    model_cfg = dict(dim=1024, depth=24, heads=16, ff_mult=4)
     if ckpt_file == "": 
-       repo_name= "E2-TTS"
-       exp_name = "E2TTS_Base"
-       ckpt_step= 1200000
-       ckpt_file = str(cached_path(f"hf://SWivid/{repo_name}/{exp_name}/model_{ckpt_step}.safetensors"))
-    
-    ema_model = load_model(UNetT, E2TTS_model_cfg, ckpt_file,vocab_file)
-
-asr_pipe = pipeline(
-    "automatic-speech-recognition",
-    model="openai/whisper-large-v3-turbo",
-    torch_dtype=torch.float16,
-    device=device,
-)
-
-def chunk_text(text, max_chars=135):
-    """
-    Splits the input text into chunks, each with a maximum number of characters.
-    Args:
-        text (str): The text to be split.
-        max_chars (int): The maximum number of characters per chunk.
-    Returns:
-        List[str]: A list of text chunks.
-    """
-    chunks = []
-    current_chunk = ""
-    # Split the text into sentences based on punctuation followed by whitespace
-    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())
-
-    return chunks
-
-    #ckpt_path = f"ckpts/{exp_name}/model_{ckpt_step}.pt" # .pt | .safetensors
-    #if not Path(ckpt_path).exists():
-        #ckpt_path = str(cached_path(f"hf://SWivid/{repo_name}/{exp_name}/model_{ckpt_step}.safetensors"))
-
-def infer_batch(ref_audio, ref_text, gen_text_batches, model, remove_silence, cross_fade_duration=0.15):
-    audio, sr = ref_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 = []
-    spectrograms = []
-
-    if len(ref_text[-1].encode('utf-8')) == 1:
-        ref_text = ref_text + " "
-    for i, gen_text in enumerate(tqdm.tqdm(gen_text_batches)):
-        # Prepare the text
-        text_list = [ref_text + gen_text]
-        final_text_list = convert_char_to_pinyin(text_list)
-
-        # Calculate duration
-        ref_audio_len = audio.shape[-1] // hop_length
-        ref_text_len = len(ref_text.encode('utf-8'))
-        gen_text_len = len(gen_text.encode('utf-8'))
-        duration = ref_audio_len + int(ref_audio_len / ref_text_len * gen_text_len / speed)
-
-        # inference
-        with torch.inference_mode():
-            generated, _ = ema_model.sample(
-                cond=audio,
-                text=final_text_list,
-                duration=duration,
-                steps=nfe_step,
-                cfg_strength=cfg_strength,
-                sway_sampling_coef=sway_sampling_coef,
-            )
-
-        generated = generated.to(torch.float32)
-        generated = generated[:, ref_audio_len:, :]
-        generated_mel_spec = generated.permute(0, 2, 1)
-        generated_wave = vocos.decode(generated_mel_spec.cpu())
-        if rms < target_rms:
-            generated_wave = generated_wave * rms / target_rms
-
-        # wav -> numpy
-        generated_wave = generated_wave.squeeze().cpu().numpy()
-        
-        generated_waves.append(generated_wave)
-        spectrograms.append(generated_mel_spec[0].cpu().numpy())
-
-    # Combine all generated waves with cross-fading
-    if cross_fade_duration <= 0:
-        # Simply concatenate
-        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]
-
-            # Calculate cross-fade samples, ensuring it does not exceed wave lengths
-            cross_fade_samples = int(cross_fade_duration * target_sample_rate)
-            cross_fade_samples = min(cross_fade_samples, len(prev_wave), len(next_wave))
-
-            if cross_fade_samples <= 0:
-                # No overlap possible, concatenate
-                final_wave = np.concatenate([prev_wave, next_wave])
-                continue
-
-            # Overlapping parts
-            prev_overlap = prev_wave[-cross_fade_samples:]
-            next_overlap = next_wave[:cross_fade_samples]
-
-            # Fade out and fade in
-            fade_out = np.linspace(1, 0, cross_fade_samples)
-            fade_in = np.linspace(0, 1, cross_fade_samples)
-
-            # Cross-faded overlap
-            cross_faded_overlap = prev_overlap * fade_out + next_overlap * fade_in
-
-            # Combine
-            new_wave = np.concatenate([
-                prev_wave[:-cross_fade_samples],
-                cross_faded_overlap,
-                next_wave[cross_fade_samples:]
-            ])
-
-            final_wave = new_wave
-
-    # Create a combined spectrogram
-    combined_spectrogram = np.concatenate(spectrograms, axis=1)
-
-    return final_wave, combined_spectrogram
-
-def process_voice(ref_audio_orig, ref_text):
-    print("Converting", ref_audio_orig)
-    with tempfile.NamedTemporaryFile(delete=False, suffix=".wav") as f:
-        aseg = AudioSegment.from_file(ref_audio_orig)
-
-        non_silent_segs = silence.split_on_silence(aseg, min_silence_len=1000, silence_thresh=-50, keep_silence=1000)
-        non_silent_wave = AudioSegment.silent(duration=0)
-        for non_silent_seg in non_silent_segs:
-            non_silent_wave += non_silent_seg
-        aseg = non_silent_wave
-
-        audio_duration = len(aseg)
-        if audio_duration > 15000:
-            print("Audio is over 15s, clipping to only first 15s.")
-            aseg = aseg[:15000]
-        aseg.export(f.name, format="wav")
-        ref_audio = f.name
-
-    if not ref_text.strip():
-        print("No reference text provided, transcribing reference audio...")
-        ref_text = asr_pipe(
-            ref_audio,
-            chunk_length_s=30,
-            batch_size=128,
-            generate_kwargs={"task": "transcribe"},
-            return_timestamps=False,
-        )["text"].strip()
-        print("Finished transcription")
-    else:
-        print("Using custom reference text...")
-    return ref_audio, ref_text    
-
-def infer(ref_audio, ref_text, gen_text, model, remove_silence, cross_fade_duration=0.15):
-    # Add the functionality to ensure it ends with ". "
-    if not ref_text.endswith(". ") and not ref_text.endswith("。"):
-        if ref_text.endswith("."):
-            ref_text += " "
-        else:
-            ref_text += ". "
-
-    # Split the input text into batches
-    audio, sr = torchaudio.load(ref_audio)
-    max_chars = int(len(ref_text.encode('utf-8')) / (audio.shape[-1] / sr) * (25 - audio.shape[-1] / sr))
-    gen_text_batches = chunk_text(gen_text, max_chars=max_chars)
-    for i, gen_text in enumerate(gen_text_batches):
-        print(f'gen_text {i}', gen_text)
-    
-    print(f"Generating audio using {model} in {len(gen_text_batches)} batches, loading models...")
-    return infer_batch((audio, sr), ref_text, gen_text_batches, model, remove_silence, cross_fade_duration)
+        repo_name= "E2-TTS"
+        exp_name = "E2TTS_Base"
+        ckpt_step= 1200000
+        ckpt_file = str(cached_path(f"hf://SWivid/{repo_name}/{exp_name}/model_{ckpt_step}.safetensors"))
+        # ckpt_path = f"ckpts/{exp_name}/model_{ckpt_step}.pt"  # .pt | .safetensors; local path
+
+print(f"Using {model}...")
+ema_model = load_model(model_cls, model_cfg, ckpt_file, vocab_file)
     
 
-def process(ref_audio, ref_text, text_gen, model, remove_silence):
+def main_process(ref_audio, ref_text, text_gen, model_obj, remove_silence):
     main_voice = {"ref_audio":ref_audio, "ref_text":ref_text}
     if "voices" not in config:
         voices = {"main": main_voice}
@@ -387,7 +139,7 @@ def process(ref_audio, ref_text, text_gen, model, remove_silence):
         voices = config["voices"]
         voices["main"] = main_voice
     for voice in voices:
-        voices[voice]['ref_audio'], voices[voice]['ref_text'] = process_voice(voices[voice]['ref_audio'], voices[voice]['ref_text'])
+        voices[voice]['ref_audio'], voices[voice]['ref_text'] = preprocess_ref_audio_text(voices[voice]['ref_audio'], voices[voice]['ref_text'])
         print("Voice:", voice)
         print("Ref_audio:", voices[voice]['ref_audio'])
         print("Ref_text:", voices[voice]['ref_text'])
@@ -407,23 +159,17 @@ def process(ref_audio, ref_text, text_gen, model, remove_silence):
         ref_audio = voices[voice]['ref_audio']
         ref_text = voices[voice]['ref_text']
         print(f"Voice: {voice}")
-        audio, spectragram = infer(ref_audio, ref_text, gen_text, model,remove_silence)
+        audio, final_sample_rate, spectragram = infer_process(ref_audio, ref_text, gen_text, model_obj)
         generated_audio_segments.append(audio)
 
     if generated_audio_segments:
         final_wave = np.concatenate(generated_audio_segments)
         with open(wave_path, "wb") as f:
-            sf.write(f.name, final_wave, target_sample_rate)
+            sf.write(f.name, final_wave, final_sample_rate)
             # Remove silence
             if remove_silence:
-                aseg = AudioSegment.from_file(f.name)
-                non_silent_segs = silence.split_on_silence(aseg, min_silence_len=1000, silence_thresh=-50, keep_silence=500)
-                non_silent_wave = AudioSegment.silent(duration=0)
-                for non_silent_seg in non_silent_segs:
-                    non_silent_wave += non_silent_seg
-                aseg = non_silent_wave
-                aseg.export(f.name, format="wav")
+                remove_silence_for_generated_wav(f.name)
             print(f.name)
 
 
-process(ref_audio, ref_text, gen_text, model, remove_silence)
+main_process(ref_audio, ref_text, gen_text, ema_model, remove_silence)

model/utils_infer.py

@@ -0,0 +1,306 @@
+# A unified script for inference process
+# Make adjustments inside functions, and consider both gradio and cli scripts if need to change func output format
+
+import re
+import tempfile
+
+import numpy as np
+import torch
+import torchaudio
+import tqdm
+from pydub import AudioSegment, silence
+from transformers import pipeline
+from vocos import Vocos
+
+from model import CFM
+from model.utils import (
+    load_checkpoint,
+    get_tokenizer,
+    convert_char_to_pinyin,
+)
+
+device = (
+    "cuda"
+    if torch.cuda.is_available()
+    else "mps" if torch.backends.mps.is_available() else "cpu"
+)
+print(f"Using {device} device")
+
+asr_pipe = pipeline(
+    "automatic-speech-recognition",
+    model="openai/whisper-large-v3-turbo",
+    torch_dtype=torch.float16,
+    device=device,
+)
+
+vocos = Vocos.from_pretrained("charactr/vocos-mel-24khz")
+
+
+# -----------------------------------------
+
+target_sample_rate = 24000
+n_mel_channels = 100
+hop_length = 256
+target_rms = 0.1
+nfe_step = 32  # 16, 32
+cfg_strength = 2.0
+ode_method = "euler"
+sway_sampling_coef = -1.0
+speed = 1.0
+fix_duration = None
+
+# -----------------------------------------
+
+
+# chunk text into smaller pieces
+
+def chunk_text(text, max_chars=135):
+    """
+    Splits the input text into chunks, each with a maximum number of characters.
+
+    Args:
+        text (str): The text to be split.
+        max_chars (int): The maximum number of characters per chunk.
+
+    Returns:
+        List[str]: A list of text chunks.
+    """
+    chunks = []
+    current_chunk = ""
+    # Split the text into sentences based on punctuation followed by whitespace
+    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())
+
+    return chunks
+
+
+# load vocoder
+
+def load_vocoder(is_local=False, local_path=""):
+    if is_local:
+        print(f"Load vocos from local path {local_path}")
+        vocos = Vocos.from_hparams(f"{local_path}/config.yaml")
+        state_dict = torch.load(f"{local_path}/pytorch_model.bin", map_location=device)
+        vocos.load_state_dict(state_dict)
+        vocos.eval()
+    else:
+        print("Download Vocos from huggingface charactr/vocos-mel-24khz")
+        vocos = Vocos.from_pretrained("charactr/vocos-mel-24khz")
+    return vocos
+
+
+# load model for inference
+
+def load_model(model_cls, model_cfg, ckpt_path, vocab_file=""):
+    
+    if vocab_file == "":
+        vocab_file = "Emilia_ZH_EN"
+        tokenizer = "pinyin"
+    else:
+        tokenizer = "custom"
+
+    print("\nvocab : ", vocab_file, tokenizer) 
+    print("tokenizer : ", tokenizer) 
+    print("model : ", ckpt_path,"\n")    
+
+    vocab_char_map, vocab_size = get_tokenizer(vocab_file, tokenizer)
+    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
+
+
+# preprocess reference audio and text
+
+def preprocess_ref_audio_text(ref_audio_orig, ref_text, show_info=print):
+    show_info("Converting audio...")
+    with tempfile.NamedTemporaryFile(delete=False, suffix=".wav") as f:
+        aseg = AudioSegment.from_file(ref_audio_orig)
+
+        non_silent_segs = silence.split_on_silence(
+            aseg, min_silence_len=1000, silence_thresh=-50, keep_silence=1000
+        )
+        non_silent_wave = AudioSegment.silent(duration=0)
+        for non_silent_seg in non_silent_segs:
+            non_silent_wave += non_silent_seg
+        aseg = non_silent_wave
+
+        audio_duration = len(aseg)
+        if audio_duration > 15000:
+            show_info("Audio is over 15s, clipping to only first 15s.")
+            aseg = aseg[:15000]
+        aseg.export(f.name, format="wav")
+        ref_audio = f.name
+
+    if not ref_text.strip():
+        show_info("No reference text provided, transcribing reference audio...")
+        ref_text = asr_pipe(
+            ref_audio,
+            chunk_length_s=30,
+            batch_size=128,
+            generate_kwargs={"task": "transcribe"},
+            return_timestamps=False,
+        )["text"].strip()
+        show_info("Finished transcription")
+    else:
+        show_info("Using custom reference text...")
+
+    # Add the functionality to ensure it ends with ". "
+    if not ref_text.endswith(". ") and not ref_text.endswith("。"):
+        if ref_text.endswith("."):
+            ref_text += " "
+        else:
+            ref_text += ". "
+
+    return ref_audio, ref_text
+
+
+# infer process: chunk text -> infer batches [i.e. infer_batch_process()]
+
+def infer_process(ref_audio, ref_text, gen_text, model_obj, cross_fade_duration=0.15, speed=speed, show_info=print, progress=tqdm):
+
+    # Split the input text into batches
+    audio, sr = torchaudio.load(ref_audio)
+    max_chars = int(len(ref_text.encode('utf-8')) / (audio.shape[-1] / sr) * (25 - audio.shape[-1] / sr))
+    gen_text_batches = chunk_text(gen_text, max_chars=max_chars)
+    for i, gen_text in enumerate(gen_text_batches):
+        print(f'gen_text {i}', gen_text)
+    
+    show_info(f"Generating audio in {len(gen_text_batches)} batches...")
+    return infer_batch_process((audio, sr), ref_text, gen_text_batches, model_obj, cross_fade_duration, speed, progress)
+
+
+# infer batches
+
+def infer_batch_process(ref_audio, ref_text, gen_text_batches, model_obj, cross_fade_duration=0.15, speed=1, progress=tqdm):
+    audio, sr = ref_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 = []
+    spectrograms = []
+
+    if len(ref_text[-1].encode('utf-8')) == 1:
+        ref_text = ref_text + " "
+    for i, gen_text in enumerate(progress.tqdm(gen_text_batches)):
+        # Prepare the text
+        text_list = [ref_text + gen_text]
+        final_text_list = convert_char_to_pinyin(text_list)
+
+        # Calculate duration
+        ref_audio_len = audio.shape[-1] // hop_length
+        ref_text_len = len(ref_text.encode('utf-8'))
+        gen_text_len = len(gen_text.encode('utf-8'))
+        duration = ref_audio_len + int(ref_audio_len / ref_text_len * gen_text_len / speed)
+
+        # inference
+        with torch.inference_mode():
+            generated, _ = model_obj.sample(
+                cond=audio,
+                text=final_text_list,
+                duration=duration,
+                steps=nfe_step,
+                cfg_strength=cfg_strength,
+                sway_sampling_coef=sway_sampling_coef,
+            )
+
+        generated = generated.to(torch.float32)
+        generated = generated[:, ref_audio_len:, :]
+        generated_mel_spec = generated.permute(0, 2, 1)
+        generated_wave = vocos.decode(generated_mel_spec.cpu())
+        if rms < target_rms:
+            generated_wave = generated_wave * rms / target_rms
+
+        # wav -> numpy
+        generated_wave = generated_wave.squeeze().cpu().numpy()
+        
+        generated_waves.append(generated_wave)
+        spectrograms.append(generated_mel_spec[0].cpu().numpy())
+
+    # Combine all generated waves with cross-fading
+    if cross_fade_duration <= 0:
+        # Simply concatenate
+        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]
+
+            # Calculate cross-fade samples, ensuring it does not exceed wave lengths
+            cross_fade_samples = int(cross_fade_duration * target_sample_rate)
+            cross_fade_samples = min(cross_fade_samples, len(prev_wave), len(next_wave))
+
+            if cross_fade_samples <= 0:
+                # No overlap possible, concatenate
+                final_wave = np.concatenate([prev_wave, next_wave])
+                continue
+
+            # Overlapping parts
+            prev_overlap = prev_wave[-cross_fade_samples:]
+            next_overlap = next_wave[:cross_fade_samples]
+
+            # Fade out and fade in
+            fade_out = np.linspace(1, 0, cross_fade_samples)
+            fade_in = np.linspace(0, 1, cross_fade_samples)
+
+            # Cross-faded overlap
+            cross_faded_overlap = prev_overlap * fade_out + next_overlap * fade_in
+
+            # Combine
+            new_wave = np.concatenate([
+                prev_wave[:-cross_fade_samples],
+                cross_faded_overlap,
+                next_wave[cross_fade_samples:]
+            ])
+
+            final_wave = new_wave
+
+    # Create a combined spectrogram
+    combined_spectrogram = np.concatenate(spectrograms, axis=1)
+
+    return final_wave, target_sample_rate, combined_spectrogram
+
+
+# remove silence from generated wav
+
+def remove_silence_for_generated_wav(filename):
+    aseg = AudioSegment.from_file(filename)
+    non_silent_segs = silence.split_on_silence(aseg, min_silence_len=1000, silence_thresh=-50, keep_silence=500)
+    non_silent_wave = AudioSegment.silent(duration=0)
+    for non_silent_seg in non_silent_segs:
+        non_silent_wave += non_silent_seg
+    aseg = non_silent_wave
+    aseg.export(filename, format="wav")