From fb925f5e6af3a7bfddd24c046fbfeec7bbcbc4ef Mon Sep 17 00:00:00 2001 From: Bryan Van Deusen Date: Sun, 19 Apr 2026 21:12:43 -0400 Subject: [PATCH] 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 --- .gitignore | 3 +- Dockerfile.ml | 16 ++++--- app/static/js/view-modal.js | 8 ++++ scripts/ensure_models.py | 84 +++++++++++++++++++++++++++++++++++++ scripts/entrypoint-ml.sh | 4 ++ 5 files changed, 105 insertions(+), 10 deletions(-) create mode 100644 scripts/ensure_models.py create mode 100755 scripts/entrypoint-ml.sh diff --git a/.gitignore b/.gitignore index 283b2dd..a59d21d 100644 --- a/.gitignore +++ b/.gitignore @@ -13,4 +13,5 @@ env/ venv/ # Claude files -.claude/ \ No newline at end of file +.claude/ +.remember/ \ No newline at end of file diff --git a/Dockerfile.ml b/Dockerfile.ml index 8abbfd4..c8ecd64 100644 --- a/Dockerfile.ml +++ b/Dockerfile.ml @@ -22,14 +22,9 @@ RUN pip install --no-cache-dir \ --index-url https://download.pytorch.org/whl/cpu \ "torch>=2.5" -# Pre-install only what the model downloader needs, then fetch models. By keeping the model-download -# layer above the full requirements install, a change to requirements-ml.txt won't invalidate the -# (slow, ~2 GB) model download layer. -RUN pip install --no-cache-dir "huggingface_hub>=0.25" -COPY scripts/download_models.py /app/scripts/download_models.py -RUN mkdir -p ${ML_MODEL_DIR} && python /app/scripts/download_models.py - -# Install the rest of the ML requirements. Model files are already on disk from the step above. +# Install ML requirements. Models are NOT baked in — the entrypoint fetches them +# at container start into the ${ML_MODEL_DIR} volume, and huggingface_hub's +# content-based integrity check makes subsequent starts a no-op. COPY requirements-ml.txt /app/requirements-ml.txt RUN pip install --no-cache-dir -r requirements-ml.txt @@ -37,6 +32,9 @@ RUN pip install --no-cache-dir -r requirements-ml.txt # config.py lives at the repo root and is imported by app/__init__.py via 'config.Config'. COPY app /app/app COPY config.py /app/config.py +COPY scripts/ensure_models.py /app/scripts/ensure_models.py +COPY scripts/entrypoint-ml.sh /app/scripts/entrypoint-ml.sh +RUN chmod +x /app/scripts/entrypoint-ml.sh && mkdir -p ${ML_MODEL_DIR} -# Celery needs these env vars set at runtime via docker-compose. +ENTRYPOINT ["/app/scripts/entrypoint-ml.sh"] CMD ["celery", "-A", "app.celery_app:celery", "worker", "--loglevel=info", "-Q", "ml", "--concurrency=1"] diff --git a/app/static/js/view-modal.js b/app/static/js/view-modal.js index ba333ed..2934da6 100644 --- a/app/static/js/view-modal.js +++ b/app/static/js/view-modal.js @@ -811,6 +811,14 @@ document.addEventListener('DOMContentLoaded', () => { modal.classList.add('active'); updateImage(index); history.replaceState({ modalIndex: index }, '', window.location.href); + // Desktop only — skip on touch devices to avoid popping the on-screen keyboard + if (tagInput && !isTouchDevice()) { + setTimeout(() => tagInput.focus(), 100); + } + } + + function isTouchDevice() { + return ('ontouchstart' in window) || (navigator.maxTouchPoints > 0); } function closeModal() { diff --git a/scripts/ensure_models.py b/scripts/ensure_models.py new file mode 100644 index 0000000..b1f2379 --- /dev/null +++ b/scripts/ensure_models.py @@ -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) diff --git a/scripts/entrypoint-ml.sh b/scripts/entrypoint-ml.sh new file mode 100755 index 0000000..353245a --- /dev/null +++ b/scripts/entrypoint-ml.sh @@ -0,0 +1,4 @@ +#!/bin/sh +set -e +python /app/scripts/ensure_models.py +exec "$@"