Update app.py

app.py

@@ -1,13 +1,13 @@
 import torch
 import spaces
 import gradio as gr
-from diffusers import FluxImg2ImgPipeline
+from diffusers import FluxInpaintPipeline
 from PIL import Image
 
-# 1. Carga del modelo optimizado en memoria
+# 1. Carga del modelo en modo Inpaint
 model_id = "black-forest-labs/FLUX.1-schnell"
-print("Cargando pipeline en GPU...")
-pipe = FluxImg2ImgPipeline.from_pretrained(
+print("Cargando pipeline de Inpainting en GPU...")
+pipe = FluxInpaintPipeline.from_pretrained(
     model_id, 
     torch_dtype=torch.bfloat16
 ).to("cuda")
@@ -16,45 +16,62 @@ pipe = FluxImg2ImgPipeline.from_pretrained(
 print("Inyectando LoRA...")
 pipe.load_lora_weights("joyfox/Kontext-Cosplay-Lora", weight_name="kontext_cosplay_lora_20000.safetensors")
 
-# 3. Función de generación corregida
+# 3. Función de generación con Inpainting
 @spaces.GPU
-def generate_cosplay(image, prompt, strength, steps):
-    if image is None:
+def generate_cosplay_inpaint(input_dict, prompt, strength, steps):
+    # input_dict contiene la imagen original y los trazos del pincel
+    if input_dict is None or input_dict["background"] is None:
         return None
         
-    # Redimensionamos para mantener estabilidad en FLUX
-    image = image.convert("RGB").resize((1024, 1024))
+    # 3.1 Procesar imagen original
+    image = input_dict["background"].convert("RGB").resize((1024, 1024))
     
-    # Trigger indispensable del LoRA
-    full_prompt = f"{prompt}, r2cos"
+    # 3.2 Construir la máscara desde los trazos del usuario
+    layers = input_dict.get("layers", [])
+    if not layers:
+        raise gr.Error("¡Debes usar el pincel para pintar sobre la ropa!")
+        
+    # Creamos una máscara negra vacía del mismo tamaño
+    mask = Image.new("L", image.size, 0)
     
-    # Eliminamos la semilla fija estricta para permitir variedad por intento,
-    # pero puedes volver a usar un generador si deseas replicabilidad exacta.
+    # Sumamos todos los trazos (capas) que el usuario dibujó
+    for layer in layers:
+        layer_resized = layer.convert("RGBA").resize((1024, 1024))
+        alpha_channel = layer_resized.split()[3] # Extraemos por dónde pasó el pincel
+        mask.paste(alpha_channel, (0, 0), alpha_channel)
+        
+    # 3.3 Configurar el prompt
+    full_prompt = f"{prompt}, r2cos"
     
-    # Ejecución con los nuevos parámetros dinámicos
+    # 3.4 Ejecutar el modelo Inpaint
     result = pipe(
         prompt=full_prompt,
         image=image,
-        strength=strength,          # Control dinámico desde la UI
-        num_inference_steps=steps,  # Elevado para dar tiempo al LoRA
+        mask_image=mask, # Aquí le pasamos la máscara a la IA
+        strength=strength,          
+        num_inference_steps=steps,  
         guidance_scale=0.0
     ).images[0]
     
     return result
 
-# 4. Interfaz de Gradio interactiva
+# 4. Interfaz interactiva con editor de imagen
 demo = gr.Interface(
-    fn=generate_cosplay, 
+    fn=generate_cosplay_inpaint, 
     inputs=[
-        gr.Image(type="pil", label="Foto de Referencia"), 
-        gr.Textbox(label="Descripción del personaje (Ej: Naruto, Cyberpunk, Samurai)"),
-        gr.Slider(minimum=0.1, maximum=1.0, value=0.7, step=0.05, label="Fuerza de Transformación (Strength)"),
-        gr.Slider(minimum=10, maximum=30, value=16, step=1, label="Pasos de Inferencia Totales (Steps)")
+        gr.ImageEditor(
+            type="pil", 
+            label="1. Sube tu foto y PINTA LA ROPA (deja la cara intacta)",
+            brush=gr.Brush(colors=["#FF0000"], color_mode="fixed") # Pincel rojo para ver bien
+        ), 
+        gr.Textbox(label="2. Descripción (Ej: detailed samurai armor, cyberpunk)"),
+        gr.Slider(minimum=0.1, maximum=1.0, value=0.85, step=0.05, label="Fuerza (Strength)"),
+        gr.Slider(minimum=10, maximum=30, value=20, step=1, label="Pasos (Steps)")
     ], 
     outputs=gr.Image(label="Resultado Cosplay"),
-    title="Kontext-Cosplay AI Pro",
-    description="Ajusta los deslizadores si la imagen cambia muy poco."
+    title="Kontext-Cosplay AI Inpaint",
+    description="Pinta de rojo la ropa y el fondo. La IA transformará solo la zona pintada, respetando la cara original."
 )
 
 if __name__ == "__main__":
-    demo.launch() 
+    demo.launch()