Update app.py

app.py

@@ -15,26 +15,30 @@ from colabfold.plot import plot_msa_v2
 def a3m_file(file):
     return "tmp.a3m"
 
-def make_uncond_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
 
-def make_cond_seq(seq_len, msa_file, model_type):
-    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')
+def predict_protein(sequence):
+    download_alphafold_params("alphafold2_ptm", Path("."))
+    results = run(
+        queries=[('evodiff_protein',sequence, None)],
+        result_dir='evodiff_protein',
+        use_templates=False,
+        num_relax=0,
+        msa_mode="mmseqs2_uniref_env",
+        model_type="alphafold2_ptm",
+        num_models=1,
+        num_recycles=1,
+        model_order=[1],
+        is_complex=False,
+        data_dir=Path("."),
+        keep_existing_results=False,
+        rank_by="auto",
+        stop_at_score=float(100),
+        zip_results=False,
+        user_agent="colabfold/google-colab-main",
+    )
 
-    return generated_sequence
+    return f"evodiff_protein/evodiff_protein_unrelaxed_rank_001_alphafold2_ptm_model_1_seed_000.pdb"
 
 def display_pdb(path_to_pdb):
     '''
@@ -65,40 +69,32 @@ def display_pdb(path_to_pdb):
                             allowpaymentrequest="" frameborder="0" srcdoc='{x}'></iframe>"""
 '''
 
-def predict_protein(sequence):
-    model_type = "alphafold2_ptm"
-    download_alphafold_params(model_type, Path("."))
-    results = run(
-        queries=queries,
-        result_dir=result_dir,
-        use_templates=use_templates,
-        custom_template_path=custom_template_path,
-        num_relax=0,
-        msa_mode=msa_mode,
-        model_type=model_type,
-        num_models=1,
-        num_recycles=1,
-        recycle_early_stop_tolerance=recycle_early_stop_tolerance,
-        num_seeds=num_seeds,
-        use_dropout=use_dropout,
-        model_order=[1],
-        is_complex=False,
-        data_dir=Path("."),
-        keep_existing_results=False,
-        rank_by="auto",
-        pair_mode=pair_mode,
-        pairing_strategy=pairing_strategy,
-        stop_at_score=float(100),
-        prediction_callback=prediction_callback,
-        dpi=dpi,
-        zip_results=False,
-        save_all=save_all,
-        max_msa=max_msa,
-        use_cluster_profile=use_cluster_profile,
-        input_features_callback=input_features_callback,
-        save_recycles=save_recycles,
-        user_agent="colabfold/google-colab-main",
-    )
+def make_uncond_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')
+
+    path_to_pdb = predict_protein(generated_sequence)
+    molhtml = display_pdb(path_to_pdb)
+
+    return generated_sequence, molhtml
+
+def make_cond_seq(seq_len, msa_file, model_type):
+    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)
+
+    return generated_sequence, molhtml
 
 usg_app = gr.Interface(
             fn=make_uncond_seq,
@@ -106,7 +102,10 @@ usg_app = gr.Interface(
                 gr.Slider(10, 100, label = "Sequence Length"),
                 gr.Dropdown(["EvoDiff-Seq-OADM 38M", "EvoDiff-D3PM-Uniform 38M"], type="value", label = "Model")
                 ],
-            outputs="text",
+            outputs=[
+                "text",
+                gr.HTML()
+            ],
             title = "Unconditional sequence generation",
             description="Generate a sequence with `EvoDiff-Seq-OADM 38M` (smaller/faster) or `EvoDiff-D3PM-Uniform 38M` (larger/slower) models."
             )
@@ -118,7 +117,10 @@ csg_app = gr.Interface(
                 gr.File(file_types=["a3m"], label = "MSA File"),
                 gr.Dropdown(["EvoDiff-MSA"], type="value", label = "Model")
                 ],
-            outputs="text",
+            outputs=[
+                "text",
+                gr.HTML()
+            ],
             # examples=[["https://github.com/microsoft/evodiff/raw/main/examples/example_files/bfd_uniclust_hits.a3m"]], 
             title = "Conditional sequence generation",
             description="Evolutionary guided sequence generation with the `EvoDiff-MSA` model."