Fix step error handling and wire up error_box

app.py

@@ -1,119 +1,103 @@
+import os
+
+# ── Set cache/config dirs ──
+hf_home = "/data/.cache/huggingface"
+yolo_cfg = "/data/ultralytics"
+os.makedirs(hf_home, exist_ok=True)
+os.makedirs(yolo_cfg, exist_ok=True)
+os.environ["HF_HOME"] = hf_home
+os.environ["YOLO_CONFIG_DIR"] = yolo_cfg
+
+from ultralytics import YOLO
+import numpy as np
+import torch
+from PIL import Image
+import cv2
+from diffusers import StableDiffusionXLInpaintPipeline
 import gradio as gr
-from inference import generate_with_lora
-from background_edit import run_background_removal_and_inpaint
-import traceback, torch, gc
-
-# ───────────────────── Helpers ─────────────────────
-def _print_trace():
-    traceback.print_exc()
-
-def safe_generate_with_lora(*a, **kw):
-    try:
-        return generate_with_lora(*a, **kw)
-    except gr.Error:
-        _print_trace()
-        raise
-    except Exception as e:
-        _print_trace()
-        raise gr.Error(f"Image generation failed: {e}")
-
-def unload_models():
-    torch.cuda.empty_cache()
-    gc.collect()
-
-def safe_run_background(*args, **kwargs):
-    try:
-        unload_models()  # free VRAM before loading the inpainting model
-        return run_background_removal_and_inpaint(*args, **kwargs)
-    except Exception as e:
-        _print_trace()
-        raise gr.Error(f"[Step 2] Background replacement failed: {type(e).__name__}: {e}")
-
-# ───────────────────── UI ─────────────────────
-shared_output  = gr.State()  # holds the Step 1 output image
-original_input = gr.State()  # holds the original upload (optional)
 
+# ---- utils ----
+def pil_to_cv2(pil_img):
+    return cv2.cvtColor(np.array(pil_img), cv2.COLOR_RGB2BGR)
+
+def cv2_to_pil(cv_img):
+    return Image.fromarray(cv2.cvtColor(cv_img, cv2.COLOR_BGR2RGB))
+
+# ---- load models ----
+yolo = YOLO("yolov8x-seg.pt")
+inpaint_pipe = StableDiffusionXLInpaintPipeline.from_pretrained(
+    "diffusers/stable-diffusion-xl-1.0-inpainting-0.1",
+    torch_dtype=torch.float16,
+    use_safetensors=True,
+    use_auth_token=os.getenv("HF_TOKEN")
+).to("cuda")
+
+# ---- processing logic ----
+def run_background_removal_and_inpaint(image, prompt, negative_prompt):
+    if image is None:
+        raise gr.Error("Please upload an image.")
+
+    img_cv = pil_to_cv2(image)
+    results = yolo(img_cv)
+
+    if not results or not results[0].masks or len(results[0].masks.data) == 0:
+        raise gr.Error("No subject detected in the image. Please upload a clearer photo.")
+
+    mask = results[0].masks.data[0].cpu().numpy()
+    binary = (mask > 0.5).astype(np.uint8)
+    background_mask = 1 - binary
+    kernel = np.ones((15, 15), np.uint8)
+    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")
+
+    result = inpaint_pipe(
+        prompt=prompt,
+        negative_prompt=negative_prompt or "",
+        image=img_pil,
+        mask_image=mask_pil,
+        guidance_scale=10,
+        num_inference_steps=40
+    ).images[0]
+
+    return result
+
+# ---- Gradio interface ----
 with gr.Blocks() as demo:
