import gradio as gr
import torch
from controlnet_aux import CannyDetector
from diffusers import FluxControlPipeline
from diffusers.utils import load_image
import random
import os
from huggingface_hub import login

# Autenticarse en HuggingFace
login(os.getenv("HUGGINGFACEHUB_TOKEN"))

# Configuración del modelo
model_id = "black-forest-labs/FLUX.1-Canny-dev"
pipe = FluxControlPipeline.from_pretrained(model_id, torch_dtype=torch.bfloat16).to("cuda")

# Procesador Canny para detección de bordes
processor = CannyDetector()

MAX_SEED = 2**32 - 1
MAX_IMAGE_SIZE = 1024

def generate_image(image, prompt, seed, random_seed, low_thresh, high_thresh, guidance_scale, num_inference_steps):
    # Control de Seed
    if random_seed:
        seed = random.randint(0, MAX_SEED)
    generator = torch.manual_seed(seed)
    
    # Cargar imagen y aplicar Canny
    control_image = processor(image, low_threshold=low_thresh, high_threshold=high_thresh, 
                              detect_resolution=MAX_IMAGE_SIZE, image_resolution=MAX_IMAGE_SIZE)
    
    # Generar imagen con Diffusers
    generated_image = pipe(
        prompt=prompt,
        control_image=control_image,
        height=MAX_IMAGE_SIZE,
        width=MAX_IMAGE_SIZE,
        num_inference_steps=num_inference_steps,
        guidance_scale=guidance_scale,
        generator=generator,
    ).images[0]
    
    return generated_image, seed

# Interfaz Gradio
with gr.Blocks() as demo:
    gr.Markdown("# Generador de Imágenes con FLUX.1 Canny 🚀")
    
    with gr.Row():
        image_input = gr.Image(type="pil", label="Sube tu imagen de control")
        result_image = gr.Image(label="Imagen Generada")
    
    with gr.Row():
        prompt = gr.Textbox(label="Prompt", placeholder="Describe la imagen que deseas generar...")
        generate_button = gr.Button("Generar Imagen")
    
    with gr.Accordion("Configuración Avanzada", open=False):
        seed = gr.Slider(label="Seed", minimum=0, maximum=MAX_SEED, step=1, value=0)
        random_seed = gr.Checkbox(label="Randomizar Seed", value=True)
        low_thresh = gr.Slider(label="Umbral Bajo Canny", minimum=0, maximum=255, step=1, value=50)
        high_thresh = gr.Slider(label="Umbral Alto Canny", minimum=0, maximum=255, step=1, value=200)
        guidance_scale = gr.Slider(label="Guidance Scale", minimum=1.0, maximum=30.0, step=0.1, value=7.5)
        num_inference_steps = gr.Slider(label="Pasos de Inferencia", minimum=1, maximum=50, step=1, value=20)
    
    generate_button.click(
        generate_image,
        inputs=[image_input, prompt, seed, random_seed, low_thresh, high_thresh, guidance_scale, num_inference_steps],
        outputs=[result_image, seed]
    )

if __name__ == "__main__":
    demo.launch()