Download checkpoint from HF hub in OcrReorderPipeline

inference.py

@@ -29,7 +29,11 @@ class OcrReorderPipeline(Pipeline):
         self.projection.to(self.device)
 
     def _sanitize_parameters(self, **kwargs):
-        return {}, {}, {}
+        # extract the pipeline 'inputs' (base64 image) and custom args
+        image = kwargs.pop("inputs", None)
+        words = kwargs.pop("words", None)
+        boxes = kwargs.pop("boxes", None)
+        return {"image": image, "words": words, "boxes": boxes}, {}, {}
 
     def preprocess(self, image, words, boxes):
         data = base64.b64decode(image)
@@ -42,7 +46,7 @@ class OcrReorderPipeline(Pipeline):
     def _forward(self, model_inputs):
         pv, ids, mask, bbox = (
             model_inputs[k].to(self.device)
-            for k in ("pixel_values","input_ids","attention_mask","bbox")
+            for k in ("pixel_values", "input_ids", "attention_mask", "bbox")
         )
         vision_out = self.model.vision_model(
             pixel_values=pv,