Fix step error handling and wire up error_box

app.py

@@ -27,7 +27,6 @@ def safe_run_background(*args, **kwargs):
         return run_background_removal_and_inpaint(*args, **kwargs)
     except Exception as e:
         _print_trace()
-        # We still raise gr.Error so the wrapper will catch & log
         raise gr.Error(f"[Step 2] Background replacement failed: {type(e).__name__}: {e}")
 
 # ───────────────────── UI ─────────────────────
@@ -62,7 +61,6 @@ with gr.Blocks() as demo:
         def _save_to_state(img):
             return {"step1": img} if img is not None else gr.skip()
 
-        # Build the click-chain and store it in `event`
         event = (
             run_btn.click(
                 fn=safe_generate_with_lora,
@@ -75,8 +73,8 @@ with gr.Blocks() as demo:
 
     # ─────────── STEP 2: Background Replacement ───────────
     with gr.Tab("Step 2: Replace Background"):
-        # This textbox will show any error messages from Step 2
-        error_box = gr.Textbox(label="Error", interactive=False, lines=2)
+        # Show formatted error messages
+        error_box = gr.Markdown(value="", visible=True)
 
         with gr.Row():
             inpaint_prompt = gr.Textbox(
@@ -95,26 +93,27 @@ with gr.Blocks() as demo:
             inpaint_btn = gr.Button("Remove Background & Inpaint", interactive=False)
 
         def guarded_inpaint(img, prompt_bg, neg_bg):
-            # If no image is in state, return an error immediately
             if img is None:
-                return None, "Error: No headshot found—please run Step 1 first."
+                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, ""  # success: show result, clear error
+                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:
-                # Log to console/Space logs
-                print(f"[Step 2 ERROR] {type(e).__name__}: {e}", flush=True)
-                return None, f"{type(e).__name__}: {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],  # wire both outputs!
+            outputs=[inpaint_result, error_box],
         )
 
-    # Enable the Step 2 button once Step 1’s chain finishes
+    # Enable Step 2 after Step 1 completes
     event.then(lambda: gr.update(interactive=True), None, inpaint_btn)
 
-    # Launch in debug mode so exceptions & prints go to your terminal/logs
     demo.launch(debug=True)