| import hashlib
|
| import os
|
| import urllib
|
| import warnings
|
| from typing import Union, List
|
| from pkg_resources import packaging
|
|
|
| import torch
|
| from PIL import Image
|
| from torchvision.transforms import Compose, Resize, ToTensor, Normalize
|
| from tqdm import tqdm
|
| import numpy as np
|
|
|
| from .build_model import build_model
|
| from .simple_tokenizer import SimpleTokenizer as _Tokenizer
|
| from torchvision.transforms import InterpolationMode
|
|
|
| if packaging.version.parse(torch.__version__) < packaging.version.parse("1.7.1"):
|
| warnings.warn("PyTorch version 1.7.1 or higher is recommended")
|
|
|
|
|
| __all__ = ["available_models", "load",
|
| "get_similarity_map", "compute_similarity"]
|
| _tokenizer = _Tokenizer()
|
|
|
| _MODELS = {
|
| "ViT-L/14@336px": "https://openaipublic.azureedge.net/clip/models/3035c92b350959924f9f00213499208652fc7ea050643e8b385c2dac08641f02/ViT-L-14-336px.pt",
|
| }
|
|
|
|
|
| def _download(
|
| url: str,
|
| cache_dir: Union[str, None] = None,
|
| ):
|
|
|
| if not cache_dir:
|
|
|
| cache_dir = os.path.expanduser("/remote-home/iot_zhouqihang/root/.cache/clip")
|
| os.makedirs(cache_dir, exist_ok=True)
|
| filename = os.path.basename(url)
|
|
|
| if 'openaipublic' in url:
|
| expected_sha256 = url.split("/")[-2]
|
| elif 'mlfoundations' in url:
|
| expected_sha256 = os.path.splitext(filename)[0].split("-")[-1]
|
| else:
|
| expected_sha256 = ''
|
|
|
| download_target = os.path.join(cache_dir, filename)
|
|
|
| if os.path.exists(download_target) and not os.path.isfile(download_target):
|
| raise RuntimeError(f"{download_target} exists and is not a regular file")
|
|
|
| if os.path.isfile(download_target):
|
| if expected_sha256:
|
| if hashlib.sha256(open(download_target, "rb").read()).hexdigest().startswith(expected_sha256):
|
| return download_target
|
| else:
|
| warnings.warn(f"{download_target} exists, but the SHA256 checksum does not match; re-downloading the file")
|
| else:
|
| return download_target
|
|
|
| with urllib.request.urlopen(url) as source, open(download_target, "wb") as output:
|
| with tqdm(total=int(source.headers.get("Content-Length")), ncols=80, unit='iB', unit_scale=True) as loop:
|
| while True:
|
| buffer = source.read(8192)
|
| if not buffer:
|
| break
|
|
|
| output.write(buffer)
|
| loop.update(len(buffer))
|
|
|
| if expected_sha256 and not hashlib.sha256(open(download_target, "rb").read()).hexdigest().startswith(expected_sha256):
|
| raise RuntimeError(f"Model has been downloaded but the SHA256 checksum does not not match")
|
|
|
| return download_target
|
|
|
|
|
| def _convert_image_to_rgb(image):
|
| return image.convert("RGB")
|
|
|
|
|
| def _transform(n_px):
|
| return Compose([
|
| Resize((n_px, n_px), interpolation=InterpolationMode.BICUBIC),
|
|
|
| _convert_image_to_rgb,
|
| ToTensor(),
|
| Normalize((0.48145466, 0.4578275, 0.40821073), (0.26862954, 0.26130258, 0.27577711)),
|
| ])
|
|
|
|
|
| def available_models() -> List[str]:
|
| """Returns the names of available CLIP models"""
|
| return list(_MODELS.keys())
|
|
|
|
|
| def load_state_dict(checkpoint_path: str, map_location='cpu'):
|
| checkpoint = torch.load(checkpoint_path, map_location=map_location)
|
| if isinstance(checkpoint, dict) and 'state_dict' in checkpoint:
|
| state_dict = checkpoint['state_dict']
|
| else:
|
| state_dict = checkpoint
|
| if next(iter(state_dict.items()))[0].startswith('module'):
|
| state_dict = {k[7:]: v for k, v in state_dict.items()}
|
| return state_dict
|
|
|
| def load_checkpoint(model, checkpoint_path, strict=True):
|
| state_dict = load_state_dict(checkpoint_path)
|
|
|
| if 'positional_embedding' in state_dict and not hasattr(model, 'positional_embedding'):
|
| state_dict = convert_to_custom_text_state_dict(state_dict)
|
| resize_pos_embed(state_dict, model)
|
| incompatible_keys = model.load_state_dict(state_dict, strict=strict)
|
| return incompatible_keys
|
|
|
| def load(name: str, device: Union[str, torch.device] = "cuda" if torch.cuda.is_available() else "cpu", design_details = None, jit: bool = False, download_root: str = None):
|
| """Load a CLIP model
|
|
|
| Parameters
|
| ----------
|
| name : str
|
| A model name listed by `clip.available_models()`, or the path to a model checkpoint containing the state_dict
|
|
|
| device : Union[str, torch.device]
|
| The device to put the loaded model
|
|
|
| jit : bool
|
| Whether to load the optimized JIT model or more hackable non-JIT model (default).
|
|
|
| download_root: str
|
| path to download the model files; by default, it uses "~/.cache/clip"
|
|
|
| Returns
|
| -------
|
| model : torch.nn.Module
|
| The CLIP model
|
|
|
| preprocess : Callable[[PIL.Image], torch.Tensor]
|
| A torchvision transform that converts a PIL image into a tensor that the returned model can take as its input
|
| """
|
| print("name", name)
|
| if name in _MODELS:
|
|
|
| model_path = _download(_MODELS[name], download_root or os.path.expanduser("/remote-home/iot_zhouqihang/root/.cache/clip"))
|
| elif os.path.isfile(name):
|
| model_path = name
|
| else:
|
| raise RuntimeError(f"Model {name} not found; available models = {available_models()}")
|
|
|
| with open(model_path, 'rb') as opened_file:
|
| try:
|
|
|
| model = torch.jit.load(opened_file, map_location=device if jit else "cpu").eval()
|
| state_dict = None
|
| except RuntimeError:
|
|
|
| if jit:
|
| warnings.warn(f"File {model_path} is not a JIT archive. Loading as a state dict instead")
|
| jit = False
|
| state_dict = torch.load(opened_file, map_location="cpu")
|
|
|
| if not jit:
|
| model = build_model(name, state_dict or model.state_dict(), design_details).to(device)
|
| if str(device) == "cpu":
|
| model.float()
|
| return model, _transform(model.visual.input_resolution)
|
|
|
|
|
| device_holder = torch.jit.trace(lambda: torch.ones([]).to(torch.device(device)), example_inputs=[])
|
| device_node = [n for n in device_holder.graph.findAllNodes("prim::Constant") if "Device" in repr(n)][-1]
|
|
|
| def patch_device(module):
|
| try:
|
| graphs = [module.graph] if hasattr(module, "graph") else []
|
| except RuntimeError:
|
| graphs = []
|
|
|
| if hasattr(module, "forward1"):
|
| graphs.append(module.forward1.graph)
|
|
|
| for graph in graphs:
|
| for node in graph.findAllNodes("prim::Constant"):
|
| if "value" in node.attributeNames() and str(node["value"]).startswith("cuda"):
|
| node.copyAttributes(device_node)
|
|
|
| model.apply(patch_device)
|
| patch_device(model.encode_image)
|
| patch_device(model.encode_text)
|
|
|
|
|
| if str(device) == "cpu":
|
| float_holder = torch.jit.trace(lambda: torch.ones([]).float(), example_inputs=[])
|
| float_input = list(float_holder.graph.findNode("aten::to").inputs())[1]
|
| float_node = float_input.node()
|
|
|
| def patch_float(module):
|
| try:
|
| graphs = [module.graph] if hasattr(module, "graph") else []
|
| except RuntimeError:
|
| graphs = []
|
|
|
| if hasattr(module, "forward1"):
|
| graphs.append(module.forward1.graph)
|
|
|
| for graph in graphs:
|
| for node in graph.findAllNodes("aten::to"):
|
| inputs = list(node.inputs())
|
| for i in [1, 2]:
|
| if inputs[i].node()["value"] == 5:
|
| inputs[i].node().copyAttributes(float_node)
|
|
|
| model.apply(patch_float)
|
| patch_float(model.encode_image)
|
| patch_float(model.encode_text)
|
|
|
| model.float()
|
|
|
| return model, _transform(model.input_resolution.item())
|
|
|
|
|
| def get_similarity_map(sm, shape):
|
| side = int(sm.shape[1] ** 0.5)
|
| sm = sm.reshape(sm.shape[0], side, side, -1).permute(0, 3, 1, 2)
|
| sm = torch.nn.functional.interpolate(sm, shape, mode='bilinear')
|
| sm = sm.permute(0, 2, 3, 1)
|
| return sm
|
|
|
|
|
| def compute_similarity(image_features, text_features, t=2):
|
| prob_1 = image_features[:, :1, :] @ text_features.t()
|
| b, n_t, n_i, c = image_features.shape[0], text_features.shape[0], image_features.shape[1], image_features.shape[2]
|
| feats = image_features.reshape(b, n_i, 1, c) * text_features.reshape(1, 1, n_t, c)
|
| similarity = feats.sum(-1)
|
| return (similarity/0.07).softmax(-1), prob_1
|
|
|
