Upload 68 files

app.py

@@ -0,0 +1,68 @@
+import gradio as gr
+from PIL import Image
+from detector import FakeImageDetector  
+
+print("正在初始化检测器，请稍候...")
+try:
+    detector = FakeImageDetector()
+    
+    print("检测器初始化完成，Web 服务准备就绪。")
+    models_loaded = True
+except Exception as e:
+    print(f"模型加载失败: {e}")
+    models_loaded = False
+    detector = None
+
+def predict_image(input_image_numpy, threshold):
+    """
+    接收 Gradio 的输入 (numpy array)，调用检测器，并返回结果。
+    """
+    if not models_loaded or detector is None:
+        return "错误：模型未能成功加载，请检查后台日志。", None
+
+    pil_image = Image.fromarray(input_image_numpy)
+    
+    result_text, score = detector.detect(pil_image, threshold)
+    
+    label_color = "red" if score > threshold else "green"
+    
+    return result_text, gr.Label(value=f"{score:.10f}", label=label_color)
+
+
+with gr.Blocks(title="伪造图像检测器", theme=gr.themes.Soft()) as demo:
+    gr.Markdown(
+        """
+        # 伪造图像检测器 (Fake Image Detector)
+        上传一张图片，模型将判断其为 **真实的 (Real)** 还是 **AI 生成的伪造图像 (Fake)**。
+        """
+    )
+    
+    with gr.Row():
+        with gr.Column(scale=1):
+            # 输入组件
+            image_input = gr.Image(type="numpy", label="上传图片", height=300)
+            # threshold_slider = gr.Slider(
+            #     minimum=0.495, maximum=0.55, value=0.499892068, step=0.0001, 
+            #     label="检测门限 (Threshold)",
+            #     info="得分低于此门限的图片被认为是伪造的"
+            # )
+            submit_btn = gr.Button("开始检测", variant="primary")
+        
+        with gr.Column(scale=1):
+            # 输出组件
+            result_output_text = gr.Textbox(label="检测结论", lines=2)
+            # 这里我们用一个临时的 Label 来显示带颜色的分数
+            result_output_score = gr.Label(label="模型原始得分")
+
+    submit_btn.click(
+        fn=predict_image, 
+        inputs=[image_input, 0.49999], 
+        outputs=[result_output_text, result_output_score]
+    )
+
+if not models_loaded:
+    print("\n由于模型加载失败，Gradio Web服务无法启动。")
+else:
+    print("正在启动 Gradio 服务...")
+    
+    demo.launch(server_name="0.0.0.0")

augmentations_clip.py

@@ -0,0 +1,255 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+#
+# This source code is licensed under the Apache License, Version 2.0
+# found in the LICENSE file in the root directory of this source tree.
+
+# import logging
+
+from torchvision import transforms
+
+import torch 
+import cv2
+from PIL import Image
+import numpy as np
+
+from my_transforms import (
+    GaussianBlur,
+    make_normalize_transform,
+    make_normalize_transform_clip,
+)
+
+def add_gaussian_noise(tensor, mean=0.0, std=0.1):
+    noise = torch.randn(tensor.size()).cuda() * std + mean
+    return tensor + noise
+
+
+
+
+class DataAugmentationCLIP(object):
+    def __init__(
+        self,
+        global_crops_scale,
+        local_crops_scale,
+        local_crops_number,
+        global_crops_size=224,
+        local_crops_size=96,
+    ):
+
+        self.source_trans = transforms.Compose([
+            # transforms.RandomCrop(224),
+            # transforms.CenterCrop(224),
+            transforms.ToTensor(),
+            make_normalize_transform_clip(),
+        ])
+
+        # self.crop = transforms.Compose([
+        #     transforms.CenterCrop(224),
+           
+        # ])
+
+        self.crop = transforms.Compose([
+    transforms.Resize(224),  # 将短边缩放到 224，长边会按比例缩放
+    transforms.RandomCrop(224),  # 然后裁剪到 224x224
+])
+
+        self.centercrop = transforms.Compose([
+            transforms.CenterCrop(224),
+           
+        ])
+
+        self.randomcrop = transforms.Compose([
+            transforms.RandomCrop(224),
+           
+        ])
+
+        self.local_crops_number = local_crops_number
+
+    def __call__(self, image):
+        output = {}
+        output["source"] = []
+
+        if np.array(image).shape[0]<224 or np.array(image).shape[1]<224:
+            crops_all = [
+             self.centercrop(image) for _ in range(self.local_crops_number)
+         ]
+        else:
+            crops_all = [
+                self.centercrop(image) for _ in range(self.local_crops_number)
+            ]
+            
+        for crops_image in crops_all:
+            output["source"].append(self.source_trans(crops_image))   #单独使用好一些
+    
+
+        output["offsets"] = ()
+
+        return output
+    
+
+class DataAugmentationDINO(object):
+    def __init__(
+        self,
+        global_crops_scale,
+        local_crops_scale,
+        local_crops_number,
+        global_crops_size=224,
+        local_crops_size=96,
+    ):
+
+        self.source_trans = transforms.Compose([
+            # transforms.RandomCrop(224),
+            # transforms.CenterCrop(224),
+            transforms.ToTensor(),
+            make_normalize_transform(),
+        ])
+
+        # self.crop = transforms.Compose([
+        #     transforms.CenterCrop(224),
+           
+        # ])
+
+        self.crop = transforms.Compose([
+    transforms.Resize(224),  # 将短边缩放到 224，长边会按比例缩放
+    transforms.CenterCrop(224),  # 然后裁剪到 224x224
+])
+
+        self.centercrop = transforms.Compose([
+            transforms.CenterCrop(224),
+           
+        ])
+
+        self.local_crops_number = local_crops_number
+
+    def __call__(self, image):
+        output = {}
+        output["source"] = []
+
+        if np.array(image).shape[0]<224 or np.array(image).shape[1]<224:
+            crops_all = [
+             self.centercrop(image) for _ in range(self.local_crops_number)
+         ]
+        else:
+            crops_all = [
+                self.centercrop(image) for _ in range(self.local_crops_number)
+            ]
+            
+        for crops_image in crops_all:
+            output["source"].append(self.source_trans(crops_image))   #单独使用好一些
+    
+
+        output["offsets"] = ()
+
+        return output
+
+
+class DataAugmentationResNet_test(object):
+    def __init__(
+        self,
+        global_crops_scale,
+        local_crops_scale,
+        local_crops_number,
+        global_crops_size=224,
+        local_crops_size=96,
+    ):
+
+        self.source_trans = transforms.Compose([
+            # transforms.RandomCrop(224),
+            # transforms.CenterCrop(224),
+            transforms.ToTensor(),
+            make_normalize_transform(),
+        ])
+
+        # self.crop = transforms.Compose([
+        #     transforms.CenterCrop(224),
+           
+        # ])
+
+        self.crop = transforms.Compose([
+    transforms.Resize(224),  # 将短边缩放到 224，长边会按比例缩放
+    transforms.CenterCrop(224),  # 然后裁剪到 224x224
+])
+
+        self.centercrop = transforms.Compose([
+            transforms.CenterCrop(224),
+           
+        ])
+
+        self.local_crops_number = local_crops_number
+
+    def __call__(self, image):
+        output = {}
+        output["source"] = []
+
+        if np.array(image).shape[0]<224 or np.array(image).shape[1]<224:
+            crops_all = [
+             self.centercrop(image) for _ in range(self.local_crops_number)
+         ]
+        else:
+            crops_all = [
+                self.centercrop(image) for _ in range(self.local_crops_number)
+            ]
+            
+        for crops_image in crops_all:
+            output["source"].append(self.source_trans(crops_image))   #单独使用好一些
+    
+
+        output["offsets"] = ()
+
+        return output
+    
+
+
+class DataAugmentationCLIP_gen(object):
+    def __init__(
+        self,
+        global_crops_scale,
+        local_crops_scale,
+        local_crops_number,
+        global_crops_size=224,
+        local_crops_size=96,
+    ):
+
+        self.source_trans = transforms.Compose([
+            # transforms.RandomCrop(224),
+            # transforms.CenterCrop(224),
+            transforms.ToTensor(),
+            make_normalize_transform_clip(),
+        ])
+
+        # self.crop = transforms.Compose([
+        #     transforms.RandomCrop(224),
+           
+        # ])
+
+        self.crop = transforms.Compose([
+    transforms.Resize(224),  # 将短边缩放到 224，长边会按比例缩放
+    transforms.CenterCrop(224),  # 然后裁剪到 224x224
+])
+
+        self.centercrop = transforms.Compose([
+            transforms.CenterCrop(224),
+           
+        ])
+
+        self.local_crops_number = local_crops_number
+
+    def __call__(self, image):
+        output = {}
+        output["source"] = []
+
+        # if np.array(image).shape[0]<224 or np.array(image).shape[1]<224:
+        #     crops_all = [
+        #      self.crop(self.centercrop(image)) for _ in range(self.local_crops_number)
+        #  ]
+        # else:
+        crops_all = [
+            self.crop(image) for _ in range(self.local_crops_number)
+        ]
+            
+        for crops_image in crops_all:
+            output["source"].append(self.source_trans(crops_image))   #单独使用好一些
+    
+
+        output["offsets"] = ()
+
+        return output

clip/.ipynb_checkpoints/clip-checkpoint.py

