Update app.py

app.py

@@ -12,66 +12,59 @@
 # if __name__ == "__main__":
 #     demo.launch()
 
-import gradio as gr
+import os
+from PIL import Image
 import torch
-from transformers import AutoProcessor, AutoModelForCausalLM
-
-# ------------------------------
-# Image-only Processor Wrapper
-# ------------------------------
-class ImageOnlyProcessor:
-    def __init__(self, processor):
-        self.image_processor = processor.image_processor
-        self.tokenizer = processor.tokenizer
-        self.feature_extractor = getattr(processor, "feature_extractor", None)
-
-    def __getattr__(self, name):
-        # Pass through to tokenizer if not found here
-        return getattr(self.tokenizer, name)
-
-# ------------------------------
-# Load Model + Processor
-# ------------------------------
-model_id = "rednote/dots-ocr"
-
-print("🔄 Loading processor...")
-_base_proc = AutoProcessor.from_pretrained(model_id, trust_remote_code=True)
-processor = ImageOnlyProcessor(_base_proc)
-
-print("🔄 Loading model (may take time)...")
-model = AutoModelForCausalLM.from_pretrained(
-    model_id,
-    trust_remote_code=True,
-    torch_dtype=torch.float16,
-    device_map="auto"
-)
+import gradio as gr
+
+# Make sure you have the local model folder (e.g., "DotsOCR") with all files from the repo
+LOCAL_MODEL_PATH = "./DotsOCR"
+
+# Import the model and processor code locally
+import sys
+sys.path.append(LOCAL_MODEL_PATH)
+
+from modeling_dots_ocr import DotsOCRForVision2Text
+from configuration_dots import DotsOCRConfig
+from transformers import PreTrainedTokenizerFast
+
+# Load tokenizer
+tokenizer = PreTrainedTokenizerFast.from_pretrained(LOCAL_MODEL_PATH)
 
-# ------------------------------
-# OCR Function
-# ------------------------------
-def run_ocr(image):
-    try:
-        # Preprocess
-        inputs = processor.image_processor(images=image, return_tensors="pt").to(model.device)
-
-        # Generate
-        generated_ids = model.generate(**inputs, max_new_tokens=256)
-        generated_text = processor.batch_decode(generated_ids, skip_special_tokens=True)[0]
-
-        return generated_text.strip()
-    except Exception as e:
-        return f"❌ Error: {str(e)}"
-
-# ------------------------------
-# Gradio UI
-# ------------------------------
+# Load model configuration
+config = DotsOCRConfig.from_pretrained(LOCAL_MODEL_PATH)
+
+# Load model
+device = "cuda" if torch.cuda.is_available() else "cpu"
+model = DotsOCRForVision2Text.from_pretrained(LOCAL_MODEL_PATH, config=config)
+model.to(device)
+model.eval()
+
+# Load only the image processor
+from transformers import AutoFeatureExtractor
+image_processor = AutoFeatureExtractor.from_pretrained(LOCAL_MODEL_PATH)
+
+def parse_document(image: Image.Image):
+    # Preprocess the image
+    inputs = image_processor(images=image, return_tensors="pt").to(device)
+
+    # Forward pass
+    with torch.no_grad():
+        output_ids = model.generate(**inputs, max_new_tokens=1024)
+    
+    # Decode output
+    text = tokenizer.batch_decode(output_ids, skip_special_tokens=True)[0]
+    return text
+
+# Gradio demo
 demo = gr.Interface(
-    fn=run_ocr,
+    fn=parse_document,
     inputs=gr.Image(type="pil"),
     outputs="text",
-    title="📖 DOTS-OCR (Image Only)",
-    description="Upload an image to extract text using the rednote/dots-ocr model (video support disabled)."
+    title="dots.ocr Document Parser",
+    description="Parse text from images using dots.ocr"
 )
 
 if __name__ == "__main__":
-    demo.launch(server_name="0.0.0.0", server_port=7860)
+    demo.launch()
+