Views
No views yet
1bash install
2pip install -r requirements.txt1import time
2import torch
3from config import *
4from PIL import Image
5from utils.utils import *
6import torch.nn.functional as F
7from meteor.load_mmamba import load_mmamba
8from meteor.load_meteor import load_meteor
9from torchvision.transforms.functional import pil_to_tensor
10
11# User prompt
12prompt_type='with_image' # text_only / with_image
13img_path='figures/demo.png'
14question='Provide the detail of the image'
15
16# loading meteor model
17mmamba = load_mmamba('BK-Lee/Meteor-Mamba').cuda()
18meteor, tok_meteor = load_meteor('BK-Lee/Meteor-MLM', bits=4)
19
20# freeze model
21freeze_model(mmamba)
22freeze_model(meteor)
23
24# Device
25device = torch.cuda.current_device()
26
27# prompt type -> input prompt
28image_token_number = int((490/14)**2)
29if prompt_type == 'with_image':
30 # Image Load
31 image = F.interpolate(pil_to_tensor(Image.open(img_path).convert("RGB")).unsqueeze(0), size=(490, 490), mode='bicubic').squeeze(0)
32 inputs = [{'image': image, 'question': question}]
33elif prompt_type=='text_only':
34 inputs = [{'question': question}]
35
36# Generate
37with torch.inference_mode():
38
39 # Meteor Mamba
40 mmamba_inputs = mmamba.eval_process(inputs=inputs, tokenizer=tok_meteor, device=device, img_token_number=image_token_number)
41 if 'image' in mmamba_inputs.keys():
42 clip_features = meteor.clip_features(mmamba_inputs['image'])
43 mmamba_inputs.update({"image_features": clip_features})
44 mmamba_outputs = mmamba(**mmamba_inputs)
45
46 # Meteor
47 meteor_inputs = meteor.eval_process(inputs=inputs, data='demo', tokenizer=tok_meteor, device=device, img_token_number=image_token_number)
48 if 'image' in mmamba_inputs.keys():
49 meteor_inputs.update({"image_features": clip_features})
50 meteor_inputs.update({"tor_features": mmamba_outputs.tor_features})
51
52 # Generation
53 generate_ids = meteor.generate(**meteor_inputs, do_sample=True, max_new_tokens=128, top_p=0.95, temperature=0.9, use_cache=True)
54
55# Text decoding
56decoded_text = tok_meteor.batch_decode(generate_ids, skip_special_tokens=True)[0].split('assistant\n')[-1].split('[U')[0].strip()
57print(decoded_text)
58
59# Paper arxiv.org/abs/2405.15574
60