diff --git a/modules/face/__init__.py b/modules/face/__init__.py index e0d76b689..5b4c3a31c 100644 --- a/modules/face/__init__.py +++ b/modules/face/__init__.py @@ -91,7 +91,9 @@ class Script(scripts.Script): def run(self, p: processing.StableDiffusionProcessing, mode, input_images, ip_model, ip_override, ip_cache, ip_strength, ip_structure, id_strength, id_conditioning, id_cache, pm_trigger, pm_strength, pm_start, fs_cache): # pylint: disable=arguments-differ, unused-argument if shared.backend != shared.Backend.DIFFUSERS: - return + return None + if mode == 'None': + return None if input_images is None or len(input_images) == 0: shared.log.error('Face: no init images') return None diff --git a/modules/face_restoration.py b/modules/face_restoration.py index d7fc5d1e9..d17191fdf 100644 --- a/modules/face_restoration.py +++ b/modules/face_restoration.py @@ -9,9 +9,9 @@ class FaceRestoration: return np_image -def restore_faces(np_image): +def restore_faces(np_image, p=None): face_restorers = [x for x in shared.face_restorers if x.name() == shared.opts.face_restoration_model or shared.opts.face_restoration_model is None] if len(face_restorers) == 0: return np_image face_restorer = face_restorers[0] - return face_restorer.restore(np_image) + return face_restorer.restore(np_image, p) diff --git a/modules/masking.py b/modules/masking.py index 89c6914d4..17495a9f1 100644 --- a/modules/masking.py +++ b/modules/masking.py @@ -397,7 +397,7 @@ def run_mask(input_image: Image.Image, input_mask: Image.Image = None, return_ty if mask_blur is not None: # compatibility with old img2img values which uses px values opts.mask_blur = round(4 * mask_blur / size, 3) if mask_padding is not None: # compatibility with old img2img values which uses px values - opts.mask_erode = 4 * mask_padding / size + opts.mask_dilate = 4 * mask_padding / size if opts.model is None or not segment_enable: mask = input_mask diff --git a/modules/postprocess/codeformer_model.py b/modules/postprocess/codeformer_model.py index 4a812fdb7..26dec124f 100644 --- a/modules/postprocess/codeformer_model.py +++ b/modules/postprocess/codeformer_model.py @@ -66,7 +66,7 @@ def setup_model(dirname): self.face_helper.face_det.to(device) # pylint: disable=no-member self.face_helper.face_parse.to(device) - def restore(self, np_image, w=None): + def restore(self, np_image, p=None, w=None): # pylint: disable=unused-argument from torchvision.transforms.functional import normalize from basicsr.utils import img2tensor, tensor2img np_image = np_image[:, :, ::-1] @@ -90,7 +90,7 @@ def setup_model(dirname): del output devices.torch_gc() except Exception as e: - shared.log.error(f'CodeForomer error: {e}') + shared.log.error(f'CodeFormer error: {e}') restored_face = tensor2img(cropped_face_t, rgb2bgr=True, min_max=(-1, 1)) restored_face = restored_face.astype('uint8') self.face_helper.add_restored_face(restored_face) diff --git a/modules/postprocess/gfpgan_model.py b/modules/postprocess/gfpgan_model.py index ee412012d..fb7ff0f5d 100644 --- a/modules/postprocess/gfpgan_model.py +++ b/modules/postprocess/gfpgan_model.py @@ -105,7 +105,7 @@ def setup_model(dirname): def name(self): return "GFPGAN" - def restore(self, np_image): + def restore(self, np_image, p=None): # pylint: disable=unused-argument return gfpgan_fix_faces(np_image) shared.face_restorers.append(FaceRestorerGFPGAN()) diff --git a/modules/processing.py b/modules/processing.py index 23d2bb5c9..0110d694e 100644 --- a/modules/processing.py +++ b/modules/processing.py @@ -262,8 +262,9 @@ def process_images_inner(p: StableDiffusionProcessing) -> Processed: shared.state.job_count = p.n_iter with devices.inference_context(), ema_scope_context(): t0 = time.time() - with devices.autocast(): - p.init(p.all_prompts, p.all_seeds, p.all_subseeds) + if not hasattr(p, 'skip_init'): + with devices.autocast(): + p.init(p.all_prompts, p.all_seeds, p.all_subseeds) extra_network_data = None debug(f'Processing inner: args={vars(p)}') for n in range(p.n_iter): @@ -340,7 +341,7 @@ def process_images_inner(p: StableDiffusionProcessing) -> Processed: p.restore_faces = orig images.save_image(Image.fromarray(x_sample), path=p.outpath_samples, basename="", seed=p.seeds[i], prompt=p.prompts[i], extension=shared.opts.samples_format, info=info, p=p, suffix="-before-face-restore") p.ops.append('face') - x_sample = face_restoration.restore_faces(x_sample) + x_sample = face_restoration.restore_faces(x_sample, p) image = Image.fromarray(x_sample) if p.scripts is not None and isinstance(p.scripts, scripts.ScriptRunner): pp = scripts.PostprocessImageArgs(image) @@ -366,7 +367,11 @@ def process_images_inner(p: StableDiffusionProcessing) -> Processed: images.save_image(image, p.outpath_samples, "", p.seeds[i], p.prompts[i], shared.opts.samples_format, info=text, p=p) # main save image if hasattr(p, 'mask_for_overlay') and p.mask_for_overlay and any([shared.opts.save_mask, shared.opts.save_mask_composite, shared.opts.return_mask, shared.opts.return_mask_composite]): image_mask = p.mask_for_overlay.convert('RGB') - image_mask_composite = Image.composite(image.convert('RGBA').convert('RGBa'), Image.new('RGBa', image.size), images.resize_image(3, p.mask_for_overlay, image.width, image.height).convert('L')).convert('RGBA') + image1 = image.convert('RGBA').convert('RGBa') + image2 = Image.new('RGBa', image.size) + mask = images.resize_image(3, p.mask_for_overlay, image.width, image.height).convert('L') + image_mask_composite = Image.composite(image1, image2, mask).convert('RGBA') + image_mask_composite.save('/tmp/composite.png') if shared.opts.save_mask: images.save_image(image_mask, p.outpath_samples, "", p.seeds[i], p.prompts[i], shared.opts.samples_format, info=text, p=p, suffix="-mask") if shared.opts.save_mask_composite: diff --git a/modules/processing_class.py b/modules/processing_class.py index d5fbe5551..27f7a161c 100644 --- a/modules/processing_class.py +++ b/modules/processing_class.py @@ -172,7 +172,7 @@ class StableDiffusionProcessing: def comment(self, text): self.comments[text] = 1 - def init(self, all_prompts, all_seeds, all_subseeds): + def init(self, all_prompts=None, all_seeds=None, all_subseeds=None): pass def sample(self, conditioning, unconditional_conditioning, seeds, subseeds, subseed_strength, prompts): @@ -230,11 +230,17 @@ class StableDiffusionProcessingTxt2Img(StableDiffusionProcessing): self.scripts = None self.script_args = [] - def init(self, all_prompts, all_seeds, all_subseeds): + def init(self, all_prompts=None, all_seeds=None, all_subseeds=None): if shared.backend == shared.Backend.DIFFUSERS: shared.sd_model = sd_models.set_diffuser_pipe(self.sd_model, sd_models.DiffusersTaskType.TEXT_2_IMAGE) self.width = self.width or 512 self.height = self.height or 512 + if all_prompts is not None: + self.all_prompts = all_prompts + if all_seeds is not None: + self.all_seeds = all_seeds + if all_subseeds is not None: + self.all_subseeds = all_subseeds def init_hr(self, scale = None, upscaler = None): scale = scale or self.hr_scale @@ -312,12 +318,19 @@ class StableDiffusionProcessingImg2Img(StableDiffusionProcessing): self.scripts = None self.script_args = [] - def init(self, all_prompts, all_seeds, all_subseeds): + def init(self, all_prompts=None, all_seeds=None, all_subseeds=None): if shared.backend == shared.Backend.DIFFUSERS and getattr(self, 'image_mask', None) is not None: shared.sd_model = sd_models.set_diffuser_pipe(self.sd_model, sd_models.DiffusersTaskType.INPAINTING) elif shared.backend == shared.Backend.DIFFUSERS and getattr(self, 'init_images', None) is not None: shared.sd_model = sd_models.set_diffuser_pipe(self.sd_model, sd_models.DiffusersTaskType.IMAGE_2_IMAGE) + if all_prompts is not None: + self.all_prompts = all_prompts + if all_seeds is not None: + self.all_seeds = all_seeds + if all_subseeds is not None: + self.all_subseeds = all_subseeds + if self.sampler_name == "PLMS": self.sampler_name = 'UniPC' if shared.backend == shared.Backend.ORIGINAL: @@ -491,3 +504,25 @@ class StableDiffusionProcessingControl(StableDiffusionProcessingImg2Img): # hypertile_set(self, hr=True) shared.state.job_count = 2 * self.n_iter shared.log.debug(f'Control hires: upscaler="{self.hr_upscaler}" upscale={scale} size={self.hr_upscale_to_x}x{self.hr_upscale_to_y}') + + +def switch_class(p: StableDiffusionProcessing, new_class: type, dct: dict = None): + import inspect + signature = inspect.signature(type(new_class).__init__, follow_wrapped=True) + possible = list(signature.parameters) + kwargs = {} + for k, v in p.__dict__.items(): + if k in possible: + kwargs[k] = v + if dct is not None: + for k, v in dct.items(): + if k in possible: + kwargs[k] = v + shared.log.debug(f"Switching class: {p.__class__} -> {new_class}") + p.__class__ = new_class + p.__init__(**kwargs) + if dct is not None: # post init set additional values + for k, v in dct.items(): + if hasattr(p, k): + setattr(p, k, v) + return p diff --git a/modules/processing_diffusers.py b/modules/processing_diffusers.py index 6ade59e8f..c84d89528 100644 --- a/modules/processing_diffusers.py +++ b/modules/processing_diffusers.py @@ -363,6 +363,8 @@ def process_diffusers(p: processing.StableDiffusionProcessing): # sanitize init_images if hasattr(p, 'init_images') and getattr(p, 'init_images', None) is None: del p.init_images + if hasattr(p, 'init_images') and not isinstance(getattr(p, 'init_images', []), list): + p.init_images = [p.init_images] if len(getattr(p, 'init_images', [])) > 0: while len(p.init_images) < len(p.prompts): p.init_images.append(p.init_images[-1]) diff --git a/modules/shared.py b/modules/shared.py index 960970fc0..12263ba38 100644 --- a/modules/shared.py +++ b/modules/shared.py @@ -668,9 +668,9 @@ options_templates.update(options_section(('postprocessing', "Postprocessing"), { "CLIP_stop_at_last_layers": OptionInfo(1, "Clip skip", gr.Slider, {"minimum": 1, "maximum": 8, "step": 1, "visible": False}), "postprocessing_sep_face_restoration": OptionInfo("

