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

# 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 Flux
    result = pipe(
        prompt=prompt,
        control_image=control_image,
        guidance_scale=guidance_scale,
        generator=generator
    ).images[0]
    
    return result, seed

# Interfaz de usuario Gradio
with gr.Blocks() as demo:
    gr.Markdown("# Generador de Imágenes con FLUX.1 Canny-dev 🚀")