rc-foundry 0.1.3__py3-none-any.whl → 0.1.4__py3-none-any.whl

foundry/inference_engines/base.py

@@ -69,15 +69,25 @@ class BaseInferenceEngine:
         self.verbose = verbose
 
         # Resolve checkpoint path
-        if '.' not in str(ckpt_path):
+        if "." not in str(ckpt_path):
             # Assume registered model
             name = str(ckpt_path)
-            assert name in REGISTERED_CHECKPOINTS, 'Checkpoint provided not and not in registered checkpoints'
+            assert (
+                name in REGISTERED_CHECKPOINTS
+            ), "Checkpoint provided not and not in registered checkpoints"
             ckpt = REGISTERED_CHECKPOINTS[name]
-            
+
             ckpt_path = ckpt.get_default_path()
-            ranked_logger.info("Using checkpoint from default installation directory, got: {}".format(str(ckpt_path)))
-            assert os.path.exists(ckpt_path), 'Invalid checkpoint: {}. And could not find checkpoint in default installation location: {}'.format(name, ckpt_path)
+            ranked_logger.info(
+                "Using checkpoint from default installation directory, got: {}".format(
+                    str(ckpt_path)
+                )
+            )
+            assert os.path.exists(
+                ckpt_path
+            ), "Invalid checkpoint: {}. And could not find checkpoint in default installation location: {}".format(
+                name, ckpt_path
+            )
         self.ckpt_path = Path(ckpt_path).resolve()
 
         # Set random seed (only if seed is not None)

foundry/inference_engines/checkpoint_registry.py

@@ -1,4 +1,5 @@
-'''Management of checkpoints'''
+"""Management of checkpoints"""
+
 import os
 from dataclasses import dataclass
 from pathlib import Path
