Update app.py

app.py

@@ -7,7 +7,7 @@ from lavis.models.protein_models.protein_function_opt import Blip2ProteinMistral
 from lavis.models.base_model import FAPMConfig
 import spaces
 import gradio as gr
-
+from esm_scripts.extract import run_demo
 
 # Load the model
 model = Blip2ProteinMistral(config=FAPMConfig(), esm_size='3b')
@@ -18,16 +18,14 @@ model.to('cuda')
 @spaces.GPU
 def generate_caption(protein, prompt):
     # Process the image and the prompt
-    print(f"system path: {os.getcwd()}")
-    with open('/home/user/app/example.fasta', 'w') as f:
-        f.write('>{}\n'.format("protein_name"))
-        f.write('{}\n'.format(protein.strip()))
-    print(f"fasta prepared")
-    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")
-    print(f"protein pt file prepared")
-    esm_emb = torch.load("/home/user/app/protein_name.pt")['representations'][36]
+    # with open('/home/user/app/example.fasta', 'w') as f:
+    #     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', 
+                       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': ['test_protein'],
+    samples = {'name': ['protein_name'],
                'image': torch.unsqueeze(esm_emb, dim=0),
                'text_input': ['none'],
                'prompt': [prompt]}