import torch import torch.nn as nn import torch.nn.functional as F import torchvision.transforms as transforms from transformers import GPT2LMHeadModel, GPT2Tokenizer import clip from PIL import Image import re import numpy as np import cv2 import gradio as gr # ====================== # Device # ====================== device = torch.device("cuda" if torch.cuda.is_available() else "cpu") # ====================== # Utils # ====================== def proc_ques(ques): return re.sub(r"([.,'!?\"()*#:;])", '', ques.lower()).replace('-', ' ').replace('/', ' ') def change_requires_grad(model, req_grad): for p in model.parameters(): p.requires_grad = req_grad def load_checkpoint(ckpt_path, epoch): model_name = f"vqae_model_{epoch}" tokenizer_name = "vqae_gpt2_tokenizer_0" tokenizer = GPT2Tokenizer.from_pretrained(ckpt_path + tokenizer_name) model = GPT2LMHeadModel.from_pretrained(ckpt_path + model_name).to(device) return tokenizer, model # ====================== # Image Encoder (CLIP) # ====================== class ImageEncoder(nn.Module): def __init__(self): super().__init__() self.encoder, _ = clip.load("ViT-B/16", device=device) def forward(self, x): with torch.no_grad(): x = x.type(self.encoder.visual.conv1.weight.dtype) x = self.encoder.visual.conv1(x) x = x.reshape(x.shape[0], x.shape[1], -1) x = x.permute(0, 2, 1) x = torch.cat([ self.encoder.visual.class_embedding.to(x.dtype) + torch.zeros(x.shape[0], 1, x.shape[-1], device=x.device), x ], dim=1) x = x + self.encoder.visual.positional_embedding.to(x.dtype) x = self.encoder.visual.ln_pre(x) x = x.permute(1, 0, 2) x = self.encoder.visual.transformer(x) x = x.permute(1, 0, 2) x = self.encoder.visual.ln_post(x[:, 1:]) return x.float() # ====================== # Sampling # ====================== def top_filtering(logits, top_k=0, top_p=0.9, filter_value=-float("Inf")): if top_k > 0: indices_to_remove = logits < torch.topk(logits, top_k)[0][..., -1, None] logits[indices_to_remove] = filter_value if top_p > 0.0: sorted_logits, sorted_indices = torch.sort(logits, descending=True) cumulative_probs = torch.cumsum(F.softmax(sorted_logits, dim=-1), dim=-1) sorted_indices_to_remove = cumulative_probs > top_p sorted_indices_to_remove[..., 1:] = sorted_indices_to_remove[..., :-1] sorted_indices_to_remove[..., 0] = 0 indices_to_remove = sorted_indices[sorted_indices_to_remove] logits[indices_to_remove] = filter_value return logits def sample_sequences(img, model, input_ids, segment_ids, tokenizer): SPECIAL_TOKENS = ['<|endoftext|>', '', '', '', ''] special_ids = tokenizer.convert_tokens_to_ids(SPECIAL_TOKENS) because_token = tokenizer.convert_tokens_to_ids('Ġbecause') img_emb = image_encoder(img) output_tokens = [] always_exp = False max_len = 20 with torch.no_grad(): for _ in range(max_len): out = model( input_ids=input_ids, token_type_ids=segment_ids, encoder_hidden_states=img_emb, use_cache=False, output_attentions=True, return_dict=True ) logits = out.logits[0, -1] / temperature logits = top_filtering(logits, top_k, top_p) probs = F.softmax(logits, dim=-1) next_token = torch.argmax(probs).unsqueeze(0) if next_token.item() in special_ids: break if not always_exp: if next_token.item() == because_token: seg = special_ids[-1] always_exp = True else: seg = special_ids[-2] else: seg = special_ids[-1] output_tokens.append(next_token.item()) input_ids = torch.cat([input_ids, next_token.unsqueeze(0)], dim=1) segment_ids = torch.cat([segment_ids, torch.tensor([[seg]], device=device)], dim=1) decoded = tokenizer.decode(output_tokens, skip_special_tokens=True).strip() return decoded, out.cross_attentions # ====================== # Load models # ====================== img_size = 224 ckpt_path = "ckpts/" load_from_epoch = 29 top_k = 0 top_p = 0.9 temperature = 1 image_encoder = ImageEncoder().to(device) change_requires_grad(image_encoder, False) tokenizer, model = load_checkpoint(ckpt_path, load_from_epoch) model.eval() img_transform = transforms.Compose([ transforms.Resize((img_size, img_size)), transforms.ToTensor(), transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]) ]) # ====================== # Input builder # ====================== def get_inputs(text, tokenizer): q_id, a_id, e_id, _ = tokenizer.convert_tokens_to_ids( ['', '', '', ''] ) tokens = tokenizer.tokenize(text) segs = [q_id] * len(tokens) ans = [tokenizer.bos_token] + tokenizer.tokenize(" the answer is") tokens += ans segs += [a_id] * len(ans) return ( torch.tensor([tokenizer.convert_tokens_to_ids(tokens)], device=device), torch.tensor([segs], device=device) ) # ====================== # Inference # ====================== def inference(raw_image, question): oimg = raw_image.convert("RGB").resize((224, 224)) img = img_transform(oimg).unsqueeze(0).to(device) text = proc_ques(question) input_ids, segment_ids = get_inputs(text, tokenizer) seq, xa_maps = sample_sequences(img, model, input_ids, segment_ids, tokenizer) # Attention map q_len = input_ids.shape[1] attn = xa_maps[-1].mean(1)[0, q_len:] mask = attn[0].reshape(14, 14).cpu().numpy() mask = cv2.resize(mask / mask.max(), oimg.size)[..., None] vis = (mask * np.array(oimg)).astype("uint8") parts = seq.split("because") answer = parts[0].strip() explanation = "because " + parts[-1].strip() if len(parts) > 1 else "" return answer, explanation, Image.fromarray(vis) # ====================== # Gradio UI # ====================== inputs = [ gr.Image(type="pil", label="Load an image"), gr.Textbox(label="Ask a question") ] outputs = [ gr.Textbox(label="Answer"), gr.Textbox(label="Textual Explanation"), gr.Image(type="pil", label="Visual Explanation") ] title = "VQA – Visual Question Answering" gr.Interface( fn=inference, inputs=inputs, outputs=outputs, title=title, api_name=False ).launch( share=True )