Update app.py

app.py

@@ -1,199 +1,188 @@
-import os
-import json
-import base64
-import random
-import requests
-import io
-
 import gradio as gr
+import torch
+from diffusers import FluxPipeline
+import random
+import json
 from PIL import Image
+import os
 
-# ===========================
-#   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.")
-
-# ===========================
-#   FUNCIÓN SELECT MODEL
-# ===========================
-def select_model(preference):
-    if preference == "Lento 🧠 (mejor calidad)":
-        return {
-            "name": "Qwen2.5-72B-Instruct",
-            "description": "Máxima calidad de razonamiento"
-        }
-    else:
-        return {
-            "name": "Meta-Llama-3.1-8B-Instruct",
-            "description": "Rápido y eficiente"
-        }
-
-# ===========================
-#    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
+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):
+
+    if seed is None or seed == -1:
+        seed = random.randint(0, 2**32 - 1)
+
+    generator = torch.Generator("cpu").manual_seed(seed)
+
+    try:
+        image = pipe(
+            prompt,
+            num_inference_steps=steps,
+            guidance_scale=guidance,
+            height=height,
+            width=width,
+            generator=generator
+        ).images[0]
+
+        return image, seed
+
+    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}
     }
 
-    if images:
-        payload["images"] = images
-
-    headers = {
-        "Authorization": f"Bearer {API_KEY}",
-        "Content-Type": "application/json"
+    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"
+        }
     }
 
-    response = requests.post(url, json=payload, headers=headers)
-    
-    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"]
+    return [examples[example_id][field] for field in [
+        "scene", "subjects", "style", "colors", "lighting",
+        "mood", "background", "camera_angle"
+    ]]
 
-# ===========================
-#   CHAT PRINCIPAL
-# ===========================
-def chat_logic(user_text, user_image, model_selection, history):
-    selected = MODELS[model_selection]
-    model_name = selected["name"]
+# -------------------------------------------------------
+# INTERFAZ GRADIO
+# -------------------------------------------------------
+with gr.Blocks(title="FLUX.2 - Generador de Imágenes") as demo:
 
-    msg = []
-    images_encoded = None
+    gr.Markdown("# 🎨 FLUX.2 - Generador de Imágenes (Optimizado para CPU, Formato 9:16)")
 
-    msg.append({"role": "system", "content": selected["role"]})
-
-    if user_image:
-        images_encoded = [encode_image(user_image)]
-        msg.append({
-            "role": "user",
-            "content": [
-                {"type": "text", "text": user_text or "Describe esta imagen."},
-                {"type": "image", "image": images_encoded[0]}
-            ]
-        })
-    else:
-        msg.append({"role": "user", "content": user_text})
-
-    reply = call_sambanova(model_name, msg, images_encoded)
-
-    history.append(("🧍 Usuario: " + (user_text or "[Imagen]"), "🤖 Modelo: " + reply))
-    return history, reply
-
-# ===========================
-#   INTERFAZ
-# ===========================
-def create_ui():
-    with gr.Blocks(css=custom_css) as demo:
-
-        gr.Markdown("# 🤖 Chat avanzado con modelos SambaNova")
+    # ----------------- 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():
+                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():
-            model_selection = gr.Dropdown(
-                choices=list(MODELS.keys()),
-                value="general_fast",
-                label="Modelo"
+            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]
             )
 
-        user_text = gr.Textbox(label="Escribe tu mensaje")
-        user_image = gr.Image(label="Sube una imagen (opcional)")
-
-        chatbox = gr.Chatbot(label="Conversación")
-        output_text = gr.Textbox(label="Respuesta del modelo")
-
-        send_btn = gr.Button("Enviar")
-
-        send_btn.click(
-            chat_logic,
-            inputs=[user_text, user_image, model_selection, chatbox],
-            outputs=[chatbox, output_text]
-        )
-
-    return demo# ===========================
-#   LANZAR INTERFAZ
-# ===========================
-demo = create_ui()
-
-if __name__ == "__main__":
-    demo.launch()
+    # 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()