(args)
| 16 | |
| 17 | |
| 18 | def predict(args): |
| 19 | # Remove generation config from model folder |
| 20 | # to read generation parameters from args |
| 21 | model_path = os.path.expanduser(args.model_path) |
| 22 | generation_config = None |
| 23 | if os.path.exists(os.path.join(model_path, 'generation_config.json')): |
| 24 | generation_config = os.path.join(model_path, '.generation_config.json') |
| 25 | os.rename(os.path.join(model_path, 'generation_config.json'), |
| 26 | generation_config) |
| 27 | |
| 28 | # Load model |
| 29 | disable_torch_init() |
| 30 | model_name = get_model_name_from_path(model_path) |
| 31 | tokenizer, model, image_processor, context_len = load_pretrained_model(model_path, args.model_base, model_name, device="mps") |
| 32 | |
| 33 | # Construct prompt |
| 34 | qs = args.prompt |
| 35 | if model.config.mm_use_im_start_end: |
| 36 | qs = DEFAULT_IM_START_TOKEN + DEFAULT_IMAGE_TOKEN + DEFAULT_IM_END_TOKEN + '\n' + qs |
| 37 | else: |
| 38 | qs = DEFAULT_IMAGE_TOKEN + '\n' + qs |
| 39 | conv = conv_templates[args.conv_mode].copy() |
| 40 | conv.append_message(conv.roles[0], qs) |
| 41 | conv.append_message(conv.roles[1], None) |
| 42 | prompt = conv.get_prompt() |
| 43 | |
| 44 | # Set the pad token id for generation |
| 45 | model.generation_config.pad_token_id = tokenizer.pad_token_id |
| 46 | |
| 47 | # Tokenize prompt |
| 48 | input_ids = tokenizer_image_token(prompt, tokenizer, IMAGE_TOKEN_INDEX, return_tensors='pt').unsqueeze(0).to(torch.device("mps")) |
| 49 | |
| 50 | # Load and preprocess image |
| 51 | image = Image.open(args.image_file).convert('RGB') |
| 52 | image_tensor = process_images([image], image_processor, model.config)[0] |
| 53 | |
| 54 | # Run inference |
| 55 | with torch.inference_mode(): |
| 56 | output_ids = model.generate( |
| 57 | input_ids, |
| 58 | images=image_tensor.unsqueeze(0).half(), |
| 59 | image_sizes=[image.size], |
| 60 | do_sample=True if args.temperature > 0 else False, |
| 61 | temperature=args.temperature, |
| 62 | top_p=args.top_p, |
| 63 | num_beams=args.num_beams, |
| 64 | max_new_tokens=256, |
| 65 | use_cache=True) |
| 66 | |
| 67 | outputs = tokenizer.batch_decode(output_ids, skip_special_tokens=True)[0].strip() |
| 68 | print(outputs) |
| 69 | |
| 70 | # Restore generation config |
| 71 | if generation_config is not None: |
| 72 | os.rename(generation_config, os.path.join(model_path, 'generation_config.json')) |
| 73 | |
| 74 | |
| 75 | if __name__ == "__main__": |
no test coverage detected