add better validation

app.py

@@ -12,7 +12,11 @@ from pathlib import Path
 
 import gradio as gr
 
-from simulation_scripts import run_md_simulation, run_relaxation_simulation
+from simulation_scripts import (
+    run_md_simulation,
+    run_relaxation_simulation,
+    validate_ase_atoms_and_login,
+)
 
 DEFAULT_MOLECULAR_REPRESENTATIONS = [
     {
@@ -515,7 +519,8 @@ def main():
                                 label="Relaxation examples!",
                             )
                         with gr.Accordion(
-                            "UMA can use different output heads for the same structures", open=False
+                            "UMA can use different output heads for the same structures",
+                            open=False,
                         ):
                             gr.Examples(
                                 examples=[
@@ -674,7 +679,6 @@ def main():
                                 label="Molecular Dynamics Examples",
                             )
 
-
                         gr.Markdown(
                             """
                             Once you understand how the UMA model can be applied to different types of molecules and materials, the final tab above will help you try it out with your own structures!
@@ -777,7 +781,9 @@ def main():
                         with gr.Column():
                             input_structure.render()
 
-                            gr.LoginButton(size="large")
+                            structure_validation = gr.Markdown()
+
+                            login_button = gr.LoginButton(size="large")
 
                             gr.Markdown(
                                 """
@@ -973,15 +979,15 @@ def main():
         )
 
         input_structure.change(
-            lambda file: gr.Button(interactive=bool(file)),
-            inputs=[input_structure],
-            outputs=[optimization_button],
+            validate_ase_atoms_and_login,
+            inputs=[input_structure, login_button],
+            outputs=[optimization_button, md_button, structure_validation],
         )
 
-        input_structure.change(
-            lambda file: gr.Button(interactive=bool(file)),
-            inputs=[input_structure],
-            outputs=[md_button],
+        login_button.change(
+            validate_ase_atoms_and_login,
+            inputs=[input_structure, login_button],
+            outputs=[optimization_button, md_button, structure_validation],
         )
 
     demo.launch()

simulation_scripts.py

@@ -42,6 +42,58 @@ MAX_ATOMS = os.environ.get("MAX_ATOMS", 2000)
 INFERENCE_ENDPOINT_URL = os.environ["INFERENCE_ENDPOINT_URL"]
 
 
+def validate_ase_atoms_and_login(structure_file, login_button_value):
+    # Validate and write the uploaded file content
+    if not structure_file:
+        return (
+            gr.Button(interactive=False),
+            gr.Button(interactive=False),
+            "Missing input structure!",
+        )
+
+    try:
+        atoms = ase.io.read(structure_file)
+    except Exception as e:
+        return (
+            gr.Button(interactive=False),
+            gr.Button(interactive=False),
+            f"Failed to load ASE input: {str(e)}!",
+        )
+
+    if len(atoms) == 0:
+        return (
+            gr.Button(interactive=False),
+            gr.Button(interactive=False),
+            "No atoms in the structure file!",
+        )
+    elif not (all(atoms.pbc) or np.all(~np.array(atoms.pbc))):
+        return (
+            gr.Button(interactive=False),
+            gr.Button(interactive=False),
+            f"Your atoms has PBC {atoms.pbc}. Mixed PBC are not supported yet - please set PBC all True or False in your structure before uploading!",
+        )
+    elif len(atoms) > MAX_ATOMS:
+        return (
+            gr.Button(interactive=False),
+            gr.Button(interactive=False),
+            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 (
+        "login_button_value" not in login_button_value
+    ):
+        return (
+            gr.Button(interactive=False),
+            gr.Button(interactive=False),
+            "You must login to huggingface to use a custom structure!",
+        )
+    else:
+        return (
+            gr.Button(interactive=True),
+            gr.Button(interactive=True),
+            "Valid structure file and you are logged in!",
+        )
+
+
 def load_check_ase_atoms(structure_file):
     # Validate and write the uploaded file content
     if not structure_file: