🐛 Fix: ensure DeepFashion2 model is downloaded before loading

app.py

@@ -1,19 +1,4 @@
 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.")
-
-
-
 import traceback
 from datetime import datetime
 import torch, gc
@@ -21,16 +6,18 @@ from PIL import Image
 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
-
-
-
-
-
+from background_edit import run_background_removal_and_inpaint, run_clothing_inpaint
 
+# Ensure DeepFashion2 model is downloaded early
+MODEL_URL = "https://huggingface.co/Bingsu/adetailer/resolve/main/deepfashion2_yolov8s-seg.pt"
+MODEL_PATH = "deepfashion2_yolov8s-seg.pt"
+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.")
 
-# ───────────────────── Helpers ─────────────────────
+# ─────────────── Helpers ───────────────
 def _print_trace():
     traceback.print_exc()
 
@@ -38,10 +25,11 @@ def unload_models():
     torch.cuda.empty_cache()
     gc.collect()
 
-def safe_generate_and_inpaint(
+def safe_generate_all_steps(
     image,
     prompt_1, neg_1, strength_1, guidance_1,
-    prompt_2, neg_2, guidance_2
+    prompt_2, neg_2, guidance_2,
+    prompt_3, neg_3, guidance_3
 ):
     try:
         if image is None:
@@ -66,100 +54,78 @@ def safe_generate_and_inpaint(
         # Step 2: Background Inpainting
         print("[INFO] Step 2: Inpainting background...", flush=True)
         unload_models()
-        result = run_background_removal_and_inpaint(
+        with_bg = run_background_removal_and_inpaint(
             image_path=path,
             prompt=prompt_2,
             negative_prompt=neg_2,
             guidance_scale=guidance_2
         )
 
-        return refined, result, ""
+        # Step 3: Clothing Inpainting
+        print("[INFO] Step 3: Inpainting clothing...", flush=True)
+        final, err = run_clothing_inpaint(
+            with_bg,
+            prompt_3,
+            neg_3,
+            guidance_3
+        )
+        if err:
+            return refined, with_bg, None, err
+        return refined, with_bg, final, ""
 
     except gr.Error as e:
-        return None, None, f"🛑 {str(e)}"
+        return None, None, None, f"🛑 {str(e)}"
     except Exception as e:
         _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)}"
+        return None, None, None, f"❌ Unexpected Error: {type(e).__name__}: {str(e)}"
 
-# ───────────────────── Gradio UI ─────────────────────
+# ─────────────── Gradio UI ───────────────
 with gr.Blocks() as demo:
-    gr.Markdown("## 🧠 Headshot + Background Generator (Full Prompt Control)")
+    gr.Markdown("## 🧠 Full Headshot + Background + Clothing Generator (One Click)")
 
     with gr.Row():
         input_image = gr.Image(type="pil", label="Upload Headshot")
 
     gr.Markdown("### Step 1: Headshot Refinement (LoRA)")
-
     with gr.Row():
         prompt_1 = gr.Textbox(label="Headshot Prompt", value="a professional headshot of a confident woman in her 30s with blonde hair")
         neg_1 = gr.Textbox(label="Headshot Negative Prompt", value="deformed, cartoon, anime, sketch, blurry, low quality")
-
     with gr.Row():
         strength_1 = gr.Slider(0.1, 1.0, value=0.2, step=0.05, label="Refinement Strength")
         guidance_1 = gr.Slider(1, 20, value=17, step=0.5, label="Guidance Scale (Headshot)")
 
     gr.Markdown("### Step 2: Background Inpainting (SDXL)")
-
     with gr.Row():
         prompt_2 = gr.Textbox(label="Background Prompt", value="modern hospital background, clean, soft lighting")
         neg_2 = gr.Textbox(label="Background Negative Prompt", value="fantasy, cartoon, cluttered, sketch")
-
     with gr.Row():
         guidance_2 = gr.Slider(1, 20, value=10, step=0.5, label="Guidance Scale (Background)")
 
-    go_btn = gr.Button("✨ Generate Refined Headshot + Background")
+    gr.Markdown("### Step 3: Clothing Replacement")
+    with gr.Row():
+        prompt_3 = gr.Textbox(label="Clothing Prompt", value="white female CEO professional blazer, clean look")
+        neg_3 = gr.Textbox(label="Clothing Negative Prompt", value="hoodie, casual wear, fantasy, cartoon, jeans, distorted, blurry")
+    with gr.Row():
+        guidance_3 = gr.Slider(1, 20, value=17.0, step=0.5, label="Clothing Guidance Scale")
+
+    go_btn = gr.Button("✨ Run Full Pipeline (All 3 Steps)")
 
     with gr.Row():
         output_refined = gr.Image(type="pil", label="Step 1: Refined Headshot")
-        output_final = gr.Image(type="pil", label="Step 2: Final Image with Background")
+        output_bg = gr.Image(type="pil", label="Step 2: With New Background")
+        output_final = gr.Image(type="pil", label="Step 3: Final with New Clothing")
 
     error_box = gr.Markdown(label="Error", value="", visible=True)
 
     go_btn.click(
-        fn=safe_generate_and_inpaint,
+        fn=safe_generate_all_steps,
         inputs=[
-            input_image, prompt_1, neg_1, strength_1, guidance_1,
-            prompt_2, neg_2, guidance_2
+            input_image,
+            prompt_1, neg_1, strength_1, guidance_1,
+            prompt_2, neg_2, guidance_2,
+            prompt_3, neg_3, guidance_3
         ],
-        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
+        outputs=[output_refined, output_bg, output_final, error_box]
     )
 
 demo.launch(debug=True)