Spaces:

hugging-apps
/

aceff-2

Running on Zero

Upload folder using huggingface_hub

bf92d6c verified 7 days ago

18.6 kB

"""AceFF-2: interactive machine-learning potential for small molecules.
Enter a molecule (XYZ block, PDB, or SMILES), pick a net charge, and run a
single-point calculation (energy + per-atom forces) and optionally an ASE LBFGS
geometry relaxation with the AceFF-2 TensorNet-2 potential
(https://huggingface.co/papers/2601.00581, Acellera, Apache-2.0).
"""

import json
import os
import time
import traceback

# Keep warp JIT kernels out of the home dir
os.environ.setdefault("WARP_CACHE_DIR", "/tmp/warp_cache")

import spaces  # must be first (before any CUDA-touching import)
import gradio as gr

MODELS = {
    "AceFF-2.0 (H B C N O F Si P S Cl Br I)": ("Acellera/AceFF-2.0", "aceff_v2.0.ckpt"),
    "AceFF-2-OMOL (adds Na, K; OMol25 data)": ("Acellera/AceFF-2-OMOL", "aceff_v2_omol.ckpt"),
}
SUPPORTED_Z = {1, 5, 6, 7, 8, 9, 14, 15, 16, 17, 35, 53}  # AceFF-2 elements
SUPPORTED_Z_OMOL = SUPPORTED_Z | {11, 19}
EV_TO_KCAL_MOL = 23.0605
MAX_ATOMS = 60
DEFAULT_MODEL = list(MODELS)[0]

_state = {}  # per-GPU-worker cache: models + ASE calculators


def _sym_z(symbol):
    from ase.data import atomic_numbers

    return int(atomic_numbers.get(symbol.title(), 0))


def _get_model(model_key):
    """Load (and cache) an AceFF checkpoint on the GPU worker via load_model."""
    if "model" in _state.get(model_key, {}):
        return _state[model_key]["model"]
    import torch
    from huggingface_hub import hf_hub_download
    from torchmdnet.models.model import load_model

    repo_id, filename = MODELS[model_key]
    path = hf_hub_download(repo_id=repo_id, filename=filename)
    model = load_model(path, derivative=False, remove_ref_energy=True)
    model.eval()
    for p in model.parameters():
        p.requires_grad = False
    model.to("cuda")
    _state.setdefault(model_key, {})["model"] = model
    return model


def _get_calculator(model_key):
    """Load (and cache) the TorchMD-Net ASE calculator on the GPU worker."""
    calcs = _state.setdefault(model_key, {})
    if "calc" in calcs:
        return calcs["calc"]
    from huggingface_hub import hf_hub_download
    from torchmdnet.calculators import TMDNETCalculator

    repo_id, filename = MODELS[model_key]
    calc = TMDNETCalculator(
        hf_hub_download(repo_id=repo_id, filename=filename), device="cuda"
    )
    calcs["calc"] = calc
    return calc


def _energy_forces(model, z, pos, q):
    """Single-point energy + forces: fresh graph, autograd.grad each call.
    The checkpoint is loaded with derivative=False so the model returns only
    the energy; forces come from torch.autograd.grad, which is repeatable
    across calls (a derivative=True model frees its graph on .backward()).
    """
    import torch

    pos = pos.detach().clone().requires_grad_(True)
    energy, _ = model(z, pos, q=q)
    grad = torch.autograd.grad(energy, pos)[0]
    return float(energy.item()), (-grad).detach().cpu().numpy()


def _parse_molecule(mol_text):
    """Parse XYZ / PDB / SMILES input into (symbols, positions Nx3 ndarray)."""
    import numpy as np
    from ase.io import read as ase_read

    raw = mol_text or ""
    if not raw.strip():
        raise ValueError(
            "Please provide a molecule: paste an XYZ block or PDB, or type a SMILES string."
        )
    nonempty = [l for l in raw.splitlines() if l.strip()]
    all_lines = raw.splitlines()

    is_pdb = any(l.startswith(("ATOM", "HETATM")) for l in nonempty)
    n_first = None
    try:
        n_first = int(all_lines[0].strip().split()[0])
    except (ValueError, IndexError):
        pass
    counts_ok = False
    if n_first is not None and n_first > 0:
        coord_lines = [
            l for l in all_lines[2:] if l.strip() and len(l.split()) >= 4
        ]
        counts_ok = len(coord_lines) >= n_first and len(all_lines) >= n_first + 2
    is_xyz = counts_ok

    if not is_xyz and not is_pdb:
        if len(nonempty) == 1 and len(nonempty[0]) < 120:
            return _smiles_to_atoms(nonempty[0])
        raise ValueError(
            "Input must be an XYZ block (atom count on line 1, comment on line 2), "
            "a PDB, or a SMILES string on a single line."
        )

    import tempfile

    with tempfile.NamedTemporaryFile(
        "w", suffix=".pdb" if is_pdb else ".xyz", delete=False
    ) as f:
        f.write(raw + "\n")
        path = f.name
    try:
        atoms = ase_read(path)
    finally:
        os.unlink(path)
    symbols = list(atoms.get_chemical_symbols())
    pos = np.array(atoms.get_positions(), dtype=np.float32)
    if len(symbols) > 1 and float(pos.std()) < 1e-6:
        raise ValueError(
            "All atoms sit at the same position — is this a valid 3D structure?"
        )
    return symbols, pos


