import gradio as gr

UPSAMPLER_THEME = gr.themes.Soft(
    primary_hue=gr.themes.colors.indigo,
    secondary_hue=gr.themes.colors.purple,
    neutral_hue=gr.themes.colors.slate,
).set(
    button_primary_background_fill="linear-gradient(135deg, #6366f1, #a855f7)",
    button_primary_background_fill_hover="linear-gradient(135deg, #5457e5, #9333ea)",
    button_primary_text_color="#ffffff",
    button_primary_border_color="*primary_500",
)

UPSAMPLER_CSS = """
footer{display:none !important}
.gradio-container{max-width:1000px !important; margin:0 auto !important}
h1,h2,h3{font-family:system-ui,-apple-system,'Segoe UI',sans-serif}
"""

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

torch.set_float32_matmul_precision("high")

birefnet = AutoModelForImageSegmentation.from_pretrained(
    "ZhengPeng7/BiRefNet", trust_remote_code=True
)
birefnet.to("cuda")

transform_image = transforms.Compose(
    [
        transforms.Resize((1024, 1024)),
        transforms.ToTensor(),
        transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),
    ]
)


def get_duration(image):
    return 15


@spaces.GPU(duration=get_duration)
def remove_background(image: Image.Image) -> Image.Image:
    """Remove the background from an image and return a transparent PNG."""
    if image is None:
        raise gr.Error("Please upload an image.")
    im = image.convert("RGB")
    # ZeroGPU packs the model to fp16; match the input dtype to the model's params.
    model_dtype = next(birefnet.parameters()).dtype
    input_images = transform_image(im).unsqueeze(0).to("cuda", model_dtype)
    with torch.no_grad():
        preds = birefnet(input_images)[-1].sigmoid().float().cpu()
    pred = preds[0].squeeze()
    mask = transforms.ToPILImage()(pred).resize(im.size)
    im.putalpha(mask)
    return im


header = """<div style="max-width:760px;margin:0 auto;text-align:center;padding:20px 16px 2px;font-family:system-ui,-apple-system,'Segoe UI',sans-serif">
  <h1 style="font-size:1.7rem;font-weight:700;margin:0 0 6px;letter-spacing:-.02em">BiRefNet Background Removal</h1>
  <p style="font-size:1rem;line-height:1.5;opacity:.6;margin:0">Remove image backgrounds and get a clean transparent PNG in seconds.</p>
</div>"""

footer = """<div style="max-width:640px;margin:2rem auto .4rem;text-align:center;font-family:system-ui,-apple-system,'Segoe UI',sans-serif">
  <p style="font-size:.85rem;line-height:1.6;opacity:.5;margin:0 0 10px">BiRefNet is one of the most accurate open-source background removal models, using bilateral reference segmentation to cut out subjects with clean edges around hair, fur, and other fine detail. Upload any photo and download a transparent PNG cutout, ready for product shots, thumbnails, and design work.</p>
  <p style="font-size:.85rem;line-height:1.6;opacity:.65;margin:0">Maintained by <a href="https://upsampler.com" target="_blank" rel="noopener" style="color:#8b7cf6;font-weight:600;text-decoration:none">Upsampler</a>. Check out the <a href="https://upsampler.com/free-background-remover-no-signup" target="_blank" rel="noopener" style="color:#8b7cf6;font-weight:600;text-decoration:none">free background remover</a>, no sign-up required.</p>
</div>"""

with gr.Blocks(title="BiRefNet - Background Removal", theme=UPSAMPLER_THEME, css=UPSAMPLER_CSS) as demo:
    gr.HTML(header)
    with gr.Row(equal_height=False):
        with gr.Column():
            image_input = gr.Image(label="Upload an image", type="pil", height=360)
            run_btn = gr.Button("Remove Background", variant="primary")
        with gr.Column():
            result = gr.Image(
                label="Result (transparent WEBP)",
                type="pil",
                format="webp",
                show_share_button=False,
                height=360,
            )
    run_btn.click(remove_background, inputs=image_input, outputs=result, api_name="image")
    gr.HTML(footer)

if __name__ == "__main__":
    demo.launch(show_error=True, mcp_server=True, ssr_mode=False)