multiple_model (#2)

- Added support for multiple pretrained models (OMol,OMat,OMol-Direct) (ba742371795a1f34a0014d9deb3f3700df9d9292)


Co-authored-by: Davide Sarpa <Nekkrad@users.noreply.huggingface.co>

app.py

@@ -6,6 +6,7 @@ import gradio as gr
 from ase.io import read, write
 from ase.io.trajectory import Trajectory
 from gradio_molecule3d import Molecule3D
+from simulation_scripts_orbmol import load_orbmol_model  # CORREGIDO
 import hashlib
 import shutil
 
@@ -110,22 +111,22 @@ def prepare_input_for_viewer(structure_file):
 from orb_models.forcefield import pretrained
 from orb_models.forcefield.calculator import ORBCalculator
 
-_MODEL_CALC = None
-def _load_orbmol_calc():
-    global _MODEL_CALC
-    if _MODEL_CALC is None:
-        orbff = pretrained.orb_v3_conservative_inf_omat(
-            device="cpu", precision="float32-high"
-        )
-        _MODEL_CALC = ORBCalculator(orbff, device="cpu")
-    return _MODEL_CALC
-
-def predict_molecule(structure_file, charge=0, spin_multiplicity=1):
+#_MODEL_CALC = None
+#def _load_orbmol_calc():
+#    global _MODEL_CALC
+#    if _MODEL_CALC is None:
+#        orbff = pretrained.orb_v3_conservative_inf_omat(
+#            device="cpu", precision="float32-high"
+#        )
+#        _MODEL_CALC = ORBCalculator(orbff, device="cpu")
+#    return _MODEL_CALC
+
+def predict_molecule(structure_file,task_name,charge=0, spin_multiplicity=1):
     """
     Single Point Energy + fuerzas (OrbMol). Acepta archivos subidos.
     """
     try:
-        calc = _load_orbmol_calc()
+        calc = load_orbmol_model(task_name)
         if not structure_file:
             return "Error: Please upload a structure file", "Error", None
 
@@ -151,7 +152,7 @@ def predict_molecule(structure_file, charge=0, spin_multiplicity=1):
         # Preparar PDB para visualización
         pdb_file = prepare_input_for_viewer(file_path)
 
-        return "\n".join(lines), "Calculation completed with OrbMol", pdb_file
+        return "\n".join(lines), f"Calculation completed with {task_name}", pdb_file
     except Exception as e:
         import traceback; traceback.print_exc()
         return f"Error during calculation: {e}", "Error", None
@@ -163,7 +164,7 @@ from simulation_scripts_orbmol import (
 )
 
 # ==== Wrappers con debug y Molecule3D ====
-def md_wrapper(structure_file, charge, spin, steps, tempK, timestep_fs, ensemble):
+def md_wrapper(structure_file,task_name, charge, spin, steps, tempK, timestep_fs, ensemble):
     try:
         if not structure_file:
             return ("Error: Please upload a structure file", None, "", "", "", None)
@@ -178,6 +179,7 @@ def md_wrapper(structure_file, charge, spin, steps, tempK, timestep_fs, ensemble
             float(timestep_fs),
             float(tempK),
             "NVT" if ensemble == "NVT" else "NVE",
+            str(task_name),
             int(charge),
             int(spin),
         )
@@ -194,7 +196,7 @@ def md_wrapper(structure_file, charge, spin, steps, tempK, timestep_fs, ensemble
         import traceback; traceback.print_exc()
         return (f"Error: {e}", None, "", "", "", None)
 
-def relax_wrapper(structure_file, steps, fmax, charge, spin, relax_cell):
+def relax_wrapper(structure_file,task_name, steps, fmax, charge, spin, relax_cell):
     try:
         if not structure_file:
             return ("Error: Please upload a structure file", None, "", "", "", None)
@@ -206,6 +208,7 @@ def relax_wrapper(structure_file, steps, fmax, charge, spin, relax_cell):
             file_path,
             int(steps),
             float(fmax),
+            str(task_name),
             int(charge),
             int(spin),
             bool(relax_cell),
@@ -238,9 +241,17 @@ with gr.Blocks(theme=gr.themes.Ocean(), title="OrbMol Demo") as demo:
                         file_types=[".xyz", ".pdb", ".cif", ".traj", ".mol", ".sdf"],
                         file_count="single"
                     )
+                    with gr.Row():
+                        task_name = gr.Radio(
+                            ["OMol", "OMat", "OMol-Direct"],
+                            value="OMol",
+                            label="Model Type",
+                            info="Choose the OrbMol model variant for the calculation."
+                        )
                     with gr.Row():
                         charge_input = gr.Slider(minimum=-10, maximum=10, value=0, step=1, label="Charge")
                         spin_input = gr.Slider(minimum=1, maximum=11, value=1, step=1, label="Spin Multiplicity")
+                    
                     run_spe = gr.Button("Run OrbMol Prediction", variant="primary")
                     
                 with gr.Column(variant="panel", min_width=500):
@@ -256,8 +267,18 @@ with gr.Blocks(theme=gr.themes.Ocean(), title="OrbMol Demo") as demo:
                             }
                         ]
                     )
