LLaVA Next と OpenVINO による視覚言語アシスタント#

%pip install -q "openvino>=2024.0.0" "nncf>=2.9.0" "torch>=2.1" "transformers>=4.39.1" "accelerate" "pillow" "gradio>=4.26" "datasets>=2.14.6" "tqdm" --extra-index-url https://download.pytorch.org/whl/cpu

from pathlib import Path 

MODEL_DIR = Path("model") 
IMAGE_ENCODER_PATH = MODEL_DIR / "image_encoder.xml" 
INPUT_EMBEDDING_PATH = MODEL_DIR / "input_embeddings.xml" 
LANGUAGE_MODEL_PATH = MODEL_DIR / "language_model.xml" 

requires_pt_model_loading = not all([p.exists() for p in [IMAGE_ENCODER_PATH, INPUT_EMBEDDING_PATH, LANGUAGE_MODEL_PATH]])

from transformers import LlavaNextProcessor, LlavaNextForConditionalGeneration 
import torch 
import gc 

processor = LlavaNextProcessor.from_pretrained("llava-hf/llava-v1.6-mistral-7b-hf") 
image_encoder_model, input_embedding_model, language_model = None, None, None 

class ImageEncoder(torch.nn.Module): 
    def __init__(self, config, vision_tower, multi_modal_projector): 
        super().__init__() 
        self.config = config 
        self.vision_tower = vision_tower 
        self.multi_modal_projector = multi_modal_projector 

    def forward(self, pixel_values): 
        batch_size, num_patches, num_channels, height, width = pixel_values.shape 
        reshaped_pixel_values = pixel_values.view(batch_size * num_patches, num_channels, height, width) 
        image_features = self.vision_tower(reshaped_pixel_values, output_hidden_states=True) 
        selected_image_feature = image_features.hidden_states[self.config.vision_feature_layer] 
        if self.config.vision_feature_select_strategy == "default": 
            selected_image_feature = selected_image_feature[:, 1:] 
        elif self.config.vision_feature_select_strategy == "full": 
            selected_image_feature = selected_image_feature 
        image_features = self.multi_modal_projector(selected_image_feature) 
        return image_features 

if requires_pt_model_loading: 
    model = LlavaNextForConditionalGeneration.from_pretrained("llava-hf/llava-v1.6-mistral-7b-hf", low_cpu_mem_usage=True) 
    model.config.save_pretrained(MODEL_DIR) 
    image_encoder_model = ImageEncoder(model.config, model.vision_tower, model.multi_modal_projector) 
    input_embedding_model = input_embedding_model = model.get_input_embeddings() 
    language_model = model.language_model 
    del model 
    gc.collect()

import ipywidgets as widgets 

to_compress_weights = widgets.Checkbox( 
    value=True, 
    description="Weights Compression", 
    disabled=False, 
) 

to_compress_weights

import nncf 

compression_configuration = { 
    "mode": nncf.CompressWeightsMode.INT4_SYM, 
    "group_size": 64, 
    "ratio": 0.6, 
} 

LANGUAGE_MODEL_PATH_INT4 = LANGUAGE_MODEL_PATH.parent / LANGUAGE_MODEL_PATH.name.replace(".xml", "-int4.xml") 
if to_compress_weights.value and not LANGUAGE_MODEL_PATH_INT4.exists(): 
    ov_model = core.read_model(LANGUAGE_MODEL_PATH) 
    ov_compressed_model = nncf.compress_weights(ov_model, **compression_configuration) 
    ov.save_model(ov_compressed_model, LANGUAGE_MODEL_PATH_INT4) 
    del ov_compressed_model 
    del ov_model 
    gc.collect()

to_quantize = widgets.Checkbox( 
    value=True, 
    description="Quantization", 
    disabled=False, 
) 

to_quantize

IMAGE_ENCODER_PATH_INT8 = IMAGE_ENCODER_PATH.parent / IMAGE_ENCODER_PATH.name.replace(".xml", "-int8.xml") 

import requests 

r = requests.get( 

url="https://raw.githubusercontent.com/openvinotoolkit/openvino_notebooks/latest/utils/skip_kernel_extension.py", 
) 
open("skip_kernel_extension.py", "w").write(r.text) 

