def __init__(self):

self.models = {}

def start(self):

self.end()

def add(self, outputs):

if isinstance(outputs, collections.abc.Mapping):

outputs = outputs.values()

elif not isinstance(outputs, (list, tuple)):

outputs = (outputs,)

for output in outputs:

if isinstance(output, (collections.abc.Mapping, list, tuple)):

self.add(output)

continue

models = []

if isinstance(output, ModelPatcher):

models.append(output)

models.extend(output.model_patches_models())

models.extend(output.get_nested_additional_models())

else:

patcher = getattr(output, "patcher", None)

if isinstance(patcher, ModelPatcher):

models.append(patcher)

get_models = getattr(output, "get_models", None)

if callable(get_models):

models.extend(get_models())

for model in models:

if not isinstance(model, ModelPatcher) or not model.is_dynamic():

continue

key = (id(model.model), model.load_device)

self.models[key] = model

model.set_in_use_by_current_prompt(True)

def end(self):

for model in self.models.values():

model.set_in_use_by_current_prompt(False)

to = model_options["transformer_options"].copy()

if "patches_replace" not in to:

to["patches_replace"] = {}

else:

to["patches_replace"] = to["patches_replace"].copy()

if name not in to["patches_replace"]:

to["patches_replace"][name] = {}

else:

to["patches_replace"][name] = to["patches_replace"][name].copy()

if transformer_index is not None:

block = (block_name, number, transformer_index)

else:

block = (block_name, number)

to["patches_replace"][name][block] = patch

model_options["transformer_options"] = to

model_options["sampler_post_cfg_function"] = model_options.get("sampler_post_cfg_function", []) + [post_cfg_function]

if disable_cfg1_optimization:

model_options["disable_cfg1_optimization"] = True

model_options["sampler_pre_cfg_function"] = model_options.get("sampler_pre_cfg_function", []) + [pre_cfg_function]

if disable_cfg1_optimization:

model_options["disable_cfg1_optimization"] = True

new_hook_patches = {}

for hook_ref in orig_hook_patches:

new_hook_patches[hook_ref] = {}

for k in orig_hook_patches[hook_ref]:

new_hook_patches[hook_ref][k] = orig_hook_patches[hook_ref][k][:]

if copy_tuples:

for i in range(len(new_hook_patches[hook_ref][k])):

new_hook_patches[hook_ref][k][i] = tuple(new_hook_patches[hook_ref][k][i])

if hasattr(m, "prev_comfy_cast_weights"):

m.comfy_cast_weights = m.prev_comfy_cast_weights

del m.prev_comfy_cast_weights

if hasattr(m, "weight_function"):

m.weight_function = []

if hasattr(m, "bias_function"):

if device is None:

return 0

memory = 0

if hasattr(m, "weight_function"):

for f in m.weight_function:

if hasattr(f, "move_to"):

memory += f.move_to(device=device)

if hasattr(m, "bias_function"):

for f in m.bias_function:

if hasattr(f, "move_to"):

memory += f.move_to(device=device)

is_lowvram_patch = True

def __init__(self, key, patches, convert_func=None, set_func=None):

self.key = key

self.patches = patches

self.convert_func = convert_func # TODO: remove

self.set_func = set_func

self.prepared_patches = None

def memory_required(self):

counter = [0]

for patch in self.patches[self.key]:

comfy.lora.prefetch_prepared_value(patch[1], counter, None, None, False)

return counter[0]

def prepare(self, destination, stream, copy=True, commit=True):

counter = [0]

prepared_patches = [

(patch[0], comfy.lora.prefetch_prepared_value(patch[1], counter, destination, stream, copy), patch[2], patch[3], patch[4])

for patch in self.patches[self.key]

]

if commit:

self.prepared_patches = prepared_patches

return prepared_patches

def clear_prepared(self):

self.prepared_patches = None

def __call__(self, weight):

