Update app.py

app.py

@@ -140,16 +140,7 @@ def handle_name(name=None, pdb_input=None, model_version="ESM3"):
     return f'{pdb_name}-Dyna1{"" if model_version == "ESM3" else "-ESM2"}'
 
 @spaces.GPU(duration=300)
-def run_model(model_version='ESM2', seq_input=None, struct_input=None, sequence_id=None):
-    if model_version == "ESM3":
-        model = ESM_model(method='esm3')
-        model.load_state_dict(torch.load('Dyna-1/model/weights/dyna1.pt'), strict=False)
-    else:  
-        model = ESM_model(method='esm2', nheads=8, nlayers=12, layer=30).to(DEVICE)
-        model.load_state_dict(torch.load('Dyna-1/model/weights/dyna1-esm2.pt'), strict=False)
-
-    model.eval()    
-    
+def run_model(model, model_version='ESM2', seq_input=None, struct_input=None, sequence_id=None):
     if model_version == "ESM3":
         logits = model((seq_input, struct_input), sequence_id)
     else:
@@ -170,6 +161,14 @@ def predict_dynamics(sequence=None, pdb_input=None, chain_id='A', use_pdb_seq=Fa
         seq_input, struct_input = None, None
         sequence = validate_sequence(sequence) if sequence else None
         protein = None
+        if model_version == "ESM3":
+            model = ESM_model(method='esm3')
+            model.load_state_dict(torch.load('Dyna-1/model/weights/dyna1.pt'), strict=False)
+        else:  
+            model = ESM_model(method='esm2', nheads=8, nlayers=12, layer=30).to(DEVICE)
+            model.load_state_dict(torch.load('Dyna-1/model/weights/dyna1-esm2.pt'), strict=False)
+
+        model.eval()
         
         if pdb_input and model_version == "ESM3":
             protein, protein_chain = process_structure(pdb_input, chain_id)