fix style switching bugs

adaface/adaface_wrapper.py

@@ -4,6 +4,7 @@ from transformers import CLIPTextModel
 from diffusers import (
     StableDiffusionPipeline,
     StableDiffusionImg2ImgPipeline,
+    StableDiffusionXLPipeline,
     StableDiffusion3Pipeline,
     #FluxPipeline,
     DDIMScheduler,
@@ -25,7 +26,7 @@ class AdaFaceWrapper(nn.Module):
                  main_unet_filepath=None, unet_types=None, extra_unet_dirpaths=None, unet_weights=None,
                  device='cuda', is_training=False):
         '''
-        pipeline_name: "text2img", "img2img", "text2img3", "flux", or None. 
+        pipeline_name: "text2img", "text2imgxl", "img2img", "text2img3", "flux", or None. 
         If None, it's used only as a face encoder, and the unet and vae are
         removed from the pipeline to release RAM.
         '''
@@ -64,6 +65,13 @@ class AdaFaceWrapper(nn.Module):
         self.encoders_num_id_vecs = self.id2ada_prompt_encoder.encoders_num_id_vecs
         self.extend_tokenizer_and_text_encoder()
 
+    def to(self, device):
+        self.device = device
+        self.id2ada_prompt_encoder.to(device)
+        self.pipeline.to(device)
+        print(f"Moved AdaFaceWrapper to {device}.")
+        return self
+    
     def initialize_pipeline(self):
         self.id2ada_prompt_encoder = create_id2ada_prompt_encoder(self.adaface_encoder_types,
                                                                   self.adaface_ckpt_paths,
@@ -95,6 +103,8 @@ class AdaFaceWrapper(nn.Module):
             PipelineClass = StableDiffusionImg2ImgPipeline
         elif self.pipeline_name == "text2img":
             PipelineClass = StableDiffusionPipeline
+        elif self.pipeline_name == "text2imgxl":
+            PipelineClass = StableDiffusionXLPipeline
         elif self.pipeline_name == "text2img3":
             PipelineClass = StableDiffusion3Pipeline
         #elif self.pipeline_name == "flux":
@@ -109,6 +119,7 @@ class AdaFaceWrapper(nn.Module):
         if self.base_model_path is None:
             base_model_path_dict = { 
                 'text2img':  'models/sd15-dste8-vae.safetensors',
+                'text2imgxl': 'stabilityai/stable-diffusion-xl-base-1.0',
                 'text2img3': 'stabilityai/stable-diffusion-3-medium-diffusers',
                 'flux':      'black-forest-labs/FLUX.1-schnell',
             }
@@ -156,7 +167,7 @@ class AdaFaceWrapper(nn.Module):
             pipeline.vae  = None
             print("Removed UNet and VAE from the pipeline.")
 
-        if self.pipeline_name not in ["text2img3", "flux"]:
+        if self.pipeline_name not in ["text2imgxl", "text2img3", "flux"]:
             noise_scheduler = DDIMScheduler(
                 num_train_timesteps=1000,
                 beta_start=0.00085,
@@ -277,6 +288,8 @@ class AdaFaceWrapper(nn.Module):
 
         return prompt
 
+    # If face_id_embs is None, then it extracts face_id_embs from the images,
+    # then map them to ada prompt embeddings.
     # avg_at_stage: 'id_emb', 'img_prompt_emb', or None.
     # avg_at_stage == ada_prompt_emb usually produces the worst results.
     # id_emb is slightly better than img_prompt_emb, but sometimes img_prompt_emb is better.
@@ -297,14 +310,18 @@ class AdaFaceWrapper(nn.Module):
         if all_adaface_subj_embs is None:
             return None
         
-        # [1, 1, 16, 768] -> [16, 768]
-        all_adaface_subj_embs = all_adaface_subj_embs.squeeze(0).squeeze(0)
+        if all_adaface_subj_embs.ndim == 4:
+            # [1, 1, 16, 768] -> [16, 768]
+            all_adaface_subj_embs = all_adaface_subj_embs.squeeze(0).squeeze(0)
+        elif all_adaface_subj_embs.ndim == 3:
+            # [1, 16, 768] -> [16, 768]
+            all_adaface_subj_embs = all_adaface_subj_embs.squeeze(0)
 
         if update_text_encoder:
             self.update_text_encoder_subj_embeddings(all_adaface_subj_embs)
         return all_adaface_subj_embs
 
-    def diffusers_encode_prompts(self, prompt, negative_prompt, device):
+    def diffusers_encode_prompts(self, prompt, plain_prompt, negative_prompt, device):
         # pooled_prompt_embeds_, negative_pooled_prompt_embeds_ are used by text2img3 and flux.
         pooled_prompt_embeds_, negative_pooled_prompt_embeds_ = None, None
 
@@ -318,17 +335,26 @@ class AdaFaceWrapper(nn.Module):
             prompt_embeds_ = prompt_embeds_.unsqueeze(0)
             negative_prompt_embeds_ = negative_prompt_embeds_.unsqueeze(0)
         else:
-            if self.pipeline_name in ["text2img3", "flux"]:
-                # prompt_embeds_, negative_prompt_embeds_: [1, 333, 4096]
-                # pooled_prompt_embeds_, negative_pooled_prompt_embeds_: [1, 2048]
+            if self.pipeline_name in ["text2imgxl", "text2img3", "flux"]:
+                prompt_2 = plain_prompt
                 # CLIP Text Encoder prompt uses a maximum sequence length of 77.
                 # T5 Text Encoder prompt uses a maximum sequence length of 256.
                 # 333 = 256 + 77.
                 prompt_t5 = prompt + "".join([", "] * 256)
-                if self.pipeline_name == "text2img3":
+
+                # prompt_embeds_, negative_prompt_embeds_: [1, 333, 4096]
+                # pooled_prompt_embeds_, negative_pooled_prompt_embeds_: [1, 2048]
+                if self.pipeline_name == "text2imgxl":
+                    prompt_embeds_, negative_prompt_embeds_, \
+                    pooled_prompt_embeds_, negative_pooled_prompt_embeds_ = \
+                        self.pipeline.encode_prompt(prompt, prompt_2, device=device, 
+                                                    num_images_per_prompt=1, 
+                                                    do_classifier_free_guidance=True,
+                                                    negative_prompt=negative_prompt)        
+                elif self.pipeline_name == "text2img3":
                     prompt_embeds_, negative_prompt_embeds_, \
                     pooled_prompt_embeds_, negative_pooled_prompt_embeds_ = \
-                        self.pipeline.encode_prompt(prompt, prompt, prompt_t5, device=device, 
+                        self.pipeline.encode_prompt(prompt, prompt_2, prompt_t5, device=device, 
                                                     num_images_per_prompt=1, 
                                                     do_classifier_free_guidance=True,
                                                     negative_prompt=negative_prompt)
@@ -362,6 +388,7 @@ class AdaFaceWrapper(nn.Module):
         if device is None:
             device = self.device
         
+        plain_prompt = prompt
         prompt = self.update_prompt(prompt, placeholder_tokens_pos=placeholder_tokens_pos)
         if verbose:
             print(f"Subject prompt:\n{prompt}")
@@ -373,11 +400,11 @@ class AdaFaceWrapper(nn.Module):
             negative_prompt      = self.update_prompt(negative_prompt0, placeholder_tokens_pos='prepend')
             null_negative_prompt = self.update_prompt(negative_prompt0, placeholder_tokens_pos='prepend',
                                                       use_null_placeholders=True)
-            if verbose:
-                print(f"do_neg_id_prompt_weight: {do_neg_id_prompt_weight}")
-                #print(f"Negative prompt:\n{negative_prompt}")
-                #print(f"Null negative prompt:\n{null_negative_prompt}")
+            '''         if verbose:
+                            print(f"Negative prompt:\n{negative_prompt}")
+                            print(f"Null negative prompt:\n{null_negative_prompt}")
 
+            '''        
         else:
             null_negative_prompt = None
 
@@ -386,11 +413,11 @@ class AdaFaceWrapper(nn.Module):
         self.pipeline.text_encoder.to(device)
 
         prompt_embeds_, negative_prompt_embeds_, pooled_prompt_embeds_, negative_pooled_prompt_embeds_ = \
-            self.diffusers_encode_prompts(prompt, negative_prompt, device)
+            self.diffusers_encode_prompts(prompt, plain_prompt, negative_prompt, device)
         
         if 0 < do_neg_id_prompt_weight < 1:
             _, negative_prompt_embeds_null, _, _ = \
-                self.diffusers_encode_prompts(prompt, null_negative_prompt, device)
+                self.diffusers_encode_prompts(prompt, plain_prompt, null_negative_prompt, device)
             negative_prompt_embeds_ = negative_prompt_embeds_ * do_neg_id_prompt_weight + \
                                       negative_prompt_embeds_null * (1 - do_neg_id_prompt_weight)
             
@@ -399,9 +426,8 @@ class AdaFaceWrapper(nn.Module):
     # ref_img_strength is used only in the img2img pipeline.
     def forward(self, noise, prompt, negative_prompt=None, 
                 placeholder_tokens_pos='append',
-                guidance_scale=6.0, 
                 do_neg_id_prompt_weight=0,
-                out_image_count=4, 
+                guidance_scale=6.0, out_image_count=4, 
                 ref_img_strength=0.8, generator=None, verbose=False):
         noise = noise.to(device=self.device, dtype=torch.float16)
 
@@ -419,7 +445,7 @@ class AdaFaceWrapper(nn.Module):
         if negative_prompt_embeds_ is not None:
             negative_prompt_embeds_ = negative_prompt_embeds_.repeat(out_image_count, 1, 1)
 
-        if self.pipeline_name == "text2img3":
+        if self.pipeline_name in ["text2imgxl", "text2img3"]:
             pooled_prompt_embeds_           = pooled_prompt_embeds_.repeat(out_image_count, 1)
             negative_pooled_prompt_embeds_  = negative_pooled_prompt_embeds_.repeat(out_image_count, 1)
 

app.py

@@ -24,7 +24,7 @@ parser = argparse.ArgumentParser()
 parser.add_argument("--adaface_encoder_types", type=str, nargs="+", default=["consistentID", "arc2face"],
                     choices=["arc2face", "consistentID"], help="Type(s) of the ID2Ada prompt encoders")
 parser.add_argument('--adaface_ckpt_path', type=str, 
-                    default='models/adaface/VGGface2_HQ_masks2024-10-13T11-21-07_zero3-ada-9000.pt')
+                    default='models/adaface/VGGface2_HQ_masks2024-10-14T16-09-24_zero3-ada-3500.pt')
 parser.add_argument('--model_style_type', type=str, default='realistic',
                     choices=["realistic", "anime", "photorealistic"], help="Type of the base model")
 parser.add_argument("--guidance_scale", type=float, default=6.0,
@@ -51,10 +51,10 @@ device = "cuda" if args.gpu is None else f"cuda:{args.gpu}"
 global adaface, id_animator
 
 base_model_path = model_style_type2base_model_path[args.model_style_type]
-id_animator = load_model(model_style_type=args.model_style_type, device=device)
+id_animator = load_model(model_style_type=args.model_style_type, device='cpu')
 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], device=device)
+                         adaface_ckpt_paths=[args.adaface_ckpt_path], device='cpu')
 
 basedir = os.getcwd()
 savedir = os.path.join(basedir,'samples')
