Create app.py

app.py

@@ -0,0 +1,87 @@
+import torch
+from torch import nn
+from torch.nn.modules.dropout import Dropout
+from torch.nn.modules.linear import Linear
+from torch.nn.modules.pooling import AdaptiveAvgPool2d
+from timm.models.efficientnet import tf_efficientnet_b7_ns
+from functools import partial
+import re
+import gradio as gr
+import os
+
+# --- 1. MODEL ARCHITECTURE ---
+encoder_params = {
+    "tf_efficientnet_b7_ns": {
+        "features": 2560,
+        "init_op": partial(tf_efficientnet_b7_ns, pretrained=False, drop_path_rate=0.2)
+    }
+}
+
+class DeepFakeClassifier(nn.Module):
+    def __init__(self, encoder="tf_efficientnet_b7_ns", dropout_rate=0.0) -> None:
+        super().__init__()
+        self.encoder = encoder_params[encoder]["init_op"]()
+        self.avg_pool = AdaptiveAvgPool2d((1, 1))
+        self.dropout = Dropout(dropout_rate)
+        self.fc = Linear(encoder_params[encoder]["features"], 1)
+
+    def forward(self, x):
+        x = self.encoder.forward_features(x)
+        x = self.avg_pool(x).flatten(1)
+        x = self.dropout(x)
+        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_777_DeepFakeClassifier_tf_efficientnet_b7_ns_0_29" 
+
+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 ---
+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
+        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
+            if y_pred.dim() == 0:
+                bpred = [float(y_pred.cpu().numpy())]
+            else:
+                bpred = y_pred.cpu().numpy().tolist()
+                
+        return {"predictions": bpred}
+    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."
+)
+
+if __name__ == "__main__":
+    interface.launch()