Update handler.py

handler.py

@@ -7,14 +7,41 @@ from PIL import Image
 import os
 
 class EndpointHandler:
-    def __init__(self):
-        model_name = 'deepseek-ai/DeepSeek-OCR'
-        self.tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
-        model = AutoModel.from_pretrained(model_name, trust_remote_code=True)
-        self.model = model.eval().cuda().to(torch.bfloat16)  # Use .cpu() if no GPU
+    def __init__(self, model_dir = 'deepseek-ai/DeepSeek-OCR'):
+        model_name = model_dir
+        
+        self.tokenizer = AutoTokenizer.from_pretrained(
+            model_path, 
+            trust_remote_code=True,
+            local_files_only=bool(model_dir)  # Only use local files if model_dir is provided
+        )
+        
+        # Check if CUDA is available
+        self.device = 'cuda' if torch.cuda.is_available() else 'cpu'
+        
+        # Load model with appropriate settings
+        model_kwargs = {
+            'trust_remote_code': True,
+            'torch_dtype': torch.bfloat16 if self.device == 'cuda' else torch.float32
+        }
+        
+        # Add flash attention if available and on CUDA
+        if self.device == 'cuda':
+            try:
+                model_kwargs['_attn_implementation'] = 'flash_attention_2'
+            except:
+                pass  # Fall back to default if flash attention not available
+        
+        self.model = AutoModel.from_pretrained(model_path, **model_kwargs)
+        self.model = self.model.eval()
+        
+        # Move to appropriate device
+        if self.device == 'cuda':
+            self.model = self.model.cuda()
     
     def __call__(self, data: Dict[str, Any]) -> str:
         try:
+            inputs = data.get("inputs")
             base64_string = inputs["base64"]
             # Remove data URL prefix if present
             if ',' in base64_string: