Fix step error handling and wire up error_box

app.py

@@ -1,130 +1,90 @@
+
 import os
 import traceback
 from datetime import datetime
 import torch, gc
 from PIL import Image
 import gradio as gr
+
 from inference import generate_with_lora
 from background_edit import run_background_removal_and_inpaint
 
-
 # ───────────────────── 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(image_path, *args, **kwargs):
+def safe_generate_and_inpaint(image, prompt, negative_prompt, strength, guidance_scale):
     try:
-        unload_models()  # free VRAM before loading the inpainting model
-        return run_background_removal_and_inpaint(image_path, *args, **kwargs)
-    except Exception as e:
-        _print_trace()
-        raise gr.Error(f"[Step 2] Background replacement failed: {type(e).__name__}: {e}")
+        if image is None:
+            raise gr.Error("Please upload an image first.")
+
+        # Step 1: Refinement
+        print("[INFO] Step 1: Generating refined image...", flush=True)
+        refined = generate_with_lora(
+            image=image,
+            prompt=prompt,
+            negative_prompt=negative_prompt,
+            strength=strength,
+            guidance_scale=guidance_scale,
+        )
 
-def _save_to_disk(img):
-    if img is None:
-        return gr.skip()
-    
-    os.makedirs("./outputs", exist_ok=True)
-    ts = datetime.now().strftime("%Y%m%d_%H%M%S")
-    path = f"./outputs/step1_result_{ts}.png"
-    img.save(path)
-    return path
+        # Save to disk
+        os.makedirs("./outputs", exist_ok=True)
+        ts = datetime.now().strftime("%Y%m%d_%H%M%S")
+        path = f"./outputs/step1_result_{ts}.png"
+        refined.save(path)
 
+        # Step 2: Background removal and inpainting
+        print("[INFO] Step 2: Inpainting background...", flush=True)
+        unload_models()
+        result = run_background_removal_and_inpaint(path, prompt, negative_prompt)
 
-# ───────────────────── UI ─────────────────────
-shared_output_path = gr.State()  # holds file path to Step 1 output
-original_input     = gr.State()  # holds the original upload (if needed)
+        return refined, result, ""
 
+    except gr.Error as e:
+        return None, None, f"🛑 {str(e)}"
+    except Exception as e:
+        _print_trace()
+        return None, None, f"❌ Unexpected Error: {type(e).__name__}: {str(e)}"
+
+# ───────────────────── Gradio UI ─────────────────────
 with gr.Blocks() as demo:
-    demo.queue()
-
-    # ─────────── 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")
-
-        event = (
-            run_btn.click(
-                fn=safe_generate_with_lora,
-                inputs=[input_image, prompt, negative_prompt, strength, guidance],
-                outputs=output_image,
-            )
-            .then(_save_to_disk, output_image, shared_output_path)
-            .then(lambda x: x, input_image, original_input)
-        )
+    gr.Markdown("## 🖼️ AI Headshot Enhancer + Background Replacer")
+    gr.Markdown("Upload a headshot, adjust the prompt, and click one button. We'll refine the image and replace the background automatically.")
+
+    with gr.Row():
+        input_image = gr.Image(type="pil", label="Upload Headshot")
 
-    # ─────────── STEP 2: Background Replacement ───────────
-    with gr.Tab("Step 2: Replace Background"):
-        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(image_path, prompt_bg, neg_bg):
-            if not image_path or not os.path.isfile(image_path):
-                return None, "**🛑 Error:** No valid headshot found — please run Step 1 first."
-
-            try:
-                print(f"[DEBUG] Loading image from: {image_path}", flush=True)
-                result = safe_run_background(image_path, prompt_bg, neg_bg)
-                return result, ""
-            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_path, inpaint_prompt, inpaint_negative],
-            outputs=[inpaint_result, error_box],
+    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="Refinement Strength")
+        guidance = gr.Slider(1, 20, value=17.0, step=0.5, label="Guidance Scale")
+
+    go_btn = gr.Button("✨ Run Full Process (Refine + Inpaint)")
+
+    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")
+
+    error_box = gr.Markdown(label="Error", value="", visible=True)
 
-    event.then(lambda: gr.update(interactive=True), None, inpaint_btn)
+    go_btn.click(
+        fn=safe_generate_and_inpaint,
+        inputs=[input_image, prompt, negative_prompt, strength, guidance],
+        outputs=[output_refined, output_final, error_box]
+    )
 
-    demo.launch(debug=True)
+demo.launch(debug=True)