Update app.py

app.py

@@ -1,150 +1,436 @@
 import os
-import time
-import json
-import base64
-import requests
 import gradio as gr
+from diffusers import DiffusionPipeline
+import torch
+import requests
+from PIL import Image
+from io import BytesIO
+import concurrent.futures
+import threading
 
-# =============================================================
-# Utilidades Generales
-# =============================================================
-def load_image(path):
-    with open(path, "rb") as f:
-        return f.read()
-
-# =============================================================
-# FLUX (Optimizado CPU)
-# =============================================================
-def flux_generate(prompt, seed, steps, width, height):
-    try:
-        payload = {
-            "prompt": prompt,
-            "seed": seed,
-            "steps": steps,
-            "width": width,
-            "height": height,
-        }
-        return "⚠ FLUX CPU: Modelo optimizado próximamente."
-    except Exception as e:
-        return f"Error FLUX: {e}"
+# ==============================
+#   CONFIGURACIÓN BASE CPU
+# ==============================
+DEVICE = "cpu"
+torch.set_grad_enabled(False)
 
-# =============================================================
-# Stable Diffusion 1.5 (Optimizado CPU, resultados realistas)
-# =============================================================
-def sd15_generate(prompt, seed, steps, width, height):
-    try:
-        payload = {
-            "prompt": prompt + ", ultra realistic, photorealistic, sharp details, 8k, high texture clarity",
-            "seed": seed,
-            "steps": steps,
-            "width": width,
-            "height": height,
-        }
-        return "⚠ SD1.5 CPU: Generador listo para integración real."
-    except Exception as e:
-        return f"Error SD1.5: {e}"
+# PARÁMETROS POR DEFECTO AJUSTADOS PARA CPU MÁS RÁPIDO
+DEFAULT_STEPS = 15  # Reducido de 20 para más velocidad
+DEFAULT_WIDTH = 512  # Reducido de 576 para menos carga
+DEFAULT_HEIGHT = 768  # Reducido de 1024 para menos carga (mantiene relación aproximada)
+
+def load_flux(model_id):
+    pipe = DiffusionPipeline.from_pretrained(model_id, torch_dtype=torch.float32)
+    pipe.to(DEVICE)
+    pipe.enable_attention_slicing()
+    return pipe
+
+# Cache de modelos
+MODEL_CACHE = {}
+
+# ==============================
+#   GENERADOR FLUX
+# ==============================
+def generate_flux(model_name, prompt, steps, guidance, width, height, seed):
+    if model_name not in MODEL_CACHE:
+        MODEL_CACHE[model_name] = load_flux(model_name)
+    pipe = MODEL_CACHE[model_name]
+
+    generator = torch.manual_seed(seed) if seed else None
+
+    image = pipe(
+        prompt=prompt,
+        num_inference_steps=steps,
+        guidance_scale=guidance,
+        width=width,
+        height=height,
+        generator=generator
+    ).images[0]
 
-# =============================================================
-# REVE CREATE (Con selector de imágenes 1–8)
-# =============================================================
-def reve_create_generate(api_key, prompt, seed, steps, num_images):
+    out = "/tmp/flux_output.png"
+    image.save(out)
+    return out
+
+# ==============================
+#   GENERADOR SD1.5
+# ==============================
+def load_sd15():
+    pipe = DiffusionPipeline.from_pretrained("runwayml/stable-diffusion-v1-5", torch_dtype=torch.float32)
+    pipe.to(DEVICE)
+    pipe.enable_attention_slicing()
+    return pipe
+
+# SD15 load único
+def generate_sd(prompt, steps, guidance, width, height, seed):
+    if "sd15" not in MODEL_CACHE:
+        MODEL_CACHE["sd15"] = load_sd15()
+    pipe = MODEL_CACHE["sd15"]
+
+    generator = torch.manual_seed(seed) if seed else None
+
+    image = pipe(
+        prompt=prompt,
+        num_inference_steps=steps,
+        guidance_scale=guidance,
+        width=width,
+        height=height,
+        generator=generator
+    ).images[0]
+
+    out = "/tmp/sd15_output.png"
+    image.save(out)
+    return out
+
+# ==============================
+#   REVE CREATE MODIFICADO (Generación múltiple con threads)
+# ==============================
+def generate_single_reve_image(prompt, key, model, index, results_list, lock, progress_callback=None):
+    """Función auxiliar para generar una sola imagen"""
     try:
