Update app.py

app.py

@@ -5,145 +5,106 @@ 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
 
-# ==== 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}")
-
-# ==== PDB creation for Molecule3D ====
-def create_pdb_for_molecule3d(atoms):
-    """Create a PDB file specifically for Molecule3D viewer"""
+# ==== util: PDB writer robusto para Molecule3D ====
+def _pdb_cache_path(prefix: str, key: str) -> str:
+    gradio_cache_dir = os.path.join(tempfile.gettempdir(), "gradio")
+    os.makedirs(gradio_cache_dir, exist_ok=True)
+    return os.path.join(gradio_cache_dir, f"{prefix}_{key}.pdb")
+
+def _write_pdb_with_fallback(atoms, pdb_path: str) -> str | None:
+    """
+    Intenta escribir PDB con ASE usando el id correcto ('proteindatabank').
+    Si falla, crea un PDB minimal válido (líneas ATOM + END) suficiente para gradio_molecule3d.
+    Devuelve la ruta si existe y tiene tamaño > 0, en caso contrario None.
+    """
     try:
-        # Create gradio cache directory  
-        cache_dir = os.path.join(tempfile.gettempdir(), "gradio")
-        os.makedirs(cache_dir, exist_ok=True)
-        print(f"Cache directory: {cache_dir}")
-        
-        # Create unique filename based on atomic positions
-        pos_hash = hashlib.md5(atoms.get_positions().tobytes()).hexdigest()[:12]
-        pdb_filename = f"mol_{pos_hash}.pdb"
-        pdb_path = os.path.join(cache_dir, pdb_filename)
-        print(f"Creating PDB: {pdb_path}")
-        
-        # Write PDB file using ASE
-        write(pdb_path, atoms, format='pdb')
-        
-        # Verify file was created and has content
-        if os.path.exists(pdb_path):
-            file_size = os.path.getsize(pdb_path)
-            print(f"PDB created successfully: {file_size} bytes")
-            
-            # Read and log first few lines for verification
-            with open(pdb_path, 'r') as f:
-                first_lines = f.read(200)
-                print(f"PDB preview: {first_lines[:100]}...")
-            
-            if file_size > 0:
-                return pdb_path
-        
-        print("PDB file was not created or is empty")
+        # writer id correcto en ASE para PDB
+        write(pdb_path, atoms, format="proteindatabank")
+    except Exception as e:
+        print(f"ASE PDB writer failed: {e} — trying manual minimal PDB")
+        try:
+            xyz = atoms.get_positions()
+            syms = atoms.get_chemical_symbols()
+            with open(pdb_path, "w") as f:
+                serial = 1
+                for (sym, (x, y, z)) in zip(syms, xyz):
+                    # Campos esenciales: 'ATOM', serial, nombre (símbolo), resname, chain, res seq, coords y elemento
+                    f.write(
+                        f"ATOM  {serial:5d} {sym:<3s} MOL A   1    "
+                        f"{x:8.3f}{y:8.3f}{z:8.3f}  1.00  0.00          {sym:>2s}\n"
+                    )
+                    serial += 1
+                f.write("END\n")
+        except Exception as e2:
+            print(f"Manual PDB fallback failed: {e2}")
+            return None
+
+    # Validación
+    if os.path.exists(pdb_path) and os.path.getsize(pdb_path) > 0:
+        print(f"✅ PDB created: {pdb_path} ({os.path.getsize(pdb_path)} bytes)")
+        try:
+            with open(pdb_path, "r") as f:
+                preview = f.read(160)
+            print(f"PDB preview: {preview[:120]}...")
+        except Exception:
+            pass
+        return pdb_path
+    print("❌ PDB file not created or empty")
+    return None
+
+# ==== Molecule3D viewer preparation CORREGIDO ====
+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):
+        print("No trajectory path provided or file doesn't exist")
         return None
-            
+    try:
+        traj = Trajectory(traj_path)
+        if len(traj) == 0:
+            print("Empty trajectory")
+            return None
+
+        print(f"Preparing viewer: {len(traj)} frames, {len(traj[-1])} atoms")
+        atoms = traj[-1]
+
+        # clave única por posiciones del último frame
+        pos_hash = hashlib.md5(atoms.get_positions().tobytes()).hexdigest()[:12]
+        pdb_path = _pdb_cache_path("molecule", pos_hash)
+
+        if os.path.exists(pdb_path) and os.path.getsize(pdb_path) > 0:
+            print(f"PDB already exists: {pdb_path}")
+            return pdb_path
+
+        return _write_pdb_with_fallback(atoms, pdb_path)
+
     except Exception as e:
