Update app.py

app.py

@@ -1,188 +1,236 @@
+import os
+import base64
+import requests
+import io
 import gradio as gr
-import torch
-from diffusers import FluxPipeline
-import random
-import json
 from PIL import Image
-import os
-
-HF_TOKEN = os.getenv("HF_TOKEN")  # Toma tu token automáticamente
-
-# -------------------------------------------------------
-# Cargar modelo FLUX.2 (optimizado para CPU)
-# -------------------------------------------------------
-def load_flux_model():
-    # Usar torch.float32 para CPU, ya que float16 puede no ser óptimo
-    dtype = torch.float32
-    pipe = FluxPipeline.from_pretrained(
-        "black-forest-labs/FLUX.2-dev",
-        torch_dtype=dtype,
-        token=HF_TOKEN
-    )
-
-    # Forzar a CPU
-    pipe.to("cpu")
-    # Optimizar para CPU: habilitar offload secuencial para ahorrar memoria
-    pipe.enable_sequential_cpu_offload()
-    # Habilitar attention slicing para reducir uso de memoria
-    pipe.enable_attention_slicing()
-
-    return pipe
 
-pipe = load_flux_model()
-
-# -------------------------------------------------------
-# GENERACIÓN SIMPLE
-# -------------------------------------------------------
-def generate_flux_image(prompt, negative_prompt="", steps=28, guidance=7.0,
-                        width=576, height=1024, seed=None):
+# ===========================
+#   CLAVES PARA SAMBANOVA
+# ===========================
+API_KEY = os.getenv("REVE_API_KEY")
+BASE_URL = "https://api.sambanova.ai/v1"
+
+if not API_KEY:
+    raise ValueError("⚠️ ERROR: No se encontró la variable REVE_API_KEY en HuggingFace Spaces.")
+
+# ===========================
+#   DICCIONARIO DE MODELOS
+# ===========================
+MODELS = {
+    "general_fast": {
+        "name": "Meta-Llama-3.1-8B-Instruct",
+        "role": "🔄 Respuestas rápidas y generales",
+        "description": "Conversación ligera y eficiente."
+    },
+    "general_smart": {
+        "name": "Meta-Llama-3.3-70B-Instruct",
+        "role": "🧠 Razonamiento profundo",
+        "description": "Análisis detallado y avanzado."
+    },
+    "coding_expert": {
+        "name": "DeepSeek-V3.1",
+        "role": "💻 Programación y debugging",
+        "description": "Ideal para desarrollo."
+    },
+    "coding_alt": {
+        "name": "DeepSeek-V3-0324",
+        "role": "⚡ Código rápido",
+        "description": "Alternativa veloz."
+    },
+    "massive_brain": {
+        "name": "gpt-oss-120b",
+        "role": "🏛️ Sabiduría masiva",
+        "description": "Problemas pesados y complejos."
+    },
+    "specialized_1": {
+        "name": "DeepSeek-V3.1-Terminus",
+        "role": "🎯 Especialista técnico",
+        "description": "Tareas científicas y avanzadas."
+    },
+    "specialized_2": {
+        "name": "Llama-3.3-Swallow-70B-Instruct-v0.4",
+        "role": "🔥 Sin censura",
+        "description": "Modelo sin restricciones."
+    },
+    "multilingual": {
+        "name": "Qwen3-32B",
+        "role": "🌍 Multilingüe",
+        "description": "Múltiples idiomas."
+    },
+    "arabic_special": {
+        "name": "ALLaM-7B-Instruct-preview",
+        "role": "🕌 Estilo árabe y moda íntima",
+        "description": "Experto en estilos del Medio Oriente."
+    },
+    "vision_expert": {
+        "name": "Llama-4-Maverick-17B-128E-Instruct",
+        "role": "👁️ Visión avanzada",
+        "description": "Análisis de imágenes."
+    }
+}
+
+# ===========================
+#   CSS PERSONALIZADO
+# ===========================
+custom_css = """
+.gradio-container {max-width: 900px !important; margin: auto;}
+.model-box {padding: 10px; border-radius: 10px; background: #1d1d1d; margin-bottom: 6px;}
+"""
+
+# ===========================
+#   CODIFICAR IMAGEN EN BASE64
+# ===========================
+def encode_image(img):
+    if img is None:
+        return None
+    buffer = io.BytesIO()
+    img.save(buffer, format="PNG")
+    return base64.b64encode(buffer.getvalue()).decode()
+
+# ===========================
+#   API CALL SAMBANOVA
+# ===========================
+def call_sambanova(model, messages, images=None):
+    url = f"{BASE_URL}/chat/completions"
+    
+    payload = {
+        "model": model,
+        "messages": messages,
+        "stream": False
+    }
 
-    if seed is None or seed == -1:
-        seed = random.randint(0, 2**32 - 1)
+    if images:
+        payload["images"] = images
 
