Update app.py

app.py

@@ -5,120 +5,163 @@ 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
 import shutil
 
-# ==== Molecule3D viewer preparation CORREGIDO ====
-def prepare_molecule_for_viewer(traj_path):
-    """Convert trajectory to format compatible with Molecule3D viewer"""
+# ==== Test de gradio_molecule3d ====
+try:
+    from gradio_molecule3d import Molecule3D
+    MOLECULE3D_AVAILABLE = True
+    print("✅ gradio_molecule3d imported successfully")
+except ImportError as e:
+    MOLECULE3D_AVAILABLE = False
+    print(f"❌ gradio_molecule3d import failed: {e}")
+
+# ==== Debug function ====
+def debug_pdb_creation(traj_path):
+    """Debug PDB creation step by step"""
+    print(f"\n=== DEBUG PDB CREATION ===")
+    print(f"Input trajectory: {traj_path}")
+    print(f"Trajectory exists: {os.path.exists(traj_path) if traj_path else False}")
+    
     if not traj_path or not os.path.exists(traj_path):
-        print("No trajectory path provided or file doesn't exist")
-        return None
+        return None, "No trajectory file"
     
     try:
+        # Step 1: Read trajectory
         traj = Trajectory(traj_path)
-        if len(traj) == 0:
-            print("Empty trajectory")
-            return None
+        print(f"Trajectory frames: {len(traj)}")
         
-        # Debug info
-        print(f"Preparing viewer: {len(traj)} frames, {len(traj[-1])} atoms")
-        
-        # Crear directorio cache de Gradio si no existe
-        gradio_cache_dir = os.path.join(tempfile.gettempdir(), "gradio")
-        os.makedirs(gradio_cache_dir, exist_ok=True)
+        if len(traj) == 0:
+            return None, "Empty trajectory"
         
-        # Crear PDB temporal PERSISTENTE en el directorio cache de Gradio
-        unique_id = hashlib.md5(str(traj_path).encode() + str(os.path.getmtime(traj_path)).encode()).hexdigest()[:12]
-        pdb_filename = f"molecule_{unique_id}.pdb"
-        pdb_path = os.path.join(gradio_cache_dir, pdb_filename)
+        # Step 2: Get atoms
+        atoms = traj[-1]
+        print(f"Atoms count: {len(atoms)}")
+        print(f"Chemical symbols: {atoms.get_chemical_symbols()[:5]}...")  # First 5
         
-        # Si ya existe, no lo recreamos
-        if os.path.exists(pdb_path) and os.path.getsize(pdb_path) > 0:
-            print(f"PDB already exists: {pdb_path}")
-            return pdb_path
+        # Step 3: Create cache directory
+        cache_dir = os.path.join(tempfile.gettempdir(), "gradio_cache")
+        os.makedirs(cache_dir, exist_ok=True)
+        print(f"Cache directory: {cache_dir}")
+        print(f"Cache dir exists: {os.path.exists(cache_dir)}")
         
-        # Tomar el último frame de la trayectoria
-        atoms = traj[-1]
+        # Step 4: Create unique filename
+        unique_id = hashlib.md5(str(traj_path).encode()).hexdigest()[:12]
+        pdb_filename = f"mol_{unique_id}.pdb"
+        pdb_path = os.path.join(cache_dir, pdb_filename)
+        print(f"Target PDB path: {pdb_path}")
         
-        # Escribir PDB con elementos en minúsculas (como requiere gradio_molecule3d)
-        try:
-            # Usar ASE para escribir PDB
-            write(pdb_path, atoms, format='pdb')
-            
-            # Post-procesar para asegurar formato correcto
-            with open(pdb_path, 'r') as f:
-                content = f.read()
-            
-            # Asegurar que los símbolos de elementos estén en formato correcto
-            lines = content.split('\n')
-            processed_lines = []
-            
-            for line in lines:
-                if line.startswith(('ATOM', 'HETATM')):
-                    # Formato PDB: columnas específicas para elemento
-                    processed_lines.append(line)
-                else:
-                    processed_lines.append(line)
-            
-            with open(pdb_path, 'w') as f:
-                f.write('\n'.join(processed_lines))
-                
-        except Exception as write_error:
-            print(f"Error writing PDB: {write_error}")
-            return None
+        # Step 5: Write PDB
+        write(pdb_path, atoms, format='pdb')
+        print(f"PDB written, exists: {os.path.exists(pdb_path)}")
         
