Update app.py

app.py

@@ -6,8 +6,9 @@ models=[
     "Nahrawy/AIorNot",
     "RishiDarkDevil/ai-image-det-resnet152",
     "arnolfokam/ai-generated-image-detector",
+    "umm-maybe/AI-image-detector",
 ]
-pipe = pipeline("image-classification", "umm-maybe/AI-image-detector")
+#pipe = pipeline("image-classification", "umm-maybe/AI-image-detector")
 
 def image_classifier(image):
     outputs = pipe(image)
@@ -28,13 +29,14 @@ def aiornot(image,mod_choose):
     
     #feature_extractor = AutoFeatureExtractor.from_pretrained("microsoft/resnet-50")
     mod=models[int(mod_choose)]
-    feature_extractor = AutoFeatureExtractor.from_pretrained("Nahrawy/AIorNot")
-    model = AutoModelForImageClassification.from_pretrained("Nahrawy/AIorNot")
+    feature_extractor = AutoFeatureExtractor.from_pretrained(mod)
+    model = AutoModelForImageClassification.from_pretrained(mod)
     
     input = feature_extractor(image, return_tensors="pt")
     with torch.no_grad():
       outputs = model(**input)
       logits = outputs.logits
+      print (logits)
     prediction = logits.argmax(-1).item()
     label = labels[prediction] 
     return label