app.py

import os
import uuid
import time
import asyncio
from threading import Thread

import gradio as gr
import spaces
import torch
import numpy as np
from PIL import Image
import cv2

from transformers import (
    Qwen2VLForConditionalGeneration,
    Qwen2_5_VLForConditionalGeneration,
    AutoModelForImageTextToText,
    AutoProcessor,
    TextIteratorStreamer,
)
from transformers.image_utils import load_image

# Constants
MAX_MAX_NEW_TOKENS = 2048
DEFAULT_MAX_NEW_TOKENS = 1024
MAX_INPUT_TOKEN_LENGTH = int(os.getenv("MAX_INPUT_TOKEN_LENGTH", "4096"))
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")

# Load public OCR models

MODEL_ID_V = "nanonets/Nanonets-OCR-s"
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.bfloat16
).to(device).eval()

MODEL_ID_X = "prithivMLmods/Qwen2-VL-OCR-2B-Instruct"
processor_x = AutoProcessor.from_pretrained(MODEL_ID_X, trust_remote_code=True)
model_x = Qwen2VLForConditionalGeneration.from_pretrained(
    MODEL_ID_X, trust_remote_code=True, torch_dtype=torch.bfloat16
).to(device).eval()

MODEL_ID_M = "reducto/RolmOCR"
processor_m = AutoProcessor.from_pretrained(MODEL_ID_M, trust_remote_code=True)
model_m = Qwen2_5_VLForConditionalGeneration.from_pretrained(
    MODEL_ID_M, trust_remote_code=True, torch_dtype=torch.bfloat16
).to(device).eval()

MODEL_ID_W = "prithivMLmods/Lh41-1042-Magellanic-7B-0711"
processor_w = AutoProcessor.from_pretrained(MODEL_ID_W, trust_remote_code=True)
model_w = Qwen2_5_VLForConditionalGeneration.from_pretrained(
    MODEL_ID_W, trust_remote_code=True, torch_dtype=torch.bfloat16
).to(device).eval()

def downsample_video(video_path):
    vidcap = cv2.VideoCapture(video_path)
    total = int(vidcap.get(cv2.CAP_PROP_FRAME_COUNT))
    fps = vidcap.get(cv2.CAP_PROP_FPS)
    frames = []
    for i in np.linspace(0, total - 1, 10, dtype=int):
        vidcap.set(cv2.CAP_PROP_POS_FRAMES, i)
        ok, img = vidcap.read()
        if ok:
            img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
            frames.append((Image.fromarray(img), round(i / fps, 2)))
    vidcap.release()
    return frames

@spaces.GPU
def generate_image(model_name, text, image, max_new_tokens, temperature, top_p, top_k, repetition_penalty):
    mapping = {
        "Nanonets-OCR-s": (processor_v, model_v),
        "Qwen2-VL-OCR-2B": (processor_x, model_x),
        "RolmOCR-7B": (processor_m, model_m),
        "Lh41-1042-Magellanic-7B-0711": (processor_w, model_w),
    }
    if model_name not in mapping:
        yield "Invalid model selected.", "Invalid model."
        return

    processor, model = mapping[model_name]
    if image is None:
        yield "Please upload an image.", ""
        return

    msg = [{"role": "user", "content": [{"type": "image", "image": image}, {"type": "text", "text": text}]}]
    prompt = processor.apply_chat_template(msg, tokenize=False, add_generation_prompt=True)
    inputs = processor(text=[prompt], images=[image], return_tensors="pt", padding=True).to(device)
    streamer = TextIteratorStreamer(processor, skip_prompt=True, skip_special_tokens=True)
    thread = Thread(target=model.generate, kwargs={**inputs, "streamer": streamer, "max_new_tokens": max_new_tokens})
    thread.start()

    out = ""
    for token in streamer:
        out += token.replace("<|im_end|>", "")
        time.sleep(0.01)
        yield out, out