def _smiles_to_atoms(smiles):
    """Embed a 3D geometry from SMILES with RDKit (ETKDG)."""
    import numpy as np
    from rdkit import Chem
    from rdkit.Chem import AllChem

    mol = Chem.MolFromSmiles(smiles)
    if mol is None:
        raise ValueError(f"Could not parse input as SMILES: {smiles!r}")
    mol = Chem.AddHs(mol)
    status = AllChem.EmbedMultipleConfs(
        mol, numConfs=1, useRandomCoords=True, randomSeed=42
    )
    confs = list(status)
    if not confs:
        raise ValueError(f"RDKit could not generate a 3D geometry for: {smiles}")
    conf = mol.GetConformer(confs[0])
    pos = np.array(conf.GetPositions(), dtype=np.float32)
    symbols = [a.GetSymbol() for a in mol.GetAtoms()]
    return symbols, pos


def _xyz_block(symbols, pos, comment=""):
    import numpy as np

    lines = [str(len(symbols)), comment]
    for s, (x, y, z) in zip(symbols, np.asarray(pos)):
        lines.append(f"{s:<2s} {x:>12.6f} {y:>12.6f} {z:>12.6f}")
    return "\n".join(lines)


def _viewer_html(title, xyz_in, xyz_out, arrows, arrow_scale_label):
    """Build the 3Dmol.js viewer as an escaped srcdoc iframe.
    Gradio injects gr.HTML values via innerHTML, and browsers refuse to run
    <script> tags inserted that way (gradio warns about it at runtime), so a
    plain inline script would leave a blank box. An <iframe srcdoc="...">
    is a real document parsed by the browser: its scripts execute normally,
    and the iframe reloads whenever the value updates, so the viewer re-renders
    after every calculation.
    """
    import html as html_mod

    has_out = bool(xyz_out)
    doc = f"""<!DOCTYPE html>
<html>
<head>
<meta charset="utf-8">
<script src="https://3Dmol.org/build/3Dmol-min.js"></script>
<style>
  html, body {{ margin: 0; padding: 0; height: 100%; background: #17191c; overflow: hidden;
               font-family: system-ui, -apple-system, sans-serif; }}
  #label {{ position: absolute; top: 8px; left: 10px; right: 10px; color: #e8eaed;
            font-size: 13px; z-index: 10; pointer-events: none; }}
  #label b {{ font-size: 14px; }}
  #label span {{ opacity: 0.7; }}
  #v {{ position: absolute; inset: 0; }}
</style>
</head>
<body>
<div id="label"><b>{title}</b>
  <span> &nbsp;· rotate: drag &nbsp;· zoom: scroll &nbsp;· translate: right-drag &nbsp;· red arrows = per-atom force vectors (&times;{arrow_scale_label})</span></div>
<div id="v"></div>
<script>
var XYZ_IN = {json.dumps(xyz_in)};
var XYZ_OUT = {json.dumps(xyz_out) if has_out else "null"};
var ARROWS = {json.dumps(arrows)};
function boot(attempts) {{
  var el = document.getElementById('v');
  if (typeof $3Dmol === 'undefined') {{
    if (attempts > 0) setTimeout(function () {{ boot(attempts - 1); }}, 400);
    return;
  }}
  var viewer = $3Dmol.createViewer({{el: el, backgroundColor: '#17191c'}});
  var mIn = viewer.addModel(XYZ_IN, 'xyz');
  mIn.setStyle({{}}, {{stick: {{radius: 0.14, colorscheme: 'cyanCarbon'}},
                      sphere: {{scale: 0.2, colorscheme: 'cyanCarbon'}}}});
  if (XYZ_OUT) {{
    var mOut = viewer.addModel(XYZ_OUT, 'xyz');
    mOut.setStyle({{}}, {{stick: {{radius: 0.10, opacity: 0.35}}}});
  }}
  for (var i = 0; i < ARROWS.length; i++) {{
    var a = ARROWS[i];
    if (a.s <= 0.01) continue;
    viewer.addArrow({{start: {{x: a.p0[0], y: a.p0[1], z: a.p0[2]}},
                      end:   {{x: a.p1[0], y: a.p1[1], z: a.p1[2]}},
                      radius: 0.05, color: 'red'}});
  }}
  viewer.zoomTo();
  viewer.render();
}}
boot(25);
</script>
</body>
</html>"""
    escaped = html_mod.escape(doc, quote=True)
    return (
        f'<iframe srcdoc="{escaped}" title="3Dmol viewer" '
        f'style="width:100%;height:450px;border:0;border-radius:8px;display:block;" '
        f'allowfullscreen></iframe>'
    )


