| import roop.globals
|
| import numpy as np
|
| import onnx
|
| import onnxruntime
|
|
|
| from roop.typing import Face, Frame
|
| from roop.utilities import resolve_relative_path
|
|
|
|
|
|
|
| class FaceSwapInsightFace():
|
| plugin_options:dict = None
|
| model_swap_insightface = None
|
|
|
| processorname = 'faceswap'
|
| type = 'swap'
|
|
|
|
|
| def Initialize(self, plugin_options:dict):
|
| if self.plugin_options is not None:
|
| if self.plugin_options["devicename"] != plugin_options["devicename"] or self.plugin_options["modelname"] != plugin_options["modelname"]:
|
| self.Release()
|
|
|
| self.plugin_options = plugin_options
|
| if self.model_swap_insightface is None:
|
| model_path = resolve_relative_path('../models/' + self.plugin_options["modelname"])
|
| graph = onnx.load(model_path).graph
|
| self.emap = onnx.numpy_helper.to_array(graph.initializer[-1])
|
| self.devicename = self.plugin_options["devicename"].replace('mps', 'cpu')
|
| self.input_mean = 0.0
|
| self.input_std = 255.0
|
|
|
| sess_options = onnxruntime.SessionOptions()
|
| sess_options.enable_cpu_mem_arena = False
|
| self.model_swap_insightface = onnxruntime.InferenceSession(model_path, sess_options, providers=roop.globals.execution_providers)
|
|
|
|
|
|
|
| def Run(self, source_face: Face, target_face: Face, temp_frame: Frame) -> Frame:
|
| latent = source_face.normed_embedding.reshape((1,-1))
|
| latent = np.dot(latent, self.emap)
|
| latent /= np.linalg.norm(latent)
|
| io_binding = self.model_swap_insightface.io_binding()
|
| io_binding.bind_cpu_input("target", temp_frame)
|
| io_binding.bind_cpu_input("source", latent)
|
| io_binding.bind_output("output", self.devicename)
|
| self.model_swap_insightface.run_with_iobinding(io_binding)
|
| ort_outs = io_binding.copy_outputs_to_cpu()[0]
|
| return ort_outs[0]
|
|
|
|
|
| def Release(self):
|
| del self.model_swap_insightface
|
| self.model_swap_insightface = None
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|