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


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}