| import numpy as np | |
| import gradio as gr | |
| from huggingface_hub import hf_hub_download | |
| import omegaconf | |
| from hydra import utils | |
| import os | |
| import torch | |
| import matplotlib.pyplot as plt | |
| from attn_helper import VITAttentionGradRollout, overlay_attn | |
| import vc_models | |
| import torchvision | |
| HF_TOKEN = os.environ['HF_ACC_TOKEN'] | |
| eai_filepath = vc_models.__file__.split('src')[0] | |
| MODEL_DIR=os.path.join(os.path.dirname(eai_filepath),'model_ckpts') | |
| if not os.path.isdir(MODEL_DIR): | |
| os.mkdir(MODEL_DIR) | |
| FILENAME = "config.yaml" | |
| BASE_MODEL_TUPLE = None | |
| LARGE_MODEL_TUPLE = None | |
| def get_model(model_name): | |
| global BASE_MODEL_TUPLE,LARGE_MODEL_TUPLE | |
| download_bin(model_name) | |
| model = None | |
| if BASE_MODEL_TUPLE is None and model_name == 'vc1-base': | |
| repo_name = "facebook/" + model_name | |
| model_cfg = omegaconf.OmegaConf.load( | |
| hf_hub_download(repo_id=repo_name, filename=FILENAME,token=HF_TOKEN) | |
| ) | |
| BASE_MODEL_TUPLE = utils.instantiate(model_cfg) | |
| BASE_MODEL_TUPLE[0].eval() | |
| model = BASE_MODEL_TUPLE | |
| elif LARGE_MODEL_TUPLE is None and model_name == 'vc1-large': | |
| repo_name = "facebook/" + model_name | |
| model_cfg = omegaconf.OmegaConf.load( | |
| hf_hub_download(repo_id=repo_name, filename=FILENAME,token=HF_TOKEN) | |
| ) | |
| LARGE_MODEL_TUPLE = utils.instantiate(model_cfg) | |
| LARGE_MODEL_TUPLE[0].eval() | |
| model = LARGE_MODEL_TUPLE | |
| elif model_name == 'vc1-base': | |
| model = BASE_MODEL_TUPLE | |
| elif model_name == 'vc1-large': | |
| model = LARGE_MODEL_TUPLE | |
| return model | |
| def download_bin(model): | |
| bin_file = "" | |
| if model == "vc1-large": | |
| bin_file = 'vc1_vitl.pth' | |
| elif model == "vc1-base": | |
| bin_file = 'vc1_vitb.pth' | |
| else: | |
| raise NameError("model not found: " + model) | |
| repo_name = 'facebook/' + model | |
| bin_path = os.path.join(MODEL_DIR,bin_file) | |
| if not os.path.isfile(bin_path): | |
| model_bin = hf_hub_download(repo_id=repo_name, filename='pytorch_model.bin',local_dir=MODEL_DIR,local_dir_use_symlinks=True,token=HF_TOKEN) | |
| os.rename(model_bin, bin_path) | |
| def run_attn(input_img, model="vc1-base",discard_ratio=0.89): | |
| download_bin(model) | |
| model, embedding_dim, transform, metadata = get_model(model) | |
| if input_img.shape[0] != 3: | |
| input_img = input_img.transpose(2, 0, 1) | |
| if(len(input_img.shape)== 3): | |
| input_img = torch.tensor(input_img).unsqueeze(0) | |
| input_img = input_img.float() | |
| resize_transform = torchvision.transforms.Resize((250,250)) | |
| input_img = resize_transform(input_img) | |
| x = transform(input_img) | |
| attention_rollout = VITAttentionGradRollout(model,head_fusion="max",discard_ratio=discard_ratio) | |
| y = model(x) | |
| mask = attention_rollout.get_attn_mask() | |
| attn_img = overlay_attn(input_img[0].permute(1,2,0),mask) | |
| return attn_img | |
| model_type = gr.Dropdown( | |
| ["vc1-base", "vc1-large"], label="Model Size", value="vc1-base") | |
| input_img = gr.Image(shape=(250,250)) | |
| discard_ratio = gr.Slider(0,1,value=0.89) | |
| output_img = gr.Image(shape=(250,250)) | |
| css = "#component-2, .input-image, .image-preview {height: 240px !important}" | |
| markdown ="This is a demo for the Visual Cortex models. When passed an image input, it displays the attention(green) of the last layer of the transformer." | |
| demo = gr.Interface(fn=run_attn, title="Visual Cortex Model", description=markdown, | |
| examples=[[os.path.join('./imgs',x),None,None]for x in os.listdir(os.path.join(os.getcwd(),'imgs')) if 'jpg' in x], | |
| inputs=[input_img,model_type,discard_ratio],outputs=output_img,css=css) | |
| demo.launch() | |