Update app.py

app.py

@@ -16,6 +16,7 @@ import cv2
 from transformers import (
     Qwen2VLForConditionalGeneration,
     Qwen2_5_VLForConditionalGeneration,
+    AutoModelForCausalLM,
     AutoModelForVision2Seq,
     AutoProcessor,
     TextIteratorStreamer,
@@ -77,6 +78,16 @@ model_x = AutoModelForVision2Seq.from_pretrained(
     torch_dtype=torch.float16
 ).to(device).eval()
 
+# Load dots.ocr
+MODEL_ID_D = "rednote-hilab/dots.ocr"
+processor_d = AutoProcessor.from_pretrained(MODEL_ID_D, trust_remote_code=True)
+model_d = AutoModelForCausalLM.from_pretrained(
+    MODEL_ID_D,
+    attn_implementation="flash_attention_2",
+    trust_remote_code=True,
+    torch_dtype=torch.float16
+).to(device).eval()
+
 # Preprocessing functions for SmolDocling-256M
 def add_random_padding(image, min_percent=0.1, max_percent=0.10):
     """Add random padding to an image based on its size."""
@@ -143,6 +154,9 @@ def generate_image(model_name: str, text: str, image: Image.Image,
     elif model_name == "Typhoon-OCR-7B":
         processor = processor_l
         model = model_l
+    elif model_name == "rednote-dots.ocr":
+        processor = processor_d
+        model = model_d
     else:
         yield "Invalid model selected.", "Invalid model selected."
         return
@@ -221,6 +235,9 @@ def generate_video(model_name: str, text: str, video_path: str,
     elif model_name == "Typhoon-OCR-7B":
         processor = processor_l
         model = model_l
+    elif model_name == "rednote-dots.ocr":
+        processor = processor_d
+        model = model_d
     else:
         yield "Invalid model selected.", "Invalid model selected."
         return
@@ -351,7 +368,7 @@ with gr.Blocks(css=css, theme="bethecloud/storj_theme") as demo:
                     formatted_output = gr.Markdown(label="(Result.md)")
             
             model_choice = gr.Radio(
-                choices=["Nanonets-OCR-s", "MonkeyOCR-Recognition", "SmolDocling-256M-preview", "Typhoon-OCR-7B"],
+                choices=["Nanonets-OCR-s", "MonkeyOCR-Recognition", "rednote-dots.ocr", "Typhoon-OCR-7B", "SmolDocling-256M-preview"],
                 label="Select Model",
                 value="Nanonets-OCR-s"
             )