import gradio as gr
from evodiff.pretrained import OA_DM_38M
from evodiff.generate import generate_oaardm

checkpoint = OA_DM_38M()
model, collater, tokenizer, scheme = checkpoint


def make_seq(seq_len):
    tokeinzed_sample, generated_sequence = generate_oaardm(model, tokenizer, seq_len, batch_size=1, device='cpu')
    return generated_sequence

# iface = gr.Interface(
#     fn=make_seq,
#     inputs=gr.Slider(10, 100),
#     outputs="text"
#     )

# iface.launch()


with gr.Blocks() as edapp:
    gr.Markdown(
        "# EvoDiff"
    )
    with gr.Row():
        gr.Markdown(
            """
            ## Unconditional sequence generation
            Generate a sequence with EvoDiff-Seq-OADM 38M
            """)
        gr.Interface(
            fn=make_seq,
            inputs=gr.Slider(10, 100),
            outputs="text"
            )

if __name__ == "__main__":
    edapp.launch()