David310/Detect_AI-generated_Image

1import hashlib2import os3import urllib4import warnings5from typing import Any, Union, List6from pkg_resources import packaging7 8import torch9from PIL import Image10from torchvision.transforms import Compose, Resize, CenterCrop, ToTensor, Normalize11from tqdm import tqdm12 13from .model import build_model14from .simple_tokenizer import SimpleTokenizer as _Tokenizer15 16try:17    from torchvision.transforms import InterpolationMode18    BICUBIC = InterpolationMode.BICUBIC19except ImportError:20    BICUBIC = Image.BICUBIC21 22 23if packaging.version.parse(torch.__version__) < packaging.version.parse("1.7.1"):24    warnings.warn("PyTorch version 1.7.1 or higher is recommended")25 26 27__all__ = ["available_models", "load", "tokenize"]28_tokenizer = _Tokenizer()29 30_MODELS = {31    "RN50": "https://openaipublic.azureedge.net/clip/models/afeb0e10f9e5a86da6080e35cf09123aca3b358a0c3e3b6c78a7b63bc04b6762/RN50.pt",32    "RN101": "https://openaipublic.azureedge.net/clip/models/8fa8567bab74a42d41c5915025a8e4538c3bdbe8804a470a72f30b0d94fab599/RN101.pt",33    "RN50x4": "https://openaipublic.azureedge.net/clip/models/7e526bd135e493cef0776de27d5f42653e6b4c8bf9e0f653bb11773263205fdd/RN50x4.pt",34    "RN50x16": "https://openaipublic.azureedge.net/clip/models/52378b407f34354e150460fe41077663dd5b39c54cd0bfd2b27167a4a06ec9aa/RN50x16.pt",35    "RN50x64": "https://openaipublic.azureedge.net/clip/models/be1cfb55d75a9666199fb2206c106743da0f6468c9d327f3e0d0a543a9919d9c/RN50x64.pt",36    "ViT-B/32": "https://openaipublic.azureedge.net/clip/models/40d365715913c9da98579312b702a82c18be219cc2a73407c4526f58eba950af/ViT-B-32.pt",37    "ViT-B/16": "https://openaipublic.azureedge.net/clip/models/5806e77cd80f8b59890b7e101eabd078d9fb84e6937f9e85e4ecb61988df416f/ViT-B-16.pt",38    "ViT-L/14": "https://openaipublic.azureedge.net/clip/models/b8cca3fd41ae0c99ba7e8951adf17d267cdb84cd88be6f7c2e0eca1737a03836/ViT-L-14.pt",39    "ViT-L/14@336px": "https://openaipublic.azureedge.net/clip/models/3035c92b350959924f9f00213499208652fc7ea050643e8b385c2dac08641f02/ViT-L-14-336px.pt",40}41 42 43def _download(url: str, root: str):44    os.makedirs(root, exist_ok=True)45    filename = os.path.basename(url)46 47    expected_sha256 = url.split("/")[-2]48    download_target = os.path.join(root, filename)49 50    if os.path.exists(download_target) and not os.path.isfile(download_target):51        raise RuntimeError(f"{download_target} exists and is not a regular file")52 53    if os.path.isfile(download_target):54        if hashlib.sha256(open(download_target, "rb").read()).hexdigest() == expected_sha256:55            return download_target56        else:57            warnings.warn(f"{download_target} exists, but the SHA256 checksum does not match; re-downloading the file")58 59    with urllib.request.urlopen(url) as source, open(download_target, "wb") as output:60        with tqdm(total=int(source.info().get("Content-Length")), ncols=80, unit='iB', unit_scale=True, unit_divisor=1024) as loop:61            while True:62                buffer = source.read(8192)63                if not buffer:64                    break65 66                output.write(buffer)67                loop.update(len(buffer))68 69    if hashlib.sha256(open(download_target, "rb").read()).hexdigest() != expected_sha256:70        raise RuntimeError("Model has been downloaded but the SHA256 checksum does not not match")71 72    return download_target73 74 75def _convert_image_to_rgb(image):76    return image.convert("RGB")77 78 79def _transform(n_px):80    return Compose([81        Resize(n_px, interpolation=BICUBIC),82        CenterCrop(n_px),83        _convert_image_to_rgb,84        ToTensor(),85        Normalize((0.48145466, 0.4578275, 0.40821073), (0.26862954, 0.26130258, 0.27577711)),86    ])87 88 89def available_models() -> List[str]:90    """Returns the names of available CLIP models"""91    return list(_MODELS.keys())92 93 94def load(name: str, device: Union[str, torch.device] = "cuda" if torch.cuda.is_available() else "cpu", jit: bool = False, download_root: str = None):95    """Load a CLIP model96 97    Parameters98    ----------99    name : str100        A model name listed by `clip.available_models()`, or the path to a model checkpoint containing the state_dict101 102    device : Union[str, torch.device]103        The device to put the loaded model104 105    jit : bool106        Whether to load the optimized JIT model or more hackable non-JIT model (default).107 108    download_root: str109        path to download the model files; by default, it uses "~/.cache/clip"110 111    Returns112    -------113    model : torch.nn.Module114        The CLIP model115 116    preprocess : Callable[[PIL.Image], torch.Tensor]117        A torchvision transform that converts a PIL image into a tensor that the returned model can take as its input118    """119    if name in _MODELS:120        model_path = _download(_MODELS[name], download_root or os.path.expanduser("~/.cache/clip"))121    elif os.path.isfile(name):122        model_path = name123    else:124        raise RuntimeError(f"Model {name} not found; available models = {available_models()}")125 126    with open(model_path, 'rb') as opened_file:127        try:128            # loading JIT archive129            model = torch.jit.load(opened_file, map_location=device if jit else "cpu").eval()130            state_dict = None131        except RuntimeError:132            # loading saved state dict133            if jit:134                warnings.warn(f"File {model_path} is not a JIT archive. Loading as a state dict instead")135                jit = False136            state_dict = torch.load(opened_file, map_location="cpu")137 138    if not jit:139        model = build_model(state_dict or model.state_dict()).to(device)140        if str(device) == "cpu":141            model.float()142        return model, _transform(model.visual.input_resolution)143 144    # patch the device names145    device_holder = torch.jit.trace(lambda: torch.ones([]).to(torch.device(device)), example_inputs=[])146    device_node = [n for n in device_holder.graph.findAllNodes("prim::Constant") if "Device" in repr(n)][-1]147 148    def patch_device(module):149        try:150            graphs = [module.graph] if hasattr(module, "graph") else []151        except RuntimeError:152            graphs = []153 154        if hasattr(module, "forward1"):155            graphs.append(module.forward1.graph)156 157        for graph in graphs:158            for node in graph.findAllNodes("prim::Constant"):159                if "value" in node.attributeNames() and str(node["value"]).startswith("cuda"):160                    node.copyAttributes(device_node)161 162    model.apply(patch_device)163    patch_device(model.encode_image)164    patch_device(model.encode_text)165 166    # patch dtype to float32 on CPU167    if str(device) == "cpu":168        float_holder = torch.jit.trace(lambda: torch.ones([]).float(), example_inputs=[])169        float_input = list(float_holder.graph.findNode("aten::to").inputs())[1]170        float_node = float_input.node()171 172        def patch_float(module):173            try:174                graphs = [module.graph] if hasattr(module, "graph") else []175            except RuntimeError:176                graphs = []177 178            if hasattr(module, "forward1"):179                graphs.append(module.forward1.graph)180 181            for graph in graphs:182                for node in graph.findAllNodes("aten::to"):183                    inputs = list(node.inputs())184                    for i in [1, 2]:  # dtype can be the second or third argument to aten::to()185                        if inputs[i].node()["value"] == 5:186                            inputs[i].node().copyAttributes(float_node)187 188        model.apply(patch_float)189        patch_float(model.encode_image)190        patch_float(model.encode_text)191 192        model.float()193 194    return model, _transform(model.input_resolution.item())195 196 197def tokenize(texts: Union[str, List[str]], context_length: int = 77, truncate: bool = False) -> Union[torch.IntTensor, torch.LongTensor]:198    """199    Returns the tokenized representation of given input string(s)200 201    Parameters202    ----------203    texts : Union[str, List[str]]204        An input string or a list of input strings to tokenize205 206    context_length : int207        The context length to use; all CLIP models use 77 as the context length208 209    truncate: bool210        Whether to truncate the text in case its encoding is longer than the context length211 212    Returns213    -------214    A two-dimensional tensor containing the resulting tokens, shape = [number of input strings, context_length].215    We return LongTensor when torch version is <1.8.0, since older index_select requires indices to be long.216    """217    if isinstance(texts, str):218        texts = [texts]219 220    sot_token = _tokenizer.encoder["<|startoftext|>"]221    eot_token = _tokenizer.encoder["<|endoftext|>"]222    all_tokens = [[sot_token] + _tokenizer.encode(text) + [eot_token] for text in texts]223    if packaging.version.parse(torch.__version__) < packaging.version.parse("1.8.0"):224        result = torch.zeros(len(all_tokens), context_length, dtype=torch.long)225    else:226        result = torch.zeros(len(all_tokens), context_length, dtype=torch.int)227 228    for i, tokens in enumerate(all_tokens):229        if len(tokens) > context_length:230            if truncate:231                tokens = tokens[:context_length]232                tokens[-1] = eot_token233            else:234                raise RuntimeError(f"Input {texts[i]} is too long for context length {context_length}")235        result[i, :len(tokens)] = torch.tensor(tokens)236 237    return result238