Skip to content

Commit

Permalink
Add files via upload
Browse files Browse the repository at this point in the history
  • Loading branch information
PRITHIVSAKTHIUR authored Jan 28, 2025
1 parent 494a82d commit 3995a2c
Show file tree
Hide file tree
Showing 3 changed files with 125 additions and 0 deletions.
14 changes: 14 additions & 0 deletions README.md
Original file line number Diff line number Diff line change
@@ -0,0 +1,14 @@
---
title: OCR
emoji: 🍍
colorFrom: blue
colorTo: yellow
sdk: gradio
sdk_version: 5.13.1
app_file: app.py
pinned: true
license: creativeml-openrail-m
short_description: Qwen VL 2B
---

Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
103 changes: 103 additions & 0 deletions app.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,103 @@
import gradio as gr
from transformers import Qwen2VLForConditionalGeneration, AutoProcessor, TextIteratorStreamer
from transformers.image_utils import load_image
from threading import Thread
import time
import torch
import spaces

# Fine-tuned for OCR-based tasks from Qwen's [ Qwen/Qwen2-VL-2B-Instruct ]
MODEL_ID = "prithivMLmods/Qwen2-VL-OCR-2B-Instruct"
processor = AutoProcessor.from_pretrained(MODEL_ID, trust_remote_code=True)
model = Qwen2VLForConditionalGeneration.from_pretrained(
MODEL_ID,
trust_remote_code=True,
torch_dtype=torch.float16
).to("cuda").eval()

@spaces.GPU
def model_inference(input_dict, history):
text = input_dict["text"]
files = input_dict["files"]

# Load images if provided
if len(files) > 1:
images = [load_image(image) for image in files]
elif len(files) == 1:
images = [load_image(files[0])]
else:
images = []

# Validate input
if text == "" and not images:
gr.Error("Please input a query and optionally image(s).")
return
if text == "" and images:
gr.Error("Please input a text query along with the image(s).")
return

# Prepare messages for the model
messages = [
{
"role": "user",
"content": [
*[{"type": "image", "image": image} for image in images],
{"type": "text", "text": text},
],
}
]

# Apply chat template and process inputs
prompt = processor.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
inputs = processor(
text=[prompt],
images=images if images else None,
return_tensors="pt",
padding=True,
).to("cuda")

# Set up streamer for real-time output
streamer = TextIteratorStreamer(processor, skip_prompt=True, skip_special_tokens=True)
generation_kwargs = dict(inputs, streamer=streamer, max_new_tokens=1024)

# Start generation in a separate thread
thread = Thread(target=model.generate, kwargs=generation_kwargs)
thread.start()

# Stream the output
buffer = ""
yield "Thinking..."
for new_text in streamer:
buffer += new_text
# Remove <|im_end|> or similar tokens from the output
buffer = buffer.replace("<|im_end|>", "")
time.sleep(0.01)
yield buffer

# Example inputs
examples = [

[{"text": "Extract JSON from the image", "files": ["example_images/document.jpg"]}],
[{"text": "summarize the letter", "files": ["examples/1.png"]}],
[{"text": "Describe the photo", "files": ["examples/3.png"]}],
[{"text": "Extract as JSON table from the table", "files": ["examples/4.jpg"]}],
[{"text": "Summarize the full image in detail", "files": ["examples/2.jpg"]}],
[{"text": "Describe this image.", "files": ["example_images/campeones.jpg"]}],
[{"text": "What is this UI about?", "files": ["example_images/s2w_example.png"]}],
[{"text": "Can you describe this image?", "files": ["example_images/newyork.jpg"]}],
[{"text": "Can you describe this image?", "files": ["example_images/dogs.jpg"]}],
[{"text": "Where do the severe droughts happen according to this diagram?", "files": ["example_images/examples_weather_events.png"]}],

]

demo = gr.ChatInterface(
fn=model_inference,
description="# **Multimodal OCR**",
examples=examples,
textbox=gr.MultimodalTextbox(label="Query Input", file_types=["image"], file_count="multiple"),
stop_btn="Stop Generation",
multimodal=True,
cache_examples=False,
)

demo.launch(debug=True)
8 changes: 8 additions & 0 deletions requirements.txt
Original file line number Diff line number Diff line change
@@ -0,0 +1,8 @@
gradio_client==1.3.0
qwen-vl-utils==0.0.2
transformers-stream-generator==0.0.4
torch==2.4.0
torchvision==0.19.0
git+https://github.com/huggingface/transformers.git
accelerate
av

0 comments on commit 3995a2c

Please sign in to comment.