MiniCPM-V2 と OpenVINO による視覚言語アシスタント#

%pip install -q "torch>=2.1" "torchvision" "timm" "transformers>=4.40" "Pillow" "gradio>=4.19" "tqdm" "sentencepiece" --extra-index-url https://download.pytorch.org/whl/cpu 
%pip install -q "openvino>=2024.2.0" "nncf>=2.11.0"

from transformers import AutoModel, AutoTokenizer 
from pathlib import Path 

model_dir = Path("model") 
text_emb_path = model_dir / "language_model/embed_tokens.xml" 
image_encoder_path = model_dir / "image_encoder.xml" 
llm_path = model_dir / "language_model/language_model.xml" 

model = None 

if not all([text_emb_path.exists(), image_encoder_path.exists(), llm_path.exists()]): 
    model = AutoModel.from_pretrained("openbmb/MiniCPM-V-2", trust_remote_code=True) 
    model.eval() 
    model.config.save_pretrained(model_dir) 
    tokenizer = AutoTokenizer.from_pretrained("openbmb/MiniCPM-V-2", trust_remote_code=True) 
    tokenizer.save_pretrained(model_dir)

from transformers.generation import GenerationMixin 
from transformers import AutoConfig, GenerationConfig 

core = ov.Core() 

class OvModelForCausalLMWithEmb(GenerationMixin): 
    def __init__(self, model_dir, device="CPU", ov_config=None, compile=True) -> None: 
        self._supports_cache_class = False 
        self.config = AutoConfig.from_pretrained(model_dir, trust_remote_code=True) 
        self.config.is_decoder = True 
        self.config.is_encoder_decoder = False 
        self.generation_config = GenerationConfig.from_model_config(self.config) 
        model_dir = Path(model_dir) 
        self.model = core.read_model(model_dir / "language_model.xml") 
        self.token_emb = core.read_model(model_dir / "embed_tokens.xml") 
        self.request = None 
        self.token_emb_request = None 
        self._device = device.upper() 
        self.device = torch.device("cpu") 
        self.ov_config = ov_config 
        self.next_beam_idx = None 
        self._past_length = None 
        self.input_names = [input_t.get_any_name() for input_t in self.model.inputs] 
        self.main_input_name = "input_ids" 
        if compile: 
            self.compile() 

    def compile(self): 
        if self.request is None: 
            self.request = core.compile_model(self.model, self._device, self.ov_config).create_infer_request() 
        self._compile_token_emb() 

    def _compile_token_emb(self): 
        if self.token_emb_request is None: 
            self.token_emb_request = core.compile_model(self.token_emb, self._device, self.ov_config) 

    def to(self, device: str): 
        if isinstance(device, str): 
            self._device = device.upper() 
            self.clear_requests() 

        return self 

    def clear_requests(self): 
        del self.request 
        del self.token_emb_request 
        self.request = None 
        self.token_emb_request = None 

    def embed_tokens(self, input_ids: torch.LongTensor): 
        self._compile_token_emb() 
        res = self.token_emb_request(input_ids, share_inputs=True) 
        return res[0] 

    def prepare_inputs( 
        self, 
        input_ids: torch.LongTensor, 
        attention_mask: Optional[torch.LongTensor] = None, 
        past_key_values: Optional[Tuple[Tuple[torch.FloatTensor]]] = None, 
        position_ids: Optional[torch.LongTensor] = None, 
        inputs_embeds: Optional[torch.FloatTensor] = None, 
        **kwargs, 
    ): 
        batch_size = input_ids.shape[0] if input_ids is not None else inputs_embeds.shape[0] 

        inputs = {} 
        # past_key_values は明示的には使用されず、モデル内で処理されます 
        if past_key_values is None:
            # これはシーケンスの最初の反復であり、すべての状態をリセットします 
            if self.request is not None: 
                self.request.reset_state() 
                # 現在の反復で使用される次の beam_idx 入力の初期値を設定し、 
                # beam_search が使用される場合は次の反復で _reorder_cache によってオプションで更新されます 
                self.next_beam_idx = np.arange(batch_size, dtype=int) 
                self._past_length = 0 
        past_len = self._get_past_length(past_key_values) 

        if inputs_embeds is None: 
            inputs_embeds = self.embed_tokens(input_ids if past_key_values is None else input_ids[:, -1:])* self.config.scale_emb 
        inputs["inputs_embeds"] = inputs_embeds 

        # 必要に応じて attention_mask 入力を追加 
        if "attention_mask" in self.input_names or "position_ids" in self.input_names: 
            if attention_mask is not None: 
                attention_mask = np.array(attention_mask) 
            else: 
                attention_mask = np.ones((inputs_embeds.shape[0], inputs_embeds.shape[1] + past_len), dtype=int) 

        if "attention_mask" in self.input_names: 
            inputs["attention_mask"] = attention_mask 

        if "position_ids" in self.input_names: 
            if position_ids is not None: 
                position_ids = np.array(position_ids) 
            else: 
                position_ids = np.cumsum(attention_mask, axis=1) - 1 
                position_ids[attention_mask == 0] = 1 
                if past_key_values: 
                    position_ids = position_ids[:, -input_ids.shape[1] :] 

            inputs["position_ids"] = position_ids 

        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) 

        return inputs 

    def forward( 
        self, 
        input_ids: torch.LongTensor, 
        attention_mask: Optional[torch.LongTensor] = None, 
        past_key_values: Optional[Tuple[Tuple[torch.FloatTensor]]] = None, 
        position_ids: Optional[torch.LongTensor] = None, 
        inputs_embeds: Optional[torch.LongTensor] = None, 
        **kwargs, 
    ): 
        self.compile() 

        inputs = self.prepare_inputs( 
            input_ids=input_ids, 
            attention_mask=attention_mask, 
            past_key_values=past_key_values, 
            position_ids=position_ids, 
            inputs_embeds=inputs_embeds, 
            **kwargs, 
        ) 

        # 推論を実行 
        self.request.start_async(inputs, share_inputs=True) 
        self.request.wait() 
        logits = self.request.get_tensor("logits").data 
        logits = torch.from_numpy(logits).to(self.device) 
        past_key_values = ((),) 
        self._past_length += inputs["inputs_embeds"].shape[1] 

        return CausalLMOutputWithPast(logits=logits, past_key_values=past_key_values) 

    # transformers.models.llama.modeling_llama.LlamaForCausalLM.prepare_inputs_for_generation からトランスフォーム 
    def prepare_inputs_for_generation(self, input_ids, past_key_values=None, inputs_embeds=None, **kwargs):
        # モデルがエンコーダー・デコーダー・モデルでデコーダーとして使用される場合、デコーダー・アテンション・マスクはオンザフライで作成されます 
        attention_mask = kwargs.get("attention_mask", None) 
        use_cache = kwargs.get("use_cache", None) 

        if past_key_values is not None: 
            past_len = self._get_past_length(past_key_values) 
            # 未処理のトークンのみ保存:
            # 1 - attention_mask の長さが input_ids の長さを超える場合、 
            # 入力の一部がキャッシュの一部として排他的に渡される設定になります 
            # (例: input_embeds を入力として渡す場合) 
            if attention_mask is not None and input_ids is not None and attention_mask.shape[1] > input_ids.shape[1]: 
                input_ids = input_ids[:, -(attention_mask.shape[1] - past_len) :] 
            # 2 - past_length が input_ids より小さい場合、input_ids はすべての入力トークンを保持します。past_length に基づいて 
            # input_ids を破棄できます 
            elif input_ids is not None and past_len < input_ids.shape[1]: 
                input_ids = input_ids[:, past_len:]
            # 3 - それ以外の場合 (past_length >= input_ids.shape[1])、input_ids には未処理のトークンのみがあると仮定 
        position_ids = kwargs.get("position_ids", None) 
        if attention_mask is not None and position_ids is None and "position_ids" in self.input_names:
            # バッチ生成のために position_ids をオンザフライで作成 
            position_ids = attention_mask.long().cumsum(-1) - 1 
            position_ids.masked_fill_(attention_mask == 0, 1) 
            if past_key_values and input_ids is not None: 
                position_ids = position_ids[:, -input_ids.shape[1] :] 

            model_inputs = { 
                "input_ids": input_ids, 
                "past_key_values": past_key_values, 
                "use_cache": use_cache, 
                "position_ids": position_ids, 
                "attention_mask": attention_mask, 
                "inputs_embeds": inputs_embeds if past_key_values is None else None, 
            } 

            return model_inputs 

        def _get_past_length(self, past_key_values=None): 
            if past_key_values is None: 
                return 0 
            return self._past_length 

    # transformers.models.gpt2.modeling_gpt2.GPT2LMHeadModel._reorder_cache から適用 
    def _reorder_cache(self, past_key_values: Tuple[Tuple[torch.Tensor]], beam_idx: torch.Tensor) -> Tuple[Tuple[torch.Tensor]]: 
        """ 
        This function is used to re-order the `past_key_values` cache if [`~PreTrainedModel.beam_search`] or 
        [`~PreTrainedModel.beam_sample`] is called. 
        This is required to match `past_key_values` with the correct beam_idx at every generation step.
        """ 
        self.next_beam_idx = np.array(beam_idx) # 次の反復で入力として使用するために beam_idx を保存 
        return past_key_values 

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

        return True 

    def __call__(self, *args, **kwargs): 
        return self.forward(*args, **kwargs)

