Merge pull request #628 from SWivid/bugfix/fix/enhancement-pr

fix/enhancement pr #399 #400 #529 #534 #591 #611
2026-07-28 13:26:59 -07:00 · 2024-12-16 16:32:55 +08:00
parent c85252f59a 81f69e57b7
commit c4c1fde6f2
16 changed files with 441 additions and 146 deletions
@@ -147,11 +147,11 @@ Note: Some model components have linting exceptions for E722 to accommodate tens
 ## 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
+- [Emilia](https://arxiv.org/abs/2407.05361), [WenetSpeech4TTS](https://arxiv.org/abs/2406.05763), [LibriTTS](https://arxiv.org/abs/1904.02882), [LJSpeech](https://keithito.com/LJ-Speech-Dataset/) 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
- [FunASR](https://github.com/modelscope/FunASR), [faster-whisper](https://github.com/SYSTRAN/faster-whisper), [UniSpeech](https://github.com/microsoft/UniSpeech) for evaluation tools
+- [torchdiffeq](https://github.com/rtqichen/torchdiffeq) as ODE solver, [Vocos](https://huggingface.co/charactr/vocos-mel-24khz) and [BigVGAN](https://github.com/NVIDIA/BigVGAN) as vocoder
+- [FunASR](https://github.com/modelscope/FunASR), [faster-whisper](https://github.com/SYSTRAN/faster-whisper), [UniSpeech](https://github.com/microsoft/UniSpeech), [SpeechMOS](https://github.com/tarepan/SpeechMOS) 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 with MLX framework by [Lucas Newman](https://github.com/lucasnewman)
@@ -28,6 +28,7 @@ model:
    ff_mult: 2
    text_dim: 512
    conv_layers: 4
+    checkpoint_activations: False  # recompute activations and save memory for extra compute
  mel_spec:
    target_sample_rate: 24000
    n_mel_channels: 100
@@ -28,6 +28,7 @@ model:
    ff_mult: 2
    text_dim: 512
    conv_layers: 4
+    checkpoint_activations: False  # recompute activations and save memory for extra compute
  mel_spec:
    target_sample_rate: 24000
    n_mel_channels: 100
@@ -39,11 +39,14 @@ Then update in the following scripts with the paths you put evaluation model ckp

 ### Objective Evaluation

-Update the path with your batch-inferenced results, and carry out WER / SIM evaluations:
+Update the path with your batch-inferenced results, and carry out WER / SIM / UTMOS evaluations:
 ```bash
-# Evaluation for Seed-TTS test set
-python src/f5_tts/eval/eval_seedtts_testset.py --gen_wav_dir <GEN_WAVE_DIR>
+# Evaluation [WER] for Seed-TTS test [ZH] set
+python src/f5_tts/eval/eval_seedtts_testset.py --eval_task wer --lang zh --gen_wav_dir <GEN_WAV_DIR> --gpu_nums 8

-# Evaluation for LibriSpeech-PC test-clean (cross-sentence)
-python src/f5_tts/eval/eval_librispeech_test_clean.py --gen_wav_dir <GEN_WAVE_DIR> --librispeech_test_clean_path <TEST_CLEAN_PATH>
-```
+# Evaluation [SIM] for LibriSpeech-PC test-clean (cross-sentence)
+python src/f5_tts/eval/eval_librispeech_test_clean.py --eval_task sim --gen_wav_dir <GEN_WAV_DIR> --librispeech_test_clean_path <TEST_CLEAN_PATH>
+
+# Evaluation [UTMOS]. --ext: Audio extension
+python src/f5_tts/eval/eval_utmos.py --audio_dir <WAV_DIR> --ext wav
+```
@@ -1,8 +1,9 @@
 # Evaluate with Librispeech test-clean, ~3s prompt to generate 4-10s audio (the way of valle/voicebox evaluation)

-import sys
-import os
 import argparse
+import json
+import os
+import sys

 sys.path.append(os.getcwd())

@@ -10,7 +11,6 @@ import multiprocessing as mp
 from importlib.resources import files

 import numpy as np
-
 from f5_tts.eval.utils_eval import (
    get_librispeech_test,
    run_asr_wer,
@@ -54,29 +54,41 @@ def main():
    wavlm_ckpt_dir = "../checkpoints/UniSpeech/wavlm_large_finetune.pth"

    # --------------------------- WER ---------------------------
+
    if eval_task == "wer":
+        wer_results = []
        wers = []
+
        with mp.Pool(processes=len(gpus)) as pool:
            args = [(rank, lang, sub_test_set, asr_ckpt_dir) for (rank, sub_test_set) in test_set]
            results = pool.map(run_asr_wer, args)
-            for wers_ in results:
-                wers.extend(wers_)
+            for r in results:
+                wer_results.extend(r)
+
+        wer_result_path = f"{gen_wav_dir}/{lang}_wer_results.jsonl"
+        with open(wer_result_path, "w") as f:
+            for line in wer_results:
+                wers.append(line["wer"])
+                json_line = json.dumps(line, ensure_ascii=False)
+                f.write(json_line + "\n")

        wer = round(np.mean(wers) * 100, 3)
        print(f"\nTotal {len(wers)} samples")
        print(f"WER      : {wer}%")
+        print(f"Results have been saved to {wer_result_path}")

    # --------------------------- SIM ---------------------------
+
    if eval_task == "sim":
-        sim_list = []
+        sims = []
        with mp.Pool(processes=len(gpus)) as pool:
            args = [(rank, sub_test_set, wavlm_ckpt_dir) for (rank, sub_test_set) in test_set]
            results = pool.map(run_sim, args)
-            for sim_ in results:
-                sim_list.extend(sim_)
+            for r in results:
+                sims.extend(r)

-        sim = round(sum(sim_list) / len(sim_list), 3)
-        print(f"\nTotal {len(sim_list)} samples")
+        sim = round(sum(sims) / len(sims), 3)
+        print(f"\nTotal {len(sims)} samples")
        print(f"SIM      : {sim}")


@@ -1,8 +1,9 @@
 # Evaluate with Seed-TTS testset

-import sys
-import os
 import argparse
+import json
+import os
+import sys

 sys.path.append(os.getcwd())

@@ -10,7 +11,6 @@ import multiprocessing as mp
 from importlib.resources import files

 import numpy as np
-
 from f5_tts.eval.utils_eval import (
    get_seed_tts_test,
    run_asr_wer,
@@ -55,28 +55,39 @@ def main():
    # --------------------------- WER ---------------------------

    if eval_task == "wer":
+        wer_results = []
        wers = []
+
        with mp.Pool(processes=len(gpus)) as pool:
            args = [(rank, lang, sub_test_set, asr_ckpt_dir) for (rank, sub_test_set) in test_set]
            results = pool.map(run_asr_wer, args)
-            for wers_ in results:
-                wers.extend(wers_)
+            for r in results:
+                wer_results.extend(r)
+
+        wer_result_path = f"{gen_wav_dir}/{lang}_wer_results.jsonl"
+        with open(wer_result_path, "w") as f:
+            for line in wer_results:
+                wers.append(line["wer"])
+                json_line = json.dumps(line, ensure_ascii=False)
+                f.write(json_line + "\n")

        wer = round(np.mean(wers) * 100, 3)
        print(f"\nTotal {len(wers)} samples")
        print(f"WER      : {wer}%")
+        print(f"Results have been saved to {wer_result_path}")

    # --------------------------- SIM ---------------------------
+
    if eval_task == "sim":
-        sim_list = []
+        sims = []
        with mp.Pool(processes=len(gpus)) as pool:
            args = [(rank, sub_test_set, wavlm_ckpt_dir) for (rank, sub_test_set) in test_set]
            results = pool.map(run_sim, args)
-            for sim_ in results:
-                sim_list.extend(sim_)
+            for r in results:
+                sims.extend(r)

-        sim = round(sum(sim_list) / len(sim_list), 3)
-        print(f"\nTotal {len(sim_list)} samples")
+        sim = round(sum(sims) / len(sims), 3)
+        print(f"\nTotal {len(sims)} samples")
        print(f"SIM      : {sim}")


@@ -0,0 +1,44 @@
+import argparse
+import json
+from pathlib import Path
+
+import librosa
+import torch
+from tqdm import tqdm
+
+
+def main():
+    parser = argparse.ArgumentParser(description="UTMOS Evaluation")
+    parser.add_argument("--audio_dir", type=str, required=True, help="Audio file path.")
+    parser.add_argument("--ext", type=str, default="wav", help="Audio extension.")
+    args = parser.parse_args()
+
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+
+    predictor = torch.hub.load("tarepan/SpeechMOS:v1.2.0", "utmos22_strong", trust_repo=True)
+    predictor = predictor.to(device)
+
+    audio_paths = list(Path(args.audio_dir).rglob(f"*.{args.ext}"))
+    utmos_results = {}
+    utmos_score = 0
+
+    for audio_path in tqdm(audio_paths, desc="Processing"):
+        wav_name = audio_path.stem
+        wav, sr = librosa.load(audio_path, sr=None, mono=True)
+        wav_tensor = torch.from_numpy(wav).to(device).unsqueeze(0)
+        score = predictor(wav_tensor, sr)
+        utmos_results[str(wav_name)] = score.item()
+        utmos_score += score.item()
+
+    avg_score = utmos_score / len(audio_paths) if len(audio_paths) > 0 else 0
+    print(f"UTMOS: {avg_score}")
+
+    utmos_result_path = Path(args.audio_dir) / "utmos_results.json"
+    with open(utmos_result_path, "w", encoding="utf-8") as f:
+        json.dump(utmos_results, f, ensure_ascii=False, indent=4)
+
+    print(f"Results have been saved to {utmos_result_path}")
+
+
+if __name__ == "__main__":
+    main()
@@ -2,6 +2,7 @@ import math
 import os
 import random
 import string
+from pathlib import Path

 import torch
 import torch.nn.functional as F
@@ -320,7 +321,7 @@ def run_asr_wer(args):
    from zhon.hanzi import punctuation

    punctuation_all = punctuation + string.punctuation
-    wers = []
+    wer_results = []

    from jiwer import compute_measures

@@ -335,8 +336,8 @@ def run_asr_wer(args):
            for segment in segments:
                hypo = hypo + " " + segment.text

-        # raw_truth = truth
-        # raw_hypo = hypo
+        raw_truth = truth
+        raw_hypo = hypo

        for x in punctuation_all:
            truth = truth.replace(x, "")
@@ -360,9 +361,16 @@ def run_asr_wer(args):
        # dele = measures["deletions"] / len(ref_list)
        # inse = measures["insertions"] / len(ref_list)

-        wers.append(wer)
+        wer_results.append(
+            {
+                "wav": Path(gen_wav).stem,
+                "truth": raw_truth,
+                "hypo": raw_hypo,
+                "wer": wer,
+            }
+        )

-    return wers
+    return wer_results


 # SIM Evaluation
@@ -381,7 +389,7 @@ def run_sim(args):
        model = model.cuda(device)
    model.eval()

-    sim_list = []
+    sims = []
    for wav1, wav2, truth in tqdm(test_set):
        wav1, sr1 = torchaudio.load(wav1)
        wav2, sr2 = torchaudio.load(wav2)
@@ -400,6 +408,6 @@ def run_sim(args):

        sim = F.cosine_similarity(emb1, emb2)[0].item()
        # print(f"VSim score between two audios: {sim:.4f} (-1.0, 1.0).")
-        sim_list.append(sim)
+        sims.append(sim)

-    return sim_list
+    return sims
@@ -64,6 +64,9 @@ f5-tts_infer-cli \
 # Choose Vocoder
 f5-tts_infer-cli --vocoder_name bigvgan --load_vocoder_from_local --ckpt_file <YOUR_CKPT_PATH, eg:ckpts/F5TTS_Base_bigvgan/model_1250000.pt>
 f5-tts_infer-cli --vocoder_name vocos --load_vocoder_from_local --ckpt_file <YOUR_CKPT_PATH, eg:ckpts/F5TTS_Base/model_1200000.safetensors>
+
+# More instructions
+f5-tts_infer-cli --help
 ```

 And a `.toml` file would help with more flexible usage.
@@ -16,31 +16,34 @@
 <!-- omit in toc -->
 ### Supported Languages
 - [Multilingual](#multilingual)
-    - [F5-TTS Base @ pretrain @ zh \& en](#f5-tts-base--pretrain--zh--en)
+    - [F5-TTS Base @ zh \& en @ F5-TTS](#f5-tts-base--zh--en--f5-tts)
 - [English](#english)
 - [Finnish](#finnish)
-    - [Finnish Common\_Voice Vox\_Populi @ finetune @ fi](#finnish-common_voice-vox_populi--finetune--fi)
+    - [F5-TTS Base @ fi @ AsmoKoskinen](#f5-tts-base--fi--asmokoskinen)
 - [French](#french)
-    - [French LibriVox @ finetune @ fr](#french-librivox--finetune--fr)
+    - [F5-TTS Base @ fr @ RASPIAUDIO](#f5-tts-base--fr--raspiaudio)
+- [Hindi](#hindi)
+    - [F5-TTS Small @ hi @ SPRINGLab](#f5-tts-small--hi--springlab)
 - [Italian](#italian)
-    - [F5-TTS Italian @ finetune @ it](#f5-tts-italian--finetune--it)
+    - [F5-TTS Base @ it @ alien79](#f5-tts-base--it--alien79)
 - [Japanese](#japanese)
-    - [F5-TTS Japanese @ pretrain/finetune @ ja](#f5-tts-japanese--pretrainfinetune--ja)
+    - [F5-TTS Base @ ja @ Jmica](#f5-tts-base--ja--jmica)
 - [Mandarin](#mandarin)
 - [Spanish](#spanish)
-    - [F5-TTS Spanish @ pretrain/finetune @ es](#f5-tts-spanish--pretrainfinetune--es)
+    - [F5-TTS Base @ es @ jpgallegoar](#f5-tts-base--es--jpgallegoar)


 ## Multilingual

-#### F5-TTS Base @ pretrain @ zh & en
+#### F5-TTS Base @ zh & en @ F5-TTS
 |Model|🤗Hugging Face|Data (Hours)|Model License|
 |:---:|:------------:|:-----------:|:-------------:|
 |F5-TTS Base|[ckpt & vocab](https://huggingface.co/SWivid/F5-TTS/tree/main/F5TTS_Base)|[Emilia 95K zh&en](https://huggingface.co/datasets/amphion/Emilia-Dataset/tree/fc71e07)|cc-by-nc-4.0|

 ```bash
-MODEL_CKPT: hf://SWivid/F5-TTS/F5TTS_Base/model_1200000.safetensors
-VOCAB_FILE: hf://SWivid/F5-TTS/F5TTS_Base/vocab.txt
+Model: hf://SWivid/F5-TTS/F5TTS_Base/model_1200000.safetensors
+Vocab: hf://SWivid/F5-TTS/F5TTS_Base/vocab.txt
+Config: {"dim": 1024, "depth": 22, "heads": 16, "ff_mult": 2, "text_dim": 512, "conv_layers": 4}
 ```

 *Other infos, e.g. Author info, Github repo, Link to some sampled results, Usage instruction, Tutorial (Blog, Video, etc.) ...*
@@ -51,27 +54,29 @@ VOCAB_FILE: hf://SWivid/F5-TTS/F5TTS_Base/vocab.txt

 ## Finnish

-#### Finnish Common_Voice Vox_Populi @ finetune @ fi
+#### F5-TTS Base @ fi @ AsmoKoskinen
 |Model|🤗Hugging Face|Data|Model License|
 |:---:|:------------:|:-----------:|:-------------:|
-|F5-TTS Finnish|[ckpt & vocab](https://huggingface.co/AsmoKoskinen/F5-TTS_Finnish_Model)|[Common Voice](https://huggingface.co/datasets/mozilla-foundation/common_voice_17_0), [Vox Populi](https://huggingface.co/datasets/facebook/voxpopuli)|cc-by-nc-4.0|
+|F5-TTS Base|[ckpt & vocab](https://huggingface.co/AsmoKoskinen/F5-TTS_Finnish_Model)|[Common Voice](https://huggingface.co/datasets/mozilla-foundation/common_voice_17_0), [Vox Populi](https://huggingface.co/datasets/facebook/voxpopuli)|cc-by-nc-4.0|

 ```bash
-MODEL_CKPT: hf://AsmoKoskinen/F5-TTS_Finnish_Model/model_common_voice_fi_vox_populi_fi_20241206.safetensors
-VOCAB_FILE: hf://AsmoKoskinen/F5-TTS_Finnish_Model/vocab.txt
+Model: hf://AsmoKoskinen/F5-TTS_Finnish_Model/model_common_voice_fi_vox_populi_fi_20241206.safetensors
+Vocab: hf://AsmoKoskinen/F5-TTS_Finnish_Model/vocab.txt
+Config: {"dim": 1024, "depth": 22, "heads": 16, "ff_mult": 2, "text_dim": 512, "conv_layers": 4}
 ```


 ## French

-#### French LibriVox @ finetune @ fr
+#### F5-TTS Base @ fr @ RASPIAUDIO
 |Model|🤗Hugging Face|Data (Hours)|Model License|
 |:---:|:------------:|:-----------:|:-------------:|
-|F5-TTS French|[ckpt & vocab](https://huggingface.co/RASPIAUDIO/F5-French-MixedSpeakers-reduced)|[LibriVox](https://librivox.org/)|cc-by-nc-4.0|
+|F5-TTS Base|[ckpt & vocab](https://huggingface.co/RASPIAUDIO/F5-French-MixedSpeakers-reduced)|[LibriVox](https://librivox.org/)|cc-by-nc-4.0|

 ```bash
-MODEL_CKPT: hf://RASPIAUDIO/F5-French-MixedSpeakers-reduced/model_last_reduced.pt
-VOCAB_FILE: hf://RASPIAUDIO/F5-French-MixedSpeakers-reduced/vocab.txt
+Model: hf://RASPIAUDIO/F5-French-MixedSpeakers-reduced/model_last_reduced.pt
+Vocab: hf://RASPIAUDIO/F5-French-MixedSpeakers-reduced/vocab.txt
+Config: {"dim": 1024, "depth": 22, "heads": 16, "ff_mult": 2, "text_dim": 512, "conv_layers": 4}
 ```

 - [Online Inference with Hugging Face Space](https://huggingface.co/spaces/RASPIAUDIO/f5-tts_french).
@@ -79,16 +84,34 @@ VOCAB_FILE: hf://RASPIAUDIO/F5-French-MixedSpeakers-reduced/vocab.txt
 - [Discussion about this training can be found here](https://github.com/SWivid/F5-TTS/issues/434).


-## Italian
+## Hindi

-#### F5-TTS Italian @ finetune @ it
-|Model|🤗Hugging Face|Data|Model License|
+#### F5-TTS Small @ hi @ SPRINGLab
+|Model|🤗Hugging Face|Data (Hours)|Model License|
 |:---:|:------------:|:-----------:|:-------------:|
-|F5-TTS Italian|[ckpt & vocab](https://huggingface.co/alien79/F5-TTS-italian)|[ylacombe/cml-tts](https://huggingface.co/datasets/ylacombe/cml-tts) |cc-by-nc-4.0|
+|F5-TTS Small|[ckpt & vocab](https://huggingface.co/SPRINGLab/F5-Hindi-24KHz)|[IndicTTS Hi](https://huggingface.co/datasets/SPRINGLab/IndicTTS-Hindi) & [IndicVoices-R Hi](https://huggingface.co/datasets/SPRINGLab/IndicVoices-R_Hindi) |cc-by-4.0|

 ```bash
-MODEL_CKPT: hf://alien79/F5-TTS-italian/model_159600.safetensors
-VOCAB_FILE: hf://alien79/F5-TTS-italian/vocab.txt
+Model: hf://SPRINGLab/F5-Hindi-24KHz/model_2500000.safetensors
+Vocab: hf://SPRINGLab/F5-Hindi-24KHz/vocab.txt
+Config: {"dim": 768, "depth": 18, "heads": 12, "ff_mult": 2, "text_dim": 512, "conv_layers": 4}
+```
+
+- Authors: SPRING Lab, Indian Institute of Technology, Madras
+- Website: https://asr.iitm.ac.in/
+
+
+## Italian
+
+#### F5-TTS Base @ it @ alien79
+|Model|🤗Hugging Face|Data|Model License|
+|:---:|:------------:|:-----------:|:-------------:|
+|F5-TTS Base|[ckpt & vocab](https://huggingface.co/alien79/F5-TTS-italian)|[ylacombe/cml-tts](https://huggingface.co/datasets/ylacombe/cml-tts) |cc-by-nc-4.0|
+
+```bash
+Model: hf://alien79/F5-TTS-italian/model_159600.safetensors
+Vocab: hf://alien79/F5-TTS-italian/vocab.txt
+Config: {"dim": 1024, "depth": 22, "heads": 16, "ff_mult": 2, "text_dim": 512, "conv_layers": 4}
 ```

 - Trained by [Mithril Man](https://github.com/MithrilMan)
@@ -98,14 +121,15 @@ VOCAB_FILE: hf://alien79/F5-TTS-italian/vocab.txt

 ## Japanese

-#### F5-TTS Japanese @ pretrain/finetune @ ja
+#### F5-TTS Base @ ja @ Jmica
 |Model|🤗Hugging Face|Data (Hours)|Model License|
 |:---:|:------------:|:-----------:|:-------------:|
-|F5-TTS Japanese|[ckpt & vocab](https://huggingface.co/Jmica/F5TTS/tree/main/JA_8500000)|[Emilia 1.7k JA](https://huggingface.co/datasets/amphion/Emilia-Dataset/tree/fc71e07) & [Galgame Dataset 5.4k](https://huggingface.co/datasets/OOPPEENN/Galgame_Dataset)|cc-by-nc-4.0|
+|F5-TTS Base|[ckpt & vocab](https://huggingface.co/Jmica/F5TTS/tree/main/JA_8500000)|[Emilia 1.7k JA](https://huggingface.co/datasets/amphion/Emilia-Dataset/tree/fc71e07) & [Galgame Dataset 5.4k](https://huggingface.co/datasets/OOPPEENN/Galgame_Dataset)|cc-by-nc-4.0|

 ```bash
-MODEL_CKPT: hf://Jmica/F5TTS/JA_8500000/model_8499660.pt
-VOCAB_FILE: hf://Jmica/F5TTS/JA_8500000/vocab_updated.txt
+Model: hf://Jmica/F5TTS/JA_8500000/model_8499660.pt
+Vocab: hf://Jmica/F5TTS/JA_8500000/vocab_updated.txt
+Config: {"dim": 1024, "depth": 22, "heads": 16, "ff_mult": 2, "text_dim": 512, "conv_layers": 4}
 ```


@@ -114,9 +138,9 @@ VOCAB_FILE: hf://Jmica/F5TTS/JA_8500000/vocab_updated.txt

 ## Spanish

-#### F5-TTS Spanish @ pretrain/finetune @ es
+#### F5-TTS Base @ es @ jpgallegoar
 |Model|🤗Hugging Face|Data (Hours)|Model License|
 |:---:|:------------:|:-----------:|:-------------:|
-|F5-TTS Spanish|[ckpt & vocab](https://huggingface.co/jpgallegoar/F5-Spanish)|[Voxpopuli](https://huggingface.co/datasets/facebook/voxpopuli) & Crowdsourced & TEDx, 218 hours|cc0-1.0|
+|F5-TTS Base|[ckpt & vocab](https://huggingface.co/jpgallegoar/F5-Spanish)|[Voxpopuli](https://huggingface.co/datasets/facebook/voxpopuli) & Crowdsourced & TEDx, 218 hours|cc0-1.0|

 - @jpgallegoar [GitHub repo](https://github.com/jpgallegoar/Spanish-F5), Jupyter Notebook and Gradio usage for Spanish model.
@@ -8,4 +8,4 @@ gen_text = "I don't really care what you call me. I've been a silent spectator,
 gen_file = ""
 remove_silence = false
 output_dir = "tests"
-output_file = "infer_cli_out.wav"
+output_file = "infer_cli_basic.wav"
@@ -8,6 +8,7 @@ gen_text = ""
 gen_file = "infer/examples/multi/story.txt"
 remove_silence = true
 output_dir = "tests"
+output_file = "infer_cli_story.wav"

 [voices.town]
 ref_audio = "infer/examples/multi/town.flac"
@@ -2,6 +2,7 @@ import argparse
 import codecs
 import os
 import re
+from datetime import datetime
 from importlib.resources import files
 from pathlib import Path

@@ -9,8 +10,17 @@ import numpy as np
 import soundfile as sf
 import tomli
 from cached_path import cached_path
+from omegaconf import OmegaConf

 from f5_tts.infer.utils_infer import (
+    mel_spec_type,
+    target_rms,
+    cross_fade_duration,
+    nfe_step,
+    cfg_strength,
+    sway_sampling_coef,
+    speed,
+    fix_duration,
    infer_process,
    load_model,
    load_vocoder,
@@ -19,6 +29,7 @@ from f5_tts.infer.utils_infer import (
 )
 from f5_tts.model import DiT, UNetT

+
 parser = argparse.ArgumentParser(
    prog="python3 infer-cli.py",
    description="Commandline interface for E2/F5 TTS with Advanced Batch Processing.",
@@ -27,74 +38,168 @@ parser = argparse.ArgumentParser(
 parser.add_argument(
    "-c",
    "--config",
-    help="Configuration file. Default=infer/examples/basic/basic.toml",
+    type=str,
    default=os.path.join(files("f5_tts").joinpath("infer/examples/basic"), "basic.toml"),
+    help="The configuration file, default see infer/examples/basic/basic.toml",
 )
+
+
+# Note. Not to provide default value here in order to read default from config file
+
 parser.add_argument(
    "-m",
    "--model",
-    help="F5-TTS | E2-TTS",
+    type=str,
+    help="The model name: F5-TTS | E2-TTS",
+)
+parser.add_argument(
+    "-mc",
+    "--model_cfg",
+    type=str,
+    help="The path to F5-TTS model config file .yaml",
 )
 parser.add_argument(
    "-p",
    "--ckpt_file",
-    help="The Checkpoint .pt",
+    type=str,
+    help="The path to model checkpoint .pt, leave blank to use default",
 )
 parser.add_argument(
    "-v",
    "--vocab_file",
-    help="The vocab .txt",
+    type=str,
+    help="The path to vocab file .txt, leave blank to use default",
+)
+parser.add_argument(
+    "-r",
+    "--ref_audio",
+    type=str,
+    help="The reference audio file.",
+)
+parser.add_argument(
+    "-s",
+    "--ref_text",
+    type=str,
+    help="The transcript/subtitle for the reference audio",
 )
-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.",
+    help="The text to make model synthesize a speech",
 )
 parser.add_argument(
    "-f",
    "--gen_file",
    type=str,
-    help="File with text to generate. Ignores --gen_text",
+    help="The file with text to generate, will ignore --gen_text",
 )
 parser.add_argument(
    "-o",
    "--output_dir",
    type=str,
-    help="Path to output folder..",
+    help="The path to output folder",
 )
 parser.add_argument(
    "-w",
    "--output_file",
    type=str,
-    help="Filename of output file..",
+    help="The name of output file",
+)
+parser.add_argument(
+    "--save_chunk",
+    action="store_true",
+    help="To save each audio chunks during inference",
 )
 parser.add_argument(
    "--remove_silence",
-    help="Remove silence.",
+    action="store_true",
+    help="To remove long silence found in ouput",
 )
-parser.add_argument("--vocoder_name", type=str, default="vocos", choices=["vocos", "bigvgan"], help="vocoder name")
 parser.add_argument(
    "--load_vocoder_from_local",
    action="store_true",
-    help="load vocoder from local. Default: ../checkpoints/charactr/vocos-mel-24khz",
+    help="To load vocoder from local dir, default to ../checkpoints/vocos-mel-24khz",
+)
+parser.add_argument(
+    "--vocoder_name",
+    type=str,
+    choices=["vocos", "bigvgan"],
+    help=f"Used vocoder name: vocos | bigvgan, default {mel_spec_type}",
+)
+parser.add_argument(
+    "--target_rms",
+    type=float,
+    help=f"Target output speech loudness normalization value, default {target_rms}",
+)
+parser.add_argument(
+    "--cross_fade_duration",
+    type=float,
+    help=f"Duration of cross-fade between audio segments in seconds, default {cross_fade_duration}",
+)
+parser.add_argument(
+    "--nfe_step",
+    type=int,
+    help=f"The number of function evaluation (denoising steps), default {nfe_step}",
+)
+parser.add_argument(
+    "--cfg_strength",
+    type=float,
+    help=f"Classifier-free guidance strength, default {cfg_strength}",
+)
+parser.add_argument(
+    "--sway_sampling_coef",
+    type=float,
+    help=f"Sway Sampling coefficient, default {sway_sampling_coef}",
 )
 parser.add_argument(
    "--speed",
    type=float,
-    default=1.0,
-    help="Adjust the speed of the audio generation (default: 1.0)",
+    help=f"The speed of the generated audio, default {speed}",
+)
+parser.add_argument(
+    "--fix_duration",
+    type=float,
+    help=f"Fix the total duration (ref and gen audios) in seconds, default {fix_duration}",
 )
 args = parser.parse_args()

+
+# config file
+
 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"]
-gen_file = args.gen_file if args.gen_file else config["gen_file"]
+
+# command-line interface parameters
+
+model = args.model or config.get("model", "F5-TTS")
+model_cfg = args.model_cfg or config.get("model_cfg", str(files("f5_tts").joinpath("configs/F5TTS_Base_train.yaml")))
+ckpt_file = args.ckpt_file or config.get("ckpt_file", "")
+vocab_file = args.vocab_file or config.get("vocab_file", "")
+
+ref_audio = args.ref_audio or config.get("ref_audio", "infer/examples/basic/basic_ref_en.wav")
+ref_text = args.ref_text or config.get("ref_text", "Some call me nature, others call me mother nature.")
+gen_text = args.gen_text or config.get("gen_text", "Here we generate something just for test.")
+gen_file = args.gen_file or config.get("gen_file", "")
+
+output_dir = args.output_dir or config.get("output_dir", "tests")
+output_file = args.output_file or config.get(
+    "output_file", f"infer_cli_{datetime.now().strftime(r'%Y%m%d_%H%M%S')}.wav"
+)
+
+save_chunk = args.save_chunk or config.get("save_chunk", False)
+remove_silence = args.remove_silence or config.get("remove_silence", False)
+load_vocoder_from_local = args.load_vocoder_from_local or config.get("load_vocoder_from_local", False)
+
+vocoder_name = args.vocoder_name or config.get("vocoder_name", mel_spec_type)
+target_rms = args.target_rms or config.get("target_rms", target_rms)
+cross_fade_duration = args.cross_fade_duration or config.get("cross_fade_duration", cross_fade_duration)
+nfe_step = args.nfe_step or config.get("nfe_step", nfe_step)
+cfg_strength = args.cfg_strength or config.get("cfg_strength", cfg_strength)
+sway_sampling_coef = args.sway_sampling_coef or config.get("sway_sampling_coef", sway_sampling_coef)
+speed = args.speed or config.get("speed", speed)
+fix_duration = args.fix_duration or config.get("fix_duration", fix_duration)
+

 # patches for pip pkg user
 if "infer/examples/" in ref_audio:
@@ -107,34 +212,39 @@ if "voices" in config:
        if "infer/examples/" in voice_ref_audio:
            config["voices"][voice]["ref_audio"] = str(files("f5_tts").joinpath(f"{voice_ref_audio}"))

+
+# ignore gen_text if gen_file provided
+
 if gen_file:
    gen_text = codecs.open(gen_file, "r", "utf-8").read()
-output_dir = args.output_dir if args.output_dir else config["output_dir"]
-output_file = args.output_file if args.output_file else config["output_file"]
-model = args.model if args.model else config["model"]
-ckpt_file = args.ckpt_file if args.ckpt_file else ""
-vocab_file = args.vocab_file if args.vocab_file else ""
-remove_silence = args.remove_silence if args.remove_silence else config["remove_silence"]
-speed = args.speed
+
+
+# output path

 wave_path = Path(output_dir) / output_file
 # spectrogram_path = Path(output_dir) / "infer_cli_out.png"
+if save_chunk:
+    output_chunk_dir = os.path.join(output_dir, f"{Path(output_file).stem}_chunks")
+    if not os.path.exists(output_chunk_dir):
+        os.makedirs(output_chunk_dir)
+
+
+# load vocoder

-vocoder_name = args.vocoder_name
-mel_spec_type = args.vocoder_name
 if vocoder_name == "vocos":
    vocoder_local_path = "../checkpoints/vocos-mel-24khz"
 elif vocoder_name == "bigvgan":
    vocoder_local_path = "../checkpoints/bigvgan_v2_24khz_100band_256x"

-vocoder = load_vocoder(vocoder_name=mel_spec_type, is_local=args.load_vocoder_from_local, local_path=vocoder_local_path)
+vocoder = load_vocoder(vocoder_name=vocoder_name, is_local=load_vocoder_from_local, local_path=vocoder_local_path)


-# load models
+# load TTS model
+
 if model == "F5-TTS":
    model_cls = DiT
-    model_cfg = dict(dim=1024, depth=22, heads=16, ff_mult=2, text_dim=512, conv_layers=4)
-    if ckpt_file == "":
+    model_cfg = OmegaConf.load(model_cfg).model.arch
+    if not ckpt_file:  # path not specified, download from repo
        if vocoder_name == "vocos":
            repo_name = "F5-TTS"
            exp_name = "F5TTS_Base"
@@ -148,22 +258,25 @@ if model == "F5-TTS":
            ckpt_file = str(cached_path(f"hf://SWivid/{repo_name}/{exp_name}/model_{ckpt_step}.pt"))

 elif model == "E2-TTS":
-    assert vocoder_name == "vocos", "E2-TTS only supports vocoder vocos"
+    assert args.model_cfg is None, "E2-TTS does not support custom model_cfg yet"
+    assert vocoder_name == "vocos", "E2-TTS only supports vocoder vocos yet"
    model_cls = UNetT
    model_cfg = dict(dim=1024, depth=24, heads=16, ff_mult=4)
-    if ckpt_file == "":
+    if not ckpt_file:  # path not specified, download from repo
        repo_name = "E2-TTS"
        exp_name = "E2TTS_Base"
        ckpt_step = 1200000
        ckpt_file = str(cached_path(f"hf://SWivid/{repo_name}/{exp_name}/model_{ckpt_step}.safetensors"))
        # ckpt_file = f"ckpts/{exp_name}/model_{ckpt_step}.pt"  # .pt | .safetensors; local path

-
 print(f"Using {model}...")
-ema_model = load_model(model_cls, model_cfg, ckpt_file, mel_spec_type=mel_spec_type, vocab_file=vocab_file)
+ema_model = load_model(model_cls, model_cfg, ckpt_file, mel_spec_type=vocoder_name, vocab_file=vocab_file)


-def main_process(ref_audio, ref_text, text_gen, model_obj, mel_spec_type, remove_silence, speed):
+# inference process
+
+
+def main():
    main_voice = {"ref_audio": ref_audio, "ref_text": ref_text}
    if "voices" not in config:
        voices = {"main": main_voice}
@@ -171,16 +284,16 @@ def main_process(ref_audio, ref_text, text_gen, model_obj, mel_spec_type, remove
        voices = config["voices"]
        voices["main"] = main_voice
    for voice in voices:
+        print("Voice:", voice)
+        print("ref_audio ", voices[voice]["ref_audio"])
        voices[voice]["ref_audio"], voices[voice]["ref_text"] = preprocess_ref_audio_text(
            voices[voice]["ref_audio"], voices[voice]["ref_text"]
        )
-        print("Voice:", voice)
-        print("Ref_audio:", voices[voice]["ref_audio"])
-        print("Ref_text:", voices[voice]["ref_text"])
+        print("ref_audio_", voices[voice]["ref_audio"], "\n\n")

    generated_audio_segments = []
    reg1 = r"(?=\[\w+\])"
-    chunks = re.split(reg1, text_gen)
+    chunks = re.split(reg1, gen_text)
    reg2 = r"\[(\w+)\]"
    for text in chunks:
        if not text.strip():
@@ -195,14 +308,35 @@ def main_process(ref_audio, ref_text, text_gen, model_obj, mel_spec_type, remove
            print(f"Voice {voice} not found, using main.")
            voice = "main"
        text = re.sub(reg2, "", text)
-        gen_text = text.strip()
-        ref_audio = voices[voice]["ref_audio"]
-        ref_text = voices[voice]["ref_text"]
+        ref_audio_ = voices[voice]["ref_audio"]
+        ref_text_ = voices[voice]["ref_text"]
+        gen_text_ = text.strip()
        print(f"Voice: {voice}")
-        audio, final_sample_rate, spectragram = infer_process(
-            ref_audio, ref_text, gen_text, model_obj, vocoder, mel_spec_type=mel_spec_type, speed=speed
+        audio_segment, final_sample_rate, spectragram = infer_process(
+            ref_audio_,
+            ref_text_,
+            gen_text_,
+            ema_model,
+            vocoder,
+            mel_spec_type=vocoder_name,
+            target_rms=target_rms,
+            cross_fade_duration=cross_fade_duration,
+            nfe_step=nfe_step,
+            cfg_strength=cfg_strength,
+            sway_sampling_coef=sway_sampling_coef,
+            speed=speed,
+            fix_duration=fix_duration,
        )
-        generated_audio_segments.append(audio)
+        generated_audio_segments.append(audio_segment)
+
+        if save_chunk:
+            if len(gen_text_) > 200:
+                gen_text_ = gen_text_[:200] + " ... "
+            sf.write(
+                os.path.join(output_chunk_dir, f"{len(generated_audio_segments)-1}_{gen_text_}.wav"),
+                audio_segment,
+                final_sample_rate,
+            )

    if generated_audio_segments:
        final_wave = np.concatenate(generated_audio_segments)
@@ -218,9 +352,5 @@ def main_process(ref_audio, ref_text, text_gen, model_obj, mel_spec_type, remove
            print(f.name)


-def main():
-    main_process(ref_audio, ref_text, gen_text, ema_model, mel_spec_type, remove_silence, speed)
-
-
 if __name__ == "__main__":
    main()
@@ -1,6 +1,7 @@
 # ruff: noqa: E402
 # Above allows ruff to ignore E402: module level import not at top of file

+import json
 import re
 import tempfile
 from collections import OrderedDict
@@ -43,6 +44,12 @@ from f5_tts.infer.utils_infer import (
 DEFAULT_TTS_MODEL = "F5-TTS"
 tts_model_choice = DEFAULT_TTS_MODEL

+DEFAULT_TTS_MODEL_CFG = [
+    "hf://SWivid/F5-TTS/F5TTS_Base/model_1200000.safetensors",
+    "hf://SWivid/F5-TTS/F5TTS_Base/vocab.txt",
+    json.dumps(dict(dim=1024, depth=22, heads=16, ff_mult=2, text_dim=512, conv_layers=4)),
+]
+

 # load models

@@ -103,7 +110,15 @@ def generate_response(messages, model, tokenizer):

@gpu_decorator
 def infer(
-    ref_audio_orig, ref_text, gen_text, model, remove_silence, cross_fade_duration=0.15, speed=1, show_info=gr.Info
+    ref_audio_orig,
+    ref_text,
+    gen_text,
+    model,
+    remove_silence,
+    cross_fade_duration=0.15,
+    nfe_step=32,
+    speed=1,
+    show_info=gr.Info,
 ):
    ref_audio, ref_text = preprocess_ref_audio_text(ref_audio_orig, ref_text, show_info=show_info)

@@ -120,7 +135,7 @@ def infer(
        global custom_ema_model, pre_custom_path
        if pre_custom_path != model[1]:
            show_info("Loading Custom TTS model...")
-            custom_ema_model = load_custom(model[1], vocab_path=model[2])
+            custom_ema_model = load_custom(model[1], vocab_path=model[2], model_cfg=model[3])
            pre_custom_path = model[1]
        ema_model = custom_ema_model

@@ -131,6 +146,7 @@ def infer(
        ema_model,
        vocoder,
        cross_fade_duration=cross_fade_duration,
+        nfe_step=nfe_step,
        speed=speed,
        show_info=show_info,
        progress=gr.Progress(),
@@ -184,6 +200,14 @@ with gr.Blocks() as app_tts:
            step=0.1,
            info="Adjust the speed of the audio.",
        )
+        nfe_slider = gr.Slider(
+            label="NFE Steps",
+            minimum=4,
+            maximum=64,
+            value=32,
+            step=2,
+            info="Set the number of denoising steps.",
+        )
        cross_fade_duration_slider = gr.Slider(
            label="Cross-Fade Duration (s)",
            minimum=0.0,
@@ -203,6 +227,7 @@ with gr.Blocks() as app_tts:
        gen_text_input,
        remove_silence,
        cross_fade_duration_slider,
+        nfe_slider,
        speed_slider,
    ):
        audio_out, spectrogram_path, ref_text_out = infer(
@@ -211,8 +236,9 @@ with gr.Blocks() as app_tts:
            gen_text_input,
            tts_model_choice,
            remove_silence,
-            cross_fade_duration_slider,
-            speed_slider,
+            cross_fade_duration=cross_fade_duration_slider,
+            nfe_step=nfe_slider,
+            speed=speed_slider,
        )
        return audio_out, spectrogram_path, gr.update(value=ref_text_out)

@@ -224,6 +250,7 @@ with gr.Blocks() as app_tts:
            gen_text_input,
            remove_silence,
            cross_fade_duration_slider,
+            nfe_slider,
            speed_slider,
        ],
        outputs=[audio_output, spectrogram_output, ref_text_input],
@@ -744,34 +771,38 @@ If you're having issues, try converting your reference audio to WAV or MP3, clip
 """
    )

-    last_used_custom = files("f5_tts").joinpath("infer/.cache/last_used_custom.txt")
+    last_used_custom = files("f5_tts").joinpath("infer/.cache/last_used_custom_model_info.txt")

    def load_last_used_custom():
        try:
-            with open(last_used_custom, "r") as f:
-                return f.read().split(",")
+            custom = []
+            with open(last_used_custom, "r", encoding="utf-8") as f:
+                for line in f:
+                    custom.append(line.strip())
+            return custom
        except FileNotFoundError:
            last_used_custom.parent.mkdir(parents=True, exist_ok=True)
-            return [
-                "hf://SWivid/F5-TTS/F5TTS_Base/model_1200000.safetensors",
-                "hf://SWivid/F5-TTS/F5TTS_Base/vocab.txt",
-            ]
+            return DEFAULT_TTS_MODEL_CFG

    def switch_tts_model(new_choice):
        global tts_model_choice
        if new_choice == "Custom":  # override in case webpage is refreshed
-            custom_ckpt_path, custom_vocab_path = load_last_used_custom()
-            tts_model_choice = ["Custom", custom_ckpt_path, custom_vocab_path]
-            return gr.update(visible=True, value=custom_ckpt_path), gr.update(visible=True, value=custom_vocab_path)
+            custom_ckpt_path, custom_vocab_path, custom_model_cfg = load_last_used_custom()
+            tts_model_choice = ["Custom", custom_ckpt_path, custom_vocab_path, json.loads(custom_model_cfg)]
+            return (
+                gr.update(visible=True, value=custom_ckpt_path),
+                gr.update(visible=True, value=custom_vocab_path),
+                gr.update(visible=True, value=custom_model_cfg),
+            )
        else:
            tts_model_choice = new_choice
-            return gr.update(visible=False), gr.update(visible=False)
+            return gr.update(visible=False), gr.update(visible=False), gr.update(visible=False)

-    def set_custom_model(custom_ckpt_path, custom_vocab_path):
+    def set_custom_model(custom_ckpt_path, custom_vocab_path, custom_model_cfg):
        global tts_model_choice
-        tts_model_choice = ["Custom", custom_ckpt_path, custom_vocab_path]
-        with open(last_used_custom, "w") as f:
-            f.write(f"{custom_ckpt_path},{custom_vocab_path}")
+        tts_model_choice = ["Custom", custom_ckpt_path, custom_vocab_path, json.loads(custom_model_cfg)]
+        with open(last_used_custom, "w", encoding="utf-8") as f:
+            f.write(custom_ckpt_path + "\n" + custom_vocab_path + "\n" + custom_model_cfg + "\n")

    with gr.Row():
        if not USING_SPACES:
@@ -783,34 +814,46 @@ If you're having issues, try converting your reference audio to WAV or MP3, clip
                choices=[DEFAULT_TTS_MODEL, "E2-TTS"], label="Choose TTS Model", value=DEFAULT_TTS_MODEL
            )
        custom_ckpt_path = gr.Dropdown(
-            choices=["hf://SWivid/F5-TTS/F5TTS_Base/model_1200000.safetensors"],
+            choices=[DEFAULT_TTS_MODEL_CFG[0]],
            value=load_last_used_custom()[0],
            allow_custom_value=True,
-            label="MODEL CKPT: local_path | hf://user_id/repo_id/model_ckpt",
+            label="Model: local_path | hf://user_id/repo_id/model_ckpt",
            visible=False,
        )
        custom_vocab_path = gr.Dropdown(
-            choices=["hf://SWivid/F5-TTS/F5TTS_Base/vocab.txt"],
+            choices=[DEFAULT_TTS_MODEL_CFG[1]],
            value=load_last_used_custom()[1],
            allow_custom_value=True,
-            label="VOCAB FILE: local_path | hf://user_id/repo_id/vocab_file",
+            label="Vocab: local_path | hf://user_id/repo_id/vocab_file",
+            visible=False,
+        )
+        custom_model_cfg = gr.Dropdown(
+            choices=[DEFAULT_TTS_MODEL_CFG[2]],
+            value=load_last_used_custom()[2],
+            allow_custom_value=True,
+            label="Config: in a dictionary form",
            visible=False,
        )

    choose_tts_model.change(
        switch_tts_model,
        inputs=[choose_tts_model],
-        outputs=[custom_ckpt_path, custom_vocab_path],
+        outputs=[custom_ckpt_path, custom_vocab_path, custom_model_cfg],
        show_progress="hidden",
    )
    custom_ckpt_path.change(
        set_custom_model,
-        inputs=[custom_ckpt_path, custom_vocab_path],
+        inputs=[custom_ckpt_path, custom_vocab_path, custom_model_cfg],
        show_progress="hidden",
    )
    custom_vocab_path.change(
        set_custom_model,
-        inputs=[custom_ckpt_path, custom_vocab_path],
+        inputs=[custom_ckpt_path, custom_vocab_path, custom_model_cfg],
+        show_progress="hidden",
+    )
+    custom_model_cfg.change(
+        set_custom_model,
+        inputs=[custom_ckpt_path, custom_vocab_path, custom_model_cfg],
        show_progress="hidden",
    )

@@ -105,6 +105,7 @@ class DiT(nn.Module):
        text_dim=None,
        conv_layers=0,
        long_skip_connection=False,
+        checkpoint_activations=False,
    ):
        super().__init__()

@@ -127,6 +128,16 @@ class DiT(nn.Module):
        self.norm_out = AdaLayerNormZero_Final(dim)  # final modulation
        self.proj_out = nn.Linear(dim, mel_dim)

+        self.checkpoint_activations = checkpoint_activations
+
+    def ckpt_wrapper(self, module):
+        # https://github.com/chuanyangjin/fast-DiT/blob/main/models.py
+        def ckpt_forward(*inputs):
+            outputs = module(*inputs)
+            return outputs
+
+        return ckpt_forward
+
    def forward(
        self,
        x: float["b n d"],  # nosied input audio  # noqa: F722
@@ -152,7 +163,10 @@ class DiT(nn.Module):
            residual = x

        for block in self.transformer_blocks:
-            x = block(x, t, mask=mask, rope=rope)
+            if self.checkpoint_activations:
+                x = torch.utils.checkpoint.checkpoint(self.ckpt_wrapper(block), x, t, mask, rope)
+            else:
+                x = block(x, t, mask=mask, rope=rope)

        if self.long_skip_connection is not None:
            x = self.long_skip_connection(torch.cat((x, residual), dim=-1))
@@ -315,7 +315,7 @@ class Trainer:
                    self.scheduler.step()
                    self.optimizer.zero_grad()

-                if self.is_main:
+                if self.is_main and self.accelerator.sync_gradients:
                    self.ema_model.update()

                global_step += 1