Update app.py

app.py

@@ -3,12 +3,52 @@ 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
 
-device = 'cuda' if torch.cuda.is_available() else 'cpu'
+# 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(
-    'joelseytre/toonout', trust_remote_code=True
-).eval().to(device)
+    "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"):
@@ -46,7 +86,7 @@ demo = gr.Interface(
     ],
     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."
+    description="99.5% accuracy on anime/cartoon images. Fine-tuned BiRefNet for anime. MIT License - free for commercial use."
 )
 
 demo.launch()