import gradio as gr
import torch
from diffusers import DiffusionPipeline
import random

# --- Configuración del Modelo ---
model_id = "NewBie-AI/NewBie-image-Exp0.1"
device = "cuda" if torch.cuda.is_available() else "cpu"

print(f"Cargando modelo en: {device}...")

# Cargamos el pipeline. 
# Si usas GPU, usamos float16 para mayor velocidad y menos memoria.
dtype = torch.float16 if device == "cuda" else torch.float32

try:
    pipe = DiffusionPipeline.from_pretrained(
        model_id, 
        torch_dtype=dtype,
        use_safetensors=True
    )
    pipe.to(device)
    print("Modelo cargado exitosamente.")
except Exception as e:
    print(f"Error cargando el modelo: {e}")

# --- Función de Generación ---
def generate_image(prompt, negative_prompt, steps, guidance_scale, width, height, seed):
    if seed == -1:
        seed = random.randint(0, 2147483647)
    
    # Configuramos el generador para reproducibilidad
    generator = torch.Generator(device).manual_seed(int(seed))
    
    print(f"Generando con semilla: {seed}")
    
    try:
        image = pipe(
            prompt=prompt,
            negative_prompt=negative_prompt,
            num_inference_steps=int(steps),
            guidance_scale=guidance_scale,
            width=int(width),
            height=int(height),
            generator=generator
        ).images[0]
        return image, seed
    except Exception as e:
        return None, f"Error: {str(e)}"

# --- Interfaz de Gradio ---
css = """
#col-container {max-width: 800px; margin-left: auto; margin-right: auto;}
"""

with gr.Blocks(css=css) as demo:
    with gr.Column(elem_id="col-container"):
        gr.Markdown(f"# 🎨 Generador de Imágenes: {model_id}")
        gr.Markdown("Escribe un prompt para generar una imagen usando el modelo NewBie-AI.")
        
        with gr.Group():
            prompt = gr.Textbox(label="Prompt (Descripción positiva)", placeholder="Un astronauta montando un caballo en marte, 4k, realista...", lines=2)
            negative_prompt = gr.Textbox(label="Prompt Negativo (Lo que NO quieres)", placeholder="borroso, deforme, mala calidad, texto...", value="bad quality, worst quality, low resolution, blurry, distorted")
            
        with gr.Row():
            run_button = gr.Button("Generar Imagen", variant="primary", scale=1)
        
        with gr.Row():
            result_image = gr.Image(label="Resultado", interactive=False)
        
        with gr.Accordion("Configuración Avanzada", open=False):
            with gr.Row():
                width = gr.Slider(label="Ancho", minimum=256, maximum=1024, step=64, value=512)
                height = gr.Slider(label="Alto", minimum=256, maximum=1024, step=64, value=512)
            
            with gr.Row():
                steps = gr.Slider(label="Pasos de Inferencia", minimum=10, maximum=100, step=1, value=25)
                guidance_scale = gr.Slider(label="Guidance Scale (Fidelidad al prompt)", minimum=1, maximum=20, step=0.5, value=7.5)
            
            seed = gr.Number(label="Semilla (Seed) - Usa -1 para aleatorio", value=-1, precision=0)
            seed_output = gr.Number(label="Semilla utilizada", interactive=False)

    # Eventos
    run_button.click(
        fn=generate_image,
        inputs=[prompt, negative_prompt, steps, guidance_scale, width, height, seed],
        outputs=[result_image, seed_output]
    )

# Lanzar la app
if __name__ == "__main__":
    demo.launch()