"""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("") + 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)