app.py

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|>', '<pad>', '<question>', '<answer>', '<explanation>']
    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(
        ['<question>', '<answer>', '<explanation>', '<labels>']
    )

    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         
)