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
from background_edit import run_clothing_inpaint

import os

MODEL_URL = "https://huggingface.co/Bingsu/adetailer/resolve/main/deepfashion2_yolov8s-seg.pt"
MODEL_PATH = "deepfashion2_yolov8s-seg.pt"

# ─── Download DeepFashion2 model if not already present ───
if not os.path.exists(MODEL_PATH):
    import urllib.request
    print("[INFO] Downloading DeepFashion2 YOLOv8 model...")
    urllib.request.urlretrieve(MODEL_URL, MODEL_PATH)
    print("[INFO] Model downloaded.")
else:
    print("[INFO] DeepFashion2 model already exists.")


# ───────────────────── Helpers ─────────────────────
def _print_trace():
    traceback.print_exc()

def unload_models():
    torch.cuda.empty_cache()
    gc.collect()

def safe_generate_and_inpaint(
    image,
    prompt_1, neg_1, strength_1, guidance_1,
    prompt_2, neg_2, guidance_2
):
    try:
        if image is None:
            raise gr.Error("Please upload an image first.")

        # Step 1: Headshot Refinement
        print("[INFO] Step 1: Refining headshot...", flush=True)
        refined = generate_with_lora(
            image=image,
            prompt=prompt_1,
            negative_prompt=neg_1,
            strength=strength_1,
            guidance_scale=guidance_1,
        )

        # Save intermediate result 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 Inpainting
        print("[INFO] Step 2: Inpainting background...", flush=True)
        unload_models()
        result = run_background_removal_and_inpaint(
            image_path=path,
            prompt=prompt_2,
            negative_prompt=neg_2,
            guidance_scale=guidance_2
        )

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

def guarded_clothing(image, prompt, neg, guidance):
    try:
        result, err = run_clothing_inpaint(image, prompt, neg, guidance)
        return result, err
    except Exception as e:
        import traceback
        traceback.print_exc()
        return None, f"❌ Unexpected Error: {type(e).__name__}: {str(e)}"

# ───────────────────── Gradio UI ─────────────────────
with gr.Blocks() as demo:
    gr.Markdown("## 🧠 Headshot + Background Generator (Full Prompt Control)")

    with gr.Row():
        input_image = gr.Image(type="pil", label="Upload Headshot")

    gr.Markdown("### Step 1: Headshot Refinement (LoRA)")

    with gr.Row():
        prompt_1 = gr.Textbox(label="Headshot Prompt", value="a professional headshot of a confident woman in her 30s with blonde hair")
        neg_1 = gr.Textbox(label="Headshot Negative Prompt", value="deformed, cartoon, anime, sketch, blurry, low quality")

    with gr.Row():
        strength_1 = gr.Slider(0.1, 1.0, value=0.2, step=0.05, label="Refinement Strength")
        guidance_1 = gr.Slider(1, 20, value=17, step=0.5, label="Guidance Scale (Headshot)")

    gr.Markdown("### Step 2: Background Inpainting (SDXL)")

    with gr.Row():
        prompt_2 = gr.Textbox(label="Background Prompt", value="modern hospital background, clean, soft lighting")
        neg_2 = gr.Textbox(label="Background Negative Prompt", value="fantasy, cartoon, cluttered, sketch")

    with gr.Row():
        guidance_2 = gr.Slider(1, 20, value=10, step=0.5, label="Guidance Scale (Background)")

    go_btn = gr.Button("✨ Generate Refined Headshot + Background")

    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)

    go_btn.click(
        fn=safe_generate_and_inpaint,
        inputs=[
            input_image, prompt_1, neg_1, strength_1, guidance_1,
            prompt_2, neg_2, guidance_2
        ],
        outputs=[output_refined, output_final, error_box]
    )

    gr.Markdown("### 👗 Step 3: Clothing Replacement")

    with gr.Row():
        clothing_prompt = gr.Textbox(
            label="Clothing Prompt",
            value="white female CEO professional blazer, clean look"
        )
        clothing_negative = gr.Textbox(
            label="Clothing Negative Prompt",
            value="hoodie, casual wear, fantasy, cartoon, jeans, distorted, blurry"
        )

    with gr.Row():
        clothing_guidance = gr.Slider(1, 20, value=17.0, step=0.5, label="Clothing Guidance Scale")

    with gr.Row():
        clothing_btn = gr.Button("🧵 Inpaint Clothing")
        clothing_output = gr.Image(type="pil", label="Step 3: Final Image with New Clothing")

    clothing_error = gr.Markdown(label="Clothing Error", value="", visible=True)


    clothing_btn.click(
        fn=guarded_clothing,
        inputs=[output_final, clothing_prompt, clothing_negative, clothing_guidance],
        outputs=[clothing_output, clothing_error],
        preprocess=False
    )

demo.launch(debug=True)