Update app.py

app.py

@@ -1,95 +1,394 @@
-WEIGHT_FILE = "final_111_DeepFakeClassifier_tf_efficientnet_b7_ns_0_36" 
+WEIGHT_FILE = os.environ.get(
+    "WEIGHT_FILE",
+    "final_111_DeepFakeClassifier_tf_efficientnet_b7_ns_0_36",  # safe default
+)
 
+"""
+================================================================================
+  VERIDEX — DeepFake Worker Space  (Generic Template)
+  ─────────────────────────────────────────────────────
+  DEPLOY INSTRUCTIONS  — zero code changes between workers
+  ──────────────────────────────────────────────────────────
+  1. Commit this IDENTICAL app.py to all 7 Worker Spaces.
+  2. Upload each worker's .pt weight file to its Space's files tab.
+  3. In each Space → Settings → Variables, set:
+
+       WEIGHT_FILE  = final_111_DeepFakeClassifier_tf_efficientnet_b7_ns_0_36
+       MODEL_CLASS  = base        # or srm / gwap  (optional, default: base)
+
+  4. That's it. No code edits required.
+
+  API CONTRACT (called by the Master UI)
+  ───────────────────────────────────────
+  Input  : a .npy file (uint8, shape [N, H, W, 3], HWC, 380×380)
+  Output : JSON  { "predictions": [float, ...], "n_frames": int }
+            OR   { "error": "...", "predictions": null }
+
+  GRADIO VERSION NOTE
+  ────────────────────
+  HF Spaces force-installs gradio==6.x at build time regardless of what
+  requirements.txt pins. This file targets Gradio 6:
+    • gr.File input passes a tempfile.SpooledTemporaryFile-backed object
+      with a .name attribute in Gradio 6 (not a plain string or dict).
+    • allow_flagging is removed (deprecated in Gradio 6; raises a warning
+      that can abort startup on strict HF runtime configs).
+================================================================================
+"""
+
+import os
+import io
+import re
+import traceback
+import logging
+
+import numpy as np
 import torch
-from torch import nn
+import torch.nn as 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 torchvision.transforms import Normalize
 from functools import partial
-import re
+
 import gradio as gr
-import numpy as np
-from torchvision.transforms import Normalize
 
