Update app.py

app.py

@@ -1,5 +1,5 @@
 import gradio as gr 
-from transformers import TrOCRProcessor, VisionEncoderDecoderModel
+from transformers import TrOCRProcessor, VisionEncoderDecoderModel, AutoProcessor, AutoModelForVision2Seq
 from PIL import Image
 import numpy as np
 import torch
@@ -9,8 +9,8 @@ processor1 = TrOCRProcessor.from_pretrained("DeepDiveDev/transformodocs-ocr")
 model1 = VisionEncoderDecoderModel.from_pretrained("DeepDiveDev/transformodocs-ocr")
 
 # Load the fallback model (allenai/olmOCR-7B-0225-preview)
-processor2 = TrOCRProcessor.from_pretrained("allenai/olmOCR-7B-0225-preview")
-model2 = VisionEncoderDecoderModel.from_pretrained("allenai/olmOCR-7B-0225-preview")
+processor2 = AutoProcessor.from_pretrained("allenai/olmOCR-7B-0225-preview")
+model2 = AutoModelForVision2Seq.from_pretrained("allenai/olmOCR-7B-0225-preview", torch_dtype=torch.float16)
 
 # Function to extract text using both models
 def extract_text(image):
@@ -32,8 +32,8 @@ def extract_text(image):
 
         # If output seems incorrect, use the fallback model
         if len(extracted_text.strip()) < 2:  # If output is too short, retry with second model
-            pixel_values = processor2(images=image, return_tensors="pt").pixel_values
-            generated_ids = model2.generate(pixel_values)
+            inputs = processor2(images=image, return_tensors="pt").pixel_values
+            generated_ids = model2.generate(inputs)
             extracted_text = processor2.batch_decode(generated_ids, skip_special_tokens=True)[0]
 
         return extracted_text
@@ -50,4 +50,4 @@ iface = gr.Interface(
     description="Upload a handwritten document and get the extracted text.",
 )
 
-iface.launch()
+iface.launch()