if torch.compiler.is_compiling():

return

if key_heads != value_heads:

raise ValueError(f"Key/value head count mismatch for GQA: {key_heads} != {value_heads}")

if query_heads == key_heads:

return 1

if query_heads % key_heads != 0:

raise ValueError(f"Query heads must be divisible by key/value heads for GQA: {query_heads} vs {key_heads}")

n_rep = gqa_repeat_factor(query_heads, k.shape[head_dim], v.shape[head_dim])

if n_rep > 1:

k = k.repeat_interleave(n_rep, dim=head_dim)

v = v.repeat_interleave(n_rep, dim=head_dim)

attn_mask = args[0] if len(args) > 0 else kwargs.get("attn_mask")

if kwargs.get("enable_gqa", False) and attn_mask is not None:

k, v = repeat_kv_for_gqa(k, v, q.shape[-3], -3)

kwargs["enable_gqa"] = False

if q.nelement() < 1024 * 128: # arbitrary number, for small inputs cudnn attention seems slower

return torch.nn.functional.scaled_dot_product_attention(q, k, v, *args, **kwargs)

attn_mask = args[0] if len(args) > 0 else kwargs.get("attn_mask")

if kwargs.get("enable_gqa", False) and attn_mask is not None and not comfy.model_management.is_nvidia():

k, v = repeat_kv_for_gqa(k, v, q.shape[-3], -3)

kwargs["enable_gqa"] = False

with sdpa_kernel(SDPA_BACKEND_PRIORITY, set_priority=True):

if kwargs.get("enable_gqa", False) and attn_mask is not None and q.shape[-3] != k.shape[-3]:

dropout_p = args[1] if len(args) > 1 else kwargs.get("dropout_p", 0.0)

is_causal = args[2] if len(args) > 2 else kwargs.get("is_causal", False)

params = torch.backends.cuda.SDPAParams(q, k, v, attn_mask, dropout_p, is_causal, True)

supports_native_gqa = (

torch.backends.cuda.can_use_flash_attention(params)

or torch.backends.cuda.can_use_cudnn_attention(params)

or torch.backends.cuda.can_use_efficient_attention(params)

)

if not supports_native_gqa:

k, v = repeat_kv_for_gqa(k, v, q.shape[-3], -3)

kwargs["enable_gqa"] = False

for param_key in param_keys:

param = getattr(s, param_key, None)

if param is not None and getattr(param, "is_meta", False):

offload_stream = None

cast_buffer = None

cast_buffer_offset = 0

if return_faulted:

fully_faulted = all(not getattr(s, param_key + "_function", []) for s in comfy_modules for param_key in ("weight", "bias"))

def ensure_offload_stream(module, required_size, check_largest):

nonlocal offload_stream

nonlocal cast_buffer

if offload_stream is None:

offload_stream = comfy.model_management.get_offload_stream(device)

if offload_stream is None or not check_largest or len(comfy_modules) != 1:

return

current_size = 0 if cast_buffer is None else cast_buffer.size()

if current_size < required_size and module is comfy.model_management.LARGEST_AIMDO_CASTED_WEIGHT[0]:

offload_stream = comfy.model_management.get_offload_stream(device)

cast_buffer = None

if required_size > comfy.model_management.LARGEST_AIMDO_CASTED_WEIGHT[1]:

comfy.model_management.LARGEST_AIMDO_CASTED_WEIGHT = (module, required_size)

def get_cast_buffer(buffer_size):

nonlocal offload_stream

nonlocal cast_buffer

nonlocal cast_buffer_offset

if buffer_size == 0:

return None

if offload_stream is None:

return torch.empty((buffer_size,), dtype=torch.uint8, device=device)

cast_buffer = comfy.model_management.get_aimdo_cast_buffer(offload_stream, device)

buffer = comfy_aimdo.torch.aimdo_to_tensor(cast_buffer.get(buffer_size, cast_buffer_offset), device)

cast_buffer_offset += buffer_size

return buffer

for s in comfy_modules:

signature = comfy_aimdo.model_vbar.vbar_fault(s._v)

resident = comfy_aimdo.model_vbar.vbar_signature_compare(signature, s._v_signature)

if return_faulted and (signature is None or not resident):

fully_faulted = False