def _fmt_metrics(e_ev, de_ev, fmax_ev, relaxed):
    import numpy as np

    fmax_kcal = fmax_ev * EV_TO_KCAL_MOL
    items = [
        ("Energy", f"{e_ev:.4f} eV ({e_ev * EV_TO_KCAL_MOL:,.1f} kcal/mol)"),
    ]
    if relaxed:
        items.append(("Δ after relaxation", f"{de_ev:+.4f} eV ({de_ev * EV_TO_KCAL_MOL:+,.1f} kcal/mol)"))
    items.append(("Max |F|", f"{fmax_ev:.3f} eV/Å ({fmax_kcal:.1f} kcal/mol/Å)"))
    return items


@spaces.GPU(duration=30)
def run(
    molecule: str,
    net_charge: float = 0,
    do_relax: bool = False,
    relax_fmax: float = 0.05,
    relax_steps: int = 200,
    model_choice: str = DEFAULT_MODEL,
    show_forces: bool = True,
):
    """Run an AceFF-2 single-point calculation (energy + forces) and optional
    geometry relaxation for one small molecule.
    Args:
        molecule: XYZ block, PDB text, or a SMILES string.
        net_charge: total molecular charge (-2..+2 supported by the model).
        do_relax: also relax the geometry with ASE LBFGS at this potential.
        relax_fmax: convergence threshold on max atomic force (eV/Å).
        relax_steps: max optimizer steps.
        model_choice: which AceFF-2 checkpoint to use.
        show_forces: draw per-atom force arrows on the 3D viewer.
    Returns:
        (status markdown, metric items, viewer HTML, relaxed XYZ text)
    """
    t_start = time.perf_counter()
    try:
        import numpy as np
        import torch
        from ase import Atoms
        from ase.optimize import LBFGS

        model_key = model_choice if model_choice in MODELS else DEFAULT_MODEL
        supported = SUPPORTED_Z_OMOL if "OMOL" in model_key else SUPPORTED_Z

        # ---- 1. parse + validate input -------------------------------------
        symbols, pos = _parse_molecule(molecule)
        if len(symbols) > MAX_ATOMS:
            raise ValueError(
                f"Molecule has {len(symbols)} atoms; this demo caps at {MAX_ATOMS}."
            )
        znums = [_sym_z(s) for s in symbols]
        bad = sorted({s for s, zn in zip(symbols, znums) if zn not in supported})
        if bad:
            elems = (
                "H, B, C, N, O, F, Na, Si, P, S, Cl, K, Br, I"
                if "OMOL" in model_key
                else "H, B, C, N, O, F, Si, P, S, Cl, Br, I"
            )
            raise ValueError(
                f"Unsupported element(s) {', '.join(bad)}. {model_key} was trained on: {elems}."
            )
        if not -5 <= float(net_charge) <= 5:
            raise ValueError("Net charge must be an integer between -2 and +2.")
        if abs(float(net_charge) - round(float(net_charge))) > 1e-9:
            raise ValueError("Net charge must be an integer (e.g. 0, -1, +2).")
        q_float = float(net_charge)
        if not -2 <= q_float <= 2:
            raise ValueError(
                "AceFF-2 was trained on total molecular charges from -2 to +2 only."
            )

        # ---- 2. single point E + F (reference path: forward + backward) ----
        model = _get_model(model_key)
        z = torch.tensor(znums, dtype=torch.long, device="cuda")
        p = torch.tensor(np.asarray(pos), dtype=torch.float32, device="cuda")
        q = torch.tensor([q_float], dtype=torch.float32, device="cuda")
        e_init_ev, f_init = _energy_forces(model, z, p, q)

        xyz_in = _xyz_block(symbols, pos, "input geometry")
        xyz_out = None
        steps_used = 0
        e_final_ev = e_init_ev
        fmax_final = float(np.abs(f_init).max()) if len(f_init) else 0.0

        # ---- 3. optional LBFGS relaxation via the ASE calculator -----------
        if do_relax:
            atoms = Atoms(symbols=symbols, positions=np.asarray(pos))
            atoms.info["charge"] = q_float
            atoms.calc = _get_calculator(model_key)
            dyn = LBFGS(atoms, logfile=None)
            dyn.run(fmax=float(relax_fmax), steps=int(relax_steps))
            steps_used = dyn.get_number_of_steps()
            e_final_ev = float(atoms.get_potential_energy())
            fmax_final = float(np.abs(atoms.get_forces()).max())
            xyz_out = _xyz_block(
                symbols, atoms.get_positions(), "relaxed geometry (AceFF-2, LBFGS)"
            )

        de_ev = e_final_ev - e_init_ev
        metrics = _fmt_metrics(e_final_ev, de_ev, fmax_final, do_relax)

        # ---- 4. viewer: arrows scale longest force to ~0.55 Å -------------
        fmax_init = float(np.abs(f_init).max()) if len(f_init) else 0.0
        scale = 0.55 / fmax_init if fmax_init > 1e-9 else 0.0
        disp_pos = np.asarray(pos)
        arrows = []
        if show_forces and scale > 0:
            for (x, y, zw), (fx, fy, fz) in zip(disp_pos, f_init):
                mag = float(np.linalg.norm([fx, fy, fz]))
                arrows.append(
                    {
                        "p0": [round(float(x), 3), round(float(y), 3), round(float(zw), 3)],
                        "p1": [
                            round(float(x + fx * scale), 3),
                            round(float(y + fy * scale), 3),
                            round(float(zw + fz * scale), 3),
                        ],
                        "s": mag,
                    }
                )
        title = (
            "Input geometry (solid) → relaxed geometry (faded)"
            if do_relax
            else "Input geometry"
        )
        viewer = _viewer_html(title, xyz_in, xyz_out, arrows, f"{1.0/scale:.2f} Å/(eV/Å)" if scale else "—")

        mode = "single point + relaxation" if do_relax else "single point"
        q_str = f"{int(q_float):+d}" if q_float == int(q_float) else f"{q_float:+g}"
        status = (
            f"**Done** — {mode} · {len(symbols)} atoms · charge {q_str}"
            + (f" · {steps_used} LBFGS steps" if do_relax else "")
            + f" · {time.perf_counter() - t_start:.1f}s"
        )
        return status, metrics, viewer, xyz_out or ""
    except Exception as e:
        print("run() error:\n", traceback.format_exc(limit=6))
        raise gr.Error(f"{type(e).__name__}: {e}")


