app.py

"""Gradio demo for PaddlePaddle/HPD-Parsing – Hierarchical Parallel Document Parsing."""

import os

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

try:
    import spaces  # MUST be before any torch/CUDA import; only present on HF ZeroGPU
    GPU_DECORATOR = spaces.GPU(duration=60)
except ImportError:
    def GPU_DECORATOR(fn):
        return fn

import torch
import gradio as gr
from transformers import AutoModel, AutoTokenizer

from hpd_postprocess import parse_blocks, blocks_to_markdown, draw_boxes_on_image

MODEL_ID = os.environ.get(
    "HPD_MODEL_PATH",
    "PaddlePaddle/HPD-Parsing",
)

# --- Image preprocessing (mirrors the repo's image_preprocess.py) ----------
import torchvision.transforms as T
from torchvision.transforms.functional import InterpolationMode
from PIL import Image

IMAGENET_MEAN, IMAGENET_STD = (0.485, 0.456, 0.406), (0.229, 0.224, 0.225)
IMAGE_SIZE = 448
MIN_DYNAMIC_PATCH = 1
MAX_DYNAMIC_PATCH = 24
USE_THUMBNAIL = True


def build_transform(input_size=IMAGE_SIZE):
    return T.Compose([
        T.Lambda(lambda img: img.convert("RGB")),
        T.Resize((input_size, input_size), interpolation=InterpolationMode.BICUBIC),
        T.ToTensor(),
        T.Normalize(IMAGENET_MEAN, IMAGENET_STD),
    ])


def get_target_ratios(min_num, max_num):
    ratios = {(i, j)
              for n in range(min_num, max_num + 1)
              for i in range(1, n + 1) for j in range(1, n + 1)
              if min_num <= i * j <= max_num}
    return sorted(ratios, key=lambda x: x[0] * x[1])


def find_closest_aspect_ratio_optim(aspect_ratio, target_ratios, width, height,
                                    image_size, top_k=3, ar_threshold=0.2):
    area = width * height
    candidates = []
    for ratio in target_ratios:
        ar_diff = abs(aspect_ratio - ratio[0] / ratio[1])
        if ar_threshold is not None and ar_diff > ar_threshold:
            continue
        area_diff = abs(area - image_size * image_size * ratio[0] * ratio[1])
        candidates.append((ratio, area_diff, ar_diff))
    if not candidates:
        for ratio in target_ratios:
            ar_diff = abs(aspect_ratio - ratio[0] / ratio[1])
            area_diff = abs(area - image_size * image_size * ratio[0] * ratio[1])
            candidates.append((ratio, area_diff, ar_diff))
    candidates.sort(key=lambda x: x[1])
    top = candidates[:top_k]
    top.sort(key=lambda x: x[2])
    return top[0][0]