prefetch = {

"signature": signature,

"resident": resident,

}

if resident:

s._prefetch = prefetch

continue

materialize_meta_param(s, ["weight", "bias"])

xfer_dest = comfy_aimdo.torch.aimdo_to_tensor(s._v, device) if signature is not None else None

cast_geometry = comfy.memory_management.tensors_to_geometries([ s.weight, s.bias ])

cast_dest = None

needs_cast = False

xfer_source = [ s.weight, s.bias ]

fast_disk = s._pin_state["fast_disk"]

subset = "weights-fast" if fast_disk else "weights"

pin = comfy.pinned_memory.get_pin(s, subset=subset)

if pin is None and not fast_disk:

loaded_pin = comfy.pinned_memory.get_pin(s, subset="weights-loaded")

if loaded_pin is not None or signature is not None:

subset = "weights-loaded"

pin = loaded_pin

if pin is not None:

xfer_source = [ pin ]

for data, geometry in zip([ s.weight, s.bias ], cast_geometry):

if data is None:

continue

if data.dtype != geometry.dtype:

needs_cast = True

cast_dest = xfer_dest

xfer_dest = None

break

dest_size = comfy.memory_management.vram_aligned_size(xfer_source)

ensure_offload_stream(s, dest_size if xfer_dest is None else 0, True)

if xfer_dest is None:

xfer_dest = get_cast_buffer(dest_size)

def cast_maybe_lowvram_patch(xfer_source, xfer_dest, stream, xfer_dest2=None):

if xfer_source is not None:

if getattr(xfer_source, "is_lowvram_patch", False):

if xfer_dest is not None:

xfer_source.prepare(xfer_dest, stream, copy=True, commit=False)

xfer_source = [ xfer_dest ]

xfer_dest = xfer_dest2

xfer_dest2 = None

elif xfer_dest2 is not None:

xfer_source.prepare(xfer_dest2, stream, copy=True, commit=False)

return

else:

return

comfy.model_management.cast_to_gathered(xfer_source, xfer_dest, non_blocking=non_blocking, stream=stream, r2=xfer_dest2)

def handle_pin(m, pin, source, dest, subset="weights", size=None):

if pin is not None:

cast_maybe_lowvram_patch([pin], dest, offload_stream)

return

if signature is None or not fast_disk or args.high_ram:

comfy.pinned_memory.pin_memory(m, subset=subset, size=size)

pin = comfy.pinned_memory.get_pin(m, subset=subset)

cast_maybe_lowvram_patch(source, pin, offload_stream, xfer_dest2=dest)

handle_pin(s, pin, xfer_source, xfer_dest, subset=subset, size=dest_size)

for param_key in ("weight", "bias"):

lowvram_source = getattr(s, param_key + "_lowvram_function", None)

if lowvram_source is not None:

ensure_offload_stream(s, cast_buffer_offset, False)

lowvram_size = lowvram_source.memory_required()

lowvram_dest = get_cast_buffer(lowvram_size)

lowvram_source.prepare(lowvram_dest, None, copy=False, commit=True)

subset = "patches-fast" if fast_disk else "patches"

pin = comfy.pinned_memory.get_pin(lowvram_source, subset=subset)

if pin is None and not fast_disk:

loaded_pin = comfy.pinned_memory.get_pin(lowvram_source, subset="patches-loaded")

if loaded_pin is not None:

subset = "patches-loaded"

pin = loaded_pin

elif signature is not None:

subset = "patches-loaded"

handle_pin(lowvram_source, pin, lowvram_source, lowvram_dest, subset=subset, size=lowvram_size)

prefetch["xfer_dest"] = xfer_dest

prefetch["cast_dest"] = cast_dest

prefetch["cast_geometry"] = cast_geometry

prefetch["needs_cast"] = needs_cast

s._prefetch = prefetch

if return_faulted:

return offload_stream, fully_faulted

prefetch = getattr(s, "_prefetch", None)

if prefetch["resident"]:

weight = s._v_weight

bias = s._v_bias

else:

xfer_dest = prefetch["xfer_dest"]

if prefetch["needs_cast"]:

cast_dest = prefetch["cast_dest"] if prefetch["cast_dest"] is not None else torch.empty((comfy.memory_management.vram_aligned_size(prefetch["cast_geometry"]),), dtype=torch.uint8, device=device)

