Diegulio commited on
Commit
a17a2b8
1 Parent(s): a2d4f77

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +6 -0
app.py CHANGED
@@ -4,6 +4,7 @@ print("gradio version:", gr.__version__)
4
 
5
  import torch
6
  from transformers import DonutProcessor, VisionEncoderDecoderModel
 
7
 
8
  processor = DonutProcessor.from_pretrained("naver-clova-ix/donut-base-finetuned-cord-v2")
9
  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
11
  device = "cuda" if torch.cuda.is_available() else "cpu"
12
  model.to(device)
13
 
 
 
 
 
 
14
  def process_document(image):
15
  # prepare encoder inputs
16
  pixel_values = processor(image, return_tensors="pt").pixel_values
 
4
 
5
  import torch
6
  from transformers import DonutProcessor, VisionEncoderDecoderModel
7
+ import multiprocessing
8
 
9
  processor = DonutProcessor.from_pretrained("naver-clova-ix/donut-base-finetuned-cord-v2")
10
  model = VisionEncoderDecoderModel.from_pretrained("naver-clova-ix/donut-base-finetuned-cord-v2")
 
12
  device = "cuda" if torch.cuda.is_available() else "cpu"
13
  model.to(device)
14
 
15
+ # Number of threads
16
+ # Set the number of threads you want to use
17
+ desired_num_threads = multiprocessing.cpu_count() # Change this value as needed
18
+ torch.set_num_threads(desired_num_threads)
19
+
20
  def process_document(image):
21
  # prepare encoder inputs
22
  pixel_values = processor(image, return_tensors="pt").pixel_values