SS-Scaffold: Secondary-Structure-Conditioned Motif Scaffolding

SS-Scaffold extends Microsoft FoldingDiff with secondary-structure (SS) conditioning and RePaint-style motif inpainting. Given a structural motif (e.g. a helix from an existing PDB), the model designs a complete protein backbone of a chosen length where the motif residues are preserved and the surrounding residues are generated by a denoising diffusion model.

Adapted from Wu, K. E. et al., FoldingDiff: Diffusion-based protein backbone generation. New work lives under ss_scaffold/; foldingdiff/ is reused largely as-is.

A pretrained checkpoint is hosted at Yavanash/ss-scaffold-v2 and is downloaded automatically on first sampling — you do not need to train.

---

Requirements

• Windows or Linux, Python 3.8 (matches environment.yml)
• CUDA-capable NVIDIA GPU. A GPU must be available; CPU sampling is not supported in practice.
• ~5 GB disk for the HF snapshot + caches

External tools (only for the full evaluation pipeline): TMalign, ProteinMPNN, ESMFold.

---

Installation (conda)

git clone https://github.com/Yavanash/ss_scaffold.git
cd ss_scaffold

conda env create -f environment.yml
conda activate foldingdiff

python -c "import foldingdiff, ss_scaffold; print('OK')"

---

Input PDB preprocessing

The motif PDB must be cleaned before it is passed to the sampler. Specifically the file must contain:

• No header / metadata records (no HEADER, TITLE, REMARK, CRYST1, SEQRES, etc.)
• Only ATOM records — no HETATM, no TER, no END, no MODEL/ENDMDL markers
• Exactly one model (one chain, one conformation)

Quick clean-up:

grep '^ATOM' raw.pdb > motif.pdb

If your file contains multiple models or alternative conformations, keep only the first model and a single altloc before running grep. Files that violate these rules will cause DSSP / biotite parsing errors at sample time.

---

Sampling (motif scaffolding)

python -m ss_scaffold.sample \
    --motif-pdb data/4HB1.pdb \
    --motif-range 10-30 \
    --total-length 80 \
    --motif-target-start 20 \
    --flank-ss C \
    --n-samples 8 \
    --out generated/run01

On first run the checkpoint and configs are pulled from Yavanash/ss-scaffold-v2 into model_cache/ss_scaffold_v2/.

Flag	Meaning
--motif-pdb	preprocessed PDB containing the motif to keep
--motif-range	half-open, 0-indexed range (e.g. 10-30 = residues 11..31). Default: whole input
--motif-class	optional SS override (H/E/C) for every motif residue. Default: per-residue DSSP
--flank-ss	SS class for non-motif residues (H/E/C). Default C (loop-biased flanks)
--total-length	length of the designed protein
--motif-target-start	where to place the motif in the new protein (default: same start index as input)
--n-samples	number of independent backbones to generate
--timesteps, --beta-schedule	must match the trained model (defaults are correct for the HF checkpoint)
--model-dir	optional local checkpoint dir — overrides the HF download
--out	output directory

Outputs in --out:

• sample_000.pdb … sample_{N-1}.pdb — designed backbones (NeRF reconstruction)
• sample_000.csv … — raw 6-angle tensors
• metadata.json — sampling parameters (consumed by the eval pipeline)

Open any sample_*.pdb in PyMOL or ChimeraX. The motif residues should match your input; the rest is generated.

---

Repository layout

foldingdiff_v2/
├── foldingdiff/         upstream code (reused) — base diffusion machinery
├── ss_scaffold/         new work (see ss_scaffold/README.md)
│   ├── sample.py        inference entry point — start here
│   ├── train.py         training entry point
│   ├── model.py         BertForSSConditionedDiffusion
│   ├── sampling.py      RePaint-style motif inpainting
│   ├── losses.py        Ramachandran-region prior
│   ├── ss_labels.py     DSSP wrapper, motif span sampling
│   └── eval/            ProteinMPNN + ESMFold designability pipeline
├── hyperparam_jsons/    ready-made model configs
├── bin/                 ad-hoc helper scripts (not part of the main pipeline)
└── jupyter/             analysis notebooks for report figures

---

Training

The HF checkpoint already provides a trained model, so most users will not need to train from scratch. This section is included for reproducibility and for anyone who wants to retrain on a different dataset, change the SS-conditioning recipe, or fine-tune from the existing checkpoint.

1. Get the data

We use CATH-S40 (the 40 %-sequence-redundant non-redundant subset), the same dataset FoldingDiff was trained on.

1. Download the archive from https://www.cathdb.info/ (typically cath-dataset-nonredundant-S40.pdb.gz).

2. Extract every domain into a single directory of .pdb files:

   mkdir -p data/cath_s40
   tar -xzvf cath-dataset-nonredundant-S40.pdb.gz -C data/cath_s40
   ls data/cath_s40 | wc -l    # expect ~30,000

The same preprocessing rules from the Input PDB preprocessing section apply — files must contain only ATOM records, no headers, no HETATM/TER, single model. CATH-S40 is mostly clean already, but a stray malformed file will crash the dataset cache build.

On the first training run, the dataset class parses every PDB into 6-angle tensors and caches the result as a .pkl next to the source code (filename includes a hash of the dataset config). DSSP labels for the SS-conditioning head are cached separately at <cache-dir>/ss_labels_v1.pkl. Both caches are large (~1–3 GB total) and are gitignored. Subsequent runs load in seconds.

2. Train

python -m ss_scaffold.train \
    --pdb-dir data/cath_s40 \
    --out runs/ss_scaffold_v1 \
    --epochs 100 \
    --batch-size 64 \
    --timesteps 1000 \
    --beta-schedule cosine \
    --rama-lambda 0.1 \
    --motif-mode mixed \
    --p-no-motif 0.3

Most useful flags (full list: python -m ss_scaffold.train --help):

Flag	Default	Meaning
--pdb-dir	required	directory of preprocessed .pdb files
--out	required	run directory; checkpoints go to <out>/checkpoints/, config.json and training_means.npy to <out>/
--epochs	100	training epochs
--batch-size	64	per-device batch size
--pad	128	sequence pad length; longer proteins are trimmed per --trim-strategy
--timesteps	1000	T (number of diffusion steps)
--beta-schedule	cosine	one of cosine / linear / quadratic
--rama-lambda	0.1	weight of the Ramachandran-region prior loss; set to 0 to disable
--motif-mode	mixed	how training motif spans are sampled (see ss_scaffold/ss_labels.py)
--p-no-motif	0.3	probability that a training example has no motif (preserves unconditional behaviour)
--lr	5e-5	base learning rate (LinearWarmup scheduler)
--pretrained-checkpoint	None	path to a .ckpt to fine-tune from
--freeze-pretrained	off	freeze the pretrained backbone and train only the new SS / motif heads

The model is BertForSSConditionedDiffusion — a small subclass of FoldingDiff's BertForDiffusion that adds an nn.Embedding for the 3-state SS label and an extra input channel for the is_motif mask. The training loop is FoldingDiff's, unchanged. Tensorboard logs are written under <out>/lightning_logs/.

3. Hardware and runtime

A full 100-epoch run takes ~12 h on a single A100 (batch 64). On a 24 GB 3090 use --batch-size 32 — same recipe otherwise. If you hit CUDA out of memory, drop --batch-size first, then --pad if necessary.

4. Use your checkpoint for sampling

Pass the resulting run directory to the sampler via --model-dir to use it instead of the HF snapshot:

python -m ss_scaffold.sample \
    --model-dir runs/ss_scaffold_v1 \
    --motif-pdb data/motif.pdb \
    --total-length 80 \
    --n-samples 8 \
    --out generated/run01

Make sure --timesteps and --beta-schedule at sample time match what the model was trained with.

---

Evaluation (optional)

Most users will not need this — sampling alone gives you the designed backbones. The evaluation pipeline is provided for the report's quantitative results and for anyone who wants to reproduce the designability numbers.

It mirrors the Yim et al. designability protocol: ProteinMPNN designs sequences for each backbone, ESMFold predicts a structure for each sequence, then scRMSD / scTM / motif-RMSD / DSSP fidelity are computed. A backbone is designable if min(scRMSD) < 2 Å and paired pLDDT > 70.

External tools required: TMalign on $PATH, ProteinMPNN checkout, ESMFold weights (downloaded by transformers on first use, ~7 GB).

python -m ss_scaffold.eval.run \
    --generated-dir generated/run01 \
    --metadata generated/run01/metadata.json \
    --out eval_out/run01 \
    --n-seqs 8 \
    --mpnn-script /path/to/ProteinMPNN/protein_mpnn_run.py

Outputs: per_sequence.csv, per_backbone.csv, summary.json, seqs/, folded/. Pass --skip-diversity if TMalign isn't available.

---

Troubleshooting

• DSSP failed (...) / biotite parse error — the input PDB wasn't preprocessed. Re-run grep '^ATOM' raw.pdb > motif.pdb and ensure only one model remains.
• CUDA out of memory during training — drop --batch-size. With --pad 128 --batch-size 32 the model fits in 16 GB.
• Motif RMSD is huge — check --motif-target-start matches where you actually want the motif placed.
• NeRF reconstruction failed for sample N — sampled angles produced an invalid geometry (rare). The CSV is still saved; rerun that sample.
• TMalign not found in eval — install from https://zhanggroup.org/TM-align/, or pass --skip-diversity.

---

Citations

@article{wu2022foldingdiff,
  title  = {Protein structure generation via folding diffusion},
  author = {Wu, Kevin E. and Yang, Kevin K. and van den Berg, Rianne and Zou, James Y. and Lu, Alex X. and Amini, Ava P.},
  year   = {2022},
  journal= {arXiv:2209.15611}
}

@inproceedings{lugmayr2022repaint,
  title    = {RePaint: Inpainting using Denoising Diffusion Probabilistic Models},
  author   = {Lugmayr, Andreas and Danelljan, Martin and Romero, Andres and Yu, Fisher and Timofte, Radu and Van Gool, Luc},
  booktitle= {CVPR},
  year     = {2022}
}

@misc{ss_scaffold_bt305,
  title = {SS-Scaffold: secondary-structure-conditioned motif scaffolding on FoldingDiff},
  author= {Singh Srayash, Vadali Yavanash},
  year  = {2026},
  note  = {BT305 course project (30 marks)}
}

---

License

Original FoldingDiff code is © Microsoft, MIT-licensed. New code under ss_scaffold/ is released under the same MIT license for academic use.