Update chat_ai.py

chat_ai.py

@@ -1,10 +1,30 @@
-import gradio as gr
+import re
 import tempfile
+
+import gradio as gr
 import numpy as np
-import torchaudio
 import soundfile as sf
-from transformers import AutoTokenizer
-from f5_tts.model import DiT
+import torchaudio
+from cached_path import cached_path
+from transformers import AutoModelForCausalLM, AutoTokenizer
+from num2words import num2words
+
+try:
+    import spaces
+
+    USING_SPACES = True
+except ImportError:
+    USING_SPACES = False
+
+
+def gpu_decorator(func):
+    if USING_SPACES:
+        return spaces.GPU(func)
+    else:
+        return func
+
+
+from f5_tts.model import DiT, UNetT
 from f5_tts.infer.utils_infer import (
     load_vocoder,
     load_model,
@@ -13,18 +33,43 @@ from f5_tts.infer.utils_infer import (
     remove_silence_for_generated_wav,
     save_spectrogram,
 )
-from num2words import num2words
-import re
 
-# Cargar vocoder y modelo
 vocoder = load_vocoder()
+
+# Cargar modelos
 F5TTS_model_cfg = dict(dim=1024, depth=22, heads=16, ff_mult=2, text_dim=512, conv_layers=4)
 F5TTS_ema_model = load_model(
-    DiT, F5TTS_model_cfg, "hf://jpgallegoar/F5-Spanish/model_1200000.safetensors"
+    DiT, F5TTS_model_cfg, str(cached_path("hf://jpgallegoar/F5-Spanish/model_1200000.safetensors"))
 )
 
+chat_model_state = None
+chat_tokenizer_state = None
+
+
+@gpu_decorator
+def generate_response(messages, model, tokenizer):
+    """Generar respuesta usando Qwen."""
+    text = tokenizer.apply_chat_template(
+        messages,
+        tokenize=False,
+        add_generation_prompt=True,
+    )
+
+    model_inputs = tokenizer([text], return_tensors="pt").to(model.device)
+    generated_ids = model.generate(
+        **model_inputs,
+        max_new_tokens=512,
+        temperature=0.7,
+        top_p=0.95,
+    )
+
+    generated_ids = [
+        output_ids[len(input_ids):] for input_ids, output_ids in zip(model_inputs.input_ids, generated_ids)
+    ]
+    return tokenizer.batch_decode(generated_ids, skip_special_tokens=True)[0]
+
+
 def traducir_numero_a_texto(texto):
-    """Convierte números a palabras en el texto."""
     texto_separado = re.sub(r'([A-Za-z])(\d)', r'\1 \2', texto)
     texto_separado = re.sub(r'(\d)([A-Za-z])', r'\1 \2', texto_separado)
 
@@ -36,24 +81,36 @@ def traducir_numero_a_texto(texto):
 
     return texto_traducido
 
+
+@gpu_decorator
 def infer(
-    ref_audio_orig, ref_text, gen_text, remove_silence=False, cross_fade_duration=0.15, speed=1.0
+    ref_audio_orig, ref_text, gen_text, model, remove_silence, cross_fade_duration=0.15, speed=1, show_info=gr.Info
 ):
-    """Realiza la inferencia para convertir texto en voz."""
-    ref_audio, ref_text = preprocess_ref_audio_text(ref_audio_orig, ref_text)
+    ref_audio, ref_text = preprocess_ref_audio_text(ref_audio_orig, ref_text, show_info=show_info)
+
+    ema_model = F5TTS_ema_model
+
+    if not gen_text.startswith(" "):
+        gen_text = " " + gen_text
+    if not gen_text.endswith(". "):
+        gen_text += ". "
 
-    gen_text = traducir_numero_a_texto(gen_text.lower())
+    gen_text = gen_text.lower()
+    gen_text = traducir_numero_a_texto(gen_text)
 
     final_wave, final_sample_rate, combined_spectrogram = infer_process(
         ref_audio,
         ref_text,
         gen_text,
-        F5TTS_ema_model,
+        ema_model,
         vocoder,
         cross_fade_duration=cross_fade_duration,
         speed=speed,
+        show_info=show_info,
+        progress=gr.Progress(),
     )
 
+    # Remover silencios
     if remove_silence:
         with tempfile.NamedTemporaryFile(delete=False, suffix=".wav") as f:
             sf.write(f.name, final_wave, final_sample_rate)
@@ -68,57 +125,72 @@ def infer(
 
     return (final_sample_rate, final_wave), spectrogram_path
 
-def tts_pipeline(ref_audio, ref_text, gen_text, remove_silence, speed):
-    """Pipeline para la interfaz de Gradio."""
-    if not ref_audio:
-        return None, "Por favor sube un audio de referencia."
 
-    try:
-        (sample_rate, audio), spectrogram_path = infer(
-            ref_audio, ref_text, gen_text, remove_silence=remove_silence, speed=speed
+with gr.Blocks() as app_tts:
+    gr.Markdown("# TTS por Lotes")
+    ref_audio_input = gr.Audio(label="Audio de Referencia", type="filepath")
+    gen_text_input = gr.Textbox(label="Texto para Generar", lines=10)
+    model_choice = gr.Radio(choices=["F5-TTS"], label="Seleccionar Modelo TTS", value="F5-TTS")
+    generate_btn = gr.Button("Sintetizar", variant="primary")
+    with gr.Accordion("Configuraciones Avanzadas", open=False):
+        ref_text_input = gr.Textbox(
+            label="Texto de Referencia",
+            info="Deja en blanco para transcribir automáticamente el audio de referencia. Si ingresas texto, sobrescribirá la transcripción automática.",
+            lines=2,
+        )
+        remove_silence = gr.Checkbox(
+            label="Eliminar Silencios",
+            info="El modelo tiende a producir silencios, especialmente en audios más largos. Podemos eliminar manualmente los silencios si es necesario. Ten en cuenta que esta es una característica experimental y puede producir resultados extraños. Esto también aumentará el tiempo de generación.",
+            value=False,
+        )
+        speed_slider = gr.Slider(
+            label="Velocidad",
+            minimum=0.3,
+            maximum=2.0,
+            value=1.0,
+            step=0.1,
+            info="Ajusta la velocidad del audio.",
+        )
+        cross_fade_duration_slider = gr.Slider(
+            label="Duración del Cross-Fade (s)",
+            minimum=0.0,
+            maximum=1.0,
+            value=0.15,
+            step=0.01,
+            info="Establece la duración del cross-fade entre clips de audio.",
         )
-        with tempfile.NamedTemporaryFile(suffix=".wav", delete=False) as tmp_audio:
-            sf.write(tmp_audio.name, audio, sample_rate)
-            return tmp_audio.name, spectrogram_path
-    except Exception as e:
-        return None, f"Error al generar audio: {str(e)}"
-
-# Crear interfaz con Gradio
-with gr.Blocks() as demo:
-    gr.Markdown("""# Conversión de Texto a Voz (TTS) en Español
-
-Convierte texto en audio en español usando un modelo de TTS. Proporciona un audio de referencia y el texto a convertir.
-
-**Instrucciones:**
-1. Sube un audio de referencia (formato WAV o MP3, de 11 a 14 segundos).
-2. Opcionalmente, ingresa el texto correspondiente al audio de referencia.
-3. Escribe el texto que deseas convertir a voz.
-4. Haz clic en "Generar Audio".
-
-*Nota: Los números en el texto serán convertidos automáticamente a palabras.*
-""")
-
-    with gr.Row():
-        ref_audio = gr.Audio(label="Audio de Referencia", type="filepath")
-        ref_text = gr.Textbox(label="Texto de Referencia (Opcional)", placeholder="Transcripción del audio de referencia")
 
-    gen_text = gr.Textbox(label="Texto para Convertir a Voz", lines=4, placeholder="Escribe aquí el texto a convertir")
+    audio_output = gr.Audio(label="Audio Sintetizado")
+    spectrogram_output = gr.Image(label="Espectrograma")
 
-    with gr.Row():
-        remove_silence = gr.Checkbox(label="Eliminar Silencios", value=False)
-        speed = gr.Slider(label="Velocidad", minimum=0.5, maximum=2.0, value=1.0, step=0.1)
+    generate_btn.click(
+        infer,
+        inputs=[
+            ref_audio_input,
+            ref_text_input,
+            gen_text_input,
+            model_choice,
+            remove_silence,
+            cross_fade_duration_slider,
+            speed_slider,
+        ],
+        outputs=[audio_output, spectrogram_output],
+    )
 
-    generate_btn = gr.Button("Generar Audio")
 
-    with gr.Row():
-        audio_output = gr.Audio(label="Audio Generado", type="filepath")
-        spectrogram_output = gr.Image(label="Espectrograma")
+with gr.Blocks() as app:
+    gr.Markdown(
+        """
+# Spanish-F5
 
-    generate_btn.click(
-        tts_pipeline,
-        inputs=[ref_audio, ref_text, gen_text, remove_silence, speed],
-        outputs=[audio_output, spectrogram_output],
+Esta es una interfaz web para F5 TTS, con un finetuning para poder hablar en castellano.
+"""
     )
+    gr.TabbedInterface(
+        [app_tts],
+        ["TTS"],
+    )
+
 
-# Ejecutar la aplicación en Spaces
-demo.launch()
+if __name__ == "__main__":
+    app.queue().launch()