""" Inference engines for SURPRISE. """ import logging from abc import ABC, abstractmethod from typing import Dict, List, Tuple import numpy as np from config import settings log = logging.getLogger("surprise.inference") class InferenceEngine(ABC): @abstractmethod def analyze(self, frames: np.ndarray) -> Dict: ... def _classify(self, surprise: float, straightness: float) -> Tuple[str, str, float]: s_low = settings.surprise_threshold_real s_high = settings.surprise_threshold_fake st_high = settings.straightness_threshold_real st_low = settings.straightness_threshold_fake if surprise < s_low and straightness > st_high: conf = 80 + (1 - surprise) * 15 return "real", "AUTHENTIC", min(conf, 96) if surprise > s_high or straightness < st_low: conf = 60 + min(surprise * 40, 35) return "fake", "AI-GENERATED", min(conf, 95) conf = 50 + abs(surprise - 0.4) * 30 return "uncertain", "INCONCLUSIVE", min(conf, 75) class MockLeWMInference(InferenceEngine): def __init__(self): log.info("MockLeWMInference initialized (no real model)") def analyze(self, frames: np.ndarray) -> Dict: N = len(frames) diffs = np.abs(frames[1:].astype(np.float32) - frames[:-1].astype(np.float32)) motion = diffs.mean(axis=(1, 2, 3)) / 255.0 variance = float(frames.std(axis=(1, 2, 3)).mean()) / 255.0 seed = int((motion.sum() + variance * 100) * 1000) % (2**32) rng = np.random.default_rng(seed) base = 0.20 + variance * 0.3 per_frame = base + rng.normal(0, 0.04, size=N) if len(motion) > 0: m_mean = motion.mean() m_std = motion.std() + 1e-6 for i in range(1, N): if motion[i - 1] > m_mean + 1.5 * m_std: per_frame[i] += 0.25 + rng.uniform(0, 0.15) per_frame = np.clip(per_frame, 0.05, 1.0) surprise = float(per_frame.mean()) straightness = float(np.clip(0.95 - per_frame.std() * 1.5, 0.5, 0.99)) threshold = float(np.quantile(per_frame, 0.95)) flagged = [int(i) for i in np.where(per_frame >= threshold)[0].tolist()][:8] verdict, label, conf = self._classify(surprise, straightness) flag_reason = self._reason(verdict, flagged) return { "verdict": verdict, "label": label, "confidence": round(float(conf), 1), "surprise_score": round(surprise, 3), "straightness": round(straightness, 3), "per_frame_surprise": [round(float(v), 3) for v in per_frame], "flagged_frames": flagged, "flag_reason": flag_reason, "backend": "mock", } @staticmethod def _reason(verdict, flagged): if verdict == "real": return "Smooth temporal dynamics; no physical violations detected" if verdict == "fake": if flagged: shown = flagged[:3] return f"Surprise spike at frames {', '.join(str(f) for f in shown)}" return "Inconsistent latent dynamics" return "Mixed signals" class RealLeWMInference(InferenceEngine): def __init__(self): import torch from models import TinyEncoder, TinyPredictor self.device = self._pick_device() log.info(f"RealLeWMInference loading from {settings.checkpoint_path} on {self.device}") ckpt = torch.load(settings.checkpoint_path, map_location=self.device, weights_only=False) cfg = ckpt.get("config", {}) self.embed_dim = cfg.get("embed_dim", 64) self.image_size = cfg.get("image_size", 64) self.encoder = TinyEncoder(embed_dim=self.embed_dim).to(self.device) self.predictor = TinyPredictor(embed_dim=self.embed_dim).to(self.device) self.encoder.load_state_dict(ckpt["encoder_state"]) self.predictor.load_state_dict(ckpt["predictor_state"]) self.encoder.eval() self.predictor.eval() log.info(f"Loaded checkpoint: embed_dim={self.embed_dim}, image_size={self.image_size}") def _pick_device(self): import torch if settings.device == "auto": return "cuda" if torch.cuda.is_available() else "cpu" return settings.device def analyze(self, frames: np.ndarray) -> Dict: import torch resized = self._resize_batch(frames, self.image_size) x = torch.from_numpy(resized).float() / 255.0 x = x.permute(0, 3, 1, 2).contiguous().to(self.device) with torch.no_grad(): z = self.encoder(x) z_pred = self.predictor(z[:-1]) per_frame = ((z_pred - z[1:]) ** 2).mean(dim=-1).cpu().numpy() v = (z[1:] - z[:-1]) v_norm = v / (v.norm(dim=-1, keepdim=True) + 1e-8) if v_norm.shape[0] > 1: cos_sim = (v_norm[:-1] * v_norm[1:]).sum(dim=-1) straightness = float(cos_sim.mean().cpu().item()) straightness = (straightness + 1.0) / 2.0 else: straightness = 0.5 surprise = float(per_frame.mean()) if len(per_frame) > 0: threshold = float(np.quantile(per_frame, 0.95)) flagged = [int(i) for i in np.where(per_frame >= threshold)[0].tolist()][:8] else: flagged = [] rescaled_surprise = min(surprise / 0.10, 1.0) verdict, label, conf = self._classify(rescaled_surprise, straightness) flag_reason = self._reason(verdict, flagged) return { "verdict": verdict, "label": label, "confidence": round(float(conf), 1), "surprise_score": round(float(rescaled_surprise), 3), "straightness": round(straightness, 3), "per_frame_surprise": [round(float(v), 4) for v in per_frame], "flagged_frames": flagged, "flag_reason": flag_reason, "raw_surprise": round(surprise, 5), "backend": "lewm", } @staticmethod def _resize_batch(frames, size): import cv2 out = np.zeros((len(frames), size, size, 3), dtype=np.uint8) for i, f in enumerate(frames): out[i] = cv2.resize(f, (size, size), interpolation=cv2.INTER_AREA) return out @staticmethod def _reason(verdict, flagged): if verdict == "real": return "Smooth temporal dynamics; no anomalies detected" if verdict == "fake": if flagged: shown = flagged[:3] return f"Surprise spike at frames {', '.join(str(f) for f in shown)}" return "Inconsistent temporal dynamics" return "Mixed signals" def get_inference_engine() -> InferenceEngine: backend = settings.model_backend if backend == "mock": return MockLeWMInference() if backend == "lewm": return RealLeWMInference() raise ValueError(f"Unknown MODEL_BACKEND: {backend!r}")