Add resource management and cancellation support

- Add automatic cleanup of old files (30 minutes)
- Free model memory after inference with gc.collect()
- Clear CUDA cache if available
- Configure queue with size limits and status updates
- Limit concurrent threads to 2 for better resource management
- Add proper error handling for file operations
- Prevent API access to avoid external requests

app.py

@@ -16,6 +16,10 @@ from huggingface_hub import hf_hub_download
 import zipfile
 import shutil
 import uuid
+import tempfile
+import atexit
+import threading
+import gc
 
 # Add current directory to path
 sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
@@ -187,6 +191,28 @@ def check_valid_mutant(sequence,mutant,AA_vocab=AA_vocab):
   if to_AA not in AA_vocab: valid=False
   return valid
 
+# Global variable to track active inference threads
+active_inferences = {}
+inference_lock = threading.Lock()
+
+def cleanup_old_files(max_age_minutes=30):
+    """Clean up old inference files"""
+    import glob
+    import time
+    current_time = time.time()
+    patterns = ["fitness_scoring_substitution_matrix_*.png", 
+                "fitness_scoring_substitution_matrix_*.csv",
+                "all_mutations_fitness_scores_*.csv"]
+    
+    for pattern in patterns:
+        for file_path in glob.glob(pattern):
+            try:
+                file_age = current_time - os.path.getmtime(file_path)
+                if file_age > max_age_minutes * 60:
+                    os.remove(file_path)
+            except:
+                pass
+
 def get_mutated_protein(sequence,mutant):
   if not check_valid_mutant(sequence,mutant):
     return "The mutant is not valid"
@@ -195,6 +221,9 @@ def get_mutated_protein(sequence,mutant):
   return ''.join(mutated_sequence)
 
 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()
+  
   # Generate unique ID for this request
   unique_id = str(uuid.uuid4())
   
@@ -275,6 +304,12 @@ def score_and_create_matrix_all_singles(sequence,mutation_range_start=None,mutat
   scores_export = scores_export[['position', 'original_AA', 'target_AA', 'mutant', 'fitness_score', 'mutated_sequence']]
   scores_export.to_csv(comprehensive_csv_path, index=False)
   csv_files.append(comprehensive_csv_path)
+  
+  # Clean up model from memory after inference
+  del model
+  gc.collect()
+  if torch.cuda.is_available():
+    torch.cuda.empty_cache()
     
   return score_heatmaps, suggest_mutations(scores), csv_files
 
@@ -383,5 +418,17 @@ with tranception_design:
     gr.Markdown("Links: <a href='https://proceedings.mlr.press/v162/notin22a.html' target='_blank'>Paper</a>  <a href='https://github.com/OATML-Markslab/Tranception' target='_blank'>Code</a>  <a href='https://sites.google.com/view/proteingym/substitutions' target='_blank'>ProteinGym</a>")
 
 if __name__ == "__main__":
-    tranception_design.queue()
-    tranception_design.launch()
+    # Configure queue for better resource management
+    tranception_design.queue(
+        max_size=10,  # Limit queue size
+        status_update_rate="auto",  # Show status updates
+        api_open=False  # Disable API to prevent external requests
+    )
+    
+    # Launch with appropriate settings for HF Spaces
+    tranception_design.launch(
+        max_threads=2,  # Limit concurrent threads
+        show_error=True,
+        server_name="0.0.0.0",
+        server_port=7860
+    )