Update to Get Donut Results

app.py

@@ -1,20 +1,34 @@
 import io
 import requests
 import gradio as gr
-# from transformers import AutoModel, AutoTokenizer
 from huggingface_hub import list_models
 from datasets import load_dataset
 from typing import List
 from PIL import Image
+import torch
+from transformers import DonutProcessor, VisionEncoderDecoderModel
+import json
+import re
+import logging
+
+# Logging configuration
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+
+# Global variables for Donut model and processor
+donut_model = None
+donut_processor = None
 
 
 def get_image_names(dataset):
     return [str(i) for i in range(len(dataset))]
 
+
 def get_image_from_dataset(index):
     image_data = dataset[int(index)]["image"]
     return image_data
 
+
 def process_image(image=None, dataset_image_index=None):
     if dataset_image_index:
         image = get_image_from_dataset(dataset_image_index)
@@ -22,19 +36,20 @@ def process_image(image=None, dataset_image_index=None):
     return image
 
 
-
 def create_interface(tag, image_indices):
-    """ Create Gradio interface"""
+    """Create Gradio interface"""
     iface = gr.Interface(
         fn=process_image,
         inputs=[
             gr.Dropdown(choices=get_collection_models(tag), label="Select Model"),
             gr.Image(type="pil", label="Upload Image"),
-            gr.Dropdown(choices=image_indices, label="Select one from MERIT Dataset test-set"),
+            gr.Dropdown(
+                choices=image_indices, label="Select one from MERIT Dataset test-set"
+            ),
         ],
         outputs=gr.Image(label="Output Image"),
         title="Saliency Visualization",
-        description="Upload your image or select one from the MERIT Dataset test-set."
+        description="Upload your image or select one from the MERIT Dataset test-set.",
     )
     return iface
 
@@ -50,19 +65,102 @@ def get_collection_models(tag: str) -> List[str]:
 
     return model_names
 
+
 def load_model(model_name: str):
     """Load a model from Hugging Face Hub."""
     model = AutoModel.from_pretrained(model_name)
     tokenizer = AutoTokenizer.from_pretrained(model_name)
     return model, tokenizer
 
-# # Example processing function
-# def process_input(text: str, model_name: str) -> str:
-#     model, tokenizer = load_model(model_name)
-#     inputs = tokenizer(text, return_tensors="pt")
-#     outputs = model(**inputs)
-#     return f"Processed output with {model_name}"
 
+def get_donut():
+    global donut_model, donut_processor
+    if donut_model is None or donut_processor is None:
+        try:
+            donut_model = VisionEncoderDecoderModel.from_pretrained(
+                "de-Rodrigo/donut-merit"
+            )
+            donut_processor = DonutProcessor.from_pretrained("de-Rodrigo/donut-merit")
+            if torch.cuda.is_available():
+                donut_model = donut_model.to("cuda")
+            logger.info("Donut model loaded successfully")
+        except Exception as e:
+            logger.error(f"Error loading Donut model: {str(e)}")
+            raise
+    return donut_model, donut_processor
+
+
+def process_image_donut(model, processor, image):
+    try:
+        if not isinstance(image, Image.Image):
+            image = Image.fromarray(image)
+
+        pixel_values = processor(image, return_tensors="pt").pixel_values
+        if torch.cuda.is_available():
+            pixel_values = pixel_values.to("cuda")
+
+        task_prompt = "<s_cord-v2>"
+        decoder_input_ids = processor.tokenizer(
+            task_prompt, add_special_tokens=False, return_tensors="pt"
+        )["input_ids"]
+
+        outputs = model.generate(
+            pixel_values,
+            decoder_input_ids=decoder_input_ids,
+            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,
+        )
+
+        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()
+
+        result = processor.token2json(sequence)
+        return json.dumps(result, indent=2)
+    except Exception as e:
+        logger.error(f"Error processing image with Donut: {str(e)}")
+        return f"Error: {str(e)}"
+
+
+def process_image(model_name, image=None, dataset_image_index=None):
+    if dataset_image_index is not None:
+        image = get_image_from_dataset(dataset_image_index)
+
+    if model_name == "de-Rodrigo/donut-merit":
+        model, processor = get_donut()
+        result = process_image_donut(model, processor, image)
+    else:
+        # Here you should implement processing for other models
+        result = f"Processing for model {model_name} not implemented"
+
+    return image, result
+
+
+if __name__ == "__main__":
+    models = get_collection_models("saliency")
+    models.append("de-Rodrigo/donut-merit")
+
+    demo = gr.Interface(
+        fn=process_image,
+        inputs=[
+            gr.Dropdown(choices=models, label="Select Model"),
+            gr.Image(type="pil", label="Upload Image"),
+            gr.Slider(minimum=0, maximum=99, step=1, label="Dataset Image Index"),
+        ],
+        outputs=[gr.Image(label="Processed Image"), gr.Textbox(label="Result")],
+        title="Document Understanding with Donut",
+        description="Upload an image or select one from the dataset to process with the selected model.",
+    )
+
+    demo.launch()
 
 dataset_name = "de-Rodrigo/merit"
 dataset = load_dataset(dataset_name, name="en-digital-seq", split="train", num_proc=8)