Fix step error handling and wire up error_box

app.py

@@ -1,103 +1,130 @@
 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
+import traceback
+from datetime impor
+import torch, gc
 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
+
+
+# ───────────────────── 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):
+    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}")
+
+def _save_to_disk(img):
+    if img is None:
+        return gr.skip()
+    
+    os.makedirs("/tmp/gradio_outputs", exist_ok=True)
+    ts = datetime.now().strftime("%Y%m%d_%H%M%S")
+    path = f"/tmp/gradio_outputs/step1_result_{ts}.png"
+    img.save(path)
+    return path
+
+
+# ───────────────────── UI ─────────────────────
+shared_output_path = gr.State()  # holds file path to Step 1 output
+original_input     = gr.State()  # holds the original upload (if needed)
 
-# ---- 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:
-    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")
-
-    with gr.Row():
-        prompt = gr.Textbox(
-            label="New Background Prompt",
-            value="modern open-plan office, soft natural light, minimalistic decor"
+    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)
         )
-        neg_prompt = gr.Textbox(
-            label="Negative Prompt",
-            value="cartoon, fantasy, dark lighting, painting, anime"
-        )
-
-    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}"
+    # ─────────── 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],
+        )
 
-    run_btn = gr.Button("Run Background Inpaint")
-    run_btn.click(
-        fn=safe_run,
-        inputs=[input_img, prompt, neg_prompt],
-        outputs=[output_img, error_box]
-    )
+    event.then(lambda: gr.update(interactive=True), None, inpaint_btn)
 
-demo.launch(debug=True)
+    demo.launch(debug=True)

app_remov.py

@@ -0,0 +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
+
+# ---- 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:
+    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")
+
+    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"
+        )
+
+    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)

background_edit.py

@@ -33,38 +33,27 @@ inpaint_pipe = StableDiffusionXLInpaintPipeline.from_pretrained(
     use_auth_token=os.getenv("HF_TOKEN")
 ).to("cuda")
 
-def run_background_removal_and_inpaint(shared_output, prompt, negative_prompt):
-
-    # Get image from shared_output
-    if isinstance(shared_output, dict):
-        image = shared_output.get("step1")
-    else:
-        image = None
-
-    if image is None:
-        raise gr.Error("Run Step 1 first to get a base image.")
+def run_background_removal_and_inpaint(image_path, prompt, negative_prompt):
+    if not image_path or not os.path.isfile(image_path):
+        raise gr.Error("No valid image found. Please run Step 1 first.")
 
+    image = Image.open(image_path).convert("RGB")
     img_cv = pil_to_cv2(image)
     results = yolo(img_cv)
 
-    # ✅ Validate YOLO detection result
     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()
-
-    # Create inpainting mask
     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)
 
-    # Resize and prepare images
     mask_pil = cv2_to_pil(inpaint_mask).resize((1024, 1024)).convert("L")
     img_pil = image.resize((1024, 1024)).convert("RGB")
 
-    # Inpaint
     result = inpaint_pipe(
         prompt=prompt,
         negative_prompt=negative_prompt or "",