ff6a512630
Working tree (including .git) was lost to an rm. Rebuilt by replaying Write/Edit/ Read/attachment events from 25 Claude sessions and 22 successful codex apply_patch blocks into one timestamp-ordered timeline. Verified against ground truth recorded in the transcripts: wc -l on 10 files and ls -l on 5 files at 2026-07-18T13:13:44Z both match exactly; 18 files are byte-identical to their newest ~/.claude/file-history blob. See HANDOFF.md for sources, gaps, and how to rebuild .venv. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
323 lines
16 KiB
Python
323 lines
16 KiB
Python
# worker_identity.py — stage 5 character identity. FastAPI :8003.
|
|
# no model of its own conceptually, but siglip2 loads in-process here (transformers, rocm torch)
|
|
# after the orchestrator has opened a siglip2 session (the mutex guarantees it's the only GPU
|
|
# resident). known characters + their reference embeddings come from the homesrv orchestrator.
|
|
import os, uuid, json
|
|
from fastapi import FastAPI
|
|
from pydantic import BaseModel
|
|
import numpy as np
|
|
import requests
|
|
import transport
|
|
|
|
app = FastAPI()
|
|
transport.install_logging(app, "identity")
|
|
SHM = "/dev/shm"
|
|
ORCH = os.environ.get("ORCHESTRATOR_URL", "http://192.168.1.104:9090")
|
|
DEFAULT_THRESHOLD = 0.85
|
|
AMBIGUOUS_MARGIN = 0.05 # two known chars within this of each other -> flag low-confidence
|
|
|
|
_siglip = None # (model, processor), lazy-loaded
|
|
|
|
# F3 confirm-before-persist: unnamed provisional characters are held here (in memory, keyed by
|
|
# session_id = one chapter) and only written to DB/S3 once seen >=2x. seen-once NPCs are never
|
|
# persisted -> no DB row, no bucket crop. named chars skip the cache and persist immediately.
|
|
# entry: {emb, crop(np), count, name, gender, appearance, occ:[(panel_id,local_id,conf)], cid}
|
|
_pending: dict[str, list[dict]] = {}
|
|
PENDING_PROMOTE_AT = 2
|
|
|
|
|
|
def _load_siglip():
|
|
global _siglip
|
|
if _siglip is None:
|
|
import torch
|
|
from transformers import AutoModel, AutoProcessor
|
|
name = "google/siglip2-so400m-patch16-384"
|
|
model = AutoModel.from_pretrained(name).to("cuda").eval()
|
|
_siglip = (model, AutoProcessor.from_pretrained(name), torch)
|
|
return _siglip
|
|
|
|
|
|
def embed_crop(img) -> np.ndarray:
|
|
"""siglip2 image embedding of a BGR/RGB crop (numpy HxWx3), L2-normalized."""
|
|
model, proc, torch = _load_siglip()
|
|
inputs = proc(images=img, return_tensors="pt").to("cuda")
|
|
with torch.no_grad():
|
|
out = model.get_image_features(**inputs)
|
|
# transformers returns a ModelOutput here (not a bare tensor): use the attention-pooled
|
|
# image embedding. fall back to mean-pooling patch tokens if a build lacks a pooler head.
|
|
feat = getattr(out, "pooler_output", None)
|
|
if feat is None:
|
|
feat = getattr(out, "last_hidden_state", out)
|
|
if hasattr(feat, "dim") and feat.dim() == 3:
|
|
feat = feat.mean(dim=1)
|
|
feat = feat.detach().cpu().numpy().reshape(-1)
|
|
return feat / (np.linalg.norm(feat) + 1e-8)
|
|
|
|
|
|
def cosine(a: np.ndarray, b: np.ndarray) -> float:
|
|
return float(np.dot(a, b) / ((np.linalg.norm(a) * np.linalg.norm(b)) + 1e-8))
|
|
|
|
|
|
def gender_ok(a: str, b: str) -> bool:
|
|
"""two people can be the same only if their DECIDED genders agree. unknown on either side is a
|
|
pass (don't over-block). the single biggest siglip cross-match error is a male crop scoring >0.85
|
|
against a female character (shared art style + panel context); this hard-blocks it."""
|
|
a, b = (a or "").strip().lower(), (b or "").strip().lower()
|
|
return not (a in ("m", "f") and b in ("m", "f") and a != b)
|
|
|
|
|
|
def match(emb: np.ndarray, known: list, threshold: float):
|
|
"""known: [{"character_id","embedding"(np)}]. returns (character_id|None, confidence, ambiguous)."""
|
|
if not known:
|
|
return None, 0.0, False
|
|
scored = sorted(((cosine(emb, k["embedding"]), k["character_id"]) for k in known), reverse=True)
|
|
best_conf, best_id = scored[0]
|
|
ambiguous = len(scored) > 1 and (best_conf - scored[1][0]) < AMBIGUOUS_MARGIN
|
|
if best_conf >= threshold:
|
|
return best_id, best_conf, ambiguous
|
|
return None, best_conf, ambiguous
|
|
|
|
|
|
def shortlist(emb: np.ndarray, known: list, k: int = 5, gender: str = None) -> list:
|
|
"""TIER-2 evidence: the top-k gender-gated known characters by cosine, best first. Cosine is now a
|
|
SHORTLISTER, not the decider — gemma /vision/resolve picks from this list. Each item carries the
|
|
full row (name/gender/description) so the resolver can build a text character-sheet. pure."""
|
|
cands = [(cosine(emb, c["embedding"]), c) for c in known if gender_ok(gender, c.get("gender"))]
|
|
cands.sort(key=lambda t: t[0], reverse=True)
|
|
return [{**c, "cosine": round(s, 3)} for s, c in cands[:k]]
|
|
|
|
|
|
def _crop_bbox(img, bbox):
|
|
# vision emits [x1, y1, x2, y2] pixel corners (gemma4's native bbox convention).
|
|
x1, y1, x2, y2 = bbox
|
|
return img[y1:y2, x1:x2]
|
|
|
|
|
|
def _save_npy(emb: np.ndarray, uri: str):
|
|
"""upload an embedding in .npy format (np.load reads it back on the known-char side)."""
|
|
tmp = f"{SHM}/emb_{uuid.uuid4().hex[:8]}.npy"
|
|
np.save(tmp, emb.astype(np.float32))
|
|
transport.put(tmp, uri)
|
|
os.remove(tmp)
|
|
|
|
|
|
def _pending_match(pend: list, emb, threshold: float, gender: str = None):
|
|
"""index of the pending entry this embedding belongs to, or None (a new provisional). candidates
|
|
of a conflicting decided gender are excluded. character_id=i keeps the returned index original. pure."""
|
|
cands = [{"character_id": i, "embedding": e["emb"]}
|
|
for i, e in enumerate(pend) if gender_ok(gender, e.get("gender"))]
|
|
idx, conf, _ = match(emb, cands, threshold)
|
|
return idx, conf
|
|
|
|
|
|
def _persist_char(manga_id, panel_id, local_id, crop, emb, name, gender, appearance) -> str:
|
|
"""upload crop + embedding to S3 and register the row via the orchestrator; return its id."""
|
|
import cv2
|
|
key = f"{manga_id}/characters/_new/{panel_id}_{local_id}"
|
|
ref_img_uri, emb_uri = f"s3://manga/{key}.png", f"s3://manga/{key}.npy"
|
|
ref_png = f"{SHM}/crop_{uuid.uuid4().hex[:8]}.png"
|
|
cv2.imwrite(ref_png, crop)
|
|
transport.put(ref_png, ref_img_uri)
|
|
os.remove(ref_png)
|
|
_save_npy(emb, emb_uri)
|
|
resp = requests.post(f"{ORCH}/characters/create", json={
|
|
"manga_id": manga_id, "name": (name or "").strip() or None,
|
|
"gender": gender or "unknown", "description": appearance or {},
|
|
"ref_image_uri": ref_img_uri, "embedding_uri": emb_uri,
|
|
}, timeout=30).json()
|
|
_known_cache.pop(manga_id, None)
|
|
return resp["character_id"]
|
|
|
|
|
|
_known_cache: dict = {} # manga_id -> known list; invalidated in _persist_char
|
|
|
|
|
|
def _load_known(manga_id: str) -> list:
|
|
"""orchestrator /characters/known -> roster rows with their embeddings loaded. A dangling
|
|
embedding_uri (row persisted but .npy never landed / bucket wiped) is skipped, not fatal —
|
|
every panel loads the full roster up front and one bad ref must not 500 the panel.
|
|
Cached per manga_id for the run; _persist_char drops the entry when the roster changes."""
|
|
if manga_id in _known_cache:
|
|
return _known_cache[manga_id]
|
|
rows = requests.get(f"{ORCH}/characters/known", params={"manga_id": manga_id}, timeout=30).json()
|
|
known = []
|
|
for c in rows:
|
|
if not c.get("embedding_uri"):
|
|
continue
|
|
try:
|
|
ref = transport.get(c["embedding_uri"], f"{SHM}/ref_{uuid.uuid4().hex[:8]}.npy")
|
|
except Exception as e:
|
|
print(f"[identity] skip {c['character_id']}: bad embedding_uri {c['embedding_uri']}: {e}", flush=True)
|
|
continue
|
|
refs = c.get("ref_image_uris") or []
|
|
if isinstance(refs, str):
|
|
try:
|
|
refs = json.loads(refs)
|
|
except (ValueError, TypeError):
|
|
refs = []
|
|
known.append({"character_id": c["character_id"], "embedding": np.load(ref),
|
|
"gender": c.get("gender"), "name": c.get("name"),
|
|
"species": c.get("species"), "description": c.get("description"),
|
|
"reference_image_uris": refs})
|
|
os.remove(ref)
|
|
_known_cache[manga_id] = known
|
|
return known
|
|
|
|
|
|
class IdentityInput(BaseModel):
|
|
panel_uri: str
|
|
panel_id: str = ""
|
|
vision_characters: list = []
|
|
manga_id: str = ""
|
|
session_id: str = ""
|
|
k: int = 5 # shortlist width for gemma's tracklet decider (204: merged from /identity/candidates)
|
|
|
|
|
|
@app.post("/identity/resolve")
|
|
async def resolve(data: IdentityInput):
|
|
"""204: does the cosine assignment/persist pass AND builds gemma's shortlist in one crop/embed pass
|
|
(was two separate endpoints, each re-cropping + re-embedding + reloading the roster per character).
|
|
Cosine still makes the provisional assignment; the orchestrator's gemma tracklet phase (/vision/resolve)
|
|
can still override it — shortlists are returned for every crop regardless of assignment outcome."""
|
|
import cv2
|
|
local = transport.get(data.panel_uri, f"{SHM}/ident_{uuid.uuid4().hex[:8]}.png")
|
|
img = cv2.imread(local)
|
|
|
|
# orchestrator /characters/known returns a plain list of character rows (embedding_uri per row).
|
|
# ponytail: per-manga threshold isn't exposed by the orchestrator yet -> default; wire a
|
|
# manga_config lookup here if tuning per title ever matters.
|
|
threshold = DEFAULT_THRESHOLD
|
|
known = _load_known(data.manga_id)
|
|
|
|
# F3: only THIS session's pending cache is relevant (session == chapter); evict any others so a
|
|
# long-lived worker doesn't leak past chapters' provisionals.
|
|
for sid in [s for s in _pending if s != data.session_id]:
|
|
_pending.pop(sid, None)
|
|
pend = _pending.setdefault(data.session_id, [])
|
|
|
|
assignments, backfill, new_chars, shortlists = [], [], [], []
|
|
for ch in data.vision_characters:
|
|
crop = _crop_bbox(img, ch["bbox"])
|
|
if crop.size == 0: # degenerate/out-of-bounds bbox -> nothing to embed, skip
|
|
continue
|
|
emb = embed_crop(crop)
|
|
|
|
# shortlist for gemma's decider, from the roster as it stood before this crop's own outcome.
|
|
sl = shortlist(emb, known, data.k, ch.get("gender"))
|
|
crop_uri = f"s3://manga/{data.manga_id}/characters/_crops/{data.panel_id}_{ch['local_id']}.png"
|
|
cp = f"{SHM}/cc_{uuid.uuid4().hex[:8]}.png"; cv2.imwrite(cp, crop)
|
|
transport.put(cp, crop_uri); os.remove(cp)
|
|
shortlists.append({"local_id": ch["local_id"], "crop_uri": crop_uri,
|
|
"candidates": [{"character_id": c["character_id"], "name": c.get("name"),
|
|
"gender": c.get("gender"), "species": c.get("species"),
|
|
"appearance": c.get("description"), "cosine": c["cosine"],
|
|
"reference_image_uris": c.get("reference_image_uris", [])}
|
|
for c in sl]})
|
|
|
|
# gender gate: never match this crop to a known character of the opposite decided gender.
|
|
g = ch.get("gender")
|
|
cands = [k for k in known if gender_ok(g, k.get("gender"))]
|
|
cid, conf, ambiguous = match(emb, cands, threshold)
|
|
name = (ch.get("name") or "").strip()
|
|
if cid is not None: # matched an already-persisted character
|
|
assignments.append({"local_id": ch["local_id"], "character_id": cid,
|
|
"confidence": round(conf, 3), "ambiguous": ambiguous})
|
|
continue
|
|
if name: # named -> persist immediately (not an NPC)
|
|
cid = _persist_char(data.manga_id, data.panel_id, ch["local_id"], crop, emb,
|
|
name, ch.get("gender"), ch.get("appearance"))
|
|
known.append({"character_id": cid, "embedding": emb, "gender": ch.get("gender")})
|
|
new_chars.append(cid)
|
|
assignments.append({"local_id": ch["local_id"], "character_id": cid,
|
|
"confidence": round(conf, 3), "ambiguous": ambiguous})
|
|
continue
|
|
# unnamed + unknown -> confirm-before-persist. match against this chapter's pending cache.
|
|
pcid, pconf = _pending_match(pend, emb, threshold, ch.get("gender"))
|
|
if pcid is None: # first sighting: hold, do not persist or assign yet
|
|
pend.append({"emb": emb, "crop": crop, "count": 1, "name": name,
|
|
"gender": ch.get("gender"), "appearance": ch.get("appearance"),
|
|
"occ": [(data.panel_id, ch["local_id"], round(conf, 3))], "cid": None})
|
|
continue
|
|
e = pend[pcid] # seen before this chapter
|
|
e["count"] += 1
|
|
e["occ"].append((data.panel_id, ch["local_id"], round(pconf, 3)))
|
|
if e["cid"] is None and e["count"] >= PENDING_PROMOTE_AT: # promote -> persist + backfill
|
|
e["cid"] = _persist_char(data.manga_id, e["occ"][0][0], e["occ"][0][1], e["crop"],
|
|
e["emb"], e["name"], e["gender"], e["appearance"])
|
|
known.append({"character_id": e["cid"], "embedding": e["emb"], "gender": e.get("gender")})
|
|
new_chars.append(e["cid"])
|
|
for (bp, bl, bc) in e["occ"][:-1]: # earlier sightings that were deferred
|
|
backfill.append({"panel_id": bp, "local_id": bl,
|
|
"character_id": e["cid"], "confidence": bc})
|
|
if e["cid"]:
|
|
assignments.append({"local_id": ch["local_id"], "character_id": e["cid"],
|
|
"confidence": round(pconf, 3), "ambiguous": False})
|
|
os.remove(local)
|
|
return {"panel_id": data.panel_id, "assignments": assignments,
|
|
"backfill": backfill, "new_characters": new_chars, "shortlists": shortlists}
|
|
|
|
|
|
@app.post("/unload")
|
|
async def unload():
|
|
"""free the resident siglip2 so the session manager can hand the GPU to the next model.
|
|
the mutex can't reclaim in-process VRAM -- only the worker holding the model can."""
|
|
global _siglip
|
|
was = _siglip is not None
|
|
if was:
|
|
torch = _siglip[2]
|
|
_siglip = None
|
|
import gc; gc.collect()
|
|
torch.cuda.empty_cache()
|
|
_pending.clear() # F3: drop any held provisionals with the model
|
|
return {"ok": True, "unloaded": was}
|
|
|
|
|
|
@app.get("/health")
|
|
async def health():
|
|
return {"status": "ok"}
|
|
|
|
|
|
if __name__ == "__main__":
|
|
# self-check: cosine + match decision. deterministic vectors, no real model.
|
|
v = lambda *xs: np.array(xs, dtype=np.float32)
|
|
a, b = v(1, 0, 0), v(1, 0, 0)
|
|
assert abs(cosine(a, b) - 1.0) < 1e-6
|
|
known = [{"character_id": "c1", "embedding": v(1, 0, 0)},
|
|
{"character_id": "c2", "embedding": v(0, 1, 0)}]
|
|
cid, conf, amb = match(v(0.99, 0.01, 0), known, 0.85)
|
|
assert cid == "c1" and conf > 0.85 and not amb, (cid, conf, amb)
|
|
cid, conf, amb = match(v(0, 0, 1), known, 0.85) # nothing close -> new
|
|
assert cid is None
|
|
# ambiguous: equidistant-ish between c1 and c2
|
|
_, _, amb = match(v(0.71, 0.70, 0), known, 0.5)
|
|
assert amb is True
|
|
# gender gate: opposite decided genders never match; unknown on either side passes.
|
|
assert gender_ok("m", "m") and gender_ok("m", "unknown") and gender_ok("", "f")
|
|
assert not gender_ok("m", "f") and not gender_ok("f", "m")
|
|
# a male crop must NOT match a female known char even at high cosine (the Choi Haeseon bug).
|
|
kn = [{"character_id": "female_char", "embedding": v(1, 0, 0), "gender": "f"}]
|
|
cands = [k for k in kn if gender_ok("m", k.get("gender"))]
|
|
assert match(v(1, 0, 0), cands, 0.85)[0] is None # gated out -> new male char, not the female id
|
|
# pending gate: a female provisional is skipped for a male crop even if embeddings are identical.
|
|
pend_g = [{"emb": v(1, 0, 0), "count": 1, "gender": "f"}]
|
|
assert _pending_match(pend_g, v(1, 0, 0), 0.85, "m")[0] is None
|
|
assert _pending_match(pend_g, v(1, 0, 0), 0.85, "f")[0] == 0
|
|
# F3 confirm-before-persist: first sighting is new (held, not persisted); a matching second
|
|
# sighting hits the same pending entry -> promotes.
|
|
pend = []
|
|
i, _ = _pending_match(pend, v(1, 0, 0), 0.85)
|
|
assert i is None # nothing pending yet -> new provisional
|
|
pend.append({"emb": v(1, 0, 0), "count": 1})
|
|
i, _ = _pending_match(pend, v(0.99, 0.02, 0), 0.85)
|
|
assert i == 0 # second sighting matches the held entry
|
|
i, _ = _pending_match(pend, v(0, 0, 1), 0.85)
|
|
assert i is None # unrelated crop -> its own new provisional
|
|
# tier-2 shortlist: top-k by cosine, gender-gated, best first; carries the row for the sheet.
|
|
kn = [{"character_id": "c1", "embedding": v(1, 0, 0), "gender": "m", "name": "Gojo"},
|
|
{"character_id": "c2", "embedding": v(0, 1, 0), "gender": "f", "name": "Choi"},
|
|
{"character_id": "c3", "embedding": v(0.9, 0.1, 0), "gender": "m", "name": "Nanami"}]
|
|
sl = shortlist(v(1, 0, 0), kn, k=2, gender="m")
|
|
assert [c["character_id"] for c in sl] == ["c1", "c3"] # female c2 gated out; c1 beats c3
|
|
assert sl[0]["cosine"] >= sl[1]["cosine"] and sl[0]["name"] == "Gojo"
|
|
print("worker_identity self-check ok")
|