@@ -11,10 +12,13 @@ def get_default_checkpoint_dir() -> Path:
     1. FOUNDRY_CHECKPOINTS_DIR environment variable
     2. ~/.foundry/checkpoints
     """
-    if "FOUNDRY_CHECKPOINTS_DIR" in os.environ and os.environ.get("FOUNDRY_CHECKPOINTS_DIR"):
+    if "FOUNDRY_CHECKPOINTS_DIR" in os.environ and os.environ.get(
+        "FOUNDRY_CHECKPOINTS_DIR"
+    ):
         return Path(os.environ["FOUNDRY_CHECKPOINTS_DIR"]).absolute()
     return Path.home() / ".foundry" / "checkpoints"
 
+
 @dataclass
 class RegisteredCheckpoint:
     url: str
@@ -28,39 +32,39 @@ class RegisteredCheckpoint:
 
 REGISTERED_CHECKPOINTS = {
     "rfd3": RegisteredCheckpoint(
-        url = "https://files.ipd.uw.edu/pub/rfd3/rfd3_foundry_2025_12_01_remapped.ckpt",
-        filename = "rfd3_latest.ckpt",
-        description = "RFdiffusion3 checkpoint",
+        url="https://files.ipd.uw.edu/pub/rfd3/rfd3_foundry_2025_12_01_remapped.ckpt",
+        filename="rfd3_latest.ckpt",
+        description="RFdiffusion3 checkpoint",
     ),
-    "rf3":  RegisteredCheckpoint(
-        url = "https://files.ipd.uw.edu/pub/rf3/rf3_foundry_01_24_latest_remapped.ckpt",
-        filename= "rf3_foundry_01_24_latest_remapped.ckpt",
-        description= "latest RF3 checkpoint trained with data until 1/2024 (expect best performance)",
+    "rf3": RegisteredCheckpoint(
+        url="https://files.ipd.uw.edu/pub/rf3/rf3_foundry_01_24_latest_remapped.ckpt",
+        filename="rf3_foundry_01_24_latest_remapped.ckpt",
+        description="latest RF3 checkpoint trained with data until 1/2024 (expect best performance)",
     ),
-    "proteinmpnn":  RegisteredCheckpoint(
-        url = "https://files.ipd.uw.edu/pub/ligandmpnn/proteinmpnn_v_48_020.pt",
-        filename= "proteinmpnn_v_48_020.pt",
-        description= "ProteinMPNN checkpoint",
+    "proteinmpnn": RegisteredCheckpoint(
+        url="https://files.ipd.uw.edu/pub/ligandmpnn/proteinmpnn_v_48_020.pt",
+        filename="proteinmpnn_v_48_020.pt",
+        description="ProteinMPNN checkpoint",
     ),
     "ligandmpnn": RegisteredCheckpoint(
-        url = "https://files.ipd.uw.edu/pub/ligandmpnn/ligandmpnn_v_32_010_25.pt",
-        filename= "ligandmpnn_v_32_010_25.pt",
-        description= "LigandMPNN checkpoint",
+        url="https://files.ipd.uw.edu/pub/ligandmpnn/ligandmpnn_v_32_010_25.pt",
+        filename="ligandmpnn_v_32_010_25.pt",
+        description="LigandMPNN checkpoint",
     ),
     # Other models
     "rf3_preprint_921": RegisteredCheckpoint(
-        url = "https://files.ipd.uw.edu/pub/rf3/rf3_foundry_09_21_preprint_remapped.ckpt",
-        filename = "rf3_foundry_09_21_preprint_remapped.ckpt",
-        description = "RF3 preprint checkpoint trained with data until 9/2021",
+        url="https://files.ipd.uw.edu/pub/rf3/rf3_foundry_09_21_preprint_remapped.ckpt",
+        filename="rf3_foundry_09_21_preprint_remapped.ckpt",
+        description="RF3 preprint checkpoint trained with data until 9/2021",
     ),
     "rf3_preprint_124": RegisteredCheckpoint(
-        url = "https://files.ipd.uw.edu/pub/rf3/rf3_foundry_01_24_preprint_remapped.ckpt",
-        filename = "rf3_foundry_01_24_preprint_remapped.ckpt",
-        description= "RF3 preprint checkpoint trained with data until 1/2024",
+        url="https://files.ipd.uw.edu/pub/rf3/rf3_foundry_01_24_preprint_remapped.ckpt",
+        filename="rf3_foundry_01_24_preprint_remapped.ckpt",
+        description="RF3 preprint checkpoint trained with data until 1/2024",
+    ),
+    "solublempnn": RegisteredCheckpoint(
+        url="https://files.ipd.uw.edu/pub/ligandmpnn/solublempnn_v_48_020.pt",
+        filename="solublempnn_v_48_020.pt",
+        description="SolubleMPNN checkpoint",
     ),
-    "solublempnn":  RegisteredCheckpoint(
-        url =  "https://files.ipd.uw.edu/pub/ligandmpnn/solublempnn_v_48_020.pt",
-        filename= "solublempnn_v_48_020.pt",
-        description= "SolubleMPNN checkpoint"
-    )
 }

foundry/version.py

@@ -28,7 +28,7 @@ version_tuple: VERSION_TUPLE
 commit_id: COMMIT_ID
 __commit_id__: COMMIT_ID
 
-__version__ = version = '0.1.3'
-__version_tuple__ = version_tuple = (0, 1, 3)
+__version__ = version = '0.1.4'
+__version_tuple__ = version_tuple = (0, 1, 4)
 
 __commit_id__ = commit_id = None

foundry_cli/download_checkpoints.py

@@ -5,7 +5,6 @@ from pathlib import Path
 from typing import Optional
 from urllib.request import urlopen
 
-import rootutils
 import typer
 from dotenv import find_dotenv, load_dotenv, set_key
 from rich.console import Console
@@ -29,6 +28,7 @@ load_dotenv(override=True)
 app = typer.Typer(help="Foundry model checkpoint installation utilities")
 console = Console()
 
+
 def download_file(url: str, dest: Path, verify_hash: Optional[str] = None) -> None:
     """Download a file with progress bar and optional hash verification.
 
@@ -81,9 +81,7 @@ def download_file(url: str, dest: Path, verify_hash: Optional[str] = None) -> No
         console.print("[green]✓[/green] Hash verification passed")
 
 
