Update app.py

app.py

@@ -4,77 +4,27 @@ 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
 import hashlib
 
-# ==== Forzar visor HTML con 3Dmol.js ====
-HAVE_MOL3D = False
-
-# ==== Fallback HTML con 3Dmol.js ====
-def traj_to_html(traj_path, width=520, height=520, interval_ms=200):
-    """
-    Lee una .traj de ASE y genera un visor HTML (3Dmol.js) con animación.
-    """
+# ==== Molecule3D viewer preparation ====
+def prepare_molecule_for_viewer(traj_path):
+    """Convert trajectory to format compatible with Molecule3D viewer"""
     if not traj_path or not os.path.exists(traj_path):
-        return "<div style='color:#b00; padding:20px;'>No trajectory file found</div>"
-
-    viewer_id = f"viewer_{abs(hash(traj_path)) % 100000}"
-
+        return None
+    
     try:
         traj = Trajectory(traj_path)
         if len(traj) == 0:
-            return "<div style='color:#555; padding:20px;'>Empty trajectory</div>"
+            return None
+        
+        # Save last frame as PDB for Molecule3D
+        temp_pdb = tempfile.NamedTemporaryFile(suffix='.pdb', delete=False)
+        write(temp_pdb.name, traj[-1], format='pdb')
+        return temp_pdb.name
     except Exception as e:
-        return f"<div style='color:#b00; padding:20px;'>Error: {e}</div>"
-
-    xyz_frames = []
-    for atoms in traj:
-        symbols = atoms.get_chemical_symbols()
-        coords = atoms.get_positions()
-        parts = [str(len(symbols)), "frame"]
-        for s, (x, y, z) in zip(symbols, coords):
-            parts.append(f"{s} {x:.6f} {y:.6f} {z:.6f}")
-        xyz_frames.append("\n".join(parts))
-
-    frames_json = str(xyz_frames).replace("'", '"')
-
-    html = f"""
-<div style="margin-bottom:10px; padding:10px; background:#f5f5f5; border-radius:5px;">
-  <strong>🧬 3D Molecular Viewer</strong> — {len(xyz_frames)} frames
-</div>
-<div id="{viewer_id}" style="width:{width}px; height:{height}px; position:relative; border:2px solid #ddd; border-radius:8px; background:#fafafa;"></div>
-<script>
-if (typeof window.$3Dmol === 'undefined') {{
-  var script = document.createElement('script');
-  script.src = 'https://3dmol.org/build/3Dmol-min.js';
-  script.onload = function() {{ setTimeout(function() {{ initViewer_{viewer_id}(); }}, 100); }};
-  document.head.appendChild(script);
-}} else {{
-  initViewer_{viewer_id}();
-}}
-function initViewer_{viewer_id}() {{
-  var el = document.getElementById("{viewer_id}");
-  if (!el || typeof $3Dmol === "undefined") return;
-  var viewer = $3Dmol.createViewer(el, {{backgroundColor: 'white'}});
-  var frames = {frames_json};
-  var currentFrame = 0;
-  function showFrame(i) {{
-    viewer.clear();
-    viewer.addModel(frames[i], "xyz");
-    viewer.setStyle({{}}, {{stick: {{}}, sphere: {{}}}});
-    viewer.zoomTo();
-    viewer.render();
-  }}
-  showFrame(0);
-  if (frames.length > 1) {{
-    setInterval(function() {{
-      currentFrame = (currentFrame + 1) % frames.length;
-      showFrame(currentFrame);
-    }}, {interval_ms});
-  }}
-}}
-</script>
-"""
-    return html
+        print(f"Error preparing molecule: {e}")
+        return None
 
 # ==== OrbMol SPE ====
 from orb_models.forcefield import pretrained
