import re
import torch
import clip
from PIL import Image


def get_activation(name, activation: dict):
    def hook(model, input, output):
        if isinstance(output, torch.Tensor):
            activation[name] = output.detach()
    return hook

def attach_intermediate_layer(model, activation: dict):
    for name, layer in model.named_modules():
        # if "" in name:
        if re.match(r"visual.transformer.resblocks.[0-9]*$", name) or re.match(r"transformer.resblocks.[0-9]*$", name):
            layer.register_forward_hook(get_activation(name, activation))


