app.py

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
import numpy as np
import torchvision.transforms as T
from torchvision.transforms.functional import InterpolationMode


from transformers import (
    Qwen2_5_VLForConditionalGeneration,
    Qwen2VLForConditionalGeneration,
    AutoModelForCausalLM,
    AutoModelForVision2Seq,
    AutoModelForImageTextToText,
    AutoModel,
    AutoProcessor,
    TextIteratorStreamer,
    AutoTokenizer,
    LlavaOnevisionForConditionalGeneration,
    LlavaOnevisionProcessor,
)

from transformers.image_utils import load_image as hf_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)

# --- InternVL3_5-2B-MPO Preprocessing Functions ---
IMAGENET_MEAN = (0.485, 0.456, 0.406)
IMAGENET_STD = (0.229, 0.224, 0.225)

def build_transform(input_size):
    MEAN, STD = IMAGENET_MEAN, IMAGENET_STD
    transform = T.Compose([
        T.Lambda(lambda img: img.convert('RGB') if img.mode != 'RGB' else img),
        T.Resize((input_size, input_size), interpolation=InterpolationMode.BICUBIC),
        T.ToTensor(),
        T.Normalize(mean=MEAN, std=STD)
    ])
    return transform

def find_closest_aspect_ratio(aspect_ratio, target_ratios, width, height, image_size):
    best_ratio_diff = float('inf')
    best_ratio = (1, 1)
    area = width * height
    for ratio in target_ratios:
        target_aspect_ratio = ratio[0] / ratio[1]
        ratio_diff = abs(aspect_ratio - target_aspect_ratio)
        if ratio_diff < best_ratio_diff:
            best_ratio_diff = ratio_diff
            best_ratio = ratio
        elif ratio_diff == best_ratio_diff:
            if area > 0.5 * image_size * image_size * ratio[0] * ratio[1]:
                best_ratio = ratio
    return best_ratio

def dynamic_preprocess(image, min_num=1, max_num=12, image_size=448, use_thumbnail=False):
    orig_width, orig_height = image.size
    aspect_ratio = orig_width / orig_height

    target_ratios = set(
        (i, j) for n in range(min_num, max_num + 1) for i in range(1, n + 1) for j in range(1, n + 1) if
        i * j <= max_num and i * j >= min_num)
    target_ratios = sorted(target_ratios, key=lambda x: x[0] * x[1])

    target_aspect_ratio = find_closest_aspect_ratio(
        aspect_ratio, target_ratios, orig_width, orig_height, image_size)

    target_width = image_size * target_aspect_ratio[0]
    target_height = image_size * target_aspect_ratio[1]
    blocks = target_aspect_ratio[0] * target_aspect_ratio[1]

    resized_img = image.resize((target_width, target_height))
    processed_images = []
    for i in range(blocks):
        box = (
            (i % (target_width // image_size)) * image_size,
            (i // (target_width // image_size)) * image_size,
            ((i % (target_width // image_size)) + 1) * image_size,
            ((i // (target_width // image_size)) + 1) * image_size
        )
        split_img = resized_img.crop(box)
        processed_images.append(split_img)
    assert len(processed_images) == blocks
    if use_thumbnail and len(processed_images) != 1:
        thumbnail_img = image.resize((image_size, image_size))
        processed_images.append(thumbnail_img)
    return processed_images

def load_image_internvl(image, input_size=448, max_num=12):
    transform = build_transform(input_size=input_size)
    images = dynamic_preprocess(image, image_size=input_size, use_thumbnail=True, max_num=max_num)
    pixel_values = [transform(img) for img in images]
    pixel_values = torch.stack(pixel_values)
    return pixel_values

# --- 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()

# --- 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()

MODEL_ID_O = "HuggingFaceTB/SmolVLM-500M-Instruct"
processor_o = AutoProcessor.from_pretrained(MODEL_ID_O, trust_remote_code=True)
model_o = AutoModelForVision2Seq.from_pretrained(
    MODEL_ID_O, trust_remote_code=True, torch_dtype=torch.float16, _attn_implementation="flash_attention_2"
).to(device).eval()

# --- New Model: llava-onevision ---
MODEL_ID_LO = "llava-hf/llava-onevision-qwen2-0.5b-ov-hf"
processor_lo = LlavaOnevisionProcessor.from_pretrained(MODEL_ID_LO)
model_lo = LlavaOnevisionForConditionalGeneration.from_pretrained(
    MODEL_ID_LO,
    torch_dtype=torch.float16
).to(device).eval()

# OpenGVLab/InternVL3_5-2B-MPO ---
MODEL_ID_IV = 'OpenGVLab/InternVL3_5-2B-MPO'
model_iv = AutoModel.from_pretrained(
    MODEL_ID_IV,
    torch_dtype=torch.bfloat16,
    trust_remote_code=True,
    device_map="auto").eval()
tokenizer_iv = AutoTokenizer.from_pretrained(MODEL_ID_IV, trust_remote_code=True, use_fast=False)


# --- 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

    # --- Special Handling for Moondream2 ---
    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
    
    # --- Special Handling for InternVL ---
    if model_name == "OpenGVLab/InternVL3_5-2B-MPO":
        pixel_values = load_image_internvl(image, max_num=12).to(torch.bfloat16).to(device)
        generation_config = dict(
            max_new_tokens=max_new_tokens,
            do_sample=True if temperature > 0 else False,
            temperature=temperature,
            top_p=top_p,
            top_k=top_k,
            repetition_penalty=repetition_penalty,
        )
        question = f"<image>\n{prompt_input}"
        response = model_iv.chat(tokenizer_iv, pixel_values, question, generation_config)
        yield response, response
        return


    processor = None
    model = None
    
    # --- Special Handling for Llava-OneVision ---
    if model_name == "llava-onevision-qwen2-0.5b-ov-hf(mini)":
        processor, model = processor_lo, model_lo
        prompt = f"<|im_start|>user <image>\n{prompt_input}<|im_end|><|im_start|>assistant"
        inputs = processor(text=prompt, images=image, return_tensors="pt").to(device, torch.float16)
    # --- Generic Handling for all other models ---
    else:
        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 == "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, model_v
        elif model_name == "SmolVLM-500M-Instruct(smol)": processor, model = processor_o, model_o
        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("""
        <div class="title" style="text-align: center">
            <h1>Tiny VLMs Lab🧪</h1>
            <p style="font-size: 1.1em; color: #6b7280; margin-bottom: 0.6em;">
                Tiny VLMs for Image Content Extraction and Understanding
            </p>
        </div>
        """)

        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)",
                             "OpenGVLab/InternVL3_5-2B-MPO", "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)", "OCRFlux-3B(ocr)", "TBAC-VLR1-3B(open-r1)", 
                             "SmolVLM-500M-Instruct(smol)", "llava-onevision-qwen2-0.5b-ov-hf(mini)"],
                    label="Select Model", value= "LFM2-VL-450M(fast)"
                )
                
                prompt_input = gr.Textbox(label="Query Input", placeholder="✦︎ Enter the prompt")
                image_input = gr.Image(label="Upload Image", type="pil", sources=['upload'])
                
                with gr.Accordion("Advanced Settings (PDF)", open=False):
                    max_new_tokens = gr.Slider(minimum=512, maximum=8192, value=2048, 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/Tiny-VLMs-Lab/discussions) | [prithivMLmods🤗](https://huggingface.co/prithivMLmods)")
                    
                    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(mcp_server=True, ssr_mode=False, show_error=True)