Views
No views yet
1
2import requests
3from PIL import Image
4
5from transformers import ViTFeatureExtractor, AutoTokenizer, FlaxVisionEncoderDecoderModel
6
7loc = "ydshieh/flax-vit-gpt2-coco-en"
8
9feature_extractor = ViTFeatureExtractor.from_pretrained(loc)
10tokenizer = AutoTokenizer.from_pretrained(loc)
11model = FlaxVisionEncoderDecoderModel.from_pretrained(loc)
12
13# We will verify our results on an image of cute cats
14url = "http://images.cocodataset.org/val2017/000000039769.jpg"
15with Image.open(requests.get(url, stream=True).raw) as img:
16 pixel_values = feature_extractor(images=img, return_tensors="np").pixel_values
17
18def generate_step(pixel_values):
19
20 output_ids = model.generate(pixel_values, max_length=16, num_beams=4).sequences
21 preds = tokenizer.batch_decode(output_ids, skip_special_tokens=True)
22 preds = [pred.strip() for pred in preds]
23
24 return preds
25
26preds = generate_step(pixel_values)
27print(preds)
28
29# should produce
30# ['a cat laying on top of a couch next to another cat']
31