(args)
| 25 | |
| 26 | |
| 27 | def main(args): |
| 28 | # Model |
| 29 | disable_torch_init() |
| 30 | |
| 31 | model_name = get_model_name_from_path(args.model_path) |
| 32 | 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) |
| 33 | |
| 34 | if "llama-2" in model_name.lower(): |
| 35 | conv_mode = "llava_llama_2" |
| 36 | elif "mistral" in model_name.lower(): |
| 37 | conv_mode = "mistral_instruct" |
| 38 | elif "v1.6-34b" in model_name.lower(): |
| 39 | conv_mode = "chatml_direct" |
| 40 | elif "v1" in model_name.lower(): |
| 41 | conv_mode = "llava_v1" |
| 42 | elif "mpt" in model_name.lower(): |
| 43 | conv_mode = "mpt" |
| 44 | else: |
| 45 | conv_mode = "llava_v0" |
| 46 | |
| 47 | if args.conv_mode is not None and conv_mode != args.conv_mode: |
| 48 | print('[WARNING] the auto inferred conversation mode is {}, while `--conv-mode` is {}, using {}'.format(conv_mode, args.conv_mode, args.conv_mode)) |
| 49 | else: |
| 50 | args.conv_mode = conv_mode |
| 51 | |
| 52 | conv = conv_templates[args.conv_mode].copy() |
| 53 | if "mpt" in model_name.lower(): |
| 54 | roles = ('user', 'assistant') |
| 55 | else: |
| 56 | roles = conv.roles |
| 57 | |
| 58 | image = load_image(args.image_file) |
| 59 | image_size = image.size |
| 60 | # Similar operation in model_worker.py |
| 61 | image_tensor = process_images([image], image_processor, model.config) |
| 62 | if type(image_tensor) is list: |
| 63 | image_tensor = [image.to(model.device, dtype=torch.float16) for image in image_tensor] |
| 64 | else: |
| 65 | image_tensor = image_tensor.to(model.device, dtype=torch.float16) |
| 66 | |
| 67 | while True: |
| 68 | try: |
| 69 | inp = input(f"{roles[0]}: ") |
| 70 | except EOFError: |
| 71 | inp = "" |
| 72 | if not inp: |
| 73 | print("exit...") |
| 74 | break |
| 75 | |
| 76 | print(f"{roles[1]}: ", end="") |
| 77 | |
| 78 | if image is not None: |
| 79 | # first message |
| 80 | if model.config.mm_use_im_start_end: |
| 81 | inp = DEFAULT_IM_START_TOKEN + DEFAULT_IMAGE_TOKEN + DEFAULT_IM_END_TOKEN + '\n' + inp |
| 82 | else: |
| 83 | inp = DEFAULT_IMAGE_TOKEN + '\n' + inp |
| 84 | image = None |
no test coverage detected