Update app.py

app.py

@@ -15,7 +15,6 @@ import cv2
 
 from transformers import (
     Qwen2VLForConditionalGeneration,
-    Glm4vForConditionalGeneration,
     Qwen2_5_VLForConditionalGeneration,
     AutoModelForImageTextToText,
     AutoProcessor,
@@ -67,10 +66,10 @@ model_g = Qwen2_5_VLForConditionalGeneration.from_pretrained(
 ).to(device).eval()
 #-----------------------------subfolder-----------------------------#
 
-# Load GLM-4.1V-9B-Thinking
-MODEL_ID_O = "THUDM/GLM-4.1V-9B-Thinking"
+# Load Pollux-OCR-Caption-2B
+MODEL_ID_O = "prithivMLmods/Pollux-OCR-Caption-2B"
 processor_o = AutoProcessor.from_pretrained(MODEL_ID_O, trust_remote_code=True)
-model_o = Glm4vForConditionalGeneration.from_pretrained(
+model_o = Qwen2_5_VLForConditionalGeneration.from_pretrained(
     MODEL_ID_O,
     trust_remote_code=True,
     torch_dtype=torch.float16
@@ -117,7 +116,7 @@ def generate_image(model_name: str, text: str, image: Image.Image,
     elif model_name == "MonkeyOCR-Recognition":
         processor = processor_g
         model = model_g
-    elif model_name == "GLM-4.1V-9B-Thinking":
+    elif model_name == "Pollux-OCR-Caption-2B":
         processor = processor_o
         model = model_o
     else:
@@ -175,7 +174,7 @@ def generate_video(model_name: str, text: str, video_path: str,
     elif model_name == "MonkeyOCR-Recognition":
         processor = processor_g
         model = model_g
-    elif model_name == "GLM-4.1V-9B-Thinking":
+    elif model_name == "Pollux-OCR-Caption-2B":
         processor = processor_o
         model = model_o
     else:
@@ -291,14 +290,14 @@ with gr.Blocks(css=css, theme="bethecloud/storj_theme") as demo:
                     markdown_output = gr.Markdown(label="(Result.Md)")
                     
             model_choice = gr.Radio(
-                choices=["docscopeOCR-7B-050425-exp", "MonkeyOCR-Recognition", "coreOCR-7B-050325-preview", "GLM-4.1V-9B-Thinking"],
+                choices=["docscopeOCR-7B-050425-exp", "MonkeyOCR-Recognition", "coreOCR-7B-050325-preview", "Pollux-OCR-Caption-2B"],
                 label="Select Model",
                 value="docscopeOCR-7B-050425-exp"
             )
             gr.Markdown("**Model Info 💻** | [Report Bug](https://huggingface.co/spaces/prithivMLmods/core-OCR/discussions)")
-            gr.Markdown("> [GLM-4.1V-9B-Thinking](https://huggingface.co/THUDM/GLM-4.1V-9B-Thinking): GLM-4.1V-9B-Thinking, designed to explore the upper limits of reasoning in vision-language models. By introducing a thinking paradigm and leveraging reinforcement learning, the model significantly enhances its capabilities. It achieves state-of-the-art performance among 10B-parameter VLMs.")
             gr.Markdown("> [docscopeOCR-7B-050425-exp](https://huggingface.co/prithivMLmods/docscopeOCR-7B-050425-exp): The docscopeOCR-7B-050425-exp model is a fine-tuned version of Qwen2.5-VL-7B-Instruct, optimized for Document-Level Optical Character Recognition (OCR), long-context vision-language understanding, and accurate image-to-text conversion with mathematical LaTeX formatting.")
             gr.Markdown("> [MonkeyOCR](https://huggingface.co/echo840/MonkeyOCR): MonkeyOCR adopts a Structure-Recognition-Relation (SRR) triplet paradigm, which simplifies the multi-tool pipeline of modular approaches while avoiding the inefficiency of using large multimodal models for full-page document processing.")
+            gr.Markdown("> [Pollux-OCR-Caption-2B](https://huggingface.co/prithivMLmods/Pollux-OCR-Caption-2B):  The Pollux-OCR-Caption-2B model is a fine-tuned version of Qwen2.5-VL architecture, optimized for Document Retrieval, Content Extraction, and Analysis Recognition. Built on top of the Qwen2.5-VL architecture.")
             gr.Markdown("> [coreOCR-7B-050325-preview](https://huggingface.co/prithivMLmods/coreOCR-7B-050325-preview): The coreOCR-7B-050325-preview model is a fine-tuned version of Qwen2-VL-7B, optimized for Document-Level Optical Character Recognition (OCR), long-context vision-language understanding, and accurate image-to-text conversion with mathematical LaTeX formatting.")
             gr.Markdown(">⚠️note: all the models in space are not guaranteed to perform well in video inference use cases.")