diff --git a/agent/docker-compose.yml b/agent/docker-compose.yml index 0f07cfa..5032669 100644 --- a/agent/docker-compose.yml +++ b/agent/docker-compose.yml @@ -37,6 +37,18 @@ services: # Crop embedder (SigLIP concept bag): float16 keeps VRAM low on a shared # desktop GPU; the model itself is announced by the server. SIGLIP_DTYPE: ${SIGLIP_DTYPE:-float16} + # Crop PROPOSERS (extra YOLO detectors → more/better concept crops). Each + # downloads its weights once (cached on the models volume) and self-disables + # if the download/load fails. Blank any one to turn it off. + # PERSON_WEIGHTS: general COCO person detector (Western/realistic figures), + # merged with the anime detector. yolo11n.pt (~6 MB, auto-downloaded). + # ANATOMY_WEIGHTS: booru_yolo anime/furry/NSFW components (~40 MB). NB the + # repo states no license — fine for private use. yolov8n_as01.pt is the + # 6 MB nano if you want lighter than yolov11m_aa22.pt. + # PANEL_WEIGHTS: mosesb comic-panel detector (Apache-2.0), "hf_repo::file". + PERSON_WEIGHTS: ${PERSON_WEIGHTS:-yolo11n.pt} + ANATOMY_WEIGHTS: ${ANATOMY_WEIGHTS:-https://github.com/aperveyev/booru_yolo/raw/main/models/yolov11m_aa22.pt} + PANEL_WEIGHTS: ${PANEL_WEIGHTS:-mosesb/best-comic-panel-detection::best.pt} volumes: # Persist the downloaded ONNX models so restarts are fast. - fc-agent-models:/models diff --git a/agent/fc_agent/client.py b/agent/fc_agent/client.py index 7003bcc..1c297c9 100644 --- a/agent/fc_agent/client.py +++ b/agent/fc_agent/client.py @@ -40,6 +40,19 @@ class FcClient: r.raise_for_status() return r.json() + def submit_embedding(self, job_id: int, embedding: list, version: str) -> dict: + """Post a whole-image SigLIP embedding (the 'embed' task) → image_record.""" + r = self.s.post( + f"{self.base}/api/gpu/jobs/submit_embedding", + json={ + "agent_id": self.agent_id, "job_id": job_id, + "embedding": embedding, "embedding_version": version, + }, + timeout=120, + ) + r.raise_for_status() + return r.json() + def heartbeat(self, job_ids: list[int]) -> None: try: self.s.post( diff --git a/agent/fc_agent/config.py b/agent/fc_agent/config.py index 42dee5c..630fc3f 100644 --- a/agent/fc_agent/config.py +++ b/agent/fc_agent/config.py @@ -18,6 +18,16 @@ class Config: # the server announces in the lease) auto_start: bool # start the worker pool on boot (so a container restart # resumes processing without anyone clicking Start) + # Crop PROPOSERS (extra YOLO detectors that say where to crop). Each weight + # spec is an ultralytics name | http(s) URL | "hf_repo::file" ("" = off). + person_weights: str # general COCO person detector (Western/realistic figs) + person_conf: float + anatomy_weights: str # booru_yolo anime/furry/NSFW components + anatomy_conf: float + panel_weights: str # comic-panel detector + panel_conf: float + max_components: int # cap anatomy component crops per frame + max_panels: int # cap panel crops per frame @classmethod def from_env(cls) -> "Config": @@ -33,4 +43,12 @@ class Config: embed_dtype=os.environ.get("SIGLIP_DTYPE", "float16"), embed_model_override=os.environ.get("EMBED_MODEL_NAME", ""), auto_start=os.environ.get("AUTO_START", "").lower() in ("1", "true", "yes"), + person_weights=os.environ.get("PERSON_WEIGHTS", "yolo11n.pt"), + person_conf=float(os.environ.get("PERSON_CONF", "0.35")), + anatomy_weights=os.environ.get("ANATOMY_WEIGHTS", ""), + anatomy_conf=float(os.environ.get("ANATOMY_CONF", "0.30")), + panel_weights=os.environ.get("PANEL_WEIGHTS", ""), + panel_conf=float(os.environ.get("PANEL_CONF", "0.30")), + max_components=int(os.environ.get("MAX_COMPONENTS", "8")), + max_panels=int(os.environ.get("MAX_PANELS", "8")), ) diff --git a/agent/fc_agent/detectors.py b/agent/fc_agent/detectors.py new file mode 100644 index 0000000..f91957f --- /dev/null +++ b/agent/fc_agent/detectors.py @@ -0,0 +1,163 @@ +"""Region PROPOSERS — small YOLO detectors that decide WHERE to crop. They run +on the agent GPU and their boxes feed the crop → SigLIP → max-over-bag pipeline: + + - person (general COCO yolo11n): full-figure boxes for realistic / Western art + the anime person-detector misses; NMS-merged with imgutils detect_person and + fed to CCIP (identity) + a concept crop. + - anatomy (booru_yolo): anime / furry / NSFW torso components (head, cat-head, + boob, hip, …) — concept crops aligned to the operator's tag vocabulary. + - panel (mosesb): a comic page → panel regions → concept crops. + +Each proposer is INDEPENDENTLY optional + guarded: a bad weight path or an +inference error disables just that proposer (logged) and never breaks the +worker, which still falls back to imgutils detection. Weights resolve from an +ultralytics builtin name ("yolo11n.pt"), an http(s) URL, or "hf_repo::file" — +cached under HF_HOME so the download happens once. +""" +import logging +import os +import threading +from pathlib import Path + +log = logging.getLogger("fc_agent.detectors") +_CACHE = Path(os.environ.get("HF_HOME", "/models")) / "yolo" + + +def _resolve(spec: str) -> str | None: + """A local weights path (downloading if needed) or an ultralytics builtin + name. None if the spec is empty/unresolvable.""" + if not spec: + return None + if "::" in spec: # hf_repo::filename + repo, _, fname = spec.partition("::") + from huggingface_hub import hf_hub_download + return hf_hub_download( + repo_id=repo, filename=fname, cache_dir=str(_CACHE) + ) + if spec.startswith(("http://", "https://")): + _CACHE.mkdir(parents=True, exist_ok=True) + dest = _CACHE / spec.rsplit("/", 1)[-1] + if not dest.is_file(): + import requests + r = requests.get(spec, timeout=300) + r.raise_for_status() + dest.write_bytes(r.content) + return str(dest) + return spec # ultralytics builtin name + + +def _iou(a, b) -> float: + ax, ay, aw, ah = a + bx, by, bw, bh = b + ix = max(0.0, min(ax + aw, bx + bw) - max(ax, bx)) + iy = max(0.0, min(ay + ah, by + bh) - max(ay, by)) + inter = ix * iy + union = aw * ah + bw * bh - inter + return inter / union if union > 0 else 0.0 + + +def nms_merge(boxes, iou_thresh: float = 0.6): + """Greedy NMS over (bbox_norm, score, label) from possibly several detectors, + so the same figure found by two of them collapses to one (higher-score) box.""" + kept = [] + for bb, sc, lb in sorted(boxes, key=lambda b: b[1], reverse=True): + if all(_iou(bb, k[0]) < iou_thresh for k in kept): + kept.append((bb, sc, lb)) + return kept + + +class YoloProposer: + """One lazily-loaded ultralytics YOLO. detect(image) → [(bbox_norm, score, + label)] with bbox normalized (x, y, w, h) in [0,1]. Self-disables on any + load/inference failure.""" + + def __init__(self, name, weights, conf=0.25, keep_labels=None): + self.name = name + self._spec = weights + self._conf = conf + self._keep = [k.lower() for k in keep_labels] if keep_labels else None + self._model = None + self._ok = True + self._lock = threading.Lock() + + def _load(self): + if self._model is not None or not self._ok: + return + with self._lock: + if self._model is not None or not self._ok: + return + try: + from ultralytics import YOLO + path = _resolve(self._spec) + if path is None: + self._ok = False + return + self._model = YOLO(path) + log.info("detector %s loaded (%s)", self.name, path) + except Exception as exc: # noqa: BLE001 + log.warning("detector %s disabled (load failed): %s", self.name, exc) + self._ok = False + + def detect(self, image): + self._load() + if self._model is None: + return [] + try: + res = self._model.predict(image, conf=self._conf, verbose=False)[0] + except Exception as exc: # noqa: BLE001 + log.warning("detector %s inference failed: %s", self.name, exc) + return [] + iw, ih = image.size + names = getattr(res, "names", None) or {} + out = [] + for b in res.boxes: + label = str(names.get(int(b.cls), int(b.cls))).lower() + if self._keep is not None and not any(k in label for k in self._keep): + continue + x0, y0, x1, y1 = (float(v) for v in b.xyxy[0].tolist()) + out.append(( + (x0 / iw, y0 / ih, (x1 - x0) / iw, (y1 - y0) / ih), + float(b.conf), label, + )) + return out + + +class Proposers: + """The agent's proposer set, built from config. Each detector is optional — + an empty weight spec leaves that proposer off.""" + + def __init__(self, cfg): + self.cfg = cfg + self._person = ( + YoloProposer("person-coco", cfg.person_weights, + conf=cfg.person_conf, keep_labels=["person"]) + if cfg.person_weights else None + ) + self._anatomy = ( + YoloProposer("anatomy", cfg.anatomy_weights, conf=cfg.anatomy_conf) + if cfg.anatomy_weights else None + ) + self._panel = ( + YoloProposer("panel", cfg.panel_weights, conf=cfg.panel_conf) + if cfg.panel_weights else None + ) + + def figures(self, image, base_boxes): + """Merge imgutils person boxes (base_boxes: [(bbox, score)]) with the + general COCO person detector → NMS'd figure boxes [(bbox, score, label)].""" + boxes = [(bb, sc if sc is not None else 1.0, "person") for bb, sc in base_boxes] + if self._person is not None: + boxes += self._person.detect(image) + return nms_merge(boxes) + + def components(self, image): + if self._anatomy is None: + return [] + items = sorted(self._anatomy.detect(image), key=lambda b: b[1], reverse=True) + return items[: self.cfg.max_components] + + def panels(self, image): + if self._panel is None: + return [] + items = sorted(self._panel.detect(image), key=lambda b: b[1], reverse=True) + return items[: self.cfg.max_panels] diff --git a/agent/fc_agent/worker.py b/agent/fc_agent/worker.py index 8b42277..2b7e412 100644 --- a/agent/fc_agent/worker.py +++ b/agent/fc_agent/worker.py @@ -11,6 +11,7 @@ orphaned work is re-picked at once rather than waiting out the lease. """ import threading +import numpy as np import requests from . import media, models @@ -75,6 +76,9 @@ class Worker: # needs it, from the model the server announces — one shared instance. self._embedder = None self._embedder_lock = threading.Lock() + # Region proposers (extra YOLO detectors) — lazily built once, shared. + self._proposers = None + self._proposers_lock = threading.Lock() # --- control ----------------------------------------------------------- def start(self): @@ -176,6 +180,15 @@ class Worker: self._embedder = CropEmbedder(model_name, self.cfg.embed_dtype) return self._embedder + def _ensure_proposers(self): + if self._proposers is not None: + return self._proposers + with self._proposers_lock: + if self._proposers is None: + from .detectors import Proposers + self._proposers = Proposers(self.cfg) + return self._proposers + def _process(self, job: dict) -> bool: """Process one job. Returns True when handled (completed, or hard-failed because the job itself is bad) and False on a TRANSPORT error (curator @@ -193,62 +206,101 @@ class Worker: else: frames = [(None, media.load_image(data))] + task = job.get("task") or "ccip" + embed_version = job.get("embed_version") or DEFAULT_EMBED_VERSION + model_name = ( + self.cfg.embed_model_override + or job.get("embed_model_name") + or DEFAULT_EMBED_MODEL + ) + + # 'embed' = WHOLE-IMAGE SigLIP embedding (re-embed the library under a + # new model, #1190) → image_record.siglip_embedding. Mean-pool video + # frames, matching the server's tag_and_embed. No regions. + if task == "embed": + embedder = self._ensure_embedder(model_name) + vecs = [embedder.embed(frame) for _, frame in frames] + if len(vecs) > 1: + vec = np.mean( + np.asarray(vecs, dtype=np.float32), axis=0 + ).tolist() + else: + vec = vecs[0] + self.client.submit_embedding(job["job_id"], vec, embed_version) + self._bump(processed=1) + return True + # task picks what to produce per crop: # 'siglip' (backfill existing images) → concept (SigLIP) regions # ONLY, so it never churns their figure/CCIP regions or the # character-reference cache. # 'ccip' / 'both' (a new image's first pass) → figure (CCIP) AND # concept (SigLIP) in one go, off the same crop. - task = job.get("task") or "ccip" want_ccip = task in ("ccip", "both") want_siglip = task in ("ccip", "siglip", "both") replace_kinds = ( - ["concept"] if task == "siglip" else ["figure", "face", "concept"] + ["concept", "panel"] if task == "siglip" + else ["figure", "face", "concept", "panel"] ) - - embed_version = job.get("embed_version") or DEFAULT_EMBED_VERSION - embedder = None - if want_siglip: - model_name = ( - self.cfg.embed_model_override - or job.get("embed_model_name") - or DEFAULT_EMBED_MODEL - ) - embedder = self._ensure_embedder(model_name) + embedder = self._ensure_embedder(model_name) if want_siglip else None + proposers = self._ensure_proposers() regions = [] ccip_ev = self.cfg.ccip_model or "ccip-default" dv = f"person-{self.cfg.detector_level}" + + def _concept(frame, bbox, t, score, detver, kind="concept"): + """A SigLIP region for one crop (None if below the size floor).""" + crop = crop_region(frame, bbox) + if crop is None: + return None + return { + "kind": kind, "bbox": list(bbox), "frame_time": t, + "score": score, "siglip_embedding": embedder.embed(crop), + "embedding_version": embed_version, "detector_version": detver, + } + for t, frame in frames: - figs = models.detect_figures(frame, self.cfg.detector_level) + # FIGURE boxes: imgutils detect_person ∪ general COCO person, + # NMS-merged → CCIP identity (+ a concept crop). Covers anime + + # Western/realistic figures. + base = models.detect_figures(frame, self.cfg.detector_level) + figs = proposers.figures(frame, base) if not figs: - figs = [((0.0, 0.0, 1.0, 1.0), None)] # whole-frame fallback - for bbox, score in figs: + figs = [((0.0, 0.0, 1.0, 1.0), 1.0, "whole")] # whole-frame fallback + for bbox, score, _label in figs: crop = crop_region(frame, bbox) if crop is None: continue if want_ccip: regions.append({ - "kind": "figure", - "bbox": list(bbox), - "frame_time": t, + "kind": "figure", "bbox": list(bbox), "frame_time": t, "score": score, "ccip_embedding": models.ccip_vector( crop, self.cfg.ccip_model or None ), - "embedding_version": ccip_ev, - "detector_version": dv, + "embedding_version": ccip_ev, "detector_version": dv, }) if want_siglip: regions.append({ - "kind": "concept", - "bbox": list(bbox), - "frame_time": t, + "kind": "concept", "bbox": list(bbox), "frame_time": t, "score": score, "siglip_embedding": embedder.embed(crop), - "embedding_version": embed_version, - "detector_version": dv, + "embedding_version": embed_version, "detector_version": dv, }) + if not want_siglip: + continue + # ANATOMY components (booru_yolo: head/cat-head/anatomy/…) → + # concept crops only (not full characters, so no CCIP). + for bbox, score, label in proposers.components(frame): + r = _concept(frame, bbox, t, score, f"booru:{label}") + if r is not None: + regions.append(r) + # PANEL crops (comic page → panels) → kind='panel' (still SigLIP). + for bbox, score, _label in proposers.panels(frame): + r = _concept(frame, bbox, t, score, "panel", kind="panel") + if r is not None: + regions.append(r) self.client.submit(job["job_id"], regions, replace_kinds) self._bump(processed=1) return True diff --git a/agent/requirements.txt b/agent/requirements.txt index 53f87bb..b71fdb0 100644 --- a/agent/requirements.txt +++ b/agent/requirements.txt @@ -7,6 +7,9 @@ onnxruntime-gpu # Dockerfile from the CUDA-12.4 wheel index so the GPU build is deterministic; # transformers loads whatever SigLIP-family model the server announces. transformers>=4.45 +# Crop PROPOSERS — small YOLO detectors (booru_yolo anatomy, COCO person, comic +# panel) that decide where to crop. Uses the torch already installed above. +ultralytics>=8.3 # Control surface + HTTP. fastapi uvicorn[standard] diff --git a/alembic/versions/0065_embedder_model_name.py b/alembic/versions/0065_embedder_model_name.py new file mode 100644 index 0000000..0a986b3 --- /dev/null +++ b/alembic/versions/0065_embedder_model_name.py @@ -0,0 +1,35 @@ +"""ml_settings: embedder_model_name (#1190 operator model swap) + +The embedder MODEL VERSION was already a setting (and stamps image_record. +siglip_model_version); the HF model NAME was env-only, so an operator couldn't +actually point the pipeline at a different embedder. Storing the name as a +setting makes the model an operator choice: set name + version → re-embed (the +GPU agent) → retrain heads. Default = the current SigLIP so400m. + +Revision ID: 0065 +Revises: 0064 +Create Date: 2026-06-30 +""" +from typing import Sequence, Union + +import sqlalchemy as sa +from alembic import op + +revision: str = "0065" +down_revision: Union[str, None] = "0064" +branch_labels: Union[str, Sequence[str], None] = None +depends_on: Union[str, Sequence[str], None] = None + + +def upgrade() -> None: + op.add_column( + "ml_settings", + sa.Column( + "embedder_model_name", sa.String(length=128), nullable=False, + server_default="google/siglip-so400m-patch14-384", + ), + ) + + +def downgrade() -> None: + op.drop_column("ml_settings", "embedder_model_name") diff --git a/alembic/versions/0066_drop_centroids.py b/alembic/versions/0066_drop_centroids.py new file mode 100644 index 0000000..d75a334 --- /dev/null +++ b/alembic/versions/0066_drop_centroids.py @@ -0,0 +1,57 @@ +"""drop the dead per-tag centroid subsystem (#1189 cleanup) + +The v2 pivot replaced per-tag SigLIP centroids with learned heads + CCIP. +Nothing read the centroids anymore — they were recomputed (on merge + a daily +beat) but never consumed for suggestions or auto-apply. Remove the storage + +its two now-unused settings columns. (The recompute tasks, beat, endpoint, +service, and UI card are removed in the same change.) + +Revision ID: 0066 +Revises: 0065 +Create Date: 2026-06-30 +""" +from typing import Sequence, Union + +import sqlalchemy as sa +from alembic import op + +revision: str = "0066" +down_revision: Union[str, None] = "0065" +branch_labels: Union[str, Sequence[str], None] = None +depends_on: Union[str, Sequence[str], None] = None + + +def upgrade() -> None: + op.drop_table("tag_reference_embedding") + op.drop_column("ml_settings", "centroid_similarity_threshold") + op.drop_column("ml_settings", "min_reference_images") + + +def downgrade() -> None: + op.add_column( + "ml_settings", + sa.Column( + "min_reference_images", sa.Integer(), nullable=False, + server_default="5", + ), + ) + op.add_column( + "ml_settings", + sa.Column( + "centroid_similarity_threshold", sa.Float(), nullable=False, + server_default="0.55", + ), + ) + op.create_table( + "tag_reference_embedding", + sa.Column("tag_id", sa.Integer(), nullable=False), + sa.Column("embedding", sa.LargeBinary(), nullable=False), + sa.Column("reference_count", sa.Integer(), nullable=False), + sa.Column("model_version", sa.String(length=128), nullable=False), + sa.Column( + "updated_at", sa.DateTime(timezone=True), + server_default=sa.func.now(), nullable=False, + ), + sa.ForeignKeyConstraint(["tag_id"], ["tag.id"], ondelete="CASCADE"), + sa.PrimaryKeyConstraint("tag_id"), + ) diff --git a/alembic/versions/0067_retire_camie_allowlist.py b/alembic/versions/0067_retire_camie_allowlist.py new file mode 100644 index 0000000..e3edd02 --- /dev/null +++ b/alembic/versions/0067_retire_camie_allowlist.py @@ -0,0 +1,66 @@ +"""retire the Camie tagger + allowlist bulk-apply (#1189) + +The v2 pivot made heads + CCIP the tag source and head auto-apply the earned +propagation. The Camie tagger ran only to feed the allowlist bulk-apply (its +predictions had no other consumer), and the allowlist was a second, un-earned +auto-apply path parallel to heads. Both are retired — drop their storage. + +(image_prediction = Camie's per-image predictions; tag_allowlist = the bulk- +apply allowlist. Nothing references INTO these tables, so the drop is clean.) + +Revision ID: 0067 +Revises: 0066 +Create Date: 2026-06-30 +""" +from typing import Sequence, Union + +import sqlalchemy as sa +from alembic import op + +revision: str = "0067" +down_revision: Union[str, None] = "0066" +branch_labels: Union[str, Sequence[str], None] = None +depends_on: Union[str, Sequence[str], None] = None + + +def upgrade() -> None: + op.drop_table("image_prediction") + op.drop_table("tag_allowlist") + + +def downgrade() -> None: + op.create_table( + "tag_allowlist", + sa.Column("tag_id", sa.Integer(), nullable=False), + sa.Column( + "min_confidence", sa.Float(), nullable=False, server_default="0.9" + ), + sa.Column( + "created_at", sa.DateTime(timezone=True), + server_default=sa.func.now(), nullable=False, + ), + sa.ForeignKeyConstraint(["tag_id"], ["tag.id"], ondelete="CASCADE"), + sa.PrimaryKeyConstraint("tag_id"), + sa.CheckConstraint( + "min_confidence >= 0 AND min_confidence <= 1", + name="ck_tag_allowlist_confidence_range", + ), + ) + op.create_table( + "image_prediction", + sa.Column("id", sa.Integer(), primary_key=True), + sa.Column("image_record_id", sa.Integer(), nullable=False), + sa.Column("raw_name", sa.String(length=255), nullable=False), + sa.Column("category", sa.String(length=32), nullable=False), + sa.Column("score", sa.Float(), nullable=False), + sa.ForeignKeyConstraint( + ["image_record_id"], ["image_record.id"], ondelete="CASCADE" + ), + ) + op.create_index( + "ix_image_prediction_image", "image_prediction", ["image_record_id"] + ) + op.create_index( + "ix_image_prediction_name_score", "image_prediction", + ["raw_name", "score"], + ) diff --git a/alembic/versions/0068_drop_dead_tagger_settings.py b/alembic/versions/0068_drop_dead_tagger_settings.py new file mode 100644 index 0000000..770676d --- /dev/null +++ b/alembic/versions/0068_drop_dead_tagger_settings.py @@ -0,0 +1,80 @@ +"""drop dead tagger/suggestion settings + columns left after Camie retirement (#1199) + +Hygiene follow-up to #1189. These were left inert to bound that change; nothing +reads them now: +- ml_settings: tagger_store_floor + tagger_model_version (only the deleted Camie + tagger used them), suggestion_threshold_character/general (already dead pre- + retirement — scoring uses per-head thresholds), video_min_tag_frames (only the + deleted video-prediction aggregator used it). +- image_record: tagger_model_version (no writer now), centroid_scores (long-dead + JSON cache, no reader). + +Revision ID: 0068 +Revises: 0067 +Create Date: 2026-06-30 +""" +from typing import Sequence, Union + +import sqlalchemy as sa +from alembic import op + +revision: str = "0068" +down_revision: Union[str, None] = "0067" +branch_labels: Union[str, Sequence[str], None] = None +depends_on: Union[str, Sequence[str], None] = None + + +def upgrade() -> None: + op.drop_column("ml_settings", "suggestion_threshold_character") + op.drop_column("ml_settings", "suggestion_threshold_general") + op.drop_column("ml_settings", "tagger_store_floor") + op.drop_column("ml_settings", "video_min_tag_frames") + op.drop_column("ml_settings", "tagger_model_version") + op.drop_column("image_record", "tagger_model_version") + op.drop_column("image_record", "centroid_scores") + + +def downgrade() -> None: + op.add_column( + "image_record", + sa.Column("centroid_scores", sa.JSON(), nullable=True), + ) + op.add_column( + "image_record", + sa.Column("tagger_model_version", sa.String(length=128), nullable=True), + ) + op.add_column( + "ml_settings", + sa.Column( + "tagger_model_version", sa.String(length=128), nullable=False, + server_default="camie-tagger-v2", + ), + ) + op.add_column( + "ml_settings", + sa.Column( + "video_min_tag_frames", sa.Integer(), nullable=False, + server_default="3", + ), + ) + op.add_column( + "ml_settings", + sa.Column( + "tagger_store_floor", sa.Float(), nullable=False, + server_default="0.7", + ), + ) + op.add_column( + "ml_settings", + sa.Column( + "suggestion_threshold_general", sa.Float(), nullable=False, + server_default="0.7", + ), + ) + op.add_column( + "ml_settings", + sa.Column( + "suggestion_threshold_character", sa.Float(), nullable=False, + server_default="0.7", + ), + ) diff --git a/backend/app/api/__init__.py b/backend/app/api/__init__.py index cdc1442..8345f2d 100644 --- a/backend/app/api/__init__.py +++ b/backend/app/api/__init__.py @@ -16,7 +16,6 @@ api_bp.add_url_rule("/health", view_func=health.get_health, methods=["GET"]) def all_blueprints() -> list[Blueprint]: from .admin import admin_bp from .aliases import aliases_bp - from .allowlist import allowlist_bp from .artist import artist_bp from .artists import artists_bp from .attachments import attachments_bp @@ -58,7 +57,6 @@ def all_blueprints() -> list[Blueprint]: cleanup_bp, import_admin_bp, suggestions_bp, - allowlist_bp, aliases_bp, tag_eval_bp, heads_bp, diff --git a/backend/app/api/allowlist.py b/backend/app/api/allowlist.py deleted file mode 100644 index 31241de..0000000 --- a/backend/app/api/allowlist.py +++ /dev/null @@ -1,84 +0,0 @@ -"""Allowlist API: list, adjust threshold, remove.""" - -from quart import Blueprint, jsonify, request - -from ..extensions import get_session -from ..models import TagAllowlist -from ..services.ml.allowlist import AllowlistService - -allowlist_bp = Blueprint("allowlist", __name__, url_prefix="/api") - - -@allowlist_bp.route("/allowlist", methods=["GET"]) -async def list_allowlist(): - async with get_session() as session: - rows = await AllowlistService(session).list_all() - return jsonify( - [ - { - "tag_id": r.tag_id, - "tag_name": r.tag_name, - "tag_kind": r.tag_kind, - "min_confidence": r.min_confidence, - "applied_count": r.applied_count, - "coverage_count": r.coverage_count, - } - for r in rows - ] - ) - - -@allowlist_bp.route("/tags//allowlist/coverage", methods=["GET"]) -async def coverage(tag_id: int): - """Live "at threshold T, a sweep would cover ~N images" projection for the - allowlist tuning dashboard. Defaults to the tag's stored threshold.""" - raw = request.args.get("threshold") - async with get_session() as session: - svc = AllowlistService(session) - if raw is not None: - try: - threshold = float(raw) - except ValueError: - return jsonify({"error": "threshold must be a float"}), 400 - if not (0 < threshold <= 1): - return jsonify({"error": "threshold must be in (0, 1]"}), 400 - else: - row = await session.get(TagAllowlist, tag_id) - if row is None: - return jsonify({"error": "not on allowlist"}), 404 - threshold = row.min_confidence - count = await svc.coverage(tag_id, threshold) - return jsonify({"count": count, "threshold": threshold}) - - -@allowlist_bp.route("/tags//allowlist", methods=["GET"]) -async def get_one(tag_id: int): - async with get_session() as session: - row = await session.get(TagAllowlist, tag_id) - if row is None: - return jsonify({"error": "not on allowlist"}), 404 - return jsonify( - {"min_confidence": row.min_confidence, "added_at": row.added_at.isoformat()} - ) - - -@allowlist_bp.route("/tags//allowlist", methods=["PATCH"]) -async def patch_threshold(tag_id: int): - body = await request.get_json() - if not body or "min_confidence" not in body: - return jsonify({"error": "min_confidence required"}), 400 - mc = float(body["min_confidence"]) - if not (0 < mc <= 1): - return jsonify({"error": "min_confidence must be in (0, 1]"}), 400 - async with get_session() as session: - await AllowlistService(session).update_threshold(tag_id, mc) - await session.commit() - return "", 204 - - -@allowlist_bp.route("/tags//allowlist", methods=["DELETE"]) -async def remove(tag_id: int): - async with get_session() as session: - await AllowlistService(session).remove(tag_id) - await session.commit() - return "", 204 diff --git a/backend/app/api/gpu.py b/backend/app/api/gpu.py index cf8fc73..09a488f 100644 --- a/backend/app/api/gpu.py +++ b/backend/app/api/gpu.py @@ -17,7 +17,6 @@ from sqlalchemy.dialects.postgresql import insert as pg_insert from ..extensions import get_session from ..models import AppSetting, GpuJob, ImageRecord, MLSettings from ..services.gallery_service import image_url -from ..services.ml.embedder import MODEL_NAME as EMBED_MODEL_NAME from ..services.ml.gpu_jobs import GpuJobService from ..services.ml.regions import RegionService @@ -138,11 +137,12 @@ async def lease(): # For video/animated: the agent samples at this cadence. "frame_interval_seconds": ml.video_frame_interval_seconds, "max_frames": ml.video_max_frames, - # The embedding model the agent must use for concept crops, so - # its region vectors land in the SAME space the heads trained in. - # Server-announced → the agent stays model-agnostic; a swap is a - # server setting + a re-embed migration, never an agent change. - "embed_model_name": EMBED_MODEL_NAME, + # The embedding model the agent must use for concept crops + the + # whole-image 'embed' task, so its vectors land in the SAME space + # the heads trained in. Server-announced FROM THE SETTING → the + # agent stays model-agnostic; an operator swap is a setting + a + # re-embed, never an agent change. + "embed_model_name": ml.embedder_model_name, "embed_version": ml.embedder_model_version, }) return jsonify({"jobs": out}) @@ -188,6 +188,34 @@ async def submit(): return jsonify({"ok": True, "stored": len(regions)}) +@gpu_bp.route("/jobs/submit_embedding", methods=["POST"]) +async def submit_embedding(): + """Store a whole-image SigLIP embedding (the 'embed' task) on image_record + + close the job. Body: {agent_id, job_id, embedding:[...], embedding_version}. + This is how the GPU agent re-embeds the library under a new model (#1190) — + much faster than the CPU ml-worker at higher resolutions.""" + body = await request.get_json(silent=True) or {} + agent_id = str(body.get("agent_id") or "agent") + job_id = body.get("job_id") + embedding = body.get("embedding") + version = body.get("embedding_version") + if job_id is None or not embedding or not version: + return jsonify({"error": "job_id, embedding, embedding_version required"}), 400 + async with get_session() as session: + if not await _agent_authed(session): + return jsonify({"error": "unauthorized"}), 401 + job = await session.get(GpuJob, int(job_id)) + if job is None or job.status != "leased" or job.lease_token != agent_id: + return jsonify({"error": "lease_invalid"}), 409 + img = await session.get(ImageRecord, job.image_record_id) + if img is not None: + img.siglip_embedding = embedding + img.siglip_model_version = version + await GpuJobService(session).complete(agent_id, int(job_id)) + await session.commit() + return jsonify({"ok": True}) + + @gpu_bp.route("/jobs/fail", methods=["POST"]) async def fail(): body = await request.get_json(silent=True) or {} diff --git a/backend/app/api/ml_admin.py b/backend/app/api/ml_admin.py index 1472770..054fa69 100644 --- a/backend/app/api/ml_admin.py +++ b/backend/app/api/ml_admin.py @@ -1,4 +1,4 @@ -"""ML admin API: settings, backfill trigger, centroid recompute trigger.""" +"""ML admin API: settings + backfill trigger.""" from quart import Blueprint, jsonify, request @@ -9,14 +9,8 @@ ml_admin_bp = Blueprint("ml_admin", __name__, url_prefix="/api/ml") _EDITABLE = ( - "suggestion_threshold_character", - "suggestion_threshold_general", - "centroid_similarity_threshold", - "min_reference_images", - "tagger_store_floor", "video_frame_interval_seconds", "video_max_frames", - "video_min_tag_frames", "head_min_positives", "head_auto_apply_precision", "head_auto_apply_enabled", @@ -24,6 +18,8 @@ _EDITABLE = ( "ccip_match_threshold", "ccip_auto_apply_enabled", "ccip_auto_apply_threshold", + "embedder_model_name", + "embedder_model_version", ) @@ -37,15 +33,8 @@ async def get_settings(): ).scalar_one() return jsonify( { - "suggestion_threshold_character": s.suggestion_threshold_character, - "suggestion_threshold_general": s.suggestion_threshold_general, - "centroid_similarity_threshold": s.centroid_similarity_threshold, - "min_reference_images": s.min_reference_images, - "tagger_store_floor": s.tagger_store_floor, "video_frame_interval_seconds": s.video_frame_interval_seconds, "video_max_frames": s.video_max_frames, - "video_min_tag_frames": s.video_min_tag_frames, - "tagger_model_version": s.tagger_model_version, "embedder_model_version": s.embedder_model_version, "head_min_positives": s.head_min_positives, "head_auto_apply_precision": s.head_auto_apply_precision, @@ -54,6 +43,7 @@ async def get_settings(): "ccip_match_threshold": s.ccip_match_threshold, "ccip_auto_apply_enabled": s.ccip_auto_apply_enabled, "ccip_auto_apply_threshold": s.ccip_auto_apply_threshold, + "embedder_model_name": s.embedder_model_name, } ) @@ -89,31 +79,12 @@ async def patch_settings(): def _validate(p: dict) -> str | None: - """Returns an error string if the proposed settings are invalid, else None. - - Invariant (plan-task #764): the per-category suggestion thresholds can't - drop below tagger_store_floor — nothing below the floor is stored, so a - lower threshold would silently surface nothing in that gap. The UI clamps - the sliders to the floor; this is the server-side backstop. - """ - floor = p["tagger_store_floor"] - if not (0.0 <= floor <= 1.0): - return "tagger_store_floor must be between 0 and 1" - for cat in ("character", "general"): - if p[f"suggestion_threshold_{cat}"] < floor: - return ( - f"suggestion_threshold_{cat} cannot be below tagger_store_floor " - f"({floor}) — predictions below the floor are not stored" - ) - # Video tagging (#747). + """Returns an error string if the proposed settings are invalid, else None.""" + # Video embedding (#747). if p["video_frame_interval_seconds"] <= 0: return "video_frame_interval_seconds must be > 0" if p["video_max_frames"] < 1: return "video_max_frames must be >= 1" - if p["video_min_tag_frames"] < 1: - return "video_min_tag_frames must be >= 1" - if p["video_min_tag_frames"] > p["video_max_frames"]: - return "video_min_tag_frames cannot exceed video_max_frames" # Head training (#114). if int(p["head_min_positives"]) < 1: return "head_min_positives must be >= 1" @@ -125,6 +96,11 @@ def _validate(p: dict) -> str | None: return "ccip_match_threshold must be between 0.5 and 0.999" if not (0.5 <= float(p["ccip_auto_apply_threshold"]) <= 0.999): return "ccip_auto_apply_threshold must be between 0.5 and 0.999" + # Embedder model swap (#1190): both must be non-empty. Changing them means a + # different embedding space — the operator must re-embed + retrain after. + for key in ("embedder_model_name", "embedder_model_version"): + if not str(p[key]).strip(): + return f"{key} must not be empty" return None @@ -134,11 +110,3 @@ async def trigger_backfill(): r = backfill.delay() return jsonify({"celery_task_id": r.id}), 202 - - -@ml_admin_bp.route("/recompute-centroids", methods=["POST"]) -async def trigger_recompute(): - from ..tasks.ml import recompute_centroids - - r = recompute_centroids.delay() - return jsonify({"celery_task_id": r.id}), 202 diff --git a/backend/app/api/suggestions.py b/backend/app/api/suggestions.py index 7a5ca1d..cf27125 100644 --- a/backend/app/api/suggestions.py +++ b/backend/app/api/suggestions.py @@ -3,31 +3,12 @@ from quart import Blueprint, jsonify, request from ..extensions import get_session -from ..models import Tag, TagAllowlist from ..services.ml.allowlist import AllowlistService from ..services.ml.suggestions import SuggestionService suggestions_bp = Blueprint("suggestions", __name__, url_prefix="/api") -async def _accept_payload(session, svc, newly_added: bool, tag_id: int) -> dict: - """Shape the accept/alias response. When accepting newly allowlists a tag, - include the coverage PROJECTION (at the tag's threshold) so the UI can show - a non-blocking "auto-applying to ~N images" toast — the actual apply runs - async via apply_allowlist_tags, so this is an estimate, not a post-hoc - count (#7).""" - payload = {"allowlisted": newly_added} - if newly_added: - tag = await session.get(Tag, tag_id) - row = await session.get(TagAllowlist, tag_id) - payload["tag_id"] = tag_id - payload["tag_name"] = tag.name if tag is not None else None - payload["projected_count"] = await svc.coverage( - tag_id, row.min_confidence if row is not None else 0.90, - ) - return payload - - @suggestions_bp.route("/images//suggestions", methods=["GET"]) async def get_suggestions(image_id: int): # ?min= overrides the configured per-category thresholds so the typed @@ -83,15 +64,9 @@ async def accept_suggestion(image_id: int): return jsonify({"error": "tag_id required"}), 400 tag_id = body["tag_id"] async with get_session() as session: - svc = AllowlistService(session) - newly_added = await svc.accept(image_id, tag_id) - payload = await _accept_payload(session, svc, newly_added, tag_id) + await AllowlistService(session).accept(image_id, tag_id) await session.commit() - if newly_added: - from ..tasks.ml import apply_allowlist_tags - - apply_allowlist_tags.delay(tag_id=tag_id) - return jsonify(payload) + return jsonify({"accepted": True, "tag_id": tag_id}) @suggestions_bp.route( @@ -104,22 +79,14 @@ async def alias_suggestion(image_id: int): return jsonify({"error": f"required: {sorted(required)}"}), 400 canonical_tag_id = body["canonical_tag_id"] async with get_session() as session: - svc = AllowlistService(session) - newly_added = await svc.add_alias_and_accept( + await AllowlistService(session).add_alias_and_accept( image_id, body["alias_string"], body["alias_category"], canonical_tag_id, ) - payload = await _accept_payload( - session, svc, newly_added, canonical_tag_id, - ) await session.commit() - if newly_added: - from ..tasks.ml import apply_allowlist_tags - - apply_allowlist_tags.delay(tag_id=canonical_tag_id) - return jsonify(payload) + return jsonify({"accepted": True, "tag_id": canonical_tag_id}) @suggestions_bp.route( diff --git a/backend/app/api/tags.py b/backend/app/api/tags.py index 23b0817..ec9990a 100644 --- a/backend/app/api/tags.py +++ b/backend/app/api/tags.py @@ -1,13 +1,14 @@ """Tags API: autocomplete, create, list/add/remove for an image.""" from quart import Blueprint, jsonify, request -from sqlalchemy import exists, select +from sqlalchemy import func, select from sqlalchemy.dialects.postgresql import insert as pg_insert from sqlalchemy.exc import IntegrityError from ..extensions import get_session -from ..models import Tag, TagKind, TagPositiveConfirmation -from ..models.tag_allowlist import TagAllowlist +from ..models import Tag, TagHead, TagKind, TagPositiveConfirmation +from ..models.tag import image_tag +from ..models.tag_suggestion_rejection import TagSuggestionRejection from ..services.bulk_tag_service import BulkTagService from ..services.ml.aliases import AliasService from ..services.series_match_service import SeriesMatchService @@ -61,6 +62,117 @@ def _parse_bulk_ids( return ids, None +# Application-source groupings (image_tag.source). HUMAN = operator signal; +# AUTO = machine-applied (heads/CCIP, + legacy Camie ml_auto). +_SOURCE_GROUPS = { + "human": ("manual", "ml_accepted"), + "manual": ("manual",), + "accepted": ("ml_accepted",), + "auto": ("head_auto", "ccip_auto", "ml_auto"), +} + + +@tags_bp.route("/tags/top", methods=["GET"]) +async def tags_top(): + """Top tags by image count — a fast indexed aggregate for ANALYSIS (not the + paged UI directory, which is alphabetical + builds previews). Params: + ?kind=general|character|fandom|… ?source=all|human|manual|accepted|auto + ?limit=50 (cap 500) ?min_count=N. → {tags:[{tag_id,name,kind,count}]} desc.""" + kind = _coerce_kind(request.args.get("kind")) + try: + limit = min(max(int(request.args.get("limit", "50")), 1), 500) + except ValueError: + return jsonify({"error": "limit must be an integer"}), 400 + min_count = None + if "min_count" in request.args: + try: + min_count = int(request.args["min_count"]) + except ValueError: + return jsonify({"error": "min_count must be an integer"}), 400 + src_vals = _SOURCE_GROUPS.get((request.args.get("source") or "all").lower()) + + cnt = func.count(image_tag.c.image_record_id) + stmt = ( + select(Tag.id, Tag.name, Tag.kind, cnt.label("count")) + .select_from(Tag) + .join(image_tag, image_tag.c.tag_id == Tag.id) + .group_by(Tag.id, Tag.name, Tag.kind) + .order_by(cnt.desc(), Tag.name.asc()) + .limit(limit) + ) + if kind is not None: + stmt = stmt.where(Tag.kind == kind) + if src_vals is not None: + stmt = stmt.where(image_tag.c.source.in_(src_vals)) + if min_count is not None: + stmt = stmt.having(cnt >= min_count) + async with get_session() as session: + rows = (await session.execute(stmt)).all() + return jsonify({"tags": [ + { + "tag_id": r.id, "name": r.name, + "kind": r.kind.value if hasattr(r.kind, "value") else str(r.kind), + "count": r.count, + } + for r in rows + ]}) + + +@tags_bp.route("/tags//stats", methods=["GET"]) +async def tag_stats(tag_id: int): + """Per-tag dataset health: total + per-source application counts (human vs + machine), rejection count, and whether a trained head exists. Read-only, + analysis-shaped — backs concept-readiness + source-split decisions.""" + async with get_session() as session: + tag = await session.get(Tag, tag_id) + if tag is None: + return jsonify({"error": "not found"}), 404 + by_source = dict( + ( + await session.execute( + select(image_tag.c.source, func.count()) + .where(image_tag.c.tag_id == tag_id) + .group_by(image_tag.c.source) + ) + ).all() + ) + rejected = ( + await session.execute( + select(func.count()) + .select_from(TagSuggestionRejection) + .where(TagSuggestionRejection.tag_id == tag_id) + ) + ).scalar_one() + has_head = ( + await session.execute( + select(func.count()) + .select_from(TagHead) + .where(TagHead.tag_id == tag_id) + ) + ).scalar_one() > 0 + human = by_source.get("manual", 0) + by_source.get("ml_accepted", 0) + auto = ( + by_source.get("head_auto", 0) + + by_source.get("ccip_auto", 0) + + by_source.get("ml_auto", 0) + ) + return jsonify({ + "tag_id": tag_id, + "name": tag.name, + "kind": tag.kind.value if hasattr(tag.kind, "value") else str(tag.kind), + "count_total": sum(by_source.values()), + "count_human": human, + "count_manual": by_source.get("manual", 0), + "count_accepted": by_source.get("ml_accepted", 0), + "count_auto": auto, + "count_head_auto": by_source.get("head_auto", 0), + "count_ccip_auto": by_source.get("ccip_auto", 0), + "count_rejected": rejected, + "by_source": by_source, + "has_head": has_head, + }) + + @tags_bp.route("/tags/autocomplete", methods=["GET"]) async def autocomplete(): q = request.args.get("q", "") @@ -297,19 +409,6 @@ async def merge_tag(source_id: int): status = 404 if "not found" in msg else 400 return jsonify({"error": msg}), status await session.commit() - target_allowlisted = await session.scalar( - select(exists().where(TagAllowlist.tag_id == result.target_id)) - ) - if target_allowlisted: - from ..tasks.ml import apply_allowlist_tags - - apply_allowlist_tags.delay(tag_id=result.target_id) - # Tag merge invalidates the target's centroid (the merged-in source - # tag's images now contribute to it). Daily list_drifted catches it - # within 24h, but eager recompute closes the suggestion-quality dip - # in the meantime. Audit 2026-06-02. - from ..tasks.ml import recompute_centroid - recompute_centroid.delay(result.target_id) return jsonify( { "target": { diff --git a/backend/app/celery_app.py b/backend/app/celery_app.py index 5778a1e..811b390 100644 --- a/backend/app/celery_app.py +++ b/backend/app/celery_app.py @@ -101,14 +101,6 @@ def make_celery() -> Celery: "task": "backend.app.tasks.ml.backfill", "schedule": 86400.0, }, - "recompute-centroids-daily": { - "task": "backend.app.tasks.ml.recompute_centroids", - "schedule": 86400.0, - }, - "apply-allowlist-sweep-daily": { - "task": "backend.app.tasks.ml.apply_allowlist_tags", - "schedule": 86400.0, - }, "train-heads-nightly": { "task": "backend.app.tasks.ml.scheduled_train_heads", "schedule": 86400.0, # passive cadence; manual retrain stays available @@ -131,6 +123,11 @@ def make_celery() -> Celery: "schedule": 86400.0, # drain the concept-crop back-catalogue + "args": ("siglip",), # retry failed embeds, no button needed }, + "enqueue-embed-backfill-daily": { + "task": "backend.app.tasks.ml.enqueue_gpu_backfill", + "schedule": 86400.0, # whole-image re-embed under the current + "args": ("embed",), # model (an operator swap) drains via agent + }, "ccip-auto-apply-daily": { "task": "backend.app.tasks.ml.scheduled_ccip_auto_apply", "schedule": 86400.0, # no-op unless ccip_auto_apply_enabled diff --git a/backend/app/models/__init__.py b/backend/app/models/__init__.py index 087ff0b..8707467 100644 --- a/backend/app/models/__init__.py +++ b/backend/app/models/__init__.py @@ -13,7 +13,6 @@ from .head_auto_apply_run import HeadAutoApplyRun from .head_metric import HeadMetric from .head_metrics_snapshot import HeadMetricsSnapshot from .head_training_run import HeadTrainingRun -from .image_prediction import ImagePrediction from .image_provenance import ImageProvenance from .image_record import ImageRecord from .image_region import ImageRegion @@ -34,11 +33,9 @@ from .subscribestar_failed_media import SubscribeStarFailedMedia from .subscribestar_seen_media import SubscribeStarSeenMedia from .tag import Tag, TagKind, image_tag from .tag_alias import TagAlias -from .tag_allowlist import TagAllowlist from .tag_eval_run import TagEvalRun from .tag_head import TagHead from .tag_positive_confirmation import TagPositiveConfirmation -from .tag_reference_embedding import TagReferenceEmbedding from .tag_suggestion_rejection import TagSuggestionRejection from .task_run import TaskRun @@ -60,7 +57,6 @@ __all__ = [ "SeriesPage", "SeriesSuggestion", "ImageRecord", - "ImagePrediction", "ImageProvenance", "ImageRegion", "Tag", @@ -79,11 +75,9 @@ __all__ = [ "HeadMetricsSnapshot", "HeadTrainingRun", "TagAlias", - "TagAllowlist", "TagEvalRun", "TagHead", "TagPositiveConfirmation", - "TagReferenceEmbedding", "TagSuggestionRejection", "TaskRun", ] diff --git a/backend/app/models/image_prediction.py b/backend/app/models/image_prediction.py deleted file mode 100644 index f532243..0000000 --- a/backend/app/models/image_prediction.py +++ /dev/null @@ -1,37 +0,0 @@ -"""ImagePrediction — one row per (image, tagger vocab prediction). - -Replaces the image_record.tagger_predictions JSON blob (#768). Storing the -raw Camie/booru vocab name (not a tag_id) preserves the suggestion read -path's semantics: raw_name → canonical Tag resolution happens at read time -via the alias map, and accepting a prediction can CREATE the Tag. The store -floor (ml_settings.tagger_store_floor) is applied at WRITE time, so only -predictions >= the floor land here. -""" - -from sqlalchemy import Float, ForeignKey, Index, String, UniqueConstraint -from sqlalchemy.orm import Mapped, mapped_column - -from .base import Base - - -class ImagePrediction(Base): - __tablename__ = "image_prediction" - __table_args__ = ( - UniqueConstraint( - "image_record_id", "raw_name", name="image_raw_name", - ), - # Per-image read (suggestion build) and the "images with tag X above - # Y" query the JSON blob never allowed. - Index("ix_image_prediction_image", "image_record_id"), - Index("ix_image_prediction_name_score", "raw_name", "score"), - ) - - id: Mapped[int] = mapped_column(primary_key=True) - image_record_id: Mapped[int] = mapped_column( - ForeignKey("image_record.id", ondelete="CASCADE"), nullable=False, - ) - # The raw tagger vocab key (booru form) — NOT a tag_id. Resolved to a - # canonical Tag at read time, exactly as the old JSON keys were. - raw_name: Mapped[str] = mapped_column(String(255), nullable=False) - category: Mapped[str] = mapped_column(String(64), nullable=False) - score: Mapped[float] = mapped_column(Float, nullable=False) diff --git a/backend/app/models/image_record.py b/backend/app/models/image_record.py index b3ba120..fc459dd 100644 --- a/backend/app/models/image_record.py +++ b/backend/app/models/image_record.py @@ -9,7 +9,6 @@ from datetime import datetime from pgvector.sqlalchemy import Vector from sqlalchemy import ( - JSON, BigInteger, DateTime, Enum, @@ -77,19 +76,13 @@ class ImageRecord(Base): ForeignKey("artist.id", ondelete="SET NULL"), nullable=True, index=True ) - # ML fields (populated by FC-2's ml-worker). Per-tag predictions live in the - # normalized image_prediction table (#768) — the tagger_predictions JSON - # column was dropped in migration 0046. tagger_model_version stays as the - # "has this been tagged / is it current?" signal the backfill sweep reads. - tagger_model_version: Mapped[str | None] = mapped_column(String(128), nullable=True) - # 1152 = SigLIP-so400m embedding dim. Swapping models in FC-2 may require - # a column-width migration. + # ML fields (populated by the ml-worker / GPU agent). 1152 = SigLIP-so400m + # embedding dim; siglip_model_version stamps which model produced it (so an + # operator model swap, #1190, can re-embed the stale rows). A different-dim + # model would need a column-width migration. siglip_embedding: Mapped[list[float] | None] = mapped_column(Vector(1152), nullable=True) siglip_model_version: Mapped[str | None] = mapped_column(String(128), nullable=True) - # Centroid score cache (populated post-tagging) - centroid_scores: Mapped[dict | None] = mapped_column(JSON, nullable=True) - created_at: Mapped[datetime] = mapped_column( DateTime(timezone=True), nullable=False, server_default=func.now() ) diff --git a/backend/app/models/image_region.py b/backend/app/models/image_region.py index 58cfa8c..7ae6d4d 100644 --- a/backend/app/models/image_region.py +++ b/backend/app/models/image_region.py @@ -31,7 +31,10 @@ class ImageRegion(Base): ForeignKey("image_record.id", ondelete="CASCADE"), index=True ) # 'frame' (a whole video frame → SigLIP bag) | 'face' | 'figure' (→ CCIP - # character id) | 'concept' (→ SigLIP head bag). + # character id) | 'concept' (→ SigLIP head bag) | 'panel' (a comic panel crop, + # also SigLIP → the bag). Free String, not an enum — proposers can add kinds + # without a migration; the bag scorer keys on a non-null siglip_embedding, not + # the kind, so any SigLIP-embedded region joins the bag. kind: Mapped[str] = mapped_column(String(16), nullable=False) # For video/animated media: the source frame's timestamp in SECONDS. NULL for # static images. Lets a video be a BAG of per-frame instances (fixes the diff --git a/backend/app/models/ml_settings.py b/backend/app/models/ml_settings.py index 9825568..1d7c729 100644 --- a/backend/app/models/ml_settings.py +++ b/backend/app/models/ml_settings.py @@ -23,46 +23,16 @@ class MLSettings(Base): __table_args__ = (CheckConstraint("id = 1", name="singleton"),) id: Mapped[int] = mapped_column(Integer, primary_key=True) - suggestion_threshold_character: Mapped[float] = mapped_column( - Float, nullable=False, default=0.70 - ) - # Default raised 0.50 → 0.70 on 2026-06-02 — operator-flagged 0.50 - # surfaced too many low-confidence picks; 0.70 keeps the rail - # signal-rich while still surfacing more than the original 0.95 - # which hid almost everything. Operator-tunable via Settings → ML. - suggestion_threshold_general: Mapped[float] = mapped_column( - Float, nullable=False, default=0.70 - ) - centroid_similarity_threshold: Mapped[float] = mapped_column( - Float, nullable=False, default=0.55 - ) - # Ingest floor: tagger predictions below this confidence are not stored - # (tagger.Tagger.infer). Default 0.70 — the suggestion path already - # filters at 0.70 and the centroid/learned path covers low-confidence - # preferred tags, so the sub-0.70 tail is redundant weight (it had - # bloated image_record's TOAST to ~100 GB; plan-task #764). Operator- - # tunable via Settings → ML; must stay ≤ the suggestion thresholds. - tagger_store_floor: Mapped[float] = mapped_column( - Float, nullable=False, default=0.70 - ) - min_reference_images: Mapped[int] = mapped_column( - Integer, nullable=False, default=5 - ) - # Video tagging (#747). Sample one frame every N seconds (fixed CADENCE, not a - # fixed count) so a tag's frame-presence reflects real screen time regardless - # of video length; cap the total so a long video can't explode into hundreds - # of inferences (the cadence stretches past the cap). A tag is kept only if it - # appears in >= video_min_tag_frames sampled frames (≈ that many × interval - # seconds on screen) — duration-independent noise rejection. Operator-tunable. + # Video embedding (#747). Sample one frame every N seconds (fixed CADENCE, not + # a fixed count) so coverage reflects real screen time regardless of length; + # cap the total so a long video can't explode into hundreds of embeds. The + # per-frame SigLIP embeddings are mean-pooled. Operator-tunable. video_frame_interval_seconds: Mapped[float] = mapped_column( Float, nullable=False, default=4.0 ) video_max_frames: Mapped[int] = mapped_column( Integer, nullable=False, default=64 ) - video_min_tag_frames: Mapped[int] = mapped_column( - Integer, nullable=False, default=3 - ) # Tagging-v2 head training (#114). The head is the suggestion source that # LEARNS from the operator's tags (replacing Camie + centroid). A concept # needs >= head_min_positives labelled images before a head is trained; @@ -101,12 +71,15 @@ class MLSettings(Base): ccip_auto_apply_threshold: Mapped[float] = mapped_column( Float, nullable=False, default=0.92 ) - tagger_model_version: Mapped[str] = mapped_column( - String(128), nullable=False, default="camie-tagger-v2" - ) embedder_model_version: Mapped[str] = mapped_column( String(128), nullable=False, default="siglip-so400m-patch14-384" ) + # The HF model NAME the embedder loads (server CPU embed + announced to the + # GPU agent in the lease). Operator-settable so the embedder is a choice, not + # a hardcode (#1190): set name + version together, then re-embed + retrain. + embedder_model_name: Mapped[str] = mapped_column( + String(128), nullable=False, default="google/siglip-so400m-patch14-384" + ) updated_at: Mapped[datetime] = mapped_column( DateTime(timezone=True), nullable=False, server_default=func.now() ) diff --git a/backend/app/models/tag_allowlist.py b/backend/app/models/tag_allowlist.py deleted file mode 100644 index 3bbbc7a..0000000 --- a/backend/app/models/tag_allowlist.py +++ /dev/null @@ -1,32 +0,0 @@ -"""TagAllowlist — tags the operator opted-in to auto-apply via maintenance.""" - -from datetime import datetime - -from sqlalchemy import CheckConstraint, DateTime, Float, ForeignKey, func -from sqlalchemy.orm import Mapped, mapped_column - -from .base import Base - - -class TagAllowlist(Base): - __tablename__ = "tag_allowlist" - # Bare name — Base.metadata's naming convention prepends ck__, - # producing the final ck_tag_allowlist_confidence_range (matches migration 0003). - __table_args__ = ( - CheckConstraint( - "min_confidence > 0 AND min_confidence <= 1", - name="confidence_range", - ), - ) - - tag_id: Mapped[int] = mapped_column( - ForeignKey("tag.id", ondelete="CASCADE"), primary_key=True - ) - # Default auto-apply threshold for a newly-accepted tag. 0.90 (lowered from - # 0.95 on operator evidence 2026-06-07: 0.95 was too strict and skipped - # confident-enough applications). Per-tag value is still tunable in the - # allowlist table; existing rows keep whatever they were stored with. - min_confidence: Mapped[float] = mapped_column(Float, nullable=False, default=0.90) - added_at: Mapped[datetime] = mapped_column( - DateTime(timezone=True), nullable=False, server_default=func.now() - ) diff --git a/backend/app/models/tag_reference_embedding.py b/backend/app/models/tag_reference_embedding.py deleted file mode 100644 index 2dc841b..0000000 --- a/backend/app/models/tag_reference_embedding.py +++ /dev/null @@ -1,23 +0,0 @@ -"""TagReferenceEmbedding — per-tag centroid (mean SigLIP embedding of members).""" - -from datetime import datetime - -from pgvector.sqlalchemy import Vector -from sqlalchemy import DateTime, ForeignKey, Integer, String, func -from sqlalchemy.orm import Mapped, mapped_column - -from .base import Base - - -class TagReferenceEmbedding(Base): - __tablename__ = "tag_reference_embedding" - - tag_id: Mapped[int] = mapped_column( - ForeignKey("tag.id", ondelete="CASCADE"), primary_key=True - ) - embedding: Mapped[list[float]] = mapped_column(Vector(1152), nullable=False) - reference_count: Mapped[int] = mapped_column(Integer, nullable=False) - model_version: Mapped[str] = mapped_column(String(128), nullable=False) - updated_at: Mapped[datetime] = mapped_column( - DateTime(timezone=True), nullable=False, server_default=func.now() - ) diff --git a/backend/app/scripts/download_models.py b/backend/app/scripts/download_models.py index 4b5e6d3..c6ff7a3 100644 --- a/backend/app/scripts/download_models.py +++ b/backend/app/scripts/download_models.py @@ -7,7 +7,6 @@ import sys from pathlib import Path MODEL_ROOT = Path(os.environ.get("ML_MODEL_DIR", "/models")) -CAMIE_REPO = os.environ.get("CAMIE_HF_REPO", "Camais03/camie-tagger-v2") SIGLIP_REPO = os.environ.get( "SIGLIP_HF_REPO", "google/siglip-so400m-patch14-384" ) @@ -24,34 +23,6 @@ def _snapshot(repo_id: str, dest: Path, allow_patterns: list[str] | None) -> Non ) -def ensure_camie() -> None: - """Fetch Camie v2 weights + metadata. - - v2 layout (HuggingFace Camais03/camie-tagger-v2): the ONNX file is - named camie-tagger-v2.onnx (not model.onnx) and tags ship inside - camie-tagger-v2-metadata.json (not selected_tags.csv). Both at root. - The repo also contains app/, game/, training/, images/ subdirs full - of setup/demo files we don't need — allow_patterns scopes the fetch - to just the inference essentials (~790 MB instead of ~2 GB). - """ - dest = MODEL_ROOT / "camie" - model_file = dest / "camie-tagger-v2.onnx" - meta_file = dest / "camie-tagger-v2-metadata.json" - if model_file.is_file() and meta_file.is_file(): - print(f"[download_models] Camie present at {dest}") - return - print(f"[download_models] Fetching {CAMIE_REPO} -> {dest}") - _snapshot( - CAMIE_REPO, dest, - [ - "camie-tagger-v2.onnx", - "camie-tagger-v2-metadata.json", - "config.json", - "config.yaml", - ], - ) - - def ensure_siglip() -> None: dest = MODEL_ROOT / "siglip" if (dest / "config.json").is_file() and any(dest.glob("*.safetensors")): @@ -62,7 +33,6 @@ def ensure_siglip() -> None: def main() -> int: - ensure_camie() ensure_siglip() print("[download_models] Done.") return 0 diff --git a/backend/app/services/cleanup_service.py b/backend/app/services/cleanup_service.py index d5f86d4..3916909 100644 --- a/backend/app/services/cleanup_service.py +++ b/backend/app/services/cleanup_service.py @@ -395,9 +395,8 @@ def delete_images( def delete_tag(session: Session, *, tag_id: int) -> dict: """Simple DELETE FROM tag WHERE id=?. - Postgres cascades the rest (image_tag, tag_alias, tag_allowlist, - tag_reference_embedding, tag_suggestion_rejection, series_page). - Returns counts BEFORE delete so the caller can surface them. + Postgres cascades the rest (image_tag, tag_alias, tag_suggestion_rejection, + series_page). Returns counts BEFORE delete so the caller can surface them. Raises LookupError if tag_id not found. """ tag = session.get(Tag, tag_id) @@ -742,8 +741,7 @@ def purge_legacy_tags(session: Session, *, dry_run: bool = False) -> dict: artist-kind tags PLUS general tags whose name matches a legacy prefix (source:*). - CASCADE on image_tag / tag_alias / tag_allowlist / - tag_reference_embedding / tag_suggestion_rejection / series_page + CASCADE on image_tag / tag_alias / tag_suggestion_rejection / series_page clears the related rows on the parent DELETE. Returns: @@ -785,23 +783,21 @@ def purge_legacy_tags(session: Session, *, dry_run: bool = False) -> dict: return result -# The Camie-suggestable CONTENT vocabulary. "Reset content tagging" wipes -# these so the operator can re-tag from scratch via auto-suggest. fandom + -# series (and series_page ordering) are deliberately NOT here — they're kept. +# The CONTENT vocabulary. "Reset content tagging" wipes these so the operator +# can re-tag from scratch. fandom + series (and series_page ordering) are +# deliberately NOT here — they're kept. RESETTABLE_TAG_KINDS = ("general", "character") def reset_content_tagging(session: Session, *, dry_run: bool = False) -> dict: """Count (dry_run) or DELETE every general + character tag so the operator - can re-tag from scratch via the Camie auto-suggest. + can re-tag from scratch (heads/CCIP repopulate suggestions). - PRESERVED: fandom + series tags and their series_page ordering, plus every - image's image_prediction rows (untouched) so suggestions - repopulate immediately. CASCADE on image_tag / tag_alias / tag_allowlist / - tag_reference_embedding / tag_suggestion_rejection clears each deleted - tag's applications + metadata. Tag.fandom_id is SET NULL, so deleting - character tags never touches the fandom rows. Irreversible except via DB - backup restore. + PRESERVED: fandom + series tags and their series_page ordering. CASCADE on + image_tag / tag_alias / tag_suggestion_rejection clears each deleted tag's + applications + metadata. Tag.fandom_id is SET NULL, so deleting character + tags never touches the fandom rows. Irreversible except via DB backup + restore. Returns: {"by_kind": {"general": N, "character": M}, diff --git a/backend/app/services/gallery_service.py b/backend/app/services/gallery_service.py index cd93804..ca2a7df 100644 --- a/backend/app/services/gallery_service.py +++ b/backend/app/services/gallery_service.py @@ -289,6 +289,75 @@ def _gallery_images(rows, artists: dict[int, dict]) -> list[GalleryImage]: ] +def _diversify_similar(src, rows, limit, *, dup_threshold=6, lam=0.55): + """Trim a nearest-cosine candidate pool down to `limit` diverse picks. + + 1. pHash collapse: drop any candidate whose perceptual hash is within + `dup_threshold` Hamming bits of the anchor or an already-kept candidate — + so a reposted banner (and the anchor's own clones) appears at most once. + 2. MMR (Maximal Marginal Relevance): greedily pick the candidate maximising + `lam * sim_to_anchor - (1 - lam) * max_sim_to_already_picked`. This keeps + the most relevant up top but pushes the selection to SPAN clusters + instead of returning 40 variations of one image. + + Falls back to nearest-order (`rows[:limit]`) on any failure or a small pool. + """ + if len(rows) <= 1: + return rows[:limit] + try: + import imagehash + import numpy as np + except Exception: + return rows[:limit] + + # --- 1. pHash near-duplicate collapse (videos/NULL phash pass through) --- + kept = [] + seen = [] + if src.phash: + try: + seen.append(imagehash.hex_to_hash(src.phash)) + except Exception: + pass + for row in rows: + ph = row[0].phash + if ph: + try: + h = imagehash.hex_to_hash(ph) + if any((h - k) <= dup_threshold for k in seen): + continue + seen.append(h) + except Exception: + pass + kept.append(row) + if len(kept) <= limit: + return kept + + # --- 2. MMR re-rank on the L2-normalised SigLIP embeddings --- + try: + a = np.asarray(src.siglip_embedding, dtype=np.float32) + a = a / (np.linalg.norm(a) or 1.0) + V = np.vstack([ + np.asarray(row[0].siglip_embedding, dtype=np.float32) for row in kept + ]) + V = V / np.clip(np.linalg.norm(V, axis=1, keepdims=True), 1e-8, None) + except Exception: + return kept[:limit] + + rel = V @ a # (N,) cosine to the anchor + n = len(kept) + picked_mask = np.zeros(n, dtype=bool) + max_sim = np.zeros(n, dtype=np.float32) # max sim to anything picked yet + order = [] + for _ in range(min(limit, n)): + scores = lam * rel - (1.0 - lam) * max_sim + scores[picked_mask] = -np.inf + i = int(np.argmax(scores)) + order.append(i) + picked_mask[i] = True + max_sim = np.maximum(max_sim, V @ V[i]) + return [kept[i] for i in order] + + async def _artists_for(session, image_ids: list[int]) -> dict[int, dict]: """Map image_id -> {"name","slug"} via the canonical image_record.artist_id (FC-2d-vii-c). Bounded by page size.""" @@ -565,14 +634,20 @@ class GalleryService: untagged: bool = False, no_artist: bool = False, date_from: datetime | None = None, date_to: datetime | None = None, ) -> list[GalleryImage] | None: - """Visual "more like this": images ranked by cosine distance to - `image_id`'s SigLIP embedding (pgvector, HNSW-indexed — alembic 0036). - No ML inference here; the embedding was computed at import. + """Visual "more like this": images near `image_id`'s SigLIP embedding + (pgvector, HNSW-indexed — alembic 0036), then DIVERSIFIED so the result + doesn't collapse into one cluster. No ML inference here. - Returns None if the source image doesn't exist (→ 404), [] if it has - no embedding (a video / not-yet-embedded). Composes with the Phase-1/2 - scope filters (AND) but REPLACES the date sort — always nearest-first, - bounded to `limit` (no cursor; distance-ranking has no date cursor). + Pure nearest-cosine piles up near-identical images — a reposted banner + fills the whole grid, and once you wander into a B&W / comic-panel + cluster every neighbour is more of the same with no way back to colour + (operator-reported 2026-06-30). So we pull a WIDER candidate pool, then: + 1. collapse near-duplicate pHashes (and drop clones of the anchor), + 2. MMR re-rank — pick for closeness-to-anchor but penalise similarity + to what's already picked, so the result SPANS clusters. + + Returns None if the source doesn't exist (→ 404), [] if it has no + embedding. Composes with the scope filters (AND); REPLACES the date sort. """ if limit < 1 or limit > 200: raise ValueError("limit must be between 1 and 200") @@ -582,6 +657,9 @@ class GalleryService: if src.siglip_embedding is None: return [] + # Over-fetch so diversification has clusters to spread across — without a + # wide pool there's nothing but the near-dupes to choose from. + pool_n = min(200, max(limit * 5, 60)) distance = ImageRecord.siglip_embedding.cosine_distance(src.siglip_embedding) eff = _effective_date_col() stmt = select(ImageRecord, Post.post_date, eff.label("eff")) @@ -597,8 +675,9 @@ class GalleryService: platform=platform, untagged=untagged, no_artist=no_artist, date_from=date_from, date_to=date_to, ) - stmt = stmt.order_by(distance.asc()).limit(limit) + stmt = stmt.order_by(distance.asc()).limit(pool_n) rows = (await self.session.execute(stmt)).all() + rows = _diversify_similar(src, rows, limit) artists = await _artists_for(self.session, [r[0].id for r in rows]) return _gallery_images(rows, artists) diff --git a/backend/app/services/importer.py b/backend/app/services/importer.py index 5634928..70bc7c8 100644 --- a/backend/app/services/importer.py +++ b/backend/app/services/importer.py @@ -1475,20 +1475,8 @@ class Importer: existing.duration_seconds = duration # #871: keep the kept copy's duration existing.thumbnail_path = None existing.integrity_status = "unknown" - existing.tagger_model_version = None existing.siglip_embedding = None existing.siglip_model_version = None - existing.centroid_scores = None - # #768: predictions also live in the normalized image_prediction table - # now — clear them so a re-imported file re-derives a fresh set. - from sqlalchemy import delete as _delete - - from ..models import ImagePrediction as _ImagePrediction - self.session.execute( - _delete(_ImagePrediction).where( - _ImagePrediction.image_record_id == existing.id - ) - ) # created_at intentionally preserved; updated_at auto-bumps. self.session.flush() self.session.commit() diff --git a/backend/app/services/ml/allowlist.py b/backend/app/services/ml/allowlist.py index 857d665..a6bee65 100644 --- a/backend/app/services/ml/allowlist.py +++ b/backend/app/services/ml/allowlist.py @@ -1,36 +1,20 @@ -"""Allowlist semantics: accepting a suggestion adds the canonical tag to -image_tag AND to tag_allowlist; per-image removal/dismiss writes a rejection. +"""Suggestion actions: accept applies the canonical tag to an image (which +feeds head training); dismiss / reject record a per-image rejection. + +(The Camie allowlist bulk-apply was retired #1189 — heads + CCIP are the tag +source, and head auto-apply is the earned propagation. Accept no longer +allowlists or fans a tag out across the library.) """ -from collections.abc import Sequence -from dataclasses import dataclass - -from sqlalchemy import and_, delete, distinct, func, or_, select +from sqlalchemy import delete from sqlalchemy.dialects.postgresql import insert from sqlalchemy.ext.asyncio import AsyncSession -from ...models import ( - ImagePrediction, - MLSettings, - Tag, - TagAlias, - TagAllowlist, - TagSuggestionRejection, -) +from ...models import TagSuggestionRejection from ...models.tag import image_tag from .aliases import AliasService -@dataclass(frozen=True) -class AllowlistRow: - tag_id: int - tag_name: str - tag_kind: str - min_confidence: float - applied_count: int # image_tag rows currently carrying this tag - coverage_count: int # images a sweep WOULD cover at min_confidence - - class AllowlistService: def __init__(self, session: AsyncSession): self.session = session @@ -39,21 +23,11 @@ class AllowlistService: async def _apply_image_tag(self, image_id: int, tag_id: int, source: str): stmt = insert(image_tag).values( image_record_id=image_id, tag_id=tag_id, source=source - ) - stmt = stmt.on_conflict_do_nothing( + ).on_conflict_do_nothing( index_elements=["image_record_id", "tag_id"] ) await self.session.execute(stmt) - async def _add_to_allowlist(self, tag_id: int) -> bool: - """Returns True if newly added (caller should kick off retro-apply).""" - exists = await self.session.get(TagAllowlist, tag_id) - if exists is not None: - return False - self.session.add(TagAllowlist(tag_id=tag_id)) - await self.session.flush() - return True - async def _clear_rejection(self, image_id: int, tag_id: int): await self.session.execute( delete(TagSuggestionRejection) @@ -61,12 +35,11 @@ class AllowlistService: .where(TagSuggestionRejection.tag_id == tag_id) ) - async def accept(self, image_id: int, tag_id: int) -> bool: - """Accept a suggestion. Returns True if the tag was newly added to - the allowlist (the API layer enqueues apply_allowlist_tags then).""" + async def accept(self, image_id: int, tag_id: int) -> None: + """Apply the accepted tag to this image (source='ml_accepted', a head + training positive) and clear any prior rejection.""" await self._apply_image_tag(image_id, tag_id, source="ml_accepted") await self._clear_rejection(image_id, tag_id) - return await self._add_to_allowlist(tag_id) async def add_alias_and_accept( self, @@ -74,17 +47,16 @@ class AllowlistService: alias_string: str, alias_category: str, canonical_tag_id: int, - ) -> bool: + ) -> None: await self.aliases.create( alias_string, alias_category, canonical_tag_id ) - return await self.accept(image_id, canonical_tag_id) + await self.accept(image_id, canonical_tag_id) async def dismiss(self, image_id: int, tag_id: int) -> None: stmt = insert(TagSuggestionRejection).values( image_record_id=image_id, tag_id=tag_id - ) - stmt = stmt.on_conflict_do_nothing( + ).on_conflict_do_nothing( index_elements=["image_record_id", "tag_id"] ) await self.session.execute(stmt) @@ -96,118 +68,11 @@ class AllowlistService: await self._clear_rejection(image_id, tag_id) async def reject_applied_tag(self, image_id: int, tag_id: int) -> None: - """Operator removed an applied tag from an image. Remove the - image_tag row AND record a rejection so the allowlist won't - re-apply it on the next maintenance sweep.""" + """Operator removed an applied tag from an image. Remove the image_tag + row AND record a rejection so head auto-apply won't re-apply it.""" await self.session.execute( image_tag.delete() .where(image_tag.c.image_record_id == image_id) .where(image_tag.c.tag_id == tag_id) ) await self.dismiss(image_id, tag_id) - - async def _store_floor(self) -> float: - return ( - await self.session.execute( - select(MLSettings.tagger_store_floor).where(MLSettings.id == 1) - ) - ).scalar_one() - - async def update_threshold( - self, tag_id: int, min_confidence: float - ) -> None: - row = await self.session.get(TagAllowlist, tag_id) - if row is not None: - # An allowlist tag can't auto-apply more permissively than the - # ingest store floor — predictions below tagger_store_floor aren't - # stored, so a lower min_confidence would behave identically to the - # floor. Clamp so the stored threshold matches actual behavior - # (#764). - floor = await self._store_floor() - row.min_confidence = max(min_confidence, floor) - - async def remove(self, tag_id: int) -> None: - await self.session.execute( - delete(TagAllowlist).where(TagAllowlist.tag_id == tag_id) - ) - - async def _coverage_match(self, tag: Tag): - """The predicate over image_prediction rows that resolve to `tag`, - mirroring tasks.ml._confidence_for_tag's resolution: a prediction whose - raw_name equals the tag name (any category), OR an alias maps - (raw_name, category) -> this tag. Returns a SQLAlchemy boolean clause. - """ - alias_rows = ( - await self.session.execute( - select(TagAlias.alias_string, TagAlias.alias_category).where( - TagAlias.canonical_tag_id == tag.id - ) - ) - ).all() - name_clause = ImagePrediction.raw_name == tag.name - alias_clauses = [ - and_( - ImagePrediction.raw_name == a, - ImagePrediction.category == c, - ) - for a, c in alias_rows - ] - return or_(name_clause, *alias_clauses) if alias_clauses else name_clause - - async def coverage(self, tag_id: int, threshold: float) -> int: - """How many distinct images a sweep WOULD cover for this tag at - `threshold`: images with a resolving prediction scoring >= threshold. - The gross candidate pool (NOT minus already-applied/rejected) — it's - the tuning signal for "lower the threshold and ~N more images qualify". - """ - tag = await self.session.get(Tag, tag_id) - if tag is None: - return 0 - match = await self._coverage_match(tag) - stmt = select( - func.count(distinct(ImagePrediction.image_record_id)) - ).where(ImagePrediction.score >= threshold, match) - return (await self.session.execute(stmt)).scalar_one() - - async def list_all(self) -> Sequence[AllowlistRow]: - stmt = ( - select( - TagAllowlist.tag_id, - Tag.name, - Tag.kind, - TagAllowlist.min_confidence, - ) - .join(Tag, Tag.id == TagAllowlist.tag_id) - .order_by(Tag.name.asc()) - ) - rows = (await self.session.execute(stmt)).all() - tag_ids = [r[0] for r in rows] - - # Applied counts in ONE grouped query (vs N per-row counts). - applied: dict[int, int] = {} - if tag_ids: - applied = dict( - ( - await self.session.execute( - select(image_tag.c.tag_id, func.count()) - .where(image_tag.c.tag_id.in_(tag_ids)) - .group_by(image_tag.c.tag_id) - ) - ).all() - ) - - result = [] - for r in rows: - # Coverage is per-tag (alias set differs); allowlist is small. - cov = await self.coverage(r[0], r[3]) - result.append( - AllowlistRow( - tag_id=r[0], - tag_name=r[1], - tag_kind=r[2].value if hasattr(r[2], "value") else str(r[2]), - min_confidence=r[3], - applied_count=applied.get(r[0], 0), - coverage_count=cov, - ) - ) - return result diff --git a/backend/app/services/ml/centroids.py b/backend/app/services/ml/centroids.py deleted file mode 100644 index e18907f..0000000 --- a/backend/app/services/ml/centroids.py +++ /dev/null @@ -1,163 +0,0 @@ -"""Tag centroids: the mean SigLIP embedding of a tag's member images. - -Powers centroid-augmented suggestions (a tag whose centroid is close to an -image's embedding becomes a suggestion even if Camie didn't predict it). -""" - -from dataclasses import dataclass - -import numpy as np -from sqlalchemy import func, select -from sqlalchemy.dialects.postgresql import insert -from sqlalchemy.ext.asyncio import AsyncSession - -from ...models import ( - ImageRecord, - MLSettings, - Tag, - TagKind, - TagReferenceEmbedding, -) -from ...models.tag import image_tag - -ELIGIBLE_KINDS = { - TagKind.character, - TagKind.fandom, - TagKind.general, - TagKind.series, -} - - -@dataclass(frozen=True) -class CentroidHit: - tag_id: int - similarity: float - - -class CentroidService: - def __init__(self, session: AsyncSession): - self.session = session - - async def _min_reference_images(self) -> int: - return ( - await self.session.execute( - select(MLSettings.min_reference_images).where(MLSettings.id == 1) - ) - ).scalar_one() - - async def _model_version(self) -> str: - """Audit 2026-06-02: SigLIP model-version stamp comes from the - DB row, not the env constant. tag_and_embed (tasks/ml.py:110) - already reads from MLSettings.embedder_model_version, so by - sourcing centroid stamps + drift checks from the same row, we - eliminate the silent-drift case the audit flagged. env - SIGLIP_MODEL_VERSION still drives which model embedder.py - loads at runtime; the version stamp is purely the operator- - controlled identifier.""" - return ( - await self.session.execute( - select(MLSettings.embedder_model_version).where(MLSettings.id == 1) - ) - ).scalar_one() - - async def recompute_for_tag(self, tag_id: int) -> bool: - """Recompute one tag's centroid. Returns True if a centroid was - written, False if skipped (ineligible kind or too few members).""" - tag = await self.session.get(Tag, tag_id) - if tag is None or tag.kind not in ELIGIBLE_KINDS: - return False - - min_refs = await self._min_reference_images() - - stmt = ( - select(ImageRecord.siglip_embedding) - .join(image_tag, image_tag.c.image_record_id == ImageRecord.id) - .where(image_tag.c.tag_id == tag_id) - .where(ImageRecord.siglip_embedding.is_not(None)) - ) - embeddings = [ - np.array(e, dtype=np.float32) - for e in (await self.session.execute(stmt)).scalars().all() - ] - if len(embeddings) < min_refs: - return False - - centroid = np.mean(np.stack(embeddings), axis=0).astype(np.float32) - model_version = await self._model_version() - - stmt = insert(TagReferenceEmbedding).values( - tag_id=tag_id, - embedding=centroid.tolist(), - reference_count=len(embeddings), - model_version=model_version, - ) - stmt = stmt.on_conflict_do_update( - index_elements=["tag_id"], - set_={ - "embedding": centroid.tolist(), - "reference_count": len(embeddings), - "model_version": model_version, - "updated_at": func.now(), - }, - ) - await self.session.execute(stmt) - return True - - async def list_drifted(self) -> list[int]: - """Tag ids whose centroid is stale: member count != reference_count, - OR no centroid row, OR centroid built on a different SigLIP version. - Only considers eligible-kind tags with embeddings present.""" - current_model_version = await self._model_version() - member_counts = ( - select( - image_tag.c.tag_id.label("tag_id"), - func.count(image_tag.c.image_record_id).label("members"), - ) - .join(ImageRecord, ImageRecord.id == image_tag.c.image_record_id) - .where(ImageRecord.siglip_embedding.is_not(None)) - .group_by(image_tag.c.tag_id) - .subquery() - ) - stmt = ( - select(Tag.id) - .join(member_counts, member_counts.c.tag_id == Tag.id) - .outerjoin( - TagReferenceEmbedding, - TagReferenceEmbedding.tag_id == Tag.id, - ) - .where(Tag.kind.in_(ELIGIBLE_KINDS)) - .where( - (TagReferenceEmbedding.tag_id.is_(None)) - | ( - TagReferenceEmbedding.reference_count - != member_counts.c.members - ) - | (TagReferenceEmbedding.model_version != current_model_version) - ) - ) - return list((await self.session.execute(stmt)).scalars().all()) - - async def find_similar_tags( - self, image_id: int, limit: int = 20 - ) -> list[CentroidHit]: - """Cosine similarity between an image's embedding and stored - centroids. Returns top-`limit` by similarity DESC. pgvector's - cosine_distance gives 1 - cosine_similarity.""" - img = await self.session.get(ImageRecord, image_id) - if img is None or img.siglip_embedding is None: - return [] - emb = img.siglip_embedding - distance = TagReferenceEmbedding.embedding.cosine_distance(emb) - stmt = ( - select( - TagReferenceEmbedding.tag_id, - (1 - distance).label("similarity"), - ) - .order_by(distance.asc()) - .limit(limit) - ) - rows = (await self.session.execute(stmt)).all() - return [ - CentroidHit(tag_id=r.tag_id, similarity=float(r.similarity)) - for r in rows - ] diff --git a/backend/app/services/ml/embedder.py b/backend/app/services/ml/embedder.py index d94f71e..d55646c 100644 --- a/backend/app/services/ml/embedder.py +++ b/backend/app/services/ml/embedder.py @@ -18,9 +18,11 @@ ImageFile.LOAD_TRUNCATED_IMAGES = True # N_replicas × this within the cores allotted to ML to avoid oversubscription. _INTRA_OP_THREADS = 4 -MODEL_NAME = os.environ.get( +DEFAULT_MODEL_NAME = os.environ.get( "SIGLIP_MODEL_NAME", "google/siglip-so400m-patch14-384" ) +# Back-compat alias (api/gpu imported this name as the fallback embedder id). +MODEL_NAME = DEFAULT_MODEL_NAME MODEL_VERSION = os.environ.get( "SIGLIP_MODEL_VERSION", "siglip-so400m-patch14-384" ) @@ -29,35 +31,42 @@ _LOCAL_DIR = Path(os.environ.get("ML_MODEL_DIR", "/models")) / "siglip" class Embedder: - def __init__(self, model_dir: Path | None = None): - self._model_dir = model_dir or _LOCAL_DIR + """Loads whatever SigLIP-family model it's given by HF NAME. For the default + model it prefers the pre-downloaded local dir (no re-download on existing + deploys); any other name resolves as an HF repo id (downloaded + cached on + first use), so an operator model swap (#1190) just works server-side.""" + + def __init__(self, model_name: str | None = None, model_dir: Path | None = None): + self.model_name = model_name or DEFAULT_MODEL_NAME + self._explicit_dir = model_dir self._model = None self._processor = None self._torch = None + def _source(self) -> str: + if self._explicit_dir is not None: + return str(self._explicit_dir) + if self.model_name == DEFAULT_MODEL_NAME and _LOCAL_DIR.exists(): + return str(_LOCAL_DIR) + return self.model_name + def load(self) -> None: if self._model is not None: return import torch - from transformers import AutoModel, SiglipImageProcessor + from transformers import AutoImageProcessor, AutoModel self._torch = torch # Bound torch's CPU thread pool (see _INTRA_OP_THREADS) so each replica # stays a predictable core consumer on a shared node. torch.set_num_threads(_INTRA_OP_THREADS) - # FC's embedder only does IMAGE inference — never text. AutoProcessor - # loads the full processor including SiglipTokenizer, which requires - # the sentencepiece library at import time even if we never call it. - # SiglipImageProcessor loads ONLY preprocessor_config.json (image - # side) and skips the tokenizer config entirely. Operator hit the - # ImportError 2026-05-25 once the ml-worker started actually running - # tag_and_embed; switching to the image-only loader avoids the - # tokenizer dep without adding ~30 MB of unused C++ build to the - # lean ml-worker image. - self._processor = SiglipImageProcessor.from_pretrained( - str(self._model_dir) - ) - self._model = AutoModel.from_pretrained(str(self._model_dir)) + # IMAGE inference only — AutoImageProcessor loads just the image side + # (preprocessor_config.json), skipping the SigLIP tokenizer + its + # sentencepiece dep (operator hit that ImportError 2026-05-25). Works + # for any SigLIP-family model, keeping the embedder model-agnostic. + src = self._source() + self._processor = AutoImageProcessor.from_pretrained(src) + self._model = AutoModel.from_pretrained(src) self._model.eval() def infer(self, image_path: Path) -> np.ndarray: @@ -74,8 +83,12 @@ class Embedder: _default_embedder: Embedder | None = None -def get_embedder() -> Embedder: +def get_embedder(model_name: str | None = None) -> Embedder: + """Cached embedder for `model_name` (default if None). Rebuilds the singleton + when the requested name changes, so an operator model swap takes effect + without restarting the worker.""" global _default_embedder - if _default_embedder is None: - _default_embedder = Embedder() + name = model_name or DEFAULT_MODEL_NAME + if _default_embedder is None or _default_embedder.model_name != name: + _default_embedder = Embedder(model_name=name) return _default_embedder diff --git a/backend/app/services/ml/heads.py b/backend/app/services/ml/heads.py index 879b922..b234897 100644 --- a/backend/app/services/ml/heads.py +++ b/backend/app/services/ml/heads.py @@ -308,25 +308,36 @@ async def score_image( import numpy as np img = await session.get(ImageRecord, image_id) - if img is None or img.siglip_embedding is None: + if img is None: return [] settings = await _settings_async(session) - heads = await _current_heads(session, settings.embedder_model_version) + cur_version = settings.embedder_model_version + heads = await _current_heads(session, cur_version) if heads["W"] is None: return [] - bag = [np.asarray(img.siglip_embedding, dtype=np.float32)] + # Only embeddings in the CURRENT model's space enter the bag. Mid model-swap + # (#1190), an image still carrying the OLD-version whole-image vector is + # skipped rather than scored by heads trained in a different space; a legacy + # NULL version is treated as current (those predate per-row stamping). + bag = [] + if img.siglip_embedding is not None and img.siglip_model_version in ( + cur_version, None, + ): + bag.append(np.asarray(img.siglip_embedding, dtype=np.float32)) region_vecs = ( await session.execute( select(ImageRegion.siglip_embedding) .where(ImageRegion.image_record_id == image_id) .where(ImageRegion.siglip_embedding.is_not(None)) - .where(ImageRegion.embedding_version == settings.embedder_model_version) + .where(ImageRegion.embedding_version == cur_version) ) ).all() for (vec,) in region_vecs: if vec is not None: bag.append(np.asarray(vec, dtype=np.float32)) + if not bag: + return [] X = np.vstack(bag) # (B, D) norms = np.linalg.norm(X, axis=1, keepdims=True) diff --git a/backend/app/services/ml/tagger.py b/backend/app/services/ml/tagger.py deleted file mode 100644 index 0f15c4a..0000000 --- a/backend/app/services/ml/tagger.py +++ /dev/null @@ -1,210 +0,0 @@ -"""Camie-tagger-v2 ONNX wrapper (CPU). - -Single-image at a time. Loaded lazily inside the ml-worker process; NOT -thread-safe — the ml queue worker runs --concurrency=1 per process (scale ML by -running multiple worker replicas, not threads). - -v2 layout reference: HuggingFace Camais03/camie-tagger-v2 root has -camie-tagger-v2.onnx (789 MB) + camie-tagger-v2-metadata.json (7.77 MB) -+ config.json. Tags ship as nested JSON, not CSV. Preprocessing and -output handling follow the published onnx_inference.py reference: -ImageNet normalize, NCHW layout, sigmoid on refined logits (output[1]). -""" - -import json -import os -from dataclasses import dataclass -from pathlib import Path - -import numpy as np -from PIL import Image, ImageFile - -# Cap inference threads (see Tagger.load) so each ml-worker replica is a bounded -# core consumer on a shared node — keep N_replicas × this within the cores -# allotted to ML so replicas don't oversubscribe the box / starve the DB. -_INTRA_OP_THREADS = 4 - -# onnxruntime lives in requirements-ml.txt only — it is NOT installed in the -# lean web image or in CI. Imported lazily inside Tagger.load() so this module -# imports fine without it (the suggestion service imports SURFACED_CATEGORIES -# from here in the web container, and CI collects the pure-logic tests). - -# Tolerate minutely-truncated source images (same rationale as IR's wd14.py: -# a few missing bytes at the JPEG EOI shouldn't block tagging the whole image). -ImageFile.LOAD_TRUNCATED_IMAGES = True - -MODEL_NAME = os.environ.get("CAMIE_MODEL_NAME", "camie-tagger-v2") -_MODEL_DIR = Path(os.environ.get("ML_MODEL_DIR", "/models")) / "camie" -_MODEL_FILE = f"{MODEL_NAME}.onnx" -_METADATA_FILE = f"{MODEL_NAME}-metadata.json" - -# Ingest floor below which predictions aren't stored (keeps the JSON compact). -# DEFAULT/fallback only — the live value is DB-backed -# (ml_settings.tagger_store_floor) and passed into infer() per call by the ml -# task. 0.70: the suggestion path already filters there and the centroid path -# covers lower-confidence preferred tags, so the sub-0.70 tail is redundant -# (it had bloated image_record's TOAST to ~100 GB; plan-task #764). -DEFAULT_STORE_FLOOR = 0.70 - -# The categories FC-2b surfaces in the UI. Others (meta/rating/year) are -# still stored but the suggestion service filters them out. -# 'artist' retired in FC-2d-vii-c — artist identity is acquisition-derived -# (image_record.artist_id), never ML-inferred. 'copyright' retired -# 2026-06-01 — operator doesn't use the copyright tag-kind; fandom is -# this app's franchise/series concept (per TagsView.vue's doc comment). -# Raw predictions for both categories still get stored at STORE_FLOOR but -# don't surface in suggestions. -SURFACED_CATEGORIES = {"character", "general"} - -# ImageNet preprocessing constants (per Camie v2 onnx_inference.py). -_IMAGENET_MEAN = np.array([0.485, 0.456, 0.406], dtype=np.float32) -_IMAGENET_STD = np.array([0.229, 0.224, 0.225], dtype=np.float32) -# Square-pad color ≈ ImageNet mean × 255 (matches reference inference). -_PAD_COLOR = (124, 116, 104) - - -@dataclass(frozen=True) -class TagPrediction: - name: str - category: str - confidence: float - - -class Tagger: - def __init__(self, model_dir: Path | None = None): - self._model_dir = model_dir or _MODEL_DIR - self._session = None # onnxruntime.InferenceSession once load()ed - self._tag_names: list[str] | None = None - self._tag_categories: list[str] | None = None - self._input_name: str | None = None - self._input_size: int = 512 - - def load(self) -> None: - if self._session is not None: - return - model_path = self._model_dir / _MODEL_FILE - meta_path = self._model_dir / _METADATA_FILE - if not model_path.is_file(): - raise RuntimeError( - f"Camie {_MODEL_FILE} missing at {model_path}. " - f"Populate /models via the ml-worker downloader." - ) - if not meta_path.is_file(): - raise RuntimeError( - f"Camie {_METADATA_FILE} missing at {meta_path}. " - f"Populate /models via the ml-worker downloader." - ) - - with open(meta_path) as f: - metadata = json.load(f) - - # Per Camie v2 onnx_inference.py: idx_to_tag is keyed by str(idx); - # tag_to_category maps tag_name -> category. Project to two parallel - # lists indexed by output position for O(1) lookup in the hot path. - ds = metadata["dataset_info"] - idx_to_tag = ds["tag_mapping"]["idx_to_tag"] - tag_to_category = ds["tag_mapping"]["tag_to_category"] - total = ds["total_tags"] - names: list[str] = [] - cats: list[str] = [] - for i in range(total): - name = idx_to_tag.get(str(i), f"unknown-{i}") - names.append(name) - cats.append(tag_to_category.get(name, "general")) - - # Input size from metadata; fall back to 512 (the v2 default). - self._input_size = int( - metadata.get("model_info", {}).get("img_size", 512) - ) - - # Lazy import — kept after the file-existence checks so the - # missing-model RuntimeError still fires first in environments - # without onnxruntime (CI / lean web image). - import onnxruntime as ort - - # Cap the intra-op thread pool. ONNX Runtime otherwise sizes it to ALL - # host cores, so on a shared node each ml-worker replica would grab every - # core and oversubscribe (and starve the co-located DB/web). Bounding it - # makes each replica a predictable core consumer — run N replicas where - # N × _INTRA_OP_THREADS stays within the cores you allot to ML. - opts = ort.SessionOptions() - opts.intra_op_num_threads = _INTRA_OP_THREADS - session = ort.InferenceSession( - str(model_path), sess_options=opts, providers=["CPUExecutionProvider"], - ) - self._input_name = session.get_inputs()[0].name - # Assign sentinels last so a partial load isn't observable. - self._tag_names = names - self._tag_categories = cats - self._session = session - - def _preprocess(self, image_path: Path) -> np.ndarray: - img = Image.open(image_path) - # Composite RGBA onto neutral so transparency doesn't bias the model. - if img.mode == "RGBA": - bg = Image.new("RGBA", img.size, (255, 255, 255, 255)) - bg.paste(img, mask=img.split()[3]) - img = bg.convert("RGB") - elif img.mode != "RGB": - img = img.convert("RGB") - - # Pad to square with ImageNet-mean color, then bicubic resize. - w, h = img.size - side = max(w, h) - square = Image.new("RGB", (side, side), _PAD_COLOR) - square.paste(img, ((side - w) // 2, (side - h) // 2)) - square = square.resize( - (self._input_size, self._input_size), Image.BICUBIC - ) - - arr = np.array(square, dtype=np.float32) / 255.0 # HWC, [0,1] - arr = (arr - _IMAGENET_MEAN) / _IMAGENET_STD # ImageNet normalize - arr = arr.transpose(2, 0, 1) # HWC -> CHW - return arr[np.newaxis, :, :, :] # NCHW - - def infer( - self, image_path: Path, *, store_floor: float = DEFAULT_STORE_FLOOR, - ) -> dict[str, TagPrediction]: - """Run Camie v2 on one image. Returns {name: TagPrediction} with - confidence >= store_floor (across all categories — the suggestion - service does category filtering later). store_floor is the DB-backed - ml_settings.tagger_store_floor, passed in by the ml task. - - v2 emits multiple outputs; we use the refined predictions - (output[1] per onnx_inference.py). Sigmoid is applied to raw - logits to produce [0,1] confidence scores. - """ - self.load() - x = self._preprocess(image_path) - outputs = self._session.run(None, {self._input_name: x}) - # Refined predictions if present (v2 emits initial + refined), - # fall back to initial for single-output forks. - logits = outputs[1] if len(outputs) > 1 else outputs[0] - # Squeeze batch dim, apply sigmoid. - probs = 1.0 / (1.0 + np.exp(-logits[0])) - results: dict[str, TagPrediction] = {} - names = self._tag_names - cats = self._tag_categories - for idx, score in enumerate(probs): - conf = float(score) - if conf < store_floor: - continue - if idx >= len(names): - # Output longer than metadata declared — shouldn't happen but - # don't crash the import pipeline if v2 metadata desynchronizes. - continue - results[names[idx]] = TagPrediction( - name=names[idx], category=cats[idx], confidence=conf - ) - return results - - -_default_tagger: Tagger | None = None - - -def get_tagger() -> Tagger: - """Process-level singleton so the ONNX session loads once per worker.""" - global _default_tagger - if _default_tagger is None: - _default_tagger = Tagger() - return _default_tagger diff --git a/backend/app/services/tag_service.py b/backend/app/services/tag_service.py index 839071c..a3ca07e 100644 --- a/backend/app/services/tag_service.py +++ b/backend/app/services/tag_service.py @@ -10,8 +10,6 @@ from sqlalchemy.dialects.postgresql import insert as pg_insert from sqlalchemy.ext.asyncio import AsyncSession from ..models import HeadMetric, Tag, TagHead, TagKind, image_tag -from ..models.tag_allowlist import TagAllowlist -from ..models.tag_reference_embedding import TagReferenceEmbedding from .db_helpers import get_or_create from .tag_query import fandom_join_alias, tag_columns @@ -304,35 +302,22 @@ class TagService: async def _keep_as_alias(self, tag_id: int) -> bool: """A merged-away tag's old name must survive as an alias iff the ML - pipeline has ever applied it OR could re-emit it (allowlisted / has - a centroid) — otherwise the proactive apply_allowlist_tags worker - would silently regenerate it. Purely-manual, ML-unknown tags are - deleted outright (no DB bloat).""" + pipeline has ever applied it (manual accept or head auto-apply) — so a + re-application or an alias remap resolves the canonical name. Purely- + manual, ML-unknown tags are deleted outright (no DB bloat).""" is_machine = await self.session.scalar( select( exists().where( and_( image_tag.c.tag_id == tag_id, image_tag.c.source.in_( - ("ml_auto", "ml_accepted", "auto") + ("ml_auto", "ml_accepted", "head_auto", "auto") ), ) ) ) ) - if is_machine: - return True - allowlisted = await self.session.scalar( - select(exists().where(TagAllowlist.tag_id == tag_id)) - ) - if allowlisted: - return True - has_centroid = await self.session.scalar( - select( - exists().where(TagReferenceEmbedding.tag_id == tag_id) - ) - ) - return bool(has_centroid) + return bool(is_machine) async def rename(self, tag_id: int, new_name: str) -> Tag: """Rename a tag. Raises TagMergeConflict if the new name collides @@ -572,8 +557,6 @@ class TagService: merged_count = await self._repoint_image_tags(source_id, target_id) await self._repoint_rejections(source_id, target_id) - await self._repoint_allowlist(source_id, target_id) - await self._repoint_embedding(source_id) await self._repoint_aliases(source_id, target_id) await self._repoint_fandom_children( source_id, target_id, source_kind @@ -639,30 +622,6 @@ class TagService: .values(tag_id=tgt) ) - async def _repoint_allowlist(self, src: int, tgt: int) -> None: - tgt_has = await self.session.scalar( - select(exists().where(TagAllowlist.tag_id == tgt)) - ) - if tgt_has: - await self.session.execute( - text("DELETE FROM tag_allowlist WHERE tag_id = :src"), - {"src": src}, - ) - else: - await self.session.execute( - update(TagAllowlist) - .where(TagAllowlist.tag_id == src) - .values(tag_id=tgt) - ) - - async def _repoint_embedding(self, src: int) -> None: - await self.session.execute( - text( - "DELETE FROM tag_reference_embedding WHERE tag_id = :src" - ), - {"src": src}, - ) - async def _repoint_aliases(self, src: int, tgt: int) -> None: from ..models.tag_alias import TagAlias diff --git a/backend/app/tasks/ml.py b/backend/app/tasks/ml.py index 84b40a2..84ca0ee 100644 --- a/backend/app/tasks/ml.py +++ b/backend/app/tasks/ml.py @@ -1,20 +1,19 @@ -"""ML Celery tasks: per-image inference, backfill discovery, centroid -recompute, allowlist auto-apply, model self-heal. +"""ML Celery tasks: per-image embedding, backfill discovery, head training, +model self-heal. -All run on the ml-worker (queue 'ml') except recompute_centroids and -apply_allowlist_tags sweeps which are 'maintenance' lane. Sync sessions -(Celery workers are sync processes), same pattern as FC-2a tasks. +All run on the ml-worker (queue 'ml'). Sync sessions (Celery workers are sync +processes), same pattern as FC-2a tasks. """ import logging from pathlib import Path from celery.exceptions import SoftTimeLimitExceeded -from sqlalchemy import delete, select +from sqlalchemy import select from sqlalchemy.exc import DBAPIError, OperationalError from ..celery_app import celery -from ..models import ImagePrediction, ImageRecord, MLSettings +from ..models import ImageRecord, MLSettings from ._sync_engine import sync_session_factory as _sync_session_factory log = logging.getLogger(__name__) @@ -46,19 +45,16 @@ def _is_video(path: Path) -> bool: time_limit=1200, # 20 min hard ) def tag_and_embed(self, image_id: int) -> dict: - """Run Camie + SigLIP on one image; store predictions + embedding; - then enqueue per-image allowlist application. + """Compute + store one image's SigLIP embedding. Video (#747): sample frames at a fixed cadence (ml_settings - video_frame_interval_seconds, capped at video_max_frames), keep a tag only if - it appears in >= video_min_tag_frames frames and average its confidence over - those frames (mean-pool, not max — kills one-frame noise); mean-pool the - SigLIP embeddings. On no-frames returns status='no_frames' (not an error). + video_frame_interval_seconds, capped at video_max_frames) and mean-pool the + per-frame SigLIP embeddings. On no-frames returns status='no_frames' (not an + error). (Camie tagging was retired #1189 — heads + CCIP are the tag source.) """ import time from ..services.ml.embedder import get_embedder - from ..services.ml.tagger import get_tagger # Phase + file context, so a timeout/crash names WHICH file and WHERE it # died instead of a bare SoftTimeLimitExceeded() (operator-flagged 2026-06-08: @@ -94,15 +90,13 @@ def tag_and_embed(self, image_id: int) -> dict: return {"status": "file_missing", "image_id": image_id} phase = "load_models" - tagger = get_tagger() - embedder = get_embedder() + embedder = get_embedder(settings.embedder_model_name) if is_vid: # Layer-3 isolation: ffprobe (a separate process) validates - # the container before we burn ~20 GPU ops sampling frames - # from it. A corrupt video that would crash the frame - # decoder is rejected cleanly here instead of taking down - # the ml-worker. Operator-flagged 2026-05-28. + # the container before we burn GPU ops sampling frames from it. + # A corrupt video that would crash the frame decoder is rejected + # cleanly here instead of taking down the ml-worker. phase = "video_probe" from ..utils import safe_probe vprobe = safe_probe.probe_video(src) @@ -115,48 +109,23 @@ def tag_and_embed(self, image_id: int) -> dict: "reason": vprobe.reason, } phase = "video_sample_frames" - t0 = time.monotonic() frames = _sample_video_frames( src, interval=settings.video_frame_interval_seconds, max_frames=settings.video_max_frames, ) - log.info( - "tag_and_embed sampled %d frame(s) in %.1fs: %s", - len(frames), time.monotonic() - t0, ctx, - ) if not frames: return {"status": "no_frames", "image_id": image_id} - phase = "video_infer" + phase = "video_embed" import numpy as np - preds = _aggregate_video_predictions( - [tagger.infer(f, store_floor=settings.tagger_store_floor) - for f in frames], - min_frames=settings.video_min_tag_frames, - ) + # Mean-pool the per-frame SigLIP embeddings into one vector. embedding = np.mean( [embedder.infer(f) for f in frames], axis=0 ).astype("float32") - log.info( - "tag_and_embed video aggregated %d tag(s) from %d frame(s) " - "(min_frames=%d): %s", - len(preds), len(frames), settings.video_min_tag_frames, ctx, - ) for f in frames: f.unlink(missing_ok=True) else: - phase = "tag" - t0 = time.monotonic() - raw = tagger.infer(src, store_floor=settings.tagger_store_floor) - log.info( - "tag_and_embed tagged in %.1fs (%d tags): %s", - time.monotonic() - t0, len(raw), ctx, - ) - preds = { - name: {"category": p.category, "confidence": p.confidence} - for name, p in raw.items() - } phase = "embed" t0 = time.monotonic() embedding = embedder.infer(src) @@ -166,28 +135,9 @@ def tag_and_embed(self, image_id: int) -> dict: ) phase = "persist" - record.tagger_model_version = settings.tagger_model_version record.siglip_embedding = embedding.tolist() record.siglip_model_version = settings.embedder_model_version session.add(record) - # Write the normalized image_prediction rows (#768) — the sole home - # for predictions now (image_record.tagger_predictions was dropped in - # migration 0046). Delete-then-insert keeps a re-tag idempotent; - # tagger_store_floor was already applied in tagger.infer, so preds is - # the >=floor set. - session.execute( - delete(ImagePrediction).where( - ImagePrediction.image_record_id == image_id - ) - ) - session.add_all([ - ImagePrediction( - image_record_id=image_id, raw_name=name, - category=p.get("category", "general"), - score=float(p.get("confidence", 0.0)), - ) - for name, p in preds.items() - ]) session.commit() except SoftTimeLimitExceeded: log.error( @@ -210,11 +160,8 @@ def tag_and_embed(self, image_id: int) -> dict: ) raise - log.info( - "tag_and_embed ok in %.1fs (%d tags): %s", _elapsed(), len(preds), ctx - ) - apply_allowlist_tags.delay(image_id=image_id) - return {"status": "ok", "image_id": image_id, "tags": len(preds)} + log.info("tag_and_embed ok in %.1fs: %s", _elapsed(), ctx) + return {"status": "ok", "image_id": image_id} def _sample_video_frames( @@ -273,68 +220,24 @@ def _sample_video_frames( return out -def _aggregate_video_predictions(per_frame: list[dict], *, min_frames: int) -> dict: - """Aggregate per-frame {name: TagPrediction} into one prediction set (#747). - - A tag is kept only if it appears (≥ the tagger store floor, already applied) - in at least `min_frames` of the sampled frames — because sampling is at a - fixed cadence, that means it was on screen for roughly min_frames×interval - seconds, so a single-frame flicker / scene-transition artifact is dropped - while a genuine scene-local tag in a long video survives. Confidence is the - MEAN over the frames where the tag appears (not max — max re-inflated the - one-frame noise this whole change exists to remove). - - `min_frames` is clamped to the number of frames actually sampled so a very - short video (1–2 frames) still tags instead of dropping everything. - """ - n = len(per_frame) - if n == 0: - return {} - threshold = max(1, min(min_frames, n)) - agg: dict[str, dict] = {} - for frame_preds in per_frame: - for name, p in frame_preds.items(): - cur = agg.get(name) - if cur is None: - agg[name] = {"category": p.category, "sum": p.confidence, "count": 1} - else: - cur["sum"] += p.confidence - cur["count"] += 1 - return { - name: {"category": v["category"], "confidence": v["sum"] / v["count"]} - for name, v in agg.items() - if v["count"] >= threshold - } - - @celery.task(name="backend.app.tasks.ml.backfill", bind=True) def backfill(self) -> int: - """Enqueue tag_and_embed for images missing predictions/embeddings for - the current model versions. Keyset pagination by id ASC (restart-safe). + """Enqueue tag_and_embed (embed-only) for images with no SigLIP embedding. + Keyset pagination by id ASC (restart-safe). + + NB: a siglip MODEL-VERSION mismatch (an operator model swap, #1190) is NOT + re-embedded here — the CPU ml-worker can't churn the library at 384/512px; + the GPU agent owns version re-embeds via the 'embed' job. """ SessionLocal = _sync_session_factory() enqueued = 0 last_id = 0 with SessionLocal() as session: - settings = session.execute( - select(MLSettings).where(MLSettings.id == 1) - ).scalar_one() while True: rows = session.execute( select(ImageRecord.id) .where(ImageRecord.id > last_id) - .where( - (ImageRecord.tagger_model_version.is_(None)) - | ( - ImageRecord.tagger_model_version - != settings.tagger_model_version - ) - | (ImageRecord.siglip_embedding.is_(None)) - | ( - ImageRecord.siglip_model_version - != settings.embedder_model_version - ) - ) + .where(ImageRecord.siglip_embedding.is_(None)) .order_by(ImageRecord.id.asc()) .limit(500) ).scalars().all() @@ -347,199 +250,6 @@ def backfill(self) -> int: return enqueued -@celery.task( - name="backend.app.tasks.ml.apply_allowlist_tags", - bind=True, - # Audit 2026-06-02 — the full-sweep mode (neither tag_id nor image_id) - # is O(images × allowlist) and legitimately runs >5 min on large - # libraries. Cap matches the maintenance queue's recovery threshold. - soft_time_limit=1800, time_limit=2100, -) -def apply_allowlist_tags(self, tag_id: int | None = None, - image_id: int | None = None) -> int: - """Retroactively apply allowlisted tags. - - Modes: - - tag_id only : scan all images for this tag. - - image_id only : scan all allowlisted tags for this image. - - both : just the (image, tag) pair. - - neither : full sweep (daily beat). - - Skips: already-applied, rejected (tag_suggestion_rejection), or - confidence below the tag's allowlist min_confidence. Applied with - source='ml_auto'. - """ - from sqlalchemy import and_ - from sqlalchemy import select as sa_select - from sqlalchemy.dialects.postgresql import insert as pg_insert - - from ..models import TagAllowlist, TagSuggestionRejection - from ..models.tag import image_tag - - SessionLocal = _sync_session_factory() - applied = 0 - with SessionLocal() as session: - allow_rows = session.execute( - sa_select(TagAllowlist.tag_id, TagAllowlist.min_confidence) - if tag_id is None - else sa_select( - TagAllowlist.tag_id, TagAllowlist.min_confidence - ).where(TagAllowlist.tag_id == tag_id) - ).all() - allow = {r[0]: r[1] for r in allow_rows} - if not allow: - return 0 - - # Images that have any predictions (#768: from image_prediction, not - # the old JSON column), optionally narrowed to one image. - img_ids_query = sa_select(ImagePrediction.image_record_id).distinct() - if image_id is not None: - img_ids_query = img_ids_query.where( - ImagePrediction.image_record_id == image_id - ) - - for (img_id,) in session.execute(img_ids_query).all(): - preds = _load_predictions_sync(session, img_id) - for a_tag_id, min_conf in allow.items(): - exists = session.execute( - sa_select(image_tag.c.tag_id).where( - and_( - image_tag.c.image_record_id == img_id, - image_tag.c.tag_id == a_tag_id, - ) - ) - ).scalar_one_or_none() - if exists is not None: - continue - rej = session.get( - TagSuggestionRejection, (img_id, a_tag_id) - ) - if rej is not None: - continue - from ..models import Tag - - tag = session.get(Tag, a_tag_id) - if tag is None: - continue - conf = _confidence_for_tag(session, tag, preds) - if conf is None or conf < min_conf: - continue - stmt = pg_insert(image_tag).values( - image_record_id=img_id, - tag_id=a_tag_id, - source="ml_auto", - ) - stmt = stmt.on_conflict_do_nothing( - index_elements=["image_record_id", "tag_id"] - ) - session.execute(stmt) - applied += 1 - session.commit() - return applied - - -def _load_predictions_sync(session, image_id: int) -> dict: - """Predictions for one image from image_prediction (#768), in the - {raw_name: {category, confidence}} shape _confidence_for_tag consumes — - keeps the allowlist resolution logic unchanged.""" - from sqlalchemy import select as sa_select - - rows = session.execute( - sa_select( - ImagePrediction.raw_name, - ImagePrediction.category, - ImagePrediction.score, - ).where(ImagePrediction.image_record_id == image_id) - ).all() - return { - r.raw_name: {"category": r.category, "confidence": r.score} - for r in rows - } - - -def _confidence_for_tag(session, tag, preds: dict) -> float | None: - """Highest confidence among predictions that resolve to `tag` — - either the prediction name equals the tag name, or an alias maps - (prediction name, category) -> tag.id. - """ - from sqlalchemy import select as sa_select - - from ..models import TagAlias - - best: float | None = None - direct = preds.get(tag.name) - if direct is not None: - best = float(direct.get("confidence", 0.0)) - alias_rows = session.execute( - sa_select(TagAlias.alias_string, TagAlias.alias_category).where( - TagAlias.canonical_tag_id == tag.id - ) - ).all() - for alias_string, alias_category in alias_rows: - p = preds.get(alias_string) - if p is None: - continue - if p.get("category") != alias_category: - continue - c = float(p.get("confidence", 0.0)) - if best is None or c > best: - best = c - return best - - -@celery.task(name="backend.app.tasks.ml.recompute_centroid", bind=True) -def recompute_centroid(self, tag_id: int) -> bool: - import asyncio - - from ..services.ml.centroids import CentroidService - from ._async_session import async_session_factory - - async def _run() -> bool: - # Per-task NullPool engine bound to THIS asyncio.run loop — the shared - # process-wide engine reuses connections across loops and raises - # "Future attached to a different loop" on every call after the first. - async_factory, async_engine = async_session_factory() - try: - async with async_factory() as session: - svc = CentroidService(session) - result = await svc.recompute_for_tag(tag_id) - await session.commit() - return result - finally: - await async_engine.dispose() - - return asyncio.run(_run()) - - -@celery.task( - name="backend.app.tasks.ml.recompute_centroids", - bind=True, - # Audit 2026-06-02 — drifted-centroid rebuild over potentially - # hundreds of tags. - soft_time_limit=1800, time_limit=2100, -) -def recompute_centroids(self) -> int: - """Daily: find drifted centroids, enqueue recompute_centroid for each.""" - import asyncio - - from ..services.ml.centroids import CentroidService - from ._async_session import async_session_factory - - async def _list() -> list[int]: - # Per-task NullPool engine bound to this loop (see recompute_centroid). - async_factory, async_engine = async_session_factory() - try: - async with async_factory() as session: - return await CentroidService(session).list_drifted() - finally: - await async_engine.dispose() - - drifted = asyncio.run(_list()) - for tid in drifted: - recompute_centroid.delay(tid) - return len(drifted) - - @celery.task( name="backend.app.tasks.ml.tag_eval_run", bind=True, @@ -750,17 +460,40 @@ def enqueue_gpu_backfill(task_name: str) -> int: job, so it picks up the back-catalogue of images that were CCIP-embedded before concept crops existed, and retries images whose concept embed failed — without re-touching their figure/CCIP regions.""" - from sqlalchemy import exists, insert, literal + from sqlalchemy import exists, insert, literal, or_ from sqlalchemy import select as sa_select - from ..models import GpuJob, ImageRecord, ImageRegion + from ..models import GpuJob, ImageRecord, ImageRegion, MLSettings SessionLocal = _sync_session_factory() with SessionLocal() as session: - if task_name == "siglip": - has_concept = exists().where( + cur_version = session.execute( + select(MLSettings.embedder_model_version).where(MLSettings.id == 1) + ).scalar_one() + if task_name == "embed": + # Whole-image GPU re-embed (#1190): images with no embedding, or one + # stamped under a DIFFERENT model version (an operator model swap). + stale = or_( + ImageRecord.siglip_embedding.is_(None), + ImageRecord.siglip_model_version.is_(None), + ImageRecord.siglip_model_version != cur_version, + ) + queued = exists().where( + GpuJob.image_record_id == ImageRecord.id, + GpuJob.task == "embed", + GpuJob.status.in_(["pending", "leased"]), + ) + sel = sa_select( + ImageRecord.id, literal("embed"), literal("pending") + ).where(stale).where(~queued) + elif task_name == "siglip": + # Concept-crop re-embed: enqueue when there's no concept region AT THE + # CURRENT model version — so a model swap re-triggers crops too, not + # only the never-embedded back-catalogue. + has_current_concept = exists().where( ImageRegion.image_record_id == ImageRecord.id, ImageRegion.kind == "concept", + ImageRegion.embedding_version == cur_version, ) queued = exists().where( GpuJob.image_record_id == ImageRecord.id, @@ -769,7 +502,7 @@ def enqueue_gpu_backfill(task_name: str) -> int: ) sel = sa_select( ImageRecord.id, literal("siglip"), literal("pending") - ).where(~has_concept).where(~queued) + ).where(~has_current_concept).where(~queued) else: already = exists().where( GpuJob.image_record_id == ImageRecord.id, diff --git a/frontend/src/components/settings/AllowlistTable.vue b/frontend/src/components/settings/AllowlistTable.vue deleted file mode 100644 index d128611..0000000 --- a/frontend/src/components/settings/AllowlistTable.vue +++ /dev/null @@ -1,120 +0,0 @@ - - - - - diff --git a/frontend/src/components/settings/CentroidRecomputeCard.vue b/frontend/src/components/settings/CentroidRecomputeCard.vue deleted file mode 100644 index f9c5df4..0000000 --- a/frontend/src/components/settings/CentroidRecomputeCard.vue +++ /dev/null @@ -1,36 +0,0 @@ - - - diff --git a/frontend/src/components/settings/GpuAgentCard.vue b/frontend/src/components/settings/GpuAgentCard.vue index dc0ed24..b0da818 100644 --- a/frontend/src/components/settings/GpuAgentCard.vue +++ b/frontend/src/components/settings/GpuAgentCard.vue @@ -106,6 +106,37 @@ reversible) — so identity tags keep flowing without review. Stricter than the suggest cut; 0.92 recommended.

