Update src/streamlit_app.py

src/streamlit_app.py

@@ -12,6 +12,8 @@ from ase.visualize import view
 import py3Dmol
 from mace.calculators import mace_mp
 from fairchem.core import pretrained_mlip, FAIRChemCalculator
+from orb_models.forcefield import pretrained
+from orb_models.forcefield.calculator import ORBCalculator
 import pandas as pd
 import yaml # Added for FairChem reference energies
 
@@ -673,7 +675,10 @@ FAIRCHEM_MODELS = {
     "ESEN SM Conserving All OMOL": "esen-sm-conserving-all-omol",
     "ESEN SM Direct All OMOL": "esen-sm-direct-all-omol"
 }
-
+# Define the available ORB models
+ORB_MODELS = {
+    "V3 OMAT Conserving": "orb_v3_conservative_inf_omat",
+}
 @st.cache_resource
 def get_mace_model(model_path, device, selected_default_dtype):
     return mace_mp(model=model_path, device=device, default_dtype=selected_default_dtype)
@@ -742,7 +747,7 @@ if atoms is not None:
 
 
 st.sidebar.markdown("## Model Selection")
-model_type = st.sidebar.radio("Select Model Type:", ["MACE", "FairChem"])
+model_type = st.sidebar.radio("Select Model Type:", ["MACE", "FairChem", "ORB"])
 
 selected_task_type = None # For FairChem UMA
 if model_type == "MACE":
@@ -762,7 +767,12 @@ if model_type == "FairChem":
             spin_multiplicity = st.sidebar.number_input("Spin Multiplicity (2S + 1)", min_value=1, max_value=11, step=2, value=int(atoms.info.get("spin",0)*2+1 if atoms.info.get("spin",0) is not None else 1)) # Assuming spin in atoms.info is S
             atoms.info["charge"] = charge
             atoms.info["spin"] = spin_multiplicity # FairChem expects multiplicity
-
+if model_type == "ORB":
+    selected_model = st.sidebar.selectbox("Select ORB Model:", list(ORB_MODELS.keys()))
+    model_path = ORB_MODELS[selected_model]
+    # if "omat" in selected_model:
+    #     st.sidebar.warning("Using model under Academic Software License (ASL) license, see [https://github.com/gabor1/ASL](https://github.com/gabor1/ASL). To use this model you accept the terms of the license.")
+    selected_default_dtype = st.sidebar.selectbox("Select Precision (default_dtype):", ['float32-high', 'float32-highest', 'float64'])
 if atoms is not None:
     if not check_atom_limit(atoms, selected_model):
         st.stop() # Stop execution if limit exceeded
@@ -842,14 +852,17 @@ if atoms is not None:
                     if model_type == "MACE":
                         # st.write("Setting up MACE calculator...")
                         calc = get_mace_model(model_path, device, selected_default_dtype)
-                    else:  # FairChem
+                    elif model_type == "FairChem":  # FairChem
                         # st.write("Setting up FairChem calculator...")
                         # Workaround for potential dtype issues when switching models
                         if device == "cpu": # Ensure torch default dtype matches if needed
                              torch.set_default_dtype(torch.float32)
                         _ = get_mace_model(MACE_MODELS["MACE MP 0a Small"], 'cpu', 'float32') # Dummy call
                         calc = get_fairchem_model(selected_model, model_path, device, selected_task_type)
-                    
+                    elif model_type == "ORB":
+                        st.write("Setting up ORB calculator...")
+                        orbff = pretrained.orb_v3_conservative_inf_omat(device=device, precision=selected_default_dtype)
+                        calc = ORBCalculator(orbff, device=device)
                     calc_atoms.calc = calc
                     
                     if task == "Energy Calculation":