Update inference.py

inference.py

@@ -1,3 +1,4 @@
+import os
 import torch
 import re
 from PIL import Image
@@ -6,16 +7,16 @@ import base64
 import io
 import json
 
-def model_fn(model_dir):
+def model_fn(model_dir, context=None):
     device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
     processor = DonutProcessor.from_pretrained(model_dir)
     model = VisionEncoderDecoderModel.from_pretrained(model_dir)
     model.to(device)
     return model, processor, device
 
-def transform_fn(model, request_body, input_content_type, output_content_type):
+def transform_fn(model, request_body, request_content_type, response_content_type, context=None):
     model, processor, device = model
-    if input_content_type == 'application/json':
+    if request_content_type == 'application/json':
         data = json.loads(request_body)
         image_data = data['inputs']
         image = Image.open(io.BytesIO(base64.b64decode(image_data))).convert("RGB")
@@ -46,6 +47,6 @@ def transform_fn(model, request_body, input_content_type, output_content_type):
         
         # Prepare the response
         prediction = {'result': decoded_text}
-        return json.dumps(prediction), output_content_type
+        return json.dumps(prediction), response_content_type
     else:
-        raise ValueError(f"Unsupported content type: {input_content_type}")
+        raise ValueError(f"Unsupported content type: {request_content_type}")