import os import gradio as gr import torch from transformers import pipeline, AutoTokenizer, AutoModelForCausalLM from diffusers import StableDiffusionPipeline, DPMSolverMultistepScheduler # ========================================== # 1. CONFIGURACIÓN Y CARGA DE MODELOS # ========================================== # Variables de entorno MODEL_TEXT_NAME = os.environ.get("MODEL_NAME", "gpt2") MODEL_IMG_ID = os.environ.get("MODEL_ID", "runwayml/stable-diffusion-v1-5") HF_TOKEN = os.environ.get("HF_TOKEN", None) # Carga del modelo de texto def load_text_generator(model_name=MODEL_TEXT_NAME): try: tokenizer = AutoTokenizer.from_pretrained(model_name) model = AutoModelForCausalLM.from_pretrained(model_name) device = 0 if torch.cuda.is_available() else -1 return pipeline("text-generation", model=model, tokenizer=tokenizer, device=device) except Exception as e: print("Error cargando modelo de texto, intentando pipeline directo:", e) device = 0 if torch.cuda.is_available() else -1 return pipeline("text-generation", model=model_name, device=device) # Carga del modelo de imagen def load_image_pipeline(model_id=MODEL_IMG_ID, token=HF_TOKEN): device = "cuda" if torch.cuda.is_available() else "cpu" try: scheduler = DPMSolverMultistepScheduler.from_pretrained(model_id, subfolder="scheduler", use_auth_token=token) if token else DPMSolverMultistepScheduler.from_pretrained(model_id, subfolder="scheduler") pipe = StableDiffusionPipeline.from_pretrained(model_id, torch_dtype=torch.float16 if device=="cuda" else torch.float32, use_auth_token=token) pipe.scheduler = scheduler if device == "cuda": pipe = pipe.to(device) pipe.enable_attention_slicing() return pipe except Exception as e: print(f"Error cargando modelo de imagen (¿falta GPU o token?): {e}") return None print("Iniciando secuencia de arranque de HECTRON...") generador_texto = load_text_generator() pipe_imagen = load_image_pipeline() print("Sistemas cargados.") # ========================================== # 2. FUNCIONES DE GENERACIÓN # ========================================== def generar_texto(prompt, max_length=150, temperature=0.7, top_k=50, top_p=0.95, num_return_sequences=1): if not prompt or prompt.strip() == "": return "Inicia la consola de HECTRON con un prompt válido." salida = generador_texto( prompt, max_length=int(max_length), temperature=float(temperature), top_k=int(top_k), top_p=float(top_p), num_return_sequences=int(num_return_sequences), do_sample=True, ) textos = [s["generated_text"] for s in salida] return textos[0] if len(textos) == 1 else "\n\n---\n\n".join(textos) def generar_imagen(prompt, negative_prompt, steps, guidance_scale, width, height, seed): if pipe_imagen is None: return None # Falla segura si no cargó el modelo if not prompt or prompt.strip() == "": return None generator = torch.Generator(device="cuda") if torch.cuda.is_available() else torch.Generator() if seed is not None and seed != "": try: generator = generator.manual_seed(int(seed)) except: pass image = pipe_imagen( prompt=prompt, negative_prompt=negative_prompt or None, num_inference_steps=int(steps), guidance_scale=float(guidance_scale), height=int(height), width=int(width), generator=generator ).images[0] return image # ========================================== # 3. INTERFAZ GRÁFICA (GRADIO) # ========================================== with gr.Blocks(title="HECTRON - Sistema Unificado", theme=gr.themes.Monochrome()) as demo: gr.Markdown("# 🧠 HECTRON: Terminal de Control Unificada") gr.Markdown("Ecosistema integrado. Selecciona el módulo operativo en las pestañas inferiores.") # PESTAÑA 1: TEXTO with gr.Tab("Módulo de Lenguaje (Texto)"): with gr.Row(): with gr.Column(scale=3): prompt_text = gr.Textbox(lines=6, label="Directiva del Sistema", placeholder="Ej. Ejecuta MODO GUARDIÁN 22:00...", value="HECTRON activado. Analiza el siguiente concepto: ") btn_text = gr.Button("Ejecutar Secuencia de Texto") with gr.Column(scale=1): max_len = gr.Slider(minimum=20, maximum=1024, value=150, step=10, label="Límite de Tokens") temp = gr.Slider(minimum=0.1, maximum=1.5, value=0.7, step=0.05, label="Temperatura") topk = gr.Slider(minimum=0, maximum=200, value=50, step=1, label="Top-k") topp = gr.Slider(minimum=0.0, maximum=1.0, value=0.95, step=0.01, label="Top-p") n_seq = gr.Slider(minimum=1, maximum=5, value=1, step=1, label="Secuencias") output_text = gr.Textbox(lines=12, label="Respuesta de HECTRON") btn_text.click(fn=generar_texto, inputs=[prompt_text, max_len, temp, topk, topp, n_seq], outputs=output_text) # PESTAÑA 2: IMAGEN with gr.Tab("Módulo de Síntesis Visual (Imagen)"): if pipe_imagen is None: gr.Markdown("⚠️ **Advertencia:** El modelo de imagen no pudo cargarse. Revisa la memoria de tu sistema o tu token de Hugging Face.") with gr.Row(): with gr.Column(scale=3): prompt_img = gr.Textbox(lines=4, label="Directiva Visual (Prompt)", value="A futuristic AI core named Hectron, cyberpunk style, glowing neon lights, highly detailed anime art style") neg_img = gr.Textbox(lines=2, label="Filtro Negativo", value="lowres, bad anatomy, text, error, blurry, deformed") btn_img = gr.Button("Renderizar Imagen") with gr.Column(scale=1): steps = gr.Slider(minimum=1, maximum=100, value=30, step=1, label="Pasos de Renderizado") guidance = gr.Slider(minimum=1.0, maximum=20.0, value=7.5, step=0.1, label="Fidelidad (Guidance)") w_img = gr.Dropdown(choices=["512","640","768"], value="512", label="Ancho (px)") h_img = gr.Dropdown(choices=["512","640","768"], value="512", label="Alto (px)") seed_img = gr.Textbox(lines=1, label="Semilla (Seed)") output_img = gr.Image(label="Renderizado de HECTRON") btn_img.click(fn=generar_imagen, inputs=[prompt_img, neg_img, steps, guidance, w_img, h_img, seed_img], outputs=output_img) if __name__ == "__main__": demo.launch(server_name="0.0.0.0", server_port=int(os.environ.get("PORT", 7860)))