-        url = "https://api.reveai.xyz/v1/generate"
-
-        payload = {
-            "prompt": prompt,
-            "seed": seed,
-            "steps": steps,
-            "num_images": num_images,
-            "width": 576,
-            "height": 1024,
-        }
-
-        headers = {
-            "Authorization": f"Bearer {api_key}",
-            "Content-Type": "application/json",
-        }
-
-        res = requests.post(url, json=payload, headers=headers)
-        data = res.json()
-
-        outputs = []
-        for img_b64 in data.get("images", []):
-            img_bytes = base64.b64decode(img_b64)
-            filename = f"reve_{int(time.time())}.png"
-            filepath = f"/tmp/{filename}"
-            with open(filepath, "wb") as f:
-                f.write(img_bytes)
-            outputs.append(filepath)
-
-        return outputs
+        url = "https://api.reveai.xyz/v1/images"
+        headers = {"Authorization": f"Bearer {key}"}
+        data = {"prompt": prompt, "model": model}
+        
+        resp = requests.post(url, json=data, headers=headers, timeout=30)
+        
+        if resp.status_code != 200:
+            print(f"Error en imagen {index+1}: {resp.status_code}")
+            return
+        
+        img_url = resp.json().get("image")
+        if not img_url:
+            return
+            
+        img_data = requests.get(img_url, timeout=30).content
+        img = Image.open(BytesIO(img_data))
+        
+        out = f"/tmp/reve_{index}_{threading.current_thread().ident}.png"
+        img.save(out)
+        
+        with lock:
+            results_list.append(out)
+        
+        # Notificar progreso si hay callback
+        if progress_callback:
+            progress_callback(index + 1)
+            
     except Exception as e:
-        return [f"Error REVE: {e}"]
+        print(f"Error generando imagen {index+1}: {e}")
 
-# =============================================================
-# INTERFAZ GRADIO (Optimizada)
-# =============================================================
+def reve_generate_multiple(prompt, key, model, num_images, progress_callback=None):
+    if not key:
+        return None
+    
+    num_images = min(num_images, 8)  # Máximo 8 imágenes
+    results = []
+    lock = threading.Lock()
+    
+    # Usamos ThreadPoolExecutor para generar imágenes en paralelo
+    with concurrent.futures.ThreadPoolExecutor(max_workers=min(num_images, 4)) as executor:
+        futures = []
+        for i in range(num_images):
+            future = executor.submit(
+                generate_single_reve_image,
+                prompt, key, model, i, results, lock, progress_callback
+            )
+            futures.append(future)
+        
+        # Esperar a que todas terminen
+        concurrent.futures.wait(futures)
+    
+    return results if results else None
+
+# ==============================
+#   UI COMPLETA
+# ==============================
 def build_ui():
     with gr.Blocks(title="BATUTO-ART MIX", theme=gr.themes.Soft()) as demo:
 
-        gr.Markdown("# 🎨 BATUTO-ART MIX – Generador Realista y Rápido en CPU")
+        gr.Markdown("""
+        # 🖼️ **BATUTO-ART MIX**
+        *Generador de imágenes con FLUX, Stable Diffusion 1.5 y REVE CREATE*
+        """)
 
         with gr.Tabs():
 
