Update app.py

app.py

@@ -1,3 +1,5 @@
+WEIGHT_FILE = "final_555_DeepFakeClassifier_tf_efficientnet_b7_ns_0_19" 
+
 import torch
 from torch import nn
 from torch.nn.modules.dropout import Dropout
@@ -9,7 +11,7 @@ import re
 import gradio as gr
 import os
 
-# --- 1. MODEL ARCHITECTURE ---
+# --- 1. MODEL ARCHITECTURE (Matches Notebook Exactly) ---
 encoder_params = {
     "tf_efficientnet_b7_ns": {
         "features": 2560,
@@ -32,39 +34,27 @@ class DeepFakeClassifier(nn.Module):
         x = self.fc(x)
         return x
 
-# --- 2. LOAD THE SPECIFIC WEIGHT FILE ---
-# >>> UPDATE THIS STRING IN EVERY SPACE TO MATCH THE UPLOADED FILE EXACTLY <<<
-WEIGHT_FILE = "final_555_DeepFakeClassifier_tf_efficientnet_b7_ns_0_19" 
-
 print(f"Booting API Worker: Loading {WEIGHT_FILE}...")
 device = torch.device('cpu')
 model = DeepFakeClassifier(encoder="tf_efficientnet_b7_ns").to(device)
 
-# Using weights_only=False to bypass PyTorch 2.6 security restriction
 checkpoint = torch.load(WEIGHT_FILE, map_location="cpu", weights_only=False)
 state_dict = checkpoint.get("state_dict", checkpoint)
 model.load_state_dict({re.sub("^module.", "", k): v for k, v in state_dict.items()}, strict=True)
 model.eval()
-print("Model loaded successfully. API Ready.")
 
-# --- 3. THE API INFERENCE FUNCTION ---
+# --- 3. API ENDPOINT ---
 def predict_tensor(tensor_file):
-    """
-    Receives a preprocessed .pt tensor file from the Master UI,
-    runs inference, and returns the confidence score array.
-    """
     if tensor_file is None:
         return {"error": "No tensor file received"}
-        
     try:
-        # Load the pre-processed tensor sent by the Master UI
+        # Load math tensor
         x = torch.load(tensor_file.name, map_location=device, weights_only=True)
-        
         with torch.no_grad():
             y_pred = model(x)
             y_pred = torch.sigmoid(y_pred.squeeze())
             
-            # Format output so it can be sent back via JSON
+            # Format to standard list
             if y_pred.dim() == 0:
                 bpred = [float(y_pred.cpu().numpy())]
             else:
@@ -74,14 +64,7 @@ def predict_tensor(tensor_file):
     except Exception as e:
         return {"error": str(e)}
 
-# --- 4. GRADIO API INTERFACE ---
-interface = gr.Interface(
-    fn=predict_tensor,
-    inputs=gr.File(label="Input Tensor (.pt)"),
-    outputs=gr.JSON(label="Prediction Array"),
-    title="DeepGuard Worker API",
-    description="Microservice endpoint for EfficientNet-B7 Deepfake Inference."
-)
+interface = gr.Interface(fn=predict_tensor, inputs=gr.File(label="Input Tensor (.pt)"), outputs=gr.JSON())
 
 if __name__ == "__main__":
     interface.launch()