Spaces:
Running on Zero
Running on Zero
File size: 2,914 Bytes
9a25493 | 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 | """Resolve the original released UMT5 assets without loading an archive in RAM."""
import os
from pathlib import Path
import shutil
import zipfile
from filelock import FileLock
from huggingface_hub import hf_hub_download
REPO = 'AlayaLab/FloodDiffusion2'
REVISION = '2402033aff49850b31c1c5be9179c2fd7cb476cb'
PREFIX = 'deps/t5_umt5-xxl-enc-bf16/'
FILES = {
'models_t5_umt5-xxl-enc-bf16.pth': (11361920418, 0x201263e2),
'google/umt5-xxl/special_tokens_map.json': (6623, 0x52b4332b),
'google/umt5-xxl/spiece.model': (4548313, 0x3d61aeda),
'google/umt5-xxl/tokenizer.json': (16837417, 0xdf89c52f),
'google/umt5-xxl/tokenizer_config.json': (61728, 0x882574af),
}
def resolve_t5():
override = os.getenv('FLOOD2_T5_DIR')
if override:
directory = Path(override).resolve()
for name, (size, _) in FILES.items():
if (directory/name).stat().st_size != size:
raise ValueError(f'Unexpected UMT5 asset size: {name}')
else:
directory = Path.home() / '.cache/flood2-live/t5'
directory.mkdir(parents=True, exist_ok=True)
with FileLock(str(directory/'.extract.lock')):
if not all((directory/name).is_file() and (directory/name).stat().st_size == size
for name,(size,_) in FILES.items()):
download_dir = directory/'download'
archive = hf_hub_download(REPO, 'deps.zip', revision=REVISION,
local_dir=download_dir)
with zipfile.ZipFile(archive) as source:
for name, (size, crc) in FILES.items():
target = directory/name
if target.is_file() and target.stat().st_size == size:
continue
info = source.getinfo(PREFIX+name)
if info.file_size != size or info.CRC != crc:
raise ValueError(f'Unexpected UMT5 archive member: {name}')
target.parent.mkdir(parents=True, exist_ok=True)
temporary = target.with_suffix(target.suffix+'.part')
with source.open(info) as src, temporary.open('wb') as dst:
shutil.copyfileobj(src, dst, length=8*1024*1024)
if temporary.stat().st_size != size:
raise ValueError(f'Incomplete UMT5 extraction: {name}')
temporary.replace(target)
# This dedicated temporary archive is no longer needed after
# extraction; preserve room for ZeroGPU's tensor offload file.
archive_path = Path(archive).resolve()
if archive_path.parent == download_dir.resolve():
archive_path.unlink()
return directory/'models_t5_umt5-xxl-enc-bf16.pth', directory/'google/umt5-xxl'
|