Update app.py

app.py

@@ -32,6 +32,44 @@ import biotite.sequence as seq
 
 from data.scripts.data_utils import modify_bfactor_biotite
 
+GLOBAL_MODEL_CACHE = {}
+
+def get_loaded_model_and_tokenizer(target_device):
+    """Loads the model once and caches it in memory."""
+    global _GLOBAL_MODEL, _GLOBAL_TOKENIZER
+    
+    if _GLOBAL_MODEL is not None and _GLOBAL_TOKENIZER is not None:
+        # Move to the requested device (ZeroGPU handles this dynamically)
+        _GLOBAL_MODEL.to(target_device)
+        return _GLOBAL_MODEL, _GLOBAL_TOKENIZER
+
+    print("First run: Initializing model and downloading/loading weights...")
+    
+    # Initialize architecture
+    model, tokenizer = PT5_classification_model(
+        half_precision=config['mixed_precision'], 
+        class_config=class_config
+    )
+    
+    # Fetch weights
+    repo_id = "Honzus24/Flexpert_weights"
+    file_weights = config['inference_args']['seq_model_path']
+    weights_path = get_weights_path(repo_id, file_weights)
+    
+    # Load weights into CPU memory first (best practice for HF Spaces)
+    state_dict = torch.load(weights_path, map_location='cpu')
+    model.load_state_dict(state_dict, strict=False)
+    model.eval()
+    
+    # Cache them globally
+    _GLOBAL_MODEL = model
+    _GLOBAL_TOKENIZER = tokenizer
+    
+    # Move to the target device
+    _GLOBAL_MODEL.to(target_device)
+    
+    return _GLOBAL_MODEL, _GLOBAL_TOKENIZER
+
 def get_first_chain_id(pdb_file):
     try:
         # Load the PDB file
@@ -156,16 +194,29 @@ def core_flex_seq(input_seq, input_file, force_cpu=False):
         target_device = config['inference_args']['device'] if torch.cuda.is_available() else 'cpu'
     config['inference_args']['device'] = target_device
 
-    model, tokenizer = PT5_classification_model(half_precision=config['mixed_precision'], class_config=class_config)
-    model.to(target_device)
+    global GLOBAL_MODEL_CACHE
+    model_key = 'seq'
     
-    repo_id = "Honzus24/Flexpert_weights"
-    file_weights = config['inference_args']['seq_model_path']
-    weights_path = get_weights_path(repo_id, file_weights)
-    
-    state_dict = torch.load(weights_path, map_location=target_device)
-    model.load_state_dict(state_dict, strict=False)
-    model.eval()
+    if model_key not in GLOBAL_MODEL_CACHE:        
+        # 1. Initialize model
+        model, tokenizer = PT5_classification_model(half_precision=config['mixed_precision'], class_config=class_config)
+        
+        # 2. Get weights path
+        repo_id = "Honzus24/Flexpert_weights"
+        file_weights = config['inference_args']['seq_model_path'] # Update for 3D if needed
+        weights_path = get_weights_path(repo_id, file_weights)
+        
+        # 3. Load weights to CPU first (Crucial for ZeroGPU Spaces compatibility)
+        state_dict = torch.load(weights_path, map_location='cpu')
+        model.load_state_dict(state_dict, strict=False)
+        model.eval()
+        
+        # 4. Save to cache
+        GLOBAL_MODEL_CACHE[model_key] = (model, tokenizer)
+
+    # Retrieve from cache and move to the dynamically assigned device
+    model, tokenizer = GLOBAL_MODEL_CACHE[model_key]
+    model.to(target_device)
     
     data_to_collate = []
     for idx, (backbone, sequence) in enumerate(zip(backbones, sequences)):
@@ -299,20 +350,28 @@ def core_flex_3d(input_file):
         target_device = config['inference_args']['device'] if torch.cuda.is_available() else 'cpu'
     config['inference_args']['device'] = target_device
 
-    model, tokenizer = PT5_classification_model(half_precision=config['mixed_precision'], class_config=class_config)
-    model.to(config['inference_args']['device'])
-    
-    repo_id = "Honzus24/Flexpert_weights"
-    print("Loading 3D model from {}".format(config['inference_args']['3d_model_path']))
-    file_weights = config['inference_args']['3d_model_path']
+    model_key = '3d'
     
-    # Get path (instant if cached)
-    weights_path = get_weights_path(repo_id, file_weights)
-    
-    # Load weights
-    state_dict = torch.load(weights_path, map_location=config['inference_args']['device'])
-    model.load_state_dict(state_dict, strict=False)
-    model.eval()
+    if model_key not in GLOBAL_MODEL_CACHE:        
+        # 1. Initialize model
+        model, tokenizer = PT5_classification_model(half_precision=config['mixed_precision'], class_config=class_config)
+        
+        # 2. Get weights path
+        repo_id = "Honzus24/Flexpert_weights"
+        file_weights = config['inference_args']['3d_model_path'] # Update for 3D if needed
+        weights_path = get_weights_path(repo_id, file_weights)
+        
+        # 3. Load weights to CPU first (Crucial for ZeroGPU Spaces compatibility)
+        state_dict = torch.load(weights_path, map_location='cpu')
+        model.load_state_dict(state_dict, strict=False)
+        model.eval()
+        
+        # 4. Save to cache
+        GLOBAL_MODEL_CACHE[model_key] = (model, tokenizer)
+
+    # Retrieve from cache and move to the dynamically assigned device
+    model, tokenizer = GLOBAL_MODEL_CACHE[model_key]
+    model.to(target_device)
     
     data_to_collate = []
     for idx, (backbone, sequence) in enumerate(zip(backbones, sequences)):