Fix step 2 error handling and wire up error_box

app.py

@@ -1,3 +1,4 @@
+
 import gradio as gr
 from inference import generate_with_lora
 from background_edit import run_background_removal_and_inpaint
@@ -62,12 +63,14 @@ with gr.Blocks() as demo:
                 inputs=[input_image, prompt, negative_prompt, strength, guidance],
                 outputs=output_image,
             )
-            .then(_save_to_state, None, shared_output)
+            .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"):
+        # … your prompts and inpaint_result …
+        error_box = gr.Textbox(label="Error", interactive=False, lines=2)
         with gr.Row():
             inpaint_prompt = gr.Textbox(
                 label="New Background Prompt",
@@ -85,14 +88,20 @@ with gr.Blocks() as demo:
             inpaint_btn = gr.Button("Remove Background & Inpaint", interactive=False)
 
         def guarded_inpaint(img, prompt_bg, neg_bg):
-            if img is None:
-                raise gr.Error("[Step 2] Error: No image to inpaint. Please run Step 1 first.")
-            return safe_run_background(img, prompt_bg, neg_bg)
+            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
+            except Exception as e:
+                # logs to console for debugging
+                print(f"[Step 2 ERROR] {type(e).__name__}: {e}", flush=True)
+                return None, f"{type(e).__name__}: {e}"   # image=None, show error
 
         inpaint_btn.click(
             fn=guarded_inpaint,
             inputs=[shared_output, inpaint_prompt, inpaint_negative],
-            outputs=inpaint_result,
+            outputs=[inpaint_result, error_box],  
         )
 
     # Enable the Step 2 button once Step 1 finishes

yolov8x-seg.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:c333d6a7aff884ec821050eacdab29871f3edd10d996cf48746a02f8e01f1b5d
+size 144101612