for pre_cast, post_cast in zip(comfy.memory_management.interpret_gathered_like([s.weight, s.bias ], xfer_dest),

comfy.memory_management.interpret_gathered_like(prefetch["cast_geometry"], cast_dest)):

if post_cast is not None:

post_cast.copy_(pre_cast)

xfer_dest = cast_dest

params = comfy.memory_management.interpret_gathered_like(prefetch["cast_geometry"], xfer_dest)

weight = params[0]

bias = params[1]

if prefetch["signature"] is not None:

s._v_weight = weight

s._v_bias = bias

s._v_signature = prefetch["signature"]

def post_cast(s, param_key, x, dtype, resident, update_weight):

lowvram_fn = getattr(s, param_key + "_lowvram_function", None)

fns = getattr(s, param_key + "_function", [])

if x is None:

return None

orig = x

def to_dequant(tensor, dtype):

tensor = tensor.to(dtype=dtype)

if isinstance(tensor, QuantizedTensor):

tensor = tensor.dequantize()

return tensor

if (return_weights and orig.dtype != dtype) or len(fns) > 0:

x = to_dequant(x, dtype)

if not resident and lowvram_fn is not None:

x = to_dequant(x, dtype if compute_dtype is None else compute_dtype)

x = lowvram_fn(x)

if (want_requant and len(fns) == 0 or update_weight):

seed = comfy.utils.string_to_seed(s.seed_key)

if isinstance(orig, QuantizedTensor):

y = orig.requantize_from_float(x, scale="recalculate", stochastic_rounding=seed)

else:

y = comfy.float.stochastic_rounding(x, orig.dtype, seed=seed)

if want_requant and len(fns) == 0:

x = y

if update_weight:

orig.copy_(y)

for f in fns:

x = f(x)

return x

update_weight = prefetch["signature"] is not None

weight = post_cast(s, "weight", weight, dtype, prefetch["resident"], update_weight)

if bias is not None:

bias = post_cast(s, "bias", bias, bias_dtype, prefetch["resident"], update_weight)

if prefetch["signature"] is not None:

prefetch["resident"] = True

# NOTE: offloadable=False is a legacy mode and if you are a custom node author reading this please pass

# offloadable=True and call uncast_bias_weight() after your last usage of the weight/bias. This

# will add async-offload support to your cast and improve performance.

if input is not None:

if dtype is None:

if isinstance(input, QuantizedTensor):

dtype = input.params.orig_dtype

else:

dtype = input.dtype

if bias_dtype is None:

bias_dtype = dtype

if device is None:

device = input.device

def format_return(result, offloadable):

weight, bias, offload_stream = result

return (weight, bias, offload_stream) if offloadable else (weight, bias)

non_blocking = comfy.model_management.device_supports_non_blocking(device)

if hasattr(s, "_v") and comfy.model_management.is_device_cpu(device):

#vbar doesn't support CPU weights, but some custom nodes have weird paths

#that might switch the layer to the CPU and expect it to work. We have to take

#a clone conservatively as we are mmapped and some SFT files are packed misaligned

#If you are a custom node author reading this, please move your layer to the GPU

#or declare your ModelPatcher as CPU in the first place.

materialize_meta_param(s, ["weight", "bias"])

weight = s.weight.to(dtype=dtype, copy=True)

if isinstance(weight, QuantizedTensor):

weight = weight.dequantize()

bias = s.bias.to(dtype=bias_dtype, copy=True) if s.bias is not None else None

return format_return((weight, bias, (None, None, None)), offloadable)

elif hasattr(s, "_v") and s.weight.device != device:

prefetched = hasattr(s, "_prefetch")

offload_stream = None

offload_device = None

if not prefetched:

offload_stream = cast_modules_with_vbar([s], dtype, device, bias_dtype, non_blocking)

comfy.model_management.sync_stream(device, offload_stream)

weight, bias = resolve_cast_module_with_vbar(s, dtype, device, bias_dtype, compute_dtype, want_requant)

if not prefetched:

if getattr(s, "_prefetch")["signature"] is not None:

offload_device = device

for param_key in ("weight", "bias"):

lowvram_fn = getattr(s, param_key + "_lowvram_function", None)

if lowvram_fn is not None:

