diff --git a/modules/sd_hijack.py b/modules/sd_hijack.py index 3d340fc9b21..14fe62c73cb 100644 --- a/modules/sd_hijack.py +++ b/modules/sd_hijack.py @@ -38,6 +38,10 @@ optimizers = [] current_optimizer: sd_hijack_optimizations.SdOptimization = None +ldm_original_forward = patches.patch(__file__, ldm.modules.diffusionmodules.openaimodel.UNetModel, "forward", sd_unet.UNetModel_forward) +sgm_original_forward = patches.patch(__file__, sgm.modules.diffusionmodules.openaimodel.UNetModel, "forward", sd_unet.UNetModel_forward) + + def list_optimizers(): new_optimizers = script_callbacks.list_optimizers_callback() @@ -255,9 +259,6 @@ def flatten(el): import modules.models.diffusion.ddpm_edit - ldm_original_forward = patches.patch(__file__, ldm.modules.diffusionmodules.openaimodel.UNetModel, "forward", sd_unet.UNetModel_forward) - sgm_original_forward = patches.patch(__file__, sgm.modules.diffusionmodules.openaimodel.UNetModel, "forward", sd_unet.UNetModel_forward) - if isinstance(m, ldm.models.diffusion.ddpm.LatentDiffusion): sd_unet.original_forward = ldm_original_forward elif isinstance(m, modules.models.diffusion.ddpm_edit.LatentDiffusion): @@ -303,11 +304,6 @@ def undo_hijack(self, m): self.layers = None self.clip = None - patches.undo(__file__, ldm.modules.diffusionmodules.openaimodel.UNetModel, "forward") - patches.undo(__file__, sgm.modules.diffusionmodules.openaimodel.UNetModel, "forward") - - sd_unet.original_forward = None - def apply_circular(self, enable): if self.circular_enabled == enable: