add link to adaface, various improvements

adaface/adaface_wrapper.py

@@ -12,6 +12,8 @@ from insightface.app import FaceAnalysis
 from adaface.arc2face_models import CLIPTextModelWrapper
 from adaface.util import get_arc2face_id_prompt_embs
 import re, os
+import sys
+sys.modules['ldm'] = sys.modules['adaface']
 
 class AdaFaceWrapper(nn.Module):
     def __init__(self, pipeline_name, base_model_path, adaface_ckpt_path, device, 
@@ -216,7 +218,7 @@ class AdaFaceWrapper(nn.Module):
         # NOTE: Since return_core_id_embs is True, id_prompt_emb is only the 16 core ID embeddings.
         # arc2face prompt template: "photo of a id person"
         # ID embeddings start from "id person ...". So there are 3 template tokens before the 16 ID embeddings.
-        faceid_embeds, id_prompt_emb \
+        face_image_count, faceid_embeds, id_prompt_emb \
             = get_arc2face_id_prompt_embs(self.face_app, self.pipeline.tokenizer, self.arc2face_text_encoder,
                                           extract_faceid_embeds=not gen_rand_face,
                                           pre_face_embs=pre_face_embs,
@@ -235,6 +237,9 @@ class AdaFaceWrapper(nn.Module):
                                           gen_neg_prompt=False, 
                                           verbose=True)
         
+        if face_image_count == 0:
+            return None
+                
         # adaface_subj_embs: [1, 1, 16, 768]. 
         # adaface_prompt_embs: [1, 77, 768] (not used).
         adaface_subj_embs, adaface_prompt_embs = \
@@ -248,7 +253,7 @@ class AdaFaceWrapper(nn.Module):
             self.update_text_encoder_subj_embs(adaface_subj_embs)
         return adaface_subj_embs
 
-    def encode_prompt(self, prompt, device="cuda", verbose=False):
+    def encode_prompt(self, prompt, negative_prompt, device="cuda", verbose=False):
         prompt = self.update_prompt(prompt)
         if verbose:
             print(f"Prompt: {prompt}")
@@ -259,14 +264,16 @@ class AdaFaceWrapper(nn.Module):
         # prompt_embeds_, negative_prompt_embeds_: [1, 77, 768]
         prompt_embeds_, negative_prompt_embeds_ = \
             self.pipeline.encode_prompt(prompt, device=device, num_images_per_prompt=1,
-                                        do_classifier_free_guidance=True, negative_prompt=self.negative_prompt)
+                                        do_classifier_free_guidance=True, negative_prompt=negative_prompt)
         return prompt_embeds_, negative_prompt_embeds_
     
     # ref_img_strength is used only in the img2img pipeline.
-    def forward(self, noise, prompt, guidance_scale=4.0, out_image_count=4, ref_img_strength=0.8, verbose=False):
+    def forward(self, noise, prompt, negative_prompt=None, guidance_scale=4.0, 
+                out_image_count=4, ref_img_strength=0.8, generator=None, verbose=False):
+        if negative_prompt is None:
+            negative_prompt = self.negative_prompt
         # prompt_embeds_, negative_prompt_embeds_: [1, 77, 768]
-        prompt_embeds_, negative_prompt_embeds_ = self.encode_prompt(prompt, device=self.device, verbose=verbose)
-
+        prompt_embeds_, negative_prompt_embeds_ = self.encode_prompt(prompt, negative_prompt, device=self.device, verbose=verbose)
         # Repeat the prompt embeddings for all images in the batch.
         prompt_embeds_          = prompt_embeds_.repeat(out_image_count, 1, 1)
         negative_prompt_embeds_ = negative_prompt_embeds_.repeat(out_image_count, 1, 1)
@@ -280,7 +287,8 @@ class AdaFaceWrapper(nn.Module):
                                num_inference_steps=self.num_inference_steps, 
                                guidance_scale=guidance_scale, 
                                num_images_per_prompt=1,
-                               strength=ref_img_strength).images
+                               strength=ref_img_strength,
+                               generator=generator).images
         # images: [BS, 3, 512, 512]
         return images
     

adaface/util.py

@@ -246,8 +246,9 @@ def get_arc2face_id_prompt_embs(face_app, clip_tokenizer, arc2face_text_encoder,
                                 input_max_length=77, noise_level=0.0, 
                                 return_core_id_embs=False,
                                 gen_neg_prompt=False, verbose=False):
+    face_image_count = 0
+
     if extract_faceid_embeds:
-        image_count = 0
         faceid_embeds = []
         if image_paths is not None:
             images_np = []
