Update donut_inference.py

donut_inference.py

@@ -26,7 +26,7 @@ def inference(image):
 
     # device = "cuda" if torch.cuda.is_available() else "cpu"
     # model.to(device)
-
+    start_time = time.time()
     outputs = model.generate(pixel_values.to(device),
                                 decoder_input_ids=decoder_input_ids.to(device),
                                 max_length=model.decoder.config.max_position_embeddings,
@@ -38,11 +38,12 @@ def inference(image):
                                 bad_words_ids=[[processor.tokenizer.unk_token_id]],
                                 return_dict_in_generate=True,
                                 output_scores=True,)
-    
+    end_time = time.time()
     sequence = processor.batch_decode(outputs.sequences)[0]
     sequence = sequence.replace(processor.tokenizer.eos_token, "").replace(processor.tokenizer.pad_token, "")
     sequence = re.sub(r"<.*?>", "", sequence, count=1).strip()  # remove first task start token
     print(processor.token2json(sequence))
+    print(f"Donut Inference time {start_time-end_time}")
     return processor.token2json(sequence)
 
 # data = inference(image)