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
from f5_tts.infer.utils_infer import preprocess_ref_audio_text

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="/home/user/app/data/Emilia_ZH_EN_pinyin/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

# Estilos CSS com tema escuro
custom_css = """
#sentences-container {
    border: 1px solid #555;
    border-radius: 4px;
    padding: 10px;
    margin-bottom: 10px;
}
.sentence-box {
    border: 1px solid #333;
    padding: 5px;
    margin-bottom: 5px;
    border-radius: 4px;
    background-color: #222;
    color: #eee;
}
body {
    background-color: #111;
    color: #eee;
}
.gradio-container {
    background-color: #111;
}
.dark {
    background-color: #333;
    color: #eee;
}
"""

with gr.Blocks(css=custom_css, theme=gr.themes.Default(primary_hue="gray", secondary_hue="gray")) 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
                f5tts_model,  # Passando a instância F5TTS_ema_model como argumento
            ):
                # 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:
                    # Usando a função infer correta aqui
                    audio_file, _, _ = f5tts_model.infer(  # Usando f5tts_model.infer
                        ref_file=ref_audio_input,
                        ref_text=ref_text_input,
                        gen_text=chunk,
                        nfe_step=nfe_slider,
                        speed=speed_slider,
                        seed=seed_input,
                        remove_silence=remove_silence,
                    )
                    audio_data, _ = torchaudio.load(audio_file)
                    audio_segments.append(audio_data.squeeze().cpu().numpy())

                # Concatenar os segmentos de áudio gerados
                if audio_segments:
                    final_audio_data = np.concatenate(audio_segments)
                    return (
                        (24000, final_audio_data),  # Áudio final - assumindo taxa de amostragem de 24000
                        None,  # Espectrograma - não estamos gerando um espectrograma aqui
                        gr.update(value=ref_text_input),  # Nenhuma mudança no Texto de Referência
                        f5tts_model.seed  # Retornando o seed da instância F5TTS_ema_model
                    )
                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")
            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,
                    F5TTS_ema_model,  # Passando a instância F5TTS_ema_model como argumento
                ],
                outputs=[
                    audio_output,
                    ref_text_input,
                    seed_output,
                ],
            )

# 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()