Method __init__

src/simplefold/wrapper.py:200–237

(
        self,
        output_dir,
        prediction_dir,
        num_steps,
        nsample_per_protein,
        tau,
        device,
        backend,
    )

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