-
-            run_spe.click(predict_molecule, [xyz_input, charge_input, spin_input], [spe_out, spe_status, spe_viewer])
+                    # Charge and Spin are only applicable to OMol and OMol-Direct
+                    task_name.input(
+                        lambda x: (
+                            (gr.Number(visible=True), gr.Number(visible=True))
+                            if x == "OMol" or x == "OMol-Direct"
+                            else (gr.Number(visible=False), gr.Number(visible=False))
+                            ),
+                            [task_name],
+                            [charge_input, spin_input],
+                            )
+                    
+            run_spe.click(predict_molecule, [xyz_input, task_name,charge_input, spin_input], [spe_out, spe_status, spe_viewer])
 
         # -------- MD --------
         with gr.Tab("Molecular Dynamics"):
@@ -271,6 +292,13 @@ with gr.Blocks(theme=gr.themes.Ocean(), title="OrbMol Demo") as demo:
                         file_types=[".xyz", ".pdb", ".cif", ".traj", ".mol", ".sdf"],
                         file_count="single"
                     )
+                    with gr.Row():
+                        task_name = gr.Radio(
+                            ["OMol", "OMat", "OMol-Direct"],
+                            value="OMol",
+                            label="Model Type",
+                            info="Choose the OrbMol model variant for the calculation."
+                        )
                     with gr.Row():
                         charge_md = gr.Slider(minimum=-10, maximum=10, value=0, step=1, label="Charge")
                         spin_md = gr.Slider(minimum=1, maximum=11, value=1, step=1, label="Spin Multiplicity")
@@ -303,10 +331,21 @@ with gr.Blocks(theme=gr.themes.Ocean(), title="OrbMol Demo") as demo:
                     md_log = gr.Textbox(label="Log", interactive=False, lines=15, max_lines=25)
                     md_script = gr.Code(label="Reproduction Script", language="python", interactive=False, lines=20, max_lines=30)
                     md_explain = gr.Markdown()
-
+                    
+                    # Charge and Spin are only applicable to OMol and OMol-Direct
+                    task_name.input(
+                        lambda x: (
+                            (gr.Number(visible=True), gr.Number(visible=True))
+                            if x == "OMol" or x == "OMol-Direct"
+                            else (gr.Number(visible=False), gr.Number(visible=False))
+                            ),
+                            [task_name],
+                            [charge_md, spin_md],
+                            )
+                    
             run_md_btn.click(
                 md_wrapper,
-                inputs=[xyz_md, charge_md, spin_md, steps_md, temp_md, timestep_md, ensemble_md],
+                inputs=[xyz_md, task_name,charge_md, spin_md, steps_md, temp_md, timestep_md, ensemble_md],
                 outputs=[md_status, md_traj, md_log, md_script, md_explain, md_viewer],
             )
 
@@ -322,8 +361,16 @@ with gr.Blocks(theme=gr.themes.Ocean(), title="OrbMol Demo") as demo:
                         file_types=[".xyz", ".pdb", ".cif", ".traj", ".mol", ".sdf"],
                         file_count="single"
                     )
-                    steps_rlx = gr.Slider(minimum=1, maximum=2000, value=300, step=1, label="Max Steps")
-                    fmax_rlx = gr.Slider(minimum=0.001, maximum=0.5, value=0.05, step=0.001, label="Fmax (eV/Å)")
+                    with gr.Row():
+                        task_name = gr.Radio(
+                            ["OMol", "OMat", "OMol-Direct"],
+                            value="OMol",
+                            label="Model Type",
+                            info="Choose the OrbMol model variant for the calculation."
+                        )
+                    with gr.Row():   
+                        steps_rlx = gr.Slider(minimum=1, maximum=2000, value=300, step=1, label="Max Steps")
+                        fmax_rlx = gr.Slider(minimum=0.001, maximum=0.5, value=0.05, step=0.001, label="Fmax (eV/Å)")
                     with gr.Row():
                         charge_rlx = gr.Slider(minimum=-10, maximum=10, value=0, step=1, label="Charge")
                         spin_rlx = gr.Slider(minimum=1, maximum=11, value=1, step=1, label="Spin")
@@ -352,14 +399,25 @@ with gr.Blocks(theme=gr.themes.Ocean(), title="OrbMol Demo") as demo:
                     rlx_script = gr.Code(label="Reproduction Script", language="python", interactive=False, lines=20, max_lines=30)
                     rlx_explain = gr.Markdown()
 
+                    # Charge and Spin are only applicable to OMol and OMol-Direct
+                    task_name.input(
+                        lambda x: (
+                            (gr.Number(visible=True), gr.Number(visible=True))
+                            if x == "OMol" or x == "OMol-Direct"
+                            else (gr.Number(visible=False), gr.Number(visible=False))
+                            ),
+                            [task_name],
+                            [charge_rlx, spin_rlx],
+                            )
+                    
             run_rlx_btn.click(
                 relax_wrapper,
-                inputs=[xyz_rlx, steps_rlx, fmax_rlx, charge_rlx, spin_rlx, relax_cell],
+                inputs=[xyz_rlx,task_name, steps_rlx, fmax_rlx, charge_rlx, spin_rlx, relax_cell],
                 outputs=[rlx_status, rlx_traj, rlx_log, rlx_script, rlx_explain, rlx_viewer],
             )
 
 print("Starting OrbMol model loading…")
