Update app.py

app.py

@@ -140,7 +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):
+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:
@@ -159,11 +159,11 @@ 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').to(DEVICE)
-            model.load_state_dict(torch.load('Dyna-1/model/weights/dyna1.pt', map_location=DEVICE), strict=False)
+            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', map_location=DEVICE), strict=False)
+            model.load_state_dict(torch.load('Dyna-1/model/weights/dyna1-esm2.pt'), strict=False)
 
         model.eval()
         
@@ -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_version, seq_input, struct_input, sequence_id)
+        logits = run_model(model, model_version, seq_input, struct_input, sequence_id)
         
         probabilities = utils.prob_adjusted(logits).cpu().detach().numpy()