|
|
|
|
| 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" |
| model_name = "llava_llada" |
| device = "cuda" |
| device_map = "cuda:0" |
| conv_template = "llada" |
| question = DEFAULT_IMAGE_TOKEN + "\nWrite a story based on the image." |
|
|
| FIM = '<|reserved_token_1|>' |
| |
| draft_answer = f''' |
| Sure, here is a segment. |
| The boy climbed the hill with his dog.{"<|mdm_mask|>"*40}{FIM} They stood in silence as the star streaked across the sky. |
| <|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() |
| |
| |
| 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') |
|
|
| model.eval() |
| model.tie_weights() |
| model.to(torch.bfloat16) |
|
|
| img = 'images/port.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, |
| |
| 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]) |
| |
|
|
