import gradio as gr
from PIL import Image
import torch
import numpy as np
from torchvision import transforms
from transformers import AutoModelForImageSegmentation
import os

print("Carregando BiRefNet_lite...")
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
use_half = device.type == 'cuda'

model = AutoModelForImageSegmentation.from_pretrained(
    "ZhengPeng7/BiRefNet_lite",
    trust_remote_code=True,
    torch_dtype=torch.float16 if use_half else torch.float32
)
model = model.to(device)
model.eval()
print(f"BiRefNet-lite carregado em {device}!")

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

def remove_background(input_image):
    try:
        if input_image is None:
            raise gr.Error("Nenhuma imagem enviada")

        img = input_image.convert('RGB')
        w, h = img.size
        print(f"Processando {w}x{h}...")

        img_tensor = transform(img).unsqueeze(0)
        if use_half:
            img_tensor = img_tensor.half()
        img_tensor = img_tensor.to(device)

        with torch.no_grad():
            preds = model(img_tensor)[-1].sigmoid().cpu().float()

        mask = preds[0].squeeze().numpy()
        mask = Image.fromarray((mask * 255).astype(np.uint8)).resize((w, h), Image.LANCZOS)

        orig = input_image.convert('RGBA')
        orig_arr = np.array(orig)
        orig_arr[:,:,3] = np.array(mask)
        result_img = Image.fromarray(orig_arr, 'RGBA')

        out_path = f"/tmp/result_{os.getpid()}.png"
        result_img.save(out_path, 'PNG')
        print("Pronto!")
        return out_path

    except gr.Error:
        raise
    except Exception as e:
        raise gr.Error(str(e))

with gr.Blocks(title="Iluminados BG Remover") as demo:
    gr.Markdown("## Iluminados BG Remover — BiRefNet-lite")
    with gr.Row():
        inp = gr.Image(label="Imagem", type="pil")
        out = gr.File(label="PNG sem fundo")
    gr.Button("Remover Fundo").click(fn=remove_background, inputs=inp, outputs=out)

demo.launch()