extend CLIP text encoder to 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

@@ -50,17 +50,33 @@ parser.add_argument("--q_lora_updates_query", type=str2bool, nargs="?", const=Tr
                          "If False, the q lora only updates query2.")
 parser.add_argument("--show_disable_adaface_checkbox", type=str2bool, nargs="?", const=True, default=False,
                     help="Whether to show the checkbox for disabling AdaFace")
+parser.add_argument('--show_ablate_prompt_embed_type', type=str2bool, nargs="?", const=True, default=False,
+                    help="Whether to show the dropdown for ablate prompt embeddings type")
 parser.add_argument('--extra_save_dir', type=str, default=None, help="Directory to save the generated images")
 parser.add_argument('--test_ui_only', type=str2bool, nargs="?", const=True, default=False,
                     help="Only test the UI layout, and skip loadding the adaface model")
+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")
 args = parser.parse_args()
 
-from huggingface_hub import snapshot_download
-large_files = ["models/*", "models/**/*"]
-snapshot_download(repo_id="adaface-neurips/adaface-models", repo_type="model", allow_patterns=large_files, local_dir=".")
+
+if is_running_on_hf_space():
+    args.device = 'cuda:0'
+    is_on_hf_space = True
+else:
+    if args.gpu is None:
+        args.device = "cuda" 
+    else:
+        args.device = f"cuda:{args.gpu}"
+    is_on_hf_space = False
+
 os.makedirs("/tmp/gradio", exist_ok=True)
+from huggingface_hub import snapshot_download
+if is_on_hf_space:
+    large_files = ["models/*", "models/**/*"]
+    snapshot_download(repo_id="adaface-neurips/adaface-models", repo_type="model", allow_patterns=large_files, local_dir=".")
 
 model_style_type2base_model_path = {
     "realistic": "models/rv51/realisticVisionV51_v51VAE_dste8.safetensors",
@@ -75,22 +91,13 @@ MAX_SEED = np.iinfo(np.int32).max
 global adaface
 adaface = None
 
-if is_running_on_hf_space():
-    args.device = 'cuda:0'
-    is_on_hf_space = True
-else:
-    if args.gpu is None:
-        args.device = "cuda" 
-    else:
-        args.device = f"cuda:{args.gpu}"
-    is_on_hf_space = False
-
 if not args.test_ui_only:
     adaface = AdaFaceWrapper(pipeline_name="text2img", base_model_path=base_model_path,
                              adaface_encoder_types=args.adaface_encoder_types, 
                              adaface_ckpt_paths=args.adaface_ckpt_path, 
                              adaface_encoder_cfg_scales=args.adaface_encoder_cfg_scales,
                              enabled_encoders=args.enabled_encoders,
+                             max_prompt_length=args.max_prompt_length,
                              unet_types=None, extra_unet_dirpaths=None, unet_weights_in_ensemble=None, 
                              unet_uses_attn_lora=args.unet_uses_attn_lora,
                              attn_lora_layer_names=args.attn_lora_layer_names,
@@ -120,7 +127,7 @@ def remove_back_to_files():
 @spaces.GPU
 def generate_image(image_paths, image_paths2, guidance_scale, perturb_std,
                    num_images, prompt, negative_prompt, gender, highlight_face, 
-                   ablate_prompt_embed_type, nonmix_prompt_emb_weight,
+                   ablate_prompt_embed_type, 
                    composition_level, seed, disable_adaface, subj_name_sig, progress=gr.Progress(track_tqdm=True)):
 
     global adaface, args
@@ -168,6 +175,12 @@ def generate_image(image_paths, image_paths2, guidance_scale, perturb_std,
         else:
             prompt = gender + ", " + prompt
 
+    if 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(ablate_prompt_embed_type) + 1
+        ablate_prompt_embed_type = f"img{ablate_prompt_embed_index}"
+
     generator = torch.Generator(device=adaface.pipeline._execution_device).manual_seed(seed)
     samples = adaface(noise, prompt, negative_prompt=negative_prompt, 
                       guidance_scale=guidance_scale, 
@@ -175,7 +188,6 @@ def generate_image(image_paths, image_paths2, guidance_scale, perturb_std,
                       repeat_prompt_for_each_encoder=(composition_level >= 1),
                       ablate_prompt_no_placeholders=disable_adaface,
                       ablate_prompt_embed_type=ablate_prompt_embed_type,
-                      nonmix_prompt_emb_weight=nonmix_prompt_emb_weight,
                       verbose=True)
 
     session_signature = ",".join(image_paths + [prompt, str(seed)])
@@ -387,14 +399,8 @@ with gr.Blocks(css=css, theme=gr.themes.Origin()) as demo:
                           minimum=0, maximum=2, step=1, value=0)
 
             ablate_prompt_embed_type = gr.Dropdown(label="Ablate prompt embeddings type",
-                                                   choices=["ada", "ada-nonmix", "img"], value="ada", visible=False,
+                                                   choices=["ada", "arc2face", "consistentID"], value="ada", visible=args.show_ablate_prompt_embed_type,
                                                    info="Use this type of prompt embeddings for ablation study")
-            
-            nonmix_prompt_emb_weight = gr.Slider(label="Weight of ada-nonmix ID embeddings",
-                                                 minimum=0.0, maximum=0.5, step=0.1, value=0,
-                                                 info="Weight of ada-nonmix ID embeddings in the prompt embeddings",
-                                                 visible=False)
-                                        
 
             subj_name_sig = gr.Textbox(
                 label="Nickname of Subject (optional; used to name saved images)", 
@@ -497,7 +503,7 @@ with gr.Blocks(css=css, theme=gr.themes.Origin()) as demo:
             'fn': generate_image,
             'inputs': [img_files, img_files2, guidance_scale, perturb_std, num_images, prompt, 
                        negative_prompt, gender, highlight_face, ablate_prompt_embed_type, 
-                       nonmix_prompt_emb_weight, composition_level, seed, disable_adaface, subj_name_sig],
+                       composition_level, seed, disable_adaface, subj_name_sig],
             'outputs': [out_gallery]
         }
         submit.click(**check_prompt_and_model_type_call_dict).success(**randomize_seed_fn_call_dict).then(**generate_image_call_dict)