app.py

import os
import gradio as gr

import torch
from foldingdiff import sampling
from foldingdiff import angles_and_coords as ac

def sample_at_length(l:int, seed:int):
 """
 Sample a single structure at the given length
 """
 torch.manual_seed(seed)
 l = int(l)
 s = sampling.sample_simple("wukevin/foldingdiff_cath", n=1, sweep_lengths=(l, l+1))[0]
 # Create a PDB file
 outdir = os.path.join(os.getcwd(), "output")
 os.makedirs(outdir, exist_ok=True)
 pdb_file = ac.create_new_chain_nerf(os.path.join(outdir, "generated.pdb"), s)
 return s, pdb_file

interface = gr.Interface(
 fn=sample_at_length,
 inputs=[
 gr.Number(value=55, label="Protein backbone length to generate", show_label=True, precision=0),
 gr.Number(value=8964, label="Random seed", show_label=True, precision=0),
 ],
 outputs=[
 gr.Dataframe(label="Generated angles defining structure", max_rows=8),
 gr.File(label="Generated structure in PDB format (cartesian coordinates)"),
 ],
)
interface.launch()