davidr99 commited on
Commit
ab844dd
·
verified ·
1 Parent(s): bbe6d77

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +2 -2
app.py CHANGED
@@ -31,8 +31,8 @@ def blackjack_ai(image):
31
 
32
  prompt = "<image><bos>extract json\n"
33
  image_file = Image.fromarray(image.astype('uint8'), 'RGB')
34
- inputs = processor(images=[image_file], text=prompt, return_tensors="pt")
35
- output = model.generate(**inputs, max_new_tokens=125)
36
  return processor.decode(output[0], skip_special_tokens=True)[14:]
37
 
38
  with gr.Blocks() as demo:
 
31
 
32
  prompt = "<image><bos>extract json\n"
33
  image_file = Image.fromarray(image.astype('uint8'), 'RGB')
34
+ inputs = processor(images=[image_file], text=prompt, return_tensors="pt").to('cuda')
35
+ output = model.generate(**inputs, max_new_tokens=125).to('cpu')
36
  return processor.decode(output[0], skip_special_tokens=True)[14:]
37
 
38
  with gr.Blocks() as demo: