init

app.py

@@ -0,0 +1,472 @@
+import gradio as gr
+import torch
+import torchaudio
+import numpy as np
+from pathlib import Path
+import tempfile
+
+# Import the DMOInference class (assuming it's in a file called dmo_inference.py)
+from infer import DMOInference
+
+def initialize_model(student_checkpoint, duration_predictor_checkpoint, model_type, device, cuda_device_id):
+    """Initialize the DMOSpeech 2 model with given checkpoints."""
+    try:
+        model = DMOInference(
+            student_checkpoint_path=student_checkpoint,
+            duration_predictor_path=duration_predictor_checkpoint,
+            device=device,
+            model_type=model_type,
+            tokenizer="pinyin",
+            dataset_name="Emilia_ZH_EN",
+            cuda_device_id=str(cuda_device_id)
+        )
+        return model, "Model initialized successfully!"
+    except Exception as e:
+        return None, f"Error initializing model: {str(e)}"
+
+def generate_speech(
+    model,
+    generation_mode,
+    prompt_audio,
+    prompt_text,
+    target_text,
+    # Duration settings
+    duration_mode,
+    manual_duration,
+    dp_softmax_range,
+    dp_temperature,
+    # Teacher-student settings
+    teacher_steps,
+    teacher_stopping_time,
+    student_start_step,
+    # Advanced settings
+    eta,
+    cfg_strength,
+    sway_coefficient,
+    # Teacher-guided specific
+    tg_switch_time,
+    tg_teacher_steps,
+    tg_student_steps
+):
+    """Generate speech using the selected mode and parameters."""
+    
+    if model is None:
+        return None, "Please initialize the model first!"
+    
+    if prompt_audio is None:
+        return None, "Please upload a reference audio!"
+    
+    if not target_text:
+        return None, "Please enter target text to generate!"
+    
+    try:
+        # Convert prompt_text to None if empty (for ASR)
+        prompt_text = prompt_text.strip() if prompt_text else None
+        
+        # Determine duration
+        if duration_mode == "automatic":
+            duration = None
+        else:
+            duration = int(manual_duration)
+        
+        # Generate based on selected mode
+        if generation_mode == "Student-Only (4 steps)":
+            # Standard DMOSpeech 2 generation
+            generated_wave = model.generate(
+                gen_text=target_text,
+                audio_path=prompt_audio,
+                prompt_text=prompt_text,
+                teacher_steps=0,  # No teacher guidance
+                student_start_step=1,
+                duration=duration,
+                dp_softmax_range=dp_softmax_range,
+                temperature=dp_temperature,
+                eta=eta,
+                cfg_strength=cfg_strength,
+                sway_coefficient=sway_coefficient,
+                verbose=True
+            )
+            
+        elif generation_mode == "Teacher-Student Distillation":
+            # Full teacher-student distillation
+            generated_wave = model.generate(
+                gen_text=target_text,
+                audio_path=prompt_audio,
+                prompt_text=prompt_text,
+                teacher_steps=teacher_steps,
+                teacher_stopping_time=teacher_stopping_time,
+                student_start_step=student_start_step,
+                duration=duration,
+                dp_softmax_range=dp_softmax_range,
+                temperature=dp_temperature,
+                eta=eta,
+                cfg_strength=cfg_strength,
+                sway_coefficient=sway_coefficient,
+                verbose=True
+            )
+            
+        elif generation_mode == "Teacher-Only":
+            # Teacher-only generation
+            generated_wave = model.generate_teacher_only(
+                gen_text=target_text,
+                audio_path=prompt_audio,
+                prompt_text=prompt_text,
+                teacher_steps=teacher_steps,
+                duration=duration,
+                eta=eta,
+                cfg_strength=cfg_strength,
+                sway_coefficient=sway_coefficient
+            )
+            
+        elif generation_mode == "Teacher-Guided Sampling":
+            # Implement teacher-guided sampling
+            # This would require implementing the teacher-guided sampling algorithm
+            # For now, we'll use the regular generation with specific parameters
+            total_teacher_steps = tg_teacher_steps
+            
+            generated_wave = model.generate(
+                gen_text=target_text,
+                audio_path=prompt_audio,
+                prompt_text=prompt_text,
+                teacher_steps=total_teacher_steps,
+                teacher_stopping_time=tg_switch_time,
+                student_start_step=1,
+                duration=duration,
+                dp_softmax_range=dp_softmax_range,
+                temperature=dp_temperature,
+                eta=eta,
+                cfg_strength=cfg_strength,
+                sway_coefficient=sway_coefficient,
+                verbose=True
+            )
+        
+        # Save generated audio
+        with tempfile.NamedTemporaryFile(delete=False, suffix=".wav") as tmp_file:
+            output_path = tmp_file.name
+            
+        # Convert to tensor and save
+        if isinstance(generated_wave, np.ndarray):
+            generated_wave = torch.from_numpy(generated_wave)
+        
+        if generated_wave.dim() == 1:
+            generated_wave = generated_wave.unsqueeze(0)
+            
+        torchaudio.save(output_path, generated_wave, 24000)
+        
+        return output_path, "Speech generated successfully!"
+        
+    except Exception as e:
+        return None, f"Error generating speech: {str(e)}"
+
+def predict_duration_only(
+    model,
+    prompt_audio,
+    prompt_text,
+    target_text,
+    dp_softmax_range,
+    dp_temperature
+):
+    """Predict duration for the target text."""
+    if model is None:
+        return "Please initialize the model first!"
+    
+    if prompt_audio is None:
+        return "Please upload a reference audio!"
+    
+    if not target_text:
+        return "Please enter target text!"
+    
+    try:
+        prompt_text = prompt_text.strip() if prompt_text else None
+        
+        predicted_duration = model.predict_duration(
+            pmt_wav_path=prompt_audio,
+            tar_text=target_text,
+            pmt_text=prompt_text,
+            dp_softmax_range=dp_softmax_range,
+            temperature=dp_temperature
+        )
+        
+        return f"Predicted duration: {predicted_duration} frames (~{predicted_duration/100:.2f} seconds)"
+        
+    except Exception as e:
+        return f"Error predicting duration: {str(e)}"
+
+# Create Gradio interface
+with gr.Blocks(title="DMOSpeech 2: Advanced Zero-Shot TTS") as demo:
+    gr.Markdown("""
+    # DMOSpeech 2: Reinforcement Learning for Duration Prediction in Metric-Optimized Speech Synthesis
+    
+    This demo showcases DMOSpeech 2, which features:
+    - **Direct metric optimization** for speaker similarity and intelligibility
+    - **RL-optimized duration prediction** for better speech quality
+    - **Teacher-guided sampling** for improved diversity
+    - **Efficient 4-step generation** while maintaining high quality
+    """)
+    
+    # Model state
+    model_state = gr.State(None)
+    
+    with gr.Tab("Model Setup"):
+        gr.Markdown("### Initialize Model")
+        with gr.Row():
+            student_checkpoint = gr.Textbox(
+                label="Student Model Checkpoint Path",
+                placeholder="/path/to/student_checkpoint.pt"
+            )
+            duration_checkpoint = gr.Textbox(
+                label="Duration Predictor Checkpoint Path",
+                placeholder="/path/to/duration_predictor.pt"
+            )
+        
+        with gr.Row():
+            model_type = gr.Dropdown(
+                choices=["F5TTS_Base", "E2TTS_Base"],
+                value="F5TTS_Base",
+                label="Model Type"
+            )
+            device = gr.Dropdown(
+                choices=["cuda", "cpu"],
+                value="cuda",
+                label="Device"
+            )
+            cuda_device_id = gr.Number(
+                value=0,
+                label="CUDA Device ID",
+                precision=0
+            )
+        
+        init_button = gr.Button("Initialize Model", variant="primary")
+        init_status = gr.Textbox(label="Initialization Status", interactive=False)
+        
+    with gr.Tab("Speech Generation"):
+        with gr.Row():
+            with gr.Column(scale=1):
+                gr.Markdown("### Input Settings")
+                
+                prompt_audio = gr.Audio(
+                    label="Reference Audio",
+                    type="filepath",
+                    sources=["upload", "microphone"]
+                )
+                
+                prompt_text = gr.Textbox(
+                    label="Reference Text (optional - will use ASR if empty)",
+                    placeholder="The text spoken in the reference audio..."
+                )
+                
+                target_text = gr.Textbox(
+                    label="Target Text to Generate",
+                    placeholder="Enter the text you want to synthesize...",
+                    lines=3
+                )
+                
+                generation_mode = gr.Radio(
+                    choices=[
+                        "Student-Only (4 steps)",
+                        "Teacher-Student Distillation",
+                        "Teacher-Only",
+                        "Teacher-Guided Sampling"
+                    ],
+                    value="Student-Only (4 steps)",
+                    label="Generation Mode"
+                )
+                
+            with gr.Column(scale=1):
+                gr.Markdown("### Duration Settings")
+                
+                duration_mode = gr.Radio(
+                    choices=["automatic", "manual"],
+                    value="automatic",
+                    label="Duration Mode"
+                )
+                
+                manual_duration = gr.Slider(
+                    minimum=100,
+                    maximum=3000,
+                    value=500,
+                    step=10,
+                    label="Manual Duration (frames)",
+                    visible=False
+                )
+                
+                dp_softmax_range = gr.Slider(
+                    minimum=0.1,
+                    maximum=1.0,
+                    value=0.7,
+                    step=0.1,
+                    label="Duration Predictor Softmax Range"
+                )
+                
+                dp_temperature = gr.Slider(
+                    minimum=0.0,
+                    maximum=2.0,
+                    value=0.0,
+                    step=0.1,
+                    label="Duration Predictor Temperature (0=argmax)"
+                )
+                
+                predict_duration_btn = gr.Button("Predict Duration Only")
+                duration_output = gr.Textbox(label="Predicted Duration", interactive=False)
+        
+        with gr.Accordion("Advanced Settings", open=False):
+            with gr.Tab("Teacher-Student Settings"):
+                teacher_steps = gr.Slider(
+                    minimum=0,
+                    maximum=32,
+                    value=16,
+                    step=1,
+                    label="Teacher Steps"
+                )
+                
+                teacher_stopping_time = gr.Slider(
+                    minimum=0.0,
+                    maximum=1.0,
+                    value=0.07,
+                    step=0.01,
+                    label="Teacher Stopping Time"
+                )
+                
+                student_start_step = gr.Slider(
+                    minimum=1,
+                    maximum=4,
+                    value=1,
+                    step=1,
+                    label="Student Start Step"
+                )
+                
+            with gr.Tab("Sampling Settings"):
+                eta = gr.Slider(
+                    minimum=0.0,
+                    maximum=1.0,
+                    value=1.0,
+                    step=0.1,
+                    label="Eta (Stochasticity: 0=DDIM, 1=DDPM)"
+                )
+                
+                cfg_strength = gr.Slider(
+                    minimum=0.0,
+                    maximum=5.0,
+                    value=2.0,
+                    step=0.1,
+                    label="CFG Strength"
+                )
+                
+                sway_coefficient = gr.Slider(
+                    minimum=-2.0,
+                    maximum=2.0,
+                    value=-1.0,
+                    step=0.1,
+                    label="Sway Sampling Coefficient"
+                )
+                
+            with gr.Tab("Teacher-Guided Settings"):
+                tg_switch_time = gr.Slider(
+                    minimum=0.1,
+                    maximum=0.5,
+                    value=0.25,
+                    step=0.05,
+                    label="Switch Time (when to transition to student)"
+                )
+                
+                tg_teacher_steps = gr.Slider(
+                    minimum=6,
+                    maximum=20,
+                    value=14,
+                    step=1,
+                    label="Teacher Steps"
+                )
+                
+                tg_student_steps = gr.Slider(
+                    minimum=1,
+                    maximum=4,
+                    value=2,
+                    step=1,
+                    label="Student Steps"
+                )
+        
+        generate_button = gr.Button("Generate Speech", variant="primary")
+        
+        with gr.Row():
+            output_audio = gr.Audio(label="Generated Speech", type="filepath")
+            generation_status = gr.Textbox(label="Generation Status", interactive=False)
+    
+    with gr.Tab("Examples & Info"):
+        gr.Markdown("""
+        ### Usage Tips:
+        
+        1. **Generation Modes:**
+           - **Student-Only (4 steps)**: Fastest, uses the distilled model with direct metric optimization
+           - **Teacher-Student Distillation**: Uses teacher guidance for initial steps
+           - **Teacher-Only**: Full quality but slower (32 steps)
+           - **Teacher-Guided Sampling**: Best balance of quality and diversity
+        
+        2. **Duration Settings:**
+           - **Automatic**: Uses RL-optimized duration predictor
+           - **Manual**: Specify exact duration in frames (100 frames ≈ 1 second)
+        
+        3. **Advanced Parameters:**
+           - **Eta**: Controls sampling stochasticity (0 = deterministic, 1 = fully stochastic)
+           - **CFG Strength**: Higher values = stronger adherence to text
+           - **Sway Coefficient**: Negative values focus on early denoising steps
+        
+        ### Key Features:
+        - ✅ 5× faster than teacher model
+        - ✅ Better WER and speaker similarity
+        - ✅ RL-optimized duration prediction
+        - ✅ Maintains prosodic diversity with teacher-guided sampling
+        """)
+    
+    # Event handlers
+    duration_mode.change(
+        lambda x: gr.update(visible=(x == "manual")),
+        inputs=[duration_mode],
+        outputs=[manual_duration]
+    )
+    
+    init_button.click(
+        lambda sc, dc, mt, d, cid: initialize_model(sc, dc, mt, d, cid),
+        inputs=[student_checkpoint, duration_checkpoint, model_type, device, cuda_device_id],
+        outputs=[model_state, init_status]
+    )
+    
+    generate_button.click(
+        generate_speech,
+        inputs=[
+            model_state,
+            generation_mode,
+            prompt_audio,
+            prompt_text,
+            target_text,
+            duration_mode,
+            manual_duration,
+            dp_softmax_range,
+            dp_temperature,
+            teacher_steps,
+            teacher_stopping_time,
+            student_start_step,
+            eta,
+            cfg_strength,
+            sway_coefficient,
+            tg_switch_time,
+            tg_teacher_steps,
+            tg_student_steps
+        ],
+        outputs=[output_audio, generation_status]
+    )
+    
+    predict_duration_btn.click(
+        predict_duration_only,
+        inputs=[
+            model_state,
+            prompt_audio,
+            prompt_text,
+            target_text,
+            dp_softmax_range,
+            dp_temperature
+        ],
+        outputs=[duration_output]
+    )
+
+if __name__ == "__main__":
+    demo.launch(share=True)

ctcmodel.py

