bvandeusen
0f35a0c484
Consolidated merge of feat/tag-suggestions branch. Original 64-commit history was lost to git-object corruption in a Nextcloud-synced checkout; this single commit captures the equivalent diff. Includes: - pgvector-backed tag suggestion infra (WD14 + SigLIP centroids, ml-worker container, Celery tasks, suggestion service, accept/reject endpoints + modal UI with green/red chip buttons) - Character/fandom integrity: title-case normalization on every write path, fandom-id backfill, maintenance task + settings button, migrations g26041901 + h26041901 to canonicalize legacy rows with case-only duplicate merging - Tag-underscores + modal polish: WD14 name canonicalization at emit + accept + add/bulk-add paths, migration i26041901 for legacy-row rename-or-merge across character/fandom/NULL kinds, suggestion-accept refresh parity via awaited loadTags, persistent chip tint
46 lines
1.4 KiB
Python
46 lines
1.4 KiB
Python
"""Pre-fetch WD14 and SigLIP model files into $ML_MODEL_DIR at image build time."""
|
|
from __future__ import annotations
|
|
import os
|
|
import sys
|
|
|
|
from huggingface_hub import hf_hub_download, snapshot_download
|
|
|
|
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')
|
|
|
|
|
|
def download_wd14() -> None:
|
|
target = os.path.join(MODEL_DIR, 'wd14')
|
|
os.makedirs(target, exist_ok=True)
|
|
hf_hub_download(repo_id=WD14_REPO, filename='model.onnx', local_dir=target, local_dir_use_symlinks=False)
|
|
hf_hub_download(repo_id=WD14_REPO, filename='selected_tags.csv', local_dir=target, local_dir_use_symlinks=False)
|
|
print(f"WD14 downloaded to {target}")
|
|
|
|
|
|
def download_siglip() -> None:
|
|
target = os.path.join(MODEL_DIR, 'siglip')
|
|
os.makedirs(target, exist_ok=True)
|
|
snapshot_download(
|
|
repo_id=SIGLIP_REPO,
|
|
local_dir=target,
|
|
local_dir_use_symlinks=False,
|
|
allow_patterns=[
|
|
'config.json',
|
|
'preprocessor_config.json',
|
|
'tokenizer_config.json',
|
|
'tokenizer.json',
|
|
'special_tokens_map.json',
|
|
'spiece.model',
|
|
'model.safetensors',
|
|
'pytorch_model.bin',
|
|
],
|
|
)
|
|
print(f"SigLIP downloaded to {target}")
|
|
|
|
|
|
if __name__ == '__main__':
|
|
download_wd14()
|
|
download_siglip()
|
|
sys.exit(0)
|