Spaces:
Running on Zero
Running on Zero
| """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' | |