-    generator = torch.Generator("cpu").manual_seed(seed)
+    headers = {
+        "Authorization": f"Bearer {API_KEY}",
+        "Content-Type": "application/json"
+    }
 
     try:
-        image = pipe(
-            prompt,
-            num_inference_steps=steps,
-            guidance_scale=guidance,
-            height=height,
-            width=width,
-            generator=generator
-        ).images[0]
-
-        return image, seed
+        response = requests.post(url, json=payload, headers=headers, timeout=30)
+        
+        if response.status_code != 200:
+            return f"⚠️ Error en API: {response.status_code}\n{response.text}"
 
+        data = response.json()
+        return data["choices"][0]["message"]["content"]
     except Exception as e:
-        return None, f"Error: {str(e)}"
-
-# -------------------------------------------------------
-# GENERACIÓN JSON
-# -------------------------------------------------------
-def generate_json_image(scene, subjects, style, colors, lighting, mood,
-                        background, camera_angle, seed=None):
-
-    json_prompt = {
-        "scene": scene,
-        "subjects": [{"description": subjects}],
-        "style": style,
-        "color_palette": [c.strip() for c in colors.split(",")] if colors else [],
-        "lighting": lighting,
-        "mood": mood,
-        "background": background,
-        "camera": {"angle": camera_angle}
-    }
-
-    prompt_text = json.dumps(json_prompt, ensure_ascii=False)
-
-    return generate_flux_image(prompt_text, steps=25, seed=seed)
-
-def load_example(example_id):
-    examples = {
-        "producto": {
-            "scene": "Fotografía profesional de producto en estudio",
-            "subjects": "Taza de café minimalista con vapor ascendiendo",
-            "style": "Fotografía de producto ultra realista",
-            "colors": "#2C2C2C, #E8E8E8, #FF6B35",
-            "lighting": "Iluminación suave de 3 puntos",
-            "mood": "Limpio y profesional",
-            "background": "Superficie de concreto pulido",
-            "camera_angle": "ángulo alto"
-        },
-        "paisaje": {
-            "scene": "Paisaje de montaña al atardecer",
-            "subjects": "Lobo solitario en la cima de la montaña",
-            "style": "Pintura digital épica",
-            "colors": "#FF6B35, #1A535C, #4ECDC4",
-            "lighting": "Luz dorada del atardecer",
-            "mood": "Épico y sereno",
-            "background": "Montañas y cielo naranja",
-            "camera_angle": "vista panorámica"
-        },
-        "retrato": {
-            "scene": "Estudio de retrato profesional",
-            "subjects": "Persona con sonrisa genuina, iluminación dramática",
-            "style": "Retrato fotográfico profesional",
-            "colors": "#2C2C2C, #F5F5F5, #8B4513",
-            "lighting": "Iluminación Rembrandt",
-            "mood": "Elegante y confiado",
-            "background": "Fondo negro mate",
-            "camera_angle": "primer plano"
-        }
-    }
-
-    return [examples[example_id][field] for field in [
-        "scene", "subjects", "style", "colors", "lighting",
-        "mood", "background", "camera_angle"
-    ]]
-
-# -------------------------------------------------------
-# INTERFAZ GRADIO
-# -------------------------------------------------------
-with gr.Blocks(title="FLUX.2 - Generador de Imágenes") as demo:
+        return f"⚠️ Error de conexión: {str(e)}"
+
+# ===========================
+#   CHAT PRINCIPAL
+# ===========================
+def chat_logic(user_text, user_image, model_selection, history):
+    if not user_text and not user_image:
+        return history, "⚠️ Por favor, escribe un mensaje o sube una imagen."
+    
+    selected = MODELS[model_selection]
+    model_name = selected["name"]
+
+    messages = []
+    images_encoded = None
+
+    # Add system message
+    messages.append({"role": "system", "content": selected["role"]})
+
+    # Handle image and text
+    if user_image:
+        images_encoded = [encode_image(user_image)]
+        messages.append({
+            "role": "user",
+            "content": [
+                {"type": "text", "text": user_text or "Describe esta imagen."},
+                {"type": "image", "image": images_encoded[0]}
+            ]
+        })
+    else:
+        messages.append({"role": "user", "content": user_text})
+
+    # Call API
+    reply = call_sambanova(model_name, messages, images_encoded)
+
+    # Update history
+    user_display = f"🧍 Usuario: {user_text}" if user_text else "🧍 Usuario: [Imagen]"
+    history.append((user_display, f"🤖 {selected['role']}: {reply}"))
+    
+    return history, reply
+
+# ===========================
+#   INTERFAZ MEJORADA
+# ===========================
+def create_ui():
+    with gr.Blocks(css=custom_css, title="Chat Multimodal SambaNova") as demo:
+        gr.Markdown("""
+        # 🤖 Chat avanzado con modelos SambaNova
+        *Soporte multimodal (texto + imágenes) con 10+ modelos especializados*
+        """)
 
