imagerepo/docs/superpowers/tag-suggestion-system-spec.md

Tag Suggestion System - Technical Spec

Overview

Add an ML-powered tag suggestion system to the existing gallery application. The system suggests tags for images in two ways:

1. Direct classification via WD14 for general tags, ratings, popular characters, and series
2. Embedding similarity for OC/niche character identification based on previously-tagged reference images

The system produces suggestions only - never auto-applies tags. All suggestions require user confirmation. Artists are explicitly out of scope; they are set at import time from source metadata.

Scale & Context

• ~62,000 existing images (anime/drawn art exclusively)
• ~300 tagged characters currently, expected to grow
• Single-user or small-user tagging workflow
• Target hardware: NVIDIA Tesla P4 GPU (8GB VRAM, compute-only, no NVDEC needed for this workload)
• Existing app is homegrown; this is an addition, not a rewrite

Out of Scope

• Artist identification (handled at import from source metadata)
• Auto-applying tags without user confirmation
• Training or fine-tuning models (inference only; "learning" happens via reference set growth)
• Multi-user permission models for suggestions

Models

Tagger: WD14 (SmilingWolf EVA02-Large variant)

• Model: SmilingWolf/wd-eva02-large-tagger-v3 (or current best available)
• Purpose: Produces confidence-scored tags across categories (general, character, copyright/series, rating, meta)
• Runs on GPU via ONNX Runtime with CUDA execution provider
• Output: tag name + category + confidence score (0.0-1.0)

Embedding Model: SigLIP or OpenCLIP

• Recommended: google/siglip-so400m-patch14-384 (strong performance, reasonable size)
• Alternative: laion/CLIP-ViT-H-14-laion2B-s32B-b79K if SigLIP integration is problematic
• Output: Fixed-dimension embedding vector per image (SigLIP SO400M: 1152-dim)
• Purpose: Enables similarity search for character matching against reference set

Model Versioning

Store the model identifier (name + version) alongside every output. Schema must support multiple model versions coexisting to allow migration without requiring immediate reprocessing of all 62K images.

Database Schema Changes

Assumes Postgres. If the existing app uses a different database, adapt accordingly - pgvector is strongly preferred for embedding storage at this scale.

New extension

CREATE EXTENSION IF NOT EXISTS vector;

New tables

-- Raw WD14 tagger output, one row per (image, tag, model_version)
CREATE TABLE image_tag_predictions (
    id BIGSERIAL PRIMARY KEY,
    image_id BIGINT NOT NULL REFERENCES images(id) ON DELETE CASCADE,
    tag_name TEXT NOT NULL,
    tag_category TEXT NOT NULL,  -- 'general', 'character', 'copyright', 'rating', 'meta'
    confidence REAL NOT NULL,
    model_version TEXT NOT NULL,  -- e.g. 'wd-eva02-large-tagger-v3'
    created_at TIMESTAMPTZ NOT NULL DEFAULT NOW()
);

CREATE INDEX idx_tag_predictions_image ON image_tag_predictions(image_id);
CREATE INDEX idx_tag_predictions_tag ON image_tag_predictions(tag_name);
CREATE INDEX idx_tag_predictions_model ON image_tag_predictions(model_version);

-- Image embeddings for similarity search
CREATE TABLE image_embeddings (
    image_id BIGINT NOT NULL REFERENCES images(id) ON DELETE CASCADE,
    model_version TEXT NOT NULL,
    embedding vector(1152) NOT NULL,  -- adjust dimension to match chosen model
    created_at TIMESTAMPTZ NOT NULL DEFAULT NOW(),
    PRIMARY KEY (image_id, model_version)
);

-- HNSW index for fast nearest-neighbor search
CREATE INDEX idx_embeddings_hnsw ON image_embeddings 
    USING hnsw (embedding vector_cosine_ops);

