diff --git a/extensions-builtin/forge_preprocessor_inpaint/scripts/preprocessor_inpaint.py b/extensions-builtin/forge_preprocessor_inpaint/scripts/preprocessor_inpaint.py index e679d33e..d3a2a486 100644 --- a/extensions-builtin/forge_preprocessor_inpaint/scripts/preprocessor_inpaint.py +++ b/extensions-builtin/forge_preprocessor_inpaint/scripts/preprocessor_inpaint.py @@ -28,6 +28,7 @@ class PreprocessorInpaintOnly(PreprocessorInpaint): self.name = 'inpaint_only' self.image = None self.mask = None + self.latent = None def process_before_every_sampling(self, process, cond, *args, **kwargs): self.image = kwargs['cond_before_inpaint_fix'][:, 0:3] @@ -55,6 +56,8 @@ class PreprocessorInpaintOnly(PreprocessorInpaint): unet.set_model_sampler_post_cfg_function(post_cfg) process.sd_model.forge_objects.unet = unet + + self.latent = latent_image return def process_after_every_sampling(self, process, params, *args, **kwargs): @@ -128,6 +131,13 @@ class PreprocessorInpaintLama(PreprocessorInpaintOnly): result = np.concatenate([fin_color, raw_mask], axis=2) return result + def process_before_every_sampling(self, process, cond, *args, **kwargs): + super().process_before_every_sampling(process, cond, *args, **kwargs) + sigma_max = process.sd_model.forge_objects.unet.model.model_sampling.sigma_max + original_noise = kwargs['noise'] + process.modified_noise = original_noise + self.latent.to(original_noise) / sigma_max.to(original_noise) + return + add_supported_preprocessor(PreprocessorInpaint()) diff --git a/modules/processing.py b/modules/processing.py index 775b0e57..a86d9b7f 100644 --- a/modules/processing.py +++ b/modules/processing.py @@ -257,6 +257,7 @@ class StableDiffusionProcessing: self.cached_c = StableDiffusionProcessing.cached_c self.extra_result_images = [] + self.modified_noise = None @property def sd_model(self): @@ -1256,6 +1257,10 @@ class StableDiffusionProcessingTxt2Img(StableDiffusionProcessing): c=conditioning, uc=unconditional_conditioning) + if self.modified_noise is not None: + x = self.modified_noise + self.modified_noise = None + samples = self.sampler.sample(self, x, conditioning, unconditional_conditioning, image_conditioning=self.txt2img_image_conditioning(x)) del x @@ -1365,6 +1370,10 @@ class StableDiffusionProcessingTxt2Img(StableDiffusionProcessing): c=self.hr_c, uc=self.hr_uc) + if self.modified_noise is not None: + noise = self.modified_noise + self.modified_noise = None + samples = self.sampler.sample_img2img(self, samples, noise, self.hr_c, self.hr_uc, steps=self.hr_second_pass_steps or self.steps, image_conditioning=image_conditioning) sd_models.apply_token_merging(self.sd_model, self.get_token_merging_ratio()) @@ -1676,6 +1685,10 @@ class StableDiffusionProcessingImg2Img(StableDiffusionProcessing): c=conditioning, uc=unconditional_conditioning) + if self.modified_noise is not None: + x = self.modified_noise + self.modified_noise = None + samples = self.sampler.sample_img2img(self, self.init_latent, x, conditioning, unconditional_conditioning, image_conditioning=self.image_conditioning) if self.mask is not None: