Update app.py

app.py

@@ -9,7 +9,7 @@ from gradio_molecule3d import Molecule3D
 from simulation_scripts_orbmol import load_orbmol_model, run_md_simulation, run_relaxation_simulation
 import hashlib
 
-# ==== Configuración Molecule3D (igual que Facebook UMA) ====
+# ==== Configuración Molecule3D ====
 DEFAULT_MOLECULAR_REPRESENTATIONS = [
     {
         "model": 0,
@@ -39,26 +39,46 @@ DEFAULT_MOLECULAR_SETTINGS = {
     "disableFog": False,
 }
 
+# ==== Conversión a PDB para Molecule3D ====
+def convert_to_pdb_for_viewer(file_path):
+    """Convierte cualquier archivo a PDB para Molecule3D"""
+    if not file_path or not os.path.exists(file_path):
+        return None
+    
+    try:
+        atoms = read(file_path)
+        
+        cache_dir = os.path.join(tempfile.gettempdir(), "gradio")
+        os.makedirs(cache_dir, exist_ok=True)
+        
+        pdb_path = os.path.join(cache_dir, f"mol_{hashlib.md5(file_path.encode()).hexdigest()[:12]}.pdb")
+        
+        write(pdb_path, atoms, format="proteindatabank")
+        
+        return pdb_path
+    except Exception as e:
+        print(f"Error converting to PDB: {e}")
+        return None
+
 # ==== OrbMol SPE ====
 def predict_molecule(structure_file, task_name, charge=0, spin_multiplicity=1):
     """Single Point Energy + fuerzas (OrbMol)"""
     try:
         calc = load_orbmol_model(task_name)
         if not structure_file:
-            return "Error: Please upload a structure file", "Error"
+            return "Error: Please upload a structure file", "Error", None
 
         file_path = structure_file
         if not os.path.exists(file_path):
-            return f"Error: File not found: {file_path}", "Error"
+            return f"Error: File not found: {file_path}", "Error", None
         if os.path.getsize(file_path) == 0:
-            return f"Error: Empty file: {file_path}", "Error"
+            return f"Error: Empty file: {file_path}", "Error", None
 
         atoms = read(file_path)
-
-        # Solo aplicar charge/spin para OMol
+        
         if task_name in ["OMol", "OMol-Direct"]:
             atoms.info = {"charge": int(charge), "spin": int(spin_multiplicity)}
-
+        
         atoms.calc = calc
         energy = atoms.get_potential_energy()
         forces = atoms.get_forces()
@@ -74,18 +94,20 @@ def predict_molecule(structure_file, task_name, charge=0, spin_multiplicity=1):
         max_force = float(np.max(np.linalg.norm(forces, axis=1)))
         lines += ["", f"Max Force: {max_force:.4f} eV/Å"]
 
-        return "\n".join(lines), f"Calculation completed with {task_name}"
-
+        pdb_file = convert_to_pdb_for_viewer(file_path)
+        
+        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"
+        return f"Error during calculation: {e}", "Error", None
 
 # ==== Wrappers MD y Relax ====
 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, "", "", "")
+            return ("Error: Please upload a structure file", None, "", "", "", None)
 
         traj_path, log_text, script_text, explanation = run_md_simulation(
             structure_file,
@@ -100,17 +122,19 @@ def md_wrapper(structure_file, task_name, charge, spin, steps, tempK, timestep_f
         )
         status = f"MD completed: {int(steps)} steps at {int(tempK)} K ({ensemble})"
 
-        return (status, traj_path, log_text, script_text, explanation)
+        pdb_file = convert_to_pdb_for_viewer(traj_path)
+        
+        return (status, traj_path, log_text, script_text, explanation, pdb_file)
 
     except Exception as e:
         import traceback
         traceback.print_exc()
-        return (f"Error: {e}", None, "", "", "")
+        return (f"Error: {e}", None, "", "", "", None)
 
 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, "", "", "")
+            return ("Error: Please upload a structure file", None, "", "", "", None)
 
         traj_path, log_text, script_text, explanation = run_relaxation_simulation(
             structure_file,
@@ -123,12 +147,14 @@ def relax_wrapper(structure_file, task_name, steps, fmax, charge, spin, relax_ce
         )
         status = f"Relaxation finished (<={int(steps)} steps, fmax={float(fmax)} eV/Å)"
 
-        return (status, traj_path, log_text, script_text, explanation)
+        pdb_file = convert_to_pdb_for_viewer(traj_path)
+        
+        return (status, traj_path, log_text, script_text, explanation, pdb_file)
 
     except Exception as e:
         import traceback
         traceback.print_exc()
-        return (f"Error: {e}", None, "", "", "")
+        return (f"Error: {e}", None, "", "", "", None)
 
 # ==== UI ====
 with gr.Blocks(theme=gr.themes.Ocean(), title="OrbMol Demo") as demo:
@@ -139,7 +165,7 @@ with gr.Blocks(theme=gr.themes.Ocean(), title="OrbMol Demo") as demo:
                 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"],
@@ -153,23 +179,19 @@ with gr.Blocks(theme=gr.themes.Ocean(), title="OrbMol Demo") as demo:
                     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,
-                        render=True,               # ← activado
-                        inputs=[xyz_input],
-                        value=lambda x: x,
-                        interactive=False
+                        config=DEFAULT_MOLECULAR_SETTINGS
                     )
-
+                    
                     task_name_spe.change(
                         lambda x: (
                             gr.update(visible=x in ["OMol", "OMol-Direct"]),
@@ -178,11 +200,11 @@ with gr.Blocks(theme=gr.themes.Ocean(), title="OrbMol Demo") as demo:
                         [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_out, spe_status, spe_viewer]
             )
 
         # -------- MD --------
@@ -190,7 +212,7 @@ with gr.Blocks(theme=gr.themes.Ocean(), title="OrbMol Demo") as demo:
             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"],
@@ -215,21 +237,17 @@ 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 = Molecule3D(
                         label="MD Result Viewer",
                         reps=DEFAULT_MOLECULAR_REPRESENTATIONS,
-                        config=DEFAULT_MOLECULAR_SETTINGS,
-                        render=True,               # ← activado
-                        inputs=[md_traj],
-                        value=lambda x: x,
-                        interactive=False
+                        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"]),
@@ -238,11 +256,11 @@ with gr.Blocks(theme=gr.themes.Ocean(), title="OrbMol Demo") as demo:
                         [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_status, md_traj, md_log, md_script, md_explain, md_viewer]
             )
 
         # -------- Relax --------
@@ -250,7 +268,7 @@ with gr.Blocks(theme=gr.themes.Ocean(), title="OrbMol Demo") as demo:
             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"],
@@ -273,17 +291,13 @@ 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 = Molecule3D(
                         label="Optimized Structure Viewer",
                         reps=DEFAULT_MOLECULAR_REPRESENTATIONS,
-                        config=DEFAULT_MOLECULAR_SETTINGS,
-                        render=True,               # ← activado
-                        inputs=[rlx_traj],
-                        value=lambda x: x,
-                        interactive=False
+                        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()
@@ -296,12 +310,12 @@ with gr.Blocks(theme=gr.themes.Ocean(), title="OrbMol Demo") as demo:
                         [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_status, rlx_traj, rlx_log, rlx_script, rlx_explain, rlx_viewer]
             )
 
 if __name__ == "__main__":
-    demo.launch(server_name="0.0.0.0", server_port=7860, show_error=True)
+    demo.launch(server_name="0.0.0.0", server_port=7860, show_error=True)