update ipadapter

lyrasd_model/module/lyrasd_ip_adapter.py

@@ -45,17 +45,11 @@ class LyraIPAdapter:
             image_encoder_path=None,
             num_ip_tokens=4,
             ip_projection_dim=None,
-            fp_ckpt=None,
-            num_fp_tokens=1,
-            fp_projection_dim=None,
         ):
         self.pipe = sd_pipe
         self.device = device
-        self.fp_ckpt = fp_ckpt
         self.ip_ckpt = ip_ckpt
-        self.num_fp_tokens = num_fp_tokens
         self.num_ip_tokens = num_ip_tokens
-        self.fp_projection_dim = fp_projection_dim
         self.ip_projection_dim = ip_projection_dim
         self.sdxl = sdxl
         self.ip_plus = ip_plus
@@ -76,10 +70,6 @@ class LyraIPAdapter:
             else:
                 self.image_proj_model = self.init_proj(self.ip_projection_dim, self.num_ip_tokens)
 
-        # face proj model
-        if self.fp_ckpt:
-            self.face_proj_model = self.init_proj(self.fp_projection_dim, self.num_fp_tokens)
-        
         self.load_ip_adapter()
         
     def init_proj_diffuser(self, state_dict):
@@ -131,16 +121,9 @@ class LyraIPAdapter:
             pretrained_path, subfolder, weight_name = parse_ckpt_path(self.ip_ckpt)
             dir_ipadapter = os.path.join(pretrained_path, "lyra_tran", subfolder, '.'.join(weight_name.split(".")[:-1]))
             unet.load_ip_adapter(dir_ipadapter, "", 1, "fp16")
-
-        if self.fp_ckpt:
-            state_dict = torch.load(self.fp_ckpt, map_location="cpu")
-            self.face_proj_model.load_state_dict(state_dict["face_proj"])
-            pretrained_path, subfolder, weight_name = parse_ckpt_path(self.fp_ckpt)
-            dir_ipadapter = os.path.join(pretrained_path, "lyra_tran", subfolder, '.'.join(weight_name.split(".")[:-1]))
-            unet.load_facein(dir_ipadapter, "fp16")
         
     @torch.inference_mode()
-    def get_image_embeds(self, image=None, face_emb=None):
+    def get_image_embeds(self, image=None):
         image_prompt_embeds, uncond_image_prompt_embeds = None, None
 
         if image is not None:
@@ -160,22 +143,11 @@ class LyraIPAdapter:
             uncond_clip_image_prompt_embeds = self.image_proj_model(uncond_clip_image_embeds)
             image_prompt_embeds = clip_image_prompt_embeds
             uncond_image_prompt_embeds = uncond_clip_image_prompt_embeds
-        
-        if face_emb is not None:
-            face_embeds = face_emb.to(self.device, dtype=torch.float16)
-            face_prompt_embeds = self.face_proj_model(face_embeds)
-            uncond_face_prompt_embeds = self.face_proj_model(torch.zeros_like(face_embeds))
-            if image_prompt_embeds is None:
-                image_prompt_embeds = face_prompt_embeds
-                uncond_image_prompt_embeds = uncond_face_prompt_embeds
-            else:
-                image_prompt_embeds = torch.cat([face_prompt_embeds, image_prompt_embeds], axis=1)
-                uncond_image_prompt_embeds = torch.cat([uncond_face_prompt_embeds, uncond_image_prompt_embeds], dim=1)       
 
         return image_prompt_embeds, uncond_image_prompt_embeds
     
     @torch.inference_mode()
-    def get_image_embeds_lyrasd(self, image=None, ip_image_embeds=None, face_emb=None, batch_size = 1, ip_scale=1.0, fp_scale=1.0, do_classifier_free_guidance=True):
+    def get_image_embeds_lyrasd(self, image=None, ip_image_embeds=None, batch_size = 1, ip_scale=1.0, do_classifier_free_guidance=True):
         dict_tensor = {}
 
         if self.ip_ckpt and ip_scale>0:
@@ -199,91 +171,4 @@ class LyraIPAdapter:
                     clip_image_embeds = torch.cat([uncond_clip_image_embeds, clip_image_embeds])
                 ip_image_embeds = self.image_proj_model(clip_image_embeds)
                 dict_tensor["ip_hidden_states"] = ip_image_embeds