-    demo.queue()  # enable batching / concurrency
-
-    # ─────────── STEP 1: Headshot Refinement ───────────
-    with gr.Tab("Step 1: Headshot Refinement"):
-        with gr.Row():
-            input_image  = gr.Image(type="pil", label="Upload Headshot")
-            output_image = gr.Image(type="pil", label="Refined Output")
-
-        with gr.Row():
-            prompt = gr.Textbox(
-                label="Prompt",
-                value="a professional corporate headshot of a confident woman in her 30s with blonde hair"
-            )
-            negative_prompt = gr.Textbox(
-                label="Negative Prompt",
-                value="deformed, cartoon, anime, illustration, painting, drawing, sketch, low resolution, blurry, out of focus, pixelated"
-            )
-
-        with gr.Row():
-            strength = gr.Slider(0.1, 1.0, value=0.20, step=0.05, label="Strength")
-            guidance = gr.Slider(1, 20, value=17.0, step=0.5, label="Guidance Scale")
-
-        run_btn = gr.Button("Generate")
-
-        def _save_to_state(img):
-            return {"step1": img} if img is not None else gr.skip()
-
-        event = (
-            run_btn.click(
-                fn=safe_generate_with_lora,
-                inputs=[input_image, prompt, negative_prompt, strength, guidance],
-                outputs=output_image,
-            )
-            .then(_save_to_state, output_image, shared_output)
-            .then(lambda x: x, input_image, original_input)
-        )
+    gr.Markdown("## 🖼️ Remove & Replace Background")
+    gr.Markdown("Upload a headshot, and describe the desired new background.")
+
+    with gr.Row():
+        input_img = gr.Image(type="pil", label="Upload Image")
+        output_img = gr.Image(type="pil", label="Result")
 
