Views
No views yet
1from PIL import Image
2import argparse
3import torch
4import os
5
6from llava.constants import IMAGE_TOKEN_INDEX, DEFAULT_IMAGE_TOKEN
7from llava.conversation import conv_templates
8from llava.model.builder import load_pretrained_model
9from llava.mm_utils import tokenizer_image_token, process_images, get_model_name_from_path
10from llava.utils import disable_torch_init
11
12@torch.inference_mode()
13def run(model_path, image_paths, prompt, num_retrievals=1):
14 '''
15 Executes MIRAGE with specified inputs to generate descriptive text based on the provided images.
16
17 Args:
18 model_path (str): Path to the MIRAGE model, e.g., 'tsunghanwu/mirage-llama3.1-8.3B'
19 image_paths (list): List of paths to image files, e.g., images in 'assets/example'
20 prompt (str): Text prompt for image description, e.g., 'Here are a set of random images in my photo album.
21 If you can find a cat, tell me what's the cat doing and what's its color.'
22 num_retrievals (int): Maximum number of images to retrieve and pass to the LMM
23
24 Returns:
25 output_text (str): Descriptive text generated by the LMM
26 output_ret (list): List of images retrieved by the model
27 '''
28 # Load the model and prepare the environment
29 model_name = get_model_name_from_path(model_path)
30 disable_torch_init()
31 model_name = os.path.expanduser(model_name)
32 tokenizer, model, image_processor, _ = \
33 load_pretrained_model(model_path=model_path, model_base=None, model_name=model_name, device="cuda")
34 model.eval_mode = True
35
36 # Process the images
37 clip_images = []
38 for image_path in image_paths:
39 image = Image.open(image_path). convert("RGB")
40 image_tensor = process_images([image], image_processor, model.config)[0]
41 image_tensor = image_tensor.to(dtype=torch.float16)
42 clip_images.append(image_tensor)
43
44 # Prepare text input and interaction
45 qformer_text_input = tokenizer(prompt, return_tensors='pt')["input_ids"].to(model.device)
46 N = len(clip_images)
47 img_str = DEFAULT_IMAGE_TOKEN * N + "\n"
48 inp = img_str + prompt
49 conv.append_message(conv.roles[0], inp)
50 conv.append_message(conv.roles[1], None)
51 prompt = conv.get_prompt()
52
53 # Generate model output
54 input_ids = tokenizer_image_token(prompt, tokenizer, IMAGE_TOKEN_INDEX, return_tensors='pt').unsqueeze(0).to(model.device)
55 tokenizer.pad_token_id = 128002
56 batch_clip_images = [torch.stack(clip_images).to(model.device)]
57
58 output_ret, output_ids = model.generate(
59 input_ids,
60 pad_token_id=tokenizer.pad_token_id,
61 clip_images=batch_clip_images,
62 qformer_text_input=qformer_text_input,
63 relevance=None,
64 num_retrieval=num_retrievals,
65 do_sample=False,
66 max_new_tokens=512,
67 use_cache=True)
68
69 # Process output
70 output_text = tokenizer.batch_decode(output_ids, skip_special_tokens=True)[0].strip()
71 if not isinstance(output_ret[0], list):
72 output_ret[0] = output_ret[0].tolist()
73 return output_text, output_ret[0]