app.py

import nltk
nltk.download('punkt_tab')
from sentence_analyzer import SentenceAnalyzer
import re
import tempfile
from collections import OrderedDict
from importlib.resources import files
import click
import gradio as gr
import numpy as np
import soundfile as sf
import torchaudio
import torch
from cached_path import cached_path
from transformers import AutoModelForCausalLM, AutoTokenizer

try:
    import spaces
    USING_SPACES = True
except ImportError:
    USING_SPACES = False

def gpu_decorator(func):
    if USING_SPACES:
        return spaces.GPU(func)
    else:
        return func

# Importando a nova API F5TTS
from f5_tts.api import F5TTS

import os
from huggingface_hub import hf_hub_download

def load_f5tts():
    # Carrega o caminho do repositório e o nome do arquivo das variáveis de ambiente
    repo_id = os.getenv("MODEL_REPO_ID", "SWivid/F5-TTS/F5TTS_Base")
    filename = os.getenv("MODEL_FILENAME", "model_1200000.safetensors")
    token = os.getenv("HUGGINGFACE_TOKEN")
    # Valida se o token está presente
    if not token:
        raise ValueError("A variável de ambiente 'HUGGINGFACE_TOKEN' não foi definida.")
    # Faz o download do modelo do repositório privado
    ckpt_path = hf_hub_download(repo_id=repo_id, filename=filename, use_auth_token=token)

    # Define as configurações do modelo (ajuste se necessário)
    F5TTS_model_cfg = dict(dim=1024, depth=22, heads=16, ff_mult=2, text_dim=512, conv_layers=4)

    # Retorna a instância da API F5TTS
    return F5TTS(
        model_type="F5-TTS",  # Ajuste o nome do modelo se necessário
        ckpt_file=ckpt_path,
        vocab_file=os.path.join(os.path.dirname(ckpt_path), "vocab.txt"), # Caminho para o arquivo vocab.txt
        device="cuda" if torch.cuda.is_available() else "cpu", # Define o dispositivo
        use_ema=True
    )

# Carregar modelo F5TTS usando a nova API
F5TTS_ema_model = load_f5tts()

# Variáveis globais para o cache
last_checkpoint = None
last_device = None
last_ema = None
tts_api = None
training_process = None  # Adicione esta linha se necessário para o seu contexto

@gpu_decorator
def infer(
    ref_audio_orig, ref_text, gen_text, remove_silence, cross_fade_duration=0.15, speed=1, nfe=32, show_info=gr.Info, seed=-1 
):
    print(nfe)
    ref_audio, ref_text = preprocess_ref_audio_text(ref_audio_orig, ref_text, show_info=show_info)
    ema_model = F5TTS_ema_model
    final_wave, final_sample_rate, combined_spectrogram = infer_process(
        ref_audio,
        ref_text.lower().strip(),
        gen_text.lower().strip(),
        ema_model,
        vocoder,
        cross_fade_duration=cross_fade_duration,
        nfe_step=nfe,
        speed=speed,
        show_info=show_info,
        progress=gr.Progress(),
        seed=seed  # Passando o seed para infer_process
    )
    # Remover silêncios
    if remove_silence:
        with tempfile.NamedTemporaryFile(delete=False, suffix=".wav") as f:
            sf.write(f.name, final_wave, final_sample_rate)
            remove_silence_for_generated_wav(f.name)
            final_wave, _ = torchaudio.load(f.name)
        final_wave = final_wave.squeeze().cpu().numpy()
    # Salvar espectrograma
    with tempfile.NamedTemporaryFile(suffix=".png", delete=False) as tmp_spectrogram:
        spectrogram_path = tmp_spectrogram.name
        save_spectrogram(combined_spectrogram, spectrogram_path)
    return (final_sample_rate, final_wave), spectrogram_path, ref_text, seed # Retornando o seed


# Estilos CSS
custom_css = """
#sentences-container {
    border: 1px solid #ddd;
    border-radius: 4px;
    padding: 10px;
    margin-bottom: 10px;
}
.sentence-box {
    border: 1px solid #eee;
    padding: 5px;
    margin-bottom: 5px;
    border-radius: 4px;
    background-color: #f9f9f9;
}
"""

