import spaces
import gradio as gr

import torch
import numpy as np
import cv2
import mediapipe as mp

from PIL import Image
from torchvision import transforms
from transformers import AutoModelForImageSegmentation

# ============================================================
# CONFIG
# ============================================================

torch.set_float32_matmul_precision("high")

DEVICE = "cuda" if torch.cuda.is_available() else "cpu"

# ------------------------------------------------------------
# PRIMARY MODELS
# ------------------------------------------------------------

GLOBAL_MODEL_ID = "ZhengPeng7/BiRefNet_HR"
MATTING_MODEL_ID = "ZhengPeng7/BiRefNet_HR-matting"

# ------------------------------------------------------------
# FALLBACK MODELS
# ------------------------------------------------------------

GLOBAL_FALLBACK_ID = "ZhengPeng7/BiRefNet"
MATTING_FALLBACK_ID = "ZhengPeng7/BiRefNet-matting"

# ------------------------------------------------------------
# INFERENCE
# ------------------------------------------------------------

HR_MAX_SIZE = 2048
FALLBACK_MAX_SIZE = 1536

# ------------------------------------------------------------
# HEAD / BEARD ROI
# ------------------------------------------------------------

HEAD_PAD_X = 0.85
HEAD_PAD_TOP = 0.90
HEAD_PAD_BOTTOM = 1.90

HEAD_UPSCALE = 1.50
HEAD_MAX_SIZE = 2048

# ------------------------------------------------------------
# NORMALIZATION
# ------------------------------------------------------------

MEAN = [0.485, 0.456, 0.406]
STD = [0.229, 0.224, 0.225]

# ============================================================
# MODEL LOADING
# ============================================================


def load_model(primary_id, fallback_id, name):

    print()
    print("=" * 70)
    print(f"LOADING {name}")
    print(f"Primary:  {primary_id}")
    print("=" * 70)

    try:

        model = AutoModelForImageSegmentation.from_pretrained(
            primary_id,
            trust_remote_code=True
        )

        model.to(DEVICE)
        model.eval()

        print(f"{name}: PRIMARY MODEL READY")

        return model, True

    except Exception as primary_error:

        print()
        print(f"WARNING: {name} primary model failed:")
        print(primary_error)
        print()
        print(f"Loading fallback: {fallback_id}")

        model = AutoModelForImageSegmentation.from_pretrained(
            fallback_id,
            trust_remote_code=True
        )

        model.to(DEVICE)
        model.eval()

        print(f"{name}: FALLBACK MODEL READY")

        return model, False


global_model, GLOBAL_HR = load_model(
    GLOBAL_MODEL_ID,
    GLOBAL_FALLBACK_ID,
    "GLOBAL SEGMENTATION"
)

matting_model, MATTING_HR = load_model(
    MATTING_MODEL_ID,
    MATTING_FALLBACK_ID,
    "BEARD / HAIR MATTING"
)

print()
print("=" * 70)
print("MODEL STATUS")
print("Global HR :", GLOBAL_HR)
print("Matte HR  :", MATTING_HR)
print("=" * 70)
print()

# ============================================================
# MEDIAPIPE
# ============================================================

mp_face_detection = mp.solutions.face_detection

face_detector = mp_face_detection.FaceDetection(
    model_selection=1,
    min_detection_confidence=0.45
)

# ============================================================
# TRANSFORM
# ============================================================


def make_transform(size):

    return transforms.Compose([
        transforms.Resize(
            size,
            interpolation=transforms.InterpolationMode.BILINEAR
        ),

        transforms.ToTensor(),

        transforms.Normalize(
            MEAN,
            STD
        )
    ])


# ============================================================
# MODEL SIZE
# ============================================================


def inference_size(h, w, max_size):

    largest = max(h, w)

    scale = min(
        max_size / float(largest),
        1.0
    )

    nh = int(
        round(
            (h * scale) / 32
        ) * 32
    )

    nw = int(
        round(
            (w * scale) / 32
        ) * 32
    )

    nh = max(32, nh)
    nw = max(32, nw)

    return nh, nw


# ============================================================
# MODEL RUNNER
# ============================================================


