Update app.py

app.py

@@ -4,6 +4,7 @@ print("gradio version:", gr.__version__)
 
 import torch
 from transformers import DonutProcessor, VisionEncoderDecoderModel
+import multiprocessing
 
 processor = DonutProcessor.from_pretrained("naver-clova-ix/donut-base-finetuned-cord-v2")
 model = VisionEncoderDecoderModel.from_pretrained("naver-clova-ix/donut-base-finetuned-cord-v2")
@@ -11,6 +12,11 @@ model = VisionEncoderDecoderModel.from_pretrained("naver-clova-ix/donut-base-fin
 device = "cuda" if torch.cuda.is_available() else "cpu"
 model.to(device)
 
+# Number of threads
+# Set the number of threads you want to use
+desired_num_threads = multiprocessing.cpu_count()  # Change this value as needed
+torch.set_num_threads(desired_num_threads)
+
 def process_document(image):
     # prepare encoder inputs
     pixel_values = processor(image, return_tensors="pt").pixel_values