File size: 5,126 Bytes
5ab5b72 | 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 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 | """Download and run the non-antibody catalog using pinned model, prep and Diffusers revisions."""
import argparse
import hashlib
import importlib.metadata
import json
from pathlib import Path
def pipeline_options(entry, scheduler):
options = dict(entry["pipeline_kwargs"])
fraction = entry.get("remaining_schedule_fraction")
if fraction is not None:
steps = scheduler.num_inference_steps
scheduler.set_timesteps(steps)
index = round((1 - fraction) * steps)
options["partial_t"] = {
name: float(scheduler.modality_timestep(name, index))
for name in scheduler.modalities
}
return options
def main():
parser = argparse.ArgumentParser(description=__doc__)
parser.add_argument("--repo", default="nvidia/RFD4-Proteina-prep")
parser.add_argument("--revision", default="main")
parser.add_argument("--tasks", nargs="+")
parser.add_argument("--seed", default=42, type=int)
parser.add_argument("--device", default="cuda:0")
parser.add_argument("--out", type=Path, default=Path("rfd4-designs"))
parser.add_argument(
"--prepare-only",
action="store_true",
help="Validate inputs on CPU without model weights.",
)
args = parser.parse_args()
import torch
from diffusers import ModularPipeline, RFD4ProteinaPipeline
from huggingface_hub import HfApi, hf_hub_download
revision = HfApi().model_info(args.repo, revision=args.revision).sha
manifest_path = hf_hub_download(
args.repo, "examples/manifest.json", revision=revision
)
manifest = json.loads(Path(manifest_path).read_text())
catalog = {entry["id"]: entry for entry in manifest["examples"]}
tasks = args.tasks if args.tasks is not None else list(catalog)
unknown = set(tasks) - catalog.keys()
if unknown:
parser.error(
f"Unknown or excluded tasks: {sorted(unknown)}. Choose from {list(catalog)}."
)
if len(tasks) != len(set(tasks)):
parser.error(
"Duplicate tasks would overwrite the same output; select each task once."
)
args.out.mkdir(parents=True, exist_ok=True)
report_path = args.out / "report.json"
destinations = [report_path] + (
[] if args.prepare_only else [args.out / f"{name}.cif" for name in tasks]
)
for path in destinations:
if path.exists() or path.is_symlink():
raise FileExistsError(
f"Refusing to overwrite {path}; choose a new --out directory."
)
helper = ModularPipeline.from_pretrained(
args.repo, revision=revision, trust_remote_code=True
)
device = "cpu" if args.prepare_only else args.device
pipe = None
if not args.prepare_only:
# Check the source revision before downloading weights; recent PR commits rename their keys.
dist = importlib.metadata.distribution("diffusers")
direct_url = json.loads(dist.read_text("direct_url.json") or "{}")
commit = direct_url.get("vcs_info", {}).get("commit_id")
if commit != manifest["diffusers_revision"]:
raise RuntimeError(
"This model snapshot requires a pinned Diffusers install. Run: pip install "
f"git+{manifest['diffusers_repository']}@{manifest['diffusers_revision']}"
)
dtype = torch.bfloat16 if torch.device(device).type == "cuda" else torch.float32
pipe = RFD4ProteinaPipeline.from_pretrained(
manifest["model_repo"], revision=manifest["model_revision"], dtype=dtype
).to(device)
report = {
"helper_revision": revision,
"model_revision": manifest["model_revision"],
"seed": args.seed,
"tasks": {},
}
for name in tasks:
entry = catalog[name]
input_path = hf_hub_download(args.repo, entry["file"], revision=revision)
if hashlib.sha256(Path(input_path).read_bytes()).hexdigest() != entry["sha256"]:
raise ValueError(f"Input digest mismatch: {name}")
prepared = helper(cif_path=input_path, seed=args.seed, device=device)
features, atoms = prepared.get("features"), prepared.get("atom_array")
result = {
"input_sha256": entry["sha256"],
"tokens": int(features["padding_mask"].shape[1]),
}
if pipe is not None:
options = pipeline_options(entry, pipe.scheduler)
output = pipe(
features,
generator=torch.Generator(device).manual_seed(args.seed),
**options,
)
path = helper(
pipeline_output=output,
features=features,
atom_array=atoms,
output_path=str(args.out / f"{name}.cif"),
output="cif_path",
)
result.update(output=path, options=options)
report["tasks"][name] = result
print(name, result, flush=True)
with report_path.open("x") as handle:
json.dump(report, handle, indent=2)
handle.write("\n")
if __name__ == "__main__":
main()
|