Spaces:
Running
Running
File size: 2,817 Bytes
5125056 5ad7715 5125056 5ad7715 5125056 5ad7715 5125056 5ad7715 5125056 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 | 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() |