MohamedRashad commited on
Commit
cfeccec
1 Parent(s): 74a2ba9

Add GPU support for text extraction

Browse files
Files changed (1) hide show
  1. app.py +2 -0
app.py CHANGED
@@ -4,6 +4,7 @@ import torch
4
  from PIL import Image
5
  from pathlib import Path
6
  from pdf2image import convert_from_path
 
7
 
8
  # Load the model and processor
9
  processor = NougatProcessor.from_pretrained("MohamedRashad/arabic-small-nougat")
@@ -14,6 +15,7 @@ model.to(device)
14
  print(f"Using {device} device")
15
  context_length = 2048
16
 
 
17
  def extract_text_from_image(image):
18
  """
19
  Extract text from PIL image
 
4
  from PIL import Image
5
  from pathlib import Path
6
  from pdf2image import convert_from_path
7
+ import spaces
8
 
9
  # Load the model and processor
10
  processor = NougatProcessor.from_pretrained("MohamedRashad/arabic-small-nougat")
 
15
  print(f"Using {device} device")
16
  context_length = 2048
17
 
18
+ @spaces.GPU
19
  def extract_text_from_image(image):
20
  """
21
  Extract text from PIL image