MaxP commited on
Commit
17a838d
1 Parent(s): 7a3063b
Files changed (1) hide show
  1. app.py +6 -4
app.py CHANGED
@@ -9,7 +9,7 @@ from torchvision import transforms
9
  import torch
10
  from PIL import Image
11
 
12
- device = 'cuda' if torch.cuda.is_available() else 'cpu'
13
  model_name = 'naver-clova-ix/donut-base-finetuned-docvqa'
14
 
15
  # Importante esta app esta pensada para que el modelo corra en CPU
@@ -27,8 +27,8 @@ def process_image(image, question):
27
  decoder_input_ids = processor.tokenizer(prompt, add_special_tokens=False, return_tensors='pt').input_ids
28
  # Defino los outputs
29
  outputs = model.generate(
30
- pixel_values.to(device).half(),
31
- decoder_input_ids=decoder_input_ids.to(device),
32
  max_length=model.decoder.config.max_position_embeddings,
33
  early_stopping=True,
34
  pad_token_id=processor.tokenizer.pad_token_id,
@@ -60,4 +60,6 @@ demo = gr.Interface(
60
  ['examples/extracto.jpg', 'cual es el telefono de centros servicios de banco galicia?'],
61
  ['examples/factura_5.jpg', 'cual es el total de la factura?'],
62
  ]
63
- )
 
 
 
9
  import torch
10
  from PIL import Image
11
 
12
+ # device = 'cuda' if torch.cuda.is_available() else 'cpu'
13
  model_name = 'naver-clova-ix/donut-base-finetuned-docvqa'
14
 
15
  # Importante esta app esta pensada para que el modelo corra en CPU
 
27
  decoder_input_ids = processor.tokenizer(prompt, add_special_tokens=False, return_tensors='pt').input_ids
28
  # Defino los outputs
29
  outputs = model.generate(
30
+ pixel_values,
31
+ decoder_input_ids=decoder_input_ids,
32
  max_length=model.decoder.config.max_position_embeddings,
33
  early_stopping=True,
34
  pad_token_id=processor.tokenizer.pad_token_id,
 
60
  ['examples/extracto.jpg', 'cual es el telefono de centros servicios de banco galicia?'],
61
  ['examples/factura_5.jpg', 'cual es el total de la factura?'],
62
  ]
63
+ )
64
+
65
+ demo.launch(inline=True)