test_model.py

#!/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()