File size: 3,312 Bytes
24ba92e | 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 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 | #!/usr/bin/env python
"""Preflight checks for the standalone RFdiffusion package."""
from __future__ import annotations
import argparse
import importlib
import os
import sys
from pathlib import Path
ROOT = Path(__file__).resolve().parents[1]
os.environ.setdefault("DGLBACKEND", "pytorch")
REQUIRED_WEIGHTS = [
"ActiveSite_ckpt.pt",
"Base_ckpt.pt",
"Base_epoch8_ckpt.pt",
"Complex_Fold_base_ckpt.pt",
"Complex_base_ckpt.pt",
"Complex_beta_ckpt.pt",
"InpaintSeq_Fold_ckpt.pt",
"InpaintSeq_ckpt.pt",
"RF_structure_prediction_weights.pt",
]
REQUIRED_FILES = [
"README.md",
"configuration.json",
"config/inference/base.yaml",
"config/inference/symmetry.yaml",
"scripts/run_inference.py",
"examples/input_pdbs/1qys.pdb",
"examples/input_pdbs/1YCR.pdb",
]
TEXT_SUFFIXES = {".py", ".yaml", ".yml", ".json", ".md", ".sh", ".txt", ".tsv"}
GENERATED_DIRS = {"outputs", ".cache", "logs"}
def fail(message: str) -> None:
print(f"[FAIL] {message}")
raise SystemExit(1)
def check_required_files() -> None:
missing = [rel for rel in REQUIRED_FILES if not (ROOT / rel).is_file()]
if missing:
fail("Missing required files: " + ", ".join(missing))
print(f"[OK] Required files present: {len(REQUIRED_FILES)}")
def check_weights(strict: bool) -> None:
missing = []
bad = []
for name in REQUIRED_WEIGHTS:
path = ROOT / "weight" / name
if not path.is_file():
missing.append(str(path.relative_to(ROOT)))
continue
if strict:
size = path.stat().st_size
head = path.read_bytes()[:128]
lfs_marker = b"version https://git-lfs" + b".github.com"
if size < 1024 * 1024 or head.startswith(lfs_marker):
bad.append(f"{path.relative_to(ROOT)} ({size} bytes)")
if missing:
fail("Missing weight files: " + ", ".join(missing))
if bad:
fail("Invalid or placeholder weight files: " + ", ".join(bad))
mode = "strict" if strict else "basic"
print(f"[OK] Weight files present ({mode}): {len(REQUIRED_WEIGHTS)}")
def iter_text_files():
for path in ROOT.rglob("*"):
if not path.is_file():
continue
if GENERATED_DIRS.intersection(path.relative_to(ROOT).parts):
continue
if path.name == ".gitattributes" or path.suffix in TEXT_SUFFIXES:
yield path
def check_imports() -> None:
sys.path.insert(0, str(ROOT))
modules = [
"onescience.utils.rfdiffusion.inference.utils",
"onescience.utils.rfdiffusion.inference.model_runners",
"onescience.models.rfdiffusion.RoseTTAFoldModel",
"onescience.models.se3_transformer",
]
for name in modules:
importlib.import_module(name)
print(f"[OK] Strict imports succeeded: {len(modules)} modules")
def main() -> None:
parser = argparse.ArgumentParser()
parser.add_argument("--strict-weights", action="store_true")
parser.add_argument("--strict-imports", action="store_true")
args = parser.parse_args()
check_required_files()
check_weights(strict=args.strict_weights)
if args.strict_imports:
check_imports()
print("[OK] RFdiffusion standalone preflight passed")
if __name__ == "__main__":
main()
|