72 lines
2.8 KiB
Python
Executable file
72 lines
2.8 KiB
Python
Executable file
#!/usr/bin/env python3
|
|
import sys
|
|
import os
|
|
sys.path.insert(0, os.path.dirname(__file__))
|
|
from embd_input import MyModel
|
|
import numpy as np
|
|
from torch import nn
|
|
import torch
|
|
from transformers import CLIPVisionModel, CLIPImageProcessor
|
|
from PIL import Image
|
|
|
|
# model parameters from 'liuhaotian/LLaVA-13b-delta-v1-1'
|
|
vision_tower = "openai/clip-vit-large-patch14"
|
|
select_hidden_state_layer = -2
|
|
# (vision_config.image_size // vision_config.patch_size) ** 2
|
|
image_token_len = (224//14)**2
|
|
|
|
class Llava:
|
|
def __init__(self, args):
|
|
self.image_processor = CLIPImageProcessor.from_pretrained(vision_tower)
|
|
self.vision_tower = CLIPVisionModel.from_pretrained(vision_tower)
|
|
self.mm_projector = nn.Linear(1024, 5120)
|
|
self.model = MyModel(["main", *args])
|
|
|
|
def load_projection(self, path):
|
|
state = torch.load(path)
|
|
self.mm_projector.load_state_dict({
|
|
"weight": state["model.mm_projector.weight"],
|
|
"bias": state["model.mm_projector.bias"]})
|
|
|
|
def chat(self, question):
|
|
self.model.eval_string("user: ")
|
|
self.model.eval_string(question)
|
|
self.model.eval_string("\nassistant: ")
|
|
return self.model.generate_with_print()
|
|
|
|
def chat_with_image(self, image, question):
|
|
with torch.no_grad():
|
|
embd_image = self.image_processor.preprocess(image, return_tensors='pt')['pixel_values'][0]
|
|
image_forward_out = self.vision_tower(embd_image.unsqueeze(0), output_hidden_states=True)
|
|
select_hidden_state = image_forward_out.hidden_states[select_hidden_state_layer]
|
|
image_feature = select_hidden_state[:, 1:]
|
|
embd_image = self.mm_projector(image_feature)
|
|
embd_image = embd_image.cpu().numpy()[0]
|
|
self.model.eval_string("user: ")
|
|
self.model.eval_token(32003-2) # im_start
|
|
self.model.eval_float(embd_image.T)
|
|
for i in range(image_token_len-embd_image.shape[0]):
|
|
self.model.eval_token(32003-3) # im_patch
|
|
self.model.eval_token(32003-1) # im_end
|
|
self.model.eval_string(question)
|
|
self.model.eval_string("\nassistant: ")
|
|
return self.model.generate_with_print()
|
|
|
|
|
|
if __name__=="__main__":
|
|
# model form liuhaotian/LLaVA-13b-delta-v1-1
|
|
a = Llava(["--model", "./models/ggml-llava-13b-v1.1.bin", "-c", "2048"])
|
|
# Extract from https://huggingface.co/liuhaotian/LLaVA-13b-delta-v1-1/blob/main/pytorch_model-00003-of-00003.bin.
|
|
# Also here can use pytorch_model-00003-of-00003.bin directly.
|
|
a.load_projection(os.path.join(
|
|
os.path.dirname(__file__) ,
|
|
"llava_projection.pth"))
|
|
respose = a.chat_with_image(
|
|
Image.open("./media/llama1-logo.png").convert('RGB'),
|
|
"what is the text in the picture?")
|
|
respose
|
|
a.chat("what is the color of it?")
|
|
|
|
|
|
|