moe-llava

1
import torch
2
from PIL import Image
3
from moellava.constants import IMAGE_TOKEN_INDEX, DEFAULT_IMAGE_TOKEN
4
from moellava.conversation import conv_templates, SeparatorStyle
5
from moellava.model.builder import load_pretrained_model
6
from moellava.utils import disable_torch_init
7
from moellava.mm_utils import tokenizer_image_token, get_model_name_from_path, KeywordsStoppingCriteria
8

9
def main():
10
    disable_torch_init()
11
    image = 'moellava/serve/examples/extreme_ironing.jpg'
12
    inp = 'What is unusual about this image?'
13
    model_path = 'LanguageBind/MoE-LLaVA-xxxxxxxxxxxxxxxx'  # choose a model
14
    device = 'cuda'
15
    load_4bit, load_8bit = False, False
16
    model_name = get_model_name_from_path(model_path)
17
    tokenizer, model, processor, context_len = load_pretrained_model(model_path, None, model_name, load_8bit, load_4bit, device=device)
18
    image_processor = processor['image']
19
    conv_mode = "phi"  # phi or qwen or stablelm
20
    conv = conv_templates[conv_mode].copy()
21
    roles = conv.roles
22

23
    image_tensor = image_processor.preprocess(Image.open(image).convert('RGB'), return_tensors='pt')['pixel_values'].to(model.device, dtype=torch.float16)
24

25

26
    print(f"{roles[1]}: {inp}")
27
    inp = DEFAULT_IMAGE_TOKEN + '\n' + inp
28
    conv.append_message(conv.roles[0], inp)
29
    conv.append_message(conv.roles[1], None)
30
    prompt = conv.get_prompt()
31
    input_ids = tokenizer_image_token(prompt, tokenizer, IMAGE_TOKEN_INDEX, return_tensors='pt').unsqueeze(0).cuda()
32
    stop_str = conv.sep if conv.sep_style != SeparatorStyle.TWO else conv.sep2
33
    keywords = [stop_str]
34
    stopping_criteria = KeywordsStoppingCriteria(keywords, tokenizer, input_ids)
35

36
    with torch.inference_mode():
37
        output_ids = model.generate(
38
            input_ids,
39
            images=image_tensor,
40
            do_sample=True,
41
            temperature=0.2,
42
            max_new_tokens=1024,
43
            use_cache=True,
44
            stopping_criteria=[stopping_criteria])
45

46
    outputs = tokenizer.decode(output_ids[0, input_ids.shape[1]:], skip_special_tokens=True).strip()
47
    print(outputs)
48

49
if __name__ == '__main__':
50
    main()
51
    '''
52
    deepspeed predict.py
53
    '''
54