durgappc/infinitetalk2
0
1# Modified from ``https://github.com/openai/CLIP'' and ``https://github.com/mlfoundations/open_clip''2# Copyright 2024-2025 The Alibaba Wan Team Authors. All rights reserved.3import logging4import math5 6import torch7import torch.nn as nn8import torch.nn.functional as F9import torchvision.transforms as T10 11from .attention import flash_attention12from .tokenizers import HuggingfaceTokenizer13from .xlm_roberta import XLMRoberta14 15__all__ = [16 'XLMRobertaCLIP',17 'clip_xlm_roberta_vit_h_14',18 'CLIPModel',19]20 21 22def pos_interpolate(pos, seq_len):23 if pos.size(1) == seq_len:24 return pos25 else:26 src_grid = int(math.sqrt(pos.size(1)))27 tar_grid = int(math.sqrt(seq_len))28 n = pos.size(1) - src_grid * src_grid29 return torch.cat([30 pos[:, :n],31 F.interpolate(32 pos[:, n:].float().reshape(1, src_grid, src_grid, -1).permute(33 0, 3, 1, 2),34 size=(tar_grid, tar_grid),35 mode='bicubic',36 align_corners=False).flatten(2).transpose(1, 2)37 ],38 dim=1)39 40 41class QuickGELU(nn.Module):42 43 def forward(self, x):44 return x * torch.sigmoid(1.702 * x)45 46 47class LayerNorm(nn.LayerNorm):48 49 def forward(self, x):50 return super().forward(x.float()).type_as(x)51 52 53class SelfAttention(nn.Module):54 55 def __init__(self,56 dim,57 num_heads,58 causal=False,59 attn_dropout=0.0,60 proj_dropout=0.0):61 assert dim % num_heads == 062 super().__init__()63 self.dim = dim64 self.num_heads = num_heads65 self.head_dim = dim // num_heads66 self.causal = causal67 self.attn_dropout = attn_dropout68 self.proj_dropout = proj_dropout69 70 # layers71 self.to_qkv = nn.Linear(dim, dim * 3)72 self.proj = nn.Linear(dim, dim)73 74 def forward(self, x):75 """76 x: [B, L, C].77 """78 b, s, c, n, d = *x.size(), self.num_heads, self.head_dim79 80 # compute query, key, value81 q, k, v = self.to_qkv(x).view(b, s, 3, n, d).unbind(2)82 83 # compute attention84 p = self.attn_dropout if self.training else 0.085 x = flash_attention(q, k, v, dropout_p=p, causal=self.causal, version=2)86 x = x.reshape(b, s, c)87 88 # output89 x = self.proj(x)90 x = F.dropout(x, self.proj_dropout, self.training)91 return x92 93 94class SwiGLU(nn.Module):95 96 def __init__(self, dim, mid_dim):97 super().__init__()98 self.dim = dim99 self.mid_dim = mid_dim100 101 # layers102 self.fc1 = nn.Linear(dim, mid_dim)103 self.fc2 = nn.Linear(dim, mid_dim)104 self.fc3 = nn.Linear(mid_dim, dim)105 106 def forward(self, x):107 x = F.silu(self.fc1(x)) * self.fc2(x)108 x = self.fc3(x)109 return x110 111 112class AttentionBlock(nn.Module):113 114 def __init__(self,115 dim,116 mlp_ratio,117 num_heads,118 post_norm=False,119 causal=False,120 activation='quick_gelu',121 attn_dropout=0.0,122 proj_dropout=0.0,123 norm_eps=1e-5):124 assert activation in ['quick_gelu', 'gelu', 'swi_glu']125 super().__init__()126 self.dim = dim127 self.mlp_ratio = mlp_ratio128 self.num_heads = num_heads129 self.post_norm = post_norm130 self.causal = causal131 self.norm_eps = norm_eps132 133 # layers134 self.norm1 = LayerNorm(dim, eps=norm_eps)135 self.attn = SelfAttention(dim, num_heads, causal, attn_dropout,136 proj_dropout)137 self.norm2 = LayerNorm(dim, eps=norm_eps)138 if activation == 'swi_glu':139 self.mlp = SwiGLU(dim, int(dim * mlp_ratio))140 else:141 self.mlp = nn.Sequential(142 nn.Linear(dim, int(dim * mlp_ratio)),143 QuickGELU() if activation == 'quick_gelu' else nn.GELU(),144 nn.Linear(int(dim * mlp_ratio), dim), nn.Dropout(proj_dropout))145 146 def forward(self, x):147 if self.post_norm:148 x = x + self.norm1(self.attn(x))149 x = x + self.norm2(self.mlp(x))150 else:151 x = x + self.attn(self.norm1(x))152 x = x + self.mlp(self.norm2(x))153 return x154 155 156class AttentionPool(nn.Module):157 158 def __init__(self,159 dim,160 mlp_ratio,161 num_heads,162 activation='gelu',163 proj_dropout=0.0,164 norm_eps=1e-5):165 assert dim % num_heads == 0166 super().__init__()167 self.dim = dim168 self.mlp_ratio = mlp_ratio169 self.num_heads = num_heads170 self.head_dim = dim // num_heads171 self.proj_dropout = proj_dropout172 self.norm_eps = norm_eps173 174 # layers175 gain = 1.0 / math.sqrt(dim)176 self.cls_embedding = nn.Parameter(gain * torch.randn(1, 1, dim))177 self.to_q = nn.Linear(dim, dim)178 self.to_kv = nn.Linear(dim, dim * 2)179 self.proj = nn.Linear(dim, dim)180 self.norm = LayerNorm(dim, eps=norm_eps)181 self.mlp = nn.Sequential(182 nn.Linear(dim, int(dim * mlp_ratio)),183 QuickGELU() if activation == 'quick_gelu' else nn.GELU(),184 nn.Linear(int(dim * mlp_ratio), dim), nn.Dropout(proj_dropout))185 186 def forward(self, x):187 """188 x: [B, L, C].189 """190 b, s, c, n, d = *x.size(), self.num_heads, self.head_dim191 192 # compute query, key, value193 q = self.to_q(self.cls_embedding).view(1, 1, n, d).expand(b, -1, -1, -1)194 k, v = self.to_kv(x).view(b, s, 2, n, d).unbind(2)195 196 # compute attention197 x = flash_attention(q, k, v, version=2)198 x = x.reshape(b, 1, c)199 200 # output201 x = self.proj(x)202 x = F.dropout(x, self.proj_dropout, self.training)203 204 # mlp205 x = x + self.mlp(self.norm(x))206 return x[:, 0]207 208 209class VisionTransformer(nn.Module):210 211 def __init__(self,212 image_size=224,213 patch_size=16,214 dim=768,215 mlp_ratio=4,216 out_dim=512,217 num_heads=12,218 num_layers=12,219 pool_type='token',220 pre_norm=True,221 post_norm=False,222 activation='quick_gelu',223 attn_dropout=0.0,224 proj_dropout=0.0,225 embedding_dropout=0.0,226 norm_eps=1e-5):227 if image_size % patch_size != 0:228 print(229 '[WARNING] image_size is not divisible by patch_size',230 flush=True)231 assert pool_type in ('token', 'token_fc', 'attn_pool')232 out_dim = out_dim or dim233 super().__init__()234 self.image_size = image_size235 self.patch_size = patch_size236 self.num_patches = (image_size // patch_size)**2237 self.dim = dim238 self.mlp_ratio = mlp_ratio239 self.out_dim = out_dim240 self.num_heads = num_heads241 self.num_layers = num_layers242 self.pool_type = pool_type243 self.post_norm = post_norm244 self.norm_eps = norm_eps245 246 # embeddings247 gain = 1.0 / math.sqrt(dim)248 self.patch_embedding = nn.Conv2d(249 3,250 dim,251 kernel_size=patch_size,252 stride=patch_size,253 bias=not pre_norm)254 if pool_type in ('token', 'token_fc'):255 self.cls_embedding = nn.Parameter(gain * torch.randn(1, 1, dim))256 self.pos_embedding = nn.Parameter(gain * torch.randn(257 1, self.num_patches +258 (1 if pool_type in ('token', 'token_fc') else 0), dim))259 self.dropout = nn.Dropout(embedding_dropout)260 261 # transformer262 self.pre_norm = LayerNorm(dim, eps=norm_eps) if pre_norm else None263 self.transformer = nn.Sequential(*[264 AttentionBlock(dim, mlp_ratio, num_heads, post_norm, False,265 activation, attn_dropout, proj_dropout, norm_eps)266 for _ in range(num_layers)267 ])268 self.post_norm = LayerNorm(dim, eps=norm_eps)269 270 # head271 if pool_type == 'token':272 self.head = nn.Parameter(gain * torch.randn(dim, out_dim))273 elif pool_type == 'token_fc':274 self.head = nn.Linear(dim, out_dim)275 elif pool_type == 'attn_pool':276 self.head = AttentionPool(dim, mlp_ratio, num_heads, activation,277 proj_dropout, norm_eps)278 279 def forward(self, x, interpolation=False, use_31_block=False):280 b = x.size(0)281 282 # embeddings283 x = self.patch_embedding(x).flatten(2).permute(0, 2, 1)284 if self.pool_type in ('token', 'token_fc'):285 x = torch.cat([self.cls_embedding.expand(b, -1, -1), x], dim=1)286 if interpolation:287 e = pos_interpolate(self.pos_embedding, x.size(1))288 else:289 e = self.pos_embedding290 x = self.dropout(x + e)291 if self.pre_norm is not None:292 x = self.pre_norm(x)293 294 # transformer295 if use_31_block:296 x = self.transformer[:-1](x)297 return x298 else:299 x = self.transformer(x)300 return x301 302 303class XLMRobertaWithHead(XLMRoberta):304 305 def __init__(self, **kwargs):306 self.out_dim = kwargs.pop('out_dim')307 super().__init__(**kwargs)308 309 # head310 mid_dim = (self.dim + self.out_dim) // 2311 self.head = nn.Sequential(312 nn.Linear(self.dim, mid_dim, bias=False), nn.GELU(),313 nn.Linear(mid_dim, self.out_dim, bias=False))314 315 def forward(self, ids):316 # xlm-roberta317 x = super().forward(ids)318 319 # average pooling320 mask = ids.ne(self.pad_id).unsqueeze(-1).to(x)321 x = (x * mask).sum(dim=1) / mask.sum(dim=1)322 323 # head324 x = self.head(x)325 return x326 327 328class XLMRobertaCLIP(nn.Module):329 330 def __init__(self,331 embed_dim=1024,332 image_size=224,333 patch_size=14,334 vision_dim=1280,335 vision_mlp_ratio=4,336 vision_heads=16,337 vision_layers=32,338 vision_pool='token',339 vision_pre_norm=True,340 vision_post_norm=False,341 activation='gelu',342 vocab_size=250002,343 max_text_len=514,344 type_size=1,345 pad_id=1,346 text_dim=1024,347 text_heads=16,348 text_layers=24,349 text_post_norm=True,350 text_dropout=0.1,351 attn_dropout=0.0,352 proj_dropout=0.0,353 embedding_dropout=0.0,354 norm_eps=1e-5):355 super().__init__()356 self.embed_dim = embed_dim357 self.image_size = image_size358 self.patch_size = patch_size359 self.vision_dim = vision_dim360 self.vision_mlp_ratio = vision_mlp_ratio361 self.vision_heads = vision_heads362 self.vision_layers = vision_layers363 self.vision_pre_norm = vision_pre_norm364 self.vision_post_norm = vision_post_norm365 self.activation = activation366 self.vocab_size = vocab_size367 self.max_text_len = max_text_len368 self.type_size = type_size369 self.pad_id = pad_id370 self.text_dim = text_dim371 self.text_heads = text_heads372 self.text_layers = text_layers373 self.text_post_norm = text_post_norm374 self.norm_eps = norm_eps375 376 # models377 self.visual = VisionTransformer(378 image_size=image_size,379 patch_size=patch_size,380 dim=vision_dim,381 mlp_ratio=vision_mlp_ratio,382 out_dim=embed_dim,383 num_heads=vision_heads,384 num_layers=vision_layers,385 pool_type=vision_pool,386 pre_norm=vision_pre_norm,387 post_norm=vision_post_norm,388 activation=activation,389 attn_dropout=attn_dropout,390 proj_dropout=proj_dropout,391 embedding_dropout=embedding_dropout,392 norm_eps=norm_eps)393 self.textual = XLMRobertaWithHead(394 vocab_size=vocab_size,395 max_seq_len=max_text_len,396 type_size=type_size,397 pad_id=pad_id,398 dim=text_dim,399 out_dim=embed_dim,400 num_heads=text_heads,401 num_layers=text_layers,402 post_norm=text_post_norm,403 dropout=text_dropout)404 self.log_scale = nn.Parameter(math.log(1 / 0.07) * torch.ones([]))405 406 def forward(self, imgs, txt_ids):407 """408 imgs: [B, 3, H, W] of torch.float32.409 - mean: [0.48145466, 0.4578275, 0.40821073]410 - std: [0.26862954, 0.26130258, 0.27577711]411 txt_ids: [B, L] of torch.long.412 Encoded by data.CLIPTokenizer.413 """414 xi = self.visual(imgs)415 xt = self.textual(txt_ids)416 return xi, xt417 418 def param_groups(self):419 groups = [{420 'params': [421 p for n, p in self.named_parameters()422 if 'norm' in n or n.endswith('bias')423 ],424 'weight_decay': 0.0425 }, {426 'params': [427 p for n, p in self.named_parameters()428 if not ('norm' in n or n.endswith('bias'))429 ]430 }]431 return groups432 433 434def _clip(pretrained=False,435 pretrained_name=None,436 model_cls=XLMRobertaCLIP,437 return_transforms=False,438 return_tokenizer=False,439 tokenizer_padding='eos',440 dtype=torch.float32,441 device='cpu',442 **kwargs):443 # init a model on device444 with torch.device(device):445 model = model_cls(**kwargs)446 447 # set device448 model = model.to(dtype=dtype, device=device)449 output = (model,)450 451 # init transforms452 if return_transforms:453 # mean and std454 if 'siglip' in pretrained_name.lower():455 mean, std = [0.5, 0.5, 0.5], [0.5, 0.5, 0.5]456 else:457 mean = [0.48145466, 0.4578275, 0.40821073]458 std = [0.26862954, 0.26130258, 0.27577711]459 460 # transforms461 transforms = T.Compose([462 T.Resize((model.image_size, model.image_size),463 interpolation=T.InterpolationMode.BICUBIC),464 T.ToTensor(),465 T.Normalize(mean=mean, std=std)466 ])467 output += (transforms,)468 return output[0] if len(output) == 1 else output469 470 471def clip_xlm_roberta_vit_h_14(472 pretrained=False,473 pretrained_name='open-clip-xlm-roberta-large-vit-huge-14',474 **kwargs):475 cfg = dict(476 embed_dim=1024,477 image_size=224,478 patch_size=14,479 vision_dim=1280,480 vision_mlp_ratio=4,481 vision_heads=16,482 vision_layers=32,483 vision_pool='token',484 activation='gelu',485 vocab_size=250002,486 max_text_len=514,487 type_size=1,488 pad_id=1,489 text_dim=1024,490 text_heads=16,491 text_layers=24,492 text_post_norm=True,493 text_dropout=0.1,494 attn_dropout=0.0,495 proj_dropout=0.0,496 embedding_dropout=0.0)497 cfg.update(**kwargs)498 return _clip(pretrained, pretrained_name, XLMRobertaCLIP, **cfg)499 500 501class CLIPModel:502 503 def __init__(self, dtype, device, checkpoint_path, tokenizer_path):504 self.dtype = dtype505 self.device = device506 self.checkpoint_path = checkpoint_path507 self.tokenizer_path = tokenizer_path508 509 # init model510 self.model, self.transforms = clip_xlm_roberta_vit_h_14(511 pretrained=False,512 return_transforms=True,513 return_tokenizer=False,514 dtype=dtype,515 device=device)516 self.model = self.model.eval().requires_grad_(False)517 logging.info(f'loading {checkpoint_path}')518 self.model.load_state_dict(519 torch.load(checkpoint_path, map_location='cpu'))520 521 # init tokenizer522 self.tokenizer = HuggingfaceTokenizer(523 name=tokenizer_path,524 seq_len=self.model.max_text_len - 2,525 clean='whitespace')526 527 def visual(self, videos):528 # preprocess529 size = (self.model.image_size,) * 2530 videos = torch.cat([531 F.interpolate(532 u.transpose(0, 1),533 size=size,534 mode='bicubic',535 align_corners=False) for u in videos536 ])537 videos = self.transforms.transforms[-1](videos.mul_(0.5).add_(0.5))538 539 # forward540 with torch.cuda.amp.autocast(dtype=self.dtype):541 out = self.model.visual(videos, use_31_block=True)542 return out543