add faceswap

CHANGELOG.md

@@ -2,8 +2,8 @@
 
 ## Update for 2023-01-15
 
-Another release with a lot more functionality in new Control module and FaceID & IPAdapter modules  
-Plus welcome additions to UI performance and accessibility and flexibility of deployment  
+Another release with a lot more functionality in the **Control** module and **FaceID/FaceSwap** & **PAdapter** modules  
+Plus welcome additions to **UI performance, usability and accessibility** and flexibility of deployment  
 And it also includes fixes for all reported issues so far  
 
 - **Control**:
@@ -48,11 +48,14 @@ And it also includes fixes for all reported issues so far
   - fix batch/folder/video modes  
   - fix processor switching within same unit  
   - fix pipeline switching between different modes  
-- [FaceID](https://huggingface.co/h94/IP-Adapter-FaceID)  
+- [FaceID/FaceSwap](https://huggingface.co/h94/IP-Adapter-FaceID)  
   - full implementation for *SD15* and *SD-XL*, to use simply select from *Scripts*  
     **Base** (93MB) uses *InsightFace* to generate face embeds and *OpenCLIP-ViT-H-14* (2.5GB) as image encoder  
     **SXDL** (1022MB) uses *InsightFace* to generate face embeds and *OpenCLIP-ViT-bigG-14* (3.7GB) as image encoder  
     **Plus** (150MB) uses *InsightFace* to generate face embeds and *CLIP-ViT-H-14-laion2B* (3.8GB) as image encoder  
+  - **FaceSwap**  
+    you can use just faceid or just faceswap or both at the same time  
+    faceid guides image generation given the input image while face swap performs face swapping at the end of generation  
   - *note*: all models are downloaded on first use  
   - enable use via api, thanks @trojaner
 - [IPAdapter](https://huggingface.co/h94/IP-Adapter)  

modules/processing.py

@@ -863,7 +863,6 @@ def process_images_inner(p: StableDiffusionProcessing) -> Processed:
     if p.scripts is not None and isinstance(p.scripts, modules.scripts.ScriptRunner):
         p.scripts.process(p)
 
-
     def get_conds_with_caching(function, required_prompts, steps, cache):
         if cache[0] is not None and (required_prompts, steps) == cache[0]:
             return cache[1]

scripts/faceid.py

@@ -5,6 +5,7 @@ import numpy as np
 import gradio as gr
 import diffusers
 import huggingface_hub as hf
+from PIL import Image
 from modules import scripts, processing, shared, devices
 
 
@@ -19,6 +20,7 @@ ip_model = None
 ip_model_name = None
 ip_model_tokens = None
 ip_model_rank = None
+swapper = None
 
 
 def dependencies():
@@ -32,6 +34,135 @@ def dependencies():
             install(pkg[0], pkg[1], ignore=True)
 
 
+def face_id(p: processing.StableDiffusionProcessing, faces, image, model, override, tokens, rank, cache, scale, structure):
+    global ip_model, ip_model_name, ip_model_tokens, ip_model_rank # pylint: disable=global-statement
+    from insightface.utils import face_align
+    from ip_adapter.ip_adapter_faceid import IPAdapterFaceID, IPAdapterFaceIDPlus, IPAdapterFaceIDXL
+
+    face_embeds = torch.from_numpy(faces[0].normed_embedding).unsqueeze(0)
+    face_image = face_align.norm_crop(image, landmark=faces[0].kps, image_size=224) # you can also segment the face
+
+    ip_ckpt = MODELS[model]
+    folder, filename = os.path.split(ip_ckpt)
+    basename, _ext = os.path.splitext(filename)
+    model_path = hf.hf_hub_download(repo_id=folder, filename=filename, cache_dir=shared.opts.diffusers_dir)
+    if model_path is None:
+        shared.log.error(f'FaceID download failed: model={model} file={ip_ckpt}')
+        return None
+
+    processing.process_init(p)
+    if override:
+        shared.sd_model.scheduler = diffusers.DDIMScheduler(
+            num_train_timesteps=1000,
+            beta_start=0.00085,
+            beta_end=0.012,
+            beta_schedule="scaled_linear",
+            clip_sample=False,
+            set_alpha_to_one=False,
+            steps_offset=1,
+        )
+    shortcut = None
+    if ip_model is None or ip_model_name != model or ip_model_tokens != tokens or ip_model_rank != rank or not cache:
+        shared.log.debug(f'FaceID load: model={model} file={ip_ckpt} tokens={tokens} rank={rank}')
+        if 'Plus' in model:
+            image_encoder_path = "laion/CLIP-ViT-H-14-laion2B-s32B-b79K"
+            ip_model = IPAdapterFaceIDPlus(
+                sd_pipe=shared.sd_model,
+                image_encoder_path=image_encoder_path,
+                ip_ckpt=model_path,
+                lora_rank=rank,
+                num_tokens=tokens,
+                device=devices.device,
+                torch_dtype=devices.dtype,
+            )
+            shortcut = 'v2' in model
+        elif 'XL' in model:
+            ip_model = IPAdapterFaceIDXL(
+                sd_pipe=shared.sd_model,
+                ip_ckpt=model_path,
+                lora_rank=rank,
+                num_tokens=tokens,
+                device=devices.device,
+                torch_dtype=devices.dtype,
+            )
+        else:
+            ip_model = IPAdapterFaceID(
+                sd_pipe=shared.sd_model,
+                ip_ckpt=model_path,
+                lora_rank=rank,
+                num_tokens=tokens,
+                device=devices.device,
+                torch_dtype=devices.dtype,
+            )
+        ip_model_name = model
+        ip_model_tokens = tokens
+        ip_model_rank = rank
+    else:
+        shared.log.debug(f'FaceID cached: model={model} file={ip_ckpt} tokens={tokens} rank={rank}')
+
+    # main generate dict
+    ip_model_dict = {
+        'prompt': p.all_prompts[0],
+        'negative_prompt': p.all_negative_prompts[0],
+        'num_samples': p.batch_size,
+        'width': p.width,
+        'height': p.height,
+        'num_inference_steps': p.steps,
+        'scale': scale,
+        'guidance_scale': p.cfg_scale,
+        'seed': int(p.all_seeds[0]),
+        'faceid_embeds': face_embeds.shape,
+    }
+
+    # optional generate dict
+    if shortcut is not None:
+        ip_model_dict['shortcut'] = shortcut
+    if 'Plus' in model:
+        ip_model_dict['s_scale'] = structure
+        ip_model_dict['face_image'] = face_image.shape
+    shared.log.debug(f'FaceID args: {ip_model_dict}')
+    if 'Plus' in model:
+        ip_model_dict['face_image'] = face_image
+    ip_model_dict['faceid_embeds'] = face_embeds
+
+    # run generate
+    images = []
+    ip_model.set_scale(scale)
+    for _i in range(p.n_iter):
+        res = ip_model.generate(**ip_model_dict)
+        if isinstance(res, list):
+            images += res
+    ip_model.set_scale(0)
+
+    if not cache:
+        ip_model = None
+        ip_model_name = None
+    devices.torch_gc()
+
+    p.extra_generation_params["IP Adapter"] = f'{basename}:{scale}'
+    return images
+
+
+def face_swap(p: processing.StableDiffusionProcessing, image, source_face):
+    import insightface.model_zoo
+    global swapper # pylint: disable=global-statement
+    if swapper is None:
+        model_path = hf.hf_hub_download(repo_id='ezioruan/inswapper_128.onnx', filename='inswapper_128.onnx', cache_dir=shared.opts.diffusers_dir)
+        router = insightface.model_zoo.model_zoo.ModelRouter(model_path)
+        swapper = router.get_model()
+
+    np_image = cv2.cvtColor(np.array(image), cv2.COLOR_RGB2BGR)
+    faces = app.get(np_image)
+
+    res = np_image.copy()
+    for target_face in faces:
+        res = swapper.get(res, target_face, source_face, paste_back=True) # pylint: disable=too-many-function-args, unexpected-keyword-arg
+
+    p.extra_generation_params["FaceSwap"] = f'{len(faces)}'
+    np_image = cv2.cvtColor(res, cv2.COLOR_BGR2RGB)
+    return Image.fromarray(np_image)
+
+
 class Script(scripts.Script):
     def title(self):
         return 'FaceID'
@@ -42,7 +173,8 @@ class Script(scripts.Script):
     # return signature is array of gradio components
     def ui(self, _is_img2img):
         with gr.Row():
-            model = gr.Dropdown(choices=list(MODELS), label='Model', value='FaceID Base')
+            mode = gr.CheckboxGroup(label='Mode', choices=['FaceID', 'FaceSwap'], value=['FaceID'])
+            model = gr.Dropdown(choices=list(MODELS), label='FaceID Model', value='FaceID Base')
         with gr.Row(visible=True):
             override = gr.Checkbox(label='Override sampler', value=True)
             cache = gr.Checkbox(label='Cache model', value=True)
@@ -54,15 +186,15 @@ class Script(scripts.Script):
             tokens = gr.Slider(label='Tokens', minimum=1, maximum=16, step=1, value=4)
         with gr.Row():
             image = gr.Image(image_mode='RGB', label='Image', source='upload', type='pil', width=512)
-        return [model, scale, image, override, rank, tokens, structure, cache]
+        return [mode, model, scale, image, override, rank, tokens, structure, cache]
 
-    def run(self, p: processing.StableDiffusionProcessing, model, scale, image, override, rank, tokens, structure, cache): # pylint: disable=arguments-differ, unused-argument
+    def run(self, p: processing.StableDiffusionProcessing, mode, model, scale, image, override, rank, tokens, structure, cache): # pylint: disable=arguments-differ, unused-argument
+        if len(mode) == 0:
+            return None
         dependencies()
         try:
             import onnxruntime
             from insightface.app import FaceAnalysis
-            from insightface.utils import face_align
-            from ip_adapter.ip_adapter_faceid import IPAdapterFaceID, IPAdapterFaceIDPlus, IPAdapterFaceIDXL
         except Exception as e:
             shared.log.error(f'FaceID: {e}')
             return None
@@ -73,7 +205,7 @@ class Script(scripts.Script):
             shared.log.error('FaceID: base model not supported')
             return None
 
-        global app, ip_model, ip_model_name, ip_model_tokens, ip_model_rank # pylint: disable=global-statement
+        global app # pylint: disable=global-statement
         if app is None:
             shared.log.debug(f"ONNX: device={onnxruntime.get_device()} providers={onnxruntime.get_available_providers()}")
             app = FaceAnalysis(name="buffalo_l", providers=['CUDAExecutionProvider', 'CPUExecutionProvider'])
@@ -84,127 +216,36 @@ class Script(scripts.Script):
             from modules.api.api import decode_base64_to_image
             image = decode_base64_to_image(image)
 
-        image = cv2.cvtColor(np.array(image), cv2.COLOR_RGB2BGR)
-        faces = app.get(image)
+        np_image = cv2.cvtColor(np.array(image), cv2.COLOR_RGB2BGR)
+        faces = app.get(np_image)
         if len(faces) == 0:
             shared.log.error('FaceID: no faces found')
             return None
-        for face in faces:
-            shared.log.debug(f'FaceID face: score={face.det_score:.2f} gender={"female" if face.gender==0 else "male"} age={face.age} bbox={face.bbox}')
-        face_embeds = torch.from_numpy(faces[0].normed_embedding).unsqueeze(0)
-        face_image = face_align.norm_crop(image, landmark=faces[0].kps, image_size=224) # you can also segment the face
-
-        ip_ckpt = MODELS[model]
-        folder, filename = os.path.split(ip_ckpt)
-        basename, _ext = os.path.splitext(filename)
-        model_path = hf.hf_hub_download(repo_id=folder, filename=filename, cache_dir=shared.opts.diffusers_dir)
-        if model_path is None:
-            shared.log.error(f'FaceID download failed: model={model} file={ip_ckpt}')
-            return None
-
-        processing.process_init(p)
-        if override:
-            shared.sd_model.scheduler = diffusers.DDIMScheduler(
-                num_train_timesteps=1000,
-                beta_start=0.00085,
-                beta_end=0.012,
-                beta_schedule="scaled_linear",
-                clip_sample=False,
-                set_alpha_to_one=False,
-                steps_offset=1,
-            )
-        shortcut = None
-        if ip_model is None or ip_model_name != model or ip_model_tokens != tokens or ip_model_rank != rank or not cache:
-            shared.log.debug(f'FaceID load: model={model} file={ip_ckpt} tokens={tokens} rank={rank}')
-            if 'Plus' in model:
-                image_encoder_path = "laion/CLIP-ViT-H-14-laion2B-s32B-b79K"
-                ip_model = IPAdapterFaceIDPlus(
-                    sd_pipe=shared.sd_model,
-                    image_encoder_path=image_encoder_path,
-                    ip_ckpt=model_path,
-                    lora_rank=rank,
-                    num_tokens=tokens,
-                    device=devices.device,
-                    torch_dtype=devices.dtype,
-                )
-                shortcut = 'v2' in model
-            elif 'XL' in model:
-                ip_model = IPAdapterFaceIDXL(
-                    sd_pipe=shared.sd_model,
-                    ip_ckpt=model_path,
-                    lora_rank=rank,
-                    num_tokens=tokens,
-                    device=devices.device,
-                    torch_dtype=devices.dtype,
-                )
-            else:
-                ip_model = IPAdapterFaceID(
-                    sd_pipe=shared.sd_model,
-                    ip_ckpt=model_path,
-                    lora_rank=rank,
-                    num_tokens=tokens,
-                    device=devices.device,
-                    torch_dtype=devices.dtype,
-                )
-            ip_model_name = model
-            ip_model_tokens = tokens
-            ip_model_rank = rank
-        else:
-            shared.log.debug(f'FaceID cached: model={model} file={ip_ckpt} tokens={tokens} rank={rank}')
-
-        # main generate dict
-        ip_model_dict = {
-            'prompt': p.all_prompts[0],
-            'negative_prompt': p.all_negative_prompts[0],
-            'num_samples': p.batch_size,
-            'width': p.width,
-            'height': p.height,
-            'num_inference_steps': p.steps,
-            'scale': scale,
-            'guidance_scale': p.cfg_scale,
-            'seed': int(p.all_seeds[0]),
-            'faceid_embeds': face_embeds.shape,
-        }
-
-        # optional generate dict
-        if shortcut is not None:
-            ip_model_dict['shortcut'] = shortcut
-        if 'Plus' in model:
-            ip_model_dict['s_scale'] = structure
-            ip_model_dict['face_image'] = face_image.shape
-        shared.log.debug(f'FaceID args: {ip_model_dict}')
-        if 'Plus' in model:
-            ip_model_dict['face_image'] = face_image
-        ip_model_dict['faceid_embeds'] = face_embeds
-
-        # run generate
-        images = []
-
-        ip_model.set_scale(scale)
-        for _i in range(p.n_iter):
-            res = ip_model.generate(**ip_model_dict)
-            if isinstance(res, list):
-                images += res
-        ip_model.set_scale(0)
-
-        if not cache:
-            ip_model = None
-            ip_model_name = None
-        devices.torch_gc()
-
-        p.extra_generation_params["IP Adapter"] = f'{basename}:{scale}'
         for i, face in enumerate(faces):
+            shared.log.debug(f'FaceID face: i={i} score={face.det_score:.2f} gender={"female" if face.gender==0 else "male"} age={face.age} bbox={face.bbox}')
             p.extra_generation_params[f"FaceID {i} score"] = f'{face.det_score:.2f}'
             p.extra_generation_params[f"FaceID {i} gender"] = "female" if face.gender==0 else "male"
             p.extra_generation_params[f"FaceID {i} age"] = face.age
 
-        processed = processing.Processed(
-            p,
-            images_list=images,
-            seed=p.seed,
-            subseed=p.subseed,
-            index_of_first_image=0,
-        )
+        images = []
+        if 'FaceID' in mode:
+            images = face_id(p, faces, np_image, model, override, tokens, rank, cache, scale, structure) # run faceid pipeline
+            processed = processing.Processed(
+                p,
+                images_list=images,
+                seed=p.seed,
+                subseed=p.subseed,
+                index_of_first_image=0,
+            )
+        else:
+            processed = processing.process_images(p) # run normal pipeline
+            images = processed.images
+        if 'FaceSwap' in mode: # replace faces as postprocess
+            processed.images = []
+            for batch_image in images:
+                swapped_image = face_swap(p, batch_image, source_face=faces[0])
+                processed.images.append(swapped_image)
+
         processed.info = processed.infotext(p, 0)
         processed.infotexts = [processed.info]
         return processed