Spaces:
Sleeping
Sleeping
| import os | |
| os.environ["TORCH_DYNAMO_DISABLE"] = "1" | |
| import tempfile | |
| import numpy as np | |
| import gradio as gr | |
| from ase.io import read, write | |
| from ase.io.trajectory import Trajectory | |
| 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}") | |
| # ==== 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): | |
| return None, "No trajectory file" | |
| try: | |
| # Step 1: Read trajectory | |
| traj = Trajectory(traj_path) | |
| print(f"Trajectory frames: {len(traj)}") | |
| if len(traj) == 0: | |
| return None, "Empty trajectory" | |
| # Step 2: Get atoms | |
| atoms = traj[-1] | |
| print(f"Atoms count: {len(atoms)}") | |
| print(f"Chemical symbols: {atoms.get_chemical_symbols()[:5]}...") # First 5 | |
| # 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)}") | |
| # 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}") | |
| # Step 5: Write PDB | |
| write(pdb_path, atoms, format='pdb') | |
| print(f"PDB written, exists: {os.path.exists(pdb_path)}") | |
| if os.path.exists(pdb_path): | |
| file_size = os.path.getsize(pdb_path) | |
| print(f"PDB file size: {file_size} bytes") | |
| # Read first few lines | |
| with open(pdb_path, 'r') as f: | |
| first_lines = f.read(300) | |
| print(f"PDB content preview:\n{first_lines}") | |
| return pdb_path, f"Success: {file_size} bytes" | |
| else: | |
| return None, "PDB file not created" | |
| except Exception as e: | |
| print(f"Exception in debug_pdb_creation: {e}") | |
| import traceback | |
| traceback.print_exc() | |
| return None, f"Error: {e}" | |
| # ==== 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 | |
| # ==== OrbMol SPE ==== | |
| from orb_models.forcefield import pretrained | |
| from orb_models.forcefield.calculator import ORBCalculator | |
| _MODEL_CALC = None | |
| def _load_orbmol_calc(): | |
| global _MODEL_CALC | |
| if _MODEL_CALC is None: | |
| orbff = pretrained.orb_v3_conservative_inf_omat( | |
| device="cpu", precision="float32-high" | |
| ) | |
| _MODEL_CALC = ORBCalculator(orbff, device="cpu") | |
| return _MODEL_CALC | |
| def predict_molecule(structure_file, charge=0, spin_multiplicity=1): | |
| """Single Point Energy + fuerzas (OrbMol)""" | |
| try: | |
| calc = _load_orbmol_calc() | |
| if not structure_file: | |
| 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 if MOLECULE3D_AVAILABLE else "" | |
| if os.path.getsize(file_path) == 0: | |
| 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)} | |
| atoms.calc = calc | |
| energy = atoms.get_potential_energy() | |
| forces = atoms.get_forces() | |
| lines = [f"Total Energy: {energy:.6f} eV", "", "Atomic Forces:"] | |
| for i, fc in enumerate(forces): | |
| lines.append(f"Atom {i+1}: [{fc[0]:.4f}, {fc[1]:.4f}, {fc[2]:.4f}] eV/Å") | |
| max_force = float(np.max(np.linalg.norm(forces, axis=1))) | |
| lines += ["", f"Max Force: {max_force:.4f} eV/Å"] | |
| # 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: | |
| error_msg = f"Error during calculation: {e}" | |
| return error_msg, "Error", str(e), None if MOLECULE3D_AVAILABLE else "" | |
| # ==== Simulaciones ==== | |
| from simulation_scripts_orbmol import ( | |
| run_md_simulation, | |
| run_relaxation_simulation, | |
| ) | |
| 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 "", "") | |
| file_path = structure_file | |
| print(f"MD Wrapper: Processing {file_path}") | |
| traj_path, log_text, script_text, explanation = run_md_simulation( | |
| file_path, | |
| int(steps), | |
| 20, | |
| float(timestep_fs), | |
| float(tempK), | |
| "NVT" if ensemble == "NVT" else "NVE", | |
| int(charge), | |
| int(spin), | |
| ) | |
| status = f"MD completed: {int(steps)} steps at {int(tempK)} K ({ensemble})" | |
| print(f"MD completed, trajectory: {traj_path}") | |
| 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 (error_msg, None, "", "", "", None if MOLECULE3D_AVAILABLE else "", str(e)) | |
| # ==== 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.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.") | |
| 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") | |
| 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, 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] | |
| 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") | |
| xyz_md = gr.File( | |
| 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") | |
| 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)") | |
| 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") | |
| run_md_btn = gr.Button("Run MD Simulation", variant="primary") | |
| 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] | |
| run_md_btn.click( | |
| md_wrapper, | |
| inputs=[xyz_md, charge_md, spin_md, steps_md, temp_md, timestep_md, ensemble_md], | |
| outputs=md_outputs | |
| ) | |
| 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) |