app.py

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):
    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)

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)
        )

    # ─────────── 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],
        )

    event.then(lambda: gr.update(interactive=True), None, inpaint_btn)

    demo.launch(debug=True)