File size: 2,353 Bytes
20df93e | 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 | """Restore the two upstream HF cache entries used by the censor workers."""
import argparse
import hashlib
import json
from pathlib import Path
import shutil
def main():
from huggingface_hub.constants import HF_HUB_CACHE
parser = argparse.ArgumentParser(description=__doc__)
parser.add_argument("--repo-dir", type=Path, default=Path(__file__).resolve().parent)
parser.add_argument("--cache-dir", type=Path, default=Path(HF_HUB_CACHE))
args = parser.parse_args()
root = args.repo_dir.resolve()
manifest = json.loads((root / "model-manifest.json").read_text(encoding="utf-8"))
models = {f["path"]: f for f in manifest["files"]}
for entry in manifest["censor_hf_cache"]:
source = root / entry["path"]
expected = models[entry["path"]]
if source.stat().st_size != expected["size_bytes"]:
raise RuntimeError("File size mismatch: " + str(source))
digest = hashlib.sha256(source.read_bytes()).hexdigest()
if digest != expected["sha256"]:
raise RuntimeError("SHA256 mismatch: " + str(source))
cache = args.cache_dir / ("models--" + entry["repo_id"].replace("/", "--"))
blob = cache / "blobs" / digest
blob.parent.mkdir(parents=True, exist_ok=True)
if blob.exists():
if hashlib.sha256(blob.read_bytes()).hexdigest() != digest:
raise RuntimeError("Existing cache blob is corrupt: " + str(blob))
else:
shutil.copyfile(source, blob)
snapshot = cache / "snapshots" / entry["revision"] / entry["filename"]
snapshot.parent.mkdir(parents=True, exist_ok=True)
if snapshot.exists():
if hashlib.sha256(snapshot.read_bytes()).hexdigest() != digest:
raise RuntimeError("Existing cache snapshot differs: " + str(snapshot))
else:
snapshot.symlink_to(blob.resolve())
ref = cache / "refs" / "main"
ref.parent.mkdir(parents=True, exist_ok=True)
ref.write_text(entry["revision"], encoding="utf-8")
print("Restored " + entry["repo_id"] + " @ " + entry["revision"])
print("Set CENSOR_MODELS_DIR=" + str(root / "auxiliary" / "censor") + " before starting ComfyUI.")
print("The cache directory must be the one used by the ComfyUI service account.")
if __name__ == "__main__":
main()
|