55 lines
14 KiB
JSON
55 lines
14 KiB
JSON
{
|
|
"summary": "The code downloads and handles pre-trained models, including CLIP, checks availability, loads optimized versions, and implements device-specific patches while tokenizing input strings. This function encodes a list of texts using a tokenizer, adds start/end tokens, returns tensor, and optionally truncates if exceeding context length.",
|
|
"details": [
|
|
{
|
|
"comment": "Importing necessary libraries and defining variables for model choices and tokenizer.",
|
|
"location": "\"/media/root/Toshiba XG3/works/CLIP/docs/src/clip/clip.py\":0-32",
|
|
"content": "import hashlib\nimport os\nimport urllib\nimport warnings\nfrom typing import Any, Union, List\nfrom pkg_resources import packaging\nimport torch\nfrom PIL import Image\nfrom torchvision.transforms import Compose, Resize, CenterCrop, ToTensor, Normalize\nfrom tqdm import tqdm\nfrom .model import build_model\nfrom .simple_tokenizer import SimpleTokenizer as _Tokenizer\ntry:\n from torchvision.transforms import InterpolationMode\n BICUBIC = InterpolationMode.BICUBIC\nexcept ImportError:\n BICUBIC = Image.BICUBIC\nif packaging.version.parse(torch.__version__) < packaging.version.parse(\"1.7.1\"):\n warnings.warn(\"PyTorch version 1.7.1 or higher is recommended\")\n__all__ = [\"available_models\", \"load\", \"tokenize\"]\n_tokenizer = _Tokenizer()\n_MODELS = {\n \"RN50\": \"https://openaipublic.azureedge.net/clip/models/afeb0e10f9e5a86da6080e35cf09123aca3b358a0c3e3b6c78a7b63bc04b6762/RN50.pt\",\n \"RN101\": \"https://openaipublic.azureedge.net/clip/models/8fa8567bab74a42d41c5915025a8e4538c3bdbe8804a470a72f30b0d94fab599/RN101.pt\",\n "
|
|
},
|
|
{
|
|
"comment": "This code defines a dictionary of URLs for different pre-trained models and includes a function _download() to download the model files.",
|
|
"location": "\"/media/root/Toshiba XG3/works/CLIP/docs/src/clip/clip.py\":32-42",
|
|
"content": "\"RN50x4\": \"https://openaipublic.azureedge.net/clip/models/7e526bd135e493cef0776de27d5f42653e6b4c8bf9e0f653bb11773263205fdd/RN50x4.pt\",\n \"RN50x16\": \"https://openaipublic.azureedge.net/clip/models/52378b407f34354e150460fe41077663dd5b39c54cd0bfd2b27167a4a06ec9aa/RN50x16.pt\",\n \"RN50x64\": \"https://openaipublic.azureedge.net/clip/models/be1cfb55d75a9666199fb2206c106743da0f6468c9d327f3e0d0a543a9919d9c/RN50x64.pt\",\n \"ViT-B/32\": \"https://openaipublic.azureedge.net/clip/models/40d365715913c9da98579312b702a82c18be219cc2a73407c4526f58eba950af/ViT-B-32.pt\",\n \"ViT-B/16\": \"https://openaipublic.azureedge.net/clip/models/5806e77cd80f8b59890b7e101eabd078d9fb84e6937f9e85e4ecb61988df416f/ViT-B-16.pt\",\n \"ViT-L/14\": \"https://openaipublic.azureedge.net/clip/models/b8cca3fd41ae0c99ba7e8951adf17d267cdb84cd88be6f7c2e0eca1737a03836/ViT-L-14.pt\",\n \"ViT-L/14@336px\": \"https://openaipublic.azureedge.net/clip/models/3035c92b350959924f9f00213499208652fc7ea050643e8b385c2dac08641f02/ViT-L-14-336px.pt\",\n}\ndef _download(url: str, root: str):"
|
|
},
|
|
{
|
|
"comment": "Creates a directory and checks if the file already exists, then downloads or verifies the file's SHA256 checksum.",
|
|
"location": "\"/media/root/Toshiba XG3/works/CLIP/docs/src/clip/clip.py\":43-66",
|
|
"content": " os.makedirs(root, exist_ok=True)\n filename = os.path.basename(url)\n expected_sha256 = url.split(\"/\")[-2]\n download_target = os.path.join(root, filename)\n if os.path.exists(download_target) and not os.path.isfile(download_target):\n raise RuntimeError(f\"{download_target} exists and is not a regular file\")\n if os.path.isfile(download_target):\n if hashlib.sha256(open(download_target, \"rb\").read()).hexdigest() == expected_sha256:\n return download_target\n else:\n warnings.warn(f\"{download_target} exists, but the SHA256 checksum does not match; re-downloading the file\")\n with urllib.request.urlopen(url) as source, open(download_target, \"wb\") as output:\n with tqdm(total=int(source.info().get(\"Content-Length\")), ncols=80, unit='iB', unit_scale=True, unit_divisor=1024) as loop:\n while True:\n buffer = source.read(8192)\n if not buffer:\n break\n output.write(buffer)\n loop.update(len(buffer))"
|
|
},
|
|
{
|
|
"comment": "Code checks the SHA256 checksum of a downloaded model and raises an error if it doesn't match. It also defines functions for image transformation, loading available CLIP models, and converting images to RGB format.",
|
|
"location": "\"/media/root/Toshiba XG3/works/CLIP/docs/src/clip/clip.py\":68-101",
|
|
"content": " if hashlib.sha256(open(download_target, \"rb\").read()).hexdigest() != expected_sha256:\n raise RuntimeError(\"Model has been downloaded but the SHA256 checksum does not not match\")\n return download_target\ndef _convert_image_to_rgb(image):\n return image.convert(\"RGB\")\ndef _transform(n_px):\n return Compose([\n Resize(n_px, interpolation=BICUBIC),\n CenterCrop(n_px),\n _convert_image_to_rgb,\n ToTensor(),\n Normalize((0.48145466, 0.4578275, 0.40821073), (0.26862954, 0.26130258, 0.27577711)),\n ])\ndef available_models() -> List[str]:\n \"\"\"Returns the names of available CLIP models\"\"\"\n return list(_MODELS.keys())\ndef load(name: str, device: Union[str, torch.device] = \"cuda\" if torch.cuda.is_available() else \"cpu\", jit: bool = False, download_root: str = None):\n \"\"\"Load a CLIP model\n Parameters\n ----------\n name : str\n A model name listed by `clip.available_models()`, or the path to a model checkpoint containing the state_dict\n device : Union[str, torch.device]"
|
|
},
|
|
{
|
|
"comment": "The code downloads the CLIP model based on the specified name and device. It checks if the model is available as a file or downloads it from the provided root path. The JIT-optimized version of the model is loaded if 'jit' is True, otherwise, the non-JIT version is used.",
|
|
"location": "\"/media/root/Toshiba XG3/works/CLIP/docs/src/clip/clip.py\":102-130",
|
|
"content": " The device to put the loaded model\n jit : bool\n Whether to load the optimized JIT model or more hackable non-JIT model (default).\n download_root: str\n path to download the model files; by default, it uses \"~/.cache/clip\"\n Returns\n -------\n model : torch.nn.Module\n The CLIP model\n preprocess : Callable[[PIL.Image], torch.Tensor]\n A torchvision transform that converts a PIL image into a tensor that the returned model can take as its input\n \"\"\"\n if name in _MODELS:\n model_path = _download(_MODELS[name], download_root or os.path.expanduser(\"~/.cache/clip\"))\n elif os.path.isfile(name):\n model_path = name\n else:\n raise RuntimeError(f\"Model {name} not found; available models = {available_models()}\")\n with open(model_path, 'rb') as opened_file:\n try:\n # loading JIT archive\n model = torch.jit.load(opened_file, map_location=device if jit else \"cpu\").eval()\n state_dict = None\n except RuntimeError:"
|
|
},
|
|
{
|
|
"comment": "Loading saved state dict and handling JIT (Just-In-Time) support.\nIf not a JIT archive, loading as a state dict instead.\nLoading model with or without JIT support depending on jit variable.\nConverting model to float if device is CPU.\nReturning the model and transformed input resolution.\nPatching device names using torch.jit.trace.",
|
|
"location": "\"/media/root/Toshiba XG3/works/CLIP/docs/src/clip/clip.py\":131-156",
|
|
"content": " # loading saved state dict\n if jit:\n warnings.warn(f\"File {model_path} is not a JIT archive. Loading as a state dict instead\")\n jit = False\n state_dict = torch.load(opened_file, map_location=\"cpu\")\n if not jit:\n model = build_model(state_dict or model.state_dict()).to(device)\n if str(device) == \"cpu\":\n model.float()\n return model, _transform(model.visual.input_resolution)\n # patch the device names\n device_holder = torch.jit.trace(lambda: torch.ones([]).to(torch.device(device)), example_inputs=[])\n device_node = [n for n in device_holder.graph.findAllNodes(\"prim::Constant\") if \"Device\" in repr(n)][-1]\n def _node_get(node: torch._C.Node, key: str):\n \"\"\"Gets attributes of a node which is polymorphic over return type.\n From https://github.com/pytorch/pytorch/pull/82628\n \"\"\"\n sel = node.kindOf(key)\n return getattr(node, sel)(key)\n def patch_device(module):\n try:"
|
|
},
|
|
{
|
|
"comment": "Applying device-specific patches to the model's graph.",
|
|
"location": "\"/media/root/Toshiba XG3/works/CLIP/docs/src/clip/clip.py\":157-183",
|
|
"content": " graphs = [module.graph] if hasattr(module, \"graph\") else []\n except RuntimeError:\n graphs = []\n if hasattr(module, \"forward1\"):\n graphs.append(module.forward1.graph)\n for graph in graphs:\n for node in graph.findAllNodes(\"prim::Constant\"):\n if \"value\" in node.attributeNames() and str(_node_get(node, \"value\")).startswith(\"cuda\"):\n node.copyAttributes(device_node)\n model.apply(patch_device)\n patch_device(model.encode_image)\n patch_device(model.encode_text)\n # patch dtype to float32 on CPU\n if str(device) == \"cpu\":\n float_holder = torch.jit.trace(lambda: torch.ones([]).float(), example_inputs=[])\n float_input = list(float_holder.graph.findNode(\"aten::to\").inputs())[1]\n float_node = float_input.node()\n def patch_float(module):\n try:\n graphs = [module.graph] if hasattr(module, \"graph\") else []\n except RuntimeError:\n graphs = []"
|
|
},
|
|
{
|
|
"comment": "The code is applying a patch to select functions in the model to convert certain types to floats, then tokenizes input strings based on context_length.",
|
|
"location": "\"/media/root/Toshiba XG3/works/CLIP/docs/src/clip/clip.py\":185-213",
|
|
"content": " if hasattr(module, \"forward1\"):\n graphs.append(module.forward1.graph)\n for graph in graphs:\n for node in graph.findAllNodes(\"aten::to\"):\n inputs = list(node.inputs())\n for i in [1, 2]: # dtype can be the second or third argument to aten::to()\n if _node_get(inputs[i].node(), \"value\") == 5:\n inputs[i].node().copyAttributes(float_node)\n model.apply(patch_float)\n patch_float(model.encode_image)\n patch_float(model.encode_text)\n model.float()\n return model, _transform(model.input_resolution.item())\ndef tokenize(texts: Union[str, List[str]], context_length: int = 77, truncate: bool = False) -> Union[torch.IntTensor, torch.LongTensor]:\n \"\"\"\n Returns the tokenized representation of given input string(s)\n Parameters\n ----------\n texts : Union[str, List[str]]\n An input string or a list of input strings to tokenize\n context_length : int"
|
|
},
|
|
{
|
|
"comment": "This function takes a list of texts and encodes them using the tokenizer. It then adds start and end of text tokens, and returns a tensor of shape [number of input strings, context_length]. If torch version is <1.8.0, it uses LongTensor for indices.",
|
|
"location": "\"/media/root/Toshiba XG3/works/CLIP/docs/src/clip/clip.py\":214-236",
|
|
"content": " The context length to use; all CLIP models use 77 as the context length\n truncate: bool\n Whether to truncate the text in case its encoding is longer than the context length\n Returns\n -------\n A two-dimensional tensor containing the resulting tokens, shape = [number of input strings, context_length].\n We return LongTensor when torch version is <1.8.0, since older index_select requires indices to be long.\n \"\"\"\n if isinstance(texts, str):\n texts = [texts]\n sot_token = _tokenizer.encoder[\"<|startoftext|>\"]\n eot_token = _tokenizer.encoder[\"<|endoftext|>\"]\n all_tokens = [[sot_token] + _tokenizer.encode(text) + [eot_token] for text in texts]\n if packaging.version.parse(torch.__version__) < packaging.version.parse(\"1.8.0\"):\n result = torch.zeros(len(all_tokens), context_length, dtype=torch.long)\n else:\n result = torch.zeros(len(all_tokens), context_length, dtype=torch.int)\n for i, tokens in enumerate(all_tokens):\n if len(tokens) > context_length:"
|
|
},
|
|
{
|
|
"comment": "If truncate is True, only keep the first context_length tokens and set the last token to eot_token. If not, raise an error if input text exceeds the context length. Store the tokens as a torch tensor in result.",
|
|
"location": "\"/media/root/Toshiba XG3/works/CLIP/docs/src/clip/clip.py\":237-244",
|
|
"content": " if truncate:\n tokens = tokens[:context_length]\n tokens[-1] = eot_token\n else:\n raise RuntimeError(f\"Input {texts[i]} is too long for context length {context_length}\")\n result[i, :len(tokens)] = torch.tensor(tokens)\n return result"
|
|
}
|
|
]
|
|
} |