Extend CLIP text encoder to support 97 tokens

adaface/adaface_wrapper.py

@@ -14,7 +14,7 @@ from diffusers import (
     LCMScheduler,
 )
 from diffusers.loaders.single_file_utils import convert_ldm_unet_checkpoint
-from adaface.util import UNetEnsemble
+from adaface.util import UNetEnsemble, extend_nn_embedding
 from adaface.face_id_to_ada_prompt import create_id2ada_prompt_encoder
 from adaface.diffusers_attn_lora_capture import set_up_attn_processors, set_up_ffn_loras, set_lora_and_capture_flags
 from safetensors.torch import load_file as safetensors_load_file
@@ -27,7 +27,7 @@ class AdaFaceWrapper(nn.Module):
                  adaface_ckpt_paths, adaface_encoder_cfg_scales=None, 
                  enabled_encoders=None, use_lcm=False, default_scheduler_name='ddim',
                  num_inference_steps=50, subject_string='z', negative_prompt=None,
-                 use_840k_vae=False, use_ds_text_encoder=False, 
+                 max_prompt_length=77, use_840k_vae=False, use_ds_text_encoder=False, 
                  main_unet_filepath=None, unet_types=None, extra_unet_dirpaths=None, unet_weights_in_ensemble=None,
                  enable_static_img_suffix_embs=None, unet_uses_attn_lora=False,
                  attn_lora_layer_names=['q', 'k', 'v', 'out'], normalize_cross_attn=False, q_lora_updates_query=False,
@@ -56,6 +56,9 @@ class AdaFaceWrapper(nn.Module):
 
         self.default_scheduler_name = default_scheduler_name
         self.num_inference_steps = num_inference_steps if not use_lcm else 4
+
+        self.max_prompt_length = max_prompt_length
+
         self.use_840k_vae = use_840k_vae
         self.use_ds_text_encoder = use_ds_text_encoder
         self.main_unet_filepath = main_unet_filepath
@@ -199,6 +202,21 @@ class AdaFaceWrapper(nn.Module):
                                                 
             pipeline.unet = unet2
 
+        # Extending prompt length is for SD 1.5 only.
+        if (self.pipeline_name == "text2img") and (self.max_prompt_length > 77):
+            # pipeline.text_encoder.text_model.embeddings.position_embedding.weight: [77, 768] -> [max_length, 768]
+            # We reuse the last EL position embeddings for the new position embeddings.
+            # If we use the "neat" way, i.e., initialize CLIPTextModel with a CLIPTextConfig with 
+            # a larger max_position_embeddings, and set ignore_mismatched_sizes=True, 
+            # then the old position embeddings won't be loaded from the pretrained ckpt, 
+            # leading to degenerated performance. 
+            EL = self.max_prompt_length - 77
+            # position_embedding.weight: [77, 768] -> [max_length, 768]
+            new_position_embedding = extend_nn_embedding(pipeline.text_encoder.text_model.embeddings.position_embedding,
+                                                         pipeline.text_encoder.text_model.embeddings.position_embedding.weight[-EL:])
+            pipeline.text_encoder.text_model.embeddings.position_embedding = new_position_embedding
+            pipeline.text_encoder.text_model.embeddings.position_ids = torch.arange(self.max_prompt_length).unsqueeze(0)
+
         if self.use_840k_vae:
             pipeline.vae = vae
             print("Replaced the VAE with the 840k-step VAE.")
@@ -715,7 +733,7 @@ class AdaFaceWrapper(nn.Module):
                 ref_img_strength=0.8, generator=None, 
                 ablate_prompt_only_placeholders=False,
                 ablate_prompt_no_placeholders=False,
-                ablate_prompt_embed_type='ada', # 'ada', 'ada-nonmix', 'img'
+                ablate_prompt_embed_type='ada', # 'ada', 'ada-nonmix', 'img1', 'img2'.
                 nonmix_prompt_emb_weight=0,
                 repeat_prompt_for_each_encoder=True,                
                 verbose=False):