-    # ─────────── STEP 2: Background Replacement ───────────
-    with gr.Tab("Step 2: Replace Background"):
-        # Show formatted error messages
-        error_box = gr.Markdown(value="", visible=True)
-
-        with gr.Row():
-            inpaint_prompt = gr.Textbox(
-                label="New Background Prompt",
-                value="modern open-plan startup office background, natural lighting, glass walls, clean design, minimalistic decor"
-            )
-            inpaint_negative = gr.Textbox(
-                label="Negative Prompt",
-                value="dark lighting, cluttered background, fantasy elements, cartoon, anime, painting, low quality, distorted shapes"
-            )
-
-        with gr.Row():
-            inpaint_result = gr.Image(type="pil", label="Inpainted Image")
-
-        with gr.Row():
-            inpaint_btn = gr.Button("Remove Background & Inpaint", interactive=False)
-
-        def guarded_inpaint(img, prompt_bg, neg_bg):
-            if img is None:
-                return None, "**🛑 Error:** No headshot found — please run Step 1 first."
-
-            try:
-                print("[DEBUG] Starting background removal and inpainting…", flush=True)
-                result = safe_run_background(img, prompt_bg, neg_bg)
-                return result, ""  # Clear error on success
-            except gr.Error as e:
-                print(f"[Step 2 gr.Error] {e}", flush=True)
-                return None, f"**🛑 Step 2 Failed:** {str(e)}"
-            except Exception as e:
-                print(f"[Step 2 UNEXPECTED ERROR] {type(e).__name__}: {e}", flush=True)
-                return None, f"**❌ Unexpected Error:** {type(e).__name__}: {e}"
-
-        inpaint_btn.click(
-            fn=guarded_inpaint,
-            inputs=[shared_output, inpaint_prompt, inpaint_negative],
-            outputs=[inpaint_result, error_box],
+    with gr.Row():
+        prompt = gr.Textbox(
+            label="New Background Prompt",
+            value="modern open-plan office, soft natural light, minimalistic decor"
+        )
+        neg_prompt = gr.Textbox(
+            label="Negative Prompt",
+            value="cartoon, fantasy, dark lighting, painting, anime"
         )
 
-    # Enable Step 2 after Step 1 completes
-    event.then(lambda: gr.update(interactive=True), None, inpaint_btn)
+    error_box = gr.Markdown()
+
+    def safe_run(img, prompt, neg_prompt):
+        try:
+            result = run_background_removal_and_inpaint(img, prompt, neg_prompt)
+            return result, ""
+        except Exception as e:
+            print(f"[ERROR] {type(e).__name__}: {e}")
+            return None, f"**❌ Error:** {type(e).__name__}: {e}"
+
+    run_btn = gr.Button("Run Background Inpaint")
+    run_btn.click(
+        fn=safe_run,
+        inputs=[input_img, prompt, neg_prompt],
+        outputs=[output_img, error_box]
+    )
 
-    demo.launch(debug=True)
+demo.launch(debug=True)

app_new.py

@@ -0,0 +1,119 @@
+import gradio as gr
+from inference import generate_with_lora
+from background_edit import run_background_removal_and_inpaint
+import traceback, torch, gc
+
+# ───────────────────── Helpers ─────────────────────
+def _print_trace():
+    traceback.print_exc()
+
+def safe_generate_with_lora(*a, **kw):
+    try:
+        return generate_with_lora(*a, **kw)
+    except gr.Error:
+        _print_trace()
+        raise
+    except Exception as e:
+        _print_trace()
+        raise gr.Error(f"Image generation failed: {e}")
+
+def unload_models():
+    torch.cuda.empty_cache()
+    gc.collect()
+
+def safe_run_background(*args, **kwargs):
+    try:
+        unload_models()  # free VRAM before loading the inpainting model
+        return run_background_removal_and_inpaint(*args, **kwargs)
+    except Exception as e:
+        _print_trace()
+        raise gr.Error(f"[Step 2] Background replacement failed: {type(e).__name__}: {e}")
+
+# ───────────────────── UI ─────────────────────
+shared_output  = gr.State()  # holds the Step 1 output image
+original_input = gr.State()  # holds the original upload (optional)
+
+with gr.Blocks() as demo:
+    demo.queue()  # enable batching / concurrency
+
+    # ─────────── STEP 1: Headshot Refinement ───────────
+    with gr.Tab("Step 1: Headshot Refinement"):
+        with gr.Row():
+            input_image  = gr.Image(type="pil", label="Upload Headshot")
+            output_image = gr.Image(type="pil", label="Refined Output")
+
+        with gr.Row():
+            prompt = gr.Textbox(
+                label="Prompt",
+                value="a professional corporate headshot of a confident woman in her 30s with blonde hair"
+            )
+            negative_prompt = gr.Textbox(
+                label="Negative Prompt",
+                value="deformed, cartoon, anime, illustration, painting, drawing, sketch, low resolution, blurry, out of focus, pixelated"
+            )
+
+        with gr.Row():
+            strength = gr.Slider(0.1, 1.0, value=0.20, step=0.05, label="Strength")
+            guidance = gr.Slider(1, 20, value=17.0, step=0.5, label="Guidance Scale")
+
+        run_btn = gr.Button("Generate")
+
+        def _save_to_state(img):
+            return {"step1": img} if img is not None else gr.skip()
+
+        event = (
+            run_btn.click(
+                fn=safe_generate_with_lora,
+                inputs=[input_image, prompt, negative_prompt, strength, guidance],
+                outputs=output_image,
+            )
+            .then(_save_to_state, output_image, shared_output)
+            .then(lambda x: x, input_image, original_input)
+        )
+
+    # ─────────── STEP 2: Background Replacement ───────────
+    with gr.Tab("Step 2: Replace Background"):
+        # Show formatted error messages
+        error_box = gr.Markdown(value="", visible=True)
+
+        with gr.Row():
+            inpaint_prompt = gr.Textbox(
+                label="New Background Prompt",
+                value="modern open-plan startup office background, natural lighting, glass walls, clean design, minimalistic decor"
+            )
+            inpaint_negative = gr.Textbox(
+                label="Negative Prompt",
+                value="dark lighting, cluttered background, fantasy elements, cartoon, anime, painting, low quality, distorted shapes"
+            )
+
+        with gr.Row():
+            inpaint_result = gr.Image(type="pil", label="Inpainted Image")
+
+        with gr.Row():
+            inpaint_btn = gr.Button("Remove Background & Inpaint", interactive=False)
+
+        def guarded_inpaint(img, prompt_bg, neg_bg):
+            if img is None:
+                return None, "**🛑 Error:** No headshot found — please run Step 1 first."
+
+            try:
+                print("[DEBUG] Starting background removal and inpainting…", flush=True)
+                result = safe_run_background(img, prompt_bg, neg_bg)
+                return result, ""  # Clear error on success
+            except gr.Error as e:
+                print(f"[Step 2 gr.Error] {e}", flush=True)
+                return None, f"**🛑 Step 2 Failed:** {str(e)}"
+            except Exception as e:
+                print(f"[Step 2 UNEXPECTED ERROR] {type(e).__name__}: {e}", flush=True)
+                return None, f"**❌ Unexpected Error:** {type(e).__name__}: {e}"
+
+        inpaint_btn.click(
+            fn=guarded_inpaint,
+            inputs=[shared_output, inpaint_prompt, inpaint_negative],
+            outputs=[inpaint_result, error_box],
+        )
+
+    # Enable Step 2 after Step 1 completes
+    event.then(lambda: gr.update(interactive=True), None, inpaint_btn)
+
+    demo.launch(debug=True)