fixed app for zero gpu

app.py

@@ -10,13 +10,14 @@ from protobind_diff.esm_inference import get_esm_embedding
 from protobind_diff.model import ModelGenerator
 from protobind_diff.data_loader import InferenceDataset
 from huggingface_hub import hf_hub_download
+import spaces
 
 # Hugging Face Hub details
 REPO_ID = "ai-gero/ProtoBind-Diff"
 MODEL_FILENAME = "model.ckpt"
 TOKENIZER_FILENAME = "tokenizer_smiles_diffusion.json"
 
-
+@spaces.GPU(duration=120) 
 def generate_smiles_for_sequence(protein_sequence: str, num_samples: int):
     """
     The main prediction function that runs the full pipeline.
@@ -27,7 +28,9 @@ def generate_smiles_for_sequence(protein_sequence: str, num_samples: int):
     if len(protein_sequence) < 10:
         raise gr.Error("Protein sequence is too short.")
 
-    print(">> inference started, attempts:", num_samples, flush=True)
+    device = "cuda" 
+    esm_model.to(device)
+    protobind_model.to(device) 
 
     with torch.no_grad():
         batch_converter = alphabet.get_batch_converter()
@@ -58,7 +61,7 @@ def generate_smiles_for_sequence(protein_sequence: str, num_samples: int):
 # --- GRADIO APP DEFINITION ---
 
 # Load models on app startup
-device = "cuda" if torch.cuda.is_available() else "cpu"
+device = "cpu"
 esm_model, alphabet = esm.pretrained.load_model_and_alphabet('esm2_t33_650M_UR50D')
 esm_model.eval()
 esm_model = esm_model.to(device)