Update app.py

app.py

@@ -7,7 +7,7 @@ from ultralytics import YOLO
 from huggingface_hub import hf_hub_download
 import torch 
 import os
-from PIL import Image
+from PIL import Image, ImageFilter, ImageOps
 
 device = "cuda" if torch.cuda.is_available() else "cpu"
 model_path = "https://huggingface.co/DreamingOracle/Quagmaform_alpha-1/blob/main/DPS_Quagmaform_Alpha1.safetensors" # Your custom model file URL
@@ -51,24 +51,66 @@ def infer(
         prompt=prompt, negative_prompt=negative_prompt, guidance_scale=guidance_scale, num_inference_steps=num_inference_steps, width=width, height=height, generator=generator, 
     ).images[0] 
     
-    # ADetailer post-processing for face enhancement
-    results = adetailer_model(image)
-    if results[0].boxes:  # If faces detected
-        for box in results[0].boxes:
-            x1, y1, x2, y2 = map(int, box.xyxy[0])
-            face = image.crop((x1, y1, x2, y2))
-            # Inpaint face area (simple example; adjust prompt/strength as needed)
-            inpainted_face = pipe(
-                prompt=prompt + ", high detail face",
-                image=face,
-                mask_image=Image.new("L", face.size, 255),  # Full mask for inpaint
-                strength=0.75,
-                num_inference_steps=20,
-                guidance_scale=7.5,
-                generator=generator
-            ).images[0]
-            image.paste(inpainted_face, (x1, y1))
-    
+    # ---------------------------
+    # ADetailer post-processing for face enhancement (with padding + soft blend)
+    # ---------------------------
+    try:
+        results = adetailer_model(image)
+        # ultralytics YOLO results typically have results[0].boxes
+        if results and len(results) and getattr(results[0], "boxes", None):
+            for box in results[0].boxes:
+                # extract bbox (xyxy)
+                x1, y1, x2, y2 = map(int, box.xyxy[0])
+
+                # 1) Pad bbox so inpaint has surrounding context
+                w = max(1, x2 - x1)
+                h = max(1, y2 - y1)
+                pad = int(max(10, 0.18 * max(w, h)))  # ~18% padding, tweak if desired
+                x1p = max(0, x1 - pad)
+                y1p = max(0, y1 - pad)
+                x2p = min(image.width, x2 + pad)
+                y2p = min(image.height, y2 + pad)
+
+                face = image.crop((x1p, y1p, x2p, y2p))
+
+                # 2) Make dimensions divisible by 8 (many pipelines prefer multiples of 8)
+                fw, fh = face.size
+                fw8 = max(8, (fw // 8) * 8)
+                fh8 = max(8, (fh // 8) * 8)
+                if (fw8, fh8) != (fw, fh):
+                    face = face.resize((fw8, fh8), Image.LANCZOS)
+
+                # 3) Create a full white mask for inpainting (255 = area to replace)
+                mask = Image.new("L", face.size, 255)
+
+                # 4) Create a blurred mask for soft blending when pasting back
+                blur_radius = max(4, int(min(face.size) / 10))  # heuristic; adjust if needed
+                paste_mask = mask.filter(ImageFilter.GaussianBlur(radius=blur_radius))
+
+                # 5) Inpaint on the crop — use gentler strength so output doesn't look pasted
+                #    (this re-uses your existing `pipe` call which previously worked for you)
+                inpaint_result = pipe(
+                    prompt=prompt + ", high detail face",
+                    image=face,
+                    mask_image=mask,
+                    strength=0.45,                # reduced from 0.75 to lower mismatch
+                    num_inference_steps=20,
+                    guidance_scale=7.5,
+                    generator=generator
+                ).images[0]
+
+                # 6) Optionally attempt a simple color-match (blend) if lighting differs.
+                #    Here we keep it simple and rely on the blurred mask to smooth edges.
+                # Ensure paste_mask is 'L'
+                if paste_mask.mode != "L":
+                    paste_mask = paste_mask.convert("L")
+
+                # 7) Paste the inpainted crop back using the soft mask for feathered edges
+                image.paste(inpaint_result, (x1p, y1p), paste_mask)
+    except Exception as e:
+        # Log exception to console for debugging but don't crash the whole run
+        print("ADetailer post-process failed:", e)
+
     # Save in selected format for download consistency
     output_path = f"generated_image.{save_format}"
     image.save(output_path, format=save_format.upper())
@@ -131,4 +173,4 @@ with gr.Blocks(css=css) as demo:
             outputs=[result, seed], 
         ) 
 if __name__ == "__main__": 
-    demo.launch()
+    demo.launch()