-# --- 1. MODEL ARCHITECTURE ---
-encoder_params = {
-    "tf_efficientnet_b7.ns_jft_in1k": { 
-        "features": 2560,
-        "init_op": partial(tf_efficientnet_b7_ns, pretrained=False, drop_path_rate=0.2)
-    }
-}
+# ── timm / efficientnet ───────────────────────────────────────────────────────
+try:
+    from timm.models.efficientnet import tf_efficientnet_b7_ns
+except ImportError:
+    # timm ≥ 0.9 moved the alias; fall back gracefully
+    import timm
+    tf_efficientnet_b7_ns = partial(timm.create_model, "tf_efficientnet_b7.ns_jft_in1k")
+
+
+logging.basicConfig(level=logging.INFO, format="%(asctime)s [WORKER] %(levelname)s %(message)s")
+logger = logging.getLogger(__name__)
+
+# ══════════════════════════════════════════════════════════════════════════════
+#  ❶  ALL CONFIG IS VIA ENV VARS — set these in each Space's Settings → Variables
+#     WEIGHT_FILE : filename of the .pt checkpoint (no extension required)
+#     MODEL_CLASS : "base" | "srm" | "gwap"  (default: base)
+#     MINI_BATCH  : frames per forward pass   (default: 8)
+#     WEIGHTS_DIR : directory containing the .pt file (default: repo root ".")
+# ══════════════════════════════════════════════════════════════════════════════
+
+MODEL_CLASS = os.environ.get("MODEL_CLASS", "base")   # "base" | "srm" | "gwap"
+MINI_BATCH  = int(os.environ.get("MINI_BATCH", "8"))   # frames per forward pass
+WEIGHTS_DIR = os.environ.get("WEIGHTS_DIR", ".")        # dir that contains the .pt
+# ══════════════════════════════════════════════════════════════════════════════
+
+# ── ImageNet normalisation ────────────────────────────────────────────────────
+IMAGENET_MEAN = [0.485, 0.456, 0.406]
+IMAGENET_STD  = [0.229, 0.224, 0.225]
+normalize_fn  = Normalize(mean=IMAGENET_MEAN, std=IMAGENET_STD)
+
+# ── EfficientNet-B7 feature size ──────────────────────────────────────────────
+ENCODER_FEATURES = 2560
+
+# ─────────────────────────────────────────────────────────────────────────────
+#  Model definitions  (identical to deepfake_det.py so checkpoints load clean)
+# ─────────────────────────────────────────────────────────────────────────────
+
+def _make_encoder():
+    return tf_efficientnet_b7_ns(pretrained=False, drop_path_rate=0.2)
+
+
+def _setup_srm_weights(input_channels: int = 3) -> torch.Tensor:
+    srm_kernel = torch.from_numpy(np.array([
+        [[0.,0.,0.,0.,0.],[0.,0.,0.,0.,0.],[0.,1.,-2.,1.,0.],[0.,0.,0.,0.,0.],[0.,0.,0.,0.,0.]],
+        [[0.,0.,0.,0.,0.],[0.,-1.,2.,-1.,0.],[0.,2.,-4.,2.,0.],[0.,-1.,2.,-1.,0.],[0.,0.,0.,0.,0.]],
+        [[-1.,2.,-2.,2.,-1.],[2.,-6.,8.,-6.,2.],[-2.,8.,-12.,8.,-2.],[2.,-6.,8.,-6.,2.],[-1.,2.,-2.,2.,-1.]],
+    ])).float()
+    srm_kernel[0] /= 2
+    srm_kernel[1] /= 4
+    srm_kernel[2] /= 12
+    return srm_kernel.view(3, 1, 5, 5).repeat(1, input_channels, 1, 1)
+
+
+def _setup_srm_layer(input_channels: int = 3) -> nn.Module:
+    weights = _setup_srm_weights(input_channels)
+    conv = nn.Conv2d(input_channels, 3, kernel_size=5, stride=1, padding=2, bias=False)
+    with torch.no_grad():
+        conv.weight = nn.Parameter(weights, requires_grad=False)
+    return conv
+
 
 class DeepFakeClassifier(nn.Module):
-    def __init__(self, encoder="tf_efficientnet_b7.ns_jft_in1k", dropout_rate=0.0) -> None:
+    def __init__(self, dropout_rate=0.0):
         super().__init__()
-        self.encoder = encoder_params[encoder]["init_op"]()
+        self.encoder  = _make_encoder()
         self.avg_pool = AdaptiveAvgPool2d((1, 1))
