app.py

"""Gradio demo for UniParser/MolParser-Mobile-V2 — molecule image → E-SMILES 2.0.

Recognizes a molecular structure image with the 10M-param MolParser-Mobile-V2
vision-encoder-decoder, then runs the authors' E-SMILES post-processing
(normalized E-SMILES, substituted SMILES, CXSMILES) and renders the parsed
structure back to a 2D drawing using the MolParser utils (Apache-2.0).
"""

import os

os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True")

import spaces  # MUST come before any CUDA-touching import
import time
import traceback

import gradio as gr
import torch
from PIL import Image
from transformers import AutoModelForImageTextToText, AutoProcessor

from molparser import utils as mutils

MODEL_ID = "UniParser/MolParser-Mobile-V2"
DTYPE = torch.float16
MAX_LENGTH = 384  # V2 max output length (was 256 in V1)

# Load model and processor at module scope; ZeroGPU intercepts .to("cuda").
processor = AutoProcessor.from_pretrained(MODEL_ID, trust_remote_code=True)
model = AutoModelForImageTextToText.from_pretrained(
    MODEL_ID,
    dtype=DTYPE,
    trust_remote_code=True,
).to("cuda").eval()
print("MolParser-Mobile-V2 loaded")


@spaces.GPU(duration=15)
def recognize_molecule(image: str):
    """Convert a molecular structure image into an E-SMILES 2.0 string.

    Runs MolParser-Mobile-V2 (greedy decoding, fp16) on the image, then applies
    the MolParser E-SMILES post-processing to produce a substituted SMILES,
    normalized E-SMILES and CXSMILES, plus an SVG rendering of the parsed
    structure.

    Args:
        image: Path to a molecular structure image (PNG/JPG).

    Returns:
        Tuple of (raw E-SMILES, substituted SMILES, normalized E-SMILES,
        CXSMILES, status info, SVG rendering of the parsed structure).
    """
    started = time.perf_counter()
    if not image:
        return ("", "", "", "", "Please upload a molecular structure image.", "")

    try:
        pil_image = Image.open(image).convert("RGB")
    except Exception as exc:
        return ("", "", "", "", f"Could not open image: {exc}", "")

    inputs = processor(images=pil_image, return_tensors="pt")
    pixel_values = inputs["pixel_values"].to("cuda", dtype=DTYPE)

    with torch.inference_mode(), torch.autocast(device_type="cuda", dtype=DTYPE):
        output_ids = model.generate(
            pixel_values=pixel_values,
            max_length=MAX_LENGTH,
            num_beams=1,
            do_sample=False,
        )

    caption = processor.batch_decode(output_ids.detach().cpu().tolist(), skip_special_tokens=True)[0]
    infer_secs = time.perf_counter() - started

    # --- Authors' post-processing (molparser.utils, Apache-2.0) ---
    smi = esmi = cxsmiles = ""
    info = ""
    try:
        post = mutils.postprocess_caption(caption)
        smi = str(post.get("smi", ""))
        esmi = str(post.get("esmi", ""))
        cxsmiles = str(post.get("cxsmiles", ""))
        flags = []
        if post.get("markush"):
            flags.append("Markush")
        if post.get("sru"):
            flags.append("SRU")
        info = f"Recognized in {infer_secs:.2f}s" + (f" · {', '.join(flags)} structure" if flags else "")
    except Exception:
        traceback.print_exc()
        info = f"Recognized in {infer_secs:.2f}s · post-processing failed (raw caption shown)"

    # --- Render the parsed structure as SVG (best effort) ---
    svg = ""
    for candidate in (esmi, caption):
        if not candidate:
            continue
        try:
            rendered = mutils.draw(candidate, output_format="svg")
            # Strip the XML prolog so the fragment embeds cleanly in HTML.
            if rendered.lstrip().startswith("<?"):
                rendered = rendered[rendered.find("?>") + 2 :]
            svg = rendered.strip()
            break
        except Exception:
            continue

    return (caption, smi, esmi, cxsmiles, info, svg)


CSS = """
#col-container { max-width: 1100px; margin: 0 auto; }
.dark .gradio-container { color: var(--body-text-color); }
.mol-svg svg { max-width: 100%; height: auto; }
"""

with gr.Blocks() as demo:
    with gr.Column(elem_id="col-container"):
        gr.Markdown(
            """
            # MolParser-Mobile-V2: Molecule Image → E-SMILES 2.0
            Upload a molecular structure image and get its **E-SMILES 2.0** representation,
            plus the substituted **SMILES**, **CXSMILES**, and a rendered 2D structure.
            [MolParser-Mobile-V2](https://huggingface.co/UniParser/MolParser-Mobile-V2) is a
            10M-parameter optical chemical structure recognition (OCSR) model — it reads
            molecule drawings (including complex Markush structures) straight from the literature.
            """
        )

        with gr.Row():
            image_input = gr.Image(
                label="Molecular structure image",
                type="filepath",
                sources=["upload", "clipboard"],
                height=320,
            )
            with gr.Column():
                info_text = gr.Markdown("", visible=False)
                raw_esmiles = gr.Textbox(
                    label="E-SMILES 2.0 (raw model output)",
                    lines=3,
                    interactive=False,
                )
                smi_out = gr.Textbox(
                    label="SMILES (abbreviations substituted)",
                    lines=2,
                    interactive=False,
                )
                with gr.Accordion("Normalized E-SMILES / CXSMILES", open=False):
                    esmi_out = gr.Textbox(
                        label="Normalized E-SMILES",
                        lines=3,
                        interactive=False,
                    )
                    cxsmiles_out = gr.Textbox(
                        label="CXSMILES",
                        lines=3,
                        interactive=False,
                    )
                run_btn = gr.Button("Recognize", variant="primary")

        with gr.Accordion("Rendered structure (parsed from the recognized E-SMILES)", open=False):
            svg_html = gr.HTML(elem_classes=["mol-svg"])

        gr.Examples(
            examples=[
                ["examples/01-ordinary-molecule.png"],
                ["examples/02-markush-substituent.png"],
                ["examples/03-connection-point.png"],
                ["examples/04-special-markush-dna.png"],
                ["examples/05-ring-attachment-one-repeat-range.png"],
                ["examples/06-ring-attachment-two-substituents.png"],
                ["examples/07-ring-attachment-complex-groups.png"],
                ["examples/08-repeat-n.png"],
                ["examples/09-polymer-original.png"],
                ["examples/10-abstract-ring-basic.png"],
                ["examples/biaryl-axial-ra-source.png"],
            ],
            inputs=[image_input],
            outputs=[raw_esmiles, smi_out, esmi_out, cxsmiles_out, info_text, svg_html],
            fn=recognize_molecule,
            cache_examples=True,
            cache_mode="lazy",
        )

    def _show_info(info: str):
        """Reveal the status line only when it has content."""
        return gr.update(value=info, visible=bool(info))

    run_btn.click(
        fn=recognize_molecule,
        inputs=image_input,
        outputs=[raw_esmiles, smi_out, esmi_out, cxsmiles_out, info_text, svg_html],
        api_name="recognize",
    ).then(
        fn=_show_info,
        inputs=info_text,
        outputs=info_text,
    )

demo.launch(theme=gr.themes.Citrus(), css=CSS, mcp_server=True)