Upload folder using huggingface_hub

README.md

@@ -9,10 +9,7 @@ All three models — RFD3, ProteinMPNN, and LigandMPNN — rely on [Foundry](htt
 ### Installation
 
 ```bash
-# Install foundry (provides model implementations + AtomWorks)
 pip install rc-foundry[all]
-
-# Install diffusers with modular pipeline support
 pip install diffusers
 ```
 
@@ -22,72 +19,114 @@ pip install diffusers
 import torch
 from diffusers import ModularPipeline
 
-# Load the pipeline
-pipe = ModularPipeline.from_pretrained(
-    "dn6/RFDiffusion-3",
-    trust_remote_code=True,
-)
-pipe.load_components(
-    device_map="cuda",
-    torch_dtype=torch.bfloat16,
-    trust_remote_code=True,
-)
+pipe = ModularPipeline.from_pretrained("dn6/RFDiffusion-3", trust_remote_code=True)
+pipe.load_components(device_map="cuda", torch_dtype=torch.bfloat16, trust_remote_code=True)
 
-# Generate a 100-residue protein backbone
 state = pipe(contigs="100")
-
-# Access coordinates
-print(state.output.xyz.shape)  # [B, L, 3]
+print(state.output.xyz.shape)  # [1, 100, 3]
 ```
 
-## Why Diffusers?
+## Workflows
 
-Wrapping RFdiffusion3 as a diffusers `ModularPipeline` gives you access to the full diffusers ecosystem out of the box:
+The active workflow is selected automatically based on which inputs you provide:
 
-### CPU Offloading
+| Workflow | Trigger inputs | What runs |
+|----------|---------------|-----------|
+| `structure_only` | `contigs` | RFdiffusion3 |
+| `structure_and_sequence` | `contigs`, `temperature` | RFdiffusion3 → MPNN |
+| `motif_structure_and_sequence` | `contigs`, `input_xyz`, `temperature` | Motif-conditioned RFdiffusion3 → MPNN |
 
-Run large models on limited VRAM by offloading components to CPU when not in use:
+### Structure Only
 
 ```python
-pipe.enable_model_cpu_offload()
 state = pipe(contigs="100")
+print(state.output.xyz.shape)  # [1, 100, 3]
 ```
 
-### Hub Integration
+### Structure + Sequence Design
 
-All models are hosted on the Hugging Face Hub. Load by repo ID, share fine-tuned variants, and version your checkpoints:
+Passing `temperature` triggers the MPNN sequence design step. Load an MPNN variant first:
 
 ```python
-pipe = ModularPipeline.from_pretrained("dn6/RFDiffusion-3", trust_remote_code=True)
+from diffusers import AutoModel
+
+mpnn = AutoModel.from_pretrained("dn6/RFDiffusion-3", subfolder="mpnn", trust_remote_code=True)
+pipe.update_components(mpnn=mpnn)
+
+state = pipe(contigs="100", temperature=0.1)
+print(state.mpnn_output.designed_sequence)  # e.g. "MKVLSEG..."
 ```
 
-### LoRA Fine-Tuning
+Three MPNN variants are available:
 
-Fine-tune RFdiffusion3 on custom datasets with LoRA — supported natively by `ModelMixin`:
+| Subfolder | Variant | Params | Description |
+|-----------|---------|--------|-------------|
+| `mpnn/` | ProteinMPNN | 1.66M | Standard protein sequence design |
+| `mpnn_ligand/` | LigandMPNN | 2.62M | Ligand-aware sequence design |
+| `mpnn_soluble/` | SolubleMPNN | 1.66M | Optimized for soluble proteins |
+
+### Motif-Conditioned Design
+
+Passing `input_xyz` enables motif conditioning — fix specific residues in place while designing the rest:
 
 ```python
-from peft import LoraConfig
+import torch
 
-lora_config = LoraConfig(r=16, lora_alpha=16, target_modules=["to_q", "to_k", "to_v"])
-pipe.transformer.add_adapter(lora_config)
+motif_coords = torch.randn(16, 3)  # [N_motif, 3]
+state = pipe(
+    contigs="A10-25/50",
+    input_xyz=motif_coords,
+    temperature=0.1,
+)
+```
+
+### Full Design Pipeline
+
+The three pipelines can be composed into a complete protein design workflow:
 
-# After training
-pipe.transformer.save_pretrained("my-rfd3-lora")
 ```
+RFD3 (design backbone) → MPNN (design sequence) → RF3 (validate fold)
+```
+
+```python
+import torch
+from diffusers import AutoModel, ModularPipeline
 
-### Composable Workflows
+# 1. Design a backbone + sequence
+design_pipe = ModularPipeline.from_pretrained("dn6/RFDiffusion-3", trust_remote_code=True)
+design_pipe.load_components(device_map="cuda", torch_dtype=torch.bfloat16, trust_remote_code=True)
+
+mpnn = AutoModel.from_pretrained("dn6/RFDiffusion-3", subfolder="mpnn", trust_remote_code=True)
+design_pipe.update_components(mpnn=mpnn)
+
+state = design_pipe(contigs="100", temperature=0.1, output_type="cif.gz", output_path="design")
+designed_sequence = state.mpnn_output.designed_sequence
+
+# 2. Validate the fold with RF3
+fold_pipe = ModularPipeline.from_pretrained("dn6/RosettaFold-3", trust_remote_code=True)
+fold_pipe.load_components(device_map="cuda", torch_dtype=torch.bfloat16, trust_remote_code=True)
+
+state = fold_pipe(sequence=designed_sequence, output_type="cif.gz", output_path="prediction")
+```
+
+> [RF3](https://www.biorxiv.org/content/10.1101/2025.08.14.670328) (RosettaFold3) is available as a separate pipeline at [`dn6/RosettaFold-3`](https://huggingface.co/dn6/RosettaFold-3).
+
+## Customizing Workflows
 
 Inspect, swap, and extend pipeline blocks at runtime:
 
 ```python
-# Inspect the pipeline
+# Inspect the pipeline structure
 print(pipe.blocks)
 
 # Swap ProteinMPNN for LigandMPNN
 mpnn = AutoModel.from_pretrained("dn6/RFDiffusion-3", subfolder="mpnn_ligand", trust_remote_code=True)
 pipe.update_components(mpnn=mpnn)
 
-# Add a custom post-processing block
+# Select a workflow explicitly
+workflow = pipe.get_workflow("structure_and_sequence")
+
+# Add a custom block
 from diffusers.modular_pipelines import ModularPipelineBlocks, PipelineState
 from diffusers.modular_pipelines.modular_pipeline_utils import InputParam, OutputParam
 
@@ -108,14 +147,11 @@ class ScoreDesignStep(ModularPipelineBlocks):
         self.set_block_state(state, block_state)
         return components, state
 
-# Insert after the decoder
 pipe._blocks.sub_blocks.insert("score", ScoreDesignStep(), index=3)
 ```
 
 ## Output Types
 
-All output types return the base tensors (`xyz`, `sequence_indices`, `sequence_logits`). The `output_type` parameter controls what additional format is produced:
-
 | `output_type` | Additional output | Writes to disk |
 |---|---|---|
 | `"tensor"` | — | — |
@@ -126,142 +162,22 @@ All output types return the base tensors (`xyz`, `sequence_indices`, `sequence_l
 CIF outputs use [AtomWorks](https://github.com/RosettaCommons/atomworks) `to_cif_file` and return [biotite](https://www.biotite-python.org/) `AtomArray` / `AtomArrayStack` objects, matching the foundry output format.
 
 ```python
-# Save as compressed CIF (matches foundry output format)
+# Save as compressed CIF
 state = pipe(contigs="100", output_type="cif.gz", output_path="design_0")
 
-# AtomArray is available directly
+# Access AtomArray directly
 atom_array = state.output.atom_array
-print(atom_array)  # biotite AtomArray with CA coords + residue names
 
-# Denoising trajectory as AtomArrayStack (one model per step)
+# Denoising trajectory
 trajectory = state.output.trajectory_stack
 
-# PDB string output
+# PDB output
 state = pipe(contigs="100", output_type="pdb", output_path="design_0.pdb")
 print(state.output.pdb_string[:200])
 ```
 
-## Models
-
-By default, `load_components` loads the RFdiffusion3 transformer and scheduler. MPNN models are optional — load them separately when you need sequence design.
-
-### RFdiffusion3 (RFD3)
-
-[RFdiffusion3](https://www.biorxiv.org/content/10.1101/2025.09.18.676967v2) is an all-atom generative model that designs protein structures via iterative denoising. Uses an EDM noise schedule with 200 steps. Loaded automatically by `load_components`.
-
-| Component | Subfolder | Description |
-|-----------|-----------|-------------|
-| `transformer` | `transformer/` | `RFDiffusionTransformerModel` (168M params) |
-| `scheduler` | `scheduler/` | `RFDiffusionScheduler` (EDM noise schedule + Euler stepping) |
-
-### ProteinMPNN / LigandMPNN
-
-[ProteinMPNN](https://www.science.org/doi/10.1126/science.add2187) and [LigandMPNN](https://www.nature.com/articles/s41592-025-02626-1) are inverse-folding models that design amino acid sequences for a given protein backbone. These are **not** loaded by default — load them with `AutoModel` and register via `update_components`:
-
-```python
-from diffusers import AutoModel
-
-mpnn = AutoModel.from_pretrained("dn6/RFDiffusion-3", subfolder="mpnn", trust_remote_code=True)
-pipe.update_components(mpnn=mpnn)
-```
-
-Three variants are available:
-
-| Subfolder | Variant | Params | Description |
-|-----------|---------|--------|-------------|
-| `mpnn/` | ProteinMPNN | 1.66M | Standard protein sequence design |
-| `mpnn_ligand/` | LigandMPNN | 2.62M | Ligand-aware sequence design |
-| `mpnn_soluble/` | SolubleMPNN | 1.66M | Optimized for soluble proteins |
-
-## Workflows
-
-The active workflow is selected automatically based on which inputs you provide. Passing `temperature` triggers the MPNN sequence design step; passing `input_xyz` enables motif conditioning.
-
-| Workflow | Trigger inputs | What runs |
-|----------|---------------|-----------|
-| `structure_only` | `contigs` | RFdiffusion3 |
-| `structure_and_sequence` | `contigs`, `temperature` | RFdiffusion3 → MPNN |
-| `motif_structure_and_sequence` | `contigs`, `input_xyz`, `temperature` | Motif-conditioned RFdiffusion3 → MPNN |
-
-> Workflows that include MPNN require loading an MPNN variant first (see above).
-
-You can also select a workflow explicitly:
-
-```python
-workflow = pipe.get_workflow("structure_and_sequence")
-```
-
-### Structure Only
-
-```python
-state = pipe(contigs="100")
-print(state.output.xyz.shape)  # [1, 100, 3]
-```
-
-### Structure + Sequence Design
-
-```python
-from diffusers import AutoModel
-
-# Load an MPNN variant and register it
-mpnn = AutoModel.from_pretrained("dn6/RFDiffusion-3", subfolder="mpnn", trust_remote_code=True)
-pipe.update_components(mpnn=mpnn)
-
-# Passing temperature triggers the MPNN step
-state = pipe(contigs="100", temperature=0.1)
-print(state.mpnn_output.designed_sequence)  # e.g. "MKVLSEG..."
-```
-
-### Motif-Conditioned Design
-
-```python
-import torch
-
-motif_coords = torch.randn(16, 3)  # [N_motif, 3]
-state = pipe(
-    contigs="A10-25/50",
-    input_xyz=motif_coords,
-    temperature=0.1,
-)
-```
-
-## Full Design Pipeline
-
-The three pipelines can be composed into a complete protein design workflow:
-
-```
-RFD3 (design backbone) → MPNN (design sequence) → RF3 (validate fold)
-```
-
-Each is a standalone `ModularPipeline` that can run independently. Here's the full end-to-end flow:
-
-```python
-import torch
-from diffusers import AutoModel, ModularPipeline
-
-# 1. Design a backbone with RFdiffusion3
-design_pipe = ModularPipeline.from_pretrained("dn6/RFDiffusion-3", trust_remote_code=True)
-design_pipe.load_components(device_map="cuda", torch_dtype=torch.bfloat16, trust_remote_code=True)
-
-mpnn = AutoModel.from_pretrained("dn6/RFDiffusion-3", subfolder="mpnn", trust_remote_code=True)
-design_pipe.update_components(mpnn=mpnn)
-
-state = design_pipe(contigs="100", temperature=0.1, output_type="cif.gz", output_path="design")
-designed_sequence = state.mpnn_output.designed_sequence
-
-# 2. Validate the design with RF3 (structure prediction)
-fold_pipe = ModularPipeline.from_pretrained("dn6/RosettaFold-3", trust_remote_code=True)
-fold_pipe.load_components(device_map="cuda", torch_dtype=torch.bfloat16, trust_remote_code=True)
-
-state = fold_pipe(sequence=designed_sequence, output_type="cif.gz", output_path="prediction")
-```
-
-> [RF3](https://www.biorxiv.org/content/10.1101/2025.08.14.670328) (RosettaFold3) is available as a separate pipeline at [`dn6/RosettaFold-3`](https://huggingface.co/dn6/RosettaFold-3).
-
 ## Citation
 
-If you use this code, please cite the relevant work:
-
 ```bibtex
 @article{butcher2025_rfdiffusion3,
     author = {Butcher, Jasper and Krishna, Rohith and Mitra, Raktim and Brent, Rafael Isaac and Li, Yanjing and Corley, Nathaniel and Kim, Paul T and Funk, Jonathan and Mathis, Simon Valentin and Salike, Saman and Muraishi, Aiko and Eisenach, Helen and Thompson, Tuscan Rock and Chen, Jie and Politanska, Yuliya and Sehgal, Enisha and Coventry, Brian and Zhang, Odin and Qiang, Bo and Didi, Kieran and Kazman, Maxwell and DiMaio, Frank and Baker, David},