Update app.py

app.py

@@ -1,13 +1,20 @@
 import gradio as gr
-from evodiff.pretrained import OA_DM_38M
-from evodiff.generate import generate_oaardm
+from evodiff.pretrained import OA_DM_38M, D3PM_UNIFORM_38M
+from evodiff.generate import generate_oaardm, generate_d3pm
 
-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')
+def make_seq(seq_len, model_type):
+    if model_type == "EvoDiff-Seq-OADM 38M":
+        checkpoint = OA_DM_38M()
+        model, collater, tokenizer, scheme = checkpoint
+        tokeinzed_sample, generated_sequence = generate_oaardm(model, tokenizer, seq_len, batch_size=1, device='cpu')
+    
+    if model_type == "EvoDiff-D3PM-Uniform 38M":
+        checkpoint = D3PM_UNIFORM_38M(return_all=True)
+        model, collater, tokenizer, scheme, timestep, Q_bar, Q = checkpoint
+        tokeinzed_sample, generated_sequence = generate_d3pm(model, tokenizer, Q, Q_bar, timestep, seq_len, batch_size=1, device='cpu')
+
     return generated_sequence
 
 # iface = gr.Interface(
@@ -24,7 +31,10 @@ with gr.Blocks() as edapp:
         gr.Markdown(
             """
             # EvoDiff
-            Created By: Microsoft Research (Sarah Alamdari,  Nitya Thakkar,  Rianne van den Berg,  Alex X. Lu, Nicolo Fusi,  ProfileAva P. Amini,  and Kevin K. Yang)
+            ## Generation of protein sequences and evolutionary alignments via discrete diffusion models
+
+            Created By: Microsoft Research [Sarah Alamdari, Nitya Thakkar, Rianne van den Berg, Alex X. Lu, Nicolo Fusi, ProfileAva P. Amini, and Kevin K. Yang]
+            
             Spaces App By: Colby T. Ford
             """
         )
@@ -36,7 +46,11 @@ with gr.Blocks() as edapp:
             """)
         gr.Interface(
             fn=make_seq,
-            inputs=gr.Slider(10, 100, label = "Sequence Length"),
+            inputs=[
+                gr.Slider(10, 100, label = "Sequence Length")
+                gr.Dropdown(["EvoDiff-Seq-OADM 38M", "EvoDiff-D3PM-Uniform 38M"], type="value"),
+                gr.()
+                ],
             outputs="text"
             )