-
-        if face_emb is not None and self.fp_ckpt and ip_scale>0:
-            face_embeds = face_emb.to(self.device, dtype=torch.float16)
-            face_prompt_embeds = self.face_proj_model(face_embeds)
-            uncond_face_prompt_embeds = self.face_proj_model(torch.zeros_like(face_embeds))
-            if do_classifier_free_guidance:
-                fp_image_embeds = torch.cat([uncond_face_prompt_embeds, face_prompt_embeds])
-            else:
-                fp_image_embeds = face_prompt_embeds
-            dict_tensor["fp_hidden_states"] = fp_image_embeds
         return dict_tensor
-
-
-if __name__ == "__main__":
-    sys.path.append("/data/home/kiokaxiao/repos/LyraSD/python/lyrasd")
-    from lyrasd_model import LyraSdXLTxt2ImgPipeline
-
-    model_path = "/data/SharedModels/SD/checkpoints/stable-diffusion-xl-base-1.0/"
-    # model_path = "/cfs-datasets/projects/VirtualIdol/models/base_model/sdxl/xxmix9realisticsdxlV1"
-    lib_path = os.environ.get("LIBLYRASD_SO")
-
-    dir_ip_adapter = "/cfs-datasets/projects/VirtualIdol/models/ip_adapter/sdxl_models/ip-adapter-plus_sdxl_vit-h.bin"
-    dir_facein = "/cfs-datasets/projects/VirtualIdol/models/FaceIn/v1/FaceIn_sdxl.bin"
-    image_encoder_path = "/cfs-datasets/projects/VirtualIdol/models/ip_adapter/models/image_encoder"
-
-    pipeline = LyraSdXLTxt2ImgPipeline(model_path, lib_path)
-    pipeline.load_ip_adapter(dir_ip_adapter, True, image_encoder_path, 16,1024, dir_facein, 1, 512)
-    # pipeline.load_ip_adapter(dir_ip_adapter, True, image_encoder_path, 16,1024, "", 1, 512)
-
-    face_emb = np.load("/data/home/kiokaxiao/repos/VidolImageDraw/girl.npy")
-    face_emb = torch.Tensor(face_emb.reshape([1,-1]))
-    ip_image = Image.open("/data/home/kiokaxiao/repos/VidolImageDraw/images/input_image.png").convert('RGB')
-
-    generator = torch.Generator("cuda").manual_seed(123)
-    batches = [2]
-    sizes = [[512, 512], [768, 768], [1024, 1024]]
-    # sizes = [[832, 640]]
-    # sizes = [[1024, 1024]]
-    running_cnt = 1
-    do_bench = False
-
-    ip_ratio = 1
-    facein_ratio = 0.6
-    extra_tensor_dict = {}
-    extra_tensor_dict = pipeline.ip_adapter_helper.get_image_embeds_lyrasd(ip_image, None, face_emb, batches[0], ip_ratio, facein_ratio)
-    param_scale_dict = {"facein_ratio": facein_ratio, "ip_ratio": ip_ratio}
-    draw_cfg = {'width': 640,
-                'num_inference_steps': 30, 
-                'height': 832, 
-                'negative_prompt': '(worst quality, low quality, 3d, 2d, cartoons, sketch), tooth, open mouth', 
-                'guidance_scale': 7, 
-                'prompt': 'xxmixgirl, masterpiece, best quality, 1girl, solo, looking at viewer, simple background, hair ornament, black eyes, portrait',
-                'output_type': 'pil',
-                'extra_tensor_dict': extra_tensor_dict,
-                "param_scale_dict": param_scale_dict}
-
-
-    def warmup(draw_cfg):
-        draw_cfg_wm = deepcopy(draw_cfg)
-        draw_cfg_wm['num_inference_steps'] = 1
-        pipeline(**draw_cfg_wm, generator= generator)
-
-    if not do_bench:
-        images = pipeline(**draw_cfg, generator= generator)
-    else:    
-        for batch in batches:
-            for height, width in sizes:
-                draw_cfg['width'] = width
-                draw_cfg['height'] = height
-                draw_cfg['num_images_per_prompt'] = batch
-                draw_cfg["num_inference_steps"] = 20
-                warmup(draw_cfg)
-                time_uses = []
-                for x in range(running_cnt):
-                    start = time.perf_counter()
-                    draw_cfg['num_images_per_prompt'] = batch
-                    generator = torch.Generator("cuda").manual_seed(123)
-                    print("draw_cfg: ", draw_cfg.keys())
-                    print("draw_cfg: ", draw_cfg)
-
-                    images = pipeline(**draw_cfg, generator= generator)
-                    time_use = time.perf_counter() - start
-                    time_uses.append(time_use)
-                print("bench", batch, width, sum(time_uses)/running_cnt, get_mem_use())
-
-    print(type(images))
-    images[0].save("t.png")
-