patches = self.prepared_patches if self.prepared_patches is not None else self.patches[self.key]

weight, set_func, convert_func = get_key_weight(model, key)

if weight is None:

return 0

model_dtype = getattr(model, "manual_cast_dtype", torch.float32)

if model_dtype is None:

model_dtype = weight.dtype

set_func = None

convert_func = None

op_keys = key.rsplit('.', 1)

if len(op_keys) < 2:

weight = comfy.utils.get_attr(model, key)

else:

op = comfy.utils.get_attr(model, op_keys[0])

try:

set_func = getattr(op, "set_{}".format(op_keys[1]))

except AttributeError:

pass

try:

convert_func = getattr(op, "convert_{}".format(op_keys[1]))

except AttributeError:

pass

weight = getattr(op, op_keys[1])

if convert_func is not None:

weight = comfy.utils.get_attr(model, key)

def __init__(self, model: 'ModelPatcher', skip_and_inject_on_exit_only=False):

self.model = model

self.was_injected = False

self.prev_skip_injection = False

self.skip_and_inject_on_exit_only = skip_and_inject_on_exit_only

def __enter__(self):

self.was_injected = False

self.prev_skip_injection = self.model.skip_injection

if self.skip_and_inject_on_exit_only:

self.model.skip_injection = True

if self.model.is_injected:

self.model.eject_model()

self.was_injected = True

def __exit__(self, *args):

if self.skip_and_inject_on_exit_only:

self.model.skip_injection = self.prev_skip_injection

self.model.inject_model()

if self.was_injected and not self.model.skip_injection:

self.model.inject_model()

def __init__(self, initial: int, minimum=0):

self.value = initial

self.minimum = minimum

# TODO: add a safe limit besides 0

def use(self, weight: torch.Tensor):

weight_size = weight.nelement() * weight.element_size()

if self.is_useable(weight_size):

self.decrement(weight_size)

return True

return False

def is_useable(self, used: int):

return self.value - used > self.minimum

def decrement(self, used: int):

def __new__(cls, model, key, tensor):

return super().__new__(cls, tensor)

def __init__(self, model, key, tensor):

self.model = model

self.key = key

@property

def device(self):

return CustomTorchDevice

#safetensors will .to() us to the cpu which we catch here to cast on demand. The returned tensor is

#then just a short lived thing in the safetensors serialization logic inside its big for loop over

#all weights getting garbage collected per-weight

def to(self, *args, **kwargs):

def __init__(self, model, key):

self.model = model

self.key = key

self.cpu_state_dict = None

def state_dict_tensor(self, state_dict_key):

if self.cpu_state_dict is None:

weight = self.model.patch_weight_to_device(self.key, device_to=self.model.load_device, return_weight=True)

self.cpu_state_dict = {k: v.to("cpu") for k, v in weight.state_dict(self.key).items()}

def __new__(cls, caster, state_dict_key, tensor):

return super().__new__(cls, tensor)

def __init__(self, caster, state_dict_key, tensor):

self.caster = caster

self.state_dict_key = state_dict_key

@property

def device(self):

return CustomTorchDevice

def to(self, *args, **kwargs):

caster = self.caster

del self.caster

def __init__(self, model, load_device, offload_device, size=0, weight_inplace_update=False, fast_disk=False):

self.size = size

self.model = model

self.fast_disk = fast_disk

if not hasattr(self.model, 'device'):

logging.debug("Model doesn't have a device attribute.")

self.model.device = offload_device

elif self.model.device is None:

self.model.device = offload_device

self.patches = {}

self.backup = {}

self.backup_buffers = {}

self.object_patches = {}

self.object_patches_backup = {}

self.weight_wrapper_patches = {}

self.model_options = {"transformer_options":{}}

self.load_device = load_device

self.offload_device = offload_device

self.weight_inplace_update = weight_inplace_update

self.force_cast_weights = False