def dynamic_preprocess(image, target_ratios, image_size=IMAGE_SIZE, use_thumbnail=USE_THUMBNAIL):
    w, h = image.size
    ratio = find_closest_aspect_ratio_optim(w / h, target_ratios, w, h, image_size)
    tw, th = image_size * ratio[0], image_size * ratio[1]
    blocks = ratio[0] * ratio[1]
    resized = image.resize((tw, th))
    cols = tw // image_size
    tiles = []
    for i in range(blocks):
        box = ((i % cols) * image_size, (i // cols) * image_size,
               ((i % cols) + 1) * image_size, ((i // cols) + 1) * image_size)
        tiles.append(resized.crop(box))
    if use_thumbnail and blocks != 1:
        tiles.append(image.resize((image_size, image_size)))
    return tiles


def load_image_from_pil(pil_image):
    """Preprocess a PIL image into the dynamic-tiling tensor the model expects."""
    image = pil_image.convert("RGB")
    min_num, max_num = MIN_DYNAMIC_PATCH, MAX_DYNAMIC_PATCH
    if USE_THUMBNAIL and max_num != 1:
        max_num += 1
    target_ratios = get_target_ratios(min_num, max_num)
    transform = build_transform(IMAGE_SIZE)
    tiles = dynamic_preprocess(image, target_ratios, IMAGE_SIZE, USE_THUMBNAIL)
    return torch.stack([transform(t) for t in tiles])


# --- Example gallery (fixed sample copied into examples/) ------------------

EXAMPLE_IMAGES_DIR = os.environ.get(
    "HPD_EXAMPLE_IMAGES_DIR",
    os.path.join(os.path.dirname(os.path.abspath(__file__)), "examples"),
)
EXAMPLE_SAMPLE_SIZE = 6


def _list_example_images(dir_path, sample_size):
    """List example image paths from the local examples directory.

    Returns an empty list if the directory doesn't exist (no example gallery
    is rendered in that case).
    """
    if not os.path.isdir(dir_path):
        return []
    supported_exts = {".png", ".jpg", ".jpeg", ".bmp", ".webp"}
    candidates = sorted(
        f for f in os.listdir(dir_path)
        if os.path.splitext(f)[1].lower() in supported_exts
    )
    return [os.path.join(dir_path, name) for name in candidates[:sample_size]]


EXAMPLE_IMAGE_PATHS = _list_example_images(EXAMPLE_IMAGES_DIR, EXAMPLE_SAMPLE_SIZE)


# --- Load model at module scope (ZeroGPU pattern) ---------------------------
print("Loading model...")

# transformers 5.x expects `all_tied_weights_keys` on PreTrainedModel subclasses.
# The custom InternVLChatModel (written for transformers 4.x) doesn't call post_init()
# which is where transformers 5.x sets this attribute. Add a default so loading works.
import transformers.modeling_utils as _mu
if 'all_tied_weights_keys' not in _mu.PreTrainedModel.__dict__:
    _mu.PreTrainedModel.all_tied_weights_keys = {}

model = AutoModel.from_pretrained(
    MODEL_ID,
    torch_dtype=torch.bfloat16,
    trust_remote_code=True,
).eval().to("cuda")

# Ensure the attribute exists on the loaded model instance (custom code may skip post_init)
if not hasattr(model, 'all_tied_weights_keys') or not isinstance(getattr(model, 'all_tied_weights_keys', None), dict):
    model.all_tied_weights_keys = {}

tokenizer = AutoTokenizer.from_pretrained(
    MODEL_ID,
    trust_remote_code=True,
    use_fast=False,
)

# Load P-MTP weights (they ship inside the checkpoint; just mark them ready)
model.load_mtp_weights()
print("Model loaded.")


# --- Inference ---------------------------------------------------------------

DEFAULT_USE_FORK = True
DEFAULT_USE_MTP = True
DEFAULT_MAX_NEW_TOKENS = 8000


@GPU_DECORATOR
def parse_document(image):
    """Parse a document image into structured text using HPD-Parsing.

    Args:
        image: The document image to parse.

    Returns:
        A tuple ``(boxed_image, markdown_text, raw_response)`` where
        ``boxed_image`` is the input image annotated with typed bounding
        boxes, ``markdown_text`` is the cleaned, tag-free markdown
        reconstruction of the parse, and ``raw_response`` is the model's
        unprocessed output (including ``<BLOCK>/<FORK>/<CHILD>`` tags).
    """
    if image is None:
        return None, "Please upload a document image first.", ""

    pixel_values = load_image_from_pil(image).to(torch.bfloat16).to("cuda")

    prompt = "document parsing with fork." if DEFAULT_USE_FORK else "document parsing."

    response = model.generate_hpd(
        tokenizer,
        pixel_values,
        prompt,
        dict(max_new_tokens=DEFAULT_MAX_NEW_TOKENS),
        use_mtp=DEFAULT_USE_MTP,
        num_speculative_tokens=6,
        batch_children=False,
    )

    blocks = parse_blocks(response)
    markdown_text = blocks_to_markdown(blocks) or response
    boxed_image = draw_boxes_on_image(image, blocks)
    return boxed_image, markdown_text, response


# --- Gradio UI ---------------------------------------------------------------

LATEX_DELIMS = [
    {"left": "$$", "right": "$$", "display": True},
    {"left": "$", "right": "$", "display": False},
    {"left": "\\(", "right": "\\)", "display": False},
    {"left": "\\[", "right": "\\]", "display": True},
]

CUSTOM_CSS = """
body, .gradio-container { font-family: "Noto Sans SC", "Microsoft YaHei", "PingFang SC", sans-serif; }
.app-header { text-align: center; max-width: 1100px; margin: 0 auto 8px !important; }
#result-tabs .tabitem { padding-top: 8px !important; }
#example-gallery img { object-fit: cover !important; }
"""

with gr.Blocks(css=CUSTOM_CSS) as demo:
    gr.Markdown(
        "# HPD-Parsing: Hierarchical Parallel Document Parsing\n"
        "Upload a document image and get structured text output. "
        "Powered by [PaddlePaddle/HPD-Parsing](https://huggingface.co/PaddlePaddle/HPD-Parsing) – "
        "a 1B-parameter VLM that achieves SOTA on OmniDocBench via hierarchical parallel decoding.",
        elem_classes=["app-header"],
    )

    with gr.Row():
        with gr.Column(scale=5):
            input_image = gr.Image(label="Document Image", type="pil")
            run_btn = gr.Button("Parse Document", variant="primary")

            if EXAMPLE_IMAGE_PATHS:
                gr.Markdown("_Click an example below to load it._")
                example_gallery = gr.Gallery(
                    value=EXAMPLE_IMAGE_PATHS,
                    columns=3,
                    height=360,
                    preview=False,
                    allow_preview=False,
                    label=None,
                    elem_id="example-gallery",
                )

                def _on_example_select(evt: gr.SelectData):
                    return EXAMPLE_IMAGE_PATHS[evt.index]

                example_gallery.select(_on_example_select, inputs=None, outputs=input_image)

        with gr.Column(scale=7):
            with gr.Tabs(elem_id="result-tabs"):
                with gr.Tab("Visualization"):
                    output_image = gr.Image(label="Detected Layout (bounding boxes)")
                with gr.Tab("Markdown Preview"):
                    output_markdown = gr.Markdown(
                        label="Parsed Output (rendered Markdown)",
                        latex_delimiters=LATEX_DELIMS,
                    )
                with gr.Tab("Raw Output"):
                    output_raw = gr.Code(label="Raw model output", language="markdown")

    run_btn.click(
        fn=parse_document,
        inputs=[input_image],
        outputs=[output_image, output_markdown, output_raw],
        api_name="parse_document",
    )

demo.launch(mcp_server=True)