Views
No views yet
git clone https://github.com/tosiyuki/LLaVA-JP.git1import requests
2import torch
3import transformers
4from PIL import Image
5
6from transformers.generation.streamers import TextStreamer
7from llava.constants import DEFAULT_IMAGE_TOKEN, IMAGE_TOKEN_INDEX
8from llava.conversation import conv_templates, SeparatorStyle
9from llava.model.llava_gpt2 import LlavaGpt2ForCausalLM
10from llava.train.arguments_dataclass import ModelArguments, DataArguments, TrainingArguments
11from llava.train.dataset import tokenizer_image_token
12
13
14if __name__ == "__main__":
15 parser = transformers.HfArgumentParser(
16 (ModelArguments, DataArguments, TrainingArguments))
17 model_args, data_args, training_args = parser.parse_args_into_dataclasses()
18 model_path = 'toshi456/llava-jp-1.3b-v1.0-620k'
19 device = "cuda" if torch.cuda.is_available() else "cpu"
20 torch_dtype = torch.bfloat16 if device=="cuda" else torch.float32
21
22 model = LlavaGpt2ForCausalLM.from_pretrained(
23 model_path,
24 low_cpu_mem_usage=True,
25 use_safetensors=True,
26 torch_dtype=torch_dtype,
27 device_map=device,
28 )
29 tokenizer = transformers.AutoTokenizer.from_pretrained(
30 model_path,
31 model_max_length=1532,
32 padding_side="right",
33 use_fast=False,
34 )
35 model.eval()
36
37 conv_mode = "v1"
38 conv = conv_templates[conv_mode].copy()
39
40 # image pre-process
41 image_url = "https://huggingface.co/rinna/bilingual-gpt-neox-4b-minigpt4/resolve/main/sample.jpg"
42 image = Image.open(requests.get(image_url, stream=True).raw).convert('RGB')
43
44 image_size = model.get_model().vision_tower.image_processor.size["height"]
45 if model.get_model().vision_tower.scales is not None:
46 image_size = model.get_model().vision_tower.image_processor.size["height"] * len(model.get_model().vision_tower.scales)
47
48 if device == "cuda":
49 image_tensor = model.get_model().vision_tower.image_processor(
50 image,
51 return_tensors='pt',
52 size={"height": image_size, "width": image_size}
53 )['pixel_values'].half().cuda().to(torch_dtype)
54 else:
55 image_tensor = model.get_model().vision_tower.image_processor(
56 image,
57 return_tensors='pt',
58 size={"height": image_size, "width": image_size}
59 )['pixel_values'].to(torch_dtype)
60
61 # create prompt
62 # ユーザー: <image>\n{prompt}
63 prompt = "猫の隣には何がありますか?"
64 inp = DEFAULT_IMAGE_TOKEN + '\n' + prompt
65 conv.append_message(conv.roles[0], inp)
66 conv.append_message(conv.roles[1], None)
67 prompt = conv.get_prompt()
68
69 input_ids = tokenizer_image_token(
70 prompt,
71 tokenizer,
72 IMAGE_TOKEN_INDEX,
73 return_tensors='pt'
74 ).unsqueeze(0)
75 if device == "cuda":
76 input_ids = input_ids.to(device)
77
78 input_ids = input_ids[:, :-1] # </sep>がinputの最後に入るので削除する
79 stop_str = conv.sep if conv.sep_style != SeparatorStyle.TWO else conv.sep2
80 keywords = [stop_str]
81 streamer = TextStreamer(tokenizer, skip_prompt=True, timeout=20.0)
82
83 # predict
84 with torch.inference_mode():
85 model.generate(
86 inputs=input_ids,
87 images=image_tensor,
88 do_sample=True,
89 temperature=0.01,
90 top_p=1.0,
91 max_new_tokens=256,
92 streamer=streamer,
93 use_cache=True,
94 )
95
96 """猫の隣にはノートパソコンがあります。"""