-def install_model(
-    model_name: str, checkpoint_dir: Path, force: bool = False
-) -> None:
+def install_model(model_name: str, checkpoint_dir: Path, force: bool = False) -> None:
     """Install a single model checkpoint.
 
     Args:
@@ -112,9 +110,7 @@ def install_model(
     )
 
     try:
-        download_file(
-            checkpoint_info.url, dest_path, checkpoint_info.sha256
-        )
+        download_file(checkpoint_info.url, dest_path, checkpoint_info.sha256)
         console.print(
             f"[green]✓[/green] Successfully installed {model_name} to {dest_path}"
         )
@@ -158,7 +154,7 @@ def install(
 
     # Expand 'all' to all available models
     if "all" in models:
-        models_to_install = ['rfd3', 'proteinmpnn', 'ligandmpnn', 'rf3']
+        models_to_install = ["rfd3", "proteinmpnn", "ligandmpnn", "rf3"]
     else:
         models_to_install = models
 
@@ -167,15 +163,16 @@ def install(
         install_model(model_name, checkpoint_dir, force)
         console.print()
 
-    set_key(
-        dotenv_path=find_dotenv(),
-        key_to_set='FOUNDRY_CHECKPOINTS_DIR',
-        value_to_set=str(checkpoint_dir),
-        export = False,
-    )
-    console.print(
-        f"Set checkpoint installation directory to: {checkpoint_dir}"
-    )
+    # Try to persist checkpoint dir to .env (optional, may not exist in Colab etc.)
+    dotenv_path = find_dotenv()
+    if dotenv_path:
+        set_key(
+            dotenv_path=dotenv_path,
+            key_to_set="FOUNDRY_CHECKPOINTS_DIR",
+            value_to_set=str(checkpoint_dir),
+            export=False,
+        )
+        console.print(f"Saved FOUNDRY_CHECKPOINTS_DIR to {dotenv_path}")
 
     console.print("[bold green]Installation complete![/bold green]")
 
@@ -209,9 +206,7 @@ def show(
 
     checkpoint_files = list(checkpoint_dir.glob("*.ckpt"))
     if not checkpoint_files:
-        console.print(
-            f"[yellow]No checkpoint files found in {checkpoint_dir}[/yellow]"
-        )
+        console.print(f"[yellow]No checkpoint files found in {checkpoint_dir}[/yellow]")
         raise typer.Exit(0)
 
     console.print(f"[bold]Installed checkpoints in {checkpoint_dir}:[/bold]\n")
@@ -247,9 +242,7 @@ def clean(
     # List files to delete
     checkpoint_files = list(checkpoint_dir.glob("*.ckpt"))
     if not checkpoint_files:
-        console.print(
-            f"[yellow]No checkpoint files found in {checkpoint_dir}[/yellow]"
-        )
+        console.print(f"[yellow]No checkpoint files found in {checkpoint_dir}[/yellow]")
         raise typer.Exit(0)
 
     console.print("[bold]Files to delete:[/bold]")

mpnn/inference_engines/mpnn.py

@@ -54,11 +54,18 @@ class MPNNInferenceEngine:
         self.out_directory = out_directory
         self.write_fasta = write_fasta
         self.write_structures = write_structures
-        
+
         # 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[self.model_type.replace('_', '')].get_default_path()) \
-            if not checkpoint_path else checkpoint_path
+        self.checkpoint_path = (
+            str(
+                REGISTERED_CHECKPOINTS[
+                    self.model_type.replace("_", "")
+                ].get_default_path()
+            )
+            if not checkpoint_path
+            else checkpoint_path
+        )
 
         # Determine the device.
         if device is not None:

{rc_foundry-0.1.3.dist-info → rc_foundry-0.1.4.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: rc-foundry
-Version: 0.1.3
+Version: 0.1.4
 Summary: Shared utilities and training infrastructure for biomolecular structure prediction models.
 Author-email: Institute for Protein Design <contact@ipd.uw.edu>
 License: BSD 3-Clause License
@@ -123,7 +123,7 @@ foundry install rfd3 ligandmpnn rf3 --checkpoint_dir  <path/to/ckpt/dir>
 > *See [models/rfd3/README.md](models/rfd3/README.md) for complete documentation.*
 
 <div align="center">
-  <img src="docs/_static/rfd3_trajectory.png" alt="RFdiffusion3 generation trajectory." width="400">
+  <img src="models/rfd3/docs/.assets/trajectory.png" alt="RFdiffusion3 generation trajectory." width="700">
 </div>
 
 ### ProteinMPNN

{rc_foundry-0.1.3.dist-info → rc_foundry-0.1.4.dist-info}/RECORD

@@ -1,7 +1,7 @@
 foundry/__init__.py,sha256=H8S1nl5v6YeW8ggn1jKy4GdtH7c-FGS-j7CqUCAEnAU,1926
 foundry/common.py,sha256=Aur8mH-CNmcUqSsw7VgaCQSW5sH1Bqf8Da91jzxPV1Y,3035
 foundry/constants.py,sha256=0n1wBKCvNuw3QaQehSbmsHYkIdaGn3tLeRFItBrdeHY,913
-foundry/version.py,sha256=q5nF98G8SoVeJqaknL0xdyxtv0egsqb0fK06_84Izu8,704
+foundry/version.py,sha256=rLCrf4heo25FJtBY-2Ap7ZuWW-5FS7sqTjsolIUuI5c,704
 foundry/callbacks/__init__.py,sha256=VsRT1e4sqlJHPcTCsfupMEx82Iz-LoOAGPpwvf_OJeE,126
 foundry/callbacks/callback.py,sha256=xZBo_suP4bLrP6gl5uJPbaXm00DXigePa6dMeDxucgg,3890
 foundry/callbacks/health_logging.py,sha256=tEtkByOlaAA7nnelxb7PbM9_dcIgOsdbxCdQY3K5pMc,16664
@@ -9,8 +9,8 @@ foundry/callbacks/metrics_logging.py,sha256=Vekzs831d-HE7TfLJZnQ45iPeG9ziQWLQaMB
 foundry/callbacks/timing_logging.py,sha256=u-r0hKp7fWOY3mLk7CcuIwHgZbhte13m5M09xNgatZA,2343
 foundry/callbacks/train_logging.py,sha256=Xs3tmZA88qLxmdSOwt-x8YKN4NKb1kVm59uptNXl4Qo,10399
 foundry/hydra/resolvers.py,sha256=xyJzo6OeWAc_LOu8RiHhX7_CRNoLZ22626AvYHXYl4U,2186
-foundry/inference_engines/base.py,sha256=qv5Gnk6NIxMpIxZ3oeOJurqMMUzBCZgfzHckb7SSzmU,8227
-foundry/inference_engines/checkpoint_registry.py,sha256=xRN3PjtmcnN7aEEhDR0MKhp1yaMDXHRXdLAGTDxi_Yk,2563
+foundry/inference_engines/base.py,sha256=ZHdlmGUqH4-p3v4RdrLH-Ps8_zalr7j5mQ4x-S53N4M,8375
+foundry/inference_engines/checkpoint_registry.py,sha256=kt2Z1JhrAjoOiEpkIIQ0sLttie1ceL8OgXUBmmyA6iw,2544
 foundry/metrics/__init__.py,sha256=qL4wwaiQ7EtR30pmZ9MCknqx909BJcNvHVmNJUaz_WM,236
 foundry/metrics/losses.py,sha256=2CLUmf7oCdFUCvgJukdNkff0FVG3BlATI-NI60TtpVY,903
 foundry/metrics/metric.py,sha256=23pKh_Ra0EcHGo5cSzYQQrUGr5zWRxeufKSJ58tfXXo,12687
@@ -34,12 +34,12 @@ foundry/utils/squashfs.py,sha256=QlcwuJyVe-QVfIOS7o1QfLhaCQPNzzox7ln4n8dcYEg,523
 foundry/utils/torch.py,sha256=OLsqoxw4CTXbGzWUHernLUT7uQjLu0tVPtD8h8747DI,11211
 foundry/utils/weights.py,sha256=btz4S02xff2vgiq4xMfiXuhK1ERafqQPtmimo1DmoWY,10381
 foundry_cli/__init__.py,sha256=0BxY2RUKJLaMXUGgypPCwlTskTEFdVnkhTR4C4ft2Kw,52
-foundry_cli/download_checkpoints.py,sha256=CDMxm3otzpb_c46AE2OwQY7aG7K_vnCJLKjP4U3b0Oc,8380
+foundry_cli/download_checkpoints.py,sha256=UCNdy4VZyJe1PH_lnVLqy-VSMuTu875mGGd99ma7fTQ,8426
 mpnn/__init__.py,sha256=hgQcXFaCbAxFrhydVAy0xj8yC7UJF-GCCFhqD0sZ7I4,57
 mpnn/inference.py,sha256=wPtGR325eVRVeesXoWtBK6b_-VcU8BZae5IfQN3-mvA,1669
 mpnn/train.py,sha256=9eQGBd3rdNF5Zr2w8oUgETbqxBavNBajtA6Vbc5zESE,10239
 mpnn/collate/feature_collator.py,sha256=LpzAFWo1VMa06dJLmfUWZsKe4xvLZjHbx4RICg2lgbQ,10510
-mpnn/inference_engines/mpnn.py,sha256=985Ce84dWKNZk_5_dk1eJScYIgBGUKUhy-xDVCu15PA,21631
+mpnn/inference_engines/mpnn.py,sha256=PmDEsIFipdk2fY57FA-vCp4evoU83DVVuUVmlViUtWk,21725
 mpnn/loss/nll_loss.py,sha256=KmdNe-BCzGYtijjappzBArQcT1gHVlJnKdY1PYQ4mhU,5947
 mpnn/metrics/nll.py,sha256=T6oMeUOEeHZzOMTH8NHFtsH9vUwLAsHQDPszzj4YKXI,15299
 mpnn/metrics/sequence_recovery.py,sha256=YDw_LmH-a3ajBYWK0mucJEQvw0_VEyxvrBN7da4vX8Q,19034
@@ -119,7 +119,7 @@ rfd3/__init__.py,sha256=2Wto2IsUIj2lGag9m_gqgdCwBNl5p21-Xnr7W_RpU3c,348
 rfd3/callbacks.py,sha256=Zjt8RiaYWquoKOwRmC_wCUbRbov-V4zd2_73zjhgDHE,2783
 rfd3/cli.py,sha256=TZpZouXGmwAMFaH8hp4r3q9tbUi1xlcN8n_r8hO2q8c,1424
 rfd3/constants.py,sha256=wLvDzrThpOrK8T3wGFNQeGrhAXOJQze8l3v_7pjIdMM,13141
-rfd3/engine.py,sha256=HCIHgB_4xFBYz7d7DGnqibnFcdC7RGXQ6qHgUJEvm2M,20889
+rfd3/engine.py,sha256=La_dB48Ewz0IdY1ocxvSWg-PXVAsySm0OGvwyz42lI8,20824
 rfd3/run_inference.py,sha256=dubMwEFkNPOq_yYf0ny37qvvEkRjNNPRFksZgmEFkUc,1520
 rfd3/train.py,sha256=rHswffIUhOae3_iYyvAiQ3jALoFuzrcRUgMlbJLinlI,7947
 rfd3/inference/datasets.py,sha256=u-2U7deHXu-iOs7doiKKynewP-NEyJfdORSTDzUSaQI,6538
@@ -173,8 +173,8 @@ rfd3/transforms/virtual_atoms.py,sha256=UpmxzPPd5FaJigcRoxgLSHHrLLOqsCvZ5PPZfQSG
 rfd3/utils/inference.py,sha256=RQp5CCy6Z6uHVZ2Mx0zmmGluYEOrASke4bABtfRjpy0,26448
 rfd3/utils/io.py,sha256=wbdjUTQkDc3RCSM7gdogA-XOKR68HeQ-cfvyN4pP90w,9849
 rfd3/utils/vizualize.py,sha256=HPlczrA3zkOuxV5X05eOvy_Oga9e3cPnFUXOEP4RR_g,11046
-rc_foundry-0.1.3.dist-info/METADATA,sha256=gY8d46JTRTfot6O9MqW34E0D1O64xNhMeVdpJrQ-G18,10578
-rc_foundry-0.1.3.dist-info/WHEEL,sha256=WLgqFyCfm_KASv4WHyYy0P3pM_m7J5L9k2skdKLirC8,87
-rc_foundry-0.1.3.dist-info/entry_points.txt,sha256=BmiWCbWGtrd_lSOFMuCLBXyo84B7Nco-alj7hB0Yw9A,130
-rc_foundry-0.1.3.dist-info/licenses/LICENSE.md,sha256=NKtPCJ7QMysFmzeDg56ZfUStvgzbq5sOvRQv7_ddZOs,1533
-rc_foundry-0.1.3.dist-info/RECORD,,
+rc_foundry-0.1.4.dist-info/METADATA,sha256=hzcS1buvLzRRAv7rPRgKwYjeNDL_iTGyR6u8CRpL-Ic,10585
+rc_foundry-0.1.4.dist-info/WHEEL,sha256=WLgqFyCfm_KASv4WHyYy0P3pM_m7J5L9k2skdKLirC8,87
+rc_foundry-0.1.4.dist-info/entry_points.txt,sha256=BmiWCbWGtrd_lSOFMuCLBXyo84B7Nco-alj7hB0Yw9A,130
+rc_foundry-0.1.4.dist-info/licenses/LICENSE.md,sha256=NKtPCJ7QMysFmzeDg56ZfUStvgzbq5sOvRQv7_ddZOs,1533
+rc_foundry-0.1.4.dist-info/RECORD,,

rfd3/engine.py

@@ -15,7 +15,6 @@ from toolz import merge_with
 
 from foundry.common import exists
 from foundry.inference_engines.base import BaseInferenceEngine
-from foundry.inference_engines.checkpoint_registry import REGISTERED_CHECKPOINTS
 from foundry.utils.alignment import weighted_rigid_align
 from foundry.utils.ddp import RankedLogger
 from rfd3.constants import SAVED_CONDITIONING_ANNOTATIONS
@@ -38,7 +37,9 @@ ranked_logger = RankedLogger(__name__, rank_zero_only=True)
 
 @dataclass(kw_only=True)
 class RFD3InferenceConfig:
-    ckpt_path: str | Path = 'rfd3'  # Defaults to foundry installation upon instantiation
+    ckpt_path: str | Path = (
+        "rfd3"  # Defaults to foundry installation upon instantiation
+    )
     diffusion_batch_size: int = 16
 
     # RFD3 specific