adaface/util.py

@@ -73,6 +73,26 @@ def calc_stats(emb_name, embeddings, mean_dim=-1):
     print("Norms: min: %.4f, max: %.4f, mean: %.4f, std: %.4f" %(norms.min(), norms.max(), norms.mean(), norms.std()))
 
 
+# new_token_embeddings: [new_num_tokens, 768].
+def extend_nn_embedding(old_nn_embedding, new_token_embeddings):
+    emb_dim         = old_nn_embedding.embedding_dim
+    num_old_tokens  = old_nn_embedding.num_embeddings
+    num_new_tokens  = new_token_embeddings.shape[0]
+    num_tokens2     = num_old_tokens + num_new_tokens
+    
+    new_nn_embedding = nn.Embedding(num_tokens2, emb_dim, 
+                                    device=old_nn_embedding.weight.device,
+                                    dtype=old_nn_embedding.weight.dtype)
+
+    old_num_tokens = old_nn_embedding.weight.shape[0]
+    # Copy the first old_num_tokens embeddings from old_nn_embedding to new_nn_embedding.
+    new_nn_embedding.weight.data[:old_num_tokens] = old_nn_embedding.weight.data
+    # Copy the new embeddings to new_nn_embedding.
+    new_nn_embedding.weight.data[old_num_tokens:] = new_token_embeddings
+
+    print(f"Extended nn.Embedding from {num_old_tokens} to {num_tokens2} tokens.")
+    return new_nn_embedding
+
 # Revised from RevGrad, by removing the grad negation.
 class ScaleGrad(torch.autograd.Function):
     @staticmethod

app.py

@@ -34,6 +34,8 @@ parser.add_argument('--num_inference_steps', type=int, default=50,
 parser.add_argument('--ablate_prompt_embed_type', type=str, default='ada',
                     choices=["ada", "arc2face", "consistentID"],
                     help="Ablate to use the image ID embs instead of Ada embs")
+parser.add_argument('--max_prompt_length', type=int, default=97, 
+                    help="Maximum length of the prompt. If > 77, the CLIP text encoder will be extended.")
 
 parser.add_argument('--gpu', type=int, default=None)
 parser.add_argument('--ip', type=str, default="0.0.0.0")
@@ -79,6 +81,7 @@ adaface = AdaFaceWrapper(pipeline_name="text2img", base_model_path=adaface_base_
                          adaface_encoder_types=args.adaface_encoder_types,
                          adaface_ckpt_paths=args.adaface_ckpt_path, device='cpu',
                          num_inference_steps=args.num_inference_steps,
+                         max_prompt_length=args.max_prompt_length,
                          is_on_hf_space=is_on_hf_space)
 
 basedir = os.getcwd()
@@ -208,7 +211,7 @@ def generate_video(image_container, uploaded_image_paths, init_img_file_paths, i
             if args.ablate_prompt_embed_type != "ada":
                 # Find the prompt_emb_type index in adaface_encoder_types
                 # adaface_encoder_types: ["consistentID", "arc2face"]
-                ablate_prompt_embed_index = args.adaface_encoder_types.index(args.ablate_prompt_embed_type)
+                ablate_prompt_embed_index = args.adaface_encoder_types.index(args.ablate_prompt_embed_type) + 1
                 ablate_prompt_embed_type = f"img{ablate_prompt_embed_index}"
             else:
                 ablate_prompt_embed_type = "ada"
@@ -270,6 +273,7 @@ def check_prompt_and_model_type(prompt, model_style_type, progress=gr.Progress()
                                  adaface_encoder_types=args.adaface_encoder_types,
                                  adaface_ckpt_paths=[args.adaface_ckpt_path], device='cpu',
                                  num_inference_steps=args.num_inference_steps,
+                                 max_prompt_length=args.max_prompt_length,                                 
                                  is_on_hf_space=is_on_hf_space)
         # Update base model type.
         args.model_style_type = model_style_type