43 lines
2.1 KiB
Python
43 lines
2.1 KiB
Python
"""Prepare project-local Hugging Face cache metadata and resumable downloads."""
|
|
import json
|
|
import subprocess
|
|
from pathlib import Path
|
|
ROOT = Path(__file__).resolve().parents[1]
|
|
CACHE = ROOT / '.hf_cache' / 'hub'
|
|
|
|
def fetch(url, target):
|
|
target.parent.mkdir(parents=True, exist_ok=True)
|
|
subprocess.run(['curl', '-fsSL', '--retry', '3', '--retry-all-errors', '--max-time', '45', url, '-o', str(target)], check=True)
|
|
|
|
manifest = []
|
|
for repo, files, aliases in [
|
|
('google-bert/bert-base-uncased', ['config.json', 'tokenizer.json', 'tokenizer_config.json', 'vocab.txt', 'model.safetensors'], ['bert-base-uncased']),
|
|
('lpiccinelli/unidepth-v2-vitl14', ['config.json', 'model.safetensors'], []),
|
|
('Rain729/Prior-Depth-Anything', ['depth_anything_v2_vitb.pth', 'prior_depth_anything_vitb.pth'], []),
|
|
]:
|
|
meta = ROOT / 'output' / (repo.replace('/', '_') + '_metadata.json')
|
|
fetch('https://hf-mirror.com/api/models/' + repo, meta)
|
|
revision = json.loads(meta.read_text())['sha']
|
|
cache = CACHE / ('models--' + repo.replace('/', '--'))
|
|
snapshot = cache / 'snapshots' / revision
|
|
snapshot.mkdir(parents=True, exist_ok=True)
|
|
(cache / 'refs').mkdir(exist_ok=True)
|
|
(cache / 'refs' / 'main').write_text(revision)
|
|
for name in files:
|
|
target = snapshot / name
|
|
url = 'https://hf-mirror.com/' + repo + '/resolve/' + revision + '/' + name
|
|
if repo.startswith('lpiccinelli/'):
|
|
local = ROOT / 'weights' / 'unidepth-v2-vitl14' / name
|
|
if not target.exists() and not target.is_symlink():
|
|
target.symlink_to(local)
|
|
elif name.endswith(('.json', '.txt')):
|
|
fetch(url, target)
|
|
else:
|
|
manifest.extend([url, ' dir=' + str(snapshot), ' out=' + name])
|
|
for alias in aliases:
|
|
alias_path = CACHE / ('models--' + alias.replace('/', '--'))
|
|
if not alias_path.exists():
|
|
alias_path.symlink_to(cache, target_is_directory=True)
|
|
(ROOT / 'output' / 'vipe_weights_downloads.txt').write_text('\n'.join(manifest) + '\n')
|
|
print('Prepared', len(manifest) // 3, 'weight downloads')
|