Update app.py

app.py

@@ -65,7 +65,7 @@ def display_pdb(path_to_pdb):
                             allowpaymentrequest="" frameborder="0" srcdoc='{x}'></iframe>"""
 '''
 
-def make_uncond_seq(seq_len, model_type):
+def make_uncond_seq(seq_len, model_type, pred_structure):
     if model_type == "EvoDiff-Seq-OADM 38M":
         checkpoint = OA_DM_38M()
         model, collater, tokenizer, scheme = checkpoint
@@ -76,21 +76,27 @@ def make_uncond_seq(seq_len, model_type):
         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')
 
-    path_to_pdb = predict_protein(generated_sequence)
-    molhtml = display_pdb(path_to_pdb)
+    if pred_structure:
+        path_to_pdb = predict_protein(generated_sequence)
+        molhtml = display_pdb(path_to_pdb)
 
-    return generated_sequence, molhtml
+        return generated_sequence, molhtml
+    else:
+        return generated_sequence
 
-def make_cond_seq(seq_len, msa_file, model_type):
+def make_cond_seq(seq_len, msa_file, model_type, pred_structure):
     if model_type == "EvoDiff-MSA":
         checkpoint = MSA_OA_DM_MAXSUB()
         model, collater, tokenizer, scheme = checkpoint
         tokeinzed_sample, generated_sequence  = generate_query_oadm_msa_simple(msa_file.name, model, tokenizer, n_sequences=64, seq_length=seq_len, device='cpu', selection_type='random')
 
-    path_to_pdb = predict_protein(generated_sequence)
-    molhtml = display_pdb(path_to_pdb)
+    if pred_structure:
+        path_to_pdb = predict_protein(generated_sequence)
+        molhtml = display_pdb(path_to_pdb)
 
-    return generated_sequence, molhtml
+        return generated_sequence, molhtml
+    else:
+        return generated_sequence
 
 usg_app = gr.Interface(
             fn=make_uncond_seq,
@@ -111,7 +117,8 @@ csg_app = gr.Interface(
             inputs=[
                 gr.Slider(10, 100, label = "Sequence Length"),
                 gr.File(file_types=["a3m"], label = "MSA File"),
-                gr.Dropdown(["EvoDiff-MSA"], value="EvoDiff-MSA", type="value", label = "Model")
+                gr.Dropdown(["EvoDiff-MSA"], value="EvoDiff-MSA", type="value", label = "Model"),
+                gr.Checkbox(value=False, label = "Predict Structure?")
                 ],
             outputs=[
                 "text",