import os
import gradio as gr
import torch
from foldingdiff import sampling
from foldingdiff import angles_and_coords as ac
def read_mol(molpath: str) -> str:
with open(molpath, "r") as fp:
lines = fp.readlines()
mol = ""
for l in lines:
mol += l
return mol
def molecule(input_pdb: str) -> str:
"""Get the string to view the given pdb in 3dmol.js"""
mol = read_mol(input_pdb)
x = (
"""
"""
)
return f""""""
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 molecule(s), pdb_file, s
interface = gr.Interface(
fn=sample_at_length,
inputs=[
gr.Number(value=80, label="Protein backbone length to generate", show_label=True, precision=0),
gr.Number(value=42, label="Random seed", show_label=True, precision=0),
],
outputs=[
gr.HTML(),
gr.File(label="Generated structure in PDB format (cartesian coordinates)"),
gr.Dataframe(label="Generated angles defining structure", max_rows=8),
],
)
interface.launch()