from typing import List, Optional 
import math 
import json 
import torch 
from torchvision import transforms 
from timm.data import IMAGENET_INCEPTION_MEAN, IMAGENET_INCEPTION_STD 
from PIL import Image 

def pad(orig_items, key, max_length=None, padding_value=0, padding_side="left"): 
    items = [] 
    if isinstance(orig_items[0][key], list): 
        assert isinstance(orig_items[0][key][0], torch.Tensor) 
        for it in orig_items: 
            for tr in it[key]: 
                items.append({key: tr})
    else: 
        assert isinstance(orig_items[0][key], torch.Tensor) 
        items = orig_items 

    batch_size = len(items) 
    shape = items[0][key].shape 
    dim = len(shape) 
    assert dim <= 3 
    if max_length is None: 
        max_length = 0 
    max_length = max(max_length, max(item[key].shape[-1] for item in items)) 
    min_length = min(item[key].shape[-1] for item in items) 
    dtype = items[0][key].dtype 

    if dim == 1: 
        return torch.cat([item[key] for item in items], dim=0) 
    elif dim == 2: 
        if max_length == min_length: 
            return torch.cat([item[key] for item in items], dim=0) 
        tensor = torch.zeros((batch_size, max_length), dtype=dtype) + padding_value 
    else: 
        tensor = torch.zeros((batch_size, max_length, shape[-1]), dtype=dtype) + padding_value 

    for i, item in enumerate(items): 
        if dim == 2: 
            if padding_side == "left": 
                tensor[i, -len(item[key][0]) :]= item[key][0].clone() 
            else: 
                tensor[i, : len(item[key][0])] = item[key][0].clone() 
        elif dim == 3: 
            if padding_side == "left": 
                tensor[i, -len(item[key][0]) :, :]= item[key][0].clone() 
            else: 
                tensor[i, : len(item[key][0]), :]= item[key][0].clone() 

    return tensor 

