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

pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "f5-tts"
-version = "1.0.6"
+version = "1.0.7"
 description = "F5-TTS: A Fairytaler that Fakes Fluent and Faithful Speech with Flow Matching"
 readme = "README.md"
 license = {text = "MIT License"}

src/f5_tts/infer/SHARED.md

@@ -44,6 +44,7 @@
 
 ```bash
 Model: hf://SWivid/F5-TTS/F5TTS_v1_Base/model_1250000.safetensors
+# A Variant Model: hf://SWivid/F5-TTS/F5TTS_v1_Base_no_zero_init/model_1250000.safetensors
 Vocab: hf://SWivid/F5-TTS/F5TTS_v1_Base/vocab.txt
 Config: {"dim": 1024, "depth": 22, "heads": 16, "ff_mult": 2, "text_dim": 512, "conv_layers": 4}
 ```

src/f5_tts/train/README.md

@@ -51,6 +51,8 @@ Discussion board for Finetuning [#57](https://github.com/SWivid/F5-TTS/discussio
 
 Gradio UI training/finetuning with `src/f5_tts/train/finetune_gradio.py` see [#143](https://github.com/SWivid/F5-TTS/discussions/143).
 
+If want to finetune with a variant version e.g. *F5TTS_v1_Base_no_zero_init*, manually download pretrained checkpoint from model weight repository and fill in the path correspondingly on web interface.
+
 If use tensorboard as logger, install it first with `pip install tensorboard`.
 
 <ins>The `use_ema = True` might be harmful for early-stage finetuned checkpoints</ins> (which goes just few updates, thus ema weights still dominated by pretrained ones), try turn it off with finetune gradio option or `load_model(..., use_ema=False)`, see if offer better results.

src/f5_tts/train/finetune_gradio.py

@@ -965,21 +965,23 @@ def calculate_train(
     )
 
 
-def extract_and_save_ema_model(checkpoint_path: str, new_checkpoint_path: str, safetensors: bool) -> str:
+def prune_checkpoint(checkpoint_path: str, new_checkpoint_path: str, save_ema: bool, safetensors: bool) -> str:
     try:
         checkpoint = torch.load(checkpoint_path, weights_only=True)
         print("Original Checkpoint Keys:", checkpoint.keys())
 
-        ema_model_state_dict = checkpoint.get("ema_model_state_dict", None)
-        if ema_model_state_dict is None:
-            return "No 'ema_model_state_dict' found in the checkpoint."
+        to_retain = "ema_model_state_dict" if save_ema else "model_state_dict"
+        try:
+            model_state_dict_to_retain = checkpoint[to_retain]
+        except KeyError:
+            return f"{to_retain} not found in the checkpoint."
 
         if safetensors:
             new_checkpoint_path = new_checkpoint_path.replace(".pt", ".safetensors")
-            save_file(ema_model_state_dict, new_checkpoint_path)
+            save_file(model_state_dict_to_retain, new_checkpoint_path)
         else:
             new_checkpoint_path = new_checkpoint_path.replace(".safetensors", ".pt")
-            new_checkpoint = {"ema_model_state_dict": ema_model_state_dict}
+            new_checkpoint = {"ema_model_state_dict": model_state_dict_to_retain}
             torch.save(new_checkpoint, new_checkpoint_path)
 
         return f"New checkpoint saved at: {new_checkpoint_path}"
@@ -1849,12 +1851,14 @@ Reduce the Base model size from 5GB to 1.3GB. The new checkpoint file prunes out
 ```""")
             txt_path_checkpoint = gr.Textbox(label="Path to Checkpoint:")
             txt_path_checkpoint_small = gr.Textbox(label="Path to Output:")
-            ch_safetensors = gr.Checkbox(label="Safetensors", value="")
+            with gr.Row():
+                ch_save_ema = gr.Checkbox(label="Save EMA checkpoint", value=True)
+                ch_safetensors = gr.Checkbox(label="Save with safetensors format", value=True)
             txt_info_reduse = gr.Textbox(label="Info", value="")
-            reduse_button = gr.Button("Reduce")
+            reduse_button = gr.Button("Prune")
             reduse_button.click(
-                fn=extract_and_save_ema_model,
-                inputs=[txt_path_checkpoint, txt_path_checkpoint_small, ch_safetensors],
+                fn=prune_checkpoint,
+                inputs=[txt_path_checkpoint, txt_path_checkpoint_small, ch_save_ema, ch_safetensors],
                 outputs=[txt_info_reduse],
             )