def run_model(
    model,
    image,
    max_size
):

    h = image.height
    w = image.width

    target_h, target_w = inference_size(
        h,
        w,
        max_size
    )

    transform = make_transform(
        (target_h, target_w)
    )

    tensor = transform(
        image
    ).unsqueeze(0).to(DEVICE)

    # --------------------------------------------------------
    # Do not force FP16 on CPU.
    # --------------------------------------------------------

    if DEVICE == "cuda":
        tensor = tensor.half()

    with torch.inference_mode():

        prediction = model(
            tensor
        )

        # Most BiRefNet variants return a list/tuple
        # of progressively refined predictions.
        if isinstance(prediction, (list, tuple)):

            prediction = prediction[-1]

        elif isinstance(prediction, dict):

            values = list(prediction.values())
            prediction = values[-1]

        prediction = prediction.sigmoid()

    alpha = (
        prediction[0]
        .squeeze()
        .float()
        .cpu()
        .numpy()
    )

    # --------------------------------------------------------
    # Back to EXACT original dimensions.
    # --------------------------------------------------------

    alpha = cv2.resize(
        alpha,
        (w, h),
        interpolation=cv2.INTER_CUBIC
    )

    return np.clip(
        alpha,
        0.0,
        1.0
    ).astype(np.float32)


# ============================================================
# FACE / BEARD REGION
# ============================================================


def detect_head_region(rgb):

    h, w = rgb.shape[:2]

    try:

        result = face_detector.process(rgb)

        if not result.detections:
            return None

        detection = max(
            result.detections,
            key=lambda d: (
                max(
                    0.0,
                    d.location_data
                    .relative_bounding_box
                    .width
                )
                *
                max(
                    0.0,
                    d.location_data
                    .relative_bounding_box
                    .height
                )
            )
        )

        box = (
            detection
            .location_data
            .relative_bounding_box
        )

        x = int(box.xmin * w)
        y = int(box.ymin * h)

        bw = int(box.width * w)
        bh = int(box.height * h)

        # ----------------------------------------------------
        # BIG REGION.
        #
        # Especially important:
        # beard extends far below normal face box.
        # ----------------------------------------------------

        padx = int(bw * HEAD_PAD_X)
        padtop = int(bh * HEAD_PAD_TOP)
        padbottom = int(bh * HEAD_PAD_BOTTOM)

        x0 = max(
            0,
            x - padx
        )

        y0 = max(
            0,
            y - padtop
        )

        x1 = min(
            w,
            x + bw + padx
        )

        y1 = min(
            h,
            y + bh + padbottom
        )

        if x1 <= x0 or y1 <= y0:
            return None

        return (
            x0,
            y0,
            x1,
            y1
        )

    except Exception as e:

        print(
            "Face detection error:",
            e
        )

        return None


# ============================================================
# ALPHA HELPERS
# ============================================================


def alpha_uncertainty(alpha):

    return np.clip(
        1.0 - np.abs(alpha - 0.5) * 2.0,
        0.0,
        1.0
    ).astype(np.float32)


def alpha_boundary(alpha):

    fg = (
        alpha > 0.015
    ).astype(np.uint8)

    kernel = cv2.getStructuringElement(
        cv2.MORPH_ELLIPSE,
        (9, 9)
    )

    dilated = cv2.dilate(
        fg,
        kernel
    )

    eroded = cv2.erode(
        fg,
        kernel
    )

    band = (
        dilated.astype(np.float32)
        -
        eroded.astype(np.float32)
    )

    uncertain = alpha_uncertainty(
        alpha
    )

    band = np.maximum(
        band,
        uncertain
    )

    return np.clip(
        band,
        0.0,
        1.0
    )


# ============================================================
# VERY LIGHT GUIDED EDGE ALIGNMENT
# ============================================================


def guided_alpha(
    alpha,
    rgb,
    radius=2,
    eps=0.00035
):

    if not hasattr(cv2, "ximgproc"):
        return alpha

    if not hasattr(
        cv2.ximgproc,
        "guidedFilter"
    ):
        return alpha

    guide = (
        rgb.astype(np.float32)
        / 255.0
    )

    try:

        filtered = cv2.ximgproc.guidedFilter(
            guide=guide,
            src=alpha.astype(np.float32),
            radius=radius,
            eps=eps
        )

        # IMPORTANT:
        # Only 20% guided correction.
        #
        # Strong guided filtering can destroy
        # individual beard hairs.

        result = (
            alpha * 0.80
            +
            filtered * 0.20
        )

        return np.clip(
            result,
            0.0,
            1.0
        ).astype(np.float32)

    except Exception:

        return alpha


