LLaVA-JP is a vision-language model that can converse about input images.
This model was trained by fine-tuning
llm-jp/llm-jp-1.3b-v1.0 using
LLaVA method.
This model was initially trained with the Vision Projector using
LLaVA-CC3M-Pretrain-595K-JA and STAIR Captions.
In the second phase, it was fine-tuned with LLaVA-Instruct-150K-JA and Japanese Visual Genome.
1import 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'
19 model_args.vision_tower = "openai/clip-vit-large-patch14-336"
20 device = "cuda" if torch.cuda.is_available() else "cpu"
21 torch_dtype = torch.bfloat16 if device=="cuda" else torch.float32
22
23 model = LlavaGpt2ForCausalLM.from_pretrained(
24 model_path,
25 low_cpu_mem_usage=True,
26 use_safetensors=True,
27 torch_dtype=torch_dtype,
28 device_map=device,
29 )
30 tokenizer = transformers.AutoTokenizer.from_pretrained(
31 model_path,
32 model_max_length=1024,
33 padding_side="right",
34 use_fast=False,
35 )
36 model.eval()
37
38 conv_mode = "v1"
39 conv = conv_templates[conv_mode].copy()
40
41 # image pre-process
42 image_url = "https://huggingface.co/rinna/bilingual-gpt-neox-4b-minigpt4/resolve/main/sample.jpg"
43 image = Image.open(requests.get(image_url, stream=True).raw).convert('RGB')
44 if device == "cuda":
45 image_tensor = model.get_model().vision_tower.image_processor(image, return_tensors='pt')['pixel_values'].half().cuda().to(torch_dtype)
46 else:
47 image_tensor = model.get_model().vision_tower.image_processor(image, return_tensors='pt')['pixel_values'].to(torch_dtype)
48
49 # create prompt
50 # ユーザー: <image>\n{prompt}
51 prompt = "猫の隣には何がありますか？"
52 inp = DEFAULT_IMAGE_TOKEN + '\n' + prompt
53 conv.append_message(conv.roles[0], inp)
54 conv.append_message(conv.roles[1], None)
55 prompt = conv.get_prompt()
56
57 input_ids = tokenizer_image_token(
58 prompt,
59 tokenizer,
60 IMAGE_TOKEN_INDEX,
61 return_tensors='pt'
62 ).unsqueeze(0)
63 if device == "cuda":
64 input_ids = input_ids.to(device)
65
66 input_ids = input_ids[:, :-1] # </sep>がinputの最後に入るので削除する
67 stop_str = conv.sep if conv.sep_style != SeparatorStyle.TWO else conv.sep2
68 keywords = [stop_str]
69 streamer = TextStreamer(tokenizer, skip_prompt=True, timeout=20.0)
70
71 # predict
72 with torch.inference_mode():
73 model.generate(
74 inputs=input_ids,
75 images=image_tensor,
76 do_sample=True,
77 temperature=0.01,
78 top_p=1.0,
79 max_new_tokens=256,
80 streamer=streamer,
81 use_cache=True,
82 )
83 """ノートパソコン"""
