feat: ml-worker self-heals models; restore modal tag-input auto-focus

ml-worker now fetches WD14 + SigLIP on container start via an entrypoint script that calls huggingface_hub with a populated local_dir. HF does a content-based integrity check, so subsequent starts are a no-op. Transient HF outages warn and continue when a local copy is already present. This drops the baked-in model layer (~4 GB) from the image so it pushes cleanly through registry proxies and no longer bloats every tagged release. view-modal.js restores the desktop-only auto-focus of the tag input on openModal (with touch-device guard). The block existed in an earlier main commit (0bf1961) but was absent from the post-corruption consolidation commit. Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
2026-04-19 21:12:43 -04:00
parent 0f35a0c484
commit fb925f5e6a
5 changed files with 105 additions and 10 deletions
@@ -0,0 +1,84 @@
+"""Ensure WD14 and SigLIP model files exist under $ML_MODEL_DIR.
+
+Run at container start by entrypoint-ml.sh. huggingface_hub performs content-based
+integrity checks: matching files are skipped, changed/missing files are fetched.
+If the remote is unreachable but a local copy is already present, log a warning
+and continue so a transient HF outage does not take a healthy worker down.
+"""
+from __future__ import annotations
+
+import logging
+import os
+import sys
+
+from huggingface_hub import hf_hub_download, snapshot_download
+from huggingface_hub.utils import HfHubHTTPError
+
+logging.basicConfig(level=logging.INFO, format='%(asctime)s %(levelname)s %(name)s: %(message)s')
+log = logging.getLogger('ensure_models')
+
+MODEL_DIR = os.environ.get('ML_MODEL_DIR', '/models')
+WD14_REPO = os.environ.get('WD14_REPO', 'SmilingWolf/wd-eva02-large-tagger-v3')
+SIGLIP_REPO = os.environ.get('SIGLIP_REPO', 'google/siglip-so400m-patch14-384')
+WD14_REVISION = os.environ.get('WD14_REVISION') or None
+SIGLIP_REVISION = os.environ.get('SIGLIP_REVISION') or None
+
+WD14_FILES = ('model.onnx', 'selected_tags.csv')
+SIGLIP_CORE_FILES = (
+    'config.json',
+    'preprocessor_config.json',
+    'tokenizer_config.json',
+    'tokenizer.json',
+    'special_tokens_map.json',
+    'spiece.model',
+    'model.safetensors',
+)
+
+
+def _all_present(dir_path: str, files: tuple[str, ...]) -> bool:
+    return all(os.path.exists(os.path.join(dir_path, f)) for f in files)
+
+
+def ensure_wd14() -> None:
+    target = os.path.join(MODEL_DIR, 'wd14')
+    os.makedirs(target, exist_ok=True)
+    try:
+        for fname in WD14_FILES:
+            hf_hub_download(
+                repo_id=WD14_REPO,
+                filename=fname,
+                revision=WD14_REVISION,
+                local_dir=target,
+            )
+        log.info('WD14 ready at %s', target)
+    except (HfHubHTTPError, OSError) as exc:
+        if _all_present(target, WD14_FILES):
+            log.warning('WD14 refresh failed (%s); continuing with local copy', exc)
+        else:
+            log.error('WD14 missing and cannot download: %s', exc)
+            raise
+
+
+def ensure_siglip() -> None:
+    target = os.path.join(MODEL_DIR, 'siglip')
+    os.makedirs(target, exist_ok=True)
+    try:
+        snapshot_download(
+            repo_id=SIGLIP_REPO,
+            revision=SIGLIP_REVISION,
+            local_dir=target,
+            allow_patterns=list(SIGLIP_CORE_FILES),
+        )
+        log.info('SigLIP ready at %s', target)
+    except (HfHubHTTPError, OSError) as exc:
+        if _all_present(target, SIGLIP_CORE_FILES):
+            log.warning('SigLIP refresh failed (%s); continuing with local copy', exc)
+        else:
+            log.error('SigLIP missing and cannot download: %s', exc)
+            raise
+
+
+if __name__ == '__main__':
+    ensure_wd14()
+    ensure_siglip()
+    sys.exit(0)
@@ -0,0 +1,4 @@
+#!/bin/sh
+set -e
+python /app/scripts/ensure_models.py
+exec "$@"