-    gr.Markdown("# 🎨 FLUX.2 - Generador de Imágenes (Optimizado para CPU, Formato 9:16)")
-
-    # ----------------- TAB SIMPLE -----------------
-    with gr.Tab("🎯 Prompt Simple"):
         with gr.Row():
-            with gr.Column():
-                prompt_simple = gr.Textbox(label="Prompt")
-                negative_prompt = gr.Textbox(label="Negative Prompt (FLUX no lo usa)")
-                steps = gr.Slider(10, 50, value=28, step=1, label="Steps")
-                guidance = gr.Slider(1, 20, value=7, step=0.5, label="Guidance")
-                width = gr.Slider(256, 1024, value=576, step=64, label="Ancho (9:16 por defecto)")
-                height = gr.Slider(256, 1024, value=1024, step=64, label="Alto (9:16 por defecto)")
-                seed = gr.Number(value=-1, label="Semilla")
-                btn_simple = gr.Button("Generar Imagen")
+            with gr.Column(scale=1):
+                model_selection = gr.Dropdown(
+                    choices=list(MODELS.keys()),
+                    value="general_fast",
+                    label="🎛️ Selecciona modelo",
+                    info="Elige el modelo según tus necesidades"
+                )
+                # Display model description
+                model_info = gr.Markdown("**Meta-Llama-3.1-8B-Instruct**: Conversación ligera y eficiente.")
 
-            with gr.Column():
-                out_img_simple = gr.Image(label="Imagen Generada")
-                out_seed_simple = gr.Number(label="Semilla Usada")
-
-    # ----------------- TAB JSON -----------------
-    with gr.Tab("📝 JSON Prompt"):
-        scene = gr.Textbox(label="Escena")
-        subjects = gr.Textbox(label="Sujetos")
-        style = gr.Textbox(label="Estilo")
-        colors = gr.Textbox(label="Colores (coma)")
-        lighting = gr.Textbox(label="Iluminación")
-        mood = gr.Textbox(label="Estado de Ánimo")
-        background = gr.Textbox(label="Fondo")
-        camera_angle = gr.Textbox(label="Ángulo de cámara")
-        json_seed = gr.Number(value=-1, label="Semilla")
-        btn_json = gr.Button("Generar con JSON")
-
-        out_img_json = gr.Image()
-        out_seed_json = gr.Number()
+        with gr.Row():
+            with gr.Column(scale=1):
+                user_image = gr.Image(
+                    label="🖼️ Sube una imagen (opcional)",
+                    type="pil",
+                    height=200
+                )
+            with gr.Column(scale=2):
+                user_text = gr.Textbox(
+                    label="💬 Escribe tu mensaje",
+                    placeholder="Escribe tu pregunta aquí...",
+                    lines=3
+                )
+
+        send_btn = gr.Button("🚀 Enviar mensaje", variant="primary")
 
         with gr.Row():
-            gr.Button("Ej Producto").click(
-                fn=lambda: load_example("producto"),
-                outputs=[scene, subjects, style, colors, lighting, mood, background, camera_angle]
-            )
-            gr.Button("Ej Paisaje").click(
-                fn=lambda: load_example("paisaje"),
-                outputs=[scene, subjects, style, colors, lighting, mood, background, camera_angle]
-            )
-            gr.Button("Ej Retrato").click(
-                fn=lambda: load_example("retrato"),
-                outputs=[scene, subjects, style, colors, lighting, mood, background, camera_angle]
-            )
-
-    # Eventos
-    btn_simple.click(
-        fn=generate_flux_image,
-        inputs=[prompt_simple, negative_prompt, steps, guidance, width, height, seed],
-        outputs=[out_img_simple, out_seed_simple]
-    )
-
-    btn_json.click(
-        fn=generate_json_image,
-        inputs=[scene, subjects, style, colors, lighting, mood, background, camera_angle, json_seed],
-        outputs=[out_img_json, out_seed_json]
-    )
-
-# Lanzar la demo (para ejecución local o en CPU)
-demo.launch()
+            with gr.Column():
+                chatbox = gr.Chatbot(
+                    label="💭 Conversación",
+                    height=400
+                )
+
+        # Update model info when selection changes
+        def update_model_info(selection):
+            selected = MODELS[selection]
+            return f"**{selected['name']}**: {selected['description']}"
+        
+        model_selection.change(
+            update_model_info,
+            inputs=[model_selection],
+            outputs=[model_info]
+        )
+
+        # Clear inputs after sending
+        def clear_inputs():
+            return None, ""
+        
+        send_btn.click(
+            chat_logic,
+            inputs=[user_text, user_image, model_selection, chatbox],
+            outputs=[chatbox, user_text]
+        ).then(
+            clear_inputs,
+            outputs=[user_image, user_text]
+        )
+
+    return demo
+
+# ===========================
+#   LANZAR INTERFAZ
+# ===========================
+demo = create_ui()
+
+if __name__ == "__main__":
+    demo.launch(share=False)