first commit

LICENSE

@@ -0,0 +1,9 @@
+MIT License
+
+Copyright 2025 ByteDance Ltd. and/or its affiliates
+
+Permission is hereby granted, free of charge, to any person obtaining a copy of this software and associated documentation files (the “Software”), to deal in the Software without restriction, including without limitation the rights to use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of the Software, and to permit persons to whom the Software is furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED “AS IS”, WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.

README_CN.md

@@ -0,0 +1,198 @@
+<div align="center">
+  <img src="./assets/dolphin.png" width="300">
+</div>
+
+<div align="center">
+  <a href="https://arxiv.org/abs/2505.14059">
+    <img src="https://img.shields.io/badge/论文-arXiv-red">
+  </a>
+  <a href="https://huggingface.co/ByteDance/Dolphin">
+    <img src="https://img.shields.io/badge/HuggingFace-Dolphin-yellow">
+  </a>
+  <a href="https://modelscope.cn/models/ByteDance/Dolphin">
+    <img src="https://img.shields.io/badge/ModelScope-Dolphin-purple">
+  </a>
+  <a href="https://huggingface.co/spaces/ByteDance/Dolphin">
+    <img src="https://img.shields.io/badge/演示-Dolphin-blue">
+  </a>
+  <a href="https://github.com/bytedance/Dolphin">
+    <img src="https://img.shields.io/badge/代码-Github-green">
+  </a>
+  <a href="https://opensource.org/licenses/MIT">
+    <img src="https://img.shields.io/badge/许可证-MIT-lightgray">
+  </a>
+  <br>
+</div>
+
+<br>
+
+<div align="center">
+  <img src="./assets/demo.gif" width="800">
+</div>
+
+# Dolphin: 基于异构锚点提示的文档图像解析
+
+Dolphin（**Do**cument Image **P**arsing via **H**eterogeneous Anchor Prompt**in**g）是一个创新的多模态文档图像解析模型，采用"分析-解析"的两阶段范式。本仓库包含Dolphin的演示代码和预训练模型。
+
+## 📑 概述
+
+由于文档图像中文本段落、图表、公式和表格等元素的复杂交织，文档图像解析具有挑战性。Dolphin通过两阶段方法解决这些挑战：
+
+1. **🔍 第一阶段**：通过按自然阅读顺序生成元素序列进行全面的页面级布局分析
+2. **🧩 第二阶段**：使用异构锚点和任务特定提示高效并行解析文档元素
+
+<div align="center">
+  <img src="./assets/framework.png" width="680">
+</div>
+
+Dolphin在多样化的页面级和元素级解析任务中取得了优异的性能，同时通过其轻量级架构和并行解析机制确保了卓越的效率。
+
+## 🚀 演示
+在 [Demo-Dolphin](http://115.190.42.15:8888/dolphin/) 上试用我们的演示。
+
+## 📅 更新日志
+- 🔥 **2025.06.30** 新增[TensorRT-LLM](https://github.com/bytedance/Dolphin/blob/master/deployment/tensorrt_llm/ReadMe.md)支持，提升推理速度！
+- 🔥 **2025.06.27** 新增[vLLM](https://github.com/bytedance/Dolphin/blob/master/deployment/vllm/ReadMe.md)支持，提升推理速度！
+- 🔥 **2025.06.13** 新增多页PDF文档解析功能。
+- 🔥 **2025.05.21** 我们的演示已在 [链接](http://115.190.42.15:8888/dolphin/) 发布。快来体验吧！
+- 🔥 **2025.05.20** Dolphin的预训练模型和推理代码已发布。
+- 🔥 **2025.05.16** 我们的论文已被ACL 2025接收。论文链接：[arXiv](https://arxiv.org/abs/2505.14059)。
+
+## 🛠️ 安装
+
+1. 克隆仓库：
+   ```bash
+   git clone https://github.com/ByteDance/Dolphin.git
+   cd Dolphin
+   ```
+
+2. 安装依赖：
+   ```bash
+   pip install -r requirements.txt
+   ```
+
+3. 使用以下选项之一下载预训练模型：
+
+   **选项A：原始模型格式（基于配置文件）**
+   
+   从 [百度网盘](https://pan.baidu.com/s/15zcARoX0CTOHKbW8bFZovQ?pwd=9rpx) 或 [Google Drive](https://drive.google.com/drive/folders/1PQJ3UutepXvunizZEw-uGaQ0BCzf-mie?usp=sharing) 下载，并将其放在 `./checkpoints` 文件夹中。
+
+   **选项B：Hugging Face模型格式**
+   
+   访问我们的Huggingface [模型卡片](https://huggingface.co/ByteDance/Dolphin)，或通过以下方式下载模型：
+   
+   ```bash
+   # 从Hugging Face Hub下载模型
+   git lfs install
+   git clone https://huggingface.co/ByteDance/Dolphin ./hf_model
+   # 或使用Hugging Face CLI
+   pip install huggingface_hub
+   huggingface-cli download ByteDance/Dolphin --local-dir ./hf_model
+   ```
+
+## ⚡ 推理
+
+Dolphin提供两个推理框架，支持两种解析粒度：
+- **页面级解析**：将整个文档页面解析为结构化的JSON和Markdown格式
+- **元素级解析**：解析单个文档元素（文本、表格、公式）
+
+### 📄 页面级解析
+
+#### 使用原始框架（基于配置文件）
+
+```bash
+# 处理单个文档图像
+python demo_page.py --config ./config/Dolphin.yaml --input_path ./demo/page_imgs/page_1.jpeg --save_dir ./results
+
+# 处理单个文档PDF
+python demo_page.py --config ./config/Dolphin.yaml --input_path ./demo/page_imgs/page_6.pdf --save_dir ./results
+
+# 处理目录中的所有文档
+python demo_page.py --config ./config/Dolphin.yaml --input_path ./demo/page_imgs --save_dir ./results
+
+# 使用自定义批次大小进行并行元素解码
+python demo_page.py --config ./config/Dolphin.yaml --input_path ./demo/page_imgs --save_dir ./results --max_batch_size 8
+```
+
+#### 使用Hugging Face框架
+
+```bash
+# 处理单个文档图像
+python demo_page_hf.py --model_path ./hf_model --input_path ./demo/page_imgs/page_1.jpeg --save_dir ./results
+
+# 处理单个文档PDF
+python demo_page_hf.py --model_path ./hf_model --input_path ./demo/page_imgs/page_6.pdf --save_dir ./results
+
+# 处理目录中的所有文档
+python demo_page_hf.py --model_path ./hf_model --input_path ./demo/page_imgs --save_dir ./results
+
+# 使用自定义批次大小进行并行元素解码
+python demo_page_hf.py --model_path ./hf_model --input_path ./demo/page_imgs --save_dir ./results --max_batch_size 16
+```
+
+### 🧩 元素级解析
+
+#### 使用原始框架（基于配置文件）
+
+```bash
+# 处理单个表格图像
+python demo_element.py --config ./config/Dolphin.yaml --input_path ./demo/element_imgs/table_1.jpeg --element_type table
+
+# 处理单个公式图像
+python demo_element.py --config ./config/Dolphin.yaml --input_path ./demo/element_imgs/line_formula.jpeg --element_type formula
+
+# 处理单个文本段落图像
+python demo_element.py --config ./config/Dolphin.yaml --input_path ./demo/element_imgs/para_1.jpg --element_type text
+```
+
+#### 使用Hugging Face框架
+
+```bash
+# 处理单个表格图像
+python demo_element_hf.py --model_path ./hf_model --input_path ./demo/element_imgs/table_1.jpeg --element_type table
+
+# 处理单个公式图像
+python demo_element_hf.py --model_path ./hf_model --input_path ./demo/element_imgs/line_formula.jpeg --element_type formula
+
+# 处理单个文本段落图像
+python demo_element_hf.py --model_path ./hf_model --input_path ./demo/element_imgs/para_1.jpg --element_type text
+```
+
+## 🌟 主要特性
+
+- 🔄 基于单一VLM的两阶段分析-解析方法
+- 📊 在文档解析任务上的优异性能
+- 🔍 自然阅读顺序元素序列生成
+- 🧩 针对不同文档元素的异构锚点提示
+- ⏱️ 高效的并行解析机制
+- 🤗 支持Hugging Face Transformers，便于集成
+
+## 📮 通知
+**征集不良案例：** 如果您遇到模型表现不佳的案例，我们非常欢迎您在issue中分享。我们正在持续优化和改进模型。
+
+## 💖 致谢
+
+我们要感谢以下开源项目为本工作提供的灵感和参考：
+- [Donut](https://github.com/clovaai/donut/)
+- [Nougat](https://github.com/facebookresearch/nougat)
+- [GOT](https://github.com/Ucas-HaoranWei/GOT-OCR2.0)
+- [MinerU](https://github.com/opendatalab/MinerU/tree/master)
+- [Swin](https://github.com/microsoft/Swin-Transformer)
+- [Hugging Face Transformers](https://github.com/huggingface/transformers)
+
+## 📝 引用
+
+如果您在研究中发现此代码有用，请使用以下BibTeX条目。
+
+```bibtex
+@article{feng2025dolphin,
+  title={Dolphin: Document Image Parsing via Heterogeneous Anchor Prompting},
+  author={Feng, Hao and Wei, Shu and Fei, Xiang and Shi, Wei and Han, Yingdong and Liao, Lei and Lu, Jinghui and Wu, Binghong and Liu, Qi and Lin, Chunhui and others},
+  journal={arXiv preprint arXiv:2505.14059},
+  year={2025}
+}
+```
+
+## 星标历史
+
+[![Star History Chart](https://api.star-history.com/svg?repos=bytedance/Dolphin&type=Date)](https://www.star-history.com/#bytedance/Dolphin&Date) 

README_PDF_APP.md

@@ -0,0 +1,117 @@
+# DOLPHIN PDF Document AI - HuggingFace Spaces App
+
+A Gradio-based web application for processing PDF documents using the DOLPHIN vision-language model. This app converts PDF files to images and processes them page by page to extract text, tables, and figures.
+
+## Features
+
+- **PDF Upload**: Upload PDF documents directly through the web interface
+- **Page-by-Page Processing**: Converts PDF pages to high-quality images and processes each individually
+- **Document Parsing**: Extracts text, tables, and figures using the DOLPHIN model
+- **Markdown Output**: Generates clean markdown with embedded images and tables
+- **Memory Optimized**: Designed for NVIDIA T4 GPU deployment on HuggingFace Spaces
+- **Progress Tracking**: Real-time progress updates during processing
+
+## Files
+
+- `gradio_pdf_app.py` - Main Gradio application with PDF processing functionality
+- `app.py` - HuggingFace Spaces entry point
+- `requirements_hf_spaces.txt` - Dependencies optimized for HF Spaces deployment
+
+## Usage
+
+### Local Development
+
+```bash
+# Install dependencies
+pip install -r requirements_hf_spaces.txt
+
+# Run the app
+python gradio_pdf_app.py
+```
+
+### HuggingFace Spaces Deployment
+
+1. Create a new HuggingFace Space with Gradio SDK
+2. Upload the following files:
+   - `app.py`
+   - `gradio_pdf_app.py`
+   - `utils/` (directory with utility functions)
+   - `requirements_hf_spaces.txt` (rename to `requirements.txt`)
+
+3. Configure the Space:
+   - **SDK**: Gradio
+   - **Hardware**: NVIDIA T4 Small (recommended)
+   - **Python Version**: 3.9+
+
+## Technical Details
+
+### Memory Optimizations
+
+- Uses `torch.float16` for GPU inference
+- Smaller batch sizes (4) for element processing
+- Memory cleanup with `torch.cuda.empty_cache()`
+- Reduced max sequence length (2048) for generation
+
+### PDF Processing Pipeline
+
+1. **PDF to Images**: Uses PyMuPDF with 2x zoom for quality
+2. **Layout Analysis**: DOLPHIN model parses document structure
+3. **Element Extraction**: Processes text, tables, and figures separately
+4. **Markdown Generation**: Converts results to formatted markdown
+5. **Gallery View**: Creates overview of all processed pages
+
+### Model Integration
+
+- Uses HuggingFace transformers implementation
+- Loads model with `device_map="auto"` for GPU optimization
+- Batch processing for improved efficiency
+- Graceful fallback to CPU if GPU unavailable
+
+## Configuration
+
+The app automatically detects and uses the DOLPHIN model:
+- Local path: `./hf_model`
+- HuggingFace Hub: `ByteDance/DOLPHIN`
+
+## Dependencies
+
+Core requirements:
+- `torch>=2.1.0` - PyTorch for model inference
+- `transformers>=4.47.0` - HuggingFace model loading
+- `gradio>=5.36.0` - Web interface
+- `pymupdf>=1.26.0` - PDF processing
+- `pillow>=9.3.0` - Image processing
+- `opencv-python-headless>=4.8.0` - Computer vision operations
+
+## Error Handling
+
+- Graceful handling of PDF conversion failures
+- Memory management for large documents
+- Progress reporting for long-running operations
+- Fallback markdown generation if converter fails
+
+## Performance Notes
+
+- Optimized for NVIDIA T4 with 16GB VRAM
+- Processing time: ~30-60 seconds per page (depends on complexity)
+- Memory usage: ~8-12GB VRAM for typical documents
+- CPU fallback available but significantly slower
+
+## Example Output
+
+The app generates:
+1. **Markdown Preview**: Rendered document with LaTeX support
+2. **Raw Markdown**: Source text for copying/editing
+3. **Page Gallery**: Visual overview of all processed pages
+4. **JSON Details**: Technical processing information
+
+## Troubleshooting
+
+- **Out of Memory**: Reduce batch size or use CPU
+- **PDF Conversion Failed**: Check PDF format compatibility
+- **Model Loading Error**: Verify model path and permissions
+- **Slow Processing**: Ensure GPU is available and configured
+
+## Credits
+
+Built on the DOLPHIN model by ByteDance. Optimized for HuggingFace Spaces deployment.

app.py

@@ -0,0 +1,27 @@
+"""
+HuggingFace Spaces entry point for DOLPHIN PDF Document AI
+"""
+
+import os
+import sys
+
+# Add the current directory to Python path for imports
+sys.path.append(os.path.dirname(os.path.abspath(__file__)))
+
+# Import and run the Gradio app
+from gradio_pdf_app import demo
+
+if __name__ == "__main__":
+    # Launch the app for HuggingFace Spaces
+    demo.launch(
+        server_name="0.0.0.0",
+        server_port=7860,
+        share=False,
+        show_error=True,
+        enable_queue=True,
+        max_threads=2,
+        # Additional HF Spaces specific settings
+        inbrowser=False,
+        show_tips=False,
+        quiet=True
+    )

chat.py

@@ -0,0 +1,198 @@
+""" 
+Copyright (c) 2025 Bytedance Ltd. and/or its affiliates
+SPDX-License-Identifier: MIT
+"""
+
+import os
+import warnings
+from collections import OrderedDict
+
+from omegaconf import ListConfig
+
+warnings.filterwarnings("ignore", category=UserWarning)
+warnings.filterwarnings("ignore", category=FutureWarning)
+os.environ.setdefault("PROTOCOL_BUFFERS_PYTHON_IMPLEMENTATION", "python")
+
+import torch
+from PIL import Image
+from transformers import PreTrainedTokenizerFast
+
+from utils.model import DonutConfig, DonutModel, SwinEncoder
+from utils.processor import DolphinProcessor
+
+
+def try_rename_lagacy_weights(ckpt, output_path=""):
+    if "state_dict" in ckpt.keys():
+        ckpt = ckpt["state_dict"]
+    if "module" in ckpt.keys():
+        ckpt = ckpt["module"]
+    new_ckpt = OrderedDict()
+    for k, v in ckpt.items():
+        if k.startswith("model."):
+            k = k[len("model.") :]
+        if k.startswith("encoder"):
+            new_ckpt["vpm" + k[len("encoder") :]] = v
+        elif k.startswith("decoder"):
+            new_ckpt["llm" + k[len("encoder") :]] = v
+        else:
+            new_ckpt[k] = v
+    if output_path:
+        torch.save(new_ckpt, output_path)
+    return new_ckpt
+
+
+def convert_listconfig_to_list(config):
+    new_config = {}
+    for k, v in config.items():
+        if isinstance(v, ListConfig):
+            new_config[k] = list(v)
+        else:
+            new_config[k] = v
+    return new_config
+
+
+class DOLPHIN:
+    def __init__(self, config, ckpt_path="") -> None:
+        self.model_args = config.model
+        self.swin_args = config.model.pop("swin_args")
+        self.swin_args = convert_listconfig_to_list(self.swin_args)
+
+        vision_tower = SwinEncoder(
+            input_size=self.swin_args["img_size"],
+            patch_size=self.swin_args["patch_size"],
+            embed_dim=self.swin_args["embed_dim"],
+            window_size=self.swin_args["window_size"],
+            encoder_layer=self.swin_args["encoder_layer"],
+            num_heads=self.swin_args["num_heads"],
+            align_long_axis=self.swin_args["align_long_axis"],
+        )
+
+        self.tokenizer = PreTrainedTokenizerFast(tokenizer_file=self.model_args.tokenizer_path)
+        self.tokenizer.pad_token = "<pad>"
+        self.tokenizer.bos_token = "<s>"
+        self.tokenizer.eos_token = "</s>"
+        self.tokenizer.unk_token = "<unk>"
+
+        if self.model_args.get("extra_answer_tokens", False):
+            # print("Allowing multitask training: adding <Answer/> to the tokenizer.")
+            prompt_end_token = " <Answer/>"
+            self.tokenizer.add_special_tokens({"additional_special_tokens": sorted(set([prompt_end_token]))})
+            self.tokenizer._prompt_end_token = prompt_end_token
+            self.tokenizer._prompt_end_token_id = self.tokenizer.convert_tokens_to_ids(prompt_end_token)
+
+        donut_config = DonutConfig(
+            decoder_layer=self.model_args.decoder_layer,
+            max_length=self.model_args.max_length,
+            max_position_embeddings=self.model_args.max_position_embeddings,
+            hidden_dimension=self.model_args.hidden_dimension,
+        )
+
+        self.model = DonutModel(config=donut_config, vision_tower=vision_tower, tokenizer=self.tokenizer)
+        if self.model_args.model_name_or_path:
+            ckpt = torch.load(self.model_args.model_name_or_path)
+            ckpt = try_rename_lagacy_weights(ckpt)
+            self.model.load_state_dict(ckpt, strict=True)
+
+        device = "cuda" if torch.cuda.is_available() else "cpu"
+        self.model.to(device)
+        self.model.eval()
+        transform_args = {
+            "input_size": self.swin_args["img_size"],
+            "max_length": self.model_args.max_length,
+        }
+        self.processor = DolphinProcessor({}, self.tokenizer, transform_args=transform_args)
+
+    def chat(
+        self,
+        question,
+        image,
+        return_raw=False,
+        return_score=False,
+        return_img_size=False,
+        only_return_img_size=False,
+        max_batch_size=16,
+    ):
+
+        def _preprocess_image(image):
+            if isinstance(image, str):
+                image = Image.open(image).convert("RGB")
+            if return_img_size or only_return_img_size:
+                image_tensor, ori_size = self.processor.process_image_for_inference(image, return_img_size=True)
+            else:
+                image_tensor = self.processor.process_image_for_inference(image, return_img_size=False)
+                ori_size = None
+            return image_tensor, ori_size
+
+        def _preprocess_prompt(question):
+            if self.model_args.get("extra_answer_tokens", False):
+                if self.tokenizer._prompt_end_token not in question:
+                    question = question + self.tokenizer._prompt_end_token
+            prompt_ids = self.processor.process_prompt_for_inference(question)
+            return prompt_ids
+
+        def _preprocess_prompt_batch(question):
+            if self.model_args.get("extra_answer_tokens", False):
+                for i in range(len(question)):
+                    if self.tokenizer._prompt_end_token not in question[i]:
+                        question[i] = question[i] + self.tokenizer._prompt_end_token
+                    if not question[i].startswith("<s>"):
+                        question[i] = "<s>" + question[i]
+            return question
+
+        def _postprocess(output, question):
+            output = output.replace("<s>", "").replace(question, "").replace("</s>", "").replace("<pad>", "")
+            if self.model_args.get("extra_answer_tokens", False):
+                output = output.split(self.tokenizer._prompt_end_token)[-1]
+            return output
+
+        if isinstance(question, list):
+            image_tensor_list = []
+            for i in image:
+                image_tensor, ori_size = _preprocess_image(i)
+                image_tensor_list.append(image_tensor)
+            image_tensor = torch.cat(image_tensor_list, dim=0)
+
+            question = _preprocess_prompt_batch(question)
+            self.processor.tokenizer.padding_side = "left"
+            prompt_ids = self.processor.tokenizer(
+                question, add_special_tokens=False, return_tensors="pt", padding=True
+            ).input_ids
+        else:
+            image_tensor, ori_size = _preprocess_image(image)
+            prompt_ids = _preprocess_prompt(question)
+
+        if only_return_img_size:
+            return ori_size
+
+        model_output_batch = []
+        for i in range(0, image_tensor.shape[0], max_batch_size):
+            image_tensor_batch = image_tensor[i : i + max_batch_size]
+            prompt_ids_batch = prompt_ids[i : i + max_batch_size]
+            model_output = self.model.inference(image_tensors=image_tensor_batch, prompt_ids=prompt_ids_batch)
+            model_output_batch.append(model_output)
+        model_output = {}
+        for k, v in model_output_batch[0].items():
+            if isinstance(v, torch.Tensor):
+                model_output[k] = sum(
+                    [v_batch[k].cpu().numpy().tolist() for v_batch in model_output_batch],
+                    [],
+                )
+            else:
+                model_output[k] = sum([v_batch[k] for v_batch in model_output_batch], [])
+
+        if return_raw:
+            if return_img_size:
+                return model_output, ori_size
+            return model_output
+        else:
+            if isinstance(question, list):
+                output = [_postprocess(model_output["repetitions"][i], question[i]) for i in range(len(question))]
+                score = model_output["scores"]
+            else:
+                output = _postprocess(model_output["repetitions"][0], question)
+                score = model_output["scores"][0]
+            if return_score:
+                return output, score
+            if return_img_size:
+                return output, ori_size
+            return output

config/Dolphin.yaml

@@ -0,0 +1,17 @@
+model:
+  model_name_or_path: "./checkpoints/dolphin_model.bin"
+  tokenizer_path: "./checkpoints/dolphin_tokenizer.json"
+  extra_answer_tokens: True   # add <Answer/> token
+  max_length: 4096
+  decoder_layer: 10
+  max_position_embeddings: 4096
+  hidden_dimension: 1024
+  swin_args:
+    name: 'swin'
+    img_size: [896, 896]
+    patch_size: 4
+    embed_dim: 128
+    align_long_axis: False
+    window_size: 7
+    encoder_layer: [2, 2, 14, 2]
+    num_heads: [4, 8, 16, 32]

demo/element_imgs/markdown/table_1.md

@@ -0,0 +1,2 @@
+<table><tr><td></td><td></td><td>100-class (top-1 acc.)</td><td>1000-class (top-1 acc.)</td></tr><tr><td colspan="2">4096-d (float)</td><td>77.1 ± 1.5</td><td>65.0</td></tr><tr><td rowspan="3">1024 bits</td><td>BP</td><td>72.9 ± 1.3</td><td>58.1</td></tr><tr><td>CBE</td><td>73.0 ± 1.3</td><td>59.2</td></tr><tr><td>SP</td><td>73.8 ± 1.3</td><td>60.1</td></tr><tr><td rowspan="4">4096 bits</td><td>threshold [1]</td><td>73.5 ± 1.4</td><td>59.1</td></tr><tr><td>BP</td><td>76.0 ± 1.5</td><td>63.2</td></tr><tr><td>CBE</td><td>75.9 ± 1.4</td><td>63.0</td></tr><tr><td>SP</td><td>76.3 ± 1.5</td><td>63.3</td></tr><tr><td>8192 bits</td><td>SP</td><td>76.8 ± 1.4</td><td>64.2</td></tr><tr><td>16384 bits</td><td>SP</td><td>77.1 ± 1.6</td><td>64.5</td></tr></table>
+

demo/element_imgs/recognition_json/table_1.json

@@ -0,0 +1,6 @@
+[
+  {
+    "label": "tab",
+    "text": "<table><tr><td></td><td></td><td>100-class (top-1 acc.)</td><td>1000-class (top-1 acc.)</td></tr><tr><td colspan=\"2\">4096-d (float)</td><td>77.1 ± 1.5</td><td>65.0</td></tr><tr><td rowspan=\"3\">1024 bits</td><td>BP</td><td>72.9 ± 1.3</td><td>58.1</td></tr><tr><td>CBE</td><td>73.0 ± 1.3</td><td>59.2</td></tr><tr><td>SP</td><td>73.8 ± 1.3</td><td>60.1</td></tr><tr><td rowspan=\"4\">4096 bits</td><td>threshold [1]</td><td>73.5 ± 1.4</td><td>59.1</td></tr><tr><td>BP</td><td>76.0 ± 1.5</td><td>63.2</td></tr><tr><td>CBE</td><td>75.9 ± 1.4</td><td>63.0</td></tr><tr><td>SP</td><td>76.3 ± 1.5</td><td>63.3</td></tr><tr><td>8192 bits</td><td>SP</td><td>76.8 ± 1.4</td><td>64.2</td></tr><tr><td>16384 bits</td><td>SP</td><td>77.1 ± 1.6</td><td>64.5</td></tr></table>"
+  }
+]

demo/page_imgs/markdown/test_page3.md

@@ -0,0 +1,22 @@
+![Figure](figures/test_page3_figure_000.png)
+
+Figure 2: (left) Scaled Dot-Product Attention. (right) Multi-Head Attention consists of several attention layers running in parallel.
+
+query with all keys, divide each by $\sqrt{d_k}$ , and apply a softmax function to obtain the weights on the values.
+
+In practice, we compute the attention function on a set of queries simultaneously, packed together into a matrix $Q$ . The keys and values are also packed together into matrices $K$ and $V$ . We compute the matrix of outputs as: $$ \\     \text{Attention}(Q, K, V) = \mathrm{softmax}(\frac{QK^T}{\sqrt{d_k}})V \\ $$
+
+The two most commonly used attention functions are additive attention [2] , and dot-product (multiplicative) attention. Dot-product attention is identical to our algorithm, except for the scaling factor of $\frac{1}{\sqrt{d_k}}$ . Additive attention computes the compatibility function using a feed-forward network with a single hidden layer. While the two are similar in theoretical complexity, dot-product attention is much faster and more space-efficient in practice, since it can be implemented using highly optimized matrix multiplication code.
+
+While for small values of $d_k$ the two mechanisms perform similarly, additive attention outperforms dot product attention without scaling for larger values of $d_k$ [ 3 ] . We suspect that for large values of $d_k$ , the dot products grow large in magnitude, pushing the softmax function into regions where it has extremely small gradients 4 To counteract this effect, we scale the dot products by $\frac{1}{\sqrt{d_k}}$ .
+
+3.2.2 Multi-Head Attention
+
+Instead of performing a single attention function with $d_{\text{model}}$ -dimensional keys, values and queries, we found it beneficial to linearly project the queries, keys and values $h$ times with different, learned linear projections to $d_k$ , $d_k$ and $d_v$ dimensions, respectively. On each of these projected versions of queries, keys and values we then perform the attention function in parallel, yielding $d_v$ -dimensional output values. These are concatenated and once again projected, resulting in the final values, as depicted in Figure 2 .
+
+Multi­head attention allows the model to jointly attend to information from different representation subspaces at different positions. With a single attention head, averaging inhibits this.
+
+${ }^{4}$ To illustrate why the dot products get large, assume that the components of $q$ and $k$ are independent random variables with mean 0 and variance 1 . Then their dot product, $q \cdot k=\sum_{i=1}^{d_{k}} q_{i} k_{i}$, has mean 0 and variance $d_{k}$.
+
+4
+

demo/page_imgs/recognition_json/page_1.json

@@ -0,0 +1,178 @@
+[
+  {
+    "label": "title",
+    "bbox": [
+      271,
+      188,
+      1194,
+      221
+    ],
+    "text": "LLaMA: Open and Efficient Foundation Language Models",
+    "reading_order": 0
+  },
+  {
+    "label": "author",
+    "bbox": [
+      313,
+      272,
+      1154,
+      317
+    ],
+    "text": "Hugo Touvron; Thibaut Lavril*, Gautier Izacard*, Xavier Martinet",
+    "reading_order": 1
+  },
+  {
+    "label": "para",
+    "bbox": [
+      269,
+      317,
+      1201,
+      425
+    ],
+    "text": "Marie-Anne Lachaux, Timothee Lacroix, Baptiste Rozière, Naman Goyal\nEric Hambro, Faisal Azhar, Aurelien Rodriguez, Armand Joulin\nEdouard Grave*Guillaume Lample*",
+    "reading_order": 2
+  },
+  {
+    "label": "para",
+    "bbox": [
+      685,
+      440,
+      795,
+      482
+    ],
+    "text": "Meta AI",
+    "reading_order": 3
+  },
+  {
+    "label": "sec",
+    "bbox": [
+      376,
+      524,
+      502,
+      565
+    ],
+    "text": "\\begin{abstract}",
+    "reading_order": 4
+  },
+  {
+    "label": "para",
+    "bbox": [
+      209,
+      586,
+      675,
+      946
+    ],
+    "text": "We introduce LLaMA, a collection of founda-\ntion language models ranging from 7B to 65B\nparameters. We train our models on trillions\nof tokens, and show that it is possible to train\nstate-of-the-art models using publicly avail-\nable datasets exclusively, without resorting\nto proprietary and inaccessible datasets. In\nparticular, LLaMA-13B outperforms GPT-3\n(175B) on most benchmarks, and LLaMA-\n65B is competitive with the best models,\nChinchilla-70B and PaLM-540B. We release\nall our models to the research community $^1$ .",
+    "reading_order": 5
+  },
+  {
+    "label": "sec",
+    "bbox": [
+      167,
+      964,
+      376,
+      1006
+    ],
+    "text": "1 Introduction",
+    "reading_order": 6
+  },
+  {
+    "label": "para",
+    "bbox": [
+      167,
+      1027,
+      718,
+      1498
+    ],
+    "text": "Large Languages Models (LLMs) trained on mas-\nsive corpora of texts have shown their ability to per-\nform new tasks from textual instructions or from a\nfew examples ( Brown et al. , 2020 ) . These few-shot\nproperties first appeared when scaling models to a\nsufficient size ( Kaplan et al. , 2020 ) , resulting in a\nline of work that focuses on further scaling these\nmodels ( Chowdhery et al. , 2022 ; Rae et al. , 2021 ) .\nThese efforts are based on the assumption that\nmore parameters will lead to better performance.\nHowever, recent work from Hoffmann et al. ( 2022 )\nshows that, for a given compute budget, the best\nperformances are not achieved by the largest mod-\nels, but by smaller models trained on more data.",
+    "reading_order": 7
+  },
+  {
+    "label": "para",
+    "bbox": [
+      167,
+      1506,
+      717,
+      1844
+    ],
+    "text": "The objective of the scaling laws from Hoff-\nmann et al. ( 2022 ) is to determine how to best\nscale the dataset and model sizes for a particular\ntraining compute budget. However, this objective\ndisregards the inference budget, which becomes\ncritical when serving a language model at scale.\nIn this context, given a target level of performance,\nthe preferred model is not the fastest to train but the\nfastest at inference, and although it may be cheaper\nto train a large model to reach a certain level of",
+    "reading_order": 8
+  },
+  {
+    "label": "para",
+    "bbox": [
+      753,
+      539,
+      1304,
+      734
+    ],
+    "text": "performance, a smaller one trained longer will\nultimately be cheaper at inference. For instance,\nalthough Hoffmann et al. ( 2022 ) recommends\ntraining a 10B model on 200B tokens, we find\nthat the performance of a 7B model continues to\nimprove even after 1T tokens.",
+    "reading_order": 9
+  },
+  {
+    "label": "para",
+    "bbox": [
+      753,
+      769,
+      1305,
+      1236
+    ],
+    "text": "The focus of this work is to train a series of\nlanguage models that achieve the best possible per-\nformance at various inference budgets, by training\non more tokens than what is typically used. The\nresulting models, called LLaMA , ranges from 7B\nto 65B parameters with competitive performance\ncompared to the best existing LLMs. For instance,\nLLaMA-13B outperforms GPT-3 on most bench-\nmarks, despite being 10 $\\times$ smaller. We believe that\nthis model will help democratize the access and\nstudy of LLMs, since it can be run on a single GPU.\nAt the higher-end of the scale, our 65B-parameter\nmodel is also competitive with the best large lan-\nguage models such as Chinchilla or PaLM-540B.",
+    "reading_order": 10
+  },
+  {
+    "label": "para",
+    "bbox": [
+      753,
+      1257,
+      1305,
+      1601
+    ],
+    "text": "Unlike Chinchilla, PaLM, or GPT-3, we only\nuse publicly available data, making our work com-\npatible with open-sourcing, while most existing\nmodels rely on data which is either not publicly\navailable or undocumented (e.g. “ Books – 2TB ” or\n“ Social media conversations ” ). There exist some\nexceptions, notably OPT ( Zhang et al. , 2022 ) ,\nGPT-NeoX ( Black et al. , 2022 ) , BLOOM ( Scao\net al. , 2022 ) and GLM ( Zeng et al. , 2022 ) , but none\nthat are competitive with PaLM-62B or Chinchilla.",
+    "reading_order": 11
+  },
+  {
+    "label": "para",
+    "bbox": [
+      753,
+      1634,
+      1304,
+      1933
+    ],
+    "text": "In the rest of this paper, we present an overview\nof the modifications we made to the transformer\narchitecture ( Vaswani et al. , 2017 ) , as well as our\ntraining method. We then report the performance of\nour models and compare with others LLMs on a set\nof standard benchmarks. Finally, we expose some\nof the biases and toxicity encoded in our models,\nusing some of the most recent benchmarks from\nthe responsible AI community.",
+    "reading_order": 12
+  },
+  {
+    "label": "fnote",
+    "bbox": [
+      167,
+      1844,
+      712,
+      1907
+    ],
+    "text": "* Equal contribution.\nCorrespondence:\n{htouvron\nthibautlav,gizacard,egrave,glample}@meta.com",
+    "reading_order": 13
+  },
+  {
+    "label": "fnote",
+    "bbox": [
+      209,
+      1907,
+      632,
+      1931
+    ],
+    "text": "https://github.com/facebookresearch/llama",
+    "reading_order": 14
+  },
+  {
+    "label": "watermark",
+    "bbox": [
+      20,
+      649,
+      83,
+      1530
+    ],
+    "text": "arXiv:2302.1397lvl [cs.CL] 27 Feb 2023",
+    "reading_order": 15
+  }
+]

demo/page_imgs/recognition_json/test_page.json

@@ -0,0 +1,47 @@
+[
+  {
+    "label": "header",
+    "bbox": [
+      291,
+      90,
+      675,
+      120
+    ],
+    "text": "Scaled Dot-Product Attention",
+    "reading_order": 0
+  },
+  {
+    "label": "fig",
+    "text": "![Figure](figures/test_page_figure_001.png)",
+    "figure_path": "figures/test_page_figure_001.png",
+    "bbox": [
+      1274,
+      105,
+      1536,
+      627
+    ],
+    "reading_order": 1
+  },
+  {
+    "label": "cap",
+    "bbox": [
+      168,
+      719,
+      1413,
+      789
+    ],
+    "text": "Figure 2: (left) Scaled Dot-Product Attention. (right) Multi-Head Attention consists of several\nattention layers running in parallel.",
+    "reading_order": 2
+  },
+  {
+    "label": "para",
+    "bbox": [
+      168,
+      858,
+      1413,
+      934
+    ],
+    "text": "query with all keys, divide each by $\\sqrt{d_{k}}$, and apply a softmax function to obtain the weights on the\nvalues.",
+    "reading_order": 3
+  }
+]

demo/page_imgs/recognition_json/test_page2.json

@@ -0,0 +1,102 @@
+[
+  {
+    "label": "fig",
+    "text": "![Figure](figures/test_page2_figure_000.png)",
+    "figure_path": "figures/test_page2_figure_000.png",
+    "bbox": [
+      394,
+      117,
+      897,
+      837
+    ],
+    "reading_order": 0
+  },
+  {
+    "label": "cap",
+    "bbox": [
+      445,
+      852,
+      856,
+      873
+    ],
+    "text": "Figure 1: The Transformer - model architecture",
+    "reading_order": 1
+  },
+  {
+    "label": "para",
+    "bbox": [
+      218,
+      920,
+      1086,
+      1044
+    ],
+    "text": "wise fully connected feed-forward network. We employ a residual connection [ 10 ] around each of\nthe two sub-layers, followed by layer normalization [ 1 ] . That is, the output of each sub-layer is\n$\\mathrm{LayerNorm}(x+\\mathrm{Sublayer}(x))$ , where $\\mathrm{Sublayer}(x)$ is the function implemented by the sub-layer\nitself. To facilitate these residual connections, all sub-layers in the model, as well as the embedding\nlayers, produce outputs of dimension $d_{\\text{model}}=512$ .",
+    "reading_order": 2
+  },
+  {
+    "label": "para",
+    "bbox": [
+      218,
+      1071,
+      1085,
+      1244
+    ],
+    "text": "The The decoder is also composed of a stack of $N=6$ identical layers. In addition to the two\nsub-layers in each encoder layer, the decoder inserts a third sub-layer, which performs multi-head\nattention over the output of the encoder stack. Similar to the encoder, we employ residual connections\naround each of the sub-layers, followed by layer normalization. We also modify the self-attention\nsub-layer in the decoder stack to prevent positions from attending to subsequent positions. This\nmasking, combined with fact that the output embeddings are offset by one position, ensures that the\npredictions for position $i$ can depend only on the known outputs at positions less than $i$ .",
+    "reading_order": 3
+  },
+  {
+    "label": "sub_sec",
+    "bbox": [
+      226,
+      1283,
+      344,
+      1305
+    ],
+    "text": "3.2 Attention",
+    "reading_order": 4
+  },
+  {
+    "label": "para",
+    "bbox": [
+      218,
+      1322,
+      1087,
+      1422
+    ],
+    "text": "An attention function can be described as mapping a query and a set of key-value pairs to an output,\nwhere the query, keys, values, and output are all vectors. The output is computed as a weighted sum\nof the values, where the weight assigned to each value is computed by a compatibility function of the\nquery with the corresponding key.",
+    "reading_order": 5
+  },
+  {
+    "label": "sub_sub_sec",
+    "bbox": [
+      218,
+      1456,
+      562,
+      1474
+    ],
+    "text": "3.2.1 Scaled Dot-Product Attention",
+    "reading_order": 6
+  },
+  {
+    "label": "para",
+    "bbox": [
+      218,
+      1498,
+      1085,
+      1546
+    ],
+    "text": "We call our particular attention \"Scaled Dot-Product Attention\" (Figure 2 ). The input consists of\nqueries and keys of dimension $d_k$ , and values of dimension $d_v$ . We compute the dot products of the",
+    "reading_order": 7
+  },
+  {
+    "label": "foot",
+    "bbox": [
+      646,
+      1590,
+      662,
+      1607
+    ],
+    "text": "3",
+    "reading_order": 8
+  }
+]

demo/page_imgs/recognition_json/test_page3.json

@@ -0,0 +1,124 @@
+[
+  {
+    "label": "fig",
+    "text": "![Figure](figures/test_page3_figure_000.png)",
+    "figure_path": "figures/test_page3_figure_000.png",
+    "bbox": [
+      331,
+      134,
+      984,
+      489
+    ],
+    "reading_order": 0
+  },
+  {
+    "label": "cap",
+    "bbox": [
+      198,
+      554,
+      1065,
+      603
+    ],
+    "text": "Figure 2: (left) Scaled Dot-Product Attention. (right) Multi-Head Attention consists of several\nattention layers running in parallel.",
+    "reading_order": 1
+  },
+  {
+    "label": "para",
+    "bbox": [
+      198,
+      652,
+      1065,
+      701
+    ],
+    "text": "query with all keys, divide each by $\\sqrt{d_k}$ , and apply a softmax function to obtain the weights on the\nvalues.",
+    "reading_order": 2
+  },
+  {
+    "label": "para",
+    "bbox": [
+      198,
+      715,
+      1065,
+      881
+    ],
+    "text": "In practice, we compute the attention function on a set of queries simultaneously, packed together\ninto a matrix $Q$ . The keys and values are also packed together into matrices $K$ and $V$ . We compute\nthe matrix of outputs as:\n\\[\n    \\text{Attention}(Q, K, V) = \\mathrm{softmax}(\\frac{QK^T}{\\sqrt{d_k}})V\n\\]",
+    "reading_order": 3
+  },
+  {
+    "label": "para",
+    "bbox": [
+      198,
+      913,
+      1068,
+      1060
+    ],
+    "text": "The two most commonly used attention functions are additive attention [2] , and dot-product (multi-\nplicative) attention. Dot-product attention is identical to our algorithm, except for the scaling factor\nof $\\frac{1}{\\sqrt{d_k}}$ . Additive attention computes the compatibility function using a feed-forward network with\na single hidden layer. While the two are similar in theoretical complexity, dot-product attention is\nmuch faster and more space-efficient in practice, since it can be implemented using highly optimized\nmatrix multiplication code.",
+    "reading_order": 4
+  },
+  {
+    "label": "para",
+    "bbox": [
+      198,
+      1074,
+      1066,
+      1175
+    ],
+    "text": "While for small values of $d_k$ the two mechanisms perform similarly, additive attention outperforms\ndot product attention without scaling for larger values of $d_k$ [ 3 ] . We suspect that for large values of\n$d_k$ , the dot products grow large in magnitude, pushing the softmax function into regions where it has\nextremely small gradients 4 To counteract this effect, we scale the dot products by $\\frac{1}{\\sqrt{d_k}}$ .",
+    "reading_order": 5
+  },
+  {
+    "label": "sub_sub_sec",
+    "bbox": [
+      198,
+      1207,
+      467,
+      1225
+    ],
+    "text": "3.2.2 Multi-Head Attention",
+    "reading_order": 6
+  },
+  {
+    "label": "para",
+    "bbox": [
+      198,
+      1253,
+      1067,
+      1395
+    ],
+    "text": "Instead of performing a single attention function with $d_{\\text{model}}$ -dimensional keys, values and queries,\nwe found it beneficial to linearly project the queries, keys and values $h$ times with different, learned\nlinear projections to $d_k$ , $d_k$ and $d_v$ dimensions, respectively. On each of these projected versions of\nqueries, keys and values we then perform the attention function in parallel, yielding $d_v$ -dimensional\noutput values. These are concatenated and once again projected, resulting in the final values, as\ndepicted in Figure 2 .",
+    "reading_order": 7
+  },
+  {
+    "label": "para",
+    "bbox": [
+      198,
+      1403,
+      1065,
+      1453
+    ],
+    "text": "Multi­head attention allows the model to jointly attend to information from different representation\nsubspaces at different positions. With a single attention head, averaging inhibits this.",
+    "reading_order": 8
+  },
+  {
+    "label": "fnote",
+    "bbox": [
+      198,
+      1485,
+      1065,
+      1535
+    ],
+    "text": "${ }^{4}$ To illustrate why the dot products get large, assume that the components of $q$ and $k$ are independent random\nvariables with mean 0 and variance 1 . Then their dot product, $q \\cdot k=\\sum_{i=1}^{d_{k}} q_{i} k_{i}$, has mean 0 and variance $d_{k}$.",
+    "reading_order": 9
+  },
+  {
+    "label": "foot",
+    "bbox": [
+      625,
+      1578,
+      641,
+      1599
+    ],
+    "text": "4",
+    "reading_order": 10
+  }
+]

demo_element.py

@@ -0,0 +1,129 @@
+""" 
+Copyright (c) 2025 Bytedance Ltd. and/or its affiliates
+SPDX-License-Identifier: MIT
+"""
+
+import argparse
+import glob
+import os
+
+from omegaconf import OmegaConf
+from PIL import Image
+
+from chat import DOLPHIN
+from utils.utils import *
+
+
+def process_element(image_path, model, element_type, save_dir=None):
+    """Process a single element image (text, table, formula)
+
+    Args:
+        image_path: Path to the element image
+        model: DOLPHIN model instance
+        element_type: Type of element ('text', 'table', 'formula')
+        save_dir: Directory to save results (default: same as input directory)
+
+    Returns:
+        Parsed content of the element and recognition results
+    """
+    # Load and prepare image
+    pil_image = Image.open(image_path).convert("RGB")
+    pil_image = crop_margin(pil_image)
+
+    # Select appropriate prompt based on element type
+    if element_type == "table":
+        prompt = "Parse the table in the image."
+        label = "tab"
+    elif element_type == "formula":
+        prompt = "Read text in the image."
+        label = "formula"
+    else:  # Default to text
+        prompt = "Read text in the image."
+        label = "text"
+
+    # Process the element
+    result = model.chat(prompt, pil_image)
+
+    # Create recognition result in the same format as the document parser
+    recognition_result = [
+        {
+            "label": label,
+            "text": result.strip(),
+        }
+    ]
+
+    # Save results if save_dir is provided
+    if save_dir:
+        save_outputs(recognition_result, image_path, save_dir)
+        print(f"Results saved to {save_dir}")
+
+    return result, recognition_result
+
+
+def main():
+    parser = argparse.ArgumentParser(description="Element-level processing using DOLPHIN model")
+    parser.add_argument("--config", default="./config/Dolphin.yaml", help="Path to configuration file")
+    parser.add_argument("--input_path", type=str, required=True, help="Path to input image or directory of images")
+    parser.add_argument(
+        "--element_type",
+        type=str,
+        choices=["text", "table", "formula"],
+        default="text",
+        help="Type of element to process (text, table, formula)",
+    )
+    parser.add_argument(
+        "--save_dir",
+        type=str,
+        default=None,
+        help="Directory to save parsing results (default: same as input directory)",
+    )
+    parser.add_argument("--print_results", action="store_true", help="Print recognition results to console")
+    args = parser.parse_args()
+
+    # Load Model
+    config = OmegaConf.load(args.config)
+    model = DOLPHIN(config)
+
+    # Set save directory
+    save_dir = args.save_dir or (
+        args.input_path if os.path.isdir(args.input_path) else os.path.dirname(args.input_path)
+    )
+    setup_output_dirs(save_dir)
+
+    # Collect Images
+    if os.path.isdir(args.input_path):
+        image_files = []
+        for ext in [".jpg", ".jpeg", ".png", ".JPG", ".JPEG", ".PNG"]:
+            image_files.extend(glob.glob(os.path.join(args.input_path, f"*{ext}")))
+        image_files = sorted(image_files)
+    else:
+        if not os.path.exists(args.input_path):
+            raise FileNotFoundError(f"Input path {args.input_path} does not exist")
+        image_files = [args.input_path]
+
+    total_samples = len(image_files)
+    print(f"\nTotal samples to process: {total_samples}")
+
+    # Process images one by one
+    for image_path in image_files:
+        print(f"\nProcessing {image_path}")
+        try:
+            result, recognition_result = process_element(
+                image_path=image_path,
+                model=model,
+                element_type=args.element_type,
+                save_dir=save_dir,
+            )
+
+            if args.print_results:
+                print("\nRecognition result:")
+                print(result)
+                print("-" * 40)
+
+        except Exception as e:
+            print(f"Error processing {image_path}: {str(e)}")
+            continue
+
+
+if __name__ == "__main__":
+    main()

demo_element_hf.py

@@ -0,0 +1,195 @@
+"""
+Copyright (c) 2025 Bytedance Ltd. and/or its affiliates
+SPDX-License-Identifier: MIT
+"""
+
+import argparse
+import glob
+import os
+
+import torch
+from PIL import Image
+from transformers import AutoProcessor, VisionEncoderDecoderModel
+
+from utils.utils import *
+
+
+class DOLPHIN:
+    def __init__(self, model_id_or_path):
+        """Initialize the Hugging Face model
+        
+        Args:
+            model_id_or_path: Path to local model or Hugging Face model ID
+        """
+        # Load model from local path or Hugging Face hub
+        self.processor = AutoProcessor.from_pretrained(model_id_or_path)
+        self.model = VisionEncoderDecoderModel.from_pretrained(model_id_or_path)
+        self.model.eval()
+        
+        # Set device and precision
+        self.device = "cuda" if torch.cuda.is_available() else "cpu"
+        self.model.to(self.device)
+        self.model = self.model.half()  # Always use half precision by default
+        
+        # set tokenizer
+        self.tokenizer = self.processor.tokenizer
+        
+    def chat(self, prompt, image):
+        """Process an image with the given prompt
+        
+        Args:
+            prompt: Text prompt to guide the model
+            image: PIL Image to process
+            
+        Returns:
+            Generated text from the model
+        """
+        # Prepare image
+        pixel_values = self.processor(image, return_tensors="pt").pixel_values
+        pixel_values = pixel_values.half()
+            
+        # Prepare prompt
+        prompt = f"<s>{prompt} <Answer/>"
+        prompt_ids = self.tokenizer(
+            prompt, 
+            add_special_tokens=False, 
+            return_tensors="pt"
+        ).input_ids.to(self.device)
+        
+        decoder_attention_mask = torch.ones_like(prompt_ids)
+        
+        # Generate text
+        outputs = self.model.generate(
+            pixel_values=pixel_values.to(self.device),
+            decoder_input_ids=prompt_ids,
+            decoder_attention_mask=decoder_attention_mask,
+            min_length=1,
+            max_length=4096,
+            pad_token_id=self.tokenizer.pad_token_id,
+            eos_token_id=self.tokenizer.eos_token_id,
+            use_cache=True,
+            bad_words_ids=[[self.tokenizer.unk_token_id]],
+            return_dict_in_generate=True,
+            do_sample=False,
+            num_beams=1,
+            repetition_penalty=1.1,
+            temperature=1.0
+        )
+        
+        # Process the output
+        sequence = self.tokenizer.batch_decode(outputs.sequences, skip_special_tokens=False)[0]
+        sequence = sequence.replace(prompt, "").replace("<pad>", "").replace("</s>", "").strip()
+        
+        return sequence
+
+def process_element(image_path, model, element_type, save_dir=None):
+    """Process a single element image (text, table, formula)
+    
+    Args:
+        image_path: Path to the element image
+        model: HFModel model instance
+        element_type: Type of element ('text', 'table', 'formula')
+        save_dir: Directory to save results (default: same as input directory)
+        
+    Returns:
+        Parsed content of the element and recognition results
+    """
+    # Load and prepare image
+    pil_image = Image.open(image_path).convert("RGB")
+    pil_image = crop_margin(pil_image)
+    
+    # Select appropriate prompt based on element type
+    if element_type == "table":
+        prompt = "Parse the table in the image."
+        label = "tab"
+    elif element_type == "formula":
+        prompt = "Read text in the image."
+        label = "formula"
+    else:  # Default to text
+        prompt = "Read text in the image."
+        label = "text"
+    
+    # Process the element
+    result = model.chat(prompt, pil_image)
+    
+    # Create recognition result in the same format as the document parser
+    recognition_result = [
+        {
+            "label": label,
+            "text": result.strip(),
+        }
+    ]
+    
+    # Save results if save_dir is provided
+    if save_dir:
+        save_outputs(recognition_result, image_path, save_dir)
+        print(f"Results saved to {save_dir}")
+    
+    return result, recognition_result
+
+
+def main():
+    parser = argparse.ArgumentParser(description="Element-level processing using DOLPHIN model")
+    parser.add_argument("--model_path", default="./hf_model", help="Path to Hugging Face model")
+    parser.add_argument("--input_path", type=str, required=True, help="Path to input image or directory of images")
+    parser.add_argument(
+        "--element_type",
+        type=str,
+        choices=["text", "table", "formula"],
+        default="text",
+        help="Type of element to process (text, table, formula)",
+    )
+    parser.add_argument(
+        "--save_dir",
+        type=str,
+        default=None,
+        help="Directory to save parsing results (default: same as input directory)",
+    )
+    parser.add_argument("--print_results", action="store_true", help="Print recognition results to console")
+    args = parser.parse_args()
+    
+    # Load Model
+    model = DOLPHIN(args.model_path)
+    
+    # Set save directory
+    save_dir = args.save_dir or (
+        args.input_path if os.path.isdir(args.input_path) else os.path.dirname(args.input_path)
+    )
+    setup_output_dirs(save_dir)
+    
+    # Collect Images
+    if os.path.isdir(args.input_path):
+        image_files = []
+        for ext in [".jpg", ".jpeg", ".png", ".JPG", ".JPEG", ".PNG"]:
+            image_files.extend(glob.glob(os.path.join(args.input_path, f"*{ext}")))
+        image_files = sorted(image_files)
+    else:
+        if not os.path.exists(args.input_path):
+            raise FileNotFoundError(f"Input path {args.input_path} does not exist")
+        image_files = [args.input_path]
+    
+    total_samples = len(image_files)
+    print(f"\nTotal samples to process: {total_samples}")
+    
+    # Process images one by one
+    for image_path in image_files:
+        print(f"\nProcessing {image_path}")
+        try:
+            result, recognition_result = process_element(
+                image_path=image_path,
+                model=model,
+                element_type=args.element_type,
+                save_dir=save_dir,
+            )
+
+            if args.print_results:
+                print("\nRecognition result:")
+                print(result)
+                print("-" * 40)
+        except Exception as e:
+            print(f"Error processing {image_path}: {str(e)}")
+            continue
+
+
+if __name__ == "__main__":
+    main()

demo_page.py

@@ -0,0 +1,247 @@
+""" 
+Copyright (c) 2025 Bytedance Ltd. and/or its affiliates
+SPDX-License-Identifier: MIT
+"""
+
+import argparse
+import glob
+import os
+
+import cv2
+from omegaconf import OmegaConf
+from PIL import Image
+
+from chat import DOLPHIN
+from utils.utils import *
+
+
+def process_document(document_path, model, save_dir, max_batch_size):
+    """Parse documents - Handles both images and PDFs"""
+    file_ext = os.path.splitext(document_path)[1].lower()
+    
+    if file_ext == '.pdf':
+        # Process PDF file
+        # Convert PDF to images
+        images = convert_pdf_to_images(document_path)
+        if not images:
+            raise Exception(f"Failed to convert PDF {document_path} to images")
+        
+        all_results = []
+        
+        # Process each page
+        for page_idx, pil_image in enumerate(images):
+            print(f"Processing page {page_idx + 1}/{len(images)}")
+            
+            # Generate output name for this page
+            base_name = os.path.splitext(os.path.basename(document_path))[0]
+            page_name = f"{base_name}_page_{page_idx + 1:03d}"
+            
+            # Process this page (don't save individual page results)
+            json_path, recognition_results = process_single_image(
+                pil_image, model, save_dir, page_name, max_batch_size, save_individual=False
+            )
+            
+            # Add page information to results
+            page_results = {
+                "page_number": page_idx + 1,
+                "elements": recognition_results
+            }
+            all_results.append(page_results)
+        
+        # Save combined results for multi-page PDF
+        combined_json_path = save_combined_pdf_results(all_results, document_path, save_dir)
+        
+        return combined_json_path, all_results
+
+    else:
+        # Process regular image file
+        pil_image = Image.open(document_path).convert("RGB")
+        base_name = os.path.splitext(os.path.basename(document_path))[0]
+        return process_single_image(pil_image, model, save_dir, base_name, max_batch_size)
+
+
+def process_single_image(image, model, save_dir, image_name, max_batch_size, save_individual=True):
+    """Process a single image (either from file or converted from PDF page)
+    
+    Args:
+        image: PIL Image object
+        model: DOLPHIN model instance
+        save_dir: Directory to save results
+        image_name: Name for the output file
+        max_batch_size: Maximum batch size for processing
+        save_individual: Whether to save individual results (False for PDF pages)
+        
+    Returns:
+        Tuple of (json_path, recognition_results)
+    """
+    # Stage 1: Page-level layout and reading order parsing
+    layout_output = model.chat("Parse the reading order of this document.", image)
+
+    # Stage 2: Element-level content parsing
+    padded_image, dims = prepare_image(image)
+    recognition_results = process_elements(layout_output, padded_image, dims, model, max_batch_size, save_dir, image_name)
+
+    # Save outputs only if requested (skip for PDF pages)
+    json_path = None
+    if save_individual:
+        # Create a dummy image path for save_outputs function
+        dummy_image_path = f"{image_name}.jpg"  # Extension doesn't matter, only basename is used
+        json_path = save_outputs(recognition_results, dummy_image_path, save_dir)
+
+    return json_path, recognition_results
+
+
+def process_elements(layout_results, padded_image, dims, model, max_batch_size, save_dir=None, image_name=None):
+    """Parse all document elements with parallel decoding"""
+    layout_results = parse_layout_string(layout_results)
+
+    text_table_elements = []  # Elements that need processing
+    figure_results = []  # Figure elements (no processing needed)
+    previous_box = None
+    reading_order = 0
+
+    # Collect elements for processing
+    for bbox, label in layout_results:
+        try:
+            # Adjust coordinates
+            x1, y1, x2, y2, orig_x1, orig_y1, orig_x2, orig_y2, previous_box = process_coordinates(
+                bbox, padded_image, dims, previous_box
+            )
+
+            # Crop and parse element
+            cropped = padded_image[y1:y2, x1:x2]
+            if cropped.size > 0 and cropped.shape[0] > 3 and cropped.shape[1] > 3:
+                if label == "fig":
+                    pil_crop = Image.fromarray(cv2.cvtColor(cropped, cv2.COLOR_BGR2RGB))
+                    
+                    figure_filename = save_figure_to_local(pil_crop, save_dir, image_name, reading_order)
+                    
+                    # For figure regions, store relative path instead of base64
+                    figure_results.append(
+                        {
+                            "label": label,
+                            "text": f"![Figure](figures/{figure_filename})",
+                            "figure_path": f"figures/{figure_filename}",
+                            "bbox": [orig_x1, orig_y1, orig_x2, orig_y2],
+                            "reading_order": reading_order,
+                        }
+                    )
+                else:
+                    # For text or table regions, prepare for parsing
+                    pil_crop = Image.fromarray(cv2.cvtColor(cropped, cv2.COLOR_BGR2RGB))
+                    prompt = "Parse the table in the image." if label == "tab" else "Read text in the image."
+                    text_table_elements.append(
+                        {
+                            "crop": pil_crop,
+                            "prompt": prompt,
+                            "label": label,
+                            "bbox": [orig_x1, orig_y1, orig_x2, orig_y2],
+                            "reading_order": reading_order,
+                        }
+                    )
+
+            reading_order += 1
+
+        except Exception as e:
+            print(f"Error processing bbox with label {label}: {str(e)}")
+            continue
+
+    # Parse text/table elements in parallel
+    recognition_results = figure_results
+    if text_table_elements:
+        crops_list = [elem["crop"] for elem in text_table_elements]
+        prompts_list = [elem["prompt"] for elem in text_table_elements]
+
+        # Inference in batch
+        batch_results = model.chat(prompts_list, crops_list, max_batch_size=max_batch_size)
+
+        # Add batch results to recognition_results
+        for i, result in enumerate(batch_results):
+            elem = text_table_elements[i]
+            recognition_results.append(
+                {
+                    "label": elem["label"],
+                    "bbox": elem["bbox"],
+                    "text": result.strip(),
+                    "reading_order": elem["reading_order"],
+                }
+            )
+
+    # Sort elements by reading order
+    recognition_results.sort(key=lambda x: x.get("reading_order", 0))
+
+    return recognition_results
+
+
+def main():
+    parser = argparse.ArgumentParser(description="Document parsing based on DOLPHIN")
+    parser.add_argument("--config", default="./config/Dolphin.yaml", help="Path to configuration file")
+    parser.add_argument("--input_path", type=str, default="./demo", help="Path to input image/PDF or directory of files")
+    parser.add_argument(
+        "--save_dir",
+        type=str,
+        default=None,
+        help="Directory to save parsing results (default: same as input directory)",
+    )
+    parser.add_argument(
+        "--max_batch_size",
+        type=int,
+        default=4,
+        help="Maximum number of document elements to parse in a single batch (default: 4)",
+    )
+    args = parser.parse_args()
+
+    # Load Model
+    config = OmegaConf.load(args.config)
+    model = DOLPHIN(config)
+
+    # Collect Document Files (images and PDFs)
+    if os.path.isdir(args.input_path):
+        # Support both image and PDF files
+        file_extensions = [".jpg", ".jpeg", ".png", ".JPG", ".JPEG", ".PNG", ".pdf", ".PDF"]
+        
+        document_files = []
+        for ext in file_extensions:
+            document_files.extend(glob.glob(os.path.join(args.input_path, f"*{ext}")))
+        document_files = sorted(document_files)
+    else:
+        if not os.path.exists(args.input_path):
+            raise FileNotFoundError(f"Input path {args.input_path} does not exist")
+        
+        # Check if it's a supported file type
+        file_ext = os.path.splitext(args.input_path)[1].lower()
+        supported_exts = ['.jpg', '.jpeg', '.png', '.pdf']
+        
+        if file_ext not in supported_exts:
+            raise ValueError(f"Unsupported file type: {file_ext}. Supported types: {supported_exts}")
+        
+        document_files = [args.input_path]
+
+    save_dir = args.save_dir or (
+        args.input_path if os.path.isdir(args.input_path) else os.path.dirname(args.input_path)
+    )
+    setup_output_dirs(save_dir)
+
+    total_samples = len(document_files)
+    print(f"\nTotal files to process: {total_samples}")
+
+    # Process All Document Files
+    for file_path in document_files:
+        print(f"\nProcessing {file_path}")
+        try:
+            json_path, recognition_results = process_document(
+                document_path=file_path,
+                model=model,
+                save_dir=save_dir,
+                max_batch_size=args.max_batch_size,
+            )
+
+            print(f"Processing completed. Results saved to {save_dir}")
+
+        except Exception as e:
+            print(f"Error processing {file_path}: {str(e)}")
+            continue
+
+
+if __name__ == "__main__":
+    main()

demo_page_hf.py

@@ -0,0 +1,365 @@
+""" 
+Copyright (c) 2025 Bytedance Ltd. and/or its affiliates
+SPDX-License-Identifier: MIT
+"""
+
+import argparse
+import glob
+import os
+
+import cv2
+import torch
+from PIL import Image
+from transformers import AutoProcessor, VisionEncoderDecoderModel
+
+from utils.utils import *
+
+
+class DOLPHIN:
+    def __init__(self, model_id_or_path):
+        """Initialize the Hugging Face model
+        
+        Args:
+            model_id_or_path: Path to local model or Hugging Face model ID
+        """
+        # Load model from local path or Hugging Face hub
+        self.processor = AutoProcessor.from_pretrained(model_id_or_path)
+        self.model = VisionEncoderDecoderModel.from_pretrained(model_id_or_path)
+        self.model.eval()
+        
+        # Set device and precision
+        self.device = "cuda" if torch.cuda.is_available() else "cpu"
+        self.model.to(self.device)
+        self.model = self.model.half()  # Always use half precision by default
+        
+        # set tokenizer
+        self.tokenizer = self.processor.tokenizer
+        
+    def chat(self, prompt, image):
+        """Process an image or batch of images with the given prompt(s)
+        
+        Args:
+            prompt: Text prompt or list of prompts to guide the model
+            image: PIL Image or list of PIL Images to process
+            
+        Returns:
+            Generated text or list of texts from the model
+        """
+        # Check if we're dealing with a batch
+        is_batch = isinstance(image, list)
+        
+        if not is_batch:
+            # Single image, wrap it in a list for consistent processing
+            images = [image]
+            prompts = [prompt]
+        else:
+            # Batch of images
+            images = image
+            prompts = prompt if isinstance(prompt, list) else [prompt] * len(images)
+        
+        # Prepare image
+        batch_inputs = self.processor(images, return_tensors="pt", padding=True)
+        batch_pixel_values = batch_inputs.pixel_values.half().to(self.device)
+        
+        # Prepare prompt
+        prompts = [f"<s>{p} <Answer/>" for p in prompts]
+        batch_prompt_inputs = self.tokenizer(
+            prompts,
+            add_special_tokens=False,
+            return_tensors="pt"
+        )
+
+        batch_prompt_ids = batch_prompt_inputs.input_ids.to(self.device)
+        batch_attention_mask = batch_prompt_inputs.attention_mask.to(self.device)
+        
+        # Generate text
+        outputs = self.model.generate(
+            pixel_values=batch_pixel_values,
+            decoder_input_ids=batch_prompt_ids,
+            decoder_attention_mask=batch_attention_mask,
+            min_length=1,
+            max_length=4096,
+            pad_token_id=self.tokenizer.pad_token_id,
+            eos_token_id=self.tokenizer.eos_token_id,
+            use_cache=True,
+            bad_words_ids=[[self.tokenizer.unk_token_id]],
+            return_dict_in_generate=True,
+            do_sample=False,
+            num_beams=1,
+            repetition_penalty=1.1,
+            temperature=1.0
+        )
+        
+        # Process output
+        sequences = self.tokenizer.batch_decode(outputs.sequences, skip_special_tokens=False)
+        
+        # Clean prompt text from output
+        results = []
+        for i, sequence in enumerate(sequences):
+            cleaned = sequence.replace(prompts[i], "").replace("<pad>", "").replace("</s>", "").strip()
+            results.append(cleaned)
+            
+        # Return a single result for single image input
+        if not is_batch:
+            return results[0]
+        return results
+
+
+def process_document(document_path, model, save_dir, max_batch_size=None):
+    """Parse documents with two stages - Handles both images and PDFs"""
+    file_ext = os.path.splitext(document_path)[1].lower()
+    
+    if file_ext == '.pdf':
+        # Process PDF file
+        # Convert PDF to images
+        images = convert_pdf_to_images(document_path)
+        if not images:
+            raise Exception(f"Failed to convert PDF {document_path} to images")
+        
+        all_results = []
+        
+        # Process each page
+        for page_idx, pil_image in enumerate(images):
+            print(f"Processing page {page_idx + 1}/{len(images)}")
+            
+            # Generate output name for this page
+            base_name = os.path.splitext(os.path.basename(document_path))[0]
+            page_name = f"{base_name}_page_{page_idx + 1:03d}"
+            
+            # Process this page (don't save individual page results)
+            json_path, recognition_results = process_single_image(
+                pil_image, model, save_dir, page_name, max_batch_size, save_individual=False
+            )
+            
+            # Add page information to results
+            page_results = {
+                "page_number": page_idx + 1,
+                "elements": recognition_results
+            }
+            all_results.append(page_results)
+        
+        # Save combined results for multi-page PDF
+        combined_json_path = save_combined_pdf_results(all_results, document_path, save_dir)
+        
+        return combined_json_path, all_results
+    
+    else:
+        # Process regular image file
+        pil_image = Image.open(document_path).convert("RGB")
+        base_name = os.path.splitext(os.path.basename(document_path))[0]
+        return process_single_image(pil_image, model, save_dir, base_name, max_batch_size)
+
+
+def process_single_image(image, model, save_dir, image_name, max_batch_size=None, save_individual=True):
+    """Process a single image (either from file or converted from PDF page)
+    
+    Args:
+        image: PIL Image object
+        model: DOLPHIN model instance
+        save_dir: Directory to save results
+        image_name: Name for the output file
+        max_batch_size: Maximum batch size for processing
+        save_individual: Whether to save individual results (False for PDF pages)
+        
+    Returns:
+        Tuple of (json_path, recognition_results)
+    """
+    # Stage 1: Page-level layout and reading order parsing
+    layout_output = model.chat("Parse the reading order of this document.", image)
+
+    # Stage 2: Element-level content parsing
+    padded_image, dims = prepare_image(image)
+    recognition_results = process_elements(layout_output, padded_image, dims, model, max_batch_size, save_dir, image_name)
+
+    # Save outputs only if requested (skip for PDF pages)
+    json_path = None
+    if save_individual:
+        # Create a dummy image path for save_outputs function
+        dummy_image_path = f"{image_name}.jpg"  # Extension doesn't matter, only basename is used
+        json_path = save_outputs(recognition_results, dummy_image_path, save_dir)
+
+    return json_path, recognition_results
+
+
+def process_elements(layout_results, padded_image, dims, model, max_batch_size, save_dir=None, image_name=None):
+    """Parse all document elements with parallel decoding"""
+    layout_results = parse_layout_string(layout_results)
+
+    # Store text and table elements separately
+    text_elements = []  # Text elements
+    table_elements = []  # Table elements
+    figure_results = []  # Image elements (no processing needed)
+    previous_box = None
+    reading_order = 0
+
+    # Collect elements to process and group by type
+    for bbox, label in layout_results:
+        try:
+            # Adjust coordinates
+            x1, y1, x2, y2, orig_x1, orig_y1, orig_x2, orig_y2, previous_box = process_coordinates(
+                bbox, padded_image, dims, previous_box
+            )
+
+            # Crop and parse element
+            cropped = padded_image[y1:y2, x1:x2]
+            if cropped.size > 0 and cropped.shape[0] > 3 and cropped.shape[1] > 3:
+                if label == "fig":
+                    pil_crop = Image.fromarray(cv2.cvtColor(cropped, cv2.COLOR_BGR2RGB))
+                    
+                    figure_filename = save_figure_to_local(pil_crop, save_dir, image_name, reading_order)
+                    
+                    # For figure regions, store relative path instead of base64
+                    figure_results.append(
+                        {
+                            "label": label,
+                            "text": f"![Figure](figures/{figure_filename})",
+                            "figure_path": f"figures/{figure_filename}",
+                            "bbox": [orig_x1, orig_y1, orig_x2, orig_y2],
+                            "reading_order": reading_order,
+                        }
+                    )
+                else:
+                    # Prepare element for parsing
+                    pil_crop = Image.fromarray(cv2.cvtColor(cropped, cv2.COLOR_BGR2RGB))
+                    element_info = {
+                        "crop": pil_crop,
+                        "label": label,
+                        "bbox": [orig_x1, orig_y1, orig_x2, orig_y2],
+                        "reading_order": reading_order,
+                    }
+                    
+                    # Group by type
+                    if label == "tab":
+                        table_elements.append(element_info)
+                    else:  # Text elements
+                        text_elements.append(element_info)
+
+            reading_order += 1
+
+        except Exception as e:
+            print(f"Error processing bbox with label {label}: {str(e)}")
+            continue
+
+    # Initialize results list
+    recognition_results = figure_results.copy()
+    
+    # Process text elements (in batches)
+    if text_elements:
+        text_results = process_element_batch(text_elements, model, "Read text in the image.", max_batch_size)
+        recognition_results.extend(text_results)
+    
+    # Process table elements (in batches)
+    if table_elements:
+        table_results = process_element_batch(table_elements, model, "Parse the table in the image.", max_batch_size)
+        recognition_results.extend(table_results)
+
+    # Sort elements by reading order
+    recognition_results.sort(key=lambda x: x.get("reading_order", 0))
+
+    return recognition_results
+
+
+def process_element_batch(elements, model, prompt, max_batch_size=None):
+    """Process elements of the same type in batches"""
+    results = []
+    
+    # Determine batch size
+    batch_size = len(elements)
+    if max_batch_size is not None and max_batch_size > 0:
+        batch_size = min(batch_size, max_batch_size)
+    
+    # Process in batches
+    for i in range(0, len(elements), batch_size):
+        batch_elements = elements[i:i+batch_size]
+        crops_list = [elem["crop"] for elem in batch_elements]
+        
+        # Use the same prompt for all elements in the batch
+        prompts_list = [prompt] * len(crops_list)
+        
+        # Batch inference
+        batch_results = model.chat(prompts_list, crops_list)
+        
+        # Add results
+        for j, result in enumerate(batch_results):
+            elem = batch_elements[j]
+            results.append({
+                "label": elem["label"],
+                "bbox": elem["bbox"],
+                "text": result.strip(),
+                "reading_order": elem["reading_order"],
+            })
+    
+    return results
+
+
+def main():
+    parser = argparse.ArgumentParser(description="Document parsing based on DOLPHIN")
+    parser.add_argument("--model_path", default="./hf_model", help="Path to Hugging Face model")
+    parser.add_argument("--input_path", type=str, default="./demo", help="Path to input image/PDF or directory of files")
+    parser.add_argument(
+        "--save_dir",
+        type=str,
+        default=None,
+        help="Directory to save parsing results (default: same as input directory)",
+    )
+    parser.add_argument(
+        "--max_batch_size",
+        type=int,
+        default=16,
+        help="Maximum number of document elements to parse in a single batch (default: 16)",
+    )
+    args = parser.parse_args()
+
+    # Load Model
+    model = DOLPHIN(args.model_path)
+
+    # Collect Document Files (images and PDFs)
+    if os.path.isdir(args.input_path):
+        # Support both image and PDF files
+        file_extensions = [".jpg", ".jpeg", ".png", ".JPG", ".JPEG", ".PNG", ".pdf", ".PDF"]
+        
+        document_files = []
+        for ext in file_extensions:
+            document_files.extend(glob.glob(os.path.join(args.input_path, f"*{ext}")))
+        document_files = sorted(document_files)
+    else:
+        if not os.path.exists(args.input_path):
+            raise FileNotFoundError(f"Input path {args.input_path} does not exist")
+        
+        # Check if it's a supported file type
+        file_ext = os.path.splitext(args.input_path)[1].lower()
+        supported_exts = ['.jpg', '.jpeg', '.png', '.pdf']
+        
+        if file_ext not in supported_exts:
+            raise ValueError(f"Unsupported file type: {file_ext}. Supported types: {supported_exts}")
+        
+        document_files = [args.input_path]
+
+    save_dir = args.save_dir or (
+        args.input_path if os.path.isdir(args.input_path) else os.path.dirname(args.input_path)
+    )
+    setup_output_dirs(save_dir)
+
+    total_samples = len(document_files)
+    print(f"\nTotal files to process: {total_samples}")
+
+    # Process All Document Files
+    for file_path in document_files:
+        print(f"\nProcessing {file_path}")
+        try:
+            json_path, recognition_results = process_document(
+                document_path=file_path,
+                model=model,
+                save_dir=save_dir,
+                max_batch_size=args.max_batch_size,
+            )
+
+            print(f"Processing completed. Results saved to {save_dir}")
+
+        except Exception as e:
+            print(f"Error processing {file_path}: {str(e)}")
+            continue
+
+
+if __name__ == "__main__":
+    main()

deployment/ReadMe.md

@@ -0,0 +1,12 @@
+<h1 align="center">
+🚀 Dolphin Inference/Serving
+</h1>
+
+## vLLM
+> [Doc](./vllm/ReadMe.md)
+
+## TensorRT-LLM
+> [Doc](./tensorrt_llm/ReadMe.md)
+
+## Others
+

deployment/tensorrt_llm/ReadMe.md

@@ -0,0 +1,89 @@
+<h1 align="center">
+🚀 Dolphin TensorRT-LLM Demo
+</h1>
+
+## ✅ Introduction
+The Dolphin model employs a **Swin Encoder + MBart Decoder** architecture. In the HuggingFace Transformers [Config](https://huggingface.co/ByteDance/Dolphin/blob/main/config.json), 
+its architectures field is specified as "VisionEncoderDecoderModel". **Dolphin**, **[Nougat](https://huggingface.co/docs/transformers/model_doc/nougat)**, and **[Donut](https://huggingface.co/docs/transformers/model_doc/donut)** share the same model architecture. TensorRT-LLM has already supported the Nougat model. 
+Following Nougat's conversion script, we have successfully implemented Dolphin on TensorRT-LLM. 
+
+**Note:** [prompt_ids](./dolphin_runner.py#L120) MUST be of **int32** type, otherwise TensorRT-LLM will produce incorrect results.
+
+## 🛠️ Installation
+> We only test TensorRT-LLM 0.18.1 on Linux.
+
+https://nvidia.github.io/TensorRT-LLM/0.18.1/installation/linux.html
+
+
+## ⚡ Offline Inference
+```
+export MODEL_NAME="Dolphin"
+
+# predict elements reading order
+python run_dolphin.py \
+    --batch_size 1 \
+    --hf_model_dir tmp/hf_models/${MODEL_NAME} \
+    --visual_engine_dir tmp/trt_engines/${MODEL_NAME}/vision_encoder \
+    --llm_engine_dir tmp/trt_engines/${MODEL_NAME}/1-gpu/bfloat16 \
+    --max_new_tokens 4096 \
+    --repetition_penalty 1.0 \
+    --input_text "Parse the reading order of this document." \
+    --image_path "../../demo/page_imgs/page_1.jpeg"
+
+# recognize text/latex
+python run_dolphin.py \
+    --batch_size 1 \
+    --hf_model_dir tmp/hf_models/${MODEL_NAME} \
+    --visual_engine_dir tmp/trt_engines/${MODEL_NAME}/vision_encoder \
+    --llm_engine_dir tmp/trt_engines/${MODEL_NAME}/1-gpu/bfloat16 \
+    --max_new_tokens 4096 \
+    --repetition_penalty 1.0 \
+    --input_text "Read text in the image." \
+    --image_path "../../demo/element_imgs/block_formula.jpeg"
+
+
+python run_dolphin.py \
+    --batch_size 1 \
+    --hf_model_dir tmp/hf_models/${MODEL_NAME} \
+    --visual_engine_dir tmp/trt_engines/${MODEL_NAME}/vision_encoder \
+    --llm_engine_dir tmp/trt_engines/${MODEL_NAME}/1-gpu/bfloat16 \
+    --max_new_tokens 4096 \
+    --repetition_penalty 1.0 \
+    --input_text "Read text in the image." \
+    --image_path "../../demo/element_imgs/para_1.jpg"
+
+# recognize table
+python run_dolphin.py \
+    --batch_size 1 \
+    --hf_model_dir tmp/hf_models/${MODEL_NAME} \
+    --visual_engine_dir tmp/trt_engines/${MODEL_NAME}/vision_encoder \
+    --llm_engine_dir tmp/trt_engines/${MODEL_NAME}/1-gpu/bfloat16 \
+    --max_new_tokens 4096 \
+    --repetition_penalty 1.0 \
+    --input_text "Parse the table in the image." \
+    --image_path "../../demo/element_imgs/table_1.jpeg"
+```
+
+
+## ⚡ Online Inference
+```
+# 1. Start Api Server
+export MODEL_NAME="Dolphin"
+
+python api_server.py \
+    --hf_model_dir tmp/hf_models/${MODEL_NAME} \
+    --visual_engine_dir tmp/trt_engines/${MODEL_NAME}/vision_encoder \
+    --llm_engine_dir tmp/trt_engines/${MODEL_NAME}/1-gpu/bfloat16 \
+    --max_batch_size 16
+
+# 2. Predict
+# predict elements reading order
+python deployment/tensorrt_llm/api_client.py --image_path ./demo/page_imgs/page_1.jpeg --prompt "Parse the reading order of this document."
+
+# recognize text/latex
+python deployment/tensorrt_llm/api_client.py --image_path ./demo/element_imgs/block_formula.jpeg --prompt "Read text in the image."
+python deployment/tensorrt_llm/api_client.py --image_path ./demo/element_imgs/para_1.jpg --prompt "Read text in the image."
+
+# recognize table
+python deployment/tensorrt_llm/api_client.py --image_path ./demo/element_imgs/table_1.jpeg --prompt "Parse the table in the image."
+```

deployment/tensorrt_llm/api_client.py

@@ -0,0 +1,100 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Example Python client for `vllm.entrypoints.api_server`
+Start the demo server:
+    python -m vllm.entrypoints.api_server --model <model_name>
+
+NOTE: The API server is used only for demonstration and simple performance
+benchmarks. It is not intended for production use.
+For production use, we recommend `vllm serve` and the OpenAI client API.
+"""
+
+import argparse
+import base64
+import json
+from argparse import Namespace
+from collections.abc import Iterable
+
+import requests
+
+
+def clear_line(n: int = 1) -> None:
+    LINE_UP = "\033[1A"
+    LINE_CLEAR = "\x1b[2K"
+    for _ in range(n):
+        print(LINE_UP, end=LINE_CLEAR, flush=True)
+
+
+def encode_image_base64(image_path: str) -> str:
+    """Encode local image to base64 format."""
+
+    with open(image_path, "rb") as f:
+        image_data = f.read()
+        result = base64.b64encode(image_data).decode("utf-8")
+
+    return result
+
+
+def post_http_request(
+        prompt: str, image_path: str, api_url: str, stream: bool = False
+) -> requests.Response:
+    headers = {"User-Agent": "Test Client"}
+    pload = {
+        "prompt": prompt,
+        "image_base64": encode_image_base64(image_path),
+    }
+    response = requests.post(api_url, headers=headers, json=pload, stream=stream)
+    return response
+
+
+def get_streaming_response(response: requests.Response) -> Iterable[list[str]]:
+    for chunk in response.iter_lines(
+            chunk_size=8192, decode_unicode=False, delimiter=b"\n"
+    ):
+        if chunk:
+            data = json.loads(chunk.decode("utf-8"))
+            output = data["text"]
+            yield output
+
+
+def get_response(response: requests.Response) -> list[str]:
+    data = json.loads(response.content)
+    output = data["text"]
+    return output
+
+
+def parse_args():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--host", type=str, default="localhost")
+    parser.add_argument("--port", type=int, default=8000)
+    parser.add_argument("--prompt", type=str, default="Parse the reading order of this document.")
+    parser.add_argument("--image_path", type=str, default="./demo/page_imgs/page_1.jpeg")
+    parser.add_argument("--stream", action="store_true")
+    return parser.parse_args()
+
+
+def main(args: Namespace):
+    prompt = args.prompt
+    image_path = args.image_path
+    api_url = f"http://{args.host}:{args.port}/generate"
+    stream = args.stream
+
+    print(f"Prompt: {prompt!r}\n", flush=True)
+    response = post_http_request(prompt, image_path, api_url, stream)
+
+    if stream:
+        num_printed_lines = 0
+        for h in get_streaming_response(response):
+            clear_line(num_printed_lines)
+            num_printed_lines = 0
+            for i, line in enumerate(h):
+                num_printed_lines += 1
+                print(f"Response {i}: {line!r}", flush=True)
+    else:
+        output = get_response(response)
+        print(f"Response: {output!r}", flush=True)
+
+
+if __name__ == "__main__":
+    args = parse_args()
+    main(args)

deployment/tensorrt_llm/api_server.py

@@ -0,0 +1,112 @@
+# copied from: https://github.com/NVIDIA/TensorRT-LLM/blob/v0.18.1/examples/apps/fastapi_server.py
+
+#!/usr/bin/env python
+import asyncio
+import base64
+import io
+import logging
+import signal
+from http import HTTPStatus
+from PIL import Image
+from typing import Optional
+
+import click
+import uvicorn
+from fastapi import FastAPI, Request
+from fastapi.responses import JSONResponse, Response
+
+from tensorrt_llm.executor import CppExecutorError, RequestError
+from dolphin_runner import DolphinRunner, InferenceConfig
+
+TIMEOUT_KEEP_ALIVE = 5  # seconds.
+
+
+async def decode_image(image_base64: str) -> Image.Image:
+    image_data = base64.b64decode(image_base64)
+    image = Image.open(io.BytesIO(image_data))
+    return image
+
+
+class LlmServer:
+    def __init__(self, runner: DolphinRunner):
+        self.runner = runner
+        self.app = FastAPI()
+        self.register_routes()
+
+    def register_routes(self):
+        self.app.add_api_route("/health", self.health, methods=["GET"])
+        self.app.add_api_route("/generate", self.generate, methods=["POST"])
+
+    async def health(self) -> Response:
+        return Response(status_code=200)
+
+    async def generate(self, request: Request) -> Response:
+        """ Generate completion for the request.
+
+        The request should be a JSON object with the following fields:
+        - prompt: the prompt to use for the generation.
+        - image_base64: the image to use for the generation.
+        """
+        request_dict = await request.json()
+
+        prompt = request_dict.pop("prompt", "")
+        logging.info(f"request prompt: {prompt}")
+        image_base64 = request_dict.pop("image_base64", "")
+        image = await decode_image(image_base64)
+
+        try:
+            output_texts = self.runner.run([prompt], [image], 4024)
+            output_texts = [texts[0] for texts in output_texts]
+            return JSONResponse({"text": output_texts[0]})
+        except RequestError as e:
+            return JSONResponse(content=str(e),
+                                status_code=HTTPStatus.BAD_REQUEST)
+        except CppExecutorError:
+            # If internal executor error is raised, shutdown the server
+            signal.raise_signal(signal.SIGINT)
+
+    async def __call__(self, host, port):
+        config = uvicorn.Config(self.app,
+                                host=host,
+                                port=port,
+                                log_level="info",
+                                timeout_keep_alive=TIMEOUT_KEEP_ALIVE)
+        await uvicorn.Server(config).serve()
+
+
+@click.command()
+@click.option("--hf_model_dir", type=str, required=True)
+@click.option("--visual_engine_dir", type=str, required=True)
+@click.option("--llm_engine_dir", type=str, required=True)
+@click.option("--max_batch_size", type=int, default=16)
+@click.option("--max_new_tokens", type=int, default=4024)
+@click.option("--host", type=str, default=None)
+@click.option("--port", type=int, default=8000)
+def entrypoint(hf_model_dir: str,
+               visual_engine_dir: str,
+               llm_engine_dir: str,
+               max_batch_size: int,
+               max_new_tokens: int,
+               host: Optional[str] = None,
+               port: int = 8000):
+    host = host or "0.0.0.0"
+    port = port or 8000
+    logging.info(f"Starting server at {host}:{port}")
+
+    config = InferenceConfig(
+        max_new_tokens=max_new_tokens,
+        batch_size=max_batch_size,
+        log_level="info",
+        hf_model_dir=hf_model_dir,
+        visual_engine_dir=visual_engine_dir,
+        llm_engine_dir=llm_engine_dir,
+    )
+
+    dolphin_runner = DolphinRunner(config)
+    server = LlmServer(runner=dolphin_runner)
+
+    asyncio.run(server(host, port))
+
+
+if __name__ == "__main__":
+    entrypoint()

deployment/tensorrt_llm/convert/build_visual_engine.py

@@ -0,0 +1,14 @@
+# copied from: https://github.com/NVIDIA/TensorRT-LLM/blob/v0.18.2/examples/multimodal/build_visual_engine.py
+
+import argparse
+
+from tensorrt_llm.tools.multimodal_builder import (VisionEngineBuilder,
+                                                   add_multimodal_arguments)
+
+if __name__ == '__main__':
+    parser = argparse.ArgumentParser()
+    parser = add_multimodal_arguments(parser)
+    args = parser.parse_args()
+
+    builder = VisionEngineBuilder(args)
+    builder.build()