Update app.py

app.py

@@ -24,9 +24,31 @@ from biotite.structure import annotate_sse
 import biotite.structure.io as strucio
 import biotite.structure.residues as residues
 import numpy as np
+from huggingface_hub import hf_hub_download, utils
 
 from data.scripts.data_utils import modify_bfactor_biotite
 
+def get_weights_path(repo_id, filename):
+    """
+    Tries to get the local path immediately. If not found, downloads it.
+    """
+    print(f"Looking for {filename} in {repo_id}...")
+    try:
+        # 1. FASTEST: Try loading entirely from local cache (no internet check)
+        return hf_hub_download(
+            repo_id=repo_id, 
+            filename=filename, 
+            local_files_only=True
+        )
+    except (utils.EntryNotFoundError, utils.LocalEntryNotFoundError, FileNotFoundError):
+        # 2. FALLBACK: If not found locally, download it (cached for next time)
+        print(f"Weights not found locally. Downloading from HF Hub...")
+        return hf_hub_download(
+            repo_id=repo_id, 
+            filename=filename, 
+            local_files_only=False
+        )
+
 def process_pdb_file(pdb_file, backbones, sequences, names):
     _name = pdb_file[:-4]
     _chain = ""
@@ -111,8 +133,14 @@ def flex_seq(input_seq, input_file):
     config['inference_args']['device'] = config['inference_args']['device'] if torch.cuda.is_available() else 'cpu'
     model, tokenizer = PT5_classification_model(half_precision=config['mixed_precision'], class_config=class_config)
     model.to(config['inference_args']['device'])
-    state_dict = torch.load(config['inference_args']['seq_model_path'], map_location=config['inference_args']['device'])
-    model.load_state_dict(state_dict, strict=False)
+    repo_id = "Honzus24/Flexpert_weights"
+    file_weights = config['inference_args']['seq_model_path']
+    
+    # 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()
 
     data_to_collate = []
@@ -240,7 +268,14 @@ def flex_3d(input_file):
 
     model.to(config['inference_args']['device'])
     print("Loading 3D model from {}".format(config['inference_args']['3d_model_path']))
-    state_dict = torch.load(config['inference_args']['3d_model_path'], map_location=config['inference_args']['device'])
+    repo_id = "Honzus24/Flexpert_weights"
+    file_weights = config['inference_args']['3d_model_path']
+    
+    # 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()