# ============================================================
# HEAD MATTE
# ============================================================


def run_head_matte(
    rgb,
    global_alpha
):

    box = detect_head_region(
        rgb
    )

    if box is None:

        print(
            "No face detected -> no local beard crop."
        )

        return global_alpha

    x0, y0, x1, y1 = box

    crop = rgb[
        y0:y1,
        x0:x1
    ]

    crop_h, crop_w = crop.shape[:2]

    if crop.size == 0:
        return global_alpha

    print(
        f"HEAD/BEARD ROI: {crop_w}x{crop_h}"
    )

    # --------------------------------------------------------
    # High-resolution local image.
    #
    # This is ONLY used to estimate alpha.
    # It is never used as final RGB.
    # --------------------------------------------------------

    scale = HEAD_UPSCALE

    high_w = min(
        HEAD_MAX_SIZE,
        max(
            crop_w,
            int(crop_w * scale)
        )
    )

    high_h = min(
        HEAD_MAX_SIZE,
        max(
            crop_h,
            int(crop_h * scale)
        )
    )

    highres = cv2.resize(
        crop,
        (high_w, high_h),
        interpolation=cv2.INTER_CUBIC
    )

    matte_input = Image.fromarray(
        highres
    )

    print(
        f"BEARD MATTE INPUT: {high_w}x{high_h}"
    )

    matte_high = run_model(
        matting_model,
        matte_input,
        (
            HR_MAX_SIZE
            if MATTING_HR
            else FALLBACK_MAX_SIZE
        )
    )

    # --------------------------------------------------------
    # Return matte to original crop dimensions.
    # --------------------------------------------------------

    matte = cv2.resize(
        matte_high,
        (crop_w, crop_h),
        interpolation=cv2.INTER_LANCZOS4
    )

    coarse = global_alpha[
        y0:y1,
        x0:x1
    ]

    # ========================================================
    # CRITICAL PART
    # ========================================================
    #
    # - global controls confident interior/background
    # - local matte controls boundary + fine hairs
    # ========================================================

    uncertain = alpha_uncertainty(
        coarse
    )

    boundary = alpha_boundary(
        coarse
    )

    local_weight = (
        uncertain * 0.70
        +
        boundary * 0.55
    )

    local_weight = np.clip(
        local_weight,
        0.0,
        1.0
    )

    # --------------------------------------------------------
    # If global is strongly confident foreground,
    # preserve global interior.
    # --------------------------------------------------------

    local_weight[
        coarse > 0.96
    ] = np.minimum(
        local_weight[
            coarse > 0.96
        ],
        0.12
    )

    # --------------------------------------------------------
    # If global is strongly background, still allow the
    # matting model to recover tiny beard/hair strands.
    # --------------------------------------------------------

    tiny_hair_zone = (
        (coarse > 0.003)
        &
        (coarse < 0.35)
    )

    local_weight[
        tiny_hair_zone
    ] = np.maximum(
        local_weight[
            tiny_hair_zone
        ],
        0.72
    )

    # --------------------------------------------------------
    # FIX: this blur was previously part of a 3-pass blur
    # stack (here + preserve_fine_alpha + micro_sharpen_alpha)
    # that softened 1-3px hair strands. Reduced further since
    # the other two passes downstream are now toned down too.
    # --------------------------------------------------------

    local_weight = cv2.GaussianBlur(
        local_weight,
        (0, 0),
        0.20
    )

    # --------------------------------------------------------
    # Matte is authoritative around beard/hair.
    # --------------------------------------------------------

    refined = (
        matte * local_weight
        +
        coarse * (1.0 - local_weight)
    )

    # --------------------------------------------------------
    # Restore very confident global interior.
    # --------------------------------------------------------

    refined[
        coarse > 0.985
    ] = coarse[
        coarse > 0.985
    ]

    # --------------------------------------------------------
    # Edge alignment is deliberately tiny.
    # --------------------------------------------------------

    refined = guided_alpha(
        refined,
        crop,
        radius=2,
        eps=0.0003
    )

    output = global_alpha.copy()

    output[
        y0:y1,
        x0:x1
    ] = refined

    return np.clip(
        output,
        0.0,
        1.0
    ).astype(np.float32)


