Update app.py

app.py

@@ -1,15 +1,18 @@
 import torch
 import spaces
 from diffusers import StableDiffusionPipeline, DDIMScheduler, AutoencoderKL
-from ip_adapter.ip_adapter_faceid import IPAdapterFaceID
+from ip_adapter.ip_adapter_faceid import IPAdapterFaceID, IPAdapterFaceIDPlus
 from huggingface_hub import hf_hub_download
 from insightface.app import FaceAnalysis
+from insightface.utils import face_align
 import gradio as gr
 import cv2
 
 base_model_path = "SG161222/Realistic_Vision_V4.0_noVAE"
 vae_model_path = "stabilityai/sd-vae-ft-mse"
-ip_ckpt = hf_hub_download(repo_id='h94/IP-Adapter-FaceID', filename="ip-adapter-faceid_sd15.bin", repo_type="model")
+image_encoder_path = "laion/CLIP-ViT-H-14-laion2B-s32B-b79K"
+ip_ckpt = hf_hub_download(repo_id="h94/IP-Adapter-FaceID", filename="ip-adapter-faceid_sd15.bin", repo_type="model")
+ip_plus_ckpt = hf_hub_download(repo_id="h94/IP-Adapter-FaceID", filename="ip-adapter-faceid-plusv2_sd15.bin", repo_type="model")
 
 device = "cuda"
 
@@ -31,25 +34,37 @@ pipe = StableDiffusionPipeline.from_pretrained(
 )
 
 ip_model = IPAdapterFaceID(pipe, ip_ckpt, device)
+ip_model_plus = IPAdapterFaceIDPlus(pipe, image_encoder_path, ip_plus_ckpt, device)
 
 @spaces.GPU(enable_queue=True)
-def generate_image(images, prompt, negative_prompt, progress=gr.Progress(track_tqdm=True)):
+def generate_image(images, prompt, negative_prompt, preserve_face_structure, progress=gr.Progress(track_tqdm=True)):
     pipe.to(device)
     app = FaceAnalysis(name="buffalo_l", providers=['CUDAExecutionProvider', 'CPUExecutionProvider'])
     app.prepare(ctx_id=0, det_size=(640, 640))
-
+    
     faceid_all_embeds = []
+    first_iteration = True
     for image in images:
         face = cv2.imread(image)
         faces = app.get(face)
         faceid_embed = torch.from_numpy(faces[0].normed_embedding).unsqueeze(0)
         faceid_all_embeds.append(faceid_embed)
-    
+        if(first_iteration):
+            face_image = face_align.norm_crop(face, landmark=faces[0].kps, image_size=224) # you can also segment the face
+            first_iteration = False
+            
     average_embedding = torch.mean(torch.stack(faceid_all_embeds, dim=0), dim=0)
     
-    image = ip_model.generate(
-        prompt=prompt, negative_prompt=negative_prompt, faceid_embeds=average_embedding, width=512, height=512, num_inference_steps=30
-    )
+    if(not preserve_face_structure):
+        image = ip_model.generate(
+            prompt=prompt, negative_prompt=negative_prompt, faceid_embeds=average_embedding,
+            width=512, height=512, num_inference_steps=30
+        )
+    else:
+        image = ip_model_plus.generate(
+            prompt=prompt, negative_prompt=negative_prompt, faceid_embeds=average_embedding,
+            face_image=face_image, shortcut=True, s_scale=1.5, width=512, height=512, num_inference_steps=30
+        )
     print(image)
     return image
 css = '''
@@ -66,7 +81,8 @@ demo = gr.Interface(
             gr.Textbox(label="Prompt",
                        info="Try something like 'a photo of a man/woman/person'",
                        placeholder="A photo of a [man/woman/person]..."),
-            gr.Textbox(label="Negative Prompt", placeholder="low quality")
+            gr.Textbox(label="Negative Prompt", placeholder="low quality"),
+            gr.Checkbox(label="Preserve Face Structure", value=False),
         ],
         outputs=[gr.Gallery(label="Generated Image")],
         title="IP-Adapter-FaceID demo",