+ ⚠️ This demo is for academic research and experiential use only.
+
+ """)
+ with gr.Row():
+ with gr.Column():
+ with gr.Accordion("IPT2V: Face Input", open=True):
+ image_input = gr.Image(label="Input Image (should contain clear face)")
+ prompt = gr.Textbox(label="Prompt (Less than 200 Words)", placeholder="Enter your prompt here", lines=5)
+ with gr.Accordion("Examples", open=False):
+ examples_component_images = gr.Examples(
+ examples_images,
+ inputs=[image_input, prompt],
+ cache_examples=False,
+ )
+
+ with gr.Group():
+ with gr.Column():
+ with gr.Row():
+ seed_param = gr.Number(
+ label="Inference Seed (Enter a positive number, -1 for random)", value=42
+ )
+ with gr.Row():
+ enable_scale = gr.Checkbox(label="Super-Resolution (720 × 480 -> 2880 × 1920)", value=False)
+ enable_rife = gr.Checkbox(label="Frame Interpolation (8fps -> 16fps)", value=False)
+ gr.Markdown(
+ "✨In this demo, we use [RIFE](https://github.com/hzwer/ECCV2022-RIFE) for frame interpolation and [Real-ESRGAN](https://github.com/xinntao/Real-ESRGAN) for upscaling(Super-Resolution)."
+ )
+
+ generate_button = gr.Button("🎬 Generate Video")
+
+ with gr.Column():
+ video_output = gr.Video(label="ConsisID Generate Video", width=720, height=480)
+ with gr.Row():
+ download_video_button = gr.File(label="📥 Download Video", visible=False)
+ download_gif_button = gr.File(label="📥 Download GIF", visible=False)
+ seed_text = gr.Number(label="Seed Used for Video Generation", visible=False)
+
+ gr.Markdown("""
+
+
+ 🎥 Video Gallery
+
+
+ |
+ The video features a woman in exquisite hybrid armor adorned with iridescent gemstones, standing amidst gently falling cherry blossoms. Her piercing yet serene gaze hints at quiet determination, as a breeze catches a loose strand of her hair. She stands in a tranquil courtyard framed by moss-covered stone walls and wooden arches, with blossoms casting soft shadows on the ground. The petals swirl around her, adding a dreamlike quality, while the blurred backdrop emphasizes her poised figure. The scene conveys elegance, strength, and tranquil readiness, capturing a moment of peace before an upcoming challenge.
+ |
+
+
+ |
+
+ The video features a baby wearing a bright superhero cape, standing confidently with arms raised in a powerful pose. The baby has a determined look on their face, with eyes wide and lips pursed in concentration, as if ready to take on a challenge. The setting appears playful, with colorful toys scattered around and a soft rug underfoot, while sunlight streams through a nearby window, highlighting the fluttering cape and adding to the impression of heroism. The overall atmosphere is lighthearted and fun, with the baby's expressions capturing a mix of innocence and an adorable attempt at bravery, as if truly ready to save the day.
+ |
+
+
+ |
+
+
+ |
+ The video features a man standing next to an airplane, engaged in a conversation on his cell phone. he is wearing sunglasses and a black top, and he appears to be talking seriously. The airplane has a green stripe running along its side, and there is a large engine visible behind his. The man seems to be standing near the entrance of the airplane, possibly preparing to board or just having disembarked. The setting suggests that he might be at an airport or a private airfield. The overall atmosphere of the video is professional and focused, with the man's attire and the presence of the airplane indicating a business or travel context.
+ |
+
+
+ |
+
+ The video features a woman with blonde hair standing on a beach near the water's edge. She is wearing a black swimsuit and appears to be enjoying her time by the sea. The sky above is clear with some clouds, and the ocean waves gently lap against the shore. The woman seems to be holding something white in her hand, possibly a piece of driftwood or a small object found on the beach. The overall atmosphere of the video is serene and relaxing, capturing the beauty of nature and the simple pleasure of being by the ocean.
+ |
+
+
+ |
+
+
+ |
+ The video features a man sitting in a red armchair, enjoying a cup of coffee or tea. he is dressed in a light-colored outfit and has long dark-haired hair. The setting appears to be indoors, with large windows providing a view of a misty or foggy coastal landscape outside. The room has a modern design with geometric structures visible in the background. There is a small round table next to the armchair, also holding a cup. The overall atmosphere suggests a calm and serene moment, possibly during a cold or rainy day by the sea.
+ |
+
+
+ |
+
+ The video shows a young boy sitting at a table, eating a piece of food. He appears to be enjoying his meal, as he takes a bite and chews it. The boy is wearing a blue shirt and has short hair. The background is dark, with some light coming from the left side of the frame. There is a straw visible on the right side of the frame, suggesting that there may be a drink next to the boy's plate. The overall atmosphere of the video seems casual and relaxed, with the focus on the boy's enjoyment of his food.
+ |
+
+
+ |
+
+
+ |
+ The video captures a boy walking along a city street, filmed in black and white on a classic 35mm camera. His expression is thoughtful, his brow slightly furrowed as if he's lost in contemplation. The film grain adds a textured, timeless quality to the image, evoking a sense of nostalgia. Around him, the cityscape is filled with vintage buildings, cobblestone sidewalks, and softly blurred figures passing by, their outlines faint and indistinct. Streetlights cast a gentle glow, while shadows play across the boy's path, adding depth to the scene. The lighting highlights the boy's subtle smile, hinting at a fleeting moment of curiosity. The overall cinematic atmosphere, complete with classic film still aesthetics and dramatic contrasts, gives the scene an evocative and introspective feel.
+ |
+
+
+ |
+
+ The video features a young man standing outdoors in a snowy park. he is wearing a colorful winter jacket with a floral pattern and a white knit hat. The background shows a snowy landscape with trees, benches, and a metal fence. The ground is covered in snow, and there is a light snowfall in the air. The man appears to be enjoying the winter weather, as he smiles and gives a thumbs-up gesture towards the camera. The overall atmosphere of the video is cheerful and festive, capturing the beauty of a snowy day in a park.
+ |
+
+
+ |
+
+
+ """)
+
+ def generate(
+ prompt,
+ image_input,
+ seed_value,
+ scale_status,
+ rife_status,
+ progress=gr.Progress(track_tqdm=True)
+ ):
+ latents, seed = infer(
+ prompt,
+ image_input,
+ num_inference_steps=50,
+ guidance_scale=7.0,
+ seed=seed_value,
+ progress=progress,
+ )
+ if scale_status:
+ latents = upscale_batch_and_concatenate(upscale_model, latents, device)
+ if rife_status:
+ latents = rife_inference_with_latents(frame_interpolation_model, latents)
+
+ batch_size = latents.shape[0]
+ batch_video_frames = []
+ for batch_idx in range(batch_size):
+ pt_image = latents[batch_idx]
+ pt_image = torch.stack([pt_image[i] for i in range(pt_image.shape[0])])
+
+ image_np = VaeImageProcessor.pt_to_numpy(pt_image)
+ image_pil = VaeImageProcessor.numpy_to_pil(image_np)
+ batch_video_frames.append(image_pil)
+
+ video_path = save_video(batch_video_frames[0], fps=math.ceil((len(batch_video_frames[0]) - 1) / 6))
+ video_update = gr.update(visible=True, value=video_path)
+ gif_path = convert_to_gif(video_path)
+ gif_update = gr.update(visible=True, value=gif_path)
+ seed_update = gr.update(visible=True, value=seed)
+
+ return video_path, video_update, gif_update, seed_update
+
+ generate_button.click(
+ generate,
+ inputs=[prompt, image_input, seed_param, enable_scale, enable_rife],
+ outputs=[video_output, download_video_button, download_gif_button, seed_text],
+ )
+
+if __name__ == "__main__":
+ demo.queue(max_size=15)
+ demo.launch()
diff --git a/asserts/example_images/1.png b/asserts/example_images/1.png
new file mode 100644
index 0000000000000000000000000000000000000000..846825db715fb8ad83e168ce76d5d032ebf548aa
--- /dev/null
+++ b/asserts/example_images/1.png
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:434856739faadf1c89bc38d8b940fcdbe027595de89645f721d8585fc2fe2459
+size 1639385
diff --git a/asserts/example_images/2.png b/asserts/example_images/2.png
new file mode 100644
index 0000000000000000000000000000000000000000..cf68b72128b70cf448e8a589e5e78f0ef10c3efc
Binary files /dev/null and b/asserts/example_images/2.png differ
diff --git a/asserts/example_images/3.png b/asserts/example_images/3.png
new file mode 100644
index 0000000000000000000000000000000000000000..422c567fa537359740343fa78b83d25b7b357d32
Binary files /dev/null and b/asserts/example_images/3.png differ
diff --git a/models/eva_clip/__init__.py b/models/eva_clip/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..fa2d014bbfe644b1e247758116bbf1b184738fe5
--- /dev/null
+++ b/models/eva_clip/__init__.py
@@ -0,0 +1,11 @@
+from .constants import OPENAI_DATASET_MEAN, OPENAI_DATASET_STD
+from .factory import create_model, create_model_and_transforms, create_model_from_pretrained, get_tokenizer, create_transforms
+from .factory import list_models, add_model_config, get_model_config, load_checkpoint
+from .loss import ClipLoss
+from .model import CLIP, CustomCLIP, CLIPTextCfg, CLIPVisionCfg,\
+ convert_weights_to_lp, convert_weights_to_fp16, trace_model, get_cast_dtype
+from .openai import load_openai_model, list_openai_models
+from .pretrained import list_pretrained, list_pretrained_models_by_tag, list_pretrained_tags_by_model,\
+ get_pretrained_url, download_pretrained_from_url, is_pretrained_cfg, get_pretrained_cfg, download_pretrained
+from .tokenizer import SimpleTokenizer, tokenize
+from .transform import image_transform
\ No newline at end of file
diff --git a/models/eva_clip/bpe_simple_vocab_16e6.txt.gz b/models/eva_clip/bpe_simple_vocab_16e6.txt.gz
new file mode 100644
index 0000000000000000000000000000000000000000..36a15856e00a06a9fbed8cdd34d2393fea4a3113
--- /dev/null
+++ b/models/eva_clip/bpe_simple_vocab_16e6.txt.gz
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:924691ac288e54409236115652ad4aa250f48203de50a9e4722a6ecd48d6804a
+size 1356917
diff --git a/models/eva_clip/constants.py b/models/eva_clip/constants.py
new file mode 100644
index 0000000000000000000000000000000000000000..a670bb3fab442baeb9af53b91c312e6982af57ee
--- /dev/null
+++ b/models/eva_clip/constants.py
@@ -0,0 +1,2 @@
+OPENAI_DATASET_MEAN = (0.48145466, 0.4578275, 0.40821073)
+OPENAI_DATASET_STD = (0.26862954, 0.26130258, 0.27577711)
diff --git a/models/eva_clip/eva_vit_model.py b/models/eva_clip/eva_vit_model.py
new file mode 100644
index 0000000000000000000000000000000000000000..51db88cf0c7b5d7a43f2be80bc59abb6c859c4b4
--- /dev/null
+++ b/models/eva_clip/eva_vit_model.py
@@ -0,0 +1,548 @@
+# --------------------------------------------------------
+# Adapted from https://github.com/microsoft/unilm/tree/master/beit
+# --------------------------------------------------------
+import math
+import os
+from functools import partial
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+try:
+ from timm.models.layers import drop_path, to_2tuple, trunc_normal_
+except:
+ from timm.layers import drop_path, to_2tuple, trunc_normal_
+
+from .transformer import PatchDropout
+from .rope import VisionRotaryEmbedding, VisionRotaryEmbeddingFast
+
+if os.getenv('ENV_TYPE') == 'deepspeed':
+ try:
+ from deepspeed.runtime.activation_checkpointing.checkpointing import checkpoint
+ except:
+ from torch.utils.checkpoint import checkpoint
+else:
+ from torch.utils.checkpoint import checkpoint
+
+try:
+ import xformers
+ import xformers.ops as xops
+ XFORMERS_IS_AVAILBLE = True
+except:
+ XFORMERS_IS_AVAILBLE = False
+
+class DropPath(nn.Module):
+ """Drop paths (Stochastic Depth) per sample (when applied in main path of residual blocks).
+ """
+ def __init__(self, drop_prob=None):
+ super(DropPath, self).__init__()
+ self.drop_prob = drop_prob
+
+ def forward(self, x):
+ return drop_path(x, self.drop_prob, self.training)
+
+ def extra_repr(self) -> str:
+ return 'p={}'.format(self.drop_prob)
+
+
+class Mlp(nn.Module):
+ def __init__(
+ self,
+ in_features,
+ hidden_features=None,
+ out_features=None,
+ act_layer=nn.GELU,
+ norm_layer=nn.LayerNorm,
+ drop=0.,
+ subln=False,
+
+ ):
+ super().__init__()
+ out_features = out_features or in_features
+ hidden_features = hidden_features or in_features
+ self.fc1 = nn.Linear(in_features, hidden_features)
+ self.act = act_layer()
+
+ self.ffn_ln = norm_layer(hidden_features) if subln else nn.Identity()
+
+ self.fc2 = nn.Linear(hidden_features, out_features)
+ self.drop = nn.Dropout(drop)
+
+ def forward(self, x):
+ x = self.fc1(x)
+ x = self.act(x)
+ # x = self.drop(x)
+ # commit this for the orignal BERT implement
+ x = self.ffn_ln(x)
+
+ x = self.fc2(x)
+ x = self.drop(x)
+ return x
+
+class SwiGLU(nn.Module):
+ def __init__(self, in_features, hidden_features=None, out_features=None, act_layer=nn.SiLU, drop=0.,
+ norm_layer=nn.LayerNorm, subln=False):
+ super().__init__()
+ out_features = out_features or in_features
+ hidden_features = hidden_features or in_features
+
+ self.w1 = nn.Linear(in_features, hidden_features)
+ self.w2 = nn.Linear(in_features, hidden_features)
+
+ self.act = act_layer()
+ self.ffn_ln = norm_layer(hidden_features) if subln else nn.Identity()
+ self.w3 = nn.Linear(hidden_features, out_features)
+
+ self.drop = nn.Dropout(drop)
+
+ def forward(self, x):
+ x1 = self.w1(x)
+ x2 = self.w2(x)
+ hidden = self.act(x1) * x2
+ x = self.ffn_ln(hidden)
+ x = self.w3(x)
+ x = self.drop(x)
+ return x
+
+class Attention(nn.Module):
+ def __init__(
+ self, dim, num_heads=8, qkv_bias=False, qk_scale=None, attn_drop=0.,
+ proj_drop=0., window_size=None, attn_head_dim=None, xattn=False, rope=None, subln=False, norm_layer=nn.LayerNorm):
+ super().__init__()
+ self.num_heads = num_heads
+ head_dim = dim // num_heads
+ if attn_head_dim is not None:
+ head_dim = attn_head_dim
+ all_head_dim = head_dim * self.num_heads
+ self.scale = qk_scale or head_dim ** -0.5
+
+ self.subln = subln
+ if self.subln:
+ self.q_proj = nn.Linear(dim, all_head_dim, bias=False)
+ self.k_proj = nn.Linear(dim, all_head_dim, bias=False)
+ self.v_proj = nn.Linear(dim, all_head_dim, bias=False)
+ else:
+ self.qkv = nn.Linear(dim, all_head_dim * 3, bias=False)
+
+ if qkv_bias:
+ self.q_bias = nn.Parameter(torch.zeros(all_head_dim))
+ self.v_bias = nn.Parameter(torch.zeros(all_head_dim))
+ else:
+ self.q_bias = None
+ self.v_bias = None
+
+ if window_size:
+ self.window_size = window_size
+ self.num_relative_distance = (2 * window_size[0] - 1) * (2 * window_size[1] - 1) + 3
+ self.relative_position_bias_table = nn.Parameter(
+ torch.zeros(self.num_relative_distance, num_heads)) # 2*Wh-1 * 2*Ww-1, nH
+ # cls to token & token 2 cls & cls to cls
+
+ # get pair-wise relative position index for each token inside the window
+ coords_h = torch.arange(window_size[0])
+ coords_w = torch.arange(window_size[1])
+ coords = torch.stack(torch.meshgrid([coords_h, coords_w])) # 2, Wh, Ww
+ coords_flatten = torch.flatten(coords, 1) # 2, Wh*Ww
+ relative_coords = coords_flatten[:, :, None] - coords_flatten[:, None, :] # 2, Wh*Ww, Wh*Ww
+ relative_coords = relative_coords.permute(1, 2, 0).contiguous() # Wh*Ww, Wh*Ww, 2
+ relative_coords[:, :, 0] += window_size[0] - 1 # shift to start from 0
+ relative_coords[:, :, 1] += window_size[1] - 1
+ relative_coords[:, :, 0] *= 2 * window_size[1] - 1
+ relative_position_index = \
+ torch.zeros(size=(window_size[0] * window_size[1] + 1, ) * 2, dtype=relative_coords.dtype)
+ relative_position_index[1:, 1:] = relative_coords.sum(-1) # Wh*Ww, Wh*Ww
+ relative_position_index[0, 0:] = self.num_relative_distance - 3
+ relative_position_index[0:, 0] = self.num_relative_distance - 2
+ relative_position_index[0, 0] = self.num_relative_distance - 1
+
+ self.register_buffer("relative_position_index", relative_position_index)
+ else:
+ self.window_size = None
+ self.relative_position_bias_table = None
+ self.relative_position_index = None
+
+ self.attn_drop = nn.Dropout(attn_drop)
+ self.inner_attn_ln = norm_layer(all_head_dim) if subln else nn.Identity()
+ # self.proj = nn.Linear(all_head_dim, all_head_dim)
+ self.proj = nn.Linear(all_head_dim, dim)
+ self.proj_drop = nn.Dropout(proj_drop)
+ self.xattn = xattn
+ self.xattn_drop = attn_drop
+
+ self.rope = rope
+
+ def forward(self, x, rel_pos_bias=None, attn_mask=None):
+ B, N, C = x.shape
+ if self.subln:
+ q = F.linear(input=x, weight=self.q_proj.weight, bias=self.q_bias)
+ k = F.linear(input=x, weight=self.k_proj.weight, bias=None)
+ v = F.linear(input=x, weight=self.v_proj.weight, bias=self.v_bias)
+
+ q = q.reshape(B, N, self.num_heads, -1).permute(0, 2, 1, 3) # B, num_heads, N, C
+ k = k.reshape(B, N, self.num_heads, -1).permute(0, 2, 1, 3)
+ v = v.reshape(B, N, self.num_heads, -1).permute(0, 2, 1, 3)
+ else:
+
+ qkv_bias = None
+ if self.q_bias is not None:
+ qkv_bias = torch.cat((self.q_bias, torch.zeros_like(self.v_bias, requires_grad=False), self.v_bias))
+
+ qkv = F.linear(input=x, weight=self.qkv.weight, bias=qkv_bias)
+ qkv = qkv.reshape(B, N, 3, self.num_heads, -1).permute(2, 0, 3, 1, 4) # 3, B, num_heads, N, C
+ q, k, v = qkv[0], qkv[1], qkv[2]
+
+ if self.rope:
+ # slightly fast impl
+ q_t = q[:, :, 1:, :]
+ ro_q_t = self.rope(q_t)
+ q = torch.cat((q[:, :, :1, :], ro_q_t), -2).type_as(v)
+
+ k_t = k[:, :, 1:, :]
+ ro_k_t = self.rope(k_t)
+ k = torch.cat((k[:, :, :1, :], ro_k_t), -2).type_as(v)
+
+ if self.xattn:
+ q = q.permute(0, 2, 1, 3) # B, num_heads, N, C -> B, N, num_heads, C
+ k = k.permute(0, 2, 1, 3)
+ v = v.permute(0, 2, 1, 3)
+
+ x = xops.memory_efficient_attention(
+ q, k, v,
+ p=self.xattn_drop,
+ scale=self.scale,
+ )
+ x = x.reshape(B, N, -1)
+ x = self.inner_attn_ln(x)
+ x = self.proj(x)
+ x = self.proj_drop(x)
+ else:
+ q = q * self.scale
+ attn = (q @ k.transpose(-2, -1))
+
+ if self.relative_position_bias_table is not None:
+ relative_position_bias = \
+ self.relative_position_bias_table[self.relative_position_index.view(-1)].view(
+ self.window_size[0] * self.window_size[1] + 1,
+ self.window_size[0] * self.window_size[1] + 1, -1) # Wh*Ww,Wh*Ww,nH
+ relative_position_bias = relative_position_bias.permute(2, 0, 1).contiguous() # nH, Wh*Ww, Wh*Ww
+ attn = attn + relative_position_bias.unsqueeze(0).type_as(attn)
+
+ if rel_pos_bias is not None:
+ attn = attn + rel_pos_bias.type_as(attn)
+
+ if attn_mask is not None:
+ attn_mask = attn_mask.bool()
+ attn = attn.masked_fill(~attn_mask[:, None, None, :], float("-inf"))
+
+ attn = attn.softmax(dim=-1)
+ attn = self.attn_drop(attn)
+
+ x = (attn @ v).transpose(1, 2).reshape(B, N, -1)
+ x = self.inner_attn_ln(x)
+ x = self.proj(x)
+ x = self.proj_drop(x)
+ return x
+
+
+class Block(nn.Module):
+
+ def __init__(self, dim, num_heads, mlp_ratio=4., qkv_bias=False, qk_scale=None, drop=0., attn_drop=0.,
+ drop_path=0., init_values=None, act_layer=nn.GELU, norm_layer=nn.LayerNorm,
+ window_size=None, attn_head_dim=None, xattn=False, rope=None, postnorm=False,
+ subln=False, naiveswiglu=False):
+ super().__init__()
+ self.norm1 = norm_layer(dim)
+ self.attn = Attention(
+ dim, num_heads=num_heads, qkv_bias=qkv_bias, qk_scale=qk_scale,
+ attn_drop=attn_drop, proj_drop=drop, window_size=window_size, attn_head_dim=attn_head_dim,
+ xattn=xattn, rope=rope, subln=subln, norm_layer=norm_layer)
+ # NOTE: drop path for stochastic depth, we shall see if this is better than dropout here
+ self.drop_path = DropPath(drop_path) if drop_path > 0. else nn.Identity()
+ self.norm2 = norm_layer(dim)
+ mlp_hidden_dim = int(dim * mlp_ratio)
+
+ if naiveswiglu:
+ self.mlp = SwiGLU(
+ in_features=dim,
+ hidden_features=mlp_hidden_dim,
+ subln=subln,
+ norm_layer=norm_layer,
+ )
+ else:
+ self.mlp = Mlp(
+ in_features=dim,
+ hidden_features=mlp_hidden_dim,
+ act_layer=act_layer,
+ subln=subln,
+ drop=drop
+ )
+
+ if init_values is not None and init_values > 0:
+ self.gamma_1 = nn.Parameter(init_values * torch.ones((dim)),requires_grad=True)
+ self.gamma_2 = nn.Parameter(init_values * torch.ones((dim)),requires_grad=True)
+ else:
+ self.gamma_1, self.gamma_2 = None, None
+
+ self.postnorm = postnorm
+
+ def forward(self, x, rel_pos_bias=None, attn_mask=None):
+ if self.gamma_1 is None:
+ if self.postnorm:
+ x = x + self.drop_path(self.norm1(self.attn(x, rel_pos_bias=rel_pos_bias, attn_mask=attn_mask)))
+ x = x + self.drop_path(self.norm2(self.mlp(x)))
+ else:
+ x = x + self.drop_path(self.attn(self.norm1(x), rel_pos_bias=rel_pos_bias, attn_mask=attn_mask))
+ x = x + self.drop_path(self.mlp(self.norm2(x)))
+ else:
+ if self.postnorm:
+ x = x + self.drop_path(self.gamma_1 * self.norm1(self.attn(x, rel_pos_bias=rel_pos_bias, attn_mask=attn_mask)))
+ x = x + self.drop_path(self.gamma_2 * self.norm2(self.mlp(x)))
+ else:
+ x = x + self.drop_path(self.gamma_1 * self.attn(self.norm1(x), rel_pos_bias=rel_pos_bias, attn_mask=attn_mask))
+ x = x + self.drop_path(self.gamma_2 * self.mlp(self.norm2(x)))
+ return x
+
+
+class PatchEmbed(nn.Module):
+ """ Image to Patch Embedding
+ """
+ def __init__(self, img_size=224, patch_size=16, in_chans=3, embed_dim=768):
+ super().__init__()
+ img_size = to_2tuple(img_size)
+ patch_size = to_2tuple(patch_size)
+ num_patches = (img_size[1] // patch_size[1]) * (img_size[0] // patch_size[0])
+ self.patch_shape = (img_size[0] // patch_size[0], img_size[1] // patch_size[1])
+ self.img_size = img_size
+ self.patch_size = patch_size
+ self.num_patches = num_patches
+
+ self.proj = nn.Conv2d(in_chans, embed_dim, kernel_size=patch_size, stride=patch_size)
+
+ def forward(self, x, **kwargs):
+ B, C, H, W = x.shape
+ # FIXME look at relaxing size constraints
+ assert H == self.img_size[0] and W == self.img_size[1], \
+ f"Input image size ({H}*{W}) doesn't match model ({self.img_size[0]}*{self.img_size[1]})."
+ x = self.proj(x).flatten(2).transpose(1, 2)
+ return x
+
+
+class RelativePositionBias(nn.Module):
+
+ def __init__(self, window_size, num_heads):
+ super().__init__()
+ self.window_size = window_size
+ self.num_relative_distance = (2 * window_size[0] - 1) * (2 * window_size[1] - 1) + 3
+ self.relative_position_bias_table = nn.Parameter(
+ torch.zeros(self.num_relative_distance, num_heads)) # 2*Wh-1 * 2*Ww-1, nH
+ # cls to token & token 2 cls & cls to cls
+
+ # get pair-wise relative position index for each token inside the window
+ coords_h = torch.arange(window_size[0])
+ coords_w = torch.arange(window_size[1])
+ coords = torch.stack(torch.meshgrid([coords_h, coords_w])) # 2, Wh, Ww
+ coords_flatten = torch.flatten(coords, 1) # 2, Wh*Ww
+ relative_coords = coords_flatten[:, :, None] - coords_flatten[:, None, :] # 2, Wh*Ww, Wh*Ww
+ relative_coords = relative_coords.permute(1, 2, 0).contiguous() # Wh*Ww, Wh*Ww, 2
+ relative_coords[:, :, 0] += window_size[0] - 1 # shift to start from 0
+ relative_coords[:, :, 1] += window_size[1] - 1
+ relative_coords[:, :, 0] *= 2 * window_size[1] - 1
+ relative_position_index = \
+ torch.zeros(size=(window_size[0] * window_size[1] + 1,) * 2, dtype=relative_coords.dtype)
+ relative_position_index[1:, 1:] = relative_coords.sum(-1) # Wh*Ww, Wh*Ww
+ relative_position_index[0, 0:] = self.num_relative_distance - 3
+ relative_position_index[0:, 0] = self.num_relative_distance - 2
+ relative_position_index[0, 0] = self.num_relative_distance - 1
+
+ self.register_buffer("relative_position_index", relative_position_index)
+
+ def forward(self):
+ relative_position_bias = \
+ self.relative_position_bias_table[self.relative_position_index.view(-1)].view(
+ self.window_size[0] * self.window_size[1] + 1,
+ self.window_size[0] * self.window_size[1] + 1, -1) # Wh*Ww,Wh*Ww,nH
+ return relative_position_bias.permute(2, 0, 1).contiguous() # nH, Wh*Ww, Wh*Ww
+
+
+class EVAVisionTransformer(nn.Module):
+ """ Vision Transformer with support for patch or hybrid CNN input stage
+ """
+ def __init__(self, img_size=224, patch_size=16, in_chans=3, num_classes=1000, embed_dim=768, depth=12,
+ num_heads=12, mlp_ratio=4., qkv_bias=False, qk_scale=None, drop_rate=0., attn_drop_rate=0.,
+ drop_path_rate=0., norm_layer=nn.LayerNorm, init_values=None, patch_dropout=0.,
+ use_abs_pos_emb=True, use_rel_pos_bias=False, use_shared_rel_pos_bias=False, rope=False,
+ use_mean_pooling=True, init_scale=0.001, grad_checkpointing=False, xattn=False, postnorm=False,
+ pt_hw_seq_len=16, intp_freq=False, naiveswiglu=False, subln=False):
+ super().__init__()
+
+ if not XFORMERS_IS_AVAILBLE:
+ xattn = False
+
+ self.image_size = img_size
+ self.num_classes = num_classes
+ self.num_features = self.embed_dim = embed_dim # num_features for consistency with other models
+
+ self.patch_embed = PatchEmbed(
+ img_size=img_size, patch_size=patch_size, in_chans=in_chans, embed_dim=embed_dim)
+ num_patches = self.patch_embed.num_patches
+
+ self.cls_token = nn.Parameter(torch.zeros(1, 1, embed_dim))
+ # self.mask_token = nn.Parameter(torch.zeros(1, 1, embed_dim))
+ if use_abs_pos_emb:
+ self.pos_embed = nn.Parameter(torch.zeros(1, num_patches + 1, embed_dim))
+ else:
+ self.pos_embed = None
+ self.pos_drop = nn.Dropout(p=drop_rate)
+
+ if use_shared_rel_pos_bias:
+ self.rel_pos_bias = RelativePositionBias(window_size=self.patch_embed.patch_shape, num_heads=num_heads)
+ else:
+ self.rel_pos_bias = None
+
+ if rope:
+ half_head_dim = embed_dim // num_heads // 2
+ hw_seq_len = img_size // patch_size
+ self.rope = VisionRotaryEmbeddingFast(
+ dim=half_head_dim,
+ pt_seq_len=pt_hw_seq_len,
+ ft_seq_len=hw_seq_len if intp_freq else None,
+ # patch_dropout=patch_dropout
+ )
+ else:
+ self.rope = None
+
+ self.naiveswiglu = naiveswiglu
+
+ dpr = [x.item() for x in torch.linspace(0, drop_path_rate, depth)] # stochastic depth decay rule
+ self.use_rel_pos_bias = use_rel_pos_bias
+ self.blocks = nn.ModuleList([
+ Block(
+ dim=embed_dim, num_heads=num_heads, mlp_ratio=mlp_ratio, qkv_bias=qkv_bias, qk_scale=qk_scale,
+ drop=drop_rate, attn_drop=attn_drop_rate, drop_path=dpr[i], norm_layer=norm_layer,
+ init_values=init_values, window_size=self.patch_embed.patch_shape if use_rel_pos_bias else None,
+ xattn=xattn, rope=self.rope, postnorm=postnorm, subln=subln, naiveswiglu=naiveswiglu)
+ for i in range(depth)])
+ self.norm = nn.Identity() if use_mean_pooling else norm_layer(embed_dim)
+ self.fc_norm = norm_layer(embed_dim) if use_mean_pooling else None
+ self.head = nn.Linear(embed_dim, num_classes) if num_classes > 0 else nn.Identity()
+
+ if self.pos_embed is not None:
+ trunc_normal_(self.pos_embed, std=.02)
+
+ trunc_normal_(self.cls_token, std=.02)
+ # trunc_normal_(self.mask_token, std=.02)
+
+ self.apply(self._init_weights)
+ self.fix_init_weight()
+
+ if isinstance(self.head, nn.Linear):
+ trunc_normal_(self.head.weight, std=.02)
+ self.head.weight.data.mul_(init_scale)
+ self.head.bias.data.mul_(init_scale)
+
+ # setting a patch_dropout of 0. would mean it is disabled and this function would be the identity fn
+ self.patch_dropout = PatchDropout(patch_dropout) if patch_dropout > 0. else nn.Identity()
+
+ self.grad_checkpointing = grad_checkpointing
+
+ def fix_init_weight(self):
+ def rescale(param, layer_id):
+ param.div_(math.sqrt(2.0 * layer_id))
+
+ for layer_id, layer in enumerate(self.blocks):
+ rescale(layer.attn.proj.weight.data, layer_id + 1)
+ if self.naiveswiglu:
+ rescale(layer.mlp.w3.weight.data, layer_id + 1)
+ else:
+ rescale(layer.mlp.fc2.weight.data, layer_id + 1)
+
+ def get_cast_dtype(self) -> torch.dtype:
+ return self.blocks[0].mlp.fc2.weight.dtype
+
+ def _init_weights(self, m):
+ if isinstance(m, nn.Linear):
+ trunc_normal_(m.weight, std=.02)
+ if m.bias is not None:
+ nn.init.constant_(m.bias, 0)
+ elif isinstance(m, nn.LayerNorm):
+ nn.init.constant_(m.bias, 0)
+ nn.init.constant_(m.weight, 1.0)
+
+ def get_num_layers(self):
+ return len(self.blocks)
+
+ def lock(self, unlocked_groups=0, freeze_bn_stats=False):
+ assert unlocked_groups == 0, 'partial locking not currently supported for this model'
+ for param in self.parameters():
+ param.requires_grad = False
+
+ @torch.jit.ignore
+ def set_grad_checkpointing(self, enable=True):
+ self.grad_checkpointing = enable
+
+ @torch.jit.ignore
+ def no_weight_decay(self):
+ return {'pos_embed', 'cls_token'}
+
+ def get_classifier(self):
+ return self.head
+
+ def reset_classifier(self, num_classes, global_pool=''):
+ self.num_classes = num_classes
+ self.head = nn.Linear(self.embed_dim, num_classes) if num_classes > 0 else nn.Identity()
+
+ def forward_features(self, x, return_all_features=False, return_hidden=False, shuffle=False):
+
+ x = self.patch_embed(x)
+ batch_size, seq_len, _ = x.size()
+
+ if shuffle:
+ idx = torch.randperm(x.shape[1]) + 1
+ zero = torch.LongTensor([0, ])
+ idx = torch.cat([zero, idx])
+ pos_embed = self.pos_embed[:, idx]
+
+ cls_tokens = self.cls_token.expand(batch_size, -1, -1) # stole cls_tokens impl from Phil Wang, thanks
+ x = torch.cat((cls_tokens, x), dim=1)
+ if shuffle:
+ x = x + pos_embed
+ elif self.pos_embed is not None:
+ x = x + self.pos_embed
+ x = self.pos_drop(x)
+
+ # a patch_dropout of 0. would mean it is disabled and this function would do nothing but return what was passed in
+ if os.getenv('RoPE') == '1':
+ if self.training and not isinstance(self.patch_dropout, nn.Identity):
+ x, patch_indices_keep = self.patch_dropout(x)
+ self.rope.forward = partial(self.rope.forward, patch_indices_keep=patch_indices_keep)
+ else:
+ self.rope.forward = partial(self.rope.forward, patch_indices_keep=None)
+ x = self.patch_dropout(x)
+ else:
+ x = self.patch_dropout(x)
+
+ rel_pos_bias = self.rel_pos_bias() if self.rel_pos_bias is not None else None
+ hidden_states = []
+ for idx, blk in enumerate(self.blocks):
+ if (0 < idx <= 20) and (idx % 4 == 0) and return_hidden:
+ hidden_states.append(x)
+ if self.grad_checkpointing:
+ x = checkpoint(blk, x, (rel_pos_bias,))
+ else:
+ x = blk(x, rel_pos_bias=rel_pos_bias)
+
+ if not return_all_features:
+ x = self.norm(x)
+ if self.fc_norm is not None:
+ return self.fc_norm(x.mean(1)), hidden_states
+ else:
+ return x[:, 0], hidden_states
+ return x
+
+ def forward(self, x, return_all_features=False, return_hidden=False, shuffle=False):
+ if return_all_features:
+ return self.forward_features(x, return_all_features, return_hidden, shuffle)
+ x, hidden_states = self.forward_features(x, return_all_features, return_hidden, shuffle)
+ x = self.head(x)
+ if return_hidden:
+ return x, hidden_states
+ return x
diff --git a/models/eva_clip/factory.py b/models/eva_clip/factory.py
new file mode 100644
index 0000000000000000000000000000000000000000..ced8999997bf374b69f846bc73ea635fe8a6eb63
--- /dev/null
+++ b/models/eva_clip/factory.py
@@ -0,0 +1,517 @@
+import json
+import logging
+import os
+import pathlib
+import re
+from copy import deepcopy
+from pathlib import Path
+from typing import Optional, Tuple, Union, Dict, Any
+import torch
+
+from .constants import OPENAI_DATASET_MEAN, OPENAI_DATASET_STD
+from .model import CLIP, CustomCLIP, convert_weights_to_lp, convert_to_custom_text_state_dict,\
+ get_cast_dtype
+from .openai import load_openai_model
+from .pretrained import is_pretrained_cfg, get_pretrained_cfg, download_pretrained, list_pretrained_tags_by_model
+from .transform import image_transform
+from .tokenizer import HFTokenizer, tokenize
+from .utils import resize_clip_pos_embed, resize_evaclip_pos_embed, resize_visual_pos_embed, resize_eva_pos_embed
+
+
+_MODEL_CONFIG_PATHS = [Path(__file__).parent / f"model_configs/"]
+_MODEL_CONFIGS = {} # directory (model_name: config) of model architecture configs
+
+
+def _natural_key(string_):
+ return [int(s) if s.isdigit() else s for s in re.split(r'(\d+)', string_.lower())]
+
+
+def _rescan_model_configs():
+ global _MODEL_CONFIGS
+
+ config_ext = ('.json',)
+ config_files = []
+ for config_path in _MODEL_CONFIG_PATHS:
+ if config_path.is_file() and config_path.suffix in config_ext:
+ config_files.append(config_path)
+ elif config_path.is_dir():
+ for ext in config_ext:
+ config_files.extend(config_path.glob(f'*{ext}'))
+
+ for cf in config_files:
+ with open(cf, "r", encoding="utf8") as f:
+ model_cfg = json.load(f)
+ if all(a in model_cfg for a in ('embed_dim', 'vision_cfg', 'text_cfg')):
+ _MODEL_CONFIGS[cf.stem] = model_cfg
+
+ _MODEL_CONFIGS = dict(sorted(_MODEL_CONFIGS.items(), key=lambda x: _natural_key(x[0])))
+
+
+_rescan_model_configs() # initial populate of model config registry
+
+
+def list_models():
+ """ enumerate available model architectures based on config files """
+ return list(_MODEL_CONFIGS.keys())
+
+
+def add_model_config(path):
+ """ add model config path or file and update registry """
+ if not isinstance(path, Path):
+ path = Path(path)
+ _MODEL_CONFIG_PATHS.append(path)
+ _rescan_model_configs()
+
+
+def get_model_config(model_name):
+ if model_name in _MODEL_CONFIGS:
+ return deepcopy(_MODEL_CONFIGS[model_name])
+ else:
+ return None
+
+
+def get_tokenizer(model_name):
+ config = get_model_config(model_name)
+ tokenizer = HFTokenizer(config['text_cfg']['hf_tokenizer_name']) if 'hf_tokenizer_name' in config['text_cfg'] else tokenize
+ return tokenizer
+
+
+# loading openai CLIP weights when is_openai=True for training
+def load_state_dict(checkpoint_path: str, map_location: str='cpu', model_key: str='model|module|state_dict', is_openai: bool=False, skip_list: list=[]):
+ if is_openai:
+ model = torch.jit.load(checkpoint_path, map_location="cpu").eval()
+ state_dict = model.state_dict()
+ for key in ["input_resolution", "context_length", "vocab_size"]:
+ state_dict.pop(key, None)
+ else:
+ checkpoint = torch.load(checkpoint_path, map_location=map_location)
+ for mk in model_key.split('|'):
+ if isinstance(checkpoint, dict) and mk in checkpoint:
+ state_dict = checkpoint[mk]
+ break
+ else:
+ state_dict = checkpoint
+ if next(iter(state_dict.items()))[0].startswith('module'):
+ state_dict = {k[7:]: v for k, v in state_dict.items()}
+
+ for k in skip_list:
+ if k in list(state_dict.keys()):
+ logging.info(f"Removing key {k} from pretrained checkpoint")
+ del state_dict[k]
+
+ if os.getenv('RoPE') == '1':
+ for k in list(state_dict.keys()):
+ if 'freqs_cos' in k or 'freqs_sin' in k:
+ del state_dict[k]
+ return state_dict
+
+
+
+def load_checkpoint(model, checkpoint_path, model_key="model|module|state_dict", strict=True):
+ state_dict = load_state_dict(checkpoint_path, model_key=model_key, is_openai=False)
+ # detect old format and make compatible with new format
+ if 'positional_embedding' in state_dict and not hasattr(model, 'positional_embedding'):
+ state_dict = convert_to_custom_text_state_dict(state_dict)
+ if 'text.logit_scale' in state_dict and hasattr(model, 'logit_scale'):
+ state_dict['logit_scale'] = state_dict['text.logit_scale']
+ del state_dict['text.logit_scale']
+
+ # resize_clip_pos_embed for CLIP and open CLIP
+ if 'visual.positional_embedding' in state_dict:
+ resize_clip_pos_embed(state_dict, model)
+ # specified to eva_vit_model
+ elif 'visual.pos_embed' in state_dict:
+ resize_evaclip_pos_embed(state_dict, model)
+
+ # resize_clip_pos_embed(state_dict, model)
+ incompatible_keys = model.load_state_dict(state_dict, strict=strict)
+ logging.info(f"incompatible_keys.missing_keys: {incompatible_keys.missing_keys}")
+ return incompatible_keys
+
+def load_clip_visual_state_dict(checkpoint_path: str, map_location: str='cpu', is_openai: bool=False, skip_list:list=[]):
+ state_dict = load_state_dict(checkpoint_path, map_location=map_location, is_openai=is_openai, skip_list=skip_list)
+
+ for k in list(state_dict.keys()):
+ if not k.startswith('visual.'):
+ del state_dict[k]
+ for k in list(state_dict.keys()):
+ if k.startswith('visual.'):
+ new_k = k[7:]
+ state_dict[new_k] = state_dict[k]
+ del state_dict[k]
+ return state_dict
+
+def load_clip_text_state_dict(checkpoint_path: str, map_location: str='cpu', is_openai: bool=False, skip_list:list=[]):
+ state_dict = load_state_dict(checkpoint_path, map_location=map_location, is_openai=is_openai, skip_list=skip_list)
+
+ for k in list(state_dict.keys()):
+ if k.startswith('visual.'):
+ del state_dict[k]
+ return state_dict
+
+def get_pretrained_tag(pretrained_model):
+ pretrained_model = pretrained_model.lower()
+ if "laion" in pretrained_model or "open_clip" in pretrained_model:
+ return "open_clip"
+ elif "openai" in pretrained_model:
+ return "clip"
+ elif "eva" in pretrained_model and "clip" in pretrained_model:
+ return "eva_clip"
+ else:
+ return "other"
+
+def load_pretrained_checkpoint(
+ model,
+ visual_checkpoint_path,
+ text_checkpoint_path,
+ strict=True,
+ visual_model=None,
+ text_model=None,
+ model_key="model|module|state_dict",
+ skip_list=[]):
+ visual_tag = get_pretrained_tag(visual_model)
+ text_tag = get_pretrained_tag(text_model)
+
+ logging.info(f"num of model state_dict keys: {len(model.state_dict().keys())}")
+ visual_incompatible_keys, text_incompatible_keys = None, None
+ if visual_checkpoint_path:
+ if visual_tag == "eva_clip" or visual_tag == "open_clip":
+ visual_state_dict = load_clip_visual_state_dict(visual_checkpoint_path, is_openai=False, skip_list=skip_list)
+ elif visual_tag == "clip":
+ visual_state_dict = load_clip_visual_state_dict(visual_checkpoint_path, is_openai=True, skip_list=skip_list)
+ else:
+ visual_state_dict = load_state_dict(visual_checkpoint_path, model_key=model_key, is_openai=False, skip_list=skip_list)
+
+ # resize_clip_pos_embed for CLIP and open CLIP
+ if 'positional_embedding' in visual_state_dict:
+ resize_visual_pos_embed(visual_state_dict, model)
+ # specified to EVA model
+ elif 'pos_embed' in visual_state_dict:
+ resize_eva_pos_embed(visual_state_dict, model)
+
+ visual_incompatible_keys = model.visual.load_state_dict(visual_state_dict, strict=strict)
+ logging.info(f"num of loaded visual_state_dict keys: {len(visual_state_dict.keys())}")
+ logging.info(f"visual_incompatible_keys.missing_keys: {visual_incompatible_keys.missing_keys}")
+
+ if text_checkpoint_path:
+ if text_tag == "eva_clip" or text_tag == "open_clip":
+ text_state_dict = load_clip_text_state_dict(text_checkpoint_path, is_openai=False, skip_list=skip_list)
+ elif text_tag == "clip":
+ text_state_dict = load_clip_text_state_dict(text_checkpoint_path, is_openai=True, skip_list=skip_list)
+ else:
+ text_state_dict = load_state_dict(visual_checkpoint_path, model_key=model_key, is_openai=False, skip_list=skip_list)
+
+ text_incompatible_keys = model.text.load_state_dict(text_state_dict, strict=strict)
+
+ logging.info(f"num of loaded text_state_dict keys: {len(text_state_dict.keys())}")
+ logging.info(f"text_incompatible_keys.missing_keys: {text_incompatible_keys.missing_keys}")
+
+ return visual_incompatible_keys, text_incompatible_keys
+
+def create_model(
+ model_name: str,
+ pretrained: Optional[str] = None,
+ precision: str = 'fp32',
+ device: Union[str, torch.device] = 'cpu',
+ jit: bool = False,
+ force_quick_gelu: bool = False,
+ force_custom_clip: bool = False,
+ force_patch_dropout: Optional[float] = None,
+ pretrained_image: str = '',
+ pretrained_text: str = '',
+ pretrained_hf: bool = True,
+ pretrained_visual_model: str = None,
+ pretrained_text_model: str = None,
+ cache_dir: Optional[str] = None,
+ skip_list: list = [],
+):
+ model_name = model_name.replace('/', '-') # for callers using old naming with / in ViT names
+ if isinstance(device, str):
+ device = torch.device(device)
+
+ if pretrained and pretrained.lower() == 'openai':
+ logging.info(f'Loading pretrained {model_name} from OpenAI.')
+ model = load_openai_model(
+ model_name,
+ precision=precision,
+ device=device,
+ jit=jit,
+ cache_dir=cache_dir,
+ )
+ else:
+ model_cfg = get_model_config(model_name)
+ if model_cfg is not None:
+ logging.info(f'Loaded {model_name} model config.')
+ else:
+ logging.error(f'Model config for {model_name} not found; available models {list_models()}.')
+ raise RuntimeError(f'Model config for {model_name} not found.')
+
+ if 'rope' in model_cfg.get('vision_cfg', {}):
+ if model_cfg['vision_cfg']['rope']:
+ os.environ['RoPE'] = "1"
+ else:
+ os.environ['RoPE'] = "0"
+
+ if force_quick_gelu:
+ # override for use of QuickGELU on non-OpenAI transformer models
+ model_cfg["quick_gelu"] = True
+
+ if force_patch_dropout is not None:
+ # override the default patch dropout value
+ model_cfg['vision_cfg']["patch_dropout"] = force_patch_dropout
+
+ cast_dtype = get_cast_dtype(precision)
+ custom_clip = model_cfg.pop('custom_text', False) or force_custom_clip or ('hf_model_name' in model_cfg['text_cfg'])
+
+
+ if custom_clip:
+ if 'hf_model_name' in model_cfg.get('text_cfg', {}):
+ model_cfg['text_cfg']['hf_model_pretrained'] = pretrained_hf
+ model = CustomCLIP(**model_cfg, cast_dtype=cast_dtype)
+ else:
+ model = CLIP(**model_cfg, cast_dtype=cast_dtype)
+
+ pretrained_cfg = {}
+ if pretrained:
+ checkpoint_path = ''
+ pretrained_cfg = get_pretrained_cfg(model_name, pretrained)
+ if pretrained_cfg:
+ checkpoint_path = download_pretrained(pretrained_cfg, cache_dir=cache_dir)
+ elif os.path.exists(pretrained):
+ checkpoint_path = pretrained
+
+ if checkpoint_path:
+ logging.info(f'Loading pretrained {model_name} weights ({pretrained}).')
+ load_checkpoint(model,
+ checkpoint_path,
+ model_key="model|module|state_dict",
+ strict=False
+ )
+ else:
+ error_str = (
+ f'Pretrained weights ({pretrained}) not found for model {model_name}.'
+ f'Available pretrained tags ({list_pretrained_tags_by_model(model_name)}.')
+ logging.warning(error_str)
+ raise RuntimeError(error_str)
+ else:
+ visual_checkpoint_path = ''
+ text_checkpoint_path = ''
+
+ if pretrained_image:
+ pretrained_visual_model = pretrained_visual_model.replace('/', '-') # for callers using old naming with / in ViT names
+ pretrained_image_cfg = get_pretrained_cfg(pretrained_visual_model, pretrained_image)
+ if 'timm_model_name' in model_cfg.get('vision_cfg', {}):
+ # pretrained weight loading for timm models set via vision_cfg
+ model_cfg['vision_cfg']['timm_model_pretrained'] = True
+ elif pretrained_image_cfg:
+ visual_checkpoint_path = download_pretrained(pretrained_image_cfg, cache_dir=cache_dir)
+ elif os.path.exists(pretrained_image):
+ visual_checkpoint_path = pretrained_image
+ else:
+ logging.warning(f'Pretrained weights ({visual_checkpoint_path}) not found for model {model_name}.visual.')
+ raise RuntimeError(f'Pretrained weights ({visual_checkpoint_path}) not found for model {model_name}.visual.')
+
+ if pretrained_text:
+ pretrained_text_model = pretrained_text_model.replace('/', '-') # for callers using old naming with / in ViT names
+ pretrained_text_cfg = get_pretrained_cfg(pretrained_text_model, pretrained_text)
+ if pretrained_image_cfg:
+ text_checkpoint_path = download_pretrained(pretrained_text_cfg, cache_dir=cache_dir)
+ elif os.path.exists(pretrained_text):
+ text_checkpoint_path = pretrained_text
+ else:
+ logging.warning(f'Pretrained weights ({text_checkpoint_path}) not found for model {model_name}.text.')
+ raise RuntimeError(f'Pretrained weights ({text_checkpoint_path}) not found for model {model_name}.text.')
+
+ if visual_checkpoint_path:
+ logging.info(f'Loading pretrained {model_name}.visual weights ({visual_checkpoint_path}).')
+ if text_checkpoint_path:
+ logging.info(f'Loading pretrained {model_name}.text weights ({text_checkpoint_path}).')
+
+ if visual_checkpoint_path or text_checkpoint_path:
+ load_pretrained_checkpoint(
+ model,
+ visual_checkpoint_path,
+ text_checkpoint_path,
+ strict=False,
+ visual_model=pretrained_visual_model,
+ text_model=pretrained_text_model,
+ model_key="model|module|state_dict",
+ skip_list=skip_list
+ )
+
+ if "fp16" in precision or "bf16" in precision:
+ logging.info(f'convert precision to {precision}')
+ model = model.to(torch.bfloat16) if 'bf16' in precision else model.to(torch.float16)
+
+ model.to(device=device)
+
+ # set image / mean metadata from pretrained_cfg if available, or use default
+ model.visual.image_mean = pretrained_cfg.get('mean', None) or OPENAI_DATASET_MEAN
+ model.visual.image_std = pretrained_cfg.get('std', None) or OPENAI_DATASET_STD
+
+ if jit:
+ model = torch.jit.script(model)
+
+ return model
+
+
+def create_model_and_transforms(
+ model_name: str,
+ pretrained: Optional[str] = None,
+ precision: str = 'fp32',
+ device: Union[str, torch.device] = 'cpu',
+ jit: bool = False,
+ force_quick_gelu: bool = False,
+ force_custom_clip: bool = False,
+ force_patch_dropout: Optional[float] = None,
+ pretrained_image: str = '',
+ pretrained_text: str = '',
+ pretrained_hf: bool = True,
+ pretrained_visual_model: str = None,
+ pretrained_text_model: str = None,
+ image_mean: Optional[Tuple[float, ...]] = None,
+ image_std: Optional[Tuple[float, ...]] = None,
+ cache_dir: Optional[str] = None,
+ skip_list: list = [],
+):
+ model = create_model(
+ model_name,
+ pretrained,
+ precision=precision,
+ device=device,
+ jit=jit,
+ force_quick_gelu=force_quick_gelu,
+ force_custom_clip=force_custom_clip,
+ force_patch_dropout=force_patch_dropout,
+ pretrained_image=pretrained_image,
+ pretrained_text=pretrained_text,
+ pretrained_hf=pretrained_hf,
+ pretrained_visual_model=pretrained_visual_model,
+ pretrained_text_model=pretrained_text_model,
+ cache_dir=cache_dir,
+ skip_list=skip_list,
+ )
+
+ image_mean = image_mean or getattr(model.visual, 'image_mean', None)
+ image_std = image_std or getattr(model.visual, 'image_std', None)
+ preprocess_train = image_transform(
+ model.visual.image_size,
+ is_train=True,
+ mean=image_mean,
+ std=image_std
+ )
+ preprocess_val = image_transform(
+ model.visual.image_size,
+ is_train=False,
+ mean=image_mean,
+ std=image_std
+ )
+
+ return model, preprocess_train, preprocess_val
+
+
+def create_transforms(
+ model_name: str,
+ pretrained: Optional[str] = None,
+ precision: str = 'fp32',
+ device: Union[str, torch.device] = 'cpu',
+ jit: bool = False,
+ force_quick_gelu: bool = False,
+ force_custom_clip: bool = False,
+ force_patch_dropout: Optional[float] = None,
+ pretrained_image: str = '',
+ pretrained_text: str = '',
+ pretrained_hf: bool = True,
+ pretrained_visual_model: str = None,
+ pretrained_text_model: str = None,
+ image_mean: Optional[Tuple[float, ...]] = None,
+ image_std: Optional[Tuple[float, ...]] = None,
+ cache_dir: Optional[str] = None,
+ skip_list: list = [],
+):
+ model = create_model(
+ model_name,
+ pretrained,
+ precision=precision,
+ device=device,
+ jit=jit,
+ force_quick_gelu=force_quick_gelu,
+ force_custom_clip=force_custom_clip,
+ force_patch_dropout=force_patch_dropout,
+ pretrained_image=pretrained_image,
+ pretrained_text=pretrained_text,
+ pretrained_hf=pretrained_hf,
+ pretrained_visual_model=pretrained_visual_model,
+ pretrained_text_model=pretrained_text_model,
+ cache_dir=cache_dir,
+ skip_list=skip_list,
+ )
+
+
+ image_mean = image_mean or getattr(model.visual, 'image_mean', None)
+ image_std = image_std or getattr(model.visual, 'image_std', None)
+ preprocess_train = image_transform(
+ model.visual.image_size,
+ is_train=True,
+ mean=image_mean,
+ std=image_std
+ )
+ preprocess_val = image_transform(
+ model.visual.image_size,
+ is_train=False,
+ mean=image_mean,
+ std=image_std
+ )
+ del model
+
+ return preprocess_train, preprocess_val
+
+def create_model_from_pretrained(
+ model_name: str,
+ pretrained: str,
+ precision: str = 'fp32',
+ device: Union[str, torch.device] = 'cpu',
+ jit: bool = False,
+ force_quick_gelu: bool = False,
+ force_custom_clip: bool = False,
+ force_patch_dropout: Optional[float] = None,
+ return_transform: bool = True,
+ image_mean: Optional[Tuple[float, ...]] = None,
+ image_std: Optional[Tuple[float, ...]] = None,
+ cache_dir: Optional[str] = None,
+ is_frozen: bool = False,
+):
+ if not is_pretrained_cfg(model_name, pretrained) and not os.path.exists(pretrained):
+ raise RuntimeError(
+ f'{pretrained} is not a valid pretrained cfg or checkpoint for {model_name}.'
+ f' Use open_clip.list_pretrained() to find one.')
+
+ model = create_model(
+ model_name,
+ pretrained,
+ precision=precision,
+ device=device,
+ jit=jit,
+ force_quick_gelu=force_quick_gelu,
+ force_custom_clip=force_custom_clip,
+ force_patch_dropout=force_patch_dropout,
+ cache_dir=cache_dir,
+ )
+
+ if is_frozen:
+ for param in model.parameters():
+ param.requires_grad = False
+
+ if not return_transform:
+ return model
+
+ image_mean = image_mean or getattr(model.visual, 'image_mean', None)
+ image_std = image_std or getattr(model.visual, 'image_std', None)
+ preprocess = image_transform(
+ model.visual.image_size,
+ is_train=False,
+ mean=image_mean,
+ std=image_std
+ )
+
+ return model, preprocess
diff --git a/models/eva_clip/hf_configs.py b/models/eva_clip/hf_configs.py
new file mode 100644
index 0000000000000000000000000000000000000000..a8c9b704db1879676aed5cef26796303b65fe987
--- /dev/null
+++ b/models/eva_clip/hf_configs.py
@@ -0,0 +1,57 @@
+# HF architecture dict:
+arch_dict = {
+ # https://huggingface.co/docs/transformers/model_doc/roberta#roberta
+ "roberta": {
+ "config_names": {
+ "context_length": "max_position_embeddings",
+ "vocab_size": "vocab_size",
+ "width": "hidden_size",
+ "heads": "num_attention_heads",
+ "layers": "num_hidden_layers",
+ "layer_attr": "layer",
+ "token_embeddings_attr": "embeddings"
+ },
+ "pooler": "mean_pooler",
+ },
+ # https://huggingface.co/docs/transformers/model_doc/xlm-roberta#transformers.XLMRobertaConfig
+ "xlm-roberta": {
+ "config_names": {
+ "context_length": "max_position_embeddings",
+ "vocab_size": "vocab_size",
+ "width": "hidden_size",
+ "heads": "num_attention_heads",
+ "layers": "num_hidden_layers",
+ "layer_attr": "layer",
+ "token_embeddings_attr": "embeddings"
+ },
+ "pooler": "mean_pooler",
+ },
+ # https://huggingface.co/docs/transformers/model_doc/mt5#mt5
+ "mt5": {
+ "config_names": {
+ # unlimited seqlen
+ # https://github.com/google-research/text-to-text-transfer-transformer/issues/273
+ # https://github.com/huggingface/transformers/blob/v4.24.0/src/transformers/models/t5/modeling_t5.py#L374
+ "context_length": "",
+ "vocab_size": "vocab_size",
+ "width": "d_model",
+ "heads": "num_heads",
+ "layers": "num_layers",
+ "layer_attr": "block",
+ "token_embeddings_attr": "embed_tokens"
+ },
+ "pooler": "mean_pooler",
+ },
+ "bert": {
+ "config_names": {
+ "context_length": "max_position_embeddings",
+ "vocab_size": "vocab_size",
+ "width": "hidden_size",
+ "heads": "num_attention_heads",
+ "layers": "num_hidden_layers",
+ "layer_attr": "layer",
+ "token_embeddings_attr": "embeddings"
+ },
+ "pooler": "mean_pooler",
+ }
+}
diff --git a/models/eva_clip/hf_model.py b/models/eva_clip/hf_model.py
new file mode 100644
index 0000000000000000000000000000000000000000..c4b9fd85b4066ba31db2bda5767ed1ce15de479d
--- /dev/null
+++ b/models/eva_clip/hf_model.py
@@ -0,0 +1,248 @@
+""" huggingface model adapter
+
+Wraps HuggingFace transformers (https://github.com/huggingface/transformers) models for use as a text tower in CLIP model.
+"""
+
+import re
+
+import torch
+import torch.nn as nn
+from torch.nn import functional as F
+from torch import TensorType
+try:
+ import transformers
+ from transformers import AutoModel, AutoModelForMaskedLM, AutoTokenizer, AutoConfig, PretrainedConfig
+ from transformers.modeling_outputs import BaseModelOutput, BaseModelOutputWithPooling, \
+ BaseModelOutputWithPoolingAndCrossAttentions
+except ImportError as e:
+ transformers = None
+
+
+ class BaseModelOutput:
+ pass
+
+
+ class PretrainedConfig:
+ pass
+
+from .hf_configs import arch_dict
+
+# utils
+def _camel2snake(s):
+ return re.sub(r'(? TensorType:
+ # image_atts = torch.ones(image_embeds.size()[:-1],dtype=torch.long).to(x.device)
+ # attn_mask = (x != self.config.pad_token_id).long()
+ # out = self.transformer(
+ # input_ids=x,
+ # attention_mask=attn_mask,
+ # encoder_hidden_states = image_embeds,
+ # encoder_attention_mask = image_atts,
+ # )
+ # pooled_out = self.pooler(out, attn_mask)
+
+ # return self.itm_proj(pooled_out)
+
+ def mask(self, input_ids, vocab_size, device, targets=None, masked_indices=None, probability_matrix=None):
+ if masked_indices is None:
+ masked_indices = torch.bernoulli(probability_matrix).bool()
+
+ masked_indices[input_ids == self.tokenizer.pad_token_id] = False
+ masked_indices[input_ids == self.tokenizer.cls_token_id] = False
+
+ if targets is not None:
+ targets[~masked_indices] = -100 # We only compute loss on masked tokens
+
+ # 80% of the time, we replace masked input tokens with tokenizer.mask_token ([MASK])
+ indices_replaced = torch.bernoulli(torch.full(input_ids.shape, 0.8)).bool() & masked_indices
+ input_ids[indices_replaced] = self.tokenizer.mask_token_id
+
+ # 10% of the time, we replace masked input tokens with random word
+ indices_random = torch.bernoulli(torch.full(input_ids.shape, 0.5)).bool() & masked_indices & ~indices_replaced
+ random_words = torch.randint(vocab_size, input_ids.shape, dtype=torch.long).to(device)
+ input_ids[indices_random] = random_words[indices_random]
+ # The rest of the time (10% of the time) we keep the masked input tokens unchanged
+
+ if targets is not None:
+ return input_ids, targets
+ else:
+ return input_ids
+
+ def forward_mlm(self, input_ids, image_embeds, mlm_probability=0.25):
+ labels = input_ids.clone()
+ attn_mask = (input_ids != self.config.pad_token_id).long()
+ image_atts = torch.ones(image_embeds.size()[:-1],dtype=torch.long).to(input_ids.device)
+ vocab_size = getattr(self.config, arch_dict[self.config.model_type]["config_names"]["vocab_size"])
+ probability_matrix = torch.full(labels.shape, mlm_probability)
+ input_ids, labels = self.mask(input_ids, vocab_size, input_ids.device, targets=labels,
+ probability_matrix = probability_matrix)
+ mlm_output = self.transformer(input_ids,
+ attention_mask = attn_mask,
+ encoder_hidden_states = image_embeds,
+ encoder_attention_mask = image_atts,
+ return_dict = True,
+ labels = labels,
+ )
+ return mlm_output.loss
+ # mlm_output = self.transformer(input_ids,
+ # attention_mask = attn_mask,
+ # encoder_hidden_states = image_embeds,
+ # encoder_attention_mask = image_atts,
+ # return_dict = True,
+ # ).last_hidden_state
+ # logits = self.mlm_proj(mlm_output)
+
+ # # logits = logits[:, :-1, :].contiguous().view(-1, vocab_size)
+ # logits = logits[:, 1:, :].contiguous().view(-1, vocab_size)
+ # labels = labels[:, 1:].contiguous().view(-1)
+
+ # mlm_loss = F.cross_entropy(
+ # logits,
+ # labels,
+ # # label_smoothing=0.1,
+ # )
+ # return mlm_loss
+
+
+ def forward(self, x:TensorType) -> TensorType:
+ attn_mask = (x != self.config.pad_token_id).long()
+ out = self.transformer(input_ids=x, attention_mask=attn_mask)
+ pooled_out = self.pooler(out, attn_mask)
+
+ return self.proj(pooled_out)
+
+ def lock(self, unlocked_layers:int=0, freeze_layer_norm:bool=True):
+ if not unlocked_layers: # full freezing
+ for n, p in self.transformer.named_parameters():
+ p.requires_grad = (not freeze_layer_norm) if "LayerNorm" in n.split(".") else False
+ return
+
+ encoder = self.transformer.encoder if hasattr(self.transformer, 'encoder') else self.transformer
+ layer_list = getattr(encoder, arch_dict[self.config.model_type]["config_names"]["layer_attr"])
+ print(f"Unlocking {unlocked_layers}/{len(layer_list) + 1} layers of hf model")
+ embeddings = getattr(
+ self.transformer, arch_dict[self.config.model_type]["config_names"]["token_embeddings_attr"])
+ modules = [embeddings, *layer_list][:-unlocked_layers]
+ # freeze layers
+ for module in modules:
+ for n, p in module.named_parameters():
+ p.requires_grad = (not freeze_layer_norm) if "LayerNorm" in n.split(".") else False
+
+
+ @torch.jit.ignore
+ def set_grad_checkpointing(self, enable=True):
+ self.transformer.gradient_checkpointing_enable()
+
+ def get_num_layers(self):
+ encoder = self.transformer.encoder if hasattr(self.transformer, 'encoder') else self.transformer
+ layer_list = getattr(encoder, arch_dict[self.config.model_type]["config_names"]["layer_attr"])
+ return len(layer_list)
+
+ def init_parameters(self):
+ pass
diff --git a/models/eva_clip/loss.py b/models/eva_clip/loss.py
new file mode 100644
index 0000000000000000000000000000000000000000..473f60d98d501067e85ace2dd089b00e249b6d17
--- /dev/null
+++ b/models/eva_clip/loss.py
@@ -0,0 +1,138 @@
+import math
+import torch
+import torch.nn as nn
+from torch.nn import functional as F
+
+try:
+ import torch.distributed.nn
+ from torch import distributed as dist
+ has_distributed = True
+except ImportError:
+ has_distributed = False
+
+try:
+ import horovod.torch as hvd
+except ImportError:
+ hvd = None
+
+from timm.loss import LabelSmoothingCrossEntropy
+
+
+def gather_features(
+ image_features,
+ text_features,
+ local_loss=False,
+ gather_with_grad=False,
+ rank=0,
+ world_size=1,
+ use_horovod=False
+):
+ assert has_distributed, 'torch.distributed did not import correctly, please use a PyTorch version with support.'
+ if use_horovod:
+ assert hvd is not None, 'Please install horovod'
+ if gather_with_grad:
+ all_image_features = hvd.allgather(image_features)
+ all_text_features = hvd.allgather(text_features)
+ else:
+ with torch.no_grad():
+ all_image_features = hvd.allgather(image_features)
+ all_text_features = hvd.allgather(text_features)
+ if not local_loss:
+ # ensure grads for local rank when all_* features don't have a gradient
+ gathered_image_features = list(all_image_features.chunk(world_size, dim=0))
+ gathered_text_features = list(all_text_features.chunk(world_size, dim=0))
+ gathered_image_features[rank] = image_features
+ gathered_text_features[rank] = text_features
+ all_image_features = torch.cat(gathered_image_features, dim=0)
+ all_text_features = torch.cat(gathered_text_features, dim=0)
+ else:
+ # We gather tensors from all gpus
+ if gather_with_grad:
+ all_image_features = torch.cat(torch.distributed.nn.all_gather(image_features), dim=0)
+ all_text_features = torch.cat(torch.distributed.nn.all_gather(text_features), dim=0)
+ # all_image_features = torch.cat(torch.distributed.nn.all_gather(image_features, async_op=True), dim=0)
+ # all_text_features = torch.cat(torch.distributed.nn.all_gather(text_features, async_op=True), dim=0)
+ else:
+ gathered_image_features = [torch.zeros_like(image_features) for _ in range(world_size)]
+ gathered_text_features = [torch.zeros_like(text_features) for _ in range(world_size)]
+ dist.all_gather(gathered_image_features, image_features)
+ dist.all_gather(gathered_text_features, text_features)
+ if not local_loss:
+ # ensure grads for local rank when all_* features don't have a gradient
+ gathered_image_features[rank] = image_features
+ gathered_text_features[rank] = text_features
+ all_image_features = torch.cat(gathered_image_features, dim=0)
+ all_text_features = torch.cat(gathered_text_features, dim=0)
+
+ return all_image_features, all_text_features
+
+
+class ClipLoss(nn.Module):
+
+ def __init__(
+ self,
+ local_loss=False,
+ gather_with_grad=False,
+ cache_labels=False,
+ rank=0,
+ world_size=1,
+ use_horovod=False,
+ smoothing=0.,
+ ):
+ super().__init__()
+ self.local_loss = local_loss
+ self.gather_with_grad = gather_with_grad
+ self.cache_labels = cache_labels
+ self.rank = rank
+ self.world_size = world_size
+ self.use_horovod = use_horovod
+ self.label_smoothing_cross_entropy = LabelSmoothingCrossEntropy(smoothing=smoothing) if smoothing > 0 else None
+
+ # cache state
+ self.prev_num_logits = 0
+ self.labels = {}
+
+ def forward(self, image_features, text_features, logit_scale=1.):
+ device = image_features.device
+ if self.world_size > 1:
+ all_image_features, all_text_features = gather_features(
+ image_features, text_features,
+ self.local_loss, self.gather_with_grad, self.rank, self.world_size, self.use_horovod)
+
+ if self.local_loss:
+ logits_per_image = logit_scale * image_features @ all_text_features.T
+ logits_per_text = logit_scale * text_features @ all_image_features.T
+ else:
+ logits_per_image = logit_scale * all_image_features @ all_text_features.T
+ logits_per_text = logits_per_image.T
+ else:
+ logits_per_image = logit_scale * image_features @ text_features.T
+ logits_per_text = logit_scale * text_features @ image_features.T
+ # calculated ground-truth and cache if enabled
+ num_logits = logits_per_image.shape[0]
+ if self.prev_num_logits != num_logits or device not in self.labels:
+ labels = torch.arange(num_logits, device=device, dtype=torch.long)
+ if self.world_size > 1 and self.local_loss:
+ labels = labels + num_logits * self.rank
+ if self.cache_labels:
+ self.labels[device] = labels
+ self.prev_num_logits = num_logits
+ else:
+ labels = self.labels[device]
+
+ if self.label_smoothing_cross_entropy:
+ total_loss = (
+ self.label_smoothing_cross_entropy(logits_per_image, labels) +
+ self.label_smoothing_cross_entropy(logits_per_text, labels)
+ ) / 2
+ else:
+ total_loss = (
+ F.cross_entropy(logits_per_image, labels) +
+ F.cross_entropy(logits_per_text, labels)
+ ) / 2
+
+ acc = None
+ i2t_acc = (logits_per_image.argmax(-1) == labels).sum() / len(logits_per_image)
+ t2i_acc = (logits_per_text.argmax(-1) == labels).sum() / len(logits_per_text)
+ acc = {"i2t": i2t_acc, "t2i": t2i_acc}
+ return total_loss, acc
\ No newline at end of file
diff --git a/models/eva_clip/model.py b/models/eva_clip/model.py
new file mode 100644
index 0000000000000000000000000000000000000000..da3bbd755799ced672385d1029ba7ce6d5215b0b
--- /dev/null
+++ b/models/eva_clip/model.py
@@ -0,0 +1,439 @@
+""" CLIP Model
+
+Adapted from https://github.com/openai/CLIP. Originally MIT License, Copyright (c) 2021 OpenAI.
+"""
+import os
+from dataclasses import dataclass
+from typing import Optional, Tuple, Union
+from functools import partial
+
+import numpy as np
+import torch
+import torch.nn.functional as F
+from torch import nn
+
+try:
+ from .hf_model import HFTextEncoder
+except:
+ HFTextEncoder = None
+from .modified_resnet import ModifiedResNet
+from .timm_model import TimmModel
+from .eva_vit_model import EVAVisionTransformer
+from .transformer import LayerNorm, QuickGELU, Attention, VisionTransformer, TextTransformer
+
+try:
+ from apex.normalization import FusedLayerNorm
+except:
+ FusedLayerNorm = LayerNorm
+ print("Please 'pip install apex'")
+
+try:
+ import xformers.ops as xops
+except ImportError:
+ xops = None
+ print("Please 'pip install xformers'")
+
+@dataclass
+class CLIPVisionCfg:
+ layers: Union[Tuple[int, int, int, int], int] = 12
+ width: int = 768
+ head_width: int = 64
+ mlp_ratio: float = 4.0
+ patch_size: int = 16
+ image_size: Union[Tuple[int, int], int] = 224
+ ls_init_value: Optional[float] = None # layer scale initial value
+ patch_dropout: float = 0. # what fraction of patches to dropout during training (0 would mean disabled and no patches dropped) - 0.5 to 0.75 recommended in the paper for optimal results
+ global_average_pool: bool = False # whether to global average pool the last embedding layer, instead of using CLS token (https://arxiv.org/abs/2205.01580)
+ drop_path_rate: Optional[float] = None # drop path rate
+ timm_model_name: str = None # a valid model name overrides layers, width, patch_size
+ timm_model_pretrained: bool = False # use (imagenet) pretrained weights for named model
+ timm_pool: str = 'avg' # feature pooling for timm model ('abs_attn', 'rot_attn', 'avg', '')
+ timm_proj: str = 'linear' # linear projection for timm model output ('linear', 'mlp', '')
+ timm_proj_bias: bool = False # enable bias final projection
+ eva_model_name: str = None # a valid eva model name overrides layers, width, patch_size
+ qkv_bias: bool = True
+ fusedLN: bool = False
+ xattn: bool = False
+ postnorm: bool = False
+ rope: bool = False
+ pt_hw_seq_len: int = 16 # 224/14
+ intp_freq: bool = False
+ naiveswiglu: bool = False
+ subln: bool = False
+
+
+@dataclass
+class CLIPTextCfg:
+ context_length: int = 77
+ vocab_size: int = 49408
+ width: int = 512
+ heads: int = 8
+ layers: int = 12
+ ls_init_value: Optional[float] = None # layer scale initial value
+ hf_model_name: str = None
+ hf_tokenizer_name: str = None
+ hf_model_pretrained: bool = True
+ proj: str = 'mlp'
+ pooler_type: str = 'mean_pooler'
+ masked_language_modeling: bool = False
+ fusedLN: bool = False
+ xattn: bool = False
+ attn_mask: bool = True
+
+def get_cast_dtype(precision: str):
+ cast_dtype = None
+ if precision == 'bf16':
+ cast_dtype = torch.bfloat16
+ elif precision == 'fp16':
+ cast_dtype = torch.float16
+ return cast_dtype
+
+
+def _build_vision_tower(
+ embed_dim: int,
+ vision_cfg: CLIPVisionCfg,
+ quick_gelu: bool = False,
+ cast_dtype: Optional[torch.dtype] = None
+):
+ if isinstance(vision_cfg, dict):
+ vision_cfg = CLIPVisionCfg(**vision_cfg)
+
+ # OpenAI models are pretrained w/ QuickGELU but native nn.GELU is both faster and more
+ # memory efficient in recent PyTorch releases (>= 1.10).
+ # NOTE: timm models always use native GELU regardless of quick_gelu flag.
+ act_layer = QuickGELU if quick_gelu else nn.GELU
+
+ if vision_cfg.eva_model_name:
+ vision_heads = vision_cfg.width // vision_cfg.head_width
+ norm_layer = LayerNorm
+
+ visual = EVAVisionTransformer(
+ img_size=vision_cfg.image_size,
+ patch_size=vision_cfg.patch_size,
+ num_classes=embed_dim,
+ use_mean_pooling=vision_cfg.global_average_pool, #False
+ init_values=vision_cfg.ls_init_value,
+ patch_dropout=vision_cfg.patch_dropout,
+ embed_dim=vision_cfg.width,
+ depth=vision_cfg.layers,
+ num_heads=vision_heads,
+ mlp_ratio=vision_cfg.mlp_ratio,
+ qkv_bias=vision_cfg.qkv_bias,
+ drop_path_rate=vision_cfg.drop_path_rate,
+ norm_layer= partial(FusedLayerNorm, eps=1e-6) if vision_cfg.fusedLN else partial(norm_layer, eps=1e-6),
+ xattn=vision_cfg.xattn,
+ rope=vision_cfg.rope,
+ postnorm=vision_cfg.postnorm,
+ pt_hw_seq_len= vision_cfg.pt_hw_seq_len, # 224/14
+ intp_freq= vision_cfg.intp_freq,
+ naiveswiglu= vision_cfg.naiveswiglu,
+ subln= vision_cfg.subln
+ )
+ elif vision_cfg.timm_model_name:
+ visual = TimmModel(
+ vision_cfg.timm_model_name,
+ pretrained=vision_cfg.timm_model_pretrained,
+ pool=vision_cfg.timm_pool,
+ proj=vision_cfg.timm_proj,
+ proj_bias=vision_cfg.timm_proj_bias,
+ embed_dim=embed_dim,
+ image_size=vision_cfg.image_size
+ )
+ act_layer = nn.GELU # so that text transformer doesn't use QuickGELU w/ timm models
+ elif isinstance(vision_cfg.layers, (tuple, list)):
+ vision_heads = vision_cfg.width * 32 // vision_cfg.head_width
+ visual = ModifiedResNet(
+ layers=vision_cfg.layers,
+ output_dim=embed_dim,
+ heads=vision_heads,
+ image_size=vision_cfg.image_size,
+ width=vision_cfg.width
+ )
+ else:
+ vision_heads = vision_cfg.width // vision_cfg.head_width
+ norm_layer = LayerNormFp32 if cast_dtype in (torch.float16, torch.bfloat16) else LayerNorm
+ visual = VisionTransformer(
+ image_size=vision_cfg.image_size,
+ patch_size=vision_cfg.patch_size,
+ width=vision_cfg.width,
+ layers=vision_cfg.layers,
+ heads=vision_heads,
+ mlp_ratio=vision_cfg.mlp_ratio,
+ ls_init_value=vision_cfg.ls_init_value,
+ patch_dropout=vision_cfg.patch_dropout,
+ global_average_pool=vision_cfg.global_average_pool,
+ output_dim=embed_dim,
+ act_layer=act_layer,
+ norm_layer=norm_layer,
+ )
+
+ return visual
+
+
+def _build_text_tower(
+ embed_dim: int,
+ text_cfg: CLIPTextCfg,
+ quick_gelu: bool = False,
+ cast_dtype: Optional[torch.dtype] = None,
+):
+ if isinstance(text_cfg, dict):
+ text_cfg = CLIPTextCfg(**text_cfg)
+
+ if text_cfg.hf_model_name:
+ text = HFTextEncoder(
+ text_cfg.hf_model_name,
+ output_dim=embed_dim,
+ tokenizer_name=text_cfg.hf_tokenizer_name,
+ proj=text_cfg.proj,
+ pooler_type=text_cfg.pooler_type,
+ masked_language_modeling=text_cfg.masked_language_modeling
+ )
+ else:
+ act_layer = QuickGELU if quick_gelu else nn.GELU
+ norm_layer = LayerNorm
+
+ text = TextTransformer(
+ context_length=text_cfg.context_length,
+ vocab_size=text_cfg.vocab_size,
+ width=text_cfg.width,
+ heads=text_cfg.heads,
+ layers=text_cfg.layers,
+ ls_init_value=text_cfg.ls_init_value,
+ output_dim=embed_dim,
+ act_layer=act_layer,
+ norm_layer= FusedLayerNorm if text_cfg.fusedLN else norm_layer,
+ xattn=text_cfg.xattn,
+ attn_mask=text_cfg.attn_mask,
+ )
+ return text
+
+class CLIP(nn.Module):
+ def __init__(
+ self,
+ embed_dim: int,
+ vision_cfg: CLIPVisionCfg,
+ text_cfg: CLIPTextCfg,
+ quick_gelu: bool = False,
+ cast_dtype: Optional[torch.dtype] = None,
+ ):
+ super().__init__()
+ self.visual = _build_vision_tower(embed_dim, vision_cfg, quick_gelu, cast_dtype)
+
+ text = _build_text_tower(embed_dim, text_cfg, quick_gelu, cast_dtype)
+ self.transformer = text.transformer
+ self.vocab_size = text.vocab_size
+ self.token_embedding = text.token_embedding
+ self.positional_embedding = text.positional_embedding
+ self.ln_final = text.ln_final
+ self.text_projection = text.text_projection
+ self.register_buffer('attn_mask', text.attn_mask, persistent=False)
+
+ self.logit_scale = nn.Parameter(torch.ones([]) * np.log(1 / 0.07))
+
+ def lock_image_tower(self, unlocked_groups=0, freeze_bn_stats=False):
+ # lock image tower as per LiT - https://arxiv.org/abs/2111.07991
+ self.visual.lock(unlocked_groups=unlocked_groups, freeze_bn_stats=freeze_bn_stats)
+
+ @torch.jit.ignore
+ def set_grad_checkpointing(self, enable=True):
+ self.visual.set_grad_checkpointing(enable)
+ self.transformer.grad_checkpointing = enable
+
+ @torch.jit.ignore
+ def no_weight_decay(self):
+ return {'logit_scale'}
+
+ def encode_image(self, image, normalize: bool = False):
+ features = self.visual(image)
+ return F.normalize(features, dim=-1) if normalize else features
+
+ def encode_text(self, text, normalize: bool = False):
+ cast_dtype = self.transformer.get_cast_dtype()
+
+ x = self.token_embedding(text).to(cast_dtype) # [batch_size, n_ctx, d_model]
+
+ x = x + self.positional_embedding.to(cast_dtype)
+ x = x.permute(1, 0, 2) # NLD -> LND
+ x = self.transformer(x, attn_mask=self.attn_mask)
+ x = x.permute(1, 0, 2) # LND -> NLD
+ x = self.ln_final(x) # [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 F.normalize(x, dim=-1) if normalize else x
+
+ def forward(self, image, text):
+ image_features = self.encode_image(image, normalize=True)
+ text_features = self.encode_text(text, normalize=True)
+ return image_features, text_features, self.logit_scale.exp()
+
+
+class CustomCLIP(nn.Module):
+ def __init__(
+ self,
+ embed_dim: int,
+ vision_cfg: CLIPVisionCfg,
+ text_cfg: CLIPTextCfg,
+ quick_gelu: bool = False,
+ cast_dtype: Optional[torch.dtype] = None,
+ itm_task: bool = False,
+ ):
+ super().__init__()
+ self.visual = _build_vision_tower(embed_dim, vision_cfg, quick_gelu, cast_dtype)
+ self.text = _build_text_tower(embed_dim, text_cfg, quick_gelu, cast_dtype)
+ self.logit_scale = nn.Parameter(torch.ones([]) * np.log(1 / 0.07))
+
+ def lock_image_tower(self, unlocked_groups=0, freeze_bn_stats=False):
+ # lock image tower as per LiT - https://arxiv.org/abs/2111.07991
+ self.visual.lock(unlocked_groups=unlocked_groups, freeze_bn_stats=freeze_bn_stats)
+
+ def lock_text_tower(self, unlocked_layers:int=0, freeze_layer_norm:bool=True):
+ self.text.lock(unlocked_layers, freeze_layer_norm)
+
+ @torch.jit.ignore
+ def set_grad_checkpointing(self, enable=True):
+ self.visual.set_grad_checkpointing(enable)
+ self.text.set_grad_checkpointing(enable)
+
+ @torch.jit.ignore
+ def no_weight_decay(self):
+ return {'logit_scale'}
+
+ def encode_image(self, image, normalize: bool = False):
+ features = self.visual(image)
+ return F.normalize(features, dim=-1) if normalize else features
+
+ def encode_text(self, text, normalize: bool = False):
+ features = self.text(text)
+ return F.normalize(features, dim=-1) if normalize else features
+
+ def forward(self, image, text):
+ image_features = self.encode_image(image, normalize=True)
+ text_features = self.encode_text(text, normalize=True)
+ return image_features, text_features, self.logit_scale.exp()
+
+
+def convert_weights_to_lp(model: nn.Module, dtype=torch.float16):
+ """Convert applicable model parameters to low-precision (bf16 or fp16)"""
+
+ def _convert_weights(l):
+
+ if isinstance(l, (nn.Conv1d, nn.Conv2d, nn.Linear)):
+ l.weight.data = l.weight.data.to(dtype)
+ if l.bias is not None:
+ l.bias.data = l.bias.data.to(dtype)
+
+ if isinstance(l, (nn.MultiheadAttention, Attention)):
+ 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, None)
+ if tensor is not None:
+ tensor.data = tensor.data.to(dtype)
+
+ if isinstance(l, nn.Parameter):
+ l.data = l.data.to(dtype)
+
+ for name in ["text_projection", "proj"]:
+ if hasattr(l, name) and isinstance(l, nn.Parameter):
+ attr = getattr(l, name, None)
+ if attr is not None:
+ attr.data = attr.data.to(dtype)
+
+ model.apply(_convert_weights)
+
+
+convert_weights_to_fp16 = convert_weights_to_lp # backwards compat
+
+
+# used to maintain checkpoint compatibility
+def convert_to_custom_text_state_dict(state_dict: dict):
+ if 'text_projection' in state_dict:
+ # old format state_dict, move text tower -> .text
+ new_state_dict = {}
+ for k, v in state_dict.items():
+ if any(k.startswith(p) for p in (
+ 'text_projection',
+ 'positional_embedding',
+ 'token_embedding',
+ 'transformer',
+ 'ln_final',
+ 'logit_scale'
+ )):
+ k = 'text.' + k
+ new_state_dict[k] = v
+ return new_state_dict
+ return state_dict
+
+
+def build_model_from_openai_state_dict(
+ state_dict: dict,
+ quick_gelu=True,
+ cast_dtype=torch.float16,
+):
+ 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_size = 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_size = 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")))
+
+ vision_cfg = CLIPVisionCfg(
+ layers=vision_layers,
+ width=vision_width,
+ patch_size=vision_patch_size,
+ image_size=image_size,
+ )
+ text_cfg = CLIPTextCfg(
+ context_length=context_length,
+ vocab_size=vocab_size,
+ width=transformer_width,
+ heads=transformer_heads,
+ layers=transformer_layers
+ )
+ model = CLIP(
+ embed_dim,
+ vision_cfg=vision_cfg,
+ text_cfg=text_cfg,
+ quick_gelu=quick_gelu, # OpenAI models were trained with QuickGELU
+ cast_dtype=cast_dtype,
+ )
+
+ for key in ["input_resolution", "context_length", "vocab_size"]:
+ state_dict.pop(key, None)
+
+ convert_weights_to_fp16(model) # OpenAI state dicts are partially converted to float16
+ model.load_state_dict(state_dict)
+ return model.eval()
+
+
+def trace_model(model, batch_size=256, device=torch.device('cpu')):
+ model.eval()
+ image_size = model.visual.image_size
+ example_images = torch.ones((batch_size, 3, image_size, image_size), device=device)
+ example_text = torch.zeros((batch_size, model.context_length), dtype=torch.int, device=device)
+ model = torch.jit.trace_module(
+ model,
+ inputs=dict(
+ forward=(example_images, example_text),
+ encode_text=(example_text,),
+ encode_image=(example_images,)
+ ))
+ model.visual.image_size = image_size
+ return model
diff --git a/models/eva_clip/model_configs/EVA01-CLIP-B-16.json b/models/eva_clip/model_configs/EVA01-CLIP-B-16.json
new file mode 100644
index 0000000000000000000000000000000000000000..aad2058003962a4ab286bf4e1ae956288af34e62
--- /dev/null
+++ b/models/eva_clip/model_configs/EVA01-CLIP-B-16.json
@@ -0,0 +1,19 @@
+{
+ "embed_dim": 512,
+ "vision_cfg": {
+ "image_size": 224,
+ "layers": 12,
+ "width": 768,
+ "patch_size": 16,
+ "eva_model_name": "eva-clip-b-16",
+ "ls_init_value": 0.1,
+ "drop_path_rate": 0.0
+ },
+ "text_cfg": {
+ "context_length": 77,
+ "vocab_size": 49408,
+ "width": 512,
+ "heads": 8,
+ "layers": 12
+ }
+}
\ No newline at end of file
diff --git a/models/eva_clip/model_configs/EVA01-CLIP-g-14-plus.json b/models/eva_clip/model_configs/EVA01-CLIP-g-14-plus.json
new file mode 100644
index 0000000000000000000000000000000000000000..100279572ff6d1bcca601f0eb526b4d4ff174c7d
--- /dev/null
+++ b/models/eva_clip/model_configs/EVA01-CLIP-g-14-plus.json
@@ -0,0 +1,24 @@
+{
+ "embed_dim": 1024,
+ "vision_cfg": {
+ "image_size": 224,
+ "layers": 40,
+ "width": 1408,
+ "head_width": 88,
+ "mlp_ratio": 4.3637,
+ "patch_size": 14,
+ "eva_model_name": "eva-clip-g-14-x",
+ "drop_path_rate": 0,
+ "xattn": true,
+ "fusedLN": true
+ },
+ "text_cfg": {
+ "context_length": 77,
+ "vocab_size": 49408,
+ "width": 1024,
+ "heads": 16,
+ "layers": 24,
+ "xattn": false,
+ "fusedLN": true
+ }
+}
\ No newline at end of file
diff --git a/models/eva_clip/model_configs/EVA01-CLIP-g-14.json b/models/eva_clip/model_configs/EVA01-CLIP-g-14.json
new file mode 100644
index 0000000000000000000000000000000000000000..5d338b4e6104241d1f0304ee82400035d5385332
--- /dev/null
+++ b/models/eva_clip/model_configs/EVA01-CLIP-g-14.json
@@ -0,0 +1,24 @@
+{
+ "embed_dim": 1024,
+ "vision_cfg": {
+ "image_size": 224,
+ "layers": 40,
+ "width": 1408,
+ "head_width": 88,
+ "mlp_ratio": 4.3637,
+ "patch_size": 14,
+ "eva_model_name": "eva-clip-g-14-x",
+ "drop_path_rate": 0.4,
+ "xattn": true,
+ "fusedLN": true
+ },
+ "text_cfg": {
+ "context_length": 77,
+ "vocab_size": 49408,
+ "width": 768,
+ "heads": 12,
+ "layers": 12,
+ "xattn": false,
+ "fusedLN": true
+ }
+}
\ No newline at end of file
diff --git a/models/eva_clip/model_configs/EVA02-CLIP-B-16.json b/models/eva_clip/model_configs/EVA02-CLIP-B-16.json
new file mode 100644
index 0000000000000000000000000000000000000000..e4a6e723f77033caa341ddf9b5be1787d64ad42c
--- /dev/null
+++ b/models/eva_clip/model_configs/EVA02-CLIP-B-16.json
@@ -0,0 +1,29 @@
+{
+ "embed_dim": 512,
+ "vision_cfg": {
+ "image_size": 224,
+ "layers": 12,
+ "width": 768,
+ "head_width": 64,
+ "patch_size": 16,
+ "mlp_ratio": 2.6667,
+ "eva_model_name": "eva-clip-b-16-X",
+ "drop_path_rate": 0.0,
+ "xattn": true,
+ "fusedLN": true,
+ "rope": true,
+ "pt_hw_seq_len": 16,
+ "intp_freq": true,
+ "naiveswiglu": true,
+ "subln": true
+ },
+ "text_cfg": {
+ "context_length": 77,
+ "vocab_size": 49408,
+ "width": 512,
+ "heads": 8,
+ "layers": 12,
+ "xattn": true,
+ "fusedLN": true
+ }
+}
\ No newline at end of file
diff --git a/models/eva_clip/model_configs/EVA02-CLIP-L-14-336.json b/models/eva_clip/model_configs/EVA02-CLIP-L-14-336.json
new file mode 100644
index 0000000000000000000000000000000000000000..3e1d124e1118911c5ad7b1ce85df195aca363ac4
--- /dev/null
+++ b/models/eva_clip/model_configs/EVA02-CLIP-L-14-336.json
@@ -0,0 +1,29 @@
+{
+ "embed_dim": 768,
+ "vision_cfg": {
+ "image_size": 336,
+ "layers": 24,
+ "width": 1024,
+ "drop_path_rate": 0,
+ "head_width": 64,
+ "mlp_ratio": 2.6667,
+ "patch_size": 14,
+ "eva_model_name": "eva-clip-l-14-336",
+ "xattn": true,
+ "fusedLN": true,
+ "rope": true,
+ "pt_hw_seq_len": 16,
+ "intp_freq": true,
+ "naiveswiglu": true,
+ "subln": true
+ },
+ "text_cfg": {
+ "context_length": 77,
+ "vocab_size": 49408,
+ "width": 768,
+ "heads": 12,
+ "layers": 12,
+ "xattn": false,
+ "fusedLN": true
+ }
+}
\ No newline at end of file
diff --git a/models/eva_clip/model_configs/EVA02-CLIP-L-14.json b/models/eva_clip/model_configs/EVA02-CLIP-L-14.json
new file mode 100644
index 0000000000000000000000000000000000000000..03b22ad3cfb92f9c843b9ec8d672e57e7a9ba4a2
--- /dev/null
+++ b/models/eva_clip/model_configs/EVA02-CLIP-L-14.json
@@ -0,0 +1,29 @@
+{
+ "embed_dim": 768,
+ "vision_cfg": {
+ "image_size": 224,
+ "layers": 24,
+ "width": 1024,
+ "drop_path_rate": 0,
+ "head_width": 64,
+ "mlp_ratio": 2.6667,
+ "patch_size": 14,
+ "eva_model_name": "eva-clip-l-14",
+ "xattn": true,
+ "fusedLN": true,
+ "rope": true,
+ "pt_hw_seq_len": 16,
+ "intp_freq": true,
+ "naiveswiglu": true,
+ "subln": true
+ },
+ "text_cfg": {
+ "context_length": 77,
+ "vocab_size": 49408,
+ "width": 768,
+ "heads": 12,
+ "layers": 12,
+ "xattn": false,
+ "fusedLN": true
+ }
+}
\ No newline at end of file
diff --git a/models/eva_clip/model_configs/EVA02-CLIP-bigE-14-plus.json b/models/eva_clip/model_configs/EVA02-CLIP-bigE-14-plus.json
new file mode 100644
index 0000000000000000000000000000000000000000..aa04e2545ac1e015daae2c10133956ce969524f7
--- /dev/null
+++ b/models/eva_clip/model_configs/EVA02-CLIP-bigE-14-plus.json
@@ -0,0 +1,25 @@
+{
+ "embed_dim": 1024,
+ "vision_cfg": {
+ "image_size": 224,
+ "layers": 64,
+ "width": 1792,
+ "head_width": 112,
+ "mlp_ratio": 8.571428571428571,
+ "patch_size": 14,
+ "eva_model_name": "eva-clip-4b-14-x",
+ "drop_path_rate": 0,
+ "xattn": true,
+ "postnorm": true,
+ "fusedLN": true
+ },
+ "text_cfg": {
+ "context_length": 77,
+ "vocab_size": 49408,
+ "width": 1280,
+ "heads": 20,
+ "layers": 32,
+ "xattn": false,
+ "fusedLN": true
+ }
+}
diff --git a/models/eva_clip/model_configs/EVA02-CLIP-bigE-14.json b/models/eva_clip/model_configs/EVA02-CLIP-bigE-14.json
new file mode 100644
index 0000000000000000000000000000000000000000..747ffccc8bd49dbb6701b58e15843b7fe3754e64
--- /dev/null
+++ b/models/eva_clip/model_configs/EVA02-CLIP-bigE-14.json
@@ -0,0 +1,25 @@
+{
+ "embed_dim": 1024,
+ "vision_cfg": {
+ "image_size": 224,
+ "layers": 64,
+ "width": 1792,
+ "head_width": 112,
+ "mlp_ratio": 8.571428571428571,
+ "patch_size": 14,
+ "eva_model_name": "eva-clip-4b-14-x",
+ "drop_path_rate": 0,
+ "xattn": true,
+ "postnorm": true,
+ "fusedLN": true
+ },
+ "text_cfg": {
+ "context_length": 77,
+ "vocab_size": 49408,
+ "width": 1024,
+ "heads": 16,
+ "layers": 24,
+ "xattn": false,
+ "fusedLN": true
+ }
+}
\ No newline at end of file
diff --git a/models/eva_clip/modified_resnet.py b/models/eva_clip/modified_resnet.py
new file mode 100644
index 0000000000000000000000000000000000000000..299080850061a0dab433322e5f8fe2a55fb4e9a2
--- /dev/null
+++ b/models/eva_clip/modified_resnet.py
@@ -0,0 +1,188 @@
+import os
+import sys
+
+import torch
+from torch import nn
+from torch.nn import functional as F
+from collections import OrderedDict
+
+current_file_path = os.path.abspath(__file__)
+project_roots = [os.path.dirname(current_file_path)]
+for project_root in project_roots:
+ sys.path.insert(0, project_root) if project_root not in sys.path else None
+
+from utils import freeze_batch_norm_2d
+
+
+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.act1 = nn.ReLU(inplace=True)
+
+ self.conv2 = nn.Conv2d(planes, planes, 3, padding=1, bias=False)
+ self.bn2 = nn.BatchNorm2d(planes)
+ self.act2 = nn.ReLU(inplace=True)
+
+ 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.act3 = 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.act1(self.bn1(self.conv1(x)))
+ out = self.act2(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.act3(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, image_size=224, width=64):
+ super().__init__()
+ self.output_dim = output_dim
+ self.image_size = image_size
+
+ # 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.act1 = nn.ReLU(inplace=True)
+ self.conv2 = nn.Conv2d(width // 2, width // 2, kernel_size=3, padding=1, bias=False)
+ self.bn2 = nn.BatchNorm2d(width // 2)
+ self.act2 = nn.ReLU(inplace=True)
+ self.conv3 = nn.Conv2d(width // 2, width, kernel_size=3, padding=1, bias=False)
+ self.bn3 = nn.BatchNorm2d(width)
+ self.act3 = nn.ReLU(inplace=True)
+ self.avgpool = nn.AvgPool2d(2)
+
+ # 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(image_size // 32, embed_dim, heads, output_dim)
+
+ self.init_parameters()
+
+ 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 init_parameters(self):
+ if self.attnpool is not None:
+ std = self.attnpool.c_proj.in_features ** -0.5
+ nn.init.normal_(self.attnpool.q_proj.weight, std=std)
+ nn.init.normal_(self.attnpool.k_proj.weight, std=std)
+ nn.init.normal_(self.attnpool.v_proj.weight, std=std)
+ nn.init.normal_(self.attnpool.c_proj.weight, std=std)
+
+ for resnet_block in [self.layer1, self.layer2, self.layer3, self.layer4]:
+ for name, param in resnet_block.named_parameters():
+ if name.endswith("bn3.weight"):
+ nn.init.zeros_(param)
+
+ def lock(self, unlocked_groups=0, freeze_bn_stats=False):
+ assert unlocked_groups == 0, 'partial locking not currently supported for this model'
+ for param in self.parameters():
+ param.requires_grad = False
+ if freeze_bn_stats:
+ freeze_batch_norm_2d(self)
+
+ @torch.jit.ignore
+ def set_grad_checkpointing(self, enable=True):
+ # FIXME support for non-transformer
+ pass
+
+ def stem(self, x):
+ x = self.act1(self.bn1(self.conv1(x)))
+ x = self.act2(self.bn2(self.conv2(x)))
+ x = self.act3(self.bn3(self.conv3(x)))
+ x = self.avgpool(x)
+ return x
+
+ def forward(self, x):
+ x = self.stem(x)
+ x = self.layer1(x)
+ x = self.layer2(x)
+ x = self.layer3(x)
+ x = self.layer4(x)
+ x = self.attnpool(x)
+
+ return x
diff --git a/models/eva_clip/openai.py b/models/eva_clip/openai.py
new file mode 100644
index 0000000000000000000000000000000000000000..cc4e13e876d6a7a3463b457e62c517cb063b1356
--- /dev/null
+++ b/models/eva_clip/openai.py
@@ -0,0 +1,144 @@
+""" OpenAI pretrained model functions
+
+Adapted from https://github.com/openai/CLIP. Originally MIT License, Copyright (c) 2021 OpenAI.
+"""
+
+import os
+import warnings
+from typing import List, Optional, Union
+
+import torch
+
+from .model import build_model_from_openai_state_dict, convert_weights_to_lp, get_cast_dtype
+from .pretrained import get_pretrained_url, list_pretrained_models_by_tag, download_pretrained_from_url
+
+__all__ = ["list_openai_models", "load_openai_model"]
+
+
+def list_openai_models() -> List[str]:
+ """Returns the names of available CLIP models"""
+ return list_pretrained_models_by_tag('openai')
+
+
+def load_openai_model(
+ name: str,
+ precision: Optional[str] = None,
+ device: Optional[Union[str, torch.device]] = None,
+ jit: bool = True,
+ cache_dir: Optional[str] = None,
+):
+ """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
+ precision: str
+ Model precision, if None defaults to 'fp32' if device == 'cpu' else 'fp16'.
+ device : Union[str, torch.device]
+ The device to put the loaded model
+ jit : bool
+ Whether to load the optimized JIT model (default) or more hackable non-JIT model.
+ cache_dir : Optional[str]
+ The directory to cache the downloaded model weights
+
+ 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 device is None:
+ device = "cuda" if torch.cuda.is_available() else "cpu"
+ if precision is None:
+ precision = 'fp32' if device == 'cpu' else 'fp16'
+
+ if get_pretrained_url(name, 'openai'):
+ model_path = download_pretrained_from_url(get_pretrained_url(name, 'openai'), cache_dir=cache_dir)
+ elif os.path.isfile(name):
+ model_path = name
+ else:
+ raise RuntimeError(f"Model {name} not found; available models = {list_openai_models()}")
+
+ 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:
+ # Build a non-jit model from the OpenAI jitted model state dict
+ cast_dtype = get_cast_dtype(precision)
+ try:
+ model = build_model_from_openai_state_dict(state_dict or model.state_dict(), cast_dtype=cast_dtype)
+ except KeyError:
+ sd = {k[7:]: v for k, v in state_dict["state_dict"].items()}
+ model = build_model_from_openai_state_dict(sd, cast_dtype=cast_dtype)
+
+ # model from OpenAI state dict is in manually cast fp16 mode, must be converted for AMP/fp32/bf16 use
+ model = model.to(device)
+ if precision.startswith('amp') or precision == 'fp32':
+ model.float()
+ elif precision == 'bf16':
+ convert_weights_to_lp(model, dtype=torch.bfloat16)
+
+ return model
+
+ # 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 (typically for CPU)
+ if precision == 'fp32':
+ 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()
+
+ # ensure image_size attr available at consistent location for both jit and non-jit
+ model.visual.image_size = model.input_resolution.item()
+ return model
diff --git a/models/eva_clip/pretrained.py b/models/eva_clip/pretrained.py
new file mode 100644
index 0000000000000000000000000000000000000000..a1e55dcf36a0e7dbd4c13b4ca2d7cb460e4c3547
--- /dev/null
+++ b/models/eva_clip/pretrained.py
@@ -0,0 +1,332 @@
+import hashlib
+import os
+import urllib
+import warnings
+from functools import partial
+from typing import Dict, Union
+
+from tqdm import tqdm
+
+try:
+ from huggingface_hub import hf_hub_download
+ _has_hf_hub = True
+except ImportError:
+ hf_hub_download = None
+ _has_hf_hub = False
+
+
+def _pcfg(url='', hf_hub='', filename='', mean=None, std=None):
+ return dict(
+ url=url,
+ hf_hub=hf_hub,
+ mean=mean,
+ std=std,
+ )
+
+_VITB32 = dict(
+ openai=_pcfg(
+ "https://openaipublic.azureedge.net/clip/models/40d365715913c9da98579312b702a82c18be219cc2a73407c4526f58eba950af/ViT-B-32.pt"),
+ laion400m_e31=_pcfg(
+ "https://github.com/mlfoundations/open_clip/releases/download/v0.2-weights/vit_b_32-quickgelu-laion400m_e31-d867053b.pt"),
+ laion400m_e32=_pcfg(
+ "https://github.com/mlfoundations/open_clip/releases/download/v0.2-weights/vit_b_32-quickgelu-laion400m_e32-46683a32.pt"),
+ laion2b_e16=_pcfg(
+ "https://github.com/mlfoundations/open_clip/releases/download/v0.2-weights/vit_b_32-laion2b_e16-af8dbd0c.pth"),
+ laion2b_s34b_b79k=_pcfg(hf_hub='laion/CLIP-ViT-B-32-laion2B-s34B-b79K/')
+)
+
+_VITB32_quickgelu = dict(
+ openai=_pcfg(
+ "https://openaipublic.azureedge.net/clip/models/40d365715913c9da98579312b702a82c18be219cc2a73407c4526f58eba950af/ViT-B-32.pt"),
+ laion400m_e31=_pcfg(
+ "https://github.com/mlfoundations/open_clip/releases/download/v0.2-weights/vit_b_32-quickgelu-laion400m_e31-d867053b.pt"),
+ laion400m_e32=_pcfg(
+ "https://github.com/mlfoundations/open_clip/releases/download/v0.2-weights/vit_b_32-quickgelu-laion400m_e32-46683a32.pt"),
+)
+
+_VITB16 = dict(
+ openai=_pcfg(
+ "https://openaipublic.azureedge.net/clip/models/5806e77cd80f8b59890b7e101eabd078d9fb84e6937f9e85e4ecb61988df416f/ViT-B-16.pt"),
+ laion400m_e31=_pcfg(
+ "https://github.com/mlfoundations/open_clip/releases/download/v0.2-weights/vit_b_16-laion400m_e31-00efa78f.pt"),
+ laion400m_e32=_pcfg(
+ "https://github.com/mlfoundations/open_clip/releases/download/v0.2-weights/vit_b_16-laion400m_e32-55e67d44.pt"),
+ laion2b_s34b_b88k=_pcfg(hf_hub='laion/CLIP-ViT-B-16-laion2B-s34B-b88K/'),
+)
+
+_EVAB16 = dict(
+ eva=_pcfg(hf_hub='QuanSun/EVA-CLIP/EVA02_B_psz14to16.pt'),
+ eva02=_pcfg(hf_hub='QuanSun/EVA-CLIP/EVA02_B_psz14to16.pt'),
+ eva_clip=_pcfg(hf_hub='QuanSun/EVA-CLIP/EVA02_CLIP_B_psz16_s8B.pt'),
+ eva02_clip=_pcfg(hf_hub='QuanSun/EVA-CLIP/EVA02_CLIP_B_psz16_s8B.pt'),
+)
+
+_VITB16_PLUS_240 = dict(
+ laion400m_e31=_pcfg(
+ "https://github.com/mlfoundations/open_clip/releases/download/v0.2-weights/vit_b_16_plus_240-laion400m_e31-8fb26589.pt"),
+ laion400m_e32=_pcfg(
+ "https://github.com/mlfoundations/open_clip/releases/download/v0.2-weights/vit_b_16_plus_240-laion400m_e32-699c4b84.pt"),
+)
+
+_VITL14 = dict(
+ openai=_pcfg(
+ "https://openaipublic.azureedge.net/clip/models/b8cca3fd41ae0c99ba7e8951adf17d267cdb84cd88be6f7c2e0eca1737a03836/ViT-L-14.pt"),
+ laion400m_e31=_pcfg(
+ "https://github.com/mlfoundations/open_clip/releases/download/v0.2-weights/vit_l_14-laion400m_e31-69988bb6.pt"),
+ laion400m_e32=_pcfg(
+ "https://github.com/mlfoundations/open_clip/releases/download/v0.2-weights/vit_l_14-laion400m_e32-3d133497.pt"),
+ laion2b_s32b_b82k=_pcfg(
+ hf_hub='laion/CLIP-ViT-L-14-laion2B-s32B-b82K/',
+ mean=(0.5, 0.5, 0.5), std=(0.5, 0.5, 0.5)),
+)
+
+_EVAL14 = dict(
+ eva=_pcfg(hf_hub='QuanSun/EVA-CLIP/EVA02_L_psz14.pt'),
+ eva02=_pcfg(hf_hub='QuanSun/EVA-CLIP/EVA02_L_psz14.pt'),
+ eva_clip=_pcfg(hf_hub='QuanSun/EVA-CLIP/EVA02_CLIP_L_psz14_s4B.pt'),
+ eva02_clip=_pcfg(hf_hub='QuanSun/EVA-CLIP/EVA02_CLIP_L_psz14_s4B.pt'),
+)
+
+_VITL14_336 = dict(
+ openai=_pcfg(
+ "https://openaipublic.azureedge.net/clip/models/3035c92b350959924f9f00213499208652fc7ea050643e8b385c2dac08641f02/ViT-L-14-336px.pt"),
+)
+
+_EVAL14_336 = dict(
+ eva_clip=_pcfg(hf_hub='QuanSun/EVA-CLIP/EVA02_CLIP_L_336_psz14_s6B.pt'),
+ eva02_clip=_pcfg(hf_hub='QuanSun/EVA-CLIP/EVA02_CLIP_L_336_psz14_s6B.pt'),
+ eva_clip_224to336=_pcfg(hf_hub='QuanSun/EVA-CLIP/EVA02_CLIP_L_psz14_224to336.pt'),
+ eva02_clip_224to336=_pcfg(hf_hub='QuanSun/EVA-CLIP/EVA02_CLIP_L_psz14_224to336.pt'),
+)
+
+_VITH14 = dict(
+ laion2b_s32b_b79k=_pcfg(hf_hub='laion/CLIP-ViT-H-14-laion2B-s32B-b79K/'),
+)
+
+_VITg14 = dict(
+ laion2b_s12b_b42k=_pcfg(hf_hub='laion/CLIP-ViT-g-14-laion2B-s12B-b42K/'),
+ laion2b_s34b_b88k=_pcfg(hf_hub='laion/CLIP-ViT-g-14-laion2B-s34B-b88K/'),
+)
+
+_EVAg14 = dict(
+ eva=_pcfg(hf_hub='QuanSun/EVA-CLIP/'),
+ eva01=_pcfg(hf_hub='QuanSun/EVA-CLIP/EVA01_g_psz14.pt'),
+ eva_clip=_pcfg(hf_hub='QuanSun/EVA-CLIP/EVA01_CLIP_g_14_psz14_s11B.pt'),
+ eva01_clip=_pcfg(hf_hub='QuanSun/EVA-CLIP/EVA01_CLIP_g_14_psz14_s11B.pt'),
+)
+
+_EVAg14_PLUS = dict(
+ eva=_pcfg(hf_hub='QuanSun/EVA-CLIP/'),
+ eva01=_pcfg(hf_hub='QuanSun/EVA-CLIP/EVA01_g_psz14.pt'),
+ eva_clip=_pcfg(hf_hub='QuanSun/EVA-CLIP/EVA01_CLIP_g_14_plus_psz14_s11B.pt'),
+ eva01_clip=_pcfg(hf_hub='QuanSun/EVA-CLIP/EVA01_CLIP_g_14_plus_psz14_s11B.pt'),
+)
+
+_VITbigG14 = dict(
+ laion2b_s39b_b160k=_pcfg(hf_hub='laion/CLIP-ViT-bigG-14-laion2B-39B-b160k/'),
+)
+
+_EVAbigE14 = dict(
+ eva=_pcfg(hf_hub='QuanSun/EVA-CLIP/EVA02_E_psz14.pt'),
+ eva02=_pcfg(hf_hub='QuanSun/EVA-CLIP/EVA02_E_psz14.pt'),
+ eva_clip=_pcfg(hf_hub='QuanSun/EVA-CLIP/EVA02_CLIP_E_psz14_s4B.pt'),
+ eva02_clip=_pcfg(hf_hub='QuanSun/EVA-CLIP/EVA02_CLIP_E_psz14_s4B.pt'),
+)
+
+_EVAbigE14_PLUS = dict(
+ eva=_pcfg(hf_hub='QuanSun/EVA-CLIP/EVA02_E_psz14.pt'),
+ eva02=_pcfg(hf_hub='QuanSun/EVA-CLIP/EVA02_E_psz14.pt'),
+ eva_clip=_pcfg(hf_hub='QuanSun/EVA-CLIP/EVA02_CLIP_E_psz14_plus_s9B.pt'),
+ eva02_clip=_pcfg(hf_hub='QuanSun/EVA-CLIP/EVA02_CLIP_E_psz14_plus_s9B.pt'),
+)
+
+
+_PRETRAINED = {
+ # "ViT-B-32": _VITB32,
+ "OpenaiCLIP-B-32": _VITB32,
+ "OpenCLIP-B-32": _VITB32,
+
+ # "ViT-B-32-quickgelu": _VITB32_quickgelu,
+ "OpenaiCLIP-B-32-quickgelu": _VITB32_quickgelu,
+ "OpenCLIP-B-32-quickgelu": _VITB32_quickgelu,
+
+ # "ViT-B-16": _VITB16,
+ "OpenaiCLIP-B-16": _VITB16,
+ "OpenCLIP-B-16": _VITB16,
+
+ "EVA02-B-16": _EVAB16,
+ "EVA02-CLIP-B-16": _EVAB16,
+
+ # "ViT-B-16-plus-240": _VITB16_PLUS_240,
+ "OpenCLIP-B-16-plus-240": _VITB16_PLUS_240,
+
+ # "ViT-L-14": _VITL14,
+ "OpenaiCLIP-L-14": _VITL14,
+ "OpenCLIP-L-14": _VITL14,
+
+ "EVA02-L-14": _EVAL14,
+ "EVA02-CLIP-L-14": _EVAL14,
+
+ # "ViT-L-14-336": _VITL14_336,
+ "OpenaiCLIP-L-14-336": _VITL14_336,
+
+ "EVA02-CLIP-L-14-336": _EVAL14_336,
+
+ # "ViT-H-14": _VITH14,
+ # "ViT-g-14": _VITg14,
+ "OpenCLIP-H-14": _VITH14,
+ "OpenCLIP-g-14": _VITg14,
+
+ "EVA01-CLIP-g-14": _EVAg14,
+ "EVA01-CLIP-g-14-plus": _EVAg14_PLUS,
+
+ # "ViT-bigG-14": _VITbigG14,
+ "OpenCLIP-bigG-14": _VITbigG14,
+
+ "EVA02-CLIP-bigE-14": _EVAbigE14,
+ "EVA02-CLIP-bigE-14-plus": _EVAbigE14_PLUS,
+}
+
+
+def _clean_tag(tag: str):
+ # normalize pretrained tags
+ return tag.lower().replace('-', '_')
+
+
+def list_pretrained(as_str: bool = False):
+ """ returns list of pretrained models
+ Returns a tuple (model_name, pretrain_tag) by default or 'name:tag' if as_str == True
+ """
+ return [':'.join([k, t]) if as_str else (k, t) for k in _PRETRAINED.keys() for t in _PRETRAINED[k].keys()]
+
+
+def list_pretrained_models_by_tag(tag: str):
+ """ return all models having the specified pretrain tag """
+ models = []
+ tag = _clean_tag(tag)
+ for k in _PRETRAINED.keys():
+ if tag in _PRETRAINED[k]:
+ models.append(k)
+ return models
+
+
+def list_pretrained_tags_by_model(model: str):
+ """ return all pretrain tags for the specified model architecture """
+ tags = []
+ if model in _PRETRAINED:
+ tags.extend(_PRETRAINED[model].keys())
+ return tags
+
+
+def is_pretrained_cfg(model: str, tag: str):
+ if model not in _PRETRAINED:
+ return False
+ return _clean_tag(tag) in _PRETRAINED[model]
+
+
+def get_pretrained_cfg(model: str, tag: str):
+ if model not in _PRETRAINED:
+ return {}
+ model_pretrained = _PRETRAINED[model]
+ return model_pretrained.get(_clean_tag(tag), {})
+
+
+def get_pretrained_url(model: str, tag: str):
+ cfg = get_pretrained_cfg(model, _clean_tag(tag))
+ return cfg.get('url', '')
+
+
+def download_pretrained_from_url(
+ url: str,
+ cache_dir: Union[str, None] = None,
+):
+ if not cache_dir:
+ cache_dir = os.path.expanduser("~/.cache/clip")
+ os.makedirs(cache_dir, exist_ok=True)
+ filename = os.path.basename(url)
+
+ if 'openaipublic' in url:
+ expected_sha256 = url.split("/")[-2]
+ elif 'mlfoundations' in url:
+ expected_sha256 = os.path.splitext(filename)[0].split("-")[-1]
+ else:
+ expected_sha256 = ''
+
+ download_target = os.path.join(cache_dir, 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 expected_sha256:
+ if hashlib.sha256(open(download_target, "rb").read()).hexdigest().startswith(expected_sha256):
+ return download_target
+ else:
+ warnings.warn(f"{download_target} exists, but the SHA256 checksum does not match; re-downloading the file")
+ else:
+ return download_target
+
+ with urllib.request.urlopen(url) as source, open(download_target, "wb") as output:
+ with tqdm(total=int(source.headers.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 expected_sha256 and not hashlib.sha256(open(download_target, "rb").read()).hexdigest().startswith(expected_sha256):
+ raise RuntimeError(f"Model has been downloaded but the SHA256 checksum does not not match")
+
+ return download_target
+
+
+def has_hf_hub(necessary=False):
+ if not _has_hf_hub and necessary:
+ # if no HF Hub module installed, and it is necessary to continue, raise error
+ raise RuntimeError(
+ 'Hugging Face hub model specified but package not installed. Run `pip install huggingface_hub`.')
+ return _has_hf_hub
+
+
+def download_pretrained_from_hf(
+ model_id: str,
+ filename: str = 'open_clip_pytorch_model.bin',
+ revision=None,
+ cache_dir: Union[str, None] = None,
+):
+ has_hf_hub(True)
+ cached_file = hf_hub_download(model_id, filename, revision=revision, cache_dir=cache_dir)
+ return cached_file
+
+
+def download_pretrained(
+ cfg: Dict,
+ force_hf_hub: bool = False,
+ cache_dir: Union[str, None] = None,
+):
+ target = ''
+ if not cfg:
+ return target
+
+ download_url = cfg.get('url', '')
+ download_hf_hub = cfg.get('hf_hub', '')
+ if download_hf_hub and force_hf_hub:
+ # use HF hub even if url exists
+ download_url = ''
+
+ if download_url:
+ target = download_pretrained_from_url(download_url, cache_dir=cache_dir)
+ elif download_hf_hub:
+ has_hf_hub(True)
+ # we assume the hf_hub entries in pretrained config combine model_id + filename in
+ # 'org/model_name/filename.pt' form. To specify just the model id w/o filename and
+ # use 'open_clip_pytorch_model.bin' default, there must be a trailing slash 'org/model_name/'.
+ model_id, filename = os.path.split(download_hf_hub)
+ if filename:
+ target = download_pretrained_from_hf(model_id, filename=filename, cache_dir=cache_dir)
+ else:
+ target = download_pretrained_from_hf(model_id, cache_dir=cache_dir)
+
+ return target
diff --git a/models/eva_clip/rope.py b/models/eva_clip/rope.py
new file mode 100644
index 0000000000000000000000000000000000000000..69030c35ea7b6b4f298daebbee5717f3fa1254ab
--- /dev/null
+++ b/models/eva_clip/rope.py
@@ -0,0 +1,137 @@
+from math import pi
+import torch
+from torch import nn
+from einops import rearrange, repeat
+import logging
+
+def broadcat(tensors, dim = -1):
+ num_tensors = len(tensors)
+ shape_lens = set(list(map(lambda t: len(t.shape), tensors)))
+ assert len(shape_lens) == 1, 'tensors must all have the same number of dimensions'
+ shape_len = list(shape_lens)[0]
+ dim = (dim + shape_len) if dim < 0 else dim
+ dims = list(zip(*map(lambda t: list(t.shape), tensors)))
+ expandable_dims = [(i, val) for i, val in enumerate(dims) if i != dim]
+ assert all([*map(lambda t: len(set(t[1])) <= 2, expandable_dims)]), 'invalid dimensions for broadcastable concatentation'
+ max_dims = list(map(lambda t: (t[0], max(t[1])), expandable_dims))
+ expanded_dims = list(map(lambda t: (t[0], (t[1],) * num_tensors), max_dims))
+ expanded_dims.insert(dim, (dim, dims[dim]))
+ expandable_shapes = list(zip(*map(lambda t: t[1], expanded_dims)))
+ tensors = list(map(lambda t: t[0].expand(*t[1]), zip(tensors, expandable_shapes)))
+ return torch.cat(tensors, dim = dim)
+
+def rotate_half(x):
+ x = rearrange(x, '... (d r) -> ... d r', r = 2)
+ x1, x2 = x.unbind(dim = -1)
+ x = torch.stack((-x2, x1), dim = -1)
+ return rearrange(x, '... d r -> ... (d r)')
+
+
+class VisionRotaryEmbedding(nn.Module):
+ def __init__(
+ self,
+ dim,
+ pt_seq_len,
+ ft_seq_len=None,
+ custom_freqs = None,
+ freqs_for = 'lang',
+ theta = 10000,
+ max_freq = 10,
+ num_freqs = 1,
+ ):
+ super().__init__()
+ if custom_freqs:
+ freqs = custom_freqs
+ elif freqs_for == 'lang':
+ freqs = 1. / (theta ** (torch.arange(0, dim, 2)[:(dim // 2)].float() / dim))
+ elif freqs_for == 'pixel':
+ freqs = torch.linspace(1., max_freq / 2, dim // 2) * pi
+ elif freqs_for == 'constant':
+ freqs = torch.ones(num_freqs).float()
+ else:
+ raise ValueError(f'unknown modality {freqs_for}')
+
+ if ft_seq_len is None: ft_seq_len = pt_seq_len
+ t = torch.arange(ft_seq_len) / ft_seq_len * pt_seq_len
+
+ freqs_h = torch.einsum('..., f -> ... f', t, freqs)
+ freqs_h = repeat(freqs_h, '... n -> ... (n r)', r = 2)
+
+ freqs_w = torch.einsum('..., f -> ... f', t, freqs)
+ freqs_w = repeat(freqs_w, '... n -> ... (n r)', r = 2)
+
+ freqs = broadcat((freqs_h[:, None, :], freqs_w[None, :, :]), dim = -1)
+
+ self.register_buffer("freqs_cos", freqs.cos())
+ self.register_buffer("freqs_sin", freqs.sin())
+
+ logging.info(f'Shape of rope freq: {self.freqs_cos.shape}')
+
+ def forward(self, t, start_index = 0):
+ rot_dim = self.freqs_cos.shape[-1]
+ end_index = start_index + rot_dim
+ assert rot_dim <= t.shape[-1], f'feature dimension {t.shape[-1]} is not of sufficient size to rotate in all the positions {rot_dim}'
+ t_left, t, t_right = t[..., :start_index], t[..., start_index:end_index], t[..., end_index:]
+ t = (t * self.freqs_cos) + (rotate_half(t) * self.freqs_sin)
+
+ return torch.cat((t_left, t, t_right), dim = -1)
+
+class VisionRotaryEmbeddingFast(nn.Module):
+ def __init__(
+ self,
+ dim,
+ pt_seq_len,
+ ft_seq_len=None,
+ custom_freqs = None,
+ freqs_for = 'lang',
+ theta = 10000,
+ max_freq = 10,
+ num_freqs = 1,
+ patch_dropout = 0.
+ ):
+ super().__init__()
+ if custom_freqs:
+ freqs = custom_freqs
+ elif freqs_for == 'lang':
+ freqs = 1. / (theta ** (torch.arange(0, dim, 2)[:(dim // 2)].float() / dim))
+ elif freqs_for == 'pixel':
+ freqs = torch.linspace(1., max_freq / 2, dim // 2) * pi
+ elif freqs_for == 'constant':
+ freqs = torch.ones(num_freqs).float()
+ else:
+ raise ValueError(f'unknown modality {freqs_for}')
+
+ if ft_seq_len is None: ft_seq_len = pt_seq_len
+ t = torch.arange(ft_seq_len) / ft_seq_len * pt_seq_len
+
+ freqs = torch.einsum('..., f -> ... f', t, freqs)
+ freqs = repeat(freqs, '... n -> ... (n r)', r = 2)
+ freqs = broadcat((freqs[:, None, :], freqs[None, :, :]), dim = -1)
+
+ freqs_cos = freqs.cos().view(-1, freqs.shape[-1])
+ freqs_sin = freqs.sin().view(-1, freqs.shape[-1])
+
+ self.patch_dropout = patch_dropout
+
+ self.register_buffer("freqs_cos", freqs_cos)
+ self.register_buffer("freqs_sin", freqs_sin)
+
+ logging.info(f'Shape of rope freq: {self.freqs_cos.shape}')
+
+ def forward(self, t, patch_indices_keep=None):
+ if patch_indices_keep is not None:
+ batch = t.size()[0]
+ batch_indices = torch.arange(batch)
+ batch_indices = batch_indices[..., None]
+
+ freqs_cos = repeat(self.freqs_cos, 'i j -> n i m j', n=t.shape[0], m=t.shape[1])
+ freqs_sin = repeat(self.freqs_sin, 'i j -> n i m j', n=t.shape[0], m=t.shape[1])
+
+ freqs_cos = freqs_cos[batch_indices, patch_indices_keep]
+ freqs_cos = rearrange(freqs_cos, 'n i m j -> n m i j')
+ freqs_sin = freqs_sin[batch_indices, patch_indices_keep]
+ freqs_sin = rearrange(freqs_sin, 'n i m j -> n m i j')
+
+ return t * freqs_cos + rotate_half(t) * freqs_sin
+
+ return t * self.freqs_cos + rotate_half(t) * self.freqs_sin
\ No newline at end of file
diff --git a/models/eva_clip/timm_model.py b/models/eva_clip/timm_model.py
new file mode 100644
index 0000000000000000000000000000000000000000..b58122c0b84fbda9e51867342823222234e17505
--- /dev/null
+++ b/models/eva_clip/timm_model.py
@@ -0,0 +1,122 @@
+""" timm model adapter
+
+Wraps timm (https://github.com/rwightman/pytorch-image-models) models for use as a vision tower in CLIP model.
+"""
+import logging
+from collections import OrderedDict
+
+import torch
+import torch.nn as nn
+
+try:
+ import timm
+ from timm.models.layers import Mlp, to_2tuple
+ try:
+ # old timm imports < 0.8.1
+ from timm.models.layers.attention_pool2d import RotAttentionPool2d
+ from timm.models.layers.attention_pool2d import AttentionPool2d as AbsAttentionPool2d
+ except ImportError:
+ # new timm imports >= 0.8.1
+ from timm.layers import RotAttentionPool2d
+ from timm.layers import AttentionPool2d as AbsAttentionPool2d
+except ImportError:
+ timm = None
+
+from .utils import freeze_batch_norm_2d
+
+
+class TimmModel(nn.Module):
+ """ timm model adapter
+ # FIXME this adapter is a work in progress, may change in ways that break weight compat
+ """
+
+ def __init__(
+ self,
+ model_name,
+ embed_dim,
+ image_size=224,
+ pool='avg',
+ proj='linear',
+ proj_bias=False,
+ drop=0.,
+ pretrained=False):
+ super().__init__()
+ if timm is None:
+ raise RuntimeError("Please `pip install timm` to use timm models.")
+
+ self.image_size = to_2tuple(image_size)
+ self.trunk = timm.create_model(model_name, pretrained=pretrained)
+ feat_size = self.trunk.default_cfg.get('pool_size', None)
+ feature_ndim = 1 if not feat_size else 2
+ if pool in ('abs_attn', 'rot_attn'):
+ assert feature_ndim == 2
+ # if attn pooling used, remove both classifier and default pool
+ self.trunk.reset_classifier(0, global_pool='')
+ else:
+ # reset global pool if pool config set, otherwise leave as network default
+ reset_kwargs = dict(global_pool=pool) if pool else {}
+ self.trunk.reset_classifier(0, **reset_kwargs)
+ prev_chs = self.trunk.num_features
+
+ head_layers = OrderedDict()
+ if pool == 'abs_attn':
+ head_layers['pool'] = AbsAttentionPool2d(prev_chs, feat_size=feat_size, out_features=embed_dim)
+ prev_chs = embed_dim
+ elif pool == 'rot_attn':
+ head_layers['pool'] = RotAttentionPool2d(prev_chs, out_features=embed_dim)
+ prev_chs = embed_dim
+ else:
+ assert proj, 'projection layer needed if non-attention pooling is used.'
+
+ # NOTE attention pool ends with a projection layer, so proj should usually be set to '' if such pooling is used
+ if proj == 'linear':
+ head_layers['drop'] = nn.Dropout(drop)
+ head_layers['proj'] = nn.Linear(prev_chs, embed_dim, bias=proj_bias)
+ elif proj == 'mlp':
+ head_layers['mlp'] = Mlp(prev_chs, 2 * embed_dim, embed_dim, drop=drop, bias=(True, proj_bias))
+
+ self.head = nn.Sequential(head_layers)
+
+ def lock(self, unlocked_groups=0, freeze_bn_stats=False):
+ """ lock modules
+ Args:
+ unlocked_groups (int): leave last n layer groups unlocked (default: 0)
+ """
+ if not unlocked_groups:
+ # lock full model
+ for param in self.trunk.parameters():
+ param.requires_grad = False
+ if freeze_bn_stats:
+ freeze_batch_norm_2d(self.trunk)
+ else:
+ # NOTE: partial freeze requires latest timm (master) branch and is subject to change
+ try:
+ # FIXME import here until API stable and in an official release
+ from timm.models.helpers import group_parameters, group_modules
+ except ImportError:
+ raise RuntimeError(
+ 'Please install latest timm `pip install git+https://github.com/rwightman/pytorch-image-models`')
+ matcher = self.trunk.group_matcher()
+ gparams = group_parameters(self.trunk, matcher)
+ max_layer_id = max(gparams.keys())
+ max_layer_id = max_layer_id - unlocked_groups
+ for group_idx in range(max_layer_id + 1):
+ group = gparams[group_idx]
+ for param in group:
+ self.trunk.get_parameter(param).requires_grad = False
+ if freeze_bn_stats:
+ gmodules = group_modules(self.trunk, matcher, reverse=True)
+ gmodules = {k for k, v in gmodules.items() if v <= max_layer_id}
+ freeze_batch_norm_2d(self.trunk, gmodules)
+
+ @torch.jit.ignore
+ def set_grad_checkpointing(self, enable=True):
+ try:
+ self.trunk.set_grad_checkpointing(enable)
+ except Exception as e:
+ logging.warning('grad checkpointing not supported for this timm image tower, continuing without...')
+
+ def forward(self, x):
+ x = self.trunk(x)
+ x = self.head(x)
+ return x
diff --git a/models/eva_clip/tokenizer.py b/models/eva_clip/tokenizer.py
new file mode 100644
index 0000000000000000000000000000000000000000..41482f82aebbf197f4ee4e6c07c845a0d69dd7d6
--- /dev/null
+++ b/models/eva_clip/tokenizer.py
@@ -0,0 +1,201 @@
+""" CLIP tokenizer
+
+Copied from https://github.com/openai/CLIP. Originally MIT License, Copyright (c) 2021 OpenAI.
+"""
+import gzip
+import html
+import os
+from functools import lru_cache
+from typing import Union, List
+
+import ftfy
+import regex as re
+import torch
+
+# https://stackoverflow.com/q/62691279
+import os
+os.environ["TOKENIZERS_PARALLELISM"] = "false"
+
+
+@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(), special_tokens=None):
+ 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+'' for v in vocab]
+ for merge in merges:
+ vocab.append(''.join(merge))
+ if not special_tokens:
+ special_tokens = ['
+