-        # Verificar que el archivo se creó correctamente
-        if os.path.exists(pdb_path) and os.path.getsize(pdb_path) > 0:
+        if os.path.exists(pdb_path):
             file_size = os.path.getsize(pdb_path)
-            print(f"PDB created successfully: {pdb_path}, size: {file_size} bytes")
+            print(f"PDB file size: {file_size} bytes")
             
-            # Verificar contenido básico del PDB
+            # Read first few lines
             with open(pdb_path, 'r') as f:
-                first_lines = f.read(200)
-                print(f"PDB content preview: {first_lines[:100]}...")
-                
-            return pdb_path
+                first_lines = f.read(300)
+                print(f"PDB content preview:\n{first_lines}")
+            
+            return pdb_path, f"Success: {file_size} bytes"
         else:
-            print("PDB file was not created properly")
-            return None
+            return None, "PDB file not created"
             
     except Exception as e:
-        print(f"Error preparing molecule for viewer: {e}")
+        print(f"Exception in debug_pdb_creation: {e}")
         import traceback
         traceback.print_exc()
-        return None
+        return None, f"Error: {e}"
 
-# ==== Función para convertir archivo inicial a PDB para SPE ====
-def prepare_input_for_viewer(structure_file):
-    """Convert input structure file to PDB for Molecule3D viewer"""
-    if not structure_file or not os.path.exists(structure_file):
-        return None
+# ==== Simple molecule viewer using HTML + 3Dmol.js ====
+def create_simple_molecule_viewer(atoms_or_pdb_path, viewer_id="mol_viewer"):
+    """Create a simple 3D molecule viewer using 3Dmol.js"""
     
