import os 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 # ============================== # CONFIGURACIÓN BASE CPU # ============================== DEVICE = "cpu" torch.set_grad_enabled(False) DEFAULT_STEPS = 15 DEFAULT_WIDTH = 512 DEFAULT_HEIGHT = 768 # Cache de modelos MODEL_CACHE = {} # ============================== # CARGA GENÉRICA DE MODELOS # ============================== def load_pipeline(model_id): if model_id not in MODEL_CACHE: pipe = DiffusionPipeline.from_pretrained( model_id, torch_dtype=torch.float32 ) pipe.to(DEVICE) pipe.enable_attention_slicing() MODEL_CACHE[model_id] = pipe return MODEL_CACHE[model_id] # ============================== # GENERACIÓN DIFFUSERS # ============================== def generate_diffusion(model_id, prompt, steps, guidance, width, height, seed): pipe = load_pipeline(model_id) generator = None if seed and int(seed) != 0: generator = torch.manual_seed(int(seed)) image = pipe( prompt=prompt, num_inference_steps=int(steps), guidance_scale=float(guidance), width=int(width), height=int(height), generator=generator, ).images[0] # Nombre único por modelo safe_name = model_id.split("/")[-1].replace(".", "_") out = f"/tmp/{safe_name}_output.png" image.save(out) return out # ============================== # REVE CREATE (PARALELO) # ============================== def generate_single_reve_image(prompt, key, model, index, results_list, lock): try: 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) except Exception as e: print(f"Error generando imagen {index+1}: {e}") def reve_generate_multiple(prompt, key, model, num_images): if not key: return None num_images = min(int(num_images), 8) results = [] lock = threading.Lock() with concurrent.futures.ThreadPoolExecutor( max_workers=min(num_images, 4) ) as executor: futures = [ executor.submit( generate_single_reve_image, prompt, key, model, i, results, lock, ) for i in range(num_images) ] concurrent.futures.wait(futures) return results if results else None # ============================== # UI COMPLETA # ============================== def build_ui(): with gr.Blocks(title="BATUTO-ART MIX") as demo: gr.Markdown( """ # 🖼️ **BATUTO-ART MIX** *Generador de imágenes con FLUX, Stable Diffusion 1.5 y REVE CREATE* """ ) with gr.Tabs(): # ============================ # 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...", ) flux_model = 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(): flux_steps = gr.Slider( 5, 50, value=DEFAULT_STEPS, step=1, label="Steps" ) flux_guidance = gr.Slider( 0, 10, value=3, step=0.1, label="Guidance Scale", ) with gr.Row(): flux_width = gr.Number( value=DEFAULT_WIDTH, label="Width", precision=0 ) flux_height = gr.Number( value=DEFAULT_HEIGHT, label="Height", precision=0 ) flux_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): flux_img = gr.Image( label="Resultado FLUX", height=400, interactive=False, ) flux_file = gr.File( label="Descargar imagen", visible=False, ) def flux_wrapper(model_id, prompt, steps, guidance, width, height, seed): if not prompt or not prompt.strip(): return None, gr.update(visible=False) try: path = generate_diffusion( model_id, prompt, steps, guidance, width, height, seed ) return path, gr.update(visible=True) except Exception as e: print(f"Error FLUX: {e}") return None, gr.update(visible=False) btn_flux.click( flux_wrapper, inputs=[ flux_model, flux_prompt, flux_steps, flux_guidance, flux_width, flux_height, flux_seed, ], outputs=[flux_file, flux_file], ) flux_file.change( lambda f: Image.open(f) if f else None, inputs=flux_file, outputs=flux_img, ) # ============================ # 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): sd_img = gr.Image( label="Resultado SD1.5", height=400, interactive=False, ) sd_file = gr.File( label="Descargar imagen", visible=False, ) def sd_wrapper(prompt, steps, guidance, width, height, seed): if not prompt or not prompt.strip(): return None, gr.update(visible=False) try: path = generate_diffusion( "runwayml/stable-diffusion-v1-5", prompt, steps, guidance, width, height, seed, ) return path, gr.update(visible=True) except Exception as e: print(f"Error SD1.5: {e}") return None, gr.update(visible=False) btn_sd.click( sd_wrapper, inputs=[ sd_prompt, sd_steps, sd_guidance, sd_width, sd_height, sd_seed, ], outputs=[sd_file, sd_file], ) sd_file.change( lambda f: Image.open(f) if f else None, inputs=sd_file, outputs=sd_img, ) # ============================ # TAB: REVE CREATE # ============================ 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", ) 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", ) num_images = gr.Slider( 1, 8, value=4, step=1, label="Cantidad de imágenes a generar", ) btn_reve = gr.Button( "🚀 Generar Imágenes", variant="primary", size="lg" ) btn_clear = gr.Button("🗑️ Limpiar", variant="secondary") progress_info = gr.Markdown("Esperando generación...") with gr.Column(scale=2): gallery = gr.Gallery( label="Imágenes generadas", show_label=True, columns=4, rows=2, height="auto", object_fit="contain", ) result_info = gr.Markdown("") 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, ) last_files = gr.State([]) def generate_and_update_gallery(prompt, key, model, n, progress=gr.Progress()): if not key: return ( [], [], "❌ Error: Ingresa una API key válida", gr.update(value="❌ Falta API key"), ) if not prompt or not prompt.strip(): return ( [], [], "❌ Error: Ingresa un prompt válido", gr.update(value="❌ Prompt vacío o inválido"), ) progress.tqdm(total=1, desc="Generando con REVE") files = reve_generate_multiple(prompt, key, model, n) if files: imgs = [(f,) for f in files] info = f"✅ Generadas {len(files)} imágenes" return imgs, files, info, gr.update( value="✅ Generación completada" ) else: return ( [], [], "❌ Error: No se pudieron generar imágenes. Revisa API key y conexión.", gr.update(value="❌ Error en la generación"), ) btn_reve.click( generate_and_update_gallery, inputs=[reve_prompt, reve_api, reve_model, num_images], outputs=[gallery, last_files, result_info, progress_info], ) def prepare_download(files): if files: return gr.update(value=files, visible=True) return gr.update(value=[], visible=False) btn_download_all.click( prepare_download, inputs=last_files, outputs=download_output ) def clear_gallery(): return ( [], [], "✅ Galería limpiada", gr.update(value="Listo para nueva generación"), ) btn_clear.click( clear_gallery, outputs=[gallery, last_files, result_info, progress_info], ).then( lambda: gr.update(visible=False), outputs=download_output, ) return demo if __name__ == "__main__": 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( share=args.share, server_name=args.server_name, server_port=args.server_port, show_error=True, )