File size: 735 Bytes
245a4de | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 |
from transformers import AutoModelForImageClassification
import sys
model_name = 'yevvonlim/DistilDIRE'
try:
print(f"Attempting to load {model_name}...")
model = AutoModelForImageClassification.from_pretrained(model_name)
print("Success!")
print("Labels:", model.config.id2label)
except Exception as e:
print(f"Error loading {model_name}: {e}")
# Check Organika as backup
model_name_2 = 'Organika/SDXL-Detector'
try:
print(f"Attempting to load {model_name_2}...")
model2 = AutoModelForImageClassification.from_pretrained(model_name_2)
print("Success Organika!")
print("Labels:", model2.config.id2label)
except Exception as e:
print(f"Error loading {model_name_2}: {e}")
|