Merge pull request #56 from RosettaCommons/fix/install-cmds

Fix installation issues

Author: Ubiquinone-dot

.env

@@ -4,11 +4,6 @@
 # - you will need to set the paths below as appropriate for your environment.
 #   We provide examples in the comments to give you an idea of the expected format.
 
-# Foundry install dir for checkpoints
-# Commented out by default since otherwise may be overridden by user export (load_dotenv(override=True) used at the moment)
-# TODO: Ensure override=False can be used.
-# FOUNDRY_CHECKPOINTS_DIR=
-
 # --- Mirrors to RCSB data ---
 
 # The `PDB_MIRROR_PATH` is a path to a local mirror of the PDB database. It's 
@@ -62,3 +57,8 @@ COLABFOLD_LOCAL_DB_PATH_CPU=
 # Network access (fallback; may cause IO-related issues)
 COLABFOLD_NET_DB_PATH_GPU=
 COLABFOLD_NET_DB_PATH_CPU=
+
+# Foundry install dir for checkpoints
+# Commented out by default since otherwise may be overridden by user export (load_dotenv(override=True) used at the moment)
+# TODO: Ensure override=False can be used.
+# FOUNDRY_CHECKPOINTS_DIR=

.gitignore

@@ -1,4 +1,6 @@
 # For docs / outputs from example notebooks;
