From 2f54bd431f49cdc1512d3b0f75366d8a4b8a02af Mon Sep 17 00:00:00 2001
From: Rocco Moretti <rmorettiase@gmail.com>
Date: Wed, 29 Jul 2026 15:21:47 -0500
Subject: [PATCH 1/2] feat: Ease MPNN command line usage

There were checks which prevented people from omitting --checkpoint_path when not using --config_json.
Additionally, simplify use of SolubleMPNN with the default parameter set.
---
 models/mpnn/src/mpnn/inference_engines/mpnn.py | 15 ++++++++++++---
 models/mpnn/src/mpnn/utils/inference.py        |  6 +-----
 2 files changed, 13 insertions(+), 8 deletions(-)

diff --git a/models/mpnn/src/mpnn/inference_engines/mpnn.py b/models/mpnn/src/mpnn/inference_engines/mpnn.py
index 50c48e15..0352be6f 100644
--- a/models/mpnn/src/mpnn/inference_engines/mpnn.py
+++ b/models/mpnn/src/mpnn/inference_engines/mpnn.py
@@ -50,8 +50,6 @@ def __init__(
         device: str | torch.device | None = None,
     ):
         # Store raw configuration
-        self.model_type = model_type
-        self.is_legacy_weights = is_legacy_weights
         self.out_directory = out_directory
         self.write_fasta = write_fasta
         self.write_structures = write_structures
@@ -61,12 +59,23 @@ def __init__(
         self.checkpoint_path = (
             str(
                 REGISTERED_CHECKPOINTS[
-                    self.model_type.replace("_", "")
+                    model_type.replace("_", "")
                 ].get_default_path()
             )
             if not checkpoint_path
             else checkpoint_path
         )
+        # The default weights sets are all legacy types
+        self.is_legacy_weights = (
+            True
+            if not checkpoint_path
+            else is_legacy_weights
+        )
+        # the soluble_mpnn type just changes the default weights -- otherwise it behaves like regular protein_mpnn
+        if model_type == "soluble_mpnn":
+            self.model_type = "protein_mpnn"
+        else:
+            self.model_type = model_type
 
         # Determine the device (supports XPU, CUDA, and CPU).
         if device is not None:
diff --git a/models/mpnn/src/mpnn/utils/inference.py b/models/mpnn/src/mpnn/utils/inference.py
index 19092aa5..e2e16ca5 100644
--- a/models/mpnn/src/mpnn/utils/inference.py
+++ b/models/mpnn/src/mpnn/utils/inference.py
@@ -139,7 +139,7 @@ def build_arg_parser() -> argparse.ArgumentParser:
     parser.add_argument(
         "--model_type",
         type=str,
-        choices=["protein_mpnn", "ligand_mpnn"],
+        choices=["protein_mpnn", "ligand_mpnn", "soluble_mpnn"],
         help="Model type to use.",
         default=MPNN_GLOBAL_INFERENCE_DEFAULTS["model_type"],
     )
@@ -593,15 +593,11 @@ def cli_to_json(args: argparse.Namespace) -> dict[str, Any]:
     # Build a single-input JSON object from CLI
     if (
         args.model_type is None
-        or args.checkpoint_path is None
-        or args.is_legacy_weights is None
         or args.structure_path is None
     ):
         raise ValueError(
             "When --config_json is not provided, "
             "--model_type, "
-            "--checkpoint_path, "
-            "--is_legacy_weights, "
             "--structure_path "
             "must all be specified."
         )

From f55267d4937f1f435a563fbff65d9bc7a96065af Mon Sep 17 00:00:00 2001
From: Rocco Moretti <rmorettiase@gmail.com>
Date: Wed, 29 Jul 2026 15:40:43 -0500
Subject: [PATCH 2/2] Ruff format

---
 models/mpnn/src/mpnn/inference_engines/mpnn.py | 12 ++----------
 models/mpnn/src/mpnn/utils/inference.py        |  5 +----
 2 files changed, 3 insertions(+), 14 deletions(-)

diff --git a/models/mpnn/src/mpnn/inference_engines/mpnn.py b/models/mpnn/src/mpnn/inference_engines/mpnn.py
index 0352be6f..576815d3 100644
--- a/models/mpnn/src/mpnn/inference_engines/mpnn.py
+++ b/models/mpnn/src/mpnn/inference_engines/mpnn.py
@@ -57,20 +57,12 @@ def __init__(
         # allow null for checkpoint path when foundry-installed
         # TODO: Currently this assumes the model type is the key in the registered path. Rework needed
         self.checkpoint_path = (
-            str(
-                REGISTERED_CHECKPOINTS[
-                    model_type.replace("_", "")
-                ].get_default_path()
-            )
+            str(REGISTERED_CHECKPOINTS[model_type.replace("_", "")].get_default_path())
             if not checkpoint_path
             else checkpoint_path
         )
         # The default weights sets are all legacy types
-        self.is_legacy_weights = (
-            True
-            if not checkpoint_path
-            else is_legacy_weights
-        )
+        self.is_legacy_weights = True if not checkpoint_path else is_legacy_weights
         # the soluble_mpnn type just changes the default weights -- otherwise it behaves like regular protein_mpnn
         if model_type == "soluble_mpnn":
             self.model_type = "protein_mpnn"
diff --git a/models/mpnn/src/mpnn/utils/inference.py b/models/mpnn/src/mpnn/utils/inference.py
index e2e16ca5..e631c383 100644
--- a/models/mpnn/src/mpnn/utils/inference.py
+++ b/models/mpnn/src/mpnn/utils/inference.py
@@ -591,10 +591,7 @@ def cli_to_json(args: argparse.Namespace) -> dict[str, Any]:
             return json.load(f)
 
     # Build a single-input JSON object from CLI
-    if (
-        args.model_type is None
-        or args.structure_path is None
-    ):
+    if args.model_type is None or args.structure_path is None:
         raise ValueError(
             "When --config_json is not provided, "
             "--model_type, "