Update app.py

app.py

@@ -1,60 +1,61 @@
-import gradio as gr
 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"))
+# Máximo valor para seeds y tamaño de imagen
+MAX_SEED = 2**32 - 1
+MAX_IMAGE_SIZE = 1024
 
-# 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")
+# Inicialización del pipeline Flux Canny Dev
+pipe = FluxControlPipeline.from_pretrained(
+    "black-forest-labs/FLUX.1-Canny-dev", torch_dtype=torch.float32
+).to("cuda" if torch.cuda.is_available() else "cpu")
 
-# Procesador Canny para detección de bordes
+# Procesador CannyDetector
 processor = CannyDetector()
 
-MAX_SEED = 2**32 - 1
-MAX_IMAGE_SIZE = 1024
-
+@spaces.GPU
 def generate_image(image, prompt, seed, random_seed, low_thresh, high_thresh, guidance_scale, num_inference_steps):
-    # Control de Seed
+    # Controlar el seed para reproducibilidad
     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)
+    # 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 Diffusers
-    generated_image = pipe(
+    # Generar imagen con el modelo Flux
+    result = pipe(
         prompt=prompt,
         control_image=control_image,
-        height=MAX_IMAGE_SIZE,
-        width=MAX_IMAGE_SIZE,
+        height=image.height,
+        width=image.width,
         num_inference_steps=num_inference_steps,
         guidance_scale=guidance_scale,
-        generator=generator,
+        generator=generator
     ).images[0]
     
-    return generated_image, seed
+    return result, seed
 
-# Interfaz Gradio
+# Interfaz de usuario Gradio
 with gr.Blocks() as demo:
-    gr.Markdown("# Generador de Imágenes con FLUX.1 Canny 🚀")
+    gr.Markdown("# Generador de Imágenes con FLUX.1 Canny-dev 🚀")
     
     with gr.Row():
-        image_input = gr.Image(type="pil", label="Sube tu imagen de control")
-        result_image = gr.Image(label="Imagen Generada")
+        input_image = gr.Image(type="pil", label="Sube tu imagen de control")
+        output_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)
@@ -62,11 +63,11 @@ with gr.Blocks() as demo:
         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]
+        inputs=[input_image, prompt, seed, random_seed, low_thresh, high_thresh, guidance_scale, num_inference_steps],
+        outputs=[output_image, seed]
     )
 
 if __name__ == "__main__":