File size: 1,001 Bytes
6807f91 | 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 | #!/usr/bin/env python3
"""
Simple test to check if the ISNet model can be loaded
"""
import torch
from transformers import AutoModelForImageSegmentation
def test_model_loading():
"""Test if the model can be loaded"""
try:
print("Attempting to load ISNet model...")
model = AutoModelForImageSegmentation.from_pretrained(
"mateenahmed/isnet-background-remover",
trust_remote_code=True
)
print("✅ Model loaded successfully!")
# Test with a dummy input
dummy_input = torch.randn(1, 3, 1024, 1024)
with torch.no_grad():
output = model(dummy_input)
print(f"✅ Model forward pass successful! Output shape: {output.shape}")
return True
except Exception as e:
print(f"❌ Error loading model: {e}")
import traceback
traceback.print_exc()
return False
if __name__ == "__main__":
test_model_loading() |