"""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 ``//`` 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)