with gr.Blocks(title="AceFF-2 · ML potential for small molecules") as demo:
    gr.Markdown(
        """
# AceFF-2 — Machine Learning Potential for Small Molecules
**AceFF** ([paper](https://huggingface.co/papers/2601.00581)) is a state-of-the-art neural
interatomic potential for drug-like molecules, built on the **TensorNet v2** architecture in
[TorchMD-Net](https://github.com/torchmd/torchmd-net) (Acellera, Apache-2.0).
Give it a molecule — **XYZ**, **PDB**, or **SMILES** — and it predicts the **potential energy**
and **per-atom forces**, and can **relax the geometry** with ASE.
Supported elements: H, B, C, N, O, F, Si, P, S, Cl, Br, I (Na, K with the OMOL checkpoint) ·
net charge −2…+2 · small molecules only (≤ 60 atoms in this demo).
        """
    )
    with gr.Row(elem_id="col-container"):
        with gr.Column(variant="panel", scale=5):
            molecule = gr.Textbox(
                label="Molecule (XYZ block, PDB, or SMILES)",
                placeholder="Paste an XYZ block or PDB, or type a SMILES like CCO…",
                lines=10,
                max_lines=18,
            )
            with gr.Row():
                net_charge = gr.Number(label="Net charge (e)", value=0, precision=0, minimum=-2, maximum=2)
                do_relax = gr.Checkbox(label="Relax geometry (LBFGS)", value=False)
                show_forces = gr.Checkbox(label="Show force arrows", value=True)
            model_choice = gr.Radio(
                choices=list(MODELS),
                value=DEFAULT_MODEL,
                label="Checkpoint",
                info="AceFF-2.0: Acellera training set · AceFF-2-OMOL: public OMol25 data, adds Na/K.",
            )
            with gr.Accordion("Relaxation settings", open=False):
                relax_fmax = gr.Number(label="Convergence: max force (eV/Å)", value=0.05)
                relax_steps = gr.Number(label="Max optimizer steps", value=200, precision=0, minimum=1, maximum=1000)
            run_btn = gr.Button("Compute energy & forces", variant="primary")

        with gr.Column(scale=7):
            status = gr.Markdown("Ready. Pick an example below or paste your own molecule.")
            metrics = gr.HighlightedText(
                label="Results",
                combine_adjacent=False,
            )
            viewer = gr.HTML(
                value="<div style='height:430px;background:#17191c;border-radius:8px;display:flex;"
                "align-items:center;justify-content:center;color:#9aa0a6;font-family:var(--font);'>"
                "3D viewer appears here after a calculation (force arrows in red).</div>"
            )
            with gr.Accordion("Relaxed geometry (XYZ) — copy or download", open=False):
                xyz_out = gr.Textbox(label="XYZ", lines=8, interactive=True)

    gr.Examples(
        examples=[
            ["examples/caffeine.xyz", 0, False, 0.05, 200, DEFAULT_MODEL, True],
            ["examples/aspirin.xyz", 0, True, 0.05, 200, DEFAULT_MODEL, True],
            ["examples/acetate.xyz", -1, True, 0.05, 200, DEFAULT_MODEL, True],
            ["examples/ethanol.xyz", 0, False, 0.05, 200, DEFAULT_MODEL, True],
            ["examples/water.xyz", 0, False, 0.05, 200, DEFAULT_MODEL, True],
        ],
        inputs=[molecule, net_charge, do_relax, relax_fmax, relax_steps, model_choice, show_forces],
        outputs=[status, metrics, viewer, xyz_out],
        fn=run,
        cache_examples=True,
        cache_mode="lazy",
        label="Examples (molecules from the AceFF usage docs; geometries embedded with RDKit)",
    )

    run_btn.click(
        run,
        inputs=[molecule, net_charge, do_relax, relax_fmax, relax_steps, model_choice, show_forces],
        outputs=[status, metrics, viewer, xyz_out],
        api_name="compute",
    )

