Views
No views yet
1git clone https://github.com/ByungKwanLee/TroL
2bash installdemo.py)1import torch
2from config import *
3from PIL import Image
4from utils.utils import *
5import torch.nn.functional as F
6from trol.load_trol import load_trol
7from torchvision.transforms.functional import pil_to_tensor
8
9# model selection
10link = "TroL-7B" # [Select One] 'TroL-1.8B' | 'TroL-3.8B' | 'TroL-7B'
11
12# User prompt
13prompt_type="with_image" # Select one option "text_only", "with_image"
14img_path='figures/demo.png'
15question="What is the troll doing? Provide the detail in the image and imagine what the event happens."
16
17# loading model
18model, tokenizer = load_trol(link=link)
19
20# cpu -> gpu
21for param in model.parameters():
22 if not param.is_cuda:
23 param.data = param.to('cuda:0')
24
25# prompt type -> input prompt
26image_token_number = None
27if prompt_type == 'with_image':
28 # Image Load
29 image = pil_to_tensor(Image.open(img_path).convert("RGB"))
30 if not "3.8B" in link:
31 image_token_number = 1225
32 image = F.interpolate(image.unsqueeze(0), size=(490, 490), mode='bicubic').squeeze(0)
33 inputs = [{'image': image, 'question': question}]
34elif prompt_type=='text_only':
35 inputs = [{'question': question}]
36
37# Generate
38with torch.inference_mode():
39 _inputs = model.eval_process(inputs=inputs,
40 data='demo',
41 tokenizer=tokenizer,
42 device='cuda:0',
43 img_token_number=image_token_number)
44 generate_ids = model.generate(**_inputs, max_new_tokens=256, use_cache=True)
45 response = output_filtering(tokenizer.batch_decode(generate_ids, skip_special_tokens=False)[0], model)
46print(response)