+demo/*/**.cif*
+demo/*/**.json
 examples/*.cif
 **.ckpt
 **.pt

README.md

@@ -11,17 +11,21 @@ All models within Foundry rely on [AtomWorks](https://github.com/RosettaCommons/
 pip install rc-foundry[all]
 ```
 
-**Downloading weights** All models can be downloaded to a target folder with:
-
+**Downloading weights** Models can be downloaded to a target folder with:
+```
+foundry install base-models --checkpoint-dir <path/to/ckpt/dir>
 ```
-foundry install all --checkpoint-dir <path/to/ckpt/dir>
+where `checkpoint-dir` will be `~/.foundry/checkpoints` by default. Once installed, foundry sets the env `FOUNDRY_CHECKPOINTS_DIR` which it will use during inference or subsequent commands to find the checkpoints. `base-models` installs the latest RFD3, RF3 and MPNN variants - you can also download all of the models supported (including multiple checkpoints of RF3) with `all`, or by listing the models sequentially (e.g. `foundry install rfd3 rf3 ...`).
+To list the registry of available checkpoints:
 ```
-This will download all the models supported (including multiple checkpoints of RF3) but as a beginner you can start with:
+foundry list-available
 ```
-foundry install rfd3 ligandmpnn rf3 --checkpoint-dir  <path/to/ckpt/dir>
+To check what you already have downloaded (defaults to `$FOUNDRY_CHECKPOINTS_DIR` if set):
+```
+foundry list-installed
 ```
 
->*See `examples/all.ipynb` for how to run each model in a notebook.*
+>*See `examples/all.ipynb` for how to run each model and design proteins end-to-end in a notebook.*
 
 ### Google Colab
 For an interactive Google Colab notebook walking through a basic design pipeline with RFD3, MPNN, and RF3, please see the [IPD Design Pipeline Tutorial](https://colab.research.google.com/drive/1ZwIMV3n9h0ZOnIXX0GyKUuoiahgifBxh?usp=sharing).
@@ -68,11 +72,7 @@ For an interactive Google Colab notebook walking through a basic design pipeline
 Install both `foundry` and models in editable mode for development:
 
 ```bash
-# Install foundry and RF3 in editable mode
-uv pip install -e . -e ./models/rf3 -e ./models/rfd3 -e ./models/mpnn
-
-# Or install only foundry (no models)
-uv pip install -e .
+uv pip install -e '.[all,dev]'
 ```
 
 This approach allows you to:

models/rfd3/README.md

@@ -10,18 +10,21 @@ both are described in more detail below.
   <img src="docs/.assets/overview.png" alt="All-atom design with RFD3">
 </p>
 
-## Get Started
-1. Install RFdiffusion3. See [Main README](../../README.md) for instructions how to install all models to run full pipeline (recommended). If you have already installed all the models skip [here](#run-inference). 
+> [!NOTE]
+> Looking for config documentation? See [here](./docs/input.md)
+
+## Getting Started
+1. Install RFdiffusion3. See [Main README](../../README.md) for instructions how to install all models to run full pipeline (recommended). If you have already installed all the models skip [here](#run-inference).
 ```bash
 pip install rc-foundry[rfd3]
 ```
 2. Download checkpoint to your desired checkpoint location.
 ```bash
-foundry install rfd3 --checkpoint-dir /path/to/ckpt/dir
+foundry install rfd3 --checkpoint-dir <path/to/ckpt/dir>
 ```
 This sets `FOUNDRY_CHECKPOINTS_DIR` and will in future look for checkpoints in that directory, allowing you to run inference without supplying the checkpoint path. The checkpoint directory is optional, defaulting to `~/.foundry/checkpoints` if unset.
 
-## Run Inference
+## Running Inference
 
 To run inference (with foundry installed in your environment, or RFD3 & Foundry src in PYTHONPATH):
 ```bash
@@ -54,15 +57,6 @@ that you are running the examples from the `foundry/models/rfd3/docs`
 directory. If you would like to run RFD3 from a different location, 
 you will need to change the path in the `.json` file(s) before running.
 
-### Install HBPLUS for hydrogen bond conditioning:
-One of the examples shows how to incorporate hydrogen bond conditioning 
-into your designs. To make use of this feature, you will need to 
-additionally complete the following steps:
-
-1. Download hbplus from here: https://www.ebi.ac.uk/thornton-srv/software/HBPLUS/download.html (available for free)
-2. Follow the installation instruction here: https://www.ebi.ac.uk/thornton-srv/software/HBPLUS/install.html
-3. Update `HBPLUS_PATH` in `foundry/.env` file with the path to your `hbplus` executable.
-
 <table>
   <tr>
     <td align="center">
@@ -90,7 +84,7 @@ additionally complete the following steps:
   </tr>
 </table>
 
-## Training:
+## Training and Fine-Tuning
 
 We make available to the community not only the weights to run RFdiffusion3 but also the complete training code, easily extendable to additional use cases. Any AtomWorks-compatible dataset (and thus, any collection of structure files) can be readily incorporated and used for training or fine-tuning.
 
@@ -112,9 +106,10 @@ RFdiffusion3 supports arbitrary datasets of structure files for training and fin
 
 After setting up Hydra configs, launch a training run:
 ```bash
-uv run python models/rfd3/src/rfd3/train.py experiment=pretrain
+uv run python models/rfd3/src/rfd3/train.py experiment=pretrain ckpt_path=<path/to/ckpt>
 ```
 
+Supplying `ckpt_path=null` (default) will start with fresh weights.
 See the [path configs](/models/rfd3/configs/paths/) to customize data input and log output directories.
 
 ### Logging Configuration
@@ -162,6 +157,17 @@ In `models/rfd3/configs/datasets/design_base.yaml` there's the shared configs fo
 
 **Training with WandB:** We strongly recommend tracking your runs via wandb. To use it, simply have your WANDB_API_KEY set and use the wandb logger. For more details see [here](wandb.ai)
 
+# Appendix
+
+## Install HBPLUS for hydrogen bond conditioning:
+One of the examples shows how to incorporate hydrogen bond conditioning 
+into your designs. To make use of this feature, you will need to 
+additionally complete the following steps:
+
+1. Download hbplus from here: https://www.ebi.ac.uk/thornton-srv/software/HBPLUS/download.html (available for free)
+2. Follow the installation instruction here: https://www.ebi.ac.uk/thornton-srv/software/HBPLUS/install.html
+3. Update `HBPLUS_PATH` in `foundry/.env` file with the path to your `hbplus` executable.
+
 ## Citation
 
 If you use this code or data in your work, please consider citing:

models/rfd3/docs/input.md

@@ -1,4 +1,4 @@
-# RFdiffusion3 — Input specification (dialect **2**)
+# RFdiffusion3 — Input Specification & Command-line arguments
 
 ---
 
@@ -30,6 +30,23 @@ JSON inputs take the following top-level structure;
 }
 ```
 
+You can then run inference at the command line with:
+```
+rfd3 design out_dir=<path/to/outdir> inputs=<path/to/inputs>
+```
+In this document, we detail the syntax of the config structure.
+
+## CLI arguments
+Key CLI arguments (from the default config) to know include:
+- `n_batches` — number of batches to generate per input key (default: 1).
+- `diffusion_batch_size` — number of diffusion samples per batch (default: 8).
+- `specification` — JSON overrides for the per-example InputSpecification (default: `{}`). For example, you can run `rfd3 design inputs=null specification.length=200` for a quick debug of creating a 200-length protein.
+- `inference_sampler.num_timesteps` — diffusion timesteps for sampling (default: 200).
+- `inference_sampler.step_scale` — scales diffusion step size; higher → less diverse, more designable (default: 1.5).
+- `low_memory_mode` — memory-efficient tokenization mode; set `True` if GPU RAM is tight (default: False).
+
+The full config of default arguments that are applied can be seen in [inference_engine/rfdiffusion3.yaml](../configs/inference_engine/rfdiffusion3.yaml)
+
 ## InputSpecification fields
 
 Below is a table of all of the inputs that the `InputSpecification` accepts. Use these fields to describe what RFdiffusion3 should do with your inputs.
@@ -88,7 +105,7 @@ select_fixed_atoms:
 ```
 
 <p align="center">
-  <img src=".assets/input_selection.png" alt="InputSelection language for foundry">
+  <img src=".assets/input_selection.png" alt="InputSelection language for foundry" width=500>
 </p>
 
 ## Unindexing specifics
@@ -166,4 +183,4 @@ Let us know if you have any additional questions, we'd be happy to answer them!
 Below is a reference for more examples of different ways you can specify inputs to select from your pdb in configs; we hope the community can find use in this flexible system for future models!
 <p align="center">
   <img src=".assets/input_selection_large.png" alt="Input selection syntax" width=650>
-</p>
+</p>

src/foundry_cli/download_checkpoints.py

@@ -29,6 +29,13 @@
 console = Console()
 
 
+def _resolve_checkpoint_dir(checkpoint_dir: Optional[Path]) -> Path:
+    """Return user-specified checkpoint dir or fall back to default."""
+    return (
+        checkpoint_dir if checkpoint_dir is not None else get_default_checkpoint_dir()
+    )
+
+
 def download_file(url: str, dest: Path, verify_hash: Optional[str] = None) -> None:
     """Download a file with progress bar and optional hash verification.
 
@@ -123,7 +130,7 @@ def install_model(model_name: str, checkpoint_dir: Path, force: bool = False) ->
 def install(
     models: list[str] = typer.Argument(
         ...,
-        help="Models to install: 'all', 'rfd3', 'rf3', 'mpnn', or combination",
+        help="Models to install: 'all', 'rfd3', 'rf3', 'mpnn', or a combination thereof",
     ),
     checkpoint_dir: Optional[Path] = typer.Option(
         None,
@@ -136,24 +143,21 @@ def install(
     ),
 ):
     """Install model checkpoints for foundry.
-
     Examples:
-
         foundry install all
-
         foundry install rfd3 rf3
-
         foundry install proteinmpnn --checkpoint-dir ./checkpoints
     """
     # Determine checkpoint directory
-    if checkpoint_dir is None:
-        checkpoint_dir = get_default_checkpoint_dir()
+    checkpoint_dir = _resolve_checkpoint_dir(checkpoint_dir)
 
     console.print(f"[bold]Checkpoint directory:[/bold] {checkpoint_dir}")
     console.print()
 
     # Expand 'all' to all available models
     if "all" in models:
+        models_to_install = list(REGISTERED_CHECKPOINTS.keys())
+    elif "base-models" in models:
         models_to_install = ["rfd3", "proteinmpnn", "ligandmpnn", "rf3"]
     else:
         models_to_install = models
@@ -177,34 +181,28 @@ def install(
     console.print("[bold green]Installation complete![/bold green]")
 
 
-@app.command(name="list")
-def list_models():
+@app.command(name="list-available")
+def list_available():
     """List available model checkpoints."""
     console.print("[bold]Available models:[/bold]\n")
     for name, info in REGISTERED_CHECKPOINTS.items():
         console.print(f"  [cyan]{name:8}[/cyan] - {info.description}")
 
 
-@app.command()
-def show(
-    checkpoint_dir: Optional[Path] = typer.Option(
-        None,
-        "--checkpoint-dir",
-        "-d",
-        help="Checkpoint directory to show",
-    ),
-):
-    """Show installed checkpoints."""
-    if checkpoint_dir is None:
-        checkpoint_dir = get_default_checkpoint_dir()
+@app.command(name="list-installed")
+def list_installed():
+    """List installed checkpoints and their sizes."""
+    checkpoint_dir = _resolve_checkpoint_dir(None)
 
     if not checkpoint_dir.exists():
         console.print(
             f"[yellow]No checkpoints directory found at {checkpoint_dir}[/yellow]"
         )
         raise typer.Exit(0)
 
-    checkpoint_files = list(checkpoint_dir.glob("*.ckpt"))
+    checkpoint_files = list(checkpoint_dir.glob("*.ckpt")) + list(
+        checkpoint_dir.glob("*.pt")
+    )
     if not checkpoint_files:
         console.print(f"[yellow]No checkpoint files found in {checkpoint_dir}[/yellow]")
         raise typer.Exit(0)
@@ -219,21 +217,14 @@ def show(
     console.print(f"\n[bold]Total:[/bold] {total_size:.2f} GB")
 
 
-@app.command()
+@app.command(name="clean")
 def clean(
-    checkpoint_dir: Optional[Path] = typer.Option(
-        None,
-        "--checkpoint-dir",
-        "-d",
-        help="Checkpoint directory to clean",
-    ),
     confirm: bool = typer.Option(
         True, "--confirm/--no-confirm", help="Ask for confirmation before deleting"
     ),
 ):
     """Remove all downloaded checkpoints."""
-    if checkpoint_dir is None:
-        checkpoint_dir = get_default_checkpoint_dir()
+    checkpoint_dir = _resolve_checkpoint_dir(None)
 
     if not checkpoint_dir.exists():
         console.print(f"[yellow]No checkpoints found at {checkpoint_dir}[/yellow]")