"""klaus-3 vision service: CLIP + all-MiniLM embeddings + LPIPS over HTTP. Runs on klaus-3 (has the GPU + torch). The Mac orchestrator calls these endpoints so it never needs torch locally. Endpoints: GET /health POST /clip/similarity form: image=, texts="a\nb\nc" -> {similarity:{text:cos}} POST /clip/zero_shot form: image=, labels="a\nb\nc" -> {scores:{label:prob}} POST /clip/caption form: image= -> {label: str} POST /embed json: {texts:[...]} -> {embeddings:[[...]]} POST /lpips form: a=, b= -> {lpips: float} Run: pip install fastapi uvicorn python-multipart open_clip_torch sentence-transformers lpips torch pillow uvicorn vision_service:app --host 0.0.0.0 --port 8090 """ from __future__ import annotations import io import numpy as np import torch from fastapi import FastAPI, File, Form, UploadFile from PIL import Image app = FastAPI(title="veil-pgd vision service") _DEVICE = "cuda" if torch.cuda.is_available() else "cpu" # --- lazy singletons so startup is fast and VRAM is only used on demand --- _clip = {"model": None, "preprocess": None, "tokenizer": None} _embed_model = None _lpips_model = None # A small open-vocabulary label set for /clip/caption fallback. _CAPTION_VOCAB = [ "dog", "cat", "car", "person", "tree", "building", "food", "flower", "bird", "boat", "bicycle", "mountain", "beach", "book", "phone", ] def _load_clip(): if _clip["model"] is None: import open_clip model, _, preprocess = open_clip.create_model_and_transforms( "ViT-B-32", pretrained="laion2b_s34b_b79k" ) model = model.to(_DEVICE).eval() _clip["model"] = model _clip["preprocess"] = preprocess _clip["tokenizer"] = open_clip.get_tokenizer("ViT-B-32") return _clip def _load_embed(): global _embed_model if _embed_model is None: from sentence_transformers import SentenceTransformer _embed_model = SentenceTransformer("all-MiniLM-L6-v2", device=_DEVICE) return _embed_model def _load_lpips(): global _lpips_model if _lpips_model is None: import lpips _lpips_model = lpips.LPIPS(net="alex").to(_DEVICE).eval() return _lpips_model def _read_image(data: bytes) -> Image.Image: return Image.open(io.BytesIO(data)).convert("RGB") @torch.no_grad() def _clip_features(image: Image.Image, texts: list[str]): c = _load_clip() img_t = c["preprocess"](image).unsqueeze(0).to(_DEVICE) tok = c["tokenizer"](texts).to(_DEVICE) img_f = c["model"].encode_image(img_t) txt_f = c["model"].encode_text(tok) img_f = img_f / img_f.norm(dim=-1, keepdim=True) txt_f = txt_f / txt_f.norm(dim=-1, keepdim=True) return img_f, txt_f @app.get("/health") def health(): return {"ok": True, "device": _DEVICE} @app.post("/clip/similarity") async def clip_similarity(image: UploadFile = File(...), texts: str = Form(...)): img = _read_image(await image.read()) text_list = [t for t in texts.split("\n") if t.strip()] img_f, txt_f = _clip_features(img, text_list) sims = (img_f @ txt_f.T).squeeze(0).tolist() return {"similarity": dict(zip(text_list, sims))} @app.post("/clip/zero_shot") async def clip_zero_shot(image: UploadFile = File(...), labels: str = Form(...)): img = _read_image(await image.read()) label_list = [t for t in labels.split("\n") if t.strip()] img_f, txt_f = _clip_features(img, label_list) logits = (100.0 * img_f @ txt_f.T).softmax(dim=-1).squeeze(0).tolist() return {"scores": dict(zip(label_list, logits))} @app.post("/clip/caption") async def clip_caption(image: UploadFile = File(...)): img = _read_image(await image.read()) img_f, txt_f = _clip_features(img, _CAPTION_VOCAB) idx = int((img_f @ txt_f.T).squeeze(0).argmax().item()) return {"label": _CAPTION_VOCAB[idx]} @app.post("/embed") async def embed(payload: dict): texts = payload.get("texts", []) model = _load_embed() vecs = model.encode(texts, normalize_embeddings=True) return {"embeddings": np.asarray(vecs).tolist()} @torch.no_grad() @app.post("/lpips") async def lpips_endpoint(a: UploadFile = File(...), b: UploadFile = File(...)): ia = _read_image(await a.read()) ib = _read_image(await b.read()).resize(ia.size) def to_t(im: Image.Image) -> torch.Tensor: arr = np.asarray(im, dtype=np.float32) / 127.5 - 1.0 return torch.from_numpy(arr).permute(2, 0, 1).unsqueeze(0).to(_DEVICE) model = _load_lpips() d = model(to_t(ia), to_t(ib)).item() return {"lpips": float(d)} # ============================================================================= # White-box ViT endpoints (P0 of the pivot). Uniform encoder interface over # OpenCLIP and SigLIP2 so the Mac can score decoy directions and probe gradients. # ============================================================================= _encoders: dict = {} class _OpenClipEncoder: """OpenCLIP ViT-* wrapper with a differentiable image path.""" def __init__(self, arch: str, pretrained: str): import open_clip self.model, _, self.preprocess = open_clip.create_model_and_transforms( arch, pretrained=pretrained ) self.model = self.model.to(_DEVICE).eval() self.tokenizer = open_clip.get_tokenizer(arch) self.input_px = 224 def image_tensor(self, image: Image.Image) -> torch.Tensor: return self.preprocess(image).unsqueeze(0).to(_DEVICE) def image_feat(self, t: torch.Tensor) -> torch.Tensor: f = self.model.encode_image(t) return f / f.norm(dim=-1, keepdim=True) def text_feat(self, texts: list[str]) -> torch.Tensor: tok = self.tokenizer(texts).to(_DEVICE) f = self.model.encode_text(tok) return f / f.norm(dim=-1, keepdim=True) class _SiglipEncoder: """HF SigLIP2 wrapper with a differentiable image path (fp32 for grads).""" def __init__(self, model_id: str, dtype=torch.float16): from transformers import AutoModel, AutoProcessor self.dtype = dtype self.model = AutoModel.from_pretrained(model_id, torch_dtype=dtype) self.model = self.model.to(_DEVICE).eval() self.processor = AutoProcessor.from_pretrained(model_id) self.input_px = getattr( self.processor.image_processor, "size", {} ).get("height", 224) def image_tensor(self, image: Image.Image) -> torch.Tensor: px = self.processor(images=image, return_tensors="pt")["pixel_values"] return px.to(_DEVICE, dtype=self.dtype) @staticmethod def _as_tensor(x): if isinstance(x, torch.Tensor): return x for attr in ("image_embeds", "text_embeds", "pooler_output", "last_hidden_state"): v = getattr(x, attr, None) if isinstance(v, torch.Tensor): return v.mean(dim=1) if v.dim() == 3 else v raise TypeError(f"cannot extract tensor from {type(x)}") def image_feat(self, t: torch.Tensor) -> torch.Tensor: f = self._as_tensor(self.model.get_image_features(pixel_values=t)) return f / f.norm(dim=-1, keepdim=True) def text_feat(self, texts: list[str]) -> torch.Tensor: tok = self.processor( text=texts, padding="max_length", max_length=64, return_tensors="pt" ).to(_DEVICE) f = self._as_tensor(self.model.get_text_features(input_ids=tok["input_ids"])) return f / f.norm(dim=-1, keepdim=True) _ENCODER_SPECS = { "openclip:ViT-B-32": ("openclip", "ViT-B-32", "laion2b_s34b_b79k"), "openclip:ViT-L-14": ("openclip", "ViT-L-14", "laion2b_s32b_b82k"), "siglip2-base": ("siglip", "google/siglip2-base-patch16-224", None), "siglip-base": ("siglip", "google/siglip-base-patch16-224", None), } def _load_encoder(model_id: str): if model_id not in _encoders: if model_id not in _ENCODER_SPECS: raise ValueError(f"unknown model_id {model_id}") kind, a, b = _ENCODER_SPECS[model_id] if kind == "openclip": _encoders[model_id] = _OpenClipEncoder(a, b) else: _encoders[model_id] = _SiglipEncoder(a) return _encoders[model_id] @app.get("/vit/health") def vit_health(): mem = {} if _DEVICE == "cuda": mem = { "allocated_mb": round(torch.cuda.memory_allocated() / 1e6, 1), "reserved_mb": round(torch.cuda.memory_reserved() / 1e6, 1), } return {"ok": True, "device": _DEVICE, "loaded": list(_encoders.keys()), "available": list(_ENCODER_SPECS.keys()), "cuda_memory": mem} @app.post("/vit/load") async def vit_load(payload: dict): mid = payload["model_id"] enc = _load_encoder(mid) return {"loaded": mid, "input_px": enc.input_px, "device": _DEVICE} @torch.no_grad() @app.post("/vit/score") async def vit_score( image: UploadFile = File(...), truth: str = Form(...), decoy: str = Form(...), model_id: str = Form("openclip:ViT-B-32"), clean: UploadFile | None = File(None), ): """Contrastive decoy-direction score for one candidate image. margin = cos(img, decoy) - cos(img, truth). If a clean image is supplied, also returns the clean margin and delta (how much the overlay moved it) plus away_from_clean = 1 - cos(img_candidate, img_clean). """ enc = _load_encoder(model_id) img = _read_image(await image.read()) it = enc.image_tensor(img) ifeat = enc.image_feat(it) tfeat = enc.text_feat([truth, decoy]) sim_truth = float((ifeat @ tfeat[0:1].T).item()) sim_decoy = float((ifeat @ tfeat[1:2].T).item()) out = {"model_id": model_id, "sim_truth": sim_truth, "sim_decoy": sim_decoy, "margin": sim_decoy - sim_truth} if clean is not None: cimg = _read_image(await clean.read()) cfeat = enc.image_feat(enc.image_tensor(cimg)) c_truth = float((cfeat @ tfeat[0:1].T).item()) c_decoy = float((cfeat @ tfeat[1:2].T).item()) out["clean_margin"] = c_decoy - c_truth out["delta_margin"] = out["margin"] - out["clean_margin"] out["away_from_clean"] = 1.0 - float((ifeat @ cfeat.T).item()) return out @torch.no_grad() @app.post("/vit/zero_shot") async def vit_zero_shot( image: UploadFile = File(...), labels: str = Form(...), model_id: str = Form("openclip:ViT-B-32"), ): enc = _load_encoder(model_id) img = _read_image(await image.read()) label_list = [t for t in labels.split("\n") if t.strip()] ifeat = enc.image_feat(enc.image_tensor(img)) tfeat = enc.text_feat(label_list) sims = (ifeat @ tfeat.T).squeeze(0).tolist() return {"similarity": dict(zip(label_list, sims))} @app.post("/vit/grad_region") async def vit_grad_region( image: UploadFile = File(...), truth: str = Form(...), decoy: str = Form(...), model_id: str = Form("openclip:ViT-B-32"), region_json: str = Form("null"), step_eps: float = Form(0.01), ): """Single backward pass of margin loss wrt the input tensor. Proves the differentiable path works: returns grad L2 (whole + in region) and whether one gradient-ascent step actually increases the decoy margin. """ import json as _json enc = _load_encoder(model_id) img = _read_image(await image.read()) t = enc.image_tensor(img).clone().detach().requires_grad_(True) tfeat = enc.text_feat([truth, decoy]).detach() ifeat = enc.image_feat(t) margin = (ifeat @ tfeat[1:2].T) - (ifeat @ tfeat[0:1].T) margin0 = float(margin.item()) enc.model.zero_grad(set_to_none=True) margin.backward() g = t.grad.detach() grad_l2 = float(g.norm().item()) region = _json.loads(region_json) grad_l2_region = None if region is not None: # map region [x0,y0,x1,y1] in original px to the input tensor grid W, H = img.size px = t.shape[-1] sx, sy = px / W, px / H x0, y0, x1, y1 = region rx0, ry0 = max(0, int(x0 * sx)), max(0, int(y0 * sy)) rx1, ry1 = min(px, int(x1 * sx)), min(px, int(y1 * sy)) if rx1 > rx0 and ry1 > ry0: grad_l2_region = float(g[..., ry0:ry1, rx0:rx1].norm().item()) # one gradient-ascent step on the margin, re-encode, check it went up with torch.no_grad(): t2 = (t + step_eps * g.sign()).detach() ifeat2 = enc.image_feat(t2) margin1 = float(((ifeat2 @ tfeat[1:2].T) - (ifeat2 @ tfeat[0:1].T)).item()) return {"model_id": model_id, "margin_before": margin0, "margin_after_step": margin1, "margin_increased": margin1 > margin0, "grad_l2": grad_l2, "grad_l2_region": grad_l2_region} # CLIP normalization constants (OpenAI/LAION share these). _CLIP_MEAN = (0.48145466, 0.4578275, 0.40821073) _CLIP_STD = (0.26862954, 0.26130258, 0.27577711) def _clip_feat_from_pixels(enc, x: torch.Tensor) -> torch.Tensor: """Differentiable CLIP image feature from a full-res [0,1] (1,3,H,W) tensor. Replicates resize+normalize (skips center crop) so gradients flow back to the original-resolution pixels and the perturbation is meaningful at full res. """ import torch.nn.functional as F x = F.interpolate(x, size=(224, 224), mode="bicubic", align_corners=False) mean = torch.tensor(_CLIP_MEAN, device=x.device).view(1, 3, 1, 1) std = torch.tensor(_CLIP_STD, device=x.device).view(1, 3, 1, 1) x = (x.clamp(0, 1) - mean) / std f = enc.model.encode_image(x) return f / f.norm(dim=-1, keepdim=True) @app.post("/vit/pgd_region") async def vit_pgd_region( image: UploadFile = File(...), truth: str = Form(...), decoy: str = Form(...), model_id: str = Form("openclip:ViT-B-32"), region_json: str = Form("null"), mask: UploadFile | None = File(None), eps: float = Form(0.0627), # L-inf budget in [0,1] px (~16/255) steps: int = Form(60), step_size: float = Form(0.0078), # ~2/255 return_image: bool = Form(True), ): """PGD that maximizes the decoy-truth margin, confined to a text box or an arbitrary editable mask. This is the *ceiling probe*: how far can ANY in-region perturbation move the encoder? mask (grayscale, >0 = editable) restricts edits to the text glyphs (stealthy ceiling); region_json alone = full box (Nightshade ceiling). Returns the optimized image so the caller can test caption flip / JPEG. """ import base64 import io as _io import json as _json if not model_id.startswith("openclip"): return {"error": "pgd_region currently supports openclip:* only"} enc = _load_encoder(model_id) img = _read_image(await image.read()) W, H = img.size x0 = (torch.from_numpy(np.asarray(img, dtype=np.float32) / 255.0) .permute(2, 0, 1).unsqueeze(0).to(_DEVICE)) # editable mask (1 where we may perturb) m = torch.zeros(1, 1, H, W, device=_DEVICE) region = _json.loads(region_json) if mask is not None: mimg = Image.open(_io.BytesIO(await mask.read())).convert("L").resize((W, H)) marr = torch.from_numpy(np.asarray(mimg, dtype=np.float32) / 255.0) m[0, 0] = (marr > 0.05).float().to(_DEVICE) elif region is not None: rx0, ry0, rx1, ry1 = [int(v) for v in region] m[0, 0, max(0, ry0):min(H, ry1), max(0, rx0):min(W, rx1)] = 1.0 else: m[:] = 1.0 editable_px = int(m.sum().item()) tfeat = enc.text_feat([truth, decoy]).detach() def margin_of(x): f = _clip_feat_from_pixels(enc, x) return (f @ tfeat[1:2].T) - (f @ tfeat[0:1].T) with torch.no_grad(): margin0 = float(margin_of(x0).item()) delta = torch.zeros_like(x0, requires_grad=True) for _ in range(int(steps)): x_adv = (x0 + delta * m).clamp(0, 1) loss = margin_of(x_adv) enc.model.zero_grad(set_to_none=True) if delta.grad is not None: delta.grad = None loss.backward() with torch.no_grad(): delta += step_size * delta.grad.sign() delta.clamp_(-eps, eps) with torch.no_grad(): x_final = (x0 + delta * m).clamp(0, 1) margin1 = float(margin_of(x_final).item()) linf = float((delta * m).abs().max().item()) out = {"model_id": model_id, "margin_before": margin0, "margin_after": margin1, "delta_margin": margin1 - margin0, "eps": eps, "steps": int(steps), "editable_px": editable_px, "linf": linf, "mode": "mask" if mask is not None else ("box" if region else "full")} if return_image: arr = (x_final.squeeze(0).permute(1, 2, 0).cpu().numpy() * 255.0) arr = arr.round().astype("uint8") buf = _io.BytesIO() Image.fromarray(arr).save(buf, format="PNG") out["image_png_b64"] = base64.b64encode(buf.getvalue()).decode() return out