Update handler.py

handler.py

@@ -6,9 +6,6 @@ from io import BytesIO
 from PIL import Image
 import os
 import tempfile
-import sys
-from io import StringIO
-from contextlib import redirect_stdout
 
 class EndpointHandler:
     def __init__(self, model_dir = 'deepseek-ai/DeepSeek-OCR'):
@@ -82,7 +79,7 @@ class EndpointHandler:
                 print(f"Image saved to: {image_path}")
                 
                 # Verify the image can be opened
-                output_file = os.path.join(temp_dir, "output.md")
+                output_file = os.path.join(temp_dir, "output-1.md")
                 try:
                     test_image = Image.open(image_path)
                     if test_image.mode != 'RGB':
@@ -92,28 +89,22 @@ class EndpointHandler:
                 except Exception as img_error:
                     return {"error": f"Invalid image: {str(img_error)}"}
 
-                output_buffer = StringIO()
+                 # Run OCR inference
+                result = self.model.infer(
+                    self.tokenizer,
+                    prompt=prompt,
+                    image_file=image_path,  # Pass the PIL Image object directly
+                    output_path=output_file,
+                    base_size=1024,
+                    image_size=640,
+                    crop_mode=True,
+                    save_results=True
+                )
                 
-                with redirect_stdout(output_buffer):
-                    
-                    # Run OCR inference
-                    result = self.model.infer(
-                        self.tokenizer,
-                        prompt=prompt,
-                        image_file=image_path,  # Pass the PIL Image object directly
-                        output_path=output_file,
-                        base_size=1024,
-                        image_size=640,
-                        crop_mode=True,
-                        save_results=True
-                    )
-                    
-                captured_output = output_buffer.getvalue()
-                
-                print(f"Captured output length: {len(captured_output)}")
-                print(f"First 500 chars of captured output: {captured_output[:500]}")
-                 
-                return captured_output
+                if os.path.exists(output_file):
+                    with open(output_file, 'r', encoding='utf-8') as f:
+                        content = f.read()
+                    return content
         
         except Exception as e:
             print(f"Error processing image: {e}")