move model run

app.py

@@ -139,7 +139,14 @@ def handle_name(name=None, pdb_input=None, model_version="ESM3"):
         pdb_name = str(random.randint(0, 100000))
     return f'{pdb_name}-Dyna1{"" if model_version == "ESM3" else "-ESM2"}'
 
-@spaces.GPU(duration=1200)
+@spaces.GPU(duration=300)
+def run_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:
+        logits = model(seq_input, sequence_id)
+    return logits
+
 def predict_dynamics(sequence=None, pdb_input=None, chain_id='A', use_pdb_seq=False, model_version="ESM3", name=None, oauth_token: Optional[str] = None):
     try:
         # Validate ESM2 requires sequence
@@ -151,7 +158,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').to(DEVICE)
             model.load_state_dict(torch.load('Dyna-1/model/weights/dyna1.pt', map_location=DEVICE), strict=False)
@@ -187,11 +193,9 @@ 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 ''))
-            
-        if model_version == "ESM3":
-            logits = model((seq_input, struct_input), sequence_id)
-        else:
-            logits = model(seq_input, sequence_id)  
+        
+        logits = run_model(model_version, seq_input, struct_input, sequence_id)
+        
         probabilities = utils.prob_adjusted(logits).cpu().detach().numpy()
         
         seq_to_use = sequence if sequence else pdb_seq if pdb_input else sequence