neuralcomputation/batik

1import argparse2import torch3from transformers import AutoTokenizer, AutoModelForCausalLM, AutoConfig, BitsAndBytesConfig4import numpy as np5 6from huggingface_hub import whoami7 8import llava9from llava.constants import IMAGE_TOKEN_INDEX, DEFAULT_IMAGE_TOKEN, DEFAULT_IM_START_TOKEN, DEFAULT_IM_END_TOKEN, DEFAULT_IMAGE_PATCH_TOKEN10from llava.conversation import conv_templates, SeparatorStyle11from llava.model.builder import load_pretrained_model12from llava.utils import disable_torch_init13from llava.mm_utils import process_images, tokenizer_image_token, get_model_name_from_path14 15from PIL import Image16 17import requests18from PIL import Image19from io import BytesIO20from transformers import TextStreamer21from tqdm import tqdm22 23import warnings24warnings.filterwarnings('ignore')25 26REPO_NAME = 'ncoria/llava-lora-vicuna-clip-5-epochs-merge'27 28def load_image(image_file):29    if image_file.startswith('http://') or image_file.startswith('https://'):30        response = requests.get(image_file)31        image = Image.open(BytesIO(response.content)).convert('RGB')32    else:33        image = Image.open(image_file).convert('RGB')34    return image35 36def load_llava_checkpoint(model_path: str):37    model_name = get_model_name_from_path(model_path)38    return load_pretrained_model(model_path, None, model_name, load_4bit=True, device="cuda")39 40def load_llava_checkpoint_hf(model_path):41    kwargs = {"device_map": "auto"}42    kwargs['load_in_4bit'] = True43    kwargs['quantization_config'] = BitsAndBytesConfig(44        load_in_4bit=True,45        bnb_4bit_compute_dtype=torch.float16,46        bnb_4bit_use_double_quant=True,47        bnb_4bit_quant_type='nf4'48    )49    tokenizer = AutoTokenizer.from_pretrained(model_path, use_fast=False)50    model = AutoModelForCausalLM.from_pretrained(model_path, low_cpu_mem_usage=True, **kwargs)51    mm_use_im_start_end = getattr(model.config, "mm_use_im_start_end", False)52    mm_use_im_patch_token = getattr(model.config, "mm_use_im_patch_token", True)53    if mm_use_im_patch_token:54        tokenizer.add_tokens([DEFAULT_IMAGE_PATCH_TOKEN], special_tokens=True)55    if mm_use_im_start_end:56        tokenizer.add_tokens([DEFAULT_IM_START_TOKEN, DEFAULT_IM_END_TOKEN], special_tokens=True)57    model.resize_token_embeddings(len(tokenizer))58 59    vision_tower = model.get_vision_tower()60    if not vision_tower.is_loaded:61        vision_tower.load_model(device_map="auto")62    image_processor = vision_tower.image_processor63    return tokenizer, model, image_processor64 65def get_llava_response(user_prompts: list[str],66                       images: list,67                       sys_prompt: str,68                       tokenizer,69                       model,70                       image_processor,71                       model_path = REPO_NAME,72                       stream_output = True):73    """74    This function returns the response from the given model. It creates a one turn conversation in which75    the only content is a system prompt and the given user message applied to each image.76 77    Parameters:78    ----------79    user_prompt : str80        The prompt sent by the user.81    images : str82        List of images from file.83    sys_prompt : str84        The prompt that sets the tone for the conversation.85    model_path : str86        The path to the merged checkpoint or base model.87 88    Returns:89    --------90    """91    # set up and load model92    model_name = get_model_name_from_path(model_path)93    temperature = 0.2 # default94    max_new_tokens = 512 # default95 96    # determine conversation type97    if "llama-2" in model_name.lower():98        conv_mode = "llava_llama_2"99    elif "mistral" in model_name.lower():100        conv_mode = "mistral_instruct"101    elif "v1.6-34b" in model_name.lower():102        conv_mode = "chatml_direct"103    elif "v1" in model_name.lower():104        conv_mode = "llava_v1"105    elif "mpt" in model_name.lower():106        conv_mode = "mpt"107    else:108        conv_mode = "llava_v0"109 110    # run clean conversation for each image111    llm_outputs = []112    for i, img in tqdm(enumerate(images)):113        # set up clean conversation114        conv = conv_templates[conv_mode].copy()115        if "mpt" in model_name.lower():116            roles = ('user', 'assistant')117        else:118            roles = conv.roles119    120        conv.system = sys_prompt121    122        # load image123        # image = load_image("../images/mouse.png") # previous method124        if isinstance(img, np.ndarray) and len(img.shape) == 2:125            img = Image.fromarray(img, 'L')126        elif isinstance(img, np.ndarray):127            img = Image.fromarray(img)128 129        image = img.convert('RGB')130        image_size = image.size131 132        # NOTE: image is simply PIL Image (.convert('RGB')), no need for temp files!133        134        # Similar operation in model_worker.py135        image_tensor = process_images([image], image_processor, model.config)136        if type(image_tensor) is list:137            image_tensor = [image.to(model.device, dtype=torch.float16) for image in image_tensor]138        else:139            image_tensor = image_tensor.to(model.device, dtype=torch.float16)140 141        # execute conversation142        inp = user_prompts[i]143        if image is not None:144            # first message145            if model.config.mm_use_im_start_end:146                inp = DEFAULT_IM_START_TOKEN + DEFAULT_IMAGE_TOKEN + DEFAULT_IM_END_TOKEN + '\n' + inp147            else:148                inp = DEFAULT_IMAGE_TOKEN + '\n' + inp149            image = None150        conv.append_message(conv.roles[0], inp)151        conv.append_message(conv.roles[1], None)152        prompt = conv.get_prompt()153        input_ids = tokenizer_image_token(prompt,154                                          tokenizer,155                                          IMAGE_TOKEN_INDEX,156                                          return_tensors='pt').unsqueeze(0).to(model.device)157        stop_str = conv.sep if conv.sep_style != SeparatorStyle.TWO else conv.sep2158        keywords = [stop_str]159        if stream_output:160            streamer = TextStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True)161        else:162            streamer = None163        164        with torch.inference_mode():165            output_ids = model.generate(166                input_ids,167                images=image_tensor,168                image_sizes=[image_size],169                do_sample=True if temperature > 0 else False,170                temperature=temperature,171                max_new_tokens=max_new_tokens,172                streamer=streamer,173                use_cache=True)174 175        outputs = tokenizer.decode(output_ids[0]).strip()176        llm_outputs.append(outputs)177    return llm_outputs178        179 180        181    182    183 184    185