Update app.py

app.py

@@ -5,6 +5,7 @@ from tensorflow import keras
 import torch
 from huggingface_hub import hf_hub_download
 from speechbrain.inference.TTS import Tacotron2
+import os
 
 # Cargar modelo Tacotron2
 tacotron2 = Tacotron2.from_hparams(
@@ -13,15 +14,51 @@ tacotron2 = Tacotron2.from_hparams(
     run_opts={"device": "cpu"}
 )
 
-# Descargar y cargar el modelo Generator entrenado
-model_path = hf_hub_download(
-    repo_id="Bmo411/WGAN",
-    filename="generator_epoch_1000.keras"
-)
-generator = keras.models.load_model(model_path, compile=False)
+# Diccionario para almacenar los modelos cargados
+loaded_models = {}
+
+# Modelos disponibles - define aquí las épocas que quieres incluir
+available_models = {
+    "Época 100": "generator_epoch_100.keras",
+    "Época 1000": "generator_epoch_250.keras",
+    "Época 4200": "generator_epoch_500.keras",
+    "Época 4700": "generator_epoch_750.keras",
+    "Época 7700": "generator_epoch_1000.keras"
+}
+
+# Función para cargar un modelo específico
+def load_generator_model(model_name):
+    if model_name in loaded_models:
+        return loaded_models[model_name]
+    
+    try:
+        model_path = hf_hub_download(
+            repo_id="Bmo411/WGAN",
+            filename=model_name
+        )
+        model = keras.models.load_model(model_path, compile=False)
+        loaded_models[model_name] = model
+        print(f"Modelo {model_name} cargado correctamente")
+        return model
+    except Exception as e:
+        print(f"Error al cargar el modelo {model_name}: {e}")
+        # Si falla la carga, intentamos usar el modelo de la época 1000 como fallback
+        try:
+            fallback_model = "generator_epoch_1000.keras"
+            model_path = hf_hub_download(
+                repo_id="Bmo411/WGAN",
+                filename=fallback_model
+            )
+            model = keras.models.load_model(model_path, compile=False)
+            loaded_models[model_name] = model  # Guardamos con el nombre original para evitar recargar
+            print(f"Usando modelo fallback {fallback_model}")
+            return model
+        except:
+            print("Error crítico al cargar modelos. No hay modelos disponibles.")
+            return None
 
 # Función para convertir texto a audio
-def text_to_audio(text):
+def text_to_audio(text, model_epoch):
     # Crear un array vacío por defecto en caso de error
     default_audio = np.zeros(8000, dtype=np.float32)
     sample_rate = 8000  # Ajusta según la configuración de tu modelo
@@ -30,6 +67,16 @@ def text_to_audio(text):
         return (sample_rate, default_audio)
     
     try:
+        # Obtener el nombre del archivo del modelo seleccionado
+        model_filename = available_models[model_epoch]
+        
+        # Cargar el modelo generador correspondiente
+        generator = load_generator_model(model_filename)
+        
+        if generator is None:
+            print("No se pudo cargar el generador")
+            return (sample_rate, default_audio)
+        
         # Convertir texto a mel-spectrograma con Tacotron2
         mel_output, _, _ = tacotron2.encode_text(text)
         mel = mel_output.detach().cpu().numpy().astype(np.float32)
@@ -91,16 +138,45 @@ def text_to_audio(text):
         return (sample_rate, default_audio)
 
 # Crear interfaz en Gradio
-interface = gr.Interface(
-    fn=text_to_audio,
-    inputs=gr.Textbox(lines=2, placeholder="Escribe nine-"),
-    outputs=gr.Audio(label="Audio generado"),
-    title="Demo de TTS con Tacotron2 + Generador",
-    description="Convierte texto en audio usando Tacotron2 + modelo Generator entrenado.",
-    examples=[["nine"], ["nine"]]
-)
+with gr.Blocks(title="Demo de TTS con Tacotron2 + Generador") as interface:
+    gr.Markdown("# Demo de TTS con Tacotron2 + Generador")
+    gr.Markdown("Convierte texto en audio usando Tacotron2 + modelo Generator entrenado en diferentes épocas.")
+    
+    with gr.Row():
+        with gr.Column(scale=3):
+            text_input = gr.Textbox(lines=2, placeholder="Escribe nine-", label="Texto a convertir")
+        with gr.Column(scale=1):
+            model_selection = gr.Dropdown(
+                choices=list(available_models.keys()),
+                value="Época 1000",
+                label="Selecciona la época del modelo"
+            )
+    
+    generate_btn = gr.Button("Generar Audio", variant="primary")
+    
+    audio_output = gr.Audio(label="Audio generado")
+    
+    # Configurar ejemplos
+    examples = gr.Examples(
+        examples=[
+            ["nine", "Época 100"],
+            ["nine", "Época 1000"],
+            ["nine", "Época 4200"]
+        ],
+        inputs=[text_input, model_selection],
+        outputs=audio_output
+    )
+    
+    # Conectar botón a la función
+    generate_btn.click(fn=text_to_audio, inputs=[text_input, model_selection], outputs=audio_output)
+    
+    # También permitir enviar con Enter desde el cuadro de texto
+    text_input.submit(fn=text_to_audio, inputs=[text_input, model_selection], outputs=audio_output)
 
 # Lanzar aplicación
 if __name__ == "__main__":
-    interface.launch(debug=True)
-    
+    # Precargamos el modelo de la época 1000 para tenerlo disponible inmediatamente
+    load_generator_model(available_models["Época 1000"])
+    
+    # Lanzamos la interfaz
+    interface.launch(debug=True)