🔧 Add clothing inpainting step (Step 3) with independent prompt and guidance

app.py

@@ -7,6 +7,22 @@ import gradio as gr
 
 from inference import generate_with_lora
 from background_edit import run_background_removal_and_inpaint
+from background_edit import run_clothing_inpaint
+
+import os
+
+MODEL_URL = "https://huggingface.co/Bingsu/adetailer/resolve/main/deepfashion2_yolov8s-seg.pt"
+MODEL_PATH = "deepfashion2_yolov8s-seg.pt"
+
+# ─── Download DeepFashion2 model if not already present ───
+if not os.path.exists(MODEL_PATH):
+    import urllib.request
+    print("[INFO] Downloading DeepFashion2 YOLOv8 model...")
+    urllib.request.urlretrieve(MODEL_URL, MODEL_PATH)
+    print("[INFO] Model downloaded.")
+else:
+    print("[INFO] DeepFashion2 model already exists.")
+
 
 # ───────────────────── Helpers ─────────────────────
 def _print_trace():
@@ -59,6 +75,15 @@ def safe_generate_and_inpaint(
         _print_trace()
         return None, None, f"❌ Unexpected Error: {type(e).__name__}: {str(e)}"
 
+def guarded_clothing(image, prompt, neg, guidance):
+    try:
+        result, err = run_clothing_inpaint(image, prompt, neg, guidance)
+        return result, err
+    except Exception as e:
+        import traceback
+        traceback.print_exc()
+        return None, f"❌ Unexpected Error: {type(e).__name__}: {str(e)}"
+
 # ───────────────────── Gradio UI ─────────────────────
 with gr.Blocks() as demo:
     gr.Markdown("## 🧠 Headshot + Background Generator (Full Prompt Control)")
@@ -102,4 +127,33 @@ with gr.Blocks() as demo:
         outputs=[output_refined, output_final, error_box]
     )
 
+    gr.Markdown("### 👗 Step 3: Clothing Replacement")
+
+    with gr.Row():
+        clothing_prompt = gr.Textbox(
+            label="Clothing Prompt",
+            value="white female CEO professional blazer, clean look"
+        )
+        clothing_negative = gr.Textbox(
+            label="Clothing Negative Prompt",
+            value="hoodie, casual wear, fantasy, cartoon, jeans, distorted, blurry"
+        )
+
+    with gr.Row():
+        clothing_guidance = gr.Slider(1, 20, value=17.0, step=0.5, label="Clothing Guidance Scale")
+
+    with gr.Row():
+        clothing_btn = gr.Button("🧵 Inpaint Clothing")
+        clothing_output = gr.Image(type="pil", label="Step 3: Final Image with New Clothing")
+
+    clothing_error = gr.Markdown(label="Clothing Error", value="", visible=True)
+
+
+    clothing_btn.click(
+        fn=guarded_clothing,
+        inputs=[output_final, clothing_prompt, clothing_negative, clothing_guidance],
+        outputs=[clothing_output, clothing_error],
+        preprocess=False
+    )
+
 demo.launch(debug=True)

background_edit.py

@@ -21,11 +21,15 @@ from diffusers import StableDiffusionXLInpaintPipeline
 from utils import pil_to_cv2, cv2_to_pil
 import gradio as gr  # ✅ Needed for error handling
 
+INPAINT_SIZE = 1024
 
+# Load clothing model
+clothing_model = YOLO("deepfashion2_yolov8s-seg.pt")
 
 # ✅ Load models once
 yolo = YOLO("yolov8x-seg.pt")
 
+
 inpaint_pipe = StableDiffusionXLInpaintPipeline.from_pretrained(
     "diffusers/stable-diffusion-xl-1.0-inpainting-0.1",
     torch_dtype=torch.float16,
@@ -51,8 +55,8 @@ def run_background_removal_and_inpaint(image_path, prompt, negative_prompt, guid
     dilated = cv2.dilate(background_mask, kernel, iterations=1)
     inpaint_mask = (dilated * 255).astype(np.uint8)
 
-    mask_pil = cv2_to_pil(inpaint_mask).resize((1024, 1024)).convert("L")
-    img_pil = image.resize((1024, 1024)).convert("RGB")
+    mask_pil = cv2_to_pil(inpaint_mask).resize((INPAINT_SIZE, INPAINT_SIZE)).convert("L")
+    img_pil = image.resize((INPAINT_SIZE, INPAINT_SIZE)).convert("RGB")
 
     result = inpaint_pipe(
         prompt=prompt,
@@ -64,3 +68,42 @@ def run_background_removal_and_inpaint(image_path, prompt, negative_prompt, guid
     ).images[0]
 
     return result
+
+def run_clothing_inpaint(image, prompt, negative_prompt, guidance):
+    try:
+        print("[INFO] Step 3: Clothing segmentation and inpainting...", flush=True)
+
+        img_cv = np.array(image.convert("RGB"))[..., ::-1]  # PIL → OpenCV BGR
+        h, w = img_cv.shape[:2]
+
+        # Segment clothing
+        results = clothing_model(img_cv)
+        masks = results[0].masks.data.cpu().numpy()
+        if len(masks) == 0:
+            raise gr.Error("No clothing detected. Try a different image.")
+
+        mask = masks[0]
+        resized_mask = cv2.resize(mask, (w, h), interpolation=cv2.INTER_NEAREST)
+        binary_mask = (resized_mask > 0.5).astype(np.uint8) * 255
+        mask_pil = Image.fromarray(binary_mask).convert("L").resize((INPAINT_SIZE, INPAINT_SIZE))
+
+        # Resize input image
+        resized_image = image.convert("RGB").resize((INPAINT_SIZE, INPAINT_SIZE))
+
+        # Inpaint clothing
+        result = inpaint_pipe(
+            prompt=prompt,
+            negative_prompt=negative_prompt,
+            image=resized_image,
+            mask_image=mask_pil,
+            guidance_scale=guidance,
+            num_inference_steps=50
+        ).images[0]
+
+        return result, ""
+    
+    except gr.Error as e:
+        return None, f"🛑 {str(e)}"
+    except Exception as e:
+        traceback.print_exc()
+        return None, f"❌ Unexpected Error: {type(e).__name__}: {str(e)}"