import spaces import json import math import os import traceback from io import BytesIO from typing import Any, Dict, List, Optional, Tuple import re import time from threading import Thread from io import BytesIO import uuid import tempfile import gradio as gr import requests import torch from PIL import Image import fitz from transformers import ( Qwen2_5_VLForConditionalGeneration, Qwen2VLForConditionalGeneration, AutoModelForCausalLM, AutoModelForVision2Seq, AutoModelForImageTextToText, AutoModel, AutoProcessor, TextIteratorStreamer, AutoTokenizer, ) from transformers.image_utils import load_image from reportlab.lib.pagesizes import A4 from reportlab.lib.styles import getSampleStyleSheet from reportlab.platypus import SimpleDocTemplate, Image as RLImage, Paragraph, Spacer from reportlab.lib.units import inch # --- Constants and Model Setup --- MAX_INPUT_TOKEN_LENGTH = 4096 device = torch.device("cuda" if torch.cuda.is_available() else "cpu") print("CUDA_VISIBLE_DEVICES=", os.environ.get("CUDA_VISIBLE_DEVICES")) print("torch.__version__ =", torch.__version__) print("torch.version.cuda =", torch.version.cuda) print("cuda available:", torch.cuda.is_available()) print("cuda device count:", torch.cuda.device_count()) if torch.cuda.is_available(): print("current device:", torch.cuda.current_device()) print("device name:", torch.cuda.get_device_name(torch.cuda.current_device())) print("Using device:", device) # --- Model Loading --- MODEL_ID_M = "LiquidAI/LFM2-VL-450M" processor_m = AutoProcessor.from_pretrained(MODEL_ID_M, trust_remote_code=True) model_m = AutoModelForImageTextToText.from_pretrained( MODEL_ID_M, trust_remote_code=True, torch_dtype=torch.float16 ).to(device).eval() MODEL_ID_T = "LiquidAI/LFM2-VL-1.6B" processor_t = AutoProcessor.from_pretrained(MODEL_ID_T, trust_remote_code=True) model_t = AutoModelForImageTextToText.from_pretrained( MODEL_ID_T, trust_remote_code=True, torch_dtype=torch.float16 ).to(device).eval() MODEL_ID_C = "HuggingFaceTB/SmolVLM-Instruct-250M" processor_c = AutoProcessor.from_pretrained(MODEL_ID_C, trust_remote_code=True) model_c = AutoModelForVision2Seq.from_pretrained( MODEL_ID_C, trust_remote_code=True, torch_dtype=torch.float16, _attn_implementation="flash_attention_2" ).to(device).eval() MODEL_ID_G = "echo840/MonkeyOCR-pro-1.2B" SUBFOLDER = "Recognition" processor_g = AutoProcessor.from_pretrained( MODEL_ID_G, trust_remote_code=True, subfolder=SUBFOLDER ) model_g = Qwen2_5_VLForConditionalGeneration.from_pretrained( MODEL_ID_G, trust_remote_code=True, subfolder=SUBFOLDER, torch_dtype=torch.float16 ).to(device).eval() MODEL_ID_I = "UCSC-VLAA/VLAA-Thinker-Qwen2VL-2B" processor_i = AutoProcessor.from_pretrained(MODEL_ID_I, trust_remote_code=True) model_i = Qwen2VLForConditionalGeneration.from_pretrained( MODEL_ID_I, trust_remote_code=True, torch_dtype=torch.float16 ).to(device).eval() MODEL_ID_A = "nanonets/Nanonets-OCR-s" processor_a = AutoProcessor.from_pretrained(MODEL_ID_A, trust_remote_code=True) model_a = Qwen2_5_VLForConditionalGeneration.from_pretrained( MODEL_ID_A, trust_remote_code=True, torch_dtype=torch.float16 ).to(device).eval() MODEL_ID_X = "prithivMLmods/Megalodon-OCR-Sync-0713" processor_x = AutoProcessor.from_pretrained(MODEL_ID_X, trust_remote_code=True) model_x = Qwen2_5_VLForConditionalGeneration.from_pretrained( MODEL_ID_X, trust_remote_code=True, torch_dtype=torch.float16 ).to(device).eval() MODEL_ID_Z = "Vchitect/ShotVL-3B" processor_z = AutoProcessor.from_pretrained(MODEL_ID_Z, trust_remote_code=True) model_z = Qwen2_5_VLForConditionalGeneration.from_pretrained( MODEL_ID_Z, trust_remote_code=True, torch_dtype=torch.float16 ).to(device).eval() # --- Moondream2 Model Loading --- MODEL_ID_MD = "vikhyatk/moondream2" REVISION_MD = "2025-06-21" moondream = AutoModelForCausalLM.from_pretrained( MODEL_ID_MD, revision=REVISION_MD, trust_remote_code=True, torch_dtype=torch.float16, device_map={"": "cuda"}, ) tokenizer_md = AutoTokenizer.from_pretrained(MODEL_ID_MD, revision=REVISION_MD) # --- Qwen2.5-VL-3B-Abliterated-Caption-it --- MODEL_ID_N = "prithivMLmods/Qwen2.5-VL-3B-Abliterated-Caption-it" processor_n = AutoProcessor.from_pretrained(MODEL_ID_N, trust_remote_code=True) model_n = Qwen2_5_VLForConditionalGeneration.from_pretrained( MODEL_ID_N, trust_remote_code=True, torch_dtype=torch.float16 ).to(device).eval() # --- LMM-R1-MGT-PerceReason --- MODEL_ID_F = "VLM-Reasoner/LMM-R1-MGT-PerceReason" processor_f = AutoProcessor.from_pretrained(MODEL_ID_F, trust_remote_code=True) model_f = Qwen2_5_VLForConditionalGeneration.from_pretrained( MODEL_ID_F, trust_remote_code=True, torch_dtype=torch.float16 ).to(device).eval() # TencentBAC/TBAC-VLR1-3B MODEL_ID_G = "TencentBAC/TBAC-VLR1-3B" processor_g = AutoProcessor.from_pretrained(MODEL_ID_G, trust_remote_code=True) model_g = Qwen2_5_VLForConditionalGeneration.from_pretrained( MODEL_ID_G, trust_remote_code=True, torch_dtype=torch.float16 ).to(device).eval() # OCRFlux-3B MODEL_ID_V = "ChatDOC/OCRFlux-3B" processor_v = AutoProcessor.from_pretrained(MODEL_ID_V, trust_remote_code=True) model_v = Qwen2_5_VLForConditionalGeneration.from_pretrained( MODEL_ID_V, trust_remote_code=True, torch_dtype=torch.float16 ).to(device).eval() # --- PDF Generation and Preview Utility Function --- def generate_and_preview_pdf(image: Image.Image, text_content: str, font_size: int, line_spacing: float, alignment: str, image_size: str): """ Generates a PDF, saves it, and then creates image previews of its pages. Returns the path to the PDF and a list of paths to the preview images. """ if image is None or not text_content or not text_content.strip(): raise gr.Error("Cannot generate PDF. Image or text content is missing.") # --- 1. Generate the PDF --- temp_dir = tempfile.gettempdir() pdf_filename = os.path.join(temp_dir, f"output_{uuid.uuid4()}.pdf") doc = SimpleDocTemplate( pdf_filename, pagesize=A4, rightMargin=inch, leftMargin=inch, topMargin=inch, bottomMargin=inch ) styles = getSampleStyleSheet() style_normal = styles["Normal"] style_normal.fontSize = int(font_size) style_normal.leading = int(font_size) * line_spacing style_normal.alignment = {"Left": 0, "Center": 1, "Right": 2, "Justified": 4}[alignment] story = [] img_buffer = BytesIO() image.save(img_buffer, format='PNG') img_buffer.seek(0) page_width, _ = A4 available_width = page_width - 2 * inch image_widths = { "Small": available_width * 0.3, "Medium": available_width * 0.6, "Large": available_width * 0.9, } img_width = image_widths[image_size] img = RLImage(img_buffer, width=img_width, height=image.height * (img_width / image.width)) story.append(img) story.append(Spacer(1, 12)) cleaned_text = re.sub(r'#+\s*', '', text_content).replace("*", "") text_paragraphs = cleaned_text.split('\n') for para in text_paragraphs: if para.strip(): story.append(Paragraph(para, style_normal)) doc.build(story) # --- 2. Render PDF pages as images for preview --- preview_images = [] try: pdf_doc = fitz.open(pdf_filename) for page_num in range(len(pdf_doc)): page = pdf_doc.load_page(page_num) pix = page.get_pixmap(dpi=150) preview_img_path = os.path.join(temp_dir, f"preview_{uuid.uuid4()}_p{page_num}.png") pix.save(preview_img_path) preview_images.append(preview_img_path) pdf_doc.close() except Exception as e: print(f"Error generating PDF preview: {e}") return pdf_filename, preview_images # --- Core Application Logic --- @spaces.GPU def process_document_stream( model_name: str, image: Image.Image, prompt_input: str, max_new_tokens: int, temperature: float, top_p: float, top_k: int, repetition_penalty: float ): """ Main generator function that handles model inference tasks with advanced generation parameters. """ if image is None: yield "Please upload an image.", "" return if not prompt_input or not prompt_input.strip(): yield "Please enter a prompt.", "" return if model_name == "Moondream2(vision)": image_embeds = moondream.encode_image(image) answer = moondream.answer_question( image_embeds=image_embeds, question=prompt_input, tokenizer=tokenizer_md ) yield answer, answer return if model_name == "LFM2-VL-450M(fast)": processor, model = processor_m, model_m elif model_name == "LFM2-VL-1.6B(fast)": processor, model = processor_t, model_t elif model_name == "ShotVL-3B(cinematic)": processor, model = processor_z, model_z elif model_name == "SmolVLM-Instruct-250M(smol)": processor, model = processor_c, model_c elif model_name == "MonkeyOCR-pro-1.2B(ocr)": processor, model = processor_g, model_g elif model_name == "VLAA-Thinker-Qwen2VL-2B(reason)": processor, model = processor_i, model_i elif model_name == "Nanonets-OCR-s(ocr)": processor, model = processor_a, model_a elif model_name == "Megalodon-OCR-Sync-0713(ocr)": processor, model = processor_x, model_x elif model_name == "Qwen2.5-VL-3B-Abliterated-Caption-it(caption)": processor, model = processor_n, model_n elif model_name == "LMM-R1-MGT-PerceReason(reason)": processor, model = processor_f, model_f elif model_name == "TBAC-VLR1-3B(open-r1)": processor, model = processor_g, model_g elif model_name == "OCRFlux-3B(ocr)": processor, model = processor_v, modelv else: yield "Invalid model selected.", "" return messages = [{"role": "user", "content": [{"type": "image", "image": image}, {"type": "text", "text": prompt_input}]}] prompt_full = processor.apply_chat_template(messages, tokenize=False, add_generation_prompt=True) inputs = processor(text=[prompt_full], images=[image], return_tensors="pt", padding=True, truncation=True, max_length=MAX_INPUT_TOKEN_LENGTH).to(device) streamer = TextIteratorStreamer(processor, skip_prompt=True, skip_special_tokens=True) generation_kwargs = { **inputs, "streamer": streamer, "max_new_tokens": max_new_tokens, "temperature": temperature, "top_p": top_p, "top_k": top_k, "repetition_penalty": repetition_penalty, "do_sample": True } thread = Thread(target=model.generate, kwargs=generation_kwargs) thread.start() buffer = "" for new_text in streamer: buffer += new_text buffer = buffer.replace("<|im_end|>", "") time.sleep(0.01) yield buffer , buffer yield buffer, buffer # --- Gradio UI Definition --- def create_gradio_interface(): """Builds and returns the Gradio web interface.""" css = """ .main-container { max-width: 1400px; margin: 0 auto; } .process-button { border: none !important; color: white !important; font-weight: bold !important; background-color: blue !important;} .process-button:hover { background-color: darkblue !important; transform: translateY(-2px) !important; box-shadow: 0 4px 8px rgba(0,0,0,0.2) !important; } #gallery { min-height: 400px; } """ with gr.Blocks(theme="bethecloud/storj_theme", css=css) as demo: gr.HTML("""

