shivamAttarkar commited on
Commit
a81fbb9
1 Parent(s): 814421e

updated code for CPU use instead of cuda

Browse files
Files changed (1) hide show
  1. app.py +2 -2
app.py CHANGED
@@ -4,8 +4,8 @@ from PIL import Image
4
 
5
  from transformers import AutoModel, AutoTokenizer
6
  tokenizer = AutoTokenizer.from_pretrained('ucaslcl/GOT-OCR2_0', trust_remote_code=True)
7
- model = AutoModel.from_pretrained('ucaslcl/GOT-OCR2_0', trust_remote_code=True, low_cpu_mem_usage=True, device_map='cuda', use_safetensors=True, pad_token_id=tokenizer.eos_token_id)
8
- model = model.eval().cuda()
9
 
10
  def OCR(image):
11
  pil_image = Image.fromarray(image)
 
4
 
5
  from transformers import AutoModel, AutoTokenizer
6
  tokenizer = AutoTokenizer.from_pretrained('ucaslcl/GOT-OCR2_0', trust_remote_code=True)
7
+ model = AutoModel.from_pretrained('ucaslcl/GOT-OCR2_0', trust_remote_code=True, low_cpu_mem_usage=True, use_safetensors=True, pad_token_id=tokenizer.eos_token_id)
8
+ model = model.eval()
9
 
10
  def OCR(image):
11
  pil_image = Image.fromarray(image)