Failing in the input stage Using Encoder and Decoder Model Architecture (DONUT model)

[swapnil-lader]

I am trying to get a model summary of the donut model but am unable to define the input for the torch summary.
###########################################################
import argparse
import gradio as gr
import torch
from PIL import Image
from donut.donut.model import DonutModel
from torchvision import models
from torchsummary import summary

def demo_process_vqa(input_img, question):
global pretrained_model, task_prompt, task_name
# pretrained_model = './donut/result/train_docvqa/20220912_103244'
# task_name = "docvqa"
# task_prompt = "<s_pdf-donut>"
input_img = Image.fromarray(input_img)
user_prompt = task_prompt.replace("{user_input}", question)
print(user_prompt)
output = pretrained_model.inference(input_img, prompt=user_prompt)["predictions"][0]
print('inf_out',output)
return output

def demo_process(input_img):
global pretrained_model, task_prompt, task_name
input_img = Image.fromarray(input_img)
output = pretrained_model.inference(image=input_img, prompt=task_prompt)["predictions"][0]
return output

parser = argparse.ArgumentParser()
parser.add_argument("--task", type=str, default="docvqa")
parser.add_argument("--pretrained_path", type=str, default="train_docvqa_for_all_atts/donut/result/train_docvqa/20220915_125713")
args, left_argv = parser.parse_known_args()

task_name = args.task
if "docvqa" == task_name:
task_prompt = "<s_taco_eiko_pdf_donut>{user_input}</s_question><s_answer>"
else: # rvlcdip, cord, ...
task_prompt = f"<s_{task_name}>"

pretrained_model = DonutModel.from_pretrained(args.pretrained_path)

if torch.cuda.is_available():
# pretrained_model.half()
device = torch.device("cuda")
pretrained_model.to(device)
else:
pretrained_model.encoder.to(torch.bfloat16)

summary(pretrained_model, [(1, 3, 1280 , 960), (1, 21),(1, 21)])

The shape of the encoder and decoder is as follows.
Encoder : torch.Size([1, 3, 1280, 960])
Decode : torch.Size([1, 21])

##Model forward architecture looks like this

    encoder_outputs = self.encoder(image_tensors)
    decoder_outputs = self.decoder(
        input_ids=decoder_input_ids,
        encoder_hidden_states=encoder_outputs,
        labels=decoder_labels,
    )
    return decoder_outputs

Can you please guide how to pass down the model input in summary?