-- Cached centroid embeddings per character, for fast suggestion queries
CREATE TABLE character_reference_embeddings (
    character_tag TEXT NOT NULL,
    model_version TEXT NOT NULL,
    centroid vector(1152) NOT NULL,
    reference_count INTEGER NOT NULL,  -- how many confirmed images contributed
    computed_at TIMESTAMPTZ NOT NULL DEFAULT NOW(),
    PRIMARY KEY (character_tag, model_version)
);

-- Tracks user decisions on suggestions for future tuning/reranking
CREATE TABLE suggestion_feedback (
    id BIGSERIAL PRIMARY KEY,
    image_id BIGINT NOT NULL REFERENCES images(id) ON DELETE CASCADE,
    tag_name TEXT NOT NULL,
    suggestion_source TEXT NOT NULL,  -- 'wd14' or 'embedding_similarity'
    confidence REAL NOT NULL,
    decision TEXT NOT NULL,  -- 'accepted' or 'rejected'
    decided_at TIMESTAMPTZ NOT NULL DEFAULT NOW()
);

CREATE INDEX idx_feedback_image ON suggestion_feedback(image_id);
CREATE INDEX idx_feedback_tag ON suggestion_feedback(tag_name);

Configuration table (or use existing config system)

Store these as tunable values, not hardcoded:

CREATE TABLE tag_suggestion_config (
    key TEXT PRIMARY KEY,
    value TEXT NOT NULL,
    description TEXT
);

-- Initial values (thresholds will need tuning)
INSERT INTO tag_suggestion_config VALUES
    ('threshold_general', '0.35', 'Min confidence for general tag suggestions'),
    ('threshold_character_wd14', '0.75', 'Min confidence for WD14 character tags'),
    ('threshold_copyright', '0.5', 'Min confidence for copyright/series tags'),
    ('threshold_rating', '0.5', 'Min confidence for rating tags'),
    ('threshold_meta', '0.5', 'Min confidence for meta tags'),
    ('threshold_embedding_character', '0.85', 'Min cosine similarity for character via embeddings'),
    ('min_reference_images_per_character', '5', 'Characters with fewer confirmed refs are excluded from embedding suggestions'),
    ('centroid_recompute_delta', '3', 'Recompute character centroid after N new confirmed references');

Inference Pipeline

On upload (per-image, near-realtime)

1. WD14 inference → insert rows into image_tag_predictions
2. Embedding model inference → insert row into image_embeddings
3. Suggestions are generated on-demand at view time (see below); no storage of suggestions themselves needed

Initial backfill (one-time, for existing 62K images)

• Background job that iterates all images without predictions/embeddings
• Batched GPU inference (WD14 supports batch sizes; 8-16 is a reasonable starting point on P4)
• Expect several hours of runtime; run overnight
• Must be resumable (track last-processed image_id, handle failures gracefully)
• Progress reporting (log every N images, estimated time remaining)

On confirmed character tag (user accepts a character tag on an image)

• Increment a counter for that character's reference set
• If delta since last centroid computation exceeds centroid_recompute_delta, schedule centroid recomputation for that character
• Recomputation: fetch all embeddings for images confirmed to have this character tag, compute mean (or medoid if desired), update character_reference_embeddings

Suggestion generation (at view time or via background job)

When generating suggestions for an image:

From WD14 predictions:

• Query image_tag_predictions for this image
• Filter by per-category confidence thresholds
• Exclude tags already confirmed on the image

From embedding similarity (for characters only):

• Fetch this image's embedding
• Query character_reference_embeddings for all characters with reference_count >= min_reference_images_per_character
• Compute cosine similarity against each character's centroid
• Filter by threshold_embedding_character
• Exclude characters already confirmed on the image
• Return top K (suggest K=5 initially)

Merging:

• WD14 character suggestions and embedding-based character suggestions may overlap; dedupe by tag name, prefer whichever has higher normalized confidence
• Return merged suggestion list grouped by category, sorted by confidence within each category

