app.py

import os
import random
import uuid
import json
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 (
    Qwen2_5_VLForConditionalGeneration,
    AutoModelForCausalLM,
    AutoProcessor,
    TextIteratorStreamer,
)
from transformers.image_utils import load_image

# Constants for text generation
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 Camel-Doc-OCR-080125
MODEL_ID_M = "prithivMLmods/Camel-Doc-OCR-080125"
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.float16).to(device).eval()

# Load OCRFlux-3B
MODEL_ID_X = "ChatDOC/OCRFlux-3B"
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()

# Load Behemoth-3B-070225
MODEL_ID_T = "prithivMLmods/Behemoth-3B-070225-post0.1"
processor_t = AutoProcessor.from_pretrained(MODEL_ID_T, trust_remote_code=True)
model_t = Qwen2_5_VLForConditionalGeneration.from_pretrained(
    MODEL_ID_T, trust_remote_code=True,
    torch_dtype=torch.float16).to(device).eval()

# Load MonkeyOCR-pro-1.2B
MODEL_ID_O = "echo840/MonkeyOCR-pro-1.2B"
SUBFOLDER = "Recognition"
processor_o = AutoProcessor.from_pretrained(MODEL_ID_O, trust_remote_code=True, subfolder=SUBFOLDER)
model_o = Qwen2_5_VLForConditionalGeneration.from_pretrained(
    MODEL_ID_O, trust_remote_code=True, subfolder=SUBFOLDER, 
    torch_dtype=torch.float16).to(device).eval()

# Load ViGoRL-MCTS-SFT-7b-Spatial
MODEL_ID_A = "gsarch/ViGoRL-MCTS-SFT-7b-Spatial"
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()

# Function to downsample video frames
def downsample_video(video_path):
    """
    Downsamples the video to evenly spaced frames.
    Each frame is returned as a PIL image along with its timestamp.
    """
    vidcap = cv2.VideoCapture(video_path)
    total_frames = int(vidcap.get(cv2.CAP_PROP_FRAME_COUNT))
    fps = vidcap.get(cv2.CAP_PROP_FPS)
    frames = []
    frame_indices = np.linspace(0, total_frames - 1, 10, dtype=int)
    for i in frame_indices:
        vidcap.set(cv2.CAP_PROP_POS_FRAMES, i)
        success, image = vidcap.read()
        if success:
            image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
            pil_image = Image.fromarray(image)
            timestamp = round(i / fps, 2)
            frames.append((pil_image, timestamp))
    vidcap.release()
    return frames

# Function to generate text responses based on image input
@spaces.GPU
def generate_image(model_name: str,
                   text: str,
                   image: Image.Image,
                   max_new_tokens: int = 1024,
                   temperature: float = 0.6,
                   top_p: float = 0.9,
                   top_k: int = 50,
                   repetition_penalty: float = 1.2):
    """
    Generates responses using the selected model for image input.
    """
    if model_name == "Camel-Doc-OCR-080125(v2)":
        processor = processor_m
        model = model_m
    elif model_name == "OCRFlux-3B":
        processor = processor_x
        model = model_x
    elif model_name == "Behemoth-3B-070225":
        processor = processor_o
        model = model_o
    elif model_name == "MonkeyOCR-pro-1.2B":
        processor = processor_t
        model = model_t
    elif model_name == "ViGoRL-MCTS-SFT-7B":
        processor = processor_a
        model = model_a
    else:
        yield "Invalid model selected.", "Invalid model selected."
        return

    if image is None:
        yield "Please upload an image.", "Please upload an image."
        return

    messages = [{
        "role": "user",
        "content": [
            {"type": "image", "image": image},
            {"type": "text", "text": text},
        ]
    }]
    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=False,
                       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
    }
    thread = Thread(target=model.generate, kwargs=generation_kwargs)
    thread.start()
    buffer = ""
    for new_text in streamer:
        buffer += new_text
        time.sleep(0.01)
        yield buffer, buffer

# Function to generate text responses based on video input
@spaces.GPU
def generate_video(model_name: str,
                   text: str,
                   video_path: str,
                   max_new_tokens: int = 1024,
                   temperature: float = 0.6,
                   top_p: float = 0.9,
                   top_k: int = 50,
                   repetition_penalty: float = 1.2):
    """
    Generates responses using the selected model for video input.
    """
    if model_name == "Camel-Doc-OCR-080125(v2)":
        processor = processor_m
        model = model_m
    elif model_name == "OCRFlux-3B":
        processor = processor_x
        model = model_x
    elif model_name == "Behemoth-3B-070225":
        processor = processor_o
        model = model_o
    elif model_name == "MonkeyOCR-pro-1.2B":
        processor = processor_t
        model = model_t
    elif model_name == "ViGoRL-MCTS-SFT-7B":
        processor = processor_a
        model = model_a
    else:
        yield "Invalid model selected.", "Invalid model selected."
        return

    if video_path is None:
        yield "Please upload a video.", "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 frame in frames:
        image, timestamp = frame
        messages[1]["content"].append({"type": "text", "text": f"Frame {timestamp}:"})
        messages[1]["content"].append({"type": "image", "image": image})
    inputs = processor.apply_chat_template(
        messages,
        tokenize=True,
        add_generation_prompt=True,
        return_dict=True,
        return_tensors="pt",
        truncation=False,
        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,
        "do_sample": True,
        "temperature": temperature,
        "top_p": top_p,
        "top_k": top_k,
        "repetition_penalty": repetition_penalty,
    }
    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