lowvram_fn.clear_prepared()

delattr(s, "_prefetch")

return format_return((weight, bias, (offload_stream, offload_device, None)), offloadable)

if offloadable and (device != s.weight.device or

(s.bias is not None and device != s.bias.device)):

offload_stream = comfy.model_management.get_offload_stream(device)

else:

offload_stream = None

bias = None

weight = None

if offload_stream is not None and not args.cuda_malloc:

cast_buffer_size = comfy.memory_management.vram_aligned_size([ s.weight, s.bias ])

cast_buffer = comfy.model_management.get_cast_buffer(offload_stream, device, cast_buffer_size, s)

#The streams can be uneven in buffer capability and reject us. Retry to get the other stream

if cast_buffer is None:

offload_stream = comfy.model_management.get_offload_stream(device)

cast_buffer = comfy.model_management.get_cast_buffer(offload_stream, device, cast_buffer_size, s)

params = comfy.memory_management.interpret_gathered_like([ s.weight, s.bias ], cast_buffer)

weight = params[0]

bias = params[1]

weight_has_function = len(s.weight_function) > 0

bias_has_function = len(s.bias_function) > 0

weight = comfy.model_management.cast_to(s.weight, None, device, non_blocking=non_blocking, copy=weight_has_function, stream=offload_stream, r=weight)

if s.bias is not None:

bias = comfy.model_management.cast_to(s.bias, None, device, non_blocking=non_blocking, copy=bias_has_function, stream=offload_stream, r=bias)

comfy.model_management.sync_stream(device, offload_stream)

bias_a = bias

weight_a = weight

if s.bias is not None:

bias = bias.to(dtype=bias_dtype)

for f in s.bias_function:

bias = f(bias)

if weight_has_function or weight.dtype != dtype:

weight = weight.to(dtype=dtype)

if isinstance(weight, QuantizedTensor):

weight = weight.dequantize()

for f in s.weight_function:

weight = f(weight)

if offload_stream is None:

return

os, weight_a, bias_a = offload_stream

device=None

#FIXME: This is really bad RTTI

if weight_a is not None and not isinstance(weight_a, torch.Tensor):

comfy_aimdo.model_vbar.vbar_unpin(s._v)

device = weight_a

if os is None:

return

if device is None:

if weight_a is not None:

device = weight_a.device

else:

if bias_a is None:

return

device = bias_a.device

# When initialized with no arguments or the first is None, the context

# will return the tuple (None, None).

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

self.slf = args[0] if len(args) else None

self.state = (None, None) if self.slf is None else cast_bias_weight(*args, **kwargs)

def __enter__(self):

result = self.state

if len(result) < 3 or result[2] is None:

# Not offloaded, immediately drop references.

self.state = self.slf = None

return result[:2]

def __exit__(self, *_args) -> None:

if self.slf is None:

return

slf, state = self.slf, self.state

self.state = self.slf = None

@staticmethod

def _zero_init_parameter(module, name):

param = getattr(module, name)

device = None if getattr(param, "is_meta", False) else param.device

setattr(module, name, torch.nn.Parameter(torch.zeros(param.shape, device=device, dtype=param.dtype), requires_grad=False))

@staticmethod

def _lazy_load_from_state_dict(module, state_dict, prefix, local_metadata,

missing_keys, unexpected_keys, weight_shape,

bias_shape=None):

assign_to_params_buffers = local_metadata.get("assign_to_params_buffers", False)

prefix_len = len(prefix)

for k, v in state_dict.items():

key = k[prefix_len:]

if key == "weight":

if not assign_to_params_buffers:

v = v.clone()

module.weight = torch.nn.Parameter(v, requires_grad=False)

elif bias_shape is not None and key == "bias" and v is not None:

if not assign_to_params_buffers:

v = v.clone()

module.bias = torch.nn.Parameter(v, requires_grad=False)

else:

unexpected_keys.append(k)

if module.weight is None:

module.weight = torch.nn.Parameter(torch.zeros(weight_shape), requires_grad=False)

missing_keys.append(prefix + "weight")

if bias_shape is not None and module.bias is None and getattr(module, "comfy_need_lazy_init_bias", False):

module.bias = torch.nn.Parameter(torch.zeros(bias_shape), requires_grad=False)

