feat: ML tag suggestions, character/fandom integrity, underscores, modal polish

Consolidated merge of feat/tag-suggestions branch. Original 64-commit history
was lost to git-object corruption in a Nextcloud-synced checkout; this single
commit captures the equivalent diff.

Includes:
- pgvector-backed tag suggestion infra (WD14 + SigLIP centroids, ml-worker
  container, Celery tasks, suggestion service, accept/reject endpoints + modal
  UI with green/red chip buttons)
- Character/fandom integrity: title-case normalization on every write path,
  fandom-id backfill, maintenance task + settings button, migrations g26041901
  + h26041901 to canonicalize legacy rows with case-only duplicate merging
- Tag-underscores + modal polish: WD14 name canonicalization at emit + accept
  + add/bulk-add paths, migration i26041901 for legacy-row rename-or-merge
  across character/fandom/NULL kinds, suggestion-accept refresh parity via
  awaited loadTags, persistent chip tint
This commit is contained in:
2026-04-19 19:50:58 -04:00
parent 69b3ddcbd0
commit 0f35a0c484
37 changed files with 8642 additions and 30 deletions
+7 -1
View File
@@ -37,7 +37,9 @@ def make_celery(app=None):
'app.tasks.scan',
'app.tasks.import_file',
'app.tasks.thumbnail',
'app.tasks.sidecar'
'app.tasks.sidecar',
'app.tasks.ml',
'app.tasks.maintenance',
]
)
@@ -69,11 +71,15 @@ def make_celery(app=None):
'app.tasks.scan.cleanup_old_tasks': {'queue': 'maintenance'},
'app.tasks.scan.update_system_stats': {'queue': 'maintenance'},
'app.tasks.scan.update_batch_stats': {'queue': 'maintenance'},
'app.tasks.maintenance.sync_character_fandoms': {'queue': 'maintenance'},
# Import tasks - handled by worker (heavy processing)
'app.tasks.import_file.*': {'queue': 'import'},
'app.tasks.thumbnail.*': {'queue': 'thumbnail'},
'app.tasks.sidecar.*': {'queue': 'sidecar'},
# ML inference tasks - handled by ml-worker
'app.tasks.ml.*': {'queue': 'ml'},
},
# Task default queue for unrouted tasks
+208 -14
View File
@@ -1,6 +1,6 @@
# app/main.py
from flask import Blueprint, render_template, flash, redirect, url_for, abort, send_from_directory, request, jsonify
from flask import Blueprint, render_template, flash, redirect, url_for, abort, send_from_directory, request, jsonify, current_app
from sqlalchemy import desc, func, text, extract, and_
from sqlalchemy.orm import joinedload, aliased
@@ -9,6 +9,7 @@ from datetime import datetime
from app.models import ImageRecord, Tag, ArchiveRecord, image_tags, ImportTask, ImportBatch, AppSettings
from app import db
from app.utils.tag_names import normalize_display_name
import os
import re
import shutil
@@ -25,8 +26,12 @@ def _parse_character_fandom(name_after_prefix):
def _ensure_fandom_tag(fandom_name):
"""Find or create a fandom tag by name. Returns the Tag object."""
full_name = f"fandom:{fandom_name}"
"""Find or create a fandom tag by name. Returns the Tag object.
Normalizes the display name so callers don't have to remember.
"""
normalized = normalize_display_name(fandom_name.strip())
full_name = f"fandom:{normalized}"
fandom_tag = Tag.query.filter_by(name=full_name).first()
if not fandom_tag:
fandom_tag = Tag(name=full_name, kind="fandom")
@@ -601,6 +606,27 @@ def settings():
return render_template('settings.html', active_tab=tab)
@main.post('/settings/maintenance/ml-backfill')
def trigger_ml_backfill():
from app.tasks.ml import backfill
backfill.apply_async(queue='ml')
return redirect(url_for('main.settings', tab='maintenance'))
@main.post('/settings/maintenance/recompute-centroids')
def trigger_recompute_all_centroids():
from app.tasks.ml import recompute_all_centroids
recompute_all_centroids.apply_async(queue='ml')
return redirect(url_for('main.settings', tab='maintenance'))
@main.post('/settings/maintenance/sync-character-fandoms')
def trigger_sync_character_fandoms():
from app.tasks.maintenance import sync_character_fandoms
sync_character_fandoms.apply_async(queue='maintenance')
return redirect(url_for('main.settings', tab='maintenance'))
# ----------------------------
# Tag add/remove endpoints
# ----------------------------
@@ -677,6 +703,114 @@ def list_tags(image_id):
])
# ---------------------------------------------------------------------------
# Tag suggestions (ML-backed)
# ---------------------------------------------------------------------------
@main.get("/image/<int:image_id>/suggestions")
def get_image_suggestions(image_id):
from app.services.tag_suggestions import get_suggestions
suggestions = get_suggestions(image_id)
return jsonify({'ok': True, 'suggestions': suggestions})
@main.post("/image/<int:image_id>/suggestions/accept")
def accept_image_suggestion(image_id):
from app.models import SuggestionFeedback
from app.services.tag_suggestions import _WD14_CATEGORY_TO_KIND, _canonicalize_wd14_name
payload = request.get_json(force=True, silent=True) or {}
tag_name = (payload.get('tag_name') or '').strip()
source = payload.get('source') or ''
category = payload.get('category') or ''
confidence = float(payload.get('confidence') or 0.0)
if not tag_name or not source or not category:
return jsonify({'ok': False, 'error': 'missing_fields'}), 400
image = ImageRecord.query.get(image_id)
if image is None:
return jsonify({'ok': False, 'error': 'image_not_found'}), 404
# Canonicalize before lookup so an already-attached "big hair" matches a
# freshly-POSTed "big_hair" (defensive — the client should already send canonical).
tag_name = _canonicalize_wd14_name(tag_name, category)
tag = Tag.query.filter_by(name=tag_name).first()
if tag is None:
kind = _WD14_CATEGORY_TO_KIND.get(category) # None for 'general'
tag = Tag(name=tag_name, kind=kind)
db.session.add(tag)
db.session.flush()
if tag not in image.tags:
image.tags.append(tag)
db.session.add(SuggestionFeedback(
image_id=image_id,
tag_name=tag_name,
suggestion_source=source,
confidence=confidence,
decision='accepted',
))
db.session.commit()
# Schedule centroid recompute if this tag's kind is eligible and delta is exceeded.
from app.services.tag_suggestions import ELIGIBLE_CENTROID_KINDS
if tag.kind in ELIGIBLE_CENTROID_KINDS:
from app.models import TagReferenceEmbedding, TagSuggestionConfig
from app.tasks.ml import recompute_centroid
from app.ml.siglip import MODEL_VERSION as SIGLIP_VER
delta_cfg = TagSuggestionConfig.query.filter_by(key='centroid_recompute_delta').first()
delta_threshold = int(delta_cfg.value) if delta_cfg else 3
current_count = (
db.session.query(db.func.count(image_tags.c.image_id))
.filter(image_tags.c.tag_id == tag.id)
.scalar()
) or 0
ref = (
TagReferenceEmbedding.query
.filter_by(tag_name=tag_name, model_version=SIGLIP_VER)
.first()
)
last_count = ref.reference_count if ref else 0
if current_count - last_count >= delta_threshold:
try:
recompute_centroid.apply_async(args=[tag_name], queue='ml')
except Exception:
# Don't fail the accept if enqueue fails
pass
return jsonify({
'ok': True,
'tag': {'id': tag.id, 'name': tag.name, 'kind': tag.kind},
})
@main.post("/image/<int:image_id>/suggestions/reject")
def reject_image_suggestion(image_id):
from app.models import SuggestionFeedback
payload = request.get_json(force=True, silent=True) or {}
tag_name = (payload.get('tag_name') or '').strip()
source = payload.get('source') or ''
confidence = float(payload.get('confidence') or 0.0)
if not tag_name or not source:
return jsonify({'ok': False, 'error': 'missing_fields'}), 400
if ImageRecord.query.get(image_id) is None:
return jsonify({'ok': False, 'error': 'image_not_found'}), 404
db.session.add(SuggestionFeedback(
image_id=image_id,
tag_name=tag_name,
suggestion_source=source,
confidence=confidence,
decision='rejected',
))
db.session.commit()
return jsonify({'ok': True})
@main.post("/image/<int:image_id>/tags/add")
def add_tag(image_id):
"""
@@ -695,8 +829,11 @@ def add_tag(image_id):
kind = name.split(":", 1)[0]
tag_name = name
else:
# No prefix — treat as a user-typed general tag. Strip underscores so
# WD14-style names pasted into the add box converge with our convention.
from app.utils.tag_names import underscores_to_spaces
kind = "user"
tag_name = name
tag_name = underscores_to_spaces(name)
img = ImageRecord.query.get_or_404(image_id)
tag = Tag.query.filter_by(name=tag_name).first()
@@ -704,16 +841,29 @@ def add_tag(image_id):
fandom_tag = None
if not tag:
tag = Tag(name=tag_name, kind=kind)
# For new character tags, parse and link fandom
# For new character tags, parse + normalize + link fandom
if kind == "character":
char_part = tag_name.split(":", 1)[1]
char_name, fandom_name = _parse_character_fandom(char_part)
norm_char = normalize_display_name(char_name)
if fandom_name:
fandom_tag = _ensure_fandom_tag(fandom_name)
tag_name = f"character:{norm_char} ({fandom_tag.name.split(':', 1)[1]})"
else:
tag_name = f"character:{norm_char}"
elif kind == "fandom":
fandom_part = tag_name.split(":", 1)[1]
tag_name = f"fandom:{normalize_display_name(fandom_part)}"
# Look up again under the normalized name — an existing row may match.
tag = Tag.query.filter_by(name=tag_name).first()
if not tag:
tag = Tag(name=tag_name, kind=kind)
if kind == "character" and fandom_name:
tag.fandom_id = fandom_tag.id
db.session.add(tag)
db.session.flush()
db.session.add(tag)
db.session.flush()
elif kind == "character" and tag.fandom_id:
fandom_tag = Tag.query.get(tag.fandom_id)
elif tag.kind == "character" and tag.fandom_id:
# Existing character tag with a fandom — grab the fandom tag for auto-apply
fandom_tag = Tag.query.get(tag.fandom_id)
@@ -786,9 +936,10 @@ def set_tag_fandom(tag_id):
fandom_id = request.form.get("fandom_id", type=int)
fandom_name = (request.form.get("fandom_name") or "").strip()
# Get the character's base name (strip existing fandom suffix)
# Get the character's base name (strip existing fandom suffix) and normalize it.
char_part = tag.name.split(":", 1)[1] if ":" in tag.name else tag.name
char_base, _ = _parse_character_fandom(char_part)
char_base = normalize_display_name(char_base)
if fandom_id:
fandom_tag = Tag.query.get(fandom_id)
@@ -796,8 +947,9 @@ def set_tag_fandom(tag_id):
return jsonify(ok=False, error="Invalid fandom tag"), 400
fandom_display = fandom_tag.name.split(":", 1)[1] if ":" in fandom_tag.name else fandom_tag.name
elif fandom_name:
# _ensure_fandom_tag normalizes internally.
fandom_tag = _ensure_fandom_tag(fandom_name)
fandom_display = fandom_name
fandom_display = fandom_tag.name.split(":", 1)[1]
else:
# Clear fandom
tag.fandom_id = None
@@ -821,6 +973,33 @@ def set_tag_fandom(tag_id):
tag.name = new_name
db.session.commit()
# Retroactive backfill: every image already tagged with this character should
# also have the fandom tag attached. Runs in a second transaction so a failed
# backfill doesn't roll back the rename.
try:
db.session.execute(
text("""
INSERT INTO image_tags (image_id, tag_id)
SELECT it.image_id, :fandom_id
FROM image_tags it
WHERE it.tag_id = :char_id
AND NOT EXISTS (
SELECT 1 FROM image_tags it2
WHERE it2.image_id = it.image_id
AND it2.tag_id = :fandom_id
)
"""),
{'char_id': tag.id, 'fandom_id': fandom_tag.id},
)
db.session.commit()
except Exception as e:
db.session.rollback()
# Rename already committed; log and continue. Sweep button recovers.
current_app.logger.warning(
"set_tag_fandom backfill failed for tag %s -> fandom %s: %s",
tag.id, fandom_tag.id, e,
)
return jsonify(ok=True, tag={
"id": tag.id,
"name": tag.name,
@@ -2112,23 +2291,38 @@ def bulk_add_tag():
if ':' in tag_name:
kind = tag_name.split(':', 1)[0]
else:
# No prefix — treat as a user-typed general tag. Strip underscores so
# WD14-style names pasted into the bulk box converge with our convention.
from app.utils.tag_names import underscores_to_spaces
kind = 'user'
tag_name = underscores_to_spaces(tag_name)
# Get or create the tag
tag = Tag.query.filter_by(name=tag_name).first()
fandom_tag = None
if not tag:
tag = Tag(name=tag_name, kind=kind)
# For new character tags, parse and link fandom
if kind == 'character':
char_part = tag_name.split(':', 1)[1]
char_name, fandom_name = _parse_character_fandom(char_part)
norm_char = normalize_display_name(char_name)
if fandom_name:
fandom_tag = _ensure_fandom_tag(fandom_name)
tag_name = f"character:{norm_char} ({fandom_tag.name.split(':', 1)[1]})"
else:
tag_name = f"character:{norm_char}"
elif kind == 'fandom':
fandom_part = tag_name.split(':', 1)[1]
tag_name = f"fandom:{normalize_display_name(fandom_part)}"
tag = Tag.query.filter_by(name=tag_name).first()
if not tag:
tag = Tag(name=tag_name, kind=kind)
if kind == 'character' and fandom_name:
tag.fandom_id = fandom_tag.id
db.session.add(tag)
db.session.flush()
db.session.add(tag)
db.session.flush()
elif tag.kind == 'character' and tag.fandom_id:
fandom_tag = Tag.query.get(tag.fandom_id)
elif tag.kind == 'character' and tag.fandom_id:
fandom_tag = Tag.query.get(tag.fandom_id)
+1
View File
@@ -0,0 +1 @@
"""Machine-learning inference wrappers used by the ml-worker."""
+58
View File
@@ -0,0 +1,58 @@
"""SigLIP SO400M image-embedding wrapper (PyTorch CPU)."""
from __future__ import annotations
import os
import numpy as np
from PIL import Image
# Defer torch/transformers imports to lazily-loaded functions to allow
# importing MODEL_VERSION in non-ML-worker contexts (e.g., web container
# centroid-recompute enqueue logic).
_torch = None
_AutoModel = None
_AutoProcessor = None
MODEL_NAME = os.environ.get('SIGLIP_MODEL_NAME', 'google/siglip-so400m-patch14-384')
MODEL_VERSION = os.environ.get('SIGLIP_MODEL_VERSION', 'siglip-so400m-patch14-384')
# Model files live flat under this directory (written by scripts/download_models.py via
# snapshot_download(local_dir=...)). We point from_pretrained at the local path directly
# so transformers bypasses its HF cache layout and doesn't need network access at load time.
_LOCAL_DIR = os.path.join(os.environ.get('ML_MODEL_DIR', '/models'), 'siglip')
_model = None
_processor = None
# NOT thread-safe. Must run in the ml-worker container with --concurrency=1.
def _load() -> None:
global _model, _processor, _torch, _AutoModel, _AutoProcessor
if _model is not None:
return
# Lazy import torch/transformers so this module can be imported in contexts
# where they're not available (e.g., web container for centroid-recompute enqueue).
import torch
from transformers import AutoModel, AutoProcessor
_torch = torch
_AutoModel = AutoModel
_AutoProcessor = AutoProcessor
_processor = _AutoProcessor.from_pretrained(_LOCAL_DIR)
_model = _AutoModel.from_pretrained(_LOCAL_DIR)
_model.eval()
def infer(image_path: str) -> np.ndarray:
"""Return a 1152-dim float32 numpy embedding for the image.
SigLIP uses a MAP-pooled vision head — the pooled output is the retrieval-ready
embedding the model was trained to produce. `get_image_features` on transformers
>= 4.45 returns a BaseModelOutputWithPooling, so pull `.pooler_output` explicitly
rather than relying on the first-field fallback from indexing.
"""
_load()
img = Image.open(image_path).convert('RGB')
with _torch.no_grad():
inputs = _processor(images=img, return_tensors='pt')
out = _model.get_image_features(**inputs)
pooled = out.pooler_output if hasattr(out, 'pooler_output') else out # (1, 1152)
return pooled[0].numpy().astype(np.float32)
+110
View File
@@ -0,0 +1,110 @@
"""WD14 EVA02-Large tagger (ONNX CPU inference)."""
from __future__ import annotations
import csv
import os
from typing import Iterable
import numpy as np
import onnxruntime as ort
from PIL import Image
MODEL_VERSION = os.environ.get('WD14_MODEL_VERSION', 'wd-eva02-large-tagger-v3')
_MODEL_DIR = os.environ.get('ML_MODEL_DIR', '/models')
_WD14_DIR = os.path.join(_MODEL_DIR, 'wd14')
_MODEL_PATH = os.path.join(_WD14_DIR, 'model.onnx')
_TAGS_PATH = os.path.join(_WD14_DIR, 'selected_tags.csv')
# WD14 selected_tags.csv uses Danbooru category ids:
# 0=general, 1=artist, 3=copyright, 4=character, 5=meta, 9=rating
_CATEGORY_MAP = {0: 'general', 1: 'artist', 3: 'copyright', 4: 'character', 5: 'meta', 9: 'rating'}
_session: ort.InferenceSession | None = None
_tag_meta: list[dict] | None = None
_input_name: str | None = None
_output_name: str | None = None
_input_size: int = 448
# NOT thread-safe. Must run in the ml-worker container with --concurrency=1.
def _load() -> None:
global _session, _tag_meta, _input_name, _output_name, _input_size
if _session is not None:
return
if not os.path.isfile(_MODEL_PATH):
raise RuntimeError(
f"WD14 model file missing at {_MODEL_PATH}. "
f"Populate the /models volume via the ml-worker downloader."
)
if not os.path.isfile(_TAGS_PATH):
raise RuntimeError(
f"WD14 selected_tags.csv missing at {_TAGS_PATH}. "
f"Populate the /models volume via the ml-worker downloader."
)
tag_meta: list[dict] = []
with open(_TAGS_PATH, newline='') as f:
reader = csv.DictReader(f)
for row in reader:
tag_meta.append({
'name': row['name'],
'category': _CATEGORY_MAP.get(int(row['category']), 'unknown'),
})
session = ort.InferenceSession(
_MODEL_PATH,
providers=['CPUExecutionProvider'],
)
_input_name = session.get_inputs()[0].name
_output_name = session.get_outputs()[0].name
# Input shape is usually [batch, H, W, 3] NHWC; pick the spatial dim
input_shape = session.get_inputs()[0].shape
for dim in input_shape:
if isinstance(dim, int) and dim > 1:
_input_size = dim
break
# Assign sentinels last so a partially-loaded state can't be observed.
_tag_meta = tag_meta
_session = session
def _preprocess(image_path: str) -> np.ndarray:
img = Image.open(image_path)
if img.mode != 'RGBA':
img = img.convert('RGBA')
# Composite onto white background so transparency doesn't bias the model
bg = Image.new('RGBA', img.size, (255, 255, 255, 255))
bg.paste(img, mask=img.split()[3] if img.mode == 'RGBA' else None)
img = bg.convert('RGB')
w, h = img.size
side = max(w, h)
square = Image.new('RGB', (side, side), (255, 255, 255))
square.paste(img, ((side - w) // 2, (side - h) // 2))
square = square.resize((_input_size, _input_size), Image.BICUBIC)
arr = np.array(square, dtype=np.float32)
# WD14 was trained on BGR
arr = arr[:, :, ::-1]
return arr[np.newaxis, :, :, :] # NHWC
def infer(image_path: str) -> list[dict]:
"""Run WD14 on one image. Returns a list of {name, category, confidence}."""
_load()
x = _preprocess(image_path)
out = _session.run([_output_name], {_input_name: x})[0][0]
results: list[dict] = []
for idx, score in enumerate(out):
meta = _tag_meta[idx]
results.append({
'name': meta['name'],
'category': meta['category'],
'confidence': float(score),
})
return results
def infer_filtered(image_path: str, min_any: float = 0.05) -> list[dict]:
"""Same as infer() but drops tags below a floor to keep DB rows reasonable."""
return [r for r in infer(image_path) if r['confidence'] >= min_any]
+60
View File
@@ -1,3 +1,5 @@
from pgvector.sqlalchemy import Vector
from . import db
# tag to object relationship table
@@ -218,3 +220,61 @@ class ImportTask(db.Model):
db.Index('ix_import_task_status_type', 'status', 'task_type'),
db.Index('ix_import_task_batch_status', 'batch_id', 'status'),
)
class ImageTagPrediction(db.Model):
__tablename__ = "image_tag_prediction"
id = db.Column(db.Integer, primary_key=True)
image_id = db.Column(db.Integer, db.ForeignKey("image_record.id", ondelete="CASCADE"), nullable=False)
tag_name = db.Column(db.Text, nullable=False)
tag_category = db.Column(db.Text, nullable=False) # general/character/copyright/rating/meta
confidence = db.Column(db.Float, nullable=False)
model_version = db.Column(db.Text, nullable=False)
created_at = db.Column(db.DateTime(timezone=True), server_default=db.func.now(), nullable=False)
__table_args__ = (
db.Index('idx_tag_predictions_image', 'image_id'),
db.Index('idx_tag_predictions_tag', 'tag_name'),
db.Index('idx_tag_predictions_model', 'model_version'),
)
class ImageEmbedding(db.Model):
__tablename__ = "image_embedding"
image_id = db.Column(db.Integer, db.ForeignKey("image_record.id", ondelete="CASCADE"), primary_key=True)
model_version = db.Column(db.Text, primary_key=True)
embedding = db.Column(Vector(1152), nullable=False)
created_at = db.Column(db.DateTime(timezone=True), server_default=db.func.now(), nullable=False)
class TagReferenceEmbedding(db.Model):
__tablename__ = "tag_reference_embedding"
tag_name = db.Column(db.Text, primary_key=True)
model_version = db.Column(db.Text, primary_key=True)
tag_kind = db.Column(db.Text, nullable=True)
centroid = db.Column(Vector(1152), nullable=False)
reference_count = db.Column(db.Integer, nullable=False)
computed_at = db.Column(db.DateTime(timezone=True), server_default=db.func.now(), nullable=False)
class SuggestionFeedback(db.Model):
__tablename__ = "suggestion_feedback"
id = db.Column(db.Integer, primary_key=True)
image_id = db.Column(db.Integer, db.ForeignKey("image_record.id", ondelete="CASCADE"), nullable=False)
tag_name = db.Column(db.Text, nullable=False)
suggestion_source = db.Column(db.Text, nullable=False) # wd14 | embedding_similarity
confidence = db.Column(db.Float, nullable=False)
decision = db.Column(db.Text, nullable=False) # accepted | rejected
decided_at = db.Column(db.DateTime(timezone=True), server_default=db.func.now(), nullable=False)
__table_args__ = (
db.Index('idx_feedback_image', 'image_id'),
db.Index('idx_feedback_tag', 'tag_name'),
)
class TagSuggestionConfig(db.Model):
__tablename__ = "tag_suggestion_config"
key = db.Column(db.Text, primary_key=True)
value = db.Column(db.Text, nullable=False)
description = db.Column(db.Text, nullable=True)
+1
View File
@@ -0,0 +1 @@
"""Read-path service modules called from Flask routes."""
+235
View File
@@ -0,0 +1,235 @@
"""Compute merged tag suggestions for an image on demand.
Sources:
- WD14 predictions filtered by per-category thresholds
- Embedding similarity vs. character centroids (for character tags only)
The result is grouped by category and annotated with whether each tag already
exists in the Tag table (the modal dims disallowed auto-creations).
"""
from __future__ import annotations
from typing import Iterable
from sqlalchemy import select
from app import db
from app.models import (
ImageTagPrediction,
ImageEmbedding,
TagReferenceEmbedding,
TagSuggestionConfig,
ImageRecord,
Tag,
image_tags,
)
# Tag kinds that receive embedding-similarity suggestions. None = general/topic tags.
# Adding 'artist' or 'series' here enables them with no other code changes.
ELIGIBLE_CENTROID_KINDS = ('character', 'fandom', None)
# Maps the WD14 prediction category ('character', 'copyright') to the kind used
# in the Tag table. 'general' maps to None and is handled separately at accept
# time. Lives here (not in main) so the canonicalization helper below can use it
# without a circular import.
_WD14_CATEGORY_TO_KIND = {
'character': 'character',
'copyright': 'fandom',
}
def _canonicalize_wd14_name(raw: str, category: str) -> str:
"""Rewrite a raw WD14 tag name to the canonical form ImageRepo persists.
- 'character' / 'copyright': title-case the display portion.
- everything else ('general', 'meta', unknown): underscore-to-space only,
preserving the user's casing for general topics.
"""
from app.utils.tag_names import normalize_display_name, underscores_to_spaces
kind = _WD14_CATEGORY_TO_KIND.get(category)
if ':' in raw:
prefix, rest = raw.split(':', 1)
transform = normalize_display_name if kind in ('character', 'fandom') else underscores_to_spaces
return f"{prefix}:{transform(rest)}"
return normalize_display_name(raw) if kind in ('character', 'fandom') else underscores_to_spaces(raw)
# Config keys with safe defaults (used if the row is missing from tag_suggestion_config).
_DEFAULTS = {
'threshold_general': 0.35,
'threshold_character_wd14': 0.75,
'threshold_copyright': 0.5,
'threshold_meta': 0.5,
'threshold_embedding': 0.85,
'min_reference_images': 5.0,
'wd14_model_version': 'wd-eva02-large-tagger-v3',
'siglip_model_version': 'siglip-so400m-patch14-384',
}
def _config() -> dict:
rows = TagSuggestionConfig.query.all()
cfg = dict(_DEFAULTS)
for r in rows:
try:
cfg[r.key] = float(r.value)
except ValueError:
cfg[r.key] = r.value
return cfg
def _category_threshold(category: str, cfg: dict) -> float:
return {
'general': cfg['threshold_general'],
'character': cfg['threshold_character_wd14'],
'copyright': cfg['threshold_copyright'],
'meta': cfg['threshold_meta'],
}.get(category, 1.01) # unknown/rating → effectively disabled (not surfaced)
def _existing_tag_names_for_image(image_id: int) -> set[str]:
"""Names of tags already attached to this image.
Names are returned as-stored; since every write path canonicalizes, any
row in the DB is already in canonical form. WD14 output is canonicalized
before lookup, so equality works.
"""
rows = (
db.session.query(Tag.name)
.join(image_tags, image_tags.c.tag_id == Tag.id)
.filter(image_tags.c.image_id == image_id)
.all()
)
return {row[0] for row in rows}
def _existing_tag_names() -> set[str]:
return {row[0] for row in db.session.query(Tag.name).all()}
def _wd14_suggestions(image_id: int, cfg: dict, already: set[str]) -> list[dict]:
wd14_ver = cfg.get('wd14_model_version', _DEFAULTS['wd14_model_version'])
preds = (
ImageTagPrediction.query
.filter_by(image_id=image_id, model_version=wd14_ver)
.all()
)
out: list[dict] = []
for p in preds:
threshold = _category_threshold(p.tag_category, cfg)
if p.confidence < threshold:
continue
canonical = _canonicalize_wd14_name(p.tag_name, p.tag_category)
if canonical in already:
continue
out.append({
'name': canonical,
'category': p.tag_category,
'confidence': p.confidence,
'source': 'wd14',
})
return out
# Maps tag kind (as stored on Tag and TagReferenceEmbedding) to the display
# category used in the modal UI's grouped suggestions.
_KIND_TO_DISPLAY_CATEGORY = {
'character': 'character',
'fandom': 'copyright',
None: 'general',
}
def _embedding_tag_suggestions(image_id: int, cfg: dict, already: set[str]) -> list[dict]:
siglip_ver = cfg.get('siglip_model_version', _DEFAULTS['siglip_model_version'])
min_refs = int(cfg.get('min_reference_images', 5))
threshold = cfg.get('threshold_embedding', 0.85)
image_emb = (
ImageEmbedding.query
.filter_by(image_id=image_id, model_version=siglip_ver)
.first()
)
if image_emb is None:
return []
# pgvector cosine distance; similarity = 1 - distance
distance = TagReferenceEmbedding.centroid.cosine_distance(image_emb.embedding)
rows = (
db.session.query(
TagReferenceEmbedding.tag_name,
TagReferenceEmbedding.tag_kind,
TagReferenceEmbedding.reference_count,
distance.label('distance'),
)
.filter(TagReferenceEmbedding.model_version == siglip_ver)
.filter(TagReferenceEmbedding.reference_count >= min_refs)
.order_by(distance)
.limit(40)
.all()
)
out: list[dict] = []
for tag_name, tag_kind, ref_count, dist in rows:
similarity = 1.0 - float(dist)
if similarity < threshold:
continue
if tag_name in already:
continue
category = _KIND_TO_DISPLAY_CATEGORY.get(tag_kind)
if category is None and tag_kind is not None:
# Unknown kind — skip defensively. (Should never happen because
# recompute_centroid only writes eligible kinds.)
continue
out.append({
'name': tag_name,
'category': category,
'confidence': similarity,
'source': 'embedding_similarity',
})
return out
def _merge(wd14: list[dict], embedding: list[dict]) -> list[dict]:
by_name: dict[str, dict] = {}
for s in wd14 + embedding:
existing = by_name.get(s['name'])
if existing is None or s['confidence'] > existing['confidence']:
by_name[s['name']] = s
return list(by_name.values())
def get_suggestions(image_id: int, top_k_per_category: int = 10) -> dict:
"""Return suggestions grouped by category for one image.
Shape:
{
'character': [{name, confidence, source, exists_in_db}, ...],
'copyright': [...],
'general': [...],
}
Ordered by confidence desc within each category. 'meta' and 'rating' are omitted.
"""
if ImageRecord.query.get(image_id) is None:
return {'character': [], 'copyright': [], 'general': []}
cfg = _config()
already = _existing_tag_names_for_image(image_id)
merged = _merge(
_wd14_suggestions(image_id, cfg, already),
_embedding_tag_suggestions(image_id, cfg, already),
)
existing_all = _existing_tag_names()
grouped: dict[str, list[dict]] = {'character': [], 'copyright': [], 'general': []}
for s in merged:
if s['category'] not in grouped:
continue # meta / rating / artist / unknown are dropped
s['exists_in_db'] = s['name'] in existing_all
grouped[s['category']].append(s)
for cat in grouped:
grouped[cat].sort(key=lambda x: x['confidence'], reverse=True)
grouped[cat] = grouped[cat][:top_k_per_category]
return grouped
+3 -2
View File
@@ -498,8 +498,9 @@
// Tag Refresh for Gallery Items
// ---------------------------
async function refreshItemTags(imageId) {
// Find the gallery item
const item = document.querySelector(`.img-clickable[data-id="${imageId}"]`);
// Only the gallery page renders tag overlays on thumbnails. Skip showcase items
// (class .masonry-item) so closing the modal doesn't inject a .tag-overlay there.
const item = document.querySelector(`.gallery-item[data-id="${imageId}"]`);
if (!item) return;
try {
+222 -12
View File
@@ -14,6 +14,11 @@ document.addEventListener('DOMContentLoaded', () => {
const tagInput = tagForm ? tagForm.querySelector('input[name="name"]') : null;
const tagAutocomplete = document.getElementById('tagAutocomplete');
const tagActionFeedback = document.getElementById('tagActionFeedback');
const provenanceSection = document.getElementById('modalProvenanceSection');
const provenanceList = document.getElementById('modalProvenanceList');
const suggestionsSection = document.getElementById('modalSuggestionsSection');
const suggestionsList = document.getElementById('modalSuggestionsList');
const suggestionsRefresh = document.getElementById('suggestionsRefreshBtn');
// Series info elements (when image IS in a series)
const seriesInfoEl = document.getElementById('modalSeriesInfo');
@@ -82,6 +87,15 @@ document.addEventListener('DOMContentLoaded', () => {
return tagEditor ? tagEditor.dataset.imageId : '';
}
function getProvenanceDisplayName(name, kind) {
// Artist tag names are "artist:name" — strip the prefix.
if (kind === 'artist' && name.startsWith('artist:')) return name.slice(7);
// Post tag names are "post:platform:artist:id" — strip "post:" only; the rest
// is a placeholder replaced by PostMetadata.title or post_id once it loads.
if (kind === 'post' && name.startsWith('post:')) return name.slice(5);
return name;
}
function renderTags(tags) {
if (!tagList) return;
tagList.innerHTML = '';
@@ -94,12 +108,206 @@ document.addEventListener('DOMContentLoaded', () => {
tagList.appendChild(chip);
});
}
const SUGGESTION_CATEGORY_ORDER = ['character', 'copyright', 'general'];
const SUGGESTION_CATEGORY_LABEL = {
character: 'Characters',
copyright: 'Fandoms',
general: 'General',
};
function clearSuggestions() {
if (suggestionsList) suggestionsList.innerHTML = '';
if (suggestionsSection) suggestionsSection.style.display = 'none';
}
function buildSuggestionChip(suggestion, imageId) {
const chip = document.createElement('div');
chip.className = 'suggestion-chip';
chip.dataset.tagName = suggestion.name;
chip.dataset.source = suggestion.source;
chip.dataset.category = suggestion.category;
chip.dataset.confidence = String(suggestion.confidence);
chip.title = `${suggestion.name} (${suggestion.source})`;
const pct = Math.round(suggestion.confidence * 100);
chip.innerHTML = `
<button type="button" class="sugg-btn sugg-accept" aria-label="Accept">✓</button>
<span class="sugg-body">
<span class="sugg-label">${escapeHtml(suggestion.name)}</span>
<span class="sugg-confidence">${pct}%</span>
</span>
<button type="button" class="sugg-btn sugg-reject" aria-label="Reject">✕</button>
`;
const acceptBtn = chip.querySelector('.sugg-accept');
const rejectBtn = chip.querySelector('.sugg-reject');
if (acceptBtn) {
acceptBtn.addEventListener('click', () => acceptSuggestion(chip, imageId));
}
if (rejectBtn) {
rejectBtn.addEventListener('click', () => rejectSuggestion(chip, imageId));
}
return chip;
}
function renderSuggestions(grouped, imageId) {
if (!suggestionsSection || !suggestionsList) return;
suggestionsList.innerHTML = '';
let total = 0;
for (const cat of SUGGESTION_CATEGORY_ORDER) {
const items = (grouped && grouped[cat]) || [];
if (items.length === 0) continue;
total += items.length;
const group = document.createElement('div');
group.className = 'suggestions-category';
group.dataset.category = cat;
const label = document.createElement('div');
label.className = 'suggestions-category-label';
label.textContent = SUGGESTION_CATEGORY_LABEL[cat] || cat;
group.appendChild(label);
const chips = document.createElement('div');
chips.className = 'suggestions-chips';
for (const item of items) {
chips.appendChild(buildSuggestionChip(item, imageId));
}
group.appendChild(chips);
suggestionsList.appendChild(group);
}
suggestionsSection.style.display = total > 0 ? '' : 'none';
}
async function loadSuggestions(imageId) {
if (!imageId || !suggestionsSection) { clearSuggestions(); return; }
try {
const r = await fetch(`/image/${imageId}/suggestions`);
const j = await r.json();
if (j.ok) renderSuggestions(j.suggestions || {}, imageId);
else clearSuggestions();
} catch {
clearSuggestions();
}
}
function animateChipOut(chip) {
chip.classList.add('leaving');
setTimeout(() => { chip.remove(); }, 200);
}
async function acceptSuggestion(chip, imageId) {
if (chip.dataset.disabled === 'true') return;
chip.dataset.disabled = 'true';
chip.style.pointerEvents = 'none';
try {
const r = await fetch(`/image/${imageId}/suggestions/accept`, {
method: 'POST',
headers: { 'Content-Type': 'application/json' },
body: JSON.stringify({
tag_name: chip.dataset.tagName,
source: chip.dataset.source,
category: chip.dataset.category,
confidence: Number(chip.dataset.confidence) || 0,
}),
});
const j = await r.json();
if (!j.ok) {
chip.dataset.disabled = 'false';
chip.style.pointerEvents = '';
if (tagActionFeedback) {
tagActionFeedback.textContent = `Couldn't accept: ${j.error || 'error'}`;
}
return;
}
if (typeof loadTags === 'function') await loadTags(imageId);
if (tagActionFeedback) tagActionFeedback.textContent = `Added ${j.tag.name}`;
if (typeof window.refreshItemTags === 'function') window.refreshItemTags(imageId);
animateChipOut(chip);
} catch {
chip.dataset.disabled = 'false';
chip.style.pointerEvents = '';
}
}
async function rejectSuggestion(chip, imageId) {
chip.style.pointerEvents = 'none';
try {
await fetch(`/image/${imageId}/suggestions/reject`, {
method: 'POST',
headers: { 'Content-Type': 'application/json' },
body: JSON.stringify({
tag_name: chip.dataset.tagName,
source: chip.dataset.source,
confidence: Number(chip.dataset.confidence) || 0,
}),
});
animateChipOut(chip);
if (tagActionFeedback) tagActionFeedback.textContent = `Rejected ${chip.dataset.tagName}`;
} catch {
chip.style.pointerEvents = '';
}
}
function renderProvenance(tags) {
if (!provenanceSection || !provenanceList) return;
provenanceList.innerHTML = '';
if (!tags || tags.length === 0) {
provenanceSection.style.display = 'none';
return;
}
tags.forEach(t => {
const chip = document.createElement('a');
chip.className = 'provenance-chip';
chip.href = `/gallery?tag=${encodeURIComponent(t.name)}`;
chip.title = t.name;
chip.dataset.tagName = t.name;
chip.dataset.kind = t.kind;
const icon = getTagIcon(t.kind);
const placeholder = escapeHtml(getProvenanceDisplayName(t.name, t.kind));
chip.innerHTML = `${icon} <span class="provenance-label">${placeholder}</span> <span class="provenance-arrow">↗</span>`;
provenanceList.appendChild(chip);
if (t.kind === 'post') upgradeProvenanceChip(chip, t.name);
});
provenanceSection.style.display = '';
}
async function upgradeProvenanceChip(chip, tagName) {
try {
const r = await fetch(`/api/post-metadata/by-tag-name/${encodeURIComponent(tagName)}`);
const j = await r.json();
if (!j.ok || !j.metadata) return;
const pm = j.metadata;
if (pm.platform) chip.dataset.platform = pm.platform;
const label = pm.title || pm.post_id;
if (label) {
const labelEl = chip.querySelector('.provenance-label');
if (labelEl) labelEl.textContent = label;
}
} catch {
// Fetch failed — chip keeps its placeholder label and neutral accent.
}
}
async function loadTags(imageId) {
if (!imageId) { renderTags([]); return; }
if (!imageId) { renderTags([]); renderProvenance([]); return; }
try {
const r = await fetch(`/image/${imageId}/tags`);
const j = await r.json();
if (j.ok) renderTags(j.tags);
if (j.ok) {
const all = j.tags || [];
// Provenance: artist first, then post — non-editable, system-tagged origin info.
const provenanceTags = [
...all.filter(t => t.kind === 'artist'),
...all.filter(t => t.kind === 'post'),
];
// Editable list excludes provenance kinds and hidden kinds (source, archive).
const hiddenFromModal = new Set(['artist', 'post', 'source', 'archive']);
const editableTags = all.filter(t => !hiddenFromModal.has(t.kind));
renderTags(editableTags);
renderProvenance(provenanceTags);
loadSuggestions(imageId);
}
} catch {
// ignore
}
@@ -109,7 +317,9 @@ document.addEventListener('DOMContentLoaded', () => {
// Gallery item tag refresh (fallback if bulk-select.js not loaded)
// ---------------------------
async function refreshItemTagsFallback(imageId) {
const item = document.querySelector(`.img-clickable[data-id="${imageId}"]`);
// Only the gallery page renders tag overlays on thumbnails. Showcase thumbnails
// use a tag-count badge instead and must not have .tag-overlay forced onto them.
const item = document.querySelector(`.gallery-item[data-id="${imageId}"]`);
if (!item) return;
try {
@@ -601,15 +811,6 @@ document.addEventListener('DOMContentLoaded', () => {
modal.classList.add('active');
updateImage(index);
history.replaceState({ modalIndex: index }, '', window.location.href);
// Focus the tag input after modal opens (desktop only - avoids keyboard popup on mobile)
if (tagInput && !isTouchDevice()) {
setTimeout(() => tagInput.focus(), 100);
}
}
// Detect touch devices to avoid auto-focus keyboard popup
function isTouchDevice() {
return ('ontouchstart' in window) || (navigator.maxTouchPoints > 0);
}
function closeModal() {
@@ -623,6 +824,8 @@ document.addEventListener('DOMContentLoaded', () => {
// clear tag UI
setEditorImageId('');
renderTags([]);
renderProvenance([]);
clearSuggestions();
hideSeriesInfo();
hideAddToSeries();
history.replaceState({}, '', window.location.pathname + window.location.search);
@@ -821,4 +1024,11 @@ document.addEventListener('DOMContentLoaded', () => {
location.reload();
}
});
if (suggestionsRefresh) {
suggestionsRefresh.addEventListener('click', () => {
const currentId = getEditorImageId();
if (currentId) loadSuggestions(currentId);
});
}
});
+176
View File
@@ -840,6 +840,182 @@ header {
color: #ef4444;
}
/* Modal provenance section — system-generated, non-editable tags (artist, post) */
.modal-provenance-section {
border-bottom: none;
}
.modal-provenance-section .provenance-list {
display: flex;
flex-direction: column;
align-items: flex-start;
gap: 0.25rem;
}
.provenance-chip {
display: inline-flex;
align-items: center;
gap: 0.4rem;
padding: 2px 8px 2px 10px;
border-radius: 999px;
background: transparent;
color: var(--text-dim);
text-decoration: none;
border: 1px solid rgba(255, 255, 255, 0.15);
border-left-width: 2px;
border-left-color: rgba(255, 255, 255, 0.4);
white-space: nowrap;
overflow: hidden;
text-overflow: ellipsis;
font: 12px/1.6 system-ui, sans-serif;
max-width: clamp(8rem, 22vw, 14rem);
transition: border-color 0.15s ease, background 0.15s ease, color 0.15s ease;
}
.provenance-chip:hover {
border-color: rgba(255, 255, 255, 0.35);
background: rgba(255, 255, 255, 0.04);
color: var(--text);
}
.provenance-chip[data-platform="patreon"] {
border-left-color: #f96854;
}
.provenance-chip[data-platform="subscribestar"] {
border-left-color: #009cde;
}
.provenance-chip[data-platform="hentaifoundry"] {
border-left-color: #c8232c;
}
.provenance-chip .provenance-arrow {
opacity: 0.6;
font-size: 0.85em;
}
/* Modal suggestions section — ML-backed chips with accept/reject */
.modal-suggestions-section .section-header {
display: flex;
justify-content: space-between;
align-items: center;
margin-bottom: 0.4rem;
}
.modal-suggestions-section .section-title {
font-size: 0.85rem;
font-weight: 600;
color: var(--text-dim);
letter-spacing: 0.04em;
text-transform: uppercase;
}
.modal-suggestions-section .suggestions-refresh {
background: transparent;
border: none;
color: var(--text-dim);
cursor: pointer;
font-size: 0.9rem;
padding: 0 0.25rem;
}
.modal-suggestions-section .suggestions-refresh:hover {
color: var(--text);
}
.suggestions-list {
display: flex;
flex-direction: column;
gap: 0.35rem;
}
.suggestions-category {
display: flex;
flex-direction: column;
gap: 0.2rem;
}
.suggestions-category-label {
font-size: 0.7rem;
color: var(--text-dim);
opacity: 0.65;
letter-spacing: 0.05em;
text-transform: uppercase;
}
.suggestions-chips {
display: flex;
flex-direction: column;
gap: 4px;
}
.suggestion-chip {
display: flex;
align-items: stretch;
gap: 6px;
min-height: 36px;
padding: 0;
border-radius: 6px;
background: rgba(255, 255, 255, 0.02);
color: var(--text-dim);
border: 1px dashed rgba(255, 255, 255, 0.18);
font: 13px/1.3 system-ui, sans-serif;
transition: border-color 0.15s ease, background 0.15s ease, color 0.15s ease, opacity 0.15s ease, transform 0.18s ease;
}
.suggestion-chip:hover {
border-color: rgba(255, 255, 255, 0.35);
background: rgba(255, 255, 255, 0.04);
color: var(--text);
}
.suggestion-chip .sugg-body {
flex: 1;
display: flex;
align-items: center;
gap: 8px;
min-width: 0;
padding: 0 2px;
}
.suggestion-chip .sugg-label {
flex: 1;
overflow: hidden;
text-overflow: ellipsis;
white-space: nowrap;
}
.suggestion-chip .sugg-confidence {
font-size: 0.72rem;
opacity: 0.6;
font-variant-numeric: tabular-nums;
flex-shrink: 0;
}
.suggestion-chip .sugg-btn {
flex: 0 0 auto;
width: 32px;
min-height: 32px;
display: inline-flex;
align-items: center;
justify-content: center;
background: transparent;
border: 1px solid rgba(255, 255, 255, 0.15);
border-radius: 4px;
color: inherit;
cursor: pointer;
padding: 0;
font-size: 0.95rem;
line-height: 1;
margin: 1px;
transition: background 0.12s ease, color 0.12s ease, border-color 0.12s ease;
}
.suggestion-chip .sugg-btn.sugg-accept {
color: #8be78b;
border-color: rgba(139, 231, 139, 0.35);
}
.suggestion-chip .sugg-btn.sugg-accept:hover {
background: rgba(139, 231, 139, 0.12);
border-color: rgba(139, 231, 139, 0.6);
}
.suggestion-chip .sugg-btn.sugg-reject {
color: #e78b8b;
border-color: rgba(231, 139, 139, 0.35);
}
.suggestion-chip .sugg-btn.sugg-reject:hover {
background: rgba(231, 139, 139, 0.12);
border-color: rgba(231, 139, 139, 0.6);
}
.suggestion-chip.leaving {
opacity: 0;
transform: translateX(6px);
}
/* Tag form with autocomplete */
.tag-form {
display: flex;
+8
View File
@@ -364,6 +364,14 @@ def import_media_file(self, task_id: int):
if sidecar_path:
apply_sidecar_metadata.delay(record.id, sidecar_path)
# Queue ML inference (WD14 tagging + SigLIP embedding)
try:
from app.tasks.ml import tag_and_embed
tag_and_embed.apply_async(args=[record.id], queue='ml')
except Exception as ml_err:
# Never let an enqueue failure roll back an otherwise-successful import.
log.warning(f"Could not enqueue tag_and_embed for image {record.id}: {ml_err}")
# Update batch stats
if task.batch_id:
from app.tasks.scan import update_batch_stats
+66
View File
@@ -0,0 +1,66 @@
"""Maintenance-queue Celery tasks. Lightweight, user-triggered, non-ML."""
import logging
from celery import shared_task
from sqlalchemy import text
from app import db
from app.models import Tag
log = logging.getLogger(__name__)
@shared_task(
name='app.tasks.maintenance.sync_character_fandoms',
soft_time_limit=120,
time_limit=180,
)
def sync_character_fandoms():
"""For every character tag with a non-null fandom_id, attach the fandom
tag to every image tagged with the character but not the fandom.
Additive only — never removes fandom tags.
Returns a summary dict with counts.
"""
chars = (
Tag.query
.filter_by(kind='character')
.filter(Tag.fandom_id.isnot(None))
.all()
)
total_added = 0
failed = 0
for char in chars:
try:
result = db.session.execute(
text("""
INSERT INTO image_tags (image_id, tag_id)
SELECT it.image_id, :fandom_id
FROM image_tags it
WHERE it.tag_id = :char_id
AND NOT EXISTS (
SELECT 1 FROM image_tags it2
WHERE it2.image_id = it.image_id
AND it2.tag_id = :fandom_id
)
"""),
{'char_id': char.id, 'fandom_id': char.fandom_id},
)
total_added += result.rowcount or 0
db.session.commit()
except Exception as e:
db.session.rollback()
failed += 1
log.warning(
"sync_character_fandoms: char tag %s (%s) failed: %s",
char.id, char.name, e,
)
log.info(
"sync_character_fandoms: scanned %d characters, added %d image↔fandom links, %d failed",
len(chars), total_added, failed,
)
return {
'characters_scanned': len(chars),
'links_added': total_added,
'characters_failed': failed,
}
+239
View File
@@ -0,0 +1,239 @@
"""ML inference Celery tasks (WD14 tagging, SigLIP embedding, centroid recompute)."""
from __future__ import annotations
import logging
import os
import time
import numpy as np
from sqlalchemy import and_, func
from sqlalchemy.dialects.postgresql import insert as pg_insert
from app.celery_app import celery
from app import db
from app.models import (
ImageRecord,
Tag,
ImageTagPrediction,
ImageEmbedding,
TagReferenceEmbedding,
TagSuggestionConfig,
image_tags,
)
log = logging.getLogger('celery.tasks.ml')
# Minimum raw WD14 confidence to bother storing. Below this, rows are noise.
WD14_STORE_FLOOR = float(os.environ.get('WD14_STORE_FLOOR', '0.05'))
def _config_value(key: str, default: str) -> str:
row = TagSuggestionConfig.query.filter_by(key=key).first()
return row.value if row else default
@celery.task(bind=True, name='app.tasks.ml.tag_and_embed',
max_retries=2, default_retry_delay=60,
soft_time_limit=120, time_limit=180)
def tag_and_embed(self, image_id: int):
"""Run WD14 + SigLIP on one image and persist predictions and embedding."""
from app.ml import wd14, siglip # lazy import so web process never loads torch
image = ImageRecord.query.get(image_id)
if image is None:
log.warning(f"tag_and_embed: image {image_id} not found")
return {'status': 'missing'}
if not os.path.exists(image.filepath):
log.warning(f"tag_and_embed: file missing at {image.filepath}")
return {'status': 'file_missing'}
try:
t0 = time.time()
raw = wd14.infer_filtered(image.filepath, min_any=WD14_STORE_FLOOR)
embedding = siglip.infer(image.filepath)
t_inf = time.time() - t0
# Remove any pre-existing predictions/embedding for this image+model_version pair
db.session.query(ImageTagPrediction).filter(
ImageTagPrediction.image_id == image_id,
ImageTagPrediction.model_version == wd14.MODEL_VERSION,
).delete(synchronize_session=False)
db.session.query(ImageEmbedding).filter(
ImageEmbedding.image_id == image_id,
ImageEmbedding.model_version == siglip.MODEL_VERSION,
).delete(synchronize_session=False)
for pred in raw:
db.session.add(ImageTagPrediction(
image_id=image_id,
tag_name=pred['name'],
tag_category=pred['category'],
confidence=pred['confidence'],
model_version=wd14.MODEL_VERSION,
))
db.session.add(ImageEmbedding(
image_id=image_id,
model_version=siglip.MODEL_VERSION,
embedding=embedding.tolist(),
))
db.session.commit()
log.info(f"tag_and_embed: image {image_id} done in {t_inf:.2f}s ({len(raw)} predictions)")
return {'status': 'ok', 'predictions': len(raw), 'duration_s': t_inf}
except Exception as e:
db.session.rollback()
log.error(f"tag_and_embed failed for image {image_id}: {e}", exc_info=True)
raise self.retry(exc=e)
@celery.task(bind=True, name='app.tasks.ml.backfill',
soft_time_limit=None, time_limit=None)
def backfill(self, batch_size: int = 50, pause_seconds: float = 0.5):
"""Enqueue tag_and_embed for every image missing predictions or embeddings for the current model versions.
Uses keyset pagination on image_record.id so the loop moves forward monotonically.
Re-running after some tag_and_embed tasks have completed is safe — those images
simply drop out of the filter on the next run.
Runs on the ml queue with concurrency=1, so enqueued tag_and_embed tasks only
start executing after this task finishes. Keyset pagination prevents the loop
from seeing its own pending enqueues as "still missing" and re-enqueueing them.
"""
from app.ml.wd14 import MODEL_VERSION as WD14_VER
from app.ml.siglip import MODEL_VERSION as SIGLIP_VER
enqueued_total = 0
last_id = 0
while True:
q = (
db.session.query(ImageRecord.id)
.outerjoin(
ImageTagPrediction,
and_(
ImageTagPrediction.image_id == ImageRecord.id,
ImageTagPrediction.model_version == WD14_VER,
),
)
.outerjoin(
ImageEmbedding,
and_(
ImageEmbedding.image_id == ImageRecord.id,
ImageEmbedding.model_version == SIGLIP_VER,
),
)
.filter(ImageRecord.id > last_id)
.filter(
(ImageTagPrediction.image_id.is_(None))
| (ImageEmbedding.image_id.is_(None))
)
.order_by(ImageRecord.id)
.limit(batch_size)
)
ids = [row[0] for row in q.all()]
if not ids:
break
for image_id in ids:
tag_and_embed.apply_async(args=[image_id], queue='ml')
enqueued_total += len(ids)
last_id = ids[-1]
log.info(f"backfill: enqueued batch of {len(ids)} (total {enqueued_total}, last_id {last_id})")
time.sleep(pause_seconds)
log.info(f"backfill: complete, enqueued {enqueued_total} images")
return {'status': 'ok', 'enqueued': enqueued_total}
@celery.task(bind=True, name='app.tasks.ml.recompute_centroid',
soft_time_limit=60, time_limit=120)
def recompute_centroid(self, tag_name: str):
"""Recompute the mean embedding for an eligible tag from its currently-associated images.
Eligible tag kinds are defined by ELIGIBLE_CENTROID_KINDS in
app.services.tag_suggestions. Tags of any other kind return early with
status='ineligible_kind' and do not write a centroid row.
"""
from app.ml.siglip import MODEL_VERSION as SIGLIP_VER
from app.services.tag_suggestions import ELIGIBLE_CENTROID_KINDS
tag = Tag.query.filter_by(name=tag_name).first()
if tag is None:
log.warning(f"recompute_centroid: tag {tag_name!r} not found")
return {'status': 'missing_tag'}
if tag.kind not in ELIGIBLE_CENTROID_KINDS:
log.info(f"recompute_centroid: tag {tag_name!r} has ineligible kind {tag.kind!r}")
return {'status': 'ineligible_kind'}
rows = (
db.session.query(ImageEmbedding.embedding)
.join(image_tags, image_tags.c.image_id == ImageEmbedding.image_id)
.filter(image_tags.c.tag_id == tag.id)
.filter(ImageEmbedding.model_version == SIGLIP_VER)
.all()
)
if not rows:
log.info(f"recompute_centroid: no embeddings yet for {tag_name}")
return {'status': 'no_embeddings'}
vectors = np.array([np.array(r[0], dtype=np.float32) for r in rows])
centroid = vectors.mean(axis=0)
stmt = pg_insert(TagReferenceEmbedding).values(
tag_name=tag_name,
tag_kind=tag.kind,
model_version=SIGLIP_VER,
centroid=centroid.tolist(),
reference_count=len(rows),
computed_at=func.now(),
)
stmt = stmt.on_conflict_do_update(
index_elements=['tag_name', 'model_version'],
set_={
'tag_kind': stmt.excluded.tag_kind,
'centroid': stmt.excluded.centroid,
'reference_count': stmt.excluded.reference_count,
'computed_at': func.now(),
},
)
db.session.execute(stmt)
db.session.commit()
log.info(f"recompute_centroid: {tag_name} (kind={tag.kind}) -> n={len(rows)}")
return {'status': 'ok', 'reference_count': len(rows), 'tag_kind': tag.kind}
@celery.task(bind=True, name='app.tasks.ml.recompute_all_centroids',
soft_time_limit=None, time_limit=None)
def recompute_all_centroids(self):
"""Enqueue recompute_centroid for every eligible tag with enough reference images.
Uses a single aggregate query to find tags with >= min_reference_images applied
images, then enqueues one recompute_centroid task per tag on the ml queue.
"""
from app.services.tag_suggestions import ELIGIBLE_CENTROID_KINDS
min_refs_row = TagSuggestionConfig.query.filter_by(key='min_reference_images').first()
min_refs = int(min_refs_row.value) if min_refs_row else 5
# Build an IS NULL / IN filter that covers ELIGIBLE_CENTROID_KINDS including None.
kinds_not_null = [k for k in ELIGIBLE_CENTROID_KINDS if k is not None]
allow_null = None in ELIGIBLE_CENTROID_KINDS
kind_filter = Tag.kind.in_(kinds_not_null)
if allow_null:
kind_filter = kind_filter | Tag.kind.is_(None)
rows = (
db.session.query(Tag.name, func.count(image_tags.c.image_id).label('n'))
.join(image_tags, image_tags.c.tag_id == Tag.id)
.filter(kind_filter)
.group_by(Tag.name)
.having(func.count(image_tags.c.image_id) >= min_refs)
.all()
)
enqueued = 0
for tag_name, n in rows:
recompute_centroid.apply_async(args=[tag_name], queue='ml')
enqueued += 1
log.info(f"recompute_all_centroids: enqueued {enqueued} tags (min_refs={min_refs})")
return {'status': 'ok', 'enqueued': enqueued, 'min_refs': min_refs}
+14
View File
@@ -47,6 +47,11 @@
</div>
</div>
<!-- SECTION: Provenance (system-generated tags: post; future: source, archive) -->
<div class="modal-sidebar-section modal-provenance-section" id="modalProvenanceSection" style="display: none;">
<div id="modalProvenanceList" class="provenance-list"></div>
</div>
<!-- SECTION: Tags -->
<div class="modal-sidebar-section modal-tags-section">
<div id="modalTagList" class="tags"></div>
@@ -60,6 +65,15 @@
<span id="tagActionFeedback" class="tag-feedback tag-feedback-inline"></span>
</div>
<!-- SECTION: Suggestions (ML-backed, modal-only) -->
<div class="modal-sidebar-section modal-suggestions-section" id="modalSuggestionsSection" style="display: none;">
<div class="section-header">
<span class="section-title">Suggestions</span>
<button type="button" class="suggestions-refresh" id="suggestionsRefreshBtn" aria-label="Refresh suggestions"></button>
</div>
<div id="modalSuggestionsList" class="suggestions-list"></div>
</div>
</div>
</div>
</div>
+29
View File
@@ -436,6 +436,35 @@
</p>
</div>
</div>
<section class="settings-section">
<h3>ML tag suggestions</h3>
<p class="settings-hint">
Run inference on all images missing predictions or embeddings for the current model versions.
Resumable and safe to re-run. Backfill enqueues work on the <code>ml</code> queue; monitor progress in the ml-worker logs.
</p>
<form action="{{ url_for('main.trigger_ml_backfill') }}" method="post" onsubmit="return confirm('Enqueue ML backfill for all unprocessed images?');">
<button type="submit" class="btn secondary-btn">Run ML backfill</button>
</form>
<p class="settings-hint" style="margin-top: 1em;">
Recompute centroids for all eligible tags (character, fandom, and general/topic) that have at least
the minimum number of reference images. Run this after the initial backfill, or any time you've
applied many tags manually and want the suggestions to catch up.
</p>
<form action="{{ url_for('main.trigger_recompute_all_centroids') }}" method="post" onsubmit="return confirm('Enqueue centroid recompute for all eligible tags?');">
<button type="submit" class="btn secondary-btn">Recompute all centroids</button>
</form>
<p class="settings-hint" style="margin-top: 1em;">
Attach each character's fandom to every image already tagged with that character.
Additive only — existing fandom tags are never removed. Run after you've set a
fandom on an existing character, or whenever you suspect drift.
</p>
<form action="{{ url_for('main.trigger_sync_character_fandoms') }}" method="post" onsubmit="return confirm('Enqueue character↔fandom sync for all tagged images?');">
<button type="submit" class="btn secondary-btn">Sync character fandoms to images</button>
</form>
</section>
</div>
</div>
</div><!-- end maintenance tab panel -->
+49
View File
@@ -0,0 +1,49 @@
"""Normalization helpers for Tag display names.
Keeps the first letter of every whitespace-separated word upper-case and
preserves internal punctuation. Called from every character/fandom write
path so typos like "misty" don't accumulate alongside "Misty".
"""
def underscores_to_spaces(s: str) -> str:
"""Replace every underscore with a space.
Used by general/null-kind write paths and as a pre-step inside
normalize_display_name so WD14's 'misty_ketchum' and a hand-typed
'Misty Ketchum' collapse onto the same tag name.
Examples:
"big_hair" -> "big hair"
"long_blue_hair" -> "long blue hair"
"" -> ""
"""
return s.replace('_', ' ') if s else s
def normalize_display_name(s: str) -> str:
"""Title-case each whitespace-separated word. Preserves internal punctuation.
Underscores are converted to spaces first so WD14-style names normalize
identically to hand-typed names.
Examples:
"misty" -> "Misty"
"misty ketchum" -> "Misty Ketchum"
"misty_ketchum" -> "Misty Ketchum"
"o'brien" -> "O'brien"
"mary-kate" -> "Mary-kate"
"" -> ""
" " -> " " (only whitespace — returned unchanged)
"""
if not s:
return s
s = underscores_to_spaces(s)
parts = s.split(' ')
out: list[str] = []
for p in parts:
if p == '':
out.append(p)
else:
out.append(p[:1].upper() + p[1:])
return ' '.join(out)