-        print(f"Error creating PDB for Molecule3D: {e}")
-        import traceback
-        traceback.print_exc()
+        print(f"Error preparing molecule for viewer: {e}")
+        import traceback; traceback.print_exc()
         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"""
-    
-    # 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;
-        }}
-        
-        if (typeof $3Dmol === 'undefined') {{
-            console.error('$3Dmol is undefined');
-            return;
-        }}
-        
-        console.log('Creating 3Dmol viewer...');
-        var viewer = $3Dmol.createViewer(element, {{
-            backgroundColor: 'white',
-            antialias: true
-        }});
-        
-        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'}}
-        }});
-        
-        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
+# ==== 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
+    try:
+        atoms = read(structure_file)
+        key = hashlib.md5(str(structure_file).encode()).hexdigest()[:12]
+        pdb_path = _pdb_cache_path("input", key)
+
+        if os.path.exists(pdb_path) and os.path.getsize(pdb_path) > 0:
+            print(f"Input PDB already exists: {pdb_path}")
+            return pdb_path
+
+        return _write_pdb_with_fallback(atoms, pdb_path)
+
+    except Exception as e:
+        print(f"Error preparing input for viewer: {e}")
+        return None
 
 # ==== OrbMol SPE ====
 from orb_models.forcefield import pretrained
@@ -159,47 +120,20 @@ def _load_orbmol_calc():
         _MODEL_CALC = ORBCalculator(orbff, device="cpu")
     return _MODEL_CALC
 
-def create_pdb_for_molecule3d(atoms):
-    """Create a PDB file specifically for Molecule3D viewer"""
-    try:
-        # Create gradio cache directory
-        cache_dir = os.path.join(tempfile.gettempdir(), "gradio")
-        os.makedirs(cache_dir, exist_ok=True)
-        
-        # Create unique filename
-        unique_id = hashlib.md5(str(atoms.get_positions().tostring()).encode()).hexdigest()[:12]
-        pdb_filename = f"mol_{unique_id}.pdb"
-        pdb_path = os.path.join(cache_dir, pdb_filename)
-        
-        # Write PDB file
-        write(pdb_path, atoms, format='pdb')
-        
-        # Verify file was created
-        if os.path.exists(pdb_path) and os.path.getsize(pdb_path) > 0:
-            print(f"✅ PDB created for Molecule3D: {pdb_path} ({os.path.getsize(pdb_path)} bytes)")
-            return pdb_path
-        else:
-            print("❌ PDB file was not created properly")
-            return None
-            
-    except Exception as e:
-        print(f"❌ Error creating PDB: {e}")
-        return None
-
 def predict_molecule(structure_file, charge=0, spin_multiplicity=1):
-    """Single Point Energy + fuerzas (OrbMol)"""
+    """
+    Single Point Energy + fuerzas (OrbMol). Acepta archivos subidos.
+    """
     try:
         calc = _load_orbmol_calc()
         if not structure_file:
-            return "Error: Please upload a structure file", "Error", "", None if MOLECULE3D_AVAILABLE else ""
+            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", "", None if MOLECULE3D_AVAILABLE else ""
-        
+            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", "", None if MOLECULE3D_AVAILABLE else ""
+            return f"Error: Empty file: {file_path}", "Error", None
 
         atoms = read(file_path)
         atoms.info = {"charge": int(charge), "spin": int(spin_multiplicity)}
@@ -214,37 +148,25 @@ 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/Å"]
 
-        # Create debug info
-        symbols = atoms.get_chemical_symbols()
-        debug_info = f"Molecule: {len(atoms)} atoms\nElements: {', '.join(set(symbols))}\nFormula: {''.join(sorted(symbols))}"
-        
-        # Create viewer
-        if MOLECULE3D_AVAILABLE:
-            pdb_path = create_pdb_for_molecule3d(atoms)
-            if pdb_path:
-                debug_info += f"\nPDB created: ✅ {os.path.basename(pdb_path)}"
-            else:
-                debug_info += f"\nPDB creation: ❌ Failed"
-            return "\n".join(lines), "Calculation completed with OrbMol", debug_info, 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
+        # Preparar PDB para visualización
+        pdb_file = prepare_input_for_viewer(file_path)
+
+        return "\n".join(lines), "Calculation completed with OrbMol", pdb_file
     except Exception as e:
-        error_msg = f"Error during calculation: {e}"
-        import traceback
-        traceback.print_exc()
-        return error_msg, "Error", str(e), None if MOLECULE3D_AVAILABLE else ""
+        import traceback; traceback.print_exc()
+        return f"Error during calculation: {e}", "Error", None
 
-# ==== Simulaciones ====
+# ==== Simulaciones (helpers) ====
 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 if MOLECULE3D_AVAILABLE else "", "")
+            return ("Error: Please upload a structure file", None, "", "", "", None)
 
         file_path = structure_file
         print(f"MD Wrapper: Processing {file_path}")
@@ -252,7 +174,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,
+            20,  # pre-relax steps
             float(timestep_fs),
             float(tempK),
             "NVT" if ensemble == "NVT" else "NVE",
@@ -260,71 +182,62 @@ def md_wrapper(structure_file, charge, spin, steps, tempK, timestep_fs, ensemble
             int(spin),
         )
         status = f"MD completed: {int(steps)} steps at {int(tempK)} K ({ensemble})"
-
         print(f"MD completed, trajectory: {traj_path}")
-        
-        if MOLECULE3D_AVAILABLE:
-            # Read last frame from trajectory
-            try:
-                traj = Trajectory(traj_path)
-                last_atoms = traj[-1]
-                print(f"Read trajectory: {len(traj)} frames, last frame has {len(last_atoms)} atoms")
-                
-                pdb_path = create_pdb_for_molecule3d(last_atoms)
-                
-                debug_info = f"Trajectory: {len(traj)} frames\nLast frame: {len(last_atoms)} atoms"
-                if pdb_path:
-                    debug_info += f"\nPDB: ✅ {os.path.basename(pdb_path)}"
-                else:
-                    debug_info += f"\nPDB: ❌ Creation failed"
-                    
-                return (status, traj_path, log_text, script_text, explanation, pdb_path, debug_info)
-                
-            except Exception as traj_error:
-                print(f"Error reading trajectory: {traj_error}")
-                debug_info = f"Trajectory error: {traj_error}"
-                return (status, traj_path, log_text, script_text, explanation, None, 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)
+
+        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:
-        error_msg = f"Error: {e}"
         print(f"MD Wrapper Error: {e}")
-        import traceback
-        traceback.print_exc()
-        return (error_msg, None, "", "", "", None if MOLECULE3D_AVAILABLE else "", str(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}")
+
+        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)
 
 # ==== UI ====
-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.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 — Single Point Energy")
-                    gr.Markdown("Upload molecular structure files for energy and force calculations.")
+                    gr.Markdown("# OrbMol — Quantum-Accurate Molecular Predictions")
+                    gr.Markdown("Upload molecular structure files (.xyz, .pdb, .cif, .traj) for energy and force calculations.")
                     
                     xyz_input = gr.File(
                         label="Upload Structure File (.xyz/.pdb/.cif/.traj)",
                         file_types=[".xyz", ".pdb", ".cif", ".traj", ".mol", ".sdf"],
                         file_count="single"
                     )
-                    
                     with gr.Row():
                         charge_input = gr.Slider(minimum=-10, maximum=10, value=0, step=1, label="Charge")
                         spin_input = gr.Slider(minimum=1, maximum=11, value=1, step=1, label="Spin Multiplicity")
@@ -333,36 +246,37 @@ with gr.Blocks(theme=gr.themes.Ocean(), title="OrbMol Demo - Debug") 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)
-                    
-                    # 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]
+                    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
+                            }
+                        ]
+                    )
 
-            run_spe.click(predict_molecule, [xyz_input, charge_input, spin_input], outputs)
+            run_spe.click(predict_molecule, [xyz_input, 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")
+                    gr.Markdown("Upload your molecular structure and configure MD parameters.")
                     
                     xyz_md = gr.File(
-                        label="Upload Structure File",
-                        file_types=[".xyz", ".pdb", ".cif", ".traj"],
+                        label="Upload Structure File (.xyz/.pdb/.cif/.traj)",
+                        file_types=[".xyz", ".pdb", ".cif", ".traj", ".mol", ".sdf"],
                         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")
+                        spin_md = gr.Slider(minimum=1, maximum=11, value=1, step=1, label="Spin Multiplicity")
                     with gr.Row():
-                        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)")
+                        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)")
                     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")
@@ -371,24 +285,81 @@ with gr.Blocks(theme=gr.themes.Ocean(), title="OrbMol Demo - Debug") 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)
-                    
-                    # 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]
+                    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()
 
             run_md_btn.click(
                 md_wrapper,
                 inputs=[xyz_md, charge_md, spin_md, steps_md, temp_md, timestep_md, ensemble_md],
-                outputs=md_outputs
+                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)
+                    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],
             )
 
 print("Starting OrbMol model loading…")
 _ = _load_orbmol_calc()
 
 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)