import gradio as gr
import torch
import spaces
from PIL import Image
from torchvision import transforms
from huggingface_hub import hf_hub_download

# Fix for BiRefNet compatibility
import transformers.configuration_utils
original_getattribute = transformers.configuration_utils.PretrainedConfig.__getattribute__
def patched_getattribute(self, key):
    if key == 'is_encoder_decoder':
        return False
    return original_getattribute(self, key)
transformers.configuration_utils.PretrainedConfig.__getattribute__ = patched_getattribute

from transformers import AutoModelForImageSegmentation

# Download ToonOut weights
print("Downloading ToonOut weights...")
weights_path = hf_hub_download(
    repo_id="joelseytre/toonout",
    filename="birefnet_finetuned_toonout.pth"
)
print(f"Weights downloaded to: {weights_path}")

# Load base BiRefNet model
print("Loading BiRefNet base model...")
model = AutoModelForImageSegmentation.from_pretrained(
    "ZhengPeng7/BiRefNet",
    trust_remote_code=True
)

# Load ToonOut fine-tuned weights
print("Applying ToonOut weights...")
state_dict = torch.load(weights_path, map_location='cpu')
clean_state_dict = {}
for k, v in state_dict.items():
    if k.startswith("module._orig_mod."):
        clean_state_dict[k[len("module._orig_mod."):]] = v
    elif k.startswith("module."):
        clean_state_dict[k[len("module."):]] = v
    else:
        clean_state_dict[k] = v

model.load_state_dict(clean_state_dict)
model.eval()
print("ToonOut model loaded successfully!")

device = 'cuda' if torch.cuda.is_available() else 'cpu'
model.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. Fine-tuned BiRefNet for anime. MIT License - free for commercial use."
)

demo.launch()