From 8ec474c75be06510112429274a7972bfbc92528e Mon Sep 17 00:00:00 2001 From: lllyasviel Date: Sun, 28 Jan 2024 10:19:19 -0800 Subject: [PATCH] Update diffusers_patcher.py --- modules_forge/diffusers_patcher.py | 9 +++++++-- 1 file changed, 7 insertions(+), 2 deletions(-) diff --git a/modules_forge/diffusers_patcher.py b/modules_forge/diffusers_patcher.py index 9dbf11d7..262bad88 100644 --- a/modules_forge/diffusers_patcher.py +++ b/modules_forge/diffusers_patcher.py @@ -1,13 +1,14 @@ -import diffusers import torch import ldm_patched.modules.ops as ops + +from diffusers.models.attention_processor import AttnProcessor2_0 from ldm_patched.modules.model_patcher import ModelPatcher from ldm_patched.modules import model_management from modules_forge.ops import use_patched_ops from transformers import modeling_utils -class DiffusersPatcher: +class DiffusersModelPatcher: def __init__(self, pipeline_class, dtype=torch.float16, *args, **kwargs): load_device = model_management.get_torch_device() offload_device = torch.device("cpu") @@ -21,6 +22,10 @@ class DiffusersPatcher: with modeling_utils.no_init_weights(): self.pipeline = pipeline_class.from_pretrained(*args, **kwargs) + if hasattr(self.pipeline, 'unet'): + self.pipeline.unet.set_attn_processor(AttnProcessor2_0()) + print('Attention optimization applied to DiffusersModelPatcher') + self.pipeline = self.pipeline.to(device=offload_device, dtype=dtype) self.pipeline.eval()