File size: 3,315 Bytes
81c1623 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 |
from llava.model.builder import load_pretrained_model
from llava.mm_utils import get_model_name_from_path, process_images, tokenizer_image_token
from llava.constants import IMAGE_TOKEN_INDEX, DEFAULT_IMAGE_TOKEN, DEFAULT_IM_START_TOKEN, DEFAULT_IM_END_TOKEN, IGNORE_INDEX
from llava.conversation import conv_templates, SeparatorStyle
from PIL import Image
import requests
import copy
import torch
import json
# pretrained = "lavida-ckpts/lavida-llada-hd-fim" # this may also work
pretrained = "lavida-ckpts/lavida-llada-hd-fim"
model_name = "llava_llada"
device = "cuda"
device_map = "cuda:0"
conv_template = "llada" # Make sure you use correct chat template for different models
question = DEFAULT_IMAGE_TOKEN + "\nWrite a poem about the image."
FIM = '<|reserved_token_1|>'
# FIM = ''
draft_answer = f'''
Sure, here is a poem about the image:
A<|mdm_mask|><|mdm_mask|><|mdm_mask|><|mdm_mask|><|mdm_mask|><|mdm_mask|><|mdm_mask|><|mdm_mask|><|mdm_mask|><|mdm_mask|><|mdm_mask|>{FIM}
Be<|mdm_mask|><|mdm_mask|><|mdm_mask|><|mdm_mask|><|mdm_mask|><|mdm_mask|><|mdm_mask|><|mdm_mask|><|mdm_mask|><|mdm_mask|><|mdm_mask|>{FIM}
Ca<|mdm_mask|><|mdm_mask|><|mdm_mask|><|mdm_mask|><|mdm_mask|><|mdm_mask|><|mdm_mask|><|mdm_mask|><|mdm_mask|><|mdm_mask|><|mdm_mask|>{FIM}
De<|mdm_mask|><|mdm_mask|><|mdm_mask|><|mdm_mask|><|mdm_mask|><|mdm_mask|><|mdm_mask|><|mdm_mask|><|mdm_mask|><|mdm_mask|><|mdm_mask|>{FIM}
<|eot_id|>'''
conv = copy.deepcopy(conv_templates[conv_template])
conv.append_message(conv.roles[0], question)
conv.append_message(conv.roles[1], None)
prompt_question = conv.get_prompt()
# breakpoint()
# exit()
vision_kwargs = None
vision_kwargs = dict(
mm_vision_tower="google/siglip-so400m-patch14-384",
mm_resampler_type=None,
mm_projector_type='mlp2x_gelu',
mm_hidden_size=1152,
use_mm_proj=True
)
tokenizer, model, image_processor, max_length = load_pretrained_model(pretrained, None, model_name, device_map=device_map,vision_kwargs=vision_kwargs,torch_dtype='bfloat16') # Add any other thing you want to pass in llava_model_args
model.eval()
model.tie_weights()
model.to(torch.bfloat16)
img = 'images/temple.png'
image = Image.open(img).convert('RGB')
image_tensor = process_images([image], image_processor, model.config)
image_tensor = [_image.to(dtype=torch.bfloat16, device=device) for _image in image_tensor]
input_ids = tokenizer_image_token(prompt_question, tokenizer, IMAGE_TOKEN_INDEX, return_tensors="pt").unsqueeze(0).to(device)
image_sizes = [image.size]
if True:
draft_tokens = tokenizer(draft_answer,return_tensors='pt').to(input_ids.device).input_ids
cont,_ = model.generate(
input_ids,
images=image_tensor,
image_sizes=image_sizes,
do_sample=False,
temperature=0.0,
max_new_tokens=256,
step_ratio=1.0,
# steps=128
tokenizer=tokenizer,
prefix_lm=True,
verbose=True,
draft_tokens=draft_tokens,
schedule='shift',
schedule_kwargs=dict(shift=1/3),
)
text_outputs = tokenizer.batch_decode(cont, skip_special_tokens=False)
text_outputs = [text_output.lstrip('!').replace('<|endoftext|>','') for text_output in text_outputs]
print(text_outputs[0])
# print(tokenizer.batch_decode(cont[0], skip_special_tokens=True))
|