interface/backend/engine.py

"""Inference engine: build SynthScope models from config, load checkpoints, and predict.

Models are built lazily on first use (backbones are large and slow to download/instantiate)
and kept in a small LRU so a shared GPU is not exhausted. All forward passes happen under
``torch.no_grad()`` with the model in fp32/eval mode.
"""
import base64
import gc
import io
import threading
import time
from collections import OrderedDict

import numpy as np
import torch
import torchvision.transforms as T
from PIL import Image

from . import config
from .registry import ModelSpec, get_spec, input_note
from . import media


def _png_b64(img: Image.Image) -> str:
    buf = io.BytesIO()
    img.save(buf, format="PNG")
    return "data:image/png;base64," + base64.b64encode(buf.getvalue()).decode("ascii")


def _saliency_rgb(sal: torch.Tensor):
    """2-D saliency -> (jet RGB uint8 HxWx3, normalized intensity HxW in [0,1]).

    Robust-normalizes by the 1st/99th percentile so a few hot pixels don't wash it out.
    """
    s = sal.detach().float().cpu().numpy()
    lo, hi = np.percentile(s, 1.0), np.percentile(s, 99.0)
    s = np.clip((s - lo) / (hi - lo + 1e-8), 0.0, 1.0)
    r = np.clip(1.5 - np.abs(4 * s - 3), 0, 1)   # jet colormap
    g = np.clip(1.5 - np.abs(4 * s - 2), 0, 1)
    b = np.clip(1.5 - np.abs(4 * s - 1), 0, 1)
    return (np.stack([r, g, b], axis=-1) * 255).astype(np.uint8), s


def _saliency_to_png_b64(sal: torch.Tensor, out_size: int = 288) -> str:
    """Render a saliency tensor as a base64 'jet' heatmap PNG."""
    rgb, _ = _saliency_rgb(sal)
    return _png_b64(Image.fromarray(rgb, "RGB").resize((out_size, out_size), Image.BILINEAR))


def _overlay_to_png_b64(sal: torch.Tensor, base: Image.Image, out_size: int = 288,
                        alpha: float = 0.6) -> str:
    """Blend the jet heatmap over the base image; per-pixel opacity = saliency intensity so cool
    regions show the image and hot regions show the artifact heat."""
    rgb, s = _saliency_rgb(sal)
    heat = Image.fromarray(rgb, "RGB").resize((out_size, out_size), Image.BILINEAR)
    base = base.convert("RGB").resize((out_size, out_size), Image.BILINEAR)
    mask = Image.fromarray((np.clip(s, 0, 1) * alpha * 255).astype(np.uint8), "L").resize(
        (out_size, out_size), Image.BILINEAR)
    return _png_b64(Image.composite(heat, base, mask))


def _geometric_base(spec, pil: Image.Image) -> Image.Image:
    """The model's input region in RGB (apply only the Resize/CenterCrop steps of the eval
    transform, skipping ToTensor/Normalize) so the overlay aligns with the saliency. Custom
    transforms (AIDE/SPAI native-res) fall back to the native image."""
    import torchvision.transforms as T
    base = pil.convert("RGB")
    for t in getattr(spec.transform, "transforms", []) or []:
        if isinstance(t, (T.Resize, T.CenterCrop)):
            base = t(base)
    return base


