Class InferenceWrapper

src/simplefold/wrapper.py:199–392

class InferenceWrapper:
    def __init__(
        self,
        output_dir,
        prediction_dir,
        num_steps,
        nsample_per_protein,
        tau,
        device,
        backend,
    ):
        self.num_steps = num_steps
        self.nsample_per_protein = nsample_per_protein
        self.tau = tau
        self.device = device
        self.backend = backend

        if self.backend == "mlx" and not MLX_AVAILABLE:
            self.backend = "torch"
            print("MLX not installed, switch to torch backend.")

        # create output directory
        output_dir = Path(output_dir)
        output_dir.mkdir(parents=True, exist_ok=True)

        # create cache directory
        cache = output_dir / "cache"
        cache.mkdir(parents=True, exist_ok=True)

        # create prediction directory
        prediction_dir = output_dir / prediction_dir
        prediction_dir.mkdir(parents=True, exist_ok=True)

        self.output_dir = output_dir
        self.cache = cache
        self.prediction_dir = prediction_dir

        self.initialize_esm_model()
        self.initialize_others()

    def initialize_esm_model(self):
        # load ESM2 model
        esm_model, esm_dict = esm_registry["esm2_3B"]()
        af2_to_esm = _af2_to_esm(esm_dict)

        if self.backend == "torch":
            esm_model = esm_model.to(self.device)
            af2_to_esm = af2_to_esm.to(self.device)
        elif self.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())))