# 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

```sql
CREATE EXTENSION IF NOT EXISTS vector;
```

### New tables

```sql
-- 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:

```sql
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.