From aafe7174d82577da7b2a94246c96bc47d12af431 Mon Sep 17 00:00:00 2001 From: lllyasviel Date: Sun, 28 Jan 2024 10:09:20 -0800 Subject: [PATCH] Create diffusers_patcher.py --- modules_forge/diffusers_patcher.py | 36 ++++++++++++++++++++++++++++++ 1 file changed, 36 insertions(+) create mode 100644 modules_forge/diffusers_patcher.py diff --git a/modules_forge/diffusers_patcher.py b/modules_forge/diffusers_patcher.py new file mode 100644 index 00000000..643f74d7 --- /dev/null +++ b/modules_forge/diffusers_patcher.py @@ -0,0 +1,36 @@ +import torch +import ldm_patched.modules.ops as ops +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: + def __init__(self, pipeline_class, dtype=torch.float16, *args, **kwargs): + load_device = model_management.get_torch_device() + offload_device = torch.device("cpu") + + if not model_management.should_use_fp16(device=load_device): + dtype = torch.float32 + + self.dtype = dtype + + with use_patched_ops(ops.manual_cast): + with modeling_utils.no_init_weights(): + self.pipeline = pipeline_class.from_pretrained(*args, **kwargs).to(dtype) + + self.pipeline.eval() + + self.patcher = ModelPatcher( + model=self.pipeline, + load_device=load_device, + offload_device=offload_device) + + def prepare_memory_before_sampling(self, batchsize, latent_width, latent_height): + area = 2 * batchsize * latent_width * latent_height + inference_memory = (((area * 0.6) / 0.9) + 1024) * (1024 * 1024) + model_management.load_models_gpu( + models=[self.pipeline], + memory_required=inference_memory + ) \ No newline at end of file