Views
No views yet
1from PIL import Image
2from transformers import NougatProcessor, VisionEncoderDecoderModel
3import torch
4
5processor = NougatProcessor.from_pretrained("powow/nougat-swe")
6model = VisionEncoderDecoderModel.from_pretrained("powow/nougat-swe")
7device = "cuda" if torch.cuda.is_available() else "cpu"
8
9model.to(device)
10
11image = Image.open(image_path)
12pixel_values = processor(image, return_tensors="pt").pixel_values
13
14outputs = model.generate(
15 pixel_values.to(device),
16 min_length=1,
17 max_new_tokens=3584,
18 bad_words_ids=[[processor.tokenizer.unk_token_id]],
19)
20
21sequence = processor.batch_decode(outputs, skip_special_tokens=True)[0]
22sequence = processor.post_process_generation(sequence, fix_markdown=False)
23
24print(sequence)