with gr.Blocks(css=custom_css) as app:
    with gr.Tabs():
        with gr.Tab("TTS Básico"):
            gr.Markdown("# TTS Básico com F5-TTS")

            # Entradas básicas
            ref_audio_input = gr.Audio(label="Áudio de Referência", type="filepath")
            gen_text_input = gr.Textbox(label="Texto para Gerar", lines=10)
            generate_btn = gr.Button("Sintetizar", variant="primary")

            # Configurações avançadas
            gr.Markdown("### Configurações Avançadas")
            with gr.Accordion("Expandir Configurações Avançadas", open=False):
                ref_text_input = gr.Textbox(
                    label="Texto de Referência",
                    info="Deixe em branco para transcrever automaticamente o áudio de referência. Se você inserir texto, ele substituirá a transcrição automática.",
                    lines=2,
                )
                remove_silence = gr.Checkbox(
                    label="Remover Silêncios",
                    info="O modelo tende a produzir silêncios, especialmente em áudios mais longos. Podemos remover manualmente os silêncios, se necessário. Isso também aumentará o tempo de geração.",
                    value=False,
                )
                speed_slider = gr.Slider(
                    label="Velocidade",
                    minimum=0.3,
                    maximum=2.0,
                    value=1.0,
                    step=0.1,
                    info="Ajuste a velocidade do áudio.",
                )
                cross_fade_duration_slider = gr.Slider(
                    label="Duração do Cross-fade (s)",
                    minimum=0.0,
                    maximum=1.0,
                    value=0.15,
                    step=0.01,
                    info="Defina a duração do cross-fade entre os clipes de áudio.",
                )
                chunk_size_slider = gr.Slider(
                    label="Número de Sentenças por Chunk",
                    minimum=1,
                    maximum=10,
                    value=1,
                    step=1,
                    info="Defina quantas sentenças serão processadas em cada chunk.",
                )
                nfe_slider = gr.Slider(
                    label="NFE",
                    minimum=16,
                    maximum=64,
                    value=32,
                    step=1,
                    info="Ajuste NFE Step.",
                )
                seed_input = gr.Number(label="Seed", value=-1, minimum=-1)  # Seed na seção avançada

            analyzer = SentenceAnalyzer()

            @gpu_decorator
            def process_chunks(
                ref_audio_input,
                ref_text_input,
                gen_text_input,
                remove_silence,
                cross_fade_duration_slider,
                speed_slider,
                nfe_slider,
                chunk_size_slider,
                seed_input,  # Passando o seed para process_chunks
            ):
                # Dividir o texto em sentenças
                sentences = analyzer.split_into_sentences(gen_text_input)

                # Agrupar sentenças em chunks
                chunks = [
                    " ".join(sentences[i : i + chunk_size_slider])
                    for i in range(0, len(sentences), chunk_size_slider)
                ]

                # Processar cada chunk
                audio_segments = []
                for chunk in chunks:
                    audio_out, spectrogram_path, ref_text_out, seed_output = infer(  # Recebendo o seed de infer
                        ref_audio_input,
                        ref_text_input,  # Utiliza o Texto de Referência como está
                        chunk,  # Processa o chunk atual
                        remove_silence,
                        cross_fade_duration_slider,
                        speed_slider,
                        nfe_slider,
                        seed=seed_input,  # Passando o seed para infer
                    )
                    sr, audio_data = audio_out
                    audio_segments.append(audio_data)

                # Concatenar os segmentos de áudio gerados
                if audio_segments:
                    final_audio_data = np.concatenate(audio_segments)
                    return (
                        (sr, final_audio_data),  # Áudio final
                        spectrogram_path,  # Espectrograma
                        gr.update(value=ref_text_out),  # Nenhuma mudança no Texto de Referência
                        seed_output  # Retornando o seed
                    )
                else:
                    gr.Warning("Nenhum áudio gerado.")
                    return None, None, gr.update(), None  # Retornando None para o seed

            # Saídas
            gr.Markdown("### Resultados")
            audio_output = gr.Audio(label="Áudio Sintetizado")
            spectrogram_output = gr.Image(label="Espectrograma")
            seed_output = gr.Text(label="Seed usada:")  # Saída do Seed

            # Associação do botão `generate_btn` à função `process_chunks`
            generate_btn.click(
                process_chunks,
                inputs=[
                    ref_audio_input,
                    ref_text_input,
                    gen_text_input,
                    remove_silence,
                    cross_fade_duration_slider,
                    speed_slider,
                    nfe_slider,
                    chunk_size_slider,
                    seed_input,  # Passando o seed como entrada
                ],
                outputs=[
                    audio_output,
                    spectrogram_output,
                    ref_text_input,  # Atualiza o texto de referência, se necessário
                    seed_output,  # Saída do Seed
                ],
            )


    # Código para iniciar a aplicação Gradio
    @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("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()