import torch
from PIL import Image
import numpy as np
import base64
import io
import gradio as gr
from your_model_imports import BiRefNet  # replace with your actual model import

# Force CPU
device = torch.device("cpu")

# Load model
birefnet = BiRefNet()  # or your model class
birefnet.to(device)
birefnet.eval()  # set evaluation mode

# Helper to convert base64 to PIL
def b64_to_pil(b64_image):
    header, data = b64_image.split(",", 1)
    img_bytes = base64.b64decode(data)
    return Image.open(io.BytesIO(img_bytes)).convert("RGBA")

# Helper to convert PIL to base64
def pil_to_b64(pil_img):
    buffered = io.BytesIO()
    pil_img.save(buffered, format="PNG")
    img_str = base64.b64encode(buffered.getvalue()).decode("utf-8")
    return f"data:image/png;base64,{img_str}"

# Background removal function
def remove_bg(image_b64):
    try:
        # Convert to PIL
        img = b64_to_pil(image_b64)

        # Convert PIL to tensor
        img_tensor = torch.from_numpy(np.array(img)).permute(2,0,1).unsqueeze(0).float() / 255.0
        img_tensor = img_tensor.to(device)

        # Run model
        with torch.no_grad():
            output_tensor = birefnet(img_tensor)

        # Convert output tensor to PIL
        output_np = (output_tensor.squeeze().permute(1,2,0).numpy() * 255).astype(np.uint8)
        output_pil = Image.fromarray(output_np)

        # Convert to base64
        return pil_to_b64(output_pil)
    except Exception as e:
        return f"ERROR: {str(e)}"

# Gradio interface
iface = gr.Interface(
    fn=remove_bg,
    inputs=gr.Image(type="pil", label="Input Image"),
    outputs=gr.Image(type="auto", label="Background Removed"),
    title="Background Remover Pixels",
    description="Removes background using CPU-only model."
)

if __name__ == "__main__":
    iface.launch()