Update app.py

app.py

@@ -7,6 +7,7 @@ from PIL import Image
 import pandas as pd
 import warnings
 import math
+import numpy as np
 
 # Suppress warnings
 warnings.filterwarnings("ignore", category=UserWarning, message="Using a slow image processor as `use_fast` is unset")
@@ -26,8 +27,8 @@ 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], device=device)
-pipe1 = pipeline("image-classification", model=models[1], device=device)
+pipe0 = pipeline("image-classification", model=models[0])
+pipe1 = pipeline("image-classification", model=models[1])
 
 # Define class names for all models
 class_names_1 = ['artificial', 'real']
@@ -36,7 +37,7 @@ class_names_3 = ['AI', 'Real']
 class_names_4 = ['AI', 'Real']
 
 def softmax(vector):
-    e = math.exp(vector - vector.max())  # for numerical stability
+    e = np.exp(vector - np.max(vector))  # for numerical stability
     return e / e.sum()
 
 @spaces.GPU(duration=10)
@@ -94,12 +95,13 @@ def predict_image(img, confidence_threshold):
     except Exception as e:
         label_2 = f"Error: {str(e)}"
     
-    # Predict using the third model
+    # Predict using the third model with softmax
     try:
-        prediction_3 = pipe0(img_pil)
-        result_3 = {}
-        for idx, result in enumerate(prediction_3):
-            result_3[class_names_3[idx]] = float(result['score'])
+        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))}
         
         # Ensure the result dictionary contains all class names
         for class_name in class_names_3:
@@ -116,12 +118,13 @@ def predict_image(img, confidence_threshold):
     except Exception as e:
         label_3 = f"Error: {str(e)}"
     
-    # Predict using the fourth model
+    # Predict using the fourth model with softmax
     try:
-        prediction_4 = pipe1(img_pil)
-        result_4 = {}
-        for idx, result in enumerate(prediction_4):
-            result_4[class_names_4[idx]] = float(result['score'])
+        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))}
         
         # Ensure the result dictionary contains all class names
         for class_name in class_names_4: