| from transformers import AutoModel, AutoTokenizer |
| import os |
| import torch |
|
|
| class OCRModel: |
| _instance = None |
| |
| def __new__(cls): |
| if cls._instance is None: |
| cls._instance = super(OCRModel, cls).__new__(cls) |
| cls._instance.initialize() |
| return cls._instance |
| |
| def initialize(self): |
| |
| model_path = os.getenv('MODEL_PATH', 'ucaslcl/GOT-OCR2_0') |
| |
| self.tokenizer = AutoTokenizer.from_pretrained( |
| model_path, |
| trust_remote_code=True, |
| local_files_only=False |
| ) |
| |
| self.model = AutoModel.from_pretrained( |
| model_path, |
| trust_remote_code=True, |
| low_cpu_mem_usage=True, |
| device_map='auto', |
| use_safetensors=True, |
| pad_token_id=self.tokenizer.eos_token_id |
| ) |
| |
| self.model = self.model.eval() |
| |
| def process_image(self, image_path): |
| try: |
| with torch.no_grad(): |
| result = self.model.chat(self.tokenizer, image_path, ocr_type='format') |
| return result |
| except Exception as e: |
| return f"Error processing image: {str(e)}" |
