Update app.py

app.py

@@ -1,13 +1,11 @@
 import os
 os.environ["TORCH_DYNAMO_DISABLE"] = "1"
-import subprocess
-import sys
-from pathlib import Path
 import tempfile
 import numpy as np
 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, run_md_simulation, run_relaxation_simulation
 import hashlib
 
@@ -158,247 +156,166 @@ def relax_wrapper(structure_file, task_name, steps, fmax, charge, spin, relax_ce
         traceback.print_exc()
         return (f"Error: {e}", None, "", "", "", None)
 
-# ==== MAIN FUNCTION ====
-def main():
-    with gr.Blocks(theme=gr.themes.Ocean(), title="OrbMol Demo") as demo:
-        with gr.Tabs():
-            # -------- HOME --------
-            with gr.Tab("Home"):
-                with gr.Row():
-                    with gr.Column(scale=1):
-                        gr.Markdown("## Learn more about OrbMol")
-                        
-                        with gr.Accordion("What is OrbMol?", open=False):
-                            gr.Markdown("""
-OrbMol is a suite of quantum-accurate machine learning models for molecular predictions. Built on the **Orb-v3 architecture**, OrbMol provides fast and accurate calculations of energies, forces, and molecular properties.
-                            """)
-                        
-                        with gr.Accordion("Available Models", open=False):
-                            gr.Markdown("""
-**OMol** and **OMol-Direct**
-- Training dataset: OMol25 (>100M calculations)
-- Level of theory: ωB97M-V/def2-TZVPD
-- Applications: biology, organic chemistry, drug discovery
-
-**OMat**
-- Training dataset: OMat24 (>100M inorganic calculations)
-- Level of theory: PBE/PBE+U
-- Applications: materials discovery, photovoltaics, alloys
-                            """)
-                        
-                        with gr.Accordion("Supported File Formats", open=False):
-                            gr.Markdown("""
-Supported formats: `.xyz`, `.pdb`, `.cif`, `.traj`, `.mol`, `.sdf`
-                            """)
-                        
-                        with gr.Accordion("Resources & Support", open=False):
-                            gr.Markdown("""
-- [Orb-v3 paper](https://arxiv.org/abs/2504.06231)
-- [Orb-Models GitHub](https://github.com/orbital-materials/orb-models)
-                            """)
+# ==== UI ====
+with gr.Blocks(theme=gr.themes.Ocean(), title="OrbMol Demo") as demo:
+    with gr.Tabs():
+        # -------- SPE --------
+        with gr.Tab("Single Point Energy"):
+            with gr.Row():
+                with gr.Column(scale=2):
+                    gr.Markdown("# OrbMol — Quantum-Accurate Molecular Predictions")
+                    gr.Markdown("**Supported formats:** .xyz, .pdb, .cif, .traj, .mol, .sdf")
                     
-                    with gr.Column(scale=2):
-                        gr.Image("logo_color_text.png",
-                                 show_share_button=False,
-                                 show_download_button=False,
-                                 show_label=False,
-                                 show_fullscreen_button=False)
-                        
-                        gr.Markdown("# OrbMol — Quantum-Accurate Molecular Predictions")
-                        gr.Markdown("""
-Welcome to OrbMol! Use the tabs above to access:
-
-1. **Single Point Energy**: Calculate energies and forces
-2. **Molecular Dynamics**: Run MD simulations  
-3. **Relaxation / Optimization**: Optimize structures
-
-Supported formats: `.xyz`, `.pdb`, `.cif`, `.traj`, `.mol`, `.sdf`
-                        """)
-            
-            # -------- SPE --------
-            with gr.Tab("Single Point Energy"):
-                with gr.Row():
-                    with gr.Column(scale=2):
-                        gr.Markdown("# OrbMol — Quantum-Accurate Molecular Predictions")
-                        gr.Markdown("**Supported formats:** .xyz, .pdb, .cif, .traj, .mol, .sdf")
-                        
-                        xyz_input = gr.File(
-                            label="Upload Structure File",
-                            file_types=[".xyz", ".pdb", ".cif", ".traj", ".mol", ".sdf"],
-                            file_count="single"
-                        )
-                        task_name_spe = gr.Radio(
-                            ["OMol", "OMat", "OMol-Direct"],
-                            value="OMol",
-                            label="Model Type"
-                        )
-                        with gr.Row():
-                            charge_input = gr.Slider(-10, 10, 0, step=1, label="Charge")
-                            spin_input = gr.Slider(1, 11, 1, step=1, label="Spin Multiplicity")
-                        
-                        run_spe = gr.Button("Run OrbMol Prediction", variant="primary")
-                        
-                    with gr.Column(variant="panel", min_width=500):
-                        spe_out = gr.Textbox(label="Energy & Forces", lines=15, interactive=False)
-                        spe_status = gr.Textbox(label="Status", interactive=False)
-                        
-                        spe_viewer = Molecule3D(
-                            label="Input Structure Viewer",
-                            reps=DEFAULT_MOLECULAR_REPRESENTATIONS,
-                            config=DEFAULT_MOLECULAR_SETTINGS
-                        )
-                        
-                        task_name_spe.change(
-                            lambda x: (
-                                gr.update(visible=x in ["OMol", "OMol-Direct"]),
-                                gr.update(visible=x in ["OMol", "OMol-Direct"])
-                            ),
-                            [task_name_spe],
-                            [charge_input, spin_input]
-                        )
-                        
-                run_spe.click(
-                    predict_molecule,
-                    [xyz_input, task_name_spe, charge_input, spin_input],
-                    [spe_out, spe_status, spe_viewer]
-                )
-
-            # -------- MD --------
-            with gr.Tab("Molecular Dynamics"):
-                with gr.Row():
-                    with gr.Column(scale=2):
-                        gr.Markdown("## Molecular Dynamics Simulation")
-                        
-                        xyz_md = gr.File(
-                            label="Upload Structure File",
-                            file_types=[".xyz", ".pdb", ".cif", ".traj", ".mol", ".sdf"],
-                            file_count="single"
-                        )
-                        task_name_md = gr.Radio(
-                            ["OMol", "OMat", "OMol-Direct"],
-                            value="OMol",
-                            label="Model Type"
-                        )
-                        with gr.Row():
-                            charge_md = gr.Slider(-10, 10, 0, step=1, label="Charge")
-                            spin_md = gr.Slider(1, 11, 1, step=1, label="Spin Multiplicity")
-                        with gr.Row():
-                            steps_md = gr.Slider(10, 2000, 100, step=10, label="Steps")
-                            temp_md = gr.Slider(10, 1500, 300, step=10, label="Temperature (K)")
-                        with gr.Row():
-                            timestep_md = gr.Slider(0.1, 5.0, 1.0, step=0.1, label="Timestep (fs)")
-                            ensemble_md = gr.Radio(["NVE", "NVT"], value="NVE", label="Ensemble")
-                        run_md_btn = gr.Button("Run MD Simulation", variant="primary")
-
-                    with gr.Column(variant="panel", min_width=520):
-                        md_status = gr.Textbox(label="MD Status", interactive=False)
-                        md_traj = gr.File(label="Trajectory (.traj)", interactive=False)
-                        
-                        md_viewer = Molecule3D(
-                            label="MD Result Viewer",
-                            reps=DEFAULT_MOLECULAR_REPRESENTATIONS,
-                            config=DEFAULT_MOLECULAR_SETTINGS
-                        )
-                        
-                        md_log = gr.Textbox(label="Log", interactive=False, lines=15)
-                        md_script = gr.Code(label="Reproduction Script", language="python", interactive=False, lines=20)
-                        md_explain = gr.Markdown()
-                        
-                        task_name_md.change(
-                            lambda x: (
-                                gr.update(visible=x in ["OMol", "OMol-Direct"]),
-                                gr.update(visible=x in ["OMol", "OMol-Direct"])
-                            ),
-                            [task_name_md],
-                            [charge_md, spin_md]
-                        )
-                        
-                run_md_btn.click(
-                    md_wrapper,
-                    [xyz_md, task_name_md, charge_md, spin_md, steps_md, temp_md, timestep_md, ensemble_md],
-                    [md_status, md_traj, md_log, md_script, md_explain, md_viewer]
-                )
+                    xyz_input = gr.File(
+                        label="Upload Structure File",
+                        file_types=[".xyz", ".pdb", ".cif", ".traj", ".mol", ".sdf"],
+                        file_count="single"
+                    )
+                    task_name_spe = gr.Radio(
+                        ["OMol", "OMat", "OMol-Direct"],
+                        value="OMol",
+                        label="Model Type"
+                    )
+                    with gr.Row():
+                        charge_input = gr.Slider(-10, 10, 0, step=1, label="Charge")
+                        spin_input = gr.Slider(1, 11, 1, step=1, label="Spin Multiplicity")
+                    
+                    run_spe = gr.Button("Run OrbMol Prediction", variant="primary")
+                    
+                with gr.Column(variant="panel", min_width=500):
+                    spe_out = gr.Textbox(label="Energy & Forces", lines=15, interactive=False)
+                    spe_status = gr.Textbox(label="Status", interactive=False)
+                    
+                    spe_viewer = Molecule3D(
+                        label="Input Structure Viewer",
+                        reps=DEFAULT_MOLECULAR_REPRESENTATIONS,
+                        config=DEFAULT_MOLECULAR_SETTINGS
+                    )
+                    
+                    task_name_spe.change(
+                        lambda x: (
+                            gr.update(visible=x in ["OMol", "OMol-Direct"]),
+                            gr.update(visible=x in ["OMol", "OMol-Direct"])
+                        ),
+                        [task_name_spe],
+                        [charge_input, spin_input]
+                    )
+                    
+            run_spe.click(
+                predict_molecule,
+                [xyz_input, task_name_spe, charge_input, spin_input],
+                [spe_out, spe_status, spe_viewer]
+            )
 
-            # -------- Relax --------
-            with gr.Tab("Relaxation / Optimization"):
-                with gr.Row():
-                    with gr.Column(scale=2):
-                        gr.Markdown("## Structure Relaxation/Optimization")
-                        
-                        xyz_rlx = gr.File(
-                            label="Upload Structure File",
-                            file_types=[".xyz", ".pdb", ".cif", ".traj", ".mol", ".sdf"],
-                            file_count="single"
-                        )
-                        task_name_rlx = gr.Radio(
-                            ["OMol", "OMat", "OMol-Direct"],
-                            value="OMol",
-                            label="Model Type"
-                        )
-                        with gr.Row():
-                            steps_rlx = gr.Slider(1, 2000, 300, step=1, label="Max Steps")
-                            fmax_rlx = gr.Slider(0.001, 0.5, 0.05, step=0.001, label="Fmax (eV/Å)")
-                        with gr.Row():
-                            charge_rlx = gr.Slider(-10, 10, 0, step=1, label="Charge")
-                            spin_rlx = gr.Slider(1, 11, 1, step=1, label="Spin")
-                        relax_cell = gr.Checkbox(False, label="Relax Unit Cell")
-                        run_rlx_btn = gr.Button("Run Optimization", variant="primary")
+        # -------- MD --------
+        with gr.Tab("Molecular Dynamics"):
+            with gr.Row():
+                with gr.Column(scale=2):
+                    gr.Markdown("## Molecular Dynamics Simulation")
+                    
+                    xyz_md = gr.File(
+                        label="Upload Structure File",
+                        file_types=[".xyz", ".pdb", ".cif", ".traj", ".mol", ".sdf"],
+                        file_count="single"
+                    )
+                    task_name_md = gr.Radio(
+                        ["OMol", "OMat", "OMol-Direct"],
+                        value="OMol",
+                        label="Model Type"
+                    )
+                    with gr.Row():
+                        charge_md = gr.Slider(-10, 10, 0, step=1, label="Charge")
+                        spin_md = gr.Slider(1, 11, 1, step=1, label="Spin Multiplicity")
+                    with gr.Row():
+                        steps_md = gr.Slider(10, 2000, 100, step=10, label="Steps")
+                        temp_md = gr.Slider(10, 1500, 300, step=10, label="Temperature (K)")
+                    with gr.Row():
+                        timestep_md = gr.Slider(0.1, 5.0, 1.0, step=0.1, label="Timestep (fs)")
+                        ensemble_md = gr.Radio(["NVE", "NVT"], value="NVE", label="Ensemble")
+                    run_md_btn = gr.Button("Run MD Simulation", variant="primary")
 
-                    with gr.Column(variant="panel", min_width=520):
-                        rlx_status = gr.Textbox(label="Status", interactive=False)
-                        rlx_traj = gr.File(label="Trajectory (.traj)", interactive=False)
-                        
-                        rlx_viewer = Molecule3D(
-                            label="Optimized Structure Viewer",
-                            reps=DEFAULT_MOLECULAR_REPRESENTATIONS,
-                            config=DEFAULT_MOLECULAR_SETTINGS
-                        )
-                        
-                        rlx_log = gr.Textbox(label="Log", interactive=False, lines=15)
-                        rlx_script = gr.Code(label="Reproduction Script", language="python", interactive=False, lines=20)
-                        rlx_explain = gr.Markdown()
+                with gr.Column(variant="panel", min_width=520):
+                    md_status = gr.Textbox(label="MD Status", interactive=False)
+                    md_traj = gr.File(label="Trajectory (.traj)", interactive=False)
+                    
+                    md_viewer = Molecule3D(
+                        label="MD Result Viewer",
+                        reps=DEFAULT_MOLECULAR_REPRESENTATIONS,
+                        config=DEFAULT_MOLECULAR_SETTINGS
+                    )
+                    
+                    md_log = gr.Textbox(label="Log", interactive=False, lines=15)
+                    md_script = gr.Code(label="Reproduction Script", language="python", interactive=False, lines=20)
+                    md_explain = gr.Markdown()
+                    
+                    task_name_md.change(
+                        lambda x: (
+                            gr.update(visible=x in ["OMol", "OMol-Direct"]),
+                            gr.update(visible=x in ["OMol", "OMol-Direct"])
+                        ),
+                        [task_name_md],
+                        [charge_md, spin_md]
+                    )
+                    
+            run_md_btn.click(
+                md_wrapper,
+                [xyz_md, task_name_md, charge_md, spin_md, steps_md, temp_md, timestep_md, ensemble_md],
+                [md_status, md_traj, md_log, md_script, md_explain, md_viewer]
+            )
 
-                        task_name_rlx.change(
-                            lambda x: (
-                                gr.update(visible=x in ["OMol", "OMol-Direct"]),
-                                gr.update(visible=x in ["OMol", "OMol-Direct"])
-                            ),
-                            [task_name_rlx],
-                            [charge_rlx, spin_rlx]
-                        )
-                        
-                run_rlx_btn.click(
-                    relax_wrapper,
-                    [xyz_rlx, task_name_rlx, steps_rlx, fmax_rlx, charge_rlx, spin_rlx, relax_cell],
-                    [rlx_status, rlx_traj, rlx_log, rlx_script, rlx_explain, rlx_viewer]
-                )
+        # -------- Relax --------
+        with gr.Tab("Relaxation / Optimization"):
+            with gr.Row():
+                with gr.Column(scale=2):
+                    gr.Markdown("## Structure Relaxation/Optimization")
+                    
+                    xyz_rlx = gr.File(
+                        label="Upload Structure File",
+                        file_types=[".xyz", ".pdb", ".cif", ".traj", ".mol", ".sdf"],
+                        file_count="single"
+                    )
+                    task_name_rlx = gr.Radio(
+                        ["OMol", "OMat", "OMol-Direct"],
+                        value="OMol",
+                        label="Model Type"
+                    )
+                    with gr.Row():
+                        steps_rlx = gr.Slider(1, 2000, 300, step=1, label="Max Steps")
+                        fmax_rlx = gr.Slider(0.001, 0.5, 0.05, step=0.001, label="Fmax (eV/Å)")
+                    with gr.Row():
+                        charge_rlx = gr.Slider(-10, 10, 0, step=1, label="Charge")
+                        spin_rlx = gr.Slider(1, 11, 1, step=1, label="Spin")
+                    relax_cell = gr.Checkbox(False, label="Relax Unit Cell")
+                    run_rlx_btn = gr.Button("Run Optimization", variant="primary")
 
-        demo.queue(default_concurrency_limit=None)
-        demo.launch(ssr_mode=False)
+                with gr.Column(variant="panel", min_width=520):
+                    rlx_status = gr.Textbox(label="Status", interactive=False)
+                    rlx_traj = gr.File(label="Trajectory (.traj)", interactive=False)
+                    
+                    rlx_viewer = Molecule3D(
+                        label="Optimized Structure Viewer",
+                        reps=DEFAULT_MOLECULAR_REPRESENTATIONS,
+                        config=DEFAULT_MOLECULAR_SETTINGS
+                    )
+                    
+                    rlx_log = gr.Textbox(label="Log", interactive=False, lines=15)
+                    rlx_script = gr.Code(label="Reproduction Script", language="python", interactive=False, lines=20)
+                    rlx_explain = gr.Markdown()
 
+                    task_name_rlx.change(
+                        lambda x: (
+                            gr.update(visible=x in ["OMol", "OMol-Direct"]),
+                            gr.update(visible=x in ["OMol", "OMol-Direct"])
+                        ),
+                        [task_name_rlx],
+                        [charge_rlx, spin_rlx]
+                    )
+                    
+            run_rlx_btn.click(
+                relax_wrapper,
+                [xyz_rlx, task_name_rlx, steps_rlx, fmax_rlx, charge_rlx, spin_rlx, relax_cell],
+                [rlx_status, rlx_traj, rlx_log, rlx_script, rlx_explain, rlx_viewer]
+            )
 
 if __name__ == "__main__":
-    # On load, build and install the gradio_molecule3d fork
-    subprocess.call(
-        ["gradio", "cc", "install"], cwd=Path(__file__).parent / "gradio_molecule3d/"
-    )
-    subprocess.call(
-        ["gradio", "cc", "build"], cwd=Path(__file__).parent / "gradio_molecule3d/"
-    )
-    subprocess.call(
-        [
-            sys.executable,
-            "-m",
-            "pip",
-            "install",
-            Path(__file__).parent
-            / "gradio_molecule3d/"
-            / "dist/gradio_molecule3d-0.0.7-py3-none-any.whl",
-        ],
-        cwd=Path(__file__).parent.parent,
-    )
-
-    # Load gradio_molecule3d only once it's built and installed
-    from gradio_molecule3d import Molecule3D
-
-    main()
+    demo.launch(server_name="0.0.0.0", server_port=7860, show_error=True)