diff --git a/modules/processing.py b/modules/processing.py index ffc3d1ac..658de32e 100644 --- a/modules/processing.py +++ b/modules/processing.py @@ -909,8 +909,19 @@ def process_images_inner(p: StableDiffusionProcessing) -> Processed: p.extra_generation_params['Noise Schedule'] = opts.sd_noise_schedule p.sd_model.alphas_cumprod = rescale_zero_terminal_snr_abar(p.sd_model.alphas_cumprod).to(shared.device) + alphas_cumprod_modifiers = p.sd_model.forge_objects.unet.model_options.get('alphas_cumprod_modifiers', []) + alphas_cumprod_backup = None + + if len(alphas_cumprod_modifiers) > 0: + alphas_cumprod_backup = p.sd_model.alphas_cumprod + for modifier in alphas_cumprod_modifiers: + p.sd_model.alphas_cumprod = modifier(p.sd_model.alphas_cumprod) + samples_ddim = p.sample(conditioning=p.c, unconditional_conditioning=p.uc, seeds=p.seeds, subseeds=p.subseeds, subseed_strength=p.subseed_strength, prompts=p.prompts) + if alphas_cumprod_backup is not None: + p.sd_model.alphas_cumprod = alphas_cumprod_backup + if p.scripts is not None: ps = scripts.PostSampleArgs(samples_ddim) p.scripts.post_sample(p, ps) diff --git a/modules_forge/unet_patcher.py b/modules_forge/unet_patcher.py index e6992206..9e832d1c 100644 --- a/modules_forge/unet_patcher.py +++ b/modules_forge/unet_patcher.py @@ -102,6 +102,10 @@ class UnetPatcher(ModelPatcher): self.append_model_option('conditioning_modifiers', modifier, ensure_uniqueness) return + def add_alphas_cumprod_modifier(self, modifier, ensure_uniqueness=False): + self.append_model_option('alphas_cumprod_modifiers', modifier, ensure_uniqueness) + return + def add_block_modifier(self, modifier, ensure_uniqueness=False): self.append_transformer_option('block_modifiers', modifier, ensure_uniqueness) return