Add Zero GPU support for H200 GPU on Hugging Face Spaces

- Import spaces module for Zero GPU support
- Add @spaces.GPU decorator to inference function with 5-minute duration
- Fix device selection to properly use GPU when available
- Use model.to(device) instead of model.cuda() for consistency
- Print GPU name when running on GPU
- Increase batch size to 50 for GPU inference
- Add spaces>=0.19.0 to requirements.txt

app.py

@@ -17,6 +17,7 @@ import zipfile
 import shutil
 import uuid
 import gc
+import spaces
 
 # Add current directory to path
 sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
@@ -224,6 +225,7 @@ def get_mutated_protein(sequence,mutant):
   mutated_sequence[int(mutant[1:-1])-1]=mutant[-1]
   return ''.join(mutated_sequence)
 
+@spaces.GPU(duration=300)  # Request GPU for up to 5 minutes
 def score_and_create_matrix_all_singles(sequence,mutation_range_start=None,mutation_range_end=None,model_type="Large",scoring_mirror=False,batch_size_inference=20,max_number_positions_per_heatmap=50,num_workers=0,AA_vocab=AA_vocab):
   # Clean up old files periodically
   cleanup_old_files()
@@ -259,14 +261,16 @@ def score_and_create_matrix_all_singles(sequence,mutation_range_start=None,mutat
     model_path = download_model_from_hf("Tranception_Medium")
     model = tranception.model_pytorch.TranceptionLMHeadModel.from_pretrained(pretrained_model_name_or_path=model_path)
   
-  # Device selection - for HF Spaces, typically CPU
+  # Device selection - Zero GPU will provide CUDA when decorated with @spaces.GPU
   if torch.cuda.is_available():
     device = torch.device("cuda")
-    model.cuda()
-    print("Inference will take place on NVIDIA GPU")
+    model = model.to(device)
+    print(f"Inference will take place on {torch.cuda.get_device_name(0)}")
+    # Increase batch size for GPU inference
+    batch_size_inference = min(batch_size_inference, 50)
   else:
     device = torch.device("cpu")
-    model.to(device)
+    model = model.to(device)
     print("Inference will take place on CPU")
     # Reduce batch size for CPU inference
     batch_size_inference = min(batch_size_inference, 10)

requirements.txt

@@ -1,6 +1,7 @@
 torch>=1.12.0
 transformers==4.17.0
 tokenizers==0.11.6
+spaces>=0.19.0
 numpy>=1.21.0,<1.24.0
 pandas>=1.3.0,<2.0.0
 scipy>=1.7.0,<1.11.0