Background Jobs

Three scheduled/triggered jobs:

1. Backfill job (one-time initially, resumable)

• Process any image missing WD14 predictions or embeddings for current model versions
• Can be re-run after model version change to update stored data

2. Centroid recomputation (triggered or nightly)

• For each character whose reference_count has changed by more than centroid_recompute_delta since last computation, recompute centroid
• Fast operation; can run frequently

3. Upload processing (per-image, on upload)

• Runs WD14 + embedding inference on newly uploaded images
• Should not block the upload response; queue for async processing

API / Integration Points

The spec doesn't prescribe exact endpoints since it depends on the existing app's conventions, but the following operations must be exposed:

• get_suggestions(image_id) -> list of suggested tags grouped by category - called when viewing an image's tag panel
• accept_suggestion(image_id, tag_name, source) - records feedback, applies tag to image, triggers centroid recomputation if character
• reject_suggestion(image_id, tag_name, source) - records feedback only
• queue_backfill() - admin operation to process any unprocessed images
• get_character_reference_stats() -> list of (character, reference_count) - admin visibility into which characters have enough references to participate in embedding suggestions

Implementation Notes & Decisions

Why not store suggestions in the database?

Suggestions are a function of (predictions + embeddings + current confirmed tags + current thresholds + current centroids). Any of these can change. Storing suggestions creates a cache invalidation problem. Instead, compute them on read - the operations are cheap once predictions and embeddings exist.

Why store raw WD14 predictions at all?

Two reasons:

1. Threshold tuning without re-running inference. Change thresholds, suggestions update immediately.
2. Model upgrade path. When moving from model v3 to v4, old predictions remain queryable until new ones are generated, avoiding a blackout period.

Why centroids instead of querying all reference images?

At 300 characters with potentially thousands of references each, querying per-image embeddings for every character at suggestion time gets expensive. Centroids reduce it to one vector per character. Trade-off: loses some nuance (a character with multiple distinct outfits is represented as one "average" point). If this becomes a problem, options include:

• Multiple centroids per character via clustering (k-medoids with small k)
• Fallback to top-K raw embedding search for characters with high intra-class variance

Not worth building upfront. Revisit if suggestion quality suffers.

Why track rejection feedback?

Stored feedback enables:

• Per-tag or per-category threshold tuning based on acceptance rates
• Future reranker training if suggestion quality plateaus
• Diagnosing which characters consistently produce bad suggestions

Low cost to collect; high value if needed later.

Embedding dimension

Adjust vector(1152) in schema to match whatever model is actually selected. SigLIP SO400M is 1152-dim; ViT-H/14 CLIP is 1024-dim; some other variants differ. Mismatch will cause pgvector errors.

GPU considerations

• P4 is compute-only (no NVDEC needed for this workload - we're doing ML inference, not video decode)
• WD14 ONNX model fits comfortably in 8GB VRAM
• SigLIP SO400M inference fits in 8GB with room to spare
• If co-located with other GPU workloads (e.g., Ollama), use CUDA_VISIBLE_DEVICES to pin to a specific GPU and avoid contention

Tuning Phase (post-implementation)

After initial deployment, expect to spend time tuning:

1. Per-category confidence thresholds (collect feedback data first, then analyze acceptance rate vs. confidence curves)
2. Per-character embedding thresholds (some characters are distinctive, some generic-looking)
3. Minimum reference count before a character enters the suggestion pool
4. Centroid recomputation cadence

Build thresholds as config values (see tag_suggestion_config table), not constants. Expose them via whatever admin interface the existing app has.

Open Questions for Implementation

• Which ONNX Runtime integration approach fits the existing app's stack (Python FastAPI, Rust, Node, etc.)?
• Is there an existing background job system in the app, or does one need to be added?
• What's the existing app's approach to config management - env vars, DB table, config file?

These should be answered by inspecting the existing codebase before implementation begins.