@@ -136,7 +86,7 @@ from simulation_scripts_orbmol import (
 def md_wrapper(structure_file, charge, spin, steps, tempK, timestep_fs, ensemble):
     try:
         if not structure_file:
-            return ("Error: Please upload a structure file", None, "", "", "", None, "")
+            return ("Error: Please upload a structure file", None, "", "", "", None)
 
         # structure_file es directamente el path del archivo
         file_path = structure_file
@@ -153,16 +103,18 @@ def md_wrapper(structure_file, charge, spin, steps, tempK, timestep_fs, ensemble
         )
         status = f"MD completed: {int(steps)} steps at {int(tempK)} K ({ensemble})"
 
-        html_value = traj_to_html(traj_path)
-        return (status, traj_path, log_text, script_text, explanation, None, html_value)
+        # Preparar archivo PDB para el visor 3D
+        pdb_file = prepare_molecule_for_viewer(traj_path)
+        
+        return (status, traj_path, log_text, script_text, explanation, pdb_file)
 
     except Exception as e:
-        return (f"Error: {e}", None, "", "", "", None, "")
+        return (f"Error: {e}", None, "", "", "", None)
 
 def relax_wrapper(structure_file, steps, fmax, charge, spin, relax_cell):
     try:
         if not structure_file:
-            return ("Error: Please upload a structure file", None, "", "", "", None, "")
+            return ("Error: Please upload a structure file", None, "", "", "", None)
 
         # structure_file es directamente el path del archivo
         file_path = structure_file
@@ -177,11 +129,13 @@ def relax_wrapper(structure_file, steps, fmax, charge, spin, relax_cell):
         )
         status = f"Relaxation finished (≤ {int(steps)} steps, fmax={float(fmax)} eV/Å)"
 
-        html_value = traj_to_html(traj_path)
-        return (status, traj_path, log_text, script_text, explanation, None, html_value)
+        # Preparar archivo PDB para el visor 3D
+        pdb_file = prepare_molecule_for_viewer(traj_path)
+        
+        return (status, traj_path, log_text, script_text, explanation, pdb_file)
 
     except Exception as e:
-        return (f"Error: {e}", None, "", "", "", None, "")
+        return (f"Error: {e}", None, "", "", "", None)
 
 # ==== UI ====
 with gr.Blocks(theme=gr.themes.Ocean(), title="OrbMol Demo") as demo:
@@ -237,8 +191,7 @@ with gr.Blocks(theme=gr.themes.Ocean(), title="OrbMol Demo") as demo:
                 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_placeholder = gr.HTML(visible=False)
-                    md_html = gr.HTML(label="Trajectory Viewer")
+                    md_viewer = Molecule3D(label="3D Structure Viewer")
                     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()
@@ -246,7 +199,7 @@ with gr.Blocks(theme=gr.themes.Ocean(), title="OrbMol Demo") as demo:
             run_md_btn.click(
                 md_wrapper,
                 inputs=[xyz_md, 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_placeholder, md_html],
+                outputs=[md_status, md_traj, md_log, md_script, md_explain, md_viewer],
             )
 
         # -------- Relax --------
@@ -273,8 +226,7 @@ with gr.Blocks(theme=gr.themes.Ocean(), title="OrbMol Demo") as demo:
                 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_placeholder = gr.HTML(visible=False)
-                    rlx_html = gr.HTML(label="Final Structure")
+                    rlx_viewer = Molecule3D(label="3D Structure Viewer")
                     rlx_log = gr.Textbox(label="Log", interactive=False, lines=15, max_lines=25)
                     rlx_script = gr.Code(label="Reproduction Script", language="python", interactive=False, lines=20, max_lines=30)
                     rlx_explain = gr.Markdown()
@@ -282,7 +234,7 @@ with gr.Blocks(theme=gr.themes.Ocean(), title="OrbMol Demo") as demo:
             run_rlx_btn.click(
                 relax_wrapper,
                 inputs=[xyz_rlx, steps_rlx, fmax_rlx, charge_rlx, spin_rlx, relax_cell],
-                outputs=[rlx_status, rlx_traj, rlx_log, rlx_script, rlx_explain, rlx_viewer_placeholder, rlx_html],
+                outputs=[rlx_status, rlx_traj, rlx_log, rlx_script, rlx_explain, rlx_viewer],
             )
 
 print("Starting OrbMol model loading…")