Update app.py

app.py

@@ -141,6 +141,15 @@ def handle_name(name=None, pdb_input=None, model_version="ESM3"):
 
 @spaces.GPU(duration=300)
 def run_model(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()    
+    
     if model_version == "ESM3":
         logits = model((seq_input, struct_input), sequence_id)
     else:
@@ -158,15 +167,6 @@ def predict_dynamics(sequence=None, pdb_input=None, chain_id='A', use_pdb_seq=Fa
             
         base_name = handle_name(name, pdb_input, model_version)
         
-        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()
-        
         seq_input, struct_input = None, None
         sequence = validate_sequence(sequence) if sequence else None
         protein = None
@@ -194,7 +194,7 @@ def predict_dynamics(sequence=None, pdb_input=None, chain_id='A', use_pdb_seq=Fa
         if not (sequence or (pdb_input and model_version == "ESM3")):
             raise ValueError('Please provide a sequence' + (' or structure input' if model_version == "ESM3" else ''))
         
-        logits = run_model(model, model_version, seq_input, struct_input, sequence_id)
+        logits = run_model(model_version, seq_input, struct_input, sequence_id)
         
         probabilities = utils.prob_adjusted(logits).cpu().detach().numpy()