Spaces:
Running on Zero
Running on Zero
| """Obtain the six external inference assets, without downloading datasets.""" | |
| from pathlib import Path | |
| import os | |
| UPSTREAM_COMMIT = "4faebb83ccda8630e27fbe2189aed1928f4623a0" | |
| SALAD_DRIVE_FOLDER = "1YuDQCgc6RJ4WlR9vt_L34nkXChC7C98_" | |
| ASSETS = ( | |
| "t2m/t2m_denoiser_vpred_vaegelu/opt.txt", | |
| "t2m/t2m_denoiser_vpred_vaegelu/model/net_best_fid.tar", | |
| "t2m/t2m_vae_gelu/opt.txt", | |
| "t2m/t2m_vae_gelu/model/net_best_fid.tar", | |
| "t2m/Comp_v6_KLD005/meta/mean.npy", | |
| "t2m/Comp_v6_KLD005/meta/std.npy", | |
| ) | |
| def ensure_assets(source_root, report=lambda message: None): | |
| source_root = Path(source_root) | |
| checkpoint = source_root / "checkpoints/t2sm/ours/epoch_0100.ckpt" | |
| if not checkpoint.is_file() or checkpoint.stat().st_size == 0: | |
| raise RuntimeError("Style-SALAD checkpoint missing. Use the complete upstream checkout.") | |
| root = source_root / "checkpoints" | |
| missing = [name for name in ASSETS if not (root / name).is_file() | |
| or (root / name).stat().st_size == 0] | |
| if not missing: | |
| return | |
| if os.environ.get("AUTO_DOWNLOAD_ASSETS", "1") != "1": | |
| raise RuntimeError("Missing assets: " + ", ".join(missing)) | |
| # Optional mirror owned by the user, with paths t2m/... at repo root. | |
| # Credentials are read by the Hub client, never from browser fields. | |
| mirror = os.environ.get("ASSET_REPO_ID", "").strip() | |
| if mirror: | |
| from huggingface_hub import hf_hub_download | |
| import shutil | |
| for name in missing: | |
| report("Obtaining model weights and statistics...") | |
| cached = hf_hub_download(repo_id=mirror, filename=name, | |
| token=os.environ.get("HF_TOKEN") or None) | |
| target = root / name | |
| target.parent.mkdir(parents=True, exist_ok=True) | |
| temporary = target.with_name(target.name + ".download") | |
| shutil.copyfile(cached, temporary) | |
| temporary.replace(target) | |
| else: | |
| import gdown | |
| report("Reading the published SALAD asset list...") | |
| entries = gdown.download_folder(id=SALAD_DRIVE_FOLDER, | |
| skip_download=True, quiet=True) | |
| if not entries: | |
| raise RuntimeError("Could not list the published weights on Google Drive.") | |
| for name in missing: | |
| suffix = name.removeprefix("t2m/") | |
| candidates = [e for e in entries | |
| if e.path.replace("\\", "/") == suffix | |
| or e.path.replace("\\", "/").endswith("/" + suffix)] | |
| if len(candidates) != 1: | |
| raise RuntimeError("Asset not uniquely identified on Drive: " + name) | |
| target = root / name | |
| target.parent.mkdir(parents=True, exist_ok=True) | |
| temporary = target.with_name(target.name + ".download") | |
| report("Downloading " + Path(name).name + "...") | |
| result = gdown.download(id=candidates[0].id, output=str(temporary), quiet=False) | |
| if not result or not temporary.is_file() or not temporary.stat().st_size: | |
| raise RuntimeError("Incomplete download: " + name) | |
| temporary.replace(target) | |
| missing = [name for name in ASSETS if not (root / name).is_file() | |
| or not (root / name).stat().st_size] | |
| if missing: | |
| raise RuntimeError("Assets still missing: " + ", ".join(missing)) | |