Update to ZeroGPU

app.py

@@ -128,4 +128,4 @@ iface = gr.Interface(
     theme=theme
 )
 
-iface.launch()
+iface.launch(share=True)

model/barlow_twins.py

@@ -9,10 +9,12 @@ import os
 import pickle
 import inspect
 from tqdm.auto import trange
+import spaces
 
 from model.base_model import BaseModel
 
 
+@spaces.GPU
 class BarlowTwins(BaseModel):
     def __init__(
             self,

model/base_model.py

@@ -2,8 +2,10 @@ from typing import Tuple, Any, Union
 import torch
 from torch import nn
 import numpy as np
+import spaces
 
 
+@spaces.GPU
 class BaseModel(nn.Module):
     def __init__(self):
         super(BaseModel, self).__init__()

model/model.py

@@ -17,6 +17,7 @@ import torch
 from typing import *
 from rdkit import RDLogger
 RDLogger.DisableLog("rdApp.*")
+import spaces
 
 from xgboost import XGBClassifier, DMatrix
 
@@ -26,7 +27,7 @@ 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()

utils/sequence.py

@@ -8,6 +8,7 @@ import concurrent.futures
 from tqdm.auto import tqdm
 import multiprocessing
 from multiprocessing import Pool
+import spaces
 
 
 ENCODERS = {
@@ -49,6 +50,7 @@ def uniprot2sequence(uniprot_id):
         return None
 
 
+@spaces.GPU
 def encode_sequences(sequences: list, encoder: str):
     if encoder not in ENCODERS.keys():
         raise ValueError(f"Invalid encoder: {encoder}")