bvandeusen
fb925f5e6a
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>
85 lines
2.8 KiB
Python
85 lines
2.8 KiB
Python
"""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)
|