From 2841a16c03a2abcd463f7820fcc0625fb43f1b15 Mon Sep 17 00:00:00 2001 From: BIT2024 Date: Tue, 20 Aug 2024 17:00:23 +0800 Subject: [PATCH] ADD file via upload --- jclip/clip.py | 187 ++++++++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 187 insertions(+) create mode 100644 jclip/clip.py diff --git a/jclip/clip.py b/jclip/clip.py new file mode 100644 index 0000000..537c62a --- /dev/null +++ b/jclip/clip.py @@ -0,0 +1,187 @@ +import hashlib +import os +import numpy as np +import urllib +import warnings +from typing import Union, List + +import jittor as jt +from tqdm import tqdm + +from .model import build_model +from .simple_tokenizer import SimpleTokenizer as _Tokenizer +from PIL import Image +from jittor.transform import CenterCrop, ImageNormalize, Compose, _setup_size, to_pil_image, resize + +__all__ = ["available_models", "load", "tokenize"] +_tokenizer = _Tokenizer() + +_MODELS = { + "RN50": + "https://openaipublic.azureedge.net/clip/models/afeb0e10f9e5a86da6080e35cf09123aca3b358a0c3e3b6c78a7b63bc04b6762/RN50.pt", + "RN101": + "https://openaipublic.azureedge.net/clip/models/8fa8567bab74a42d41c5915025a8e4538c3bdbe8804a470a72f30b0d94fab599/RN101.pt", + "RN50x4": + "https://openaipublic.azureedge.net/clip/models/7e526bd135e493cef0776de27d5f42653e6b4c8bf9e0f653bb11773263205fdd/RN50x4.pt", + "RN50x16": + "https://openaipublic.azureedge.net/clip/models/52378b407f34354e150460fe41077663dd5b39c54cd0bfd2b27167a4a06ec9aa/RN50x16.pt", + "RN50x64": + "https://openaipublic.azureedge.net/clip/models/be1cfb55d75a9666199fb2206c106743da0f6468c9d327f3e0d0a543a9919d9c/RN50x64.pt", + "ViT-B/32": + "https://openaipublic.azureedge.net/clip/models/40d365715913c9da98579312b702a82c18be219cc2a73407c4526f58eba950af/ViT-B-32.pt", + "ViT-B/16": + "https://openaipublic.azureedge.net/clip/models/5806e77cd80f8b59890b7e101eabd078d9fb84e6937f9e85e4ecb61988df416f/ViT-B-16.pt", + "ViT-L/14": + "https://openaipublic.azureedge.net/clip/models/b8cca3fd41ae0c99ba7e8951adf17d267cdb84cd88be6f7c2e0eca1737a03836/ViT-L-14.pt", + "ViT-L/14@336px": + "https://openaipublic.azureedge.net/clip/models/3035c92b350959924f9f00213499208652fc7ea050643e8b385c2dac08641f02/ViT-L-14-336px.pt", +} + + +def _download(url: str, root: str): + os.makedirs(root, exist_ok=True) + filename = os.path.basename(url) + + expected_sha256 = url.split("/")[-2] + download_target = os.path.join(root, 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 hashlib.sha256(open(download_target, + "rb").read()).hexdigest() == expected_sha256: + return download_target + else: + warnings.warn( + f"{download_target} exists, but the SHA256 checksum does not match; re-downloading the file" + ) + + with urllib.request.urlopen(url) as source, open(download_target, + "wb") as output: + with tqdm(total=int(source.info().get("Content-Length")), + ncols=80, + unit='iB', + unit_scale=True, + unit_divisor=1024) as loop: + while True: + buffer = source.read(8192) + if not buffer: + break + + output.write(buffer) + loop.update(len(buffer)) + + if hashlib.sha256(open(download_target, + "rb").read()).hexdigest() != expected_sha256: + raise RuntimeError( + "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 to_tensor(data): + return jt.Var(data) + + +class ImageToTensor(object): + + def __call__(self, input): + input = np.asarray(input) + if len(input.shape) < 3: + input = np.expand_dims(input, -1) + return to_tensor(input) + + +class Resize: + + def __init__(self, size, mode=Image.BILINEAR): + if isinstance(size, int): + self.size = size + else: + self.size = _setup_size( + size, + error_msg="If size is a sequence, it should have 2 values") + self.mode = mode + + def __call__(self, img: Image.Image): + if not isinstance(img, Image.Image): + img = to_pil_image(img) + if isinstance(self.size, int): + w, h = img.size + + short, long = (w, h) if w <= h else (h, w) + if short == self.size: + return img + + new_short, new_long = self.size, int(self.size * long / short) + new_w, new_h = (new_short, new_long) if w <= h else (new_long, + new_short) + size = (new_h, new_w) + return resize(img, size, self.mode) + + +def _transform(n_px): + return Compose([ + Resize(n_px, mode=Image.BICUBIC), + CenterCrop(n_px), _convert_image_to_rgb, + ImageNormalize((0.48145466, 0.4578275, 0.40821073), + (0.26862954, 0.26130258, 0.27577711)), + ImageToTensor() + ]) + + +def available_models() -> List[str]: + """Returns the names of available CLIP models""" + return list(_MODELS.keys()) + + +def load(name, download_root=None): + if name in _MODELS: + model_path = _download( + _MODELS[name], download_root + or os.path.expanduser("~/.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: + state_dict = jt.load(model_path) + + model = build_model(state_dict) + return model, _transform(model.visual.input_resolution) + + +def tokenize(texts: Union[str, List[str]], + context_length: int = 77, + truncate: bool = False): + 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] + + result = jt.zeros((len(all_tokens), context_length), dtype=jt.int64) + + 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)] = jt.Var(tokens) + + return result