Update handler.py

handler.py

@@ -51,26 +51,33 @@ class EndpointHandler:
             image_data = base64.b64decode(base64_string)
             image = Image.open(BytesIO(image_data))
         
-        # Convert to RGB if necessary (handles PNG, JPEG, etc.)
+            # Convert to RGB if necessary (handles PNG, JPEG, etc.)
             if image.mode != 'RGB':
                 image = image.convert('RGB')
         
-        # Define the prompt for Markdown conversion
+            # Define the prompt for Markdown conversion
             prompt = "<image>\n<|grounding|>Convert the document to markdown."
         
-        # Run OCR inference
-            result = self.model.infer(
-                self.tokenizer,
-                prompt=prompt,
-                image_file=image,  # Pass PIL Image directly
-                output_path=output_path,
-                base_size=1024,
-                image_size=640,
-                crop_mode=True,
-                save_results=output_path is not None
-            )
-        
-            return result
+            # Run OCR inference
+            with tempfile.TemporaryDirectory() as temp_dir:
+                output_path = temp_dir  # Now output_path is defined
+                
+                # Define the prompt for Markdown conversion
+                prompt = "<image>\n<|grounding|>Convert the document to markdown."
+            
+                # Run OCR inference
+                result = self.model.infer(
+                    self.tokenizer,
+                    prompt=prompt,
+                    image_file=image,
+                    output_path=output_path,  # Use temporary directory
+                    base_size=1024,
+                    image_size=640,
+                    crop_mode=True,
+                    save_results=True  # Save to temporary directory
+                )
+            
+                return result
         
         except Exception as e:
             print(f"Error processing image: {e}")