Agregando más modelos: chat, traducción y optimizaciones de velocidad

app.py

@@ -11,13 +11,42 @@ import base64
 MODELS = {
     "text": {
         "microsoft/DialoGPT-medium": "Chat conversacional",
+        "microsoft/DialoGPT-large": "Chat conversacional avanzado",
+        "microsoft/DialoGPT-small": "Chat conversacional rápido",
         "gpt2": "Generación de texto",
+        "gpt2-medium": "GPT-2 mediano",
+        "gpt2-large": "GPT-2 grande",
         "distilgpt2": "GPT-2 optimizado",
-        "EleutherAI/gpt-neo-125M": "GPT-Neo pequeño"
+        "EleutherAI/gpt-neo-125M": "GPT-Neo pequeño",
+        "EleutherAI/gpt-neo-1.3B": "GPT-Neo mediano",
+        "microsoft/DialoGPT-medium": "Chat conversacional",
+        "facebook/opt-125m": "OPT pequeño",
+        "facebook/opt-350m": "OPT mediano",
+        "bigscience/bloom-560m": "BLOOM multilingüe",
+        "bigscience/bloom-1b1": "BLOOM grande",
+        "microsoft/DialoGPT-medium": "Chat conversacional",
+        "Helsinki-NLP/opus-mt-es-en": "Traductor español-inglés",
+        "Helsinki-NLP/opus-mt-en-es": "Traductor inglés-español"
     },
     "image": {
         "runwayml/stable-diffusion-v1-5": "Stable Diffusion v1.5",
-        "CompVis/stable-diffusion-v1-4": "Stable Diffusion v1.4"
+        "CompVis/stable-diffusion-v1-4": "Stable Diffusion v1.4",
+        "stabilityai/stable-diffusion-2-1": "Stable Diffusion 2.1",
+        "stabilityai/stable-diffusion-xl-base-1.0": "SDXL Base",
+        "stabilityai/stable-diffusion-xl-refiner-1.0": "SDXL Refiner",
+        "prompthero/openjourney": "Midjourney style",
+        "dreamlike-art/dreamlike-photoreal-2.0": "Fotorealista",
+        "nitrosocke/Ghibli-Diffusion": "Estilo Studio Ghibli",
+        "nitrosocke/mo-di-diffusion": "Estilo moderno",
+        "CompVis/stable-diffusion-v1-4": "Stable Diffusion v1.4",
+        "runwayml/stable-diffusion-v1-5": "Stable Diffusion v1.5"
+    },
+    "chat": {
+        "microsoft/DialoGPT-medium": "Chat conversacional",
+        "microsoft/DialoGPT-large": "Chat conversacional avanzado",
+        "microsoft/DialoGPT-small": "Chat conversacional rápido",
+        "facebook/opt-350m": "OPT conversacional",
+        "bigscience/bloom-560m": "BLOOM multilingüe"
     }
 }
 
