Update app.py

app.py

@@ -15,6 +15,9 @@ 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.to('cuda')
+# model_esm.eval()
 
 @spaces.GPU
 def generate_caption(protein, prompt):
@@ -37,19 +40,15 @@ def generate_caption(protein, prompt):
     print("dataset prepared")
     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
+        dataset, collate_fn=model.alphabet.get_batch_converter(truncation_seq_length), batch_sampler=batches
     )
     print(f"Read sequences")
     return_contacts = "contacts" in include
 
-    assert all(-(model_esm.num_layers + 1) <= i <= model_esm.num_layers for i in repr_layers)
-    repr_layers = [(i + model_esm.num_layers + 1) % (model_esm.num_layers + 1) for i in repr_layers]
+    assert all(-(model.model_esm.num_layers + 1) <= i <= model.model_esm.num_layers for i in repr_layers)
+    repr_layers = [(i + model.model_esm.num_layers + 1) % (model.model_esm.num_layers + 1) for i in repr_layers]
 
     with torch.no_grad():
         for batch_idx, (labels, strs, toks) in enumerate(data_loader):
@@ -58,8 +57,7 @@ def generate_caption(protein, prompt):
             )
             if torch.cuda.is_available():
                 toks = toks.to(device="cuda", non_blocking=True)
-            out = model_esm(toks, repr_layers=repr_layers, return_contacts=return_contacts)
-            del model_esm
+            out = model.model_esm(toks, repr_layers=repr_layers, return_contacts=return_contacts)
             logits = out["logits"].to(device="cpu")
             representations = {
                 layer: t.to(device="cpu") for layer, t in out["representations"].items()