File size: 6,217 Bytes
75b53c6 ceb0adf 75b53c6 0a34f04 75b53c6 0a34f04 75b53c6 0a34f04 75b53c6 0a34f04 75b53c6 ceb0adf 75b53c6 0a34f04 75b53c6 0a34f04 75b53c6 0a34f04 75b53c6 | 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 | """Run `RFD4ProteinaPipeline` on the AtomWorks Studio example inputs and write annotated CIFs plus a JSON report.
```bash
python run_examples.py --checkpoint nvidia/RFD4-Proteina-Diffusers --inputs-dir inputs --out outputs
```
"""
import argparse
import json
import time
from pathlib import Path
import torch
from diffusers import RFD4ProteinaPipeline
from rfd4_proteina_prep import prepare_cif, write_cif
# Example id -> (input file, pipeline arguments), for the inputs this repository ships under
# `examples/inputs/`. `examples/manifest.json` is the catalog of record; `check_against_manifest()`
# below fails the run if the two disagree, because they have drifted apart before -- this table
# still listed two antibody tasks whose inputs the repository does not ship.
TASKS = {
"protein_binder": ("protein_binder.cif", {}),
"enzyme_design": ("enzyme_design.cif", {}),
"motif_scaffolding": ("motif_scaffolding.cif", {}),
"small_molecule_binder": ("small_molecule_binder.cif", {}),
"unconditional": ("unconditional.cif", {}),
"inverse_folding": (
"inverse_folding.cif",
{"inverse_folding": True, "partial_t_bb_ca": 0.0, "partial_t_local_latents": 0.0},
),
"partial_diffusion": ("partial_diffusion.cif", {"partial_fraction": 0.4}),
}
# The tensors the pipeline takes by name; everything else prepare_cif returns is conditioning.
# `coords_nm` belongs here: partial diffusion and inverse folding read it as a named argument, and
# a full run ignores it, so passing it always is correct for every task.
REQUIRED = [
"padding_mask",
"is_in_context_conditioning",
"supports_local_latent",
"token_type",
"atom37_mask",
"coords_nm",
]
def check_against_manifest():
"""Fail loudly if TASKS and the shipped manifest disagree about which inputs exist.
Ids and files only, deliberately. The manifest's `pipeline_kwargs` are written for the pinned
Diffusers revision that `run_hub_examples.py` enforces, which still takes one `partial_t` dict;
this script runs against whatever Diffusers is installed, where the start times are separate
arguments. Comparing the options would fail on that difference alone.
"""
manifest_path = Path(__file__).resolve().parent / "manifest.json"
if not manifest_path.is_file():
return
catalog = {entry["id"]: entry for entry in json.loads(manifest_path.read_text())["examples"]}
for name, (file, options) in TASKS.items():
entry = catalog.get(name)
if entry is None:
raise SystemExit(f"{name} is not in {manifest_path}; it may have been excluded upstream.")
if Path(entry["file"]).name != file:
raise SystemExit(f"{name} reads {file}, but the manifest names {entry['file']}.")
fraction = entry.get("remaining_schedule_fraction")
if (fraction is not None) != ("partial_fraction" in options):
raise SystemExit(f"{name} and the manifest disagree about whether it starts part-way through.")
if fraction is not None and fraction != options["partial_fraction"]:
raise SystemExit(f"{name} starts at {options['partial_fraction']}, the manifest says {fraction}.")
def main():
parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
parser.add_argument("--checkpoint", default="nvidia/RFD4-Proteina-Diffusers")
parser.add_argument("--inputs-dir", type=Path, required=True)
parser.add_argument("--out", type=Path, required=True)
parser.add_argument("--tasks", nargs="*", default=list(TASKS))
parser.add_argument("--seed", type=int, default=42)
parser.add_argument("--dtype", choices=["bfloat16", "float32"], default="bfloat16")
args = parser.parse_args()
check_against_manifest()
unknown = set(args.tasks) - TASKS.keys()
if unknown:
parser.error(f"Unknown tasks: {sorted(unknown)}. Choose from {list(TASKS)}.")
args.out.mkdir(parents=True, exist_ok=True)
device = torch.device("cuda:0")
pipe = RFD4ProteinaPipeline.from_pretrained(args.checkpoint, dtype=getattr(torch, args.dtype)).to(device)
steps = pipe.scheduler.num_inference_steps
report = {}
for name in args.tasks:
file, options = TASKS[name]
options = dict(options)
fraction = options.pop("partial_fraction", None)
if fraction is not None:
# Start from the noised input at the point of the schedule where `fraction` of the trajectory remains.
pipe.scheduler.set_timesteps(steps)
index = round((1 - fraction) * steps)
# One start time per modality, each its own argument: `partial_t_bb_ca`, `partial_t_local_latents`.
for modality in pipe.scheduler.modalities:
options[f"partial_t_{modality}"] = float(pipe.scheduler.modality_timestep(modality, index))
features, atoms = prepare_cif(args.inputs_dir / file, seed=args.seed, device=device)
generator = torch.Generator(device).manual_seed(args.seed)
torch.cuda.synchronize()
start = time.time()
# The pipeline takes the masks and molecule tensors by name; the annotations go in one
# mapping. Passing `features` positionally binds the whole dict to `padding_mask`.
inputs = {name: features[name] for name in REQUIRED}
conditioning = {name: value for name, value in features.items() if name not in REQUIRED}
output = pipe(**inputs, conditioning_features=conditioning, generator=generator, **options)
torch.cuda.synchronize()
seconds = time.time() - start
destination = args.out / f"{name}.cif"
destination.unlink(missing_ok=True)
write_cif(output, features, atoms, destination)
report[name] = {
"input": file,
"tokens": int(features["padding_mask"].shape[1]),
"designed_tokens": int(output.emit_mask.sum()),
"seconds": round(seconds, 1),
"options": {k: v for k, v in options.items()},
}
print(name, report[name], flush=True)
(args.out / "report.json").write_text(json.dumps(report, indent=2) + "\n")
if __name__ == "__main__":
main()
|