# Local AI Image Detector

A Chrome Manifest V3 extension that detects AI-generated images on any webpage,
entirely inside the browser. No cloud inference, no external APIs, no local
server — every tensor operation runs as a **hand-written WebGPU compute
kernel** written directly against the WebGPU API (no ONNX Runtime, no
transformers.js, no WASM inference runtime).

## How it works

- **Model**: [`buildborderless/CommunityForensics-DeepfakeDet-ViT`](https://huggingface.co/buildborderless/CommunityForensics-DeepfakeDet-ViT)
  — a ViT-Small/16 vision transformer trained on 2.7M images across 4,803
  generators ([Community Forensics, CVPR 2025](https://arxiv.org/abs/2411.04125)),
  fine-tuned as a binary real/fake classifier with a single sigmoid logit.
- **One-time download**: on first use, the extension downloads the model's
  `model.safetensors` weights (~87MB) directly from Hugging Face and caches
  them in IndexedDB. Every run after that reads only from the local cache —
  the extension never re-downloads or fetches any other model/weight file.
- **Inference**: the entire ViT forward pass (patch embedding, 12 transformer
  encoder blocks, classifier head) is implemented from scratch as WGSL
  compute shaders in [`extension/src/lib/shaders.js`](extension/src/lib/shaders.js)
  and orchestrated in [`extension/src/lib/vit-engine.js`](extension/src/lib/vit-engine.js) —
  a register-tiled batched GEMM, layernorm, exact GELU, multi-head attention
  and softmax, all running as native WebGPU pipelines.
- **Multi-scale scoring**: the detector reads pixel-level generator
  fingerprints, so it is very sensitive to what resampling an image has been
  through. Each image is scored at several views (the processor-faithful
  resize-and-crop, plus 384x384 centre crops at 1/1, 1/2 and 1/4 scale) and
  the strongest response wins — see
  [`preprocess.js`](extension/src/lib/preprocess.js).
- **Calibration**: the model's raw sigmoid is not centred on 0.5; its decision
  boundary sits near 1%. A monotonic logit-space shift maps that operating
  point onto the displayed 65% threshold, so the badge number means what it
  says — see [`calibration.js`](extension/src/lib/calibration.js).
- **On-page detection**: a content script finds `<img>` elements as they
  scroll into view and sends their URL to the background worker, which
  forwards it to the offscreen document. The offscreen document — an
  extension page with the same `host_permissions`-granted, CORS-free fetch
  privilege as the background worker — fetches the image bytes itself,
  decodes them, and runs WebGPU inference, so the image data never leaves
  the browser and never crosses `chrome.runtime` messaging as a large binary
  payload (only URL strings and small JSON results are passed between
  contexts).

## Performance

Measured on an Apple Silicon Mac via [`tools/verify.html`](tools/verify.html),
which runs the shipped engine in the browser:

| | per forward pass |
| --- | ---: |
| First inference (pipeline warm-up) | ~42 ms |
| Steady state | **17-20 ms** |

That is the full 384x384 ViT-Small/16 forward pass — 577 tokens, 12 encoder
blocks, ~15.3 GMAC — in fp32, with no quantization. Weight load from the
IndexedDB cache is a one-off ~90 ms. Because each image is scored at several
views, wall-clock cost per image is roughly 20 ms times the number of views
(2-4, depending on the image's size).

For reference, the same model as a Q8-quantized ONNX graph on the WASM
backend runs ~136-143 ms per image, so the hand-written WebGPU path is
roughly 7x faster while keeping full fp32 precision.

### Where the speed comes from

The forward pass is recorded once, at weight-load time, into a static list of
(pipeline, bind group, workgroup counts); per-image work is then just "upload
the image, replay the plan, read back one float".

- The whole model runs as **one compute pass** rather than ~150 begin/end pairs.
- Every buffer, uniform and bind group is allocated once at load, not per image
  (previously ~150 `createBuffer` + `createBindGroup` calls *per image*).
- **One generic strided batched GEMM** replaces five separate kernels. It is
  register-tiled: a 16x16 workgroup produces a 64x64 output tile, 16 outputs per
  thread, so each shared-memory read feeds 4 MACs instead of 1. Shared tiles are
  k-major and padded to kill bank conflicts.
- Bias, GELU and the residual add are **fused into the GEMM epilogue**, removing
  the separate elementwise passes and two activation buffers.
- **Q, K and V are one fused GEMM**; strides read the Q/K/V slices back out, and
  read V transposed for the attention product, with no transpose pass.
- The last encoder block runs its MLP on the CLS row only, and the final
  layernorm on one row instead of 577.

Together these took the forward pass from ~67 ms to ~17 ms with **bit-identical
output** — the engine still agrees with the CPU reference to every digit printed
(see [Verification](#verification)).

## Verification

[`tools/`](tools/) contains an offline oracle: `reference-forward.mjs` is a
plain-JS fp32 CPU implementation of the same forward pass, with no dependencies.
Both it and `verify.html` can be fed a fixed pseudo-random input, so the CPU and
GPU paths see bit-identical tensors:

| | P(fake) |
| --- | --- |
| CPU reference (Node) | `0.891299` |
| WebGPU engine | `0.891299` |

Re-run this after touching any kernel. See [`tools/README.md`](tools/README.md).

## Architecture

```
content script (finds <img> elements, overlays badges)
       │  chrome.runtime messages (URL string only)
       ▼
background service worker (manages offscreen doc lifecycle, routes messages)
       │  chrome.runtime messages (URL string only)
       ▼
offscreen document (fetches image, WebGPU device, model cache, inference engine)
```

An offscreen document is used (rather than running WebGPU in the service
worker or content script) because it's the standard MV3 pattern for
DOM/GPU-dependent background work: it has a real `navigator.gpu` + canvas
context and persists independently of any particular tab.

## Build & install

No build step or bundler is required — the extension is plain ES modules
loaded directly by Chrome.

1. Clone this repository.
2. Open `chrome://extensions` in Chrome.
3. Enable **Developer mode** (top-right toggle).
4. Click **Load unpacked** and select the `extension/` directory in this repo.
5. Visit any webpage. On the first image encountered, the extension
   downloads the model weights once (~87MB, progress shown in the extension
   popup) and caches them in IndexedDB. After that, everything runs offline.

Requires a Chrome version with WebGPU support (Chrome 121+ recommended;
`chrome://gpu` should show WebGPU as available). No other dependencies.

## Confidence scores & the 65% threshold

Every analyzed image gets a badge showing a single number: the probability the
image is AI-generated. Only the colour changes across the bounty's 65%
threshold.

- **Red**: `P(AI) >= 65%` — the image is also blurred, reversibly. Hovering
  reveals it temporarily; clicking reveals it for good.
- **Yellow**: between 35% and 65%.
- **Green**: `P(AI) <= 35%`.

The displayed number is the calibrated score. The raw model output is kept in
the badge tooltip and the console log for debugging.

## Privacy

Image bytes are fetched by the extension (needed to run the model on them)
but are never sent anywhere except through the local WebGPU pipeline running
on your own GPU. No image, URL, or inference result is transmitted to any
server controlled by this project or anyone else.

## Repository layout

```
extension/
  manifest.json           MV3 manifest
  src/
    background.js          service worker: routing, image fetch, offscreen lifecycle
    content.js              finds images on the page, requests analysis, renders badges
    content.css              badge styling
    offscreen.html/.js       WebGPU device, model download/cache, inference
    popup.html/.js           status popup
    lib/
      model-config.js         architecture + preprocessing constants
      safetensors.js          minimal .safetensors parser
      model-cache.js          IndexedDB weight cache
      preprocess.js           resize/crop/normalize (CLIP-style), multi-scale views
      calibration.js          maps the raw sigmoid onto the displayed threshold
      shaders.js               WGSL compute kernels
      vit-engine.js            forward-pass orchestration
tools/                     offline CPU oracle + in-browser verification harness
```

## License

MIT — see [LICENSE](LICENSE).