feat(fc2b): importer enqueues tag_and_embed + ml-worker model self-heal

import_media_file now enqueues tag_and_embed alongside generate_thumbnail
after a successful import. scripts/download_models.py snapshots Camie +
SigLIP into /models, idempotent (skips when present). The ml-worker
entrypoint runs it before starting the Celery worker so a fresh /models
volume self-heals on first boot. Downloader tests are pure-logic (no
network in CI).

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
This commit is contained in:
2026-05-15 07:45:09 -04:00
parent bb48845268
commit b68a382b60
5 changed files with 104 additions and 1 deletions
+1
View File
@@ -0,0 +1 @@
"""Operational scripts run from the container entrypoint or CLI."""
+53
View File
@@ -0,0 +1,53 @@
"""Self-heal model weights into /models. Idempotent — only missing files
are fetched. Called by the ml-worker entrypoint before Celery starts.
"""
import os
import sys
from pathlib import Path
MODEL_ROOT = Path(os.environ.get("ML_MODEL_DIR", "/models"))
CAMIE_REPO = os.environ.get("CAMIE_HF_REPO", "Camais03/camie-tagger-v2")
SIGLIP_REPO = os.environ.get(
"SIGLIP_HF_REPO", "google/siglip-so400m-patch14-384"
)
def _snapshot(repo_id: str, dest: Path, allow_patterns: list[str] | None) -> None:
from huggingface_hub import snapshot_download
dest.mkdir(parents=True, exist_ok=True)
snapshot_download(
repo_id=repo_id,
local_dir=str(dest),
allow_patterns=allow_patterns,
)
def ensure_camie() -> None:
dest = MODEL_ROOT / "camie"
if (dest / "model.onnx").is_file() and (dest / "selected_tags.csv").is_file():
print(f"[download_models] Camie present at {dest}")
return
print(f"[download_models] Fetching {CAMIE_REPO} -> {dest}")
_snapshot(CAMIE_REPO, dest, ["model.onnx", "selected_tags.csv", "*.json"])
def ensure_siglip() -> None:
dest = MODEL_ROOT / "siglip"
if (dest / "config.json").is_file() and any(dest.glob("*.safetensors")):
print(f"[download_models] SigLIP present at {dest}")
return
print(f"[download_models] Fetching {SIGLIP_REPO} -> {dest}")
_snapshot(SIGLIP_REPO, dest, None)
def main() -> int:
ensure_camie()
ensure_siglip()
print("[download_models] Done.")
return 0
if __name__ == "__main__":
sys.exit(main())
+3 -1
View File
@@ -95,11 +95,13 @@ def import_media_file(self, import_task_id: int) -> dict:
session.add(task)
session.commit()
# Enqueue the thumbnail task for newly imported images.
# Enqueue the thumbnail + ML tasks for newly imported images.
if result.status == "imported" and result.image_id is not None:
from .ml import tag_and_embed
from .thumbnail import generate_thumbnail
generate_thumbnail.delay(result.image_id)
tag_and_embed.delay(result.image_id)
# If this was the last task in the batch, mark the batch complete.
remaining = session.execute(
+2
View File
@@ -37,6 +37,8 @@ case "$ROLE" in
;;
ml-worker)
echo "[entrypoint] Ensuring ML models present in /models..."
python -m backend.app.scripts.download_models
echo "[entrypoint] Starting ML Celery worker (ml queue)"
exec celery -A backend.app.celery_app:celery worker \
--loglevel=info \
+45
View File
@@ -0,0 +1,45 @@
"""download_models tests. No network in CI; we test the 'already present →
skip' short-circuit by faking the expected files, and that main() wires
both ensure_* calls.
"""
from unittest.mock import patch
from backend.app.scripts import download_models as dm
def test_ensure_camie_skips_when_present(tmp_path, monkeypatch):
monkeypatch.setattr(dm, "MODEL_ROOT", tmp_path)
camie = tmp_path / "camie"
camie.mkdir(parents=True)
(camie / "model.onnx").write_bytes(b"x")
(camie / "selected_tags.csv").write_text("tag_id,name,category,count\n")
with patch.object(dm, "_snapshot") as snap:
dm.ensure_camie()
snap.assert_not_called()
def test_ensure_camie_downloads_when_missing(tmp_path, monkeypatch):
monkeypatch.setattr(dm, "MODEL_ROOT", tmp_path)
with patch.object(dm, "_snapshot") as snap:
dm.ensure_camie()
snap.assert_called_once()
def test_ensure_siglip_skips_when_present(tmp_path, monkeypatch):
monkeypatch.setattr(dm, "MODEL_ROOT", tmp_path)
sig = tmp_path / "siglip"
sig.mkdir(parents=True)
(sig / "config.json").write_text("{}")
(sig / "model.safetensors").write_bytes(b"x")
with patch.object(dm, "_snapshot") as snap:
dm.ensure_siglip()
snap.assert_not_called()
def test_main_calls_both(monkeypatch):
calls = []
monkeypatch.setattr(dm, "ensure_camie", lambda: calls.append("camie"))
monkeypatch.setattr(dm, "ensure_siglip", lambda: calls.append("siglip"))
assert dm.main() == 0
assert calls == ["camie", "siglip"]