-    try:
-        atoms = read(structure_file)
+    # If it's a path, read atoms
+    if isinstance(atoms_or_pdb_path, str) and os.path.exists(atoms_or_pdb_path):
+        try:
+            atoms = read(atoms_or_pdb_path)
+        except:
+            return "<div style='color:red;'>Error reading file</div>"
+    else:
+        atoms = atoms_or_pdb_path
+    
+    # Convert to XYZ format
+    symbols = atoms.get_chemical_symbols()
+    coords = atoms.get_positions()
+    
+    xyz_content = f"{len(symbols)}\nMolecule\n"
+    for s, (x, y, z) in zip(symbols, coords):
+        xyz_content += f"{s} {x:.6f} {y:.6f} {z:.6f}\n"
+    
+    # Escape for JavaScript
+    xyz_escaped = xyz_content.replace('\n', '\\n').replace('"', '\\"')
+    
+    html = f"""
+<div style="margin:10px; padding:15px; border:2px solid #ddd; border-radius:10px; background:#f9f9f9;">
+    <div style="margin-bottom:10px; font-weight:bold; color:#333;">
+        🧬 Molecule Viewer ({len(symbols)} atoms)
+    </div>
+    <div id="{viewer_id}" style="width:100%; height:400px; background:white; border:1px solid #ccc; border-radius:5px;"></div>
+</div>
+
+<script>
+// Load 3Dmol if not already loaded
+if (typeof window.$3Dmol === 'undefined') {{
+    console.log('Loading 3Dmol.js...');
+    var script = document.createElement('script');
+    script.src = 'https://3dmol.org/build/3Dmol-min.js';
+    script.onload = function() {{
+        console.log('3Dmol.js loaded, initializing viewer...');
+        setTimeout(function() {{ init_{viewer_id}(); }}, 100);
+    }};
+    document.head.appendChild(script);
+}} else {{
+    console.log('3Dmol.js already available, initializing viewer...');
+    init_{viewer_id}();
+}}
+
+function init_{viewer_id}() {{
+    try {{
+        var element = document.getElementById('{viewer_id}');
+        if (!element) {{
+            console.error('Element {viewer_id} not found');
+            return;
+        }}
         
-        # Crear directorio cache de Gradio si no existe
-        gradio_cache_dir = os.path.join(tempfile.gettempdir(), "gradio")
-        os.makedirs(gradio_cache_dir, exist_ok=True)
+        if (typeof $3Dmol === 'undefined') {{
+            console.error('$3Dmol is undefined');
+            return;
+        }}
         
-        # Crear PDB para visualización
-        unique_id = hashlib.md5(str(structure_file).encode()).hexdigest()[:12]
-        pdb_filename = f"input_{unique_id}.pdb"
-        pdb_path = os.path.join(gradio_cache_dir, pdb_filename)
+        console.log('Creating 3Dmol viewer...');
+        var viewer = $3Dmol.createViewer(element, {{
+            backgroundColor: 'white',
+            antialias: true
+        }});
         
-        write(pdb_path, atoms, format='pdb')
+        var xyz_data = "{xyz_escaped}";
+        console.log('Adding molecule model...');
+        viewer.addModel(xyz_data, "xyz");
+        viewer.setStyle({{}}, {{
+            stick: {{radius: 0.2, colorscheme: 'Jmol'}},
+            sphere: {{radius: 0.4, colorscheme: 'Jmol'}}
+        }});
         
-        if os.path.exists(pdb_path) and os.path.getsize(pdb_path) > 0:
-            print(f"Input PDB created: {pdb_path}")
-            return pdb_path
-        else:
-            return None
-            
-    except Exception as e:
-        print(f"Error preparing input for viewer: {e}")
-        return None
+        viewer.zoomTo();
+        viewer.render();
+        console.log('Molecule viewer initialized successfully');
+        
+    }} catch (error) {{
+        console.error('Error initializing molecule viewer:', error);
+        document.getElementById('{viewer_id}').innerHTML = '<div style="color:red; padding:20px;">Error loading 3D viewer: ' + error.message + '</div>';
+    }}
+}}
+</script>
+"""
+    return html
 
 # ==== OrbMol SPE ====
 from orb_models.forcefield import pretrained
@@ -135,21 +178,19 @@ def _load_orbmol_calc():
     return _MODEL_CALC
 
 def predict_molecule(structure_file, charge=0, spin_multiplicity=1):
-    """
-    Single Point Energy + fuerzas (OrbMol). Acepta archivos subidos.
-    """
+    """Single Point Energy + fuerzas (OrbMol)"""
     try:
         calc = _load_orbmol_calc()
         if not structure_file:
-            return "Error: Please upload a structure file", "Error", None
+            return "Error: Please upload a structure file", "Error", "", None if MOLECULE3D_AVAILABLE else ""
 
         file_path = structure_file
 
         if not os.path.exists(file_path):
-            return f"Error: File not found: {file_path}", "Error", None
+            return f"Error: File not found: {file_path}", "Error", "", None if MOLECULE3D_AVAILABLE else ""
         
         if os.path.getsize(file_path) == 0:
-            return f"Error: Empty file: {file_path}", "Error", None
+            return f"Error: Empty file: {file_path}", "Error", "", None if MOLECULE3D_AVAILABLE else ""
 
         atoms = read(file_path)
         atoms.info = {"charge": int(charge), "spin": int(spin_multiplicity)}
@@ -164,24 +205,30 @@ def predict_molecule(structure_file, charge=0, spin_multiplicity=1):
         max_force = float(np.max(np.linalg.norm(forces, axis=1)))
         lines += ["", f"Max Force: {max_force:.4f} eV/Å"]
 