missing_keys.append(prefix + "bias")

class Linear(torch.nn.Linear, CastWeightBiasOp):

def __init__(self, in_features, out_features, bias=True, device=None, dtype=None):

# don't trust subclasses that BYO state dict loader to call us.

if (not comfy.memory_management.aimdo_enabled

or type(self)._load_from_state_dict is not disable_weight_init.Linear._load_from_state_dict):

super().__init__(in_features, out_features, bias, device, dtype)

return

# Issue is with `torch.empty` still reserving the full memory for the layer.

# Windows doesn't over-commit memory so without this, We are momentarily commit

# charged for the weight even though we might zero-copy it when we load the

# state dict. If the commit charge exceeds the ceiling we can destabilize the

# system.

torch.nn.Module.__init__(self)

self.in_features = in_features

self.out_features = out_features

self.weight = None

self.bias = None

self.comfy_need_lazy_init_bias=bias

self.weight_comfy_model_dtype = dtype

self.bias_comfy_model_dtype = dtype

def _load_from_state_dict(self, state_dict, prefix, local_metadata,

strict, missing_keys, unexpected_keys, error_msgs):

if (not comfy.memory_management.aimdo_enabled

or type(self)._load_from_state_dict is not disable_weight_init.Linear._load_from_state_dict):

return super()._load_from_state_dict(state_dict, prefix, local_metadata, strict,

missing_keys, unexpected_keys, error_msgs)

disable_weight_init._lazy_load_from_state_dict(

self,

state_dict,

prefix,

local_metadata,

missing_keys,

unexpected_keys,

weight_shape=(self.in_features, self.out_features),

bias_shape=(self.out_features,),

)

def reset_parameters(self):

return None

def forward_comfy_cast_weights(self, input):

with CastBiasWeightContext(self, input, offloadable=True) as (weight, bias):

return torch.nn.functional.linear(input, weight, bias)

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

run_every_op()

if self.comfy_cast_weights or len(self.weight_function) > 0 or len(self.bias_function) > 0:

return self.forward_comfy_cast_weights(*args, **kwargs)

else:

return super().forward(*args, **kwargs)

class Conv1d(torch.nn.Conv1d, CastWeightBiasOp):

def reset_parameters(self):

return None

def forward_comfy_cast_weights(self, input):

with CastBiasWeightContext(self, input, offloadable=True) as (weight, bias):

return self._conv_forward(input, weight, bias)

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

run_every_op()

if self.comfy_cast_weights or len(self.weight_function) > 0 or len(self.bias_function) > 0:

return self.forward_comfy_cast_weights(*args, **kwargs)

else:

return super().forward(*args, **kwargs)

class Conv2d(torch.nn.Conv2d, CastWeightBiasOp):

def reset_parameters(self):

return None

def forward_comfy_cast_weights(self, input):

with CastBiasWeightContext(self, input, offloadable=True) as (weight, bias):

return self._conv_forward(input, weight, bias)

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

run_every_op()

if self.comfy_cast_weights or len(self.weight_function) > 0 or len(self.bias_function) > 0:

return self.forward_comfy_cast_weights(*args, **kwargs)

else:

return super().forward(*args, **kwargs)

class Conv3d(torch.nn.Conv3d, CastWeightBiasOp):

def reset_parameters(self):

return None

def _conv_forward(self, input, weight, bias, autopad=None, *args, **kwargs):

if autopad == "causal_zero":

weight = weight[:, :, -input.shape[2]:, :, :]

if NVIDIA_MEMORY_CONV_BUG_WORKAROUND and weight.dtype in (torch.float16, torch.bfloat16):

out = torch.cudnn_convolution(input, weight, self.padding, self.stride, self.dilation, self.groups, benchmark=False, deterministic=False, allow_tf32=True)

if bias is not None:

out += bias.reshape((1, -1) + (1,) * (out.ndim - 2))

return out

else:

return super()._conv_forward(input, weight, bias, *args, **kwargs)

def forward_comfy_cast_weights(self, input, autopad=None):

with CastBiasWeightContext(self, input, offloadable=True) as (weight, bias):

return self._conv_forward(input, weight, bias, autopad=autopad)

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

run_every_op()

if self.comfy_cast_weights or len(self.weight_function) > 0 or len(self.bias_function) > 0 or "autopad" in kwargs:

