Spaces:
Runtime error
Runtime error
File size: 2,285 Bytes
4c16f3d | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 | 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
# cuda_options = {"arena_extend_strategy": "kSameAsRequested", 'cudnn_conv_algo_search': 'DEFAULT'}
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
|