# ============================================================
# FINE DETAIL RECOVERY
# ============================================================


def preserve_fine_alpha(alpha):

    # --------------------------------------------------------
    # FIX: sigma reduced from 0.55 -> 0.35 and strength halved.
    # This was the second of three stacked blurs that were
    # collectively softening thin hair strands. It still
    # recovers real high-frequency detail, just less
    # aggressively, so it doesn't fight the matte's own edges.
    # --------------------------------------------------------

    smooth = cv2.GaussianBlur(
        alpha,
        (0, 0),
        0.35
    )

    detail = (
        alpha
        -
        smooth
    )

    uncertainty = alpha_uncertainty(
        alpha
    )

    strength = (
        0.05
        +
        0.05 * uncertainty
    )

    result = (
        alpha
        +
        detail * strength
    )

    return np.clip(
        result,
        0.0,
        1.0
    ).astype(np.float32)


# ============================================================
# MICRO EDGE SHARPEN
# ============================================================


def micro_sharpen_alpha(alpha):

    # --------------------------------------------------------
    # FIX: this was the third blur in the stack. Sigma reduced
    # and strength reduced — this stage was contributing very
    # little sharpening but a full extra blur pass of softening
    # to every hair edge.
    # --------------------------------------------------------

    blurred = cv2.GaussianBlur(
        alpha,
        (0, 0),
        0.25
    )

    detail = (
        alpha
        -
        blurred
    )

    uncertainty = alpha_uncertainty(
        alpha
    )

    result = (
        alpha
        +
        detail
        *
        (
            0.04
            *
            uncertainty
        )
    )

    return np.clip(
        result,
        0.0,
        1.0
    ).astype(np.float32)


# ============================================================
# ALPHA CLEANUP
# ============================================================


def cleanup_alpha(alpha):

    result = np.clip(
        alpha,
        0.0,
        1.0
    ).astype(np.float32)

    # Only remove almost-zero numerical noise.
    result[
        result < 0.002
    ] = 0.0

    result[
        result > 0.999
    ] = 1.0

    return result


# ============================================================
# BACKGROUND ESTIMATION
# ============================================================


def estimate_background_color(
    rgb,
    alpha
):

    h, w = alpha.shape

    border = max(
        4,
        int(min(h, w) * 0.025)
    )

    border_mask = np.zeros(
        (h, w),
        dtype=bool
    )

    border_mask[
        :border,
        :
    ] = True

    border_mask[
        -border:,
        :
    ] = True

    border_mask[
        :,
        :border
    ] = True

    border_mask[
        :,
        -border:
    ] = True

    mask = (
        border_mask
        &
        (alpha < 0.02)
    )

    pixels = rgb[
        mask
    ]

    if len(pixels) < 100:

        pixels = rgb[
            alpha < 0.02
        ]

    if len(pixels) < 100:

        return np.array(
            [255.0, 255.0, 255.0],
            dtype=np.float32
        )

    return np.median(
        pixels.astype(np.float32),
        axis=0
    )


# ============================================================
# LOCAL BACKGROUND COLOR
# ============================================================


def local_background_color(
    rgb,
    alpha
):
    """
    Estimate background around each pixel.

    This is better than using one global color because
    real photos may have gradients/shadows.
    """

    background_mask = (
        alpha < 0.025
    ).astype(np.float32)

    weight = cv2.GaussianBlur(
        background_mask,
        (0, 0),
        7.0
    )

    weight = np.maximum(
        weight,
        1e-5
    )

    rgb_float = rgb.astype(
        np.float32
    )

    channels = []

    for c in range(3):

        channel = (
            rgb_float[..., c]
            *
            background_mask
        )

        smooth_channel = cv2.GaussianBlur(
            channel,
            (0, 0),
            7.0
        )

        local = (
            smooth_channel
            /
            weight
        )

        channels.append(
            local
        )

    bg = np.stack(
        channels,
        axis=2
    )

    global_bg = estimate_background_color(
        rgb,
        alpha
    )

    # If local estimate has no real background nearby,
    # fall back toward global estimate.
    bg = np.where(
        np.isfinite(bg),
        bg,
        global_bg.reshape(1, 1, 3)
    )

    bg = np.clip(
        bg,
        0,
        255
    )

    return bg.astype(np.float32)