-        self.dropout = Dropout(dropout_rate)
-        self.fc = Linear(encoder_params[encoder]["features"], 1)
+        self.dropout  = Dropout(dropout_rate)
+        self.fc       = Linear(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
-
-print(f"Booting API Worker: Loading {WEIGHT_FILE}...")
-device = torch.device('cpu')
-model = DeepFakeClassifier(encoder="tf_efficientnet_b7.ns_jft_in1k").to(device)
-
-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()
-
-# --- 3. PREPROCESSING UTILS ---
-mean = [0.485, 0.456, 0.406]
-std = [0.229, 0.224, 0.225]
-normalize_transform = Normalize(mean, std)
-
-# --- 4. API ENDPOINT ---
-def predict_tensor(tensor_file):
-    if tensor_file is None:
-        return {"error": "No file received"}
+        return self.fc(x)
+
+
+class DeepFakeClassifierSRM(nn.Module):
+    def __init__(self, dropout_rate=0.5):
+        super().__init__()
+        self.encoder  = _make_encoder()
+        self.avg_pool = AdaptiveAvgPool2d((1, 1))
+        self.srm_conv = _setup_srm_layer(3)
+        self.dropout  = Dropout(dropout_rate)
+        self.fc       = Linear(ENCODER_FEATURES, 1)
+
+    def forward(self, x):
+        noise = self.srm_conv(x)
+        x = self.encoder.forward_features(noise)
+        x = self.avg_pool(x).flatten(1)
+        x = self.dropout(x)
+        return self.fc(x)
+
+
+class _GWAP(nn.Module):
+    def __init__(self, features: int):
+        super().__init__()
+        self.conv = nn.Conv2d(features, 1, kernel_size=1, bias=True)
+
+    def forward(self, x):
+        w = self.conv(x).sigmoid().exp()
+        w = w / w.sum(dim=[2, 3], keepdim=True)
+        return (w * x).sum(dim=[2, 3], keepdim=False)
+
+
+class DeepFakeClassifierGWAP(nn.Module):
+    def __init__(self, dropout_rate=0.5):
+        super().__init__()
+        self.encoder  = _make_encoder()
+        self.avg_pool = _GWAP(ENCODER_FEATURES)
+        self.dropout  = Dropout(dropout_rate)
+        self.fc       = Linear(ENCODER_FEATURES, 1)
+
+    def forward(self, x):
+        x = self.encoder.forward_features(x)
+        x = self.avg_pool(x)
+        x = self.dropout(x)
+        return self.fc(x)
+
+
+_MODEL_MAP = {
+    "base": DeepFakeClassifier,
+    "srm":  DeepFakeClassifierSRM,
+    "gwap": DeepFakeClassifierGWAP,
+}
+
+
+# ─────────────────────────────────────────────────────────────────────────────
+#  Model loading  (runs once at startup)
+# ────────────────────────────────��────────────────────────────────────────────
+
+def load_model() -> nn.Module:
+    device  = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    cls     = _MODEL_MAP.get(MODEL_CLASS, DeepFakeClassifier)
+    model   = cls().to(device)
+
+    weight_path = os.path.join(WEIGHTS_DIR, WEIGHT_FILE)
+    # Allow common extensions in case the file was renamed
+    if not os.path.exists(weight_path):
+        for ext in (".pt", ".pth", ".bin"):
+            if os.path.exists(weight_path + ext):
+                weight_path = weight_path + ext
+                break
+
+    if not os.path.exists(weight_path):
+        raise FileNotFoundError(
+            f"Weight file not found: {weight_path}\n"
+            f"Files present in '{WEIGHTS_DIR}': {os.listdir(WEIGHTS_DIR)}"
+        )
+
+    logger.info(f"Loading weights from: {weight_path}")
+
+    # PyTorch 2.6+ requires weights_only=False for pickled checkpoints; also
+    # use map_location='cpu' so the model loads on any machine regardless of
+    # how it was saved.
+    checkpoint  = torch.load(weight_path, map_location="cpu", weights_only=False)
+    state_dict  = checkpoint.get("state_dict", checkpoint)
+
+    # Strip "module." prefix added by DataParallel / DistributedDataParallel
+    cleaned = {re.sub(r"^module\.", "", k): v for k, v in state_dict.items()}
+    model.load_state_dict(cleaned, strict=True)
+    model.eval()
+
+    # FP16 halves VRAM; safe on both CUDA and CPU
+    model = model.half()
+
+    logger.info(f"Model ready — class={cls.__name__}, device={device}, fp16=True")
+    return model, device
+
+
+try:
+    MODEL, DEVICE = load_model()
+    LOAD_ERROR    = None
+except Exception as exc:
+    MODEL      = None
+    DEVICE     = None
+    LOAD_ERROR = traceback.format_exc()
+    logger.error(f"MODEL LOAD FAILED:\n{LOAD_ERROR}")
+
+
+# ─────────────────────────────────────────────────────────────────────────────
+#  Inference helper
+# ─────────────────────────────────────────────────────────────────────────────
+
+def _preprocess_npy(npy_input) -> torch.Tensor:
+    """
+    Load a uint8 HWC .npy face-batch, convert to normalised float CHW tensor.
+
+    Gradio version compatibility matrix
+    ─────────────────────────────────────
+    Gradio 4 : passes a plain string filepath  "/tmp/gradio/.../faces.npy"
+    Gradio 4 : may wrap in dict {"path": "...", "orig_name": "..."}
+    Gradio 6 : passes a tempfile.SpooledTemporaryFile (file-like with .name)
+               OR a gradio.FileData dataclass with a .path attribute
+
+    We resolve all four forms to a final file path or file-like object
+    that np.load() can consume.
+    """
+    npy_path = None   # will hold a string path if resolvable
+    file_obj = None   # will hold a file-like if path is unavailable
+
+    # ── Form 1: plain string ──────────────────────────────────────────────────
+    if isinstance(npy_input, str):
+        npy_path = npy_input
+
+    # ── Form 2: Gradio 4 dict {"path": ..., "orig_name": ...} ────────────────
+    elif isinstance(npy_input, dict):
+        npy_path = (
+            npy_input.get("path")
+            or npy_input.get("name")
+            or next(iter(npy_input.values()), None)
+        )
+
+    # ── Form 3: Gradio 6 dataclass (has .path attribute) ─────────────────────
+    elif hasattr(npy_input, "path"):
+        npy_path = npy_input.path
+
+    # ── Form 4: file-like object (SpooledTemporaryFile, BytesIO, etc.) ────────
+    elif hasattr(npy_input, "read"):
+        # Try to get the backing file path first (avoids reading into RAM twice)
+        backing = getattr(npy_input, "name", None)
+        if backing and isinstance(backing, str) and os.path.exists(backing):
+            npy_path = backing
+        else:
+            file_obj = npy_input
+
+    else:
+        raise TypeError(
+            f"Cannot resolve npy input of type {type(npy_input)}: {npy_input!r}"
+        )
+
+    # ── Load the array ─────────────────────────────────────────────────────────
+    def _load(src):
+        try:
+            return np.load(src, allow_pickle=False)
+        except ValueError:
+            # Legacy pickled .npy — seek back to start if file-like
+            if hasattr(src, "seek"):
+                src.seek(0)
+            return np.load(src, allow_pickle=True)
+
+    if npy_path is not None:
+        if not os.path.exists(npy_path):
+            raise FileNotFoundError(f"NPY payload not found at: {npy_path}")
+        faces_uint8 = _load(npy_path)
+    else:
+        faces_uint8 = _load(file_obj)
+
+    # ── Validate shape ─────────────────────────────────────────────────────────
+    if faces_uint8.ndim != 4 or faces_uint8.shape[3] != 3:
+        raise ValueError(
+            f"Expected uint8 array shape (N, H, W, 3), got {faces_uint8.shape}"
+        )
+
+    # Convert: uint8 HWC  →  float32 CHW  →  normalised
+    tensor = torch.from_numpy(faces_uint8).float()   # [N, H, W, 3]
+    tensor = tensor.permute(0, 3, 1, 2)              # [N, 3, H, W]
+    # Normalise each frame in-place
+    for i in range(tensor.shape[0]):
+        tensor[i] = normalize_fn(tensor[i] / 255.0)
+
+    return tensor  # float32, shape [N, 3, H, W]
+
+
+def run_inference(tensor: torch.Tensor) -> list:
+    """
+    Forward-pass the pre-processed face tensor through the model in
+    mini-batches of size MINI_BATCH to avoid OOM on 16 GB RAM spaces.
+    Returns a flat Python list of per-frame fake-probabilities [0, 1].
+    """
+    predictions = []
+    n = tensor.shape[0]
+
+    with torch.no_grad():
+        for start in range(0, n, MINI_BATCH):
+            batch = tensor[start : start + MINI_BATCH]
+            batch = batch.to(DEVICE).half()     # fp16 matches model dtype
+
+            logits = MODEL(batch)               # [B, 1]
+            probs  = torch.sigmoid(logits.squeeze(-1))  # [B]
+            predictions.extend(probs.cpu().float().tolist())
+
+    return predictions
+
+
+# ─────────────────────────────────────────────────────────────────────────────
+#  Gradio endpoint  (headless — no UI blocks, purely an API)
+# ─────────────────────────────────────────────────────────────────────────────
+
+def predict(npy_file) -> dict:
+    """
+    Gradio API endpoint.
+
+    Parameters
+    ----------
+    npy_file : str | dict
+        Filepath (or Gradio file dict) pointing to the .npy face batch.
+
+    Returns
+    -------
+    dict with keys:
+        predictions : list[float] | None
+        n_frames    : int
+        error       : str | None
+    """
+    if MODEL is None:
+        msg = f"Model failed to load at startup:\n{LOAD_ERROR}"
+        logger.error(msg)
+        return {"predictions": None, "n_frames": 0, "error": msg}
+
     try:
-        # Handle Gradio 4 filepath string formatting safely
-        file_path = tensor_file if isinstance(tensor_file, str) else tensor_file.name
-        
-        # 1. Load the raw numpy array (Bypasses PyTorch security & saves bandwidth)
-        x_np = np.load(file_path)
-        x = torch.tensor(x_np).float()
-        
-        # 2. Permute NHWC to NCHW
-        x = x.permute((0, 3, 1, 2))
-        
-        # 3. Normalize locally on the worker
-        for i in range(len(x)):
-            x[i] = normalize_transform(x[i] / 255.)
-            
-        # 4. MINI-BATCHING: Process 8 frames at a time to prevent 16GB RAM Crash
-        mini_batch_size = 8
-        all_preds = []
-        
-        with torch.no_grad():
-            for i in range(0, len(x), mini_batch_size):
-                chunk = x[i:i+mini_batch_size]
-                # Cast to float16 ONLY if running on local GPU, else stay float32 for HF CPU
-                chunk = chunk.half() if torch.cuda.is_available() else chunk
-                y_chunk = model(chunk)
-                y_chunk = torch.sigmoid(y_chunk.squeeze())
-                
-                if y_chunk.dim() == 0:
-                    all_preds.append(float(y_chunk.cpu().numpy()))
-                else:
-                    all_preds.extend(y_chunk.cpu().numpy().tolist())
-                    
-        return {"predictions": all_preds}
-    except Exception as e:
-        return {"error": str(e)}
-
-interface = gr.Interface(fn=predict_tensor, inputs=gr.File(label="Input Tensor (.npy)"), outputs=gr.JSON())
+        tensor      = _preprocess_npy(npy_file)
+        n_frames    = tensor.shape[0]
+        predictions = run_inference(tensor)
+        logger.info(f"Inference OK — frames={n_frames}, mean_pred={np.mean(predictions):.4f}")
+        return {"predictions": predictions, "n_frames": n_frames, "error": None}
+
+    except Exception:
+        err = traceback.format_exc()
+        logger.error(f"Inference failed:\n{err}")
+        return {"predictions": None, "n_frames": 0, "error": err}
+
+
+# ─────────────────────────────────────────────────────────────────────��───────
+#  Launch
+# ─────────────────────────────────────────────────────────────────────────────
+
+demo = gr.Interface(
+    fn=predict,
+    inputs=gr.File(label="Face batch (.npy)", file_types=[".npy"]),
+    outputs=gr.JSON(label="Worker prediction"),
+    title=f"VERIDEX Worker — {WEIGHT_FILE}",
+    description=(
+        "Headless inference worker. "
+        "POST a uint8 .npy face-batch; receive per-frame fake probabilities."
+    ),
+    # allow_flagging removed: deprecated in Gradio 5, gone in Gradio 6
+)
 
 if __name__ == "__main__":
-    interface.launch()
+    demo.launch(server_name="0.0.0.0", server_port=7860, show_error=True)