Spaces:
Running on Zero
Running on Zero
File size: 5,420 Bytes
7fabddf | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 | """Portable, checksum-verified Hub assets and scene archives."""
import hashlib
import json
import os
from pathlib import Path, PurePosixPath
import shutil
import tarfile
import tempfile
def sha256(path):
digest = hashlib.sha256()
with Path(path).open('rb') as handle:
for block in iter(lambda: handle.read(8 * 1024 * 1024), b''):
digest.update(block)
return digest.hexdigest()
def relative_path(value):
path = PurePosixPath(value)
if not value or path.is_absolute() or '..' in path.parts or '\\' in value:
raise ValueError(f'Unsafe release path: {value!r}')
return path
def describe(path, root):
path = Path(path)
if path.is_symlink() or not path.is_file():
raise ValueError(f'Expected a regular file: {path}')
return {'path': path.relative_to(root).as_posix(),
'bytes': path.stat().st_size, 'sha256': sha256(path)}
def fetch(repo_id, entry, destination, revision, install=True):
from huggingface_hub import hf_hub_download
relative = relative_path(entry['path'])
destination = Path(destination)
destination.mkdir(parents=True, exist_ok=True)
# HF maintains resumable downloads in its cache. Verify before copying into data.
source = Path(hf_hub_download(repo_id, str(relative), repo_type='dataset',
revision=revision, cache_dir=str(destination / '.hf-cache')))
if source.stat().st_size != entry['bytes'] or sha256(source) != entry['sha256']:
raise ValueError(f'Checksum mismatch: {relative}')
if not install:
return source
target = destination / relative
if not target.resolve().is_relative_to(destination.resolve()):
raise ValueError(f'Destination escapes output directory: {target}')
target.parent.mkdir(parents=True, exist_ok=True)
if target.is_symlink():
raise ValueError(f'Refusing to overwrite symlink: {target}')
if target.is_file() and target.stat().st_size == entry['bytes'] and sha256(target) == entry['sha256']:
return target
temporary = target.with_name(target.name + '.partial')
shutil.copyfile(source, temporary)
os.replace(temporary, target)
return target
def load_manifest(repo_id, output, revision='main'):
from huggingface_hub import HfApi, hf_hub_download
# Pin all files to one immutable revision, including the manifest.
revision = HfApi().dataset_info(repo_id, revision=revision).sha
output = Path(output)
output.mkdir(parents=True, exist_ok=True)
path = hf_hub_download(repo_id, 'manifest.json', repo_type='dataset',
revision=revision, cache_dir=str(output / '.hf-cache'))
manifest = json.loads(Path(path).read_text())
if manifest.get('format_version') != 1:
raise ValueError('Unsupported Urbex release manifest')
return manifest, revision
def unpack(archive, destination, prefix, expected):
"""Validate into a temporary tree; replace only the selected scene/package."""
import zstandard
destination = Path(destination)
prefix = relative_path(prefix)
target = destination / prefix
if not target.resolve().is_relative_to(destination.resolve()):
raise ValueError(f'Destination escapes output directory: {target}')
marker = target / '.urbex-download.json'
if marker.is_file():
saved = json.loads(marker.read_text())
files = [p for p in target.rglob('*') if p.is_file() and p != marker]
if (saved.get('sha256') == expected['sha256'] and
len(files) == expected['file_count'] and
sum(p.stat().st_size for p in files) == expected['unpacked_bytes']):
return target
with tempfile.TemporaryDirectory(prefix='.urbex-extract-', dir=destination) as temporary:
stage = Path(temporary)
names = set()
size = 0
with Path(archive).open('rb') as source:
with zstandard.ZstdDecompressor().stream_reader(source) as reader:
with tarfile.open(fileobj=reader, mode='r|') as handle:
for member in handle:
path = relative_path(member.name)
if not path.is_relative_to(prefix) or not member.isfile() or str(path) in names:
raise ValueError(f'Invalid archive member: {member.name}')
names.add(str(path))
size += member.size
if len(names) > expected['file_count'] or size > expected['unpacked_bytes']:
raise ValueError('Archive exceeds declared contents')
handle.extract(member, path=stage, filter='data')
if len(names) != expected['file_count'] or size != expected['unpacked_bytes']:
raise ValueError('Archive contents do not match the manifest')
new = stage / prefix
(new / marker.name).write_text(json.dumps({'sha256': expected['sha256']}) + '\n')
target.parent.mkdir(parents=True, exist_ok=True)
if target.is_symlink():
raise ValueError(f'Refusing to replace symlink: {target}')
backup = stage / 'previous'
if target.exists():
target.rename(backup)
try:
new.rename(target)
except BaseException:
if backup.exists():
backup.rename(target)
raise
return target
|