Fix step 2 error handling and wire up error_box

app.py

@@ -1,37 +1,43 @@
-
 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 _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
+    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}")
+        _print_trace()
+        raise gr.Error(f"Image generation failed: {e}")
 
-def unload_models(): torch.cuda.empty_cache(); gc.collect()
+def unload_models():
+    torch.cuda.empty_cache()
+    gc.collect()
 
 def safe_run_background(*args, **kwargs):
     try:
-        unload_models()  # free VRAM before running background edit
+        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}")
-
+        # 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 ─────────────────────
-shared_output  = gr.State()
-original_input = gr.State()
+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()
+    demo.queue()  # enable batching / concurrency
 
-    # ─────────── STEP 1: Headshot Refinement ───────────
+    # ─────────── STEP 1: Headshot Refinement ───────────
     with gr.Tab("Step 1: Headshot Refinement"):
         with gr.Row():
             input_image  = gr.Image(type="pil", label="Upload Headshot")
@@ -56,7 +62,7 @@ with gr.Blocks() as demo:
         def _save_to_state(img):
             return {"step1": img} if img is not None else gr.skip()
 
-        # Build the *single* click chain and keep the handle in `event`
+        # Build the click-chain and store it in `event`
         event = (
             run_btn.click(
                 fn=safe_generate_with_lora,
@@ -67,14 +73,15 @@ with gr.Blocks() as demo:
             .then(lambda x: x, input_image, original_input)
         )
 
-    # ─────────── STEP 2: Background Replacement ───────────
+    # ─────────── STEP 2: Background Replacement ───────────
     with gr.Tab("Step 2: Replace Background"):
-        # … your prompts and inpaint_result …
+        # This textbox will show any error messages from Step 2
         error_box = gr.Textbox(label="Error", interactive=False, lines=2)
+
         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"
+                value="modern open-plan startup office background, natural lighting, glass walls, clean design, minimalistic decor"
             )
             inpaint_negative = gr.Textbox(
                 label="Negative Prompt",
@@ -88,23 +95,26 @@ 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."
             try:
-                if img is None:
-                    raise ValueError("No image to inpaint—run Step 1 first.")
-                out = safe_run_background(img, prompt_bg, neg_bg)
-                return out, ""                            # success
+                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
             except Exception as e:
-                # logs to console for debugging
+                # Log to console/Space logs
                 print(f"[Step 2 ERROR] {type(e).__name__}: {e}", flush=True)
-                return None, f"{type(e).__name__}: {e}"   # image=None, show error
+                return None, f"{type(e).__name__}: {e}"
 
         inpaint_btn.click(
             fn=guarded_inpaint,
             inputs=[shared_output, inpaint_prompt, inpaint_negative],
-            outputs=[inpaint_result, error_box],  
+            outputs=[inpaint_result, error_box],  # wire both outputs!
         )
 
-    # Enable the Step 2 button once Step 1 finishes
+    # Enable the Step 2 button once Step 1’s chain finishes
     event.then(lambda: gr.update(interactive=True), None, inpaint_btn)
 
-demo.launch(debug=True)
+    # Launch in debug mode so exceptions & prints go to your terminal/logs
+    demo.launch(debug=True)