Upload 12 files

Dockerfile

@@ -0,0 +1,25 @@
+FROM pytorch/pytorch:2.4.0-cuda12.4-cudnn9-devel
+
+USER root
+
+ARG DEBIAN_FRONTEND=noninteractive
+
+LABEL github_repo="https://github.com/SWivid/F5-TTS"
+
+RUN set -x \
+    && apt-get update \
+    && apt-get -y install wget curl man git less openssl libssl-dev unzip unar build-essential aria2 tmux vim \
+    && apt-get install -y openssh-server sox libsox-fmt-all libsox-fmt-mp3 libsndfile1-dev ffmpeg \
+    && rm -rf /var/lib/apt/lists/* \
+    && apt-get clean
+
+WORKDIR /workspace
+
+RUN git clone https://github.com/SWivid/F5-TTS.git \
+    && cd F5-TTS \
+    && pip install --no-cache-dir -r requirements.txt \
+    && pip install --no-cache-dir -r requirements_eval.txt
+
+ENV SHELL=/bin/bash
+
+WORKDIR /workspace/F5-TTS

app.py

@@ -0,0 +1,794 @@
+import re
+import torch
+import torchaudio
+import gradio as gr
+import numpy as np
+import tempfile
+from einops import rearrange
+from vocos import Vocos
+from pydub import AudioSegment, silence
+from model import CFM, UNetT, DiT, MMDiT
+from cached_path import cached_path
+from model.utils import (
+    load_checkpoint,
+    get_tokenizer,
+    convert_char_to_pinyin,
+    save_spectrogram,
+)
+from transformers import pipeline
+import click
+import soundfile as sf
+
+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
+
+device = (
+    "cuda"
+    if torch.cuda.is_available()
+    else "mps" if torch.backends.mps.is_available() else "cpu"
+)
+
+print(f"Using {device} device")
+
+pipe = pipeline(
+    "automatic-speech-recognition",
+    model="openai/whisper-large-v3-turbo",
+    torch_dtype=torch.float16,
+    device=device,
+)
+vocos = Vocos.from_pretrained("charactr/vocos-mel-24khz")
+
+# --------------------- Settings -------------------- #
+
+target_sample_rate = 24000
+n_mel_channels = 100
+hop_length = 256
+target_rms = 0.1
+nfe_step = 32  # 16, 32
+cfg_strength = 2.0
+ode_method = "euler"
+sway_sampling_coef = -1.0
+speed = 1.0
+fix_duration = None
+
+
+def load_model(repo_name, exp_name, model_cls, model_cfg, ckpt_step):
+    ckpt_path = str(cached_path(f"hf://SWivid/{repo_name}/{exp_name}/model_{ckpt_step}.safetensors"))
+    # ckpt_path = f"ckpts/{exp_name}/model_{ckpt_step}.pt"  # .pt | .safetensors
+    vocab_char_map, vocab_size = get_tokenizer("Emilia_ZH_EN", "pinyin")
+    model = CFM(
+        transformer=model_cls(
+            **model_cfg, text_num_embeds=vocab_size, mel_dim=n_mel_channels
+        ),
+        mel_spec_kwargs=dict(
+            target_sample_rate=target_sample_rate,
+            n_mel_channels=n_mel_channels,
+            hop_length=hop_length,
+        ),
+        odeint_kwargs=dict(
+            method=ode_method,
+        ),
+        vocab_char_map=vocab_char_map,
+    ).to(device)
+
+    model = load_checkpoint(model, ckpt_path, device, use_ema = True)
+
+    return model
+
+
+# load models
+F5TTS_model_cfg = dict(
+    dim=1024, depth=22, heads=16, ff_mult=2, text_dim=512, conv_layers=4
+)
+E2TTS_model_cfg = dict(dim=1024, depth=24, heads=16, ff_mult=4)
+
+F5TTS_ema_model = load_model(
+    "F5-TTS", "F5TTS_Base", DiT, F5TTS_model_cfg, 1200000
+)
+E2TTS_ema_model = load_model(
+    "E2-TTS", "E2TTS_Base", UNetT, E2TTS_model_cfg, 1200000
+)
+
+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
+
+@gpu_decorator
+def infer_batch(ref_audio, ref_text, gen_text_batches, exp_name, remove_silence, cross_fade_duration=0.15, progress=gr.Progress()):
+    if exp_name == "F5-TTS":
+        ema_model = F5TTS_ema_model
+    elif exp_name == "E2-TTS":
+        ema_model = E2TTS_ema_model
+
+    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 + " "
+    for i, gen_text in enumerate(progress.tqdm(gen_text_batches)):
+        # Prepare the text
+        text_list = [ref_text + gen_text]
+        final_text_list = convert_char_to_pinyin(text_list)
+
+        # Calculate duration
+        ref_audio_len = audio.shape[-1] // hop_length
+        zh_pause_punc = r"。，、；：？！"
+        ref_text_len = len(ref_text.encode('utf-8')) + 3 * len(re.findall(zh_pause_punc, ref_text))
+        gen_text_len = len(gen_text.encode('utf-8')) + 3 * len(re.findall(zh_pause_punc, gen_text))
+        duration = ref_audio_len + int(ref_audio_len / ref_text_len * gen_text_len / speed)
+
+        # inference
+        with torch.inference_mode():
+            generated, _ = ema_model.sample(
+                cond=audio,
+                text=final_text_list,
+                duration=duration,
+                steps=nfe_step,
+                cfg_strength=cfg_strength,
+                sway_sampling_coef=sway_sampling_coef,
+            )
+
+        generated = generated[:, ref_audio_len:, :]
+        generated_mel_spec = rearrange(generated, "1 n d -> 1 d n")
+        generated_wave = vocos.decode(generated_mel_spec.cpu())
+        if rms < target_rms:
+            generated_wave = generated_wave * rms / target_rms
+
+        # wav -> numpy
+        generated_wave = generated_wave.squeeze().cpu().numpy()
+        
+        generated_waves.append(generated_wave)
+        spectrograms.append(generated_mel_spec[0].cpu().numpy())
+
+    # Combine all generated waves with cross-fading
+    if cross_fade_duration <= 0:
+        # Simply concatenate
+        final_wave = np.concatenate(generated_waves)
+    else:
+        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
+
+    # Remove silence
+    if remove_silence:
+        with tempfile.NamedTemporaryFile(delete=False, suffix=".wav") as f:
+            sf.write(f.name, final_wave, target_sample_rate)
+            aseg = AudioSegment.from_file(f.name)
+            non_silent_segs = silence.split_on_silence(aseg, min_silence_len=1000, silence_thresh=-50, keep_silence=500)
+            non_silent_wave = AudioSegment.silent(duration=0)
+            for non_silent_seg in non_silent_segs:
+                non_silent_wave += non_silent_seg
+            aseg = non_silent_wave
+            aseg.export(f.name, format="wav")
+            final_wave, _ = torchaudio.load(f.name)
+        final_wave = final_wave.squeeze().cpu().numpy()
+
+    # Create a combined spectrogram
+    combined_spectrogram = np.concatenate(spectrograms, axis=1)
+    
+    with tempfile.NamedTemporaryFile(suffix=".png", delete=False) as tmp_spectrogram:
+        spectrogram_path = tmp_spectrogram.name
+        save_spectrogram(combined_spectrogram, spectrogram_path)
+
+    return (target_sample_rate, final_wave), spectrogram_path
+
+@gpu_decorator
+def infer(ref_audio_orig, ref_text, gen_text, exp_name, remove_silence, cross_fade_duration=0.15):
+
+    print(gen_text)
+
+    gr.Info("Converting audio...")
+    with tempfile.NamedTemporaryFile(delete=False, suffix=".wav") as f:
+        aseg = AudioSegment.from_file(ref_audio_orig)
+
+        non_silent_segs = silence.split_on_silence(
+            aseg, min_silence_len=1000, silence_thresh=-50, keep_silence=1000
+        )
+        non_silent_wave = AudioSegment.silent(duration=0)
+        for non_silent_seg in non_silent_segs:
+            non_silent_wave += non_silent_seg
+        aseg = non_silent_wave
+
+        audio_duration = len(aseg)
+        if audio_duration > 15000:
+            gr.Warning("Audio is over 15s, clipping to only first 15s.")
+            aseg = aseg[:15000]
+        aseg.export(f.name, format="wav")
+        ref_audio = f.name
+
+    if not ref_text.strip():
+        gr.Info("No reference text provided, transcribing reference audio...")
+        ref_text = pipe(
+            ref_audio,
+            chunk_length_s=30,
+            batch_size=128,
+            generate_kwargs={"task": "transcribe"},
+            return_timestamps=False,
+        )["text"].strip()
+        gr.Info("Finished transcription")
+    else:
+        gr.Info("Using custom reference text...")
+
+    # Add the functionality to ensure it ends with ". "
+    if not ref_text.endswith(". "):
+        if ref_text.endswith("."):
+            ref_text += " "
+        else:
+            ref_text += ". "
+
+    audio, sr = torchaudio.load(ref_audio)
+
+    # Use the new chunk_text function to split gen_text
+    max_chars = int(len(ref_text.encode('utf-8')) / (audio.shape[-1] / sr) * (25 - audio.shape[-1] / sr))
+    gen_text_batches = chunk_text(gen_text, max_chars=max_chars)
+    print('ref_text', ref_text)
+    for i, batch_text in enumerate(gen_text_batches):
+        print(f'gen_text {i}', batch_text)
+    
+    gr.Info(f"Generating audio using {exp_name} in {len(gen_text_batches)} batches")
+    return infer_batch((audio, sr), ref_text, gen_text_batches, exp_name, remove_silence, cross_fade_duration)
+
+
+@gpu_decorator
+def generate_podcast(script, speaker1_name, ref_audio1, ref_text1, speaker2_name, ref_audio2, ref_text2, exp_name, remove_silence):
+    # Split the script into speaker blocks
+    speaker_pattern = re.compile(f"^({re.escape(speaker1_name)}|{re.escape(speaker2_name)}):", re.MULTILINE)
+    speaker_blocks = speaker_pattern.split(script)[1:]  # Skip the first empty element
+    
+    generated_audio_segments = []
+    
+    for i in range(0, len(speaker_blocks), 2):
+        speaker = speaker_blocks[i]
+        text = speaker_blocks[i+1].strip()
+        
+        # Determine which speaker is talking
+        if speaker == speaker1_name:
+            ref_audio = ref_audio1
+            ref_text = ref_text1
+        elif speaker == speaker2_name:
+            ref_audio = ref_audio2
+            ref_text = ref_text2
+        else:
+            continue  # Skip if the speaker is neither speaker1 nor speaker2
+        
+        # Generate audio for this block
+        audio, _ = infer(ref_audio, ref_text, text, exp_name, remove_silence)
+        
+        # Convert the generated audio to a numpy array
+        sr, audio_data = audio
+        
+        # Save the audio data as a WAV file
+        with tempfile.NamedTemporaryFile(suffix=".wav", delete=False) as temp_file:
+            sf.write(temp_file.name, audio_data, sr)
+            audio_segment = AudioSegment.from_wav(temp_file.name)
+        
+        generated_audio_segments.append(audio_segment)
+        
+        # Add a short pause between speakers
+        pause = AudioSegment.silent(duration=500)  # 500ms pause
+        generated_audio_segments.append(pause)
+    
+    # Concatenate all audio segments
+    final_podcast = sum(generated_audio_segments)
+    
+    # Export the final podcast
+    with tempfile.NamedTemporaryFile(suffix=".wav", delete=False) as temp_file:
+        podcast_path = temp_file.name
+        final_podcast.export(podcast_path, format="wav")
+    
+    return podcast_path
+
+def parse_speechtypes_text(gen_text):
+    # Pattern to find (Emotion)
+    pattern = r'\((.*?)\)'
+
+    # Split the text by the pattern
+    tokens = re.split(pattern, gen_text)
+
+    segments = []
+
+    current_emotion = 'Regular'
+
+    for i in range(len(tokens)):
+        if i % 2 == 0:
+            # This is text
+            text = tokens[i].strip()
+            if text:
+                segments.append({'emotion': current_emotion, 'text': text})
+        else:
+            # This is emotion
+            emotion = tokens[i].strip()
+            current_emotion = emotion
+
+    return segments
+
+def update_speed(new_speed):
+    global speed
+    speed = new_speed
+    return f"Speed set to: {speed}"
+
+with gr.Blocks() as app_credits:
+    gr.Markdown("""
+# Credits
+
+* [mrfakename](https://github.com/fakerybakery) for the original [online demo](https://huggingface.co/spaces/mrfakename/E2-F5-TTS)
+* [RootingInLoad](https://github.com/RootingInLoad) for the podcast generation
+* [jpgallegoar](https://github.com/jpgallegoar) for multiple speech-type generation
+""")
+with gr.Blocks() as app_tts:
+    gr.Markdown("# Batched TTS")
+    ref_audio_input = gr.Audio(label="Reference Audio", type="filepath")
+    gen_text_input = gr.Textbox(label="Text to Generate", lines=10)
+    model_choice = gr.Radio(
+        choices=["F5-TTS", "E2-TTS"], label="Choose TTS Model", value="F5-TTS"
+    )
+    generate_btn = gr.Button("Synthesize", variant="primary")
+    with gr.Accordion("Advanced Settings", open=False):
+        ref_text_input = gr.Textbox(
+            label="Reference Text",
+            info="Leave blank to automatically transcribe the reference audio. If you enter text it will override automatic transcription.",
+            lines=2,
+        )
+        remove_silence = gr.Checkbox(
+            label="Remove Silences",
+            info="The model tends to produce silences, especially on longer audio. We can manually remove silences if needed. Note that this is an experimental feature and may produce strange results. This will also increase generation time.",
+            value=False,
+        )
+        speed_slider = gr.Slider(
+            label="Speed",
+            minimum=0.3,
+            maximum=2.0,
+            value=speed,
+            step=0.1,
+            info="Adjust the speed of the audio.",
+        )
+        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.",
+        )
+    speed_slider.change(update_speed, inputs=speed_slider)
+
+    audio_output = gr.Audio(label="Synthesized Audio")
+    spectrogram_output = gr.Image(label="Spectrogram")
+
+    generate_btn.click(
+        infer,
+        inputs=[
+            ref_audio_input,
+            ref_text_input,
+            gen_text_input,
+            model_choice,
+            remove_silence,
+            cross_fade_duration_slider,
+        ],
+        outputs=[audio_output, spectrogram_output],
+    )
+    
+with gr.Blocks() as app_podcast:
+    gr.Markdown("# Podcast Generation")
+    speaker1_name = gr.Textbox(label="Speaker 1 Name")
+    ref_audio_input1 = gr.Audio(label="Reference Audio (Speaker 1)", type="filepath")
+    ref_text_input1 = gr.Textbox(label="Reference Text (Speaker 1)", lines=2)
+    
+    speaker2_name = gr.Textbox(label="Speaker 2 Name")
+    ref_audio_input2 = gr.Audio(label="Reference Audio (Speaker 2)", type="filepath")
+    ref_text_input2 = gr.Textbox(label="Reference Text (Speaker 2)", lines=2)
+    
+    script_input = gr.Textbox(label="Podcast Script", lines=10, 
+                                placeholder="Enter the script with speaker names at the start of each block, e.g.:\nSean: How did you start studying...\n\nMeghan: I came to my interest in technology...\nIt was a long journey...\n\nSean: That's fascinating. Can you elaborate...")
+    
+    podcast_model_choice = gr.Radio(
+        choices=["F5-TTS", "E2-TTS"], label="Choose TTS Model", value="F5-TTS"
+    )
+    podcast_remove_silence = gr.Checkbox(
+        label="Remove Silences",
+        value=True,
+    )
+    generate_podcast_btn = gr.Button("Generate Podcast", variant="primary")
+    podcast_output = gr.Audio(label="Generated Podcast")
+
+    def podcast_generation(script, speaker1, ref_audio1, ref_text1, speaker2, ref_audio2, ref_text2, model, remove_silence):
+        return generate_podcast(script, speaker1, ref_audio1, ref_text1, speaker2, ref_audio2, ref_text2, model, remove_silence)
+
+    generate_podcast_btn.click(
+        podcast_generation,
+        inputs=[
+            script_input,
+            speaker1_name,
+            ref_audio_input1,
+            ref_text_input1,
+            speaker2_name,
+            ref_audio_input2,
+            ref_text_input2,
+            podcast_model_choice,
+            podcast_remove_silence,
+        ],
+        outputs=podcast_output,
+    )
+
+def parse_emotional_text(gen_text):
+    # Pattern to find (Emotion)
+    pattern = r'\((.*?)\)'
+
+    # Split the text by the pattern
+    tokens = re.split(pattern, gen_text)
+
+    segments = []
+
+    current_emotion = 'Regular'
+
+    for i in range(len(tokens)):
+        if i % 2 == 0:
+            # This is text
+            text = tokens[i].strip()
+            if text:
+                segments.append({'emotion': current_emotion, 'text': text})
+        else:
+            # This is emotion
+            emotion = tokens[i].strip()
+            current_emotion = emotion
+
+    return segments
+
+with gr.Blocks() as app_emotional:
+    # New section for emotional generation
+    gr.Markdown(
+        """
+    # Multiple Speech-Type Generation
+
+    This section allows you to upload different audio clips for each speech type. 'Regular' emotion is mandatory. You can add additional speech types by clicking the "Add Speech Type" button. Enter your text in the format shown below, and the system will generate speech using the appropriate emotions. If unspecified, the model will use the regular speech type. The current speech type will be used until the next speech type is specified.
+
+    **Example Input:**
+
+    (Regular) Hello, I'd like to order a sandwich please. (Surprised) What do you mean you're out of bread? (Sad) I really wanted a sandwich though... (Angry) You know what, darn you and your little shop, you suck! (Whisper) I'll just go back home and cry now. (Shouting) Why me?!
+    """
+    )
+
+    gr.Markdown("Upload different audio clips for each speech type. 'Regular' emotion is mandatory. You can add additional speech types by clicking the 'Add Speech Type' button.")
+
+    # Regular speech type (mandatory)
+    with gr.Row():
+        regular_name = gr.Textbox(value='Regular', label='Speech Type Name', interactive=False)
+        regular_audio = gr.Audio(label='Regular Reference Audio', type='filepath')
+        regular_ref_text = gr.Textbox(label='Reference Text (Regular)', lines=2)
+
+    # Additional speech types (up to 99 more)
+    max_speech_types = 100
+    speech_type_names = []
+    speech_type_audios = []
+    speech_type_ref_texts = []
+    speech_type_delete_btns = []
+
+    for i in range(max_speech_types - 1):
+        with gr.Row():
+            name_input = gr.Textbox(label='Speech Type Name', visible=False)
+            audio_input = gr.Audio(label='Reference Audio', type='filepath', visible=False)
+            ref_text_input = gr.Textbox(label='Reference Text', lines=2, visible=False)
+            delete_btn = gr.Button("Delete", variant="secondary", visible=False)
+        speech_type_names.append(name_input)
+        speech_type_audios.append(audio_input)
+        speech_type_ref_texts.append(ref_text_input)
+        speech_type_delete_btns.append(delete_btn)
+
+    # Button to add speech type
+    add_speech_type_btn = gr.Button("Add Speech Type")
+
+    # Keep track of current number of speech types
+    speech_type_count = gr.State(value=0)
+
+    # Function to add a speech type
+    def add_speech_type_fn(speech_type_count):
+        if speech_type_count < max_speech_types - 1:
+            speech_type_count += 1
+            # Prepare updates for the components
+            name_updates = []
+            audio_updates = []
+            ref_text_updates = []
+            delete_btn_updates = []
+            for i in range(max_speech_types - 1):
+                if i < speech_type_count:
+                    name_updates.append(gr.update(visible=True))
+                    audio_updates.append(gr.update(visible=True))
+                    ref_text_updates.append(gr.update(visible=True))
+                    delete_btn_updates.append(gr.update(visible=True))
+                else:
+                    name_updates.append(gr.update())
+                    audio_updates.append(gr.update())
+                    ref_text_updates.append(gr.update())
+                    delete_btn_updates.append(gr.update())
+        else:
+            # Optionally, show a warning
+            # gr.Warning("Maximum number of speech types reached.")
+            name_updates = [gr.update() for _ in range(max_speech_types - 1)]
+            audio_updates = [gr.update() for _ in range(max_speech_types - 1)]
+            ref_text_updates = [gr.update() for _ in range(max_speech_types - 1)]
+            delete_btn_updates = [gr.update() for _ in range(max_speech_types - 1)]
+        return [speech_type_count] + name_updates + audio_updates + ref_text_updates + delete_btn_updates
+
+    add_speech_type_btn.click(
+        add_speech_type_fn,
+        inputs=speech_type_count,
+        outputs=[speech_type_count] + speech_type_names + speech_type_audios + speech_type_ref_texts + speech_type_delete_btns
+    )
+
+    # Function to delete a speech type
+    def make_delete_speech_type_fn(index):
+        def delete_speech_type_fn(speech_type_count):
+            # Prepare updates
+            name_updates = []
+            audio_updates = []
+            ref_text_updates = []
+            delete_btn_updates = []
+
+            for i in range(max_speech_types - 1):
+                if i == index:
+                    name_updates.append(gr.update(visible=False, value=''))
+                    audio_updates.append(gr.update(visible=False, value=None))
+                    ref_text_updates.append(gr.update(visible=False, value=''))
+                    delete_btn_updates.append(gr.update(visible=False))
+                else:
+                    name_updates.append(gr.update())
+                    audio_updates.append(gr.update())
+                    ref_text_updates.append(gr.update())
+                    delete_btn_updates.append(gr.update())
+
+            speech_type_count = max(0, speech_type_count - 1)
+
+            return [speech_type_count] + name_updates + audio_updates + ref_text_updates + delete_btn_updates
+
+        return delete_speech_type_fn
+
+    for i, delete_btn in enumerate(speech_type_delete_btns):
+        delete_fn = make_delete_speech_type_fn(i)
+        delete_btn.click(
+            delete_fn,
+            inputs=speech_type_count,
+            outputs=[speech_type_count] + speech_type_names + speech_type_audios + speech_type_ref_texts + speech_type_delete_btns
+        )
+
+    # Text input for the prompt
+    gen_text_input_emotional = gr.Textbox(label="Text to Generate", lines=10)
+
+    # Model choice
+    model_choice_emotional = gr.Radio(
+        choices=["F5-TTS", "E2-TTS"], label="Choose TTS Model", value="F5-TTS"
+    )
+
+    with gr.Accordion("Advanced Settings", open=False):
+        remove_silence_emotional = gr.Checkbox(
+            label="Remove Silences",
+            value=True,
+        )
+
+    # Generate button
+    generate_emotional_btn = gr.Button("Generate Emotional Speech", variant="primary")
+
+    # Output audio
+    audio_output_emotional = gr.Audio(label="Synthesized Audio")
+    @gpu_decorator
+    def generate_emotional_speech(
+        regular_audio,
+        regular_ref_text,
+        gen_text,
+        *args,
+    ):
+        num_additional_speech_types = max_speech_types - 1
+        speech_type_names_list = args[:num_additional_speech_types]
+        speech_type_audios_list = args[num_additional_speech_types:2 * num_additional_speech_types]
+        speech_type_ref_texts_list = args[2 * num_additional_speech_types:3 * num_additional_speech_types]
+        model_choice = args[3 * num_additional_speech_types]
+        remove_silence = args[3 * num_additional_speech_types + 1]
+
+        # Collect the speech types and their audios into a dict
+        speech_types = {'Regular': {'audio': regular_audio, 'ref_text': regular_ref_text}}
+
+        for name_input, audio_input, ref_text_input in zip(speech_type_names_list, speech_type_audios_list, speech_type_ref_texts_list):
+            if name_input and audio_input:
+                speech_types[name_input] = {'audio': audio_input, 'ref_text': ref_text_input}
+
+        # Parse the gen_text into segments
+        segments = parse_speechtypes_text(gen_text)
+
+        # For each segment, generate speech
+        generated_audio_segments = []
+        current_emotion = 'Regular'
+
+        for segment in segments:
+            emotion = segment['emotion']
+            text = segment['text']
+
+            if emotion in speech_types:
+                current_emotion = emotion
+            else:
+                # If emotion not available, default to Regular
+                current_emotion = 'Regular'
+
+            ref_audio = speech_types[current_emotion]['audio']
+            ref_text = speech_types[current_emotion].get('ref_text', '')
+
+            # Generate speech for this segment
+            audio, _ = infer(ref_audio, ref_text, text, model_choice, remove_silence, 0)
+            sr, audio_data = audio
+
+            generated_audio_segments.append(audio_data)
+
+        # Concatenate all audio segments
+        if generated_audio_segments:
+            final_audio_data = np.concatenate(generated_audio_segments)
+            return (sr, final_audio_data)
+        else:
+            gr.Warning("No audio generated.")
+            return None
+
+    generate_emotional_btn.click(
+        generate_emotional_speech,
+        inputs=[
+            regular_audio,
+            regular_ref_text,
+            gen_text_input_emotional,
+        ] + speech_type_names + speech_type_audios + speech_type_ref_texts + [
+            model_choice_emotional,
+            remove_silence_emotional,
+        ],
+        outputs=audio_output_emotional,
+    )
+
+    # Validation function to disable Generate button if speech types are missing
+    def validate_speech_types(
+        gen_text,
+        regular_name,
+        *args
+    ):
+        num_additional_speech_types = max_speech_types - 1
+        speech_type_names_list = args[:num_additional_speech_types]
+
+        # Collect the speech types names
+        speech_types_available = set()
+        if regular_name:
+            speech_types_available.add(regular_name)
+        for name_input in speech_type_names_list:
+            if name_input:
+                speech_types_available.add(name_input)
+
+        # Parse the gen_text to get the speech types used
+        segments = parse_emotional_text(gen_text)
+        speech_types_in_text = set(segment['emotion'] for segment in segments)
+
+        # Check if all speech types in text are available
+        missing_speech_types = speech_types_in_text - speech_types_available
+
+        if missing_speech_types:
+            # Disable the generate button
+            return gr.update(interactive=False)
+        else:
+            # Enable the generate button
+            return gr.update(interactive=True)
+
+    gen_text_input_emotional.change(
+        validate_speech_types,
+        inputs=[gen_text_input_emotional, regular_name] + speech_type_names,
+        outputs=generate_emotional_btn
+    )
+with gr.Blocks() as app:
+    gr.Markdown(
+        """
+# E2/F5 TTS
+
+This is a local web UI for 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 support English and Chinese.
+
+If you're having issues, try converting your reference audio to WAV or MP3, clipping it to 15s, and shortening your prompt.
+
+**NOTE: Reference text will be automatically transcribed with Whisper if not provided. For best results, keep your reference clips short (<15s). Ensure the audio is fully uploaded before generating.**
+"""
+    )
+    gr.TabbedInterface([app_tts, app_podcast, app_emotional, app_credits], ["TTS", "Podcast", "Multi-Style", "Credits"])
+
+@click.command()
+@click.option("--port", "-p", default=None, type=int, help="Port to run the app on")
+@click.option("--host", "-H", default=None, help="Host to run the app on")
+@click.option(
+    "--share",
+    "-s",
+    default=False,
+    is_flag=True,
+    help="Share the app via Gradio share link",
+)
+@click.option("--api", "-a", default=True, is_flag=True, help="Allow API access")
+def main(port, host, share, api):
+    global app
+    print(f"Starting app...")
+    app.queue(api_open=api).launch(
+        server_name=host, server_port=port, share=share, show_api=api
+    )
+
+
+if __name__ == "__main__":
+    if not USING_SPACES:
+        main()
+    else:
+        app.queue().launch()

finetune-cli.py

@@ -0,0 +1,108 @@
+import argparse
+from model import CFM, UNetT, DiT, MMDiT, Trainer
+from model.utils import get_tokenizer
+from model.dataset import load_dataset
+from cached_path import cached_path
+import shutil,os
+# -------------------------- Dataset Settings --------------------------- #
+target_sample_rate = 24000
+n_mel_channels = 100
+hop_length = 256
+
+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)
+
+# -------------------------- Argument Parsing --------------------------- #
+def parse_args():
+    parser = argparse.ArgumentParser(description='Train CFM Model')
+
+    parser.add_argument('--exp_name', type=str, default="F5TTS_Base", choices=["F5TTS_Base", "E2TTS_Base"],help='Experiment name')
+    parser.add_argument('--dataset_name', type=str, default="Emilia_ZH_EN", help='Name of the dataset to use')
+    parser.add_argument('--learning_rate', type=float, default=1e-4, help='Learning rate for training')
+    parser.add_argument('--batch_size_per_gpu', type=int, default=256, help='Batch size per GPU')
+    parser.add_argument('--batch_size_type', type=str, default="frame", choices=["frame", "sample"],help='Batch size type')
+    parser.add_argument('--max_samples', type=int, default=16, help='Max sequences per batch')
+    parser.add_argument('--grad_accumulation_steps', type=int, default=1,help='Gradient accumulation steps')
+    parser.add_argument('--max_grad_norm', type=float, default=1.0, help='Max gradient norm for clipping')
+    parser.add_argument('--epochs', type=int, default=10, help='Number of training epochs')
+    parser.add_argument('--num_warmup_updates', type=int, default=5, help='Warmup steps')
+    parser.add_argument('--save_per_updates', type=int, default=10, help='Save checkpoint every X steps')
+    parser.add_argument('--last_per_steps', type=int, default=10, help='Save last checkpoint every X steps')
+    parser.add_argument('--finetune', type=bool, default=True, help='Use Finetune')
+    
+    return parser.parse_args()
+
+# -------------------------- Training Settings -------------------------- #
+
+def main():
+    args = parse_args()
+    
+
+    # Model parameters based on experiment name
+    if args.exp_name == "F5TTS_Base":
+        wandb_resume_id = None
+        model_cls = DiT
+        model_cfg = dict(dim=1024, depth=22, heads=16, ff_mult=2, text_dim=512, conv_layers=4)
+        if args.finetune:
+           ckpt_path = str(cached_path(f"hf://SWivid/F5-TTS/F5TTS_Base/model_1200000.pt"))
+    elif args.exp_name == "E2TTS_Base":
+        wandb_resume_id = None
+        model_cls = UNetT
+        model_cfg = dict(dim=1024, depth=24, heads=16, ff_mult=4)
+        if args.finetune:
+           ckpt_path = str(cached_path(f"hf://SWivid/E2-TTS/E2TTS_Base/model_1200000.pt"))
+           
+    if args.finetune:
+       path_ckpt = os.path.join("ckpts",args.dataset_name)
+       if os.path.isdir(path_ckpt)==False:
+          os.makedirs(path_ckpt,exist_ok=True)
+          shutil.copy2(ckpt_path,os.path.join(path_ckpt,os.path.basename(ckpt_path)))
+    
+    checkpoint_path=os.path.join("ckpts",args.dataset_name)
+
+    # Use the dataset_name provided in the command line
+    tokenizer_path = args.dataset_name if tokenizer != "custom" else tokenizer_path
+    vocab_char_map, vocab_size = get_tokenizer(tokenizer_path, tokenizer)
+
+    mel_spec_kwargs = dict(
+        target_sample_rate=target_sample_rate,
+        n_mel_channels=n_mel_channels,
+        hop_length=hop_length,
+    )
+
+    e2tts = CFM(
+        transformer=model_cls(
+            **model_cfg,
+            text_num_embeds=vocab_size,
+            mel_dim=n_mel_channels
+        ),
+        mel_spec_kwargs=mel_spec_kwargs,
+        vocab_char_map=vocab_char_map,
+    )
+
+    trainer = Trainer(
+        e2tts,
+        args.epochs,
+        args.learning_rate,
+        num_warmup_updates=args.num_warmup_updates,
+        save_per_updates=args.save_per_updates,
+        checkpoint_path=checkpoint_path,
+        batch_size=args.batch_size_per_gpu,
+        batch_size_type=args.batch_size_type,
+        max_samples=args.max_samples,
+        grad_accumulation_steps=args.grad_accumulation_steps,
+        max_grad_norm=args.max_grad_norm,
+        wandb_project="CFM-TTS",
+        wandb_run_name=args.exp_name,
+        wandb_resume_id=wandb_resume_id,
+        last_per_steps=args.last_per_steps,
+    )
+
+    train_dataset = load_dataset(args.dataset_name, tokenizer, mel_spec_kwargs=mel_spec_kwargs)
+    trainer.train(train_dataset,
+                  resumable_with_seed=666  # seed for shuffling dataset
+                  )
+
+
+if __name__ == '__main__':
+    main()

finetune_gradio.py

@@ -0,0 +1,734 @@
+import os,sys
+
+from transformers import pipeline
+import gradio as gr
+import torch
+import click
+import torchaudio
+from glob import glob
+import librosa
+import numpy as np
+from scipy.io import wavfile
+import shutil
+import time
+
+import json
+from model.utils import convert_char_to_pinyin
+import signal
+import psutil
+import platform
+import subprocess
+from datasets.arrow_writer import ArrowWriter
+
+import json
+
+training_process = None    
+system = platform.system()
+python_executable = sys.executable or "python"
+
+path_data="data"
+
+device = (
+"cuda"
+    if torch.cuda.is_available()
+    else "mps" if torch.backends.mps.is_available() else "cpu"
+)
+
+pipe = None
+
+# Load metadata
+def get_audio_duration(audio_path):
+    """Calculate the duration of an audio file."""
+    audio, sample_rate = torchaudio.load(audio_path)
+    num_channels = audio.shape[0]  
+    return audio.shape[1] / (sample_rate * num_channels)
+
+def clear_text(text):
+    """Clean and prepare text by lowering the case and stripping whitespace."""
+    return text.lower().strip()
+
+def get_rms(y,frame_length=2048,hop_length=512,pad_mode="constant",): # https://github.com/RVC-Boss/GPT-SoVITS/blob/main/tools/slicer2.py
+    padding = (int(frame_length // 2), int(frame_length // 2))
+    y = np.pad(y, padding, mode=pad_mode)
+
+    axis = -1
+    # put our new within-frame axis at the end for now
+    out_strides = y.strides + tuple([y.strides[axis]])
+    # Reduce the shape on the framing axis
+    x_shape_trimmed = list(y.shape)
+    x_shape_trimmed[axis] -= frame_length - 1
+    out_shape = tuple(x_shape_trimmed) + tuple([frame_length])
+    xw = np.lib.stride_tricks.as_strided(y, shape=out_shape, strides=out_strides)
+    if axis < 0:
+        target_axis = axis - 1
+    else:
+        target_axis = axis + 1
+    xw = np.moveaxis(xw, -1, target_axis)
+    # Downsample along the target axis
+    slices = [slice(None)] * xw.ndim
+    slices[axis] = slice(0, None, hop_length)
+    x = xw[tuple(slices)]
+
+    # Calculate power
+    power = np.mean(np.abs(x) ** 2, axis=-2, keepdims=True)
+
+    return np.sqrt(power)
+
+class Slicer: # https://github.com/RVC-Boss/GPT-SoVITS/blob/main/tools/slicer2.py
+    def __init__(
+        self,
+        sr: int,
+        threshold: float = -40.0,
+        min_length: int = 2000,
+        min_interval: int = 300,
+        hop_size: int = 20,
+        max_sil_kept: int = 2000,
+    ):
+        if not min_length >= min_interval >= hop_size:
+            raise ValueError(
+                "The following condition must be satisfied: min_length >= min_interval >= hop_size"
+            )
+        if not max_sil_kept >= hop_size:
+            raise ValueError(
+                "The following condition must be satisfied: max_sil_kept >= hop_size"
+            )
+        min_interval = sr * min_interval / 1000
+        self.threshold = 10 ** (threshold / 20.0)
+        self.hop_size = round(sr * hop_size / 1000)
+        self.win_size = min(round(min_interval), 4 * self.hop_size)
+        self.min_length = round(sr * min_length / 1000 / self.hop_size)
+        self.min_interval = round(min_interval / self.hop_size)
+        self.max_sil_kept = round(sr * max_sil_kept / 1000 / self.hop_size)
+
+    def _apply_slice(self, waveform, begin, end):
+        if len(waveform.shape) > 1:
+            return waveform[
+                :, begin * self.hop_size : min(waveform.shape[1], end * self.hop_size)
+            ]
+        else:
+            return waveform[
+                begin * self.hop_size : min(waveform.shape[0], end * self.hop_size)
+            ]
+
+    # @timeit
+    def slice(self, waveform):
+        if len(waveform.shape) > 1:
+            samples = waveform.mean(axis=0)
+        else:
+            samples = waveform
+        if samples.shape[0] <= self.min_length:
+            return [waveform]
+        rms_list = get_rms(
+            y=samples, frame_length=self.win_size, hop_length=self.hop_size
+        ).squeeze(0)
+        sil_tags = []
+        silence_start = None
+        clip_start = 0
+        for i, rms in enumerate(rms_list):
+            # Keep looping while frame is silent.
+            if rms < self.threshold:
+                # Record start of silent frames.
+                if silence_start is None:
+                    silence_start = i
+                continue
+            # Keep looping while frame is not silent and silence start has not been recorded.
+            if silence_start is None:
+                continue
+            # Clear recorded silence start if interval is not enough or clip is too short
+            is_leading_silence = silence_start == 0 and i > self.max_sil_kept
+            need_slice_middle = (
+                i - silence_start >= self.min_interval
+                and i - clip_start >= self.min_length
+            )
+            if not is_leading_silence and not need_slice_middle:
+                silence_start = None
+                continue
+            # Need slicing. Record the range of silent frames to be removed.
+            if i - silence_start <= self.max_sil_kept:
+                pos = rms_list[silence_start : i + 1].argmin() + silence_start
+                if silence_start == 0:
+                    sil_tags.append((0, pos))
+                else:
+                    sil_tags.append((pos, pos))
+                clip_start = pos
+            elif i - silence_start <= self.max_sil_kept * 2:
+                pos = rms_list[
+                    i - self.max_sil_kept : silence_start + self.max_sil_kept + 1
+                ].argmin()
+                pos += i - self.max_sil_kept
+                pos_l = (
+                    rms_list[
+                        silence_start : silence_start + self.max_sil_kept + 1
+                    ].argmin()
+                    + silence_start
+                )
+                pos_r = (
+                    rms_list[i - self.max_sil_kept : i + 1].argmin()
+                    + i
+                    - self.max_sil_kept
+                )
+                if silence_start == 0:
+                    sil_tags.append((0, pos_r))
+                    clip_start = pos_r
+                else:
+                    sil_tags.append((min(pos_l, pos), max(pos_r, pos)))
+                    clip_start = max(pos_r, pos)
+            else:
+                pos_l = (
+                    rms_list[
+                        silence_start : silence_start + self.max_sil_kept + 1
+                    ].argmin()
+                    + silence_start
+                )
+                pos_r = (
+                    rms_list[i - self.max_sil_kept : i + 1].argmin()
+                    + i
+                    - self.max_sil_kept
+                )
+                if silence_start == 0:
+                    sil_tags.append((0, pos_r))
+                else:
+                    sil_tags.append((pos_l, pos_r))
+                clip_start = pos_r
+            silence_start = None
+        # Deal with trailing silence.
+        total_frames = rms_list.shape[0]
+        if (
+            silence_start is not None
+            and total_frames - silence_start >= self.min_interval
+        ):
+            silence_end = min(total_frames, silence_start + self.max_sil_kept)
+            pos = rms_list[silence_start : silence_end + 1].argmin() + silence_start
+            sil_tags.append((pos, total_frames + 1))
+        # Apply and return slices.
+        ####音频+起始时间+终止时间
+        if len(sil_tags) == 0:
+            return [[waveform,0,int(total_frames*self.hop_size)]]
+        else:
+            chunks = []
+            if sil_tags[0][0] > 0:
+                chunks.append([self._apply_slice(waveform, 0, sil_tags[0][0]),0,int(sil_tags[0][0]*self.hop_size)])
+            for i in range(len(sil_tags) - 1):
+                chunks.append(
+                    [self._apply_slice(waveform, sil_tags[i][1], sil_tags[i + 1][0]),int(sil_tags[i][1]*self.hop_size),int(sil_tags[i + 1][0]*self.hop_size)]
+                )
+            if sil_tags[-1][1] < total_frames:
+                chunks.append(
+                    [self._apply_slice(waveform, sil_tags[-1][1], total_frames),int(sil_tags[-1][1]*self.hop_size),int(total_frames*self.hop_size)]
+                )
+            return chunks
+
+#terminal
+def terminate_process_tree(pid, including_parent=True):  
+    try:
+        parent = psutil.Process(pid)
+    except psutil.NoSuchProcess:
+        # Process already terminated
+        return
+
+    children = parent.children(recursive=True)
+    for child in children:
+        try:
+            os.kill(child.pid, signal.SIGTERM)  # or signal.SIGKILL
+        except OSError:
+            pass
+    if including_parent:
+        try:
+            os.kill(parent.pid, signal.SIGTERM)  # or signal.SIGKILL
+        except OSError:
+            pass
+
+def terminate_process(pid):
+    if system == "Windows":
+        cmd = f"taskkill /t /f /pid {pid}"
+        os.system(cmd)
+    else:
+        terminate_process_tree(pid)
+
+def start_training(dataset_name="",     
+    exp_name="F5TTS_Base",            
+    learning_rate=1e-4,               
+    batch_size_per_gpu=400,           
+    batch_size_type="frame",          
+    max_samples=64,                   
+    grad_accumulation_steps=1,        
+    max_grad_norm=1.0,                
+    epochs=11,                        
+    num_warmup_updates=200,           
+    save_per_updates=400,             
+    last_per_steps=800,               
+    finetune=True,                    
+    ):
+
+
+    global training_process
+
+    path_project = os.path.join(path_data, dataset_name + "_pinyin")
+
+    if os.path.isdir(path_project)==False:
+        yield f"There is not project with name {dataset_name}",gr.update(interactive=True),gr.update(interactive=False)
+        return
+
+    file_raw = os.path.join(path_project,"raw.arrow")
+    if os.path.isfile(file_raw)==False:
+       yield f"There is no file {file_raw}",gr.update(interactive=True),gr.update(interactive=False)
+       return    
+
+    # Check if a training process is already running
+    if training_process is not None:
+        return "Train run already!",gr.update(interactive=False),gr.update(interactive=True)
+
+    yield "start train",gr.update(interactive=False),gr.update(interactive=False)
+
+    # Command to run the training script with the specified arguments
+    cmd = f"accelerate launch finetune-cli.py --exp_name {exp_name} " \
+          f"--learning_rate {learning_rate} " \
+          f"--batch_size_per_gpu {batch_size_per_gpu} " \
+          f"--batch_size_type {batch_size_type} " \
+          f"--max_samples {max_samples} " \
+          f"--grad_accumulation_steps {grad_accumulation_steps} " \
+          f"--max_grad_norm {max_grad_norm} " \
+          f"--epochs {epochs} " \
+          f"--num_warmup_updates {num_warmup_updates} " \
+          f"--save_per_updates {save_per_updates} " \
+          f"--last_per_steps {last_per_steps} " \
+          f"--dataset_name {dataset_name}"  
+    if finetune:cmd += f" --finetune {finetune}"   
+
+    print(cmd)
+        
+    try:
+      # Start the training process
+      training_process = subprocess.Popen(cmd, shell=True)
+
+      time.sleep(5) 
+      yield "check terminal for wandb",gr.update(interactive=False),gr.update(interactive=True) 
+      
+      # Wait for the training process to finish
+      training_process.wait()
+      time.sleep(1)
+      
+      if training_process is None: 
+         text_info = 'train stop'
+      else:     
+         text_info = "train complete !"
+
+    except Exception as e:  # Catch all exceptions
+        # Ensure that we reset the training process variable in case of an error
+        text_info=f"An error occurred: {str(e)}"
+    
+    training_process=None
+
+    yield text_info,gr.update(interactive=True),gr.update(interactive=False)
+
+def stop_training():
+    global training_process
+    if training_process is None:return f"Train not run !",gr.update(interactive=True),gr.update(interactive=False)
+    terminate_process_tree(training_process.pid)
+    training_process = None
+    return 'train stop',gr.update(interactive=True),gr.update(interactive=False)
+
+def create_data_project(name):
+    name+="_pinyin"
+    os.makedirs(os.path.join(path_data,name),exist_ok=True)
+    os.makedirs(os.path.join(path_data,name,"dataset"),exist_ok=True)
+    
+def transcribe(file_audio,language="english"):
+    global pipe
+
+    if pipe is None:
+       pipe = pipeline("automatic-speech-recognition",model="openai/whisper-large-v3-turbo", torch_dtype=torch.float16,device=device)
+
+    text_transcribe = pipe(
+        file_audio,
+        chunk_length_s=30,
+        batch_size=128,
+        generate_kwargs={"task": "transcribe","language": language},
+        return_timestamps=False,
+    )["text"].strip()
+    return text_transcribe
+
+def transcribe_all(name_project,audio_files,language,user=False,progress=gr.Progress()):
+    name_project+="_pinyin"
+    path_project= os.path.join(path_data,name_project)
+    path_dataset = os.path.join(path_project,"dataset")
+    path_project_wavs = os.path.join(path_project,"wavs")
+    file_metadata = os.path.join(path_project,"metadata.csv")
+
+    if audio_files is None:return "You need to load an audio file."
+
+    if os.path.isdir(path_project_wavs):
+       shutil.rmtree(path_project_wavs)
+
+    if os.path.isfile(file_metadata):
+       os.remove(file_metadata)
+
+    os.makedirs(path_project_wavs,exist_ok=True)
+    
+    if user:
+       file_audios = [file for format in ('*.wav', '*.ogg', '*.opus', '*.mp3', '*.flac') for file in glob(os.path.join(path_dataset, format))]
+       if file_audios==[]:return "No audio file was found in the dataset."
+    else:
+       file_audios = audio_files
+       
+
+    alpha = 0.5
+    _max = 1.0
+    slicer = Slicer(24000)
+
+    num = 0
+    error_num = 0
+    data=""
+    for file_audio in progress.tqdm(file_audios, desc="transcribe files",total=len((file_audios))):
+        
+        audio, _ = librosa.load(file_audio, sr=24000, mono=True) 
+
+        list_slicer=slicer.slice(audio)
+        for chunk, start, end in progress.tqdm(list_slicer,total=len(list_slicer), desc="slicer files"): 
+            
+            name_segment = os.path.join(f"segment_{num}")
+            file_segment = os.path.join(path_project_wavs, f"{name_segment}.wav")    
+            
+            tmp_max = np.abs(chunk).max()
+            if(tmp_max>1):chunk/=tmp_max
+            chunk = (chunk / tmp_max * (_max * alpha)) + (1 - alpha) * chunk
+            wavfile.write(file_segment,24000, (chunk * 32767).astype(np.int16))
+            
+            try:
+               text=transcribe(file_segment,language)
+               text = text.lower().strip().replace('"',"")
+
+               data+= f"{name_segment}|{text}\n"
+
+               num+=1
+            except:
+               error_num +=1
+
+    with open(file_metadata,"w",encoding="utf-8") as f:
+        f.write(data)
+    
+    if error_num!=[]:
+       error_text=f"\nerror files : {error_num}"    
+    else:
+       error_text=""
+       
+    return f"transcribe complete samples : {num}\npath : {path_project_wavs}{error_text}"
+
+def format_seconds_to_hms(seconds):
+    hours = int(seconds / 3600)
+    minutes = int((seconds % 3600) / 60)
+    seconds = seconds % 60
+    return "{:02d}:{:02d}:{:02d}".format(hours, minutes, int(seconds))
+
+def create_metadata(name_project,progress=gr.Progress()):
+    name_project+="_pinyin"
+    path_project= os.path.join(path_data,name_project)
+    path_project_wavs = os.path.join(path_project,"wavs")
+    file_metadata = os.path.join(path_project,"metadata.csv")
+    file_raw = os.path.join(path_project,"raw.arrow")
+    file_duration = os.path.join(path_project,"duration.json")
+    file_vocab = os.path.join(path_project,"vocab.txt")
+
+    if os.path.isfile(file_metadata)==False: return "The file was not found in " + file_metadata 
+    
+    with open(file_metadata,"r",encoding="utf-8") as f:
+        data=f.read()
+    
+    audio_path_list=[]
+    text_list=[]
+    duration_list=[]
+    
+    count=data.split("\n")
+    lenght=0
+    result=[]
+    error_files=[]
+    for line in progress.tqdm(data.split("\n"),total=count):
+        sp_line=line.split("|")
+        if len(sp_line)!=2:continue
+        name_audio,text = sp_line[:2]
+
+        file_audio = os.path.join(path_project_wavs, name_audio + ".wav")
+
+        if os.path.isfile(file_audio)==False:
+            error_files.append(file_audio)
+            continue
+
+        duraction = get_audio_duration(file_audio)
+        if duraction<2 and duraction>15:continue
+        if len(text)<4:continue
+
+        text = clear_text(text)
+        text = convert_char_to_pinyin([text], polyphone = True)[0]
+
+        audio_path_list.append(file_audio)
+        duration_list.append(duraction)
+        text_list.append(text)
+      
+        result.append({"audio_path": file_audio, "text": text, "duration": duraction})
+
+        lenght+=duraction
+
+    if duration_list==[]:
+        error_files_text="\n".join(error_files)
+        return f"Error: No audio files found in the specified path : \n{error_files_text}"
+    
+    min_second = round(min(duration_list),2)   
+    max_second = round(max(duration_list),2)
+
+    with ArrowWriter(path=file_raw, writer_batch_size=1) as writer:
+        for line in progress.tqdm(result,total=len(result), desc=f"prepare data"):
+            writer.write(line)
+
+    with open(file_duration, 'w', encoding='utf-8') as f:
+        json.dump({"duration": duration_list}, f, ensure_ascii=False)
+ 
+    file_vocab_finetune = "data/Emilia_ZH_EN_pinyin/vocab.txt"    
+    if os.path.isfile(file_vocab_finetune==False):return "Error: Vocabulary file 'Emilia_ZH_EN_pinyin' not found!"
+    shutil.copy2(file_vocab_finetune, file_vocab)
+    
+    if error_files!=[]:
+       error_text="error files\n" + "\n".join(error_files)
+    else:
+       error_text=""
+        
+    return f"prepare complete \nsamples : {len(text_list)}\ntime data : {format_seconds_to_hms(lenght)}\nmin sec : {min_second}\nmax sec : {max_second}\nfile_arrow : {file_raw}\n{error_text}"
+
+def check_user(value):
+    return gr.update(visible=not value),gr.update(visible=value)
+
+def calculate_train(name_project,batch_size_type,max_samples,learning_rate,num_warmup_updates,save_per_updates,last_per_steps,finetune):
+    name_project+="_pinyin"
+    path_project= os.path.join(path_data,name_project)
+    file_duraction = os.path.join(path_project,"duration.json")
+
+    with open(file_duraction, 'r') as file:
+        data = json.load(file)  
+    
+    duration_list = data['duration']
+
+    samples = len(duration_list)
+
+    if torch.cuda.is_available():
+        gpu_properties = torch.cuda.get_device_properties(0)
+        total_memory = gpu_properties.total_memory / (1024 ** 3)  
+    elif torch.backends.mps.is_available():
+        total_memory = psutil.virtual_memory().available / (1024 ** 3)
+    
+    if batch_size_type=="frame":
+       batch = int(total_memory * 0.5)
+       batch = (lambda num: num + 1 if num % 2 != 0 else num)(batch)
+       batch_size_per_gpu = int(38400 / batch )
+    else:
+       batch_size_per_gpu = int(total_memory / 8)   
+       batch_size_per_gpu = (lambda num: num + 1 if num % 2 != 0 else num)(batch_size_per_gpu)
+       batch = batch_size_per_gpu
+
+    if batch_size_per_gpu<=0:batch_size_per_gpu=1
+
+    if samples<64:
+       max_samples = int(samples * 0.25) 
+    else:
+       max_samples = 64    
+    
+    num_warmup_updates = int(samples * 0.10) 
+    save_per_updates = int(samples * 0.25) 
+    last_per_steps =int(save_per_updates * 5)
+    
+    max_samples = (lambda num: num + 1 if num % 2 != 0 else num)(max_samples)
+    num_warmup_updates = (lambda num: num + 1 if num % 2 != 0 else num)(num_warmup_updates)
+    save_per_updates = (lambda num: num + 1 if num % 2 != 0 else num)(save_per_updates)
+    last_per_steps = (lambda num: num + 1 if num % 2 != 0 else num)(last_per_steps)
+
+    if finetune:learning_rate=1e-4
+    else:learning_rate=7.5e-5
+
+    return batch_size_per_gpu,max_samples,num_warmup_updates,save_per_updates,last_per_steps,samples,learning_rate
+
+def extract_and_save_ema_model(checkpoint_path: str, new_checkpoint_path: str) -> None:
+    try:
+        checkpoint = torch.load(checkpoint_path)
+        print("Original Checkpoint Keys:", checkpoint.keys())
+        
+        ema_model_state_dict = checkpoint.get('ema_model_state_dict', None)
+
+        if ema_model_state_dict is not None:
+            new_checkpoint = {'ema_model_state_dict': ema_model_state_dict}
+            torch.save(new_checkpoint, new_checkpoint_path)
+            return f"New checkpoint saved at: {new_checkpoint_path}"
+        else:
+            return "No 'ema_model_state_dict' found in the checkpoint."
+
+    except Exception as e:
+        return f"An error occurred: {e}"
+
+def vocab_check(project_name):
+    name_project = project_name + "_pinyin"
+    path_project = os.path.join(path_data, name_project)
+
+    file_metadata = os.path.join(path_project, "metadata.csv")
+   
+    file_vocab="data/Emilia_ZH_EN_pinyin/vocab.txt"
+    if os.path.isfile(file_vocab)==False:
+        return f"the file {file_vocab} not found !"
+    
+    with open(file_vocab,"r",encoding="utf-8") as f:
+         data=f.read()
+
+    vocab = data.split("\n")
+
+    if os.path.isfile(file_metadata)==False:
+        return f"the file {file_metadata} not found !"
+
+    with open(file_metadata,"r",encoding="utf-8") as f:
+         data=f.read()
+
+    miss_symbols=[]
+    miss_symbols_keep={}
+    for item in data.split("\n"):
+         sp=item.split("|")
+         if len(sp)!=2:continue
+         text=sp[1].lower().strip()
+
+         for t in text:
+             if (t in vocab)==False and (t in miss_symbols_keep)==False:
+                miss_symbols.append(t) 
+                miss_symbols_keep[t]=t
+
+            
+    if miss_symbols==[]:info ="You can train using your language !"     
+    else:info = f"The following symbols are missing in your language : {len(miss_symbols)}\n\n" + "\n".join(miss_symbols)   
+
+    return info
+
+
+
+with gr.Blocks() as app:
+
+    with gr.Row():
+         project_name=gr.Textbox(label="project name",value="my_speak")
+         bt_create=gr.Button("create new project")
+    
+    bt_create.click(fn=create_data_project,inputs=[project_name])
+
+    with gr.Tabs():
+         
+
+         with gr.TabItem("transcribe Data"):
+
+
+              ch_manual = gr.Checkbox(label="user",value=False)
+
+              mark_info_transcribe=gr.Markdown(
+                         """```plaintext    
+     Place your 'wavs' folder and 'metadata.csv' file in the {your_project_name}' directory. 
+                 
+     my_speak/
+     │
+     └── dataset/
+         ├── audio1.wav
+         └── audio2.wav
+         ...
+     ```""",visible=False)
+
+              audio_speaker = gr.File(label="voice",type="filepath",file_count="multiple")
+              txt_lang = gr.Text(label="Language",value="english")
+              bt_transcribe=bt_create=gr.Button("transcribe")
+              txt_info_transcribe=gr.Text(label="info",value="")
+              bt_transcribe.click(fn=transcribe_all,inputs=[project_name,audio_speaker,txt_lang,ch_manual],outputs=[txt_info_transcribe])
+              ch_manual.change(fn=check_user,inputs=[ch_manual],outputs=[audio_speaker,mark_info_transcribe])
+         
+         with gr.TabItem("prepare Data"):
+              gr.Markdown(
+                         """```plaintext    
+     place all your wavs folder and your metadata.csv file in {your name project}                                 
+     my_speak/
+     │
+     ├── wavs/
+     │   ├── audio1.wav
+     │   └── audio2.wav
+     |   ...
+     │
+     └── metadata.csv
+      
+     file format metadata.csv
+
+     audio1|text1
+     audio2|text1
+     ...
+
+     ```""")
+
+              bt_prepare=bt_create=gr.Button("prepare")
+              txt_info_prepare=gr.Text(label="info",value="")
+              bt_prepare.click(fn=create_metadata,inputs=[project_name],outputs=[txt_info_prepare])
+         
+         with gr.TabItem("train Data"):
+              
+              with gr.Row():    
+                   bt_calculate=bt_create=gr.Button("Auto Settings") 
+                   ch_finetune=bt_create=gr.Checkbox(label="finetune",value=True)   
+                   lb_samples = gr.Label(label="samples")  
+                   batch_size_type = gr.Radio(label="Batch Size Type", choices=["frame", "sample"], value="frame")  
+                
+              with gr.Row():     
+                   exp_name = gr.Radio(label="Model", choices=["F5TTS_Base", "E2TTS_Base"], value="F5TTS_Base")
+                   learning_rate = gr.Number(label="Learning Rate", value=1e-4, step=1e-4)
+                   
+              with gr.Row():
+                   batch_size_per_gpu = gr.Number(label="Batch Size per GPU", value=1000)
+                   max_samples = gr.Number(label="Max Samples", value=16)
+                   
+              with gr.Row():
+                   grad_accumulation_steps = gr.Number(label="Gradient Accumulation Steps", value=1)
+                   max_grad_norm = gr.Number(label="Max Gradient Norm", value=1.0)
+                   
+              with gr.Row():
+                   epochs = gr.Number(label="Epochs", value=10)
+                   num_warmup_updates = gr.Number(label="Warmup Updates", value=5)
+
+              with gr.Row():     
+                   save_per_updates = gr.Number(label="Save per Updates", value=10)
+                   last_per_steps = gr.Number(label="Last per Steps", value=50)
+       
+              with gr.Row():
+                   start_button = gr.Button("Start Training")
+                   stop_button = gr.Button("Stop Training",interactive=False)
+         
+              txt_info_train=gr.Text(label="info",value="")
+              start_button.click(fn=start_training,inputs=[project_name,exp_name,learning_rate,batch_size_per_gpu,batch_size_type,max_samples,grad_accumulation_steps,max_grad_norm,epochs,num_warmup_updates,save_per_updates,last_per_steps,ch_finetune],outputs=[txt_info_train,start_button,stop_button])
+              stop_button.click(fn=stop_training,outputs=[txt_info_train,start_button,stop_button])
+              bt_calculate.click(fn=calculate_train,inputs=[project_name,batch_size_type,max_samples,learning_rate,num_warmup_updates,save_per_updates,last_per_steps,ch_finetune],outputs=[batch_size_per_gpu,max_samples,num_warmup_updates,save_per_updates,last_per_steps,lb_samples,learning_rate]) 
+         
+         with gr.TabItem("reduse checkpoint"):
+              txt_path_checkpoint = gr.Text(label="path checkpoint :") 
+              txt_path_checkpoint_small = gr.Text(label="path output :") 
+              txt_info_reduse = gr.Text(label="info",value="")
+              reduse_button = gr.Button("reduse") 
+              reduse_button.click(fn=extract_and_save_ema_model,inputs=[txt_path_checkpoint,txt_path_checkpoint_small],outputs=[txt_info_reduse])
+
+         with gr.TabItem("vocab check experiment"):
+              check_button = gr.Button("check vocab") 
+              txt_info_check=gr.Text(label="info",value="")
+              check_button.click(fn=vocab_check,inputs=[project_name],outputs=[txt_info_check])
+       
+
+@click.command()
+@click.option("--port", "-p", default=None, type=int, help="Port to run the app on")
+@click.option("--host", "-H", default=None, help="Host to run the app on")
+@click.option(
+    "--share",
+    "-s",
+    default=False,
+    is_flag=True,
+    help="Share the app via Gradio share link",
+)
+@click.option("--api", "-a", default=True, is_flag=True, help="Allow API access")
+def main(port, host, share, api):
+    global app
+    print(f"Starting app...")
+    app.queue(api_open=api).launch(
+        server_name=host, server_port=port, share=share, show_api=api
+    )
+
+if __name__ == "__main__":
+    main()

gradio_app.py

@@ -0,0 +1,824 @@
+import os
+import re
+import torch
+import torchaudio
+import gradio as gr
+import numpy as np
+import tempfile
+from einops import rearrange
+from vocos import Vocos
+from pydub import AudioSegment, silence
+from model import CFM, UNetT, DiT, MMDiT
+from cached_path import cached_path
+from model.utils import (
+    load_checkpoint,
+    get_tokenizer,
+    convert_char_to_pinyin,
+    save_spectrogram,
+)
+from transformers import pipeline
+import librosa
+import click
+import soundfile as sf
+
+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
+
+
+
+SPLIT_WORDS = [
+    "but", "however", "nevertheless", "yet", "still",
+    "therefore", "thus", "hence", "consequently",
+    "moreover", "furthermore", "additionally",
+    "meanwhile", "alternatively", "otherwise",
+    "namely", "specifically", "for example", "such as",
+    "in fact", "indeed", "notably",
+    "in contrast", "on the other hand", "conversely",
+    "in conclusion", "to summarize", "finally"
+]
+
+device = (
+    "cuda"
+    if torch.cuda.is_available()
+    else "mps" if torch.backends.mps.is_available() else "cpu"
+)
+
+print(f"Using {device} device")
+
+pipe = pipeline(
+    "automatic-speech-recognition",
+    model="openai/whisper-large-v3-turbo",
+    torch_dtype=torch.float16,
+    device=device,
+)
+vocos = Vocos.from_pretrained("charactr/vocos-mel-24khz")
+
+# --------------------- Settings -------------------- #
+
+target_sample_rate = 24000
+n_mel_channels = 100
+hop_length = 256
+target_rms = 0.1
+nfe_step = 32  # 16, 32
+cfg_strength = 2.0
+ode_method = "euler"
+sway_sampling_coef = -1.0
+speed = 1.0
+# fix_duration = 27  # None or float (duration in seconds)
+fix_duration = None
+
+
+def load_model(repo_name, exp_name, model_cls, model_cfg, ckpt_step):
+    ckpt_path = str(cached_path(f"hf://SWivid/{repo_name}/{exp_name}/model_{ckpt_step}.safetensors"))
+    # ckpt_path = f"ckpts/{exp_name}/model_{ckpt_step}.pt"  # .pt | .safetensors
+    vocab_char_map, vocab_size = get_tokenizer("Emilia_ZH_EN", "pinyin")
+    model = CFM(
+        transformer=model_cls(
+            **model_cfg, text_num_embeds=vocab_size, mel_dim=n_mel_channels
+        ),
+        mel_spec_kwargs=dict(
+            target_sample_rate=target_sample_rate,
+            n_mel_channels=n_mel_channels,
+            hop_length=hop_length,
+        ),
+        odeint_kwargs=dict(
+            method=ode_method,
+        ),
+        vocab_char_map=vocab_char_map,
+    ).to(device)
+
+    model = load_checkpoint(model, ckpt_path, device, use_ema = True)
+
+    return model
+
+
+# load models
+F5TTS_model_cfg = dict(
+    dim=1024, depth=22, heads=16, ff_mult=2, text_dim=512, conv_layers=4
+)
+E2TTS_model_cfg = dict(dim=1024, depth=24, heads=16, ff_mult=4)
+
+F5TTS_ema_model = load_model(
+    "F5-TTS", "F5TTS_Base", DiT, F5TTS_model_cfg, 1200000
+)
+E2TTS_ema_model = load_model(
+    "E2-TTS", "E2TTS_Base", UNetT, E2TTS_model_cfg, 1200000
+)
+
+def split_text_into_batches(text, max_chars=200, split_words=SPLIT_WORDS):
+    if len(text.encode('utf-8')) <= max_chars:
+        return [text]
+    if text[-1] not in ['。', '.', '!', '！', '?', '？']:
+        text += '.'
+        
+    sentences = re.split('([。.!?！？])', text)
+    sentences = [''.join(i) for i in zip(sentences[0::2], sentences[1::2])]
+    
+    batches = []
+    current_batch = ""
+    
+    def split_by_words(text):
+        words = text.split()
+        current_word_part = ""
+        word_batches = []
+        for word in words:
+            if len(current_word_part.encode('utf-8')) + len(word.encode('utf-8')) + 1 <= max_chars:
+                current_word_part += word + ' '
+            else:
+                if current_word_part:
+                    # Try to find a suitable split word
+                    for split_word in split_words:
+                        split_index = current_word_part.rfind(' ' + split_word + ' ')
+                        if split_index != -1:
+                            word_batches.append(current_word_part[:split_index].strip())
+                            current_word_part = current_word_part[split_index:].strip() + ' '
+                            break
+                    else:
+                        # If no suitable split word found, just append the current part
+                        word_batches.append(current_word_part.strip())
+                        current_word_part = ""
+                current_word_part += word + ' '
+        if current_word_part:
+            word_batches.append(current_word_part.strip())
+        return word_batches
+
+    for sentence in sentences:
+        if len(current_batch.encode('utf-8')) + len(sentence.encode('utf-8')) <= max_chars:
+            current_batch += sentence
+        else:
+            # If adding this sentence would exceed the limit
+            if current_batch:
+                batches.append(current_batch)
+                current_batch = ""
+            
+            # If the sentence itself is longer than max_chars, split it
+            if len(sentence.encode('utf-8')) > max_chars:
+                # First, try to split by colon
+                colon_parts = sentence.split(':')
+                if len(colon_parts) > 1:
+                    for part in colon_parts:
+                        if len(part.encode('utf-8')) <= max_chars:
+                            batches.append(part)
+                        else:
+                            # If colon part is still too long, split by comma
+                            comma_parts = re.split('[,，]', part)
+                            if len(comma_parts) > 1:
+                                current_comma_part = ""
+                                for comma_part in comma_parts:
+                                    if len(current_comma_part.encode('utf-8')) + len(comma_part.encode('utf-8')) <= max_chars:
+                                        current_comma_part += comma_part + ','
+                                    else:
+                                        if current_comma_part:
+                                            batches.append(current_comma_part.rstrip(','))
+                                        current_comma_part = comma_part + ','
+                                if current_comma_part:
+                                    batches.append(current_comma_part.rstrip(','))
+                            else:
+                                # If no comma, split by words
+                                batches.extend(split_by_words(part))
+                else:
+                    # If no colon, split by comma
+                    comma_parts = re.split('[,，]', sentence)
+                    if len(comma_parts) > 1:
+                        current_comma_part = ""
+                        for comma_part in comma_parts:
+                            if len(current_comma_part.encode('utf-8')) + len(comma_part.encode('utf-8')) <= max_chars:
+                                current_comma_part += comma_part + ','
+                            else:
+                                if current_comma_part:
+                                    batches.append(current_comma_part.rstrip(','))
+                                current_comma_part = comma_part + ','
+                        if current_comma_part:
+                            batches.append(current_comma_part.rstrip(','))
+                    else:
+                        # If no comma, split by words
+                        batches.extend(split_by_words(sentence))
+            else:
+                current_batch = sentence
+    
+    if current_batch:
+        batches.append(current_batch)
+    
+    return batches
+
+def infer_batch(ref_audio, ref_text, gen_text_batches, exp_name, remove_silence, progress=gr.Progress()):
+    if exp_name == "F5-TTS":
+        ema_model = F5TTS_ema_model
+    elif exp_name == "E2-TTS":
+        ema_model = E2TTS_ema_model
+
+    audio, sr = ref_audio
+    if audio.shape[0] > 1:
+        audio = torch.mean(audio, dim=0, keepdim=True)
+
+    rms = torch.sqrt(torch.mean(torch.square(audio)))
+    if rms < target_rms:
+        audio = audio * target_rms / rms
+    if sr != target_sample_rate:
+        resampler = torchaudio.transforms.Resample(sr, target_sample_rate)
+        audio = resampler(audio)
+    audio = audio.to(device)
+
+    generated_waves = []
+    spectrograms = []
+
+    for i, gen_text in enumerate(progress.tqdm(gen_text_batches)):
+        # Prepare the text
+        if len(ref_text[-1].encode('utf-8')) == 1:
+            ref_text = ref_text + " "
+        text_list = [ref_text + gen_text]
+        final_text_list = convert_char_to_pinyin(text_list)
+
+        # Calculate duration
+        ref_audio_len = audio.shape[-1] // hop_length
+        zh_pause_punc = r"。，、；：？！"
+        ref_text_len = len(ref_text.encode('utf-8')) + 3 * len(re.findall(zh_pause_punc, ref_text))
+        gen_text_len = len(gen_text.encode('utf-8')) + 3 * len(re.findall(zh_pause_punc, gen_text))
+        duration = ref_audio_len + int(ref_audio_len / ref_text_len * gen_text_len / speed)
+
+        # inference
+        with torch.inference_mode():
+            generated, _ = ema_model.sample(
+                cond=audio,
+                text=final_text_list,
+                duration=duration,
+                steps=nfe_step,
+                cfg_strength=cfg_strength,
+                sway_sampling_coef=sway_sampling_coef,
+            )
+
+        generated = generated[:, ref_audio_len:, :]
+        generated_mel_spec = rearrange(generated, "1 n d -> 1 d n")
+        generated_wave = vocos.decode(generated_mel_spec.cpu())
+        if rms < target_rms:
+            generated_wave = generated_wave * rms / target_rms
+
+        # wav -> numpy
+        generated_wave = generated_wave.squeeze().cpu().numpy()
+        
+        generated_waves.append(generated_wave)
+        spectrograms.append(generated_mel_spec[0].cpu().numpy())
+
+    # Combine all generated waves
+    final_wave = np.concatenate(generated_waves)
+
+    # Remove silence
+    if remove_silence:
+        with tempfile.NamedTemporaryFile(delete=False, suffix=".wav") as f:
+            sf.write(f.name, final_wave, target_sample_rate)
+            aseg = AudioSegment.from_file(f.name)
+            non_silent_segs = silence.split_on_silence(aseg, min_silence_len=1000, silence_thresh=-50, keep_silence=500)
+            non_silent_wave = AudioSegment.silent(duration=0)
+            for non_silent_seg in non_silent_segs:
+                non_silent_wave += non_silent_seg
+            aseg = non_silent_wave
+            aseg.export(f.name, format="wav")
+            final_wave, _ = torchaudio.load(f.name)
+        final_wave = final_wave.squeeze().cpu().numpy()
+
+    # Create a combined spectrogram
+    combined_spectrogram = np.concatenate(spectrograms, axis=1)
+    
+    with tempfile.NamedTemporaryFile(suffix=".png", delete=False) as tmp_spectrogram:
+        spectrogram_path = tmp_spectrogram.name
+        save_spectrogram(combined_spectrogram, spectrogram_path)
+
+    return (target_sample_rate, final_wave), spectrogram_path
+
+def infer(ref_audio_orig, ref_text, gen_text, exp_name, remove_silence, custom_split_words=''):
+    if not custom_split_words.strip():
+        custom_words = [word.strip() for word in custom_split_words.split(',')]
+        global SPLIT_WORDS
+        SPLIT_WORDS = custom_words
+
+    print(gen_text)
+
+    gr.Info("Converting audio...")
+    with tempfile.NamedTemporaryFile(delete=False, suffix=".wav") as f:
+        aseg = AudioSegment.from_file(ref_audio_orig)
+
+        non_silent_segs = silence.split_on_silence(aseg, min_silence_len=1000, silence_thresh=-50, keep_silence=500)
+        non_silent_wave = AudioSegment.silent(duration=0)
+        for non_silent_seg in non_silent_segs:
+            non_silent_wave += non_silent_seg
+        aseg = non_silent_wave
+
+        audio_duration = len(aseg)
+        if audio_duration > 15000:
+            gr.Warning("Audio is over 15s, clipping to only first 15s.")
+            aseg = aseg[:15000]
+        aseg.export(f.name, format="wav")
+        ref_audio = f.name
+
+    if not ref_text.strip():
+        gr.Info("No reference text provided, transcribing reference audio...")
+        ref_text = pipe(
+            ref_audio,
+            chunk_length_s=30,
+            batch_size=128,
+            generate_kwargs={"task": "transcribe"},
+            return_timestamps=False,
+        )["text"].strip()
+        gr.Info("Finished transcription")
+    else:
+        gr.Info("Using custom reference text...")
+
+    # Split the input text into batches
+    audio, sr = torchaudio.load(ref_audio)
+    max_chars = int(len(ref_text.encode('utf-8')) / (audio.shape[-1] / sr) * (30 - audio.shape[-1] / sr))
+    gen_text_batches = split_text_into_batches(gen_text, max_chars=max_chars)
+    print('ref_text', ref_text)
+    for i, gen_text in enumerate(gen_text_batches):
+        print(f'gen_text {i}', gen_text)
+    
+    gr.Info(f"Generating audio using {exp_name} in {len(gen_text_batches)} batches")
+    return infer_batch((audio, sr), ref_text, gen_text_batches, exp_name, remove_silence)
+    
+def generate_podcast(script, speaker1_name, ref_audio1, ref_text1, speaker2_name, ref_audio2, ref_text2, exp_name, remove_silence):
+    # Split the script into speaker blocks
+    speaker_pattern = re.compile(f"^({re.escape(speaker1_name)}|{re.escape(speaker2_name)}):", re.MULTILINE)
+    speaker_blocks = speaker_pattern.split(script)[1:]  # Skip the first empty element
+    
+    generated_audio_segments = []
+    
+    for i in range(0, len(speaker_blocks), 2):
+        speaker = speaker_blocks[i]
+        text = speaker_blocks[i+1].strip()
+        
+        # Determine which speaker is talking
+        if speaker == speaker1_name:
+            ref_audio = ref_audio1
+            ref_text = ref_text1
+        elif speaker == speaker2_name:
+            ref_audio = ref_audio2
+            ref_text = ref_text2
+        else:
+            continue  # Skip if the speaker is neither speaker1 nor speaker2
+        
+        # Generate audio for this block
+        audio, _ = infer(ref_audio, ref_text, text, exp_name, remove_silence)
+        
+        # Convert the generated audio to a numpy array
+        sr, audio_data = audio
+        
+        # Save the audio data as a WAV file
+        with tempfile.NamedTemporaryFile(suffix=".wav", delete=False) as temp_file:
+            sf.write(temp_file.name, audio_data, sr)
+            audio_segment = AudioSegment.from_wav(temp_file.name)
+        
+        generated_audio_segments.append(audio_segment)
+        
+        # Add a short pause between speakers
+        pause = AudioSegment.silent(duration=500)  # 500ms pause
+        generated_audio_segments.append(pause)
+    
+    # Concatenate all audio segments
+    final_podcast = sum(generated_audio_segments)
+    
+    # Export the final podcast
+    with tempfile.NamedTemporaryFile(suffix=".wav", delete=False) as temp_file:
+        podcast_path = temp_file.name
+        final_podcast.export(podcast_path, format="wav")
+    
+    return podcast_path
+
+def parse_speechtypes_text(gen_text):
+    # Pattern to find (Emotion)
+    pattern = r'\((.*?)\)'
+
+    # Split the text by the pattern
+    tokens = re.split(pattern, gen_text)
+
+    segments = []
+
+    current_emotion = 'Regular'
+
+    for i in range(len(tokens)):
+        if i % 2 == 0:
+            # This is text
+            text = tokens[i].strip()
+            if text:
+                segments.append({'emotion': current_emotion, 'text': text})
+        else:
+            # This is emotion
+            emotion = tokens[i].strip()
+            current_emotion = emotion
+
+    return segments
+
+def update_speed(new_speed):
+    global speed
+    speed = new_speed
+    return f"Speed set to: {speed}"
+
+with gr.Blocks() as app_credits:
+    gr.Markdown("""
+# Credits
+
+* [mrfakename](https://github.com/fakerybakery) for the original [online demo](https://huggingface.co/spaces/mrfakename/E2-F5-TTS)
+* [RootingInLoad](https://github.com/RootingInLoad) for the podcast generation
+""")
+with gr.Blocks() as app_tts:
+    gr.Markdown("# Batched TTS")
+    ref_audio_input = gr.Audio(label="Reference Audio", type="filepath")
+    gen_text_input = gr.Textbox(label="Text to Generate", lines=10)
+    model_choice = gr.Radio(
+        choices=["F5-TTS", "E2-TTS"], label="Choose TTS Model", value="F5-TTS"
+    )
+    generate_btn = gr.Button("Synthesize", variant="primary")
+    with gr.Accordion("Advanced Settings", open=False):
+        ref_text_input = gr.Textbox(
+            label="Reference Text",
+            info="Leave blank to automatically transcribe the reference audio. If you enter text it will override automatic transcription.",
+            lines=2,
+        )
+        remove_silence = gr.Checkbox(
+            label="Remove Silences",
+            info="The model tends to produce silences, especially on longer audio. We can manually remove silences if needed. Note that this is an experimental feature and may produce strange results. This will also increase generation time.",
+            value=True,
+        )
+        split_words_input = gr.Textbox(
+            label="Custom Split Words",
+            info="Enter custom words to split on, separated by commas. Leave blank to use default list.",
+            lines=2,
+        )
+        speed_slider = gr.Slider(
+            label="Speed",
+            minimum=0.3,
+            maximum=2.0,
+            value=speed,
+            step=0.1,
+            info="Adjust the speed of the audio.",
+        )
+    speed_slider.change(update_speed, inputs=speed_slider)
+
+    audio_output = gr.Audio(label="Synthesized Audio")
+    spectrogram_output = gr.Image(label="Spectrogram")
+
+    generate_btn.click(
+        infer,
+        inputs=[
+            ref_audio_input,
+            ref_text_input,
+            gen_text_input,
+            model_choice,
+            remove_silence,
+            split_words_input,
+        ],
+        outputs=[audio_output, spectrogram_output],
+    )
+    
+with gr.Blocks() as app_podcast:
+    gr.Markdown("# Podcast Generation")
+    speaker1_name = gr.Textbox(label="Speaker 1 Name")
+    ref_audio_input1 = gr.Audio(label="Reference Audio (Speaker 1)", type="filepath")
+    ref_text_input1 = gr.Textbox(label="Reference Text (Speaker 1)", lines=2)
+    
+    speaker2_name = gr.Textbox(label="Speaker 2 Name")
+    ref_audio_input2 = gr.Audio(label="Reference Audio (Speaker 2)", type="filepath")
+    ref_text_input2 = gr.Textbox(label="Reference Text (Speaker 2)", lines=2)
+    
+    script_input = gr.Textbox(label="Podcast Script", lines=10, 
+                                placeholder="Enter the script with speaker names at the start of each block, e.g.:\nSean: How did you start studying...\n\nMeghan: I came to my interest in technology...\nIt was a long journey...\n\nSean: That's fascinating. Can you elaborate...")
+    
+    podcast_model_choice = gr.Radio(
+        choices=["F5-TTS", "E2-TTS"], label="Choose TTS Model", value="F5-TTS"
+    )
+    podcast_remove_silence = gr.Checkbox(
+        label="Remove Silences",
+        value=True,
+    )
+    generate_podcast_btn = gr.Button("Generate Podcast", variant="primary")
+    podcast_output = gr.Audio(label="Generated Podcast")
+
+    def podcast_generation(script, speaker1, ref_audio1, ref_text1, speaker2, ref_audio2, ref_text2, model, remove_silence):
+        return generate_podcast(script, speaker1, ref_audio1, ref_text1, speaker2, ref_audio2, ref_text2, model, remove_silence)
+
+    generate_podcast_btn.click(
+        podcast_generation,
+        inputs=[
+            script_input,
+            speaker1_name,
+            ref_audio_input1,
+            ref_text_input1,
+            speaker2_name,
+            ref_audio_input2,
+            ref_text_input2,
+            podcast_model_choice,
+            podcast_remove_silence,
+        ],
+        outputs=podcast_output,
+    )
+
+def parse_emotional_text(gen_text):
+    # Pattern to find (Emotion)
+    pattern = r'\((.*?)\)'
+
+    # Split the text by the pattern
+    tokens = re.split(pattern, gen_text)
+
+    segments = []
+
+    current_emotion = 'Regular'
+
+    for i in range(len(tokens)):
+        if i % 2 == 0:
+            # This is text
+            text = tokens[i].strip()
+            if text:
+                segments.append({'emotion': current_emotion, 'text': text})
+        else:
+            # This is emotion
+            emotion = tokens[i].strip()
+            current_emotion = emotion
+
+    return segments
+
+with gr.Blocks() as app_emotional:
+    # New section for emotional generation
+    gr.Markdown(
+        """
+    # Multiple Speech-Type Generation
+
+    This section allows you to upload different audio clips for each speech type. 'Regular' emotion is mandatory. You can add additional speech types by clicking the "Add Speech Type" button. Enter your text in the format shown below, and the system will generate speech using the appropriate emotions. If unspecified, the model will use the regular speech type. The current speech type will be used until the next speech type is specified.
+
+    **Example Input:**
+
+    (Regular) Hello, I'd like to order a sandwich please. (Surprised) What do you mean you're out of bread? (Sad) I really wanted a sandwich though... (Angry) You know what, darn you and your little shop, you suck! (Whisper) I'll just go back home and cry now. (Shouting) Why me?!
+    """
+    )
+
+    gr.Markdown("Upload different audio clips for each speech type. 'Regular' emotion is mandatory. You can add additional speech types by clicking the 'Add Speech Type' button.")
+
+    # Regular speech type (mandatory)
+    with gr.Row():
+        regular_name = gr.Textbox(value='Regular', label='Speech Type Name', interactive=False)
+        regular_audio = gr.Audio(label='Regular Reference Audio', type='filepath')
+        regular_ref_text = gr.Textbox(label='Reference Text (Regular)', lines=2)
+
+    # Additional speech types (up to 9 more)
+    max_speech_types = 10
+    speech_type_names = []
+    speech_type_audios = []
+    speech_type_ref_texts = []
+    speech_type_delete_btns = []
+
+    for i in range(max_speech_types - 1):
+        with gr.Row():
+            name_input = gr.Textbox(label='Speech Type Name', visible=False)
+            audio_input = gr.Audio(label='Reference Audio', type='filepath', visible=False)
+            ref_text_input = gr.Textbox(label='Reference Text', lines=2, visible=False)
+            delete_btn = gr.Button("Delete", variant="secondary", visible=False)
+        speech_type_names.append(name_input)
+        speech_type_audios.append(audio_input)
+        speech_type_ref_texts.append(ref_text_input)
+        speech_type_delete_btns.append(delete_btn)
+
+    # Button to add speech type
+    add_speech_type_btn = gr.Button("Add Speech Type")
+
+    # Keep track of current number of speech types
+    speech_type_count = gr.State(value=0)
+
+    # Function to add a speech type
+    def add_speech_type_fn(speech_type_count):
+        if speech_type_count < max_speech_types - 1:
+            speech_type_count += 1
+            # Prepare updates for the components
+            name_updates = []
+            audio_updates = []
+            ref_text_updates = []
+            delete_btn_updates = []
+            for i in range(max_speech_types - 1):
+                if i < speech_type_count:
+                    name_updates.append(gr.update(visible=True))
+                    audio_updates.append(gr.update(visible=True))
+                    ref_text_updates.append(gr.update(visible=True))
+                    delete_btn_updates.append(gr.update(visible=True))
+                else:
+                    name_updates.append(gr.update())
+                    audio_updates.append(gr.update())
+                    ref_text_updates.append(gr.update())
+                    delete_btn_updates.append(gr.update())
+        else:
+            # Optionally, show a warning
+            # gr.Warning("Maximum number of speech types reached.")
+            name_updates = [gr.update() for _ in range(max_speech_types - 1)]
+            audio_updates = [gr.update() for _ in range(max_speech_types - 1)]
+            ref_text_updates = [gr.update() for _ in range(max_speech_types - 1)]
+            delete_btn_updates = [gr.update() for _ in range(max_speech_types - 1)]
+        return [speech_type_count] + name_updates + audio_updates + ref_text_updates + delete_btn_updates
+
+    add_speech_type_btn.click(
+        add_speech_type_fn,
+        inputs=speech_type_count,
+        outputs=[speech_type_count] + speech_type_names + speech_type_audios + speech_type_ref_texts + speech_type_delete_btns
+    )
+
+    # Function to delete a speech type
+    def make_delete_speech_type_fn(index):
+        def delete_speech_type_fn(speech_type_count):
+            # Prepare updates
+            name_updates = []
+            audio_updates = []
+            ref_text_updates = []
+            delete_btn_updates = []
+
+            for i in range(max_speech_types - 1):
+                if i == index:
+                    name_updates.append(gr.update(visible=False, value=''))
+                    audio_updates.append(gr.update(visible=False, value=None))
+                    ref_text_updates.append(gr.update(visible=False, value=''))
+                    delete_btn_updates.append(gr.update(visible=False))
+                else:
+                    name_updates.append(gr.update())
+                    audio_updates.append(gr.update())
+                    ref_text_updates.append(gr.update())
+                    delete_btn_updates.append(gr.update())
+
+            speech_type_count = max(0, speech_type_count - 1)
+
+            return [speech_type_count] + name_updates + audio_updates + ref_text_updates + delete_btn_updates
+
+        return delete_speech_type_fn
+
+    for i, delete_btn in enumerate(speech_type_delete_btns):
+        delete_fn = make_delete_speech_type_fn(i)
+        delete_btn.click(
+            delete_fn,
+            inputs=speech_type_count,
+            outputs=[speech_type_count] + speech_type_names + speech_type_audios + speech_type_ref_texts + speech_type_delete_btns
+        )
+
+    # Text input for the prompt
+    gen_text_input_emotional = gr.Textbox(label="Text to Generate", lines=10)
+
+    # Model choice
+    model_choice_emotional = gr.Radio(
+        choices=["F5-TTS", "E2-TTS"], label="Choose TTS Model", value="F5-TTS"
+    )
+
+    with gr.Accordion("Advanced Settings", open=False):
+        remove_silence_emotional = gr.Checkbox(
+            label="Remove Silences",
+            value=True,
+        )
+
+    # Generate button
+    generate_emotional_btn = gr.Button("Generate Emotional Speech", variant="primary")
+
+    # Output audio
+    audio_output_emotional = gr.Audio(label="Synthesized Audio")
+
+    def generate_emotional_speech(
+        regular_audio,
+        regular_ref_text,
+        gen_text,
+        *args,
+    ):
+        num_additional_speech_types = max_speech_types - 1
+        speech_type_names_list = args[:num_additional_speech_types]
+        speech_type_audios_list = args[num_additional_speech_types:2 * num_additional_speech_types]
+        speech_type_ref_texts_list = args[2 * num_additional_speech_types:3 * num_additional_speech_types]
+        model_choice = args[3 * num_additional_speech_types]
+        remove_silence = args[3 * num_additional_speech_types + 1]
+
+        # Collect the speech types and their audios into a dict
+        speech_types = {'Regular': {'audio': regular_audio, 'ref_text': regular_ref_text}}
+
+        for name_input, audio_input, ref_text_input in zip(speech_type_names_list, speech_type_audios_list, speech_type_ref_texts_list):
+            if name_input and audio_input:
+                speech_types[name_input] = {'audio': audio_input, 'ref_text': ref_text_input}
+
+        # Parse the gen_text into segments
+        segments = parse_speechtypes_text(gen_text)
+
+        # For each segment, generate speech
+        generated_audio_segments = []
+        current_emotion = 'Regular'
+
+        for segment in segments:
+            emotion = segment['emotion']
+            text = segment['text']
+
+            if emotion in speech_types:
+                current_emotion = emotion
+            else:
+                # If emotion not available, default to Regular
+                current_emotion = 'Regular'
+
+            ref_audio = speech_types[current_emotion]['audio']
+            ref_text = speech_types[current_emotion].get('ref_text', '')
+
+            # Generate speech for this segment
+            audio, _ = infer(ref_audio, ref_text, text, model_choice, remove_silence, "")
+            sr, audio_data = audio
+
+            generated_audio_segments.append(audio_data)
+
+        # Concatenate all audio segments
+        if generated_audio_segments:
+            final_audio_data = np.concatenate(generated_audio_segments)
+            return (sr, final_audio_data)
+        else:
+            gr.Warning("No audio generated.")
+            return None
+
+    generate_emotional_btn.click(
+        generate_emotional_speech,
+        inputs=[
+            regular_audio,
+            regular_ref_text,
+            gen_text_input_emotional,
+        ] + speech_type_names + speech_type_audios + speech_type_ref_texts + [
+            model_choice_emotional,
+            remove_silence_emotional,
+        ],
+        outputs=audio_output_emotional,
+    )
+
+    # Validation function to disable Generate button if speech types are missing
+    def validate_speech_types(
+        gen_text,
+        regular_name,
+        *args
+    ):
+        num_additional_speech_types = max_speech_types - 1
+        speech_type_names_list = args[:num_additional_speech_types]
+
+        # Collect the speech types names
+        speech_types_available = set()
+        if regular_name:
+            speech_types_available.add(regular_name)
+        for name_input in speech_type_names_list:
+            if name_input:
+                speech_types_available.add(name_input)
+
+        # Parse the gen_text to get the speech types used
+        segments = parse_emotional_text(gen_text)
+        speech_types_in_text = set(segment['emotion'] for segment in segments)
+
+        # Check if all speech types in text are available
+        missing_speech_types = speech_types_in_text - speech_types_available
+
+        if missing_speech_types:
+            # Disable the generate button
+            return gr.update(interactive=False)
+        else:
+            # Enable the generate button
+            return gr.update(interactive=True)
+
+    gen_text_input_emotional.change(
+        validate_speech_types,
+        inputs=[gen_text_input_emotional, regular_name] + speech_type_names,
+        outputs=generate_emotional_btn
+    )
+with gr.Blocks() as app:
+    gr.Markdown(
+        """
+# E2/F5 TTS
+
+This is a local web UI for 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 support English and Chinese.
+
+If you're having issues, try converting your reference audio to WAV or MP3, clipping it to 15s, and shortening your prompt.
+
+**NOTE: Reference text will be automatically transcribed with Whisper if not provided. For best results, keep your reference clips short (<15s). Ensure the audio is fully uploaded before generating.**
+"""
+    )
+    gr.TabbedInterface([app_tts, app_podcast, app_emotional, app_credits], ["TTS", "Podcast", "Multi-Style", "Credits"])
+
+@click.command()
+@click.option("--port", "-p", default=None, type=int, help="Port to run the app on")
+@click.option("--host", "-H", default=None, help="Host to run the app on")
+@click.option(
+    "--share",
+    "-s",
+    default=False,
+    is_flag=True,
+    help="Share the app via Gradio share link",
+)
+@click.option("--api", "-a", default=True, is_flag=True, help="Allow API access")
+def main(port, host, share, api):
+    global app
+    print(f"Starting app...")
+    app.queue(api_open=api).launch(
+        server_name=host, server_port=port, share=share, show_api=api
+    )
+
+
+if __name__ == "__main__":
+    main()

inference-cli.py

@@ -0,0 +1,428 @@
+import argparse
+import codecs
+import re
+import tempfile
+from pathlib import Path
+
+import numpy as np
+import soundfile as sf
+import tomli
+import torch
+import torchaudio
+import tqdm
+from cached_path import cached_path
+from einops import rearrange
+from pydub import AudioSegment, silence
+from transformers import pipeline
+from vocos import Vocos
+
+from model import CFM, DiT, MMDiT, UNetT
+from model.utils import (convert_char_to_pinyin, get_tokenizer,
+                         load_checkpoint, save_spectrogram)
+
+parser = argparse.ArgumentParser(
+    prog="python3 inference-cli.py",
+    description="Commandline interface for E2/F5 TTS with Advanced Batch Processing.",
+    epilog="Specify  options above  to override  one or more settings from config.",
+)
+parser.add_argument(
+    "-c",
+    "--config",
+    help="Configuration file. Default=cli-config.toml",
+    default="inference-cli.toml",
+)
+parser.add_argument(
+    "-m",
+    "--model",
+    help="F5-TTS | E2-TTS",
+)
+parser.add_argument(
+    "-p",
+    "--ckpt_file",
+    help="The Checkpoint .pt",
+)
+parser.add_argument(
+    "-v",
+    "--vocab_file",
+    help="The vocab .txt",
+)
+parser.add_argument(
+    "-r",
+    "--ref_audio",
+    type=str,
+    help="Reference audio file < 15 seconds."
+)
+parser.add_argument(
+    "-s",
+    "--ref_text",
+    type=str,
+    default="666",
+    help="Subtitle for the reference audio."
+)
+parser.add_argument(
+    "-t",
+    "--gen_text",
+    type=str,
+    help="Text to generate.",
+)
+parser.add_argument(
+    "-f",
+    "--gen_file",
+    type=str,
+    help="File with text to generate. Ignores --text",
+)
+parser.add_argument(
+    "-o",
+    "--output_dir",
+    type=str,
+    help="Path to output folder..",
+)
+parser.add_argument(
+    "--remove_silence",
+    help="Remove silence.",
+)
+parser.add_argument(
+    "--load_vocoder_from_local",
+    action="store_true",
+    help="load vocoder from local. Default: ../checkpoints/charactr/vocos-mel-24khz",
+)
+args = parser.parse_args()
+
+config = tomli.load(open(args.config, "rb"))
+
+ref_audio = args.ref_audio if args.ref_audio else config["ref_audio"]
+ref_text = args.ref_text if args.ref_text != "666" else config["ref_text"]
+gen_text = args.gen_text if args.gen_text else config["gen_text"]
+gen_file = args.gen_file if args.gen_file else config["gen_file"]
+if gen_file:
+    gen_text = codecs.open(gen_file, "r", "utf-8").read()
+output_dir = args.output_dir if args.output_dir else config["output_dir"]
+model = args.model if args.model else config["model"]
+ckpt_file = args.ckpt_file if args.ckpt_file else ""
+vocab_file = args.vocab_file if args.vocab_file else ""
+remove_silence = args.remove_silence if args.remove_silence else config["remove_silence"]
+wave_path = Path(output_dir)/"out.wav"
+spectrogram_path = Path(output_dir)/"out.png"
+vocos_local_path = "../checkpoints/charactr/vocos-mel-24khz"
+
+device = (
+    "cuda"
+    if torch.cuda.is_available()
+    else "mps" if torch.backends.mps.is_available() else "cpu"
+)
+
+if args.load_vocoder_from_local:
+    print(f"Load vocos from local path {vocos_local_path}")
+    vocos = Vocos.from_hparams(f"{vocos_local_path}/config.yaml")
+    state_dict = torch.load(f"{vocos_local_path}/pytorch_model.bin", map_location=device)
+    vocos.load_state_dict(state_dict)
+    vocos.eval()
+else:
+    print("Donwload Vocos from huggingface charactr/vocos-mel-24khz")
+    vocos = Vocos.from_pretrained("charactr/vocos-mel-24khz")
+
+print(f"Using {device} device")
+
+# --------------------- Settings -------------------- #
+
+target_sample_rate = 24000
+n_mel_channels = 100
+hop_length = 256
+target_rms = 0.1
+nfe_step = 32  # 16, 32
+cfg_strength = 2.0
+ode_method = "euler"
+sway_sampling_coef = -1.0
+speed = 1.0
+# fix_duration = 27  # None or float (duration in seconds)
+fix_duration = None
+
+def load_model(model_cls, model_cfg, ckpt_path,file_vocab):
+    
+    if file_vocab=="":
+        file_vocab="Emilia_ZH_EN"
+        tokenizer="pinyin"
+    else:
+        tokenizer="custom"
+
+    print("\nvocab : ",vocab_file,tokenizer) 
+    print("tokenizer : ",tokenizer) 
+    print("model : ",ckpt_path,"\n")    
+
+    vocab_char_map, vocab_size = get_tokenizer(file_vocab, tokenizer)
+    model = CFM(
+        transformer=model_cls(
+            **model_cfg, text_num_embeds=vocab_size, mel_dim=n_mel_channels
+        ),
+        mel_spec_kwargs=dict(
+            target_sample_rate=target_sample_rate,
+            n_mel_channels=n_mel_channels,
+            hop_length=hop_length,
+        ),
+        odeint_kwargs=dict(
+            method=ode_method,
+        ),
+        vocab_char_map=vocab_char_map,
+    ).to(device)
+
+    model = load_checkpoint(model, ckpt_path, device, use_ema = True)
+
+    return model
+
+# load models
+F5TTS_model_cfg = dict(
+    dim=1024, depth=22, heads=16, ff_mult=2, text_dim=512, conv_layers=4
+)
+E2TTS_model_cfg = dict(dim=1024, depth=24, heads=16, ff_mult=4)
+
+def 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
+
+    #ckpt_path = f"ckpts/{exp_name}/model_{ckpt_step}.pt" # .pt | .safetensors
+    #if not Path(ckpt_path).exists():
+        #ckpt_path = str(cached_path(f"hf://SWivid/{repo_name}/{exp_name}/model_{ckpt_step}.safetensors"))
+
+def infer_batch(ref_audio, ref_text, gen_text_batches, model,ckpt_file,file_vocab, remove_silence, cross_fade_duration=0.15):
+    if model == "F5-TTS":
+
+        if ckpt_file == "": 
+           repo_name= "F5-TTS"
+           exp_name = "F5TTS_Base"
+           ckpt_step= 1200000
+           ckpt_file = str(cached_path(f"hf://SWivid/{repo_name}/{exp_name}/model_{ckpt_step}.safetensors"))
+
+        ema_model = load_model(DiT, F5TTS_model_cfg, ckpt_file,file_vocab)
+
+    elif model == "E2-TTS":
+        if ckpt_file == "": 
+           repo_name= "E2-TTS"
+           exp_name = "E2TTS_Base"
+           ckpt_step= 1200000
+           ckpt_file = str(cached_path(f"hf://SWivid/{repo_name}/{exp_name}/model_{ckpt_step}.safetensors"))
+        
+        ema_model = load_model(UNetT, E2TTS_model_cfg, ckpt_file,file_vocab)
+
+    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 + " "
+    for i, gen_text in enumerate(tqdm.tqdm(gen_text_batches)):
+        # Prepare the text
+        text_list = [ref_text + gen_text]
+        final_text_list = convert_char_to_pinyin(text_list)
+
+        # Calculate duration
+        ref_audio_len = audio.shape[-1] // hop_length
+        zh_pause_punc = r"。，、；：？！"
+        ref_text_len = len(ref_text.encode('utf-8')) + 3 * len(re.findall(zh_pause_punc, ref_text))
+        gen_text_len = len(gen_text.encode('utf-8')) + 3 * len(re.findall(zh_pause_punc, gen_text))
+        duration = ref_audio_len + int(ref_audio_len / ref_text_len * gen_text_len / speed)
+
+        # inference
+        with torch.inference_mode():
+            generated, _ = ema_model.sample(
+                cond=audio,
+                text=final_text_list,
+                duration=duration,
+                steps=nfe_step,
+                cfg_strength=cfg_strength,
+                sway_sampling_coef=sway_sampling_coef,
+            )
+
+        generated = generated[:, ref_audio_len:, :]
+        generated_mel_spec = rearrange(generated, "1 n d -> 1 d n")
+        generated_wave = vocos.decode(generated_mel_spec.cpu())
+        if rms < target_rms:
+            generated_wave = generated_wave * rms / target_rms
+
+        # wav -> numpy
+        generated_wave = generated_wave.squeeze().cpu().numpy()
+        
+        generated_waves.append(generated_wave)
+        spectrograms.append(generated_mel_spec[0].cpu().numpy())
+
+    # Combine all generated waves with cross-fading
+    if cross_fade_duration <= 0:
+        # Simply concatenate
+        final_wave = np.concatenate(generated_waves)
+    else:
+        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)
+
+    return final_wave, combined_spectrogram
+
+def process_voice(ref_audio_orig, ref_text):
+    print("Converting audio...")
+    with tempfile.NamedTemporaryFile(delete=False, suffix=".wav") as f:
+        aseg = AudioSegment.from_file(ref_audio_orig)
+
+        non_silent_segs = silence.split_on_silence(aseg, min_silence_len=1000, silence_thresh=-50, keep_silence=1000)
+        non_silent_wave = AudioSegment.silent(duration=0)
+        for non_silent_seg in non_silent_segs:
+            non_silent_wave += non_silent_seg
+        aseg = non_silent_wave
+
+        audio_duration = len(aseg)
+        if audio_duration > 15000:
+            print("Audio is over 15s, clipping to only first 15s.")
+            aseg = aseg[:15000]
+        aseg.export(f.name, format="wav")
+        ref_audio = f.name
+
+    if not ref_text.strip():
+        print("No reference text provided, transcribing reference audio...")
+        pipe = pipeline(
+            "automatic-speech-recognition",
+            model="openai/whisper-large-v3-turbo",
+            torch_dtype=torch.float16,
+            device=device,
+        )
+        ref_text = pipe(
+            ref_audio,
+            chunk_length_s=30,
+            batch_size=128,
+            generate_kwargs={"task": "transcribe"},
+            return_timestamps=False,
+        )["text"].strip()
+        print("Finished transcription")
+    else:
+        print("Using custom reference text...")
+    return ref_audio, ref_text    
+
+def infer(ref_audio, ref_text, gen_text, model,ckpt_file,file_vocab, remove_silence, cross_fade_duration=0.15):
+    print(gen_text)
+    # Add the functionality to ensure it ends with ". "
+    if not ref_text.endswith(". ") and not ref_text.endswith("。"):
+        if ref_text.endswith("."):
+            ref_text += " "
+        else:
+            ref_text += ". "
+
+    # Split the input text into batches
+    audio, sr = torchaudio.load(ref_audio)
+    max_chars = int(len(ref_text.encode('utf-8')) / (audio.shape[-1] / sr) * (25 - audio.shape[-1] / sr))
+    gen_text_batches = chunk_text(gen_text, max_chars=max_chars)
+    print('ref_text', ref_text)
+    for i, gen_text in enumerate(gen_text_batches):
+        print(f'gen_text {i}', gen_text)
+    
+    print(f"Generating audio using {model} in {len(gen_text_batches)} batches, loading models...")
+    return infer_batch((audio, sr), ref_text, gen_text_batches, model,ckpt_file,file_vocab, remove_silence, cross_fade_duration)
+    
+
+def process(ref_audio, ref_text, text_gen, model,ckpt_file,file_vocab, remove_silence):
+    main_voice = {"ref_audio":ref_audio, "ref_text":ref_text}
+    if "voices" not in config:
+        voices = {"main": main_voice}
+    else:
+        voices = config["voices"]
+        voices["main"] = main_voice
+    for voice in voices:
+        voices[voice]['ref_audio'], voices[voice]['ref_text'] = process_voice(voices[voice]['ref_audio'], voices[voice]['ref_text'])
+
+    generated_audio_segments = []
+    reg1 = r'(?=\[\w+\])'
+    chunks = re.split(reg1, text_gen)
+    reg2 = r'\[(\w+)\]'
+    for text in chunks:
+        match = re.match(reg2, text)
+        if not match or voice not in voices:
+            voice = "main"
+        else:
+            voice = match[1]
+        text = re.sub(reg2, "", text)
+        gen_text = text.strip()
+        ref_audio = voices[voice]['ref_audio']
+        ref_text = voices[voice]['ref_text']
+        print(f"Voice: {voice}")
+        audio, spectragram = infer(ref_audio, ref_text, gen_text, model,ckpt_file,file_vocab, remove_silence)
+        generated_audio_segments.append(audio)
+
+    if generated_audio_segments:
+        final_wave = np.concatenate(generated_audio_segments)
+        with open(wave_path, "wb") as f:
+            sf.write(f.name, final_wave, target_sample_rate)
+            # Remove silence
+            if remove_silence:
+                aseg = AudioSegment.from_file(f.name)
+                non_silent_segs = silence.split_on_silence(aseg, min_silence_len=1000, silence_thresh=-50, keep_silence=500)
+                non_silent_wave = AudioSegment.silent(duration=0)
+                for non_silent_seg in non_silent_segs:
+                    non_silent_wave += non_silent_seg
+                aseg = non_silent_wave
+                aseg.export(f.name, format="wav")
+            print(f.name)
+
+
+process(ref_audio, ref_text, gen_text, model,ckpt_file,vocab_file, remove_silence)

inference-cli.toml

@@ -0,0 +1,10 @@
+# F5-TTS | E2-TTS
+model = "F5-TTS"
+ref_audio = "tests/ref_audio/test_en_1_ref_short.wav"
+# If an empty "", transcribes the reference audio automatically.
+ref_text = "Some call me nature, others call me mother nature."
+gen_text = "I don't really care what you call me. I've been a silent spectator, watching species evolve, empires rise and fall. But always remember, I am mighty and enduring. Respect me and I'll nurture you; ignore me and you shall face the consequences."
+# File with text to generate. Ignores the text above.
+gen_file = ""
+remove_silence = false
+output_dir = "tests"

requirements.txt

@@ -0,0 +1,23 @@
+accelerate>=0.33.0
+cached_path
+click
+datasets
+einops>=0.8.0
+einx>=0.3.0
+ema_pytorch>=0.5.2
+gradio
+jieba
+librosa
+matplotlib
+numpy<=1.26.4
+pydub
+pypinyin
+safetensors
+soundfile
+tomli
+torchdiffeq
+tqdm>=4.65.0
+transformers
+vocos
+wandb
+x_transformers>=1.31.14

requirements_eval.txt

@@ -0,0 +1,5 @@
+faster_whisper
+funasr
+jiwer
+zhconv
+zhon

speech_edit.py

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

train.py

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