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("