self.patches_uuid = uuid.uuid4()

self.parent = None

self.pinned = set()

self.attachments: dict[str] = {}

self.additional_models: dict[str, list[ModelPatcher]] = {}

self.callbacks: dict[str, dict[str, list[Callable]]] = CallbacksMP.init_callbacks()

self.wrappers: dict[str, dict[str, list[Callable]]] = WrappersMP.init_wrappers()

self.is_injected = False

self.skip_injection = False

self.injections: dict[str, list[PatcherInjection]] = {}

self.hook_patches: dict[comfy.hooks._HookRef] = {}

self.hook_patches_backup: dict[comfy.hooks._HookRef] = None

self.hook_backup: dict[str, tuple[torch.Tensor, torch.device]] = {}

self.cached_hook_patches: dict[comfy.hooks.HookGroup, dict[str, torch.Tensor]] = {}

self.current_hooks: Optional[comfy.hooks.HookGroup] = None

self.forced_hooks: Optional[comfy.hooks.HookGroup] = None # NOTE: only used for CLIP at this time

self.is_clip = False

self.hook_mode = comfy.hooks.EnumHookMode.MaxSpeed

self.cached_patcher_init: tuple[Callable, tuple] | tuple[Callable, tuple, int] | None = None

self.is_multigpu_base_clone = False

self.clone_base_uuid = uuid.uuid4()

if not hasattr(self.model, 'model_loaded_weight_memory'):

self.model.model_loaded_weight_memory = 0

if not hasattr(self.model, 'lowvram_patch_counter'):

self.model.lowvram_patch_counter = 0

if not hasattr(self.model, 'model_lowvram'):

self.model.model_lowvram = False

if not hasattr(self.model, 'current_weight_patches_uuid'):

self.model.current_weight_patches_uuid = None

if not hasattr(self.model, 'model_offload_buffer_memory'):

self.model.model_offload_buffer_memory = 0

def is_dynamic(self):

return False

def model_size(self):

if self.size > 0:

return self.size

self.size = comfy.model_management.module_size(self.model)

return self.size

def loaded_size(self):

return self.model.model_loaded_weight_memory

def lowvram_patch_counter(self):

return self.model.lowvram_patch_counter

def get_free_memory(self, device):

#Prioritize batching (incl. CFG/conds etc) over keeping the model resident. In

#the vast majority of setups a little bit of offloading on the giant model more

#than pays for CFG. So return everything both torch and Aimdo could give us

aimdo_mem = 0

if comfy.memory_management.aimdo_enabled:

aimdo_device = device.index if getattr(device, "type", None) == "cuda" else None

aimdo_mem = comfy_aimdo.model_vbar.vbars_analyze(aimdo_device)

return comfy.model_management.get_free_memory(device) + aimdo_mem

def get_clone_model_override(self):

return self.model, (self.backup, self.backup_buffers, self.object_patches_backup, self.pinned)

def clone(self, disable_dynamic=False, model_override=None, force_deepcopy=False):

class_ = self.__class__

if self.is_dynamic() and disable_dynamic or force_deepcopy:

if self.is_dynamic() and disable_dynamic:

class_ = ModelPatcher

if model_override is None:

if self.cached_patcher_init is None:

raise RuntimeError("Cannot create non-dynamic delegate: cached_patcher_init is not initialized.")

temp_model_patcher = self.cached_patcher_init[0](*self.cached_patcher_init[1], disable_dynamic=True)

if len(self.cached_patcher_init) > 2:

temp_model_patcher = temp_model_patcher[self.cached_patcher_init[2]]

model_override = temp_model_patcher.get_clone_model_override()

if model_override is None:

model_override = self.get_clone_model_override()

n = class_(model_override[0], self.load_device, self.offload_device, self.model_size(), weight_inplace_update=self.weight_inplace_update, fast_disk=self.fast_disk)

n.patches = {}

for k in self.patches:

