Update app.py

app.py

@@ -3,6 +3,9 @@ from transformers import AutoTokenizer, EsmForProteinFolding
 from transformers.models.esm.openfold_utils.protein import to_pdb, Protein as OFProtein
 from transformers.models.esm.openfold_utils.feats import atom14_to_atom37
 import torch
+from logging import getLogger
+
+logger = getLogger(__name__)
 
 def convert_outputs_to_pdb(outputs):
     final_atom_positions = atom14_to_atom37(outputs["positions"][-1], outputs)
@@ -27,7 +30,7 @@ def convert_outputs_to_pdb(outputs):
     return pdbs
 
 def fold_prot_locally(sequence):
-    print("Folding:", sequence)
+    logger.info("Folding: " + sequence)
     tokenized_input = tokenizer([sequence], return_tensors="pt", add_special_tokens=False)['input_ids'].cuda()
 
     with torch.no_grad():
@@ -36,7 +39,9 @@ def fold_prot_locally(sequence):
     return pdb
 
 tokenizer = AutoTokenizer.from_pretrained("facebook/esmfold_v1")
-model = EsmForProteinFolding.from_pretrained("facebook/esmfold_v1", low_cpu_mem_usage=True).cuda().half()
+model = EsmForProteinFolding.from_pretrained("facebook/esmfold_v1", low_cpu_mem_usage=True).cuda()
+model.esm = model.esm.half()
+torch.backends.cuda.matmul.allow_tf32 = True
 
 iface = gr.Interface(fn=fold_prot_locally, inputs="text", outputs="text")
 iface.launch()