major change

.gitignore

@@ -0,0 +1,200 @@
+# Created by https://www.toptal.com/developers/gitignore/api/visualstudiocode,python
+# Edit at https://www.toptal.com/developers/gitignore?templates=visualstudiocode,python
+
+### Python ###
+# Byte-compiled / optimized / DLL files
+__pycache__/
+*.py[cod]
+*$py.class
+
+# C extensions
+*.so
+
+models/*.onnx
+
+.vscode
+.mypy_cache
+
+# Distribution / packaging
+.Python
+build/
+develop-eggs/
+dist/
+downloads/
+eggs/
+.eggs/
+lib/
+lib64/
+parts/
+sdist/
+var/
+wheels/
+share/python-wheels/
+*.egg-info/
+.installed.cfg
+*.egg
+MANIFEST
+
+# PyInstaller
+#  Usually these files are written by a python script from a template
+#  before PyInstaller builds the exe, so as to inject date/other infos into it.
+*.manifest
+*.spec
+
+# Installer logs
+pip-log.txt
+pip-delete-this-directory.txt
+
+# Unit test / coverage reports
+htmlcov/
+.tox/
+.nox/
+.coverage
+.coverage.*
+.cache
+nosetests.xml
+coverage.xml
+*.cover
+*.py,cover
+.hypothesis/
+.pytest_cache/
+cover/
+
+# Translations
+*.mo
+*.pot
+
+# Django stuff:
+*.log
+local_settings.py
+db.sqlite3
+db.sqlite3-journal
+
+# Flask stuff:
+instance/
+.webassets-cache
+
+# Scrapy stuff:
+.scrapy
+
+# Sphinx documentation
+docs/_build/
+
+# PyBuilder
+.pybuilder/
+target/
+
+# Jupyter Notebook
+.ipynb_checkpoints
+
+# IPython
+profile_default/
+ipython_config.py
+
+# pyenv
+#   For a library or package, you might want to ignore these files since the code is
+#   intended to run in multiple environments; otherwise, check them in:
+# .python-version
+
+# pipenv
+#   According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control.
+#   However, in case of collaboration, if having platform-specific dependencies or dependencies
+#   having no cross-platform support, pipenv may install dependencies that don't work, or not
+#   install all needed dependencies.
+#Pipfile.lock
+
+# poetry
+#   Similar to Pipfile.lock, it is generally recommended to include poetry.lock in version control.
+#   This is especially recommended for binary packages to ensure reproducibility, and is more
+#   commonly ignored for libraries.
+#   https://python-poetry.org/docs/basic-usage/#commit-your-poetrylock-file-to-version-control
+#poetry.lock
+
+# pdm
+#   Similar to Pipfile.lock, it is generally recommended to include pdm.lock in version control.
+#pdm.lock
+#   pdm stores project-wide configurations in .pdm.toml, but it is recommended to not include it
+#   in version control.
+#   https://pdm.fming.dev/#use-with-ide
+.pdm.toml
+
+# PEP 582; used by e.g. github.com/David-OConnor/pyflow and github.com/pdm-project/pdm
+__pypackages__/
+
+# Celery stuff
+celerybeat-schedule
+celerybeat.pid
+
+# SageMath parsed files
+*.sage.py
+
+# Environments
+.env
+.venv
+env/
+venv/
+ENV/
+env.bak/
+venv.bak/
+
+# Spyder project settings
+.spyderproject
+.spyproject
+
+# Rope project settings
+.ropeproject
+
+# mkdocs documentation
+/site
+
+# mypy
+.mypy_cache/
+.dmypy.json
+dmypy.json
+
+# Pyre type checker
+.pyre/
+
+# pytype static type analyzer
+.pytype/
+
+# Cython debug symbols
+cython_debug/
+
+# PyCharm
+#  JetBrains specific template is maintained in a separate JetBrains.gitignore that can
+#  be found at https://github.com/github/gitignore/blob/main/Global/JetBrains.gitignore
+#  and can be added to the global gitignore or merged into this file.  For a more nuclear
+#  option (not recommended) you can uncomment the following to ignore the entire idea folder.
+#.idea/
+
+### Python Patch ###
+# Poetry local configuration file - https://python-poetry.org/docs/configuration/#local-configuration
+poetry.toml
+
+# ruff
+.ruff_cache/
+
+# LSP config files
+pyrightconfig.json
+
+### VisualStudioCode ###
+.vscode/*
+!.vscode/settings.json
+!.vscode/tasks.json
+!.vscode/launch.json
+!.vscode/extensions.json
+!.vscode/*.code-snippets
+
+# Local History for Visual Studio Code
+.history/
+
+# Built Visual Studio Code Extensions
+*.vsix
+
+### VisualStudioCode Patch ###
+# Ignore all local history of files
+.history
+.ionide
+
+# End of https://www.toptal.com/developers/gitignore/api/visualstudiocode,python

CHANGELOG.md

@@ -0,0 +1,19 @@
+## 0.1.0 :
+
+### Major :
++ add multiple face support
++ add face blending support (will blend sources faces)
++ add face similarity evaluation (will compare face to a reference)
+    + add filters to discard images that are not rated similar enough to reference image and source images
++ add face tools tab
+    + face extraction tool
+    + face builder tool : will build a face model that can be reused
++ add faces models
+
+### Minor :
+
+Improve performance by not reprocessing source face each time
+
+### Breaking changes
+
+base64 and api not supported anymore (will be reintroduced in the future)

install.py

@@ -8,6 +8,7 @@ import urllib.request
 req_file = os.path.join(os.path.dirname(os.path.realpath(__file__)), "requirements.txt")
 
 models_dir = os.path.abspath("models/roop")
+faces_dir = os.path.abspath("models/roop/faces")
 model_url = "https://huggingface.co/henryruhs/roop/resolve/main/inswapper_128.onnx"
 model_name = os.path.basename(model_url)
 model_path = os.path.join(models_dir, model_name)
@@ -21,9 +22,13 @@ def download(url, path):
 if not os.path.exists(models_dir):
     os.makedirs(models_dir)
 
+if not os.path.exists(faces_dir):
+    os.makedirs(faces_dir)
+
 if not os.path.exists(model_path):
     download(model_url, model_path)
 
+
 print("Checking roop requirements")
 with open(req_file) as file:
     for package in file:

javascript/faceswap.js

@@ -0,0 +1,4 @@
+window.onbeforeunload = function() {
+    // Prevent the stable diffusion window from being closed by mistake
+    return "Are you sure ?";
+};

main.py

@@ -0,0 +1,5 @@
+from scripts.faceswap import FaceSwapScript
+
+
+f = FaceSwapScript()
+print(f.ui(True))

models/Put_the_model_here.txt

@@ -0,0 +1 @@
+The model file required is "inswapper_128.onnx".Mirrors are given the roop project [installation guide](https://github.com/s0md3v/roop/wiki/1.-Installation).

requirements.txt

@@ -2,5 +2,7 @@ insightface==0.7.3
 onnx==1.14.0
 onnxruntime==1.15.0
 opencv-python==4.7.0.72
+dill==0.3.6
+pandas
 ifnude
-cython
+cython

scripts/cimage.py

@@ -6,4 +6,4 @@ def convert_to_sd(img):
     chunks = detect(img)
     for chunk in chunks:
         shapes.append(chunk["score"] > 0.7)
-    return [any(shapes), tempfile.NamedTemporaryFile(delete=False, suffix=".png")]
+    return any(shapes)

scripts/faceswap.py

@@ -1,84 +1,370 @@
+import glob
+import importlib
+import json
 import os
+from dataclasses import dataclass, fields
+from pprint import pformat, pprint
+from typing import Dict, List, Set, Tuple, Union
+from scripts.cimage import convert_to_sd
+import cv2
+import dill as pickle
 import gradio as gr
 import modules.scripts as scripts
-from modules.upscaler import Upscaler, UpscalerData
-from modules import scripts, shared, images, scripts_postprocessing
-from modules.processing import (
-    StableDiffusionProcessing,
-    StableDiffusionProcessingImg2Img,
-)
-from modules.shared import cmd_opts, opts, state
-from PIL import Image
-import glob
+import numpy as np
+import onnx
+import pandas as pd
+import torch
+from insightface.app.common import Face
+from modules import script_callbacks, scripts, shared
 from modules.face_restoration import FaceRestoration
+from modules.images import save_image
+from modules.processing import (Processed, StableDiffusionProcessing,
+                                StableDiffusionProcessingImg2Img,
+                                StableDiffusionProcessingTxt2Img)
+from modules.shared import cmd_opts, opts, state
+from modules.upscaler import Upscaler, UpscalerData
+from onnx import numpy_helper
+from PIL import Image
 
+import scripts.swapper as swapper
 from scripts.roop_logging import logger
-from scripts.swapper import UpscaleOptions, swap_face, ImageResult
 from scripts.roop_version import version_flag
-import os
+from scripts.imgutils import (create_square_image, cv2_to_pil, pil_to_cv2,
+                              pil_to_torch, torch_to_pil)
+from scripts.upscaling import UpscaleOptions, upscale_image
 
+EXTENSION_PATH=os.path.join("extensions","sd-webui-roop")
 
 def get_models():
-    models_path = os.path.join(scripts.basedir(), "models" + os.path.sep + "roop" + os.path.sep + "*")
+    models_path = os.path.join(
+        scripts.basedir(), EXTENSION_PATH,"models","*"
+    )
     models = glob.glob(models_path)
+    models_path = os.path.join(scripts.basedir(), "models", "roop", "*")
+    models += glob.glob(models_path)
     models = [x for x in models if x.endswith(".onnx") or x.endswith(".pth")]
     return models
 
+def get_faces():
+    faces_path = os.path.join(scripts.basedir(), "models", "roop", "faces","*.pkl")
+    faces = glob.glob(faces_path)
+    return ["None"] + faces
+
+@dataclass
+class FaceSwapUnitSettings:
+    source_img: Image.Image
+    source_face : str
+    _batch_files: gr.components.File
+    blend_faces: bool
+    enable: bool
+    same_gender: bool
+    min_sim: float
+    min_ref_sim: float
+    _faces_index: int
+    swap_in_source: bool
+    swap_in_generated: bool
+
+    @staticmethod
+    def get_unit_configuration(unit: int, components):
+        fields_count = len(fields(FaceSwapUnitSettings))
+        return FaceSwapUnitSettings(
+            *components[unit * fields_count : unit * fields_count + fields_count]
+        )
+
+    @property
+    def faces_index(self):
+        faces_index = {
+            int(x) for x in self._faces_index.strip(",").split(",") if x.isnumeric()
+        }
+        if len(faces_index) == 0:
+            return {0}
+
+        return faces_index
+
+    @property
+    def batch_files(self):
+        return self._batch_files or []
+    
+    @property
+    def reference_face(self) :
+        if not hasattr(self,"_reference_face") :
+            if self.source_face and self.source_face != "None" :
+                with open(self.source_face, "rb") as file:
+                    logger.info(f"loading pickle {file.name}")
+                    face = Face(pickle.load(file))
+                    self._reference_face = face
+            elif self.source_img is not None :
+                source_img = pil_to_cv2(self.source_img)
+                self._reference_face =  swapper.get_or_default(swapper.get_faces(source_img), 0, None)    
+            else :
+                logger.info("You need at least one face")
+                self._reference_face = None
+                
+        return self._reference_face
+    
+    @property
+    def faces(self) :
+        if self.batch_files is not None and not hasattr(self,"_faces") :
+            self._faces = [self.reference_face] if self.reference_face is not None else []
+            for file in self.batch_files :
+                img = Image.open(file.name)
+                face = swapper.get_or_default(swapper.get_faces(pil_to_cv2(img)), 0, None)
+                if face is not None :
+                    self._faces.append(face)
+        return self._faces
+
+    @property
+    def blended_faces(self):
+        if not hasattr(self,"_blended_faces") :
+            self._blended_faces = swapper.blend_faces(self.faces)
+        return self._blended_faces
+
+
+def compare(img1, img2):
+    if img1 is not None and img2 is not None:
+        return swapper.compare_faces(img1, img2)
+
+    return "You need 2 images to compare"
+
+import tempfile
+def extract_faces(files, extract_path,  face_restorer_name, face_restorer_visibility,upscaler_name,upscaler_scale, upscaler_visibility):
+    if not extract_path :
+        tempfile.mkdtemp()
+    if files is not None:
+        images = []
+        for file in files :
+            img = Image.open(file.name).convert("RGB")
+            faces = swapper.get_faces(pil_to_cv2(img))
+            if faces:
+                face_images = []
+                for face in faces:
+                    bbox = face.bbox.astype(int)
+                    x_min, y_min, x_max, y_max = bbox
+                    face_image = img.crop((x_min, y_min, x_max, y_max))
+                    if face_restorer_name or face_restorer_visibility:
+                        scale = 1 if face_image.width > 512 else 512//face_image.width
+                        face_image = upscale_image(face_image, UpscaleOptions(face_restorer_name=face_restorer_name, restorer_visibility=face_restorer_visibility, upscaler_name=upscaler_name, upscale_visibility=upscaler_visibility, scale=scale))
+                    path = tempfile.NamedTemporaryFile(delete=False,suffix=".png",dir=extract_path).name
+                    face_image.save(path)
+                    face_images.append(path)
+                images+= face_images
+        return images
+    return None
+
+
+
+def save(batch_files, name):
+    batch_files = batch_files or []
+    print("Build", name, [x.name for x in batch_files])
+    faces = swapper.get_faces_from_img_files(batch_files)
+    blended_face = swapper.blend_faces(faces)
+    preview_path = os.path.join(
+        scripts.basedir(), "extensions", "sd-webui-roop", "references"
+    )
+    faces_path = os.path.join(scripts.basedir(), "models", "roop","faces")
+
+    target_img = None
+    if blended_face:
+        if blended_face["gender"] == 0:
+            target_img = Image.open(os.path.join(preview_path, "woman.png"))
+        else:
+            target_img = Image.open(os.path.join(preview_path, "man.png"))
+
+        if name == "":
+            name = "default_name"
+        pprint(blended_face)
+        result = swapper.swap_face(blended_face, blended_face, target_img, get_models()[0])
+        result_image = upscale_image(result.image, UpscaleOptions(face_restorer_name="CodeFormer", restorer_visibility=1))
+        
+        file_path = os.path.join(faces_path, f"{name}.pkl")
+        file_number = 1
+        while os.path.exists(file_path):
+            file_path = os.path.join(faces_path, f"{name}_{file_number}.pkl")
+            file_number += 1
+        result_image.save(file_path+".png")
+        with open(file_path, "wb") as file:
+            pickle.dump({"embedding" :blended_face.embedding, "gender" :blended_face.gender, "age" :blended_face.age},file)
+        try :
+            with open(file_path, "rb") as file:
+                data = Face(pickle.load(file))
+                print(data)
+        except Exception as e :
+            print(e)
+        return result_image
+
+    print("No face found")
+
+    return target_img
+
+
+
+
+def explore(model_path):
+    data = {
+        'Node Name': [],
+        'Op Type': [],
+        'Inputs': [],
+        'Outputs': [],
+        'Attributes': []
+    }
+    if model_path:
+        model = onnx.load(model_path)
+        for node in model.graph.node:
+            data['Node Name'].append(pformat(node.name))
+            data['Op Type'].append(pformat(node.op_type))
+            data['Inputs'].append(pformat(node.input))
+            data['Outputs'].append(pformat(node.output))
+            attributes = []
+            for attr in node.attribute:
+                attr_name = attr.name
+                attr_value = attr.t
+                attributes.append("{} = {}".format(pformat(attr_name), pformat(attr_value)))
+            data['Attributes'].append(attributes)
+
+    df = pd.DataFrame(data)
+    return df
+
+def upscaler_ui():
+    with gr.Tab(f"Upscaler"):
+        with gr.Row():
+            face_restorer_name = gr.Radio(
+                label="Restore Face",
+                choices=["None"] + [x.name() for x in shared.face_restorers],
+                value=shared.face_restorers[0].name(),
+                type="value",
+            )
+            face_restorer_visibility = gr.Slider(
+                0, 1, 1, step=0.1, label="Restore visibility"
+            )
+        upscaler_name = gr.inputs.Dropdown(
+            choices=[upscaler.name for upscaler in shared.sd_upscalers],
+            label="Upscaler",
+        )
+        upscaler_scale = gr.Slider(1, 8, 1, step=0.1, label="Upscaler scale")
+        upscaler_visibility = gr.Slider(
+            0, 1, 1, step=0.1, label="Upscaler visibility (if scale = 1)"
+        )
+    return [
+        face_restorer_name,
+        face_restorer_visibility,
+        upscaler_name,
+        upscaler_scale,
+        upscaler_visibility,
+    ]
+
+def tools_ui():
+    models = get_models()
+    with gr.Tab("Tools"):
+        with gr.Tab("Build"):
+            with gr.Row():
+                batch_files = gr.components.File(
+                    type="file",
+                    file_count="multiple",
+                    label="Batch Sources Images",
+                    optional=True,
+                )
+                preview = gr.components.Image(type="pil", label="Preview", interactive=False)
+            name = gr.Textbox(
+                value="Face",
+                placeholder="Name of the character",
+                label="Name of the character",
+            )
+            generate_checkpoint_btn = gr.Button("Save")
+        with gr.Tab("Compare"):
+            with gr.Row():
+                img1 = gr.components.Image(type="pil", label="Face 1")
+                img2 = gr.components.Image(type="pil", label="Face 2")
+            compare_btn = gr.Button("Compare")
+            compare_result_text = gr.Textbox(
+                interactive=False, label="Similarity", value="0"
+            )
+        with gr.Tab("Extract"):
+            with gr.Row():
+                extracted_source_files = gr.components.File(
+                    type="file",
+                    file_count="multiple",
+                    label="Batch Sources Images",
+                    optional=True,
+                )
+                extracted_faces =  gr.Gallery(
+                                        label="Extracted faces", show_label=False
+                                    ).style(columns=[2], rows=[2])
+            extract_save_path = gr.Textbox(label="Destination Directory", value="")
+            extract_btn = gr.Button("Extract")
+        with gr.Tab("Explore Model"):
+            model = gr.inputs.Dropdown(
+                choices=models,
+                label="Model not found, please download one and reload automatic 1111",
+            )            
+            explore_btn = gr.Button("Explore")
+            explore_result_text = gr.Dataframe(
+                interactive=False, label="Explored"
+            )
+        upscale_options = upscaler_ui()
+
+    explore_btn.click(explore, inputs=[model], outputs=[explore_result_text])  
+    compare_btn.click(compare, inputs=[img1, img2], outputs=[compare_result_text])
+    generate_checkpoint_btn.click(save, inputs=[batch_files, name], outputs=[preview])
+    extract_btn.click(extract_faces, inputs=[extracted_source_files, extract_save_path]+upscale_options, outputs=[extracted_faces])  
+
+def on_ui_tabs() :
+    with gr.Blocks(analytics_enabled=False) as ui_faceswap:
+        tools_ui()
+    return [(ui_faceswap, "FaceTools", "roop_tab")] 
+
+script_callbacks.on_ui_tabs(on_ui_tabs)
 
 class FaceSwapScript(scripts.Script):
+    units_count = 3
+
     def title(self):
         return f"roop"
 
     def show(self, is_img2img):
         return scripts.AlwaysVisible
 
-    def ui(self, is_img2img):
-        with gr.Accordion(f"roop {version_flag}", open=False):
+    def faceswap_unit_ui(self, is_img2img, unit_num=1):
+        with gr.Tab(f"Face {unit_num}"):
             with gr.Column():
-                img = gr.inputs.Image(type="pil")
-                enable = gr.Checkbox(False, placeholder="enable", label="Enable")
+                with gr.Row():
+                    img = gr.components.Image(type="pil", label="Reference")
+                    batch_files = gr.components.File(
+                        type="file",
+                        file_count="multiple",
+                        label="Batch Sources Images",
+                        optional=True,
+                    )
+                with gr.Row() :
+                    face = gr.inputs.Dropdown(
+                        choices=get_faces(),
+                        label="Face Checkpoint",
+                    )
+                    refresh = gr.Button(value='↻', variant='tool')
+                    def refresh_fn(selected):
+                        return gr.Dropdown.update(value=selected, choices=get_faces())
+                    refresh.click(fn=refresh_fn,inputs=face, outputs=face)
+
+                with gr.Row():
+                    enable = gr.Checkbox(False, placeholder="enable", label="Enable")
+                    same_gender = gr.Checkbox(
+                        False, placeholder="Same Gender", label="Same Gender"
+                    )
+                    blend_faces = gr.Checkbox(
+                        True, placeholder="Blend Faces", label="Blend Faces ((Source|Checkpoint)+References = 1)"
+                    )
+                min_sim = gr.Slider(0, 1, 0, step=0.01, label="Min similarity")
+                min_ref_sim = gr.Slider(
+                    0, 1, 0, step=0.01, label="Min reference similarity"
+                )
                 faces_index = gr.Textbox(
                     value="0",
-                    placeholder="Which face to swap (comma separated), start from 0",
+                    placeholder="Which face to swap (comma separated), start from 0 (by gender if same_gender is enabled)",
                     label="Comma separated face number(s)",
                 )
-                with gr.Row():
-                    face_restorer_name = gr.Radio(
-                        label="Restore Face",
-                        choices=["None"] + [x.name() for x in shared.face_restorers],
-                        value=shared.face_restorers[0].name(),
-                        type="value",
-                    )
-                    face_restorer_visibility = gr.Slider(
-                        0, 1, 1, step=0.1, label="Restore visibility"
-                    )
-                upscaler_name = gr.inputs.Dropdown(
-                    choices=[upscaler.name for upscaler in shared.sd_upscalers],
-                    label="Upscaler",
-                )
-                upscaler_scale = gr.Slider(1, 8, 1, step=0.1, label="Upscaler scale")
-                upscaler_visibility = gr.Slider(
-                    0, 1, 1, step=0.1, label="Upscaler visibility (if scale = 1)"
-                )
-
-                models = get_models()
-                if len(models) == 0:
-                    logger.warning(
-                        "You should at least have one model in models directory, please read the doc here : https://github.com/s0md3v/sd-webui-roop/"
-                    )
-                    model = gr.inputs.Dropdown(
-                        choices=models,
-                        label="Model not found, please download one and reload automatic 1111",
-                    )
-                else:
-                    model = gr.inputs.Dropdown(
-                        choices=models, label="Model", default=models[0]
-                    )
-
                 swap_in_source = gr.Checkbox(
                     False,
                     placeholder="Swap face in source image",
-                    label="Swap in source image",
+                    label="Swap in source image (must be blended)",
                     visible=is_img2img,
                 )
                 swap_in_generated = gr.Checkbox(
@@ -87,108 +373,140 @@ class FaceSwapScript(scripts.Script):
                     label="Swap in generated image",
                     visible=is_img2img,
                 )
-
         return [
             img,
+            face,
+            batch_files,
+            blend_faces,
             enable,
+            same_gender,
+            min_sim,
+            min_ref_sim,
             faces_index,
-            model,
-            face_restorer_name,
-            face_restorer_visibility,
-            upscaler_name,
-            upscaler_scale,
-            upscaler_visibility,
             swap_in_source,
             swap_in_generated,
         ]
 
-    @property
-    def upscaler(self) -> UpscalerData:
-        for upscaler in shared.sd_upscalers:
-            if upscaler.name == self.upscaler_name:
-                return upscaler
-        return None
-
-    @property
-    def face_restorer(self) -> FaceRestoration:
-        for face_restorer in shared.face_restorers:
-            if face_restorer.name() == self.face_restorer_name:
-                return face_restorer
-        return None
-
-    @property
-    def upscale_options(self) -> UpscaleOptions:
-        return UpscaleOptions(
-            scale=self.upscaler_scale,
-            upscaler=self.upscaler,
-            face_restorer=self.face_restorer,
-            upscale_visibility=self.upscaler_visibility,
-            restorer_visibility=self.face_restorer_visibility,
-        )
-
-    def process(
-        self,
-        p: StableDiffusionProcessing,
-        img,
-        enable,
-        faces_index,
-        model,
-        face_restorer_name,
-        face_restorer_visibility,
-        upscaler_name,
-        upscaler_scale,
-        upscaler_visibility,
-        swap_in_source,
-        swap_in_generated,
-    ):
-        self.source = img
-        self.face_restorer_name = face_restorer_name
-        self.upscaler_scale = upscaler_scale
-        self.upscaler_visibility = upscaler_visibility
-        self.face_restorer_visibility = face_restorer_visibility
-        self.enable = enable
-        self.upscaler_name = upscaler_name
-        self.swap_in_generated = swap_in_generated
-        self.model = model
-        self.faces_index = {
-            int(x) for x in faces_index.strip(",").split(",") if x.isnumeric()
-        }
-        if len(self.faces_index) == 0:
-            self.faces_index = {0}
-        if self.enable:
-            if self.source is not None:
-                if isinstance(p, StableDiffusionProcessingImg2Img) and swap_in_source:
-                    logger.info(f"roop enabled, face index %s", self.faces_index)
-
-                    for i in range(len(p.init_images)):
-                        logger.info(f"Swap in source %s", i)
-                        result = swap_face(
-                            self.source,
-                            p.init_images[i],
-                            faces_index=self.faces_index,
-                            model=self.model,
-                            upscale_options=self.upscale_options,
-                        )
-                        p.init_images[i] = result.image()
-            else:
-                logger.error(f"Please provide a source face")
-
-    def postprocess_batch(self, *args, **kwargs):
-        if self.enable:
-            return images
-
-    def postprocess_image(self, p, script_pp: scripts.PostprocessImageArgs, *args):
-        if self.enable and self.swap_in_generated:
-            if self.source is not None:
-                image: Image.Image = script_pp.image
-                result: ImageResult = swap_face(
-                    self.source,
-                    image,
-                    faces_index=self.faces_index,
-                    model=self.model,
-                    upscale_options=self.upscale_options,
+    def configuration_ui(self, is_img2img):
+        with gr.Tab(f"Settings"):
+            models = get_models()
+            show_unmodified = gr.Checkbox(
+                False,
+                placeholder="Show Unmodified",
+                label="Show Unmodified (original)",
+            )
+            if len(models) == 0:
+                logger.warning(
+                    "You should at least have one model in models directory, please read the doc here : https://github.com/s0md3v/sd-webui-roop"
                 )
-                pp = scripts_postprocessing.PostprocessedImage(result.image())
-                pp.info = {}
-                p.extra_generation_params.update(pp.info)
-                script_pp.image = pp.image
+                model = gr.inputs.Dropdown(
+                    choices=models,
+                    label="Model not found, please download one and reload automatic 1111",
+                )
+            else:
+                model = gr.inputs.Dropdown(
+                    choices=models, label="Model", default=models[0]
+                )
+        return [show_unmodified, model]
+
+    def ui(self, is_img2img):
+        with gr.Accordion(f"Roop {version_flag}", open=False):
+            components = []
+            for i in range(1, self.units_count + 1):
+                components += self.faceswap_unit_ui(is_img2img, i)
+            upscaler = upscaler_ui()
+            configuration = self.configuration_ui(is_img2img)
+            tools_ui()
+        return components + upscaler + configuration
+
+    def process(self, p: StableDiffusionProcessing, *components):
+        self.units: List[FaceSwapUnitSettings] = []
+        for i in range(0, self.units_count):
+            self.units += [FaceSwapUnitSettings.get_unit_configuration(i, components)]
+
+        for i, u in enumerate(self.units):
+            print(i, u)
+
+        len_conf: int = len(fields(FaceSwapUnitSettings))
+        shift: int = self.units_count * len_conf
+        self.upscale_options = UpscaleOptions(
+            *components[shift : shift + len(fields(UpscaleOptions))]
+        )
+        print(self.upscale_options)
+        self.model = components[-1]
+        self.show_unmodified = components[-2]
+
+        if isinstance(p, StableDiffusionProcessingImg2Img):
+            if any([u.enable for u in self.units]):
+                init_images = p.init_images
+                for unit in self.units:
+                    if unit.enable and unit.swap_in_source :
+                        (init_images, result_infos) = self.process_images_unit(unit, init_images)
+                        logger.info(f"processed init image: {len(init_images)}, {len(result_infos)}")
+
+                p.init_images = init_images
+
+
+    def process_images_unit(self, unit, images, infos = None, processed = None)  :
+        if unit.enable :
+            result_images = []
+            result_infos = []
+            if not infos :
+                infos = [None] * len(images)
+            for i, (img, info) in enumerate(zip(images, infos)):
+                if convert_to_sd(img) :
+                    return(images,infos)
+                if not unit.blend_faces :
+                    src_faces = unit.faces
+                    logger.info(f"will generate {len(src_faces)} images")
+                else :
+                    logger.info("blend all faces together")
+                    src_faces = [unit.blended_faces]
+                if not processed or img.width == processed.width and img.height == processed.height :
+                    for i,src_face in enumerate(src_faces):
+                        logger.info(f"Process face {i}")
+                        result: swapper.ImageResult = swapper.swap_face(
+                            unit.reference_face if unit.reference_face is not None else src_face,
+                            src_face,
+                            img,
+                            faces_index=unit.faces_index,
+                            model=self.model,
+                            same_gender=unit.same_gender,
+                        )
+                        if result.similarity and all([result.similarity.values()!=0]+[x >= unit.min_sim for x in result.similarity.values()]) and all([result.ref_similarity.values()!=0]+[x >= unit.min_ref_sim for x in result.ref_similarity.values()]):
+                            result_infos.append(f"{info}, similarity = {result.similarity}, ref_similarity = {result.ref_similarity}")
+                            result_images.append(result.image)
+                        else:
+                            logger.info(
+                                f"skip, similarity to low, sim = {result.similarity} (target {unit.min_sim}) ref sim = {result.ref_similarity} (target = {unit.min_ref_sim})"
+                            )
+            logger.info(f"{len(result_images)} images processed")
+            return (result_images, result_infos)
+        return (images, infos)
+
+    def postprocess(self, p, processed: Processed, *args):
+        orig_images = processed.images
+        orig_infos = processed.infotexts
+
+        if any([u.enable for u in self.units]):
+            result_images = processed.images[:]
+            result_infos = processed.infotexts
+            for unit in self.units:
+                if unit.enable and unit.swap_in_generated :
+                    (result_images, result_infos) = self.process_images_unit(unit, result_images, result_infos, processed)
+                    logger.info(f"processed : {len(result_images)}, {len(result_infos)}")
+
+            for i, img in enumerate(result_images):
+                if self.upscale_options is not None:
+                    result_images[i] = upscale_image(img, self.upscale_options)
+
+            if len(result_images) > 1:
+                result_images.append(create_square_image(result_images))
+
+            processed.images = result_images
+            processed.infotexts = result_infos
+            
+            if self.show_unmodified:
+                processed.images += orig_images
+                processed.infotexts+= orig_infos
+

scripts/imgutils.py

@@ -0,0 +1,64 @@
+from PIL import Image
+import cv2
+import numpy as np
+from math import isqrt, ceil
+
+def pil_to_cv2(pil_img):
+    return cv2.cvtColor(np.array(pil_img), cv2.COLOR_RGB2BGR)
+
+
+def cv2_to_pil(cv2_img):
+    return Image.fromarray(cv2.cvtColor(cv2_img, cv2.COLOR_BGR2RGB))
+
+
+
+def torch_to_pil(images):
+    """
+    Convert a numpy image or a batch of images to a PIL image.
+    """
+    images = images.cpu().permute(0, 2, 3, 1).numpy()
+    if images.ndim == 3:
+        images = images[None, ...]
+    images = (images * 255).round().astype("uint8")
+    pil_images = [Image.fromarray(image) for image in images]
+    return pil_images
+
+
+def pil_to_torch(pil_images):
+    """
+    Convert a PIL image or a list of PIL images to a torch tensor or a batch of torch tensors.
+    """
+    if isinstance(pil_images, list):
+        numpy_images = [np.array(image) for image in pil_images]
+        torch_images = torch.from_numpy(np.stack(numpy_images)).permute(0, 3, 1, 2)
+        return torch_images
+
+    numpy_image = np.array(pil_images)
+    torch_image = torch.from_numpy(numpy_image).permute(2, 0, 1)
+    return torch_image
+
+
+
+def create_square_image(image_list):
+    size = None
+    for image in image_list:
+        if size is None:
+            size = image.size
+        elif image.size != size:
+            raise ValueError("Not same size images")
+
+    num_images = len(image_list)
+    rows = isqrt(num_images)
+    cols = ceil(num_images / rows)
+
+    square_size = (cols * size[0], rows * size[1])
+
+    square_image = Image.new("RGB", square_size)
+
+    for i, image in enumerate(image_list):
+        row = i // cols
+        col = i % cols
+
+        square_image.paste(image, (col * size[0], row * size[1]))
+
+    return square_image

scripts/roop_logging.py

@@ -24,7 +24,7 @@ class ColoredFormatter(logging.Formatter):
 
 
 # Create a new logger
-logger = logging.getLogger("roop")
+logger = logging.getLogger("Roop")
 logger.propagate = False
 
 # Add handler if we don't have one.

scripts/roop_version.py

@@ -1,5 +1,5 @@
-version_flag = "v0.0.2"
+version_flag = "v0.1.0"
 
 from scripts.roop_logging import logger
 
-logger.info(f"roop {version_flag}")
+logger.info(f"Roop {version_flag}")

scripts/swapper.py

@@ -1,7 +1,5 @@
 import copy
-import math
 import os
-import tempfile
 from dataclasses import dataclass
 from typing import List, Union, Dict, Set, Tuple
 
@@ -11,22 +9,32 @@ from PIL import Image
 
 import insightface
 import onnxruntime
-from scripts.cimage import convert_to_sd
-
-from modules.face_restoration import FaceRestoration, restore_faces
-from modules.upscaler import Upscaler, UpscalerData
 from scripts.roop_logging import logger
+from pprint import pprint
+from sklearn.metrics.pairwise import cosine_similarity
 
-providers = onnxruntime.get_available_providers()
+from scripts.imgutils import pil_to_cv2, cv2_to_pil
+
+providers = ["CPUExecutionProvider"]
 
 
-@dataclass
-class UpscaleOptions:
-    scale: int = 1
-    upscaler: UpscalerData = None
-    upscale_visibility: float = 0.5
-    face_restorer: FaceRestoration = None
-    restorer_visibility: float = 0.5
+def cosine_similarity_face(face1, face2) -> float:
+    vec1 = face1.embedding.reshape(1, -1)
+    vec2 = face2.embedding.reshape(1, -1)
+    return max(0, cosine_similarity(vec1, vec2)[0, 0])
+
+
+ANALYSIS_MODEL = None
+
+
+def getAnalysisModel():
+    global ANALYSIS_MODEL
+    if ANALYSIS_MODEL is None:
+        ANALYSIS_MODEL = insightface.app.FaceAnalysis(
+            name="buffalo_l", providers=providers
+        )
+    return ANALYSIS_MODEL
+
 
 FS_MODEL = None
 CURRENT_FS_MODEL_PATH = None
@@ -42,101 +50,119 @@ def getFaceSwapModel(model_path: str):
     return FS_MODEL
 
 
-def upscale_image(image: Image, upscale_options: UpscaleOptions):
-    result_image = image
-    if upscale_options.upscaler is not None and upscale_options.upscaler.name != "None":
-        original_image = result_image.copy()
-        logger.info(
-            "Upscale with %s scale = %s",
-            upscale_options.upscaler.name,
-            upscale_options.scale,
-        )
-        result_image = upscale_options.upscaler.scaler.upscale(
-            image, upscale_options.scale, upscale_options.upscaler.data_path
-        )
-        if upscale_options.scale == 1:
-            result_image = Image.blend(
-                original_image, result_image, upscale_options.upscale_visibility
-            )
-
-    if upscale_options.face_restorer is not None:
-        original_image = result_image.copy()
-        logger.info("Restore face with %s", upscale_options.face_restorer.name())
-        numpy_image = np.array(result_image)
-        numpy_image = upscale_options.face_restorer.restore(numpy_image)
-        restored_image = Image.fromarray(numpy_image)
-        result_image = Image.blend(
-            original_image, restored_image, upscale_options.restorer_visibility
-        )
-
-    return result_image
-
-
-def get_face_single(img_data: np.ndarray, face_index=0, det_size=(640, 640)):
-    face_analyser = insightface.app.FaceAnalysis(name="buffalo_l", providers=providers)
+def get_faces(img_data: np.ndarray, det_size=(640, 640)):
+    face_analyser = copy.deepcopy(getAnalysisModel())
     face_analyser.prepare(ctx_id=0, det_size=det_size)
     face = face_analyser.get(img_data)
 
     if len(face) == 0 and det_size[0] > 320 and det_size[1] > 320:
         det_size_half = (det_size[0] // 2, det_size[1] // 2)
-        return get_face_single(img_data, face_index=face_index, det_size=det_size_half)
+        return get_faces(img_data, det_size=det_size_half)
 
     try:
-        return sorted(face, key=lambda x: x.bbox[0])[face_index]
+        return sorted(face, key=lambda x: x.bbox[0])
     except IndexError:
         return None
 
 
+def compare_faces(img1: Image.Image, img2: Image.Image) -> float:
+    face1 = get_or_default(get_faces(pil_to_cv2(img1)), 0, None)
+    face2 = get_or_default(get_faces(pil_to_cv2(img2)), 0, None)
+
+    if face1 is not None and face2 is not None:
+        return cosine_similarity_face(face1, face2)
+    return -1
+
+
 @dataclass
 class ImageResult:
-    path: Union[str, None] = None
-    similarity: Union[Dict[int, float], None] = None  # face, 0..1
+    image: Image.Image
+    similarity: Dict[int, float]  # face, 0..1
+    ref_similarity: Dict[int, float]  # face, 0..1
 
-    def image(self) -> Union[Image.Image, None]:
-        if self.path:
-            return Image.open(self.path)
-        return None
+
+def get_or_default(l, index, default):
+    return l[index] if index < len(l) else default
+
+def get_faces_from_img_files(files) :
+    faces = []
+    if len(files) > 0 :
+        for file in files :
+            print("open", file.name)
+            img = Image.open(file.name)
+            face = get_or_default(get_faces(pil_to_cv2(img)), 0, None)
+            if face is not None :
+                faces.append(face)
+    return faces
+
+
+def blend_faces(faces) :
+    embeddings = [face.embedding for face in faces]
+    if len(embeddings)> 0 :
+        embedding_shape = embeddings[0].shape
+        for embedding in embeddings:
+            if embedding.shape != embedding_shape:
+                raise ValueError("embedding shape mismatch")
+
+        blended_embedding = np.mean(embeddings, axis=0)
+        blended = faces[0]
+        blended.embedding = blended_embedding
+        return blended
+    return None
 
 
 def swap_face(
-    source_img: Image.Image,
+    reference_face: np.ndarray,
+    source_face: np.ndarray,
     target_img: Image.Image,
-    model: Union[str, None] = None,
+    model: str,
     faces_index: Set[int] = {0},
-    upscale_options: Union[UpscaleOptions, None] = None,
+    same_gender=True,
 ) -> ImageResult:
-    result_image = target_img
-    converted = convert_to_sd(target_img)
-    scale, fn = converted[0], converted[1]
-    if model is not None and not scale:
-        if isinstance(source_img, str):  # source_img is a base64 string
-            import base64, io
-            if 'base64,' in source_img:  # check if the base64 string has a data URL scheme
-                base64_data = source_img.split('base64,')[-1]
-                img_bytes = base64.b64decode(base64_data)
-            else:
-                # if no data URL scheme, just decode
-                img_bytes = base64.b64decode(source_img)
-            source_img = Image.open(io.BytesIO(img_bytes))
-        source_img = cv2.cvtColor(np.array(source_img), cv2.COLOR_RGB2BGR)
-        target_img = cv2.cvtColor(np.array(target_img), cv2.COLOR_RGB2BGR)
-        source_face = get_face_single(source_img, face_index=0)
-        if source_face is not None:
-            result = target_img
-            model_path = os.path.join(os.path.abspath(os.path.dirname(__file__)), model)
-            face_swapper = getFaceSwapModel(model_path)
+    return_result = ImageResult(target_img, {}, {})
+    target_img = cv2.cvtColor(np.array(target_img), cv2.COLOR_RGB2BGR)
+    gender = source_face["gender"]
+    print("Source Gender ", gender)
+    if source_face is not None:
+        result = target_img
+        model_path = os.path.join(os.path.abspath(os.path.dirname(__file__)), model)
+        face_swapper = getFaceSwapModel(model_path)
+        target_faces = get_faces(target_img)
+        print("Target faces count", len(target_faces))
 
-            for face_num in faces_index:
-                target_face = get_face_single(target_img, face_index=face_num)
-                if target_face is not None:
-                    result = face_swapper.get(result, target_face, source_face)
-                else:
-                    logger.info(f"No target face found for {face_num}")
+        if same_gender:
+            target_faces = [x for x in target_faces if x["gender"] == gender]
+            print("Target Gender Matches count", len(target_faces))
 
-            result_image = Image.fromarray(cv2.cvtColor(result, cv2.COLOR_BGR2RGB))
-            if upscale_options is not None:
-                result_image = upscale_image(result_image, upscale_options)
-        else:
-            logger.info("No source face found")
-    result_image.save(fn.name)
-    return ImageResult(path=fn.name)
+        for i, swapped_face in enumerate(target_faces):
+            logger.info(f"swap face {i}")
+            if i in faces_index:
+                result = face_swapper.get(result, swapped_face, source_face)
+
+        result_image = Image.fromarray(cv2.cvtColor(result, cv2.COLOR_BGR2RGB))
+        return_result.image = result_image
+
+        try:
+            result_faces = get_faces(
+                cv2.cvtColor(np.array(result_image), cv2.COLOR_RGB2BGR)
+            )
+            if same_gender:
+                result_faces = [x for x in result_faces if x["gender"] == gender]
+
+            for i, swapped_face in enumerate(result_faces):
+                logger.info(f"compare face {i}")
+                if i in faces_index and i < len(target_faces):
+                    return_result.similarity[i] = cosine_similarity_face(
+                        source_face, swapped_face
+                    )
+                    return_result.ref_similarity[i] = cosine_similarity_face(
+                        reference_face, swapped_face
+                    )
+
+                print("similarity", return_result.similarity)
+                print("ref similarity", return_result.ref_similarity)
+
+        except Exception as e:
+            logger.error(str(e))
+
+    return return_result

scripts/upscaling.py

@@ -0,0 +1,64 @@
+from modules.face_restoration import FaceRestoration
+from modules.upscaler import UpscalerData
+from dataclasses import dataclass
+from typing import List, Union, Dict, Set, Tuple
+from scripts.roop_logging import logger
+from PIL import Image
+import numpy as np
+from modules import scripts, shared
+
+
+@dataclass
+class UpscaleOptions:
+    face_restorer_name: str = ""
+    restorer_visibility: float = 0.5
+    upscaler_name: str = ""
+    scale: int = 1
+    upscale_visibility: float = 0.5
+
+    @property
+    def upscaler(self) -> UpscalerData:
+        for upscaler in shared.sd_upscalers:
+            if upscaler.name == self.upscaler_name:
+                return upscaler
+        return None
+
+    @property
+    def face_restorer(self) -> FaceRestoration:
+        for face_restorer in shared.face_restorers:
+            if face_restorer.name() == self.face_restorer_name:
+                return face_restorer
+        return None
+
+
+def upscale_image(image: Image.Image, upscale_options: UpscaleOptions):
+    result_image = image
+    try :
+        if upscale_options.upscaler is not None and upscale_options.upscaler.name != "None":
+            original_image = result_image.copy()
+            logger.info(
+                "Upscale with %s scale = %s",
+                upscale_options.upscaler.name,
+                upscale_options.scale,
+            )
+            result_image = upscale_options.upscaler.scaler.upscale(
+                image, upscale_options.scale, upscale_options.upscaler.data_path
+            )
+            if upscale_options.scale == 1:
+                result_image = Image.blend(
+                    original_image, result_image, upscale_options.upscale_visibility
+                )
+
+        if upscale_options.face_restorer is not None:
+            original_image = result_image.copy()
+            logger.info("Restore face with %s", upscale_options.face_restorer.name())
+            numpy_image = np.array(result_image)
+            numpy_image = upscale_options.face_restorer.restore(numpy_image)
+            restored_image = Image.fromarray(numpy_image)
+            result_image = Image.blend(
+                original_image, restored_image, upscale_options.restorer_visibility
+            )
+    except Exception as e:
+        logger.info("Failed to upscale %s", e)
+
+    return result_image