return self.forward_comfy_cast_weights(*args, **kwargs)

else:

return super().forward(*args, **kwargs)

class GroupNorm(torch.nn.GroupNorm, CastWeightBiasOp):

def reset_parameters(self):

return None

def forward_comfy_cast_weights(self, input):

with CastBiasWeightContext(self, input, offloadable=True) as (weight, bias):

return torch.nn.functional.group_norm(input, self.num_groups, weight, bias, self.eps)

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

run_every_op()

if self.comfy_cast_weights or len(self.weight_function) > 0 or len(self.bias_function) > 0:

return self.forward_comfy_cast_weights(*args, **kwargs)

else:

return super().forward(*args, **kwargs)

class BatchNorm2d(torch.nn.BatchNorm2d, CastWeightBiasOp):

def reset_parameters(self):

return None

def forward_comfy_cast_weights(self, input):

with CastBiasWeightContext(self, input, offloadable=True) as (weight, bias):

running_mean = self.running_mean.to(device=input.device, dtype=weight.dtype) if self.running_mean is not None else None

running_var = self.running_var.to(device=input.device, dtype=weight.dtype) if self.running_var is not None else None

return torch.nn.functional.batch_norm(input, running_mean, running_var, weight, bias, self.training, self.momentum, self.eps)

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

run_every_op()

if self.comfy_cast_weights or len(self.weight_function) > 0 or len(self.bias_function) > 0:

return self.forward_comfy_cast_weights(*args, **kwargs)

else:

return super().forward(*args, **kwargs)

class LayerNorm(torch.nn.LayerNorm, CastWeightBiasOp):

def reset_parameters(self):

return None

def forward_comfy_cast_weights(self, input):

with CastBiasWeightContext(self if self.weight is not None else None, input, offloadable=True) as (weight, bias):

return torch.nn.functional.layer_norm(input, self.normalized_shape, weight, bias, self.eps)

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

run_every_op()

if self.comfy_cast_weights or len(self.weight_function) > 0 or len(self.bias_function) > 0:

return self.forward_comfy_cast_weights(*args, **kwargs)

else:

return super().forward(*args, **kwargs)

class RMSNorm(torch.nn.RMSNorm, CastWeightBiasOp):

def reset_parameters(self):

self.bias = None

return None

def forward_comfy_cast_weights(self, input):

with CastBiasWeightContext(self if self.weight is not None else None, input, offloadable=True) as (weight, bias):

return torch.nn.functional.rms_norm(input, self.normalized_shape, weight, self.eps)

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

run_every_op()

if self.comfy_cast_weights or len(self.weight_function) > 0 or len(self.bias_function) > 0:

return self.forward_comfy_cast_weights(*args, **kwargs)

else:

return super().forward(*args, **kwargs)

class ConvTranspose2d(torch.nn.ConvTranspose2d, CastWeightBiasOp):

def reset_parameters(self):

return None

def forward_comfy_cast_weights(self, input, output_size=None):

num_spatial_dims = 2

output_padding = self._output_padding(

input, output_size, self.stride, self.padding, self.kernel_size,

num_spatial_dims, self.dilation)

with CastBiasWeightContext(self, input, offloadable=True) as (weight, bias):

return torch.nn.functional.conv_transpose2d(

input, weight, bias, self.stride, self.padding,

output_padding, self.groups, self.dilation)

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

run_every_op()

if self.comfy_cast_weights or len(self.weight_function) > 0 or len(self.bias_function) > 0:

return self.forward_comfy_cast_weights(*args, **kwargs)

else:

return super().forward(*args, **kwargs)

class ConvTranspose1d(torch.nn.ConvTranspose1d, CastWeightBiasOp):

def reset_parameters(self):

return None

def forward_comfy_cast_weights(self, input, output_size=None):

num_spatial_dims = 1

output_padding = self._output_padding(

input, output_size, self.stride, self.padding, self.kernel_size,

num_spatial_dims, self.dilation)

with CastBiasWeightContext(self, input, offloadable=True) as (weight, bias):

return torch.nn.functional.conv_transpose1d(

input, weight, bias, self.stride, self.padding,

output_padding, self.groups, self.dilation)

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

run_every_op()

