Update app/utils.py

app/utils.py

@@ -18,22 +18,35 @@ class OCRModel:
 
     def initialize(self):
         try:
-            logger.info("Initializing OCR model...")
             
-            # Model initilization
-            self.tokenizer = AutoTokenizer.from_pretrained('stepfun-ai/GOT-OCR-2.0-hf', trust_remote_code=True)
+            logger.info("Initializing OCR model...")
+            # Try loading with use_fast=False if the fast tokenizer fails
+            try:
+                self.tokenizer = AutoTokenizer.from_pretrained(
+                    'stepfun-ai/GOT-OCR-2.0-hf', 
+                    trust_remote_code=True,
+                    use_fast=False  # Try with slow tokenizer
+                )
+            except Exception as e:
+                logger.warning(f"Fast tokenizer failed, trying alternative: {str(e)}")
+                self.tokenizer = AutoTokenizer.from_pretrained(
+                    'stepfun-ai/GOT-OCR-2.0-hf', 
+                    trust_remote_code=True,
+                    use_fast=False
+                )
+                
             self.model = AutoModel.from_pretrained(
                 'stepfun-ai/GOT-OCR-2.0-hf',
                 trust_remote_code=True,
                 low_cpu_mem_usage=True,
-                device_map='cuda',
+                device_map='auto',  # Let transformers decide the best device
                 use_safetensors=True,
                 pad_token_id=self.tokenizer.eos_token_id
             )
             
-            # choose cpu
-            self.device = "cuda"  
-            self.model = self.model.eval().cpu()
+            # Let the model decide device placement
+            self.device = "cuda" if torch.cuda.is_available() else "cpu"
+            self.model = self.model.eval().to(self.device)
             
             logger.info("Model initialization completed successfully")