@spaces.GPU
def generate_video(model_name, text, video_path, max_new_tokens, temperature, top_p, top_k, repetition_penalty):
    mapping = {
        "Nanonets-OCR-s": (processor_v, model_v),
        "Qwen2-VL-OCR-2B": (processor_x, model_x),
        "RolmOCR-7B": (processor_m, model_m),
        "Lh41-1042-Magellanic-7B-0711": (processor_w, model_w),
    }
    if model_name not in mapping:
        yield "Invalid model selected.", "Invalid model."
        return

    processor, model = mapping[model_name]
    if video_path is None:
        yield "Please upload a video.", ""
        return

    frames = downsample_video(video_path)
    messages = [{"role": "system", "content": [{"type": "text", "text": "You are a helpful assistant."}]},
                {"role": "user", "content": [{"type": "text", "text": text}]}]
    for img, ts in frames:
        messages[1]["content"].append({"type": "text", "text": f"Frame {ts}:"})
        messages[1]["content"].append({"type": "image", "image": img})

    inputs = processor.apply_chat_template(messages, tokenize=True, add_generation_prompt=True,
                                           return_tensors="pt").to(device)
    streamer = TextIteratorStreamer(processor, skip_prompt=True, skip_special_tokens=True)
    thread = Thread(target=model.generate, kwargs={**inputs,
                                                  "streamer": streamer,
                                                  "max_new_tokens": max_new_tokens,
                                                  "temperature": temperature,
                                                  "top_p": top_p,
                                                  "top_k": top_k,
                                                  "repetition_penalty": repetition_penalty})
    thread.start()
    out = ""
    for token in streamer:
        out += token.replace("<|im_end|>", "")
        time.sleep(0.01)
        yield out, out

# Examples
image_examples = [
    ["Extract the content", "images/4.png"],
    ["Explain the scene", "images/3.jpg"],
    ["Perform OCR on the image", "images/1.jpg"],
]
video_examples = [
    ["Explain the Ad in Detail", "videos/1.mp4"],
]

css = """
.submit-btn { background-color: #2980b9 !important; color: white !important; }
.submit-btn:hover { background-color: #3498db !important; }
.canvas-output { border: 2px solid #4682B4; border-radius: 10px; padding: 20px; }
"""

with gr.Blocks(css=css, theme="bethecloud/storj_theme") as demo:
    gr.Markdown("# **Multimodal OCR**")
    with gr.Row():
        with gr.Column():
            with gr.Tabs():
                with gr.TabItem("Image Inference"):
                    img_q = gr.Textbox(label="Query Input", placeholder="Enter prompt")
                    img_up = gr.Image(type="pil", label="Upload Image")
                    img_btn = gr.Button("Submit", elem_classes="submit-btn")
                    gr.Examples(examples=image_examples, inputs=[img_q, img_up])
                with gr.TabItem("Video Inference"):
                    vid_q = gr.Textbox(label="Query Input")
                    vid_up = gr.Video(label="Upload Video")
                    vid_btn = gr.Button("Submit", elem_classes="submit-btn")
                    gr.Examples(examples=video_examples, inputs=[vid_q, vid_up])
        with gr.Column(elem_classes="canvas-output"):
            gr.Markdown("## Output")
            out_raw = gr.Textbox(interactive=False, lines=2, show_copy_button=True)
            with gr.Accordion("Formatted Output", open=False):
                out_md = gr.Markdown()

        model_choice = gr.Radio(
            choices=["Nanonets-OCR-s", "Qwen2-VL-OCR-2B", "RolmOCR-7B", "Lh41-1042-Magellanic-7B-0711"],
            label="Select Model",
            value="Nanonets-OCR-s"
        )

    img_btn.click(generate_image, inputs=[model_choice, img_q, img_up,
                                          gr.Slider(1, MAX_MAX_NEW_TOKENS, value=DEFAULT_MAX_NEW_TOKENS),
                                          gr.Slider(0.1,4.0,value=0.6),
                                          gr.Slider(0.05,1.0,value=0.9),
                                          gr.Slider(1,1000,value=50),
                                          gr.Slider(1.0,2.0,value=1.2)],
                  outputs=[out_raw, out_md])

    vid_btn.click(generate_video, inputs=[model_choice, vid_q, vid_up,
                                          gr.Slider(1, MAX_MAX_NEW_TOKENS, value=DEFAULT_MAX_NEW_TOKENS),
                                          gr.Slider(0.1,4.0,value=0.6),
                                          gr.Slider(0.05,1.0,value=0.9),
                                          gr.Slider(1,1000,value=50),
                                          gr.Slider(1.0,2.0,value=1.2)],
                  outputs=[out_raw, out_md])

if __name__ == "__main__":
    demo.queue(max_size=30).launch(share=True, mcp_server=True, ssr_mode=False, show_error=True)