-_ = _load_orbmol_calc()
+#_ = load_orbmol_model(task_name)
 
 if __name__ == "__main__":
     demo.launch(server_name="0.0.0.0", server_port=7860, show_error=True)

simulation_scripts_orbmol.py

@@ -26,16 +26,37 @@ from orb_models.forcefield.calculator import ORBCalculator
 # Global model
 # -----------------------------
 _model_calc: ORBCalculator | None = None
-
-def load_orbmol_model(device: str = "cpu", precision: str = "float32-high") -> ORBCalculator:
-    """Load OrbMol once and reuse the same calculator."""
-    global _model_calc
-    if _model_calc is None:
-        orbff = pretrained.orb_v3_conservative_inf_omat(
-            device=device,
-            precision=precision,
-        )
+_current_task_name = None
+# Necesario para script de reproducción
+model_name = {
+    "OMol": "orb_v3_conservative_omol", 
+    "OMat": "orb_v3_conservative_inf_omat", 
+    "OMol-Direct": "orb_v3_direct_omol"}
+
+
+def load_orbmol_model(task_name,device: str = "cpu", precision: str = "float32-high") -> ORBCalculator:
+    """Load OrbMol calculator, switches only if another model is required."""
+    global _model_calc, _current_task_name
+    if _model_calc is None or _current_task_name != task_name:
+        if task_name == "OMol":
+            orbff= pretrained.orb_v3_conservative_omol(
+                device=device,
+                precision=precision,
+            )
+        elif task_name == "OMat":
+            orbff = pretrained.orb_v3_conservative_inf_omat(
+                device=device,
+                precision=precision,
+            )
+        elif task_name == "OMol-Direct":
+            orbff = pretrained.orb_v3_direct_omol(
+                device=device,
+                precision=precision,
+            )
+        else:
+            raise ValueError(f"Unknown task_name: {task_name}")
         _model_calc = ORBCalculator(orbff, device=device)
+        _current_task_name = task_name
     return _model_calc
 
 # -----------------------------
@@ -79,7 +100,7 @@ def run_md_simulation(
     md_timestep,
     temperature_k,
     md_ensemble,
-    task_name="OMol",  # Siempre OMol para OrbMol
+    task_name,  
     total_charge=0,
     spin_multiplicity=1,
     explanation: str | None = None,
@@ -103,7 +124,7 @@ def run_md_simulation(
         atoms.info["spin"] = spin_multiplicity
 
         # AQUÍ EL CAMBIO: OrbMol en lugar de HFEndpointCalculator
-        calc = load_orbmol_model()
+        calc = load_orbmol_model(task_name)
         atoms.calc = calc
 
         # Progress callback si existe
@@ -164,6 +185,7 @@ def run_md_simulation(
         dyn.run(num_steps)
 
         # Script de reproducción (estilo Facebook pero con OrbMol)
+        
         reproduction_script = f"""
 import ase.io
 from ase.md.velocitydistribution import MaxwellBoltzmannDistribution
@@ -181,7 +203,7 @@ atoms = ase.io.read('input_file.traj')
 atoms.info["charge"] = {total_charge}
 atoms.info["spin"] = {spin_multiplicity}
 # Set up the OrbMol calculator
-orbff = pretrained.orb_v3_conservative_inf_omat(device='cpu', precision='float32-high')
+orbff = pretrained.{model_name[task_name]}(device='cpu', precision='float32-high')
 atoms.calc = ORBCalculator(orbff, device='cpu')
 # Do a quick pre-relaxation to make sure the system is stable
 opt = LBFGS(atoms, trajectory="relaxation_output.traj")
@@ -227,7 +249,7 @@ def run_relaxation_simulation(
     structure_file,
     num_steps,
     fmax,
-    task_name="OMol",  # Siempre OMol para OrbMol
+    task_name,
     total_charge: float = 0,
     spin_multiplicity: float = 1,
     relax_unit_cell=False,
@@ -252,7 +274,7 @@ def run_relaxation_simulation(
         atoms.info["spin"] = spin_multiplicity
 
         # AQUÍ EL CAMBIO: OrbMol en lugar de HFEndpointCalculator
-        calc = load_orbmol_model()
+        calc = load_orbmol_model(task_name)
         atoms.calc = calc
 
         # Archivos temporales
@@ -294,7 +316,7 @@ atoms = ase.io.read('input_file.traj')
 atoms.info["charge"] = {total_charge}
 atoms.info["spin"] = {spin_multiplicity}
 # Set up the OrbMol calculator
-orbff = pretrained.orb_v3_conservative_inf_omat(device='cpu', precision='float32-high')
+orbff = pretrained.{model_name[task_name]}(device='cpu', precision='float32-high')
 atoms.calc = ORBCalculator(orbff, device='cpu')
 # Initialize the optimizer
 relax_unit_cell = {relax_unit_cell}