%load_ext skip_kernel_extension

import torch 
from transformers.generation import GenerationConfig, GenerationMixin 
from transformers.modeling_outputs import CausalLMOutputWithPast 
from transformers import AutoConfig 
from transformers.models.llava_next.modeling_llava_next import ( 
    get_anyres_image_grid_shape, 
    unpad_image, 
) 
import openvino as ov 

class OVLlavaForCausalLM(GenerationMixin): 
    def __init__( 
        self, 
        core, 
        image_encoder_path, 
        input_embedding_path, 
        language_model_path, 
        device, 
    ): 
        self.image_encoder = core.compile_model(core.read_model(image_encoder_path), device) 
        self.input_embeddings = core.compile_model(core.read_model(input_embedding_path), device) 
        self.model = core.read_model(language_model_path) 
        self.input_names = {key.get_any_name(): idx for idx, key in enumerate(self.model.inputs)} 
        self.output_names = {idx: key for idx, key in enumerate(self.model.outputs)} 
        self.key_value_input_names = [key for key in list(self.input_names) if key not in ["beam_idx", "inputs_embeds", "attention_mask", "position_ids"]] 
        self.key_value_output_names = [key for key in list(self.output_names)[1:]] 
        self.stateful = len(self.key_value_input_names) == 0 
        compiled_model = core.compile_model(self.model, device) 
        self.request = compiled_model.create_infer_request() 
        self.config = AutoConfig.from_pretrained(Path(language_model_path).parent) 
        self.generation_config = GenerationConfig.from_model_config(self.config) 
        self.main_input_name = "input_ids" 
        self.device = torch.device("cpu") 
        self.num_pkv = 2 
        self.next_beam_idx = None 
        self.image_newline = torch.zeros(self.config.text_config.hidden_size, dtype=torch.float32) 
        self.pad_token_id = self.config.pad_token_id if self.config.pad_token_id is not None else -1 
        self.past_len = 0 
        self._supports_cache_class = False 

