Update inference.py

inference.py

@@ -4,9 +4,6 @@ from PIL import Image
 from transformers import DonutProcessor, VisionEncoderDecoderModel
 import io
 import json
-import logging
-
-logging.basicConfig(level=logging.INFO)
 
 def model_fn(model_dir, context=None):
     device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
@@ -18,16 +15,16 @@ def model_fn(model_dir, context=None):
 def input_fn(input_data, content_type, context=None):
     """Deserialize the input data."""
     logging.info("Entering input_fn")
-    if content_type == 'application/json':
+    if content_type == 'application/x-image' or content_type == 'application/octet-stream':
         image = Image.open(io.BytesIO(input_data))
         return image
     else:
         raise ValueError(f"Unsupported content type: {content_type}")
 
-def predict_fn(data, model, context=None):
+def predict_fn(data, model_data, context=None):
     """Apply the model to the input data."""
     logging.info("Entering predict_fn")
-    model, processor, device = model
+    model, processor, device = model_data
 
     # Preprocess the image
     pixel_values = processor(data, return_tensors="pt").pixel_values.to(device)
@@ -60,6 +57,8 @@ def output_fn(prediction, accept):
     """Serialize the prediction output."""
     logging.info("Entering output_fn")
     if accept == 'application/json':
-        return json.dumps(prediction)
+        return json.dumps(prediction), 'application/json'
     else:
         raise ValueError(f"Unsupported response content type: {accept}")
+
+