q-future/Co-Instruct

1import argparse2import torch3 4from mplug_owl2.constants import IMAGE_TOKEN_INDEX, DEFAULT_IMAGE_TOKEN5from mplug_owl2.conversation import conv_templates, SeparatorStyle6from mplug_owl2.model.builder import load_pretrained_model7from mplug_owl2.mm_utils import process_images, tokenizer_image_token, get_model_name_from_path, KeywordsStoppingCriteria8 9from PIL import Image10 11import requests12from PIL import Image13from io import BytesIO14from transformers import TextStreamer15 16 17def disable_torch_init():18    """19    Disable the redundant torch default initialization to accelerate model creation.20    """21    import torch22    setattr(torch.nn.Linear, "reset_parameters", lambda self: None)23    setattr(torch.nn.LayerNorm, "reset_parameters", lambda self: None)24 25 26def load_image(image_file):27    if image_file.startswith('http://') or image_file.startswith('https://'):28        response = requests.get(image_file)29        image = Image.open(BytesIO(response.content)).convert('RGB')30    else:31        image = Image.open(image_file).convert('RGB')32    return image33 34 35def main(args):36    # Model37    disable_torch_init()38 39    model_name = get_model_name_from_path(args.model_path)40    tokenizer, model, image_processor, context_len = load_pretrained_model(args.model_path, args.model_base, model_name, args.load_8bit, args.load_4bit, device=args.device)41 42    conv_mode = "mplug_owl2"43 44    if args.conv_mode is not None and conv_mode != args.conv_mode:45        print('[WARNING] the auto inferred conversation mode is {}, while `--conv-mode` is {}, using {}'.format(conv_mode, args.conv_mode, args.conv_mode))46    else:47        args.conv_mode = conv_mode48 49    conv = conv_templates[args.conv_mode].copy()50    roles = conv.roles51 52    image = load_image(args.image_file)53    # Similar operation in model_worker.py54    image_tensor = process_images([image], image_processor, args)55    if type(image_tensor) is list:56        image_tensor = [image.to(model.device, dtype=torch.float16) for image in image_tensor]57    else:58        image_tensor = image_tensor.to(model.device, dtype=torch.float16)59 60    while True:61        try:62            inp = input(f"{roles[0]}: ")63        except EOFError:64            inp = ""65        if not inp:66            print("exit...")67            break68 69        print(f"{roles[1]}: ", end="")70 71        if image is not None:72            # first message73            inp = DEFAULT_IMAGE_TOKEN + inp74            conv.append_message(conv.roles[0], inp)75            image = None76        else:77            # later messages78            conv.append_message(conv.roles[0], inp)79        conv.append_message(conv.roles[1], None)80        prompt = conv.get_prompt()81 82        input_ids = tokenizer_image_token(prompt, tokenizer, IMAGE_TOKEN_INDEX, return_tensors='pt').unsqueeze(0).to(model.device)83        stop_str = conv.sep if conv.sep_style not in [SeparatorStyle.TWO, SeparatorStyle.TWO_NO_SYS] else conv.sep284        keywords = [stop_str]85        stopping_criteria = KeywordsStoppingCriteria(keywords, tokenizer, input_ids)86        streamer = TextStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True)87 88        with torch.inference_mode():89            output_ids = model.generate(90                input_ids,91                images=image_tensor,92                do_sample=True,93                temperature=args.temperature,94                max_new_tokens=args.max_new_tokens,95                streamer=streamer,96                use_cache=True,97                stopping_criteria=[stopping_criteria])98 99        outputs = tokenizer.decode(output_ids[0, input_ids.shape[1]:]).strip()100        conv.messages[-1][-1] = outputs101 102        if args.debug:103            print("\n", {"prompt": prompt, "outputs": outputs}, "\n")104 105 106if __name__ == "__main__":107    parser = argparse.ArgumentParser()108    parser.add_argument("--model-path", type=str, default="facebook/opt-350m")109    parser.add_argument("--model-base", type=str, default=None)110    parser.add_argument("--image-file", type=str, required=True)111    parser.add_argument("--device", type=str, default="cuda")112    parser.add_argument("--conv-mode", type=str, default=None)113    parser.add_argument("--temperature", type=float, default=0.2)114    parser.add_argument("--max-new-tokens", type=int, default=512)115    parser.add_argument("--load-8bit", action="store_true")116    parser.add_argument("--load-4bit", action="store_true")117    parser.add_argument("--debug", action="store_true")118    parser.add_argument("--image-aspect-ratio", type=str, default='pad')119    args = parser.parse_args()120    main(args)