update consisid

Author: SHYuanBest

src/diffusers/models/transformers/consisid_transformer_3d.py

@@ -75,9 +75,9 @@ def reshape_tensor(x, heads):
 class PerceiverAttention(nn.Module):
     """
     Implements the Perceiver attention mechanism with multi-head attention.
-    
+
     This layer takes two inputs: 'x' (image features) and 'latents' (latent features),
-    applying multi-head attention to both and producing an output tensor with the same 
+    applying multi-head attention to both and producing an output tensor with the same
     dimension as the input tensor 'x'.
 
     Args:
@@ -522,19 +522,19 @@ class ConsisIDTransformer3DModel(ModelMixin, ConfigMixin, PeftAdapterMixin):
             which can help reduce computational overhead.
         LFE_num_tokens (`int`, defaults to `32`):
             The number of tokens to use in the Local Facial Extractor (LFE).
-            This module is responsible for capturing high frequency representations 
+            This module is responsible for capturing high frequency representations
             of the face.
         LFE_output_dim (`int`, defaults to `768`):
             The output dimension of the Local Facial Extractor (LFE) module.
-            This dimension determines the size of the feature vectors produced 
+            This dimension determines the size of the feature vectors produced
             by the LFE module.
         LFE_heads (`int`, defaults to `12`):
             The number of attention heads used in the Local Facial Extractor (LFE) module.
-            More heads may improve the ability to capture diverse features, but 
+            More heads may improve the ability to capture diverse features, but
             can also increase computational complexity.
         local_face_scale (`float`, defaults to `1.0`):
-            A scaling factor used to adjust the importance of local facial features 
-            in the model. This can influence how strongly the model focuses on 
+            A scaling factor used to adjust the importance of local facial features
+            in the model. This can influence how strongly the model focuses on
             high frequency face-related content.
     """
 

src/diffusers/pipelines/consisid/util_clip/__init__.py

@@ -1,11 +1,36 @@
 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 .factory import (
+    add_model_config,
+    create_model,
+    create_model_and_transforms,
+    create_model_from_pretrained,
+    create_transforms,
+    get_model_config,
+    get_tokenizer,
+    list_models,
+    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 .model import (
+    CLIP,
+    CLIPTextCfg,
+    CLIPVisionCfg,
+    CustomCLIP,
+    convert_weights_to_fp16,
+    convert_weights_to_lp,
+    get_cast_dtype,
+    trace_model,
+)
+from .openai import list_openai_models, load_openai_model
+from .pretrained import (
+    download_pretrained,
+    download_pretrained_from_url,
+    get_pretrained_cfg,
+    get_pretrained_url,
+    is_pretrained_cfg,
+    list_pretrained,
+    list_pretrained_models_by_tag,
+    list_pretrained_tags_by_model,
+)
 from .tokenizer import SimpleTokenizer, tokenize
-from .transform import image_transform
+from .transform import image_transform

src/diffusers/pipelines/consisid/util_clip/eva_vit_model.py

@@ -4,16 +4,20 @@
 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 .rope import VisionRotaryEmbeddingFast
 from .transformer import PatchDropout
-from .rope import VisionRotaryEmbedding, VisionRotaryEmbeddingFast
+
 
 if os.getenv('ENV_TYPE') == 'deepspeed':
     try:
@@ -24,7 +28,6 @@
     from torch.utils.checkpoint import checkpoint
 
 try:
-    import xformers
     import xformers.ops as xops
     XFORMERS_IS_AVAILBLE = True
 except:
@@ -39,19 +42,19 @@ def __init__(self, drop_prob=None):
 
     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, 
+        self,
+        in_features,
+        hidden_features=None,
+        out_features=None,
+        act_layer=nn.GELU,
+        norm_layer=nn.LayerNorm,
         drop=0.,
         subln=False,
 
@@ -71,15 +74,15 @@ def forward(self, x):
         x = self.fc1(x)
         x = self.act(x)
         # x = self.drop(x)
-        # commit this for the orignal BERT implement 
+        # 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., 
+    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
@@ -91,7 +94,7 @@ def __init__(self, in_features, hidden_features=None, out_features=None, act_lay
         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):
@@ -172,20 +175,20 @@ def __init__(
 
     def forward(self, x, rel_pos_bias=None, attn_mask=None):
         B, N, C = x.shape
-        if self.subln: 
+        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: 
+            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]
@@ -232,7 +235,7 @@ def forward(self, x, rel_pos_bias=None, attn_mask=None):
             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)
 
@@ -262,15 +265,15 @@ def __init__(self, dim, num_heads, mlp_ratio=4., qkv_bias=False, qk_scale=None,
 
         if naiveswiglu:
             self.mlp = SwiGLU(
-                in_features=dim, 
-                hidden_features=mlp_hidden_dim, 
+                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, 
+                in_features=dim,
+                hidden_features=mlp_hidden_dim,
                 act_layer=act_layer,
                 subln=subln,
                 drop=drop
@@ -407,7 +410,7 @@ def __init__(self, img_size=224, patch_size=16, in_chans=3, num_classes=1000, em
                 ft_seq_len=hw_seq_len if intp_freq else None,
                 # patch_dropout=patch_dropout
             )
-        else: 
+        else:
             self.rope = None
 
         self.naiveswiglu = naiveswiglu
@@ -469,7 +472,7 @@ def _init_weights(self, m):
 
     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():
@@ -491,7 +494,7 @@ def reset_classifier(self, num_classes, global_pool=''):
         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()
 

src/diffusers/pipelines/consisid/util_clip/factory.py

@@ -1,24 +1,23 @@
 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
+from typing import Optional, Tuple, Union
+
 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 .model import CLIP, CustomCLIP, 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 .pretrained import download_pretrained, get_pretrained_cfg, is_pretrained_cfg, list_pretrained_tags_by_model
 from .tokenizer import HFTokenizer, tokenize
-from .utils import resize_clip_pos_embed, resize_evaclip_pos_embed, resize_visual_pos_embed, resize_eva_pos_embed
+from .transform import image_transform
+from .utils import resize_clip_pos_embed, resize_eva_pos_embed, resize_evaclip_pos_embed, resize_visual_pos_embed
 
 
-_MODEL_CONFIG_PATHS = [Path(__file__).parent / f"model_configs/"]
+_MODEL_CONFIG_PATHS = [Path(__file__).parent / "model_configs/"]
 _MODEL_CONFIGS = {}  # directory (model_name: config) of model architecture configs
 
 
@@ -93,7 +92,7 @@ def load_state_dict(checkpoint_path: str, map_location: str='cpu', model_key: st
                 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")
@@ -181,7 +180,7 @@ def load_pretrained_checkpoint(
             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)
@@ -202,7 +201,7 @@ def load_pretrained_checkpoint(
             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}")
 
@@ -255,7 +254,7 @@ def create_model(
         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
@@ -286,7 +285,7 @@ def create_model(
                                checkpoint_path,
                                model_key="model|module|state_dict",
                                strict=False
-                               ) 
+                               )
             else:
                 error_str = (
                     f'Pretrained weights ({pretrained}) not found for model {model_name}.'
@@ -296,7 +295,7 @@ def create_model(
         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)
@@ -321,7 +320,7 @@ def create_model(
                 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:
@@ -338,7 +337,7 @@ def create_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)