@@ -0,0 +1,225 @@
+import hashlib
+import os
+import urllib
+import warnings
+from typing import Union, List
+
+import torch
+from PIL import Image
+from torchvision.transforms import Compose, Resize, CenterCrop, ToTensor, Normalize
+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 torch.__version__.split(".") < ["1", "7", "1"]:
+    warnings.warn("PyTorch version 1.7.1 or higher is recommended")
+
+
+__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",
+    "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",
+}
+
+
+def _download(url: str, root: str = os.path.expanduser("~/.cache/clip")):
+    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) 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(f"Model has been downloaded but the SHA256 checksum does not not match")
+
+    return download_target
+
+
+def _transform(n_px):
+    return Compose([
+        Resize(n_px, interpolation=BICUBIC),
+        CenterCrop(n_px),
+        lambda image: image.convert("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(name: str, device: Union[str, torch.device] = "cuda" if torch.cuda.is_available() else "cpu", jit=False):
+    """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).
+
+    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
+    """
+    '''
+    if name in _MODELS:
+        model_path = _download(_MODELS[name])
+    elif os.path.isfile(name):
+        model_path = name
+    else:
+        raise RuntimeError(f"Model {name} not found; available models = {available_models()}")
+    '''
+    model_path = '/model/4DaiRui/pretrained_ood/ViT-B-16.pt'
+    
+
+    try:
+        # loading JIT archive
+        model = torch.jit.load(model_path, map_location=device if jit else "cpu").eval()
+        state_dict = None
+    except RuntimeError:
+        # loading saved state dict
+        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(model_path, map_location="cpu")
+
+    if not jit:
+        model = build_model(state_dict or model.state_dict()).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(texts: Union[str, List[str]], context_length: int = 77, truncate: bool = False) -> 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]
+    """
+    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 = torch.zeros(len(all_tokens), context_length, dtype=torch.long)
+
+    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

clip/.ipynb_checkpoints/model-checkpoint.py

@@ -0,0 +1,432 @@
+from collections import OrderedDict
+from typing import Tuple, Union
+
+import numpy as np
+import torch
+import torch.nn.functional as F
+from torch import nn
+
+
+class Bottleneck(nn.Module):
+    expansion = 4
+
+    def __init__(self, inplanes, planes, stride=1):
+        super().__init__()
+
+        # all conv layers have stride 1. an avgpool is performed after the second convolution when stride > 1
+        self.conv1 = nn.Conv2d(inplanes, planes, 1, bias=False)
+        self.bn1 = nn.BatchNorm2d(planes)
+
+        self.conv2 = nn.Conv2d(planes, planes, 3, padding=1, bias=False)
+        self.bn2 = nn.BatchNorm2d(planes)
+
+        self.avgpool = nn.AvgPool2d(stride) if stride > 1 else nn.Identity()
+
+        self.conv3 = nn.Conv2d(planes, planes * self.expansion, 1, bias=False)
+        self.bn3 = nn.BatchNorm2d(planes * self.expansion)
+
+        self.relu = nn.ReLU(inplace=True)
+        self.downsample = None
+        self.stride = stride
+
+        if stride > 1 or inplanes != planes * Bottleneck.expansion:
+            # downsampling layer is prepended with an avgpool, and the subsequent convolution has stride 1
+            self.downsample = nn.Sequential(OrderedDict([
+                ("-1", nn.AvgPool2d(stride)),
+                ("0", nn.Conv2d(inplanes, planes * self.expansion, 1, stride=1, bias=False)),
+                ("1", nn.BatchNorm2d(planes * self.expansion))
+            ]))
+
+    def forward(self, x: torch.Tensor):
+        identity = x
+
+        out = self.relu(self.bn1(self.conv1(x)))
+        out = self.relu(self.bn2(self.conv2(out)))
+        out = self.avgpool(out)
+        out = self.bn3(self.conv3(out))
+
+        if self.downsample is not None:
+            identity = self.downsample(x)
+
+        out += identity
+        out = self.relu(out)
+        return out
+
+
+class AttentionPool2d(nn.Module):
+    def __init__(self, spacial_dim: int, embed_dim: int, num_heads: int, output_dim: int = None):
+        super().__init__()
+        self.positional_embedding = nn.Parameter(torch.randn(spacial_dim ** 2 + 1, embed_dim) / embed_dim ** 0.5)
+        self.k_proj = nn.Linear(embed_dim, embed_dim)
+        self.q_proj = nn.Linear(embed_dim, embed_dim)
+        self.v_proj = nn.Linear(embed_dim, embed_dim)
+        self.c_proj = nn.Linear(embed_dim, output_dim or embed_dim)
+        self.num_heads = num_heads
+
+    def forward(self, x):
+        x = x.reshape(x.shape[0], x.shape[1], x.shape[2] * x.shape[3]).permute(2, 0, 1)  # NCHW -> (HW)NC
+        x = torch.cat([x.mean(dim=0, keepdim=True), x], dim=0)  # (HW+1)NC
+        x = x + self.positional_embedding[:, None, :].to(x.dtype)  # (HW+1)NC
+        x, _ = F.multi_head_attention_forward(
+            query=x, key=x, value=x,
+            embed_dim_to_check=x.shape[-1],
+            num_heads=self.num_heads,
+            q_proj_weight=self.q_proj.weight,
+            k_proj_weight=self.k_proj.weight,
+            v_proj_weight=self.v_proj.weight,
+            in_proj_weight=None,
+            in_proj_bias=torch.cat([self.q_proj.bias, self.k_proj.bias, self.v_proj.bias]),
+            bias_k=None,
+            bias_v=None,
+            add_zero_attn=False,
+            dropout_p=0,
+            out_proj_weight=self.c_proj.weight,
+            out_proj_bias=self.c_proj.bias,
+            use_separate_proj_weight=True,
+            training=self.training,
+            need_weights=False
+        )
+
+        return x[0]
+
+
+class ModifiedResNet(nn.Module):
+    """
+    A ResNet class that is similar to torchvision's but contains the following changes:
+    - There are now 3 "stem" convolutions as opposed to 1, with an average pool instead of a max pool.
+    - Performs anti-aliasing strided convolutions, where an avgpool is prepended to convolutions with stride > 1
+    - The final pooling layer is a QKV attention instead of an average pool
+    """
+
+    def __init__(self, layers, output_dim, heads, input_resolution=224, width=64):
+        super().__init__()
+        self.output_dim = output_dim
+        self.input_resolution = input_resolution
+
+        # the 3-layer stem
+        self.conv1 = nn.Conv2d(3, width // 2, kernel_size=3, stride=2, padding=1, bias=False)
+        self.bn1 = nn.BatchNorm2d(width // 2)
+        self.conv2 = nn.Conv2d(width // 2, width // 2, kernel_size=3, padding=1, bias=False)
+        self.bn2 = nn.BatchNorm2d(width // 2)
+        self.conv3 = nn.Conv2d(width // 2, width, kernel_size=3, padding=1, bias=False)
+        self.bn3 = nn.BatchNorm2d(width)
+        self.avgpool = nn.AvgPool2d(2)
+        self.relu = nn.ReLU(inplace=True)
+
+        # residual layers
+        self._inplanes = width  # this is a *mutable* variable used during construction
+        self.layer1 = self._make_layer(width, layers[0])
+        self.layer2 = self._make_layer(width * 2, layers[1], stride=2)
+        self.layer3 = self._make_layer(width * 4, layers[2], stride=2)
+        self.layer4 = self._make_layer(width * 8, layers[3], stride=2)
+
+        embed_dim = width * 32  # the ResNet feature dimension
+        self.attnpool = AttentionPool2d(input_resolution // 32, embed_dim, heads, output_dim)
+
+    def _make_layer(self, planes, blocks, stride=1):
+        layers = [Bottleneck(self._inplanes, planes, stride)]
+
+        self._inplanes = planes * Bottleneck.expansion
+        for _ in range(1, blocks):
+            layers.append(Bottleneck(self._inplanes, planes))
+
+        return nn.Sequential(*layers)
+
+    def forward(self, x):
+        def stem(x):
+            for conv, bn in [(self.conv1, self.bn1), (self.conv2, self.bn2), (self.conv3, self.bn3)]:
+                x = self.relu(bn(conv(x)))
+            x = self.avgpool(x)
+            return x
+
+        x = x.type(self.conv1.weight.dtype)
+        x = stem(x)
+        x = self.layer1(x)
+        x = self.layer2(x)
+        x = self.layer3(x)
+        x = self.layer4(x)
+        x = self.attnpool(x)
+
+        return x
+
+
+class LayerNorm(nn.LayerNorm):
+    """Subclass torch's LayerNorm to handle fp16."""
+
+    def forward(self, x: torch.Tensor):
+        orig_type = x.dtype
+        ret = super().forward(x.type(torch.float32))
+        return ret.type(orig_type)
+
+
+class QuickGELU(nn.Module):
+    def forward(self, x: torch.Tensor):
+        return x * torch.sigmoid(1.702 * x)
+
+
+class ResidualAttentionBlock(nn.Module):
+    def __init__(self, d_model: int, n_head: int, attn_mask: torch.Tensor = None):
+        super().__init__()
+
+        self.attn = nn.MultiheadAttention(d_model, n_head)
+        self.ln_1 = LayerNorm(d_model)
+        self.mlp = nn.Sequential(OrderedDict([
+            ("c_fc", nn.Linear(d_model, d_model * 4)),
+            ("gelu", QuickGELU()),
+            ("c_proj", nn.Linear(d_model * 4, d_model))
+        ]))
+        self.ln_2 = LayerNorm(d_model)
+        self.attn_mask = attn_mask
+
+    def attention(self, x: torch.Tensor):
+        self.attn_mask = self.attn_mask.to(dtype=x.dtype, device=x.device) if self.attn_mask is not None else None
+        return self.attn(x, x, x, need_weights=False, attn_mask=self.attn_mask)[0]
+
+    def forward(self, x: torch.Tensor):
+        x = x + self.attention(self.ln_1(x))
+        x = x + self.mlp(self.ln_2(x))
+        return x
+
+
+class Transformer(nn.Module):
+    def __init__(self, width: int, layers: int, heads: int, attn_mask: torch.Tensor = None):
+        super().__init__()
+        self.width = width
+        self.layers = layers
+        self.resblocks = nn.Sequential(*[ResidualAttentionBlock(width, heads, attn_mask) for _ in range(layers)])
+
+    def forward(self, x: torch.Tensor):
+        return self.resblocks(x)
+
+
+class VisionTransformer(nn.Module):
+    def __init__(self, input_resolution: int, patch_size: int, width: int, layers: int, heads: int, output_dim: int):
+        super().__init__()
+        self.input_resolution = input_resolution
+        self.output_dim = output_dim
+        self.conv1 = nn.Conv2d(in_channels=3, out_channels=width, kernel_size=patch_size, stride=patch_size, bias=False)
+
+        scale = width ** -0.5
+        self.class_embedding = nn.Parameter(scale * torch.randn(width))
+        self.positional_embedding = nn.Parameter(scale * torch.randn((input_resolution // patch_size) ** 2 + 1, width))
+        self.ln_pre = LayerNorm(width)
+
+        self.transformer = Transformer(width, layers, heads)
+
+        self.ln_post = LayerNorm(width)
+        self.proj = nn.Parameter(scale * torch.randn(width, output_dim))
+
+    def forward(self, x: torch.Tensor):
+        x = self.conv1(x)  # shape = [*, width, grid, grid]
+        x = x.reshape(x.shape[0], x.shape[1], -1)  # shape = [*, width, grid ** 2]
+        x = x.permute(0, 2, 1)  # shape = [*, grid ** 2, width]
+        x = torch.cat([self.class_embedding.to(x.dtype) + torch.zeros(x.shape[0], 1, x.shape[-1], dtype=x.dtype, device=x.device), x], dim=1)  # shape = [*, grid ** 2 + 1, width]
+        x = x + self.positional_embedding.to(x.dtype)
+        x = self.ln_pre(x)
+
+        x = x.permute(1, 0, 2)  # NLD -> LND
+        x = self.transformer(x)
+        x = x.permute(1, 0, 2)  # LND -> NLD
+
+        x = self.ln_post(x[:, 0, :])
+
+        if self.proj is not None:
+            x = x @ self.proj
+
+        return x
+
+
+class CLIP(nn.Module):
+    def __init__(self,
+                 embed_dim: int,
+                 # vision
+                 image_resolution: int,
+                 vision_layers: Union[Tuple[int, int, int, int], int],
+                 vision_width: int,
+                 vision_patch_size: int,
+                 # text
+                 context_length: int,
+                 vocab_size: int,
+                 transformer_width: int,
+                 transformer_heads: int,
+                 transformer_layers: int
+                 ):
+        super().__init__()
+
+        self.context_length = context_length
+
+        if isinstance(vision_layers, (tuple, list)):
+            vision_heads = vision_width * 32 // 64
+            self.visual = ModifiedResNet(
+                layers=vision_layers,
+                output_dim=embed_dim,
+                heads=vision_heads,
+                input_resolution=image_resolution,
+                width=vision_width
+            )
+        else:
+            vision_heads = vision_width // 64
+            self.visual = VisionTransformer(
+                input_resolution=image_resolution,
+                patch_size=vision_patch_size,
+                width=vision_width,
+                layers=vision_layers,
+                heads=vision_heads,
+                output_dim=embed_dim
+            )
+
+        self.transformer = Transformer(
+            width=transformer_width,
+            layers=transformer_layers,
+            heads=transformer_heads,
+            attn_mask=self.build_attention_mask()
+        )
+
+        self.vocab_size = vocab_size
+        self.token_embedding = nn.Embedding(vocab_size, transformer_width)
+        self.positional_embedding = nn.Parameter(torch.empty(self.context_length, transformer_width))
+        self.ln_final = LayerNorm(transformer_width)
+
+        self.text_projection = nn.Parameter(torch.empty(transformer_width, embed_dim))
+        self.logit_scale = nn.Parameter(torch.ones([]) * np.log(1 / 0.07))
+
+        self.initialize_parameters()
+
+    def initialize_parameters(self):
+        nn.init.normal_(self.token_embedding.weight, std=0.02)
+        nn.init.normal_(self.positional_embedding, std=0.01)
+
+        if isinstance(self.visual, ModifiedResNet):
+            if self.visual.attnpool is not None:
+                std = self.visual.attnpool.c_proj.in_features ** -0.5
+                nn.init.normal_(self.visual.attnpool.q_proj.weight, std=std)
+                nn.init.normal_(self.visual.attnpool.k_proj.weight, std=std)
+                nn.init.normal_(self.visual.attnpool.v_proj.weight, std=std)
+                nn.init.normal_(self.visual.attnpool.c_proj.weight, std=std)
+
+            for resnet_block in [self.visual.layer1, self.visual.layer2, self.visual.layer3, self.visual.layer4]:
+                for name, param in resnet_block.named_parameters():
+                    if name.endswith("bn3.weight"):
+                        nn.init.zeros_(param)
+
+        proj_std = (self.transformer.width ** -0.5) * ((2 * self.transformer.layers) ** -0.5)
+        attn_std = self.transformer.width ** -0.5
+        fc_std = (2 * self.transformer.width) ** -0.5
+        for block in self.transformer.resblocks:
+            nn.init.normal_(block.attn.in_proj_weight, std=attn_std)
+            nn.init.normal_(block.attn.out_proj.weight, std=proj_std)
+            nn.init.normal_(block.mlp.c_fc.weight, std=fc_std)
+            nn.init.normal_(block.mlp.c_proj.weight, std=proj_std)
+
+        if self.text_projection is not None:
+            nn.init.normal_(self.text_projection, std=self.transformer.width ** -0.5)
+
+    def build_attention_mask(self):
+        # lazily create causal attention mask, with full attention between the vision tokens
+        # pytorch uses additive attention mask; fill with -inf
+        mask = torch.empty(self.context_length, self.context_length)
+        mask.fill_(float("-inf"))
+        mask.triu_(1)  # zero out the lower diagonal
+        return mask
+
+    @property
+    def dtype(self):
+        return self.visual.conv1.weight.dtype
+
+    def encode_image(self, image):
+        return self.visual(image.type(self.dtype))
+
+    def encode_text(self, text):
+        x = self.token_embedding(text).type(self.dtype)  # [batch_size, n_ctx, d_model]
+
+        x = x + self.positional_embedding.type(self.dtype)
+        x = x.permute(1, 0, 2)  # NLD -> LND
+        x = self.transformer(x)
+        x = x.permute(1, 0, 2)  # LND -> NLD
+        x = self.ln_final(x).type(self.dtype)
+
+        # x.shape = [batch_size, n_ctx, transformer.width]
+        # take features from the eot embedding (eot_token is the highest number in each sequence)
+        x = x[torch.arange(x.shape[0]), text.argmax(dim=-1)] @ self.text_projection
+
+        return x
+
+    def forward(self, image, text):
+        image_features = self.encode_image(image)
+        text_features = self.encode_text(text)
+
+        # normalized features
+        image_features = image_features / image_features.norm(dim=-1, keepdim=True)
+        text_features = text_features / text_features.norm(dim=-1, keepdim=True)
+
+        # cosine similarity as logits
+        logit_scale = self.logit_scale.exp()
+        logits_per_image = logit_scale * image_features @ text_features.t()
+        logits_per_text = logit_scale * text_features @ image_features.t()
+
+        # shape = [global_batch_size, global_batch_size]
+        return logits_per_image, logits_per_text
+
+
+def convert_weights(model: nn.Module):
+    """Convert applicable model parameters to fp16"""
+
+    def _convert_weights_to_fp16(l):
+        if isinstance(l, (nn.Conv1d, nn.Conv2d, nn.Linear)):
+            l.weight.data = l.weight.data.half()
+            if l.bias is not None:
+                l.bias.data = l.bias.data.half()
+
+        if isinstance(l, nn.MultiheadAttention):
+            for attr in [*[f"{s}_proj_weight" for s in ["in", "q", "k", "v"]], "in_proj_bias", "bias_k", "bias_v"]:
+                tensor = getattr(l, attr)
+                if tensor is not None:
+                    tensor.data = tensor.data.half()
+
+        for name in ["text_projection", "proj"]:
+            if hasattr(l, name):
+                attr = getattr(l, name)
+                if attr is not None:
+                    attr.data = attr.data.half()
+
+    model.apply(_convert_weights_to_fp16)
+
+
+def build_model(state_dict: dict):
+    vit = "visual.proj" in state_dict
+
+    if vit:
+        vision_width = state_dict["visual.conv1.weight"].shape[0]
+        vision_layers = len([k for k in state_dict.keys() if k.startswith("visual.") and k.endswith(".attn.in_proj_weight")])
+        vision_patch_size = state_dict["visual.conv1.weight"].shape[-1]
+        grid_size = round((state_dict["visual.positional_embedding"].shape[0] - 1) ** 0.5)
+        image_resolution = vision_patch_size * grid_size
+    else:
+        counts: list = [len(set(k.split(".")[2] for k in state_dict if k.startswith(f"visual.layer{b}"))) for b in [1, 2, 3, 4]]
+        vision_layers = tuple(counts)
+        vision_width = state_dict["visual.layer1.0.conv1.weight"].shape[0]
+        output_width = round((state_dict["visual.attnpool.positional_embedding"].shape[0] - 1) ** 0.5)
+        vision_patch_size = None
+        assert output_width ** 2 + 1 == state_dict["visual.attnpool.positional_embedding"].shape[0]
+        image_resolution = output_width * 32
+
+    embed_dim = state_dict["text_projection"].shape[1]
+    context_length = state_dict["positional_embedding"].shape[0]
+    vocab_size = state_dict["token_embedding.weight"].shape[0]
+    transformer_width = state_dict["ln_final.weight"].shape[0]
+    transformer_heads = transformer_width // 64
+    transformer_layers = len(set(k.split(".")[2] for k in state_dict if k.startswith(f"transformer.resblocks")))
+
+    model = CLIP(
+        embed_dim,
+        image_resolution, vision_layers, vision_width, vision_patch_size,
+        context_length, vocab_size, transformer_width, transformer_heads, transformer_layers
+    )
+
+    for key in ["input_resolution", "context_length", "vocab_size"]:
+        if key in state_dict:
+            del state_dict[key]
+
+    convert_weights(model)
+    model.load_state_dict(state_dict)
+    return model.eval()

clip/__init__.py

@@ -0,0 +1 @@
+from .clip import *

clip/bpe_simple_vocab_16e6.txt.gz

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:924691ac288e54409236115652ad4aa250f48203de50a9e4722a6ecd48d6804a
+size 1356917

clip/clip.py

@@ -0,0 +1,228 @@
+import hashlib
+import os
+import urllib
+import warnings
+from typing import Union, List
+
+import torch
+from PIL import Image
+from torchvision.transforms import Compose, Resize, CenterCrop, ToTensor, Normalize
+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 torch.__version__.split(".") < ["1", "7", "1"]:
+    warnings.warn("PyTorch version 1.7.1 or higher is recommended")
+
+
+__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",
+    "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.path.expanduser("~/.cache/clip")):
+    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) 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(f"Model has been downloaded but the SHA256 checksum does not not match")
+
+    return download_target
+
+
+def _transform(n_px):
+    return Compose([
+        Resize(n_px, interpolation=BICUBIC),
+        CenterCrop(n_px),
+        lambda image: image.convert("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(name: str, device: Union[str, torch.device] = "cuda" if torch.cuda.is_available() else "cpu", jit=False):
+    """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).
+
+    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
+    """
+    
+    if name in _MODELS:
+        model_path = _download(_MODELS[name])
+    elif os.path.isfile(name):
+        model_path = name
+    else:
+        raise RuntimeError(f"Model {name} not found; available models = {available_models()}")
+    
+    #model_path = 'E:/code/lsn/clip/RN50.pt'
+    # model_path = 'E:/code/lsn/clip/ViT-B-16.pt'
+
+  
+    try:
+        # loading JIT archive
+        model = torch.jit.load(model_path, map_location=device if jit else "cpu").eval()
+        state_dict = None
+    except RuntimeError:
+        # loading saved state dict
+        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(model_path, map_location="cpu")
+
+    if not jit:
+        model = build_model(state_dict or model.state_dict()).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(texts: Union[str, List[str]], context_length: int = 77, truncate: bool = False) -> 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]
+    """
+    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 = torch.zeros(len(all_tokens), context_length, dtype=torch.long)
+
+    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

clip/model.py

@@ -0,0 +1,432 @@
+from collections import OrderedDict
+from typing import Tuple, Union
+
+import numpy as np
+import torch
+import torch.nn.functional as F
+from torch import nn
+
+
+class Bottleneck(nn.Module):
+    expansion = 4
+
+    def __init__(self, inplanes, planes, stride=1):
+        super().__init__()
+
+        # all conv layers have stride 1. an avgpool is performed after the second convolution when stride > 1
+        self.conv1 = nn.Conv2d(inplanes, planes, 1, bias=False)
+        self.bn1 = nn.BatchNorm2d(planes)
+
+        self.conv2 = nn.Conv2d(planes, planes, 3, padding=1, bias=False)
+        self.bn2 = nn.BatchNorm2d(planes)
+
+        self.avgpool = nn.AvgPool2d(stride) if stride > 1 else nn.Identity()
+
+        self.conv3 = nn.Conv2d(planes, planes * self.expansion, 1, bias=False)
+        self.bn3 = nn.BatchNorm2d(planes * self.expansion)
+
+        self.relu = nn.ReLU(inplace=True)
+        self.downsample = None
+        self.stride = stride
+
+        if stride > 1 or inplanes != planes * Bottleneck.expansion:
+            # downsampling layer is prepended with an avgpool, and the subsequent convolution has stride 1
+            self.downsample = nn.Sequential(OrderedDict([
+                ("-1", nn.AvgPool2d(stride)),
+                ("0", nn.Conv2d(inplanes, planes * self.expansion, 1, stride=1, bias=False)),
+                ("1", nn.BatchNorm2d(planes * self.expansion))
+            ]))
+
+    def forward(self, x: torch.Tensor):
+        identity = x
+
+        out = self.relu(self.bn1(self.conv1(x)))
+        out = self.relu(self.bn2(self.conv2(out)))
+        out = self.avgpool(out)
+        out = self.bn3(self.conv3(out))
+
+        if self.downsample is not None:
+            identity = self.downsample(x)
+
+        out += identity
+        out = self.relu(out)
+        return out
+
+
+class AttentionPool2d(nn.Module):
+    def __init__(self, spacial_dim: int, embed_dim: int, num_heads: int, output_dim: int = None):
+        super().__init__()
+        self.positional_embedding = nn.Parameter(torch.randn(spacial_dim ** 2 + 1, embed_dim) / embed_dim ** 0.5)
+        self.k_proj = nn.Linear(embed_dim, embed_dim)
+        self.q_proj = nn.Linear(embed_dim, embed_dim)
+        self.v_proj = nn.Linear(embed_dim, embed_dim)
+        self.c_proj = nn.Linear(embed_dim, output_dim or embed_dim)
+        self.num_heads = num_heads
+
+    def forward(self, x):
+        x = x.reshape(x.shape[0], x.shape[1], x.shape[2] * x.shape[3]).permute(2, 0, 1)  # NCHW -> (HW)NC
+        x = torch.cat([x.mean(dim=0, keepdim=True), x], dim=0)  # (HW+1)NC
+        x = x + self.positional_embedding[:, None, :].to(x.dtype)  # (HW+1)NC
+        x, _ = F.multi_head_attention_forward(
+            query=x, key=x, value=x,
+            embed_dim_to_check=x.shape[-1],
+            num_heads=self.num_heads,
+            q_proj_weight=self.q_proj.weight,
+            k_proj_weight=self.k_proj.weight,
+            v_proj_weight=self.v_proj.weight,
+            in_proj_weight=None,
+            in_proj_bias=torch.cat([self.q_proj.bias, self.k_proj.bias, self.v_proj.bias]),
+            bias_k=None,
+            bias_v=None,
+            add_zero_attn=False,
+            dropout_p=0,
+            out_proj_weight=self.c_proj.weight,
+            out_proj_bias=self.c_proj.bias,
+            use_separate_proj_weight=True,
+            training=self.training,
+            need_weights=False
+        )
+
+        return x[0]
+
+
+class ModifiedResNet(nn.Module):
+    """
+    A ResNet class that is similar to torchvision's but contains the following changes:
+    - There are now 3 "stem" convolutions as opposed to 1, with an average pool instead of a max pool.
+    - Performs anti-aliasing strided convolutions, where an avgpool is prepended to convolutions with stride > 1
+    - The final pooling layer is a QKV attention instead of an average pool
+    """
+
+    def __init__(self, layers, output_dim, heads, input_resolution=224, width=64):
+        super().__init__()
+        self.output_dim = output_dim
+        self.input_resolution = input_resolution
+
+        # the 3-layer stem
+        self.conv1 = nn.Conv2d(3, width // 2, kernel_size=3, stride=2, padding=1, bias=False)
+        self.bn1 = nn.BatchNorm2d(width // 2)
+        self.conv2 = nn.Conv2d(width // 2, width // 2, kernel_size=3, padding=1, bias=False)
+        self.bn2 = nn.BatchNorm2d(width // 2)
+        self.conv3 = nn.Conv2d(width // 2, width, kernel_size=3, padding=1, bias=False)
+        self.bn3 = nn.BatchNorm2d(width)
+        self.avgpool = nn.AvgPool2d(2)
+        self.relu = nn.ReLU(inplace=True)
+
+        # residual layers
+        self._inplanes = width  # this is a *mutable* variable used during construction
+        self.layer1 = self._make_layer(width, layers[0])
+        self.layer2 = self._make_layer(width * 2, layers[1], stride=2)
+        self.layer3 = self._make_layer(width * 4, layers[2], stride=2)
+        self.layer4 = self._make_layer(width * 8, layers[3], stride=2)
+
+        embed_dim = width * 32  # the ResNet feature dimension
+        self.attnpool = AttentionPool2d(input_resolution // 32, embed_dim, heads, output_dim)
+
+    def _make_layer(self, planes, blocks, stride=1):
+        layers = [Bottleneck(self._inplanes, planes, stride)]
+
+        self._inplanes = planes * Bottleneck.expansion
+        for _ in range(1, blocks):
+            layers.append(Bottleneck(self._inplanes, planes))
+
+        return nn.Sequential(*layers)
+
+    def forward(self, x):
+        def stem(x):
+            for conv, bn in [(self.conv1, self.bn1), (self.conv2, self.bn2), (self.conv3, self.bn3)]:
+                x = self.relu(bn(conv(x)))
+            x = self.avgpool(x)
+            return x
+
+        x = x.type(self.conv1.weight.dtype)
+        x = stem(x)
+        x = self.layer1(x)
+        x = self.layer2(x)
+        x = self.layer3(x)
+        x = self.layer4(x)
+        x = self.attnpool(x)
+
+        return x
+
+
+class LayerNorm(nn.LayerNorm):
+    """Subclass torch's LayerNorm to handle fp16."""
+
+    def forward(self, x: torch.Tensor):
+        orig_type = x.dtype
+        ret = super().forward(x.type(torch.float32))
+        return ret.type(orig_type)
+
+
+class QuickGELU(nn.Module):
+    def forward(self, x: torch.Tensor):
+        return x * torch.sigmoid(1.702 * x)
+
+
+class ResidualAttentionBlock(nn.Module):
+    def __init__(self, d_model: int, n_head: int, attn_mask: torch.Tensor = None):
+        super().__init__()
+
+        self.attn = nn.MultiheadAttention(d_model, n_head)
+        self.ln_1 = LayerNorm(d_model)
+        self.mlp = nn.Sequential(OrderedDict([
+            ("c_fc", nn.Linear(d_model, d_model * 4)),
+            ("gelu", QuickGELU()),
+            ("c_proj", nn.Linear(d_model * 4, d_model))
+        ]))
+        self.ln_2 = LayerNorm(d_model)
+        self.attn_mask = attn_mask
+
+    def attention(self, x: torch.Tensor):
+        self.attn_mask = self.attn_mask.to(dtype=x.dtype, device=x.device) if self.attn_mask is not None else None
+        return self.attn(x, x, x, need_weights=False, attn_mask=self.attn_mask)[0]
+
+    def forward(self, x: torch.Tensor):
+        x = x + self.attention(self.ln_1(x))
+        x = x + self.mlp(self.ln_2(x))
+        return x
+
+
+class Transformer(nn.Module):
+    def __init__(self, width: int, layers: int, heads: int, attn_mask: torch.Tensor = None):
+        super().__init__()
+        self.width = width
+        self.layers = layers
+        self.resblocks = nn.Sequential(*[ResidualAttentionBlock(width, heads, attn_mask) for _ in range(layers)])
+
+    def forward(self, x: torch.Tensor):
+        return self.resblocks(x)
+
+
+class VisionTransformer(nn.Module):
+    def __init__(self, input_resolution: int, patch_size: int, width: int, layers: int, heads: int, output_dim: int):
+        super().__init__()
+        self.input_resolution = input_resolution
+        self.output_dim = output_dim
+        self.conv1 = nn.Conv2d(in_channels=3, out_channels=width, kernel_size=patch_size, stride=patch_size, bias=False)
+
+        scale = width ** -0.5
+        self.class_embedding = nn.Parameter(scale * torch.randn(width))
+        self.positional_embedding = nn.Parameter(scale * torch.randn((input_resolution // patch_size) ** 2 + 1, width))
+        self.ln_pre = LayerNorm(width)
+
+        self.transformer = Transformer(width, layers, heads)
+
+        self.ln_post = LayerNorm(width)
+        self.proj = nn.Parameter(scale * torch.randn(width, output_dim))
+
+    def forward(self, x: torch.Tensor):
+        x = self.conv1(x)  # shape = [*, width, grid, grid]
+        x = x.reshape(x.shape[0], x.shape[1], -1)  # shape = [*, width, grid ** 2]
+        x = x.permute(0, 2, 1)  # shape = [*, grid ** 2, width]
+        x = torch.cat([self.class_embedding.to(x.dtype) + torch.zeros(x.shape[0], 1, x.shape[-1], dtype=x.dtype, device=x.device), x], dim=1)  # shape = [*, grid ** 2 + 1, width]
+        x = x + self.positional_embedding.to(x.dtype)
+        x = self.ln_pre(x)
+
+        x = x.permute(1, 0, 2)  # NLD -> LND
+        x = self.transformer(x)
+        x = x.permute(1, 0, 2)  # LND -> NLD
+
+        x = self.ln_post(x[:, 0, :])
+
+        # if self.proj is not None:
+        #     x = x @ self.proj
+
+        return x
+
+
+class CLIP(nn.Module):
+    def __init__(self,
+                 embed_dim: int,
+                 # vision
+                 image_resolution: int,
+                 vision_layers: Union[Tuple[int, int, int, int], int],
+                 vision_width: int,
+                 vision_patch_size: int,
+                 # text
+                 context_length: int,
+                 vocab_size: int,
+                 transformer_width: int,
+                 transformer_heads: int,
+                 transformer_layers: int
+                 ):
+        super().__init__()
+
+        self.context_length = context_length
+
+        if isinstance(vision_layers, (tuple, list)):
+            vision_heads = vision_width * 32 // 64
+            self.visual = ModifiedResNet(
+                layers=vision_layers,
+                output_dim=embed_dim,
+                heads=vision_heads,
+                input_resolution=image_resolution,
+                width=vision_width
+            )
+        else:
+            vision_heads = vision_width // 64
+            self.visual = VisionTransformer(
+                input_resolution=image_resolution,
+                patch_size=vision_patch_size,
+                width=vision_width,
+                layers=vision_layers,
+                heads=vision_heads,
+                output_dim=embed_dim
+            )
+
+        self.transformer = Transformer(
+            width=transformer_width,
+            layers=transformer_layers,
+            heads=transformer_heads,
+            attn_mask=self.build_attention_mask()
+        )
+
+        self.vocab_size = vocab_size
+        self.token_embedding = nn.Embedding(vocab_size, transformer_width)
+        self.positional_embedding = nn.Parameter(torch.empty(self.context_length, transformer_width))
+        self.ln_final = LayerNorm(transformer_width)
+
+        self.text_projection = nn.Parameter(torch.empty(transformer_width, embed_dim))
+        self.logit_scale = nn.Parameter(torch.ones([]) * np.log(1 / 0.07))
+
+        self.initialize_parameters()
+
+    def initialize_parameters(self):
+        nn.init.normal_(self.token_embedding.weight, std=0.02)
+        nn.init.normal_(self.positional_embedding, std=0.01)
+
+        if isinstance(self.visual, ModifiedResNet):
+            if self.visual.attnpool is not None:
+                std = self.visual.attnpool.c_proj.in_features ** -0.5
+                nn.init.normal_(self.visual.attnpool.q_proj.weight, std=std)
+                nn.init.normal_(self.visual.attnpool.k_proj.weight, std=std)
+                nn.init.normal_(self.visual.attnpool.v_proj.weight, std=std)
+                nn.init.normal_(self.visual.attnpool.c_proj.weight, std=std)
+
+            for resnet_block in [self.visual.layer1, self.visual.layer2, self.visual.layer3, self.visual.layer4]:
+                for name, param in resnet_block.named_parameters():
+                    if name.endswith("bn3.weight"):
+                        nn.init.zeros_(param)
+
+        proj_std = (self.transformer.width ** -0.5) * ((2 * self.transformer.layers) ** -0.5)
+        attn_std = self.transformer.width ** -0.5
+        fc_std = (2 * self.transformer.width) ** -0.5
+        for block in self.transformer.resblocks:
+            nn.init.normal_(block.attn.in_proj_weight, std=attn_std)
+            nn.init.normal_(block.attn.out_proj.weight, std=proj_std)
+            nn.init.normal_(block.mlp.c_fc.weight, std=fc_std)
+            nn.init.normal_(block.mlp.c_proj.weight, std=proj_std)
+
+        if self.text_projection is not None:
+            nn.init.normal_(self.text_projection, std=self.transformer.width ** -0.5)
+
+    def build_attention_mask(self):
+        # lazily create causal attention mask, with full attention between the vision tokens
+        # pytorch uses additive attention mask; fill with -inf
+        mask = torch.empty(self.context_length, self.context_length)
+        mask.fill_(float("-inf"))
+        mask.triu_(1)  # zero out the lower diagonal
+        return mask
+
+    @property
+    def dtype(self):
+        return self.visual.conv1.weight.dtype
+
+    def encode_image(self, image):
+        return self.visual(image.type(self.dtype))
+
+    def encode_text(self, text):
+        x = self.token_embedding(text).type(self.dtype)  # [batch_size, n_ctx, d_model]
+
+        x = x + self.positional_embedding.type(self.dtype)
+        x = x.permute(1, 0, 2)  # NLD -> LND
+        x = self.transformer(x)
+        x = x.permute(1, 0, 2)  # LND -> NLD
+        x = self.ln_final(x).type(self.dtype)
+
+        # x.shape = [batch_size, n_ctx, transformer.width]
+        # take features from the eot embedding (eot_token is the highest number in each sequence)
+        x = x[torch.arange(x.shape[0]), text.argmax(dim=-1)] @ self.text_projection
+
+        return x
+
+    def forward(self, image, text):
+        image_features = self.encode_image(image)
+        text_features = self.encode_text(text)
+
+        # normalized features
+        image_features = image_features / image_features.norm(dim=-1, keepdim=True)
+        text_features = text_features / text_features.norm(dim=-1, keepdim=True)
+
+        # cosine similarity as logits
+        logit_scale = self.logit_scale.exp()
+        logits_per_image = logit_scale * image_features @ text_features.t()
+        logits_per_text = logit_scale * text_features @ image_features.t()
+
+        # shape = [global_batch_size, global_batch_size]
+        return logits_per_image, logits_per_text
+
+
+def convert_weights(model: nn.Module):
+    """Convert applicable model parameters to fp16"""
+
+    def _convert_weights_to_fp16(l):
+        if isinstance(l, (nn.Conv1d, nn.Conv2d, nn.Linear)):
+            l.weight.data = l.weight.data.half()
+            if l.bias is not None:
+                l.bias.data = l.bias.data.half()
+
+        if isinstance(l, nn.MultiheadAttention):
+            for attr in [*[f"{s}_proj_weight" for s in ["in", "q", "k", "v"]], "in_proj_bias", "bias_k", "bias_v"]:
+                tensor = getattr(l, attr)
+                if tensor is not None:
+                    tensor.data = tensor.data.half()
+
+        for name in ["text_projection", "proj"]:
+            if hasattr(l, name):
+                attr = getattr(l, name)
+                if attr is not None:
+                    attr.data = attr.data.half()
+
+    model.apply(_convert_weights_to_fp16)
+
+
+def build_model(state_dict: dict):
+    vit = "visual.proj" in state_dict
+
+    if vit:
+        vision_width = state_dict["visual.conv1.weight"].shape[0]
+        vision_layers = len([k for k in state_dict.keys() if k.startswith("visual.") and k.endswith(".attn.in_proj_weight")])
+        vision_patch_size = state_dict["visual.conv1.weight"].shape[-1]
+        grid_size = round((state_dict["visual.positional_embedding"].shape[0] - 1) ** 0.5)
+        image_resolution = vision_patch_size * grid_size
+    else:
+        counts: list = [len(set(k.split(".")[2] for k in state_dict if k.startswith(f"visual.layer{b}"))) for b in [1, 2, 3, 4]]
+        vision_layers = tuple(counts)
+        vision_width = state_dict["visual.layer1.0.conv1.weight"].shape[0]
+        output_width = round((state_dict["visual.attnpool.positional_embedding"].shape[0] - 1) ** 0.5)
+        vision_patch_size = None
+        assert output_width ** 2 + 1 == state_dict["visual.attnpool.positional_embedding"].shape[0]
+        image_resolution = output_width * 32
+
+    embed_dim = state_dict["text_projection"].shape[1]
+    context_length = state_dict["positional_embedding"].shape[0]
+    vocab_size = state_dict["token_embedding.weight"].shape[0]
+    transformer_width = state_dict["ln_final.weight"].shape[0]
+    transformer_heads = transformer_width // 64
+    transformer_layers = len(set(k.split(".")[2] for k in state_dict if k.startswith(f"transformer.resblocks")))
+
+    model = CLIP(
+        embed_dim,
+        image_resolution, vision_layers, vision_width, vision_patch_size,
+        context_length, vocab_size, transformer_width, transformer_heads, transformer_layers
+    )
+
+    for key in ["input_resolution", "context_length", "vocab_size"]:
+        if key in state_dict:
+            del state_dict[key]
+
+    convert_weights(model)
+    model.load_state_dict(state_dict)
+    return model.eval()

clip/simple_tokenizer.py

@@ -0,0 +1,132 @@
+import gzip
+import html
+import os
+from functools import lru_cache
+
+import ftfy
+import regex as re
+
+
+@lru_cache()
+def default_bpe():
+    return os.path.join(os.path.dirname(os.path.abspath(__file__)), "bpe_simple_vocab_16e6.txt.gz")
+
+
+@lru_cache()
+def bytes_to_unicode():
+    """
+    Returns list of utf-8 byte and a corresponding list of unicode strings.
+    The reversible bpe codes work on unicode strings.
+    This means you need a large # of unicode characters in your vocab if you want to avoid UNKs.
+    When you're at something like a 10B token dataset you end up needing around 5K for decent coverage.
+    This is a signficant percentage of your normal, say, 32K bpe vocab.
+    To avoid that, we want lookup tables between utf-8 bytes and unicode strings.
+    And avoids mapping to whitespace/control characters the bpe code barfs on.
+    """
+    bs = list(range(ord("!"), ord("~")+1))+list(range(ord("¡"), ord("¬")+1))+list(range(ord("®"), ord("ÿ")+1))
+    cs = bs[:]
+    n = 0
+    for b in range(2**8):
+        if b not in bs:
+            bs.append(b)
+            cs.append(2**8+n)
+            n += 1
+    cs = [chr(n) for n in cs]
+    return dict(zip(bs, cs))
+
+
+def get_pairs(word):
+    """Return set of symbol pairs in a word.
+    Word is represented as tuple of symbols (symbols being variable-length strings).
+    """
+    pairs = set()
+    prev_char = word[0]
+    for char in word[1:]:
+        pairs.add((prev_char, char))
+        prev_char = char
+    return pairs
+
+
+def basic_clean(text):
+    text = ftfy.fix_text(text)
+    text = html.unescape(html.unescape(text))
+    return text.strip()
+
+
+def whitespace_clean(text):
+    text = re.sub(r'\s+', ' ', text)
+    text = text.strip()
+    return text
+
+
+class SimpleTokenizer(object):
+    def __init__(self, bpe_path: str = default_bpe()):
+        self.byte_encoder = bytes_to_unicode()
+        self.byte_decoder = {v: k for k, v in self.byte_encoder.items()}
+        merges = gzip.open(bpe_path).read().decode("utf-8").split('\n')
+        merges = merges[1:49152-256-2+1]
+        merges = [tuple(merge.split()) for merge in merges]
+        vocab = list(bytes_to_unicode().values())
+        vocab = vocab + [v+'</w>' for v in vocab]
+        for merge in merges:
+            vocab.append(''.join(merge))
+        vocab.extend(['<|startoftext|>', '<|endoftext|>'])
+        self.encoder = dict(zip(vocab, range(len(vocab))))
+        self.decoder = {v: k for k, v in self.encoder.items()}
+        self.bpe_ranks = dict(zip(merges, range(len(merges))))
+        self.cache = {'<|startoftext|>': '<|startoftext|>', '<|endoftext|>': '<|endoftext|>'}
+        self.pat = re.compile(r"""<\|startoftext\|>|<\|endoftext\|>|'s|'t|'re|'ve|'m|'ll|'d|[\p{L}]+|[\p{N}]|[^\s\p{L}\p{N}]+""", re.IGNORECASE)
+
+    def bpe(self, token):
+        if token in self.cache:
+            return self.cache[token]
+        word = tuple(token[:-1]) + ( token[-1] + '</w>',)
+        pairs = get_pairs(word)
+
+        if not pairs:
+            return token+'</w>'
+
+        while True:
+            bigram = min(pairs, key = lambda pair: self.bpe_ranks.get(pair, float('inf')))
+            if bigram not in self.bpe_ranks:
+                break
+            first, second = bigram
+            new_word = []
+            i = 0
+            while i < len(word):
+                try:
+                    j = word.index(first, i)
+                    new_word.extend(word[i:j])
+                    i = j
+                except:
+                    new_word.extend(word[i:])
+                    break
+
+                if word[i] == first and i < len(word)-1 and word[i+1] == second:
+                    new_word.append(first+second)
+                    i += 2
+                else:
+                    new_word.append(word[i])
+                    i += 1
+            new_word = tuple(new_word)
+            word = new_word
+            if len(word) == 1:
+                break
+            else:
+                pairs = get_pairs(word)
+        word = ' '.join(word)
+        self.cache[token] = word
+        return word
+
+    def encode(self, text):
+        bpe_tokens = []
+        text = whitespace_clean(basic_clean(text)).lower()
+        for token in re.findall(self.pat, text):
+            token = ''.join(self.byte_encoder[b] for b in token.encode('utf-8'))
+            bpe_tokens.extend(self.encoder[bpe_token] for bpe_token in self.bpe(token).split(' '))
+        return bpe_tokens
+
+    def decode(self, tokens):
+        text = ''.join([self.decoder[token] for token in tokens])
+        text = bytearray([self.byte_decoder[c] for c in text]).decode('utf-8', errors="replace").replace('</w>', ' ')
+        return text

config.py

@@ -0,0 +1,48 @@
+'''This file configures the training procedure because handling arguments in every single function is so exhaustive for
+research purposes. Don't try this code if you are a software engineer.'''
+
+# device settings
+device = 'cuda' # or 'cpu'
+import torch
+torch.cuda.set_device(0)
+
+# data settings
+dataset_path = "dummy_dataset"
+class_name = "dummy_class"
+modelname = "dummy_test"
+
+img_size = (448, 448)
+img_dims = [3] + list(img_size)
+
+# transformation settings
+transf_rotations = True
+transf_brightness = 0.0
+transf_contrast = 0.0
+transf_saturation = 0.0
+norm_mean, norm_std = [0.485, 0.456, 0.406], [0.229, 0.224, 0.225]
+
+# network hyperparameters
+n_scales = 3 # number of scales at which features are extracted, img_size is the highest - others are //2, //4,...
+clamp_alpha = 3 # see paper equation 2 for explanation
+n_coupling_blocks = 2
+fc_internal = 4096 # number of neurons in hidden layers of s-t-networks
+dropout = 0# dropout in s-t-networks
+lr_init = 2e-4
+n_feat = 256 * n_scales # do not change except you change the feature extractor
+
+# dataloader parameters
+n_transforms = 4 # number of transformations per sample in training
+n_transforms_test = 64 # number of transformations per sample in testing
+batch_size = 24 # actual batch size is this value multiplied by n_transforms(_test)
+batch_size_test = batch_size * n_transforms // n_transforms_test
+
+# total epochs = meta_epochs * sub_epochs
+# evaluation after <sub_epochs> epochs
+meta_epochs = 24
+sub_epochs = 8
+
+# output settings
+verbose = True
+grad_map_viz = False
+hide_tqdm_bar = True
+save_model = True

detector.py

@@ -0,0 +1,86 @@
+
+import torch
+import clip  
+from PIL import Image
+from torch.cuda.amp import autocast as autocast
+from huggingface_hub import hf_hub_download 
+import spaces 
+
+from model import flow_model 
+from augmentations_clip import DataAugmentationCLIP as DataAugmentationCLIP_test
+
+MODEL_REPO_ID = "davjoython/flow_fake"
+FLOW_MODEL_FILENAME = "flow_fake_detector_centercrop_v4.pth" 
+CLIP_MODEL_FILENAME = "my_clip_ViT-L-14.pt" 
+class FakeImageDetector:
+
+    def __init__(self):
+        
+        self.device = "cuda" if torch.cuda.is_available() else "cpu"
+        print(f"检测器初始化在 CPU 上，运行时将使用 {self.device}")
+
+        print(f"正在从 {MODEL_REPO_ID} 下载 CLIP 模型...")
+        clip_model_path = hf_hub_download(
+            repo_id=MODEL_REPO_ID,
+            filename=CLIP_MODEL_FILENAME
+        )
+        print("CLIP 模型已下载。")
+        self.clip_model, _ = clip.load(clip_model_path, device="cpu") 
+        self.clip_model.eval()
+        print("CLIP 模型已加载到 CPU。")
+
+        print(f"正在从 {MODEL_REPO_ID} 下载 Flow 模型...")
+        flow_model_path = hf_hub_download(
+            repo_id=MODEL_REPO_ID,
+            filename=FLOW_MODEL_FILENAME
+        )
+        print("Flow 模型已下载。")
+        self.flow = flow_model()
+        self.flow.load_state_dict(torch.load(flow_model_path, map_location="cpu"))
+        self.flow = self.flow.to("cpu")
+        self.flow.eval()
+        print("Flow 模型已加载到 CPU。")
+        
+        print("模型加载完成。")
+
+        self.transform = DataAugmentationCLIP_test(
+            (0.9, 1.0), (0.05, 0.4), 1,
+            global_crops_size=224, local_crops_size=96,
+        )
+
+    @spaces.GPU(duration=10) 
+    def detect(self, image_pil, threshold=0.5):
+        
+        if not isinstance(image_pil, Image.Image):
+            raise TypeError("输入必须是 PIL Image 对象")
+
+        img_rgb = image_pil.convert("RGB")
+        
+        current_device = "cuda" if torch.cuda.is_available() else "cpu"
+        
+        flow_model_gpu = self.flow.to(current_device)
+        clip_model_gpu = self.clip_model.to(current_device)
+        
+        transformed_img_dict = self.transform(img_rgb)
+        img_tensor = transformed_img_dict["source"][0].unsqueeze(0).to(current_device)
+
+        with torch.no_grad():
+            if current_device == "cuda":
+                with autocast():
+                    embedding = clip_model_gpu.visual(img_tensor.half())
+                    z = flow_model_gpu(embedding)
+                    score = 1 - torch.sigmoid(torch.mean(z.float()**2 / 10000, dim=1)).item()
+            else:
+                embedding = clip_model_gpu.visual(img_tensor)
+                z = flow_model_gpu(embedding.float()) 
+                score = 1 - torch.sigmoid(torch.mean(z.float()**2 / 10000, dim=1)).item()
+
+        if current_device == "cuda":
+            torch.cuda.empty_cache()
+
+        if score > threshold:
+            result_text = f"结论: 伪造的 (Fake)\n分数: {score:.10f}"
+        else:
+            result_text = f"结论: 真实的 (Real)\n分数: {score:.10f}"
+            
+        return result_text, score

freia_funcs.py

@@ -0,0 +1,473 @@
+'''This Code is based on the FrEIA Framework, source: https://github.com/VLL-HD/FrEIA
+It is a assembly of the necessary modules/functions from FrEIA that are needed for our purposes.'''
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch.autograd import Variable
+from math import exp
+import numpy as np
+VERBOSE = False
+
+
+class dummy_data:
+    def __init__(self, *dims):
+        self.dims = dims
+
+    @property
+    def shape(self):
+        return self.dims
+
+class F_fully_connected(nn.Module):
+    '''Fully connected tranformation, not reversible, but used below.'''
+
+    def __init__(self, size_in, size, internal_size=None, dropout=0.0):
+        super(F_fully_connected, self).__init__()
+        if not internal_size:
+            internal_size = 2*size
+
+        self.d1 = nn.Dropout(p=dropout)
+        self.d2 = nn.Dropout(p=dropout)
+        self.d2b = nn.Dropout(p=dropout)
+
+        self.fc1 = nn.Linear(size_in, internal_size)
+        self.fc2 = nn.Linear(internal_size, internal_size)
+        self.fc2b = nn.Linear(internal_size, internal_size)
+        self.fc3 = nn.Linear(internal_size, size)
+
+        self.nl1 = nn.ReLU()
+        self.nl2 = nn.ReLU()
+        self.nl2b = nn.ReLU()
+
+        self.bn = nn.BatchNorm1d(size_in)
+
+
+    def forward(self, x):
+        out = self.nl1(self.d1(self.fc1(x)))
+        out = self.nl2(self.d2(self.fc2(out)))
+        out = self.nl2b(self.d2b(self.fc2b(out)))
+        out = self.fc3(out)
+        return out
+
+class permute_layer(nn.Module):
+    '''permutes input vector in a random but fixed way'''
+
+    def __init__(self, dims_in, seed):
+        super(permute_layer, self).__init__()
+        self.in_channels = dims_in[0][0]
+
+        np.random.seed(seed)
+        self.perm = np.random.permutation(self.in_channels)
+        np.random.seed()
+
+        self.perm_inv = np.zeros_like(self.perm)
+        for i, p in enumerate(self.perm):
+            self.perm_inv[p] = i
+
+        self.perm = torch.LongTensor(self.perm)
+        self.perm_inv = torch.LongTensor(self.perm_inv)
+
+    def forward(self, x, rev=False):
+        if not rev:
+            return [x[0][:, self.perm]]
+        else:
+            return [x[0][:, self.perm_inv]]
+
+    def jacobian(self, x, rev=False):
+        # TODO: use batch size, set as nn.Parameter so cuda() works
+        return 0.
+
+    def output_dims(self, input_dims):
+        assert len(input_dims) == 1, "Can only use 1 input"
+        return input_dims
+
+
+
+class glow_coupling_layer(nn.Module):
+    def __init__(self, dims_in, F_class=F_fully_connected, F_args={},
+                 clamp=5.):
+        super(glow_coupling_layer, self).__init__()
+        channels = dims_in[0][0]
+        self.ndims = len(dims_in[0])
+
+        self.split_len1 = channels // 2
+        self.split_len2 = channels - channels // 2
+
+        self.clamp = clamp
+        self.max_s = exp(clamp)
+        self.min_s = exp(-clamp)
+
+        self.s1 = F_class(self.split_len1, self.split_len2*2, **F_args)
+        self.s2 = F_class(self.split_len2, self.split_len1*2, **F_args)
+
+    def e(self, s):
+        return torch.exp(self.log_e(s))
+
+    def log_e(self, s):
+        return self.clamp * 0.636 * torch.atan(s / self.clamp)
+
+    def forward(self, x, rev=False):
+        x1, x2 = (x[0].narrow(1, 0, self.split_len1),
+                  x[0].narrow(1, self.split_len1, self.split_len2))
+
+        if not rev:
+            r2 = self.s2(x2)
+            s2, t2 = r2[:, :self.split_len1], r2[:, self.split_len1:]
+            #print(s2.shape, x1.shape, t2.shape)
+            y1 = self.e(s2) * x1 + t2
+
+            r1 = self.s1(y1)
+            s1, t1 = r1[:, :self.split_len2], r1[:, self.split_len2:]
+            y2 = self.e(s1) * x2 + t1
+
+        else:  # names of x and y are swapped!
+            r1 = self.s1(x1)
+            s1, t1 = r1[:, :self.split_len2], r1[:, self.split_len2:]
+            y2 = (x2 - t1) / self.e(s1)
+
+            r2 = self.s2(y2)
+            s2, t2 = r2[:, :self.split_len1], r2[:, self.split_len1:]
+            y1 = (x1 - t2) / self.e(s2)
+        y = torch.cat((y1, y2), 1)
+        y = torch.clamp(y, -1e6, 1e6)
+        return [y]
+
+    def jacobian(self, x, rev=False):
+        x1, x2 = (x[0].narrow(1, 0, self.split_len1),
+                  x[0].narrow(1, self.split_len1, self.split_len2))
+
+        if not rev:
+            r2 = self.s2(x2)
+            s2, t2 = r2[:, :self.split_len1], r2[:, self.split_len1:]
+            y1 = self.e(s2) * x1 + t2
+
+            r1 = self.s1(y1)
+            s1, t1 = r1[:, :self.split_len2], r1[:, self.split_len2:]
+
+        else:  # names of x and y are swapped!
+            r1 = self.s1(x1)
+            s1, t1 = r1[:, :self.split_len2], r1[:, self.split_len2:]
+            y2 = (x2 - t1) / self.e(s1)
+
+            r2 = self.s2(y2)
+            s2, t2 = r2[:, :self.split_len1], r2[:, self.split_len1:]
+
+        jac = (torch.sum(self.log_e(s1), dim=1)
+               + torch.sum(self.log_e(s2), dim=1))
+        for i in range(self.ndims-1):
+            jac = torch.sum(jac, dim=1)
+
+        return jac
+
+    def output_dims(self, input_dims):
+        assert len(input_dims) == 1, "Can only use 1 input"
+        return input_dims
+
+class Node:
+    '''The Node class represents one transformation in the graph, with an
+    arbitrary number of in- and outputs.'''
+    def __init__(self, inputs, module_type, module_args, name=None):
+        self.inputs = inputs
+        self.outputs = []
+        self.module_type = module_type
+        self.module_args = module_args
+
+        self.input_dims, self.module = None, None
+        self.computed = None
+        self.computed_rev = None
+        self.id = None
+
+        if name:
+            self.name = name
+        else:
+            self.name = hex(id(self))[-6:]
+        for i in range(255):
+            exec('self.out{0} = (self, {0})'.format(i))
+
+    def build_modules(self, verbose=VERBOSE):
+        ''' Returns a list with the dimension of each output of this node,
+        recursively calling build_modules of the nodes connected to the input.
+        Use this information to initialize the pytorch nn.Module of this node.
+        '''
+
+        if not self.input_dims:  # Only do it if this hasn't been computed yet
+            self.input_dims = [n.build_modules(verbose=verbose)[c]
+                               for n, c in self.inputs]
+            try:
+                self.module = self.module_type(self.input_dims,
+                                               **self.module_args)
+            except Exception as e:
+                print('Error in node %s' % (self.name))
+                raise e
+
+            if verbose:
+                print("Node %s has following input dimensions:" % (self.name))
+                for d, (n, c) in zip(self.input_dims, self.inputs):
+                    print("\t Output #%i of node %s:" % (c, n.name), d)
+                print()
+
+            self.output_dims = self.module.output_dims(self.input_dims)
+            self.n_outputs = len(self.output_dims)
+
+        return self.output_dims
+
+    def run_forward(self, op_list):
+        '''Determine the order of operations needed to reach this node. Calls
+        run_forward of parent nodes recursively. Each operation is appended to
+        the global list op_list, in the form (node ID, input variable IDs,
+        output variable IDs)'''
+
+        if not self.computed:
+
+            # Compute all nodes which provide inputs, filter out the
+            # channels you need
+            self.input_vars = []
+            for i, (n, c) in enumerate(self.inputs):
+                self.input_vars.append(n.run_forward(op_list)[c])
+                # Register youself as an output in the input node
+                n.outputs.append((self, i))
+
+            # All outputs could now be computed
+            self.computed = [(self.id, i) for i in range(self.n_outputs)]
+            op_list.append((self.id, self.input_vars, self.computed))
+
+        # Return the variables you have computed (this happens mulitple times
+        # without recomputing if called repeatedly)
+        return self.computed
+
+    def run_backward(self, op_list):
+        '''See run_forward, this is the same, only for the reverse computation.
+        Need to call run_forward first, otherwise this function will not
+        work'''
+
+        assert len(self.outputs) > 0, "Call run_forward first"
+        if not self.computed_rev:
+
+            # These are the input variables that must be computed first
+            output_vars = [(self.id, i) for i in range(self.n_outputs)]
+
+            # Recursively compute these
+            for n, c in self.outputs:
+                n.run_backward(op_list)
+
+            # The variables that this node computes are the input variables
+            # from the forward pass
+            self.computed_rev = self.input_vars
+            op_list.append((self.id, output_vars, self.computed_rev))
+
+        return self.computed_rev
+
+
+class InputNode(Node):
+    '''Special type of node that represents the input data of the whole net (or
+    ouput when running reverse)'''
+
+    def __init__(self, *dims, name='node'):
+        self.name = name
+        self.data = dummy_data(*dims)
+        self.outputs = []
+        self.module = None
+        self.computed_rev = None
+        self.n_outputs = 1
+        self.input_vars = []
+        self.out0 = (self, 0)
+
+    def build_modules(self, verbose=VERBOSE):
+        return [self.data.shape]
+
+    def run_forward(self, op_list):
+        return [(self.id, 0)]
+
+
+class OutputNode(Node):
+    '''Special type of node that represents the output of the whole net (of the
+    input when running in reverse)'''
+    class dummy(nn.Module):
+
+        def __init__(self, *args):
+            super(OutputNode.dummy, self).__init__()
+
+        def __call__(*args):
+            return args
+
+        def output_dims(*args):
+            return args
+
+    def __init__(self, inputs, name='node'):
+        self.module_type, self.module_args = self.dummy, {}
+        self.output_dims = []
+        self.inputs = inputs
+        self.input_dims, self.module = None, None
+        self.computed = None
+        self.id = None
+        self.name = name
+
+        for c, inp in enumerate(self.inputs):
+            inp[0].outputs.append((self, c))
+
+    def run_backward(self, op_list):
+        return [(self.id, 0)]
+
+
+class ReversibleGraphNet(nn.Module):
+    '''This class represents the invertible net itself. It is a subclass of
+    torch.nn.Module and supports the same methods. The forward method has an
+    additional option 'rev', whith which the net can be computed in reverse.'''
+
+    def __init__(self, node_list, ind_in=None, ind_out=None, verbose=False):
+        '''node_list should be a list of all nodes involved, and ind_in,
+        ind_out are the indexes of the special nodes InputNode and OutputNode
+        in this list.'''
+        super(ReversibleGraphNet, self).__init__()
+
+        # Gather lists of input and output nodes
+        if ind_in is not None:
+            if isinstance(ind_in, int):
+                self.ind_in = list([ind_in])
+            else:
+                self.ind_in = ind_in
+        else:
+            self.ind_in = [i for i in range(len(node_list))
+                           if isinstance(node_list[i], InputNode)]
+            assert len(self.ind_in) > 0, "No input nodes specified."
+        if ind_out is not None:
+            if isinstance(ind_out, int):
+                self.ind_out = list([ind_out])
+            else:
+                self.ind_out = ind_out
+        else:
+            self.ind_out = [i for i in range(len(node_list))
+                            if isinstance(node_list[i], OutputNode)]
+            assert len(self.ind_out) > 0, "No output nodes specified."
+
+        self.return_vars = []
+        self.input_vars = []
+
+        # Assign each node a unique ID
+        self.node_list = node_list
+        for i, n in enumerate(node_list):
+            n.id = i
+
+        # Recursively build the nodes nn.Modules and determine order of
+        # operations
+        ops = []
+        for i in self.ind_out:
+            node_list[i].build_modules(verbose=verbose)
+            node_list[i].run_forward(ops)
+
+        # create list of Pytorch variables that are used
+        variables = set()
+        for o in ops:
+            variables = variables.union(set(o[1] + o[2]))
+        self.variables_ind = list(variables)
+
+        self.indexed_ops = self.ops_to_indexed(ops)
+
+        self.module_list = nn.ModuleList([n.module for n in node_list])
+        self.variable_list = [Variable(requires_grad=True) for v in variables]
+
+        # Find out the order of operations for reverse calculations
+        ops_rev = []
+        for i in self.ind_in:
+            node_list[i].run_backward(ops_rev)
+        self.indexed_ops_rev = self.ops_to_indexed(ops_rev)
+
+    def ops_to_indexed(self, ops):
+        '''Helper function to translate the list of variables (origin ID, channel),
+        to variable IDs.'''
+        result = []
+
+        for o in ops:
+            try:
+                vars_in = [self.variables_ind.index(v) for v in o[1]]
+            except ValueError:
+                vars_in = -1
+
+            vars_out = [self.variables_ind.index(v) for v in o[2]]
+
+            # Collect input/output nodes in separate lists, but don't add to
+            # indexed ops
+            if o[0] in self.ind_out:
+                self.return_vars.append(self.variables_ind.index(o[1][0]))
+                continue
+            if o[0] in self.ind_in:
+                self.input_vars.append(self.variables_ind.index(o[1][0]))
+                continue
+
+            result.append((o[0], vars_in, vars_out))
+
+        # Sort input/output variables so they correspond to initial node list
+        # order
+        self.return_vars.sort(key=lambda i: self.variables_ind[i][0])
+        self.input_vars.sort(key=lambda i: self.variables_ind[i][0])
+
+        return result
+
+    def forward(self, x, rev=False):
+        '''Forward or backward computation of the whole net.'''
+        if rev:
+            use_list = self.indexed_ops_rev
+            input_vars, output_vars = self.return_vars, self.input_vars
+        else:
+            use_list = self.indexed_ops
+            input_vars, output_vars = self.input_vars, self.return_vars
+
+        if isinstance(x, (list, tuple)):
+            assert len(x) == len(input_vars), (
+                f"Got list of {len(x)} input tensors for "
+                f"{'inverse' if rev else 'forward'} pass, but expected "
+                f"{len(input_vars)}."
+            )
+            for i in range(len(input_vars)):
+                self.variable_list[input_vars[i]] = x[i]
+        else:
+            assert len(input_vars) == 1, (f"Got single input tensor for "
+                                          f"{'inverse' if rev else 'forward'} "
+                                          f"pass, but expected list of "
+                                          f"{len(input_vars)}.")
+            self.variable_list[input_vars[0]] = x
+
+        for o in use_list:
+            try:
+                results = self.module_list[o[0]]([self.variable_list[i]
+                                                  for i in o[1]], rev=rev)
+            except TypeError:
+                raise RuntimeError("Are you sure all used Nodes are in the "
+                                   "Node list?")
+            for i, r in zip(o[2], results):
+                self.variable_list[i] = r
+            # self.variable_list[o[2][0]] = self.variable_list[o[1][0]]
+
+        out = [self.variable_list[output_vars[i]]
+               for i in range(len(output_vars))]
+        if len(out) == 1:
+            return out[0]
+        else:
+            return out
+
+    def jacobian(self, x=None, rev=False, run_forward=True):
+        '''Compute the jacobian determinant of the whole net.'''
+        jacobian = 0
+
+        if rev:
+            use_list = self.indexed_ops_rev
+        else:
+            use_list = self.indexed_ops
+
+        if run_forward:
+            if x is None:
+                raise RuntimeError("You need to provide an input if you want "
+                                   "to run a forward pass")
+            self.forward(x, rev=rev)
+        jacobian_list = list()
+        for o in use_list:
+            try:
+                node_jac = self.module_list[o[0]].jacobian(
+                    [self.variable_list[i] for i in o[1]], rev=rev
+                )
+                jacobian += node_jac
+                jacobian_list.append(jacobian)
+            except TypeError:
+                raise RuntimeError("Are you sure all used Nodes are in the "
+                                   "Node list?")
+
+        return jacobian

loralib/__init__.py

@@ -0,0 +1,2 @@
+from .layers import *
+from .utils import *

loralib/easymultiheadattention.py

@@ -0,0 +1,124 @@
+import torch.nn as nn
+import torch
+import torch.nn.functional as F
+
+"""
+Source : https://github.com/KyanChen/MakeMultiHeadNaive/blob/master/main.py
+"""
+
+class PlainMultiHeadAttention(nn.Module):
+    def __init__(
+            self,
+            existing_mha: nn.MultiheadAttention):
+        super().__init__()
+        
+        self.dropout = 0 # this module is not used to retrain the main block
+        self.embed_dim = existing_mha.embed_dim
+        self.kdim = existing_mha.kdim
+        self.vdim = existing_mha.vdim
+        self._qkv_same_embed_dim = existing_mha._qkv_same_embed_dim
+        self.num_heads = existing_mha.num_heads
+        self.batch_first = existing_mha.batch_first
+        self.head_dim = existing_mha.head_dim
+        self.qkv = nn.Linear(self.embed_dim, self.embed_dim * 3, bias=existing_mha.in_proj_bias is not None)
+        self.proj = nn.Linear(self.embed_dim, self.embed_dim, bias=existing_mha.out_proj.bias is not None)
+
+        # Initialize parameters
+        with torch.no_grad():
+            self.qkv.weight.data.copy_(existing_mha.in_proj_weight.data)
+            if self.qkv.bias is not None:
+                self.qkv.bias.data.copy_(existing_mha.in_proj_bias.data) 
+            self.proj.weight.data.copy_(existing_mha.out_proj.weight.data)
+            if self.proj.bias is not None:
+                self.proj.bias.data.copy_(existing_mha.out_proj.bias.data) 
+
+        self.scaled_dot_product_attention = F.scaled_dot_product_attention
+        
+    def forward(
+            self,
+            query,
+            key,
+            value,
+            key_padding_mask=None,
+            need_weights=True,
+            attn_mask=None,
+            average_attn_weights=True,
+            is_causal=False):
+
+        if attn_mask is not None and is_causal:
+            raise AssertionError("Only allow causal mask or attn_mask")
+        is_batched = query.dim() == 3
+        key_padding_mask = F._canonical_mask(
+            mask=key_padding_mask,
+            mask_name="key_padding_mask",
+            other_type=F._none_or_dtype(attn_mask),
+            other_name="attn_mask",
+            target_type=query.dtype
+        )
+
+        if self.batch_first and is_batched:
+            if key is value:
+                if query is key:
+                    query = key = value = query.transpose(1, 0)
+                else:
+                    query, key = [x.transpose(1, 0) for x in (query, key)]
+                    value = key
+            else:
+                query, key, value = [x.transpose(1, 0) for x in (query, key, value)]
+
+        tgt_len, bsz, embed_dim = query.shape
+        src_len, _, _ = key.shape
+
+        E = query.size(-1)
+        qkv = self.qkv(query)
+        qkv = qkv.unflatten(-1, (3, E)).unsqueeze(0).transpose(0, -2).squeeze(-2).contiguous()
+        q, k, v = qkv[0], qkv[1], qkv[2]
+
+        attn_mask = F._canonical_mask(
+            mask=attn_mask,
+            mask_name="attn_mask",
+            other_type=F._none_or_dtype(key_padding_mask),
+            other_name="key_padding_mask",
+            target_type=q.dtype,
+            check_other=False,
+        )
+
+        if attn_mask is not None:
+            # ensure attn_mask's dim is 3
+            if attn_mask.dim() == 2:
+                correct_2d_size = (tgt_len, src_len)
+                if attn_mask.shape != correct_2d_size:
+                    raise RuntimeError(
+                        f"The shape of the 2D attn_mask is {attn_mask.shape}, but should be {correct_2d_size}.")
+                attn_mask = attn_mask.unsqueeze(0)
+            elif attn_mask.dim() == 3:
+                correct_3d_size = (bsz * self.num_heads, tgt_len, src_len)
+                if attn_mask.shape != correct_3d_size:
+                    raise RuntimeError(
+                        f"The shape of the 3D attn_mask is {attn_mask.shape}, but should be {correct_3d_size}.")
+            else:
+                raise RuntimeError(f"attn_mask's dimension {attn_mask.dim()} is not supported")
+
+        if attn_mask is not None:
+            if attn_mask.size(0) == 1 and attn_mask.dim() == 3:
+                attn_mask = attn_mask.unsqueeze(0)
+            else:
+                attn_mask = attn_mask.view(bsz, self.num_heads, -1, src_len)
+
+        dropout_p = self.dropout if self.training else 0.
+
+        q = q.view(tgt_len, bsz * self.num_heads, self.head_dim).transpose(0, 1)
+        k = k.view(src_len, bsz * self.num_heads, self.head_dim).transpose(0, 1)
+        v = v.view(src_len, bsz * self.num_heads, self.head_dim).transpose(0, 1)
+        src_len = k.size(1)
+        q = q.view(bsz, self.num_heads, tgt_len, self.head_dim)
+        k = k.view(bsz, self.num_heads, src_len, self.head_dim)
+        v = v.view(bsz, self.num_heads, src_len, self.head_dim)
+
+        attn_output = self.scaled_dot_product_attention(q, k, v, attn_mask, dropout_p, is_causal)
+        attn_output = attn_output.permute(2, 0, 1, 3).contiguous().view(bsz * tgt_len, embed_dim)
+        attn_output = self.proj(attn_output)
+        attn_output = attn_output.view(tgt_len, bsz, attn_output.size(1))
+        if self.batch_first and is_batched:
+            return attn_output.transpose(1, 0), None
+        return attn_output, None

loralib/layers.py

@@ -0,0 +1,598 @@
+#  ------------------------------------------------------------------------------------------
+#  This code is reconstructed based on loralib (https://github.com/microsoft/LoRA) by Baijiong Lin.
+#  ------------------------------------------------------------------------------------------
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+import math
+from typing import Optional, List
+
+def set_param(curr_mod, name, param=None, mode='update'):
+    r"""Refer to https://github.com/Baijiong-Lin/MOML/blob/main/MTL/utils.py"""
+    if '.' in name:
+        n = name.split('.')
+        module_name = n[0]
+        rest = '.'.join(n[1:])
+        for name, mod in curr_mod.named_children():
+            if module_name == name:
+                return set_param(mod, rest, param, mode=mode)
+    else:
+        if mode == 'update':
+            delattr(curr_mod, name)
+            setattr(curr_mod, name, param)
+        elif mode == 'get':
+            if hasattr(curr_mod, name):
+                p = getattr(curr_mod, name)
+                return p
+
+class LoRALayer():
+    def __init__(
+        self, 
+        r: int, 
+        lora_alpha: int, 
+        fan_in_fan_out: bool = False,
+        dropout_rate:float = 0,
+    ):
+        self.r = r
+        self.lora_alpha = lora_alpha
+        self.dropout_rate = dropout_rate
+        if self.r > 0:
+            #self.scaling = self.lora_alpha / self.r
+            self.scaling = self.lora_alpha/math.sqrt(self.r) # 
+        # Mark the weight as unmerged
+        self.merged = False
+        # Set this to True if the layer to replace stores weight like (fan_in, fan_out)
+        self.fan_in_fan_out = fan_in_fan_out
+        # define params that require LoRA {'param_name': 'lora_name'}
+        self.params_with_lora = {}
+
+    def register_lora_param(self):
+        r"""Register LoRA matrix"""
+        for param_name, lora_name in self.params_with_lora.items():
+            assert len(eval(f'self.{param_name}').size()) == 2
+            self.register_parameter(f'{lora_name}_lora_A', 
+                nn.Parameter(eval(f'self.{param_name}').new_zeros((self.r, eval(f'self.{param_name}').size()[1])))
+                )
+            self.register_parameter(f'{lora_name}_lora_B', 
+                nn.Parameter(eval(f'self.{param_name}').new_zeros((eval(f'self.{param_name}').size()[0], self.r)))
+                )
+                
+            eval(f'self.{param_name}').requires_grad = False
+
+    def init_lora_param(self):
+        for param_name, lora_name in self.params_with_lora.items():
+            if hasattr(self, f'{lora_name}_lora_A'):
+                # initialize A the same way as the default for nn.Linear and B to zero
+                nn.init.kaiming_uniform_(eval(f'self.{lora_name}_lora_A'), a=math.sqrt(5))
+                nn.init.zeros_(eval(f'self.{lora_name}_lora_B'))
+
+    def transpose(self, w: torch.Tensor):
+        return w.transpose(0, 1) if self.fan_in_fan_out else w
+
+    def merge_BA(self, param_name: str):
+        lora_name = self.params_with_lora[param_name]
+        return self.transpose((eval(f'self.{lora_name}_lora_B') @ eval(f'self.{lora_name}_lora_A')).view(eval(f'self.{param_name}').shape))
+
+    
+   
+    
+    def merge_lora_param(self):
+        r"""p_new = p + scaling * B @ A and keep differentiable to A and B"""
+        for param_name, lora_name in self.params_with_lora.items():
+            p = set_param(self, param_name, mode='get')
+            # detach() is very important here
+            
+            p_new = p.detach() + self.merge_BA(param_name) * self.scaling
+            set_param(self, param_name, param=p_new, mode='update')
+
+    def add_lora_data(self):
+        r"""NOT differentiable"""
+        for param_name, lora_name in self.params_with_lora.items():
+            eval(f'self.{param_name}').data += self.merge_BA(param_name) * self.scaling
+    
+    def sub_lora_data(self):
+        r"""NOT differentiable"""
+        for param_name, lora_name in self.params_with_lora.items():
+            eval(f'self.{param_name}').data -= self.merge_BA(param_name) * self.scaling
+            
+    
+    def lora_train(self, mode: bool = True):
+        if mode:
+            if self.merged and self.r > 0:
+            # Make sure that the weights are not merged
+                self.sub_lora_data()
+            self.merged = False
+        else:
+            if not self.merged and self.r > 0:
+            # Merge the weights and mark it
+                self.add_lora_data()
+            self.merged = True 
+
+
+class Embedding(nn.Embedding, LoRALayer):
+    # LoRA implemented in a Embedding layer
+    def __init__(
+        self,
+        num_embeddings: int,
+        embedding_dim: int,
+        r: int = 0,
+        lora_alpha: int = 1,
+        **kwargs
+    ):
+        nn.Embedding.__init__(self, num_embeddings, embedding_dim, **kwargs)
+        LoRALayer.__init__(self, r=r, lora_alpha=lora_alpha)
+
+        self.params_with_lora = {'weight': 'w'}
+        if r > 0:
+            self.register_lora_param()
+        nn.Embedding.reset_parameters(self)
+        self.init_lora_param()
+
+    def init_lora_param(self):
+        if hasattr(self, 'w_lora_A'):
+            # initialize A the same way as the default for nn.Linear and B to zero
+            nn.init.zeros_(self.w_lora_A)
+            nn.init.normal_(self.w_lora_B)
+
+    def train(self, mode: bool = True):
+        nn.Embedding.train(self, mode)
+        self.lora_train(mode)
+        
+    def forward(self, x: torch.Tensor, **kwargs):
+
+        if self.r > 0 and not self.merged:
+            self.merge_lora_param()
+            result = nn.Embedding.forward(self, x, **kwargs)
+            self.sub_lora_data()
+            return result
+        else:
+            return nn.Embedding.forward(self, x, **kwargs)
+
+class LinearLoRA(nn.Linear, LoRALayer):
+    # LoRA implemented in a Linear layer
+    def __init__(
+        self, 
+        existing_linear: nn.Linear,
+        r: int = 0, 
+        lora_alpha: int = 1, 
+        fan_in_fan_out: bool = False,
+        dropout_rate = 0.,
+        **kwargs
+    ):
+        super().__init__(
+            in_features=existing_linear.in_features, 
+            out_features=existing_linear.out_features)
+        
+        self.load_state_dict(existing_linear.state_dict())
+        LoRALayer.__init__(self, r=r, lora_alpha=lora_alpha, fan_in_fan_out=fan_in_fan_out)
+
+        # Actual trainable parameters
+        self.params_with_lora = {'weight': 'w'}
+        if r > 0:
+            self.register_lora_param()
+        self.init_lora_param()
+        self.weight.data = self.transpose(self.weight.data)
+        if dropout_rate > 0:
+            self.dropout = nn.Dropout(dropout_rate)
+        else:
+            self.dropout = None
+
+    def train(self, mode: bool = True):
+        super().train(mode)     
+        self.lora_train(mode)
+
+        
+    def forward(self, x: torch.Tensor, **kwargs):
+        
+        if self.dropout is None: # do as before
+            if self.r > 0 and not self.merged:
+                self.merge_lora_param()
+                result = nn.Linear.forward(self, x, **kwargs)
+                self.sub_lora_data()
+                return result
+            else:
+                return nn.Linear.forward(self, x, **kwargs)
+            
+        # Compute the original linear transformation
+        original_output = nn.Linear.forward(self, x)
+
+        if self.training and self.dropout.p > 0:
+            x = self.dropout(x)
+        
+        if self.r > 0 and not self.merged:
+            lora_adjustment = torch.matmul(x,self.merge_BA('weight').transpose(0, 1)) * self.scaling 
+            result = original_output + lora_adjustment
+        else:
+            result = original_output
+        return result
+
+class Conv1d(nn.Conv1d, LoRALayer):
+    # LoRA implemented in a Conv1d layer
+    def __init__(
+        self, 
+        in_channels: int, 
+        out_channels: int,
+        kernel_size: int,
+        r: int = 0, 
+        lora_alpha: int = 1, 
+        **kwargs
+    ):
+        nn.Conv1d.__init__(self, in_channels, out_channels, kernel_size, **kwargs)
+        LoRALayer.__init__(self, r=r, lora_alpha=lora_alpha)
+
+        assert type(kernel_size) is int
+        # Actual trainable parameters
+        self.params_with_lora = {'weight': 'w'}
+        if r > 0:
+            self.w_lora_A = nn.Parameter(
+                self.weight.new_zeros((r*kernel_size, in_channels*kernel_size))
+            )
+            self.w_lora_B = nn.Parameter(
+                self.weight.new_zeros((out_channels//self.groups*kernel_size, r*kernel_size))
+            )
+            # Freezing the pre-trained weight matrix
+            self.weight.requires_grad = False
+        nn.Conv1d.reset_parameters(self)
+        self.init_lora_param()
+
+    def train(self, mode: bool = True):
+        nn.Conv1d.train(self, mode)      
+        self.lora_train(mode)
+
+    def forward(self, x: torch.Tensor, **kwargs):
+
+        if self.r > 0 and not self.merged:
+            self.merge_lora_param()
+            result = nn.Conv1d.forward(self, x, **kwargs)
+            self.sub_lora_data()
+            return result
+        else:
+            return nn.Conv1d.forward(self, x, **kwargs)
+
+class Conv2d(nn.Conv2d, LoRALayer):
+    # LoRA implemented in a Conv2d layer
+    def __init__(
+        self, 
+        in_channels: int, 
+        out_channels: int,
+        kernel_size: int,
+        r: int = 0, 
+        lora_alpha: int = 1, 
+        **kwargs
+    ):
+        nn.Conv2d.__init__(self, in_channels, out_channels, kernel_size, **kwargs)
+        LoRALayer.__init__(self, r=r, lora_alpha=lora_alpha)
+
+        assert type(kernel_size) is int
+        # Actual trainable parameters
+        self.params_with_lora = {'weight': 'w'}
+        if r > 0:
+            self.w_lora_A = nn.Parameter(
+                self.weight.new_zeros((r*kernel_size, in_channels*kernel_size))
+            )
+            self.w_lora_B = nn.Parameter(
+                self.weight.new_zeros((out_channels//self.groups*kernel_size, r*kernel_size))
+            )
+            # Freezing the pre-trained weight matrix
+            self.weight.requires_grad = False
+        nn.Conv2d.reset_parameters(self)
+        self.init_lora_param()
+
+    def train(self, mode: bool = True):
+        nn.Conv2d.train(self, mode)      
+        self.lora_train(mode)
+
+    def forward(self, x: torch.Tensor, **kwargs):
+
+        if self.r > 0 and not self.merged:
+            self.merge_lora_param()
+            result = nn.Conv2d.forward(self, x, **kwargs)
+            self.sub_lora_data()
+            return result
+        else:
+            return nn.Conv2d.forward(self, x, **kwargs)
+
+class Conv3d(nn.Conv3d, LoRALayer):
+    # LoRA implemented in a Conv3d layer
+    def __init__(
+        self, 
+        in_channels: int, 
+        out_channels: int,
+        kernel_size: int,
+        r: int = 0, 
+        lora_alpha: int = 1, 
+        **kwargs
+    ):
+        nn.Conv3d.__init__(self, in_channels, out_channels, kernel_size, **kwargs)
+        LoRALayer.__init__(self, r=r, lora_alpha=lora_alpha)
+
+        assert type(kernel_size) is int
+        # Actual trainable parameters
+        self.params_with_lora = {'weight': 'w'}
+        if r > 0:
+            self.w_lora_A = nn.Parameter(
+                self.weight.new_zeros((r*kernel_size, in_channels*kernel_size))
+            )
+            self.w_lora_B = nn.Parameter(
+                self.weight.new_zeros((out_channels//self.groups*kernel_size, r*kernel_size))
+            )
+            # Freezing the pre-trained weight matrix
+            self.weight.requires_grad = False
+        nn.Conv3d.reset_parameters(self)
+        self.init_lora_param()
+
+    def train(self, mode: bool = True):
+        nn.Conv3d.train(self, mode)      
+        self.lora_train(mode)
+
+    def forward(self, x: torch.Tensor, **kwargs):
+
+        if self.r > 0 and not self.merged:
+            self.merge_lora_param()
+            result = nn.Conv3d.forward(self, x, **kwargs)
+            self.sub_lora_data()
+            return result
+        else:
+            return nn.Conv3d.forward(self, x, **kwargs)
+
+
+class PlainMultiheadAttentionLoRA(nn.Module):
+    def __init__(
+            self,
+            existing_mha: nn.MultiheadAttention,
+            enable_lora: list = ['q', 'k', 'v', 'o'],
+            r: int = 0, 
+            lora_alpha: int = 1, 
+            dropout_rate:float = 0.,
+            **kwargs
+        ):
+        super().__init__()
+        
+        self.dropout = 0 # this module is not used to retrain the main block
+        self.embed_dim = existing_mha.embed_dim
+        self.kdim = existing_mha.kdim
+        self.vdim = existing_mha.vdim
+        self._qkv_same_embed_dim = existing_mha._qkv_same_embed_dim
+        self.num_heads = existing_mha.num_heads
+        self.batch_first = existing_mha.batch_first
+        self.head_dim = existing_mha.head_dim
+        #self.qkv = nn.Linear(self.embed_dim, self.embed_dim * 3, bias=existing_mha.in_proj_bias is not None)
+        self.q_proj = nn.Linear(self.embed_dim, self.embed_dim, bias=existing_mha.in_proj_bias is not None)
+        self.k_proj = nn.Linear(self.embed_dim, self.embed_dim, bias=existing_mha.in_proj_bias is not None)
+        self.v_proj = nn.Linear(self.embed_dim, self.embed_dim, bias=existing_mha.in_proj_bias is not None)
+        self.proj = nn.Linear(self.embed_dim, self.embed_dim, bias=existing_mha.out_proj.bias is not None)
+
+        # Initialize parameters
+        with torch.no_grad():
+            
+            # Extract the existing weights and biases
+            existing_weight = existing_mha.in_proj_weight.data
+            existing_bias = existing_mha.in_proj_bias.data if existing_mha.in_proj_bias is not None else None
+
+            # Initialize q_proj
+            self.q_proj.weight.data.copy_(existing_weight[:self.embed_dim, :])
+            if existing_bias is not None:
+                self.q_proj.bias.data.copy_(existing_bias[:self.embed_dim])
+
+            # Initialize k_proj
+            self.k_proj.weight.data.copy_(existing_weight[self.embed_dim:2*self.embed_dim, :])
+            if existing_bias is not None:
+                self.k_proj.bias.data.copy_(existing_bias[self.embed_dim:2*self.embed_dim])
+
+            # Initialize v_proj
+            self.v_proj.weight.data.copy_(existing_weight[2*self.embed_dim:, :])
+            if existing_bias is not None:
+                self.v_proj.bias.data.copy_(existing_bias[2*self.embed_dim:])
+
+            # Initialize proj
+            self.proj.weight.data.copy_(existing_mha.out_proj.weight.data)
+            if self.proj.bias is not None:
+                self.proj.bias.data.copy_(existing_mha.out_proj.bias.data)
+
+        self.scaled_dot_product_attention = F.scaled_dot_product_attention
+        
+        
+        LoRALayer.__init__(self, r=r, lora_alpha=lora_alpha, dropout_rate=dropout_rate)
+        
+        # Init qkv as a new lora linear layer 
+        for item in enable_lora:
+            if item == 'q':
+                self.q_proj = LinearLoRA(self.q_proj,
+                                         r=r,
+                                         lora_alpha=lora_alpha,
+                                         fan_in_fan_out=False,
+                                         dropout_rate = dropout_rate)
+            elif item == 'k':
+                self.k_proj = LinearLoRA(self.k_proj,
+                                         r=r,
+                                         lora_alpha=lora_alpha,
+                                         fan_in_fan_out=False,
+                                         dropout_rate = dropout_rate)
+            elif item == 'v':
+                self.v_proj = LinearLoRA(self.v_proj,
+                                         r=r,
+                                         lora_alpha=lora_alpha,
+                                         fan_in_fan_out=False,
+                                         dropout_rate = dropout_rate)
+            elif item == 'o':
+                self.proj = LinearLoRA(self.proj,
+                                         r=r,
+                                         lora_alpha=lora_alpha,
+                                         fan_in_fan_out=False,
+                                         dropout_rate = dropout_rate)
+        
+    def forward_module(
+            self,
+            query,
+            key,
+            value,
+            key_padding_mask=None,
+            need_weights=True,
+            attn_mask=None,
+            average_attn_weights=True,
+            is_causal=False):
+
+        if attn_mask is not None and is_causal:
+            raise AssertionError("Only allow causal mask or attn_mask")
+        is_batched = query.dim() == 3
+        key_padding_mask = F._canonical_mask(
+            mask=key_padding_mask,
+            mask_name="key_padding_mask",
+            other_type=F._none_or_dtype(attn_mask),
+            other_name="attn_mask",
+            target_type=query.dtype
+        )
+
+        if self.batch_first and is_batched:
+            if key is value:
+                if query is key:
+                    query = key = value = query.transpose(1, 0)
+                else:
+                    query, key = [x.transpose(1, 0) for x in (query, key)]
+                    value = key
+            else:
+                query, key, value = [x.transpose(1, 0) for x in (query, key, value)]
+
+        tgt_len, bsz, embed_dim = query.shape
+        src_len, _, _ = key.shape
+        """
+        E = query.size(-1)
+        qkv = self.qkv(query)
+        qkv = qkv.unflatten(-1, (3, E)).unsqueeze(0).transpose(0, -2).squeeze(-2).contiguous()
+        q, k, v = qkv[0], qkv[1], qkv[2]
+        """
+        
+        q = self.q_proj(query)
+        k = self.k_proj(key)
+        v = self.v_proj(value)
+
+        attn_mask = F._canonical_mask(
+            mask=attn_mask,
+            mask_name="attn_mask",
+            other_type=F._none_or_dtype(key_padding_mask),
+            other_name="key_padding_mask",
+            target_type=q.dtype,
+            check_other=False,
+        )
+
+        if attn_mask is not None:
+            # ensure attn_mask's dim is 3
+            if attn_mask.dim() == 2:
+                correct_2d_size = (tgt_len, src_len)
+                if attn_mask.shape != correct_2d_size:
+                    raise RuntimeError(
+                        f"The shape of the 2D attn_mask is {attn_mask.shape}, but should be {correct_2d_size}.")
+                attn_mask = attn_mask.unsqueeze(0)
+            elif attn_mask.dim() == 3:
+                correct_3d_size = (bsz * self.num_heads, tgt_len, src_len)
+                if attn_mask.shape != correct_3d_size:
+                    raise RuntimeError(
+                        f"The shape of the 3D attn_mask is {attn_mask.shape}, but should be {correct_3d_size}.")
+            else:
+                raise RuntimeError(f"attn_mask's dimension {attn_mask.dim()} is not supported")
+
+        if attn_mask is not None:
+            if attn_mask.size(0) == 1 and attn_mask.dim() == 3:
+                attn_mask = attn_mask.unsqueeze(0)
+            else:
+                attn_mask = attn_mask.view(bsz, self.num_heads, -1, src_len)
+
+        dropout_p = self.dropout if self.training else 0.
+
+        q = q.view(tgt_len, bsz * self.num_heads, self.head_dim).transpose(0, 1)
+        k = k.view(src_len, bsz * self.num_heads, self.head_dim).transpose(0, 1)
+        v = v.view(src_len, bsz * self.num_heads, self.head_dim).transpose(0, 1)
+        src_len = k.size(1)
+        q = q.view(bsz, self.num_heads, tgt_len, self.head_dim)
+        k = k.view(bsz, self.num_heads, src_len, self.head_dim)
+        v = v.view(bsz, self.num_heads, src_len, self.head_dim)
+
+        attn_output = self.scaled_dot_product_attention(q, k, v, attn_mask, dropout_p, is_causal)
+        attn_output = attn_output.permute(2, 0, 1, 3).contiguous().view(bsz * tgt_len, embed_dim)
+        attn_output = self.proj(attn_output)
+        attn_output = attn_output.view(tgt_len, bsz, attn_output.size(1))
+        if self.batch_first and is_batched:
+            return attn_output.transpose(1, 0), None
+        return attn_output, None  
+
+    def train(self, mode: bool = True):
+        super().train(mode)
+        #self.lora_train(mode)  
+
+    def forward(self,
+            query: torch.Tensor,
+            key: torch.Tensor,
+            value: torch.Tensor,
+            **kwargs):
+        
+
+        return self.forward_module(query, key, value, **kwargs) 
+        
+
+
+class MergedLinear(nn.Linear, LoRALayer):
+    # LoRA implemented in a dense layer
+    def __init__(
+        self, 
+        in_features: int, 
+        out_features: int, 
+        r: int = 0, 
+        lora_alpha: int = 1, 
+        enable_lora: List[bool] = [False],
+        fan_in_fan_out: bool = False,
+        **kwargs
+    ):
+        nn.Linear.__init__(self, in_features, out_features, **kwargs)
+        LoRALayer.__init__(self, r=r, lora_alpha=lora_alpha)
+
+        assert out_features % len(enable_lora) == 0, \
+            'The length of enable_lora must divide out_features'
+        self.enable_lora = enable_lora
+        # Actual trainable parameters
+        self.params_with_lora = {'weight': 'w'}
+        if r > 0 and any(enable_lora):
+            self.w_lora_A = nn.Parameter(
+                self.weight.new_zeros((r * sum(enable_lora), in_features)))
+            self.w_lora_B = nn.Parameter(
+                self.weight.new_zeros((out_features // len(enable_lora) * sum(enable_lora), r))
+            ) # weights for Conv1D with groups=sum(enable_lora)
+            # Freezing the pre-trained weight matrix
+            self.weight.requires_grad = False
+            # Compute the indices
+            self.lora_ind = self.weight.new_zeros(
+                (out_features, ), dtype=torch.bool
+            ).view(len(enable_lora), -1)
+            self.lora_ind[enable_lora, :] = True
+            self.lora_ind = self.lora_ind.view(-1)
+        nn.Linear.reset_parameters(self)
+        self.init_lora_param()
+        self.weight.data = self.transpose(self.weight.data)
+
+    def zero_pad(self, x):
+        result = x.new_zeros((len(self.lora_ind), *x.shape[1:]))
+        result[self.lora_ind] = x
+        return result
+
+    def merge_BA(self, param_name: str):
+        lora_name = self.params_with_lora[param_name]
+        delta_w = F.conv1d(
+            eval(f'self.{lora_name}_lora_A').unsqueeze(0), 
+            eval(f'self.{lora_name}_lora_B').unsqueeze(-1), 
+            groups=sum(self.enable_lora)
+        ).squeeze(0)
+        return self.transpose(self.zero_pad(delta_w))
+
+    def train(self, mode: bool = True):
+        nn.Linear.train(self, mode)
+        self.lora_train(mode)        
+
+    def forward(self, x: torch.Tensor, **kwargs):
+
+        if self.r > 0 and not self.merged:
+            self.merge_lora_param()
+            result = nn.Linear.forward(self, x, **kwargs)
+            self.sub_lora_data()
+            return result
+        else:
+            return nn.Linear.forward(self, x, **kwargs)

loralib/utils.py

@@ -0,0 +1,236 @@
+import os
+
+import torch
+import torch.nn as nn
+
+from typing import Dict
+
+from .layers import LoRALayer, PlainMultiheadAttentionLoRA
+
+INDEX_POSITIONS_TEXT = {
+    'top1': [11],
+    'top2': [10, 11],
+    'top3': [9, 10, 11],
+    'bottom': [0, 1, 2, 3],
+    'mid': [4, 5, 6, 7],
+    'up': [8, 9, 10, 11],
+    'half-up': [6, 7, 8, 9, 10, 11],
+    'half-bottom': [0, 1, 2, 3, 4, 5],
+    'all': [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11]}
+
+
+INDEX_POSITIONS_VISION = {
+    'ViT-B/16': {
+        'top': [11],
+        'top3': [9, 10, 11],
+        'bottom': [0, 1, 2, 3],
+        'mid': [4, 5, 6, 7],
+        'up': [8, 9, 10, 11],
+        'half-up': [6, 7, 8, 9, 10, 11],
+        'half-bottom': [0, 1, 2, 3, 4, 5],
+        'all': [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11]},
+    'ViT-B/32': {
+        'bottom': [0, 1, 2, 3],
+        'mid': [4, 5, 6, 7],
+        'up': [8, 9, 10, 11],
+        'half-up': [6, 7, 8, 9, 10, 11],
+        'half-bottom': [0, 1, 2, 3, 4, 5],
+        'all': [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11]},
+
+    'ViT-L/14': {
+        'half-up': [12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23],
+        'half-bottom': [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11],
+        'all': [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23]}
+}
+
+
+def mark_only_lora_as_trainable(model: nn.Module, bias: str = 'none') -> None:
+    for n, p in model.named_parameters():
+        if 'lora_' not in n:
+            p.requires_grad = False
+    if bias == 'none':
+        return
+    elif bias == 'all':
+        for n, p in model.named_parameters():
+            if 'bias' in n:
+                p.requires_grad = True
+    elif bias == 'lora_only':
+        for m in model.modules():
+            if isinstance(m, LoRALayer) and \
+                    hasattr(m, 'bias') and \
+                    m.bias is not None:
+                m.bias.requires_grad = True
+    else:
+        raise NotImplementedError
+
+
+def lora_state_dict(model: nn.Module, bias: str = 'none') -> Dict[str, torch.Tensor]:
+    my_state_dict = model.state_dict()
+    if bias == 'none':
+        return {k: my_state_dict[k] for k in my_state_dict if 'lora_' in k}
+    elif bias == 'all':
+        return {k: my_state_dict[k] for k in my_state_dict if 'lora_' in k or 'bias' in k}
+    elif bias == 'lora_only':
+        to_return = {}
+        for k in my_state_dict:
+            if 'lora_' in k:
+                to_return[k] = my_state_dict[k]
+                bias_name = k.split('lora_')[0]+'bias'
+                if bias_name in my_state_dict:
+                    to_return[bias_name] = my_state_dict[bias_name]
+        return to_return
+    else:
+        raise NotImplementedError
+
+
+def get_lora_parameters(model, bias='none'):
+    params = []
+    for name, param in model.named_parameters():
+        if bias == 'none':
+            if 'lora_' in name:
+                params.append(param)
+        elif bias == 'all':
+            if 'lora_' in name or 'bias' in name:
+                params.append(param)
+        elif bias == 'lora_only':
+            if 'lora_' in name:
+                params.append(param)
+                bias_name = name.split('lora_')[0] + 'bias'
+                if bias_name in model.state_dict():
+                    bias_param = dict(model.named_parameters())[bias_name]
+                    params.append(bias_param)
+        else:
+            raise NotImplementedError
+    return params
+
+
+def apply_lora(args, clip_model):
+    list_lora_layers = []
+    if args.encoder == 'text' or args.encoder == 'both':
+        indices = INDEX_POSITIONS_TEXT[args.position]
+        text_encoder = clip_model.transformer
+        for i, block in enumerate(text_encoder.resblocks):
+            print(f"Residual Attention Block {i}: {block}")
+            if i in indices:
+                for name, submodule in block.named_children():
+                    if isinstance(submodule, nn.MultiheadAttention):
+                        new_multi_head_lora = PlainMultiheadAttentionLoRA(
+                            submodule, enable_lora=args.params, r=args.r, lora_alpha=args.alpha, dropout_rate=args.dropout_rate)
+                        setattr(block, name, new_multi_head_lora)
+                        list_lora_layers.append(new_multi_head_lora)
+
+    if args.encoder == 'vision' or args.encoder == 'both':
+        indices = INDEX_POSITIONS_VISION[args.backbone][args.position]
+        vision_encoder = clip_model.visual.transformer
+        for i, block in enumerate(vision_encoder.resblocks):
+            print(f"Residual Attention Block {i}: {block}")
+            if i in indices:
+                for name, submodule in block.named_children():
+                    if isinstance(submodule, nn.MultiheadAttention):
+                        new_multi_head_lora = PlainMultiheadAttentionLoRA(
+                            submodule, enable_lora=args.params, r=args.r, lora_alpha=args.alpha, dropout_rate=args.dropout_rate)
+                        setattr(block, name, new_multi_head_lora)
+                        list_lora_layers.append(new_multi_head_lora)
+    return list_lora_layers
+
+
+def save_lora(args, list_lora_layers):
+    weights = {}
+    for i, layer in enumerate(list_lora_layers):
+        layer_weights = {}
+        if 'q' in args.params:
+            layer_weights['q_proj'] = {
+                'w_lora_A': layer.q_proj.w_lora_A.data,
+                'w_lora_B': layer.q_proj.w_lora_B.data
+            }
+        if 'k' in args.params:
+            layer_weights['k_proj'] = {
+                'w_lora_A': layer.k_proj.w_lora_A.data,
+                'w_lora_B': layer.k_proj.w_lora_B.data
+            }
+        if 'v' in args.params:
+            layer_weights['v_proj'] = {
+                'w_lora_A': layer.v_proj.w_lora_A.data,
+                'w_lora_B': layer.v_proj.w_lora_B.data
+            }
+        if 'o' in args.params:
+            layer_weights['proj'] = {
+                'w_lora_A': layer.proj.w_lora_A.data,
+                'w_lora_B': layer.proj.w_lora_B.data
+            }
+
+        weights[f'layer_{i}'] = layer_weights
+
+    metadata = {
+        'r': args.r,
+        'alpha': args.alpha,
+        'encoder': args.encoder,
+        'params': args.params,
+        'position': args.position
+    }
+
+    save_data = {
+        'weights': weights,
+        'metadata': metadata
+    }
+
+    # to manage names like ViT-B/16
+    backbone = args.backbone.replace('/', '').replace('-', '').lower()
+    save_dir = f'{args.save_path}/{backbone}/{args.dataset}/{args.shots}shots/seed{args.seed}'
+    os.makedirs(save_dir, exist_ok=True)
+
+    save_path = f'{save_dir}/{args.filename}.pt'
+    torch.save(save_data, save_path)
+    print(f'LoRA weights saved to {save_path}')
+
+
+def load_lora(args, list_lora_layers):
+    # to manage names like ViT-B/16
+    backbone = args.backbone.replace('/', '').replace('-', '').lower()
+    load_path = f'{args.save_path}/{backbone}/{args.dataset}/{args.shots}shots/seed{args.seed}/{args.filename}.pt'
+
+    if not os.path.exists(load_path):
+        raise FileNotFoundError(f'File {load_path} does not exist.')
+
+    loaded_data = torch.load(load_path)
+
+    metadata = loaded_data['metadata']
+    if metadata['r'] != args.r:
+        raise ValueError(
+            f"r mismatch: expected {args.r}, found {metadata['r']}")
+    if metadata['alpha'] != args.alpha:
+        raise ValueError(
+            f"alpha mismatch: expected {args.alpha}, found {metadata['alpha']}")
+    if metadata['encoder'] != args.encoder:
+        raise ValueError(
+            f"Encoder mismatch: expected {args.encoder}, found {metadata['encoder']}")
+    if metadata['params'] != args.params:
+        raise ValueError(
+            f"Params mismatch: expected {args.params}, found {metadata['params']}")
+    if metadata['position'] != args.position:
+        raise ValueError(
+            f"Position mismatch: expected {args.position}, found {metadata['position']}")
+
+    weights = loaded_data['weights']
+    for i, layer in enumerate(list_lora_layers):
+        layer_weights = weights[f'layer_{i}']
+        if 'q' in args.params and 'q_proj' in layer_weights:
+            layer.q_proj.w_lora_A.data.copy_(
+                layer_weights['q_proj']['w_lora_A'])
+            layer.q_proj.w_lora_B.data.copy_(
+                layer_weights['q_proj']['w_lora_B'])
+        if 'k' in args.params and 'k_proj' in layer_weights:
+            layer.k_proj.w_lora_A.data.copy_(
+                layer_weights['k_proj']['w_lora_A'])
+            layer.k_proj.w_lora_B.data.copy_(
+                layer_weights['k_proj']['w_lora_B'])
+        if 'v' in args.params and 'v_proj' in layer_weights:
+            layer.v_proj.w_lora_A.data.copy_(
+                layer_weights['v_proj']['w_lora_A'])
+            layer.v_proj.w_lora_B.data.copy_(
+                layer_weights['v_proj']['w_lora_B'])
+        if 'o' in args.params and 'proj' in layer_weights:
+            layer.proj.w_lora_A.data.copy_(layer_weights['proj']['w_lora_A'])
+            layer.proj.w_lora_B.data.copy_(layer_weights['proj']['w_lora_B'])
+
+    print(f'LoRA weights loaded from {load_path}')

model.py

@@ -0,0 +1,95 @@
+import numpy as np
+import os
+import torch
+import torch.nn.functional as F
+from torch import nn
+from torchvision.models import alexnet
+
+import config as c
+from freia_funcs import permute_layer, glow_coupling_layer, F_fully_connected, ReversibleGraphNet, OutputNode, \
+    InputNode, Node
+
+WEIGHT_DIR = './weights'
+MODEL_DIR = './models'
+
+
+def nf_head(input_dim=c.n_feat):
+    nodes = list()
+    nodes.append(InputNode(input_dim, name='input'))
+    for k in range(c.n_coupling_blocks):
+        nodes.append(Node([nodes[-1].out0], permute_layer, {'seed': k}, name=F'permute_{k}'))
+        nodes.append(Node([nodes[-1].out0], glow_coupling_layer,
+                          {'clamp': c.clamp_alpha, 'F_class': F_fully_connected,
+                           'F_args': {'internal_size': c.fc_internal, 'dropout': c.dropout}},
+                          name=F'fc_{k}'))
+    nodes.append(OutputNode([nodes[-1].out0], name='output'))
+    coder = ReversibleGraphNet(nodes)
+    return coder
+
+
+class flow_model(nn.Module):
+    def __init__(self):
+        super(flow_model, self).__init__()
+
+        self.nf = nf_head(input_dim = 1024)
+
+    def forward(self, x):
+        z = self.nf(x)
+        return z
+    
+class flow_model_multi_fc(nn.Module):
+    def __init__(self):
+        super(flow_model_multi_fc, self).__init__()
+        self.fc1 = torch.nn.Linear(1024, 512)
+        self.relu = torch.nn.LeakyReLU(0.2)
+        self.fc2 = torch.nn.Linear(512, 256)
+
+        self.nf = nf_head(input_dim = 256)
+
+    def forward(self, x):
+        res_x = self.fc2(self.relu((self.fc1(x)))) 
+        z = self.nf(res_x)
+        return z
+    
+
+class DifferNet(nn.Module):
+    def __init__(self):
+        super(DifferNet, self).__init__()
+        self.feature_extractor = alexnet(pretrained=True)
+        self.nf = nf_head()
+
+    def forward(self, x):
+        y_cat = list()
+
+        for s in range(c.n_scales):
+            x_scaled = F.interpolate(x, size=c.img_size[0] // (2 ** s)) if s > 0 else x
+            feat_s = self.feature_extractor.features(x_scaled)
+            y_cat.append(torch.mean(feat_s, dim=(2, 3)))
+
+        y = torch.cat(y_cat, dim=1)
+        z = self.nf(y)
+        return z
+
+
+def save_model(model, filename):
+    if not os.path.exists(MODEL_DIR):
+        os.makedirs(MODEL_DIR)
+    torch.save(model, os.path.join(MODEL_DIR, filename))
+
+
+def load_model(filename):
+    path = os.path.join(MODEL_DIR, filename)
+    model = torch.load(path)
+    return model
+
+
+def save_weights(model, filename):
+    if not os.path.exists(WEIGHT_DIR):
+        os.makedirs(WEIGHT_DIR)
+    torch.save(model.state_dict(), os.path.join(WEIGHT_DIR, filename))
+
+
+def load_weights(model, filename):
+    path = os.path.join(WEIGHT_DIR, filename)
+    model.load_state_dict(torch.load(path))
+    return model

models/__init__.py

@@ -0,0 +1,43 @@
+from .clip_models import CLIPModel
+from .imagenet_models import ImagenetModel
+
+
+VALID_NAMES = [
+    'Imagenet:resnet18',
+    'Imagenet:resnet34',
+    'Imagenet:resnet50',
+    'Imagenet:resnet101',
+    'Imagenet:resnet152',
+    'Imagenet:vgg11',
+    'Imagenet:vgg19',
+    'Imagenet:swin-b',
+    'Imagenet:swin-s',
+    'Imagenet:swin-t',
+    'Imagenet:vit_b_16',
+    'Imagenet:vit_b_32',
+    'Imagenet:vit_l_16',
+    'Imagenet:vit_l_32',
+
+    'CLIP:RN50', 
+    'CLIP:RN101', 
+    'CLIP:RN50x4', 
+    'CLIP:RN50x16', 
+    'CLIP:RN50x64', 
+    'CLIP:ViT-B/32', 
+    'CLIP:ViT-B/16', 
+    'CLIP:ViT-L/14', 
+    'CLIP:ViT-L/14@336px',
+]
+
+
+
+
+
+def get_model(name):
+    assert name in VALID_NAMES
+    if name.startswith("Imagenet:"):
+        return ImagenetModel(name[9:]) 
+    elif name.startswith("CLIP:"):
+        return CLIPModel(name[5:])  
+    else:
+        assert False 

models/clip/__init__.py

@@ -0,0 +1 @@
+from .clip import *