-        # Preparar PDB para visualización
-        pdb_file = prepare_input_for_viewer(file_path)
-
-        return "\n".join(lines), "Calculation completed with OrbMol", pdb_file
+        # Debug info
+        debug_info = f"Molecule: {len(atoms)} atoms, symbols: {', '.join(set(atoms.get_chemical_symbols()))}"
+        
+        # Create viewer
+        if MOLECULE3D_AVAILABLE:
+            pdb_path, debug_msg = debug_pdb_creation(None)  # We'll create from atoms directly
+            return "\n".join(lines), "Calculation completed with OrbMol", debug_info + f"\nDebug: {debug_msg}", pdb_path
+        else:
+            html_viewer = create_simple_molecule_viewer(atoms, f"spe_viewer_{hash(file_path) % 10000}")
+            return "\n".join(lines), "Calculation completed with OrbMol", debug_info, html_viewer
     except Exception as e:
-        return f"Error during calculation: {e}", "Error", None
+        error_msg = f"Error during calculation: {e}"
+        return error_msg, "Error", str(e), None if MOLECULE3D_AVAILABLE else ""
 
-# ==== Simulaciones (helpers) ====
+# ==== Simulaciones ====
 from simulation_scripts_orbmol import (
     run_md_simulation,
     run_relaxation_simulation,
 )
 
-# ==== Wrappers con debug y Molecule3D ====
 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 if MOLECULE3D_AVAILABLE else "", "")
 
         file_path = structure_file
         print(f"MD Wrapper: Processing {file_path}")
