| import torch |
| import random |
| import logging |
| import open_clip |
| import torch.nn as nn |
| import src.clip.clip as CLIP |
|
|
|
|
| def load_model(architecture, **kargs): |
| |
| |
| params = architecture.split('|') |
|
|
| clip_arch = params[0] |
| if (len(params) == 1): |
| model, transform = CLIP.load( |
| clip_arch, |
| "cpu", |
| **kargs |
| ) |
|
|
| elif (len(params) == 2): |
| open_pretrain = params[1] |
|
|
| _model, _, _ = open_clip.create_model_and_transforms( |
| clip_arch.replace("/", "-"), |
| pretrained=open_pretrain, |
| device="cpu" |
| ) |
|
|
| model, transform = CLIP.load( |
| _model.state_dict(), |
| "cpu", |
| **kargs |
| ) |
| del _model |
|
|
| elif (len(params) > 2): |
| raise NotImplementedError() |
|
|
| return model, transform |
|
|
|
|
| class VideoAttrExtractor(nn.Module): |
| def __init__( |
| self, |
| architecture, |
| text_embed, |
| store_attrs=[], |
| attn_record=False, |
| pretrain=None |
| ): |
| super().__init__() |
| self.model, self.transform = load_model( |
| architecture, |
| store_attrs=store_attrs, |
| attn_record=attn_record |
| ) |
| self.model = self.model.visual.float() |
|
|
| if (pretrain): |
| logging.info("Loading image encoder pretrain weights...") |
| state_dict = torch.load(pretrain, "cpu") |
| try: |
| self.model.load_state_dict(state_dict, strict=True) |
| except: |
| conflicts = self.model.load_state_dict(state_dict, strict=False) |
| logging.warning( |
| f"during visual pretrain weights loading, disabling strict mode with conflicts:\n{conflicts}" |
| ) |
|
|
| self.model.requires_grad_(False) |
|
|
| if not text_embed: |
| self.model.proj = None |
| self.feat_dim = self.model.transformer.width |
| else: |
| self.feat_dim = self.model.output_dim |
|
|
| @property |
| def n_px(self): |
| return self.model.input_resolution |
|
|
| @property |
| def n_layers(self): |
| return self.model.transformer.layers |
|
|
| @property |
| def n_heads(self): |
| return self.model.transformer.heads |
|
|
| @property |
| def patch_size(self): |
| return self.model.patch_size |
|
|
| @property |
| def patch_num(self): |
| return self.model.patch_num |
|
|
| @property |
| def n_patch(self): |
| return int(self.n_px // self.patch_size) |
|
|
| @property |
| def embed_dim(self): |
| return self.feat_dim |
|
|
| def forward(self, x): |
| b, t = x.shape[:2] |
|
|
| |
| embeds = self.model(x) |
| |
| layer_attrs = [] |
| for blk in self.model.transformer.resblocks: |
| attrs = blk.pop_attr() |
| layer_attrs.append(attrs) |
| return dict( |
| layer_attrs=layer_attrs, |
| embeds=embeds |
| ) |
|
|
| def train(self, mode=True): |
| self.model.eval() |
| return self |
|
|
|
|
| if __name__ == "__main__": |
| VideoAttrExtractor( |
| "ViT-L/14|laion2b_s32b_b82k", |
| text_embed=False |
| ) |
|
|
