Update app.py

app.py

@@ -27,19 +27,24 @@ clf_2 = pipeline("image-classification", model=model_2_path)
 
 # Load additional models
 models = ["Organika/sdxl-detector", "cmckinle/sdxl-flux-detector"]
-pipe0 = pipeline("image-classification", model=models[0])
-pipe1 = pipeline("image-classification", model=models[1])
+
+# Load the third and fourth models
+feature_extractor_3 = AutoFeatureExtractor.from_pretrained(models[0])
+model_3 = AutoModelForImageClassification.from_pretrained(models[0]).to(device)
+
+feature_extractor_4 = AutoFeatureExtractor.from_pretrained(models[1])
+model_4 = AutoModelForImageClassification.from_pretrained(models[1]).to(device)
 
 # Define class names for all models
 class_names_1 = ['artificial', 'real']
 class_names_2 = ['AI Image', 'Real Image']
-class_names_3 = ['AI', 'Real']
-class_names_4 = ['AI', 'Real']
+labels_3 = ['AI', 'Real']
+labels_4 = ['AI', 'Real']
 
 def softmax(vector):
     e = np.exp(vector - np.max(vector))  # for numerical stability
     return e / e.sum()
-
+    
 @spaces.GPU(duration=10)
 def predict_image(img, confidence_threshold):
     # Ensure the image is a PIL Image
@@ -97,14 +102,19 @@ def predict_image(img, confidence_threshold):
     
     # Predict using the third model with softmax
     try:
+        inputs_3 = feature_extractor_3(img_pil, return_tensors="pt").to(device)
         with torch.no_grad():
-            outputs = pipe0(img_pil)
-            logits = outputs[0].logits if isinstance(outputs, list) else outputs.logits
-            probabilities = softmax(logits.cpu().numpy())
-            result_3 = {class_names_3[idx]: float(probabilities[idx]) for idx in range(len(class_names_3))}
+            outputs_3 = model_3(**inputs_3)
+            logits_3 = outputs_3.logits
+            probabilities_3 = softmax(logits_3.cpu().numpy()[0])
+        
+        result_3 = {
+            labels_3[0]: float(probabilities_3[0]),  # AI
+            labels_3[1]: float(probabilities_3[1])   # Real
+        }
         
         # Ensure the result dictionary contains all class names
-        for class_name in class_names_3:
+        for class_name in labels_3:
             if class_name not in result_3:
                 result_3[class_name] = 0.0
         
@@ -120,14 +130,19 @@ def predict_image(img, confidence_threshold):
     
     # Predict using the fourth model with softmax
     try:
+        inputs_4 = feature_extractor_4(img_pil, return_tensors="pt").to(device)
         with torch.no_grad():
-            outputs = pipe1(img_pil)
-            logits = outputs[0].logits if isinstance(outputs, list) else outputs.logits
-            probabilities = softmax(logits.cpu().numpy())
-            result_4 = {class_names_4[idx]: float(probabilities[idx]) for idx in range(len(class_names_4))}
+            outputs_4 = model_4(**inputs_4)
+            logits_4 = outputs_4.logits
+            probabilities_4 = softmax(logits_4.cpu().numpy()[0])
+        
+        result_4 = {
+            labels_4[0]: float(probabilities_4[0]),  # AI
+            labels_4[1]: float(probabilities_4[1])   # Real
+        }
         
         # Ensure the result dictionary contains all class names
-        for class_name in class_names_4:
+        for class_name in labels_4:
             if class_name not in result_4:
                 result_4[class_name] = 0.0
         
@@ -163,4 +178,4 @@ iface = gr.Interface(
     outputs=label,
     title="AI Generated Classification"
 )
-iface.launch()
+iface.launch()