n.patches[k] = self.patches[k][:]

n.patches_uuid = self.patches_uuid

n.object_patches = self.object_patches.copy()

n.weight_wrapper_patches = self.weight_wrapper_patches.copy()

n.model_options = comfy.utils.deepcopy_list_dict(self.model_options)

n.parent = self

n.force_cast_weights = self.force_cast_weights

n.backup, n.backup_buffers, n.object_patches_backup, n.pinned = model_override[1]

# attachments

n.attachments = {}

for k in self.attachments:

if hasattr(self.attachments[k], "on_model_patcher_clone"):

n.attachments[k] = self.attachments[k].on_model_patcher_clone()

else:

n.attachments[k] = self.attachments[k]

# additional models

for k, c in self.additional_models.items():

n.additional_models[k] = [x.clone() for x in c]

# callbacks

for k, c in self.callbacks.items():

n.callbacks[k] = {}

for k1, c1 in c.items():

n.callbacks[k][k1] = c1.copy()

# sample wrappers

for k, w in self.wrappers.items():

n.wrappers[k] = {}

for k1, w1 in w.items():

n.wrappers[k][k1] = w1.copy()

# injection

n.is_injected = self.is_injected

n.skip_injection = self.skip_injection

for k, i in self.injections.items():

n.injections[k] = i.copy()

# hooks

n.hook_patches = create_hook_patches_clone(self.hook_patches)

n.hook_patches_backup = create_hook_patches_clone(self.hook_patches_backup) if self.hook_patches_backup else self.hook_patches_backup

for group in self.cached_hook_patches:

n.cached_hook_patches[group] = {}

for k in self.cached_hook_patches[group]:

n.cached_hook_patches[group][k] = self.cached_hook_patches[group][k]

n.hook_backup = self.hook_backup

n.current_hooks = self.current_hooks.clone() if self.current_hooks else self.current_hooks

n.forced_hooks = self.forced_hooks.clone() if self.forced_hooks else self.forced_hooks

n.is_clip = self.is_clip

n.hook_mode = self.hook_mode

n.cached_patcher_init = self.cached_patcher_init

n.is_multigpu_base_clone = self.is_multigpu_base_clone

n.clone_base_uuid = self.clone_base_uuid

for callback in self.get_all_callbacks(CallbacksMP.ON_CLONE):

callback(self, n)

return n

def deepclone_multigpu(self, new_load_device=None, models_cache: dict[uuid.UUID,ModelPatcher]=None):

logging.info(f"Creating deepclone of {self.model.__class__.__name__} for {new_load_device if new_load_device else self.load_device}.")

if self.cached_patcher_init is None:

raise RuntimeError(

f"Cannot create multigpu deepclone of {self.model.__class__.__name__}: "

"the loader that produced this model does not support multigpu "

"(cached_patcher_init is not initialized). Use a core loader "

"(CheckpointLoaderSimple, UNETLoader, CLIPLoader/DualCLIPLoader, VAELoader), "

"or have the custom loader register a cached_patcher_init factory."

)

comfy.model_management.unload_model_and_clones(self)

# Produce a freshly-loaded patcher from the loader factory so the multigpu

# clone owns its own untainted model weights (rather than relying on

# copy.deepcopy of an already-patched/already-loaded module).

temp_model_patcher: ModelPatcher | list[ModelPatcher] = self.cached_patcher_init[0](*self.cached_patcher_init[1])

if len(self.cached_patcher_init) > 2:

temp_model_patcher = temp_model_patcher[self.cached_patcher_init[2]]

# Override clone()'s normal "share self.model + share backup containers" with

# the pristine model from temp_model_patcher plus empty backup containers --

# the fresh model has no patches applied, so any deepcopy of self's stale

# backup/object_patches_backup/pinned would just propagate dead state that

# no longer corresponds to anything in n.model.

model_override = (temp_model_patcher.model, ({}, {}, {}, set()))

