fluxdev/stable-diffusion-webui-forge

1import torch2import ldm_patched.modules.ops as ops3 4from ldm_patched.modules.model_patcher import ModelPatcher5from ldm_patched.modules import model_management6from transformers import modeling_utils7 8 9class DiffusersModelPatcher:10    def __init__(self, pipeline_class, dtype=torch.float16, *args, **kwargs):11        load_device = model_management.get_torch_device()12        offload_device = torch.device("cpu")13 14        if not model_management.should_use_fp16(device=load_device):15            dtype = torch.float3216 17        self.dtype = dtype18 19        with ops.use_patched_ops(ops.manual_cast):20            with modeling_utils.no_init_weights():21                self.pipeline = pipeline_class.from_pretrained(*args, **kwargs)22 23        if hasattr(self.pipeline, 'unet'):24            if hasattr(self.pipeline.unet, 'set_attn_processor'):25                from diffusers.models.attention_processor import AttnProcessor2_026                self.pipeline.unet.set_attn_processor(AttnProcessor2_0())27                print('Attention optimization applied to DiffusersModelPatcher')28 29        self.pipeline = self.pipeline.to(device=offload_device)30 31        if self.dtype == torch.float16:32            self.pipeline = self.pipeline.half()33 34        self.pipeline.eval()35 36        self.patcher = ModelPatcher(37            model=self.pipeline,38            load_device=load_device,39            offload_device=offload_device)40 41    def prepare_memory_before_sampling(self, batchsize, latent_width, latent_height):42        area = 2 * batchsize * latent_width * latent_height43        inference_memory = (((area * 0.6) / 0.9) + 1024) * (1024 * 1024)44        model_management.load_models_gpu(45            models=[self.patcher],46            memory_required=inference_memory47        )48 49    def move_tensor_to_current_device(self, x):50        return x.to(device=self.patcher.current_device, dtype=self.dtype)51