fid input parsing and cuda

app.py

@@ -21,6 +21,7 @@ from huggingface_hub.utils import GatedRepoError, RepositoryNotFoundError
 from loguru import logger
 from matplotlib.cm import ScalarMappable
 from matplotlib.colors import Normalize
+from scipy.spatial.distance import cdist
 
 from rocketshp import RocketSHP
 from rocketshp import load_sequence as get_sequence_features
@@ -163,9 +164,11 @@ def predict_rocketshp(
             if structure_code == "":
                 raise gr.Error("Structure input is required for the selected model.")
 
+            structure_code = structure_code.strip().upper()
+
             structure_tmp_dir = tempfile.TemporaryDirectory()
             structure_file = rcsb.fetch(
-                structure_code,
+                structure_code.strip(),
                 "pdb",
                 target_path=structure_tmp_dir.name,
             )
@@ -247,7 +250,14 @@ def predict_rocketshp(
     rmsf = dynamics_pred["rmsf"].squeeze().cpu().numpy()
     gcc_lmi = dynamics_pred["gcc_lmi"].squeeze().cpu().numpy()
     shp = dynamics_pred["shp"].squeeze().cpu().numpy()
-    ca_dist = dynamics_pred["ca_dist"].squeeze().cpu().numpy()
+
+    if is_sequence_model:
+        ca_dist = dynamics_pred["ca_dist"].squeeze().cpu().numpy()
+    else:
+        ca_struct = structure[bs.filter_amino_acids(structure)]
+        ca_struct = structure[structure.atom_name == "CA"]
+        ca_dist = cdist(ca_struct.coord, ca_struct.coord)
+        ca_dist /= 10.0  # convert to nm
 
     fig, plot_file_name = plot_predictions(
         rmsf,
@@ -439,12 +449,10 @@ def check_user_access(oauth_token: gr.OAuthToken | None):
         raise gr.Error("Please log in to use this Space")
 
     token = oauth_token.token
-    print(token)
 
     try:
         # Try to access a file from your private repo
         api = HfApi(token=token)
-        print(api)
 
         # Test access by trying to get repo info
         _ = api.repo_info(
@@ -531,6 +539,7 @@ with rocketshp_gradio:
     sequence = gr.State([])
 
     gr.LoginButton()
+
     model_variant = gr.Dropdown(
         label="Select RocketSHP Model",
         choices=["latest", "v1_seq", "v1_mini"],
@@ -645,4 +654,4 @@ with rocketshp_gradio:
 
 
 if __name__ == "__main__":
-    rocketshp_gradio.launch(share=False)
+    rocketshp_gradio.launch(share=True)