n = self.clone(model_override=model_override)

# clone() copies hook_backup by reference from self; reset since model is pristine.

n.hook_backup = {}

# set load device, if present

if new_load_device is not None:

n.load_device = new_load_device

# Ensure any per-device bookkeeping (e.g. ModelPatcherDynamic.dynamic_pins)

# has an entry for n.load_device on the freshly-loaded n.model. temp_model_patcher's

# __init__ only registered its own (default) load_device.

if hasattr(n, "register_load_device"):

n.register_load_device(n.load_device)

# multigpu clone should not have multigpu additional_models entry

n.remove_additional_models("multigpu")

# multigpu_clone all stored additional_models; make sure circular references are properly handled

if models_cache is None:

models_cache = {}

for key, model_list in n.additional_models.items():

for i in range(len(model_list)):

add_model = n.additional_models[key][i]

if add_model.clone_base_uuid not in models_cache:

models_cache[add_model.clone_base_uuid] = add_model.deepclone_multigpu(new_load_device=new_load_device, models_cache=models_cache)

n.additional_models[key][i] = models_cache[add_model.clone_base_uuid]

for callback in self.get_all_callbacks(CallbacksMP.ON_DEEPCLONE_MULTIGPU):

callback(self, n)

return n

def match_multigpu_clones(self):

multigpu_models = self.get_additional_models_with_key("multigpu")

if len(multigpu_models) > 0:

new_multigpu_models = []

for mm in multigpu_models:

# clone main model, but bring over relevant props from existing multigpu clone

n = self.clone(model_override=mm.get_clone_model_override())

n.load_device = mm.load_device

n.hook_backup = mm.hook_backup

n.is_multigpu_base_clone = mm.is_multigpu_base_clone

n.remove_additional_models("multigpu")

orig_additional_models: dict[str, list[ModelPatcher]] = comfy.patcher_extension.copy_nested_dicts(n.additional_models)

n.additional_models = comfy.patcher_extension.copy_nested_dicts(mm.additional_models)

# figure out which additional models are not present in multigpu clone

models_cache = {}

for mm_add_model in mm.get_additional_models():

models_cache[mm_add_model.clone_base_uuid] = mm_add_model

remove_models_uuids = set(list(models_cache.keys()))

for key, model_list in orig_additional_models.items():

for orig_add_model in model_list:

if orig_add_model.clone_base_uuid not in models_cache:

models_cache[orig_add_model.clone_base_uuid] = orig_add_model.deepclone_multigpu(new_load_device=n.load_device, models_cache=models_cache)

existing_list = n.get_additional_models_with_key(key)

existing_list.append(models_cache[orig_add_model.clone_base_uuid])

n.set_additional_models(key, existing_list)

if orig_add_model.clone_base_uuid in remove_models_uuids:

remove_models_uuids.remove(orig_add_model.clone_base_uuid)

# remove duplicate additional models

for key, model_list in n.additional_models.items():

new_model_list = [x for x in model_list if x.clone_base_uuid not in remove_models_uuids]

n.set_additional_models(key, new_model_list)

for callback in self.get_all_callbacks(CallbacksMP.ON_MATCH_MULTIGPU_CLONES):

callback(self, n)

new_multigpu_models.append(n)

self.set_additional_models("multigpu", new_multigpu_models)

def is_clone(self, other):

if hasattr(other, 'model') and self.model is other.model:

return True

return False

def clone_has_same_weights(self, clone: ModelPatcher, allow_multigpu=False):

if allow_multigpu:

if self.clone_base_uuid != clone.clone_base_uuid:

return False

else:

if not self.is_clone(clone):

return False

if self.current_hooks != clone.current_hooks:

return False

if self.forced_hooks != clone.forced_hooks:

return False

if self.hook_patches.keys() != clone.hook_patches.keys():

return False

if self.attachments.keys() != clone.attachments.keys():

