imagerepo/app/tasks/ml.py

"""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'))

# Number of frames to sample from each video for ML inference.
VIDEO_ML_FRAMES = int(os.environ.get('VIDEO_ML_FRAMES', '10'))


def _config_value(key: str, default: str) -> str:
    row = TagSuggestionConfig.query.filter_by(key=key).first()
    return row.value if row else default


def _run_video_inference(image, wd14, siglip) -> tuple[list[dict], np.ndarray] | None:
    """Sample VIDEO_ML_FRAMES frames from the video and aggregate predictions.

    WD14 predictions are aggregated by taking the max confidence per (name, category)
    across frames — a character appearing clearly in one frame shouldn't be diluted by
    frames where it's absent. SigLIP embeddings are mean-pooled, which preserves cosine
    distance behavior since the pgvector index normalizes as needed.

    Returns (predictions, embedding) or None if frame extraction fails.
    """
    import shutil
    from app.utils.image_importer import extract_video_frames

    frame_paths = extract_video_frames(image.filepath, count=VIDEO_ML_FRAMES)
    if not frame_paths:
        return None

    tmpdir = os.path.dirname(frame_paths[0])
    try:
        best: dict[tuple[str, str], dict] = {}
        embeddings: list[np.ndarray] = []
        for fp in frame_paths:
            raw = wd14.infer_filtered(fp, min_any=WD14_STORE_FLOOR)
            for pred in raw:
                key = (pred['name'], pred['category'])
                prev = best.get(key)
                if prev is None or pred['confidence'] > prev['confidence']:
                    best[key] = pred
            embeddings.append(siglip.infer(fp))

        if not embeddings:
            return None
        mean_emb = np.stack(embeddings).mean(axis=0).astype(np.float32)
        return list(best.values()), mean_emb
    finally:
        shutil.rmtree(tmpdir, ignore_errors=True)


def _ensure_video_thumb(image) -> None:
    """Generate a thumbnail if the video record is missing one, and persist the path."""
    from app.utils.image_importer import generate_video_thumbnail_mirrored

    if image.thumb_path and os.path.exists(image.thumb_path):
        return
    new_thumb = generate_video_thumbnail_mirrored(image.filepath)
    if new_thumb:
        image.thumb_path = new_thumb


@celery.task(bind=True, name='app.tasks.ml.tag_and_embed',
             max_retries=2, default_retry_delay=60,
             soft_time_limit=240, time_limit=360)
def tag_and_embed(self, image_id: int):
    """Run WD14 + SigLIP on one image (or sampled video frames) and persist results."""
    from app.ml import wd14, siglip  # lazy import so web process never loads torch
    from app.utils.image_importer import VIDEO_EXTS

    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'}

    # Skip flagged-corrupt rows. The verifier already saw the file fail
    # structurally; running inference would either crash or produce noise.
    # 'unknown' (pre-sweep / mid-import) still proceeds.
    if image.integrity_status not in ('ok', 'unknown'):
        log.info(
            f"tag_and_embed: skipping image {image_id} — integrity={image.integrity_status}"
        )
        return {'status': 'skipped_integrity', 'integrity': image.integrity_status}

    is_video = image.filepath.lower().endswith(VIDEO_EXTS)

    try:
        t0 = time.time()
        if is_video:
            result = _run_video_inference(image, wd14, siglip)
            if result is None:
                log.warning(f"tag_and_embed: video {image_id} produced no frames")
                return {'status': 'no_frames'}
            raw, embedding = result
            _ensure_video_thumb(image)
        else:
            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()
        kind = 'video' if is_video else 'image'
        log.info(f"tag_and_embed: {kind} {image_id} done in {t_inf:.2f}s ({len(raw)} predictions)")
        return {'status': 'ok', 'predictions': len(raw), 'duration_s': t_inf, 'kind': kind}
    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(ImageRecord.integrity_status.in_(('ok', 'unknown')))
            .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}