if self.comfy_cast_weights or len(self.weight_function) > 0 or len(self.bias_function) > 0:

return self.forward_comfy_cast_weights(*args, **kwargs)

else:

return super().forward(*args, **kwargs)

class Embedding(torch.nn.Embedding, CastWeightBiasOp):

def __init__(self, num_embeddings, embedding_dim, padding_idx=None, max_norm=None,

norm_type=2.0, scale_grad_by_freq=False, sparse=False, _weight=None,

_freeze=False, device=None, dtype=None):

# don't trust subclasses that BYO state dict loader to call us.

if (not comfy.memory_management.aimdo_enabled

or type(self)._load_from_state_dict is not disable_weight_init.Embedding._load_from_state_dict):

super().__init__(num_embeddings, embedding_dim, padding_idx, max_norm,

norm_type, scale_grad_by_freq, sparse, _weight,

_freeze, device, dtype)

return

torch.nn.Module.__init__(self)

self.num_embeddings = num_embeddings

self.embedding_dim = embedding_dim

self.padding_idx = padding_idx

self.max_norm = max_norm

self.norm_type = norm_type

self.scale_grad_by_freq = scale_grad_by_freq

self.sparse = sparse

# Keep shape/dtype visible for module introspection without reserving storage.

embedding_dtype = dtype if dtype is not None else torch.get_default_dtype()

self.weight = torch.nn.Parameter(

torch.empty((num_embeddings, embedding_dim), device="meta", dtype=embedding_dtype),

requires_grad=False,

)

self.bias = None

self.weight_comfy_model_dtype = dtype

def _load_from_state_dict(self, state_dict, prefix, local_metadata,

strict, missing_keys, unexpected_keys, error_msgs):

if (not comfy.memory_management.aimdo_enabled

or type(self)._load_from_state_dict is not disable_weight_init.Embedding._load_from_state_dict):

return super()._load_from_state_dict(state_dict, prefix, local_metadata, strict,

missing_keys, unexpected_keys, error_msgs)

disable_weight_init._lazy_load_from_state_dict(

self,

state_dict,

prefix,

local_metadata,

missing_keys,

unexpected_keys,

weight_shape=(self.num_embeddings, self.embedding_dim),

)

def reset_parameters(self):

self.bias = None

return None

def forward_comfy_cast_weights(self, input, out_dtype=None):

output_dtype = out_dtype

if self.weight.dtype == torch.float16 or self.weight.dtype == torch.bfloat16:

out_dtype = None

with CastBiasWeightContext(self, device=input.device, dtype=out_dtype, offloadable=True) as (weight, bias):

return torch.nn.functional.embedding(input, weight, self.padding_idx, self.max_norm, self.norm_type, self.scale_grad_by_freq, self.sparse).to(dtype=output_dtype)

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

run_every_op()

if self.comfy_cast_weights or len(self.weight_function) > 0 or len(self.bias_function) > 0:

return self.forward_comfy_cast_weights(*args, **kwargs)

else:

if "out_dtype" in kwargs:

kwargs.pop("out_dtype")

return super().forward(*args, **kwargs)

@classmethod

def conv_nd(s, dims, *args, **kwargs):

if dims == 2:

return s.Conv2d(*args, **kwargs)

elif dims == 3:

return s.Conv3d(*args, **kwargs)

else:

class Linear(disable_weight_init.Linear):

comfy_cast_weights = True

class Conv1d(disable_weight_init.Conv1d):

comfy_cast_weights = True

class Conv2d(disable_weight_init.Conv2d):

comfy_cast_weights = True

class Conv3d(disable_weight_init.Conv3d):

comfy_cast_weights = True

class BatchNorm2d(disable_weight_init.BatchNorm2d):

comfy_cast_weights = True

class GroupNorm(disable_weight_init.GroupNorm):

comfy_cast_weights = True

class LayerNorm(disable_weight_init.LayerNorm):

comfy_cast_weights = True

class ConvTranspose2d(disable_weight_init.ConvTranspose2d):

comfy_cast_weights = True

class ConvTranspose1d(disable_weight_init.ConvTranspose1d):

comfy_cast_weights = True

class RMSNorm(disable_weight_init.RMSNorm):

comfy_cast_weights = True

class Embedding(disable_weight_init.Embedding):