return False

if self.additional_models.keys() != clone.additional_models.keys():

return False

for key in self.callbacks:

if len(self.callbacks[key]) != len(clone.callbacks[key]):

return False

for key in self.wrappers:

if len(self.wrappers[key]) != len(clone.wrappers[key]):

return False

if self.injections.keys() != clone.injections.keys():

return False

if len(self.patches) == 0 and len(clone.patches) == 0:

return True

if self.patches_uuid == clone.patches_uuid:

if len(self.patches) != len(clone.patches):

logging.warning("WARNING: something went wrong, same patch uuid but different length of patches.")

else:

return True

def memory_required(self, input_shape):

return self.model.memory_required(input_shape=input_shape)

def disable_model_cfg1_optimization(self):

self.model_options["disable_cfg1_optimization"] = True

def set_model_sampler_cfg_function(self, sampler_cfg_function, disable_cfg1_optimization=False):

if len(inspect.signature(sampler_cfg_function).parameters) == 3:

self.model_options["sampler_cfg_function"] = lambda args: sampler_cfg_function(args["cond"], args["uncond"], args["cond_scale"]) #Old way

else:

self.model_options["sampler_cfg_function"] = sampler_cfg_function

if disable_cfg1_optimization:

self.disable_model_cfg1_optimization()

def set_model_sampler_post_cfg_function(self, post_cfg_function, disable_cfg1_optimization=False):

self.model_options = set_model_options_post_cfg_function(self.model_options, post_cfg_function, disable_cfg1_optimization)

def set_model_sampler_pre_cfg_function(self, pre_cfg_function, disable_cfg1_optimization=False):

self.model_options = set_model_options_pre_cfg_function(self.model_options, pre_cfg_function, disable_cfg1_optimization)

def set_model_sampler_calc_cond_batch_function(self, sampler_calc_cond_batch_function):

self.model_options["sampler_calc_cond_batch_function"] = sampler_calc_cond_batch_function

def set_model_unet_function_wrapper(self, unet_wrapper_function: UnetWrapperFunction):

self.model_options["model_function_wrapper"] = unet_wrapper_function

def set_model_denoise_mask_function(self, denoise_mask_function):

self.model_options["denoise_mask_function"] = denoise_mask_function

def set_model_patch(self, patch, name):

to = self.model_options["transformer_options"]

if "patches" not in to:

to["patches"] = {}

to["patches"][name] = to["patches"].get(name, []) + [patch]

def set_model_patch_replace(self, patch, name, block_name, number, transformer_index=None):

self.model_options = set_model_options_patch_replace(self.model_options, patch, name, block_name, number, transformer_index=transformer_index)

def set_model_attn1_patch(self, patch):

self.set_model_patch(patch, "attn1_patch")

def set_model_attn2_patch(self, patch):

self.set_model_patch(patch, "attn2_patch")

def set_model_attn1_replace(self, patch, block_name, number, transformer_index=None):

self.set_model_patch_replace(patch, "attn1", block_name, number, transformer_index)

def set_model_attn2_replace(self, patch, block_name, number, transformer_index=None):

self.set_model_patch_replace(patch, "attn2", block_name, number, transformer_index)

def set_model_attn1_output_patch(self, patch):

self.set_model_patch(patch, "attn1_output_patch")

def set_model_attn2_output_patch(self, patch):

self.set_model_patch(patch, "attn2_output_patch")

def set_model_optimized_attention(self, optimized_attention):

def optimized_attention_override(_, *args, **kwargs):

return optimized_attention(*args, **kwargs)

if hasattr(optimized_attention, "container_function") and optimized_attention.container_function is not None:

optimized_attention_override.container_function = optimized_attention.container_function

self.model_options["transformer_options"]["optimized_attention_override"] = optimized_attention_override

def set_model_input_block_patch(self, patch):

self.set_model_patch(patch, "input_block_patch")

