File size: 3,429 Bytes
91c1e9e | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 | """Run PRIMA on ordered image groups and text prompts."""
import argparse
import json
from pathlib import Path
from tools.config import parse_with_config, validate_paths
def parse_args():
from tools.model_loader import DEFAULT_HUB_MODEL, add_model_arguments
parser = argparse.ArgumentParser(description=__doc__)
add_model_arguments(parser, merged_checkpoint=DEFAULT_HUB_MODEL)
parser.add_argument('--image_paths', nargs='+', help='One image group for CLI inference.')
parser.add_argument('--prompts', nargs='+', help='Prompts paired with image groups in the config.')
parser.add_argument('--output_save_path', default='outputs/inference')
parser.add_argument('--max_new_tokens', type=int, default=1024)
return parse_with_config(parser)
def main():
args = parse_args()
validate_paths(args)
if not args.image_paths or not args.prompts:
raise ValueError('Provide image_paths and prompts in the config or CLI.')
groups = args.image_paths
if all(isinstance(item, str) for item in groups):
groups = [groups]
if len(groups) != len(args.prompts):
raise ValueError('Each image group must have exactly one prompt.')
if any(not isinstance(group, list) or not 1 <= len(group) <= args.max_num_images for group in groups):
raise ValueError('Each image group must be a nonempty list within max_num_images.')
for group in groups:
for path in group:
if not Path(path).is_file():
raise ValueError(f'Image does not exist: {path}')
import numpy as np
import torch
from model.llava import conversation
from tools.chat import Chat
from tools.model_loader import load_bundle
from tools.outputs import image_mask_slots
model, tokenizer, qformer_tokenizer = load_bundle(args)
chat = Chat(model, tokenizer, qformer_tokenizer, args, device=f'cuda:{args.local_rank}')
output = Path(args.output_save_path)
output.mkdir(parents=True, exist_ok=True)
for index, (images, question) in enumerate(zip(groups, args.prompts)):
turn = output / f'turn{index}'
turn.mkdir(exist_ok=True)
conv = conversation.conv_templates[args.conv_type].copy()
encoded = []
for image in images:
chat.upload_img(image, conv, encoded)
chat.ask(question, conv, num_images=len(images))
with torch.inference_mode():
answer, masks, _ = chat.answer(conv, encoded, question,
max_new_tokens=args.max_new_tokens,
max_length=args.model_max_length + args.max_new_tokens)
mask_paths = []
for image_index, predicted in enumerate(image_mask_slots(masks, len(images))):
paths = []
if predicted is not None:
for mask_index, mask in enumerate(predicted):
name = f'image{image_index + 1}_mask{mask_index + 1}.npy'
np.save(turn / name, (mask > 0).cpu().numpy())
paths.append(name)
mask_paths.append(paths)
record = {'images': images, 'prompt': question, 'response': answer, 'masks': mask_paths}
(turn / 'dialogue.json').write_text(json.dumps(record, indent=2) + '\n')
print(answer)
if __name__ == '__main__':
try:
main()
except ValueError as error:
raise SystemExit(str(error))
|