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

README.md

@@ -1,12 +1,12 @@
 ---
-title: E2/F5 TTS
+title: F5-TTS
 emoji: 🗣️
 colorFrom: green
 colorTo: green
 sdk: gradio
 app_file: app.py
 pinned: true
-short_description: 'E2-TTS & F5-TTS: Zero-Shot Voice Cloning (Unofficial Demo)'
+short_description: 'F5-TTS & E2-TTS: Zero-Shot Voice Cloning (Unofficial Demo)'
 sdk_version: 5.1.0
 ---
 

README_REPO.md

@@ -0,0 +1,196 @@
+# F5-TTS: A Fairytaler that Fakes Fluent and Faithful Speech with Flow Matching
+
+[![python](https://img.shields.io/badge/Python-3.10-brightgreen)](https://github.com/SWivid/F5-TTS)
+[![arXiv](https://img.shields.io/badge/arXiv-2410.06885-b31b1b.svg?logo=arXiv)](https://arxiv.org/abs/2410.06885)
+[![demo](https://img.shields.io/badge/GitHub-Demo%20page-blue.svg)](https://swivid.github.io/F5-TTS/)
+[![space](https://img.shields.io/badge/🤗-Space%20demo-yellow)](https://huggingface.co/spaces/mrfakename/E2-F5-TTS)
+
+**F5-TTS**: Diffusion Transformer with ConvNeXt V2, faster trained and inference.
+
+**E2 TTS**: Flat-UNet Transformer, closest reproduction.
+
+**Sway Sampling**: Inference-time flow step sampling strategy, greatly improves performance
+
+## Installation
+
+Clone the repository:
+
+```bash
+git clone https://github.com/SWivid/F5-TTS.git
+cd F5-TTS
+```
+
+Install torch with your CUDA version, e.g. :
+
+```bash
+pip install torch==2.3.0+cu118 --extra-index-url https://download.pytorch.org/whl/cu118
+pip install torchaudio==2.3.0+cu118 --extra-index-url https://download.pytorch.org/whl/cu118
+```
+
+Install other packages:
+
+```bash
+pip install -r requirements.txt
+```
+
+## Prepare Dataset
+
+Example data processing scripts for Emilia and Wenetspeech4TTS, and you may tailor your own one along with a Dataset class in `model/dataset.py`.
+
+```bash
+# prepare custom dataset up to your need
+# download corresponding dataset first, and fill in the path in scripts
+
+# Prepare the Emilia dataset
+python scripts/prepare_emilia.py
+
+# Prepare the Wenetspeech4TTS dataset
+python scripts/prepare_wenetspeech4tts.py
+```
+
+## Training
+
+Once your datasets are prepared, you can start the training process.
+
+```bash
+# setup accelerate config, e.g. use multi-gpu ddp, fp16
+# will be to: ~/.cache/huggingface/accelerate/default_config.yaml     
+accelerate config
+accelerate launch train.py
+```
+An initial guidance on Finetuning [#57](https://github.com/SWivid/F5-TTS/discussions/57).
+
+## Inference
+
+To run inference with pretrained models, download the checkpoints from [🤗 Hugging Face](https://huggingface.co/SWivid/F5-TTS), or automatically downloaded with `inference-cli` and `gradio_app`.
+
+Currently support 30s for a single generation, which is the **TOTAL** length of prompt audio and the generated. Batch inference with chunks is supported by `inference-cli` and `gradio_app`. 
+- To avoid possible inference failures, make sure you have seen through the following instructions.
+- A longer prompt audio allows shorter generated output. The part longer than 30s cannot be generated properly. Consider using a prompt audio <15s.
+- Uppercased letters will be uttered letter by letter, so use lowercased letters for normal words. 
+- Add some spaces (blank: " ") or punctuations (e.g. "," ".") to explicitly introduce some pauses. If first few words skipped in code-switched generation (cuz different speed with different languages), this might help.
+
+### CLI Inference
+
+Either you can specify everything in `inference-cli.toml` or override with flags. Leave `--ref_text ""` will have ASR model transcribe the reference audio automatically (use extra GPU memory). If encounter network error, consider use local ckpt, just set `ckpt_path` in `inference-cli.py`
+
+```bash
+python inference-cli.py \
+--model "F5-TTS" \
+--ref_audio "tests/ref_audio/test_en_1_ref_short.wav" \
+--ref_text "Some call me nature, others call me mother nature." \
+--gen_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."
+
+python inference-cli.py \
+--model "E2-TTS" \
+--ref_audio "tests/ref_audio/test_zh_1_ref_short.wav" \
+--ref_text "对，这就是我，万人敬仰的太乙真人。" \
+--gen_text "突然，身边一阵笑声。我看着他们，意气风发地挺直了胸膛，甩了甩那稍显肉感的双臂，轻笑道，我身上的肉，是为了掩饰我爆棚的魅力，否则，岂不吓坏了你们呢？"
+```
+
+### Gradio App
+Currently supported features:
+- Chunk inference
+- Podcast Generation
+- Multiple Speech-Type Generation
+
+You can launch a Gradio app (web interface) to launch a GUI for inference (will load ckpt from Huggingface, you may set `ckpt_path` to local file in `gradio_app.py`). Currently load ASR model, F5-TTS and E2 TTS all in once, thus use more GPU memory than `inference-cli`.
+
+```bash
+python gradio_app.py
+```
+
+You can specify the port/host:
+
+```bash
+python gradio_app.py --port 7860 --host 0.0.0.0
+```
+
+Or launch a share link:
+
+```bash
+python gradio_app.py --share
+```
+
+### Speech Editing
+
+To test speech editing capabilities, use the following command.
+
+```bash
+python speech_edit.py
+```
+
+## Evaluation
+
+### Prepare Test Datasets
+
+1. Seed-TTS test set: Download from [seed-tts-eval](https://github.com/BytedanceSpeech/seed-tts-eval).
+2. LibriSpeech test-clean: Download from [OpenSLR](http://www.openslr.org/12/).
+3. Unzip the downloaded datasets and place them in the data/ directory.
+4. Update the path for the test-clean data in `scripts/eval_infer_batch.py`
+5. Our filtered LibriSpeech-PC 4-10s subset is already under data/ in this repo
+
+### Batch Inference for Test Set
+
+To run batch inference for evaluations, execute the following commands:
+
+```bash
+# batch inference for evaluations
+accelerate config  # if not set before
+bash scripts/eval_infer_batch.sh
+```
+
+### Download Evaluation Model Checkpoints
+
+1. Chinese ASR Model: [Paraformer-zh](https://huggingface.co/funasr/paraformer-zh)
+2. English ASR Model: [Faster-Whisper](https://huggingface.co/Systran/faster-whisper-large-v3)
+3. WavLM Model: Download from [Google Drive](https://drive.google.com/file/d/1-aE1NfzpRCLxA4GUxX9ITI3F9LlbtEGP/view).
+
+### Objective Evaluation
+
+**Some Notes**
+
+For faster-whisper with CUDA 11:
+
+```bash
+pip install --force-reinstall ctranslate2==3.24.0
+```
+
+(Recommended) To avoid possible ASR failures, such as abnormal repetitions in output:
+
+```bash
+pip install faster-whisper==0.10.1
+```
+
+Update the path with your batch-inferenced results, and carry out WER / SIM evaluations:
+```bash
+# Evaluation for Seed-TTS test set
+python scripts/eval_seedtts_testset.py
+
+# Evaluation for LibriSpeech-PC test-clean (cross-sentence)
+python scripts/eval_librispeech_test_clean.py
+```
+
+## Acknowledgements
+
+- [E2-TTS](https://arxiv.org/abs/2406.18009) brilliant work, simple and effective
+- [Emilia](https://arxiv.org/abs/2407.05361), [WenetSpeech4TTS](https://arxiv.org/abs/2406.05763) valuable datasets
+- [lucidrains](https://github.com/lucidrains) initial CFM structure with also [bfs18](https://github.com/bfs18) for discussion
+- [SD3](https://arxiv.org/abs/2403.03206) & [Hugging Face diffusers](https://github.com/huggingface/diffusers) DiT and MMDiT code structure
+- [torchdiffeq](https://github.com/rtqichen/torchdiffeq) as ODE solver, [Vocos](https://huggingface.co/charactr/vocos-mel-24khz) as vocoder
+- [mrfakename](https://x.com/realmrfakename) huggingface space demo ~
+- [FunASR](https://github.com/modelscope/FunASR), [faster-whisper](https://github.com/SYSTRAN/faster-whisper), [UniSpeech](https://github.com/microsoft/UniSpeech) for evaluation tools
+- [ctc-forced-aligner](https://github.com/MahmoudAshraf97/ctc-forced-aligner) for speech edit test
+
+## Citation
+```
+@article{chen-etal-2024-f5tts,
+      title={F5-TTS: A Fairytaler that Fakes Fluent and Faithful Speech with Flow Matching}, 
+      author={Yushen Chen and Zhikang Niu and Ziyang Ma and Keqi Deng and Chunhui Wang and Jian Zhao and Kai Yu and Xie Chen},
+      journal={arXiv preprint arXiv:2410.06885},
+      year={2024},
+}
+```
+## License
+
+Our code is released under MIT License.

app.py

@@ -6,28 +6,53 @@ import gradio as gr
 import numpy as np
 import tempfile
 from einops import rearrange
-from ema_pytorch import EMA
 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 (
-    get_tokenizer, 
-    convert_char_to_pinyin, 
+    load_checkpoint,
+    get_tokenizer,
+    convert_char_to_pinyin,
     save_spectrogram,
 )
 from transformers import pipeline
-import spaces
 import librosa
+import click
 import soundfile as sf
-from txtsplit import txtsplit
-from detoxify import Detoxify
 
+try:
+    import spaces
+    USING_SPACES = True
+except ImportError:
+    USING_SPACES = False
 
-device = "cuda" if torch.cuda.is_available() else "mps" if torch.backends.mps.is_available() else "cpu"
+def gpu_decorator(func):
+    if USING_SPACES:
+        return spaces.GPU(func)
+    else:
+        return func
+
+
+
+SPLIT_WORDS = [
+    "but", "however", "nevertheless", "yet", "still",
+    "therefore", "thus", "hence", "consequently",
+    "moreover", "furthermore", "additionally",
+    "meanwhile", "alternatively", "otherwise",
+    "namely", "specifically", "for example", "such as",
+    "in fact", "indeed", "notably",
+    "in contrast", "on the other hand", "conversely",
+    "in conclusion", "to summarize", "finally"
+]
 
-model = Detoxify('original', device=device)
+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",
@@ -35,6 +60,7 @@ pipe = pipeline(
     torch_dtype=torch.float16,
     device=device,
 )
+vocos = Vocos.from_pretrained("charactr/vocos-mel-24khz")
 
 # --------------------- Settings -------------------- #
 
@@ -44,20 +70,20 @@ hop_length = 256
 target_rms = 0.1
 nfe_step = 32  # 16, 32
 cfg_strength = 2.0
-ode_method = 'euler'
+ode_method = "euler"
 sway_sampling_coef = -1.0
 speed = 1.0
 # fix_duration = 27  # None or float (duration in seconds)
 fix_duration = None
 
+
 def load_model(repo_name, exp_name, model_cls, model_cfg, ckpt_step):
-    checkpoint = torch.load(str(cached_path(f"hf://SWivid/{repo_name}/{exp_name}/model_{ckpt_step}.pt")), map_location=device)
+    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
+            **model_cfg, text_num_embeds=vocab_size, mel_dim=n_mel_channels
         ),
         mel_spec_kwargs=dict(
             target_sample_rate=target_sample_rate,
@@ -70,64 +96,130 @@ def load_model(repo_name, exp_name, model_cls, model_cfg, ckpt_step):
         vocab_char_map=vocab_char_map,
     ).to(device)
 
-    ema_model = EMA(model, include_online_model=False).to(device)
-    ema_model.load_state_dict(checkpoint['ema_model_state_dict'])
-    ema_model.copy_params_from_ema_to_model()
+    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)
+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)
+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
+)
 
-@spaces.GPU
-def infer(ref_audio_orig, ref_text, gen_text, exp_name, remove_silence, progress = gr.Progress()):
-    print(gen_text)
-    if model.predict(gen_text)['toxicity'] > 0.8:
-        print("Flagged for toxicity:", gen_text)
-        raise gr.Error("Your text was flagged for toxicity, please try again with a different text.")
-    gr.Info("Converting audio...")
-    with tempfile.NamedTemporaryFile(delete=False, suffix=".wav") as f:
-        aseg = AudioSegment.from_file(ref_audio_orig)
-        # remove long silence in reference audio
-        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
-        # Convert to mono
-        aseg = aseg.set_channels(1)
-        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
+def split_text_into_batches(text, max_chars=200, split_words=SPLIT_WORDS):
+    if len(text.encode('utf-8')) <= max_chars:
+        return [text]
+    if text[-1] not in ['。', '.', '!', '！', '?', '？']:
+        text += '.'
+        
+    sentences = re.split('([。.!?！？])', text)
+    sentences = [''.join(i) for i in zip(sentences[0::2], sentences[1::2])]
+    
+    batches = []
+    current_batch = ""
+    
+    def split_by_words(text):
+        words = text.split()
+        current_word_part = ""
+        word_batches = []
+        for word in words:
+            if len(current_word_part.encode('utf-8')) + len(word.encode('utf-8')) + 1 <= max_chars:
+                current_word_part += word + ' '
+            else:
+                if current_word_part:
+                    # Try to find a suitable split word
+                    for split_word in split_words:
+                        split_index = current_word_part.rfind(' ' + split_word + ' ')
+                        if split_index != -1:
+                            word_batches.append(current_word_part[:split_index].strip())
+                            current_word_part = current_word_part[split_index:].strip() + ' '
+                            break
+                    else:
+                        # If no suitable split word found, just append the current part
+                        word_batches.append(current_word_part.strip())
+                        current_word_part = ""
+                current_word_part += word + ' '
+        if current_word_part:
+            word_batches.append(current_word_part.strip())
+        return word_batches
+
+    for sentence in sentences:
+        if len(current_batch.encode('utf-8')) + len(sentence.encode('utf-8')) <= max_chars:
+            current_batch += sentence
+        else:
+            # If adding this sentence would exceed the limit
+            if current_batch:
+                batches.append(current_batch)
+                current_batch = ""
+            
+            # If the sentence itself is longer than max_chars, split it
+            if len(sentence.encode('utf-8')) > max_chars:
+                # First, try to split by colon
+                colon_parts = sentence.split(':')
+                if len(colon_parts) > 1:
+                    for part in colon_parts:
+                        if len(part.encode('utf-8')) <= max_chars:
+                            batches.append(part)
+                        else:
+                            # If colon part is still too long, split by comma
+                            comma_parts = re.split('[,，]', part)
+                            if len(comma_parts) > 1:
+                                current_comma_part = ""
+                                for comma_part in comma_parts:
+                                    if len(current_comma_part.encode('utf-8')) + len(comma_part.encode('utf-8')) <= max_chars:
+                                        current_comma_part += comma_part + ','
+                                    else:
+                                        if current_comma_part:
+                                            batches.append(current_comma_part.rstrip(','))
+                                        current_comma_part = comma_part + ','
+                                if current_comma_part:
+                                    batches.append(current_comma_part.rstrip(','))
+                            else:
+                                # If no comma, split by words
+                                batches.extend(split_by_words(part))
+                else:
+                    # If no colon, split by comma
+                    comma_parts = re.split('[,，]', sentence)
+                    if len(comma_parts) > 1:
+                        current_comma_part = ""
+                        for comma_part in comma_parts:
+                            if len(current_comma_part.encode('utf-8')) + len(comma_part.encode('utf-8')) <= max_chars:
+                                current_comma_part += comma_part + ','
+                            else:
+                                if current_comma_part:
+                                    batches.append(current_comma_part.rstrip(','))
+                                current_comma_part = comma_part + ','
+                        if current_comma_part:
+                            batches.append(current_comma_part.rstrip(','))
+                    else:
+                        # If no comma, split by words
+                        batches.extend(split_by_words(sentence))
+            else:
+                current_batch = sentence
+    
+    if current_batch:
+        batches.append(current_batch)
+    
+    return batches
+
+def infer_batch(ref_audio, ref_text, gen_text_batches, exp_name, remove_silence, progress=gr.Progress()):
     if exp_name == "F5-TTS":
         ema_model = F5TTS_ema_model
     elif exp_name == "E2-TTS":
         ema_model = E2TTS_ema_model
-    
-    if not ref_text.strip():
-        gr.Info("No reference text provided, transcribing reference audio...")
-        ref_text = outputs = 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...")
-    audio, sr = torchaudio.load(ref_audio)
-    max_chars = int(len(ref_text) / (audio.shape[-1] / sr) * (30 - audio.shape[-1] / sr))
-    # Audio
+
+    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
@@ -135,28 +227,25 @@ def infer(ref_audio_orig, ref_text, gen_text, exp_name, remove_silence, progress
         resampler = torchaudio.transforms.Resample(sr, target_sample_rate)
         audio = resampler(audio)
     audio = audio.to(device)
-    # Chunk
-    chunks = txtsplit(gen_text, 0.7*max_chars, 0.9*max_chars)
-    results = []
-    generated_mel_specs = []
-    for chunk in progress.tqdm(chunks):
+
+    generated_waves = []
+    spectrograms = []
+
+    for i, gen_text in enumerate(progress.tqdm(gen_text_batches)):
         # Prepare the text
-        text_list = [ref_text + chunk]
+        if len(ref_text[-1].encode('utf-8')) == 1:
+            ref_text = ref_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
-        # if fix_duration is not None:
-        #     duration = int(fix_duration * target_sample_rate / hop_length)
-        # else:
         zh_pause_punc = r"。，、；：？！"
         ref_text_len = len(ref_text.encode('utf-8')) + 3 * len(re.findall(zh_pause_punc, ref_text))
         gen_text_len = len(gen_text.encode('utf-8')) + 3 * len(re.findall(zh_pause_punc, gen_text))
-        chunk = len(chunk.encode('utf-8')) + 3 * len(re.findall(zh_pause_punc, gen_text))
         duration = ref_audio_len + int(ref_audio_len / ref_text_len * gen_text_len / speed)
-    
+
         # inference
-        gr.Info(f"Generating audio using {exp_name}")
         with torch.inference_mode():
             generated, _ = ema_model.sample(
                 cond=audio,
@@ -166,29 +255,26 @@ def infer(ref_audio_orig, ref_text, gen_text, exp_name, remove_silence, progress
                 cfg_strength=cfg_strength,
                 sway_sampling_coef=sway_sampling_coef,
             )
-    
+
         generated = generated[:, ref_audio_len:, :]
-        generated_mel_spec = rearrange(generated, '1 n d -> 1 d n')
-        gr.Info("Running vocoder")
-        vocos = Vocos.from_pretrained("charactr/vocos-mel-24khz")
+        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
-    
+
         # wav -> numpy
         generated_wave = generated_wave.squeeze().cpu().numpy()
-        results.append(generated_wave)
-    generated_wave = np.concatenate(results)
+        
+        generated_waves.append(generated_wave)
+        spectrograms.append(generated_mel_spec[0].cpu().numpy())
+
+    # Combine all generated waves
+    final_wave = np.concatenate(generated_waves)
+
+    # Remove silence
     if remove_silence:
-        gr.Info("Removing audio silences... This may take a moment")
-        # non_silent_intervals = librosa.effects.split(generated_wave, top_db=30)
-        # non_silent_wave = np.array([])
-        # for interval in non_silent_intervals:
-        #     start, end = interval
-        #     non_silent_wave = np.concatenate([non_silent_wave, generated_wave[start:end]])
-        # generated_wave = non_silent_wave
         with tempfile.NamedTemporaryFile(delete=False, suffix=".wav") as f:
-            sf.write(f.name, generated_wave, target_sample_rate)
+            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)
@@ -196,65 +282,543 @@ def infer(ref_audio_orig, ref_text, gen_text, exp_name, remove_silence, progress
                 non_silent_wave += non_silent_seg
             aseg = non_silent_wave
             aseg.export(f.name, format="wav")
-            generated_wave, _ = torchaudio.load(f.name)
-        generated_wave = generated_wave.squeeze().cpu().numpy()
+            final_wave, _ = torchaudio.load(f.name)
+        final_wave = final_wave.squeeze().cpu().numpy()
 
+    # Create a combined spectrogram
+    combined_spectrogram = np.concatenate(spectrograms, axis=1)
+    
+    with tempfile.NamedTemporaryFile(suffix=".png", delete=False) as tmp_spectrogram:
+        spectrogram_path = tmp_spectrogram.name
+        save_spectrogram(combined_spectrogram, spectrogram_path)
 
-    # spectogram
-    # with tempfile.NamedTemporaryFile(suffix=".png", delete=False) as tmp_spectrogram:
-    #     spectrogram_path = tmp_spectrogram.name
-    #     save_spectrogram(generated_mel_spec[0].cpu().numpy(), spectrogram_path)
+    return (target_sample_rate, final_wave), spectrogram_path
 
-    return (target_sample_rate, generated_wave)
+def infer(ref_audio_orig, ref_text, gen_text, exp_name, remove_silence, custom_split_words=''):
+    if not custom_split_words.strip():
+        custom_words = [word.strip() for word in custom_split_words.split(',')]
+        global SPLIT_WORDS
+        SPLIT_WORDS = custom_words
 
-with gr.Blocks() as app:
-    gr.Markdown("""
-# E2/F5 TTS
+    print(gen_text)
 
-This is an unofficial E2/F5 TTS demo. This demo supports the following TTS models:
+    gr.Info("Converting audio...")
+    with tempfile.NamedTemporaryFile(delete=False, suffix=".wav") as f:
+        aseg = AudioSegment.from_file(ref_audio_orig)
 
-* [E2-TTS](https://arxiv.org/abs/2406.18009) (Embarrassingly Easy Fully Non-Autoregressive Zero-Shot TTS)
-* [F5-TTS](https://arxiv.org/abs/2410.06885) (A Fairytaler that Fakes Fluent and Faithful Speech with Flow Matching)
+        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
 
-This 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).
+        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
 
-The checkpoints support English and Chinese.
+    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...")
 
-If 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).
+    # 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) * (30 - audio.shape[-1] / sr))
+    gen_text_batches = split_text_into_batches(gen_text, max_chars=max_chars)
+    print('ref_text', ref_text)
+    for i, gen_text in enumerate(gen_text_batches):
+        print(f'gen_text {i}', gen_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)
+    
+def generate_podcast(script, speaker1_name, ref_audio1, ref_text1, speaker2_name, ref_audio2, ref_text2, exp_name, 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
+    
+    generated_audio_segments = []
+    
+    for i in range(0, len(speaker_blocks), 2):
+        speaker = speaker_blocks[i]
+        text = speaker_blocks[i+1].strip()
+        
+        # Determine which speaker is talking
+        if speaker == speaker1_name:
+            ref_audio = ref_audio1
+            ref_text = ref_text1
+        elif speaker == speaker2_name:
+            ref_audio = ref_audio2
+            ref_text = ref_text2
+        else:
+            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)
+        
+        # Convert the generated audio to a numpy array
+        sr, audio_data = audio
+        
+        # Save the audio data as a WAV file
+        with tempfile.NamedTemporaryFile(suffix=".wav", delete=False) as temp_file:
+            sf.write(temp_file.name, audio_data, sr)
+            audio_segment = AudioSegment.from_wav(temp_file.name)
+        
+        generated_audio_segments.append(audio_segment)
+        
+        # Add a short pause between speakers
+        pause = AudioSegment.silent(duration=500)  # 500ms pause
+        generated_audio_segments.append(pause)
+    
+    # Concatenate all audio segments
+    final_podcast = sum(generated_audio_segments)
+    
+    # Export the final podcast
+    with tempfile.NamedTemporaryFile(suffix=".wav", delete=False) as temp_file:
+        podcast_path = temp_file.name
+        final_podcast.export(podcast_path, format="wav")
+    
+    return podcast_path
 
-The model is licensed under the CC-BY-NC license, this demo cannot be used for commercial purposes.
+def parse_speechtypes_text(gen_text):
+    # Pattern to find (Emotion)
+    pattern = r'\((.*?)\)'
 
-**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.**
-""")
+    # Split the text by the pattern
+    tokens = re.split(pattern, gen_text)
+
+    segments = []
+
+    current_emotion = 'Regular'
 
+    for i in range(len(tokens)):
+        if i % 2 == 0:
+            # This is text
+            text = tokens[i].strip()
+            if text:
+                segments.append({'emotion': current_emotion, 'text': text})
+        else:
+            # This is emotion
+            emotion = tokens[i].strip()
+            current_emotion = emotion
+
+    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("""
+# Credits
+
+* [mrfakename](https://github.com/fakerybakery) for the original [online demo](https://huggingface.co/spaces/mrfakename/E2-F5-TTS)
+* [RootingInLoad](https://github.com/RootingInLoad) for the podcast generation
+""")
+with gr.Blocks() as app_tts:
+    gr.Markdown("# Batched TTS")
     ref_audio_input = gr.Audio(label="Reference Audio", type="filepath")
-    gen_text_input = gr.Textbox(label="Text to Generate (longer text will use chunking)", lines=4)
-    model_choice = gr.Radio(choices=["F5-TTS", "E2-TTS"], label="Choose TTS Model", value="F5-TTS")
+    gen_text_input = gr.Textbox(label="Text to Generate", lines=10)
+    model_choice = gr.Radio(
+        choices=["F5-TTS", "E2-TTS"], label="Choose TTS Model", value="F5-TTS"
+    )
     generate_btn = gr.Button("Synthesize", variant="primary")
     with gr.Accordion("Advanced Settings", open=False):
-        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)
-        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)
-    
+        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,
+        )
+        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,
+        )
+        split_words_input = gr.Textbox(
+            label="Custom Split Words",
+            info="Enter custom words to split on, separated by commas. Leave blank to use default list.",
+            lines=2,
+        )
+        speed_slider = gr.Slider(
+            label="Speed",
+            minimum=0.3,
+            maximum=2.0,
+            value=speed,
+            step=0.1,
+            info="Adjust the speed of the audio.",
+        )
+    speed_slider.change(update_speed, inputs=speed_slider)
+
     audio_output = gr.Audio(label="Synthesized Audio")
-    # spectrogram_output = gr.Image(label="Spectrogram")
+    spectrogram_output = gr.Image(label="Spectrogram")
+
+    generate_btn.click(
+        infer,
+        inputs=[
+            ref_audio_input,
+            ref_text_input,
+            gen_text_input,
+            model_choice,
+            remove_silence,
+            split_words_input,
+        ],
+        outputs=[audio_output, spectrogram_output],
+    )
+    
+with gr.Blocks() as app_podcast:
+    gr.Markdown("# Podcast Generation")
+    speaker1_name = gr.Textbox(label="Speaker 1 Name")
+    ref_audio_input1 = gr.Audio(label="Reference Audio (Speaker 1)", type="filepath")
+    ref_text_input1 = gr.Textbox(label="Reference Text (Speaker 1)", lines=2)
+    
+    speaker2_name = gr.Textbox(label="Speaker 2 Name")
+    ref_audio_input2 = gr.Audio(label="Reference Audio (Speaker 2)", type="filepath")
+    ref_text_input2 = gr.Textbox(label="Reference Text (Speaker 2)", lines=2)
+    
+    script_input = gr.Textbox(label="Podcast Script", lines=10, 
+                                placeholder="Enter the script with speaker names at the start of each block, e.g.:\nSean: How did you start studying...\n\nMeghan: I came to my interest in technology...\nIt was a long journey...\n\nSean: That's fascinating. Can you elaborate...")
+    
+    podcast_model_choice = gr.Radio(
+        choices=["F5-TTS", "E2-TTS"], label="Choose TTS Model", value="F5-TTS"
+    )
+    podcast_remove_silence = gr.Checkbox(
+        label="Remove Silences",
+        value=True,
+    )
+    generate_podcast_btn = gr.Button("Generate Podcast", variant="primary")
+    podcast_output = gr.Audio(label="Generated Podcast")
+
+    def podcast_generation(script, speaker1, ref_audio1, ref_text1, speaker2, ref_audio2, ref_text2, model, remove_silence):
+        return generate_podcast(script, speaker1, ref_audio1, ref_text1, speaker2, ref_audio2, ref_text2, model, remove_silence)
+
+    generate_podcast_btn.click(
+        podcast_generation,
+        inputs=[
+            script_input,
+            speaker1_name,
+            ref_audio_input1,
+            ref_text_input1,
+            speaker2_name,
+            ref_audio_input2,
+            ref_text_input2,
+            podcast_model_choice,
+            podcast_remove_silence,
+        ],
+        outputs=podcast_output,
+    )
+
+def parse_emotional_text(gen_text):
+    # Pattern to find (Emotion)
+    pattern = r'\((.*?)\)'
+
+    # Split the text by the pattern
+    tokens = re.split(pattern, gen_text)
+
+    segments = []
+
+    current_emotion = 'Regular'
+
+    for i in range(len(tokens)):
+        if i % 2 == 0:
+            # This is text
+            text = tokens[i].strip()
+            if text:
+                segments.append({'emotion': current_emotion, 'text': text})
+        else:
+            # This is emotion
+            emotion = tokens[i].strip()
+            current_emotion = emotion
+
+    return segments
+
+with gr.Blocks() as app_emotional:
+    # New section for emotional generation
+    gr.Markdown(
+        """
+    # Multiple Speech-Type Generation
+
+    This section allows you to upload different audio clips for each speech type. 'Regular' emotion is mandatory. You can add additional speech types by clicking the "Add Speech Type" button. Enter your text in the format shown below, and the system will generate speech using the appropriate emotions. If unspecified, the model will use the regular speech type. The current speech type will be used until the next speech type is specified.
+
+    **Example Input:**
+
+    (Regular) Hello, I'd like to order a sandwich please. (Surprised) What do you mean you're out of bread? (Sad) I really wanted a sandwich though... (Angry) You know what, darn you and your little shop, you suck! (Whisper) I'll just go back home and cry now. (Shouting) Why me?!
+    """
+    )
+
+    gr.Markdown("Upload different audio clips for each speech type. 'Regular' emotion is mandatory. You can add additional speech types by clicking the 'Add Speech Type' button.")
+
+    # Regular speech type (mandatory)
+    with gr.Row():
+        regular_name = gr.Textbox(value='Regular', label='Speech Type Name', interactive=False)
+        regular_audio = gr.Audio(label='Regular Reference Audio', type='filepath')
+        regular_ref_text = gr.Textbox(label='Reference Text (Regular)', lines=2)
+
+    # Additional speech types (up to 9 more)
+    max_speech_types = 10
+    speech_type_names = []
+    speech_type_audios = []
+    speech_type_ref_texts = []
+    speech_type_delete_btns = []
+
+    for i in range(max_speech_types - 1):
+        with gr.Row():
+            name_input = gr.Textbox(label='Speech Type Name', visible=False)
+            audio_input = gr.Audio(label='Reference Audio', type='filepath', visible=False)
+            ref_text_input = gr.Textbox(label='Reference Text', lines=2, visible=False)
+            delete_btn = gr.Button("Delete", variant="secondary", visible=False)
+        speech_type_names.append(name_input)
+        speech_type_audios.append(audio_input)
+        speech_type_ref_texts.append(ref_text_input)
+        speech_type_delete_btns.append(delete_btn)
+
+    # Button to add speech type
+    add_speech_type_btn = gr.Button("Add Speech Type")
+
+    # Keep track of current number of speech types
+    speech_type_count = gr.State(value=0)
+
+    # Function to add a speech type
+    def add_speech_type_fn(speech_type_count):
+        if speech_type_count < max_speech_types - 1:
+            speech_type_count += 1
+            # Prepare updates for the components
+            name_updates = []
+            audio_updates = []
+            ref_text_updates = []
+            delete_btn_updates = []
+            for i in range(max_speech_types - 1):
+                if i < speech_type_count:
+                    name_updates.append(gr.update(visible=True))
+                    audio_updates.append(gr.update(visible=True))
+                    ref_text_updates.append(gr.update(visible=True))
+                    delete_btn_updates.append(gr.update(visible=True))
+                else:
+                    name_updates.append(gr.update())
+                    audio_updates.append(gr.update())
+                    ref_text_updates.append(gr.update())
+                    delete_btn_updates.append(gr.update())
+        else:
+            # Optionally, show a warning
+            # gr.Warning("Maximum number of speech types reached.")
+            name_updates = [gr.update() for _ in range(max_speech_types - 1)]
+            audio_updates = [gr.update() for _ in range(max_speech_types - 1)]
+            ref_text_updates = [gr.update() for _ in range(max_speech_types - 1)]
+            delete_btn_updates = [gr.update() for _ in range(max_speech_types - 1)]
+        return [speech_type_count] + name_updates + audio_updates + ref_text_updates + delete_btn_updates
+
+    add_speech_type_btn.click(
+        add_speech_type_fn,
+        inputs=speech_type_count,
+        outputs=[speech_type_count] + speech_type_names + speech_type_audios + speech_type_ref_texts + speech_type_delete_btns
+    )
+
+    # Function to delete a speech type
+    def make_delete_speech_type_fn(index):
+        def delete_speech_type_fn(speech_type_count):
+            # Prepare updates
+            name_updates = []
+            audio_updates = []
+            ref_text_updates = []
+            delete_btn_updates = []
+
+            for i in range(max_speech_types - 1):
+                if i == index:
+                    name_updates.append(gr.update(visible=False, value=''))
+                    audio_updates.append(gr.update(visible=False, value=None))
+                    ref_text_updates.append(gr.update(visible=False, value=''))
+                    delete_btn_updates.append(gr.update(visible=False))
+                else:
+                    name_updates.append(gr.update())
+                    audio_updates.append(gr.update())
+                    ref_text_updates.append(gr.update())
+                    delete_btn_updates.append(gr.update())
+
+            speech_type_count = max(0, speech_type_count - 1)
+
+            return [speech_type_count] + name_updates + audio_updates + ref_text_updates + delete_btn_updates
+
+        return delete_speech_type_fn
+
+    for i, delete_btn in enumerate(speech_type_delete_btns):
+        delete_fn = make_delete_speech_type_fn(i)
+        delete_btn.click(
+            delete_fn,
+            inputs=speech_type_count,
+            outputs=[speech_type_count] + speech_type_names + speech_type_audios + speech_type_ref_texts + speech_type_delete_btns
+        )
+
+    # Text input for the prompt
+    gen_text_input_emotional = gr.Textbox(label="Text to Generate", lines=10)
+
+    # Model choice
+    model_choice_emotional = gr.Radio(
+        choices=["F5-TTS", "E2-TTS"], label="Choose TTS Model", value="F5-TTS"
+    )
 
-    generate_btn.click(infer, inputs=[ref_audio_input, ref_text_input, gen_text_input, model_choice, remove_silence], outputs=[audio_output])
-    gr.Markdown("""
-## Run Locally
+    with gr.Accordion("Advanced Settings", open=False):
+        remove_silence_emotional = gr.Checkbox(
+            label="Remove Silences",
+            value=True,
+        )
+
+    # Generate button
+    generate_emotional_btn = gr.Button("Generate Emotional Speech", variant="primary")
+
+    # Output audio
+    audio_output_emotional = gr.Audio(label="Synthesized Audio")
+
+    def generate_emotional_speech(
+        regular_audio,
+        regular_ref_text,
+        gen_text,
+        *args,
+    ):
+        num_additional_speech_types = max_speech_types - 1
+        speech_type_names_list = args[:num_additional_speech_types]
+        speech_type_audios_list = args[num_additional_speech_types:2 * num_additional_speech_types]
+        speech_type_ref_texts_list = args[2 * num_additional_speech_types:3 * num_additional_speech_types]
+        model_choice = args[3 * num_additional_speech_types]
+        remove_silence = args[3 * num_additional_speech_types + 1]
+
+        # Collect the speech types and their audios into a dict
+        speech_types = {'Regular': {'audio': regular_audio, 'ref_text': regular_ref_text}}
+
+        for name_input, audio_input, ref_text_input in zip(speech_type_names_list, speech_type_audios_list, speech_type_ref_texts_list):
+            if name_input and audio_input:
+                speech_types[name_input] = {'audio': audio_input, 'ref_text': ref_text_input}
+
+        # Parse the gen_text into segments
+        segments = parse_speechtypes_text(gen_text)
+
+        # For each segment, generate speech
+        generated_audio_segments = []
+        current_emotion = 'Regular'
+
+        for segment in segments:
+            emotion = segment['emotion']
+            text = segment['text']
+
+            if emotion in speech_types:
+                current_emotion = emotion
+            else:
+                # If emotion not available, default to Regular
+                current_emotion = 'Regular'
+
+            ref_audio = speech_types[current_emotion]['audio']
+            ref_text = speech_types[current_emotion].get('ref_text', '')
+
+            # Generate speech for this segment
+            audio, _ = infer(ref_audio, ref_text, text, model_choice, remove_silence, "")
+            sr, audio_data = audio
+
+            generated_audio_segments.append(audio_data)
+
+        # Concatenate all audio segments
+        if generated_audio_segments:
+            final_audio_data = np.concatenate(generated_audio_segments)
+            return (sr, final_audio_data)
+        else:
+            gr.Warning("No audio generated.")
+            return None
+
+    generate_emotional_btn.click(
+        generate_emotional_speech,
+        inputs=[
+            regular_audio,
+            regular_ref_text,
+            gen_text_input_emotional,
+        ] + speech_type_names + speech_type_audios + speech_type_ref_texts + [
+            model_choice_emotional,
+            remove_silence_emotional,
+        ],
+        outputs=audio_output_emotional,
+    )
+
+    # Validation function to disable Generate button if speech types are missing
+    def validate_speech_types(
+        gen_text,
+        regular_name,
+        *args
+    ):
+        num_additional_speech_types = max_speech_types - 1
+        speech_type_names_list = args[:num_additional_speech_types]
+
+        # Collect the speech types names
+        speech_types_available = set()
+        if regular_name:
+            speech_types_available.add(regular_name)
+        for name_input in speech_type_names_list:
+            if name_input:
+                speech_types_available.add(name_input)
+
+        # Parse the gen_text to get the speech types used
+        segments = parse_emotional_text(gen_text)
+        speech_types_in_text = set(segment['emotion'] for segment in segments)
+
+        # Check if all speech types in text are available
+        missing_speech_types = speech_types_in_text - speech_types_available
+
+        if missing_speech_types:
+            # Disable the generate button
+            return gr.update(interactive=False)
+        else:
+            # Enable the generate button
+            return gr.update(interactive=True)
+
+    gen_text_input_emotional.change(
+        validate_speech_types,
+        inputs=[gen_text_input_emotional, regular_name] + speech_type_names,
+        outputs=generate_emotional_btn
+    )
+with gr.Blocks() as app:
+    gr.Markdown(
+        """
+# E2/F5 TTS
 
-Run this demo locally on CPU, CUDA, or MPS/Apple Silicon (requires macOS >= 14):
+This is a local web UI for F5 TTS with advanced batch processing support. This app supports the following TTS models:
 
-First, ensure `ffmpeg` is installed.
+* [F5-TTS](https://arxiv.org/abs/2410.06885) (A Fairytaler that Fakes Fluent and Faithful Speech with Flow Matching)
+* [E2 TTS](https://arxiv.org/abs/2406.18009) (Embarrassingly Easy Fully Non-Autoregressive Zero-Shot TTS)
 
-```bash
-git clone https://huggingface.co/spaces/mrfakename/E2-F5-TTS
-cd E2-F5-TTS
-python -m pip install -r requirements.txt
-python app_local.py
-```
+The checkpoints support English and Chinese.
 
-""")
-    gr.Markdown("Unofficial demo by [mrfakename](https://x.com/realmrfakename)")
+If you're having issues, try converting your reference audio to WAV or MP3, clipping it to 15s, and shortening your prompt.
+
+**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.**
+"""
+    )
+    gr.TabbedInterface([app_tts, app_podcast, app_emotional, app_credits], ["TTS", "Podcast", "Multi-Style", "Credits"])
+
+@click.command()
+@click.option("--port", "-p", default=None, type=int, help="Port to run the app on")
+@click.option("--host", "-H", default=None, help="Host to run the app on")
+@click.option(
+    "--share",
+    "-s",
+    default=False,
+    is_flag=True,
+    help="Share the app via Gradio share link",
+)
+@click.option("--api", "-a", default=True, is_flag=True, help="Allow API access")
+def main(port, host, share, api):
+    global app
+    print(f"Starting app...")
+    app.queue(api_open=api).launch(
+        server_name=host, server_port=port, share=share, show_api=api
+    )
 
 
-app.queue().launch()
+if __name__ == "__main__":
+    main()

inference-cli.py

@@ -0,0 +1,378 @@
+import re
+import torch
+import torchaudio
+import numpy as np
+import tempfile
+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,
+)
+from transformers import pipeline
+import soundfile as sf
+import tomli
+import argparse
+import tqdm
+from pathlib import Path
+
+parser = argparse.ArgumentParser(
+    prog="python3 inference-cli.py",
+    description="Commandline interface for E2/F5 TTS with Advanced Batch Processing.",
+    epilog="Specify  options above  to override  one or more settings from config.",
+)
+parser.add_argument(
+    "-c",
+    "--config",
+    help="Configuration file. Default=cli-config.toml",
+    default="inference-cli.toml",
+)
+parser.add_argument(
+    "-m",
+    "--model",
+    help="F5-TTS | E2-TTS",
+)
+parser.add_argument(
+    "-r",
+    "--ref_audio",
+    type=str,
+    help="Reference audio file < 15 seconds."
+)
+parser.add_argument(
+    "-s",
+    "--ref_text",
+    type=str,
+    default="666",
+    help="Subtitle for the reference audio."
+)
+parser.add_argument(
+    "-t",
+    "--gen_text",
+    type=str,
+    help="Text to generate.",
+)
+parser.add_argument(
+    "-o",
+    "--output_dir",
+    type=str,
+    help="Path to output folder..",
+)
+parser.add_argument(
+    "--remove_silence",
+    help="Remove silence.",
+)
+args = parser.parse_args()
+
+config = tomli.load(open(args.config, "rb"))
+
+ref_audio = args.ref_audio if args.ref_audio else config["ref_audio"]
+ref_text = args.ref_text if args.ref_text != "666" else config["ref_text"]
+gen_text = args.gen_text if args.gen_text else config["gen_text"]
+output_dir = args.output_dir if args.output_dir else config["output_dir"]
+model = args.model if args.model else config["model"]
+remove_silence = args.remove_silence if args.remove_silence else config["remove_silence"]
+wave_path = Path(output_dir)/"out.wav"
+spectrogram_path = Path(output_dir)/"out.png"
+
+SPLIT_WORDS = [
+    "but", "however", "nevertheless", "yet", "still",
+    "therefore", "thus", "hence", "consequently",
+    "moreover", "furthermore", "additionally",
+    "meanwhile", "alternatively", "otherwise",
+    "namely", "specifically", "for example", "such as",
+    "in fact", "indeed", "notably",
+    "in contrast", "on the other hand", "conversely",
+    "in conclusion", "to summarize", "finally"
+]
+
+device = (
+    "cuda"
+    if torch.cuda.is_available()
+    else "mps" if torch.backends.mps.is_available() else "cpu"
+)
+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
+
+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
+
+
+# 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)
+
+def split_text_into_batches(text, max_chars=200, split_words=SPLIT_WORDS):
+    if len(text.encode('utf-8')) <= max_chars:
+        return [text]
+    if text[-1] not in ['。', '.', '!', '！', '?', '？']:
+        text += '.'
+        
+    sentences = re.split('([。.!?！？])', text)
+    sentences = [''.join(i) for i in zip(sentences[0::2], sentences[1::2])]
+    
+    batches = []
+    current_batch = ""
+    
+    def split_by_words(text):
+        words = text.split()
+        current_word_part = ""
+        word_batches = []
+        for word in words:
+            if len(current_word_part.encode('utf-8')) + len(word.encode('utf-8')) + 1 <= max_chars:
+                current_word_part += word + ' '
+            else:
+                if current_word_part:
+                    # Try to find a suitable split word
+                    for split_word in split_words:
+                        split_index = current_word_part.rfind(' ' + split_word + ' ')
+                        if split_index != -1:
+                            word_batches.append(current_word_part[:split_index].strip())
+                            current_word_part = current_word_part[split_index:].strip() + ' '
+                            break
+                    else:
+                        # If no suitable split word found, just append the current part
+                        word_batches.append(current_word_part.strip())
+                        current_word_part = ""
+                current_word_part += word + ' '
+        if current_word_part:
+            word_batches.append(current_word_part.strip())
+        return word_batches
+
+    for sentence in sentences:
+        if len(current_batch.encode('utf-8')) + len(sentence.encode('utf-8')) <= max_chars:
+            current_batch += sentence
+        else:
+            # If adding this sentence would exceed the limit
+            if current_batch:
+                batches.append(current_batch)
+                current_batch = ""
+            
+            # If the sentence itself is longer than max_chars, split it
+            if len(sentence.encode('utf-8')) > max_chars:
+                # First, try to split by colon
+                colon_parts = sentence.split(':')
+                if len(colon_parts) > 1:
+                    for part in colon_parts:
+                        if len(part.encode('utf-8')) <= max_chars:
+                            batches.append(part)
+                        else:
+                            # If colon part is still too long, split by comma
+                            comma_parts = re.split('[,，]', part)
+                            if len(comma_parts) > 1:
+                                current_comma_part = ""
+                                for comma_part in comma_parts:
+                                    if len(current_comma_part.encode('utf-8')) + len(comma_part.encode('utf-8')) <= max_chars:
+                                        current_comma_part += comma_part + ','
+                                    else:
+                                        if current_comma_part:
+                                            batches.append(current_comma_part.rstrip(','))
+                                        current_comma_part = comma_part + ','
+                                if current_comma_part:
+                                    batches.append(current_comma_part.rstrip(','))
+                            else:
+                                # If no comma, split by words
+                                batches.extend(split_by_words(part))
+                else:
+                    # If no colon, split by comma
+                    comma_parts = re.split('[,，]', sentence)
+                    if len(comma_parts) > 1:
+                        current_comma_part = ""
+                        for comma_part in comma_parts:
+                            if len(current_comma_part.encode('utf-8')) + len(comma_part.encode('utf-8')) <= max_chars:
+                                current_comma_part += comma_part + ','
+                            else:
+                                if current_comma_part:
+                                    batches.append(current_comma_part.rstrip(','))
+                                current_comma_part = comma_part + ','
+                        if current_comma_part:
+                            batches.append(current_comma_part.rstrip(','))
+                    else:
+                        # If no comma, split by words
+                        batches.extend(split_by_words(sentence))
+            else:
+                current_batch = sentence
+    
+    if current_batch:
+        batches.append(current_batch)
+    
+    return batches
+
+def infer_batch(ref_audio, ref_text, gen_text_batches, model, remove_silence):
+    if model == "F5-TTS":
+        ema_model = load_model(model, "F5TTS_Base", DiT, F5TTS_model_cfg, 1200000)
+    elif model == "E2-TTS":
+        ema_model = load_model(model, "E2TTS_Base", UNetT, E2TTS_model_cfg, 1200000)
+
+    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 = []
+
+    for i, gen_text in enumerate(tqdm.tqdm(gen_text_batches)):
+        # Prepare the text
+        if len(ref_text[-1].encode('utf-8')) == 1:
+            ref_text = ref_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
+        zh_pause_punc = r"。，、；：？！"
+        ref_text_len = len(ref_text.encode('utf-8')) + 3 * len(re.findall(zh_pause_punc, ref_text))
+        gen_text_len = len(gen_text.encode('utf-8')) + 3 * len(re.findall(zh_pause_punc, gen_text))
+        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[:, ref_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
+
+        # 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
+    final_wave = np.concatenate(generated_waves)
+
+    with open(wave_path, "wb") as f:
+        sf.write(f.name, final_wave, target_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")
+        print(f.name)
+
+    # Create a combined spectrogram
+    combined_spectrogram = np.concatenate(spectrograms, axis=1)
+    save_spectrogram(combined_spectrogram, spectrogram_path)
+    print(spectrogram_path)
+
+
+def infer(ref_audio_orig, ref_text, gen_text, model, remove_silence, custom_split_words):
+    if not custom_split_words.strip():
+        custom_words = [word.strip() for word in custom_split_words.split(',')]
+        global SPLIT_WORDS
+        SPLIT_WORDS = custom_words
+
+    print(gen_text)
+
+    print("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=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
+
+        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...")
+        pipe = pipeline(
+            "automatic-speech-recognition",
+            model="openai/whisper-large-v3-turbo",
+            torch_dtype=torch.float16,
+            device=device,
+        )
+        ref_text = 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...")
+
+    # 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) * (30 - audio.shape[-1] / sr))
+    gen_text_batches = split_text_into_batches(gen_text, max_chars=max_chars)
+    print('ref_text', ref_text)
+    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)
+    
+
+infer(ref_audio, ref_text, gen_text, model, remove_silence, ",".join(SPLIT_WORDS))

inference-cli.toml

@@ -0,0 +1,8 @@
+# F5-TTS | E2-TTS
+model = "F5-TTS"
+ref_audio = "tests/ref_audio/test_en_1_ref_short.wav"
+# If an empty "", transcribes the reference audio automatically.
+ref_text = "Some call me nature, others call me mother nature."
+gen_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."
+remove_silence = true
+output_dir = "tests"

model/cfm.py

@@ -95,6 +95,7 @@ class CFM(nn.Module):
         no_ref_audio = False,
         duplicate_test = False,
         t_inter = 0.1,
+        edit_mask = None,
     ):
         self.eval()
 
@@ -125,6 +126,8 @@ class CFM(nn.Module):
         # duration
 
         cond_mask = lens_to_mask(lens)
+        if edit_mask is not None:
+            cond_mask = cond_mask & edit_mask
 
         if isinstance(duration, int):
             duration = torch.full((batch,), duration, device = device, dtype = torch.long)
@@ -142,7 +145,10 @@ class CFM(nn.Module):
         cond_mask = rearrange(cond_mask, '... -> ... 1')
         step_cond = torch.where(cond_mask, cond, torch.zeros_like(cond))  # allow direct control (cut cond audio) with lens passed in
 
-        mask = lens_to_mask(duration)
+        if batch > 1:
+            mask = lens_to_mask(duration)
+        else:  # save memory and speed up, as single inference need no mask currently
+            mask = None
 
         # test for no ref audio
         if no_ref_audio:

model/dataset.py

@@ -188,7 +188,7 @@ def load_dataset(
         dataset_type: str = "CustomDataset", 
         audio_type: str = "raw", 
         mel_spec_kwargs: dict = dict()
-        ) -> CustomDataset | HFDataset:
+        ) -> CustomDataset:
     
     print("Loading dataset ...")
 

model/trainer.py

@@ -138,19 +138,24 @@ class Trainer:
         if "model_last.pt" in os.listdir(self.checkpoint_path):
             latest_checkpoint = "model_last.pt"
         else:
-            latest_checkpoint = sorted(os.listdir(self.checkpoint_path), key=lambda x: int(''.join(filter(str.isdigit, x))))[-1]
+            latest_checkpoint = sorted([f for f in os.listdir(self.checkpoint_path) if f.endswith('.pt')], key=lambda x: int(''.join(filter(str.isdigit, x))))[-1]
         # checkpoint = torch.load(f"{self.checkpoint_path}/{latest_checkpoint}", map_location=self.accelerator.device)  # rather use accelerator.load_state ಥ_ಥ
         checkpoint = torch.load(f"{self.checkpoint_path}/{latest_checkpoint}", map_location="cpu")
-        self.accelerator.unwrap_model(self.model).load_state_dict(checkpoint['model_state_dict'])
-        self.accelerator.unwrap_model(self.optimizer).load_state_dict(checkpoint['optimizer_state_dict'])
 
         if self.is_main:
             self.ema_model.load_state_dict(checkpoint['ema_model_state_dict'])
 
-        if self.scheduler:
-            self.scheduler.load_state_dict(checkpoint['scheduler_state_dict'])
-        
-        step = checkpoint['step']
+        if 'step' in checkpoint:
+            self.accelerator.unwrap_model(self.model).load_state_dict(checkpoint['model_state_dict'])
+            self.accelerator.unwrap_model(self.optimizer).load_state_dict(checkpoint['optimizer_state_dict'])
+            if self.scheduler:
+                self.scheduler.load_state_dict(checkpoint['scheduler_state_dict'])
+            step = checkpoint['step']
+        else:
+            checkpoint['model_state_dict'] = {k.replace("ema_model.", ""): v for k, v in checkpoint['ema_model_state_dict'].items() if k not in ["initted", "step"]}
+            self.accelerator.unwrap_model(self.model).load_state_dict(checkpoint['model_state_dict'])
+            step = 0
+
         del checkpoint; gc.collect()
         return step
 
@@ -163,16 +168,16 @@ class Trainer:
             generator = None
 
         if self.batch_size_type == "sample":
-            train_dataloader = DataLoader(train_dataset, collate_fn=collate_fn, num_workers=num_workers, pin_memory=True, 
+            train_dataloader = DataLoader(train_dataset, collate_fn=collate_fn, num_workers=num_workers, pin_memory=True, persistent_workers=True,
                                           batch_size=self.batch_size, shuffle=True, generator=generator)
         elif self.batch_size_type == "frame":
             self.accelerator.even_batches = False
             sampler = SequentialSampler(train_dataset)
             batch_sampler = DynamicBatchSampler(sampler, self.batch_size, max_samples=self.max_samples, random_seed=resumable_with_seed, drop_last=False)
-            train_dataloader = DataLoader(train_dataset, collate_fn=collate_fn, num_workers=num_workers, pin_memory=True,
+            train_dataloader = DataLoader(train_dataset, collate_fn=collate_fn, num_workers=num_workers, pin_memory=True, persistent_workers=True,
                                           batch_sampler=batch_sampler)
         else:
-            raise ValueError(f"batch_size_type must be either 'sample' or 'frame', but recieved {self.batch_size_type}")
+            raise ValueError(f"batch_size_type must be either 'sample' or 'frame', but received {self.batch_size_type}")
         
         #  accelerator.prepare() dispatches batches to devices;
         #  which means the length of dataloader calculated before, should consider the number of devices

model/utils.py

@@ -134,7 +134,7 @@ def get_tokenizer(dataset_name, tokenizer: str = "pinyin"):
                 - if use "byte", set to 256 (unicode byte range) 
     ''' 
     if tokenizer in ["pinyin", "char"]:
-        with open (f"data/{dataset_name}_{tokenizer}/vocab.txt", "r") as f:
+        with open (f"data/{dataset_name}_{tokenizer}/vocab.txt", "r", encoding="utf-8") as f:
             vocab_char_map = {}
             for i, char in enumerate(f):
                 vocab_char_map[char[:-1]] = i
@@ -153,9 +153,11 @@ def get_tokenizer(dataset_name, tokenizer: str = "pinyin"):
 def convert_char_to_pinyin(text_list, polyphone = True):
     final_text_list = []
     god_knows_why_en_testset_contains_zh_quote = str.maketrans({'“': '"', '”': '"', '‘': "'", '’': "'"})  # in case librispeech (orig no-pc) test-clean
+    custom_trans = str.maketrans({';': ','})  # add custom trans here, to address oov
     for text in text_list:
         char_list = []
         text = text.translate(god_knows_why_en_testset_contains_zh_quote)
+        text = text.translate(custom_trans)
         for seg in jieba.cut(text):
             seg_byte_len = len(bytes(seg, 'UTF-8'))
             if seg_byte_len == len(seg):  # if pure alphabets and symbols
@@ -273,6 +275,8 @@ def get_inference_prompt(
             ref_audio = resampler(ref_audio)
 
         # Text
+        if len(prompt_text[-1].encode('utf-8')) == 1:
+            prompt_text = prompt_text + " "
         text = [prompt_text + gt_text]
         if tokenizer == "pinyin":
             text_list = convert_char_to_pinyin(text, polyphone = polyphone)
@@ -292,8 +296,8 @@ def get_inference_prompt(
             # ref_audio = gt_audio
         else:
             zh_pause_punc = r"。，、；：？！"
-            ref_text_len = len(prompt_text) + len(re.findall(zh_pause_punc, prompt_text))
-            gen_text_len = len(gt_text) + len(re.findall(zh_pause_punc, gt_text))
+            ref_text_len = len(prompt_text.encode('utf-8')) + 3 * len(re.findall(zh_pause_punc, prompt_text))
+            gen_text_len = len(gt_text.encode('utf-8')) + 3 * len(re.findall(zh_pause_punc, gt_text))
             total_mel_len = ref_mel_len + int(ref_mel_len / ref_text_len * gen_text_len / speed)
 
         # to mel spectrogram
@@ -543,3 +547,28 @@ def repetition_found(text, length = 2, tolerance = 10):
         if count > tolerance:
             return True
     return False
+
+
+# load model checkpoint for inference
+
+def load_checkpoint(model, ckpt_path, device, use_ema = True):
+    from ema_pytorch import EMA
+
+    ckpt_type = ckpt_path.split(".")[-1]
+    if ckpt_type == "safetensors":
+        from safetensors.torch import load_file
+        checkpoint = load_file(ckpt_path, device=device)
+    else:
+        checkpoint = torch.load(ckpt_path, map_location=device)
+
+    if use_ema == True:
+        ema_model = EMA(model, include_online_model = False).to(device)
+        if ckpt_type == "safetensors":
+            ema_model.load_state_dict(checkpoint)
+        else:
+            ema_model.load_state_dict(checkpoint['ema_model_state_dict'])
+        ema_model.copy_params_from_ema_to_model()
+    else:
+        model.load_state_dict(checkpoint['model_state_dict'])
+        
+    return model

requirements.txt

@@ -1,17 +1,24 @@
 accelerate>=0.33.0
+cached_path
+click
 datasets
 einops>=0.8.0
 einx>=0.3.0
 ema_pytorch>=0.5.2
 faster_whisper
 funasr
+gradio
 jieba
 jiwer
 librosa
 matplotlib
+numpy==1.23.5
+pydub
 pypinyin
-torch>=2.0
-torchaudio>=2.3.0
+safetensors
+soundfile
+# torch>=2.0
+# torchaudio>=2.3.0
 torchdiffeq
 tqdm>=4.65.0
 transformers
@@ -19,9 +26,4 @@ vocos
 wandb
 x_transformers>=1.31.14
 zhconv
-zhon
-cached_path
-pydub
-txtsplit
-detoxify
-soundfile
+zhon

scripts/eval_infer_batch.py

@@ -0,0 +1,199 @@
+import sys, os
+sys.path.append(os.getcwd())
+
+import time
+import random
+from tqdm import tqdm
+import argparse
+
+import torch
+import torchaudio
+from accelerate import Accelerator
+from einops import rearrange
+from vocos import Vocos
+
+from model import CFM, UNetT, DiT
+from model.utils import (
+    load_checkpoint,
+    get_tokenizer, 
+    get_seedtts_testset_metainfo, 
+    get_librispeech_test_clean_metainfo, 
+    get_inference_prompt,
+)
+
+accelerator = Accelerator()
+device = f"cuda:{accelerator.process_index}"
+
+
+# --------------------- Dataset Settings -------------------- #
+
+target_sample_rate = 24000
+n_mel_channels = 100
+hop_length = 256
+target_rms = 0.1
+
+tokenizer = "pinyin"
+
+
+# ---------------------- infer setting ---------------------- #
+
+parser = argparse.ArgumentParser(description="batch inference")
+
+parser.add_argument('-s', '--seed', default=None, type=int)
+parser.add_argument('-d', '--dataset', default="Emilia_ZH_EN")
+parser.add_argument('-n', '--expname', required=True)
+parser.add_argument('-c', '--ckptstep', default=1200000, type=int)
+
+parser.add_argument('-nfe', '--nfestep', default=32, type=int)
+parser.add_argument('-o', '--odemethod', default="euler")
+parser.add_argument('-ss', '--swaysampling', default=-1, type=float)
+
+parser.add_argument('-t', '--testset', required=True)
+
+args = parser.parse_args()
+
+
+seed = args.seed
+dataset_name = args.dataset
+exp_name = args.expname
+ckpt_step = args.ckptstep
+ckpt_path = f"ckpts/{exp_name}/model_{ckpt_step}.pt"
+
+nfe_step = args.nfestep
+ode_method = args.odemethod
+sway_sampling_coef = args.swaysampling
+
+testset = args.testset
+
+
+infer_batch_size = 1  # max frames. 1 for ddp single inference (recommended)
+cfg_strength = 2.
+speed = 1.
+use_truth_duration = False
+no_ref_audio = False
+
+
+if exp_name == "F5TTS_Base":
+    model_cls = DiT
+    model_cfg = dict(dim = 1024, depth = 22, heads = 16, ff_mult = 2, text_dim = 512, conv_layers = 4)
+
+elif exp_name == "E2TTS_Base":
+    model_cls = UNetT
+    model_cfg = dict(dim = 1024, depth = 24, heads = 16, ff_mult = 4)
+
+
+if testset == "ls_pc_test_clean":
+    metalst = "data/librispeech_pc_test_clean_cross_sentence.lst"
+    librispeech_test_clean_path = "<SOME_PATH>/LibriSpeech/test-clean"  # test-clean path
+    metainfo = get_librispeech_test_clean_metainfo(metalst, librispeech_test_clean_path)
+    
+elif testset == "seedtts_test_zh":
+    metalst = "data/seedtts_testset/zh/meta.lst"
+    metainfo = get_seedtts_testset_metainfo(metalst)
+
+elif testset == "seedtts_test_en":
+    metalst = "data/seedtts_testset/en/meta.lst"
+    metainfo = get_seedtts_testset_metainfo(metalst)
+
+
+# path to save genereted wavs
+if seed is None: seed = random.randint(-10000, 10000)
+output_dir = f"results/{exp_name}_{ckpt_step}/{testset}/" \
+    f"seed{seed}_{ode_method}_nfe{nfe_step}" \
+    f"{f'_ss{sway_sampling_coef}' if sway_sampling_coef else ''}" \
+    f"_cfg{cfg_strength}_speed{speed}" \
+    f"{'_gt-dur' if use_truth_duration else ''}" \
+    f"{'_no-ref-audio' if no_ref_audio else ''}"
+
+
+# -------------------------------------------------#
+
+use_ema = True
+
+prompts_all = get_inference_prompt(
+    metainfo, 
+    speed = speed, 
+    tokenizer = tokenizer, 
+    target_sample_rate = target_sample_rate, 
+    n_mel_channels = n_mel_channels,
+    hop_length = hop_length,
+    target_rms = target_rms,
+    use_truth_duration = use_truth_duration,
+    infer_batch_size = infer_batch_size,
+)
+
+# Vocoder model
+local = False
+if local:
+    vocos_local_path = "../checkpoints/charactr/vocos-mel-24khz"
+    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:
+    vocos = Vocos.from_pretrained("charactr/vocos-mel-24khz")
+
+# Tokenizer
+vocab_char_map, vocab_size = get_tokenizer(dataset_name, tokenizer)
+
+# Model
+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 = use_ema)
+
+if not os.path.exists(output_dir) and accelerator.is_main_process:
+    os.makedirs(output_dir)
+
+# start batch inference
+accelerator.wait_for_everyone()
+start = time.time()
+
+with accelerator.split_between_processes(prompts_all) as prompts:
+
+    for prompt in tqdm(prompts, disable=not accelerator.is_local_main_process):
+        utts, ref_rms_list, ref_mels, ref_mel_lens, total_mel_lens, final_text_list = prompt
+        ref_mels = ref_mels.to(device)
+        ref_mel_lens = torch.tensor(ref_mel_lens, dtype = torch.long).to(device)
+        total_mel_lens = torch.tensor(total_mel_lens, dtype = torch.long).to(device)
+        
+        # Inference
+        with torch.inference_mode():
+            generated, _ = model.sample(
+                cond = ref_mels,
+                text = final_text_list,
+                duration = total_mel_lens,
+                lens = ref_mel_lens,
+                steps = nfe_step,
+                cfg_strength = cfg_strength,
+                sway_sampling_coef = sway_sampling_coef,
+                no_ref_audio = no_ref_audio,
+                seed = seed,
+            )
+        # Final result
+        for i, gen in enumerate(generated):
+            gen = gen[ref_mel_lens[i]:total_mel_lens[i], :].unsqueeze(0)
+            gen_mel_spec = rearrange(gen, '1 n d -> 1 d n')
+            generated_wave = vocos.decode(gen_mel_spec.cpu())
+            if ref_rms_list[i] < target_rms:
+                generated_wave = generated_wave * ref_rms_list[i] / target_rms
+            torchaudio.save(f"{output_dir}/{utts[i]}.wav", generated_wave, target_sample_rate)
+
+accelerator.wait_for_everyone()
+if accelerator.is_main_process:
+    timediff = time.time() - start
+    print(f"Done batch inference in {timediff / 60 :.2f} minutes.")

scripts/eval_infer_batch.sh

@@ -0,0 +1,13 @@
+#!/bin/bash
+
+# e.g. F5-TTS, 16 NFE
+accelerate launch scripts/eval_infer_batch.py -n "F5TTS_Base" -t "seedtts_test_zh" -nfe 16
+accelerate launch scripts/eval_infer_batch.py -n "F5TTS_Base" -t "seedtts_test_en" -nfe 16
+accelerate launch scripts/eval_infer_batch.py -n "F5TTS_Base" -t "ls_pc_test_clean" -nfe 16
+
+# e.g. Vanilla E2 TTS, 32 NFE
+accelerate launch scripts/eval_infer_batch.py -n "E2TTS_Base" -t "seedtts_test_zh" -o "midpoint" -ss 0
+accelerate launch scripts/eval_infer_batch.py -n "E2TTS_Base" -t "seedtts_test_en" -o "midpoint" -ss 0
+accelerate launch scripts/eval_infer_batch.py -n "E2TTS_Base" -t "ls_pc_test_clean" -o "midpoint" -ss 0
+
+# etc.

speech_edit.py

@@ -0,0 +1,182 @@
+import os
+
+import torch
+import torch.nn.functional as F
+import torchaudio
+from einops import rearrange
+from vocos import Vocos
+
+from model import CFM, UNetT, DiT, MMDiT
+from model.utils import (
+    load_checkpoint,
+    get_tokenizer, 
+    convert_char_to_pinyin, 
+    save_spectrogram,
+)
+
+device = "cuda" if torch.cuda.is_available() else "mps" if torch.backends.mps.is_available() else "cpu"
+
+
+# --------------------- Dataset Settings -------------------- #
+
+target_sample_rate = 24000
+n_mel_channels = 100
+hop_length = 256
+target_rms = 0.1
+
+tokenizer = "pinyin"
+dataset_name = "Emilia_ZH_EN"
+
+
+# ---------------------- infer setting ---------------------- #
+
+seed = None  # int | None
+
+exp_name = "F5TTS_Base"  # F5TTS_Base | E2TTS_Base
+ckpt_step = 1200000
+
+nfe_step = 32  # 16, 32
+cfg_strength = 2.
+ode_method = 'euler'  # euler | midpoint
+sway_sampling_coef = -1.
+speed = 1.
+
+if exp_name == "F5TTS_Base":
+    model_cls = DiT
+    model_cfg = dict(dim = 1024, depth = 22, heads = 16, ff_mult = 2, text_dim = 512, conv_layers = 4)
+
+elif exp_name == "E2TTS_Base":
+    model_cls = UNetT
+    model_cfg = dict(dim = 1024, depth = 24, heads = 16, ff_mult = 4)
+
+ckpt_path = f"ckpts/{exp_name}/model_{ckpt_step}.pt"
+output_dir = "tests"
+
+# [leverage https://github.com/MahmoudAshraf97/ctc-forced-aligner to get char level alignment]
+# pip install git+https://github.com/MahmoudAshraf97/ctc-forced-aligner.git
+# [write the origin_text into a file, e.g. tests/test_edit.txt]
+# ctc-forced-aligner --audio_path "tests/ref_audio/test_en_1_ref_short.wav" --text_path "tests/test_edit.txt" --language "zho" --romanize --split_size "char"
+# [result will be saved at same path of audio file]
+# [--language "zho" for Chinese, "eng" for English]
+# [if local ckpt, set --alignment_model "../checkpoints/mms-300m-1130-forced-aligner"]
+
+audio_to_edit = "tests/ref_audio/test_en_1_ref_short.wav"
+origin_text = "Some call me nature, others call me mother nature."
+target_text = "Some call me optimist, others call me realist."
+parts_to_edit = [[1.42, 2.44], [4.04, 4.9], ]  # stard_ends of "nature" & "mother nature", in seconds
+fix_duration = [1.2, 1, ]  # fix duration for "optimist" & "realist", in seconds
+
+# audio_to_edit = "tests/ref_audio/test_zh_1_ref_short.wav"
+# origin_text = "对，这就是我，万人敬仰的太乙真人。"
+# target_text = "对，那就是你，万人敬仰的太白金星。"
+# parts_to_edit = [[0.84, 1.4], [1.92, 2.4], [4.26, 6.26], ]
+# fix_duration = None  # use origin text duration
+
+
+# -------------------------------------------------#
+
+use_ema = True
+
+if not os.path.exists(output_dir):
+    os.makedirs(output_dir)
+
+# Vocoder model
+local = False
+if local:
+    vocos_local_path = "../checkpoints/charactr/vocos-mel-24khz"
+    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:
+    vocos = Vocos.from_pretrained("charactr/vocos-mel-24khz")
+
+# Tokenizer
+vocab_char_map, vocab_size = get_tokenizer(dataset_name, tokenizer)
+
+# Model
+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 = use_ema)
+
+# Audio
+audio, sr = torchaudio.load(audio_to_edit)
+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)
+offset = 0
+audio_ = torch.zeros(1, 0)
+edit_mask = torch.zeros(1, 0, dtype=torch.bool)
+for part in parts_to_edit:
+    start, end = part
+    part_dur = end - start if fix_duration is None else fix_duration.pop(0)
+    part_dur = part_dur * target_sample_rate
+    start = start * target_sample_rate
+    audio_ = torch.cat((audio_, audio[:, round(offset):round(start)], torch.zeros(1, round(part_dur))), dim = -1)
+    edit_mask = torch.cat((edit_mask, 
+                           torch.ones(1, round((start - offset) / hop_length), dtype = torch.bool), 
+                           torch.zeros(1, round(part_dur / hop_length), dtype = torch.bool)
+                           ), dim = -1)
+    offset = end * target_sample_rate
+# audio = torch.cat((audio_, audio[:, round(offset):]), dim = -1)
+edit_mask = F.pad(edit_mask, (0, audio.shape[-1] // hop_length - edit_mask.shape[-1] + 1), value = True)
+audio = audio.to(device)
+edit_mask = edit_mask.to(device)
+
+# Text
+text_list = [target_text]
+if tokenizer == "pinyin":
+    final_text_list = convert_char_to_pinyin(text_list)
+else:
+    final_text_list = [text_list]
+print(f"text  : {text_list}")
+print(f"pinyin: {final_text_list}")
+
+# Duration
+ref_audio_len = 0
+duration = audio.shape[-1] // hop_length
+
+# Inference
+with torch.inference_mode():
+    generated, trajectory = model.sample(
+        cond = audio,
+        text = final_text_list,
+        duration = duration,
+        steps = nfe_step,
+        cfg_strength = cfg_strength,
+        sway_sampling_coef = sway_sampling_coef,
+        seed = seed,
+        edit_mask = edit_mask,
+    )
+print(f"Generated mel: {generated.shape}")
+
+# Final result
+generated = generated[:, ref_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
+
+save_spectrogram(generated_mel_spec[0].cpu().numpy(), f"{output_dir}/test_single_edit.png")
+torchaudio.save(f"{output_dir}/test_single_edit.wav", generated_wave, target_sample_rate)
+print(f"Generated wav: {generated_wave.shape}")

train.py

@@ -0,0 +1,91 @@
+from model import CFM, UNetT, DiT, MMDiT, Trainer
+from model.utils import get_tokenizer
+from model.dataset import load_dataset
+
+
+# -------------------------- Dataset Settings --------------------------- #
+
+target_sample_rate = 24000
+n_mel_channels = 100
+hop_length = 256
+
+tokenizer = "pinyin"
+dataset_name = "Emilia_ZH_EN"
+
+
+# -------------------------- Training Settings -------------------------- #
+
+exp_name = "F5TTS_Base"  # F5TTS_Base | E2TTS_Base
+
+learning_rate = 7.5e-5
+
+batch_size_per_gpu = 38400  # 8 GPUs, 8 * 38400 = 307200
+batch_size_type = "frame"  # "frame" or "sample"
+max_samples = 64  # max sequences per batch if use frame-wise batch_size. we set 32 for small models, 64 for base models
+grad_accumulation_steps = 1  # note: updates = steps / grad_accumulation_steps
+max_grad_norm = 1.
+
+epochs = 11  # use linear decay, thus epochs control the slope
+num_warmup_updates = 20000  # warmup steps
+save_per_updates = 50000  # save checkpoint per steps
+last_per_steps = 5000  # save last checkpoint per steps
+
+# model params
+if exp_name == "F5TTS_Base":
+    wandb_resume_id = None
+    model_cls = DiT
+    model_cfg = dict(dim = 1024, depth = 22, heads = 16, ff_mult = 2, text_dim = 512, conv_layers = 4)
+elif exp_name == "E2TTS_Base":
+    wandb_resume_id = None
+    model_cls = UNetT
+    model_cfg = dict(dim = 1024, depth = 24, heads = 16, ff_mult = 4)
+
+
+# ----------------------------------------------------------------------- #
+
+def main():
+
+    vocab_char_map, vocab_size = get_tokenizer(dataset_name, tokenizer)
+
+    mel_spec_kwargs = dict(
+            target_sample_rate = target_sample_rate, 
+            n_mel_channels = n_mel_channels,
+            hop_length = hop_length,
+        )
+    
+    e2tts = CFM(
+        transformer = model_cls(
+            **model_cfg,
+            text_num_embeds = vocab_size, 
+            mel_dim = n_mel_channels
+        ),
+        mel_spec_kwargs = mel_spec_kwargs,
+        vocab_char_map = vocab_char_map,
+    )
+
+    trainer = Trainer(
+        e2tts,
+        epochs, 
+        learning_rate,
+        num_warmup_updates = num_warmup_updates,
+        save_per_updates = save_per_updates, 
+        checkpoint_path = f'ckpts/{exp_name}',
+        batch_size = batch_size_per_gpu, 
+        batch_size_type = batch_size_type,
+        max_samples = max_samples,
+        grad_accumulation_steps = grad_accumulation_steps,
+        max_grad_norm = max_grad_norm,
+        wandb_project = "CFM-TTS",
+        wandb_run_name = exp_name,
+        wandb_resume_id = wandb_resume_id,
+        last_per_steps = last_per_steps,
+    )
+
+    train_dataset = load_dataset(dataset_name, tokenizer, mel_spec_kwargs=mel_spec_kwargs)
+    trainer.train(train_dataset, 
+                  resumable_with_seed = 666 # seed for shuffling dataset
+                  )
+
+
+if __name__ == '__main__':
+    main()