Update for ZeroGPU

model/model.py

@@ -27,7 +27,6 @@ from model.barlow_twins import BarlowTwins
 from utils.sequence import uniprot2sequence, encode_sequences
 
 
-@spaces.GPU
 class DTIModel:
     def __init__(self, bt_model_path: str, gbm_model_path: str, encoder: str = "prost_t5"):
         self.bt_model = BarlowTwins()
@@ -68,7 +67,8 @@ class DTIModel:
     def _encode_smiles_mult(self, smiles: List[str], radius: int = 2, bits: int = 1024, features: bool = False):
         morgan = [self._encode_smiles(s, radius, bits, features) for s in smiles]
         return np.array(morgan)
-
+    
+    @spaces.GPU
     def _encode_sequence(self, sequence: str):
         # Clear torch cache
         torch.cuda.empty_cache()
@@ -88,10 +88,12 @@ class DTIModel:
                 print(e)
                 return None
 
+    @spaces.GPU
     def _encode_sequence_mult(self, sequences: List[str]):
         seq = [self._encode_sequence(sequence) for sequence in sequences]
         return np.array(seq)
 
+    @spaces.GPU
     def __predict_pair(self, drug_emb: np.ndarray, target_emb: np.ndarray, pred_leaf: bool):
         if drug_emb.shape[0] < target_emb.shape[0]:
             drug_emb = np.tile(drug_emb, (len(target_emb), 1))