-            # =============================================================
-            # FLUX TAB
-            # =============================================================
-            with gr.TabItem("FLUX Optimizado"):
-                flux_prompt = gr.Textbox(label="Prompt")
-                flux_seed = gr.Number(label="Seed", value=1234)
-                flux_steps = gr.Slider(5, 60, value=20, step=1, label="Steps")
-                flux_btn = gr.Button("Generar")
-                flux_output = gr.Textbox(label="Resultado")
-
-                flux_btn.click(
-                    fn=flux_generate,
-                    inputs=[flux_prompt, flux_seed, flux_steps, gr.Number(value=576), gr.Number(value=1024)],
-                    outputs=[flux_output],
+            # ============================
+            # TAB: FLUX
+            # ============================
+            with gr.Tab("FLUX.2 / 1-Schnell"):
+                with gr.Row():
+                    with gr.Column(scale=1):
+                        flux_prompt = gr.Textbox(
+                            label="Prompt", 
+                            lines=3,
+                            placeholder="Describe la imagen que quieres generar con FLUX..."
+                        )
+                        model_select = gr.Dropdown([
+                            "black-forest-labs/FLUX.1-schnell",
+                            "black-forest-labs/FLUX.1-dev",
+                            "black-forest-labs/FLUX.2-dev"
+                        ], value="black-forest-labs/FLUX.1-schnell", label="Modelo FLUX")
+                        
+                        with gr.Row():
+                            steps = gr.Slider(5, 50, value=DEFAULT_STEPS, step=1, label="Steps")
+                            guidance = gr.Slider(0, 10, value=3, step=0.1, label="Guidance Scale")
+                        
+                        with gr.Row():
+                            width = gr.Number(value=DEFAULT_WIDTH, label="Width", precision=0)
+                            height = gr.Number(value=DEFAULT_HEIGHT, label="Height", precision=0)
+                        
+                        seed = gr.Number(value=0, label="Seed (0 = aleatorio)", precision=0)
+                        
+                        btn_flux = gr.Button("✨ Generar Imagen", variant="primary", size="lg")
+                        
+                    with gr.Column(scale=1):
+                        out_flux_img = gr.Image(
+                            label="Resultado FLUX",
+                            height=400,
+                            interactive=False
+                        )
+                        out_flux_file = gr.File(
+                            label="Descargar imagen",
+                            visible=False
+                        )
+                
+                # Acción del botón FLUX
+                def generate_flux_wrapper(model_name, prompt, steps, guidance, width, height, seed):
+                    if not prompt.strip():
+                        return None, "❌ Error: Ingresa un prompt válido"
+                    try:
+                        file_path = generate_flux(
+                            model_name, prompt, int(steps), float(guidance), 
+                            int(width), int(height), int(seed)
+                        )
+                        return file_path, gr.update(visible=True)
+                    except Exception as e:
+                        return None, f"❌ Error: {str(e)}"
+                
+                btn_flux.click(
+                    fn=generate_flux_wrapper,
+                    inputs=[model_select, flux_prompt, steps, guidance, width, height, seed],
+                    outputs=[out_flux_file, out_flux_file]
+                )
+                
+                # Mostrar imagen automáticamente
+                out_flux_file.change(
+                    fn=lambda f: Image.open(f) if f else None, 
+                    inputs=[out_flux_file], 
+                    outputs=[out_flux_img]
                 )
 
-            # =============================================================
-            # SD 1.5 TAB
-            # =============================================================
-            with gr.TabItem("Stable Diffusion 1.5 CPU"):
-                sd_prompt = gr.Textbox(label="Prompt")
-                sd_seed = gr.Number(label="Seed", value=1234)
-                sd_steps = gr.Slider(5, 50, value=20, step=1, label="Steps")
-                sd_btn = gr.Button("Generar Imagen")
-                sd_output = gr.Textbox(label="Resultado")
-
-                sd_btn.click(
-                    fn=sd15_generate,
-                    inputs=[sd_prompt, sd_seed, sd_steps, gr.Number(value=576), gr.Number(value=1024)],
-                    outputs=[sd_output],
+            # ============================
+            # TAB: SD1.5
+            # ============================
+            with gr.Tab("Stable Diffusion 1.5"):
+                with gr.Row():
+                    with gr.Column(scale=1):
+                        sd_prompt = gr.Textbox(
+                            label="Prompt", 
+                            lines=3,
+                            placeholder="Describe la imagen que quieres generar con SD1.5..."
+                        )
+                        
+                        with gr.Row():
+                            sd_steps = gr.Slider(5, 50, value=DEFAULT_STEPS, step=1, label="Steps")
+                            sd_guidance = gr.Slider(0, 10, value=3, step=0.1, label="Guidance Scale")
+                        
+                        with gr.Row():
+                            sd_width = gr.Number(value=DEFAULT_WIDTH, label="Width", precision=0)
+                            sd_height = gr.Number(value=DEFAULT_HEIGHT, label="Height", precision=0)
+                        
+                        sd_seed = gr.Number(value=0, label="Seed (0 = aleatorio)", precision=0)
+                        
+                        btn_sd = gr.Button("✨ Generar Imagen", variant="primary", size="lg")
+                        
+                    with gr.Column(scale=1):
+                        out_sd_img = gr.Image(
+                            label="Resultado SD1.5",
+                            height=400,
+                            interactive=False
+                        )
+                        out_sd_file = gr.File(
+                            label="Descargar imagen",
+                            visible=False
+                        )
+                
+                # Acción del botón SD1.5
+                def generate_sd_wrapper(prompt, steps, guidance, width, height, seed):
+                    if not prompt.strip():
+                        return None, "❌ Error: Ingresa un prompt válido"
+                    try:
+                        file_path = generate_sd(
+                            prompt, int(steps), float(guidance), 
+                            int(width), int(height), int(seed)
+                        )
+                        return file_path, gr.update(visible=True)
+                    except Exception as e:
+                        return None, f"❌ Error: {str(e)}"
+                
+                btn_sd.click(
+                    fn=generate_sd_wrapper,
+                    inputs=[sd_prompt, sd_steps, sd_guidance, sd_width, sd_height, sd_seed],
+                    outputs=[out_sd_file, out_sd_file]
+                )
+                
+                out_sd_file.change(
+                    fn=lambda f: Image.open(f) if f else None, 
+                    inputs=[out_sd_file], 
+                    outputs=[out_sd_img]
                 )
 
-            # =============================================================
-            # REVE CREATE TAB (Selector 1–8 imágenes)
-            # =============================================================
-            with gr.TabItem("REVE CREATE Avanzado"):
-                reve_api_key = gr.Textbox(label="API Key de REVE", type="password")
-                reve_prompt = gr.Textbox(label="Prompt")
-                reve_seed = gr.Number(label="Seed", value=1234)
-                reve_steps = gr.Slider(5, 50, value=20, step=1, label="Steps")
-                reve_num = gr.Slider(1, 8, value=1, step=1, label="Cantidad de Imágenes (1–8)")
-
-                reve_btn = gr.Button("Generar con REVE")
-                reve_gallery = gr.Gallery(label="Imágenes Generadas", columns=2, height=500)
-
-                reve_btn.click(
-                    fn=lambda key, p, s, st, n: reve_create_generate(key, p, s, st, int(n)),
-                    inputs=[reve_api_key, reve_prompt, reve_seed, reve_steps, reve_num],
-                    outputs=[reve_gallery],
+            # ============================
+            # TAB: REVE CREATE MODIFICADA
+            # ============================
+            with gr.Tab("REVE CREATE"):
+                with gr.Row():
+                    with gr.Column(scale=1):
+                        reve_api = gr.Textbox(
+                            label="API Key REVE", 
+                            type="password",
+                            placeholder="Ingresa tu API key de REVE",
+                            info="Necesitas una clave API válida de REVE"
+                        )
+                        reve_prompt = gr.Textbox(
+                            label="Prompt", 
+                            lines=3,
+                            placeholder="Describe la imagen que quieres generar..."
+                        )
+                        reve_model = gr.Dropdown([
+                            "reve-1",
+                            "reve-2",
+                            "reve-fast"
+                        ], value="reve-fast", label="Modelo REVE")
+                        
+                        # Slider para cantidad de imágenes
+                        num_images_slider = gr.Slider(
+                            minimum=1, 
+                            maximum=8, 
+                            value=4, 
+                            step=1,
+                            label="Cantidad de imágenes a generar"
+                        )
+                        
+                        with gr.Row():
+                            btn_reve = gr.Button("🚀 Generar Imágenes", variant="primary", size="lg")
+                            btn_clear = gr.Button("🗑️ Limpiar", variant="secondary")
+                        
+                        # Indicador de progreso
+                        progress_info = gr.Markdown("Esperando generación...")
+                        
+                    with gr.Column(scale=2):
+                        # Galería para mostrar múltiples imágenes
+                        gallery = gr.Gallery(
+                            label="Imágenes generadas",
+                            show_label=True,
+                            columns=4,
+                            rows=2,
+                            height="auto",
+                            object_fit="contain"
+                        )
+                        
+                        # Información de resultados
+                        result_info = gr.Markdown("")
+                        
+                        # Botón de descarga
+                        with gr.Row():
+                            btn_download_all = gr.Button(
+                                "📥 Descargar todas las imágenes", 
+                                size="lg",
+                                variant="secondary"
+                            )
+                            download_output = gr.File(
+                                label="Archivos descargables", 
+                                file_count="multiple",
+                                visible=False
+                            )
+                
+                # Estado para almacenar las rutas de archivos
+                last_files = gr.State([])
+                
+                # Función para generar múltiples imágenes con progreso
+                def generate_and_update_gallery(prompt, key, model, num_images, progress=gr.Progress()):
+                    if not key:
+                        return [], [], "❌ Error: Ingresa una API key válida", progress_info.update(value="")
+                    
+                    if not prompt.strip():
+                        return [], [], "❌ Error: Ingresa un prompt válido", progress_info.update(value="")
+                    
+                    progress_info.update(value="⏳ Iniciando generación...")
+                    
+                    # Función de callback para progreso
+                    def update_progress(current):
+                        progress_info.update(value=f"⏳ Generando imagen {current}/{num_images}...")
+                    
+                    files = reve_generate_multiple(prompt, key, model, num_images, update_progress)
+                    
+                    if files:
+                        # Crear lista de imágenes para la galería
+                        images = [(file,) for file in files]
+                        info_text = f"✅ Generadas {len(files)} de {num_images} imágenes"
+                        progress_info.update(value="✅ Generación completada")
+                        return images, files, info_text
+                    else:
+                        progress_info.update(value="❌ Error en la generación")
+                        return [], [], "❌ Error: No se pudieron generar imágenes. Verifica tu API key y conexión."
+                
+                # Acción del botón de generación
+                btn_reve.click(
+                    fn=generate_and_update_gallery,
+                    inputs=[reve_prompt, reve_api, reve_model, num_images_slider],
+                    outputs=[gallery, last_files, result_info]
+                )
+                
+                # Acción del botón de descarga
+                def prepare_download(files):
+                    if files:
+                        return gr.update(value=files, visible=True)
+                    return gr.update(value=[], visible=False)
+                
+                btn_download_all.click(
+                    fn=prepare_download,
+                    inputs=[last_files],
+                    outputs=[download_output]
+                )
+                
+                # Acción del botón de limpiar
+                def clear_gallery():
+                    return [], [], "✅ Galería limpiada", gr.update(value="Listo para nueva generación")
+                
+                btn_clear.click(
+                    fn=clear_gallery,
+                    outputs=[gallery, last_files, result_info, progress_info]
+                ).then(
+                    fn=lambda: gr.update(visible=False),
+                    outputs=[download_output]
                 )
 
     return demo
 
-# Ejecutar local o en HuggingFace
+# ==============================
+#   EJECUCIÓN PRINCIPAL
+# ==============================
 if __name__ == "__main__":
+    # Configuración para HuggingFace Spaces
+    import argparse
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--share", action="store_true", help="Crear enlace público")
+    parser.add_argument("--server-name", type=str, default="0.0.0.0", help="Dirección del servidor")
+    parser.add_argument("--server-port", type=int, default=7860, help="Puerto del servidor")
+    args = parser.parse_args()
+    
+    print("=" * 50)
+    print("🚀 Iniciando BATUTO-ART MIX")
+    print("=" * 50)
+    print(f"📱 Dispositivo: {DEVICE}")
+    print(f"⚡ Steps por defecto: {DEFAULT_STEPS}")
+    print(f"📐 Resolución por defecto: {DEFAULT_WIDTH}x{DEFAULT_HEIGHT}")
+    print("=" * 50)
+    
     demo = build_ui()
-    demo.launch()
+    demo.launch(
+        share=args.share,
+        server_name=args.server_name,
+        server_port=args.server_port,
+        show_error=True
+    )