Spaces:
Sleeping
Sleeping
| import gc | |
| import importlib.util | |
| import inspect | |
| import re | |
| import tempfile | |
| import time | |
| from io import BytesIO | |
| from pathlib import Path | |
| from threading import Lock, Thread | |
| import gradio as gr | |
| import spaces | |
| import torch | |
| import transformers | |
| from PIL import Image, ImageOps | |
| try: | |
| import fitz | |
| except ImportError: | |
| fitz = None | |
| from transformers import ( | |
| AutoModelForImageTextToText, | |
| AutoProcessor, | |
| TextIteratorStreamer, | |
| ) | |
| MAX_MAX_NEW_TOKENS = 8192 | |
| DEFAULT_MAX_NEW_TOKENS = 4096 | |
| MODEL_PATH = "zai-org/GLM-OCR" | |
| MIN_TRANSFORMERS_VERSION = "5.0.0" | |
| UPGRADE_TRANSFORMERS_CMD = 'pip install -U "transformers>=5.0.0"' | |
| MODEL_CARD_TRANSFORMERS_CMD = "pip install git+https://github.com/huggingface/transformers.git" | |
| BASE_DIR = Path(__file__).resolve().parent | |
| DATA_DIR = BASE_DIR / "data" | |
| INPUT_MODES = ["Upload PDF", "PDF from data", "All PDFs in data"] | |
| DEFAULT_INPUT_MODE = "Upload PDF" | |
| device = torch.device("cuda" if torch.cuda.is_available() else "cpu") | |
| print("Using device:", device) | |
| processor = None | |
| model = None | |
| model_load_error = None | |
| model_init_lock = Lock() | |
| def parse_version_triplet(version_text): | |
| parts = [int(part) for part in re.findall(r"\d+", str(version_text))[:3]] | |
| while len(parts) < 3: | |
| parts.append(0) | |
| return tuple(parts) | |
| def build_glm_ocr_dependency_error(): | |
| transformers_version = getattr(transformers, "__version__", "unknown") | |
| try: | |
| import torchvision # noqa: F401 | |
| except ImportError: | |
| torch_version = getattr(torch, "__version__", "unknown") | |
| return ( | |
| "GLM-OCR requires torchvision, but it is not installed in the active environment. " | |
| f"Install a torchvision build compatible with torch {torch_version} " | |
| "(for example, `pip install torchvision`) and restart the app." | |
| ) | |
| except Exception as exc: | |
| return ( | |
| "GLM-OCR could not import torchvision from the active environment. " | |
| f"Resolve the torchvision installation issue ({exc}) and restart the app." | |
| ) | |
| has_glm_ocr_support = importlib.util.find_spec("transformers.models.glm_ocr") is not None | |
| if has_glm_ocr_support: | |
| return None | |
| version_triplet = parse_version_triplet(transformers_version) | |
| minimum_triplet = parse_version_triplet(MIN_TRANSFORMERS_VERSION) | |
| if version_triplet and version_triplet < minimum_triplet: | |
| return ( | |
| f"GLM-OCR support is not available in transformers {transformers_version}. " | |
| f"Upgrade to {MIN_TRANSFORMERS_VERSION}+ with `{UPGRADE_TRANSFORMERS_CMD}` or use the " | |
| f"model card recommendation `{MODEL_CARD_TRANSFORMERS_CMD}`, then restart the app." | |
| ) | |
| return ( | |
| f"GLM-OCR support is not available in the installed transformers build ({transformers_version}). " | |
| f"Upgrade transformers with `{UPGRADE_TRANSFORMERS_CMD}` or use the model card recommendation " | |
| f"`{MODEL_CARD_TRANSFORMERS_CMD}`, then restart the app." | |
| ) | |
| def load_glm_ocr_components(): | |
| global processor, model, model_load_error | |
| if model_load_error is not None: | |
| raise RuntimeError(model_load_error) | |
| if processor is not None and model is not None: | |
| return processor, model | |
| with model_init_lock: | |
| if model_load_error is not None: | |
| raise RuntimeError(model_load_error) | |
| if processor is not None and model is not None: | |
| return processor, model | |
| dependency_error = build_glm_ocr_dependency_error() | |
| if dependency_error is not None: | |
| model_load_error = dependency_error | |
| print(model_load_error) | |
| raise RuntimeError(model_load_error) | |
| try: | |
| processor = AutoProcessor.from_pretrained(MODEL_PATH, trust_remote_code=True) | |
| model_kwargs = { | |
| "pretrained_model_name_or_path": MODEL_PATH, | |
| "torch_dtype": torch.bfloat16 if torch.cuda.is_available() else torch.float32, | |
| "trust_remote_code": True, | |
| } | |
| if torch.cuda.is_available(): | |
| model_kwargs["device_map"] = "auto" | |
| model = AutoModelForImageTextToText.from_pretrained(**model_kwargs).eval() | |
| if not torch.cuda.is_available(): | |
| model = model.to(device) | |
| except Exception as exc: | |
| model_load_error = ( | |
| f"Failed to load {MODEL_PATH}: {exc}. " | |
| "If the error mentions an unrecognized processor or model type, upgrade transformers with " | |
| f"`{UPGRADE_TRANSFORMERS_CMD}` or follow the model card recommendation " | |
| f"`{MODEL_CARD_TRANSFORMERS_CMD}`. Restart the app after upgrading." | |
| ) | |
| print(model_load_error) | |
| raise RuntimeError(model_load_error) from exc | |
| return processor, model | |
| TASK_PROMPTS = { | |
| "Text": "Text Recognition:", | |
| "Formula": "Formula Recognition:", | |
| "Table": "Table Recognition:", | |
| } | |
| TASK_CHOICES = list(TASK_PROMPTS.keys()) | |
| def list_data_pdfs(): | |
| if not DATA_DIR.exists(): | |
| return [] | |
| return sorted( | |
| [path for path in DATA_DIR.iterdir() if path.is_file() and path.suffix.lower() == ".pdf"], | |
| key=lambda path: path.name.lower(), | |
| ) | |
| def build_data_folder_note(): | |
| pdf_paths = list_data_pdfs() | |
| folder_line = f"Data folder: `{DATA_DIR}`" | |
| if not pdf_paths: | |
| return folder_line + "\n\nNo PDF files were found." | |
| lines = "\n".join(f"- `{path.name}`" for path in pdf_paths) | |
| return folder_line + "\n\nAvailable PDFs:\n" + lines | |
| def refresh_pdf_dropdown(): | |
| pdf_names = [path.name for path in list_data_pdfs()] | |
| value = pdf_names[0] if pdf_names else None | |
| return gr.update(choices=pdf_names, value=value), build_data_folder_note() | |
| def resolve_data_pdf_path(pdf_name): | |
| if not pdf_name: | |
| raise ValueError("Please choose a PDF from the data folder.") | |
| for path in list_data_pdfs(): | |
| if path.name == pdf_name: | |
| return path | |
| raise ValueError(f"Could not find `{pdf_name}` in `{DATA_DIR}`.") | |
| def resolve_uploaded_pdf_path(uploaded_pdf): | |
| if not uploaded_pdf: | |
| raise ValueError("Please upload a PDF first.") | |
| pdf_path = Path(str(uploaded_pdf)) | |
| if not pdf_path.exists(): | |
| raise ValueError("The uploaded PDF could not be read. Please upload it again.") | |
| if pdf_path.suffix.lower() != ".pdf": | |
| raise ValueError("Please upload a PDF file.") | |
| return pdf_path | |
| def cache_uploaded_pdf(uploaded_pdf): | |
| if not uploaded_pdf: | |
| return None, "No PDF uploaded yet." | |
| pdf_path = resolve_uploaded_pdf_path(uploaded_pdf) | |
| return str(pdf_path), f"Selected upload: `{pdf_path.name}`" | |
| def normalize_image(image: Image.Image): | |
| if image.mode in ("RGBA", "LA", "P"): | |
| image = image.convert("RGB") | |
| return ImageOps.exif_transpose(image) | |
| def pdf_page_to_image(pdf_path, page_number): | |
| if fitz is None: | |
| raise RuntimeError("PDF support requires PyMuPDF (`fitz`) to be installed.") | |
| page_number = int(page_number) | |
| if page_number < 1: | |
| raise ValueError("Page number must be 1 or higher.") | |
| with fitz.open(pdf_path) as document: | |
| if page_number > document.page_count: | |
| raise ValueError( | |
| f"Page {page_number} is outside the page count for {Path(pdf_path).name} " | |
| f"({document.page_count} pages)." | |
| ) | |
| page = document.load_page(page_number - 1) | |
| pixmap = page.get_pixmap(matrix=fitz.Matrix(2, 2), alpha=False) | |
| image = Image.open(BytesIO(pixmap.tobytes("png"))).convert("RGB") | |
| return ImageOps.exif_transpose(image) | |
| def get_pdf_page_count(pdf_path): | |
| if fitz is None: | |
| raise RuntimeError("PDF support requires PyMuPDF (`fitz`) to be installed.") | |
| with fitz.open(pdf_path) as document: | |
| return document.page_count | |
| def parse_page_selection(page_selection, page_count): | |
| selection = str(page_selection or "1").strip().lower() | |
| if not selection: | |
| selection = "1" | |
| if selection == "all": | |
| return list(range(1, page_count + 1)) | |
| pages = [] | |
| for raw_part in selection.split(","): | |
| part = raw_part.strip() | |
| if not part: | |
| continue | |
| if "-" in part: | |
| start_str, end_str = [token.strip() for token in part.split("-", 1)] | |
| start_page = int(start_str) | |
| end_page = int(end_str) | |
| if start_page > end_page: | |
| raise ValueError(f"Invalid page range: {part}") | |
| pages.extend(range(start_page, end_page + 1)) | |
| else: | |
| pages.append(int(part)) | |
| if not pages: | |
| raise ValueError("Please provide page numbers such as `1`, `1,3`, `2-5`, or `all`.") | |
| deduped_pages = [] | |
| seen = set() | |
| for page in pages: | |
| if page < 1 or page > page_count: | |
| raise ValueError(f"Page {page} is outside the PDF page count ({page_count}).") | |
| if page not in seen: | |
| deduped_pages.append(page) | |
| seen.add(page) | |
| return deduped_pages | |
| def join_output_blocks(blocks): | |
| return "\n\n".join(block.strip() for block in blocks if str(block).strip()).strip() | |
| def stream_pdf_text(pdf_path, task, page_selection, max_new_tokens): | |
| page_count = get_pdf_page_count(pdf_path) | |
| selected_pages = parse_page_selection(page_selection, page_count) | |
| include_page_headers = len(selected_pages) > 1 | |
| combined_sections = [] | |
| for page_number in selected_pages: | |
| page_title = f"Page {page_number}" | |
| try: | |
| image = pdf_page_to_image(pdf_path, page_number) | |
| page_text = "" | |
| for chunk in process_image_stream( | |
| image=image, | |
| task=task, | |
| max_new_tokens=max_new_tokens, | |
| ): | |
| page_text = chunk.strip() | |
| current_block = f"{page_title}\n{page_text}" if include_page_headers else page_text | |
| yield join_output_blocks(combined_sections + [current_block]) | |
| if not page_text.strip(): | |
| page_text = "[ERROR] No output was generated." | |
| except Exception as exc: | |
| page_text = f"[ERROR] {str(exc)}" | |
| final_block = f"{page_title}\n{page_text}" if include_page_headers else page_text | |
| combined_sections.append(final_block) | |
| yield join_output_blocks(combined_sections) | |
| def stream_data_folder_pdfs(task, page_selection, max_new_tokens): | |
| pdf_paths = list_data_pdfs() | |
| if not pdf_paths: | |
| yield f"[ERROR] No PDF files were found in `{DATA_DIR}`." | |
| return | |
| combined_files = [] | |
| for path in pdf_paths: | |
| pdf_text = "" | |
| try: | |
| for partial in stream_pdf_text(path, task, page_selection, max_new_tokens): | |
| pdf_text = partial | |
| yield join_output_blocks(combined_files + [f"File: {path.name}\n{pdf_text}"]) | |
| if not pdf_text.strip(): | |
| pdf_text = "[ERROR] No output was generated." | |
| file_block = f"File: {path.name}\n{pdf_text}" | |
| except Exception as exc: | |
| file_block = f"File: {path.name}\n[ERROR] {str(exc)}" | |
| yield join_output_blocks(combined_files + [file_block]) | |
| combined_files.append(file_block) | |
| yield join_output_blocks(combined_files) | |
| def calc_timeout_generic(*args, **kwargs): | |
| gpu_timeout = kwargs.get("gpu_timeout", None) | |
| if gpu_timeout is None and args: | |
| gpu_timeout = args[-1] | |
| try: | |
| return int(gpu_timeout) | |
| except Exception: | |
| return 60 | |
| def process_image_stream(image, task, max_new_tokens=DEFAULT_MAX_NEW_TOKENS, gpu_timeout=60): | |
| del gpu_timeout | |
| tmp_path = None | |
| try: | |
| if image is None: | |
| yield "[ERROR] Please upload an image first." | |
| return | |
| if task not in TASK_PROMPTS: | |
| yield "[ERROR] Invalid OCR task selected." | |
| return | |
| processor_obj, model_obj = load_glm_ocr_components() | |
| image = normalize_image(image) | |
| tmp = tempfile.NamedTemporaryFile(delete=False, suffix=".png") | |
| image.save(tmp.name, "PNG") | |
| tmp_path = tmp.name | |
| tmp.close() | |
| prompt = TASK_PROMPTS[task] | |
| messages = [ | |
| { | |
| "role": "user", | |
| "content": [ | |
| {"type": "image", "url": tmp_path}, | |
| {"type": "text", "text": prompt}, | |
| ], | |
| } | |
| ] | |
| inputs = processor_obj.apply_chat_template( | |
| messages, | |
| tokenize=True, | |
| add_generation_prompt=True, | |
| return_dict=True, | |
| return_tensors="pt", | |
| ) | |
| inputs.pop("token_type_ids", None) | |
| inputs = {key: value.to(model_obj.device) if hasattr(value, "to") else value for key, value in inputs.items()} | |
| streamer = TextIteratorStreamer( | |
| processor_obj.tokenizer if hasattr(processor_obj, "tokenizer") else processor_obj, | |
| skip_prompt=True, | |
| skip_special_tokens=True, | |
| ) | |
| generation_error = {"error": None} | |
| generation_kwargs = { | |
| **inputs, | |
| "streamer": streamer, | |
| "max_new_tokens": int(max_new_tokens), | |
| } | |
| def run_generation(): | |
| try: | |
| model_obj.generate(**generation_kwargs) | |
| except Exception as exc: | |
| generation_error["error"] = exc | |
| try: | |
| streamer.end() | |
| except Exception: | |
| pass | |
| thread = Thread(target=run_generation, daemon=True) | |
| thread.start() | |
| buffer = "" | |
| for new_text in streamer: | |
| buffer += new_text | |
| time.sleep(0.01) | |
| yield buffer.strip() | |
| thread.join(timeout=1.0) | |
| if generation_error["error"] is not None: | |
| error_message = f"[ERROR] Inference failed: {generation_error['error']}" | |
| if buffer.strip(): | |
| yield buffer.strip() + "\n\n" + error_message | |
| else: | |
| yield error_message | |
| return | |
| if not buffer.strip(): | |
| yield "[ERROR] No output was generated." | |
| except Exception as exc: | |
| yield f"[ERROR] {str(exc)}" | |
| finally: | |
| if tmp_path and Path(tmp_path).exists(): | |
| try: | |
| Path(tmp_path).unlink() | |
| except Exception: | |
| pass | |
| gc.collect() | |
| if torch.cuda.is_available(): | |
| torch.cuda.empty_cache() | |
| def toggle_input_mode(mode): | |
| mode = str(mode or DEFAULT_INPUT_MODE) | |
| if mode == "Upload PDF": | |
| help_text = "Click the upload button to choose a PDF, then enter pages like `1`, `1,3`, `2-5`, or `all`." | |
| return ( | |
| gr.update(visible=True), | |
| gr.update(visible=True), | |
| gr.update(visible=False), | |
| gr.update(visible=True), | |
| gr.update(visible=False), | |
| gr.update(visible=False), | |
| gr.update(value=help_text), | |
| ) | |
| if mode == "All PDFs in data": | |
| help_text = ( | |
| "Run OCR on every PDF in the data folder. Page selections such as `1`, `1,3`, `2-5`, " | |
| "or `all` are applied to each file." | |
| ) | |
| return ( | |
| gr.update(visible=False), | |
| gr.update(visible=False), | |
| gr.update(visible=False), | |
| gr.update(visible=True), | |
| gr.update(visible=True), | |
| gr.update(visible=True), | |
| gr.update(value=help_text), | |
| ) | |
| help_text = ( | |
| "Choose one PDF from the data folder and enter pages like `1`, `1,3`, `2-5`, or `all`." | |
| ) | |
| return ( | |
| gr.update(visible=False), | |
| gr.update(visible=False), | |
| gr.update(visible=True), | |
| gr.update(visible=True), | |
| gr.update(visible=True), | |
| gr.update(visible=True), | |
| gr.update(value=help_text), | |
| ) | |
| def run_router(input_mode, task, uploaded_pdf, pdf_name, page_selection, max_new_tokens_v, gpu_timeout_v): | |
| try: | |
| mode = str(input_mode or DEFAULT_INPUT_MODE) | |
| if mode == "All PDFs in data": | |
| yield from stream_data_folder_pdfs( | |
| task=task, | |
| page_selection=page_selection, | |
| max_new_tokens=max_new_tokens_v, | |
| ) | |
| return | |
| if mode == "Upload PDF": | |
| pdf_path = resolve_uploaded_pdf_path(uploaded_pdf) | |
| else: | |
| pdf_path = resolve_data_pdf_path(pdf_name) | |
| yield from stream_pdf_text( | |
| pdf_path=pdf_path, | |
| task=task, | |
| page_selection=page_selection, | |
| max_new_tokens=max_new_tokens_v, | |
| ) | |
| except Exception as exc: | |
| yield f"[ERROR] {str(exc)}" | |
| available_pdfs = [path.name for path in list_data_pdfs()] | |
| default_pdf = available_pdfs[0] if available_pdfs else None | |
| with gr.Blocks(title="GLM-OCR") as demo: | |
| gr.Markdown("# GLM-OCR") | |
| gr.Markdown("Upload a PDF with the button below or run OCR on PDFs stored in `glmocr/data`.") | |
| with gr.Row(): | |
| input_mode = gr.Radio( | |
| choices=INPUT_MODES, | |
| value=DEFAULT_INPUT_MODE, | |
| label="Input mode", | |
| ) | |
| task = gr.Radio( | |
| choices=TASK_CHOICES, | |
| value="Text", | |
| label="OCR task", | |
| ) | |
| input_help = gr.Markdown( | |
| "Click the upload button to choose a PDF, then enter pages like `1`, `1,3`, `2-5`, or `all`." | |
| ) | |
| with gr.Column(): | |
| uploaded_pdf_state = gr.State(None) | |
| upload_pdf_btn = gr.UploadButton( | |
| "Upload PDF", | |
| file_types=[".pdf"], | |
| file_count="single", | |
| type="filepath", | |
| visible=True, | |
| ) | |
| uploaded_pdf_status = gr.Markdown("No PDF uploaded yet.", visible=True) | |
| pdf_dropdown = gr.Dropdown( | |
| choices=available_pdfs, | |
| value=default_pdf, | |
| label="PDF from data folder", | |
| visible=False, | |
| ) | |
| page_selection = gr.Textbox( | |
| value="1", | |
| label="Pages", | |
| placeholder="Examples: 1, 1,3, 2-5, all", | |
| visible=True, | |
| ) | |
| data_folder_note = gr.Markdown(build_data_folder_note(), visible=False) | |
| refresh_pdfs_btn = gr.Button("Refresh PDF list", visible=False) | |
| with gr.Row(): | |
| max_new_tokens = gr.Slider( | |
| minimum=1, | |
| maximum=MAX_MAX_NEW_TOKENS, | |
| step=1, | |
| value=DEFAULT_MAX_NEW_TOKENS, | |
| label="Max new tokens", | |
| ) | |
| gpu_duration_state = gr.Number( | |
| value=60, | |
| precision=0, | |
| label="GPU duration (seconds)", | |
| ) | |
| run_btn = gr.Button("Run OCR", variant="primary") | |
| result = gr.Textbox(label="OCR output", lines=24, max_lines=40) | |
| demo.load( | |
| fn=refresh_pdf_dropdown, | |
| inputs=None, | |
| outputs=[pdf_dropdown, data_folder_note], | |
| queue=False, | |
| ) | |
| input_mode.change( | |
| fn=toggle_input_mode, | |
| inputs=[input_mode], | |
| outputs=[ | |
| upload_pdf_btn, | |
| uploaded_pdf_status, | |
| pdf_dropdown, | |
| page_selection, | |
| data_folder_note, | |
| refresh_pdfs_btn, | |
| input_help, | |
| ], | |
| queue=False, | |
| ) | |
| upload_pdf_btn.upload( | |
| fn=cache_uploaded_pdf, | |
| inputs=[upload_pdf_btn], | |
| outputs=[uploaded_pdf_state, uploaded_pdf_status], | |
| queue=False, | |
| ) | |
| refresh_pdfs_btn.click( | |
| fn=refresh_pdf_dropdown, | |
| inputs=None, | |
| outputs=[pdf_dropdown, data_folder_note], | |
| queue=False, | |
| ) | |
| run_btn.click( | |
| fn=run_router, | |
| inputs=[ | |
| input_mode, | |
| task, | |
| uploaded_pdf_state, | |
| pdf_dropdown, | |
| page_selection, | |
| max_new_tokens, | |
| gpu_duration_state, | |
| ], | |
| outputs=[result], | |
| ) | |
| if __name__ == "__main__": | |
| launch_signature = inspect.signature(demo.launch) | |
| launch_kwargs = { | |
| "show_error": True, | |
| } | |
| if "ssr_mode" in launch_signature.parameters: | |
| launch_kwargs["ssr_mode"] = False | |
| if "mcp_server" in launch_signature.parameters: | |
| launch_kwargs["mcp_server"] = True | |
| demo.queue(max_size=50).launch(**launch_kwargs) | |