# ============================================================
# BACKGROUND COLOR DECONTAMINATION  (FIXED)
# ============================================================


def decontaminate_edges(
    rgb,
    alpha
):
    """
    Remove background color contamination from semi-transparent
    edge pixels.

    Formula:
        C = alpha * F + (1-alpha) * B
    therefore:
        F = (C - (1-alpha)B) / alpha

    ------------------------------------------------------------
    THE BUG THAT WAS HERE:

    The original code applied strong correction to mid-alpha
    pixels but then explicitly SUPPRESSED correction on the
    lowest-alpha pixels:

        strength[a < 0.10] *= 0.35

    and used a high safe_alpha floor (0.16), which under-divides
    the recovered color for anything below that.

    But thin beard/hair flyaways sit exactly in that a < 0.10
    range. That's precisely where the blue/cyan background tint
    was showing up — the math correctly detected the
    contamination, then was told to only partially fix it.

    THE FIX:
    - Lower the safe_alpha floor so correction isn't artificially
      weakened for real hair pixels.
    - Do NOT crush "strength" for low alpha. Instead correct at
      (near) full strength, then denoise the recovered color in
      the low-alpha band with a small median filter. This keeps
      the instability under control without leaving background
      color baked into the hair.
    ------------------------------------------------------------
    """

    image = rgb.astype(
        np.float32
    )

    a = np.clip(
        alpha.astype(np.float32),
        0.0,
        1.0
    )

    background = local_background_color(
        rgb,
        alpha
    )

    a3 = a[..., None]

    # --------------------------------------------------------
    # Edge zone. Widened slightly on the low end so very thin
    # flyaway hairs (alpha ~0.03-0.06) are included.
    # --------------------------------------------------------

    edge = (
        (a > 0.03)
        &
        (a < 0.94)
    )

    if not np.any(edge):

        return rgb

    # --------------------------------------------------------
    # FIX: lower floor (was 0.16). Let the correction actually
    # happen for thin hair instead of pretending alpha is higher
    # than it is.
    # --------------------------------------------------------

    safe_alpha = np.maximum(
        a3,
        0.06
    )

    recovered = (
        image
        -
        (1.0 - a3)
        *
        background
    ) / safe_alpha

    recovered = np.clip(
        recovered,
        0.0,
        255.0
    )

    # --------------------------------------------------------
    # FIX: denoise the recovered color ONLY in the low-alpha
    # band. Division by a small number amplifies per-pixel
    # sensor noise; a 3x3 median kills that noise without
    # re-introducing background color (median doesn't blend
    # in neighboring background pixels' color the way a
    # Gaussian blur would).
    # --------------------------------------------------------

    low_alpha_mask = (a < 0.20) & edge

    if np.any(low_alpha_mask):

        recovered_u8 = np.clip(
            recovered, 0, 255
        ).astype(np.uint8)

        recovered_denoised = cv2.medianBlur(
            recovered_u8,
            3
        ).astype(np.float32)

        recovered = np.where(
            low_alpha_mask[..., None],
            recovered_denoised,
            recovered
        )

    # --------------------------------------------------------
    # Determine how strongly to correct.
    #
    # FIX: strength now RISES toward low alpha instead of being
    # cut down. Low-alpha pixels are the ones most likely to be
    # majority background, so they need the most correction, not
    # the least.
    # --------------------------------------------------------

    low_alpha_strength = np.clip(
        (0.94 - a) / 0.70,
        0.0,
        1.0
    )

    distance = np.linalg.norm(
        image - background,
        axis=2
    )

    similarity = np.exp(
        -(
            distance / 85.0
        ) ** 2
    )

    strength = (
        low_alpha_strength
        *
        (
            0.35
            +
            0.55 * similarity
        )
    )

    # Allow near-full correction now (was capped at 0.58).
    strength = np.clip(
        strength,
        0.0,
        0.95
    )

    # Strong foreground should barely change.
    strength[
        a > 0.82
    ] *= 0.20

    # --------------------------------------------------------
    # FIX: removed the old line that crushed strength for
    # a < 0.10 (`strength[a < 0.10] *= 0.35`). That line was
    # the main cause of the residual blue fringe. The median
    # denoise above handles the instability instead.
    # --------------------------------------------------------

    strength3 = strength[..., None]

    corrected = (
        image * (1.0 - strength3)
        +
        recovered * strength3
    )

    output = np.where(
        edge[..., None],
        corrected,
        image
    )

    return np.clip(
        output,
        0,
        255
    ).astype(np.uint8)