+ + +
Embedding model (advanced)
+
+ + +
+ Save model + Re-embed library (GPU) +
+

+ Changing the model means a DIFFERENT embedding space. After saving a new + model + version, run Re-embed library (the GPU agent re-embeds + whole images + concept crops), then Retrain heads. Suggestions + degrade until both finish. SigLIP 2 (google/siglip2-so400m-patch16-512, + version siglip2-so400m-patch16-512) is a 1152-d drop-in at + 512px — no schema change. +

+
@@ -131,6 +162,10 @@ const savingThreshold = ref(false) const autoApply = ref(true) const autoThreshold = ref(0.92) const savingAuto = ref(false) +const modelName = ref('') +const modelVersion = ref('') +const savingModel = ref(false) +const reembedding = ref(false) const queue = ref({ pending: 0, leased: 0, done: 0, error: 0 }) let pollTimer = null @@ -157,9 +192,42 @@ onMounted(async () => { autoApply.value = ml.settings.ccip_auto_apply_enabled autoThreshold.value = ml.settings.ccip_auto_apply_threshold } + if (ml.settings?.embedder_model_name != null) { + modelName.value = ml.settings.embedder_model_name + modelVersion.value = ml.settings.embedder_model_version + } } catch { /* non-fatal */ } }) +async function onSaveModel() { + savingModel.value = true + try { + await ml.patchSettings({ + embedder_model_name: modelName.value.trim(), + embedder_model_version: modelVersion.value.trim(), + }) + toast({ text: 'Embedding model saved — now Re-embed library, then Retrain heads', type: 'success' }) + } catch (e) { + toast({ text: `Could not save model: ${e.message}`, type: 'error' }) + } finally { + savingModel.value = false + } +} + +async function onReembed() { + reembedding.value = true + try { + await store.backfill('embed') + await store.backfill('siglip') + toast({ text: 'Queued whole-image + concept re-embed — run the agent, then Retrain heads', type: 'success' }) + await refreshQueue() + } catch (e) { + toast({ text: `Could not queue re-embed: ${e.message}`, type: 'error' }) + } finally { + reembedding.value = false + } +} + async function onSaveAuto() { savingAuto.value = true try { diff --git a/frontend/src/components/settings/MLBackfillCard.vue b/frontend/src/components/settings/MLBackfillCard.vue index e4e2fb4..d0e9c21 100644 --- a/frontend/src/components/settings/MLBackfillCard.vue +++ b/frontend/src/components/settings/MLBackfillCard.vue @@ -2,12 +2,13 @@

- Re-run Camie + SigLIP on images missing predictions or embeddings - for the current model versions. Safe to re-run. + Compute the SigLIP embedding for any image that doesn't have one yet + (CPU). Safe to re-run. To re-embed under a NEW model, use the GPU + agent's "Re-embed library" instead.

mdi-refresh Run backfill now diff --git a/frontend/src/components/settings/MLThresholdSliders.vue b/frontend/src/components/settings/MLThresholdSliders.vue index eb1a8de..49c4a59 100644 --- a/frontend/src/components/settings/MLThresholdSliders.vue +++ b/frontend/src/components/settings/MLThresholdSliders.vue @@ -1,70 +1,30 @@