diff --git a/README.md b/README.md index 252d774..aff7b13 100644 --- a/README.md +++ b/README.md @@ -1,4 +1,4 @@ -# Roop-GE 0.2.3b for StableDiffusion +# Roop-GE 0.2.3 for StableDiffusion ### NSFW (uncensored) version (use it on your own responsibility) of [original sd-webui-roop](https://github.com/s0md3v/sd-webui-roop) with a lot of improvements > GE (Gourieff Edition), aka "NSFW-Roop" diff --git a/example/api_example.py b/example/api_example.py index f8d75c7..0b3eac1 100644 --- a/example/api_example.py +++ b/example/api_example.py @@ -35,6 +35,7 @@ args=[ 1, #10 Upscaler visibility (if scale = 1) False, #11 Swap in source image True, #12 Swap in generated image + 1, #13 Console Log Level (0 - min, 1 - med or 2 - max) ] # The args for roop-ge can be found by diff --git a/requirements.txt b/requirements.txt index 661461e..059457a 100644 --- a/requirements.txt +++ b/requirements.txt @@ -2,4 +2,3 @@ insightface==0.7.3 onnx==1.14.0 onnxruntime==1.15.0 opencv-python==4.7.0.72 -diffusers==0.17.1 diff --git a/scripts/console_log_patch.py b/scripts/console_log_patch.py new file mode 100644 index 0000000..1e737e3 --- /dev/null +++ b/scripts/console_log_patch.py @@ -0,0 +1,120 @@ +import os.path as osp +import glob +import logging +import insightface +from insightface.model_zoo.model_zoo import ModelRouter, PickableInferenceSession +from insightface.model_zoo.retinaface import RetinaFace +from insightface.model_zoo.landmark import Landmark +from insightface.model_zoo.attribute import Attribute +from insightface.model_zoo.inswapper import INSwapper +from insightface.model_zoo.arcface_onnx import ArcFaceONNX +from insightface.app import FaceAnalysis +from insightface.utils import DEFAULT_MP_NAME, ensure_available +from insightface.model_zoo import model_zoo +import onnxruntime +import onnx +from onnx import numpy_helper +from scripts.logger import logger + + +def patched_get_model(self, **kwargs): + session = PickableInferenceSession(self.onnx_file, **kwargs) + inputs = session.get_inputs() + input_cfg = inputs[0] + input_shape = input_cfg.shape + outputs = session.get_outputs() + + if len(outputs) >= 5: + return RetinaFace(model_file=self.onnx_file, session=session) + elif input_shape[2] == 192 and input_shape[3] == 192: + return Landmark(model_file=self.onnx_file, session=session) + elif input_shape[2] == 96 and input_shape[3] == 96: + return Attribute(model_file=self.onnx_file, session=session) + elif len(inputs) == 2 and input_shape[2] == 128 and input_shape[3] == 128: + return INSwapper(model_file=self.onnx_file, session=session) + elif input_shape[2] == input_shape[3] and input_shape[2] >= 112 and input_shape[2] % 16 == 0: + return ArcFaceONNX(model_file=self.onnx_file, session=session) + else: + return None + + +def patched_faceanalysis_init(self, name=DEFAULT_MP_NAME, root='~/.insightface', allowed_modules=None, **kwargs): + onnxruntime.set_default_logger_severity(3) + self.models = {} + self.model_dir = ensure_available('models', name, root=root) + onnx_files = glob.glob(osp.join(self.model_dir, '*.onnx')) + onnx_files = sorted(onnx_files) + for onnx_file in onnx_files: + model = model_zoo.get_model(onnx_file, **kwargs) + if model is None: + print('model not recognized:', onnx_file) + elif allowed_modules is not None and model.taskname not in allowed_modules: + print('model ignore:', onnx_file, model.taskname) + del model + elif model.taskname not in self.models and (allowed_modules is None or model.taskname in allowed_modules): + self.models[model.taskname] = model + else: + print('duplicated model task type, ignore:', onnx_file, model.taskname) + del model + assert 'detection' in self.models + self.det_model = self.models['detection'] + + +def patched_faceanalysis_prepare(self, ctx_id, det_thresh=0.5, det_size=(640, 640)): + self.det_thresh = det_thresh + assert det_size is not None + self.det_size = det_size + for taskname, model in self.models.items(): + if taskname == 'detection': + model.prepare(ctx_id, input_size=det_size, det_thresh=det_thresh) + else: + model.prepare(ctx_id) + + +def patched_inswapper_init(self, model_file=None, session=None): + self.model_file = model_file + self.session = session + model = onnx.load(self.model_file) + graph = model.graph + self.emap = numpy_helper.to_array(graph.initializer[-1]) + self.input_mean = 0.0 + self.input_std = 255.0 + if self.session is None: + self.session = onnxruntime.InferenceSession(self.model_file, None) + inputs = self.session.get_inputs() + self.input_names = [] + for inp in inputs: + self.input_names.append(inp.name) + outputs = self.session.get_outputs() + output_names = [] + for out in outputs: + output_names.append(out.name) + self.output_names = output_names + assert len(self.output_names) == 1 + input_cfg = inputs[0] + input_shape = input_cfg.shape + self.input_shape = input_shape + self.input_size = tuple(input_shape[2:4][::-1]) + + +def patch_insightface(get_model, faceanalysis_init, faceanalysis_prepare, inswapper_init): + insightface.model_zoo.model_zoo.ModelRouter.get_model = get_model + insightface.app.FaceAnalysis.__init__ = faceanalysis_init + insightface.app.FaceAnalysis.prepare = faceanalysis_prepare + insightface.model_zoo.inswapper.INSwapper.__init__ = inswapper_init + + +original_functions = [ModelRouter.get_model, FaceAnalysis.__init__, FaceAnalysis.prepare, INSwapper.__init__] +patched_functions = [patched_get_model, patched_faceanalysis_init, patched_faceanalysis_prepare, patched_inswapper_init] + + +def apply_logging_patch(console_logging_level): + if console_logging_level == 0: + patch_insightface(*patched_functions) + logger.setLevel(logging.WARNING) + elif console_logging_level == 1: + patch_insightface(*patched_functions) + logger.setLevel(logging.INFO) + elif console_logging_level == 2: + patch_insightface(*original_functions) + logger.setLevel(logging.INFO) diff --git a/scripts/faceswap.py b/scripts/faceswap.py index b29963b..aadc7ae 100644 --- a/scripts/faceswap.py +++ b/scripts/faceswap.py @@ -11,9 +11,10 @@ from PIL import Image import glob from modules.face_restoration import FaceRestoration -from scripts.roop_logging import logger +from scripts.logger import logger from scripts.swapper import UpscaleOptions, swap_face -from scripts.roop_version import version_flag +from scripts.version import version_flag +from scripts.console_log_patch import apply_logging_patch import os def get_models(): @@ -34,7 +35,7 @@ class FaceSwapScript(scripts.Script): with gr.Accordion(f"Roop-GE {version_flag}", open=False): with gr.Column(): img = gr.inputs.Image(type="pil") - enable = gr.Checkbox(False, placeholder="enable", label="Enable") + enable = gr.Checkbox(False, label="Enable") source_faces_index = gr.Textbox( value="0", placeholder="Which face(s) to use as source (comma separated)", @@ -57,7 +58,6 @@ class FaceSwapScript(scripts.Script): ) restore_first = gr.Checkbox( True, - placeholder="Restore face, then Upscale", label="1. Restore face -> 2. Upscale (-Uncheck- if you want vice versa)", ) upscaler_name = gr.inputs.Dropdown( @@ -76,7 +76,7 @@ class FaceSwapScript(scripts.Script): ) model = gr.inputs.Dropdown( choices=models, - label="Model not found, please download one and reload automatic 1111", + label="Model not found, please download one and reload WebUI", ) else: model = gr.inputs.Dropdown( @@ -85,16 +85,20 @@ class FaceSwapScript(scripts.Script): swap_in_source = gr.Checkbox( False, - placeholder="Swap face in source image", label="Swap in source image", visible=is_img2img, ) swap_in_generated = gr.Checkbox( True, - placeholder="Swap face in generated image", label="Swap in generated image", visible=is_img2img, ) + console_logging_level = gr.Radio( + ["Minimum", "Medium", "Maximum"], + value="Medium", + label="Console Log Level", + type="index", + ) return [ img, @@ -110,6 +114,7 @@ class FaceSwapScript(scripts.Script): upscaler_visibility, swap_in_source, swap_in_generated, + console_logging_level ] @@ -154,6 +159,7 @@ class FaceSwapScript(scripts.Script): upscaler_visibility, swap_in_source, swap_in_generated, + console_logging_level, ): self.source = img self.face_restorer_name = face_restorer_name @@ -165,6 +171,7 @@ class FaceSwapScript(scripts.Script): self.upscaler_name = upscaler_name self.swap_in_generated = swap_in_generated self.model = model + self.console_logging_level = console_logging_level self.source_faces_index = [ int(x) for x in source_faces_index.strip(",").split(",") if x.isnumeric() ] @@ -177,6 +184,7 @@ class FaceSwapScript(scripts.Script): self.faces_index = [0] if self.enable: if self.source is not None: + apply_logging_patch(console_logging_level) if isinstance(p, StableDiffusionProcessingImg2Img) and swap_in_source: logger.info(f"Roop-GE is enabled, face index %s", self.faces_index) diff --git a/scripts/roop_logging.py b/scripts/logger.py similarity index 96% rename from scripts/roop_logging.py rename to scripts/logger.py index 08aca39..55529c3 100644 --- a/scripts/roop_logging.py +++ b/scripts/logger.py @@ -1,41 +1,41 @@ -import logging -import copy -import sys - -from modules import shared - - -class ColoredFormatter(logging.Formatter): - COLORS = { - "DEBUG": "\033[0;36m", # CYAN - "INFO": "\033[0;32m", # GREEN - "WARNING": "\033[0;33m", # YELLOW - "ERROR": "\033[0;31m", # RED - "CRITICAL": "\033[0;37;41m", # WHITE ON RED - "RESET": "\033[0m", # RESET COLOR - } - - def format(self, record): - colored_record = copy.copy(record) - levelname = colored_record.levelname - seq = self.COLORS.get(levelname, self.COLORS["RESET"]) - colored_record.levelname = f"{seq}{levelname}{self.COLORS['RESET']}" - return super().format(colored_record) - - -# Create a new logger -logger = logging.getLogger("Roop-GE") -logger.propagate = False - -# Add handler if we don't have one. -if not logger.handlers: - handler = logging.StreamHandler(sys.stdout) - handler.setFormatter( - ColoredFormatter("%(asctime)s - %(name)s - %(levelname)s - %(message)s") - ) - logger.addHandler(handler) - -# Configure logger -loglevel_string = getattr(shared.cmd_opts, "roop_loglevel", "INFO") -loglevel = getattr(logging, loglevel_string.upper(), "info") -logger.setLevel(loglevel) +import logging +import copy +import sys + +from modules import shared + + +class ColoredFormatter(logging.Formatter): + COLORS = { + "DEBUG": "\033[0;36m", # CYAN + "INFO": "\033[0;32m", # GREEN + "WARNING": "\033[0;33m", # YELLOW + "ERROR": "\033[0;31m", # RED + "CRITICAL": "\033[0;37;41m", # WHITE ON RED + "RESET": "\033[0m", # RESET COLOR + } + + def format(self, record): + colored_record = copy.copy(record) + levelname = colored_record.levelname + seq = self.COLORS.get(levelname, self.COLORS["RESET"]) + colored_record.levelname = f"{seq}{levelname}{self.COLORS['RESET']}" + return super().format(colored_record) + + +# Create a new logger +logger = logging.getLogger("Roop-GE") +logger.propagate = False + +# Add handler if we don't have one. +if not logger.handlers: + handler = logging.StreamHandler(sys.stdout) + handler.setFormatter( + ColoredFormatter("%(asctime)s - %(name)s - %(levelname)s - %(message)s") + ) + logger.addHandler(handler) + +# Configure logger +loglevel_string = getattr(shared.cmd_opts, "roop_loglevel", "INFO") +loglevel = getattr(logging, loglevel_string.upper(), "info") +logger.setLevel(loglevel) diff --git a/scripts/roop_version.py b/scripts/roop_version.py deleted file mode 100644 index 91a70ec..0000000 --- a/scripts/roop_version.py +++ /dev/null @@ -1,5 +0,0 @@ -version_flag = "v0.2.3b" - -from scripts.roop_logging import logger - -logger.info(f"Roop-GE {version_flag}") diff --git a/scripts/swapper.py b/scripts/swapper.py index 585dcb9..e125769 100644 --- a/scripts/swapper.py +++ b/scripts/swapper.py @@ -12,7 +12,7 @@ import onnxruntime from modules.face_restoration import FaceRestoration from modules.upscaler import UpscalerData -from scripts.roop_logging import logger +from scripts.logger import logger providers = onnxruntime.get_available_providers() diff --git a/scripts/version.py b/scripts/version.py new file mode 100644 index 0000000..12b2421 --- /dev/null +++ b/scripts/version.py @@ -0,0 +1,5 @@ +version_flag = "v0.2.3" + +from scripts.logger import logger + +logger.info(f"Roop-GE {version_flag}")