# ============================================================
# ADDITIONAL BLUE-SPILL SAFETY  (FIXED)
# ============================================================


def remove_background_tint(
    rgb,
    alpha
):
    """
    Conservative second pass for strongly colored backgrounds
    (blue/green screen etc).

    FIX: previously this ALSO crushed correction for the
    lowest-alpha pixels (`strength[a < 0.08] *= 0.25`), stacking
    on top of the same bug in decontaminate_edges. Removed that
    line so this pass actually helps rather than mostly no-oping
    on exactly the pixels that show the tint.
    """

    image = rgb.astype(
        np.float32
    )

    a = alpha.astype(
        np.float32
    )

    bg = estimate_background_color(
        rgb,
        alpha
    )

    bg_mean = float(
        np.mean(bg)
    )

    chroma = float(
        np.max(bg) - np.min(bg)
    )

    # If background isn't strongly colored, use only tiny
    # correction. If it IS strongly colored (like your blue
    # studio background), allow more correction than before.
    if chroma < 18.0:
        max_strength = 0.10
    else:
        max_strength = 0.32

    bg3 = bg.reshape(
        1,
        1,
        3
    )

    distance = np.linalg.norm(
        image - bg3,
        axis=2
    )

    similarity = np.exp(
        -(
            distance / 95.0
        ) ** 2
    )

    edge_weight = np.clip(
        (0.85 - a) / 0.70,
        0.0,
        1.0
    )

    strength = (
        similarity
        *
        edge_weight
        *
        max_strength
    )

    # Do not touch solid beard/body.
    strength[
        a > 0.80
    ] = 0.0

    # --------------------------------------------------------
    # FIX: removed `strength[a < 0.08] *= 0.25`. That was
    # suppressing correction on exactly the thinnest hairs,
    # doubling up with the same bug in decontaminate_edges.
    # --------------------------------------------------------

    strength3 = strength[..., None]

    corrected = (
        image
        +
        (
            image - bg3
        )
        *
        strength3
    )

    corrected = np.clip(
        corrected,
        0,
        255
    )

    return corrected.astype(
        np.uint8
    )


# ============================================================
# FINAL EDGE PROTECTION
# ============================================================


def protect_fine_hair(
    original_rgb,
    cleaned_rgb,
    alpha
):
    """
    Prevent the color cleanup stage from changing solid
    foreground pixels. Hair/beard edge pixels remain corrected,
    while solid subject pixels stay essentially original.
    """

    a = alpha.astype(
        np.float32
    )

    original = original_rgb.astype(
        np.float32
    )

    cleaned = cleaned_rgb.astype(
        np.float32
    )

    strength = np.clip(
        (0.92 - a) / 0.75,
        0.0,
        1.0
    )

    strength[
        a > 0.86
    ] = 0.0

    strength *= 0.85

    strength3 = strength[..., None]

    result = (
        original * (1.0 - strength3)
        +
        cleaned * strength3
    )

    return np.clip(
        result,
        0,
        255
    ).astype(np.uint8)


# ============================================================
# FINAL ALPHA
# ============================================================


def finalize_alpha(alpha):

    alpha = np.clip(
        alpha,
        0.0,
        1.0
    ).astype(np.float32)

    alpha[
        alpha < 0.0015
    ] = 0.0

    alpha[
        alpha > 0.9995
    ] = 1.0

    return alpha


# ============================================================
# MAIN PIPELINE
# ============================================================


