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 ViGoRL-MCTS-SFT-3b-Spatial MODEL_ID_P = "gsarch/ViGoRL-MCTS-SFT-3b-Spatial" processor_p = AutoProcessor.from_pretrained(MODEL_ID_P, trust_remote_code=True) model_p = Qwen2_5_VLForConditionalGeneration.from_pretrained( MODEL_ID_P, 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 elif model_name == "ViGoRL-MCTS-SFT-3B": processor = processor_p model = model_p 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 elif model_name == "ViGoRL-MCTS-SFT-3B": processor = processor_p model = model_p 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") #download_btn = gr.Button("Download Result.md") model_choice = gr.Radio(choices=[ "Camel-Doc-OCR-080125(v2)", "OCRFlux-3B", "ViGoRL-MCTS-SFT-7B", "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.") # 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)