Update app.py

app.py

@@ -1,46 +1,24 @@
 import gradio as gr
-import numpy as np
-import random
-
-# import spaces #[uncomment to use ZeroGPU]
-from transformers import DonutProcessor, VisionEncoderDecoderModel
-from PIL import Image
 import torch
+from PIL import Image
+from transformers import DonutProcessor, VisionEncoderDecoderModel
 
 device = "cuda" if torch.cuda.is_available() else "cpu"
 model_repo_id = "selvakumarcts/sk_invoice_receipts"
 
+# Load model and processor
 processor = DonutProcessor.from_pretrained(model_repo_id)
 model = VisionEncoderDecoderModel.from_pretrained(model_repo_id).to(device)
 
-
-if torch.cuda.is_available():
-    torch_dtype = torch.float16
-else:
-    torch_dtype = torch.float32
-
-pipe = DiffusionPipeline.from_pretrained(model_repo_id, torch_dtype=torch_dtype)
-pipe = pipe.to(device)
-
-MAX_SEED = np.iinfo(np.int32).max
-MAX_IMAGE_SIZE = 1024
-
-
-# @spaces.GPU #[uncomment to use ZeroGPU]
-def infer(image, progress=gr.Progress(track_tqdm=True)):
-    # Preprocess the image
+# Inference function
+def infer(image):
     image = image.convert("RGB")
     pixel_values = processor(image, return_tensors="pt").pixel_values.to(device)
-
-    # Run the model
     output = model.generate(pixel_values, max_length=512)
-
-    # Decode the output (JSON)
     result = processor.batch_decode(output, skip_special_tokens=True)[0]
-
     return result
 
-
+# UI
 with gr.Blocks() as demo:
     gr.Markdown(" # Invoice/Receipt Reader (Donut Model)")
     with gr.Column():