import gradio as gr

import torch
from foldingdiff import sampling

def sample_at_length(l:int, seed:int):
    torch.manual_seed(seed)
    l = int(l)
    s = sampling.sample_simple("wukevin/foldingdiff_cath", n=1, sweep_lengths=(l, l+1))
    return s[0]

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(),
)
interface.launch()