def can_generate(self): 
    """Returns True to validate the check that the model using `GenerationMixin.generate()` can indeed generate.""" 
        return True 

    def __call__( 
        self, 
        input_ids: torch.LongTensor, 
        pixel_values: torch.Tensor, 
        attention_mask: Optional[torch.LongTensor] = None, 
        past_key_values: Optional[Tuple[Tuple[torch.FloatTensor]]] = None, 
        position_ids: Optional[torch.LongTensor] = None, 
        image_sizes=None, 
        **kwargs, 
    ) -> CausalLMOutputWithPast: 
        return self.forward( 
            input_ids, 
            pixel_values, 
            attention_mask, 
            past_key_values, 
            position_ids, 
            image_sizes, 
            **kwargs, 
        ) 

    def forward( 
        self, 
        input_ids: torch.LongTensor, 
        pixel_values: torch.Tensor, 
        attention_mask: Optional[torch.LongTensor] = None, 
        past_key_values: Optional[Tuple[Tuple[torch.FloatTensor]]] = None, 
        position_ids: Optional[torch.LongTensor] = None, 
        image_sizes=None, 
        **kwargs, 
    ) -> CausalLMOutputWithPast: 
        """General inference method""" 
        inputs = {} 
        if past_key_values is not None: 
            inputs = {} 
            if not self.stateful: 
                past_key_values = tuple(past_key_value for pkv_per_layer in past_key_values for past_key_value in pkv_per_layer) 
                # past_key_values をデコーダー入力に追加 
                inputs = dict(zip(self.key_value_input_names, past_key_values)) 
            # input_ids = np.array(input_ids)[:, -1:] 
            inputs_embeds = self.input_embeddings(input_ids)[0] 
            inputs["inputs_embeds"] = inputs_embeds 
            # inputs["attention_mask"] = attention_mask 
            if "beam_idx" in self.input_names: 
                inputs["beam_idx"] = self.next_beam_idx if self.next_beam_idx is not None else np.arange(batch_size, dtype=int) 

            if not self.stateful: 
                first_layer_past_key_value = torch.from_numpy(past_key_values[0][0][:, :, :, 0]) 
            else: 
                first_layer_past_key_value = torch.from_numpy(self.request.query_state()[0].state.data[:, :, :, 0]) 

            # ランダムエラーを避けるために、head_dim (-2) のすべての次元を合計:
            # https://github.com/huggingface/transformers/pull/28032#issuecomment-1863691941 
            batch_index, non_attended_tokens = torch.where(first_layer_past_key_value.float().sum(-2) == 0) 

            # ターゲット長を取得 
            target_length = input_ids.shape[1] 
            past_length = first_layer_past_key_value.shape[-1] 

            extended_attention_mask = torch.ones( 
                (attention_mask.shape[0], past_length), 
                dtype=attention_mask.dtype, 
                device=attention_mask.device, 
            ) 

            # 無視できるトークンのみをフィルタリングします。 
            # これは、最初の反復でキャッシュがすでに十分に大きい Llava + Fused モジュールを使用する場合、 
            # またはカスタムキャッシュを渡す場合に発生する可能性があります。 
            valid_indices = non_attended_tokens < extended_attention_mask.size(-1) 
            new_batch_index = batch_index[valid_indices] 
            new_non_attended_tokens = non_attended_tokens[valid_indices] 

            # 必要のない場所をゼロにする 
            extended_attention_mask[new_batch_index, new_non_attended_tokens] = 0 

            attention_mask = torch.cat((extended_attention_mask, attention_mask[:, -target_length:]), dim=1) 
            position_ids = torch.sum(attention_mask, dim=1).unsqueeze(-1) - 1 
            inputs["attention_mask"] = attention_mask 
            inputs["position_ids"] = position_ids
        else: 
            inputs = self.prepare_multimodal_input(input_ids, pixel_values, attention_mask, position_ids, image_sizes) 

        # 推論を実行 
        self.request.start_async(inputs, share_inputs=True) 
        self.request.wait() 

        logits = torch.from_numpy(self.request.get_tensor(self.output_names[0]).data) 

        if not self.stateful:
            # 長さが等しいタプル: レイヤー数 * デコーダーレイヤーごとの past_key_value の数 (2 は自己注意レイヤーに対応) 
            past_key_values = tuple(self.request.get_tensor(key).data for key in self.key_value_output_names) 
            # 長さが `n_layers` のタプルのタプル。各タプルの長さは 2 (自己注意の k/v) に等しい 
            past_key_values = tuple(past_key_values[i : i + self.num_pkv] for i in range(0, len(past_key_values), self.num_pkv)) 
        else: 
            past_key_values = ((),) 
        self.past_len += inputs["inputs_embeds"].shape[1] 
        return CausalLMOutputWithPast(logits=logits, past_key_values=past_key_values) 

    def prepare_multimodal_input(self, input_ids, pixel_values, attention_mask, position_ids, image_sizes=None): 
        """Preprocessing function for embedding multimodal data""」 
        inputs = {} 
        inputs_embeds = torch.from_numpy(self.input_embeddings(input_ids)[0]) 
        batch_size = input_ids.shape[0] 
        if not self.stateful: 
            for input_name in self.key_value_input_names: 
                model_inputs = self.modeget_anyres_image_grid_shapel.input(input_name) 
                shape = model_inputs.get_partial_shape() 
                shape[0] = batch_size 
                if shape[2].is_dynamic: 
                    shape[2] = 0 
                else: 
                    shape[1] = 0 
                inputs[input_name] = ov.Tensor(model_inputs.get_element_type(), shape.get_shape()) 
        else: 
            self.past_len = 0 
            self.request.reset_state() 
            # 現在の反復で使用される次の beam_idx 入力の初期値を設定し、 
            # beam_search が使用される場合は次の反復で _reorder_cache によってオプションで更新されます 
            self.next_beam_idx = np.arange(batch_size, dtype=int) 

        if "beam_idx" in self.input_names: 
            inputs["beam_idx"] = self.next_beam_idx if self.next_beam_idx is not None else np.arange(batch_size, dtype=int) 
        if pixel_values is None: 
            inputs["inputs_embeds"] = inputs_embeds 
            inputs["attention_mask"] = attention_mask 
            if position_ids is None: 
                position_ids = torch.cumsum(attention_mask, axis=1) - 1 
                position_ids[attention_mask == 0] = 1 
            inputs["position_ids"] = position_ids 
        res = self.image_encoder(pixel_values) 
        image_features = torch.from_numpy(res[0]) 
        split_sizes = [image.shape[0] for image in pixel_values] 
        image_features = torch.split(image_features, split_sizes, dim=0) 

        # 注: multimodal_patch_merge_type == "spatial_unpad" のみをサポートします 
        height = width = self.config.vision_config.image_size // self.config.vision_config.patch_size 

        new_image_features = [] 
        for image_idx, image_feature in enumerate(image_features): 
            if image_feature.shape[0] > 1: 
                base_image_feature = image_feature[0] 
                image_feature = image_feature[1:] 

                if height * width != base_image_feature.shape[0]: 
                    raise ValueError("The number of patches is not consistent with the image size.") 
                num_patch_height, num_patch_width = get_anyres_image_grid_shape( 
                    image_sizes[image_idx], 
                    self.config.image_grid_pinpoints, 
                    self.config.vision_config.image_size, 
                ) 
                image_feature = image_feature.view(num_patch_height, num_patch_width, height, width, -1) 

                image_feature = image_feature.permute(4, 0, 2, 1, 3).contiguous() 
                image_feature = image_feature.flatten(1, 2).flatten(2, 3) 
                image_feature = unpad_image(image_feature, image_sizes[image_idx]) 
                image_feature = torch.cat( 
                    ( 
                        image_feature, 
                        self.image_newline[:, None, None].expand(*image_feature.shape[:-1], 1), 
                    ), 
                    dim=-1, 
                ) 
                image_feature = image_feature.flatten(1, 2).transpose(0, 1) 
                image_feature = torch.cat((base_image_feature, image_feature), dim=0) 
            else: 
                image_feature = image_feature[0] 
                image_feature = torch.cat((image_feature, self.image_newline[None]), dim=0) 
            new_image_features.append(image_feature) 
        image_features = torch.stack(new_image_features, dim=0) 

        ( 
            inputs_embeds, 
            attention_mask, 
            position_ids, 
        ) = self._merge_input_ids_with_image_features(image_features, inputs_embeds, input_ids, attention_mask, None) 
        inputs["inputs_embeds"] = inputs_embeds 
        inputs["attention_mask"] = attention_mask 
        inputs["position_ids"] = position_ids 

        return inputs 

    def _merge_input_ids_with_image_features(self, image_features, inputs_embeds, input_ids, attention_mask, labels): 
        num_images, num_image_patches, embed_dim = image_features.shape 
        batch_size, sequence_length = input_ids.shape 
        left_padding = not torch.sum(input_ids[:, -1] == torch.tensor(self.pad_token_id)) 
        # 1. 特別な画像トークンがどこにあるかを知るマスクを作成 
        special_image_token_mask = input_ids == self.config.image_token_index 
        num_special_image_tokens = torch.sum(special_image_token_mask, dim=-1) 
        # 最大埋め込み次元を計算 
        max_embed_dim = (num_special_image_tokens.max() * (num_image_patches - 1)) + sequence_length 
        batch_indices, non_image_indices = torch.where(input_ids != self.config.image_token_index) 

        # 2. テキストを書き込む位置を計算 
        # 結合された画像からテキストへのシーケンス内のテキストトークンの新しい位置を計算 
        # special_image_token_mask` は画像トークンを識別。各画像トークンは `nb_text_tokens_per_images - 1` 個のテキストトークンに置き換えられます。 
        # `torch.cumsum` は、各画像トークンが後続のテキストトークンの位置をどのようにシフトするかを計算
        # `cumsum`  は本質的にインデックスを 1 ずつ増やすため、ゼロベースのインデックスを調整するために 1 を指定します。 
        new_token_positions = torch.cumsum((special_image_token_mask * (num_image_patches - 1) + 1), -1) - 1 
        nb_image_pad = max_embed_dim - 1 - new_token_positions[:, -1] 
        if left_padding: 
            new_token_positions += nb_image_pad[:, None] # 左パディングのオフセット 
        text_to_overwrite = new_token_positions[batch_indices, non_image_indices] 

        # 3. 最大位置までパディングされた完全な埋め込みを作成 
        final_embedding = torch.zeros( 
            batch_size, 
            max_embed_dim, 
            embed_dim, 
            dtype=inputs_embeds.dtype, 
            device=inputs_embeds.device, 
        ) 
        final_attention_mask = torch.zeros( 
            batch_size, 
            max_embed_dim, 
            dtype=attention_mask.dtype, 
            device=inputs_embeds.device, 
        ) 
        # ビジョンモデルまたは言語モデルが CPU にオフロードされている場合は、 
        # 対応するテンソルを正しいターゲットデバイスに手動で設定する必要があります 
        target_device = inputs_embeds.device 
        batch_indices, non_image_indices, text_to_overwrite = ( 
            batch_indices.to(target_device), 
            non_image_indices.to(target_device), 
            text_to_overwrite.to(target_device), 
        ) 
        attention_mask = attention_mask.to(target_device) 

        # 4. マスクに基づいて埋め込みを埋めます。["hey" "<image>", "how", "are"] がある場合、テキストについては 
        #  [0, 577, 578, 579] に、画像機能については [1:576] にコピーをインデックスする必要があります。 
        final_embedding[batch_indices, text_to_overwrite] = inputs_embeds[batch_indices, non_image_indices] 
        final_attention_mask[batch_indices, text_to_overwrite] = attention_mask[batch_indices, non_image_indices] 
        if labels is not None: 
            final_labels[batch_indices, text_to_overwrite] = labels[batch_indices, non_image_indices] 

        # 5. 画像に対応する埋め込みを入力。まだゼロになっているものはすべて埋めます 
        image_to_overwrite = torch.all(final_embedding == 0, dim=-1) 
        image_to_overwrite &= image_to_overwrite.cumsum(-1) - 1 >= nb_image_pad[:, None].to(target_device) 
        if image_to_overwrite.sum() != image_features.shape[:-1].numel(): 
            raise ValueError( 
                f"The input provided to the model are wrong.The number of image tokens is {torch.sum(special_image_token_mask)} while" 
                f" the number of image given to the model is {num_images}. This prevents correct indexing and breaks batch generation."            ) 

        final_embedding[image_to_overwrite] = image_features.contiguous().reshape(-1, embed_dim).to(target_device) 
        final_attention_mask |= image_to_overwrite 
        position_ids = (final_attention_mask.cumsum(-1) - 1).masked_fill_((final_attention_mask == 0), 1) 

        # 6. 後で past_key_value 値を使用して未対応のトークンを決定するため、パディング位置での埋め込みをマスク 
        batch_indices, pad_indices = torch.where(input_ids == self.pad_token_id) 
        indices_to_mask = new_token_positions[batch_indices, pad_indices] 

        final_embedding[batch_indices, indices_to_mask] = 0 

        return final_embedding, final_attention_mask, position_ids 

    def prepare_inputs_for_generation( 
        self, 
        input_ids, 
        past_key_values=None, 
        inputs_embeds=None, 
        pixel_values=None, 
        image_sizes=None, 
        attention_mask=None, 
        **kwargs, 
    ): 
        if past_key_values is not None: 
            if not self.stateful: 
                cache_length = past_length = past_key_values[0][0].shape[2] 
            else: 
                cache_length = past_length = self.past_len 

            # 未処理のトークンのみ保存:
            # 1 - attention_mask の長さが input_ids の長さを超える場合、 
            # 入力の一部がキャッシュの一部として排他的に渡される設定になります 
            # (例: input_embeds を入力として渡す場合) 
            if attention_mask is not None and attention_mask.shape[1] > input_ids.shape[1]: 
                input_ids = input_ids[:, -(attention_mask.shape[1] - past_length) :] 
            # 2 - past_length が input_ids より小さい場合、input_ids はすべての入力トークンを保持します。 
            # past_length.llava に基づいて input_ids を破棄できます 
            elif past_length < input_ids.shape[1]: 
                input_ids = input_ids[:, past_length:]
            # 3 - それ以外の場合 (past_length >= input_ids.shape[1])、input_ids には未処理のトークンのみがあると仮定 
            elif self.config.image_token_index in input_ids: 
                input_ids = input_ids[:, input_ids.shape[1] - 1 :] 
            # キャッシュが保持できる数を超えるトークンが検出された場合、キャッシュにはサイズ制限があります。 
            # 対応する値は入力の一部ではないため、古いアテンション値を破棄します 
            if cache_length < past_length and attention_mask is not None: 
                attention_mask = attention_mask[:, -(cache_length + input_ids.shape[1]) :] 

        position_ids = kwargs.get("position_ids", None) 
        if attention_mask is not None and position_ids is None:
            # バッチ gllavaenerationsubset_siz の position_ids をオンザフライで作成 
            position_ids = attention_mask.long().cumsum(-1) - 1 
            position_ids.masked_fill_(attention_mask == 0, 1) 
            if past_key_values: 
                position_ids = position_ids[:, -input_ids.shape[1] :]
        # `inputs_embeds` が渡された場合、最初の生成ステップでのみ使用 
        if inputs_embeds is not None and past_key_values is None: 
            model_inputs = {"inputs_embeds": inputs_embeds} 
        else: 
            model_inputs = {"input_ids": input_ids} 

        model_inputs.update( 
            { 
                "position_ids": position_ids, 
                "past_key_values": past_key_values, 
                "use_cache": kwargs.get("use_cache"), 
                "attention_mask": attention_mask, 
                "pixel_values": pixel_values, 
                "image_sizes": image_sizes, 
            } 
        ) 
        return model_inputs

import gradio as gr 
from transformers import TextIteratorStreamer 
from threading import Thread 
from PIL import Image 
import torch 

example_image_urls = [ 
    ( 
        "https://github.com/openvinotoolkit/openvino_notebooks/assets/29454499/1d6a0188-5613-418d-a1fd-4560aae1d907", 
        "bee.jpg",
    ), 
    ( 
        "https://github.com/openvinotoolkit/openvino_notebooks/assets/29454499/6cc7feeb-0721-4b5d-8791-2576ed9d2863", 
        "baklava.png", 
    ), 
] 
for url, file_name in example_image_urls:
     Image.open(requests.get(url, stream=True).raw).save(file_name) 

def bot_streaming(message, history): 
    print(message) 
    if message["files"]: 
        image = message["files"][-1]["path"] if isinstance(message["files"][-1], dict) else message["files"][-1] 
    else:
        # このターンに画像がアップロードされていない場合は、 
        # タプル内に保存されている過去のターンの画像を探し、最後のものを取得 
        for hist in history: 
            if isinstance(hist[0], tuple): 
                image = hist[0][0] 

    if image is None: 
        gr.Error("You need to upload an image for LLaVA to work.") 
    prompt = f"[INST] <image>\n{message['text']} [/INST]" 
    image = Image.open(image).convert("RGB") 
    inputs = processor(prompt, image, return_tensors="pt") 

    streamer = TextIteratorStreamer(processor, **{"skip_special_tokens": True}) 
    generation_kwargs = dict(inputs, streamer=streamer, max_new_tokens=100) 

    thread = Thread(target=ov_llava_model.generate, kwargs=generation_kwargs) 
    thread.start() 

    text_prompt = f"[INST] \n{message['text']} [/INST]" 

    buffer = "" 
    for new_text in streamer: 
        buffer += new_text 
        generated_text_without_prompt = buffer[len(text_prompt) :] 
        yield generated_text_without_prompt 

demo = gr.ChatInterface( 
    fn=bot_streaming, 
    title="LLaVA NeXT", 
    examples=[ 
        {"text": "What is on the flower?", "files": ["./bee.jpg"]}, 
        {"text": "How to make this pastry?", "files": ["./baklava.png"]}, 
    ], 
    description="Try [LLaVA NeXT](https://huggingface.co/docs/transformers/main/en/model_doc/llava_next) in this demo using OpenVINO. Upload an image and start chatting about it, or simply try one of the examples below. If you don't upload an image, you will receive an error.", 
    stop_btn="Stop Generation", 
    multimodal=True, 
) 

try: 
    demo.launch(debug=False) 
except Exception: 
    demo.launch(debug=False, share=True)