Update app.py

app.py

@@ -8,6 +8,12 @@ from lavis.models.base_model import FAPMConfig
 import spaces
 import gradio as gr
 from esm_scripts.extract import run_demo
+from esm import pretrained
+
+
+model_esm, alphabet = pretrained.load_model_and_alphabet('esm2_t36_3B_UR50D')
+model_esm.eval()
+model_esm = model.cuda()
 
 # Load the model
 model = Blip2ProteinMistral(config=FAPMConfig(), esm_size='3b')
@@ -22,7 +28,8 @@ def generate_caption(protein, prompt):
     #     f.write('>{}\n'.format("protein_name"))
     #     f.write('{}\n'.format(protein.strip()))
     # os.system("python esm_scripts/extract.py esm2_t36_3B_UR50D /home/user/app/example.fasta /home/user/app --repr_layers 36 --truncation_seq_length 1024 --include per_tok")
-    esm_emb = run_demo(protein_name='protein_name', protein_seq=protein, model_location='esm2_t36_3B_UR50D', 
+    esm_emb = run_demo(protein_name='protein_name', protein_seq=protein, 
+                       model=model_esm, alphabet=alphabet, 
                        include='per_tok', repr_layers=36, truncation_seq_length=1024)
     esm_emb = F.pad(esm_emb.t(), (0, 1024 - len(esm_emb))).t().to('cuda')
     samples = {'name': ['protein_name'],