fixing secure mode

app.py

@@ -15,6 +15,7 @@ class_names = ['Fake', 'Real']  # Fix the incorrect mapping
 
 # Load the trained model
 def load_model(model_path, device):
+    print("load_model")
     model = resnet18_custom(weights=None)
     num_ftrs = model.fc.in_features
     model.fc = nn.Linear(num_ftrs, len(class_names))  # Assuming 2 classes: Fake and Real
@@ -36,6 +37,9 @@ def load_secure_model(model):
     )
     return secure_model
 
+model = load_model('models/deepfake_detection_model.pth', 'cpu')
+secure_model = load_secure_model(model)
+
 # Image preprocessing (match with the transforms used during training)
 data_transform = transforms.Compose([
     transforms.Resize((224, 224)),
@@ -44,17 +48,7 @@ data_transform = transforms.Compose([
 
 # Prediction function
 def predict(image, mode):
-    # Device configuration
-    device = torch.device(
-        "cuda:0" if torch.cuda.is_available() else
-        "mps" if torch.backends.mps.is_available() else
-        "cpu"
-    )
-
-    print(f"Device: {device}")
-    # Load model
-    model_path = 'models/deepfake_detection_model.pth'
-    model = load_model(model_path, device)
+    device = 'cpu'
 
     # Apply transformations to the input image
     image = Image.open(image).convert('RGB')
@@ -69,11 +63,9 @@ def predict(image, mode):
             outputs = model(image)
         elif mode == "Secure":
             # Secure mode (e.g., running multiple times for higher confidence)
-            secure_model = load_secure_model(model)
             detached_input = image.detach().numpy()
-            outputs = secure_model(detached_input, fhe="simulate")
+            outputs = torch.from_numpy(secure_model.forward(detached_input, fhe="simulate"))
         
-        print(outputs)
         _, preds = torch.max(outputs, 1)
         elapsed_time = time.time() - start_time
 
@@ -96,4 +88,4 @@ iface = gr.Interface(
 )
 
 if __name__ == "__main__":
-    iface.launch(share=True)      
+    iface.launch(share=True)