crystantine/sd-reactor-node

1import os.path as osp2import glob3import logging4import insightface5from insightface.model_zoo.model_zoo import ModelRouter, PickableInferenceSession6from insightface.model_zoo.retinaface import RetinaFace7from insightface.model_zoo.landmark import Landmark8from insightface.model_zoo.attribute import Attribute9from insightface.model_zoo.inswapper import INSwapper10from insightface.model_zoo.arcface_onnx import ArcFaceONNX11from insightface.app import FaceAnalysis12from insightface.utils import DEFAULT_MP_NAME, ensure_available13from insightface.model_zoo import model_zoo14import onnxruntime15import onnx16from onnx import numpy_helper17from scripts.logger import logger18 19 20def patched_get_model(self, **kwargs):21    session = PickableInferenceSession(self.onnx_file, **kwargs)22    inputs = session.get_inputs()23    input_cfg = inputs[0]24    input_shape = input_cfg.shape25    outputs = session.get_outputs()26 27    if len(outputs) >= 5:28        return RetinaFace(model_file=self.onnx_file, session=session)29    elif input_shape[2] == 192 and input_shape[3] == 192:30        return Landmark(model_file=self.onnx_file, session=session)31    elif input_shape[2] == 96 and input_shape[3] == 96:32        return Attribute(model_file=self.onnx_file, session=session)33    elif len(inputs) == 2 and input_shape[2] == 128 and input_shape[3] == 128:34        return INSwapper(model_file=self.onnx_file, session=session)35    elif input_shape[2] == input_shape[3] and input_shape[2] >= 112 and input_shape[2] % 16 == 0:36        return ArcFaceONNX(model_file=self.onnx_file, session=session)37    else:38        return None39 40 41def patched_faceanalysis_init(self, name=DEFAULT_MP_NAME, root='~/.insightface', allowed_modules=None, **kwargs):42    onnxruntime.set_default_logger_severity(3)43    self.models = {}44    self.model_dir = ensure_available('models', name, root=root)45    onnx_files = glob.glob(osp.join(self.model_dir, '*.onnx'))46    onnx_files = sorted(onnx_files)47    for onnx_file in onnx_files:48        model = model_zoo.get_model(onnx_file, **kwargs)49        if model is None:50            print('model not recognized:', onnx_file)51        elif allowed_modules is not None and model.taskname not in allowed_modules:52            print('model ignore:', onnx_file, model.taskname)53            del model54        elif model.taskname not in self.models and (allowed_modules is None or model.taskname in allowed_modules):55            self.models[model.taskname] = model56        else:57            print('duplicated model task type, ignore:', onnx_file, model.taskname)58            del model59    assert 'detection' in self.models60    self.det_model = self.models['detection']61 62 63def patched_faceanalysis_prepare(self, ctx_id, det_thresh=0.5, det_size=(640, 640)):64    self.det_thresh = det_thresh65    assert det_size is not None66    self.det_size = det_size67    for taskname, model in self.models.items():68        if taskname == 'detection':69            model.prepare(ctx_id, input_size=det_size, det_thresh=det_thresh)70        else:71            model.prepare(ctx_id)72 73 74def patched_inswapper_init(self, model_file=None, session=None):75    self.model_file = model_file76    self.session = session77    model = onnx.load(self.model_file)78    graph = model.graph79    self.emap = numpy_helper.to_array(graph.initializer[-1])80    self.input_mean = 0.081    self.input_std = 255.082    if self.session is None:83        self.session = onnxruntime.InferenceSession(self.model_file, None)84    inputs = self.session.get_inputs()85    self.input_names = []86    for inp in inputs:87        self.input_names.append(inp.name)88    outputs = self.session.get_outputs()89    output_names = []90    for out in outputs:91        output_names.append(out.name)92    self.output_names = output_names93    assert len(self.output_names) == 194    input_cfg = inputs[0]95    input_shape = input_cfg.shape96    self.input_shape = input_shape97    self.input_size = tuple(input_shape[2:4][::-1])98 99 100def patch_insightface(get_model, faceanalysis_init, faceanalysis_prepare, inswapper_init):101    insightface.model_zoo.model_zoo.ModelRouter.get_model = get_model102    insightface.app.FaceAnalysis.__init__ = faceanalysis_init103    insightface.app.FaceAnalysis.prepare = faceanalysis_prepare104    insightface.model_zoo.inswapper.INSwapper.__init__ = inswapper_init105 106 107original_functions = [ModelRouter.get_model, FaceAnalysis.__init__, FaceAnalysis.prepare, INSwapper.__init__]108patched_functions = [patched_get_model, patched_faceanalysis_init, patched_faceanalysis_prepare, patched_inswapper_init]109 110 111def apply_logging_patch(console_log_level):112    if console_log_level == 0:113        patch_insightface(*patched_functions)114        logger.setLevel(logging.WARNING)115    elif console_log_level == 1:116        patch_insightface(*patched_functions)117        logger.setLevel(logging.INFO)118    elif console_log_level == 2:119        patch_insightface(*original_functions)120        logger.setLevel(logging.INFO)121