# ---------------------------------------------------------------------------
# Model construction (mirrors trainer.SynthScopeModelTrainer._build_model dispatch, with lazy
# imports so unused heavy dependencies are never loaded).
# ---------------------------------------------------------------------------
def _build_model(cfg):
    model_type = cfg.ModelConfig.get("model_type", "dino").lower()
    if model_type == "safe":
        from models.image.safe.safe import SynthScopeSAFEModel
        return SynthScopeSAFEModel(cfg)
    if model_type == "fat_former":
        from models.image.fat_former.fat_former import SynthScopeFatFormerModel
        return SynthScopeFatFormerModel(cfg)
    if model_type == "co_spy":
        from models.image.co_spy.co_spy import SynthScopeCoSpyModel
        return SynthScopeCoSpyModel(cfg)
    if model_type == "aide":
        from models.image.aide.aide import SynthScopeAIDEModel
        return SynthScopeAIDEModel(cfg)
    if model_type == "spai":
        from models.image.spai.spai import SPAI
        return SPAI(cfg)
    if model_type == "vib_net":
        from models.image.vib_net.vib_net import SynthScopeVIBNetModel
        return SynthScopeVIBNetModel(cfg)
    if model_type == "npr":
        from models.image.npr.npr import NPR
        return NPR(cfg)
    if model_type == "pgc":
        from models.image.pgc.pgc import SynthScopePGCModel
        return SynthScopePGCModel(cfg)
    if model_type == "cure":
        from models.image.cure.cure import SynthScopeCuReModel
        return SynthScopeCuReModel(cfg)
    if model_type == "stall":
        from models.video.stall.stall import SynthScopeSTALLModel
        return SynthScopeSTALLModel(cfg)
    if model_type == "d3":
        from models.video.d3.d3 import SynthScopeD3Model
        return SynthScopeD3Model(cfg)
    if model_type == "restrav":
        from models.video.restrav.restrav import SynthScopeReStraVModel
        return SynthScopeReStraVModel(cfg)
    if model_type == "cnn":
        from models.image.cnn_detection.cnn_detection import SynthScopeCNNModel
        return SynthScopeCNNModel(cfg)
    if model_type == "adapted_dino":
        from models.image.adapted_dino.adapted_dino import SynthScopeAdaptedDINOModel
        return SynthScopeAdaptedDINOModel(cfg)
    if model_type == "dinov2":
        from models.image.dino_detection.dino_detection import SynthScopeDINOModel
        return SynthScopeDINOModel(cfg)
    if model_type == "dinov3":
        from models.image.dino_detection.dino_detection import SynthScopeDINOv3Model
        return SynthScopeDINOv3Model(cfg)
    if model_type == "clip" or cfg.ModelConfig.backbone_name.startswith("CLIP:"):
        from models.image.clip_detection.clip_detection import SynthScopeCLIPModel
        return SynthScopeCLIPModel(cfg)
    raise ValueError(f"Invalid model_type: {model_type}")


def _load_checkpoint(model, spec: ModelSpec):
    """Load weights for ``spec`` into ``model`` (already constructed). Returns a status string."""
    path = config.CHECKPOINTS_DIR / spec.checkpoint
    if not path.exists():
        raise FileNotFoundError(
            f"Checkpoint not found: {path}. Train or place the checkpoint, then retry."
        )
    raw = torch.load(str(path), map_location="cpu", weights_only=False)
    sd = raw["model"] if isinstance(raw, dict) and "model" in raw else raw
    if spec.remap is not None:
        sd = spec.remap(sd)
    missing, unexpected = model.load_state_dict(sd, strict=False)
    # Buffers that are legitimately reconstructed in __init__ / not in official dumps.
    benign = {"clip_mean", "clip_std", "imnet_mean", "imnet_std"}
    real_missing = [
        k for k in missing
        if not k.startswith("criterion") and not k.startswith("dwt_filter.")
        and k.split(".")[-1] not in benign
    ]
    if spec.ckpt_kind.startswith("official"):
        # Official-weights loads are verified to be exact; flag any genuine mismatch loudly.
        if real_missing or unexpected:
            raise RuntimeError(
                f"{spec.key}: checkpoint mismatch missing={real_missing[:8]} "
                f"unexpected={unexpected[:8]}"
            )
    return (
        f"loaded {len(sd)} tensors from {path.name} "
        f"(missing={len(real_missing)}, unexpected={len(unexpected)})"
    )


