src/veil_pgd/optimizer/tier_a.py

"""Tier A: free surrogate search (klaus-3 models). No API cost.

Stage 1 coarse grid over structural knobs, then Stage 2 evolutionary/hill-climb
over the fine knobs, scored by surrogate label distance-to-truth with scraper
preprocessing baked in and stealth as a soft penalty. Returns gate-passing
candidates ranked by surrogate fitness.
"""

from __future__ import annotations

import random
from concurrent.futures import ThreadPoolExecutor

from PIL import Image

from veil_pgd.config import Settings
from veil_pgd.fitness.embed import Embedder
from veil_pgd.fitness.objective import combined_fitness
from veil_pgd.fitness.semantic import aggregate, embedding_distance
from veil_pgd.optimizer import search_space as ss
from veil_pgd.render import render
from veil_pgd.robustness import scraper_sim
from veil_pgd.stealth import evaluate_stealth
from veil_pgd.targets.base import LabelPrompt
from veil_pgd.types import Candidate, RenderSpec


def cfg_workers(settings: Settings) -> int:
    """Concurrency for cross-spec scoring. Bounded so we don't blow the shared
    GPU's per-slot KV cache on klaus-3."""
    return max(2, int(getattr(settings.optimizer, "tier_a_workers", 4)))


class TierA:
    def __init__(
        self,
        settings: Settings,
        surrogates: list,
        clip,
        embedder: Embedder,
        lpips_fn=None,
        stealth_thresholds=None,
    ):
        self.s = settings
        self.surrogates = surrogates
        self.clip = clip
        self.embedder = embedder
        self.lpips_fn = lpips_fn
        self.thresholds = stealth_thresholds or settings.stealth
        self.prompt = LabelPrompt()
        # Independent specs pipeline over HTTP to the (separate) llama-servers.
        # Surrogates are queried sequentially *within* a spec to avoid a nested
        # pool deadlock; cross-spec concurrency provides the overlap.
        self._pool = ThreadPoolExecutor(max_workers=cfg_workers(settings))

    def _score_spec(self, image: Image.Image, spec: RenderSpec, truth: str) -> Candidate:
        rendered = render(image, spec)
        probe = scraper_sim(rendered)

        dists: list[float] = []
        preds: dict[str, str] = {}
        for m in self.surrogates:
            res = m.label(probe, self.prompt)
            preds[m.name] = res.parsed_label
            dists.append(embedding_distance(self.embedder, res.parsed_label, truth))
        distance = aggregate(dists, self.s.aggregation)

        # reachability: how close is the decoy to what surrogates actually said
        reach_vals = [
            1.0 - embedding_distance(self.embedder, p, spec.text) for p in preds.values()
        ]
        reachability = aggregate(reach_vals, "mean") if reach_vals else 0.0

        # stealth measured on the (clean) rendered image within the text region
        stealth = evaluate_stealth(
            image, rendered, thresholds=self.thresholds, lpips_fn=self.lpips_fn
        )

        fitness = combined_fitness(distance, reachability, stealth, self.thresholds)
        return Candidate(
            spec=spec,
            surrogate_fitness=fitness,
            per_model_distance={k: d for k, d in zip(preds, dists)},
            reachability=reachability,
            stealth=stealth,
            notes={"surrogate_preds": preds, "distance_to_truth": distance},
        )

    def search(
        self, image: Image.Image, truth: str, decoys: list[str]
    ) -> list[Candidate]:
        cfg = self.s.optimizer
        rng = random.Random(0)

        # Stage 1: coarse grid over decoys x positions x color x font size.
        positions = ss.default_positions()
        specs = ss.grid_specs(decoys, positions, cfg)

        scored = list(
            self._pool.map(lambda sp: self._score_spec(image, sp, truth), specs)
        )
        scored.sort(key=lambda c: c.surrogate_fitness, reverse=True)

        # Stage 2: evolutionary refine seeded from the best grid cells.
        population = scored[: cfg.evo_population]
        best = population[0].surrogate_fitness if population else 0.0
        stale = 0
        for _ in range(cfg.evo_generations):
            child_specs = [ss.mutate(p.spec, cfg, rng) for p in population]
            children = list(
                self._pool.map(lambda sp: self._score_spec(image, sp, truth), child_specs)
            )
            population = sorted(
                population + children, key=lambda c: c.surrogate_fitness, reverse=True
            )[: cfg.evo_population]
            top = population[0].surrogate_fitness
            if top - best < cfg.early_stop_epsilon:
                stale += 1
                if stale >= cfg.early_stop_patience:
                    break
            else:
                stale = 0
            best = max(best, top)

        # Hard stealth gate; rank survivors by surrogate fitness.
        survivors = [c for c in population if c.stealth and c.stealth.passed]
        survivors.sort(key=lambda c: c.surrogate_fitness, reverse=True)
        return survivors[: self.s.budget.top_k_to_validate]

    def close(self) -> None:
        self._pool.shutdown(wait=False, cancel_futures=True)