@@ -25,16 +54,25 @@ MODELS = {
 model_cache = {}
 
 def load_text_model(model_name):
-    """Cargar modelo de texto"""
+    """Cargar modelo de texto con soporte para diferentes tipos"""
     if model_name not in model_cache:
         print(f"Cargando modelo de texto: {model_name}")
-        tokenizer = AutoTokenizer.from_pretrained(model_name)
-        model = AutoModelForCausalLM.from_pretrained(model_name)
         
-        # Configurar para chat si es DialoGPT
-        if "dialogpt" in model_name.lower():
-            tokenizer.pad_token = tokenizer.eos_token
-            model.config.pad_token_id = model.config.eos_token_id
+        # Detectar tipo de modelo
+        if "opus-mt" in model_name.lower():
+            # Modelo de traducción
+            from transformers import MarianMTModel, MarianTokenizer
+            tokenizer = MarianTokenizer.from_pretrained(model_name)
+            model = MarianMTModel.from_pretrained(model_name)
+        else:
+            # Modelo de generación de texto
+            tokenizer = AutoTokenizer.from_pretrained(model_name)
+            model = AutoModelForCausalLM.from_pretrained(model_name)
+            
+            # Configurar para chat si es DialoGPT
+            if "dialogpt" in model_name.lower():
+                tokenizer.pad_token = tokenizer.eos_token
+                model.config.pad_token_id = model.config.eos_token_id
         
         model_cache[model_name] = {
             "tokenizer": tokenizer,
@@ -45,16 +83,21 @@ def load_text_model(model_name):
     return model_cache[model_name]
 
 def load_image_model(model_name):
-    """Cargar modelo de imagen"""
+    """Cargar modelo de imagen - optimizado para velocidad"""
     if model_name not in model_cache:
         print(f"Cargando modelo de imagen: {model_name}")
+        
+        # Optimizaciones para CPU y velocidad
         pipe = StableDiffusionPipeline.from_pretrained(
             model_name,
-            torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32
+            torch_dtype=torch.float32,  # Usar float32 para CPU
+            safety_checker=None,  # Desactivar safety checker para velocidad
+            requires_safety_checker=False
         )
         
-        if torch.cuda.is_available():
-            pipe = pipe.to("cuda")
+        # Optimizaciones adicionales
+        pipe.enable_attention_slicing()  # Reducir uso de memoria
+        pipe.enable_sequential_cpu_offload()  # Optimizar para CPU
         
         model_cache[model_name] = {
             "pipeline": pipe,
@@ -64,32 +107,40 @@ def load_image_model(model_name):
     return model_cache[model_name]
 
 def generate_text(prompt, model_name, max_length=100):
-    """Generar texto con el modelo seleccionado"""
+    """Generar texto con el modelo seleccionado - mejorado para diferentes tipos"""
     try:
         model_data = load_text_model(model_name)
         tokenizer = model_data["tokenizer"]
         model = model_data["model"]
         
-        # Preparar input
-        inputs = tokenizer.encode(prompt, return_tensors="pt")
-        
-        # Generar
-        with torch.no_grad():
-            outputs = model.generate(
-                inputs,
-                max_length=max_length,
-                num_return_sequences=1,
-                temperature=0.7,
-                do_sample=True,
-                pad_token_id=tokenizer.eos_token_id
-            )
-        
-        # Decodificar respuesta
-        response = tokenizer.decode(outputs[0], skip_special_tokens=True)
-        
-        # Para DialoGPT, extraer solo la respuesta del asistente
-        if "dialogpt" in model_name.lower():
-            response = response.replace(prompt, "").strip()
+        # Detectar si es modelo de traducción
+        if "opus-mt" in model_name.lower():
+            # Traducción
+            inputs = tokenizer.encode(prompt, return_tensors="pt", max_length=512, truncation=True)
+            with torch.no_grad():
+                outputs = model.generate(inputs, max_length=max_length, num_beams=4, early_stopping=True)
+            response = tokenizer.decode(outputs[0], skip_special_tokens=True)
+        else:
+            # Generación de texto
+            inputs = tokenizer.encode(prompt, return_tensors="pt")
+            
+            # Generar
+            with torch.no_grad():
+                outputs = model.generate(
+                    inputs,
+                    max_length=max_length,
+                    num_return_sequences=1,
+                    temperature=0.7,
+                    do_sample=True,
+                    pad_token_id=tokenizer.eos_token_id
+                )
+            
+            # Decodificar respuesta
+            response = tokenizer.decode(outputs[0], skip_special_tokens=True)
+            
+            # Para DialoGPT, extraer solo la respuesta del asistente
+            if "dialogpt" in model_name.lower():
+                response = response.replace(prompt, "").strip()
         
         return response
     
@@ -97,16 +148,22 @@ def generate_text(prompt, model_name, max_length=100):
         return f"Error generando texto: {str(e)}"
 
 def generate_image(prompt, model_name, num_inference_steps=20):
-    """Generar imagen con el modelo seleccionado"""
+    """Generar imagen con el modelo seleccionado - optimizado para velocidad"""
     try:
         model_data = load_image_model(model_name)
         pipeline = model_data["pipeline"]
         
-        # Generar imagen
+        # Optimizaciones para velocidad
+        if num_inference_steps > 20:
+            num_inference_steps = 20  # Limitar a máximo 20 pasos para velocidad
+        
+        # Generar imagen con configuración optimizada
         image = pipeline(
             prompt,
             num_inference_steps=num_inference_steps,
-            guidance_scale=7.5
+            guidance_scale=7.0,  # Reducido de 7.5 para velocidad
+            height=512,  # Tamaño fijo para consistencia
+            width=512
         ).images[0]
         
         return image
@@ -207,7 +264,7 @@ with gr.Blocks(title="Modelos Libres de IA", theme=gr.themes.Soft()) as demo:
             with gr.Row():
                 with gr.Column():
                     chat_model = gr.Dropdown(
-                        choices=["microsoft/DialoGPT-medium"],
+                        choices=list(MODELS["chat"].keys()),
                         value="microsoft/DialoGPT-medium",
                         label="Modelo de Chat"
                     )
@@ -237,6 +294,35 @@ with gr.Blocks(title="Modelos Libres de IA", theme=gr.themes.Soft()) as demo:
                 outputs=[chatbot]
             )
         
+        # Tab de Traducción
+        with gr.TabItem("🌐 Traducción"):
+            with gr.Row():
+                with gr.Column():
+                    translate_model = gr.Dropdown(
+                        choices=["Helsinki-NLP/opus-mt-es-en", "Helsinki-NLP/opus-mt-en-es"],
+                        value="Helsinki-NLP/opus-mt-es-en",
+                        label="Modelo de Traducción"
+                    )
+                    translate_text = gr.Textbox(
+                        label="Texto a traducir",
+                        placeholder="Escribe el texto que quieres traducir...",
+                        lines=3
+                    )
+                    translate_btn = gr.Button("Traducir", variant="primary")
+                
+                with gr.Column():
+                    translate_output = gr.Textbox(
+                        label="Traducción",
+                        lines=3,
+                        interactive=False
+                    )
+            
+            translate_btn.click(
+                generate_text,
+                inputs=[translate_text, translate_model, gr.Slider(value=100, visible=False)],
+                outputs=translate_output
+            )
+        
         # Tab de Generación de Imágenes
         with gr.TabItem("🎨 Generación de Imágenes"):
             with gr.Row():
@@ -254,7 +340,7 @@ with gr.Blocks(title="Modelos Libres de IA", theme=gr.themes.Soft()) as demo:
                     steps = gr.Slider(
                         minimum=10,
                         maximum=50,
-                        value=20,
+                        value=15,
                         step=5,
                         label="Pasos de inferencia"
                     )