revert esm structure prediction

app.py

@@ -242,21 +242,22 @@ def display_protein():
     sequence = st.text_input('Enter the amino-acid sequence of the query protein target', value='HXHVWPVQDAKARFSEFLDACITEGPQIVSRRGAEEAVLVPIGEWRRLQAAA', placeholder='HXHVWPVQDAKARFSEFLDACITEGPQIVSRRGAEEAVLVPIGEWRRLQAAA')
     
     if sequence:
-        #model = esm.pretrained.esmfold_v1()
-        #model = model.eval().cuda()
-        
-        #with torch.no_grad():
-        #    output = model.infer_pdb(sequence)
-        
-        #with open("result.pdb", "w") as f:
-        #    f.write(output)
+        model = esm.pretrained.esmfold_v1()
+        model = model.eval().cuda()
         
+        with torch.no_grad():
+            output = model.infer_pdb(sequence)
+            st.write(output)
         
-        #struct = bsio.load_structure("result.pdb", extra_fields=["b_factor"])
-        #print(struct.b_factor.mean())
+        with open("result.pdb", "w") as f:
+            f.write(output)
         
-        #st.write(output)
         
+        struct = bsio.load_structure("result.pdb", extra_fields=["b_factor"])
+        print(struct.b_factor.mean())
+
+
+    """
         model, alphabet = esm.pretrained.esm2_t33_650M_UR50D()
         batch_converter = alphabet.get_batch_converter()
         batch_labels, batch_strs, batch_tokens = batch_converter([("protein1", sequence),])
@@ -311,6 +312,7 @@ def display_protein():
             st.session_state['xq'] = st.session_state.model
     
     # example proteins ["HXHVWPVQDAKARFSEFLDACITEGPQIVSRRGAEEAVLVPIGEWRRLQAAA"], ["AHKLFIGGLPNYLNDDQVKELLTSFGPLKAFNLVKDSATGLSKGYAFCEYVDINVTDQAIAGLNGMQLGDKKLLVQRASVGAKNA"]
+    """
 
 def display_context():
     st.markdown('## Display context')