Update app.py

app.py

@@ -14,13 +14,13 @@ from esm import pretrained, FastaBatchedDataset
 
 
 # Load the model
-# model = Blip2ProteinMistral(config=FAPMConfig(), esm_size='3b')
-# model.load_checkpoint("model/checkpoint_mf2.pth")
+model = Blip2ProteinMistral(config=FAPMConfig(), esm_size='3b')
+model.load_checkpoint("model/checkpoint_mf2.pth")
 # model.to('cuda')
 
-# model_esm, alphabet = pretrained.load_model_and_alphabet('esm2_t36_3B_UR50D')
+model_esm, alphabet = pretrained.load_model_and_alphabet('esm2_t36_3B_UR50D')
 # model_esm.to('cuda')
-# model_esm.eval()
+model_esm.eval()
 
 
 # tokenizer = EsmTokenizer.from_pretrained("facebook/esm2_t36_3B_UR50D")
@@ -51,9 +51,7 @@ def generate_caption(protein, prompt):
     batches = dataset.get_batch_indices(toks_per_batch, extra_toks_per_seq=1)
     print("batches prepared")
 
-    model_esm, alphabet = pretrained.load_model_and_alphabet('esm2_t36_3B_UR50D')
     model_esm.to('cuda')
-    model_esm.eval()
 
     data_loader = torch.utils.data.DataLoader(
         dataset, collate_fn=alphabet.get_batch_converter(truncation_seq_length), batch_sampler=batches
@@ -115,9 +113,7 @@ def generate_caption(protein, prompt):
                'prompt': [prompt]}
 
     del model_esm
-
-    model = Blip2ProteinMistral(config=FAPMConfig(), esm_size='3b')
-    model.load_checkpoint("model/checkpoint_mf2.pth")
+    
     model.to('cuda')
     # Generate the output
     prediction = model.generate(samples, length_penalty=0., num_beams=15, num_captions=10, temperature=1.,