Update app.py

app.py

@@ -1,24 +1,67 @@
 import gradio as gr
-import easyocr
 import time
+import spaces
+from PIL import Image
+from transformers import Qwen2VLForConditionalGeneration, AutoProcessor
+from qwen_vl_utils import process_vision_info
+import torch
+import uuid
 
+
+model_name = "oddadmix/Qari-OCR-0.1-VL-2B-Instruct"
+model = Qwen2VLForConditionalGeneration.from_pretrained(
+                model_name,
+                torch_dtype="auto",
+                device_map="cuda"
+            )
+processor = AutoProcessor.from_pretrained(model_name)
+max_tokens = 2000
+
+@spaces.GPU
 def perform_ocr(image):
-    reader = easyocr.Reader(['en', 'ar'])  # Add more languages if needed
-    result = reader.readtext(image, detail=0)
-    
-    def stream_results():
-        for text in result:
-            yield text
-            time.sleep(0.5)  # Simulate streaming delay
+    src = uuid.uuid4() + ".png"
+
+    prompt = "Below is the image of one page of a document, as well as some raw textual content that was previously extracted for it. Just return the plain text representation of this document as if you were reading it naturally. Do not hallucinate."
+    image.save(src)
 
-    return stream_results
+    messages = [
+        {
+            "role": "user",
+            "content": [
+                {"type": "image", "image": f"file://{src}"},
+                {"type": "text", "text": prompt},
+            ],
+        }
+    ]
+    text = processor.apply_chat_template(
+        messages, tokenize=False, add_generation_prompt=True
+    )
+    image_inputs, video_inputs = process_vision_info(messages)
+    inputs = processor(
+        text=[text],
+        images=image_inputs,
+        videos=video_inputs,
+        padding=True,
+        return_tensors="pt",
+    )
+    inputs = inputs.to("cuda")
+    generated_ids = model.generate(**inputs, max_new_tokens=max_tokens)
+    generated_ids_trimmed = [
+        out_ids[len(in_ids) :] for in_ids, out_ids in zip(inputs.input_ids, generated_ids)
+    ]
+    output_text = processor.batch_decode(
+        generated_ids_trimmed, skip_special_tokens=True, clean_up_tokenization_spaces=False
+    )[0]
+    os.remove(src)
+    print(output_text)
+    return output_text
 
 iface = gr.Interface(
     fn=perform_ocr,
     inputs=gr.Image(type="numpy"),
     outputs=gr.Textbox(),
     live=True,
-    title="OCR Streaming App",
+    title="Qari Arabic OCR",
     description="Upload an image to extract text in real-time."
 )