check uma access on file upload

app.py

@@ -12,11 +12,8 @@ from pathlib import Path
 
 import gradio as gr
 
-from simulation_scripts import (
-    run_md_simulation,
-    run_relaxation_simulation,
-    validate_ase_atoms_and_login,
-)
+from simulation_scripts import (run_md_simulation, run_relaxation_simulation,
+                                validate_ase_atoms_and_login)
 
 DEFAULT_MOLECULAR_REPRESENTATIONS = [
     {

gradio_molecule3d/backend/gradio_molecule3d/molecule3d.py

@@ -9,8 +9,8 @@ from typing import TYPE_CHECKING, Any, Callable, Literal, Sequence
 
 import gradio_client.utils as client_utils
 from gradio_client import handle_file
-from gradio_client.documentation import document
 
+import gradio as gr
 from gradio import processing_utils
 from gradio.components.base import Component
 from gradio.data_classes import FileData, ListFiles
@@ -107,7 +107,7 @@ def convert_file_to_pdb(file_path: str | Path, gradio_cache: str | Path) -> str:
         structures = ase.io.read(file_path, ':')
     except Exception as e:
         # Bad upload structure, no need to visualize
-        return None
+        raise gr.Error(f'Error parsing file with ase: {str(e)}')
 
     if all(structures[0].pbc):
         # find the minimum number of repeats in each unit cell direction to meet at least 20 angstroms
@@ -322,11 +322,14 @@ class Molecule3D(Component):
             )
         else:
             value = convert_file_to_pdb(str(value), self.GRADIO_CACHE)
-            return FileData(
-                path=value,
-                orig_name=Path(value).name,
-                size=Path(value).stat().st_size,
-            )
+            if value is not None:
+                return FileData(
+                    path=value,
+                    orig_name=Path(value).name,
+                    size=Path(value).stat().st_size,
+                )
+            else:
+                return None
 
     def process_example(self, value: str | list | None) -> str:
         if value is None:

hf_calculator.py

@@ -29,6 +29,14 @@ def hash_save_file(atoms: ase.Atoms, task_name, path: Path | str):
     return
 
 
+def validate_uma_access(oauth_token):
+    try:
+        hf_hub.HfApi().auth_check(repo_id="facebook/UMA", token=oauth_token.token)
+        return True
+    except (hf_hub.errors.HfHubHTTPError, AttributeError):
+        return False
+
+
 class HFEndpointCalculator(Calculator):
     # A simple calculator that uses the Hugging Face Inference Endpoints to run
 
@@ -47,15 +55,12 @@ class HFEndpointCalculator(Calculator):
         # If we have an example structure, we don't need to check for authentication
         # Otherwise, we need to check if the user is authenticated and has gated access to the UMA models
         if not example:
-            try:
-                hf_hub.HfApi().auth_check(
-                    repo_id="facebook/UMA", token=oauth_token.token
-                )
+            if validate_uma_access(oauth_token):
                 try:
                     hash_save_file(atoms, task_name, "/data/custom_inputs/")
                 except FileNotFoundError:
                     pass
-            except (hf_hub.errors.HfHubHTTPError, AttributeError):
+            else:
                 raise gr.Error(
                     "You need to log in to HF and have gated model access to UMA before running your own simulations!"
                 )

simulation_scripts.py

@@ -23,7 +23,7 @@ from ase.md.velocitydistribution import MaxwellBoltzmannDistribution
 from ase.md.verlet import VelocityVerlet
 from ase.optimize import LBFGS
 
-from hf_calculator import HFEndpointCalculator
+from hf_calculator import HFEndpointCalculator, validate_uma_access
 
 
 def hash_file(file_path):
@@ -43,7 +43,9 @@ INFERENCE_ENDPOINT_URL = os.environ["INFERENCE_ENDPOINT_URL"]
 
 
 def validate_ase_atoms_and_login(
-    structure_file: dict | str, login_button_value: str
+    structure_file: dict | str,
+    login_button_value: str,
+    oauth_token: gr.OAuthToken | None,
 ) -> tuple[gr.Button, gr.Button, str]:
     # Validate and write the uploaded file content
     if not structure_file:
@@ -84,7 +86,8 @@ def validate_ase_atoms_and_login(
             f"Structure file contains {len(atoms)}, which is more than {MAX_ATOMS} atoms. Please use a smaller structure for this demo, or run this on a local machine!",
         )
     elif (hash_file(structure_file) not in EXAMPLE_FILE_HASHES) and (
-        "Logout" not in login_button_value
+        ("Logout" not in login_button_value)
+        or not validate_uma_access(oauth_token=oauth_token)
     ):
         return (
             gr.Button(interactive=False),