Face Restoration

", "", gr.HTML), - "face_restoration_model": OptionInfo("CodeFormer", "Face restoration model", gr.Radio, lambda: {"choices": [x.name() for x in face_restorers]}), + "face_restoration_model": OptionInfo("Face HiRes", "Face restoration model", gr.Radio, lambda: {"choices": [x.name() for x in face_restorers]}), "code_former_weight": OptionInfo(0.2, "CodeFormer weight parameter", gr.Slider, {"minimum": 0, "maximum": 1, "step": 0.01}), - "face_restoration_unload": OptionInfo(False, "Move face restoration model to CPU when complete"), + "face_restoration_unload": OptionInfo(False, "Move model to CPU when complete"), "postprocessing_sep_upscalers": OptionInfo("

Upscaling

", "", gr.HTML), "upscaler_unload": OptionInfo(False, "Unload upscaler after processing"), diff --git a/modules/ui_sections.py b/modules/ui_sections.py index e2cac4054..b2a186996 100644 --- a/modules/ui_sections.py +++ b/modules/ui_sections.py @@ -140,7 +140,7 @@ def create_advanced_inputs(tab): gr.HTML('
') with gr.Row(elem_id=f"{tab}_advanced_options"): full_quality = gr.Checkbox(label='Full quality', value=True, elem_id=f"{tab}_full_quality") - restore_faces = gr.Checkbox(label='Face restore', value=False, visible=len(shared.face_restorers) > 1, elem_id=f"{tab}_restore_faces") + restore_faces = gr.Checkbox(label='Face restore', value=False, elem_id=f"{tab}_restore_faces") tiling = gr.Checkbox(label='Tiling', value=False, elem_id=f"{tab}_tiling", visible=True) return cfg_scale, clip_skip, image_cfg_scale, diffusers_guidance_rescale, diffusers_sag_scale, cfg_end, full_quality, restore_faces, tiling @@ -226,7 +226,7 @@ def create_hires_inputs(tab): hr_sampler_index = gr.Dropdown(label='Secondary sampler', elem_id=f"{tab}_sampling_alt", choices=[x.name for x in sd_samplers.samplers], value='Default', type="index") with gr.Row(elem_id=f"{tab}_hires_row2"): hr_second_pass_steps = gr.Slider(minimum=0, maximum=99, step=1, label='HiRes steps', elem_id=f"{tab}_steps_alt", value=20) - denoising_strength = gr.Slider(minimum=0.0, maximum=0.99, step=0.01, label='Strength', value=0.5, elem_id=f"{tab}_denoising_strength") + denoising_strength = gr.Slider(minimum=0.0, maximum=0.99, step=0.01, label='Strength', value=0.3, elem_id=f"{tab}_denoising_strength") with gr.Group(visible=shared.backend == shared.Backend.DIFFUSERS): with gr.Row(elem_id=f"{tab}_refiner_row1", variant="compact"): refiner_start = gr.Slider(minimum=0.0, maximum=1.0, step=0.05, label='Refiner start', value=0.8, elem_id=f"{tab}_refiner_start") diff --git a/scripts/face-details.py b/scripts/face-details.py new file mode 100644 index 000000000..d2b706fc1 --- /dev/null +++ b/scripts/face-details.py @@ -0,0 +1,138 @@ +import os +import numpy as np +from PIL import Image, ImageDraw +from modules import shared, paths, devices, modelloader, processing, processing_class, face_restoration + + +class YoLoResult: + """Class face result""" + def __init__(self, score: float, box: list[int], mask: Image.Image = None, size: float = 0): + self.score = score + self.box = box + self.mask = mask + self.size = size + + +class FaceRestorerYolo(face_restoration.FaceRestoration): + def name(self): + return "Face HiRes" + + def __init__(self): + self.model = None + self.model_dir = os.path.join(paths.models_path, 'yolo') + self.model_name = 'yolov8n-face.pt' + self.model_url = 'https://github.com/akanametov/yolov8-face/releases/download/v0.0.0/yolov8n-face.pt' + + def predict( + self, + image: Image.Image, + offload: bool = False, + conf: float = 0.5, + iou: float = 0.5, + imgsz: int = 640, + half: bool = True, + device = 'cuda', + n: int = 5, + augment: bool = True, + agnostic: bool = False, + retina: bool = False, + mask: bool = True, + ) -> list[YoLoResult]: + + self.model.to(devices.device) + predictions = self.model.predict( + source=[image], + stream=False, + verbose=False, + conf=conf, + iou=iou, + imgsz=imgsz, + half=half, + device=device, + max_det=n, + augment=augment, + agnostic_nms=agnostic, + retina_masks=retina, + ) + if offload: + self.model.to('cpu') + result = [] + for prediction in predictions: + boxes = prediction.boxes.xyxy.detach().int().cpu().numpy() if prediction.boxes is not None else [] + scores = prediction.boxes.conf.detach().float().cpu().numpy() if prediction.boxes is not None else [] + for score, box in zip(scores, boxes): + box = box.tolist() + mask_image = None + size = (box[2] - box[0]) * (box[3] - box[1]) / (image.width * image.height) + if mask: + mask_image = image.copy() + mask_image = Image.new('L', image.size, 0) + draw = ImageDraw.Draw(mask_image) + draw.rectangle(box, fill="white", outline=None, width=0) + result.append(YoLoResult(score=score, box=box, mask=mask_image, size=size)) + return result + + def load(self): + if self.model is None: + model_files = modelloader.load_models(model_path=self.model_dir, model_url=self.model_url, download_name=self.model_name) + for f in model_files: + if self.model_name in f: + shared.log.info(f'Loading: type=FaceHires model={f}') + from ultralytics import YOLO # pylint: disable=import-outside-toplevel + self.model = YOLO(f) + + def restore(self, np_image, p: processing.StableDiffusionProcessing = None): + if np_image is None or hasattr(p, 'facehires'): + return np_image + self.load() + if self.model is None: + shared.log.error(f"Model load: type=FaceHires model={self.model_name} dir={self.model_dir} url={self.model_url}") + return np_image + image = Image.fromarray(np_image) + faces = self.predict(image, mask=True, device=devices.device, offload=shared.opts.face_restoration_unload) + if len(faces) == 0: + return np_image + + # create backups + orig_apply_overlay = shared.opts.mask_apply_overlay + orig_p = p.__dict__.copy() + orig_cls = p.__class__ + + pp = None + p.facehires = True # set flag to avoid recursion + shared.opts.data['mask_apply_overlay'] = True + p = processing_class.switch_class(p, processing.StableDiffusionProcessingImg2Img) + + for face in faces: + if face.mask is None: + continue + if face.size < 0.0002 or face.size > 0.8: + shared.log.debug(f'Face HiRes skip: {face.__dict__}') + continue + p.init_images = [image] + p.image_mask = [face.mask] + p.inpaint_full_res = True + p.inpainting_mask_invert = 0 + p.inpainting_fill = 1 # no fill + p.denoising_strength = orig_p.get('denoising_strength', 0.3) + # TODO facehires expose as tunable + p.mask_blur = 10 + p.inpaint_full_res_padding = 15 + p.restore_faces = True + shared.log.debug(f'Face HiRes: {face.__dict__} strength={p.denoising_strength} blur={p.mask_blur} padding={p.inpaint_full_res_padding}') + pp = processing.process_images_inner(p) + p.overlay_images = None # skip applying overlay twice + if pp is not None and pp.images is not None and len(pp.images) > 0: + image = pp.images[0] + + # restore pipeline + p = processing_class.switch_class(p, orig_cls, orig_p) + shared.opts.data['mask_apply_overlay'] = orig_apply_overlay + if pp is not None and pp.images is not None and len(pp.images) > 0: + image = pp.images[0] + np_image = np.array(image) + return np_image + + +yolo = FaceRestorerYolo() +shared.face_restorers.append(yolo) diff --git a/wiki b/wiki index 3051fdf2f..b707b4e2b 160000 --- a/wiki +++ b/wiki @@ -1 +1 @@ -Subproject commit 3051fdf2ff096ed8a8a98e2e5152383f4f9064db +Subproject commit b707b4e2b588d5d4809a9e97491acfcc66acc40d