@@ -0,0 +1,221 @@
+from torch import nn
+import torch 
+import copy
+
+from pathlib import Path
+from torchaudio.models import Conformer
+
+
+from f5_tts.model.utils import default
+from f5_tts.model.utils import exists
+from f5_tts.model.utils import list_str_to_idx
+from f5_tts.model.utils import list_str_to_tensor
+from f5_tts.model.utils import lens_to_mask
+from f5_tts.model.utils import mask_from_frac_lengths
+
+
+from f5_tts.model.utils import (
+    default,
+    exists,
+    list_str_to_idx,
+    list_str_to_tensor,
+    lens_to_mask,
+    mask_from_frac_lengths,
+)
+
+class ResBlock(nn.Module):
+    def __init__(self, hidden_dim, n_conv=3, dropout_p=0.2):
+        super().__init__()
+        self._n_groups = 8
+        self.blocks = nn.ModuleList([
+            self._get_conv(hidden_dim, dilation=3**i, dropout_p=dropout_p)
+            for i in range(n_conv)])
+
+
+    def forward(self, x):
+        for block in self.blocks:
+            res = x
+            x = block(x)
+            x += res
+        return x
+
+    def _get_conv(self, hidden_dim, dilation, dropout_p=0.2):
+        layers = [
+            nn.Conv1d(hidden_dim, hidden_dim, kernel_size=3, padding=dilation, dilation=dilation),
+            nn.ReLU(),
+            nn.GroupNorm(num_groups=self._n_groups, num_channels=hidden_dim),
+            nn.Dropout(p=dropout_p),
+            nn.Conv1d(hidden_dim, hidden_dim, kernel_size=3, padding=1, dilation=1),
+            nn.ReLU(),
+            nn.Dropout(p=dropout_p)
+        ]
+        return nn.Sequential(*layers)
+
+
+class ConformerCTC(nn.Module):
+    def __init__(self,
+                 vocab_size,
+                 mel_dim=100, 
+                 num_heads=8, 
+                 d_hid=512, 
+                 nlayers=6):
+        super().__init__()
+        
+        self.mel_proj = nn.Conv1d(mel_dim, d_hid, kernel_size=3, padding=1)
+        
+        self.d_hid = d_hid
+        
+        self.resblock1 = nn.Sequential(
+                ResBlock(d_hid),
+                nn.GroupNorm(num_groups=1, num_channels=d_hid)
+            )
+        
+        self.resblock2 = nn.Sequential(
+                ResBlock(d_hid),
+                nn.GroupNorm(num_groups=1, num_channels=d_hid)
+            )
+        
+
+        self.conf_pre = torch.nn.ModuleList(
+            [Conformer(
+             input_dim=d_hid,
+             num_heads=num_heads,
+             ffn_dim=d_hid * 2,
+             num_layers=1,
+             depthwise_conv_kernel_size=15,
+             use_group_norm=True,)
+                for _ in range(nlayers // 2)
+            ]
+        )
+        
+        self.conf_after = torch.nn.ModuleList(
+            [Conformer(
+             input_dim=d_hid,
+             num_heads=num_heads,
+             ffn_dim=d_hid * 2,
+             num_layers=1,
+             depthwise_conv_kernel_size=7,
+             use_group_norm=True,)
+                for _ in range(nlayers // 2)
+            ]
+        )
+
+        self.out = nn.Linear(d_hid, 1 + vocab_size) # 1 for blank
+
+        self.ctc_loss = nn.CTCLoss(blank=vocab_size, zero_infinity=True).cuda()
+
+                
+    def forward(self, latent, text=None, text_lens=None):
+        layers = []
+
+        x = self.mel_proj(latent.transpose(-1, -2)).transpose(-1, -2)
+
+        x = x.transpose(1, 2)
+        layers.append(nn.functional.avg_pool1d(x, 4))
+        # x = x.transpose(1, 2)
+
+        x = self.resblock1(x)
+        x = nn.functional.avg_pool1d(x, 2)
+        layers.append(nn.functional.avg_pool1d(x, 2))
+        x = self.resblock2(x)
+        x = nn.functional.avg_pool1d(x, 2)
+        layers.append(x)
+
+        x = x.transpose(1, 2)
+
+        batch_size, time_steps, _ = x.shape
+        # Create a dummy lengths tensor (all sequences are assumed to be full length).
+        input_lengths = torch.full((batch_size,), time_steps, device=x.device, dtype=torch.int64)
+
+        for layer in (self.conf_pre):
+            x, _ = layer(x, input_lengths)
+            layers.append(x.transpose(1, 2))
+
+        for layer in (self.conf_after):
+            x, _ = layer(x, input_lengths)
+            layers.append(x.transpose(1, 2))
+
+        x = self.out(x)
+
+        if text_lens is not None and text is not None:
+            loss = self.ctc_loss(x.log_softmax(dim=2).transpose(0, 1), text, input_lengths, text_lens)
+            return x, layers, loss
+        else:
+            return x, layers
+
+
+if __name__ == "__main__":
+    from f5_tts.model.utils import get_tokenizer
+
+
+    bsz = 16
+    
+    tokenizer = "pinyin"  # 'pinyin', 'char', or 'custom'
+    tokenizer_path = None  # if tokenizer = 'custom', define the path to the tokenizer you want to use (should be vocab.txt)
+    dataset_name = "Emilia_ZH_EN"
+    if tokenizer == "custom":
+        tokenizer_path = tokenizer_path
+    else:
+        tokenizer_path = dataset_name
+    vocab_char_map, vocab_size = get_tokenizer(tokenizer_path, tokenizer)
+    
+    model = ConformerCTC(vocab_size, mel_dim=80, num_heads=8, d_hid=512, nlayers=6).cuda()
+    
+    text = ["hello world"] * bsz
+    lens = torch.randint(1, 1000, (bsz,)).cuda()
+    inp = torch.randn(bsz, lens.max(), 80).cuda()
+    
+    batch, seq_len, dtype, device = *inp.shape[:2], inp.dtype, inp.device
+    
+    # handle text as string
+    text_lens = torch.tensor([len(t) for t in text], device=device)
+    if isinstance(text, list):
+        if exists(vocab_char_map):
+            text = list_str_to_idx(text, vocab_char_map).to(device)
+        else:
+            text = list_str_to_tensor(text).to(device)
+        assert text.shape[0] == batch
+
+    # lens and mask
+    if not exists(lens):
+        lens = torch.full((batch,), seq_len, device=device)
+
+    out, layers, loss = model(inp, text_lens)
+
+    print(out.shape)
+    print(out)
+    print(len(layers))
+    print(torch.stack(layers, axis=1).shape)
+    print(loss)
+
+    probs = out.softmax(dim=2)  # Convert logits to probabilities
+
+    # Greedy decoding
+    best_path = torch.argmax(probs, dim=2)
+
+    decoded_sequences = []
+    blank_idx = vocab_size
+
+    char_vocab_map = list(vocab_char_map.keys())
+
+
+    for batch in best_path:
+        decoded_sequence = []
+        previous_token = None
+
+        for token in batch:
+            if token != previous_token:  # Collapse repeated tokens
+                if token != blank_idx:  # Ignore blank tokens
+                    decoded_sequence.append(token.item())
+            previous_token = token
+
+        decoded_sequences.append(decoded_sequence)
+
+    # Convert token indices to characters
+    decoded_texts = [''.join([char_vocab_map[token] for token in sequence]) for sequence in decoded_sequences]
+    gt_texts = []
+    for i in range(text_lens.size(0)):
+        gt_texts.append(''.join([char_vocab_map[token] for token in text[i, :text_lens[i]]]))
+    
+    print(decoded_texts)
+    print(gt_texts)

demo.ipynb

@@ -0,0 +1,294 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "id": "93caa9f6",
+   "metadata": {},
+   "source": [
+    "## Basic Inference (teacher-guided, 8 steps)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "72e5ced3",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from infer import DMOInference\n",
+    "import IPython.display as ipd\n",
+    "import torchaudio\n",
+    "import time\n",
+    "\n",
+    "# Initialize the model\n",
+    "tts = DMOInference(\n",
+    "    student_checkpoint_path=\"../ckpts/model_85000.pt\", \n",
+    "    duration_predictor_path=\"../ckpts/model_1500.pt\",\n",
+    "    device=\"cuda\",\n",
+    "    model_type=\"F5TTS_Base\"\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "bc5b2634",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "prompt_audio = \"f5_tts/infer/examples/basic/basic_ref_en.wav\"\n",
+    "\n",
+    "ref_text = \"Some call me nature, others call me mother nature.\"\n",
+    "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.\"\n",
+    "\n",
+    "start_time = time.time()\n",
+    "# Generate with default settings\n",
+    "generated_audio = tts.generate(\n",
+    "    gen_text=gen_text,\n",
+    "    audio_path=prompt_audio,\n",
+    "    prompt_text=ref_text)\n",
+    "end_time = time.time()\n",
+    "\n",
+    "processing_time = end_time - start_time\n",
+    "audio_duration = generated_audio.shape[-1] / 24000\n",
+    "rtf = processing_time / audio_duration\n",
+    "\n",
+    "print('\\n--------\\n')\n",
+    "print('Prompt Audio: ')\n",
+    "display(ipd.Audio(prompt_audio, rate=24000))\n",
+    "print('Generated Audio: ')\n",
+    "display(ipd.Audio(generated_audio, rate=24000))\n",
+    "\n",
+    "print(f\"  RTF: {rtf:.2f}x ({1/rtf:.2f}x speed)\")\n",
+    "print(f\"  Processing: {processing_time:.2f}s for {audio_duration:.2f}s audio\")\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "d65b4bde",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "prompt_audio = \"f5_tts/infer/examples/basic/basic_ref_zh.wav\"\n",
+    "\n",
+    "ref_text = \"对，这就是我，万人敬仰的太乙真人。\"\n",
+    "gen_text = '突然，身边一阵笑声。我看着他们，意气风发地挺直了胸膛，甩了甩那稍显肉感的双臂，轻笑道：\"我身上的肉，是为了掩饰我爆棚的魅力，否则，岂不吓坏了你们呢？\"'\n",
+    "\n",
+    "start_time = time.time()\n",
+    "# Generate with default settings\n",
+    "generated_audio = tts.generate(\n",
+    "    gen_text=gen_text,\n",
+    "    audio_path=prompt_audio,\n",
+    "    prompt_text=ref_text\n",
+    ")\n",
+    "end_time = time.time()\n",
+    "\n",
+    "processing_time = end_time - start_time\n",
+    "audio_duration = generated_audio.shape[-1] / 24000\n",
+    "rtf = processing_time / audio_duration\n",
+    "\n",
+    "print('\\n--------\\n')\n",
+    "print('Prompt Audio: ')\n",
+    "display(ipd.Audio(prompt_audio, rate=24000))\n",
+    "print('Generated Audio: ')\n",
+    "display(ipd.Audio(generated_audio, rate=24000))\n",
+    "\n",
+    "print(f\"  RTF: {rtf:.2f}x ({1/rtf:.2f}x speed)\")\n",
+    "print(f\"  Processing: {processing_time:.2f}s for {audio_duration:.2f}s audio\")\n"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "a0b8df0f",
+   "metadata": {},
+   "source": [
+    "## Comparision between different sampling configurations"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "3962e78c",
+   "metadata": {},
+   "source": [
+    "#### Student only (4 steps)\n",
+    "\n",
+    "Need to set `teacher_steps` and `student_start_step` to 0 to enable full student sampling."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "53a232a1",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "prompt_audio = \"f5_tts/infer/examples/basic/basic_ref_zh.wav\"\n",
+    "\n",
+    "ref_text = \"对，这就是我，万人敬仰的太乙真人。\"\n",
+    "gen_text = '突然，身边一阵笑声。我看着他们，意气风发地挺直了胸膛，甩了甩那稍显肉感的双臂，轻笑道：\"我身上的肉，是为了掩饰我爆棚的魅力，否则，岂不吓坏了你们呢？\"'\n",
+    "\n",
+    "start_time = time.time()\n",
+    "# Generate with default settings\n",
+    "generated_audio = tts.generate(\n",
+    "    gen_text=gen_text,\n",
+    "    audio_path=prompt_audio,\n",
+    "    prompt_text=ref_text,\n",
+    "    teacher_steps=0, # set this to 0 for no teachr sampling\n",
+    "    student_start_step=0, # set this to 0 for full student sampling\n",
+    ")\n",
+    "end_time = time.time()\n",
+    "\n",
+    "processing_time = end_time - start_time\n",
+    "audio_duration = generated_audio.shape[-1] / 24000\n",
+    "rtf = processing_time / audio_duration\n",
+    "\n",
+    "print('\\n--------\\n')\n",
+    "print('Prompt Audio: ')\n",
+    "display(ipd.Audio(prompt_audio, rate=24000))\n",
+    "print('Generated Audio: ')\n",
+    "display(ipd.Audio(generated_audio, rate=24000))\n",
+    "\n",
+    "print(f\"  RTF: {rtf:.2f}x ({1/rtf:.2f}x speed)\")\n",
+    "print(f\"  Processing: {processing_time:.2f}s for {audio_duration:.2f}s audio\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "5cb24808",
+   "metadata": {},
+   "source": [
+    "#### More teacher steps (16 steps)\n",
+    "\n",
+    "Now we use 14 steps from the teacher and 2 steps from the student to have higher diversity (16 steps total)."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "10b3620e",
+   "metadata": {
+    "scrolled": false
+   },
+   "outputs": [],
+   "source": [
+    "prompt_audio = \"f5_tts/infer/examples/basic/basic_ref_zh.wav\"\n",
+    "\n",
+    "ref_text = \"对，这就是我，万人敬仰的太乙真人。\"\n",
+    "gen_text = '突然，身边一阵笑声。我看着他们，意气风发地挺直了胸膛，甩了甩那稍显肉感的双臂，轻笑道：\"我身上的肉，是为了掩饰我爆棚的魅力，否则，岂不吓坏了你们呢？\"'\n",
+    "\n",
+    "start_time = time.time()\n",
+    "# Generate with default settings\n",
+    "generated_audio = tts.generate(\n",
+    "    gen_text=gen_text,\n",
+    "    audio_path=prompt_audio,\n",
+    "    prompt_text=ref_text,\n",
+    "    teacher_steps=24, \n",
+    "    teacher_stopping_time=0.3, # 0.25 means students go for the last two steps (0.26ish, 0.6ish)\n",
+    "    student_start_step=2, # only two steps for students\n",
+    "    verbose=True # see the number of steps used\n",
+    ")\n",
+    "end_time = time.time()\n",
+    "\n",
+    "processing_time = end_time - start_time\n",
+    "audio_duration = generated_audio.shape[-1] / 24000\n",
+    "rtf = processing_time / audio_duration\n",
+    "\n",
+    "print('\\n--------\\n')\n",
+    "print('Prompt Audio: ')\n",
+    "display(ipd.Audio(prompt_audio, rate=24000))\n",
+    "print('Generated Audio: ')\n",
+    "display(ipd.Audio(generated_audio, rate=24000))\n",
+    "\n",
+    "print(f\"  RTF: {rtf:.2f}x ({1/rtf:.2f}x speed)\")\n",
+    "print(f\"  Processing: {processing_time:.2f}s for {audio_duration:.2f}s audio\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "70fd3cb1",
+   "metadata": {},
+   "source": [
+    "#### Stochastic duration \n",
+    "\n",
+    "Introduce even more diversity by adding randomness to the duration"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "bb39d323",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "prompt_audio = \"f5_tts/infer/examples/basic/basic_ref_zh.wav\"\n",
+    "\n",
+    "ref_text = \"对，这就是我，万人敬仰的太乙真人。\"\n",
+    "gen_text = '突然，身边一阵笑声。我看着他们，意气风发地挺直了胸膛，甩了甩那稍显肉感的双臂，轻笑道：\"我身上的肉，是为了掩饰我爆棚的魅力，否则，岂不吓坏了你们呢？\"'\n",
+    "\n",
+    "start_time = time.time()\n",
+    "# Generate with default settings\n",
+    "generated_audio = tts.generate(\n",
+    "    gen_text=gen_text,\n",
+    "    audio_path=prompt_audio,\n",
+    "    prompt_text=ref_text,\n",
+    "    teacher_steps=24, \n",
+    "    teacher_stopping_time=0.25, # 0.25 means students go for the last two steps (0.26ish, 0.6ish)\n",
+    "    student_start_step=2, # only two steps for students\n",
+    "    temperature=0.8, # set some temperature for duration sampling \n",
+    ")\n",
+    "end_time = time.time()\n",
+    "\n",
+    "processing_time = end_time - start_time\n",
+    "audio_duration = generated_audio.shape[-1] / 24000\n",
+    "rtf = processing_time / audio_duration\n",
+    "\n",
+    "print('\\n--------\\n')\n",
+    "print('Prompt Audio: ')\n",
+    "display(ipd.Audio(prompt_audio, rate=24000))\n",
+    "print('Generated Audio: ')\n",
+    "display(ipd.Audio(generated_audio, rate=24000))\n",
+    "\n",
+    "print(f\"  RTF: {rtf:.2f}x ({1/rtf:.2f}x speed)\")\n",
+    "print(f\"  Processing: {processing_time:.2f}s for {audio_duration:.2f}s audio\")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "7f109ed3",
+   "metadata": {},
+   "outputs": [],
+   "source": []
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "1cc13360",
+   "metadata": {},
+   "outputs": [],
+   "source": []
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "Python 3 (ipykernel)",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.10.12"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}

discriminator_conformer.py

@@ -0,0 +1,202 @@
+# part of the code is borrowed from https://github.com/lawlict/ECAPA-TDNN
+
+from __future__ import annotations
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import torchaudio.transforms as trans
+from pathlib import Path
+from torchaudio.models import Conformer
+
+from f5_tts.model.utils import (
+    default,
+    exists,
+    list_str_to_idx,
+    list_str_to_tensor,
+    lens_to_mask,
+    mask_from_frac_lengths,
+)
+
+class ResBlock(nn.Module):
+    def __init__(self, hidden_dim, n_conv=3, dropout_p=0.2):
+        super().__init__()
+        self._n_groups = 8
+        self.blocks = nn.ModuleList([
+            self._get_conv(hidden_dim, dilation=3**i, dropout_p=dropout_p)
+            for i in range(n_conv)])
+
+
+    def forward(self, x):
+        for block in self.blocks:
+            res = x
+            x = block(x)
+            x += res
+        return x
+
+    def _get_conv(self, hidden_dim, dilation, dropout_p=0.2):
+        layers = [
+            nn.Conv1d(hidden_dim, hidden_dim, kernel_size=3, padding=dilation, dilation=dilation),
+            nn.ReLU(),
+            nn.GroupNorm(num_groups=self._n_groups, num_channels=hidden_dim),
+            nn.Dropout(p=dropout_p),
+            nn.Conv1d(hidden_dim, hidden_dim, kernel_size=3, padding=1, dilation=1),
+            nn.ReLU(),
+            nn.Dropout(p=dropout_p)
+        ]
+        return nn.Sequential(*layers)
+
+class ConformerDiscirminator(nn.Module):
+    def __init__(self, input_dim, channels=512, num_layers=3, num_heads=8, depthwise_conv_kernel_size=15, use_group_norm=True):
+        super().__init__()
+        
+        self.input_layer = nn.Conv1d(input_dim, channels, kernel_size=3, padding=1)
+
+        self.resblock1 = nn.Sequential(
+                ResBlock(channels),
+                nn.GroupNorm(num_groups=1, num_channels=channels)
+            )
+        
+        self.resblock2 = nn.Sequential(
+                ResBlock(channels),
+                nn.GroupNorm(num_groups=1, num_channels=channels)
+            )
+
+        self.conformer1 = Conformer(**{"input_dim": channels, 
+                "num_heads": num_heads, 
+                "ffn_dim": channels * 2, 
+                "num_layers": 1, 
+                "depthwise_conv_kernel_size": depthwise_conv_kernel_size // 2,
+                "use_group_norm": use_group_norm})
+
+        self.conformer2 = Conformer(**{"input_dim": channels, 
+                "num_heads": num_heads, 
+                "ffn_dim": channels * 2, 
+                "num_layers": num_layers - 1, 
+                "depthwise_conv_kernel_size": depthwise_conv_kernel_size,
+                "use_group_norm": use_group_norm})
+        
+        self.linear = nn.Conv1d(channels, 1, kernel_size=1)
+
+    def forward(self, x):
+        # x = torch.stack(x, dim=1).transpose(-1, -2).flatten(start_dim=1, end_dim=2)
+        x = torch.cat(x, dim=-1)
+        x = x.transpose(1, 2)
+
+        x = self.input_layer(x)
+
+        x = self.resblock1(x)
+        x = nn.functional.avg_pool1d(x, 2)
+        x = self.resblock2(x)
+        x = nn.functional.avg_pool1d(x, 2)
+        
+        # Transpose to (B, T, C) for the conformer.
+        x = x.transpose(1, 2)
+        batch_size, time_steps, _ = x.shape
+        # Create a dummy lengths tensor (all sequences are assumed to be full length).
+        lengths = torch.full((batch_size,), time_steps, device=x.device, dtype=torch.int64)
+        # The built-in Conformer returns (output, output_lengths); we discard lengths.
+
+        x, _ = self.conformer1(x, lengths)
+        x, _ = self.conformer2(x, lengths)
+        # Transpose back to (B, C, T).
+        x = x.transpose(1, 2)
+
+        # out = self.bn(self.pooling(out))
+        out = self.linear(x).squeeze(1)
+
+        return out
+
+if __name__ == "__main__":
+    from f5_tts.model.utils import get_tokenizer
+    from f5_tts.model import DiT
+
+    bsz = 2
+    
+    tokenizer = "pinyin"  # 'pinyin', 'char', or 'custom'
+    tokenizer_path = None  # if tokenizer = 'custom', define the path to the tokenizer you want to use (should be vocab.txt)
+    dataset_name = "Emilia_ZH_EN"
+    if tokenizer == "custom":
+        tokenizer_path = tokenizer_path
+    else:
+        tokenizer_path = dataset_name
+    vocab_char_map, vocab_size = get_tokenizer(tokenizer_path, tokenizer)
+    
+    
+    fake_unet = DiT(dim=1024, depth=22, heads=16, ff_mult=2, text_dim=512, conv_layers=4, text_num_embeds=vocab_size, mel_dim=80)
+
+    fake_unet = fake_unet.cuda()
+
+    text = ["hello world"] * bsz
+    lens = torch.randint(1, 1000, (bsz,)).cuda()
+    inp = torch.randn(bsz, lens.max(), 80).cuda()
+    
+    batch, seq_len, dtype, device = *inp.shape[:2], inp.dtype, inp.device
+
+    batch, seq_len, dtype, device = *inp.shape[:2], inp.dtype, inp.device
+    
+    # handle text as string
+    if isinstance(text, list):
+        if exists(vocab_char_map):
+            text = list_str_to_idx(text, vocab_char_map).to(device)
+        else:
+            text = list_str_to_tensor(text).to(device)
+        assert text.shape[0] == batch
+
+    # lens and mask
+    if not exists(lens):
+        lens = torch.full((batch,), seq_len, device=device)
+
+    mask = lens_to_mask(lens, length=seq_len)  # useless here, as collate_fn will pad to max length in batch
+    frac_lengths_mask = (0.7, 1.0)
+    
+    # get a random span to mask out for training conditionally
+    frac_lengths = torch.zeros((batch,), device=device).float().uniform_(*frac_lengths_mask)
+    rand_span_mask = mask_from_frac_lengths(lens, frac_lengths)
+                
+    if exists(mask):
+        rand_span_mask &= mask
+
+    # Sample a time
+    time = torch.rand((batch,), dtype=dtype, device=device)
+
+    x1 = inp
+    x0 = torch.randn_like(x1)
+    t = time.unsqueeze(-1).unsqueeze(-1)
+    
+    phi = (1 - t) * x0 + t * x1
+    flow = x1 - x0
+    cond = torch.where(rand_span_mask[..., None], torch.zeros_like(x1), x1)
+
+    layers = fake_unet(
+        x=phi, 
+        cond=cond,
+        text=text, 
+        time=time, 
+        drop_audio_cond=False,
+        drop_text=False,
+        classify_mode=True
+    )
+
+    # layers = torch.stack(layers, dim=1).transpose(-1, -2).flatten(start_dim=1, end_dim=2)
+    # print(layers.shape)
+
+    from ctcmodel import ConformerCTC
+    ctcmodel = ConformerCTC(vocab_size=vocab_size, mel_dim=80, num_heads=8, d_hid=512, nlayers=6).cuda()
+    real_out, layer = ctcmodel(inp)
+    layer = layer[-3:] # only use the last 3 layers
+    layer = [F.interpolate(l, mode='nearest', scale_factor=4).transpose(-1, -2) for l in layer]
+    if layer[0].size(1) < layers[0].size(1):
+        layer = [F.pad(l, (0, 0, 0, layers[0].size(1) - l.size(1))) for l in layer]
+    
+    layers = layer + layers
+
+    model = ConformerDiscirminator(input_dim=23 * 1024 + 3 * 512, 
+                            channels=512
+                            )
+    
+
+    model = model.cuda()
+    print(model)
+    out = model(layers)
+    print(out.shape)

dmd_trainer.py

@@ -0,0 +1,535 @@
+from __future__ import annotations
+
+import os
+import gc
+from tqdm import tqdm
+import wandb
+
+import torch
+import torch.nn as nn
+from torch.optim import AdamW
+from torch.utils.data import DataLoader, Dataset, SequentialSampler
+from torch.optim.lr_scheduler import LinearLR, SequentialLR
+
+from accelerate import Accelerator
+from accelerate.utils import DistributedDataParallelKwargs
+
+from unimodel import UniModel
+from f5_tts.model import CFM
+from f5_tts.model.utils import exists, default
+from f5_tts.model.dataset import DynamicBatchSampler, collate_fn
+
+
+# trainer
+
+import math
+
+class RunningStats:
+    def __init__(self):
+        self.count = 0
+        self.mean = 0.0
+        self.M2 = 0.0  # Sum of squared differences from the current mean
+
+    def update(self, x):
+        """Update the running statistics with a new value x."""
+        self.count += 1
+        delta = x - self.mean
+        self.mean += delta / self.count
+        delta2 = x - self.mean
+        self.M2 += delta * delta2
+
+    @property
+    def variance(self):
+        """Return the sample variance. Returns NaN if fewer than two samples."""
+        return self.M2 / (self.count - 1) if self.count > 1 else float('nan')
+
+    @property
+    def std(self):
+        """Return the sample standard deviation."""
+        return math.sqrt(self.variance)
+
+
+
+class Trainer:
+    def __init__(
+        self,
+        model: UniModel,
+        epochs,
+        learning_rate,
+        num_warmup_updates=20000,
+        save_per_updates=1000,
+        checkpoint_path=None,
+        batch_size=32,
+        batch_size_type: str = "sample",
+        max_samples=32,
+        grad_accumulation_steps=1,
+        max_grad_norm=1.0,
+        noise_scheduler: str | None = None,
+        duration_predictor: torch.nn.Module | None = None,
+        wandb_project="test_e2-tts",
+        wandb_run_name="test_run",
+        wandb_resume_id: str = None,
+        last_per_steps=None,
+        log_step=1000,
+        accelerate_kwargs: dict = dict(),
+        bnb_optimizer: bool = False,
+        scale: float = 1.0,
+        
+        # training parameters for DMDSpeech
+        num_student_step: int = 1,
+        gen_update_ratio: int = 5,
+        lambda_discriminator_loss: float = 1.0,
+        lambda_generator_loss: float = 1.0,
+        lambda_ctc_loss: float = 1.0,
+        lambda_sim_loss: float = 1.0,
+
+        num_GAN: int = 5000,
+        num_D: int = 500,
+        num_ctc: int = 5000,
+        num_sim: int = 10000,
+        num_simu: int = 1000,
+    ):
+        ddp_kwargs = DistributedDataParallelKwargs(find_unused_parameters=True)
+
+        logger = "wandb" if wandb.api.api_key else None
+        print(f"Using logger: {logger}")
+
+        self.accelerator = Accelerator(
+            log_with=logger,
+            kwargs_handlers=[ddp_kwargs],
+            gradient_accumulation_steps=grad_accumulation_steps,
+            **accelerate_kwargs,
+        )
+
+        if logger == "wandb":
+            if exists(wandb_resume_id):
+                init_kwargs = {"wandb": {"resume": "allow", "name": wandb_run_name, "id": wandb_resume_id}}
+            else:
+                init_kwargs = {"wandb": {"resume": "allow", "name": wandb_run_name}}
+            self.accelerator.init_trackers(
+                project_name=wandb_project,
+                init_kwargs=init_kwargs,
+                config={
+                    "epochs": epochs,
+                    "learning_rate": learning_rate,
+                    "num_warmup_updates": num_warmup_updates,
+                    "batch_size": batch_size,
+                    "batch_size_type": batch_size_type,
+                    "max_samples": max_samples,
+                    "grad_accumulation_steps": grad_accumulation_steps,
+                    "max_grad_norm": max_grad_norm,
+                    "gpus": self.accelerator.num_processes,
+                    "noise_scheduler": noise_scheduler,
+                },
+            )
+
+        self.model = model
+
+        self.scale = scale
+
+        self.epochs = epochs
+        self.num_warmup_updates = num_warmup_updates
+        self.save_per_updates = save_per_updates
+        self.last_per_steps = default(last_per_steps, save_per_updates * grad_accumulation_steps)
+        self.checkpoint_path = default(checkpoint_path, "ckpts/test_e2-tts")
+
+        self.batch_size = batch_size
+        self.batch_size_type = batch_size_type
+        self.max_samples = max_samples
+        self.grad_accumulation_steps = grad_accumulation_steps
+        self.max_grad_norm = max_grad_norm
+
+        self.noise_scheduler = noise_scheduler
+
+        self.duration_predictor = duration_predictor
+        
+        self.log_step = log_step
+
+        self.gen_update_ratio = gen_update_ratio # number of generator updates per guidance (fake score function and discriminator) update
+        self.lambda_discriminator_loss = lambda_discriminator_loss # weight for discriminator loss (L_adv)
+        self.lambda_generator_loss = lambda_generator_loss # weight for generator loss (L_adv)
+        self.lambda_ctc_loss = lambda_ctc_loss # weight for ctc loss
+        self.lambda_sim_loss = lambda_sim_loss # weight for similarity loss
+        
+        # create distillation schedule for student model
+        self.student_steps = (
+                torch.linspace(0.0, 1.0, num_student_step + 1)[:-1])
+        
+        self.GAN = model.guidance_model.gen_cls_loss # whether to use GAN training
+        self.num_GAN = num_GAN # number of steps before adversarial training
+        self.num_D = num_D # number of steps to train the discriminator before adversarial training 
+        self.num_ctc = num_ctc # number of steps before CTC training
+        self.num_sim = num_sim # number of steps before similarity training
+        self.num_simu = num_simu # number of steps before using simulated data
+
+        # Assuming `self.model.fake_unet.parameters()` and `self.model.guidance_model.parameters()` are accessible
+        if bnb_optimizer:
+            import bitsandbytes as bnb
+            self.optimizer_generator = bnb.optim.AdamW8bit(self.model.feedforward_model.parameters(), lr=learning_rate)
+            self.optimizer_guidance = bnb.optim.AdamW8bit(self.model.guidance_model.parameters(), lr=learning_rate)
+        else:
+            self.optimizer_generator = AdamW(self.model.feedforward_model.parameters(), lr=learning_rate, eps=1e-7)
+            self.optimizer_guidance = AdamW(self.model.guidance_model.parameters(), lr=learning_rate, eps=1e-7)
+
+        self.model, self.optimizer_generator, self.optimizer_guidance = self.accelerator.prepare(self.model, self.optimizer_generator, self.optimizer_guidance)
+
+        self.generator_norm = RunningStats()
+        self.guidance_norm = RunningStats()
+
+    
+    @property
+    def is_main(self):
+        return self.accelerator.is_main_process
+
+    def save_checkpoint(self, step, last=False):
+        self.accelerator.wait_for_everyone()
+        if self.is_main:
+            checkpoint = dict(
+                model_state_dict=self.accelerator.unwrap_model(self.model).state_dict(),
+                optimizer_generator_state_dict=self.accelerator.unwrap_model(self.optimizer_generator).state_dict(),
+                optimizer_guidance_state_dict=self.accelerator.unwrap_model(self.optimizer_guidance).state_dict(),
+                scheduler_generator_state_dict=self.scheduler_generator.state_dict(),
+                scheduler_guidance_state_dict=self.scheduler_guidance.state_dict(),
+                step=step,
+            )
+
+            if not os.path.exists(self.checkpoint_path):
+                os.makedirs(self.checkpoint_path)
+            if last:
+                self.accelerator.save(checkpoint, f"{self.checkpoint_path}/model_last.pt")
+                print(f"Saved last checkpoint at step {step}")
+            else:
+                self.accelerator.save(checkpoint, f"{self.checkpoint_path}/model_{step}.pt")
+
+    def load_checkpoint(self):
+        if (
+            not exists(self.checkpoint_path)
+            or not os.path.exists(self.checkpoint_path)
+            or not os.listdir(self.checkpoint_path)
+        ):
+            return 0
+
+        self.accelerator.wait_for_everyone()
+        if "model_last.pt" in os.listdir(self.checkpoint_path):
+            latest_checkpoint = "model_last.pt"
+        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}", weights_only=True, map_location="cpu")
+
+        self.accelerator.unwrap_model(self.model).load_state_dict(checkpoint["model_state_dict"], strict=False)
+        # self.accelerator.unwrap_model(self.optimizer_generator).load_state_dict(checkpoint["optimizer_generator_state_dict"])
+        # self.accelerator.unwrap_model(self.optimizer_guidance).load_state_dict(checkpoint["optimizer_guidance_state_dict"])
+        # if self.scheduler_guidance:
+        #     self.scheduler_guidance.load_state_dict(checkpoint["scheduler_guidance_state_dict"])
+        # if self.scheduler_generator:
+        #     self.scheduler_generator.load_state_dict(checkpoint["scheduler_generator_state_dict"])
+        step = checkpoint["step"]
+
+        del checkpoint
+        gc.collect()
+        return step
+    
+
+    def train(self, train_dataset: Dataset, num_workers=64, resumable_with_seed: int = None, vocoder: nn.Module = None):
+        if exists(resumable_with_seed):
+            generator = torch.Generator()
+            generator.manual_seed(resumable_with_seed)
+        else:
+            generator = None
+
+        if self.batch_size_type == "sample":
+            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,
+                persistent_workers=True,
+                batch_sampler=batch_sampler,
+            )
+        else:
+            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
+        warmup_steps = (
+            self.num_warmup_updates * self.accelerator.num_processes
+        )
+        
+        # consider a fixed warmup steps while using accelerate multi-gpu ddp
+        # otherwise by default with split_batches=False, warmup steps change with num_processes
+        total_steps = len(train_dataloader) * self.epochs / self.grad_accumulation_steps
+        decay_steps = total_steps - warmup_steps
+        
+        warmup_scheduler_generator = LinearLR(self.optimizer_generator, start_factor=1e-8, end_factor=1.0, total_iters=warmup_steps // (self.gen_update_ratio * self.grad_accumulation_steps))
+        decay_scheduler_generator = LinearLR(self.optimizer_generator, start_factor=1.0, end_factor=1e-8, total_iters=decay_steps // (self.gen_update_ratio * self.grad_accumulation_steps))
+        self.scheduler_generator = SequentialLR(self.optimizer_generator, schedulers=[warmup_scheduler_generator, decay_scheduler_generator], milestones=[warmup_steps // (self.gen_update_ratio * self.grad_accumulation_steps)])
+
+        warmup_scheduler_guidance = LinearLR(self.optimizer_guidance, start_factor=1e-8, end_factor=1.0, total_iters=warmup_steps)
+        decay_scheduler_guidance = LinearLR(self.optimizer_guidance, start_factor=1.0, end_factor=1e-8, total_iters=decay_steps)
+        self.scheduler_guidance = SequentialLR(self.optimizer_guidance, schedulers=[warmup_scheduler_guidance, decay_scheduler_guidance], milestones=[warmup_steps])
+
+        train_dataloader, self.scheduler_generator, self.scheduler_guidance = self.accelerator.prepare(
+            train_dataloader, self.scheduler_generator, self.scheduler_guidance
+        )  # actual steps = 1 gpu steps / gpus
+        start_step = self.load_checkpoint()
+        global_step = start_step
+
+        if exists(resumable_with_seed):
+            orig_epoch_step = len(train_dataloader)
+            skipped_epoch = int(start_step // orig_epoch_step)
+            skipped_batch = start_step % orig_epoch_step
+            skipped_dataloader = self.accelerator.skip_first_batches(train_dataloader, num_batches=skipped_batch)
+        else:
+            skipped_epoch = 0
+
+        for epoch in range(skipped_epoch, self.epochs):
+            self.model.train()
+            if exists(resumable_with_seed) and epoch == skipped_epoch:
+                progress_bar = tqdm(
+                    skipped_dataloader,
+                    desc=f"Epoch {epoch+1}/{self.epochs}",
+                    unit="step",
+                    disable=not self.accelerator.is_local_main_process,
+                    initial=skipped_batch,
+                    total=orig_epoch_step,
+                )
+            else:
+                progress_bar = tqdm(
+                    train_dataloader,
+                    desc=f"Epoch {epoch+1}/{self.epochs}",
+                    unit="step",
+                    disable=not self.accelerator.is_local_main_process,
+                )
+
+            for batch in progress_bar:
+                update_generator = global_step % self.gen_update_ratio == 0
+                        
+                with self.accelerator.accumulate(self.model):
+                    metrics = {}
+                    text_inputs = batch["text"]
+                    mel_spec = batch["mel"].permute(0, 2, 1)
+                    mel_lengths = batch["mel_lengths"]
+                    
+                    mel_spec = mel_spec / self.scale
+                    
+                    guidance_loss_dict, guidance_log_dict = self.model(inp=mel_spec, 
+                                                                text=text_inputs, 
+                                                                lens=mel_lengths, 
+                                                                student_steps=self.student_steps,
+                                                                update_generator=False,
+                                                                use_simulated=global_step >= self.num_simu,
+                                                                )
+
+                    # if self.GAN and update_generator:
+                    #     # only add discriminator loss if GAN is enabled and generator is being updated
+                    #     guidance_cls_loss = guidance_loss_dict["guidance_cls_loss"] * (self.lambda_discriminator_loss if global_step >= self.num_GAN and update_generator else 0)
+                    #     metrics['loss/discriminator_loss'] = guidance_loss_dict["guidance_cls_loss"]
+                    #     self.accelerator.backward(guidance_cls_loss, retain_graph=True)
+                        
+                    #     if self.max_grad_norm > 0 and self.accelerator.sync_gradients:
+                    #         metrics['grad_norm_guidance'] = self.accelerator.clip_grad_norm_(self.model.parameters(), self.max_grad_norm)
+
+                    guidance_loss = 0
+                    guidance_loss += guidance_loss_dict["loss_fake_mean"]
+                    metrics['loss/fake_score'] = guidance_loss_dict["loss_fake_mean"]
+                    metrics["loss/guidance_loss"] = guidance_loss
+
+                    if self.GAN and update_generator:
+                        # only add discriminator loss if GAN is enabled and generator is being updated
+                        guidance_cls_loss = guidance_loss_dict["guidance_cls_loss"] * (self.lambda_discriminator_loss if global_step >= self.num_GAN and update_generator else 0)
+                        metrics['loss/discriminator_loss'] = guidance_loss_dict["guidance_cls_loss"]
+
+                        guidance_loss += guidance_cls_loss
+                    
+                    self.accelerator.backward(guidance_loss)
+
+                    if self.max_grad_norm > 0 and self.accelerator.sync_gradients:
+                        metrics['grad_norm_guidance'] = self.accelerator.clip_grad_norm_(self.model.parameters(), self.max_grad_norm)
+
+                        # if self.guidance_norm.count < 100:
+                        #     self.guidance_norm.update(metrics['grad_norm_guidance'])
+
+                        # if metrics['grad_norm_guidance'] > self.guidance_norm.mean + 5 * self.guidance_norm.std:
+                        #     self.optimizer_generator.zero_grad()
+                        #     self.optimizer_guidance.zero_grad()
+                        #     print("Gradient explosion detected. Skipping batch.")
+                        # elif self.guidance_norm.count >= 100:
+                        #     self.guidance_norm.update(metrics['grad_norm_guidance'])
+
+
+                    self.optimizer_guidance.step()
+                    self.scheduler_guidance.step()
+                    self.optimizer_guidance.zero_grad()
+                    self.optimizer_generator.zero_grad()  # zero out the generator's gradient as well
+                    
+                    if update_generator:
+                        generator_loss_dict, generator_log_dict = self.model(inp=mel_spec, 
+                                                                        text=text_inputs, 
+                                                                        lens=mel_lengths, 
+                                                                        student_steps=self.student_steps,
+                                                                        update_generator=True,
+                                                                        use_simulated=global_step >= self.num_ctc,
+                                                                        )
+                        # if self.GAN:
+                        #     gen_cls_loss = generator_loss_dict["gen_cls_loss"] * (self.lambda_generator_loss if global_step >= (self.num_GAN + self.num_D) and update_generator else 0)
+                        #     metrics["loss/gen_cls_loss"] = generator_loss_dict["gen_cls_loss"]
+
+                        #     self.accelerator.backward(gen_cls_loss, retain_graph=True)
+
+                        #     if self.max_grad_norm > 0 and self.accelerator.sync_gradients:
+                        #         metrics['grad_norm_generator'] = self.accelerator.clip_grad_norm_(self.model.parameters(), self.max_grad_norm)
+
+                        generator_loss = 0
+                        generator_loss += generator_loss_dict["loss_dm"]
+                        if "loss_mse" in generator_loss_dict:
+                            generator_loss += generator_loss_dict["loss_mse"] 
+                        generator_loss += generator_loss_dict["loss_ctc"] * (self.lambda_ctc_loss if global_step >= self.num_ctc else 0)
+                        generator_loss += generator_loss_dict["loss_sim"] * (self.lambda_sim_loss if global_step >= self.num_sim else 0)
+                        generator_loss += generator_loss_dict["loss_kl"] * (self.lambda_ctc_loss if global_step >= self.num_ctc else 0)
+                        if self.GAN:
+                            gen_cls_loss = generator_loss_dict["gen_cls_loss"] * (self.lambda_generator_loss if global_step >= (self.num_GAN + self.num_D) and update_generator else 0)
+                            metrics["loss/gen_cls_loss"] = generator_loss_dict["gen_cls_loss"]
+                            generator_loss += gen_cls_loss
+
+                        metrics['loss/dm_loss'] = generator_loss_dict["loss_dm"]
+                        metrics['loss/ctc_loss'] = generator_loss_dict["loss_ctc"]
+
+                        metrics['loss/similarity_loss'] = generator_loss_dict["loss_sim"]
+                        metrics['loss/generator_loss'] = generator_loss
+                        
+                        if "loss_mse" in generator_loss_dict and generator_loss_dict["loss_mse"] != 0:
+                            metrics['loss/mse_loss'] = generator_loss_dict["loss_mse"]
+                        if "loss_kl" in generator_loss_dict and generator_loss_dict["loss_kl"] != 0:
+                            metrics['loss/kl_loss'] = generator_loss_dict["loss_kl"]
+
+                        self.accelerator.backward(generator_loss)
+
+                        if self.max_grad_norm > 0 and self.accelerator.sync_gradients:
+                            metrics['grad_norm_generator'] = self.accelerator.clip_grad_norm_(self.model.parameters(), self.max_grad_norm)
+                            # self.generator_norm.update(metrics['grad_norm_generator'])
+                            
+                            # if metrics['grad_norm_generator'] > self.generator_norm.mean + 15 * self.generator_norm.std:
+                            #     self.optimizer_generator.zero_grad()
+                            #     self.optimizer_guidance.zero_grad()
+                            #     update_generator = False
+                            #     print("Gradient explosion detected. Skipping batch.")
+
+                        if update_generator:
+                            self.optimizer_generator.step()
+                            self.scheduler_generator.step()
+                            self.optimizer_generator.zero_grad()
+                            self.optimizer_guidance.zero_grad()  # zero out the guidance's gradient as well
+
+
+                global_step += 1
+
+                if self.accelerator.is_local_main_process:
+                    self.accelerator.log({**metrics,
+                                          "lr_generator": self.scheduler_generator.get_last_lr()[0],
+                                          "lr_guidance": self.scheduler_guidance.get_last_lr()[0],
+                                          }
+                                         , step=global_step)
+                
+                if global_step % self.log_step == 0 and self.accelerator.is_local_main_process and vocoder is not None:
+                    # log the first batch of the epoch
+                    with torch.no_grad():
+                        generator_input = generator_log_dict['generator_input'][0].unsqueeze(0).permute(0, 2, 1) * self.scale
+                        generator_input = vocoder.decode(generator_input.float().cpu())
+                        generator_input = wandb.Audio(
+                            generator_input.float().numpy().squeeze(),
+                            sample_rate=24000,
+                            caption="time: " + str(generator_log_dict['time'][0].float().cpu().numpy())
+                        )
+
+                        generator_output = generator_log_dict['generator_output'][0].unsqueeze(0).permute(0, 2, 1) * self.scale
+                        generator_output = vocoder.decode(generator_output.float().cpu())
+                        generator_output = wandb.Audio(
+                            generator_output.float().numpy().squeeze(),
+                            sample_rate=24000,
+                            caption="time: " + str(generator_log_dict['time'][0].float().cpu().numpy())
+                        )
+                        
+                        generator_cond = generator_log_dict['generator_cond'][0].unsqueeze(0).permute(0, 2, 1) * self.scale
+                        generator_cond = vocoder.decode(generator_cond.float().cpu())
+                        generator_cond = wandb.Audio(
+                            generator_cond.float().numpy().squeeze(),
+                            sample_rate=24000,
+                            caption="time: " + str(generator_log_dict['time'][0].float().cpu().numpy())
+                        )
+                        
+                        ground_truth = generator_log_dict['ground_truth'][0].unsqueeze(0).permute(0, 2, 1) * self.scale
+                        ground_truth = vocoder.decode(ground_truth.float().cpu())
+                        ground_truth = wandb.Audio(
+                            ground_truth.float().numpy().squeeze(),
+                            sample_rate=24000,
+                            caption="time: " + str(generator_log_dict['time'][0].float().cpu().numpy())
+                        )
+                        
+                        dmtrain_noisy_inp = generator_log_dict['dmtrain_noisy_inp'][0].unsqueeze(0).permute(0, 2, 1) * self.scale
+                        dmtrain_noisy_inp = vocoder.decode(dmtrain_noisy_inp.float().cpu())
+                        dmtrain_noisy_inp = wandb.Audio(
+                            dmtrain_noisy_inp.float().numpy().squeeze(),
+                            sample_rate=24000,
+                            caption="dmtrain_time: " + str(generator_log_dict['dmtrain_time'][0].float().cpu().numpy())
+                        )
+                        
+                        dmtrain_pred_real_image = generator_log_dict['dmtrain_pred_real_image'][0].unsqueeze(0).permute(0, 2, 1) * self.scale
+                        dmtrain_pred_real_image = vocoder.decode(dmtrain_pred_real_image.float().cpu())
+                        dmtrain_pred_real_image = wandb.Audio(
+                            dmtrain_pred_real_image.float().numpy().squeeze(),
+                            sample_rate=24000,
+                            caption="dmtrain_time: " + str(generator_log_dict['dmtrain_time'][0].float().cpu().numpy())
+                        )
+                        
+                        dmtrain_pred_fake_image = generator_log_dict['dmtrain_pred_fake_image'][0].unsqueeze(0).permute(0, 2, 1) * self.scale
+                        dmtrain_pred_fake_image = vocoder.decode(dmtrain_pred_fake_image.float().cpu())
+                        dmtrain_pred_fake_image = wandb.Audio(
+                            dmtrain_pred_fake_image.float().numpy().squeeze(),
+                            sample_rate=24000,
+                            caption="dmtrain_time: " + str(generator_log_dict['dmtrain_time'][0].float().cpu().numpy())
+                        )
+                        
+                                                
+                        self.accelerator.log({"noisy_input": generator_input, 
+                                              "output": generator_output,
+                                                "cond": generator_cond,
+                                                "ground_truth": ground_truth,
+                                                "dmtrain_noisy_inp": dmtrain_noisy_inp,
+                                                "dmtrain_pred_real_image": dmtrain_pred_real_image,
+                                                "dmtrain_pred_fake_image": dmtrain_pred_fake_image,
+                                                
+                                             }, step=global_step)
+
+                progress_bar.set_postfix(step=str(global_step), metrics=metrics)
+
+                if global_step % (self.save_per_updates * self.grad_accumulation_steps) == 0:
+                    self.save_checkpoint(global_step)
+
+                if global_step % self.last_per_steps == 0:
+                    self.save_checkpoint(global_step, last=True)
+
+        self.save_checkpoint(global_step, last=True)
+
+        self.accelerator.end_training()
+        
+        

duration_predictor.py

@@ -0,0 +1,84 @@
+import torch
+import torch.nn as nn
+
+# from tts_encode import tts_encode
+
+def calculate_remaining_lengths(mel_lengths):
+    B = mel_lengths.shape[0]
+    max_L = mel_lengths.max().item()  # Get the maximum length in the batch
+
+    # Create a range tensor: shape (max_L,), [0, 1, 2, ..., max_L-1]
+    range_tensor = torch.arange(max_L, device=mel_lengths.device).expand(B, max_L)
+
+    # Compute targets using broadcasting: (L-1) - range_tensor
+    remain_lengths = (mel_lengths[:, None] - 1 - range_tensor).clamp(min=0)
+
+    return remain_lengths
+
+
+class PositionalEncoding(nn.Module):
+    def __init__(self, hidden_dim, max_len=4096):
+        super().__init__()
+        pe = torch.zeros(max_len, hidden_dim)
+        position = torch.arange(0, max_len, dtype=torch.float).unsqueeze(1)
+        div_term = torch.exp(torch.arange(0, hidden_dim, 2).float() * (-torch.log(torch.tensor(10000.0)) / hidden_dim))
+        pe[:, 0::2] = torch.sin(position * div_term)
+        pe[:, 1::2] = torch.cos(position * div_term)
+        self.pe = pe.unsqueeze(0)  # Shape: (1, max_len, hidden_dim)
+
+    def forward(self, x):
+        x = x + self.pe[:, :x.size(1)].to(x.device)
+        return x
+
+
+class SpeechLengthPredictor(nn.Module):
+
+    def __init__(self, 
+        vocab_size=2545, n_mel=100, hidden_dim=256, 
+        n_text_layer=4, n_cross_layer=4, n_head=8,
+        output_dim=1,
+    ):
+        super().__init__()
+        
+        # Text Encoder: Embedding + Transformer Layers
+        self.text_embedder = nn.Embedding(vocab_size+1, hidden_dim, padding_idx=vocab_size)
+        self.text_pe = PositionalEncoding(hidden_dim)
+        encoder_layer = nn.TransformerEncoderLayer(
+            d_model=hidden_dim, nhead=n_head, dim_feedforward=hidden_dim*2, batch_first=True
+        )
+        self.text_encoder = nn.TransformerEncoder(encoder_layer, num_layers=n_text_layer)
+        
+        # Mel Spectrogram Embedder
+        self.mel_embedder = nn.Linear(n_mel, hidden_dim)
+        self.mel_pe = PositionalEncoding(hidden_dim)
+
+        # Transformer Decoder Layers with Cross-Attention in Every Layer
+        decoder_layer = nn.TransformerDecoderLayer(
+            d_model=hidden_dim, nhead=n_head, dim_feedforward=hidden_dim*2, batch_first=True
+        )
+        self.decoder = nn.TransformerDecoder(decoder_layer, num_layers=n_cross_layer)
+        
+        # Final Classification Layer
+        self.predictor = nn.Linear(hidden_dim, output_dim)
+
+    def forward(self, text_ids, mel):
+        # Encode text
+        text_embedded = self.text_pe(self.text_embedder(text_ids))
+        text_features = self.text_encoder(text_embedded) # (B, L_text, D)
+        
+        # Encode Mel spectrogram
+        mel_features = self.mel_pe(self.mel_embedder(mel))  # (B, L_mel, D)
+        
+        # Causal Masking for Decoder
+        seq_len = mel_features.size(1)
+        causal_mask = torch.triu(torch.ones(seq_len, seq_len), diagonal=1).bool().to(mel.device)
+        # causal_mask = torch.triu(
+        #     torch.full((seq_len, seq_len), float('-inf'), device=mel.device), diagonal=1
+        # )
+
+        # Transformer Decoder with Cross-Attention in Each Layer
+        decoder_out = self.decoder(mel_features, text_features, tgt_mask=causal_mask)
+        
+        # Length Prediction
+        length_logits = self.predictor(decoder_out).squeeze(-1)
+        return length_logits

duration_trainer.py

@@ -0,0 +1,563 @@
+from __future__ import annotations
+
+import gc
+import os
+
+import math
+
+import torch
+import torchaudio
+import wandb
+from accelerate import Accelerator
+from accelerate.utils import DistributedDataParallelKwargs
+from ema_pytorch import EMA
+from torch.optim import AdamW
+from torch.optim.lr_scheduler import LinearLR, SequentialLR
+from torch.utils.data import DataLoader, Dataset, SequentialSampler, Subset  # <-- Added Subset import
+from tqdm import tqdm
+
+import torch.nn.functional as F
+
+from f5_tts.model import CFM
+from f5_tts.model.dataset import collate_fn, DynamicBatchSampler
+from f5_tts.model.utils import default, exists
+
+from duration_predictor import calculate_remaining_lengths
+
+# trainer
+
+from f5_tts.model.utils import (
+    default,
+    exists,
+    list_str_to_idx,
+    list_str_to_tensor,
+    lens_to_mask,
+    mask_from_frac_lengths,
+)
+
+SAMPLE_RATE = 24_000
+
+
+def masked_l1_loss(est_lengths, tar_lengths):
+    first_zero_idx = (tar_lengths == 0).int().argmax(dim=1)  
+    B, L = tar_lengths.shape
+    range_tensor = torch.arange(L, device=tar_lengths.device).expand(B, L)  
+    mask = range_tensor <= first_zero_idx[:, None]  # Include the first 0
+    loss = F.l1_loss(est_lengths, tar_lengths, reduction='none')  # (B, L)
+    loss = loss * mask  # Zero out ignored positions
+    loss = loss.sum() / mask.sum()  # Normalize by valid elements
+    return loss
+
+
+def masked_cross_entropy_loss(est_length_logits, tar_length_labels):
+    first_zero_idx = (tar_length_labels == 0).int().argmax(dim=1)
+    B, L = tar_length_labels.shape
+    range_tensor = torch.arange(L, device=tar_length_labels.device).expand(B, L)
+    mask = range_tensor <= first_zero_idx[:, None]  # Include the first 0
+    loss = F.cross_entropy(
+        est_length_logits.reshape(-1, est_length_logits.size(-1)), 
+        tar_length_labels.reshape(-1), 
+        reduction='none'
+    ).reshape(B, L)
+    loss = loss * mask
+    loss = loss.sum() / mask.sum()
+    return loss
+
+
+class Trainer:
+    def __init__(
+        self,
+        model,
+        vocab_size,
+        vocab_char_map,
+        process_token_to_id=True,
+        loss_fn='L1',
+        lambda_L1=1,
+        gumbel_tau=0.5,
+        n_class=301,
+        n_frame_per_class=10,
+        epochs=15,
+        learning_rate=1e-4,
+        num_warmup_updates=20000,
+        save_per_updates=1000,
+        checkpoint_path=None,
+        batch_size=32,
+        batch_size_type: str = "sample",
+        max_samples=32,
+        grad_accumulation_steps=1,
+        max_grad_norm=1.0,
+        logger: str | None = "wandb",  # "wandb" | "tensorboard" | None
+        wandb_project="test_e2-tts",
+        wandb_run_name="test_run",
+        wandb_resume_id: str = None,
+        last_per_steps=None,
+        accelerate_kwargs: dict = dict(),
+        ema_kwargs: dict = dict(),
+    ):
+        ddp_kwargs = DistributedDataParallelKwargs(find_unused_parameters=False)
+
+        if logger == "wandb" and not wandb.api.api_key:
+            logger = None
+        print(f"Using logger: {logger}")
+
+        self.accelerator = Accelerator(
+            log_with=logger if logger == "wandb" else None,
+            kwargs_handlers=[ddp_kwargs],
+            gradient_accumulation_steps=grad_accumulation_steps,
+            **accelerate_kwargs,
+        )
+
+        self.logger = logger
+        if self.logger == "wandb":
+            if exists(wandb_resume_id):
+                init_kwargs = {"wandb": {"resume": "allow", "name": wandb_run_name, "id": wandb_resume_id}}
+            else:
+                init_kwargs = {"wandb": {"resume": "allow", "name": wandb_run_name}}
+
+            self.accelerator.init_trackers(
+                project_name=wandb_project,
+                init_kwargs=init_kwargs,
+                config={
+                    "epochs": epochs,
+                    "learning_rate": learning_rate,
+                    "num_warmup_updates": num_warmup_updates,
+                    "batch_size": batch_size,
+                    "batch_size_type": batch_size_type,
+                    "max_samples": max_samples,
+                    "grad_accumulation_steps": grad_accumulation_steps,
+                    "max_grad_norm": max_grad_norm,
+                    "gpus": self.accelerator.num_processes,
+                },
+            )
+
+        elif self.logger == "tensorboard":
+            from torch.utils.tensorboard import SummaryWriter
+
+            self.writer = SummaryWriter(log_dir=f"runs/{wandb_run_name}")
+
+        self.model = model
+        self.vocab_size = vocab_size
+        self.vocab_char_map = vocab_char_map
+        self.process_token_to_id = process_token_to_id
+        assert loss_fn in ['L1', 'CE', 'L1_and_CE']
+        self.loss_fn = loss_fn
+        self.lambda_L1 = lambda_L1
+        self.n_class = n_class
+        self.n_frame_per_class = n_frame_per_class
+        self.gumbel_tau = gumbel_tau
+
+        self.epochs = epochs
+        self.num_warmup_updates = num_warmup_updates
+        self.save_per_updates = save_per_updates
+        self.last_per_steps = default(last_per_steps, save_per_updates * grad_accumulation_steps)
+        self.checkpoint_path = default(checkpoint_path, "ckpts/test_e2-tts")
+
+        self.batch_size = batch_size
+        self.batch_size_type = batch_size_type
+        self.max_samples = max_samples
+        self.grad_accumulation_steps = grad_accumulation_steps
+        self.max_grad_norm = max_grad_norm
+
+        if bnb_optimizer:
+            import bitsandbytes as bnb
+
+            self.optimizer = bnb.optim.AdamW8bit(model.parameters(), lr=learning_rate)
+        else:
+            self.optimizer = AdamW(model.parameters(), lr=learning_rate)
+        self.model, self.optimizer = self.accelerator.prepare(self.model, self.optimizer)
+        
+    @property
+    def is_main(self):
+        return self.accelerator.is_main_process
+
+    def save_checkpoint(self, step, last=False):
+        self.accelerator.wait_for_everyone()
+        if self.is_main:  
+            checkpoint = dict(
+                model_state_dict=self.accelerator.unwrap_model(self.model).state_dict(),
+                optimizer_state_dict=self.accelerator.unwrap_model(self.optimizer).state_dict(),
+                scheduler_state_dict=self.scheduler.state_dict(),
+                step=step,
+            )
+            if not os.path.exists(self.checkpoint_path):
+                os.makedirs(self.checkpoint_path)
+            if last:
+                self.accelerator.save(checkpoint, f"{self.checkpoint_path}/model_last.pt")
+            else:
+                self.accelerator.save(checkpoint, f"{self.checkpoint_path}/model_{step}.pt")
+
+    def load_checkpoint(self):
+        if (
+            not exists(self.checkpoint_path)
+            or not os.path.exists(self.checkpoint_path)
+            or not any(filename.endswith(".pt") for filename in os.listdir(self.checkpoint_path))
+        ):
+            return 0
+
+        self.accelerator.wait_for_everyone()
+        if "model_last.pt" in os.listdir(self.checkpoint_path):
+            latest_checkpoint = "model_last.pt"
+        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]
+
+        print(f'To load from {latest_checkpoint}.')
+
+        # 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}", weights_only=True, map_location="cpu")
+
+        print(f'Loaded from {latest_checkpoint}.')
+
+        if "step" in checkpoint:
+            # patch for backward compatibility, 305e3ea
+            for key in ["mel_spec.mel_stft.mel_scale.fb", "mel_spec.mel_stft.spectrogram.window"]:
+                if key in checkpoint["model_state_dict"]:
+                    del checkpoint["model_state_dict"][key]
+
+            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()
+
+        print(f'Exit load_checkpoint.')
+
+        return step
+
+
+    def validate(self, valid_dataloader, global_step):
+        """
+        Runs evaluation on the validation set, computes the average loss,
+        and logs the average validation loss along with the CTC decoded strings.
+        """
+        self.model.eval()
+        total_valid_loss = 0.0
+        total_sec_error = 0.0
+        count = 0
+        # Iterate over the validation dataloader
+        with torch.no_grad():
+            for batch in valid_dataloader:
+                # Inputs
+                mel = batch['mel'].permute(0, 2, 1) # (B, L_mel, D)
+                text = batch['text']
+
+                if self.process_token_to_id:
+                    text_ids = list_str_to_idx(text, self.vocab_char_map).to(mel.device)
+                    text_ids = text_ids.masked_fill(text_ids==-1, self.vocab_size)
+                else:
+                    text_ids = text
+
+                # Targets
+                mel_lengths = batch['mel_lengths']
+                tar_lengths = calculate_remaining_lengths(mel_lengths)
+                predictions = self.model(text_ids=text_ids, mel=mel)
+
+                if self.loss_fn == 'L1':
+                    est_lengths = predictions
+                    loss = masked_l1_loss(
+                        est_lengths=est_lengths, tar_lengths=tar_lengths
+                    )
+                    frame_error = loss
+
+                elif self.loss_fn == 'CE':
+                    tar_length_labels = (tar_lengths // self.n_frame_per_class) \
+                        .clamp(min=0, max=self.n_class-1) # [0, 1, ..., n_class-1]
+                    est_length_logtis = predictions
+                    est_length_labels = torch.argmax(est_length_logtis, dim=-1)
+                    loss = masked_cross_entropy_loss(
+                        est_length_logits=est_length_logtis, tar_length_labels=tar_length_labels
+                    )
+                    est_lengths = est_length_labels * self.n_frame_per_class
+                    frame_error = masked_l1_loss(
+                        est_lengths=est_lengths, tar_lengths=tar_lengths
+                    )
+
+                elif self.loss_fn == 'L1_and_CE':
+                    tar_length_labels = (tar_lengths // self.n_frame_per_class) \
+                        .clamp(min=0, max=self.n_class-1) # [0, 1, ..., n_class-1]
+                    est_length_logtis = predictions
+                    est_length_1hots = F.gumbel_softmax(
+                        est_length_logtis, tau=self.gumbel_tau, hard=True, dim=-1
+                    )
+                    length_values = torch.arange(
+                        self.n_class, device=est_length_1hots.device
+                    ).float() * self.n_frame_per_class
+                    est_lengths = (est_length_1hots * length_values).sum(-1)
+
+                    loss_CE = masked_cross_entropy_loss(
+                        est_length_logits=est_length_logtis, tar_length_labels=tar_length_labels
+                    )
+
+                    loss_L1 = masked_l1_loss(
+                        est_lengths=est_lengths, tar_lengths=tar_lengths
+                    )
+
+                    loss = loss_CE + self.lambda_L1 * loss_L1
+
+                    frame_error = loss_L1
+
+                else:
+                    raise NotImplementedError(self.loss_fn)
+
+                sec_error = frame_error * 256 / 24000
+                total_sec_error += sec_error.item()
+                total_valid_loss += loss.item()
+                count += 1
+
+        avg_valid_loss = total_valid_loss / count if count > 0 else 0.0
+        avg_valid_sec_error = total_sec_error / count if count > 0 else 0.0
+        # Log validation metrics
+        self.accelerator.log(
+            {
+                f"valid_loss": avg_valid_loss,
+                f"valid_sec_error": avg_valid_sec_error
+            }, 
+            step=global_step
+        )
+ 
+        self.model.train()
+
+
+    def train(self, train_dataset: Dataset, valid_dataset: Dataset,
+        num_workers=64, resumable_with_seed: int = None):
+        if exists(resumable_with_seed):
+            generator = torch.Generator()
+            generator.manual_seed(resumable_with_seed)
+        else:
+            generator = None
+
+        # Create training dataloader using the appropriate batching strategy
+        if self.batch_size_type == "sample":
+            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,
+            )
+            # Create validation dataloader (always sequential, no shuffling)
+            valid_dataloader = DataLoader(
+                valid_dataset,
+                collate_fn=collate_fn,
+                num_workers=num_workers,
+                pin_memory=True,
+                batch_size=self.batch_size,
+                shuffle=False,
+            )
+
+        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,
+                persistent_workers=True,
+                batch_sampler=batch_sampler,
+            )
+
+            sampler = SequentialSampler(valid_dataset)
+            batch_sampler = DynamicBatchSampler(
+                sampler, self.batch_size, max_samples=self.max_samples, random_seed=resumable_with_seed, drop_last=False
+            )
+            # Create validation dataloader (always sequential, no shuffling)
+            valid_dataloader = DataLoader(
+                valid_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 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
+        warmup_steps = (
+            self.num_warmup_updates * self.accelerator.num_processes
+        )  # consider a fixed warmup steps while using accelerate multi-gpu ddp
+        # otherwise by default with split_batches=False, warmup steps change with num_processes
+        total_steps = len(train_dataloader) * self.epochs / self.grad_accumulation_steps
+        decay_steps = total_steps - warmup_steps
+        warmup_scheduler = LinearLR(self.optimizer, start_factor=1e-8, end_factor=1.0, total_iters=warmup_steps)
+        decay_scheduler = LinearLR(self.optimizer, start_factor=1.0, end_factor=1e-8, total_iters=decay_steps)
+        self.scheduler = SequentialLR(
+            self.optimizer, schedulers=[warmup_scheduler, decay_scheduler], milestones=[warmup_steps]
+        )
+        train_dataloader, self.scheduler = self.accelerator.prepare(
+            train_dataloader, self.scheduler
+        )  # actual steps = 1 gpu steps / gpus
+        start_step = self.load_checkpoint()
+        global_step = start_step
+
+        valid_dataloader = self.accelerator.prepare(valid_dataloader)
+
+        if exists(resumable_with_seed):
+            orig_epoch_step = len(train_dataloader)
+            skipped_epoch = int(start_step // orig_epoch_step)
+            skipped_batch = start_step % orig_epoch_step
+            skipped_dataloader = self.accelerator.skip_first_batches(train_dataloader, num_batches=skipped_batch)
+        else:
+            skipped_epoch = 0
+
+        for epoch in range(skipped_epoch, self.epochs):
+            self.model.train()
+            if exists(resumable_with_seed) and epoch == skipped_epoch:
+                progress_bar = tqdm(
+                    skipped_dataloader,
+                    desc=f"Epoch {epoch+1}/{self.epochs}",
+                    unit="step",
+                    disable=not self.accelerator.is_local_main_process,
+                    initial=skipped_batch,
+                    total=orig_epoch_step,
+                )
+            else:
+                progress_bar = tqdm(
+                    train_dataloader,
+                    desc=f"Epoch {epoch+1}/{self.epochs}",
+                    unit="step",
+                    disable=not self.accelerator.is_local_main_process,
+                )
+
+            for batch in progress_bar:
+                with self.accelerator.accumulate(self.model):
+                    # Inputs
+                    mel = batch['mel'].permute(0, 2, 1) # (B, L_mel, D)
+                    text = batch['text']
+
+                    if self.process_token_to_id:
+                        text_ids = list_str_to_idx(text, self.vocab_char_map).to(mel.device)
+                        text_ids = text_ids.masked_fill(text_ids==-1, self.vocab_size)
+                    else:
+                        text_ids = text
+
+                    # Targets
+                    mel_lengths = batch['mel_lengths']
+                    tar_lengths = calculate_remaining_lengths(mel_lengths)
+                    predictions = self.model(text_ids=text_ids, mel=mel)
+
+                    if self.loss_fn == 'L1':
+                        est_lengths = predictions
+                        loss = masked_l1_loss(
+                            est_lengths=est_lengths, tar_lengths=tar_lengths
+                        )
+
+                        with torch.no_grad():
+                            frame_error = loss
+                            sec_error = frame_error * 256 / 24000
+
+                        log_dict = {
+                            'loss': loss.item(), 
+                            'loss_L1': loss.item(), 
+                            'sec_error': sec_error.item(),
+                            'lr': self.scheduler.get_last_lr()[0]
+                         }
+
+                    elif self.loss_fn == 'CE':
+                        tar_length_labels = (tar_lengths // self.n_frame_per_class) \
+                            .clamp(min=0, max=self.n_class-1) # [0, 1, ..., n_class-1]
+                        est_length_logtis = predictions
+                        est_length_labels = torch.argmax(est_length_logtis, dim=-1)
+                        loss = masked_cross_entropy_loss(
+                            est_length_logits=est_length_logtis, tar_length_labels=tar_length_labels
+                        )
+                        with torch.no_grad():
+                            est_lengths = est_length_labels * self.n_frame_per_class
+                            frame_error = masked_l1_loss(
+                                est_lengths=est_lengths, tar_lengths=tar_lengths
+                            )
+                            sec_error = frame_error * 256 / 24000
+
+                        log_dict = {
+                            'loss': loss.item(), 
+                            'loss_CE': loss.item(), 
+                            'sec_error': sec_error.item(),
+                            'lr': self.scheduler.get_last_lr()[0]
+                         }
+
+                    elif self.loss_fn == 'L1_and_CE':
+                        tar_length_labels = (tar_lengths // self.n_frame_per_class) \
+                            .clamp(min=0, max=self.n_class-1) # [0, 1, ..., n_class-1]
+                        est_length_logtis = predictions
+                        est_length_1hots = F.gumbel_softmax(
+                            est_length_logtis, tau=self.gumbel_tau, hard=True, dim=-1
+                        )
+                        length_values = torch.arange(
+                            self.n_class, device=est_length_1hots.device
+                        ).float() * self.n_frame_per_class
+                        est_lengths = (est_length_1hots * length_values).sum(-1)
+
+                        loss_CE = masked_cross_entropy_loss(
+                            est_length_logits=est_length_logtis, tar_length_labels=tar_length_labels
+                        )
+
+                        loss_L1 =  masked_l1_loss(
+                            est_lengths=est_lengths, tar_lengths=tar_lengths
+                        )
+
+                        loss = loss_CE + self.lambda_L1 * loss_L1
+
+                        with torch.no_grad():
+                            frame_error = loss_L1
+                            sec_error = frame_error * 256 / 24000
+
+                        log_dict = {
+                            'loss': loss.item(), 
+                            'loss_L1': loss_L1.item(), 
+                            'loss_CE': loss_CE.item(), 
+                            'sec_error': sec_error.item(),
+                            'lr': self.scheduler.get_last_lr()[0]
+                        }
+
+                    else:
+                        raise NotImplementedError(self.loss_fn)
+
+        
+                    self.accelerator.backward(loss)
+
+                    if self.max_grad_norm > 0 and self.accelerator.sync_gradients:
+                        self.accelerator.clip_grad_norm_(self.model.parameters(), self.max_grad_norm)
+
+                    self.optimizer.step()
+                    self.scheduler.step()
+                    self.optimizer.zero_grad()
+
+                global_step += 1
+
+                if self.accelerator.is_local_main_process:
+                    self.accelerator.log(log_dict, step=global_step)
+                progress_bar.set_postfix(step=str(global_step), loss=loss.item())
+
+                if global_step % (self.save_per_updates * self.grad_accumulation_steps) == 0:
+                    self.save_checkpoint(global_step)
+                    # if self.log_samples and self.accelerator.is_local_main_process:
+                    # Run validation at the end of each epoch (only on the main process)
+                    if self.accelerator.is_local_main_process:
+                        self.validate(valid_dataloader, global_step)
+                # if global_step % self.last_per_steps == 0:
+                #     self.save_checkpoint(global_step, last=True)
+
+        self.save_checkpoint(global_step, last=True)
+        self.accelerator.end_training()

duration_trainer_with_prompt.py

@@ -0,0 +1,452 @@
+from __future__ import annotations
+
+import gc
+import os
+
+import math
+
+import torch
+import torchaudio
+import wandb
+from accelerate import Accelerator
+from accelerate.utils import DistributedDataParallelKwargs
+from ema_pytorch import EMA
+from torch.optim import AdamW
+from torch.optim.lr_scheduler import LinearLR, SequentialLR
+from torch.utils.data import DataLoader, Dataset, SequentialSampler, Subset  # <-- Added Subset import
+from tqdm import tqdm
+
+import torch.nn.functional as F
+
+from f5_tts.model import CFM
+from f5_tts.model.dataset import collate_fn, DynamicBatchSampler
+from f5_tts.model.utils import default, exists
+
+# trainer
+
+from f5_tts.model.utils import (
+    default,
+    exists,
+    list_str_to_idx,
+    list_str_to_tensor,
+    lens_to_mask,
+    mask_from_frac_lengths,
+)
+
+SAMPLE_RATE = 24_000
+
+
+class Trainer:
+    def __init__(
+        self,
+        model,
+        vocab_size,
+        vocab_char_map,
+        process_token_to_id=True,
+        loss_fn='L1',
+        lambda_L1=1,
+        gumbel_tau=0.5,
+        n_class=301,
+        n_frame_per_class=10,
+        epochs=15,
+        learning_rate=1e-4,
+        num_warmup_updates=20000,
+        save_per_updates=1000,
+        checkpoint_path=None,
+        batch_size=32,
+        batch_size_type: str = "sample",
+        max_samples=32,
+        grad_accumulation_steps=1,
+        max_grad_norm=1.0,
+        logger: str | None = "wandb",  # "wandb" | "tensorboard" | None
+        wandb_project="test_e2-tts",
+        wandb_run_name="test_run",
+        wandb_resume_id: str = None,
+        last_per_steps=None,
+        accelerate_kwargs: dict = dict(),
+        ema_kwargs: dict = dict(),
+        bnb_optimizer: bool = False,
+    ):
+        ddp_kwargs = DistributedDataParallelKwargs(find_unused_parameters=False)
+
+        if logger == "wandb" and not wandb.api.api_key:
+            logger = None
+        print(f"Using logger: {logger}")
+
+        self.accelerator = Accelerator(
+            log_with=logger if logger == "wandb" else None,
+            kwargs_handlers=[ddp_kwargs],
+            gradient_accumulation_steps=grad_accumulation_steps,
+            **accelerate_kwargs,
+        )
+
+        self.logger = logger
+        if self.logger == "wandb":
+            if exists(wandb_resume_id):
+                init_kwargs = {"wandb": {"resume": "allow", "name": wandb_run_name, "id": wandb_resume_id}}
+            else:
+                init_kwargs = {"wandb": {"resume": "allow", "name": wandb_run_name}}
+
+            self.accelerator.init_trackers(
+                project_name=wandb_project,
+                init_kwargs=init_kwargs,
+                config={
+                    "epochs": epochs,
+                    "learning_rate": learning_rate,
+                    "num_warmup_updates": num_warmup_updates,
+                    "batch_size": batch_size,
+                    "batch_size_type": batch_size_type,
+                    "max_samples": max_samples,
+                    "grad_accumulation_steps": grad_accumulation_steps,
+                    "max_grad_norm": max_grad_norm,
+                    "gpus": self.accelerator.num_processes,
+                },
+            )
+
+        elif self.logger == "tensorboard":
+            from torch.utils.tensorboard import SummaryWriter
+
+            self.writer = SummaryWriter(log_dir=f"runs/{wandb_run_name}")
+
+        self.model = model
+        self.vocab_size = vocab_size
+        self.vocab_char_map = vocab_char_map
+        self.process_token_to_id = process_token_to_id
+        assert loss_fn in ['L1', 'CE', 'L1_and_CE']
+        self.loss_fn = loss_fn
+        self.lambda_L1 = lambda_L1
+        self.n_class = n_class
+        self.n_frame_per_class = n_frame_per_class
+        self.gumbel_tau = gumbel_tau
+
+        self.epochs = epochs
+        self.num_warmup_updates = num_warmup_updates
+        self.save_per_updates = save_per_updates
+        self.last_per_steps = default(last_per_steps, save_per_updates * grad_accumulation_steps)
+        self.checkpoint_path = default(checkpoint_path, "ckpts/test_e2-tts")
+
+        self.batch_size = batch_size
+        self.batch_size_type = batch_size_type
+        self.max_samples = max_samples
+        self.grad_accumulation_steps = grad_accumulation_steps
+        self.max_grad_norm = max_grad_norm
+
+        if bnb_optimizer:
+            import bitsandbytes as bnb
+
+            self.optimizer = bnb.optim.AdamW8bit(model.parameters(), lr=learning_rate)
+        else:
+            self.optimizer = AdamW(model.parameters(), lr=learning_rate)
+        self.model, self.optimizer = self.accelerator.prepare(self.model, self.optimizer)
+        
+    @property
+    def is_main(self):
+        return self.accelerator.is_main_process
+
+    def save_checkpoint(self, step, last=False):
+        self.accelerator.wait_for_everyone()
+        if self.is_main:  
+            checkpoint = dict(
+                model_state_dict=self.accelerator.unwrap_model(self.model).state_dict(),
+                optimizer_state_dict=self.accelerator.unwrap_model(self.optimizer).state_dict(),
+                scheduler_state_dict=self.scheduler.state_dict(),
+                step=step,
+            )
+            if not os.path.exists(self.checkpoint_path):
+                os.makedirs(self.checkpoint_path)
+            if last:
+                self.accelerator.save(checkpoint, f"{self.checkpoint_path}/model_last.pt")
+            else:
+                self.accelerator.save(checkpoint, f"{self.checkpoint_path}/model_{step}.pt")
+
+    def load_checkpoint(self):
+        if (
+            not exists(self.checkpoint_path)
+            or not os.path.exists(self.checkpoint_path)
+            or not any(filename.endswith(".pt") for filename in os.listdir(self.checkpoint_path))
+        ):
+            return 0
+
+        self.accelerator.wait_for_everyone()
+        if "model_last.pt" in os.listdir(self.checkpoint_path):
+            latest_checkpoint = "model_last.pt"
+        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]
+
+        print(f'To load from {latest_checkpoint}.')
+
+        # 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}", weights_only=True, map_location="cpu")
+
+        print(f'Loaded from {latest_checkpoint}.')
+
+        if "step" in checkpoint:
+            # patch for backward compatibility, 305e3ea
+            for key in ["mel_spec.mel_stft.mel_scale.fb", "mel_spec.mel_stft.spectrogram.window"]:
+                if key in checkpoint["model_state_dict"]:
+                    del checkpoint["model_state_dict"][key]
+
+            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()
+
+        print(f'Exit load_checkpoint.')
+
+        return step
+
+
+    def validate(self, valid_dataloader, global_step):
+        """
+        Runs evaluation on the validation set, computes the average loss,
+        and logs the average validation loss along with the CTC decoded strings.
+        """
+        self.model.eval()
+        total_valid_loss = 0.0
+        total_sec_error = 0.0
+        count = 0
+
+        # Iterate over the validation dataloader
+        with torch.no_grad():
+            for batch in valid_dataloader:
+
+                # Inputs
+                prompt_mel = batch['pmt_mel_specs'].permute(0, 2, 1) # (B, L_mel, D)
+                prompt_text = batch['pmt_text']
+                text = batch['text']
+
+                target_ids = list_str_to_idx(text, self.vocab_char_map).to(prompt_mel.device)
+                target_ids = target_ids.masked_fill(target_ids==-1, vocab_size)
+
+                prompt_ids = list_str_to_idx(prompt_text, self.vocab_char_map).to(prompt_mel.device)
+                prompt_ids = prompt_ids.masked_fill(prompt_ids==-1, vocab_size)
+
+                # Targets
+                tar_lengths = batch['mel_lengths']
+
+                # Forward
+                predictions = SLP(target_ids=target_ids, prompt_ids=prompt_ids, prompt_mel=prompt_mel) # (B, C)
+
+                if self.loss_fn == 'CE':
+                    tar_length_labels = (tar_lengths // self.n_frame_per_class) \
+                        .clamp(min=0, max=self.n_class-1) # [0, 1, ..., n_class-1]
+                    est_length_logtis = predictions
+                    est_length_labels = torch.argmax(est_length_logtis, dim=-1)
+                    loss = F.cross_entropy(est_length_logtis, tar_length_labels)
+                    
+                    est_lengths = est_length_labels * self.n_frame_per_class
+                    frame_error = (est_lengths.float() - tar_lengths.float()).abs().mean()
+                    sec_error = frame_error * 256 / 24000
+
+                total_sec_error += sec_error.item()
+                total_valid_loss += loss.item()
+                count += 1
+
+        avg_valid_loss = total_valid_loss / count if count > 0 else 0.0
+        avg_valid_sec_error = total_sec_error / count if count > 0 else 0.0
+
+        # Log validation metrics
+        self.accelerator.log(
+            {
+                f"valid_loss": avg_valid_loss,
+                f"valid_sec_error": avg_valid_sec_error
+            }, 
+            step=global_step
+        )
+ 
+        self.model.train()
+
+
+    def train(self, train_dataset: Dataset, valid_dataset: Dataset,
+        num_workers=64, resumable_with_seed: int = None):
+        if exists(resumable_with_seed):
+            generator = torch.Generator()
+            generator.manual_seed(resumable_with_seed)
+        else:
+            generator = None
+
+        # Create training dataloader using the appropriate batching strategy
+        if self.batch_size_type == "sample":
+            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,
+            )
+            # Create validation dataloader (always sequential, no shuffling)
+            valid_dataloader = DataLoader(
+                valid_dataset,
+                collate_fn=collate_fn,
+                num_workers=num_workers,
+                pin_memory=True,
+                batch_size=self.batch_size,
+                shuffle=False,
+            )
+
+        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,
+                persistent_workers=True,
+                batch_sampler=batch_sampler,
+            )
+
+            sampler = SequentialSampler(valid_dataset)
+            batch_sampler = DynamicBatchSampler(
+                sampler, self.batch_size, max_samples=self.max_samples, random_seed=resumable_with_seed, drop_last=False
+            )
+            # Create validation dataloader (always sequential, no shuffling)
+            valid_dataloader = DataLoader(
+                valid_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 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
+        warmup_steps = (
+            self.num_warmup_updates * self.accelerator.num_processes
+        )  # consider a fixed warmup steps while using accelerate multi-gpu ddp
+        # otherwise by default with split_batches=False, warmup steps change with num_processes
+        total_steps = len(train_dataloader) * self.epochs / self.grad_accumulation_steps
+        decay_steps = total_steps - warmup_steps
+        warmup_scheduler = LinearLR(self.optimizer, start_factor=1e-8, end_factor=1.0, total_iters=warmup_steps)
+        decay_scheduler = LinearLR(self.optimizer, start_factor=1.0, end_factor=1e-8, total_iters=decay_steps)
+        self.scheduler = SequentialLR(
+            self.optimizer, schedulers=[warmup_scheduler, decay_scheduler], milestones=[warmup_steps]
+        )
+        train_dataloader, self.scheduler = self.accelerator.prepare(
+            train_dataloader, self.scheduler
+        )  # actual steps = 1 gpu steps / gpus
+        start_step = self.load_checkpoint()
+        global_step = start_step
+
+        valid_dataloader = self.accelerator.prepare(valid_dataloader)
+
+        if exists(resumable_with_seed):
+            orig_epoch_step = len(train_dataloader)
+            skipped_epoch = int(start_step // orig_epoch_step)
+            skipped_batch = start_step % orig_epoch_step
+            skipped_dataloader = self.accelerator.skip_first_batches(train_dataloader, num_batches=skipped_batch)
+        else:
+            skipped_epoch = 0
+
+        for epoch in range(skipped_epoch, self.epochs):
+            self.model.train()
+            if exists(resumable_with_seed) and epoch == skipped_epoch:
+                progress_bar = tqdm(
+                    skipped_dataloader,
+                    desc=f"Epoch {epoch+1}/{self.epochs}",
+                    unit="step",
+                    disable=not self.accelerator.is_local_main_process,
+                    initial=skipped_batch,
+                    total=orig_epoch_step,
+                )
+            else:
+                progress_bar = tqdm(
+                    train_dataloader,
+                    desc=f"Epoch {epoch+1}/{self.epochs}",
+                    unit="step",
+                    disable=not self.accelerator.is_local_main_process,
+                )
+
+            for batch in progress_bar:
+                with self.accelerator.accumulate(self.model):
+                    # Inputs
+                    prompt_mel = batch['pmt_mel_specs'].permute(0, 2, 1) # (B, L_mel, D)
+                    prompt_text = batch['pmt_text']
+                    text = batch['text']
+
+                    target_ids = list_str_to_idx(text, self.vocab_char_map).to(prompt_mel.device)
+                    target_ids = target_ids.masked_fill(target_ids==-1, vocab_size)
+
+                    prompt_ids = list_str_to_idx(prompt_text, self.vocab_char_map).to(prompt_mel.device)
+                    prompt_ids = prompt_ids.masked_fill(prompt_ids==-1, vocab_size)
+
+                    # Targets
+                    tar_lengths = batch['mel_lengths']
+
+                    # Forward
+                    predictions = SLP(target_ids=target_ids, prompt_ids=prompt_ids, prompt_mel=prompt_mel) # (B, C)
+
+                    if self.loss_fn == 'CE':
+                        tar_length_labels = (tar_lengths // self.n_frame_per_class) \
+                            .clamp(min=0, max=self.n_class-1) # [0, 1, ..., n_class-1]
+                        est_length_logtis = predictions
+                        est_length_labels = torch.argmax(est_length_logtis, dim=-1)
+                        loss = F.cross_entropy(est_length_logtis, tar_length_labels)
+                        
+                        with torch.no_grad():
+                            est_lengths = est_length_labels * self.n_frame_per_class
+                            frame_error = (est_lengths.float() - tar_lengths.float()).abs().mean()
+                            sec_error = frame_error * 256 / 24000
+
+                        log_dict = {
+                            'loss': loss.item(), 
+                            'loss_CE': loss.item(), 
+                            'sec_error': sec_error.item(),
+                            'lr': self.scheduler.get_last_lr()[0]
+                         }
+
+                    else:
+                        raise NotImplementedError(self.loss_fn)
+
+        
+                    self.accelerator.backward(loss)
+
+                    if self.max_grad_norm > 0 and self.accelerator.sync_gradients:
+                        self.accelerator.clip_grad_norm_(self.model.parameters(), self.max_grad_norm)
+
+                    self.optimizer.step()
+                    self.scheduler.step()
+                    self.optimizer.zero_grad()
+
+                global_step += 1
+
+                if self.accelerator.is_local_main_process:
+                    self.accelerator.log(log_dict, step=global_step)
+                progress_bar.set_postfix(step=str(global_step), loss=loss.item())
+
+                if global_step % (self.save_per_updates * self.grad_accumulation_steps) == 0:
+                    self.save_checkpoint(global_step)
+                    # if self.log_samples and self.accelerator.is_local_main_process:
+                    # Run validation at the end of each epoch (only on the main process)
+                    if self.accelerator.is_local_main_process:
+                        self.validate(valid_dataloader, global_step)
+                # if global_step % self.last_per_steps == 0:
+                #     self.save_checkpoint(global_step, last=True)
+
+        self.save_checkpoint(global_step, last=True)
+        self.accelerator.end_training()

ecapa_tdnn.py

@@ -0,0 +1,268 @@
+# part of the code is borrowed from https://github.com/lawlict/ECAPA-TDNN
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import torchaudio.transforms as trans
+from ctcmodel import ConformerCTC
+# from ctcmodel_nopool import ConformerCTC as ConformerCTCNoPool
+from pathlib import Path
+
+''' Res2Conv1d + BatchNorm1d + ReLU
+'''
+
+
+class Res2Conv1dReluBn(nn.Module):
+    '''
+    in_channels == out_channels == channels
+    '''
+
+    def __init__(self, channels, kernel_size=1, stride=1, padding=0, dilation=1, bias=True, scale=4):
+        super().__init__()
+        assert channels % scale == 0, "{} % {} != 0".format(channels, scale)
+        self.scale = scale
+        self.width = channels // scale
+        self.nums = scale if scale == 1 else scale - 1
+
+        self.convs = []
+        self.bns = []
+        for i in range(self.nums):
+            self.convs.append(nn.Conv1d(self.width, self.width, kernel_size, stride, padding, dilation, bias=bias))
+            self.bns.append(nn.BatchNorm1d(self.width))
+        self.convs = nn.ModuleList(self.convs)
+        self.bns = nn.ModuleList(self.bns)
+
+    def forward(self, x):
+        out = []
+        spx = torch.split(x, self.width, 1)
+        for i in range(self.nums):
+            if i == 0:
+                sp = spx[i]
+            else:
+                sp = sp + spx[i]
+            # Order: conv -> relu -> bn
+            sp = self.convs[i](sp)
+            sp = self.bns[i](F.relu(sp))
+            out.append(sp)
+        if self.scale != 1:
+            out.append(spx[self.nums])
+        out = torch.cat(out, dim=1)
+
+        return out
+
+
+''' Conv1d + BatchNorm1d + ReLU
+'''
+
+
+class Conv1dReluBn(nn.Module):
+    def __init__(self, in_channels, out_channels, kernel_size=1, stride=1, padding=0, dilation=1, bias=True):
+        super().__init__()
+        self.conv = nn.Conv1d(in_channels, out_channels, kernel_size, stride, padding, dilation, bias=bias)
+        self.bn = nn.BatchNorm1d(out_channels)
+
+    def forward(self, x):
+        return self.bn(F.relu(self.conv(x)))
+
+
+''' The SE connection of 1D case.
+'''
+
+
+class SE_Connect(nn.Module):
+    def __init__(self, channels, se_bottleneck_dim=128):
+        super().__init__()
+        self.linear1 = nn.Linear(channels, se_bottleneck_dim)
+        self.linear2 = nn.Linear(se_bottleneck_dim, channels)
+
+    def forward(self, x):
+        out = x.mean(dim=2)
+        out = F.relu(self.linear1(out))
+        out = torch.sigmoid(self.linear2(out))
+        out = x * out.unsqueeze(2)
+
+        return out
+
+
+''' SE-Res2Block of the ECAPA-TDNN architecture.
+'''
+
+class SE_Res2Block(nn.Module):
+    def __init__(self, in_channels, out_channels, kernel_size, stride, padding, dilation, scale, se_bottleneck_dim):
+        super().__init__()
+        self.Conv1dReluBn1 = Conv1dReluBn(in_channels, out_channels, kernel_size=1, stride=1, padding=0)
+        self.Res2Conv1dReluBn = Res2Conv1dReluBn(out_channels, kernel_size, stride, padding, dilation, scale=scale)
+        self.Conv1dReluBn2 = Conv1dReluBn(out_channels, out_channels, kernel_size=1, stride=1, padding=0)
+        self.SE_Connect = SE_Connect(out_channels, se_bottleneck_dim)
+
+        self.shortcut = None
+        if in_channels != out_channels:
+            self.shortcut = nn.Conv1d(
+                in_channels=in_channels,
+                out_channels=out_channels,
+                kernel_size=1,
+            )
+
+    def forward(self, x):
+        residual = x
+        if self.shortcut:
+            residual = self.shortcut(x)
+
+        x = self.Conv1dReluBn1(x)
+        x = self.Res2Conv1dReluBn(x)
+        x = self.Conv1dReluBn2(x)
+        x = self.SE_Connect(x)
+
+        return x + residual
+
+
+''' Attentive weighted mean and standard deviation pooling.
+'''
+
+class AttentiveStatsPool(nn.Module):
+    def __init__(self, in_dim, attention_channels=128, global_context_att=False):
+        super().__init__()
+        self.global_context_att = global_context_att
+
+        # Use Conv1d with stride == 1 rather than Linear, then we don't need to transpose inputs.
+        if global_context_att:
+            self.linear1 = nn.Conv1d(in_dim * 3, attention_channels, kernel_size=1)  # equals W and b in the paper
+        else:
+            self.linear1 = nn.Conv1d(in_dim, attention_channels, kernel_size=1)  # equals W and b in the paper
+        self.linear2 = nn.Conv1d(attention_channels, in_dim, kernel_size=1)  # equals V and k in the paper
+
+    def forward(self, x):
+
+        if self.global_context_att:
+            context_mean = torch.mean(x, dim=-1, keepdim=True).expand_as(x)
+            context_std = torch.sqrt(torch.var(x, dim=-1, keepdim=True) + 1e-10).expand_as(x)
+            x_in = torch.cat((x, context_mean, context_std), dim=1)
+        else:
+            x_in = x
+
+        # DON'T use ReLU here! In experiments, I find ReLU hard to converge.
+        alpha = torch.tanh(self.linear1(x_in))
+        # alpha = F.relu(self.linear1(x_in))
+        alpha = torch.softmax(self.linear2(alpha), dim=2)
+        mean = torch.sum(alpha * x, dim=2)
+        residuals = torch.sum(alpha * (x ** 2), dim=2) - mean ** 2
+        std = torch.sqrt(residuals.clamp(min=1e-9))
+        return torch.cat([mean, std], dim=1)
+
+
+class ECAPA_TDNN(nn.Module):
+    def __init__(self, channels=512, emb_dim=512, 
+        global_context_att=False, use_fp16=True,
+        ctc_cls=ConformerCTC,
+        ctc_path='/data4/F5TTS/ckpts/F5TTS_norm_ASR_vocos_pinyin_Emilia_ZH_EN/model_last.pt',
+        ctc_args={'vocab_size': 2545, 'mel_dim': 100, 'num_heads': 8, 'd_hid': 512, 'nlayers': 6},
+        ctc_no_grad=False
+    ):
+        super().__init__()
+        if ctc_path != None:
+            ctc_path = Path(ctc_path)
+            model = ctc_cls(**ctc_args)
+            state_dict = torch.load(ctc_path, map_location='cpu')
+            model.load_state_dict(state_dict['model_state_dict'])
+            print(f"Initialized pretrained ConformerCTC backbone from {ctc_path}.")
+        else:
+            raise ValueError(ctc_path)
+
+        self.ctc_model = model
+        self.ctc_model.out.requires_grad_(False)
+    
+        if ctc_cls == ConformerCTC:
+            self.feat_num = ctc_args['nlayers'] + 2 + 1
+        # elif ctc_cls == ConformerCTCNoPool:
+        #     self.feat_num = ctc_args['nlayers'] + 1
+        else:
+            raise ValueError(ctc_cls)
+        feat_dim = ctc_args['d_hid']
+
+        self.emb_dim = emb_dim
+        
+        self.feature_weight = nn.Parameter(torch.zeros(self.feat_num))
+        self.instance_norm = nn.InstanceNorm1d(feat_dim)
+
+        # self.channels = [channels] * 4 + [channels * 3]
+        self.channels = [channels] * 4 + [1536]
+
+        self.layer1 = Conv1dReluBn(feat_dim, self.channels[0], kernel_size=5, padding=2)
+        self.layer2 = SE_Res2Block(self.channels[0], self.channels[1], kernel_size=3, stride=1, padding=2, dilation=2, scale=8, se_bottleneck_dim=128)
+        self.layer3 = SE_Res2Block(self.channels[1], self.channels[2], kernel_size=3, stride=1, padding=3, dilation=3, scale=8, se_bottleneck_dim=128)
+        self.layer4 = SE_Res2Block(self.channels[2], self.channels[3], kernel_size=3, stride=1, padding=4, dilation=4, scale=8, se_bottleneck_dim=128)
+
+        # self.conv = nn.Conv1d(self.channels[-1], self.channels[-1], kernel_size=1)
+        cat_channels = channels * 3
+        self.conv = nn.Conv1d(cat_channels, self.channels[-1], kernel_size=1)
+        self.pooling = AttentiveStatsPool(self.channels[-1], attention_channels=128, global_context_att=global_context_att)
+        self.bn = nn.BatchNorm1d(self.channels[-1] * 2)
+        self.linear = nn.Linear(self.channels[-1] * 2, emb_dim)
+
+        if ctc_no_grad:
+            for param in self.ctc_model.parameters():
+                param.requires_grad = False
+            self.ctc_model = self.ctc_model.eval()
+        else:
+            self.ctc_model = self.ctc_model.train()
+        self.ctc_no_grad = ctc_no_grad
+        print('ctc_no_grad: ', self.ctc_no_grad)
+
+    def forward(self, latent, input_lengths,  return_asr=False):
+        if self.ctc_no_grad:
+            with torch.no_grad():
+                asr, h = self.ctc_model(latent, input_lengths)
+        else:
+            asr, h = self.ctc_model(latent, input_lengths)
+        
+        x = torch.stack(h, dim=0)
+        norm_weights = F.softmax(self.feature_weight, dim=-1).unsqueeze(-1).unsqueeze(-1).unsqueeze(-1)
+        x = (norm_weights * x).sum(dim=0)
+        x = x + 1e-6
+        # x = torch.transpose(x, 1, 2) + 1e-6
+            
+        x = self.instance_norm(x)
+        # x = torch.transpose(x, 1, 2)
+
+        out1 = self.layer1(x)
+        out2 = self.layer2(out1)
+        out3 = self.layer3(out2)
+        out4 = self.layer4(out3)
+
+        out = torch.cat([out2, out3, out4], dim=1)
+        out = F.relu(self.conv(out))
+        out = self.bn(self.pooling(out))
+        out = self.linear(out)
+
+        if return_asr:
+            return out, asr
+        return out
+
+if __name__ == "__main__":
+    from diffspeech.ldm.model import DiT
+    from diffspeech.data.collate import get_mask_from_lengths
+    from diffspeech.tools.text.vocab import IPA
+
+    bsz = 3
+
+    # Sample ipa
+    ipa_lens = torch.randint(10, 50, (bsz,)).cuda()
+    ipa_mask = get_mask_from_lengths(ipa_lens).cuda()
+    ipa = torch.randint(0, len(IPA.vocab), (bsz, ipa_mask.size(-1))).cuda()
+
+    # Sample latent
+    latent_lens = torch.randint(50, 250, (bsz,)).cuda()
+    latent_mask = get_mask_from_lengths(latent_lens).cuda()
+    latent = torch.randn(bsz, latent_mask.size(-1), 64).cuda()
+
+    # Sample prompt
+    prompt_mask = get_mask_from_lengths(
+        (latent_lens * 0.25).long(), max_len=latent_mask.size(-1)
+    ).cuda()
+    prompt_latent = latent * prompt_mask.unsqueeze(-1)
+
+    model = ECAPA_TDNN(emb_dim=512).cuda()
+
+    emb = model(latent, latent_mask.sum(axis=-1))
+
+    print(emb.shape)

f5_tts/api.py

@@ -0,0 +1,164 @@
+import random
+import sys
+from importlib.resources import files
+
+import soundfile as sf
+import tqdm
+from cached_path import cached_path
+from hydra.utils import get_class
+from omegaconf import OmegaConf
+
+from f5_tts.infer.utils_infer import (
+    infer_process,
+    load_model,
+    load_vocoder,
+    preprocess_ref_audio_text,
+    remove_silence_for_generated_wav,
+    save_spectrogram,
+    transcribe,
+)
+from f5_tts.model.utils import seed_everything
+
+
+class F5TTS:
+    def __init__(
+        self,
+        model="F5TTS_v1_Base",
+        ckpt_file="",
+        vocab_file="",
+        ode_method="euler",
+        use_ema=True,
+        vocoder_local_path=None,
+        device=None,
+        hf_cache_dir=None,
+    ):
+        model_cfg = OmegaConf.load(str(files("f5_tts").joinpath(f"configs/{model}.yaml")))
+        model_cls = get_class(f"f5_tts.model.{model_cfg.model.backbone}")
+        model_arc = model_cfg.model.arch
+
+        self.mel_spec_type = model_cfg.model.mel_spec.mel_spec_type
+        self.target_sample_rate = model_cfg.model.mel_spec.target_sample_rate
+
+        self.ode_method = ode_method
+        self.use_ema = use_ema
+
+        if device is not None:
+            self.device = device
+        else:
+            import torch
+
+            self.device = (
+                "cuda"
+                if torch.cuda.is_available()
+                else "xpu"
+                if torch.xpu.is_available()
+                else "mps"
+                if torch.backends.mps.is_available()
+                else "cpu"
+            )
+
+        # Load models
+        self.vocoder = load_vocoder(
+            self.mel_spec_type, vocoder_local_path is not None, vocoder_local_path, self.device, hf_cache_dir
+        )
+
+        repo_name, ckpt_step, ckpt_type = "F5-TTS", 1250000, "safetensors"
+
+        # override for previous models
+        if model == "F5TTS_Base":
+            if self.mel_spec_type == "vocos":
+                ckpt_step = 1200000
+            elif self.mel_spec_type == "bigvgan":
+                model = "F5TTS_Base_bigvgan"
+                ckpt_type = "pt"
+        elif model == "E2TTS_Base":
+            repo_name = "E2-TTS"
+            ckpt_step = 1200000
+
+        if not ckpt_file:
+            ckpt_file = str(
+                cached_path(f"hf://SWivid/{repo_name}/{model}/model_{ckpt_step}.{ckpt_type}", cache_dir=hf_cache_dir)
+            )
+        self.ema_model = load_model(
+            model_cls, model_arc, ckpt_file, self.mel_spec_type, vocab_file, self.ode_method, self.use_ema, self.device
+        )
+
+    def transcribe(self, ref_audio, language=None):
+        return transcribe(ref_audio, language)
+
+    def export_wav(self, wav, file_wave, remove_silence=False):
+        sf.write(file_wave, wav, self.target_sample_rate)
+
+        if remove_silence:
+            remove_silence_for_generated_wav(file_wave)
+
+    def export_spectrogram(self, spec, file_spec):
+        save_spectrogram(spec, file_spec)
+
+    def infer(
+        self,
+        ref_file,
+        ref_text,
+        gen_text,
+        show_info=print,
+        progress=tqdm,
+        target_rms=0.1,
+        cross_fade_duration=0.15,
+        sway_sampling_coef=-1,
+        cfg_strength=2,
+        nfe_step=32,
+        speed=1.0,
+        fix_duration=None,
+        remove_silence=False,
+        file_wave=None,
+        file_spec=None,
+        seed=None,
+    ):
+        if seed is None:
+            seed = random.randint(0, sys.maxsize)
+        seed_everything(seed)
+        self.seed = seed
+
+        ref_file, ref_text = preprocess_ref_audio_text(ref_file, ref_text)
+
+        wav, sr, spec = infer_process(
+            ref_file,
+            ref_text,
+            gen_text,
+            self.ema_model,
+            self.vocoder,
+            self.mel_spec_type,
+            show_info=show_info,
+            progress=progress,
+            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,
+            device=self.device,
+        )
+
+        if file_wave is not None:
+            self.export_wav(wav, file_wave, remove_silence)
+
+        if file_spec is not None:
+            self.export_spectrogram(spec, file_spec)
+
+        return wav, sr, spec
+
+
+if __name__ == "__main__":
+    f5tts = F5TTS()
+
+    wav, sr, spec = f5tts.infer(
+        ref_file=str(files("f5_tts").joinpath("infer/examples/basic/basic_ref_en.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.""",
+        file_wave=str(files("f5_tts").joinpath("../../tests/api_out.wav")),
+        file_spec=str(files("f5_tts").joinpath("../../tests/api_out.png")),
+        seed=None,
+    )
+
+    print("seed :", f5tts.seed)

f5_tts/configs/E2TTS_Base.yaml

@@ -0,0 +1,49 @@
+hydra:
+  run:
+    dir: ckpts/${model.name}_${model.mel_spec.mel_spec_type}_${model.tokenizer}_${datasets.name}/${now:%Y-%m-%d}/${now:%H-%M-%S}
+
+datasets:
+  name: Emilia_ZH_EN  # dataset name
+  batch_size_per_gpu: 38400  # 8 GPUs, 8 * 38400 = 307200
+  batch_size_type: frame  # frame | sample
+  max_samples: 64  # max sequences per batch if use frame-wise batch_size. we set 32 for small models, 64 for base models
+  num_workers: 16
+
+optim:
+  epochs: 11
+  learning_rate: 7.5e-5
+  num_warmup_updates: 20000  # warmup updates
+  grad_accumulation_steps: 1  # note: updates = steps / grad_accumulation_steps
+  max_grad_norm: 1.0  # gradient clipping
+  bnb_optimizer: False  # use bnb 8bit AdamW optimizer or not
+
+model:
+  name: E2TTS_Base
+  tokenizer: pinyin
+  tokenizer_path: null  # if 'custom' tokenizer, define the path want to use (should be vocab.txt)
+  backbone: UNetT
+  arch:
+    dim: 1024
+    depth: 24
+    heads: 16
+    ff_mult: 4
+    text_mask_padding: False
+    pe_attn_head: 1
+  mel_spec:
+    target_sample_rate: 24000
+    n_mel_channels: 100
+    hop_length: 256
+    win_length: 1024
+    n_fft: 1024
+    mel_spec_type: vocos  # vocos | bigvgan
+  vocoder:
+    is_local: False  # use local offline ckpt or not
+    local_path: null  # local vocoder path
+
+ckpts:
+  logger: wandb  # wandb | tensorboard | null
+  log_samples: True  # infer random sample per save checkpoint. wip, normal to fail with extra long samples
+  save_per_updates: 50000  # save checkpoint per updates
+  keep_last_n_checkpoints: -1  # -1 to keep all, 0 to not save intermediate, > 0 to keep last N checkpoints
+  last_per_updates: 5000  # save last checkpoint per updates
+  save_dir: ckpts/${model.name}_${model.mel_spec.mel_spec_type}_${model.tokenizer}_${datasets.name}

f5_tts/configs/E2TTS_Small.yaml

@@ -0,0 +1,49 @@
+hydra:
+  run:
+    dir: ckpts/${model.name}_${model.mel_spec.mel_spec_type}_${model.tokenizer}_${datasets.name}/${now:%Y-%m-%d}/${now:%H-%M-%S}
+
+datasets:
+  name: Emilia_ZH_EN
+  batch_size_per_gpu: 38400  # 8 GPUs, 8 * 38400 = 307200
+  batch_size_type: frame  # frame | sample
+  max_samples: 64  # max sequences per batch if use frame-wise batch_size. we set 32 for small models, 64 for base models
+  num_workers: 16
+
+optim:
+  epochs: 11
+  learning_rate: 7.5e-5
+  num_warmup_updates: 20000  # warmup updates
+  grad_accumulation_steps: 1  # note: updates = steps / grad_accumulation_steps
+  max_grad_norm: 1.0
+  bnb_optimizer: False  
+
+model:
+  name: E2TTS_Small
+  tokenizer: pinyin
+  tokenizer_path: null  # if 'custom' tokenizer, define the path want to use (should be vocab.txt)
+  backbone: UNetT
+  arch:
+    dim: 768
+    depth: 20
+    heads: 12
+    ff_mult: 4
+    text_mask_padding: False
+    pe_attn_head: 1
+  mel_spec:
+    target_sample_rate: 24000
+    n_mel_channels: 100
+    hop_length: 256
+    win_length: 1024
+    n_fft: 1024
+    mel_spec_type: vocos  # vocos | bigvgan
+  vocoder:
+    is_local: False  # use local offline ckpt or not
+    local_path: null  # local vocoder path
+
+ckpts:
+  logger: wandb  # wandb | tensorboard | null
+  log_samples: True  # infer random sample per save checkpoint. wip, normal to fail with extra long samples
+  save_per_updates: 50000  # save checkpoint per updates
+  keep_last_n_checkpoints: -1  # -1 to keep all, 0 to not save intermediate, > 0 to keep last N checkpoints
+  last_per_updates: 5000  # save last checkpoint per updates
+  save_dir: ckpts/${model.name}_${model.mel_spec.mel_spec_type}_${model.tokenizer}_${datasets.name}

f5_tts/configs/F5TTS_Base.yaml

@@ -0,0 +1,54 @@
+hydra:
+  run:
+    dir: ckpts/${model.name}_${model.mel_spec.mel_spec_type}_${model.tokenizer}_${datasets.name}/${now:%Y-%m-%d}/${now:%H-%M-%S}
+
+datasets:
+  name: Emilia_ZH_EN  # dataset name
+  batch_size_per_gpu: 38400  # 8 GPUs, 8 * 38400 = 307200
+  batch_size_type: frame  # frame | sample
+  max_samples: 64  # max sequences per batch if use frame-wise batch_size. we set 32 for small models, 64 for base models
+  num_workers: 16
+
+optim:
+  epochs: 11
+  learning_rate: 7.5e-5
+  num_warmup_updates: 20000  # warmup updates
+  grad_accumulation_steps: 1  # note: updates = steps / grad_accumulation_steps
+  max_grad_norm: 1.0  # gradient clipping
+  bnb_optimizer: False  # use bnb 8bit AdamW optimizer or not
+
+model:
+  name: F5TTS_Base  # model name
+  tokenizer: pinyin  # tokenizer type
+  tokenizer_path: null  # if 'custom' tokenizer, define the path want to use (should be vocab.txt)
+  backbone: DiT
+  arch:
+    dim: 1024
+    depth: 22
+    heads: 16
+    ff_mult: 2
+    text_dim: 512
+    text_mask_padding: False
+    conv_layers: 4
+    pe_attn_head: 1
+    attn_backend: torch  # torch | flash_attn
+    attn_mask_enabled: False
+    checkpoint_activations: False  # recompute activations and save memory for extra compute
+  mel_spec:
+    target_sample_rate: 24000
+    n_mel_channels: 100
+    hop_length: 256
+    win_length: 1024
+    n_fft: 1024
+    mel_spec_type: vocos  # vocos | bigvgan
+  vocoder:
+    is_local: False  # use local offline ckpt or not
+    local_path: null  # local vocoder path
+
+ckpts:
+  logger: wandb  # wandb | tensorboard | null
+  log_samples: True  # infer random sample per save checkpoint. wip, normal to fail with extra long samples
+  save_per_updates: 50000  # save checkpoint per updates
+  keep_last_n_checkpoints: -1  # -1 to keep all, 0 to not save intermediate, > 0 to keep last N checkpoints
+  last_per_updates: 5000  # save last checkpoint per updates
+  save_dir: ckpts/${model.name}_${model.mel_spec.mel_spec_type}_${model.tokenizer}_${datasets.name}

f5_tts/configs/F5TTS_Small.yaml

@@ -0,0 +1,54 @@
+hydra:
+  run:
+    dir: ckpts/${model.name}_${model.mel_spec.mel_spec_type}_${model.tokenizer}_${datasets.name}/${now:%Y-%m-%d}/${now:%H-%M-%S}
+
+datasets:
+  name: Emilia_ZH_EN
+  batch_size_per_gpu: 38400  # 8 GPUs, 8 * 38400 = 307200
+  batch_size_type: frame  # frame | sample
+  max_samples: 64  # max sequences per batch if use frame-wise batch_size. we set 32 for small models, 64 for base models
+  num_workers: 16
+
+optim:
+  epochs: 11  # only suitable for Emilia, if you want to train it on LibriTTS, set epoch 686
+  learning_rate: 7.5e-5
+  num_warmup_updates: 20000  # warmup updates
+  grad_accumulation_steps: 1  # note: updates = steps / grad_accumulation_steps
+  max_grad_norm: 1.0  # gradient clipping
+  bnb_optimizer: False  # use bnb 8bit AdamW optimizer or not
+
+model:
+  name: F5TTS_Small
+  tokenizer: pinyin
+  tokenizer_path: null  # if 'custom' tokenizer, define the path want to use (should be vocab.txt)
+  backbone: DiT
+  arch:
+    dim: 768
+    depth: 18
+    heads: 12
+    ff_mult: 2
+    text_dim: 512
+    text_mask_padding: False
+    conv_layers: 4
+    pe_attn_head: 1
+    attn_backend: torch  # torch | flash_attn
+    attn_mask_enabled: False
+    checkpoint_activations: False  # recompute activations and save memory for extra compute
+  mel_spec:
+    target_sample_rate: 24000
+    n_mel_channels: 100
+    hop_length: 256
+    win_length: 1024
+    n_fft: 1024
+    mel_spec_type: vocos  # vocos | bigvgan
+  vocoder:
+    is_local: False  # use local offline ckpt or not
+    local_path: null  # local vocoder path
+
+ckpts:
+  logger: wandb  # wandb | tensorboard | null
+  log_samples: True  # infer random sample per save checkpoint. wip, normal to fail with extra long samples
+  save_per_updates: 50000  # save checkpoint per updates
+  keep_last_n_checkpoints: -1  # -1 to keep all, 0 to not save intermediate, > 0 to keep last N checkpoints
+  last_per_updates: 5000  # save last checkpoint per updates
+  save_dir: ckpts/${model.name}_${model.mel_spec.mel_spec_type}_${model.tokenizer}_${datasets.name}

f5_tts/configs/F5TTS_v1_Base.yaml

@@ -0,0 +1,55 @@
+hydra:
+  run:
+    dir: ckpts/${model.name}_${model.mel_spec.mel_spec_type}_${model.tokenizer}_${datasets.name}/${now:%Y-%m-%d}/${now:%H-%M-%S}
+
+datasets:
+  name: Emilia_ZH_EN  # dataset name
+  batch_size_per_gpu: 38400  # 8 GPUs, 8 * 38400 = 307200
+  batch_size_type: frame  # frame | sample
+  max_samples: 64  # max sequences per batch if use frame-wise batch_size. we set 32 for small models, 64 for base models
+  num_workers: 16
+
+optim:
+  epochs: 11
+  learning_rate: 7.5e-5
+  num_warmup_updates: 20000  # warmup updates
+  grad_accumulation_steps: 1  # note: updates = steps / grad_accumulation_steps
+  max_grad_norm: 1.0  # gradient clipping
+  bnb_optimizer: False  # use bnb 8bit AdamW optimizer or not
+
+model:
+  name: F5TTS_v1_Base  # model name
+  tokenizer: pinyin  # tokenizer type
+  tokenizer_path: null  # if 'custom' tokenizer, define the path want to use (should be vocab.txt)
+  backbone: DiT
+  arch:
+    dim: 1024
+    depth: 22
+    heads: 16
+    ff_mult: 2
+    text_dim: 512
+    text_mask_padding: True
+    qk_norm: null  # null | rms_norm
+    conv_layers: 4
+    pe_attn_head: null
+    attn_backend: torch  # torch | flash_attn
+    attn_mask_enabled: False
+    checkpoint_activations: False  # recompute activations and save memory for extra compute
+  mel_spec:
+    target_sample_rate: 24000
+    n_mel_channels: 100
+    hop_length: 256
+    win_length: 1024
+    n_fft: 1024
+    mel_spec_type: vocos  # vocos | bigvgan
+  vocoder:
+    is_local: False  # use local offline ckpt or not
+    local_path: null  # local vocoder path
+
+ckpts:
+  logger: wandb  # wandb | tensorboard | null
+  log_samples: True  # infer random sample per save checkpoint. wip, normal to fail with extra long samples
+  save_per_updates: 50000  # save checkpoint per updates
+  keep_last_n_checkpoints: -1  # -1 to keep all, 0 to not save intermediate, > 0 to keep last N checkpoints
+  last_per_updates: 5000  # save last checkpoint per updates
+  save_dir: ckpts/${model.name}_${model.mel_spec.mel_spec_type}_${model.tokenizer}_${datasets.name}

f5_tts/eval/README.md

@@ -0,0 +1,52 @@
+
+# Evaluation
+
+Install packages for evaluation:
+
+```bash
+pip install -e .[eval]
+```
+
+## Generating Samples for Evaluation
+
+### Prepare Test Datasets
+
+1. *Seed-TTS testset*: 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 *LibriSpeech test-clean* data in `src/f5_tts/eval/eval_infer_batch.py`
+5. Our filtered LibriSpeech-PC 4-10s subset: `data/librispeech_pc_test_clean_cross_sentence.lst`
+
+### 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 src/f5_tts/eval/eval_infer_batch.sh
+```
+
+## Objective Evaluation on Generated Results
+
+### 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).
+
+Then update in the following scripts with the paths you put evaluation model ckpts to.
+
+### Objective Evaluation
+
+Update the path with your batch-inferenced results, and carry out WER / SIM / UTMOS evaluations:
+```bash
+# 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 [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
+```

f5_tts/eval/ecapa_tdnn.py

@@ -0,0 +1,331 @@
+# just for speaker similarity evaluation, third-party code
+
+# From https://github.com/microsoft/UniSpeech/blob/main/downstreams/speaker_verification/models/
+# part of the code is borrowed from https://github.com/lawlict/ECAPA-TDNN
+
+import os
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+""" Res2Conv1d + BatchNorm1d + ReLU
+"""
+
+
+class Res2Conv1dReluBn(nn.Module):
+    """
+    in_channels == out_channels == channels
+    """
+
+    def __init__(self, channels, kernel_size=1, stride=1, padding=0, dilation=1, bias=True, scale=4):
+        super().__init__()
+        assert channels % scale == 0, "{} % {} != 0".format(channels, scale)
+        self.scale = scale
+        self.width = channels // scale
+        self.nums = scale if scale == 1 else scale - 1
+
+        self.convs = []
+        self.bns = []
+        for i in range(self.nums):
+            self.convs.append(nn.Conv1d(self.width, self.width, kernel_size, stride, padding, dilation, bias=bias))
+            self.bns.append(nn.BatchNorm1d(self.width))
+        self.convs = nn.ModuleList(self.convs)
+        self.bns = nn.ModuleList(self.bns)
+
+    def forward(self, x):
+        out = []
+        spx = torch.split(x, self.width, 1)
+        for i in range(self.nums):
+            if i == 0:
+                sp = spx[i]
+            else:
+                sp = sp + spx[i]
+            # Order: conv -> relu -> bn
+            sp = self.convs[i](sp)
+            sp = self.bns[i](F.relu(sp))
+            out.append(sp)
+        if self.scale != 1:
+            out.append(spx[self.nums])
+        out = torch.cat(out, dim=1)
+
+        return out
+
+
+""" Conv1d + BatchNorm1d + ReLU
+"""
+
+
+class Conv1dReluBn(nn.Module):
+    def __init__(self, in_channels, out_channels, kernel_size=1, stride=1, padding=0, dilation=1, bias=True):
+        super().__init__()
+        self.conv = nn.Conv1d(in_channels, out_channels, kernel_size, stride, padding, dilation, bias=bias)
+        self.bn = nn.BatchNorm1d(out_channels)
+
+    def forward(self, x):
+        return self.bn(F.relu(self.conv(x)))
+
+
+""" The SE connection of 1D case.
+"""
+
+
+class SE_Connect(nn.Module):
+    def __init__(self, channels, se_bottleneck_dim=128):
+        super().__init__()
+        self.linear1 = nn.Linear(channels, se_bottleneck_dim)
+        self.linear2 = nn.Linear(se_bottleneck_dim, channels)
+
+    def forward(self, x):
+        out = x.mean(dim=2)
+        out = F.relu(self.linear1(out))
+        out = torch.sigmoid(self.linear2(out))
+        out = x * out.unsqueeze(2)
+
+        return out
+
+
+""" SE-Res2Block of the ECAPA-TDNN architecture.
+"""
+
+# def SE_Res2Block(channels, kernel_size, stride, padding, dilation, scale):
+#     return nn.Sequential(
+#         Conv1dReluBn(channels, 512, kernel_size=1, stride=1, padding=0),
+#         Res2Conv1dReluBn(512, kernel_size, stride, padding, dilation, scale=scale),
+#         Conv1dReluBn(512, channels, kernel_size=1, stride=1, padding=0),
+#         SE_Connect(channels)
+#     )
+
+
+class SE_Res2Block(nn.Module):
+    def __init__(self, in_channels, out_channels, kernel_size, stride, padding, dilation, scale, se_bottleneck_dim):
+        super().__init__()
+        self.Conv1dReluBn1 = Conv1dReluBn(in_channels, out_channels, kernel_size=1, stride=1, padding=0)
+        self.Res2Conv1dReluBn = Res2Conv1dReluBn(out_channels, kernel_size, stride, padding, dilation, scale=scale)
+        self.Conv1dReluBn2 = Conv1dReluBn(out_channels, out_channels, kernel_size=1, stride=1, padding=0)
+        self.SE_Connect = SE_Connect(out_channels, se_bottleneck_dim)
+
+        self.shortcut = None
+        if in_channels != out_channels:
+            self.shortcut = nn.Conv1d(
+                in_channels=in_channels,
+                out_channels=out_channels,
+                kernel_size=1,
+            )
+
+    def forward(self, x):
+        residual = x
+        if self.shortcut:
+            residual = self.shortcut(x)
+
+        x = self.Conv1dReluBn1(x)
+        x = self.Res2Conv1dReluBn(x)
+        x = self.Conv1dReluBn2(x)
+        x = self.SE_Connect(x)
+
+        return x + residual
+
+
+""" Attentive weighted mean and standard deviation pooling.
+"""
+
+
+class AttentiveStatsPool(nn.Module):
+    def __init__(self, in_dim, attention_channels=128, global_context_att=False):
+        super().__init__()
+        self.global_context_att = global_context_att
+
+        # Use Conv1d with stride == 1 rather than Linear, then we don't need to transpose inputs.
+        if global_context_att:
+            self.linear1 = nn.Conv1d(in_dim * 3, attention_channels, kernel_size=1)  # equals W and b in the paper
+        else:
+            self.linear1 = nn.Conv1d(in_dim, attention_channels, kernel_size=1)  # equals W and b in the paper
+        self.linear2 = nn.Conv1d(attention_channels, in_dim, kernel_size=1)  # equals V and k in the paper
+
+    def forward(self, x):
+        if self.global_context_att:
+            context_mean = torch.mean(x, dim=-1, keepdim=True).expand_as(x)
+            context_std = torch.sqrt(torch.var(x, dim=-1, keepdim=True) + 1e-10).expand_as(x)
+            x_in = torch.cat((x, context_mean, context_std), dim=1)
+        else:
+            x_in = x
+
+        # DON'T use ReLU here! In experiments, I find ReLU hard to converge.
+        alpha = torch.tanh(self.linear1(x_in))
+        # alpha = F.relu(self.linear1(x_in))
+        alpha = torch.softmax(self.linear2(alpha), dim=2)
+        mean = torch.sum(alpha * x, dim=2)
+        residuals = torch.sum(alpha * (x**2), dim=2) - mean**2
+        std = torch.sqrt(residuals.clamp(min=1e-9))
+        return torch.cat([mean, std], dim=1)
+
+
+class ECAPA_TDNN(nn.Module):
+    def __init__(
+        self,
+        feat_dim=80,
+        channels=512,
+        emb_dim=192,
+        global_context_att=False,
+        feat_type="wavlm_large",
+        sr=16000,
+        feature_selection="hidden_states",
+        update_extract=False,
+        config_path=None,
+    ):
+        super().__init__()
+
+        self.feat_type = feat_type
+        self.feature_selection = feature_selection
+        self.update_extract = update_extract
+        self.sr = sr
+
+        torch.hub._validate_not_a_forked_repo = lambda a, b, c: True
+        try:
+            local_s3prl_path = os.path.expanduser("~/.cache/torch/hub/s3prl_s3prl_main")
+            self.feature_extract = torch.hub.load(local_s3prl_path, feat_type, source="local", config_path=config_path)
+        except:  # noqa: E722
+            self.feature_extract = torch.hub.load("s3prl/s3prl", feat_type)
+
+        if len(self.feature_extract.model.encoder.layers) == 24 and hasattr(
+            self.feature_extract.model.encoder.layers[23].self_attn, "fp32_attention"
+        ):
+            self.feature_extract.model.encoder.layers[23].self_attn.fp32_attention = False
+        if len(self.feature_extract.model.encoder.layers) == 24 and hasattr(
+            self.feature_extract.model.encoder.layers[11].self_attn, "fp32_attention"
+        ):
+            self.feature_extract.model.encoder.layers[11].self_attn.fp32_attention = False
+
+        self.feat_num = self.get_feat_num()
+        self.feature_weight = nn.Parameter(torch.zeros(self.feat_num))
+
+        if feat_type != "fbank" and feat_type != "mfcc":
+            freeze_list = ["final_proj", "label_embs_concat", "mask_emb", "project_q", "quantizer"]
+            for name, param in self.feature_extract.named_parameters():
+                for freeze_val in freeze_list:
+                    if freeze_val in name:
+                        param.requires_grad = False
+                        break
+
+        if not self.update_extract:
+            for param in self.feature_extract.parameters():
+                param.requires_grad = False
+
+        self.instance_norm = nn.InstanceNorm1d(feat_dim)
+        # self.channels = [channels] * 4 + [channels * 3]
+        self.channels = [channels] * 4 + [1536]
+
+        self.layer1 = Conv1dReluBn(feat_dim, self.channels[0], kernel_size=5, padding=2)
+        self.layer2 = SE_Res2Block(
+            self.channels[0],
+            self.channels[1],
+            kernel_size=3,
+            stride=1,
+            padding=2,
+            dilation=2,
+            scale=8,
+            se_bottleneck_dim=128,
+        )
+        self.layer3 = SE_Res2Block(
+            self.channels[1],
+            self.channels[2],
+            kernel_size=3,
+            stride=1,
+            padding=3,
+            dilation=3,
+            scale=8,
+            se_bottleneck_dim=128,
+        )
+        self.layer4 = SE_Res2Block(
+            self.channels[2],
+            self.channels[3],
+            kernel_size=3,
+            stride=1,
+            padding=4,
+            dilation=4,
+            scale=8,
+            se_bottleneck_dim=128,
+        )
+
+        # self.conv = nn.Conv1d(self.channels[-1], self.channels[-1], kernel_size=1)
+        cat_channels = channels * 3
+        self.conv = nn.Conv1d(cat_channels, self.channels[-1], kernel_size=1)
+        self.pooling = AttentiveStatsPool(
+            self.channels[-1], attention_channels=128, global_context_att=global_context_att
+        )
+        self.bn = nn.BatchNorm1d(self.channels[-1] * 2)
+        self.linear = nn.Linear(self.channels[-1] * 2, emb_dim)
+
+    def get_feat_num(self):
+        self.feature_extract.eval()
+        wav = [torch.randn(self.sr).to(next(self.feature_extract.parameters()).device)]
+        with torch.no_grad():
+            features = self.feature_extract(wav)
+        select_feature = features[self.feature_selection]
+        if isinstance(select_feature, (list, tuple)):
+            return len(select_feature)
+        else:
+            return 1
+
+    def get_feat(self, x):
+        if self.update_extract:
+            x = self.feature_extract([sample for sample in x])
+        else:
+            with torch.no_grad():
+                if self.feat_type == "fbank" or self.feat_type == "mfcc":
+                    x = self.feature_extract(x) + 1e-6  # B x feat_dim x time_len
+                else:
+                    x = self.feature_extract([sample for sample in x])
+
+        if self.feat_type == "fbank":
+            x = x.log()
+
+        if self.feat_type != "fbank" and self.feat_type != "mfcc":
+            x = x[self.feature_selection]
+            if isinstance(x, (list, tuple)):
+                x = torch.stack(x, dim=0)
+            else:
+                x = x.unsqueeze(0)
+            norm_weights = F.softmax(self.feature_weight, dim=-1).unsqueeze(-1).unsqueeze(-1).unsqueeze(-1)
+            x = (norm_weights * x).sum(dim=0)
+            x = torch.transpose(x, 1, 2) + 1e-6
+
+        x = self.instance_norm(x)
+        return x
+
+    def forward(self, x):
+        x = self.get_feat(x)
+
+        out1 = self.layer1(x)
+        out2 = self.layer2(out1)
+        out3 = self.layer3(out2)
+        out4 = self.layer4(out3)
+
+        out = torch.cat([out2, out3, out4], dim=1)
+        out = F.relu(self.conv(out))
+        out = self.bn(self.pooling(out))
+        out = self.linear(out)
+
+        return out
+
+
+def ECAPA_TDNN_SMALL(
+    feat_dim,
+    emb_dim=256,
+    feat_type="wavlm_large",
+    sr=16000,
+    feature_selection="hidden_states",
+    update_extract=False,
+    config_path=None,
+):
+    return ECAPA_TDNN(
+        feat_dim=feat_dim,
+        channels=512,
+        emb_dim=emb_dim,
+        feat_type=feat_type,
+        sr=sr,
+        feature_selection=feature_selection,
+        update_extract=update_extract,
+        config_path=config_path,
+    )

f5_tts/eval/eval_infer_batch.py

@@ -0,0 +1,210 @@
+import os
+import sys
+
+
+sys.path.append(os.getcwd())
+
+import argparse
+import time
+from importlib.resources import files
+
+import torch
+import torchaudio
+from accelerate import Accelerator
+from hydra.utils import get_class
+from omegaconf import OmegaConf
+from tqdm import tqdm
+
+from f5_tts.eval.utils_eval import (
+    get_inference_prompt,
+    get_librispeech_test_clean_metainfo,
+    get_seedtts_testset_metainfo,
+)
+from f5_tts.infer.utils_infer import load_checkpoint, load_vocoder
+from f5_tts.model import CFM
+from f5_tts.model.utils import get_tokenizer
+
+
+accelerator = Accelerator()
+device = f"cuda:{accelerator.process_index}"
+
+
+use_ema = True
+target_rms = 0.1
+
+
+rel_path = str(files("f5_tts").joinpath("../../"))
+
+
+def main():
+    parser = argparse.ArgumentParser(description="batch inference")
+
+    parser.add_argument("-s", "--seed", default=None, type=int)
+    parser.add_argument("-n", "--expname", required=True)
+    parser.add_argument("-c", "--ckptstep", default=1250000, 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
+    exp_name = args.expname
+    ckpt_step = args.ckptstep
+
+    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.0
+    speed = 1.0
+    use_truth_duration = False
+    no_ref_audio = False
+
+    model_cfg = OmegaConf.load(str(files("f5_tts").joinpath(f"configs/{exp_name}.yaml")))
+    model_cls = get_class(f"f5_tts.model.{model_cfg.model.backbone}")
+    model_arc = model_cfg.model.arch
+
+    dataset_name = model_cfg.datasets.name
+    tokenizer = model_cfg.model.tokenizer
+
+    mel_spec_type = model_cfg.model.mel_spec.mel_spec_type
+    target_sample_rate = model_cfg.model.mel_spec.target_sample_rate
+    n_mel_channels = model_cfg.model.mel_spec.n_mel_channels
+    hop_length = model_cfg.model.mel_spec.hop_length
+    win_length = model_cfg.model.mel_spec.win_length
+    n_fft = model_cfg.model.mel_spec.n_fft
+
+    if testset == "ls_pc_test_clean":
+        metalst = rel_path + "/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 = rel_path + "/data/seedtts_testset/zh/meta.lst"
+        metainfo = get_seedtts_testset_metainfo(metalst)
+
+    elif testset == "seedtts_test_en":
+        metalst = rel_path + "/data/seedtts_testset/en/meta.lst"
+        metainfo = get_seedtts_testset_metainfo(metalst)
+
+    # path to save genereted wavs
+    output_dir = (
+        f"{rel_path}/"
+        f"results/{exp_name}_{ckpt_step}/{testset}/"
+        f"seed{seed}_{ode_method}_nfe{nfe_step}_{mel_spec_type}"
+        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 ''}"
+    )
+
+    # -------------------------------------------------#
+
+    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,
+        mel_spec_type=mel_spec_type,
+        target_rms=target_rms,
+        use_truth_duration=use_truth_duration,
+        infer_batch_size=infer_batch_size,
+    )
+
+    # Vocoder model
+    local = False
+    if mel_spec_type == "vocos":
+        vocoder_local_path = "../checkpoints/charactr/vocos-mel-24khz"
+    elif mel_spec_type == "bigvgan":
+        vocoder_local_path = "../checkpoints/bigvgan_v2_24khz_100band_256x"
+    vocoder = load_vocoder(vocoder_name=mel_spec_type, is_local=local, local_path=vocoder_local_path)
+
+    # Tokenizer
+    vocab_char_map, vocab_size = get_tokenizer(dataset_name, tokenizer)
+
+    # Model
+    model = CFM(
+        transformer=model_cls(**model_arc, text_num_embeds=vocab_size, mel_dim=n_mel_channels),
+        mel_spec_kwargs=dict(
+            n_fft=n_fft,
+            hop_length=hop_length,
+            win_length=win_length,
+            n_mel_channels=n_mel_channels,
+            target_sample_rate=target_sample_rate,
+            mel_spec_type=mel_spec_type,
+        ),
+        odeint_kwargs=dict(
+            method=ode_method,
+        ),
+        vocab_char_map=vocab_char_map,
+    ).to(device)
+
+    ckpt_prefix = rel_path + f"/ckpts/{exp_name}/model_{ckpt_step}"
+    if os.path.exists(ckpt_prefix + ".pt"):
+        ckpt_path = ckpt_prefix + ".pt"
+    elif os.path.exists(ckpt_prefix + ".safetensors"):
+        ckpt_path = ckpt_prefix + ".safetensors"
+    else:
+        print("Loading from self-organized training checkpoints rather than released pretrained.")
+        ckpt_path = rel_path + f"/{model_cfg.ckpts.save_dir}/model_{ckpt_step}.pt"
+
+    dtype = torch.float32 if mel_spec_type == "bigvgan" else None
+    model = load_checkpoint(model, ckpt_path, device, dtype=dtype, 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 = gen.permute(0, 2, 1).to(torch.float32)
+                    if mel_spec_type == "vocos":
+                        generated_wave = vocoder.decode(gen_mel_spec).cpu()
+                    elif mel_spec_type == "bigvgan":
+                        generated_wave = vocoder(gen_mel_spec).squeeze(0).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.")
+
+
+if __name__ == "__main__":
+    main()

f5_tts/eval/eval_infer_batch.sh

@@ -0,0 +1,18 @@
+#!/bin/bash
+
+# e.g. F5-TTS, 16 NFE
+accelerate launch src/f5_tts/eval/eval_infer_batch.py -s 0 -n "F5TTS_v1_Base" -t "seedtts_test_zh" -nfe 16
+accelerate launch src/f5_tts/eval/eval_infer_batch.py -s 0 -n "F5TTS_v1_Base" -t "seedtts_test_en" -nfe 16
+accelerate launch src/f5_tts/eval/eval_infer_batch.py -s 0 -n "F5TTS_v1_Base" -t "ls_pc_test_clean" -nfe 16
+
+# e.g. Vanilla E2 TTS, 32 NFE
+accelerate launch src/f5_tts/eval/eval_infer_batch.py -s 0 -n "E2TTS_Base" -c 1200000 -t "seedtts_test_zh" -o "midpoint" -ss 0
+accelerate launch src/f5_tts/eval/eval_infer_batch.py -s 0 -n "E2TTS_Base" -c 1200000 -t "seedtts_test_en" -o "midpoint" -ss 0
+accelerate launch src/f5_tts/eval/eval_infer_batch.py -s 0 -n "E2TTS_Base" -c 1200000 -t "ls_pc_test_clean" -o "midpoint" -ss 0
+
+# e.g. evaluate F5-TTS 16 NFE result on Seed-TTS test-zh
+python src/f5_tts/eval/eval_seedtts_testset.py -e wer -l zh --gen_wav_dir results/F5TTS_v1_Base_1250000/seedtts_test_zh/seed0_euler_nfe32_vocos_ss-1_cfg2.0_speed1.0 --gpu_nums 8
+python src/f5_tts/eval/eval_seedtts_testset.py -e sim -l zh --gen_wav_dir results/F5TTS_v1_Base_1250000/seedtts_test_zh/seed0_euler_nfe32_vocos_ss-1_cfg2.0_speed1.0 --gpu_nums 8
+python src/f5_tts/eval/eval_utmos.py --audio_dir results/F5TTS_v1_Base_1250000/seedtts_test_zh/seed0_euler_nfe32_vocos_ss-1_cfg2.0_speed1.0
+
+# etc.

f5_tts/eval/eval_librispeech_test_clean.py

@@ -0,0 +1,89 @@
+# Evaluate with Librispeech test-clean, ~3s prompt to generate 4-10s audio (the way of valle/voicebox evaluation)
+
+import argparse
+import json
+import os
+import sys
+
+
+sys.path.append(os.getcwd())
+
+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, run_sim
+
+
+rel_path = str(files("f5_tts").joinpath("../../"))
+
+
+def get_args():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("-e", "--eval_task", type=str, default="wer", choices=["sim", "wer"])
+    parser.add_argument("-l", "--lang", type=str, default="en")
+    parser.add_argument("-g", "--gen_wav_dir", type=str, required=True)
+    parser.add_argument("-p", "--librispeech_test_clean_path", type=str, required=True)
+    parser.add_argument("-n", "--gpu_nums", type=int, default=8, help="Number of GPUs to use")
+    parser.add_argument("--local", action="store_true", help="Use local custom checkpoint directory")
+    return parser.parse_args()
+
+
+def main():
+    args = get_args()
+    eval_task = args.eval_task
+    lang = args.lang
+    librispeech_test_clean_path = args.librispeech_test_clean_path  # test-clean path
+    gen_wav_dir = args.gen_wav_dir
+    metalst = rel_path + "/data/librispeech_pc_test_clean_cross_sentence.lst"
+
+    gpus = list(range(args.gpu_nums))
+    test_set = get_librispeech_test(metalst, gen_wav_dir, gpus, librispeech_test_clean_path)
+
+    ## In LibriSpeech, some speakers utilized varying voice characteristics for different characters in the book,
+    ## leading to a low similarity for the ground truth in some cases.
+    # test_set = get_librispeech_test(metalst, gen_wav_dir, gpus, librispeech_test_clean_path, eval_ground_truth = True)  # eval ground truth
+
+    local = args.local
+    if local:  # use local custom checkpoint dir
+        asr_ckpt_dir = "../checkpoints/Systran/faster-whisper-large-v3"
+    else:
+        asr_ckpt_dir = ""  # auto download to cache dir
+    wavlm_ckpt_dir = "../checkpoints/UniSpeech/wavlm_large_finetune.pth"
+
+    # --------------------------------------------------------------------------
+
+    full_results = []
+    metrics = []
+
+    if eval_task == "wer":
+        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 r in results:
+                full_results.extend(r)
+    elif eval_task == "sim":
+        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 r in results:
+                full_results.extend(r)
+    else:
+        raise ValueError(f"Unknown metric type: {eval_task}")
+
+    result_path = f"{gen_wav_dir}/_{eval_task}_results.jsonl"
+    with open(result_path, "w") as f:
+        for line in full_results:
+            metrics.append(line[eval_task])
+            f.write(json.dumps(line, ensure_ascii=False) + "\n")
+        metric = round(np.mean(metrics), 5)
+        f.write(f"\n{eval_task.upper()}: {metric}\n")
+
+    print(f"\nTotal {len(metrics)} samples")
+    print(f"{eval_task.upper()}: {metric}")
+    print(f"{eval_task.upper()} results saved to {result_path}")
+
+
+if __name__ == "__main__":
+    main()

f5_tts/eval/eval_seedtts_testset.py

@@ -0,0 +1,88 @@
+# Evaluate with Seed-TTS testset
+
+import argparse
+import json
+import os
+import sys
+
+
+sys.path.append(os.getcwd())
+
+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, run_sim
+
+
+rel_path = str(files("f5_tts").joinpath("../../"))
+
+
+def get_args():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("-e", "--eval_task", type=str, default="wer", choices=["sim", "wer"])
+    parser.add_argument("-l", "--lang", type=str, default="en", choices=["zh", "en"])
+    parser.add_argument("-g", "--gen_wav_dir", type=str, required=True)
+    parser.add_argument("-n", "--gpu_nums", type=int, default=8, help="Number of GPUs to use")
+    parser.add_argument("--local", action="store_true", help="Use local custom checkpoint directory")
+    return parser.parse_args()
+
+
+def main():
+    args = get_args()
+    eval_task = args.eval_task
+    lang = args.lang
+    gen_wav_dir = args.gen_wav_dir
+    metalst = rel_path + f"/data/seedtts_testset/{lang}/meta.lst"  # seed-tts testset
+
+    # NOTE. paraformer-zh result will be slightly different according to the number of gpus, cuz batchsize is different
+    #       zh 1.254 seems a result of 4 workers wer_seed_tts
+    gpus = list(range(args.gpu_nums))
+    test_set = get_seed_tts_test(metalst, gen_wav_dir, gpus)
+
+    local = args.local
+    if local:  # use local custom checkpoint dir
+        if lang == "zh":
+            asr_ckpt_dir = "../checkpoints/funasr"  # paraformer-zh dir under funasr
+        elif lang == "en":
+            asr_ckpt_dir = "../checkpoints/Systran/faster-whisper-large-v3"
+    else:
+        asr_ckpt_dir = ""  # auto download to cache dir
+    wavlm_ckpt_dir = "../checkpoints/UniSpeech/wavlm_large_finetune.pth"
+
+    # --------------------------------------------------------------------------
+
+    full_results = []
+    metrics = []
+
+    if eval_task == "wer":
+        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 r in results:
+                full_results.extend(r)
+    elif eval_task == "sim":
+        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 r in results:
+                full_results.extend(r)
+    else:
+        raise ValueError(f"Unknown metric type: {eval_task}")
+
+    result_path = f"{gen_wav_dir}/_{eval_task}_results.jsonl"
+    with open(result_path, "w") as f:
+        for line in full_results:
+            metrics.append(line[eval_task])
+            f.write(json.dumps(line, ensure_ascii=False) + "\n")
+        metric = round(np.mean(metrics), 5)
+        f.write(f"\n{eval_task.upper()}: {metric}\n")
+
+    print(f"\nTotal {len(metrics)} samples")
+    print(f"{eval_task.upper()}: {metric}")
+    print(f"{eval_task.upper()} results saved to {result_path}")
+
+
+if __name__ == "__main__":
+    main()

f5_tts/eval/eval_utmos.py

@@ -0,0 +1,42 @@
+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 "xpu" if torch.xpu.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_score = 0
+
+    utmos_result_path = Path(args.audio_dir) / "_utmos_results.jsonl"
+    with open(utmos_result_path, "w", encoding="utf-8") as f:
+        for audio_path in tqdm(audio_paths, desc="Processing"):
+            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)
+            line = {}
+            line["wav"], line["utmos"] = str(audio_path.stem), score.item()
+            utmos_score += score.item()
+            f.write(json.dumps(line, ensure_ascii=False) + "\n")
+        avg_score = utmos_score / len(audio_paths) if len(audio_paths) > 0 else 0
+        f.write(f"\nUTMOS: {avg_score:.4f}\n")
+
+    print(f"UTMOS: {avg_score:.4f}")
+    print(f"UTMOS results saved to {utmos_result_path}")
+
+
+if __name__ == "__main__":
+    main()

f5_tts/eval/utils_eval.py

@@ -0,0 +1,419 @@
+import math
+import os
+import random
+import string
+from pathlib import Path
+
+import torch
+import torch.nn.functional as F
+import torchaudio
+from tqdm import tqdm
+
+from f5_tts.eval.ecapa_tdnn import ECAPA_TDNN_SMALL
+from f5_tts.model.modules import MelSpec
+from f5_tts.model.utils import convert_char_to_pinyin
+
+
+# seedtts testset metainfo: utt, prompt_text, prompt_wav, gt_text, gt_wav
+def get_seedtts_testset_metainfo(metalst):
+    f = open(metalst)
+    lines = f.readlines()
+    f.close()
+    metainfo = []
+    for line in lines:
+        if len(line.strip().split("|")) == 5:
+            utt, prompt_text, prompt_wav, gt_text, gt_wav = line.strip().split("|")
+        elif len(line.strip().split("|")) == 4:
+            utt, prompt_text, prompt_wav, gt_text = line.strip().split("|")
+            gt_wav = os.path.join(os.path.dirname(metalst), "wavs", utt + ".wav")
+        if not os.path.isabs(prompt_wav):
+            prompt_wav = os.path.join(os.path.dirname(metalst), prompt_wav)
+        metainfo.append((utt, prompt_text, prompt_wav, gt_text, gt_wav))
+    return metainfo
+
+
+# librispeech test-clean metainfo: gen_utt, ref_txt, ref_wav, gen_txt, gen_wav
+def get_librispeech_test_clean_metainfo(metalst, librispeech_test_clean_path):
+    f = open(metalst)
+    lines = f.readlines()
+    f.close()
+    metainfo = []
+    for line in lines:
+        ref_utt, ref_dur, ref_txt, gen_utt, gen_dur, gen_txt = line.strip().split("\t")
+
+        # ref_txt = ref_txt[0] + ref_txt[1:].lower() + '.'  # if use librispeech test-clean (no-pc)
+        ref_spk_id, ref_chaptr_id, _ = ref_utt.split("-")
+        ref_wav = os.path.join(librispeech_test_clean_path, ref_spk_id, ref_chaptr_id, ref_utt + ".flac")
+
+        # gen_txt = gen_txt[0] + gen_txt[1:].lower() + '.'  # if use librispeech test-clean (no-pc)
+        gen_spk_id, gen_chaptr_id, _ = gen_utt.split("-")
+        gen_wav = os.path.join(librispeech_test_clean_path, gen_spk_id, gen_chaptr_id, gen_utt + ".flac")
+
+        metainfo.append((gen_utt, ref_txt, ref_wav, " " + gen_txt, gen_wav))
+
+    return metainfo
+
+
+# padded to max length mel batch
+def padded_mel_batch(ref_mels):
+    max_mel_length = torch.LongTensor([mel.shape[-1] for mel in ref_mels]).amax()
+    padded_ref_mels = []
+    for mel in ref_mels:
+        padded_ref_mel = F.pad(mel, (0, max_mel_length - mel.shape[-1]), value=0)
+        padded_ref_mels.append(padded_ref_mel)
+    padded_ref_mels = torch.stack(padded_ref_mels)
+    padded_ref_mels = padded_ref_mels.permute(0, 2, 1)
+    return padded_ref_mels
+
+
+# get prompts from metainfo containing: utt, prompt_text, prompt_wav, gt_text, gt_wav
+
+
+def get_inference_prompt(
+    metainfo,
+    speed=1.0,
+    tokenizer="pinyin",
+    polyphone=True,
+    target_sample_rate=24000,
+    n_fft=1024,
+    win_length=1024,
+    n_mel_channels=100,
+    hop_length=256,
+    mel_spec_type="vocos",
+    target_rms=0.1,
+    use_truth_duration=False,
+    infer_batch_size=1,
+    num_buckets=200,
+    min_secs=3,
+    max_secs=40,
+):
+    prompts_all = []
+
+    min_tokens = min_secs * target_sample_rate // hop_length
+    max_tokens = max_secs * target_sample_rate // hop_length
+
+    batch_accum = [0] * num_buckets
+    utts, ref_rms_list, ref_mels, ref_mel_lens, total_mel_lens, final_text_list = (
+        [[] for _ in range(num_buckets)] for _ in range(6)
+    )
+
+    mel_spectrogram = MelSpec(
+        n_fft=n_fft,
+        hop_length=hop_length,
+        win_length=win_length,
+        n_mel_channels=n_mel_channels,
+        target_sample_rate=target_sample_rate,
+        mel_spec_type=mel_spec_type,
+    )
+
+    for utt, prompt_text, prompt_wav, gt_text, gt_wav in tqdm(metainfo, desc="Processing prompts..."):
+        # Audio
+        ref_audio, ref_sr = torchaudio.load(prompt_wav)
+        ref_rms = torch.sqrt(torch.mean(torch.square(ref_audio)))
+        if ref_rms < target_rms:
+            ref_audio = ref_audio * target_rms / ref_rms
+        assert ref_audio.shape[-1] > 5000, f"Empty prompt wav: {prompt_wav}, or torchaudio backend issue."
+        if ref_sr != target_sample_rate:
+            resampler = torchaudio.transforms.Resample(ref_sr, target_sample_rate)
+            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)
+        else:
+            text_list = text
+
+        # to mel spectrogram
+        ref_mel = mel_spectrogram(ref_audio)
+        ref_mel = ref_mel.squeeze(0)
+
+        # Duration, mel frame length
+        ref_mel_len = ref_mel.shape[-1]
+
+        if use_truth_duration:
+            gt_audio, gt_sr = torchaudio.load(gt_wav)
+            if gt_sr != target_sample_rate:
+                resampler = torchaudio.transforms.Resample(gt_sr, target_sample_rate)
+                gt_audio = resampler(gt_audio)
+            total_mel_len = ref_mel_len + int(gt_audio.shape[-1] / hop_length / speed)
+
+            # # test vocoder resynthesis
+            # ref_audio = gt_audio
+        else:
+            ref_text_len = len(prompt_text.encode("utf-8"))
+            gen_text_len = len(gt_text.encode("utf-8"))
+            total_mel_len = ref_mel_len + int(ref_mel_len / ref_text_len * gen_text_len / speed)
+
+        # deal with batch
+        assert infer_batch_size > 0, "infer_batch_size should be greater than 0."
+        assert min_tokens <= total_mel_len <= max_tokens, (
+            f"Audio {utt} has duration {total_mel_len * hop_length // target_sample_rate}s out of range [{min_secs}, {max_secs}]."
+        )
+        bucket_i = math.floor((total_mel_len - min_tokens) / (max_tokens - min_tokens + 1) * num_buckets)
+
+        utts[bucket_i].append(utt)
+        ref_rms_list[bucket_i].append(ref_rms)
+        ref_mels[bucket_i].append(ref_mel)
+        ref_mel_lens[bucket_i].append(ref_mel_len)
+        total_mel_lens[bucket_i].append(total_mel_len)
+        final_text_list[bucket_i].extend(text_list)
+
+        batch_accum[bucket_i] += total_mel_len
+
+        if batch_accum[bucket_i] >= infer_batch_size:
+            # print(f"\n{len(ref_mels[bucket_i][0][0])}\n{ref_mel_lens[bucket_i]}\n{total_mel_lens[bucket_i]}")
+            prompts_all.append(
+                (
+                    utts[bucket_i],
+                    ref_rms_list[bucket_i],
+                    padded_mel_batch(ref_mels[bucket_i]),
+                    ref_mel_lens[bucket_i],
+                    total_mel_lens[bucket_i],
+                    final_text_list[bucket_i],
+                )
+            )
+            batch_accum[bucket_i] = 0
+            (
+                utts[bucket_i],
+                ref_rms_list[bucket_i],
+                ref_mels[bucket_i],
+                ref_mel_lens[bucket_i],
+                total_mel_lens[bucket_i],
+                final_text_list[bucket_i],
+            ) = [], [], [], [], [], []
+
+    # add residual
+    for bucket_i, bucket_frames in enumerate(batch_accum):
+        if bucket_frames > 0:
+            prompts_all.append(
+                (
+                    utts[bucket_i],
+                    ref_rms_list[bucket_i],
+                    padded_mel_batch(ref_mels[bucket_i]),
+                    ref_mel_lens[bucket_i],
+                    total_mel_lens[bucket_i],
+                    final_text_list[bucket_i],
+                )
+            )
+    # not only leave easy work for last workers
+    random.seed(666)
+    random.shuffle(prompts_all)
+
+    return prompts_all
+
+
+# get wav_res_ref_text of seed-tts test metalst
+# https://github.com/BytedanceSpeech/seed-tts-eval
+
+
+def get_seed_tts_test(metalst, gen_wav_dir, gpus):
+    f = open(metalst)
+    lines = f.readlines()
+    f.close()
+
+    test_set_ = []
+    for line in tqdm(lines):
+        if len(line.strip().split("|")) == 5:
+            utt, prompt_text, prompt_wav, gt_text, gt_wav = line.strip().split("|")
+        elif len(line.strip().split("|")) == 4:
+            utt, prompt_text, prompt_wav, gt_text = line.strip().split("|")
+
+        if not os.path.exists(os.path.join(gen_wav_dir, utt + ".wav")):
+            continue
+        gen_wav = os.path.join(gen_wav_dir, utt + ".wav")
+        if not os.path.isabs(prompt_wav):
+            prompt_wav = os.path.join(os.path.dirname(metalst), prompt_wav)
+
+        test_set_.append((gen_wav, prompt_wav, gt_text))
+
+    num_jobs = len(gpus)
+    if num_jobs == 1:
+        return [(gpus[0], test_set_)]
+
+    wav_per_job = len(test_set_) // num_jobs + 1
+    test_set = []
+    for i in range(num_jobs):
+        test_set.append((gpus[i], test_set_[i * wav_per_job : (i + 1) * wav_per_job]))
+
+    return test_set
+
+
+# get librispeech test-clean cross sentence test
+
+
+def get_librispeech_test(metalst, gen_wav_dir, gpus, librispeech_test_clean_path, eval_ground_truth=False):
+    f = open(metalst)
+    lines = f.readlines()
+    f.close()
+
+    test_set_ = []
+    for line in tqdm(lines):
+        ref_utt, ref_dur, ref_txt, gen_utt, gen_dur, gen_txt = line.strip().split("\t")
+
+        if eval_ground_truth:
+            gen_spk_id, gen_chaptr_id, _ = gen_utt.split("-")
+            gen_wav = os.path.join(librispeech_test_clean_path, gen_spk_id, gen_chaptr_id, gen_utt + ".flac")
+        else:
+            if not os.path.exists(os.path.join(gen_wav_dir, gen_utt + ".wav")):
+                raise FileNotFoundError(f"Generated wav not found: {gen_utt}")
+            gen_wav = os.path.join(gen_wav_dir, gen_utt + ".wav")
+
+        ref_spk_id, ref_chaptr_id, _ = ref_utt.split("-")
+        ref_wav = os.path.join(librispeech_test_clean_path, ref_spk_id, ref_chaptr_id, ref_utt + ".flac")
+
+        test_set_.append((gen_wav, ref_wav, gen_txt))
+
+    num_jobs = len(gpus)
+    if num_jobs == 1:
+        return [(gpus[0], test_set_)]
+
+    wav_per_job = len(test_set_) // num_jobs + 1
+    test_set = []
+    for i in range(num_jobs):
+        test_set.append((gpus[i], test_set_[i * wav_per_job : (i + 1) * wav_per_job]))
+
+    return test_set
+
+
+# load asr model
+
+
+def load_asr_model(lang, ckpt_dir=""):
+    if lang == "zh":
+        from funasr import AutoModel
+
+        model = AutoModel(
+            model=os.path.join(ckpt_dir, "paraformer-zh"),
+            # vad_model = os.path.join(ckpt_dir, "fsmn-vad"),
+            # punc_model = os.path.join(ckpt_dir, "ct-punc"),
+            # spk_model = os.path.join(ckpt_dir, "cam++"),
+            disable_update=True,
+        )  # following seed-tts setting
+    elif lang == "en":
+        from faster_whisper import WhisperModel
+
+        model_size = "large-v3" if ckpt_dir == "" else ckpt_dir
+        model = WhisperModel(model_size, device="cuda", compute_type="float16")
+    return model
+
+
+# WER Evaluation, the way Seed-TTS does
+
+
+def run_asr_wer(args):
+    rank, lang, test_set, ckpt_dir = args
+
+    if lang == "zh":
+        import zhconv
+
+        torch.cuda.set_device(rank)
+    elif lang == "en":
+        os.environ["CUDA_VISIBLE_DEVICES"] = str(rank)
+    else:
+        raise NotImplementedError(
+            "lang support only 'zh' (funasr paraformer-zh), 'en' (faster-whisper-large-v3), for now."
+        )
+
+    asr_model = load_asr_model(lang, ckpt_dir=ckpt_dir)
+
+    from zhon.hanzi import punctuation
+
+    punctuation_all = punctuation + string.punctuation
+    wer_results = []
+
+    from jiwer import compute_measures
+
+    for gen_wav, prompt_wav, truth in tqdm(test_set):
+        if lang == "zh":
+            res = asr_model.generate(input=gen_wav, batch_size_s=300, disable_pbar=True)
+            hypo = res[0]["text"]
+            hypo = zhconv.convert(hypo, "zh-cn")
+        elif lang == "en":
+            segments, _ = asr_model.transcribe(gen_wav, beam_size=5, language="en")
+            hypo = ""
+            for segment in segments:
+                hypo = hypo + " " + segment.text
+
+        raw_truth = truth
+        raw_hypo = hypo
+
+        for x in punctuation_all:
+            truth = truth.replace(x, "")
+            hypo = hypo.replace(x, "")
+
+        truth = truth.replace("  ", " ")
+        hypo = hypo.replace("  ", " ")
+
+        if lang == "zh":
+            truth = " ".join([x for x in truth])
+            hypo = " ".join([x for x in hypo])
+        elif lang == "en":
+            truth = truth.lower()
+            hypo = hypo.lower()
+
+        measures = compute_measures(truth, hypo)
+        wer = measures["wer"]
+
+        # ref_list = truth.split(" ")
+        # subs = measures["substitutions"] / len(ref_list)
+        # dele = measures["deletions"] / len(ref_list)
+        # inse = measures["insertions"] / len(ref_list)
+
+        wer_results.append(
+            {
+                "wav": Path(gen_wav).stem,
+                "truth": raw_truth,
+                "hypo": raw_hypo,
+                "wer": wer,
+            }
+        )
+
+    return wer_results
+
+
+# SIM Evaluation
+
+
+def run_sim(args):
+    rank, test_set, ckpt_dir = 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, 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
+    if use_gpu:
+        model = model.cuda(device)
+    model.eval()
+
+    sim_results = []
+    for gen_wav, prompt_wav, truth in tqdm(test_set):
+        wav1, sr1 = torchaudio.load(gen_wav)
+        wav2, sr2 = torchaudio.load(prompt_wav)
+
+        resample1 = torchaudio.transforms.Resample(orig_freq=sr1, new_freq=16000)
+        resample2 = torchaudio.transforms.Resample(orig_freq=sr2, new_freq=16000)
+        wav1 = resample1(wav1)
+        wav2 = resample2(wav2)
+
+        if use_gpu:
+            wav1 = wav1.cuda(device)
+            wav2 = wav2.cuda(device)
+        with torch.no_grad():
+            emb1 = model(wav1)
+            emb2 = model(wav2)
+
+        sim = F.cosine_similarity(emb1, emb2)[0].item()
+        # print(f"VSim score between two audios: {sim:.4f} (-1.0, 1.0).")
+        sim_results.append(
+            {
+                "wav": Path(gen_wav).stem,
+                "sim": sim,
+            }
+        )
+
+    return sim_results

f5_tts/infer/README.md

@@ -0,0 +1,177 @@
+# 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 will be automatically downloaded when running inference scripts.
+
+**More checkpoints with whole community efforts can be found in [SHARED.md](SHARED.md), supporting more languages.**
+
+Currently support **30s for a single** generation, which is the **total length** (same logic if `fix_duration`) including both prompt and output audio. However, `infer_cli` and `infer_gradio` will automatically do chunk generation for longer text. Long reference audio will be **clip short to ~12s**.
+
+To avoid possible inference failures, make sure you have seen through the following instructions.
+
+- Use reference audio <12s and leave proper silence space (e.g. 1s) at the end. Otherwise there is a risk of truncating in the middle of word, leading to suboptimal generation.
+- <ins>Uppercased letters</ins> (best with form like K.F.C.) will be uttered letter by letter, and lowercased letters used for common words. 
+- Add some spaces (blank: " ") or punctuations (e.g. "," ".") <ins>to explicitly introduce some pauses</ins>.
+- If English punctuation marks the end of a sentence, make sure there is a space " " after it. Otherwise not regarded as when chunk.
+- <ins>Preprocess numbers</ins> to Chinese letters if you want to have them read in Chinese, otherwise in English.
+- If the generation output is blank (pure silence), <ins>check for FFmpeg installation</ins>.
+- Try <ins>turn off `use_ema` if using an early-stage</ins> finetuned checkpoint (which goes just few updates).
+
+
+## Gradio App
+
+Currently supported features:
+
+- Basic TTS with Chunk Inference
+- Multi-Style / Multi-Speaker Generation
+- Voice Chat powered by Qwen2.5-3B-Instruct
+- [Custom inference with more language support](SHARED.md)
+
+The cli command `f5-tts_infer-gradio` equals to `python src/f5_tts/infer/infer_gradio.py`, which launches a Gradio APP (web interface) for inference.
+
+The script will load model checkpoints from Huggingface. You can also manually download files and update the path to `load_model()` in `infer_gradio.py`. Currently only load TTS models first, will load ASR model to do transcription if `ref_text` not provided, will load LLM model if use Voice Chat.
+
+More flags options:
+
+```bash
+# Automatically launch the interface in the default web browser
+f5-tts_infer-gradio --inbrowser
+
+# Set the root path of the application, if it's not served from the root ("/") of the domain
+# For example, if the application is served at "https://example.com/myapp"
+f5-tts_infer-gradio --root_path "/myapp"
+```
+
+Could also be used as a component for larger application:
+```python
+import gradio as gr
+from f5_tts.infer.infer_gradio import app
+
+with gr.Blocks() as main_app:
+    gr.Markdown("# This is an example of using F5-TTS within a bigger Gradio app")
+
+    # ... other Gradio components
+
+    app.render()
+
+main_app.launch()
+```
+
+
+## CLI Inference
+
+The cli command `f5-tts_infer-cli` equals to `python src/f5_tts/infer/infer_cli.py`, which is a command line tool for inference.
+
+The script will load model checkpoints from Huggingface. You can also manually download files and use `--ckpt_file` to specify the model you want to load, or directly update in `infer_cli.py`.
+
+For change vocab.txt use `--vocab_file` to provide your `vocab.txt` file.
+
+Basically you can inference with flags:
+```bash
+# Leave --ref_text "" will have ASR model transcribe (extra GPU memory usage)
+f5-tts_infer-cli \
+--model F5TTS_v1_Base \
+--ref_audio "ref_audio.wav" \
+--ref_text "The content, subtitle or transcription of reference audio." \
+--gen_text "Some text you want TTS model generate for you."
+
+# Use BigVGAN as vocoder. Currently only support F5TTS_Base. 
+f5-tts_infer-cli --model F5TTS_Base --vocoder_name bigvgan --load_vocoder_from_local
+
+# Use custom path checkpoint, e.g.
+f5-tts_infer-cli --ckpt_file ckpts/F5TTS_v1_Base/model_1250000.safetensors
+
+# More instructions
+f5-tts_infer-cli --help
+```
+
+And a `.toml` file would help with more flexible usage.
+
+```bash
+f5-tts_infer-cli -c custom.toml
+```
+
+For example, you can use `.toml` to pass in variables, refer to `src/f5_tts/infer/examples/basic/basic.toml`:
+
+```toml
+# F5TTS_v1_Base | E2TTS_Base
+model = "F5TTS_v1_Base"
+ref_audio = "infer/examples/basic/basic_ref_en.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."
+# File with text to generate. Ignores the text above.
+gen_file = ""
+remove_silence = false
+output_dir = "tests"
+```
+
+You can also leverage `.toml` file to do multi-style generation, refer to `src/f5_tts/infer/examples/multi/story.toml`.
+
+```toml
+# F5TTS_v1_Base | E2TTS_Base
+model = "F5TTS_v1_Base"
+ref_audio = "infer/examples/multi/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 = "infer/examples/multi/story.txt"
+remove_silence = true
+output_dir = "tests"
+
+[voices.town]
+ref_audio = "infer/examples/multi/town.flac"
+ref_text = ""
+
+[voices.country]
+ref_audio = "infer/examples/multi/country.flac"
+ref_text = ""
+```
+You should mark the voice with `[main]` `[town]` `[country]` whenever you want to change voice, refer to `src/f5_tts/infer/examples/multi/story.txt`.
+
+## API Usage
+
+```python
+from importlib.resources import files
+from f5_tts.api import F5TTS
+
+f5tts = F5TTS()
+wav, sr, spec = f5tts.infer(
+    ref_file=str(files("f5_tts").joinpath("infer/examples/basic/basic_ref_en.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.""",
+    file_wave=str(files("f5_tts").joinpath("../../tests/api_out.wav")),
+    file_spec=str(files("f5_tts").joinpath("../../tests/api_out.png")),
+    seed=None,
+)
+```
+Check [api.py](../api.py) for more details.
+
+## TensorRT-LLM Deployment
+
+See [detailed instructions](../runtime/triton_trtllm/README.md) for more information.
+
+## Socket Real-time Service
+
+Real-time voice output with chunk stream:
+
+```bash
+# Start socket server
+python src/f5_tts/socket_server.py
+
+# If PyAudio not installed
+sudo apt-get install portaudio19-dev
+pip install pyaudio
+
+# Communicate with socket client
+python src/f5_tts/socket_client.py
+```
+
+## Speech Editing
+
+To test speech editing capabilities, use the following command:
+
+```bash
+python src/f5_tts/infer/speech_edit.py
+```
+

f5_tts/infer/SHARED.md

@@ -0,0 +1,193 @@
+<!-- omit in toc -->
+# Shared Model Cards
+
+<!-- omit in toc -->
+### **Prerequisites of using**
+- This document is serving as a quick lookup table for the community training/finetuning result, with various language support.
+- The models in this repository are open source and are based on voluntary contributions from contributors.
+- The use of models must be conditioned on respect for the respective creators. The convenience brought comes from their efforts.
+
+<!-- omit in toc -->
+### **Welcome to share here**
+- Have a pretrained/finetuned result: model checkpoint (pruned best to facilitate inference, i.e. leave only `ema_model_state_dict`) and corresponding vocab file (for tokenization).
+- Host a public [huggingface model repository](https://huggingface.co/new) and upload the model related files.
+- Make a pull request adding a model card to the current page, i.e. `src\f5_tts\infer\SHARED.md`.
+
+<!-- omit in toc -->
+### Supported Languages
+- [Multilingual](#multilingual)
+    - [F5-TTS v1 v0 Base @ zh \& en @ F5-TTS](#f5-tts-v1-v0-base--zh--en--f5-tts)
+- [English](#english)
+- [Finnish](#finnish)
+    - [F5-TTS Base @ fi @ AsmoKoskinen](#f5-tts-base--fi--asmokoskinen)
+- [French](#french)
+    - [F5-TTS Base @ fr @ RASPIAUDIO](#f5-tts-base--fr--raspiaudio)
+- [German](#german)
+    - [F5-TTS Base @ de @ hvoss-techfak](#f5-tts-base--de--hvoss-techfak)
+- [Hindi](#hindi)
+    - [F5-TTS Small @ hi @ SPRINGLab](#f5-tts-small--hi--springlab)
+- [Italian](#italian)
+    - [F5-TTS Base @ it @ alien79](#f5-tts-base--it--alien79)
+- [Japanese](#japanese)
+    - [F5-TTS Base @ ja @ Jmica](#f5-tts-base--ja--jmica)
+- [Mandarin](#mandarin)
+- [Russian](#russian)
+    - [F5-TTS Base @ ru @ HotDro4illa](#f5-tts-base--ru--hotdro4illa)
+- [Spanish](#spanish)
+    - [F5-TTS Base @ es @ jpgallegoar](#f5-tts-base--es--jpgallegoar)
+
+
+## Multilingual
+
+#### F5-TTS v1 v0 Base @ zh & en @ F5-TTS
+|Model|🤗Hugging Face|Data (Hours)|Model License|
+|:---:|:------------:|:-----------:|:-------------:|
+|F5-TTS v1 Base|[ckpt & vocab](https://huggingface.co/SWivid/F5-TTS/tree/main/F5TTS_v1_Base)|[Emilia 95K zh&en](https://huggingface.co/datasets/amphion/Emilia-Dataset/tree/fc71e07)|cc-by-nc-4.0|
+
+```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}
+```
+
+|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: 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, "text_mask_padding": False, "conv_layers": 4, "pe_attn_head": 1}
+```
+
+*Other infos, e.g. Author info, Github repo, Link to some sampled results, Usage instruction, Tutorial (Blog, Video, etc.) ...*
+
+
+## English
+
+
+## Finnish
+
+#### F5-TTS Base @ fi @ AsmoKoskinen
+|Model|🤗Hugging Face|Data|Model License|
+|:---:|:------------:|:-----------:|:-------------:|
+|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: 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, "text_mask_padding": False, "conv_layers": 4, "pe_attn_head": 1}
+```
+
+
+## French
+
+#### F5-TTS Base @ fr @ RASPIAUDIO
+|Model|🤗Hugging Face|Data (Hours)|Model License|
+|:---:|:------------:|:-----------:|:-------------:|
+|F5-TTS Base|[ckpt & vocab](https://huggingface.co/RASPIAUDIO/F5-French-MixedSpeakers-reduced)|[LibriVox](https://librivox.org/)|cc-by-nc-4.0|
+
+```bash
+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, "text_mask_padding": False, "conv_layers": 4, "pe_attn_head": 1}
+```
+
+- [Online Inference with Hugging Face Space](https://huggingface.co/spaces/RASPIAUDIO/f5-tts_french).
+- [Tutorial video to train a new language model](https://www.youtube.com/watch?v=UO4usaOojys).
+- [Discussion about this training can be found here](https://github.com/SWivid/F5-TTS/issues/434).
+
+
+## German
+
+#### F5-TTS Base @ de @ hvoss-techfak
+|Model|🤗Hugging Face|Data (Hours)|Model License|
+|:---:|:------------:|:-----------:|:-------------:|
+|F5-TTS Base|[ckpt & vocab](https://huggingface.co/hvoss-techfak/F5-TTS-German)|[Mozilla Common Voice 19.0](https://commonvoice.mozilla.org/en/datasets) & 800 hours Crowdsourced |cc-by-nc-4.0|
+
+```bash
+Model: hf://hvoss-techfak/F5-TTS-German/model_f5tts_german.pt
+Vocab: hf://hvoss-techfak/F5-TTS-German/vocab.txt
+Config: {"dim": 1024, "depth": 22, "heads": 16, "ff_mult": 2, "text_dim": 512, "text_mask_padding": False, "conv_layers": 4, "pe_attn_head": 1}
+```
+
+- Finetuned by [@hvoss-techfak](https://github.com/hvoss-techfak)
+
+
+## Hindi
+
+#### F5-TTS Small @ hi @ SPRINGLab
+|Model|🤗Hugging Face|Data (Hours)|Model License|
+|:---:|:------------:|:-----------:|:-------------:|
+|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: 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, "text_mask_padding": False, "conv_layers": 4, "pe_attn_head": 1}
+```
+
+- 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, "text_mask_padding": False, "conv_layers": 4, "pe_attn_head": 1}
+```
+
+- Trained by [Mithril Man](https://github.com/MithrilMan)
+- Model details on [hf project home](https://huggingface.co/alien79/F5-TTS-italian)
+- Open to collaborations to further improve the model
+
+
+## Japanese
+
+#### F5-TTS Base @ ja @ Jmica
+|Model|🤗Hugging Face|Data (Hours)|Model License|
+|:---:|:------------:|:-----------:|:-------------:|
+|F5-TTS Base|[ckpt & vocab](https://huggingface.co/Jmica/F5TTS/tree/main/JA_21999120)|[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: hf://Jmica/F5TTS/JA_21999120/model_21999120.pt
+Vocab: hf://Jmica/F5TTS/JA_21999120/vocab_japanese.txt
+Config: {"dim": 1024, "depth": 22, "heads": 16, "ff_mult": 2, "text_dim": 512, "text_mask_padding": False, "conv_layers": 4, "pe_attn_head": 1}
+```
+
+
+## Mandarin
+
+
+## Russian
+
+#### F5-TTS Base @ ru @ HotDro4illa
+|Model|🤗Hugging Face|Data (Hours)|Model License|
+|:---:|:------------:|:-----------:|:-------------:|
+|F5-TTS Base|[ckpt & vocab](https://huggingface.co/hotstone228/F5-TTS-Russian)|[Common voice](https://huggingface.co/datasets/mozilla-foundation/common_voice_17_0)|cc-by-nc-4.0|
+
+```bash
+Model: hf://hotstone228/F5-TTS-Russian/model_last.safetensors
+Vocab: hf://hotstone228/F5-TTS-Russian/vocab.txt
+Config: {"dim": 1024, "depth": 22, "heads": 16, "ff_mult": 2, "text_dim": 512, "text_mask_padding": False, "conv_layers": 4, "pe_attn_head": 1}
+```
+- Finetuned by [HotDro4illa](https://github.com/HotDro4illa)
+- Any improvements are welcome
+
+
+## Spanish
+
+#### F5-TTS Base @ es @ jpgallegoar
+|Model|🤗Hugging Face|Data (Hours)|Model License|
+|:---:|:------------:|:-----------:|:-------------:|
+|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.

f5_tts/infer/examples/basic/basic.toml

@@ -0,0 +1,11 @@
+# F5TTS_v1_Base | E2TTS_Base
+model = "F5TTS_v1_Base"
+ref_audio = "infer/examples/basic/basic_ref_en.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."
+# File with text to generate. Ignores the text above.
+gen_file = ""
+remove_silence = false
+output_dir = "tests"
+output_file = "infer_cli_basic.wav"

f5_tts/infer/examples/multi/story.toml

@@ -0,0 +1,20 @@
+# F5TTS_v1_Base | E2TTS_Base
+model = "F5TTS_v1_Base"
+ref_audio = "infer/examples/multi/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 = "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"
+ref_text = ""
+speed = 0.8  # will ignore global speed
+
+[voices.country]
+ref_audio = "infer/examples/multi/country.flac"
+ref_text = ""

f5_tts/infer/examples/multi/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."

f5_tts/infer/examples/vocab.txt

@@ -0,0 +1,2545 @@
+ 
+!
+"
+#
+$
+%
+&
+'
+(
+)
+*
++
+,
+-
+.
+/
+0
+1
+2
+3
+4
+5
+6
+7
+8
+9
+:
+;
+=
+>
+?
+@
+A
+B
+C
+D
+E
+F
+G
+H
+I
+J
+K
+L
+M
+N
+O
+P
+Q
+R
+S
+T
+U
+V
+W
+X
+Y
+Z
+[
+\
+]
+_
+a
+a1
+ai1
+ai2
+ai3
+ai4
+an1
+an3
+an4
+ang1
+ang2
+ang4
+ao1
+ao2
+ao3
+ao4
+b
+ba
+ba1
+ba2
+ba3
+ba4
+bai1
+bai2
+bai3
+bai4
+ban1
+ban2
+ban3
+ban4
+bang1
+bang2
+bang3
+bang4
+bao1
+bao2
+bao3
+bao4
+bei
+bei1
+bei2
+bei3
+bei4
+ben1
+ben2
+ben3
+ben4
+beng
+beng1
+beng2
+beng3
+beng4
+bi1
+bi2
+bi3
+bi4
+bian1
+bian2
+bian3
+bian4
+biao1
+biao2
+biao3
+bie1
+bie2
+bie3
+bie4
+bin1
+bin4
+bing1
+bing2
+bing3
+bing4
+bo
+bo1
+bo2
+bo3
+bo4
+bu2
+bu3
+bu4
+c
+ca1
+cai1
+cai2
+cai3
+cai4
+can1
+can2
+can3
+can4
+cang1
+cang2
+cao1
+cao2
+cao3
+ce4
+cen1
+cen2
+ceng1
+ceng2
+ceng4
+cha1
+cha2
+cha3
+cha4
+chai1
+chai2
+chan1
+chan2
+chan3
+chan4
+chang1
+chang2
+chang3
+chang4
+chao1
+chao2
+chao3
+che1
+che2
+che3
+che4
+chen1
+chen2
+chen3
+chen4
+cheng1
+cheng2
+cheng3
+cheng4
+chi1
+chi2
+chi3
+chi4
+chong1
+chong2
+chong3
+chong4
+chou1
+chou2
+chou3
+chou4
+chu1
+chu2
+chu3
+chu4
+chua1
+chuai1
+chuai2
+chuai3
+chuai4
+chuan1
+chuan2
+chuan3
+chuan4
+chuang1
+chuang2
+chuang3
+chuang4
+chui1
+chui2
+chun1
+chun2
+chun3
+chuo1
+chuo4
+ci1
+ci2
+ci3
+ci4
+cong1
+cong2
+cou4
+cu1
+cu4
+cuan1
+cuan2
+cuan4
+cui1
+cui3
+cui4
+cun1
+cun2
+cun4
+cuo1
+cuo2
+cuo4
+d
+da
+da1
+da2
+da3
+da4
+dai1
+dai2
+dai3
+dai4
+dan1
+dan2
+dan3
+dan4
+dang1
+dang2
+dang3
+dang4
+dao1
+dao2
+dao3
+dao4
+de
+de1
+de2
+dei3
+den4
+deng1
+deng2
+deng3
+deng4
+di1
+di2
+di3
+di4
+dia3
+dian1
+dian2
+dian3
+dian4
+diao1
+diao3
+diao4
+die1
+die2
+die4
+ding1
+ding2
+ding3
+ding4
+diu1
+dong1
+dong3
+dong4
+dou1
+dou2
+dou3
+dou4
+du1
+du2
+du3
+du4
+duan1
+duan2
+duan3
+duan4
+dui1
+dui4
+dun1
+dun3
+dun4
+duo1
+duo2
+duo3
+duo4
+e
+e1
+e2
+e3
+e4
+ei2
+en1
+en4
+er
+er2
+er3
+er4
+f
+fa1
+fa2
+fa3
+fa4
+fan1
+fan2
+fan3
+fan4
+fang1
+fang2
+fang3
+fang4
+fei1
+fei2
+fei3
+fei4
+fen1
+fen2
+fen3
+fen4
+feng1
+feng2
+feng3
+feng4
+fo2
+fou2
+fou3
+fu1
+fu2
+fu3
+fu4
+g
+ga1
+ga2
+ga3
+ga4
+gai1
+gai2
+gai3
+gai4
+gan1
+gan2
+gan3
+gan4
+gang1
+gang2
+gang3
+gang4
+gao1
+gao2
+gao3
+gao4
+ge1
+ge2
+ge3
+ge4
+gei2
+gei3
+gen1
+gen2
+gen3
+gen4
+geng1
+geng3
+geng4
+gong1
+gong3
+gong4
+gou1
+gou2
+gou3
+gou4
+gu
+gu1
+gu2
+gu3
+gu4
+gua1
+gua2
+gua3
+gua4
+guai1
+guai2
+guai3
+guai4
+guan1
+guan2
+guan3
+guan4
+guang1
+guang2
+guang3
+guang4
+gui1
+gui2
+gui3
+gui4
+gun3
+gun4
+guo1
+guo2
+guo3
+guo4
+h
+ha1
+ha2
+ha3
+hai1
+hai2
+hai3
+hai4
+han1
+han2
+han3
+han4
+hang1
+hang2
+hang4
+hao1
+hao2
+hao3
+hao4
+he1
+he2
+he4
+hei1
+hen2
+hen3
+hen4
+heng1
+heng2
+heng4
+hong1
+hong2
+hong3
+hong4
+hou1
+hou2
+hou3
+hou4
+hu1
+hu2
+hu3
+hu4
+hua1
+hua2
+hua4
+huai2
+huai4
+huan1
+huan2
+huan3
+huan4
+huang1
+huang2
+huang3
+huang4
+hui1
+hui2
+hui3
+hui4
+hun1
+hun2
+hun4
+huo
+huo1
+huo2
+huo3
+huo4
+i
+j
+ji1
+ji2
+ji3
+ji4
+jia
+jia1
+jia2
+jia3
+jia4
+jian1
+jian2
+jian3
+jian4
+jiang1
+jiang2
+jiang3
+jiang4
+jiao1
+jiao2
+jiao3
+jiao4
+jie1
+jie2
+jie3
+jie4
+jin1
+jin2
+jin3
+jin4
+jing1
+jing2
+jing3
+jing4
+jiong3
+jiu1
+jiu2
+jiu3
+jiu4
+ju1
+ju2
+ju3
+ju4
+juan1
+juan2
+juan3
+juan4
+jue1
+jue2
+jue4
+jun1
+jun4
+k
+ka1
+ka2
+ka3
+kai1
+kai2
+kai3
+kai4
+kan1
+kan2
+kan3
+kan4
+kang1
+kang2
+kang4
+kao1
+kao2
+kao3
+kao4
+ke1
+ke2
+ke3
+ke4
+ken3
+keng1
+kong1
+kong3
+kong4
+kou1
+kou2
+kou3
+kou4
+ku1
+ku2
+ku3
+ku4
+kua1
+kua3
+kua4
+kuai3
+kuai4
+kuan1
+kuan2
+kuan3
+kuang1
+kuang2
+kuang4
+kui1
+kui2
+kui3
+kui4
+kun1
+kun3
+kun4
+kuo4
+l
+la
+la1
+la2
+la3
+la4
+lai2
+lai4
+lan2
+lan3
+lan4
+lang1
+lang2
+lang3
+lang4
+lao1
+lao2
+lao3
+lao4
+le
+le1
+le4
+lei
+lei1
+lei2
+lei3
+lei4
+leng1
+leng2
+leng3
+leng4
+li
+li1
+li2
+li3
+li4
+lia3
+lian2
+lian3
+lian4
+liang2
+liang3
+liang4
+liao1
+liao2
+liao3
+liao4
+lie1
+lie2
+lie3
+lie4
+lin1
+lin2
+lin3
+lin4
+ling2
+ling3
+ling4
+liu1
+liu2
+liu3
+liu4
+long1
+long2
+long3
+long4
+lou1
+lou2
+lou3
+lou4
+lu1
+lu2
+lu3
+lu4
+luan2
+luan3
+luan4
+lun1
+lun2
+lun4
+luo1
+luo2
+luo3
+luo4
+lv2
+lv3
+lv4
+lve3
+lve4
+m
+ma
+ma1
+ma2
+ma3
+ma4
+mai2
+mai3
+mai4
+man1
+man2
+man3
+man4
+mang2
+mang3
+mao1
+mao2
+mao3
+mao4
+me
+mei2
+mei3
+mei4
+men
+men1
+men2
+men4
+meng
+meng1
+meng2
+meng3
+meng4
+mi1
+mi2
+mi3
+mi4
+mian2
+mian3
+mian4
+miao1
+miao2
+miao3
+miao4
+mie1
+mie4
+min2
+min3
+ming2
+ming3
+ming4
+miu4
+mo1
+mo2
+mo3
+mo4
+mou1
+mou2
+mou3
+mu2
+mu3
+mu4
+n
+n2
+na1
+na2
+na3
+na4
+nai2
+nai3
+nai4
+nan1
+nan2
+nan3
+nan4
+nang1
+nang2
+nang3
+nao1
+nao2
+nao3
+nao4
+ne
+ne2
+ne4
+nei3
+nei4
+nen4
+neng2
+ni1
+ni2
+ni3
+ni4
+nian1
+nian2
+nian3
+nian4
+niang2
+niang4
+niao2
+niao3
+niao4
+nie1
+nie4
+nin2
+ning2
+ning3
+ning4
+niu1
+niu2
+niu3
+niu4
+nong2
+nong4
+nou4
+nu2
+nu3
+nu4
+nuan3
+nuo2
+nuo4
+nv2
+nv3
+nve4
+o
+o1
+o2
+ou1
+ou2
+ou3
+ou4
+p
+pa1
+pa2
+pa4
+pai1
+pai2
+pai3
+pai4
+pan1
+pan2
+pan4
+pang1
+pang2
+pang4
+pao1
+pao2
+pao3
+pao4
+pei1
+pei2
+pei4
+pen1
+pen2
+pen4
+peng1
+peng2
+peng3
+peng4
+pi1
+pi2
+pi3
+pi4
+pian1
+pian2
+pian4
+piao1
+piao2
+piao3
+piao4
+pie1
+pie2
+pie3
+pin1
+pin2
+pin3
+pin4
+ping1
+ping2
+po1
+po2
+po3
+po4
+pou1
+pu1
+pu2
+pu3
+pu4
+q
+qi1
+qi2
+qi3
+qi4
+qia1
+qia3
+qia4
+qian1
+qian2
+qian3
+qian4
+qiang1
+qiang2
+qiang3
+qiang4
+qiao1
+qiao2
+qiao3
+qiao4
+qie1
+qie2
+qie3
+qie4
+qin1
+qin2
+qin3
+qin4
+qing1
+qing2
+qing3
+qing4
+qiong1
+qiong2
+qiu1
+qiu2
+qiu3
+qu1
+qu2
+qu3
+qu4
+quan1
+quan2
+quan3
+quan4
+que1
+que2
+que4
+qun2
+r
+ran2
+ran3
+rang1
+rang2
+rang3
+rang4
+rao2
+rao3
+rao4
+re2
+re3
+re4
+ren2
+ren3
+ren4
+reng1
+reng2
+ri4
+rong1
+rong2
+rong3
+rou2
+rou4
+ru2
+ru3
+ru4
+ruan2
+ruan3
+rui3
+rui4
+run4
+ruo4
+s
+sa1
+sa2
+sa3
+sa4
+sai1
+sai4
+san1
+san2
+san3
+san4
+sang1
+sang3
+sang4
+sao1
+sao2
+sao3
+sao4
+se4
+sen1
+seng1
+sha1
+sha2
+sha3
+sha4
+shai1
+shai2
+shai3
+shai4
+shan1
+shan3
+shan4
+shang
+shang1
+shang3
+shang4
+shao1
+shao2
+shao3
+shao4
+she1
+she2
+she3
+she4
+shei2
+shen1
+shen2
+shen3
+shen4
+sheng1
+sheng2
+sheng3
+sheng4
+shi
+shi1
+shi2
+shi3
+shi4
+shou1
+shou2
+shou3
+shou4
+shu1
+shu2
+shu3
+shu4
+shua1
+shua2
+shua3
+shua4
+shuai1
+shuai3
+shuai4
+shuan1
+shuan4
+shuang1
+shuang3
+shui2
+shui3
+shui4
+shun3
+shun4
+shuo1
+shuo4
+si1
+si2
+si3
+si4
+song1
+song3
+song4
+sou1
+sou3
+sou4
+su1
+su2
+su4
+suan1
+suan4
+sui1
+sui2
+sui3
+sui4
+sun1
+sun3
+suo
+suo1
+suo2
+suo3
+t
+ta1
+ta2
+ta3
+ta4
+tai1
+tai2
+tai4
+tan1
+tan2
+tan3
+tan4
+tang1
+tang2
+tang3
+tang4
+tao1
+tao2
+tao3
+tao4
+te4
+teng2
+ti1
+ti2
+ti3
+ti4
+tian1
+tian2
+tian3
+tiao1
+tiao2
+tiao3
+tiao4
+tie1
+tie2
+tie3
+tie4
+ting1
+ting2
+ting3
+tong1
+tong2
+tong3
+tong4
+tou
+tou1
+tou2
+tou4
+tu1
+tu2
+tu3
+tu4
+tuan1
+tuan2
+tui1
+tui2
+tui3
+tui4
+tun1
+tun2
+tun4
+tuo1
+tuo2
+tuo3
+tuo4
+u
+v
+w
+wa
+wa1
+wa2
+wa3
+wa4
+wai1
+wai3
+wai4
+wan1
+wan2
+wan3
+wan4
+wang1
+wang2
+wang3
+wang4
+wei1
+wei2
+wei3
+wei4
+wen1
+wen2
+wen3
+wen4
+weng1
+weng4
+wo1
+wo2
+wo3
+wo4
+wu1
+wu2
+wu3
+wu4
+x
+xi1
+xi2
+xi3
+xi4
+xia1
+xia2
+xia4
+xian1
+xian2
+xian3
+xian4
+xiang1
+xiang2
+xiang3
+xiang4
+xiao1
+xiao2
+xiao3
+xiao4
+xie1
+xie2
+xie3
+xie4
+xin1
+xin2
+xin4
+xing1
+xing2
+xing3
+xing4
+xiong1
+xiong2
+xiu1
+xiu3
+xiu4
+xu
+xu1
+xu2
+xu3
+xu4
+xuan1
+xuan2
+xuan3
+xuan4
+xue1
+xue2
+xue3
+xue4
+xun1
+xun2
+xun4
+y
+ya
+ya1
+ya2
+ya3
+ya4
+yan1
+yan2
+yan3
+yan4
+yang1
+yang2
+yang3
+yang4
+yao1
+yao2
+yao3
+yao4
+ye1
+ye2
+ye3
+ye4
+yi
+yi1
+yi2
+yi3
+yi4
+yin1
+yin2
+yin3
+yin4
+ying1
+ying2
+ying3
+ying4
+yo1
+yong1
+yong2
+yong3
+yong4
+you1
+you2
+you3
+you4
+yu1
+yu2
+yu3
+yu4
+yuan1
+yuan2
+yuan3
+yuan4
+yue1
+yue4
+yun1
+yun2
+yun3
+yun4
+z
+za1
+za2
+za3
+zai1
+zai3
+zai4
+zan1
+zan2
+zan3
+zan4
+zang1
+zang4
+zao1
+zao2
+zao3
+zao4
+ze2
+ze4
+zei2
+zen3
+zeng1
+zeng4
+zha1
+zha2
+zha3
+zha4
+zhai1
+zhai2
+zhai3
+zhai4
+zhan1
+zhan2
+zhan3
+zhan4
+zhang1
+zhang2
+zhang3
+zhang4
+zhao1
+zhao2
+zhao3
+zhao4
+zhe
+zhe1
+zhe2
+zhe3
+zhe4
+zhen1
+zhen2
+zhen3
+zhen4
+zheng1
+zheng2
+zheng3
+zheng4
+zhi1
+zhi2
+zhi3
+zhi4
+zhong1
+zhong2
+zhong3
+zhong4
+zhou1
+zhou2
+zhou3
+zhou4
+zhu1
+zhu2
+zhu3
+zhu4
+zhua1
+zhua2
+zhua3
+zhuai1
+zhuai3
+zhuai4
+zhuan1
+zhuan2
+zhuan3
+zhuan4
+zhuang1
+zhuang4
+zhui1
+zhui4
+zhun1
+zhun2
+zhun3
+zhuo1
+zhuo2
+zi
+zi1
+zi2
+zi3
+zi4
+zong1
+zong2
+zong3
+zong4
+zou1
+zou2
+zou3
+zou4
+zu1
+zu2
+zu3
+zuan1
+zuan3
+zuan4
+zui2
+zui3
+zui4
+zun1
+zuo
+zuo1
+zuo2
+zuo3
+zuo4
+{
+~
+¡
+¢
+£
+¥
+§
+¨
+©
+«
+®
+¯
+°
+±
+²
+³
+´
+µ
+·
+¹
+º
+»
+¼
+½
+¾
+¿
+À
+Á
+Â
+Ã
+Ä
+Å
+Æ
+Ç
+È
+É
+Ê
+Í
+Î
+Ñ
+Ó
+Ö
+×
+Ø
+Ú
+Ü
+Ý
+Þ
+ß
+à
+á
+â
+ã
+ä
+å
+æ
+ç
+è
+é
+ê
+ë
+ì
+í
+î
+ï
+ð
+ñ
+ò
+ó
+ô
+õ
+ö
+ø
+ù
+ú
+û
+ü
+ý
+Ā
+ā
+ă
+ą
+ć
+Č
+č
+Đ
+đ
+ē
+ė
+ę
+ě
+ĝ
+ğ
+ħ
+ī
+į
+İ
+ı
+Ł
+ł
+ń
+ņ
+ň
+ŋ
+Ō
+ō
+ő
+œ
+ř
+Ś
+ś
+Ş
+ş
+Š
+š
+Ť
+ť
+ũ
+ū
+ź
+Ż
+ż
+Ž
+ž
+ơ
+ư
+ǎ
+ǐ
+ǒ
+ǔ
+ǚ
+ș
+ț
+ɑ
+ɔ
+ɕ
+ə
+ɛ
+ɜ
+ɡ
+ɣ
+ɪ
+ɫ
+ɴ
+ɹ
+ɾ
+ʃ
+ʊ
+ʌ
+ʒ
+ʔ
+ʰ
+ʷ
+ʻ
+ʾ
+ʿ
+ˈ
+ː
+˙
+˜
+ˢ
+́
+̅
+Α
+Β
+Δ
+Ε
+Θ
+Κ
+Λ
+Μ
+Ξ
+Π
+Σ
+Τ
+Φ
+Χ
+Ψ
+Ω
+ά
+έ
+ή
+ί
+α
+β
+γ
+δ
+ε
+ζ
+η
+θ
+ι
+κ
+λ
+μ
+ν
+ξ
+ο
+π
+ρ
+ς
+σ
+τ
+υ
+φ
+χ
+ψ
+ω
+ϊ
+ό
+ύ
+ώ
+ϕ
+ϵ
+Ё
+А
+Б
+В
+Г
+Д
+Е
+Ж
+З
+И
+Й
+К
+Л
+М
+Н
+О
+П
+Р
+С
+Т
+У
+Ф
+Х
+Ц
+Ч
+Ш
+Щ
+Ы
+Ь
+Э
+Ю
+Я
+а
+б
+в
+г
+д
+е
+ж
+з
+и
+й
+к
+л
+м
+н
+о
+п
+р
+с
+т
+у
+ф
+х
+ц
+ч
+ш
+щ
+ъ
+ы
+ь
+э
+ю
+я
+ё
+і
+ְ
+ִ
+ֵ
+ֶ
+ַ
+ָ
+ֹ
+ּ
+־
+ׁ
+א
+ב
+ג
+ד
+ה
+ו
+ז
+ח
+ט
+י
+כ
+ל
+ם
+מ
+ן
+נ
+ס
+ע
+פ
+ק
+ר
+ש
+ת
+أ
+ب
+ة
+ت
+ج
+ح
+د
+ر
+ز
+س
+ص
+ط
+ع
+ق
+ك
+ل
+م
+ن
+ه
+و
+ي
+َ
+ُ
+ِ
+ْ
+ก
+ข
+ง
+จ
+ต
+ท
+น
+ป
+ย
+ร
+ว
+ส
+ห
+อ
+ฮ
+ั
+า
+ี
+ึ
+โ
+ใ
+ไ
+่
+้
+์
+ḍ
+Ḥ
+ḥ
+ṁ
+ṃ
+ṅ
+ṇ
+Ṛ
+ṛ
+Ṣ
+ṣ
+Ṭ
+ṭ
+ạ
+ả
+Ấ
+ấ
+ầ
+ậ
+ắ
+ằ
+ẻ
+ẽ
+ế
+ề
+ể
+ễ
+ệ
+ị
+ọ
+ỏ
+ố
+ồ
+ộ
+ớ
+ờ
+ở
+ụ
+ủ
+ứ
+ữ
+ἀ
+ἁ
+Ἀ
+ἐ
+ἔ
+ἰ
+ἱ
+ὀ
+ὁ
+ὐ
+ὲ
+ὸ
+���
+᾽
+ῆ
+ῇ
+ῶ
+‎
+‑
+‒
+–
+—
+―
+‖
+†
+‡
+•
+…
+‧
+‬
+′
+″
+⁄
+⁡
+⁰
+⁴
+⁵
+⁶
+⁷
+⁸
+⁹
+₁
+₂
+₃
+€
+₱
+₹
+₽
+℃
+ℏ
+ℓ
+№
+ℝ
+™
+⅓
+⅔
+⅛
+→
+∂
+∈
+∑
+−
+∗
+√
+∞
+∫
+≈
+≠
+≡
+≤
+≥
+⋅
+⋯
+█
+♪
+⟨
+⟩
+、
+。
+《
+》
+「
+」
+【
+】
+あ
+う
+え
+お
+か
+が
+き
+ぎ
+く
+ぐ
+け
+げ
+こ
+ご
+さ
+し
+じ
+す
+ず
+せ
+ぜ
+そ
+ぞ
+た
+だ
+ち
+っ
+つ
+で
+と
+ど
+な
+に
+ね
+の
+は
+ば
+ひ
+ぶ
+へ
+べ
+ま
+み
+む
+め
+も
+ゃ
+や
+ゆ
+ょ
+よ
+ら
+り
+る
+れ
+ろ
+わ
+を
+ん
+ァ
+ア
+ィ
+イ
+ウ
+ェ
+エ
+オ
+カ
+ガ
+キ
+ク
+ケ
+ゲ
+コ
+ゴ
+サ
+ザ
+シ
+ジ
+ス
+ズ
+セ
+ゾ
+タ
+ダ
+チ
+ッ
+ツ
+テ
+デ
+ト
+ド
+ナ
+ニ
+ネ
+ノ
+バ
+パ
+ビ
+ピ
+フ
+プ
+ヘ
+ベ
+ペ
+ホ
+ボ
+ポ
+マ
+ミ
+ム
+メ
+モ
+ャ
+ヤ
+ュ
+ユ
+ョ
+ヨ
+ラ
+リ
+ル
+レ
+ロ
+ワ
+ン
+・
+ー
+ㄋ
+ㄍ
+ㄎ
+ㄏ
+ㄓ
+ㄕ
+ㄚ
+ㄜ
+ㄟ
+ㄤ
+ㄥ
+ㄧ
+ㄱ
+ㄴ
+ㄷ
+ㄹ
+ㅁ
+ㅂ
+ㅅ
+ㅈ
+ㅍ
+ㅎ
+ㅏ
+ㅓ
+ㅗ
+ㅜ
+ㅡ
+ㅣ
+㗎
+가
+각
+간
+갈
+감
+갑
+갓
+갔
+강
+같
+개
+거
+건
+걸
+겁
+것
+겉
+게
+겠
+겨
+결
+겼
+경
+계
+고
+곤
+골
+곱
+공
+과
+관
+광
+교
+구
+국
+굴
+귀
+귄
+그
+근
+글
+금
+기
+긴
+길
+까
+깍
+깔
+깜
+깨
+께
+꼬
+꼭
+꽃
+꾸
+꿔
+끔
+끗
+끝
+끼
+나
+난
+날
+남
+납
+내
+냐
+냥
+너
+넘
+넣
+네
+녁
+년
+녕
+노
+녹
+놀
+누
+눈
+느
+는
+늘
+니
+님
+닙
+다
+닥
+단
+달
+닭
+당
+대
+더
+덕
+던
+덥
+데
+도
+독
+동
+돼
+됐
+되
+된
+될
+두
+둑
+둥
+드
+들
+등
+디
+따
+딱
+딸
+땅
+때
+떤
+떨
+떻
+또
+똑
+뚱
+뛰
+뜻
+띠
+라
+락
+란
+람
+랍
+랑
+래
+랜
+러
+런
+럼
+렇
+레
+려
+력
+렵
+렸
+로
+록
+롬
+루
+르
+른
+를
+름
+릉
+리
+릴
+림
+마
+막
+만
+많
+말
+맑
+맙
+맛
+매
+머
+먹
+멍
+메
+면
+명
+몇
+모
+목
+몸
+못
+무
+문
+물
+뭐
+뭘
+미
+민
+밌
+밑
+바
+박
+밖
+반
+받
+발
+밤
+밥
+방
+배
+백
+밸
+뱀
+버
+번
+벌
+벚
+베
+벼
+벽
+별
+병
+보
+복
+본
+볼
+봐
+봤
+부
+분
+불
+비
+빔
+빛
+빠
+빨
+뼈
+뽀
+뿅
+쁘
+사
+산
+살
+삼
+샀
+상
+새
+색
+생
+서
+선
+설
+섭
+섰
+성
+세
+셔
+션
+셨
+소
+속
+손
+송
+수
+숙
+순
+술
+숫
+숭
+숲
+쉬
+쉽
+스
+슨
+습
+슷
+시
+식
+신
+실
+싫
+심
+십
+싶
+싸
+써
+쓰
+쓴
+씌
+씨
+씩
+씬
+아
+악
+안
+않
+알
+야
+약
+얀
+양
+얘
+어
+언
+얼
+엄
+업
+없
+었
+엉
+에
+여
+역
+연
+염
+엽
+영
+옆
+예
+옛
+오
+온
+올
+옷
+옹
+와
+왔
+왜
+요
+욕
+용
+우
+운
+울
+웃
+워
+원
+월
+웠
+위
+윙
+유
+육
+윤
+으
+은
+을
+음
+응
+의
+이
+익
+인
+일
+읽
+임
+입
+있
+자
+작
+잔
+잖
+잘
+잡
+잤
+장
+재
+저
+전
+점
+정
+제
+져
+졌
+조
+족
+좀
+종
+좋
+죠
+주
+준
+줄
+중
+줘
+즈
+즐
+즘
+지
+진
+집
+짜
+짝
+쩌
+쪼
+쪽
+쫌
+쭈
+쯔
+찌
+찍
+차
+착
+찾
+책
+처
+천
+철
+체
+쳐
+쳤
+초
+촌
+추
+출
+춤
+춥
+춰
+치
+친
+칠
+침
+칩
+칼
+커
+켓
+코
+콩
+쿠
+퀴
+크
+큰
+큽
+키
+킨
+타
+태
+터
+턴
+털
+테
+토
+통
+투
+트
+특
+튼
+틀
+티
+팀
+파
+팔
+패
+페
+펜
+펭
+평
+포
+폭
+표
+품
+풍
+프
+플
+피
+필
+하
+학
+한
+할
+함
+합
+항
+해
+햇
+했
+행
+허
+험
+형
+혜
+호
+혼
+홀
+화
+회
+획
+후
+휴
+흐
+흔
+희
+히
+힘
+ﷺ
+ﷻ
+！
+，
+？
+�
+𠮶

f5_tts/infer/infer_cli.py

@@ -0,0 +1,383 @@
+import argparse
+import codecs
+import os
+import re
+from datetime import datetime
+from importlib.resources import files
+from pathlib import Path
+
+import numpy as np
+import soundfile as sf
+import tomli
+from cached_path import cached_path
+from hydra.utils import get_class
+from omegaconf import OmegaConf
+from unidecode import unidecode
+
+from f5_tts.infer.utils_infer import (
+    cfg_strength,
+    cross_fade_duration,
+    device,
+    fix_duration,
+    infer_process,
+    load_model,
+    load_vocoder,
+    mel_spec_type,
+    nfe_step,
+    preprocess_ref_audio_text,
+    remove_silence_for_generated_wav,
+    speed,
+    sway_sampling_coef,
+    target_rms,
+)
+
+
+parser = argparse.ArgumentParser(
+    prog="python3 infer-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",
+    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",
+    type=str,
+    help="The model name: F5TTS_v1_Base | F5TTS_Base | E2TTS_Base | etc.",
+)
+parser.add_argument(
+    "-mc",
+    "--model_cfg",
+    type=str,
+    help="The path to F5-TTS model config file .yaml",
+)
+parser.add_argument(
+    "-p",
+    "--ckpt_file",
+    type=str,
+    help="The path to model checkpoint .pt, leave blank to use default",
+)
+parser.add_argument(
+    "-v",
+    "--vocab_file",
+    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(
+    "-t",
+    "--gen_text",
+    type=str,
+    help="The text to make model synthesize a speech",
+)
+parser.add_argument(
+    "-f",
+    "--gen_file",
+    type=str,
+    help="The file with text to generate, will ignore --gen_text",
+)
+parser.add_argument(
+    "-o",
+    "--output_dir",
+    type=str,
+    help="The path to output folder",
+)
+parser.add_argument(
+    "-w",
+    "--output_file",
+    type=str,
+    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(
+    "--no_legacy_text",
+    action="store_false",
+    help="Not to use lossy ASCII transliterations of unicode text in saved file names.",
+)
+parser.add_argument(
+    "--remove_silence",
+    action="store_true",
+    help="To remove long silence found in ouput",
+)
+parser.add_argument(
+    "--load_vocoder_from_local",
+    action="store_true",
+    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,
+    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}",
+)
+parser.add_argument(
+    "--device",
+    type=str,
+    help="Specify the device to run on",
+)
+args = parser.parse_args()
+
+
+# config file
+
+config = tomli.load(open(args.config, "rb"))
+
+
+# command-line interface parameters
+
+model = args.model or config.get("model", "F5TTS_v1_Base")
+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
+    if args.ref_text is not None
+    else 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)
+use_legacy_text = args.no_legacy_text or config.get("no_legacy_text", False)  # no_legacy_text is a store_false arg
+if save_chunk and use_legacy_text:
+    print(
+        "\nWarning to --save_chunk: lossy ASCII transliterations of unicode text for legacy (.wav) file names, --no_legacy_text to disable.\n"
+    )
+
+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)
+device = args.device or config.get("device", device)
+
+
+# patches for pip pkg user
+if "infer/examples/" in ref_audio:
+    ref_audio = str(files("f5_tts").joinpath(f"{ref_audio}"))
+if "infer/examples/" in gen_file:
+    gen_file = str(files("f5_tts").joinpath(f"{gen_file}"))
+if "voices" in config:
+    for voice in config["voices"]:
+        voice_ref_audio = config["voices"][voice]["ref_audio"]
+        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 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
+
+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=vocoder_name, is_local=load_vocoder_from_local, local_path=vocoder_local_path, device=device
+)
+
+
+# load TTS model
+
+model_cfg = OmegaConf.load(
+    args.model_cfg or config.get("model_cfg", str(files("f5_tts").joinpath(f"configs/{model}.yaml")))
+)
+model_cls = get_class(f"f5_tts.model.{model_cfg.model.backbone}")
+model_arc = model_cfg.model.arch
+
+repo_name, ckpt_step, ckpt_type = "F5-TTS", 1250000, "safetensors"
+
+if model != "F5TTS_Base":
+    assert vocoder_name == model_cfg.model.mel_spec.mel_spec_type
+
+# override for previous models
+if model == "F5TTS_Base":
+    if vocoder_name == "vocos":
+        ckpt_step = 1200000
+    elif vocoder_name == "bigvgan":
+        model = "F5TTS_Base_bigvgan"
+        ckpt_type = "pt"
+elif model == "E2TTS_Base":
+    repo_name = "E2-TTS"
+    ckpt_step = 1200000
+
+if not ckpt_file:
+    ckpt_file = str(cached_path(f"hf://SWivid/{repo_name}/{model}/model_{ckpt_step}.{ckpt_type}"))
+
+print(f"Using {model}...")
+ema_model = load_model(
+    model_cls, model_arc, ckpt_file, mel_spec_type=vocoder_name, vocab_file=vocab_file, device=device
+)
+
+
+# inference process
+
+
+def main():
+    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:
+        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("ref_audio_", voices[voice]["ref_audio"], "\n\n")
+
+    generated_audio_segments = []
+    reg1 = r"(?=\[\w+\])"
+    chunks = re.split(reg1, gen_text)
+    reg2 = r"\[(\w+)\]"
+    for text in chunks:
+        if not text.strip():
+            continue
+        match = re.match(reg2, text)
+        if match:
+            voice = match[1]
+        else:
+            print("No voice tag found, using main.")
+            voice = "main"
+        if voice not in voices:
+            print(f"Voice {voice} not found, using main.")
+            voice = "main"
+        text = re.sub(reg2, "", text)
+        ref_audio_ = voices[voice]["ref_audio"]
+        ref_text_ = voices[voice]["ref_text"]
+        local_speed = voices[voice].get("speed", speed)
+        gen_text_ = text.strip()
+        print(f"Voice: {voice}")
+        audio_segment, final_sample_rate, spectrogram = 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=local_speed,
+            fix_duration=fix_duration,
+            device=device,
+        )
+        generated_audio_segments.append(audio_segment)
+
+        if save_chunk:
+            if len(gen_text_) > 200:
+                gen_text_ = gen_text_[:200] + " ... "
+            if use_legacy_text:
+                gen_text_ = unidecode(gen_text_)
+            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)
+
+        if not os.path.exists(output_dir):
+            os.makedirs(output_dir)
+
+        with open(wave_path, "wb") as f:
+            sf.write(f.name, final_wave, final_sample_rate)
+            # Remove silence
+            if remove_silence:
+                remove_silence_for_generated_wav(f.name)
+            print(f.name)
+
+
+if __name__ == "__main__":
+    main()

f5_tts/infer/infer_gradio.py

@@ -0,0 +1,1121 @@
+# ruff: noqa: E402
+# Above allows ruff to ignore E402: module level import not at top of file
+
+import gc
+import json
+import os
+import re
+import tempfile
+from collections import OrderedDict
+from functools import lru_cache
+from importlib.resources import files
+
+import click
+import gradio as gr
+import numpy as np
+import soundfile as sf
+import torch
+import torchaudio
+from cached_path import cached_path
+from transformers import AutoModelForCausalLM, AutoTokenizer
+
+
+try:
+    import spaces
+
+    USING_SPACES = True
+except ImportError:
+    USING_SPACES = False
+
+
+def gpu_decorator(func):
+    if USING_SPACES:
+        return spaces.GPU(func)
+    else:
+        return func
+
+
+from f5_tts.infer.utils_infer import (
+    infer_process,
+    load_model,
+    load_vocoder,
+    preprocess_ref_audio_text,
+    remove_silence_for_generated_wav,
+    save_spectrogram,
+    tempfile_kwargs,
+)
+from f5_tts.model import DiT, UNetT
+
+
+DEFAULT_TTS_MODEL = "F5-TTS_v1"
+tts_model_choice = DEFAULT_TTS_MODEL
+
+DEFAULT_TTS_MODEL_CFG = [
+    "hf://SWivid/F5-TTS/F5TTS_v1_Base/model_1250000.safetensors",
+    "hf://SWivid/F5-TTS/F5TTS_v1_Base/vocab.txt",
+    json.dumps(dict(dim=1024, depth=22, heads=16, ff_mult=2, text_dim=512, conv_layers=4)),
+]
+
+
+# load models
+
+vocoder = load_vocoder()
+
+
+def load_f5tts():
+    ckpt_path = str(cached_path(DEFAULT_TTS_MODEL_CFG[0]))
+    F5TTS_model_cfg = json.loads(DEFAULT_TTS_MODEL_CFG[2])
+    return load_model(DiT, F5TTS_model_cfg, ckpt_path)
+
+
+def load_e2tts():
+    ckpt_path = str(cached_path("hf://SWivid/E2-TTS/E2TTS_Base/model_1200000.safetensors"))
+    E2TTS_model_cfg = dict(dim=1024, depth=24, heads=16, ff_mult=4, text_mask_padding=False, pe_attn_head=1)
+    return load_model(UNetT, E2TTS_model_cfg, ckpt_path)
+
+
+def load_custom(ckpt_path: str, vocab_path="", model_cfg=None):
+    ckpt_path, vocab_path = ckpt_path.strip(), vocab_path.strip()
+    if ckpt_path.startswith("hf://"):
+        ckpt_path = str(cached_path(ckpt_path))
+    if vocab_path.startswith("hf://"):
+        vocab_path = str(cached_path(vocab_path))
+    if model_cfg is None:
+        model_cfg = json.loads(DEFAULT_TTS_MODEL_CFG[2])
+    elif isinstance(model_cfg, str):
+        model_cfg = json.loads(model_cfg)
+    return load_model(DiT, model_cfg, ckpt_path, vocab_file=vocab_path)
+
+
+F5TTS_ema_model = load_f5tts()
+E2TTS_ema_model = load_e2tts() if USING_SPACES else None
+custom_ema_model, pre_custom_path = None, ""
+
+chat_model_state = None
+chat_tokenizer_state = None
+
+
+@gpu_decorator
+def chat_model_inference(messages, model, tokenizer):
+    """Generate response using Qwen"""
+    text = tokenizer.apply_chat_template(
+        messages,
+        tokenize=False,
+        add_generation_prompt=True,
+    )
+
+    model_inputs = tokenizer([text], return_tensors="pt").to(model.device)
+    generated_ids = model.generate(
+        **model_inputs,
+        max_new_tokens=512,
+        temperature=0.7,
+        top_p=0.95,
+    )
+
+    generated_ids = [
+        output_ids[len(input_ids) :] for input_ids, output_ids in zip(model_inputs.input_ids, generated_ids)
+    ]
+    return tokenizer.batch_decode(generated_ids, skip_special_tokens=True)[0]
+
+
+@gpu_decorator
+def load_text_from_file(file):
+    if file:
+        with open(file, "r", encoding="utf-8") as f:
+            text = f.read().strip()
+    else:
+        text = ""
+    return gr.update(value=text)
+
+
+@lru_cache(maxsize=1000)  # NOTE. need to ensure params of infer() hashable
+@gpu_decorator
+def infer(
+    ref_audio_orig,
+    ref_text,
+    gen_text,
+    model,
+    remove_silence,
+    seed,
+    cross_fade_duration=0.15,
+    nfe_step=32,
+    speed=1,
+    show_info=gr.Info,
+):
+    if not ref_audio_orig:
+        gr.Warning("Please provide reference audio.")
+        return gr.update(), gr.update(), ref_text
+
+    # Set inference seed
+    if seed < 0 or seed > 2**31 - 1:
+        gr.Warning("Seed must in range 0 ~ 2147483647. Using random seed instead.")
+        seed = np.random.randint(0, 2**31 - 1)
+    torch.manual_seed(seed)
+    used_seed = seed
+
+    if not gen_text.strip():
+        gr.Warning("Please enter text to generate or upload a text file.")
+        return gr.update(), gr.update(), ref_text
+
+    ref_audio, ref_text = preprocess_ref_audio_text(ref_audio_orig, ref_text, show_info=show_info)
+
+    if model == DEFAULT_TTS_MODEL:
+        ema_model = F5TTS_ema_model
+    elif model == "E2-TTS":
+        global E2TTS_ema_model
+        if E2TTS_ema_model is None:
+            show_info("Loading E2-TTS model...")
+            E2TTS_ema_model = load_e2tts()
+        ema_model = E2TTS_ema_model
+    elif isinstance(model, tuple) and model[0] == "Custom":
+        assert not USING_SPACES, "Only official checkpoints allowed in Spaces."
+        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], model_cfg=model[3])
+            pre_custom_path = model[1]
+        ema_model = custom_ema_model
+
+    final_wave, final_sample_rate, combined_spectrogram = infer_process(
+        ref_audio,
+        ref_text,
+        gen_text,
+        ema_model,
+        vocoder,
+        cross_fade_duration=cross_fade_duration,
+        nfe_step=nfe_step,
+        speed=speed,
+        show_info=show_info,
+        progress=gr.Progress(),
+    )
+
+    # Remove silence
+    if remove_silence:
+        with tempfile.NamedTemporaryFile(suffix=".wav", **tempfile_kwargs) as f:
+            temp_path = f.name
+        try:
+            sf.write(temp_path, final_wave, final_sample_rate)
+            remove_silence_for_generated_wav(f.name)
+            final_wave, _ = torchaudio.load(f.name)
+        finally:
+            os.unlink(temp_path)
+        final_wave = final_wave.squeeze().cpu().numpy()
+
+    # Save the spectrogram
+    with tempfile.NamedTemporaryFile(suffix=".png", **tempfile_kwargs) as tmp_spectrogram:
+        spectrogram_path = tmp_spectrogram.name
+    save_spectrogram(combined_spectrogram, spectrogram_path)
+
+    return (final_sample_rate, final_wave), spectrogram_path, ref_text, used_seed
+
+
+with gr.Blocks() as app_tts:
+    gr.Markdown("# Batched TTS")
+    ref_audio_input = gr.Audio(label="Reference Audio", type="filepath")
+    with gr.Row():
+        gen_text_input = gr.Textbox(
+            label="Text to Generate",
+            lines=10,
+            max_lines=40,
+            scale=4,
+        )
+        gen_text_file = gr.File(label="Load Text to Generate from File (.txt)", file_types=[".txt"], scale=1)
+    generate_btn = gr.Button("Synthesize", variant="primary")
+    with gr.Accordion("Advanced Settings", open=False):
+        with gr.Row():
+            ref_text_input = gr.Textbox(
+                label="Reference Text",
+                info="Leave blank to automatically transcribe the reference audio. If you enter text or upload a file, it will override automatic transcription.",
+                lines=2,
+                scale=4,
+            )
+            ref_text_file = gr.File(label="Load Reference Text from File (.txt)", file_types=[".txt"], scale=1)
+        with gr.Row():
+            randomize_seed = gr.Checkbox(
+                label="Randomize Seed",
+                info="Check to use a random seed for each generation. Uncheck to use the seed specified.",
+                value=True,
+                scale=3,
+            )
+            seed_input = gr.Number(show_label=False, value=0, precision=0, scale=1)
+            with gr.Column(scale=4):
+                remove_silence = gr.Checkbox(
+                    label="Remove Silences",
+                    info="If undesired long silence(s) produced, turn on to automatically detect and crop.",
+                    value=False,
+                )
+        speed_slider = gr.Slider(
+            label="Speed",
+            minimum=0.3,
+            maximum=2.0,
+            value=1.0,
+            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,
+            maximum=1.0,
+            value=0.15,
+            step=0.01,
+            info="Set the duration of the cross-fade between audio clips.",
+        )
+
+    audio_output = gr.Audio(label="Synthesized Audio")
+    spectrogram_output = gr.Image(label="Spectrogram")
+
+    @gpu_decorator
+    def basic_tts(
+        ref_audio_input,
+        ref_text_input,
+        gen_text_input,
+        remove_silence,
+        randomize_seed,
+        seed_input,
+        cross_fade_duration_slider,
+        nfe_slider,
+        speed_slider,
+    ):
+        if randomize_seed:
+            seed_input = np.random.randint(0, 2**31 - 1)
+
+        audio_out, spectrogram_path, ref_text_out, used_seed = infer(
+            ref_audio_input,
+            ref_text_input,
+            gen_text_input,
+            tts_model_choice,
+            remove_silence,
+            seed=seed_input,
+            cross_fade_duration=cross_fade_duration_slider,
+            nfe_step=nfe_slider,
+            speed=speed_slider,
+        )
+        return audio_out, spectrogram_path, ref_text_out, used_seed
+
+    gen_text_file.upload(
+        load_text_from_file,
+        inputs=[gen_text_file],
+        outputs=[gen_text_input],
+    )
+
+    ref_text_file.upload(
+        load_text_from_file,
+        inputs=[ref_text_file],
+        outputs=[ref_text_input],
+    )
+
+    ref_audio_input.clear(
+        lambda: [None, None],
+        None,
+        [ref_text_input, ref_text_file],
+    )
+
+    generate_btn.click(
+        basic_tts,
+        inputs=[
+            ref_audio_input,
+            ref_text_input,
+            gen_text_input,
+            remove_silence,
+            randomize_seed,
+            seed_input,
+            cross_fade_duration_slider,
+            nfe_slider,
+            speed_slider,
+        ],
+        outputs=[audio_output, spectrogram_output, ref_text_input, seed_input],
+    )
+
+
+def parse_speechtypes_text(gen_text):
+    # Pattern to find {str} or {"name": str, "seed": int, "speed": float}
+    pattern = r"(\{.*?\})"
+
+    # Split the text by the pattern
+    tokens = re.split(pattern, gen_text)
+
+    segments = []
+
+    current_type_dict = {
+        "name": "Regular",
+        "seed": -1,
+        "speed": 1.0,
+    }
+
+    for i in range(len(tokens)):
+        if i % 2 == 0:
+            # This is text
+            text = tokens[i].strip()
+            if text:
+                current_type_dict["text"] = text
+                segments.append(current_type_dict)
+        else:
+            # This is type
+            type_str = tokens[i].strip()
+            try:  # if type dict
+                current_type_dict = json.loads(type_str)
+            except json.decoder.JSONDecodeError:
+                type_str = type_str[1:-1]  # remove brace {}
+                current_type_dict = {"name": type_str, "seed": -1, "speed": 1.0}
+
+    return segments
+
+
+with gr.Blocks() as app_multistyle:
+    # New section for multistyle generation
+    gr.Markdown(
+        """
+    # Multiple Speech-Type Generation
+
+    This section allows you to generate multiple speech types or multiple people's voices. Enter your text in the format shown below, or upload a .txt file with the same format. The system will generate speech using the appropriate type. If unspecified, the model will use the regular speech type. The current speech type will be used until the next speech type is specified.
+    """
+    )
+
+    with gr.Row():
+        gr.Markdown(
+            """
+            **Example Input:** <br>
+            {Regular} Hello, I'd like to order a sandwich please. <br>
+            {Surprised} What do you mean you're out of bread? <br>
+            {Sad} I really wanted a sandwich though... <br>
+            {Angry} You know what, darn you and your little shop! <br>
+            {Whisper} I'll just go back home and cry now. <br>
+            {Shouting} Why me?!
+            """
+        )
+
+        gr.Markdown(
+            """
+            **Example Input 2:** <br>
+            {"name": "Speaker1_Happy", "seed": -1, "speed": 1} Hello, I'd like to order a sandwich please. <br>
+            {"name": "Speaker2_Regular", "seed": -1, "speed": 1} Sorry, we're out of bread. <br>
+            {"name": "Speaker1_Sad", "seed": -1, "speed": 1} I really wanted a sandwich though... <br>
+            {"name": "Speaker2_Whisper", "seed": -1, "speed": 1} I'll give you the last one I was hiding.
+            """
+        )
+
+    gr.Markdown(
+        'Upload different audio clips for each speech type. The first speech type is mandatory. You can add additional speech types by clicking the "Add Speech Type" button.'
+    )
+
+    # Regular speech type (mandatory)
+    with gr.Row(variant="compact") as regular_row:
+        with gr.Column(scale=1, min_width=160):
+            regular_name = gr.Textbox(value="Regular", label="Speech Type Name")
+            regular_insert = gr.Button("Insert Label", variant="secondary")
+        with gr.Column(scale=3):
+            regular_audio = gr.Audio(label="Regular Reference Audio", type="filepath")
+        with gr.Column(scale=3):
+            regular_ref_text = gr.Textbox(label="Reference Text (Regular)", lines=4)
+            with gr.Row():
+                regular_seed_slider = gr.Slider(
+                    show_label=False, minimum=-1, maximum=999, value=-1, step=1, info="Seed, -1 for random"
+                )
+                regular_speed_slider = gr.Slider(
+                    show_label=False, minimum=0.3, maximum=2.0, value=1.0, step=0.1, info="Adjust the speed"
+                )
+        with gr.Column(scale=1, min_width=160):
+            regular_ref_text_file = gr.File(label="Load Reference Text from File (.txt)", file_types=[".txt"])
+
+    # Regular speech type (max 100)
+    max_speech_types = 100
+    speech_type_rows = [regular_row]
+    speech_type_names = [regular_name]
+    speech_type_audios = [regular_audio]
+    speech_type_ref_texts = [regular_ref_text]
+    speech_type_ref_text_files = [regular_ref_text_file]
+    speech_type_seeds = [regular_seed_slider]
+    speech_type_speeds = [regular_speed_slider]
+    speech_type_delete_btns = [None]
+    speech_type_insert_btns = [regular_insert]
+
+    # Additional speech types (99 more)
+    for i in range(max_speech_types - 1):
+        with gr.Row(variant="compact", visible=False) as row:
+            with gr.Column(scale=1, min_width=160):
+                name_input = gr.Textbox(label="Speech Type Name")
+                insert_btn = gr.Button("Insert Label", variant="secondary")
+                delete_btn = gr.Button("Delete Type", variant="stop")
+            with gr.Column(scale=3):
+                audio_input = gr.Audio(label="Reference Audio", type="filepath")
+            with gr.Column(scale=3):
+                ref_text_input = gr.Textbox(label="Reference Text", lines=4)
+                with gr.Row():
+                    seed_input = gr.Slider(
+                        show_label=False, minimum=-1, maximum=999, value=-1, step=1, info="Seed. -1 for random"
+                    )
+                    speed_input = gr.Slider(
+                        show_label=False, minimum=0.3, maximum=2.0, value=1.0, step=0.1, info="Adjust the speed"
+                    )
+            with gr.Column(scale=1, min_width=160):
+                ref_text_file_input = gr.File(label="Load Reference Text from File (.txt)", file_types=[".txt"])
+        speech_type_rows.append(row)
+        speech_type_names.append(name_input)
+        speech_type_audios.append(audio_input)
+        speech_type_ref_texts.append(ref_text_input)
+        speech_type_ref_text_files.append(ref_text_file_input)
+        speech_type_seeds.append(seed_input)
+        speech_type_speeds.append(speed_input)
+        speech_type_delete_btns.append(delete_btn)
+        speech_type_insert_btns.append(insert_btn)
+
+    # Global logic for all speech types
+    for i in range(max_speech_types):
+        speech_type_audios[i].clear(
+            lambda: [None, None],
+            None,
+            [speech_type_ref_texts[i], speech_type_ref_text_files[i]],
+        )
+        speech_type_ref_text_files[i].upload(
+            load_text_from_file,
+            inputs=[speech_type_ref_text_files[i]],
+            outputs=[speech_type_ref_texts[i]],
+        )
+
+    # Button to add speech type
+    add_speech_type_btn = gr.Button("Add Speech Type")
+
+    # Keep track of autoincrement of speech types, no roll back
+    speech_type_count = 1
+
+    # Function to add a speech type
+    def add_speech_type_fn():
+        row_updates = [gr.update() for _ in range(max_speech_types)]
+        global speech_type_count
+        if speech_type_count < max_speech_types:
+            row_updates[speech_type_count] = gr.update(visible=True)
+            speech_type_count += 1
+        else:
+            gr.Warning("Exhausted maximum number of speech types. Consider restart the app.")
+        return row_updates
+
+    add_speech_type_btn.click(add_speech_type_fn, outputs=speech_type_rows)
+
+    # Function to delete a speech type
+    def delete_speech_type_fn():
+        return gr.update(visible=False), None, None, None, None
+
+    # Update delete button clicks and ref text file changes
+    for i in range(1, len(speech_type_delete_btns)):
+        speech_type_delete_btns[i].click(
+            delete_speech_type_fn,
+            outputs=[
+                speech_type_rows[i],
+                speech_type_names[i],
+                speech_type_audios[i],
+                speech_type_ref_texts[i],
+                speech_type_ref_text_files[i],
+            ],
+        )
+
+    # Text input for the prompt
+    with gr.Row():
+        gen_text_input_multistyle = gr.Textbox(
+            label="Text to Generate",
+            lines=10,
+            max_lines=40,
+            scale=4,
+            placeholder="Enter the script with speaker names (or emotion types) at the start of each block, e.g.:\n\n{Regular} Hello, I'd like to order a sandwich please.\n{Surprised} What do you mean you're out of bread?\n{Sad} I really wanted a sandwich though...\n{Angry} You know what, darn you and your little shop!\n{Whisper} I'll just go back home and cry now.\n{Shouting} Why me?!",
+        )
+        gen_text_file_multistyle = gr.File(label="Load Text to Generate from File (.txt)", file_types=[".txt"], scale=1)
+
+    def make_insert_speech_type_fn(index):
+        def insert_speech_type_fn(current_text, speech_type_name, speech_type_seed, speech_type_speed):
+            current_text = current_text or ""
+            if not speech_type_name:
+                gr.Warning("Please enter speech type name before insert.")
+                return current_text
+            speech_type_dict = {
+                "name": speech_type_name,
+                "seed": speech_type_seed,
+                "speed": speech_type_speed,
+            }
+            updated_text = current_text + json.dumps(speech_type_dict) + " "
+            return updated_text
+
+        return insert_speech_type_fn
+
+    for i, insert_btn in enumerate(speech_type_insert_btns):
+        insert_fn = make_insert_speech_type_fn(i)
+        insert_btn.click(
+            insert_fn,
+            inputs=[gen_text_input_multistyle, speech_type_names[i], speech_type_seeds[i], speech_type_speeds[i]],
+            outputs=gen_text_input_multistyle,
+        )
+
+    with gr.Accordion("Advanced Settings", open=True):
+        with gr.Row():
+            with gr.Column():
+                show_cherrypick_multistyle = gr.Checkbox(
+                    label="Show Cherry-pick Interface",
+                    info="Turn on to show interface, picking seeds from previous generations.",
+                    value=False,
+                )
+            with gr.Column():
+                remove_silence_multistyle = gr.Checkbox(
+                    label="Remove Silences",
+                    info="Turn on to automatically detect and crop long silences.",
+                    value=True,
+                )
+
+    # Generate button
+    generate_multistyle_btn = gr.Button("Generate Multi-Style Speech", variant="primary")
+
+    # Output audio
+    audio_output_multistyle = gr.Audio(label="Synthesized Audio")
+
+    # Used seed gallery
+    cherrypick_interface_multistyle = gr.Textbox(
+        label="Cherry-pick Interface",
+        lines=10,
+        max_lines=40,
+        show_copy_button=True,
+        interactive=False,
+        visible=False,
+    )
+
+    # Logic control to show/hide the cherrypick interface
+    show_cherrypick_multistyle.change(
+        lambda is_visible: gr.update(visible=is_visible),
+        show_cherrypick_multistyle,
+        cherrypick_interface_multistyle,
+    )
+
+    # Function to load text to generate from file
+    gen_text_file_multistyle.upload(
+        load_text_from_file,
+        inputs=[gen_text_file_multistyle],
+        outputs=[gen_text_input_multistyle],
+    )
+
+    @gpu_decorator
+    def generate_multistyle_speech(
+        gen_text,
+        *args,
+    ):
+        speech_type_names_list = args[:max_speech_types]
+        speech_type_audios_list = args[max_speech_types : 2 * max_speech_types]
+        speech_type_ref_texts_list = args[2 * max_speech_types : 3 * max_speech_types]
+        remove_silence = args[3 * max_speech_types]
+        # Collect the speech types and their audios into a dict
+        speech_types = OrderedDict()
+
+        ref_text_idx = 0
+        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}
+            else:
+                speech_types[f"@{ref_text_idx}@"] = {"audio": "", "ref_text": ""}
+            ref_text_idx += 1
+
+        # Parse the gen_text into segments
+        segments = parse_speechtypes_text(gen_text)
+
+        # For each segment, generate speech
+        generated_audio_segments = []
+        current_type_name = "Regular"
+        inference_meta_data = ""
+
+        for segment in segments:
+            name = segment["name"]
+            seed_input = segment["seed"]
+            speed = segment["speed"]
+            text = segment["text"]
+
+            if name in speech_types:
+                current_type_name = name
+            else:
+                gr.Warning(f"Type {name} is not available, will use Regular as default.")
+                current_type_name = "Regular"
+
+            try:
+                ref_audio = speech_types[current_type_name]["audio"]
+            except KeyError:
+                gr.Warning(f"Please provide reference audio for type {current_type_name}.")
+                return [None] + [speech_types[name]["ref_text"] for name in speech_types] + [None]
+            ref_text = speech_types[current_type_name].get("ref_text", "")
+
+            if seed_input == -1:
+                seed_input = np.random.randint(0, 2**31 - 1)
+
+            # Generate or retrieve speech for this segment
+            audio_out, _, ref_text_out, used_seed = infer(
+                ref_audio,
+                ref_text,
+                text,
+                tts_model_choice,
+                remove_silence,
+                seed=seed_input,
+                cross_fade_duration=0,
+                speed=speed,
+                show_info=print,  # no pull to top when generating
+            )
+            sr, audio_data = audio_out
+
+            generated_audio_segments.append(audio_data)
+            speech_types[current_type_name]["ref_text"] = ref_text_out
+            inference_meta_data += json.dumps(dict(name=name, seed=used_seed, speed=speed)) + f" {text}\n"
+
+        # Concatenate all audio segments
+        if generated_audio_segments:
+            final_audio_data = np.concatenate(generated_audio_segments)
+            return (
+                [(sr, final_audio_data)]
+                + [speech_types[name]["ref_text"] for name in speech_types]
+                + [inference_meta_data]
+            )
+        else:
+            gr.Warning("No audio generated.")
+            return [None] + [speech_types[name]["ref_text"] for name in speech_types] + [None]
+
+    generate_multistyle_btn.click(
+        generate_multistyle_speech,
+        inputs=[
+            gen_text_input_multistyle,
+        ]
+        + speech_type_names
+        + speech_type_audios
+        + speech_type_ref_texts
+        + [
+            remove_silence_multistyle,
+        ],
+        outputs=[audio_output_multistyle] + speech_type_ref_texts + [cherrypick_interface_multistyle],
+    )
+
+    # Validation function to disable Generate button if speech types are missing
+    def validate_speech_types(gen_text, regular_name, *args):
+        speech_type_names_list = args
+
+        # 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_speechtypes_text(gen_text)
+        speech_types_in_text = set(segment["name"] 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_multistyle.change(
+        validate_speech_types,
+        inputs=[gen_text_input_multistyle, regular_name] + speech_type_names,
+        outputs=generate_multistyle_btn,
+    )
+
+
+with gr.Blocks() as app_chat:
+    gr.Markdown(
+        """
+# Voice Chat
+Have a conversation with an AI using your reference voice!
+1. Upload a reference audio clip and optionally its transcript (via text or .txt file).
+2. Load the chat model.
+3. Record your message through your microphone or type it.
+4. The AI will respond using the reference voice.
+"""
+    )
+
+    chat_model_name_list = [
+        "Qwen/Qwen2.5-3B-Instruct",
+        "microsoft/Phi-4-mini-instruct",
+    ]
+
+    @gpu_decorator
+    def load_chat_model(chat_model_name):
+        show_info = gr.Info
+        global chat_model_state, chat_tokenizer_state
+        if chat_model_state is not None:
+            chat_model_state = None
+            chat_tokenizer_state = None
+            gc.collect()
+            torch.cuda.empty_cache()
+
+        show_info(f"Loading chat model: {chat_model_name}")
+        chat_model_state = AutoModelForCausalLM.from_pretrained(chat_model_name, torch_dtype="auto", device_map="auto")
+        chat_tokenizer_state = AutoTokenizer.from_pretrained(chat_model_name)
+        show_info(f"Chat model {chat_model_name} loaded successfully!")
+
+        return gr.update(visible=False), gr.update(visible=True)
+
+    if USING_SPACES:
+        load_chat_model(chat_model_name_list[0])
+
+    chat_model_name_input = gr.Dropdown(
+        choices=chat_model_name_list,
+        value=chat_model_name_list[0],
+        label="Chat Model Name",
+        info="Enter the name of a HuggingFace chat model",
+        allow_custom_value=not USING_SPACES,
+    )
+    load_chat_model_btn = gr.Button("Load Chat Model", variant="primary", visible=not USING_SPACES)
+    chat_interface_container = gr.Column(visible=USING_SPACES)
+
+    chat_model_name_input.change(
+        lambda: gr.update(visible=True),
+        None,
+        load_chat_model_btn,
+        show_progress="hidden",
+    )
+    load_chat_model_btn.click(
+        load_chat_model, inputs=[chat_model_name_input], outputs=[load_chat_model_btn, chat_interface_container]
+    )
+
+    with chat_interface_container:
+        with gr.Row():
+            with gr.Column():
+                ref_audio_chat = gr.Audio(label="Reference Audio", type="filepath")
+            with gr.Column():
+                with gr.Accordion("Advanced Settings", open=False):
+                    with gr.Row():
+                        ref_text_chat = gr.Textbox(
+                            label="Reference Text",
+                            info="Optional: Leave blank to auto-transcribe",
+                            lines=2,
+                            scale=3,
+                        )
+                        ref_text_file_chat = gr.File(
+                            label="Load Reference Text from File (.txt)", file_types=[".txt"], scale=1
+                        )
+                    with gr.Row():
+                        randomize_seed_chat = gr.Checkbox(
+                            label="Randomize Seed",
+                            value=True,
+                            info="Uncheck to use the seed specified.",
+                            scale=3,
+                        )
+                        seed_input_chat = gr.Number(show_label=False, value=0, precision=0, scale=1)
+                    remove_silence_chat = gr.Checkbox(
+                        label="Remove Silences",
+                        value=True,
+                    )
+                    system_prompt_chat = gr.Textbox(
+                        label="System Prompt",
+                        value="You are not an AI assistant, you are whoever the user says you are. You must stay in character. Keep your responses concise since they will be spoken out loud.",
+                        lines=2,
+                    )
+
+        chatbot_interface = gr.Chatbot(label="Conversation", type="messages")
+
+        with gr.Row():
+            with gr.Column():
+                audio_input_chat = gr.Microphone(
+                    label="Speak your message",
+                    type="filepath",
+                )
+                audio_output_chat = gr.Audio(autoplay=True)
+            with gr.Column():
+                text_input_chat = gr.Textbox(
+                    label="Type your message",
+                    lines=1,
+                )
+                send_btn_chat = gr.Button("Send Message")
+                clear_btn_chat = gr.Button("Clear Conversation")
+
+        # Modify process_audio_input to generate user input
+        @gpu_decorator
+        def process_audio_input(conv_state, audio_path, text):
+            """Handle audio or text input from user"""
+
+            if not audio_path and not text.strip():
+                return conv_state
+
+            if audio_path:
+                text = preprocess_ref_audio_text(audio_path, text)[1]
+            if not text.strip():
+                return conv_state
+
+            conv_state.append({"role": "user", "content": text})
+            return conv_state
+
+        # Use model and tokenizer from state to get text response
+        @gpu_decorator
+        def generate_text_response(conv_state, system_prompt):
+            """Generate text response from AI"""
+
+            system_prompt_state = [{"role": "system", "content": system_prompt}]
+            response = chat_model_inference(system_prompt_state + conv_state, chat_model_state, chat_tokenizer_state)
+
+            conv_state.append({"role": "assistant", "content": response})
+            return conv_state
+
+        @gpu_decorator
+        def generate_audio_response(conv_state, ref_audio, ref_text, remove_silence, randomize_seed, seed_input):
+            """Generate TTS audio for AI response"""
+            if not conv_state or not ref_audio:
+                return None, ref_text, seed_input
+
+            last_ai_response = conv_state[-1]["content"]
+            if not last_ai_response or conv_state[-1]["role"] != "assistant":
+                return None, ref_text, seed_input
+
+            if randomize_seed:
+                seed_input = np.random.randint(0, 2**31 - 1)
+
+            audio_result, _, ref_text_out, used_seed = infer(
+                ref_audio,
+                ref_text,
+                last_ai_response,
+                tts_model_choice,
+                remove_silence,
+                seed=seed_input,
+                cross_fade_duration=0.15,
+                speed=1.0,
+                show_info=print,  # show_info=print no pull to top when generating
+            )
+            return audio_result, ref_text_out, used_seed
+
+        def clear_conversation():
+            """Reset the conversation"""
+            return [], None
+
+        ref_text_file_chat.upload(
+            load_text_from_file,
+            inputs=[ref_text_file_chat],
+            outputs=[ref_text_chat],
+        )
+
+        for user_operation in [audio_input_chat.stop_recording, text_input_chat.submit, send_btn_chat.click]:
+            user_operation(
+                process_audio_input,
+                inputs=[chatbot_interface, audio_input_chat, text_input_chat],
+                outputs=[chatbot_interface],
+            ).then(
+                generate_text_response,
+                inputs=[chatbot_interface, system_prompt_chat],
+                outputs=[chatbot_interface],
+            ).then(
+                generate_audio_response,
+                inputs=[
+                    chatbot_interface,
+                    ref_audio_chat,
+                    ref_text_chat,
+                    remove_silence_chat,
+                    randomize_seed_chat,
+                    seed_input_chat,
+                ],
+                outputs=[audio_output_chat, ref_text_chat, seed_input_chat],
+            ).then(
+                lambda: [None, None],
+                None,
+                [audio_input_chat, text_input_chat],
+            )
+
+        # Handle clear button or system prompt change and reset conversation
+        for user_operation in [clear_btn_chat.click, system_prompt_chat.change, chatbot_interface.clear]:
+            user_operation(
+                clear_conversation,
+                outputs=[chatbot_interface, audio_output_chat],
+            )
+
+
+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 initial chunk generation and podcast app exploration
+* [jpgallegoar](https://github.com/jpgallegoar) for multiple speech-type generation & voice chat
+""")
+
+
+with gr.Blocks() as app:
+    gr.Markdown(
+        f"""
+# E2/F5 TTS
+
+This is {"a local web UI for [F5 TTS](https://github.com/SWivid/F5-TTS)" if not USING_SPACES else "an online demo for [F5-TTS](https://github.com/SWivid/F5-TTS)"} with advanced batch processing support. This app supports the following TTS models:
+
+* [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)
+
+The checkpoints currently support English and Chinese.
+
+If you're having issues, try converting your reference audio to WAV or MP3, clipping it to 12s with  ✂  in the bottom right corner (otherwise might have non-optimal auto-trimmed result).
+
+**NOTE: Reference text will be automatically transcribed with Whisper if not provided. For best results, keep your reference clips short (<12s). Ensure the audio is fully uploaded before generating.**
+"""
+    )
+
+    last_used_custom = files("f5_tts").joinpath("infer/.cache/last_used_custom_model_info_v1.txt")
+
+    def load_last_used_custom():
+        try:
+            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 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, custom_model_cfg = load_last_used_custom()
+            tts_model_choice = ("Custom", custom_ckpt_path, custom_vocab_path, 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), gr.update(visible=False)
+
+    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, 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:
+            choose_tts_model = gr.Radio(
+                choices=[DEFAULT_TTS_MODEL, "E2-TTS", "Custom"], label="Choose TTS Model", value=DEFAULT_TTS_MODEL
+            )
+        else:
+            choose_tts_model = gr.Radio(
+                choices=[DEFAULT_TTS_MODEL, "E2-TTS"], label="Choose TTS Model", value=DEFAULT_TTS_MODEL
+            )
+        custom_ckpt_path = gr.Dropdown(
+            choices=[DEFAULT_TTS_MODEL_CFG[0]],
+            value=load_last_used_custom()[0],
+            allow_custom_value=True,
+            label="Model: local_path | hf://user_id/repo_id/model_ckpt",
+            visible=False,
+        )
+        custom_vocab_path = gr.Dropdown(
+            choices=[DEFAULT_TTS_MODEL_CFG[1]],
+            value=load_last_used_custom()[1],
+            allow_custom_value=True,
+            label="Vocab: local_path | hf://user_id/repo_id/vocab_file",
+            visible=False,
+        )
+        custom_model_cfg = gr.Dropdown(
+            choices=[
+                DEFAULT_TTS_MODEL_CFG[2],
+                json.dumps(
+                    dict(
+                        dim=1024,
+                        depth=22,
+                        heads=16,
+                        ff_mult=2,
+                        text_dim=512,
+                        text_mask_padding=False,
+                        conv_layers=4,
+                        pe_attn_head=1,
+                    )
+                ),
+                json.dumps(
+                    dict(
+                        dim=768,
+                        depth=18,
+                        heads=12,
+                        ff_mult=2,
+                        text_dim=512,
+                        text_mask_padding=False,
+                        conv_layers=4,
+                        pe_attn_head=1,
+                    )
+                ),
+            ],
+            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, custom_model_cfg],
+        show_progress="hidden",
+    )
+    custom_ckpt_path.change(
+        set_custom_model,
+        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, 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",
+    )
+
+    gr.TabbedInterface(
+        [app_tts, app_multistyle, app_chat, app_credits],
+        ["Basic-TTS", "Multi-Speech", "Voice-Chat", "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")
+@click.option(
+    "--root_path",
+    "-r",
+    default=None,
+    type=str,
+    help='The root path (or "mount point") of the application, if it\'s not served from the root ("/") of the domain. Often used when the application is behind a reverse proxy that forwards requests to the application, e.g. set "/myapp" or full URL for application served at "https://example.com/myapp".',
+)
+@click.option(
+    "--inbrowser",
+    "-i",
+    is_flag=True,
+    default=False,
+    help="Automatically launch the interface in the default web browser",
+)
+def main(port, host, share, api, root_path, inbrowser):
+    global app
+    print("Starting app...")
+    app.queue(api_open=api).launch(
+        server_name=host,
+        server_port=port,
+        share=share,
+        show_api=api,
+        root_path=root_path,
+        inbrowser=inbrowser,
+    )
+
+
+if __name__ == "__main__":
+    if not USING_SPACES:
+        main()
+    else:
+        app.queue().launch()

f5_tts/infer/speech_edit.py

@@ -0,0 +1,205 @@
+import os
+
+
+os.environ["PYTORCH_ENABLE_MPS_FALLBACK"] = "1"  # for MPS device compatibility
+
+from importlib.resources import files
+
+import torch
+import torch.nn.functional as F
+import torchaudio
+from cached_path import cached_path
+from hydra.utils import get_class
+from omegaconf import OmegaConf
+
+from f5_tts.infer.utils_infer import load_checkpoint, load_vocoder, save_spectrogram
+from f5_tts.model import CFM
+from f5_tts.model.utils import convert_char_to_pinyin, get_tokenizer
+
+
+device = (
+    "cuda"
+    if torch.cuda.is_available()
+    else "xpu"
+    if torch.xpu.is_available()
+    else "mps"
+    if torch.backends.mps.is_available()
+    else "cpu"
+)
+
+
+# ---------------------- infer setting ---------------------- #
+
+seed = None  # int | None
+
+exp_name = "F5TTS_v1_Base"  # F5TTS_v1_Base | E2TTS_Base
+ckpt_step = 1250000
+
+nfe_step = 32  # 16, 32
+cfg_strength = 2.0
+ode_method = "euler"  # euler | midpoint
+sway_sampling_coef = -1.0
+speed = 1.0
+target_rms = 0.1
+
+
+model_cfg = OmegaConf.load(str(files("f5_tts").joinpath(f"configs/{exp_name}.yaml")))
+model_cls = get_class(f"f5_tts.model.{model_cfg.model.backbone}")
+model_arc = model_cfg.model.arch
+
+dataset_name = model_cfg.datasets.name
+tokenizer = model_cfg.model.tokenizer
+
+mel_spec_type = model_cfg.model.mel_spec.mel_spec_type
+target_sample_rate = model_cfg.model.mel_spec.target_sample_rate
+n_mel_channels = model_cfg.model.mel_spec.n_mel_channels
+hop_length = model_cfg.model.mel_spec.hop_length
+win_length = model_cfg.model.mel_spec.win_length
+n_fft = model_cfg.model.mel_spec.n_fft
+
+
+# ckpt_path = str(files("f5_tts").joinpath("../../")) + f"/ckpts/{exp_name}/model_{ckpt_step}.safetensors"
+ckpt_path = str(cached_path(f"hf://SWivid/F5-TTS/{exp_name}/model_{ckpt_step}.safetensors"))
+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 "src/f5_tts/infer/examples/basic/basic_ref_en.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 = str(files("f5_tts").joinpath("infer/examples/basic/basic_ref_en.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 = "src/f5_tts/infer/examples/basic/basic_ref_zh.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 mel_spec_type == "vocos":
+    vocoder_local_path = "../checkpoints/charactr/vocos-mel-24khz"
+elif mel_spec_type == "bigvgan":
+    vocoder_local_path = "../checkpoints/bigvgan_v2_24khz_100band_256x"
+vocoder = load_vocoder(vocoder_name=mel_spec_type, is_local=local, local_path=vocoder_local_path)
+
+# Tokenizer
+vocab_char_map, vocab_size = get_tokenizer(dataset_name, tokenizer)
+
+# Model
+model = CFM(
+    transformer=model_cls(**model_arc, text_num_embeds=vocab_size, mel_dim=n_mel_channels),
+    mel_spec_kwargs=dict(
+        n_fft=n_fft,
+        hop_length=hop_length,
+        win_length=win_length,
+        n_mel_channels=n_mel_channels,
+        target_sample_rate=target_sample_rate,
+        mel_spec_type=mel_spec_type,
+    ),
+    odeint_kwargs=dict(
+        method=ode_method,
+    ),
+    vocab_char_map=vocab_char_map,
+).to(device)
+
+dtype = torch.float32 if mel_spec_type == "bigvgan" else None
+model = load_checkpoint(model, ckpt_path, device, dtype=dtype, 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.to(torch.float32)
+    generated = generated[:, ref_audio_len:, :]
+    gen_mel_spec = generated.permute(0, 2, 1)
+    if mel_spec_type == "vocos":
+        generated_wave = vocoder.decode(gen_mel_spec).cpu()
+    elif mel_spec_type == "bigvgan":
+        generated_wave = vocoder(gen_mel_spec).squeeze(0).cpu()
+
+    if rms < target_rms:
+        generated_wave = generated_wave * rms / target_rms
+
+    save_spectrogram(gen_mel_spec[0].cpu().numpy(), f"{output_dir}/speech_edit_out.png")
+    torchaudio.save(f"{output_dir}/speech_edit_out.wav", generated_wave, target_sample_rate)
+    print(f"Generated wav: {generated_wave.shape}")

f5_tts/infer/utils_infer.py

@@ -0,0 +1,605 @@
+# A unified script for inference process
+# Make adjustments inside functions, and consider both gradio and cli scripts if need to change func output format
+import os
+import sys
+from concurrent.futures import ThreadPoolExecutor
+
+
+os.environ["PYTORCH_ENABLE_MPS_FALLBACK"] = "1"  # for MPS device compatibility
+sys.path.append(f"{os.path.dirname(os.path.abspath(__file__))}/../../third_party/BigVGAN/")
+
+import hashlib
+import re
+import tempfile
+from importlib.resources import files
+
+import matplotlib
+
+
+matplotlib.use("Agg")
+
+import matplotlib.pylab as plt
+import numpy as np
+import torch
+import torchaudio
+import tqdm
+from huggingface_hub import hf_hub_download
+from pydub import AudioSegment, silence
+from transformers import pipeline
+from vocos import Vocos
+
+from f5_tts.model import CFM
+from f5_tts.model.utils import convert_char_to_pinyin, get_tokenizer
+
+
+_ref_audio_cache = {}
+_ref_text_cache = {}
+
+device = (
+    "cuda"
+    if torch.cuda.is_available()
+    else "xpu"
+    if torch.xpu.is_available()
+    else "mps"
+    if torch.backends.mps.is_available()
+    else "cpu"
+)
+
+tempfile_kwargs = {"delete_on_close": False} if sys.version_info >= (3, 12) else {"delete": False}
+
+# -----------------------------------------
+
+target_sample_rate = 24000
+n_mel_channels = 100
+hop_length = 256
+win_length = 1024
+n_fft = 1024
+mel_spec_type = "vocos"
+target_rms = 0.1
+cross_fade_duration = 0.15
+ode_method = "euler"
+nfe_step = 32  # 16, 32
+cfg_strength = 2.0
+sway_sampling_coef = -1.0
+speed = 1.0
+fix_duration = None
+
+# -----------------------------------------
+
+
+# chunk text into smaller pieces
+
+
+def chunk_text(text, max_chars=135):
+    """
+    Splits the input text into chunks, each with a maximum number of characters.
+
+    Args:
+        text (str): The text to be split.
+        max_chars (int): The maximum number of characters per chunk.
+
+    Returns:
+        List[str]: A list of text chunks.
+    """
+    chunks = []
+    current_chunk = ""
+    # Split the text into sentences based on punctuation followed by whitespace
+    sentences = re.split(r"(?<=[;:,.!?])\s+|(?<=[；：，。！？])", text)
+
+    for sentence in sentences:
+        if len(current_chunk.encode("utf-8")) + len(sentence.encode("utf-8")) <= max_chars:
+            current_chunk += sentence + " " if sentence and len(sentence[-1].encode("utf-8")) == 1 else sentence
+        else:
+            if current_chunk:
+                chunks.append(current_chunk.strip())
+            current_chunk = sentence + " " if sentence and len(sentence[-1].encode("utf-8")) == 1 else sentence
+
+    if current_chunk:
+        chunks.append(current_chunk.strip())
+
+    return chunks
+
+
+# load vocoder
+def load_vocoder(vocoder_name="vocos", is_local=False, local_path="", device=device, hf_cache_dir=None):
+    if vocoder_name == "vocos":
+        # vocoder = Vocos.from_pretrained("charactr/vocos-mel-24khz").to(device)
+        if is_local:
+            print(f"Load vocos from local path {local_path}")
+            config_path = f"{local_path}/config.yaml"
+            model_path = f"{local_path}/pytorch_model.bin"
+        else:
+            print("Download Vocos from huggingface charactr/vocos-mel-24khz")
+            repo_id = "charactr/vocos-mel-24khz"
+            config_path = hf_hub_download(repo_id=repo_id, cache_dir=hf_cache_dir, filename="config.yaml")
+            model_path = hf_hub_download(repo_id=repo_id, cache_dir=hf_cache_dir, filename="pytorch_model.bin")
+        vocoder = Vocos.from_hparams(config_path)
+        state_dict = torch.load(model_path, map_location="cpu", weights_only=True)
+        from vocos.feature_extractors import EncodecFeatures
+
+        if isinstance(vocoder.feature_extractor, EncodecFeatures):
+            encodec_parameters = {
+                "feature_extractor.encodec." + key: value
+                for key, value in vocoder.feature_extractor.encodec.state_dict().items()
+            }
+            state_dict.update(encodec_parameters)
+        vocoder.load_state_dict(state_dict)
+        vocoder = vocoder.eval().to(device)
+    elif vocoder_name == "bigvgan":
+        try:
+            from third_party.BigVGAN import bigvgan
+        except ImportError:
+            print("You need to follow the README to init submodule and change the BigVGAN source code.")
+        if is_local:
+            # download generator from https://huggingface.co/nvidia/bigvgan_v2_24khz_100band_256x/tree/main
+            vocoder = bigvgan.BigVGAN.from_pretrained(local_path, use_cuda_kernel=False)
+        else:
+            vocoder = bigvgan.BigVGAN.from_pretrained(
+                "nvidia/bigvgan_v2_24khz_100band_256x", use_cuda_kernel=False, cache_dir=hf_cache_dir
+            )
+
+        vocoder.remove_weight_norm()
+        vocoder = vocoder.eval().to(device)
+    return vocoder
+
+
+# load asr pipeline
+
+asr_pipe = None
+
+
+def initialize_asr_pipeline(device: str = device, dtype=None):
+    if dtype is None:
+        dtype = (
+            torch.float16
+            if "cuda" in device
+            and torch.cuda.get_device_properties(device).major >= 7
+            and not torch.cuda.get_device_name().endswith("[ZLUDA]")
+            else torch.float32
+        )
+    global asr_pipe
+    asr_pipe = pipeline(
+        "automatic-speech-recognition",
+        model="openai/whisper-large-v3-turbo",
+        torch_dtype=dtype,
+        device=device,
+    )
+
+
+# transcribe
+
+
+def transcribe(ref_audio, language=None):
+    global asr_pipe
+    if asr_pipe is None:
+        initialize_asr_pipeline(device=device)
+    return asr_pipe(
+        ref_audio,
+        chunk_length_s=30,
+        batch_size=128,
+        generate_kwargs={"task": "transcribe", "language": language} if language else {"task": "transcribe"},
+        return_timestamps=False,
+    )["text"].strip()
+
+
+# load model checkpoint for inference
+
+
+def load_checkpoint(model, ckpt_path, device: str, dtype=None, use_ema=True):
+    if dtype is None:
+        dtype = (
+            torch.float16
+            if "cuda" in device
+            and torch.cuda.get_device_properties(device).major >= 7
+            and not torch.cuda.get_device_name().endswith("[ZLUDA]")
+            else torch.float32
+        )
+    model = model.to(dtype)
+
+    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, weights_only=True)
+
+    if use_ema:
+        if ckpt_type == "safetensors":
+            checkpoint = {"ema_model_state_dict": checkpoint}
+        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"]
+        }
+
+        # patch for backward compatibility, 305e3ea
+        for key in ["mel_spec.mel_stft.mel_scale.fb", "mel_spec.mel_stft.spectrogram.window"]:
+            if key in checkpoint["model_state_dict"]:
+                del checkpoint["model_state_dict"][key]
+
+        model.load_state_dict(checkpoint["model_state_dict"])
+    else:
+        if ckpt_type == "safetensors":
+            checkpoint = {"model_state_dict": checkpoint}
+        model.load_state_dict(checkpoint["model_state_dict"])
+
+    del checkpoint
+    torch.cuda.empty_cache()
+
+    return model.to(device)
+
+
+# load model for inference
+
+
+def load_model(
+    model_cls,
+    model_cfg,
+    ckpt_path,
+    mel_spec_type=mel_spec_type,
+    vocab_file="",
+    ode_method=ode_method,
+    use_ema=True,
+    device=device,
+):
+    if vocab_file == "":
+        vocab_file = str(files("f5_tts").joinpath("infer/examples/vocab.txt"))
+    tokenizer = "custom"
+
+    print("\nvocab : ", vocab_file)
+    print("token : ", tokenizer)
+    print("model : ", ckpt_path, "\n")
+
+    vocab_char_map, vocab_size = get_tokenizer(vocab_file, tokenizer)
+    model = CFM(
+        transformer=model_cls(**model_cfg, text_num_embeds=vocab_size, mel_dim=n_mel_channels),
+        mel_spec_kwargs=dict(
+            n_fft=n_fft,
+            hop_length=hop_length,
+            win_length=win_length,
+            n_mel_channels=n_mel_channels,
+            target_sample_rate=target_sample_rate,
+            mel_spec_type=mel_spec_type,
+        ),
+        odeint_kwargs=dict(
+            method=ode_method,
+        ),
+        vocab_char_map=vocab_char_map,
+    ).to(device)
+
+    dtype = torch.float32 if mel_spec_type == "bigvgan" else None
+    model = load_checkpoint(model, ckpt_path, device, dtype=dtype, use_ema=use_ema)
+
+    return model
+
+
+def remove_silence_edges(audio, silence_threshold=-42):
+    # Remove silence from the start
+    non_silent_start_idx = silence.detect_leading_silence(audio, silence_threshold=silence_threshold)
+    audio = audio[non_silent_start_idx:]
+
+    # Remove silence from the end
+    non_silent_end_duration = audio.duration_seconds
+    for ms in reversed(audio):
+        if ms.dBFS > silence_threshold:
+            break
+        non_silent_end_duration -= 0.001
+    trimmed_audio = audio[: int(non_silent_end_duration * 1000)]
+
+    return trimmed_audio
+
+
+# preprocess reference audio and text
+
+
+def preprocess_ref_audio_text(ref_audio_orig, ref_text, show_info=print):
+    show_info("Converting audio...")
+
+    # Compute a hash of the reference audio file
+    with open(ref_audio_orig, "rb") as audio_file:
+        audio_data = audio_file.read()
+        audio_hash = hashlib.md5(audio_data).hexdigest()
+
+    global _ref_audio_cache
+
+    if audio_hash in _ref_audio_cache:
+        show_info("Using cached preprocessed reference audio...")
+        ref_audio = _ref_audio_cache[audio_hash]
+
+    else:  # first pass, do preprocess
+        with tempfile.NamedTemporaryFile(suffix=".wav", **tempfile_kwargs) as f:
+            temp_path = f.name
+
+        aseg = AudioSegment.from_file(ref_audio_orig)
+
+        # 1. try to find long silence for clipping
+        non_silent_segs = silence.split_on_silence(
+            aseg, min_silence_len=1000, silence_thresh=-50, keep_silence=1000, seek_step=10
+        )
+        non_silent_wave = AudioSegment.silent(duration=0)
+        for non_silent_seg in non_silent_segs:
+            if len(non_silent_wave) > 6000 and len(non_silent_wave + non_silent_seg) > 12000:
+                show_info("Audio is over 12s, clipping short. (1)")
+                break
+            non_silent_wave += non_silent_seg
+
+        # 2. try to find short silence for clipping if 1. failed
+        if len(non_silent_wave) > 12000:
+            non_silent_segs = silence.split_on_silence(
+                aseg, min_silence_len=100, silence_thresh=-40, keep_silence=1000, seek_step=10
+            )
+            non_silent_wave = AudioSegment.silent(duration=0)
+            for non_silent_seg in non_silent_segs:
+                if len(non_silent_wave) > 6000 and len(non_silent_wave + non_silent_seg) > 12000:
+                    show_info("Audio is over 12s, clipping short. (2)")
+                    break
+                non_silent_wave += non_silent_seg
+
+        aseg = non_silent_wave
+
+        # 3. if no proper silence found for clipping
+        if len(aseg) > 12000:
+            aseg = aseg[:12000]
+            show_info("Audio is over 12s, clipping short. (3)")
+
+        aseg = remove_silence_edges(aseg) + AudioSegment.silent(duration=50)
+        aseg.export(temp_path, format="wav")
+        ref_audio = temp_path
+
+        # Cache the processed reference audio
+        _ref_audio_cache[audio_hash] = ref_audio
+
+    if not ref_text.strip():
+        global _ref_text_cache
+        if audio_hash in _ref_text_cache:
+            # Use cached asr transcription
+            show_info("Using cached reference text...")
+            ref_text = _ref_text_cache[audio_hash]
+        else:
+            show_info("No reference text provided, transcribing reference audio...")
+            ref_text = transcribe(ref_audio)
+            # Cache the transcribed text (not caching custom ref_text, enabling users to do manual tweak)
+            _ref_text_cache[audio_hash] = ref_text
+    else:
+        show_info("Using custom reference text...")
+
+    # Ensure ref_text ends with a proper sentence-ending punctuation
+    if not ref_text.endswith(". ") and not ref_text.endswith("。"):
+        if ref_text.endswith("."):
+            ref_text += " "
+        else:
+            ref_text += ". "
+
+    print("\nref_text  ", ref_text)
+
+    return ref_audio, ref_text
+
+
+# infer process: chunk text -> infer batches [i.e. infer_batch_process()]
+
+
+def infer_process(
+    ref_audio,
+    ref_text,
+    gen_text,
+    model_obj,
+    vocoder,
+    mel_spec_type=mel_spec_type,
+    show_info=print,
+    progress=tqdm,
+    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,
+    device=device,
+):
+    # 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) * (22 - audio.shape[-1] / sr) * speed)
+    gen_text_batches = chunk_text(gen_text, max_chars=max_chars)
+    for i, gen_text in enumerate(gen_text_batches):
+        print(f"gen_text {i}", gen_text)
+    print("\n")
+
+    show_info(f"Generating audio in {len(gen_text_batches)} batches...")
+    return next(
+        infer_batch_process(
+            (audio, sr),
+            ref_text,
+            gen_text_batches,
+            model_obj,
+            vocoder,
+            mel_spec_type=mel_spec_type,
+            progress=progress,
+            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,
+            device=device,
+        )
+    )
+
+
+# infer batches
+
+
+def infer_batch_process(
+    ref_audio,
+    ref_text,
+    gen_text_batches,
+    model_obj,
+    vocoder,
+    mel_spec_type="vocos",
+    progress=tqdm,
+    target_rms=0.1,
+    cross_fade_duration=0.15,
+    nfe_step=32,
+    cfg_strength=2.0,
+    sway_sampling_coef=-1,
+    speed=1,
+    fix_duration=None,
+    device=None,
+    streaming=False,
+    chunk_size=2048,
+):
+    audio, sr = ref_audio
+    if audio.shape[0] > 1:
+        audio = torch.mean(audio, dim=0, keepdim=True)
+
+    rms = torch.sqrt(torch.mean(torch.square(audio)))
+    if rms < target_rms:
+        audio = audio * target_rms / rms
+    if sr != target_sample_rate:
+        resampler = torchaudio.transforms.Resample(sr, target_sample_rate)
+        audio = resampler(audio)
+    audio = audio.to(device)
+
+    generated_waves = []
+    spectrograms = []
+
+    if len(ref_text[-1].encode("utf-8")) == 1:
+        ref_text = ref_text + " "
+
+    def process_batch(gen_text):
+        local_speed = speed
+        if len(gen_text.encode("utf-8")) < 10:
+            local_speed = 0.3
+
+        # Prepare the text
+        text_list = [ref_text + gen_text]
+        final_text_list = convert_char_to_pinyin(text_list)
+
+        ref_audio_len = audio.shape[-1] // hop_length
+        if fix_duration is not None:
+            duration = int(fix_duration * target_sample_rate / hop_length)
+        else:
+            # Calculate duration
+            ref_text_len = len(ref_text.encode("utf-8"))
+            gen_text_len = len(gen_text.encode("utf-8"))
+            duration = ref_audio_len + int(ref_audio_len / ref_text_len * gen_text_len / local_speed)
+
+        # inference
+        with torch.inference_mode():
+            generated, _ = model_obj.sample(
+                cond=audio,
+                text=final_text_list,
+                duration=duration,
+                steps=nfe_step,
+                cfg_strength=cfg_strength,
+                sway_sampling_coef=sway_sampling_coef,
+            )
+            del _
+
+            generated = generated.to(torch.float32)  # generated mel spectrogram
+            generated = generated[:, ref_audio_len:, :]
+            generated = generated.permute(0, 2, 1)
+            if mel_spec_type == "vocos":
+                generated_wave = vocoder.decode(generated)
+            elif mel_spec_type == "bigvgan":
+                generated_wave = vocoder(generated)
+            if rms < target_rms:
+                generated_wave = generated_wave * rms / target_rms
+
+            # wav -> numpy
+            generated_wave = generated_wave.squeeze().cpu().numpy()
+
+            if streaming:
+                for j in range(0, len(generated_wave), chunk_size):
+                    yield generated_wave[j : j + chunk_size], target_sample_rate
+            else:
+                generated_cpu = generated[0].cpu().numpy()
+                del generated
+                yield generated_wave, generated_cpu
+
+    if streaming:
+        for gen_text in progress.tqdm(gen_text_batches) if progress is not None else gen_text_batches:
+            for chunk in process_batch(gen_text):
+                yield chunk
+    else:
+        with ThreadPoolExecutor() as executor:
+            futures = [executor.submit(process_batch, gen_text) for gen_text in gen_text_batches]
+            for future in progress.tqdm(futures) if progress is not None else futures:
+                result = future.result()
+                if result:
+                    generated_wave, generated_mel_spec = next(result)
+                    generated_waves.append(generated_wave)
+                    spectrograms.append(generated_mel_spec)
+
+        if generated_waves:
+            if cross_fade_duration <= 0:
+                # Simply concatenate
+                final_wave = np.concatenate(generated_waves)
+            else:
+                # Combine all generated waves with cross-fading
+                final_wave = generated_waves[0]
+                for i in range(1, len(generated_waves)):
+                    prev_wave = final_wave
+                    next_wave = generated_waves[i]
+
+                    # Calculate cross-fade samples, ensuring it does not exceed wave lengths
+                    cross_fade_samples = int(cross_fade_duration * target_sample_rate)
+                    cross_fade_samples = min(cross_fade_samples, len(prev_wave), len(next_wave))
+
+                    if cross_fade_samples <= 0:
+                        # No overlap possible, concatenate
+                        final_wave = np.concatenate([prev_wave, next_wave])
+                        continue
+
+                    # Overlapping parts
+                    prev_overlap = prev_wave[-cross_fade_samples:]
+                    next_overlap = next_wave[:cross_fade_samples]
+
+                    # Fade out and fade in
+                    fade_out = np.linspace(1, 0, cross_fade_samples)
+                    fade_in = np.linspace(0, 1, cross_fade_samples)
+
+                    # Cross-faded overlap
+                    cross_faded_overlap = prev_overlap * fade_out + next_overlap * fade_in
+
+                    # Combine
+                    new_wave = np.concatenate(
+                        [prev_wave[:-cross_fade_samples], cross_faded_overlap, next_wave[cross_fade_samples:]]
+                    )
+
+                    final_wave = new_wave
+
+            # Create a combined spectrogram
+            combined_spectrogram = np.concatenate(spectrograms, axis=1)
+
+            yield final_wave, target_sample_rate, combined_spectrogram
+
+        else:
+            yield None, target_sample_rate, None
+
+
+# remove silence from generated wav
+
+
+def remove_silence_for_generated_wav(filename):
+    aseg = AudioSegment.from_file(filename)
+    non_silent_segs = silence.split_on_silence(
+        aseg, min_silence_len=1000, silence_thresh=-50, keep_silence=500, seek_step=10
+    )
+    non_silent_wave = AudioSegment.silent(duration=0)
+    for non_silent_seg in non_silent_segs:
+        non_silent_wave += non_silent_seg
+    aseg = non_silent_wave
+    aseg.export(filename, format="wav")
+
+
+# save spectrogram
+
+
+def save_spectrogram(spectrogram, path):
+    plt.figure(figsize=(12, 4))
+    plt.imshow(spectrogram, origin="lower", aspect="auto")
+    plt.colorbar()
+    plt.savefig(path)
+    plt.close()

f5_tts/model/__init__.py

@@ -0,0 +1,10 @@
+from f5_tts.model.cfm import CFM
+
+from f5_tts.model.backbones.unett import UNetT
+from f5_tts.model.backbones.dit import DiT
+from f5_tts.model.backbones.mmdit import MMDiT
+
+from f5_tts.model.trainer import Trainer
+
+
+__all__ = ["CFM", "UNetT", "DiT", "MMDiT", "Trainer"]

f5_tts/model/backbones/README.md

@@ -0,0 +1,20 @@
+## Backbones quick introduction
+
+
+### unett.py
+- flat unet transformer
+- structure same as in e2-tts & voicebox paper except using rotary pos emb
+- update: allow possible abs pos emb & convnextv2 blocks for embedded text before concat
+
+### dit.py
+- adaln-zero dit
+- embedded timestep as condition
+- concatted noised_input + masked_cond + embedded_text, linear proj in
+- possible abs pos emb & convnextv2 blocks for embedded text before concat
+- possible long skip connection (first layer to last layer)
+
+### mmdit.py
+- sd3 structure
+- timestep as condition
+- left stream: text embedded and applied a abs pos emb
+- right stream: masked_cond & noised_input concatted and with same conv pos emb as unett

f5_tts/model/backbones/dit.py

@@ -0,0 +1,204 @@
+"""
+ein notation:
+b - batch
+n - sequence
+nt - text sequence
+nw - raw wave length
+d - dimension
+"""
+
+from __future__ import annotations
+
+import torch
+from torch import nn
+import torch.nn.functional as F
+
+from x_transformers.x_transformers import RotaryEmbedding
+
+from f5_tts.model.modules import (
+    TimestepEmbedding,
+    ConvNeXtV2Block,
+    ConvPositionEmbedding,
+    DiTBlock,
+    AdaLayerNormZero_Final,
+    precompute_freqs_cis,
+    get_pos_embed_indices,
+)
+
+
+# Text embedding
+
+
+class TextEmbedding(nn.Module):
+    def __init__(self, text_num_embeds, text_dim, conv_layers=0, conv_mult=2):
+        super().__init__()
+        self.text_embed = nn.Embedding(text_num_embeds + 1, text_dim)  # use 0 as filler token
+
+        if conv_layers > 0:
+            self.extra_modeling = True
+            self.precompute_max_pos = 4096  # ~44s of 24khz audio
+            self.register_buffer("freqs_cis", precompute_freqs_cis(text_dim, self.precompute_max_pos), persistent=False)
+            self.text_blocks = nn.Sequential(
+                *[ConvNeXtV2Block(text_dim, text_dim * conv_mult) for _ in range(conv_layers)]
+            )
+        else:
+            self.extra_modeling = False
+
+    def forward(self, text: int["b nt"], seq_len, drop_text=False):  # noqa: F722
+        text = text + 1  # use 0 as filler token. preprocess of batch pad -1, see list_str_to_idx()
+        text = text[:, :seq_len]  # curtail if character tokens are more than the mel spec tokens
+        batch, text_len = text.shape[0], text.shape[1]
+        text = F.pad(text, (0, seq_len - text_len), value=0)
+
+        if drop_text:  # cfg for text
+            text = torch.zeros_like(text)
+        
+        text = self.text_embed(text)  # b n -> b n d
+
+        # possible extra modeling
+        if self.extra_modeling:
+            # sinus pos emb
+            batch_start = torch.zeros((batch,), dtype=torch.long)
+            pos_idx = get_pos_embed_indices(batch_start, seq_len, max_pos=self.precompute_max_pos)
+            text_pos_embed = self.freqs_cis[pos_idx]
+            text = text + text_pos_embed
+
+            # convnextv2 blocks
+            text = self.text_blocks(text)
+
+        return text
+
+
+# noised input audio and context mixing embedding
+
+
+class InputEmbedding(nn.Module):
+    def __init__(self, mel_dim, text_dim, out_dim):
+        super().__init__()
+        self.proj = nn.Linear(mel_dim * 2 + text_dim, out_dim)
+        self.conv_pos_embed = ConvPositionEmbedding(dim=out_dim)
+
+    def forward(self, x: float["b n d"], cond: float["b n d"], text_embed: float["b n d"], drop_audio_cond=False):  # noqa: F722
+        if drop_audio_cond:  # cfg for cond audio
+            cond = torch.zeros_like(cond)
+
+        x = self.proj(torch.cat((x, cond, text_embed), dim=-1))
+        x = self.conv_pos_embed(x) + x
+        return x
+
+
+# Transformer backbone using DiT blocks
+
+
+class DiT(nn.Module):
+    def __init__(
+        self,
+        *,
+        dim,
+        depth=8,
+        heads=8,
+        dim_head=64,
+        dropout=0.1,
+        ff_mult=4,
+        mel_dim=100,
+        text_num_embeds=256,
+        text_dim=None,
+        conv_layers=0,
+        long_skip_connection=False,
+        checkpoint_activations=False,
+        second_time=False,
+    ):
+        super().__init__()
+
+        self.time_embed = TimestepEmbedding(dim)
+        if second_time:
+            self.time_embed2 = TimestepEmbedding(dim)
+            # Zero-init the weights and biases of the first and last Linear layers in time_mlp
+            nn.init.zeros_(self.time_embed2.time_mlp[0].weight)  # First Linear layer weights
+            nn.init.zeros_(self.time_embed2.time_mlp[0].bias)    # First Linear layer bias
+            nn.init.zeros_(self.time_embed2.time_mlp[-1].weight) # Last Linear layer weights
+            nn.init.zeros_(self.time_embed2.time_mlp[-1].bias)   # Last Linear layer bias
+        else:
+            self.time_embed2 = None
+        
+        if text_dim is None:
+            text_dim = mel_dim
+        self.vocab_size = text_num_embeds
+        self.text_embed = TextEmbedding(text_num_embeds, text_dim, conv_layers=conv_layers)
+        self.input_embed = InputEmbedding(mel_dim, text_dim, dim)
+
+        self.rotary_embed = RotaryEmbedding(dim_head)
+
+        self.dim = dim
+        self.depth = depth
+
+        self.transformer_blocks = nn.ModuleList(
+            [DiTBlock(dim=dim, heads=heads, dim_head=dim_head, ff_mult=ff_mult, dropout=dropout) for _ in range(depth)]
+        )
+        self.long_skip_connection = nn.Linear(dim * 2, dim, bias=False) if long_skip_connection else None
+
+        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
+        cond: float["b n d"],  # masked cond audio  # noqa: F722
+        text: int["b nt"],  # text  # noqa: F722
+        time: float["b"] | float[""],  # time step  # noqa: F821 F722
+        drop_audio_cond,  # cfg for cond audio
+        drop_text,  # cfg for text
+        mask: bool["b n"] | None = None,  # noqa: F722
+        second_time: float["b"] | float[""] = None,  # noqa: F821 F722
+        classify_mode: bool = False,  # noqa: F821
+    ):
+        batch, seq_len = x.shape[0], x.shape[1]
+        if time.ndim == 0:
+            time = time.repeat(batch)
+
+        # t: conditioning time, c: context (text + masked cond audio), x: noised input audio
+        t = self.time_embed(time)
+        if second_time is not None and self.time_embed2 is not None:
+            t2 = self.time_embed2(second_time)
+            t = t + t2
+        
+        text_embed = self.text_embed(text, seq_len, drop_text=drop_text)
+        x = self.input_embed(x, cond, text_embed, drop_audio_cond=drop_audio_cond)
+
+        rope = self.rotary_embed.forward_from_seq_len(seq_len)
+
+        if self.long_skip_connection is not None:
+            residual = x
+
+        if classify_mode:
+            layers = [x]
+
+        for block in self.transformer_blocks:
+            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 classify_mode:
+                layers.append(x)
+
+        if self.long_skip_connection is not None:
+            x = self.long_skip_connection(torch.cat((x, residual), dim=-1))
+
+        x = self.norm_out(x, t)
+        output = self.proj_out(x)
+
+        if classify_mode:
+            return layers
+
+        return output

f5_tts/model/backbones/mmdit.py

@@ -0,0 +1,336 @@
+"""
+ein notation:
+b - batch
+n - sequence
+nt - text sequence
+nw - raw wave length
+d - dimension
+"""
+
+from __future__ import annotations
+
+import torch
+from torch import nn
+import torch.nn.functional as F
+
+from x_transformers.x_transformers import RotaryEmbedding
+
+from f5_tts.model.modules import (
+    TimestepEmbedding,
+    ConvPositionEmbedding,
+    MMDiTBlock,
+    DiTBlock,
+    AdaLayerNormZero_Final,
+    precompute_freqs_cis,
+    get_pos_embed_indices,
+)
+
+from f5_tts.model.utils import (
+    default,
+    exists,
+    lens_to_mask,
+    list_str_to_idx,
+    list_str_to_tensor,
+    mask_from_frac_lengths,
+)
+# text embedding
+
+
+class TextEmbedding(nn.Module):
+    def __init__(self, out_dim, text_num_embeds):
+        super().__init__()
+        self.text_embed = nn.Embedding(text_num_embeds + 1, out_dim)  # will use 0 as filler token
+
+        self.precompute_max_pos = 1024
+        self.register_buffer("freqs_cis", precompute_freqs_cis(out_dim, self.precompute_max_pos), persistent=False)
+
+    def forward(self, text: int["b nt"], drop_text=False) -> int["b nt d"]:  # noqa: F722
+        text = text + 1
+        if drop_text:
+            text = torch.zeros_like(text)
+        text = self.text_embed(text)
+
+        # sinus pos emb
+        batch_start = torch.zeros((text.shape[0],), dtype=torch.long)
+        batch_text_len = text.shape[1]
+        pos_idx = get_pos_embed_indices(batch_start, batch_text_len, max_pos=self.precompute_max_pos)
+        text_pos_embed = self.freqs_cis[pos_idx]
+
+        text = text + text_pos_embed
+
+        return text
+
+
+# noised input & masked cond audio embedding
+
+
+class AudioEmbedding(nn.Module):
+    def __init__(self, in_dim, out_dim):
+        super().__init__()
+        self.linear = nn.Linear(2 * in_dim, out_dim)
+        self.conv_pos_embed = ConvPositionEmbedding(out_dim)
+
+    def forward(self, x: float["b n d"], cond: float["b n d"], drop_audio_cond=False):  # noqa: F722
+        if drop_audio_cond:
+            cond = torch.zeros_like(cond)
+        x = torch.cat((x, cond), dim=-1)
+        x = self.linear(x)
+        x = self.conv_pos_embed(x) + x
+        return x
+
+
+# Transformer backbone using MM-DiT blocks
+
+
+class MMDiT(nn.Module):
+    def __init__(
+        self,
+        *,
+        dim,
+        text_depth=4,
+        depth=8,
+        heads=8,
+        dim_head=64,
+        dropout=0.1,
+        ff_mult=4,
+        text_num_embeds=256,
+        mel_dim=100,
+        checkpoint_activations=False,
+        text_encoder=True,
+
+    ):
+        super().__init__()
+
+        self.time_embed = TimestepEmbedding(dim)
+        if text_encoder:
+            self.text_encoder = TextEncoder(text_num_embeds=text_num_embeds, 
+                                        text_dim=dim,
+                                        depth=text_depth,
+                                        heads=heads,
+                                        dim_head=dim_head,
+                                        ff_mult=ff_mult,
+                                        dropout=dropout)
+        else:
+            self.text_encoder = None
+            self.text_embed = TextEmbedding(dim, text_num_embeds)
+        
+        self.audio_embed = AudioEmbedding(mel_dim, dim)
+
+        self.rotary_embed = RotaryEmbedding(dim_head)
+
+        self.dim = dim
+        self.depth = depth
+
+        self.transformer_blocks = nn.ModuleList(
+            [
+                MMDiTBlock(
+                    dim=dim,
+                    heads=heads,
+                    dim_head=dim_head,
+                    dropout=dropout,
+                    ff_mult=ff_mult,
+                    context_pre_only=i == depth - 1,
+                )
+                for i in range(depth)
+            ]
+        )
+        self.norm_out = AdaLayerNormZero_Final(dim)  # final modulation
+        self.proj_out = nn.Linear(dim, mel_dim)
+        
+        self.checkpoint_activations = checkpoint_activations
+
+
+    def forward(
+        self,
+        x: float["b n d"],  # nosied input audio  # noqa: F722
+        cond: float["b n d"],  # masked cond audio  # noqa: F722
+        text: int["b nt"],  # text  # noqa: F722
+        time: float["b"] | float[""],  # time step  # noqa: F821 F722
+        drop_audio_cond,  # cfg for cond audio
+        drop_text,  # cfg for text
+        mask: bool["b n"] | None = None,  # noqa: F722
+        text_mask: bool["b nt"] | None = None,  # noqa: F722
+    ):
+        batch = x.shape[0]
+        if time.ndim == 0:
+            time = time.repeat(batch)
+
+        # t: conditioning (time), c: context (text + masked cond audio), x: noised input audio
+        t = self.time_embed(time)
+        if self.text_encoder is not None:
+            c = self.text_encoder(text, t, mask=text_mask, drop_text=drop_text)
+        else:
+            c = self.text_embed(text, drop_text=drop_text)
+        
+        x = self.audio_embed(x, cond, drop_audio_cond=drop_audio_cond)
+
+        seq_len = x.shape[1]
+        text_len = text.shape[1]
+        rope_audio = self.rotary_embed.forward_from_seq_len(seq_len)
+        rope_text = self.rotary_embed.forward_from_seq_len(text_len)
+        
+        # if mask is not None:
+        #     rope_audio = self.rotary_embed.forward_from_seq_len(seq_len + 1)
+            
+        #     dummy_token = torch.zeros((x.shape[0], 1, x.shape[-1]), device=x.device, dtype=x.dtype)
+        #     x = torch.cat([x, dummy_token], dim=1)  # shape is now [b, nw+1, d]
+            
+        #     # pad the mask so that new dummy token is always masked out
+        #     # mask: [b, nw] -> [b, nw+1]
+        #     false_col = torch.zeros((x.shape[0], 1), dtype=torch.bool, device=x.device)
+        #     mask = torch.cat([mask, false_col], dim=1)
+                
+        # if text_mask is not None:
+        #     rope_text = self.rotary_embed.forward_from_seq_len(text_len + 1)
+
+        #     dummy_token = torch.zeros((c.shape[0], 1, c.shape[-1]), device=c.device, dtype=c.dtype)
+        #     c = torch.cat([c, dummy_token], dim=1)  # shape is now [b, nt+1, d]
+        
+        #     # pad the text mask so that new dummy token is always masked out
+        #     # text_mask: [b, nt] -> [b, nt+1]
+        #     false_col = torch.zeros((c.shape[0], 1), dtype=torch.bool, device=c.device)
+        #     text_mask = torch.cat([text_mask, false_col], dim=1)
+                        
+        for block in self.transformer_blocks:
+            c, x = block(x, c, t, mask=mask, src_mask=text_mask, rope=rope_audio, c_rope=rope_text)
+
+        x = self.norm_out(x, t)
+        output = self.proj_out(x)
+        
+
+        return output
+
+class TextEncoder(nn.Module):
+    def __init__(
+        self,
+        text_num_embeds: int,
+        text_dim: int = 512,
+        depth: int = 4,
+        heads: int = 8,
+        dim_head: int = 64,
+        ff_mult: int = 4,
+        dropout: float = 0.1,
+    ):
+        """
+        A simple text encoder: an embedding layer + multiple DiTBlocks or any other
+        transformer blocks for text-only self-attention.
+        """
+        super().__init__()
+        # Embeddings
+        self.text_embed = TextEmbedding(text_dim, text_num_embeds)
+        self.rotary_embed = RotaryEmbedding(dim_head)
+        
+        # Example stack of DiTBlocks or any custom blocks
+        self.transformer_blocks = nn.ModuleList(
+            [
+                DiTBlock(
+                    dim=text_dim,
+                    heads=heads,
+                    dim_head=dim_head,
+                    ff_mult=ff_mult,
+                    dropout=dropout,
+                )
+                for _ in range(depth)
+            ]
+        )
+
+    def forward(
+        self,
+        text: int["b nt"],  # noqa: F821
+        time: float["b"] | float[""],  # time step  # noqa: F821 F722
+        mask: bool["b nt"] | None = None,  # noqa: F821 F722
+        drop_text: bool = False
+    ):
+        """
+        Encode text into hidden states of shape [b, nt, d].
+        """
+        batch, seq_len, device = text.shape[0], text.shape[1], text.device
+
+        if drop_text:
+            text = torch.zeros_like(text)
+
+        # Basic embedding
+        hidden_states = self.text_embed(text, seq_len)  # [b, nt, d]
+        
+        # lens and mask
+        rope = self.rotary_embed.forward_from_seq_len(seq_len)
+
+        # Pass through self-attention blocks
+        for block in self.transformer_blocks:
+            # Here, you likely want standard self-attn, so no cross-attn
+            hidden_states = block(
+                x=hidden_states,
+                t=time,       # no time embedding for the text encoder by default
+                mask=mask,    # or pass a text mask if needed
+                rope=rope     # pass a rope if you want rotary embeddings for text
+            )
+        return hidden_states
+
+if __name__ == "__main__":
+    from f5_tts.model.utils import get_tokenizer
+
+    bsz = 16
+    
+    tokenizer = "pinyin"  # 'pinyin', 'char', or 'custom'
+    tokenizer_path = None  # if tokenizer = 'custom', define the path to the tokenizer you want to use (should be vocab.txt)
+    dataset_name = "Emilia_ZH_EN"
+    if tokenizer == "custom":
+        tokenizer_path = tokenizer_path
+    else:
+        tokenizer_path = dataset_name
+    vocab_char_map, vocab_size = get_tokenizer(tokenizer_path, tokenizer)
+    
+    text = ["hello world"] * bsz
+    text_lens = torch.ones((bsz, ), dtype=torch.long) * len("hello world")
+    text_lens[-1] = 5
+    device = "cuda"
+    batch = bsz
+    time_embed = TimestepEmbedding(512).to(device)
+    
+    
+    # handle text as string
+    if isinstance(text, list):
+        if exists(vocab_char_map):
+            text = list_str_to_idx(text, vocab_char_map).to(device)
+        else:
+            text = list_str_to_tensor(text).to(device)
+        assert text.shape[0] == batch  
+        
+    time = torch.rand((batch,), device=device)
+    text_mask = lens_to_mask(text_lens).to(device)
+
+    # # test text encoder
+    # text_encoder = TextEncoder(
+    #     text_num_embeds=vocab_size,
+    #     text_dim=512,
+    #     depth=4,
+    #     heads=8,
+    #     dim_head=64,
+    #     ff_mult=4,
+    #     dropout=0.1
+    # ).to('cuda')
+    # hidden_states = text_encoder(text, time_embed(time), mask)
+    # print(hidden_states.shape)  # [bsz, seq_len, text_dim]
+    
+    # test MMDiT
+    mel_dim = 80
+    model = MMDiT(
+        dim=512,
+        text_depth=4,
+        depth=8,
+        heads=8,
+        dim_head=64,
+        dropout=0.1,
+        ff_mult=4,
+        text_num_embeds=vocab_size,
+        mel_dim=mel_dim
+    ).to(device)
+    
+    x = torch.rand((batch, 100, mel_dim), device=device)
+    cond = torch.rand((batch, 100, mel_dim), device=device)
+    lens = torch.ones((batch,), dtype=torch.long) * 100
+    mask = lens_to_mask(lens).to(device)
+    
+    output = model(x, cond, text, time, drop_audio_cond=False, drop_text=False, mask=mask, text_mask=text_mask)
+    
+    print(output.shape)  # [bsz, seq_len, mel_dim]

f5_tts/model/backbones/unett.py

@@ -0,0 +1,219 @@
+"""
+ein notation:
+b - batch
+n - sequence
+nt - text sequence
+nw - raw wave length
+d - dimension
+"""
+
+from __future__ import annotations
+from typing import Literal
+
+import torch
+from torch import nn
+import torch.nn.functional as F
+
+from x_transformers import RMSNorm
+from x_transformers.x_transformers import RotaryEmbedding
+
+from f5_tts.model.modules import (
+    TimestepEmbedding,
+    ConvNeXtV2Block,
+    ConvPositionEmbedding,
+    Attention,
+    AttnProcessor,
+    FeedForward,
+    precompute_freqs_cis,
+    get_pos_embed_indices,
+)
+
+
+# Text embedding
+
+
+class TextEmbedding(nn.Module):
+    def __init__(self, text_num_embeds, text_dim, conv_layers=0, conv_mult=2):
+        super().__init__()
+        self.text_embed = nn.Embedding(text_num_embeds + 1, text_dim)  # use 0 as filler token
+
+        if conv_layers > 0:
+            self.extra_modeling = True
+            self.precompute_max_pos = 4096  # ~44s of 24khz audio
+            self.register_buffer("freqs_cis", precompute_freqs_cis(text_dim, self.precompute_max_pos), persistent=False)
+            self.text_blocks = nn.Sequential(
+                *[ConvNeXtV2Block(text_dim, text_dim * conv_mult) for _ in range(conv_layers)]
+            )
+        else:
+            self.extra_modeling = False
+
+    def forward(self, text: int["b nt"], seq_len, drop_text=False):  # noqa: F722
+        text = text + 1  # use 0 as filler token. preprocess of batch pad -1, see list_str_to_idx()
+        text = text[:, :seq_len]  # curtail if character tokens are more than the mel spec tokens
+        batch, text_len = text.shape[0], text.shape[1]
+        text = F.pad(text, (0, seq_len - text_len), value=0)
+
+        if drop_text:  # cfg for text
+            text = torch.zeros_like(text)
+
+        text = self.text_embed(text)  # b n -> b n d
+
+        # possible extra modeling
+        if self.extra_modeling:
+            # sinus pos emb
+            batch_start = torch.zeros((batch,), dtype=torch.long)
+            pos_idx = get_pos_embed_indices(batch_start, seq_len, max_pos=self.precompute_max_pos)
+            text_pos_embed = self.freqs_cis[pos_idx]
+            text = text + text_pos_embed
+
+            # convnextv2 blocks
+            text = self.text_blocks(text)
+
+        return text
+
+
+# noised input audio and context mixing embedding
+
+
+class InputEmbedding(nn.Module):
+    def __init__(self, mel_dim, text_dim, out_dim):
+        super().__init__()
+        self.proj = nn.Linear(mel_dim * 2 + text_dim, out_dim)
+        self.conv_pos_embed = ConvPositionEmbedding(dim=out_dim)
+
+    def forward(self, x: float["b n d"], cond: float["b n d"], text_embed: float["b n d"], drop_audio_cond=False):  # noqa: F722
+        if drop_audio_cond:  # cfg for cond audio
+            cond = torch.zeros_like(cond)
+
+        x = self.proj(torch.cat((x, cond, text_embed), dim=-1))
+        x = self.conv_pos_embed(x) + x
+        return x
+
+
+# Flat UNet Transformer backbone
+
+
+class UNetT(nn.Module):
+    def __init__(
+        self,
+        *,
+        dim,
+        depth=8,
+        heads=8,
+        dim_head=64,
+        dropout=0.1,
+        ff_mult=4,
+        mel_dim=100,
+        text_num_embeds=256,
+        text_dim=None,
+        conv_layers=0,
+        skip_connect_type: Literal["add", "concat", "none"] = "concat",
+    ):
+        super().__init__()
+        assert depth % 2 == 0, "UNet-Transformer's depth should be even."
+
+        self.time_embed = TimestepEmbedding(dim)
+        if text_dim is None:
+            text_dim = mel_dim
+        self.text_embed = TextEmbedding(text_num_embeds, text_dim, conv_layers=conv_layers)
+        self.input_embed = InputEmbedding(mel_dim, text_dim, dim)
+
+        self.rotary_embed = RotaryEmbedding(dim_head)
+
+        # transformer layers & skip connections
+
+        self.dim = dim
+        self.skip_connect_type = skip_connect_type
+        needs_skip_proj = skip_connect_type == "concat"
+
+        self.depth = depth
+        self.layers = nn.ModuleList([])
+
+        for idx in range(depth):
+            is_later_half = idx >= (depth // 2)
+
+            attn_norm = RMSNorm(dim)
+            attn = Attention(
+                processor=AttnProcessor(),
+                dim=dim,
+                heads=heads,
+                dim_head=dim_head,
+                dropout=dropout,
+            )
+
+            ff_norm = RMSNorm(dim)
+            ff = FeedForward(dim=dim, mult=ff_mult, dropout=dropout, approximate="tanh")
+
+            skip_proj = nn.Linear(dim * 2, dim, bias=False) if needs_skip_proj and is_later_half else None
+
+            self.layers.append(
+                nn.ModuleList(
+                    [
+                        skip_proj,
+                        attn_norm,
+                        attn,
+                        ff_norm,
+                        ff,
+                    ]
+                )
+            )
+
+        self.norm_out = RMSNorm(dim)
+        self.proj_out = nn.Linear(dim, mel_dim)
+
+    def forward(
+        self,
+        x: float["b n d"],  # nosied input audio  # noqa: F722
+        cond: float["b n d"],  # masked cond audio  # noqa: F722
+        text: int["b nt"],  # text  # noqa: F722
+        time: float["b"] | float[""],  # time step  # noqa: F821 F722
+        drop_audio_cond,  # cfg for cond audio
+        drop_text,  # cfg for text
+        mask: bool["b n"] | None = None,  # noqa: F722
+    ):
+        batch, seq_len = x.shape[0], x.shape[1]
+        if time.ndim == 0:
+            time = time.repeat(batch)
+
+        # t: conditioning time, c: context (text + masked cond audio), x: noised input audio
+        t = self.time_embed(time)
+        text_embed = self.text_embed(text, seq_len, drop_text=drop_text)
+        x = self.input_embed(x, cond, text_embed, drop_audio_cond=drop_audio_cond)
+
+        # postfix time t to input x, [b n d] -> [b n+1 d]
+        x = torch.cat([t.unsqueeze(1), x], dim=1)  # pack t to x
+        if mask is not None:
+            mask = F.pad(mask, (1, 0), value=1)
+
+        rope = self.rotary_embed.forward_from_seq_len(seq_len + 1)
+
+        # flat unet transformer
+        skip_connect_type = self.skip_connect_type
+        skips = []
+        for idx, (maybe_skip_proj, attn_norm, attn, ff_norm, ff) in enumerate(self.layers):
+            layer = idx + 1
+
+            # skip connection logic
+            is_first_half = layer <= (self.depth // 2)
+            is_later_half = not is_first_half
+
+            if is_first_half:
+                skips.append(x)
+
+            if is_later_half:
+                skip = skips.pop()
+                if skip_connect_type == "concat":
+                    x = torch.cat((x, skip), dim=-1)
+                    x = maybe_skip_proj(x)
+                elif skip_connect_type == "add":
+                    x = x + skip
+
+            # attention and feedforward blocks
+            x = attn(attn_norm(x), rope=rope, mask=mask) + x
+            x = ff(ff_norm(x)) + x
+
+        assert len(skips) == 0
+
+        x = self.norm_out(x)[:, 1:, :]  # unpack t from x
+
+        return self.proj_out(x)