import gradio as gr
import torch
import spaces
from PIL import Image
from torchvision import transforms
from transformers import AutoModelForImageSegmentation

device = 'cuda' if torch.cuda.is_available() else 'cpu'
model = AutoModelForImageSegmentation.from_pretrained(
    'joelseytre/toonout', trust_remote_code=True
).eval().to(device)

@spaces.GPU
def remove_background(image, resolution="1024x1024"):
    if image is None:
        return None
    
    res = int(resolution.split('x')[0])
    if image.mode != 'RGB':
        image = image.convert('RGB')
    
    original_size = image.size
    transform = transforms.Compose([
        transforms.Resize((res, res)),
        transforms.ToTensor(),
        transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
    ])
    
    input_tensor = transform(image).unsqueeze(0).to(device)
    
    with torch.no_grad():
        preds = model(input_tensor)[-1].sigmoid().cpu()
    
    pred = preds[0].squeeze()
    mask = transforms.ToPILImage()(pred).resize(original_size)
    
    output = image.copy()
    output.putalpha(mask)
    return output

demo = gr.Interface(
    fn=remove_background,
    inputs=[
        gr.Image(type="pil", label="Upload Image"),
        gr.Dropdown(["512x512", "1024x1024", "2048x2048"], value="1024x1024", label="Resolution")
    ],
    outputs=gr.Image(type="pil", label="Result"),
    title="ToonOut - Anime Background Removal",
    description="99.5% accuracy on anime/cartoon images. MIT License - free for commercial use."
)

demo.launch()