Update app.py

app.py

@@ -17,7 +17,7 @@ from orb_models.forcefield import pretrained
 from orb_models.forcefield.calculator import ORBCalculator
 
 # -----------------------------
-# Global model
+# Global OrbMol model
 # -----------------------------
 model_calc = None
 
@@ -74,7 +74,7 @@ def predict_molecule(xyz_content, charge=0, spin_multiplicity=1):
         return f"❌ Error during calculation: {str(e)}", "Error"
 
 # -----------------------------
-# Helper: convert trajectory → HTML animation
+# Trajectory → HTML for Gradio
 # -----------------------------
 def traj_to_html(traj_file):
     traj = Trajectory(traj_file)
@@ -89,10 +89,10 @@ def traj_to_html(traj_file):
     view.setStyle({"stick": {}})
     view.zoomTo()
     view.animate({"loop": "forward"})
-    return view._make_html()
+    return view.render()  # ✅ en Gradio funciona mejor que _make_html()
 
 # -----------------------------
-# Molecular dynamics simulation
+# MD simulation
 # -----------------------------
 def run_md(xyz_content, charge=0, spin_multiplicity=1, steps=100, temperature=300, timestep=1.0):
     try:
@@ -111,11 +111,11 @@ def run_md(xyz_content, charge=0, spin_multiplicity=1, steps=100, temperature=30
         atoms.info = {"charge": int(charge), "spin": int(spin_multiplicity)}
         atoms.calc = calc
 
-        # Pre-relaxation
+        # Pre-relax
         opt = LBFGS(atoms)
         opt.run(fmax=0.05, steps=20)
 
-        # Velocities
+        # Initialize velocities
         MaxwellBoltzmannDistribution(atoms, temperature_K=2 * temperature)
 
         # MD setup
@@ -135,20 +135,20 @@ def run_md(xyz_content, charge=0, spin_multiplicity=1, steps=100, temperature=30
         return f"❌ Error during MD simulation: {str(e)}", ""
 
 # -----------------------------
-# Gradio UI
+# Examples
 # -----------------------------
 examples = [
     ["""2
 Hydrogen molecule
 H 0.0 0.0 0.0
 H 0.0 0.0 0.74""", 0, 1],
-    
+
     ["""3
 Water molecule
 O 0.0000 0.0000 0.0000
 H 0.7571 0.0000 0.5864
 H -0.7571 0.0000 0.5864""", 0, 1],
-    
+
     ["""5
 Methane
 C 0.0000 0.0000 0.0000
@@ -158,25 +158,34 @@ H -0.3630 -0.5133 0.8887
 H -0.3630 -0.5133 -0.8887""", 0, 1]
 ]
 
+# -----------------------------
+# Gradio UI
+# -----------------------------
 with gr.Blocks(theme=gr.themes.Ocean(), title="OrbMol + MD Demo") as demo:
-    gr.Markdown("# OrbMol Demo with Molecular Dynamics")
-
     with gr.Tab("Single Point Energy"):
-        xyz_input = gr.Textbox(label="XYZ Coordinates", lines=12)
-        charge_input = gr.Slider(value=0, minimum=-10, maximum=10, step=1, label="Charge")
-        spin_input = gr.Slider(value=1, minimum=1, maximum=11, step=1, label="Spin Multiplicity")
-        run_btn = gr.Button("Run OrbMol Calculation")
-        results_output = gr.Textbox(label="Results", lines=15)
-        status_output = gr.Textbox(label="Status")
+        with gr.Row():
+            with gr.Column(scale=2):
+                with gr.Column(variant="panel"):
+                    gr.Markdown("# OrbMol Demo - Quantum-Accurate Molecular Predictions")
+                    xyz_input = gr.Textbox(
+                        label="XYZ Coordinates", lines=12,
+                        placeholder="Paste XYZ here"
+                    )
+                    with gr.Row():
+                        charge_input = gr.Slider(value=0, label="Charge", minimum=-10, maximum=10, step=1)
+                        spin_input = gr.Slider(value=1, maximum=11, minimum=1, step=1, label="Spin Multiplicity")
+                    run_btn = gr.Button("Run OrbMol Prediction", variant="primary", size="lg")
+            with gr.Column(variant="panel", min_width=500):
+                gr.Markdown("## Results")
+                results_output = gr.Textbox(label="Energy & Forces", lines=15, interactive=False)
+                status_output = gr.Textbox(label="Status", interactive=False, max_lines=1)
+
+        gr.Examples(examples=examples, inputs=[xyz_input, charge_input, spin_input])
+
         run_btn.click(
             predict_molecule,
             inputs=[xyz_input, charge_input, spin_input],
-            outputs=[results_output, status_output],
-        )
-        gr.Examples(
-            examples=examples,
-            inputs=[xyz_input, charge_input, spin_input],
-            label="Try these examples"
+            outputs=[results_output, status_output]
         )
 
     with gr.Tab("Molecular Dynamics"):
@@ -186,7 +195,7 @@ with gr.Blocks(theme=gr.themes.Ocean(), title="OrbMol + MD Demo") as demo:
         steps_input = gr.Slider(value=100, minimum=10, maximum=1000, step=10, label="Steps")
         temp_input = gr.Slider(value=300, minimum=10, maximum=1000, step=10, label="Temperature (K)")
         timestep_input = gr.Slider(value=1.0, minimum=0.1, maximum=5.0, step=0.1, label="Timestep (fs)")
-        run_md_btn = gr.Button("Run MD Simulation")
+        run_md_btn = gr.Button("Run MD Simulation", variant="primary")
         md_status = gr.Textbox(label="MD Status", lines=2)
         md_view = gr.HTML()
         run_md_btn.click(