import torch
from PIL import Image
from transformers import AutoProcessor, openclipVisionModel
import matplotlib.pyplot as plt
import numpy as np

# Load your model
processor = AutoProcessor.from_pretrained("google/openclip-so400m-patch14-384")
model = openclipVisionModel.from_pretrained("google/openclip-so400m-patch14-384")
model.load_state_dict(torch.load("your_finetuned_openclip.pt"))

# Set up to get attention maps
model.eval()
model.vision_model.encoder.config.output_attentions = True

# Load image
img = Image.open("car_product_photo.jpg")
inputs = processor(images=img, return_tensors="pt")

# Forward pass WITH attention
with torch.no_grad():
    outputs = model.vision_model(**inputs, output_attentions=True)

    # Get attention weights from last layer
    # Shape: (batch, num_heads, seq_len, seq_len)
    attention_weights = outputs.attentions[-1]

    # Average across heads and batch
    attention_map = attention_weights[0].mean(dim=0)  # (seq_len, seq_len)

    # Reshape back to image space
    # openclip uses patch embedding, so we need to reshape
    H, W = 384, 384  # input image size
    patch_size = 14
    num_patches_h = H // patch_size  # 27
    num_patches_w = W // patch_size  # 27

    # Take the attention to the [CLS] token (first token)
    cls_attention = attention_map[0, 1:]  # Ignore self-attention to CLS
    cls_attention = cls_attention.reshape(num_patches_h, num_patches_w)

    # Upsample to image size
    cls_attention_upsampled = torch.nn.functional.interpolate(
        cls_attention.unsqueeze(0).unsqueeze(0),
        size=(H, W),
        mode='bilinear'
    ).squeeze()

    # Normalize to 0-1
    cls_attention_upsampled = (cls_attention_upsampled - cls_attention_upsampled.min()) / \
                              (cls_attention_upsampled.max() - cls_attention_upsampled.min())

    return cls_attention_upsampled.numpy()

# Visualize
heatmap = get_attention_heatmap(img)
plt.imshow(img)
plt.imshow(heatmap, alpha=0.4, cmap='jet')
plt.show()