Merge branch 'SWivid:main' into main

README.md

@@ -2,8 +2,9 @@
 
 [![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)
+[![demo](https://img.shields.io/badge/GitHub-Demo%20page-orange.svg)](https://swivid.github.io/F5-TTS/)
+[![hfspace](https://img.shields.io/badge/🤗-Space%20demo-yellow)](https://huggingface.co/spaces/mrfakename/E2-F5-TTS)
+[![msspace](https://img.shields.io/badge/🤖-Space%20demo-blue)](https://modelscope.cn/studios/modelscope/E2-F5-TTS)
 [![lab](https://img.shields.io/badge/X--LANCE-Lab-grey?labelColor=lightgrey)](https://x-lance.sjtu.edu.cn/)
 <img src="https://github.com/user-attachments/assets/12d7749c-071a-427c-81bf-b87b91def670" alt="Watermark" style="width: 40px; height: auto">
 
@@ -52,7 +53,7 @@ python scripts/prepare_emilia.py
 python scripts/prepare_wenetspeech4tts.py
 ```
 
-## Training
+## Training & Finetuning
 
 Once your datasets are prepared, you can start the training process.
 
@@ -64,9 +65,11 @@ accelerate launch train.py
 ```
 An initial guidance on Finetuning [#57](https://github.com/SWivid/F5-TTS/discussions/57).
 
+Gradio UI finetuning with `finetune_gradio.py` see [#143](https://github.com/SWivid/F5-TTS/discussions/143).
+
 ## Inference
 
-The pretrained model checkpoints can be reached at [🤗 Hugging Face](https://huggingface.co/SWivid/F5-TTS) and [⭐ Model Scope](https://www.modelscope.cn/models/SWivid/F5-TTS_Emilia-ZH-EN), or automatically downloaded with `inference-cli` and `gradio_app`.
+The pretrained model checkpoints can be reached at [🤗 Hugging Face](https://huggingface.co/SWivid/F5-TTS) and [🤖 Model Scope](https://www.modelscope.cn/models/SWivid/F5-TTS_Emilia-ZH-EN), 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.
@@ -90,6 +93,9 @@ python inference-cli.py \
 --ref_audio "tests/ref_audio/test_zh_1_ref_short.wav" \
 --ref_text "对，这就是我，万人敬仰的太乙真人。" \
 --gen_text "突然，身边一阵笑声。我看着他们，意气风发地挺直了胸膛，甩了甩那稍显肉感的双臂，轻笑道，我身上的肉，是为了掩饰我爆棚的魅力，否则，岂不吓坏了你们呢？"
+
+# Multi voice
+python inference-cli.py -c samples/story.toml
 ```
 
 ### Gradio App
@@ -188,11 +194,13 @@ python scripts/eval_librispeech_test_clean.py
 - [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
+- [mrfakename](https://x.com/realmrfakename) huggingface space demo ~
+- [f5-tts-mlx](https://github.com/lucasnewman/f5-tts-mlx/tree/main) Implementation of F5-TTS, with the MLX framework.
 
 ## Citation
+If our work and codebase is useful for you, please cite as:
 ```
 @article{chen-etal-2024-f5tts,
       title={F5-TTS: A Fairytaler that Fakes Fluent and Faithful Speech with Flow Matching}, 

finetune_gradio.py

@@ -339,7 +339,7 @@ def transcribe(file_audio,language="english"):
     )["text"].strip()
     return text_transcribe
 
-def transcribe_all(name_project,audio_file,language,user=False,progress=gr.Progress()):
+def transcribe_all(name_project,audio_files,language,user=False,progress=gr.Progress()):
     name_project+="_pinyin"
     path_project= os.path.join(path_data,name_project)
     path_dataset = os.path.join(path_project,"dataset")
@@ -357,7 +357,7 @@ def transcribe_all(name_project,audio_file,language,user=False,progress=gr.Progr
     if user:
        file_audios = [file for format in ('*.wav', '*.ogg', '*.opus', '*.mp3', '*.flac') for file in glob(os.path.join(path_dataset, format))]
     else:
-       file_audios = [audio_file]
+       file_audios = audio_files
 
     print([file_audios])   
 
@@ -580,7 +580,7 @@ with gr.Blocks() as app:
          ...
      ```""",visible=False)
 
-              audio_speaker = gr.Audio(label="voice",type="filepath")
+              audio_speaker = gr.File(label="voice",type="filepath",file_count="multiple")
               txt_lang = gr.Text(label="Language",value="english")
               bt_transcribe=bt_create=gr.Button("transcribe")
               txt_info_transcribe=gr.Text(label="info",value="")
@@ -686,4 +686,4 @@ if __name__ == "__main__":
     #transcribe_all(name)
     #create_metadata(name)
 
-    main()
+    main()

gradio_app.py

@@ -532,8 +532,8 @@ with gr.Blocks() as app_emotional:
         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
+    # Additional speech types (up to 99 more)
+    max_speech_types = 100
     speech_type_names = []
     speech_type_audios = []
     speech_type_ref_texts = []

inference-cli.py

@@ -282,29 +282,12 @@ def infer_batch(ref_audio, ref_text, gen_text_batches, model, remove_silence, cr
 
             final_wave = new_wave
 
-    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, cross_fade_duration=0.15):
-
-    print(gen_text)
+    return final_wave, combined_spectrogram
 
+def process_voice(ref_audio_orig, ref_text):
     print("Converting audio...")
     with tempfile.NamedTemporaryFile(delete=False, suffix=".wav") as f:
         aseg = AudioSegment.from_file(ref_audio_orig)
@@ -340,7 +323,10 @@ def infer(ref_audio_orig, ref_text, gen_text, model, remove_silence, cross_fade_
         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):
+    print(gen_text)
     # Add the functionality to ensure it ends with ". "
     if not ref_text.endswith(". ") and not ref_text.endswith("。"):
         if ref_text.endswith("."):
@@ -360,4 +346,47 @@ def infer(ref_audio_orig, ref_text, gen_text, model, remove_silence, cross_fade_
     return infer_batch((audio, sr), ref_text, gen_text_batches, model, remove_silence, cross_fade_duration)
     
 
-infer(ref_audio, ref_text, gen_text, model, remove_silence)
+def process(ref_audio, ref_text, text_gen, model, remove_silence):
+    main_voice = {"ref_audio":ref_audio, "ref_text":ref_text}
+    if "voices" not in config:
+        voices = {"main": main_voice}
+    else:
+        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'])
+
+    generated_audio_segments = []
+    reg1 = r'(?=\[\w+\])'
+    chunks = re.split(reg1, text_gen)
+    reg2 = r'\[(\w+)\]'
+    for text in chunks:
+        match = re.match(reg2, text)
+        if not match or voice not in voices:
+            voice = "main"
+        else:
+            voice = match[1]
+        text = re.sub(reg2, "", text)
+        gen_text = text.strip()
+        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)
+        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)
+            # 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)
+
+process(ref_audio, ref_text, gen_text, model, remove_silence)

model/trainer.py

@@ -140,7 +140,7 @@ class Trainer:
         else:
             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")
+        checkpoint = torch.load(f"{self.checkpoint_path}/{latest_checkpoint}", weights_only=True, map_location="cpu")
 
         if self.is_main:
             self.ema_model.load_state_dict(checkpoint['ema_model_state_dict'])

model/utils.py

@@ -509,7 +509,7 @@ def run_sim(args):
     device = f"cuda:{rank}"
 
     model = ECAPA_TDNN_SMALL(feat_dim=1024, feat_type='wavlm_large', config_path=None)
-    state_dict = torch.load(ckpt_dir, map_location=lambda storage, loc: storage)
+    state_dict = torch.load(ckpt_dir, weights_only=True, map_location=lambda storage, loc: storage)
     model.load_state_dict(state_dict['model'], strict=False)
 
     use_gpu=True if torch.cuda.is_available() else False
@@ -565,7 +565,7 @@ def load_checkpoint(model, ckpt_path, device, use_ema = True):
         from safetensors.torch import load_file
         checkpoint = load_file(ckpt_path, device=device)
     else:
-        checkpoint = torch.load(ckpt_path, map_location=device)
+        checkpoint = torch.load(ckpt_path, weights_only=True, map_location=device)
 
     if use_ema == True:
         ema_model = EMA(model, include_online_model = False).to(device)

samples/story.toml

@@ -0,0 +1,19 @@
+# F5-TTS | E2-TTS
+model = "F5-TTS"
+ref_audio = "samples/main.flac"
+# If an empty "", transcribes the reference audio automatically.
+ref_text = ""
+gen_text = ""
+# File with text to generate. Ignores the text above.
+gen_file = "samples/story.txt"
+remove_silence = true
+output_dir = "samples"
+
+[voices.town]
+ref_audio = "samples/town.flac"
+ref_text = ""
+
+[voices.country]
+ref_audio = "samples/country.flac"
+ref_text = ""
+

samples/story.txt

@@ -0,0 +1 @@
+A Town Mouse and a Country Mouse were acquaintances, and the Country Mouse one day invited his friend to come and see him at his home in the fields. The Town Mouse came, and they sat down to a dinner of barleycorns and roots, the latter of which had a distinctly earthy flavour. The fare was not much to the taste of the guest, and presently he broke out with [town] “My poor dear friend, you live here no better than the ants. Now, you should just see how I fare! My larder is a regular horn of plenty. You must come and stay with me, and I promise you you shall live on the fat of the land.” [main] So when he returned to town he took the Country Mouse with him, and showed him into a larder containing flour and oatmeal and figs and honey and dates. The Country Mouse had never seen anything like it, and sat down to enjoy the luxuries his friend provided: but before they had well begun, the door of the larder opened and someone came in. The two Mice scampered off and hid themselves in a narrow and exceedingly uncomfortable hole. Presently, when all was quiet, they ventured out again; but someone else came in, and off they scuttled again. This was too much for the visitor. [country] “Goodbye,” [main] said he, [country] “I’m off. You live in the lap of luxury, I can see, but you are surrounded by dangers; whereas at home I can enjoy my simple dinner of roots and corn in peace.”

scripts/eval_infer_batch.py

@@ -127,7 +127,7 @@ 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)
+    state_dict = torch.load(f"{vocos_local_path}/pytorch_model.bin", weights_only=True, map_location=device)
     vocos.load_state_dict(state_dict)
     vocos.eval()
 else:

scripts/prepare_csv_wavs.py

@@ -0,0 +1,132 @@
+import sys, os
+sys.path.append(os.getcwd())
+
+from pathlib import Path
+import json
+import shutil
+import argparse
+
+import csv
+import torchaudio
+from tqdm import tqdm
+from datasets.arrow_writer import ArrowWriter
+
+from model.utils import (
+    convert_char_to_pinyin,
+)
+
+PRETRAINED_VOCAB_PATH = Path(__file__).parent.parent / "data/Emilia_ZH_EN_pinyin/vocab.txt"
+
+def is_csv_wavs_format(input_dataset_dir):
+    fpath = Path(input_dataset_dir)
+    metadata = fpath / "metadata.csv"
+    wavs = fpath / 'wavs'
+    return metadata.exists() and metadata.is_file() and wavs.exists() and wavs.is_dir()
+
+
+def prepare_csv_wavs_dir(input_dir):
+    assert is_csv_wavs_format(input_dir), f"not csv_wavs format: {input_dir}"
+    input_dir = Path(input_dir)
+    metadata_path = input_dir / "metadata.csv"
+    audio_path_text_pairs = read_audio_text_pairs(metadata_path.as_posix())
+
+    sub_result, durations = [], []
+    vocab_set = set()
+    polyphone = True
+    for audio_path, text in audio_path_text_pairs:
+        if not Path(audio_path).exists():
+            print(f"audio {audio_path} not found, skipping")
+            continue
+        audio_duration = get_audio_duration(audio_path)
+        # assume tokenizer = "pinyin"  ("pinyin" | "char")
+        text = convert_char_to_pinyin([text], polyphone=polyphone)[0]
+        sub_result.append({"audio_path": audio_path, "text": text, "duration": audio_duration})
+        durations.append(audio_duration)
+        vocab_set.update(list(text))
+
+    return sub_result, durations, vocab_set
+
+def get_audio_duration(audio_path):
+    audio, sample_rate = torchaudio.load(audio_path)
+    num_channels = audio.shape[0]
+    return audio.shape[1] / (sample_rate * num_channels)
+
+def read_audio_text_pairs(csv_file_path):
+    audio_text_pairs = []
+
+    parent = Path(csv_file_path).parent
+    with open(csv_file_path, mode='r', newline='', encoding='utf-8') as csvfile:
+        reader = csv.reader(csvfile, delimiter='|')
+        next(reader)  # Skip the header row
+        for row in reader:
+            if len(row) >= 2:
+                audio_file = row[0].strip()  # First column: audio file path
+                text = row[1].strip()          # Second column: text
+                audio_file_path = parent / audio_file
+                audio_text_pairs.append((audio_file_path.as_posix(), text))
+
+    return audio_text_pairs
+
+
+def save_prepped_dataset(out_dir, result, duration_list, text_vocab_set, is_finetune):
+    out_dir = Path(out_dir)
+    # save preprocessed dataset to disk
+    out_dir.mkdir(exist_ok=True, parents=True)
+    print(f"\nSaving to {out_dir} ...")
+
+    # dataset = Dataset.from_dict({"audio_path": audio_path_list, "text": text_list, "duration": duration_list})  # oom
+    # dataset.save_to_disk(f"data/{dataset_name}/raw", max_shard_size="2GB")
+    raw_arrow_path = out_dir / "raw.arrow"
+    with ArrowWriter(path=raw_arrow_path.as_posix(), writer_batch_size=1) as writer:
+        for line in tqdm(result, desc=f"Writing to raw.arrow ..."):
+            writer.write(line)
+
+    # dup a json separately saving duration in case for DynamicBatchSampler ease
+    dur_json_path = out_dir / "duration.json"
+    with open(dur_json_path.as_posix(), 'w', encoding='utf-8') as f:
+        json.dump({"duration": duration_list}, f, ensure_ascii=False)
+
+    # vocab map, i.e. tokenizer
+    # add alphabets and symbols (optional, if plan to ft on de/fr etc.)
+    # if tokenizer == "pinyin":
+    #     text_vocab_set.update([chr(i) for i in range(32, 127)] + [chr(i) for i in range(192, 256)])
+    voca_out_path = out_dir / "vocab.txt"
+    with open(voca_out_path.as_posix(), "w") as f:
+        for vocab in sorted(text_vocab_set):
+            f.write(vocab + "\n")
+
+    if is_finetune:
+        file_vocab_finetune = PRETRAINED_VOCAB_PATH.as_posix()
+        shutil.copy2(file_vocab_finetune, voca_out_path)
+    else:
+        with open(voca_out_path, "w") as f:
+            for vocab in sorted(text_vocab_set):
+                f.write(vocab + "\n")
+
+    dataset_name = out_dir.stem
+    print(f"\nFor {dataset_name}, sample count: {len(result)}")
+    print(f"For {dataset_name}, vocab size is: {len(text_vocab_set)}")
+    print(f"For {dataset_name}, total {sum(duration_list)/3600:.2f} hours")
+
+
+def prepare_and_save_set(inp_dir, out_dir, is_finetune: bool = True):
+    if is_finetune:
+        assert PRETRAINED_VOCAB_PATH.exists(), f"pretrained vocab.txt not found: {PRETRAINED_VOCAB_PATH}"
+    sub_result, durations, vocab_set = prepare_csv_wavs_dir(inp_dir)
+    save_prepped_dataset(out_dir, sub_result, durations, vocab_set, is_finetune)
+
+
+def cli():
+    # finetune: python scripts/prepare_csv_wavs.py /path/to/input_dir /path/to/output_dir_pinyin
+    # pretrain: python scripts/prepare_csv_wavs.py /path/to/output_dir_pinyin --pretrain
+    parser = argparse.ArgumentParser(description="Prepare and save dataset.")
+    parser.add_argument('inp_dir', type=str, help="Input directory containing the data.")
+    parser.add_argument('out_dir', type=str, help="Output directory to save the prepared data.")
+    parser.add_argument('--pretrain', action='store_true', help="Enable for new pretrain, otherwise is a fine-tune")
+
+    args = parser.parse_args()
+
+    prepare_and_save_set(args.inp_dir, args.out_dir, is_finetune=not args.pretrain)
+
+if __name__ == "__main__":
+    cli()

speech_edit.py

@@ -85,8 +85,9 @@ 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)
+    state_dict = torch.load(f"{vocos_local_path}/pytorch_model.bin", weights_only=True, map_location=device)
     vocos.load_state_dict(state_dict)
+    
     vocos.eval()
 else:
     vocos = Vocos.from_pretrained("charactr/vocos-mel-24khz")