Spaces:
Running on CPU Upgrade
Running on CPU Upgrade
File size: 3,137 Bytes
cd24b4a | 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 | """Assemble the central checkpoint directory used by the web interface.
- Symlinks the already-available OFFICIAL checkpoints (FatFormer, AIDE, NPR) to the canonical
filenames the registry expects.
- Writes a MANIFEST.{json,md} summarizing, per model, whether a checkpoint is present and how
it was obtained (official weights / trained-from-scratch / training-free).
Run: conda run -n SynthScope python interface/scripts/collect_checkpoints.py
Override the destination with SYNTHSCOPE_CKPT_DIR.
"""
import json
import os
import sys
import time
from pathlib import Path
sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
from backend import config # noqa: E402 (sets paths)
from backend.registry import list_specs # noqa: E402
# Known source locations of the official weights on this node.
REPRO_CKPTS = Path("/hddstore/ssalehi/SynthScope/repro/ckpts")
OFFICIAL_SOURCES = {
"fat_former.pth": REPRO_CKPTS / "fatformer_official.pth",
"aide.pth": REPRO_CKPTS / "aide_progan_train.pth",
"npr.pth": REPRO_CKPTS / "NPR_model_epoch_last_3090.pth",
}
def link(dest: Path, src: Path) -> str:
if not src.exists():
return f"source missing: {src}"
if dest.exists() or dest.is_symlink():
dest.unlink()
dest.symlink_to(src)
return f"linked -> {src}"
def main():
dst = config.CHECKPOINTS_DIR
dst.mkdir(parents=True, exist_ok=True)
print(f"checkpoints dir: {dst}")
for name, src in OFFICIAL_SOURCES.items():
print(f" {name:18s} {link(dst / name, src)}")
manifest = {"checkpoints_dir": str(dst), "generated": time.strftime("%Y-%m-%d %H:%M:%S"),
"models": []}
for spec in list_specs():
if spec.ckpt_kind == "none":
status, present = "training-free", True
else:
p = dst / spec.checkpoint
present = p.exists()
if spec.ckpt_kind.startswith("official"):
status = "official weights" if present else "MISSING (official)"
else:
status = "trained from scratch" if present else "MISSING (needs training)"
manifest["models"].append({
"key": spec.key, "name": spec.name, "modality": spec.modality,
"checkpoint": spec.checkpoint, "present": present, "status": status,
})
(dst / "MANIFEST.json").write_text(json.dumps(manifest, indent=2))
lines = ["# SynthScope checkpoints manifest", "",
f"Directory: `{dst}` ", f"Generated: {manifest['generated']}", "",
"| Model | Modality | Checkpoint | Present | Status |",
"| --- | --- | --- | --- | --- |"]
for m in manifest["models"]:
lines.append(f"| {m['key']} | {m['modality']} | {m['checkpoint'] or '—'} | "
f"{'✅' if m['present'] else '❌'} | {m['status']} |")
(dst / "MANIFEST.md").write_text("\n".join(lines) + "\n")
print(f"\nwrote {dst/'MANIFEST.json'} and {dst/'MANIFEST.md'}")
present = sum(m["present"] for m in manifest["models"])
print(f"present: {present}/{len(manifest['models'])} models")
if __name__ == "__main__":
main()
|