Create app.py

app.py

@@ -0,0 +1,411 @@
+# app.py - BATUTO-ART MIX (Optimizado para HF Spaces - CPU)
+import os
+import io
+import time
+import random
+import threading
+import requests
+from pathlib import Path
+from typing import Optional
+from PIL import Image
+import gradio as gr
+
+# Diffusers imports (algunas pueden no usarse en CPU si no están instaladas)
+try:
+    import torch
+    from diffusers import FluxPipeline, StableDiffusionPipeline
+except Exception as e:
+    # En HF Spaces instala diffusers + accelerate + transformers y torch a la versión adecuada
+    print("Warning: diffusers/torch no cargados:", e)
+
+# -----------------------
+# Config y rutas
+# -----------------------
+OUT_DIR = Path("/tmp/batuto_art_mix")
+OUT_DIR.mkdir(parents=True, exist_ok=True)
+
+# Defaults 9:16
+DEFAULT_WIDTH = 576
+DEFAULT_HEIGHT = 1024
+
+# Hugging Face token (si quieres usar from env para algunos modelos)
+HF_TOKEN = os.getenv("HF_TOKEN", "")
+
+# REVE API defaults (si quieres usar la que mencionaste ajusta aquí)
+REVE_API_URL = "https://api.reve.com/v1/image/create"
+
+# -----------------------
+# Cache de pipelines (global)
+# -----------------------
+_PIPELINE_CACHE = {
+    "flux": {},
+    "sd": None
+}
+_PIPELINE_LOCK = threading.Lock()
+
+# -----------------------
+# Utilities
+# -----------------------
+def save_image_to_file(img: Image.Image, prefix="batuto", fmt="png"):
+    ts = int(time.time() * 1000)
+    fn = OUT_DIR / f"{prefix}_{ts}.png"
+    img.save(fn, format="PNG")
+    return str(fn)
+
+def pil_image_to_bytes(img: Image.Image):
+    bio = io.BytesIO()
+    img.save(bio, format="PNG")
+    bio.seek(0)
+    return bio
+
+def random_seed():
+    return random.randint(0, 2**31 - 1)
+
+# -----------------------
+# Loaders: FLUX
+# -----------------------
+def load_flux_pipeline(model_id: str):
+    """
+    Carga (y cachea) un FluxPipeline para CPU.
+    Usa torch.float32 para compatibilidad CPU.
+    """
+    if "FluxPipeline" not in globals():
+        raise RuntimeError("diffusers FluxPipeline no está disponible en este entorno.")
+    with _PIPELINE_LOCK:
+        if model_id in _PIPELINE_CACHE["flux"]:
+            return _PIPELINE_CACHE["flux"][model_id]
+        # Intentar cargar
+        dtype = torch.float32
+        kwargs = {}
+        if HF_TOKEN:
+            kwargs["use_auth_token"] = HF_TOKEN
+        pipe = FluxPipeline.from_pretrained(model_id, torch_dtype=dtype, **kwargs)
+        # Forzar CPU
+        pipe.to("cpu")
+        # Optimizaciones para CPU
+        try:
+            pipe.enable_sequential_cpu_offload()
+        except Exception:
+            pass
+        try:
+            pipe.enable_attention_slicing()
+        except Exception:
+            pass
+        _PIPELINE_CACHE["flux"][model_id] = pipe
+        return pipe
+
+# -----------------------
+# Loaders: Stable Diffusion (advertencia: pesado en CPU)
+# -----------------------
+def load_sd_pipeline(model_id: str = "runwayml/stable-diffusion-v1-5"):
+    if "StableDiffusionPipeline" not in globals():
+        raise RuntimeError("StableDiffusionPipeline no está disponible en este entorno.")
+    with _PIPELINE_LOCK:
+        if _PIPELINE_CACHE["sd"] is not None:
+            return _PIPELINE_CACHE["sd"]
+        dtype = torch.float32
+        kwargs = dict()
+        if HF_TOKEN:
+            kwargs["use_auth_token"] = HF_TOKEN
+        pipe = StableDiffusionPipeline.from_pretrained(model_id, torch_dtype=dtype, safety_checker=None, **kwargs)
+        pipe.to("cpu")
+        try:
+            pipe.enable_attention_slicing()
+        except Exception:
+            pass
+        _PIPELINE_CACHE["sd"] = pipe
+        return pipe
+
+# -----------------------
+# Generación: FLUX
+# -----------------------
+def generate_flux(model_id: str, prompt: str, negative_prompt: str,
+                  steps: int, guidance: float, width: int, height: int,
+                  seed: int):
+    if not prompt or prompt.strip() == "":
+        return None, None, "Prompt vacío"
+
+    try:
+        pipe = load_flux_pipeline(model_id)
+    except Exception as e:
+        return None, None, f"Error cargando FLUX: {e}"
+
+    if seed is None or int(seed) == -1:
+        seed = random_seed()
+    gen = torch.Generator(device="cpu").manual_seed(int(seed))
+
+    try:
+        out = pipe(
+            prompt,
+            negative_prompt=negative_prompt if negative_prompt else None,
+            num_inference_steps=int(steps),
+            guidance_scale=float(guidance),
+            height=int(height),
+            width=int(width),
+            generator=gen
+        )
+        img = out.images[0]
+        path = save_image_to_file(img, prefix="flux")
+        return img, path, str(seed)
+    except Exception as e:
+        return None, None, f"Error generación FLUX: {e}"
+
+# -----------------------
+# Generación: Stable Diffusion (ligera, en CPU puede ser muy lenta)
+# -----------------------
+def generate_sd(prompt: str, negative_prompt: str, steps: int, guidance: float,
+                width: int, height: int, seed: int):
+    if not prompt or prompt.strip() == "":
+        return None, None, "Prompt vacío"
+    try:
+        pipe = load_sd_pipeline()
+    except Exception as e:
+        return None, None, f"Error cargando SD: {e}"
+
+    if seed is None or int(seed) == -1:
+        seed = random_seed()
+    gen = torch.Generator(device="cpu").manual_seed(int(seed))
+
+    try:
+        out = pipe(
+            prompt=prompt,
+            negative_prompt=negative_prompt if negative_prompt else None,
+            num_inference_steps=int(steps),
+            guidance_scale=float(guidance),
+            height=int(height),
+            width=int(width),
+            generator=gen
+        )
+        img = out.images[0]
+        path = save_image_to_file(img, prefix="sd")
+        return img, path, str(seed)
+    except Exception as e:
+        return None, None, f"Error generación SD: {e}"
+
+# -----------------------
+# Generación: REVE Create (API)
+# -----------------------
+def generate_reve(prompt: str, api_key: str, ratio: str, version: str, num_images: int,
+                  width: int, height: int):
+    """
+    Llama a REVE API. Se asume que la API acepta:
+    payload = {"prompt": prompt, "aspect_ratio": ratio, "version": version}
+    Responde con base64 en 'image' en JSON (según ejemplo previo).
+    """
+    if not api_key or api_key.strip() == "":
+        return [], "Falta API Key"
+
+    payload = {
+        "prompt": prompt,
+        "aspect_ratio": ratio,
+        "version": version or "latest",
+        "num_images": num_images,
+        "width": width,
+        "height": height
+    }
+    headers = {
+        "Authorization": f"Bearer {api_key}",
+        "Accept": "application/json",
+        "Content-Type": "application/json"
+    }
+    results = []
+    errors = []
+    for i in range(int(num_images)):
+        try:
+            resp = requests.post(REVE_API_URL, headers=headers, json=payload, timeout=60)
+            if resp.status_code == 200:
+                data = resp.json()
+                if "image" in data:
+                    import base64
+                    img_bytes = base64.b64decode(data["image"])
+                    img = Image.open(io.BytesIO(img_bytes)).convert("RGBA").convert("RGB")
+                    path = save_image_to_file(img, prefix="reve")
+                    results.append((img, path))
+                else:
+                    errors.append(f"Respuesta sin 'image': {data.keys()}")
+            else:
+                errors.append(f"{resp.status_code} {resp.text[:200]}")
+        except Exception as e:
+            errors.append(str(e))
+    if errors and not results:
+        return [], " | ".join(errors[:3])
+    return results, f"{len(results)} generadas, errores: {len(errors)}"
+
+# -----------------------
+# UI - Gradio
+# -----------------------
+def build_ui():
+    with gr.Blocks(title="BATUTO-ART MIX", theme=gr.themes.Soft()) as demo:
+        gr.Markdown("# 🎨 BATUTO-ART MIX — Generador (Optimizado para CPU, 9:16)")
+
+        # Estado global de historial
+        historial_state = gr.State([])
+
+        with gr.Tabs():
+            # ----------------- TAB FLUX (multi-model selector) -----------------
+            with gr.TabItem("FLUX Models"):
+                with gr.Row():
+                    with gr.Column(scale=2):
+                        model_flux = gr.Dropdown(
+                            choices=[
+                                "black-forest-labs/FLUX.1-dev",
+                                "black-forest-labs/FLUX.1-schnell",
+                                "black-forest-labs/FLUX.2-dev"
+                            ],
+                            label="Selecciona Modelo FLUX",
+                            value="black-forest-labs/FLUX.1-schnell"
+                        )
+                        prompt_flux = gr.Textbox(label="Prompt", lines=4, placeholder="Describe la escena...")
+                        negative_flux = gr.Textbox(label="Negative Prompt (opcional)", lines=2)
+                        with gr.Row():
+                            steps_flux = gr.Slider(4, 50, value=12, step=1, label="Steps")
+                            guidance_flux = gr.Slider(1.0, 20.0, value=4.0, step=0.5, label="Guidance")
+                        with gr.Row():
+                            width_flux = gr.Number(value=DEFAULT_WIDTH, label="Width (px)")
+                            height_flux = gr.Number(value=DEFAULT_HEIGHT, label="Height (px)")
+                        seed_flux = gr.Number(value=-1, label="Seed (-1 random)")
+                        with gr.Row():
+                            gen_flux_btn = gr.Button("✨ Generar (FLUX)", variant="primary")
+                            clear_flux_btn = gr.Button("🧹 Limpiar Prompt", variant="secondary")
+                            download_flux_btn = gr.Button("⬇️ Guardar Última", variant="secondary")
+                        flux_status = gr.Markdown("")
+
+                    with gr.Column(scale=3):
+                        out_flux_img = gr.Image(label="Imagen FLUX", type="pil", show_label=True)
+                        out_flux_seed = gr.Textbox(label="Semilla usada", interactive=False)
+                        flux_last_file = gr.Textbox(visible=False)  # internal store last path
+                        # Galería/Historial
+                        flux_gallery = gr.Gallery(label="Historial FLUX (sesión)", columns=4, object_fit="cover", height=300)
+
+                # Conexiones FLUX
+                def _clear_flux():
+                    return "", ""
+                clear_flux_btn.click(fn=_clear_flux, inputs=None, outputs=[prompt_flux, negative_flux])
+
+                def _download_last(flux_last):
+                    if flux_last:
+                        return gr.File.update(value=flux_last)
+                    return None
+
+                download_file_flux = gr.File(label="Descargar imagen", visible=True)
+                download_flux_btn.click(fn=_download_last, inputs=[flux_last_file], outputs=[download_file_flux])
+
+                def _gen_flux(model_id, prompt, negative, steps, guidance, width, height, seed, hist):
+                    img, path, info = generate_flux(model_id, prompt, negative, steps, guidance, width, height, seed)
+                    if img is None:
+                        return None, "", "", hist, f"❌ {info}"
+                    # actualizar historial
+                    new_hist = (hist or []) + [img]
+                    return img, str(info), path, new_hist, f"✅ Generada | seed: {info}"
+
+                gen_flux_btn.click(
+                    fn=_gen_flux,
+                    inputs=[model_flux, prompt_flux, negative_flux, steps_flux, guidance_flux, width_flux, height_flux, seed_flux, historial_state],
+                    outputs=[out_flux_img, out_flux_seed, flux_last_file, flux_gallery, flux_status]
+                ).then(lambda h: h, inputs=[historial_state], outputs=[historial_state])
+
+            # ----------------- TAB Stable Diffusion -----------------
+            with gr.TabItem("Stable Diffusion (opcional)"):
+                gr.Markdown("**Aviso:** Stable Diffusion en CPU puede ser muy lento. Úsalo si no tienes prisa.")
+                with gr.Row():
+                    with gr.Column(scale=2):
+                        prompt_sd = gr.Textbox(label="Prompt", lines=4)
+                        negative_sd = gr.Textbox(label="Negative Prompt (opcional)", lines=2)
+                        with gr.Row():
+                            steps_sd = gr.Slider(10, 50, value=20, step=1, label="Steps")
+                            guidance_sd = gr.Slider(1.0, 20.0, value=7.5, step=0.5)
+                        with gr.Row():
+                            width_sd = gr.Number(value=DEFAULT_WIDTH, label="Width")
+                            height_sd = gr.Number(value=DEFAULT_HEIGHT, label="Height")
+                        seed_sd = gr.Number(value=-1, label="Seed (-1 random)")
+                        with gr.Row():
+                            gen_sd_btn = gr.Button("✨ Generar (SD)", variant="primary")
+                            clear_sd_btn = gr.Button("🧹 Limpiar Prompt", variant="secondary")
+                            download_sd_btn = gr.Button("⬇️ Guardar Última", variant="secondary")
+                        sd_status = gr.Markdown()
+                    with gr.Column(scale=3):
+                        out_sd_img = gr.Image(label="Imagen SD", type="pil")
+                        out_sd_seed = gr.Textbox(label="Semilla usada", interactive=False)
+                        sd_last_file = gr.Textbox(visible=False)
+                        sd_gallery = gr.Gallery(label="Historial SD (sesión)", columns=4, object_fit="cover", height=300)
+
+                clear_sd_btn.click(fn=lambda: ("", ""), inputs=None, outputs=[prompt_sd, negative_sd])
+
+                download_file_sd = gr.File(label="Descargar imagen SD", visible=True)
+                download_sd_btn.click(fn=lambda p: gr.File.update(value=p) if p else None, inputs=[sd_last_file], outputs=[download_file_sd])
+
+                def _gen_sd(prompt, negative, steps, guidance, width, height, seed, hist):
+                    img, path, info = generate_sd(prompt, negative, steps, guidance, width, height, seed)
+                    if img is None:
+                        return None, "", "", hist, f"❌ {info}"
+                    new_hist = (hist or []) + [img]
+                    return img, str(info), path, new_hist, f"✅ Generada | seed: {info}"
+
+                gen_sd_btn.click(
+                    fn=_gen_sd,
+                    inputs=[prompt_sd, negative_sd, steps_sd, guidance_sd, width_sd, height_sd, seed_sd, historial_state],
+                    outputs=[out_sd_img, out_sd_seed, sd_last_file, sd_gallery, sd_status]
+                ).then(lambda h: h, inputs=[historial_state], outputs=[historial_state])
+
+            # ----------------- TAB REVE CREATE -----------------
+            with gr.TabItem("REVE Create (API)"):
+                gr.Markdown("Aquí puedes pegar/actualizar tu **API Key** para REVE y generar desde el front-end.")
+                with gr.Row():
+                    with gr.Column(scale=2):
+                        reve_api_input = gr.Textbox(label="🔑 REVE API Key (se puede pegar o dejar vacío para usar env var)", type="password")
+                        reve_prompt = gr.Textbox(label="Prompt REVE", lines=4)
+                        with gr.Row():
+                            reve_ratio = gr.Dropdown(choices=["9:16", "16:9", "1:1"], value="9:16", label="Aspect Ratio")
+                            reve_num = gr.Slider(1, 4, value=1, step=1, label="Cantidad")
+                        reve_version = gr.Textbox(label="Version (dejar 'latest' o especificar)", value="latest")
+                        with gr.Row():
+                            reve_width = gr.Number(value=DEFAULT_WIDTH, label="Width")
+                            reve_height = gr.Number(value=DEFAULT_HEIGHT, label="Height")
+                        reve_gen_btn = gr.Button("🚀 Generar REVE")
+                        reve_clear_btn = gr.Button("🧹 Limpiar Prompt REVE", variant="secondary")
+                        reve_status = gr.Markdown()
+                    with gr.Column(scale=3):
+                        reve_gallery = gr.Gallery(label="Resultados REVE", columns=3, object_fit="cover", height=300)
+                        reve_last_file = gr.Textbox(visible=False)
+                        reve_download = gr.File(label="Descargar REVE", visible=True)
+
+                reve_clear_btn.click(fn=lambda: "", inputs=None, outputs=[reve_prompt])
+
+                def _gen_reve(prompt, api_key_ui, ratio, version, num, width, height, hist):
+                    api_key = api_key_ui.strip() if api_key_ui and api_key_ui.strip() else os.getenv("REVE_API_KEY", "")
+                    res, msg = generate_reve(prompt, api_key, ratio, version, int(num), int(width), int(height))
+                    if isinstance(res, list) and res:
+                        imgs = [r[0] for r in res]
+                        # paths maybe multiple -> store last
+                        last_path = res[-1][1]
+                        new_hist = (hist or []) + imgs
+                        return imgs, last_path, new_hist, f"✅ {msg}"
+                    else:
+                        return [], "", hist, f"❌ {msg}"
+
+                reve_gen_btn.click(
+                    fn=_gen_reve,
+                    inputs=[reve_prompt, reve_api_input, reve_ratio, reve_version, reve_num, reve_width, reve_height, historial_state],
+                    outputs=[reve_gallery, reve_last_file, historial_state, reve_status]
+                )
+
+                reve_download.click(fn=lambda p: gr.File.update(value=p) if p else None, inputs=[reve_last_file], outputs=[reve_download])
+
+        # Footer / recomendaciones
+        gr.Markdown("---")
+        gr.Markdown("""
+        **Notas de despliegue:**
+        - Recomendado para **HF Spaces (CPU)**: usa `black-forest-labs/FLUX.1-schnell` por default.
+        - Stable Diffusion en CPU será **muy lento** y puede hacer timeouts en Spaces; úsalo con cuidado.
+        - REVE Create requiere tu API Key — puedes pegarla en la caja o ponerla en `REVE_API_KEY` en los Secrets del Space.
+        - Tamaño 9:16 por defecto: 576×1024 (balance calidad / memoria).
+        """)
+
+    return demo
+
+# -----------------------
+# Run
+# -----------------------
+if __name__ == "__main__":
+    demo = build_ui()
+    demo.launch(server_name="0.0.0.0", server_port=int(os.environ.get("PORT", 7860)))