def set_model_input_block_patch_after_skip(self, patch):

self.set_model_patch(patch, "input_block_patch_after_skip")

def set_model_output_block_patch(self, patch):

self.set_model_patch(patch, "output_block_patch")

def set_model_emb_patch(self, patch):

self.set_model_patch(patch, "emb_patch")

def set_model_forward_timestep_embed_patch(self, patch):

self.set_model_patch(patch, "forward_timestep_embed_patch")

def set_model_double_block_patch(self, patch):

self.set_model_patch(patch, "double_block")

def set_model_post_input_patch(self, patch):

self.set_model_patch(patch, "post_input")

def set_model_noise_refiner_patch(self, patch):

self.set_model_patch(patch, "noise_refiner")

def set_model_middle_block_after_patch(self, patch):

self.set_model_patch(patch, "middle_block_after_patch")

def set_model_rope_options(self, scale_x, shift_x, scale_y, shift_y, scale_t, shift_t, **kwargs):

rope_options = self.model_options["transformer_options"].get("rope_options", {})

rope_options["scale_x"] = scale_x

rope_options["scale_y"] = scale_y

rope_options["scale_t"] = scale_t

rope_options["shift_x"] = shift_x

rope_options["shift_y"] = shift_y

rope_options["shift_t"] = shift_t

self.model_options["transformer_options"]["rope_options"] = rope_options

def add_object_patch(self, name, obj):

self.object_patches[name] = obj

def set_model_compute_dtype(self, dtype):

self.add_object_patch("manual_cast_dtype", dtype)

if dtype is not None:

self.force_cast_weights = True

self.patches_uuid = uuid.uuid4() #TODO: optimize by preventing a full model reload for this

def add_weight_wrapper(self, name, function):

self.weight_wrapper_patches[name] = self.weight_wrapper_patches.get(name, []) + [function]

self.patches_uuid = uuid.uuid4()

def get_model_object(self, name: str) -> torch.nn.Module:

"""Retrieves a nested attribute from an object using dot notation considering

if name in self.object_patches:

return self.object_patches[name]

else:

if name in self.object_patches_backup:

return self.object_patches_backup[name]

else:

return comfy.utils.get_attr(self.model, name)

def model_patches_to(self, device):

to = self.model_options["transformer_options"]

if "patches" in to:

patches = to["patches"]

for name in patches:

patch_list = patches[name]

for i in range(len(patch_list)):

if hasattr(patch_list[i], "to"):

patch_list[i] = patch_list[i].to(device)

if "patches_replace" in to:

patches = to["patches_replace"]

for name in patches:

patch_list = patches[name]

for k in patch_list:

if hasattr(patch_list[k], "to"):

patch_list[k] = patch_list[k].to(device)

if "model_function_wrapper" in self.model_options:

wrap_func = self.model_options["model_function_wrapper"]

if hasattr(wrap_func, "to"):

self.model_options["model_function_wrapper"] = wrap_func.to(device)

def model_patches_models(self):

to = self.model_options["transformer_options"]

models = []

if "patches" in to:

patches = to["patches"]

for name in patches:

patch_list = patches[name]

for i in range(len(patch_list)):

if hasattr(patch_list[i], "models"):

models += patch_list[i].models()

if "patches_replace" in to:

patches = to["patches_replace"]

for name in patches:

patch_list = patches[name]

for k in patch_list:

if hasattr(patch_list[k], "models"):

models += patch_list[k].models()

if "model_function_wrapper" in self.model_options:

wrap_func = self.model_options["model_function_wrapper"]

if hasattr(wrap_func, "models"):

models += wrap_func.models()

return models

def model_patches_call_function(self, function_name="cleanup", arguments={}):

to = self.model_options["transformer_options"]

if "patches" in to:

patches = to["patches"]

for name in patches:

patch_list = patches[name]

for i in range(len(patch_list)):

if hasattr(patch_list[i], function_name):

getattr(patch_list[i], function_name)(**arguments)

if "patches_replace" in to:

patches = to["patches_replace"]

for name in patches:

patch_list = patches[name]

for k in patch_list:

if hasattr(patch_list[k], function_name):

getattr(patch_list[k], function_name)(**arguments)

if "model_function_wrapper" in self.model_options:

wrap_func = self.model_options["model_function_wrapper"]

if hasattr(wrap_func, function_name):

getattr(wrap_func, function_name)(**arguments)

def model_dtype(self):

if hasattr(self.model, "get_dtype"):

return self.model.get_dtype()

def add_patches(self, patches, strength_patch=1.0, strength_model=1.0):

with self.use_ejected():

p = set()

model_sd = self.model.state_dict()

for k in patches:

offset = None

function = None

if isinstance(k, str):

key = k

else:

offset = k[1]

key = k[0]

if len(k) > 2:

function = k[2]

if key in model_sd:

p.add(k)

current_patches = self.patches.get(key, [])

current_patches.append((strength_patch, patches[k], strength_model, offset, function))

self.patches[key] = current_patches

self.patches_uuid = uuid.uuid4()

return list(p)

def get_key_patches(self, filter_prefix=None):

model_sd = self.model_state_dict()

p = {}

for k in model_sd:

if filter_prefix is not None:

if not k.startswith(filter_prefix):

continue

bk = self.backup.get(k, None)

hbk = self.hook_backup.get(k, None)

weight, set_func, convert_func = get_key_weight(self.model, k)

if not isinstance(weight, torch.Tensor):

continue

if bk is not None:

weight = bk.weight

if hbk is not None:

weight = hbk[0]

if convert_func is None:

convert_func = lambda a, **kwargs: a

if k in self.patches:

p[k] = [(weight, convert_func)] + self.patches[k]

else:

p[k] = [(weight, convert_func)]

return p

def model_state_dict(self, filter_prefix=None):

with self.use_ejected():

sd = self.model.state_dict()

keys = list(sd.keys())

if filter_prefix is not None:

for k in keys:

if not k.startswith(filter_prefix):

sd.pop(k)

return sd

def patch_weight_to_device(self, key, device_to=None, inplace_update=False, return_weight=False, force_cast=False):

weight, set_func, convert_func = get_key_weight(self.model, key)

if key not in self.patches and not force_cast:

return weight

inplace_update = self.weight_inplace_update or inplace_update

if key not in self.backup and not return_weight:

self.backup[key] = collections.namedtuple('Dimension', ['weight', 'inplace_update'])(weight.to(device=self.offload_device, copy=inplace_update), inplace_update)

temp_dtype = comfy.model_management.lora_compute_dtype(device_to) if key in self.patches else None

if device_to is not None:

temp_weight = comfy.model_management.cast_to_device(weight, device_to, temp_dtype, copy=True)

else:

temp_weight = weight.to(temp_dtype, copy=True)

if convert_func is not None:

temp_weight = convert_func(temp_weight, inplace=True)

out_weight = comfy.lora.calculate_weight(self.patches[key], temp_weight, key) if key in self.patches else temp_weight

if set_func is None:

if key in self.patches:

out_weight = comfy.float.stochastic_rounding(out_weight, weight.dtype, seed=comfy.utils.string_to_seed(key))

if return_weight:

return out_weight

elif inplace_update:

comfy.utils.copy_to_param(self.model, key, out_weight)

else:

comfy.utils.set_attr_param(self.model, key, out_weight)

else:

return set_func(out_weight, inplace_update=inplace_update, seed=comfy.utils.string_to_seed(key), return_weight=return_weight)

def pin_weight_to_device(self, key):

weight, set_func, convert_func = get_key_weight(self.model, key)

if comfy.model_management.pin_memory(weight):

self.pinned.add(key)

def unpin_weight(self, key):