fluxdev/stable-diffusion-webui-forge

1import copy2import torch3 4from ldm_patched.modules.model_patcher import ModelPatcher5from ldm_patched.modules.sample import convert_cond6from ldm_patched.modules.samplers import encode_model_conds7 8 9class UnetPatcher(ModelPatcher):10    def __init__(self, *args, **kwargs):11        super().__init__(*args, **kwargs)12        self.controlnet_linked_list = None13        self.extra_preserved_memory_during_sampling = 014        self.extra_model_patchers_during_sampling = []15        self.extra_concat_condition = None16 17    def clone(self):18        n = UnetPatcher(self.model, self.load_device, self.offload_device, self.size, self.current_device,19                        weight_inplace_update=self.weight_inplace_update)20 21        n.patches = {}22        for k in self.patches:23            n.patches[k] = self.patches[k][:]24 25        n.object_patches = self.object_patches.copy()26        n.model_options = copy.deepcopy(self.model_options)27        n.model_keys = self.model_keys28        n.controlnet_linked_list = self.controlnet_linked_list29        n.extra_preserved_memory_during_sampling = self.extra_preserved_memory_during_sampling30        n.extra_model_patchers_during_sampling = self.extra_model_patchers_during_sampling.copy()31        n.extra_concat_condition = self.extra_concat_condition32        return n33 34    def add_extra_preserved_memory_during_sampling(self, memory_in_bytes: int):35        # Use this to ask Forge to preserve a certain amount of memory during sampling.36        # If GPU VRAM is 8 GB, and memory_in_bytes is 2GB, i.e., memory_in_bytes = 2 * 1024 * 1024 * 102437        # Then the sampling will always use less than 6GB memory by dynamically offload modules to CPU RAM.38        # You can estimate this using model_management.module_size(any_pytorch_model) to get size of any pytorch models.39        self.extra_preserved_memory_during_sampling += memory_in_bytes40        return41 42    def add_extra_model_patcher_during_sampling(self, model_patcher: ModelPatcher):43        # Use this to ask Forge to move extra model patchers to GPU during sampling.44        # This method will manage GPU memory perfectly.45        self.extra_model_patchers_during_sampling.append(model_patcher)46        return47 48    def add_extra_torch_module_during_sampling(self, m: torch.nn.Module, cast_to_unet_dtype: bool = True):49        # Use this method to bind an extra torch.nn.Module to this UNet during sampling.50        # This model `m` will be delegated to Forge memory management system.51        # `m` will be loaded to GPU everytime when sampling starts.52        # `m` will be unloaded if necessary.53        # `m` will influence Forge's judgement about use GPU memory or54        # capacity and decide whether to use module offload to make user's batch size larger.55        # Use cast_to_unet_dtype if you want `m` to have same dtype with unet during sampling.56 57        if cast_to_unet_dtype:58            m.to(self.model.diffusion_model.dtype)59 60        patcher = ModelPatcher(model=m, load_device=self.load_device, offload_device=self.offload_device)61 62        self.add_extra_model_patcher_during_sampling(patcher)63        return patcher64 65    def add_patched_controlnet(self, cnet):66        cnet.set_previous_controlnet(self.controlnet_linked_list)67        self.controlnet_linked_list = cnet68        return69 70    def list_controlnets(self):71        results = []72        pointer = self.controlnet_linked_list73        while pointer is not None:74            results.append(pointer)75            pointer = pointer.previous_controlnet76        return results77 78    def append_model_option(self, k, v, ensure_uniqueness=False):79        if k not in self.model_options:80            self.model_options[k] = []81 82        if ensure_uniqueness and v in self.model_options[k]:83            return84 85        self.model_options[k].append(v)86        return87 88    def append_transformer_option(self, k, v, ensure_uniqueness=False):89        if 'transformer_options' not in self.model_options:90            self.model_options['transformer_options'] = {}91 92        to = self.model_options['transformer_options']93 94        if k not in to:95            to[k] = []96 97        if ensure_uniqueness and v in to[k]:98            return99 100        to[k].append(v)101        return102 103    def set_transformer_option(self, k, v):104        if 'transformer_options' not in self.model_options:105            self.model_options['transformer_options'] = {}106 107        self.model_options['transformer_options'][k] = v108        return109 110    def add_conditioning_modifier(self, modifier, ensure_uniqueness=False):111        self.append_model_option('conditioning_modifiers', modifier, ensure_uniqueness)112        return113 114    def add_sampler_pre_cfg_function(self, modifier, ensure_uniqueness=False):115        self.append_model_option('sampler_pre_cfg_function', modifier, ensure_uniqueness)116        return117 118    def set_memory_peak_estimation_modifier(self, modifier):119        self.model_options['memory_peak_estimation_modifier'] = modifier120        return121 122    def add_alphas_cumprod_modifier(self, modifier, ensure_uniqueness=False):123        """124 125        For some reasons, this function only works in A1111's Script.process_batch(self, p, *args, **kwargs)126 127        For example, below is a worked modification:128 129        class ExampleScript(scripts.Script):130 131            def process_batch(self, p, *args, **kwargs):132                unet = p.sd_model.forge_objects.unet.clone()133 134                def modifier(x):135                    return x ** 0.5136 137                unet.add_alphas_cumprod_modifier(modifier)138                p.sd_model.forge_objects.unet = unet139 140                return141 142        This add_alphas_cumprod_modifier is the only patch option that should be used in process_batch()143        All other patch options should be called in process_before_every_sampling()144 145        """146 147        self.append_model_option('alphas_cumprod_modifiers', modifier, ensure_uniqueness)148        return149 150    def add_block_modifier(self, modifier, ensure_uniqueness=False):151        self.append_transformer_option('block_modifiers', modifier, ensure_uniqueness)152        return153 154    def add_block_inner_modifier(self, modifier, ensure_uniqueness=False):155        self.append_transformer_option('block_inner_modifiers', modifier, ensure_uniqueness)156        return157 158    def add_controlnet_conditioning_modifier(self, modifier, ensure_uniqueness=False):159        self.append_transformer_option('controlnet_conditioning_modifiers', modifier, ensure_uniqueness)160        return161 162    def set_controlnet_model_function_wrapper(self, wrapper):163        self.set_transformer_option('controlnet_model_function_wrapper', wrapper)164        return165 166    def set_model_replace_all(self, patch, target="attn1"):167        for block_name in ['input', 'middle', 'output']:168            for number in range(16):169                for transformer_index in range(16):170                    self.set_model_patch_replace(patch, target, block_name, number, transformer_index)171        return172 173    def encode_conds_after_clip(self, conds, noise, prompt_type="positive"):174        return encode_model_conds(175            model_function=self.model.extra_conds,176            conds=convert_cond(conds),177            noise=noise,178            device=noise.device,179            prompt_type=prompt_type180        )181 182    def load_frozen_patcher(self, state_dict, strength):183        patch_dict = {}184        for k, w in state_dict.items():185            model_key, patch_type, weight_index = k.split('::')186            if model_key not in patch_dict:187                patch_dict[model_key] = {}188            if patch_type not in patch_dict[model_key]:189                patch_dict[model_key][patch_type] = [None] * 16190            patch_dict[model_key][patch_type][int(weight_index)] = w191 192        patch_flat = {}193        for model_key, v in patch_dict.items():194            for patch_type, weight_list in v.items():195                patch_flat[model_key] = (patch_type, weight_list)196 197        self.add_patches(patches=patch_flat, strength_patch=float(strength), strength_model=1.0)198        return199