Update app.py

app.py

@@ -84,72 +84,27 @@ def inpaint(
     import torch
     import torch.nn.functional as F
     import numpy as np
+    from PIL import Image
 
     image = image.convert("RGB")
     mask = mask.convert("L")
     width, height = calculate_optimal_dimensions(image)
 
-    pipe.to("cuda")
-
-    callback_on_step_end = None
-    callback_on_step_end_tensor_inputs = None
-
-    # Si se activa la opción, preparar latentes originales para preservar zonas sin máscara
-    if preserve_unmasked:
-        np_img = np.array(image).astype(np.float32) / 255.0
-        img_t = torch.from_numpy(np_img).permute(2, 0, 1).unsqueeze(0).to(pipe.device)
-        img_t = F.interpolate(img_t, size=(height, width), mode='bilinear', align_corners=False)
-        img_t = (img_t * 2 - 1).to(dtype=pipe.vae.dtype)
-    
-        np_mask = np.array(mask).astype(np.float32) / 255.0
-        mask_t = torch.from_numpy(np_mask).unsqueeze(0).unsqueeze(0).to(pipe.device)
-        mask_t = F.interpolate(mask_t, size=(height, width), mode='nearest')
+    # Guardar imagen original redimensionada para el blending final
+    original_resized = image.resize((width, height), Image.Resampling.LANCZOS)
     
-        with torch.no_grad():
-            latents_orig = pipe.vae.encode(img_t).latent_dist.sample()
-            scaling = getattr(pipe.vae.config, "scaling_factor", getattr(pipe, "vae_scale_factor", 0.13025))
-            latents_orig = latents_orig * scaling
-    
-        # Ajustar máscara al tamaño de los latentes
-        latent_height = latents_orig.shape[2]
-        latent_width = latents_orig.shape[3]
-        mask_t = F.interpolate(mask_t, size=(latent_height, latent_width), mode="nearest")
-    
-        def callback_on_step_end(pipe_self, i, t, callback_kwargs):
-            latents = callback_kwargs.get("latents", None)
-            if latents is not None:
-                # Verificar que tengamos 4 dimensiones [batch, channels, height, width]
-                if latents.dim() != 4:
-                    print(f"⚠️ Warning: latents has {latents.dim()} dimensions, expected 4")
-                    return callback_kwargs
-                
-                # Ajustar dinámicamente los tamaños al del tensor actual
-                current_height = latents.shape[2]
-                current_width = latents.shape[3]
-                
-                if mask_t.shape[-2:] != (current_height, current_width):
-                    resized_mask = F.interpolate(mask_t, size=(current_height, current_width), mode="nearest")
-                else:
-                    resized_mask = mask_t
-        
-                if latents_orig.shape[-2:] != (current_height, current_width):
-                    resized_latents_orig = F.interpolate(latents_orig, size=(current_height, current_width), mode="nearest")
-                else:
-                    resized_latents_orig = latents_orig
-        
-                # Mezclar solo en las áreas no enmascaradas
-                latents = latents * resized_mask + resized_latents_orig * (1 - resized_mask)
-                callback_kwargs["latents"] = latents
-        
-            return callback_kwargs
+    # Crear máscara suavizada para mejor blending
+    mask_resized = mask.resize((width, height), Image.Resampling.LANCZOS)
+    mask_array = np.array(mask_resized).astype(np.float32) / 255.0
     
-        callback_on_step_end_tensor_inputs = ["latents"]
-    
-    else:
-        callback_on_step_end = None
-        callback_on_step_end_tensor_inputs = None
+    # Aplicar un pequeño blur a la máscara para transiciones más suaves
+    from scipy.ndimage import gaussian_filter
+    mask_blurred = gaussian_filter(mask_array, sigma=2.0)
+    mask_blurred = np.clip(mask_blurred, 0, 1)
 
-    # Ejecutar pipeline
+    pipe.to("cuda")
+
+    # Ejecutar pipeline SIN callback (dejamos que el modelo haga su trabajo)
     result = pipe(
         image=image,
         mask_image=mask,
@@ -160,10 +115,23 @@ def inpaint(
         guidance_scale=guidance_scale,
         strength=strength,
         generator=torch.Generator(device="cuda").manual_seed(seed),
-        callback_on_step_end=callback_on_step_end,
-        callback_on_step_end_tensor_inputs=callback_on_step_end_tensor_inputs,
     ).images[0]
 
+    # Si preserve_unmasked está activado, hacer blending en espacio de píxeles
+    if preserve_unmasked:
+        # Convertir a arrays numpy
+        result_array = np.array(result).astype(np.float32)
+        original_array = np.array(original_resized).astype(np.float32)
+        
+        # Expandir máscara a 3 canales (RGB)
+        mask_3channel = np.stack([mask_blurred] * 3, axis=-1)
+        
+        # Blending: donde mask=1 (blanco) usamos result, donde mask=0 (negro) usamos original
+        blended = result_array * mask_3channel + original_array * (1 - mask_3channel)
+        
+        # Convertir de vuelta a imagen
+        result = Image.fromarray(blended.astype(np.uint8), mode='RGB')
+
     return result.convert("RGBA"), prompt, seed
 
 def inpaint_api(