| """ Conversion functions for 3rd part state-dicts and non-torch native checkpoint formats. |
| """ |
| from typing import Union |
|
|
| import torch |
| import numpy as np |
|
|
| from .model import CLIP, CustomTextCLIP |
| from .transformer import TextTransformer, Transformer |
|
|
|
|
| @torch.no_grad() |
| def load_big_vision_weights(model: CustomTextCLIP, checkpoint_path: str): |
| """ Load weights from .npz checkpoints for official Google big_vision image-text models |
| |
| Currently the SigLIP source models are supported and a CustomTextCLIP destination model |
| w/ timm image encoder. |
| """ |
| from timm.layers import resample_patch_embed, resample_abs_pos_embed |
|
|
| def _n2p(w, t=True): |
| if w.ndim == 4 and w.shape[0] == w.shape[1] == w.shape[2] == 1: |
| w = w.flatten() |
| if t: |
| if w.ndim == 4: |
| w = w.transpose([3, 2, 0, 1]) |
| elif w.ndim == 3: |
| w = w.transpose([2, 0, 1]) |
| elif w.ndim == 2: |
| w = w.transpose([1, 0]) |
| return torch.from_numpy(w) |
|
|
| w = np.load(checkpoint_path) |
| interpolation = 'bilinear' |
| antialias = False |
|
|
| def _convert_timm_img(module, prefix): |
| embed_conv_w = _n2p(w[f'{prefix}embedding/kernel']) |
| if embed_conv_w.shape[-2:] != module.patch_embed.proj.weight.shape[-2:]: |
| embed_conv_w = resample_patch_embed( |
| embed_conv_w, |
| module.patch_embed.proj.weight.shape[-2:], |
| interpolation=interpolation, |
| antialias=antialias, |
| verbose=True, |
| ) |
| module.patch_embed.proj.weight.copy_(embed_conv_w) |
| module.patch_embed.proj.bias.copy_(_n2p(w[f'{prefix}embedding/bias'])) |
|
|
| if module.cls_token is not None: |
| module.cls_token.copy_(_n2p(w[f'{prefix}cls'], t=False)) |
|
|
| pos_embed_w = _n2p(w[f'{prefix}pos_embedding'], t=False) |
| if pos_embed_w.shape != module.pos_embed.shape: |
| assert False, f'{pos_embed_w.shape}, {module.pos_embed.shape}' |
| num_prefix_tokens = 0 if getattr(module, 'no_embed_class', False) else getattr(module, 'num_prefix_tokens', 1) |
| pos_embed_w = resample_abs_pos_embed( |
| pos_embed_w, |
| new_size=module.patch_embed.grid_size, |
| num_prefix_tokens=num_prefix_tokens, |
| interpolation=interpolation, |
| antialias=antialias, |
| verbose=True, |
| ) |
| module.pos_embed.copy_(pos_embed_w) |
|
|
| mha_sub, b_sub, ln1_sub = (0, 0, 1) |
| for i, block in enumerate(module.blocks.children()): |
| block_prefix = f'{prefix}Transformer/encoderblock_{i}/' |
| mha_prefix = block_prefix + f'MultiHeadDotProductAttention_{mha_sub}/' |
| block.norm1.weight.copy_(_n2p(w[f'{block_prefix}LayerNorm_0/scale'])) |
| block.norm1.bias.copy_(_n2p(w[f'{block_prefix}LayerNorm_0/bias'])) |
| block.attn.qkv.weight.copy_(torch.cat([ |
| _n2p(w[f'{mha_prefix}{n}/kernel'], t=False).flatten(1).T for n in ('query', 'key', 'value')])) |
| block.attn.qkv.bias.copy_(torch.cat([ |
| _n2p(w[f'{mha_prefix}{n}/bias'], t=False).reshape(-1) for n in ('query', 'key', 'value')])) |
| block.attn.proj.weight.copy_(_n2p(w[f'{mha_prefix}out/kernel']).flatten(1)) |
| block.attn.proj.bias.copy_(_n2p(w[f'{mha_prefix}out/bias'])) |
| for r in range(2): |
| getattr(block.mlp, f'fc{r + 1}').weight.copy_(_n2p(w[f'{block_prefix}MlpBlock_{b_sub}/Dense_{r}/kernel'])) |
| getattr(block.mlp, f'fc{r + 1}').bias.copy_(_n2p(w[f'{block_prefix}MlpBlock_{b_sub}/Dense_{r}/bias'])) |
| block.norm2.weight.copy_(_n2p(w[f'{block_prefix}LayerNorm_{ln1_sub}/scale'])) |
| block.norm2.bias.copy_(_n2p(w[f'{block_prefix}LayerNorm_{ln1_sub}/bias'])) |
|
|
| module.norm.weight.copy_(_n2p(w[f'{prefix}Transformer/encoder_norm/scale'])) |
| module.norm.bias.copy_(_n2p(w[f'{prefix}Transformer/encoder_norm/bias'])) |
|
|
| if module.attn_pool is not None: |
| block_prefix = f'{prefix}MAPHead_0/' |
| mha_prefix = block_prefix + f'MultiHeadDotProductAttention_0/' |
| module.attn_pool.latent.copy_(_n2p(w[f'{block_prefix}probe'], t=False)) |
| module.attn_pool.q.weight.copy_(_n2p(w[f'{mha_prefix}query/kernel'], t=False).flatten(1).T) |
| module.attn_pool.q.bias.copy_(_n2p(w[f'{mha_prefix}query/bias'], t=False).reshape(-1)) |
| module.attn_pool.kv.weight.copy_(torch.cat([ |
| _n2p(w[f'{mha_prefix}{n}/kernel'], t=False).flatten(1).T for n in ('key', 'value')])) |
| module.attn_pool.kv.bias.copy_(torch.cat([ |
| _n2p(w[f'{mha_prefix}{n}/bias'], t=False).reshape(-1) for n in ('key', 'value')])) |
| module.attn_pool.proj.weight.copy_(_n2p(w[f'{mha_prefix}out/kernel']).flatten(1)) |
| module.attn_pool.proj.bias.copy_(_n2p(w[f'{mha_prefix}out/bias'])) |
| module.attn_pool.norm.weight.copy_(_n2p(w[f'{block_prefix}LayerNorm_0/scale'])) |
| module.attn_pool.norm.bias.copy_(_n2p(w[f'{block_prefix}LayerNorm_0/bias'])) |
| for r in range(2): |
| getattr(module.attn_pool.mlp, f'fc{r + 1}').weight.copy_(_n2p(w[f'{block_prefix}MlpBlock_0/Dense_{r}/kernel'])) |
| getattr(module.attn_pool.mlp, f'fc{r + 1}').bias.copy_(_n2p(w[f'{block_prefix}MlpBlock_0/Dense_{r}/bias'])) |
|
|
| def _convert_openclip_transformer(module: Transformer, prefix): |
| for i, block in enumerate(module.resblocks.children()): |
| block_prefix = f'{prefix}encoderblock_{i}/' |
| mha_prefix = block_prefix + f'MultiHeadDotProductAttention_0/' |
| block.ln_1.weight.copy_(_n2p(w[f'{block_prefix}LayerNorm_0/scale'])) |
| block.ln_1.bias.copy_(_n2p(w[f'{block_prefix}LayerNorm_0/bias'])) |
| block.attn.in_proj_weight.copy_(torch.cat([ |
| _n2p(w[f'{mha_prefix}{n}/kernel'], t=False).flatten(1).T for n in ('query', 'key', 'value')])) |
| block.attn.in_proj_bias.copy_(torch.cat([ |
| _n2p(w[f'{mha_prefix}{n}/bias'], t=False).reshape(-1) for n in ('query', 'key', 'value')])) |
| block.attn.out_proj.weight.copy_(_n2p(w[f'{mha_prefix}out/kernel']).flatten(1)) |
| block.attn.out_proj.bias.copy_(_n2p(w[f'{mha_prefix}out/bias'])) |
| block.ln_2.weight.copy_(_n2p(w[f'{block_prefix}LayerNorm_1/scale'])) |
| block.ln_2.bias.copy_(_n2p(w[f'{block_prefix}LayerNorm_1/bias'])) |
| block.mlp.c_fc.weight.copy_(_n2p(w[f'{block_prefix}MlpBlock_0/Dense_0/kernel'])) |
| block.mlp.c_fc.bias.copy_(_n2p(w[f'{block_prefix}MlpBlock_0/Dense_0/bias'])) |
| block.mlp.c_proj.weight.copy_(_n2p(w[f'{block_prefix}MlpBlock_0/Dense_1/kernel'])) |
| block.mlp.c_proj.bias.copy_(_n2p(w[f'{block_prefix}MlpBlock_0/Dense_1/bias'])) |
|
|
| def _convert_openclip_txt(module: TextTransformer, prefix): |
| module.token_embedding.weight.copy_(_n2p(w[f'{prefix}Embed_0/embedding'], t=False)) |
| pos_embed_w = _n2p(w[f'{prefix}pos_embedding'], t=False).squeeze(0) |
| module.positional_embedding.copy_(pos_embed_w) |
| _convert_openclip_transformer(module.transformer, prefix=prefix + 'Encoder_0/') |
| module.ln_final.weight.copy_(_n2p(w[f'{prefix}Encoder_0/encoder_norm/scale'])) |
| module.ln_final.bias.copy_(_n2p(w[f'{prefix}Encoder_0/encoder_norm/bias'])) |
| module.text_projection.weight.copy_(_n2p(w[f'{prefix}head/kernel'])) |
| module.text_projection.bias.copy_(_n2p(w[f'{prefix}head/bias'])) |
|
|
| _convert_timm_img(model.visual.trunk, 'params/img/') |
| _convert_openclip_txt(model.text, 'params/txt/') |
| model.logit_bias.copy_(_n2p(w['params/b'])[0]) |
| model.logit_scale.copy_(_n2p(w['params/t'])[0]) |
|
|
|
|
| @torch.no_grad() |
| def convert_mobile_clip_state_dict(model: CustomTextCLIP, state_dict, fastvit = True): |
|
|
| def _convert_timm_img(state_dict): |
| if fastvit: |
| from timm.models.fastvit import checkpoint_filter_fn |
| else: |
| from timm.models.vision_transformer_hybrid import checkpoint_filter_fn |
| timm_state_dict = checkpoint_filter_fn(state_dict, model.visual.trunk) |
| timm_state_dict = {'visual.trunk.' + k: v for k, v in timm_state_dict.items()} |
| return timm_state_dict |
|
|
| def _convert_openclip_txt(state_dict, prefix='text_encoder.'): |
| text_dict = {} |
| for k, v in state_dict.items(): |
| if not k.startswith(prefix): |
| continue |
| k = k.replace(prefix, '') |
| k = k.replace('projection_layer', 'text_projection') |
| k = k.replace('embedding_layer', 'token_embedding') |
| if k.startswith('positional_embedding.pos_embed.pos_embed'): |
| k = k.replace('positional_embedding.pos_embed.pos_embed', 'positional_embedding') |
| v = v.squeeze() |
| k = k.replace('final_layer_norm', 'ln_final') |
| k = k.replace('pre_norm_mha.0', 'ln_1') |
| k = k.replace('pre_norm_mha.1', 'attn') |
| k = k.replace('pre_norm_ffn.0', 'ln_2') |
| k = k.replace('pre_norm_ffn.1', 'mlp.c_fc') |
| k = k.replace('pre_norm_ffn.4', 'mlp.c_proj') |
| k = k.replace('qkv_proj.weight', 'in_proj_weight') |
| k = k.replace('qkv_proj.bias', 'in_proj_bias') |
| k = k.replace('transformer.', 'transformer.resblocks.') |
| text_dict['text.' + k] = v |
| return text_dict |
|
|
| image_dict = _convert_timm_img(state_dict) |
| text_dict = _convert_openclip_txt(state_dict) |
| out_dict = {**image_dict, **text_dict} |
| out_dict['logit_scale'] = state_dict['logit_scale'] |
| return out_dict |
|
|
|
|
| def convert_state_dict(model: Union[CustomTextCLIP, CLIP], state_dict): |
| if 'image_encoder.model.patch_embed.0.rbr_conv.0.conv.weight' in state_dict: |
| |
| state_dict = convert_mobile_clip_state_dict(model, state_dict) |
| if 'image_encoder.model.patch_emb.0.block.conv.weight' in state_dict: |
| |
| state_dict = convert_mobile_clip_state_dict(model, state_dict, fastvit=False) |
| return state_dict |
|
|
