| """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: |
| |
| 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() |
|
|
