Update app.py

app.py

@@ -10,14 +10,10 @@ import gradio as gr
 from esm_scripts.extract import run_demo
 from esm import pretrained, FastaBatchedDataset
 
-
-model_esm, alphabet = pretrained.load_model_and_alphabet('esm2_t36_3B_UR50D')
-model_esm.to('cuda')
-
 # Load the model
-# model = Blip2ProteinMistral(config=FAPMConfig(), esm_size='3b')
-# model.load_checkpoint("model/checkpoint_mf2.pth")
-# model.to('cuda')
+model = Blip2ProteinMistral(config=FAPMConfig(), esm_size='3b')
+model.load_checkpoint("model/checkpoint_mf2.pth")
+model.to('cuda')
 
 
 @spaces.GPU
@@ -46,9 +42,13 @@ def generate_caption(protein, prompt):
     )
     print(f"Read sequences")
     return_contacts = "contacts" in include
+
+    model_esm, alphabet = pretrained.load_model_and_alphabet('esm2_t36_3B_UR50D')
+    model_esm.to('cuda')
+    model_esm.eval()
     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]
-    model_esm.eval()
+
     with torch.no_grad():
         for batch_idx, (labels, strs, toks) in enumerate(data_loader):
             print(
@@ -57,6 +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
             logits = out["logits"].to(device="cpu")
             representations = {
                 layer: t.to(device="cpu") for layer, t in out["representations"].items()
@@ -94,10 +95,10 @@ def generate_caption(protein, prompt):
                'text_input': ['none'],
                'prompt': [prompt]}
     # Generate the output
-    # prediction = model.generate(samples, length_penalty=0., num_beams=15, num_captions=10, temperature=1., repetition_penalty=1.0)
+    prediction = model.generate(samples, length_penalty=0., num_beams=15, num_captions=10, temperature=1., repetition_penalty=1.0)
 
-    # return prediction
-    return "test"
+    return prediction
+    # return "test"
 
 # Define the FAPM interface
 description = """Quick demonstration of the FAPM model for protein function prediction. Upload an protein sequence to generate a function description. Modify the Prompt to provide the taxonomy information.