@@ -80,11 +80,14 @@ def get_clicked_image(data: gr.SelectData):
     
 @spaces.GPU
 def gen_init_images(uploaded_image_paths, prompt, guidance_scale, do_neg_id_prompt_weight, out_image_count=4):
-    global adaface, id_animator
     if uploaded_image_paths is None:
         print("No image uploaded")
         return None, None, None
 
+    global adaface, id_animator
+    adaface.to(device)
+    id_animator.to(device)
+
     # uploaded_image_paths is a list of tuples:
     # [('/tmp/gradio/249981e66a7c665aaaf1c7eaeb24949af4366c88/jensen huang.jpg', None)]
     # Extract the file paths.
@@ -132,7 +135,10 @@ def generate_video(image_container, uploaded_image_paths, init_img_file_paths, i
                    seed, attn_scale, image_embed_cfg_begin_scale, image_embed_cfg_end_scale,
                    is_adaface_enabled, adaface_ckpt_path, adaface_power_scale, 
                    id_animator_anneal_steps, progress=gr.Progress(track_tqdm=True)):
+    
     global adaface, id_animator
+    adaface.to(device)
+    id_animator.to(device)
 
     if prompt is None:
         prompt = ""
@@ -220,7 +226,7 @@ def check_prompt_and_model_type(prompt, model_style_type):
         id_animator = load_model(model_style_type=model_style_type, device=device)
         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], device=device)