class Engine:
    """Lazily builds, caches, and runs the detectors."""

    def __init__(self, device: str | None = None):
        # ``device`` lets a caller (e.g. the ``synthscope`` Python API) pin a specific device;
        # the web backend constructs it with no argument and auto-detects via config.device().
        self.device = device or config.device()
        self._cache: "OrderedDict[str, torch.nn.Module]" = OrderedDict()
        self._status: dict[str, str] = {}
        self._lock = threading.Lock()

    # -- availability -------------------------------------------------------
    def checkpoint_path(self, spec: ModelSpec):
        if spec.checkpoint is None:
            return None
        return config.CHECKPOINTS_DIR / spec.checkpoint

    def is_available(self, spec: ModelSpec) -> bool:
        if spec.key == "stall":  # training-free, but needs the gated DINOv3 ViT-L/16 weights
            import os
            from .registry import DINOV3_WEIGHTS
            return os.path.exists(os.environ.get("DINO_V3_WEIGHTS") or DINOV3_WEIGHTS)
        if spec.ckpt_kind == "none":  # training-free (D3)
            return True
        p = self.checkpoint_path(spec)
        return bool(p and p.exists())

    # -- model lifecycle ----------------------------------------------------
    def release_gpu(self) -> None:
        """Drop every resident detector (they reload lazily on the next request) so a heavy local image generator,
        which runs in its own process and wants most of the GPU, does not collide with them."""
        with self._lock:
            self._cache.clear()
        gc.collect()
        if self.device == "cuda":
            torch.cuda.empty_cache()

    def _get_model(self, spec: ModelSpec) -> torch.nn.Module:
        with self._lock:
            if spec.key in self._cache:
                self._cache.move_to_end(spec.key)
                return self._cache[spec.key]

            # Evict LRU entries if at capacity.
            while len(self._cache) >= config.MAX_RESIDENT_MODELS:
                old_key, old_model = self._cache.popitem(last=False)
                del old_model
                if self.device == "cuda":
                    torch.cuda.empty_cache()

            from config_utils import compose_config

            cfg = compose_config(spec.overrides)
            model = _build_model(cfg)
            if spec.ckpt_kind != "none":
                self._status[spec.key] = _load_checkpoint(model, spec)
            else:
                self._status[spec.key] = "training-free (no checkpoint)"
            model = model.to(self.device).float().eval()
            self._cache[spec.key] = model
            return model

    # -- prediction helpers -------------------------------------------------
    @torch.no_grad()
    def _fake_prob(self, model, spec: ModelSpec, x: torch.Tensor) -> torch.Tensor:
        """Return P(fake) per sample for an image batch ``x`` [B,3,H,W]."""
        kind = spec.predict_kind
        if kind == "ce":
            return model(x).softmax(dim=1)[:, 1]
        if kind == "bce":
            return torch.sigmoid(model(x).reshape(x.shape[0]))
        if kind == "bce_npr":
            logits = model.forward(model.NPR_transform(x)).reshape(x.shape[0])
            return torch.sigmoid(logits)
        if kind == "vib":
            out = model(x)
            logit = out[1] if isinstance(out, (tuple, list)) else out
            return torch.sigmoid(logit.reshape(x.shape[0]))
        raise ValueError(f"Unexpected image predict_kind: {kind}")

    def _d3_calibration(self):
        """Load (and cache) the D3 score->probability calibration if present."""
        if not hasattr(self, "_d3_cal"):
            import json
            # A file in the checkpoint folder overrides the one shipped with the model (fitted on GenVideo-Val).
            p = config.CHECKPOINTS_DIR / "d3_calibration.json"
            if not p.exists():
                p = config.REPO_ROOT / "models" / "video" / "d3" / "d3_calibration.json"
            self._d3_cal = json.loads(p.read_text()) if p.exists() else None
        return self._d3_cal

    @torch.no_grad()
    def _video_fake_prob(self, model, spec: ModelSpec, clip: torch.Tensor) -> torch.Tensor:
        """Return P(fake) for a single normalized clip tensor [T,3,H,W]."""
        frames = clip.unsqueeze(0).to(self.device)  # [1,T,3,H,W]
        if spec.key == "d3":
            cal = self._d3_calibration()
            if cal is not None:
                s = model.score(frames)  # [1] raw volatility score
                z = (s - cal["threshold"]) / cal["scale"]
                if not cal["higher_is_fake"]:
                    z = -z
                return torch.sigmoid(z)
        out = model(frames)  # [1,2]
        if spec.key == "stall":
            # STALL returns [real_prob, fake_prob] (already calibrated percentiles); use directly.
            return out[:, 1]
        return out.softmax(dim=1)[:, 1]

    # -- public API ---------------------------------------------------------
    @torch.no_grad()
    def _interpret_map(self, model, spec: ModelSpec, x: torch.Tensor, pil=None):
        """Return ``{"map": heatmap_b64, "overlay": overlay_b64}`` for the detector's artifact map,
        or None. ``map`` is the bare jet heatmap; ``overlay`` blends it on the model-input image.

        Best-effort + isolated: any failure returns None and never affects the prediction.
        ``x`` is the already-transformed [1,3,H,W] input; ``pil`` is the native image (some maps,
        e.g. AIDE's native-res DCT energy, are computed from the original resolution).
        """
        kind = getattr(spec, "interpret", "") or ""
        if not kind:
            return None
        try:
            if kind == "safe_hf":          # SAFE: DWT HH high-frequency band magnitude
                hh = model._preprocess_high_frequency(x)        # [1,3,H,W]
                sal = hh.abs().mean(dim=1)[0]                   # [H,W]
            elif kind == "npr_residual":   # NPR: neighboring-pixel residual magnitude
                res = model.NPR_transform(x)                    # [1,3,H,W]
                sal = res.abs().mean(dim=1)[0]                  # [H,W]
            elif kind == "pgc_peak":       # PGC: peak-guided residual fake-evidence map
                sal = model.peak_residual_map(x)               # [Hr,Wr]
            elif kind == "aide_dct":       # AIDE: native-res DCT per-window energy (patch selection)
                if pil is None:
                    return None
                sal = spec.transform.dct_energy_map(pil)       # [nH,Wn]
            elif kind == "spai_patch":     # SPAI: per-patch spectral fake-evidence over the 224-grid
                sal = model.patch_fake_map(x)                  # [nH,nW] or None (five-crop/large grid)
                if sal is None:
                    return None
            else:
                return None
            base = _geometric_base(spec, pil) if pil is not None else None
            out = {"map": _saliency_to_png_b64(sal)}
            if base is not None:
                out["overlay"] = _overlay_to_png_b64(sal, base)
            return out
        except Exception:  # noqa: BLE001 — interpretability is optional, never break predict
            return None

    @staticmethod
    def _model_input_data_url(spec: ModelSpec, x: torch.Tensor, max_side: int = 1280) -> str:
        """The image exactly as the network receives it (after this detector's preprocessing), as a PNG data URL.
        Normalization is undone for display; multi-view inputs (AIDE) are laid out side by side."""
        t = x[0].detach().float().cpu()
        if t.dim() == 4:
            t = torch.cat(list(t), dim=2)
        norm = next((s for s in getattr(spec.transform, "transforms", []) if isinstance(s, T.Normalize)), None)
        if norm is not None:
            t = t * torch.tensor(norm.std)[:, None, None] + torch.tensor(norm.mean)[:, None, None]
        elif spec.key == "cure":                       # CuRe maps pixels to (x - 127.5) / 127.5
            t = t * 0.5 + 0.5
        if t.min() < 0 or t.max() > 1:                 # unknown scaling: stretch for display only
            t = (t - t.min()) / (t.max() - t.min() + 1e-8)
        img = Image.fromarray((t.clamp(0, 1) * 255).round().byte().permute(1, 2, 0).numpy())
        if max(img.size) > max_side:
            k = max_side / max(img.size)
            img = img.resize((max(1, int(img.width * k)), max(1, int(img.height * k))), Image.LANCZOS)
        buf = io.BytesIO(); img.save(buf, "PNG")
        return "data:image/png;base64," + base64.b64encode(buf.getvalue()).decode()

    def predict_image(self, spec: ModelSpec, pil: Image.Image) -> dict:
        model = self._get_model(spec)
        x = spec.transform(pil).unsqueeze(0).to(self.device)
        t0 = time.time()
        p = float(self._fake_prob(model, spec, x).item())
        res = self._result(spec, p, time.time() - t0)
        res["orig_size"] = [pil.width, pil.height]
        res["input_size"] = [int(x.shape[-1]), int(x.shape[-2])]     # what the network actually received
        res["input_note"] = input_note(spec)
        try:
            res["input_image"] = self._model_input_data_url(spec, x)
        except Exception:                                           # display-only: never fail a prediction over it
            pass
        amap = self._interpret_map(model, spec, x, pil)
        if amap is not None:
            res["interpret_map"] = amap["map"]
            if amap.get("overlay"):
                res["interpret_overlay"] = amap["overlay"]
            res["interpret_kind"] = spec.interpret
        return res

    def predict_video_frames(self, spec: ModelSpec, frames: list[Image.Image]) -> dict:
        """Run an *image* detector on sampled video frames and aggregate (mean P(fake))."""
        model = self._get_model(spec)
        batch = torch.stack([spec.transform(f) for f in frames]).to(self.device)
        t0 = time.time()
        per_frame = self._fake_prob(model, spec, batch).detach().float().cpu().tolist()
        elapsed = time.time() - t0
        mean_p = sum(per_frame) / len(per_frame)
        res = self._result(spec, mean_p, elapsed)
        res["per_frame"] = [round(v, 6) for v in per_frame]
        res["aggregation"] = "mean over sampled frames"
        return res

    def video_clip(self, spec: ModelSpec, path) -> torch.Tensor:
        """Decode and normalize a clip exactly as ``spec`` prescribes: ``[T, 3, H, W]``."""
        clip = media.video_clip_tensor(
            path, num_frames=spec.num_frames, fps=spec.fps, frame_size=spec.frame_size,
            window_center=spec.window_center, interp=spec.resize_interp,
            edge_crop=getattr(spec, "edge_crop", 0.0),
            edge_crop_long_side=getattr(spec, "edge_crop_long_side", False),
            short_clip_frames=getattr(spec, "short_clip_frames", 0),
            fps_nearest=getattr(spec, "fps_nearest", False), bgr=getattr(spec, "bgr", False),
        )
        mean, std = spec.video_norm
        return (clip - torch.tensor(mean).view(1, 3, 1, 1)) / torch.tensor(std).view(1, 3, 1, 1)

    def predict_video_clip(self, spec: ModelSpec, path) -> dict:
        """Run a native *video* detector on a decoded clip."""
        model = self._get_model(spec)
        clip = self.video_clip(spec, path)
        t0 = time.time()
        p = float(self._video_fake_prob(model, spec, clip).item())
        res = self._result(spec, p, time.time() - t0)
        res["num_frames"] = int(clip.shape[0])
        return res

    def _result(self, spec: ModelSpec, fake_prob: float, elapsed: float) -> dict:
        fake_prob = max(0.0, min(1.0, fake_prob))
        return {
            "model": spec.key,
            "model_name": spec.name,
            "fake_probability": round(fake_prob, 6),
            "real_probability": round(1.0 - fake_prob, 6),
            "verdict": "fake" if fake_prob >= 0.5 else "real",
            "confidence": round(max(fake_prob, 1.0 - fake_prob), 6),
            "elapsed_ms": round(elapsed * 1000, 1),
            "checkpoint_status": self._status.get(spec.key, ""),
        }


_ENGINE: Engine | None = None


def get_engine() -> Engine:
    global _ENGINE
    if _ENGINE is None:
        _ENGINE = Engine()
    return _ENGINE