Update app.py

app.py

@@ -133,11 +133,13 @@ model_v = Qwen2_5_VLForConditionalGeneration.from_pretrained(
 ).to(device).eval()
 
 # Load PaddleOCR-VL
-MODEL_ID_P = "strangervisionhf/paddle"
-processor_p = AutoProcessor.from_pretrained(MODEL_ID_P, trust_remote_code=True)
+MODEL_ID_P = "PaddlePaddle/PaddleOCR-VL"
+SUBFOLDER_P = "PaddleOCR-VL-0.9B"
+processor_p = AutoProcessor.from_pretrained(MODEL_ID_P, trust_remote_code=True, subfolder=SUBFOLDER_P)
 model_p = AutoModelForCausalLM.from_pretrained(
     MODEL_ID_P,
     trust_remote_code=True,
+    subfolder=SUBFOLDER_P,
     torch_dtype=torch.float16
 ).to(device).eval()
 
@@ -157,7 +159,8 @@ def generate_image(model_name: str, text: str, image: Image.Image,
     if model_name == "Nanonets-OCR2-3B":
         processor = processor_v
         model = model_v
-
+        
+        # Nanonets/Qwen-VL format: content is a list of dicts
         messages = [{
             "role": "user",
             "content": [
@@ -172,7 +175,7 @@ def generate_image(model_name: str, text: str, image: Image.Image,
             images=[image],
             return_tensors="pt",
             padding=True).to(device)
-            
+
         streamer = TextIteratorStreamer(processor, skip_prompt=True, skip_special_tokens=True)
         generation_kwargs = {
             **inputs,
@@ -196,42 +199,39 @@ def generate_image(model_name: str, text: str, image: Image.Image,
     elif model_name == "PaddleOCR-VL":
         processor = processor_p
         model = model_p
-        
-        # FIX: PaddleOCR-VL expects a simple string content, not a list of dicts.
+
+        # PaddleOCR-VL format: content is a simple string
         messages = [{"role": "user", "content": text}]
         prompt_full = processor.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
+        
+        inputs = processor(text=[prompt_full], images=[image], return_tensors="pt").to(device)
 
-        inputs = processor(
-            text=[prompt_full], 
-            images=[image], 
-            return_tensors="pt"
-        ).to(device)
-
+        # Use generation parameters from the reference script for best results
         generation_kwargs = {
             **inputs,
             "max_new_tokens": max_new_tokens,
-            "do_sample": False,
+            "do_sample": False, 
             "use_cache": True,
         }
         
         with torch.inference_mode():
             generated_ids = model.generate(**generation_kwargs)
-
+            
         resp = processor.batch_decode(generated_ids, skip_special_tokens=True)[0]
-        # Clean the output by removing the prompt
+        # Extract only the generated part of the response
         answer = resp.split(prompt_full)[-1].strip()
         yield answer, answer
-
+        
     else:
         yield "Invalid model selected.", "Invalid model selected."
         return
 
 
-# Define examples for image inference
+# Define examples for image inference, tailored for both models
 image_examples = [
-    ["Extract the full page.", "images/ocr.png"],
-    ["OCR:", "images/4.png"], # Example prompt for PaddleOCR-VL
-    ["Table Recognition:", "images/0.png"] # Example prompt for PaddleOCR-VL
+    ["OCR:", "images/ocr.png"],
+    ["Table Recognition:", "images/4.png"],
+    ["Extract the content from this image.", "images/0.png"]
 ]
 
 
@@ -240,7 +240,7 @@ with gr.Blocks(css=css, theme=steel_blue_theme) as demo:
     gr.Markdown("# **Multimodal OCR**", elem_id="main-title")
     with gr.Row():
         with gr.Column(scale=2):
-            image_query = gr.Textbox(label="Query Input", placeholder="Enter your query here... (e.g., 'OCR:')")
+            image_query = gr.Textbox(label="Query Input", placeholder="Enter query or task (e.g., 'OCR:')")
             image_upload = gr.Image(type="pil", label="Upload Image", height=290)
 
             image_submit = gr.Button("Submit", variant="primary")