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

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

# Instalar dependencias en Hugging Face Spaces
os.system("pip install -U diffusers controlnet_aux mediapipe")

# Máximo valor para seeds y tamaño de imagen
MAX_SEED = 2**32 - 1
MAX_IMAGE_SIZE = 1024

# Inicialización del pipeline Flux Canny Dev
pipe = FluxControlPipeline.from_pretrained(
    "black-forest-labs/FLUX.1-Canny-dev", torch_dtype=torch.float16
).to("cuda" if torch.cuda.is_available() else "cpu")

# Procesador Canny
processor = CannyDetector()

@spaces.GPU
def generate_image(image, prompt, seed, random_seed, low_thresh, high_thresh, guidance_scale, num_inference_steps):
    # Controlar el seed para reproducibilidad
    if random_seed:
        seed = random.randint(0, MAX_SEED)
    generator = torch.manual_seed(seed)
    
    # Aplicar CannyDetector a la imagen de control
    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 el modelo ControlNet
    result = pipe(
        prompt=prompt,
        control_image=control_image,
        height=512,
        width=512,
        guidance_scale=guidance_scale,
        num_inference_steps=num_inference_steps,
        generator=generator
    ).images[0]
    
    return result, 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()