Update src/streamlit_app.py

src/streamlit_app.py

@@ -14,6 +14,7 @@ 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
+from sevenn.calculator import SevenNetCalculator
 import pandas as pd
 import yaml # Added for FairChem reference energies
 import subprocess
@@ -718,8 +719,14 @@ ORB_MODELS = {
 }
 # Define the available MatterSim models
 MATTERSIM_MODELS = {
-    "V1 SMALL: MatterSim-v1.0.0-1M.pth",
-    "V1 LARGE: MatterSim-v1.0.0-5M.pth"
+    "V1 SMALL": "MatterSim-v1.0.0-1M.pth",
+    "V1 LARGE": "MatterSim-v1.0.0-5M.pth"
+}
+SEVEN_NET_MODELS = {
+    "7net-0": "7net-0",
+    "7net-l3i5": "7net-l3i5",
+    "7net-omat": "7net-omat",
+    "7net-mf-ompa": "7net-mf-ompa"
 }
 @st.cache_resource
 def get_mace_model(model_path, device, selected_default_dtype):
@@ -792,7 +799,7 @@ st.sidebar.markdown("## Model Selection")
 if mattersim_available:
     model_type = st.sidebar.radio("Select Model Type:", ["MACE", "FairChem", "ORB", "MatterSim"])
 else:
-    model_type = st.sidebar.radio("Select Model Type:", ["MACE", "FairChem", "ORB"])
+    model_type = st.sidebar.radio("Select Model Type:", ["MACE", "FairChem", "ORB", "SEVEN_NET"])
 
 selected_task_type = None # For FairChem UMA
 if model_type == "MACE":
@@ -821,6 +828,9 @@ if model_type == "ORB":
 if model_type == "MatterSim":
     selected_model = st.sidebar.selectbox("Select MatterSim Model:", list(MATTERSIM_MODELS.keys()))
     model_path = MATTERSIM_MODELS[selected_model]
+if model_type == "SEVEN_NET":
+    selected_model = st.sidebar.selectbox("Select SEVENNET Model:", list(SEVEN_NET_MODELS.keys()))
+    model_path = SEVEN_NET_MODELS[selected_model]
 if atoms is not None:
     if not check_atom_limit(atoms, selected_model):
         st.stop() # Stop execution if limit exceeded
@@ -914,6 +924,9 @@ if atoms is not None:
                     elif model_type == "MatterSim":
                         st.write("Setting up MatterSim calculator...")
                         calc = MatterSimCalculator(load_path=model_path, device=device)
+                    elif model_type == "SEVEN_NET":
+                        st.write("Setting up SEVENNET calculator...")
+                        calc = SevenNetD3Calculator(model=model_path, device=device)
                     calc_atoms.calc = calc
                     
                     if task == "Energy Calculation":