try esm fold

app.py

@@ -191,14 +191,20 @@ def display_protein():
     sequence = st.text_input('Enter the amino-acid sequence of the query protein target', value='HXHVWPVQDAKARFSEFLDACITEGPQIVSRRGAEEAVLVPIGEWRRLQAAA', placeholder='HXHVWPVQDAKARFSEFLDACITEGPQIVSRRGAEEAVLVPIGEWRRLQAAA')
     
     if sequence:
-        model, alphabet = esm.pretrained.esm1_t6_43M_UR50S()
-
-        batch_converter = alphabet.get_batch_converter()
-        inputs = batch_converter([sequence])
-        embeddings = model.embed(inputs)
-        coordinates = embeddings["representations"][0].detach().numpy()
+        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)
+        
+        
+        #struct = bsio.load_structure("result.pdb", extra_fields=["b_factor"])
+        #print(struct.b_factor.mean())
         
-        st.write(coordinates)
+        st.write(output)
     
     """
     sequence = st.text_input("Enter the amino-acid sequence of the query protein target", value="HXHVWPVQDAKARFSEFLDACITEGPQIVSRRGAEEAVLVPIGEWRRLQAAA", placeholder="HXHVWPVQDAKARFSEFLDACITEGPQIVSRRGAEEAVLVPIGEWRRLQAAA")