# # For licensing see accompanying LICENSE file. # Copyright (c) 2025 Apple Inc. Licensed under MIT License. # import os import sys import torch import hydra import omegaconf import argparse import numpy as np from copy import deepcopy from pathlib import Path from itertools import starmap import lightning.pytorch as pl ROOT = Path(__file__).resolve().parents[1] CONFIG_DIR = ROOT / "config" WEIGHT_DIR = ROOT / "weight" if str(ROOT) not in sys.path: sys.path.insert(0, str(ROOT)) from models.simplefold.flow import LinearPath from models.simplefold.torch.sampler import EMSampler from onescience.datapipes.simplefold.processor.protein_processor import ProteinDataProcessor from onescience.utils.simplefold.datamodule_utils import process_one_inference_structure from onescience.utils.simplefold.esm_utils import _af2_to_esm, esm_registry import models.esm.pretrained as esm_pretrained from onescience.utils.simplefold.boltz_utils import process_structure, save_structure from onescience.utils.simplefold.fasta_utils import process_fastas, check_fasta_inputs from onescience.datapipes.boltz_data_pipeline.feature.featurizer import BoltzFeaturizer from onescience.datapipes.boltz_data_pipeline.tokenize.boltz_protein import BoltzTokenizer try: import mlx.core as mx from mlx.utils import tree_unflatten, tree_flatten from models.simplefold.mlx.sampler import EMSampler as EMSamplerMLX from models.simplefold.mlx.esm_network import ESM2 as ESM2MLX from onescience.utils.simplefold.mlx_utils import map_torch_to_mlx, map_plddt_torch_to_mlx MLX_AVAILABLE = True except: MLX_AVAILABLE = False print("MLX not installed, skip importing MLX related packages.") SUPPORTED_MODELS = { "simplefold_100M", "simplefold_360M", "simplefold_700M", "simplefold_1.1B", "simplefold_1.6B", "simplefold_3B", } ESM2_MODEL_NAME = "esm2_t36_3B_UR50D" MIN_REAL_WEIGHT_BYTES = 1024 def _resolve_path(path: str | os.PathLike[str], base: Path = ROOT) -> Path: path = Path(path) return path if path.is_absolute() else base / path def _require_local_file(path: Path, description: str, min_bytes: int = MIN_REAL_WEIGHT_BYTES) -> Path: if not path.exists(): raise FileNotFoundError( f"Missing {description}: {path}\n" f"Please place the file under {WEIGHT_DIR} or pass the corresponding environment/CLI path." ) if path.stat().st_size < min_bytes: raise RuntimeError( f"{description} looks like a link/placeholder rather than a real weight file: {path}\n" "Replace it with the real file before running inference." ) return path def load_local_or_hub_esm2_3b(): candidate_paths = [] esm_model_path = os.getenv("SIMPLEFOLD_ESM2_MODEL_PATH") if esm_model_path: candidate_paths.append(esm_model_path) candidate_paths.extend( [ WEIGHT_DIR / "esm_models" / f"{ESM2_MODEL_NAME}.pt", WEIGHT_DIR / f"{ESM2_MODEL_NAME}.pt", ] ) seen_paths = set() for candidate_path in candidate_paths: if candidate_path in seen_paths: continue seen_paths.add(candidate_path) candidate_path = Path(candidate_path) regression_path = candidate_path.with_name( f"{candidate_path.stem}-contact-regression.pt" ) if candidate_path.exists() and regression_path.exists(): _require_local_file(candidate_path, "ESM-2 3B model weight") _require_local_file(regression_path, "ESM-2 contact regression weight") print(f"Loading ESM-2 3B weights from local path: {candidate_path}") return esm_pretrained.load_model_and_alphabet(str(candidate_path)) raise FileNotFoundError( "Local ESM-2 3B weights were not found. Put both files here:\n" f" {WEIGHT_DIR / 'esm_models' / f'{ESM2_MODEL_NAME}.pt'}\n" f" {WEIGHT_DIR / 'esm_models' / f'{ESM2_MODEL_NAME}-contact-regression.pt'}\n" "or set SIMPLEFOLD_ESM2_MODEL_PATH to the .pt file. Remote download is disabled in this package." ) def resolve_ccd_path(cache: Path) -> Path: candidate_paths = [] ccd_path = os.getenv("SIMPLEFOLD_CCD_PATH") if ccd_path: candidate_paths.append(Path(ccd_path)) candidate_paths.append(WEIGHT_DIR / "ccd.pkl") seen_paths = set() for candidate_path in candidate_paths: resolved_path = str(candidate_path) if resolved_path in seen_paths: continue seen_paths.add(resolved_path) if candidate_path.exists(): _require_local_file(candidate_path, "CCD dictionary") print(f"Using local CCD dictionary: {candidate_path}") return candidate_path raise FileNotFoundError( f"Missing CCD dictionary. Put ccd.pkl under {WEIGHT_DIR} or set SIMPLEFOLD_CCD_PATH." ) def get_config_path(relative_path): """Resolve a SimpleFold config from the bundled config directory.""" config_subpath = relative_path.replace("configs/", "") config_path = CONFIG_DIR / config_subpath if not config_path.is_file(): raise FileNotFoundError(f"Could not find config file: {config_path}") return str(config_path) def initialize_folding_model(args): # define folding model simplefold_model = args.simplefold_model if simplefold_model not in SUPPORTED_MODELS: raise ValueError(f"Unsupported model {simplefold_model!r}. Choose one of {sorted(SUPPORTED_MODELS)}") # create checkpoint directory ckpt_dir = _resolve_path(args.ckpt_dir) ckpt_path = os.path.join(ckpt_dir, f"{simplefold_model}.ckpt") # create folding model ckpt_path = os.path.join(ckpt_dir, f"{simplefold_model}.ckpt") _require_local_file(Path(ckpt_path), f"{simplefold_model} checkpoint") cfg_path = get_config_path(f"configs/model/architecture/foldingdit_{simplefold_model[11:]}.yaml") checkpoint = torch.load(ckpt_path, map_location="cpu", weights_only=False) # load model checkpoint if args.backend == 'torch': device = torch.device("cuda" if torch.cuda.is_available() else "cpu") model_config = omegaconf.OmegaConf.load(cfg_path) model = hydra.utils.instantiate(model_config) model.load_state_dict(checkpoint, strict=True) model = model.to(device) elif args.backend == 'mlx': device = "cpu" # replace torch implementations with mlx with open(cfg_path, "r") as f: yaml_str = f.read() yaml_str = yaml_str.replace('torch', 'mlx') model_config = omegaconf.OmegaConf.create(yaml_str) model = hydra.utils.instantiate(model_config) mlx_state_dict = {k: mx.array(v) for k, v in starmap(map_torch_to_mlx, checkpoint.items()) if k is not None} model.update(tree_unflatten(list(mlx_state_dict.items()))) print(f"Folding model {simplefold_model} loaded.") print(f"Using device: {device}.") model.eval() return model, device def initialize_plddt_module(args, device): if not args.plddt: return None, None # load pLDDT module if specified ckpt_dir = _resolve_path(args.ckpt_dir) plddt_ckpt_path = ckpt_dir / "plddt.ckpt" if not plddt_ckpt_path.exists(): plddt_ckpt_path = ckpt_dir / "plddt_module_1.6B.ckpt" _require_local_file(plddt_ckpt_path, "pLDDT checkpoint") plddt_module_path = get_config_path("configs/model/architecture/plddt_module.yaml") plddt_checkpoint = torch.load(plddt_ckpt_path, map_location="cpu", weights_only=False) if args.backend == "torch": plddt_config = omegaconf.OmegaConf.load(plddt_module_path) plddt_out_module = hydra.utils.instantiate(plddt_config) plddt_out_module.load_state_dict(plddt_checkpoint, strict=True) plddt_out_module = plddt_out_module.to(device) elif args.backend == "mlx": # replace torch implementations with mlx with open(plddt_module_path, "r") as f: yaml_str = f.read() yaml_str = yaml_str.replace('torch', 'mlx') plddt_config = omegaconf.OmegaConf.create(yaml_str) plddt_out_module = hydra.utils.instantiate(plddt_config) mlx_state_dict = {k: mx.array(v) for k, v in starmap(map_plddt_torch_to_mlx, plddt_checkpoint.items()) if k is not None} plddt_out_module.update(tree_unflatten(list(mlx_state_dict.items()))) plddt_out_module.eval() print(f"pLDDT output module loaded with {args.backend} backend.") plddt_latent_ckpt_path = ckpt_dir / "simplefold_1.6B.ckpt" _require_local_file(plddt_latent_ckpt_path, "SimpleFold 1.6B pLDDT latent checkpoint") plddt_latent_config_path = get_config_path("configs/model/architecture/foldingdit_1.6B.yaml") plddt_latent_checkpoint = torch.load(plddt_latent_ckpt_path, map_location="cpu", weights_only=False) if args.backend == "torch": plddt_latent_config = omegaconf.OmegaConf.load(plddt_latent_config_path) plddt_latent_module = hydra.utils.instantiate(plddt_latent_config) plddt_latent_module.load_state_dict(plddt_latent_checkpoint, strict=True) plddt_latent_module = plddt_latent_module.to(device) elif args.backend == "mlx": # replace torch implementations with mlx with open(plddt_latent_config_path, "r") as f: yaml_str = f.read() yaml_str = yaml_str.replace('torch', 'mlx') plddt_latent_config = omegaconf.OmegaConf.create(yaml_str) plddt_latent_module = hydra.utils.instantiate(plddt_latent_config) mlx_state_dict = {k: mx.array(v) for k, v in starmap(map_torch_to_mlx, plddt_latent_checkpoint.items()) if k is not None} plddt_latent_module.update(tree_unflatten(list(mlx_state_dict.items()))) plddt_latent_module.eval() print(f"pLDDT latent module loaded with {args.backend} backend.") return plddt_latent_module, plddt_out_module def initialize_esm_model(args, device): # load ESM2 model esm_model, esm_dict = load_local_or_hub_esm2_3b() af2_to_esm = _af2_to_esm(esm_dict) if args.backend == 'torch': esm_model = esm_model.to(device) af2_to_esm = af2_to_esm.to(device) elif args.backend == 'mlx': esm_model_mlx = ESM2MLX(num_layers=36, embed_dim=2560, attention_heads=40) esm_state_dict_torch = esm_model.cpu().state_dict() esm_state_dict_torch = {k: mx.array(v) for k, v in starmap(map_torch_to_mlx, esm_state_dict_torch.items()) if k is not None} esm_model_mlx.update(tree_unflatten(list(esm_state_dict_torch.items()))) esm_model = esm_model_mlx print(f"pLM ESM-3B loaded with {args.backend} backend.") esm_model.eval() return esm_model, esm_dict, af2_to_esm def initialize_others(args, device): # prepare data tokenizer, featurizer, and processor tokenizer = BoltzTokenizer() featurizer = BoltzFeaturizer() processor = ProteinDataProcessor( device=device, scale=16.0, ref_scale=5.0, multiplicity=1, inference_multiplicity=args.nsample_per_protein, backend=args.backend, ) # define flow process and sampler flow = LinearPath() if args.backend == "torch": sampler_cls = EMSampler elif args.backend == "mlx": sampler_cls = EMSamplerMLX sampler = sampler_cls( num_timesteps=args.num_steps, t_start=1e-4, tau=args.tau, log_timesteps=True, w_cutoff=0.99, ) return tokenizer, featurizer, processor, flow, sampler def generate_structure( args, batch, sampler, flow, processor, model, plddt_latent_module, plddt_out_module, device ): # run inference for target protein if args.backend == "torch": noise = torch.randn_like(batch['coords']).to(device) elif args.backend == "mlx": noise = mx.random.normal(batch['coords'].shape) out_dict = sampler.sample(model, flow, noise, batch) if args.plddt: if args.backend == "torch": t = torch.ones(batch['coords'].shape[0], device=device) # use unscaled coords to extract latent for pLDDT prediction out_feat = plddt_latent_module( out_dict["denoised_coords"].detach(), t, batch) plddt_out_dict = plddt_out_module( out_feat["latent"].detach(), batch, ) elif args.backend == "mlx": t = mx.ones(batch['coords'].shape[0]) # use unscaled coords to extract latent for pLDDT prediction out_feat = plddt_latent_module( out_dict["denoised_coords"], t, batch) plddt_out_dict = plddt_out_module( out_feat["latent"], batch, ) # scale pLDDT to [0, 100] plddts = plddt_out_dict["plddt"] * 100.0 else: plddts = None out_dict = processor.postprocess(out_dict, batch) # sampled_coord = out_dict['denoised_coords'].detach() if args.backend == "torch": sampled_coord = out_dict['denoised_coords'].detach() else: sampled_coord = out_dict['denoised_coords'] pad_mask = batch['atom_pad_mask'] return sampled_coord, pad_mask, plddts def predict_structures_from_fastas(args): # create output directories output_dir = _resolve_path(args.output_dir) output_dir.mkdir(parents=True, exist_ok=True) prediction_dir = output_dir / f"predictions_{args.simplefold_model}" prediction_dir.mkdir(parents=True, exist_ok=True) cache = output_dir / "cache" cache.mkdir(parents=True, exist_ok=True) # set random seed for reproducibility pl.seed_everything(args.seed, workers=True) if args.backend == "mlx" and not MLX_AVAILABLE: args.backend = "torch" print("MLX not available, switch to torch backend.") # initialize models model, device = initialize_folding_model(args) plddt_latent_module, plddt_out_module = initialize_plddt_module(args, device) esm_model, esm_dict, af2_to_esm = initialize_esm_model(args, device) # initialize other components tokenizer, featurizer, processor, flow, sampler = initialize_others(args, device) # process fasta files to input format ccd_path = resolve_ccd_path(cache) data = check_fasta_inputs(_resolve_path(args.fasta_path)) if not data: raise ValueError("No valid input files found. Please check the input directory.") process_fastas( data=data, out_dir=output_dir, ccd_path=ccd_path, ) for struct_file in output_dir.glob("structures/*.npz"): record_file = output_dir / "records" / f"{struct_file.stem}.json" # prepare the target protein data for inference batch, structure, record = process_one_inference_structure( struct_file, record_file, tokenizer, featurizer, processor, esm_model, esm_dict, af2_to_esm, ) sampled_coord, pad_mask, plddts = generate_structure( args, batch, sampler, flow, processor, model, plddt_latent_module, plddt_out_module, device ) for i in range(args.nsample_per_protein): sampled_coord_i = sampled_coord[i] pad_mask_i = pad_mask[i] # save the generated structure structure_save = process_structure( deepcopy(structure), sampled_coord_i, pad_mask_i, record, backend=args.backend ) outname = f"{record.id}_sampled_{i}" save_structure( structure_save, prediction_dir, outname, output_format=args.output_format, plddts=plddts[i] if plddts is not None else None )