# Define examples for image and video inference
image_examples = [
    ["Explain the essence of the image.", "assets/images/B.jpg"],
    ["Extract the content.", "assets/images/1.png"],
    ["Describe the safety of the action shown in the image.", "assets/images/C.jpg"],
    ["Caption the image.", "assets/images/A.jpg"],
    ["Make this into a table for the README.md file.", "assets/images/2.jpg"],
    ["Extract the table content from the image.", "assets/images/3.png"],
    ["Perform OCR on the image.", "assets/images/4.jpg"]
]

video_examples = [
    ["Explain the video in detail.", "assets/videos/a.mp4"],
    ["Explain the video in detail.", "assets/videos/b.mp4"]
]

#css
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;
}
"""

# Create the Gradio Interface
with gr.Blocks(css=css, theme="bethecloud/storj_theme") as demo:
    gr.Markdown(
        "# **[Multimodal OCR Outpost](https://huggingface.co/collections/prithivMLmods/multimodal-implementations-67c9982ea04b39f0608badb0)**"
    )
    with gr.Row():
        with gr.Column():
            with gr.Tabs():
                with gr.TabItem("Image Inference"):
                    image_query = gr.Textbox(
                        label="Query Input",
                        placeholder="Enter your query here...")
                    image_upload = gr.Image(type="pil", label="Image")
                    image_submit = gr.Button("Submit",
                                             elem_classes="submit-btn")
                    gr.Examples(examples=image_examples,
                                inputs=[image_query, image_upload])
                with gr.TabItem("Video Inference"):
                    video_query = gr.Textbox(
                        label="Query Input",
                        placeholder="Enter your query here...")
                    video_upload = gr.Video(label="Video")
                    video_submit = gr.Button("Submit",
                                             elem_classes="submit-btn")
                    gr.Examples(examples=video_examples,
                                inputs=[video_query, video_upload])

            with gr.Accordion("Advanced options", open=False):
                max_new_tokens = gr.Slider(label="Max new tokens",
                                           minimum=1,
                                           maximum=MAX_MAX_NEW_TOKENS,
                                           step=1,
                                           value=DEFAULT_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)

        with gr.Column():
            with gr.Column(elem_classes="canvas-output"):
                gr.Markdown("## Output")
                output = gr.Textbox(label="Raw Output Stream",
                                    interactive=False,
                                    lines=2, show_copy_button=True)
                with gr.Accordion("(Result.md)", open=False):
                    markdown_output = gr.Markdown(
                        label="markup.md")

            model_choice = gr.Radio(choices=[
                 "Camel-Doc-OCR-080125(v2)", "OCRFlux-3B",
                 "ViGoRL-MCTS-SFT-3B", "Behemoth-3B-070225",
                 "MonkeyOCR-pro-1.2B"],
                                    label="Select Model",
                                    value="Camel-Doc-OCR-080125(v2)")
            
            gr.Markdown("**Model Info 💻** | [Report Bug](https://huggingface.co/spaces/prithivMLmods/Multimodal-OCR-Outpost/discussions)")
            gr.Markdown("> Camel-Doc-OCR-080125 is a specialized vision-language model, fine-tuned from Qwen2.5-VL-7B-Instruct, and excels at document retrieval, content extraction, and analysis recognition for both structured and unstructured digital documents. OCRFlux-3B is a 3B-parameter vision-language model optimized for high-quality OCR on PDFs and images, excelling in converting documents to clean Markdown text and supporting features like cross-page table/paragraph merging.")    
            gr.Markdown("> Both ViGoRL-MCTS-SFT-3b-Spatial and 7b-Spatial are vision-language models that use multi-turn visually grounded reinforcement learning for precise spatial reasoning and visual grounding, with the 3b and 7b variants differing mainly in their architectural size for fine-grained visual tasks.")
            gr.Markdown("> Behemoth-3B-070225-post0.1 is an advanced 3B parameter model tailored for extensive multimodal comprehension, document parsing, and possibly generalized OCR/vision-language tasks. MonkeyOCR-pro-1.2B is a lightweight OCR model focusing on high-accuracy text extraction from images and scanned documents, suitable for resource-constrained environments.")    
            gr.Markdown("> ⚠️ Note: Models in this space may not perform well on video inference tasks.")
            
    # Define the submit button actions
    image_submit.click(fn=generate_image,
                       inputs=[
                           model_choice, image_query, image_upload,
                           max_new_tokens, temperature, top_p, top_k,
                           repetition_penalty
                       ],
                       outputs=[output, markdown_output])
    video_submit.click(fn=generate_video,
                       inputs=[
                           model_choice, video_query, video_upload,
                           max_new_tokens, temperature, top_p, top_k,
                           repetition_penalty
                       ],
                       outputs=[output, markdown_output])
        
if __name__ == "__main__":
    demo.queue(max_size=30).launch(share=True, mcp_server=True, ssr_mode=False, show_error=True)