Tiny VLMs LabπŸ§ͺ

Advanced Vision-Language Model for Image Content and Layout Extraction

""") with gr.Row(): # Left Column (Inputs) with gr.Column(scale=1): model_choice = gr.Dropdown( choices=["LFM2-VL-450M(fast)", "LFM2-VL-1.6B(fast)", "SmolVLM-Instruct-250M(smol)", "Moondream2(vision)", "ShotVL-3B(cinematic)", "Megalodon-OCR-Sync-0713(ocr)", "VLAA-Thinker-Qwen2VL-2B(reason)", "MonkeyOCR-pro-1.2B(ocr)", "Qwen2.5-VL-3B-Abliterated-Caption-it(caption)", "Nanonets-OCR-s(ocr)", "LMM-R1-MGT-PerceReason(reason)", "TBAC-VLR1-3B(open-r1)", "OCRFlux-3B(ocr)"], label="Select Model", value= "LFM2-VL-450M(fast)" ) prompt_input = gr.Textbox(label="Query Input", placeholder="✦︎ Enter your query", value="Describe the image!") image_input = gr.Image(label="Upload Image", type="pil", sources=['upload']) with gr.Accordion("Advanced Settings", open=False): max_new_tokens = gr.Slider(minimum=512, maximum=8192, value=4096, step=256, label="Max New Tokens") temperature = gr.Slider(label="Temperature", minimum=0.1, maximum=4.0, step=0.1, value=0.6) top_p = gr.Slider(label="Top-p (nucleus sampling)", minimum=0.05, maximum=1.0, step=0.05, value=0.9) top_k = gr.Slider(label="Top-k", minimum=1, maximum=1000, step=1, value=50) repetition_penalty = gr.Slider(label="Repetition penalty", minimum=1.0, maximum=2.0, step=0.05, value=1.2) gr.Markdown("### PDF Export Settings") font_size = gr.Dropdown(choices=["8", "10", "12", "14", "16", "18"], value="12", label="Font Size") line_spacing = gr.Dropdown(choices=[1.0, 1.15, 1.5, 2.0], value=1.15, label="Line Spacing") alignment = gr.Dropdown(choices=["Left", "Center", "Right", "Justified"], value="Justified", label="Text Alignment") image_size = gr.Dropdown(choices=["Small", "Medium", "Large"], value="Medium", label="Image Size in PDF") process_btn = gr.Button("πŸš€ Process Image", variant="primary", elem_classes=["process-button"], size="lg") clear_btn = gr.Button("πŸ—‘οΈ Clear All", variant="secondary") # Right Column (Outputs) with gr.Column(scale=2): with gr.Tabs() as tabs: with gr.Tab("πŸ“ Extracted Content"): raw_output_stream = gr.Textbox(label="Raw Model Output Stream", interactive=False, lines=15, show_copy_button=True) with gr.Row(): examples = gr.Examples( examples=["examples/1.png", "examples/2.png", "examples/3.png", "examples/4.png", "examples/5.png", "examples/6.png"], inputs=image_input, label="Examples" ) gr.Markdown("[Report-BugπŸ’»](https://huggingface.co/spaces/prithivMLmods/OCR-Comparator/discussions)") with gr.Tab("πŸ“° README.md"): with gr.Accordion("(Result.md)", open=True): markdown_output = gr.Markdown() with gr.Tab("πŸ“‹ PDF Preview"): generate_pdf_btn = gr.Button("πŸ“„ Generate PDF & Render", variant="primary") pdf_output_file = gr.File(label="Download Generated PDF", interactive=False) pdf_preview_gallery = gr.Gallery(label="PDF Page Preview", show_label=True, elem_id="gallery", columns=2, object_fit="contain", height="auto") # Event Handlers def clear_all_outputs(): return None, "", "Raw output will appear here.", "", None, None process_btn.click( fn=process_document_stream, inputs=[model_choice, image_input, prompt_input, max_new_tokens, temperature, top_p, top_k, repetition_penalty], outputs=[raw_output_stream, markdown_output] ) generate_pdf_btn.click( fn=generate_and_preview_pdf, inputs=[image_input, raw_output_stream, font_size, line_spacing, alignment, image_size], outputs=[pdf_output_file, pdf_preview_gallery] ) clear_btn.click( clear_all_outputs, outputs=[image_input, prompt_input, raw_output_stream, markdown_output, pdf_output_file, pdf_preview_gallery] ) return demo if __name__ == "__main__": demo = create_gradio_interface() demo.queue(max_size=50).launch(share=True, ssr_mode=False, show_error=True)