@spaces.GPU(duration=180)
def remove_background(
    image
):

    if image is None:

        raise gr.Error(
            "Please upload a photo first."
        )

    try:

        # ====================================================
        # ORIGINAL
        # ====================================================

        original = image.convert(
            "RGB"
        )

        rgb = np.array(
            original,
            dtype=np.uint8
        )

        h, w = rgb.shape[:2]

        print()
        print("=" * 70)
        print(f"INPUT: {w}x{h}")
        print("=" * 70)

        # ====================================================
        # STAGE 1 — GLOBAL SEGMENTATION
        # ====================================================

        print("STAGE 1: Global subject segmentation")

        global_alpha = run_model(
            global_model,
            original,
            (
                HR_MAX_SIZE
                if GLOBAL_HR
                else FALLBACK_MAX_SIZE
            )
        )

        global_alpha = guided_alpha(
            global_alpha,
            rgb,
            radius=2,
            eps=0.0005
        )

        # ====================================================
        # STAGE 2 — BEARD / HAIR MATTE
        # ====================================================

        print("STAGE 2: High-resolution beard/hair matting")

        alpha = run_head_matte(
            rgb,
            global_alpha
        )

        # ====================================================
        # STAGE 3 — FINE DETAIL
        # ====================================================

        print("STAGE 3: Fine hair/beard detail preservation")

        alpha = preserve_fine_alpha(
            alpha
        )

        # ====================================================
        # STAGE 4 — MICRO SHARPEN
        # ====================================================

        print("STAGE 4: Micro alpha sharpening")

        alpha = micro_sharpen_alpha(
            alpha
        )

        # ====================================================
        # STAGE 5 — FINAL ALPHA
        # ====================================================

        alpha = cleanup_alpha(
            alpha
        )

        alpha = finalize_alpha(
            alpha
        )

        # ====================================================
        # STAGE 6 — BACKGROUND COLOR REMOVAL (fixed)
        # ====================================================

        print("STAGE 5: Background color decontamination")

        clean_rgb = decontaminate_edges(
            rgb,
            alpha
        )

        # ====================================================
        # STAGE 7 — COLOR-SPILL SAFETY (fixed)
        # ====================================================

        print("STAGE 6: Fine background-tint cleanup")

        clean_rgb = remove_background_tint(
            clean_rgb,
            alpha
        )

        # ====================================================
        # STAGE 8 — PROTECT SUBJECT
        # ====================================================

        clean_rgb = protect_fine_hair(
            rgb,
            clean_rgb,
            alpha
        )

        # ====================================================
        # FINAL PNG
        # ====================================================

        print("Creating final transparent PNG...")

        result = Image.fromarray(
            clean_rgb,
            mode="RGB"
        )

        alpha_image = Image.fromarray(
            np.clip(
                alpha * 255.0,
                0,
                255
            ).astype(np.uint8),
            mode="L"
        )

        result.putalpha(
            alpha_image
        )

        print("DONE.")
        print("=" * 70)

        return result

    except Exception as e:

        import traceback

        traceback.print_exc()

        raise gr.Error(
            f"Background removal failed: {e}"
        )


# ============================================================
# UI
# ============================================================

css = """
#header {
    text-align: center;
    padding: 26px 0 10px;
}

#header h1 {
    font-size: 32px;
    font-weight: 700;

    background: linear-gradient(
        135deg,
        #6366f1,
        #06b6d4
    );

    -webkit-background-clip: text;
    -webkit-text-fill-color: transparent;

    margin-bottom: 5px;
}

#header p {
    color: #888;
    font-size: 14px;
}

#run-btn {
    background: linear-gradient(
        135deg,
        #6366f1,
        #06b6d4
    ) !important;

    color: white !important;
    font-weight: 600 !important;
    border: none !important;
}
"""

with gr.Blocks(
    title="Peace Network BG Remover"
) as demo:

    gr.HTML(
        """
        <div id="header">
            <h1>
                Peace Network BG Remover
            </h1>

            <p>
                High-resolution background removal
                with fine hair & beard preservation.
            </p>
        </div>
        """
    )

    with gr.Row():

        inp = gr.Image(
            type="pil",
            label="Upload photo"
        )

        out = gr.Image(
            type="pil",
            label="Transparent PNG",
            format="png"
        )

    btn = gr.Button(
        "Remove Background",
        variant="primary",
        elem_id="run-btn"
    )

    btn.click(
        remove_background,
        inputs=inp,
        outputs=out
    )

# ============================================================
# LAUNCH
# ============================================================

demo.queue().launch(
    css=css
)