def slice_image(image, max_slice_nums=9, scale_resolution=448, patch_size=14, never_split=False): 
    original_size = image.size 
    original_width, original_height = original_size 
    log_ratio = math.log(original_width / original_height) 
    ratio = original_width * original_height / (scale_resolution * scale_resolution) 
    multiple = min(math.ceil(ratio), max_slice_nums) 

    source_image = None 
    best_grid = None 
    patches = [] 

    if multiple <= 1 or never_split:
        # スライスやアップサンプリングは不要 
        best_size = find_best_resize(original_size, scale_resolution, patch_size, allow_upscale=True) 
        source_image = image.resize(best_size, Image.Resampling.BICUBIC) 
    else: 
        candidate_split_grids_nums = [] 
        for i in [multiple - 1, multiple, multiple + 1]: 
            if i == 1 or i > max_slice_nums: 
                continue 
            candidate_split_grids_nums.append(i) 

        # ソース画像、ダウンサンプリング、patch_size で割った値 
        best_resize = find_best_resize(original_size, scale_resolution, patch_size) 
        source_image = image.copy().resize(best_resize, Image.Resampling.BICUBIC) 
        candidate_grids = [] 

        # 最適なグリッドを検出  
        for split_grids_nums in candidate_split_grids_nums: 
            m = 1 
            while m <= split_grids_nums: 
                if split_grids_nums % m == 0: 
                    candidate_grids.append([m, split_grids_nums // m]) 
                m += 1 

        best_grid = [1, 1] 
        min_error = float("inf") 
        for grid in candidate_grids: 
            error = abs(log_ratio - math.log(grid[0] / grid[1])) 
            if error < min_error: 
                best_grid = grid 
                min_error = error 

        refine_size = get_refine_size(original_size, best_grid, scale_resolution, patch_size, allow_upscale=True) 

        refine_image = image.resize(refine_size, Image.Resampling.BICUBIC) 
        patches = split_to_patches(refine_image, best_grid) 

    return source_image, patches, best_grid 

def ensure_divide(length, patch_size): 
    return max(round(length / patch_size) * patch_size, patch_size) 

def find_best_resize(original_size, scale_resolution, patch_size, allow_upscale=False): 
    width, height = original_size 
    if (width * height > scale_resolution * scale_resolution) or allow_upscale: 
        r = width / height 
        height = int(scale_resolution / math.sqrt(r)) 
        width = int(height * r) 
    best_width = ensure_divide(width, patch_size) 
    best_height = ensure_divide(height, patch_size) 
    return (best_width, best_height) 

def get_refine_size(original_size, grid, scale_resolution, patch_size, allow_upscale=False): 
    width, height = original_size 
    grid_x, grid_y = grid 

    refine_width = ensure_divide(width, grid_x) 
    refine_height = ensure_divide(height, grid_y) 
    grid_width = refine_width / grid_x 
    grid_height = refine_height / grid_y 

    best_grid_size = find_best_resize( 
        (grid_width, grid_height), 
        scale_resolution, 
        patch_size, 
        allow_upscale=allow_upscale, 
    ) 

    refine_size = (best_grid_size[0] * grid_x, best_grid_size[1] * grid_y) 

    return refine_size 

def split_to_patches(image, grid): 
    patches = [] 
    width, height = image.size 
    grid_x = int(width / grid[0]) 
    grid_y = int(height / grid[1]) 

    for i in range(0, height, grid_y): 
        images = [] 
        for j in range(0, width, grid_x): 
            box = (j, i, j + grid_x, i + grid_y) 
            patch = image.crop(box) 
            images.append(patch) 
        patches.append(images) 

    return patches 

def get_grid_placeholder(tokenizer, grid, query_num): 
    image_placeholder = tokenizer.im_start + tokenizer.unk_token * query_num + tokenizer.im_end 

    cols = grid[0] 
    rows = grid[1] 
    slices = [] 
    for i in range(rows): 
        lines = [] 
        for j in range(cols): 
            lines.append(image_placeholder) 
        slices.append("".join(lines)) 
    slice_placeholder = tokenizer.slice_start + "\n".join(slices) + tokenizer.slice_end 
    return slice_placeholder 

class OvMiniCPMVModel: 
    def __init__(self, config, vpm, llm, tokenizer) -> None: 
        self.config = config 
        self.vpm = vpm 
        self.llm = llm 
        self.transform = self.init_transform() 
        self.tokenizer = tokenizer 
        self.device = torch.device("cpu") 

    def init_transform(self): 
        return transforms.Compose( 
            [ 
                transforms.ToTensor(), 
                transforms.Normalize(mean=IMAGENET_INCEPTION_MEAN, std=IMAGENET_INCEPTION_STD), 
            ] 
        ) 

    def get_vision_embedding(self, pixel_values): 
        res = [] 
        for pixel_value in pixel_values: 
            h, w = pixel_value.shape[-2:] 
            tgt_size = torch.from_numpy(np.array([math.ceil(h / self.config.patch_size), math.ceil(w / self.config.patch_size)])) 
            vision_embedding = self.vpm([pixel_value.unsqueeze(0), tgt_size])[0] 
            res.append(vision_embedding) 
        return np.vstack(res) 

    def get_vllm_embedding(self, data): 
        if "vision_hidden_states" not in data: 
            pixel_values_list = data["pixel_values"] 
            vision_hidden_states = [] 
            for pixel_values in pixel_values_list: 
                if len(pixel_values) > 0: 
                    vision_hidden_states.append(torch.from_numpy(self.get_vision_embedding(pixel_values))) 
                else
                    : vision_hidden_states.append([]) 
        else: 
            vision_hidden_states = data["vision_hidden_states"] 

        vllm_embedding = torch.from_numpy(self.llm.embed_tokens(data["input_ids"])) * self.llm.config.scale_emb 
        bs = len(data["input_ids"]) 
        for i in range(bs): 
            cur_vs_hs = vision_hidden_states[i] 
            if len(cur_vs_hs) > 0: 
                cur_vllm_emb = vllm_embedding[i] 
                cur_image_bound = data["image_bound"][i] 
                if len(cur_image_bound) > 0: 
                    image_indices = torch.stack([torch.arange(r[0], r[1], dtype=torch.long) for r in cur_image_bound]) 

                    cur_vllm_emb.scatter_( 
                        0, 
                        image_indices.view(-1, 1).repeat(1, 
                        cur_vllm_emb.shape[-1]), cur_vs_hs.view(-1, cur_vs_hs.shape[-1]), 
                    ) 

        return vllm_embedding 

    def forward(self, data, **kwargs): 
        vllm_embedding = self.get_vllm_embedding(data) 
        position_ids = data["position_ids"] 
        if position_ids.dtype != torch.int64: 
            position_ids = position_ids.long() 

        return self.llm(input_ids=None, position_ids=position_ids, inputs_embeds=vllm_embedding, **kwargs) 

    def _convert_to_tensors(self, tokenizer, input_str, max_inp_length: Optional[int] = None): 
        if tokenizer.add_bos_token: 
            input_ids = tokenizer.encode(input_str) 
        else: 
            input_ids = [tokenizer.bos_id] + tokenizer.encode(input_str) 
        if max_inp_length is not None: 
            input_ids = input_ids[:max_inp_length] 
        input_ids = torch.tensor(input_ids, dtype=torch.int32) 

        image_start_tokens = torch.where(input_ids == tokenizer.im_start_id)[0] 
        # 跳过 im_start 
        image_start_tokens += 1 
        image_end_tokens = torch.where(input_ids == tokenizer.im_end_id)[0] 
        valid_image_nums = max(len(image_start_tokens), len(image_end_tokens)) 
        image_bound = torch.hstack( 
            [ 
                image_start_tokens[:valid_image_nums].unsqueeze(-1), 
                image_end_tokens[:valid_image_nums].unsqueeze(-1), 
            ] 
        ) 

        model_input = {} 
        model_input["input_ids"] = input_ids.unsqueeze(0) 
        model_input["image_bound"] = image_bound 

        return model_input 

    def _process_list(self, tokenizer, data_list: List[str], max_inp_length: Optional[int] = None): 
        pad_keys = ["input_ids"] 
        input_tensors = [] 
        for data in data_list: 
            input_tensors.append(self._convert_to_tensors(tokenizer, data, max_inp_length)) 
        padded = {} 
        for key in pad_keys: 
            padded[key] = pad(input_tensors, key, padding_side="left").to(self.device) 
        padded["image_bound"] = [i["image_bound"] for i in input_tensors] 
        return padded 

    def _decode(self, inputs_embeds, tokenizer, **kwargs): 
        output = self.llm.generate(inputs_embeds=inputs_embeds, pad_token_id=0, eos_token_id=tokenizer.eos_token_id, **kwargs) 
        return self._decode_text(output, tokenizer) 

    def _decode_text(self, result_ids, tokenizer): 
        result_text = [] 
        for result in result_ids: 
            result = result[result != 0] 
            if result[0] == tokenizer.bos_id: 
                result = result[1:] 
            if result[-1] == tokenizer.eos_id: 
                result = result[:-1] 
            result_text.append(tokenizer.decode(result).strip()) 
        return result_text 

    def slice_image(self, image): 
        return slice_image( 
            image, 
            self.config.max_slice_nums, 
            self.config.scale_resolution, 
            self.config.patch_size, 
        ) 

    def get_slice_image_placeholder(self, image, tokenizer): 
        image_placeholder = tokenizer.im_start + tokenizer.unk_token * self.config.query_num + tokenizer.im_end 

        slice_images = [] 

        source_image, patches, best_grid = slice_image( 
            image, 
            self.config.max_slice_nums, 
            self.config.scale_resolution, 
            self.config.patch_size, 
        ) 

        slice_images.append(source_image) 
        final_placeholder = image_placeholder 

        if len(patches) > 0: 
            for i in range(len(patches)): 
                for j in range(len(patches[0])): 
                    slice_images.append(patches[i][j]) 

            final_placeholder += get_grid_placeholder(tokenizer, best_grid, self.config.query_num) 
        return slice_images, final_placeholder 

    def generate(self, data_list=None, img_list=None, tokenizer=None, max_inp_length: Optional[int] = None, vision_hidden_states=None, **kwargs): 
        assert data_list is not None 
        bs = len(data_list) 
        if img_list is None: 
            img_list = [[] for i in range(bs)] 
        assert bs == len(img_list) 

        model_inputs = self._process_list(tokenizer, data_list, max_inp_length) 

        if vision_hidden_states is None: 
            pixel_values = [] 
            for i in range(bs): 
                img_inps = [] 
                for img in img_list[i]: 
                    img_inps.append(self.transform(img).to(self.device)) 
                if img_inps: 
                    pixel_values.append(img_inps) 
                else: 
                    pixel_values.append([]) 
            model_inputs["pixel_values"] = pixel_values 
        else: 
            model_inputs["vision_hidden_states"] = vision_hidden_states 

        with torch.inference_mode(): 
            model_inputs["inputs_embeds"] = self.get_vllm_embedding(model_inputs) 

            result = self._decode(model_inputs["inputs_embeds"], tokenizer, **kwargs) 

        return result 

    def chat(self, image, msgs, context, tokenizer, vision_hidden_states=None, max_new_tokens=1024, sampling=True, max_inp_length=2048, **kwargs): 
        if isinstance(msgs, str): 
            msgs = json.loads(msgs) 
        # プロンプトメッセージ 
        prompt = "" 
        for i, msg in enumerate(msgs): 
            role = msg["role"] 
            content = msg["content"] 
            assert role in ["user", "assistant"] 
            if i == 0: 
                if image is None: 
                    images = [] 
                else: 
                    assert role == "user", "The role of first msg should be user" 
                    if self.config.slice_mode: 
                        images, final_placeholder = self.get_slice_image_placeholder(image, tokenizer) content = final_placeholder + "\n" + content 
                    else: 
                        images = [image] 
                        content = tokenizer.im_start + tokenizer.unk_token * self.config.query_num + tokenizer.im_end + "\n" + content 
            prompt += "<用户>" if role == "user" else "<AI>" 
            prompt += content 
        prompt += "<AI>" 
        final_input = prompt 

        if sampling: 
            generation_config = { 
                "top_p": 0.8, 
                "top_k": 100, 
                "temperature": 0.7, 
                "do_sample": True, 
                "repetition_penalty": 1.05, 
                "streamer": None, 
            } 
        else: 
            generation_config = { 
                "num_beams": 3, 
                "repetition_penalty": 1.2, 
                "streamer": None, 
            } 

        generation_config.update((k, kwargs[k]) for k in generation_config.keys() & kwargs.keys()) 

        with torch.inference_mode(): 
            res = self.generate( 
                data_list=[final_input], 
                max_inp_length=max_inp_length, 
                img_list=[images], 
                tokenizer=tokenizer, 
                max_new_tokens=max_new_tokens, 
                vision_hidden_states=vision_hidden_states, 
                **generation_config 
            ) 
        answer = res[0] 
        context = msgs.copy() 
        context.append({"role": "assistant", "content": answer}) 

        return answer, context, generation_config

import gradio as gr 
import traceback 
import re 
from transformers import TextIteratorStreamer 
from threading import Thread 

ERROR_MSG = "Error, please retry" 
model_name = "MiniCPM-V 2.0" 

form_radio = {"choices": ["Beam Search", "Sampling"], "value": "Sampling", "interactive": True, "label": "Decode Type"} 
# ビームフォーム 
num_beams_slider = {"minimum": 0, "maximum": 5, "value": 3, "step": 1, "interactive": True, "label": "Num Beams"} 
repetition_penalty_slider = {"minimum": 0, "maximum": 3, "value": 1.2, "step": 0.01, "interactive": True, "label": "Repetition Penalty"} 
repetition_penalty_slider2 = {"minimum": 0, "maximum": 3, "value": 1.05, "step": 0.01, "interactive": True, "label": "Repetition Penalty"} 
max_new_tokens_slider = {"minimum": 1, "maximum": 4096, "value": 1024, "step": 1, "interactive": True, "label": "Max New Tokens"} 

top_p_slider = {"minimum": 0, "maximum": 1, "value": 0.8, "step": 0.05, "interactive": True, "label": "Top P"} 
top_k_slider = {"minimum": 0, "maximum": 200, "value": 100, "step": 1, "interactive": True, "label": "Top K"} 
temperature_slider = {"minimum": 0, "maximum": 2, "value": 0.7, "step": 0.05, "interactive": True, "label": "Temperature"} 

def create_component(params, comp="Slider"): 
    if comp == "Slider": 
        return gr.Slider( 
            minimum=params["minimum"], 
            maximum=params["maximum"], 
            value=params["value"], 
            step=params["step"], 
            interactive=params["interactive"], 
            label=params["label"], 
        ) 
    elif comp == "Radio": 
        return gr.Radio(choices=params["choices"], value=params["value"], interactive=params["interactive"], label=params["label"]) 
    elif comp == "Button": 
        return gr.Button(value=params["value"], interactive=True) 

def chat(img, msgs, ctx, params=None, vision_hidden_states=None): 
    default_params = {"num_beams": 3, "repetition_penalty": 1.2, "max_new_tokens": 1024} 
    if params is None: 
        params = default_params 
    if img is None: 
        return -1, "Error, invalid image, please upload a new image", None, None 
    try: 
        image = img.convert("RGB") 
        streamer = TextIteratorStreamer(tokenizer, **{"skip_special_tokens": True}) 
        generation_params = {"image": image, "msgs": msgs, "context": None, "tokenizer": tokenizer, "streamer": streamer, **params} 
        thread = Thread(target=model.chat, kwargs=generation_params) 
        thread.start() 

        buffer = "" 

        for res in streamer: 
            res = re.sub(r"(<box>.*</box>)", "", 
            res) res = res.replace("<ref>", "") 
            res = res.replace("</ref>", "") 
            res = res.replace("<box>", "") 
            new_text = res.replace("</box>", "") 
            buffer += new_text 
            yield -1, buffer, None, None 
    except Exception as err: 
        print(err) 
        traceback.print_exc() 
        return -1, ERROR_MSG, None, None 

def upload_img(image, _chatbot, _app_session): 
    image = Image.fromarray(image) 

    _app_session["sts"] = None 
    _app_session["ctx"] = [] 
    _app_session["img"] = image 
    _chatbot.append(("", "Image uploaded successfully, you can talk to me now")) 
    return _chatbot, _app_session 

def respond(_question, _chat_bot, _app_cfg, params_form, num_beams, repetition_penalty, repetition_penalty_2, top_p, top_k, temperature): 
    if _app_cfg.get("ctx", None) is None:
        _chat_bot.append((_question, "Please upload an image to start")) 
        return "", _chat_bot, _app_cfg 

    _context = _app_cfg["ctx"].copy() 
    if _context:
         _context.append({"role": "user", "content": _question}) 
    else:
         _context = [{"role": "user", "content": _question}] 

    if params_form == "Beam Search": 
        params = {"sampling": False, "num_beams": num_beams, "repetition_penalty": repetition_penalty, "max_new_tokens": 896} 
    else: 
        params = { 
            "sampling": True, 
            "top_p": top_p, 
            "top_k": top_k, 
            "temperature": temperature, 
            "repetition_penalty": repetition_penalty_2, 
            "max_new_tokens": 896, 
        } 

    _context.append({"role": "assistant", "content": ""}) 
    _chat_bot.append([_question, ""]) 
    for code, _answer, _, sts in chat(_app_cfg["img"], _context, None, params):
        _context[-1]["content"] = _answer 
        _chat_bot[-1][-1] = _answer 

        if code == 0:
            _app_cfg["ctx"] = _context 
            _app_cfg["sts"] = sts 
        yield "", _chat_bot, _app_cfg 

def regenerate_button_clicked(_question, _chat_bot, _app_cfg, params_form, num_beams, repetition_penalty, repetition_penalty_2, top_p, top_k, temperature): 
    if len(_chat_bot) <= 1:
        _chat_bot.append(("Regenerate", "No question for regeneration.")) 
        return "", _chat_bot, _app_cfg 
    elif _chat_bot[-1][0] == "Regenerate": 
        return "", _chat_bot, _app_cfg 
    else:
        _question = _chat_bot[-1][0] 
        _chat_bot = _chat_bot[:-1] 
        _app_cfg["ctx"] = _app_cfg["ctx"][:-2] 
    for text, _chatbot, _app_cfg in respond( 
        _question, _chat_bot, _app_cfg, params_form, num_beams, repetition_penalty, repetition_penalty_2, top_p, top_k, temperature 
    ): 
        yield text, _chatbot, _app_cfg 

with gr.Blocks() as demo: 
    with gr.Row(): 
        with gr.Column(scale=1, min_width=300): 
            params_form = create_component(form_radio, comp="Radio") 
            with gr.Accordion("Beam Search") as beams_according: 
                num_beams = create_component(num_beams_slider) 
                repetition_penalty = create_component(repetition_penalty_slider) 
            with gr.Accordion("Sampling") as sampling_according: 
                top_p = create_component(top_p_slider) 
                top_k = create_component(top_k_slider) 
                temperature = create_component(temperature_slider) 
                repetition_penalty_2 = create_component(repetition_penalty_slider2) 
            regenerate = create_component({"value": "Regenerate"}, comp="Button") 
        with gr.Column(scale=3, min_width=500): 
            app_session = gr.State({"sts": None, "ctx": None, "img": None}) 
            bt_pic = gr.Image(label="Upload an image to start") 
            chat_bot = gr.Chatbot(label=f"Chat with {model_name}") 
            txt_message = gr.Textbox(label="Input text") 

            regenerate.click( 
                regenerate_button_clicked, 
                [txt_message, chat_bot, app_session, params_form, num_beams, repetition_penalty, repetition_penalty_2, top_p, top_k, temperature], 
                [txt_message, chat_bot, app_session], 
            ) 
            txt_message.submit( 
                respond, 
                [txt_message, chat_bot, app_session, params_form, num_beams, repetition_penalty, repetition_penalty_2, top_p, top_k, temperature], 
            [txt_message, chat_bot, app_session], 
            )
             bt_pic.upload(lambda: None, None, chat_bot, queue=False).then(upload_img, inputs=[bt_pic, chat_bot, app_session], outputs=[chat_bot, app_session]) 

try: 
    demo.launch(debug=False) 
except Exception: 
    demo.launch(debug=False, share=True) 
# リモートで起動する場合は、server_name と server_port を指定 
# demo.launch(server_name='your server name', server_port='server port in int') 
# 詳細はドキュメントをご覧ください: https://gradio.app/docs/