Spaces:
Running
on
Zero
Running
on
Zero
File size: 3,004 Bytes
2f3144c |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 |
import gradio as gr
from byaldi import RAGMultiModalModel
from transformers import Qwen2VLForConditionalGeneration, AutoProcessor
from qwen_vl_utils import process_vision_info
import torch
from PIL import Image
import os
import traceback
import spaces # Ensure import for GPU management
# Load the Byaldi and Qwen2-VL models without using .cuda()
rag_model = RAGMultiModalModel.from_pretrained("vidore/colpali")
qwen_model = Qwen2VLForConditionalGeneration.from_pretrained(
"Qwen/Qwen2-VL-2B-Instruct", trust_remote_code=True, torch_dtype=torch.bfloat16
)
# Processor for Qwen2-VL
processor = AutoProcessor.from_pretrained("Qwen/Qwen2-VL-2B-Instruct", trust_remote_code=True)
@spaces.GPU # Decorate the function for GPU management
def ocr_and_extract(image, text_query):
try:
# Save the uploaded image temporarily
temp_image_path = "temp_image.jpg"
image.save(temp_image_path)
# Index the image with Byaldi
rag_model.index(
input_path=temp_image_path,
index_name="image_index",
store_collection_with_index=False,
overwrite=True
)
# Perform the search query on the indexed image
results = rag_model.search(text_query, k=1)
# Prepare the input for Qwen2-VL
image_data = Image.open(temp_image_path)
messages = [
{
"role": "user",
"content": [
{"type": "image", "image": image_data},
{"type": "text", "text": text_query},
],
}
]
# Process the message and prepare for Qwen2-VL
text_input = processor.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
image_inputs, _ = process_vision_info(messages)
inputs = processor(
text=[text_input],
images=image_inputs,
padding=True,
return_tensors="pt",
)
# Move the Qwen2-VL model and inputs to GPU
qwen_model.to("cuda")
inputs = {k: v.to("cuda") for k, v in inputs.items()}
# Generate the output with Qwen2-VL
generated_ids = qwen_model.generate(**inputs, max_new_tokens=50)
output_text = processor.batch_decode(
generated_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False
)
# Clean up the temporary file
os.remove(temp_image_path)
return output_text[0]
except Exception as e:
error_message = str(e)
traceback.print_exc()
return f"Error: {error_message}"
# Gradio interface for image input
iface = gr.Interface(
fn=ocr_and_extract,
inputs=[
gr.Image(type="pil"),
gr.Textbox(label="Enter your query (optional)"),
],
outputs="text",
title="Image OCR with Byaldi + Qwen2-VL",
description="Upload an image (JPEG/PNG) containing Hindi and English text for OCR.",
)
# Launch the Gradio app
iface.launch() |