Update handler.py

handler.py

@@ -6,6 +6,9 @@ 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'):
@@ -79,6 +82,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")
                 try:
                     test_image = Image.open(image_path)
                     if test_image.mode != 'RGB':
@@ -87,21 +91,29 @@ class EndpointHandler:
                     print(f"Image verified: {test_image.size}, mode: {test_image.mode}")
                 except Exception as img_error:
                     return {"error": f"Invalid image: {str(img_error)}"}
-                print(f"Image path: {image_path}")
-                # Run OCR inference
-                result = self.model.infer(
-                    self.tokenizer,
-                    prompt=prompt,
-                    image_file=image_path,  # Pass the PIL Image object directly
-                    #output_path=temp_dir,
-                    base_size=1024,
-                    image_size=640,
-                    #crop_mode=True,
-                    #save_results=False
-                )
-                print(f"Inference completed. Result type: {type(result)}")
-                print(f"Result value: {result}")
-                return result
+
+                output_buffer = StringIO()
+                
+                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
         
         except Exception as e:
             print(f"Error processing image: {e}")