171 lines
5.4 KiB
Python
171 lines
5.4 KiB
Python
# These codes are copied from modelscope revision c58451baead80d83281f063d12fb377fad415257
|
|
# CLIP
|
|
# Adapted from https://github.com/openai/CLIP.
|
|
# Originally MIT License, Copyright (c) 2021 OpenAI.
|
|
|
|
import hashlib
|
|
import os
|
|
import urllib
|
|
import warnings
|
|
from typing import Any, List, Union
|
|
|
|
import torch
|
|
from PIL import Image
|
|
from pkg_resources import packaging
|
|
from torchvision.transforms import (CenterCrop, Compose, Normalize, Resize,
|
|
ToTensor)
|
|
from tqdm import tqdm
|
|
|
|
from .model import build_model
|
|
from .simple_tokenizer import SimpleTokenizer as _Tokenizer
|
|
|
|
try:
|
|
from torchvision.transforms import InterpolationMode
|
|
BICUBIC = InterpolationMode.BICUBIC
|
|
except ImportError:
|
|
BICUBIC = Image.BICUBIC
|
|
|
|
if packaging.version.parse(
|
|
torch.__version__) < packaging.version.parse('1.7.1'):
|
|
warnings.warn('PyTorch version 1.7.1 or higher is recommended')
|
|
__all__ = ['load', 'tokenize']
|
|
|
|
|
|
def _convert_image_to_rgb(image):
|
|
return image.convert('RGB')
|
|
|
|
|
|
def _transform(n_px):
|
|
return Compose([
|
|
Resize(n_px, interpolation=BICUBIC),
|
|
CenterCrop(n_px),
|
|
_convert_image_to_rgb,
|
|
ToTensor(),
|
|
Normalize((0.48145466, 0.4578275, 0.40821073),
|
|
(0.26862954, 0.26130258, 0.27577711)),
|
|
])
|
|
|
|
|
|
def load(name: str,
|
|
device: Union[str, torch.device] = 'cuda'
|
|
if torch.cuda.is_available() else 'cpu',
|
|
jit: bool = False,
|
|
root: str = None):
|
|
|
|
if not jit:
|
|
model = build_model().to(device)
|
|
if str(device) == 'cpu':
|
|
model.float()
|
|
return model, _transform(model.visual.input_resolution)
|
|
|
|
# patch the device names
|
|
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)
|
|
|
|
# patch dtype to float32 on CPU
|
|
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
|
|
]: # dtype can be the second or third argument to aten::to()
|
|
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 tokenize(
|
|
_tokenizer,
|
|
texts: Union[str, List[str]],
|
|
context_length: int = 77,
|
|
truncate: bool = False) -> Union[torch.IntTensor, torch.LongTensor]:
|
|
"""
|
|
Returns the tokenized representation of given input string(s)
|
|
|
|
Parameters
|
|
----------
|
|
texts : Union[str, List[str]]
|
|
An input string or a list of input strings to tokenize
|
|
|
|
context_length : int
|
|
The context length to use; all CLIP models use 77 as the context length
|
|
|
|
truncate: bool
|
|
Whether to truncate the text in case its encoding is longer than the context length
|
|
|
|
Returns
|
|
-------
|
|
A two-dimensional tensor containing the resulting tokens, shape = [number of input strings, context_length].
|
|
We return LongTensor when torch version is <1.8.0, since older index_select requires indices to be long.
|
|
"""
|
|
if isinstance(texts, str):
|
|
texts = [texts]
|
|
|
|
sot_token = _tokenizer.encoder['<|startoftext|>']
|
|
eot_token = _tokenizer.encoder['<|endoftext|>']
|
|
all_tokens = [[sot_token] + _tokenizer.encode(text) + [eot_token]
|
|
for text in texts]
|
|
if packaging.version.parse(
|
|
torch.__version__) < packaging.version.parse('1.8.0'):
|
|
result = torch.zeros(len(all_tokens), context_length, dtype=torch.long)
|
|
else:
|
|
result = torch.zeros(len(all_tokens), context_length, dtype=torch.int)
|
|
|
|
for i, tokens in enumerate(all_tokens):
|
|
if len(tokens) > context_length:
|
|
if truncate:
|
|
tokens = tokens[:context_length]
|
|
tokens[-1] = eot_token
|
|
else:
|
|
raise RuntimeError(
|
|
f'Input {texts[i]} is too long for context length {context_length}'
|
|
)
|
|
result[i, :len(tokens)] = torch.tensor(tokens)
|
|
|
|
return result
|