from io import BytesIO from PIL import Image import gradio as gr import re import torch from transformers import DonutProcessor, VisionEncoderDecoderModel from transformers import AutoProcessor, PaliGemmaProcessor, PaliGemmaForConditionalGeneration from transformers import AutoModelForVision2Seq from huggingface_hub import InferenceClient import base64 device = "cuda" if torch.cuda.is_available() else "cpu" model_choices = [ "idefics2", "paligemma", "donut" ] def load_donut_model(): processor = DonutProcessor.from_pretrained("naver-clova-ix/donut-base-finetuned-docvqa") model = VisionEncoderDecoderModel.from_pretrained("naver-clova-ix/donut-base-finetuned-docvqa") model.to(device) return model, processor def load_paligemma_docvqa(): # model_id = "google/paligemma-3b-ft-docvqa-896" model_id = "google/paligemma-3b-mix-448" processor = AutoProcessor.from_pretrained(model_id) model = PaliGemmaForConditionalGeneration.from_pretrained(model_id) model.to(device) return model, processor def load_idefics_docvqa(): model_id = "HuggingFaceM4/idefics2-8b" processor = AutoProcessor.from_pretrained(model_id) model = AutoModelForVision2Seq.from_pretrained(model_id) model.to(device) return model, processor def load_models(): # load donut donut_model, donut_processor = load_donut_model() print("donut downloaded") # #load paligemma pg_model, pg_processor = load_paligemma_docvqa() print("paligemma downloaded") return {"donut":[donut_model, donut_processor], "paligemma": [pg_model, pg_processor] } loaded_models = load_models() print("models loaded") def base64_encoded_image(image_array): im = Image.fromarray(image_array) buffered = BytesIO() im.save(buffered, format="PNG") image_bytes = buffered.getvalue() image_base64 = base64.b64encode(image_bytes).decode('ascii') return image_base64 def inference_calling_idefics(image_array, question): model_id = "HuggingFaceM4/idefics2-8b" client = InferenceClient(model=model_id) image_base64 = base64_encoded_image(image_array) image_info = f"data:image/png;base64,{image_base64}" prompt = f"![]({image_info}){question}\n\n" response = client.text_generation(prompt) return response def process_document_donut(image_array, question): model, processor = loaded_models.get("donut") # prepare encoder inputs pixel_values = processor(image_array, return_tensors="pt").pixel_values # prepare decoder inputs task_prompt = "{user_input}" prompt = task_prompt.replace("{user_input}", question) decoder_input_ids = processor.tokenizer(prompt, add_special_tokens=False, return_tensors="pt").input_ids # generate answer outputs = model.generate( pixel_values.to(device), decoder_input_ids=decoder_input_ids.to(device), max_length=model.decoder.config.max_position_embeddings, early_stopping=True, pad_token_id=processor.tokenizer.pad_token_id, eos_token_id=processor.tokenizer.eos_token_id, use_cache=True, num_beams=1, bad_words_ids=[[processor.tokenizer.unk_token_id]], return_dict_in_generate=True, ) # postprocess 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 op = processor.token2json(sequence) op = op.get("answer", str(op)) return op def process_document_pg(image_array, question): print("qustion :", question) print("called loaded model") model, processor = loaded_models.get("paligemma") print("converting inputs") inputs = processor(images=image_array, text=question, return_tensors="pt").to(device) print("get predictions") predictions = model.generate(**inputs, max_new_tokens=100) print("returning decoding") return processor.decode(predictions[0], skip_special_tokens=True)[len(question):].lstrip("\n") def process_document_idf(image_array, question): model, processor = loaded_models.get("idefics") inputs = processor(images=image_array, text=question, return_tensors="pt") #.to(device) predictions = model.generate(**inputs, max_new_tokens=100) return processor.decode(predictions[0], skip_special_tokens=True)[len(question):].lstrip("\n") def generate_answer_donut(image_array, question): try: print("processing document - donut") answer = process_document_donut(image_array, question) print(answer) return answer except Exception as e: print(e) gr.Warning("There is some issue, please try again later.") return "sorry :(" def generate_answer_idefics(image_array, question): try: print("processing document - idf2") # answer = process_document_idf(image_array, question) answer = inference_calling_idefics(image_array, question) print(answer) return answer except Exception as e: print(e) gr.Warning("There is some issue, please try again later.") return "sorry :(" def generate_answer_paligemma(image_array, question): try: print("processing document - pg") answer = process_document_pg(image_array, question) print(answer) return answer except Exception as e: print(e) gr.Warning("There is some issue, please try again later.") return "sorry :(" def generate_answers(image_path, question, selected_model=model_choices[0]): print("selected model: ", selected_model) try: if selected_model == "donut": print("generate answers donut") answer = generate_answer_donut(image_path, question) elif selected_model == "paligemma": print("generate answers pg") answer = generate_answer_paligemma(image_path, question) else: print("generate answers idf2") answer = generate_answer_idefics(image_path, question) return [answer] #[donut_answer, pg_answer, idf_answer] except Exception as e: print(e) gr.Warning("There is some issue, please try again later.") return ["sorry :("] def greet(name, shame, game): return "Hello " + shame + "!!" INTRO_TEXT = """## VQA demo\n\n VQA task models comparison This space is to compare multiple models on visual document question answering. \n\n """ with gr.Blocks(css="style.css") as demo: gr.Markdown(INTRO_TEXT) # with gr.Tab("Text Generation"): with gr.Column(): image = gr.Image(label="Input Image") question = gr.Text(label="Question") selected_model = gr.Radio(model_choices, label="Model", info="Select the model you want to run") outputs_answer = gr.Text(label="Answer generated by the selected model") run_button = gr.Button() inputs = [ image, question, selected_model ] outputs = [ outputs_answer ] run_button.click( fn=generate_answers, inputs=inputs, outputs=outputs, ) examples = [["images/sample_vendor_contract.png", "Agreement is between whom?"], ["images/apple-10k-form.png", "What were the EMEA revenues in 2017?"], ["images/infographic.png", "What is workforce in UPS?"], ["images/omr1.png", "What was the food quality of hospitality tent?"], ["images/omr2.png", "What is efficiency rating?"], ["images/omr3.png", "What is the selected reason code?"], ["images/omr4.png", "What is the product classification?"], ["images/cupon code 2.png", "The coupon code is adressed to whom?"], ["images/cupon code 2.png", "What is coupon expiration date?"], ["images/cupon code 2.png", "What is assigned code?"], ["images/completion form.png", "What is date posting completed?"], ["images/sender_receiver.png", "What is the fax phone number of the sender?"], ["images/marketing research.png", "What is the current available balance?"], ["images/toxicity.png", "What is the reported date?"], ["images/handwriting.png", "What is the contribution amount per pay period?"], ] gr.Examples( examples=examples, inputs=inputs, ) if __name__ == "__main__": demo.queue(max_size=10).launch(debug=True)