| 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): |
| self.tokenizer = AutoTokenizer.from_pretrained('ucaslcl/GOT-OCR2_0', trust_remote_code=True) |
| self.model = AutoModel.from_pretrained( |
| 'ucaslcl/GOT-OCR2_0', |
| trust_remote_code=True, |
| low_cpu_mem_usage=True, |
| device_map='cuda' if torch.cuda.is_available() else 'cpu', |
| use_safetensors=True, |
| pad_token_id=self.tokenizer.eos_token_id |
| ) |
| self.model = self.model.eval() |
| if torch.cuda.is_available(): |
| self.model = self.model.cuda() |
| |
| def process_image(self, image_path): |
| try: |
| result = self.model.chat(self.tokenizer, image_path, ocr_type='format') |
| return result |
| except Exception as e: |
| return str(e) |