@@ -189,7 +236,7 @@ def md_wrapper(structure_file, charge, spin, steps, tempK, timestep_fs, ensemble
         traj_path, log_text, script_text, explanation = run_md_simulation(
             file_path,
             int(steps),
-            20,  # pre-relax steps
+            20,
             float(timestep_fs),
             float(tempK),
             "NVT" if ensemble == "NVT" else "NVE",
@@ -200,59 +247,44 @@ def md_wrapper(structure_file, charge, spin, steps, tempK, timestep_fs, ensemble
 
         print(f"MD completed, trajectory: {traj_path}")
         
-        # Preparar archivo PDB para el visor Molecule3D
-        pdb_file = prepare_molecule_for_viewer(traj_path)
-        print(f"PDB file for Molecule3D: {pdb_file}")
-        
-        return (status, traj_path, log_text, script_text, explanation, pdb_file)
+        if MOLECULE3D_AVAILABLE:
+            pdb_path, debug_msg = debug_pdb_creation(traj_path)
+            debug_info = f"Trajectory: {traj_path}\nDebug: {debug_msg}"
+            return (status, traj_path, log_text, script_text, explanation, pdb_path, debug_info)
+        else:
+            # Create HTML viewer from last frame
+            traj = Trajectory(traj_path)
+            last_atoms = traj[-1]
+            html_viewer = create_simple_molecule_viewer(last_atoms, f"md_viewer_{hash(traj_path) % 10000}")
+            debug_info = f"Trajectory: {traj_path}, frames: {len(traj)}"
+            return (status, traj_path, log_text, script_text, explanation, html_viewer, debug_info)
 
     except Exception as e:
+        error_msg = f"Error: {e}"
         print(f"MD Wrapper Error: {e}")
         import traceback
         traceback.print_exc()
-        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)
-
-        file_path = structure_file
-        print(f"Relax Wrapper: Processing {file_path}")
-
-        traj_path, log_text, script_text, explanation = run_relaxation_simulation(
-            file_path,
-            int(steps),
-            float(fmax),
-            int(charge),
-            int(spin),
-            bool(relax_cell),
-        )
-        status = f"Relaxation finished (≤ {int(steps)} steps, fmax={float(fmax)} eV/Å)"
-
-        print(f"Relaxation completed, trajectory: {traj_path}")
-        
-        # Preparar archivo PDB para el visor Molecule3D
-        pdb_file = prepare_molecule_for_viewer(traj_path)
-        print(f"PDB file for Molecule3D: {pdb_file}")
-        
-        return (status, traj_path, log_text, script_text, explanation, pdb_file)
-
-    except Exception as e:
-        print(f"Relax Wrapper Error: {e}")
-        import traceback
-        traceback.print_exc()
-        return (f"Error: {e}", None, "", "", "", None)
+        return (error_msg, None, "", "", "", None if MOLECULE3D_AVAILABLE else "", str(e))
 
 # ==== UI ====
-with gr.Blocks(theme=gr.themes.Ocean(), title="OrbMol Demo") as demo:
+with gr.Blocks(theme=gr.themes.Ocean(), title="OrbMol Demo - Debug") as demo:
+    
+    # Header with debug info
+    gr.Markdown(f"""
+    # OrbMol Demo - Debug Version
+    
+    **System Status:**
+    - gradio_molecule3d available: {'✅ YES' if MOLECULE3D_AVAILABLE else '❌ NO'}
+    - Fallback viewer: {'❌ NO' if MOLECULE3D_AVAILABLE else '✅ HTML + 3Dmol.js'}
+    """)
+    
     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("Upload molecular structure files (.xyz, .pdb, .cif, .traj) for energy and force calculations.")
+                    gr.Markdown("## OrbMol — Single Point Energy")
+                    gr.Markdown("Upload molecular structure files for energy and force calculations.")
                     
                     xyz_input = gr.File(
                         label="Upload Structure File (.xyz/.pdb/.cif/.traj)",
@@ -268,47 +300,36 @@ with gr.Blocks(theme=gr.themes.Ocean(), title="OrbMol Demo") as demo:
                 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, max_lines=1)
+                    spe_debug = gr.Textbox(label="Debug Info", interactive=False, max_lines=3)
                     
-                    # Viewer para SPE
-                    spe_viewer = Molecule3D(
-                        label="Input Structure Viewer",
-                        reps=[
-                            {
-                                "model": 0,
-                                "chain": "",
-                                "resname": "",
-                                "style": "stick",
-                                "color": "whiteCarbon",
-                                "residue_range": "",
-                                "around": 0,
-                                "byres": False,
-                                "visible": True,
-                                "opacity": 1.0
-                            }
-                        ]
-                    )
+                    # Conditional viewer
+                    if MOLECULE3D_AVAILABLE:
+                        spe_viewer = Molecule3D(label="Structure Viewer (Molecule3D)")
+                        outputs = [spe_out, spe_status, spe_debug, spe_viewer]
+                    else:
+                        spe_viewer = gr.HTML(label="Structure Viewer (3Dmol.js)")
+                        outputs = [spe_out, spe_status, spe_debug, spe_viewer]
 
-            run_spe.click(predict_molecule, [xyz_input, charge_input, spin_input], [spe_out, spe_status, spe_viewer])
+            run_spe.click(predict_molecule, [xyz_input, charge_input, spin_input], outputs)
 
         # -------- MD --------
         with gr.Tab("Molecular Dynamics"):
             with gr.Row():
                 with gr.Column(scale=2):
                     gr.Markdown("## Molecular Dynamics Simulation")
-                    gr.Markdown("Upload your molecular structure and configure MD parameters.")
                     
                     xyz_md = gr.File(
-                        label="Upload Structure File (.xyz/.pdb/.cif/.traj)",
-                        file_types=[".xyz", ".pdb", ".cif", ".traj", ".mol", ".sdf"],
+                        label="Upload Structure File",
+                        file_types=[".xyz", ".pdb", ".cif", ".traj"],
                         file_count="single"
                     )
                     
                     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")
+                        spin_md = gr.Slider(minimum=1, maximum=11, value=1, step=1, label="Spin")
                     with gr.Row():
-                        steps_md = gr.Slider(minimum=10, maximum=2000, value=100, step=10, label="Steps")
-                        temp_md = gr.Slider(minimum=10, maximum=1500, value=300, step=10, label="Temperature (K)")
+                        steps_md = gr.Slider(minimum=10, maximum=500, value=100, step=10, label="Steps")
+                        temp_md = gr.Slider(minimum=10, maximum=1000, value=300, step=10, label="Temperature (K)")
                     with gr.Row():
                         timestep_md = gr.Slider(minimum=0.1, maximum=5.0, value=1.0, step=0.1, label="Timestep (fs)")
                         ensemble_md = gr.Radio(["NVE", "NVT"], value="NVE", label="Ensemble")
@@ -317,112 +338,20 @@ 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_debug = gr.Textbox(label="Debug Info", interactive=False, max_lines=3)
                     
-                    # Molecule3D viewer para MD
-                    md_viewer = Molecule3D(
-                        label="Final Structure Viewer (Last MD Frame)",
-                        reps=[
-                            {
-                                "model": 0,
-                                "chain": "",
-                                "resname": "",
-                                "style": "stick",
-                                "color": "whiteCarbon",
-                                "residue_range": "",
-                                "around": 0,
-                                "byres": False,
-                                "visible": True,
-                                "opacity": 1.0
-                            },
-                            {
-                                "model": 0,
-                                "chain": "",
-                                "resname": "",
-                                "style": "sphere",
-                                "color": "whiteCarbon",
-                                "residue_range": "",
-                                "around": 0,
-                                "byres": False,
-                                "visible": True,
-                                "opacity": 0.7
-                            }
-                        ]
-                    )
-                    
-                    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()
+                    # Conditional viewer
+                    if MOLECULE3D_AVAILABLE:
+                        md_viewer = Molecule3D(label="Final Structure (Molecule3D)")
+                        md_outputs = [md_status, md_traj, gr.Textbox(visible=False), gr.Code(visible=False), gr.Markdown(visible=False), md_viewer, md_debug]
+                    else:
+                        md_viewer = gr.HTML(label="Trajectory Animation (3Dmol.js)")
+                        md_outputs = [md_status, md_traj, gr.Textbox(visible=False), gr.Code(visible=False), gr.Markdown(visible=False), md_viewer, md_debug]
 
             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],
-            )
-
-        # -------- Relax --------
-        with gr.Tab("Relaxation / Optimization"):
-            with gr.Row():
-                with gr.Column(scale=2):
-                    gr.Markdown("## Structure Relaxation/Optimization")
-                    gr.Markdown("Upload your molecular structure for geometry optimization.")
-                    
-                    xyz_rlx = gr.File(
-                        label="Upload Structure File (.xyz/.pdb/.cif/.traj)",
-                        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():
-                        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")
-                    relax_cell = gr.Checkbox(False, label="Relax Unit Cell")
-                    run_rlx_btn = gr.Button("Run Optimization", 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)
-                    
-                    # Molecule3D viewer para Relaxation
-                    rlx_viewer = Molecule3D(
-                        label="Optimized Structure Viewer",
-                        reps=[
-                            {
-                                "model": 0,
-                                "chain": "",
-                                "resname": "",
-                                "style": "stick",
-                                "color": "whiteCarbon",
-                                "residue_range": "",
-                                "around": 0,
-                                "byres": False,
-                                "visible": True,
-                                "opacity": 1.0
-                            },
-                            {
-                                "model": 0,
-                                "chain": "",
-                                "resname": "",
-                                "style": "sphere",
-                                "color": "whiteCarbon",
-                                "residue_range": "",
-                                "around": 0,
-                                "byres": False,
-                                "visible": True,
-                                "opacity": 0.7
-                            }
-                        ]
-                    )
-                    
-                    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()
-
-            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],
+                outputs=md_outputs
             )
 
 print("Starting OrbMol model loading…")