demo.launch(mcp_server=True, theme=gr.themes.Citrus(),
            css="""
#col-container { max-width: 1280px; margin: 0 auto; }
.dark .gradio-container { color: var(--body-text-color); }
""")

	"""AceFF-2: interactive machine-learning potential for small molecules.

	Enter a molecule (XYZ block, PDB, or SMILES), pick a net charge, and run a
	single-point calculation (energy + per-atom forces) and optionally an ASE LBFGS
	geometry relaxation with the AceFF-2 TensorNet-2 potential
	(https://huggingface.co/papers/2601.00581, Acellera, Apache-2.0).
	"""

	import json
	import os
	import time
	import traceback

	# Keep warp JIT kernels out of the home dir
	os.environ.setdefault("WARP_CACHE_DIR", "/tmp/warp_cache")

	import spaces # must be first (before any CUDA-touching import)
	import gradio as gr

	MODELS = {
	"AceFF-2.0 (H B C N O F Si P S Cl Br I)": ("Acellera/AceFF-2.0", "aceff_v2.0.ckpt"),
	"AceFF-2-OMOL (adds Na, K; OMol25 data)": ("Acellera/AceFF-2-OMOL", "aceff_v2_omol.ckpt"),
	}
	SUPPORTED_Z = {1, 5, 6, 7, 8, 9, 14, 15, 16, 17, 35, 53} # AceFF-2 elements
	SUPPORTED_Z_OMOL = SUPPORTED_Z \| {11, 19}
	EV_TO_KCAL_MOL = 23.0605
	MAX_ATOMS = 60
	DEFAULT_MODEL = list(MODELS)[0]

	_state = {} # per-GPU-worker cache: models + ASE calculators


	def _sym_z(symbol):
	from ase.data import atomic_numbers

	return int(atomic_numbers.get(symbol.title(), 0))


	def _get_model(model_key):
	"""Load (and cache) an AceFF checkpoint on the GPU worker via load_model."""
	if "model" in _state.get(model_key, {}):
	return _state[model_key]["model"]
	import torch
	from huggingface_hub import hf_hub_download
	from torchmdnet.models.model import load_model

	repo_id, filename = MODELS[model_key]
	path = hf_hub_download(repo_id=repo_id, filename=filename)
	model = load_model(path, derivative=False, remove_ref_energy=True)
	model.eval()
	for p in model.parameters():
	p.requires_grad = False
	model.to("cuda")
	_state.setdefault(model_key, {})["model"] = model
	return model


	def _get_calculator(model_key):
	"""Load (and cache) the TorchMD-Net ASE calculator on the GPU worker."""
	calcs = _state.setdefault(model_key, {})
	if "calc" in calcs:
	return calcs["calc"]
	from huggingface_hub import hf_hub_download
	from torchmdnet.calculators import TMDNETCalculator

	repo_id, filename = MODELS[model_key]
	calc = TMDNETCalculator(
	hf_hub_download(repo_id=repo_id, filename=filename), device="cuda"
	)
	calcs["calc"] = calc
	return calc


	def _energy_forces(model, z, pos, q):
	"""Single-point energy + forces: fresh graph, autograd.grad each call.

	The checkpoint is loaded with derivative=False so the model returns only
	the energy; forces come from torch.autograd.grad, which is repeatable
	across calls (a derivative=True model frees its graph on .backward()).
	"""
	import torch

	pos = pos.detach().clone().requires_grad_(True)
	energy, _ = model(z, pos, q=q)
	grad = torch.autograd.grad(energy, pos)[0]
	return float(energy.item()), (-grad).detach().cpu().numpy()


	def _parse_molecule(mol_text):
	"""Parse XYZ / PDB / SMILES input into (symbols, positions Nx3 ndarray)."""
	import numpy as np
	from ase.io import read as ase_read

	raw = mol_text or ""
	if not raw.strip():
	raise ValueError(
	"Please provide a molecule: paste an XYZ block or PDB, or type a SMILES string."
	)
	nonempty = [l for l in raw.splitlines() if l.strip()]
	all_lines = raw.splitlines()

	is_pdb = any(l.startswith(("ATOM", "HETATM")) for l in nonempty)
	n_first = None
	try:
	n_first = int(all_lines[0].strip().split()[0])
	except (ValueError, IndexError):
	pass
	counts_ok = False
	if n_first is not None and n_first > 0:
	coord_lines = [
	l for l in all_lines[2:] if l.strip() and len(l.split()) >= 4
	]
	counts_ok = len(coord_lines) >= n_first and len(all_lines) >= n_first + 2
	is_xyz = counts_ok

	if not is_xyz and not is_pdb:
	if len(nonempty) == 1 and len(nonempty[0]) < 120:
	return _smiles_to_atoms(nonempty[0])
	raise ValueError(
	"Input must be an XYZ block (atom count on line 1, comment on line 2), "
	"a PDB, or a SMILES string on a single line."
	)

	import tempfile

	with tempfile.NamedTemporaryFile(
	"w", suffix=".pdb" if is_pdb else ".xyz", delete=False
	) as f:
	f.write(raw + "\n")
	path = f.name
	try:
	atoms = ase_read(path)
	finally:
	os.unlink(path)
	symbols = list(atoms.get_chemical_symbols())
	pos = np.array(atoms.get_positions(), dtype=np.float32)
	if len(symbols) > 1 and float(pos.std()) < 1e-6:
	raise ValueError(
	"All atoms sit at the same position — is this a valid 3D structure?"
	)
	return symbols, pos


	def _smiles_to_atoms(smiles):
	"""Embed a 3D geometry from SMILES with RDKit (ETKDG)."""
	import numpy as np
	from rdkit import Chem
	from rdkit.Chem import AllChem

	mol = Chem.MolFromSmiles(smiles)
	if mol is None:
	raise ValueError(f"Could not parse input as SMILES: {smiles!r}")
	mol = Chem.AddHs(mol)
	status = AllChem.EmbedMultipleConfs(
	mol, numConfs=1, useRandomCoords=True, randomSeed=42
	)
	confs = list(status)
	if not confs:
	raise ValueError(f"RDKit could not generate a 3D geometry for: {smiles}")
	conf = mol.GetConformer(confs[0])
	pos = np.array(conf.GetPositions(), dtype=np.float32)
	symbols = [a.GetSymbol() for a in mol.GetAtoms()]
	return symbols, pos


	def _xyz_block(symbols, pos, comment=""):
	import numpy as np

	lines = [str(len(symbols)), comment]
	for s, (x, y, z) in zip(symbols, np.asarray(pos)):
	lines.append(f"{s:<2s} {x:>12.6f} {y:>12.6f} {z:>12.6f}")
	return "\n".join(lines)


	def _viewer_html(title, xyz_in, xyz_out, arrows, arrow_scale_label):
	"""Build the 3Dmol.js viewer as an escaped srcdoc iframe.

	Gradio injects gr.HTML values via innerHTML, and browsers refuse to run
	<script> tags inserted that way (gradio warns about it at runtime), so a
	plain inline script would leave a blank box. An <iframe srcdoc="...">
	is a real document parsed by the browser: its scripts execute normally,
	and the iframe reloads whenever the value updates, so the viewer re-renders
	after every calculation.
	"""
	import html as html_mod

	has_out = bool(xyz_out)
	doc = f"""<!DOCTYPE html>
	<html>
	<head>
	<meta charset="utf-8">
	<script src="https://3Dmol.org/build/3Dmol-min.js"></script>
	<style>
	html, body {{ margin: 0; padding: 0; height: 100%; background: #17191c; overflow: hidden;
	font-family: system-ui, -apple-system, sans-serif; }}
	#label {{ position: absolute; top: 8px; left: 10px; right: 10px; color: #e8eaed;
	font-size: 13px; z-index: 10; pointer-events: none; }}
	#label b {{ font-size: 14px; }}
	#label span {{ opacity: 0.7; }}
	#v {{ position: absolute; inset: 0; }}
	</style>
	</head>
	<body>
	<div id="label"><b>{title}</b>
	<span>  · rotate: drag  · zoom: scroll  · translate: right-drag  · red arrows = per-atom force vectors (×{arrow_scale_label})</span></div>
	<div id="v"></div>
	<script>
	var XYZ_IN = {json.dumps(xyz_in)};
	var XYZ_OUT = {json.dumps(xyz_out) if has_out else "null"};
	var ARROWS = {json.dumps(arrows)};
	function boot(attempts) {{
	var el = document.getElementById('v');
	if (typeof $3Dmol === 'undefined') {{
	if (attempts > 0) setTimeout(function () {{ boot(attempts - 1); }}, 400);
	return;
	}}
	var viewer = $3Dmol.createViewer({{el: el, backgroundColor: '#17191c'}});
	var mIn = viewer.addModel(XYZ_IN, 'xyz');
	mIn.setStyle({{}}, {{stick: {{radius: 0.14, colorscheme: 'cyanCarbon'}},
	sphere: {{scale: 0.2, colorscheme: 'cyanCarbon'}}}});
	if (XYZ_OUT) {{
	var mOut = viewer.addModel(XYZ_OUT, 'xyz');
	mOut.setStyle({{}}, {{stick: {{radius: 0.10, opacity: 0.35}}}});
	}}
	for (var i = 0; i < ARROWS.length; i++) {{
	var a = ARROWS[i];
	if (a.s <= 0.01) continue;
	viewer.addArrow({{start: {{x: a.p0[0], y: a.p0[1], z: a.p0[2]}},
	end: {{x: a.p1[0], y: a.p1[1], z: a.p1[2]}},
	radius: 0.05, color: 'red'}});
	}}
	viewer.zoomTo();
	viewer.render();
	}}
	boot(25);
	</script>
	</body>
	</html>"""
	escaped = html_mod.escape(doc, quote=True)
	return (
	f'<iframe srcdoc="{escaped}" title="3Dmol viewer" '
	f'style="width:100%;height:450px;border:0;border-radius:8px;display:block;" '
	f'allowfullscreen></iframe>'
	)


	def _fmt_metrics(e_ev, de_ev, fmax_ev, relaxed):
	import numpy as np

	fmax_kcal = fmax_ev * EV_TO_KCAL_MOL
	items = [
	("Energy", f"{e_ev:.4f} eV ({e_ev * EV_TO_KCAL_MOL:,.1f} kcal/mol)"),
	]
	if relaxed:
	items.append(("Δ after relaxation", f"{de_ev:+.4f} eV ({de_ev * EV_TO_KCAL_MOL:+,.1f} kcal/mol)"))
	items.append(("Max \|F\|", f"{fmax_ev:.3f} eV/Å ({fmax_kcal:.1f} kcal/mol/Å)"))
	return items


	@spaces.GPU(duration=30)
	def run(
	molecule: str,
	net_charge: float = 0,
	do_relax: bool = False,
	relax_fmax: float = 0.05,
	relax_steps: int = 200,
	model_choice: str = DEFAULT_MODEL,
	show_forces: bool = True,
	):
	"""Run an AceFF-2 single-point calculation (energy + forces) and optional
	geometry relaxation for one small molecule.

	Args:
	molecule: XYZ block, PDB text, or a SMILES string.
	net_charge: total molecular charge (-2..+2 supported by the model).
	do_relax: also relax the geometry with ASE LBFGS at this potential.
	relax_fmax: convergence threshold on max atomic force (eV/Å).
	relax_steps: max optimizer steps.
	model_choice: which AceFF-2 checkpoint to use.
	show_forces: draw per-atom force arrows on the 3D viewer.

	Returns:
	(status markdown, metric items, viewer HTML, relaxed XYZ text)
	"""
	t_start = time.perf_counter()
	try:
	import numpy as np
	import torch
	from ase import Atoms
	from ase.optimize import LBFGS

	model_key = model_choice if model_choice in MODELS else DEFAULT_MODEL
	supported = SUPPORTED_Z_OMOL if "OMOL" in model_key else SUPPORTED_Z

	# ---- 1. parse + validate input -------------------------------------
	symbols, pos = _parse_molecule(molecule)
	if len(symbols) > MAX_ATOMS:
	raise ValueError(
	f"Molecule has {len(symbols)} atoms; this demo caps at {MAX_ATOMS}."
	)
	znums = [_sym_z(s) for s in symbols]
	bad = sorted({s for s, zn in zip(symbols, znums) if zn not in supported})
	if bad:
	elems = (
	"H, B, C, N, O, F, Na, Si, P, S, Cl, K, Br, I"
	if "OMOL" in model_key
	else "H, B, C, N, O, F, Si, P, S, Cl, Br, I"
	)
	raise ValueError(
	f"Unsupported element(s) {', '.join(bad)}. {model_key} was trained on: {elems}."
	)
	if not -5 <= float(net_charge) <= 5:
	raise ValueError("Net charge must be an integer between -2 and +2.")
	if abs(float(net_charge) - round(float(net_charge))) > 1e-9:
	raise ValueError("Net charge must be an integer (e.g. 0, -1, +2).")
	q_float = float(net_charge)
	if not -2 <= q_float <= 2:
	raise ValueError(
	"AceFF-2 was trained on total molecular charges from -2 to +2 only."
	)

	# ---- 2. single point E + F (reference path: forward + backward) ----
	model = _get_model(model_key)
	z = torch.tensor(znums, dtype=torch.long, device="cuda")
	p = torch.tensor(np.asarray(pos), dtype=torch.float32, device="cuda")
	q = torch.tensor([q_float], dtype=torch.float32, device="cuda")
	e_init_ev, f_init = _energy_forces(model, z, p, q)

	xyz_in = _xyz_block(symbols, pos, "input geometry")
	xyz_out = None
	steps_used = 0
	e_final_ev = e_init_ev
	fmax_final = float(np.abs(f_init).max()) if len(f_init) else 0.0

	# ---- 3. optional LBFGS relaxation via the ASE calculator -----------
	if do_relax:
	atoms = Atoms(symbols=symbols, positions=np.asarray(pos))
	atoms.info["charge"] = q_float
	atoms.calc = _get_calculator(model_key)
	dyn = LBFGS(atoms, logfile=None)
	dyn.run(fmax=float(relax_fmax), steps=int(relax_steps))
	steps_used = dyn.get_number_of_steps()
	e_final_ev = float(atoms.get_potential_energy())
	fmax_final = float(np.abs(atoms.get_forces()).max())
	xyz_out = _xyz_block(
	symbols, atoms.get_positions(), "relaxed geometry (AceFF-2, LBFGS)"
	)

	de_ev = e_final_ev - e_init_ev
	metrics = _fmt_metrics(e_final_ev, de_ev, fmax_final, do_relax)

	# ---- 4. viewer: arrows scale longest force to ~0.55 Å -------------
	fmax_init = float(np.abs(f_init).max()) if len(f_init) else 0.0
	scale = 0.55 / fmax_init if fmax_init > 1e-9 else 0.0
	disp_pos = np.asarray(pos)
	arrows = []
	if show_forces and scale > 0:
	for (x, y, zw), (fx, fy, fz) in zip(disp_pos, f_init):
	mag = float(np.linalg.norm([fx, fy, fz]))
	arrows.append(
	{
	"p0": [round(float(x), 3), round(float(y), 3), round(float(zw), 3)],
	"p1": [
	round(float(x + fx * scale), 3),
	round(float(y + fy * scale), 3),
	round(float(zw + fz * scale), 3),
	],
	"s": mag,
	}
	)
	title = (
	"Input geometry (solid) → relaxed geometry (faded)"
	if do_relax
	else "Input geometry"
	)
	viewer = _viewer_html(title, xyz_in, xyz_out, arrows, f"{1.0/scale:.2f} Å/(eV/Å)" if scale else "—")

	mode = "single point + relaxation" if do_relax else "single point"
	q_str = f"{int(q_float):+d}" if q_float == int(q_float) else f"{q_float:+g}"
	status = (
	f"Done — {mode} · {len(symbols)} atoms · charge {q_str}"
	+ (f" · {steps_used} LBFGS steps" if do_relax else "")
	+ f" · {time.perf_counter() - t_start:.1f}s"
	)
	return status, metrics, viewer, xyz_out or ""
	except Exception as e:
	print("run() error:\n", traceback.format_exc(limit=6))
	raise gr.Error(f"{type(e).__name__}: {e}")


	with gr.Blocks(title="AceFF-2 · ML potential for small molecules") as demo:
	gr.Markdown(
	"""
	# AceFF-2 — Machine Learning Potential for Small Molecules

	AceFF ([paper](https://huggingface.co/papers/2601.00581)) is a state-of-the-art neural
	interatomic potential for drug-like molecules, built on the TensorNet v2 architecture in
	[TorchMD-Net](https://github.com/torchmd/torchmd-net) (Acellera, Apache-2.0).
	Give it a molecule — XYZ, PDB, or SMILES — and it predicts the potential energy
	and per-atom forces, and can relax the geometry with ASE.

	Supported elements: H, B, C, N, O, F, Si, P, S, Cl, Br, I (Na, K with the OMOL checkpoint) ·
	net charge −2…+2 · small molecules only (≤ 60 atoms in this demo).
	"""
	)
	with gr.Row(elem_id="col-container"):
	with gr.Column(variant="panel", scale=5):
	molecule = gr.Textbox(
	label="Molecule (XYZ block, PDB, or SMILES)",
	placeholder="Paste an XYZ block or PDB, or type a SMILES like CCO…",
	lines=10,
	max_lines=18,
	)
	with gr.Row():
	net_charge = gr.Number(label="Net charge (e)", value=0, precision=0, minimum=-2, maximum=2)
	do_relax = gr.Checkbox(label="Relax geometry (LBFGS)", value=False)
	show_forces = gr.Checkbox(label="Show force arrows", value=True)
	model_choice = gr.Radio(
	choices=list(MODELS),
	value=DEFAULT_MODEL,
	label="Checkpoint",
	info="AceFF-2.0: Acellera training set · AceFF-2-OMOL: public OMol25 data, adds Na/K.",
	)
	with gr.Accordion("Relaxation settings", open=False):
	relax_fmax = gr.Number(label="Convergence: max force (eV/Å)", value=0.05)
	relax_steps = gr.Number(label="Max optimizer steps", value=200, precision=0, minimum=1, maximum=1000)
	run_btn = gr.Button("Compute energy & forces", variant="primary")

	with gr.Column(scale=7):
	status = gr.Markdown("Ready. Pick an example below or paste your own molecule.")
	metrics = gr.HighlightedText(
	label="Results",
	combine_adjacent=False,
	)
	viewer = gr.HTML(
	value="<div style='height:430px;background:#17191c;border-radius:8px;display:flex;"
	"align-items:center;justify-content:center;color:#9aa0a6;font-family:var(--font);'>"
	"3D viewer appears here after a calculation (force arrows in red).</div>"
	)
	with gr.Accordion("Relaxed geometry (XYZ) — copy or download", open=False):
	xyz_out = gr.Textbox(label="XYZ", lines=8, interactive=True)

	gr.Examples(
	examples=[
	["examples/caffeine.xyz", 0, False, 0.05, 200, DEFAULT_MODEL, True],
	["examples/aspirin.xyz", 0, True, 0.05, 200, DEFAULT_MODEL, True],
	["examples/acetate.xyz", -1, True, 0.05, 200, DEFAULT_MODEL, True],
	["examples/ethanol.xyz", 0, False, 0.05, 200, DEFAULT_MODEL, True],
	["examples/water.xyz", 0, False, 0.05, 200, DEFAULT_MODEL, True],
	],
	inputs=[molecule, net_charge, do_relax, relax_fmax, relax_steps, model_choice, show_forces],
	outputs=[status, metrics, viewer, xyz_out],
	fn=run,
	cache_examples=True,
	cache_mode="lazy",
	label="Examples (molecules from the AceFF usage docs; geometries embedded with RDKit)",
	)

	run_btn.click(
	run,
	inputs=[molecule, net_charge, do_relax, relax_fmax, relax_steps, model_choice, show_forces],
	outputs=[status, metrics, viewer, xyz_out],
	api_name="compute",
	)

	demo.launch(mcp_server=True, theme=gr.themes.Citrus(),
	css="""
	#col-container { max-width: 1280px; margin: 0 auto; }
	.dark .gradio-container { color: var(--body-text-color); }
	""")