@@ -275,13 +276,13 @@ def get_arc2face_id_prompt_embs(face_app, clip_tokenizer, arc2face_text_encoder,
             face_info = sorted(face_infos, key=lambda x:(x['bbox'][2]-x['bbox'][0])*x['bbox'][3]-x['bbox'][1])[-1]
             # Each faceid_embed: [1, 512]
             faceid_embeds.append(torch.from_numpy(face_info.normed_embedding).unsqueeze(0))
-            image_count += 1
+            face_image_count += 1
 
         if verbose:
             if image_folder is not None:
-                print(f"Extracted ID embeddings from {image_count} images in {image_folder}")
+                print(f"Extracted ID embeddings from {face_image_count} images in {image_folder}")
             else:
-                print(f"Extracted ID embeddings from {image_count} images")
+                print(f"Extracted ID embeddings from {face_image_count} images")
 
         if len(faceid_embeds) == 0:
             print("No face detected. Use a random face instead.")
@@ -335,7 +336,7 @@ def get_arc2face_id_prompt_embs(face_app, clip_tokenizer, arc2face_text_encoder,
 
         #if extract_faceid_embeds:
         #    arc2face_neg_prompt_emb = arc2face_neg_prompt_emb.repeat(id_batch_size, 1, 1)
-        return faceid_embeds, arc2face_pos_prompt_emb, arc2face_neg_prompt_emb
+        return face_image_count, faceid_embeds, arc2face_pos_prompt_emb, arc2face_neg_prompt_emb
     else:
-        return faceid_embeds, arc2face_pos_prompt_emb
+        return face_image_count, faceid_embeds, arc2face_pos_prompt_emb
     

app.py

@@ -76,8 +76,13 @@ def gen_init_images(uploaded_image_paths, prompt, adaface_id_cfg_scale, out_imag
     # [('/tmp/gradio/249981e66a7c665aaaf1c7eaeb24949af4366c88/jensen huang.jpg', None)]
     # Extract the file paths.
     uploaded_image_paths = [path[0] for path in uploaded_image_paths]
-    adaface.generate_adaface_embeddings(image_folder=None, image_paths=uploaded_image_paths,
-                                        out_id_embs_scale=adaface_id_cfg_scale, update_text_encoder=True)
+    adaface_subj_embs = \
+        adaface.generate_adaface_embeddings(image_folder=None, image_paths=uploaded_image_paths,
+                                            out_id_embs_scale=adaface_id_cfg_scale, update_text_encoder=True)
+    
+    if adaface_subj_embs is None:
+        raise gr.Error(f"Failed to detect any faces! Please try with other images")
+        
     # Generate two images each time for the user to select from.
     noise = torch.randn(out_image_count, 3, 512, 512)
     # samples: A list of PIL Image instances.
@@ -163,7 +168,7 @@ def generate_image(image_container, uploaded_image_paths, init_img_file_paths, i
     save_videos_grid(sample, save_sample_path)
     return save_sample_path
 
-def validate(prompt):
+def validate_prompt(prompt):
     if not prompt:
         raise gr.Error("Prompt cannot be blank")
 
@@ -229,7 +234,11 @@ with gr.Blocks(css=css) as demo:
     ❗️❗️❗️**Tips:**
     - You can upload one or more subject images for generating ID-specific video.
     - Try different parameter combinations for the best generation quality.
-    - Technical explanations and demo videos: [Readme](https://huggingface.co/spaces/adaface-neurips/adaface-animate/blob/main/README2.md).
+    - Usage explanations and demos: [Readme](https://huggingface.co/spaces/adaface-neurips/adaface-animate/blob/main/README2.md).
+    - AdaFace Text-to-Image: <a href="https://huggingface.co/spaces/adaface-neurips/adaface" style="display: inline-flex; align-items: center;">
+  AdaFace 
+  <img src="https://img.shields.io/badge/%F0%9F%A4%97%20Hugging%20Face-Spaces-yellow" alt="Hugging Face Spaces" style="margin-left: 5px;">
+  </a>
         """
     )
 
@@ -401,7 +410,7 @@ with gr.Blocks(css=css) as demo:
                        outputs=[uploaded_init_img_gallery, init_img_files, init_clear_button_column])
         uploaded_init_img_gallery.select(fn=get_clicked_image, inputs=None, outputs=init_img_selected_idx)
 
-        submit.click(fn=validate,
+        submit.click(fn=validate_prompt,
                      inputs=[prompt],outputs=None).success(
             fn=randomize_seed_fn,
             inputs=[seed, randomize_seed],

infer.py

@@ -64,7 +64,8 @@ def load_model(base_model_type="sar", adaface_base_model_type="sar",
             # scheduler=DPMSolverMultistepScheduler(**OmegaConf.to_container(inference_config.DPMSolver_scheduler_kwargs)
             # scheduler=EulerAncestralDiscreteScheduler(**OmegaConf.to_container(inference_config.noise_scheduler_kwargs)
             # scheduler=EulerAncestralDiscreteScheduler(beta_start=0.00085, beta_end=0.012, beta_schedule="linear",steps_offset=1
-                                   ),torch_dtype=torch.float16,
+                                   ),
+            torch_dtype=torch.float16,
             ).to(device=device)
     
     pipeline = load_weights(