+                                 adaface_ckpt_paths=[args.adaface_ckpt_path], device='cpu')
         # Update base model type.
         args.model_style_type = model_style_type
 
@@ -243,7 +249,7 @@ with gr.Blocks(css=css, theme=gr.themes.Origin()) as demo:
 
 ❗️**Tips**❗️
 - You can upload one or more subject images for generating ID-specific video.
-- If the face dominates the video frames, try decreasing the guidance scale.
+- If the face dominates the video frames, try increasing the 'Weight of ID prompt in the negative prompt'.
 - If the face loses focus, try increasing the guidance scale.
 - If the motion is weird, e.g., the prompt is "... running", try increasing the number of sampling steps.
 - Usage explanations and demos: [Readme](https://huggingface.co/spaces/adaface-neurips/adaface-animate/blob/main/README2.md).
@@ -344,7 +350,7 @@ with gr.Blocks(css=css, theme=gr.themes.Origin()) as demo:
                 maximum=0.9,
                 step=0.1,
                 value=args.do_neg_id_prompt_weight,
-                visible=False
+                visible=True
             )
 
             seed = gr.Slider(

faceadapter/face_adapter.py

@@ -1,11 +1,8 @@
-import os
 from typing import List
 
 import torch
-from diffusers import StableDiffusionPipeline
-from diffusers.pipelines.controlnet import MultiControlNetModel
+from torch import nn
 from PIL import Image
-from safetensors import safe_open
 from transformers import CLIPImageProcessor, CLIPVisionModelWithProjection
 from .attention_processor import LoRAFaceAttnProcessor
 
@@ -78,6 +75,13 @@ class FaceAdapterLora:
 
         self.load_face_adapter()
 
+    def to(self, device):
+        self.device = device
+        self.pipe = self.pipe.to(device)
+        self.image_encoder = self.image_encoder.to(device)
+        self.image_proj_model = self.image_proj_model.to(device)
+        return self
+    
     def init_proj(self):
         image_proj_model = ImageProjModel(
             cross_attention_dim=self.pipe.unet.config.cross_attention_dim,

models/adaface/{VGGface2_HQ_masks2024-10-13T11-21-07_zero3-ada-9000.pt → VGGface2_HQ_masks2024-10-14T16-09-24_zero3-ada-3500.pt}

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:c8d4dae6291aa783c03de6e54997a79b7d61f7bdc2457afb6303bcfc290d0956
-size 1814921242
+oid sha256:d991773c169ecfa01ce1ec365d7791d9b732212b35118d0795f21cc0c1a99e77
+size 1814921754